Fuse writing KV buffer into rope kernel (part 2: srt)

[JeremieMelo]

Motivation

Fuse set_kv_buffer to sgl-kernel rope function, only for trtllm_mha attention

---

(below is from @fzyzcjy)

speed may be suboptimal (I have not done any ncu profile or thorough optimization), but anyway it is faster than non-fused

tests/test_rotary_embedding.py ..................                                                                  [100%]

=================================================== 18 passed in 1.90s ===================================================
bench_rotary_embedding:
    batch_size  seq_len  save_kv_cache  SGL Kernel
0            1        1          False       2.788
1            1        1           True       3.119
2           32        1          False       3.052
3           32        1           True       3.299
4          128        1          False       3.167
5          128        1           True       3.463
6          512        1          False       3.925
7          512        1           True       4.091
8            2      512          False       5.643
9            2      512           True       6.371
10           4     4096          False      46.858
11           4     4096           True      57.892

acc looks good: 20B TP4, reasoning low, gpt_oss.evals gpqa: 55.1%

acc from @BBuf (he checked this agree with main):

| model_name                                              |   (‘metric’, ‘mmlu’) |
|:--------------------------------------------------------|---------------------:|
| o4-mini-with-chat-completion-and-4k-gen_20250811_143850 |                0.866 |

speedup: @BBuf test to be 345->355

[gemini-code-assist]

Summary of Changes

Hello @JeremieMelo, 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!

I've implemented a significant performance optimization by fusing the Key-Value (KV) cache buffer saving operation directly into the Rotary Positional Embedding (RoPE) computation. This change is specifically applied to the trtllm_mha attention mechanism. The primary goal is to reduce overhead and improve efficiency by performing these two related operations within a single, optimized kernel, rather than as separate steps. This involves modifications across the Python and C++ layers of the attention and rotary embedding components, ensuring seamless integration and leveraging CUDA stream synchronization for better performance, especially in CUDA graph capture mode.

Highlights

• KV Buffer and RoPE Fusion: I have fused the set_kv_buffer operation, which is responsible for saving key and value tensors to the KV cache, directly into the rotary positional embedding (RoPE) computation. This optimization specifically targets the trtllm_mha attention mechanism.
• Streamlined KV Cache Updates: The previous explicit calls to set_kv_buffer in the trtllm_mha_backend have been removed. The KV cache update now occurs implicitly as part of the RotaryEmbedding's forward_cuda method, which calls a modified C++ kernel.
• C++ Kernel Enhancements for Fused Operations: The underlying C++ CUDA kernel for applying rotary embeddings (apply_rope_pos_ids_cos_sin_cache) has been enhanced to directly handle the saving of the processed K and V tensors to the KV cache buffers. This includes handling data type conversions, scaling, and optimized memory copies, particularly for CUDA graph capture mode using an alternate stream for overlap.
• Integration into Model Architecture: The GPTAttention and GPTLMHeadModel classes have been updated to pass necessary context, such as the attention layer, forward batch information, and the start_layer for memory pooling, to the rotary embedding function, enabling the fused operation.

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 fuses the set_kv_buffer operation into the RoPE kernel for trtllm_mha attention, which is a good optimization. My review focuses on the correctness of the implementation. I've found a critical issue in the CUDA stream synchronization logic that could lead to race conditions or runtime errors, and a bug in how scaling factors are passed. I've also suggested some code cleanup by removing commented-out code blocks.

[gemini-code-assist]

The stream synchronization logic here is incorrect and will lead to runtime errors or race conditions. cudaStreamWaitEvent is called with a nullptr event, which is invalid. You need to create and use a cudaEvent_t to properly synchronize between the main and alternate streams. The original Python code stream.wait_stream(other_stream) uses an event pool internally, which needs to be replicated here.

Here is a suggestion for a correct implementation:

    if (is_capture_mode && alt_stream_ptr != 0) {
      cudaStream_t alt_stream = reinterpret_cast<cudaStream_t>(alt_stream_ptr);
      cudaStream_t main_stream = stream;
      cudaEvent_t event;
      cudaEventCreateWithFlags(&event, cudaEventDisableTiming);

      // Wait for main stream to complete RoPE
      cudaEventRecord(event, main_stream);
      cudaStreamWaitEvent(alt_stream, event, 0);

      // Copy K on main stream
      k_buffer_ptr.copy_(k_rope, /*non_blocking=*/true);

      // Copy V on alternate stream
      {
        at::cuda::CUDAStreamGuard guard(at::cuda::getStreamFromExternal(alt_stream, device.index()));
        v_buffer_ptr.copy_(v, /*non_blocking=*/true);
      }

      // Main stream waits for alt stream
      cudaEventRecord(event, alt_stream);
      cudaStreamWaitEvent(main_stream, event, 0);
      cudaEventDestroy(event);
    } else {

[gemini-code-assist]

This commented-out code should be removed as its functionality has been fused into the RoPE kernel.

[gemini-code-assist]

This commented-out code should be removed as its functionality has been fused into the RoPE kernel.

[gemini-code-assist]

This large block of commented-out code, which seems to be the original implementation from memory_pool.set_kv_buffer, should be removed to improve code clarity.

[fzyzcjy]

(don't forget to re-enable these code and do branching at correct places)

[fzyzcjy]

maybe we should not pass such objects to this API

what about e.g.

def apply_rope_with_cos_sin_cache_inplace(
  ...,
  # in non-fused version we do `k_buffer[loc] = data` etc
  k_buffer: Tensor, v_buffer: Tensor, loc: Tensor,
)

and if none, it means we do not save kv cache; if non-none then we need to save

[fzyzcjy]

confused why do we have a new file instead of minor modifications

[JeremieMelo]

I am unsure about the best practice for such modifications, e.g., ensuring compatibility. I currently use this style to avoid touching any old files/functions/logic, and introduce fully standalone files to decouple. Any suggestion to make it more compatible /extensible for future updates is appreciated.

[fzyzcjy]

I personally think the fuse is only a dozen lines thus can be inlined. note that you can use c++ template and constexpr if there are overheads

[fzyzcjy]

maybe we should not put layer,forward_batch etc into such a low-level. I personally suggest to put the k_buffer_ptr, v_buffer_ptr, etc, down to here. If too verbose, maybe make a simple @DataClass to pack them

[zhyncs]

@JeremieMelo please fix the conflicts