diff --git a/benchmarks/flashinfer_benchmark.py b/benchmarks/flashinfer_benchmark.py index 45a95f07067..fc23c942f63 100644 --- a/benchmarks/flashinfer_benchmark.py +++ b/benchmarks/flashinfer_benchmark.py @@ -6,7 +6,6 @@ from routines.flashinfer_benchmark_utils import ( benchmark_apis, full_output_columns, - output_column_dict, ) @@ -75,9 +74,11 @@ def run_test(args): if args.output_path is not None: with open(args.output_path, "a") as fout: for cur_res in res: - for key in output_column_dict["general"]: - # Only set from args if the routine hasn't already set a value - # This preserves routine-specific formatting while providing defaults + for key in full_output_columns: + # Backfill every output column the routine didn't set: from + # args when available, else "". Covers columns belonging to + # other routines (e.g. attention's s_qo) that would otherwise + # KeyError below. Routine-set values are preserved. if key not in cur_res or cur_res[key] == "": cur_res[key] = getattr(args, key, "") diff --git a/benchmarks/routines/flashinfer_benchmark_utils.py b/benchmarks/routines/flashinfer_benchmark_utils.py index 8f6a4b40507..2c244e7f029 100644 --- a/benchmarks/routines/flashinfer_benchmark_utils.py +++ b/benchmarks/routines/flashinfer_benchmark_utils.py @@ -199,6 +199,7 @@ "bmm_fp8", "bmm_mxfp8", "mm_fp4", + "mm_bf16_fp4", "mm_mxfp8", "mm_bf16", "bmm_bf16", diff --git a/benchmarks/routines/gemm.py b/benchmarks/routines/gemm.py index d5a96411fe0..804dfa1f798 100644 --- a/benchmarks/routines/gemm.py +++ b/benchmarks/routines/gemm.py @@ -44,6 +44,8 @@ def run_gemm_test(args): return testBmmMxfp8(args) elif args.routine == "mm_fp4": return testMmFp4(args) + elif args.routine == "mm_bf16_fp4": + return testMmBf16Fp4(args) elif args.routine == "mm_mxfp8": return testMmMxfp8(args) elif args.routine == "mm_bf16": @@ -200,6 +202,11 @@ def parse_gemm_args(line, parser): args.input_dtype = "bfloat16" if not has_mat2_dtype_arg: args.mat2_dtype = "bfloat16" + if args.routine == "mm_bf16_fp4": + if not has_backends_arg: + args.backends = ["cute-dsl"] + if not has_input_dtype_arg: + args.input_dtype = "bfloat16" if args.verbose >= 1: print(f"[INFO] {args = }") return args @@ -1326,6 +1333,233 @@ def run_backend( return res +# E2M1 (FP4) value table, signed (codes 0-7 positive, 8-15 negative), matching +# ``flashinfer.nvfp4_quantize``. +_E2M1_VALUES_FP32 = ( + 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, + -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0, +) # fmt: skip + + +def _dequantize_bf16_fp4_ref(b, b_descale, alpha, n, k, block_size): + """PyTorch implementation of swizzled nvfp4 dequantization to fp32.""" + import torch + from flashinfer.gemm.gemm_bf16_fp4 import _unswizzle_sf_128x4 + + device = b.device + k_sf = k // block_size + lut = torch.tensor(_E2M1_VALUES_FP32, dtype=torch.float32, device=device) + b_int = b.to(torch.int64) + codes = torch.stack([b_int & 0xF, (b_int >> 4) & 0xF], dim=-1).reshape(n, k) + values = lut[codes] + sf = _unswizzle_sf_128x4(b_descale, n, k_sf).view(torch.float8_e4m3fn) + sf_expanded = sf.to(torch.float32).repeat_interleave(block_size, dim=1) + weight = values * sf_expanded + if alpha is not None: + weight = weight * alpha.to(torch.float32) + return weight + + +def testMmBf16Fp4(args): + """Benchmark mm_bf16_fp4 (bf16 activation x FP4 weight, bf16 x fp4). + + Weights are produced by ``flashinfer.nvfp4_quantize(sfLayout=layout_128x4)`` + -- the same format the new API expects. + + Constraints: + * N must be divisible by 64 (the weight-prepack / kernel N tile). + The 128x4 SF swizzle only pads N to 128; prepare_bf16_fp4_weights + unswizzles the padded tail, so any N % 64 == 0 works. + * K must be divisible by 16 (FP4 block size). + * input_dtype must be bfloat16 (the only A dtype currently + supported; fp16 deferred). + + Refcheck uses an fp32 dequant + matmul. + """ + if args.verbose >= 1: + print("[INFO] Running testMmBf16Fp4") + print(f"[INFO] FlashInfer version: {flashinfer.__version__}") + + device = get_device(args) + if args.generate_repro_command: + print( + f"[INFO] To reproduce this test case, run the following command: {args.repro_command}" + ) + + m, n, k = args.m, args.n, args.k + input_dtype = dtype_str_to_torch_dtype(args.input_dtype) + out_dtype = dtype_str_to_torch_dtype(args.out_dtype) + backends = args.backends + run_refcheck = args.refcheck + is_cuda_graph_compatible = not args.no_cuda_graph + + if input_dtype != torch.bfloat16: + raise ValueError( + f"mm_bf16_fp4 benchmark requires input_dtype=bfloat16, got {args.input_dtype}" + ) + if out_dtype not in (torch.bfloat16, torch.float16): + raise ValueError( + f"mm_bf16_fp4 benchmark requires out_dtype in (bfloat16, float16), got {args.out_dtype}" + ) + if n % 64 != 0: + # N must be a multiple of the 64-wide N tile used by the weight + # prepack and the cute-dsl kernel. The 128x4 SF swizzle only *pads* + # N to 128, and prepare_bf16_fp4_weights unswizzles the padded + # tail, so any n % 64 == 0 works. + raise ValueError("mm_bf16_fp4 benchmark requires n % 64 == 0") + if k % 16 != 0: + raise ValueError("mm_bf16_fp4 benchmark requires k % 16 == 0 (FP4 block size)") + + torch.manual_seed(args.random_seed) + a = torch.randn((m, k), device=device, dtype=input_dtype) * 0.5 + w = torch.randn((n, k), device=device, dtype=input_dtype) * 0.1 + g_b = (448 * 6) / w.float().abs().nan_to_num().max() + b_fp4, b_sf = flashinfer.nvfp4_quantize( + w, + g_b, + sfLayout=flashinfer.SfLayout.layout_128x4, + do_shuffle=False, + backend="cute-dsl", + ) + alpha = torch.tensor([1.0 / g_b.item()], device=device, dtype=torch.float32) + + if args.verbose >= 2: + print(f"[VVERBOSE] {a.shape = } {a.dtype = }") + print(f"[VVERBOSE] {b_fp4.shape = } {b_fp4.dtype = }") + print(f"[VVERBOSE] {b_sf.shape = } {b_sf.dtype = }") + + # Per-backend prep + runner closures. Prep is one-shot and not timed. + backend_runners = {} + + def make_runner(b_p, sf_p, alpha_p, backend): + def run(a): + # mm_bf16_fp4: bf16 activation a against the prepared FP4 + # weight; b_p/sf_p are prepare_bf16_fp4_weights outputs. + return flashinfer.mm_bf16_fp4( + a, + b_p, + sf_p, + alpha_p, + backend=backend, + out_dtype=out_dtype, + block_size=16, + enable_pdl=args.enable_pdl, + ) + + return run + + backends_to_remove = [] + for backend in backends: + try: + b_p, sf_p, alpha_p = flashinfer.prepare_bf16_fp4_weights( + b_fp4, b_sf, alpha, backend=backend + ) + runner = make_runner(b_p, sf_p, alpha_p, backend) + runner(a) + backend_runners[backend] = runner + except Exception as e: + print( + f"[INFO] {backend} backend does not support this configuration: {type(e).__name__}: {e}" + ) + backends_to_remove.append(backend) + + for backend in backends_to_remove: + backends.remove(backend) + + if len(backends) == 0: + print("[ERROR] No backends passed validation. Exiting.") + return + + autotune_supported_backends = ["cudnn", "cute-dsl"] + cache_path = getattr(args, "autotune_cache", None) + if getattr(args, "autotune", False): + warmup_iters = ( + args.dry_run_iters if args.dry_run_iters and args.dry_run_iters > 0 else 10 + ) + for cur_backend in backends: + if cur_backend in autotune_supported_backends: + if args.verbose >= 1: + print( + f"[INFO] Autotune warmup for mm_bf16_fp4 {cur_backend}: " + f"{warmup_iters} iters" + ) + with autotune(True, cache=cache_path): + for _ in range(warmup_iters): + backend_runners[cur_backend](a) + elif cache_path: + with autotune(False, cache=cache_path): + pass + + ref = None + if run_refcheck: + weight_fp32 = _dequantize_bf16_fp4_ref(b_fp4, b_sf, alpha, n, k, 16) + ref = (a.float() @ weight_fp32.T).to(out_dtype) + + res = [] + flops = 2 * m * n * k + bytes_accessed = ( + m * k * input_dtype.itemsize + + (k // 2) * n # FP4 weight (uint8, 2 codes / byte) + + (k // 16) * n # FP8-E4M3 per-block SF + + m * n * out_dtype.itemsize + ) + + refcheck_tol = dict(rtol=1.5e-2, atol=1.5e-2) + for backend in backends: + runner = backend_runners[backend] + if run_refcheck: + out = runner(a) + try: + torch.testing.assert_close(out, ref, **refcheck_tol) + except AssertionError as e: + if args.allow_output_mismatch: + print(f"[WARNING] {backend} output mismatch vs fp32 ref: {e}") + else: + raise + + timing = bench_gpu_time( + fn=runner, + dry_run_iters=args.dry_run_iters, + repeat_iters=args.num_iters, + sleep_after_run=True, # GEMMs are very MMA-heavy, so prefer sleep to reduce throttling. + enable_cupti=args.use_cupti, + use_cuda_graph=is_cuda_graph_compatible, + cold_l2_cache=True, + input_args=(a,), + ) + median_time = float(np.median(timing)) + std_time = float(np.std(timing)) + tflops = flops / median_time / 1e9 + tb_per_sec = bytes_accessed / median_time / 1e9 + backend_name = backend + ( + "_autotune" + if ( + getattr(args, "autotune", False) + and backend in autotune_supported_backends + ) + else "" + ) + print_perf_metrics(backend_name, median_time, std_time, tflops, tb_per_sec) + res.append( + { + "routine": args.routine, + "median_time": median_time, + "std_time": std_time, + "tflops": tflops, + "tb_per_sec": tb_per_sec, + "backend": backend_name, + "resolved_backend": backend, + "m": m, + "n": n, + "k": k, + "input_dtype": str(input_dtype).split(".")[-1], + "out_dtype": str(out_dtype).split(".")[-1], + "case_tag": args.case_tag, + } + ) + return res + + def testMmMxfp8(args): """ Test mm_mxfp8 API. diff --git a/flashinfer/__init__.py b/flashinfer/__init__.py index 94136ba26cc..b1a11e523e1 100644 --- a/flashinfer/__init__.py +++ b/flashinfer/__init__.py @@ -105,6 +105,8 @@ from .gemm import bmm_mxfp8 as bmm_mxfp8 from .gemm import mm_bf16 as mm_bf16 from .gemm import mm_fp4 as mm_fp4 +from .gemm import mm_bf16_fp4 as mm_bf16_fp4 +from .gemm import prepare_bf16_fp4_weights as prepare_bf16_fp4_weights from .gemm import mm_fp8 as mm_fp8 from .gemm import mm_mxfp8 as mm_mxfp8 from .gemm import tgv_gemm_sm100 as tgv_gemm_sm100 diff --git a/flashinfer/cute_dsl/fp4_common.py b/flashinfer/cute_dsl/fp4_common.py index 1400a315856..0a3e7127e8c 100644 --- a/flashinfer/cute_dsl/fp4_common.py +++ b/flashinfer/cute_dsl/fp4_common.py @@ -283,6 +283,33 @@ def get_smem_ptr_as_int32( return elem_ptr.toint(loc=loc, ip=ip) +@dsl_user_op +def ld_shared_v2_u32(smem_addr: Int32, *, loc=None, ip=None) -> Tuple[Uint32, Uint32]: + """Load 64 bits (2 x uint32) from shared memory via ld.shared.v2.u32. + + Args: + smem_addr: 32-bit shared memory address (from get_smem_ptr_as_int32). + Caller is responsible for ensuring 8-byte alignment. + + Returns: + 2 Uint32 values (8 bytes total). + """ + result = llvm.inline_asm( + llvm.StructType.get_literal([T.i32(), T.i32()]), + [Int32(smem_addr).ir_value(loc=loc, ip=ip)], + "ld.shared.v2.u32 {$0, $1}, [$2];", + "=r,=r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + v0 = llvm.extractvalue(T.i32(), result, [0], loc=loc, ip=ip) + v1 = llvm.extractvalue(T.i32(), result, [1], loc=loc, ip=ip) + return Uint32(v0), Uint32(v1) + + @dsl_user_op def ld_shared_v4_u32( smem_addr: Int32, *, loc=None, ip=None @@ -1996,6 +2023,35 @@ def cvt_e4m3_to_f32_via_f16(fp8_val: Uint32, *, loc=None, ip=None) -> Float32: ) +@dsl_user_op +def cvt_s0e5m3_to_f16x2_broadcast(fp8_val: Uint32, *, loc=None, ip=None) -> Uint32: + """Convert one S0E5M3 scale byte to an f16x2 with the scale in both lanes. + + S0E5M3 (sign-0, 5-exp, 3-mantissa) is a host-side reformat of the per-block + E4M3 scale, rebiased to fp16's bias (exp 7->15, i.e. byte += 0x40) so the + bits line up with fp16 directly: ``f16(byte) = byte << 7`` (exp -> bits 14-10, + the 3 mantissa bits -> bits 9-7, low mantissa and sign = 0). Broadcasting to + both f16x2 lanes is then ``byte * 0x00800080`` = ``(byte<<7) | (byte<<23)`` -- + a single ``mul.lo.u32`` (the two shifted copies never overlap, so the mul is + exactly the OR). Replaces the 3-op E4M3 path (cvt.u16 + cvt.f16x2.e4m3x2 + + prmt) with 1 op, and the per-block scale stays 1 byte (memory-neutral). + Numerically exact for normal E4M3 scales (both formats carry 3 mantissa bits). + """ + return Uint32( + llvm.inline_asm( + T.i32(), + [Uint32(fp8_val).ir_value(loc=loc, ip=ip)], + "mul.lo.u32 $0, $1, 0x00800080;", + "=r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + @dsl_user_op def fp4_decode_4bytes( packed_u32: Uint32, *, loc=None, ip=None diff --git a/flashinfer/gemm/__init__.py b/flashinfer/gemm/__init__.py index 780fd64f510..0ac7a63b04a 100644 --- a/flashinfer/gemm/__init__.py +++ b/flashinfer/gemm/__init__.py @@ -23,6 +23,11 @@ from .gemm_base import group_gemm_fp8_nt_groupwise as group_gemm_fp8_nt_groupwise from .gemm_base import fp8_blockscale_gemm_sm90 as fp8_blockscale_gemm_sm90 +from .gemm_bf16_fp4 import ( + mm_bf16_fp4 as mm_bf16_fp4, + prepare_bf16_fp4_weights as prepare_bf16_fp4_weights, +) + from .routergemm import ( mm_M1_16_K6144_N256 as mm_M1_16_K6144_N256, mm_M1_16_K7168_N128 as mm_M1_16_K7168_N128, @@ -101,6 +106,8 @@ "gemm_fp8_nt_groupwise", "group_gemm_fp8_nt_groupwise", "fp8_blockscale_gemm_sm90", + "mm_bf16_fp4", + "prepare_bf16_fp4_weights", "mm_M1_16_K6144_N256", "mm_M1_16_K7168_N128", "mm_M1_16_K7168_N256", diff --git a/flashinfer/gemm/gemm_bf16_fp4.py b/flashinfer/gemm/gemm_bf16_fp4.py new file mode 100644 index 00000000000..3d0b3674ce4 --- /dev/null +++ b/flashinfer/gemm/gemm_bf16_fp4.py @@ -0,0 +1,309 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 by FlashInfer team. +# SPDX-License-Identifier: Apache-2.0 +"""BF16 x FP4 (W4A16) dense GEMM public API. + +Backend implementation details live in respective submodules +gemm_bf16_fp4_cudnn and gemm_bf16_fp4_cute_dsl. +""" + +from typing import Literal, Optional, Tuple + +import torch + +from ..api_logging import flashinfer_api +from ..trace.templates.gemm import mm_bf16_fp4_trace_dispatch +from ..utils import backend_requirement, supported_compute_capability + +from .gemm_base import ( + CUDNN_AVAILABLE, + _check_cudnn_fp4_availability, + _check_cute_dsl_availability, +) + +if CUDNN_AVAILABLE: + import cudnn + + +# Earliest cuDNN backend version supporting bf16 x fp4 GEMM +_CUDNN_BF16_FP4_MIN_BACKEND_VERSION = 92301 + + +def _check_mm_bf16_fp4_problem_size( + a: torch.Tensor, + b: torch.Tensor, + b_descale: torch.Tensor, + alpha: Optional[torch.Tensor] = None, + *, + backend: Literal["cudnn", "cute-dsl"], + out_dtype: Optional[torch.dtype] = None, + out: Optional[torch.Tensor] = None, + block_size: int = 16, + enable_pdl: bool = True, +): + if a.dim() != 2: + raise ValueError(f"a must be 2-D (M, K); got shape {tuple(a.shape)}") + if a.dtype != torch.bfloat16: + raise TypeError( + f"a must be bfloat16; got {a.dtype}. fp16 support is not implemented yet." + ) + if out_dtype is not None and out_dtype not in (torch.bfloat16, torch.float16): + raise ValueError(f"out_dtype must be bfloat16 or float16; got {out_dtype}") + if block_size != 16: + raise ValueError(f"block_size must be 16 for FP4; got {block_size}") + if alpha is not None and alpha.device != a.device: + raise ValueError( + f"alpha must be on the same device as a ({a.device}); got {alpha.device}" + ) + return True + + +@supported_compute_capability([100, 103, 110, 120, 121]) +def _cudnn_bf16_fp4_requirement( + a: torch.Tensor, + b: torch.Tensor, + b_descale: torch.Tensor, + alpha: Optional[torch.Tensor] = None, + *, + backend: Literal["cudnn", "cute-dsl"], + out_dtype: Optional[torch.dtype] = None, + out: Optional[torch.Tensor] = None, + block_size: int = 16, + enable_pdl: bool = True, +): + """cuDNN backend: requires a cuDNN build with FP4 block-scale support. + + Raises ``ValueError`` (not RuntimeError) so ``backend="auto"`` skips + this backend instead of aborting (auto only catches ValueError). + """ + if b.dtype != torch.uint8: + raise ValueError( + f"cudnn bf16 x fp4 expects the uint8 prepared weight from " + f"prepare_bf16_fp4_weights(..., backend='cudnn'); got {b.dtype}." + ) + _check_cudnn_fp4_availability() + + backend_version = cudnn.backend_version() + if backend_version < _CUDNN_BF16_FP4_MIN_BACKEND_VERSION: + raise ValueError( + f"cuDNN bf16 x fp4 GEMM requires backend version >= " + f"{_CUDNN_BF16_FP4_MIN_BACKEND_VERSION} (9.23.1), found {backend_version}. " + ) + return True + + +@supported_compute_capability([100, 103, 110, 120, 121]) +def _cute_dsl_bf16_fp4_requirement( + a: torch.Tensor, + b: torch.Tensor, + b_descale: torch.Tensor, + alpha: Optional[torch.Tensor] = None, + *, + backend: Literal["cudnn", "cute-dsl"], + out_dtype: Optional[torch.dtype] = None, + out: Optional[torch.Tensor] = None, + block_size: int = 16, + enable_pdl: bool = True, +): + if b.dtype != torch.int32: + raise ValueError( + f"cute-dsl bf16 x fp4 expects the int32 tile-packed weight from " + f"prepare_bf16_fp4_weights(..., backend='cute-dsl'); got {b.dtype}." + ) + _check_cute_dsl_availability() + return True + + +# ============================================================================= +# Public dispatchers +# ============================================================================= + + +@flashinfer_api +def prepare_bf16_fp4_weights( + b: torch.Tensor, + b_descale: torch.Tensor, + alpha: Optional[torch.Tensor] = None, + *, + backend: Literal["cudnn", "cute-dsl"], + block_size: int = 16, +) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """Prepare FP4 weights for the bf16 x fp4 GEMM, for a specific backend. + + The caller is expected to start with weights in the canonical format + that :func:`flashinfer.nvfp4_quantize` produces with + ``sfLayout=layout_128x4``: + + * ``b`` is ``(N, K // 2)`` ``uint8`` with two FP4 codes packed per + byte (low nibble = K=2i, high nibble = K=2i+1). + * ``b_descale`` is the 128x4-swizzled FP8-E4M3 per-block scales, + either as a 1-D byte buffer or a 2-D tensor. + + Each backend transforms these into whatever layout its compute kernel + expects. The returned ``(b, b_descale, alpha)`` tuple must be passed + back to :func:`flashinfer.mm_bf16_fp4` with the *same* ``backend`` -- + the shapes / dtypes may not match other backends' expectations. + + Args: + b: ``(N, K // 2)`` ``uint8`` packed FP4 weight. + b_descale: 128x4-swizzled FP8-E4M3 scale factors from + ``nvfp4_quantize``. Either 1-D byte buffer or 2-D tensor. + alpha: Optional ``(1,) float32`` global scalar. Pass ``None`` + (default) for implicit ``alpha=1.0``. Returned unchanged; + forward the returned tuple to :func:`flashinfer.mm_bf16_fp4`. + backend: Identifier of a supported backend (``"cudnn"`` or + ``"cute-dsl"``). + block_size: SF block size. Always 16 for FP4. + + Returns: + ``(b_prepared, b_descale_prepared, alpha_prepared)`` -- pass all + three to :func:`flashinfer.mm_bf16_fp4` with the same ``backend``. + + Raises: + ValueError: ``backend`` is unknown, or an input has an invalid + shape (``b`` not 2-D, ``K`` not a multiple of ``block_size``, + or ``alpha`` not shape ``(1,)``). + TypeError: ``b`` is not ``uint8`` or ``alpha`` is not ``float32``. + """ + if b.dim() != 2: + raise ValueError(f"b must be 2-D (N, K/2); got shape {tuple(b.shape)}") + if b.dtype != torch.uint8: + raise TypeError(f"b must be uint8; got {b.dtype}") + k = int(b.shape[1]) * 2 + if k % block_size != 0: + raise ValueError(f"K={k} must be a multiple of block_size={block_size}") + n = int(b.shape[0]) + k_sf = k // block_size + expected_sf_bytes = ((n + 127) // 128) * ((k_sf + 3) // 4) * 512 + sf_bytes = b_descale.numel() * b_descale.element_size() + if sf_bytes < expected_sf_bytes: + raise ValueError( + f"b_descale has {sf_bytes} bytes but the 128x4-swizzled layout for " + f"N={n}, K_sf={k_sf} requires at least {expected_sf_bytes}" + ) + if alpha is not None: + if alpha.dim() != 1 or alpha.shape[0] != 1: + raise ValueError(f"alpha must be shape (1,); got {tuple(alpha.shape)}") + if alpha.dtype != torch.float32: + raise TypeError(f"alpha must be float32; got {alpha.dtype}") + if backend == "cudnn": + from .gemm_bf16_fp4_cudnn import _prepare_cudnn + + return _prepare_cudnn(b, b_descale, alpha, block_size) + if backend == "cute-dsl": + from .gemm_bf16_fp4_cute_dsl import _prepare_cute_dsl + + return _prepare_cute_dsl(b, b_descale, alpha, block_size) + raise ValueError(f"Unknown backend {backend!r}. Supported: 'cudnn', 'cute-dsl'.") + + +@backend_requirement( + { + "cudnn": _cudnn_bf16_fp4_requirement, + "cute-dsl": _cute_dsl_bf16_fp4_requirement, + }, + common_check=_check_mm_bf16_fp4_problem_size, +) +@flashinfer_api(trace=mm_bf16_fp4_trace_dispatch) +def mm_bf16_fp4( + a: torch.Tensor, + b: torch.Tensor, + b_descale: torch.Tensor, + alpha: Optional[torch.Tensor] = None, + *, + backend: Literal["cudnn", "cute-dsl"], + out_dtype: Optional[torch.dtype] = None, + out: Optional[torch.Tensor] = None, + block_size: int = 16, + enable_pdl: bool = True, +) -> torch.Tensor: + """BF16 x FP4 GEMM: ``out = (a @ dequant(b).T) * alpha``. + + Intended to support **W4A16** workloads (4-bit weights, 16-bit activations) + nvfp4 weights must be prepared for ``backend`` by + :func:`prepare_bf16_fp4_weights`. ``b``, ``b_descale``, and ``alpha``. + + Example: + .. code-block:: python + + # 1) Prepare weights for a backend (once, at model load). + b_p, sf_p, alpha_p = flashinfer.prepare_bf16_fp4_weights( + b, b_descale, alpha, backend="cute-dsl", + ) + # 2) Run the GEMM with the *same* backend tag. + out = flashinfer.mm_bf16_fp4( + a, b_p, sf_p, alpha_p, backend="cute-dsl", + ) + + Args: + a: ``(M, K)`` activation matrix in ``torch.bfloat16``. This is + the only currently supported activation dtype; fp16 support + can be added when needed. + b: Prepared weight tensor (backend-specific layout). + b_descale: Prepared scale-factor tensor (backend-specific layout). + alpha: Optional ``(1,) float32`` global scalar. Pass through + whatever ``prepare_bf16_fp4_weights`` returned -- it may be + ``None`` if the backend folded it into ``b_descale``. + backend: Same identifier passed to ``prepare_bf16_fp4_weights``. + out_dtype: Output dtype. Defaults to ``a.dtype`` (``bfloat16``). + out: Optional preallocated ``(M, N)`` output tensor. + block_size: SF block size. Always 16 for FP4. + enable_pdl: Enable Programmatic Dependent Launch + + Returns: + ``(M, N)`` tensor of ``out_dtype``. + """ + out_dtype = out_dtype or a.dtype + if backend == "cudnn": + from .gemm_bf16_fp4_cudnn import _compute_cudnn + + return _compute_cudnn(a, b, b_descale, alpha, out_dtype, out, block_size) + if backend == "cute-dsl": + from .gemm_bf16_fp4_cute_dsl import _compute_cute_dsl + + return _compute_cute_dsl( + a, b, b_descale, alpha, out_dtype, out, block_size, enable_pdl=enable_pdl + ) + raise ValueError(f"Unknown backend {backend!r}. Supported: 'cudnn', 'cute-dsl'.") + + +# ============================================================================= +# Shared SF utility +# ============================================================================= +# +# Used by both backends' prepare paths (cuDNN and cute-dsl) to turn the +# canonical 128x4-swizzled SF into a linear ``(N, K_sf)`` layout. + + +def _unswizzle_sf_128x4(sf_swizzled: torch.Tensor, n: int, k_sf: int) -> torch.Tensor: + """Reverse the 128x4 SF swizzle into a flat ``(N, K_sf)`` byte tensor. + + The swizzle stores SF in 512-byte blocks each holding 128 N-rows x 4 + K_sf-cols. The byte address of logical ``(n, k_sf)`` is:: + + offset = ((n // 128) * sf_pad_blocks + k_sf // 4) * 512 + + (n % 32) * 16 + + ((n % 128) // 32) * 4 + + (k_sf % 4) + + where ``sf_pad_blocks = ceil(k_sf, 4) // 4`` accounts for K_sf + padding inside each 128-row N block. + """ + device = sf_swizzled.device + sf_flat = sf_swizzled.contiguous().view(torch.uint8).view(-1) + sf_pad_blocks = (k_sf + 3) // 4 # ceil_div(k_sf, 4) + n_idx = torch.arange(n, device=device, dtype=torch.int64) + k_idx = torch.arange(k_sf, device=device, dtype=torch.int64) + n_grid, k_grid = torch.meshgrid(n_idx, k_idx, indexing="ij") + offsets = ( + ((n_grid // 128) * sf_pad_blocks + (k_grid // 4)) * 512 + + (n_grid % 32) * 16 + + ((n_grid % 128) // 32) * 4 + + (k_grid % 4) + ) + return sf_flat[offsets] + + +__all__ = [ + "prepare_bf16_fp4_weights", + "mm_bf16_fp4", +] diff --git a/flashinfer/gemm/gemm_bf16_fp4_cudnn.py b/flashinfer/gemm/gemm_bf16_fp4_cudnn.py new file mode 100644 index 00000000000..0a73e28b724 --- /dev/null +++ b/flashinfer/gemm/gemm_bf16_fp4_cudnn.py @@ -0,0 +1,567 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 by FlashInfer team. +# SPDX-License-Identifier: Apache-2.0 +"""cuDNN backend for the bf16 x fp4 GEMM (graph build / execute / runner).""" + +import functools +from typing import List, Optional, Tuple + +import torch + +from ..autotuner import ( + AutoTuner, + ConstraintSpec, + DynamicTensorSpec, + OptimizationProfile, + TunableRunner, + TuningConfig, +) +from ..fused_moe.utils import ( + get_hybrid_num_tokens_buckets, + map_to_hybrid_bucket_uncapped, +) +from ..utils import _get_cache_buf, get_native_fp4_dtype + +from .gemm_base import ( + CUDNN_AVAILABLE, + DEFAULT_WORKSPACE_SIZE, + UIDs, + _check_cudnn_fp4_availability, + _get_cudnn_handle, + _get_cudnn_override_shape_workspace_size, + _get_cudnn_workspace_size, + _torch_data_type_to_cudnn_data_type, + is_cudnn_override_shape_available, +) +from .gemm_bf16_fp4 import _unswizzle_sf_128x4 + +if CUDNN_AVAILABLE: + import cudnn + +# Sentinel "cache M" for override-shape graphs (any value works; this one +# covers typical LLM inference shapes). Kept local to avoid importing a +# private constant from gemm_base. +_OVERRIDE_SHAPE_CACHE_M = 8192 + + +def _bf16_fp4_b_descale_layout(batch, n, k, block_size): + """Return ``(dim, stride, reordering_type)`` for the B scale-factor tensor.""" + k_sf = k // block_size + dim = (batch, k_sf, n) + stride = (k_sf * n, 1, k_sf) + return dim, stride, cudnn.tensor_reordering.NONE + + +def _build_bf16_fp4_graph_common( + graph, + a_cudnn_tensor, + b_cudnn_tensor, + block_descale_b_cudnn_tensor, + a_type, + o_type, + block_size, + alpha_is_not_none, +): + """Shared graph body: dequant(B) -> A @ dequant(B) -> optional alpha.""" + dequant_b_tensor = graph.block_scale_dequantize( + b_cudnn_tensor, + block_descale_b_cudnn_tensor, + block_size=[block_size, 1], + name="dequant_b", + ) + dequant_b_tensor.set_data_type(a_type) + + c_tensor = graph.matmul( + a_cudnn_tensor, + dequant_b_tensor, + compute_data_type=cudnn.data_type.FLOAT, + name="gemm", + ) + c_tensor.set_data_type(cudnn.data_type.FLOAT) + + c_final_cudnn_tensor = c_tensor + if alpha_is_not_none: + global_scale_cudnn_tensor = graph.tensor( + name="global_scale", + dim=(1, 1, 1), + stride=(1, 1, 1), + data_type=cudnn.data_type.FLOAT, + ) + c_final_cudnn_tensor = graph.mul( + name="scale_mul", + a=c_tensor, + b=global_scale_cudnn_tensor, + compute_data_type=cudnn.data_type.FLOAT, + ) + global_scale_cudnn_tensor.set_uid(UIDs.ALPHA_UID.value) + + c_final_cudnn_tensor.set_name("c_final").set_output(True).set_data_type(o_type) + + a_cudnn_tensor.set_uid(UIDs.A_UID.value) + b_cudnn_tensor.set_uid(UIDs.B_UID.value) + block_descale_b_cudnn_tensor.set_uid(UIDs.BLOCK_DESCALE_B_UID.value) + c_final_cudnn_tensor.set_uid(UIDs.O_UID.value) + return c_final_cudnn_tensor + + +@functools.lru_cache(maxsize=1024) +def build_cudnn_bf16_fp4_graph( + batch, + m, + n, + k, + a_type, + o_type, + block_size, + device, + alpha_is_not_none, + use_nvfp4, + policy=None, +): + """Build a fixed-shape cuDNN bf16 x fp4 GEMM graph (no override-shape).""" + _check_cudnn_fp4_availability() + if policy is None: + policy = cudnn.build_plan_policy.HEURISTICS_CHOICE + + scale_type = cudnn.data_type.FP8_E4M3 if use_nvfp4 else cudnn.data_type.FP8_E8M0 + + a_shape = (batch, m, k) + a_stride = (m * k, k, 1) + # b weight bytes are row-major (N, K); present as column-major (K, N). + b_shape = (batch, k, n) + b_stride = (k * n, 1, k) + b_descale_shape, b_descale_stride, b_descale_reordering = ( + _bf16_fp4_b_descale_layout(batch, n, k, block_size) + ) + + stream = torch.cuda.current_stream(device) + with cudnn.graph(_get_cudnn_handle(device, stream)) as (graph, _): + a_cudnn_tensor = graph.tensor( + name="a", dim=a_shape, stride=a_stride, data_type=a_type + ) + b_cudnn_tensor = graph.tensor( + name="b", dim=b_shape, stride=b_stride, data_type=cudnn.data_type.FP4_E2M1 + ) + block_descale_b_cudnn_tensor = graph.tensor( + name="block_descale_b", + dim=b_descale_shape, + stride=b_descale_stride, + data_type=scale_type, + reordering_type=b_descale_reordering, + ) + + _build_bf16_fp4_graph_common( + graph, + a_cudnn_tensor, + b_cudnn_tensor, + block_descale_b_cudnn_tensor, + a_type, + o_type, + block_size, + alpha_is_not_none, + ) + + graph.validate() + graph.build_operation_graph() + graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + graph.check_support() + graph.build_plans(policy) + return graph + + +@functools.lru_cache(maxsize=1024) +def build_cudnn_bf16_fp4_graph_override_shape( + batch, + n, + k, + a_type, + o_type, + block_size, + device, + alpha_is_not_none, + use_nvfp4, + cache_m: int = _OVERRIDE_SHAPE_CACHE_M, + policy=None, +): + """Build a cuDNN bf16 x fp4 GEMM graph with override-shape support.""" + _check_cudnn_fp4_availability() + if policy is None: + policy = cudnn.build_plan_policy.HEURISTICS_CHOICE + + scale_type = cudnn.data_type.FP8_E4M3 if use_nvfp4 else cudnn.data_type.FP8_E8M0 + + a_shape = [batch, cache_m, k] + a_stride = [cache_m * k, k, 1] + b_shape = [batch, k, n] + b_stride = [k * n, 1, k] + b_descale_shape, b_descale_stride, b_descale_reordering = ( + _bf16_fp4_b_descale_layout(batch, n, k, block_size) + ) + + stream = torch.cuda.current_stream(device) + graph = cudnn.pygraph( + io_data_type=cudnn.data_type.FLOAT, + intermediate_data_type=cudnn.data_type.FLOAT, + compute_data_type=cudnn.data_type.FLOAT, + handle=_get_cudnn_handle(device, stream), + is_override_shape_enabled=True, + ) + + a_cudnn_tensor = graph.tensor( + name="a", dim=a_shape, stride=a_stride, data_type=a_type + ) + b_cudnn_tensor = graph.tensor( + name="b", dim=b_shape, stride=b_stride, data_type=cudnn.data_type.FP4_E2M1 + ) + block_descale_b_cudnn_tensor = graph.tensor( + name="block_descale_b", + dim=b_descale_shape, + stride=b_descale_stride, + data_type=scale_type, + reordering_type=b_descale_reordering, + ) + + _build_bf16_fp4_graph_common( + graph, + a_cudnn_tensor, + b_cudnn_tensor, + block_descale_b_cudnn_tensor, + a_type, + o_type, + block_size, + alpha_is_not_none, + ) + + graph.validate() + graph.build_operation_graph() + graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + graph.check_support() + graph.build_plans(policy) + return graph + + +def _bf16_fp4_variant_pack(a, b, b_descale, alpha, out): + """Build the {uid: tensor} variant pack shared by both execute paths.""" + variant_pack = { + UIDs.A_UID.value: a, + UIDs.B_UID.value: b.view(get_native_fp4_dtype()), + UIDs.BLOCK_DESCALE_B_UID.value: b_descale, + UIDs.O_UID.value: out, + } + if alpha is not None: + variant_pack[UIDs.ALPHA_UID.value] = alpha.view(torch.float) + return variant_pack + + +def execute_cudnn_bf16_fp4_graph( + graph, + a, + b, + b_descale, + alpha, + out, + workspace_buffer, + tactic: int = -1, +): + variant_pack = _bf16_fp4_variant_pack(a, b, b_descale, alpha, out) + + workspace_size = _get_cudnn_workspace_size(graph, tactic) + if workspace_buffer.numel() < workspace_size: + workspace_buffer.resize_(workspace_size) + + stream = torch.cuda.current_stream(a.device) + handle = _get_cudnn_handle(a.device, stream) + if tactic == -1: + graph.execute(variant_pack, workspace_buffer, handle=handle) + else: + graph.execute_plan_at_index( + variant_pack, workspace_buffer, tactic, handle=handle + ) + + +def execute_cudnn_bf16_fp4_graph_override_shape( + graph, + a, + b, + b_descale, + alpha, + out, + workspace_buffer, + block_size: int = 16, + tactic: int = 0, +): + """Execute the bf16 x fp4 graph, overriding A / output to the real M.""" + m, k = int(a.shape[0]), int(a.shape[1]) + n = int(b.shape[0]) + batch = 1 + + a_shape = (batch, m, k) + a_stride = (m * k, k, 1) + b_shape = (batch, k, n) + b_stride = (k * n, 1, k) + b_descale_shape, b_descale_stride, _ = _bf16_fp4_b_descale_layout( + batch, n, k, block_size + ) + out_shape = (batch, m, n) + out_stride = (m * n, n, 1) + + variant_pack = _bf16_fp4_variant_pack(a, b, b_descale, alpha, out) + + override_uids = [ + UIDs.A_UID.value, + UIDs.B_UID.value, + UIDs.BLOCK_DESCALE_B_UID.value, + UIDs.O_UID.value, + ] + override_shapes = [a_shape, b_shape, b_descale_shape, out_shape] + override_strides = [a_stride, b_stride, b_descale_stride, out_stride] + + stream = torch.cuda.current_stream(a.device) + cudnn_handle = _get_cudnn_handle(a.device, stream) + + workspace_size = _get_cudnn_override_shape_workspace_size( + graph, tactic, cudnn_handle, override_uids, override_shapes, override_strides + ) + if workspace_buffer.numel() < workspace_size: + workspace_buffer.resize_(workspace_size) + + graph.execute_plan_at_index( + variant_pack, + workspace_buffer, + tactic, + handle=cudnn_handle, + override_uids=override_uids, + override_shapes=override_shapes, + override_strides=override_strides, + ) + + +# Autotuner sweeps M (token count) of the bf16 activation ``a`` +_BF16_FP4_TUNING_CONFIG = TuningConfig( + dynamic_tensor_specs=( + DynamicTensorSpec( + (0,), # a_tensor_index + (0,), # M dimension + get_hybrid_num_tokens_buckets, + map_to_hybrid_bucket_uncapped, + ), + ), + constraint_specs=( + ConstraintSpec( + 5, # out_tensor_index follows M + 0, + lambda shapes: shapes[0][0], + ), + ), +) + + +def _cudnn_bf16_fp4_runner(tuning_config): + """Build a ``CudnnBf16Fp4Runner`` bound to the active tuning config.""" + m_bucket_mapper = AutoTuner.get().get_effective_map_to_tuning_buckets( + tuning_config, spec_idx=0 + ) + + class CudnnBf16Fp4Runner(TunableRunner): + def __init__(self): + super().__init__() + self._m_bucket_mapper = m_bucket_mapper + self._use_override_shape = is_cudnn_override_shape_available() + + def get_cache_key_extras(self, inputs: List[torch.Tensor]) -> tuple: + _, _, _, alpha, out_dtype, _, block_size, use_nvfp4, _ = inputs + return (out_dtype, block_size, use_nvfp4, alpha is not None) + + def _get_override_graph(self, a, b, alpha, out_dtype, block_size, use_nvfp4): + actual_m, k = int(a.shape[0]), int(a.shape[1]) + n = int(b.shape[0]) + cache_m = self._m_bucket_mapper(actual_m) + return build_cudnn_bf16_fp4_graph_override_shape( + batch=1, + n=n, + k=k, + a_type=_torch_data_type_to_cudnn_data_type(a.dtype), + o_type=_torch_data_type_to_cudnn_data_type(out_dtype), + block_size=block_size, + device=a.device, + alpha_is_not_none=alpha is not None, + use_nvfp4=use_nvfp4, + cache_m=cache_m, + policy=cudnn.build_plan_policy.ALL, + ) + + def get_valid_tactics( + self, + inputs: List[torch.Tensor], + profile: OptimizationProfile, + ) -> List[int]: + ( + a, + b, + b_descale, + alpha, + out_dtype, + out, + block_size, + use_nvfp4, + workspace_buffer, + ) = inputs + if self._use_override_shape: + graph = self._get_override_graph( + a, b, alpha, out_dtype, block_size, use_nvfp4 + ) + else: + graph = build_cudnn_bf16_fp4_graph( + batch=1, + m=int(a.shape[0]), + n=int(b.shape[0]), + k=int(a.shape[1]), + a_type=_torch_data_type_to_cudnn_data_type(a.dtype), + o_type=_torch_data_type_to_cudnn_data_type(out_dtype), + block_size=block_size, + device=a.device, + alpha_is_not_none=alpha is not None, + use_nvfp4=use_nvfp4, + policy=cudnn.build_plan_policy.HEURISTICS_CHOICE, + ) + return list(range(graph.get_execution_plan_count())) + + def forward( + self, + inputs: List[torch.Tensor], + tactic: int = -1, + do_preparation: bool = False, + **kwargs, + ) -> torch.Tensor: + ( + a, + b, + b_descale, + alpha, + out_dtype, + out, + block_size, + use_nvfp4, + workspace_buffer, + ) = inputs + if self._use_override_shape: + graph = self._get_override_graph( + a, b, alpha, out_dtype, block_size, use_nvfp4 + ) + execute_cudnn_bf16_fp4_graph_override_shape( + graph, + a, + b, + b_descale, + alpha, + out, + workspace_buffer, + block_size=block_size, + tactic=max(tactic, 0), + ) + else: + graph = build_cudnn_bf16_fp4_graph( + batch=1, + m=int(a.shape[0]), + n=int(b.shape[0]), + k=int(a.shape[1]), + a_type=_torch_data_type_to_cudnn_data_type(a.dtype), + o_type=_torch_data_type_to_cudnn_data_type(out_dtype), + block_size=block_size, + device=a.device, + alpha_is_not_none=alpha is not None, + use_nvfp4=use_nvfp4, + policy=cudnn.build_plan_policy.HEURISTICS_CHOICE, + ) + execute_cudnn_bf16_fp4_graph( + graph, + a, + b, + b_descale, + alpha, + out, + workspace_buffer, + tactic=-1, + ) + return out + + return CudnnBf16Fp4Runner() + + +def _prepare_cudnn( + b: torch.Tensor, + b_descale: torch.Tensor, + alpha: Optional[torch.Tensor], + block_size: int, +) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """cuDNN-backend prep. + + The weight bytes ``(N, K//2)`` are already in the layout the cuDNN graph + consumes (see the module banner). The scale factor, however, must be + *non-swizzled* for the cuDNN bf16 x fp4 path (cuDNN does not support the + 128x4-swizzled SF layout), so we unswizzle the canonical 128x4 SF into a + linear ``(N, K // block_size)`` FP8-E4M3 tensor. + """ + n = int(b.shape[0]) + k = int(b.shape[1]) * 2 + k_sf = k // block_size + # (N, K_sf) uint8 bytes, each byte an FP8-E4M3 per-block scale. + linear_sf = _unswizzle_sf_128x4(b_descale, n, k_sf).contiguous() + return b, linear_sf.view(torch.float8_e4m3fn), alpha + + +def _compute_cudnn( + a: torch.Tensor, + b: torch.Tensor, + b_descale: torch.Tensor, + alpha: Optional[torch.Tensor], + out_dtype: torch.dtype, + out: Optional[torch.Tensor], + block_size: int, +) -> torch.Tensor: + """cuDNN-backend compute with autotuning over the M (token) dimension.""" + n = int(b.shape[0]) + k = int(b.shape[1]) * 2 + if a.shape[1] != k: + raise ValueError( + f"a.shape[1]={a.shape[1]} but k inferred from b.shape={tuple(b.shape)} " + f"is {k}" + ) + + if out is None: + out = torch.empty((a.shape[0], n), device=a.device, dtype=out_dtype) + else: + if tuple(out.shape) != (a.shape[0], n): + raise ValueError( + f"out shape {tuple(out.shape)} != expected {(a.shape[0], n)}" + ) + if out.dtype != out_dtype: + raise TypeError(f"out dtype {out.dtype} != requested out_dtype {out_dtype}") + + workspace_buffer = _get_cache_buf( + "mm_bf16_fp4_workspace", DEFAULT_WORKSPACE_SIZE, a.device + ) + + tuning_config = _BF16_FP4_TUNING_CONFIG + tuner = AutoTuner.get() + runner = _cudnn_bf16_fp4_runner(tuning_config) + + use_nvfp4 = True + inputs = [ + a, + b, + b_descale, + alpha, + out_dtype, + out, + block_size, + use_nvfp4, + workspace_buffer, + ] + chosen_runner, tactic = tuner.choose_one( + "bf16_fp4_gemm", + [runner], + tuning_config, + inputs, + ) + chosen_runner(inputs=inputs, tactic=tactic) + return out diff --git a/flashinfer/gemm/gemm_bf16_fp4_cute_dsl.py b/flashinfer/gemm/gemm_bf16_fp4_cute_dsl.py new file mode 100644 index 00000000000..51bb135e96f --- /dev/null +++ b/flashinfer/gemm/gemm_bf16_fp4_cute_dsl.py @@ -0,0 +1,486 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 by FlashInfer team. +# SPDX-License-Identifier: Apache-2.0 +"""CuTe-DSL backend for the bf16 x fp4 GEMM (weight repack / kernel launch).""" + +from typing import List, Optional, Tuple, cast + +import torch + +from ..autotuner import ( + AutoTuner, + ConstraintSpec, + DynamicTensorSpec, + OptimizationProfile, + TunableRunner, + TuningConfig, +) +from ..fused_moe.utils import ( + get_hybrid_num_tokens_buckets, + map_to_hybrid_bucket_uncapped, +) +from .gemm_base import _TORCH_TO_CUTLASS_DTYPE_ATTR, _check_cute_dsl_availability +from .gemm_bf16_fp4 import _unswizzle_sf_128x4 + +_BF16_FP4_ALPHA_ONE_CACHE: dict = {} + + +def _prepare_bf16_fp4_alpha( + alpha: Optional[torch.Tensor], device: torch.device +) -> torch.Tensor: + """Normalize ``alpha`` to a ``(1,) float32`` tensor for the kernel.""" + if alpha is None: + cached = _BF16_FP4_ALPHA_ONE_CACHE.get(device) + if cached is None: + cached = torch.tensor([1.0], dtype=torch.float32, device=device) + _BF16_FP4_ALPHA_ONE_CACHE[device] = cached + return cached + if alpha.dim() == 0: + return alpha.to(device=device, dtype=torch.float32).unsqueeze(0) + return alpha.to(device=device, dtype=torch.float32).reshape(1) + + +def _select_bf16_fp4_tile_shape( + m: int, n: int, k: int +) -> Tuple[Tuple[int, int, int], Tuple[int, int, int]]: + """Pick a CTA tile shape AND MMA atom_layout for the cute-DSL bf16 x fp4 kernel. + + Returns ``(tile_shape_mnk, atom_layout)``. + + Tile shape selection: + tile_M choice + * M <= 16 (and tile_K=128 path): use tile_M=16 with atom_layout + (1,2,1). Halves wasted M-rows vs tile_M=32, and a 1-M-warp + layout removes the duplicate dequant that (2,2,1) suffers from. + * 16 < M <= 32: use tile_M=32 with atom_layout (2,2,1). Smaller + MMA + epilogue waste than tile_M=64. + * M > 32: use tile_M=64 with atom_layout (2,2,1) -- standard tile, + more rows to amortize across. + + tile_K choice + * K % 128 == 0: tile_K=128 (halves K-tile count and barrier + overhead). + * Otherwise: tile_K=64. + + Why atom_layout differs: + * (2,2,1) (default for tile_M >= 32): 4 MMA warps as 2 M x 2 N -- + well-tested cute layout, but the 2 M-warps redundantly dequant + the same B values into their own register files (~50% waste in + dequant compute). + * (1,2,1) (used for tile_M=16): 2 MMA warps as 1 M x 2 N -- no + M-warp duplication. Permutation_m = 16, so tile_M must be 16. + """ + tile_k = 128 if k % 128 == 0 else 64 + if m <= 16 and tile_k == 128: + return ((16, 64, 128), (1, 2, 1)) + if m <= 32: + return ((32, 64, tile_k), (2, 2, 1)) + return ((64, 64, tile_k), (2, 2, 1)) + + +_CUTE_DSL_MM_BF16_FP4_KERNEL_CACHE: dict = {} + + +def _get_cute_dsl_bf16_fp4_gemm( + tile_shape_mnk: Tuple[int, int, int], + a_dtype: torch.dtype, + c_dtype: torch.dtype, + atom_layout: Tuple[int, int, int] = (2, 2, 1), + pipeline_depth: int = 1, + use_fp16_mma: int = 1, + enable_pdl: bool = True, + tile_swizzle: int = 1, +): + # Normalize to a tuple (callers may pass a list) so the cache key is hashable. + atom_layout = cast(Tuple[int, int, int], tuple(atom_layout)) + pipeline_depth = int(pipeline_depth) + use_fp16_mma = int(use_fp16_mma) + enable_pdl = bool(enable_pdl) + tile_swizzle = int(tile_swizzle) + cache_key = ( + tile_shape_mnk, + a_dtype, + c_dtype, + atom_layout, + pipeline_depth, + use_fp16_mma, + enable_pdl, + tile_swizzle, + ) + cached = _CUTE_DSL_MM_BF16_FP4_KERNEL_CACHE.get(cache_key) + if cached is not None: + return cached + + _check_cute_dsl_availability() + + import cutlass + import cutlass.cute as cute + from flashinfer.cute_dsl.utils import get_max_active_clusters + + from .kernels.cute_dsl.dense_gemm_bf16_fp4_blackwell import ( + BlackwellDenseGemmBf16Fp4Kernel, + ) + + a_cutlass_dtype = getattr(cutlass, _TORCH_TO_CUTLASS_DTYPE_ATTR[a_dtype]) + c_cutlass_dtype = getattr(cutlass, _TORCH_TO_CUTLASS_DTYPE_ATTR[c_dtype]) + + sym_m = cute.sym_int() + sym_k = cute.sym_int() + sym_n = cute.sym_int() + sym_k_tiles = cute.sym_int() + sym_n_packed = cute.sym_int() + + a_fake = cute.runtime.make_fake_compact_tensor( + a_cutlass_dtype, (sym_m, sym_k), stride_order=(1, 0), assumed_align=16 + ) + b_packed_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, + (sym_k_tiles, sym_n_packed), + stride_order=(1, 0), + assumed_align=16, + ) + b_sf_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Uint8, (sym_k_tiles, sym_n), stride_order=(1, 0), assumed_align=16 + ) + c_fake = cute.runtime.make_fake_compact_tensor( + c_cutlass_dtype, (sym_m, sym_n), stride_order=(1, 0), assumed_align=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) + + gemm = BlackwellDenseGemmBf16Fp4Kernel( + acc_dtype=cutlass.Float32, + tile_shape_mnk=tile_shape_mnk, + atom_layout=atom_layout, + pipeline_depth=pipeline_depth, + use_fp16_mma=use_fp16_mma, + enable_pdl=enable_pdl, + tile_swizzle=tile_swizzle, + ) + max_active_clusters = get_max_active_clusters(1) + + compiled = cute.compile( + gemm.wrapper, + a_fake, + b_packed_fake, + b_sf_fake, + c_fake, + alpha_fake, + 1, # l (batch) + max_active_clusters, + stream_fake, + options="--opt-level 2 --enable-tvm-ffi", + ) + + _CUTE_DSL_MM_BF16_FP4_KERNEL_CACHE[cache_key] = compiled + return compiled + + +def _e4m3_to_s0e5m3(sf_u8: torch.Tensor) -> torch.Tensor: + """Reformat a uint8 tensor of E4M3 scale bytes to S0E5M3 bytes. + Used in cute-dsl backend for faster in-kernel scale decode. + """ + f16 = sf_u8.contiguous().view(torch.float8_e4m3fn).to(torch.float16) + bits = f16.view(torch.int16).to(torch.int32) & 0xFFFF + return ((bits >> 7) & 0xFF).to(torch.uint8) + + +_CUTE_DSL_PACK_TILE_K: int = 16 # K-tile size = MMA K-block size +_CUTE_DSL_PACK_TILE_N: int = 64 # N-tile size = kernel tile_N +_CUTE_DSL_PACK_INTS_PER_TILE: int = 128 # int32s per (16K x 64N) repack block + + +def _cute_dsl_pack_fp4_weight(b: torch.Tensor) -> torch.Tensor: + """Repack a packed FP4 weight for the bf16 x fp4 cute-DSL kernel.""" + if b.dtype != torch.uint8: + b = b.view(torch.uint8) + + k_half, n = b.shape + k = k_half * 2 + if k % _CUTE_DSL_PACK_TILE_K != 0: + raise ValueError(f"K must be a multiple of {_CUTE_DSL_PACK_TILE_K} (got K={k})") + if n % _CUTE_DSL_PACK_TILE_N != 0: + raise ValueError(f"N must be a multiple of {_CUTE_DSL_PACK_TILE_N} (got N={n})") + + device = b.device + k_tiles = k // _CUTE_DSL_PACK_TILE_K + n_tiles = n // _CUTE_DSL_PACK_TILE_N + k_half_per_tile = _CUTE_DSL_PACK_TILE_K // 2 # 8 packed K-rows per tile + + u32_pos = torch.arange( + _CUTE_DSL_PACK_INTS_PER_TILE, device=device, dtype=torch.long + ) + u32_idx_local = u32_pos % 2 + lane = (u32_pos // 2) % 32 + n_warp_idx = u32_pos // 64 + + tc_col = lane // 4 # in [0, 8) + tc_row_half = lane % 4 # tc_row = tc_row_half * 2 in {0, 2, 4, 6} + base_n = n_warp_idx * 8 + tc_col # in [0, 16) + + byte_k_half_offset = torch.tensor([0, 4, 0, 4], device=device, dtype=torch.long) + n_offset_stack = torch.tensor( + [[0, 0, 16, 16], [32, 32, 48, 48]], device=device, dtype=torch.long + ) + byte_n_offset = n_offset_stack[u32_idx_local] # (128, 4) + + # Source byte within the (8, 64) tile for each (u32_pos, byte_idx). + k_half_in_tile = tc_row_half[:, None] + byte_k_half_offset[None, :] # (128, 4) + n_in_tile = base_n[:, None] + byte_n_offset # (128, 4) + within_idx = (k_half_in_tile * _CUTE_DSL_PACK_TILE_N + n_in_tile).reshape( + -1 + ) # (512,) flat index into a row-major (8, 64) tile + + # (K/2, N) -> (K_tiles, 8, N_tiles, 64) -> (K_tiles, N_tiles, 8*64) so the + # 512 source bytes of each tile are contiguous, then gather them in + # (u32_pos, byte_idx) order. + tile_bytes = ( + b.reshape(k_tiles, k_half_per_tile, n_tiles, _CUTE_DSL_PACK_TILE_N) + .permute(0, 2, 1, 3) + .reshape(k_tiles, n_tiles, k_half_per_tile * _CUTE_DSL_PACK_TILE_N) + ) + gathered = tile_bytes[:, :, within_idx].reshape( + k_tiles, n_tiles, _CUTE_DSL_PACK_INTS_PER_TILE, 4 + ) + + # Each 4 consecutive bytes are one little-endian int32 (byte 0 = bits 0-7), + # exactly what the kernel's 32-bit loads read -- reinterpret in place. + return gathered.view(torch.int32).reshape( + k_tiles, n_tiles * _CUTE_DSL_PACK_INTS_PER_TILE + ) + + +def _prepare_cute_dsl( + b: torch.Tensor, + b_descale: torch.Tensor, + alpha: Optional[torch.Tensor], + block_size: int, +) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """cute-DSL-backend prep: repack the weight + unswizzle the SF. + + Produces the bespoke layout the cute-DSL kernel consumes: + * weight: ``(K // 16, N * 2)`` int32 (see :func:`_cute_dsl_pack_fp4_weight`). + * SF: ``(K // block_size, N)`` uint8 -- per-block scales reformatted to + S0E5M3, the format the cute-DSL kernel decodes. + ``alpha`` is passed through unchanged (the compute step normalizes it + to a ``(1,) float32`` scalar). Pair the returned tensors with + ``mm_bf16_fp4(a, b, b_descale, alpha, backend='cute-dsl')``. + """ + n = int(b.shape[0]) + k = int(b.shape[1]) * 2 + k_sf = k // block_size + + b_kn = b.t().contiguous() + b_packed = _cute_dsl_pack_fp4_weight(b_kn) # (K//16, N*2) int32 + + linear_sf = _unswizzle_sf_128x4(b_descale, n, k_sf) # (N, K_sf) uint8 + sf_ksf_n = linear_sf.t().contiguous() # (K_sf, N) uint8 (E4M3) + sf_ksf_n = _e4m3_to_s0e5m3(sf_ksf_n) # -> S0E5M3 + return b_packed, sf_ksf_n, alpha + + +def _bf16_fp4_cute_dsl_tactic_configs( + n: int, k: int +) -> List[Tuple[Tuple[int, int, int], Tuple[int, int, int], int, int, int]]: + """Enumerate cute-DSL tactic configs for a given ``(N, K)``. + + Returns a list of ``(tile_shape_mnk, atom_layout, pipeline_depth, + use_fp16_mma)`` tuples. + """ + tile_k = 128 if k % 128 == 0 else 64 + + # (tile_M, atom_layout) shapes the kernel is designed/validated for, at the + # default tile_N=64; a tile_N=128 variant is added below for very large N. + tile_m_atoms: List[Tuple[int, Tuple[int, int, int]]] = [] + if tile_k == 128: + tile_m_atoms.append((16, (1, 2, 1))) + tile_m_atoms.append((32, (2, 2, 1))) + tile_m_atoms.append((64, (2, 2, 1))) + + configs: List[Tuple[Tuple[int, int, int], Tuple[int, int, int], int, int, int]] = [] + seen = set() + + def add(tile_m, atom, pdepth, fp16, tile_n=64, tk=None, swz=1): + cfg = ( + (tile_m, tile_n, tile_k if tk is None else tk), + atom, + pdepth, + fp16, + swz, + ) + key = (cfg[0], cfg[1], pdepth, fp16, swz) + if key not in seen: + seen.add(key) + configs.append(cfg) + + base_tile_m, base_atom = tile_m_atoms[0] + add(base_tile_m, base_atom, 1, 1) # 0: baseline + add(base_tile_m, base_atom, 0, 1) # no dequant prefetch (helps short-K) + for tile_m, atom in tile_m_atoms[1:]: + add(tile_m, atom, 1, 1) + + # tile_N=128 halves the (m,n)-tile count but needs large wave count. + if tile_k == 128 and n >= 12288 and n % 128 == 0: + add(base_tile_m, base_atom, 1, 1, tile_n=128) + + # tile_K=64 has more ab stages, but requires larger problem size. + if tile_k == 128 and n >= 8192: + add(base_tile_m, base_atom, 1, 1, tile_n=64, tk=64) + + # tile_M=128 (taller M tile, atom (2,2,1)) -- the large-M *prefill* lever. + if tile_k == 128: + add(128, (2, 2, 1), 1, 1) + + # Threadblock swizzle (tile_swizzle=8) -- for large-M prefill. + if tile_k == 128 and n * k >= 16 * 1024 * 1024: + add(64, (2, 2, 1), 1, 1, swz=8) + if tile_k == 128: + add(128, (2, 2, 1), 1, 1, swz=8) + + # tile_N=128 (with tile_M=64, atom (2,2,1)) -- large shapes. + if tile_k == 128 and n % 128 == 0 and n >= 4096: + add(64, (2, 2, 1), 1, 1, tile_n=128, swz=8) + add(64, (2, 2, 1), 1, 1, tile_n=128, swz=1) + + return configs + + +_BF16_FP4_CUTE_DSL_TUNING_CONFIG = TuningConfig( + dynamic_tensor_specs=( + DynamicTensorSpec( + (0,), # a_tensor_index + (0,), # M dimension + get_hybrid_num_tokens_buckets, + map_to_hybrid_bucket_uncapped, + ), + ), + constraint_specs=( + ConstraintSpec( + 5, # out_tensor_index follows M + 0, + lambda shapes: shapes[0][0], + ), + ), +) + + +def _cute_dsl_bf16_fp4_runner(enable_pdl: bool = True) -> TunableRunner: + """Build a ``CuteDslBf16Fp4Runner`` for the cute-DSL bf16 x fp4 GEMM.""" + + class CuteDslBf16Fp4Runner(TunableRunner): + def get_cache_key_extras(self, inputs: List[torch.Tensor]) -> tuple: + a, b, _, _, out_dtype, _, block_size = inputs + n = int(b.shape[1]) // 2 + k = int(b.shape[0]) * int(block_size) + return (out_dtype, n, k) + + def get_valid_tactics( + self, + inputs: List[torch.Tensor], + profile: OptimizationProfile, + ) -> List[int]: + _, b, _, _, _, _, block_size = inputs + n = int(b.shape[1]) // 2 + k = int(b.shape[0]) * int(block_size) + return list(range(len(_bf16_fp4_cute_dsl_tactic_configs(n, k)))) + + def forward( + self, + inputs: List[torch.Tensor], + tactic: int = -1, + do_preparation: bool = False, + **kwargs, + ) -> torch.Tensor: + a, b, b_sf_u8, alpha_for_launch, out_dtype, out, block_size = inputs + n = int(b.shape[1]) // 2 + k = int(b.shape[0]) * int(block_size) + m = int(a.shape[0]) + if tactic < 0: + # Fallback == pre-autotuner heuristic (M-aware), default knobs. + tile_shape_mnk, atom_layout = _select_bf16_fp4_tile_shape(m, n, k) + pipeline_depth, use_fp16_mma, tile_swizzle = 1, 1, 1 + else: + ( + tile_shape_mnk, + atom_layout, + pipeline_depth, + use_fp16_mma, + tile_swizzle, + ) = _bf16_fp4_cute_dsl_tactic_configs(n, k)[tactic] + compiled = _get_cute_dsl_bf16_fp4_gemm( + tile_shape_mnk, + a.dtype, + out_dtype, + atom_layout, + pipeline_depth, + use_fp16_mma, + enable_pdl=enable_pdl, + tile_swizzle=tile_swizzle, + ) + compiled(a, b, b_sf_u8, out, alpha_for_launch) + return out + + return CuteDslBf16Fp4Runner() + + +def _compute_cute_dsl( + a: torch.Tensor, + b: torch.Tensor, + b_descale: torch.Tensor, + alpha: Optional[torch.Tensor], + out_dtype: torch.dtype, + out: Optional[torch.Tensor], + block_size: int, + enable_pdl: bool = True, +) -> torch.Tensor: + """cute-DSL-backend compute: dispatch to the compiled Blackwell kernel. + + ``b`` is the packed ``(K // 16, N * 2)`` int32 weight and + ``b_descale`` the ``(K // block_size, N)`` uint8 SF in S0E5M3 format + (reformatted from FP8-E4M3 by :func:`_e4m3_to_s0e5m3`) returned by + :func:`_prepare_cute_dsl`. + """ + if b.dtype != torch.int32: + raise TypeError( + f"cute-dsl backend expects the packed int32 weight from " + f"prepare_bf16_fp4_weights(..., backend='cute-dsl'); got {b.dtype}." + ) + if out_dtype != a.dtype: + raise NotImplementedError( + f"cute-dsl backend requires out_dtype == a.dtype (got " + f"out_dtype={out_dtype}, a.dtype={a.dtype}). Use the cudnn " + f"backend for a mismatched output dtype." + ) + k_tiles = int(b.shape[0]) + n = int(b.shape[1]) // 2 + k = k_tiles * block_size + m = int(a.shape[0]) + if a.shape[1] != k: + raise ValueError( + f"a.shape[1]={a.shape[1]} but k inferred from prepared b.shape=" + f"{tuple(b.shape)} is {k}" + ) + + if out is None: + out = torch.empty((m, n), device=a.device, dtype=out_dtype) + else: + if tuple(out.shape) != (m, n): + raise ValueError(f"out shape {tuple(out.shape)} != expected {(m, n)}") + if out.dtype != out_dtype: + raise TypeError(f"out dtype {out.dtype} != requested out_dtype {out_dtype}") + + b_sf_u8 = b_descale.view(torch.uint8).contiguous() + alpha_for_launch = _prepare_bf16_fp4_alpha(alpha, a.device) + + tuner = AutoTuner.get() + runner = _cute_dsl_bf16_fp4_runner(enable_pdl=enable_pdl) + inputs = [a, b, b_sf_u8, alpha_for_launch, out_dtype, out, block_size] + chosen_runner, tactic = tuner.choose_one( + "bf16_fp4_cute_dsl_gemm", + [runner], + _BF16_FP4_CUTE_DSL_TUNING_CONFIG, + inputs, + ) + chosen_runner(inputs=inputs, tactic=tactic) + return out diff --git a/flashinfer/gemm/kernels/cute_dsl/__init__.py b/flashinfer/gemm/kernels/cute_dsl/__init__.py new file mode 100644 index 00000000000..84c5c70d321 --- /dev/null +++ b/flashinfer/gemm/kernels/cute_dsl/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2025 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/flashinfer/gemm/kernels/cute_dsl/dense_gemm_bf16_fp4_blackwell.py b/flashinfer/gemm/kernels/cute_dsl/dense_gemm_bf16_fp4_blackwell.py new file mode 100644 index 00000000000..c51a13c5aaa --- /dev/null +++ b/flashinfer/gemm/kernels/cute_dsl/dense_gemm_bf16_fp4_blackwell.py @@ -0,0 +1,1464 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +"""bf16 x fp4 dense GEMM for Blackwell (SM100/103/120/121). + +Built on top of ``dense_gemm_bf16_blackwell.py`` from cutlass examples at +https://github.com/NVIDIA/cutlass/blob/main/examples/python/CuTeDSL/cute/blackwell_geforce/kernel/dense_gemm/dense_gemm.py +""" + +from typing import Optional, Tuple, Type + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import cutlass.utils.hopper_helpers as sm90_utils +from cutlass import Float32, Int32, Uint32 + +from ....cute_dsl.fp4_common import ( + cvt_s0e5m3_to_f16x2_broadcast, + f16x2_to_f32x2, + fp4_decode_4bytes, + get_smem_ptr_as_int32, + half2_mul, + ld_shared_v2_u32, +) +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm +from cutlass._mlir.extras import types as T +from cutlass.cutlass_dsl import dsl_user_op + + +@dsl_user_op +def cvt_bf16x2_to_f16x2_via_f32(packed_bf16x2: Uint32, *, loc=None, ip=None) -> Uint32: + """Packed bf16x2 (u32) -> f16x2 (u32) via f32 intermediate. + + Used by the fp16-MMA path (use_fp16_mma=1) to convert A's ldmatrix + output (bf16 bit pattern from sA) to fp16 bit pattern for the fp16 + MMA inputs. Direct ``cvt.rn.f16x2.bf16x2`` is rejected by ptxas on + sm_100a / CUDA 13.1; this packed via-f32 path is 4 PTX instrs per + pair (mov.b32 unpack + 2x cvt.f32.bf16 + 1x cvt.rn.f16x2.f32). + """ + return Uint32( + llvm.inline_asm( + T.i32(), + [Uint32(packed_bf16x2).ir_value(loc=loc, ip=ip)], + """ + { + .reg .b16 b_lo, b_hi; + .reg .f32 f_lo, f_hi; + mov.b32 {b_lo, b_hi}, $1; + cvt.f32.bf16 f_lo, b_lo; + cvt.f32.bf16 f_hi, b_hi; + cvt.rn.f16x2.f32 $0, f_hi, f_lo; + } + """, + "=r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def f16x2_unpack( + packed_h2: Uint32, *, loc=None, ip=None +) -> Tuple["cutlass.Float16", "cutlass.Float16"]: + """Unpack f16x2 (u32) into (f16_lo, f16_hi). Free at HW level -- + just a register-rename (mov.b32 {h_lo, h_hi}, packed). Used by the + fp16-MMA path to bypass the f16->f32->bf16 cvt chain that the + default bf16-MMA path needs after hmul2.""" + from cutlass import Float16 + + res = llvm.inline_asm( + ir.Type.parse("!llvm.struct<(f16, f16)>"), + [Uint32(packed_h2).ir_value(loc=loc, ip=ip)], + "mov.b32 {$0, $1}, $2;", + "=h,=h,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return ( + Float16(llvm.extractvalue(T.f16(), res, [0], loc=loc, ip=ip)), + Float16(llvm.extractvalue(T.f16(), res, [1], loc=loc, ip=ip)), + ) + + +# FP4 weight packing constants (must match _cute_dsl_pack_fp4_weight in +# flashinfer/gemm/gemm_bf16_fp4_cute_dsl.py). +_PACK_TILE_K: cutlass.Constexpr = 16 +_PACK_TILE_N: cutlass.Constexpr = 64 +_PACK_INTS_PER_TILE: cutlass.Constexpr = 128 # 128 int32 per (16K x 64N) block + + +class BlackwellDenseGemmBf16Fp4Kernel: + """Warp-MMA dense GEMM for Blackwell, FP4-weight A bf16/fp16 input. + + A: (M, K, L) bf16/fp16. + B: (K // 16, N * 2, L) int32 -- packed FP4 (see prepare). + B_sf: (K // 16, N, L) uint8 -- per-group scales (S0E5M3). + alpha: (1,) fp32 -- global scalar scale. + C: (M, N, L) bf16/fp16; fp32 accumulator cast at write. + + """ + + GROUP_SIZE: cutlass.Constexpr = 16 + + def __init__( + self, + acc_dtype, + tile_shape_mnk, + epi_stage: int = 4, + pipeline_depth: int = 1, + atom_layout: Tuple[int, int, int] = (2, 2, 1), + epi_tile_override: Optional[Tuple[int, int]] = None, + # 1 = fp16 MMA (default): MmaF16BF16Op uses Float16. Dequant writes + # fp16 directly to tCrB (skipping the f16->f32->bf16 cvt chain + # the bf16 path needs after hmul2). A is bf16 in SMEM -> + # ldmatrix into bf16 staging fragment -> in-register packed + # bf16->fp16 cvt -> fp16 tCrA. Slightly *more* accurate than + # bf16 MMA for well-behaved inputs since fp16's 10-bit mantissa + # beats bf16's 7-bit at the multiply step; accumulator stays + # fp32 in both modes. + # 0 = bf16 MMA: original path. Safer for workloads with very + # large activation magnitudes (|A| > ~30000) since bf16's + # wider exponent range avoids saturation in the A cvt. + use_fp16_mma: int = 1, + enable_pdl: bool = True, + tile_swizzle: int = 1, + raster_along_m: bool = True, + ): + """bf16 x fp4 kernel. + + Args: + acc_dtype: accumulator dtype (always Float32 for this kernel). + tile_shape_mnk: CTA tile shape. + epi_stage: TMA-store pipeline depth. Default 4 balances + cross-tile overlap (epi_stage > 1 lets next-tile compute + overlap with this-tile store) against SMEM available for + ab_stage (deeper ab_stage hides TMA-load latency, which + dominates the small-M shape's stalls). At M=4 with the + default 64-CTA grid on 148 SMs, each SM gets at most + one tile, so epi_stage > 1 doesn't help much locally; + we keep 2-4 to preserve persistent-scheduler benefits + at larger M. + """ + self.acc_dtype = acc_dtype + self.cluster_shape_mnk = (1, 1, 1) + self.tile_shape_mnk = tuple(tile_shape_mnk) + self.epi_stage_target = int(epi_stage) + if self.epi_stage_target < 1: + raise ValueError(f"epi_stage must be >= 1 (got {epi_stage})") + self.pipeline_depth = int(pipeline_depth) + self.use_fp16_mma = int(use_fp16_mma) + self.tile_swizzle = int(tile_swizzle) + self.raster_along_m = bool(raster_along_m) + self.enable_pdl = bool(enable_pdl) + # Optional override for the epilogue tile shape. + self.epi_tile_override = ( + tuple(epi_tile_override) if epi_tile_override is not None else None + ) + self.tiled_mma = None + # num_mcast_ctas_a / num_mcast_ctas_b are derived from the cluster + # shape in __call__ before use; left unset here so their type is + # inferred from that assignment (matches the sibling kernels). + self.is_a_mcast = False + self.is_b_mcast = False + + if self.tile_shape_mnk[1] % _PACK_TILE_N != 0: + raise ValueError( + f"bf16 x fp4 requires tile_N % {_PACK_TILE_N} == 0 " + f"(got tile_N={self.tile_shape_mnk[1]})" + ) + if self.tile_shape_mnk[2] % _PACK_TILE_K != 0: + raise ValueError( + f"bf16 x fp4 requires tile_K % {_PACK_TILE_K} == 0 " + f"(got tile_K={self.tile_shape_mnk[2]})" + ) + + self.occupancy = 1 + # 2x2 atom layout: 4 MMA warps arranged as 2 M-warps x 2 N-warps. + self.atom_layout = tuple(atom_layout) + if self.atom_layout not in ((2, 2, 1), (1, 2, 1)): + raise ValueError( + f"Unsupported atom_layout {self.atom_layout!r}; " + "expected (2,2,1) or (1,2,1)" + ) + self.num_mma_warps = ( + self.atom_layout[0] * self.atom_layout[1] * self.atom_layout[2] + ) + self.num_dma_warps = 1 + self.num_threads_per_warp = 32 + self.threads_per_cta = ( + self.num_mma_warps + self.num_dma_warps + ) * self.num_threads_per_warp + # SM100/103 expose >= SM120 SMEM/CTA + self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_120") + + self.ab_stage = None + self.epi_stage = None + + self.a_smem_layout_staged = None + self.b_smem_layout_staged = None + self.epi_smem_layout_staged = None + self.buffer_align_bytes = 1024 + + self.epilog_sync_barrier = pipeline.NamedBarrier( + barrier_id=2, + num_threads=self.num_mma_warps * self.num_threads_per_warp, + ) + self.load_register_requirement = 40 + self.mma_register_requirement = 232 + + def _setup_attributes(self): + self.mma_inst_mnk = (16, 8, 16) + # MmaF16BF16Op accepts ab_dtype in {Float16, BFloat16}. We pick via + # b_compute_dtype so use_fp16_mma=1 swings the whole MMA to fp16 + # (both A and B fragments will be fp16-typed). + op = cute.nvgpu.warp.MmaF16BF16Op( + self.b_compute_dtype, + self.acc_dtype, + self.mma_inst_mnk, + ) + tC = cute.make_layout(self.atom_layout) + permutation_mnk = ( + self.atom_layout[0] * self.mma_inst_mnk[0], + # *2 trick: each warp covers two atom-N tiles in one ldmatrix.x4 + self.atom_layout[1] * self.mma_inst_mnk[1] * 2, + self.atom_layout[2] * self.mma_inst_mnk[2], + ) + self.tiled_mma = cute.make_tiled_mma( + op, + tC, + permutation_mnk=permutation_mnk, + ) + + self.cta_layout_mnk = cute.make_layout(self.cluster_shape_mnk) + + self.num_mcast_ctas_a = self.cluster_shape_mnk[1] + self.num_mcast_ctas_b = self.cluster_shape_mnk[0] + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + + self.epi_tile = sm90_utils.compute_tile_shape_or_override( + self.tile_shape_mnk, + self.c_dtype, + is_cooperative=False, + epi_tile_override=self.epi_tile_override, + ) + + # B-side smem is packed int32 (4 bytes per logical FP4 + # pair). Stage budget uses int32 B + uint8 scales; bf16 phantom + # layout is only used by partition_B/make_fragment_B. + self.ab_stage, self.epi_stage = self._compute_stages( + self.tile_shape_mnk, + self.a_dtype, + self.epi_tile, + self.c_dtype, + self.smem_capacity, + self.occupancy, + self.GROUP_SIZE, + self.epi_stage_target, + ) + + if self.ab_stage == 0: + raise RuntimeError( + "ab_stage == 0: not enough shared memory for this tile shape " + f"({self.tile_shape_mnk}) at occupancy {self.occupancy}." + ) + + ( + self.a_smem_layout_staged, + self.b_packed_smem_layout_staged, + self.b_sf_smem_layout_staged, + self.b_bf16_logical_layout, + self.epi_smem_layout_staged, + ) = self._make_smem_layouts( + self.tile_shape_mnk, + self.epi_tile, + self.a_dtype, + self.a_layout, + self.b_compute_dtype, + self.b_layout_compute, + self.ab_stage, + self.c_dtype, + self.c_layout, + self.epi_stage, + self.GROUP_SIZE, + ) + + @cute.jit + def __call__( + self, + a: cute.Tensor, + b_packed: cute.Tensor, + b_sf: cute.Tensor, + c: cute.Tensor, + alpha: cute.Tensor, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + ): + from cutlass import Float16 + + self.a_dtype = a.element_type + # MMA operand dtype: bf16 by default matches A, or fp16 when + # use_fp16_mma=1 (lets us skip the f16->f32->bf16 cvt chain on B). + if cutlass.const_expr(self.use_fp16_mma == 1): + self.b_compute_dtype = Float16 + else: + self.b_compute_dtype = a.element_type + self.c_dtype = c.element_type + + self.a_layout = utils.LayoutEnum.from_tensor(a) + self.b_layout_compute = utils.LayoutEnum.ROW_MAJOR + self.c_layout = utils.LayoutEnum.from_tensor(c) + + if cutlass.const_expr(self.a_dtype.width != 16): + raise TypeError(f"a_dtype must be 16-bit (bf16/fp16), got {self.a_dtype}") + if cutlass.const_expr(self.a_dtype != self.c_dtype): + raise TypeError( + f"a_dtype and c_dtype must match, got {self.a_dtype} vs {self.c_dtype}" + ) + + self._setup_attributes() + + tma_atom_a, tma_tensor_a = self._make_tma_atoms_and_tensors( + a, + self.a_smem_layout_staged, + (self.tile_shape_mnk[0], self.tile_shape_mnk[2]), + 1, + ) + + # Packed B TMA: tile is (tile_K // 16, 2 * tile_N) int32. + b_packed_tma_tile = ( + self.tile_shape_mnk[2] // _PACK_TILE_K, + 2 * self.tile_shape_mnk[1], + ) + tma_atom_b_packed, tma_tensor_b_packed = self._make_tma_atoms_and_tensors( + b_packed, + self.b_packed_smem_layout_staged, + b_packed_tma_tile, + 1, + ) + + # B scale TMA: tile is (tile_K // group_size, tile_N) uint8. + # Small per-tile load -- this replaces ``tile_K // group_size`` + # gmem-direct loads per thread per K-block in the dequant path. + b_sf_tma_tile = ( + self.tile_shape_mnk[2] // self.GROUP_SIZE, + self.tile_shape_mnk[1], + ) + tma_atom_b_sf, tma_tensor_b_sf = self._make_tma_atoms_and_tensors( + b_sf, + self.b_sf_smem_layout_staged, + b_sf_tma_tile, + 1, + ) + + tma_atom_c, tma_tensor_c = self._make_tma_store_atoms_and_tensors( + c, + self.epi_smem_layout_staged, + self.epi_tile, + ) + + tile_sched_params, grid = self._compute_grid( + c, + self.tile_shape_mnk, + max_active_clusters, + self.tile_swizzle, + self.raster_along_m, + ) + + @cute.struct + class SharedStorage: + mainloop_pipeline_array_ptr: cute.struct.MemRange[ + cutlass.Int64, self.ab_stage * 2 + ] + sA: cute.struct.Align[ + cute.struct.MemRange[ + self.a_dtype, cute.cosize(self.a_smem_layout_staged) + ], + self.buffer_align_bytes, + ] + sB_packed: cute.struct.Align[ + cute.struct.MemRange[ + cutlass.Int32, cute.cosize(self.b_packed_smem_layout_staged) + ], + self.buffer_align_bytes, + ] + sB_sf: cute.struct.Align[ + cute.struct.MemRange[ + cutlass.Uint8, cute.cosize(self.b_sf_smem_layout_staged) + ], + self.buffer_align_bytes, + ] + sC: cute.struct.Align[ + cute.struct.MemRange[ + self.c_dtype, cute.cosize(self.epi_smem_layout_staged) + ], + self.buffer_align_bytes, + ] + + self.shared_storage = SharedStorage + + self.kernel( + tma_atom_a, + tma_tensor_a, + tma_atom_b_packed, + tma_tensor_b_packed, + tma_atom_b_sf, + tma_tensor_b_sf, + tma_atom_c, + tma_tensor_c, + alpha, + self.tiled_mma, + self.cta_layout_mnk, + self.a_smem_layout_staged, + self.b_packed_smem_layout_staged, + self.b_sf_smem_layout_staged, + self.b_bf16_logical_layout, + self.epi_smem_layout_staged, + tile_sched_params, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=[1, 1, 1], + stream=stream, + use_pdl=self.enable_pdl, + ) + return + + @cute.kernel + def kernel( + self, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b_packed: cute.CopyAtom, + mB_packed_kn: cute.Tensor, + tma_atom_b_sf: cute.CopyAtom, + mB_sf_kn: cute.Tensor, + tma_atom_c: cute.CopyAtom, + mC_mnl: cute.Tensor, + mAlpha: cute.Tensor, + tiled_mma: cute.TiledMma, + cta_layout_mnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_packed_smem_layout_staged: cute.Layout, + b_sf_smem_layout_staged: cute.Layout, + b_bf16_logical_layout: cute.ComposedLayout, + epi_smem_layout_staged: cute.ComposedLayout, + tile_sched_params: utils.PersistentTileSchedulerParams, + ): + tidx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + # Prefetch TMA descriptors from warp 0. + if warp_idx == 0: + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_a) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_b_packed) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_b_sf) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_c) + + cta_rank_in_cluster = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + cluster_coord_mnk = cta_layout_mnk.get_flat_coord(cta_rank_in_cluster) + + a_mcast_mask = cute.make_layout_image_mask( + cta_layout_mnk, cluster_coord_mnk, mode=1 + ) + b_mcast_mask = cute.make_layout_image_mask( + cta_layout_mnk, cluster_coord_mnk, mode=0 + ) + a_mcast_mask = a_mcast_mask if self.is_a_mcast else 0 + b_mcast_mask = b_mcast_mask if self.is_b_mcast else 0 + + a_smem_layout = cute.slice_(a_smem_layout_staged, (None, None, 0)) + b_packed_smem_layout = cute.slice_(b_packed_smem_layout_staged, (None, None, 0)) + b_sf_smem_layout = cute.slice_(b_sf_smem_layout_staged, (None, None, 0)) + tma_copy_bytes = ( + cute.size_in_bytes(self.a_dtype, a_smem_layout) + + cute.size_in_bytes(cutlass.Int32, b_packed_smem_layout) + + cute.size_in_bytes(cutlass.Uint8, b_sf_smem_layout) + ) + + smem = cutlass.utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + mainloop_pipeline_array_ptr = storage.mainloop_pipeline_array_ptr.data_ptr() + + mainloop_pipeline_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread + ) + mcast_size = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + consumer_arrive_cnt = mcast_size * self.num_mma_warps + mainloop_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, consumer_arrive_cnt + ) + + cta_layout_vmnk = cute.make_layout((1, *cta_layout_mnk.shape)) + mainloop_pipeline = pipeline.PipelineTmaAsync.create( + num_stages=self.ab_stage, + producer_group=mainloop_pipeline_producer_group, + consumer_group=mainloop_pipeline_consumer_group, + tx_count=tma_copy_bytes, + barrier_storage=mainloop_pipeline_array_ptr, + cta_layout_vmnk=cta_layout_vmnk, + ) + + if cute.size(self.cluster_shape_mnk) > 1: + cute.arch.cluster_arrive_relaxed() + + sA = storage.sA.get_tensor( + a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner + ) + # Both B-side tensors are plain (non-swizzled) staged layouts. + sB_packed = storage.sB_packed.get_tensor(b_packed_smem_layout_staged) + sB_sf = storage.sB_sf.get_tensor(b_sf_smem_layout_staged) + sC = storage.sC.get_tensor( + epi_smem_layout_staged.outer, swizzle=epi_smem_layout_staged.inner + ) + + gA_mkl = cute.local_tile( + mA_mkl, + cute.slice_(self.tile_shape_mnk, (None, 0, None)), + (None, None, None), + ) + # Packed B: (K // 16, N * 2, L); tile = (tile_K // 16, 2 * tile_N). + b_packed_tile_shape = ( + self.tile_shape_mnk[2] // _PACK_TILE_K, + 2 * self.tile_shape_mnk[1], + ) + gB_packed_kn = cute.local_tile( + mB_packed_kn, + b_packed_tile_shape, + (None, None, None), + ) + # Scales: (K // 16, N, L); per-tile slice has shape + # (tile_K // group_size, tile_N) = (tile_K // 16, tile_N). Dequant + # indexes directly per K-block + per N-coord. + gB_sf_kn = cute.local_tile( + mB_sf_kn, + (self.tile_shape_mnk[2] // self.GROUP_SIZE, self.tile_shape_mnk[1]), + (None, None, None), + ) + gC_mnl = cute.local_tile( + mC_mnl, + cute.slice_(self.tile_shape_mnk, (None, None, 0)), + (None, None, None), + ) + + thr_mma = tiled_mma.get_slice(tidx) + + # TMA partition for A: (m, k) -> per-CTA partition. + a_cta_layout = cute.make_layout(cute.slice_(cta_layout_mnk, (0, None, 0)).shape) + a_cta_crd = cluster_coord_mnk[1] + tAsA, tAgA = cute.nvgpu.cpasync.tma_partition( + tma_atom_a, + a_cta_crd, + a_cta_layout, + cute.group_modes(sA, 0, 2), + cute.group_modes(gA_mkl, 0, 2), + ) + + # TMA partition for B (packed int32). + b_cta_layout = cute.make_layout(cute.slice_(cta_layout_mnk, (None, 0, 0)).shape) + b_cta_crd = cluster_coord_mnk[0] + tBpacked_s, tBpacked_g = cute.nvgpu.cpasync.tma_partition( + tma_atom_b_packed, + b_cta_crd, + b_cta_layout, + cute.group_modes(sB_packed, 0, 2), + cute.group_modes(gB_packed_kn, 0, 2), + ) + + # TMA partition for B_sf (uint8 scales). Shares the b_cta_crd + # (no N multicast since cluster_shape_mnk = (1,1,1)). + tBsf_s, tBsf_g = cute.nvgpu.cpasync.tma_partition( + tma_atom_b_sf, + b_cta_crd, + b_cta_layout, + cute.group_modes(sB_sf, 0, 2), + cute.group_modes(gB_sf_kn, 0, 2), + ) + + # B partition uses a phantom bf16 layout on top of the sB_packed + # int32 storage (recast pointer + bf16 logical layout). This + # gives ``partition_B``/``make_fragment_B`` the right fragment + # shape; the data is never read through this view -- we always + # decode FP4 ourselves. + sB_phantom = cute.make_tensor( + cute.recast_ptr(sB_packed.iterator, dtype=self.b_compute_dtype), + b_bf16_logical_layout, + ) + tCsA = thr_mma.partition_A(sA) + tCsB_phantom = thr_mma.partition_B(sB_phantom) + tCrA = tiled_mma.make_fragment_A(tCsA[None, None, None, 0]) + tCrB = tiled_mma.make_fragment_B(tCsB_phantom[None, None, None, 0]) + + tCgC = thr_mma.partition_C(gC_mnl) + acc_shape = tCgC.shape[:3] + accumulators = cute.make_rmem_tensor(acc_shape, self.acc_dtype) + + if cute.size(self.cluster_shape_mnk) > 1: + cute.arch.cluster_wait() + else: + pipeline.sync(barrier_id=1) + + # PDL bookend (start): wait for the prior grid to finish so this + # kernel's TMA loads see the producer kernel's writes. + cute.arch.griddepcontrol_wait() + + k_tile_cnt = cute.size(gA_mkl, mode=[3]) + + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + mainloop_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.ab_stage + ) + mainloop_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.ab_stage + ) + + # MMA warp group: warps [0, num_mma_warps) compute. + if warp_idx < self.num_mma_warps: + cute.arch.setmaxregister_increase(self.mma_register_requirement) + + num_k_blocks = cute.size(tCrA, mode=[2]) + + # ldmatrix is only for A. B is filled in-register from sB_packed + # via FP4 decode + per-group scale. + # + # use_fp16_mma=1: tCrA is fp16-typed (matches the fp16 MMA), but + # sA is bf16 in SMEM. ldmatrix into a bf16 staging fragment, + # then per-K-block convert bf16 -> fp16 in-register before MMA + # reads tCrA. See the inline cvt at each ldmatrix site. + from cutlass import BFloat16 + + atom_copy_ldmatrix_A = cute.make_copy_atom( + cute.nvgpu.warp.LdMatrix8x8x16bOp(self.a_layout.is_m_major_a(), 4), + self.a_dtype, + ) + smem_tiled_copy_A = cute.make_tiled_copy_A(atom_copy_ldmatrix_A, tiled_mma) + thr_copy_ldmatrix_A = smem_tiled_copy_A.get_slice(tidx) + tCsA_copy_view = thr_copy_ldmatrix_A.partition_S(sA) + if cutlass.const_expr(self.use_fp16_mma == 1): + tCrA_bf16 = cute.make_fragment_like(tCrA, BFloat16) + tCrA_copy_view = thr_copy_ldmatrix_A.retile(tCrA_bf16) + else: + tCrA_bf16 = tCrA + tCrA_copy_view = thr_copy_ldmatrix_A.retile(tCrA) + + alpha_val = Float32(mAlpha[Int32(0)]) + + while work_tile.is_valid_tile: + tile_coord_mnl = work_tile.tile_idx + gC_mnl_slice = gC_mnl[(None, None, *tile_coord_mnl)] + accumulators.fill(0.0) + + mainloop_consumer_state.reset_count() + + peek_ab_full_status = cutlass.Boolean(1) + if mainloop_consumer_state.count < k_tile_cnt: + peek_ab_full_status = mainloop_pipeline.consumer_try_wait( + mainloop_consumer_state + ) + + mainloop_pipeline.consumer_wait( + mainloop_consumer_state, peek_ab_full_status + ) + tCsA_p = tCsA_copy_view[None, None, None, mainloop_consumer_state.index] + + # Prologue: prefetch tCrA[0] and tCrB[0] only when running with + # pipeline_depth >= 1. With depth=0 the dequant for block k + # happens inside the same iteration as gemm(k) -- no prefetch. + if cutlass.const_expr(self.pipeline_depth >= 1): + cute.copy( + smem_tiled_copy_A, + tCsA_p[None, None, 0], + tCrA_copy_view[None, None, 0], + ) + if cutlass.const_expr(self.use_fp16_mma == 1): + self._cvt_a_bf16_to_fp16_one_k_block(tCrA, tCrA_bf16, 0) + self._dequant_b_to_register( + sB_packed, + sB_sf, + tCrB, + tidx, + mainloop_consumer_state.index, + mainloop_consumer_state.count, + 0, + ) + + for _k_tile in cutlass.range(0, k_tile_cnt - 1, 1, unroll=1): + for k_block_idx in cutlass.range_constexpr(num_k_blocks): + k_block_next = ( + 0 if k_block_idx + 1 == num_k_blocks else k_block_idx + 1 + ) + + if cutlass.const_expr(self.pipeline_depth == 0): + # No-prefetch: dequant(k) just before gemm(k), + # all on the CURRENT K-tile's stage. Release + + # advance + wait happens AFTER gemm of the last + # block so we don't drop SMEM under our own read. + cute.copy( + smem_tiled_copy_A, + tCsA_p[None, None, k_block_idx], + tCrA_copy_view[None, None, k_block_idx], + ) + if cutlass.const_expr(self.use_fp16_mma == 1): + self._cvt_a_bf16_to_fp16_one_k_block( + tCrA, tCrA_bf16, k_block_idx + ) + self._dequant_b_to_register( + sB_packed, + sB_sf, + tCrB, + tidx, + mainloop_consumer_state.index, + mainloop_consumer_state.count, + k_block_idx, + ) + cute.gemm( + tiled_mma, + accumulators, + tCrA[None, None, k_block_idx], + tCrB[None, None, k_block_idx], + accumulators, + ) + if k_block_idx == num_k_blocks - 1: + mainloop_pipeline.consumer_release( + mainloop_consumer_state + ) + mainloop_consumer_state.advance() + peek_ab_full_status = cutlass.Boolean(1) + peek_ab_full_status = ( + mainloop_pipeline.consumer_try_wait( + mainloop_consumer_state + ) + ) + tCsA_p = tCsA_copy_view[ + None, None, None, mainloop_consumer_state.index + ] + mainloop_pipeline.consumer_wait( + mainloop_consumer_state, peek_ab_full_status + ) + else: + # 1-stage prefetch: dequant(k+1) while gemm(k). + if k_block_idx == num_k_blocks - 1: + mainloop_pipeline.consumer_release( + mainloop_consumer_state + ) + mainloop_consumer_state.advance() + + peek_ab_full_status = cutlass.Boolean(1) + peek_ab_full_status = ( + mainloop_pipeline.consumer_try_wait( + mainloop_consumer_state + ) + ) + + tCsA_p = tCsA_copy_view[ + None, None, None, mainloop_consumer_state.index + ] + mainloop_pipeline.consumer_wait( + mainloop_consumer_state, peek_ab_full_status + ) + + cute.copy( + smem_tiled_copy_A, + tCsA_p[None, None, k_block_next], + tCrA_copy_view[None, None, k_block_next], + ) + if cutlass.const_expr(self.use_fp16_mma == 1): + self._cvt_a_bf16_to_fp16_one_k_block( + tCrA, tCrA_bf16, k_block_next + ) + self._dequant_b_to_register( + sB_packed, + sB_sf, + tCrB, + tidx, + mainloop_consumer_state.index, + mainloop_consumer_state.count, + k_block_next, + ) + cute.gemm( + tiled_mma, + accumulators, + tCrA[None, None, k_block_idx], + tCrB[None, None, k_block_idx], + accumulators, + ) + # Hoist out last k_tile (no further loads after the last k_block) + for k_block_idx in cutlass.range_constexpr(num_k_blocks): + k_block_next = ( + 0 if k_block_idx + 1 == num_k_blocks else k_block_idx + 1 + ) + + if cutlass.const_expr(self.pipeline_depth == 0): + # No-prefetch path for last K-tile. Release happens + # AFTER gemm of the last block (kernel exits then). + cute.copy( + smem_tiled_copy_A, + tCsA_p[None, None, k_block_idx], + tCrA_copy_view[None, None, k_block_idx], + ) + if cutlass.const_expr(self.use_fp16_mma == 1): + self._cvt_a_bf16_to_fp16_one_k_block( + tCrA, tCrA_bf16, k_block_idx + ) + self._dequant_b_to_register( + sB_packed, + sB_sf, + tCrB, + tidx, + mainloop_consumer_state.index, + mainloop_consumer_state.count, + k_block_idx, + ) + cute.gemm( + tiled_mma, + accumulators, + tCrA[None, None, k_block_idx], + tCrB[None, None, k_block_idx], + accumulators, + ) + if k_block_idx == num_k_blocks - 1: + mainloop_pipeline.consumer_release(mainloop_consumer_state) + mainloop_consumer_state.advance() + else: + # 1-stage prefetch path for last K-tile. + if k_block_idx == num_k_blocks - 1: + mainloop_pipeline.consumer_release(mainloop_consumer_state) + mainloop_consumer_state.advance() + + if k_block_next > 0: + cute.copy( + smem_tiled_copy_A, + tCsA_p[None, None, k_block_next], + tCrA_copy_view[None, None, k_block_next], + ) + if cutlass.const_expr(self.use_fp16_mma == 1): + self._cvt_a_bf16_to_fp16_one_k_block( + tCrA, tCrA_bf16, k_block_next + ) + self._dequant_b_to_register( + sB_packed, + sB_sf, + tCrB, + tidx, + mainloop_consumer_state.index, + mainloop_consumer_state.count, + k_block_next, + ) + cute.gemm( + tiled_mma, + accumulators, + tCrA[None, None, k_block_idx], + tCrB[None, None, k_block_idx], + accumulators, + ) + + # Epilogue: accumulator -> smem -> gmem via R2S (StMatrix.x4) + # + TMA bulk store. + copy_atom_r2s = sm90_utils.sm90_get_smem_store_op( + self.c_layout, + elem_ty_d=self.c_dtype, + elem_ty_acc=self.acc_dtype, + ) + + copy_atom_C = cute.make_copy_atom( + cute.nvgpu.warp.StMatrix8x8x16bOp( + self.c_layout.is_m_major_c(), + 4, + ), + self.c_dtype, + ) + + tiled_copy_C_Atom = cute.make_tiled_copy_C_atom(copy_atom_C, tiled_mma) + + tiled_copy_r2s = cute.make_tiled_copy_S( + copy_atom_r2s, + tiled_copy_C_Atom, + ) + + thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) + tRS_sD = thr_copy_r2s.partition_D(sC) + tRS_rAcc = tiled_copy_r2s.retile(accumulators) + + rD_shape = cute.shape(thr_copy_r2s.partition_S(sC)) + tRS_rD_layout = cute.make_layout(rD_shape[:3]) + tRS_rD = cute.make_rmem_tensor(tRS_rD_layout.shape, self.acc_dtype) + + sepi_for_tma_partition = cute.group_modes(sC, 0, 2) + tcgc_for_tma_partition = cute.zipped_divide(gC_mnl_slice, self.epi_tile) + + bSG_sD, bSG_gD = cute.nvgpu.cpasync.tma_partition( + tma_atom_c, + 0, + cute.make_layout(1), + sepi_for_tma_partition, + tcgc_for_tma_partition, + ) + + # Epilogue: iterate (epi_m, epi_n) explicitly + # and use (mma_m, mma_n) mode indexing into tRS_rAcc so + # the loop works for any epi_tile_m / epi_tile_n. Also + # supports OOB-iteration skipping when m_actual < tile_M. + epi_rest_m = cute.size(tcgc_for_tma_partition, mode=[1, 0]) + epi_rest_n = cute.size(tcgc_for_tma_partition, mode=[1, 1]) + epi_tile_m = self.epi_tile[0] + epi_tile_n = self.epi_tile[1] + # mma_tile_{m,n} = per-mma-atom (M,N) size. tRS_rAcc has + # shape (atom_v, mma_m, mma_n); modes 1, 2 give the atom + # counts in M, N. + mma_tile_m = self.tile_shape_mnk[0] // cute.size(tRS_rAcc, mode=[1]) + mma_tile_n = self.tile_shape_mnk[1] // cute.size(tRS_rAcc, mode=[2]) + MmaMPerEpiM = epi_tile_m // mma_tile_m + MmaNPerEpiN = epi_tile_n // mma_tile_n + + tma_store_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.num_mma_warps * self.num_threads_per_warp, + ) + tma_store_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.epi_stage, + producer_group=tma_store_producer_group, + ) + + # Skip OOB epilogue iterations when actual M < tile_M. + m_actual = cute.size(mC_mnl, mode=[0]) + cta_m_offset = tile_coord_mnl[0] * Int32(self.tile_shape_mnk[0]) + + # kept_count cycles in lockstep with the TMA-store pipeline. + # epi_buffer = kept_count % num_stages stays in sync with + # producer_commit/acquire calls, regardless of how many + # iterations are skipped. + kept_count = 0 + for epi_n in cutlass.range_constexpr(epi_rest_n): + for epi_m in cutlass.range_constexpr(epi_rest_m): + epi_m_global_start = cta_m_offset + Int32(epi_m * epi_tile_m) + if epi_m_global_start < m_actual: + # Copy this epi-tile's slice of acc -> tRS_rD + # using b12x-style (mma_m, mma_n) indexing. + for mma_n_in_epi in cutlass.range_constexpr(MmaNPerEpiN): + for mma_m_in_epi in cutlass.range_constexpr( + MmaMPerEpiM + ): + mma_n = epi_n * MmaNPerEpiN + mma_n_in_epi + mma_m = epi_m * MmaMPerEpiM + mma_m_in_epi + tRS_rD_slice = tRS_rD[ + (None, mma_m_in_epi, mma_n_in_epi) + ] + tRS_rAcc_slice = tRS_rAcc[(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 + ] + + tRS_rD_out = cute.make_rmem_tensor( + tRS_rD_layout.shape, self.c_dtype + ) + # Apply the global alpha here, once, on the + # fp32 accumulator (hoisted out of the + # per-K-block dequant): one rounding instead + # of folding alpha into every B scale. + acc_vec = tRS_rD.load() + tRS_rD_out.store((alpha_val * acc_vec).to(self.c_dtype)) + + epi_buffer = kept_count % cute.size(tRS_sD, mode=[3]) + cute.copy( + tiled_copy_r2s, + tRS_rD_out, + tRS_sD[(None, None, None, epi_buffer)], + ) + cute.arch.fence_proxy("async.shared", space="cta") + self.epilog_sync_barrier.arrive_and_wait() + + if warp_idx == 0: + cute.copy( + tma_atom_c, + bSG_sD[(None, epi_buffer)], + bSG_gD[(None, (epi_m, epi_n))], + ) + tma_store_pipeline.producer_commit() + tma_store_pipeline.producer_acquire() + kept_count = kept_count + 1 + + tma_store_pipeline.producer_tail() + + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + # Single DMA warp: issues all 3 TMA descriptors (A, B_packed, B_sf) + # back-to-back into the same stage barrier per K-tile. + elif warp_idx == self.num_mma_warps: + cute.arch.setmaxregister_decrease(self.load_register_requirement) + while work_tile.is_valid_tile: + tile_coord_mnl = work_tile.tile_idx + tAgA_mkl = tAgA[(None, tile_coord_mnl[0], None, tile_coord_mnl[2])] + tBpacked_g_kn = tBpacked_g[ + (None, None, tile_coord_mnl[1], tile_coord_mnl[2]) + ] + tBsf_g_kn = tBsf_g[(None, None, tile_coord_mnl[1], tile_coord_mnl[2])] + mainloop_producer_state.reset_count() + for _k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + mainloop_pipeline.producer_acquire(mainloop_producer_state) + barrier_ptr = mainloop_pipeline.producer_get_barrier( + mainloop_producer_state + ) + + tAgA_k = tAgA_mkl[(None, mainloop_producer_state.count)] + tAsA_pipe = tAsA[(None, mainloop_producer_state.index)] + cute.copy( + tma_atom_a, + tAgA_k, + tAsA_pipe, + tma_bar_ptr=barrier_ptr, + mcast_mask=a_mcast_mask, + ) + + tBpacked_g_k = tBpacked_g_kn[(None, mainloop_producer_state.count)] + tBpacked_s_pipe = tBpacked_s[(None, mainloop_producer_state.index)] + cute.copy( + tma_atom_b_packed, + tBpacked_g_k, + tBpacked_s_pipe, + tma_bar_ptr=barrier_ptr, + mcast_mask=b_mcast_mask, + ) + + tBsf_g_k = tBsf_g_kn[(None, mainloop_producer_state.count)] + tBsf_s_pipe = tBsf_s[(None, mainloop_producer_state.index)] + cute.copy( + tma_atom_b_sf, + tBsf_g_k, + tBsf_s_pipe, + tma_bar_ptr=barrier_ptr, + mcast_mask=b_mcast_mask, + ) + + mainloop_pipeline.producer_commit(mainloop_producer_state) + mainloop_producer_state.advance() + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + mainloop_pipeline.producer_tail(mainloop_producer_state) + cute.arch.griddepcontrol_launch_dependents() + return + + @staticmethod + def _compute_stages( + tile_shape_mnk: Tuple[int, int, int], + a_dtype: Type[cutlass.Numeric], + epi_tile: Tuple[int, int], + c_dtype: Type[cutlass.Numeric], + smem_capacity: int, + occupancy: int, + group_size: int = 16, + epi_stage: int = 4, + ) -> Tuple[int, int]: + """Stage budget accounting for A + B (packed int32) + B_sf (uint8). + + Per (16K x 64N) packed block: 128 int32 = 512 bytes. + Per (group_size K x tile_N N) scale block: tile_N bytes. + """ + c_bytes_per_stage = cute.size(epi_tile) * c_dtype.width // 8 + epi_bytes = c_bytes_per_stage * epi_stage + + a_shape = cute.slice_(tile_shape_mnk, (None, 0, None)) + a_bytes_per_stage = cute.size(a_shape) * a_dtype.width // 8 + + # B side: packed int32, 128 int32 per (16K x 64N) block. + packed_blocks_per_tile = (tile_shape_mnk[2] // _PACK_TILE_K) * ( + tile_shape_mnk[1] // _PACK_TILE_N + ) + b_packed_bytes_per_stage = packed_blocks_per_tile * _PACK_INTS_PER_TILE * 4 + + # Scale tile: 1 byte per (K-group, N). Small compared to B_packed. + b_sf_bytes_per_stage = (tile_shape_mnk[2] // group_size) * tile_shape_mnk[1] + + ab_bytes_per_stage = ( + a_bytes_per_stage + b_packed_bytes_per_stage + b_sf_bytes_per_stage + ) + mbar_helpers_bytes = 1024 + + ab_stage = ( + (smem_capacity - occupancy * 1024) // occupancy + - mbar_helpers_bytes + - epi_bytes + ) // ab_bytes_per_stage + return ab_stage, epi_stage + + @staticmethod + def _make_smem_layouts( + tile_shape_mnk: Tuple[int, int, int], + epi_tile: Tuple[int, int], + a_dtype: Type[cutlass.Numeric], + a_layout: cute.Layout, + b_compute_dtype: Type[cutlass.Numeric], + b_layout_compute: cute.Layout, + ab_stage: int, + c_dtype: Type[cutlass.Numeric], + c_layout: cute.Layout, + epi_stage: int, + group_size: int, + ): + """Returns (sA, sB_packed, sB_sf, b_bf16_phantom, sC) layouts. + + ``sB_packed_layout`` is a plain (non-swizzled) staged int32 layout + ``(tile_K // 16, 2 * tile_N, ab_stage)``. + + ``sB_sf_layout`` is a plain staged uint8 layout + ``(tile_K // group_size, tile_N, ab_stage)`` -- one byte per + (K-group, N) cell. Loaded via TMA once per K-tile, read from + SMEM in the dequant inner loop (replaces the gmem-direct path). + + ``b_bf16_logical_layout`` is the phantom bf16 layout used only by + ``partition_B`` / ``make_fragment_B`` for fragment-shape + determination -- there is no real bf16 SMEM allocation. + """ + a_smem_layout_staged = sm90_utils.make_smem_layout_a( + a_layout, + tile_shape_mnk, + a_dtype, + ab_stage, + ) + + # sB_packed: (tile_K // 16) rows x (2 * tile_N) int32 cols x stage. + b_packed_smem_layout_staged = cute.make_ordered_layout( + ( + tile_shape_mnk[2] // _PACK_TILE_K, + 2 * tile_shape_mnk[1], + ab_stage, + ), + order=(1, 0, 2), + ) + + # sB_sf: (tile_K // group_size) rows x tile_N uint8 cols x stage. + # Order (1, 0, 2) -> N innermost (= contiguous load lane), K next, + # stage outermost. Matches the gmem layout. + b_sf_smem_layout_staged = cute.make_ordered_layout( + ( + tile_shape_mnk[2] // group_size, + tile_shape_mnk[1], + ab_stage, + ), + order=(1, 0, 2), + ) + + # bf16 phantom layout for partition_B / make_fragment_B. + b_bf16_logical_layout = sm90_utils.make_smem_layout_b( + b_layout_compute, + tile_shape_mnk, + b_compute_dtype, + ab_stage, + ) + + epi_smem_layout_staged = sm90_utils.make_smem_layout_epi( + c_dtype, + c_layout, + epi_tile, + epi_stage, + ) + return ( + a_smem_layout_staged, + b_packed_smem_layout_staged, + b_sf_smem_layout_staged, + b_bf16_logical_layout, + epi_smem_layout_staged, + ) + + @staticmethod + def _compute_grid( + c: cute.Tensor, + tile_shape_mnk: Tuple[int, int, int], + max_active_clusters: cutlass.Constexpr, + tile_swizzle: cutlass.Constexpr = 1, + raster_along_m: cutlass.Constexpr = True, + ): + c_shape = cute.slice_(tile_shape_mnk, (None, None, 0)) + gc = cute.zipped_divide(c, tiler=c_shape) + num_ctas_mnl = gc[(0, (None, None, None))].shape + cluster_shape_mnl = (1, 1, 1) + tile_sched_params = utils.PersistentTileSchedulerParams( + num_ctas_mnl, + cluster_shape_mnl, + swizzle_size=tile_swizzle, + raster_along_m=raster_along_m, + ) + grid = utils.StaticPersistentTileScheduler.get_grid_shape( + tile_sched_params, max_active_clusters + ) + return tile_sched_params, grid + + @staticmethod + def _make_tma_store_atoms_and_tensors( + tensor_c: cute.Tensor, + epi_smem_layout_staged: cute.ComposedLayout, + epi_tile: Tuple[int, int], + ) -> Tuple[cute.CopyAtom, cute.Tensor]: + epi_smem_layout = cute.slice_(epi_smem_layout_staged, (None, None, 0)) + tma_atom_c, tma_tensor_c = cute.nvgpu.cpasync.make_tiled_tma_atom( + cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp(), + tensor_c, + epi_smem_layout, + epi_tile, + ) + return tma_atom_c, tma_tensor_c + + @staticmethod + def _make_tma_atoms_and_tensors( + tensor: cute.Tensor, + smem_layout_staged: cute.ComposedLayout, + smem_tile: Tuple[int, int], + mcast_dim: int, + ) -> Tuple[cute.CopyAtom, cute.Tensor]: + op = ( + cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp() + if mcast_dim == 1 + else cute.nvgpu.cpasync.CopyBulkTensorTileG2SMulticastOp() + ) + smem_layout = cute.slice_(smem_layout_staged, (None, None, 0)) + tma_atom, tma_tensor = cute.nvgpu.cpasync.make_tiled_tma_atom( + op, + tensor, + smem_layout, + smem_tile, + num_multicast=mcast_dim, + ) + return tma_atom, tma_tensor + + @cute.jit + def _cvt_a_bf16_to_fp16_one_k_block( + self, + tCrA_dst, + tCrA_bf16_src, + k_block: cutlass.Constexpr, + ): + """In-register bf16 -> fp16 cvt for one K-block of A. + + Recasts both fragments to Uint32 (each u32 packs 2 16-bit elems) + and applies `cvt_bf16x2_to_f16x2_via_f32` per pair. The packed + narrowing cvt (`cvt.rn.f16x2.f32`) combines what would be two + scalar `cvt.rn.f16.f32` instructions if we used cute's default + `.to(Float16)` lowering. + + bf16/f16 typed inline-asm constraints don't compile on sm_100a + (NVVM rejects them), so we keep everything in Uint32 pair + representation via cute.recast_tensor. + """ + bf_u32 = cute.recast_tensor(tCrA_bf16_src[None, None, k_block], Uint32) + fp_u32 = cute.recast_tensor(tCrA_dst[None, None, k_block], Uint32) + n_pairs = cute.size(bf_u32) + for i in cutlass.range_constexpr(n_pairs): + fp_u32[i] = cvt_bf16x2_to_f16x2_via_f32(Uint32(bf_u32[i])) + + @cute.jit + def _dequant_b_to_register( + self, + sB_packed: cute.Tensor, + sB_sf: cute.Tensor, + tCrB: cute.Tensor, + tidx: Int32, + stage_idx: Int32, + k_tile_idx: Int32, + k_block_idx: cutlass.Constexpr, + ): + """Decode 2 int32 per thread per K-block into 16 fp16 fragment slots. + + Per the MMA partition for atom_layout (2,2,1) + *2 on N: + + tc_row = (lane % 4) * 2 in {0, 2, 4, 6} + tc_col = lane // 4 in [0, 8) + n_warp_idx = warp_idx // 2 in {0, 1} + base_n = n_warp_idx * 8 + tc_col in [0, 16) + + Each thread covers: + K = {tc_row, tc_row+1, tc_row+8, tc_row+9} + N = {base_n, base_n+16, base_n+32, base_n+48} + + Two int32s per K-block per thread; sB_packed offsets: + u32_0 @ sB_packed[k_block_idx, n_warp_idx * 64 + lane * 2 + 0, stage] + u32_1 @ sB_packed[k_block_idx, n_warp_idx * 64 + lane * 2 + 1, stage] + + Byte layout inside each int32 (see ``_cute_dsl_pack_fp4_weight`` in + ``flashinfer/gemm/gemm_bf16_fp4_cute_dsl.py``): + u32_0: + byte 0: K=tc_row, tc_row+1 at N=base_n -> (mma_i=0,1, nn=0) + byte 1: K=tc_row+8, tc_row+9 at N=base_n -> (mma_i=2,3, nn=0) + byte 2: K=tc_row, tc_row+1 at N=base_n+16 -> (mma_i=0,1, nn=1) + byte 3: K=tc_row+8, tc_row+9 at N=base_n+16 -> (mma_i=2,3, nn=1) + u32_1 mirrors with N=base_n+32 and N=base_n+48 (nn=2, nn=3). + """ + lane = tidx % Int32(32) + warp = tidx // Int32(32) + # n_warp_idx maps warp -> N-stripe. With (2,2,1) the 4 warps are + # arranged as 2 M-warps * 2 N-warps, so n_warp_idx = warp // 2. + # With (1,2,1) there are 2 warps total (no M-warp dim), so each + # warp is its own N-stripe: n_warp_idx = warp. In general + # n_warp_idx = warp // atom_layout[0]. + n_warp_idx = warp // Int32(self.atom_layout[0]) + tc_col = lane // Int32(4) + base_n_in_tile = n_warp_idx * Int32(8) + tc_col + + # tile_N is built from tile_N // 64 independent packed 64-N blocks. + # The hand-coded byte->fragment mapping below covers exactly one such + # 64-N block (4 nn slots); for tile_N=128 we loop it over both blocks. + # Block n_blk lives at sB_packed int32 columns [n_blk*128, +128), sB_sf + # N-rows [n_blk*64, +64), and writes tCrB nn in [n_blk*4, +4). For + # tile_N=64 this loops once. + num_n_blocks = cutlass.const_expr(self.tile_shape_mnk[1] // 64) + + # _write_hmul2 captures tCrB / k_block_idx; nn and mma-row passed per + # write. fp16 MMA writes fp16 directly; bf16 MMA goes via fp32 (no + # packed f16x2 -> bf16x2 cvt on sm_100a). + if cutlass.const_expr(self.use_fp16_mma == 1): + + def _write_hmul2(h2, scale_h2, mma_i_low, nn): + scaled_h2 = half2_mul(h2, scale_h2) + f_lo, f_hi = f16x2_unpack(scaled_h2) + tCrB[mma_i_low, nn, k_block_idx] = f_lo + tCrB[mma_i_low + 1, nn, k_block_idx] = f_hi + else: + + def _write_hmul2(h2, scale_h2, mma_i_low, nn): + scaled_h2 = half2_mul(h2, scale_h2) + f_lo, f_hi = f16x2_to_f32x2(scaled_h2) + tCrB[mma_i_low, nn, k_block_idx] = f_lo.to(self.b_compute_dtype) + tCrB[mma_i_low + 1, nn, k_block_idx] = f_hi.to(self.b_compute_dtype) + + for n_blk in cutlass.range_constexpr(num_n_blocks): + n_col_off = Int32(n_blk * 128) # int32 column offset of this 64-N block + n_sf_off = Int32(n_blk * 64) # SF N-row offset of this 64-N block + nn0 = n_blk * 4 # base tCrB N-fragment slot for this block + + # Single ld.shared.v2.u32 (8-byte load): u32_pos_base is always even + # (= n_blk*128 + n_warp*64 + lane*2), so the offset is 8-byte aligned. + u32_pos_base = n_col_off + n_warp_idx * Int32(64) + lane * Int32(2) + smem_addr_b = get_smem_ptr_as_int32( + sB_packed, sB_packed.layout((k_block_idx, u32_pos_base, stage_idx)) + ) + u32_0, u32_1 = ld_shared_v2_u32(smem_addr_b) + + # Per-group scale loads (S0E5M3). 4 distinct N positions per block. + sf_n = base_n_in_tile + n_sf_off + sf_byte_0 = Uint32(sB_sf[k_block_idx, sf_n + Int32(0), stage_idx]) + sf_byte_1 = Uint32(sB_sf[k_block_idx, sf_n + Int32(16), stage_idx]) + sf_byte_2 = Uint32(sB_sf[k_block_idx, sf_n + Int32(32), stage_idx]) + sf_byte_3 = Uint32(sB_sf[k_block_idx, sf_n + Int32(48), stage_idx]) + + h0_a, h0_b, h0_c, h0_d = fp4_decode_4bytes(u32_0) + h1_a, h1_b, h1_c, h1_d = fp4_decode_4bytes(u32_1) + + # S0E5M3 scale -> f16x2 broadcast in one mul.lo.u32. alpha is + # applied once in the epilogue, so the scale never needs an f32 form. + sc_n0 = cvt_s0e5m3_to_f16x2_broadcast(sf_byte_0) + sc_n1 = cvt_s0e5m3_to_f16x2_broadcast(sf_byte_1) + sc_n2 = cvt_s0e5m3_to_f16x2_broadcast(sf_byte_2) + sc_n3 = cvt_s0e5m3_to_f16x2_broadcast(sf_byte_3) + _write_hmul2(h0_a, sc_n0, 0, nn0 + 0) + _write_hmul2(h0_b, sc_n0, 2, nn0 + 0) + _write_hmul2(h0_c, sc_n1, 0, nn0 + 1) + _write_hmul2(h0_d, sc_n1, 2, nn0 + 1) + _write_hmul2(h1_a, sc_n2, 0, nn0 + 2) + _write_hmul2(h1_b, sc_n2, 2, nn0 + 2) + _write_hmul2(h1_c, sc_n3, 0, nn0 + 3) + _write_hmul2(h1_d, sc_n3, 2, nn0 + 3) + + @cute.jit + def wrapper( + self, + mA: cute.Tensor, + mB_packed: cute.Tensor, + mB_sf: cute.Tensor, + mC: cute.Tensor, + mAlpha: cute.Tensor, + l: cutlass.Constexpr, + max_active_clusters: cutlass.Constexpr, + current_stream, + ): + """bf16 x fp4 wrapper for the FlashInfer compile interface. + + Args: + mA: (m, k) input tensor A, bf16 or fp16. + mB_packed: (k // 16, n * 2) int32 -- packed FP4. + mB_sf: (k // 16, n) uint8 -- FP8-E4M3 per-group scales. + mC: (m, n) output tensor C, bf16 or fp16. + mAlpha: (1,) fp32 global scale. + l: batch dimension (Constexpr); typically 1. + max_active_clusters: Constexpr from get_max_active_clusters(1). + current_stream: CUDA stream (TVM-FFI fake stream). + """ + m = cute.size(mA, mode=[0]) + k = cute.size(mA, mode=[1]) + n = cute.size(mC, mode=[1]) + k_tiles = k // _PACK_TILE_K + n_packed = 2 * n + k_sf_groups = k // self.GROUP_SIZE + + a_tensor = cute.make_tensor( + mA.iterator, + layout=cute.make_ordered_layout((m, k, l), order=(1, 0, 2)), + ) + b_packed_tensor = cute.make_tensor( + mB_packed.iterator, + layout=cute.make_ordered_layout((k_tiles, n_packed, l), order=(1, 0, 2)), + ) + b_sf_tensor = cute.make_tensor( + mB_sf.iterator, + layout=cute.make_ordered_layout((k_sf_groups, n, l), order=(1, 0, 2)), + ) + c_tensor = cute.make_tensor( + mC.iterator, + layout=cute.make_ordered_layout((m, n, l), order=(1, 0, 2)), + ) + + self( + a_tensor, + b_packed_tensor, + b_sf_tensor, + c_tensor, + mAlpha, + max_active_clusters, + current_stream, + ) diff --git a/flashinfer/trace/templates/gemm.py b/flashinfer/trace/templates/gemm.py index 9733a946ee5..e13736dc1af 100644 --- a/flashinfer/trace/templates/gemm.py +++ b/flashinfer/trace/templates/gemm.py @@ -153,18 +153,11 @@ def _mm_mxfp8_reference(A, B, a_descale, b_descale): def _mm_fp4_reference(A, B, a_descale, b_descale, block_size=16): - """Dequantize FP4 inputs and compute C = A @ B. - - A and B are fp4 e2m1fn values packed two-per-byte as uint8. - a_descale: [M, K//block_size], b_descale: [K, N//block_size]. - The reference unpacks the nibbles and applies the block scales. - """ + """Dequantize FP4 inputs and compute C = A @ B.""" def _unpack_fp4(packed, rows, cols): - # Each byte holds two fp4 nibbles (low nibble = first element). lo = (packed & 0x0F).to(torch.float32) hi = ((packed >> 4) & 0x0F).to(torch.float32) - # Interleave low/high nibbles along the last dimension. out = torch.stack([lo, hi], dim=-1).reshape(rows, cols) return out @@ -468,6 +461,312 @@ def _mm_fp4_init( ) +# ── BF16 x FP4 GEMM (mm_bf16_fp4, weight-only) ────────────────────────────── +# +# ``mm_bf16_fp4`` consumes *prepared* weights whose layout is +# backend-specific (see ``flashinfer.prepare_bf16_fp4_weights``), so each +# backend gets its own template and ``mm_fp4``'s ``trace=`` is a dispatch +# callable (same pattern as the trtllm MoE routing templates). + +# E2M1 value table (low 3 bits magnitude, bit 3 sign). +_E2M1_VALUES = ( + 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, + -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0, +) # fmt: skip + + +def _bf16_fp4_matmul(a, weight_kn, alpha): + """Final fp32 matmul shared by both bf16 x fp4 references. + + ``weight_kn`` is the dequantized fp32 weight in (K, N) layout. + """ + if alpha is not None: + weight_kn = weight_kn * alpha.to(torch.float32) + return (a.to(torch.float32) @ weight_kn).to(a.dtype) + + +def _mm_bf16_fp4_cudnn_reference(a, b, b_descale, alpha=None, block_size=16): + """Reference for the cuDNN-prepared layout. + + b: [N, K//2] uint8, two FP4 codes per byte. + b_descale: [N, K//block_size] float8_e4m3fn per-block scales (linear). + """ + n, k_half = b.shape + k = k_half * 2 + lut = torch.tensor(_E2M1_VALUES, dtype=torch.float32, device=b.device) + b_int = b.to(torch.int64) + codes = torch.stack([b_int & 0xF, (b_int >> 4) & 0xF], dim=-1).reshape(n, k) + sf = b_descale.to(torch.float32).repeat_interleave(block_size, dim=1) + return _bf16_fp4_matmul(a, (lut[codes] * sf).T, alpha) + + +def _mm_bf16_fp4_cute_dsl_reference(a, b, b_descale, alpha=None, block_size=16): + """Reference for the cute-DSL-prepared layout. + + b: [K//16, N*2] int32 -- FP4 bytes permuted into (16K x 64N) MMA tiles + of 128 int32 each (inverts + ``flashinfer.gemm.gemm_bf16_fp4_cute_dsl._cute_dsl_pack_fp4_weight``). + b_descale: [K//block_size, N] uint8 -- S0E5M3 per-block scales + (fp16 value = byte << 7 reinterpreted as fp16 bits). + """ + device = b.device + k_sf, n = b_descale.shape + k = k_sf * block_size + k_tiles, n_tiles = k // 16, n // 64 + + # Rebuild the within-tile byte permutation of _cute_dsl_pack_fp4_weight. + u32_pos = torch.arange(128, device=device, dtype=torch.long) + lane = (u32_pos // 2) % 32 + base_n = (u32_pos // 64) * 8 + lane // 4 + k_half_in_tile = (lane % 4)[:, None] + torch.tensor( + [0, 4, 0, 4], device=device, dtype=torch.long + ) + n_in_tile = ( + base_n[:, None] + + torch.tensor( + [[0, 0, 16, 16], [32, 32, 48, 48]], device=device, dtype=torch.long + )[u32_pos % 2] + ) + within_idx = (k_half_in_tile * 64 + n_in_tile).reshape(-1) # (512,) permutation + + # Invert: scatter each tile's 512 gathered bytes back to row-major (8, 64). + gathered = ( + b.reshape(k_tiles, n_tiles, 128, 1) + .view(torch.uint8) + .reshape(k_tiles, n_tiles, 512) + ) + tile_bytes = torch.empty_like(gathered) + tile_bytes[:, :, within_idx] = gathered + b_kn = ( + tile_bytes.reshape(k_tiles, n_tiles, 8, 64) + .permute(0, 2, 1, 3) + .reshape(k // 2, n) + ) + + lut = torch.tensor(_E2M1_VALUES, dtype=torch.float32, device=device) + b_int = b_kn.to(torch.int64) + codes = torch.stack([b_int & 0xF, (b_int >> 4) & 0xF], dim=1).reshape(k, n) + # S0E5M3 -> fp16: the byte is the top 8 bits of the fp16 bit pattern. + sf = ( + (b_descale.to(torch.int16) << 7) + .view(torch.float16) + .to(torch.float32) + .repeat_interleave(block_size, dim=0) + ) + return _bf16_fp4_matmul(a, lut[codes] * sf, alpha) + + +def _mm_bf16_fp4_cudnn_init( + *, + M: int, + N: int = 2048, + K: int = 7168, + block_size: int = 16, + K_div_2: int = 0, # derived + K_div_block_size: int = 0, # derived + device: str = "cuda", + seed: int = 0, +): + """Build inputs for ``flashinfer.mm_bf16_fp4`` (cuDNN backend). + + Sourced from ``tests/gemm/test_mm_bf16_fp4.py``: quantize a randn + bf16 weight via ``flashinfer.nvfp4_quantize`` (layout_128x4), then + repack with ``prepare_bf16_fp4_weights``. Requires SM100+ at + runtime; CPU smoke tests skip. + """ + del K_div_2, K_div_block_size + from flashinfer import ( # noqa: PLC0415 + mm_bf16_fp4, + nvfp4_quantize, + prepare_bf16_fp4_weights, + ) + from flashinfer.quantization.fp4_quantization import SfLayout # noqa: PLC0415 + + if not torch.cuda.is_available() or torch.device(device).type != "cuda": + raise NotImplementedError("mm_bf16_fp4 init requires a CUDA device") + major, minor = torch.cuda.get_device_capability(torch.device(device)) + if not mm_bf16_fp4.is_backend_supported("cudnn", major * 10 + minor): + raise NotImplementedError(f"mm_bf16_fp4 is not supported on SM{major}{minor}") + + torch.manual_seed(seed) + a = torch.randn(M, K, dtype=torch.bfloat16, device=device) + w = torch.randn(N, K, dtype=torch.bfloat16, device=device) + g_w = (448.0 * 6.0) / w.float().abs().nan_to_num().max() + b_fp4, b_sf = nvfp4_quantize( + w, g_w, sfLayout=SfLayout.layout_128x4, do_shuffle=False, backend="cute-dsl" + ) + alpha = torch.tensor([1.0 / g_w.item()], dtype=torch.float32, device=device) + b_p, sf_p, alpha_p = prepare_bf16_fp4_weights( + b_fp4, b_sf, alpha, backend="cudnn", block_size=block_size + ) + return { + "a": a, + "b": b_p, + "b_descale": sf_p, + "alpha": alpha_p, + "backend": "cudnn", + "block_size": int(block_size), + } + + +def _mm_bf16_fp4_cute_dsl_init( + *, + M: int, + N: int = 2048, + K: int = 7168, + block_size: int = 16, + K_div_16: int = 0, # derived + K_div_block_size: int = 0, # derived + N_mul_2: int = 0, # derived + device: str = "cuda", + seed: int = 0, +): + """Build inputs for ``flashinfer.mm_bf16_fp4`` (cute-DSL backend). + + Sourced from ``tests/gemm/test_mm_bf16_fp4.py``: quantize a randn + bf16 weight via ``flashinfer.nvfp4_quantize`` (layout_128x4), then + repack with ``prepare_bf16_fp4_weights``. Requires SM100+ at + runtime; CPU smoke tests skip. + """ + del K_div_16, K_div_block_size, N_mul_2 + from flashinfer import ( # noqa: PLC0415 + mm_bf16_fp4, + nvfp4_quantize, + prepare_bf16_fp4_weights, + ) + from flashinfer.quantization.fp4_quantization import SfLayout # noqa: PLC0415 + + if not torch.cuda.is_available() or torch.device(device).type != "cuda": + raise NotImplementedError("mm_bf16_fp4 init requires a CUDA device") + major, minor = torch.cuda.get_device_capability(torch.device(device)) + if not mm_bf16_fp4.is_backend_supported("cute-dsl", major * 10 + minor): + raise NotImplementedError(f"mm_bf16_fp4 is not supported on SM{major}{minor}") + + torch.manual_seed(seed) + a = torch.randn(M, K, dtype=torch.bfloat16, device=device) + w = torch.randn(N, K, dtype=torch.bfloat16, device=device) + g_w = (448.0 * 6.0) / w.float().abs().nan_to_num().max() + b_fp4, b_sf = nvfp4_quantize( + w, g_w, sfLayout=SfLayout.layout_128x4, do_shuffle=False, backend="cute-dsl" + ) + alpha = torch.tensor([1.0 / g_w.item()], dtype=torch.float32, device=device) + b_p, sf_p, alpha_p = prepare_bf16_fp4_weights( + b_fp4, b_sf, alpha, backend="cute-dsl", block_size=block_size + ) + return { + "a": a, + "b": b_p, + "b_descale": sf_p, + "alpha": alpha_p, + "backend": "cute-dsl", + "block_size": int(block_size), + } + + +mm_bf16_fp4_cudnn_trace = TraceTemplate( + op_type="gemm_bf16_fp4", + name_prefix="mm_bf16_fp4_cudnn", + description=( + "bf16 x fp4 GEMM C = (A @ dequant(B).T) * alpha, cuDNN-prepared weights. " + "A is bf16; B is fp4 (e2m1fn_x2 packed as uint8) with fp8-e4m3 " + "per-block scales in linear [N, K//block_size] layout." + ), + axes={ + "M": Var(), + "N": Const(), + "K": Const(), + "block_size": Const(description="FP4 quantization block size (16 for nvfp4)."), + }, + inputs={ + "A": Tensor(["M", "K"], param="a", description="Activation, bfloat16."), + "B": Tensor( + ["N", "K_div_2"], + param="b", + description="Weight, fp4 e2m1fn_x2 packed as uint8, [N, K//2].", + ), + "b_descale": Tensor( + ["N", "K_div_block_size"], + description="Per-block scales, float8_e4m3fn, [N, K//block_size].", + ), + "alpha": Tensor( + ["1"], + optional=True, + description="Optional global scale, float32, shape (1,).", + ), + "block_size": Scalar("int32", description="FP4 block size (always 16)."), + }, + outputs={ + "C": Tensor(["M", "N"], dtype_from="a"), + }, + tags=["status:verified", "quantization:fp4"], + reference=_mm_bf16_fp4_cudnn_reference, + check=_fp4_gemm_check, + init=_mm_bf16_fp4_cudnn_init, +) + +mm_bf16_fp4_cute_dsl_trace = TraceTemplate( + op_type="gemm_bf16_fp4", + name_prefix="mm_bf16_fp4_cute_dsl", + description=( + "bf16 x fp4 GEMM C = (A @ dequant(B).T) * alpha, cute-DSL-prepared weights. " + "A is bf16; B is fp4 repacked into (16K x 64N) MMA tiles as int32 " + "[K//16, N*2] with S0E5M3 per-block scales [K//block_size, N]." + ), + axes={ + "M": Var(), + "N": Const(), + "K": Const(), + "block_size": Const(description="FP4 quantization block size (16 for nvfp4)."), + }, + inputs={ + "A": Tensor(["M", "K"], param="a", description="Activation, bfloat16."), + "B": Tensor( + ["K_div_16", "N_mul_2"], + param="b", + description="Weight, fp4 tile-packed as int32, [K//16, N*2].", + ), + "b_descale": Tensor( + ["K_div_block_size", "N"], + description="Per-block scales, S0E5M3 as uint8, [K//block_size, N].", + ), + "alpha": Tensor( + ["1"], + optional=True, + description="Optional global scale, float32, shape (1,).", + ), + "block_size": Scalar("int32", description="FP4 block size (always 16)."), + }, + outputs={ + "C": Tensor(["M", "N"], dtype_from="a"), + }, + tags=["status:verified", "quantization:fp4"], + reference=_mm_bf16_fp4_cute_dsl_reference, + check=_fp4_gemm_check, + init=_mm_bf16_fp4_cute_dsl_init, +) + + +def mm_bf16_fp4_trace_dispatch(**kwargs): + """Return the TraceTemplate for an ``mm_bf16_fp4`` call by backend. + + The prepared weight layout differs per backend (int32 -> cute-dsl, + uint8 -> cudnn), so each gets its own template. Pass as + ``trace=mm_bf16_fp4_trace_dispatch`` to ``@flashinfer_api``. + """ + b = kwargs.get("b") + if b is not None and b.dtype == torch.int32: + return mm_bf16_fp4_cute_dsl_trace + return mm_bf16_fp4_cudnn_trace + + +# Expose the templates so _attach_fi_trace auto-registers them for the +# consistency tests (same pattern as the MoE routing dispatchers). +mm_bf16_fp4_trace_dispatch.templates = [ # type: ignore[attr-defined] + mm_bf16_fp4_cudnn_trace, + mm_bf16_fp4_cute_dsl_trace, +] + + # ── Batched matmuls (BMM) ──────────────────────────────────────────────────── diff --git a/tests/gemm/test_mm_bf16_fp4.py b/tests/gemm/test_mm_bf16_fp4.py new file mode 100644 index 00000000000..73e97c27690 --- /dev/null +++ b/tests/gemm/test_mm_bf16_fp4.py @@ -0,0 +1,355 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 by FlashInfer team. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the BF16 x FP4 GEMM API ``mm_bf16_fp4``.""" + +import pytest +import torch + +import flashinfer +from flashinfer import mm_bf16_fp4, prepare_bf16_fp4_weights +from flashinfer.autotuner import autotune +from flashinfer.gemm.gemm_bf16_fp4 import ( + _CUDNN_BF16_FP4_MIN_BACKEND_VERSION, + _unswizzle_sf_128x4, +) +from flashinfer.utils import get_compute_capability + + +# E2M1 (FP4) value table, signed (codes 0-7 positive, 8-15 negative), matching +# ``flashinfer.nvfp4_quantize``. +_E2M1_VALUES_FP32 = ( + 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, + -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0, +) # fmt: skip + + +def _dequantize_bf16_fp4_torch(b, b_descale, alpha, n, k, block_size): + """PyTorch implementation of swizzled nvfp4 dequantization to fp32.""" + device = b.device + k_sf = k // block_size + lut = torch.tensor(_E2M1_VALUES_FP32, dtype=torch.float32, device=device) + b_int = b.to(torch.int64) + codes = torch.stack([b_int & 0xF, (b_int >> 4) & 0xF], dim=-1).reshape(n, k) + values = lut[codes] + sf = _unswizzle_sf_128x4(b_descale, n, k_sf).view(torch.float8_e4m3fn) + sf_expanded = sf.to(torch.float32).repeat_interleave(block_size, dim=1) + weight = values * sf_expanded + if alpha is not None: + weight = weight * alpha.to(torch.float32) + return weight + + +# ============================================================================= +# Backend + shape grids +# ============================================================================= + + +# Backends covered by the cross-backend contract tests. New backends get +# appended here as they land. +ALL_BACKENDS = ["cudnn", "cute-dsl"] + + +def _skip_if_backend_unavailable(backend: str) -> None: + """Skip the current test if ``backend`` can't run on this device.""" + device = torch.device("cuda") + cc = get_compute_capability(device) + cc_number = cc[0] * 10 + cc[1] + if not mm_bf16_fp4.is_backend_supported(backend, cc_number): + pytest.skip(f"{backend} not supported on compute capability {cc_number}") + if backend == "cudnn": + try: + import cudnn + except ImportError: + pytest.skip("cuDNN not available") + if cudnn.backend_version() < _CUDNN_BF16_FP4_MIN_BACKEND_VERSION: + pytest.skip( + f"cuDNN bf16 x fp4 needs backend >= {_CUDNN_BF16_FP4_MIN_BACKEND_VERSION}, " + f"found {cudnn.backend_version()}" + ) + + +def _skip_if_compute_capability_unsupported() -> None: + """Skip the current test if no bf16 x fp4 backend supports this device.""" + cc = get_compute_capability(torch.device("cuda")) + cc_number = cc[0] * 10 + cc[1] + if not mm_bf16_fp4.is_backend_supported("cudnn", cc_number): + pytest.skip(f"mm_bf16_fp4 not supported on compute capability {cc_number}") + + +PROBLEM_SIZES = [ + # tiny: smoke / minimum valid shapes + (1, 128, 128), + (1, 256, 512), + (4, 256, 512), + (16, 256, 256), + # mid: typical decode at a few model widths + (1, 1024, 1024), + (4, 1024, 1024), + (16, 1024, 1024), + (64, 1024, 1024), + # large: realistic model-layer N/K, sweep of M + (1, 4096, 4096), + (4, 4096, 4096), + (16, 4096, 4096), + (64, 4096, 4096), + (128, 4096, 4096), + (256, 4096, 4096), + (512, 4096, 4096), + # ---- non-power-of-2 shapes (N, K multiples of 64; mixed tile_K=64/128) ---- + # small M (decode), odd M and odd N/K + (1, 192, 192), + (2, 320, 256), + (3, 448, 320), + (5, 576, 192), + (7, 704, 256), + (11, 832, 384), + (13, 960, 512), + (17, 1088, 576), + (6, 1216, 640), + (9, 1344, 768), + (1, 2112, 1024), + (4, 2688, 1344), + (15, 1600, 896), + # mid M + (48, 192, 1024), + (96, 320, 768), + (100, 576, 1152), + (127, 704, 960), + (192, 832, 1024), + (200, 1088, 1280), + (250, 1216, 1024), + (160, 2560, 2112), + (96, 3072, 1344), + (48, 1856, 1536), + # large M (prefill) + (384, 1088, 1024), + (500, 1344, 1152), + (768, 2112, 2048), + (1000, 1600, 1536), + (1500, 2560, 1024), + (2048, 2688, 2688), + (3000, 1088, 768), + (1024, 4160, 2048), + (640, 6144, 1024), + (2000, 3200, 1024), + # skinny / wide extremes + (1, 5120, 256), + (7, 4160, 192), + (13, 11008, 128), + (64, 2112, 2112), + (256, 832, 1344), + (333, 1600, 640), + (17, 3072, 3072), +] + +# Default shape for API sanity tests. +# (alpha=None, out_dtype override, preallocated out, K-mismatch). +SMOKE_MNK = (16, 1024, 1024) + +ATOL = 1.5e-2 +RTOL = 1.5e-2 + + +def _assert_close_to_reference(out: torch.Tensor, ref: torch.Tensor, backend: str): + """Compare a backend's output against the fp32-accurate reference.""" + out_f = out.float().reshape(-1) + ref_f = ref.float().reshape(-1) + ref_norm = torch.linalg.vector_norm(ref_f).clamp_min(1e-6) + rel_l2 = (torch.linalg.vector_norm(out_f - ref_f) / ref_norm).item() + cos = torch.nn.functional.cosine_similarity(out_f, ref_f, dim=0).item() + assert rel_l2 < 2e-2, f"{backend}: relative L2 error {rel_l2:.4f} exceeds 2e-2" + assert cos > 0.999, f"{backend}: cosine similarity {cos:.6f} below 0.999" + + +# ============================================================================= +# Helpers +# ============================================================================= + + +def _make_random_fp4_weights( + n: int, k: int, device: torch.device +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Quantize a random matrix to NVFP4 + return (b_fp4, b_sf, alpha). + + Mirrors the canonical caller pattern: a user does + ``b_fp4, b_sf = flashinfer.nvfp4_quantize(mat2, g_b, ...)`` and pairs + that with ``alpha = 1 / g_b``. + """ + mat2 = torch.randn((n, k), device=device, dtype=torch.bfloat16) + g_b = (448 * 6) / mat2.float().abs().nan_to_num().max() + b_fp4, b_sf = flashinfer.nvfp4_quantize( + mat2, + g_b, + sfLayout=flashinfer.SfLayout.layout_128x4, + do_shuffle=False, + backend="cute-dsl", + ) + alpha = torch.tensor([1.0 / g_b.item()], device=device, dtype=torch.float32) + return b_fp4, b_sf, alpha + + +# ============================================================================= +# Cross-backend numerical / behaviour contract +# ============================================================================= + + +@pytest.mark.parametrize("backend", ALL_BACKENDS) +@pytest.mark.parametrize("auto_tuning", [False, True]) +@pytest.mark.parametrize("m,n,k", PROBLEM_SIZES) +def test_backend_matches_handwritten_dequant_matmul(auto_tuning, backend, m, n, k): + """Backend output must match a hand-rolled fp32 dequant + matmul. + + Reference = ``(a.float() @ dequant(b).T).to(bf16)``. Every backend + is expected to produce numerically equivalent output (up to ~1 bf16 + ULP). Run with ``auto_tuning`` both off (fallback tactic) and on + (so the autotuner's selected tactic is exercised too). + """ + _skip_if_backend_unavailable(backend) + device = torch.device("cuda") + torch.manual_seed(0) + a = torch.randn((m, k), device=device, dtype=torch.bfloat16) + b_fp4, b_sf, alpha = _make_random_fp4_weights(n, k, device) + + b_p, sf_p, alpha_p = prepare_bf16_fp4_weights(b_fp4, b_sf, alpha, backend=backend) + with autotune(auto_tuning): + out = mm_bf16_fp4(a, b_p, sf_p, alpha_p, backend=backend) + + weight_fp32 = _dequantize_bf16_fp4_torch(b_fp4, b_sf, alpha, n, k, 16) + ref = (a.float() @ weight_fp32.T).to(torch.bfloat16) + + _assert_close_to_reference(out, ref, backend) + assert out.shape == (m, n) + assert out.dtype == torch.bfloat16 + + +@pytest.mark.parametrize("backend", ALL_BACKENDS) +@pytest.mark.parametrize("auto_tuning", [False, True]) +def test_backend_alpha_none_equals_alpha_one(auto_tuning, backend): + """alpha=None must produce identical output to alpha=tensor([1.0]).""" + _skip_if_backend_unavailable(backend) + device = torch.device("cuda") + m, n, k = SMOKE_MNK + a = torch.randn((m, k), device=device, dtype=torch.bfloat16) + b_fp4, b_sf, _ = _make_random_fp4_weights(n, k, device) + + b1, sf1, a1 = prepare_bf16_fp4_weights( + b_fp4, + b_sf, + torch.ones(1, device=device, dtype=torch.float32), + backend=backend, + ) + b0, sf0, a0 = prepare_bf16_fp4_weights(b_fp4, b_sf, None, backend=backend) + with autotune(auto_tuning): + out_one = mm_bf16_fp4(a, b1, sf1, a1, backend=backend) + out_none = mm_bf16_fp4(a, b0, sf0, a0, backend=backend) + + torch.testing.assert_close(out_none, out_one, atol=ATOL, rtol=RTOL) + + +@pytest.mark.parametrize("backend", ALL_BACKENDS) +def test_backend_out_dtype_override(backend): + """out_dtype kwarg controls return dtype independently of a.dtype.""" + _skip_if_backend_unavailable(backend) + if backend == "cute-dsl": + # The cute-dsl kernel's MMA path requires out_dtype == a.dtype, so + # it cannot emit fp16 from a bf16 activation (see _compute_cute_dsl). + pytest.skip("cute-dsl requires out_dtype == a.dtype") + device = torch.device("cuda") + m, n, k = SMOKE_MNK + a = torch.randn((m, k), device=device, dtype=torch.bfloat16) + b_fp4, b_sf, alpha = _make_random_fp4_weights(n, k, device) + b_p, sf_p, alpha_p = prepare_bf16_fp4_weights(b_fp4, b_sf, alpha, backend=backend) + out = mm_bf16_fp4( + a, + b_p, + sf_p, + alpha_p, + backend=backend, + out_dtype=torch.float16, + ) + assert out.dtype == torch.float16 + + +@pytest.mark.parametrize("backend", ALL_BACKENDS) +def test_backend_preallocated_out(backend): + """Caller-provided out tensor is written in place.""" + _skip_if_backend_unavailable(backend) + device = torch.device("cuda") + m, n, k = SMOKE_MNK + a = torch.randn((m, k), device=device, dtype=torch.bfloat16) + b_fp4, b_sf, alpha = _make_random_fp4_weights(n, k, device) + b_p, sf_p, alpha_p = prepare_bf16_fp4_weights(b_fp4, b_sf, alpha, backend=backend) + out = torch.empty((m, n), device=device, dtype=torch.bfloat16) + out_ptr_before = out.data_ptr() + returned = mm_bf16_fp4( + a, + b_p, + sf_p, + alpha_p, + backend=backend, + out=out, + ) + assert returned.data_ptr() == out_ptr_before + ref = mm_bf16_fp4(a, b_p, sf_p, alpha_p, backend=backend) + torch.testing.assert_close(returned, ref, atol=ATOL, rtol=RTOL) + + +@pytest.mark.parametrize("backend", ALL_BACKENDS) +def test_backend_shape_mismatch_raises(backend): + """K of a must match K inferred from prepared b.""" + _skip_if_backend_unavailable(backend) + device = torch.device("cuda") + m, n, k = SMOKE_MNK + b_fp4, b_sf, alpha = _make_random_fp4_weights(n, k, device) + b_p, sf_p, alpha_p = prepare_bf16_fp4_weights(b_fp4, b_sf, alpha, backend=backend) + a_wrong_k = torch.randn((m, k * 2), device=device, dtype=torch.bfloat16) + with pytest.raises(ValueError): + mm_bf16_fp4( + a_wrong_k, + b_p, + sf_p, + alpha_p, + backend=backend, + ) + + +# ============================================================================= +# Dispatcher-level input validation +# ============================================================================= +# +# These checks fire before any backend-specific code runs, so they're +# not parametrized over backend. + + +@pytest.mark.parametrize("bad_dtype", [torch.float32, torch.float16]) +def test_a_dtype_must_be_bfloat16(bad_dtype): + """Only bfloat16 activations are supported (fp16 deferred).""" + _skip_if_compute_capability_unsupported() + device = torch.device("cuda") + b_fp4, b_sf, alpha = _make_random_fp4_weights(64, 128, device) + b_p, sf_p, alpha_p = prepare_bf16_fp4_weights( + b_fp4, b_sf, alpha, backend="cute-dsl" + ) + a_bad = torch.randn((4, 128), device=device, dtype=bad_dtype) + with pytest.raises(TypeError): + mm_bf16_fp4(a_bad, b_p, sf_p, alpha_p, backend="cute-dsl") + + +def test_b_dtype_must_be_uint8_in_prepare(): + """Prepare rejects non-uint8 B.""" + _skip_if_compute_capability_unsupported() + device = torch.device("cuda") + b_bad = torch.zeros((64, 64), device=device, dtype=torch.int32) + b_descale = torch.zeros((4096,), device=device, dtype=torch.uint8) + with pytest.raises(TypeError): + prepare_bf16_fp4_weights(b_bad, b_descale, None, backend="cute-dsl") + + +def test_alpha_dtype_must_be_float32(): + """Prepare rejects non-fp32 alpha.""" + _skip_if_compute_capability_unsupported() + device = torch.device("cuda") + b_fp4, b_sf, _ = _make_random_fp4_weights(64, 128, device) + alpha_bad = torch.ones(1, device=device, dtype=torch.bfloat16) + with pytest.raises(TypeError): + prepare_bf16_fp4_weights(b_fp4, b_sf, alpha_bad, backend="cute-dsl") diff --git a/tests/trace/example.py b/tests/trace/example.py index 7e0e1094522..9a04d51b9cc 100644 --- a/tests/trace/example.py +++ b/tests/trace/example.py @@ -39,6 +39,8 @@ merge_states_h32_d128.json mla_paged_decode_h16_ckv512_kpe64_ps1.json mla_paged_decode_h16_ckv512_kpe64_ps64.json +mm_bf16_fp4_cudnn_N2048_K7168_block_size16.json +mm_bf16_fp4_cute_dsl_N2048_K7168_block_size16.json moe_fp4_block_scale_default_routing_topk8_e32_h7168_i2048.json moe_fp4_block_scale_ds_routing_topk8_e32_h7168_i2048_ng8_kg4.json moe_fp4_block_scale_llama4_routing_topk1_e32_h7168_i2048.json @@ -311,6 +313,32 @@ except Exception: pass # Requires Blackwell (SM100+) +# ── GEMM bf16 x fp4: mm_bf16_fp4 (weight-only) ────────────────────────────── +# Blackwell SM100+: M×7168@2048×7168, block=16. b/b_descale shapes are the +# *prepared* layouts (prepare_bf16_fp4_weights). +try: + M, K, N, BSW = 128, 7168, 2048, 16 + a_w4 = torch.zeros(M, K, dtype=torch.bfloat16, device=device) + alpha_w4 = torch.ones(1, dtype=torch.float32, device=device) + # cuDNN layout: canonical packed weight + linear fp8 scales. + b_w4 = torch.zeros(N, K // 2, dtype=torch.uint8, device=device) + sf_w4 = torch.ones(N, K // BSW, dtype=torch.float8_e4m3fn, device=device) + flashinfer.mm_bf16_fp4(a_w4, b_w4, sf_w4, alpha_w4, backend="cudnn", block_size=BSW) +except Exception: + pass # Requires Blackwell (SM100+) and cuDNN >= 9.23.1 +try: + M, K, N, BSW = 128, 7168, 2048, 16 + a_w4 = torch.zeros(M, K, dtype=torch.bfloat16, device=device) + alpha_w4 = torch.ones(1, dtype=torch.float32, device=device) + # cute-DSL layout: tile-packed int32 weight + S0E5M3 uint8 scales. + b_w4 = torch.zeros(K // 16, N * 2, dtype=torch.int32, device=device) + sf_w4 = torch.ones(K // BSW, N, dtype=torch.uint8, device=device) + flashinfer.mm_bf16_fp4( + a_w4, b_w4, sf_w4, alpha_w4, backend="cute-dsl", block_size=BSW + ) +except Exception: + pass # Requires Blackwell (SM100+) + # ── GQA paged decode (Llama-3.1-8B, h=32/kv=8/d=128) ──────────────────────── num_qo, num_kv, head_dim, batch_size = 32, 8, 128, 32 diff --git a/tests/trace/fi_trace_out/mm_bf16_fp4_cudnn_N2048_K7168_block_size16.json b/tests/trace/fi_trace_out/mm_bf16_fp4_cudnn_N2048_K7168_block_size16.json new file mode 100644 index 00000000000..d84cdd22545 --- /dev/null +++ b/tests/trace/fi_trace_out/mm_bf16_fp4_cudnn_N2048_K7168_block_size16.json @@ -0,0 +1,79 @@ +{ + "name": "mm_bf16_fp4_cudnn_N2048_K7168_block_size16", + "description": "bf16 x fp4 GEMM C = (A @ dequant(B).T) * alpha, cuDNN-prepared weights. A is bf16; B is fp4 (e2m1fn_x2 packed as uint8) with fp8-e4m3 per-block scales in linear [N, K//block_size] layout.", + "op_type": "gemm_bf16_fp4", + "tags": [ + "fi_api:flashinfer.gemm.gemm_bf16_fp4.mm_bf16_fp4", + "status:verified", + "quantization:fp4" + ], + "axes": { + "M": { + "type": "var" + }, + "N": { + "type": "const", + "value": 2048 + }, + "K": { + "type": "const", + "value": 7168 + }, + "block_size": { + "type": "const", + "value": 16, + "description": "FP4 quantization block size (16 for nvfp4)." + } + }, + "inputs": { + "A": { + "shape": [ + "M", + "K" + ], + "dtype": "bfloat16", + "description": "Activation, bfloat16." + }, + "B": { + "shape": [ + "N", + "K_div_2" + ], + "dtype": "uint8", + "description": "Weight, fp4 e2m1fn_x2 packed as uint8, [N, K//2]." + }, + "b_descale": { + "shape": [ + "N", + "K_div_block_size" + ], + "dtype": "float8_e4m3fn", + "description": "Per-block scales, float8_e4m3fn, [N, K//block_size]." + }, + "alpha": { + "shape": [ + "1" + ], + "dtype": "float32", + "optional": true, + "description": "Optional global scale, float32, shape (1,)." + }, + "block_size": { + "shape": null, + "dtype": "int32", + "description": "FP4 block size (always 16)." + } + }, + "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_bf16_fp4_cudnn_reference(a, b, b_descale, alpha=None, block_size=16):\n \"\"\"Reference for the cuDNN-prepared layout.\n\n b: [N, K//2] uint8, two FP4 codes per byte (low nibble = even K).\n b_descale: [N, K//block_size] float8_e4m3fn per-block scales (linear).\n \"\"\"\n n, k_half = b.shape\n k = k_half * 2\n lut = torch.tensor(_E2M1_VALUES, dtype=torch.float32, device=b.device)\n b_int = b.to(torch.int64)\n codes = torch.stack([b_int & 0xF, (b_int >> 4) & 0xF], dim=-1).reshape(n, k)\n sf = b_descale.to(torch.float32).repeat_interleave(block_size, dim=1)\n return _bf16_fp4_matmul(a, (lut[codes] * sf).T, alpha)\n", + "check": "def _fp4_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.97,\n):\n from flashinfer.trace import default_check\n\n # Matches tests/gemm/test_mm_fp4.py.\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_fp4_cudnn_init(\n *,\n M: int,\n N: int = 2048,\n K: int = 7168,\n block_size: int = 16,\n K_div_2: int = 0, # derived\n K_div_block_size: int = 0, # derived\n device: str = \"cuda\",\n seed: int = 0,\n):\n \"\"\"Build inputs for ``flashinfer.mm_bf16_fp4`` (cuDNN backend).\n\n Sourced from ``tests/gemm/test_mm_bf16_fp4.py``: quantize a randn\n bf16 weight via ``flashinfer.nvfp4_quantize`` (layout_128x4), then\n repack with ``prepare_bf16_fp4_weights``. Requires SM100+ at\n runtime; CPU smoke tests skip.\n \"\"\"\n del K_div_2, K_div_block_size\n from flashinfer import ( # noqa: PLC0415\n mm_bf16_fp4,\n nvfp4_quantize,\n prepare_bf16_fp4_weights,\n )\n from flashinfer.quantization.fp4_quantization import SfLayout # noqa: PLC0415\n\n if not torch.cuda.is_available() or torch.device(device).type != \"cuda\":\n raise NotImplementedError(\"mm_bf16_fp4 init requires a CUDA device\")\n major, minor = torch.cuda.get_device_capability(torch.device(device))\n if not mm_bf16_fp4.is_backend_supported(\"cudnn\", major * 10 + minor):\n raise NotImplementedError(f\"mm_bf16_fp4 is not supported on SM{major}{minor}\")\n\n torch.manual_seed(seed)\n a = torch.randn(M, K, dtype=torch.bfloat16, device=device)\n w = torch.randn(N, K, dtype=torch.bfloat16, device=device)\n g_w = (448.0 * 6.0) / w.float().abs().nan_to_num().max()\n b_fp4, b_sf = nvfp4_quantize(\n w, g_w, sfLayout=SfLayout.layout_128x4, do_shuffle=False, backend=\"cute-dsl\"\n )\n alpha = torch.tensor([1.0 / g_w.item()], dtype=torch.float32, device=device)\n b_p, sf_p, alpha_p = prepare_bf16_fp4_weights(\n b_fp4, b_sf, alpha, backend=\"cudnn\", block_size=block_size\n )\n return {\n \"a\": a,\n \"b\": b_p,\n \"b_descale\": sf_p,\n \"alpha\": alpha_p,\n \"backend\": \"cudnn\",\n \"block_size\": int(block_size),\n }\n" +} diff --git a/tests/trace/fi_trace_out/mm_bf16_fp4_cute_dsl_N2048_K7168_block_size16.json b/tests/trace/fi_trace_out/mm_bf16_fp4_cute_dsl_N2048_K7168_block_size16.json new file mode 100644 index 00000000000..5b6e51ec5b0 --- /dev/null +++ b/tests/trace/fi_trace_out/mm_bf16_fp4_cute_dsl_N2048_K7168_block_size16.json @@ -0,0 +1,79 @@ +{ + "name": "mm_bf16_fp4_cute_dsl_N2048_K7168_block_size16", + "description": "bf16 x fp4 GEMM C = (A @ dequant(B).T) * alpha, cute-DSL-prepared weights. A is bf16; B is fp4 repacked into (16K x 64N) MMA tiles as int32 [K//16, N*2] with S0E5M3 per-block scales [K//block_size, N].", + "op_type": "gemm_bf16_fp4", + "tags": [ + "fi_api:flashinfer.gemm.gemm_bf16_fp4.mm_bf16_fp4", + "status:verified", + "quantization:fp4" + ], + "axes": { + "M": { + "type": "var" + }, + "N": { + "type": "const", + "value": 2048 + }, + "K": { + "type": "const", + "value": 7168 + }, + "block_size": { + "type": "const", + "value": 16, + "description": "FP4 quantization block size (16 for nvfp4)." + } + }, + "inputs": { + "A": { + "shape": [ + "M", + "K" + ], + "dtype": "bfloat16", + "description": "Activation, bfloat16." + }, + "B": { + "shape": [ + "K_div_16", + "N_mul_2" + ], + "dtype": "int32", + "description": "Weight, fp4 tile-packed as int32, [K//16, N*2]." + }, + "b_descale": { + "shape": [ + "K_div_block_size", + "N" + ], + "dtype": "uint8", + "description": "Per-block scales, S0E5M3 as uint8, [K//block_size, N]." + }, + "alpha": { + "shape": [ + "1" + ], + "dtype": "float32", + "optional": true, + "description": "Optional global scale, float32, shape (1,)." + }, + "block_size": { + "shape": null, + "dtype": "int32", + "description": "FP4 block size (always 16)." + } + }, + "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_bf16_fp4_cute_dsl_reference(a, b, b_descale, alpha=None, block_size=16):\n \"\"\"Reference for the cute-DSL-prepared layout.\n\n b: [K//16, N*2] int32 -- FP4 bytes permuted into (16K x 64N) MMA tiles\n of 128 int32 each (inverts\n ``flashinfer.gemm.gemm_bf16_fp4._cute_dsl_pack_fp4_weight``).\n b_descale: [K//block_size, N] uint8 -- S0E5M3 per-block scales\n (fp16 value = byte << 7 reinterpreted as fp16 bits).\n \"\"\"\n device = b.device\n k_sf, n = b_descale.shape\n k = k_sf * block_size\n k_tiles, n_tiles = k // 16, n // 64\n\n # Rebuild the within-tile byte permutation of _cute_dsl_pack_fp4_weight.\n u32_pos = torch.arange(128, device=device, dtype=torch.long)\n lane = (u32_pos // 2) % 32\n base_n = (u32_pos // 64) * 8 + lane // 4\n k_half_in_tile = (lane % 4)[:, None] + torch.tensor(\n [0, 4, 0, 4], device=device, dtype=torch.long\n )\n n_in_tile = (\n base_n[:, None]\n + torch.tensor(\n [[0, 0, 16, 16], [32, 32, 48, 48]], device=device, dtype=torch.long\n )[u32_pos % 2]\n )\n within_idx = (k_half_in_tile * 64 + n_in_tile).reshape(-1) # (512,) permutation\n\n # Invert: scatter each tile's 512 gathered bytes back to row-major (8, 64).\n gathered = (\n b.reshape(k_tiles, n_tiles, 128, 1)\n .view(torch.uint8)\n .reshape(k_tiles, n_tiles, 512)\n )\n tile_bytes = torch.empty_like(gathered)\n tile_bytes[:, :, within_idx] = gathered\n b_kn = (\n tile_bytes.reshape(k_tiles, n_tiles, 8, 64)\n .permute(0, 2, 1, 3)\n .reshape(k // 2, n)\n )\n\n lut = torch.tensor(_E2M1_VALUES, dtype=torch.float32, device=device)\n b_int = b_kn.to(torch.int64)\n codes = torch.stack([b_int & 0xF, (b_int >> 4) & 0xF], dim=1).reshape(k, n)\n # S0E5M3 -> fp16: the byte is the top 8 bits of the fp16 bit pattern.\n sf = (\n (b_descale.to(torch.int16) << 7)\n .view(torch.float16)\n .to(torch.float32)\n .repeat_interleave(block_size, dim=0)\n )\n return _bf16_fp4_matmul(a, lut[codes] * sf, alpha)\n", + "check": "def _fp4_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.97,\n):\n from flashinfer.trace import default_check\n\n # Matches tests/gemm/test_mm_fp4.py.\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_fp4_cute_dsl_init(\n *,\n M: int,\n N: int = 2048,\n K: int = 7168,\n block_size: int = 16,\n K_div_16: int = 0, # derived\n K_div_block_size: int = 0, # derived\n N_mul_2: int = 0, # derived\n device: str = \"cuda\",\n seed: int = 0,\n):\n \"\"\"Build inputs for ``flashinfer.mm_bf16_fp4`` (cute-DSL backend).\n\n Sourced from ``tests/gemm/test_mm_bf16_fp4.py``: quantize a randn\n bf16 weight via ``flashinfer.nvfp4_quantize`` (layout_128x4), then\n repack with ``prepare_bf16_fp4_weights``. Requires SM100+ at\n runtime; CPU smoke tests skip.\n \"\"\"\n del K_div_16, K_div_block_size, N_mul_2\n from flashinfer import ( # noqa: PLC0415\n mm_bf16_fp4,\n nvfp4_quantize,\n prepare_bf16_fp4_weights,\n )\n from flashinfer.quantization.fp4_quantization import SfLayout # noqa: PLC0415\n\n if not torch.cuda.is_available() or torch.device(device).type != \"cuda\":\n raise NotImplementedError(\"mm_bf16_fp4 init requires a CUDA device\")\n major, minor = torch.cuda.get_device_capability(torch.device(device))\n if not mm_bf16_fp4.is_backend_supported(\"cute-dsl\", major * 10 + minor):\n raise NotImplementedError(f\"mm_bf16_fp4 is not supported on SM{major}{minor}\")\n\n torch.manual_seed(seed)\n a = torch.randn(M, K, dtype=torch.bfloat16, device=device)\n w = torch.randn(N, K, dtype=torch.bfloat16, device=device)\n g_w = (448.0 * 6.0) / w.float().abs().nan_to_num().max()\n b_fp4, b_sf = nvfp4_quantize(\n w, g_w, sfLayout=SfLayout.layout_128x4, do_shuffle=False, backend=\"cute-dsl\"\n )\n alpha = torch.tensor([1.0 / g_w.item()], dtype=torch.float32, device=device)\n b_p, sf_p, alpha_p = prepare_bf16_fp4_weights(\n b_fp4, b_sf, alpha, backend=\"cute-dsl\", block_size=block_size\n )\n return {\n \"a\": a,\n \"b\": b_p,\n \"b_descale\": sf_p,\n \"alpha\": alpha_p,\n \"backend\": \"cute-dsl\",\n \"block_size\": int(block_size),\n }\n" +} diff --git a/tests/trace/test_mm_bf16_fp4_reference_correctness.py b/tests/trace/test_mm_bf16_fp4_reference_correctness.py new file mode 100644 index 00000000000..b81b3aecb03 --- /dev/null +++ b/tests/trace/test_mm_bf16_fp4_reference_correctness.py @@ -0,0 +1,63 @@ +"""Reference correctness test for the mm_bf16_fp4 trace API.""" + +import pytest +import torch + +from tests.trace.reference_utils import ( + _assert_finite, + _check, +) + + +@pytest.mark.parametrize("backend", ["cudnn", "cute-dsl"]) +@pytest.mark.parametrize( + "shape_kwargs", [dict(M=32, N=1024, K=1024), dict(M=16, N=2048, K=512)] +) +def test_mm_bf16_fp4_reference_correctness(backend, shape_kwargs): + """flashinfer.mm_bf16_fp4 kernel vs reference (dequant + matmul). + + The trace inits build *prepared* (backend-specific) weights via + ``prepare_bf16_fp4_weights``; each backend's reference dequantizes + that prepared layout directly (the cute-dsl one inverts the MMA tile + permutation and decodes S0E5M3 scales). + """ + import flashinfer + from flashinfer.trace.templates.gemm import ( + mm_bf16_fp4_cudnn_trace, + mm_bf16_fp4_cute_dsl_trace, + ) + + tpl = { + "cudnn": mm_bf16_fp4_cudnn_trace, + "cute-dsl": mm_bf16_fp4_cute_dsl_trace, + }[backend] + try: + inputs = tpl.init(**shape_kwargs) + api = flashinfer.mm_bf16_fp4( + inputs["a"], + inputs["b"], + inputs["b_descale"], + inputs["alpha"], + backend=backend, + block_size=inputs["block_size"], + ) + except Exception as exc: + pytest.skip(f"mm_bf16_fp4 ({backend}) unavailable: {exc}") + _assert_finite(inputs["a"]) + ref = tpl.reference( + inputs["a"], + inputs["b"], + inputs["b_descale"], + inputs["alpha"], + block_size=inputs["block_size"], + ) + _assert_finite(api, ref) + _check( + tpl, + ref.to(api.dtype), + api, + max_mismatch_pct=100.0, + min_cos_sim=0.99, + ) + if torch.cuda.is_available(): + torch.cuda.synchronize()