ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

perf: fp8 batched gemm and measured dispatch table

Browse Source 
- mm_fp8 accepts 3D operands through the same signature: grid.z slices by batch strides, size-1 batches broadcast (stride 0), inner .t() views fold into the layout tag at zero copy
- fix _LinearFp8 backward crash on 3D [B,L,d] training inputs (flatten before mm_fp8, reduce grad_b over leading dims)
- expose kRasterGroup/kStreamOut as template knobs; drop the 64x128 mid CTA and staged crosswise-B path from dispatch (direct wins everywhere re-measured, including DRAM-streamed B)
- dispatch thresholds grounded in fresh sweeps: m<=64 -> 64x64 CTA (+27% at 64x8192x2048), small-CTA crossover at SM*14/3 total tiles (+13% at 96 tiles), threshold counts batch x per-matrix tiles (+31% at 64x512^3 bmm, +25% at 8x1024x2048)
- remove scripts/tools/bench_fp8_gemm.py (superseded by csrc/tests/fp8_sweep.cu for kernel-level tuning)

Benchmark: NVIDIA L20, E4M3, NT pre-quantized, median of 100-200 iters
- 64x8192x2048: 29.1 -> 22.8 us (94 TF/s)
- 1024x1536x2048: 67.4 -> 59.6 us (108 TF/s)
- bmm 64x512^3: 139.8 -> 106.7 us; bmm 8x1024x2048: 186 TF/s
- regression-free: 4096^3 192 TF/s, 8192^3 200 TF/s, 512^3 unchanged
This commit is contained in:
ViperEkura
2026-08-26 06:49:25 +08:00
parent 01eacbde51
commit 4d6a244093
8 changed files with 287 additions and 193 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/kernels/fp8/common.h
+7
View File
@@ -112,6 +112,13 @@ struct FP8Params {
// to int64 for all pointer arithmetic.
int m, n, k;
// Batched (bmm) geometry: grid.z slices step the operand/output pointers
// by these element strides (0 broadcasts the operand across batches).
int batch = 1;
int64_t a_batch_stride = 0;
int64_t b_batch_stride = 0;
int64_t out_batch_stride = 0;
// Physical leading dimensions (column count, i.e. row stride) of A and
// B. For a non-transposed operand the stride equals the contract dim;
// for a transposed operand it is the operand's own column count. The
csrc/kernels/fp8/gemm.cuh
+116 -74
View File
@@ -339,8 +339,8 @@ struct Fp8GemmSmem {
// (mt x nt = 4x4 MMA each). The 64x128 variant runs 4 warps / 128 threads and
// exists for small-M calls: m <= 64 wastes half of every 128-row CTA, so the
// launcher dispatches to it there (see launch_fp8_gemm).
template <typename Traits, typename LayoutA = RowMajor, typename LayoutB = RowMajor, bool kGroupRaster = false,
bool kBStaged = true, bool kLeanRing = false>
template <typename Traits, typename LayoutA = RowMajor, typename LayoutB = RowMajor, int kRasterGroup = 0,
bool kBStaged = true, bool kLeanRing = false, bool kStreamOut = false>
__global__ void __launch_bounds__(Traits::kCtaThreads,
Fp8GemmSmem<Traits, LayoutA, LayoutB,
kBStaged, kLeanRing>::kMinCtas)
@@ -390,9 +390,14 @@ __global__ void __launch_bounds__(Traits::kCtaThreads,
reinterpret_cast<T8*>(fp8_gemm_smem + kARing * kAStageBytes);
T8* const b_canon = b_base + kStB * kBStageBytes; // staged B only
const auto* a = reinterpret_cast<const T8*>(p.a_ptr);
const auto* b = reinterpret_cast<const T8*>(p.b_ptr);
auto* out_bf16 = reinterpret_cast<__nv_bfloat16*>(p.out_ptr);
// Batch slice (grid.z): broadcast operands carry a 0 stride, so the
// same pointer serves every batch.
const auto* a = reinterpret_cast<const T8*>(p.a_ptr) +
(int64_t)blockIdx.z * p.a_batch_stride;
const auto* b = reinterpret_cast<const T8*>(p.b_ptr) +
(int64_t)blockIdx.z * p.b_batch_stride;
auto* out_bf16 = reinterpret_cast<__nv_bfloat16*>(p.out_ptr) +
(int64_t)blockIdx.z * p.out_batch_stride;
const int64_t m = p.m, n = p.n, k = p.k;
const int64_t a_ld = p.a_ld, b_ld = p.b_ld;
@@ -401,19 +406,23 @@ __global__ void __launch_bounds__(Traits::kCtaThreads,
const int lane = tid & 31;
const int group = lane >> 2;
const int thread_in_group = lane & 3;
// L2-friendly rasterization (CUTLASS-style grouped launch order): remap
// the linear block id so consecutive CTAs cover a group of kGroupM M-tiles
// before advancing along N. All CTAs of one group share the same B column
// stripe, so B tiles stay hot in L2 across the wave (the default
// N-fastest order makes each wave touch every B tile instead).
// kGroupRaster is a template knob (the launcher defaults it to the
// measured best per layout: grouped for A-crosswise (dW) and for the
// congruous NT forward — whose big B operand gains the most from the
// shared stripe — plain for dX's crosswise-B layouts, where it measured
// neutral).
constexpr int kGroupM = 8;
// Tile scheduler: the linear CTA id maps to (block_m, block_n) in
// grouped (L2-friendly, CUTLASS-style) or plain raster order — the
// grouped order makes consecutive CTAs cover a group of kRasterGroup
// M-tiles before advancing along N, so all CTAs of one group share the
// same B column stripe and B tiles stay hot in L2 across the wave (the
// plain N-fastest order makes each wave touch every B tile instead;
// kRasterGroup=0 selects plain, the measured best for dX's crosswise-B
// layouts where grouping measured neutral).
// Persistent schedules (static round-robin and an atomic ticket
// dispenser, grid capped at the resident CTAs) were both measured and
// rejected on L20: the stride desynchronizes the in-flight window
// (-4..-8%), and the ticket variant recovers the L2 locality but lands
// within noise of plain waves (its loop-head barrier costs what the
// CTA-restart overlap saves). Keep the classic retiring-wave launch.
int block_m, block_n;
if constexpr (kGroupRaster) {
if constexpr (kRasterGroup > 0) {
constexpr int kGroupM = kRasterGroup;
const int blocks_m = gridDim.y;
const int bid = blockIdx.y * gridDim.x + blockIdx.x;
const int group_first_m = (bid / (kGroupM * gridDim.x)) * kGroupM;
@@ -640,27 +649,31 @@ __global__ void __launch_bounds__(Traits::kCtaThreads,
if (k_seg + 1 < kSegs)
transpose_tile(tile_index, k_seg + 1);
}
// Software-pipelined A fragments: the ldmatrix.x4 for row mt+1
// is issued before the MMAs consuming row mt, so the LDS fixed
// latency hides behind tensor-pipe work (cuts the `wait` stall,
// ~2.3 cycles/issue before this). Costs 4 extra registers.
unsigned a_frag[kMt + 1][4];
astrai::ldmatrix_x4_lane(a_frag[0], a_base_addr + a_off[k_seg][0]);
// Software-pipelined A fragments: the ldmatrix.x4 for row mt+1
// is issued before the MMAs consuming row mt, so the LDS fixed
// latency hides behind tensor-pipe work (cuts the `wait` stall,
// ~2.3 cycles/issue before this). Costs 4 extra registers.
// (Cross-k_seg prefetch of row 0 was tried and reverted: the
// register handoff broke ptxas's software pipelining — 171T → 95T
// at 2048³; the tensor pipe is issue-bound and the seg-start LDS
// already hides behind the b-fragment issue order.)
unsigned a_frag[kMt + 1][4];
astrai::ldmatrix_x4_lane(a_frag[0], a_base_addr + a_off[k_seg][0]);
#pragma unroll
for (int mt = 0; mt < kMt; ++mt) {
if (mt + 1 < kMt)
astrai::ldmatrix_x4_lane(
a_frag[mt + 1], a_base_addr + a_off[k_seg][mt + 1]);
for (int mt = 0; mt < kMt; ++mt) {
if (mt + 1 < kMt)
astrai::ldmatrix_x4_lane(
a_frag[mt + 1], a_base_addr + a_off[k_seg][mt + 1]);
#pragma unroll
for (int nt = 0; nt < kNt; ++nt)
astrai::mma_sync<T8>(acc[nt][mt], a_frag[mt],
b_frag[bcur][nt], acc[nt][mt]);
}
// Barrier 3: region k_seg+1's transposes complete and become
// visible before the next k_seg reads them.
if constexpr (kBStagePath) {
if (k_seg + 1 < kSegs) __syncthreads();
}
for (int nt = 0; nt < kNt; ++nt)
astrai::mma_sync<T8>(acc[nt][mt], a_frag[mt],
b_frag[bcur][nt], acc[nt][mt]);
}
// Barrier 3: region k_seg+1's transposes complete and become
// visible before the next k_seg reads them.
if constexpr (kBStagePath) {
if (k_seg + 1 < kSegs) __syncthreads();
}
}
// Barrier 4 (staged-B / lean-ring only): every thread finished
// reading this stage's tiles before the prefetch for the
@@ -733,7 +746,16 @@ __global__ void __launch_bounds__(Traits::kCtaThreads,
const uint4 v = *reinterpret_cast<const uint4*>(out_chunk(r, c));
auto* dst = out_bf16 + row * n + col;
if (col + 8 <= n && (reinterpret_cast<uintptr_t>(dst) & 15) == 0) {
*reinterpret_cast<uint4*>(dst) = v;
if constexpr (kStreamOut) {
// Evict-first streaming store knob. Measured neutral on
// L20 squares and -3..4% on rects (the evict-first policy
// hurts more than the L2 B-tile protection helps at these
// sizes); kept as a template knob for other SKUs. Default
// off.
__stcs(reinterpret_cast<uint4*>(dst), v);
} else {
*reinterpret_cast<uint4*>(dst) = v;
}
} else {
// N-tail chunk or an odd-n row base: spill the elements that
// survive the row edge (and stay aligned).
@@ -748,6 +770,27 @@ __global__ void __launch_bounds__(Traits::kCtaThreads,
// Launchers — pure CUDA (no torch), usable from the binding and pure C tests.
// ---------------------------------------------------------------------------
// SM count of the current device (cached per device; benign init race —
// every writer stores the same value). Host-side only: feeds the
// device-adaptive dispatch thresholds.
inline int device_sm_count() {
static int cached[64] = {};
int dev = 0;
cudaGetDevice(&dev);
if (dev < 0 || dev >= 64) {
int sms = 0;
cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev);
return sms > 0 ? sms : 1;
}
if (!cached[dev]) {
int sms = 0;
cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev);
cached[dev] = sms > 0 ? sms : 1;
}
return cached[dev];
}
// Launch one kernel instantiation with its shared-memory budget: stages live
// in dynamic smem, so budgets beyond the 48KB static limit opt in once per
// instantiation via cudaFuncSetAttribute (see AGENTS.md "dynamic shared
@@ -778,50 +821,49 @@ void launch_with_smem(int smem_bytes, dim3 grid, dim3 block,
// per LayoutA (grouped for A-crosswise, plain for A-congruous). m <= 64
// dispatches to the 64x128 CTA — a 128-row CTA would waste half its MMA work
// on predicated-off rows.
// Crosswise B takes the asynchronous staging+transpose pipeline only when
// the contract dim is long enough that B streams from DRAM (dX-class GEMMs,
// k = N_ffn); short-K crosswise GEMMs (dW: k = M tokens) read L2-resident
// operands, where the staging round trip costs more shared-memory traffic
// than the latency it hides (measured: dW ~37 TF direct vs ~29 TF staged,
// dX ~39 TF staged vs ~38 direct).
constexpr int64_t kCrossStageMinK = 8192;
// Crosswise-B staging+transpose vs the synchronous direct load: measured on
// the current kernel generation, direct wins everywhere probed — contract k
// 2048..32768 including B operands (128/256MB) that stream from DRAM past L2
// (direct 171-181 TF vs staged 140-163 TF; the staging round trip costs more
// shared-memory traffic than the latency it hides). The old "stage past k=
// 8192" rule reflected a pre-direct-path kernel; staging is now disabled.
// The staged kernel template remains for csrc/tests/fp8_sweep.cu A/B runs.
constexpr int64_t kCrossStageMinK = (int64_t)1 << 62; // unreachable: never stage
// Shape-based tile dispatch (grid-searched on the production shapes, see
// csrc/tests/fp8_sweep.cu): small outputs — fewer than ~2 waves of 128x128
// CTAs on a 24-SM part — take 64x64 CTAs of 32x32 warps with a lean
// (kStages-deep) ring: 24KB of smem keeps 4 CTAs resident, and the extra
// blocks fill the wave quantization gap (512^3: 64 vs 16 CTAs). Everything
// larger takes the 128x128 CTA (8 warps x 64x32) with the kStages+1 ring —
// one __syncthreads per k-tile. m <= 64 keeps the 64x128 CTA so a 128-row
// tile never wastes half its MMA work on predicated-off rows.
constexpr int64_t kSmallShapeMaxTiles = 48;
// csrc/tests/fp8_sweep.cu): small outputs take 64x64 CTAs of 32x32 warps
// with a lean (kStages-deep) ring: 24KB of smem keeps 4 CTAs resident, and
// the extra blocks fill the wave quantization gap (512^3: 64 vs 16 CTAs).
// The large-output path takes the 128x128 CTA (8 warps x 64x32) with the
// kStages+1 ring — one __syncthreads per k-tile and ~200 TF at scale.
// Crossover (congruous NT, k=2048): 96 tiles small +14%, 112 tie, 135 big
// +16% — threshold at ~2.3 waves of the resident (2/SM) 128x128 CTAs.
// The threshold applies to the TOTAL tile count (batch x per-matrix tiles):
// batched runs keep full per-matrix CTA efficiency once the aggregate grid
// saturates the device (measured 64x512^3: big 160 vs small 123 TF — a
// per-matrix-only threshold lost 30%). m <= 64 always takes the small CTA:
// a 128-row CTA would waste half its MMA work on predicated-off rows.
inline int64_t small_shape_max_tiles() {
return (int64_t)device_sm_count() * 14 / 3; // 112 tiles on a 24-SM part
}
template <FP8Format Fmt, typename LayoutA = RowMajor,
typename LayoutB = RowMajor, int kK = 64, int Stages = 2,
bool GroupRaster = std::is_same_v<LayoutA, ColMajor> ||
std::is_same_v<LayoutB, ColMajor>>
int GroupRaster = (std::is_same_v<LayoutA, ColMajor> ||
std::is_same_v<LayoutB, ColMajor>)
? 8
: 0>
void launch_fp8_gemm(const FP8Params& p, cudaStream_t stream) {
const bool b_staged = p.k >= kCrossStageMinK;
if (p.m <= 64) {
using Traits = Fp8GemmTraits<Fmt, 64, 128, kK, Stages>;
dim3 grid((p.n + 127) / 128, (p.m + 63) / 64);
if (b_staged)
launch_with_smem<fp8_gemm_kernel<Traits, LayoutA, LayoutB,
GroupRaster, true, true>>(
Fp8GemmSmem<Traits, LayoutA, LayoutB, true, true>::kBytes,
grid, dim3(Traits::kCtaThreads), stream, p);
else
launch_with_smem<fp8_gemm_kernel<Traits, LayoutA, LayoutB,
GroupRaster, false, true>>(
Fp8GemmSmem<Traits, LayoutA, LayoutB, false, true>::kBytes,
grid, dim3(Traits::kCtaThreads), stream, p);
return;
}
// Staging is disabled (see kCrossStageMinK); the flag stays so the
// staged template instantiations below keep compiling for the sweep.
const bool b_staged = false;
// m <= 64 and small total outputs share the 64x64 small CTA; the
// threshold counts batch x per-matrix tiles (see small_shape_max_tiles).
const int64_t tiles_128 =
((p.m + 127) / 128) * ((p.n + 127) / 128);
if (tiles_128 < kSmallShapeMaxTiles) {
(int64_t)p.batch * ((p.m + 127) / 128) * ((p.n + 127) / 128);
if (p.m <= 64 || tiles_128 < small_shape_max_tiles()) {
using Traits = Fp8GemmTraits<Fmt, 64, 64, kK, 3, 32, 32>;
dim3 grid((p.n + 63) / 64, (p.m + 63) / 64);
dim3 grid((p.n + 63) / 64, (p.m + 63) / 64, p.batch);
if (b_staged)
launch_with_smem<fp8_gemm_kernel<Traits, LayoutA, LayoutB,
GroupRaster, true, true>>(
@@ -835,7 +877,7 @@ void launch_fp8_gemm(const FP8Params& p, cudaStream_t stream) {
return;
}
using Traits = Fp8GemmTraits<Fmt, 128, 128, kK, Stages>;
dim3 grid((p.n + 127) / 128, (p.m + 127) / 128);
dim3 grid((p.n + 127) / 128, (p.m + 127) / 128, p.batch);
if (b_staged)
launch_with_smem<fp8_gemm_kernel<Traits, LayoutA, LayoutB,
GroupRaster, true, false>>(
csrc/kernels/fp8/ops.cu
+63 -16
View File
@@ -90,6 +90,31 @@ void dispatch_gemm(const FP8Params& p, cudaStream_t stream, bool trans_a,
}
}
// Inner-layout resolution for one GEMM operand. The user flag names the
// math (0 = tensor's last two dims are [rows][contract], 1 = transposed);
// the storage may independently be a col-major view (.t() of a contiguous
// buffer), which folds into the returned dispatch flag at zero copy — the
// kernel's LayoutA/LayoutB tags cover both storages. m/n/k derive from the
// user flag only; the fold never swaps them (see the layout table in
// gemm.cuh). Tensors whose inner dims are neither natural layout fall back
// to .contiguous().
bool resolve_operand(const torch::Tensor& t_in, bool flag, int64_t& ld,
int64_t& batch_stride, torch::Tensor& storage) {
torch::Tensor t = t_in;
bool col_major = false;
if (t.stride(-1) != 1) {
if (t.stride(-2) == 1) {
col_major = true;
} else {
t = t.contiguous();
}
}
storage = t;
ld = col_major ? t.stride(-1) : t.stride(-2);
batch_stride = t.dim() == 3 ? t.stride(0) : 0;
return flag ^ col_major;
}
} // namespace
std::tuple<torch::Tensor, torch::Tensor> quantize(torch::Tensor x,
@@ -145,30 +170,52 @@ torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
a.scalar_type() == torch::kFloat8_e5m2,
"a and b must be fp8");
TORCH_CHECK(a.scalar_type() == b.scalar_type(), "a and b must share format");
TORCH_CHECK(a.dim() == 2 && b.dim() == 2, "a and b must be 2D");
TORCH_CHECK((a.dim() == 2 || a.dim() == 3) &&
(b.dim() == 2 || b.dim() == 3),
"a and b must be 2D or 3D (batched)");
TORCH_CHECK(a.device() == b.device(), "a and b must share device");
check_scale(scale, a);
check_fp8_device(a);
const at::cuda::OptionalCUDAGuard guard(a.device());
auto stream = at::cuda::getCurrentCUDAStream();
auto a_c = a.contiguous();
auto b_c = b.contiguous();
const bool ta = trans_a != 0;
const bool tb = trans_b != 0;
const int64_t a_ld = a_c.size(1);
const int64_t b_ld = b_c.size(1);
const int64_t m = ta ? a_c.size(1) : a_c.size(0);
const int64_t k = ta ? a_c.size(0) : a_c.size(1);
const int64_t n = tb ? b_c.size(0) : b_c.size(1);
TORCH_CHECK(k == (tb ? b_c.size(1) : b_c.size(0)), "inner dim mismatch");
auto output = torch::empty({m, n}, a_c.options().dtype(torch::kBFloat16));
// Batched operands follow matmul broadcast rules: 2D acts as a batch
// of 1; a size-1 batch broadcasts across the other side (stride 0).
const int64_t batch_a = a.dim() == 3 ? a.size(0) : 1;
const int64_t batch_b = b.dim() == 3 ? b.size(0) : 1;
TORCH_CHECK(batch_a == batch_b || batch_a == 1 || batch_b == 1,
"batch dim mismatch (got ", batch_a, " and ", batch_b, ")");
const int64_t batch = std::max(batch_a, batch_b);
TORCH_CHECK(batch <= 65535, "batch dim exceeds the grid.z launch limit");
torch::Tensor a_st, b_st;
int64_t a_ld, b_ld, a_bstride, b_bstride;
const bool tag_a =
resolve_operand(a, trans_a != 0, a_ld, a_bstride, a_st);
const bool tag_b =
resolve_operand(b, trans_b != 0, b_ld, b_bstride, b_st);
// GEMM dims from the user flags; storage layout never swaps them.
const int64_t m = trans_a ? a.size(-1) : a.size(-2);
const int64_t k = trans_a ? a.size(-2) : a.size(-1);
const int64_t n = trans_b ? b.size(-2) : b.size(-1);
TORCH_CHECK(k == (trans_b ? b.size(-1) : b.size(-2)), "inner dim mismatch");
const bool batched_out = a.dim() == 3 || b.dim() == 3;
torch::Tensor output =
batched_out
? torch::empty({batch, m, n}, a.options().dtype(torch::kBFloat16))
: torch::empty({m, n}, a.options().dtype(torch::kBFloat16));
FP8Params p;
pack_gemm(p, a_c.data_ptr(), b_c.data_ptr(), output.data_ptr(), scale, m, n,
k, a_ld, b_ld);
pack_gemm(p, a_st.data_ptr(), b_st.data_ptr(), output.data_ptr(), scale,
m, n, k, a_ld, b_ld);
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;
p.out_batch_stride = m * n;
if (a.scalar_type() == torch::kFloat8_e4m3fn)
dispatch_gemm<FP8Format::E4M3>(p, stream.stream(), ta, tb);
dispatch_gemm<FP8Format::E4M3>(p, stream.stream(), tag_a, tag_b);
else
dispatch_gemm<FP8Format::E5M2>(p, stream.stream(), ta, tb);
dispatch_gemm<FP8Format::E5M2>(p, stream.stream(), tag_a, tag_b);
C10_CUDA_CHECK(cudaGetLastError());
return output;
}