minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 10 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
docs/diffusion/environment_variables.md
+1 -1
View File
@@ -35,7 +35,7 @@
| Environment Variable | Default | Description |
|----------------------|---------|-------------|
| `SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND` | not set | FlashInfer FP4 GEMM backend for generic NVFP4 fallback |
| `SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND` | not set | Optional FlashInfer FP4 GEMM backend override for diffusion NVFP4. When unset, SGLang defaults to `flashinfer_trtllm`. |
## Caching Acceleration
docs/diffusion/quantization.md
+4 -10
View File
@@ -125,7 +125,7 @@ official `black-forest-labs/FLUX.2-dev-NVFP4` repo.
| `FP8` | `Qwen/Qwen-Image-Edit-2511` | `--transformer-path` | `lmsys/qwen-image-edit-modelopt-fp8-sglang-transformer` | TI2I edit path, BF16-vs-FP8 image comparison, H100 benchmark | shares `QwenImageTransformer2DModel` with Qwen Image and uses the same Qwen Image FP8 fallback preset |
| `NVFP4` | `black-forest-labs/FLUX.1-dev` | `--transformer-path` | `lmsys/flux1-dev-modelopt-nvfp4-sglang-transformer` | mixed BF16+NVFP4 transformer override, correctness validation, 4x RTX 5090 benchmark, torch-profiler trace | use `build_modelopt_nvfp4_transformer.py`; validated builder keeps selected FLUX.1 modules in BF16 and sets `swap_weight_nibbles=false` |
| `NVFP4` | `black-forest-labs/FLUX.2-dev` | `--transformer-weights-path` | `black-forest-labs/FLUX.2-dev-NVFP4` | packed-QKV load path | official raw export repo; validated packed export detection and runtime layout handling |
| `NVFP4` | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `--transformer-path` | `lmsys/wan22-t2v-a14b-modelopt-nvfp4-sglang-transformer` | primary `transformer` quantized with ModelOpt NVFP4, `transformer_2` kept BF16 | primary-transformer-only path; keep `transformer_2` on the base checkpoint, and current B200/Blackwell bring-up uses `SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND=cudnn` |
| `NVFP4` | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `--transformer-path` | `lmsys/wan22-t2v-a14b-modelopt-nvfp4-sglang-transformer` | primary `transformer` quantized with ModelOpt NVFP4, `transformer_2` kept BF16 | primary-transformer-only path; keep `transformer_2` on the base checkpoint; the default FP4 GEMM backend is `flashinfer_trtllm` |
These nine checkpoints are also the intended case set for the B200 diffusion
CI job (`multimodal-gen-test-1-b200`).
@@ -261,7 +261,6 @@ For a dual-transformer Wan2.2 export where only the primary `transformer`
was quantized:
```bash
SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND=cudnn \
sglang generate \
--model-path Wan-AI/Wan2.2-T2V-A14B-Diffusers \
--transformer-path lmsys/wan22-t2v-a14b-modelopt-nvfp4-sglang-transformer \
@@ -279,20 +278,15 @@ sglang generate \
primary `--transformer-path` override targets only `transformer`. Use a
per-component override such as `--transformer-2-path` only when you
intentionally want a non-default `transformer_2`.
- On Blackwell, the validated Wan2.2 ModelOpt NVFP4 path currently prefers
FlashInfer FP4 GEMM via
`SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND=cudnn`.
- This environment-variable override is a current workaround for NVFP4 cases
where the default sglang JIT/CUTLASS `sm100` path rejects a large-M shape at
`can_implement()`. The intended long-term fix is to add a validated CUTLASS
fallback for those shapes rather than rely on the override.
- On Blackwell, the diffusion ModelOpt NVFP4 path defaults to FlashInfer
TensorRT-LLM FP4 GEMM (`flashinfer_trtllm`).
- Direct `--model-path` loading is a compatibility path for FLUX.2 NVFP4-style
repos or local directories.
- If `--transformer-weights-path` is provided explicitly, it takes precedence
over the compatibility `--model-path` flow.
- For local directories, SGLang first looks for `*-mixed.safetensors`, then
falls back to loading from the directory.
- To force the generic diffusion ModelOpt FP4 path onto a specific FlashInfer
- To force the diffusion ModelOpt FP4 path onto a different FlashInfer
backend, set `SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND`. Supported values
include `flashinfer_cudnn`, `flashinfer_cutlass`, and `flashinfer_trtllm`.
- On disk, the quantization config stays `quant_method=modelopt` with
docs_new/docs/sglang-diffusion/environment_variables.mdx
+1 -1
View File
@@ -150,7 +150,7 @@ description: "Configure SGLang diffusion behavior with environment variables."
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>not set</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>FlashInfer FP4 GEMM backend for generic NVFP4 fallback</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Optional FlashInfer FP4 GEMM backend override for diffusion NVFP4. When unset, SGLang defaults to <code>flashinfer_trtllm</code>.</td>
</tr>
</tbody>
</table>
docs_new/docs/sglang-diffusion/quantization.mdx
+8 -13
View File
@@ -210,7 +210,7 @@ official full Diffusers repos, and the FLUX.2 NVFP4 entry keeps the official
<td><code>--model-path</code></td>
<td><code>nvidia/Wan2.2-T2V-A14B-Diffusers-NVFP4</code></td>
<td>full Diffusers repo with ModelOpt NVFP4 Wan2.2 components</td>
<td>current B200/Blackwell bring-up uses <code>SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND=trtllm</code></td>
<td>default FP4 GEMM backend is <code>flashinfer_trtllm</code></td>
</tr>
</tbody>
</table>
@@ -327,7 +327,6 @@ sglang generate \
For Wan2.2 NVFP4:
```bash
SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND=trtllm \
sglang generate \
--model-path nvidia/Wan2.2-T2V-A14B-Diffusers-NVFP4 \
--prompt "a fox walking through neon rain" \
@@ -340,23 +339,19 @@ sglang generate \
directories that already include `config.json`.
- Use `--transformer-weights-path` for raw NVFP4 exports, individual
safetensors files, or repo layouts that should be treated as weights first.
- For legacy mixed Wan2.2 transformer overrides, the primary
`--transformer-path` override targets only `transformer`. Use a per-component
override such as `--transformer-2-path` only when you intentionally want a
non-default `transformer_2`.
- On Blackwell, the validated Wan2.2 ModelOpt NVFP4 path currently prefers
FlashInfer FP4 GEMM via
`SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND=trtllm`.
- This environment-variable override selects the validated Wan2.2 NVFP4
full-repo path on Blackwell while the other NVFP4 CI cases continue to use
the generic `cudnn` backend.
- For dual-transformer pipelines such as `Wan2.2-T2V-A14B-Diffusers`, the
primary `--transformer-path` override targets only `transformer`. Use a
per-component override such as `--transformer-2-path` only when you
intentionally want a non-default `transformer_2`.
- On Blackwell, the diffusion ModelOpt NVFP4 path defaults to FlashInfer
TensorRT-LLM FP4 GEMM (`flashinfer_trtllm`).
- Direct `--model-path` loading is a compatibility path for FLUX.2 NVFP4-style
repos or local directories.
- If `--transformer-weights-path` is provided explicitly, it takes precedence
over the compatibility `--model-path` flow.
- For local directories, SGLang first looks for `*-mixed.safetensors`, then
falls back to loading from the directory.
- To force the generic diffusion ModelOpt FP4 path onto a specific FlashInfer
- To force the diffusion ModelOpt FP4 path onto a different FlashInfer
backend, set `SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND`. Supported values
include `flashinfer_cudnn`, `flashinfer_cutlass`, and `flashinfer_trtllm`.
- On disk, the quantization config stays `quant_method=modelopt` with
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(