From c0a12dbd690ef723de2f2ebd1d13f55376afa85f Mon Sep 17 00:00:00 2001 From: Anthony Chang <27950904+rosenrodt@users.noreply.github.com> Date: Sat, 8 Aug 2026 10:38:14 +0800 Subject: [PATCH 1/7] Add fused SM120 SVDQuant CuTe DSL support Extend the existing NVFP4 SVDQuant API to SM120/SM121 with a fused CuTe DSL implementation while retaining the compositional backend as a differential oracle. Changes - Add fused residual, BF16 low-rank correction, bias, and automatic SM120 dispatch - Preserve explicit cute-dsl-unfused behavior and the existing SM100 CUTLASS contract - Add focused GEMM, linear, trace-template, and generated-trace coverage Validation - Preserve the original focused SM120, SM100, and trace acceptance changes - Run applicable repository commit hooks during history reconstruction Result - SM120/SM121 callers use the existing public SVDQuant APIs - Fused and unfused implementations are independently selectable for validation --- flashinfer/gemm/gemm_svdquant.py | 381 ++++++++++++++- .../dense_blockscaled_gemm_sm120_b12x.py | 83 +++- flashinfer/trace/templates/gemm.py | 28 +- tests/gemm/test_nvfp4_svdquant_gemm.py | 434 +++++++++++++++++- tests/trace/example.py | 13 +- ..._N3072_K_packed1536_SF_B589824_rank32.json | 106 +++++ ..._K3072_K_packed1536_SF_B589824_rank32.json | 110 +++++ .../quantize_nvfp4_smooth_N3072.json | 71 +++ .../test_fi_trace_template_consistency.py | 13 + 9 files changed, 1182 insertions(+), 57 deletions(-) create mode 100644 tests/trace/fi_trace_out/gemm_nvfp4_svdquant_N3072_K_packed1536_SF_B589824_rank32.json create mode 100644 tests/trace/fi_trace_out/linear_nvfp4_svdquant_N3072_K3072_K_packed1536_SF_B589824_rank32.json create mode 100644 tests/trace/fi_trace_out/quantize_nvfp4_smooth_N3072.json diff --git a/flashinfer/gemm/gemm_svdquant.py b/flashinfer/gemm/gemm_svdquant.py index b3a5d01595a..a5a0e8ad83d 100644 --- a/flashinfer/gemm/gemm_svdquant.py +++ b/flashinfer/gemm/gemm_svdquant.py @@ -32,6 +32,7 @@ get_hybrid_num_tokens_buckets, map_to_hybrid_bucket_uncapped, ) +from ..jit.cpp_ext import get_cuda_version from ..jit.gemm import gen_gemm_sm100_module_cutlass_nvfp4_svdquant from ..trace.templates.gemm import ( mm_nvfp4_svdquant_trace, @@ -42,6 +43,7 @@ _get_cache_buf, backend_requirement, device_support_pdl, + get_device_sm_count, supported_compute_capability, ) @@ -52,6 +54,8 @@ # rank granularity (CollectiveMmaLoRA::LoRaK); ranks 32-128 are validated. SVDQUANT_LORA_RANK_GRANULARITY = 32 +_SM120_SVDQUANT_KERNEL_CACHE: dict[tuple, object] = {} + def _pad_up(x: int, y: int) -> int: return (x + y - 1) // y * y @@ -62,6 +66,200 @@ def _swizzled_sf_size(rows: int, sf_cols: int) -> int: return _pad_up(rows, 128) * _pad_up(sf_cols, 4) +def _view_128x4_sf(sf: torch.Tensor, rows: int, sf_cols: int) -> torch.Tensor: + """Restore a flat public scale buffer to its padded 128x4 storage view.""" + size = _swizzled_sf_size(rows, sf_cols) + return sf.reshape(-1)[:size].view(_pad_up(rows, 128), _pad_up(sf_cols, 4)) + + +def _compile_sm120_nvfp4_svdquant( + *, + device: torch.device, + rank: int, + with_bias: bool, + mma_tiler_mn: Tuple[int, int], + swap_ab: bool, + sf_m: int, + sf_n: int, + sf_k: int, + enable_pdl: bool, +): + """Compile one fused SM120 NVFP4 + rank-r BF16 epilogue specialization.""" + device_index = device.index + if device_index is None: + device_index = torch.cuda.current_device() + + from ..cute_dsl.utils import get_max_active_clusters + + max_active_clusters = get_max_active_clusters(1) + cache_key = ( + device_index, + rank, + with_bias, + mma_tiler_mn, + swap_ab, + max_active_clusters, + enable_pdl, + ) + if cache_key in _SM120_SVDQUANT_KERNEL_CACHE: + return _SM120_SVDQUANT_KERNEL_CACHE[cache_key] + + import cutlass + import cutlass.cute as cute + + from cutlass.cute.runtime import make_ptr + + from ..jit.cute_dsl_core import build_and_load_cute_dsl_kernel + from .kernels import dense_blockscaled_gemm_sm120_b12x + from .kernels.dense_blockscaled_gemm_sm120_b12x import ( + Sm120B12xBlockScaledDenseGemmKernel, + ) + + gemm = Sm120B12xBlockScaledDenseGemmKernel( + 16, + mma_tiler_mn, + (1, 1), + use_prefetch=False, + enable_pdl=enable_pdl, + swap_ab=swap_ab, + ) + + def compile_kernel(): + sym_m = cute.sym_int() + sym_k = cute.sym_int() + sym_n = cute.sym_int() + a_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Uint8, + (sym_m, sym_k), + stride_order=(1, 0), + assumed_align=32, + ) + b_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Uint8, + (sym_n, sym_k), + stride_order=(1, 0), + assumed_align=32, + ) + c_fake = cute.runtime.make_fake_compact_tensor( + cutlass.BFloat16, + (sym_m, sym_n), + stride_order=(1, 0), + assumed_align=16, + ) + d_fake = cute.runtime.make_fake_compact_tensor( + cutlass.BFloat16, + (sym_m, rank), + stride_order=(1, 0), + assumed_align=16, + ) + l1_fake = cute.runtime.make_fake_compact_tensor( + cutlass.BFloat16, + (sym_n, rank), + stride_order=(1, 0), + assumed_align=16, + ) + bias_fake = ( + cute.runtime.make_fake_compact_tensor( + cutlass.BFloat16, (sym_n,), assumed_align=16 + ) + if with_bias + else None + ) + a_sf_ptr = make_ptr(cutlass.Float8E4M3FN, 16, cute.AddressSpace.gmem, 16) + b_sf_ptr = make_ptr(cutlass.Float8E4M3FN, 16, cute.AddressSpace.gmem, 16) + alpha_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, (1,), assumed_align=4 + ) + stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + return cute.compile( + gemm.wrapper, + a_fake, + b_fake, + c_fake, + sf_m, + sf_n, + sf_k, + 1, + a_sf_ptr, + b_sf_ptr, + alpha_fake, + max_active_clusters, + stream_fake, + False, + d_fake, + l1_fake, + bias_fake, + options="--opt-level 2 --enable-tvm-ffi", + ) + + kernel_name = ( + f"r{rank}_bias{int(with_bias)}_t{mma_tiler_mn[0]}x{mma_tiler_mn[1]}" + f"_swap{int(swap_ab)}_mac{max_active_clusters}_pdl{int(enable_pdl)}" + ) + compiled = build_and_load_cute_dsl_kernel( + "mm_nvfp4_svdquant_sm120", + kernel_name, + compile_kernel, + extra_key_files=(__file__, dense_blockscaled_gemm_sm120_b12x.__file__), + ) + _SM120_SVDQUANT_KERNEL_CACHE[cache_key] = compiled + return compiled + + +def _mm_nvfp4_svdquant_sm120_fused( + a: torch.Tensor, + b: torch.Tensor, + a_sf: torch.Tensor, + b_sf: torch.Tensor, + alpha: torch.Tensor, + d: torch.Tensor, + l1: torch.Tensor, + bias: Optional[torch.Tensor], + out: torch.Tensor, + enable_pdl: bool, +) -> torch.Tensor: + from .kernels.dense_blockscaled_gemm_sm120_b12x import ( + _select_default_dense_gemm_plan, + ) + + m, k_packed = a.shape + n = b.shape[0] + real_k = k_packed * 2 + sf_m = (m + 127) // 128 + sf_n = (n + 127) // 128 + sf_k = (real_k // 16 + 3) // 4 + plan = _select_default_dense_gemm_plan( + m, n, real_k, get_device_sm_count(a.device), expected_m=m + ) + compiled = _compile_sm120_nvfp4_svdquant( + device=a.device, + rank=d.shape[1], + with_bias=bias is not None, + mma_tiler_mn=plan.mma_tiler_mn, + swap_ab=plan.swap_ab, + sf_m=sf_m, + sf_n=sf_n, + sf_k=sf_k, + enable_pdl=enable_pdl, + ) + args = [ + a, + b, + out, + sf_m, + sf_n, + sf_k, + a_sf.data_ptr(), + b_sf.data_ptr(), + alpha.reshape(1), + d, + l1, + bias, + ] + compiled(*args) + return out + + @functools.cache def get_nvfp4_svdquant_module(): """JIT-build and load the SM100 CUTLASS NVFP4 SVDQuant module.""" @@ -147,6 +345,31 @@ def _cutlass_nvfp4_svdquant_requirement(*args, **kwargs): return True +@supported_compute_capability([120, 121]) +def _cute_dsl_nvfp4_svdquant_requirement(*args, **kwargs): + if get_cuda_version().major < 13: + raise ValueError( + "SM120 SVDQuant CuTe DSL support requires CUDA 13 or later. " + f"Current CUDA version: {get_cuda_version()}." + ) + from ..cute_dsl import is_cute_dsl_available + + if not is_cute_dsl_available(): + raise ValueError( + "SM120 SVDQuant CuTe DSL support requires CuTe DSL, but it is not " + "available in the current environment." + ) + return True + + +def _heuristic_func_nvfp4_svdquant( + suitable_backends: List[str], *args, **kwargs +) -> List[str]: + # The backend requirements are architecture-disjoint, so retaining their order + # is sufficient and keeps automatic dispatch deterministic. + return suitable_backends + + def _check_mm_nvfp4_svdquant_problem( a: torch.Tensor, b: torch.Tensor, @@ -157,13 +380,15 @@ def _check_mm_nvfp4_svdquant_problem( l1: torch.Tensor, bias: Optional[torch.Tensor] = None, out: Optional[torch.Tensor] = None, - backend: Literal["cutlass"] = "cutlass", + backend: Literal["cutlass", "cute-dsl", "cute-dsl-unfused", "auto"] = "auto", enable_pdl: Optional[bool] = None, ): if a.ndim != 2 or b.ndim != 2: raise ValueError("a and b must be 2-D packed-e2m1 (uint8) tensors") if a.dtype != torch.uint8 or b.dtype != torch.uint8: raise ValueError("a and b must be uint8 (two e2m1 values per byte)") + if not a.is_contiguous() or not b.is_contiguous(): + raise ValueError("a and b must be contiguous") m, k_packed = a.shape n = b.shape[0] if b.shape[1] != k_packed: @@ -190,12 +415,43 @@ def _check_mm_nvfp4_svdquant_problem( ) if d.dtype != torch.bfloat16 or l1.dtype != torch.bfloat16: raise ValueError("d and l1 must be bf16") + if not d.is_contiguous() or not l1.is_contiguous(): + raise ValueError("d and l1 must be contiguous") + if a_sf.dtype != torch.uint8 or b_sf.dtype != torch.uint8: + raise ValueError("a_sf and b_sf must be uint8 (ue4m3 block scales)") + expected_a_sf = _swizzled_sf_size(m, k // 16) + expected_b_sf = _swizzled_sf_size(n, k // 16) + if a_sf.numel() < expected_a_sf or b_sf.numel() < expected_b_sf: + raise ValueError( + "128x4 scale buffers are too small: " + f"a_sf has {a_sf.numel()} elements (need {expected_a_sf}), " + f"b_sf has {b_sf.numel()} elements (need {expected_b_sf})" + ) + if not a_sf.is_contiguous() or not b_sf.is_contiguous(): + raise ValueError("a_sf and b_sf must be contiguous") + if alpha.dtype != torch.float32 or alpha.numel() != 1: + raise ValueError("alpha must be a float32 device scalar") + if bias is not None and (bias.shape != (n,) or bias.dtype != torch.bfloat16): + raise ValueError(f"bias must have shape ({n},) and dtype bf16") + if out is not None: + if out.shape != (m, n) or out.dtype != torch.bfloat16: + raise ValueError(f"out must have shape ({m}, {n}) and dtype bf16") + if not out.is_contiguous(): + raise ValueError("out must be contiguous") + tensors = (b, a_sf, b_sf, alpha, d, l1, bias, out) + if any(t is not None and t.device != a.device for t in tensors): + raise ValueError("all SVDQuant tensors must be on the same device") return True @backend_requirement( - {"cutlass": _cutlass_nvfp4_svdquant_requirement}, + { + "cutlass": _cutlass_nvfp4_svdquant_requirement, + "cute-dsl": _cute_dsl_nvfp4_svdquant_requirement, + "cute-dsl-unfused": _cute_dsl_nvfp4_svdquant_requirement, + }, common_check=_check_mm_nvfp4_svdquant_problem, + heuristic_func=_heuristic_func_nvfp4_svdquant, ) @flashinfer_api(trace=mm_nvfp4_svdquant_trace) def mm_nvfp4_svdquant( @@ -208,18 +464,19 @@ def mm_nvfp4_svdquant( l1: torch.Tensor, bias: Optional[torch.Tensor] = None, out: Optional[torch.Tensor] = None, - backend: Literal["cutlass"] = "cutlass", + backend: Literal["cutlass", "cute-dsl", "cute-dsl-unfused", "auto"] = "auto", enable_pdl: Optional[bool] = None, ) -> torch.Tensor: - r"""SVDQuant fused NVFP4 GEMM (SM100): ``out = alpha * (a @ bᵀ) + d @ l1ᵀ [+ bias]``. + r"""SVDQuant NVFP4 GEMM: ``out = alpha * (a @ bᵀ + d @ l1ᵀ) [+ bias]``. - The block-scaled NVFP4 residual GEMM is fused with the rank-r BF16 LoRA-up correction - ``d @ l1ᵀ``, computed by a second BF16 tcgen05 MMA into the same accumulator after the - NVFP4 K-loop, plus an optional fused per-column bias. The LoRA rank ``r`` is inferred - from the ``d``/``l1`` shapes and must be a positive multiple of 32 (ranks 32-128 are - validated). ``1/alpha`` must be folded into ``l1`` by the caller - (``l1 = svdquant_lora_b / alpha``) so the epilogue ``out = alpha * acc + bias`` yields - the correction unscaled. + On SM100/SM103, CUTLASS fuses the block-scaled NVFP4 residual GEMM with the rank-r + BF16 LoRA-up correction and optional bias. On SM120/SM121, ``"cute-dsl"`` fuses the + correction and bias into the b12x CuTe DSL kernel's FP32 accumulator epilogue, while + ``"cute-dsl-unfused"`` retains the compositional implementation as a differential + oracle and fallback. The LoRA rank ``r`` is inferred from the ``d``/``l1`` shapes and must + be a positive multiple of 32 (ranks 32-128 are validated). ``1/alpha`` must be folded + into ``l1`` by the caller (``l1 = svdquant_lora_b / alpha``), so both backends yield + the correction at its original scale. Parameters ---------- @@ -243,11 +500,15 @@ def mm_nvfp4_svdquant( l1: torch.Tensor LoRA-up weight pre-divided by alpha, shape ``(n, r)`` bf16 (same rank as ``d``). bias: Optional[torch.Tensor] - Optional per-column bias, shape ``(n,)`` bf16, fused in the epilogue. + Optional per-column bias, shape ``(n,)`` bf16. Fused by CUTLASS and the + SM120/SM121 CuTe DSL kernel. out: Optional[torch.Tensor] Output tensor, shape ``(m, n)`` bf16; allocated when ``None``. - backend: Literal["cutlass"] - Only the CUTLASS backend exists. + backend: Literal["cutlass", "cute-dsl", "cute-dsl-unfused", "auto"] + ``"cutlass"`` selects the fused SM100/SM103 implementation; + ``"cute-dsl"`` selects the fused SM120/SM121 implementation; + ``"cute-dsl-unfused"`` selects its compositional reference path; + ``"auto"`` (default) selects by compute capability. enable_pdl: Optional[bool] Whether to launch with Programmatic Dependent Launch. Defaults to the device default. @@ -260,6 +521,44 @@ def mm_nvfp4_svdquant( enable_pdl = device_support_pdl(a.device) if out is None: out = torch.empty(a.shape[0], b.shape[0], dtype=torch.bfloat16, device=a.device) + + if backend == "auto": + backend = mm_nvfp4_svdquant.suitable_auto_backends[0] + + if backend == "cute-dsl": + return _mm_nvfp4_svdquant_sm120_fused( + a, b, a_sf, b_sf, alpha, d, l1, bias, out, enable_pdl + ) + + if backend == "cute-dsl-unfused": + from .gemm_base import mm_fp4 + + m, k_packed = a.shape + n = b.shape[0] + sf_cols = k_packed * 2 // 16 + a_sf_2d = _view_128x4_sf(a_sf, m, sf_cols) + b_sf_2d = _view_128x4_sf(b_sf, n, sf_cols) + mm_fp4( + a, + b.T, + a_sf_2d, + b_sf_2d.T, + alpha, + torch.bfloat16, + out, + block_size=16, + use_8x4_sf_layout=False, + backend="b12x", + use_nvfp4=True, + enable_pdl=enable_pdl, + ) + correction = torch.mm(d, l1.T) + correction.mul_(alpha) + out.add_(correction) + if bias is not None: + out.add_(bias) + return out + workspace_buffer = _get_cache_buf( "nvfp4_svdquant_gemm_workspace", DEFAULT_WORKSPACE_SIZE, a.device ) @@ -282,7 +581,7 @@ def _check_nvfp4_quantize_smooth_problem( pre_quant_scale: torch.Tensor, global_scale: torch.Tensor, enable_pdl: Optional[bool] = None, - backend: Literal["cutlass"] = "cutlass", + backend: Literal["cutlass", "cute-dsl", "auto"] = "auto", ): if x.ndim != 2: raise ValueError(f"x must be [m, n], got {tuple(x.shape)}") @@ -300,8 +599,12 @@ def _check_nvfp4_quantize_smooth_problem( @backend_requirement( - {"cutlass": _cutlass_nvfp4_svdquant_requirement}, + { + "cutlass": _cutlass_nvfp4_svdquant_requirement, + "cute-dsl": _cute_dsl_nvfp4_svdquant_requirement, + }, common_check=_check_nvfp4_quantize_smooth_problem, + heuristic_func=_heuristic_func_nvfp4_svdquant, ) @flashinfer_api(trace=nvfp4_quantize_smooth_trace) def nvfp4_quantize_smooth( @@ -309,13 +612,14 @@ def nvfp4_quantize_smooth( pre_quant_scale: torch.Tensor, global_scale: torch.Tensor, enable_pdl: Optional[bool] = None, - backend: Literal["cutlass"] = "cutlass", + backend: Literal["cutlass", "cute-dsl", "auto"] = "auto", ) -> Tuple[torch.Tensor, torch.Tensor]: - r"""Fused smooth + NVFP4 quantize: ``(xq, sf) = nvfp4-quantize(x * pre_quant_scale)``. + r"""Smooth + NVFP4 quantize: ``(xq, sf) = nvfp4-quantize(x * pre_quant_scale)``. - Applies the SVDQuant per-input-channel smoothing scale and NVFP4-quantizes in one pass - over the input; the result is byte-identical to quantizing ``x * pre_quant_scale`` with - the stock NVFP4 quantizer (ue4m3 block scales, 128x4 swizzled layout, SF vector size 16). + The SM100/SM103 CUTLASS backend applies the SVDQuant per-input-channel smoothing scale + and NVFP4-quantizes in one pass. The SM120/SM121 CuTe DSL compatibility backend first + materializes the BF16 smoothed input and then invokes the CuTe DSL NVFP4 quantizer. + Both use ue4m3 block scales, the 128x4 swizzled layout, and SF vector size 16. Parameters ---------- @@ -327,8 +631,10 @@ def nvfp4_quantize_smooth( Global scale, float32 device scalar: ``(448 * 6) / (x * pre_quant_scale).abs().max()``. enable_pdl: Optional[bool] Whether to launch with Programmatic Dependent Launch. Defaults to the device default. - backend: Literal["cutlass"] - Only the CUDA backend exists. + backend: Literal["cutlass", "cute-dsl", "auto"] + ``"cutlass"`` selects fused smoothing and quantization on SM100/SM103; + ``"cute-dsl"`` selects unfused smoothing plus CuTe DSL quantization on + SM120/SM121; ``"auto"`` (default) selects by compute capability. Returns ------- @@ -340,6 +646,23 @@ def nvfp4_quantize_smooth( """ if enable_pdl is None: enable_pdl = device_support_pdl(x.device) + if backend == "auto": + backend = nvfp4_quantize_smooth.suitable_auto_backends[0] + if backend == "cute-dsl": + from ..quantization.fp4_quantization import nvfp4_quantize + from ..tllm_enums import SfLayout + + xq, sf = nvfp4_quantize( + (x * pre_quant_scale).to(torch.bfloat16), + global_scale, + sfLayout=SfLayout.layout_128x4, + do_shuffle=False, + sf_vec_size=16, + enable_pdl=enable_pdl, + backend="cute-dsl", + ) + return xq.view(torch.uint8), sf.view(torch.uint8).reshape(-1) + m, n = x.shape module = get_nvfp4_svdquant_module() xq = torch.empty(m, n // 2, dtype=torch.uint8, device=x.device) @@ -360,6 +683,7 @@ def svdquant_linear( global_scale: torch.Tensor, bias: Optional[torch.Tensor] = None, enable_pdl: Optional[bool] = None, + backend: Literal["cutlass", "cute-dsl", "cute-dsl-unfused", "auto"] = "auto", ) -> torch.Tensor: r"""The full SVDQuant linear operator: ``y = x_hat @ (R + L1 @ L2)ᵀ [+ bias]`` where ``x_hat = x * pre_quant_scale`` and ``R`` is the NVFP4-quantized residual weight. @@ -397,14 +721,22 @@ def svdquant_linear( Optional per-column bias, shape ``(n,)`` bf16. enable_pdl: Optional[bool] Whether to launch with Programmatic Dependent Launch. Defaults to the device default. + backend: Literal["cutlass", "cute-dsl", "cute-dsl-unfused", "auto"] + Backend forwarded to smooth quantization and SVDQuant GEMM. Defaults to + architecture-based automatic selection. Returns ------- out: torch.Tensor Output tensor, shape ``(m, n)`` bf16. """ + quantize_backend = "cute-dsl" if backend == "cute-dsl-unfused" else backend xq, x_sf = nvfp4_quantize_smooth( - x, pre_quant_scale, global_scale, enable_pdl=enable_pdl + x, + pre_quant_scale, + global_scale, + enable_pdl=enable_pdl, + backend=quantize_backend, ) down = torch.mm(x, l2t_smoothed) return mm_nvfp4_svdquant( @@ -417,4 +749,5 @@ def svdquant_linear( l1_scaled, bias=bias, enable_pdl=enable_pdl, + backend=backend, ) diff --git a/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py b/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py index c779f3ec21c..41b835aca44 100644 --- a/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py +++ b/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py @@ -41,7 +41,7 @@ import cutlass.utils.blockscaled_layout as blockscaled_utils import cutlass.utils.hopper_helpers as sm90_utils import logging -from cutlass import Int32, Int64 +from cutlass import Float32, Int32, Int64 from cutlass.cute.arch import griddepcontrol_launch_dependents, griddepcontrol_wait from cutlass.cute.nvgpu import cpasync from cutlass.cute.nvgpu.warp.mma import Field as WarpField @@ -409,6 +409,9 @@ def __call__( alpha: cute.Tensor, max_active_clusters: cutlass.Constexpr, stream: cuda.CUstream, + svdquant_d: Optional[cute.Tensor] = None, + svdquant_l1: Optional[cute.Tensor] = None, + svdquant_bias: Optional[cute.Tensor] = None, epilogue_op: cutlass.Constexpr = lambda x: x, ): """Execute the GEMM operation. @@ -423,6 +426,9 @@ def __call__( max_active_clusters: Max active clusters stream: CUDA stream epilogue_op: Elementwise epilogue function + svdquant_d: Optional BF16 LoRA-down output, shape (M, rank) + svdquant_l1: Optional BF16 scaled LoRA-up weight, shape (N, rank) + svdquant_bias: Optional BF16 per-column bias, shape (N,) """ # Setup static attributes self.a_dtype = a.element_type @@ -560,6 +566,9 @@ class SharedStorage: tile_sched_params, epilogue_op, alpha, + svdquant_d, + svdquant_l1, + svdquant_bias, ).launch( grid=grid, block=[self.threads_per_cta, 1, 1], @@ -907,6 +916,9 @@ def kernel( tile_sched_params: utils.PersistentTileSchedulerParams, epilogue_op: cutlass.Constexpr, alpha: cute.Tensor, + svdquant_d: Optional[cute.Tensor], + svdquant_l1: Optional[cute.Tensor], + svdquant_bias: Optional[cute.Tensor], ): # Keep alpha in FP32 for precision alpha_value = alpha[0].to(cutlass.Float32) @@ -1596,6 +1608,69 @@ def kernel( accumulators[None, _mt, _nt], ) + # SVDQuant fusion: accumulate the rank-r BF16 correction into the + # same FP32 registers as the NVFP4 residual before the epilogue. + # The first SM120 implementation uses scalar BF16 dot products; + # this keeps one device launch and the exact public contract while + # leaving BF16 warp-MMA staging as a transparent optimization. + if cutlass.const_expr(svdquant_d is not None): + acc_mn = _reshape_acc_to_mn( + accumulators, + transpose=self.swap_ab, + ) + c_identity = cute.make_identity_tensor( + ( + self.tile_shape_mnk[1], + self.tile_shape_mnk[0], + ) + if cutlass.const_expr(self.swap_ab) + else ( + self.tile_shape_mnk[0], + self.tile_shape_mnk[1], + ) + ) + coord_mn = _reshape_acc_to_mn( + thr_mma.partition_C(c_identity), + transpose=self.swap_ab, + ) + rank = cute.size(svdquant_d, mode=[1]) + for acc_m in cutlass.range_constexpr(cute.size(acc_mn.shape[0])): + for acc_n in cutlass.range_constexpr( + cute.size(acc_mn.shape[1]) + ): + coord = coord_mn[acc_m, acc_n] + if cutlass.const_expr(self.swap_ab): + m_local = coord[1] + n_local = coord[0] + else: + m_local = coord[0] + n_local = coord[1] + m_coord = ( + tile_coord_mnl[0] * Int32(self.tile_shape_mnk[0]) + + m_local + ) + n_coord = ( + tile_coord_mnl[1] * Int32(self.tile_shape_mnk[1]) + + n_local + ) + if m_coord < Int32( + directC_mnl.shape[0] + ) and n_coord < Int32(directC_mnl.shape[1]): + correction = Float32(0.0) + for rank_idx in cutlass.range(rank, unroll=1): + correction += svdquant_d[(m_coord, rank_idx)].to( + Float32 + ) * svdquant_l1[(n_coord, rank_idx)].to(Float32) + if cutlass.const_expr(svdquant_bias is not None): + # The common epilogue multiplies the whole FP32 + # accumulator by alpha. Pre-dividing bias here + # retains out = alpha*(residual+correction)+bias. + correction += ( + svdquant_bias[(n_coord,)].to(Float32) + / alpha_value + ) + acc_mn[acc_m, acc_n] += correction + if cutlass.const_expr(self.swap_ab): acc_mn = _reshape_acc_to_mn(accumulators, transpose=True) c_identity = cute.make_identity_tensor( @@ -2544,6 +2619,9 @@ def wrapper( max_active_clusters: cutlass.Constexpr, current_stream, swap_ab: cutlass.Constexpr = False, + svdquant_d: Optional[cute.Tensor] = None, + svdquant_l1: Optional[cute.Tensor] = None, + svdquant_bias: Optional[cute.Tensor] = None, epilogue_op: cutlass.Constexpr = lambda x: x, ): """Wrapper matching the SM100 compile interface.""" @@ -2602,6 +2680,9 @@ def wrapper( alpha_tensor, max_active_clusters, current_stream, + svdquant_d, + svdquant_l1, + svdquant_bias, epilogue_op, ) diff --git a/flashinfer/trace/templates/gemm.py b/flashinfer/trace/templates/gemm.py index ce283b2096e..aeddba293e4 100644 --- a/flashinfer/trace/templates/gemm.py +++ b/flashinfer/trace/templates/gemm.py @@ -2134,7 +2134,7 @@ def _split_into_indptr(total: int) -> torch.Tensor: ) -# ── SVDQuant fused NVFP4 GEMM (SM100) ──────────────────────────────────────── +# ── SVDQuant fused NVFP4 GEMM (SM100 and SM120) ────────────────────────────── def _mm_nvfp4_svdquant_init( @@ -2142,6 +2142,7 @@ def _mm_nvfp4_svdquant_init( M: int, N: int = 3072, K: int = 3072, + SF_A: int = 0, device: str = "cuda", seed: int = 0, ): @@ -2155,6 +2156,8 @@ def _mm_nvfp4_svdquant_init( """ from flashinfer import nvfp4_quantize_smooth # noqa: PLC0415 + del SF_A # output-only / derived axis + torch.manual_seed(seed) rank = 32 x = torch.randn(M, K, dtype=torch.bfloat16, device=device) @@ -2191,15 +2194,18 @@ def _mm_nvfp4_svdquant_init( mm_nvfp4_svdquant_trace = TraceTemplate( op_type="gemm_nvfp4_svdquant", description=( - "SVDQuant fused NVFP4 GEMM (SM100): out = alpha * (a @ bᵀ) + d @ l1ᵀ. " - "The block-scaled NVFP4 residual GEMM fused with a rank-r BF16 LoRA-up " - "correction in the same accumulator; 1/alpha is pre-folded into l1." + "SVDQuant NVFP4 GEMM: out = alpha * (a @ bᵀ + d @ l1ᵀ). " + "SM100/SM103 use fused CUTLASS; SM120/SM121 use fused CuTe DSL, with " + "an explicit cute-dsl-unfused oracle. 1/alpha is " + "pre-folded into l1." ), axes={ "M": Var(), "N": Const(), "K_packed": Const(description="K / 2 (two e2m1 values per byte)."), - "SF_A": Const(description="128x4-swizzled activation scale buffer size."), + "SF_A": Var( + description="128x4-swizzled activation scale buffer size derived from M and K." + ), "SF_B": Const(description="128x4-swizzled weight scale buffer size."), "rank": Const(description="LoRA rank, a positive multiple of 32."), }, @@ -2243,6 +2249,9 @@ def _mm_nvfp4_svdquant_init( outputs={ "out": Tensor(["M", "N"], dtype="bfloat16"), }, + constraints=[ + "SF_A == ((M + 127) // 128) * 128 * (((K_packed * 2 // 16) + 3) // 4) * 4", + ], tags=["quantization:fp4"], init=_mm_nvfp4_svdquant_init, ) @@ -2274,9 +2283,10 @@ def _nvfp4_quantize_smooth_init( nvfp4_quantize_smooth_trace = TraceTemplate( op_type="quantize_nvfp4_smooth", description=( - "Fused smooth + NVFP4 quantize: (xq, sf) = nvfp4-quantize(x * pre_quant_scale). " - "Byte-identical to smoothing followed by the stock NVFP4 quantizer " - "(ue4m3 block scales, 128x4 swizzled layout, SF vector size 16)." + "Smooth + NVFP4 quantize: (xq, sf) = nvfp4-quantize(x * pre_quant_scale). " + "SM100/SM103 fuse the operations; SM120/SM121 materialize the BF16 " + "smoothed input before CuTe DSL quantization. Both use ue4m3 block " + "scales, 128x4 swizzled layout, and SF vector size 16." ), axes={ "M": Var(), @@ -2359,7 +2369,7 @@ def _svdquant_linear_init( description=( "Full SVDQuant linear: y = (x * pre_quant_scale) @ (R + L1 @ L2)ᵀ where R is the " "NVFP4-quantized residual weight — smooth-quantize, BF16 rank-r down-projection, " - "and the fused NVFP4 residual + LoRA-up GEMM." + "and the architecture-selected NVFP4 residual + LoRA-up GEMM." ), axes={ "M": Var(), diff --git a/tests/gemm/test_nvfp4_svdquant_gemm.py b/tests/gemm/test_nvfp4_svdquant_gemm.py index 71df43db3bb..353816f8f48 100644 --- a/tests/gemm/test_nvfp4_svdquant_gemm.py +++ b/tests/gemm/test_nvfp4_svdquant_gemm.py @@ -1,14 +1,12 @@ -"""Tests for the SM100 NVFP4 SVDQuant fused GEMM ops (Blackwell): +"""Tests for the NVFP4 SVDQuant GEMM ops (Blackwell): -- mm_nvfp4_svdquant : out = alpha * (a @ bT) + d @ l1T [+ bias], the block-scaled - NVFP4 residual GEMM fused with the rank-r BF16 LoRA-up - (r a positive multiple of 32; 32-128 covered here). -- nvfp4_quantize_smooth : NVFP4-quantize(x * pre_quant_scale), byte-identical to the - stock quantizer run on the pre-smoothed input. -- svdquant_linear : the full quantize -> LoRA-down -> fused GEMM chain. +- mm_nvfp4_svdquant : out = alpha * (a @ bT) + d @ l1T [+ bias], with rank r a + positive multiple of 32 (32-128 covered here). +- nvfp4_quantize_smooth : NVFP4-quantize(x * pre_quant_scale). +- svdquant_linear : the full quantize -> LoRA-down -> residual/correction chain. -The unfused reference for the residual is flashinfer.mm_fp4 (cutlass backend) on the -same quantized operands, plus the LoRA correction computed in fp32. +SM100/SM103 use fused CUTLASS. SM120/SM121 use fused CuTe DSL, with an explicit +``cute-dsl-unfused`` composition retained as a differential oracle. """ import pytest @@ -33,6 +31,19 @@ _RANK = SVDQUANT_LORA_RANK_GRANULARITY # base rank == the collective's rank granularity +def test_nvfp4_svdquant_backend_arch_support(): + for api in (mm_nvfp4_svdquant, nvfp4_quantize_smooth): + assert api.is_backend_supported("cutlass", 100) + assert api.is_backend_supported("cutlass", 103) + assert not api.is_backend_supported("cutlass", 120) + assert api.is_backend_supported("cute-dsl", 120) + assert api.is_backend_supported("cute-dsl", 121) + assert not api.is_backend_supported("cute-dsl", 100) + assert mm_nvfp4_svdquant.is_backend_supported("cute-dsl-unfused", 120) + assert mm_nvfp4_svdquant.is_backend_supported("cute-dsl-unfused", 121) + assert not mm_nvfp4_svdquant.is_backend_supported("cute-dsl-unfused", 100) + + def _skip_unless_sm100(): compute_capability = get_compute_capability(torch.device(device="cuda")) if compute_capability[0] != 10: @@ -42,6 +53,15 @@ def _skip_unless_sm100(): ) +def _skip_unless_sm120(): + compute_capability = get_compute_capability(torch.device(device="cuda")) + if compute_capability[0] != 12: + pytest.skip( + "SM120 SVDQuant CuTe DSL tests require SM120-class GPUs, " + f"got compute capability {compute_capability}." + ) + + def _sqnr_db(ref: torch.Tensor, got: torch.Tensor) -> float: err = (ref - got).float() noise = (err**2).mean() @@ -50,20 +70,24 @@ def _sqnr_db(ref: torch.Tensor, got: torch.Tensor) -> float: return float(10 * torch.log10((ref.float() ** 2).mean() / noise)) -def _nvfp4_quantize_128x4(t: torch.Tensor): +def _nvfp4_quantize_128x4(t: torch.Tensor, backend="cuda"): """Stock NVFP4 quantization (ue4m3 block scales, 128x4 swizzled layout). Returns (packed e2m1 uint8 [r, c/2], swizzled sf uint8 2-D, global scale f32 [1]). """ global_sf = ((448.0 * 6.0) / t.float().abs().nan_to_num().max()).reshape(1) tq, sf = nvfp4_quantize( - t, global_sf, sfLayout=SfLayout.layout_128x4, do_shuffle=False + t, + global_sf, + sfLayout=SfLayout.layout_128x4, + do_shuffle=False, + backend=backend, ) return tq.view(torch.uint8), sf.view(torch.uint8), global_sf -def _mm_fp4_residual(xq, wq, x_sf, w_sf, alpha): - """Unfused reference residual alpha * (a @ bT) via the stock cutlass NVFP4 GEMM.""" +def _mm_fp4_residual(xq, wq, x_sf, w_sf, alpha, backend="cutlass"): + """Reference residual alpha * (a @ bT) via a generic NVFP4 GEMM backend.""" out = torch.empty(xq.shape[0], wq.shape[0], dtype=torch.bfloat16, device=xq.device) mm_fp4( xq, @@ -75,18 +99,55 @@ def _mm_fp4_residual(xq, wq, x_sf, w_sf, alpha): out, block_size=16, use_8x4_sf_layout=False, - backend="cutlass", + backend=backend, use_nvfp4=True, ) return out.float() -def _make_gemm_problem(m, n, k, rank=_RANK, device="cuda"): +def _sm120_unfused_reference(p, use_bias): + """Reproduce the exact BF16 operation order of the SM120 unfused oracle.""" + out = torch.empty( + p["xq"].shape[0], + p["wq"].shape[0], + dtype=torch.bfloat16, + device=p["xq"].device, + ) + mm_fp4( + p["xq"], + p["wq"].T, + p["x_sf"], + p["w_sf"].T, + p["alpha"], + torch.bfloat16, + out, + block_size=16, + use_8x4_sf_layout=False, + backend="b12x", + use_nvfp4=True, + ) + correction = torch.mm(p["d"], p["l1_scaled"].T) + correction.mul_(p["alpha"]) + out.add_(correction) + if use_bias: + out.add_(p["bias"]) + return out + + +def _make_gemm_problem( + m, + n, + k, + rank=_RANK, + device="cuda", + quant_backend="cuda", + residual_backend="cutlass", +): """Quantized operands and fp32 references for out = alpha*(a@bT) + D@L1T [+ bias].""" x = torch.randn(m, k, dtype=torch.bfloat16, device=device) / (k**0.25) w = torch.randn(n, k, dtype=torch.bfloat16, device=device) / (k**0.25) - xq, x_sf, gx = _nvfp4_quantize_128x4(x) - wq, w_sf, gw = _nvfp4_quantize_128x4(w) + xq, x_sf, gx = _nvfp4_quantize_128x4(x, backend=quant_backend) + wq, w_sf, gw = _nvfp4_quantize_128x4(w, backend=quant_backend) alpha = (1.0 / (gx * gw)).reshape(1).float() d = torch.randn(m, rank, dtype=torch.bfloat16, device=device) / (rank**0.25) l1 = torch.randn(n, rank, dtype=torch.bfloat16, device=device) / (rank**0.25) @@ -95,7 +156,10 @@ def _make_gemm_problem(m, n, k, rank=_RANK, device="cuda"): l1_scaled = (l1.float() / alpha).to(torch.bfloat16).contiguous() bias = torch.randn(n, dtype=torch.bfloat16, device=device).contiguous() - ref = _mm_fp4_residual(xq, wq, x_sf, w_sf, alpha) + d.float() @ l1.float().t() + ref = ( + _mm_fp4_residual(xq, wq, x_sf, w_sf, alpha, backend=residual_backend) + + d.float() @ l1.float().t() + ) return { "xq": xq, "wq": wq, @@ -357,6 +421,340 @@ def test_svdquant_linear_matches_reference(use_bias, rank): assert _sqnr_db(ref, out.float()) > 40.0 +def test_nvfp4_quantize_smooth_sm120_cute_dsl(): + _skip_unless_sm120() + torch.manual_seed(0) + m, k = 129, 256 + x = torch.randn(m, k, dtype=torch.bfloat16, device="cuda") + pqs = ( + (1.0 + 0.3 * torch.randn(k, dtype=torch.bfloat16, device="cuda")) + .abs() + .contiguous() + ) + smoothed = (x * pqs).to(torch.bfloat16) + global_sf = ( + ((448.0 * 6.0) / smoothed.float().abs().nan_to_num().max()) + .reshape(1) + .contiguous() + ) + xq_ref, sf_ref = nvfp4_quantize( + smoothed, + global_sf, + sfLayout=SfLayout.layout_128x4, + do_shuffle=False, + backend="cute-dsl", + ) + + xq, sf = nvfp4_quantize_smooth(x, pqs, global_sf, backend="cute-dsl") + xq_auto, sf_auto = nvfp4_quantize_smooth(x, pqs, global_sf) + + assert xq.dtype == torch.uint8 and xq.shape == (m, k // 2) + assert sf.dtype == torch.uint8 and sf.ndim == 1 + assert torch.equal(xq, xq_ref.view(torch.uint8)) + assert torch.equal(sf, sf_ref.view(torch.uint8).reshape(-1)) + assert torch.equal(xq_auto, xq) + assert torch.equal(sf_auto, sf) + + +@pytest.mark.parametrize("use_bias", [False, True]) +def test_mm_nvfp4_svdquant_sm120_fused(use_bias): + _skip_unless_sm120() + torch.manual_seed(0) + m, n, k, rank = 129, 256, 256, 32 + p = _make_gemm_problem( + m, + n, + k, + rank=rank, + quant_backend="cute-dsl", + residual_backend="b12x", + ) + bias = p["bias"] if use_bias else None + expected = _sm120_unfused_reference(p, use_bias) + + out_buffer = torch.full((m, n), float("nan"), dtype=torch.bfloat16, device="cuda") + out = mm_nvfp4_svdquant( + p["xq"], + p["wq"], + p["x_sf_flat"], + p["w_sf_flat"], + p["alpha"], + p["d"], + p["l1_scaled"], + bias=bias, + out=out_buffer, + backend="cute-dsl", + ) + out_auto = mm_nvfp4_svdquant( + p["xq"], + p["wq"], + p["x_sf_flat"], + p["w_sf_flat"], + p["alpha"], + p["d"], + p["l1_scaled"], + bias=bias, + ) + + assert out.data_ptr() == out_buffer.data_ptr() + # The fused path accumulates the BF16 rank correction in FP32 before its + # single BF16 store, whereas the oracle rounds the residual and correction + # in separate launches. Compare numerically, not bitwise. + assert _sqnr_db(expected.float(), out.float()) > 35.0 + assert torch.equal(out_auto, out) + fp32_ref = p["ref_bias"] if use_bias else p["ref"] + assert _sqnr_db(fp32_ref, out.float()) > 35.0 + + +def test_mm_nvfp4_svdquant_sm120_unfused_oracle(): + _skip_unless_sm120() + torch.manual_seed(0) + p = _make_gemm_problem( + 33, + 128, + 128, + rank=32, + quant_backend="cute-dsl", + residual_backend="b12x", + ) + expected = _sm120_unfused_reference(p, True) + out = mm_nvfp4_svdquant( + p["xq"], + p["wq"], + p["x_sf_flat"], + p["w_sf_flat"], + p["alpha"], + p["d"], + p["l1_scaled"], + bias=p["bias"], + backend="cute-dsl-unfused", + ) + assert torch.equal(out, expected) + + +@pytest.mark.parametrize("rank", [32, 64, 96, 128]) +def test_mm_nvfp4_svdquant_sm120_fused_rank_chunks(rank): + _skip_unless_sm120() + torch.manual_seed(rank) + p = _make_gemm_problem( + 33, + 128, + 128, + rank=rank, + quant_backend="cute-dsl", + residual_backend="b12x", + ) + out = mm_nvfp4_svdquant( + p["xq"], + p["wq"], + p["x_sf_flat"], + p["w_sf_flat"], + p["alpha"], + p["d"], + p["l1_scaled"], + bias=p["bias"], + backend="cute-dsl", + ) + assert _sqnr_db(p["ref_bias"], out.float()) > 35.0 + + +@pytest.mark.parametrize( + "m,n,k,rank", + [ + (1, 512, 4096, 32), # (64, 32) tile with swap_ab=True + (33, 160, 192, 64), # partial N tile and ragged K mainloop + ], +) +def test_mm_nvfp4_svdquant_sm120_fused_boundary_plans(m, n, k, rank): + _skip_unless_sm120() + torch.manual_seed(m + n + k + rank) + p = _make_gemm_problem( + m, + n, + k, + rank=rank, + quant_backend="cute-dsl", + residual_backend="b12x", + ) + out = mm_nvfp4_svdquant( + p["xq"], + p["wq"], + p["x_sf_flat"], + p["w_sf_flat"], + p["alpha"], + p["d"], + p["l1_scaled"], + bias=p["bias"], + backend="cute-dsl", + ) + assert _sqnr_db(p["ref_bias"], out.float()) > 35.0 + + +@pytest.mark.parametrize("alpha_shape", [(), (1, 1)]) +def test_mm_nvfp4_svdquant_sm120_normalizes_alpha_shape(alpha_shape): + _skip_unless_sm120() + torch.manual_seed(0) + p = _make_gemm_problem( + 33, + 128, + 128, + rank=32, + quant_backend="cute-dsl", + residual_backend="b12x", + ) + expected = mm_nvfp4_svdquant( + p["xq"], + p["wq"], + p["x_sf_flat"], + p["w_sf_flat"], + p["alpha"], + p["d"], + p["l1_scaled"], + bias=p["bias"], + backend="cute-dsl", + ) + out = mm_nvfp4_svdquant( + p["xq"], + p["wq"], + p["x_sf_flat"], + p["w_sf_flat"], + p["alpha"].reshape(alpha_shape), + p["d"], + p["l1_scaled"], + bias=p["bias"], + backend="cute-dsl", + ) + assert torch.equal(out, expected) + + +def test_mm_nvfp4_svdquant_rejects_noncontiguous_packed_inputs(): + _skip_unless_sm120() + torch.manual_seed(0) + p = _make_gemm_problem( + 33, + 128, + 128, + rank=32, + quant_backend="cute-dsl", + residual_backend="b12x", + ) + storage = torch.empty( + p["xq"].shape[0], + p["xq"].shape[1] * 2, + dtype=torch.uint8, + device="cuda", + ) + storage[:, ::2].copy_(p["xq"]) + a_noncontiguous = storage[:, ::2] + assert not a_noncontiguous.is_contiguous() + with pytest.raises(ValueError, match="a and b must be contiguous"): + mm_nvfp4_svdquant( + a_noncontiguous, + p["wq"], + p["x_sf_flat"], + p["w_sf_flat"], + p["alpha"], + p["d"], + p["l1_scaled"], + backend="cute-dsl", + ) + + +def test_mm_nvfp4_svdquant_sm120_fused_does_not_call_torch_mm(monkeypatch): + _skip_unless_sm120() + torch.manual_seed(0) + p = _make_gemm_problem( + 33, + 128, + 128, + rank=32, + quant_backend="cute-dsl", + residual_backend="b12x", + ) + + def fail_torch_mm(*args, **kwargs): + raise AssertionError("the fused SM120 SVDQuant path called torch.mm") + + monkeypatch.setattr(torch, "mm", fail_torch_mm) + out = mm_nvfp4_svdquant( + p["xq"], + p["wq"], + p["x_sf_flat"], + p["w_sf_flat"], + p["alpha"], + p["d"], + p["l1_scaled"], + backend="cute-dsl", + ) + assert torch.isfinite(out).all() + + +@pytest.mark.parametrize("use_bias", [False, True]) +def test_svdquant_linear_sm120_fused(use_bias): + _skip_unless_sm120() + torch.manual_seed(0) + m, n, k, rank = 129, 256, 256, 32 + x = torch.randn(m, k, dtype=torch.bfloat16, device="cuda") / (k**0.25) + pqs = ( + (1.0 + 0.3 * torch.randn(k, dtype=torch.bfloat16, device="cuda")) + .abs() + .contiguous() + ) + smoothed = (x * pqs).to(torch.bfloat16) + global_sf = ( + ((448.0 * 6.0) / smoothed.float().abs().nan_to_num().max()) + .reshape(1) + .contiguous() + ) + w = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") / (k**0.25) + wq, w_sf, gw = _nvfp4_quantize_128x4(w, backend="cute-dsl") + alpha = (1.0 / (global_sf * gw)).reshape(1).float() + lora_a = torch.randn(rank, k, dtype=torch.bfloat16, device="cuda") / (k**0.25) + l2t_smoothed = (pqs.unsqueeze(1) * lora_a.T).contiguous() + lora_b = torch.randn(n, rank, dtype=torch.bfloat16, device="cuda") / (rank**0.25) + l1_scaled = (lora_b.float() / alpha).to(torch.bfloat16).contiguous() + bias = ( + torch.randn(n, dtype=torch.bfloat16, device="cuda").contiguous() + if use_bias + else None + ) + + out = svdquant_linear( + x, + wq, + w_sf.reshape(-1), + alpha, + pqs, + l2t_smoothed, + l1_scaled, + global_sf, + bias=bias, + ) + + xq, x_sf = nvfp4_quantize( + smoothed, + global_sf, + sfLayout=SfLayout.layout_128x4, + do_shuffle=False, + backend="cute-dsl", + ) + residual = _mm_fp4_residual( + xq.view(torch.uint8), + wq, + x_sf.view(torch.uint8), + w_sf, + alpha, + backend="b12x", + ) + down = torch.mm(x, l2t_smoothed) + ref = residual + down.float() @ lora_b.float().T + if bias is not None: + ref.add_(bias.float()) + + assert out.shape == (m, n) and out.dtype == torch.bfloat16 + assert _sqnr_db(ref, out.float()) > 35.0 + + @pytest.mark.parametrize("rank", [32, 128]) def test_mm_nvfp4_svdquant_cuda_graph(rank): _skip_unless_sm100() diff --git a/tests/trace/example.py b/tests/trace/example.py index be8c7a3c775..c4a8ae516be 100644 --- a/tests/trace/example.py +++ b/tests/trace/example.py @@ -30,6 +30,7 @@ gemm_fp8_N1536_K7168.json gemm_fp8_nt_groupwise_n1536_k7168.json gemm_mxfp8_N4096_K4096.json +gemm_nvfp4_svdquant_N3072_K_packed1536_SF_A24576_SF_B589824_rank32.json gemma_fused_add_rmsnorm_h4608.json gemma_rmsnorm_h4608.json gelu_and_mul_h16384.json @@ -40,6 +41,7 @@ gqa_ragged_h32_kv8_d128.json layernorm_h768.json layernorm_quant_h768.json +linear_nvfp4_svdquant_N3072_K3072_K_packed1536_SF_B589824_rank32.json merge_state_h32_d128.json merge_state_in_place_h32_d128.json merge_states_h32_d128.json @@ -78,6 +80,7 @@ prims_ts_block_sparse_h8_kv8_d128_qb64_kb64.json prims_ts_paged_block_sparse_combined_h8_kv8_d128_qb64_kb64_ps64.json prims_ts_paged_block_sparse_tuple_h8_kv8_d128_qb64_kb64_ps64.json +quantize_nvfp4_smooth_N3072.json rmsnorm_h4096.json rmsnorm_h7168.json rmsnorm_quant_h7168.json @@ -468,7 +471,7 @@ except Exception: pass # Requires Blackwell (SM100+) -# ── SVDQuant fused NVFP4 GEMM (Blackwell SM100: M×3072@3072×3072, rank 32) ── +# ── SVDQuant fused NVFP4 GEMM (Blackwell: M×3072@3072×3072, rank 32) ───────── try: M, K, N, RANK = 128, 3072, 3072, 32 a_svdq = torch.zeros(M, K // 2, dtype=torch.uint8, device=device) @@ -484,9 +487,9 @@ a_svdq, b_svdq, a_sf_svdq, b_sf_svdq, alpha_svdq, d_svdq, l1_svdq ) except Exception: - pass # Requires Blackwell (SM100) + pass # Requires SM100/SM103 CUTLASS or SM120/SM121 CuTe DSL support -# ── SVDQuant smooth-quantize + composed linear (Blackwell SM100) ───────────── +# ── SVDQuant smooth-quantize + composed linear (Blackwell) ────────────────── try: M, K, N, RANK = 128, 3072, 3072, 32 x_sq = torch.zeros(M, K, dtype=torch.bfloat16, device=device) @@ -494,7 +497,7 @@ gs_sq = torch.ones(1, dtype=torch.float32, device=device) flashinfer.gemm.nvfp4_quantize_smooth(x_sq, pqs_sq, gs_sq) except Exception: - pass # Requires Blackwell (SM100) + pass # Requires SM100/SM103 CUTLASS or SM120/SM121 CuTe DSL support try: M, K, N, RANK = 128, 3072, 3072, 32 @@ -512,7 +515,7 @@ x_sl, w_sl, wsf_sl, alpha_sl, pqs_sl, l2t_sl, l1_sl, gs_sl ) except Exception: - pass # Requires Blackwell (SM100) + pass # Requires SM100/SM103 CUTLASS or SM120/SM121 CuTe DSL support # ── GEMM bf16 x fp4: mm_bf16_fp4 (weight-only) ────────────────────────────── # Blackwell SM100+: M×7168@2048×7168, block=16. b/b_descale shapes are the diff --git a/tests/trace/fi_trace_out/gemm_nvfp4_svdquant_N3072_K_packed1536_SF_B589824_rank32.json b/tests/trace/fi_trace_out/gemm_nvfp4_svdquant_N3072_K_packed1536_SF_B589824_rank32.json new file mode 100644 index 00000000000..c000b9e762a --- /dev/null +++ b/tests/trace/fi_trace_out/gemm_nvfp4_svdquant_N3072_K_packed1536_SF_B589824_rank32.json @@ -0,0 +1,106 @@ +{ + "name": "gemm_nvfp4_svdquant_N3072_K_packed1536_SF_B589824_rank32", + "description": "SVDQuant NVFP4 GEMM: out = alpha * (a @ b\u1d40 + d @ l1\u1d40). SM100/SM103 use fused CUTLASS; SM120/SM121 use fused CuTe DSL, with an explicit cute-dsl-unfused oracle. 1/alpha is pre-folded into l1.", + "op_type": "gemm_nvfp4_svdquant", + "tags": [ + "fi_api:flashinfer.gemm.gemm_svdquant.mm_nvfp4_svdquant", + "quantization:fp4" + ], + "axes": { + "M": { + "type": "var" + }, + "N": { + "type": "const", + "value": 3072 + }, + "K_packed": { + "type": "const", + "value": 1536, + "description": "K / 2 (two e2m1 values per byte)." + }, + "SF_A": { + "type": "var", + "description": "128x4-swizzled activation scale buffer size derived from M and K." + }, + "SF_B": { + "type": "const", + "value": 589824, + "description": "128x4-swizzled weight scale buffer size." + }, + "rank": { + "type": "const", + "value": 32, + "description": "LoRA rank, a positive multiple of 32." + } + }, + "constraints": [ + "SF_A == ((M + 127) // 128) * 128 * (((K_packed * 2 // 16) + 3) // 4) * 4" + ], + "inputs": { + "a": { + "shape": [ + "M", + "K_packed" + ], + "dtype": "uint8", + "description": "Smooth-quantized activation, packed e2m1 as uint8." + }, + "b": { + "shape": [ + "N", + "K_packed" + ], + "dtype": "uint8", + "description": "NVFP4 residual weight, packed e2m1 as uint8, row-major." + }, + "a_sf": { + "shape": [ + "SF_A" + ], + "dtype": "uint8", + "description": "Activation block scales, ue4m3 as uint8, 128x4 swizzled." + }, + "b_sf": { + "shape": [ + "SF_B" + ], + "dtype": "uint8", + "description": "Weight block scales, ue4m3 as uint8, 128x4 swizzled." + }, + "alpha": { + "shape": [ + "1" + ], + "dtype": "float32", + "description": "Per-tensor residual dequant scale, float32 device scalar." + }, + "d": { + "shape": [ + "M", + "rank" + ], + "dtype": "bfloat16", + "description": "LoRA-down output x_hat @ L2\u1d40, bf16." + }, + "l1": { + "shape": [ + "N", + "rank" + ], + "dtype": "bfloat16", + "description": "LoRA-up weight pre-divided by alpha, bf16." + } + }, + "outputs": { + "out": { + "shape": [ + "M", + "N" + ], + "dtype": "bfloat16" + } + }, + "check": "def standard_check(\n reference_outputs: Any,\n actual_outputs: Any,\n *,\n rtol: Optional[float] = None,\n atol: Optional[float] = None,\n max_mismatch_pct: float = 0.0,\n min_cos_sim: Optional[float] = 1.0 - 1e-3,\n) -> bool:\n \"\"\"Default trace correctness check used when a template does not override it.\"\"\"\n from flashinfer.trace import default_check\n\n return default_check(\n reference_outputs,\n actual_outputs,\n rtol=rtol,\n atol=atol,\n max_mismatch_pct=max_mismatch_pct,\n min_cos_sim=min_cos_sim,\n )\n", + "init": "from __future__ import annotations\nimport math\nimport torch\n\n# ----- shared init helpers -----\n# Copyright (c) 2025 by FlashInfer team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Shared helpers used by ``TraceTemplate.init`` functions.\n\nThis module contains the small set of input-construction patterns that\nrecur across many templates (paged-KV cache index arrays, ragged indptr,\nRoPE pos_ids and cos/sin caches, sampling probs). Each helper is short and\ndocumented; init functions in ``templates/.py`` call into here so\nthe per-template init bodies stay focused on shape/dtype, not boilerplate.\n\nThe full source of this module is **inlined into every dumped JSON's\n``\"init\"`` field** by ``flashinfer/trace/template.py:_render_init_source``,\nso downstream consumers don't need flashinfer installed to re-run the init\nsnippets.\n\"\"\"\n\n\nfrom typing import Optional, Tuple\n\nimport torch\n\n\ndef make_paged_kv_indices(\n batch_size: int,\n num_pages_per_seq: int,\n page_size: int,\n *,\n device: str = \"cuda\",\n) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n \"\"\"Return ``(kv_indptr, kv_indices, kv_last_page_len)`` for a uniform batch.\n\n Every sequence is assigned exactly ``num_pages_per_seq`` pages, fully\n populated (last-page length == page_size).\n\n Invariants\n ----------\n - ``kv_indptr.shape == (batch_size + 1,)``, dtype int32, monotonic, [0]=0.\n - ``kv_indices == arange(0, batch_size * num_pages_per_seq)``, int32.\n - ``kv_last_page_len == full(batch_size, page_size)``, int32.\n \"\"\"\n total_pages = batch_size * num_pages_per_seq\n kv_indptr = (\n torch.arange(batch_size + 1, dtype=torch.int32, device=device)\n * num_pages_per_seq\n )\n kv_indices = torch.arange(total_pages, dtype=torch.int32, device=device)\n kv_last_page_len = torch.full(\n (batch_size,), page_size, dtype=torch.int32, device=device\n )\n return kv_indptr, kv_indices, kv_last_page_len\n\n\ndef make_ragged_indptr(\n seg_lens,\n *,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.int32,\n) -> torch.Tensor:\n \"\"\"Return cumulative-sum ``indptr`` of length ``len(seg_lens)+1``.\n\n ``seg_lens`` may be a list / tuple / 1-D tensor of segment lengths.\n \"\"\"\n if isinstance(seg_lens, torch.Tensor):\n lens = seg_lens.to(device=device, dtype=dtype)\n else:\n lens = torch.tensor(list(seg_lens), dtype=dtype, device=device)\n indptr = torch.zeros(lens.numel() + 1, dtype=dtype, device=device)\n indptr[1:] = torch.cumsum(lens, dim=0).to(dtype)\n return indptr\n\n\ndef make_uniform_qo_indptr(\n batch_size: int,\n qo_len: int,\n *,\n device: str = \"cuda\",\n) -> torch.Tensor:\n \"\"\"Return ``[0, qo_len, 2*qo_len, ..., batch_size*qo_len]`` int32.\"\"\"\n return torch.arange(batch_size + 1, dtype=torch.int32, device=device) * qo_len\n\n\ndef make_pos_ids(\n nnz: int,\n max_seq_len: Optional[int] = None,\n *,\n device: str = \"cuda\",\n) -> torch.Tensor:\n \"\"\"Return ``[0, 1, ..., nnz-1] (% max_seq_len)`` as int32 on ``device``.\n\n If ``max_seq_len`` is None, no wrapping is applied.\n \"\"\"\n pos = torch.arange(nnz, dtype=torch.int32, device=device)\n if max_seq_len is not None:\n pos = pos % max_seq_len\n return pos\n\n\ndef make_rope_cos_sin_cache(\n max_seq_len: int,\n rope_dim: int,\n *,\n base: float = 1e4,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Return concatenated ``[cos | sin]`` cache of shape ``[max_seq_len, rope_dim]``.\"\"\"\n t = torch.arange(max_seq_len, dtype=torch.float32, device=device)\n inv = 1.0 / (\n base\n ** (torch.arange(0, rope_dim, 2, dtype=torch.float32, device=device) / rope_dim)\n )\n freqs = t.unsqueeze(-1) * inv.unsqueeze(0)\n cache = torch.cat([torch.cos(freqs), torch.sin(freqs)], dim=-1)\n return cache.to(dtype)\n\n\ndef make_probs(\n batch_size: int,\n vocab_size: int,\n *,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Return a ``[batch_size, vocab_size]`` probability distribution.\n\n Uses ``softmax(randn(...))`` so each row sums to 1.0. This mirrors the\n pattern used throughout ``tests/utils/test_sampling.py``.\n \"\"\"\n return torch.softmax(\n torch.randn(batch_size, vocab_size, dtype=torch.float32, device=device),\n dim=-1,\n ).to(dtype)\n\n\ndef make_logits(\n batch_size: int,\n vocab_size: int,\n *,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Return ``randn(batch_size, vocab_size)`` logits.\"\"\"\n return torch.randn(batch_size, vocab_size, dtype=dtype, device=device)\n\n\ndef fp8_safe_randn(\n *shape: int,\n scale: float = 0.1,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.bfloat16,\n) -> torch.Tensor:\n \"\"\"``randn(*shape) * scale`` \u2014 keeps values in the FP8/FP4 representable range.\n\n Tests for fp8/fp4 paths typically multiply ``randn`` by 0.1 to avoid\n saturation when quantizing. Use this helper to mirror that convention.\n \"\"\"\n return (torch.randn(*shape, dtype=dtype, device=device) * scale).to(dtype)\n\n\ndef per_tensor_fp8_quantize(\n x: torch.Tensor,\n *,\n fp8_dtype: torch.dtype = torch.float8_e4m3fn,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Per-tensor FP8 quantization, mirroring ``tests/utils_fp8.py:to_float8``.\n\n Returns ``(x_fp8, inv_scale)`` where ``inv_scale`` is the dequant\n multiplier (``float \u2248 fp8 * inv_scale``).\n \"\"\"\n finfo = torch.finfo(fp8_dtype)\n amax = x.abs().amax().clamp(min=1e-12)\n scale = finfo.max / amax\n x_q = (x.float() * scale).clamp(min=finfo.min, max=finfo.max).to(fp8_dtype)\n return x_q, scale.float().reciprocal()\n\n\ndef fp8_block_quant_1d(\n x_bf16: torch.Tensor,\n block: int = 128,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Quantize ``[T, H]`` activations into FP8 with per-``(token, block)``\n column-block scales. Returns ``(x_fp8, scales)`` where\n ``scales`` has shape ``[T, H // block]``.\n\n Mirrors ``_fp8_block_quant_1d`` in\n ``tests/moe/test_dpsk_fused_moe_fp8.py``.\n \"\"\"\n assert x_bf16.dim() == 2\n T, H = x_bf16.shape\n assert H % block == 0\n nb = H // block\n finfo = torch.finfo(torch.float8_e4m3fn)\n max_fp8 = finfo.max\n x_f32 = x_bf16.to(torch.float32)\n x_fp8 = torch.empty((T, H), dtype=torch.float8_e4m3fn, device=x_bf16.device)\n scales = torch.empty((T, nb), dtype=torch.float32, device=x_bf16.device)\n for j in range(nb):\n sl = slice(j * block, (j + 1) * block)\n blk = x_f32[:, sl]\n amax = torch.amax(torch.abs(blk), dim=1)\n s = torch.where(amax > 0, amax / max_fp8, torch.ones_like(amax))\n x_fp8[:, sl] = (blk / s.unsqueeze(1)).to(torch.float8_e4m3fn)\n scales[:, j] = s\n return x_fp8, scales\n\n\ndef fp8_block_quant_2d(\n w_bf16: torch.Tensor,\n block: int = 128,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Quantize weights ``[..., R, C]`` with 2-D ``block \u00d7 block`` scales.\n\n Returns ``(w_fp8, scales)`` where ``scales`` has shape\n ``[..., R // block, C // block]``. Mirrors ``_fp8_block_quant_2d`` in\n ``tests/moe/test_dpsk_fused_moe_fp8.py``.\n \"\"\"\n assert w_bf16.dim() >= 2\n *prefix, R, C = w_bf16.shape\n assert R % block == 0 and C % block == 0\n nb_r, nb_c = R // block, C // block\n finfo = torch.finfo(torch.float8_e4m3fn)\n max_fp8 = finfo.max\n w_f32 = w_bf16.to(torch.float32).contiguous()\n prefix_ndim = len(prefix)\n reshaped = w_f32.reshape(*prefix, nb_r, block, nb_c, block)\n permute_dims = tuple(range(prefix_ndim)) + (\n prefix_ndim,\n prefix_ndim + 2,\n prefix_ndim + 1,\n prefix_ndim + 3,\n )\n blocks = reshaped.permute(permute_dims).contiguous()\n amax = torch.amax(torch.abs(blocks), dim=(-1, -2))\n scales = torch.where(\n amax > 0, amax / max_fp8, torch.ones_like(amax, dtype=torch.float32)\n )\n q_blocks = (blocks / scales.unsqueeze(-1).unsqueeze(-1)).to(torch.float8_e4m3fn)\n inv_permute = [0] * (prefix_ndim + 4)\n for i, p in enumerate(permute_dims):\n inv_permute[p] = i\n w_fp8 = q_blocks.permute(*inv_permute).reshape(*prefix, R, C).contiguous()\n return w_fp8, scales\n\n\n__all__ = [\n \"make_paged_kv_indices\",\n \"make_ragged_indptr\",\n \"make_uniform_qo_indptr\",\n \"make_pos_ids\",\n \"make_rope_cos_sin_cache\",\n \"make_probs\",\n \"make_logits\",\n \"fp8_safe_randn\",\n \"per_tensor_fp8_quantize\",\n \"fp8_block_quant_1d\",\n \"fp8_block_quant_2d\",\n]\n\n# ----- init -----\ndef _mm_nvfp4_svdquant_init(\n *,\n M: int,\n N: int = 3072,\n K: int = 3072,\n SF_A: int = 0,\n device: str = \"cuda\",\n seed: int = 0,\n):\n \"\"\"Build inputs for ``flashinfer.mm_nvfp4_svdquant``.\n\n Mirrors the SVDQuant linear decomposition W ~= R + L1 @ L2 at Qwen-Image\n shapes: the residual weight is NVFP4-quantized (via ``nvfp4_quantize_smooth``\n with a unit smoothing scale, which is byte-identical to the stock NVFP4\n quantizer), the activation is smooth-quantized, and the rank-32 LoRA factors\n follow the host-side folding contract (``d = x_hat @ L2\u1d40``, ``l1 = L1 / alpha``).\n \"\"\"\n from flashinfer import nvfp4_quantize_smooth # noqa: PLC0415\n\n del SF_A # output-only / derived axis\n\n torch.manual_seed(seed)\n rank = 32\n x = torch.randn(M, K, dtype=torch.bfloat16, device=device)\n w = torch.randn(N, K, dtype=torch.bfloat16, device=device) / math.sqrt(K)\n pqs = torch.rand(K, dtype=torch.bfloat16, device=device) + 0.5\n l1 = torch.randn(N, rank, dtype=torch.bfloat16, device=device) / math.sqrt(rank)\n l2 = torch.randn(rank, K, dtype=torch.bfloat16, device=device) / math.sqrt(K)\n\n x_hat = (x.float() * pqs.float()).to(torch.bfloat16)\n x_gs = (\n ((448 * 6) / x_hat.float().abs().nan_to_num().max())\n .to(torch.float32)\n .reshape(1)\n )\n w_gs = ((448 * 6) / w.float().abs().nan_to_num().max()).to(torch.float32).reshape(1)\n ones = torch.ones(K, dtype=torch.bfloat16, device=device)\n\n a, a_sf = nvfp4_quantize_smooth(x, pqs, x_gs)\n b, b_sf = nvfp4_quantize_smooth(w, ones, w_gs)\n alpha = (1.0 / (x_gs * w_gs)).to(torch.float32).reshape(1)\n d = torch.mm(x_hat, l2.t().contiguous().to(torch.bfloat16))\n l1_scaled = (l1.float() / alpha.item()).to(torch.bfloat16)\n return {\n \"a\": a,\n \"b\": b,\n \"a_sf\": a_sf,\n \"b_sf\": b_sf,\n \"alpha\": alpha,\n \"d\": d,\n \"l1\": l1_scaled,\n }\n" +} diff --git a/tests/trace/fi_trace_out/linear_nvfp4_svdquant_N3072_K3072_K_packed1536_SF_B589824_rank32.json b/tests/trace/fi_trace_out/linear_nvfp4_svdquant_N3072_K3072_K_packed1536_SF_B589824_rank32.json new file mode 100644 index 00000000000..ef95ba2d810 --- /dev/null +++ b/tests/trace/fi_trace_out/linear_nvfp4_svdquant_N3072_K3072_K_packed1536_SF_B589824_rank32.json @@ -0,0 +1,110 @@ +{ + "name": "linear_nvfp4_svdquant_N3072_K3072_K_packed1536_SF_B589824_rank32", + "description": "Full SVDQuant linear: y = (x * pre_quant_scale) @ (R + L1 @ L2)\u1d40 where R is the NVFP4-quantized residual weight \u2014 smooth-quantize, BF16 rank-r down-projection, and the architecture-selected NVFP4 residual + LoRA-up GEMM.", + "op_type": "linear_nvfp4_svdquant", + "tags": [ + "fi_api:flashinfer.gemm.gemm_svdquant.svdquant_linear", + "quantization:fp4" + ], + "axes": { + "M": { + "type": "var" + }, + "N": { + "type": "const", + "value": 3072 + }, + "K": { + "type": "const", + "value": 3072 + }, + "K_packed": { + "type": "const", + "value": 1536, + "description": "K / 2 (two e2m1 values per byte)." + }, + "SF_B": { + "type": "const", + "value": 589824, + "description": "128x4-swizzled weight scale buffer size." + }, + "rank": { + "type": "const", + "value": 32, + "description": "LoRA rank, a positive multiple of 32." + } + }, + "inputs": { + "x": { + "shape": [ + "M", + "K" + ], + "dtype": "bfloat16", + "description": "Input activation, bf16." + }, + "weight_fp4": { + "shape": [ + "N", + "K_packed" + ], + "dtype": "uint8", + "description": "NVFP4 residual weight, packed e2m1 as uint8." + }, + "weight_sf": { + "shape": [ + "SF_B" + ], + "dtype": "uint8", + "description": "Weight block scales, ue4m3 as uint8, 128x4 swizzled." + }, + "alpha": { + "shape": [ + "1" + ], + "dtype": "float32", + "description": "Per-tensor residual dequant scale, float32 device scalar." + }, + "pre_quant_scale": { + "shape": [ + "K" + ], + "dtype": "bfloat16", + "description": "Per-input-channel smoothing scale, bf16." + }, + "l2t_smoothed": { + "shape": [ + "K", + "rank" + ], + "dtype": "bfloat16", + "description": "pre_quant_scale[:, None] * L2\u1d40, bf16." + }, + "l1_scaled": { + "shape": [ + "N", + "rank" + ], + "dtype": "bfloat16", + "description": "L1 / alpha, bf16." + }, + "global_scale": { + "shape": [ + "1" + ], + "dtype": "float32", + "description": "Activation global scale, float32 device scalar." + } + }, + "outputs": { + "out": { + "shape": [ + "M", + "N" + ], + "dtype": "bfloat16" + } + }, + "check": "def standard_check(\n reference_outputs: Any,\n actual_outputs: Any,\n *,\n rtol: Optional[float] = None,\n atol: Optional[float] = None,\n max_mismatch_pct: float = 0.0,\n min_cos_sim: Optional[float] = 1.0 - 1e-3,\n) -> bool:\n \"\"\"Default trace correctness check used when a template does not override it.\"\"\"\n from flashinfer.trace import default_check\n\n return default_check(\n reference_outputs,\n actual_outputs,\n rtol=rtol,\n atol=atol,\n max_mismatch_pct=max_mismatch_pct,\n min_cos_sim=min_cos_sim,\n )\n", + "init": "from __future__ import annotations\nimport math\nimport torch\n\n# ----- shared init helpers -----\n# Copyright (c) 2025 by FlashInfer team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Shared helpers used by ``TraceTemplate.init`` functions.\n\nThis module contains the small set of input-construction patterns that\nrecur across many templates (paged-KV cache index arrays, ragged indptr,\nRoPE pos_ids and cos/sin caches, sampling probs). Each helper is short and\ndocumented; init functions in ``templates/.py`` call into here so\nthe per-template init bodies stay focused on shape/dtype, not boilerplate.\n\nThe full source of this module is **inlined into every dumped JSON's\n``\"init\"`` field** by ``flashinfer/trace/template.py:_render_init_source``,\nso downstream consumers don't need flashinfer installed to re-run the init\nsnippets.\n\"\"\"\n\n\nfrom typing import Optional, Tuple\n\nimport torch\n\n\ndef make_paged_kv_indices(\n batch_size: int,\n num_pages_per_seq: int,\n page_size: int,\n *,\n device: str = \"cuda\",\n) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n \"\"\"Return ``(kv_indptr, kv_indices, kv_last_page_len)`` for a uniform batch.\n\n Every sequence is assigned exactly ``num_pages_per_seq`` pages, fully\n populated (last-page length == page_size).\n\n Invariants\n ----------\n - ``kv_indptr.shape == (batch_size + 1,)``, dtype int32, monotonic, [0]=0.\n - ``kv_indices == arange(0, batch_size * num_pages_per_seq)``, int32.\n - ``kv_last_page_len == full(batch_size, page_size)``, int32.\n \"\"\"\n total_pages = batch_size * num_pages_per_seq\n kv_indptr = (\n torch.arange(batch_size + 1, dtype=torch.int32, device=device)\n * num_pages_per_seq\n )\n kv_indices = torch.arange(total_pages, dtype=torch.int32, device=device)\n kv_last_page_len = torch.full(\n (batch_size,), page_size, dtype=torch.int32, device=device\n )\n return kv_indptr, kv_indices, kv_last_page_len\n\n\ndef make_ragged_indptr(\n seg_lens,\n *,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.int32,\n) -> torch.Tensor:\n \"\"\"Return cumulative-sum ``indptr`` of length ``len(seg_lens)+1``.\n\n ``seg_lens`` may be a list / tuple / 1-D tensor of segment lengths.\n \"\"\"\n if isinstance(seg_lens, torch.Tensor):\n lens = seg_lens.to(device=device, dtype=dtype)\n else:\n lens = torch.tensor(list(seg_lens), dtype=dtype, device=device)\n indptr = torch.zeros(lens.numel() + 1, dtype=dtype, device=device)\n indptr[1:] = torch.cumsum(lens, dim=0).to(dtype)\n return indptr\n\n\ndef make_uniform_qo_indptr(\n batch_size: int,\n qo_len: int,\n *,\n device: str = \"cuda\",\n) -> torch.Tensor:\n \"\"\"Return ``[0, qo_len, 2*qo_len, ..., batch_size*qo_len]`` int32.\"\"\"\n return torch.arange(batch_size + 1, dtype=torch.int32, device=device) * qo_len\n\n\ndef make_pos_ids(\n nnz: int,\n max_seq_len: Optional[int] = None,\n *,\n device: str = \"cuda\",\n) -> torch.Tensor:\n \"\"\"Return ``[0, 1, ..., nnz-1] (% max_seq_len)`` as int32 on ``device``.\n\n If ``max_seq_len`` is None, no wrapping is applied.\n \"\"\"\n pos = torch.arange(nnz, dtype=torch.int32, device=device)\n if max_seq_len is not None:\n pos = pos % max_seq_len\n return pos\n\n\ndef make_rope_cos_sin_cache(\n max_seq_len: int,\n rope_dim: int,\n *,\n base: float = 1e4,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Return concatenated ``[cos | sin]`` cache of shape ``[max_seq_len, rope_dim]``.\"\"\"\n t = torch.arange(max_seq_len, dtype=torch.float32, device=device)\n inv = 1.0 / (\n base\n ** (torch.arange(0, rope_dim, 2, dtype=torch.float32, device=device) / rope_dim)\n )\n freqs = t.unsqueeze(-1) * inv.unsqueeze(0)\n cache = torch.cat([torch.cos(freqs), torch.sin(freqs)], dim=-1)\n return cache.to(dtype)\n\n\ndef make_probs(\n batch_size: int,\n vocab_size: int,\n *,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Return a ``[batch_size, vocab_size]`` probability distribution.\n\n Uses ``softmax(randn(...))`` so each row sums to 1.0. This mirrors the\n pattern used throughout ``tests/utils/test_sampling.py``.\n \"\"\"\n return torch.softmax(\n torch.randn(batch_size, vocab_size, dtype=torch.float32, device=device),\n dim=-1,\n ).to(dtype)\n\n\ndef make_logits(\n batch_size: int,\n vocab_size: int,\n *,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Return ``randn(batch_size, vocab_size)`` logits.\"\"\"\n return torch.randn(batch_size, vocab_size, dtype=dtype, device=device)\n\n\ndef fp8_safe_randn(\n *shape: int,\n scale: float = 0.1,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.bfloat16,\n) -> torch.Tensor:\n \"\"\"``randn(*shape) * scale`` \u2014 keeps values in the FP8/FP4 representable range.\n\n Tests for fp8/fp4 paths typically multiply ``randn`` by 0.1 to avoid\n saturation when quantizing. Use this helper to mirror that convention.\n \"\"\"\n return (torch.randn(*shape, dtype=dtype, device=device) * scale).to(dtype)\n\n\ndef per_tensor_fp8_quantize(\n x: torch.Tensor,\n *,\n fp8_dtype: torch.dtype = torch.float8_e4m3fn,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Per-tensor FP8 quantization, mirroring ``tests/utils_fp8.py:to_float8``.\n\n Returns ``(x_fp8, inv_scale)`` where ``inv_scale`` is the dequant\n multiplier (``float \u2248 fp8 * inv_scale``).\n \"\"\"\n finfo = torch.finfo(fp8_dtype)\n amax = x.abs().amax().clamp(min=1e-12)\n scale = finfo.max / amax\n x_q = (x.float() * scale).clamp(min=finfo.min, max=finfo.max).to(fp8_dtype)\n return x_q, scale.float().reciprocal()\n\n\ndef fp8_block_quant_1d(\n x_bf16: torch.Tensor,\n block: int = 128,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Quantize ``[T, H]`` activations into FP8 with per-``(token, block)``\n column-block scales. Returns ``(x_fp8, scales)`` where\n ``scales`` has shape ``[T, H // block]``.\n\n Mirrors ``_fp8_block_quant_1d`` in\n ``tests/moe/test_dpsk_fused_moe_fp8.py``.\n \"\"\"\n assert x_bf16.dim() == 2\n T, H = x_bf16.shape\n assert H % block == 0\n nb = H // block\n finfo = torch.finfo(torch.float8_e4m3fn)\n max_fp8 = finfo.max\n x_f32 = x_bf16.to(torch.float32)\n x_fp8 = torch.empty((T, H), dtype=torch.float8_e4m3fn, device=x_bf16.device)\n scales = torch.empty((T, nb), dtype=torch.float32, device=x_bf16.device)\n for j in range(nb):\n sl = slice(j * block, (j + 1) * block)\n blk = x_f32[:, sl]\n amax = torch.amax(torch.abs(blk), dim=1)\n s = torch.where(amax > 0, amax / max_fp8, torch.ones_like(amax))\n x_fp8[:, sl] = (blk / s.unsqueeze(1)).to(torch.float8_e4m3fn)\n scales[:, j] = s\n return x_fp8, scales\n\n\ndef fp8_block_quant_2d(\n w_bf16: torch.Tensor,\n block: int = 128,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Quantize weights ``[..., R, C]`` with 2-D ``block \u00d7 block`` scales.\n\n Returns ``(w_fp8, scales)`` where ``scales`` has shape\n ``[..., R // block, C // block]``. Mirrors ``_fp8_block_quant_2d`` in\n ``tests/moe/test_dpsk_fused_moe_fp8.py``.\n \"\"\"\n assert w_bf16.dim() >= 2\n *prefix, R, C = w_bf16.shape\n assert R % block == 0 and C % block == 0\n nb_r, nb_c = R // block, C // block\n finfo = torch.finfo(torch.float8_e4m3fn)\n max_fp8 = finfo.max\n w_f32 = w_bf16.to(torch.float32).contiguous()\n prefix_ndim = len(prefix)\n reshaped = w_f32.reshape(*prefix, nb_r, block, nb_c, block)\n permute_dims = tuple(range(prefix_ndim)) + (\n prefix_ndim,\n prefix_ndim + 2,\n prefix_ndim + 1,\n prefix_ndim + 3,\n )\n blocks = reshaped.permute(permute_dims).contiguous()\n amax = torch.amax(torch.abs(blocks), dim=(-1, -2))\n scales = torch.where(\n amax > 0, amax / max_fp8, torch.ones_like(amax, dtype=torch.float32)\n )\n q_blocks = (blocks / scales.unsqueeze(-1).unsqueeze(-1)).to(torch.float8_e4m3fn)\n inv_permute = [0] * (prefix_ndim + 4)\n for i, p in enumerate(permute_dims):\n inv_permute[p] = i\n w_fp8 = q_blocks.permute(*inv_permute).reshape(*prefix, R, C).contiguous()\n return w_fp8, scales\n\n\n__all__ = [\n \"make_paged_kv_indices\",\n \"make_ragged_indptr\",\n \"make_uniform_qo_indptr\",\n \"make_pos_ids\",\n \"make_rope_cos_sin_cache\",\n \"make_probs\",\n \"make_logits\",\n \"fp8_safe_randn\",\n \"per_tensor_fp8_quantize\",\n \"fp8_block_quant_1d\",\n \"fp8_block_quant_2d\",\n]\n\n# ----- init -----\ndef _svdquant_linear_init(\n *,\n M: int,\n N: int = 3072,\n K: int = 3072,\n device: str = \"cuda\",\n seed: int = 0,\n):\n \"\"\"Build inputs for ``flashinfer.svdquant_linear`` (full SVDQuant linear chain).\"\"\"\n from flashinfer import nvfp4_quantize_smooth # noqa: PLC0415\n\n torch.manual_seed(seed)\n rank = 32\n x = torch.randn(M, K, dtype=torch.bfloat16, device=device)\n w = torch.randn(N, K, dtype=torch.bfloat16, device=device) / math.sqrt(K)\n pqs = torch.rand(K, dtype=torch.bfloat16, device=device) + 0.5\n l1 = torch.randn(N, rank, dtype=torch.bfloat16, device=device) / math.sqrt(rank)\n l2 = torch.randn(rank, K, dtype=torch.bfloat16, device=device) / math.sqrt(K)\n\n x_hat = (x.float() * pqs.float()).to(torch.bfloat16)\n x_gs = (\n ((448 * 6) / x_hat.float().abs().nan_to_num().max())\n .to(torch.float32)\n .reshape(1)\n )\n w_gs = ((448 * 6) / w.float().abs().nan_to_num().max()).to(torch.float32).reshape(1)\n ones = torch.ones(K, dtype=torch.bfloat16, device=device)\n\n weight_fp4, weight_sf = nvfp4_quantize_smooth(w, ones, w_gs)\n alpha = (1.0 / (x_gs * w_gs)).to(torch.float32).reshape(1)\n l2t_smoothed = (pqs.float()[:, None] * l2.float().t()).to(torch.bfloat16)\n l1_scaled = (l1.float() / alpha.item()).to(torch.bfloat16)\n return {\n \"x\": x,\n \"weight_fp4\": weight_fp4,\n \"weight_sf\": weight_sf,\n \"alpha\": alpha,\n \"pre_quant_scale\": pqs,\n \"l2t_smoothed\": l2t_smoothed,\n \"l1_scaled\": l1_scaled,\n \"global_scale\": x_gs,\n }\n" +} diff --git a/tests/trace/fi_trace_out/quantize_nvfp4_smooth_N3072.json b/tests/trace/fi_trace_out/quantize_nvfp4_smooth_N3072.json new file mode 100644 index 00000000000..7db08a189d9 --- /dev/null +++ b/tests/trace/fi_trace_out/quantize_nvfp4_smooth_N3072.json @@ -0,0 +1,71 @@ +{ + "name": "quantize_nvfp4_smooth_N3072", + "description": "Smooth + NVFP4 quantize: (xq, sf) = nvfp4-quantize(x * pre_quant_scale). SM100/SM103 fuse the operations; SM120/SM121 materialize the BF16 smoothed input before CuTe DSL quantization. Both use ue4m3 block scales, 128x4 swizzled layout, and SF vector size 16.", + "op_type": "quantize_nvfp4_smooth", + "tags": [ + "fi_api:flashinfer.gemm.gemm_svdquant.nvfp4_quantize_smooth", + "quantization:fp4" + ], + "axes": { + "M": { + "type": "var" + }, + "N": { + "type": "const", + "value": 3072 + }, + "N_half": { + "type": "var", + "description": "N / 2 (two e2m1 values per byte)." + }, + "SF": { + "type": "var", + "description": "128x4-swizzled scale buffer size derived from M and N." + } + }, + "constraints": [ + "N_half == N // 2", + "SF == ((M + 127) // 128) * 128 * ((N // 16 + 3) // 4) * 4" + ], + "inputs": { + "x": { + "shape": [ + "M", + "N" + ], + "dtype": "bfloat16", + "description": "Input activation, bf16." + }, + "pre_quant_scale": { + "shape": [ + "N" + ], + "dtype": "bfloat16", + "description": "Per-input-channel smoothing scale, bf16." + }, + "global_scale": { + "shape": [ + "1" + ], + "dtype": "float32", + "description": "Global scale, float32 device scalar." + } + }, + "outputs": { + "xq": { + "shape": [ + "M", + "N_half" + ], + "dtype": "uint8" + }, + "sf": { + "shape": [ + "SF" + ], + "dtype": "uint8" + } + }, + "check": "def standard_check(\n reference_outputs: Any,\n actual_outputs: Any,\n *,\n rtol: Optional[float] = None,\n atol: Optional[float] = None,\n max_mismatch_pct: float = 0.0,\n min_cos_sim: Optional[float] = 1.0 - 1e-3,\n) -> bool:\n \"\"\"Default trace correctness check used when a template does not override it.\"\"\"\n from flashinfer.trace import default_check\n\n return default_check(\n reference_outputs,\n actual_outputs,\n rtol=rtol,\n atol=atol,\n max_mismatch_pct=max_mismatch_pct,\n min_cos_sim=min_cos_sim,\n )\n", + "init": "from __future__ import annotations\nimport math\nimport torch\n\n# ----- shared init helpers -----\n# Copyright (c) 2025 by FlashInfer team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Shared helpers used by ``TraceTemplate.init`` functions.\n\nThis module contains the small set of input-construction patterns that\nrecur across many templates (paged-KV cache index arrays, ragged indptr,\nRoPE pos_ids and cos/sin caches, sampling probs). Each helper is short and\ndocumented; init functions in ``templates/.py`` call into here so\nthe per-template init bodies stay focused on shape/dtype, not boilerplate.\n\nThe full source of this module is **inlined into every dumped JSON's\n``\"init\"`` field** by ``flashinfer/trace/template.py:_render_init_source``,\nso downstream consumers don't need flashinfer installed to re-run the init\nsnippets.\n\"\"\"\n\n\nfrom typing import Optional, Tuple\n\nimport torch\n\n\ndef make_paged_kv_indices(\n batch_size: int,\n num_pages_per_seq: int,\n page_size: int,\n *,\n device: str = \"cuda\",\n) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n \"\"\"Return ``(kv_indptr, kv_indices, kv_last_page_len)`` for a uniform batch.\n\n Every sequence is assigned exactly ``num_pages_per_seq`` pages, fully\n populated (last-page length == page_size).\n\n Invariants\n ----------\n - ``kv_indptr.shape == (batch_size + 1,)``, dtype int32, monotonic, [0]=0.\n - ``kv_indices == arange(0, batch_size * num_pages_per_seq)``, int32.\n - ``kv_last_page_len == full(batch_size, page_size)``, int32.\n \"\"\"\n total_pages = batch_size * num_pages_per_seq\n kv_indptr = (\n torch.arange(batch_size + 1, dtype=torch.int32, device=device)\n * num_pages_per_seq\n )\n kv_indices = torch.arange(total_pages, dtype=torch.int32, device=device)\n kv_last_page_len = torch.full(\n (batch_size,), page_size, dtype=torch.int32, device=device\n )\n return kv_indptr, kv_indices, kv_last_page_len\n\n\ndef make_ragged_indptr(\n seg_lens,\n *,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.int32,\n) -> torch.Tensor:\n \"\"\"Return cumulative-sum ``indptr`` of length ``len(seg_lens)+1``.\n\n ``seg_lens`` may be a list / tuple / 1-D tensor of segment lengths.\n \"\"\"\n if isinstance(seg_lens, torch.Tensor):\n lens = seg_lens.to(device=device, dtype=dtype)\n else:\n lens = torch.tensor(list(seg_lens), dtype=dtype, device=device)\n indptr = torch.zeros(lens.numel() + 1, dtype=dtype, device=device)\n indptr[1:] = torch.cumsum(lens, dim=0).to(dtype)\n return indptr\n\n\ndef make_uniform_qo_indptr(\n batch_size: int,\n qo_len: int,\n *,\n device: str = \"cuda\",\n) -> torch.Tensor:\n \"\"\"Return ``[0, qo_len, 2*qo_len, ..., batch_size*qo_len]`` int32.\"\"\"\n return torch.arange(batch_size + 1, dtype=torch.int32, device=device) * qo_len\n\n\ndef make_pos_ids(\n nnz: int,\n max_seq_len: Optional[int] = None,\n *,\n device: str = \"cuda\",\n) -> torch.Tensor:\n \"\"\"Return ``[0, 1, ..., nnz-1] (% max_seq_len)`` as int32 on ``device``.\n\n If ``max_seq_len`` is None, no wrapping is applied.\n \"\"\"\n pos = torch.arange(nnz, dtype=torch.int32, device=device)\n if max_seq_len is not None:\n pos = pos % max_seq_len\n return pos\n\n\ndef make_rope_cos_sin_cache(\n max_seq_len: int,\n rope_dim: int,\n *,\n base: float = 1e4,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Return concatenated ``[cos | sin]`` cache of shape ``[max_seq_len, rope_dim]``.\"\"\"\n t = torch.arange(max_seq_len, dtype=torch.float32, device=device)\n inv = 1.0 / (\n base\n ** (torch.arange(0, rope_dim, 2, dtype=torch.float32, device=device) / rope_dim)\n )\n freqs = t.unsqueeze(-1) * inv.unsqueeze(0)\n cache = torch.cat([torch.cos(freqs), torch.sin(freqs)], dim=-1)\n return cache.to(dtype)\n\n\ndef make_probs(\n batch_size: int,\n vocab_size: int,\n *,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Return a ``[batch_size, vocab_size]`` probability distribution.\n\n Uses ``softmax(randn(...))`` so each row sums to 1.0. This mirrors the\n pattern used throughout ``tests/utils/test_sampling.py``.\n \"\"\"\n return torch.softmax(\n torch.randn(batch_size, vocab_size, dtype=torch.float32, device=device),\n dim=-1,\n ).to(dtype)\n\n\ndef make_logits(\n batch_size: int,\n vocab_size: int,\n *,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Return ``randn(batch_size, vocab_size)`` logits.\"\"\"\n return torch.randn(batch_size, vocab_size, dtype=dtype, device=device)\n\n\ndef fp8_safe_randn(\n *shape: int,\n scale: float = 0.1,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.bfloat16,\n) -> torch.Tensor:\n \"\"\"``randn(*shape) * scale`` \u2014 keeps values in the FP8/FP4 representable range.\n\n Tests for fp8/fp4 paths typically multiply ``randn`` by 0.1 to avoid\n saturation when quantizing. Use this helper to mirror that convention.\n \"\"\"\n return (torch.randn(*shape, dtype=dtype, device=device) * scale).to(dtype)\n\n\ndef per_tensor_fp8_quantize(\n x: torch.Tensor,\n *,\n fp8_dtype: torch.dtype = torch.float8_e4m3fn,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Per-tensor FP8 quantization, mirroring ``tests/utils_fp8.py:to_float8``.\n\n Returns ``(x_fp8, inv_scale)`` where ``inv_scale`` is the dequant\n multiplier (``float \u2248 fp8 * inv_scale``).\n \"\"\"\n finfo = torch.finfo(fp8_dtype)\n amax = x.abs().amax().clamp(min=1e-12)\n scale = finfo.max / amax\n x_q = (x.float() * scale).clamp(min=finfo.min, max=finfo.max).to(fp8_dtype)\n return x_q, scale.float().reciprocal()\n\n\ndef fp8_block_quant_1d(\n x_bf16: torch.Tensor,\n block: int = 128,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Quantize ``[T, H]`` activations into FP8 with per-``(token, block)``\n column-block scales. Returns ``(x_fp8, scales)`` where\n ``scales`` has shape ``[T, H // block]``.\n\n Mirrors ``_fp8_block_quant_1d`` in\n ``tests/moe/test_dpsk_fused_moe_fp8.py``.\n \"\"\"\n assert x_bf16.dim() == 2\n T, H = x_bf16.shape\n assert H % block == 0\n nb = H // block\n finfo = torch.finfo(torch.float8_e4m3fn)\n max_fp8 = finfo.max\n x_f32 = x_bf16.to(torch.float32)\n x_fp8 = torch.empty((T, H), dtype=torch.float8_e4m3fn, device=x_bf16.device)\n scales = torch.empty((T, nb), dtype=torch.float32, device=x_bf16.device)\n for j in range(nb):\n sl = slice(j * block, (j + 1) * block)\n blk = x_f32[:, sl]\n amax = torch.amax(torch.abs(blk), dim=1)\n s = torch.where(amax > 0, amax / max_fp8, torch.ones_like(amax))\n x_fp8[:, sl] = (blk / s.unsqueeze(1)).to(torch.float8_e4m3fn)\n scales[:, j] = s\n return x_fp8, scales\n\n\ndef fp8_block_quant_2d(\n w_bf16: torch.Tensor,\n block: int = 128,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Quantize weights ``[..., R, C]`` with 2-D ``block \u00d7 block`` scales.\n\n Returns ``(w_fp8, scales)`` where ``scales`` has shape\n ``[..., R // block, C // block]``. Mirrors ``_fp8_block_quant_2d`` in\n ``tests/moe/test_dpsk_fused_moe_fp8.py``.\n \"\"\"\n assert w_bf16.dim() >= 2\n *prefix, R, C = w_bf16.shape\n assert R % block == 0 and C % block == 0\n nb_r, nb_c = R // block, C // block\n finfo = torch.finfo(torch.float8_e4m3fn)\n max_fp8 = finfo.max\n w_f32 = w_bf16.to(torch.float32).contiguous()\n prefix_ndim = len(prefix)\n reshaped = w_f32.reshape(*prefix, nb_r, block, nb_c, block)\n permute_dims = tuple(range(prefix_ndim)) + (\n prefix_ndim,\n prefix_ndim + 2,\n prefix_ndim + 1,\n prefix_ndim + 3,\n )\n blocks = reshaped.permute(permute_dims).contiguous()\n amax = torch.amax(torch.abs(blocks), dim=(-1, -2))\n scales = torch.where(\n amax > 0, amax / max_fp8, torch.ones_like(amax, dtype=torch.float32)\n )\n q_blocks = (blocks / scales.unsqueeze(-1).unsqueeze(-1)).to(torch.float8_e4m3fn)\n inv_permute = [0] * (prefix_ndim + 4)\n for i, p in enumerate(permute_dims):\n inv_permute[p] = i\n w_fp8 = q_blocks.permute(*inv_permute).reshape(*prefix, R, C).contiguous()\n return w_fp8, scales\n\n\n__all__ = [\n \"make_paged_kv_indices\",\n \"make_ragged_indptr\",\n \"make_uniform_qo_indptr\",\n \"make_pos_ids\",\n \"make_rope_cos_sin_cache\",\n \"make_probs\",\n \"make_logits\",\n \"fp8_safe_randn\",\n \"per_tensor_fp8_quantize\",\n \"fp8_block_quant_1d\",\n \"fp8_block_quant_2d\",\n]\n\n# ----- init -----\ndef _nvfp4_quantize_smooth_init(\n *,\n M: int,\n N: int = 3072,\n N_half: int = 0,\n SF: int = 0,\n device: str = \"cuda\",\n seed: int = 0,\n):\n \"\"\"Build inputs for ``flashinfer.nvfp4_quantize_smooth``.\"\"\"\n del N_half, SF # output-only / derived axes\n torch.manual_seed(seed)\n x = torch.randn(M, N, dtype=torch.bfloat16, device=device)\n pqs = torch.rand(N, dtype=torch.bfloat16, device=device) + 0.5\n x_hat = (x.float() * pqs.float()).to(torch.bfloat16)\n gs = (\n ((448 * 6) / x_hat.float().abs().nan_to_num().max())\n .to(torch.float32)\n .reshape(1)\n )\n return {\"x\": x, \"pre_quant_scale\": pqs, \"global_scale\": gs}\n" +} diff --git a/tests/trace/test_fi_trace_template_consistency.py b/tests/trace/test_fi_trace_template_consistency.py index f966732cd04..6897662789d 100644 --- a/tests/trace/test_fi_trace_template_consistency.py +++ b/tests/trace/test_fi_trace_template_consistency.py @@ -54,6 +54,19 @@ # --------------------------------------------------------------------------- +def test_svdquant_trace_activation_scale_tracks_variable_m(): + from flashinfer.trace.templates.gemm import mm_nvfp4_svdquant_trace + + assert isinstance(mm_nvfp4_svdquant_trace.axes["M"], Var) + assert isinstance(mm_nvfp4_svdquant_trace.axes["SF_A"], Var) + assert any( + constraint.startswith("SF_A ==") + and "M" in constraint + and "K_packed" in constraint + for constraint in mm_nvfp4_svdquant_trace.constraints + ) + + def _resolved_param(json_key: str, descriptor) -> str: """Return the function-parameter name that descriptor maps to.""" p = getattr(descriptor, "param", None) From c035d91e11a643645fb66d060aacf55e2518a85a Mon Sep 17 00:00:00 2001 From: Anthony Chang <27950904+rosenrodt@users.noreply.github.com> Date: Sat, 8 Aug 2026 10:38:59 +0800 Subject: [PATCH 2/7] Optimize complete SM120 SVDQuant execution Reduce end-to-end SVDQuant linear overhead and align SM120 behavior with the established SM100 alpha and epilogue contract. Changes - Fuse activation smoothing into NVFP4 quantization and vectorize BF16 correction staging - Add implementation and tactic autotuning with flattened storage and cache contracts - Extend dense BF16 baselines, benchmark reporting, trace integration, and focused coverage Validation - Preserve the original complete-linear, boundary-tactic, quantizer, and trace acceptance changes - Run applicable repository commit hooks during history reconstruction Result - Complete SM120 SVDQuant uses the fused quantizer and tensor-core low-rank path - Backend selection and alpha normalization match the public cross-architecture contract --- benchmarks/bench_nvfp4_svdquant_gemm.py | 229 ++++++++-- flashinfer/gemm/gemm_svdquant.py | 343 +++++++++++++-- .../dense_blockscaled_gemm_sm120_b12x.py | 416 ++++++++++++++---- .../quantization/kernels/nvfp4_quantize.py | 125 +++++- .../quantization_cute_dsl_utils.py | 44 ++ flashinfer/trace/templates/gemm.py | 4 +- tests/gemm/test_nvfp4_svdquant_gemm.py | 90 +++- tests/trace/example.py | 1 + .../fi_trace_out/gemm_bf16_N32_K3072.json | 51 +++ .../quantize_nvfp4_smooth_N3072.json | 2 +- 10 files changed, 1148 insertions(+), 157 deletions(-) create mode 100644 tests/trace/fi_trace_out/gemm_bf16_N32_K3072.json diff --git a/benchmarks/bench_nvfp4_svdquant_gemm.py b/benchmarks/bench_nvfp4_svdquant_gemm.py index 77691aed323..7a0b5122a91 100644 --- a/benchmarks/bench_nvfp4_svdquant_gemm.py +++ b/benchmarks/bench_nvfp4_svdquant_gemm.py @@ -1,17 +1,33 @@ #!/usr/bin/env python3 -"""Benchmark for the SM100 NVFP4 SVDQuant fused GEMM (Qwen-Image linear shapes). +"""Benchmark NVFP4 SVDQuant on SM100/SM103 and SM120/SM121 GPUs. -For every (n, k) x m problem this script times three things after autotuning: - 1. mm_nvfp4_svdquant : fused residual NVFP4 GEMM + rank-r BF16 LoRA-up + bias - 2. svdquant_linear : the full chain (nvfp4_quantize_smooth -> bf16 LoRA-down +For every (n, k) x m problem this script times five things after autotuning: + 1. mm_nvfp4_svdquant : selected SVDQuant implementation; auto tunes fused + versus unfused on SM120, or can be explicitly overridden + 2. unfused oracle : the same operation composed from separate SM120 kernels + 3. svdquant_linear : the full chain (nvfp4_quantize_smooth -> bf16 LoRA-down GEMM -> fused GEMM) - 3. mm_fp4 (cutlass) : the stock NVFP4 GEMM on the same residual operands + 4. mm_fp4 : the stock NVFP4 GEMM on the same residual operands (no LoRA correction), as the lower-bound baseline + 5. bf16 linear : conventional dense BF16 linear GEMM + bias on the + unquantized activation and weight + +The reported algorithmic TFLOPS/s count matmul operations only: + * fused GEMM: 2*m*n*k + 2*m*n*rank + * svdquant_linear: fused GEMM + 2*m*k*rank (LoRA-down) + * mm_fp4: 2*m*n*k + * bf16 linear: 2*m*n*k + +Quantization, alpha scaling, and bias addition are timed where applicable but +are not included in the operation count. The LoRA rank defaults to 32; pass e.g. --ranks 32,64,96,128 to sweep. +The SVDQuant backend defaults to auto; pass --svdquant-backend fused or +--svdquant-backend unfused to override implementation selection. Timing uses flashinfer.testing bench_gpu_time (CUPTI preferred, automatic -fallback to CUDA events). +fallback to CUDA events). CUDA Graph replay is enabled by default; pass +--no-cuda-graph to measure eager launches instead. """ import argparse @@ -39,6 +55,7 @@ def _build_case(m, n, k, rank, device): """Build all operands for one problem once (outside the timed region).""" + quantize_backend = "cute-dsl" if get_compute_capability(device)[0] == 12 else "cuda" x = torch.randn(m, k, dtype=torch.bfloat16, device=device) / (k**0.25) pqs = ( (1.0 + 0.3 * torch.randn(k, dtype=torch.bfloat16, device=device)) @@ -54,14 +71,24 @@ def _build_case(m, n, k, rank, device): w = torch.randn(n, k, dtype=torch.bfloat16, device=device) / (k**0.25) gw = ((448.0 * 6.0) / w.float().abs().nan_to_num().max()).reshape(1) - wq, w_sf = nvfp4_quantize(w, gw, sfLayout=SfLayout.layout_128x4, do_shuffle=False) + wq, w_sf = nvfp4_quantize( + w, + gw, + sfLayout=SfLayout.layout_128x4, + do_shuffle=False, + backend=quantize_backend, + ) wq = wq.view(torch.uint8) w_sf = w_sf.view(torch.uint8) alpha = (1.0 / (global_sf * gw)).reshape(1).float() # Quantized activation (byte-identical to nvfp4_quantize_smooth(x, pqs, gs)). xq, x_sf = nvfp4_quantize( - smoothed, global_sf, sfLayout=SfLayout.layout_128x4, do_shuffle=False + smoothed, + global_sf, + sfLayout=SfLayout.layout_128x4, + do_shuffle=False, + backend=quantize_backend, ) xq = xq.view(torch.uint8) x_sf = x_sf.view(torch.uint8) @@ -81,6 +108,7 @@ def _build_case(m, n, k, rank, device): "x_sf": x_sf, # 2-D swizzled layout (mm_fp4 convention) "x_sf_flat": x_sf.reshape(-1), # 1-D buffer (fused-kernel convention) "wq": wq, + "w": w, "w_sf": w_sf, "w_sf_flat": w_sf.reshape(-1), "alpha": alpha, @@ -90,6 +118,7 @@ def _build_case(m, n, k, rank, device): "bias": bias, "out_fused": torch.empty(m, n, dtype=torch.bfloat16, device=device), "out_fp4": torch.empty(m, n, dtype=torch.bfloat16, device=device), + "out_bf16": torch.empty(m, n, dtype=torch.bfloat16, device=device), } @@ -97,10 +126,35 @@ def _median_us(times_ms): return float(np.median(times_ms) * 1000.0) -def bench_one(m, n, k, rank, device): +def _matmul_flops(m, n, k, rank): + """Return (fused, full-linear, residual-only) algorithmic FLOP counts.""" + residual_flops = 2 * m * n * k + lora_up_flops = 2 * m * n * rank + lora_down_flops = 2 * m * k * rank + fused_flops = residual_flops + lora_up_flops + return fused_flops, fused_flops + lora_down_flops, residual_flops + + +def _tflops_per_sec(flops, latency_us): + """Convert an operation count and latency in microseconds to TFLOPS/s.""" + return flops / (latency_us * 1e6) if latency_us > 0 else float("nan") + + +def bench_one( + m, + n, + k, + rank, + device, + mm_fp4_backend, + svdquant_backend, + unfused_backend=None, + use_cuda_graph=True, + cold_l2_cache=False, +): c = _build_case(m, n, k, rank, device) - def run_fused(): + def run_selected(): mm_nvfp4_svdquant( c["xq"], c["wq"], @@ -111,6 +165,7 @@ def run_fused(): c["l1_scaled"], bias=c["bias"], out=c["out_fused"], + backend=svdquant_backend, ) def run_linear(): @@ -124,6 +179,21 @@ def run_linear(): c["l1_scaled"], c["global_sf"], bias=c["bias"], + backend=svdquant_backend, + ) + + def run_unfused(): + mm_nvfp4_svdquant( + c["xq"], + c["wq"], + c["x_sf_flat"], + c["w_sf_flat"], + c["alpha"], + c["d"], + c["l1_scaled"], + bias=c["bias"], + out=c["out_fused"], + backend=unfused_backend, ) def run_mm_fp4(): @@ -137,73 +207,176 @@ def run_mm_fp4(): c["out_fp4"], block_size=16, use_8x4_sf_layout=False, - backend="cutlass", + backend=mm_fp4_backend, use_nvfp4=True, ) + def run_bf16_linear(): + torch.addmm( + c["bias"], + c["x"], + c["w"].T, + out=c["out_bf16"], + ) + # Tune once; subsequent calls replay the best tactic from the tuner cache. with autotune(True): for _ in range(3): - run_fused() + run_selected() + if unfused_backend is not None: + run_unfused() run_linear() run_mm_fp4() + run_bf16_linear() torch.cuda.synchronize() bench_kwargs = dict( dry_run_time_ms=100, repeat_time_ms=500, - use_cuda_graph=True, + use_cuda_graph=use_cuda_graph, enable_cupti=True, - cold_l2_cache=True, + cold_l2_cache=cold_l2_cache, + ) + selected_us = _median_us(bench_gpu_time(run_selected, **bench_kwargs)) + unfused_us = ( + _median_us(bench_gpu_time(run_unfused, **bench_kwargs)) + if unfused_backend is not None + else float("nan") ) - fused_us = _median_us(bench_gpu_time(run_fused, **bench_kwargs)) linear_us = _median_us(bench_gpu_time(run_linear, **bench_kwargs)) mm_fp4_us = _median_us(bench_gpu_time(run_mm_fp4, **bench_kwargs)) + bf16_linear_us = _median_us(bench_gpu_time(run_bf16_linear, **bench_kwargs)) - return fused_us, linear_us, mm_fp4_us + return selected_us, unfused_us, linear_us, mm_fp4_us, bf16_linear_us def main(): parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--nk-shapes", + type=lambda s: [tuple(map(int, shape.split("x"))) for shape in s.split(",")], + default=NK_SHAPES, + help="comma-separated NxK shapes (default: 3072x3072,12288x3072,3072x12288)", + ) + parser.add_argument( + "--m-values", + type=lambda s: [int(m) for m in s.split(",")], + default=M_VALUES, + help="comma-separated M values (default: 4096,6889,9216,16384)", + ) parser.add_argument( "--ranks", type=lambda s: [int(r) for r in s.split(",")], default=[SVDQUANT_LORA_RANK_GRANULARITY], help="comma-separated LoRA ranks to sweep (positive multiples of 32)", ) + parser.add_argument( + "--svdquant-backend", + choices=("auto", "fused", "unfused"), + default="auto", + help=( + "SVDQuant implementation policy: auto tunes fused versus unfused on " + "SM120; fused or unfused forces that implementation (default: auto)" + ), + ) + parser.add_argument( + "--cuda-graph", + action=argparse.BooleanOptionalAction, + default=True, + help="benchmark CUDA Graph replay (default); use --no-cuda-graph for eager", + ) + parser.add_argument( + "--cold-l2-cache", + action=argparse.BooleanOptionalAction, + default=False, + help="request cold-L2 timing (default: warm L2)", + ) args = parser.parse_args() if not torch.cuda.is_available(): - print("CUDA is not available; this benchmark requires an SM100-class GPU.") + print("CUDA is not available; this benchmark requires a Blackwell GPU.") sys.exit(1) major, minor = get_compute_capability(torch.device(device="cuda")) - if major != 10: + if (major, minor) not in ((10, 0), (10, 3), (12, 0), (12, 1)): print( - "NVFP4 SVDQuant kernels require SM100-class GPUs (Blackwell); " + "NVFP4 SVDQuant kernels require SM100/SM103 or SM120/SM121; " f"got SM{major}{minor}. Exiting." ) sys.exit(1) torch.manual_seed(0) device = torch.device("cuda") + mm_fp4_backend = "cutlass" if major == 10 else "b12x" + unfused_backend = "cute-dsl-unfused" if major == 12 else None + if major == 12: + svdquant_backend = { + "auto": "auto", + "fused": "cute-dsl", + "unfused": "cute-dsl-unfused", + }[args.svdquant_backend] + else: + if args.svdquant_backend == "unfused": + parser.error("--svdquant-backend unfused is only available on SM120/SM121") + svdquant_backend = "cutlass" if args.svdquant_backend == "fused" else "auto" print(f"Device: {torch.cuda.get_device_name(device)} (SM{major}{minor})") - print("Timing: median GPU time in us (CUPTI preferred, CUDA-event fallback)\n") + print(f"mm_fp4 baseline backend: {mm_fp4_backend}") + print(f"SVDQuant implementation policy: {args.svdquant_backend}") + print("BF16 linear baseline: torch.addmm (PyTorch-selected CUDA backend)") + print(f"unfused oracle backend: {unfused_backend or 'not available'}") + print(f"execution mode: {'CUDA graph' if args.cuda_graph else 'eager'}") + print(f"L2 mode: {'cold' if args.cold_l2_cache else 'warm'}") + print("Timing: median GPU time in us (CUPTI preferred, CUDA-event fallback)") + print("TFLOPS/s: algorithmic matmul operations; see module docstring\n") header = ( - f"{'n':>6} {'k':>6} {'m':>6} {'rank':>5} | {'fused GEMM':>12} " - f"{'svdq linear':>12} {'mm_fp4':>12} | {'fused/mm_fp4':>12}" + f"{'n':>6} {'k':>6} {'m':>6} {'rank':>5} | " + f"{'selected us':>11} {'selected TF/s':>13} | " + f"{'unfused us':>11} {'unfused TF/s':>13} {'unfused/selected':>17} | " + f"{'linear us':>10} {'linear TF/s':>11} | " + f"{'mm_fp4 us':>10} {'mm_fp4 TF/s':>11} | {'mm_fp4/selected':>15}" + f" | {'bf16 us':>9} {'bf16 TF/s':>10}" ) print(header) print("-" * len(header)) for rank in args.ranks: - for n, k in NK_SHAPES: - for m in M_VALUES: - fused_us, linear_us, mm_fp4_us = bench_one(m, n, k, rank, device) - ratio = fused_us / mm_fp4_us if mm_fp4_us > 0 else float("nan") + for n, k in args.nk_shapes: + for m in args.m_values: + selected_us, unfused_us, linear_us, mm_fp4_us, bf16_linear_us = ( + bench_one( + m, + n, + k, + rank, + device, + mm_fp4_backend, + svdquant_backend, + unfused_backend, + args.cuda_graph, + args.cold_l2_cache, + ) + ) + fused_flops, linear_flops, mm_fp4_flops = _matmul_flops(m, n, k, rank) + selected_tflops = _tflops_per_sec(fused_flops, selected_us) + unfused_tflops = _tflops_per_sec(fused_flops, unfused_us) + linear_tflops = _tflops_per_sec(linear_flops, linear_us) + mm_fp4_tflops = _tflops_per_sec(mm_fp4_flops, mm_fp4_us) + bf16_linear_tflops = _tflops_per_sec(mm_fp4_flops, bf16_linear_us) + unfused_to_selected = ( + unfused_us / selected_us if selected_us > 0 else float("nan") + ) + mm_fp4_to_selected = ( + mm_fp4_us / selected_us if selected_us > 0 else float("nan") + ) print( - f"{n:>6} {k:>6} {m:>6} {rank:>5} | {fused_us:>12.2f} " - f"{linear_us:>12.2f} {mm_fp4_us:>12.2f} | {ratio:>12.3f}" + f"{n:>6} {k:>6} {m:>6} {rank:>5} | " + f"{selected_us:>11.2f} {selected_tflops:>13.2f} | " + f"{unfused_us:>11.2f} {unfused_tflops:>13.2f} " + f"{unfused_to_selected:>17.3f} | " + f"{linear_us:>10.2f} {linear_tflops:>11.2f} | " + f"{mm_fp4_us:>10.2f} {mm_fp4_tflops:>11.2f} | " + f"{mm_fp4_to_selected:>15.3f} | " + f"{bf16_linear_us:>9.2f} {bf16_linear_tflops:>10.2f}" ) print("-" * len(header)) diff --git a/flashinfer/gemm/gemm_svdquant.py b/flashinfer/gemm/gemm_svdquant.py index a5a0e8ad83d..bd81f9a8c8a 100644 --- a/flashinfer/gemm/gemm_svdquant.py +++ b/flashinfer/gemm/gemm_svdquant.py @@ -15,6 +15,7 @@ """ import functools +from dataclasses import replace from typing import List, Literal, Optional, Tuple import torch @@ -79,6 +80,7 @@ def _compile_sm120_nvfp4_svdquant( with_bias: bool, mma_tiler_mn: Tuple[int, int], swap_ab: bool, + use_prefetch: bool, sf_m: int, sf_n: int, sf_k: int, @@ -98,6 +100,7 @@ def _compile_sm120_nvfp4_svdquant( with_bias, mma_tiler_mn, swap_ab, + use_prefetch, max_active_clusters, enable_pdl, ) @@ -119,7 +122,7 @@ def _compile_sm120_nvfp4_svdquant( 16, mma_tiler_mn, (1, 1), - use_prefetch=False, + use_prefetch=use_prefetch, enable_pdl=enable_pdl, swap_ab=swap_ab, ) @@ -194,7 +197,8 @@ def compile_kernel(): kernel_name = ( f"r{rank}_bias{int(with_bias)}_t{mma_tiler_mn[0]}x{mma_tiler_mn[1]}" - f"_swap{int(swap_ab)}_mac{max_active_clusters}_pdl{int(enable_pdl)}" + f"_swap{int(swap_ab)}_pf{int(use_prefetch)}_mac{max_active_clusters}" + f"_pdl{int(enable_pdl)}" ) compiled = build_and_load_cute_dsl_kernel( "mm_nvfp4_svdquant_sm120", @@ -217,6 +221,7 @@ def _mm_nvfp4_svdquant_sm120_fused( bias: Optional[torch.Tensor], out: torch.Tensor, enable_pdl: bool, + tactic=None, ) -> torch.Tensor: from .kernels.dense_blockscaled_gemm_sm120_b12x import ( _select_default_dense_gemm_plan, @@ -228,15 +233,19 @@ def _mm_nvfp4_svdquant_sm120_fused( sf_m = (m + 127) // 128 sf_n = (n + 127) // 128 sf_k = (real_k // 16 + 3) // 4 - plan = _select_default_dense_gemm_plan( - m, n, real_k, get_device_sm_count(a.device), expected_m=m - ) + if tactic is None or tactic == -1: + plan = _select_default_dense_gemm_plan( + m, n, real_k, get_device_sm_count(a.device), expected_m=m + ) + tactic = (plan.mma_tiler_mn, plan.swap_ab, False) + mma_tiler_mn, swap_ab, use_prefetch = tactic compiled = _compile_sm120_nvfp4_svdquant( device=a.device, rank=d.shape[1], with_bias=bias is not None, - mma_tiler_mn=plan.mma_tiler_mn, - swap_ab=plan.swap_ab, + mma_tiler_mn=mma_tiler_mn, + swap_ab=swap_ab, + use_prefetch=use_prefetch, sf_m=sf_m, sf_n=sf_n, sf_k=sf_k, @@ -260,6 +269,227 @@ def _mm_nvfp4_svdquant_sm120_fused( return out +def _mm_nvfp4_svdquant_sm120_unfused( + a: torch.Tensor, + b: torch.Tensor, + a_sf: torch.Tensor, + b_sf: torch.Tensor, + alpha: torch.Tensor, + d: torch.Tensor, + l1: torch.Tensor, + bias: Optional[torch.Tensor], + out: torch.Tensor, + enable_pdl: bool, +) -> torch.Tensor: + from .gemm_base import mm_fp4 + + m, k_packed = a.shape + n = b.shape[0] + sf_cols = k_packed * 2 // 16 + a_sf_2d = _view_128x4_sf(a_sf, m, sf_cols) + b_sf_2d = _view_128x4_sf(b_sf, n, sf_cols) + mm_fp4( + a, + b.T, + a_sf_2d, + b_sf_2d.T, + alpha, + torch.bfloat16, + out, + block_size=16, + use_8x4_sf_layout=False, + backend="b12x", + use_nvfp4=True, + enable_pdl=enable_pdl, + ) + correction = torch.mm(d, l1.T) + correction.mul_(alpha) + out.add_(correction) + if bias is not None: + out.add_(bias) + return out + + +def _sm120_nvfp4_svdquant_runner(enable_pdl: bool): + class Sm120Nvfp4SvdquantRunner(TunableRunner): + def get_valid_tactics( + self, + inputs: List[torch.Tensor], + profile: OptimizationProfile, + ) -> list: + import cutlass + + from ..cute_dsl.utils import torch_to_cutlass_dtype + from .kernels.dense_blockscaled_gemm_sm120_b12x import ( + Sm120B12xBlockScaledDenseGemmKernel, + _select_default_dense_gemm_plan, + ) + + a, b, _, _, _, d, _, _, out = inputs + m, k_packed = a.shape + n = b.shape[0] + real_k = k_packed * 2 + c_dtype = torch_to_cutlass_dtype(out.dtype) + tactics = [] + + def _add(mma_tiler_mn, swap_ab): + if not Sm120B12xBlockScaledDenseGemmKernel.can_implement( + cutlass.Float4E2M1FN, + cutlass.Float8E4M3FN, + 16, + c_dtype, + mma_tiler_mn, + (1, 1), + n, + real_k, + 1, + "k", + "k", + "n", + swap_ab=swap_ab, + svdquant_rank=d.shape[1], + ): + return + for use_prefetch in (False, True): + tactic = (mma_tiler_mn, swap_ab, use_prefetch) + if tactic not in tactics: + tactics.append(tactic) + + for mma_tiler_mn in [(64, 64), (64, 128), (128, 64), (128, 128)]: + _add(mma_tiler_mn, swap_ab=False) + + plan = _select_default_dense_gemm_plan( + m, n, real_k, get_device_sm_count(a.device), expected_m=m + ) + _add(plan.mma_tiler_mn, plan.swap_ab) + return tactics + + def forward( + self, + inputs: List[torch.Tensor], + tactic=None, + do_preparation: bool = False, + **kwargs, + ): + a, b, a_sf, b_sf, alpha, d, l1, bias, out = inputs + return _mm_nvfp4_svdquant_sm120_fused( + a, + b, + a_sf, + b_sf, + alpha, + d, + l1, + bias, + out, + enable_pdl, + tactic, + ) + + return Sm120Nvfp4SvdquantRunner() + + +def _sm120_nvfp4_svdquant_unfused_runner(enable_pdl: bool): + # Flatten the b12x runner into the outer fused-vs-unfused tuner. Calling + # mm_fp4 here would start a nested AutoTuner while the outer runner is under + # CUDA Graph capture; on a cold cache that leaves tensor initialization in + # the capture and fails before the unfused candidate can be profiled. + from .gemm_base import _b12x_gemm_fp4_runner + + fp4_runner = _b12x_gemm_fp4_runner( + 12, + 0, + enable_pdl, + torch.bfloat16, + True, + ) + workspace_buffers: dict[torch.device, torch.Tensor] = {} + + class Sm120Nvfp4SvdquantUnfusedRunner(TunableRunner): + def _fp4_inputs(self, inputs: List[torch.Tensor]) -> list: + a, b, a_sf, b_sf, alpha, _, _, _, out = inputs + m, k_packed = a.shape + n = b.shape[0] + sf_cols = k_packed * 2 // 16 + workspace_buffer = workspace_buffers.get(a.device) + if workspace_buffer is None: + workspace_buffer = _get_cache_buf( + "mm_fp4_workspace", + DEFAULT_WORKSPACE_SIZE, + a.device, + ) + workspace_buffers[a.device] = workspace_buffer + return [ + a, + b.T, + _view_128x4_sf(a_sf, m, sf_cols), + _view_128x4_sf(b_sf, n, sf_cols).T, + alpha, + torch.bfloat16, + out, + 16, + True, + workspace_buffer, + ] + + def get_valid_tactics( + self, + inputs: List[torch.Tensor], + profile: OptimizationProfile, + ) -> list: + return fp4_runner.get_valid_tactics(self._fp4_inputs(inputs), profile) + + def forward( + self, + inputs: List[torch.Tensor], + tactic=-1, + do_preparation: bool = False, + **kwargs, + ): + _, _, _, _, alpha, d, l1, bias, out = inputs + fp4_runner(inputs=self._fp4_inputs(inputs), tactic=tactic) + correction = torch.mm(d, l1.T) + correction.mul_(alpha) + out.add_(correction) + if bias is not None: + out.add_(bias) + return out + + return Sm120Nvfp4SvdquantUnfusedRunner() + + +_SM120_NVFP4_SVDQUANT_TUNING_CONFIG = TuningConfig( + use_cuda_graph=True, + use_cold_l2_cache=True, + dynamic_tensor_specs=( + DynamicTensorSpec( + (0,), + (0,), + get_hybrid_num_tokens_buckets, + map_to_hybrid_bucket_uncapped, + ), + ), + constraint_specs=( + ConstraintSpec( + 2, + 0, + lambda shapes: _swizzled_sf_size(shapes[0][0], shapes[0][1] * 2 // 16), + ), + ConstraintSpec(5, 0, lambda shapes: shapes[0][0]), + ConstraintSpec(8, 0, lambda shapes: shapes[0][0]), + ), +) + +# The unfused candidate is a composition of FP4 GEMM, BF16 GEMM, and pointwise +# kernels. Profiling it inside the same CUDA Graph tuner as the single fused +# kernel can leave capture active after a failed tactic. Keep graph profiling +# for fused tactic selection, but compare fused versus unfused in eager mode. +_SM120_NVFP4_SVDQUANT_IMPLEMENTATION_TUNING_CONFIG = replace( + _SM120_NVFP4_SVDQUANT_TUNING_CONFIG, + use_cuda_graph=False, +) + + @functools.cache def get_nvfp4_svdquant_module(): """JIT-build and load the SM100 CUTLASS NVFP4 SVDQuant module.""" @@ -508,7 +738,8 @@ def mm_nvfp4_svdquant( ``"cutlass"`` selects the fused SM100/SM103 implementation; ``"cute-dsl"`` selects the fused SM120/SM121 implementation; ``"cute-dsl-unfused"`` selects its compositional reference path; - ``"auto"`` (default) selects by compute capability. + ``"auto"`` (default) selects by compute capability and, on SM120/SM121, + autotunes across both the fused and unfused implementations. enable_pdl: Optional[bool] Whether to launch with Programmatic Dependent Launch. Defaults to the device default. @@ -522,42 +753,45 @@ def mm_nvfp4_svdquant( if out is None: out = torch.empty(a.shape[0], b.shape[0], dtype=torch.bfloat16, device=a.device) + tune_sm120_implementations = False if backend == "auto": backend = mm_nvfp4_svdquant.suitable_auto_backends[0] + tune_sm120_implementations = backend == "cute-dsl" if backend == "cute-dsl": - return _mm_nvfp4_svdquant_sm120_fused( - a, b, a_sf, b_sf, alpha, d, l1, bias, out, enable_pdl + inputs = [a, b, a_sf, b_sf, alpha, d, l1, bias, out] + runners = [_sm120_nvfp4_svdquant_runner(enable_pdl)] + custom_op = "nvfp4_svdquant_gemm_sm120" + if tune_sm120_implementations: + runners.append(_sm120_nvfp4_svdquant_unfused_runner(enable_pdl)) + custom_op = "nvfp4_svdquant_gemm_sm120_auto" + tuning_config = ( + _SM120_NVFP4_SVDQUANT_IMPLEMENTATION_TUNING_CONFIG + if tune_sm120_implementations + else _SM120_NVFP4_SVDQUANT_TUNING_CONFIG + ) + runner, tactic = AutoTuner.get().choose_one( + custom_op, + runners, + tuning_config, + inputs, ) + runner(inputs=inputs, tactic=tactic) + return out if backend == "cute-dsl-unfused": - from .gemm_base import mm_fp4 - - m, k_packed = a.shape - n = b.shape[0] - sf_cols = k_packed * 2 // 16 - a_sf_2d = _view_128x4_sf(a_sf, m, sf_cols) - b_sf_2d = _view_128x4_sf(b_sf, n, sf_cols) - mm_fp4( + return _mm_nvfp4_svdquant_sm120_unfused( a, - b.T, - a_sf_2d, - b_sf_2d.T, + b, + a_sf, + b_sf, alpha, - torch.bfloat16, + d, + l1, + bias, out, - block_size=16, - use_8x4_sf_layout=False, - backend="b12x", - use_nvfp4=True, - enable_pdl=enable_pdl, + enable_pdl, ) - correction = torch.mm(d, l1.T) - correction.mul_(alpha) - out.add_(correction) - if bias is not None: - out.add_(bias) - return out workspace_buffer = _get_cache_buf( "nvfp4_svdquant_gemm_workspace", DEFAULT_WORKSPACE_SIZE, a.device @@ -617,8 +851,8 @@ def nvfp4_quantize_smooth( r"""Smooth + NVFP4 quantize: ``(xq, sf) = nvfp4-quantize(x * pre_quant_scale)``. The SM100/SM103 CUTLASS backend applies the SVDQuant per-input-channel smoothing scale - and NVFP4-quantizes in one pass. The SM120/SM121 CuTe DSL compatibility backend first - materializes the BF16 smoothed input and then invokes the CuTe DSL NVFP4 quantizer. + and NVFP4-quantizes in one pass. The SM120/SM121 CuTe DSL backend also + applies smoothing inside the NVFP4 quantizer, avoiding a BF16 intermediate. Both use ue4m3 block scales, the 128x4 swizzled layout, and SF vector size 16. Parameters @@ -633,7 +867,7 @@ def nvfp4_quantize_smooth( Whether to launch with Programmatic Dependent Launch. Defaults to the device default. backend: Literal["cutlass", "cute-dsl", "auto"] ``"cutlass"`` selects fused smoothing and quantization on SM100/SM103; - ``"cute-dsl"`` selects unfused smoothing plus CuTe DSL quantization on + ``"cute-dsl"`` selects fused smoothing plus CuTe DSL quantization on SM120/SM121; ``"auto"`` (default) selects by compute capability. Returns @@ -649,17 +883,15 @@ def nvfp4_quantize_smooth( if backend == "auto": backend = nvfp4_quantize_smooth.suitable_auto_backends[0] if backend == "cute-dsl": - from ..quantization.fp4_quantization import nvfp4_quantize - from ..tllm_enums import SfLayout + from ..quantization.kernels.nvfp4_quantize import ( + nvfp4_quantize_smooth_cute_dsl, + ) - xq, sf = nvfp4_quantize( - (x * pre_quant_scale).to(torch.bfloat16), + xq, sf = nvfp4_quantize_smooth_cute_dsl( + x, + pre_quant_scale, global_scale, - sfLayout=SfLayout.layout_128x4, - do_shuffle=False, - sf_vec_size=16, enable_pdl=enable_pdl, - backend="cute-dsl", ) return xq.view(torch.uint8), sf.view(torch.uint8).reshape(-1) @@ -691,7 +923,8 @@ def svdquant_linear( Runs the three-step chain this library's kernels are designed for: 1. ``xq, x_sf = nvfp4_quantize_smooth(x, pre_quant_scale, global_scale)`` - 2. ``down = x @ l2t_smoothed`` (plain BF16 GEMM; ``l2t_smoothed = pre_quant_scale[:, None] * L2ᵀ``) + 2. ``down = x @ l2t_smoothed`` (BF16 tensor-core GEMM; + ``l2t_smoothed = pre_quant_scale[:, None] * L2ᵀ``) 3. ``mm_nvfp4_svdquant(xq, weight_fp4, x_sf, weight_sf, alpha, down, l1_scaled, bias)`` The invariant per-layer transforms must be prepared offline by the caller: @@ -738,7 +971,25 @@ def svdquant_linear( enable_pdl=enable_pdl, backend=quantize_backend, ) - down = torch.mm(x, l2t_smoothed) + use_sm120_cute_dsl = quantize_backend == "cute-dsl" or ( + quantize_backend == "auto" + and nvfp4_quantize_smooth.suitable_auto_backends[0] == "cute-dsl" + ) + if use_sm120_cute_dsl: + # The rank-r projection is small enough that PyTorch's generic BF16 + # dispatch leaves substantial launch/selection overhead on SM120. + # FlashInfer's cuDNN runner selects and caches the BF16 tensor-core + # engine for this exact MxKxR problem. Keep torch.mm as the optional- + # dependency fallback; it also preserves the existing SM100 path. + from .gemm_base import CUDNN_AVAILABLE, mm_bf16 + + down = ( + mm_bf16(x, l2t_smoothed, backend="cudnn") + if CUDNN_AVAILABLE + else torch.mm(x, l2t_smoothed) + ) + else: + down = torch.mm(x, l2t_smoothed) return mm_nvfp4_svdquant( xq, weight_fp4, diff --git a/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py b/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py index 41b835aca44..2a6a10a141d 100644 --- a/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py +++ b/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py @@ -30,7 +30,7 @@ # and adapted for the current Blackwell GeForce target. from dataclasses import dataclass -from typing import Literal, Optional, Tuple +from typing import Any, Literal, Optional, Tuple import cuda.bindings.driver as cuda import cutlass @@ -41,7 +41,7 @@ import cutlass.utils.blockscaled_layout as blockscaled_utils import cutlass.utils.hopper_helpers as sm90_utils import logging -from cutlass import Float32, Int32, Int64 +from cutlass import BFloat16, Float32, Int32, Int64 from cutlass.cute.arch import griddepcontrol_launch_dependents, griddepcontrol_wait from cutlass.cute.nvgpu import cpasync from cutlass.cute.nvgpu.warp.mma import Field as WarpField @@ -286,6 +286,8 @@ def __init__( self.epi_stage = None self.a_smem_layout_staged = None self.b_smem_layout_staged = None + self.svdquant_a_smem_layout = None + self.svdquant_b_smem_layout = None self.epi_smem_layout_staged = None self.buffer_align_bytes = 1024 @@ -336,6 +338,41 @@ def _setup_attributes(self): ) # Bare atom for manual unroll workaround (avoids hasAuxTensor address space bug) self.mma_atom = cute.make_mma_atom(mma_op) + if cutlass.const_expr(self.svdquant_enabled): + svdquant_mma_op = cute.nvgpu.warp.MmaF16BF16Op( + BFloat16, + self.acc_dtype, + (16, 8, 16), + ) + self.svdquant_tiled_mma = cute.make_tiled_mma( + svdquant_mma_op, + atom_layout, + permutation_mnk=( + permutation_mnk[0], + permutation_mnk[1], + 16, + ), + ) + self.svdquant_a_smem_layout = sm90_utils.make_smem_layout_a( + utils.LayoutEnum.ROW_MAJOR, + ( + self.mma_tile_shape_mnk[0], + self.mma_tile_shape_mnk[1], + 16, + ), + BFloat16, + 1, + ) + self.svdquant_b_smem_layout = sm90_utils.make_smem_layout_b( + utils.LayoutEnum.ROW_MAJOR, + ( + self.mma_tile_shape_mnk[0], + self.mma_tile_shape_mnk[1], + 16, + ), + BFloat16, + 1, + ) # Compute atom loop bounds from tile shape and atom/layout shape # MMA atom: m16n8k64 for FP4. mma_m, mma_n, mma_k = 16, 8, self.mma_k @@ -371,6 +408,9 @@ def _setup_attributes(self): self.c_dtype, self.smem_capacity, self.occupancy, + (self.mma_tile_shape_mnk[0] + self.mma_tile_shape_mnk[1]) * 16 * 2 + if self.svdquant_enabled + else 0, ) assert self.epi_stage > 0, ( @@ -397,6 +437,9 @@ def _setup_attributes(self): self.sf_vec_size, self.tiled_mma, ) + if cutlass.const_expr(not self.svdquant_enabled): + self.svdquant_a_smem_layout = self.a_smem_layout_staged + self.svdquant_b_smem_layout = self.b_smem_layout_staged @cute.jit def __call__( @@ -440,6 +483,7 @@ def __call__( self.a_layout = utils.LayoutEnum.from_tensor(a) self.b_layout = utils.LayoutEnum.from_tensor(b) self.c_layout = utils.LayoutEnum.from_tensor(c) + self.svdquant_enabled = svdquant_d is not None if cutlass.const_expr(self.a_dtype != self.b_dtype): raise TypeError(f"Type mismatch: {self.a_dtype} != {self.b_dtype}") @@ -502,7 +546,7 @@ def __call__( ) @cute.struct - class SharedStorage: + class DenseSharedStorage: mainloop_pipeline_array_ptr: cute.struct.MemRange[ cutlass.Int64, self.ab_stage * 2 ] @@ -537,7 +581,53 @@ class SharedStorage: self.buffer_align_bytes, ] - self.shared_storage = SharedStorage + @cute.struct + class SvdquantSharedStorage: + dense: DenseSharedStorage + sSvdquantA: cute.struct.Align[ + cute.struct.MemRange[ + BFloat16, + cute.cosize(self.svdquant_a_smem_layout), + ], + 128, + ] + sSvdquantB: cute.struct.Align[ + cute.struct.MemRange[ + BFloat16, + cute.cosize(self.svdquant_b_smem_layout), + ], + 128, + ] + + @property + def mainloop_pipeline_array_ptr(self): + return self.dense.mainloop_pipeline_array_ptr + + @property + def sA(self): + return self.dense.sA + + @property + def sB(self): + return self.dense.sB + + @property + def sSFA(self): + return self.dense.sSFA + + @property + def sSFB(self): + return self.dense.sSFB + + @property + def sC(self): + return self.dense.sC + + shared_storage: Any = DenseSharedStorage + if cutlass.const_expr(self.svdquant_enabled): + shared_storage = SvdquantSharedStorage + + self.shared_storage = shared_storage self.kernel( tma_atom_a, @@ -563,6 +653,9 @@ class SharedStorage: self.sfa_smem_layout_staged, self.sfb_smem_layout_staged, self.epi_smem_layout_staged, + self.svdquant_tiled_mma if self.svdquant_enabled else None, + self.svdquant_a_smem_layout if self.svdquant_enabled else None, + self.svdquant_b_smem_layout if self.svdquant_enabled else None, tile_sched_params, epilogue_op, alpha, @@ -818,7 +911,31 @@ def _make_scale_tiled_copy( ) @cute.jit - def _predicate_cpasync_rows( + def _make_svdquant_gmem_tiled_copy( + self, + tile_rows: cutlass.Constexpr[int], + ) -> cute.TiledCopy: + copy_bits = 128 + copy_elems = copy_bits // BFloat16.width + threads_per_row = 16 // copy_elems + copy_rows = min( + tile_rows, + self.num_mma_warps * self.num_threads_per_warp // threads_per_row, + ) + copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + BFloat16, + num_bits_per_copy=copy_bits, + ) + thread_layout = cute.make_ordered_layout( + (copy_rows, threads_per_row), + order=(1, 0), + ) + value_layout = cute.make_layout((1, copy_elems)) + return cute.make_tiled_copy_tv(copy_atom, thread_layout, value_layout) + + @cute.jit + def _predicate_tiled_copy_rows( self, tCc: cute.Tensor, row_limit: Int32, @@ -839,6 +956,25 @@ def _predicate_cpasync_rows( tPred[rest_v, 0, rest_k] = tCc[(0, rest_v), 0, rest_k][0] < row_limit return tPred + @cute.jit + def _svdquant_tiled_copy_2d( + self, + tiled_copy: cute.TiledCopy, + gmem_tensor: cute.Tensor, + smem_tensor: cute.Tensor, + coord_tensor: cute.Tensor, + tidx: Int32, + copy_threads: cutlass.Constexpr[int], + row_limit: Int32, + ) -> None: + thr_copy = tiled_copy.get_slice(tidx) + tG = thr_copy.partition_S(gmem_tensor) + tS = thr_copy.partition_D(smem_tensor) + tC = thr_copy.partition_S(coord_tensor) + tP = self._predicate_tiled_copy_rows(tC, row_limit) + if tidx < copy_threads: + cute.copy(tiled_copy, tG, tS, pred=tP) + @cute.jit def _cpasync_copy_2d( self, @@ -850,7 +986,7 @@ def _cpasync_copy_2d( predicate_rows: cutlass.Constexpr[bool], ) -> None: if cutlass.const_expr(predicate_rows): - tP = self._predicate_cpasync_rows(tC, row_limit) + tP = self._predicate_tiled_copy_rows(tC, row_limit) for rest_m in cutlass.range_constexpr(cute.size(tS.shape[1])): if cutlass.const_expr(predicate_rows): cute.copy( @@ -913,6 +1049,9 @@ def kernel( sfa_smem_layout_staged: cute.Layout, sfb_smem_layout_staged: cute.Layout, epi_smem_layout_staged: cute.ComposedLayout, + svdquant_tiled_mma: Optional[cute.TiledMma], + svdquant_a_smem_layout: Optional[cute.ComposedLayout], + svdquant_b_smem_layout: Optional[cute.ComposedLayout], tile_sched_params: utils.PersistentTileSchedulerParams, epilogue_op: cutlass.Constexpr, alpha: cute.Tensor, @@ -1021,6 +1160,13 @@ def kernel( ) sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged) sSFB = storage.sSFB.get_tensor(sfb_smem_layout_staged) + if cutlass.const_expr(svdquant_d is not None): + sSvdquantA = storage.sSvdquantA.get_tensor( + svdquant_a_smem_layout.outer, swizzle=svdquant_a_smem_layout.inner + ) + sSvdquantB = storage.sSvdquantB.get_tensor( + svdquant_b_smem_layout.outer, swizzle=svdquant_b_smem_layout.inner + ) # Local_tile partition global tensors gA_mkl = cute.local_tile( @@ -1328,6 +1474,41 @@ def kernel( producer_group=tma_store_producer_group, ) + if cutlass.const_expr(svdquant_d is not None): + svdquant_thr_mma = svdquant_tiled_mma.get_slice(tidx) + tCsSvdquantA = svdquant_thr_mma.partition_A(sSvdquantA) + tCsSvdquantB = svdquant_thr_mma.partition_B(sSvdquantB) + tCrSvdquantA = svdquant_tiled_mma.make_fragment_A( + tCsSvdquantA[None, None, None, 0] + ) + tCrSvdquantB = svdquant_tiled_mma.make_fragment_B( + tCsSvdquantB[None, None, None, 0] + ) + svdquant_copy_atom_a = cute.make_copy_atom( + cute.nvgpu.warp.LdMatrix8x8x16bOp(False, 4), BFloat16 + ) + svdquant_copy_atom_b = cute.make_copy_atom( + cute.nvgpu.warp.LdMatrix8x8x16bOp(False, 4), BFloat16 + ) + svdquant_smem_copy_a = cute.make_tiled_copy_A( + svdquant_copy_atom_a, svdquant_tiled_mma + ) + svdquant_smem_copy_b = cute.make_tiled_copy_B( + svdquant_copy_atom_b, svdquant_tiled_mma + ) + svdquant_thr_copy_a = svdquant_smem_copy_a.get_slice(tidx) + svdquant_thr_copy_b = svdquant_smem_copy_b.get_slice(tidx) + tCsSvdquantA_copy = svdquant_thr_copy_a.partition_S(sSvdquantA) + tCsSvdquantB_copy = svdquant_thr_copy_b.partition_S(sSvdquantB) + tCrSvdquantA_copy = svdquant_thr_copy_a.retile(tCrSvdquantA) + tCrSvdquantB_copy = svdquant_thr_copy_b.retile(tCrSvdquantB) + svdquant_gmem_tiled_copy_a = self._make_svdquant_gmem_tiled_copy( + self.mma_tile_shape_mnk[0] + ) + svdquant_gmem_tiled_copy_b = self._make_svdquant_gmem_tiled_copy( + self.mma_tile_shape_mnk[1] + ) + while work_tile.is_valid_tile: tile_coord_mnl = work_tile.tile_idx gC_mnl_slice = gC_mnl[(None, None, *tile_coord_mnl)] @@ -1608,68 +1789,126 @@ def kernel( accumulators[None, _mt, _nt], ) - # SVDQuant fusion: accumulate the rank-r BF16 correction into the - # same FP32 registers as the NVFP4 residual before the epilogue. - # The first SM120 implementation uses scalar BF16 dot products; - # this keeps one device launch and the exact public contract while - # leaving BF16 warp-MMA staging as a transparent optimization. + # SVDQuant fusion: stage rank-16 BF16 chunks cooperatively, then + # accumulate them with warp MMA directly into the FP32 NVFP4 + # accumulator fragment. swap_ab also swaps the correction's A/B + # operands, so both MMA paths retain identical C-fragment ownership. + # Match the SM100 CUTLASS contract: svdquant_l1 is supplied as + # L1 / alpha. The common epilogue's alpha multiply therefore + # restores the unscaled correction: + # alpha * (residual + D @ (L1 / alpha).T) + # = alpha * residual + D @ L1.T. if cutlass.const_expr(svdquant_d is not None): - acc_mn = _reshape_acc_to_mn( - accumulators, - transpose=self.swap_ab, + rank = cute.size(svdquant_d, mode=[1]) + svdquant_a_rows = self.mma_tile_shape_mnk[0] + svdquant_b_rows = self.mma_tile_shape_mnk[1] + svdquant_threads_per_row = 2 + svdquant_copy_row_capacity = ( + self.num_mma_warps + * self.num_threads_per_warp + // svdquant_threads_per_row ) - c_identity = cute.make_identity_tensor( - ( - self.tile_shape_mnk[1], - self.tile_shape_mnk[0], - ) - if cutlass.const_expr(self.swap_ab) - else ( - self.tile_shape_mnk[0], - self.tile_shape_mnk[1], + svdquant_a_copy_threads = ( + min(svdquant_a_rows, svdquant_copy_row_capacity) + * svdquant_threads_per_row + ) + svdquant_b_copy_threads = ( + min(svdquant_b_rows, svdquant_copy_row_capacity) + * svdquant_threads_per_row + ) + tile_m, tile_n, _ = tile_coord_mnl + tile_rows_m, tile_rows_n, _ = self.tile_shape_mnk + svdquant_a_plan = ( + svdquant_d, + tile_m, + tile_rows_m, + ) + svdquant_b_plan = ( + svdquant_l1, + tile_n, + tile_rows_n, + ) + if cutlass.const_expr(self.swap_ab): + svdquant_a_plan, svdquant_b_plan = ( + svdquant_b_plan, + svdquant_a_plan, ) + ( + svdquant_a_source, + svdquant_a_tile_coord, + svdquant_a_tile_rows, + ) = svdquant_a_plan + ( + svdquant_b_source, + svdquant_b_tile_coord, + svdquant_b_tile_rows, + ) = svdquant_b_plan + svdquant_a_tiler = (svdquant_a_tile_rows, 16) + svdquant_b_tiler = (svdquant_b_tile_rows, 16) + gSvdquantA = cute.local_tile( + svdquant_a_source, + svdquant_a_tiler, + (svdquant_a_tile_coord, None), ) - coord_mn = _reshape_acc_to_mn( - thr_mma.partition_C(c_identity), - transpose=self.swap_ab, + gSvdquantB = cute.local_tile( + svdquant_b_source, + svdquant_b_tiler, + (svdquant_b_tile_coord, None), ) - rank = cute.size(svdquant_d, mode=[1]) - for acc_m in cutlass.range_constexpr(cute.size(acc_mn.shape[0])): - for acc_n in cutlass.range_constexpr( - cute.size(acc_mn.shape[1]) - ): - coord = coord_mn[acc_m, acc_n] - if cutlass.const_expr(self.swap_ab): - m_local = coord[1] - n_local = coord[0] - else: - m_local = coord[0] - n_local = coord[1] - m_coord = ( - tile_coord_mnl[0] * Int32(self.tile_shape_mnk[0]) - + m_local - ) - n_coord = ( - tile_coord_mnl[1] * Int32(self.tile_shape_mnk[1]) - + n_local - ) - if m_coord < Int32( - directC_mnl.shape[0] - ) and n_coord < Int32(directC_mnl.shape[1]): - correction = Float32(0.0) - for rank_idx in cutlass.range(rank, unroll=1): - correction += svdquant_d[(m_coord, rank_idx)].to( - Float32 - ) * svdquant_l1[(n_coord, rank_idx)].to(Float32) - if cutlass.const_expr(svdquant_bias is not None): - # The common epilogue multiplies the whole FP32 - # accumulator by alpha. Pre-dividing bias here - # retains out = alpha*(residual+correction)+bias. - correction += ( - svdquant_bias[(n_coord,)].to(Float32) - / alpha_value - ) - acc_mn[acc_m, acc_n] += correction + cSvdquantA = cute.local_tile( + cute.make_identity_tensor(cute.shape(svdquant_a_source)), + svdquant_a_tiler, + (svdquant_a_tile_coord, None), + ) + cSvdquantB = cute.local_tile( + cute.make_identity_tensor(cute.shape(svdquant_b_source)), + svdquant_b_tiler, + (svdquant_b_tile_coord, None), + ) + for rank_tile in cutlass.range_constexpr(rank // 16): + self._svdquant_tiled_copy_2d( + svdquant_gmem_tiled_copy_a, + gSvdquantA[(None, None, rank_tile)], + sSvdquantA[(None, None, 0)], + cSvdquantA[(None, None, rank_tile)], + tidx, + svdquant_a_copy_threads, + Int32(cute.size(svdquant_a_source, mode=[0])), + ) + self._svdquant_tiled_copy_2d( + svdquant_gmem_tiled_copy_b, + gSvdquantB[(None, None, rank_tile)], + sSvdquantB[(None, None, 0)], + cSvdquantB[(None, None, rank_tile)], + tidx, + svdquant_b_copy_threads, + Int32(cute.size(svdquant_b_source, mode=[0])), + ) + + # All MMA warps must see the cooperatively staged rank tile. + self.mma_sync_barrier.arrive_and_wait() + cute.copy( + svdquant_smem_copy_a, + tCsSvdquantA_copy[None, None, None, 0], + tCrSvdquantA_copy, + ) + cute.copy( + svdquant_smem_copy_b, + tCsSvdquantB_copy[None, None, None, 0], + tCrSvdquantB_copy, + ) + cute.gemm( + svdquant_tiled_mma, + accumulators, + tCrSvdquantA, + tCrSvdquantB, + accumulators, + ) + # Keep early warps from overwriting SMEM still read by peers. + self.mma_sync_barrier.arrive_and_wait() + + # Bias remains an elementwise output-epilogue contribution. + # The low-rank matrix product above is exclusively BF16 MMA. if cutlass.const_expr(self.swap_ab): acc_mn = _reshape_acc_to_mn(accumulators, transpose=True) @@ -1696,17 +1935,16 @@ def kernel( if m_coord < Int32( directC_mnl.shape[0] ) and n_coord < Int32(directC_mnl.shape[1]): + acc_value = alpha_value * acc_mn[acc_m, acc_n] + if cutlass.const_expr(svdquant_bias is not None): + acc_value += svdquant_bias[(n_coord,)].to(Float32) directC_mnl[ ( m_coord, n_coord, tile_coord_mnl[2], ) - ] = epilogue_op( - (alpha_value * acc_mn[acc_m, acc_n]).to( - self.c_dtype - ) - ) + ] = epilogue_op(acc_value.to(self.c_dtype)) if cutlass.const_expr(self.single_work_tile_per_cta): work_tile = WorkTileInfo( work_tile.tile_idx, @@ -1755,6 +1993,11 @@ def kernel( thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) tRS_sD = thr_copy_r2s.partition_D(sC) tRS_rAcc = tiled_copy_r2s.retile(accumulators) + if cutlass.const_expr(svdquant_bias is not None): + c_identity = cute.make_identity_tensor( + cute.slice_(self.tile_shape_mnk, (None, None, 0)) + ) + tRS_cD = tiled_copy_r2s.retile(thr_mma.partition_C(c_identity)) rD_shape = cute.shape(thr_copy_r2s.partition_S(sC)) tRS_rD_layout = cute.make_layout(rD_shape[:3]) @@ -1794,12 +2037,29 @@ def kernel( (None, mma_m_in_epi, mma_n_in_epi) ] tRS_rAcc_slice = tRS_rAcc[(None, mma_m, mma_n)] + if cutlass.const_expr(svdquant_bias is not None): + tRS_cD_slice = tRS_cD[(None, mma_m, mma_n)] for elem_idx in cutlass.range_constexpr( cute.size(tRS_rD_slice) ): - tRS_rD_slice[elem_idx] = tRS_rAcc_slice[ - elem_idx - ] + acc_value = tRS_rAcc_slice[elem_idx] + if cutlass.const_expr( + svdquant_bias is not None + ): + coord = tRS_cD_slice[elem_idx] + n_coord = ( + tile_coord_mnl[1] + * Int32(self.tile_shape_mnk[1]) + + coord[1] + ) + if n_coord < Int32(directC_mnl.shape[1]): + acc_value = ( + alpha_value * acc_value + + svdquant_bias[(n_coord,)].to( + Float32 + ) + ) + tRS_rD_slice[elem_idx] = acc_value gmem_coord = (epi_m, epi_n) if cutlass.const_expr(self.split_k_slices > 1): @@ -1932,9 +2192,12 @@ def kernel( acc_vec = tRS_rD.load() # Multiply alpha in FP32 before converting to c_dtype # to avoid overflow when c_dtype is FP16 - acc_vec = epilogue_op( - (alpha_value * acc_vec).to(self.c_dtype) - ) + if cutlass.const_expr(svdquant_bias is not None): + acc_vec = epilogue_op(acc_vec.to(self.c_dtype)) + else: + acc_vec = epilogue_op( + (alpha_value * acc_vec).to(self.c_dtype) + ) tRS_rD_out.store(acc_vec) # Register to shared memory @@ -2343,6 +2606,7 @@ def _compute_stages( c_dtype, smem_capacity: int, occupancy: int, + svdquant_bytes: int, ) -> tuple: epi_stage_max = (tile_shape_mnk[1] // epi_tile[1]) * ( tile_shape_mnk[0] // epi_tile[0] @@ -2367,6 +2631,7 @@ def _compute_stages( (smem_capacity - occupancy * 1024) // occupancy - mbar_helpers_bytes - epi_bytes + - svdquant_bytes ) // (ab_bytes_per_stage + sf_bytes_per_stage) ab_stage = max(1, min(raw_ab_stage, 4)) if tile_shape_mnk[0] in (16, 64) and tile_shape_mnk[1] == 128: @@ -2540,6 +2805,7 @@ def can_implement( *, load_path: str = "tma", swap_ab: bool = False, + svdquant_rank: Optional[int] = None, ) -> bool: # The current target only supports cluster (1,1) if cluster_shape_mn != (1, 1): @@ -2561,6 +2827,8 @@ def can_implement( return False if load_path == "cpasync" and (sf_vec_size != 16 or l != 1): return False + if svdquant_rank is not None and svdquant_rank % 16 != 0: + return False # SF smem still allocates full 128-element blocks even when the live # MMA tile uses only 16 or 32 rows or columns. if is_mxfp8: diff --git a/flashinfer/quantization/kernels/nvfp4_quantize.py b/flashinfer/quantization/kernels/nvfp4_quantize.py index fae7167b5d0..c99edad99aa 100644 --- a/flashinfer/quantization/kernels/nvfp4_quantize.py +++ b/flashinfer/quantization/kernels/nvfp4_quantize.py @@ -76,6 +76,7 @@ bfloat2x8_to_e2m1x16_packed, process_nvfp4_block_half, process_nvfp4_block_bfloat, + process_nvfp4_block_bfloat_smooth, process_nvfp4_block_fp8, process_nvfp4_silu_block_half, process_nvfp4_silu_block_bfloat, @@ -399,6 +400,7 @@ def __init__( disable_fp4_quant_fast_math: bool = False, nvfp4_4over6_config: NVFP44Over6Config | None = None, silu_and_mul: bool = False, + smooth_quant: bool = False, ): self.dtype = dtype self.K = K @@ -412,9 +414,13 @@ def __init__( self.sf_is_8x4 = sf_layout == SF_LAYOUT_8x4 # SwiGLU uses a 2*K-wide input and does not support FP8. self.silu_and_mul = silu_and_mul + self.smooth_quant = smooth_quant assert not (silu_and_mul and self.is_fp8), ( "SwiGLU fusion does not support fp8 input" ) + assert not (smooth_quant and (not self.is_bfloat16 or silu_and_mul)), ( + "smooth quantization requires plain BF16 input" + ) assert K % NVFP4_SF_VEC_SIZE == 0 self.num_sf_blocks_per_row = K // NVFP4_SF_VEC_SIZE @@ -444,13 +450,28 @@ def _compute_sf_offset( @cute.jit def _process_block( - self, row_input, elem_base: Int32, global_scale: Float32, row_amax: Float32 + self, + row_input, + pre_quant_scale, + elem_base: Int32, + global_scale: Float32, + row_amax: Float32, ): """Quantize one 16-element block, optionally fusing silu(gate) * up. For the SwiGLU path the up block sits self.K columns after the gate block in the same (2*K-wide) input row. """ + if cutlass.const_expr(self.smooth_quant): + return process_nvfp4_block_bfloat_smooth( + row_input, + pre_quant_scale, + elem_base, + global_scale, + self.disable_fp4_quant_fast_math, + self.nvfp4_4over6_config, + row_amax, + ) return _dispatch_process_nvfp4_block( row_input, elem_base, @@ -474,9 +495,18 @@ def __call__( padded_M: Int32, num_blocks: Int32, global_scale: Union[Float32, cute.Tensor], + pre_quant_scale: cute.Tensor, stream, ): - self.kernel(mInput, mOutput, mScales, M, padded_M, global_scale).launch( + self.kernel( + mInput, + mOutput, + mScales, + M, + padded_M, + global_scale, + pre_quant_scale, + ).launch( grid=[num_blocks, 1, 1], block=[self.num_threads, 1, 1], max_number_threads=[_MAX_THREADS_PER_BLOCK, 1, 1], @@ -495,6 +525,7 @@ def kernel( M: Int32, padded_M: Int32, global_scale: Union[Float32, cute.Tensor], + pre_quant_scale: cute.Tensor, ): """ Row-based kernel for swizzled layout. @@ -549,7 +580,11 @@ def kernel( # Apply optional SwiGLU and quantize the block. scale_fp8, packed64 = self._process_block( - row_input, elem_base, global_scale_value, row_amax + row_input, + pre_quant_scale, + elem_base, + global_scale_value, + row_amax, ) # Write scale factor using swizzled indexing @@ -611,7 +646,11 @@ def kernel( # Apply optional SwiGLU and quantize the block. scale_fp8, packed64 = self._process_block( - row_input, elem_base, global_scale_value, row_amax + row_input, + pre_quant_scale, + elem_base, + global_scale_value, + row_amax, ) # Write scale factor using swizzled indexing @@ -1443,6 +1482,7 @@ def _nvfp4_kernel_name( silu_and_mul: bool, nvfp4_4over6_config: NVFP44Over6Config | None = None, global_scale_is_tensor: bool = True, + smooth_quant: bool = False, ) -> str: """Specialization name within the nvfp4_quantize module, encoding every parameter that affects codegen. @@ -1452,6 +1492,8 @@ def _nvfp4_kernel_name( name += "_host_sf" if silu_and_mul: name += "_silu" + if smooth_quant: + name += "_smooth" if disable_fp4_quant_fast_math: name += "_nofastmath" if nvfp4_4over6_config is not None: @@ -1471,6 +1513,7 @@ def _get_compiled_kernel_nvfp4( nvfp4_4over6_config: NVFP44Over6Config | None = None, silu_and_mul: bool = False, global_scale_is_tensor: bool = True, + smooth_quant: bool = False, ) -> Tuple[Callable, int]: """ Get or compile NVFP4 kernel with TVM-FFI. @@ -1495,6 +1538,8 @@ def _get_compiled_kernel_nvfp4( "float8_e4m3fn": cutlass.Float8E4M3FN, } cutlass_dtype = _dtype_map[dtype_key] + if smooth_quant and (dtype_key != "bfloat16" or sf_layout == SF_LAYOUT_LINEAR): + raise ValueError("smooth quantization requires BF16 and a swizzled SF layout") # Use symbolic M for dynamic batch sizes sym_m = cute.sym_int() @@ -1518,6 +1563,9 @@ def _get_compiled_kernel_nvfp4( if global_scale_is_tensor else Float32(1.0) ) + pre_quant_scale_fake = cute.runtime.make_fake_compact_tensor( + cutlass_dtype, (K,), assumed_align=16 + ) stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) if sf_layout == SF_LAYOUT_LINEAR: @@ -1568,6 +1616,7 @@ def _get_compiled_kernel_nvfp4( disable_fp4_quant_fast_math=disable_fp4_quant_fast_math, nvfp4_4over6_config=nvfp4_4over6_config, silu_and_mul=silu_and_mul, + smooth_quant=smooth_quant, ) compiled_kernel = build_and_load_cute_dsl_kernel( @@ -1582,6 +1631,7 @@ def _get_compiled_kernel_nvfp4( silu_and_mul, nvfp4_4over6_config, global_scale_is_tensor, + smooth_quant, ), lambda: cute.compile( swizzled_obj, @@ -1592,6 +1642,7 @@ def _get_compiled_kernel_nvfp4( Int32(128), # Dummy padded_M Int32(1), # Dummy num_blocks global_scale_fake, + pre_quant_scale_fake, stream_fake, options="--enable-tvm-ffi", ), @@ -1995,6 +2046,9 @@ def nvfp4_quantize_cute_dsl( padded_m, num_blocks, global_scale_arg, + input[0] + if m > 0 + else torch.empty(k, dtype=input.dtype, device=input.device), ) # Reshape using padded_sf_cols: for swizzled layouts the buffer includes @@ -2004,6 +2058,68 @@ def nvfp4_quantize_cute_dsl( return fp4_output, scale_output +def nvfp4_quantize_smooth_cute_dsl( + input: torch.Tensor, + pre_quant_scale: torch.Tensor, + global_scale: torch.Tensor, + enable_pdl: bool | None = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Fuse BF16 channel smoothing into the 128x4 NVFP4 quantizer.""" + from ...utils import device_support_pdl + + if input.ndim != 2 or input.dtype != torch.bfloat16 or not input.is_cuda: + raise ValueError("input must be a 2-D CUDA BF16 tensor") + m, k = input.shape + if k % NVFP4_SF_VEC_SIZE != 0: + raise ValueError(f"K ({k}) must be divisible by {NVFP4_SF_VEC_SIZE}") + if ( + pre_quant_scale.dtype != torch.bfloat16 + or not pre_quant_scale.is_cuda + or pre_quant_scale.numel() != k + ): + raise ValueError("pre_quant_scale must be a CUDA BF16 tensor with K elements") + + input = input.contiguous() + pre_quant_scale = pre_quant_scale.reshape(k).contiguous() + global_scale_arg = global_scale.float().reshape(1).contiguous().to(input.device) + enable_pdl = device_support_pdl(input.device) if enable_pdl is not False else False + + num_sf_blocks_per_row = k // NVFP4_SF_VEC_SIZE + padded_m = _round_up(m, ROW_TILE_SIZE) + padded_sf_cols = _round_up(num_sf_blocks_per_row, 4) + scale_output_size = padded_m * padded_sf_cols + + kernel_fn, rows_per_block = _get_compiled_kernel_nvfp4( + "bfloat16", + k, + SF_LAYOUT_128x4, + enable_pdl, + _env_flag_enabled("FLASHINFER_DISABLE_FP4_QUANT_FAST_MATH"), + current_nvfp4_4over6_config(), + global_scale_is_tensor=True, + smooth_quant=True, + ) + num_blocks = min( + (padded_m + rows_per_block - 1) // rows_per_block, + get_num_sm(input.device) * _BLOCKS_PER_SM, + ) + fp4_output = torch.empty(m, k // 2, dtype=torch.uint8, device=input.device) + scale_output = torch.empty( + scale_output_size, dtype=torch.uint8, device=input.device + ) + kernel_fn( + input, + fp4_output, + scale_output, + m, + padded_m, + num_blocks, + global_scale_arg, + pre_quant_scale, + ) + return fp4_output, scale_output.reshape(-1, padded_sf_cols) + + def silu_and_mul_nvfp4_quantize_cute_dsl( input: torch.Tensor, global_scale: torch.Tensor, @@ -2167,6 +2283,7 @@ def silu_and_mul_nvfp4_quantize_cute_dsl( padded_m, num_blocks, global_scale_tensor, + input[0, :k], ) # Return scales in their layout-specific 2D shape. diff --git a/flashinfer/quantization/quantization_cute_dsl_utils.py b/flashinfer/quantization/quantization_cute_dsl_utils.py index e65ea429ca6..2494363377e 100644 --- a/flashinfer/quantization/quantization_cute_dsl_utils.py +++ b/flashinfer/quantization/quantization_cute_dsl_utils.py @@ -1956,6 +1956,49 @@ def process_nvfp4_block_bfloat( ) +@cute.jit +def process_nvfp4_block_bfloat_smooth( + row_tensor, + pre_quant_scale, + elem_base: Int32, + global_scale: Float32, + disable_fp4_quant_fast_math: bool = False, + nvfp4_4over6_config: NVFP44Over6Config | None = None, + row_amax: Float32 | None = None, +) -> tuple: + """Smooth and quantize one BF16 block without materializing ``x * scale``.""" + from ..cute_dsl.fp4_common import ( + bfloat2_mul, + get_ptr_as_int64, + ld_global_v4_u32, + ) + + ptr0 = get_ptr_as_int64(row_tensor, elem_base) + ptr1 = get_ptr_as_int64(row_tensor, elem_base + Int32(8)) + scale_ptr0 = get_ptr_as_int64(pre_quant_scale, elem_base) + scale_ptr1 = get_ptr_as_int64(pre_quant_scale, elem_base + Int32(8)) + + h0, h1, h2, h3 = ld_global_v4_u32(ptr0) + h4, h5, h6, h7 = ld_global_v4_u32(ptr1) + s0, s1, s2, s3 = ld_global_v4_u32(scale_ptr0) + s4, s5, s6, s7 = ld_global_v4_u32(scale_ptr1) + + return _quantize_nvfp4_from_h2x8_bfloat( + bfloat2_mul(h0, s0), + bfloat2_mul(h1, s1), + bfloat2_mul(h2, s2), + bfloat2_mul(h3, s3), + bfloat2_mul(h4, s4), + bfloat2_mul(h5, s5), + bfloat2_mul(h6, s6), + bfloat2_mul(h7, s7), + global_scale, + disable_fp4_quant_fast_math, + nvfp4_4over6_config, + row_amax, + ) + + @cute.jit def process_nvfp4_silu_block_half( row_tensor, @@ -2299,6 +2342,7 @@ def process_nvfp4_block_fp8( "bfloat2x8_to_e2m1x16_packed", "process_nvfp4_block_half", "process_nvfp4_block_bfloat", + "process_nvfp4_block_bfloat_smooth", "process_nvfp4_silu_block_half", "process_nvfp4_silu_block_bfloat", # High-level helper functions (NVFP4 - FP8 input) diff --git a/flashinfer/trace/templates/gemm.py b/flashinfer/trace/templates/gemm.py index aeddba293e4..36a0c15f742 100644 --- a/flashinfer/trace/templates/gemm.py +++ b/flashinfer/trace/templates/gemm.py @@ -2284,8 +2284,8 @@ def _nvfp4_quantize_smooth_init( op_type="quantize_nvfp4_smooth", description=( "Smooth + NVFP4 quantize: (xq, sf) = nvfp4-quantize(x * pre_quant_scale). " - "SM100/SM103 fuse the operations; SM120/SM121 materialize the BF16 " - "smoothed input before CuTe DSL quantization. Both use ue4m3 block " + "SM100/SM103 and SM120/SM121 fuse smoothing into quantization; the " + "SM120/SM121 path uses CuTe DSL. Both use ue4m3 block " "scales, 128x4 swizzled layout, and SF vector size 16." ), axes={ diff --git a/tests/gemm/test_nvfp4_svdquant_gemm.py b/tests/gemm/test_nvfp4_svdquant_gemm.py index 353816f8f48..96845086828 100644 --- a/tests/gemm/test_nvfp4_svdquant_gemm.py +++ b/tests/gemm/test_nvfp4_svdquant_gemm.py @@ -44,6 +44,36 @@ def test_nvfp4_svdquant_backend_arch_support(): assert not mm_nvfp4_svdquant.is_backend_supported("cute-dsl-unfused", 100) +def test_sm120_svdquant_can_implement_rejects_ragged_rank(): + _skip_unless_sm120() + import cutlass + + from flashinfer.gemm.kernels.dense_blockscaled_gemm_sm120_b12x import ( + Sm120B12xBlockScaledDenseGemmKernel, + ) + + common_args = ( + cutlass.Float4E2M1FN, + cutlass.Float8E4M3FN, + 16, + cutlass.BFloat16, + (64, 64), + (1, 1), + 128, + 128, + 1, + "k", + "k", + "n", + ) + assert Sm120B12xBlockScaledDenseGemmKernel.can_implement( + *common_args, svdquant_rank=32 + ) + assert not Sm120B12xBlockScaledDenseGemmKernel.can_implement( + *common_args, svdquant_rank=18 + ) + + def _skip_unless_sm100(): compute_capability = get_compute_capability(torch.device(device="cuda")) if compute_capability[0] != 10: @@ -501,11 +531,53 @@ def test_mm_nvfp4_svdquant_sm120_fused(use_bias): # single BF16 store, whereas the oracle rounds the residual and correction # in separate launches. Compare numerically, not bitwise. assert _sqnr_db(expected.float(), out.float()) > 35.0 - assert torch.equal(out_auto, out) + assert _sqnr_db(expected.float(), out_auto.float()) > 35.0 fp32_ref = p["ref_bias"] if use_bias else p["ref"] assert _sqnr_db(fp32_ref, out.float()) > 35.0 +@pytest.mark.parametrize("backend", ["cute-dsl", "auto"]) +def test_mm_nvfp4_svdquant_sm120_autotuned_replay(backend): + _skip_unless_sm120() + torch.manual_seed(0) + p = _make_gemm_problem( + 129, + 256, + 256, + rank=32, + quant_backend="cute-dsl", + residual_backend="b12x", + ) + + with autotune(True): + out = mm_nvfp4_svdquant( + p["xq"], + p["wq"], + p["x_sf_flat"], + p["w_sf_flat"], + p["alpha"], + p["d"], + p["l1_scaled"], + bias=p["bias"], + backend=backend, + ) + assert _sqnr_db(p["ref_bias"], out.float()) > 35.0 + + # Replay outside the tuning context must reuse the selected tactic. + out_replay = mm_nvfp4_svdquant( + p["xq"], + p["wq"], + p["x_sf_flat"], + p["w_sf_flat"], + p["alpha"], + p["d"], + p["l1_scaled"], + bias=p["bias"], + backend=backend, + ) + assert torch.equal(out_replay, out) + + def test_mm_nvfp4_svdquant_sm120_unfused_oracle(): _skip_unless_sm120() torch.manual_seed(0) @@ -690,8 +762,18 @@ def fail_torch_mm(*args, **kwargs): @pytest.mark.parametrize("use_bias", [False, True]) -def test_svdquant_linear_sm120_fused(use_bias): +def test_svdquant_linear_sm120_fused(use_bias, monkeypatch): _skip_unless_sm120() + import flashinfer.gemm.gemm_base as gemm_base + + mm_bf16_backends = [] + original_mm_bf16 = gemm_base.mm_bf16 + + def recording_mm_bf16(*args, **kwargs): + mm_bf16_backends.append(kwargs.get("backend")) + return original_mm_bf16(*args, **kwargs) + + monkeypatch.setattr(gemm_base, "mm_bf16", recording_mm_bf16) torch.manual_seed(0) m, n, k, rank = 129, 256, 256, 32 x = torch.randn(m, k, dtype=torch.bfloat16, device="cuda") / (k**0.25) @@ -753,6 +835,10 @@ def test_svdquant_linear_sm120_fused(use_bias): assert out.shape == (m, n) and out.dtype == torch.bfloat16 assert _sqnr_db(ref, out.float()) > 35.0 + if gemm_base.CUDNN_AVAILABLE: + assert mm_bf16_backends == ["cudnn"] + else: + assert mm_bf16_backends == [] @pytest.mark.parametrize("rank", [32, 128]) diff --git a/tests/trace/example.py b/tests/trace/example.py index c4a8ae516be..65ca299ae4c 100644 --- a/tests/trace/example.py +++ b/tests/trace/example.py @@ -25,6 +25,7 @@ packed_kda_decode_h12_d128.json fused_kda_decode_h12_d128.json gemm_bf16_N256_K7168.json +gemm_bf16_N32_K3072.json gemm_bf16_N4096_K4096.json gemm_fp4_N2048_K7168_block_size16.json gemm_fp8_N1536_K7168.json diff --git a/tests/trace/fi_trace_out/gemm_bf16_N32_K3072.json b/tests/trace/fi_trace_out/gemm_bf16_N32_K3072.json new file mode 100644 index 00000000000..441b9a210c2 --- /dev/null +++ b/tests/trace/fi_trace_out/gemm_bf16_N32_K3072.json @@ -0,0 +1,51 @@ +{ + "name": "gemm_bf16_N32_K3072", + "description": "General matrix multiply (GEMM) C = A @ B (B is column-major [K, N]).", + "op_type": "gemm_bf16", + "tags": [ + "fi_api:flashinfer.gemm.gemm_base.mm_bf16", + "status:verified" + ], + "axes": { + "M": { + "type": "var" + }, + "N": { + "type": "const", + "value": 32 + }, + "K": { + "type": "const", + "value": 3072 + } + }, + "inputs": { + "A": { + "shape": [ + "M", + "K" + ], + "dtype": "bfloat16" + }, + "B": { + "shape": [ + "K", + "N" + ], + "dtype": "bfloat16", + "description": "Weight matrix in column-major layout (physical shape [K, N])." + } + }, + "outputs": { + "C": { + "shape": [ + "M", + "N" + ], + "dtype": "bfloat16" + } + }, + "reference": "from __future__ import annotations\nimport math\nimport torch\nimport torch.nn.functional as F\n\ndef _mm_reference(A, B):\n # B is physically [K, N] (column-major weight), so C = A @ B.\n return torch.matmul(A, B)\n", + "check": "def _gemm_check(\n reference_outputs,\n actual_outputs,\n *,\n rtol=None,\n atol=None,\n max_mismatch_pct=100.0,\n min_cos_sim=0.99,\n):\n from flashinfer.trace import default_check\n\n # Matches tests/gemm/test_mm_bf16.py, test_mm_fp8.py, test_bmm_bf16.py,\n # and test_bmm_fp8.py, which gate these kernels by cosine similarity.\n return default_check(\n reference_outputs,\n actual_outputs,\n rtol=rtol,\n atol=atol,\n max_mismatch_pct=max_mismatch_pct,\n min_cos_sim=min_cos_sim,\n )\n", + "init": "from __future__ import annotations\nimport math\nimport torch\n\n# ----- shared init helpers -----\n# Copyright (c) 2025 by FlashInfer team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Shared helpers used by ``TraceTemplate.init`` functions.\n\nThis module contains the small set of input-construction patterns that\nrecur across many templates (paged-KV cache index arrays, ragged indptr,\nRoPE pos_ids and cos/sin caches, sampling probs). Each helper is short and\ndocumented; init functions in ``templates/.py`` call into here so\nthe per-template init bodies stay focused on shape/dtype, not boilerplate.\n\nThe full source of this module is **inlined into every dumped JSON's\n``\"init\"`` field** by ``flashinfer/trace/template.py:_render_init_source``,\nso downstream consumers don't need flashinfer installed to re-run the init\nsnippets.\n\"\"\"\n\n\nfrom typing import Optional, Tuple\n\nimport torch\n\n\ndef make_paged_kv_indices(\n batch_size: int,\n num_pages_per_seq: int,\n page_size: int,\n *,\n device: str = \"cuda\",\n) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n \"\"\"Return ``(kv_indptr, kv_indices, kv_last_page_len)`` for a uniform batch.\n\n Every sequence is assigned exactly ``num_pages_per_seq`` pages, fully\n populated (last-page length == page_size).\n\n Invariants\n ----------\n - ``kv_indptr.shape == (batch_size + 1,)``, dtype int32, monotonic, [0]=0.\n - ``kv_indices == arange(0, batch_size * num_pages_per_seq)``, int32.\n - ``kv_last_page_len == full(batch_size, page_size)``, int32.\n \"\"\"\n total_pages = batch_size * num_pages_per_seq\n kv_indptr = (\n torch.arange(batch_size + 1, dtype=torch.int32, device=device)\n * num_pages_per_seq\n )\n kv_indices = torch.arange(total_pages, dtype=torch.int32, device=device)\n kv_last_page_len = torch.full(\n (batch_size,), page_size, dtype=torch.int32, device=device\n )\n return kv_indptr, kv_indices, kv_last_page_len\n\n\ndef make_ragged_indptr(\n seg_lens,\n *,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.int32,\n) -> torch.Tensor:\n \"\"\"Return cumulative-sum ``indptr`` of length ``len(seg_lens)+1``.\n\n ``seg_lens`` may be a list / tuple / 1-D tensor of segment lengths.\n \"\"\"\n if isinstance(seg_lens, torch.Tensor):\n lens = seg_lens.to(device=device, dtype=dtype)\n else:\n lens = torch.tensor(list(seg_lens), dtype=dtype, device=device)\n indptr = torch.zeros(lens.numel() + 1, dtype=dtype, device=device)\n indptr[1:] = torch.cumsum(lens, dim=0).to(dtype)\n return indptr\n\n\ndef make_uniform_qo_indptr(\n batch_size: int,\n qo_len: int,\n *,\n device: str = \"cuda\",\n) -> torch.Tensor:\n \"\"\"Return ``[0, qo_len, 2*qo_len, ..., batch_size*qo_len]`` int32.\"\"\"\n return torch.arange(batch_size + 1, dtype=torch.int32, device=device) * qo_len\n\n\ndef make_pos_ids(\n nnz: int,\n max_seq_len: Optional[int] = None,\n *,\n device: str = \"cuda\",\n) -> torch.Tensor:\n \"\"\"Return ``[0, 1, ..., nnz-1] (% max_seq_len)`` as int32 on ``device``.\n\n If ``max_seq_len`` is None, no wrapping is applied.\n \"\"\"\n pos = torch.arange(nnz, dtype=torch.int32, device=device)\n if max_seq_len is not None:\n pos = pos % max_seq_len\n return pos\n\n\ndef make_rope_cos_sin_cache(\n max_seq_len: int,\n rope_dim: int,\n *,\n base: float = 1e4,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Return concatenated ``[cos | sin]`` cache of shape ``[max_seq_len, rope_dim]``.\"\"\"\n t = torch.arange(max_seq_len, dtype=torch.float32, device=device)\n inv = 1.0 / (\n base\n ** (torch.arange(0, rope_dim, 2, dtype=torch.float32, device=device) / rope_dim)\n )\n freqs = t.unsqueeze(-1) * inv.unsqueeze(0)\n cache = torch.cat([torch.cos(freqs), torch.sin(freqs)], dim=-1)\n return cache.to(dtype)\n\n\ndef make_probs(\n batch_size: int,\n vocab_size: int,\n *,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Return a ``[batch_size, vocab_size]`` probability distribution.\n\n Uses ``softmax(randn(...))`` so each row sums to 1.0. This mirrors the\n pattern used throughout ``tests/utils/test_sampling.py``.\n \"\"\"\n return torch.softmax(\n torch.randn(batch_size, vocab_size, dtype=torch.float32, device=device),\n dim=-1,\n ).to(dtype)\n\n\ndef make_logits(\n batch_size: int,\n vocab_size: int,\n *,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Return ``randn(batch_size, vocab_size)`` logits.\"\"\"\n return torch.randn(batch_size, vocab_size, dtype=dtype, device=device)\n\n\ndef fp8_safe_randn(\n *shape: int,\n scale: float = 0.1,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.bfloat16,\n) -> torch.Tensor:\n \"\"\"``randn(*shape) * scale`` \u2014 keeps values in the FP8/FP4 representable range.\n\n Tests for fp8/fp4 paths typically multiply ``randn`` by 0.1 to avoid\n saturation when quantizing. Use this helper to mirror that convention.\n \"\"\"\n return (torch.randn(*shape, dtype=dtype, device=device) * scale).to(dtype)\n\n\ndef per_tensor_fp8_quantize(\n x: torch.Tensor,\n *,\n fp8_dtype: torch.dtype = torch.float8_e4m3fn,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Per-tensor FP8 quantization, mirroring ``tests/utils_fp8.py:to_float8``.\n\n Returns ``(x_fp8, inv_scale)`` where ``inv_scale`` is the dequant\n multiplier (``float \u2248 fp8 * inv_scale``).\n \"\"\"\n finfo = torch.finfo(fp8_dtype)\n amax = x.abs().amax().clamp(min=1e-12)\n scale = finfo.max / amax\n x_q = (x.float() * scale).clamp(min=finfo.min, max=finfo.max).to(fp8_dtype)\n return x_q, scale.float().reciprocal()\n\n\ndef fp8_block_quant_1d(\n x_bf16: torch.Tensor,\n block: int = 128,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Quantize ``[T, H]`` activations into FP8 with per-``(token, block)``\n column-block scales. Returns ``(x_fp8, scales)`` where\n ``scales`` has shape ``[T, H // block]``.\n\n Mirrors ``_fp8_block_quant_1d`` in\n ``tests/moe/test_dpsk_fused_moe_fp8.py``.\n \"\"\"\n assert x_bf16.dim() == 2\n T, H = x_bf16.shape\n assert H % block == 0\n nb = H // block\n finfo = torch.finfo(torch.float8_e4m3fn)\n max_fp8 = finfo.max\n x_f32 = x_bf16.to(torch.float32)\n x_fp8 = torch.empty((T, H), dtype=torch.float8_e4m3fn, device=x_bf16.device)\n scales = torch.empty((T, nb), dtype=torch.float32, device=x_bf16.device)\n for j in range(nb):\n sl = slice(j * block, (j + 1) * block)\n blk = x_f32[:, sl]\n amax = torch.amax(torch.abs(blk), dim=1)\n s = torch.where(amax > 0, amax / max_fp8, torch.ones_like(amax))\n x_fp8[:, sl] = (blk / s.unsqueeze(1)).to(torch.float8_e4m3fn)\n scales[:, j] = s\n return x_fp8, scales\n\n\ndef fp8_block_quant_2d(\n w_bf16: torch.Tensor,\n block: int = 128,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Quantize weights ``[..., R, C]`` with 2-D ``block \u00d7 block`` scales.\n\n Returns ``(w_fp8, scales)`` where ``scales`` has shape\n ``[..., R // block, C // block]``. Mirrors ``_fp8_block_quant_2d`` in\n ``tests/moe/test_dpsk_fused_moe_fp8.py``.\n \"\"\"\n assert w_bf16.dim() >= 2\n *prefix, R, C = w_bf16.shape\n assert R % block == 0 and C % block == 0\n nb_r, nb_c = R // block, C // block\n finfo = torch.finfo(torch.float8_e4m3fn)\n max_fp8 = finfo.max\n w_f32 = w_bf16.to(torch.float32).contiguous()\n prefix_ndim = len(prefix)\n reshaped = w_f32.reshape(*prefix, nb_r, block, nb_c, block)\n permute_dims = tuple(range(prefix_ndim)) + (\n prefix_ndim,\n prefix_ndim + 2,\n prefix_ndim + 1,\n prefix_ndim + 3,\n )\n blocks = reshaped.permute(permute_dims).contiguous()\n amax = torch.amax(torch.abs(blocks), dim=(-1, -2))\n scales = torch.where(\n amax > 0, amax / max_fp8, torch.ones_like(amax, dtype=torch.float32)\n )\n q_blocks = (blocks / scales.unsqueeze(-1).unsqueeze(-1)).to(torch.float8_e4m3fn)\n inv_permute = [0] * (prefix_ndim + 4)\n for i, p in enumerate(permute_dims):\n inv_permute[p] = i\n w_fp8 = q_blocks.permute(*inv_permute).reshape(*prefix, R, C).contiguous()\n return w_fp8, scales\n\n\n__all__ = [\n \"make_paged_kv_indices\",\n \"make_ragged_indptr\",\n \"make_uniform_qo_indptr\",\n \"make_pos_ids\",\n \"make_rope_cos_sin_cache\",\n \"make_probs\",\n \"make_logits\",\n \"fp8_safe_randn\",\n \"per_tensor_fp8_quantize\",\n \"fp8_block_quant_1d\",\n \"fp8_block_quant_2d\",\n]\n\n# ----- init -----\ndef _mm_bf16_init(\n *,\n M: int,\n N: int = 4096,\n K: int = 4096,\n device: str = \"cuda\",\n seed: int = 0,\n):\n \"\"\"Build inputs for ``flashinfer.mm_bf16``.\n\n ``B`` is constructed as ``randn(N, K).T`` to get column-major [K, N]\n matching the example call.\n \"\"\"\n torch.manual_seed(seed)\n a = torch.randn(M, K, dtype=torch.bfloat16, device=device)\n b = torch.randn(N, K, dtype=torch.bfloat16, device=device).T # [K, N] col-major\n return {\"a\": a, \"b\": b}\n" +} diff --git a/tests/trace/fi_trace_out/quantize_nvfp4_smooth_N3072.json b/tests/trace/fi_trace_out/quantize_nvfp4_smooth_N3072.json index 7db08a189d9..15ac94c724a 100644 --- a/tests/trace/fi_trace_out/quantize_nvfp4_smooth_N3072.json +++ b/tests/trace/fi_trace_out/quantize_nvfp4_smooth_N3072.json @@ -1,6 +1,6 @@ { "name": "quantize_nvfp4_smooth_N3072", - "description": "Smooth + NVFP4 quantize: (xq, sf) = nvfp4-quantize(x * pre_quant_scale). SM100/SM103 fuse the operations; SM120/SM121 materialize the BF16 smoothed input before CuTe DSL quantization. Both use ue4m3 block scales, 128x4 swizzled layout, and SF vector size 16.", + "description": "Smooth + NVFP4 quantize: (xq, sf) = nvfp4-quantize(x * pre_quant_scale). SM100/SM103 and SM120/SM121 fuse smoothing into quantization; the SM120/SM121 path uses CuTe DSL. Both use ue4m3 block scales, 128x4 swizzled layout, and SF vector size 16.", "op_type": "quantize_nvfp4_smooth", "tags": [ "fi_api:flashinfer.gemm.gemm_svdquant.nvfp4_quantize_smooth", From 40f9e56a3afa183863276ad19ec90ee26f747309 Mon Sep 17 00:00:00 2001 From: Anthony Chang <27950904+rosenrodt@users.noreply.github.com> Date: Sat, 8 Aug 2026 10:39:32 +0800 Subject: [PATCH 3/7] Streamline the SM120 SVDQuant correction pipeline Make the production LoRA-up path share the residual mainloop stage ring and tune its tile geometry without reserving a separate shared-memory pipeline. Changes - Match BF16 correction tiles to released NVFP4 A/B stage capacity and alias their storage - Reuse the mainloop producer-consumer barriers and autotune M/N/K tile choices - Use torch MM for LoRA-down and report the final cold-L2 benchmark comparisons Validation - Preserve the original deadlock, correctness, custom-shape, and RTX PRO 6000 performance acceptance changes - Run applicable repository commit hooks during history reconstruction Result - The correction pipeline avoids redundant barriers and shared-memory reservations - Fused execution remains faster than the unfused oracle across the recorded production shapes --- benchmarks/bench_nvfp4_svdquant_gemm.py | 171 ++++-- flashinfer/gemm/gemm_svdquant.py | 67 ++- .../dense_blockscaled_gemm_sm120_b12x.py | 520 +++++++++--------- tests/gemm/test_nvfp4_svdquant_gemm.py | 61 +- 4 files changed, 454 insertions(+), 365 deletions(-) diff --git a/benchmarks/bench_nvfp4_svdquant_gemm.py b/benchmarks/bench_nvfp4_svdquant_gemm.py index 7a0b5122a91..10186056215 100644 --- a/benchmarks/bench_nvfp4_svdquant_gemm.py +++ b/benchmarks/bench_nvfp4_svdquant_gemm.py @@ -1,22 +1,20 @@ #!/usr/bin/env python3 """Benchmark NVFP4 SVDQuant on SM100/SM103 and SM120/SM121 GPUs. -For every (n, k) x m problem this script times five things after autotuning: - 1. mm_nvfp4_svdquant : selected SVDQuant implementation; auto tunes fused - versus unfused on SM120, or can be explicitly overridden - 2. unfused oracle : the same operation composed from separate SM120 kernels +For every (n, k) x m problem this script times six things after autotuning: + 1. fused : the fused residual + LoRA-up SVDQuant kernel + 2. unfused : the same operation composed from separate SM120 kernels 3. svdquant_linear : the full chain (nvfp4_quantize_smooth -> bf16 LoRA-down - GEMM -> fused GEMM) - 4. mm_fp4 : the stock NVFP4 GEMM on the same residual operands + GEMM -> selected fused/unfused output implementation) + 4. residual_fp4 : the stock NVFP4 GEMM on the same residual operands (no LoRA correction), as the lower-bound baseline - 5. bf16 linear : conventional dense BF16 linear GEMM + bias on the - unquantized activation and weight + 5. BF16 GEMM : dense BF16 residual GEMM + bias; timed but hidden + 6. FP8 per-tensor : dense FP8 residual GEMM; timed but hidden The reported algorithmic TFLOPS/s count matmul operations only: * fused GEMM: 2*m*n*k + 2*m*n*rank * svdquant_linear: fused GEMM + 2*m*k*rank (LoRA-down) - * mm_fp4: 2*m*n*k - * bf16 linear: 2*m*n*k + * residual_fp4: 2*m*n*k Quantization, alpha scaling, and bias addition are timed where applicable but are not included in the operation count. @@ -39,6 +37,7 @@ from flashinfer import ( SfLayout, autotune, + bmm_fp8, mm_fp4, mm_nvfp4_svdquant, nvfp4_quantize, @@ -53,6 +52,16 @@ M_VALUES = [4096, 6889, 9216, 16384] +def _to_float8(x, dtype=torch.float8_e4m3fn): + """Quantize one tensor with a single scale, outside the timed region.""" + finfo = torch.finfo(dtype) + min_value, max_value = x.aminmax() + amax = torch.maximum(min_value.abs(), max_value.abs()).clamp(min=1e-12) + scale = finfo.max / amax + quantized = (x * scale).clamp(min=finfo.min, max=finfo.max).to(dtype) + return quantized, scale.reciprocal().float() + + def _build_case(m, n, k, rank, device): """Build all operands for one problem once (outside the timed region).""" quantize_backend = "cute-dsl" if get_compute_capability(device)[0] == 12 else "cuda" @@ -99,6 +108,10 @@ def _build_case(m, n, k, rank, device): l1_scaled = (lora_b.float() / alpha).to(torch.bfloat16).contiguous() d = torch.mm(x, l2t_smoothed) # LoRA-down output for the fused-GEMM-only path bias = torch.randn(n, dtype=torch.bfloat16, device=device).contiguous() + x_fp8, x_fp8_scale = _to_float8(x) + # bmm_fp8 expects B as column-major [batch, k, n]. Quantizing W.T retains + # its column-major stride, while the singleton batch dimension is a view. + w_fp8, w_fp8_scale = _to_float8(w.T) return { "x": x, @@ -119,6 +132,11 @@ def _build_case(m, n, k, rank, device): "out_fused": torch.empty(m, n, dtype=torch.bfloat16, device=device), "out_fp4": torch.empty(m, n, dtype=torch.bfloat16, device=device), "out_bf16": torch.empty(m, n, dtype=torch.bfloat16, device=device), + "x_fp8": x_fp8.unsqueeze(0), + "x_fp8_scale": x_fp8_scale, + "w_fp8": w_fp8.unsqueeze(0), + "w_fp8_scale": w_fp8_scale, + "out_fp8": torch.empty(1, m, n, dtype=torch.bfloat16, device=device), } @@ -154,7 +172,9 @@ def bench_one( ): c = _build_case(m, n, k, rank, device) - def run_selected(): + fused_backend = "cute-dsl" if get_compute_capability(device)[0] == 12 else "cutlass" + + def run_fused(): mm_nvfp4_svdquant( c["xq"], c["wq"], @@ -165,10 +185,10 @@ def run_selected(): c["l1_scaled"], bias=c["bias"], out=c["out_fused"], - backend=svdquant_backend, + backend=fused_backend, ) - def run_linear(): + def run_svdquant_linear(): svdquant_linear( c["x"], c["wq"], @@ -211,7 +231,7 @@ def run_mm_fp4(): use_nvfp4=True, ) - def run_bf16_linear(): + def run_residual_gemm(): torch.addmm( c["bias"], c["x"], @@ -219,15 +239,27 @@ def run_bf16_linear(): out=c["out_bf16"], ) + def run_fp8_per_tensor(): + bmm_fp8( + c["x_fp8"], + c["w_fp8"], + c["x_fp8_scale"], + c["w_fp8_scale"], + torch.bfloat16, + out=c["out_fp8"], + backend="auto", + ) + # Tune once; subsequent calls replay the best tactic from the tuner cache. with autotune(True): for _ in range(3): - run_selected() + run_fused() if unfused_backend is not None: run_unfused() - run_linear() + run_svdquant_linear() run_mm_fp4() - run_bf16_linear() + run_residual_gemm() + run_fp8_per_tensor() torch.cuda.synchronize() bench_kwargs = dict( @@ -237,17 +269,25 @@ def run_bf16_linear(): enable_cupti=True, cold_l2_cache=cold_l2_cache, ) - selected_us = _median_us(bench_gpu_time(run_selected, **bench_kwargs)) + fused_us = _median_us(bench_gpu_time(run_fused, **bench_kwargs)) unfused_us = ( _median_us(bench_gpu_time(run_unfused, **bench_kwargs)) if unfused_backend is not None else float("nan") ) - linear_us = _median_us(bench_gpu_time(run_linear, **bench_kwargs)) + svdquant_linear_us = _median_us(bench_gpu_time(run_svdquant_linear, **bench_kwargs)) mm_fp4_us = _median_us(bench_gpu_time(run_mm_fp4, **bench_kwargs)) - bf16_linear_us = _median_us(bench_gpu_time(run_bf16_linear, **bench_kwargs)) - - return selected_us, unfused_us, linear_us, mm_fp4_us, bf16_linear_us + residual_gemm_us = _median_us(bench_gpu_time(run_residual_gemm, **bench_kwargs)) + fp8_per_tensor_us = _median_us(bench_gpu_time(run_fp8_per_tensor, **bench_kwargs)) + + return ( + fused_us, + unfused_us, + svdquant_linear_us, + mm_fp4_us, + residual_gemm_us, + fp8_per_tensor_us, + ) def main(): @@ -321,7 +361,11 @@ def main(): print(f"Device: {torch.cuda.get_device_name(device)} (SM{major}{minor})") print(f"mm_fp4 baseline backend: {mm_fp4_backend}") print(f"SVDQuant implementation policy: {args.svdquant_backend}") - print("BF16 linear baseline: torch.addmm (PyTorch-selected CUDA backend)") + print("BF16 residual_gemm baseline: torch.addmm (PyTorch-selected CUDA backend)") + print( + "FP8 per-tensor baseline: bmm_fp8 backend=auto " + "(scales and quantization excluded)" + ) print(f"unfused oracle backend: {unfused_backend or 'not available'}") print(f"execution mode: {'CUDA graph' if args.cuda_graph else 'eager'}") print(f"L2 mode: {'cold' if args.cold_l2_cache else 'warm'}") @@ -329,12 +373,10 @@ def main(): print("TFLOPS/s: algorithmic matmul operations; see module docstring\n") header = ( - f"{'n':>6} {'k':>6} {'m':>6} {'rank':>5} | " - f"{'selected us':>11} {'selected TF/s':>13} | " - f"{'unfused us':>11} {'unfused TF/s':>13} {'unfused/selected':>17} | " - f"{'linear us':>10} {'linear TF/s':>11} | " - f"{'mm_fp4 us':>10} {'mm_fp4 TF/s':>11} | {'mm_fp4/selected':>15}" - f" | {'bf16 us':>9} {'bf16 TF/s':>10}" + f"{'n':>6} {'k':>6} {'m':>6} {'R':>5} | " + f"{'svdquant_linear us':>18} {'gain vs BF16':>13} {'gain vs FP8':>12} | " + f"{'residual_fp4/fused':>20} {'residual_fp4 TF/s/us':>22} | " + f"{'fusion gain':>11} {'fused TF/s/us':>17} {'unfused TF/s/us':>19}" ) print(header) print("-" * len(header)) @@ -342,41 +384,56 @@ def main(): for rank in args.ranks: for n, k in args.nk_shapes: for m in args.m_values: - selected_us, unfused_us, linear_us, mm_fp4_us, bf16_linear_us = ( - bench_one( - m, - n, - k, - rank, - device, - mm_fp4_backend, - svdquant_backend, - unfused_backend, - args.cuda_graph, - args.cold_l2_cache, - ) + ( + fused_us, + unfused_us, + svdquant_linear_us, + mm_fp4_us, + residual_gemm_us, + fp8_per_tensor_us, + ) = bench_one( + m, + n, + k, + rank, + device, + mm_fp4_backend, + svdquant_backend, + unfused_backend, + args.cuda_graph, + args.cold_l2_cache, ) - fused_flops, linear_flops, mm_fp4_flops = _matmul_flops(m, n, k, rank) - selected_tflops = _tflops_per_sec(fused_flops, selected_us) + fused_flops, _, mm_fp4_flops = _matmul_flops(m, n, k, rank) + fused_tflops = _tflops_per_sec(fused_flops, fused_us) unfused_tflops = _tflops_per_sec(fused_flops, unfused_us) - linear_tflops = _tflops_per_sec(linear_flops, linear_us) mm_fp4_tflops = _tflops_per_sec(mm_fp4_flops, mm_fp4_us) - bf16_linear_tflops = _tflops_per_sec(mm_fp4_flops, bf16_linear_us) - unfused_to_selected = ( - unfused_us / selected_us if selected_us > 0 else float("nan") + gain_over_bf16 = ( + (residual_gemm_us / svdquant_linear_us - 1.0) * 100.0 + if svdquant_linear_us > 0 + else float("nan") + ) + gain_over_fp8 = ( + (fp8_per_tensor_us / svdquant_linear_us - 1.0) * 100.0 + if svdquant_linear_us > 0 + else float("nan") + ) + fusion_efficiency = ( + mm_fp4_us / fused_us if fused_us > 0 else float("nan") ) - mm_fp4_to_selected = ( - mm_fp4_us / selected_us if selected_us > 0 else float("nan") + fusion_gain = ( + (unfused_us / fused_us - 1.0) * 100.0 + if fused_us > 0 + else float("nan") ) print( f"{n:>6} {k:>6} {m:>6} {rank:>5} | " - f"{selected_us:>11.2f} {selected_tflops:>13.2f} | " - f"{unfused_us:>11.2f} {unfused_tflops:>13.2f} " - f"{unfused_to_selected:>17.3f} | " - f"{linear_us:>10.2f} {linear_tflops:>11.2f} | " - f"{mm_fp4_us:>10.2f} {mm_fp4_tflops:>11.2f} | " - f"{mm_fp4_to_selected:>15.3f} | " - f"{bf16_linear_us:>9.2f} {bf16_linear_tflops:>10.2f}" + f"{svdquant_linear_us:>18.2f} {gain_over_bf16:>12.1f}% " + f"{gain_over_fp8:>11.1f}% | " + f"{fusion_efficiency:>20.3f} " + f"{mm_fp4_tflops:>10.2f}/{mm_fp4_us:<9.2f} | " + f"{fusion_gain:>10.1f}% " + f"{fused_tflops:>8.2f}/{fused_us:<8.2f} " + f"{unfused_tflops:>9.2f}/{unfused_us:<9.2f}" ) print("-" * len(header)) diff --git a/flashinfer/gemm/gemm_svdquant.py b/flashinfer/gemm/gemm_svdquant.py index bd81f9a8c8a..d6092487517 100644 --- a/flashinfer/gemm/gemm_svdquant.py +++ b/flashinfer/gemm/gemm_svdquant.py @@ -79,6 +79,7 @@ def _compile_sm120_nvfp4_svdquant( rank: int, with_bias: bool, mma_tiler_mn: Tuple[int, int], + tile_k: int, swap_ab: bool, use_prefetch: bool, sf_m: int, @@ -99,6 +100,7 @@ def _compile_sm120_nvfp4_svdquant( rank, with_bias, mma_tiler_mn, + tile_k, swap_ab, use_prefetch, max_active_clusters, @@ -122,6 +124,7 @@ def _compile_sm120_nvfp4_svdquant( 16, mma_tiler_mn, (1, 1), + tile_k=tile_k, use_prefetch=use_prefetch, enable_pdl=enable_pdl, swap_ab=swap_ab, @@ -197,7 +200,7 @@ def compile_kernel(): kernel_name = ( f"r{rank}_bias{int(with_bias)}_t{mma_tiler_mn[0]}x{mma_tiler_mn[1]}" - f"_swap{int(swap_ab)}_pf{int(use_prefetch)}_mac{max_active_clusters}" + f"x{tile_k}_swap{int(swap_ab)}_pf{int(use_prefetch)}_mac{max_active_clusters}" f"_pdl{int(enable_pdl)}" ) compiled = build_and_load_cute_dsl_kernel( @@ -237,13 +240,14 @@ def _mm_nvfp4_svdquant_sm120_fused( plan = _select_default_dense_gemm_plan( m, n, real_k, get_device_sm_count(a.device), expected_m=m ) - tactic = (plan.mma_tiler_mn, plan.swap_ab, False) - mma_tiler_mn, swap_ab, use_prefetch = tactic + tactic = (plan.mma_tiler_mn, 128, plan.swap_ab, False) + mma_tiler_mn, tile_k, swap_ab, use_prefetch = tactic compiled = _compile_sm120_nvfp4_svdquant( device=a.device, rank=d.shape[1], with_bias=bias is not None, mma_tiler_mn=mma_tiler_mn, + tile_k=tile_k, swap_ab=swap_ab, use_prefetch=use_prefetch, sf_m=sf_m, @@ -332,7 +336,12 @@ def get_valid_tactics( c_dtype = torch_to_cutlass_dtype(out.dtype) tactics = [] - def _add(mma_tiler_mn, swap_ab): + def _add( + mma_tiler_mn, + tile_k, + swap_ab, + prefetch_candidates=(False, True), + ): if not Sm120B12xBlockScaledDenseGemmKernel.can_implement( cutlass.Float4E2M1FN, cutlass.Float8E4M3FN, @@ -348,20 +357,40 @@ def _add(mma_tiler_mn, swap_ab): "n", swap_ab=swap_ab, svdquant_rank=d.shape[1], + mainloop_tile_k=tile_k, ): return - for use_prefetch in (False, True): - tactic = (mma_tiler_mn, swap_ab, use_prefetch) + for use_prefetch in prefetch_candidates: + tactic = ( + mma_tiler_mn, + tile_k, + swap_ab, + use_prefetch, + ) if tactic not in tactics: tactics.append(tactic) - for mma_tiler_mn in [(64, 64), (64, 128), (128, 64), (128, 128)]: - _add(mma_tiler_mn, swap_ab=False) + for mma_tiler_mn in ((64, 64), (64, 128), (128, 64), (128, 128)): + _add(mma_tiler_mn, 128, swap_ab=False) plan = _select_default_dense_gemm_plan( m, n, real_k, get_device_sm_count(a.device), expected_m=m ) - _add(plan.mma_tiler_mn, plan.swap_ab) + _add(plan.mma_tiler_mn, 128, plan.swap_ab) + for tile_k in (64, 256): + _add( + plan.mma_tiler_mn, + tile_k, + plan.swap_ab, + prefetch_candidates=(False,), + ) + if m >= 256 and n >= 64: + _add( + (256, 64), + 128, + swap_ab=False, + prefetch_candidates=(True,), + ) return tactics def forward( @@ -971,25 +1000,7 @@ def svdquant_linear( enable_pdl=enable_pdl, backend=quantize_backend, ) - use_sm120_cute_dsl = quantize_backend == "cute-dsl" or ( - quantize_backend == "auto" - and nvfp4_quantize_smooth.suitable_auto_backends[0] == "cute-dsl" - ) - if use_sm120_cute_dsl: - # The rank-r projection is small enough that PyTorch's generic BF16 - # dispatch leaves substantial launch/selection overhead on SM120. - # FlashInfer's cuDNN runner selects and caches the BF16 tensor-core - # engine for this exact MxKxR problem. Keep torch.mm as the optional- - # dependency fallback; it also preserves the existing SM100 path. - from .gemm_base import CUDNN_AVAILABLE, mm_bf16 - - down = ( - mm_bf16(x, l2t_smoothed, backend="cudnn") - if CUDNN_AVAILABLE - else torch.mm(x, l2t_smoothed) - ) - else: - down = torch.mm(x, l2t_smoothed) + down = torch.mm(x, l2t_smoothed) return mm_nvfp4_svdquant( xq, weight_fp4, diff --git a/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py b/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py index 2a6a10a141d..3f77b5243aa 100644 --- a/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py +++ b/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py @@ -30,7 +30,7 @@ # and adapted for the current Blackwell GeForce target. from dataclasses import dataclass -from typing import Any, Literal, Optional, Tuple +from typing import Literal, Optional, Tuple import cuda.bindings.driver as cuda import cutlass @@ -288,14 +288,13 @@ def __init__( self.b_smem_layout_staged = None self.svdquant_a_smem_layout = None self.svdquant_b_smem_layout = None + self.svdquant_dtype = BFloat16 + self.svdquant_rank = None + self.svdquant_tile_k = None self.epi_smem_layout_staged = None self.buffer_align_bytes = 1024 - self.mma_sync_barrier = pipeline.NamedBarrier( - barrier_id=1, - num_threads=self.num_mma_warps * self.num_threads_per_warp, - ) self.epilog_sync_barrier = pipeline.NamedBarrier( barrier_id=2, num_threads=self.num_mma_warps * self.num_threads_per_warp, @@ -338,41 +337,6 @@ def _setup_attributes(self): ) # Bare atom for manual unroll workaround (avoids hasAuxTensor address space bug) self.mma_atom = cute.make_mma_atom(mma_op) - if cutlass.const_expr(self.svdquant_enabled): - svdquant_mma_op = cute.nvgpu.warp.MmaF16BF16Op( - BFloat16, - self.acc_dtype, - (16, 8, 16), - ) - self.svdquant_tiled_mma = cute.make_tiled_mma( - svdquant_mma_op, - atom_layout, - permutation_mnk=( - permutation_mnk[0], - permutation_mnk[1], - 16, - ), - ) - self.svdquant_a_smem_layout = sm90_utils.make_smem_layout_a( - utils.LayoutEnum.ROW_MAJOR, - ( - self.mma_tile_shape_mnk[0], - self.mma_tile_shape_mnk[1], - 16, - ), - BFloat16, - 1, - ) - self.svdquant_b_smem_layout = sm90_utils.make_smem_layout_b( - utils.LayoutEnum.ROW_MAJOR, - ( - self.mma_tile_shape_mnk[0], - self.mma_tile_shape_mnk[1], - 16, - ), - BFloat16, - 1, - ) # Compute atom loop bounds from tile shape and atom/layout shape # MMA atom: m16n8k64 for FP4. mma_m, mma_n, mma_k = 16, 8, self.mma_k @@ -408,15 +372,55 @@ def _setup_attributes(self): self.c_dtype, self.smem_capacity, self.occupancy, - (self.mma_tile_shape_mnk[0] + self.mma_tile_shape_mnk[1]) * 16 * 2 - if self.svdquant_enabled - else 0, ) assert self.epi_stage > 0, ( "epi_stage <= 0, not enough shared memory. This configuration will be skipped." ) + if cutlass.const_expr(self.svdquant_enabled): + # One BF16 correction tile aliases exactly one NVFP4 mainloop A/B + # stage. Keeping this invariant lets the correction reuse the same + # circular producer/consumer pipeline without stage grouping. + rank_elements_per_stage = self.tile_shape_mnk[2] // ( + self.svdquant_dtype.width // self.a_dtype.width + ) + self.svdquant_tile_k = rank_elements_per_stage + svdquant_mma_op = cute.nvgpu.warp.MmaF16BF16Op( + self.svdquant_dtype, + self.acc_dtype, + (16, 8, 16), + ) + self.svdquant_tiled_mma = cute.make_tiled_mma( + svdquant_mma_op, + atom_layout, + permutation_mnk=( + permutation_mnk[0], + permutation_mnk[1], + self.svdquant_tile_k, + ), + ) + self.svdquant_a_smem_layout = sm90_utils.make_smem_layout_a( + utils.LayoutEnum.ROW_MAJOR, + ( + self.mma_tile_shape_mnk[0], + self.mma_tile_shape_mnk[1], + self.svdquant_tile_k, + ), + self.svdquant_dtype, + self.ab_stage, + ) + self.svdquant_b_smem_layout = sm90_utils.make_smem_layout_b( + utils.LayoutEnum.ROW_MAJOR, + ( + self.mma_tile_shape_mnk[0], + self.mma_tile_shape_mnk[1], + self.svdquant_tile_k, + ), + self.svdquant_dtype, + self.ab_stage, + ) + ( self.a_smem_layout_staged, self.b_smem_layout_staged, @@ -484,9 +488,18 @@ def __call__( self.b_layout = utils.LayoutEnum.from_tensor(b) self.c_layout = utils.LayoutEnum.from_tensor(c) self.svdquant_enabled = svdquant_d is not None + if cutlass.const_expr(self.svdquant_enabled): + self.svdquant_dtype = svdquant_d.element_type + self.svdquant_rank = cute.size(svdquant_d, mode=[1]) if cutlass.const_expr(self.a_dtype != self.b_dtype): raise TypeError(f"Type mismatch: {self.a_dtype} != {self.b_dtype}") + if cutlass.const_expr( + self.svdquant_enabled and self.svdquant_dtype != BFloat16 + ): + raise TypeError( + f"SVDQuant rank operands must be BF16, got {self.svdquant_dtype}." + ) self._setup_attributes() @@ -536,6 +549,28 @@ def __call__( self.epi_smem_layout_staged, self.epi_tile, ) + if cutlass.const_expr(self.svdquant_enabled): + svdquant_a = svdquant_l1 if self.swap_ab else svdquant_d + svdquant_b = svdquant_d if self.swap_ab else svdquant_l1 + tma_atom_svdquant_a, tma_tensor_svdquant_a = ( + self._make_tma_atoms_and_tensors( + svdquant_a, + self.svdquant_a_smem_layout, + (self.mma_tile_shape_mnk[0], self.svdquant_tile_k), + 1, + ) + ) + tma_atom_svdquant_b, tma_tensor_svdquant_b = ( + self._make_tma_atoms_and_tensors( + svdquant_b, + self.svdquant_b_smem_layout, + (self.mma_tile_shape_mnk[1], self.svdquant_tile_k), + 1, + ) + ) + else: + tma_atom_svdquant_a, tma_tensor_svdquant_a = tma_atom_a, tma_tensor_a + tma_atom_svdquant_b, tma_tensor_svdquant_b = tma_atom_b, tma_tensor_b tile_sched_params, grid = self._compute_grid( c, @@ -546,7 +581,7 @@ def __call__( ) @cute.struct - class DenseSharedStorage: + class SharedStorage: mainloop_pipeline_array_ptr: cute.struct.MemRange[ cutlass.Int64, self.ab_stage * 2 ] @@ -581,53 +616,7 @@ class DenseSharedStorage: self.buffer_align_bytes, ] - @cute.struct - class SvdquantSharedStorage: - dense: DenseSharedStorage - sSvdquantA: cute.struct.Align[ - cute.struct.MemRange[ - BFloat16, - cute.cosize(self.svdquant_a_smem_layout), - ], - 128, - ] - sSvdquantB: cute.struct.Align[ - cute.struct.MemRange[ - BFloat16, - cute.cosize(self.svdquant_b_smem_layout), - ], - 128, - ] - - @property - def mainloop_pipeline_array_ptr(self): - return self.dense.mainloop_pipeline_array_ptr - - @property - def sA(self): - return self.dense.sA - - @property - def sB(self): - return self.dense.sB - - @property - def sSFA(self): - return self.dense.sSFA - - @property - def sSFB(self): - return self.dense.sSFB - - @property - def sC(self): - return self.dense.sC - - shared_storage: Any = DenseSharedStorage - if cutlass.const_expr(self.svdquant_enabled): - shared_storage = SvdquantSharedStorage - - self.shared_storage = shared_storage + self.shared_storage = SharedStorage self.kernel( tma_atom_a, @@ -656,6 +645,10 @@ def sC(self): self.svdquant_tiled_mma if self.svdquant_enabled else None, self.svdquant_a_smem_layout if self.svdquant_enabled else None, self.svdquant_b_smem_layout if self.svdquant_enabled else None, + tma_atom_svdquant_a, + tma_tensor_svdquant_a, + tma_atom_svdquant_b, + tma_tensor_svdquant_b, tile_sched_params, epilogue_op, alpha, @@ -910,30 +903,6 @@ def _make_scale_tiled_copy( cute.make_layout((copy_bits // dtype.width,)), ) - @cute.jit - def _make_svdquant_gmem_tiled_copy( - self, - tile_rows: cutlass.Constexpr[int], - ) -> cute.TiledCopy: - copy_bits = 128 - copy_elems = copy_bits // BFloat16.width - threads_per_row = 16 // copy_elems - copy_rows = min( - tile_rows, - self.num_mma_warps * self.num_threads_per_warp // threads_per_row, - ) - copy_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), - BFloat16, - num_bits_per_copy=copy_bits, - ) - thread_layout = cute.make_ordered_layout( - (copy_rows, threads_per_row), - order=(1, 0), - ) - value_layout = cute.make_layout((1, copy_elems)) - return cute.make_tiled_copy_tv(copy_atom, thread_layout, value_layout) - @cute.jit def _predicate_tiled_copy_rows( self, @@ -956,25 +925,6 @@ def _predicate_tiled_copy_rows( tPred[rest_v, 0, rest_k] = tCc[(0, rest_v), 0, rest_k][0] < row_limit return tPred - @cute.jit - def _svdquant_tiled_copy_2d( - self, - tiled_copy: cute.TiledCopy, - gmem_tensor: cute.Tensor, - smem_tensor: cute.Tensor, - coord_tensor: cute.Tensor, - tidx: Int32, - copy_threads: cutlass.Constexpr[int], - row_limit: Int32, - ) -> None: - thr_copy = tiled_copy.get_slice(tidx) - tG = thr_copy.partition_S(gmem_tensor) - tS = thr_copy.partition_D(smem_tensor) - tC = thr_copy.partition_S(coord_tensor) - tP = self._predicate_tiled_copy_rows(tC, row_limit) - if tidx < copy_threads: - cute.copy(tiled_copy, tG, tS, pred=tP) - @cute.jit def _cpasync_copy_2d( self, @@ -1052,6 +1002,10 @@ def kernel( svdquant_tiled_mma: Optional[cute.TiledMma], svdquant_a_smem_layout: Optional[cute.ComposedLayout], svdquant_b_smem_layout: Optional[cute.ComposedLayout], + tma_atom_svdquant_a: cute.CopyAtom, + mSvdquantA: cute.Tensor, + tma_atom_svdquant_b: cute.CopyAtom, + mSvdquantB: cute.Tensor, tile_sched_params: utils.PersistentTileSchedulerParams, epilogue_op: cutlass.Constexpr, alpha: cute.Tensor, @@ -1082,6 +1036,9 @@ def kernel( cpasync.prefetch_descriptor(tma_atom_sfb) if cutlass.const_expr(not self.use_m1_non_tma_c): cpasync.prefetch_descriptor(tma_atom_c) + if cutlass.const_expr(svdquant_d is not None): + cpasync.prefetch_descriptor(tma_atom_svdquant_a) + cpasync.prefetch_descriptor(tma_atom_svdquant_b) cta_rank_in_cluster = cute.arch.make_warp_uniform( cute.arch.block_idx_in_cluster() @@ -1092,19 +1049,22 @@ def kernel( b_smem_layout = cute.slice_(b_smem_layout_staged, (None, None, 0)) sfa_smem_layout = cute.slice_(sfa_smem_layout_staged, (None, None, 0)) sfb_smem_layout = cute.slice_(sfb_smem_layout_staged, (None, None, 0)) - if cutlass.const_expr(self.use_m1_non_tma_sfa): - tma_copy_bytes = cute.size_in_bytes( - self.b_dtype, b_smem_layout - ) + cute.size_in_bytes(self.sf_dtype, sfb_smem_layout) - if cutlass.const_expr(not self.use_m1_non_tma_a): - tma_copy_bytes += cute.size_in_bytes(self.a_dtype, a_smem_layout) - else: - tma_copy_bytes = ( - cute.size_in_bytes(self.a_dtype, a_smem_layout) - + cute.size_in_bytes(self.b_dtype, b_smem_layout) - + cute.size_in_bytes(self.sf_dtype, sfa_smem_layout) - + cute.size_in_bytes(self.sf_dtype, sfb_smem_layout) + ab_copy_bytes = cute.size_in_bytes(self.b_dtype, b_smem_layout) + if cutlass.const_expr(not self.use_m1_non_tma_a): + ab_copy_bytes += cute.size_in_bytes(self.a_dtype, a_smem_layout) + scale_copy_bytes = cute.size_in_bytes(self.sf_dtype, sfb_smem_layout) + if cutlass.const_expr(not self.use_m1_non_tma_sfa): + scale_copy_bytes += cute.size_in_bytes(self.sf_dtype, sfa_smem_layout) + tma_copy_bytes = ab_copy_bytes + scale_copy_bytes + if cutlass.const_expr(svdquant_d is not None): + svdquant_copy_bytes = cute.size_in_bytes( + self.svdquant_dtype, + cute.slice_(svdquant_a_smem_layout, (None, None, 0)), + ) + cute.size_in_bytes( + self.svdquant_dtype, + cute.slice_(svdquant_b_smem_layout, (None, None, 0)), ) + assert svdquant_copy_bytes == ab_copy_bytes # Allocate shared memory smem = cutlass.utils.SmemAllocator() @@ -1161,11 +1121,34 @@ def kernel( sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged) sSFB = storage.sSFB.get_tensor(sfb_smem_layout_staged) if cutlass.const_expr(svdquant_d is not None): - sSvdquantA = storage.sSvdquantA.get_tensor( - svdquant_a_smem_layout.outer, swizzle=svdquant_a_smem_layout.inner + # The residual mainloop has released these staged A/B regions + # before the correction executes. Reinterpret their full byte + # capacity as BF16 instead of reserving separate correction SMEM. + svdquant_a_storage = ( + storage.sB if cutlass.const_expr(self.swap_ab) else storage.sA ) - sSvdquantB = storage.sSvdquantB.get_tensor( - svdquant_b_smem_layout.outer, swizzle=svdquant_b_smem_layout.inner + svdquant_b_storage = ( + storage.sA if cutlass.const_expr(self.swap_ab) else storage.sB + ) + sSvdquantA = svdquant_a_storage.get_tensor( + svdquant_a_smem_layout.outer, + swizzle=svdquant_a_smem_layout.inner, + dtype=self.svdquant_dtype, + ) + sSvdquantB = svdquant_b_storage.get_tensor( + svdquant_b_smem_layout.outer, + swizzle=svdquant_b_smem_layout.inner, + dtype=self.svdquant_dtype, + ) + gSvdquantA = cute.local_tile( + mSvdquantA, + (self.mma_tile_shape_mnk[0], self.svdquant_tile_k), + (None, None, None), + ) + gSvdquantB = cute.local_tile( + mSvdquantB, + (self.mma_tile_shape_mnk[1], self.svdquant_tile_k), + (None, None, None), ) # Local_tile partition global tensors @@ -1244,6 +1227,22 @@ def kernel( cute.group_modes(gB_nkl, 0, 2), ) + if cutlass.const_expr(svdquant_d is not None): + tSvdquantAs, tSvdquantAg = cpasync.tma_partition( + tma_atom_svdquant_a, + a_cta_crd, + a_cta_layout, + cute.group_modes(sSvdquantA, 0, 2), + cute.group_modes(gSvdquantA, 0, 2), + ) + tSvdquantBs, tSvdquantBg = cpasync.tma_partition( + tma_atom_svdquant_b, + b_cta_crd, + b_cta_layout, + cute.group_modes(sSvdquantB, 0, 2), + cute.group_modes(gSvdquantB, 0, 2), + ) + # TMA partitions for SFA if cutlass.const_expr(self.load_path == "tma" and not self.use_m1_non_tma_sfa): tAsSFA, tAgSFA = cpasync.tma_partition( @@ -1485,10 +1484,12 @@ def kernel( tCsSvdquantB[None, None, None, 0] ) svdquant_copy_atom_a = cute.make_copy_atom( - cute.nvgpu.warp.LdMatrix8x8x16bOp(False, 4), BFloat16 + cute.nvgpu.warp.LdMatrix8x8x16bOp(False, 4), + self.svdquant_dtype, ) svdquant_copy_atom_b = cute.make_copy_atom( - cute.nvgpu.warp.LdMatrix8x8x16bOp(False, 4), BFloat16 + cute.nvgpu.warp.LdMatrix8x8x16bOp(False, 4), + self.svdquant_dtype, ) svdquant_smem_copy_a = cute.make_tiled_copy_A( svdquant_copy_atom_a, svdquant_tiled_mma @@ -1502,13 +1503,6 @@ def kernel( tCsSvdquantB_copy = svdquant_thr_copy_b.partition_S(sSvdquantB) tCrSvdquantA_copy = svdquant_thr_copy_a.retile(tCrSvdquantA) tCrSvdquantB_copy = svdquant_thr_copy_b.retile(tCrSvdquantB) - svdquant_gmem_tiled_copy_a = self._make_svdquant_gmem_tiled_copy( - self.mma_tile_shape_mnk[0] - ) - svdquant_gmem_tiled_copy_b = self._make_svdquant_gmem_tiled_copy( - self.mma_tile_shape_mnk[1] - ) - while work_tile.is_valid_tile: tile_coord_mnl = work_tile.tile_idx gC_mnl_slice = gC_mnl[(None, None, *tile_coord_mnl)] @@ -1789,10 +1783,12 @@ def kernel( accumulators[None, _mt, _nt], ) - # SVDQuant fusion: stage rank-16 BF16 chunks cooperatively, then - # accumulate them with warp MMA directly into the FP32 NVFP4 - # accumulator fragment. swap_ab also swaps the correction's A/B - # operands, so both MMA paths retain identical C-fragment ownership. + # SVDQuant fusion: the load warp stages one + # mainloop-byte-equivalent BF16 rank tile at a time into the + # released A/B mainloop storage. The MMA warps consume each + # published tile directly into the FP32 NVFP4 accumulator. + # swap_ab also swaps the correction's A/B operands, so both MMA + # paths retain identical C-fragment ownership. # Match the SM100 CUTLASS contract: svdquant_l1 is supplied as # L1 / alpha. The common epilogue's alpha multiply therefore # restores the unscaled correction: @@ -1800,101 +1796,22 @@ def kernel( # = alpha * residual + D @ L1.T. if cutlass.const_expr(svdquant_d is not None): rank = cute.size(svdquant_d, mode=[1]) - svdquant_a_rows = self.mma_tile_shape_mnk[0] - svdquant_b_rows = self.mma_tile_shape_mnk[1] - svdquant_threads_per_row = 2 - svdquant_copy_row_capacity = ( - self.num_mma_warps - * self.num_threads_per_warp - // svdquant_threads_per_row - ) - svdquant_a_copy_threads = ( - min(svdquant_a_rows, svdquant_copy_row_capacity) - * svdquant_threads_per_row - ) - svdquant_b_copy_threads = ( - min(svdquant_b_rows, svdquant_copy_row_capacity) - * svdquant_threads_per_row - ) - tile_m, tile_n, _ = tile_coord_mnl - tile_rows_m, tile_rows_n, _ = self.tile_shape_mnk - svdquant_a_plan = ( - svdquant_d, - tile_m, - tile_rows_m, - ) - svdquant_b_plan = ( - svdquant_l1, - tile_n, - tile_rows_n, - ) - if cutlass.const_expr(self.swap_ab): - svdquant_a_plan, svdquant_b_plan = ( - svdquant_b_plan, - svdquant_a_plan, - ) - ( - svdquant_a_source, - svdquant_a_tile_coord, - svdquant_a_tile_rows, - ) = svdquant_a_plan - ( - svdquant_b_source, - svdquant_b_tile_coord, - svdquant_b_tile_rows, - ) = svdquant_b_plan - svdquant_a_tiler = (svdquant_a_tile_rows, 16) - svdquant_b_tiler = (svdquant_b_tile_rows, 16) - gSvdquantA = cute.local_tile( - svdquant_a_source, - svdquant_a_tiler, - (svdquant_a_tile_coord, None), - ) - gSvdquantB = cute.local_tile( - svdquant_b_source, - svdquant_b_tiler, - (svdquant_b_tile_coord, None), - ) - cSvdquantA = cute.local_tile( - cute.make_identity_tensor(cute.shape(svdquant_a_source)), - svdquant_a_tiler, - (svdquant_a_tile_coord, None), - ) - cSvdquantB = cute.local_tile( - cute.make_identity_tensor(cute.shape(svdquant_b_source)), - svdquant_b_tiler, - (svdquant_b_tile_coord, None), - ) - for rank_tile in cutlass.range_constexpr(rank // 16): - self._svdquant_tiled_copy_2d( - svdquant_gmem_tiled_copy_a, - gSvdquantA[(None, None, rank_tile)], - sSvdquantA[(None, None, 0)], - cSvdquantA[(None, None, rank_tile)], - tidx, - svdquant_a_copy_threads, - Int32(cute.size(svdquant_a_source, mode=[0])), - ) - self._svdquant_tiled_copy_2d( - svdquant_gmem_tiled_copy_b, - gSvdquantB[(None, None, rank_tile)], - sSvdquantB[(None, None, 0)], - cSvdquantB[(None, None, rank_tile)], - tidx, - svdquant_b_copy_threads, - Int32(cute.size(svdquant_b_source, mode=[0])), - ) - - # All MMA warps must see the cooperatively staged rank tile. - self.mma_sync_barrier.arrive_and_wait() + for _rank_tile in cutlass.range_constexpr( + rank // self.svdquant_tile_k + ): + mainloop_pipeline.consumer_wait(mainloop_consumer_state) cute.copy( svdquant_smem_copy_a, - tCsSvdquantA_copy[None, None, None, 0], + tCsSvdquantA_copy[ + None, None, None, mainloop_consumer_state.index + ], tCrSvdquantA_copy, ) cute.copy( svdquant_smem_copy_b, - tCsSvdquantB_copy[None, None, None, 0], + tCsSvdquantB_copy[ + None, None, None, mainloop_consumer_state.index + ], tCrSvdquantB_copy, ) cute.gemm( @@ -1904,12 +1821,11 @@ def kernel( tCrSvdquantB, accumulators, ) - # Keep early warps from overwriting SMEM still read by peers. - self.mma_sync_barrier.arrive_and_wait() + mainloop_pipeline.consumer_release(mainloop_consumer_state) + mainloop_consumer_state.advance() # Bias remains an elementwise output-epilogue contribution. # The low-rank matrix product above is exclusively BF16 MMA. - if cutlass.const_expr(self.swap_ab): acc_mn = _reshape_acc_to_mn(accumulators, transpose=True) c_identity = cute.make_identity_tensor( @@ -2287,7 +2203,6 @@ def kernel( elif warp_idx == self.tma_load_warp_id: cute.arch.setmaxregister_decrease(self.load_register_requirement) - while work_tile.is_valid_tile: tile_coord_mnl = work_tile.tile_idx if cutlass.const_expr( @@ -2323,6 +2238,9 @@ def kernel( k_tile_global = k_tile_start + mainloop_producer_state.count if cutlass.const_expr(self.load_path == "tma"): + sf_producer_barrier = mainloop_pipeline.producer_get_barrier( + mainloop_producer_state + ) tBgB_k = tBgB_nkl[(None, k_tile_global)] tBsB_pipe = tBsB[(None, mainloop_producer_state.index)] if cutlass.const_expr(not self.use_m1_non_tma_a): @@ -2552,9 +2470,7 @@ def kernel( tma_atom_sfa, tAgSFA_k, tAsSFA_pipe, - tma_bar_ptr=mainloop_pipeline.producer_get_barrier( - mainloop_producer_state - ), + tma_bar_ptr=sf_producer_barrier, ) if cutlass.const_expr(self.load_path == "tma"): cute.copy( @@ -2569,9 +2485,7 @@ def kernel( tma_atom_sfb, tBgSFB_k, tBsSFB_pipe, - tma_bar_ptr=mainloop_pipeline.producer_get_barrier( - mainloop_producer_state - ), + tma_bar_ptr=sf_producer_barrier, ) if cutlass.const_expr(self.load_path == "cpasync"): cute.arch.cp_async_commit_group() @@ -2579,6 +2493,58 @@ def kernel( mainloop_pipeline.producer_commit(mainloop_producer_state) mainloop_producer_state.advance() + if cutlass.const_expr(svdquant_d is not None): + tile_m, tile_n, _ = tile_coord_mnl + if cutlass.const_expr(self.swap_ab): + tile_m, tile_n = tile_n, tile_m + tSvdquantAg_tile = tSvdquantAg[ + (None, tile_m, None, tile_coord_mnl[2]) + ] + tSvdquantBg_tile = tSvdquantBg[ + (None, tile_n, None, tile_coord_mnl[2]) + ] + + rank = cute.size(svdquant_d, mode=[1]) + for rank_tile in cutlass.range_constexpr( + rank // self.svdquant_tile_k + ): + # Correction TMA reuses the residual pipeline's stage + # ring but transfers A/B only. Arm its raw transaction + # barrier with the exact aliased A/B byte count rather + # than the residual A+B+SFA+SFB count. + mainloop_pipeline.sync_object_empty.wait( + mainloop_producer_state.index, + mainloop_producer_state.phase, + ) + mainloop_pipeline.sync_object_full.arrive_and_expect_tx( + mainloop_producer_state.index, + ab_copy_bytes, + ) + tSvdquantAs_pipe = tSvdquantAs[ + (None, mainloop_producer_state.index) + ] + tSvdquantBs_pipe = tSvdquantBs[ + (None, mainloop_producer_state.index) + ] + cute.copy( + tma_atom_svdquant_a, + tSvdquantAg_tile[(None, rank_tile)], + tSvdquantAs_pipe, + tma_bar_ptr=mainloop_pipeline.producer_get_barrier( + mainloop_producer_state + ), + ) + cute.copy( + tma_atom_svdquant_b, + tSvdquantBg_tile[(None, rank_tile)], + tSvdquantBs_pipe, + tma_bar_ptr=mainloop_pipeline.producer_get_barrier( + mainloop_producer_state + ), + ) + mainloop_pipeline.producer_commit(mainloop_producer_state) + mainloop_producer_state.advance() + if cutlass.const_expr(self.single_work_tile_per_cta): work_tile = WorkTileInfo( work_tile.tile_idx, @@ -2606,7 +2572,6 @@ def _compute_stages( c_dtype, smem_capacity: int, occupancy: int, - svdquant_bytes: int, ) -> tuple: epi_stage_max = (tile_shape_mnk[1] // epi_tile[1]) * ( tile_shape_mnk[0] // epi_tile[0] @@ -2631,7 +2596,6 @@ def _compute_stages( (smem_capacity - occupancy * 1024) // occupancy - mbar_helpers_bytes - epi_bytes - - svdquant_bytes ) // (ab_bytes_per_stage + sf_bytes_per_stage) ab_stage = max(1, min(raw_ab_stage, 4)) if tile_shape_mnk[0] in (16, 64) and tile_shape_mnk[1] == 128: @@ -2806,6 +2770,7 @@ def can_implement( load_path: str = "tma", swap_ab: bool = False, svdquant_rank: Optional[int] = None, + mainloop_tile_k: Optional[int] = None, ) -> bool: # The current target only supports cluster (1,1) if cluster_shape_mn != (1, 1): @@ -2827,8 +2792,21 @@ def can_implement( return False if load_path == "cpasync" and (sf_vec_size != 16 or l != 1): return False - if svdquant_rank is not None and svdquant_rank % 16 != 0: - return False + if svdquant_rank is not None: + if is_mxfp8: + return False + if load_path != "tma": + return False + if mainloop_tile_k is None: + return False + rank_elements_per_stage = mainloop_tile_k // ( + cutlass.BFloat16.width // ab_dtype.width + ) + if ( + svdquant_rank < rank_elements_per_stage + or svdquant_rank % rank_elements_per_stage != 0 + ): + return False # SF smem still allocates full 128-element blocks even when the live # MMA tile uses only 16 or 32 rows or columns. if is_mxfp8: @@ -2938,6 +2916,16 @@ def wrapper( order=(2, 1, 4, 0, 3, 5), ), ) + if cutlass.const_expr(svdquant_d is not None): + rank = cute.size(svdquant_d, mode=[1]) + svdquant_d = cute.make_tensor( + svdquant_d.iterator, + layout=cute.make_ordered_layout((m, rank, l), order=(1, 0, 2)), + ) + svdquant_l1 = cute.make_tensor( + svdquant_l1.iterator, + layout=cute.make_ordered_layout((n, rank, l), order=(1, 0, 2)), + ) self( a_tensor, diff --git a/tests/gemm/test_nvfp4_svdquant_gemm.py b/tests/gemm/test_nvfp4_svdquant_gemm.py index 96845086828..31654e9f227 100644 --- a/tests/gemm/test_nvfp4_svdquant_gemm.py +++ b/tests/gemm/test_nvfp4_svdquant_gemm.py @@ -67,10 +67,16 @@ def test_sm120_svdquant_can_implement_rejects_ragged_rank(): "n", ) assert Sm120B12xBlockScaledDenseGemmKernel.can_implement( - *common_args, svdquant_rank=32 + *common_args, svdquant_rank=32, mainloop_tile_k=128 ) assert not Sm120B12xBlockScaledDenseGemmKernel.can_implement( - *common_args, svdquant_rank=18 + *common_args, svdquant_rank=18, mainloop_tile_k=128 + ) + assert not Sm120B12xBlockScaledDenseGemmKernel.can_implement( + *common_args, svdquant_rank=32, mainloop_tile_k=256 + ) + assert Sm120B12xBlockScaledDenseGemmKernel.can_implement( + *common_args, svdquant_rank=64, mainloop_tile_k=256 ) @@ -536,6 +542,37 @@ def test_mm_nvfp4_svdquant_sm120_fused(use_bias): assert _sqnr_db(fp32_ref, out.float()) > 35.0 +@pytest.mark.parametrize("tile_k,rank", [(64, 32), (128, 32), (256, 64)]) +def test_mm_nvfp4_svdquant_sm120_large_m_tile(tile_k, rank): + _skip_unless_sm120() + from flashinfer.gemm.gemm_svdquant import _mm_nvfp4_svdquant_sm120_fused + + torch.manual_seed(tile_k) + p = _make_gemm_problem( + 257, + 128, + 512, + rank=rank, + quant_backend="cute-dsl", + residual_backend="b12x", + ) + out = torch.empty(257, 128, dtype=torch.bfloat16, device="cuda") + _mm_nvfp4_svdquant_sm120_fused( + p["xq"], + p["wq"], + p["x_sf_flat"], + p["w_sf_flat"], + p["alpha"], + p["d"], + p["l1_scaled"], + p["bias"], + out, + device_support_pdl(torch.device("cuda")), + tactic=((256, 64), tile_k, False, False), + ) + assert _sqnr_db(p["ref_bias"], out.float()) > 35.0 + + @pytest.mark.parametrize("backend", ["cute-dsl", "auto"]) def test_mm_nvfp4_svdquant_sm120_autotuned_replay(backend): _skip_unless_sm120() @@ -764,16 +801,15 @@ def fail_torch_mm(*args, **kwargs): @pytest.mark.parametrize("use_bias", [False, True]) def test_svdquant_linear_sm120_fused(use_bias, monkeypatch): _skip_unless_sm120() - import flashinfer.gemm.gemm_base as gemm_base - - mm_bf16_backends = [] - original_mm_bf16 = gemm_base.mm_bf16 + torch_mm_calls = 0 + original_torch_mm = torch.mm - def recording_mm_bf16(*args, **kwargs): - mm_bf16_backends.append(kwargs.get("backend")) - return original_mm_bf16(*args, **kwargs) + def recording_torch_mm(*args, **kwargs): + nonlocal torch_mm_calls + torch_mm_calls += 1 + return original_torch_mm(*args, **kwargs) - monkeypatch.setattr(gemm_base, "mm_bf16", recording_mm_bf16) + monkeypatch.setattr(torch, "mm", recording_torch_mm) torch.manual_seed(0) m, n, k, rank = 129, 256, 256, 32 x = torch.randn(m, k, dtype=torch.bfloat16, device="cuda") / (k**0.25) @@ -812,6 +848,7 @@ def recording_mm_bf16(*args, **kwargs): global_sf, bias=bias, ) + assert torch_mm_calls == 1 xq, x_sf = nvfp4_quantize( smoothed, @@ -835,10 +872,6 @@ def recording_mm_bf16(*args, **kwargs): assert out.shape == (m, n) and out.dtype == torch.bfloat16 assert _sqnr_db(ref, out.float()) > 35.0 - if gemm_base.CUDNN_AVAILABLE: - assert mm_bf16_backends == ["cudnn"] - else: - assert mm_bf16_backends == [] @pytest.mark.parametrize("rank", [32, 128]) From cf9558e3b46bce8fa748932018c9e4309c1140e5 Mon Sep 17 00:00:00 2001 From: Anthony Chang <27950904+rosenrodt@users.noreply.github.com> Date: Sat, 8 Aug 2026 10:39:57 +0800 Subject: [PATCH 4/7] Add compile-time IKET diagnostics for SM120 SVDQuant Expose balanced warp-phase ranges for profiling while keeping diagnostic operations out of ordinary production specializations. Changes - Instrument main load, LoRA load, main MMA, LoRA MMA, and epilogue phases - Standardize the tile_k admission term - Guard all ranges with a default-off constexpr and separate _iket0/_iket1 cache identities Validation - Preserve the original five-phase trace, default-off correctness, and same-node performance A/B changes - Run applicable repository commit hooks during history reconstruction Result - IKET profiling remains explicitly available for diagnosis - Default SM120 kernels contain no active tracing operations and show no measurable regression --- flashinfer/gemm/gemm_svdquant.py | 9 ++- .../dense_blockscaled_gemm_sm120_b12x.py | 59 ++++++++++++++++++- tests/gemm/test_nvfp4_svdquant_gemm.py | 13 ++-- 3 files changed, 72 insertions(+), 9 deletions(-) diff --git a/flashinfer/gemm/gemm_svdquant.py b/flashinfer/gemm/gemm_svdquant.py index d6092487517..c7975d1abaa 100644 --- a/flashinfer/gemm/gemm_svdquant.py +++ b/flashinfer/gemm/gemm_svdquant.py @@ -86,6 +86,7 @@ def _compile_sm120_nvfp4_svdquant( sf_n: int, sf_k: int, enable_pdl: bool, + enable_iket: bool = False, ): """Compile one fused SM120 NVFP4 + rank-r BF16 epilogue specialization.""" device_index = device.index @@ -105,6 +106,7 @@ def _compile_sm120_nvfp4_svdquant( use_prefetch, max_active_clusters, enable_pdl, + enable_iket, ) if cache_key in _SM120_SVDQUANT_KERNEL_CACHE: return _SM120_SVDQUANT_KERNEL_CACHE[cache_key] @@ -128,6 +130,7 @@ def _compile_sm120_nvfp4_svdquant( use_prefetch=use_prefetch, enable_pdl=enable_pdl, swap_ab=swap_ab, + enable_iket=enable_iket, ) def compile_kernel(): @@ -201,7 +204,7 @@ def compile_kernel(): kernel_name = ( f"r{rank}_bias{int(with_bias)}_t{mma_tiler_mn[0]}x{mma_tiler_mn[1]}" f"x{tile_k}_swap{int(swap_ab)}_pf{int(use_prefetch)}_mac{max_active_clusters}" - f"_pdl{int(enable_pdl)}" + f"_pdl{int(enable_pdl)}_iket{int(enable_iket)}" ) compiled = build_and_load_cute_dsl_kernel( "mm_nvfp4_svdquant_sm120", @@ -225,6 +228,7 @@ def _mm_nvfp4_svdquant_sm120_fused( out: torch.Tensor, enable_pdl: bool, tactic=None, + enable_iket: bool = False, ) -> torch.Tensor: from .kernels.dense_blockscaled_gemm_sm120_b12x import ( _select_default_dense_gemm_plan, @@ -254,6 +258,7 @@ def _mm_nvfp4_svdquant_sm120_fused( sf_n=sf_n, sf_k=sf_k, enable_pdl=enable_pdl, + enable_iket=enable_iket, ) args = [ a, @@ -357,7 +362,7 @@ def _add( "n", swap_ab=swap_ab, svdquant_rank=d.shape[1], - mainloop_tile_k=tile_k, + tile_k=tile_k, ): return for use_prefetch in prefetch_candidates: diff --git a/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py b/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py index 3f77b5243aa..54f245345bb 100644 --- a/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py +++ b/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py @@ -30,6 +30,7 @@ # and adapted for the current Blackwell GeForce target. from dataclasses import dataclass +import importlib from typing import Literal, Optional, Tuple import cuda.bindings.driver as cuda @@ -49,6 +50,32 @@ from cutlass.utils.static_persistent_tile_scheduler import WorkTileInfo from cutlass._mlir.dialects import llvm + +class _IketShim: + """No-op IKET markers for CuTe DSL builds without IKET support.""" + + @staticmethod + def range_push(_name): + return None + + @staticmethod + def range_pop(): + return None + + +def _load_iket(): + # Keep the optional experimental import out of CuTe's AST import replay: + # normal CUDA 13.0 builds raise NotImplementedError while importing it. + for module_name in ("cutlass.cute.experimental", "cutlass.cute"): + try: + return importlib.import_module(module_name).iket + except Exception: # pragma: no cover - availability is environment-specific + pass + return _IketShim() + + +iket = _load_iket() + from flashinfer.cute_dsl.utils import ( sm120_make_smem_layout_sfa, sm120_make_smem_layout_sfb, @@ -229,6 +256,7 @@ def __init__( use_m1_non_tma_sfa: bool = False, load_path: Literal["tma", "cpasync"] = "tma", swap_ab: bool = False, + enable_iket: bool = False, ): self.acc_dtype = cutlass.Float32 self.sf_vec_size = sf_vec_size @@ -258,6 +286,9 @@ def __init__( self.use_m1_non_tma_sfa = use_m1_non_tma_sfa self.load_path = load_path self.swap_ab = swap_ab + # IKET is a compile-time diagnostic. Every marker is guarded with + # const_expr so the default specialization contains no tracing ops. + self.enable_iket = enable_iket mma_atom_mn = (self.mma_tile_shape_mnk[0], self.mma_tile_shape_mnk[1]) if mma_atom_mn in ((16, 64), (16, 128)): self.atom_shape = (1, 2, 1) @@ -1589,6 +1620,8 @@ def kernel( accumulators.fill(0.0) # Pipelined MAINLOOP + if cutlass.const_expr(self.enable_iket): + iket.range_push("mma_main") mainloop_consumer_state.reset_count() peek_ab_full_status = cutlass.Boolean(1) @@ -1782,6 +1815,8 @@ def kernel( tCrB[None, _nt, k_block_idx], accumulators[None, _mt, _nt], ) + if cutlass.const_expr(self.enable_iket): + iket.range_pop() # SVDQuant fusion: the load warp stages one # mainloop-byte-equivalent BF16 rank tile at a time into the @@ -1795,6 +1830,8 @@ def kernel( # alpha * (residual + D @ (L1 / alpha).T) # = alpha * residual + D @ L1.T. if cutlass.const_expr(svdquant_d is not None): + if cutlass.const_expr(self.enable_iket): + iket.range_push("mma_lora") rank = cute.size(svdquant_d, mode=[1]) for _rank_tile in cutlass.range_constexpr( rank // self.svdquant_tile_k @@ -1823,9 +1860,13 @@ def kernel( ) mainloop_pipeline.consumer_release(mainloop_consumer_state) mainloop_consumer_state.advance() + if cutlass.const_expr(self.enable_iket): + iket.range_pop() # Bias remains an elementwise output-epilogue contribution. # The low-rank matrix product above is exclusively BF16 MMA. + if cutlass.const_expr(self.enable_iket): + iket.range_push("epilogue") if cutlass.const_expr(self.swap_ab): acc_mn = _reshape_acc_to_mn(accumulators, transpose=True) c_identity = cute.make_identity_tensor( @@ -1861,6 +1902,8 @@ def kernel( tile_coord_mnl[2], ) ] = epilogue_op(acc_value.to(self.c_dtype)) + if cutlass.const_expr(self.enable_iket): + iket.range_pop() if cutlass.const_expr(self.single_work_tile_per_cta): work_tile = WorkTileInfo( work_tile.tile_idx, @@ -2184,6 +2227,8 @@ def kernel( # the next work tile's mainloop. tma_store_pipeline.producer_commit() + if cutlass.const_expr(self.enable_iket): + iket.range_pop() # Advance to the next work tile if cutlass.const_expr(self.single_work_tile_per_cta): work_tile = WorkTileInfo( @@ -2205,6 +2250,8 @@ def kernel( cute.arch.setmaxregister_decrease(self.load_register_requirement) while work_tile.is_valid_tile: tile_coord_mnl = work_tile.tile_idx + if cutlass.const_expr(self.enable_iket): + iket.range_push("load_warp_main") if cutlass.const_expr( self.load_path == "tma" and not self.use_m1_non_tma_a ): @@ -2492,8 +2539,12 @@ def kernel( cute.arch.cp_async_wait_group(0) mainloop_pipeline.producer_commit(mainloop_producer_state) mainloop_producer_state.advance() + if cutlass.const_expr(self.enable_iket): + iket.range_pop() if cutlass.const_expr(svdquant_d is not None): + if cutlass.const_expr(self.enable_iket): + iket.range_push("load_warp_lora") tile_m, tile_n, _ = tile_coord_mnl if cutlass.const_expr(self.swap_ab): tile_m, tile_n = tile_n, tile_m @@ -2544,6 +2595,8 @@ def kernel( ) mainloop_pipeline.producer_commit(mainloop_producer_state) mainloop_producer_state.advance() + if cutlass.const_expr(self.enable_iket): + iket.range_pop() if cutlass.const_expr(self.single_work_tile_per_cta): work_tile = WorkTileInfo( @@ -2770,7 +2823,7 @@ def can_implement( load_path: str = "tma", swap_ab: bool = False, svdquant_rank: Optional[int] = None, - mainloop_tile_k: Optional[int] = None, + tile_k: Optional[int] = None, ) -> bool: # The current target only supports cluster (1,1) if cluster_shape_mn != (1, 1): @@ -2797,9 +2850,9 @@ def can_implement( return False if load_path != "tma": return False - if mainloop_tile_k is None: + if tile_k is None: return False - rank_elements_per_stage = mainloop_tile_k // ( + rank_elements_per_stage = tile_k // ( cutlass.BFloat16.width // ab_dtype.width ) if ( diff --git a/tests/gemm/test_nvfp4_svdquant_gemm.py b/tests/gemm/test_nvfp4_svdquant_gemm.py index 31654e9f227..1a3fba15b3a 100644 --- a/tests/gemm/test_nvfp4_svdquant_gemm.py +++ b/tests/gemm/test_nvfp4_svdquant_gemm.py @@ -52,6 +52,11 @@ def test_sm120_svdquant_can_implement_rejects_ragged_rank(): Sm120B12xBlockScaledDenseGemmKernel, ) + assert not Sm120B12xBlockScaledDenseGemmKernel(16, (64, 64), (1, 1)).enable_iket + assert Sm120B12xBlockScaledDenseGemmKernel( + 16, (64, 64), (1, 1), enable_iket=True + ).enable_iket + common_args = ( cutlass.Float4E2M1FN, cutlass.Float8E4M3FN, @@ -67,16 +72,16 @@ def test_sm120_svdquant_can_implement_rejects_ragged_rank(): "n", ) assert Sm120B12xBlockScaledDenseGemmKernel.can_implement( - *common_args, svdquant_rank=32, mainloop_tile_k=128 + *common_args, svdquant_rank=32, tile_k=128 ) assert not Sm120B12xBlockScaledDenseGemmKernel.can_implement( - *common_args, svdquant_rank=18, mainloop_tile_k=128 + *common_args, svdquant_rank=18, tile_k=128 ) assert not Sm120B12xBlockScaledDenseGemmKernel.can_implement( - *common_args, svdquant_rank=32, mainloop_tile_k=256 + *common_args, svdquant_rank=32, tile_k=256 ) assert Sm120B12xBlockScaledDenseGemmKernel.can_implement( - *common_args, svdquant_rank=64, mainloop_tile_k=256 + *common_args, svdquant_rank=64, tile_k=256 ) From aa923d4c758ac55677cff3edab844a048b55cae6 Mon Sep 17 00:00:00 2001 From: Anthony Chang <27950904+rosenrodt@users.noreply.github.com> Date: Sun, 9 Aug 2026 20:47:05 +0800 Subject: [PATCH 5/7] Support SM120 SVDQuant on CUDA 12.9 Permit the validated CuTe DSL SVDQuant and generic B12x NVFP4 paths on CUDA 12.9 while preserving the existing SM100 CUTLASS dispatch. Incorporate focused review corrections to backend ordering, scalar-alpha documentation, kernel compatibility checks, and SVDQuant trace axes. Changes - lower SM120 SVDQuant and B12x admission to CUDA 12.9 - prefer B12x for automatic SM120 FP4 dispatch on CUDA 12.9 - document fused-first SM120 auto-dispatch and exact scalar-alpha contract - reject incompatible M=1 non-TMA A loads for SVDQuant and diagnose copy-size mismatches - model SF_A and SF_B as derived trace variables while retaining K_packed as the K specialization axis - regenerate the SVDQuant trace golden and align its documented filename Validation - pytest -q tests/gemm/test_nvfp4_svdquant_gemm.py (23 passed, 47 skipped on CUDA 12.9 RTX 5080) - focused generic B12x GPU cases (6 passed on CUDA 12.9 RTX 5080) - pytest -q tests/gemm/test_nvfp4_svdquant_gemm.py -k 'sm120 or backend_arch_support' -x (22 passed, 48 deselected on RTX 5080, CUDA 13.2, CUTLASS DSL 4.5.2) - pytest -q tests/trace/test_fi_trace_template_consistency.py (465 passed) - pytest -q tests/trace/test_fi_trace.py (32 passed) - SVDQuant trace initializer executed on RTX 5080 with expected packed and scale shapes - pre-commit run -a - git diff --check Result - SM120 SVDQuant and B12x execute successfully with CUDA 12.9 - SM100 SVDQuant remains on the unchanged architecture-gated CUTLASS path - SVDQuant trace names contain only independent constant specialization axes - reviewed compatibility assumptions now fail early with actionable diagnostics --- flashinfer/gemm/gemm_base.py | 37 +++++++++++++------ flashinfer/gemm/gemm_svdquant.py | 17 ++++++--- .../dense_blockscaled_gemm_sm120_b12x.py | 8 +++- flashinfer/trace/templates/gemm.py | 8 +++- tests/gemm/test_mm_fp4.py | 12 ++++-- tests/trace/example.py | 2 +- ...4_svdquant_N3072_K_packed1536_rank32.json} | 12 +++--- 7 files changed, 65 insertions(+), 31 deletions(-) rename tests/trace/fi_trace_out/{gemm_nvfp4_svdquant_N3072_K_packed1536_SF_B589824_rank32.json => gemm_nvfp4_svdquant_N3072_K_packed1536_rank32.json} (85%) diff --git a/flashinfer/gemm/gemm_base.py b/flashinfer/gemm/gemm_base.py index 94c14aa55d7..6934e71d9f7 100644 --- a/flashinfer/gemm/gemm_base.py +++ b/flashinfer/gemm/gemm_base.py @@ -23,9 +23,11 @@ from types import SimpleNamespace from typing import Callable, List, Literal, Optional, Tuple -from flashinfer.trtllm_low_latency_gemm import trtllm_low_latency_gemm +from packaging.version import Version import torch +from flashinfer.trtllm_low_latency_gemm import trtllm_low_latency_gemm + from ..api_logging import flashinfer_api from ..trace.templates.gemm import ( batch_deepgemm_fp8_nt_groupwise_trace, @@ -102,6 +104,8 @@ from .routergemm import get_tinygemm2_module +_MIN_B12X_CUDA_VERSION = Version("12.9") + logger = logging.getLogger(__name__) CUDNN_AVAILABLE = False @@ -6194,11 +6198,13 @@ def _b12x_gemm_fp4_requirement( use_nvfp4: bool = True, enable_pdl: bool = True, # unused ): - # b12x backend requires CUDA 13+ and 128x4 scale factor layout. - if get_cuda_version().major < 13: + cuda_version = get_cuda_version() + min_cuda_version = _MIN_B12X_CUDA_VERSION if use_nvfp4 else Version("13.0") + if cuda_version < min_cuda_version: raise ValueError( - "b12x FP4 GEMM requires CUDA 13 or later. " - f"Current CUDA version: {get_cuda_version()}." + f"b12x {'NVFP4' if use_nvfp4 else 'MXFP4'} GEMM requires " + f"CUDA {min_cuda_version} or later. " + f"Current CUDA version: {cuda_version}." ) if use_8x4_sf_layout: raise ValueError("b12x FP4 GEMM only supports 128x4 scale factor layout.") @@ -6887,24 +6893,33 @@ def _heuristic_func_mm_fp4( - On SM100 (B200) - use cudnn (faster based on benchmarks). """ - cuda_major = get_cuda_version().major + cuda_version = get_cuda_version() # Get compute capability to distinguish between SM100 (10.0) and SM103 (10.3) major, minor = get_compute_capability(a.device) is_sm107 = major == 10 and minor == 7 is_sm103 = major == 10 and minor == 3 is_sm120 = major == 12 and minor == 0 - # SM120 + CUDA 13: prefer b12x for both NVFP4 and MXFP4. SM121 (GB10) is - # intentionally excluded -- b12x is supported there as an explicit backend, - # but cutlass/cudnn are faster in most cases, so `auto` keeps using them. - if is_sm120 and cuda_major >= 13: + # SM120 prefers b12x from CUDA 12.9 for NVFP4 and CUDA 13 for MXFP4. + # SM121 (GB10) is intentionally excluded from automatic selection because + # cutlass/cudnn are faster in most cases; b12x remains explicitly selectable. + b12x_cuda_supported = ( + cuda_version >= _MIN_B12X_CUDA_VERSION + if use_nvfp4 + else cuda_version.major >= 13 + ) + if is_sm120 and b12x_cuda_supported: return [c for c in ("b12x", "cutlass", "cudnn") if c in suitable_backends] candidate_backends: Tuple[str, ...] # If cuda version is 13 or greater and cudnn version is 9.15 or greater: # On SM103 (B300), cutlass is more performant than cudnn. # On SM100 (B200), cudnn is more performant than cutlass. - if CUDNN_AVAILABLE and cuda_major >= 13 and cudnn.backend_version() >= 91500: + if ( + CUDNN_AVAILABLE + and cuda_version.major >= 13 + and cudnn.backend_version() >= 91500 + ): if is_sm103: candidate_backends = ("cutlass", "cudnn") elif is_sm107: diff --git a/flashinfer/gemm/gemm_svdquant.py b/flashinfer/gemm/gemm_svdquant.py index c7975d1abaa..e26ab89c781 100644 --- a/flashinfer/gemm/gemm_svdquant.py +++ b/flashinfer/gemm/gemm_svdquant.py @@ -18,6 +18,7 @@ from dataclasses import replace from typing import List, Literal, Optional, Tuple +from packaging.version import Version import torch from ..api_logging import flashinfer_api @@ -56,6 +57,7 @@ SVDQUANT_LORA_RANK_GRANULARITY = 32 _SM120_SVDQUANT_KERNEL_CACHE: dict[tuple, object] = {} +_MIN_SM120_SVDQUANT_CUDA_VERSION = Version("12.9") def _pad_up(x: int, y: int) -> int: @@ -611,10 +613,11 @@ def _cutlass_nvfp4_svdquant_requirement(*args, **kwargs): @supported_compute_capability([120, 121]) def _cute_dsl_nvfp4_svdquant_requirement(*args, **kwargs): - if get_cuda_version().major < 13: + cuda_version = get_cuda_version() + if cuda_version < _MIN_SM120_SVDQUANT_CUDA_VERSION: raise ValueError( - "SM120 SVDQuant CuTe DSL support requires CUDA 13 or later. " - f"Current CUDA version: {get_cuda_version()}." + "SM120 SVDQuant CuTe DSL support requires CUDA 12.9 or later. " + f"Current CUDA version: {cuda_version}." ) from ..cute_dsl import is_cute_dsl_available @@ -629,8 +632,9 @@ def _cute_dsl_nvfp4_svdquant_requirement(*args, **kwargs): def _heuristic_func_nvfp4_svdquant( suitable_backends: List[str], *args, **kwargs ) -> List[str]: - # The backend requirements are architecture-disjoint, so retaining their order - # is sufficient and keeps automatic dispatch deterministic. + # Preserve backend_checks order: on SM120/SM121, cute-dsl precedes + # cute-dsl-unfused when both are supported, which selects the fused-first + # implementation tuning configuration. return suitable_backends @@ -757,7 +761,8 @@ def mm_nvfp4_svdquant( b_sf: torch.Tensor Weight block scales, same layout as ``a_sf`` with ``n`` rows. alpha: torch.Tensor - Per-tensor residual dequantization scale, float32, device scalar (``numel >= 1``). + Per-tensor residual dequantization scale, float32 device scalar with + exactly one element (``numel == 1``). d: torch.Tensor LoRA-down output ``x_hat @ L2ᵀ``, shape ``(m, r)`` bf16, contiguous and 16-byte aligned (TMA). Compute it as ``x @ (pre_quant_scale[:, None] * L2ᵀ)`` in bf16. diff --git a/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py b/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py index 54f245345bb..6b6c6266b8e 100644 --- a/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py +++ b/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py @@ -1088,6 +1088,9 @@ def kernel( scale_copy_bytes += cute.size_in_bytes(self.sf_dtype, sfa_smem_layout) tma_copy_bytes = ab_copy_bytes + scale_copy_bytes if cutlass.const_expr(svdquant_d is not None): + assert not self.use_m1_non_tma_a, ( + "SVDQuant does not support use_m1_non_tma_a=True" + ) svdquant_copy_bytes = cute.size_in_bytes( self.svdquant_dtype, cute.slice_(svdquant_a_smem_layout, (None, None, 0)), @@ -1095,7 +1098,10 @@ def kernel( self.svdquant_dtype, cute.slice_(svdquant_b_smem_layout, (None, None, 0)), ) - assert svdquant_copy_bytes == ab_copy_bytes + assert svdquant_copy_bytes == ab_copy_bytes, ( + "SVDQuant copy-size mismatch: correction bytes must equal " + "the mainloop A/B copy bytes" + ) # Allocate shared memory smem = cutlass.utils.SmemAllocator() diff --git a/flashinfer/trace/templates/gemm.py b/flashinfer/trace/templates/gemm.py index 36a0c15f742..f552dbabc5a 100644 --- a/flashinfer/trace/templates/gemm.py +++ b/flashinfer/trace/templates/gemm.py @@ -2143,6 +2143,7 @@ def _mm_nvfp4_svdquant_init( N: int = 3072, K: int = 3072, SF_A: int = 0, + SF_B: int = 0, device: str = "cuda", seed: int = 0, ): @@ -2156,7 +2157,7 @@ def _mm_nvfp4_svdquant_init( """ from flashinfer import nvfp4_quantize_smooth # noqa: PLC0415 - del SF_A # output-only / derived axis + del SF_A, SF_B # derived axes torch.manual_seed(seed) rank = 32 @@ -2206,7 +2207,9 @@ def _mm_nvfp4_svdquant_init( "SF_A": Var( description="128x4-swizzled activation scale buffer size derived from M and K." ), - "SF_B": Const(description="128x4-swizzled weight scale buffer size."), + "SF_B": Var( + description="128x4-swizzled weight scale buffer size derived from N and K." + ), "rank": Const(description="LoRA rank, a positive multiple of 32."), }, inputs={ @@ -2251,6 +2254,7 @@ def _mm_nvfp4_svdquant_init( }, constraints=[ "SF_A == ((M + 127) // 128) * 128 * (((K_packed * 2 // 16) + 3) // 4) * 4", + "SF_B == ((N + 127) // 128) * 128 * (((K_packed * 2 // 16) + 3) // 4) * 4", ], tags=["quantization:fp4"], init=_mm_nvfp4_svdquant_init, diff --git a/tests/gemm/test_mm_fp4.py b/tests/gemm/test_mm_fp4.py index 5af59d6c04b..66f4f7a5c0b 100644 --- a/tests/gemm/test_mm_fp4.py +++ b/tests/gemm/test_mm_fp4.py @@ -49,8 +49,12 @@ def _test_mm_fp4( pytest.skip("b12x backend only supports 128x4 SF layout") if compute_capability[0] != 12: pytest.skip("b12x backend only supports SM120/SM121 GPUs.") - if torch.version.cuda and int(torch.version.cuda.split(".")[0]) < 13: - pytest.skip("b12x backend requires CUDA 13+.") + min_cuda_version = "12.9" if use_nvfp4 else "13.0" + if not version_at_least(torch.version.cuda, min_cuda_version): + pytest.skip( + f"b12x {'NVFP4' if use_nvfp4 else 'MXFP4'} backend requires " + f"CUDA {min_cuda_version}+." + ) if not use_128x4_sf_layout and backend != "trtllm": pytest.skip("Skipping test for non-trtllm fp4 with use_128x4_sf_layout=False") if not use_nvfp4 and backend not in ["cudnn", "auto", "cute-dsl", "b12x"]: @@ -195,9 +199,9 @@ def test_mm_fp4_b12x_ragged_k(k, auto_tuning): def test_mm_fp4_b12x_misaligned_k_raises(): device = torch.device("cuda") if not ( - is_sm12x_supported(device) and version_at_least(torch.version.cuda, "13.0") + is_sm12x_supported(device) and version_at_least(torch.version.cuda, "12.9") ): - pytest.skip("b12x backend requires SM120/SM121 + CUDA 13+.") + pytest.skip("b12x backend requires SM120/SM121 + CUDA 12.9+.") m, n, k = 64, 512, 112 # k % 32 == 16 _, _, a_fp4, a_s, b_fp4, b_s, alpha = _nvfp4_operands(m, n, k) res = torch.empty([m, n], device="cuda", dtype=torch.bfloat16) diff --git a/tests/trace/example.py b/tests/trace/example.py index 65ca299ae4c..d6f59d56985 100644 --- a/tests/trace/example.py +++ b/tests/trace/example.py @@ -31,7 +31,7 @@ gemm_fp8_N1536_K7168.json gemm_fp8_nt_groupwise_n1536_k7168.json gemm_mxfp8_N4096_K4096.json -gemm_nvfp4_svdquant_N3072_K_packed1536_SF_A24576_SF_B589824_rank32.json +gemm_nvfp4_svdquant_N3072_K_packed1536_rank32.json gemma_fused_add_rmsnorm_h4608.json gemma_rmsnorm_h4608.json gelu_and_mul_h16384.json diff --git a/tests/trace/fi_trace_out/gemm_nvfp4_svdquant_N3072_K_packed1536_SF_B589824_rank32.json b/tests/trace/fi_trace_out/gemm_nvfp4_svdquant_N3072_K_packed1536_rank32.json similarity index 85% rename from tests/trace/fi_trace_out/gemm_nvfp4_svdquant_N3072_K_packed1536_SF_B589824_rank32.json rename to tests/trace/fi_trace_out/gemm_nvfp4_svdquant_N3072_K_packed1536_rank32.json index c000b9e762a..99b6b5843cd 100644 --- a/tests/trace/fi_trace_out/gemm_nvfp4_svdquant_N3072_K_packed1536_SF_B589824_rank32.json +++ b/tests/trace/fi_trace_out/gemm_nvfp4_svdquant_N3072_K_packed1536_rank32.json @@ -1,5 +1,5 @@ { - "name": "gemm_nvfp4_svdquant_N3072_K_packed1536_SF_B589824_rank32", + "name": "gemm_nvfp4_svdquant_N3072_K_packed1536_rank32", "description": "SVDQuant NVFP4 GEMM: out = alpha * (a @ b\u1d40 + d @ l1\u1d40). SM100/SM103 use fused CUTLASS; SM120/SM121 use fused CuTe DSL, with an explicit cute-dsl-unfused oracle. 1/alpha is pre-folded into l1.", "op_type": "gemm_nvfp4_svdquant", "tags": [ @@ -24,9 +24,8 @@ "description": "128x4-swizzled activation scale buffer size derived from M and K." }, "SF_B": { - "type": "const", - "value": 589824, - "description": "128x4-swizzled weight scale buffer size." + "type": "var", + "description": "128x4-swizzled weight scale buffer size derived from N and K." }, "rank": { "type": "const", @@ -35,7 +34,8 @@ } }, "constraints": [ - "SF_A == ((M + 127) // 128) * 128 * (((K_packed * 2 // 16) + 3) // 4) * 4" + "SF_A == ((M + 127) // 128) * 128 * (((K_packed * 2 // 16) + 3) // 4) * 4", + "SF_B == ((N + 127) // 128) * 128 * (((K_packed * 2 // 16) + 3) // 4) * 4" ], "inputs": { "a": { @@ -102,5 +102,5 @@ } }, "check": "def standard_check(\n reference_outputs: Any,\n actual_outputs: Any,\n *,\n rtol: Optional[float] = None,\n atol: Optional[float] = None,\n max_mismatch_pct: float = 0.0,\n min_cos_sim: Optional[float] = 1.0 - 1e-3,\n) -> bool:\n \"\"\"Default trace correctness check used when a template does not override it.\"\"\"\n from flashinfer.trace import default_check\n\n return default_check(\n reference_outputs,\n actual_outputs,\n rtol=rtol,\n atol=atol,\n max_mismatch_pct=max_mismatch_pct,\n min_cos_sim=min_cos_sim,\n )\n", - "init": "from __future__ import annotations\nimport math\nimport torch\n\n# ----- shared init helpers -----\n# Copyright (c) 2025 by FlashInfer team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Shared helpers used by ``TraceTemplate.init`` functions.\n\nThis module contains the small set of input-construction patterns that\nrecur across many templates (paged-KV cache index arrays, ragged indptr,\nRoPE pos_ids and cos/sin caches, sampling probs). Each helper is short and\ndocumented; init functions in ``templates/.py`` call into here so\nthe per-template init bodies stay focused on shape/dtype, not boilerplate.\n\nThe full source of this module is **inlined into every dumped JSON's\n``\"init\"`` field** by ``flashinfer/trace/template.py:_render_init_source``,\nso downstream consumers don't need flashinfer installed to re-run the init\nsnippets.\n\"\"\"\n\n\nfrom typing import Optional, Tuple\n\nimport torch\n\n\ndef make_paged_kv_indices(\n batch_size: int,\n num_pages_per_seq: int,\n page_size: int,\n *,\n device: str = \"cuda\",\n) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n \"\"\"Return ``(kv_indptr, kv_indices, kv_last_page_len)`` for a uniform batch.\n\n Every sequence is assigned exactly ``num_pages_per_seq`` pages, fully\n populated (last-page length == page_size).\n\n Invariants\n ----------\n - ``kv_indptr.shape == (batch_size + 1,)``, dtype int32, monotonic, [0]=0.\n - ``kv_indices == arange(0, batch_size * num_pages_per_seq)``, int32.\n - ``kv_last_page_len == full(batch_size, page_size)``, int32.\n \"\"\"\n total_pages = batch_size * num_pages_per_seq\n kv_indptr = (\n torch.arange(batch_size + 1, dtype=torch.int32, device=device)\n * num_pages_per_seq\n )\n kv_indices = torch.arange(total_pages, dtype=torch.int32, device=device)\n kv_last_page_len = torch.full(\n (batch_size,), page_size, dtype=torch.int32, device=device\n )\n return kv_indptr, kv_indices, kv_last_page_len\n\n\ndef make_ragged_indptr(\n seg_lens,\n *,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.int32,\n) -> torch.Tensor:\n \"\"\"Return cumulative-sum ``indptr`` of length ``len(seg_lens)+1``.\n\n ``seg_lens`` may be a list / tuple / 1-D tensor of segment lengths.\n \"\"\"\n if isinstance(seg_lens, torch.Tensor):\n lens = seg_lens.to(device=device, dtype=dtype)\n else:\n lens = torch.tensor(list(seg_lens), dtype=dtype, device=device)\n indptr = torch.zeros(lens.numel() + 1, dtype=dtype, device=device)\n indptr[1:] = torch.cumsum(lens, dim=0).to(dtype)\n return indptr\n\n\ndef make_uniform_qo_indptr(\n batch_size: int,\n qo_len: int,\n *,\n device: str = \"cuda\",\n) -> torch.Tensor:\n \"\"\"Return ``[0, qo_len, 2*qo_len, ..., batch_size*qo_len]`` int32.\"\"\"\n return torch.arange(batch_size + 1, dtype=torch.int32, device=device) * qo_len\n\n\ndef make_pos_ids(\n nnz: int,\n max_seq_len: Optional[int] = None,\n *,\n device: str = \"cuda\",\n) -> torch.Tensor:\n \"\"\"Return ``[0, 1, ..., nnz-1] (% max_seq_len)`` as int32 on ``device``.\n\n If ``max_seq_len`` is None, no wrapping is applied.\n \"\"\"\n pos = torch.arange(nnz, dtype=torch.int32, device=device)\n if max_seq_len is not None:\n pos = pos % max_seq_len\n return pos\n\n\ndef make_rope_cos_sin_cache(\n max_seq_len: int,\n rope_dim: int,\n *,\n base: float = 1e4,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Return concatenated ``[cos | sin]`` cache of shape ``[max_seq_len, rope_dim]``.\"\"\"\n t = torch.arange(max_seq_len, dtype=torch.float32, device=device)\n inv = 1.0 / (\n base\n ** (torch.arange(0, rope_dim, 2, dtype=torch.float32, device=device) / rope_dim)\n )\n freqs = t.unsqueeze(-1) * inv.unsqueeze(0)\n cache = torch.cat([torch.cos(freqs), torch.sin(freqs)], dim=-1)\n return cache.to(dtype)\n\n\ndef make_probs(\n batch_size: int,\n vocab_size: int,\n *,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Return a ``[batch_size, vocab_size]`` probability distribution.\n\n Uses ``softmax(randn(...))`` so each row sums to 1.0. This mirrors the\n pattern used throughout ``tests/utils/test_sampling.py``.\n \"\"\"\n return torch.softmax(\n torch.randn(batch_size, vocab_size, dtype=torch.float32, device=device),\n dim=-1,\n ).to(dtype)\n\n\ndef make_logits(\n batch_size: int,\n vocab_size: int,\n *,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Return ``randn(batch_size, vocab_size)`` logits.\"\"\"\n return torch.randn(batch_size, vocab_size, dtype=dtype, device=device)\n\n\ndef fp8_safe_randn(\n *shape: int,\n scale: float = 0.1,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.bfloat16,\n) -> torch.Tensor:\n \"\"\"``randn(*shape) * scale`` \u2014 keeps values in the FP8/FP4 representable range.\n\n Tests for fp8/fp4 paths typically multiply ``randn`` by 0.1 to avoid\n saturation when quantizing. Use this helper to mirror that convention.\n \"\"\"\n return (torch.randn(*shape, dtype=dtype, device=device) * scale).to(dtype)\n\n\ndef per_tensor_fp8_quantize(\n x: torch.Tensor,\n *,\n fp8_dtype: torch.dtype = torch.float8_e4m3fn,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Per-tensor FP8 quantization, mirroring ``tests/utils_fp8.py:to_float8``.\n\n Returns ``(x_fp8, inv_scale)`` where ``inv_scale`` is the dequant\n multiplier (``float \u2248 fp8 * inv_scale``).\n \"\"\"\n finfo = torch.finfo(fp8_dtype)\n amax = x.abs().amax().clamp(min=1e-12)\n scale = finfo.max / amax\n x_q = (x.float() * scale).clamp(min=finfo.min, max=finfo.max).to(fp8_dtype)\n return x_q, scale.float().reciprocal()\n\n\ndef fp8_block_quant_1d(\n x_bf16: torch.Tensor,\n block: int = 128,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Quantize ``[T, H]`` activations into FP8 with per-``(token, block)``\n column-block scales. Returns ``(x_fp8, scales)`` where\n ``scales`` has shape ``[T, H // block]``.\n\n Mirrors ``_fp8_block_quant_1d`` in\n ``tests/moe/test_dpsk_fused_moe_fp8.py``.\n \"\"\"\n assert x_bf16.dim() == 2\n T, H = x_bf16.shape\n assert H % block == 0\n nb = H // block\n finfo = torch.finfo(torch.float8_e4m3fn)\n max_fp8 = finfo.max\n x_f32 = x_bf16.to(torch.float32)\n x_fp8 = torch.empty((T, H), dtype=torch.float8_e4m3fn, device=x_bf16.device)\n scales = torch.empty((T, nb), dtype=torch.float32, device=x_bf16.device)\n for j in range(nb):\n sl = slice(j * block, (j + 1) * block)\n blk = x_f32[:, sl]\n amax = torch.amax(torch.abs(blk), dim=1)\n s = torch.where(amax > 0, amax / max_fp8, torch.ones_like(amax))\n x_fp8[:, sl] = (blk / s.unsqueeze(1)).to(torch.float8_e4m3fn)\n scales[:, j] = s\n return x_fp8, scales\n\n\ndef fp8_block_quant_2d(\n w_bf16: torch.Tensor,\n block: int = 128,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Quantize weights ``[..., R, C]`` with 2-D ``block \u00d7 block`` scales.\n\n Returns ``(w_fp8, scales)`` where ``scales`` has shape\n ``[..., R // block, C // block]``. Mirrors ``_fp8_block_quant_2d`` in\n ``tests/moe/test_dpsk_fused_moe_fp8.py``.\n \"\"\"\n assert w_bf16.dim() >= 2\n *prefix, R, C = w_bf16.shape\n assert R % block == 0 and C % block == 0\n nb_r, nb_c = R // block, C // block\n finfo = torch.finfo(torch.float8_e4m3fn)\n max_fp8 = finfo.max\n w_f32 = w_bf16.to(torch.float32).contiguous()\n prefix_ndim = len(prefix)\n reshaped = w_f32.reshape(*prefix, nb_r, block, nb_c, block)\n permute_dims = tuple(range(prefix_ndim)) + (\n prefix_ndim,\n prefix_ndim + 2,\n prefix_ndim + 1,\n prefix_ndim + 3,\n )\n blocks = reshaped.permute(permute_dims).contiguous()\n amax = torch.amax(torch.abs(blocks), dim=(-1, -2))\n scales = torch.where(\n amax > 0, amax / max_fp8, torch.ones_like(amax, dtype=torch.float32)\n )\n q_blocks = (blocks / scales.unsqueeze(-1).unsqueeze(-1)).to(torch.float8_e4m3fn)\n inv_permute = [0] * (prefix_ndim + 4)\n for i, p in enumerate(permute_dims):\n inv_permute[p] = i\n w_fp8 = q_blocks.permute(*inv_permute).reshape(*prefix, R, C).contiguous()\n return w_fp8, scales\n\n\n__all__ = [\n \"make_paged_kv_indices\",\n \"make_ragged_indptr\",\n \"make_uniform_qo_indptr\",\n \"make_pos_ids\",\n \"make_rope_cos_sin_cache\",\n \"make_probs\",\n \"make_logits\",\n \"fp8_safe_randn\",\n \"per_tensor_fp8_quantize\",\n \"fp8_block_quant_1d\",\n \"fp8_block_quant_2d\",\n]\n\n# ----- init -----\ndef _mm_nvfp4_svdquant_init(\n *,\n M: int,\n N: int = 3072,\n K: int = 3072,\n SF_A: int = 0,\n device: str = \"cuda\",\n seed: int = 0,\n):\n \"\"\"Build inputs for ``flashinfer.mm_nvfp4_svdquant``.\n\n Mirrors the SVDQuant linear decomposition W ~= R + L1 @ L2 at Qwen-Image\n shapes: the residual weight is NVFP4-quantized (via ``nvfp4_quantize_smooth``\n with a unit smoothing scale, which is byte-identical to the stock NVFP4\n quantizer), the activation is smooth-quantized, and the rank-32 LoRA factors\n follow the host-side folding contract (``d = x_hat @ L2\u1d40``, ``l1 = L1 / alpha``).\n \"\"\"\n from flashinfer import nvfp4_quantize_smooth # noqa: PLC0415\n\n del SF_A # output-only / derived axis\n\n torch.manual_seed(seed)\n rank = 32\n x = torch.randn(M, K, dtype=torch.bfloat16, device=device)\n w = torch.randn(N, K, dtype=torch.bfloat16, device=device) / math.sqrt(K)\n pqs = torch.rand(K, dtype=torch.bfloat16, device=device) + 0.5\n l1 = torch.randn(N, rank, dtype=torch.bfloat16, device=device) / math.sqrt(rank)\n l2 = torch.randn(rank, K, dtype=torch.bfloat16, device=device) / math.sqrt(K)\n\n x_hat = (x.float() * pqs.float()).to(torch.bfloat16)\n x_gs = (\n ((448 * 6) / x_hat.float().abs().nan_to_num().max())\n .to(torch.float32)\n .reshape(1)\n )\n w_gs = ((448 * 6) / w.float().abs().nan_to_num().max()).to(torch.float32).reshape(1)\n ones = torch.ones(K, dtype=torch.bfloat16, device=device)\n\n a, a_sf = nvfp4_quantize_smooth(x, pqs, x_gs)\n b, b_sf = nvfp4_quantize_smooth(w, ones, w_gs)\n alpha = (1.0 / (x_gs * w_gs)).to(torch.float32).reshape(1)\n d = torch.mm(x_hat, l2.t().contiguous().to(torch.bfloat16))\n l1_scaled = (l1.float() / alpha.item()).to(torch.bfloat16)\n return {\n \"a\": a,\n \"b\": b,\n \"a_sf\": a_sf,\n \"b_sf\": b_sf,\n \"alpha\": alpha,\n \"d\": d,\n \"l1\": l1_scaled,\n }\n" + "init": "from __future__ import annotations\nimport math\nimport torch\n\n# ----- shared init helpers -----\n# Copyright (c) 2025 by FlashInfer team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Shared helpers used by ``TraceTemplate.init`` functions.\n\nThis module contains the small set of input-construction patterns that\nrecur across many templates (paged-KV cache index arrays, ragged indptr,\nRoPE pos_ids and cos/sin caches, sampling probs). Each helper is short and\ndocumented; init functions in ``templates/.py`` call into here so\nthe per-template init bodies stay focused on shape/dtype, not boilerplate.\n\nThe full source of this module is **inlined into every dumped JSON's\n``\"init\"`` field** by ``flashinfer/trace/template.py:_render_init_source``,\nso downstream consumers don't need flashinfer installed to re-run the init\nsnippets.\n\"\"\"\n\n\nfrom typing import Optional, Tuple\n\nimport torch\n\n\ndef make_paged_kv_indices(\n batch_size: int,\n num_pages_per_seq: int,\n page_size: int,\n *,\n device: str = \"cuda\",\n) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n \"\"\"Return ``(kv_indptr, kv_indices, kv_last_page_len)`` for a uniform batch.\n\n Every sequence is assigned exactly ``num_pages_per_seq`` pages, fully\n populated (last-page length == page_size).\n\n Invariants\n ----------\n - ``kv_indptr.shape == (batch_size + 1,)``, dtype int32, monotonic, [0]=0.\n - ``kv_indices == arange(0, batch_size * num_pages_per_seq)``, int32.\n - ``kv_last_page_len == full(batch_size, page_size)``, int32.\n \"\"\"\n total_pages = batch_size * num_pages_per_seq\n kv_indptr = (\n torch.arange(batch_size + 1, dtype=torch.int32, device=device)\n * num_pages_per_seq\n )\n kv_indices = torch.arange(total_pages, dtype=torch.int32, device=device)\n kv_last_page_len = torch.full(\n (batch_size,), page_size, dtype=torch.int32, device=device\n )\n return kv_indptr, kv_indices, kv_last_page_len\n\n\ndef make_ragged_indptr(\n seg_lens,\n *,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.int32,\n) -> torch.Tensor:\n \"\"\"Return cumulative-sum ``indptr`` of length ``len(seg_lens)+1``.\n\n ``seg_lens`` may be a list / tuple / 1-D tensor of segment lengths.\n \"\"\"\n if isinstance(seg_lens, torch.Tensor):\n lens = seg_lens.to(device=device, dtype=dtype)\n else:\n lens = torch.tensor(list(seg_lens), dtype=dtype, device=device)\n indptr = torch.zeros(lens.numel() + 1, dtype=dtype, device=device)\n indptr[1:] = torch.cumsum(lens, dim=0).to(dtype)\n return indptr\n\n\ndef make_uniform_qo_indptr(\n batch_size: int,\n qo_len: int,\n *,\n device: str = \"cuda\",\n) -> torch.Tensor:\n \"\"\"Return ``[0, qo_len, 2*qo_len, ..., batch_size*qo_len]`` int32.\"\"\"\n return torch.arange(batch_size + 1, dtype=torch.int32, device=device) * qo_len\n\n\ndef make_pos_ids(\n nnz: int,\n max_seq_len: Optional[int] = None,\n *,\n device: str = \"cuda\",\n) -> torch.Tensor:\n \"\"\"Return ``[0, 1, ..., nnz-1] (% max_seq_len)`` as int32 on ``device``.\n\n If ``max_seq_len`` is None, no wrapping is applied.\n \"\"\"\n pos = torch.arange(nnz, dtype=torch.int32, device=device)\n if max_seq_len is not None:\n pos = pos % max_seq_len\n return pos\n\n\ndef make_rope_cos_sin_cache(\n max_seq_len: int,\n rope_dim: int,\n *,\n base: float = 1e4,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Return concatenated ``[cos | sin]`` cache of shape ``[max_seq_len, rope_dim]``.\"\"\"\n t = torch.arange(max_seq_len, dtype=torch.float32, device=device)\n inv = 1.0 / (\n base\n ** (torch.arange(0, rope_dim, 2, dtype=torch.float32, device=device) / rope_dim)\n )\n freqs = t.unsqueeze(-1) * inv.unsqueeze(0)\n cache = torch.cat([torch.cos(freqs), torch.sin(freqs)], dim=-1)\n return cache.to(dtype)\n\n\ndef make_probs(\n batch_size: int,\n vocab_size: int,\n *,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Return a ``[batch_size, vocab_size]`` probability distribution.\n\n Uses ``softmax(randn(...))`` so each row sums to 1.0. This mirrors the\n pattern used throughout ``tests/utils/test_sampling.py``.\n \"\"\"\n return torch.softmax(\n torch.randn(batch_size, vocab_size, dtype=torch.float32, device=device),\n dim=-1,\n ).to(dtype)\n\n\ndef make_logits(\n batch_size: int,\n vocab_size: int,\n *,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Return ``randn(batch_size, vocab_size)`` logits.\"\"\"\n return torch.randn(batch_size, vocab_size, dtype=dtype, device=device)\n\n\ndef fp8_safe_randn(\n *shape: int,\n scale: float = 0.1,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.bfloat16,\n) -> torch.Tensor:\n \"\"\"``randn(*shape) * scale`` \u2014 keeps values in the FP8/FP4 representable range.\n\n Tests for fp8/fp4 paths typically multiply ``randn`` by 0.1 to avoid\n saturation when quantizing. Use this helper to mirror that convention.\n \"\"\"\n return (torch.randn(*shape, dtype=dtype, device=device) * scale).to(dtype)\n\n\ndef per_tensor_fp8_quantize(\n x: torch.Tensor,\n *,\n fp8_dtype: torch.dtype = torch.float8_e4m3fn,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Per-tensor FP8 quantization, mirroring ``tests/utils_fp8.py:to_float8``.\n\n Returns ``(x_fp8, inv_scale)`` where ``inv_scale`` is the dequant\n multiplier (``float \u2248 fp8 * inv_scale``).\n \"\"\"\n finfo = torch.finfo(fp8_dtype)\n amax = x.abs().amax().clamp(min=1e-12)\n scale = finfo.max / amax\n x_q = (x.float() * scale).clamp(min=finfo.min, max=finfo.max).to(fp8_dtype)\n return x_q, scale.float().reciprocal()\n\n\ndef fp8_block_quant_1d(\n x_bf16: torch.Tensor,\n block: int = 128,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Quantize ``[T, H]`` activations into FP8 with per-``(token, block)``\n column-block scales. Returns ``(x_fp8, scales)`` where\n ``scales`` has shape ``[T, H // block]``.\n\n Mirrors ``_fp8_block_quant_1d`` in\n ``tests/moe/test_dpsk_fused_moe_fp8.py``.\n \"\"\"\n assert x_bf16.dim() == 2\n T, H = x_bf16.shape\n assert H % block == 0\n nb = H // block\n finfo = torch.finfo(torch.float8_e4m3fn)\n max_fp8 = finfo.max\n x_f32 = x_bf16.to(torch.float32)\n x_fp8 = torch.empty((T, H), dtype=torch.float8_e4m3fn, device=x_bf16.device)\n scales = torch.empty((T, nb), dtype=torch.float32, device=x_bf16.device)\n for j in range(nb):\n sl = slice(j * block, (j + 1) * block)\n blk = x_f32[:, sl]\n amax = torch.amax(torch.abs(blk), dim=1)\n s = torch.where(amax > 0, amax / max_fp8, torch.ones_like(amax))\n x_fp8[:, sl] = (blk / s.unsqueeze(1)).to(torch.float8_e4m3fn)\n scales[:, j] = s\n return x_fp8, scales\n\n\ndef fp8_block_quant_2d(\n w_bf16: torch.Tensor,\n block: int = 128,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Quantize weights ``[..., R, C]`` with 2-D ``block \u00d7 block`` scales.\n\n Returns ``(w_fp8, scales)`` where ``scales`` has shape\n ``[..., R // block, C // block]``. Mirrors ``_fp8_block_quant_2d`` in\n ``tests/moe/test_dpsk_fused_moe_fp8.py``.\n \"\"\"\n assert w_bf16.dim() >= 2\n *prefix, R, C = w_bf16.shape\n assert R % block == 0 and C % block == 0\n nb_r, nb_c = R // block, C // block\n finfo = torch.finfo(torch.float8_e4m3fn)\n max_fp8 = finfo.max\n w_f32 = w_bf16.to(torch.float32).contiguous()\n prefix_ndim = len(prefix)\n reshaped = w_f32.reshape(*prefix, nb_r, block, nb_c, block)\n permute_dims = tuple(range(prefix_ndim)) + (\n prefix_ndim,\n prefix_ndim + 2,\n prefix_ndim + 1,\n prefix_ndim + 3,\n )\n blocks = reshaped.permute(permute_dims).contiguous()\n amax = torch.amax(torch.abs(blocks), dim=(-1, -2))\n scales = torch.where(\n amax > 0, amax / max_fp8, torch.ones_like(amax, dtype=torch.float32)\n )\n q_blocks = (blocks / scales.unsqueeze(-1).unsqueeze(-1)).to(torch.float8_e4m3fn)\n inv_permute = [0] * (prefix_ndim + 4)\n for i, p in enumerate(permute_dims):\n inv_permute[p] = i\n w_fp8 = q_blocks.permute(*inv_permute).reshape(*prefix, R, C).contiguous()\n return w_fp8, scales\n\n\n__all__ = [\n \"make_paged_kv_indices\",\n \"make_ragged_indptr\",\n \"make_uniform_qo_indptr\",\n \"make_pos_ids\",\n \"make_rope_cos_sin_cache\",\n \"make_probs\",\n \"make_logits\",\n \"fp8_safe_randn\",\n \"per_tensor_fp8_quantize\",\n \"fp8_block_quant_1d\",\n \"fp8_block_quant_2d\",\n]\n\n# ----- init -----\ndef _mm_nvfp4_svdquant_init(\n *,\n M: int,\n N: int = 3072,\n K: int = 3072,\n SF_A: int = 0,\n SF_B: int = 0,\n device: str = \"cuda\",\n seed: int = 0,\n):\n \"\"\"Build inputs for ``flashinfer.mm_nvfp4_svdquant``.\n\n Mirrors the SVDQuant linear decomposition W ~= R + L1 @ L2 at Qwen-Image\n shapes: the residual weight is NVFP4-quantized (via ``nvfp4_quantize_smooth``\n with a unit smoothing scale, which is byte-identical to the stock NVFP4\n quantizer), the activation is smooth-quantized, and the rank-32 LoRA factors\n follow the host-side folding contract (``d = x_hat @ L2\u1d40``, ``l1 = L1 / alpha``).\n \"\"\"\n from flashinfer import nvfp4_quantize_smooth # noqa: PLC0415\n\n del SF_A, SF_B # derived axes\n\n torch.manual_seed(seed)\n rank = 32\n x = torch.randn(M, K, dtype=torch.bfloat16, device=device)\n w = torch.randn(N, K, dtype=torch.bfloat16, device=device) / math.sqrt(K)\n pqs = torch.rand(K, dtype=torch.bfloat16, device=device) + 0.5\n l1 = torch.randn(N, rank, dtype=torch.bfloat16, device=device) / math.sqrt(rank)\n l2 = torch.randn(rank, K, dtype=torch.bfloat16, device=device) / math.sqrt(K)\n\n x_hat = (x.float() * pqs.float()).to(torch.bfloat16)\n x_gs = (\n ((448 * 6) / x_hat.float().abs().nan_to_num().max())\n .to(torch.float32)\n .reshape(1)\n )\n w_gs = ((448 * 6) / w.float().abs().nan_to_num().max()).to(torch.float32).reshape(1)\n ones = torch.ones(K, dtype=torch.bfloat16, device=device)\n\n a, a_sf = nvfp4_quantize_smooth(x, pqs, x_gs)\n b, b_sf = nvfp4_quantize_smooth(w, ones, w_gs)\n alpha = (1.0 / (x_gs * w_gs)).to(torch.float32).reshape(1)\n d = torch.mm(x_hat, l2.t().contiguous().to(torch.bfloat16))\n l1_scaled = (l1.float() / alpha.item()).to(torch.bfloat16)\n return {\n \"a\": a,\n \"b\": b,\n \"a_sf\": a_sf,\n \"b_sf\": b_sf,\n \"alpha\": alpha,\n \"d\": d,\n \"l1\": l1_scaled,\n }\n" } From 900918dd1e383adac5ec8ee746adbe2bb64e5191 Mon Sep 17 00:00:00 2001 From: Anthony Chang <27950904+rosenrodt@users.noreply.github.com> Date: Tue, 18 Aug 2026 14:45:06 +0800 Subject: [PATCH 6/7] Address SVDQuant review feedback Preserve the historical alpha contract while normalizing every backend launch to a one-element scalar view. Remove the dead SM120 prefetch tactic dimension and harden cache, trace, quantization, graph, accuracy, and benchmark behavior. Changes - align and empty-guard smooth quantization inputs - separate SM100 prefetch from SM120 tactics and strengthen cache identities - add pooled-alpha, CUDA Graph, accuracy, trace, and benchmark coverage Validation - pre-commit run -a - python -m pytest -v tests/trace/ - python -m pytest -q tests/jit/test_cute_dsl_cache.py - python -m pytest -q tests/gemm/test_nvfp4_svdquant_gemm.py Result - 1197 trace tests passed with 182 skipped - 34 SM120 SVDQuant tests passed with 48 skipped - corrected RTX PRO 6000 cold-L2 CUDA Graph benchmark completed 12 rows --- benchmarks/bench_nvfp4_svdquant_gemm.py | 183 +++++++++++------ flashinfer/gemm/gemm_svdquant.py | 132 +++++++----- .../quantization/kernels/nvfp4_quantize.py | 34 ++- flashinfer/trace/templates/gemm.py | 9 +- tests/gemm/test_nvfp4_svdquant_gemm.py | 194 +++++++++++++++--- tests/jit/test_cute_dsl_cache.py | 116 +++++++++-- tests/trace/example.py | 3 +- .../fi_trace_out/gemm_bf16_N32_K3072.json | 51 ----- ...uant_N3072_K3072_K_packed1536_rank32.json} | 12 +- 9 files changed, 524 insertions(+), 210 deletions(-) delete mode 100644 tests/trace/fi_trace_out/gemm_bf16_N32_K3072.json rename tests/trace/fi_trace_out/{linear_nvfp4_svdquant_N3072_K3072_K_packed1536_SF_B589824_rank32.json => linear_nvfp4_svdquant_N3072_K3072_K_packed1536_rank32.json} (86%) diff --git a/benchmarks/bench_nvfp4_svdquant_gemm.py b/benchmarks/bench_nvfp4_svdquant_gemm.py index 10186056215..b58fd6b653d 100644 --- a/benchmarks/bench_nvfp4_svdquant_gemm.py +++ b/benchmarks/bench_nvfp4_svdquant_gemm.py @@ -112,6 +112,9 @@ def _build_case(m, n, k, rank, device): # bmm_fp8 expects B as column-major [batch, k, n]. Quantizing W.T retains # its column-major stride, while the singleton batch dimension is a view. w_fp8, w_fp8_scale = _to_float8(w.T) + assert w_fp8.stride(0) == 1, ( + f"bmm_fp8 requires a column-major B; got strides {w_fp8.stride()}" + ) return { "x": x, @@ -174,92 +177,133 @@ def bench_one( fused_backend = "cute-dsl" if get_compute_capability(device)[0] == 12 else "cutlass" - def run_fused(): + def run_fused(xq, wq, x_sf, w_sf, alpha, d, l1, bias, out): mm_nvfp4_svdquant( - c["xq"], - c["wq"], - c["x_sf_flat"], - c["w_sf_flat"], - c["alpha"], - c["d"], - c["l1_scaled"], - bias=c["bias"], - out=c["out_fused"], + xq, + wq, + x_sf, + w_sf, + alpha, + d, + l1, + bias=bias, + out=out, backend=fused_backend, ) - def run_svdquant_linear(): + def run_svdquant_linear( + x, wq, w_sf, alpha, pqs, l2t_smoothed, l1_scaled, global_sf, bias + ): svdquant_linear( - c["x"], - c["wq"], - c["w_sf_flat"], - c["alpha"], - c["pqs"], - c["l2t_smoothed"], - c["l1_scaled"], - c["global_sf"], - bias=c["bias"], + x, + wq, + w_sf, + alpha, + pqs, + l2t_smoothed, + l1_scaled, + global_sf, + bias=bias, backend=svdquant_backend, ) - def run_unfused(): + def run_unfused(xq, wq, x_sf, w_sf, alpha, d, l1, bias, out): mm_nvfp4_svdquant( - c["xq"], - c["wq"], - c["x_sf_flat"], - c["w_sf_flat"], - c["alpha"], - c["d"], - c["l1_scaled"], - bias=c["bias"], - out=c["out_fused"], + xq, + wq, + x_sf, + w_sf, + alpha, + d, + l1, + bias=bias, + out=out, backend=unfused_backend, ) - def run_mm_fp4(): + def run_mm_fp4(xq, wq, x_sf, w_sf, alpha, out): mm_fp4( - c["xq"], - c["wq"].T, - c["x_sf"], - c["w_sf"].T, - c["alpha"], + xq, + wq.T, + x_sf, + w_sf.T, + alpha, torch.bfloat16, - c["out_fp4"], + out, block_size=16, use_8x4_sf_layout=False, backend=mm_fp4_backend, use_nvfp4=True, ) - def run_residual_gemm(): + def run_residual_gemm(x, w, bias, out): torch.addmm( - c["bias"], - c["x"], - c["w"].T, - out=c["out_bf16"], + bias, + x, + w.T, + out=out, ) - def run_fp8_per_tensor(): + def run_fp8_per_tensor(x, w, x_scale, w_scale, out): bmm_fp8( - c["x_fp8"], - c["w_fp8"], - c["x_fp8_scale"], - c["w_fp8_scale"], + x, + w, + x_scale, + w_scale, torch.bfloat16, - out=c["out_fp8"], + out=out, backend="auto", ) + fused_args = ( + c["xq"], + c["wq"], + c["x_sf_flat"], + c["w_sf_flat"], + c["alpha"], + c["d"], + c["l1_scaled"], + c["bias"], + c["out_fused"], + ) + svdquant_linear_args = ( + c["x"], + c["wq"], + c["w_sf_flat"], + c["alpha"], + c["pqs"], + c["l2t_smoothed"], + c["l1_scaled"], + c["global_sf"], + c["bias"], + ) + mm_fp4_args = ( + c["xq"], + c["wq"], + c["x_sf"], + c["w_sf"], + c["alpha"], + c["out_fp4"], + ) + residual_gemm_args = (c["x"], c["w"], c["bias"], c["out_bf16"]) + fp8_args = ( + c["x_fp8"], + c["w_fp8"], + c["x_fp8_scale"], + c["w_fp8_scale"], + c["out_fp8"], + ) + # Tune once; subsequent calls replay the best tactic from the tuner cache. with autotune(True): for _ in range(3): - run_fused() + run_fused(*fused_args) if unfused_backend is not None: - run_unfused() - run_svdquant_linear() - run_mm_fp4() - run_residual_gemm() - run_fp8_per_tensor() + run_unfused(*fused_args) + run_svdquant_linear(*svdquant_linear_args) + run_mm_fp4(*mm_fp4_args) + run_residual_gemm(*residual_gemm_args) + run_fp8_per_tensor(*fp8_args) torch.cuda.synchronize() bench_kwargs = dict( @@ -269,16 +313,28 @@ def run_fp8_per_tensor(): enable_cupti=True, cold_l2_cache=cold_l2_cache, ) - fused_us = _median_us(bench_gpu_time(run_fused, **bench_kwargs)) + fused_us = _median_us( + bench_gpu_time(run_fused, input_args=fused_args, **bench_kwargs) + ) unfused_us = ( - _median_us(bench_gpu_time(run_unfused, **bench_kwargs)) + _median_us(bench_gpu_time(run_unfused, input_args=fused_args, **bench_kwargs)) if unfused_backend is not None else float("nan") ) - svdquant_linear_us = _median_us(bench_gpu_time(run_svdquant_linear, **bench_kwargs)) - mm_fp4_us = _median_us(bench_gpu_time(run_mm_fp4, **bench_kwargs)) - residual_gemm_us = _median_us(bench_gpu_time(run_residual_gemm, **bench_kwargs)) - fp8_per_tensor_us = _median_us(bench_gpu_time(run_fp8_per_tensor, **bench_kwargs)) + svdquant_linear_us = _median_us( + bench_gpu_time( + run_svdquant_linear, input_args=svdquant_linear_args, **bench_kwargs + ) + ) + mm_fp4_us = _median_us( + bench_gpu_time(run_mm_fp4, input_args=mm_fp4_args, **bench_kwargs) + ) + residual_gemm_us = _median_us( + bench_gpu_time(run_residual_gemm, input_args=residual_gemm_args, **bench_kwargs) + ) + fp8_per_tensor_us = _median_us( + bench_gpu_time(run_fp8_per_tensor, input_args=fp8_args, **bench_kwargs) + ) return ( fused_us, @@ -361,10 +417,13 @@ def main(): print(f"Device: {torch.cuda.get_device_name(device)} (SM{major}{minor})") print(f"mm_fp4 baseline backend: {mm_fp4_backend}") print(f"SVDQuant implementation policy: {args.svdquant_backend}") - print("BF16 residual_gemm baseline: torch.addmm (PyTorch-selected CUDA backend)") + print( + "BF16 residual_gemm baseline: torch.addmm (PyTorch-selected CUDA backend; " + "residual + bias only, no LoRA-down or LoRA-up correction)" + ) print( "FP8 per-tensor baseline: bmm_fp8 backend=auto " - "(scales and quantization excluded)" + "(residual only; scales, quantization, LoRA-down, and LoRA-up excluded)" ) print(f"unfused oracle backend: {unfused_backend or 'not available'}") print(f"execution mode: {'CUDA graph' if args.cuda_graph else 'eager'}") diff --git a/flashinfer/gemm/gemm_svdquant.py b/flashinfer/gemm/gemm_svdquant.py index e26ab89c781..71b62502fc3 100644 --- a/flashinfer/gemm/gemm_svdquant.py +++ b/flashinfer/gemm/gemm_svdquant.py @@ -75,6 +75,37 @@ def _view_128x4_sf(sf: torch.Tensor, rows: int, sf_cols: int) -> torch.Tensor: return sf.reshape(-1)[:size].view(_pad_up(rows, 128), _pad_up(sf_cols, 4)) +def _svdquant_kernel_source_files() -> Tuple[str, ...]: + """Sources whose device-code changes invalidate the SM120 disk cache.""" + from ..cute_dsl import utils as cute_dsl_utils + from .kernels import dense_blockscaled_gemm_sm120_b12x + + return ( + __file__, + dense_blockscaled_gemm_sm120_b12x.__file__, + cute_dsl_utils.__file__, + ) + + +def _sm120_svdquant_kernel_name( + *, + rank: int, + with_bias: bool, + mma_tiler_mn: Tuple[int, int], + tile_k: int, + swap_ab: bool, + max_active_clusters: int, + enable_pdl: bool, + enable_iket: bool, +) -> str: + """Return a symbol-safe name encoding every SM120 codegen parameter.""" + return ( + f"r{rank}_bias{int(with_bias)}_t{mma_tiler_mn[0]}x{mma_tiler_mn[1]}" + f"x{tile_k}_swap{int(swap_ab)}_mac{max_active_clusters}" + f"_pdl{int(enable_pdl)}_iket{int(enable_iket)}" + ) + + def _compile_sm120_nvfp4_svdquant( *, device: torch.device, @@ -83,7 +114,6 @@ def _compile_sm120_nvfp4_svdquant( mma_tiler_mn: Tuple[int, int], tile_k: int, swap_ab: bool, - use_prefetch: bool, sf_m: int, sf_n: int, sf_k: int, @@ -105,7 +135,6 @@ def _compile_sm120_nvfp4_svdquant( mma_tiler_mn, tile_k, swap_ab, - use_prefetch, max_active_clusters, enable_pdl, enable_iket, @@ -119,7 +148,6 @@ def _compile_sm120_nvfp4_svdquant( from cutlass.cute.runtime import make_ptr from ..jit.cute_dsl_core import build_and_load_cute_dsl_kernel - from .kernels import dense_blockscaled_gemm_sm120_b12x from .kernels.dense_blockscaled_gemm_sm120_b12x import ( Sm120B12xBlockScaledDenseGemmKernel, ) @@ -129,7 +157,9 @@ def _compile_sm120_nvfp4_svdquant( mma_tiler_mn, (1, 1), tile_k=tile_k, - use_prefetch=use_prefetch, + # The shared b12x constructor retains this generic-GEMM knob, but the + # SM120 SVDQuant kernel has no prefetch dataflow specialization. + use_prefetch=False, enable_pdl=enable_pdl, swap_ab=swap_ab, enable_iket=enable_iket, @@ -203,16 +233,21 @@ def compile_kernel(): options="--opt-level 2 --enable-tvm-ffi", ) - kernel_name = ( - f"r{rank}_bias{int(with_bias)}_t{mma_tiler_mn[0]}x{mma_tiler_mn[1]}" - f"x{tile_k}_swap{int(swap_ab)}_pf{int(use_prefetch)}_mac{max_active_clusters}" - f"_pdl{int(enable_pdl)}_iket{int(enable_iket)}" + kernel_name = _sm120_svdquant_kernel_name( + rank=rank, + with_bias=with_bias, + mma_tiler_mn=mma_tiler_mn, + tile_k=tile_k, + swap_ab=swap_ab, + max_active_clusters=max_active_clusters, + enable_pdl=enable_pdl, + enable_iket=enable_iket, ) compiled = build_and_load_cute_dsl_kernel( "mm_nvfp4_svdquant_sm120", kernel_name, compile_kernel, - extra_key_files=(__file__, dense_blockscaled_gemm_sm120_b12x.__file__), + extra_key_files=_svdquant_kernel_source_files(), ) _SM120_SVDQUANT_KERNEL_CACHE[cache_key] = compiled return compiled @@ -246,8 +281,8 @@ def _mm_nvfp4_svdquant_sm120_fused( plan = _select_default_dense_gemm_plan( m, n, real_k, get_device_sm_count(a.device), expected_m=m ) - tactic = (plan.mma_tiler_mn, 128, plan.swap_ab, False) - mma_tiler_mn, tile_k, swap_ab, use_prefetch = tactic + tactic = (plan.mma_tiler_mn, 128, plan.swap_ab) + mma_tiler_mn, tile_k, swap_ab = tactic compiled = _compile_sm120_nvfp4_svdquant( device=a.device, rank=d.shape[1], @@ -255,7 +290,6 @@ def _mm_nvfp4_svdquant_sm120_fused( mma_tiler_mn=mma_tiler_mn, tile_k=tile_k, swap_ab=swap_ab, - use_prefetch=use_prefetch, sf_m=sf_m, sf_n=sf_n, sf_k=sf_k, @@ -271,7 +305,7 @@ def _mm_nvfp4_svdquant_sm120_fused( sf_k, a_sf.data_ptr(), b_sf.data_ptr(), - alpha.reshape(1), + alpha, d, l1, bias, @@ -323,6 +357,9 @@ def _mm_nvfp4_svdquant_sm120_unfused( def _sm120_nvfp4_svdquant_runner(enable_pdl: bool): class Sm120Nvfp4SvdquantRunner(TunableRunner): + def get_cache_key_extras(self, inputs: List[torch.Tensor]) -> tuple: + return (enable_pdl,) + def get_valid_tactics( self, inputs: List[torch.Tensor], @@ -347,7 +384,6 @@ def _add( mma_tiler_mn, tile_k, swap_ab, - prefetch_candidates=(False, True), ): if not Sm120B12xBlockScaledDenseGemmKernel.can_implement( cutlass.Float4E2M1FN, @@ -367,15 +403,9 @@ def _add( tile_k=tile_k, ): return - for use_prefetch in prefetch_candidates: - tactic = ( - mma_tiler_mn, - tile_k, - swap_ab, - use_prefetch, - ) - if tactic not in tactics: - tactics.append(tactic) + tactic = (mma_tiler_mn, tile_k, swap_ab) + if tactic not in tactics: + tactics.append(tactic) for mma_tiler_mn in ((64, 64), (64, 128), (128, 64), (128, 128)): _add(mma_tiler_mn, 128, swap_ab=False) @@ -385,19 +415,9 @@ def _add( ) _add(plan.mma_tiler_mn, 128, plan.swap_ab) for tile_k in (64, 256): - _add( - plan.mma_tiler_mn, - tile_k, - plan.swap_ab, - prefetch_candidates=(False,), - ) + _add(plan.mma_tiler_mn, tile_k, plan.swap_ab) if m >= 256 and n >= 64: - _add( - (256, 64), - 128, - swap_ab=False, - prefetch_candidates=(True,), - ) + _add((256, 64), 128, swap_ab=False) return tactics def forward( @@ -442,6 +462,9 @@ def _sm120_nvfp4_svdquant_unfused_runner(enable_pdl: bool): workspace_buffers: dict[torch.device, torch.Tensor] = {} class Sm120Nvfp4SvdquantUnfusedRunner(TunableRunner): + def get_cache_key_extras(self, inputs: List[torch.Tensor]) -> tuple: + return (enable_pdl,) + def _fp4_inputs(self, inputs: List[torch.Tensor]) -> list: a, b, a_sf, b_sf, alpha, _, _, _, out = inputs m, k_packed = a.shape @@ -536,6 +559,9 @@ def _nvfp4_svdquant_gemm_runner(enable_pdl: bool): module = get_nvfp4_svdquant_module() class Nvfp4SvdquantGemmRunner(TunableRunner): + def get_cache_key_extras(self, inputs: List[torch.Tensor]) -> tuple: + return (enable_pdl,) + def get_valid_tactics( self, inputs: List[torch.Tensor], @@ -697,8 +723,8 @@ def _check_mm_nvfp4_svdquant_problem( ) if not a_sf.is_contiguous() or not b_sf.is_contiguous(): raise ValueError("a_sf and b_sf must be contiguous") - if alpha.dtype != torch.float32 or alpha.numel() != 1: - raise ValueError("alpha must be a float32 device scalar") + if alpha.dtype != torch.float32 or alpha.numel() < 1: + raise ValueError("alpha must be a non-empty float32 device tensor") if bias is not None and (bias.shape != (n,) or bias.dtype != torch.bfloat16): raise ValueError(f"bias must have shape ({n},) and dtype bf16") if out is not None: @@ -740,8 +766,9 @@ def mm_nvfp4_svdquant( On SM100/SM103, CUTLASS fuses the block-scaled NVFP4 residual GEMM with the rank-r BF16 LoRA-up correction and optional bias. On SM120/SM121, ``"cute-dsl"`` fuses the correction and bias into the b12x CuTe DSL kernel's FP32 accumulator epilogue, while - ``"cute-dsl-unfused"`` retains the compositional implementation as a differential - oracle and fallback. The LoRA rank ``r`` is inferred from the ``d``/``l1`` shapes and must + ``"cute-dsl-unfused"`` retains the compositional implementation as a + differential oracle and optional autotuning candidate. The LoRA rank ``r`` is inferred + from the ``d``/``l1`` shapes and must be a positive multiple of 32 (ranks 32-128 are validated). ``1/alpha`` must be folded into ``l1`` by the caller (``l1 = svdquant_lora_b / alpha``), so both backends yield the correction at its original scale. @@ -761,8 +788,9 @@ def mm_nvfp4_svdquant( b_sf: torch.Tensor Weight block scales, same layout as ``a_sf`` with ``n`` rows. alpha: torch.Tensor - Per-tensor residual dequantization scale, float32 device scalar with - exactly one element (``numel == 1``). + Per-tensor residual dequantization scale in a non-empty float32 device + tensor. For compatibility with pooled scalar buffers, only the first + element is consumed; backend runners receive a one-element view. d: torch.Tensor LoRA-down output ``x_hat @ L2ᵀ``, shape ``(m, r)`` bf16, contiguous and 16-byte aligned (TMA). Compute it as ``x @ (pre_quant_scale[:, None] * L2ᵀ)`` in bf16. @@ -777,8 +805,9 @@ def mm_nvfp4_svdquant( ``"cutlass"`` selects the fused SM100/SM103 implementation; ``"cute-dsl"`` selects the fused SM120/SM121 implementation; ``"cute-dsl-unfused"`` selects its compositional reference path; - ``"auto"`` (default) selects by compute capability and, on SM120/SM121, - autotunes across both the fused and unfused implementations. + ``"auto"`` (default) selects by compute capability. On SM120/SM121, + fused and unfused are compared only while autotuning is enabled; + otherwise the fused-first runner is selected. enable_pdl: Optional[bool] Whether to launch with Programmatic Dependent Launch. Defaults to the device default. @@ -791,6 +820,9 @@ def mm_nvfp4_svdquant( enable_pdl = device_support_pdl(a.device) if out is None: out = torch.empty(a.shape[0], b.shape[0], dtype=torch.bfloat16, device=a.device) + # Preserve the historical public numel>=1 contract while specializing all + # backend kernels and autotune keys on one scalar device element. + alpha_scalar = alpha.reshape(-1)[:1] tune_sm120_implementations = False if backend == "auto": @@ -798,7 +830,7 @@ def mm_nvfp4_svdquant( tune_sm120_implementations = backend == "cute-dsl" if backend == "cute-dsl": - inputs = [a, b, a_sf, b_sf, alpha, d, l1, bias, out] + inputs = [a, b, a_sf, b_sf, alpha_scalar, d, l1, bias, out] runners = [_sm120_nvfp4_svdquant_runner(enable_pdl)] custom_op = "nvfp4_svdquant_gemm_sm120" if tune_sm120_implementations: @@ -824,7 +856,7 @@ def mm_nvfp4_svdquant( b, a_sf, b_sf, - alpha, + alpha_scalar, d, l1, bias, @@ -838,7 +870,7 @@ def mm_nvfp4_svdquant( tuner = AutoTuner.get() runners = [_nvfp4_svdquant_gemm_runner(enable_pdl)] - inputs = [a, b, a_sf, b_sf, alpha, d, l1, bias, out, workspace_buffer] + inputs = [a, b, a_sf, b_sf, alpha_scalar, d, l1, bias, out, workspace_buffer] runner, tactic = tuner.choose_one( "nvfp4_svdquant_gemm", runners, @@ -921,6 +953,14 @@ def nvfp4_quantize_smooth( enable_pdl = device_support_pdl(x.device) if backend == "auto": backend = nvfp4_quantize_smooth.suitable_auto_backends[0] + # Both backends consume vectorized BF16 operands. A contiguous storage-offset + # view can still be misaligned, so materialize only the exceptional case. + x = x.contiguous() + pre_quant_scale = pre_quant_scale.reshape(x.shape[1]).contiguous() + if x.data_ptr() % 16 != 0: + x = x.clone() + if pre_quant_scale.data_ptr() % 16 != 0: + pre_quant_scale = pre_quant_scale.clone() if backend == "cute-dsl": from ..quantization.kernels.nvfp4_quantize import ( nvfp4_quantize_smooth_cute_dsl, diff --git a/flashinfer/quantization/kernels/nvfp4_quantize.py b/flashinfer/quantization/kernels/nvfp4_quantize.py index c99edad99aa..4d155352af7 100644 --- a/flashinfer/quantization/kernels/nvfp4_quantize.py +++ b/flashinfer/quantization/kernels/nvfp4_quantize.py @@ -1889,6 +1889,19 @@ def nvfp4_quantize_cute_dsl( f"K ({k}) must be divisible by NVFP4_SF_VEC_SIZE={NVFP4_SF_VEC_SIZE}" ) + # Return explicit empty shapes before compiling or launching. In + # particular, K == 0 would make _compute_optimal_threads divide by zero. + if m == 0 or k == 0: + num_sf_blocks_per_row = k // NVFP4_SF_VEC_SIZE + if sf_layout == SF_LAYOUT_LINEAR: + padded_sf_cols = num_sf_blocks_per_row + else: + padded_sf_cols = ((num_sf_blocks_per_row + 3) // 4) * 4 + return ( + torch.empty((m, k // 2), dtype=torch.uint8, device=input.device), + torch.empty((m, padded_sf_cols), dtype=torch.uint8, device=input.device), + ) + input = input.contiguous() _torch_to_dtype_key = { @@ -2046,9 +2059,7 @@ def nvfp4_quantize_cute_dsl( padded_m, num_blocks, global_scale_arg, - input[0] - if m > 0 - else torch.empty(k, dtype=input.dtype, device=input.device), + input[0], ) # Reshape using padded_sf_cols: for swizzled layouts the buffer includes @@ -2081,6 +2092,12 @@ def nvfp4_quantize_smooth_cute_dsl( input = input.contiguous() pre_quant_scale = pre_quant_scale.reshape(k).contiguous() + # Vectorized BF16 loads require physical alignment; contiguous + # storage-offset views can still have a misaligned data pointer. + if input.data_ptr() % 16 != 0: + input = input.clone() + if pre_quant_scale.data_ptr() % 16 != 0: + pre_quant_scale = pre_quant_scale.clone() global_scale_arg = global_scale.float().reshape(1).contiguous().to(input.device) enable_pdl = device_support_pdl(input.device) if enable_pdl is not False else False @@ -2089,6 +2106,13 @@ def nvfp4_quantize_smooth_cute_dsl( padded_sf_cols = _round_up(num_sf_blocks_per_row, 4) scale_output_size = padded_m * padded_sf_cols + fp4_output = torch.empty(m, k // 2, dtype=torch.uint8, device=input.device) + scale_output = torch.empty( + scale_output_size, dtype=torch.uint8, device=input.device + ) + if m == 0 or k == 0: + return fp4_output, scale_output.reshape(m, padded_sf_cols) + kernel_fn, rows_per_block = _get_compiled_kernel_nvfp4( "bfloat16", k, @@ -2103,10 +2127,6 @@ def nvfp4_quantize_smooth_cute_dsl( (padded_m + rows_per_block - 1) // rows_per_block, get_num_sm(input.device) * _BLOCKS_PER_SM, ) - fp4_output = torch.empty(m, k // 2, dtype=torch.uint8, device=input.device) - scale_output = torch.empty( - scale_output_size, dtype=torch.uint8, device=input.device - ) kernel_fn( input, fp4_output, diff --git a/flashinfer/trace/templates/gemm.py b/flashinfer/trace/templates/gemm.py index f552dbabc5a..46a30c4fca8 100644 --- a/flashinfer/trace/templates/gemm.py +++ b/flashinfer/trace/templates/gemm.py @@ -2327,6 +2327,7 @@ def _nvfp4_quantize_smooth_init( def _svdquant_linear_init( *, M: int, + SF_B: int = 0, N: int = 3072, K: int = 3072, device: str = "cuda", @@ -2335,6 +2336,7 @@ def _svdquant_linear_init( """Build inputs for ``flashinfer.svdquant_linear`` (full SVDQuant linear chain).""" from flashinfer import nvfp4_quantize_smooth # noqa: PLC0415 + del SF_B # derived axis torch.manual_seed(seed) rank = 32 x = torch.randn(M, K, dtype=torch.bfloat16, device=device) @@ -2380,7 +2382,9 @@ def _svdquant_linear_init( "N": Const(), "K": Const(), "K_packed": Const(description="K / 2 (two e2m1 values per byte)."), - "SF_B": Const(description="128x4-swizzled weight scale buffer size."), + "SF_B": Var( + description="128x4-swizzled weight scale buffer size derived from N and K." + ), "rank": Const(description="LoRA rank, a positive multiple of 32."), }, inputs={ @@ -2424,6 +2428,9 @@ def _svdquant_linear_init( outputs={ "out": Tensor(["M", "N"], dtype="bfloat16"), }, + constraints=[ + "SF_B == ((N + 127) // 128) * 128 * (((K // 16) + 3) // 4) * 4", + ], tags=["quantization:fp4"], init=_svdquant_linear_init, ) diff --git a/tests/gemm/test_nvfp4_svdquant_gemm.py b/tests/gemm/test_nvfp4_svdquant_gemm.py index 1a3fba15b3a..31e3e936dca 100644 --- a/tests/gemm/test_nvfp4_svdquant_gemm.py +++ b/tests/gemm/test_nvfp4_svdquant_gemm.py @@ -44,10 +44,8 @@ def test_nvfp4_svdquant_backend_arch_support(): assert not mm_nvfp4_svdquant.is_backend_supported("cute-dsl-unfused", 100) -def test_sm120_svdquant_can_implement_rejects_ragged_rank(): - _skip_unless_sm120() - import cutlass - +def test_sm120_svdquant_kernel_iket_flag_defaults_off(): + pytest.importorskip("cutlass") from flashinfer.gemm.kernels.dense_blockscaled_gemm_sm120_b12x import ( Sm120B12xBlockScaledDenseGemmKernel, ) @@ -57,6 +55,13 @@ def test_sm120_svdquant_can_implement_rejects_ragged_rank(): 16, (64, 64), (1, 1), enable_iket=True ).enable_iket + +def test_sm120_svdquant_can_implement_rejects_ragged_rank(): + cutlass = pytest.importorskip("cutlass") + from flashinfer.gemm.kernels.dense_blockscaled_gemm_sm120_b12x import ( + Sm120B12xBlockScaledDenseGemmKernel, + ) + common_args = ( cutlass.Float4E2M1FN, cutlass.Float8E4M3FN, @@ -111,6 +116,20 @@ def _sqnr_db(ref: torch.Tensor, got: torch.Tensor) -> float: return float(10 * torch.log10((ref.float() ** 2).mean() / noise)) +def _assert_sm120_accuracy(ref: torch.Tensor, got: torch.Tensor) -> None: + """Guard aggregate quality and localized spikes with measured margin. + + The review sweep measured a 53.27 dB SQNR floor and a 0.347% maximum + error/reference-peak ceiling; historical tile coverage reached 48.99 dB. + """ + ref_f32 = ref.float() + got_f32 = got.float() + assert _sqnr_db(ref_f32, got_f32) > 45.0 + peak = ref_f32.abs().amax().clamp_min(torch.finfo(torch.float32).tiny) + normalized_max_error = (ref_f32 - got_f32).abs().amax() / peak + assert normalized_max_error < 0.01 + + def _nvfp4_quantize_128x4(t: torch.Tensor, backend="cuda"): """Stock NVFP4 quantization (ue4m3 block scales, 128x4 swizzled layout). @@ -362,6 +381,36 @@ def test_mm_nvfp4_svdquant_rejects_bad_rank(): ) +def test_mm_nvfp4_svdquant_sm100_pooled_alpha_uses_first_element(): + _skip_unless_sm100() + torch.manual_seed(0) + p = _make_gemm_problem(129, 3072, 3072, rank=32) + expected = mm_nvfp4_svdquant( + p["xq"], + p["wq"], + p["x_sf_flat"], + p["w_sf_flat"], + p["alpha"], + p["d"], + p["l1_scaled"], + bias=p["bias"], + backend="cutlass", + ) + pooled_alpha = torch.cat([p["alpha"], torch.tensor([2.0, 3.0, 4.0], device="cuda")]) + actual = mm_nvfp4_svdquant( + p["xq"], + p["wq"], + p["x_sf_flat"], + p["w_sf_flat"], + pooled_alpha, + p["d"], + p["l1_scaled"], + bias=p["bias"], + backend="cutlass", + ) + assert torch.equal(actual, expected) + + @pytest.mark.parametrize("m", [129, 6912]) @pytest.mark.parametrize("rank", [32, 96]) def test_mm_nvfp4_svdquant_autotuned(m, rank): @@ -462,10 +511,10 @@ def test_svdquant_linear_matches_reference(use_bias, rank): assert _sqnr_db(ref, out.float()) > 40.0 -def test_nvfp4_quantize_smooth_sm120_cute_dsl(): +@pytest.mark.parametrize("m,k", [(129, 256), (129, 12288), (70, 3088)]) +def test_nvfp4_quantize_smooth_sm120_cute_dsl(m, k): _skip_unless_sm120() torch.manual_seed(0) - m, k = 129, 256 x = torch.randn(m, k, dtype=torch.bfloat16, device="cuda") pqs = ( (1.0 + 0.3 * torch.randn(k, dtype=torch.bfloat16, device="cuda")) @@ -497,6 +546,50 @@ def test_nvfp4_quantize_smooth_sm120_cute_dsl(): assert torch.equal(sf_auto, sf) +def test_nvfp4_quantize_smooth_sm120_misaligned_inputs(): + _skip_unless_sm120() + torch.manual_seed(0) + m, k = 70, 256 + x_storage = torch.randn(m * k + 1, dtype=torch.bfloat16, device="cuda") + x = x_storage[1:].view(m, k) + pqs_storage = torch.randn(k + 1, dtype=torch.bfloat16, device="cuda").abs() + pqs = pqs_storage[1:] + assert x.is_contiguous() and x.data_ptr() % 16 != 0 + assert pqs.is_contiguous() and pqs.data_ptr() % 16 != 0 + smoothed = (x * pqs).to(torch.bfloat16) + global_sf = ((448.0 * 6.0) / smoothed.float().abs().max()).reshape(1) + expected_q, expected_sf = nvfp4_quantize( + smoothed, + global_sf, + sfLayout=SfLayout.layout_128x4, + do_shuffle=False, + backend="cute-dsl", + ) + actual_q, actual_sf = nvfp4_quantize_smooth(x, pqs, global_sf, backend="cute-dsl") + assert torch.equal(actual_q, expected_q.view(torch.uint8)) + assert torch.equal(actual_sf, expected_sf.view(torch.uint8).reshape(-1)) + + +@pytest.mark.parametrize("m,k", [(0, 256), (4, 0)]) +def test_nvfp4_quantizers_sm120_empty_inputs(m, k): + _skip_unless_sm120() + from flashinfer.quantization.kernels.nvfp4_quantize import ( + nvfp4_quantize_cute_dsl, + ) + + x = torch.empty(m, k, dtype=torch.bfloat16, device="cuda") + pqs = torch.ones(k, dtype=torch.bfloat16, device="cuda") + global_sf = torch.ones(1, dtype=torch.float32, device="cuda") + xq, sf = nvfp4_quantize_smooth(x, pqs, global_sf, backend="cute-dsl") + plain_xq, plain_sf = nvfp4_quantize_cute_dsl(x, global_sf) + + assert xq.shape == (m, k // 2) and xq.dtype == torch.uint8 + assert sf.shape == (0,) and sf.dtype == torch.uint8 + assert plain_xq.shape == (m, k // 2) and plain_xq.dtype == torch.uint8 + assert plain_sf.shape == (m, ((k // 16 + 3) // 4) * 4) + assert plain_sf.dtype == torch.uint8 + + @pytest.mark.parametrize("use_bias", [False, True]) def test_mm_nvfp4_svdquant_sm120_fused(use_bias): _skip_unless_sm120() @@ -541,18 +634,25 @@ def test_mm_nvfp4_svdquant_sm120_fused(use_bias): # The fused path accumulates the BF16 rank correction in FP32 before its # single BF16 store, whereas the oracle rounds the residual and correction # in separate launches. Compare numerically, not bitwise. - assert _sqnr_db(expected.float(), out.float()) > 35.0 - assert _sqnr_db(expected.float(), out_auto.float()) > 35.0 + _assert_sm120_accuracy(expected, out) + _assert_sm120_accuracy(expected, out_auto) fp32_ref = p["ref_bias"] if use_bias else p["ref"] - assert _sqnr_db(fp32_ref, out.float()) > 35.0 + _assert_sm120_accuracy(fp32_ref, out) -@pytest.mark.parametrize("tile_k,rank", [(64, 32), (128, 32), (256, 64)]) -def test_mm_nvfp4_svdquant_sm120_large_m_tile(tile_k, rank): +@pytest.mark.parametrize( + "tactic,rank", + [ + (((64, 64), 128, False), 32), + (((128, 128), 128, False), 64), + (((256, 64), 128, False), 32), + ], +) +def test_mm_nvfp4_svdquant_sm120_large_m_tactic(tactic, rank): _skip_unless_sm120() from flashinfer.gemm.gemm_svdquant import _mm_nvfp4_svdquant_sm120_fused - torch.manual_seed(tile_k) + torch.manual_seed(tactic[0][0]) p = _make_gemm_problem( 257, 128, @@ -573,9 +673,9 @@ def test_mm_nvfp4_svdquant_sm120_large_m_tile(tile_k, rank): p["bias"], out, device_support_pdl(torch.device("cuda")), - tactic=((256, 64), tile_k, False, False), + tactic=tactic, ) - assert _sqnr_db(p["ref_bias"], out.float()) > 35.0 + _assert_sm120_accuracy(p["ref_bias"], out) @pytest.mark.parametrize("backend", ["cute-dsl", "auto"]) @@ -603,7 +703,7 @@ def test_mm_nvfp4_svdquant_sm120_autotuned_replay(backend): bias=p["bias"], backend=backend, ) - assert _sqnr_db(p["ref_bias"], out.float()) > 35.0 + _assert_sm120_accuracy(p["ref_bias"], out) # Replay outside the tuning context must reuse the selected tactic. out_replay = mm_nvfp4_svdquant( @@ -669,7 +769,7 @@ def test_mm_nvfp4_svdquant_sm120_fused_rank_chunks(rank): bias=p["bias"], backend="cute-dsl", ) - assert _sqnr_db(p["ref_bias"], out.float()) > 35.0 + _assert_sm120_accuracy(p["ref_bias"], out) @pytest.mark.parametrize( @@ -701,11 +801,12 @@ def test_mm_nvfp4_svdquant_sm120_fused_boundary_plans(m, n, k, rank): bias=p["bias"], backend="cute-dsl", ) - assert _sqnr_db(p["ref_bias"], out.float()) > 35.0 + _assert_sm120_accuracy(p["ref_bias"], out) -@pytest.mark.parametrize("alpha_shape", [(), (1, 1)]) -def test_mm_nvfp4_svdquant_sm120_normalizes_alpha_shape(alpha_shape): +@pytest.mark.parametrize("alpha_shape", [(), (1, 1), (4,)]) +@pytest.mark.parametrize("backend", ["cute-dsl", "cute-dsl-unfused"]) +def test_mm_nvfp4_svdquant_sm120_normalizes_alpha_shape(alpha_shape, backend): _skip_unless_sm120() torch.manual_seed(0) p = _make_gemm_problem( @@ -725,18 +826,23 @@ def test_mm_nvfp4_svdquant_sm120_normalizes_alpha_shape(alpha_shape): p["d"], p["l1_scaled"], bias=p["bias"], - backend="cute-dsl", + backend=backend, + ) + alpha = ( + p["alpha"].reshape(alpha_shape) + if alpha_shape != (4,) + else torch.cat([p["alpha"], torch.tensor([2.0, 3.0, 4.0], device="cuda")]) ) out = mm_nvfp4_svdquant( p["xq"], p["wq"], p["x_sf_flat"], p["w_sf_flat"], - p["alpha"].reshape(alpha_shape), + alpha, p["d"], p["l1_scaled"], bias=p["bias"], - backend="cute-dsl", + backend=backend, ) assert torch.equal(out, expected) @@ -852,6 +958,7 @@ def recording_torch_mm(*args, **kwargs): l1_scaled, global_sf, bias=bias, + backend="cute-dsl", ) assert torch_mm_calls == 1 @@ -876,7 +983,7 @@ def recording_torch_mm(*args, **kwargs): ref.add_(bias.float()) assert out.shape == (m, n) and out.dtype == torch.bfloat16 - assert _sqnr_db(ref, out.float()) > 35.0 + _assert_sm120_accuracy(ref, out) @pytest.mark.parametrize("rank", [32, 128]) @@ -954,5 +1061,46 @@ def run(out_tensor): assert torch.equal(out_graph, out_eager) +@pytest.mark.parametrize("backend", ["cute-dsl", "auto"]) +def test_mm_nvfp4_svdquant_sm120_cuda_graph_replay(backend): + _skip_unless_sm120() + torch.manual_seed(0) + p = _make_gemm_problem( + 129, + 256, + 256, + rank=32, + quant_backend="cute-dsl", + residual_backend="b12x", + ) + out = torch.empty(129, 256, dtype=torch.bfloat16, device="cuda") + + def run(): + mm_nvfp4_svdquant( + p["xq"], + p["wq"], + p["x_sf_flat"], + p["w_sf_flat"], + p["alpha"], + p["d"], + p["l1_scaled"], + bias=p["bias"], + out=out, + backend=backend, + ) + + # Compile and tune before capture; capture must only replay the cached path. + with autotune(True): + run() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + out.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + _assert_sm120_accuracy(p["ref_bias"], out) + + if __name__ == "__main__": pytest.main([__file__]) diff --git a/tests/jit/test_cute_dsl_cache.py b/tests/jit/test_cute_dsl_cache.py index 3bf5712975b..1838263d719 100644 --- a/tests/jit/test_cute_dsl_cache.py +++ b/tests/jit/test_cute_dsl_cache.py @@ -39,6 +39,13 @@ pytest.importorskip("cutlass") +from flashinfer.gemm.gemm_svdquant import ( # noqa: E402 + _sm120_nvfp4_svdquant_runner, + _sm120_nvfp4_svdquant_unfused_runner, + _sm120_svdquant_kernel_name, + _svdquant_kernel_source_files, +) + from flashinfer.quantization.kernels.nvfp4_quantize import ( # noqa: E402 SF_LAYOUT_8x4, SF_LAYOUT_128x4, @@ -62,7 +69,7 @@ # ever introduced. NVFP4_NON_CODEGEN_PARAMS: set = set() -# A baseline argument set and, for each argument, a distinct alternative. +# Baseline argument sets shared by the individual one-at-a-time perturbation tests. NVFP4_NAME_BASELINE = { "variant": "swizzled", "dtype_key": "bfloat16", @@ -73,17 +80,18 @@ "silu_and_mul": False, "nvfp4_4over6_config": None, "global_scale_is_tensor": True, + "smooth_quant": False, } -NVFP4_NAME_PERTURBED = { - "variant": "linear", - "dtype_key": "float16", - "K": 2048, - "sf_layout": SF_LAYOUT_8x4, + +SVDQUANT_NAME_BASELINE = { + "rank": 32, + "with_bias": False, + "mma_tiler_mn": (64, 64), + "tile_k": 128, + "swap_ab": False, + "max_active_clusters": 1, "enable_pdl": False, - "disable_fp4_quant_fast_math": True, - "silu_and_mul": True, - "nvfp4_4over6_config": NVFP44Over6Config(), - "global_scale_is_tensor": False, + "enable_iket": False, } @@ -105,21 +113,103 @@ def test_nvfp4_kernel_name_signature_covers_codegen_params(getter): ) -@pytest.mark.parametrize("param", sorted(NVFP4_NAME_BASELINE)) -def test_nvfp4_kernel_name_varies_with_every_argument(param): +@pytest.mark.parametrize( + "param,alternate", + [ + pytest.param("variant", "linear", id="variant"), + pytest.param("dtype_key", "float16", id="dtype_key"), + pytest.param("K", 2048, id="K"), + pytest.param("sf_layout", SF_LAYOUT_8x4, id="sf_layout"), + pytest.param("enable_pdl", False, id="enable_pdl"), + pytest.param( + "disable_fp4_quant_fast_math", + True, + id="disable_fp4_quant_fast_math", + ), + pytest.param("silu_and_mul", True, id="silu_and_mul"), + pytest.param( + "nvfp4_4over6_config", + NVFP44Over6Config(), + id="nvfp4_4over6_config", + ), + pytest.param("global_scale_is_tensor", False, id="global_scale_is_tensor"), + pytest.param("smooth_quant", True, id="smooth_quant"), + ], +) +def test_nvfp4_kernel_name_varies_with_every_argument(param, alternate): """Changing any single argument must change the kernel name. Catches arguments that the name function accepts but ignores. """ baseline_name = _nvfp4_kernel_name(**NVFP4_NAME_BASELINE) kwargs = dict(NVFP4_NAME_BASELINE) - kwargs[param] = NVFP4_NAME_PERTURBED[param] + kwargs[param] = alternate assert _nvfp4_kernel_name(**kwargs) != baseline_name, ( f"_nvfp4_kernel_name ignores argument {param!r}: two different " "kernel specializations would collide on one cache artifact." ) +@pytest.mark.parametrize( + "param,alternate", + [ + pytest.param("rank", 64, id="rank"), + pytest.param("with_bias", True, id="with_bias"), + pytest.param("mma_tiler_mn", (128, 64), id="mma_tiler_mn"), + pytest.param("tile_k", 256, id="tile_k"), + pytest.param("swap_ab", True, id="swap_ab"), + pytest.param("max_active_clusters", 2, id="max_active_clusters"), + pytest.param("enable_pdl", True, id="enable_pdl"), + pytest.param("enable_iket", True, id="enable_iket"), + ], +) +def test_sm120_svdquant_kernel_name_varies_with_every_argument(param, alternate): + baseline_name = _sm120_svdquant_kernel_name(**SVDQUANT_NAME_BASELINE) + kwargs = dict(SVDQUANT_NAME_BASELINE) + kwargs[param] = alternate + assert _sm120_svdquant_kernel_name(**kwargs) != baseline_name + + +def test_sm120_svdquant_kernel_name_is_symbol_safe(): + name = _sm120_svdquant_kernel_name( + **{**SVDQUANT_NAME_BASELINE, "mma_tiler_mn": (128, 64), "enable_iket": True} + ) + assert re.fullmatch(r"[A-Za-z0-9_]+", name), name + + +def test_sm120_svdquant_source_fingerprint_covers_layout_helpers(): + from flashinfer.cute_dsl import utils as cute_dsl_utils + from flashinfer.gemm.kernels import dense_blockscaled_gemm_sm120_b12x + + sources = _svdquant_kernel_source_files() + assert cute_dsl_utils.__file__ in sources + assert dense_blockscaled_gemm_sm120_b12x.__file__ in sources + + +@pytest.mark.parametrize( + "runner_factory", + [ + _sm120_nvfp4_svdquant_runner, + _sm120_nvfp4_svdquant_unfused_runner, + ], +) +def test_svdquant_autotune_cache_distinguishes_pdl(runner_factory): + disabled = runner_factory(False) + enabled = runner_factory(True) + assert disabled.get_cache_key_extras([]) == (False,) + assert enabled.get_cache_key_extras([]) == (True,) + + +def test_sm100_svdquant_autotune_cache_distinguishes_pdl(monkeypatch): + from flashinfer.gemm import gemm_svdquant + + monkeypatch.setattr(gemm_svdquant, "get_nvfp4_svdquant_module", object) + disabled = gemm_svdquant._nvfp4_svdquant_gemm_runner(False) + enabled = gemm_svdquant._nvfp4_svdquant_gemm_runner(True) + assert disabled.get_cache_key_extras([]) == (False,) + assert enabled.get_cache_key_extras([]) == (True,) + + @pytest.mark.parametrize( "config", [ diff --git a/tests/trace/example.py b/tests/trace/example.py index d6f59d56985..fea787c9f9a 100644 --- a/tests/trace/example.py +++ b/tests/trace/example.py @@ -25,7 +25,6 @@ packed_kda_decode_h12_d128.json fused_kda_decode_h12_d128.json gemm_bf16_N256_K7168.json -gemm_bf16_N32_K3072.json gemm_bf16_N4096_K4096.json gemm_fp4_N2048_K7168_block_size16.json gemm_fp8_N1536_K7168.json @@ -42,7 +41,7 @@ gqa_ragged_h32_kv8_d128.json layernorm_h768.json layernorm_quant_h768.json -linear_nvfp4_svdquant_N3072_K3072_K_packed1536_SF_B589824_rank32.json +linear_nvfp4_svdquant_N3072_K3072_K_packed1536_rank32.json merge_state_h32_d128.json merge_state_in_place_h32_d128.json merge_states_h32_d128.json diff --git a/tests/trace/fi_trace_out/gemm_bf16_N32_K3072.json b/tests/trace/fi_trace_out/gemm_bf16_N32_K3072.json deleted file mode 100644 index 441b9a210c2..00000000000 --- a/tests/trace/fi_trace_out/gemm_bf16_N32_K3072.json +++ /dev/null @@ -1,51 +0,0 @@ -{ - "name": "gemm_bf16_N32_K3072", - "description": "General matrix multiply (GEMM) C = A @ B (B is column-major [K, N]).", - "op_type": "gemm_bf16", - "tags": [ - "fi_api:flashinfer.gemm.gemm_base.mm_bf16", - "status:verified" - ], - "axes": { - "M": { - "type": "var" - }, - "N": { - "type": "const", - "value": 32 - }, - "K": { - "type": "const", - "value": 3072 - } - }, - "inputs": { - "A": { - "shape": [ - "M", - "K" - ], - "dtype": "bfloat16" - }, - "B": { - "shape": [ - "K", - "N" - ], - "dtype": "bfloat16", - "description": "Weight matrix in column-major layout (physical shape [K, N])." - } - }, - "outputs": { - "C": { - "shape": [ - "M", - "N" - ], - "dtype": "bfloat16" - } - }, - "reference": "from __future__ import annotations\nimport math\nimport torch\nimport torch.nn.functional as F\n\ndef _mm_reference(A, B):\n # B is physically [K, N] (column-major weight), so C = A @ B.\n return torch.matmul(A, B)\n", - "check": "def _gemm_check(\n reference_outputs,\n actual_outputs,\n *,\n rtol=None,\n atol=None,\n max_mismatch_pct=100.0,\n min_cos_sim=0.99,\n):\n from flashinfer.trace import default_check\n\n # Matches tests/gemm/test_mm_bf16.py, test_mm_fp8.py, test_bmm_bf16.py,\n # and test_bmm_fp8.py, which gate these kernels by cosine similarity.\n return default_check(\n reference_outputs,\n actual_outputs,\n rtol=rtol,\n atol=atol,\n max_mismatch_pct=max_mismatch_pct,\n min_cos_sim=min_cos_sim,\n )\n", - "init": "from __future__ import annotations\nimport math\nimport torch\n\n# ----- shared init helpers -----\n# Copyright (c) 2025 by FlashInfer team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Shared helpers used by ``TraceTemplate.init`` functions.\n\nThis module contains the small set of input-construction patterns that\nrecur across many templates (paged-KV cache index arrays, ragged indptr,\nRoPE pos_ids and cos/sin caches, sampling probs). Each helper is short and\ndocumented; init functions in ``templates/.py`` call into here so\nthe per-template init bodies stay focused on shape/dtype, not boilerplate.\n\nThe full source of this module is **inlined into every dumped JSON's\n``\"init\"`` field** by ``flashinfer/trace/template.py:_render_init_source``,\nso downstream consumers don't need flashinfer installed to re-run the init\nsnippets.\n\"\"\"\n\n\nfrom typing import Optional, Tuple\n\nimport torch\n\n\ndef make_paged_kv_indices(\n batch_size: int,\n num_pages_per_seq: int,\n page_size: int,\n *,\n device: str = \"cuda\",\n) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n \"\"\"Return ``(kv_indptr, kv_indices, kv_last_page_len)`` for a uniform batch.\n\n Every sequence is assigned exactly ``num_pages_per_seq`` pages, fully\n populated (last-page length == page_size).\n\n Invariants\n ----------\n - ``kv_indptr.shape == (batch_size + 1,)``, dtype int32, monotonic, [0]=0.\n - ``kv_indices == arange(0, batch_size * num_pages_per_seq)``, int32.\n - ``kv_last_page_len == full(batch_size, page_size)``, int32.\n \"\"\"\n total_pages = batch_size * num_pages_per_seq\n kv_indptr = (\n torch.arange(batch_size + 1, dtype=torch.int32, device=device)\n * num_pages_per_seq\n )\n kv_indices = torch.arange(total_pages, dtype=torch.int32, device=device)\n kv_last_page_len = torch.full(\n (batch_size,), page_size, dtype=torch.int32, device=device\n )\n return kv_indptr, kv_indices, kv_last_page_len\n\n\ndef make_ragged_indptr(\n seg_lens,\n *,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.int32,\n) -> torch.Tensor:\n \"\"\"Return cumulative-sum ``indptr`` of length ``len(seg_lens)+1``.\n\n ``seg_lens`` may be a list / tuple / 1-D tensor of segment lengths.\n \"\"\"\n if isinstance(seg_lens, torch.Tensor):\n lens = seg_lens.to(device=device, dtype=dtype)\n else:\n lens = torch.tensor(list(seg_lens), dtype=dtype, device=device)\n indptr = torch.zeros(lens.numel() + 1, dtype=dtype, device=device)\n indptr[1:] = torch.cumsum(lens, dim=0).to(dtype)\n return indptr\n\n\ndef make_uniform_qo_indptr(\n batch_size: int,\n qo_len: int,\n *,\n device: str = \"cuda\",\n) -> torch.Tensor:\n \"\"\"Return ``[0, qo_len, 2*qo_len, ..., batch_size*qo_len]`` int32.\"\"\"\n return torch.arange(batch_size + 1, dtype=torch.int32, device=device) * qo_len\n\n\ndef make_pos_ids(\n nnz: int,\n max_seq_len: Optional[int] = None,\n *,\n device: str = \"cuda\",\n) -> torch.Tensor:\n \"\"\"Return ``[0, 1, ..., nnz-1] (% max_seq_len)`` as int32 on ``device``.\n\n If ``max_seq_len`` is None, no wrapping is applied.\n \"\"\"\n pos = torch.arange(nnz, dtype=torch.int32, device=device)\n if max_seq_len is not None:\n pos = pos % max_seq_len\n return pos\n\n\ndef make_rope_cos_sin_cache(\n max_seq_len: int,\n rope_dim: int,\n *,\n base: float = 1e4,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Return concatenated ``[cos | sin]`` cache of shape ``[max_seq_len, rope_dim]``.\"\"\"\n t = torch.arange(max_seq_len, dtype=torch.float32, device=device)\n inv = 1.0 / (\n base\n ** (torch.arange(0, rope_dim, 2, dtype=torch.float32, device=device) / rope_dim)\n )\n freqs = t.unsqueeze(-1) * inv.unsqueeze(0)\n cache = torch.cat([torch.cos(freqs), torch.sin(freqs)], dim=-1)\n return cache.to(dtype)\n\n\ndef make_probs(\n batch_size: int,\n vocab_size: int,\n *,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Return a ``[batch_size, vocab_size]`` probability distribution.\n\n Uses ``softmax(randn(...))`` so each row sums to 1.0. This mirrors the\n pattern used throughout ``tests/utils/test_sampling.py``.\n \"\"\"\n return torch.softmax(\n torch.randn(batch_size, vocab_size, dtype=torch.float32, device=device),\n dim=-1,\n ).to(dtype)\n\n\ndef make_logits(\n batch_size: int,\n vocab_size: int,\n *,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Return ``randn(batch_size, vocab_size)`` logits.\"\"\"\n return torch.randn(batch_size, vocab_size, dtype=dtype, device=device)\n\n\ndef fp8_safe_randn(\n *shape: int,\n scale: float = 0.1,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.bfloat16,\n) -> torch.Tensor:\n \"\"\"``randn(*shape) * scale`` \u2014 keeps values in the FP8/FP4 representable range.\n\n Tests for fp8/fp4 paths typically multiply ``randn`` by 0.1 to avoid\n saturation when quantizing. Use this helper to mirror that convention.\n \"\"\"\n return (torch.randn(*shape, dtype=dtype, device=device) * scale).to(dtype)\n\n\ndef per_tensor_fp8_quantize(\n x: torch.Tensor,\n *,\n fp8_dtype: torch.dtype = torch.float8_e4m3fn,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Per-tensor FP8 quantization, mirroring ``tests/utils_fp8.py:to_float8``.\n\n Returns ``(x_fp8, inv_scale)`` where ``inv_scale`` is the dequant\n multiplier (``float \u2248 fp8 * inv_scale``).\n \"\"\"\n finfo = torch.finfo(fp8_dtype)\n amax = x.abs().amax().clamp(min=1e-12)\n scale = finfo.max / amax\n x_q = (x.float() * scale).clamp(min=finfo.min, max=finfo.max).to(fp8_dtype)\n return x_q, scale.float().reciprocal()\n\n\ndef fp8_block_quant_1d(\n x_bf16: torch.Tensor,\n block: int = 128,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Quantize ``[T, H]`` activations into FP8 with per-``(token, block)``\n column-block scales. Returns ``(x_fp8, scales)`` where\n ``scales`` has shape ``[T, H // block]``.\n\n Mirrors ``_fp8_block_quant_1d`` in\n ``tests/moe/test_dpsk_fused_moe_fp8.py``.\n \"\"\"\n assert x_bf16.dim() == 2\n T, H = x_bf16.shape\n assert H % block == 0\n nb = H // block\n finfo = torch.finfo(torch.float8_e4m3fn)\n max_fp8 = finfo.max\n x_f32 = x_bf16.to(torch.float32)\n x_fp8 = torch.empty((T, H), dtype=torch.float8_e4m3fn, device=x_bf16.device)\n scales = torch.empty((T, nb), dtype=torch.float32, device=x_bf16.device)\n for j in range(nb):\n sl = slice(j * block, (j + 1) * block)\n blk = x_f32[:, sl]\n amax = torch.amax(torch.abs(blk), dim=1)\n s = torch.where(amax > 0, amax / max_fp8, torch.ones_like(amax))\n x_fp8[:, sl] = (blk / s.unsqueeze(1)).to(torch.float8_e4m3fn)\n scales[:, j] = s\n return x_fp8, scales\n\n\ndef fp8_block_quant_2d(\n w_bf16: torch.Tensor,\n block: int = 128,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Quantize weights ``[..., R, C]`` with 2-D ``block \u00d7 block`` scales.\n\n Returns ``(w_fp8, scales)`` where ``scales`` has shape\n ``[..., R // block, C // block]``. Mirrors ``_fp8_block_quant_2d`` in\n ``tests/moe/test_dpsk_fused_moe_fp8.py``.\n \"\"\"\n assert w_bf16.dim() >= 2\n *prefix, R, C = w_bf16.shape\n assert R % block == 0 and C % block == 0\n nb_r, nb_c = R // block, C // block\n finfo = torch.finfo(torch.float8_e4m3fn)\n max_fp8 = finfo.max\n w_f32 = w_bf16.to(torch.float32).contiguous()\n prefix_ndim = len(prefix)\n reshaped = w_f32.reshape(*prefix, nb_r, block, nb_c, block)\n permute_dims = tuple(range(prefix_ndim)) + (\n prefix_ndim,\n prefix_ndim + 2,\n prefix_ndim + 1,\n prefix_ndim + 3,\n )\n blocks = reshaped.permute(permute_dims).contiguous()\n amax = torch.amax(torch.abs(blocks), dim=(-1, -2))\n scales = torch.where(\n amax > 0, amax / max_fp8, torch.ones_like(amax, dtype=torch.float32)\n )\n q_blocks = (blocks / scales.unsqueeze(-1).unsqueeze(-1)).to(torch.float8_e4m3fn)\n inv_permute = [0] * (prefix_ndim + 4)\n for i, p in enumerate(permute_dims):\n inv_permute[p] = i\n w_fp8 = q_blocks.permute(*inv_permute).reshape(*prefix, R, C).contiguous()\n return w_fp8, scales\n\n\n__all__ = [\n \"make_paged_kv_indices\",\n \"make_ragged_indptr\",\n \"make_uniform_qo_indptr\",\n \"make_pos_ids\",\n \"make_rope_cos_sin_cache\",\n \"make_probs\",\n \"make_logits\",\n \"fp8_safe_randn\",\n \"per_tensor_fp8_quantize\",\n \"fp8_block_quant_1d\",\n \"fp8_block_quant_2d\",\n]\n\n# ----- init -----\ndef _mm_bf16_init(\n *,\n M: int,\n N: int = 4096,\n K: int = 4096,\n device: str = \"cuda\",\n seed: int = 0,\n):\n \"\"\"Build inputs for ``flashinfer.mm_bf16``.\n\n ``B`` is constructed as ``randn(N, K).T`` to get column-major [K, N]\n matching the example call.\n \"\"\"\n torch.manual_seed(seed)\n a = torch.randn(M, K, dtype=torch.bfloat16, device=device)\n b = torch.randn(N, K, dtype=torch.bfloat16, device=device).T # [K, N] col-major\n return {\"a\": a, \"b\": b}\n" -} diff --git a/tests/trace/fi_trace_out/linear_nvfp4_svdquant_N3072_K3072_K_packed1536_SF_B589824_rank32.json b/tests/trace/fi_trace_out/linear_nvfp4_svdquant_N3072_K3072_K_packed1536_rank32.json similarity index 86% rename from tests/trace/fi_trace_out/linear_nvfp4_svdquant_N3072_K3072_K_packed1536_SF_B589824_rank32.json rename to tests/trace/fi_trace_out/linear_nvfp4_svdquant_N3072_K3072_K_packed1536_rank32.json index ef95ba2d810..6d2065c2906 100644 --- a/tests/trace/fi_trace_out/linear_nvfp4_svdquant_N3072_K3072_K_packed1536_SF_B589824_rank32.json +++ b/tests/trace/fi_trace_out/linear_nvfp4_svdquant_N3072_K3072_K_packed1536_rank32.json @@ -1,5 +1,5 @@ { - "name": "linear_nvfp4_svdquant_N3072_K3072_K_packed1536_SF_B589824_rank32", + "name": "linear_nvfp4_svdquant_N3072_K3072_K_packed1536_rank32", "description": "Full SVDQuant linear: y = (x * pre_quant_scale) @ (R + L1 @ L2)\u1d40 where R is the NVFP4-quantized residual weight \u2014 smooth-quantize, BF16 rank-r down-projection, and the architecture-selected NVFP4 residual + LoRA-up GEMM.", "op_type": "linear_nvfp4_svdquant", "tags": [ @@ -24,9 +24,8 @@ "description": "K / 2 (two e2m1 values per byte)." }, "SF_B": { - "type": "const", - "value": 589824, - "description": "128x4-swizzled weight scale buffer size." + "type": "var", + "description": "128x4-swizzled weight scale buffer size derived from N and K." }, "rank": { "type": "const", @@ -34,6 +33,9 @@ "description": "LoRA rank, a positive multiple of 32." } }, + "constraints": [ + "SF_B == ((N + 127) // 128) * 128 * (((K // 16) + 3) // 4) * 4" + ], "inputs": { "x": { "shape": [ @@ -106,5 +108,5 @@ } }, "check": "def standard_check(\n reference_outputs: Any,\n actual_outputs: Any,\n *,\n rtol: Optional[float] = None,\n atol: Optional[float] = None,\n max_mismatch_pct: float = 0.0,\n min_cos_sim: Optional[float] = 1.0 - 1e-3,\n) -> bool:\n \"\"\"Default trace correctness check used when a template does not override it.\"\"\"\n from flashinfer.trace import default_check\n\n return default_check(\n reference_outputs,\n actual_outputs,\n rtol=rtol,\n atol=atol,\n max_mismatch_pct=max_mismatch_pct,\n min_cos_sim=min_cos_sim,\n )\n", - "init": "from __future__ import annotations\nimport math\nimport torch\n\n# ----- shared init helpers -----\n# Copyright (c) 2025 by FlashInfer team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Shared helpers used by ``TraceTemplate.init`` functions.\n\nThis module contains the small set of input-construction patterns that\nrecur across many templates (paged-KV cache index arrays, ragged indptr,\nRoPE pos_ids and cos/sin caches, sampling probs). Each helper is short and\ndocumented; init functions in ``templates/.py`` call into here so\nthe per-template init bodies stay focused on shape/dtype, not boilerplate.\n\nThe full source of this module is **inlined into every dumped JSON's\n``\"init\"`` field** by ``flashinfer/trace/template.py:_render_init_source``,\nso downstream consumers don't need flashinfer installed to re-run the init\nsnippets.\n\"\"\"\n\n\nfrom typing import Optional, Tuple\n\nimport torch\n\n\ndef make_paged_kv_indices(\n batch_size: int,\n num_pages_per_seq: int,\n page_size: int,\n *,\n device: str = \"cuda\",\n) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n \"\"\"Return ``(kv_indptr, kv_indices, kv_last_page_len)`` for a uniform batch.\n\n Every sequence is assigned exactly ``num_pages_per_seq`` pages, fully\n populated (last-page length == page_size).\n\n Invariants\n ----------\n - ``kv_indptr.shape == (batch_size + 1,)``, dtype int32, monotonic, [0]=0.\n - ``kv_indices == arange(0, batch_size * num_pages_per_seq)``, int32.\n - ``kv_last_page_len == full(batch_size, page_size)``, int32.\n \"\"\"\n total_pages = batch_size * num_pages_per_seq\n kv_indptr = (\n torch.arange(batch_size + 1, dtype=torch.int32, device=device)\n * num_pages_per_seq\n )\n kv_indices = torch.arange(total_pages, dtype=torch.int32, device=device)\n kv_last_page_len = torch.full(\n (batch_size,), page_size, dtype=torch.int32, device=device\n )\n return kv_indptr, kv_indices, kv_last_page_len\n\n\ndef make_ragged_indptr(\n seg_lens,\n *,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.int32,\n) -> torch.Tensor:\n \"\"\"Return cumulative-sum ``indptr`` of length ``len(seg_lens)+1``.\n\n ``seg_lens`` may be a list / tuple / 1-D tensor of segment lengths.\n \"\"\"\n if isinstance(seg_lens, torch.Tensor):\n lens = seg_lens.to(device=device, dtype=dtype)\n else:\n lens = torch.tensor(list(seg_lens), dtype=dtype, device=device)\n indptr = torch.zeros(lens.numel() + 1, dtype=dtype, device=device)\n indptr[1:] = torch.cumsum(lens, dim=0).to(dtype)\n return indptr\n\n\ndef make_uniform_qo_indptr(\n batch_size: int,\n qo_len: int,\n *,\n device: str = \"cuda\",\n) -> torch.Tensor:\n \"\"\"Return ``[0, qo_len, 2*qo_len, ..., batch_size*qo_len]`` int32.\"\"\"\n return torch.arange(batch_size + 1, dtype=torch.int32, device=device) * qo_len\n\n\ndef make_pos_ids(\n nnz: int,\n max_seq_len: Optional[int] = None,\n *,\n device: str = \"cuda\",\n) -> torch.Tensor:\n \"\"\"Return ``[0, 1, ..., nnz-1] (% max_seq_len)`` as int32 on ``device``.\n\n If ``max_seq_len`` is None, no wrapping is applied.\n \"\"\"\n pos = torch.arange(nnz, dtype=torch.int32, device=device)\n if max_seq_len is not None:\n pos = pos % max_seq_len\n return pos\n\n\ndef make_rope_cos_sin_cache(\n max_seq_len: int,\n rope_dim: int,\n *,\n base: float = 1e4,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Return concatenated ``[cos | sin]`` cache of shape ``[max_seq_len, rope_dim]``.\"\"\"\n t = torch.arange(max_seq_len, dtype=torch.float32, device=device)\n inv = 1.0 / (\n base\n ** (torch.arange(0, rope_dim, 2, dtype=torch.float32, device=device) / rope_dim)\n )\n freqs = t.unsqueeze(-1) * inv.unsqueeze(0)\n cache = torch.cat([torch.cos(freqs), torch.sin(freqs)], dim=-1)\n return cache.to(dtype)\n\n\ndef make_probs(\n batch_size: int,\n vocab_size: int,\n *,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Return a ``[batch_size, vocab_size]`` probability distribution.\n\n Uses ``softmax(randn(...))`` so each row sums to 1.0. This mirrors the\n pattern used throughout ``tests/utils/test_sampling.py``.\n \"\"\"\n return torch.softmax(\n torch.randn(batch_size, vocab_size, dtype=torch.float32, device=device),\n dim=-1,\n ).to(dtype)\n\n\ndef make_logits(\n batch_size: int,\n vocab_size: int,\n *,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Return ``randn(batch_size, vocab_size)`` logits.\"\"\"\n return torch.randn(batch_size, vocab_size, dtype=dtype, device=device)\n\n\ndef fp8_safe_randn(\n *shape: int,\n scale: float = 0.1,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.bfloat16,\n) -> torch.Tensor:\n \"\"\"``randn(*shape) * scale`` \u2014 keeps values in the FP8/FP4 representable range.\n\n Tests for fp8/fp4 paths typically multiply ``randn`` by 0.1 to avoid\n saturation when quantizing. Use this helper to mirror that convention.\n \"\"\"\n return (torch.randn(*shape, dtype=dtype, device=device) * scale).to(dtype)\n\n\ndef per_tensor_fp8_quantize(\n x: torch.Tensor,\n *,\n fp8_dtype: torch.dtype = torch.float8_e4m3fn,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Per-tensor FP8 quantization, mirroring ``tests/utils_fp8.py:to_float8``.\n\n Returns ``(x_fp8, inv_scale)`` where ``inv_scale`` is the dequant\n multiplier (``float \u2248 fp8 * inv_scale``).\n \"\"\"\n finfo = torch.finfo(fp8_dtype)\n amax = x.abs().amax().clamp(min=1e-12)\n scale = finfo.max / amax\n x_q = (x.float() * scale).clamp(min=finfo.min, max=finfo.max).to(fp8_dtype)\n return x_q, scale.float().reciprocal()\n\n\ndef fp8_block_quant_1d(\n x_bf16: torch.Tensor,\n block: int = 128,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Quantize ``[T, H]`` activations into FP8 with per-``(token, block)``\n column-block scales. Returns ``(x_fp8, scales)`` where\n ``scales`` has shape ``[T, H // block]``.\n\n Mirrors ``_fp8_block_quant_1d`` in\n ``tests/moe/test_dpsk_fused_moe_fp8.py``.\n \"\"\"\n assert x_bf16.dim() == 2\n T, H = x_bf16.shape\n assert H % block == 0\n nb = H // block\n finfo = torch.finfo(torch.float8_e4m3fn)\n max_fp8 = finfo.max\n x_f32 = x_bf16.to(torch.float32)\n x_fp8 = torch.empty((T, H), dtype=torch.float8_e4m3fn, device=x_bf16.device)\n scales = torch.empty((T, nb), dtype=torch.float32, device=x_bf16.device)\n for j in range(nb):\n sl = slice(j * block, (j + 1) * block)\n blk = x_f32[:, sl]\n amax = torch.amax(torch.abs(blk), dim=1)\n s = torch.where(amax > 0, amax / max_fp8, torch.ones_like(amax))\n x_fp8[:, sl] = (blk / s.unsqueeze(1)).to(torch.float8_e4m3fn)\n scales[:, j] = s\n return x_fp8, scales\n\n\ndef fp8_block_quant_2d(\n w_bf16: torch.Tensor,\n block: int = 128,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Quantize weights ``[..., R, C]`` with 2-D ``block \u00d7 block`` scales.\n\n Returns ``(w_fp8, scales)`` where ``scales`` has shape\n ``[..., R // block, C // block]``. Mirrors ``_fp8_block_quant_2d`` in\n ``tests/moe/test_dpsk_fused_moe_fp8.py``.\n \"\"\"\n assert w_bf16.dim() >= 2\n *prefix, R, C = w_bf16.shape\n assert R % block == 0 and C % block == 0\n nb_r, nb_c = R // block, C // block\n finfo = torch.finfo(torch.float8_e4m3fn)\n max_fp8 = finfo.max\n w_f32 = w_bf16.to(torch.float32).contiguous()\n prefix_ndim = len(prefix)\n reshaped = w_f32.reshape(*prefix, nb_r, block, nb_c, block)\n permute_dims = tuple(range(prefix_ndim)) + (\n prefix_ndim,\n prefix_ndim + 2,\n prefix_ndim + 1,\n prefix_ndim + 3,\n )\n blocks = reshaped.permute(permute_dims).contiguous()\n amax = torch.amax(torch.abs(blocks), dim=(-1, -2))\n scales = torch.where(\n amax > 0, amax / max_fp8, torch.ones_like(amax, dtype=torch.float32)\n )\n q_blocks = (blocks / scales.unsqueeze(-1).unsqueeze(-1)).to(torch.float8_e4m3fn)\n inv_permute = [0] * (prefix_ndim + 4)\n for i, p in enumerate(permute_dims):\n inv_permute[p] = i\n w_fp8 = q_blocks.permute(*inv_permute).reshape(*prefix, R, C).contiguous()\n return w_fp8, scales\n\n\n__all__ = [\n \"make_paged_kv_indices\",\n \"make_ragged_indptr\",\n \"make_uniform_qo_indptr\",\n \"make_pos_ids\",\n \"make_rope_cos_sin_cache\",\n \"make_probs\",\n \"make_logits\",\n \"fp8_safe_randn\",\n \"per_tensor_fp8_quantize\",\n \"fp8_block_quant_1d\",\n \"fp8_block_quant_2d\",\n]\n\n# ----- init -----\ndef _svdquant_linear_init(\n *,\n M: int,\n N: int = 3072,\n K: int = 3072,\n device: str = \"cuda\",\n seed: int = 0,\n):\n \"\"\"Build inputs for ``flashinfer.svdquant_linear`` (full SVDQuant linear chain).\"\"\"\n from flashinfer import nvfp4_quantize_smooth # noqa: PLC0415\n\n torch.manual_seed(seed)\n rank = 32\n x = torch.randn(M, K, dtype=torch.bfloat16, device=device)\n w = torch.randn(N, K, dtype=torch.bfloat16, device=device) / math.sqrt(K)\n pqs = torch.rand(K, dtype=torch.bfloat16, device=device) + 0.5\n l1 = torch.randn(N, rank, dtype=torch.bfloat16, device=device) / math.sqrt(rank)\n l2 = torch.randn(rank, K, dtype=torch.bfloat16, device=device) / math.sqrt(K)\n\n x_hat = (x.float() * pqs.float()).to(torch.bfloat16)\n x_gs = (\n ((448 * 6) / x_hat.float().abs().nan_to_num().max())\n .to(torch.float32)\n .reshape(1)\n )\n w_gs = ((448 * 6) / w.float().abs().nan_to_num().max()).to(torch.float32).reshape(1)\n ones = torch.ones(K, dtype=torch.bfloat16, device=device)\n\n weight_fp4, weight_sf = nvfp4_quantize_smooth(w, ones, w_gs)\n alpha = (1.0 / (x_gs * w_gs)).to(torch.float32).reshape(1)\n l2t_smoothed = (pqs.float()[:, None] * l2.float().t()).to(torch.bfloat16)\n l1_scaled = (l1.float() / alpha.item()).to(torch.bfloat16)\n return {\n \"x\": x,\n \"weight_fp4\": weight_fp4,\n \"weight_sf\": weight_sf,\n \"alpha\": alpha,\n \"pre_quant_scale\": pqs,\n \"l2t_smoothed\": l2t_smoothed,\n \"l1_scaled\": l1_scaled,\n \"global_scale\": x_gs,\n }\n" + "init": "from __future__ import annotations\nimport math\nimport torch\n\n# ----- shared init helpers -----\n# Copyright (c) 2025 by FlashInfer team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Shared helpers used by ``TraceTemplate.init`` functions.\n\nThis module contains the small set of input-construction patterns that\nrecur across many templates (paged-KV cache index arrays, ragged indptr,\nRoPE pos_ids and cos/sin caches, sampling probs). Each helper is short and\ndocumented; init functions in ``templates/.py`` call into here so\nthe per-template init bodies stay focused on shape/dtype, not boilerplate.\n\nThe full source of this module is **inlined into every dumped JSON's\n``\"init\"`` field** by ``flashinfer/trace/template.py:_render_init_source``,\nso downstream consumers don't need flashinfer installed to re-run the init\nsnippets.\n\"\"\"\n\n\nfrom typing import Optional, Tuple\n\nimport torch\n\n\ndef make_paged_kv_indices(\n batch_size: int,\n num_pages_per_seq: int,\n page_size: int,\n *,\n device: str = \"cuda\",\n) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n \"\"\"Return ``(kv_indptr, kv_indices, kv_last_page_len)`` for a uniform batch.\n\n Every sequence is assigned exactly ``num_pages_per_seq`` pages, fully\n populated (last-page length == page_size).\n\n Invariants\n ----------\n - ``kv_indptr.shape == (batch_size + 1,)``, dtype int32, monotonic, [0]=0.\n - ``kv_indices == arange(0, batch_size * num_pages_per_seq)``, int32.\n - ``kv_last_page_len == full(batch_size, page_size)``, int32.\n \"\"\"\n total_pages = batch_size * num_pages_per_seq\n kv_indptr = (\n torch.arange(batch_size + 1, dtype=torch.int32, device=device)\n * num_pages_per_seq\n )\n kv_indices = torch.arange(total_pages, dtype=torch.int32, device=device)\n kv_last_page_len = torch.full(\n (batch_size,), page_size, dtype=torch.int32, device=device\n )\n return kv_indptr, kv_indices, kv_last_page_len\n\n\ndef make_ragged_indptr(\n seg_lens,\n *,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.int32,\n) -> torch.Tensor:\n \"\"\"Return cumulative-sum ``indptr`` of length ``len(seg_lens)+1``.\n\n ``seg_lens`` may be a list / tuple / 1-D tensor of segment lengths.\n \"\"\"\n if isinstance(seg_lens, torch.Tensor):\n lens = seg_lens.to(device=device, dtype=dtype)\n else:\n lens = torch.tensor(list(seg_lens), dtype=dtype, device=device)\n indptr = torch.zeros(lens.numel() + 1, dtype=dtype, device=device)\n indptr[1:] = torch.cumsum(lens, dim=0).to(dtype)\n return indptr\n\n\ndef make_uniform_qo_indptr(\n batch_size: int,\n qo_len: int,\n *,\n device: str = \"cuda\",\n) -> torch.Tensor:\n \"\"\"Return ``[0, qo_len, 2*qo_len, ..., batch_size*qo_len]`` int32.\"\"\"\n return torch.arange(batch_size + 1, dtype=torch.int32, device=device) * qo_len\n\n\ndef make_pos_ids(\n nnz: int,\n max_seq_len: Optional[int] = None,\n *,\n device: str = \"cuda\",\n) -> torch.Tensor:\n \"\"\"Return ``[0, 1, ..., nnz-1] (% max_seq_len)`` as int32 on ``device``.\n\n If ``max_seq_len`` is None, no wrapping is applied.\n \"\"\"\n pos = torch.arange(nnz, dtype=torch.int32, device=device)\n if max_seq_len is not None:\n pos = pos % max_seq_len\n return pos\n\n\ndef make_rope_cos_sin_cache(\n max_seq_len: int,\n rope_dim: int,\n *,\n base: float = 1e4,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Return concatenated ``[cos | sin]`` cache of shape ``[max_seq_len, rope_dim]``.\"\"\"\n t = torch.arange(max_seq_len, dtype=torch.float32, device=device)\n inv = 1.0 / (\n base\n ** (torch.arange(0, rope_dim, 2, dtype=torch.float32, device=device) / rope_dim)\n )\n freqs = t.unsqueeze(-1) * inv.unsqueeze(0)\n cache = torch.cat([torch.cos(freqs), torch.sin(freqs)], dim=-1)\n return cache.to(dtype)\n\n\ndef make_probs(\n batch_size: int,\n vocab_size: int,\n *,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Return a ``[batch_size, vocab_size]`` probability distribution.\n\n Uses ``softmax(randn(...))`` so each row sums to 1.0. This mirrors the\n pattern used throughout ``tests/utils/test_sampling.py``.\n \"\"\"\n return torch.softmax(\n torch.randn(batch_size, vocab_size, dtype=torch.float32, device=device),\n dim=-1,\n ).to(dtype)\n\n\ndef make_logits(\n batch_size: int,\n vocab_size: int,\n *,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Return ``randn(batch_size, vocab_size)`` logits.\"\"\"\n return torch.randn(batch_size, vocab_size, dtype=dtype, device=device)\n\n\ndef fp8_safe_randn(\n *shape: int,\n scale: float = 0.1,\n device: str = \"cuda\",\n dtype: torch.dtype = torch.bfloat16,\n) -> torch.Tensor:\n \"\"\"``randn(*shape) * scale`` \u2014 keeps values in the FP8/FP4 representable range.\n\n Tests for fp8/fp4 paths typically multiply ``randn`` by 0.1 to avoid\n saturation when quantizing. Use this helper to mirror that convention.\n \"\"\"\n return (torch.randn(*shape, dtype=dtype, device=device) * scale).to(dtype)\n\n\ndef per_tensor_fp8_quantize(\n x: torch.Tensor,\n *,\n fp8_dtype: torch.dtype = torch.float8_e4m3fn,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Per-tensor FP8 quantization, mirroring ``tests/utils_fp8.py:to_float8``.\n\n Returns ``(x_fp8, inv_scale)`` where ``inv_scale`` is the dequant\n multiplier (``float \u2248 fp8 * inv_scale``).\n \"\"\"\n finfo = torch.finfo(fp8_dtype)\n amax = x.abs().amax().clamp(min=1e-12)\n scale = finfo.max / amax\n x_q = (x.float() * scale).clamp(min=finfo.min, max=finfo.max).to(fp8_dtype)\n return x_q, scale.float().reciprocal()\n\n\ndef fp8_block_quant_1d(\n x_bf16: torch.Tensor,\n block: int = 128,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Quantize ``[T, H]`` activations into FP8 with per-``(token, block)``\n column-block scales. Returns ``(x_fp8, scales)`` where\n ``scales`` has shape ``[T, H // block]``.\n\n Mirrors ``_fp8_block_quant_1d`` in\n ``tests/moe/test_dpsk_fused_moe_fp8.py``.\n \"\"\"\n assert x_bf16.dim() == 2\n T, H = x_bf16.shape\n assert H % block == 0\n nb = H // block\n finfo = torch.finfo(torch.float8_e4m3fn)\n max_fp8 = finfo.max\n x_f32 = x_bf16.to(torch.float32)\n x_fp8 = torch.empty((T, H), dtype=torch.float8_e4m3fn, device=x_bf16.device)\n scales = torch.empty((T, nb), dtype=torch.float32, device=x_bf16.device)\n for j in range(nb):\n sl = slice(j * block, (j + 1) * block)\n blk = x_f32[:, sl]\n amax = torch.amax(torch.abs(blk), dim=1)\n s = torch.where(amax > 0, amax / max_fp8, torch.ones_like(amax))\n x_fp8[:, sl] = (blk / s.unsqueeze(1)).to(torch.float8_e4m3fn)\n scales[:, j] = s\n return x_fp8, scales\n\n\ndef fp8_block_quant_2d(\n w_bf16: torch.Tensor,\n block: int = 128,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Quantize weights ``[..., R, C]`` with 2-D ``block \u00d7 block`` scales.\n\n Returns ``(w_fp8, scales)`` where ``scales`` has shape\n ``[..., R // block, C // block]``. Mirrors ``_fp8_block_quant_2d`` in\n ``tests/moe/test_dpsk_fused_moe_fp8.py``.\n \"\"\"\n assert w_bf16.dim() >= 2\n *prefix, R, C = w_bf16.shape\n assert R % block == 0 and C % block == 0\n nb_r, nb_c = R // block, C // block\n finfo = torch.finfo(torch.float8_e4m3fn)\n max_fp8 = finfo.max\n w_f32 = w_bf16.to(torch.float32).contiguous()\n prefix_ndim = len(prefix)\n reshaped = w_f32.reshape(*prefix, nb_r, block, nb_c, block)\n permute_dims = tuple(range(prefix_ndim)) + (\n prefix_ndim,\n prefix_ndim + 2,\n prefix_ndim + 1,\n prefix_ndim + 3,\n )\n blocks = reshaped.permute(permute_dims).contiguous()\n amax = torch.amax(torch.abs(blocks), dim=(-1, -2))\n scales = torch.where(\n amax > 0, amax / max_fp8, torch.ones_like(amax, dtype=torch.float32)\n )\n q_blocks = (blocks / scales.unsqueeze(-1).unsqueeze(-1)).to(torch.float8_e4m3fn)\n inv_permute = [0] * (prefix_ndim + 4)\n for i, p in enumerate(permute_dims):\n inv_permute[p] = i\n w_fp8 = q_blocks.permute(*inv_permute).reshape(*prefix, R, C).contiguous()\n return w_fp8, scales\n\n\n__all__ = [\n \"make_paged_kv_indices\",\n \"make_ragged_indptr\",\n \"make_uniform_qo_indptr\",\n \"make_pos_ids\",\n \"make_rope_cos_sin_cache\",\n \"make_probs\",\n \"make_logits\",\n \"fp8_safe_randn\",\n \"per_tensor_fp8_quantize\",\n \"fp8_block_quant_1d\",\n \"fp8_block_quant_2d\",\n]\n\n# ----- init -----\ndef _svdquant_linear_init(\n *,\n M: int,\n SF_B: int = 0,\n N: int = 3072,\n K: int = 3072,\n device: str = \"cuda\",\n seed: int = 0,\n):\n \"\"\"Build inputs for ``flashinfer.svdquant_linear`` (full SVDQuant linear chain).\"\"\"\n from flashinfer import nvfp4_quantize_smooth # noqa: PLC0415\n\n del SF_B # derived axis\n torch.manual_seed(seed)\n rank = 32\n x = torch.randn(M, K, dtype=torch.bfloat16, device=device)\n w = torch.randn(N, K, dtype=torch.bfloat16, device=device) / math.sqrt(K)\n pqs = torch.rand(K, dtype=torch.bfloat16, device=device) + 0.5\n l1 = torch.randn(N, rank, dtype=torch.bfloat16, device=device) / math.sqrt(rank)\n l2 = torch.randn(rank, K, dtype=torch.bfloat16, device=device) / math.sqrt(K)\n\n x_hat = (x.float() * pqs.float()).to(torch.bfloat16)\n x_gs = (\n ((448 * 6) / x_hat.float().abs().nan_to_num().max())\n .to(torch.float32)\n .reshape(1)\n )\n w_gs = ((448 * 6) / w.float().abs().nan_to_num().max()).to(torch.float32).reshape(1)\n ones = torch.ones(K, dtype=torch.bfloat16, device=device)\n\n weight_fp4, weight_sf = nvfp4_quantize_smooth(w, ones, w_gs)\n alpha = (1.0 / (x_gs * w_gs)).to(torch.float32).reshape(1)\n l2t_smoothed = (pqs.float()[:, None] * l2.float().t()).to(torch.bfloat16)\n l1_scaled = (l1.float() / alpha.item()).to(torch.bfloat16)\n return {\n \"x\": x,\n \"weight_fp4\": weight_fp4,\n \"weight_sf\": weight_sf,\n \"alpha\": alpha,\n \"pre_quant_scale\": pqs,\n \"l2t_smoothed\": l2t_smoothed,\n \"l1_scaled\": l1_scaled,\n \"global_scale\": x_gs,\n }\n" } From 0373d0cae53b657063a418ef67cf9d400f50edb0 Mon Sep 17 00:00:00 2001 From: Anthony Chang <27950904+rosenrodt@users.noreply.github.com> Date: Tue, 18 Aug 2026 15:07:17 +0800 Subject: [PATCH 7/7] Reconcile SM120 FP4 scale layouts Preserve the SVDQuant-tested NVFP4 SFA mode grouping while retaining upstream's trailing-K normalization for the new MXFP4 b12x path. Changes - select SFA rank normalization by scale-vector format - document why NVFP4 and MXFP4 collapse different modes Validation - pre-commit run -a - python -m pytest -q tests/gemm/test_nvfp4_svdquant_gemm.py - focused b12x ragged-K, short-K, NVFP4, MXFP4, and MXFP4-alpha checks Result - 34 SVDQuant cases passed with 48 skipped - 6 upstream b12x regression cases and 6 representative format cases passed --- .../gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py b/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py index 6b6c6266b8e..4451824ad68 100644 --- a/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py +++ b/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py @@ -726,7 +726,15 @@ def _partition_fragment_SFA( thr_vmk = (thr_vmnk[0], (thr_vmnk[1], thr_vmnk[3])) partitioned_sfa = thr_tensor[thr_vmk, (None, None)] partitioned_sfa = cute.group_modes(cute.flatten(partitioned_sfa), 0, 2) - partitioned_sfa = _collapse_to_vmk(partitioned_sfa) + if cutlass.const_expr( + self.sf_vec_size == 16 and cute.rank(partitioned_sfa) > 3 + ): + # NVFP4's extra SFA mode belongs to the MN coordinate for wide-M + # tiles; retain the SVDQuant-tested grouping. MXFP4 instead carries + # an extra trailing K mode and uses the upstream normalization. + partitioned_sfa = cute.group_modes(partitioned_sfa, 1, 3) + else: + partitioned_sfa = _collapse_to_vmk(partitioned_sfa) return cute.make_fragment_like(partitioned_sfa) def _partition_fragment_SFB(