[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>

tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py

@@ -88,6 +88,81 @@ def _make_common_metadata():
     )
 
 
+def _make_mixed_common_metadata():
+    # req0: query_len 1 (decode), req1: query_len 4 (prefill) -> decode-first
+    query_start_loc = torch.tensor([0, 1, 5], dtype=torch.int32, device="cpu")
+    seq_lens = torch.tensor([16, 10], dtype=torch.int32, device="cpu")
+    return CommonAttentionMetadata(
+        query_start_loc=query_start_loc,
+        query_start_loc_cpu=query_start_loc,
+        seq_lens=seq_lens,
+        _seq_lens_cpu=seq_lens,
+        num_reqs=2,
+        num_actual_tokens=5,
+        max_query_len=4,
+        max_seq_len=16,
+        block_table_tensor=torch.arange(16, dtype=torch.int32, device="cpu").view(2, 8),
+        slot_mapping=torch.arange(5, dtype=torch.int64, device="cpu"),
+    )
+
+
+def _patch_build_deps(monkeypatch, events=None):
+    """Stub the aiter kernel, triton helper and CUDA sync so ``build()`` runs
+    on CPU."""
+
+    def fake_generate_sparse_seqlen_triton(
+        query_lens, seq_lens, cu_query_lens, topk_token, num_tokens, max_query_len
+    ):
+        return torch.zeros(num_tokens, dtype=torch.int32, device="cpu")
+
+    fake_aiter = _FakeAiter("aiter")
+    fake_aiter.get_mla_metadata_v1 = Mock(side_effect=lambda *a, **k: None)
+    monkeypatch.setitem(sys.modules, "aiter", fake_aiter)
+    monkeypatch.setattr(
+        sparse_mod, "generate_sparse_seqlen_triton", fake_generate_sparse_seqlen_triton
+    )
+    monkeypatch.setattr(
+        sparse_mod.torch.cuda,
+        "current_stream",
+        lambda device=None: SimpleNamespace(
+            synchronize=lambda: events.append("sync") if events is not None else None
+        ),
+    )
+
+
+def test_build_populates_decode_only_split_fields(monkeypatch):
+    """Decode-only batch: all reqs count as decodes, prefill fields default."""
+    builder = _make_builder()
+    _patch_build_deps(monkeypatch)
+
+    md = builder.build(
+        common_prefix_len=0, common_attn_metadata=_make_common_metadata()
+    )
+
+    assert md.num_decodes == 2
+    assert md.num_prefills == 0
+    assert md.num_decode_tokens == 2
+    assert md.prefill_max_seq_len == 0
+    assert md.prefill is None
+
+
+def test_build_populates_mixed_split_fields(monkeypatch):
+    """Mixed decode+prefill batch: split is reported, prefill fields stay
+    default because this impl always runs the MQA path."""
+    builder = _make_builder()
+    _patch_build_deps(monkeypatch)
+
+    md = builder.build(
+        common_prefix_len=0, common_attn_metadata=_make_mixed_common_metadata()
+    )
+
+    assert md.num_decodes == 1
+    assert md.num_prefills == 1
+    assert md.num_decode_tokens == 1
+    assert md.prefill_max_seq_len == 0
+    assert md.prefill is None
+
+
 def test_sparse_persistent_metadata_syncs_only_after_recompute(monkeypatch):
     builder = _make_builder()
     common_metadata = _make_common_metadata()

vllm/model_executor/layers/attention/mla_attention.py

@@ -515,29 +515,36 @@ class MLAAttention(nn.Module, AttentionLayerBase):
         compilation_config.static_forward_context[prefix] = self
 
         self.prefill_backend: MLAPrefillBackend | None
-        try:
-            prefill_backend_cls = get_mla_prefill_backend(vllm_config)
-        except ValueError:
-            if (
-                not self.impl.is_sparse
-                or vllm_config.attention_config.mla_prefill_backend is not None
-            ):
-                raise
+        if self.impl.is_sparse and not self.impl.supports_dense_mha_prefill:
             logger.warning_once(
-                "No MLA prefill backend supports this model; sparse MLA will use the "
-                "top-k MQA path only (no dense-MHA prefill)."
+                "Sparse MLA impl has no dense-MHA prefill path; using the top-k "
+                "MQA path only."
             )
             self.prefill_backend = None
         else:
-            self.prefill_backend = prefill_backend_cls(
-                num_heads=self.num_heads,
-                scale=self.scale,
-                kv_lora_rank=self.kv_lora_rank,
-                qk_nope_head_dim=self.qk_nope_head_dim,
-                qk_rope_head_dim=self.qk_rope_head_dim,
-                v_head_dim=self.v_head_dim,
-                vllm_config=vllm_config,
-            )
+            try:
+                prefill_backend_cls = get_mla_prefill_backend(vllm_config)
+            except ValueError:
+                if (
+                    not self.impl.is_sparse
+                    or vllm_config.attention_config.mla_prefill_backend is not None
+                ):
+                    raise
+                logger.warning_once(
+                    "No MLA prefill backend supports this model; sparse MLA will "
+                    "use the top-k MQA path only (no dense-MHA prefill)."
+                )
+                self.prefill_backend = None
+            else:
+                self.prefill_backend = prefill_backend_cls(
+                    num_heads=self.num_heads,
+                    scale=self.scale,
+                    kv_lora_rank=self.kv_lora_rank,
+                    qk_nope_head_dim=self.qk_nope_head_dim,
+                    qk_rope_head_dim=self.qk_rope_head_dim,
+                    v_head_dim=self.v_head_dim,
+                    vllm_config=vllm_config,
+                )
 
         self.kv_cache = torch.tensor([])
 

vllm/model_executor/models/deepseek_v2.py

@@ -851,11 +851,16 @@ def _try_load_fp8_indexer_wk(
     # We have both weight and scale: dequantize FP8 to BF16.
     weight_fp8, scale_inv = entry["weight"], entry["scale"]
     del buf[layer_prefix]
-    block_size = weight_fp8.shape[1] // scale_inv.shape[1]
+    if scale_inv.ndim == 1:
+        # Per-channel scale: one scale per row of [out, in]
+        group_shape = GroupShape(1, weight_fp8.shape[1])
+    else:
+        block_size = weight_fp8.shape[1] // scale_inv.shape[1]
+        group_shape = GroupShape(block_size, block_size)
     weight_bf16 = scaled_dequantize(
         weight_fp8,
         scale_inv,
-        group_shape=GroupShape(block_size, block_size),
+        group_shape=group_shape,
         out_dtype=torch.bfloat16,
     )
 

vllm/v1/attention/backend.py

@@ -803,6 +803,10 @@ class AttentionImplBase(ABC, Generic[T]):
     # route between the dense-MHA prefill and sparse-MQA paths.
     is_sparse: ClassVar[bool] = False
 
+    # Whether this impl provides a dense-MHA prefill path (forward_mha). Sparse
+    # impls without one run the top-k MQA path for all requests.
+    supports_dense_mha_prefill: ClassVar[bool] = True
+
     # Required attributes that all impls should have
     num_heads: int
     head_size: int

vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py

@@ -30,6 +30,7 @@ from vllm.v1.attention.backend import (
 from vllm.v1.attention.backends.mla.rocm_aiter_mla import (
     AiterMLAHelper,
 )
+from vllm.v1.attention.backends.utils import split_decodes_and_prefills
 from vllm.v1.kv_cache_interface import AttentionSpec
 from vllm.v1.worker.workspace import current_workspace_manager
 
@@ -333,6 +334,14 @@ class ROCMAiterMLASparseMetadata(AttentionMetadata):
     block_size: int = 1
     topk_tokens: int = 2048
 
+    # Fields read by the shared MLA forward. This impl has no dense-MHA prefill
+    # path (supports_dense_mha_prefill=False), so it always runs the MQA path;
+    num_decodes: int = 0
+    num_prefills: int = 0
+    num_decode_tokens: int = 0
+    prefill_max_seq_len: int = 0
+    prefill: object = None
+
     # Persistent MLA metadata (only populated when persistent mode is enabled,
     # i.e. when the aiter sparse decode kernel supports work-stealing splits).
     work_meta_data: torch.Tensor | None = None
@@ -488,6 +497,10 @@ class ROCMAiterMLASparseMetadataBuilder(
         fast_build: bool = False,
     ) -> ROCMAiterMLASparseMetadata:
         num_tokens = common_attn_metadata.num_actual_tokens
+        (num_decodes, num_prefills, num_decode_tokens, _) = split_decodes_and_prefills(
+            common_attn_metadata,
+            decode_threshold=self.reorder_batch_threshold or 1,
+        )
         starts = np.asarray(common_attn_metadata.query_start_loc_cpu, dtype=np.int32)
         seg_lengths = np.diff(starts)
         req_id_per_token = np.repeat(
@@ -604,6 +617,9 @@ class ROCMAiterMLASparseMetadataBuilder(
             block_size=self.kv_cache_spec.block_size,
             attn_out_dtype=self.model_dtype,
             topk_tokens=self.topk_tokens,
+            num_decodes=num_decodes,
+            num_prefills=num_prefills,
+            num_decode_tokens=num_decode_tokens,
             qo_indptr=qo_indptr,
             paged_kv_last_page_len=paged_kv_last_page_len,
             paged_kv_indices=paged_kv_indices,
@@ -649,6 +665,7 @@ def reference_mla_sparse_prefill(
 
 class ROCMAiterMLASparseImpl(MLAAttentionImpl[ROCMAiterMLASparseMetadata]):
     is_sparse = True
+    supports_dense_mha_prefill = False
 
     def __init__(
         self,