flashinfer.gemm.prepare_nvfp4_lora_adapter

flashinfer.gemm.prepare_nvfp4_lora_adapter(lora_b: Tensor, lora_scale: Tensor, alpha: Tensor, *, out: PreparedNvFp4LoraAdapter | None = None) → PreparedNvFp4LoraAdapter

Prepare (fold) one LoRA adapter for the NVFP4 prepared path: b_folded = bf16(lora_b * lora_scale / alpha).

This is the adapter-switch cost of the prepared path (one small elementwise kernel, counted as the l1 fold of mm_nvfp4_svdquant() is counted); the per-call path then does no host read and no fold. The fold kernel is PyTorch’s own elementwise mul (bf16 x fp32 scalar in fp32 opmath, IEEE round-to-nearest-even to bf16; PyTorch is built without --use_fast_math): ratio = div.rn.f32(fp32(lora_scale), fp32(alpha[0])), b_folded[i] = bf16_rne(fp32(lora_b[i]) * ratio).

Parameters:

  • lora_b (torch.Tensor) – LoRA-up weight, bf16 [N, r] row-major, contiguous (the raw path’s lora_b).

  • lora_scale (torch.Tensor) – fp32 device tensor of shape (1,) – s_lora as on the raw path (0 folds to the zero adapter).

  • alpha (torch.Tensor) – fp32 device tensor of shape (1,): the base scale alpha_base / (a_global_scale * b_global_scale) the call passes as alpha – the same tensor object (alpha is fixed per prepared adapter). alpha == 0 is refused (s / alpha is undefined); the alpha == 0 LoRA-only call is the raw path.

  • out (PreparedNvFp4LoraAdapter, optional) – Refold IN PLACE into out.b_folded with NO host read (graph-capturable; capturing this refold plus the call is the adapter-switch protocol inside a CUDA graph): the same fp32 ratio computed on the device, the same fp32 product and the same single bf16 rounding – bitwise the eager form. alpha must be out’s own alpha tensor; lora_scale may be a new tensor (the adapter then carries it); lora_b must keep out’s [N, r] and device. The source pins are refreshed.

Returns:

A new adapter (out=None: the single documented host read of the adapter switch – one device-to-host copy of (lora_scale, alpha) – checks alpha != 0 and a finite ratio, else ValueError), or out.

Return type:

PreparedNvFp4LoraAdapter