Fix RoPE kernel illegal memory access for MLA models (separate q/k head strides)

[he-yufeng]

Motivation

The fused RoPE kernel crashes with illegal memory access (IMA) on GLM5 FP8 because it uses a single head_stride_bytes (derived from q) for both q and k heads. MLA models have different head dimensions for q and k:

• q shape: [8, 64, 64] → head stride 256 elements
• k shape: [8, 1, 64] → head stride 64 elements

Using q's stride (256) to index into k's memory causes out-of-bounds access.

Fix

Split head_stride_bytes into q_head_stride_bytes and k_head_stride_bytes in FusedRopeParams. The kernel selects the correct stride based on whether it's processing a q or k head. The pre-offset trick for k_ptr is updated to use k's head stride consistently.

Both the non-fused (fused_rope_kernel) and fused (fused_rope_store_kernel) paths are updated. The fused path's RuntimeCheck is also split to validate both strides independently.

Fixes #20559

[gemini-code-assist]

Summary of Changes

Hello, 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 resolves a critical illegal memory access bug in the fused RoPE kernel that occurred in MLA models with asymmetric query and key head dimensions. The fix involves refactoring the kernel to use separate head stride parameters for query and key, ensuring correct memory access and improving the stability and compatibility of the RoPE implementation across diverse model architectures.

Highlights

• Separate Q/K Head Strides: Introduced distinct q_head_stride_bytes and k_head_stride_bytes in FusedRopeParams to correctly handle models with differing query and key head dimensions, preventing illegal memory access.
• Kernel Logic Update: Modified both fused_rope_kernel and fused_rope_store_kernel to dynamically select the appropriate head stride (query or key) based on the tensor being processed.
• Offset Calculation Correction: Updated the pre-offset calculation for key (k) pointers and cache writes to consistently use k_head_stride_bytes, ensuring accurate memory addressing.
• Runtime Validation Enhancement: Split the RuntimeCheck in the fused path to independently validate rope_dim against both q_head_stride and k_head_stride, improving error detection for misconfigured models.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• python/sglang/jit_kernel/csrc/elementwise/rope.cuh
  • Modified the FusedRopeParams struct to include separate q_head_stride_bytes and k_head_stride_bytes.
  • Updated the fused_rope_kernel and fused_rope_store_kernel to unpack and utilize the new separate head stride parameters.
  • Implemented conditional logic within the kernels to select the correct head stride based on whether a query or key head is being processed.
  • Adjusted offset calculations for key pointers and cache writes to correctly use k_head_stride_bytes.
  • Refactored the FusedRopeKernel class to use separate symbolic sizes for query and key head strides in tensor matching and parameter calculation.
  • Updated RuntimeCheck assertions to validate rope_dim against both q_head_stride and k_head_stride independently.

Activity

• No specific activity (comments, reviews, progress updates) has been recorded for this pull request yet.

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 by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

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 pull request 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. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

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 addresses a critical illegal memory access bug in the RoPE kernel for MLA models. The root cause was the use of a single head stride for both Q and K tensors, which can have different head dimensions in MLA models. The fix correctly introduces separate q_head_stride_bytes and k_head_stride_bytes and applies them consistently throughout the non-fused and fused RoPE kernels, including pointer offsets and cache access. The changes are logical and effectively resolve the reported issue. I have one minor suggestion to improve code consistency.

[gemini-code-assist]

For consistency and to avoid code duplication, it's better to reuse the k_offset variable, which is calculated just a few lines above with the correct k_head_stride_bytes. This also makes the intent clearer, as both k_ptr and k_cache are offset by the same amount for the same reason.

        .k_cache = pointer::offset(k_cache.data_ptr(), -k_offset),

[DarkSharpness]

hi. my question is:

1. do you have a command to quickly reproduce the IMA? i'm still not sure about the cause of [Bug] [ROPE] IMA and ivalid logic in rope.cuh #20559 . please refer to my comments in that issue (in that case, i do not think stride is used. k head count is always 1)
2. could you please provide benchmark results before & after this PR? you may run rope benchmark in jit kernel. thanks

[he-yufeng]

Thanks @DarkSharpness — you're right that with num_kv_heads == 1, the pre-offset and head_id * head_stride_bytes cancel out, so the stride mismatch wouldn't actually cause an OOB access for k.

I don't have access to the GLM5 model or the specific hardware to reproduce the IMA. My analysis was based on the stride values in the issue (q head stride 256 vs k head stride 64), but I missed the cancellation when there's only 1 k head.

Given your comment on #20559 that you plan to write a new kernel with broader stride support, should I close this PR? The split-stride change is still correct in principle (handles the general case where num_kv_heads > 1 with different strides), but if the IMA root cause is different, it's not solving the reported problem.

[DarkSharpness]

Thanks @DarkSharpness — you're right that with num_kv_heads == 1, the pre-offset and head_id * head_stride_bytes cancel out, so the stride mismatch wouldn't actually cause an OOB access for k.

I don't have access to the GLM5 model or the specific hardware to reproduce the IMA. My analysis was based on the stride values in the issue (q head stride 256 vs k head stride 64), but I missed the cancellation when there's only 1 k head.

Given your comment on #20559 that you plan to write a new kernel with broader stride support, should I close this PR? The split-stride change is still correct in principle (handles the general case where num_kv_heads > 1 with different strides), but if the IMA root cause is different, it's not solving the reported problem.

Could you please post full benchmark result of this PR? As long as there is no performance drawback, we can simply adopt this PR. Thanks a lot.

We did not support different stride because for q k because typically they should be contiguous in num head dimension.

[he-yufeng]

Thanks! I don't have GPU access to run the benchmark locally. The change only adds one extra field to the params struct and a ternary select in the kernel (load_q ? q_head_stride : k_head_stride), so I'd expect negligible overhead — but I understand you'd want numbers to confirm.

If someone with GPU access could run python sglang/jit_kernel/benchmark/bench_rope.py before and after this commit, that would settle it. Happy to adjust anything based on the results.

[DarkSharpness]

Thanks! I don't have GPU access to run the benchmark locally. The change only adds one extra field to the params struct and a ternary select in the kernel (load_q ? q_head_stride : k_head_stride), so I'd expect negligible overhead — but I understand you'd want numbers to confirm.

If someone with GPU access could run python sglang/jit_kernel/benchmark/bench_rope.py before and after this commit, that would settle it. Happy to adjust anything based on the results.

I can run the benchmark for you some time later. I also expect there should be no much overhead.

[he-yufeng]

Closing this — the IMA root cause has been identified by @xu-yfei in #20559 as a kvcache index OOB in nsa_indexer.py:543 under CP scenarios, not the RoPE stride mismatch. Fix is tracked in #18280.

The stride separation change is still correct in principle but isn't solving the reported problem. Thanks @DarkSharpness for the review.