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

[Diffusion] Fix weight scale swizzle and add large-M kernel config for FLUX.2-dev-NVFP4 (#22064)

Browse Source
This commit is contained in:
CHEN Xi
2026-04-04 11:50:30 +08:00
committed by GitHub
parent fe92f3563c
commit 31c9d8e885
3 changed files with 30 additions and 217 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_sm100.cuh
+24 -3
View File
@@ -53,7 +53,7 @@ const dim3 KernelConfigM256<T>::preferred_cluster(2, 4, 1);
template <typename T>
const dim3 KernelConfigM256<T>::fallback_cluster(2, 1, 1);
// Default config(half_t/bfloat16_t) for M > 256
// Config(half_t/bfloat16_t) for 256 < M <= 1024
template <typename T>
struct KernelConfigDefault {
using OutputType = T;
@@ -66,10 +66,27 @@ struct KernelConfigDefault {
const static dim3 fallback_cluster;
};
template <typename T>
const dim3 KernelConfigDefault<T>::preferred_cluster(4, 4, 1);
const dim3 KernelConfigDefault<T>::preferred_cluster(2, 4, 1);
template <typename T>
const dim3 KernelConfigDefault<T>::fallback_cluster(2, 1, 1);
// Config(half_t/bfloat16_t) for M > 1024: 1x4 cluster reduces M-tail waste.
template <typename T>
struct KernelConfigLargeM {
using OutputType = T;
using MmaTileShape = Shape<_256, _256, _256>;
using ClusterShape = Shape<int, int, _1>;
using EpilogueTile = Shape<_128, _64>;
using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized2Sm;
using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized2SmNvf4Sm100;
const static dim3 preferred_cluster;
const static dim3 fallback_cluster;
};
template <typename T>
const dim3 KernelConfigLargeM<T>::preferred_cluster(1, 4, 1);
template <typename T>
const dim3 KernelConfigLargeM<T>::fallback_cluster(1, 2, 1);
struct KernelConfigFp32 {
using OutputType = float;
using MmaTileShape = Shape<_128, _128, _256>;
@@ -261,8 +278,12 @@ void cutlassFp4GemmDispatchSm100(
runGemm<Fp4GemmSm100<KernelConfigM128<OutType>>>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
} else if (m <= 256) {
runGemm<Fp4GemmSm100<KernelConfigM256<OutType>>>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
} else {
} else if (m <= 1024) {
// m in (256, 1024]: 2x4 cluster balances SM occupancy and data reuse
runGemm<Fp4GemmSm100<KernelConfigDefault<OutType>>>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
} else {
// m in (1024, inf): 1x4 cluster eliminates M-tail waste for FLUX-class shapes
runGemm<Fp4GemmSm100<KernelConfigLargeM<OutType>>>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
}
}
python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py
+6 -167
View File
@@ -43,16 +43,6 @@ def _get_fp4_gemm_op():
return current_platform.get_modelopt_fp4_gemm_op()
@lru_cache(maxsize=1)
def _get_comfy_kitchen_cuda_backend():
try:
import comfy_kitchen.backends.cuda as ck_cuda
return ck_cuda
except Exception:
return None
class ModelOptQuantConfig(QuantizationConfig):
def __init__(
self,
@@ -224,20 +214,7 @@ class ModelOptFp4Config(ModelOptQuantConfig):
return False
def get_quant_method(self, layer: torch.nn.Module, prefix: str):
should_use_best_perf_kit = getattr(
current_platform, "should_use_modelopt_fp4_best_performance_kit", None
)
warn_missing_best_perf_kit = getattr(
current_platform, "warn_if_modelopt_fp4_best_performance_kit_missing", None
)
if callable(should_use_best_perf_kit) and should_use_best_perf_kit():
linear_cls = ComfyUIFp4LinearMethod
else:
if callable(warn_missing_best_perf_kit):
warn_missing_best_perf_kit()
linear_cls = ModelOptFp4LinearMethod
return self._get_quant_method(layer, prefix, Linear=linear_cls)
return self._get_quant_method(layer, prefix, Linear=ModelOptFp4LinearMethod)
class ModelOptFp4LinearMethod(LinearMethodBase):
@@ -349,6 +326,11 @@ class ModelOptFp4LinearMethod(LinearMethodBase):
K_padded = round_up(K, 4)
padded_scales = torch.zeros((B, M_padded, K_padded), dtype=scales.dtype)
padded_scales[:B, :M, :K] = scales
# Blockwise interleave for CUTLASS TMA layout required by CUTLASS kernel
padded_scales = padded_scales.reshape(
B, M_padded // 128, 4, 32, K_padded // 4, 4
)
padded_scales = padded_scales.permute(0, 1, 4, 3, 2, 5)
padded_scales = padded_scales.contiguous().cuda()
padded_scales = (
padded_scales.reshape(M_padded, K_padded)
@@ -417,146 +399,3 @@ class ModelOptFp4LinearMethod(LinearMethodBase):
if bias is not None:
out = out + bias
return out.view(*output_shape)
class ComfyUIFp4LinearMethod(LinearMethodBase):
"""NVFP4 linear method using comfy-kitchen cuBLAS kernels (Blackwell)."""
def __init__(self, quant_config: ModelOptFp4Config):
self.quant_config = quant_config
def create_weights(
self,
layer: torch.nn.Module,
input_size_per_partition: int,
output_partition_sizes: List[int],
input_size: int,
output_size: int,
params_dtype: torch.dtype,
**extra_weight_attrs,
):
del input_size, output_size
if not self.quant_config.is_checkpoint_nvfp4_serialized:
raise ValueError(
"NVFP4 quantization was selected, "
"dynamic quantization is not supported."
)
if input_size_per_partition % 16 != 0:
raise ValueError(
f"Unsupported model when input features size is {input_size_per_partition}, "
"not multiple of 16, for NVFP4 quantization."
)
output_size_per_partition = sum(output_partition_sizes)
weight_loader = extra_weight_attrs.get("weight_loader")
layer.logical_widths = output_partition_sizes
layer.input_size_per_partition = input_size_per_partition
layer.output_size_per_partition = output_size_per_partition
weight = ModelWeightParameter(
data=torch.empty(
output_size_per_partition,
input_size_per_partition // 2,
dtype=torch.uint8,
),
input_dim=1,
output_dim=0,
weight_loader=weight_loader,
)
layer.register_parameter("weight", weight)
input_scale = PerTensorScaleParameter(
data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
weight_loader=weight_loader,
)
set_weight_attrs(input_scale, {"missing_param_init": "ones"})
layer.register_parameter("input_scale", input_scale)
weight_scale_2 = PerTensorScaleParameter(
data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
weight_loader=weight_loader,
)
layer.register_parameter("weight_scale_2", weight_scale_2)
weight_scale = ModelWeightParameter(
data=torch.empty(
output_size_per_partition,
input_size_per_partition // self.quant_config.group_size,
dtype=torch.float8_e4m3fn,
),
input_dim=1,
output_dim=0,
weight_loader=weight_loader,
)
layer.register_parameter("weight_scale", weight_scale)
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
from comfy_kitchen.float_utils import from_blocked, to_blocked
input_scale = layer.input_scale.max().to(torch.float32)
weight_scale_2 = layer.weight_scale_2.max().to(torch.float32)
copy_or_rebind_param(layer, "input_scale_ck", input_scale.cuda())
copy_or_rebind_param(layer, "weight_scale_2_ck", weight_scale_2.cuda())
layer.output_size_per_partition = layer.weight.shape[0]
copy_or_rebind_param(layer, "weight", layer.weight.data.contiguous().cuda())
# Checkpoint block scales are already in cuBLAS tiled layout.
# Pad to (roundup(N, 128), roundup(K//16, 4)) if needed.
scales = layer.weight_scale.data
N, Ks = scales.shape
N_padded = round_up(N, 128)
Ks_padded = round_up(Ks, 4)
if N == N_padded and Ks == Ks_padded:
weight_scale_ck = scales.cuda()
else:
scales_rm = from_blocked(scales, num_rows=N, num_cols=Ks)
padded_rm = torch.zeros((N_padded, Ks_padded), dtype=scales.dtype)
padded_rm[:N, :Ks] = scales_rm
weight_scale_ck = to_blocked(padded_rm, flatten=False).cuda()
copy_or_rebind_param(layer, "weight_scale_ck", weight_scale_ck)
def apply(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: Optional[torch.Tensor] = None,
) -> torch.Tensor:
ck_cuda = _get_comfy_kitchen_cuda_backend()
if ck_cuda is None:
raise RuntimeError(
"comfy_kitchen is not available. "
"Install it to use ComfyUIFp4LinearMethod."
)
output_dtype = x.dtype
input_shape = x.shape
x_2d = x.view(-1, input_shape[-1]) # [M, K]
M = x_2d.shape[0]
output_size = layer.output_size_per_partition
output_shape = list(input_shape[:-1]) + [output_size]
if not x_2d.is_contiguous():
x_2d = x_2d.contiguous()
x_fp4, x_block_scale = ck_cuda.quantize_nvfp4(
x_2d, layer.input_scale_ck, pad_16x=True
)
out = ck_cuda.scaled_mm_nvfp4(
x_fp4,
layer.weight,
tensor_scale_a=layer.input_scale_ck,
tensor_scale_b=layer.weight_scale_2_ck,
block_scale_a=x_block_scale,
block_scale_b=layer.weight_scale_ck,
bias=bias,
out_dtype=output_dtype,
)
out = out[:M, :output_size].contiguous()
return out.view(*output_shape)
python/sglang/multimodal_gen/runtime/platforms/cuda.py
-47
View File
@@ -145,53 +145,6 @@ class CudaPlatformBase(Platform):
except ImportError:
return None, None
@classmethod
@lru_cache(maxsize=1)
def has_modelopt_fp4_best_performance_kit(cls) -> bool:
try:
import comfy_kitchen.backends.cuda # noqa: F401
return True
except Exception:
return False
@classmethod
@lru_cache(maxsize=1)
def can_use_modelopt_fp4_best_performance_kit(cls) -> bool:
if not cls.is_blackwell() or not cls.has_modelopt_fp4_best_performance_kit():
return False
try:
import comfy_kitchen.backends.cuda as ck_cuda
device = cls.get_local_torch_device()
x = torch.zeros((16, 16), dtype=torch.bfloat16, device=device)
scale = torch.ones((), dtype=torch.float32, device=device)
ck_cuda.quantize_nvfp4(x, scale, pad_16x=True)
return True
except Exception as e:
logger.warning(
"best performance kit (comfy-kitchen) is installed but unusable on "
"this system (%s). Blackwell NVFP4 will fall back to the generic "
"ModelOpt FP4 path.",
e,
)
return False
@classmethod
def should_use_modelopt_fp4_best_performance_kit(cls) -> bool:
return cls.can_use_modelopt_fp4_best_performance_kit()
@classmethod
@lru_cache(maxsize=1)
def warn_if_modelopt_fp4_best_performance_kit_missing(cls) -> None:
if cls.is_blackwell() and not cls.has_modelopt_fp4_best_performance_kit():
logger.warning(
"best performance kit (comfy-kitchen) is not installed. "
"Blackwell NVFP4 will fall back to the generic ModelOpt FP4 path. "
"Install it with `pip install comfy-kitchen[cublas]`."
)
@classmethod
def is_full_nvlink(cls, device_ids: list[int]) -> bool:
raise NotImplementedError