diff --git a/docs/diffusion/environment_variables.md b/docs/diffusion/environment_variables.md index 745c84af2..a9ba2d250 100644 --- a/docs/diffusion/environment_variables.md +++ b/docs/diffusion/environment_variables.md @@ -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 diff --git a/docs/diffusion/quantization.md b/docs/diffusion/quantization.md index 4f2e988cf..dfcd3b3f3 100644 --- a/docs/diffusion/quantization.md +++ b/docs/diffusion/quantization.md @@ -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 diff --git a/docs_new/docs/sglang-diffusion/environment_variables.mdx b/docs_new/docs/sglang-diffusion/environment_variables.mdx index 8ade9a7ca..b864521e9 100644 --- a/docs_new/docs/sglang-diffusion/environment_variables.mdx +++ b/docs_new/docs/sglang-diffusion/environment_variables.mdx @@ -150,7 +150,7 @@ description: "Configure SGLang diffusion behavior with environment variables." SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND not set - FlashInfer FP4 GEMM backend for generic NVFP4 fallback + Optional FlashInfer FP4 GEMM backend override for diffusion NVFP4. When unset, SGLang defaults to flashinfer_trtllm. diff --git a/docs_new/docs/sglang-diffusion/quantization.mdx b/docs_new/docs/sglang-diffusion/quantization.mdx index f4ab9d0f5..8fedb80b2 100644 --- a/docs_new/docs/sglang-diffusion/quantization.mdx +++ b/docs_new/docs/sglang-diffusion/quantization.mdx @@ -210,7 +210,7 @@ official full Diffusers repos, and the FLUX.2 NVFP4 entry keeps the official --model-path nvidia/Wan2.2-T2V-A14B-Diffusers-NVFP4 full Diffusers repo with ModelOpt NVFP4 Wan2.2 components - current B200/Blackwell bring-up uses SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND=trtllm + default FP4 GEMM backend is flashinfer_trtllm @@ -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 diff --git a/python/sglang/jit_kernel/tests/diffusion/test_diffusion_nvfp4_scaled_mm.py b/python/sglang/jit_kernel/tests/diffusion/test_diffusion_nvfp4_scaled_mm.py index 1214497d7..9c016158b 100644 --- a/python/sglang/jit_kernel/tests/diffusion/test_diffusion_nvfp4_scaled_mm.py +++ b/python/sglang/jit_kernel/tests/diffusion/test_diffusion_nvfp4_scaled_mm.py @@ -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", diff --git a/python/sglang/multimodal_gen/envs.py b/python/sglang/multimodal_gen/envs.py index febd85b60..a5341a777 100644 --- a/python/sglang/multimodal_gen/envs.py +++ b/python/sglang/multimodal_gen/envs.py @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py b/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py index 2f3d7c33e..1a3b76e98 100755 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py @@ -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) diff --git a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py index 204608c24..58f167c94 100644 --- a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py +++ b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py @@ -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) diff --git a/python/sglang/multimodal_gen/runtime/platforms/cuda.py b/python/sglang/multimodal_gen/runtime/platforms/cuda.py index 32b8212a8..cd83f78be 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/cuda.py +++ b/python/sglang/multimodal_gen/runtime/platforms/cuda.py @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py b/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py index cfd4eb684..168227e7b 100644 --- a/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py +++ b/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py @@ -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: diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index 5d0bb5617..e3b47a22b 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py @@ -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(