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

refactor: namespace csrc kernels and extract common helpers

Browse Source 
- attention family -> astrai::attention; fp8 family -> astrai::fp8
- new common/reduce.cuh (warp/group reductions, atomic_max_float)
- new common/cp_async.cuh (predicated cp_async_16, commit/wait group)
- move MAX_SPLITS into attention/common.h; delete warp_utils.cuh
- .cu bindings and pure C tests open family namespaces via using
This commit is contained in:
ViperEkura
2026-08-24 14:49:55 +08:00
parent 34471252ab
commit 31ca357c61
21 changed files with 243 additions and 139 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/kernels/fp8/common.h
+9
View File
@@ -5,6 +5,12 @@
#include <cuda_runtime.h>
#include <cstdint>
// Pure POD/traits header — no .cuh/CUDA-kernel includes; raw __nv_* type
// spellings only.
namespace astrai {
namespace fp8 {
// Compile-time FP8 format: E4M3 (forward / high precision, max 448) or
// E5M2 (gradient / large dynamic range, max 57344).
enum class FP8Format : int {
@@ -64,3 +70,6 @@ struct FP8Params {
int total;
};
} // namespace fp8
} // namespace astrai
csrc/kernels/fp8/gemm.cuh
+11 -59
View File
@@ -10,8 +10,11 @@
#include <type_traits>
#include "common.h"
#include "../common/cp_async.cuh"
#include "../common/mma.cuh"
#include "../common/reduce.cuh"
namespace astrai {
namespace fp8 {
// m16n8k32 (see astrai::mma_shape<fp8 type>::k in common/mma.cuh)
@@ -35,61 +38,9 @@ struct fp8_input<FP8Format::E5M2> {
// FP8 MMA lives in the shared astrai::mma_sync template (common/mma.cuh);
// instantiate it with fp8_input<Fmt>::type. Accumulates in-place: callers
// pass the same accumulator array as both `d` and `c`.
__device__ __forceinline__ void atomic_max_float(float* destination,
float value) {
if (destination)
atomicMax(reinterpret_cast<unsigned*>(destination),
__float_as_uint(value));
}
__device__ __forceinline__ float warp_reduce_max(float value) {
#pragma unroll
for (int offset = 16; offset; offset >>= 1) {
value = fmaxf(value, __shfl_xor_sync(0xffffffffu, value, offset));
}
return value;
}
// One thread moves sixteen FP8 values (16 bytes) via cp.async.
template <typename T>
__device__ __forceinline__ void cp_async_16b(T* destination,
const T* source, bool valid) {
const unsigned shared_address = __cvta_generic_to_shared(destination);
const uint4* source_vec = reinterpret_cast<const uint4*>(source);
asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;"
:: "r"(shared_address), "l"(source_vec),
"r"(valid ? 16 : 0));
}
// PTX requires wait_group's operand to be an immediate value. Keep it as a
// template argument so the stage policy remains compile-time configurable.
template <int KeepGroups>
__device__ __forceinline__ void cp_async_wait_group() {
static_assert(KeepGroups >= 0 && KeepGroups <= 7,
"cp.async.wait_group supports immediates in [0, 7]");
asm volatile("cp.async.wait_group %0;" :: "n"(KeepGroups));
}
template <int MaxKeepGroups>
__device__ __forceinline__ void cp_async_wait_group_dispatch(int keep_groups) {
static_assert(MaxKeepGroups >= 0 && MaxKeepGroups <= 7,
"cp.async.wait_group supports immediates in [0, 7]");
if (keep_groups == MaxKeepGroups) {
cp_async_wait_group<MaxKeepGroups>();
} else if constexpr (MaxKeepGroups > 0) {
cp_async_wait_group_dispatch<MaxKeepGroups - 1>(keep_groups);
} else {
cp_async_wait_group<0>();
}
}
template <int Stages>
__device__ __forceinline__ void cp_async_commit_group() {
static_assert(Stages >= 1 && Stages <= 8,
"FP8 GEMM stages must be in the range [1, 8]");
asm volatile("cp.async.commit_group;");
}
// warp_reduce_max / atomic_max_float (quantize amax) live in
// common/reduce.cuh; the cp.async pipeline primitives (predicated 16-byte
// copy, commit_group, wait_group + runtime dispatch) in common/cp_async.cuh.
// ---------------------------------------------------------------------------
// Quantize kernel: BF16 -> FP8 (E4M3 or E5M2), fused amax over raw values.
@@ -252,7 +203,7 @@ __device__ __forceinline__ void load_operand_tile(
const bool full = k_base + c + 15 < contract;
if (row < rows && full &&
(reinterpret_cast<uintptr_t>(src) & 15) == 0) {
cp_async_16b(dst, src, true);
astrai::cp_async_16(dst, src, true);
} else {
#pragma unroll
for (int i = 0; i < 16; ++i)
@@ -381,7 +332,7 @@ __global__ void __launch_bounds__(kWarps * 32, 2) fp8_gemm_kernel(FP8Params p) {
for (int stage = 0; stage < kStages; ++stage) {
if (stage < tile_count) {
load_tile(stage, static_cast<int64_t>(stage) * kK);
cp_async_commit_group<kStages>();
astrai::cp_async_commit_group();
}
}
@@ -393,7 +344,7 @@ __global__ void __launch_bounds__(kWarps * 32, 2) fp8_gemm_kernel(FP8Params p) {
// oldest group (the current stage) ready for consumption.
const int keep_groups =
remaining < kStages - 1 ? static_cast<int>(remaining) : kStages - 1;
cp_async_wait_group_dispatch<kStages - 1>(keep_groups);
astrai::cp_async_wait_group_dispatch<kStages - 1>(keep_groups);
// Barrier 1: every thread's cp.async for this stage is complete
// before any thread reads tiles written by other threads.
__syncthreads();
@@ -438,7 +389,7 @@ __global__ void __launch_bounds__(kWarps * 32, 2) fp8_gemm_kernel(FP8Params p) {
__syncthreads();
if (tile_index + kStages < tile_count) {
load_tile(stage, (tile_index + kStages) * kK);
cp_async_commit_group<kStages>();
astrai::cp_async_commit_group();
}
}
@@ -517,3 +468,4 @@ void launch_fp8_gemm(const FP8Params& p, cudaStream_t stream) {
}
} // namespace fp8
} // namespace astrai
csrc/kernels/fp8/ops.cu
+14 -11
View File
@@ -14,10 +14,13 @@
#include "gemm.cuh"
#include "../common/device.cuh"
using namespace astrai::fp8;
namespace {
// FP8Format / FP8Params live in the global namespace (common.h); the
// launchers live in fp8:: (gemm.cuh).
// FP8Format / FP8Params and the launchers live in astrai::fp8 (common.h /
// gemm.cuh); this TU opens the using-directive above so the binding reads
// them unqualified.
void check_fp8_device(const torch::Tensor& tensor) {
static std::mutex mutex;
@@ -102,7 +105,7 @@ void launch_gemm_variant(const FP8Params& p, cudaStream_t stream) {
constexpr bool out_fp8 = (Variant & 4) != 0;
constexpr bool trans_a = (Variant & 2) != 0;
constexpr bool trans_b = (Variant & 1) != 0;
fp8::launch_fp8_gemm<Fmt, out_fp8, trans_a, trans_b>(p, stream);
launch_fp8_gemm<Fmt, out_fp8, trans_a, trans_b>(p, stream);
}
template <FP8Format Fmt>
@@ -151,9 +154,9 @@ std::tuple<torch::Tensor, torch::Tensor> quantize_bf16(torch::Tensor x,
pack_quantize_params(p, x_c.data_ptr(), x8.data_ptr(), scale, &amax,
x_c.numel());
if (fmt) {
fp8::launch_fp8_quantize<FP8Format::E5M2>(p, stream.stream());
launch_fp8_quantize<FP8Format::E5M2>(p, stream.stream());
} else {
fp8::launch_fp8_quantize<FP8Format::E4M3>(p, stream.stream());
launch_fp8_quantize<FP8Format::E4M3>(p, stream.stream());
}
C10_CUDA_CHECK(cudaGetLastError());
return {x8, amax};
@@ -259,9 +262,9 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> linear_forward_fp8(
pack_quantize_params(qp, src.data_ptr(), dst.data_ptr(), scale, amax,
src.numel());
if (fmt) {
fp8::launch_fp8_quantize<FP8Format::E5M2>(qp, stream.stream());
launch_fp8_quantize<FP8Format::E5M2>(qp, stream.stream());
} else {
fp8::launch_fp8_quantize<FP8Format::E4M3>(qp, stream.stream());
launch_fp8_quantize<FP8Format::E4M3>(qp, stream.stream());
}
};
quantize(x_c, x8, sx, &amax_x);
@@ -273,10 +276,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) {
fp8::launch_fp8_gemm<FP8Format::E5M2, false, false, true>(
launch_fp8_gemm<FP8Format::E5M2, false, false, true>(
p, stream.stream());
} else {
fp8::launch_fp8_gemm<FP8Format::E4M3, false, false, true>(
launch_fp8_gemm<FP8Format::E4M3, false, false, true>(
p, stream.stream());
}
C10_CUDA_CHECK(cudaGetLastError());
@@ -327,9 +330,9 @@ linear_backward_fp8(torch::Tensor g, torch::Tensor x, torch::Tensor w,
pack_quantize_params(qp, src.data_ptr(), dst.data_ptr(), scale, amax,
src.numel());
if (fmt) {
fp8::launch_fp8_quantize<FP8Format::E5M2>(qp, stream.stream());
launch_fp8_quantize<FP8Format::E5M2>(qp, stream.stream());
} else {
fp8::launch_fp8_quantize<FP8Format::E4M3>(qp, stream.stream());
launch_fp8_quantize<FP8Format::E4M3>(qp, stream.stream());
}
};
// Four-layout backward: the gradient and activation tensors keep their