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'] = 'cutlass', enable_pdl: bool | None = None) → Tensor

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

The block-scaled NVFP4 residual GEMM is fused with the rank-r BF16 LoRA-up correction d @ l1ᵀ, computed by a second BF16 tcgen05 MMA into the same accumulator after the NVFP4 K-loop, plus an optional fused per-column bias. The LoRA rank r is inferred from the d/l1 shapes and must be a positive multiple of 32 (ranks 32-128 are validated). 1/alpha must be folded into l1 by the caller (l1 = svdquant_lora_b / alpha) so the epilogue out = alpha * acc + bias yields the correction unscaled.

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, float32, device scalar (numel >= 1).

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

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

  • bias (Optional[torch.Tensor]) – Optional per-column bias, shape (n,) bf16, fused in the epilogue.

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

  • backend (Literal["cutlass"]) – Only the CUTLASS backend exists.

  • 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