[Performance]: ROPE + KV-Cache-Write + pre-attn prepare-ops fusion

[alexm-redhat]

Proposal to improve performance

Performance analysis of DeepSeekR1 and GPTOSS showed that there is non-trivial amount of overheads in the operations that come before the Attention kernel. Specifically, these are ROPE + KV-Cache-Write + attn elementwise/copy prepare ops (and sometimes small matrix multiplies for MLA).

Here is a concrete breakdown for DeepSeekR1-FP8 batch-size 32 on 8xB200 GPUs:

Before the attention kernel call (fmha::kernel), we can see 4 triton ops (ROPE), then cache_and_concat_mla, and then a torch BMM (matrix_multiply - K-Proj) followed by a bunch of elementwise and copy. The total time these ops take is about the same as the whole attention kernel, so we propose to fuse all of them to reduce HBM traffic bandwidth, since these ops are memory-bandwidth by their nature.

A similar fusion is needed for GPTOSS as well, here a breakdown:

This one is a bit simpler, since the attention kernel does not expose elementwise ops (or matmul) like in DeepSeekR1 model above. However, there is still the series of triton ops (ROPE) followed by reshape_and_cache_kernel.

NVIDIA has a fused ROPE+KV-Cache-write+attn prepare ops kernel in TRT (@pavanimajety is part of the related team) and would be good to understand the following:

1. Does current work supports both reshape_and_cache and concat_and_cache_mla?
2. What about quantization types? Can we do static per-tensor, dynamic per-token, and dynamic per-group (block quant)?
3. @Lucas Wilkinson and @ProExpertProg talked about how to expose the cache and quant and there's actually some ongoing work to refactor the attention layer into multiple layers so that MLA/SDPA/etc. don't all have to go through the same unified_attention op.

• Issue here: [Refactor]: Make an common MLAAttention Layer and custom OP #24620
• That means concat_and_cache_mla will be exposed in the layer and will appear in the graph.
• For quants, instead of adding more methods and layers (my original proposal), we should implement a method similar to output quant fusion. We can query the attention backend whether input quant de-fusion is supported. If yes, we add a quant to the graph outside unified attention. This is now handled in the attention layer: [torch.compile] Make Query Quantization Fusable #24914. This won't change the semantic meaning of unified_attention as much because we're just changing the dtype of the input. We might need to figure out how to expose the scale for the quant layer to the graph, if we can't make it visible in the traced graph we can just add a new custom op that extracts it from the attention layer.

PR tracking:

• MLAAttention extraction: #25103
• Separate kvcache from unified_attention: #25954
• Add concat_cache_mla_rope CUDA kernel: #25774

[pavanimajety]

The Kernel in question from FI is: flashinfer-ai/flashinfer@fc88829

It currently only supports per-tensor quantization for the query(code ref). k_cache quantization and v_cache quantization(in case of GPT-OSS / LLama 4) needs to be currently handled on the framework side.

What needs to be done:
This FI kernel needs to be extended to support quantization of Q, K, and V tensors, as well as in-place updates of the KV cache, in a generalized manner. It would do this by accepting the framework's KV cache layout, flexibly handling varying head sizes and cache partitions required by each model—assuming RoPE-based encoding is performed in the same way. (See similar support for MLA in TRTLLM implementations).

Model	num_attention_heads	num_kv_heads	Head Dim	RoPE Math/Cache Split	q_rope/q_nope & kv split	Quantization Support
GPT-OSS	64	8	64	Standard RoPE per Q	No explicit q_nope/q_rope distinction	mxfp4, per-query only
Llama 4 (Scout)	128	128	64 RoPE 128 NoPE	Chunked RoPE, full NoPE	Distinct q_rope/q_nope, hybrid layers	FP8/int4
DeepSeek R1	128	128	64 qk_rope 128 qk_nope 128 v	RoPE/NoPE interleaved	q_rope/q_nope, k_rope/k_nope distinct	FP8 dynamic, full support

Post Slack discussion-
My plan is to

1. Resurrect [Perf] Use FlashInfer RoPE for RotaryEmbedding.forward_cuda when available #21126 and wrap RoPE as a single custom op
2. Add the Kernel support in Flashinfer for the combined fusion kernel.
3. Add the Fusion pass in vLLM for doing Rope + QKV Quant + KV Cache Update operations for the three highlighted models.

I think (1) and (2) can be parallelized. I will discuss with @nvpohanh tackle the tasks.

[ZJY0516]

I will implement the fusion pass (task 3) once the kernel is merged.

[pavanimajety]

Hey all, update:

Rope + KV Cache + Quant Fusion for DSR1 and more

What's finished/in review:

• Separate Quant from attention (Enabled for Flashinfer, triton and Flash Attn): PR
• Separate KV Cache and Attention (Enabled for Flash Attention): PR
• Rope + Quant Kernel for DeepSeek MLA in Flashinfer.

In Progress:

• Rope + KV Cache Update Fusion with Flash Attention 2 Backend(WIP) Pavani
• Rope + Quant + KV Cache update Cuda Kernel - Flashinfer MLA RoPE + quantization fused kernel: shape generalization for MHA / GQA flashinfer-ai/flashinfer#1924

What’s needed(in dependent order):

• Separate Quant from MLAAttention
• Separate KV Cache update and MLAAttention
• Separate KV Cache update and Flashinfer MLA attention to be compatable for kv-cache-dtype=fp8
• Fusion pass to replace Rope + Quant + KV Cache for Prefill MHA in DeepSeek
• Fusion pass to replace Rope + Quant + KV Cache for Flashinfer MLA Attention
• Fusion pass to replace Rope + Quant + KV Cache for Flashinfer MHA Attention