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:
+33
-3
@@ -164,7 +164,7 @@ std::tuple<torch::Tensor, torch::Tensor> quantize(torch::Tensor x,
|
||||
}
|
||||
|
||||
torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
|
||||
int64_t trans_a, int64_t trans_b) {
|
||||
int64_t trans_a, int64_t trans_b, torch::Tensor bias) {
|
||||
TORCH_CHECK(a.is_cuda() && b.is_cuda(), "CUDA tensors required");
|
||||
TORCH_CHECK(a.scalar_type() == torch::kFloat8_e4m3fn ||
|
||||
a.scalar_type() == torch::kFloat8_e5m2,
|
||||
@@ -208,6 +208,16 @@ torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
|
||||
FP8Params p;
|
||||
pack_gemm(p, a_st.data_ptr(), b_st.data_ptr(), output.data_ptr(), scale,
|
||||
m, n, k, a_ld, b_ld);
|
||||
// Fused epilogue bias (bf16, broadcast over rows and batches). An
|
||||
// undefined or 0-element tensor keeps the plain scaled output.
|
||||
if (bias.defined() && bias.numel() > 0) {
|
||||
TORCH_CHECK(bias.is_cuda() && bias.scalar_type() == torch::kBFloat16,
|
||||
"fp8 gemm bias must be a CUDA bf16 tensor");
|
||||
TORCH_CHECK(bias.dim() == 1 && bias.size(0) == n,
|
||||
"fp8 gemm bias must be 1D of length n=", n);
|
||||
TORCH_CHECK(bias.is_contiguous(), "fp8 gemm bias must be contiguous");
|
||||
p.bias_ptr = bias.data_ptr();
|
||||
}
|
||||
p.batch = static_cast<int>(batch);
|
||||
p.a_batch_stride = (batch_a == 1 && batch > 1) ? 0 : a_bstride;
|
||||
p.b_batch_stride = (batch_b == 1 && batch > 1) ? 0 : b_bstride;
|
||||
@@ -220,9 +230,29 @@ torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
|
||||
return output;
|
||||
}
|
||||
|
||||
// mm_fp8 binding: Python None and an omitted argument both mean "no bias"
|
||||
// (resolved to an undefined tensor here, so every Python layer can pass its
|
||||
// bias argument through untouched instead of normalizing it host-side).
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("quantize", &quantize, py::arg("x"), py::arg("scale"),
|
||||
py::arg("fmt"));
|
||||
m.def("mm_fp8", &mm_fp8, py::arg("a"), py::arg("b"), py::arg("scale"),
|
||||
py::arg("trans_a") = 0, py::arg("trans_b") = 0);
|
||||
m.def(
|
||||
"mm_fp8",
|
||||
[](torch::Tensor a, torch::Tensor b, torch::Tensor scale,
|
||||
int64_t trans_a, int64_t trans_b, py::object bias) {
|
||||
torch::Tensor t;
|
||||
if (!bias.is_none()) {
|
||||
// (py::isinstance<torch::Tensor> is false for real tensors
|
||||
// here — torch's caster registers no pybind type info — so
|
||||
// validate by attempting the cast itself.)
|
||||
try {
|
||||
t = bias.cast<torch::Tensor>();
|
||||
} catch (const py::cast_error&) {
|
||||
TORCH_CHECK(false, "bias must be a torch.Tensor or None");
|
||||
}
|
||||
}
|
||||
return mm_fp8(a, b, scale, trans_a, trans_b, t);
|
||||
},
|
||||
py::arg("a"), py::arg("b"), py::arg("scale"), py::arg("trans_a") = 0,
|
||||
py::arg("trans_b") = 0, py::arg("bias") = py::none());
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user