[roborluo]

Overview

Makes the per-layer Q pad + KV insert of DeepSeek-V4.1 mega-attention layers 1.9-2.15x faster on eager steps of 1024 tokens or more at TP2, with bit-identical outputs. On layers with 32 live heads padded to 64 (TP2 of a 64-head model), those steps
now run the fused insert op in its existing KV-only mode and build the zero-padded, chunk-interleaved Q with a contiguous Triton copy, instead of padding Q inside the fused op.
Part of vllm-project#57448, on top of FlashMLA mega attention (vllm-project#56935) and the interleaved V4.1 KV records (vllm-project#61371); related: vllm-project#58872, vllm-project#55654, vllm-project#60076 (see Not a duplicate).

Claims

• 1.9-2.15x faster Q pad + KV insert per layer on eager steps of 1024-16384 tokens (2.10x at 8192 tokens; table below). With 32 live heads padded to 64, the Q half of the fused op reads 32 KiB and writes 64 KiB per token, against about 1.8 KiB for
  its KV half, and it writes Q one (token, head) slot per warp, in 16-element pieces 64 * 16 elements apart. The copy writes each token's 32 rows of 1024 elements contiguously.
• Bit-identical results: the padded Q and the sliding-window KV cache bytes equal those of the fused launch (for NaN-free Q), so model outputs do not change.
• Narrow, automatic dispatch: only eager steps with at least 1024 query tokens on mega-attention layers with 32 of 64 padded heads take the new path. Captured CUDA graphs, smaller steps, other shards (TP1/TP4/TP8 of a 64-head model), other attention
  backends and other platforms run exactly as before. There is no new flag, env var or config option.
• No new sync or memory: the per-step decision reads only host metadata, and the padded Q tensor has the shape the op would have allocated itself. A full GSM8K run against a server started with VLLM_GPU_SYNC_CHECK=error detects no host sync.

Validation

Kernel time of the Q pad + KV insert per layer, fused op vs KV-only launch + copy (bf16, 32 live / 64 padded heads, python bench_q_pad_copy.py below: each arm replayed as a CUDA graph and timed with CUDA events, median; correctness asserted before
timing; two runs agree to within 0.1%):

Tokens	Speedup
512	1.61x (below the 1024-token threshold: keeps the fused op)
1024	1.87x
2048	1.93x
4096	2.06x
8192	2.10x
16384	2.15x

The shipped copy launch config (BLOCK_ROWS=4, 4 warps) is the fastest of the 12 configurations swept at 2048 tokens and tied fastest at 8192.

bench_q_pad_copy.py

"""Fused Q pad + KV insert vs KV-only insert + _pad_interleaved_q (32 of 64 heads)."""

import statistics

import torch
from flashinfer.testing import bench_gpu_time_with_cupti

from vllm.models.deepseek_v41.attention import (
    _pad_interleaved_q,
    _pad_interleaved_q_kernel,
)
from vllm.triton_utils import triton

LIVE, PADDED, D, ROPE, BLOCK, REC = 32, 64, 512, 64, 64, 528
op = torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert
# After q_head_padded: eps, block_size, then the mega-attention layer's flags
# apply_q_norm, kv_mxfp8, apply_q_rope and is_q_interleaved.
TAIL = (1e-6, BLOCK, False, True, False, True)


def bench_us(fn):
    times = bench_gpu_time_with_cupti(fn, use_cuda_graph=True, cold_l2_cache=True)
    return statistics.median(times) * 1e3


    q = torch.randn(n, LIVE, D, dtype=torch.bfloat16)
    kv = torch.randn(n, D, dtype=torch.bfloat16)
    pos, slots = torch.arange(n), torch.arange(n)
    cos_sin = torch.randn(n, ROPE)  # values do not affect timing
    cache_a = torch.zeros(n // BLOCK + 1, BLOCK * REC, dtype=torch.uint8)
    cache_b = torch.zeros_like(cache_a)
    out = torch.empty(n, PADDED, D, dtype=torch.bfloat16)

    def fused():
        return op(q, kv, cache_a, slots, pos, cos_sin, PADDED, *TAIL)

    def kv_only():
        op(q, kv, cache_b, slots, pos, cos_sin, 0, *TAIL)

    def copy():
        _pad_interleaved_q(q, out)

    # Correctness before timing: same cache bytes and Q bits as the fused launch.
    want = fused()
    kv_only()
    copy()
    assert torch.equal(cache_a, cache_b)
    assert torch.equal(out.view(torch.int16), want.view(torch.int16))

    t_fused, t_kv, t_copy = bench_us(fused), bench_us(kv_only), bench_us(copy)
    t_new = bench_us(lambda: (kv_only(), copy()))
    # Every byte read or written once: Q in, padded Q out, KV row, position,
    # slot and cos/sin row in, KV record out.
    nbytes = n * ((LIVE + PADDED) * D * 2 + D * 2 + 16 + ROPE * 4 + REC)
    print(
        f"{n:6d} {t_fused:9.1f} {t_kv:11.1f} {t_copy:8.1f} {t_new:11.1f} "
        f"{t_fused / t_new:7.2f}x {nbytes / t_fused / 1e3:11.0f} "
        f"{nbytes / t_new / 1e3:9.0f}"
    )


def sweep(n):
    """Time the copy alone per launch config; the PR ships BLOCK_ROWS=4, 4 warps."""
    q = torch.randn(n, LIVE, D, dtype=torch.bfloat16)
    want = torch.empty(n, PADDED, D, dtype=torch.bfloat16)
    _pad_interleaved_q(q, want)
    rows = n * (D // 16)
    for block_rows in (1, 2, 4, 8):
        for num_warps in (2, 4, 8):
            out = torch.empty_like(want)
            run = _pad_interleaved_q_kernel[(triton.cdiv(rows, block_rows),)]
            args = (q, out, rows)
            kwargs = dict(
                LIVE=LIVE * 16,
                PAD=(PADDED - LIVE) * 16,
                BLOCK_ROWS=block_rows,
                num_warps=num_warps,
            )
            run(*args, **kwargs)
            assert torch.equal(out.view(torch.int16), want.view(torch.int16))
            t = bench_us(lambda run=run, args=args, kwargs=kwargs: run(*args, **kwargs))
            print(f"{n:6d} {block_rows:10d} {num_warps:9d} {t:8.1f}")


torch.set_default_device("cuda")
torch.manual_seed(0)
print(torch.cuda.get_device_name(), "torch", torch.__version__)
print("tokens  fused_us  kv_only_us  copy_us  kv+copy_us  speedup  GB/s(fused, new)")
for n in (512, 1024, 2048, 4096, 8192, 16384):
    compare(n)
print("tokens  BLOCK_ROWS  num_warps  copy_us")
for n in (2048, 8192):
    sweep(n)

In serving (TP2, mega attention selected explicitly, Engram offload and indexer settings as in the DeepSeek-V4.1-Flash recipe). 8192-token steps are above the default CUDA-graph capture sizes, so they run eagerly and take the new path: the torch
profiler of an 8192-token prefill shows _pad_interleaved_q_kernel (server flag --profiler-config '{"profiler": "torch", "torch_profiler_dir": "./prof"}', then the bench command below with --profile --num-prompts 2).

vllm serve deepseek-ai/DeepSeek-V4.1-Flash --tokenizer-mode deepseek_v41 \
  --tensor-parallel-size 2 --kv-cache-dtype fp8 \
  --attention-config '{"backend":"FLASHMLA_MEGA_ATTN_DSV41","indexer_kv_dtype":"mxfp4","indexer_sparse_logits":true}' \
  --engram-config '{"cpu_offload":true}' \
  --max-model-len 16384 --max-num-batched-tokens 8192 --no-enable-prefix-caching

vllm bench serve --backend vllm --model deepseek-ai/DeepSeek-V4.1-Flash \
  --dataset-name random --random-input-len 8192 --random-output-len 1 \
  --num-prompts 128 --max-concurrency 16

Correctness: unit tests (CUDA GPU with the _C extension built):

pytest -v "tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py::test_q_interleaved_pad_copy_matches_fused_insert"
pytest -v tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
pytest -v "tests/kernels/test_compressor_kv_cache.py::test_v41_pcp_swa_cache_matches_non_pcp_writer"

Test	Result
test_q_interleaved_pad_copy_matches_fused_insert (new)	6 passed
test_fused_deepseek_v4_qnorm_rope_kv_insert.py, whole file	221 passed
test_v41_pcp_swa_cache_matches_non_pcp_writer	2 passed

The new test runs the KV-only launch plus _pad_interleaved_q and the fused launch on the same input at 17, 1024 and 2049 tokens (both sides of the op's NUM_TOKEN_CUTOFF), with and without 5 trailing DP-padding rows that have no slot. It requires
identical cache bytes and identical padded-Q bits, with the copy writing into a NaN-filled buffer.
Lint: ruff check and ruff format --check (0.14.0), typos, py_compile and the local pre-commit hooks (check_spdx_header, check_forbidden_imports, check_torch_cuda, check_boolean_context_manager, check_init_lazy_imports,
validate_config, check_test_tethering) pass on the four changed files, and so does the pre-commit mypy hook (tools/pre_commit/mypy.py, Python 3.11).

CI coverage: the new test is in tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py, which :nvidia: (H100) DeepSeek V4 Kernels, :nvidia: (B200) DeepSeek V4 Kernels and the :amd: (MI355 DPX) DeepSeek V4 Fused Kernels mirror run. This
PR's test change triggers them, but their source_file_dependencies do not list vllm/models/deepseek_v41/attention.py, so a later change to the copy alone would not re-run the test; I can add it if wanted. test_compressor_kv_cache.py runs in
:nvidia: (B200) Miscellaneous Kernels. No CI job serves mega attention at TP2 (the DeepSeek-V4.1 GSM8K config in CI is TP4, where the copy is off), so the serving check above is a manual run.

Details

Motivation. Mega attention applies Q RoPE itself and reads Q chunk-interleaved: per token, 32 chunks of [heads, 16] with the padding heads at the tail of every chunk. The V4.1 layer norms qr before wq_b, so the Q half of
fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert is a pure zero-pad in that layout. With 32 live heads padded to 64 it reads 32 KiB and writes 64 KiB per token, against about 1.8 KiB for the KV half sharing the launch, so on large steps the op's time
is mostly the Q pad. The op processes Q one (token, head) slot per warp, and in this layout a head's 16-element lane chunks sit padded_heads * 16 elements apart; the copy instead writes each (token, chunk) row of 1024 elements contiguously.

Design.

• _pads_q_by_copy is set once per layer: the layer reads Q interleaved (accepts_unnormed_unroped_query), has 32 live heads padded to 64, and max_num_batched_tokens >= 1024. Such layers register _pad_interleaved_q for JIT warmup, so the first
  long prefill does not compile it.
• Per step, _fused_qnorm_rope_kv_insert takes the copy when the layer qualifies, the step has at least 1024 tokens and the current stream is not capturing. The op then runs with q_head_padded=0, the KV-only mode TP1 mega attention already uses on
  main: it schedules no Q slots and runs the same per-token RoPE, MXFP8 quantization and paged insert, skipping DP-padding rows as before (fused_deepseek_v4_kv_rope_insert launches the same KV kernel; passing 0 keeps the existing call).
  _pad_interleaved_q then writes the padded Q into a fresh [num_tokens, 64, 512] tensor.
• _pad_interleaved_q is a Triton kernel over num_tokens * 32 rows, one per (token, 16-dim chunk): 512 live elements copied, then 512 zeros. It goes through triton_kernel_dispatcher_with_warmup like other warmed Triton launchers.
• 1024 equals the op's NUM_TOKEN_CUTOFF, where it switches to its reduced-grid kernel; smaller steps keep one launch. The copy would also win at 512 tokens (1.61x above), but steps that small normally run inside captured CUDA graphs, which keep the
  fused pad.
• The branch lives in the shared _fused_qnorm_rope_kv_insert, next to the existing interleaved-Q handling; a mega-attention override would have to duplicate that method's KV-insert setup.

Exactness. With norm and RoPE off, the op's live-Q path converts bf16 to fp32 and back with no arithmetic in between, which reproduces every non-NaN value bit for bit, and it stores +0.0 in the padding slots. The copy moves the bits unchanged and
stores the same +0.0. The KV half runs the same code in both modes, and the test compares the cache bytes with rtol=0, atol=0.

Tradeoffs and limitations.

• Eligible steps issue two launches instead of one; the extra Triton launch is host-side overhead that large prefill steps hide.
• Only the 32-of-64 shard qualifies: with 16 or 8 live heads a chunk's padding is 768 or 896 columns, not the power of two tl.arange needs, so TP4/TP8 keep the fused pad.
• Work captured into a CUDA graph keeps the fused pad, so the copy applies only where this method runs outside stream capture: by default, steps above the largest captured size; under Model Runner V1, whose wide eager region spans this method, also
  captured sizes of 1024 tokens or more.
• Deployments that select another attention backend (for example FlashInfer sparse MLA) are unaffected.
• The CUDA op is unchanged; reworking its interleaved Q stores would keep a single launch, but in a kernel shared by every DeepSeek-V4/V4.1 backend.

Compatibility.

• Backends and platforms: the FlashMLA sparse, FlashInfer and ROCm layers do not read Q interleaved, so _pads_q_by_copy is false for them; mega attention itself is SM10x-only. DeepSeek-V4 (non-4.1) code is untouched.
• TP sizes (64-head model): TP1 (64 live heads) already skips the pad; TP2 takes the copy; TP4 and TP8 are unchanged.
• Prefill context parallelism ([Feature][PCP] Add DeepSeek-V4.1-Flash support vllm-project/vllm#59857) selects the FlashMLA layer, so it never reaches the copy; the stand-in layer of test_v41_pcp_swa_cache_matches_non_pcp_writer now declares _pads_q_by_copy=False, since that test calls
  _fused_qnorm_rope_kv_insert on it. The SM90-only token-sharded fused_wqa_wkv ([Perf][DSv4.1] Split fused_wqa_wkv by token across TP ranks on SM90 vllm-project/vllm#58872) never meets mega attention, and forward_kv (decoder replay) builds no Q.
• CUDA graphs: capture keeps the fused pad, so captured graphs are unchanged. Batch invariance: both paths produce the same bits, so taking the copy (which depends on the step's token count) changes no result.
• CPU/GPU sync: none added; the decision reads q.shape[0], a flag set at init and the stream's capture state.

Not a duplicate. On 2026-10-18 I searched open PRs and PRs since 2026-08-15 for "pad interleaved q", "qnorm rope kv insert", "Q pad mega" and "heads padded 64 decode"; none moves the mega-attention Q pad out of the fused insert. Related work: vllm-project#56935
added mega attention and the fused interleaved pad this PR bypasses on large eager steps; vllm-project#61371 changed the V4.1 KV record, which both launch modes share; vllm-project#58872 added the token-sharded fused_wqa_wkv whose _SHARD_FUSED_WQA_WKV_MIN_TOKENS sits next
to the new constant; vllm-project#60076 (open) switches the ROCm DSpark context-KV insert to the existing KV-only fused_deepseek_v4_kv_rope_insert and touches neither mega attention nor its Q pad (it and this PR edit different tests in
tests/kernels/test_compressor_kv_cache.py). vllm-project#55654 (open) adds KV-only insert ops (fused_deepseek_v4_kv_rope_quant_insert, and fused_deepseek_v4_kv_norm_rope_quant_insert with kv_norm folded in) so DeepSeek-V4's DSpark _insert_context_kv no
longer builds a dummy query. This PR adds no KV op and touches neither DSpark nor DeepSeek-V4: it calls the existing q_head_padded=0 mode of the fused op from the V4.1 attention layer. That path also sees DP-padding rows (more KV rows than slots),
which the fused op's KV-only mode skips and vllm-project#55654's ops reject. The only file both PRs touch is tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py, where each adds tests.

AI assistance. AI assistance (Claude Code) was used to develop this change, port it onto current main, run vLLM's pr-checklist skill and draft this description. I reviewed every changed line and ran the validation reported above.

---

Pull Request Checklist

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.