[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>
This commit is contained in:
Libin Tang
2026-08-25 02:05:35 +00:00
committed by GitHub
co-authored by Kunshang Ji
parent 7b7e5cbfcd
commit cbe3966f9c
@@ -88,6 +88,23 @@ class XPUMLASparseMetadata(AttentionMetadata):
block_size: int = 1
topk_tokens: int = 2048
# The shared MLA layer (`mla_attention.py::forward_impl`) reads these
# decode/prefill counts unconditionally for every MLA metadata (it asserts
# `num_decodes`/`num_prefills`/`num_decode_tokens is not None` and uses
# `num_decode_tokens` to split MQA vs dense-MHA tokens). The CUDA sparse
# backends carry them via `SparseMLACommonMetadataBuilder`; this XPU backend
# builds its own metadata and previously omitted them, so a sparse-MLA
# (DeepSeek / GLM DSA) run on XPU crashed with
# `'XPUMLASparseMetadata' object has no attribute 'num_decode_tokens'`.
# This backend serves both prefill and decode through the top-k sparse MQA
# path (see `forward_mqa`), so all tokens are routed as "decode"
# (`num_decode_tokens == num_actual_tokens`, `num_prefills == 0`); that keeps
# the shared layer's `num_mha_tokens` at 0 and never enters the dense-MHA
# prefill branch (which needs prefill-only fields this backend lacks).
num_decodes: int = 0
num_prefills: int = 0
num_decode_tokens: int = 0
@dataclass
class XPUMLASparseMetadataBuilder(AttentionMetadataBuilder[XPUMLASparseMetadata]):
@@ -157,6 +174,11 @@ class XPUMLASparseMetadataBuilder(AttentionMetadataBuilder[XPUMLASparseMetadata]
req_id_per_token=req_id_per_token,
block_size=self.kv_cache_spec.block_size,
topk_tokens=self.topk_tokens,
# Route every token through the sparse MQA path (see the field
# definitions above); this backend has no dense-MHA prefill.
num_decodes=common_attn_metadata.num_reqs,
num_prefills=0,
num_decode_tokens=common_attn_metadata.num_actual_tokens,
)
return metadata