mirror of
https://github.com/unslothai/unsloth.git
synced 2026-10-25 07:51:01 +00:00
NVFP4: decode weight tiles inside the GEMM for small-batch forward and dX (#13840)
## Problem `nvfp4_linear` (compressed-tensors `nvfp4-pack-quantized` W4A16 checkpoints) dequantizes the full packed weight to a 16 bit copy on every forward and every input-gradient call, then runs `torch.matmul` on it. For decode and small batches that temporary dominates: a 4096 x 4096 layer allocates 33.7 MB per call for 16 activation rows. ## Change `unsloth/kernels/nvfp4.py`: - Two tiled Triton GEMMs, `X @ W^T` (forward) and `dY @ W` (dX), sharing one tile decoder: the weight is decoded one K tile at a time inside the GEMM, never materialized. fp32 accumulation. - Same rounding as the dense path: each element is the E2M1 value times (block scale / global scale) in fp32, cast once to the activation dtype. E2M1 is decoded by placing the 3 magnitude bits as fp16 exponent / mantissa bits (exact upcast, then an exact power of two rescale); E4M3 scale bytes use the existing `_e4m3_to_float`. - When the output has few tiles the reduction is split into fp32 partials (capped at 1/8 of the 16 bit weight size) and summed in a fixed order by a small kernel, so results are deterministic. - Any strides, masked tails, int64 offsets past 2^31 elements. The switch sits inside the existing compile-opaque forward / backward ops, so `torch.compile` behaviour is unchanged. - Dispatch: the fused path is used only up to a measured activation row count per GPU (sm80: 64, sm89: 256, sm120: 128). Larger inputs, shapes whose N or K is not a multiple of 64 (padded tiles lose to dense), and GPUs without a measured win (T4, H100, B200, sm103), ROCm, XPU and fp32 activations keep the dense path exactly as today. `UNSLOTH_NVFP4_FUSED=0` forces dense, `=1` forces fused, `UNSLOTH_NVFP4_FUSED_MAX_ROWS` overrides the row limit; both are read on the first NVFP4 call. On GPUs without a row limit the added check costs about 1 us per call before the unchanged dense path. ## Performance Fused vs dense (dequantize + `torch.matmul`) speedup, Colab, shapes 4096 x 4096, 14336 x 4096, 4096 x 14336, forward and dX, interleaved with a second dense arm as the A/A (median gap: L4 0.08%, A100 0%, RTX PRO 6000 0.01%): | GPU | M=1 | 8 | 32 | 64 | 128 | 256 | 512 | 2048 | |---|---|---|---|---|---|---|---|---| | L4 (sm89) | 2.9-5.6 | 3.0-5.5 | 2.6-5.6 | 2.0-4.4 | 1.9-3.3 | 1.24-1.84 | 0.88-1.09 | 0.69-0.79 | | A100 40GB (sm80) | 1.8-2.8 | 1.8-2.7 | 1.6-2.5 | 1.12-1.45 | 0.71-0.94 | 0.41-0.72 | 0.33-0.61 | 0.34-0.45 | | RTX PRO 6000 (sm120) | 1.9-3.8 | 2.0-3.7 | 1.5-2.5 | 1.24-1.83 | 1.14-1.47 | 0.87-1.07 | 0.76-0.87 | 0.61-0.68 | | B200 (sm100) | 0.63-1.03 | 0.60-1.03 | 0.58-0.79 | 0.50-0.67 | 0.41-0.81 | 0.42-0.87 | 0.32-0.57 | 0.22-0.26 | | T4 (sm75) | 0.07-0.75 | 0.07-0.65 | 0.04-0.07 | 0.02-0.03 | 0.01 | 0.01 | < 0.01 | < 0.01 | Each GPU's row limit is the last row count that won on every shape and op; B200 and T4 lose everywhere, so they keep the dense path and no default case is slower than today. Re-timed after the edge case fixes in this PR (masked padding, int64 offsets, scale bound), same Colab GPUs, 5 interleaved rounds with an A/A twin: kernel time matches the first version within noise (geomean 0.999 L4, 1.004 A100, 1.001 RTX PRO 6000). Inside the row gate the worst end to end fused / dense ratio is 0.80 on L4, 0.90 on A100 and 0.82 on RTX PRO 6000; just past it A100 at 128 rows is 1.06 to 1.53 and RTX PRO 6000 at 256 rows is 0.94 to 1.14, so the limits stand. A per-call host cost the fixes added (strided offset bound, about 2 us) is skipped for contiguous operands, back to the first version's per-call time. Unaligned shapes (4096 x 4112, 4100 x 4096) are correct through the masked path but up to 2.8x slower than dense in dX, hence the 64 multiple requirement. Peak memory (4096 x 4096, 16 rows): 2.2 MB forward and 4.3 MB dX, against 33.7 MB on main. ## Numerics Against an fp64 oracle (B200, 4096 to 14336 wide, M 1 to 2048) the fused and dense paths have the same relative error (bf16 1.62e-3 to 1.66e-3, fp16 2.07e-4 to 2.09e-4 against the rounded weight); 97.7% to 100% of outputs are bit-identical to dense, the rest differ by about one output ulp (accumulation order). ## Tests New `tests/test_nvfp4_fused_gemm.py` (48): all 16 E2M1 codes exact through both kernels; every finite E4M3 block scale against global scales 1e-3 and 3.7e4 matching the dense weight bit for bit; NaN scales; tails (K = 40, 48, 272, 400, 1040, odd M and N); split reduction and its determinism; strided / transposed activations, 3D inputs, sliced packed weights, fp32 block scales; fp16 and bf16; kill switch, row limit and fallbacks; autograd engaging both kernels; a 4-step AdamW LoRA parity run fused vs dense; `torch.compile` with no graph breaks; peak memory; offset width from the strided span, operands on two devices, degenerate operands, extra scale columns, double backward, inf scale in a partial last group, zero global scale with a partial row tile, column-strided E4M3 scales, unaligned shapes kept dense, flags read once. On B200 (torch 2.14 / Triton 3.8, two GPUs) 48 passed with `tests/test_triton_fp_fusion.py` 4 passed. The original 38 fail on main. With `tests/test_nvfp4_kernel.py`: 65 passed on torch 2.7.1 / Triton 3.3 and 38 passed + 1 skipped (no compressed-tensors) on torch 2.14 / Triton 3.8; Colab torch 2.11 / Triton 3.6: 38 / 38 on L4, A100 and RTX PRO 6000, 23 / 23 on T4 (fp16). Not covered: H100 and sm103 timing (dense by default there), torch 2.6 / Triton 3.2. The `fast_lora` training path (`matmul_lora`) still dequantizes through `fast_dequantize`; moving it over is a separate change. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han
GitHub
parent
0737d7dc4f
commit
fa15e69126
2 changed files with 927 additions and 6 deletions
472
tests/test_nvfp4_fused_gemm.py
Normal file
472
tests/test_nvfp4_fused_gemm.py
Normal file
|
|
@ -0,0 +1,472 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team.
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
torch = pytest.importorskip("torch")
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.gpu,
|
||||
pytest.mark.skipif(not torch.cuda.is_available(), reason = "needs CUDA"),
|
||||
pytest.mark.skipif(torch.version.hip is not None, reason = "the fused path is CUDA only"),
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture(scope = "module")
|
||||
def N():
|
||||
os.environ.setdefault("UNSLOTH_IS_PRESENT", "1")
|
||||
from unsloth.kernels import nvfp4
|
||||
return nvfp4
|
||||
|
||||
|
||||
@pytest.fixture(autouse = True)
|
||||
def live_env(N, monkeypatch):
|
||||
# nvfp4 reads its flags once per process; these tests flip them between calls.
|
||||
cached = N._fused_env
|
||||
monkeypatch.setattr(N, "_fused_env", N._read_fused_env)
|
||||
return cached
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fused(monkeypatch):
|
||||
monkeypatch.setenv("UNSLOTH_NVFP4_FUSED", "1")
|
||||
|
||||
|
||||
def _dtypes():
|
||||
if torch.cuda.is_available() and torch.cuda.get_device_capability() >= (8, 0):
|
||||
return [torch.bfloat16, torch.float16]
|
||||
return [torch.float16]
|
||||
|
||||
|
||||
def _random_nvfp4(
|
||||
rows,
|
||||
cols,
|
||||
global_scale = 7.3,
|
||||
seed = 0,
|
||||
):
|
||||
g = torch.Generator(device = "cpu").manual_seed(seed)
|
||||
packed = torch.randint(0, 256, (rows, cols // 2), dtype = torch.uint8, generator = g)
|
||||
scale = (torch.rand(rows, -(-cols // 16), generator = g) * 3 + 0.01).to(torch.float8_e4m3fn)
|
||||
return packed.cuda(), scale.cuda(), torch.tensor([global_scale], device = "cuda")
|
||||
|
||||
|
||||
def _calls(N):
|
||||
return dict(N._FUSED_CALLS)
|
||||
|
||||
|
||||
def _dense(N, X, packed, scale, gs, monkeypatch):
|
||||
monkeypatch.setenv("UNSLOTH_NVFP4_FUSED", "0")
|
||||
before = _calls(N)
|
||||
Y = N._nvfp4_matmul_fwd(X, packed, scale, gs)
|
||||
assert _calls(N) == before
|
||||
monkeypatch.setenv("UNSLOTH_NVFP4_FUSED", "1")
|
||||
return Y
|
||||
|
||||
|
||||
def _fwd(N, X, packed, scale, gs):
|
||||
before = N._FUSED_CALLS["fwd"]
|
||||
Y = N._nvfp4_matmul_fwd(X, packed, scale, gs)
|
||||
assert N._FUSED_CALLS["fwd"] == before + 1, "fused forward did not run"
|
||||
return Y
|
||||
|
||||
|
||||
def _dx(N, G, packed, scale, gs):
|
||||
before = N._FUSED_CALLS["bwd"]
|
||||
dX = N._nvfp4_matmul_bwd(G, packed, scale, gs)
|
||||
assert N._FUSED_CALLS["bwd"] == before + 1, "fused dX did not run"
|
||||
return dX
|
||||
|
||||
|
||||
def _same(a, b):
|
||||
assert a.shape == b.shape and a.dtype == b.dtype
|
||||
assert torch.equal(a.isnan(), b.isnan())
|
||||
assert torch.equal(a.nan_to_num(), b.nan_to_num())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", _dtypes())
|
||||
def test_every_e2m1_code_decodes_exactly(N, fused, dtype):
|
||||
# Identity activations read the decoded weight back exactly through both kernels.
|
||||
packed = torch.arange(256, dtype = torch.uint8, device = "cuda").reshape(16, 16)
|
||||
scale = torch.ones(16, 2, device = "cuda").to(torch.float8_e4m3fn)
|
||||
gs = torch.ones(1, device = "cuda")
|
||||
W = N.nvfp4_dequantize(packed, scale, gs, dtype)
|
||||
lut = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]
|
||||
expect = torch.tensor(
|
||||
[[(-1 if c & 8 else 1) * lut[c & 7] for c in (b & 15, b >> 4)] for b in range(256)]
|
||||
).reshape(16, 32)
|
||||
assert torch.equal(W.float().cpu(), expect)
|
||||
_same(_fwd(N, torch.eye(32, device = "cuda", dtype = dtype), packed, scale, gs), W.t().contiguous())
|
||||
_same(_dx(N, torch.eye(16, device = "cuda", dtype = dtype), packed, scale, gs), W)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", _dtypes())
|
||||
@pytest.mark.parametrize("global_scale", [1.0, 1e-3, 3.7e4, 0.37])
|
||||
def test_scale_extremes_match_the_dense_weight(N, fused, monkeypatch, dtype, global_scale):
|
||||
# Every finite E4M3 scale vs tiny/huge global scales: bit-exact with dense (fp16 inf gives NaN on both).
|
||||
codes = torch.tensor([c for c in range(256) if c & 127 != 127], dtype = torch.uint8)
|
||||
rows = 32
|
||||
scale = (
|
||||
codes[torch.arange(rows * 16) % len(codes)]
|
||||
.reshape(rows, 16)
|
||||
.cuda()
|
||||
.view(torch.float8_e4m3fn)
|
||||
)
|
||||
packed = torch.full((rows, 128), 0x7F, dtype = torch.uint8, device = "cuda") # -6.0, 6.0
|
||||
packed[:, 1::2] = 0x1A # -1.0, 0.5
|
||||
gs = torch.tensor([global_scale], device = "cuda")
|
||||
W = N.nvfp4_dequantize(packed, scale, gs, dtype)
|
||||
finite = bool(W.isfinite().all())
|
||||
for col in (0, 1, 2, 3, 17, 255):
|
||||
X = torch.zeros(1, 256, device = "cuda", dtype = dtype)
|
||||
X[0, col] = 1
|
||||
Y = _fwd(N, X, packed, scale, gs)
|
||||
_same(Y, W[:, col][None] if finite else _dense(N, X, packed, scale, gs, monkeypatch))
|
||||
for row in (0, 5, 31):
|
||||
G = torch.zeros(1, rows, device = "cuda", dtype = dtype)
|
||||
G[0, row] = 1
|
||||
dX = _dx(N, G, packed, scale, gs)
|
||||
_same(dX, W[row][None] if finite else G @ W)
|
||||
|
||||
|
||||
def test_nan_block_scale_propagates_like_dense(N, fused, monkeypatch):
|
||||
dtype = _dtypes()[0]
|
||||
packed, scale, gs = _random_nvfp4(64, 128)
|
||||
raw = scale.view(torch.uint8).clone()
|
||||
raw[3, 2] = 0x7F
|
||||
raw[9, 0] = 0xFF
|
||||
scale = raw.view(torch.float8_e4m3fn)
|
||||
X = torch.randn(8, 128, device = "cuda", dtype = dtype)
|
||||
Y = _fwd(N, X, packed, scale, gs)
|
||||
ref = _dense(N, X, packed, scale, gs, monkeypatch)
|
||||
assert torch.equal(Y.isnan(), ref.isnan()) and Y.isnan().any()
|
||||
|
||||
|
||||
def _close_to_dense(a, ref):
|
||||
# Accumulation order differs from cuBLAS, so cancellation-sized outputs can move by an output ulp.
|
||||
torch.testing.assert_close(a, ref, atol = 2**-7 * ref.abs().max().item(), rtol = 1.6e-2)
|
||||
|
||||
|
||||
def _check(a, ref, exact):
|
||||
err_a, err_ref = _oracle_errors(a, ref, exact)
|
||||
assert err_a <= 1.25 * err_ref + 1e-6, (err_a, err_ref)
|
||||
_close_to_dense(a, ref)
|
||||
|
||||
|
||||
def _oracle_errors(Y, ref, exact):
|
||||
exact = exact.double()
|
||||
return (Y.double() - exact).norm().item() / exact.norm().item(), (
|
||||
ref.double() - exact
|
||||
).norm().item() / exact.norm().item()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", _dtypes())
|
||||
@pytest.mark.parametrize(
|
||||
"M, N_out, K",
|
||||
[(1, 96, 48), (5, 80, 40), (16, 64, 256), (37, 130, 400), (129, 257, 272), (300, 64, 1040)],
|
||||
)
|
||||
def test_forward_and_dx_match_dense_and_oracle(N, fused, monkeypatch, dtype, M, N_out, K):
|
||||
packed, scale, gs = _random_nvfp4(N_out, K, global_scale = 0.37, seed = M)
|
||||
W = N.nvfp4_dequantize(packed, scale, gs, dtype)
|
||||
X = torch.randn(M, K, device = "cuda", dtype = dtype)
|
||||
G = torch.randn(M, N_out, device = "cuda", dtype = dtype)
|
||||
|
||||
Y = _fwd(N, X, packed, scale, gs)
|
||||
ref = _dense(N, X, packed, scale, gs, monkeypatch)
|
||||
err_fused, err_dense = _oracle_errors(Y, ref, X.double() @ W.double().t())
|
||||
assert err_fused <= 1.25 * err_dense + 1e-6, (err_fused, err_dense)
|
||||
_close_to_dense(Y, ref)
|
||||
|
||||
dX = _dx(N, G, packed, scale, gs)
|
||||
monkeypatch.setenv("UNSLOTH_NVFP4_FUSED", "0")
|
||||
dref = N._nvfp4_matmul_bwd(G, packed, scale, gs)
|
||||
err_fused, err_dense = _oracle_errors(dX, dref, G.double() @ W.double())
|
||||
assert err_fused <= 1.25 * err_dense + 1e-6, (err_fused, err_dense)
|
||||
_close_to_dense(dX, dref)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", _dtypes())
|
||||
@pytest.mark.parametrize("M, N_out, K", [(4, 128, 4096), (70, 96, 3000), (3, 4096, 200)])
|
||||
def test_split_reduction_matches_unsplit(N, fused, monkeypatch, dtype, M, N_out, K):
|
||||
packed, scale, gs = _random_nvfp4(N_out, K, seed = 7)
|
||||
X = torch.randn(M, K, device = "cuda", dtype = dtype)
|
||||
G = torch.randn(M, N_out, device = "cuda", dtype = dtype)
|
||||
splits = (
|
||||
N._fused_config(M, N_out, K, X.device, False)[-1],
|
||||
N._fused_config(M, K, N_out, X.device, True)[-1],
|
||||
)
|
||||
assert max(splits) > 1, splits
|
||||
Y, dX = _fwd(N, X, packed, scale, gs), _dx(N, G, packed, scale, gs)
|
||||
monkeypatch.setattr(N, "_CONFIG_HOOK", lambda *a: N._fused_config(*a)[:-1] + (1,))
|
||||
_close_to_dense(Y, _fwd(N, X, packed, scale, gs))
|
||||
_close_to_dense(dX, _dx(N, G, packed, scale, gs))
|
||||
W = N.nvfp4_dequantize(packed, scale, gs, dtype)
|
||||
_close_to_dense(Y, X @ W.t())
|
||||
_close_to_dense(dX, G @ W)
|
||||
monkeypatch.undo()
|
||||
monkeypatch.setenv("UNSLOTH_NVFP4_FUSED", "1")
|
||||
assert torch.equal(Y, _fwd(N, X, packed, scale, gs)) and torch.equal(
|
||||
dX, _dx(N, G, packed, scale, gs)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", _dtypes())
|
||||
def test_strided_operands(N, fused, monkeypatch, dtype):
|
||||
big_packed, big_scale, gs = _random_nvfp4(72, 512, seed = 3)
|
||||
packed, scale = big_packed[:, :96], big_scale[:, :12]
|
||||
assert not packed.is_contiguous()
|
||||
X = torch.randn(192, 33, device = "cuda", dtype = dtype).t()
|
||||
assert X.stride(-1) != 1
|
||||
W = N.nvfp4_dequantize(packed, scale, gs, dtype).double()
|
||||
Y = _fwd(N, X, packed, scale, gs)
|
||||
_check(Y, _dense(N, X, packed, scale, gs, monkeypatch), X.double() @ W.t())
|
||||
X3 = torch.randn(3, 7, 192, device = "cuda", dtype = dtype)
|
||||
Y3 = _fwd(N, X3, packed, scale, gs)
|
||||
assert Y3.shape == (3, 7, 72)
|
||||
_check(Y3, _dense(N, X3, packed, scale, gs, monkeypatch), X3.double() @ W.t())
|
||||
G3 = torch.randn(3, 7, 72, device = "cuda", dtype = dtype)
|
||||
dX3 = _dx(N, G3, packed, scale, gs)
|
||||
assert dX3.shape == (3, 7, 192)
|
||||
_check(dX3, G3 @ W.to(dtype), G3.double() @ W)
|
||||
|
||||
|
||||
def test_float32_block_scales(N, fused, monkeypatch):
|
||||
dtype = _dtypes()[0]
|
||||
packed, scale, gs = _random_nvfp4(48, 160, seed = 4)
|
||||
scale = scale.float()
|
||||
X = torch.randn(9, 160, device = "cuda", dtype = dtype)
|
||||
W = N.nvfp4_dequantize(packed, scale, gs, dtype).double()
|
||||
_check(
|
||||
_fwd(N, X, packed, scale, gs),
|
||||
_dense(N, X, packed, scale, gs, monkeypatch),
|
||||
X.double() @ W.t(),
|
||||
)
|
||||
|
||||
|
||||
def test_kill_switch_and_fallbacks(N, monkeypatch):
|
||||
dtype = _dtypes()[0]
|
||||
packed, scale, gs = _random_nvfp4(64, 128)
|
||||
X = torch.randn(4, 128, device = "cuda", dtype = dtype)
|
||||
for env in ("0", "false"):
|
||||
monkeypatch.setenv("UNSLOTH_NVFP4_FUSED", env)
|
||||
before = _calls(N)
|
||||
N.nvfp4_linear(X, packed, scale, gs)
|
||||
assert _calls(N) == before
|
||||
monkeypatch.delenv("UNSLOTH_NVFP4_FUSED")
|
||||
monkeypatch.setenv("UNSLOTH_NVFP4_FUSED_MAX_ROWS", "4")
|
||||
before = N._FUSED_CALLS["fwd"]
|
||||
N.nvfp4_linear(X, packed, scale, gs)
|
||||
assert N._FUSED_CALLS["fwd"] == before + 1
|
||||
before = N._FUSED_CALLS["fwd"]
|
||||
N.nvfp4_linear(torch.cat([X, X]), packed, scale, gs)
|
||||
assert N._FUSED_CALLS["fwd"] == before # 8 rows: past the gate, dense
|
||||
before = _calls(N)
|
||||
N.nvfp4_linear(X.float(), packed, scale, gs)
|
||||
assert _calls(N) == before
|
||||
with pytest.raises(ValueError):
|
||||
N.nvfp4_linear(X, packed, scale[:, :2], gs)
|
||||
|
||||
|
||||
def test_degenerate_operands_stay_dense(N, fused):
|
||||
dtype = _dtypes()[0]
|
||||
packed, scale, gs = _random_nvfp4(64, 128)
|
||||
X = torch.randn(4, 128, device = "cuda", dtype = dtype)
|
||||
assert N._fused_ok(X, packed, scale, gs, 128)
|
||||
assert not N._fused_ok(X, packed, scale, gs[:0], 128) # empty global scale
|
||||
assert not N._fused_ok(X, packed[:0], scale[:0], gs, 128) # zero output features
|
||||
|
||||
|
||||
def test_extra_scale_columns_are_ignored(N, fused, monkeypatch):
|
||||
dtype = _dtypes()[0]
|
||||
packed, scale, gs = _random_nvfp4(64, 96) # K = 96 leaves a partial last K tile
|
||||
wide = torch.cat([scale, torch.full_like(scale[:, :2], float("nan"))], dim = 1)
|
||||
X = torch.randn(4, 96, device = "cuda", dtype = dtype)
|
||||
before = _calls(N)
|
||||
Y = N._nvfp4_matmul_fwd(X, packed, wide, gs)
|
||||
assert _calls(N)["fwd"] == before["fwd"] + 1
|
||||
assert torch.equal(Y, N._nvfp4_matmul_fwd(X, packed, scale, gs))
|
||||
G = torch.randn(4, 64, device = "cuda", dtype = dtype)
|
||||
assert torch.equal(N._nvfp4_matmul_bwd(G, packed, wide, gs), N._nvfp4_matmul_bwd(G, packed, scale, gs))
|
||||
|
||||
|
||||
def test_inf_scale_in_the_last_group_matches_dense(N, fused, monkeypatch):
|
||||
dtype = _dtypes()[0]
|
||||
packed, scale, gs = _random_nvfp4(64, 104) # the last group is half padding
|
||||
scale = scale.float()
|
||||
scale[:, -1] = float("inf")
|
||||
packed[:, 48:] = 0x22 # E2M1 1.0 in the last group: its weights are +inf, never 0 * inf
|
||||
X = torch.randn(4, 104, device = "cuda", dtype = dtype).abs() + 0.1
|
||||
Y = _fwd(N, X, packed, scale, gs)
|
||||
ref = _dense(N, X, packed, scale, gs, monkeypatch)
|
||||
assert torch.isinf(ref).all() and not torch.isnan(ref).any()
|
||||
assert torch.equal(Y, ref)
|
||||
|
||||
|
||||
def test_zero_global_scale_with_a_partial_row_tile_matches_dense(N, fused, monkeypatch):
|
||||
dtype = _dtypes()[0]
|
||||
packed, scale, gs = _random_nvfp4(72, 128, global_scale = 0.0) # 72 rows: a partial N tile in dX
|
||||
packed[:] = 0x22
|
||||
G = torch.rand(4, 72, device = "cuda", dtype = dtype) + 0.1
|
||||
before = _calls(N)
|
||||
dX = N._nvfp4_matmul_bwd(G, packed, scale, gs)
|
||||
assert _calls(N)["bwd"] == before["bwd"] + 1
|
||||
monkeypatch.setenv("UNSLOTH_NVFP4_FUSED", "0")
|
||||
ref = N._nvfp4_matmul_bwd(G, packed, scale, gs)
|
||||
assert torch.isinf(ref).all() and torch.equal(dX, ref)
|
||||
|
||||
|
||||
def test_column_strided_e4m3_scales(N, fused, monkeypatch):
|
||||
dtype = _dtypes()[0]
|
||||
packed, scale, gs = _random_nvfp4(64, 128)
|
||||
strided = torch.stack([scale, scale], dim = -1).reshape(scale.shape[0], -1)[:, ::2]
|
||||
assert strided.stride(-1) == 2 and torch.equal(strided.view(torch.uint8), scale.view(torch.uint8))
|
||||
X = torch.randn(4, 128, device = "cuda", dtype = dtype)
|
||||
assert torch.equal(_fwd(N, X, packed, strided, gs), _fwd(N, X, packed, scale, gs))
|
||||
|
||||
|
||||
def test_double_backward_takes_the_dense_dx(N, fused):
|
||||
dtype = _dtypes()[0]
|
||||
packed, scale, gs = _random_nvfp4(64, 128)
|
||||
X = torch.randn(4, 128, device = "cuda", dtype = dtype, requires_grad = True)
|
||||
out = N.nvfp4_linear(X, packed, scale, gs)
|
||||
(gX,) = torch.autograd.grad(out.float().square().sum(), X, create_graph = True)
|
||||
assert gX.requires_grad
|
||||
(ggX,) = torch.autograd.grad(gX.float().sum(), X)
|
||||
assert torch.isfinite(ggX).all() and ggX.abs().sum() > 0
|
||||
|
||||
|
||||
def test_row_gate_keeps_unaligned_shapes_dense(N, monkeypatch):
|
||||
monkeypatch.delenv("UNSLOTH_NVFP4_FUSED", raising = False)
|
||||
monkeypatch.setenv("UNSLOTH_NVFP4_FUSED_MAX_ROWS", "64")
|
||||
dtype = _dtypes()[0]
|
||||
for (n, k), want in (((64, 128), True), ((72, 128), False), ((64, 80), False)):
|
||||
packed, scale, gs = _random_nvfp4(n, k)
|
||||
X = torch.randn(4, k, device = "cuda", dtype = dtype)
|
||||
assert N._fused_ok(X, packed, scale, gs, k) == want, (n, k)
|
||||
monkeypatch.setenv("UNSLOTH_NVFP4_FUSED", "1")
|
||||
assert N._fused_ok(X, packed, scale, gs, k)
|
||||
monkeypatch.delenv("UNSLOTH_NVFP4_FUSED")
|
||||
|
||||
def test_flags_are_read_once_per_process(N, live_env, monkeypatch):
|
||||
monkeypatch.setattr(N, "_FUSED_ENV", None)
|
||||
monkeypatch.setenv("UNSLOTH_NVFP4_FUSED", " Off ")
|
||||
monkeypatch.setenv("UNSLOTH_NVFP4_FUSED_MAX_ROWS", "32")
|
||||
assert live_env() == ("off", 32)
|
||||
monkeypatch.setenv("UNSLOTH_NVFP4_FUSED", "1")
|
||||
assert live_env() == ("off", 32)
|
||||
|
||||
def test_offset_width_follows_the_strided_span(N):
|
||||
small = torch.empty((2, 4), device = "meta")
|
||||
assert not N._needs_int64(small)
|
||||
# 8 elements, but row 1 starts 2**31 elements in: int32 offsets would wrap.
|
||||
assert N._needs_int64(small.as_strided((2, 4), (2**31, 1)))
|
||||
|
||||
|
||||
@pytest.mark.skipif(torch.cuda.device_count() < 2, reason = "needs two GPUs")
|
||||
def test_operands_on_another_device_stay_dense(N, monkeypatch):
|
||||
monkeypatch.setenv("UNSLOTH_NVFP4_FUSED_MAX_ROWS", "64")
|
||||
dtype = _dtypes()[0]
|
||||
packed, scale, gs = _random_nvfp4(64, 128)
|
||||
X = torch.randn(4, 128, device = "cuda:1", dtype = dtype)
|
||||
assert N._fused_ok(X.to("cuda:0"), packed, scale, gs, 128)
|
||||
assert not N._fused_ok(X, packed, scale, gs, 128)
|
||||
assert not N._fused_ok(X.to("cuda:0"), packed, scale, gs.to("cuda:1"), 128)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("with_bias", [False, True])
|
||||
def test_autograd_uses_both_fused_kernels(N, fused, monkeypatch, with_bias):
|
||||
dtype = _dtypes()[0]
|
||||
packed, scale, gs = _random_nvfp4(192, 128)
|
||||
X = torch.randn(2, 5, 128, device = "cuda", dtype = dtype, requires_grad = True)
|
||||
dY = torch.randn(2, 5, 192, device = "cuda", dtype = dtype)
|
||||
bias = torch.randn(192, device = "cuda", dtype = dtype) if with_bias else None
|
||||
before = _calls(N)
|
||||
out = N.nvfp4_linear(X, packed, scale, gs, bias)
|
||||
out.backward(dY)
|
||||
assert _calls(N) == {"fwd": before["fwd"] + 1, "bwd": before["bwd"] + 1}
|
||||
W = N.nvfp4_dequantize(packed, scale, gs, dtype)
|
||||
Xr = X.detach().clone().requires_grad_(True)
|
||||
ref = Xr @ W.t() + (bias if with_bias else 0)
|
||||
ref.backward(dY)
|
||||
Wd = W.double()
|
||||
_check(out, ref, Xr.double() @ Wd.t() + (bias.double() if with_bias else 0))
|
||||
_check(X.grad, Xr.grad, dY.double() @ Wd)
|
||||
|
||||
|
||||
def test_no_full_weight_temporary(N, fused, monkeypatch):
|
||||
dtype = _dtypes()[0]
|
||||
rows, cols = 4096, 4096
|
||||
packed, scale, gs = _random_nvfp4(rows, cols, seed = 5)
|
||||
dense_bytes = rows * cols * 2
|
||||
X = torch.randn(16, cols, device = "cuda", dtype = dtype)
|
||||
G = torch.randn(16, rows, device = "cuda", dtype = dtype)
|
||||
|
||||
def peak(fn):
|
||||
torch.cuda.synchronize()
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
base = torch.cuda.memory_allocated()
|
||||
fn()
|
||||
torch.cuda.synchronize()
|
||||
return torch.cuda.max_memory_allocated() - base
|
||||
|
||||
fwd = peak(lambda: N._nvfp4_matmul_fwd(X, packed, scale, gs))
|
||||
bwd = peak(lambda: N._nvfp4_matmul_bwd(G, packed, scale, gs))
|
||||
monkeypatch.setenv("UNSLOTH_NVFP4_FUSED", "0")
|
||||
dense = peak(lambda: N._nvfp4_matmul_fwd(X, packed, scale, gs))
|
||||
assert dense >= dense_bytes
|
||||
assert fwd < dense_bytes // 4 and bwd < dense_bytes // 4, (fwd, bwd, dense)
|
||||
|
||||
|
||||
def test_lora_training_steps_match_dense(N, monkeypatch):
|
||||
dtype = _dtypes()[0]
|
||||
packed, scale, gs = _random_nvfp4(256, 192, global_scale = 37.5, seed = 6)
|
||||
data = [torch.randn(4, 24, 192, device = "cuda", dtype = dtype) for _ in range(4)]
|
||||
|
||||
def train(mode):
|
||||
monkeypatch.setenv("UNSLOTH_NVFP4_FUSED", mode)
|
||||
torch.manual_seed(0)
|
||||
A = torch.nn.Parameter(torch.randn(8, 192, device = "cuda") * 0.02)
|
||||
B = torch.nn.Parameter(torch.randn(256, 8, device = "cuda") * 0.02)
|
||||
opt = torch.optim.AdamW([A, B], lr = 1e-3)
|
||||
before = _calls(N)
|
||||
losses = []
|
||||
for X in data:
|
||||
X = X.clone().requires_grad_(True)
|
||||
h = N.nvfp4_linear(X, packed, scale, gs) + (X @ A.t().to(dtype)) @ B.t().to(dtype) * 2.0
|
||||
loss = h.float().square().mean() + X.float().mean()
|
||||
loss.backward()
|
||||
opt.step()
|
||||
opt.zero_grad()
|
||||
losses.append(loss.item())
|
||||
engaged = _calls(N)["fwd"] - before["fwd"], _calls(N)["bwd"] - before["bwd"]
|
||||
return losses, A.detach().float(), B.detach().float(), engaged
|
||||
|
||||
lf, Af, Bf, ef = train("1")
|
||||
ld, Ad, Bd, ed = train("0")
|
||||
assert ef == (4, 4) and ed == (0, 0)
|
||||
for a, b in zip(lf, ld):
|
||||
assert abs(a - b) <= 2e-3 * abs(b)
|
||||
assert (Af - Ad).norm() / Ad.norm() < 2e-2
|
||||
assert (Bf - Bd).norm() / Bd.norm() < 2e-2
|
||||
|
||||
|
||||
def test_fullgraph_compile_runs_the_fused_kernels(N, fused):
|
||||
import torch._dynamo
|
||||
|
||||
dtype = _dtypes()[0]
|
||||
packed, scale, gs = _random_nvfp4(384, 256)
|
||||
step = lambda X: N.nvfp4_linear(X, packed, scale, gs).float().square().mean()
|
||||
X = torch.randn(64, 256, device = "cuda", dtype = dtype, requires_grad = True)
|
||||
step(X).backward()
|
||||
eager = X.grad.clone()
|
||||
X.grad = None
|
||||
torch._dynamo.reset()
|
||||
assert torch._dynamo.explain(step)(X).graph_break_count == 0
|
||||
torch._dynamo.reset()
|
||||
compiled = torch.compile(step, fullgraph = True)
|
||||
before = _calls(N)
|
||||
for rows in (16, 48):
|
||||
compiled(torch.randn(rows, 256, device = "cuda", dtype = dtype, requires_grad = True)).backward()
|
||||
compiled(X).backward()
|
||||
assert _calls(N)["fwd"] >= before["fwd"] + 3 and _calls(N)["bwd"] >= before["bwd"] + 3
|
||||
torch.testing.assert_close(X.grad, eager)
|
||||
|
|
@ -2,6 +2,9 @@
|
|||
# Copyright 2026-present the Unsloth AI Inc. team.
|
||||
"""W4A16 NVFP4: packed E2M1 weights (compressed-tensors `nvfp4-pack-quantized`) dequantized to 16 bit on the fly."""
|
||||
|
||||
import functools
|
||||
import os
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
|
@ -59,8 +62,7 @@ def _e2m1_to_float(code):
|
|||
|
||||
@triton.jit
|
||||
def _e4m3_to_float(b):
|
||||
# Triton has no float8e4nv before sm_89 (A100, T4), so the group scale arrives as raw bytes: sign, 4 exponent
|
||||
# bits (bias 7), 3 mantissa bits; exponent 0 is subnormal (m * 2^-9) and 0x7F / 0xFF are NaN, as in torch.
|
||||
# No float8e4nv before sm_89: decode E4M3 bytes by hand (bias 7, exp 0 subnormal, 0x7F/0xFF NaN, as torch).
|
||||
exp = (b >> 3) & 15
|
||||
man = b & 7
|
||||
normal = (((exp + 120) << 23) | (man << 20)).to(tl.float32, bitcast = True)
|
||||
|
|
@ -102,7 +104,6 @@ def _nvfp4_dequant_kernel(
|
|||
)
|
||||
scale = _e4m3_to_float(scale.to(tl.int32)) if SCALE_E4M3_BYTES else scale.to(tl.float32)
|
||||
group_scale = tl.math.div_rn(scale, tl.load(global_scale_ptr).to(tl.float32))
|
||||
# Low nibble is the even column: join + reshape interleaves (lo, hi) into contiguous output columns.
|
||||
values = tl.reshape(
|
||||
tl.join(_e2m1_to_float(byte & 15), _e2m1_to_float(byte >> 4)), (BLOCK_R, BLOCK_C // 8, 16)
|
||||
)
|
||||
|
|
@ -142,7 +143,7 @@ def _nvfp4_dequantize_triton(
|
|||
BLOCK_R = BLOCK_R,
|
||||
BLOCK_C = BLOCK_C,
|
||||
SCALE_E4M3_BYTES = e4m3,
|
||||
# int32 offsets wrap once the output passes 2^31 elements; int64 only then, as it costs ~2%.
|
||||
# int32 offsets wrap past 2^31 elements; int64 only then (~2% slower).
|
||||
INT64_OFFSETS = rows * half * 2 >= 2**31,
|
||||
num_warps = 4,
|
||||
enable_fp_fusion = True,
|
||||
|
|
@ -171,8 +172,449 @@ def nvfp4_dequantize(
|
|||
return _nvfp4_dequantize_torch(packed, scale, global_scale, dtype)
|
||||
|
||||
|
||||
# Opaque to torch.compile as whole matmuls: if the dequantized weight were a graph value, AOTAutograd would save it for
|
||||
# backward in every layer (3.6x the NVFP4 storage each), which is exactly the memory W4A16 exists to avoid.
|
||||
@triton.jit
|
||||
def _e2m1_to_float_bits(code):
|
||||
# Magnitude bits as fp16 exp (2) + mantissa (1) = E2M1 * 2^-14 (0.5 a subnormal); upcast and rescale are exact.
|
||||
bits = ((code & 7) << 9) | ((code & 8) << 12)
|
||||
return bits.to(tl.int16).to(tl.float16, bitcast = True).to(tl.float32) * 16384.0
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _nvfp4_weight_tile(
|
||||
packed_ptr,
|
||||
scale_ptr,
|
||||
global_scale,
|
||||
rows,
|
||||
k0,
|
||||
n_rows,
|
||||
half_cols,
|
||||
scale_cols,
|
||||
stride_pr,
|
||||
stride_pc,
|
||||
stride_sr,
|
||||
stride_sc,
|
||||
BLOCK_R: tl.constexpr,
|
||||
BLOCK_K: tl.constexpr,
|
||||
SCALE_E4M3_BYTES: tl.constexpr,
|
||||
OUT_DTYPE: tl.constexpr,
|
||||
INT64_OFFSETS: tl.constexpr,
|
||||
MASK_PAD: tl.constexpr,
|
||||
):
|
||||
# Rounded like dense: E2M1 * (block scale / global scale) in fp32, then one cast; out-of-range decodes to 0.
|
||||
c = k0 // 2 + tl.arange(0, BLOCK_K // 2)
|
||||
g = k0 // 16 + tl.arange(0, BLOCK_K // 16)
|
||||
if INT64_OFFSETS:
|
||||
c = c.to(tl.int64)
|
||||
g = g.to(tl.int64)
|
||||
row_ok = rows[:, None] < n_rows
|
||||
byte = tl.load(
|
||||
packed_ptr + rows[:, None] * stride_pr + c[None, :] * stride_pc,
|
||||
mask = row_ok & (c[None, :] < half_cols),
|
||||
other = 0,
|
||||
).to(tl.int32)
|
||||
scale = tl.load(
|
||||
scale_ptr + rows[:, None] * stride_sr + g[None, :] * stride_sc,
|
||||
mask = row_ok & (g[None, :] < scale_cols),
|
||||
other = 0,
|
||||
)
|
||||
scale = _e4m3_to_float(scale.to(tl.int32)) if SCALE_E4M3_BYTES else scale.to(tl.float32)
|
||||
group_scale = tl.math.div_rn(scale, global_scale)
|
||||
values = tl.reshape(
|
||||
tl.join(_e2m1_to_float_bits(byte & 15), _e2m1_to_float_bits(byte >> 4)),
|
||||
(BLOCK_R, BLOCK_K // 16, 16),
|
||||
)
|
||||
w = tl.reshape(values * group_scale[:, :, None], (BLOCK_R, BLOCK_K))
|
||||
if MASK_PAD: # a partial tile: padding stays 0 even when scaling gives 0 * inf or 0 / 0
|
||||
w = tl.where(row_ok & ((k0 + tl.arange(0, BLOCK_K))[None, :] < half_cols * 2), w, 0.0)
|
||||
return w.to(OUT_DTYPE)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _nvfp4_matmul_fwd_kernel(
|
||||
x_ptr,
|
||||
packed_ptr,
|
||||
scale_ptr,
|
||||
global_scale_ptr,
|
||||
y_ptr,
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
half_cols,
|
||||
scale_cols,
|
||||
stride_xm,
|
||||
stride_xk,
|
||||
stride_pr,
|
||||
stride_pc,
|
||||
stride_sr,
|
||||
stride_sc,
|
||||
stride_ym,
|
||||
stride_yn,
|
||||
stride_ys,
|
||||
k_per_split,
|
||||
BLOCK_M: tl.constexpr,
|
||||
BLOCK_N: tl.constexpr,
|
||||
BLOCK_K: tl.constexpr,
|
||||
GROUP_M: tl.constexpr,
|
||||
SCALE_E4M3_BYTES: tl.constexpr,
|
||||
INT64_OFFSETS: tl.constexpr,
|
||||
MASK_PAD: tl.constexpr,
|
||||
):
|
||||
# Split K writes fp32 partials [split, M, N], summed in a fixed order afterwards (deterministic).
|
||||
pid = tl.program_id(0)
|
||||
split_id = tl.program_id(1)
|
||||
if INT64_OFFSETS:
|
||||
split_id = split_id.to(tl.int64)
|
||||
k_start = tl.program_id(1) * k_per_split
|
||||
k_end = min(k_start + k_per_split, K)
|
||||
num_m = tl.cdiv(M, BLOCK_M)
|
||||
num_n = tl.cdiv(N, BLOCK_N)
|
||||
group = GROUP_M * num_n
|
||||
first_m = (pid // group) * GROUP_M
|
||||
size_m = min(num_m - first_m, GROUP_M)
|
||||
pid_m = first_m + (pid % group) % size_m
|
||||
pid_n = (pid % group) // size_m
|
||||
rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
|
||||
rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
|
||||
if INT64_OFFSETS:
|
||||
rm = rm.to(tl.int64)
|
||||
rn = rn.to(tl.int64)
|
||||
global_scale = tl.load(global_scale_ptr).to(tl.float32)
|
||||
acc = tl.zeros((BLOCK_M, BLOCK_N), dtype = tl.float32)
|
||||
for k0 in range(k_start, k_end, BLOCK_K):
|
||||
rk = k0 + tl.arange(0, BLOCK_K)
|
||||
if INT64_OFFSETS:
|
||||
rk = rk.to(tl.int64)
|
||||
x = tl.load(
|
||||
x_ptr + rm[:, None] * stride_xm + rk[None, :] * stride_xk,
|
||||
mask = (rm[:, None] < M) & (rk[None, :] < k_end),
|
||||
other = 0.0,
|
||||
)
|
||||
w = _nvfp4_weight_tile(
|
||||
packed_ptr,
|
||||
scale_ptr,
|
||||
global_scale,
|
||||
rn,
|
||||
k0,
|
||||
N,
|
||||
half_cols,
|
||||
scale_cols,
|
||||
stride_pr,
|
||||
stride_pc,
|
||||
stride_sr,
|
||||
stride_sc,
|
||||
BLOCK_N,
|
||||
BLOCK_K,
|
||||
SCALE_E4M3_BYTES,
|
||||
x_ptr.dtype.element_ty,
|
||||
INT64_OFFSETS,
|
||||
MASK_PAD,
|
||||
)
|
||||
acc = tl.dot(x, tl.trans(w), acc)
|
||||
tl.store(
|
||||
y_ptr + split_id * stride_ys + rm[:, None] * stride_ym + rn[None, :] * stride_yn,
|
||||
acc.to(y_ptr.dtype.element_ty),
|
||||
mask = (rm[:, None] < M) & (rn[None, :] < N),
|
||||
)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _nvfp4_matmul_dx_kernel(
|
||||
g_ptr,
|
||||
packed_ptr,
|
||||
scale_ptr,
|
||||
global_scale_ptr,
|
||||
dx_ptr,
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
half_cols,
|
||||
scale_cols,
|
||||
stride_gm,
|
||||
stride_gn,
|
||||
stride_pr,
|
||||
stride_pc,
|
||||
stride_sr,
|
||||
stride_sc,
|
||||
stride_dm,
|
||||
stride_dk,
|
||||
stride_ds,
|
||||
n_per_split,
|
||||
BLOCK_M: tl.constexpr,
|
||||
BLOCK_N: tl.constexpr,
|
||||
BLOCK_K: tl.constexpr,
|
||||
GROUP_M: tl.constexpr,
|
||||
SCALE_E4M3_BYTES: tl.constexpr,
|
||||
INT64_OFFSETS: tl.constexpr,
|
||||
MASK_PAD: tl.constexpr,
|
||||
):
|
||||
pid = tl.program_id(0)
|
||||
split_id = tl.program_id(1)
|
||||
if INT64_OFFSETS:
|
||||
split_id = split_id.to(tl.int64)
|
||||
n_start = tl.program_id(1) * n_per_split
|
||||
n_end = min(n_start + n_per_split, N)
|
||||
num_m = tl.cdiv(M, BLOCK_M)
|
||||
num_k = tl.cdiv(K, BLOCK_K)
|
||||
group = GROUP_M * num_k
|
||||
first_m = (pid // group) * GROUP_M
|
||||
size_m = min(num_m - first_m, GROUP_M)
|
||||
pid_m = first_m + (pid % group) % size_m
|
||||
pid_k = (pid % group) // size_m
|
||||
rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
|
||||
rk = pid_k * BLOCK_K + tl.arange(0, BLOCK_K)
|
||||
if INT64_OFFSETS:
|
||||
rm = rm.to(tl.int64)
|
||||
rk = rk.to(tl.int64)
|
||||
global_scale = tl.load(global_scale_ptr).to(tl.float32)
|
||||
acc = tl.zeros((BLOCK_M, BLOCK_K), dtype = tl.float32)
|
||||
for n0 in range(n_start, n_end, BLOCK_N):
|
||||
rn = n0 + tl.arange(0, BLOCK_N)
|
||||
if INT64_OFFSETS:
|
||||
rn = rn.to(tl.int64)
|
||||
g = tl.load(
|
||||
g_ptr + rm[:, None] * stride_gm + rn[None, :] * stride_gn,
|
||||
mask = (rm[:, None] < M) & (rn[None, :] < n_end),
|
||||
other = 0.0,
|
||||
)
|
||||
w = _nvfp4_weight_tile(
|
||||
packed_ptr,
|
||||
scale_ptr,
|
||||
global_scale,
|
||||
rn,
|
||||
pid_k * BLOCK_K,
|
||||
n_end,
|
||||
half_cols,
|
||||
scale_cols,
|
||||
stride_pr,
|
||||
stride_pc,
|
||||
stride_sr,
|
||||
stride_sc,
|
||||
BLOCK_N,
|
||||
BLOCK_K,
|
||||
SCALE_E4M3_BYTES,
|
||||
g_ptr.dtype.element_ty,
|
||||
INT64_OFFSETS,
|
||||
MASK_PAD,
|
||||
)
|
||||
acc = tl.dot(g, w, acc)
|
||||
tl.store(
|
||||
dx_ptr + split_id * stride_ds + rm[:, None] * stride_dm + rk[None, :] * stride_dk,
|
||||
acc.to(dx_ptr.dtype.element_ty),
|
||||
mask = (rm[:, None] < M) & (rk[None, :] < K),
|
||||
)
|
||||
|
||||
|
||||
_FUSED_CALLS = {"fwd": 0, "bwd": 0}
|
||||
|
||||
|
||||
def _read_fused_env():
|
||||
# UNSLOTH_NVFP4_FUSED: 0 dense, 1 fused, unset = per-GPU row gate (UNSLOTH_NVFP4_FUSED_MAX_ROWS overrides it).
|
||||
rows = os.environ.get("UNSLOTH_NVFP4_FUSED_MAX_ROWS")
|
||||
return os.environ.get("UNSLOTH_NVFP4_FUSED", "auto").strip().lower(), int(rows) if rows else None
|
||||
|
||||
|
||||
_FUSED_ENV = None
|
||||
|
||||
|
||||
def _fused_env():
|
||||
# Read once: os.environ lookups would add about 1 us to every NVFP4 linear.
|
||||
global _FUSED_ENV
|
||||
if _FUSED_ENV is None:
|
||||
_FUSED_ENV = _read_fused_env()
|
||||
return _FUSED_ENV
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize = None)
|
||||
def _device_capability(index):
|
||||
return torch.cuda.get_device_capability(index)
|
||||
|
||||
|
||||
_FUSED_OFF = ("0", "false", "off", "no")
|
||||
_FUSED_FORCE = ("1", "true", "on", "yes", "force")
|
||||
_IS_HIP = torch.version.hip is not None
|
||||
|
||||
|
||||
def _fused_ok(A, packed, scale, global_scale, inner):
|
||||
# Anything unusual keeps the dense path (and its errors).
|
||||
mode, rows = _fused_env()
|
||||
if mode in _FUSED_OFF or _IS_HIP or not A.is_cuda:
|
||||
return False
|
||||
index = A.get_device()
|
||||
if mode not in _FUSED_FORCE:
|
||||
limit = rows if rows is not None else _FUSED_MAX_ROWS.get(_device_capability(index), 0)
|
||||
if A.numel() > limit * inner: # unmeasured GPUs (limit 0) leave here
|
||||
return False
|
||||
if A.dtype not in (torch.float16, torch.bfloat16) or A.numel() == 0 or A.shape[-1] != inner:
|
||||
return False
|
||||
if A.dtype == torch.bfloat16 and _device_capability(index) < (8, 0):
|
||||
return False
|
||||
return (
|
||||
packed.device == A.device
|
||||
and scale.device == A.device
|
||||
and global_scale.device == A.device
|
||||
and global_scale.numel() >= 1
|
||||
and packed.dtype == torch.uint8
|
||||
and packed.dim() == 2
|
||||
and packed.shape[0] > 0
|
||||
and (mode in _FUSED_FORCE or (inner % 64 == 0 and packed.shape[0] % 64 == 0))
|
||||
and not _is_transposed_view(packed)
|
||||
and scale.dtype in (torch.float8_e4m3fn, torch.float32, torch.float16, torch.bfloat16)
|
||||
and scale.dim() == 2
|
||||
and scale.shape[0] == packed.shape[0]
|
||||
and packed.shape[1] * 2 <= scale.shape[1] * NVFP4_GROUP_SIZE
|
||||
)
|
||||
|
||||
|
||||
# Max rows where fused beats dense (measured at N, K multiples of 64); unmeasured GPUs stay dense.
|
||||
_FUSED_MAX_ROWS = {(8, 0): 64, (8, 9): 256, (12, 0): 128}
|
||||
|
||||
|
||||
def _scale_args(scale):
|
||||
e4m3 = scale.dtype == torch.float8_e4m3fn
|
||||
return (scale.view(torch.uint8) if e4m3 else scale), e4m3
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize = None)
|
||||
def _sm_count(index):
|
||||
return torch.cuda.get_device_properties(index).multi_processor_count
|
||||
|
||||
|
||||
def _fused_config(
|
||||
M,
|
||||
n_out,
|
||||
k_red,
|
||||
device,
|
||||
dx = False,
|
||||
):
|
||||
"""(BLOCK_M, BLOCK_OUT, BLOCK_RED, num_warps, num_stages, split) for an [M, k_red] @ [k_red, n_out] product."""
|
||||
index = device.index or 0
|
||||
stages = 2 if _device_capability(index) < (8, 0) else 3
|
||||
if M <= 16:
|
||||
bm, bo, br, warps = (16, 128, 64, 4) if dx else (16, 64, 128, 4)
|
||||
elif M <= 64:
|
||||
bm, bo, br, warps = (32, 128, 64, 4) if dx else (32, 64, 128, 8)
|
||||
elif M <= 256:
|
||||
bm, bo, br, warps = (64, 128, 64, 4) if dx else (128, 64, 128, 8)
|
||||
else:
|
||||
bm, bo, br, warps = 128, 64, 64, 4
|
||||
tiles = triton.cdiv(M, bm) * triton.cdiv(n_out, bo)
|
||||
# Split under four waves; fp32 partials capped at 1/8 of the 16 bit weight.
|
||||
split = 1
|
||||
while (
|
||||
tiles * split < 4 * _sm_count(index)
|
||||
and k_red >= split * 2 * 4 * br
|
||||
and 16 * split * 2 * M <= k_red
|
||||
and split < 16
|
||||
):
|
||||
split *= 2
|
||||
return bm, bo, br, warps, stages, split
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _split_sum_kernel(partial_ptr, out_ptr, numel, SPLIT: tl.constexpr, BLOCK: tl.constexpr):
|
||||
# Fixed-order sum of the split partials (determinism), one cast.
|
||||
offs = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
|
||||
mask = offs < numel
|
||||
acc = tl.load(partial_ptr + offs, mask = mask, other = 0.0)
|
||||
for i in tl.static_range(1, SPLIT):
|
||||
acc += tl.load(partial_ptr + i * numel + offs, mask = mask, other = 0.0)
|
||||
tl.store(out_ptr + offs, acc.to(out_ptr.dtype.element_ty), mask = mask)
|
||||
|
||||
|
||||
def _split_out(M, n_out, split, dtype, device):
|
||||
out = torch.empty((M, n_out), dtype = dtype, device = device)
|
||||
if split == 1:
|
||||
return out, out, 0
|
||||
partial = torch.empty((split, M, n_out), dtype = torch.float32, device = device)
|
||||
return out, partial, partial.stride(0)
|
||||
|
||||
|
||||
def _split_sum(partial, out):
|
||||
if partial is out:
|
||||
return out
|
||||
numel = out.numel()
|
||||
_split_sum_kernel[(triton.cdiv(numel, 1024),)](
|
||||
partial, out, numel, SPLIT = partial.shape[0], BLOCK = 1024, enable_fp_fusion = True
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
def _needs_int64(*tensors):
|
||||
# Largest reachable element offset, not numel: a strided view can reach past 2**31 with fewer elements.
|
||||
for t in tensors:
|
||||
if t.is_contiguous():
|
||||
if t.numel() > 2**31:
|
||||
return True
|
||||
elif sum((n - 1) * abs(st) for n, st in zip(t.shape, t.stride())) >= 2**31:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _scale_groups(packed, scale):
|
||||
# Scales past the logical width are ignored, as dense slices them off.
|
||||
return min(scale.shape[1], -(-packed.shape[1] // 8))
|
||||
|
||||
|
||||
def _per_split(k_red, split, block):
|
||||
return triton.cdiv(triton.cdiv(k_red, split), block) * block
|
||||
|
||||
|
||||
def _nvfp4_matmul_fwd_fused(X, packed, scale, global_scale):
|
||||
K = X.shape[-1]
|
||||
N = packed.shape[0]
|
||||
X2 = X.reshape(-1, K)
|
||||
M = X2.shape[0]
|
||||
scale, e4m3 = _scale_args(scale)
|
||||
BM, BN, BK, warps, stages, split = _CONFIG_HOOK(M, N, K, X.device, False)
|
||||
Y, buf, stride_s = _split_out(M, N, split, X.dtype, X.device)
|
||||
int64 = _needs_int64(X2, buf, packed, scale)
|
||||
grid = (triton.cdiv(M, BM) * triton.cdiv(N, BN), split)
|
||||
with _fp8_triton_device_context(packed):
|
||||
_nvfp4_matmul_fwd_kernel[grid](
|
||||
X2, packed, scale, global_scale, buf,
|
||||
M, N, K, packed.shape[1], _scale_groups(packed, scale),
|
||||
X2.stride(0), X2.stride(1), packed.stride(0), packed.stride(1),
|
||||
scale.stride(0), scale.stride(1), buf.stride(-2), buf.stride(-1), stride_s,
|
||||
_per_split(K, split, BK),
|
||||
BLOCK_M = BM, BLOCK_N = BN, BLOCK_K = BK, GROUP_M = 8,
|
||||
SCALE_E4M3_BYTES = e4m3, INT64_OFFSETS = int64, MASK_PAD = N % BN != 0 or K % BK != 0,
|
||||
num_warps = warps, num_stages = stages, enable_fp_fusion = True,
|
||||
) # fmt: skip
|
||||
Y = _split_sum(buf, Y)
|
||||
_FUSED_CALLS["fwd"] += 1
|
||||
return Y.view(*X.shape[:-1], N)
|
||||
|
||||
|
||||
def _nvfp4_matmul_dx_fused(grad, packed, scale, global_scale):
|
||||
N = grad.shape[-1]
|
||||
K = packed.shape[1] * 2
|
||||
G = grad.reshape(-1, N)
|
||||
M = G.shape[0]
|
||||
scale, e4m3 = _scale_args(scale)
|
||||
BM, BK, BN, warps, stages, split = _CONFIG_HOOK(M, K, N, grad.device, True)
|
||||
dX, buf, stride_s = _split_out(M, K, split, grad.dtype, grad.device)
|
||||
int64 = _needs_int64(G, buf, packed, scale)
|
||||
grid = (triton.cdiv(M, BM) * triton.cdiv(K, BK), split)
|
||||
with _fp8_triton_device_context(packed):
|
||||
_nvfp4_matmul_dx_kernel[grid](
|
||||
G, packed, scale, global_scale, buf,
|
||||
M, N, K, packed.shape[1], _scale_groups(packed, scale),
|
||||
G.stride(0), G.stride(1), packed.stride(0), packed.stride(1),
|
||||
scale.stride(0), scale.stride(1), buf.stride(-2), buf.stride(-1), stride_s,
|
||||
_per_split(N, split, BN),
|
||||
BLOCK_M = BM, BLOCK_N = BN, BLOCK_K = BK, GROUP_M = 8,
|
||||
SCALE_E4M3_BYTES = e4m3, INT64_OFFSETS = int64, MASK_PAD = N % BN != 0 or K % BK != 0,
|
||||
num_warps = warps, num_stages = stages, enable_fp_fusion = True,
|
||||
) # fmt: skip
|
||||
dX = _split_sum(buf, dX)
|
||||
_FUSED_CALLS["bwd"] += 1
|
||||
return dX.view(*grad.shape[:-1], K)
|
||||
|
||||
|
||||
_CONFIG_HOOK = _fused_config
|
||||
|
||||
|
||||
# Opaque to torch.compile, else AOTAutograd saves every layer's dequantized weight for backward.
|
||||
@_opaque_under_compile(
|
||||
"nvfp4_matmul_fwd",
|
||||
lambda X, packed, scale, global_scale: X.new_empty(X.shape[:-1] + (packed.shape[0],)),
|
||||
|
|
@ -180,6 +622,8 @@ def nvfp4_dequantize(
|
|||
def _nvfp4_matmul_fwd(
|
||||
X: torch.Tensor, packed: torch.Tensor, scale: torch.Tensor, global_scale: torch.Tensor
|
||||
) -> torch.Tensor:
|
||||
if _fused_ok(X, packed, scale, global_scale, packed.shape[1] * 2):
|
||||
return _nvfp4_matmul_fwd_fused(X, packed, scale, global_scale.reshape(-1)[:1])
|
||||
return torch.matmul(X, nvfp4_dequantize(packed, scale, global_scale, X.dtype).t())
|
||||
|
||||
|
||||
|
|
@ -192,6 +636,11 @@ def _nvfp4_matmul_fwd(
|
|||
def _nvfp4_matmul_bwd(
|
||||
grad: torch.Tensor, packed: torch.Tensor, scale: torch.Tensor, global_scale: torch.Tensor
|
||||
) -> torch.Tensor:
|
||||
# create_graph: the raw launch has no autograd history, so differentiate through the dense matmul.
|
||||
if not (grad.requires_grad and torch.is_grad_enabled()) and _fused_ok(
|
||||
grad, packed, scale, global_scale, packed.shape[0]
|
||||
):
|
||||
return _nvfp4_matmul_dx_fused(grad, packed, scale, global_scale.reshape(-1)[:1])
|
||||
return torch.matmul(grad, nvfp4_dequantize(packed, scale, global_scale, grad.dtype))
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue