From d5067af064852bbc7c3579f250edd1cc29dc8994 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Mon, 24 Aug 2026 15:27:52 +0800 Subject: [PATCH] refactor: harden param PODs and CUTLASS-style fp8 layout tags - NSDMI null/-1 defaults for AttentionParams/FP8Params pointer+flag members: partially packed structs can no longer hold garbage non-null pointers that gate optional paths (root cause class of the paged test bug); still aggregates, still trivially copyable - move per-lane ldmatrix wrappers (ldsm_x2/x4) from fp8/gemm.cuh to common/mma.cuh as ldmatrix_x2_lane/x4_lane, next to the single-address variants - DEVICE_FORCEINLINE macro in common/mma.cuh (matches layout_policies.cuh, internal linkage) - frag_addr now delegates to tile_at: the swizzle math has one source - operand layouts as CUTLASS-style RowMajor/ColMajor tags threaded from launch_fp8_gemm through the kernel to load_operand_tile; B's operand view via transpose_layout_t; call sites read instead of --- csrc/kernels/attention/common.h | 42 +++++++------ csrc/kernels/common/mma.cuh | 32 ++++++++-- csrc/kernels/fp8/common.h | 52 ++++++++++++---- csrc/kernels/fp8/gemm.cuh | 101 ++++++++++++++------------------ csrc/kernels/fp8/ops.cu | 15 +++-- 5 files changed, 147 insertions(+), 95 deletions(-) diff --git a/csrc/kernels/attention/common.h b/csrc/kernels/attention/common.h index 6fe3ab0..88ae20b 100644 --- a/csrc/kernels/attention/common.h +++ b/csrc/kernels/attention/common.h @@ -24,6 +24,14 @@ constexpr int MAX_SPLITS = 32; // layout_policies.cuh); a given call only touches the fields of one mode, so // this is a POD shared by both paths rather than two parallel structs that // drift out of sync. +// +// Pointer/flag members carry default member initializers: the pointers gate +// optional paths via null checks (new_k_ptr, mask, o_part, ...), so a stack +// `AttentionParams p;` left partially packed must never see garbage +// non-null pointers or a garbage use_mask/causal_offset — that class of bug +// reads through wild addresses. NSDMI keeps the struct an aggregate (C++17) +// and trivially copyable, so `= {}`, memcpy-style packing and by-value kernel +// params all behave exactly as before. template struct AttentionParams { // Shape @@ -37,17 +45,17 @@ struct AttentionParams { // Attention behavior float scale; // -1 = non-causal; >=0 = absolute position of first Q token - int causal_offset; - int use_mask; + int causal_offset = -1; + int use_mask = 0; // pointers - const T* __restrict__ q_ptr; - const T* __restrict__ k_ptr; - const T* __restrict__ v_ptr; - const T* __restrict__ new_k_ptr; - const T* __restrict__ new_v_ptr; - T* __restrict__ o_ptr; - const bool* __restrict__ mask; + const T* __restrict__ q_ptr = nullptr; + const T* __restrict__ k_ptr = nullptr; + const T* __restrict__ v_ptr = nullptr; + const T* __restrict__ new_k_ptr = nullptr; + const T* __restrict__ new_v_ptr = nullptr; + T* __restrict__ o_ptr = nullptr; + const bool* __restrict__ mask = nullptr; // strides int q_b_stride; @@ -68,19 +76,19 @@ struct AttentionParams { int mask_l_stride; // Paged K/V addressing - const int* __restrict__ req_to_token; // [num_reqs, max_context_len] - const int* __restrict__ req_pool_indices; // [batch] - const int* __restrict__ kv_indptr; // [batch + 1] - const int* __restrict__ qo_indptr; // [batch + 1] or nullptr for decode - const int* __restrict__ q_tile_to_batch; // [num_q_tiles], prefill only - const int* __restrict__ q_tile_to_index; // [num_q_tiles], prefill only + const int* __restrict__ req_to_token = nullptr; // [num_reqs, max_context_len] + const int* __restrict__ req_pool_indices = nullptr; // [batch] + const int* __restrict__ kv_indptr = nullptr; // [batch + 1] + const int* __restrict__ qo_indptr = nullptr; // [batch + 1] or nullptr for decode + const int* __restrict__ q_tile_to_batch = nullptr; // [num_q_tiles], prefill only + const int* __restrict__ q_tile_to_index = nullptr; // [num_q_tiles], prefill only int num_q_tiles; int max_context_len; // req_to_token stride (dim 1) // Decode split-KV workspace int num_splits; - AT* __restrict__ o_part; - AT* __restrict__ ml_part; + AT* __restrict__ o_part = nullptr; + AT* __restrict__ ml_part = nullptr; }; } // namespace attention diff --git a/csrc/kernels/common/mma.cuh b/csrc/kernels/common/mma.cuh index afbcfd8..44d631c 100644 --- a/csrc/kernels/common/mma.cuh +++ b/csrc/kernels/common/mma.cuh @@ -15,6 +15,9 @@ #include #include + +#define DEVICE_FORCEINLINE static __device__ __forceinline__ + namespace astrai { // Compute capability of the current compilation pass: 0 in the host pass, @@ -59,9 +62,9 @@ struct mma_shape<__nv_fp8_e5m2> { // below `mma_shape::min_arch` is a **compile error** — the instruction // does not exist there, and a silent no-op would produce wrong results. template -__device__ __forceinline__ void mma_sync(float d[4], const unsigned a[4], - const unsigned b[2], - const float c[4]) { +DEVICE_FORCEINLINE void mma_sync(float d[4], const unsigned a[4], + const unsigned b[2], + const float c[4]) { static_assert(ASTRAI_DEVICE_ARCH == 0 || ASTRAI_DEVICE_ARCH >= mma_shape::min_arch, "mma_sync: this MMA shape requires a newer compute " @@ -115,7 +118,7 @@ __device__ __forceinline__ void mma_sync(float d[4], const unsigned a[4], // --------------------------------------------------------------------------- template -__device__ __forceinline__ void ldmatrix_x2(unsigned r[2], const T* p) { +DEVICE_FORCEINLINE void ldmatrix_x2(unsigned r[2], const T* p) { const unsigned a = __cvta_generic_to_shared(p); if constexpr (Trans) { asm volatile( @@ -132,7 +135,7 @@ __device__ __forceinline__ void ldmatrix_x2(unsigned r[2], const T* p) { // Four matrices at p, p+128, p+256, p+384 bytes (16-byte row stride). template -__device__ __forceinline__ void ldmatrix_x4(unsigned r[4], const T* p) { +DEVICE_FORCEINLINE void ldmatrix_x4(unsigned r[4], const T* p) { const unsigned a = __cvta_generic_to_shared(p); asm volatile( "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];" @@ -140,4 +143,23 @@ __device__ __forceinline__ void ldmatrix_x4(unsigned r[4], const T* p) { : "r"(a)); } +// Per-lane-address variants: the caller supplies a raw shared-memory address +// per lane instead of one common pointer. Use when the fragment tiles are +// XOR-swizzled per 16B chunk so each lane must compute its own row and chunk +// address (see fp8/gemm.cuh's frag_addr + lane selectors for the m16n8k32 +// operand layouts). +DEVICE_FORCEINLINE void ldmatrix_x2_lane(unsigned r[2], + unsigned addr) { + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];" + : "=r"(r[0]), "=r"(r[1]) + : "r"(addr)); +} + +DEVICE_FORCEINLINE void ldmatrix_x4_lane(unsigned r[4], + unsigned addr) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(addr)); +} + } // namespace astrai diff --git a/csrc/kernels/fp8/common.h b/csrc/kernels/fp8/common.h index 85aaabf..bc8ac04 100644 --- a/csrc/kernels/fp8/common.h +++ b/csrc/kernels/fp8/common.h @@ -18,6 +18,34 @@ enum class FP8Format : int { E5M2 = 1, }; +// Operand memory layouts as types (CUTLASS-style tags). The tag names the +// storage order of the raw buffer relative to the operand's canonical GEMM +// matrix — A is [M][K], B is [K][N]: +// A RowMajor = [M][K] storage (K-contiguous rows; the default) +// A ColMajor = [K][M] storage (M-contiguous; A^T) +// B RowMajor = [K][N] storage (N-contiguous; the plain a @ b operand) +// B ColMajor = [N][K] storage (K-contiguous; the nn.Linear weight layout) +// Empty tags: selection happens by type at compile time (see load_operand_tile). +struct RowMajor {}; +struct ColMajor {}; + +// Transpose of a layout tag: the same buffer with the rows and contract dims +// swapped. B's tag is relative to the canonical [K][N] GEMM matrix, so the +// stage-load (which views any operand as [rows][contract]) sees the transposed +// tag — this trait makes that inversion explicit. +template +struct transpose_layout; +template <> +struct transpose_layout { + using type = ColMajor; +}; +template <> +struct transpose_layout { + using type = RowMajor; +}; +template +using transpose_layout_t = typename transpose_layout::type; + // Compile-time tile configuration, mirroring KernelTraits in the attention kernels. `Fmt` selects the FP8 conversion // and the MMA PTX mnemonic; the remaining parameters shape the CTA tile and @@ -38,23 +66,27 @@ struct Fp8GemmTraits { // Unified GEMM parameter POD, mirroring AttentionParams: one struct flows // through quantize / fused / pre-quantized kernels. Each kernel touches only // the fields it needs; buffers are raw pointers packed by the torch binding. +// Pointer members default to null (same NSDMI rationale as AttentionParams: +// bias / amax / out_scale gate optional paths via null checks, so a partially +// packed struct must never hold garbage non-null pointers). Still an +// aggregate, still trivially copyable. struct FP8Params { // Inputs: a/b are BF16 for the fused (quantize-in-GEMM) path, FP8 for // the pre-quantized path. Scales are quantization steps (device scalars). - const void* __restrict__ a_ptr; - const void* __restrict__ b_ptr; - const float* __restrict__ scale_a; - const float* __restrict__ scale_b; + const void* __restrict__ a_ptr = nullptr; + const void* __restrict__ b_ptr = nullptr; + const float* __restrict__ scale_a = nullptr; + const float* __restrict__ scale_b = nullptr; // Output: BF16 or FP8 (E4M3). out_scale is the output quantization step // (FP8 output only). - void* __restrict__ out_ptr; - const float* __restrict__ out_scale; + void* __restrict__ out_ptr = nullptr; + const float* __restrict__ out_scale = nullptr; // Fused forward extras: bias (may be null) and amax slots (may be null). - const __nv_bfloat16* __restrict__ bias; - float* __restrict__ amax_a; - float* __restrict__ amax_b; + const __nv_bfloat16* __restrict__ bias = nullptr; + float* __restrict__ amax_a = nullptr; + float* __restrict__ amax_b = nullptr; // Shapes. total is only used by the elementwise quantize kernel. `int` // covers every realistic LLM shape; the kernels promote to int64 for all @@ -65,7 +97,7 @@ struct FP8Params { // For a non-transposed operand the stride equals the contract dim; for a // transposed operand it is the operand's own column count. The binding // packs these so the kernel reads both buffers either naturally or - // transposed depending on TransA/TransB. + // transposed depending on the LayoutA/LayoutB tags (see gemm.cuh). int a_ld, b_ld; int total; diff --git a/csrc/kernels/fp8/gemm.cuh b/csrc/kernels/fp8/gemm.cuh index d970bae..ede2943 100644 --- a/csrc/kernels/fp8/gemm.cuh +++ b/csrc/kernels/fp8/gemm.cuh @@ -139,17 +139,18 @@ __device__ __forceinline__ T8* tile_at(T8* tile, int row, int col) { // Stage-load one GEMM operand into the canonical flat [rows * K] shared tile // (addressing via tile_at, so stores land in the swizzled layout). The // transpose is folded into the staging step via a CUTLASS-style crosswise -// layout: the congruous case copies 16-byte K-contiguous runs with cp.async, -// while the transposed case reads 16-byte runs along the operand's contiguous -// (non-contract) dim and scatters them across the tile's rows. `block_row` is -// this block's origin in the operand's row dim; the caller restricts which -// threads invoke it (all threads for A, the first 128 for B). -template +// layout: RowMajor (stored [rows][contract]) copies 16-byte K-contiguous runs +// with cp.async, while ColMajor (stored [contract][rows]) reads 16-byte runs +// along the operand's contiguous non-contract dim and scatters them across +// the tile's rows. `block_row` is this block's origin in the operand's row +// dim; the caller restricts which threads invoke it (all threads for A, the +// first 128 for B). +template __device__ __forceinline__ void load_operand_tile( T8* tile, const T8* __restrict__ operand, int64_t rows, int64_t contract, int64_t ld, int tid, int64_t k_base, int64_t block_row) { - if constexpr (Trans) { + if constexpr (std::is_same_v) { // Operand stored [contract][rows]: contiguous along the non-contract dim. const int rg = tid >> 5; // Rows / 16 row-groups const int kl = tid & 31; // K lanes @@ -220,50 +221,27 @@ __device__ __forceinline__ void load_operand_tile( // in-kernel transpose of the operands (the binding handles transposes). // --------------------------------------------------------------------------- -// ldmatrix with per-lane addresses (unlike common/mma.cuh's single-address -// helpers, the fragment tiles here are XOR-swizzled per 16B chunk, so each -// lane computes its own row/chunk address). Layout contract for fp8 -// m16n8k32 (values packed two-per-b16 slot, K-contiguous rows): -// x4 (A fragment): lane i points at tile row (i>>3 & 1)*8 + (i&7) of -// chunk (k_seg*2 + (i>>4)); reg j = matrix j = [row g][tig*4..+3] in -// the order (rows 0-7 c, rows 8-15 c, rows 0-7 c+1, rows 8-15 c+1) — -// exactly the mma.sync A operand layout. -// x2 (B fragment): lane i points at tile row (i&7) of chunk -// (k_seg*2 + ((i>>3) & 1)); reg j = [row(n) g][tig*4..+3] chunk c/c+1 -// — exactly the mma.sync B operand layout (col operand, K-contiguous). -__device__ __forceinline__ void ldsm_x2(unsigned r[2], unsigned addr) { - asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];" - : "=r"(r[0]), "=r"(r[1]) - : "r"(addr)); -} - -__device__ __forceinline__ void ldsm_x4(unsigned r[4], unsigned addr) { - asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];" - : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) - : "r"(addr)); -} - // Swizzled 16B-chunk address (tile_at's layout) as a raw shared-memory -// pointer for ldmatrix. Requires kK == 32 (2 chunks/row swizzle). +// pointer for ldmatrix. Requires kK == 32 (2 chunks/row swizzle). The chunk +// XOR itself lives only in tile_at; this wrapper just converts the element +// address it returns. template __device__ __forceinline__ unsigned frag_addr(const T8* tile, int row, int chunk) { static_assert(kK == 32, "fragment swizzle offsets assume kK == 32"); - return __cvta_generic_to_shared( - tile + row * kK + (((chunk ^ ((row >> 2) & 1)) << 4))); + return __cvta_generic_to_shared(tile_at(tile, row, chunk << 4)); } -// TransA / TransB select the operand memory layout. The kernel always computes +// LayoutA / LayoutB tag the operands' storage (CUTLASS-style, see common.h): +// A RowMajor = [M][K] / ColMajor = [K][M]; B RowMajor = [K][N] / +// ColMajor = [N][K]. The kernel always computes // out[m][n] = sum_p tileA[m][p] * tileB[n][p] // with the tiles materialized in the canonical [M][kK] / [N][kK] layout, so the -// MMA fragments are read identically regardless of layout. The two flags only +// MMA fragments are read identically regardless of layout. The tags only // change how the stage-load gathers the operand from global memory: -// TransA: tileA[m][p] = a[p*a_ld + m] (A stored [K][M], i.e. A^T) -// else a[m*a_ld + p] (A stored [M][K]) -// TransB: tileB[n][p] = b[n*b_ld + p] (B stored [N][K]) -// else b[p*b_ld + n] (B stored [K][N], read transposed) -template +// A ColMajor: tileA[m][p] = a[p*a_ld + m]; A RowMajor: a[m*a_ld + p] +// B RowMajor: tileB[n][p] = b[p*b_ld + n]; B ColMajor: b[n*b_ld + p] +template __global__ void __launch_bounds__(kWarps * 32, 2) fp8_gemm_kernel(FP8Params p) { using T8 = std::conditional_t; constexpr int kBlockM = Traits::kBlockM; @@ -306,22 +284,31 @@ __global__ void __launch_bounds__(kWarps * 32, 2) fp8_gemm_kernel(FP8Params p) { // Both operands are staged into the canonical [M][kK] / [N][kK] shared // tiles regardless of their global layout (see load_operand_tile), so the - // MMA fragment reads below stay unchanged across the four layout flags. - // Each 128x32 tile is 256 16B chunks: one per thread. + // MMA fragment reads below stay unchanged across the four layout + // combinations. Each 128x32 tile is 256 16B chunks: one per thread. + // A's tag already names the operand view ([M][K] = [rows][contract]); + // B's tag is relative to the canonical [K][N], so the stage-load sees its + // transpose (transpose_layout_t, see common.h). auto load_tile = [&](int stage, int64_t k_base) { - // load_operand_tile's `Trans` means "the operand's contiguous dim is - // the non-contract dim" (crosswise load). For A that is TransA; for B - // the storage flag is inverted (TransB=true stores B as [N][K], i.e. - // K-contiguous, which is the congruous case). - load_operand_tile( + load_operand_tile( a_smem[stage], a, m, k, a_ld, tid, k_base, blockIdx.y * kBlockM); - load_operand_tile( + load_operand_tile>( b_smem[stage], b, n, k, b_ld, tid, k_base, blockIdx.x * kBlockN); }; const int64_t tile_count = (k + kK - 1) / kK; - // Per-lane ldmatrix row/chunk selectors (see ldsm_x2/ldsm_x4 contract). + // Per-lane ldmatrix row/chunk selectors for common/mma.cuh's + // ldmatrix_*_lane (the fragment tiles are XOR-swizzled per 16B chunk, so + // each lane computes its own row/chunk address). Layout contract for fp8 + // m16n8k32 (values packed two-per-b16 slot, K-contiguous rows): + // x4 (A fragment): lane i points at tile row (i>>3 & 1)*8 + (i&7) of + // chunk (k_seg*2 + (i>>4)); reg j = matrix j = [row g][tig*4..+3] in + // the order (rows 0-7 c, rows 8-15 c, rows 0-7 c+1, rows 8-15 c+1) — + // exactly the mma.sync A operand layout. + // x2 (B fragment): lane i points at tile row (i&7) of chunk + // (k_seg*2 + ((i>>3) & 1)); reg j = [row(n) g][tig*4..+3] chunk c/c+1 + // — exactly the mma.sync B operand layout (col operand, K-contiguous). const int r7 = lane & 7; // row within the 8-row matrix 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) @@ -358,7 +345,7 @@ __global__ void __launch_bounds__(kWarps * 32, 2) fp8_gemm_kernel(FP8Params p) { #pragma unroll for (int nt = 0; nt < 4; ++nt) { const int row = b_row0 + nt * 8 + r7; - ldsm_x2(b_frag[nt], + astrai::ldmatrix_x2_lane(b_frag[nt], frag_addr(b_smem[stage], row, k_seg * 2 + rh8)); } @@ -367,13 +354,13 @@ __global__ void __launch_bounds__(kWarps * 32, 2) fp8_gemm_kernel(FP8Params p) { // 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]; - ldsm_x4(a_frag[0], + astrai::ldmatrix_x4_lane(a_frag[0], frag_addr(a_smem[stage], a_row0 + rh8 * 8 + r7, k_seg * 2 + rh16)); #pragma unroll for (int mt = 0; mt < 4; ++mt) { if (mt < 3) - ldsm_x4(a_frag[mt + 1], + astrai::ldmatrix_x4_lane(a_frag[mt + 1], frag_addr( a_smem[stage], a_row0 + (mt + 1) * 16 + rh8 * 8 + r7, @@ -456,15 +443,15 @@ void launch_fp8_quantize(const FP8Params& p, cudaStream_t stream) { // Pre-quantized GEMM tile config: 128x128 CTA (8 warps x 64x32 warp tiles), // K=32, 3-stage pipeline (24KB smem -> 2 CTAs/SM). The wide warp tile plus // ldmatrix fragments lifts the LSU-issue bound of the old 128x64 config. -// Stages remains an explicit template override for tuning. TransA/TransB +// Stages remains an explicit template override for tuning. LayoutA/LayoutB // mirror the kernel template (defaults keep the NN layout: out = a @ b). -template +template void launch_fp8_gemm(const FP8Params& p, cudaStream_t stream) { using Traits = Fp8GemmTraits; dim3 grid((p.n + Traits::kBlockN - 1) / Traits::kBlockN, (p.m + Traits::kBlockM - 1) / Traits::kBlockM); - fp8_gemm_kernel<<>>(p); + fp8_gemm_kernel<<>>(p); } } // namespace fp8 diff --git a/csrc/kernels/fp8/ops.cu b/csrc/kernels/fp8/ops.cu index c54c587..2f462f1 100644 --- a/csrc/kernels/fp8/ops.cu +++ b/csrc/kernels/fp8/ops.cu @@ -103,9 +103,11 @@ void launch_gemm_variant(const FP8Params& p, cudaStream_t stream) { static_assert(Variant >= 0 && Variant < 8, "invalid FP8 GEMM dispatch variant"); constexpr bool out_fp8 = (Variant & 4) != 0; - constexpr bool trans_a = (Variant & 2) != 0; - constexpr bool trans_b = (Variant & 1) != 0; - launch_fp8_gemm(p, stream); + // Variant bits 1/0 = trans_a/trans_b -> CUTLASS-style layout tags + // (trans_a ? A ColMajor : RowMajor, same for B; see common.h). + using LayoutA = std::conditional_t<(Variant & 2) != 0, ColMajor, RowMajor>; + using LayoutB = std::conditional_t<(Variant & 1) != 0, ColMajor, RowMajor>; + launch_fp8_gemm(p, stream); } template @@ -276,10 +278,10 @@ std::tuple linear_forward_fp8( pack_gemm_params(p, x8.data_ptr(), w8.data_ptr(), out.data_ptr(), sx, sw, nullptr, m, n, k, k, k); if (fmt) { - launch_fp8_gemm( + launch_fp8_gemm( p, stream.stream()); } else { - launch_fp8_gemm( + launch_fp8_gemm( p, stream.stream()); } C10_CUDA_CHECK(cudaGetLastError()); @@ -337,7 +339,8 @@ linear_backward_fp8(torch::Tensor g, torch::Tensor x, torch::Tensor w, }; // Four-layout backward: the gradient and activation tensors keep their // natural row-major layout, and the kernel reads them transposed where the - // GEMM needs it (TransA / TransB). No torch-level `.transpose().contiguous()` + // GEMM needs it (the ColMajor layout tags pick the crosswise stage-load). + // No torch-level `.transpose().contiguous()` // copies are required — dX uses g8 [M,N] as A with w8 [N,K] read transposed // as B; dW uses g8 transposed as A with x8 transposed as B. // g is quantized once (amax_g measured here); both GEMMs share g8.