flashinfer.gemm.mm_fp8_lora

flashinfer.gemm.mm_fp8_lora(a: Tensor, b: Tensor, a_scale: Tensor, b_scale: Tensor, d: Tensor | None = None, lora_b: Tensor | None = None, lora_scale: Tensor | 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

FP8 (E4M3, per-tensor scales) base GEMM with a runtime BF16 LoRA correction.

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

The base operands are consumed as mm_fp8() consumes them (float8_e4m3fn codes with one fp32 scale each; a global alpha_base is carried in b_scale). The correction is computed from the BF16 factors as given – never re-rounded or pre-scaled – in fp32 inside the same kernel, and the result is rounded once to bf16. The base scales and lora_scale are independent fp32 device scalars.

Parameters:

  • a (torch.Tensor) – Quantized activations, float8_e4m3fn [M, K], row-major.

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

  • a_scale (torch.Tensor) – One-element fp32 device tensors (per-tensor dequantization scales).

  • b_scale (torch.Tensor) – One-element fp32 device tensors (per-tensor dequantization scales).

  • 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, independent of the base scales. 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.

  • 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: SM90 (H100 / H200), SM100 (B200) and SM103 (B300); SM120 / SM121 are not supported.

  • Shape rules: M >= 1, N a positive multiple of 8; K a multiple of 128 on SM100 / SM103 and of 16 on SM90 (where additionally M * N <= 2**32); the LoRA rank r = d.shape[1] must be one of (32, 64, 128). Violations raise ValueError naming the violated requirement. On SM100 / SM103 the pair-tile width is resolved from the shape and the device SM count; it is not a caller knob. On SM100 / SM103 long large-K runs (at least 10 tiles per resident CTA pair, K >= 16384, a proper block inside the tile grid) run the dynamic-tile-scheduler program with the blocked raster (fp8_blackwell_l2_point); on SM103 the longest of them (at least 20 waves of 512 x 256 tiles over the resident pairs, K >= 32768) run it on the 512-row pair tile (fp8_blackwell_tile_m); every other call runs the static persistent program.

  • Tensor parallelism (host contract): column-parallel (N sharded) ranks each hold their N shard of b / 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, the full d and 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.

  • CUDA graphs: 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.