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

refactor: harden param PODs and CUTLASS-style fp8 layout tags

Browse Source 
- 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>
This commit is contained in:
ViperEkura
2026-08-24 15:27:52 +08:00
parent f6db546578
commit d5067af064
5 changed files with 147 additions and 95 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/kernels/common/mma.cuh
+27 -5
View File
@@ -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,9 +62,9 @@ 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],
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<InT>::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 <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