[ROCm] Add UE8M0 scale packing for Triton silu_mul_quant (#37835)

Signed-off-by: Andreas Karatzas <akaratza@amd.com>

tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py

@@ -239,12 +239,8 @@ def test_silu_mul_fp8_quant_deep_gemm(E: int, T: int, H: int, fp8_type: torch.dt
     scale_fmts = [
         DeepGemmQuantScaleFMT.FLOAT32,
         DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0,
+        DeepGemmQuantScaleFMT.UE8M0,
     ]
-    # UE8M0 (int32 packed) scales require the C++ kernel which is
-    # not available on ROCm (#ifndef USE_ROCM).
-    # https://github.com/ROCm/aiter/issues/2420
-    if current_platform.is_cuda():
-        scale_fmts.append(DeepGemmQuantScaleFMT.UE8M0)
 
     # Run the SiLU V2 kernel
     for scale_fmt in scale_fmts:

vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py

@@ -220,6 +220,22 @@ def persistent_masked_m_silu_mul_quant(
         # Triton fallback for ROCm and XPU -- the C++ kernel is guarded by
         # #ifndef USE_ROCM in activation_kernels.cu.
         # https://github.com/ROCm/aiter/issues/2420
+        # For UE8M0 (int32 packed scales), compute with float32 scales
+        # first, then pack afterwards.
+        is_packed_ue8m0 = quant_scale_fmt == DeepGemmQuantScaleFMT.UE8M0
+        if is_packed_ue8m0:
+            f32_shape, f32_strides, _ = scales_shape_stride_dtype(
+                E, T, G, DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0
+            )
+            y_s_f32 = torch.empty_strided(
+                f32_shape,
+                f32_strides,
+                dtype=torch.float32,
+                device=y.device,
+            )
+        else:
+            y_s_f32 = y_s
+
         stride_cnt_e = tokens_per_expert.stride()[0]
 
         # Static grid over experts and H-groups.
@@ -231,14 +247,16 @@ def persistent_masked_m_silu_mul_quant(
 
         fp8_min, fp8_max = get_fp8_min_max()
         eps: float = 1e-10
-        assert y_s.dtype == torch.float32, (
-            "_silu_mul_fp8_quant_deep_gemm Triton fallback does not "
-            f"support {y_s.dtype} scales. Only torch.float32 supported."
-        )
+        if not is_packed_ue8m0:
+            assert y_s.dtype == torch.float32, (
+                "_silu_mul_fp8_quant_deep_gemm Triton fallback does not "
+                f"support {y_s.dtype} scales. Only torch.float32 supported."
+            )
+        f32_strides = y_s_f32.stride()
         _silu_mul_fp8_quant_deep_gemm[grid](
             y,
             y_q,
-            y_s,
+            y_s_f32,
             tokens_per_expert,
             H,
             group_size,
@@ -248,9 +266,9 @@ def persistent_masked_m_silu_mul_quant(
             stride_yq_e,
             stride_yq_t,
             stride_yq_h,
-            ys_strides[0],
-            ys_strides[1],
-            ys_strides[2],
+            f32_strides[0],
+            f32_strides[1],
+            f32_strides[2],
             stride_cnt_e,
             eps,
             fp8_min,
@@ -261,6 +279,23 @@ def persistent_masked_m_silu_mul_quant(
             num_warps=1,
         )
 
+        if is_packed_ue8m0:
+            # Pack float32 scales into int32 UE8M0 format:
+            # extract exponent bits (bits 30:23) from float32.
+            E_dim, T_dim, G_dim = y_s_f32.shape
+            y_s_cont = y_s_f32.contiguous()
+            i32_pad = round_up(G_dim, 4) - G_dim
+            y_s_u8 = (y_s_cont.view(torch.int32) >> 23).to(torch.uint8)
+            if i32_pad > 0:
+                y_s_u8 = torch.nn.functional.pad(y_s_u8, (0, i32_pad))
+            # y_s has shape (E, T, G//4) with stride (T*G//4, 1, T)
+            packed = y_s_u8.view(torch.int32)
+            # Copy with matching strides
+            for e_idx in range(E_dim):
+                nt = tokens_per_expert[e_idx].item()
+                if nt > 0:
+                    y_s[e_idx, :nt].copy_(packed[e_idx, :nt])
+
     return y_q, y_s