fix(sm120): address DeepGEMM review blockers

Author: parkerpang

rtp_llm/model_loader/per_block_fp8_quant_weight.py

@@ -100,7 +100,7 @@ def cast_to_fp8(x: torch.Tensor):
     return x.to(torch.float8_e4m3fn)
 
 
-def per_block_cast_to_fp8(
+def per_block_cast_to_fp8_grouped(
     x: torch.Tensor, group_size: int
 ) -> tuple[torch.Tensor, torch.Tensor]:
     is_2d = x.dim() == 2
@@ -247,6 +247,9 @@ def create_w8a8_fp8_per_block_weight(
 
 
 class PerBlockFp8Weight(CompositeWeight, QuantWeight):
+    def _uses_direct_ue8m0(self) -> bool:
+        return False
+
     w8a8_weight_list: Dict[str, str] = {
         W.attn_qkv_w: W.attn_qkv_s,
         W.attn_o_w: W.attn_o_s,
@@ -838,11 +841,19 @@ def _postprocess(
             )
             # kernel_weight, scale_weight = load_config.exported_device.convert_fp8_weight_params(kernel_weight, scale_weight)
 
-            # Online SM100/SM120 loading can already produce the final packed
+            # Online SM120 loading can already produce the final packed
             # UE8M0 representation directly from source weights.  Legacy/pre-quantized
             # inputs still arrive with floating-point block scales and require
             # the old dequantize/requantize conversion here.
-            if is_deep_gemm_e8m0_used() and scale_weight.dtype != torch.int32:
+            if self._uses_direct_ue8m0():
+                from rtp_llm.models_py.kernels.cuda.fp8_kernel import (
+                    pack_weight_scale_ue8m0,
+                )
+
+                scale_weight = pack_weight_scale_ue8m0(
+                    scale_weight, kernel_weight.shape[-2]
+                )
+            elif is_deep_gemm_e8m0_used():
                 kernel_weight, scale_weight = requant_weight_ue8m0(
                     kernel_weight, scale_weight
                 )
@@ -896,6 +907,16 @@ def __init__(
         self.kernel = kernel
         self.scale = scale
 
+    def _uses_direct_ue8m0(self) -> bool:
+        from rtp_llm.models_py.utils.arch import is_sm12x
+
+        return (
+            self.scale is not None
+            and self.kernel.name not in (W.moe_w1, W.moe_w2)
+            and self.group_size == 128
+            and is_sm12x()
+        )
+
     def _load_raw_tensor(
         self,
         tensor_source: TensorSource,
@@ -907,26 +928,23 @@ def _load_raw_tensor(
             tensor_source, layer_id, device, load_config
         )
 
-        from rtp_llm.models_py.kernels.cuda.deepgemm_wrapper import (
-            is_deep_gemm_e8m0_used,
-        )
-
-        is_dense_weight = self.kernel.name not in (W.moe_w1, W.moe_w2)
-        direct_ue8m0 = (
-            self.scale is not None and is_dense_weight and is_deep_gemm_e8m0_used()
-        )
-        if direct_ue8m0 and self.group_size != 128:
-            raise ValueError(
-                "SM100/SM120 DeepGEMM packed UE8M0 requires group_size=128, "
-                f"got {self.group_size} for {self.kernel.name}"
-            )
+        direct_ue8m0 = self._uses_direct_ue8m0()
+        if (
+            self.scale is not None
+            and self.kernel.name not in (W.moe_w1, W.moe_w2)
+            and not direct_ue8m0
+        ):
+            from rtp_llm.models_py.utils.arch import is_sm12x
 
+            if is_sm12x() and self.group_size != 128:
+                raise ValueError(
+                    "SM120 DeepGEMM packed UE8M0 requires group_size=128, "
+                    f"got {self.group_size} for {self.kernel.name}"
+                )
         res = {}
         scale = None
         if direct_ue8m0:
-            from rtp_llm.models_py.kernels.cuda.fp8_kernel import (
-                quant_weight_ue8m0_packed,
-            )
+            from rtp_llm.models_py.kernels.cuda.fp8_kernel import quant_weight_ue8m0
 
             source_weight = kernel.get(self.kernel.name)
             if source_weight.dim() != 2:
@@ -936,11 +954,11 @@ def _load_raw_tensor(
                     f"{self.kernel.name}"
                 )
             source_weight = source_weight.T
-            quant_kernel, scale = quant_weight_ue8m0_packed(
-                source_weight.contiguous().to(device)
+            quant_kernel, scale = quant_weight_ue8m0(
+                source_weight.contiguous().to(device), [128, 128]
             )
         elif self.scale:
-            quant_kernel, scale = per_block_cast_to_fp8(
+            quant_kernel, scale = per_block_cast_to_fp8_grouped(
                 kernel.get(self.kernel.name), self.group_size
             )
             if quant_kernel.dim() == 2:
@@ -956,10 +974,10 @@ def _load_raw_tensor(
         res = {self.kernel.name: quant_kernel.contiguous().to(device)}
         if self.scale:
             scale = scale.T if scale.dim() == 2 and not direct_ue8m0 else scale
-            # Packed UE8M0 scales intentionally use a non-contiguous TMA
-            # layout (stride(-2) == 1).  Do not normalize that layout here.
-            scale = scale.to(device) if direct_ue8m0 else scale.contiguous().to(device)
-            res.update({self.scale.name: scale})
+            # Keep ordinary [N/128, K/128] float scales through TP/DP/EP
+            # splitting. _postprocess packs each rank's local scale into the
+            # non-contiguous DeepGEMM TMA layout.
+            res.update({self.scale.name: scale.contiguous().to(device)})
 
         return res
 

rtp_llm/models_py/kernels/cuda/deepgemm_wrapper.py

@@ -1,4 +1,5 @@
 import functools
+import threading
 from contextlib import contextmanager
 from typing import Any, Callable, Generator, List, NoReturn, Optional, Tuple
 
@@ -24,6 +25,7 @@
     "transpose_packed_fp4",
     "tf32_hc_prenorm_gemm",
     "has_deep_gemm",
+    "has_deep_gemm_mk_alignment",
     "is_deep_gemm_e8m0_used",
     "configure_deep_gemm_num_sms",
     "configure_deep_gemm_mk_alignment",
@@ -80,6 +82,7 @@
 _cast_back_from_fp4_impl: Callable[..., Any] | None = None
 _transpose_packed_fp4_impl: Callable[..., Any] | None = None
 _tf32_hc_prenorm_gemm_impl: Callable[..., Any] | None = None
+_mk_alignment_lock = threading.Lock()
 
 
 @functools.cache
@@ -88,10 +91,46 @@ def has_deep_gemm() -> bool:
     return has_module("deep_gemm")
 
 
+@functools.cache
+def has_deep_gemm_mk_alignment() -> bool:
+    """Whether DeepGEMM exposes the SM120 contiguous-layout tuning API."""
+    if not has_deep_gemm():
+        return False
+    import deep_gemm
+
+    return all(
+        hasattr(deep_gemm, name)
+        for name in (
+            "get_theoretical_mk_alignment_for_contiguous_layout",
+            "get_mk_alignment_for_contiguous_layout",
+            "set_mk_alignment_for_contiguous_layout",
+        )
+    )
+
+
+def normalize_paged_mqa_context_lens(context_lens: torch.Tensor) -> torch.Tensor:
+    """Adapt paged-MQA context lengths to the installed DeepGEMM contract.
+
+    DeepGEMM 2.5 uses ``[batch, queries]`` while the cuda12 DeepGEMM 2.2
+    package accepts the legacy one-dimensional decode shape. The mk-alignment
+    API is a stable capability marker introduced with the 2.5 package.
+    """
+    if has_deep_gemm_mk_alignment():
+        if context_lens.dim() == 1:
+            return context_lens.unsqueeze(-1)
+        if context_lens.dim() == 2:
+            return context_lens
+        raise ValueError(
+            "DeepGEMM 2.5 paged MQA context_lens must be 1D or 2D, "
+            f"got shape={tuple(context_lens.shape)}"
+        )
+    return context_lens
+
+
 @functools.cache
 def is_deep_gemm_e8m0_used() -> bool:
     # Blackwell SM100 and SM120 DeepGEMM kernels consume packed UE8M0 scales.
-    # SM120 support requires the vLLM-pinned DeepGEMM build (a6b593d or newer).
+    # SM120 support is pinned in deps to DeepGEMM 2.5.0+d7d5eca.cu129.
     return torch.cuda.get_device_capability()[0] in (10, 12)
 
 
@@ -118,34 +157,51 @@ def configure_deep_gemm_num_sms(num_sms: int) -> Generator[None, None, None]:
 def get_theoretical_mk_alignment_for_contiguous_layout(
     expected_m: int, num_groups: int
 ) -> int:
-    """Return DeepGEMM's SM-specific BLOCK_M for a grouped call."""
-    import deep_gemm
+    """Return BLOCK_M for a contiguous grouped call.
 
-    try:
-        return int(
-            deep_gemm.get_theoretical_mk_alignment_for_contiguous_layout(
-                expected_m, num_groups
-            )
+    ``expected_m`` is the total routed-token count. This differs from masked
+    grouped GEMM, whose similarly named argument is a per-expert estimate.
+    """
+    if not has_deep_gemm_mk_alignment():
+        raise RuntimeError(
+            "SM120 contiguous grouped GEMM requires DeepGEMM >= 2.5.0 "
+            "with mk-alignment APIs"
         )
-    except TypeError:
-        return int(
-            deep_gemm.get_theoretical_mk_alignment_for_contiguous_layout(
-                (expected_m + num_groups - 1) // num_groups
-            )
+    import deep_gemm
+
+    return int(
+        deep_gemm.get_theoretical_mk_alignment_for_contiguous_layout(
+            expected_m, num_groups
         )
+    )
 
 
 @contextmanager
 def configure_deep_gemm_mk_alignment(alignment: int) -> Generator[None, None, None]:
     """Match DeepGEMM's scheduler BLOCK_M to contiguous workspace padding."""
+    if not has_deep_gemm_mk_alignment():
+        raise RuntimeError(
+            "SM120 contiguous grouped GEMM requires DeepGEMM >= 2.5.0 "
+            "with mk-alignment APIs"
+        )
     import deep_gemm
 
-    original = deep_gemm.get_mk_alignment_for_contiguous_layout()
-    deep_gemm.set_mk_alignment_for_contiguous_layout(alignment)
-    try:
-        yield
-    finally:
-        deep_gemm.set_mk_alignment_for_contiguous_layout(original)
+    # The alignment is process-global. Serialize set/launch/restore so two
+    # concurrent forwards cannot launch with each other's scheduler setting.
+    with _mk_alignment_lock:
+        original = deep_gemm.get_mk_alignment_for_contiguous_layout()
+        deep_gemm.set_mk_alignment_for_contiguous_layout(alignment)
+        configured = deep_gemm.get_mk_alignment_for_contiguous_layout()
+        if configured != alignment:
+            deep_gemm.set_mk_alignment_for_contiguous_layout(original)
+            raise RuntimeError(
+                "DeepGEMM rejected contiguous-layout mk alignment: "
+                f"requested={alignment}, configured={configured}"
+            )
+        try:
+            yield
+        finally:
+            deep_gemm.set_mk_alignment_for_contiguous_layout(original)
 
 
 def _missing_deep_gemm() -> NoReturn:

rtp_llm/models_py/kernels/cuda/fp8_kernel/__init__.py

@@ -6,9 +6,10 @@
     create_per_token_group_quant_fp8_output_scale,
     cutlass_moe_mm_fp8_scaled,
     get_best_config_swap_ab,
+    pack_weight_scale_ue8m0,
     per_block_cast_to_fp8,
     per_token_cast_to_fp8,
-    quant_weight_ue8m0_packed,
+    quant_weight_ue8m0,
     requant_weight_ue8m0,
     scaled_fp8_per_tensor_quant,
     scaled_fp8_per_token_quant,
@@ -21,9 +22,10 @@
     "scaled_fp8_per_token_quant",
     "cutlass_moe_mm_fp8_scaled",
     "get_best_config_swap_ab",
+    "pack_weight_scale_ue8m0",
     "per_token_cast_to_fp8",
     "per_block_cast_to_fp8",
-    "quant_weight_ue8m0_packed",
+    "quant_weight_ue8m0",
     "requant_weight_ue8m0",
     "create_per_token_group_quant_fp8_output_scale",
 ]

rtp_llm/models_py/kernels/cuda/fp8_kernel/fp8_kernel.py

@@ -430,39 +430,16 @@ def quant_weight_ue8m0(
     return out_w, out_s
 
 
-def quant_weight_ue8m0_packed(
-    weight_dequant: torch.Tensor,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """Quantize source weights once to FP8 with packed UE8M0 scales.
-
-    This is the native SM100/SM120 DeepGEMM representation.  Keeping the
-    original floating-point tensor until this operation avoids the lossy
-    float-scale FP8 -> BF16 -> UE8M0 FP8 requantization path.
-    """
-    if weight_dequant.dim() != 2:
+def pack_weight_scale_ue8m0(
+    weight_scale: torch.Tensor, weight_rows: int
+) -> torch.Tensor:
+    """Pack split-local UE8M0 block scales into DeepGEMM's TMA layout."""
+    if weight_scale.dtype != torch.float32 or weight_scale.dim() != 2:
         raise ValueError(
-            "Direct packed UE8M0 weight quantization requires a 2D tensor, "
-            f"got shape {tuple(weight_dequant.shape)}"
-        )
-
-    # Limit floating-point temporaries during online loading.  Quantizing the
-    # complete matrix in one call materializes both a padded source and a
-    # scaled floating-point tensor in addition to the master weight.  Chunking
-    # on a 128-row boundary preserves exactly the same quantization blocks.
-    row_chunk_size = 1024
-    n, k = weight_dequant.shape
-    out_w = torch.empty((n, k), dtype=fp8_dtype, device=weight_dequant.device)
-    scale_chunks = []
-    for row_start in range(0, n, row_chunk_size):
-        row_end = min(row_start + row_chunk_size, n)
-        chunk_w, chunk_s = per_block_cast_to_fp8(
-            weight_dequant[row_start:row_end], use_ue8m0=True
+            "UE8M0 block scales must be a 2D float32 tensor before packing, "
+            f"got shape={tuple(weight_scale.shape)}, dtype={weight_scale.dtype}"
         )
-        out_w[row_start:row_end].copy_(chunk_w)
-        scale_chunks.append(chunk_s)
-
-    out_s = torch.cat(scale_chunks, dim=0)
-    return out_w, _transform_scale_ue8m0(out_s, mn=out_w.shape[-2])
+    return _transform_scale_ue8m0(weight_scale, mn=weight_rows)
 
 
 def requant_weight_ue8m0(

rtp_llm/models_py/modules/base/cuda/indexer_op.py

@@ -385,10 +385,9 @@ def _get_topk_paged(
             attention_inputs.kv_cache_kernel_block_id_device.shape[1] * self.blocksize
         )
 
-        # DeepGEMM paged MQA represents the number of queries decoded together
-        # as the second dimension. RTP-LLM currently decodes one query per
-        # sequence, so expand [batch] to [batch, 1].
-        context_lens = fmha_params.kvlen_d.view(-1, 1)
+        # DeepGEMM 2.5 requires [batch, queries], while the cuda12 2.2 package
+        # retains the legacy one-dimensional decode contract.
+        context_lens = deep_gemm.normalize_paged_mqa_context_lens(fmha_params.kvlen_d)
         schedule_metadata = deep_gemm.get_paged_mqa_logits_metadata(
             context_lens,
             self.blocksize,