flashinfer.gemm.svdquant_linear

flashinfer.gemm.svdquant_linear(x: Tensor, weight_fp4: Tensor, weight_sf: Tensor, alpha: Tensor, pre_quant_scale: Tensor, l2t_smoothed: Tensor, l1_scaled: Tensor, global_scale: Tensor, bias: Tensor | None = None, enable_pdl: bool | None = None) → Tensor

The full SVDQuant linear operator: y = x_hat @ (R + L1 @ L2)ᵀ [+ bias] where x_hat = x * pre_quant_scale and R is the NVFP4-quantized residual weight.

Runs the three-step chain this library’s kernels are designed for:

  1. xq, x_sf = nvfp4_quantize_smooth(x, pre_quant_scale, global_scale)

  2. down = x @ l2t_smoothed (plain BF16 GEMM; l2t_smoothed = pre_quant_scale[:, None] * L2ᵀ)

  3. mm_nvfp4_svdquant(xq, weight_fp4, x_sf, weight_sf, alpha, down, l1_scaled, bias)

The invariant per-layer transforms must be prepared offline by the caller: l2t_smoothed = (pre_quant_scale[:, None] * svdquant_lora_a.T).to(bf16) with shape (k, r) and l1_scaled = (svdquant_lora_b / alpha).to(bf16) with shape (n, r), where the LoRA rank r is a positive multiple of 32.

Parameters:

  • x (torch.Tensor) – Input activation, shape (m, k) bf16.

  • weight_fp4 (torch.Tensor) – NVFP4 residual weight, shape (n, k // 2) uint8 (packed e2m1).

  • weight_sf (torch.Tensor) – Weight block scales, uint8 (ue4m3), 128x4 swizzled layout.

  • alpha (torch.Tensor) – Per-tensor residual dequantization scale, float32 device scalar.

  • pre_quant_scale (torch.Tensor) – Per-input-channel smoothing scale, shape (k,) bf16.

  • l2t_smoothed (torch.Tensor) – pre_quant_scale[:, None] * L2ᵀ, shape (k, r) bf16.

  • l1_scaled (torch.Tensor) – L1 / alpha, shape (n, r) bf16.

  • global_scale (torch.Tensor) – Activation global scale, float32 device scalar.

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

  • 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