From 52e1c24744bf4efe75fe976e26596ae1c9f279e2 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Mon, 31 Aug 2026 21:33:02 +0800 Subject: [PATCH] [Diffusion] Fuse FLUX.2 token concatenation and NVFP4 quantization (#37141) --- .../csrc/diffusion/flux2_token_cat_nvfp4.cuh | 120 ++++++++++++++++++ python/sglang/kernels/jit/utils/deps.py | 20 +++ .../sglang/kernels/ops/diffusion/__init__.py | 8 ++ .../layout/flux2_token_cat_nvfp4_jit.py | 101 +++++++++++++++ .../layers/quantization/modelopt_quant.py | 34 +++++ .../runtime/models/dits/flux_2.py | 34 ++++- .../kernels/bench_flux2_token_cat_nvfp4.py | 90 +++++++++++++ .../ops/diffusion/test_model_fast_paths.py | 35 +++++ .../test_diffusion_nvfp4_scaled_mm.py | 35 +++++ 9 files changed, 474 insertions(+), 3 deletions(-) create mode 100644 python/sglang/kernels/jit/csrc/diffusion/flux2_token_cat_nvfp4.cuh create mode 100644 python/sglang/kernels/ops/diffusion/layout/flux2_token_cat_nvfp4_jit.py create mode 100644 test/manual/kernels/bench_flux2_token_cat_nvfp4.py diff --git a/python/sglang/kernels/jit/csrc/diffusion/flux2_token_cat_nvfp4.cuh b/python/sglang/kernels/jit/csrc/diffusion/flux2_token_cat_nvfp4.cuh new file mode 100644 index 000000000..faffee37b --- /dev/null +++ b/python/sglang/kernels/jit/csrc/diffusion/flux2_token_cat_nvfp4.cuh @@ -0,0 +1,120 @@ +// FLUX.2 single-block [attention | MLP] concatenation + NVFP4 quantization. + +#pragma once + +#include + +#include +#include + +#ifndef FLT_MAX +#define FLT_MAX __FLT_MAX__ +#endif +#include + +#include + +namespace sglang { + +namespace flux2_token_cat_nvfp4 { + +constexpr int kAttentionHidden = 6144; +constexpr int kMlpHidden = 18432; +constexpr int kOutputHidden = kAttentionHidden + kMlpHidden; +constexpr int kGroupSize = 16; +constexpr int kScaleColumns = kOutputHidden / kGroupSize; +constexpr int kPackedColumns = kOutputHidden / 2; +constexpr int kAttentionGroups = kAttentionHidden / kGroupSize; +constexpr int kThreads = 256; +constexpr int kColumnTiles = (kScaleColumns + kThreads - 1) / kThreads; +constexpr int kMaxRows = 65408; + +static_assert(kScaleColumns == 1536); +static_assert(kColumnTiles == 6); + +struct Params { + void* quantized; + void* quant_scales; + const void* attention; + const void* mlp; + const void* global_scale; + uint32_t num_rows; +}; + +__global__ void kernel(const Params __grid_constant__ params) { + using namespace device; + using Vec = AlignedVector; + + const int row = blockIdx.y; + const int group = blockIdx.x * kThreads + threadIdx.x; + if (group >= kScaleColumns) { + return; + } + + const int64_t scale_offset = tensorrt_llm::kernels::get_sf_out_offset_128x4(row, group, kScaleColumns); + auto* quant_scales = static_cast(params.quant_scales); + if (row >= params.num_rows) { + quant_scales[scale_offset] = 0; + return; + } + + Vec input; + if (group < kAttentionGroups) { + input.load(static_cast(params.attention) + int64_t(row) * kAttentionHidden + group * kGroupSize); + } else { + const int mlp_group = group - kAttentionGroups; + input.load(static_cast(params.mlp) + int64_t(row) * kMlpHidden + mlp_group * kGroupSize); + } + + tensorrt_llm::kernels::PackedVec<__nv_bfloat16, kGroupSize> quant_vec; + auto* quant_values = reinterpret_cast<__nv_bfloat16*>(&quant_vec); +#pragma unroll + for (int element = 0; element < kGroupSize; ++element) { + quant_values[element] = static_cast<__nv_bfloat16>(input[element]); + } + + const float global_scale = *static_cast(params.global_scale); + const uint64_t packed = tensorrt_llm::kernels::cvt_warp_fp16_to_fp4<__nv_bfloat16, kGroupSize, kGroupSize, false>( + quant_vec, global_scale, quant_scales + scale_offset); + static_cast(params.quantized)[int64_t(row) * kScaleColumns + group] = packed; +} + +struct Kernel { + static void + run(tvm::ffi::TensorView quantized, + tvm::ffi::TensorView quant_scales, + tvm::ffi::TensorView attention, + tvm::ffi::TensorView mlp, + tvm::ffi::TensorView global_scale) { + using namespace host; + auto rows = SymbolicSize{"rows"}; + auto padded_rows = SymbolicSize{"padded_rows"}; + auto device = SymbolicDevice{}; + device.set_options(); + + TensorMatcher({rows, kAttentionHidden}).with_dtype().with_device(device).verify(attention); + TensorMatcher({rows, kMlpHidden}).with_dtype().with_device(device).verify(mlp); + TensorMatcher({rows, kPackedColumns}).with_dtype().with_device(device).verify(quantized); + TensorMatcher({padded_rows, kScaleColumns}).with_dtype().with_device(device).verify(quant_scales); + TensorMatcher({1}).with_dtype().with_device(device).verify(global_scale); + RuntimeCheck(rows.unwrap() > 0, "rows must be positive"); + RuntimeCheck(rows.unwrap() <= kMaxRows, "rows exceed the CUDA grid.y limit"); + const uint32_t row_count = static_cast(rows.unwrap()); + const uint32_t padded_row_count = div_ceil(row_count, uint32_t(128)) * 128; + RuntimeCheck(padded_rows.unwrap() == padded_row_count, "quant scale rows must be padded to 128"); + + const auto params = Params{ + .quantized = quantized.data_ptr(), + .quant_scales = quant_scales.data_ptr(), + .attention = attention.data_ptr(), + .mlp = mlp.data_ptr(), + .global_scale = global_scale.data_ptr(), + .num_rows = row_count, + }; + LaunchKernel(dim3(kColumnTiles, padded_row_count), kThreads, device.unwrap())(kernel, params); + } +}; + +} // namespace flux2_token_cat_nvfp4 + +} // namespace sglang diff --git a/python/sglang/kernels/jit/utils/deps.py b/python/sglang/kernels/jit/utils/deps.py index 16ce9b22f..5fd81f489 100644 --- a/python/sglang/kernels/jit/utils/deps.py +++ b/python/sglang/kernels/jit/utils/deps.py @@ -57,6 +57,26 @@ def get_flashinfer_include_paths() -> List[str]: return include_paths +@register_dependency("flashinfer_nv_internal") +def get_flashinfer_nv_internal_include_paths() -> List[str]: + flashinfer_root = _find_package_root("flashinfer") + if flashinfer_root is None: + raise RuntimeError( + "Cannot find flashinfer package. Please install flashinfer to get " + "the required NVFP4 headers for JIT compilation." + ) + + internal_root = flashinfer_root / "data" / "csrc" / "nv_internal" + candidates = [internal_root, internal_root / "include"] + for path in candidates: + if not path.exists(): + raise RuntimeError( + f"Required FlashInfer NVFP4 header path {path} was not found. " + "Please install a FlashInfer build with nv_internal headers." + ) + return [str(path) for path in candidates] + + def get_mathdx_root() -> Optional[pathlib.Path]: """Locate the NVIDIA Math-DX install (cuBLASDx headers). diff --git a/python/sglang/kernels/ops/diffusion/__init__.py b/python/sglang/kernels/ops/diffusion/__init__.py index 540232d38..892129a2f 100644 --- a/python/sglang/kernels/ops/diffusion/__init__.py +++ b/python/sglang/kernels/ops/diffusion/__init__.py @@ -333,6 +333,13 @@ _SPECS: tuple[tuple[str, KernelBackend, str, frozenset, str], ...] = ( _CUDA, "Wan causal VAE main + DupUp3D(src).", ), + ( + "diffusion.flux2_token_cat_nvfp4", + KernelBackend.JIT, + "layout.flux2_token_cat_nvfp4_jit:try_flux2_token_cat_nvfp4", + _CUDA, + "FLUX.2 single-block token concatenation + NVFP4 quantization.", + ), ) for _op, _backend, _target, _caps, _description in _SPECS: @@ -453,6 +460,7 @@ _EXPORTS: dict[str, str] = { "fused_scatter_to_padded": "layout.varlen_pack_pad_triton", "cat_pad_channels_last_3d": "layout.wan_causal_cache_triton", "dup_up3d_add": "layout.wan_causal_cache_triton", + "try_flux2_token_cat_nvfp4": "layout.flux2_token_cat_nvfp4_jit", # Fusion-site policy: quality gate, first-sight verification, mount "BitExactFusionGate": "sites.bitexact_gate", "flashinfer_rmsnorm_diagnostic_hint": "sites.bitexact_gate", diff --git a/python/sglang/kernels/ops/diffusion/layout/flux2_token_cat_nvfp4_jit.py b/python/sglang/kernels/ops/diffusion/layout/flux2_token_cat_nvfp4_jit.py new file mode 100644 index 000000000..a65354815 --- /dev/null +++ b/python/sglang/kernels/ops/diffusion/layout/flux2_token_cat_nvfp4_jit.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +import os +from typing import TYPE_CHECKING + +import torch + +from sglang.kernels.jit.utils import cache_once, load_jit + +if TYPE_CHECKING: + from tvm_ffi.module import Module + + +_ATTENTION_HIDDEN = 6144 +_MLP_HIDDEN = 18432 +_OUTPUT_HIDDEN = _ATTENTION_HIDDEN + _MLP_HIDDEN +_ALIGNMENT = 32 +_MAX_ROWS = 65408 # Largest multiple of 128 accepted by CUDA grid.y. + + +def _env_enabled(name: str) -> bool: + return os.getenv(name, "").strip().lower() not in {"", "0", "false", "off", "no"} + + +def _is_dense_bf16(x: torch.Tensor, hidden: int) -> bool: + return ( + isinstance(x, torch.Tensor) + and x.is_cuda + and x.dtype == torch.bfloat16 + and x.dim() == 3 + and x.shape[0] == 1 + and x.shape[-1] == hidden + and x.is_contiguous() + and x.numel() > 0 + and x.data_ptr() % _ALIGNMENT == 0 + ) + + +@cache_once +def _module() -> Module: + return load_jit( + "flux2_token_cat_nvfp4", + cuda_files=["diffusion/flux2_token_cat_nvfp4.cuh"], + cuda_wrappers=[("run", "flux2_token_cat_nvfp4::Kernel::run")], + extra_cuda_cflags=["-DENABLE_BF16", "-DENABLE_FP4"], + extra_dependencies=["flashinfer", "flashinfer_nv_internal"], + ) + + +def try_flux2_token_cat_nvfp4( + attention: torch.Tensor, + mlp: torch.Tensor, + global_scale: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor] | None: + """Concatenate FLUX.2 single-block branches directly into NVFP4.""" + if ( + torch.compiler.is_compiling() + or not _is_dense_bf16(attention, _ATTENTION_HIDDEN) + or not _is_dense_bf16(mlp, _MLP_HIDDEN) + or mlp.device != attention.device + or attention.shape[:-1] != mlp.shape[:-1] + or torch.cuda.is_current_stream_capturing() + or _env_enabled("FLASHINFER_DISABLE_FP4_QUANT_FAST_MATH") + or _env_enabled("TRTLLM_DISABLE_FP4_QUANT_FAST_MATH") + or _env_enabled("FLASHINFER_NVFP4_4OVER6") + or torch.cuda.get_device_capability(attention.device) != (10, 3) + ): + return None + if not ( + isinstance(global_scale, torch.Tensor) + and global_scale.is_cuda + and global_scale.device == attention.device + and global_scale.dtype == torch.float32 + and global_scale.numel() == 1 + and global_scale.is_contiguous() + ): + return None + + rows = attention.numel() // _ATTENTION_HIDDEN + if rows > _MAX_ROWS: + return None + padded_rows = (rows + 127) // 128 * 128 + quantized = torch.empty( + (rows, _OUTPUT_HIDDEN // 2), dtype=torch.uint8, device=attention.device + ) + quant_scales = torch.empty( + (padded_rows, _OUTPUT_HIDDEN // 16), + dtype=torch.uint8, + device=attention.device, + ) + _module().run( + quantized, + quant_scales, + attention.view(-1, _ATTENTION_HIDDEN), + mlp.view(-1, _MLP_HIDDEN), + global_scale.reshape(1), + ) + return quantized, quant_scales + + +__all__ = ["try_flux2_token_cat_nvfp4"] 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 1b703f50d..b11684405 100755 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py @@ -741,3 +741,37 @@ class ModelOptFp4LinearMethod(LinearMethodBase): if bias is not None: out = out + bias return out.view(*output_shape) + + +def apply_nvfp4_gemm_prequantized( + layer: torch.nn.Module, + x_fp4: torch.Tensor, + x_scale_interleaved: torch.Tensor, + output_dtype: torch.dtype, + bias: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Run a ModelOpt NVFP4 GEMM from already packed activations and scales.""" + 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 x_scale_interleaved.dtype == torch.uint8: + x_scale_interleaved = x_scale_interleaved.view(torch.float8_e4m3fn) + if w_scale_interleaved.dtype == torch.uint8: + w_scale_interleaved = w_scale_interleaved.view(torch.float8_e4m3fn) + + fp4_gemm, flashinfer_backend = _get_fp4_gemm_op() + if fp4_gemm is None: + raise RuntimeError("No FP4 GEMM kernel available. Install flashinfer.") + out = fp4_gemm( + x_fp4, + w.T, + x_scale_interleaved, + w_scale_interleaved.T, + layer.alpha, + output_dtype, + backend=flashinfer_backend, + ) + out = slice_nvfp4_output(out, layer.output_size_per_partition) + return out + bias if bias is not None else out diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py index 50d0561e4..588194eb6 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py @@ -28,6 +28,7 @@ from sglang.kernels.ops.diffusion import ( fused_packed_silu_mul_bitexact, is_plain_layer_norm, residual_gate_add, + try_flux2_token_cat_nvfp4, ) from sglang.multimodal_gen.configs.models.dits.flux import FluxConfig from sglang.multimodal_gen.runtime.distributed import ( @@ -58,6 +59,8 @@ from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config impor ) from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import ( ModelOptFp4Config, + ModelOptFp4LinearMethod, + apply_nvfp4_gemm_prequantized, ) from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( NDRotaryEmbedding, @@ -575,6 +578,15 @@ class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin): quant_config=quant_config, prefix=f"{prefix}.to_out" if prefix else "to_out", ) + self._enable_nvfp4_token_cat = False + capability = current_platform.get_device_capability() + if ( + self.tp_size == 1 + and capability is not None + and (capability.major, capability.minor) == (10, 3) + and isinstance(self.to_out.quant_method, ModelOptFp4LinearMethod) + ): + self._enable_nvfp4_token_cat = True if self.tp_size > 1: self._patch_to_out_weight_loader() @@ -676,9 +688,25 @@ class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin): # Handle the feedforward (FF) logic mlp_hidden_states = self.mlp_act_fn(mlp_hidden_states) - # Concatenate and parallel output projection - hidden_states = torch.cat([hidden_states, mlp_hidden_states], dim=-1) - hidden_states, _ = self.to_out(hidden_states) + # Concatenate and parallel output projection. On SM103 NVFP4 the + # producer writes the concatenated packed values and swizzled scales + # directly, avoiding a full-width BF16 cat materialization. + output_shape = (*hidden_states.shape[:-1], self.out_dim) + packed = None + if self._enable_nvfp4_token_cat: + packed = try_flux2_token_cat_nvfp4( + hidden_states, mlp_hidden_states, self.to_out.input_scale_inv + ) + if packed is None: + hidden_states = torch.cat([hidden_states, mlp_hidden_states], dim=-1) + hidden_states, _ = self.to_out(hidden_states) + else: + hidden_states = apply_nvfp4_gemm_prequantized( + self.to_out, + *packed, + output_dtype=hidden_states.dtype, + bias=self.to_out.bias, + ).view(*output_shape) return hidden_states diff --git a/test/manual/kernels/bench_flux2_token_cat_nvfp4.py b/test/manual/kernels/bench_flux2_token_cat_nvfp4.py new file mode 100644 index 000000000..d1d398ba0 --- /dev/null +++ b/test/manual/kernels/bench_flux2_token_cat_nvfp4.py @@ -0,0 +1,90 @@ +import time + +import flashinfer +import torch + +from sglang.kernels.ops.diffusion import try_flux2_token_cat_nvfp4 + + +def _benchmark(fn, iterations: int = 100) -> float: + for _ in range(10): + fn() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iterations): + fn() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) * 1000 / iterations + + +def _benchmark_wall(fn, iterations: int = 100) -> float: + for _ in range(10): + fn() + torch.cuda.synchronize() + start = time.perf_counter_ns() + for _ in range(iterations): + fn() + torch.cuda.synchronize() + return (time.perf_counter_ns() - start) / iterations / 1000 + + +def _run_case(token_count: int) -> None: + generator = torch.Generator(device="cuda") + generator.manual_seed(20260830 + token_count) + attention = torch.randn( + 1, + token_count, + 6144, + device="cuda", + dtype=torch.bfloat16, + generator=generator, + ) + mlp = torch.randn( + 1, + token_count, + 18432, + device="cuda", + dtype=torch.bfloat16, + generator=generator, + ) + global_scale = torch.tensor(0.625, device="cuda", dtype=torch.float32) + + def baseline(): + return flashinfer.fp4_quantize( + torch.cat([attention, mlp], dim=-1).view(-1, 24576), global_scale + ) + + def fused(): + result = try_flux2_token_cat_nvfp4(attention, mlp, global_scale) + assert result is not None + return result + + expected = baseline() + actual = fused() + exact = [torch.equal(lhs, rhs) for lhs, rhs in zip(actual, expected)] + baseline_us = _benchmark(baseline) + fused_us = _benchmark(fused) + baseline_wall_us = _benchmark_wall(baseline) + fused_wall_us = _benchmark_wall(fused) + print( + { + "tokens": token_count, + "baseline_us": baseline_us, + "fused_us": fused_us, + "speedup": baseline_us / fused_us, + "baseline_wall_us": baseline_wall_us, + "fused_wall_us": fused_wall_us, + "wall_speedup": baseline_wall_us / fused_wall_us, + "exact": exact, + } + ) + + +if __name__ == "__main__": + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 3): + raise RuntimeError("This benchmark requires an NVIDIA Blackwell SM103 GPU") + for tokens in (17, 512, 4096, 4608): + _run_case(tokens) diff --git a/test/registered/kernels/ops/diffusion/test_model_fast_paths.py b/test/registered/kernels/ops/diffusion/test_model_fast_paths.py index 2bf36dd64..ab7b2509e 100644 --- a/test/registered/kernels/ops/diffusion/test_model_fast_paths.py +++ b/test/registered/kernels/ops/diffusion/test_model_fast_paths.py @@ -51,6 +51,7 @@ from sglang.kernels.ops.diffusion import ( mount_fused_ln_modulate, mount_hunyuan_qknorm, mount_ltx2_rms_norm_modulate, + try_flux2_token_cat_nvfp4, unmount_hunyuan_qknorm, unmount_ltx2_rms_norm_modulate, wan_rmsnorm_silu, @@ -371,6 +372,40 @@ class TestFlux2EagerFusions(CustomTestCase): self.assertTrue(torch.equal(actual, expected)) self.assertEqual(len(flux2._FLUX2_SWIGLU_SIGS), 1) + @pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 3), + reason="FLUX.2 token-cat NVFP4 requires Blackwell SM103", + ) + def test_token_cat_nvfp4_matches_flashinfer(self): + import flashinfer + + torch.manual_seed(20260830) + attention = torch.randn(1, 17, 6144, device="cuda", dtype=torch.bfloat16) + mlp = torch.randn(1, 17, 18432, device="cuda", dtype=torch.bfloat16) + global_scale = torch.tensor(0.625, device="cuda", dtype=torch.float32) + + expected_fp4, expected_scales = flashinfer.fp4_quantize( + torch.cat([attention, mlp], dim=-1).view(-1, 24576), global_scale + ) + actual = try_flux2_token_cat_nvfp4(attention, mlp, global_scale) + + self.assertIsNotNone(actual) + actual_fp4, actual_scales = actual + self.assertTrue(torch.equal(actual_fp4, expected_fp4)) + self.assertTrue( + torch.equal( + actual_scales.view(torch.uint8), expected_scales.view(torch.uint8) + ) + ) + + def test_token_cat_nvfp4_falls_back_while_compiling(self): + attention = torch.empty(1, 1, 6144, device="cuda", dtype=torch.bfloat16) + mlp = torch.empty(1, 1, 18432, device="cuda", dtype=torch.bfloat16) + global_scale = torch.ones(1, device="cuda", dtype=torch.float32) + + with patch("torch.compiler.is_compiling", return_value=True): + self.assertIsNone(try_flux2_token_cat_nvfp4(attention, mlp, global_scale)) + # ------------------------------------------------------------------------- # Qwen-Image -- reuse timestep-only modulation across serial CFG branches diff --git a/test/registered/kernels/ops/quantization/test_diffusion_nvfp4_scaled_mm.py b/test/registered/kernels/ops/quantization/test_diffusion_nvfp4_scaled_mm.py index f7e15ef90..873a90954 100644 --- a/test/registered/kernels/ops/quantization/test_diffusion_nvfp4_scaled_mm.py +++ b/test/registered/kernels/ops/quantization/test_diffusion_nvfp4_scaled_mm.py @@ -10,6 +10,7 @@ from sglang.multimodal_gen.runtime.layers.quantization import ( from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import ( ModelOptFp4Config, ModelOptFp4LinearMethod, + apply_nvfp4_gemm_prequantized, ) from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.srt.layers.quantization.modelopt_quant import pad_nvfp4_weight @@ -363,6 +364,40 @@ 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", +) +@pytest.mark.parametrize( + "backend", [None, "flashinfer_trtllm"], ids=["default", "flashinfer_trtllm"] +) +def test_prequantized_input_matches_regular_apply( + monkeypatch: pytest.MonkeyPatch, backend: str | None +) -> None: + _set_diffusion_fp4_backend(monkeypatch, backend) + m, n, k = 19, 150, 80 + generator = torch.Generator(device=DEVICE) + generator.manual_seed(20260831) + 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) + weight_fp4, weight_scale_linear = _quantize_weight_for_checkpoint( + weight, weight_global_scale + ) + method, layer = _build_layer( + weight_fp4, weight_scale_linear, input_global_scale, weight_global_scale + ) + + expected = method.apply(layer, x) + x_fp4, x_scale_interleaved = flashinfer.fp4_quantize(x, input_global_scale) + actual = apply_nvfp4_gemm_prequantized( + layer, x_fp4, x_scale_interleaved, output_dtype=x.dtype + ) + + assert torch.equal(actual, expected) + + @pytest.mark.skipif( not _nvfp4_supported(), reason="Diffusion NVFP4 scaled mm correctness requires Blackwell GPUs",