[ROCm] [BugFix] Fix Quark GLM-5.2 Checkpoint inference: indexer wk per-channel FP8 dequant + missing sparse-MLA metadata fields (#48886)
Signed-off-by: Colin Zeng <Colin.Zeng@amd.com> Signed-off-by: ColinZ22 <Colin.Zeng@amd.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com> Co-authored-by: fanxingran <xingran.fan@amd.com>
Colin Z committed
Jul 27, 2026 at 15:38 UTC
28158b2fc3640fa18063b85b3357f41ea9919d7d
5 files changed
+129
-21
tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py
+75
@@ -88,6 +88,81 @@ def _make_common_metadata():
88
)
89
90
91
+def _make_mixed_common_metadata():
92
+ # req0: query_len 1 (decode), req1: query_len 4 (prefill) -> decode-first
93
+ query_start_loc = torch.tensor([0, 1, 5], dtype=torch.int32, device="cpu")
94
+ seq_lens = torch.tensor([16, 10], dtype=torch.int32, device="cpu")
95
+ return CommonAttentionMetadata(
96
+ query_start_loc=query_start_loc,
97
+ query_start_loc_cpu=query_start_loc,
98
+ seq_lens=seq_lens,
99
+ _seq_lens_cpu=seq_lens,
100
+ num_reqs=2,
101
+ num_actual_tokens=5,
102
+ max_query_len=4,
103
+ max_seq_len=16,
104
+ block_table_tensor=torch.arange(16, dtype=torch.int32, device="cpu").view(2, 8),
105
+ slot_mapping=torch.arange(5, dtype=torch.int64, device="cpu"),
106
+ )
107
+
108
+
109
+def _patch_build_deps(monkeypatch, events=None):
110
+ """Stub the aiter kernel, triton helper and CUDA sync so ``build()`` runs
111
+ on CPU."""
112
+
113
+ def fake_generate_sparse_seqlen_triton(
114
+ query_lens, seq_lens, cu_query_lens, topk_token, num_tokens, max_query_len
115
+ ):
116
+ return torch.zeros(num_tokens, dtype=torch.int32, device="cpu")
117
+
118
+ fake_aiter = _FakeAiter("aiter")
119
+ fake_aiter.get_mla_metadata_v1 = Mock(side_effect=lambda *a, **k: None)
120
+ monkeypatch.setitem(sys.modules, "aiter", fake_aiter)
121
+ monkeypatch.setattr(
122
+ sparse_mod, "generate_sparse_seqlen_triton", fake_generate_sparse_seqlen_triton
123
+ )
124
+ monkeypatch.setattr(
125
+ sparse_mod.torch.cuda,
126
+ "current_stream",
127
+ lambda device=None: SimpleNamespace(
128
+ synchronize=lambda: events.append("sync") if events is not None else None
129
+ ),
130
+ )
131
+
132
+
133
+def test_build_populates_decode_only_split_fields(monkeypatch):
134
+ """Decode-only batch: all reqs count as decodes, prefill fields default."""
135
+ builder = _make_builder()
136
+ _patch_build_deps(monkeypatch)
137
+
138
+ md = builder.build(
139
+ common_prefix_len=0, common_attn_metadata=_make_common_metadata()
140
+ )
141
+
142
+ assert md.num_decodes == 2
143
+ assert md.num_prefills == 0
144
+ assert md.num_decode_tokens == 2
145
+ assert md.prefill_max_seq_len == 0
146
+ assert md.prefill is None
147
+
148
+
149
+def test_build_populates_mixed_split_fields(monkeypatch):
150
+ """Mixed decode+prefill batch: split is reported, prefill fields stay
151
+ default because this impl always runs the MQA path."""
152
+ builder = _make_builder()
153
+ _patch_build_deps(monkeypatch)
154
+
155
+ md = builder.build(
156
+ common_prefix_len=0, common_attn_metadata=_make_mixed_common_metadata()
157
+ )
158
+
159
+ assert md.num_decodes == 1
160
+ assert md.num_prefills == 1
161
+ assert md.num_decode_tokens == 1
162
+ assert md.prefill_max_seq_len == 0
163
+ assert md.prefill is None
164
+
165
+
166
def test_sparse_persistent_metadata_syncs_only_after_recompute(monkeypatch):
167
builder = _make_builder()
168
common_metadata = _make_common_metadata()
vllm/model_executor/layers/attention/mla_attention.py
+26
-19
@@ -515,29 +515,36 @@ class MLAAttention(nn.Module, AttentionLayerBase):
515
compilation_config.static_forward_context[prefix] = self
516
517
self.prefill_backend: MLAPrefillBackend | None
518
- try:
519
- prefill_backend_cls = get_mla_prefill_backend(vllm_config)
520
- except ValueError:
521
- if (
522
- not self.impl.is_sparse
523
- or vllm_config.attention_config.mla_prefill_backend is not None
524
- ):
525
- raise
518
+ if self.impl.is_sparse and not self.impl.supports_dense_mha_prefill:
519
logger.warning_once(
527
- "No MLA prefill backend supports this model; sparse MLA will use the "
528
- "top-k MQA path only (no dense-MHA prefill)."
520
+ "Sparse MLA impl has no dense-MHA prefill path; using the top-k "
521
+ "MQA path only."
522
)
523
self.prefill_backend = None
524
else:
532
- self.prefill_backend = prefill_backend_cls(
533
- num_heads=self.num_heads,
534
- scale=self.scale,
535
- kv_lora_rank=self.kv_lora_rank,
536
- qk_nope_head_dim=self.qk_nope_head_dim,
537
- qk_rope_head_dim=self.qk_rope_head_dim,
538
- v_head_dim=self.v_head_dim,
539
- vllm_config=vllm_config,
540
- )
525
+ try:
526
+ prefill_backend_cls = get_mla_prefill_backend(vllm_config)
527
+ except ValueError:
528
+ if (
529
+ not self.impl.is_sparse
530
+ or vllm_config.attention_config.mla_prefill_backend is not None
531
+ ):
532
+ raise
533
+ logger.warning_once(
534
+ "No MLA prefill backend supports this model; sparse MLA will "
535
+ "use the top-k MQA path only (no dense-MHA prefill)."
536
+ )
537
+ self.prefill_backend = None
538
+ else:
539
+ self.prefill_backend = prefill_backend_cls(
540
+ num_heads=self.num_heads,
541
+ scale=self.scale,
542
+ kv_lora_rank=self.kv_lora_rank,
543
+ qk_nope_head_dim=self.qk_nope_head_dim,
544
+ qk_rope_head_dim=self.qk_rope_head_dim,
545
+ v_head_dim=self.v_head_dim,
546
+ vllm_config=vllm_config,
547
+ )
548
549
self.kv_cache = torch.tensor([])
550
vllm/model_executor/models/deepseek_v2.py
+7
-2
@@ -851,11 +851,16 @@ def _try_load_fp8_indexer_wk(
851
# We have both weight and scale: dequantize FP8 to BF16.
852
weight_fp8, scale_inv = entry["weight"], entry["scale"]
853
del buf[layer_prefix]
854
- block_size = weight_fp8.shape[1] // scale_inv.shape[1]
854
+ if scale_inv.ndim == 1:
855
+ # Per-channel scale: one scale per row of [out, in]
856
+ group_shape = GroupShape(1, weight_fp8.shape[1])
857
+ else:
858
+ block_size = weight_fp8.shape[1] // scale_inv.shape[1]
859
+ group_shape = GroupShape(block_size, block_size)
860
weight_bf16 = scaled_dequantize(
861
weight_fp8,
862
scale_inv,
858
- group_shape=GroupShape(block_size, block_size),
863
+ group_shape=group_shape,
864
out_dtype=torch.bfloat16,
865
)
866
vllm/v1/attention/backend.py
+4
@@ -803,6 +803,10 @@ class AttentionImplBase(ABC, Generic[T]):
803
# route between the dense-MHA prefill and sparse-MQA paths.
804
is_sparse: ClassVar[bool] = False
805
806
+ # Whether this impl provides a dense-MHA prefill path (forward_mha). Sparse
807
+ # impls without one run the top-k MQA path for all requests.
808
+ supports_dense_mha_prefill: ClassVar[bool] = True
809
+
810
# Required attributes that all impls should have
811
num_heads: int
812
head_size: int
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
+17
@@ -30,6 +30,7 @@ from vllm.v1.attention.backend import (
30
from vllm.v1.attention.backends.mla.rocm_aiter_mla import (
31
AiterMLAHelper,
32
)
33
+from vllm.v1.attention.backends.utils import split_decodes_and_prefills
34
from vllm.v1.kv_cache_interface import AttentionSpec
35
from vllm.v1.worker.workspace import current_workspace_manager
36
@@ -333,6 +334,14 @@ class ROCMAiterMLASparseMetadata(AttentionMetadata):
334
block_size: int = 1
335
topk_tokens: int = 2048
336
337
+ # Fields read by the shared MLA forward. This impl has no dense-MHA prefill
338
+ # path (supports_dense_mha_prefill=False), so it always runs the MQA path;
339
+ num_decodes: int = 0
340
+ num_prefills: int = 0
341
+ num_decode_tokens: int = 0
342
+ prefill_max_seq_len: int = 0
343
+ prefill: object = None
344
+
345
# Persistent MLA metadata (only populated when persistent mode is enabled,
346
# i.e. when the aiter sparse decode kernel supports work-stealing splits).
347
work_meta_data: torch.Tensor | None = None
@@ -488,6 +497,10 @@ class ROCMAiterMLASparseMetadataBuilder(
497
fast_build: bool = False,
498
) -> ROCMAiterMLASparseMetadata:
499
num_tokens = common_attn_metadata.num_actual_tokens
500
+ (num_decodes, num_prefills, num_decode_tokens, _) = split_decodes_and_prefills(
501
+ common_attn_metadata,
502
+ decode_threshold=self.reorder_batch_threshold or 1,
503
+ )
504
starts = np.asarray(common_attn_metadata.query_start_loc_cpu, dtype=np.int32)
505
seg_lengths = np.diff(starts)
506
req_id_per_token = np.repeat(
@@ -604,6 +617,9 @@ class ROCMAiterMLASparseMetadataBuilder(
617
block_size=self.kv_cache_spec.block_size,
618
attn_out_dtype=self.model_dtype,
619
topk_tokens=self.topk_tokens,
620
+ num_decodes=num_decodes,
621
+ num_prefills=num_prefills,
622
+ num_decode_tokens=num_decode_tokens,
623
qo_indptr=qo_indptr,
624
paged_kv_last_page_len=paged_kv_last_page_len,
625
paged_kv_indices=paged_kv_indices,
@@ -649,6 +665,7 @@ def reference_mla_sparse_prefill(
665
666
class ROCMAiterMLASparseImpl(MLAAttentionImpl[ROCMAiterMLASparseMetadata]):
667
is_sparse = True
668
+ supports_dense_mha_prefill = False
669
670
def __init__(
671
self,