Repository navigation
Conversation
…8 on gfx950 - Raise the native 96-head limit from qlen 6 to 8 in the planner, the asm remap, the opus route and the Python table - The qh32_qseqlen4 tile takes one token per 128-row tile, so qlen 7 and 8 need no kernel change, only the dispatch - Before, 96 heads at qlen 7-8 folded into six 16-head pseudo-batches; Kimi-K3 DSpark block 7 verifies 8 tokens - test_mla_shape_support pins native at qlen 6-8 and the fold at 9, and checks the launch agrees with the verdict there - The round-robin test also sweeps 96 heads at qlen 7
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
One backend per PR: PR title tags & labels: |
WAITING ON AUTHORPR Self-ReviewRevision Waiting for the PR author. The readiness report will appear after the self-review answers are complete for this revision. Questions to answer in the PR description
Edit the Missing or invalid answers
Shadow mode reports what would block but keeps the status successful. Set |
…eference - test_native_96_head_fp8_decode_matches_reference runs the native route against an fp32 causal MLA reference on gfx950 - It fails on main, where 96 heads at qlen 7-8 still fold, and passes here - _run_decode returns its output and inputs for the comparison - Format test_mla_persistent_round_robin.py's default head list with Black
…robin CP - Without CP the fold beats native from batch 16 up and under a fixed split cap at long context, so it stays there - mla_decode_head_plan_v1_2 and mla_decode_head_plan take cp_round_robin; the planner and the head-plan query pass it - The asm remap and the opus route still accept 96 heads at qlen <= 8: without CP the planner folds before them - get_mla_metadata_info_v1 cannot tell CP apart, so it sizes for both plans: 8 qo tiles per request at qlen 8, not 48 - Document that is_cp_round_robin must match g_kv_indptr and that decode_update_mla_metadata_v1 is not for CP - test_mla_shape_support: verdicts, planner agreement and launches with and without CP, a CP reference check, sizing
yuychang
changed the title
yuychang
changed the title
There was a problem hiding this comment.
🟡 Changes recommended
The metadata-size test uses gfx950-specific expectations without GPU gating, causing it to reject correct legacy sizing on gfx942.
1 open issue
What changed in this PR
Enables native 96-head FP8 MLA decode for query lengths 7–8 under round-robin context parallelism on gfx950, avoiding head-folding overhead while preserving the non-CP route.
Changes:
- Passes the CP flag through head planning and extends ASM/OPUS routing.
- Sizes metadata for both native and folded plans and documents the metadata/decode flag contract.
- Adds boundary, dispatch, metadata-sizing, and numerical-reference coverage.
| File | Description |
|---|---|
op_tests/test_mla_shape_support.py |
Adds CP boundary, sizing, dispatch, and reference checks. |
op_tests/test_mla_persistent_round_robin.py |
Adds 96-head query-length-7 coverage to the default sweep. |
csrc/py_itfs_cu/asm_mla.cu |
Extends the 96-head ASM route through query length 8. |
csrc/kernels/mla/opus_mla_decode_fwd.cu |
Extends the matching OPUS route. |
csrc/kernels/mla/metadata/v1_2_device.cuh |
Passes round-robin CP into device metadata head planning. |
csrc/kernels/mla/metadata.cu |
Threads the CP flag into shape-support planning. |
csrc/include/mla_decode_shape.h |
Gates native 96-head query lengths 7–8 on gfx950 FP8 round-robin CP. |
aiter/ops/attention.py |
Sizes metadata across both plans and clarifies CP usage constraints. |
Review effort: Balanced
…obin CP mode - get_mla_metadata_v1 records the is_cp_round_robin it planned with on work_metadata_ptrs - mla_decode_fwd raises ValueError when g_kv_indptr disagrees and that changes the head plan (96 fp8 heads, qlen 7-8) - Callers that agree, or that pass another tensor as work_meta_data, are unaffected - test_decode_rejects_metadata_planned_for_the_other_cp_mode plans for one mode and launches in the other
- The expected 7/8 qo tiles per request only hold where 96 fp8 heads run natively, which is gfx950 only. - On gfx942 the legacy bound keeps 42/48 tiles at qlen 7/8, so the assertion rejected correct sizing on MI300X.
- The incremental update grows every request's KV by one token and plans heads without CP, so it miscomputes CP metadata. - For 96 fp8 heads at qlen 7 and 8 it would also rewrite a native plan as the 6x16 fold, which mla_decode_fwd cannot detect. - Raise ValueError when work_metadata_ptrs was planned with is_cp_round_robin=True instead of only documenting the limit. - Add a test that plans 96-head fp8 qlen-8 metadata under CP and expects the update to raise.
There was a problem hiding this comment.
🔵 Needs a closer look
The gfx950 native CP route and reduced metadata bounds require final-revision GPU validation, since reported sweeps predate follow-up commits, include tolerance failures, and cannot be rerun in this environment.
No open issues 🎉
2 resolved since last review
Review effort: Balanced
- Build the head-plan argument tuple once in mla_decode_fwd so the CP-mismatch recheck cannot drift from the launch plan - Note in get_mla_metadata_v1 that the planned CP flag lives on the tensor object, so a view or copy skips the check
There was a problem hiding this comment.
🟡 Changes recommended
Deep-copied buffer views lose the CP-mode record, causing valid metadata to be rejected at the 96-head FP8 CP boundary; copied-storage identity tracking and tests for both copy/view compositions are needed.
1 open issue
Review effort: Balanced
- mla_decode_head_plan_v1_2 takes both main's page_size and this PR's cp_round_robin, in the outer planner's order - test_mla_shape_support: _fold forwards both page_size and the round-robin CP kwargs; keep get_gfx for the gfx950 gates
…hanges the head plan - A clone has no recorded is_cp_round_robin; mla_decode_fwd without g_kv_indptr ran it as non-CP metadata - decode_update_mla_metadata_v1 did the same, so a cloned CP plan for 96 fp8 heads at qlen 7-8 got the fold ratio 6 - Both now raise ValueError there, since either plan could be wrong; other shapes are unchanged - A view of a deep copy finds the recorded mode on its base, so it no longer reads as unknown - Tests cover cloned metadata in both consumers, a view of a deep copy, and a deep copy of a view
…ut round-robin CP - The native 96-head qlen 7-8 route is gfx950 only, and its own tests skip on gfx942 - Pin arch=gfx942 and assert fold x6 at qlen 1-9 for e4m3fnuz and bf16, with and without CP - Needs no GPU, so every CI runner covers the gfx950-only gate
|
gfx942 (MI300X) validation: validate-kernel-pr PASS. On gfx942 96 heads stay folded x6 at every qlen, identical on base and head; new New commits for gfx942: d6e5dc3. Reports were generated with the repo's validate-kernel-pr skill on a local head, so |
…copy - A deep copy carries its mode cell on the tensor but not on its new storage - Replanning the copy through a view made a second cell on the storage and left the copy's own cell stale - A later copy.deepcopy of the copy then recorded the old mode, so 96 fp8 heads at qlen 7-8 could use the wrong plan - Adopt the tensor's (or same-address base's) existing cell when registering it on the storage - Add copy, view-replan, copy cases for False to True and True to False; both fail without the change - The buffer a copy was made from still keeps its own mode, and a clone still has none
There was a problem hiding this comment.
🟡 Changes recommended
CP-mode records can remain stale after in-place overwrites or be incorrectly inherited by offset slices, causing decode to select incompatible execution and potentially bypass incremental-update safeguards.
2 open issues
Review effort: Balanced
- The lookup fell back to the base's record for any view, so workspace[2:4] read the record of workspace[0:2] - Copying native CP metadata into such a slice then saw a False record and could pick the folded head plan - Use the base's record only when the view starts at the base's own address, as the registering side already does - Document that a copy_ into a planned buffer keeps that buffer's own record; plan the destination instead - Add an offset-slice case: it has no record until planned, and planning it leaves the base's record alone - The case fails without the change; no code in the repo copies planned metadata with copy_
There was a problem hiding this comment.
🔵 Needs a closer look
The gfx950 dispatch and metadata/workspace changes require exact-head GPU correctness and latency validation, which prior run summaries do not establish and this environment cannot perform.
No open issues 🎉
2 resolved since last review
Review effort: Balanced
Summary
qh32_qseqlen4kernel.Motivation
Changes
csrc/include/mla_decode_shape.h,metadata.cu,v1_2_device.cuh)cp_round_robininput. 96 FP8 heads stay native up to qlen 8 only under round-robin CP (was qlen ≤ 6).asm_mla.cu,opus_mla_decode_fwd.cu)get_mla_metadata_info_v1)get_mla_metadata_v1,mla_decode_fwd)get_mla_metadata_v1recordsis_cp_round_robinwithwork_metadata_ptrs(on its storage, so views and deep copies keep it).mla_decode_fwdruns the head plan the record names;g_kv_indptronly selects the CP kernel. It raisesValueErroronly where CP changes the head plan and the mode is unknown (a clone/copy_of the buffer) or CP-planned metadata gets nog_kv_indptr.decode_update_mla_metadata_v1ValueErroron metadata planned withis_cp_round_robin=Trueonly where CP changes the head plan (96 fp8 heads, qlen 7-8, gfx950); elsewhere it updates as before (KV grows by1 - num_reject_tokens).Native-route condition (v1_2 / fast-mode planner):
Performance (if applicable)
Kernel setup (author): MI355X, DCP8 rank 0, query length 8, FP8 Q/KV, page size 1, LSE, AITER-planned KV splits. Median of 20 complete
mla_decode_fwdcalls, each after a 1 GiB cache flush. SGLang's padding copy is excluded. Baseline and PR alternate twice on one GPU; the table shows the mean; the two rounds differ by at most 3.9%.Re-measured at
9ee92f116(review run): MI355X (gfx950), one GPU,op_tests/test_mla_persistent_round_robin.py -d fp8 -kvd fp8 -n 96,7 96,8 -b 1 16 56 -c 16384 65536 -cpw 8 -ms 32 -lse,cp:rank_uscolumn, base274a74024(merge base) vs head, 2 alternating rounds, mean. Same-harness per-rank latency, no cache flush; not comparable to the table above.Testing
Model accuracy (GSM8K): SGLang Kimi-K3, TP8/DCP8, DSpark block 7, FP8 KV, AITER-planned splits; sgl-eval 0.1.3, K3 preset, all 1,319 questions, real acceptance.
Kernel checks found native output bitwise equal to the first 96 heads of padded output at every benchmark point.
Documentation
Docstrings state the CP metadata/decode flag contract and the incremental-update limit. User and performance guides are not changed.
Dependencies
Breaking Changes
ValueErrors on misuse that was silently wrong before:mla_decode_fwd, at 96 fp8 heads and qlen 7-8 only, when the planned CP mode is unknown (a cloned buffer) withg_kv_indptr, or CP-planned metadata gets nog_kv_indptr.decode_update_mla_metadata_v1, on metadata planned withis_cp_round_robin=Trueat 96 fp8 heads, qlen 7-8.PR Self-Review
Reviewed revision
9ee92f1164acfc61a0a8148c123644e8bc422e29Behavior and intent
Review findings (AITER
review-pr+code-review, atcd1d49c4e)mla_decode_fwdrepeated the 10-argument launch-plan call for the CP-mismatch check9ee92f116(shared argument tuple)work_metadata_ptrstensor object, so a view or copy skips the checkc08e32485: the mode is recorded on the buffer's storage; an unknown mode under CP raises instead of running a mismatched plandecode_update_mla_metadata_v1rejects all CP metadatac08e32485: rejected only where CP changes the head plancd1d49c4eValidation and unverified scope
At
9ee92f116on MI355X (gfx950), one GPU:validate-kernel-pronop_tests/test_mla_shape_support.py, basef60ee45dd(main) vs headpytest op_tests/test_mla_shape_support.py op_tests/test_mla_decode_update_metadata.pytest_mla_persistent_round_robin.py -n 96,6 96,7 96,8 128,8 -b 1 32 -c 1200 8192 -cpw 8 -ms 32 -lseEarlier author validation (at
0e47dbf88/cd1d49c4e):AITER_MLA_USE_OPUS=1), 96 heads at qlen 7–8 and 128 at qlen 8cd1d49c4etest_mla_decode_update_metadata.pyAITER_ENABLE_EXPERIMENTAL, the affected CP shape now selects ASM instead of HK.PR-specific question
is_cp_round_robin=Trueto get the native 96-head CP route; metadata planned without it still runs the fold under CP, as before.mla_decode_fwdraises only for an unknown mode (cloned buffer) at the affected shapes.Changed op tests
op_tests/test_mla_shape_support.py: CP/non-CP boundary verdicts, planner/dispatch agreement, metadata bounds, rejection of metadata planned for the other CP mode, rejection of CP metadata bydecode_update_mla_metadata_v1, and output/LSE against an FP32 causal reference on CP rank 3 of 8.op_tests/test_mla_persistent_round_robin.py: adds(96, 7)to the default sweep, so the new lower boundary runs through the CP pipeline.