Commit 48ccd0b
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
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
97 | 97 | | |
98 | 98 | | |
99 | 99 | | |
| 100 | + | |
| 101 | + | |
| 102 | + | |
| 103 | + | |
| 104 | + | |
| 105 | + | |
| 106 | + | |
| 107 | + | |
| 108 | + | |
| 109 | + | |
| 110 | + | |
| 111 | + | |
| 112 | + | |
| 113 | + | |
| 114 | + | |
| 115 | + | |
| 116 | + | |
100 | 117 | | |
101 | 118 | | |
102 | 119 | | |
| |||
166 | 183 | | |
167 | 184 | | |
168 | 185 | | |
| 186 | + | |
| 187 | + | |
| 188 | + | |
| 189 | + | |
| 190 | + | |
169 | 191 | | |
170 | 192 | | |
171 | 193 | | |
| |||
0 commit comments