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:
@@ -103,8 +103,12 @@ struct FP8QuantizeParams {
|
||||
struct FP8Params {
|
||||
// Inputs: a/b are FP8 for the pre-quantized path. Scales are
|
||||
// quantization steps (device scalars).
|
||||
// Optional bf16 bias broadcast over output rows (fused into the epilogue
|
||||
// before the bf16 rounding, so it adds in fp32 — one rounding fewer than
|
||||
// the separate out + bias elementwise kernel it replaces). Null disables.
|
||||
const void* __restrict__ a_ptr = nullptr;
|
||||
const void* __restrict__ b_ptr = nullptr;
|
||||
const void* __restrict__ bias_ptr = nullptr;
|
||||
void* __restrict__ out_ptr = nullptr;
|
||||
|
||||
const float* __restrict__ scale = nullptr;
|
||||
|
||||
@@ -613,6 +613,13 @@ __global__ void __launch_bounds__(Traits::kCtaThreads,
|
||||
// scatter and the gather conflict-free: a lane quad's chunk and the 8
|
||||
// rows of one gather phase map to distinct 4-bank groups.
|
||||
const float output_scale = scale;
|
||||
// Fused bias (idea B): added to the fp32 accumulator before the single
|
||||
// bf16 rounding — one fewer rounding than the out + bias elementwise
|
||||
// pass this replaces, and no extra kernel launch / m*n round-trip. The
|
||||
// per-lane loads (2 per nt, kMt-times re-read) are L1 broadcasts; rows
|
||||
// past the N edge skip the load (their smem slots never copy out).
|
||||
const __nv_bfloat16* bias =
|
||||
reinterpret_cast<const __nv_bfloat16*>(p.bias_ptr);
|
||||
__nv_bfloat16* tile_out = reinterpret_cast<__nv_bfloat16*>(fp8_gemm_smem);
|
||||
constexpr int kRowChunks = kBlockN / 8; // 16B chunks per tile row
|
||||
static_assert(kBlockM * kBlockN * 2 <=
|
||||
@@ -624,9 +631,15 @@ __global__ void __launch_bounds__(Traits::kCtaThreads,
|
||||
((c ^ (r & (kRowChunks - 1))) * 8);
|
||||
};
|
||||
const int local_col0 = warp_n * Traits::kWarpN + thread_in_group * 2;
|
||||
const int64_t bias_col0 = (int64_t)block_n * kBlockN;
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < kNt; ++nt) {
|
||||
const int col = local_col0 + nt * 8;
|
||||
const int64_t gcol = bias_col0 + col;
|
||||
const float b0 =
|
||||
bias && gcol < n ? __bfloat162float(bias[gcol]) : 0.0f;
|
||||
const float b1 =
|
||||
bias && gcol + 1 < n ? __bfloat162float(bias[gcol + 1]) : 0.0f;
|
||||
#pragma unroll
|
||||
for (int mt = 0; mt < kMt; ++mt) {
|
||||
const int r0 = warp_m * Traits::kWarpM + group + mt * 16;
|
||||
@@ -635,12 +648,12 @@ __global__ void __launch_bounds__(Traits::kCtaThreads,
|
||||
// m16n8 output, columns tig*2 and tig*2+1 inside one 16B chunk.
|
||||
const int off = col & 7; // element offset within the chunk
|
||||
*reinterpret_cast<__nv_bfloat162*>(out_chunk(r0, col >> 3) + off) =
|
||||
__floats2bfloat162_rn(tile_acc[0] * output_scale,
|
||||
tile_acc[1] * output_scale);
|
||||
__floats2bfloat162_rn(tile_acc[0] * output_scale + b0,
|
||||
tile_acc[1] * output_scale + b1);
|
||||
*reinterpret_cast<__nv_bfloat162*>(out_chunk(r0 + 8, col >> 3) +
|
||||
off) =
|
||||
__floats2bfloat162_rn(tile_acc[2] * output_scale,
|
||||
tile_acc[3] * output_scale);
|
||||
__floats2bfloat162_rn(tile_acc[2] * output_scale + b0,
|
||||
tile_acc[3] * output_scale + b1);
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
@@ -751,13 +764,15 @@ void launch_with_smem(int smem_bytes, dim3 grid, dim3 block,
|
||||
// congruous NT): the 128x128 CTA wins inside one full wave (81 tiles: big
|
||||
// +24%) and from ~1.5 waves up (144: +23%, 256: +39%, 2048^3 123->171 TF),
|
||||
// but loses inside the quantization dip just past one wave (100 tiles =
|
||||
// 1.09 waves: big -8%) where the finer 64x64 grid fills the tail. Below
|
||||
// 3/4 wave the small CTA's extra residency wins or ties (64 tiles: tie).
|
||||
// So: big CTA iff tiles are in [3/4, 1] wave or >= 7/5 waves.
|
||||
// 1.09 waves: big -8%) where the finer 64x64 grid fills the tail. With the
|
||||
// interior fast loop on the big CTA the sub-wave boundary moved down: 63-64
|
||||
// tiles already favor it (63-tile rect +8%, 1024^3 +2%) while 49 tiles
|
||||
// stays small-CTA territory, so the big band opens at 5/8 wave instead of
|
||||
// 3/4.
|
||||
inline bool prefer_small_cta(int64_t tiles_128, int64_t m) {
|
||||
if (m <= 64) return true;
|
||||
const int64_t waves = device_sm_count();
|
||||
if (tiles_128 >= waves - waves / 4 && tiles_128 <= waves) return false;
|
||||
if (tiles_128 >= waves * 5 / 8 && tiles_128 <= waves) return false;
|
||||
return tiles_128 < waves + waves * 2 / 5;
|
||||
}
|
||||
|
||||
@@ -806,8 +821,18 @@ void launch_fp8_gemm(const FP8Params& p, cudaStream_t stream) {
|
||||
}
|
||||
using Traits = Fp8GemmTraits<Fmt, 128, 128, kK, Stages>;
|
||||
dim3 grid((p.n + 127) / 128, (p.m + 127) / 128, p.batch);
|
||||
// Interior-loop specialization on the big CTA as well: with the base-pair
|
||||
// fragment addressing the doubled mainloop no longer spills, and the
|
||||
// predication-free loads win across the band (measured, L20: 1024^3
|
||||
// 98->103T, 2048^3 172->177T, 8192^3 201->205T, 896x1280 124->135T; the
|
||||
// pre-base-pair attempt regressed ~3% at 131 regs). Only congruous
|
||||
// layouts can enter fast_cta, so crosswise (TN) instantiations keep the
|
||||
// single generic body — no dead second loop in their I-cache.
|
||||
constexpr bool kBigFast = !std::is_same_v<LayoutA, ColMajor> &&
|
||||
!std::is_same_v<LayoutB, RowMajor>;
|
||||
launch_with_smem<
|
||||
fp8_gemm_kernel<Traits, LayoutA, LayoutB, GroupRaster, false, false>>(
|
||||
fp8_gemm_kernel<Traits, LayoutA, LayoutB, GroupRaster, false, false,
|
||||
kBigFast>>(
|
||||
Fp8GemmSmem<Traits, LayoutA, LayoutB, false>::kBytes, grid,
|
||||
dim3(Traits::kCtaThreads), stream, p);
|
||||
}
|
||||
|
||||
+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