[feat]: Add NVFP4QAT linear layer (Attn-QAT 3/12) (hao-ai-lab#1350)

Co-authored-by: jzhang38 <42993249+jzhang38@users.noreply.github.com>
Co-authored-by: RandNMR73 <99706358+RandNMR73@users.noreply.github.com>

Author: 3 people

fastvideo/layers/fp4linear.py

@@ -0,0 +1,116 @@
+from typing import Any
+
+import torch
+
+try:
+    import flashinfer
+except ImportError:
+    flashinfer = None
+
+
+def _require_flashinfer() -> Any:
+    if flashinfer is None:
+        raise ImportError("flashinfer is required for FP4 linear layers. "
+                          "Please install flashinfer to use this path.")
+    return flashinfer
+
+
+@torch.compile
+def _global_sf(t: torch.Tensor) -> torch.Tensor:
+    maxabs = t.float().abs().nan_to_num().max()
+    maxabs = maxabs.clamp(min=1e-12)
+    return (448.0 * 6.0) / maxabs
+
+
+class _LinearFWD4BWD16Fn(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, x, weight, bias, backend="cutlass", block_size=16, use_128x4_sf_layout=True):
+        flashinfer_mod = _require_flashinfer()
+
+        # assert activation dtype
+        if x.dtype not in (torch.float16, torch.bfloat16):
+            x = x.to(dtype=torch.bfloat16)
+
+        # cast params (can be fp32) to activation dtype for quantization
+        weight_cast = weight.to(dtype=x.dtype)
+        bias_cast = bias.to(dtype=x.dtype) if bias is not None else None
+
+        # shapes
+        orig_shape = x.shape
+        k = weight_cast.shape[1]
+        n = weight_cast.shape[0]
+        x2d = x.reshape(-1, k).contiguous()
+        M = x2d.shape[0]
+
+        out2d = torch.empty((M, n), device=x.device, dtype=x.dtype)
+
+        a_sf_layout = (flashinfer_mod.SfLayout.layout_128x4
+                       if use_128x4_sf_layout else flashinfer_mod.SfLayout.layout_8x4)
+        global_sf_a = _global_sf(x2d)
+        global_sf_b = _global_sf(weight_cast)
+
+        a_fp4, a_inv_s = flashinfer_mod.nvfp4_quantize(
+            x2d,
+            global_sf_a,
+            sfLayout=a_sf_layout,
+            do_shuffle=False,
+        )
+        b_fp4, b_inv_s = flashinfer_mod.nvfp4_quantize(
+            weight_cast,
+            global_sf_b,
+            sfLayout=flashinfer_mod.SfLayout.layout_128x4,
+            do_shuffle=False,
+        )
+
+        alpha = 1.0 / (global_sf_a * global_sf_b)
+
+        flashinfer_mod.mm_fp4(
+            a_fp4,
+            b_fp4.T,
+            a_inv_s,
+            b_inv_s.T,
+            alpha,
+            x.dtype,
+            out2d,
+            block_size=block_size,
+            use_8x4_sf_layout=(not use_128x4_sf_layout),
+            backend=backend,
+        )
+
+        if bias_cast is not None:
+            out2d.add_(bias_cast)
+
+        # save tensors for backward (keep original dtypes)
+        ctx.save_for_backward(x2d, weight, bias)
+        ctx.k = k
+        ctx.n = n
+        ctx.orig_shape = orig_shape
+        return out2d.reshape(*orig_shape[:-1], n)
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        x2d, weight, bias = ctx.saved_tensors
+        M = x2d.shape[0]
+        n = ctx.n
+
+        grad_out_2d = grad_out.reshape(M, n).contiguous()
+
+        # cast to grad dtype for matmuls
+        weight_cast = weight.to(dtype=grad_out.dtype)
+        x_cast = x2d.to(dtype=grad_out.dtype)
+
+        grad_x = grad_out_2d.matmul(weight_cast).reshape(*ctx.orig_shape)
+        grad_w = grad_out_2d.t().matmul(x_cast)
+        grad_b = grad_out_2d.sum(dim=0) if bias is not None else None
+
+        # None for the three extra forward args
+        return grad_x, grad_w, grad_b, None, None, None
+
+
+def fp4_linear_forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None]:
+    # pass config **positionally**; autograd.Function.apply ignores kwargs
+    bias = self.bias if not self.skip_bias_add else None
+    output = _LinearFWD4BWD16Fn.apply(x, self.weight, bias, "cutlass", 16, True)
+    output_bias = self.bias if self.skip_bias_add else None
+    return output, output_bias