minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: zekai <zekai01@antgroup.com>
This commit is contained in:
yuyu5333
2026-07-16 14:24:34 -07:00
committed by GitHub
co-authored by zekai
parent 68f4de162d
commit bff489284b
12 changed files with 561 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/csrc/gemm/per_tensor_quant_fp8.cuh
+35 -12
View File
@@ -94,9 +94,14 @@ __global__ void per_tensor_quant_fp8_kernel(
}
}
template <bool kIsStatic, typename DType>
void per_tensor_quant_fp8(tvm::ffi::TensorView input, tvm::ffi::TensorView output_q, tvm::ffi::TensorView output_s) {
template <bool kIsStatic, bool kSkipQuant, typename DType>
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}) //
.with_dtype<fp8_e4m3_t>()
.with_device(device)
.verify(output_q);
if constexpr (!kSkipQuant) {
TensorMatcher({N}) //
.with_dtype<fp8_e4m3_t>()
.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())(
per_tensor_quant_fp8_kernel<DType, fp8_e4m3_t>,
static_cast<const DType*>(input.data_ptr()),
static_cast<fp8_e4m3_t*>(output_q.data_ptr()),
static_cast<const float*>(output_s.data_ptr()),
static_cast<int64_t>(num_elements));
if constexpr (!kSkipQuant) {
LaunchKernel(num_blocks, kBlockSize, device.unwrap())(
per_tensor_quant_fp8_kernel<DType, fp8_e4m3_t>,
static_cast<const DType*>(input.data_ptr()),
static_cast<fp8_e4m3_t*>(output_q.data_ptr()),
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
+32
View File
@@ -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
+62
View File
@@ -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
+20 -4
View File
@@ -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 -8
View File
@@ -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
# but no input_activations provided in the config
# should be w8a16fp8 w8a16fp8 can also run for cases where
# there is an input_quant but it is ignored
# When activation quant format is set but no
# input_activations provided: valid for w8a16fp8 (FLOAT
# weights) and pack-quantized without activation quant
# (INT weights, W4A16 style).
if not input_activations:
assert (
target_scheme_map[target]["weights"].type
== QuantizationType.FLOAT
)
weight_type = target_scheme_map[target]["weights"].type
if weight_type == QuantizationType.INT:
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
+2
View File
@@ -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
+323
View File
@@ -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
+1
View File
@@ -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
+3 -1
View File
@@ -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
+22
View File
@@ -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
+3 -2
View File
@@ -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":
disable_reason = "Deepseek V3/R1 W4AFP8 model uses different quant method for routed experts and shared experts."
elif is_wint4afp8_or_wint4a16_config(self.quant_config):
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
+14 -7
View File
@@ -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)
if name not in params_dict:
resolved_name = name.replace(weight_name, param_name)
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 (