chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[FIX] Always support TP > 4 for FP4 Gemm (#17300)

Browse Source
This commit is contained in:
danielafrimi
2026-02-05 09:10:26 +02:00
committed by GitHub
parent 368936a62b
commit 3f1df322f9
1 changed files with 154 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

160
python/sglang/srt/layers/quantization/modelopt_quant.py
View File

@@ -158,6 +158,103 @@ CUTEDSL_MOE_SCALAR_INPUT_SCALE = get_bool_env_var(
"SGLANG_CUTEDSL_MOE_SCALAR_INPUT_SCALE", "true"
)
# FP4 GEMM alignment constant - CUTLASS/FlashInfer kernels require dimensions divisible by 32
FP4_GEMM_ALIGNMENT = 32
def round_up_to_multiple(x: int, m: int) -> int:
"""Round up x to the nearest multiple of m."""
return (x + m - 1) // m * m
def pad_nvfp4_weight(
weight: torch.Tensor,
n_alignment: int = FP4_GEMM_ALIGNMENT,
k_alignment: int = FP4_GEMM_ALIGNMENT,
) -> tuple[torch.Tensor, int]:
"""
Pad packed NVFP4 weights to satisfy alignment constraints for FP4 GEMM kernels.
Different backends have different alignment requirements:
- CUTLASS/cuDNN: N % 32 == 0, K % 32 == 0
- TRTLLM: N % 128 == 0 (for shuffle_matrix_sf_a), K padding handled separately
Args:
weight: Packed FP4 weight tensor of shape [N, K//2] (2 FP4 values per byte)
n_alignment: Required alignment for N dimension (default 32, use 128 for TRTLLM)
k_alignment: Required alignment for K dimension (default 32, use 0 to skip)
Returns:
Tuple of (padded_weight, weights_padding_cols) where weights_padding_cols
is the number of columns added for K-dimension padding (in bytes).
"""
weight_current_rows = weight.shape[0] # N dimension
weight_current_col_bytes = weight.shape[1] # K//2 (packed)
# Calculate padding for N dimension (rows)
pad_rows = 0
if n_alignment > 0 and weight_current_rows % n_alignment != 0:
total_rows = round_up_to_multiple(weight_current_rows, n_alignment)
pad_rows = total_rows - weight_current_rows
# Calculate padding for K dimension (columns)
# 2 FP4 items are packed per byte in the input dimension
weight_current_col_elements = weight_current_col_bytes * 2
pad_cols_bytes = 0
if k_alignment > 0 and weight_current_col_elements % k_alignment != 0:
total_cols = round_up_to_multiple(weight_current_col_elements, k_alignment)
pad_cols = total_cols - weight_current_col_elements
# pad_cols is in elements, but padding is in bytes (2 elements per byte)
pad_cols_bytes = pad_cols // 2
# Apply padding in a single operation if needed
# For 2D tensor, pad argument is (pad_left, pad_right, pad_top, pad_bottom)
if pad_rows > 0 or pad_cols_bytes > 0:
weight = torch.nn.functional.pad(
weight, (0, pad_cols_bytes, 0, pad_rows)
).contiguous()
return weight, pad_cols_bytes
def pad_nvfp4_activation_for_cutlass(
x_fp4: torch.Tensor,
weights_padding_cols: int,
) -> torch.Tensor:
"""
Pad packed FP4 activations to match the K-dimension padding applied to weights.
Args:
x_fp4: Packed FP4 activation tensor
weights_padding_cols: Number of padding columns (in bytes) from weight padding
Returns:
Padded activation tensor
"""
if weights_padding_cols > 0:
return torch.nn.functional.pad(x_fp4, (0, weights_padding_cols)).contiguous()
return x_fp4
def slice_nvfp4_output(
out: torch.Tensor,
output_size: int,
) -> torch.Tensor:
"""
Slice the output tensor to remove padding in N dimension if weight was padded.
Args:
out: Output tensor from FP4 GEMM
output_size: Original output size before padding
Returns:
Sliced output tensor with padding removed
"""
if out.shape[-1] != output_size:
return out[..., :output_size].contiguous()
return out
# TODO make it true by default when the DeepEP PR is merged
MOE_NVFP4_DISPATCH = envs.SGLANG_MOE_NVFP4_DISPATCH.get()
# Supported activation schemes for the current configuration
@@ -1059,26 +1156,66 @@ class ModelOptFp4LinearMethod(LinearMethodBase):
layer.input_scale_inv = Parameter(
(1 / input_scale_2).to(torch.float32), requires_grad=False
)
# Store original output size before any padding
layer.output_size_per_partition = layer.weight.shape[0]
if get_fp4_gemm_runner_backend().is_flashinfer_trtllm():
# FlashInfer TRTLLM FP4 GEMM requires a different weight layout.
# FlashInfer provides nvfp4_quantize to quantize + shuffle the
# layout but we use our own quantization so we have to call
# shuffles ourselves.
#
# Alignment requirements:
# - shuffle_matrix_a: weight.shape[0] (N) % 32 == 0
# - shuffle_matrix_sf_a: scale.shape[0] (N) % 128 == 0, scale.shape[1] (K/16) % 4 == 0
# We pad N to multiple of 128 and K/16 to multiple of 4.
from flashinfer import shuffle_matrix_a, shuffle_matrix_sf_a
weight = layer.weight
# Pad weight N dimension to 128
weight, _ = pad_nvfp4_weight(
layer.weight.data, n_alignment=128, k_alignment=0
)
# Pad scale N dimension to match weight
scale = layer.weight_scale
if scale.shape[0] != weight.shape[0]:
pad_n = weight.shape[0] - scale.shape[0]
scale = torch.nn.functional.pad(scale, (0, 0, 0, pad_n))
# Pad K dimension: scale K/16 must be multiple of 4
scale_k = scale.shape[1] # K/16
weights_padding_cols = 0
if scale_k % 4 != 0:
padded_scale_k = round_up_to_multiple(scale_k, 4)
pad_scale_k = padded_scale_k - scale_k
# Pad scale K/16 dimension
scale = torch.nn.functional.pad(scale, (0, pad_scale_k, 0, 0))
# Pad weight K/2 dimension correspondingly (K/2 = K/16 * 8)
pad_weight_k = pad_scale_k * 8
weight = torch.nn.functional.pad(weight, (0, pad_weight_k, 0, 0))
# Store K padding for activation padding in apply()
weights_padding_cols = pad_weight_k
# Shuffle for TRTLLM layout
epilogue_tile_m = 128
shuffled_scale_shape = scale.shape
weight = shuffle_matrix_a(weight.view(torch.uint8), epilogue_tile_m)
scale = (
shuffle_matrix_sf_a(scale.view(torch.uint8), epilogue_tile_m)
.reshape(scale.shape)
.reshape(shuffled_scale_shape)
.view(torch.float8_e4m3fn)
)
layer.weight_scale_interleaved = Parameter(scale, requires_grad=False)
layer.weight = Parameter(weight, requires_grad=False)
layer.weights_padding_cols = weights_padding_cols
return
# Pad weights for CUTLASS/FlashInfer kernel alignment (K and N divisible by 32)
weight, weights_padding_cols = pad_nvfp4_weight(layer.weight.data)
layer.weights_padding_cols = weights_padding_cols
layer.weight = Parameter(weight, requires_grad=False)
# Pad and blockwise interleave weight_scale
scales = layer.weight_scale
scale_ndim = scales.ndim
@@ -1086,9 +1223,8 @@ class ModelOptFp4LinearMethod(LinearMethodBase):
scales = scales.unsqueeze(0)
assert scales.ndim == 3
B, M, K = scales.shape
round_up_multiple = lambda x, m: (x + m - 1) // m * m
M_padded = round_up_multiple(M, 128)
K_padded = round_up_multiple(K, 4)
M_padded = round_up_to_multiple(M, 128)
K_padded = round_up_to_multiple(K, 4)
padded_scales = torch.zeros((B, M_padded, K_padded), dtype=scales.dtype)
padded_scales[:B, :M, :K] = scales
batches, rows, cols = padded_scales.shape
@@ -1112,8 +1248,11 @@ class ModelOptFp4LinearMethod(LinearMethodBase):
) -> torch.Tensor:
output_dtype = x.dtype
x_m, _ = x.shape
# Get original output size (before padding) and padded weight size
output_size = layer.output_size_per_partition
w_n, _ = layer.weight.shape
output_shape = [x_m, w_n]
output_shape = [x_m, output_size]
# Quantize BF16 or FP16 to (FP4 and interleaved block scale)
x_fp4, x_scale_interleaved = fp4_quantize(x, layer.input_scale_inv)
@@ -1123,11 +1262,16 @@ class ModelOptFp4LinearMethod(LinearMethodBase):
assert layer.weight_scale_interleaved.dtype == torch.float8_e4m3fn
assert layer.alpha.dtype == torch.float32
# Pad activations to match weight K-dimension padding
weights_padding_cols = getattr(layer, "weights_padding_cols", 0)
x_fp4 = pad_nvfp4_activation_for_cutlass(x_fp4, weights_padding_cols)
w = layer.weight
w_scale_interleaved = layer.weight_scale_interleaved
if enable_flashinfer_fp4_gemm:
w = layer.weight.T
w_scale_interleaved = layer.weight_scale_interleaved.T
out = fp4_gemm(
x_fp4,
w,
@@ -1137,6 +1281,10 @@ class ModelOptFp4LinearMethod(LinearMethodBase):
output_dtype,
w_n,
)
# Slice output to remove N-dimension padding
out = slice_nvfp4_output(out, output_size)
if bias is not None:
out = out + bias
return out.view(*output_shape)