Repository navigation
[ROCm][Perf] Drop two dead per-step ops from the AITER sparse MLA metadata build - #58354
Merged
dllehr-amd merged 3 commits intoOct 14, 2026
Merged
dllehr-amd merged 3 commits into
dllehr-amd merged 3 commits into
Conversation
1 task done
ZhengGong-amd
marked this pull request as ready for review
September 30, 2026 11:16
ZhengGong-amd
requested review from
AndreasKaratzas and
tjtanaa
as code owners
September 30, 2026 11:16
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
ZhengGong-amd
force-pushed
the
rocm-sparse-mla-metadata-dead-ops
branch
2 times, most recently
from
September 30, 2026 12:10
1853574 to
b88ba14
Compare
simondanielsson
left a comment
simondanielsson
left a comment
Contributor
There was a problem hiding this comment.
LGTM, can we attach gsm8k before and after just to validate we're not breaking anything (even though I doubt it :) )
Contributor
Author
e2e results, per runMI355X (gfx950), torch Every run starts a fresh server and waits for vllm serve zai-org/GLM-5.3-Flash --tensor-parallel-size 2 --kv-cache-dtype auto --max-model-len 8192 \
--no-enable-prefix-caching --gpu-memory-utilization 0.90 --trust-remote-code
vllm bench serve --backend vllm --model zai-org/GLM-5.3-Flash --dataset-name random --random-input-len 1031 \
--random-output-len 509 --random-range-ratio 0.25 --num-prompts 509 --max-concurrency 61 \
--request-rate inf --ignore-eos --seed 1234
|
Contributor
Author
|
@simondanielsson — GSM8K on the only end-to-end model that reaches this GSM8K resultsMI355X (gfx950), torch Each arm starts a fresh server, waits for vllm serve zai-org/GLM-5.3-Flash --tensor-parallel-size 2 --kv-cache-dtype auto \
--max-model-len 4096 --no-enable-prefix-caching --gpu-memory-utilization 0.90 --trust-remote-code
lm_eval --model local-completions --tasks gsm8k --num_fewshot 5 --batch_size 1 --seed 1234 \
--model_args model=/shared_nfs/models/GLM-5.3-Flash,base_url=http://127.0.0.1:PORT/v1/completions,\
tokenizer=/shared_nfs/models/GLM-5.3-Flash,tokenizer_backend=huggingface,trust_remote_code=True,\
tokenized_requests=False,num_concurrent=64,max_retries=5,max_gen_toks=256,timeout=3600 \
--log_samples
Per-sample strict-match agreement with main: 1235 / 1319 (84 differ). |
simondanielsson
approved these changes
Oct 1, 2026
simondanielsson
left a comment
simondanielsson
left a comment
Contributor
There was a problem hiding this comment.
LGTM
I'm aware this small perf change can't really be measured e2e but imo it's good we're cleaning up the metadata build (as e.g. for GLM-5.3-Flash it's taking a non-trivial amount of time per step)
simondanielsson
added
the
ready
|
✅ @ZhengGong-amd, CI is now available for this PR.
|
Contributor
|
/ci run |
|
✅ Triggered Buildkite CI #93880 for commit |
Contributor
Author
|
Thank you @simondanielsson ! The CI test passed. Could you please merge this PR if you think this is ready :-) |
Collaborator
CI selector (shadow): 3 test steps (10 jobs) instead of 75 (99 jobs)Shadow mode: this changes nothing about what CI runs. It shows what the evidence-based selector would pick for this PR, next to today's rules. How it works. Feedback welcome: reply here if it would skip a step this change needs, or runs something unrelated.
Selector would run (3)
Would skip (today's rules run them) (74)
Would add (today's rules do not run them) (2)
AMD mirrors: would skip (67)
AMD mirrors: would add (3)
2 changed files · base |
Contributor
|
/ci run |
|
❌ This PR is 1 commit behind upstream |
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
…adata build Generated-by: Hyperloom Co-authored-by: Claude Signed-off-by: ZhengGong-amd <zgong@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
ZhengGong-amd
force-pushed
the
rocm-sparse-mla-metadata-dead-ops
branch
from
October 13, 2026 01:54
abc4855 to
6586df2
Compare
Contributor
Author
|
/ci run |
|
✅ Triggered Buildkite CI #94351 for commit |
Contributor
Author
|
/ci run |
|
✅ Triggered Buildkite CI #94426 for commit |
Contributor
Author
|
/ci run |
|
✅ Triggered Buildkite CI #94455 for commit |
AndreasKaratzas
pushed a commit
to AndreasKaratzas/vllm
that referenced
this pull request
Oct 16, 2026
…adata build (vllm-project#58354) Signed-off-by: ZhengGong-amd <zgong@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: Andreas Karatzas <AndreasKaratzas@users.noreply.github.com>
michael-gpu
pushed a commit
to practical-inference/vllm-ampere-extended
that referenced
this pull request
Oct 20, 2026
…adata build (vllm-project#58354) Signed-off-by: ZhengGong-amd <zgong@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
5 of 6 tasks
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
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.
Purpose
ROCMAiterMLASparseMetadataBuilder.build()runs once per KV-cache group per scheduler step and launches two GPU ops whose results are not needed:query_lens.build()computesquery_start_loc[1:] - query_start_loc[:-1], butgenerate_sparse_seqlen_tritononly reads its size and device; the kernel derives per-request lengths fromcu_query_lens. The parameter is dropped andnum_seqsis taken ascu_query_lens.size(0) - 1.fill_.paged_kv_indptr[n + 1:].fill_(paged_kv_indptr[n])passes a view of the tensor being filled, so ATen clones the value first. A broadcastcopy_from the same element writes the same values without the clone.x[n+1:].fill_(x[n])fill_→clone(empty_strided+ copy) →copy_t.copy_(x[n].expand_as(t))copy_qsl[1:] - qsl[:-1]sub(+ alloc)Each
build()launches two fewer kernels and makes two fewer device allocations, with no change in output. #57885 removed a similar dead op from the DSA indexer builder.On ROCm, sparse-MLA layers that do not bring their own backend select
ROCM_AITER_MLA_SPARSE. Of 18 current checkpoints probed on main (dummy weights, 8 layers, TP=1), three reach thisbuild():build()per decode stepglm_moe_dsa)DeepseekV32MLAAttention→ROCMAiterMLASparseBackendglm_moe_dsa)glm5_next, DSA every 4th layer)MLAAttention→ROCMAiterMLASparseBackendDeepSeek-V4-Flash/Pro and V4.1-Flash use their own builders, and the other 12 have no sparse MLA. CUDA never selects this backend.
Test Plan
pytest tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py. The newtest_build_fills_paged_kv_indptr_tailpoisons thepaged_kv_indptrtail and checks that it reads as the running total; it fails if the tail copy is removed.paged_kv_indptrbuffer (tail poisoned first), the returned metadata tensors and the AITER work/reduce metadata.build()on a builder constructed from a realVllmConfig(max_num_batched_tokens=8192, top-k 2048, decode B ∈ {1, 7, 29, 61, 127}). Both implementations alternate on one builder, order flipped every pair, 300 warmup + 3000 pairs, in 3 processes. Steady: contexts > top-k, so the persistent-metadata key hits. Growing: the key misses, soget_mla_metadata_v1and a stream sync run on every build.--kv-cache-dtype auto(fp8 KV faults on main withHSA_STATUS_ERROR_MEMORY_FAULT, before and independent of this PR), a fresh server per run, 4 runs per arm with the first discarded.vllm serve zai-org/GLM-5.3-Flash --tensor-parallel-size 2 --max-model-len 8192 --no-enable-prefix-caching --trust-remote-code vllm bench serve --model zai-org/GLM-5.3-Flash --dataset-name random --random-input-len 1031 --random-output-len 509 \ --random-range-ratio 0.25 --num-prompts 509 --max-concurrency 61 --request-rate inf --ignore-eos --seed 1234--kv-cache-dtype auto(fp8 KV faults on main withHSA_STATUS_ERROR_MEMORY_FAULT, before and independent of this PR), a fresh server per run, 4 runs per arm with the first discarded.vllm serve zai-org/GLM-5.3-Flash --tensor-parallel-size 2 --max-model-len 8192 --no-enable-prefix-caching --trust-remote-code vllm bench serve --model zai-org/GLM-5.3-Flash --dataset-name random --random-input-len 1031 --random-output-len 509 \ --random-range-ratio 0.25 --num-prompts 509 --max-concurrency 61 --request-rate inf --ignore-eos --seed 1234lm_eval --model local-completions --tasks gsm8k --num_fewshot 5, GLM-5.3-Flash, TP=2,max_model_len=4096,--kv-cache-dtype auto,seed=1234.Test Result
Unit: 5 passed on MI355X (gfx950), main
863475eb9c. Output hashes match between the arms for every config and batch size.Micro, MI355X (gfx950), torch 2.12.0, main
c961121519:The saving is a flat ~10.5 µs per call, independent of batch size and decode path; this PR is faster in 88–99 % of pairs. A blocked cross-process A/B on main
72e7874fa6(4 runs per arm, the first discarded, median of runs 2–4) agrees: 10.6–12.3 µs saved when steady, 13.0–15.8 µs when growing, 12.9–13.4 µs on the non-persistent path, with identical output hashes. At one or two builds per decode step this is under 0.1 % of a GLM-5.3 decode step, below end-to-end noise.e2e, GLM-5.3-Flash (Triton sparse path), TP=2, 509 prompts in 1031 / out 509, concurrency 61, main
72e7874fa6. The difference is within run-to-run spread, so no end-to-end speedup is claimed; per-run numbers are in the first comment.GSM8K (strict-match), GLM-5.3-Flash, TP=2, main
72e7874fa6:Per-sample strict-match agreement with main is 1235/1319.
Found by Hyperloom; reviewed and measured by hand.