[Gluon][MLA] Size NUM_KV_SPLITS from the page table, not the optional min_kv_seq_len hint - #4507
Conversation
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
|
The bh16bn64 split count is bounded by triton.cdiv(min_kv_seq_len, BLOCK_N). min_kv_seq_len is an optional wrapper argument that defaults to 1, so any caller that does not populate it gets cdiv(1, 64) == 1 and therefore NUM_KV_SPLITS == 1. Stage-1 then launches only batch_size workgroups, each walking its entire sequence. At Kimi-K3 decode on MI355X (TP8, 12 heads/rank, batch 8, 327,600-token context) that is 8 workgroups on a 256-CU part: 3.1% occupancy, 363 GB/s effective KV bandwidth against 5.60-5.69 TB/s attainable, and _mla_gluon holding 85.9% of decode GPU time. Take the per-request context from the page table instead. Its shape is the allocated context width, which is what stage-1 actually walks, and it is host-known and shape-stable: no .item()/.cpu()/.max()/synchronize and no Python branch on a device value, so the constexpr is identical at HIP-graph capture time and on every replay. min_kv_seq_len remains accepted and is still honoured by the bh64 regime. Also replace the hardcoded 256-workgroup budget with the existing get_num_sms() helper from aiter.ops.triton.utils.device_info, which honours the CU_NUM override and is the same value the tuning dispatch keys are built from. On a 256-CU part this term is unchanged. Over-splitting a short request stays correct by the kernel's existing floor arithmetic: seq_len // NUM_KV_SPLITS is 0 for the leading splits, which early-return at the `split_kv_start >= split_kv_end` guard without writing partials; the last split absorbs the whole request and the stage-2 reduce re-derives the same division. Measured on 8x MI355X, Kimi-K3 TP8, bf16 KV, batch 8 / ctx 327,600: isolated kernel 8.31 ms -> 0.4474 ms (18.6x) in-situ per call 8241.6 us -> 270.7 us (30.4x) decode GPU share 85.9% -> 16.7% end-to-end 32.48 -> 155.85 tok/s (4.80x) median TPOT 234.9 -> 43.0 ms GSM8K unchanged at 0.9250 Scope: bh16bn64 only. bh16bn128 (batch_size == 1, fp8 KV) and bh64 keep their existing min_kv_seq_len bounds; bh64 in particular asserts on min_kv_seq_len for its gl.assume(num_iter > 3) invariant. Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: amd-ethany <ethany@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
5f71d07 to
db0e270
Compare
|
How do I reproduce these perf numbers? |
Hi @vgokhale Here is the recipe we were given: docker run --entrypoint vllm --device=/dev/kfd --device=/dev/dri and on the benchmark: aiperf profile --model moonshotai/Kimi-K3 --tokenizer moonshotai/Kimi-K3 You could reproduce the baseline with the above cmd, then patch the code in this PR additionally, PR#4509 is also optional. |
|
Can you refer to #4555 ? Seems that one supersedes this one? |
yes, we have deprecated using kv_seq_len for num_kv_split. |
|
@Fangzhou-Ai @Dewei-Wang-sh |
4 participants
mla_gluon()sizes its split-KV parallelism in thebh16bn64regime frommin_kv_seq_len, akeyword argument that defaults to
1. Any caller that does not populate it silently getsNUM_KV_SPLITS == 1and a launch ofbatch_sizeworkgroups, regardless of context length. This PRderives the split count from the page table instead, which is host-known, capture-safe, and cannot
be defeated by an unset argument.
On a Kimi-K3 TP8 serving workload this is worth 2.63x down to 1.31x inter-token latency across
concurrency 1–48 at 68k context, and 4.80x end-to-end at 327,600-token context, with GSM8K
unchanged. Host-side wrapper only; both
@gluon.jitkernels are byte-identical.Motivation
The split count is bounded by
triton.cdiv(min_kv_seq_len, BLOCK_N).min_kv_seq_lenis optionaland defaults to
1, so a caller that does not set it getscdiv(1, 64) == 1and thereforeNUM_KV_SPLITS == 1. Stage-1 then launches onlybatch_sizeworkgroups, each walking its entiresequence.
At Kimi-K3 decode on MI355X (TP8, 12 heads/rank, batch 8, 327,600-token context) that is 8
workgroups on a 256-CU part: 3.1% occupancy, 363 GB/s effective KV bandwidth against 5.60–5.69 TB/s
attainable, and
_mla_gluonholding 85.9% of decode GPU time.The failure is silent. There is no warning and no assertion, and because
NUM_KV_SPLITS == 1takes afast path that writes
odirectly, the stage-2 reduce disappears from the profile entirely — so theusual signal that split-KV is misconfigured is absent precisely when it is misconfigured. vLLM's ROCm
aiter-MLA backend is such a caller today: it declares the field on its decode metadata and never
assigns it.
Technical Details
Take the per-request context from the page table. Its shape is the allocated context width, which is
what stage-1 actually walks, and it is host-known and shape-stable — no
.item()/.cpu()/.max()/ synchronize, and no Python branch on a device value — so theconstexpris identical atHIP-graph capture time and on every replay.
min_kv_seq_lenremains accepted and is still honoured by thebh64regime.This also replaces the hardcoded 256-workgroup budget with the existing
get_num_sms()helper fromaiter.ops.triton.utils.device_info— the module this file already imports from. It honours theCU_NUMoverride and is the same value the tuning dispatch keys are built from. On a 256-CU part theterm is unchanged.
Correctness with ragged batches. Over-splitting a short request is safe on the kernel's existing
floor arithmetic, which this patch does not touch: splits that come out empty early-return at the
split_kv_start >= split_kv_endguard without writing their partials, and the stage-2 reducere-derives the same bound from the same
seq_info, so uninitialised slots are never read. The guardis evaluated per request, so a batch mixing
seq_len=1withseq_len=327600is safe.Determinism. Splitting changes the order of the KV reduction, so output is no longer
bit-identical to the unsplit path. This is float association, not a different computation: the
split-K online-softmax merge is the same reduction, and accuracy is gated — the frozen fp32 oracle
passes at tol 2e-2 and GSM8K is unchanged at 0.9250. A fixed launch shape stays deterministic
run-to-run. What it does cost is batch invariance:
NUM_KV_SPLITSdepends onbatch_size, so thesame request can differ in the last bits at a different batch size, where
bh16bn64previously alwaystook one split. Note that
mla_gluon()has nonum_kv_splitsargument to pin it — that knob is onaiter.mla's non-Gluon entry point, which is whereVLLM_BATCH_INVARIANTpins it. Happy to add anoverride here if maintainers want one.
Scope.
bh16bn64only.bh16bn128(batch_size == 1, fp8 KV) andbh64keep their existingmin_kv_seq_lenbounds;bh64in particular asserts onmin_kv_seq_lenfor itsgl.assume(num_iter > 3)invariant. No GPU code changes — both@gluon.jitkernels arebyte-identical and the diff is entirely in the host-side wrapper.
Test Plan
aiter's own MLA tests. Both configs run on the rebased branch and pass — every
checkAllcloseverdict green, 1 on the first and 10 on the second:
Worth being precise about what this does and does not show.
test_absorb_decode_gluon_bh16callsmla_gluonwithmin_kv_seq_len=ctx_lens, i.e. it populates the hint correctly, so it does notreproduce the defect — it is a regression check rather than a demonstration. That makes the
informative comparison the one against the branch's own pre-patch parent, run back to back in the
same container on the same GPU:
-c 327600 -b 8 -n 12,1-c 16384 -b 64 128 -n 64,1 128,1The two agree to within 0.5%, which is the expected and desired result: when a caller supplies
min_kv_seq_lencorrectly,cdiv(pt_tokens_per_req, BLOCK_N)resolves to the same bound, so thepatch selects the identical split count and nothing moves. What it establishes is that the
page-table-derived value agrees with a correctly-supplied hint, and that the large-
nhead/large-batch shapes are unaffected.
The verdict counts differ because the second command enumerates four configurations (two batch sizes
× two head counts) and checks several implementations in each. Only 3 of the 11 verdicts exercise
bh16bn64, the one kernel this patch affects — the small-nheadconfig and the twonhead=64ones;nhead=128dispatches elsewhere. The remaining verdicts coveraiter_asmandgluon_mlaand showonly that nothing adjacent broke.
Isolated kernel. Device-event timing inside the deployment graph context, fresh subprocess per
case, 3 random draws, harness spread ≤ 0.53%. Correctness gated on every run against a frozen,
implementation-independent fp32 online-softmax reference at tol 2e-2, plus a parity check against the
unmodified production kernel and a capture-once / replay-many check over boundary shapes.
Two adversarial ragged cases under replay:
seq_lens = [1, 63, 64, 65, 4096, 4097, 65536, 327600]— straddlesBLOCK_N=64, bottoms out at 1seq_lens = [308736, 308737, 312831, 312832, 312833, 320000, 327599, 327600]— straddles theshared-prefix and block boundaries
End-to-end, concurrency sweep (primary). 8x MI355X, Kimi-K3 TP8, bf16 KV,
aiperfagainst alive vLLM server. Agentic-proxy workload: 8 shared prefixes of 63,240 tokens plus 4,760 unique input
tokens per request (68,089-token mean measured ISL), 350 output tokens pinned with
ignore_eosandmin/max_tokens, prefix caching on,cudagraph_mode=FULL_AND_PIECEWISE, seed 42. Concurrency 1, 8,16, 24, 48 with 5, 40, 80, 120, 240 requests respectively. Each leg is a fresh server on the same
node in the same container; 4
aiperfprofile runs per point with a 15 s cooldown, 100 trial recordstotal. Candidate legs were confirmed by server-log grep to have loaded the patched module and the
baseline leg confirmed not to have.
Ratios are formed within a trial index and then averaged, so prefix-cache history is matched on
both sides of every ratio.
End-to-end, long context (secondary). vLLM serving benchmark: ISL 18,264 plus a 308,736-token
shared random prefix (327,600-token prompts), OSL 1,200, concurrency 8, 16 prompts, seed 42, prefix
caching on. Reference and candidate legs on the same node in the same container, serialized on a GPU
lock, both to completion. A separate locked baseline measured 34.897 tok/s over 3 timed rounds with
0.06% spread.
Accuracy. GSM8K, 200 problems, run on both legs of the same paired session.
Test Result
1. Concurrency sweep at 68k context — this patch alone
Inter-token latency, and the input throughput that follows from it. Brackets are the min–max over
trial indices.
NUM_KV_SPLITSis what this patch selects; stock is always 1.Every baseline trial was worse than every candidate trial at every concurrency — the ITL P50
distributions are disjoint, Welch p<1e-12 throughout. Work parity was checked rather than assumed:
completed requests, mean ISL and mean OSL are identical across legs at every point, with zero errors
anywhere.
The speedups above are trial-index-paired ratios, but they do not depend on that choice: taking the
ratio of means instead agrees to three decimals at every point except concurrency 8 — the
prefix-cache boundary point — where it differs by 0.02x on ITL and 0.05x on throughput.
2. Long context, single point — batch 8 / ctx 327,600
* the two in-situ rows come from a profile that also carried a follow-up grid/epilogue change, so
they are not attributable to this patch alone. Every other row is this patch by itself.
Correctness: frozen fp32 oracle pass, parity vs the unmodified kernel pass, cudagraph
replay over boundary shapes pass, both ragged cases above pass. 1 of 12 greedy probes matched
byte-for-byte, as expected from the reduction-order change.
On the bandwidth row: 6.75 TB/s sits above the 5.60–5.69 TB/s this box actually attains (two
independent 8 GB read/copy probes past the LLC; the 8.0 TB/s nameplate is ~40% optimistic). That is
not an error — the 256 MiB LLC is already servicing part of the shared prefix, so the kernel moves
fewer distinct bytes than the naive 3.019 GB per call.
Caveats — both measurements
Shared-prefix warm-up is controlled for in both — the sweep pairs on trial index, the long-context
point matches cold-to-cold and warm-to-warm.
AI assistance
This was found by GEAK, AMD's agentic GPU-kernel optimization
framework, running an end-to-end optimization pass against a live Kimi-K3 TP8 server on 8x MI355X.
Everything in this PR was verified manually before submission.
Submission Checklist