flashinfer.gemm.mm_nvfp4_lora

flashinfer.gemm.mm_nvfp4_lora(a: Tensor, b: Tensor, a_sf: Tensor, b_sf: Tensor, alpha: Tensor, d: Tensor | None = None, lora_b: Tensor | None = None, lora_scale: Tensor | None = None, *, lora_adapter: PreparedNvFp4LoraAdapter | None = None, bias: Tensor | None = None, out: Tensor | None = None, out_dtype: dtype = torch.bfloat16, backend: Literal['cake', 'auto'] = 'cake', enable_pdl: bool | None = None) → Tensor

NVFP4 base GEMM with a runtime BF16 LoRA correction.

out = alpha * (dequant(a) @ dequant(b)^T) + lora_scale * (d @ lora_b^T) [+ bias]

The NVFP4 operands are consumed exactly as mm_fp4() does (E2M1 codes packed two per byte, UE4M3 block scales of 16 along K in the 128x4 swizzled layout produced by fp4_quantize()); the correction is computed in fp32 inside the same kernel and the result is rounded once to bf16. The adapter arrives in one of two forms:

  • Raw path (lora_b + lora_scale): the BF16 factors are consumed as given – d (the LoRA-down output x @ A^T) and lora_b (the LoRA-up weight) are never re-rounded or pre-scaled – and alpha / lora_scale are independent fp32 device scalars (no folding of either into the weights). Swapping the adapter is swapping pointers (0-preparation switching).

  • Prepared path (lora_adapter=, from prepare_nvfp4_lora_adapter()): the adapter was folded once at the switch, b_folded = bf16(lora_b * lora_scale / alpha); the kernel accumulates the base product first and applies alpha once in the fp32 epilogue – the numerics of mm_nvfp4_svdquant() (lora_b re-rounded once to bf16; on the measured contract rows SQNR 53.8 dB vs 55.6 dB on the raw path, max abs error vs fp64 0.038-0.042 vs 0.031). Every prepared call runs the prepared-path program (the raw dispatch table is untouched) and keeps the snapshot checks of PreparedNvFp4LoraAdapter. Passing both forms raises ValueError; neither form is the adapter-off call.

Parameters:

  • a (torch.Tensor) – Quantized activations, uint8 [M, K // 2] (two E2M1 codes per byte), row-major.

  • b (torch.Tensor) – Quantized base weights, uint8 [N, K // 2], row-major (a @ b^T).

  • a_sf (torch.Tensor) – UE4M3 block scales of a / b as uint8 in the 128x4 swizzled layout (rows padded to 128); larger contiguous backing buffers are accepted.

  • b_sf (torch.Tensor) – UE4M3 block scales of a / b as uint8 in the 128x4 swizzled layout (rows padded to 128); larger contiguous backing buffers are accepted.

  • alpha (torch.Tensor) – fp32 device tensor whose first element is the base scale alpha_base / (a_global_scale * b_global_scale).

  • d (torch.Tensor or None) – LoRA-down output, bf16 [M, r] row-major, contiguous (16-byte aligned rows).

  • lora_b (torch.Tensor or None) – LoRA-up weight, bf16 [N, r] row-major, contiguous. A [r, N] factor must be transposed and made contiguous by the caller.

  • lora_scale (torch.Tensor or None) – One-element fp32 device tensor s_lora (the LoRA alpha / r and any other multiplier), independent of alpha. 0 and None both take the zero-adapter fast path (no LoRA work, bitwise identical outputs): lora_scale=None disables the adapter even if d / lora_b are passed; pass a device scalar to apply it.

  • lora_adapter (PreparedNvFp4LoraAdapter, optional) – The prepared path: an adapter from prepare_nvfp4_lora_adapter() (lora_b and lora_scale must then be None; d is still required). Its alpha must be this call’s alpha tensor (re-prepare when alpha changes) and its source tensors must not have moved since the prepare (StaleAdapterError otherwise). A prepared adapter with lora_scale == 0 produces the adapter-off output, bitwise.

  • bias (torch.Tensor, optional) – bf16 [N], added in fp32 before the single rounding.

  • out (torch.Tensor, optional) – Caller-owned bf16 [M, N] output; allocated when None.

  • out_dtype (torch.dtype) – Only torch.bfloat16 is supported.

  • backend ({"cake", "auto"}) – The generated Cake programs are the only implementation.

  • enable_pdl (bool, optional) – The programs are built with programmatic dependent launch; None or True.

Returns:

out (bf16 [M, N]).

Return type:

torch.Tensor

Notes

  • Architectures: SM100 (B200) and SM103 (B300) only; SM120 / SM121 are not supported (the kernel uses tcgen05 / TMEM).

  • Shape rules: M >= 1, N >= 1, K a multiple of 256; the LoRA rank r = d.shape[1] must be one of (32, 64, 128). Violations raise ValueError naming the violated requirement.

  • Tensor parallelism (host contract): column-parallel (N sharded) ranks each hold their N shard of b / b_sf / lora_b / bias and the full d, and each adds its own correction; row-parallel (K sharded) ranks hold their K shard of a / b (codes and scales), the full d (reduced by the caller or computed from the full x) and the full lora_b, and only the owner rank passes lora_scale / bias (the others pass lora_scale = 0 and bias = None) so the correction is counted exactly once before the all-reduce of the base partials. On the prepared path every rank prepares its own adapter from its shard (the non-owner ranks of a row-parallel layout with lora_scale = 0, which folds to the zero adapter).

  • CUDA graphs, raw path: a captured call keeps the adapter pointers captured. Swapping d / lora_b / lora_scale for new tensors after capture does not affect replays; write the new adapter into the captured tensors (copy_) instead. Prepared path: a captured call keeps the prepared adapter captured; the adapter switch is prepare_nvfp4_lora_adapter(lora_b, lora_scale, alpha, out=adapter) captured right before the call (an in-place refold without a host read), after which a copy_ into the captured lora_b / lora_scale is folded by every replay. Outside a graph, a copy_ into a source tensor makes the adapter stale (StaleAdapterError) until it is re-prepared.