perf: fast interior loop on the big cta and fused epilogue bias
- re-enable kFastLoop on the 128x128 CTA for congruous layouts: the base-pair fragment addressing freed the registers the old offset tables spilled, and the predication-free interior loop now wins across the band (fast body 142 SASS instr with zero predicated fallback vs 719/136 generic; 128 regs, no spill) - move the big/small CTA dispatch boundary from 3/4 to 5/8 wave: with the fast big-CTA loop the crossover sits between 49 and 63 tiles (63-tile rect +8%, 1024^3 now takes the big CTA) - fuse the linear bias into the GEMM epilogue: FP8Params.bias_ptr adds in fp32 before the single bf16 rounding, replacing the separate out + bias elementwise pass; guarded loads keep N tails exact and batch broadcast falls out of the row-major layout - resolve Python None bias in the pybind layer (py::object + cast) so ops/fp8.py and fp8.py pass the argument through untouched; drop the _empty_bias sentinel machinery - add fused-bias tests covering odd N tails, no-bias parity and batched broadcast Benchmark: L20 (sm_89), CUDA-graph e2e. Big-CTA fast loop + dispatch: 1024^3 102.6->106.3T, 1152^3 128.5->133.3T, 2048^3 173.8->178.2T, 3072^3 180.2->185.3T, 8192^3 196.2->197.7T. Bias fusion (with-bias GEMM vs unfused out + bias): 1024^3 90.5->106.1T (+17%), 2048^3 162.2->178.3T (+10%), 4096^3 178.2->191.1T (+7%). Fused bias differs from the split path by <=1 bf16 ulp and is closer to the fp64 reference. 596 tests pass.
This commit is contained in:
@@ -116,6 +116,40 @@ def test_mm_fp8_transposed_operands(trans_a, trans_b):
|
||||
torch.testing.assert_close(out, expected, atol=0.125, rtol=0.01)
|
||||
|
||||
|
||||
@skip_no_fp8
|
||||
@pytest.mark.parametrize("bias_on", [False, True])
|
||||
def test_mm_fp8_fused_bias(bias_on):
|
||||
"""Epilogue-fused bias matches the unfused out + bias reference (single
|
||||
fp32 rounding vs the reference's double rounding keeps it within 1 ulp),
|
||||
including N-tail columns and batched broadcast."""
|
||||
torch.manual_seed(31)
|
||||
m, n, k = 19, 13, 37 # odd n exercises the guarded bias loads
|
||||
a = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
|
||||
b = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
|
||||
sa, sb = _scale(a), _scale(b)
|
||||
a8, _ = quantize(a, sa.reciprocal(), "e4m3")
|
||||
b8, _ = quantize(b, sb.reciprocal(), "e4m3")
|
||||
bias = torch.randn(n, device="cuda", dtype=torch.bfloat16)
|
||||
|
||||
out = mm_fp8(a8, b8, sa * sb, trans_b=True, bias=bias if bias_on else None)
|
||||
base = (_quantize(a, sa) @ _quantize(b, sb).t() * sa * sb).to(torch.bfloat16)
|
||||
expected = base + bias if bias_on else base
|
||||
# bias is O(1) against O(sqrt(k)) accumulators: absolute tolerance rules
|
||||
torch.testing.assert_close(out, expected, atol=0.13, rtol=0.01)
|
||||
|
||||
# Batched broadcast: bias applies to every batch slice (each slice gets
|
||||
# its own reference from its own operand values).
|
||||
ab = torch.randn(3, m, k, device="cuda", dtype=torch.bfloat16)
|
||||
ab8, _ = quantize(ab, sa.reciprocal(), "e4m3")
|
||||
outb = mm_fp8(ab8, b8, sa * sb, trans_b=True, bias=bias)
|
||||
assert outb.shape == (3, m, n)
|
||||
for i in range(3):
|
||||
expected_b = (_quantize(ab[i], sa) @ _quantize(b, sb).t() * sa * sb).to(
|
||||
torch.bfloat16
|
||||
) + bias
|
||||
torch.testing.assert_close(outb[i], expected_b, atol=0.13, rtol=0.01)
|
||||
|
||||
|
||||
@skip_no_fp8
|
||||
@pytest.mark.parametrize("trans_a", [False, True])
|
||||
@pytest.mark.parametrize("trans_b", [False, True])
|
||||
|
||||
Reference in New Issue
Block a user