[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:
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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user