flashinfer.gemm.mm_nvfp4_svdquant

flashinfer.gemm.mm_nvfp4_svdquant(a: Tensor, b: Tensor, a_sf: Tensor, b_sf: Tensor, alpha: Tensor, d: Tensor, l1: Tensor, bias: Tensor | None = None, out: Tensor | None = None, backend: Literal['cutlass', 'cute-dsl', 'cute-dsl-unfused', 'auto'] = 'auto', enable_pdl: bool | None = None) → Tensor

SVDQuant NVFP4 GEMM: out = alpha * (a @ bᵀ + d @ l1ᵀ) [+ bias].

On SM100/SM103, CUTLASS fuses the block-scaled NVFP4 residual GEMM with the rank-r BF16 LoRA-up correction and optional bias. On SM120/SM121, "cute-dsl" fuses the correction and bias into the b12x CuTe DSL kernel’s FP32 accumulator epilogue, while "cute-dsl-unfused" retains the compositional implementation as a differential oracle and optional autotuning candidate. The LoRA rank r is inferred from the d/l1 shapes and must be a positive multiple of 32 (ranks 32-128 are validated). 1/alpha must be folded into l1 by the caller (l1 = svdquant_lora_b / alpha), so both backends yield the correction at its original scale.

Parameters:

  • a (torch.Tensor) – Quantized activation, shape (m, k // 2) uint8 (packed e2m1), row-major. Produce it with nvfp4_quantize_smooth() (which folds the SVDQuant pre_quant_scale into the quantization).

  • b (torch.Tensor) – Quantized residual weight, shape (n, k // 2) uint8 (packed e2m1), row-major (i.e. the GEMM computes a @ bᵀ).

  • a_sf (torch.Tensor) – Activation block scales, uint8 (ue4m3) in the 128x4 swizzled layout, numel >= ceil(m / 128) * 128 * ceil(k / 16 / 4) * 4.

  • b_sf (torch.Tensor) – Weight block scales, same layout as a_sf with n rows.

  • alpha (torch.Tensor) – Per-tensor residual dequantization scale in a non-empty float32 device tensor. For compatibility with pooled scalar buffers, only the first element is consumed; backend runners receive a one-element view.

  • d (torch.Tensor) – LoRA-down output x_hat @ L2ᵀ, shape (m, r) bf16, contiguous and 16-byte aligned (TMA). Compute it as x @ (pre_quant_scale[:, None] * L2ᵀ) in bf16.

  • l1 (torch.Tensor) – LoRA-up weight pre-divided by alpha, shape (n, r) bf16 (same rank as d).

  • bias (Optional[torch.Tensor]) – Optional per-column bias, shape (n,) bf16. Fused by CUTLASS and the SM120/SM121 CuTe DSL kernel.

  • out (Optional[torch.Tensor]) – Output tensor, shape (m, n) bf16; allocated when None.

  • backend (Literal["cutlass", "cute-dsl", "cute-dsl-unfused", "auto"]) – "cutlass" selects the fused SM100/SM103 implementation; "cute-dsl" selects the fused SM120/SM121 implementation; "cute-dsl-unfused" selects its compositional reference path; "auto" (default) selects by compute capability. On SM120/SM121, fused and unfused are compared only while autotuning is enabled; otherwise the fused-first runner is selected.

  • enable_pdl (Optional[bool]) – Whether to launch with Programmatic Dependent Launch. Defaults to the device default.

Returns:

out – Output tensor, shape (m, n) bf16.

Return type:

torch.Tensor