[Diffusion] Fuse FLUX.2 token concatenation and NVFP4 quantization (#37141)

2026-08-31 21:33:02 +08:00
parent d60d658f5f
commit 52e1c24744
9 changed files with 474 additions and 3 deletions
@@ -0,0 +1,120 @@
+// FLUX.2 single-block [attention | MLP] concatenation + NVFP4 quantization.
+
+#pragma once
+
+#include <sgl_kernel/tensor.h>
+
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+
+#ifndef FLT_MAX
+#define FLT_MAX __FLT_MAX__
+#endif
+#include <tensorrt_llm/kernels/quantization_utils.cuh>
+
+#include <cstdint>
+
+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<bf16_t, kGroupSize>;
+
+  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<uint8_t*>(params.quant_scales);
+  if (row >= params.num_rows) {
+    quant_scales[scale_offset] = 0;
+    return;
+  }
+
+  Vec input;
+  if (group < kAttentionGroups) {
+    input.load(static_cast<const bf16_t*>(params.attention) + int64_t(row) * kAttentionHidden + group * kGroupSize);
+  } else {
+    const int mlp_group = group - kAttentionGroups;
+    input.load(static_cast<const bf16_t*>(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<const float*>(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<uint64_t*>(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<kDLCUDA>();
+
+    TensorMatcher({rows, kAttentionHidden}).with_dtype<bf16_t>().with_device(device).verify(attention);
+    TensorMatcher({rows, kMlpHidden}).with_dtype<bf16_t>().with_device(device).verify(mlp);
+    TensorMatcher({rows, kPackedColumns}).with_dtype<uint8_t>().with_device(device).verify(quantized);
+    TensorMatcher({padded_rows, kScaleColumns}).with_dtype<uint8_t>().with_device(device).verify(quant_scales);
+    TensorMatcher({1}).with_dtype<fp32_t>().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<uint32_t>(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
@@ -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).

@@ -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",
@@ -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"]
@@ -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
@@ -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