perf: speed up fp8 gemm across small and large shapes
- parameterize warp tile (WarpM/WarpN) in Fp8GemmTraits; MMA loops, fragment arrays and epilogue scale with kMt/kNt instead of the fixed 64x32/4x4, enabling cuBLAS-style 64x64 CTAs of 32x32 warps - dispatch by output tiling (grid-searched via csrc/tests/fp8_sweep.cu): fewer than 48 output tiles take 64x64/32x32 with a lean ring (4 CTAs/SM fill the wave-quantization gap: 512^3 goes 16 -> 64 CTAs); larger shapes keep 128x128 with the kStages+1 ring - kStages+1 canonic ring rotation drops the post-compute barrier on the congruous path (one __syncthreads per k-tile); LeanRing keeps the kStages ring for the small CTA; direct-crosswise operands always rotate kStages+1 (their prefetch issues right after barrier 1 and would race a lean ring - caught by the pure C layout suite) - stage the bf16 epilogue through the reclaimed operand smem: swizzled scatter + barrier + coalesced 16B copy-out replaces 8 disjoint 16B per-warp segments (~50% write efficiency before) - hoist per-lane ldmatrix swizzle offsets out of the mainloop (stage-relative table + ring-base add) so the innermost loop stops recomputing IMAD/LOP3 address chains - bypass the torch.library dispatch for real CUDA tensors in quantize/mm_fp8 wrappers (~5us/call, ~40% of a 512-wide call's wall time); fake/subclass tensors keep the custom_op route vs the previous kernel + python path, wall clock on NT squares: 512^3 52 -> 13us (4.0x, 5.2 -> 20.5 TF, now 1.36x cuBLAS _scaled_mm), 1024^3 1.05x, 2048^3 1.02x (46.9 -> 48.2 TF kernel-only); correctness: 4 layouts x 6 shapes pure C suite PASS, 588 pytest PASS
This commit is contained in:
@@ -48,25 +48,36 @@ using transpose_layout_t = typename transpose_layout<Layout>::type;
|
||||
|
||||
// Compile-time tile configuration, mirroring KernelTraits<HEAD_DIM, BC,
|
||||
// WARPS, STAGES> in the attention kernels. `Fmt` selects the FP8 conversion
|
||||
// and the MMA PTX mnemonic; the remaining parameters shape the CTA tile and
|
||||
// the cp.async pipeline depth.
|
||||
template <FP8Format Fmt, int BlockM, int BlockN, int K, int Stages>
|
||||
// and the MMA PTX mnemonic; the remaining parameters shape the CTA tile, the
|
||||
// warp tile (WarpM x WarpN — e.g. 64x32 on the 128x128 CTA, or 32x32 on the
|
||||
// cuBLAS-style 64x64 small CTA that lifts small-shape occupancy) and the
|
||||
// cp.async pipeline depth.
|
||||
template <FP8Format Fmt, int BlockM, int BlockN, int K, int Stages,
|
||||
int WarpM = 64, int WarpN = 32>
|
||||
struct Fp8GemmTraits {
|
||||
static constexpr FP8Format kFormat = Fmt;
|
||||
static constexpr int kBlockM = BlockM;
|
||||
static constexpr int kBlockN = BlockN;
|
||||
static constexpr int kK = K;
|
||||
static constexpr int kStages = Stages;
|
||||
static constexpr int kWarpM = WarpM;
|
||||
static constexpr int kWarpN = WarpN;
|
||||
static constexpr bool kIsE5M2 = (Fmt == FP8Format::E5M2);
|
||||
static constexpr __nv_fp8_interpretation_t kNvFormat =
|
||||
kIsE5M2 ? __NV_E5M2 : __NV_E4M3;
|
||||
static constexpr float kFp8Max = kIsE5M2 ? 57344.0f : 448.0f;
|
||||
|
||||
// Derived launch geometry: 64x32 warp tiles give the CTA thread count.
|
||||
// The shared-memory budget is layout-aware (crosswise operands add K-
|
||||
// major staging + a canonical buffer), so it lives in Fp8GemmSmem in
|
||||
// gemm.cuh together with the resident-CTA hint for __launch_bounds__.
|
||||
static constexpr int kCtaThreads = (BlockM / 64) * (BlockN / 32) * 32;
|
||||
// Derived launch geometry: WarpM x WarpN warp tiles tile the CTA. The
|
||||
// shared-memory budget is layout-aware (crosswise operands add K-major
|
||||
// staging + a canonical buffer), so it lives in Fp8GemmSmem in gemm.cuh
|
||||
// together with the resident-CTA hint for __launch_bounds__.
|
||||
static constexpr int kWarpsM = BlockM / WarpM;
|
||||
static constexpr int kWarpsN = BlockN / WarpN;
|
||||
static constexpr int kCtaThreads = kWarpsM * kWarpsN * 32;
|
||||
static_assert(kWarpsM * WarpM == BlockM && kWarpsN * WarpN == BlockN,
|
||||
"warp tiles must exactly tile the CTA");
|
||||
static_assert(WarpM % 16 == 0 && WarpN % 8 == 0,
|
||||
"warp tile must be a multiple of the m16n8 MMA shape");
|
||||
};
|
||||
|
||||
// Quantize-kernel parameter POD: float input (bf16 / fp16 / fp32) -> FP8
|
||||
|
||||
+236
-119
@@ -290,16 +290,19 @@ transpose_crosswise_region(T8* tile, const T8* staging, int idx, int quad0) {
|
||||
}
|
||||
}
|
||||
|
||||
// Layout-aware shared-memory budget and occupancy hint. A congruous operand
|
||||
// needs its kStages rotating canonical buffers; a staged-crosswise operand
|
||||
// (crosswise B with kBStaged) needs kStages K-major staging buffers plus ONE
|
||||
// canonical buffer (rewritten every tile by the in-kernel transpose); a
|
||||
// direct-crosswise operand rotates kStages+1 canonical buffers so its load
|
||||
// can run ahead of the compute phase (see the kernel's pipelining note).
|
||||
// Layout-aware shared-memory budget and occupancy hint. Canonic rings hold
|
||||
// kStages+1 buffers (LeanRing=false): the load for tile i+kStages targets
|
||||
// slot (i-1)%(kStages+1) — already consumed — so the pure-congruous path
|
||||
// needs no post-compute barrier (one __syncthreads per k-tile). LeanRing
|
||||
// keeps the ring at kStages buffers for small CTAs whose occupancy comes
|
||||
// from more resident CTAs (less smem) rather than a deeper rotation; it
|
||||
// brings back barrier 4. A staged-crosswise B always costs kStages K-major
|
||||
// staging buffers + one canonical buffer.
|
||||
// The 48KB static-smem watermark picks the resident-CTA hint for
|
||||
// __launch_bounds__ (sm_89: 100KB smem per SM, so two CTAs fit while each
|
||||
// stays within the static budget).
|
||||
template <typename Traits, typename LayoutA, typename LayoutB, bool StagedB>
|
||||
template <typename Traits, typename LayoutA, typename LayoutB, bool StagedB,
|
||||
bool LeanRing = false>
|
||||
struct Fp8GemmSmem {
|
||||
// Crosswise = the stage-load's view: A's tag directly, B's transposed.
|
||||
// A-crosswise always loads direct (L2-typical activations); B-crosswise
|
||||
@@ -309,11 +312,17 @@ struct Fp8GemmSmem {
|
||||
static constexpr bool kBStagePath = kCrossB && StagedB;
|
||||
static constexpr bool kDirectA = kCrossA;
|
||||
static constexpr bool kDirectB = kCrossB && !kBStagePath;
|
||||
// LeanRing shrinks only the congruous (async) operand rings; a direct
|
||||
// operand's ring stays kStages+1 deep (see the kernel's ring note).
|
||||
static constexpr int kARing = kDirectA ? Traits::kStages + 1
|
||||
: Traits::kStages + !LeanRing;
|
||||
static constexpr int kBRing =
|
||||
kBStagePath ? Traits::kStages + 1
|
||||
: (kDirectB ? Traits::kStages + 1
|
||||
: Traits::kStages + !LeanRing);
|
||||
static constexpr int kBytes =
|
||||
(kDirectA ? Traits::kStages + 1 : Traits::kStages) *
|
||||
Traits::kBlockM * Traits::kK +
|
||||
(kDirectB || kBStagePath ? Traits::kStages + 1 : Traits::kStages) *
|
||||
Traits::kBlockN * Traits::kK;
|
||||
kARing * Traits::kBlockM * Traits::kK +
|
||||
kBRing * Traits::kBlockN * Traits::kK;
|
||||
static constexpr int kMinCtas = kBytes <= 48 * 1024 ? 2 : 1;
|
||||
};
|
||||
|
||||
@@ -331,10 +340,10 @@ struct Fp8GemmSmem {
|
||||
// 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 kBStaged = true, bool kLeanRing = false>
|
||||
__global__ void __launch_bounds__(Traits::kCtaThreads,
|
||||
Fp8GemmSmem<Traits, LayoutA, LayoutB,
|
||||
kBStaged>::kMinCtas)
|
||||
kBStaged, kLeanRing>::kMinCtas)
|
||||
fp8_gemm_kernel(FP8Params p) {
|
||||
using T8 = std::conditional_t<Traits::kIsE5M2, __nv_fp8_e5m2, __nv_fp8_e4m3>;
|
||||
constexpr int kBlockM = Traits::kBlockM;
|
||||
@@ -357,19 +366,25 @@ __global__ void __launch_bounds__(Traits::kCtaThreads,
|
||||
// shared memory so deep pipelines (kStages * (kBlockM + kBlockN) * kK >
|
||||
// 48KB static limit) opt in via cudaFuncSetAttribute in the launcher.
|
||||
extern __shared__ __align__(16) char fp8_gemm_smem[];
|
||||
// Per operand: congruous = kStages rotating canonical buffers; direct-
|
||||
// crosswise = kStages+1 of them (the load for tile i+kStages targets
|
||||
// buffer (i-1)%(kStages+1) — the one compute(i-1) finished reading at
|
||||
// the previous barrier — so it issues right after barrier 1 and its
|
||||
// global-load latency overlaps the MMA phase below); staged-crosswise
|
||||
// (B) = kStages K-major staging buffers (filled by cp.async, one per
|
||||
// tile in flight) followed by one canonical buffer the per-tile
|
||||
// transpose rewrites.
|
||||
// Per operand: congruous = kStages+1 rotating canonical buffers — the
|
||||
// load for tile i+kStages targets slot (i-1)%(kStages+1), which compute
|
||||
// finished reading before this iteration's barrier 1, so NO post-compute
|
||||
// barrier is needed on the pure-congruous path (one __syncthreads per
|
||||
// k-tile, the classic multistage rotation); direct-crosswise rotates the
|
||||
// same kStages+1 ring for the same reason; staged-crosswise (B) keeps
|
||||
// kStages K-major staging buffers (filled by cp.async) plus ONE canonical
|
||||
// buffer the per-tile transpose rewrites (its barrier structure keeps
|
||||
// barrier 4).
|
||||
constexpr int kAStageBytes = kBlockM * kK;
|
||||
constexpr int kBStageBytes = kBlockN * kK;
|
||||
constexpr int kARing = kDirectA ? kStages + 1 : kStages; // A canonic ring
|
||||
constexpr int kBRing = kDirectB ? kStages + 1 : kStages; // B canonic ring
|
||||
constexpr int kStB = kStages; // B staging ring size (see above)
|
||||
// Direct-crosswise operands always rotate kStages+1 buffers: their
|
||||
// prefetch issues right after barrier 1 (targeting the slot compute(i-1)
|
||||
// released), so a kStages-deep lean ring would race the in-flight MMA
|
||||
// reads. The lean ring applies only to congruous operands, whose cp.async
|
||||
// prefetch sits behind the restored barrier 4.
|
||||
constexpr int kARing = kDirectA ? kStages + 1 : kStages + !kLeanRing;
|
||||
constexpr int kBRing = kDirectB ? kStages + 1 : kStages + !kLeanRing;
|
||||
constexpr int kStB = kStages; // B staging ring size
|
||||
T8* const a_base = reinterpret_cast<T8*>(fp8_gemm_smem);
|
||||
T8* const b_base =
|
||||
reinterpret_cast<T8*>(fp8_gemm_smem + kARing * kAStageBytes);
|
||||
@@ -410,17 +425,24 @@ __global__ void __launch_bounds__(Traits::kCtaThreads,
|
||||
block_m = blockIdx.y;
|
||||
block_n = blockIdx.x;
|
||||
}
|
||||
// 128x128 CTA = 8 warps as 2x4 warp tiles of 64x32 (mt x nt = 4x4 MMA).
|
||||
constexpr int warps_n = kBlockN / 32;
|
||||
const int warp_m = warp / warps_n;
|
||||
const int warp_n = warp % warps_n;
|
||||
const int64_t row_base = (int64_t)block_m * kBlockM + warp_m * 64 + group;
|
||||
const int64_t output_col =
|
||||
(int64_t)block_n * kBlockN + warp_n * 32 + thread_in_group * 2;
|
||||
const int a_row0 = warp_m * 64; // + mt * 16 in the loop
|
||||
const int b_row0 = warp_n * 32; // + nt * 8
|
||||
// CTA = (BlockM/WarpM) x (BlockN/WarpN) warps of WarpM x WarpN tiles,
|
||||
// each warp computing (WarpM/16) x (WarpN/8) m16n8k32 MMAs (mt x nt).
|
||||
// The default 128x128 CTA runs 8 warps of 64x32 (mt x nt = 4x4); the
|
||||
// small-shape path uses 64x64 CTAs of 32x32 warps (cuBLAS-style) so more
|
||||
// CTAs fit per SM (see launch_fp8_gemm).
|
||||
constexpr int kMt = Traits::kWarpM / 16; // 16-row MMA tiles per warp
|
||||
constexpr int kNt = Traits::kWarpN / 8; // 8-col MMA tiles per warp
|
||||
const int warp_m = warp / Traits::kWarpsN;
|
||||
const int warp_n = warp % Traits::kWarpsN;
|
||||
const int64_t row_base =
|
||||
(int64_t)block_m * kBlockM + warp_m * Traits::kWarpM + group;
|
||||
const int64_t output_col = (int64_t)block_n * kBlockN +
|
||||
warp_n * Traits::kWarpN +
|
||||
thread_in_group * 2;
|
||||
const int a_row0 = warp_m * Traits::kWarpM; // + mt * 16 in the loop
|
||||
const int b_row0 = warp_n * Traits::kWarpN; // + nt * 8
|
||||
const float scale = *p.scale;
|
||||
float acc[4][4][4] = {}; // [nt][mt][acc]
|
||||
float acc[kNt][kMt][4] = {}; // [nt][mt][acc]
|
||||
|
||||
// Both operands end up in the canonical [M][kK] / [N][kK] shared tiles
|
||||
// the MMA fragments read, regardless of their global layout. A's tag
|
||||
@@ -498,6 +520,45 @@ __global__ void __launch_bounds__(Traits::kCtaThreads,
|
||||
const int rh8 = (lane >> 3) & 1; // +8 rows (A: lanes 8-15, 24-31)
|
||||
const int rh16 = lane >> 4; // +1 chunk (A: lanes 16-31; B uses rh8)
|
||||
|
||||
// Precomputed per-lane fragment offsets (stage-relative): the XOR
|
||||
// swizzle inside tile_at depends only on (row, chunk) — never on the
|
||||
// ring slot or tile_index — so every lane's ldmatrix address is its
|
||||
// stage base plus one of these fixed offsets. Building the table once,
|
||||
// outside the mainloop, removes the per-k_seg swizzle arithmetic
|
||||
// (IMAD/LOP3 chains) from the innermost loop; the SASS compute window
|
||||
// was ~36% integer address math before this.
|
||||
constexpr int kSegs = kK / kMmaK;
|
||||
unsigned a_off[kSegs][kMt]; // stage-relative byte offsets
|
||||
unsigned b_off[kSegs][kNt];
|
||||
{
|
||||
// The probe addresses are converted and immediately rebased to the
|
||||
// stage origin, so the table holds pure offsets to add to any ring
|
||||
// slot's converted base (double-adding the base was the bug here).
|
||||
const unsigned a0 = __cvta_generic_to_shared(a_base);
|
||||
#pragma unroll
|
||||
for (int s = 0; s < kSegs; ++s) {
|
||||
#pragma unroll
|
||||
for (int mt = 0; mt < kMt; ++mt)
|
||||
a_off[s][mt] =
|
||||
__cvta_generic_to_shared(
|
||||
tile_at<kK>(a_base, a_row0 + mt * 16 + rh8 * 8 + r7,
|
||||
(s * 2 + rh16) * 16)) -
|
||||
a0;
|
||||
}
|
||||
const T8* b_probe = kBStagePath ? b_canon : b_base;
|
||||
const unsigned b0 = __cvta_generic_to_shared(b_probe);
|
||||
#pragma unroll
|
||||
for (int s = 0; s < kSegs; ++s) {
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < kNt; ++nt)
|
||||
b_off[s][nt] =
|
||||
__cvta_generic_to_shared(
|
||||
tile_at<kK>(b_probe, b_row0 + nt * 8 + r7,
|
||||
(s * 2 + rh8) * 16)) -
|
||||
b0;
|
||||
}
|
||||
}
|
||||
|
||||
// Prime the pipeline. Each committed group occupies one circular shared
|
||||
// memory stage; the loop also handles K dimensions smaller than kStages.
|
||||
// Direct loads run synchronously here (back to back with their commit);
|
||||
@@ -532,7 +593,6 @@ __global__ void __launch_bounds__(Traits::kCtaThreads,
|
||||
// a time so each region's transpose overlaps the previous region's
|
||||
// MMA sequence (the transposes are pure shared-memory traffic — B's
|
||||
// global path stayed fully asynchronous above).
|
||||
constexpr int kSegs = kK / kMmaK;
|
||||
if constexpr (kBStagePath) {
|
||||
transpose_tile(tile_index, 0);
|
||||
// Barrier 2: region 0 visible to every thread before its
|
||||
@@ -544,41 +604,35 @@ __global__ void __launch_bounds__(Traits::kCtaThreads,
|
||||
const T8* b_tile = kBStagePath
|
||||
? b_canon
|
||||
: b_base + (size_t)(tile_index % kBRing) * kBStageBytes;
|
||||
const unsigned a_base_addr = __cvta_generic_to_shared(a_tile);
|
||||
const unsigned b_base_addr = __cvta_generic_to_shared(b_tile);
|
||||
|
||||
// 4 ldmatrix.x2 (B) + 4 ldmatrix.x4 (A) feed 16 mma.sync per k_seg —
|
||||
// 0.5 load instructions per MMA, versus 4.5 scalar LDS per MMA in
|
||||
// the 128x64-tile version (the kernel was LSU-issue-bound there).
|
||||
// B fragments double-buffer across k_segs while B is congruous (no
|
||||
// region writes in flight); a crosswise B reloads per k_seg after
|
||||
// the region's transpose became visible.
|
||||
unsigned b_frag[2][4][2];
|
||||
// kNt ldmatrix.x2 (B) + kMt ldmatrix.x4 (A) feed kMt*kNt*2 mma.sync
|
||||
// per k_seg — 0.5 load instructions per MMA, versus 4.5 scalar LDS
|
||||
// per MMA in the 128x64-tile version (the kernel was LSU-issue-bound
|
||||
// there). B fragments double-buffer across k_segs while B is
|
||||
// congruous (no region writes in flight); a crosswise B reloads per
|
||||
// k_seg after the region's transpose became visible.
|
||||
unsigned b_frag[2][kNt][2];
|
||||
if constexpr (!kBStagePath) {
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 4; ++nt) {
|
||||
const int row = b_row0 + nt * 8 + r7;
|
||||
for (int nt = 0; nt < kNt; ++nt)
|
||||
astrai::ldmatrix_x2_lane(b_frag[0][nt],
|
||||
frag_addr<T8, kK>(b_tile, row, rh8));
|
||||
}
|
||||
b_base_addr + b_off[0][nt]);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int k_seg = 0; k_seg < kSegs; ++k_seg) {
|
||||
const int bcur = k_seg & 1, bnext = bcur ^ 1;
|
||||
if constexpr (kBStagePath) {
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 4; ++nt) {
|
||||
const int row = b_row0 + nt * 8 + r7;
|
||||
for (int nt = 0; nt < kNt; ++nt)
|
||||
astrai::ldmatrix_x2_lane(
|
||||
b_frag[bcur][nt],
|
||||
frag_addr<T8, kK>(b_tile, row, k_seg * 2 + rh8));
|
||||
}
|
||||
b_frag[bcur][nt], b_base_addr + b_off[k_seg][nt]);
|
||||
} else if (k_seg + 1 < kSegs) {
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 4; ++nt) {
|
||||
const int row = b_row0 + nt * 8 + r7;
|
||||
for (int nt = 0; nt < kNt; ++nt)
|
||||
astrai::ldmatrix_x2_lane(
|
||||
b_frag[bnext][nt],
|
||||
frag_addr<T8, kK>(b_tile, row, (k_seg + 1) * 2 + rh8));
|
||||
}
|
||||
b_frag[bnext][nt], b_base_addr + b_off[k_seg + 1][nt]);
|
||||
}
|
||||
// Region k_seg+1's transpose overlaps this region's MMA work
|
||||
// (disjoint canonical regions, no race).
|
||||
@@ -590,22 +644,17 @@ __global__ void __launch_bounds__(Traits::kCtaThreads,
|
||||
// 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[5][4];
|
||||
astrai::ldmatrix_x4_lane(
|
||||
a_frag[0], frag_addr<T8, kK>(a_tile, a_row0 + rh8 * 8 + r7,
|
||||
k_seg * 2 + rh16));
|
||||
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 < 4; ++mt) {
|
||||
if (mt < 3)
|
||||
for (int mt = 0; mt < kMt; ++mt) {
|
||||
if (mt + 1 < kMt)
|
||||
astrai::ldmatrix_x4_lane(
|
||||
a_frag[mt + 1],
|
||||
frag_addr<T8, kK>(a_tile,
|
||||
a_row0 + (mt + 1) * 16 + rh8 * 8 + r7,
|
||||
k_seg * 2 + rh16));
|
||||
a_frag[mt + 1], a_base_addr + a_off[k_seg][mt + 1]);
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 4; ++nt)
|
||||
astrai::mma_sync<T8>(acc[nt][mt], a_frag[mt], b_frag[bcur][nt],
|
||||
acc[nt][mt]);
|
||||
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.
|
||||
@@ -613,45 +662,84 @@ __global__ void __launch_bounds__(Traits::kCtaThreads,
|
||||
if (k_seg + 1 < kSegs) __syncthreads();
|
||||
}
|
||||
}
|
||||
// Barrier 4: every thread finished reading this stage's tiles before
|
||||
// the prefetch for the (i+kStages)-th tile overwrites them (and the
|
||||
// next iteration's transposes rewrite the canonical buffer).
|
||||
__syncthreads();
|
||||
// Barrier 4 (staged-B / lean-ring only): every thread finished
|
||||
// reading this stage's tiles before the prefetch for the
|
||||
// (i+kStages)-th tile overwrites them (and the next iteration's
|
||||
// transposes rewrite the canonical buffer). With the kStages+1
|
||||
// canonic rotation the prefetch targets the slot compute(i-1)
|
||||
// released before barrier 1, so the pure-congruous path skips this
|
||||
// barrier entirely — one __syncthreads per k-tile.
|
||||
if constexpr (kBStagePath || kLeanRing) __syncthreads();
|
||||
if (tile_index + kStages < tile_count) {
|
||||
load_async(tile_index + kStages);
|
||||
astrai::cp_async_commit_group();
|
||||
}
|
||||
}
|
||||
|
||||
// Direct bf16 epilogue through the operand shared memory: the A/B rings
|
||||
// are dead once the mainloop ends, so their space stages the output tile
|
||||
// (kBlockM x kBlockN bf16, always <= the ring budget). Threads first
|
||||
// scatter their accumulators into the tile (STS.32 of bf16x2 pairs), a
|
||||
// barrier makes the tile coherent, then the whole CTA copies it out in
|
||||
// fully-coalesced 16B chunks. The direct per-thread stores this replaces
|
||||
// hit 8 disjoint 16B segments per warp (rows are n*2 bytes apart), ~50%
|
||||
// write efficiency — measurable at 2048+ where the epilogue is ~8% of
|
||||
// runtime. The 16B-chunk XOR swizzle (chunk index ^ row) keeps both the
|
||||
// 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;
|
||||
__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 <=
|
||||
kARing * kBlockM * kK + kBRing * kBlockN * kK,
|
||||
"output tile must fit the reclaimed operand smem");
|
||||
// Swizzled address of one 16B chunk (row r, chunk c) of the tile.
|
||||
auto out_chunk = [&](int r, int c) -> __nv_bfloat16* {
|
||||
return tile_out + (size_t)r * kBlockN +
|
||||
((c ^ (r & (kRowChunks - 1))) * 8);
|
||||
};
|
||||
const int local_col0 = warp_n * Traits::kWarpN + thread_in_group * 2;
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 4; ++nt) {
|
||||
const int64_t col = output_col + nt * 8;
|
||||
// Per-row store: FP8 packs two adjacent columns into one 16-bit
|
||||
// write, BF16 into one 32-bit __nv_bfloat162 (single cvt+pack
|
||||
// instruction); boundary or unaligned columns fall back to scalar
|
||||
// converts so a pack never crosses the row edge or misaligns.
|
||||
auto store_out = [&](int64_t row, float v0, float v1) {
|
||||
if (row >= m) return;
|
||||
const float r0 = v0 * output_scale;
|
||||
const float r1 = v1 * output_scale;
|
||||
auto* dst = out_bf16 + row * n + col;
|
||||
if (col + 1 < n && (reinterpret_cast<uintptr_t>(dst) & 3) == 0) {
|
||||
*reinterpret_cast<__nv_bfloat162*>(dst) =
|
||||
__floats2bfloat162_rn(r0, r1);
|
||||
} else {
|
||||
dst[0] = __float2bfloat16(r0);
|
||||
if (col + 1 < n) dst[1] = __float2bfloat16(r1);
|
||||
}
|
||||
};
|
||||
for (int nt = 0; nt < kNt; ++nt) {
|
||||
const int col = local_col0 + nt * 8;
|
||||
#pragma unroll
|
||||
for (int mt = 0; mt < 4; ++mt) {
|
||||
const int64_t row0 = row_base + mt * 16;
|
||||
float* tile_acc = acc[nt][mt];
|
||||
if (col < n) {
|
||||
store_out(row0, tile_acc[0], tile_acc[1]);
|
||||
store_out(row0 + 8, tile_acc[2], tile_acc[3]);
|
||||
}
|
||||
for (int mt = 0; mt < kMt; ++mt) {
|
||||
const int r0 = warp_m * Traits::kWarpM + group + mt * 16;
|
||||
const float* tile_acc = acc[nt][mt];
|
||||
// Two bf16x2 stores per accumulator tile: rows g and g+8 of the
|
||||
// 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);
|
||||
*reinterpret_cast<__nv_bfloat162*>(out_chunk(r0 + 8, col >> 3) +
|
||||
off) =
|
||||
__floats2bfloat162_rn(tile_acc[2] * output_scale,
|
||||
tile_acc[3] * output_scale);
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
// Coalesced copy-out: thread -> one 16B chunk; consecutive threads walk
|
||||
// a row so each global transaction covers a full 128B line.
|
||||
const int64_t row0_global = (int64_t)block_m * kBlockM;
|
||||
const int64_t col0_global = (int64_t)block_n * kBlockN;
|
||||
constexpr int kTotalChunks = kBlockM * kRowChunks;
|
||||
for (int idx = tid; idx < kTotalChunks; idx += kCtaThreads) {
|
||||
const int r = idx / kRowChunks;
|
||||
const int c = idx % kRowChunks;
|
||||
const int64_t row = row0_global + r;
|
||||
if (row >= m) break; // rows are consecutive: nothing left in range
|
||||
const int64_t col = col0_global + (int64_t)c * 8;
|
||||
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;
|
||||
} else {
|
||||
// N-tail chunk or an odd-n row base: spill the elements that
|
||||
// survive the row edge (and stay aligned).
|
||||
const __nv_bfloat16* elems =
|
||||
reinterpret_cast<const __nv_bfloat16*>(&v);
|
||||
for (int e = 0; e < 8 && col + e < n; ++e) dst[e] = elems[e];
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -698,37 +786,66 @@ void launch_with_smem(int smem_bytes, dim3 grid, dim3 block,
|
||||
// dX ~39 TF staged vs ~38 direct).
|
||||
constexpr int64_t kCrossStageMinK = 8192;
|
||||
|
||||
// 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;
|
||||
|
||||
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>>
|
||||
bool GroupRaster = std::is_same_v<LayoutA, ColMajor> ||
|
||||
std::is_same_v<LayoutB, ColMajor>>
|
||||
void launch_fp8_gemm(const FP8Params& p, cudaStream_t stream) {
|
||||
dim3 grid((p.n + 127) / 128, (p.m + 127) / 128);
|
||||
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>>(
|
||||
Fp8GemmSmem<Traits, LayoutA, LayoutB, true>::kBytes, grid,
|
||||
dim3(Traits::kCtaThreads), stream, p);
|
||||
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>>(
|
||||
Fp8GemmSmem<Traits, LayoutA, LayoutB, false>::kBytes, grid,
|
||||
dim3(Traits::kCtaThreads), stream, p);
|
||||
} else {
|
||||
using Traits = Fp8GemmTraits<Fmt, 128, 128, kK, Stages>;
|
||||
if (b_staged)
|
||||
launch_with_smem<fp8_gemm_kernel<Traits, LayoutA, LayoutB,
|
||||
GroupRaster, true>>(
|
||||
Fp8GemmSmem<Traits, LayoutA, LayoutB, true>::kBytes, grid,
|
||||
dim3(Traits::kCtaThreads), stream, p);
|
||||
else
|
||||
launch_with_smem<fp8_gemm_kernel<Traits, LayoutA, LayoutB,
|
||||
GroupRaster, false>>(
|
||||
Fp8GemmSmem<Traits, LayoutA, LayoutB, false>::kBytes, grid,
|
||||
dim3(Traits::kCtaThreads), stream, p);
|
||||
GroupRaster, false, true>>(
|
||||
Fp8GemmSmem<Traits, LayoutA, LayoutB, false, true>::kBytes,
|
||||
grid, dim3(Traits::kCtaThreads), stream, p);
|
||||
return;
|
||||
}
|
||||
const int64_t tiles_128 =
|
||||
((p.m + 127) / 128) * ((p.n + 127) / 128);
|
||||
if (tiles_128 < kSmallShapeMaxTiles) {
|
||||
using Traits = Fp8GemmTraits<Fmt, 64, 64, kK, 3, 32, 32>;
|
||||
dim3 grid((p.n + 63) / 64, (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;
|
||||
}
|
||||
using Traits = Fp8GemmTraits<Fmt, 128, 128, kK, Stages>;
|
||||
dim3 grid((p.n + 127) / 128, (p.m + 127) / 128);
|
||||
if (b_staged)
|
||||
launch_with_smem<fp8_gemm_kernel<Traits, LayoutA, LayoutB,
|
||||
GroupRaster, true, false>>(
|
||||
Fp8GemmSmem<Traits, LayoutA, LayoutB, true, false>::kBytes, grid,
|
||||
dim3(Traits::kCtaThreads), stream, p);
|
||||
else
|
||||
launch_with_smem<fp8_gemm_kernel<Traits, LayoutA, LayoutB,
|
||||
GroupRaster, false, false>>(
|
||||
Fp8GemmSmem<Traits, LayoutA, LayoutB, false, false>::kBytes, grid,
|
||||
dim3(Traits::kCtaThreads), stream, p);
|
||||
}
|
||||
|
||||
} // namespace fp8
|
||||
|
||||
@@ -0,0 +1,139 @@
|
||||
/*
|
||||
FP8 GEMM config sweep — pure C, no torch. Times (BM, BN, WarpM, WarpN, kK,
|
||||
Stages, raster) tile configurations across the production square shapes so
|
||||
the launcher's shape dispatch table is grounded in measurements.
|
||||
|
||||
nvcc -I csrc -arch=sm_89 -std=c++17 -O3 --use_fast_math \
|
||||
csrc/tests/fp8_sweep.cu -o /tmp/fp8_sweep && /tmp/fp8_sweep [iters] [sizes...]
|
||||
*/
|
||||
|
||||
#include "test_utils.cuh"
|
||||
|
||||
#include <cuda_fp8.h>
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cuda_runtime.h>
|
||||
#include <vector>
|
||||
|
||||
#include "../kernels/fp8/gemm.cuh"
|
||||
|
||||
using namespace astrai::fp8;
|
||||
|
||||
namespace {
|
||||
|
||||
struct BenchData {
|
||||
__nv_fp8_e4m3 *da, *db;
|
||||
__nv_bfloat16* dout;
|
||||
float* dscale;
|
||||
};
|
||||
|
||||
template <int BM, int BN, int WM, int WN, int kK, int Stages, bool GroupRaster,
|
||||
bool LeanRing = false>
|
||||
float bench_config(BenchData& d, int m, int n, int k, int iters) {
|
||||
FP8Params p = {};
|
||||
p.a_ptr = d.da;
|
||||
p.b_ptr = d.db;
|
||||
p.out_ptr = d.dout;
|
||||
p.scale = d.dscale;
|
||||
p.m = m;
|
||||
p.n = n;
|
||||
p.k = k;
|
||||
p.a_ld = k;
|
||||
p.b_ld = k;
|
||||
|
||||
using Traits = Fp8GemmTraits<FP8Format::E4M3, BM, BN, kK, Stages, WM, WN>;
|
||||
using Smem = Fp8GemmSmem<Traits, RowMajor, ColMajor, false, LeanRing>;
|
||||
dim3 grid((n + BN - 1) / BN, (m + BM - 1) / BM);
|
||||
dim3 block(Traits::kCtaThreads);
|
||||
const int smem = Smem::kBytes;
|
||||
|
||||
auto launch = [&] {
|
||||
launch_with_smem<fp8_gemm_kernel<Traits, RowMajor, ColMajor, GroupRaster,
|
||||
false, LeanRing>>(
|
||||
smem, grid, block, 0, p);
|
||||
};
|
||||
launch();
|
||||
CUDA_CHECK(cudaGetLastError());
|
||||
CUDA_CHECK(cudaDeviceSynchronize());
|
||||
|
||||
cudaEvent_t start, end;
|
||||
cudaEventCreate(&start);
|
||||
cudaEventCreate(&end);
|
||||
for (int i = 0; i < 3; ++i) launch();
|
||||
cudaDeviceSynchronize();
|
||||
cudaEventRecord(start);
|
||||
for (int i = 0; i < iters; ++i) launch();
|
||||
cudaEventRecord(end);
|
||||
cudaEventSynchronize(end);
|
||||
float ms = 0;
|
||||
cudaEventElapsedTime(&ms, start, end);
|
||||
cudaEventDestroy(start);
|
||||
cudaEventDestroy(end);
|
||||
return ms / iters;
|
||||
}
|
||||
|
||||
// One named config column.
|
||||
struct Col {
|
||||
const char* name;
|
||||
float (*fn)(BenchData&, int, int, int, int);
|
||||
};
|
||||
|
||||
template <int BM, int BN, int WM, int WN, int kK, int Stages, bool R,
|
||||
bool Lean = false>
|
||||
float run(BenchData& d, int m, int n, int k, int iters) {
|
||||
return bench_config<BM, BN, WM, WN, kK, Stages, R, Lean>(d, m, n, k, iters);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
const int iters = argc > 1 ? atoi(argv[1]) : 50;
|
||||
int sizes[] = {512, 1024, 2048, 4096, 8192, 0, 0, 0};
|
||||
for (int i = 2; i < argc && i < 10; ++i) sizes[i - 2] = atoi(argv[i]);
|
||||
|
||||
Col cols[] = {
|
||||
{"128ring3", &run<128, 128, 64, 32, 64, 2, true>},
|
||||
{"128ring3s3", &run<128, 128, 64, 32, 64, 3, true>},
|
||||
{"128lean", &run<128, 128, 64, 32, 64, 2, true, true>},
|
||||
{"64r5s4", &run<64, 64, 32, 32, 64, 4, false>},
|
||||
{"64lean4", &run<64, 64, 32, 32, 64, 4, false, true>},
|
||||
{"64lean5", &run<64, 64, 32, 32, 64, 5, false, true>},
|
||||
{"64lean3", &run<64, 64, 32, 32, 64, 3, false, true>},
|
||||
};
|
||||
const int ncols = sizeof(cols) / sizeof(cols[0]);
|
||||
|
||||
printf("%6s |", "shape");
|
||||
for (auto& c : cols) printf(" %11s |", c.name);
|
||||
printf("\n");
|
||||
|
||||
for (int s : sizes) {
|
||||
if (s <= 0) continue;
|
||||
const int m = s, n = s, k = s;
|
||||
BenchData d;
|
||||
std::vector<__nv_fp8_e4m3> a((size_t)m * k), b((size_t)n * k);
|
||||
for (auto& v : a) v = __nv_fp8_e4m3(randf());
|
||||
for (auto& v : b) v = __nv_fp8_e4m3(randf());
|
||||
cudaMalloc(&d.da, a.size());
|
||||
cudaMalloc(&d.db, b.size());
|
||||
cudaMalloc(&d.dout, (size_t)m * n * 2);
|
||||
cudaMalloc(&d.dscale, 4);
|
||||
const float one = 1.0f;
|
||||
cudaMemcpy(d.dscale, &one, 4, cudaMemcpyHostToDevice);
|
||||
cudaMemcpy(d.da, a.data(), a.size(), cudaMemcpyHostToDevice);
|
||||
cudaMemcpy(d.db, b.data(), b.size(), cudaMemcpyHostToDevice);
|
||||
|
||||
const double flops = 2.0 * m * n * k;
|
||||
printf("%6d |", s);
|
||||
for (auto& c : cols) {
|
||||
const float ms = c.fn(d, m, n, k, iters);
|
||||
printf(" %5.2fus %4.1fT |", ms * 1000,
|
||||
flops / (ms * 1e-3) / 1e12);
|
||||
}
|
||||
printf("\n");
|
||||
cudaFree(d.da);
|
||||
cudaFree(d.db);
|
||||
cudaFree(d.dout);
|
||||
cudaFree(d.dscale);
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
Reference in New Issue
Block a user