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 <Fmt, false, RowMajor, ColMajor> instead of <Fmt, false, false, true>

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<T> 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<typename T, typename AT = float>
 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 causal_offset = -1;
-    int use_mask;
+    int use_mask = 0;
 
     // pointers
-    const T* __restrict__ q_ptr;
+    const T* __restrict__ q_ptr = nullptr;
-    const T* __restrict__ k_ptr;
+    const T* __restrict__ k_ptr = nullptr;
-    const T* __restrict__ v_ptr;
+    const T* __restrict__ v_ptr = nullptr;
-    const T* __restrict__ new_k_ptr;
+    const T* __restrict__ new_k_ptr = nullptr;
-    const T* __restrict__ new_v_ptr;
+    const T* __restrict__ new_v_ptr = nullptr;
-    T* __restrict__ o_ptr;
+    T* __restrict__ o_ptr = nullptr;
-    const bool* __restrict__ mask;
+    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_to_token = nullptr;      // [num_reqs, max_context_len]
-    const int* __restrict__ req_pool_indices;     // [batch]
+    const int* __restrict__ req_pool_indices = nullptr;  // [batch]
-    const int* __restrict__ kv_indptr;             // [batch + 1]
+    const int* __restrict__ kv_indptr = nullptr;         // [batch + 1]
-    const int* __restrict__ qo_indptr;             // [batch + 1] or nullptr for decode
+    const int* __restrict__ qo_indptr = nullptr;         // [batch + 1] or nullptr for decode
-    const int* __restrict__ q_tile_to_batch;        // [num_q_tiles], prefill only
+    const int* __restrict__ q_tile_to_batch = nullptr;   // [num_q_tiles], prefill only
-    const int* __restrict__ q_tile_to_index;        // [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__ o_part = nullptr;
-    AT* __restrict__ ml_part;
+    AT* __restrict__ ml_part = nullptr;
 };
 
 }  // namespace attention

csrc/kernels/common/mma.cuh

@@ -15,6 +15,9 @@
 #include <cuda_runtime.h>
 #include <type_traits>
 
+
+#define DEVICE_FORCEINLINE static __device__ __forceinline__
+
 namespace astrai {
 
 // Compute capability of the current compilation pass: 0 in the host pass,
@@ -59,7 +62,7 @@ struct mma_shape<__nv_fp8_e5m2> {
 // below `mma_shape<InT>::min_arch` is a **compile error** — the instruction
 // does not exist there, and a silent no-op would produce wrong results.
 template <typename InT>
-__device__ __forceinline__ void mma_sync(float d[4], const unsigned a[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 ||
@@ -115,7 +118,7 @@ __device__ __forceinline__ void mma_sync(float d[4], const unsigned a[4],
 // ---------------------------------------------------------------------------
 
 template <typename T, bool Trans = false>
-__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 <typename T>
-__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

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 <typename Layout>
+struct transpose_layout;
+template <>
+struct transpose_layout<RowMajor> {
+    using type = ColMajor;
+};
+template <>
+struct transpose_layout<ColMajor> {
+    using type = RowMajor;
+};
+template <typename Layout>
+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
@@ -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__ a_ptr = nullptr;
-    const void* __restrict__ b_ptr;
+    const void* __restrict__ b_ptr = nullptr;
-    const float* __restrict__ scale_a;
+    const float* __restrict__ scale_a = nullptr;
-    const float* __restrict__ scale_b;
+    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;
+    void* __restrict__ out_ptr = nullptr;
-    const float* __restrict__ out_scale;
+    const float* __restrict__ out_scale = nullptr;
 
     // Fused forward extras: bias (may be null) and amax slots (may be null).
-    const __nv_bfloat16* __restrict__ bias;
+    const __nv_bfloat16* __restrict__ bias = nullptr;
-    float* __restrict__ amax_a;
+    float* __restrict__ amax_a = nullptr;
-    float* __restrict__ amax_b;
+    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;

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,
+// layout: RowMajor (stored [rows][contract]) copies 16-byte K-contiguous runs
-// while the transposed case reads 16-byte runs along the operand's contiguous
+// with cp.async, while ColMajor (stored [contract][rows]) reads 16-byte runs
-// (non-contract) dim and scatters them across the tile's rows. `block_row` is
+// along the operand's contiguous non-contract dim and scatters them across
-// this block's origin in the operand's row dim; the caller restricts which
+// the tile's rows. `block_row` is this block's origin in the operand's row
-// threads invoke it (all threads for A, the first 128 for B).
+// dim; the caller restricts which threads invoke it (all threads for A, the
-template <typename T8, int K, bool Trans>
+// first 128 for B).
+template <typename T8, int K, typename Layout>
 __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<Layout, ColMajor>) {
         // 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 <typename T8, int kK>
 __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(
+    return __cvta_generic_to_shared(tile_at<kK>(tile, row, chunk << 4));
-        tile + row * kK + (((chunk ^ ((row >> 2) & 1)) << 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)
+//   A ColMajor: tileA[m][p] = a[p*a_ld + m];  A RowMajor: a[m*a_ld + p]
-//           else        a[m*a_ld + p]   (A stored [M][K])
+//   B RowMajor: tileB[n][p] = b[p*b_ld + n];  B ColMajor: b[n*b_ld + p]
-//   TransB: tileB[n][p] = b[n*b_ld + p]  (B stored [N][K])
+template <typename Traits, bool OutFp8 = false, typename LayoutA = RowMajor, typename LayoutB = RowMajor>
-//           else        b[p*b_ld + n]   (B stored [K][N], read transposed)
-template <typename Traits, bool OutFp8 = false, bool TransA = false,
-          bool TransB = false>
 __global__ void __launch_bounds__(kWarps * 32, 2) fp8_gemm_kernel(FP8Params p) {
     using T8 = std::conditional_t<Traits::kIsE5M2, __nv_fp8_e5m2, __nv_fp8_e4m3>;
     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.
+    // MMA fragment reads below stay unchanged across the four layout
-    // Each 128x32 tile is 256 16B chunks: one per thread.
+    // 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
+        load_operand_tile<T8, kK, LayoutA>(
-        // 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<T8, kK, TransA>(
             a_smem[stage], a, m, k, a_ld, tid, k_base, blockIdx.y * kBlockM);
-        load_operand_tile<T8, kK, !TransB>(
+        load_operand_tile<T8, kK, transpose_layout_t<LayoutB>>(
             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<T8, kK>(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<T8, kK>(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<T8, kK>(
                                 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 <FP8Format Fmt, bool OutFp8 = false, bool TransA = false,
+template <FP8Format Fmt, bool OutFp8 = false, typename LayoutA = RowMajor,
-          bool TransB = false, int Stages = 3>
+          typename LayoutB = RowMajor, int Stages = 3>
 void launch_fp8_gemm(const FP8Params& p, cudaStream_t stream) {
     using Traits = Fp8GemmTraits<Fmt, 128, 128, 32, Stages>;
     dim3 grid((p.n + Traits::kBlockN - 1) / Traits::kBlockN,
               (p.m + Traits::kBlockM - 1) / Traits::kBlockM);
-    fp8_gemm_kernel<Traits, OutFp8, TransA, TransB><<<grid, kWarps * 32, 0, stream>>>(p);
+    fp8_gemm_kernel<Traits, OutFp8, LayoutA, LayoutB><<<grid, kWarps * 32, 0, stream>>>(p);
 }
 
 }  // namespace fp8

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;
+    // Variant bits 1/0 = trans_a/trans_b -> CUTLASS-style layout tags
-    constexpr bool trans_b = (Variant & 1) != 0;
+    // (trans_a ? A ColMajor : RowMajor, same for B; see common.h).
-    launch_fp8_gemm<Fmt, out_fp8, trans_a, trans_b>(p, stream);
+    using LayoutA = std::conditional_t<(Variant & 2) != 0, ColMajor, RowMajor>;
+    using LayoutB = std::conditional_t<(Variant & 1) != 0, ColMajor, RowMajor>;
+    launch_fp8_gemm<Fmt, out_fp8, LayoutA, LayoutB>(p, stream);
 }
 
 template <FP8Format Fmt>
@@ -276,10 +278,10 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> 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<FP8Format::E5M2, false, false, true>(
+        launch_fp8_gemm<FP8Format::E5M2, false, RowMajor, ColMajor>(
             p, stream.stream());
     } else {
-        launch_fp8_gemm<FP8Format::E4M3, false, false, true>(
+        launch_fp8_gemm<FP8Format::E4M3, false, RowMajor, ColMajor>(
             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.