CharlesCNorton
fp4-train: add fp4_mm, the native Blackwell NVFP4 tensor-core matmul (10x over the dequant path)
3be46d1 | import pytest | |
| import torch | |
| import kernels | |
| fp4 = kernels.get_kernel("phanerozoic/fp4-train", version=1, trust_remote_code=True) | |
| requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") | |
| E2M1 = None | |
| def _ref_qdq(x): | |
| grid = torch.tensor([0., .5, 1., 1.5, 2., 3., 4., 6.], device=x.device) | |
| M, K = x.shape | |
| amax = x.reshape(M, K // 16, 16).abs().amax(-1).clamp(min=1e-12) | |
| s = (amax / 6).to(torch.float8_e4m3fn).float().repeat_interleave(16, 1) | |
| xs = x / s | |
| mag = xs.abs().clamp(max=6) | |
| idx = torch.bucketize(mag, grid, right=False) | |
| hi = grid[idx.clamp(max=7)]; lo = grid[(idx - 1).clamp(min=0)] | |
| return (torch.where((mag - lo) < (hi - mag), lo, hi) * xs.sign()) * s | |
| def test_encoding_matches_reference(): | |
| x = torch.randn(256, 512, device="cuda") | |
| p, s = fp4.quantize(x) | |
| d = fp4.dequantize(p, s, *x.shape).float() | |
| assert (d - _ref_qdq(x)).abs().max().item() < 1e-3 | |
| def test_stochastic_rounding_unbiased(): | |
| xb = torch.empty(4096, 16, device="cuda"); xb[:, 0] = 6.0; xb[:, 1:] = 2.3 | |
| acc = torch.zeros_like(xb) | |
| for stp in range(48): | |
| p, s = fp4.quantize(xb, stochastic=True, step=stp) | |
| acc += fp4.dequantize(p, s, *xb.shape).float() | |
| assert abs((acc / 48)[:, 1:].mean().item() - 2.3) < 0.03 | |
| def test_error_feedback_cancels_drift(): | |
| xb = torch.empty(4096, 16, device="cuda"); xb[:, 0] = 6.0; xb[:, 1:] = 2.3 | |
| res = torch.zeros_like(xb, dtype=torch.float32); acc = torch.zeros_like(xb) | |
| for stp in range(48): | |
| p, s = fp4.quantize(xb, stochastic=False, residual=res, step=stp) | |
| acc += fp4.dequantize(p, s, *xb.shape).float() | |
| assert abs((acc / 48)[:, 1:].mean().item() - 2.3) < 0.03 | |
| def test_fp4_linear_trains(): | |
| torch.manual_seed(0); d = 256 | |
| W = [torch.randn(d, d, device="cuda", requires_grad=True) for _ in range(2)] | |
| X = torch.randn(512, d, device="cuda") | |
| Y = torch.tanh(X @ (torch.randn(d, d, device="cuda") / d**0.5)) | |
| opt = torch.optim.Adam(W, lr=3e-3) | |
| losses = [] | |
| for stp in range(150): | |
| h = X | |
| for i, w in enumerate(W): | |
| h = fp4.fp4_linear(h, w, stochastic=True, step=stp * 2 + i) | |
| if i < len(W) - 1: | |
| h = torch.relu(h) | |
| loss = (h - Y).pow(2).mean() | |
| opt.zero_grad(); loss.backward(); opt.step(); losses.append(loss.item()) | |
| assert losses[-1] < losses[0] * 0.5 | |
| def _kqdq(x, sr, res, step): | |
| p, s = fp4.quantize(x.contiguous(), stochastic=sr, residual=res, seed=0, step=step) | |
| return fp4.dequantize(p, s, *x.shape).float() | |
| def test_weight_storage_ef_beats_rtn(): | |
| # weights stored in NVFP4 and updated in place; sub-ULP updates. Round-to-nearest | |
| # stalls; error feedback recovers near-bf16 loss. | |
| torch.manual_seed(2); d, n = 128, 2048 | |
| Wt = [torch.randn(d, d, device="cuda") / d**0.5 for _ in range(2)] | |
| X = torch.randn(n, d, device="cuda"); Y = torch.relu(X @ Wt[0].t()) @ Wt[1].t() | |
| W0 = [torch.randn(d, d, device="cuda") / d**0.5 for _ in range(2)] | |
| def run(mode, steps=1000, lr=1e-3): | |
| W = [w.clone() for w in W0] | |
| res = [torch.zeros_like(w) for w in W] if mode == "ef" else [None, None] | |
| if mode != "bf16": | |
| W = [_kqdq(w, mode == "sr", r, 0) for w, r in zip(W, res)] | |
| for t in range(steps): | |
| h = torch.relu(X @ W[0].t()); e = h @ W[1].t() - Y | |
| g1 = 2.0 / n * (e.t() @ h); gh = (e @ W[1]) * (X @ W[0].t() > 0).float() | |
| g0 = 2.0 / n * (gh.t() @ X) | |
| for i, g in enumerate((g0, g1)): | |
| tgt = W[i] - lr * g | |
| W[i] = tgt if mode == "bf16" else _kqdq(tgt, mode == "sr", res[i], t + 1) | |
| return (torch.relu(X @ W[0].t()) @ W[1].t() - Y).pow(2).mean().item() | |
| bf16, rtn, ef = run("bf16"), run("rtn"), run("ef") | |
| assert ef < rtn * 0.5 # error feedback decisively beats round-to-nearest | |
| assert ef < bf16 * 1.3 # and recovers near-bf16 loss | |
| def test_fp4_mm_hardware_matches_fp4_matmul(): | |
| # native NVFP4 tensor-core GEMM is bit-close to the dequant path (both are the | |
| # FP4-operand product); requires a Blackwell FP4-capable GPU. | |
| if torch.cuda.get_device_capability(0)[0] < 10: | |
| pytest.skip("native FP4 _scaled_mm needs Blackwell (cc >= 10)") | |
| a = torch.randn(256, 1024, device="cuda") | |
| b = torch.randn(512, 1024, device="cuda") | |
| hw = fp4.fp4_mm(a, b).float() | |
| ref = fp4.fp4_matmul(a, b).float() | |
| assert (hw - ref).abs().max().item() / (ref.abs().max().item() + 1e-9) < 0.05 | |