[perf] wire FA and FlashMLA for sm90 GLM5Next NoPE SparseMLA - #55385
Open
JaredforReal wants to merge 1 commit into
Open
JaredforReal wants to merge 1 commit into
JaredforReal wants to merge 1 commit into
Conversation
Contributor
|
Important Draft PR not reviewedDraft PRs are not automatically reviewed by default.
To automatically review draft PRs, update your CodeRabbit configuration: reviews:
auto_review:
drafts: trueThanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
This was referenced Sep 6, 2026
JaredforReal
added a commit
to JaredforReal/vllm
that referenced
this pull request
Sep 7, 2026
…tention prefill backend GLM-5.3-Flash (qk_nope 256, qk_rope 0, v 256) was not in FlashAttnPrefillBackend.supports_mla_dimensions, so sparse MLA fell back to "No MLA prefill backend supports this model" and every prefill token went through the per-token top-k MQA kernel. The kernels for qk_head_dim 256 / v_head_dim 256 are the ones already used for the (192, 64, 256) layout. Same 7 lines as in the SM90 wiring PR (vllm-project#55385); kept as a separate commit so it can be dropped once that lands. Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Jared Wen <jaredwen@inferact.ai>
This was referenced Sep 7, 2026
JaredforReal
force-pushed
the
wire_sm90_kernels
branch
from
September 8, 2026 13:44
325aba7 to
db053be
Compare
JaredforReal
marked this pull request as ready for review
September 8, 2026 13:46
JaredforReal
requested review from
Harry-Chen,
LucasWilkinson,
pavanimajety and
tlrmchlsmth
as code owners
September 8, 2026 13:46
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
There is a correctness risk in the FlashMLA sparse tuple-input path (actual_num_heads handling) and the newly introduced head_size=512 NoPE behavior needs targeted test coverage to prevent regressions.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
This PR extends vLLM’s sparse-MLA attention backend support for SM90 to cover GLM5Next “NoPE” geometry, and updates pinned external dependencies to pick up newer FlashAttention/FlashMLA implementations.
Changes:
- Add GLM5Next NoPE MLA dimension support to the FlashAttention prefill backend selector.
- Extend FlashMLA sparse backend to advertise
head_size=512, reject quantized KV-cache formats for that geometry, and handlerope_dim==0(NoPE) query concatenation. - Bump pinned commits for
vllm-project/flash-attentionandvllm-project/FlashMLA.
File summaries
| File | Description |
|---|---|
| vllm/v1/attention/backends/mla/prefill/flash_attn.py | Adds GLM5Next NoPE MLA dimensions to FlashAttn prefill backend support list. |
| vllm/v1/attention/backends/mla/flashmla_sparse.py | Adds head_size=512 support and NoPE handling in FlashMLA sparse forward path plus a guard for quantized KV-cache formats. |
| cmake/external_projects/vllm_flash_attn.cmake | Updates the pinned FlashAttention commit used by the build. |
| cmake/external_projects/flashmla.cmake | Updates the pinned FlashMLA commit used by the build. |
Review details
- Files reviewed: 4/4 changed files
- Comments generated: 2
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
1001
to
1012
| actual_num_heads = self.num_heads | ||
| if isinstance(q, tuple): | ||
| ql_nope, q_pe = q | ||
| q = self.q_concat_buffer[: ql_nope.shape[0]] | ||
| ops.concat_mla_q(ql_nope, q_pe, q) | ||
| if q_pe.size(-1) == 0: | ||
| # NoPE (GLM5Next): concat_mla_q requires rope_dim == 64, | ||
| # copy directly into the head-padded buffer instead. | ||
| q[:, : ql_nope.shape[1]].copy_(ql_nope) | ||
| else: | ||
| ops.concat_mla_q(ql_nope, q_pe, q) | ||
| else: | ||
| actual_num_heads = q.shape[1] |
Comment on lines
137
to
+139
| def get_supported_head_sizes(cls) -> list[int]: | ||
| # DeepSeek V3.2 layout: 512 NoPE + 64 RoPE = 576. | ||
| return [576] | ||
| # DeepSeek V3.2: 512 NoPE + 64 RoPE = 576; GLM5Next NoPE: 512. | ||
| return [576, 512] |
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
| # DeepSeek V3.2 layout: 512 NoPE + 64 RoPE = 576. | ||
| return [576] | ||
| # DeepSeek V3.2: 512 NoPE + 64 RoPE = 576; GLM5Next NoPE: 512. | ||
| return [576, 512] |
Contributor
There was a problem hiding this comment.
this is unconditionally done and doesn't guard on compute capability.
i think it may reorder priority for sm_100 backend as follows:
| capability | heads/rank | before PR | after PR |
|---|---|---|---|
| sm100 | 16 (TP=4) | FLASHINFER_MLA_SPARSE | FLASHINFER_MLA_SPARSE |
| sm100 | 32 (TP=2) | FLASHINFER_MLA_SPARSE | FLASHMLA_SPARSE |
| sm100 | 64 (TP=1) | FLASHINFER_MLA_SPARSE | FLASHMLA_SPARSE |
| sm90 | any | FLASHINFER_MLA_SPARSE_SM90 | FLASHINFER_MLA_SPARSE_SM90 |
is this intended?
Contributor
|
i get different numbers than you for the backends. could you please recheck? Setup: zai-org/GLM-5.3-Flash on 4x H200 (sm90), TP=4 (16 heads per rank), vLLM at W1: 32k in / 512 out, conc 32, 96 req
W2: 4k in / 1k out, conc 128, 384 req
Bench script#!/usr/bin/env bash
# Benchmark one sparse MLA attention backend on GLM-5.3-Flash (TP=4) with the two
# workloads from the PR #55385 description. Run under the GPU reservation tool:
# chg run --gpu-ids 0,1,2,3 --timeout 40m -- scripts/bench_sparse_backend.sh <BACKEND> <PORT> <RESULT_DIR>
# e.g. scripts/bench_sparse_backend.sh FLASHMLA_SPARSE 8802 logs/bench/FLASHMLA_SPARSE
set -x
BACKEND=${1:?backend}
PORT=${2:?port}
DIR=${3:?result dir}
ROOT=/home/simon-veitner-redhat/dev/glm5-sparse-mla-sm90
mkdir -p "$DIR"; DIR=$(realpath "$DIR"); cd "$DIR"
source $ROOT/.venv/bin/activate
export TMPDIR=/home/simon-veitner-redhat/tmp; mkdir -p $TMPDIR
echo "CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES"
vllm serve zai-org/GLM-5.3-Flash -tp 4 --max-model-len 36000 --attention-backend $BACKEND --port $PORT --no-enable-log-requests > server.log 2>&1 &
SPID=$!
READY=0
for i in $(seq 1 90); do
if curl -s localhost:$PORT/v1/models > /dev/null 2>&1; then READY=1; break; fi
if ! kill -0 $SPID 2>/dev/null; then echo "SERVER DIED"; break; fi
sleep 10
done
echo "READY=$READY after $((i*10))s"
if [ "$READY" != "1" ]; then tail -40 server.log; kill $SPID 2>/dev/null; sleep 20; kill -9 $SPID 2>/dev/null; exit 1; fi
grep -n "Using AttentionBackendEnum\|MLA prefill backend" server.log | head -3
# W1: long context, 32k in / 512 out, concurrency 32, 96 requests
vllm bench serve --model zai-org/GLM-5.3-Flash --port $PORT --dataset-name random \
--random-input-len 32768 --random-output-len 512 --num-prompts 96 --max-concurrency 32 \
--ignore-eos --save-result --result-dir "$DIR" --result-filename w1.json > w1.log 2>&1
echo "W1 exit $?"; tail -35 w1.log
# W2: mid context, 4k in / 1k out, concurrency 128, 384 requests
vllm bench serve --model zai-org/GLM-5.3-Flash --port $PORT --dataset-name random \
--random-input-len 4096 --random-output-len 1024 --num-prompts 384 --max-concurrency 128 \
--ignore-eos --save-result --result-dir "$DIR" --result-filename w2.json > w2.log 2>&1
echo "W2 exit $?"; tail -35 w2.log
kill $SPID 2>/dev/null; sleep 25
pgrep -u $(id -u) -f "port $PORT" && { pkill -9 -u $(id -u) -f "port $PORT"; sleep 10; }
pgrep -u $(id -u) -f "port $PORT" && echo "LEFTOVER" || echo "CLEAN" |
Contributor
|
do we want to change priority in this PR? if so, we should adjust it. |
JaredforReal
force-pushed
the
wire_sm90_kernels
branch
from
September 10, 2026 04:44
db053be to
5222ab9
Compare
4 tasks
- FlashMLA sparse: admit head_size 512 (bf16 only) and skip the concat_mla_q call when the rope part is empty - MLA prefill: whitelist the GLM5Next NoPE dims (256, 0, 256) - cmake: pin vllm-flash-attn and FlashMLA to the NoPE support commits Signed-off-by: JaredforReal <w13431838023@gmail.com> Co-authored-by: Kimi Code <noreply@moonshot.cn>
JaredforReal
force-pushed
the
wire_sm90_kernels
branch
from
September 13, 2026 06:23
5222ab9 to
8e0b383
Compare
|
@JaredforReal there's a follow-up on your fork branch: JaredforReal#22 Two things:
Once it's merged into wire_sm90_kernels this PR picks it up automatically and pytest tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py |
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.
Wire up two additional sparse-MLA attention backends for GLM5Next NoPE (head_size=512, e.g. GLM-5.3-Flash) on SM90:
FLASHMLA_SPARSEandFLASH_ATTN_MLA_SPARSE. FlashInfer SM90 remains the default; the new backends are opt-in via--attention-backend.Changes
flashmla_sparse.py: admit head_size 512 (bf16 KV only; quantized DS-MLA cache formats stay 576-only) and skip theconcat_mla_qcall when the rope part is empty.prefill/flash_attn.py: whitelist the GLM5Next NoPE dims(qk_nope=256, rope=0, v=256)for the dense-MHA prefill path (same kernel dims as DeepSeek's(192, 64, 256)).cmake/external_projects/vllm_flash_attn.cmake: pin vllm-flash-attn to9cd61de(Rope dim flash-attention#172, merged: FA3 OnlyQv forward for head_size==0). The FlashMLA pin is unchanged: the existing pin6bc4941already supports d_qk=512 sparse prefill/decode, and [Feature] add support NoPE head dim 512 for GLM5Next model on SM90 FlashMLA#17 (dense 512 decode) lands independently.No duplication of existing PRs: the backend wiring is new; the FA3 kernel support landed in vllm-project/flash-attention#172 and FlashMLA dense-512 in vllm-project/FlashMLA#17.
Benchmarks
Setup:
zai-org/GLM-5.3-Flash(bf16, ~306 GB), 4× H200 (SM90), TP=4, FlashInfer 0.6.18.vllm bench serve, random dataset,--ignore-eos,--max-model-len 36000, warmup with a different seed, two fresh-server runs per backend (cross-rep variance ±0.1%). Both workloads exceedindex_topk(2048), so prefill and decode both exercise each backend's sparse-MQA kernel.W1: 32k in / 512 out, concurrency 32 (96 requests)
W2: 4k in / 1k out, concurrency 128 (384 requests)
Independently reproduced by @simon-veitner-redhat with matching magnitudes. An earlier version of this table (showing +16~32%) was measured with the default 1M
--max-model-len, which inflates FlashInfer's host-side plan time — corrected above.Caveat: with
--no-enable-prefix-caching+ reduced--max-model-len, this hybrid model deterministically hits an IMA in the generic_compute_slot_mappings_kernel(vllm/v1/worker/gpu/block_table.py) under the FA backend only. Not related to the attention kernels; tracked as a separate core-infra issue.Tests
hopper/test_only_qv.pyin vllm-project/flash-attention (8/8).tests/v1/attention/test_sparse_mla_backends.pypasses.AI assistance was used for the implementation and benchmarking.