460f6e398e
Co-authored-by: jzhang38 <42993249+jzhang38@users.noreply.github.com> Co-authored-by: RandNMR73 <99706358+RandNMR73@users.noreply.github.com>
117 lines
3.5 KiB
Python
117 lines
3.5 KiB
Python
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
|