ffma.py

"""FFMA matmul helpers for Pascal sm_61.

Background: under `autocast(bf16)` PyTorch lowers `x @ W` to
`magma_sgemmEx_kernel<float, __nv_bfloat16, ...>`. That kernel does a
serialised bf16->fp32 convert pass before the fp32 multiply-add, which
on sm_61 (no bf16 tensor cores) is ~25% slower than calling plain
cuBLAS `sgemm` directly. We pin the matmul kernel to fp32 so PyTorch
dispatches real FFMA `sgemm`, but the autograd tape only ever sees
the original (bf16) `x` and `W` — the fp32 cast is a transient
kernel-input, freed before forward returns. Backward re-casts the
saved bf16 tensors to fp32 just for the two backward GEMMs.

Memory profile is identical to a plain bf16 `addmm`:
  - tape carries bf16 x and bf16 W (no fp32 doubles)
  - forward output is bf16
  - transient fp32 working buffers for x, W, y exist only inside
    the autograd Function call and are freed when it returns.
"""

import torch


def _bf16_prefers_native_mm(x: torch.Tensor, W: torch.Tensor) -> bool:
    """Return True when bf16 matmul should stay on the native kernel path.

    FFMA forcing exists for old GPUs where bf16 matmul lowers poorly and a
    transient fp32 sgemm is faster. On newer architectures with native bf16
    acceleration, forcing fp32 throws away the better path.
    """
    if x.dtype != torch.bfloat16:
        return False
    if not x.is_cuda or not W.is_cuda:
        return False
    if not hasattr(torch.cuda, "is_bf16_supported"):
        return False
    try:
        if not torch.cuda.is_bf16_supported():
            return False
        major, _minor = torch.cuda.get_device_capability(x.device)
    except Exception:
        return False
    return major >= 8


class _FFMAAddmm(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, W, b):
        device_type = 'cuda' if x.is_cuda else 'cpu'
        with torch.amp.autocast(device_type, enabled=False):
            x32 = x.float() if x.dtype != torch.float32 else x
            W32 = W.float() if W.dtype != torch.float32 else W
            if x32.dim() > 2:
                x2 = x32.reshape(-1, x32.shape[-1])
                y32 = torch.mm(x2, W32)
            else:
                y32 = torch.mm(x32, W32)
            if b is not None:
                b32 = b.float() if b.dtype != torch.float32 else b
                y32 = y32 + b32
            y = y32.to(x.dtype) if x.dtype != torch.float32 else y32
            if x.dim() > 2:
                y = y.reshape(*x.shape[:-1], W.shape[1])
        # Save the ORIGINAL bf16 tensors — never the fp32 casts.
        ctx.save_for_backward(x, W)
        ctx.has_bias = b is not None
        ctx.x_dtype = x.dtype
        ctx.W_dtype = W.dtype
        ctx.b_dtype = b.dtype if b is not None else None
        return y

    @staticmethod
    def backward(ctx, gy):
        x, W = ctx.saved_tensors
        device_type = 'cuda' if gy.is_cuda else 'cpu'
        with torch.amp.autocast(device_type, enabled=False):
            gy32 = gy.float() if gy.dtype != torch.float32 else gy
            x32 = x.float() if x.dtype != torch.float32 else x
            W32 = W.float() if W.dtype != torch.float32 else W
            gy2 = gy32.reshape(-1, gy32.shape[-1]) if gy32.dim() > 2 else gy32
            x2 = x32.reshape(-1, x32.shape[-1]) if x32.dim() > 2 else x32
            gx2 = torch.mm(gy2, W32.t())
            gW = torch.mm(x2.t(), gy2)
            if gy32.dim() > 2:
                gx = gx2.reshape(*gy32.shape[:-1], W32.shape[0])
            else:
                gx = gx2
            gb = gy2.sum(dim=0) if ctx.has_bias else None
        if gx.dtype != ctx.x_dtype:
            gx = gx.to(ctx.x_dtype)
        if gW.dtype != ctx.W_dtype:
            gW = gW.to(ctx.W_dtype)
        if gb is not None and gb.dtype != ctx.b_dtype:
            gb = gb.to(ctx.b_dtype)
        return gx, gW, gb


def ffma_addmm(x: torch.Tensor, W: torch.Tensor, b) -> torch.Tensor:
    """`x @ W + b` routed through fp32 cuBLAS `sgemm` (FFMA) on Pascal.

    `W` has shape `[in, out]` (NOT the `nn.Linear` `[out, in]` layout).
    Bias may be `None`. Tape carries bf16 — see module docstring.
    """
    if _bf16_prefers_native_mm(x, W):
        y = torch.matmul(x, W)
        if b is not None:
            y = y + b
        return y
    return _FFMAAddmm.apply(x, W, b)


def ffma_linear(proj: torch.nn.Linear, x: torch.Tensor) -> torch.Tensor:
    """Drop-in replacement for `proj(x)` that lands on FFMA `sgemm`.

    `proj.weight` has shape `[out, in]`. We pass `proj.weight.t()` (a
    differentiable view, no copy) so `_FFMAAddmm` sees `[in, out]` and
    autograd flows the weight gradient back through the `t()` view.
    """
    return _FFMAAddmm.apply(x, proj.weight.t(), proj.bias)