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

[Diffusion] Default NVFP4 backend to FlashInfer TRTLLM (#25523)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-05-25 18:14:06 +08:00
committed by GitHub
parent c05756da7a
commit 533ef41112
11 changed files with 143 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/tests/diffusion/test_diffusion_nvfp4_scaled_mm.py
+58 -1
View File
@@ -136,6 +136,7 @@ def _build_layer(
weight_global_scale: torch.Tensor,
*,
weight_scale_device: torch.device | str | None = None,
checkpoint_weight_scale_layout: str = "linear",
) -> tuple[ModelOptFp4LinearMethod, torch.nn.Module]:
output_size, input_size_half = weight_fp4.shape
input_size = input_size_half * 2
@@ -144,6 +145,7 @@ def _build_layer(
is_checkpoint_nvfp4_serialized=True,
group_size=BLOCK_SIZE,
swap_weight_nibbles=True,
checkpoint_weight_scale_layout=checkpoint_weight_scale_layout,
)
)
layer = torch.nn.Module()
@@ -179,7 +181,11 @@ def _build_layer(
expected_weight, _ = pad_nvfp4_weight(
weight_fp4, n_alignment=128, k_alignment=0
)
expected_scale = weight_scale_linear
expected_scale = (
_swizzled_to_linear(weight_scale_linear, output_size, input_size)
if checkpoint_weight_scale_layout == "swizzled"
else weight_scale_linear
)
if expected_scale.shape[0] != expected_weight.shape[0]:
pad_n = expected_weight.shape[0] - expected_scale.shape[0]
expected_scale = torch.nn.functional.pad(expected_scale, (0, 0, 0, pad_n))
@@ -370,6 +376,57 @@ def test_flux2_shape_correctness_flashinfer_trtllm(
assert diff < DEEPGEMM_FP4_MAX_DIFF, f"{m=}, {n=}, {k=}, {diff=:.6f}"
@pytest.mark.skipif(
not _nvfp4_supported(),
reason="Diffusion NVFP4 scaled mm correctness requires Blackwell GPUs",
)
def test_flux2_swizzled_scale_checkpoint_flashinfer_trtllm_matches_cudnn(
monkeypatch: pytest.MonkeyPatch,
) -> None:
_set_diffusion_fp4_backend(monkeypatch, "flashinfer_trtllm")
m, n, k = FLUX2_PROJECTION_SHAPE
generator = torch.Generator(device=DEVICE)
generator.manual_seed(20260517 + m + n + k)
x = torch.randn((m, k), device=DEVICE, dtype=DTYPE, generator=generator)
weight = torch.randn((n, k), device=DEVICE, dtype=DTYPE, generator=generator)
input_global_scale = _make_global_scale(x)
weight_global_scale = _make_global_scale(weight)
alpha = (1.0 / (input_global_scale * weight_global_scale)).to(torch.float32)
x_fp4, x_scale_swizzled = flashinfer.fp4_quantize(x, input_global_scale)
weight_fp4, weight_scale_swizzled = flashinfer.fp4_quantize(
weight, weight_global_scale
)
if x_scale_swizzled.dtype == torch.uint8:
x_scale_swizzled = x_scale_swizzled.view(torch.float8_e4m3fn)
if weight_scale_swizzled.dtype == torch.uint8:
weight_scale_swizzled = weight_scale_swizzled.view(torch.float8_e4m3fn)
method, layer = _build_layer(
weight_fp4,
weight_scale_swizzled,
input_global_scale,
weight_global_scale,
checkpoint_weight_scale_layout="swizzled",
)
actual = method.apply(layer, x)
expected = flashinfer.mm_fp4(
x_fp4,
weight_fp4.t(),
x_scale_swizzled,
weight_scale_swizzled.t(),
alpha,
DTYPE,
backend="cudnn",
)
diff = _calc_diff(actual, expected)
assert diff < DEEPGEMM_FP4_MAX_DIFF, f"{m=}, {n=}, {k=}, {diff=:.6f}"
@pytest.mark.skipif(
not _nvfp4_supported(),
reason="Diffusion NVFP4 scaled mm correctness requires Blackwell GPUs",
python/sglang/multimodal_gen/envs.py
+1
View File
@@ -283,6 +283,7 @@ environment_variables: dict[str, Callable[[], Any]] = {
"SGLANG_USE_RUNAI_MODEL_STREAMER", "true"
),
# FlashInfer FP4 GEMM backend override for diffusion NVFP4.
# When unset, diffusion ModelOpt NVFP4 defaults to flashinfer_trtllm.
# Supported values:
# - auto
# - flashinfer_cudnn
python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py
+35 -3
View File
@@ -66,6 +66,29 @@ def _prepare_nvfp4_weight_bytes(
return ((weight >> 4) | (weight << 4)).contiguous()
def _swizzled_nvfp4_scales_to_linear(scales: torch.Tensor) -> torch.Tensor:
"""Convert FlashInfer/CUTLASS-swizzled FP4 scales back to row-major layout."""
scale_ndim = scales.ndim
if scale_ndim == 2:
scales = scales.unsqueeze(0)
assert scales.ndim == 3
B, M, K = scales.shape
M_padded = round_up(M, 128)
K_padded = round_up(K, 4)
if M != M_padded or K != K_padded:
padded = torch.zeros(
(B, M_padded, K_padded), dtype=scales.dtype, device=scales.device
)
padded[:B, :M, :K] = scales
scales = padded
linear = scales.reshape(B, M_padded // 128, K_padded // 4, 32, 4, 4)
linear = linear.permute(0, 1, 4, 3, 2, 5).contiguous()
linear = linear.reshape(B, M_padded, K_padded)[:, :M, :K]
return linear.squeeze(0) if scale_ndim == 2 else linear
def _require_flashinfer():
if flashinfer is None:
raise RuntimeError(
@@ -203,6 +226,7 @@ class ModelOptFp4Config(ModelOptQuantConfig):
packed_modules_mapping: Optional[Dict[str, List[str]]] = None,
checkpoint_uses_packed_qkv: bool = False,
swap_weight_nibbles: bool = False,
checkpoint_weight_scale_layout: str = "linear",
) -> None:
super().__init__(exclude_modules, packed_modules_mapping)
self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized
@@ -214,6 +238,7 @@ class ModelOptFp4Config(ModelOptQuantConfig):
self.group_size = group_size
self.checkpoint_uses_packed_qkv = checkpoint_uses_packed_qkv
self.swap_weight_nibbles = swap_weight_nibbles
self.checkpoint_weight_scale_layout = checkpoint_weight_scale_layout
@classmethod
def get_name(cls) -> str:
@@ -311,6 +336,9 @@ class ModelOptFp4Config(ModelOptQuantConfig):
packed_modules_mapping=config.get("packed_modules_mapping"),
checkpoint_uses_packed_qkv=config.get("checkpoint_uses_packed_qkv", False),
swap_weight_nibbles=swap_weight_nibbles,
checkpoint_weight_scale_layout=config.get(
"checkpoint_weight_scale_layout", "linear"
),
)
def get_quant_method(self, layer: torch.nn.Module, prefix: str):
@@ -405,7 +433,7 @@ class ModelOptFp8LinearMethod(LinearMethodBase):
class ModelOptFp4LinearMethod(LinearMethodBase):
"""NVFP4 linear method using CUTLASS FP4 GEMM."""
"""NVFP4 linear method using the selected FP4 GEMM backend."""
def __init__(self, quant_config: ModelOptFp4Config):
self.quant_config = quant_config
@@ -504,13 +532,18 @@ class ModelOptFp4LinearMethod(LinearMethodBase):
self.quant_config, "swap_weight_nibbles", False
),
)
scales = layer.weight_scale
if (
getattr(self.quant_config, "checkpoint_weight_scale_layout", "linear")
== "swizzled"
):
scales = _swizzled_nvfp4_scales_to_linear(scales)
_, flashinfer_backend = _get_fp4_gemm_op()
if flashinfer_backend == "trtllm":
flashinfer_ops = _require_flashinfer()
weight, _ = pad_nvfp4_weight(w_swapped, n_alignment=128, k_alignment=0)
scales = layer.weight_scale
if scales.shape[0] != weight.shape[0]:
pad_n = weight.shape[0] - scales.shape[0]
scales = torch.nn.functional.pad(scales, (0, 0, 0, pad_n))
@@ -550,7 +583,6 @@ class ModelOptFp4LinearMethod(LinearMethodBase):
layer.weights_padding_cols = weights_padding_cols
copy_or_rebind_param(layer, "weight", weight)
scales = layer.weight_scale
scale_ndim = scales.ndim
if scale_ndim == 2:
scales = scales.unsqueeze(0)
python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py
+11
View File
@@ -94,6 +94,17 @@ def _merge_modelopt_fp4_configs(
inferred_config.swap_weight_nibbles = getattr(
inferred_config, "swap_weight_nibbles", False
) or getattr(existing_config, "swap_weight_nibbles", False)
existing_scale_layout = getattr(
existing_config, "checkpoint_weight_scale_layout", "linear"
)
inferred_scale_layout = getattr(
inferred_config, "checkpoint_weight_scale_layout", "linear"
)
inferred_config.checkpoint_weight_scale_layout = (
existing_scale_layout
if inferred_scale_layout == "linear" and existing_scale_layout != "linear"
else inferred_scale_layout
)
if getattr(inferred_config, "group_size", None) is None:
inferred_config.group_size = getattr(existing_config, "group_size", None)
python/sglang/multimodal_gen/runtime/platforms/cuda.py
+13 -25
View File
@@ -124,7 +124,7 @@ class CudaPlatformBase(Platform):
@lru_cache(maxsize=1)
def get_modelopt_flashinfer_fp4_backend(cls) -> str:
backend = envs.SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND
default_backend = "cudnn" if cls.is_blackwell() else "auto"
default_backend = "trtllm"
if backend is None:
return default_backend
@@ -151,35 +151,23 @@ class CudaPlatformBase(Platform):
@lru_cache(maxsize=1)
def get_modelopt_fp4_gemm_op(cls) -> tuple[Callable | None, str | None]:
requested_backend = envs.SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND
prefer_flashinfer = requested_backend is not None
# TODO: Remove this explicit FlashInfer preference once the sm100 CUTLASS
# LargeM dispatch grows a validated fallback for Blackwell NVFP4 shapes
# such as Wan2.2's large-M attention projections.
if prefer_flashinfer:
try:
from flashinfer import mm_fp4 as flashinfer_mm_fp4
return flashinfer_mm_fp4, cls.get_modelopt_flashinfer_fp4_backend()
except ImportError:
logger.warning(
"Requested SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND=%r "
"but flashinfer.mm_fp4 is unavailable. Falling back to "
"cutlass.",
requested_backend,
)
try:
from sgl_kernel import cutlass_scaled_fp4_mm as cutlass_fp4_gemm
return cutlass_fp4_gemm, None
except ImportError:
pass
try:
from flashinfer import mm_fp4 as flashinfer_mm_fp4
return flashinfer_mm_fp4, cls.get_modelopt_flashinfer_fp4_backend()
except ImportError:
logger.warning(
"Requested SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND=%r "
"but flashinfer.mm_fp4 is unavailable. Falling back to "
"cutlass.",
requested_backend or "flashinfer_trtllm (default)",
)
try:
from sgl_kernel import cutlass_scaled_fp4_mm as cutlass_fp4_gemm
return cutlass_fp4_gemm, None
except ImportError:
return None, None
python/sglang/multimodal_gen/runtime/utils/quantization_utils.py
+9 -1
View File
@@ -438,14 +438,22 @@ def _build_nvfp4_config_from_safetensors_files(
"group_size": group_size,
"ignore": exclude_modules,
"checkpoint_uses_packed_qkv": checkpoint_uses_packed_qkv,
# The official FLUX.2 mixed NVFP4 export is detected by its
# packed QKV tensors and stores block scales in the
# FlashInfer/CUTLASS-swizzled layout. SGLang-converted
# transformer repos keep the linear layout.
"checkpoint_weight_scale_layout": (
"swizzled" if checkpoint_uses_packed_qkv else "linear"
),
}
)
logger.info(
"Built NVFP4 quant config from %d safetensors: group_size=%d, %d excluded modules, packed_qkv=%s",
"Built NVFP4 quant config from %d safetensors: group_size=%d, %d excluded modules, packed_qkv=%s, scale_layout=%s",
len(files_with_nvfp4_signal),
group_size,
len(exclude_modules),
checkpoint_uses_packed_qkv,
getattr(result, "checkpoint_weight_scale_layout", "linear"),
)
return result
except Exception as e:
python/sglang/multimodal_gen/test/server/testcase_configs.py
+2 -4
View File
@@ -451,10 +451,8 @@ MODELOPT_QWEN_IMAGE_EDIT_FP8_TRANSFORMER = (
MODELOPT_FLUX1_NVFP4_TRANSFORMER = "lmsys/flux1-dev-modelopt-nvfp4-sglang-transformer"
MODELOPT_FLUX2_NVFP4_WEIGHTS = "black-forest-labs/FLUX.2-dev-NVFP4"
MODELOPT_WAN22_NVFP4_MODEL = "nvidia/Wan2.2-T2V-A14B-Diffusers-NVFP4"
MODELOPT_NVFP4_B200_ENV_VARS = {"SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND": "cudnn"}
MODELOPT_WAN22_NVFP4_B200_ENV_VARS = {
"SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND": "trtllm"
}
MODELOPT_NVFP4_B200_ENV_VARS = {}
MODELOPT_WAN22_NVFP4_B200_ENV_VARS = {}
def _make_modelopt_ci_case(