feat: Fused rope fp8 quantize kernel for MLA

[yzh119]

📌 Description

Fusing RoPE + fp8 quantization kernel to prepare input for fp8 mla kernel.
Reference:
https://github.com/NVIDIA/TensorRT-LLM/blob/0df758ec9f8409410bac8b60d117374054391c2d/cpp/tensorrt_llm/kernels/mlaKernels.cu#L358

🔍 Related Issues

🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit (or used your preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

• Tests have been added or updated as needed.
• All tests are passing (unittest, etc.).

Reviewer Notes

[gemini-code-assist]

Summary of Changes

Hello @yzh119, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a significant performance optimization by implementing a new fused CUDA kernel that combines the application of Rotary Positional Embeddings (RoPE) with FP8 quantization. This specialized kernel is designed to efficiently prepare query and key/value tensors for subsequent FP8 Multi-Layer Attention (MLA) operations, aiming to reduce computational overhead and memory bandwidth requirements in large language models.

Highlights

• Fused RoPE and FP8 Quantization: Introduces a new fused kernel that combines Rotary Positional Embedding (RoPE) application with FP8 quantization. This is specifically designed to prepare inputs for FP8 Multi-Layer Attention (MLA) kernels, optimizing the data pipeline for efficient model inference.
• Performance Optimization: The primary goal of this change is to improve computational efficiency and reduce memory bandwidth requirements by fusing two sequential operations (RoPE and quantization) into a single, optimized CUDA kernel.
• CUDA Kernel Implementation: A new CUDA kernel, MLARopeQuantizeKernel, is implemented to perform the fused operation directly on the GPU. It handles both the RoPE and non-RoPE components of query and key/value tensors, supporting both interleaved and non-interleaved RoPE application, and applies quantization scales to output FP8 tensors.
• PyTorch Integration: The new fused operation is exposed to the PyTorch framework through a mla_rope_quantize function, allowing it to be seamlessly integrated and called from Python code within the FlashInfer library.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments or fill out our survey to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a new fused RoPE and FP8 quantization kernel for MLA, which is a great addition for performance. However, the current implementation has several critical issues that will prevent it from compiling and running correctly. These include undefined variables, incorrect function calls, and flawed logic within the new CUDA kernel. I've provided detailed feedback and suggestions to address these points. Once fixed, this will be a solid contribution.

[gemini-code-assist]

The variable q is not defined in this scope. This appears to be a copy-paste error and will cause a compilation failure. You should use one of the input tensors, such as q_rope_in, to get the device.

  const c10::cuda::OptionalCUDAGuard device_guard(q_rope_in.device());

[gemini-code-assist]

This is calling BatchQKApplyRotaryPosIdsCosSinCache, but the function signature and arguments do not match. It seems you intended to call the new flashinfer::MLARopeQuantize function. This will cause a compilation error.

        cudaError_t status = flashinfer::MLARopeQuantize(
            static_cast<c_type*>(q_rope_in.data_ptr()), static_cast<c_type*>(k_rope_in.data_ptr()),
            static_cast<c_type*>(q_nope_in.data_ptr()), static_cast<c_type*>(k_nope_in.data_ptr()),
            static_cast<c_quant_type*>(q_rope_out.data_ptr()),
            static_cast<c_quant_type*>(k_rope_out.data_ptr()),
            static_cast<c_quant_type*>(q_nope_out.data_ptr()),
            static_cast<c_quant_type*>(k_nope_out.data_ptr()),
            static_cast<c_idtype*>(pos_ids.data_ptr()), nnz, num_heads, q_rope_in_stride_n,
            q_rope_in_stride_h, q_nope_in_stride_n, q_nope_in_stride_h, q_rope_out_stride_n,
            q_rope_out_stride_h, q_nope_out_stride_n, q_nope_out_stride_h, k_rope_in_stride,
            k_nope_in_stride, k_rope_out_stride, k_nope_out_stride, quant_scale_q, quant_scale_kv,
            interleave, stream);
        TORCH_CHECK(status == cudaSuccess,
                    "MLARopeQuantize failed with error code " +
                        std::string(cudaGetErrorString(status)));

[gemini-code-assist]

The variable bdy is not defined within the kernel's scope. This will cause a compilation error. You should use blockDim.y instead, as it holds the value of bdy passed during kernel launch.

  if (bx * blockDim.y + ty < nnz) {

[gemini-code-assist]

The cos_sin_cache variable is used here but is not defined in the MLARopeQuantizeKernel's signature. It needs to be passed as a kernel argument. This will also require updating the calling function MLARopeQuantize to accept and pass cos_sin_cache.

[gemini-code-assist]

The quantization logic here is incorrect and will lead to incorrect results and race conditions.

1. The loop at lines 402-404 writes to q_rope_out_ptr[i] without a thread-specific offset, causing multiple threads in a warp to write to the same memory locations.
2. The cast_store on line 405 overwrites the result of the loop without applying the quant_scale_q.

The correct approach is to scale the q_rope_vec vector and then store it. This same issue exists for k_rope_vec (lines 422-426), k_nope_vec (lines 439-443), and q_nope_vec (lines 457-461).

#pragma unroll
      for (uint32_t i = 0; i < vec_size; ++i) {
        q_rope_vec[i] *= quant_scale_q;
      }
      q_rope_vec.cast_store(q_rope_out_ptr + tx * vec_size);

[gemini-code-assist]

std::max requires two arguments, but only one is provided. This will cause a compilation error. Please correct the expression to be a valid constexpr value.

    constexpr uint32_t vec_size = 16 / sizeof(DType);

[gemini-code-assist]

The template arguments for MLARopeQuantizeKernel are incorrect. rotary_dim is a runtime value, not a template parameter. Additionally, the QuantType template parameter is missing. This will cause a compilation error.

    auto kernel = MLARopeQuantizeKernel<INTERLEAVE, vec_size, bdx, DType, IdType, QuantType>;

[gemini-code-assist]

There is a missing semicolon between the dim3 declarations, which will cause a compilation error.

    dim3 nblks(nblks_x, num_heads + 8 + 1 + num_heads * 8); dim3 nthrs(bdx, bdy);

[gemini-code-assist]

The tensor k_rope_out is 2D with shape (nnz, 64), as specified by the check on line 300. Its second dimension is the rotary dimension (64), not num_heads. This check incorrectly constrains num_heads to be 64. Since k_rope_out does not have a head dimension, this check should be removed.

[pavanimajety]

Nit: Are you also supporting just applying rope and not quantizing? I think this function must make the quantize_dtype argument not optional or explicitly convey through function name that it is rope_quantize_to_fp8

[cyx-6]

updated

[fzyzcjy]

Looking forward to this