Skip to content

Commit 48ccd0b

Browse files
committed
[XPU] Add sparse-MLA decode/prefill metadata for the shared MLA layer
The shared MLA layer (mla_attention.py::forward_impl) reads num_decodes/num_prefills/num_decode_tokens on every MLA metadata (it asserts they are not None and uses num_decode_tokens to split MQA vs dense-MHA tokens). The CUDA sparse backends supply these via SparseMLACommonMetadataBuilder, but the XPU sparse backend builds its own metadata and omitted them, so a sparse-MLA (DeepSeek DSA / GLM) run on XPU crashed with: 'XPUMLASparseMetadata' object has no attribute 'num_decode_tokens' Add the fields and populate them so all tokens route through the top-k sparse MQA path (num_decode_tokens == num_actual_tokens, num_prefills == 0), which this backend already implements for both prefill and decode. This keeps the shared layer's num_mha_tokens at 0 and never enters the dense-MHA prefill branch (whose prefill-only fields this backend does not provide). XPU-only; does not touch CUDA/ROCm/CPU backends. The ragged-N FP8 block-scale GEMM fix that previously accompanied this change is now covered upstream by #52118, so it is dropped here. Signed-off-by: Libin Tang <libin.tang@intel.com>
1 parent ac7509e commit 48ccd0b

1 file changed

Lines changed: 22 additions & 0 deletions

File tree

vllm/v1/attention/backends/mla/xpu_mla_sparse.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,23 @@ class XPUMLASparseMetadata(AttentionMetadata):
9797
block_size: int = 1
9898
topk_tokens: int = 2048
9999

100+
# The shared MLA layer (`mla_attention.py::forward_impl`) reads these
101+
# decode/prefill counts unconditionally for every MLA metadata (it asserts
102+
# `num_decodes`/`num_prefills`/`num_decode_tokens is not None` and uses
103+
# `num_decode_tokens` to split MQA vs dense-MHA tokens). The CUDA sparse
104+
# backends carry them via `SparseMLACommonMetadataBuilder`; this XPU backend
105+
# builds its own metadata and previously omitted them, so a sparse-MLA
106+
# (DeepSeek / GLM DSA) run on XPU crashed with
107+
# `'XPUMLASparseMetadata' object has no attribute 'num_decode_tokens'`.
108+
# This backend serves both prefill and decode through the top-k sparse MQA
109+
# path (see `forward_mqa`), so all tokens are routed as "decode"
110+
# (`num_decode_tokens == num_actual_tokens`, `num_prefills == 0`); that keeps
111+
# the shared layer's `num_mha_tokens` at 0 and never enters the dense-MHA
112+
# prefill branch (which needs prefill-only fields this backend lacks).
113+
num_decodes: int = 0
114+
num_prefills: int = 0
115+
num_decode_tokens: int = 0
116+
100117

101118
@dataclass
102119
class XPUMLASparseMetadataBuilder(AttentionMetadataBuilder[XPUMLASparseMetadata]):
@@ -166,6 +183,11 @@ def build(
166183
req_id_per_token=req_id_per_token,
167184
block_size=self.kv_cache_spec.block_size,
168185
topk_tokens=self.topk_tokens,
186+
# Route every token through the sparse MQA path (see the field
187+
# definitions above); this backend has no dense-MHA prefill.
188+
num_decodes=common_attn_metadata.num_reqs,
189+
num_prefills=0,
190+
num_decode_tokens=common_attn_metadata.num_actual_tokens,
169191
)
170192
return metadata
171193

0 commit comments

Comments
 (0)