GLM-5.2-FP8 with DSA enablement

[jkaniecki]

Enable GLM-5.2 with DSA (DeepSeek Sparse Attention) on HPU

What changed:

hpu_sparse_attn_indexer.py (new): per-request Q·K BF16 scoring + torch.topk to select top-2048 KV cache slots per decode step
hpu_attn.py: forward_mqa_sparse — gathers top-K latent KV, decompresses, runs masked MLA decode attention
oot_mla.py: dispatches to forward_mqa_sparse when use_sparse=True
deepseek_v2.py: BF16 indexer cache, HPU-safe Indexer.forward, SparseAttnIndexer dispatch
platform.py / hpu_model_runner.py: DSA routing + indexer cache layer index fix
Validated on 8× Gaudi3 (GLM-5.2-FP8, TP=8):

GSM8K 5-shot (1319 samples): 93.6% exact match
AIME 2026 - 100% - on pair with model card's target (99,2 % mean for multiple reruns)

[Copilot]

🟡 Changes recommended

The sparse indexer may return partially uninitialized top-k index buffers (risking out-of-range cache gathers) and the sparse decode softmax can produce NaNs for fully-masked rows.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Enables GLM-5.2 FP8 execution with DSA (Dynamic Sparse Attention) on Intel Gaudi HPU by wiring an HPU BF16-based sparse-attention indexer into the DeepSeek/GLM MLA path and adding a sparse MLA decode implementation.

Changes:

• Adds an HPU SparseAttnIndexer implementation that scores Q·K in BF16/FP32 and selects top-K cache slots via torch.topk.
• Routes DeepSeek/GLM indexer and MLA attention to use the HPU sparse path (including BF16 indexer cache handling and SparseAttnIndexer dispatch).
• Extends the HPU MLA attention backend to support sparse decode and updates MLA KV-cache binding to handle HPU’s per-layer KV-cache tuple.

File summaries

File	Description
vllm_gaudi/platform.py	Enables sparse+MLA backend selection logging and adds an HPU platform hook for runner KV-cache behavior.
vllm_gaudi/ops/hpu_sparse_attn_indexer.py	Introduces the HPU BF16 sparse-attention indexer that produces top-K physical slot indices per request.
vllm_gaudi/models/deepseek_v2.py	Monkey-patches DeepSeek indexer cache/forward and SparseAttnIndexer to use the HPU indexer implementation.
vllm_gaudi/attention/oot_mla.py	Enables sparse routing for MLA decode and overrides KV-cache binding to keep the HPU tuple intact.
vllm_gaudi/attention/backends/hpu_attn.py	Adds forward_mqa_sparse to perform sparse decode attention over selected cache slots.

Review details

Suppressed comments (1)

vllm_gaudi/ops/hpu_sparse_attn_indexer.py:54

• In the seq_lens is None fallback, the code only fills :topk columns of topk_indices_buffer and returns without initializing/padding the remaining columns. The sparse attention decode gathers using the full buffer width, so any uninitialized columns can contain stale/out-of-range slot indices and lead to invalid cache reads or indexing errors. Pad the remaining columns with a valid slot index (e.g., last valid slot) when topk < self.topk_tokens.

        topk = min(self.topk_tokens, all_slots.shape[0])
        self.topk_indices_buffer[:batch_size, :topk] = (
            all_slots[:topk].unsqueeze(0).expand(batch_size, -1))
        return self.topk_indices_buffer

• Files reviewed: 5/5 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.

[pawel-olejniczak]

Reviewed at 867254b against VLLM_STABLE_COMMIT 2cf0a691.

The main thing to resolve before this lands: I do not think the sparse path can execute on this branch. There are three independent gates, any one of which routes every DSA decode back to dense MLA - see the inline comments on oot_mla.py:168, hpu_sparse_attn_indexer.py:27 and hpu_sparse_attn_indexer.py:47. The first is that this backport is missing the hpu_model_runner.py change that the main-branch sibling carries; the other five files here are byte-identical to it.

That also affects how I read the validation numbers. Dense MLA is exact for sequences up to index_topk tokens, so GSM8K and AIME agreeing with the model card is consistent with the sparse path never running, and does not by itself show that it did. Something that separates the two paths would help, for example a log line or counter proving forward_mqa_sparse was entered, or an eval at a context length well past index_topk.

Two other asks:

• Performance numbers. DSA is a latency and throughput feature, and the description only reports accuracy, while the dense-fallback PR it supersedes reported median TTFT 1.82 s / TPOT 88.5 ms / 353.7 tok/s. Once the scoring path is actually reached it is a Python loop over the batch with a per-request int(seq_lens[i].item()) host sync, inside @torch.compiler.disable. Could you post TPOT and output throughput at a long context next to that dense baseline?
• pre-commit will fail as is: the pinned yapf 0.43.0 rewrites 12 hunks across four of the five files. Pinned ruff 0.11.7 is clean.

[pawel-olejniczak]

On this branch I believe this early return is the only path this function ever takes. get_kv_cache_spec in vllm_gaudi/v1/worker/hpu_model_runner.py:1717 handles only MambaBase, Attention and MLAAttention, and DeepseekV32IndexerCache is none of those - it is a plain AttentionLayerBase (upstream vllm/model_executor/models/deepseek_v2.py:604). So the indexer cache layer never gets a spec, never gets an allocation, and self.k_cache.kv_cache keeps the torch.tensor([]) it is initialised with at deepseek_v2.py:609. kv_cache.numel() == 0 is then true on every step: the index_copy_ below never inserts K, and the returned indices are always 0..topk_tokens-1.

The main-branch sibling (#1760) adds exactly the missing branch to get_kv_cache_spec, and this PR is otherwise byte-identical to it in all five shared files. The new check_runner_kv_caches_multi_layer no-op in platform.py only has an effect once the indexer cache is in kv_caches, which also points at that change being intended here.

Please cherry-pick the elif isinstance(attn_module, AttentionLayerBase) branch from #1760. One thing worth confirming when you do: it adds a second entry per DSA layer to runner_kv_caches, which is the case upstream's own TODO at vllm/v1/worker/utils.py:529-541 describes as unanalyzed - have you checked that nothing on the HPU side indexes runner_kv_caches positionally?

[jkaniecki]

runner_kv_caches is not consumed by any function / read by any other variable.
kv_caches is also binded to model_runner.kv_Caches, which is consumed by defragmenter. I switched off the defrag for this model (and cont pa) in features.py - we don't want it to manipulate indexer cache. model_runner.kv_caches is also used in hpu_worker.py:L531-550, but there only first item from kv cache is taken and it's not indexer cache.

[iboiko-habana]

@jkaniecki do we need README and examples update?

[jkaniecki]

Reviewed at 867254b against VLLM_STABLE_COMMIT 2cf0a691.

The main thing to resolve before this lands: I do not think the sparse path can execute on this branch. There are three independent gates, any one of which routes every DSA decode back to dense MLA - see the inline comments on oot_mla.py:168, hpu_sparse_attn_indexer.py:27 and hpu_sparse_attn_indexer.py:47. The first is that this backport is missing the hpu_model_runner.py change that the main-branch sibling carries; the other five files here are byte-identical to it.

That also affects how I read the validation numbers. Dense MLA is exact for sequences up to index_topk tokens, so GSM8K and AIME agreeing with the model card is consistent with the sparse path never running, and does not by itself show that it did. Something that separates the two paths would help, for example a log line or counter proving forward_mqa_sparse was entered, or an eval at a context length well past index_topk.

Two other asks:

• Performance numbers. DSA is a latency and throughput feature, and the description only reports accuracy, while the dense-fallback PR it supersedes reported median TTFT 1.82 s / TPOT 88.5 ms / 353.7 tok/s. Once the scoring path is actually reached it is a Python loop over the batch with a per-request int(seq_lens[i].item()) host sync, inside @torch.compiler.disable. Could you post TPOT and output throughput at a long context next to that dense baseline?
• pre-commit will fail as is: the pinned yapf 0.43.0 rewrites 12 hunks across four of the five files. Pinned ruff 0.11.7 is clean.

@pawel-olejniczak Thanks for the complex review, I've changed the code to enable DSA properly, confirmed it is now launching. About testing - for sure GSM8K was only a sanity check, not showing real DSA capabilities. AIME 2026 is a bit different - some samples overcome this 2k topk threshold and switch on the DSA, so it's really working. I've also ran tests up to 128k context to see long-context capabilities using "needle" type task - if was successful in all runs.

About perf numbers - it's only a functional enablement for now, so not sharing perf results, I haven't focus on them. The MLA fallback PR may have some perf numbers (even better than DSA version due to switching t.compile off) but in fact it enables different model - GLM5.2 was trained and should be executed with DSA. Without it, we cannot be sure about it's accuracy for samples larger than top_k. Dense MLA version can also show good acc at multiple sample ranges, but it cannot be said that it's the same acc that GLM5.2 should have then.