[AMD] ROCm: route W4A16 MoE to Triton and fix packed-weight loading (#17863)

python/sglang/srt/layers/moe/fused_moe_triton/layer.py

@@ -679,6 +679,7 @@ class FusedMoE(torch.nn.Module):
                 in [
                     "CompressedTensorsWNA16MarlinMoEMethod",
                     "CompressedTensorsWNA16MoEMethod",
+                    "CompressedTensorsWNA16TritonMoEMethod",
                 ]
             )
             else loaded_weight
@@ -892,7 +893,10 @@ class FusedMoE(torch.nn.Module):
             loaded_weight.t().contiguous()
             if (
                 self.quant_method.__class__.__name__
-                == "CompressedTensorsWNA16MoEMethod"
+                in [
+                    "CompressedTensorsWNA16MoEMethod",
+                    "CompressedTensorsWNA16TritonMoEMethod",
+                ]
             )
             else loaded_weight
         )

python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py

@@ -144,6 +144,11 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
                         "Using CompressedTensorsMxInt4MoEMethod with flashinfer_trtllm backend"
                     )
                     return CompressedTensorsMxInt4MoEMethod(quant_config)
+                elif _is_hip:
+                    logger.info_once(
+                        "Using CompressedTensorsWNA16TritonMoEMethod (ROCm)"
+                    )
+                    return CompressedTensorsWNA16TritonMoEMethod(quant_config)
                 else:
                     logger.info_once("Using CompressedTensorsWNA16MarlinMoEMethod")
                     return CompressedTensorsWNA16MoEMethod(quant_config)
@@ -1379,6 +1384,70 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
         return StandardCombineInput(hidden_states=output)
 
 
+class CompressedTensorsWNA16TritonMoEMethod(CompressedTensorsWNA16MoEMethod):
+    """ROCm/HIP-compatible W4A16 MoE method using Triton kernels instead of Marlin.
+
+    Inherits weight creation from CompressedTensorsWNA16MoEMethod but converts
+    weights to the uint8-packed format expected by the Triton fused MoE kernel
+    instead of the Marlin-specific format.
+    """
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        if getattr(layer, "is_triton_converted", False):
+            return
+
+        num_experts = layer.w13_weight_packed.shape[0]
+
+        # Convert w13 weights: [E, K//8, N] int32 -> [E, N, K//2] uint8
+        w13 = layer.w13_weight_packed.data
+        w13 = w13.transpose(1, 2).contiguous().view(torch.uint8)
+        layer.w13_weight_packed = torch.nn.Parameter(w13, requires_grad=False)
+
+        # Convert w2 weights: [E, K//8, N] int32 -> [E, N, K//2] uint8
+        w2 = layer.w2_weight_packed.data
+        w2 = w2.transpose(1, 2).contiguous().view(torch.uint8)
+        layer.w2_weight_packed = torch.nn.Parameter(w2, requires_grad=False)
+
+        # Convert w13 scales: [E, K//group_size, N] -> [E, N, K//group_size]
+        w13_scale = layer.w13_weight_scale.data
+        w13_scale = w13_scale.transpose(1, 2).contiguous()
+        layer.w13_weight_scale = torch.nn.Parameter(w13_scale, requires_grad=False)
+
+        # Convert w2 scales: [E, K//group_size, N] -> [E, N, K//group_size]
+        w2_scale = layer.w2_weight_scale.data
+        w2_scale = w2_scale.transpose(1, 2).contiguous()
+        layer.w2_weight_scale = torch.nn.Parameter(w2_scale, requires_grad=False)
+
+        layer.is_triton_converted = True
+
+    def create_moe_runner(
+        self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
+    ):
+        self.moe_runner_config = moe_runner_config
+        self.runner = MoeRunner(MoeRunnerBackend.TRITON, moe_runner_config)
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        dispatch_output: "StandardDispatchOutput",
+    ) -> "CombineInput":
+        from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo
+
+        assert (
+            self.moe_runner_config.activation == "silu"
+        ), "Only SiLU activation is supported."
+
+        quant_info = TritonMoeQuantInfo(
+            w13_weight=layer.w13_weight_packed,
+            w2_weight=layer.w2_weight_packed,
+            use_int4_w4a16=True,
+            w13_scale=layer.w13_weight_scale,
+            w2_scale=layer.w2_weight_scale,
+            block_shape=[0, self.group_size],
+        )
+        return self.runner.run(dispatch_output, quant_info)
+
+
 class NPUCompressedTensorsW4A8Int8DynamicMoEMethod(CompressedTensorsMoEMethod):
 
     ### TODO: Get rid of code duplication with python/sglang/srt/modelslim/modelslim_moe.py @OrangeRedeng @TamirBaydasov