[Feature] TRITON_MLA_SPARSE backend for SM8x/11x/12x DSA Sparse MLA Support

[haosdent]

Purpose

Closes #38006. Enables sparse MLA models (GLM-5, DeepSeek-V3.2) on SM80 (A100/A800) and SM121 (GB10/DGX Spark), where DeepGEMM / FlashMLA-Sparse / FlashInfer-MLA-Sparse are unavailable.

Changes

1. Dispatch guard. is_deep_gemm_supported() (SM90+ check) replaces has_deep_gemm() in sparse_attn_indexer.py / indexer.py. Stops DeepGEMM kernels from being invoked on SM80/SM121.
2. Triton fp8_mqa_logits for the indexer. mqa_logits_triton.py reproduces DeepGEMM's prefill + paged MQA logits. Prefill takes bf16 q/k (pre-decoded from FP8 in the Python wrapper) and feeds a straight tl.dot; paged decode keeps a 256-entry bf16 LUT for in-kernel FP8 decode. K-side scale applied to the fp32 dot output, per-row K-tile early-exit on the chunked-prefill path.
3. TRITON_MLA_SPARSE backend. triton_mla_sparse_kernel.py adds a split-KV decode with N-way online-softmax merge plus a single-pass fast path. Autotune is warmed at init using indexer-derived (n_head, head_dim). Masked-out sentinel is -1e30 to avoid NaN from (-inf) − (-inf) on all-masked tiles.
4. Cudagraph support. TritonMLASparseMetadataBuilder advertises AttentionCGSupport.UNIFORM_BATCH; flips A100 TP=8 back to FULL_AND_PIECEWISE.
5. MXFP4 link stubs. mxfp4_experts_quant / silu_and_mul_mxfp4_experts_quant stubs in nvfp4_quant_entry.cu. Real impls are SM10.x-only in CMake but torch_bindings.cpp references them unconditionally, which breaks source builds on SM 8.x.

Benchmarks

8×A100 SXM TP=8, lukealonso/GLM-5.1-NVFP4, single prompt, decode 200 tokens. cold = first request on a fresh prompt; warm = repeat (prefix cache hit):

context	prefix	TTFT	TPOT median	output tok/s
short (1,744 in)	cold	0.72 s	21.7 ms	39.7
short (1,744 in)	warm	0.18 s	21.8 ms	44.3
mid (65,536 in)	cold	37.0 s	24.3 ms	4.8
long (127,744 in)	cold	98.1 s	25.2 ms	1.9
long (127,744 in)	warm	0.60 s	25.3 ms	35.5

Tests

• tests/kernels/attention/test_mqa_logits_triton.py — 41 cases (DeepGEMM reference + clean/dirty clean_logits + 256-byte FP8 decode).
• tests/kernels/attention/test_triton_mla_sparse_kernel.py — 53 cases (split vs single-pass + auto-heuristic + short-prefill no-NaN).

Limitations (follow-up): BF16 KV cache only on SM80/SM121; VLLM_BATCH_INVARIANT should force num_kv_splits=1 — not wired.

[mergify]

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

[gemini-code-assist]

Code Review

This pull request introduces the TRITON_MLA_SPARSE attention backend, providing a Triton-based fallback for sparse MLA on GPUs like NVIDIA Ampere. It also refactors FP8 MQA logit fallbacks into a new module and updates the sparse attention indexer to use these PyTorch implementations when DeepGEMM is unsupported. A review comment suggests moving a module-level import to the top of the file to comply with PEP 8 guidelines.

[gemini-code-assist]

To adhere to PEP 8 guidelines, module-level imports should be placed at the top of the file. Please move this import statement to the top of the module, for example, after the from vllm.platforms import current_platform import. This improves code readability and consistency.

References

1. PEP 8: E402 module level import not at top of file. Imports should be at the top of the file, just after any module comments and docstrings, and before module globals and constants. ^(link)

[ZJY0516]

Add PyTorch fallback for indexer MQA logits — Created a shared mqa_logits_fallback.py module with fp8_mqa_logits_torch and fp8_paged_mqa_logits_torch implementations (extracted from existing ROCm fallback code). The sparse attention indexer now dispatches to these fallbacks when DeepGEMM is not supported.

FYI, we don’t plan to support a torch native mqa_logits implementation. I also question whether it’s necessary to support sparse MLA on SM80.

[ehfd]

Add PyTorch fallback for indexer MQA logits — Created a shared mqa_logits_fallback.py module with fp8_mqa_logits_torch and fp8_paged_mqa_logits_torch implementations (extracted from existing ROCm fallback code). The sparse attention indexer now dispatches to these fallbacks when DeepGEMM is not supported.

FYI, we don’t plan to support a torch native mqa_logits implementation. I also question whether it’s necessary to support sparse MLA on SM80.

@ZJY0516

Update: This PR is no longer a torch-native mqa_logits implementation and instead use Triton.

Ampere is still extremely commonly used... We need this for DS3.2 or GLM-5. TRITON_MLA_SPARSE is possible, so we need an option for that.

[ehfd]

PyTorch fallback for indexer logits is slower than DeepGEMM — a Triton-based variant can follow

Perhaps we can integrate Triton now?

[workcode-del]

Hello, after I modified the code according to your PR, the GLM5 model service started normally. However, the response speed is very slow, with only about 3 tokens being responded to per second. My device is also an A800 with 80G of storage capacity. Is this normal?
@haosdent

[ehfd]

@workcode-del I believe that it's only expected to be anywhere remotely fast when no PyTorch fallbacks exist.

[workcode-del]

@workcode-del I believe that it's only expected to be anywhere remotely fast when no PyTorch fallbacks exist.

Could you please explain how to achieve the condition where there are no PyTorch fallbacks?

[mergify]

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

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

[ehfd]

@workcode-del
Further development is necessary.

[ehfd]

#37968, #35271, #38076, #36519

Probably containing the reason why needs-rebase was added.

[ianlevesque]

I was able to use this patch to run GLM-5.1 on an 8-node DGX Spark cluster. Performance is obviously not stellar (~5 t/s) but it's a great first step with compatibility.

[haosdent]

Thanks @ianlevesque , your 8 x DGX Spark is incredible! I just add new triton kernels to try to address the performance issue, may you test again when you are available?

[ianlevesque]

@haosdent retried with the new patch, it did improve to 10 t/s or so.

  Depth=0 (fresh context)

  ┌────────┬────────────────────────┬───────────────┐
  │  Test  │    Throughput (t/s)    │   TTFT (ms)   │
  ├────────┼────────────────────────┼───────────────┤
  │ pp512  │ 425–452                │ 1,059–1,091   │
  ├────────┼────────────────────────┼───────────────┤
  │ pp2048 │ 606–630                │ 2,892–2,971   │
  ├────────┼────────────────────────┼───────────────┤
  │ pp8192 │ 723–730                │ 10,041–10,090 │
  ├────────┼────────────────────────┼───────────────┤
  │ tg32   │ 9.91–10.29 (peak 11.0) │ —             │
  ├────────┼────────────────────────┼───────────────┤
  │ tg128  │ 9.91–10.29 (peak 11.0) │ —             │
  ├────────┼────────────────────────┼───────────────┤
  │ tg512  │ 9.85–10.23 (peak 11.0) │ —             │
  └────────┴────────────────────────┴───────────────┘

  Depth=2048 (warm context)

  ┌────────┬──────────────────┬───────────────┐
  │  Test  │ Throughput (t/s) │   TTFT (ms)   │
  ├────────┼──────────────────┼───────────────┤
  │ pp512  │ 667–697          │ 3,250–3,403   │
  ├────────┼──────────────────┼───────────────┤
  │ pp2048 │ 656–689          │ 5,282–5,618   │
  ├────────┼──────────────────┼───────────────┤
  │ pp8192 │ 717–731          │ 12,373–12,800 │
  ├────────┼──────────────────┼───────────────┤
  │ tg32   │ 9.18–9.77        │ —             │
  ├────────┼──────────────────┼───────────────┤
  │ tg128  │ 9.49–9.79        │ —             │
  ├────────┼──────────────────┼───────────────┤
  │ tg512  │ 9.51–9.70        │ —             │
  └────────┴──────────────────┴───────────────┘

[Ph0enix89]

I am able to run it on 40 GB A100 GPUs. Using the latest main with the PR on top of it. Details below. There are two main issues:

1. Tool calling doesn't work. Queries via curl work fine. However whenever I try to use it in Hermes eventually there's a timeout. In the logs there is Process group watchdog thread terminated with exception: CUDA error: an illegal memory access was encountered. Has anyone run into it? Any recommendations? I tried launching with --chat-template-content-format string but that doesn't help. There are some reports that tool calling works without pipeline parallelism but these come from AMD ROCM systems so it's not clear if it's relevant.
2. KV cache quantization is not available. It's been mentioned that the problem is that fp8e4nv is hardcoded and not supported. AI suggested that the code should be adjusted to use 'fp8e5 for SM80. However it's not clear how much effort it would take to implement. In the current main there is also commit b588f66dc ("[GLM5.2 Perf] fused_indexer_q_rope_quant triton kernel") that also uses hardcoded fp8e4nv for some optimizations. It has to be reverted, otherwise vllm fails during startup on A100 (when trying to load GLM-5.2). Since the A100s I have access to have only 40 GBs doubling KV cache would be very handy, that's probably the biggest throughput bottleneck.

As for the steps I had to combine a bunch of scattered pieces to make it work on my system:

1. Checkout main
2. Revert b588f66dc ("[GLM5.2 Perf] fused_indexer_q_rope_quant triton kernel")
3. Follow steps in this guide to modify the code
4. Add this patch
5. Install the patched version with VLLM_USE_PRECOMPILED=1 uv pip install -e .
6. Install runai-streamer with pip3 install vllm[runai]. Not strictly necessary but without it loading the model takes forever.
7. Finally launch it with:

vllm serve "$MODEL" \
        --served-model-name glm-5.2 \
        --kv-cache-dtype auto \
        --tensor-parallel-size ${TP_SIZE} \
        --pipeline-parallel-size ${PP_SIZE} \
        --distributed-executor-backend mp \
        --nnodes ${NNODES} \
        --node-rank "${PROCID}" \
        --master-addr "$HEAD_HOST_IP" \
        --no-async-scheduling \
        --tool-call-parser glm47 \
        --enable-auto-tool-choice \
        --reasoning-parser glm45 \
        --gpu-memory-utilization 0.93 \
        --load-format runai_streamer \
        --compilation-config.cudagraph_mode=PIECEWISE \
        --disable-custom-all-reduce \
        --chat-template-content-format string \
        --block-size 128 \
        ${HEADLESS_FLAG}
7) It took a while to figure out that `--block-size 128` is necessary. Without it vllm fails to start with `RuntimeError: Worker failed with error 'No common block size for 16. ', please check the stack trace above for the root cause`.

These are my steps that result in a working setup without tool calling support. Perhaps it helps some people. Any tips are welcome.

[ghostplant]

I am able to run it on 40 GB A100 GPUs. Using the latest main with the PR on top of it. Details below. There are two main issues:

1. Tool calling doesn't work. Queries via curl work fine. However whenever I try to use it in Hermes eventually there's a timeout. In the logs there is Process group watchdog thread terminated with exception: CUDA error: an illegal memory access was encountered. Has anyone run into it? Any recommendations? I tried launching with --chat-template-content-format string but that doesn't help. There are some reports that tool calling works without pipeline parallelism but these come from AMD ROCM systems so it's not clear if it's relevant.
2. KV cache quantization is not available. It's been mentioned that the problem is that fp8e4nv is hardcoded and not supported. AI suggested that the code should be adjusted to use 'fp8e5 for SM80. However it's not clear how much effort it would take to implement. In the current main there is also commit b588f66dc ("[GLM5.2 Perf] fused_indexer_q_rope_quant triton kernel") that also uses hardcoded fp8e4nv for some optimizations. It has to be reverted, otherwise vllm fails during startup on A100 (when trying to load GLM-5.2). Since the A100s I have access to have only 40 GBs doubling KV cache would be very handy, that's probably the biggest throughput bottleneck.

As for the steps I had to combine a bunch of scattered pieces to make it work on my system:

1. Checkout main
2. Revert b588f66dc ("[GLM5.2 Perf] fused_indexer_q_rope_quant triton kernel")
3. Follow steps in this guide to modify the code
4. Add this patch
5. Install the patched version with VLLM_USE_PRECOMPILED=1 uv pip install -e .
6. Install runai-streamer with pip3 install vllm[runai]. Not strictly necessary but without it loading the model takes forever.
7. Finally launch it with:

vllm serve "$MODEL" \
        --served-model-name glm-5.2 \
        --kv-cache-dtype auto \
        --tensor-parallel-size ${TP_SIZE} \
        --pipeline-parallel-size ${PP_SIZE} \
        --distributed-executor-backend mp \
        --nnodes ${NNODES} \
        --node-rank "${PROCID}" \
        --master-addr "$HEAD_HOST_IP" \
        --no-async-scheduling \
        --tool-call-parser glm47 \
        --enable-auto-tool-choice \
        --reasoning-parser glm45 \
        --gpu-memory-utilization 0.93 \
        --load-format runai_streamer \
        --compilation-config.cudagraph_mode=PIECEWISE \
        --disable-custom-all-reduce \
        --chat-template-content-format string \
        --block-size 128 \
        ${HEADLESS_FLAG}
7) It took a while to figure out that `--block-size 128` is necessary. Without it vllm fails to start with `RuntimeError: Worker failed with error 'No common block size for 16. ', please check the stack trace above for the root cause`.

These are my steps that result in a working setup without tool calling support. Perhaps it helps some people. Any tips are welcome.

Do you run this with 16 A100s?

[Ph0enix89]

Do you run this with 16 A100s?

64. 16 nodes with 4 GPUs each. Not the most efficient config but it is what it is.

[ehfd]

#43477 was merged and a maintainer said on Slack that they were open to merging this one after that PR.

Integration of GLM-5.2 and DeepSeek V4 support, as well as a rebase is desired.

@haosdent

[lzf-tech]

#43477 was merged and a maintainer said on Slack that they were open to merging this one after that PR.

Integration of GLM-5.2 and DeepSeek V4 support, as well as a rebase is desired.

@haosdent

It seems @haosdent hasn’t been active on this project for some time. Really looking forward to the GLM-5.2 and DeepSeek V4 support once the rebase is done.

[halexan]

Any progress?

[ehfd]

An option is for somebody to create a new PR on top of this one (crediting the original author with Co-Authored-By:) and make available temporary images to use.

If you can be responsive, we can get it merged, even.

[thomaslwang]

Following up on the suggestion above (@ehfd) — I've opened #47629, a rebase & takeover of this PR onto current main, with @haosdent's original commit and authorship preserved (thank you for the great work — the fix-38006-6 image has been serving GLM-5.1 on my box for weeks).

What's in it beyond the rebase:

• The indexer dispatch is now a three-way XPU → DeepGEMM → Triton fallback chain, preserving main's newer XPU path, skip_topk_buffer_clear, and DCP handling.
• Backend priority per the maintainer request: TRITON_MLA_SPARSE is appended after FLASH_ATTN_MLA_SPARSE/FLASHMLA_SPARSE, so SM90+ keeps native sparse backends and only SM80/SM121 falls through to Triton.
• New fix: use_fused_indexer_q ([GLM5.2 Perf] fused_indexer_q_rope_quant triton kernel, 1.9% ~ 3.3% E2E Throughput improvement. #46862) is gated on SM89+ — its Triton kernel stores fp8e4nv, which doesn't compile on SM80. This is the A100 startup crash @Ph0enix89 hit; the unfused per_token_group_quant_fp8 path is used on older archs instead.

Validation on 8×A800-80G PCIe (SM80), GLM-5.2 NVFP4, TP=8, bf16 KV cache:

• tests/kernels/attention/test_mqa_logits_triton.py: 41 passed; test_triton_mla_sparse_kernel.py: 53 passed
• End-to-end serve on current main boots clean with the default cudagraph mode (FULL_AND_PIECEWISE) — no capture crash and no --block-size 128 workaround needed anymore; the backend's built-in autotune warmup covers it
• Tool calling works via the glm47 parser (single node, TP-only)
• Decode throughput: ~32 tok/s single-stream, ~100 tok/s aggregate at 4-way concurrency

One note for anyone reproducing the earlier setups in this thread: the IndexShare patch from haosdent#7 is no longer needed — current main's index_topk_freq/index_skip_topk_offset logic already covers GLM-5.2's shared-indexer layout.

I'll push a prebuilt image for immediate use and link it here shortly. Happy to iterate quickly on review feedback — the 8×A800 box is available for validation runs.

[ehfd]

@thomaslwang I greatly appreciate the work!!!

However, a few questions.

(1) I would like to see pipeline parallelism working with tool/reasoning parsers, because pipeline parallelism allows 1M context on 8x A100 (GLM-5.x is an MLA model; so KV cache is duplicated on all GPUs with tensor parallelism, which pipeline parallelism solves). Were you able to confirm this?

(2) Note that SM12x now works through #43477 on fp8 (but not bfloat16) KV cache. This means that FLASHINFER_MLA_SPARSE_SM120 should come first in priority, then this one, since this one can provide bfloat16 KV cache.

[thomaslwang]

As promised — a prebuilt image for immediate use on SM80 while #47629 goes through review:

docker pull openguardrails/vllm-glm52-sm80:435f82d61-pr38476

It follows @RefalMachine's recipe from this thread: vllm/vllm-openai:cu129-nightly-9037498c2 (June 22 main, commit 435f82d) + this PR cherry-picked, conflicts resolved to the three-way XPU → DeepGEMM → Triton dispatch. Entry point is the OpenAI API server; example (8×A100/A800-80G, GLM-5.2 NVFP4):

docker run -d --gpus all --ipc host --shm-size 32g -p 8000:8000 \
  -v /path/to/models:/models:ro \
  openguardrails/vllm-glm52-sm80:435f82d61-pr38476 \
  --model /models/GLM-5.2-NVFP4 \
  --tensor-parallel-size 8 --enable-expert-parallel \
  --trust-remote-code --dtype bfloat16 --kv-cache-dtype bfloat16 \
  --max-model-len 131072 --gpu-memory-utilization 0.95 --max-num-seqs 4 \
  --no-async-scheduling --tool-call-parser glm47 --reasoning-parser glm45 \
  --enable-auto-tool-choice --chat-template-content-format=string

Verified on 8×A800-80G PCIe: TRITON_MLA_SPARSE selected, 202k-token KV cache at 131k max-model-len, tool calling OK, ~33 tok/s single-stream / ~107 tok/s at 4-way concurrency.

For latest main, use #47629 instead (VLLM_USE_PRECOMPILED=1 install works — the diff is Python-only on SM80).

[thomaslwang]

@ehfd Tested PP as requested. Everything below is from 8×A800-80G PCIe (SM80), GLM-5.2 NVFP4, bf16 KV cache, branch #47629.

(1) PP + tool/reasoning parsers: works. Single node, TP=2 × PP=4, --tool-call-parser glm47 --reasoning-parser glm45: tool calls come out well-formed (finish_reason=tool_calls), including after long contexts — needle retrieval at 79k prompt tokens is accurate and two concurrent 53k-token requests both retrieved their needles and emitted correct tool calls.

However, I can reproduce the crash @Ph0enix89 reported, and it is neither tool-parser- nor multi-node-related: single-node TP=2×PP=4 with concurrent long prefills (~50–90k tokens) nondeterministically hits Triton Error [CUDA]: an illegal memory access was encountered, surfacing at the indexer's build_prefill_chunk_metadata kernel launch. It behaves like a race: with CUDA_LAUNCH_BLOCKING=1 the identical workload is fully stable (which also matches the "eventually, under load" pattern), and TP-only (TP=8) is unaffected. Note the prefill-chunk metadata path is shared with the DeepGEMM route, so the race may be latent in the indexer metadata pipeline generally and merely easier to hit with the slower Triton kernels. Tracking as a known issue on #47629 while I bisect it. Interim workarounds: TP-only, or PP + CUDA_LAUNCH_BLOCKING=1 (~15% slower prefill in my quick check: 32s vs 28s for a 79k prompt).

Capacity with bf16 KV (--block-size 128):

Config	max-model-len that fits
TP=2 × PP=4 @ 0.96 util	524,288 boots; vLLM estimates ceiling ≈ 605k
TP=1 × PP=8 @ 0.95 util	estimated ceiling 990,848 — 0.5% short of 1M

The final blocker for a true 1M is not KV: it's the decode-path logits buffer ([batch, max_seq_len] fp32 ≈ 8 GiB at 1M). FP8 KV on SM80 (needs an fp8e5-style path, as discussed earlier in this thread) or chunking that buffer would clear 1M; both feel like follow-ups rather than blockers for this PR.

Also addressed on the branch since your comment:

• (2) SM12x priority now lists FLASHINFER_MLA_SPARSE_SM120 first (FP8 KV) with TRITON_MLA_SPARSE appended as the BF16-KV fallback — thanks, the rebase had dropped SM12x coverage entirely.
• The No common block size for 16. startup failure under PP is root-caused and fixed: the backend inherited MultipleOf(1) and auto-selection settled on 16 while the indexer requires 64; it now declares MultipleOf(64) (auto works, --block-size 128 still honored — 128 also measurably lowers profile-time peak memory at long contexts).

[thomaslwang]

Root cause found for the PP illegal-memory-access crash (@Ph0enix89) — fix up at #47644.

It's a pinned-buffer reuse race in the core model runner, not in the sparse-MLA kernels: prepare_inputs_event (which guards reuse of the pinned CPU input tensors between steps) was only created when async scheduling is enabled. The PP batch queue creates the same overlap — input prep for step N+1 rewrites seq_lens_cpu / query_start_loc_cpu / block-table staging while step N's non_blocking H2D copies from those buffers may still be pending — so step N's kernels can consume step N+1's values. The DSA indexer sizes its prefill-chunk buffers from step N's CPU values but indexes with the device values → OOB → IMA. That's why:

• it only bites with PP + --no-async-scheduling (exactly what all the recipes in this thread use — async scheduling creates the guard event as a side effect),
• TP-only is immune (max_concurrent_batches == 1),
• CUDA_LAUNCH_BLOCKING=1 makes it vanish, and
• it's "eventually, under load" — the window scales with how far the GPU lags the CPU, which is huge with SM80 Triton prefill.

Validation on 8×A800 (TP=2×PP=4, GLM-5.2 NVFP4): the 9-concurrent-long-prefill stress that reliably crashed within ~2 minutes now passes 27/27 across 3 rounds with correct outputs.

@RefalMachine your NaN → !!! degradation under sustained load may be the same race in its silent form (corrupted inputs without OOB) — if you can, try your workload with the #47644 patch; it's a ~5-line cherry-pick.

[ehfd]

I think it's possible that #38476 (comment) is related to #41623 or #42426. Great!

[ehfd]

decode-path logits buffer
chunking that buffer would clear 1M

@thomaslwang Is this a large task? I don't think we need to use FP8 KV cache, but if this is not a big work, it could be integrated (since a separate PR can also take ages to get reviewed).

[thomaslwang]

@ehfd On the two issues:

• [Bug]: Kimi-K2.6 intermittently outputs only "!!!!!!!!!!" in reasoning field with content null #42426 (Kimi-K2.6 !!!!!! + NaN logprobs) has exactly the same signature as @RefalMachine's report — intermittent, load/uptime-correlated, NaNs in the response path, restart clears it. If that deployment runs PP without async scheduling it's very plausibly the same pinned-buffer race; I've asked on the issue for their parallel config and pointed them at [Bugfix][Core] Sync reused pinned input buffers under PP batch queue (fixes IMA with sparse MLA + PP) #47644.
• [Bug]: Decode Context Parallelism (--decode-context-parallel-size) output drift and gibberish in v0.21.0 and latest nightly #41623 (DCP output drift) goes through a different mechanism (DCP doesn't use the batch queue), so I wouldn't claim it — though silent input corruption is exactly the kind of thing that produces drift, so worth re-testing once [Bugfix][Core] Sync reused pinned input buffers under PP batch queue (fixes IMA with sparse MLA + PP) #47644 lands.

On the decode logits buffer: agreed, no need for FP8 KV. Chunking is tractable — the global top-k is exactly recoverable from per-chunk top-k candidates (any token in the global top-2048 is in its chunk's top-2048), so the plan is: window the KV in ~128k-token chunks, run the paged logits kernel per window, per-chunk topk → (score, global-id) candidates, then one final top-k over the concatenated candidates. Bounded fp32 footprint regardless of context length, and it's what the memory profiler sees, so 1M should boot at TP=1×PP=8 with bf16 KV. I'll add it to #47629 as a threshold-gated path (single-shot behavior unchanged below the threshold) and validate on the 8×A800 — will report back here with 1M boot + long-context correctness results.

[thomaslwang]

@ehfd Update on the 1M question — with a correction first, then good news.

Correction: I earlier attributed the 1M shortfall to the indexer's decode logits buffer. That was wrong — at realistic decode concurrency that buffer is only megabytes (batch × max_seq_len × fp32 with batch ≈ 1–2 at these context lengths). The per-stage memory logs show the real constraint: the last PP stage, which carries the LM head plus the sampler's vocab-sized activations (~150k vocab). No chunking work is needed at all.

The zero-code fix: give the last stage fewer layers with the existing VLLM_PP_LAYER_PARTITION env. Measured on 8×A800-80G PCIe, GLM-5.2 NVFP4, bf16 KV cache, branch #47629 + the #47644 race fix:

VLLM_PP_LAYER_PARTITION="11,10,10,10,10,10,10,7" \
vllm serve nvidia/GLM-5.2-NVFP4 \
  --tensor-parallel-size 1 --pipeline-parallel-size 8 \
  --max-model-len 1048576 --block-size 128 \
  --gpu-memory-utilization 0.96 \
  --dtype bfloat16 --kv-cache-dtype bfloat16 \
  --no-async-scheduling --max-num-seqs 2 \
  --tool-call-parser glm47 --reasoning-parser glm45 --enable-auto-tool-choice

Results:

• Boots at max-model-len 1,048,576 — GPU KV cache size: 1,062,656 tokens, 1.01× concurrency at full 1M
• Needle retrieval at 250,857 prompt tokens (needle buried at ~200k depth): exact answer, 108 s end-to-end prefill+decode
• One caveat that doubles as evidence for [Bugfix][Core] Sync reused pinned input buffers under PP batch queue (fixes IMA with sparse MLA + PP) #47644: on the same config without the race fix, even a single 250k-token request crashes with the illegal-memory-access race (chunked prefill of one long request is enough to overlap steps in the PP batch queue). With the fix applied it's clean.

So: 1M context on 8×A100/A800 with bf16 KV is real once #47629 + #47644 land — no FP8 KV cache and no new kernels required. I'll fold the partition recipe into the #47629 description.

[ehfd]

@thomaslwang Spectacular. Thank you!

I didn't know that there was a trick like this.

[Ph0enix89]

@thomaslwang Thank you for all the work. Not sure if we should continue here or under the new PR.

I took both of your new branches. Initial results looked promising. Tooling worked, things seemed stable. However later weird behavior emerged. On one instance one session kept working fine while starting any new session against the same vllm instance produced gibberish.

Then on the other instance I had a session that got stuck without producing any output.

In the logs there is nothing out of the ordinary, no errors, occasional warnings. And it claims to be producing tokens but the agent doesn't get anything from it.

My command matches yours for the most part. I have TP=4 and PP=16.

Also what is the purpose of setting VLLM_PP_LAYER_PARTITION? And what's the risk of not setting it?

[ehfd]

@Ph0enix89 Did you apply #47644 as well?

Also what is the purpose of setting VLLM_PP_LAYER_PARTITION? And what's the risk of not setting it?

This is just for increasing KV cache capacity.

[ehfd]

@thomaslwang If there was a container which serves the specific #47629 + #47644, combination, I think more feedback can come in.

[RefalMachine]

Root cause found for the PP illegal-memory-access crash (@Ph0enix89) — fix up at #47644.

It's a pinned-buffer reuse race in the core model runner, not in the sparse-MLA kernels: prepare_inputs_event (which guards reuse of the pinned CPU input tensors between steps) was only created when async scheduling is enabled. The PP batch queue creates the same overlap — input prep for step N+1 rewrites seq_lens_cpu / query_start_loc_cpu / block-table staging while step N's non_blocking H2D copies from those buffers may still be pending — so step N's kernels can consume step N+1's values. The DSA indexer sizes its prefill-chunk buffers from step N's CPU values but indexes with the device values → OOB → IMA. That's why:

• it only bites with PP + --no-async-scheduling (exactly what all the recipes in this thread use — async scheduling creates the guard event as a side effect),
• TP-only is immune (max_concurrent_batches == 1),
• CUDA_LAUNCH_BLOCKING=1 makes it vanish, and
• it's "eventually, under load" — the window scales with how far the GPU lags the CPU, which is huge with SM80 Triton prefill.

Validation on 8×A800 (TP=2×PP=4, GLM-5.2 NVFP4): the 9-concurrent-long-prefill stress that reliably crashed within ~2 minutes now passes 27/27 across 3 rounds with correct outputs.

@RefalMachine your NaN → !!! degradation under sustained load may be the same race in its silent form (corrupted inputs without OOB) — if you can, try your workload with the #47644 patch; it's a ~5-line cherry-pick.

Hi! I'm testing a model with PP=1 and TP=8 + --no-async-scheduling. Should this fix have any impact under these conditions?

By the way, thanks for the ready-to-use Docker container, but unfortunately, I'm also having the same issue with !!! tokens.

[Kasempiternal]

Following up on the IndexShare note above — that GLM-5.2 sm80 bring-up is now a multi-day production deployment: BF16, 32×A100, TP4×PP8, sustained ~130–155 tok/s decode at ~24 concurrent requests. So this backend has real production mileage on Ampere now, for the "is sm80 sparse-MLA worth supporting?" question — for us it's the only way to run GLM-5.2 on A100 at all.

One pointer for the #47629 takeover: I've posted a production illegal-memory-access report on #47644 (the PP pinned-buffer race) — 13/381 serve instances, decode-time under concurrency, on this exact BF16 TP4×PP8 sm80 config — plus an offer to A/B the fix across our fleet. Since the end-to-end GLM-5.2 validation here is NVFP4/TP8, the BF16 + PP8 + sm80 path we run is complementary coverage; happy to be a test site for it.

[lzf-tech]

Great work on this PR! I've tested it on my 8×A800 80G PCIE cluster and share my runtime observations from the log screenshot attached:

1. Prefix Cache Hit Rate steadily stays at ~93.8%–93.9% under continuous chat completions traffic, the prefix cache optimization works extremely well for repeated prompt scenarios.
2. Performance metrics:
   • Avg prompt throughput peaks at 115.8 tokens/s, stable range 17–52 tokens/s under single running request.
   • Avg generation throughput maintains 16–34 tokens/s most of the time, only drops temporarily to ~7–11 tokens/s at request drain.
3. Runtime GPU KV cache usage hovers between 23%–26.5% with 1 concurrent running request, no waiting requests during the whole test window, stable memory footprint without leak.
4. API stability: All core /v1/chat/completions endpoints return 200 OK consistently; only irrelevant custom test GET routes hit 404 which has no impact on LLM serving.

No crash, hanging or OOM issues observed during long-running continuous serving test. Thank you for this valuable optimization for prefix caching and engine throughput!

[4blacktea]

Thank you so much for this exciting work! I would like to know if this PR supports DeepSeek-V4-Flash on 8×A100.

[ehfd]

Thank you so much for this exciting work! I would like to know if this PR supports DeepSeek-V4-Flash on 8×A100.

Not yet. We want stable support first on GLM-5.x and DeepSeek-V3.2 architectures.