[ROCm][MLA][DCP] Pad the AITER MLA decode head count to a native tile - #54899
Open
okorzh-amd wants to merge 1 commit into
Open
okorzh-amd wants to merge 1 commit into
okorzh-amd wants to merge 1 commit into
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
okorzh-amd
force-pushed
the
okorzh/rocm-aiter-mla-native-tile-pad
branch
2 times, most recently
from
September 2, 2026 05:50
c403adf to
7eb3f1f
Compare
Member
|
Make sure as well that a 3 line comment or something is added in the envs.py to describe your new env var and its usage (when it should be used and when not) |
AITER's MLA metadata planner folds a head count it does not natively
support down to 16 heads:
// csrc/kernels/mla/metadata/v1_2_device.cuh (v0.1.21.post1)
if (!natively_supported && (num_heads % 16 == 0)) {
qk_batch_ratio = num_heads / 16;
num_heads = 16;
num_batches *= qk_batch_ratio;
}
Every pseudo-batch of that fold re-reads the whole KV cache -- the batch
index is divided by qk_batch_ratio and the KV cursor only advances on
sub-head 0 -- so a folded decode costs num_heads/16 passes over the cache.
MLA decode is KV-bandwidth bound, so removing the fold is worth far more
than the extra query heads a pad costs.
natively_supported is not a head-count set. It is a disjunction over
(arch_id, q_is_fp8, kv_is_fp8, num_heads, max_seqlen_qo), and it has
changed in every AITER release vLLM has shipped: v0.1.20 collapsed the
qo-gated 32-head clauses, v0.1.21 (ROCm/aiter#4964) added
gfx950 && fp8 && 96 heads && max_seqlen_qo <= 6. So this does not hardcode
a tile list. It mirrors the clauses in Python and admits each one only if
its literal is present in the shipped JIT source, extending the probe
idiom already in this file (_aiter_mla_native_h24_metadata_supported,
vllm-project#51647). A clause AITER removes stops being claimed here with no vLLM
change; a clause AITER adds is picked up the same way.
Two things gate a pad target beyond the planner's verdict:
- reduce.cu's HEAD_DIM 512 instantiation list. The reducer and the planner
dispatch independently -- the same split that forced the two-probe H24
check -- and the planner's blanket gfx950/bf16 clause claims every head
count while the reducer stops at 128.
- the shipped asm kernel table, hsa/<arch>/mla/mla_asm.csv. A native
planner verdict does not imply mla_decode_fwd can dispatch:
get_heuristic_kernel_mla filters on the LSE flag exactly, and gfx942
ships gqa=128 at lse=0 only. A DCP rank asks for the LSE, so padding to
128 there would be "cannot find suitable kernel" at the first decode.
gqa=64 ships both flags, so the same rank can still pad 48 -> 64. All
four config remaps in asm_mla.cu are guarded on gfx950, where every
persistent fp8 shape lands on gqa 16 or 32 and both carry LSE variants,
so the table is consulted for gfx942 only.
The padded count is resolved once in the metadata builder's __init__,
below the dtype block, and carried on
AiterMLADecodeMetadata.padded_num_heads. It has to be a per-run constant:
max_seqlen_qo varies per build() and a head count that changed between
passes would resize o and break cudagraph capture. A target is accepted
only if it is native at *every* query length the run can produce, not at
the largest: the fold factor is num_heads/16 with no qlen term, so a
target that is non-native at some reachable qlen folds harder there than
not padding at all.
Measured on 8x MI355X (gfx950), AITER v0.1.21.post1, fp8 q + fp8 kv,
bs=52, qlen=1, 16384 KV tokens/request, kernel identity from rocprofv3:
heads AITER time padded to delta
16 native 136.6 us -- --
32 native 153.6 us -- --
48 fold 3x 349.0 us 64, 171.4 us -50.9%
64 native 171.4 us -- --
80 fold 5x 440.1 us 96, 231.7 us -47.4%
96 native 231.7 us not padded --
112 fold 7x 658.2 us 128, 247.2 us -62.4%
128 native 247.2 us -- --
The reachable win is 48 decode heads on gfx950 with an fp8 KV cache:
Kimi-K3 (96 MLA heads) at TP8+DCP4, TP4+DCP2, TP16+DCP8 or TP2+DCP1.
80 and 112 need a model with 80/112 total MLA heads; none exists in tree,
so those rows are mechanism evidence only.
What this explicitly does not do:
- It does not pad 96 heads on a current AITER. v0.1.21+ takes gqa=96
natively on gfx950/fp8 at max_seqlen_qo <= 6, and
docker/Dockerfile.rocm_base pins v0.1.21.post1 (vllm-project#52826). Padding
96 -> 128 there measures 231.7 -> 247.2 us, a 6.7% regression. An
earlier revision of this change did pad it, against v0.1.19 where 96
folded 6x; the probe is what stops that from silently rotting again.
- It does not pad under a bf16 KV cache on gfx950. The blanket
(gfx950 && !q_fp8 && !kv_fp8) clause marks every head count native, so
there is no fold to remove and a pad would be pure waste plus a
materialized q.repeat(...).contiguous(). That is the default
--kv-cache-dtype auto configuration.
- It claims no clause on an arch AITER does not name. Each clause matches
its arch guard, not just the head-count fragment, so a future AITER that
narrows one cannot produce a false positive on the other arch, and an
arch outside {gfx942, gfx950} gets the unpadded rule.
- It changes nothing at or above 129 heads, and nothing for non-multiples
of 16 (120 heads already 16-aligns to 128, which is native). AITER's
python dispatcher folds only nhead in range(32, 128+1, 16) and asserts
otherwise, so there is no target above 128; the resolver logs
warning_once instead of silently accepting the shape.
- Sparse/DSA MLA and the FP8 PS prefill are untouched. Both are different
kernel pairs with no qk_batch_ratio fold to remove.
Env: VLLM_ROCM_AITER_MLA_PAD_TO_NATIVE_SHAPE=auto|off|force, default off.
"off" is bit-identical to the previous next-multiple-of-16 rule, asserted
over every head count in 1..256. "auto" pads only when the installed AITER
reports the target native for this arch, KV dtype and every reachable
query length, has a kernel for it, and the pad stays within
_AITER_MAX_PAD_RATIO. "force" ignores that cost cap. The default flip to
"auto" is deferred until there is a gfx942 sweep and a short-KV point --
the measured curve is one operating point, and the fold/pad crossover
moves toward the fold at short KV.
The padded lanes never escape the backend: get_mla_unpadded_o and
get_mla_unpadded_lse slice them off before forward() returns, so nothing
padded reaches the DCP output combine or the LSE merge. Both now take the
resolved count explicitly, since pad and unpad branch on
m % num_heads == 0 independently and a disagreement returns wrong values
rather than raising.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Oxana Korzh <okorzh@amd.com>
okorzh-amd
force-pushed
the
okorzh/rocm-aiter-mla-native-tile-pad
branch
from
September 9, 2026 17:40
7eb3f1f to
ca03384
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The gfx950 asm persistent decode has kernels for head tiles {16,32,64,128}. aiter accepts any 16-aligned count, but a non-native one is folded:
(aiter csrc/kernels/mla/metadata/v1_2_device.cuh, mirrored in aiter/mla.py), and every sub-pass of the fold re-reads the same KV: the batch index is divided by qk_batch_ratio and the KV cursor only advances on sub-head 0.
get_actual_mla_num_heads rounds to the next multiple of 16, so a count that is already 16-aligned but not native passes straight into the fold. Under decode context parallelism the query is all-gathered to num_heads * dcp_world_size, which makes this easy to hit: Kimi-K3 at TP8/DCP8 gathers 128 = 96 heads and pays a 6x KV re-read on the largest kernel in the model. DeepSeek at TP8/DCP8 gathers 168 = 128 and is native, which is why this has gone unnoticed.
Rounding up to the next native tile instead trades a few zero query heads of compute for the entire re-read.
Measured on 8x MI355X, aiter 0.1.19, fp8 q + fp8 kv, bs=52 qlen=1, 16384 KV tokens/request, kernel identity read from rocprofv3:
NH=16 qh16_qseqlen1_gqaratio16_lse_ps 104.32 us FETCH 497.3 MB (1.01x)
NH=96 qh16_qseqlen1_gqaratio16_lse_ps 464.14 us FETCH 2970.9 MB (6.05x)
NH=128 qh32_qseqlen4_gqaratio32_lse_ps 185.26 us FETCH 517.2 MB (1.05x)
The FETCH_SIZE counter is the direct evidence of the re-read; the 1x reference is 52 * 16384 * 576 B = 490.73 MB.
End to end on Kimi-K3 DCP8/TP8, conc 52, ISL 131072, one image and two runs differing only by this knob, 8 ranks each:
off: qh16_gqaratio16, MLA 459.38 us mean (8-rank spread 4.2%), 18.9% of GPU
on : qh32_qseqlen4_gqaratio32, MLA 169.37 us (spread 2.7%), 9.6% of GPU
MLA kernel time -63.1%; median ITL 52.66 -> 43.33 ms (-17.7%); completed requests and total output tokens identical in both arms. The step-time delta (-17.3 ms/pass) is larger than the MLA delta (-7.26 ms/pass) because collective time also fell -- mscclKernel_Sum mean 57.06 -> 28.54 us at an unchanged call count -- i.e. ranks stopped waiting on the slowest MLA. Only the -7.26 is the kernel.
Off by default: it changes the kernel selected for every non-native head count, and only 96 -> 128 has been measured. 48 -> 64, 80 -> 128 and 112 -> 128 follow the same mechanism but are not measured here, and the non-gfx950 / bf16-KV combinations are not covered by aiter's native list at all.
The padded lanes never escape the backend: get_mla_unpadded_o and get_mla_unpadded_lse slice them off before forward() returns, so nothing padded reaches the DCP output combine or the LSE merge. Tests cover the rounding table and the pad/unpad round trip for o and lse.
Complementary to ROCm/aiter#4964, which teaches aiter to take gqa=96 natively (routing it to the same qh32_qseqlen4_gqaratio32 kernel) rather than padding to 128, worth a further ~15 us/call. That is merged to aiter main but is in no released aiter tag; this change works on shipped aiter.
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.