[XPU] Fix sparse-MLA metadata sync (#52066)
Signed-off-by: Libin Tang <libin.tang@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Libin Tang committed
Aug 24, 2026 at 19:05 UTC
cbe3966f9cc62848d6271a92819e51470f74fbbd
1 file changed
+22
vllm/v1/attention/backends/mla/xpu_mla_sparse.py
+22
@@ -88,6 +88,23 @@ class XPUMLASparseMetadata(AttentionMetadata):
88
block_size: int = 1
89
topk_tokens: int = 2048
90
91
+ # The shared MLA layer (`mla_attention.py::forward_impl`) reads these
92
+ # decode/prefill counts unconditionally for every MLA metadata (it asserts
93
+ # `num_decodes`/`num_prefills`/`num_decode_tokens is not None` and uses
94
+ # `num_decode_tokens` to split MQA vs dense-MHA tokens). The CUDA sparse
95
+ # backends carry them via `SparseMLACommonMetadataBuilder`; this XPU backend
96
+ # builds its own metadata and previously omitted them, so a sparse-MLA
97
+ # (DeepSeek / GLM DSA) run on XPU crashed with
98
+ # `'XPUMLASparseMetadata' object has no attribute 'num_decode_tokens'`.
99
+ # This backend serves both prefill and decode through the top-k sparse MQA
100
+ # path (see `forward_mqa`), so all tokens are routed as "decode"
101
+ # (`num_decode_tokens == num_actual_tokens`, `num_prefills == 0`); that keeps
102
+ # the shared layer's `num_mha_tokens` at 0 and never enters the dense-MHA
103
+ # prefill branch (which needs prefill-only fields this backend lacks).
104
+ num_decodes: int = 0
105
+ num_prefills: int = 0
106
+ num_decode_tokens: int = 0
107
+
108
109
@dataclass
110
class XPUMLASparseMetadataBuilder(AttentionMetadataBuilder[XPUMLASparseMetadata]):
@@ -157,6 +174,11 @@ class XPUMLASparseMetadataBuilder(AttentionMetadataBuilder[XPUMLASparseMetadata]
174
req_id_per_token=req_id_per_token,
175
block_size=self.kv_cache_spec.block_size,
176
topk_tokens=self.topk_tokens,
177
+ # Route every token through the sparse MQA path (see the field
178
+ # definitions above); this backend has no dense-MHA prefill.
179
+ num_decodes=common_attn_metadata.num_reqs,
180
+ num_prefills=0,
181
+ num_decode_tokens=common_attn_metadata.num_actual_tokens,
182
)
183
return metadata
184