[akhilg-nv]

[by Claude Code]

Motivation

On SM100/SM103 (B200/B300), DeepSeek-V4.1 Flash can run attention through FlashMLA MegaAttention
(--dsv4-attn-backend flashmla_mega, FP8 KV cache), whose kernels take the query packed as
(tokens, 32, padded_heads, 16) with the heads padded to 64. With attention TP > 1 each GPU holds fewer heads (16 of
64 at TP4), and pack_query rebuilt the packed tensor in each of the 40 attention layers: a zero-fill of the whole
padded tensor, three quarters of it padding at TP4, then a copy of the real heads. This PR gives the MegaAttention
backend one packed-query buffer, zeroed once when it is first used and never reallocated. Each layer writes only its
real heads into a row slice of it, and nothing writes the padding, so it stays zero. Its size is fixed at startup
from the largest forward the server bounds there: a prefill chunk, a CUDA-graph bucket, or a decode or verify batch.
Eager forwards, captured CUDA graphs and the decoder tail graphs (CUDA graphs of the late layers that eager prefill
chunks replay) pack into this one buffer, so its address never changes. A forward with more rows, such as a long
prompt with chunked prefill disabled, packs into fresh zeros as on main and leaves the buffer as it is. The packed
values are bitwise identical and no kernel changes. On current main, MegaAttention is the default in PD prefill and
decode servers and in aggregated servers without attention TP; aggregated servers with attention TP > 1 default to
flashmla since #46263 and use this path with --dsv4-attn-backend flashmla_mega.

Impact

• Accuracy: the packed query is bitwise equal to the old pack (B200 unit tests: 8 to 128 heads, 1 to 8,192
  rows, eager, in CUDA-graph replays with eager forwards in between, and in a forward larger than the buffer).
  sgl-eval GSM8K (200 questions) on 4x B200 TP4: 97.0% and 97.5% with the PR's first revision, 97.5% on main.
• Performance, microbenchmark (B200, 40 packs per forward, 16 heads padded to 64): decode of 8 tokens in a CUDA
  graph 145.4 -> 98.3 us; an 8K-token prefill chunk 11,494 -> 5,990 us in a CUDA graph and 11,717 -> 6,072 us eager.
• Performance, end to end: 4x B200 TP4 low-latency recipe on MegaAttention, 32K-token prompts at concurrency 1:
  TTFT 285.0 -> 275.5 ms (-3.3%) with the PR's first revision; 289.7 -> 271.0 ms (-6.5%) with this revision in the
  memory run below, a smaller sample on a later main. Decode TPOT showed no regression with the first revision.
• Memory: 4x B200 TP4 with 16K-token chunks: the buffer is 1 GiB per GPU and stays allocated, so torch-allocated
  memory is main + 1.0 GiB at all times. The peak during a prefill is main + 1.0 GiB, as with the previous revision of
  this PR. The KV pool is unchanged: the buffer is allocated after the pool is sized.

Modifications

+66 / -6 runtime code (2 files), +352 / -3 tests (1 file). No kernel changes.

• DeepseekV4MegaAttnBackend.packed_query_buffer(rows, padded_heads, dtype) (deepseek_v4_mega_backend.py) returns
  rows [:rows] of the backend's buffer. The first call allocates it with torch.zeros (asserted outside CUDA-graph
  capture). A call with more rows than the buffer returns None, and the caller packs into fresh zeros as on main;
  the buffer is never reallocated, because captured graphs keep its address.

• Size, computed in __init__ by max_packed_query_rows: rows x padded heads x 1 KiB (bf16), where rows is
  the largest of

  • a decode or verify batch: max(max running requests, decode CUDA-graph max bs), padded for MLP sync, x
    speculative draft tokens (1 without speculation);
  • except on a PD decode server, which runs no extend forward: the prefill chunk bound
    (max_prefill_buffer_tokens(), 0 with chunked prefill disabled), the prefill CUDA-graph max bs, and
    SGLANG_DSV4_DECODER_TAIL_GRAPH_MAX_ROWS (2,048 by default; the decoder tail graphs capture every bucket up
    to it).

  At TP4 with 16K-token chunks that is 16,384 x 64 x 1 KiB = 1 GiB per GPU; on a PD decode server with 256 running
  requests and 6 draft tokens, 1,536 x 64 x 1 KiB = 96 MiB. With chunked prefill disabled it covers the decode and
  verify batches and the CUDA-graph buckets, and longer prompts pack into fresh zeros.

• Allocated on the first pack, in practice the first warmup forward at startup, after the KV pool is sized.
  Ranks that never pack allocate nothing: those whose wq_b rows are permuted for MegaAttention (64 or 128 local
  heads with a FlashInfer MXFP8 wq_b), and DSpark draft backends.

• pack_query(q, padded_heads, out=None) (deepseek_v4_mega_backend.py): given out, it writes only the real heads'
  slots. Without out it works as before.

• MQALayer._forward_mega (deepseek_v4.py) packs into attn_backend.packed_query_buffer(...), or into fresh
  zeros when that returns None.

Why the padding stays zero. The buffer is zeroed when it is allocated, and nothing writes its padding after that:

• pack_query writes only [:, :, :heads].
• The MegaAttention kernels only read q. sgl-flash-mla 1.0.0 has two kernel sets: the legacy one, and the one for
  the token-interleaved V4.1 cache added in [Perf] dsv41: Serve token-interleaved V4.1 KV caches with sgl-flash-mla #45893. Both load q into registers with ld.global.nc and never store
  to it (core_attn/kernel.cuh), and test_mega_kernels_leave_q_unchanged checks on SM100 that every kernel of both
  sets leaves q bitwise unchanged at 16 -> 64, 64 -> 64 and 128 -> 128 heads.
• Every layer of a model has the same head count, so the pad positions never move.

Why one buffer per backend is safe. Packs run in _forward_mega on the forward's stream, and CUDA graphs replay
one at a time on that stream; each pack writes every real-head slot of its rows before the attention reads them.
Backends that can run at the same time on different streams (PD multiplexing's prefill and decode SM groups) are
separate instances with separate buffers. DeepSeek-V4.1 rejects two-batch overlap and pipeline parallelism.

Related: pack_query comes from #43922.

Tests, in test/registered/kernels/ops/attention/dsv4/test_mega_attention.py (expected values from the fresh-zeros
pack written out in the test, and hand-derived row bounds):

• TestMegaAttention.test_mega_kernels_leave_q_unchanged, on SM100: each MegaAttention kernel (decode with the
  extra cache, decode without it, prefill), in both kernel sets, leaves its packed query bitwise equal to the
  fresh-zeros pack, at 16 -> 64, 64 -> 64 and 128 -> 128 heads.
• TestMegaPackedQueryBuffer, on any CUDA GPU:
  • Packs equal the fresh-zeros pack bitwise: 8, 16, 32 and 64 heads padded to 64, and 128 to 128, at 8, 48, 3,
    64 (the buffer's capacity) and 1 rows; 16 -> 64 at 8,192, 1,000 and 4,096 rows; 3 layers each. Each pack lands
    in the buffer's storage and the padding is zero at the end. The allocator is filled with NaN before the buffer is
    allocated, so a buffer that is not zeroed fails.
  • CUDA graphs and eager forwards share the buffer: an eager warmup on a side stream allocates it, 64- and
    40-row graphs are captured, then three rounds of replays with new inputs run with a 17-row eager forward between
    replays. Every pack is bitwise equal, the address never changes, and the padding stays zero.
  • A forward larger than the buffer packs into fresh zeros: with a 48-row buffer and the allocator filled with
    NaN, a 49-row forward is bitwise equal to the fresh-zeros pack and does not use the buffer's storage; the buffer
    keeps its address, its contents and its zero padding, and the next 48-row forward packs into it again.
  • Two backends keep separate buffers: a second backend's pack leaves the first backend's packed query intact.
• TestMaxPackedQueryRows: the real max_packed_query_rows on a published config, six hand-derived cases: the
  chunk dominates; the 2,048-row tail bucket outgrows a 1,024-token chunk; an 8,192-token prefill-graph bucket
  outgrows a 4,096-token chunk; chunked prefill disabled, where no chunk term enters and the 2,048-row tail bucket
  sets the size; a PD decode server with 1,000 eager requests x 4 tokens; 50 requests padded to 56 for MLP sync.

Each assertion was checked against a broken version of the code on B200, and each broken version fails exactly the
tests named:

• a pack that writes a padded slot: the bitwise, graph, larger-forward and separate-buffers tests;
• a buffer one row short: the bitwise and larger-forward tests;
• torch.empty instead of torch.zeros: the bitwise, graph and larger-forward tests;
• backends sharing one buffer: the separate-buffers and bitwise tests;
• _forward_mega packing into fresh zeros again: the bitwise and larger-forward tests;
• a larger forward that reallocates the buffer, raises, or is handed the too-short buffer slice: the larger-forward
  test;
• the decode and verify term of the row bound one row short, the bound without the tail-bucket term, without the
  MLP-sync padding, or with a context-length-sized prefill term when chunked prefill is disabled: the bound test;
• a kernel that adds 1 to its query's last head slot after it runs, for each of the three kernels in each kernel
  set (6 versions): exactly that kernel's 3 subtests of the kernel test.

Accuracy Tests

The packed query is bitwise identical and the kernels do not change. To check the integration, B200 TP4 servers ran
the DeepSeek-V4.1 Flash cookbook low-latency recipe on main (535aeae) and on this PR's first revision, when TP4
servers defaulted to MegaAttention: TP4, EP4 (the experts split over the same 4 GPUs), DSpark speculative decoding
(DeepSeek-V4.1's bundled draft model) and decode CUDA graphs. A second pair of servers added the prefill CUDA-graph
flags of the DeepSeek-V4.1 CI tests (--cuda-graph-backend-prefill breakable --cuda-graph-max-bs-prefill 8192 --cuda-graph-max-seq-len-prefill 16384 --chunked-prefill-size 8192).

sgl-eval gsm8k, B200 TP4	main	this PR (first revision)
Cookbook low-latency recipe (prefill graphs off)	97.5%	97.0%
Same, plus breakable prefill graphs	97.5%	97.5%

• Harness: sgl-eval 0.1.3 in its own virtualenv, with the DeepSeek-V4.1 CI settings (dsv41_flash_common.py):
  200 examples, 64 threads, max_tokens 4096, thinking off (no chat template kwargs), temperature 0.0, top_p 0.95.
• Outputs are not deterministic from run to run at 64 concurrent requests, because batch composition changes; the
  0.5-point gap is one question.
• This revision changes where the packed query lives, not its values: the unit tests compare every pack, eager and
  replayed, with the old pack bitwise. In the memory run below, every server of this revision served all its requests,
  and the DSpark accept length was 3.45 and 3.31 against 3.35 and 3.31 on main.

Unit tests, one B200, this revision (on main bf40ca9)	this PR	main bf40ca9
test_mega_attention.py	8 passed, 0 skipped	-
test_flashmla_v41.py (also calls pack_query)	7 passed	7 passed
unit/model_executor/runner/test_early_verify_metadata.py	4 passed	4 passed
attention/unittests/dsv4/test_deepseek_v4.py	27 passed	27 passed
test_num_token_non_padded_localization.py (RowBucketGraphRunner)	14 passed	14 passed

Speed Tests and Profiling

Where on the Pareto curve, and why a static workload. The change removes a zero-fill per attention layer whose
size is rows x padded heads x 1 KiB, so its saving depends only on how many rows a forward has, not on the request
mix. It is largest in prefill chunks and matters most where a chunk's time is the latency: long prompts at low
concurrency (TTFT at the low-latency end). In decode it saves about 1 us per layer. A static workload measures that
directly; the end-to-end checks below are the low-latency end of the 4x B200 TP4 recipe.

The microbenchmark and the memory run below measured this PR before a forward larger than the buffer fell back to
fresh zeros (it raised instead). No forward in either is larger than the buffer, so they ran the same buffer path as
this revision.

Microbenchmark on one B200: each tree's own MQALayer._forward_mega, called through DeepseekV4Model.forward
with the attention call stubbed out, 40 packs per forward of a (tokens, 16, 512) bf16 query into the 64-head layout.
Exact checks against main's pack run first, eagerly and on CUDA-graph replays with new inputs, and all passed. One
process per tree in the order main, previous revision, this PR, this PR, previous revision, main; 5 rounds per
process; us per forward, median of each tree's 10 rounds:

B200, us per 40-layer forward	main	previous revision	this PR
decode, 8 tokens, CUDA graph	145.4	98.3	98.3
decode, 8 tokens, eager	567.0	474.7	432.5
extend, 8K-token chunk, CUDA graph	11,494	5,990	5,990
extend, 8K-token chunk, eager	11,717	6,211	6,072

• Saved per layer against main: 1.2 us at 8 tokens and 138 us for an 8K-token chunk, in CUDA graphs.
• Against the previous revision (one buffer per forward, freed when the forward ended): equal in CUDA graphs.
  Eager, the 8K-token forward is 139 us faster, about one 512 MiB allocation and zero-fill that it no longer does;
  eager decode is 42 us faster in both processes of each tree (host time; not profiled further).
• Sizes: the 8-token decode pack is 8 x 64 x 512 x 2 B = 0.5 MiB per layer, so it stays in L2 when the 40 packs
  run back to back; in serving, the projection GEMM writes q right before the pack and attention reads it right
  after. The 8K-token pack (512 MiB) is far larger than L2.
• The first revision of this PR measured the same pack change on B300 and VR200 with a CUDA graph that zero-filled
  the buffer once per replay: decode of 8 tokens 126.5 -> 90.1 us (B300) and 124.6 -> 83.7 us (VR200), an 8K-token
  chunk 11,201 -> 5,897 us and 8,699 -> 6,332 us.

End to end (the PR's first revision). 4x B200, the cookbook low-latency recipe (TP4, EP4, DSpark, decode CUDA
graphs, chunked prefill 16,384), on main 55f23a8, when TP4 servers defaulted to MegaAttention (on current main,
add --dsv4-attn-backend flashmla_mega). Main and the PR on one node in the order main, PR, PR, main:

python3 -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V4.1-Flash --trust-remote-code --tp 4 --ep-size 4 \
  --reasoning-parser auto --tool-call-parser auto --mem-fraction-static 0.8 --speculative-algorithm DSPARK \
  --speculative-dspark-block-size 5 --cuda-graph-max-bs-decode 64
# TTFT: 2 warmup, then 16 prompts
python3 -m sglang.benchmark.serving --backend sglang --dataset-name random --random-input-len 32768 \
  --random-output-len 1 --random-range-ratio 1.0 --tokenize-prompt --num-prompts 16 --max-concurrency 1 --seed 1
# TPOT
python3 -m sglang.benchmark.serving --backend sglang --dataset-name sharegpt --num-prompts 256 \
  --sharegpt-output-len 512 --max-concurrency 8 --seed 1

	PR / main, pair 1	PR / main, pair 2	second main / first main
TTFT, 32K-token prompts (median of 16, concurrency 1)	0.967	0.967	1.001
TPOT, ShareGPT 256 requests at concurrency 8	0.955	0.980	0.950

• TTFT is 284.8 / 285.1 ms on main and 275.3 / 275.7 ms with the PR.
• Decode: no regression. The first main arm was the slow outlier (1,106 tok/s against 1,820-1,829 in the other
  three), so the decode gain is not resolved.
• Zero-fill kernels per decode step in a profile: 45 on main, 6 with the first revision.

End to end, this revision (from the memory run below). Main 68e57f4, the previous revision and this PR, one
server each in the order main, previous, this, this, previous, main, 6 prompts of 32K tokens per server at
concurrency 1 (one or two of each server's first two requests took 370-500 ms, in every arm):

median TTFT, ms	main	previous revision	this PR
first server / second server	288.5 / 290.8	269.6 / 272.1	270.0 / 271.9

• Without each server's slow first requests, the per-request values do not overlap: 282-291 ms on main, 265-275 ms
  with either revision.
• This is a smaller sample than the first revision's run and saves more (about 19 ms per prompt against 9.5 ms);
  the two runs used different main commits, and the difference is not explained.

Memory

In this recipe every 16K-token prefill chunk runs eagerly, and its packed query is 16,384 x 64 x 512 x 2 B = 1 GiB
per GPU. Main allocated it per layer and freed it after the layer's attention; the previous revision allocated it
once per eager forward and freed it when the forward ended; this PR allocates it once and keeps it.

4x B200 TP4, the recipe above plus --dsv4-attn-backend flashmla_mega --enable-metrics, one node, a fresh server per
run, in the order main 68e57f4, previous revision, this PR, this PR, previous revision, main. A hook in every
scheduler records torch.cuda memory once a second and resets the peaks at each phase; nvidia-smi is read every
200 ms. Phases: 6 x 32K-token prompts at concurrency 1 (as above with --num-prompts 6), 5 s idle, ShareGPT 64 x 512
at concurrency 8 (--num-prompts 64), idle. Rank 0, MiB:

	main	previous revision	this PR	this PR - main	this PR - previous
torch allocated after start-up	98,285.8	98,284.5	99,309.6	+1,023.8	+1,025.1
peak allocated during the 32K prefills	103,266.2	104,288.8	104,287.9	+1,021.7	-0.9
allocated 5 s after the prefills	99,809.0	99,808.0	100,830.7	+1,021.7	+1,022.7
allocated after the decode phase	99,615.1	99,614.2	100,636.8	+1,021.7	+1,022.6
reserved after the prefills (its peak)	109,488	111,798	111,770	+2,282	-28
reserved after the decode phase	103,882 / 103,878	103,738	104,896	+1,014 / +1,018	+1,158
nvidia-smi memory.used, peak during the prefills	162,421	164,717	164,689	+2,268	-28
KV pool, max_total_num_tokens	2,712,576	2,712,576	2,712,576	0	0

• Each value repeats to 0.1 MiB between the two servers of its tree, except where two values are given. Ranks 1-3
  are within 1 MiB of rank 0 for allocated memory and within 26 MiB for reserved memory after the prefills. Every
  server resolved flashmla_mega and served 6/6 and 64/64 requests.
• The buffer is the whole difference in allocated memory: this PR holds 1 GiB more than main at all times, where
  the previous revision held it only during an eager forward.
• Prefill peak +1 GiB, as with the previous revision: a chunk-sized packed query lives through the forward,
  where main's per-layer copy lived only through each layer's attention.
• Reserved and nvidia-smi memory.used are about 2.3 GiB above main during the prefills, for both revisions.
  The extra is cached by PyTorch's allocator and free for later PyTorch allocations; we have not attributed it to
  specific allocations. During the decode phase the reserved memory falls back in every server, and this PR ends at
  main + 1 GiB.
• The KV pool is unchanged: the buffer is allocated in the first warmup forward, after the pool is sized. The
  free GPU memory the server logs after start-up (available_gpu_mem) is 26.37 GB with this PR against 27.35 GB on
  main.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ⏳ Run #42127536643
Latest PR Test (Extra): ❌ Run #42127536579
Latest PR Test (AMD ROCm 10): ⏳ Run #42127536517

Tip

Base CI only covers the core cases, and its workload may not exercise your change. Use /rerun-test to run the tests your change touches, or copy this prompt for an agent:

Find the registered tests this PR affects and run them with /rerun-test, following https://github.com/sgl-project/sglang/blob/main/.agents/skills/rerun-test/SKILL.md

[akhilg-nv]

/tag-and-rerun-ci
[by Claude Code]

[akhilg-nv]

/tag-and-rerun-ci
[by Claude Code]

[nvpohanh]

[by Claude Code] This review is done by agents. Please feel free to disagree if they do not make sense.

[akhilg-nv]

/rerun-failed-ci
[by Claude Code]

[weireweire]

The optimization make sense. But since the padded head slots are never written, could we keep one persistent zeroed buffer on the backend (sized to the max rows, zeroed once at init) and slice it everywhere? That would need much less code than the per-metadata buffers, sharing chain and release logic, and it avoids allocating and
freeing the buffer on every forward.

[nvpohanh]

@akhilg-nv Could you address @weireweire 's comment above?

[akhilg-nv]

[by Claude Code] @weireweire Agreed, and done: DeepseekV4MegaAttnBackend now owns one packed-query buffer, zeroed at its first use (a startup warmup forward) and never reallocated, sized at startup to the largest forward the server bounds there (a prefill chunk or CUDA-graph bucket, or max running requests x draft tokens on a PD decode server); every eager and captured pack writes its real heads into a row slice of it. A forward with more rows, such as a long prompt with chunked prefill disabled, gets no buffer and packs into fresh zeros as on main, so the buffer's address never changes and flashmla_mega keeps working without chunked prefill. The per-metadata buffers, the sharing between kept graphs, the decoder-tail special cases and the per-forward release are gone: the runtime diff is +66/-6 in 2 files against +82/-6 in 3 before, and the tests drop from +594 to +352 lines. One choice to flag: the buffer is allocated on the first pack rather than in __init__, because ranks whose wq_b rows are permuted (64 or 128 local heads with a FlashInfer MXFP8 wq_b) never pack and should not hold it. On 4x B200 (V4.1-Flash TP4, 16K-token chunks, 32K-token prompts) the KV pool is unchanged and torch-allocated memory is main + 1 GiB at all times (the previous revision held that 1 GiB only during a forward, with the same prefill peak), 32K-prompt TTFT is 271 ms against 290 ms on main and equal to the previous revision, and in a 40-layer microbenchmark CUDA graphs match the previous revision while eager forwards are 2-9% faster; details are in the updated description. cc @nvpohanh

Edit: the first version of this comment had a larger forward raise, which made flashmla_mega require chunked prefill; it now packs into fresh zeros instead, so auto and the --dsv4-attn-backend help text stay as on main.

[akhilg-nv]

/rerun-failed-ci
[by Claude Code]