[ROCm][MLA] Use saturated Gluon split bucket for CUDA graphs

[LiuYinfeng01]

Summary

This PR changes the default AiterMLADecodeMetadata.min_kv_seq_len from 1 to 16384.

The change is intentionally small, but the reason is specific to the interaction between vLLM FULL CUDA Graph capture and AITER's host-side Gluon split selection:

• AITER uses min_kv_seq_len to choose the split-KV launch topology in its Python wrapper.
• vLLM currently leaves this field at its default value.
• FULL CUDA Graph capture runs the wrapper once with dummy decode metadata and then replays captured GPU work without running the wrapper again.
• With the old default, the captured graph is permanently specialized to the one-split path, even when the real request has a long KV sequence.

This PR only fixes that vLLM integration issue. Split-policy tuning and stage-1/stage-2 kernel optimization remain AITER responsibilities.

How the issue was found

I first found this from Kimi-K3 TP8 end-to-end serving performance on gfx950. The 100K-input / 1K-output, concurrency-1 workload measured roughly 87–90 ms TPOT, while an isolated 12-head Gluon configuration with sufficient split parallelism indicated that the decode path should be much faster.

The initial investigation focused on AITER. ROCm/aiter#4450 tuned the 12-head BF16 Gluon split policy and tested context buckets for the Kimi-K3 TP8 shape. Reducing the split count also reduced the cost of the existing serial stage-2 reduction, so the first apparent solution was an AITER split-policy change.

During the review of ROCm/aiter#4450, @Dewei-Wang-sh and I spent an afternoon tracing this behavior from the AITER wrapper back into vLLM. His key review direction was to pass the split-sizing information from the caller. He explained that the existing AITER interface follows the Moonshot integration: Moonshot obtains a meaningful minimum KV length in its runtime and passes it to AITER. In that contract, 1 is a fallback default; it is not intended to describe every real long-context decode request.

This exposed the integration mismatch. vLLM was leaving the fallback at 1. More importantly, when FULL CUDA Graph was enabled, vLLM performed warmup/capture with a dummy sequence length of 1. Dewei helped me follow the execution path and confirm that real decode replay did not run the host wrapper again. Therefore, even an attempted runtime assignment could not reach the Python code that selects the launch topology. This discussion and the subsequent call-chain tracing moved the fix from AITER to the vLLM caller. Thank you to Dewei for the detailed review and for helping identify the actual boundary of the problem.

The relevant AITER host calculation is in mla_gluon.py:

NUM_KV_SPLITS = max(
    1,
    min(
        256 // (batch_size * qlen * NUM_M_BLOCKS),
        triton.cdiv(min_kv_seq_len, BLOCK_N),
    ),
)

For the BLOCK_N=64 path, capture with min_kv_seq_len=1 gives cdiv(1, 64)=1, so the host calculation collapses to one KV split. The captured graph then permanently records that one-split launch. With the vLLM capture hint set to 16384, the same term becomes cdiv(16384, 64)=256; it no longer forces the policy to one split, and AITER can apply its workgroup budget and small-head policy to capture a saturated, non-one-split topology. The real sequence length is still read from device metadata during replay; 16384 is only the host-side capture hint.

In vLLM, the field was declared with a default of 1, and _build_decode() constructed AiterMLADecodeMetadata without assigning it. forward_mqa() then passed that unchanged value to mla_gluon():

• min_kv_seq_len metadata default
• _build_decode() metadata construction
• forward_mqa() passing the host hint to AITER

Therefore, even a real 10K or 100K request reached the AITER wrapper with min_kv_seq_len=1 in normal execution.

Why assigning the real runtime length did not fix FULL CUDA Graph

I next tested assigning the real value while constructing decode metadata:

min_kv_seq_len=int(seq_lens_for_kernel.min())

For an actual request this computes the expected value, for example approximately 10K or 100K. It works conceptually in eager execution because the AITER Python wrapper runs for every call and can select a new split count.

It does not work for FULL CUDA Graph capture.

For pure decode, vLLM's dummy run has query length 1. The capture path consequently builds dummy sequence lengths from max_query_len, which is also 1:

• pure-decode dummy batch construction
• dummy attention sequence length selection

The resulting sequence is:

capture startup:
  max_query_len = 1
  dummy seq_len = 1
  min_kv_seq_len = 1
  AITER wrapper selects NUM_KV_SPLITS = 1
  CUDA Graph records the one-split grid and the no-stage-2 fast path

real request:
  device metadata is updated to the real KV length
  FULL graph replay starts
  AITER Python wrapper is not called again
  captured split count, grid, workspace, and stage topology remain unchanged

Instrumented FULL CUDA Graph trace

To verify this rather than infer it from performance, I instrumented an isolated vLLM/AITER copy and ran one Kimi-K3 TP8 request with 100K input tokens and 4 output tokens. Logging was restricted to TP rank 0. The complete capture call chain was:

worker_busy_loop
-> compile_or_warm_up_model
-> capture_model
-> _capture_cudagraphs
-> _warmup_and_capture
-> _dummy_run
-> self.model
-> BreakableCUDAGraphWrapper._capture
-> Kimi-K3 MLA layer
-> unified_mla_attention_with_output
-> forward_mqa
-> mla_gluon

The following excerpts are sanitized to retain only the events relevant to MLA graph capture and replay.

Eager warmup

MLA_TRACE event=dummy_metadata mode=NONE seq_lens=1
MLA_TRACE event=aiter_metadata_build seq_lens=[1] static_seq_info_ptr=137855956567552
MLA_TRACE event=forward_mqa_call stream_capturing=False host_min_kv_seq_len=1
MLA_TRACE event=mla_wrapper_split stream_capturing=False block_n=64 num_kv_splits=1
MLA_TRACE event=mla_stage1_launch stream_capturing=False grid=(1, 1, 1) logits_shape=(1, 1, 12, 1, 512) mid_lse_shape=None
MLA_TRACE event=mla_stage2_skipped stream_capturing=False reason=num_kv_splits_is_1

FULL graph capture

MLA_TRACE event=dummy_metadata mode=FULL graph_capture_request=True seq_lens=1
MLA_TRACE event=model_capture_enter mode=FULL
MLA_TRACE event=forward_mqa_call stream_capturing=True host_min_kv_seq_len=1 seq_info_ptr=137855956567552
MLA_TRACE event=mla_wrapper_split stream_capturing=True block_n=64 num_kv_splits=1
MLA_TRACE event=mla_stage1_launch stream_capturing=True grid=(1, 1, 1) logits_shape=(1, 1, 12, 1, 512) mid_lse_shape=None
MLA_TRACE event=mla_stage2_skipped stream_capturing=True reason=num_kv_splits_is_1
MLA_TRACE event=model_capture_exit mode=FULL graph_segments=1 eager_breaks=0

At this point capture has fixed all of the host-selected topology:

• NUM_KV_SPLITS=1
• stage-1 grid (1, 1, 1)
• workspace/logits shape (1, 1, 12, 1, 512)
• mid_lse=None
• no stage-2 GPU node exists in the graph

Real 100K decode replay

MLA_TRACE event=runtime_dispatch mode=FULL computed_before_step=[100000]
MLA_TRACE event=aiter_metadata_build seq_lens=[100001] static_seq_info_ptr=137855956567552
MLA_TRACE event=model_graph_replay mode=FULL count=1 segments=1 eager_breaks=0
MLA_TRACE event=aiter_metadata_build seq_lens=[100002] static_seq_info_ptr=137855956567552
MLA_TRACE event=model_graph_replay mode=FULL count=2 segments=1 eager_breaks=0
MLA_TRACE event=aiter_metadata_build seq_lens=[100003] static_seq_info_ptr=137855956567552
MLA_TRACE event=model_graph_replay mode=FULL count=3 segments=1 eager_breaks=0

The capture-time seq_info_ptr and all three runtime static_seq_info_ptr values are the same: 137855956567552. This proves that the real metadata was copied into the same static GPU buffer used by the graph. However, there were no forward_mqa_call, mla_wrapper_split, or Python stage-launch events between the three replay steps.

Why the wrapper is not re-entered

This is not a JIT-cache effect. During capture, BreakableCUDAGraphWrapper executes the model and records the graph. On later calls, its capture-versus-replay dispatch selects _replay(), which calls only entry.capture.replay(). The capture artifact then replays its saved segments.

This run reported graph_segments=1, eager_breaks=0, so the only saved segment was GPU graph replay; it was not a Python model, attention, or AITER-wrapper invocation. A PIECEWISE graph can re-enter Python when attention is registered as an eager break, but this FULL graph had no such break. JIT kernel caching also does not skip the Python wrapper and should not be confused with graph replay.

The distinction is therefore:

• Dynamic during replay: the real sequence length, because the captured kernel reads the latest values from the static seq_info GPU buffer.
• Frozen at capture: NUM_KV_SPLITS, workspace shape, stage-1 grid, mid_lse allocation, and whether stage 2 exists, because those are host-side decisions made before the graph is recorded.

This directly explains why a host-only split policy still regressed to approximately 87–88 ms: FULL capture recorded the split=1 topology and replay could not change it.

Diagnostic timing note: this instrumented 100K-to-4-token run measured 131.33 ms TPOT, but the instrumentation added GPU-to-CPU synchronization, stack collection, and synchronous printing. That number is only call-chain evidence and must not be compared with the uninstrumented 23.49 ms result below. The raw logs are approximately 640 KB and contain unrelated runtime configuration, so only the focused, sanitized trace is included here.

The runtime assignment also requires a device-to-host synchronization for .min()/.item(). It increased the measured TPOT from approximately 87–88 ms to approximately 89–90 ms while still replaying the one-split graph. It was therefore both insufficient and slower.

Why use 16,384 for capture

Changing vLLM to capture a separate graph for every context bucket would be a much larger architectural change. Moving split activation back to the device is an AITER policy decision and conflicts with the reviewer-requested host-side interface. Per-request device-to-host synchronization is not viable on the decode path.

The smallest vLLM-side integration is therefore to provide one representative, saturated capture hint. 16384 was selected so capture does not fall into the one-split topology, while avoiding a context-length graph ladder or runtime synchronization.

The value is a capture hint, not an assertion that every request has a 16K context. The kernel still reads the real per-request sequence length from seq_info during replay. The hint determines only the host-selected split topology that CUDA Graph records.

The target workload matrix includes 8K/1K, 60K/1K, and 100K/1K cases. The fixed-hint end-to-end rerun completed so far is the 100K/1K case; the shorter fixed-hint cases have not yet been rerun end to end.

Runtime-owner feedback and related AITER work

The AITER runtime discussion clarified that two separate problems had been mixed together:

1. vLLM integration correctness/performance: the caller leaves the host hint at 1 and FULL graph capture freezes that topology.
2. AITER kernel efficiency: with a valid large split count, the original stage-2 reduction becomes serial and expensive.

The related AITER PRs explore different kernel-side solutions:

• ROCm/aiter#4450 started the 12-head split-policy investigation and production-like scattered-page validation.
• ROCm/aiter#4507 derives host launch sizing from page-table shape rather than the optional hint.
• ROCm/aiter#4509 adds split-major stage-1 scheduling and a more parallel stage-2 reduction on top of ROCm/aiter#4507.
• ROCm/aiter#4555, from the AITER runtime team, directly optimizes stage-2 by loading and reducing a block of partial outputs at once instead of walking one split at a time.

The runtime-owner measurement makes the separation important:

Configuration	Splits	Stage 1	Stage 2	Total
Original serial reduction	256	34.7 us	74.3 us	109.0 us
Split-bucket workaround	96	43.1 us	26.1 us	69.2 us
ROCm/aiter#4555 blocked reduction	256	34.7 us	10.7 us	46.9 us

This shows that reducing splits helped partly because it avoided an inefficient stage-2 implementation. ROCm/aiter#4555 is the better kernel optimization: it preserves enough stage-1 workgroups to fill the CUs and fixes stage-2 itself. The AITER runtime team will own that optimization.

Accordingly, this PR does not claim that 48 or 96 splits is the final AITER policy. It only ensures that vLLM does not silently force every host-side policy to the pathological one-split capture.

End-to-end validation

Kimi-K3 TP8 on 8x MI355X/gfx950, BF16 KV, FULL CUDA Graph, 100K input / 1K output, concurrency 1:

Configuration	TPOT
Previous host-side capture with default 1	approximately 87.8 ms
Runtime min() metadata experiment	approximately 89.2 ms
This PR, fixed 16K capture hint	23.49 ms
Device-side runtime-policy reference	approximately 23.0 ms

The fixed capture hint removes the regression without changing graph dispatch or adding decode-path synchronization.

Scope question for maintainers

The remaining concern is short-context behavior. This change assumes that the small-head Gluon path is primarily valuable for the long-context workloads for which split parallelism is required. A fixed 16K hint may over-split very short requests.

AITER has empty-split guards, and ROCm/aiter#4450 validated explicit split counts larger than sequence lengths at boundaries including sequence length 1, so this is expected to be correct. However, the short-context performance trade-off still deserves explicit confirmation.

Would maintainers prefer:

1. this single representative capture hint for the Gluon path;
2. a more narrowly scoped hint for gfx950/small-head deployments;
3. a configurable vLLM capture hint; or
4. relying on an AITER-side solution such as page-table sizing or device-side active split selection?

Given the runtime-owner direction toward ROCm/aiter#4555 for kernel performance, my preference is to keep this PR narrowly focused on preventing min_kv_seq_len=1 from being frozen by vLLM FULL CUDA Graph capture.

Change

-    min_kv_seq_len: int = 1
+    min_kv_seq_len: int = 16384

No graph-dispatch, replay-path, or AITER kernel code is changed.

Test plan

• python3 -m py_compile vllm/v1/attention/backends/mla/rocm_aiter_mla.py
• Kimi-K3 TP8, gfx950, 100K input / 1K output / concurrency 1
• Confirmed FULL capture enters forward_mqa()/mla_gluon() once and runtime decode uses graph replay without re-entering either wrapper
• Confirmed runtime sequence metadata updates the same static GPU seq_info buffer used by the captured graph
• Re-run the fixed capture hint at 8K/1K and 60K/1K to quantify short/mid-context performance

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[LiuYinfeng01]

@hongxiayang @Fangzhou-Ai @tjtanaa @AndreasKaratzas Hi, could you please review this PR when you have a chance? I’d really appreciate your feedback. Thank you!

[LiuYinfeng01]

@Fangzhou-Ai Hi. Yesterday you raised the question of whether forcing min_kv_seq_len=16384 for every KV length would regress short-context performance. I measured 1, 64, 128, 256, 512, 1K, and 2K inputs against the actual CUDA-graph baseline (min_kv_seq_len=1) and confirmed TPOT regressions of 4.72%, 4.73%, 5.51%, 6.50%, 5.52%, 3.88%, and 0.32%, respectively.

I therefore updated the proposal to runtime-select five Gluon split buckets: 1/16/48/64/128. The policy is <=64 -> 1, 65..4096 -> 16, 4097..16384 -> 48, 16385..65536 -> 64, and >65536 -> 128. This remains ROCm/AITER-only. I also removed the per-layer child CUDA-graph replays: each token now uses the 25 parent graph replays plus 24 eager MLA calls instead of 25 parent + 24 child graph replays.

Decode TPOT comparison (Kimi-K3, TP8, concurrency=1, 1K output)

Speedup is fixed-1 TPOT / variant TPOT; TPOT reduction is positive when faster. Fixed-1 is the previously collected 5-request mean. Per request, fixed-128 and the final dynamic implementation were measured once as requested (the existing 5-request fixed-128 result is retained for 100K).

Input	Dynamic bucket	Fixed-1 TPOT (ms)	Fixed-128 TPOT (ms)	Fixed-128 speedup	Fixed-128 TPOT reduction	Dynamic TPOT (ms)	Dynamic speedup	Dynamic TPOT reduction
1	1	20.647	21.001	0.983x	-1.72%	20.397	1.012x	+1.21%
64	1	20.684	21.075	0.981x	-1.89%	20.867	0.991x	-0.88%
128	16	20.790	21.271	0.977x	-2.32%	20.661	1.006x	+0.62%
256	16	20.800	21.278	0.978x	-2.30%	20.638	1.008x	+0.78%
512	16	20.937	21.015	0.996x	-0.38%	20.656	1.014x	+1.34%
1,000	16	21.248	20.965	1.014x	+1.33%	20.621	1.030x	+2.95%
2,000	16	21.932	21.036	1.043x	+4.09%	20.793	1.055x	+5.19%
4,096	16	23.380	21.062	1.110x	+9.91%	20.931	1.117x	+10.47%
8,192	48	25.973	21.093	1.231x	+18.79%	21.032	1.235x	+19.03%
16,384	48	31.339	21.099	1.485x	+32.67%	21.175	1.480x	+32.43%
32,768	48	42.176	21.250	1.985x	+49.62%	21.341	1.976x	+49.40%
65,536	64	64.022	21.467	2.982x	+66.47%	21.841	2.931x	+65.89%
100,000	128	86.751	22.058	3.933x	+74.57%	22.037	3.937x	+74.60%

The one-shot 64-token dynamic point is 0.88% slower than fixed-1, within the original <=1% TPOT gate; I am calling this out rather than claiming a strict speedup at every point. Every other final dynamic point is faster than fixed-1, reaching 3.937x at 100K.

Validation: the 10 newly added bucket-selection and eager-dispatch tests pass; syntax, whitespace, and IDE lint checks also pass.

[mergify]

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

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

[LiuYinfeng01]

Already fixed in aiter pr 4555 ROCm/aiter#4555. So closed this PR.