[zhangxin81]

Overview

Use native SM90 BF16 sparse FlashMLA for nonempty 8/16-head batches on full H20 (78 SMs) and H100 SXM 80GB (132 SMs). This reuses the ETAP kernel integrated by #67175; 32-head and other-device cutoffs remain.

Full-weight H20 evaluation/serving was not run because only one H20 was available. H20 coverage uses kernel tests and a two-layer random-weight LLM fixture.

Claims

• Avoid 64-head kernel work for measured small 8/16-head BF16 sparse MQA calls, retaining output strides and unused output rows.
• Cold-L2 changed full-width kernel shapes: H20 geometric mean 3.333x (100 shapes), H100 1.411x (120 shapes). These are kernel speedups.
• GLM-5.3-Flash H100 TP8 serving: 9216-token input/concurrency1 improves about 2.6% in paired output throughput. Other workload median changes are small and variable; no general model-throughput gain or output-equivalence claim is made.

Validation

Check	Original	Candidate
Extended regression, each H20/H100	16 failed, 128 passed	144 passed
Related attention suites, each GPU	baseline config-cache issue reproduced before attention	281 passed,657 existing skips
Kernel output/LSE comparisons	padded reference	444 shapes/device passed; existing tolerances
H20 two-layer random-weight BF16 DeepSeek-V3.2 LLM fixture	48 tokens	all match; 196608 full-vocabulary LP comparisons pass 0.01; final-repeat max delta 0.000379562
GLM TP8 GSM8K, 1319 questions / five-shot / max 512	1219 correct,92.4185%	1213 correct,91.9636%;30 correct-to-wrong,24 wrong-to-correct
GLM TP8 graph smoke	64 tokens	all match;common top-5 LP max delta 1.499683; generated-token LP max delta 0.168752
GLM TP4 eager smoke, real 16 heads	64 tokens;704 padded calls/rank	all tokens match;LP max delta 1.357752;704 ETAP calls/rank

GLM uses the official FP8 checkpoint revision 46a050bec13002ad29ee795c16e61cf4fa3c4a9c with BF16 KV. Every rank has 11 affected layers. VLLM_GPU_SYNC_CHECK=error remains enabled. Real-input FP32/padded/native probes pass 176 cases at existing output/LSE tolerances; actual/native replay outputs are identical. A same-weight eager native/padded/native-repeat intervention reproduces greedy/LP changes, while native-repeat is identical. The full checkpoint fails the toy fixture's 0.01 LP comparison; no tolerance was widened. GSM8K point estimate is lower by 0.4549 percentage points; paired-discordance p=0.496617 does not establish equivalence.

Serving uses TP8, batch budget 128, maxseq 128, maxlen 16384, BF16 KV, graphs, chunked prefill, no prefix cache, no FlashInfer autotune, memory 0.85, normal engine multiprocessing, Python frontend. Three paired rounds cover 24 workloads, 960 requests, 122880 output tokens; all counts/zero failures/fixed lengths/equal input counts pass. Temperature is not supplied, so the model/server default is used; seed 42 controls random input sampling.

Input / concurrency	Base output tokens/s median	Native median	Paired median ratio [range over 3 rounds]
128 / 1	174.07	175.42	1.00618 [1.00306, 1.00887]
2048 / 16	205.51	198.52	1.00200 [0.96598, 1.03001]
9216 / 1	45.34	46.56	1.02583 [1.02530, 1.15008]
9216 / 64	53.84	54.29	1.00014 [0.97340, 1.20807]

Concurrency 16/64 has large variation with gains and regressions. No server JIT-monitor warnings or preemptions were found. A reused client's logger reports closed-stream errors during input sampling before warmup/timing on 18 invocations; timing placement was audited, original logs retained, and a separate fix is supplied for future runs. A port preflight failure and an I/O-overlapping second round were retained/excluded explicitly; the latter was rerun after copying stopped.

The extended regression was installed in both source trees. On H20 and H100 the original selector fails the 16 newly expected small-batch native cases while actual output/LSE comparisons pass; the candidate passes all 144 cases. The failure identifies padded work through overwritten unused output heads, not output corruption on the baseline. Representative H100 output:

test_flashmla_bf16_preserves_unused_output_heads[False-True-512-8-single] FAILED
> assert torch.all(out[:, actual_num_heads:] == 123)
E AssertionError: assert tensor(False, device='cuda:0')
# original: 16 failed, 128 passed, 1927 deselected, 17 warnings in 159.16s
# candidate: 144 passed, 1927 deselected, 17 warnings in 150.15s

.venv/bin/python -m pytest tests/v1/attention/test_sparse_mla_backends.py \
  -k flashmla_bf16_preserves_unused_output_heads -v
.venv/bin/python -m pytest \
  tests/v1/attention/test_sparse_mla_backends.py \
  tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py \
  tests/kernels/attention/test_flashmla_sparse.py \
  -k 'flashmla and not bf16_preserves_unused_output_heads' -v
.venv/bin/python benchmarks/kernels/benchmark_flashmla_small_heads.py \
  --heads 8 16 32 --dims 512 576 --topk 128 640 2048 8192 8320 \
  --tokens 1 8 32 64 128 132 --rounds 3 --repeats 30 --output h100.csv

For H20 use tokens 1/8/32/64/78/128. Changed-file pre-commit (Ruff, mypy, Markdown, SPDX, test tethering and others) and git diff --check pass. Current-source CMake completed 470 steps; both arms use FlashMLA b18ed744 and identical binary SHA256 276f51f87eaa0ca1a852c43a6e3d2c0d27a1d85be92d250c2a0ba4f9a528a1f2. Large matrices use FlashInfer 0.7.2.post1; model/serving and the 24 exact-width 2176 supplemental shapes per GPU use 0.7.4. These matrices pass: H100 geometric mean 1.431x; H20 3.796x across all shapes and 3.785x over the 20 shapes with changed dispatch. The four H20 query-128 shapes already used native kernels.

Details

Device name/SM count is read at construction, excluding MIG partitions by SM count rather than physical name alone. Selection reads shape/device metadata and adds no GPU tensor host read. At 32 heads retain tokens>2*SM; other SM90 devices retain tokens>SM for 8/16 heads. BF16 native/padded reduction orders can change rounding and greedy output; FP8 KV-cache decode dispatch is unchanged.

H100 32-head T1,D576,K128 regresses from 8.000 us padded to 10.945 us native, supporting the retained cutoff. Empty-key 16-head H100 cases regress 0.048–0.096 us; no host synchronization is added to branch on key counts.

Duplicate searches found no equivalent common ETAP threshold PR. #57673 is a separate Triton DSv4.1 backend; #59418/#62941 target its decode, #67338 refactors routing metadata, #65403 changes index-group workspace/masked MHA, #59319/#60220 change FP8 mixed-batch workspace/staging. DeepSeek-V4.1 uses its dedicated attention class and already selects native BF16 prefill heads on SM90; GLM exercises this common implementation.

AI assistance was used to implement the change and run experiments. The human submitter confirmed line-by-line review and personally running relevant tests. Full H20 checkpoint/serving coverage remains a documented limitation. Fresh-agent documentation discovery passed locally and on the public PR branch. Rendered-preview discovery is pending: the upstream pre-run check stops lint and the docs build before execution because this author has 2 merged PRs rather than 4 and the PR lacks the required review labels. See the pre-run log and docs build. Local changed-file lint passed; these upstream jobs have not evaluated the code or rendered docs.

Full-model and serving reproduction

The measured base is 1786f260e3d68fa4d5a97191ddbffd7179edc96a. Both Python trees use the same source-built extensions; only the dispatch file differs. Download zai-org/GLM-5.3-Flash revision 46a050bec13002ad29ee795c16e61cf4fa3c4a9c and set MODEL_PATH to the local checkpoint. Model evaluation uses the repository's _build_gsm8k_prompts(1319, 5) and _score_gsm8k helpers, raw completions, and the following configuration. Run on each arm with VLLM_GPU_SYNC_CHECK=error, VLLM_ENABLE_V1_MULTIPROCESSING=0, and NCCL_NVLS_ENABLE=0:

import os
import time
from vllm import LLM, SamplingParams
from tests.evals.gsm8k.gsm8k_eval import _build_gsm8k_prompts, _score_gsm8k

llm = LLM(
    model=os.environ["MODEL_PATH"], dtype="bfloat16", kv_cache_dtype="bfloat16",
    tokenizer_mode="glm5_next", tensor_parallel_size=8, pipeline_parallel_size=1,
    distributed_executor_backend="mp", seed=42, gpu_memory_utilization=0.85,
    max_model_len=16384, max_num_batched_tokens=128, max_num_seqs=128,
    enable_prefix_caching=False, enable_chunked_prefill=True, enforce_eager=False,
    enable_flashinfer_autotune=False, limit_mm_per_prompt={"image": 0, "video": 0},
    attention_config={"backend": "FLASHMLA_SPARSE", "sparse_mla_force_mqa": False},
)
smoke = ["Question: What is 2 + 2?\nAnswer:",
         "Question: A shop has 12 apples and sells 5. How many remain?\nAnswer:"]
smoke_sampling = SamplingParams(temperature=0, seed=42, max_tokens=32,
                                ignore_eos=True, logprobs=5)
llm.generate(smoke, smoke_sampling, use_tqdm=False)  # warmup
llm.generate(smoke, smoke_sampling, use_tqdm=False)  # measured/profiled smoke
prompts, labels = _build_gsm8k_prompts(1319, 5)
sampling = SamplingParams(temperature=0, seed=42, max_tokens=512,
                         stop=["Question", "Assistant:", "<|separator|>"])
start = time.perf_counter()
outputs = llm.generate(prompts, sampling)
metrics = _score_gsm8k(
    [o.outputs[0].text for o in outputs],
    [len(o.outputs[0].token_ids) for o in outputs], labels, 5, 512,
    time.perf_counter() - start,
)
print(metrics)

For the TP4 coverage check, use TP4/PP1, eager execution, kv_cache_memory_bytes=512*1024**2, memory utilization 0.965, max length 512, batch-token budget 64 and max sequences 1. Run the same smoke prompts only. This is a coverage check, not a TP4 serving benchmark.

Serving uses normal engine multiprocessing and Python frontend/benchmark (VLLM_ENABLE_V1_MULTIPROCESSING=1, VLLM_USE_RUST_FRONTEND=0, VLLM_USE_RUST_BENCH=0), with VLLM_GPU_SYNC_CHECK=error and NCCL_NVLS_ENABLE=0:

.venv/bin/python -m vllm.entrypoints.cli.main serve "$MODEL_PATH" \
  --host 127.0.0.1 --port 48123 --tokenizer-mode glm5_next \
  --tensor-parallel-size 8 --distributed-executor-backend mp --seed 42 \
  --dtype bfloat16 --kv-cache-dtype bfloat16 --gpu-memory-utilization 0.85 \
  --max-model-len 16384 --max-num-batched-tokens 128 --max-num-seqs 128 \
  --no-enable-prefix-caching --enable-chunked-prefill \
  --no-enable-flashinfer-autotune --attention-backend FLASHMLA_SPARSE \
  --limit-mm-per-prompt '{"image":0,"video":0}'

The measured client imports vllm.entrypoints.cli.main.main once per server arm and invokes the CLI for each (input length, concurrency, requests) tuple: (128,1,16), (2048,16,64), (9216,1,16), (9216,64,64). The CLI resets seeds for each invocation. Equivalent per-workload arguments are:

.venv/bin/python -m vllm.entrypoints.cli.main bench serve \
  --backend vllm --base-url http://127.0.0.1:48123 --model "$MODEL_PATH" \
  --tokenizer-mode glm5_next --dataset-name random \
  --random-input-len "$INPUT_LENGTH" --random-output-len 128 \
  --random-range-ratio 0.0 --num-prompts "$REQUESTS" --request-rate inf \
  --max-concurrency "$CONCURRENCY" --ignore-eos --seed 42 --num-warmups 4 \
  --save-result --result-dir results --result-filename "$RESULT_FILE" \
  --percentile-metrics ttft,tpot,itl,e2el --disable-tqdm

Run three paired rounds in base/candidate, candidate/base, base/candidate order. Round 0 uses shared source/runtime; rounds 1/2 use byte-identical local Python caches, balanced within each pair, with original binary origins retained. Startup/import/warmup time is outside measurement. An interrupted round overlapping cache-copy I/O was archived and excluded; both arms were rerun after copying stopped. Original result/log files, numerical probes and failed setup logs are retained locally rather than added to the source diff.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[mergify]

Documentation preview: https://vllm--67468.org.readthedocs.build/en/67468/

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @zhangxin81.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork