[Kernel] Add native CUDA fused RoPE + KV cache write op, opt-in for flash_attn

[A1c0r-Z]

[Kernel] Add native CUDA fused RoPE + KV cache write op, opt-in for flash_attn

Part of #43224 #43503 (the rope + kv_cache_write bullet under the CUDA fusion-pass inventory).

Why this PR

On NVIDIA today, RopeKVCacheFusionPass is effectively a no-op: no CUDA AttentionImpl subclass implements fused_rope_kvcache_supported(), and the pass_config.fuse_rope_kvcache flag is gated is_rocm()-only in vllm/config/compilation.py. This PR closes both gaps:

1. Adds a native CUDA fused kernel (csrc/rope_kvcache_fusion_kernels.cu).
2. Opts FlashAttentionImpl into the existing RopeKVCacheFusionPass machinery.
3. Relaxes the platform gate to is_cuda_alike() so CUDA users can opt in. The flag still defaults to off; this PR doesn't change user behavior unless someone explicitly sets compilation_config.pass_config.fuse_rope_kvcache=True.

AI-assisted implementation. I've read the core kernel (csrc/rope_kvcache_fusion_kernels.cu) and the backend wiring (vllm/v1/attention/backends/flash_attn.py) end-to-end, run the test suite and benchmark, and am responsible for defending the design under review. Additional files (bindings, Python wrapper, tests) are mechanical compared to those two and were spot-checked rather than read line-by-line.

Relationship to #43224

Step 1 of two: fill the CUDA kernel gap so RopeKVCacheFusionPass does something on NVIDIA. Step 2 (separate PR) will rewire model code to call the fused op directly per the #42597 pattern, then remove the FX pass. Happy to fold both into one PR if preferred.

What's in scope

New kernel file, C++ bindings, Python wrapper, FlashAttentionImpl opt-in, platform-gate relax (is_rocm() → is_cuda_alike()), parameterized correctness test, and a benchmark. See the Files changed tab.

Support matrix (first PR)

• scalar_t (q/k/v): bf16, fp16.
• cache_t: bf16, fp16, fp8_e4m3, fp8_e5m2.
• IS_NEOX: both branches.
• Layout: flash NHD only ([num_blocks, block_size, num_kv_heads, head_size]). Strides validated by host-wrapper TORCH_CHECKs.
• Scales: per-tensor only (k_scale.numel() == 1).

Deferred to follow-ups

• FP8 per-token / per-head k/v scales
• NVFP4 cache support (SM100+)
• Flashinfer backend wiring (Support Flashinfer rope+quant+cache update fusion kernel for TRTLLM attention #36858 covers the flashinfer side)
• Vectorized cache writes via vectorize_with_alignment — tried in this PR, regressed small-N by 30–60% (register pressure / occupancy hit); reverted to scalar. Could revisit with explicit uint4 stores.
• Manual call-site rewire in model code (per [RFC]: Porting compiler fusions to manual fusion #43224's manual-fusion direction — Step 2 above)

Kernel design

• 1 CTA per token. Block size = min(max(rope_work, cache_work), 512). Mirrors concat_and_cache_mla_rope_fused_kernel's shape.
• Phase 1: in-place RoPE on Q, fp32 intermediates.
• Phase 2: in-place RoPE on K, fp32 intermediates.
• __syncthreads().
• Phase 3: if slot_idx >= 0, write rotated K and V into the paged cache (optionally FP8-quantizing via fp8::scaled_convert).

See the header comment of csrc/rope_kvcache_fusion_kernels.cu for design notes (fp32 RoPE math, scalar Phase-3 writes, bit-identical guarantee).

How to enable

Two flags must be set together — there's already a warning in compilation.py near the platform gate prompting this:

LLM(
    model=...,
    compilation_config={
        "pass_config": {"fuse_rope_kvcache": True},
        "splitting_ops": [],  # so the FX pass sees RoPE and the KV-cache update in one subgraph
    },
)

The two-flag opt-in is friction; unifying these into a single user-visible default is a reasonable follow-up but out of scope here.

Test plan

# Correctness — 64 cases bit-identical vs unfused
pytest tests/kernels/attention/test_fused_rope_kvcache.py -v
# → 64 passed in 18s

# Benchmark — H200, 2000 iters / 50 warmup, scalar build
python benchmarks/kernels/benchmark_fused_rope_kvcache.py --iters 2000 --warmup 50

E2E sanity script (TinyLlama-1.1B-Chat-v1.0 + fuse_rope_kvcache=True + splitting_ops=[], greedy decode 32 tokens on "The capital of France is") run as three subprocesses for the three-way control table below. Expected DEBUG output:

DEBUG fusion/rope_kvcache_fusion.py:268  Replaced 22 patterns
DEBUG vllm_inductor_pass.py:84            RopeKVCacheFusionPass completed in 887.7 ms

Correctness assertion (bit-identical)

The test asserts torch.testing.assert_close(..., rtol=0, atol=0) across the matrix:

• dtype_cache ∈ {(bf16, auto), (fp16, auto), (bf16, fp8_e4m3), (fp16, fp8_e4m3)}
• head_config ∈ {MHA(32,32,128), GQA(32,8,128)}
• is_neox ∈ {True, False}
• num_tokens ∈ {1, 8, 128, 2048}

The test samples slot_mapping without replacement — duplicate slots would cause "last write wins" non-determinism between the fused and unfused paths under non-deterministic CUDA block scheduling. The production v1 scheduler is already constrained to produce unique non-negative slot ids within a single forward pass (each request gets a disjoint block allocation; prefix-cache sharing is read-only; CUDA-graph padding uses -1 which our kernel skips). Our kernel inherits the same uniqueness assumption as the existing reshape_and_cache_flash — no new exposure.

Benchmark — H200, 2000 iters / 50 warmup, median-of-3 independent runs

dtype	cache	heads	N=1	N=8	N=128	N=2048
bf16	auto	MHA(32/32/128)	1.71×	1.72×	1.75×	1.05×
bf16	auto	GQA(32/8/128)	1.67×	1.71×	1.76×	1.47×
fp16	auto	MHA(32/32/128)	1.74×	1.77×	1.74×	1.06×
fp16	auto	GQA(32/8/128)	1.72×	1.71×	1.75×	1.49×
bf16	fp8_e4m3	MHA(32/32/128)	1.43×	1.34×	1.27×	1.03×
bf16	fp8_e4m3	GQA(32/8/128)	1.72×	1.72×	1.72×	1.56×
fp16	fp8_e4m3	MHA(32/32/128)	1.72×	1.74×	1.64×	1.07×
fp16	fp8_e4m3	GQA(32/8/128)	1.74×	1.74×	1.74×	1.66×

Every cell ≥ 1.03×. Sweet spot is decode at small/medium N (consistent 1.6–1.77×) and prefill on GQA (1.4–1.7×). The MHA-at-large-N row (1.03–1.07×) is HBM-bandwidth-bound — fusion only saves the K re-read between the two unfused kernels, which is ~10% of total traffic on MHA; the measurement matches that physical ceiling.

E2E sanity (three-way control)

TinyLlama-1.1B-Chat, CUDA_VISIBLE_DEVICES=1, greedy decode (temperature=0, seed=0), prompt "The capital of France is", 32 generated tokens. Three configs run as separate subprocesses for clean isolation.

Run	fuse_rope_kvcache	splitting_ops	First-32 tokens (text)	Marker
baseline	False	default	' Paris.\n\n2. B. The capital of France is Paris...'	0
split_only	False	[]	' Paris.\n\n2. B. C. The capital of Canada is Ottawa...'	0
full	True	[]	identical to split_only	1540

Marker counts were obtained via temporary NVTX/print instrumentation in do_rope_and_kv_cache_update during validation; that instrumentation is not part of the submitted code.

split_only == full token-for-token; the divergence from baseline is entirely caused by splitting_ops=[] (inductor sees a different graph → different fusion choices → different reduction orders in attention → small fp diffs accumulate over autoregressive decoding). The fused op is bit-identical to the unfused op when the rest of the compile config is held constant. The marker count = 1540 confirms the runtime call path is live (~70 forward passes × 22 decoder layers).

What I haven't done in this PR

• Llama-3.2-1B + GSM8K accuracy run. The model is gated and this box has no HF token; substituted TinyLlama-1.1B-Chat for the integration check. Happy to run GSM8K on a non-gated model (e.g. Llama-2-7B) if maintainers prefer real accuracy numbers before merge.
• Vectorized cache writes (see "Deferred").
• Manual call-site rewire in model code (Step 2 above).

[gemini-code-assist]

Code Review

This pull request introduces a fused CUDA kernel that combines Rotary Positional Embedding (RoPE) with the KV cache update process for the flash NHD layout. The implementation includes the CUDA kernel, Python bindings, and integration into the FlashAttention backend. Feedback identifies two critical correctness issues: the kernel incorrectly uses integer types for floating-point conversions in the FP8 path, and the backend lacks a check to ensure the KV cache layout is NHD, which could lead to memory corruption.

[gemini-code-assist]

The kernel passes raw_kv_scalar_t (which is an integer type like uint16_t for FP16/BF16 inputs) to fp8::scaled_convert. This causes the conversion to perform integer-to-float casting on the bit representation of the RoPE output, resulting in incorrect values. It should use the floating-point type qk_t and access elements directly via k_src[i] and v_src[i] to ensure correct floating-point conversion.

    for (int i = threadIdx.x; i < n_elems; i += blockDim.x) {
      k_dst[i] = fp8::scaled_convert<cache_t, qk_t, kv_dt>(
          k_src[i], k_scale_val);
      v_dst[i] = fp8::scaled_convert<cache_t, qk_t, kv_dt>(
          v_src[i], v_scale_val);
    }

[gemini-code-assist]

The fused_rope_kvcache_supported method assumes that the HND layout is not produced by this backend. However, FlashAttentionBackend supports both NHD and HND layouts depending on configuration. The fused kernel implementation in rope_kvcache_fusion_kernels.cu assumes a contiguous [num_kv_heads, head_size] block per slot, which is only true for the NHD layout. Running this on an HND layout will lead to silent memory corruption. An explicit check for get_kv_cache_layout() == "NHD" should be added.

Suggested change

	return self.kv_cache_dtype != "nvfp4"
	return self.kv_cache_dtype != "nvfp4" and get_kv_cache_layout() == "NHD"

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @A1c0r-Z.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @A1c0r-Z.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @A1c0r-Z.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[Utopia-V]

Hi @A1c0r-Z @mgoin — I independently started working on this CUDA RoPE + KV-cache fusion gap against current vLLM main before discovering this PR and #43503. I do not want to open a competing PR without coordinating here first.

My current-main prototype follows the same broad native-kernel + FlashAttention integration direction, but the port surfaced two correctness/integration details that may be useful to this effort:

• The launch domain needs to cover the RoPE token count, which can be larger than slot_mapping.numel() under graph padding. Tokens beyond the cache-update domain still need their Q/K rotation.
• When a per-layer slot mapping is absent, the fallback needs to retain RoPE-only behavior rather than skipping the combined hook entirely.

I also have current-main SM89 operator benchmarks and compile/V2 PIECEWISE route evidence, and I am continuing to turn the prototype into an upstream-reviewable candidate.

Which continuation would you prefer?

1. Contribute the fixes and current-main rebase to this PR;
2. Prepare a current-main successor while preserving attribution to this work; or
3. Split out the native op and align the integration directly with the manual-fusion direction in [RFC]: Porting compiler fusions to manual fusion #43224 / [Feature]: Porting RopeKVCacheFusionPass to manual fusion #43503.

I am happy to adapt to maintainer direction. In the meantime I will keep validating the current-main implementation locally rather than opening a duplicate PR.

[Utopia-V]

Correction to the FP8 point in my comment above: after tracing this PR's exact scaled_convert<uint8_t, uint16_t> specialization, the uint16_t value is deliberately interpreted as the raw half bits through half_to_float; it is not converted as an ordinary integer. That concern is therefore not established and should be disregarded. Sorry for the noise.

The T_rope versus T_cache launch-domain issue and the missing-slot RoPE-only fallback point remain applicable to the current PR head.

[A1c0r-Z]

Closing this.

Two reasons, both structural rather than fixable by a rebase.

The KV-cache layout moved. get_kv_cache_shape now returns (num_blocks, num_kv_heads, block_size, 2 * head_size) with K and V packed into the last dim, so this PR's kv_cache.unbind(1) is wrong, and the kernel's host contract (key_cache.stride(2) == head_size, stride(3) == 1) rejects the transpose(1, 2).split(...) views that are the correct way to split it. It fails end to end on any real model. That's a kernel rewrite, not a rebase.

More importantly it is on the wrong side of #43503. This PR is Step 1 — filling the backend hook so RopeKVCacheFusionPass does something on CUDA — and its own description promised a Step 2 that would "rewire model code to call the fused op directly ... then remove the FX pass". I never wrote Step 2. #52363 is that Step 2, with a Step 1 that handles the packed layout, the NHD/HND strides and the separate RoPE/cache token domains correctly.

Superseded by #52363.