[Perf] dsv41: Mega-mHC and the wo_a einsum write MXFP8 for wqkv_a and wo_b

python/sglang/kernels/ops/layernorm/mhc_mega.py

@@ -26,13 +26,16 @@ def mhc_mega_boundary(
     hc_eps,
     norm_eps,
     sinkhorn_iters,
+    y_fp8=None,
+    y_gemm_sf=None,
 ):
     """Update the residual, collapse with carried pre, and compute next stats.
 
     The collapse uses ``pre`` from the preceding sublayer, not the new pre
     produced from the updated residual. Preserve this DSv4.1 shifted ordering.
     Used by the opt-in SM10x prefill path. Before CUDA graph capture, warm up
     on the capture stream to initialize DeepGEMM's per-stream barriers.
+    ``y_fp8`` / ``y_gemm_sf`` also receive the normed rows as a DeepGEMM MXFP8 pair.
     """
     rows, hidden = x.shape
     hc = residual.shape[1]
@@ -64,5 +67,7 @@ def mhc_mega_boundary(
         new_post_mix=next_post,
         new_comb_res_mix=next_comb,
         y_bf16=normalized,
+        y_fp8=y_fp8,
+        y_gemm_sf=y_gemm_sf,
     )
     return updated, normalized, (next_pre.squeeze(-1), next_post.squeeze(-1), next_comb)

python/sglang/srt/layers/linear.py

@@ -265,6 +265,9 @@ def __init__(
             params_dtype = torch.get_default_dtype()
         self.params_dtype = params_dtype
         self.quant_config = quant_config
+        # A model sets this before weights load to ask for the DeepGEMM-packed
+        # MXFP8 input path; Fp8LinearMethod keeps it True only once prepared.
+        self.deepgemm_mxfp8_input = False
         if quant_config is None:
             from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
 

python/sglang/srt/layers/quantization/fp8.py

@@ -62,13 +62,15 @@
     can_serve_block_fp8_as_mxfp8,
     cutlass_fp8_supported,
     deepgemm_w8a8_block_fp8_linear_with_fallback,
+    deepgemm_w8a8_mxfp8_linear_with_fallback,
     dense_mxscale_wanted,
     dispatch_block_fp8_mxfp8_linear,
     dispatch_w8a8_block_fp8_linear,
     dispatch_w8a8_mxfp8_linear,
     input_to_float8,
     mxfp8_group_quantize,
     normalize_e4m3fn_to_e4m3fnuz,
+    pack_mxfp8_weight_scale_for_deepgemm,
     prepare_xpu_block_scale_for_scaled_mm,
     requant_block_scale_ue8m0_for_deepgemm,
     resolve_block_fp8_mxfp8_backend,
@@ -968,6 +970,38 @@ def _prepare_block_fp8_as_mxfp8(self, layer: Module) -> None:
         self._process_mxfp8_linear_weight_scale(layer, scale_u8=scale_u8)
         layer.block_fp8_mxfp8_ready = True
 
+    def prepare_deepgemm_mxfp8_input(self, layer: Module) -> bool:
+        """Pack the weight scale a FlashInfer layer needs for a DeepGEMM-packed input pair."""
+        from sglang.srt.layers import deep_gemm_wrapper
+
+        backend = self.mxfp8_dense_backend
+        if not (
+            _is_cuda
+            and get_platform().is_blackwell
+            and backend is not None
+            and (backend.is_flashinfer_cutedsl() or backend.is_flashinfer_cutlass())
+            and deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
+            and deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0
+        ):
+            return False
+        n, k = layer.weight.shape
+        if n % 64 != 0 or k % 128 != 0:
+            return False
+        if self.use_mxfp8:
+            scale_u8 = layer.weight_scale_inv.data.view(torch.uint8).reshape(n, k // 32)
+        elif self.block_fp8_as_mxfp8 and layer.block_fp8_mxfp8_ready:
+            scale_u8 = block_fp8_scale_to_mxfp8_e8m0(
+                layer.weight_scale_inv.data, (n, k), self.weight_block_size
+            )
+        else:
+            return False
+        copy_or_rebind_param(
+            layer,
+            "weight_scale_inv_deepgemm",
+            pack_mxfp8_weight_scale_for_deepgemm(scale_u8, n, k),
+        )
+        return True
+
     def _process_mxfp8_linear_weight_scale(
         self, layer: Module, scale_u8: Optional[torch.Tensor] = None
     ) -> None:
@@ -1054,22 +1088,15 @@ def _process_mxfp8_linear_weight_scale(
                     "weight_scale_inv_swizzled",
                     block_scale_interleave(scale_u8.contiguous()).contiguous(),
                 )
-            scale_fp32 = (
-                (scale_u8.contiguous().view(-1).to(torch.int32) << 23)
-                .view(torch.float32)
-                .view(n, k // 32)
-            )
             if DEEPGEMM_SCALE_UE8M0:
                 # Pre-packed; GEMM must be called with disable_ue8m0_cast=True.
-                import deep_gemm.utils.layout
-
+                scale_packed = pack_mxfp8_weight_scale_for_deepgemm(scale_u8, n, k)
+            else:
                 scale_packed = (
-                    deep_gemm.utils.layout.get_mn_major_tma_aligned_packed_ue8m0_tensor(
-                        scale_fp32
-                    )
+                    (scale_u8.contiguous().view(-1).to(torch.int32) << 23)
+                    .view(torch.float32)
+                    .view(n, k // 32)
                 )
-            else:
-                scale_packed = scale_fp32
             copy_or_rebind_param(layer, "weight_scale_inv_deepgemm", scale_packed)
 
     def _quantize_mxfp8_weights(self, layer: Module) -> None:
@@ -1093,6 +1120,8 @@ def _quantize_mxfp8_weights(self, layer: Module) -> None:
     def process_weights_after_loading(self, layer: Module) -> None:
         if self.block_quant:
             self.process_weights_after_loading_block_quant(layer)
+            if layer.deepgemm_mxfp8_input:
+                layer.deepgemm_mxfp8_input = self.prepare_deepgemm_mxfp8_input(layer)
         else:
             layer.weight = Parameter(layer.weight.data, requires_grad=False)
 
@@ -1285,6 +1314,20 @@ def apply(
                     "Mxfp8SwizzledInput needs a layer with an MXFP8 view on a "
                     "FlashInfer CUTLASS / CuTe-DSL backend"
                 )
+        elif (
+            mxfp8_view
+            and isinstance(x, tuple)
+            and x[1].dtype == torch.int32
+            and not self.mxfp8_dense_backend.is_deep_gemm()
+        ):
+            # DeepGEMM-packed activation scales; see LinearBase.deepgemm_mxfp8_input.
+            return deepgemm_w8a8_mxfp8_linear_with_fallback(
+                x[0],
+                layer.weight,
+                layer.weight_scale_inv_deepgemm,
+                input_scale=x[1],
+                bias=bias,
+            )
         elif self.block_fp8_as_mxfp8 and isinstance(x, tuple):
             # A legacy (q, scale) block-fp8 pair keeps the block kernel.
             mxfp8_view = False

python/sglang/srt/layers/quantization/fp8_utils.py

@@ -979,10 +979,26 @@ def block_fp8_scale_to_mxfp8_e8m0(
     return e8m0.repeat_interleave(block_n, dim=0)[:n].contiguous()
 
 
+def pack_mxfp8_weight_scale_for_deepgemm(
+    scale_u8: torch.Tensor, n: int, k: int
+) -> torch.Tensor:
+    """UE8M0 bytes [N, K // 32] -> DeepGEMM's MN-major TMA-aligned packed scale."""
+    import deep_gemm.utils.layout
+
+    scale_fp32 = (
+        (scale_u8.contiguous().view(-1).to(torch.int32) << 23)
+        .view(torch.float32)
+        .view(n, k // 32)
+    )
+    return deep_gemm.utils.layout.get_mn_major_tma_aligned_packed_ue8m0_tensor(
+        scale_fp32
+    )
+
+
 def dispatch_w8a8_mxfp8_linear() -> Callable:
     backend = resolve_mxfp8_dense_gemm_backend()
     if backend.is_deep_gemm():
-        return _deepgemm_w8a8_mxfp8_linear_with_fallback
+        return deepgemm_w8a8_mxfp8_linear_with_fallback
     elif backend.is_flashinfer_trtllm():
         return partial(flashinfer_mxfp8_blockscaled_linear, backend="trtllm")
     elif backend.is_flashinfer_cutlass():
@@ -999,7 +1015,7 @@ def dispatch_w8a8_mxfp8_linear() -> Callable:
     return dot_scaled_mxfp8_blockscaled_linear
 
 
-def _deepgemm_w8a8_mxfp8_linear_with_fallback(
+def deepgemm_w8a8_mxfp8_linear_with_fallback(
     input: torch.Tensor,
     weight: torch.Tensor,
     weight_scale: torch.Tensor,

python/sglang/srt/models/deepseek_v4.py

@@ -665,17 +665,20 @@ def _fp8_wo_a_einsum(
     weight: torch.Tensor,
     weight_scale: torch.Tensor,
     recipe: Tuple[int, int, int],
-) -> torch.Tensor:
+    out: Optional[torch.Tensor | Tuple[torch.Tensor, torch.Tensor]] = None,
+) -> torch.Tensor | Tuple[torch.Tensor, torch.Tensor]:
+    """``out`` is a bf16 tensor or an (fp8, DeepGEMM-packed scale) pair."""
     import deep_gemm
 
-    groups, rank = weight.shape[:2]
-    output = torch.empty(
-        o_fp8.shape[0], groups, rank, device=o_fp8.device, dtype=torch.bfloat16
-    )
+    if out is None:
+        groups, rank = weight.shape[:2]
+        out = torch.empty(
+            o_fp8.shape[0], groups, rank, device=o_fp8.device, dtype=torch.bfloat16
+        )
     deep_gemm.fp8_einsum(
-        "bhr,hdr->bhd", (o_fp8, o_s), (weight, weight_scale), output, recipe=recipe
+        "bhr,hdr->bhd", (o_fp8, o_s), (weight, weight_scale), out, recipe=recipe
     )
-    return output
+    return out
 
 
 def _apply_gguf_grouped_wo_a(
@@ -1395,6 +1398,12 @@ def __init__(
             self.mega_cos_sin_cache = precompute_cos_sin_table(**self.rope_kwargs)
         self.mega_wo_a = None
         self.mega_wq_b_permuted = False
+        if self.use_mega_attention:
+            self.wo_b.deepgemm_mxfp8_input = (
+                self.n_local_groups * self.o_lora_rank == self.wo_b.weight.shape[-1]
+            )
+            if self.fuse_wqa_wkv:
+                self.wqkv_a.deepgemm_mxfp8_input = True
 
         self.use_fused_qk_norm_rope = _is_hip and _hip.use_fused_qk_norm_rope(self)
         # Static eligibility; token count and wo_b's output format are checked
@@ -1461,6 +1470,9 @@ def setup_mega_attention(self) -> None:
         wo_a = self.wo_a.weight.view(self.n_local_groups, self.o_lora_rank, -1)
         self.mega_wo_a = quantize_wo_a(interleave_output_columns(wo_a))
 
+    def accepts_deepgemm_mxfp8_input(self) -> bool:
+        return self.fuse_wqa_wkv and self.wqkv_a.deepgemm_mxfp8_input
+
     def _wq_b_uses_flashinfer_mxfp8(self) -> bool:
         method = self.wq_b.quant_method
         backend = getattr(method, "mxfp8_dense_backend", None)
@@ -1491,15 +1503,18 @@ def _forward_mega(
             o = q.new_empty((q.shape[0], self.n_local_groups, self.o_lora_rank))
             bcg_deepseek_v4_mega_attention_with_output(self, q, o, attn_sink)
             return o
-        return self._mega_core(q, attn_sink, forward_batch, attn_backend)
+        return self._mega_core(
+            q, attn_sink, forward_batch, attn_backend, allow_mxfp8_output=True
+        )
 
     def _mega_core(
         self,
         q: torch.Tensor,
         attn_sink: torch.Tensor,
         forward_batch: ForwardBatch,
         attn_backend,
-    ) -> torch.Tensor:
+        allow_mxfp8_output: bool = False,
+    ) -> torch.Tensor | Tuple[torch.Tensor, torch.Tensor]:
         data, scale = attn_backend.forward(
             q=q,
             k=q,
@@ -1512,9 +1527,25 @@ def _mega_core(
             cos_sin_cache=self.mega_cos_sin_cache,
         )
         groups = self.n_local_groups
-        return _fp8_wo_a_einsum(
-            data[:, :groups], scale[:, :groups], *self.mega_wo_a, recipe=(1, 1, 32)
-        )
+        rows = data.shape[0]
+        o_fp8, o_s = data[:, :groups], scale[:, :groups]
+        if (
+            allow_mxfp8_output
+            and self.wo_b.deepgemm_mxfp8_input
+            and mhc.fused_mxfp8_rows(rows)
+        ):
+            y, y_sf = mhc.empty_deepgemm_mxfp8_pair(
+                rows, groups * self.o_lora_rank, data.device
+            )
+            _fp8_wo_a_einsum(
+                o_fp8,
+                o_s,
+                *self.mega_wo_a,
+                recipe=(1, 1, 32),
+                out=(y.view(rows, groups, -1), y_sf),
+            )
+            return y, y_sf
+        return _fp8_wo_a_einsum(o_fp8, o_s, *self.mega_wo_a, recipe=(1, 1, 32))
 
     def _get_npu_rope_position_cache(
         self,

python/sglang/srt/models/deepseek_v4_boundary.py

@@ -186,6 +186,8 @@ def next_boundary(self):
             accepts_mxfp8=self.attention
             and self.layer.self_attn.accepts_mxfp8_swizzled_input(),
             hc=hc,
+            accepts_mxfp8_deepgemm=self.attention
+            and self.layer.self_attn.accepts_deepgemm_mxfp8_input(),
         )
 
     def _after_post(self, residual, norm):