[Feature] Support DeepSeek-V4 Wint4Abf16 and Win4Afp8. (#25763)

Co-authored-by: zekai <zekai01@antgroup.com>

python/sglang/jit_kernel/csrc/gemm/per_tensor_quant_fp8.cuh

@@ -94,9 +94,14 @@ __global__ void per_tensor_quant_fp8_kernel(
   }
 }
 
-template <bool kIsStatic, typename DType>
+template <bool kIsStatic, bool kSkipQuant, typename DType>
-void per_tensor_quant_fp8(tvm::ffi::TensorView input, tvm::ffi::TensorView output_q, tvm::ffi::TensorView output_s) {
+void per_tensor_quant_fp8_impl(
+    tvm::ffi::TensorView input, tvm::ffi::TensorView output_q, tvm::ffi::TensorView output_s) {
   using namespace host;
+  static_assert(
+      !(kIsStatic && kSkipQuant),
+      "kIsStatic+kSkipQuant=no work. Use per_tensor_absmax_fp8 for absmax-only "
+      "or per_tensor_quant_fp8<is_static=true> for static-scale quant.");
 
   auto device = SymbolicDevice{};
   auto N = SymbolicSize{"num_elements"};
@@ -106,10 +111,12 @@ void per_tensor_quant_fp8(tvm::ffi::TensorView input, tvm::ffi::TensorView outpu
       .with_dtype<DType>()
       .with_device(device)
       .verify(input);
-  TensorMatcher({N})  //
+  if constexpr (!kSkipQuant) {
-      .with_dtype<fp8_e4m3_t>()
+    TensorMatcher({N})  //
-      .with_device(device)
+        .with_dtype<fp8_e4m3_t>()
-      .verify(output_q);
+        .with_device(device)
+        .verify(output_q);
+  }
   TensorMatcher({1})  //
       .with_dtype<float>()
       .with_device(device)
@@ -128,12 +135,28 @@ void per_tensor_quant_fp8(tvm::ffi::TensorView input, tvm::ffi::TensorView outpu
         static_cast<int64_t>(num_elements));
   }
 
-  LaunchKernel(num_blocks, kBlockSize, device.unwrap())(
+  if constexpr (!kSkipQuant) {
-      per_tensor_quant_fp8_kernel<DType, fp8_e4m3_t>,
+    LaunchKernel(num_blocks, kBlockSize, device.unwrap())(
-      static_cast<const DType*>(input.data_ptr()),
+        per_tensor_quant_fp8_kernel<DType, fp8_e4m3_t>,
-      static_cast<fp8_e4m3_t*>(output_q.data_ptr()),
+        static_cast<const DType*>(input.data_ptr()),
-      static_cast<const float*>(output_s.data_ptr()),
+        static_cast<fp8_e4m3_t*>(output_q.data_ptr()),
-      static_cast<int64_t>(num_elements));
+        static_cast<const float*>(output_s.data_ptr()),
+        static_cast<int64_t>(num_elements));
+  }
+}
+
+template <bool kIsStatic, typename DType>
+void per_tensor_quant_fp8(tvm::ffi::TensorView input, tvm::ffi::TensorView output_q, tvm::ffi::TensorView output_s) {
+  per_tensor_quant_fp8_impl<kIsStatic, /*kSkipQuant=*/false, DType>(input, output_q, output_s);
+}
+
+// output_q is statically unused when kSkipQuant=true; reuse output_s as a zero-cost placeholder.
+template <typename DType>
+void per_tensor_absmax_fp8(tvm::ffi::TensorView input, tvm::ffi::TensorView output_s) {
+  per_tensor_quant_fp8_impl<
+      /*kIsStatic=*/false,
+      /*kSkipQuant=*/true,
+      DType>(input, output_s, output_s);
 }
 
 }  // namespace

python/sglang/jit_kernel/per_tensor_quant_fp8.py

@@ -43,3 +43,35 @@ def per_tensor_quant_fp8(
     """
     module = _jit_per_tensor_quant_fp8_module(is_static, input.dtype)
     module.per_tensor_quant_fp8(input.view(-1), output_q.view(-1), output_s.view(-1))
+
+
+@cache_once
+def _jit_per_tensor_absmax_fp8_module(dtype: torch.dtype) -> Module:
+    args = make_cpp_args(dtype)
+    return load_jit(
+        "per_tensor_absmax_fp8",
+        *args,
+        cuda_files=["gemm/per_tensor_quant_fp8.cuh"],
+        cuda_wrappers=[("per_tensor_absmax_fp8", f"per_tensor_absmax_fp8<{args}>")],
+    )
+
+
+@register_custom_op(
+    op_name="per_tensor_absmax_fp8",
+    mutates_args=["output_s"],
+)
+def per_tensor_absmax_fp8(
+    input: torch.Tensor,
+    output_s: torch.Tensor,
+) -> None:
+    """Compute scale = max(abs(input)) / fp8_e4m3_max via atomic-max reduction.
+
+    The caller must zero-initialise ``output_s`` before the call (the kernel
+    uses ``atomic_max`` across blocks, so starting from 0 is required).
+
+    Args:
+        input: Input tensor (float16, bfloat16, or float32). Any shape.
+        output_s: Pre-allocated float32 tensor of shape (1,), zero-initialised.
+    """
+    module = _jit_per_tensor_absmax_fp8_module(input.dtype)
+    module.per_tensor_absmax_fp8(input.view(-1), output_s.view(-1))

python/sglang/kernels/ops/moe/ep_moe_kernels.py

@@ -699,6 +699,68 @@ def silu_and_mul_masked_fwd(
     return output
 
 
+@triton.jit
+def silu_mul_dynamic_scale_triton_kernel_for_cutlass_moe(
+    input_ptr,
+    scale_ptr,
+    num_tokens_tensor_ptr,
+    intermediate_size,
+    fp8_max,
+    BLOCK_SIZE: tl.constexpr,
+    NUM_STAGES: tl.constexpr,
+):
+    num_tokens = tl.load(num_tokens_tensor_ptr)
+    numel = num_tokens * intermediate_size
+    gate_ptr = input_ptr
+    up_ptr = input_ptr + intermediate_size
+
+    start_idx = tl.program_id(0) * BLOCK_SIZE
+    step = tl.num_programs(0) * BLOCK_SIZE
+    absmax = 0.0
+
+    for id in tl.range(start_idx, numel, step, num_stages=NUM_STAGES):
+        ids = id + tl.arange(0, BLOCK_SIZE)
+        token_ids = ids // intermediate_size
+        mask = ids < numel
+
+        offs = ids + token_ids * intermediate_size
+        gate = tl.load(gate_ptr + offs, mask=mask, other=0.0).to(tl.float32)
+        up = tl.load(up_ptr + offs, mask=mask, other=0.0).to(tl.float32)
+        gate_up = gate / (1 + tl.exp(-gate)) * up
+        absmax = tl.maximum(absmax, tl.max(tl.abs(gate_up)))
+
+    absmax = tl.maximum(absmax, 1e-10)
+    tl.atomic_max(scale_ptr, absmax / fp8_max)
+
+
+def silu_mul_dynamic_tensorwise_quant_for_cutlass_moe(
+    input: torch.Tensor,
+    output: torch.Tensor,
+    scale: torch.Tensor,
+    num_tokens_tensor: torch.Tensor,
+    expected_num_tokens: int,
+    intermediate_size: int,
+):
+    grid, block_dim = _get_launch_config_1d(
+        input.device, expected_num_tokens * intermediate_size
+    )
+    scale.zero_()
+    fp8_max = torch.finfo(torch.float8_e4m3fn).max
+
+    silu_mul_dynamic_scale_triton_kernel_for_cutlass_moe[grid](
+        input_ptr=input,
+        scale_ptr=scale,
+        num_tokens_tensor_ptr=num_tokens_tensor,
+        intermediate_size=intermediate_size,
+        fp8_max=fp8_max,
+        BLOCK_SIZE=block_dim,
+        NUM_STAGES=3,
+    )
+    silu_mul_static_tensorwise_quant_for_cutlass_moe(
+        input, output, scale, num_tokens_tensor, expected_num_tokens, intermediate_size
+    )
+
+
 @triton.jit
 def silu_mul_static_tensorwise_quant_triton_kernel_for_cutlass_moe(
     input_ptr,

python/sglang/srt/layers/moe/cutlass_w4a8_moe.py

@@ -22,7 +22,10 @@ if _is_cuda:
 else:
     from sgl_kernel import silu_and_mul
 
-from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8
+from sglang.jit_kernel.per_tensor_quant_fp8 import (
+    per_tensor_absmax_fp8,
+    per_tensor_quant_fp8,
+)
 from sglang.kernels.ops.moe.ep_moe_kernels import (
     cutlass_w4_run_moe_ep_preproess,
     deepep_ll_get_cutlass_w4a8_moe_mm_data,
@@ -33,6 +36,7 @@ from sglang.kernels.ops.moe.ep_moe_kernels import (
     post_reorder_for_cutlass_moe,
     pre_reorder_for_cutlass_moe,
     silu_and_mul_masked_post_per_tensor_quant_fwd,
+    silu_mul_dynamic_tensorwise_quant_for_cutlass_moe,
     silu_mul_static_tensorwise_quant_for_cutlass_moe,
 )
 
@@ -137,6 +141,11 @@ def cutlass_w4a8_moe(
         dtype=torch.float8_e4m3fn,
     )
 
+    # TODO: fuse per_tensor_absmax_fp8 and pre_reorder_for_cutlass_moe
+    if a1_scale is None:
+        a1_scale = torch.zeros(1, dtype=torch.float32, device=device)
+        per_tensor_absmax_fp8(a, a1_scale)
+
     pre_reorder_for_cutlass_moe(
         a,
         gateup_input,
@@ -188,9 +197,16 @@ def cutlass_w4a8_moe(
     intermediate_q = torch.empty(
         (m * topk, n), dtype=torch.float8_e4m3fn, device=device
     )
-    silu_mul_static_tensorwise_quant_for_cutlass_moe(
+
-        c1, intermediate_q, a2_scale.float(), expert_offsets[-1:], m * topk, n
+    if a2_scale is None:
-    )
+        a2_scale = torch.zeros(1, dtype=torch.float32, device=device)
+        silu_mul_dynamic_tensorwise_quant_for_cutlass_moe(
+            c1, intermediate_q, a2_scale, expert_offsets[-1:], m * topk, n
+        )
+    else:
+        silu_mul_static_tensorwise_quant_for_cutlass_moe(
+            c1, intermediate_q, a2_scale.float(), expert_offsets[-1:], m * topk, n
+        )
 
     cutlass_w4a8_moe_mm(
         c2,

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

@@ -44,6 +44,7 @@ from sglang.srt.layers.quantization.compressed_tensors.schemes import (
     CompressedTensorsMxInt4MoE,
     CompressedTensorsW4A4Fp4,
     CompressedTensorsW4A4Nvfp4MoE,
+    CompressedTensorsW4AFP8MoE,
     CompressedTensorsW8A8Fp8,
     CompressedTensorsW8A8Fp8MoE,
     CompressedTensorsW8A8Int8,
@@ -306,15 +307,16 @@ class CompressedTensorsConfig(QuantizationConfig):
                 target_scheme_map[target]["input_activations"] = None
                 if is_activation_quantization_format(quant_format):
                     input_activations = quant_config.get("input_activations")
-                    # The only case where we have activation quant supported
+                    # When activation quant format is set but no
-                    # but no input_activations provided in the config
+                    # input_activations provided: valid for w8a16fp8 (FLOAT
-                    # should be w8a16fp8 w8a16fp8 can also run for cases where
+                    # weights) and pack-quantized without activation quant
-                    # there is an input_quant but it is ignored
+                    # (INT weights, W4A16 style).
                     if not input_activations:
-                        assert (
+                        weight_type = target_scheme_map[target]["weights"].type
-                            target_scheme_map[target]["weights"].type
+                        if weight_type == QuantizationType.INT:
-                            == QuantizationType.FLOAT
+                            pass
-                        )
+                        else:
+                            assert weight_type == QuantizationType.FLOAT
                     else:
                         target_scheme_map[target]["input_activations"] = (
                             QuantizationArgs.model_validate(  # noqa: E501
@@ -365,6 +367,33 @@ class CompressedTensorsConfig(QuantizationConfig):
             and is_dynamic
         )
 
+    def _is_wint4afp8(self, weight_quant: BaseModel, input_quant: BaseModel) -> bool:
+        """Detect W4AFP8: packed INT4 weights + 8-bit dynamic per-token activations."""
+        if weight_quant is None or input_quant is None:
+            return False
+        return (
+            self.quant_format == CompressionFormat.pack_quantized.value
+            and weight_quant.num_bits == 4
+            and weight_quant.type == QuantizationType.INT
+            and weight_quant.symmetric
+            and not weight_quant.dynamic
+            and input_quant.num_bits == 8
+            and input_quant.type in [QuantizationType.FLOAT, QuantizationType.INT]
+            and input_quant.dynamic  # currently not support static input scales
+        )
+
+    def _is_wint4abf16(self, weight_quant: BaseModel, input_quant: BaseModel) -> bool:
+        """Detect W4A16: packed INT4 weights with no activation quantization (activations stay BF16)."""
+        if weight_quant is None or input_quant is not None:
+            return False
+        return (
+            self.quant_format == CompressionFormat.pack_quantized.value
+            and weight_quant.num_bits == 4
+            and weight_quant.type == QuantizationType.INT
+            and weight_quant.symmetric
+            and not weight_quant.dynamic
+        )
+
     def _is_static_tensor_w8a8(
         self, weight_quant: BaseModel, input_quant: BaseModel
     ) -> bool:
@@ -721,6 +750,13 @@ class CompressedTensorsConfig(QuantizationConfig):
                 raise NotImplementedError(
                     f"The W8A8Int8 Fused MoE scheme is implemented only for NPU for now."
                 )
+        elif self._is_wint4afp8(weight_quant, input_quant):
+            # On NPU prefer the dedicated NPU W4A8Int8 path when activations are INT8.
+            if _is_npu and self._is_dynamic_token_w4a8(weight_quant, input_quant):
+                logger.info_once("Using NPUCompressedTensorsW4A8Int8DynamicMoE")
+                return NPUCompressedTensorsW4A8Int8DynamicMoE(self)
+            logger.info_once("Using CompressedTensorsW4AFP8MoE")
+            return CompressedTensorsW4AFP8MoE(self, weight_quant, input_quant)
         elif self._is_dynamic_token_w4a8(weight_quant, input_quant):
             if _is_npu:
                 logger.info_once("Using NPUCompressedTensorsW4A8Int8DynamicMoE")

python/sglang/srt/layers/quantization/compressed_tensors/schemes/__init__.py

@@ -7,6 +7,7 @@ from .compressed_tensors_scheme import (
 from .compressed_tensors_w4a4_mxint4_moe import CompressedTensorsMxInt4MoE
 from .compressed_tensors_w4a4_nvfp4 import CompressedTensorsW4A4Fp4
 from .compressed_tensors_w4a4_nvfp4_moe import CompressedTensorsW4A4Nvfp4MoE
+from .compressed_tensors_w4a8_fp8_moe import CompressedTensorsW4AFP8MoE
 from .compressed_tensors_w4a8_int8_moe import NPUCompressedTensorsW4A8Int8DynamicMoE
 from .compressed_tensors_w8a8_fp8 import CompressedTensorsW8A8Fp8
 from .compressed_tensors_w8a8_fp8_moe import CompressedTensorsW8A8Fp8MoE
@@ -41,4 +42,5 @@ __all__ = [
     "CompressedTensorsW4A4Nvfp4MoE",
     "NPUCompressedTensorsW4A8Int8DynamicMoE",
     "CompressedTensorsMxInt4MoE",
+    "CompressedTensorsW4AFP8MoE",
 ]

python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_fp8_moe.py

@@ -0,0 +1,323 @@
+"""W4AFP8 MoE scheme: INT4 group-quantized weights + 8-bit dynamic activations.
+
+Loads INT4 weights from compressed-tensors pack-quantized format,
+converts to CUTLASS W4A8 layout, and runs CUTLASS grouped GEMM
+with dynamic 8-bit (FP8 or INT8) activation quantization.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+import torch
+from compressed_tensors import CompressionFormat
+
+from sglang.srt.layers.moe import MoeRunnerConfig
+from sglang.srt.layers.quantization.compressed_tensors.schemes import (
+    CompressedTensorsMoEScheme,
+)
+from sglang.srt.layers.quantization.w4afp8 import interleave_scales
+from sglang.srt.utils import set_weight_attrs
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.moe.token_dispatcher import (
+        CombineInput,
+        StandardDispatchOutput,
+    )
+    from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors import (
+        CompressedTensorsConfig,
+    )
+
+logger = logging.getLogger(__name__)
+
+__all__ = ["CompressedTensorsW4AFP8MoE"]
+
+
+def _unpack_repack_int32_to_cutlass_int8(
+    weight_packed: torch.Tensor, num_bits: int
+) -> torch.Tensor:
+    """Convert compressed-tensors pack_to_int32 format to CUTLASS int8-packed format.
+
+    pack_to_int32 stores 8 unsigned-offset int4 values per int32.
+    CUTLASS expects pairs of signed int4 values packed into int8
+    (low nibble = even index, high nibble = odd index, two's complement).
+
+    Args:
+        weight_packed: [E, N, K // pack_factor] int32  (pack_factor = 32 // num_bits)
+        num_bits: quantization bit width (e.g. 4)
+
+    Returns:
+        [E, N, K // 2] int8 in CUTLASS layout
+    """
+    pack_factor = 32 // num_bits
+    mask = (1 << num_bits) - 1
+    offset = 1 << (num_bits - 1)
+    pair_factor = pack_factor // 2
+
+    # Repack directly into CUTLASS int8 without materializing full unpacked int32.
+    # This reduces peak memory from O(pack_factor) large int32 buffers to O(1) temporaries.
+    out = torch.empty(
+        (*weight_packed.shape[:-1], weight_packed.shape[-1], pair_factor),
+        dtype=torch.int8,
+        device=weight_packed.device,
+    )
+    for pair_idx in range(pair_factor):
+        low_shift = num_bits * (2 * pair_idx)
+        high_shift = low_shift + num_bits
+
+        low_nibbles = ((weight_packed >> low_shift) & mask) - offset
+        high_nibbles = ((weight_packed >> high_shift) & mask) - offset
+        out[..., pair_idx] = ((high_nibbles << 4) | (low_nibbles & 0x0F)).to(torch.int8)
+
+    return out.flatten(-2).contiguous()
+
+
+class CompressedTensorsW4AFP8MoE(CompressedTensorsMoEScheme):
+    """W4AFP8 MoE: INT4 weights (pack-quantized) + dynamic per-token 8-bit activations,
+    using CUTLASS W4A8 grouped GEMM kernel."""
+
+    def __init__(
+        self,
+        quant_config: CompressedTensorsConfig,
+        weight_quant,
+        input_quant,
+    ):
+        self.quant_config = quant_config
+        config = self.quant_config.target_scheme_map["Linear"].get("weights")
+        self.num_bits = config.num_bits
+        self.packed_factor = 32 // config.num_bits
+        self.group_size = config.group_size
+        self.weight_quant = weight_quant
+        self.input_quant = input_quant
+
+        assert config.symmetric, "Only symmetric quantization is supported"
+        assert (
+            self.quant_config.quant_format == CompressionFormat.pack_quantized.value
+        ), f"W4AFP8MoE requires pack-quantized format, got {self.quant_config.quant_format}"
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 90
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        num_experts: int,
+        hidden_size: int,
+        intermediate_size_per_partition: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+
+        from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported
+
+        # Weights in checkpoint (non-transposed) layout: [E, N, K // pack_factor]
+        # This matches the pack-quantized checkpoint format directly.
+        w13_weight = torch.nn.Parameter(
+            torch.empty(
+                num_experts,
+                2 * intermediate_size_per_partition,
+                hidden_size // self.packed_factor,
+                dtype=torch.int32,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w13_weight_packed", w13_weight)
+        set_weight_attrs(w13_weight, extra_weight_attrs)
+
+        w2_weight = torch.nn.Parameter(
+            torch.empty(
+                num_experts,
+                hidden_size,
+                intermediate_size_per_partition // self.packed_factor,
+                dtype=torch.int32,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w2_weight_packed", w2_weight)
+        set_weight_attrs(w2_weight, extra_weight_attrs)
+
+        # Scales: [E, N, K // group_size]
+        num_groups_w13 = hidden_size // self.group_size
+        num_groups_w2 = intermediate_size_per_partition // self.group_size
+
+        extra_weight_attrs.update(
+            {"quant_method": FusedMoeWeightScaleSupported.GROUP.value}
+        )
+
+        w13_scale = torch.nn.Parameter(
+            torch.zeros(
+                num_experts,
+                2 * intermediate_size_per_partition,
+                num_groups_w13,
+                dtype=torch.float32,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w13_weight_scale", w13_scale)
+        set_weight_attrs(w13_scale, extra_weight_attrs)
+
+        w2_scale = torch.nn.Parameter(
+            torch.zeros(
+                num_experts,
+                hidden_size,
+                num_groups_w2,
+                dtype=torch.float32,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w2_weight_scale", w2_scale)
+        set_weight_attrs(w2_scale, extra_weight_attrs)
+
+        # Placeholder params to accept checkpoint tensors that we don't use.
+        # Without these, the weight loader warns "not found in params_dict".
+        for name, shape in [
+            ("w13_weight_shape", (num_experts, 2)),
+            ("w2_weight_shape", (num_experts, 2)),
+        ]:
+            p = torch.nn.Parameter(
+                torch.empty(shape, dtype=torch.int32), requires_grad=False
+            )
+            layer.register_parameter(name, p)
+            set_weight_attrs(p, extra_weight_attrs)
+
+        self._init_cutlass_buffers(
+            num_experts,
+            hidden_size,
+            intermediate_size_per_partition,
+            layer.w13_weight_packed.device,
+        )
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        """Convert pack-quantized INT32 weights to CUTLASS INT8-packed format.
+
+        Unpacks weights from compressed-tensors pack_to_int32 format and repacks
+        into CUTLASS int8 layout, then interleaves scales.
+          INT32 [E, N, K//8] → unpack signed int4 → repack INT8 [E, N, K//2]
+          Scales [E, N, K//gs] → interleaved bfloat16
+        """
+        if getattr(layer, "is_w4afp8_converted", False):
+            return
+
+        dtype = torch.bfloat16
+        device = layer.w2_weight_packed.device
+
+        # TODO: currently only support per tensor quant.
+        layer.a13_scale = None
+        layer.a2_scale = None
+
+        w13 = _unpack_repack_int32_to_cutlass_int8(
+            layer.w13_weight_packed.data, self.num_bits
+        )
+        layer.w13_weight_packed = torch.nn.Parameter(w13, requires_grad=False)
+
+        w2 = _unpack_repack_int32_to_cutlass_int8(
+            layer.w2_weight_packed.data, self.num_bits
+        )
+        layer.w2_weight_packed = torch.nn.Parameter(w2, requires_grad=False)
+
+        w13_weight_scale = layer.w13_weight_scale.to(dtype)
+        w13_weight_scale = interleave_scales(w13_weight_scale)
+        layer.w13_weight_scale = torch.nn.Parameter(
+            w13_weight_scale, requires_grad=False
+        )
+
+        w2_weight_scale = layer.w2_weight_scale.to(dtype)
+        w2_weight_scale = interleave_scales(w2_weight_scale)
+        layer.w2_weight_scale = torch.nn.Parameter(w2_weight_scale, requires_grad=False)
+
+        layer.is_w4afp8_converted = True
+
+    def create_moe_runner(
+        self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
+    ):
+        self.moe_runner_config = moe_runner_config
+
+    def _init_cutlass_buffers(
+        self,
+        num_experts: int,
+        hidden_size: int,
+        intermediate_size: int,
+        device: torch.device,
+    ):
+        """Pre-allocate stride and workspace tensors for CUTLASS grouped GEMM."""
+        self.a_strides1 = torch.full(
+            (num_experts, 3), hidden_size, device=device, dtype=torch.int64
+        )
+        self.c_strides1 = torch.full(
+            (num_experts, 3),
+            2 * intermediate_size,
+            device=device,
+            dtype=torch.int64,
+        )
+        self.a_strides2 = torch.full(
+            (num_experts, 3), intermediate_size, device=device, dtype=torch.int64
+        )
+        self.c_strides2 = torch.full(
+            (num_experts, 3), hidden_size, device=device, dtype=torch.int64
+        )
+        self.b_strides1 = self.a_strides1
+        self.s_strides13 = self.c_strides1
+        self.b_strides2 = self.a_strides2
+        self.s_strides2 = self.c_strides2
+
+        self.expert_offsets = torch.empty(
+            num_experts + 1, dtype=torch.int32, device=device
+        )
+        self.problem_sizes1 = torch.empty(
+            (num_experts, 3), dtype=torch.int32, device=device
+        )
+        self.problem_sizes2 = torch.empty(
+            (num_experts, 3), dtype=torch.int32, device=device
+        )
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        dispatch_output: StandardDispatchOutput,
+    ) -> CombineInput:
+        """Run the forward pass for W4AFP8 MoE."""
+        return self.apply_weights(layer, dispatch_output)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        dispatch_output: StandardDispatchOutput,
+    ) -> CombineInput:
+        from sglang.srt.layers.moe.cutlass_w4a8_moe import cutlass_w4a8_moe
+        from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
+
+        assert (
+            self.moe_runner_config.activation == "silu"
+        ), "Only SiLU activation is supported."
+
+        x = dispatch_output.hidden_states
+        topk_output = dispatch_output.topk_output
+        topk_weights, topk_ids, _ = topk_output
+
+        # TODO: currently, group_size is hardcoded to 128 in the cutlass_w4a8_moe kernel.
+        output = cutlass_w4a8_moe(
+            x,
+            layer.w13_weight_packed,
+            layer.w2_weight_packed,
+            layer.w13_weight_scale,
+            layer.w2_weight_scale,
+            topk_weights,
+            topk_ids,
+            self.a_strides1,
+            self.b_strides1,
+            self.c_strides1,
+            self.a_strides2,
+            self.b_strides2,
+            self.c_strides2,
+            self.s_strides13,
+            self.s_strides2,
+            self.expert_offsets,
+            self.problem_sizes1,
+            self.problem_sizes2,
+            layer.a13_scale,
+            layer.a2_scale,
+            routed_scaling_factor=self.moe_runner_config.routed_scaling_factor or 1.0,
+        )
+        return StandardCombineInput(hidden_states=output)

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

@@ -16,6 +16,7 @@ def is_activation_quantization_format(format: str) -> bool:
         CompressionFormat.int_quantized.value,
         CompressionFormat.float_quantized.value,
         CompressionFormat.nvfp4_pack_quantized.value,
+        CompressionFormat.pack_quantized.value,
     ]
     return format in _ACTIVATION_QUANTIZATION_FORMATS
 

python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py

@@ -58,6 +58,7 @@ from sglang.srt.models.deepseek_common.utils import (
     _use_aiter_gfx95,
     awq_dequantize_func,
     enable_nextn_moe_bf16_cast_to_fp8,
+    is_wint4afp8_or_wint4a16_config,
 )
 from sglang.srt.utils import bind_or_assign, get_bool_env_var, log_info_on_rank0
 
@@ -182,7 +183,7 @@ class DeepseekV2WeightLoaderMixin:
         # Params for special naming rules in mixed-precision models, for example:
         # model.layers.xx.mlp.experts.xx.w1.input_scale. For details,
         # see https://huggingface.co/Barrrrry/DeepSeek-R1-W4AFP8/blob/main.
-        if self.quant_config and self.quant_config.get_name() == "w4afp8":
+        if is_wint4afp8_or_wint4a16_config(self.quant_config):
             expert_params_mapping += FusedMoE.make_expert_input_scale_params_mapping(
                 num_experts=self.config.n_routed_experts
             )
@@ -203,6 +204,7 @@ class DeepseekV2WeightLoaderMixin:
             futures = []
             params_dict = dict(self.named_parameters())
             weight_names = []
+
             for name, loaded_weight in weights:
                 use_async_loading = should_async_load(loaded_weight)
                 layer_id = get_layer_id(name)

python/sglang/srt/models/deepseek_common/utils.py

@@ -112,6 +112,28 @@ def enable_nextn_moe_bf16_cast_to_fp8(
     )
 
 
+def is_wint4afp8_or_wint4a16_config(
+    quant_config: Optional[QuantizationConfig],
+) -> bool:
+    if quant_config is None:
+        return False
+    if quant_config.get_name() == "w4afp8":
+        return True
+
+    from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors import (
+        CompressedTensorsConfig,
+    )
+
+    if not isinstance(quant_config, CompressedTensorsConfig):
+        return False
+    linear_scheme = quant_config.target_scheme_map.get("Linear", {})
+    weight_quant = linear_scheme.get("weights")
+    input_quant = linear_scheme.get("input_activations")
+    return quant_config._is_wint4afp8(
+        weight_quant, input_quant
+    ) or quant_config._is_wint4abf16(weight_quant, input_quant)
+
+
 def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float:
     if scale <= 1:
         return 1.0

python/sglang/srt/models/deepseek_v2.py

@@ -179,6 +179,7 @@ from sglang.srt.models.deepseek_common.utils import (
     _use_aiter,
     _use_aiter_bpreshuffle_gfx95,
     _use_aiter_gfx95,
+    is_wint4afp8_or_wint4a16_config,
 )
 from sglang.srt.runtime_context import (
     get_flags,
@@ -2808,8 +2809,8 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
                 "Only Deepseek V3/R1 on AMD-platform with capability >= gfx942(MI30x) "
                 "can use shared experts fusion optimization under expert parallelism."
             )
-        elif self.quant_config and self.quant_config.get_name() == "w4afp8":
+        elif is_wint4afp8_or_wint4a16_config(self.quant_config):
-            disable_reason = "Deepseek V3/R1 W4AFP8 model uses different quant method for routed experts and shared experts."
+            disable_reason = "Deepseek V3/R1 W4AFP8/W4A16 model uses different quant method for routed experts and shared experts."
 
         if disable_reason is not None:
             from sglang.srt.arg_groups.overrides import declare_load_time_override

python/sglang/srt/models/deepseek_v4.py

@@ -125,7 +125,10 @@ from sglang.srt.models.dbrx import ReplicatedLinear
 from sglang.srt.models.deepseek_common.amd.deepseek_v4_fused_mhc import (
     try_fused_hc_post_pre,
 )
-from sglang.srt.models.deepseek_common.utils import _use_aiter_bpreshuffle_gfx95
+from sglang.srt.models.deepseek_common.utils import (
+    _use_aiter_bpreshuffle_gfx95,
+    is_wint4afp8_or_wint4a16_config,
+)
 from sglang.srt.models.deepseek_v2 import (
     ParallelLMHead,
     _is_cuda,
@@ -2743,7 +2746,7 @@ class DeepseekV4ForCausalLM(nn.Module):
             num_experts=self.config.n_routed_experts + self.num_fused_shared_experts,
         )
 
-        if self.quant_config and self.quant_config.get_name() == "w4afp8":
+        if is_wint4afp8_or_wint4a16_config(self.quant_config):
             expert_params_mapping += FusedMoE.make_expert_input_scale_params_mapping(
                 num_experts=self.config.n_routed_experts
             )
@@ -2879,16 +2882,18 @@ class DeepseekV4ForCausalLM(nn.Module):
                         loaded_params.add(name)
                         break
                     else:
+                        skip_unmaterialized_expert_param = False
                         for mapping in expert_params_mapping:
                             param_name, weight_name, expert_id, shard_id = mapping
                             if weight_name not in name:
                                 continue
                             if _is_npu:
                                 name = name.replace("weight_packed", "weight")
-                            name = name.replace(weight_name, param_name)
+                            resolved_name = name.replace(weight_name, param_name)
-                            if name not in params_dict:
+                            if resolved_name not in params_dict:
+                                skip_unmaterialized_expert_param = True
                                 continue
-                            param = params_dict[name]
+                            param = params_dict[resolved_name]
                             weight_loader = param.weight_loader
                             maybe_executor_submit(
                                 executor=executor,
@@ -2898,16 +2903,18 @@ class DeepseekV4ForCausalLM(nn.Module):
                                 func_args=(
                                     param,
                                     loaded_weight,
-                                    name,
+                                    resolved_name,
                                 ),
                                 func_kwargs={
                                     "shard_id": shard_id,
                                     "expert_id": expert_id,
                                 },
                             )
-                            loaded_params.add(name)
+                            loaded_params.add(resolved_name)
                             break
                         else:
+                            if skip_unmaterialized_expert_param:
+                                continue
                             if name.endswith(".bias") and name not in params_dict:
                                 continue
                             if (