refactor: reorganize CUDA kernels into per-family directories

- move attention kernels to csrc/kernels/attention/ and rotary to rotary/
- add shared common/mma.cuh (mma_sync, ldmatrix) and device.cuh (sm checks)
- split fp8_mm into three-layer fp8/common.h, gemm.cuh, mm.cu
- fix fused FP8 GEMM ldmatrix lane indexing to fix OOB shared reads
- update extension ops, loader, and kernel tests
This commit is contained in:
2026-08-22 20:40:31 +08:00
parent cb21af38ba
commit 16a55bb474
30 changed files with 1956 additions and 1235 deletions
+23
View File
@@ -0,0 +1,23 @@
// Pure-CUDA device helpers shared across kernel families (no torch).
//
// Family-local headers under kernels/<family>/ own their POD params and
// strategy traits; anything cross-cutting (compute-capability checks, device
// constants) lives here.
#pragma once
namespace astrai {
// Compute-capability comparison: is the device at least (major, minor)?
inline bool sm_at_least(int device_major, int device_minor, int major,
int minor) {
return device_major > major ||
(device_major == major && device_minor >= minor);
}
// FP8 tensor-core MMA (`mma.sync.aligned.m16n8k32` with fp8 inputs) exists on
// Ada (sm_89) and Hopper (sm_90+); sm_80 has no fp8 instructions.
inline constexpr int kMinSmForFp8Major = 8;
inline constexpr int kMinSmForFp8Minor = 9;
} // namespace astrai