refactor: root kernel includes at csrc/kernels

- add the kernels directory to CMake target include paths and drop all ../-relative includes in kernel sources
- reference shared primitives as common/*.cuh and the fp8 type header as fp8/common.h
- update standalone test nvcc commands in file headers and cuda_kernels.md to -I csrc/kernels
This commit is contained in:
2026-09-01 14:43:38 +08:00
parent 63f23a4454
commit 1c17e80882
15 changed files with 23 additions and 22 deletions
+1
View File
@@ -90,6 +90,7 @@ foreach(i RANGE ${_kernel_last})
target_compile_definitions(${name} PRIVATE TORCH_EXTENSION_NAME=${name})
target_include_directories(${name} PRIVATE
"${CMAKE_CURRENT_SOURCE_DIR}/kernels"
"${TORCH_HOME}/include"
"${TORCH_HOME}/include/torch/csrc/api/include"
"${PYTHON_INCLUDE_DIR}")
+1 -1
View File
@@ -2,8 +2,8 @@
#include <cuda_bf16.h>
#include <float.h>
#include "common.h"
#include "common/reduce.cuh"
#include "layout_policies.cuh"
#include "../common/reduce.cuh"
namespace astrai {
namespace attention {
+2 -2
View File
@@ -3,8 +3,8 @@
#include <cuda_fp16.h>
#include <cuda_runtime.h>
#include "../common/cp_async.cuh"
#include "../common/mma.cuh"
#include "common/cp_async.cuh"
#include "common/mma.cuh"
// Predicated cp.async (4-operand form) requires CUDA 11.2+.
// bf16 mma.sync requires sm_80+ (guarded at build time by ASTRAI_NO_MMA).
+1 -1
View File
@@ -2,8 +2,8 @@
#include <cfloat>
#include <cuda_bf16.h>
#include "common.h"
#include "common/reduce.cuh"
#include "layout_policies.cuh"
#include "../common/reduce.cuh"
namespace astrai {
namespace attention {
+1 -1
View File
@@ -11,8 +11,8 @@
#include <cuda_runtime.h>
#include <type_traits>
#include "../common/cp_async.cuh"
#include "common.h"
#include "common/cp_async.cuh"
#include "gemm/epilogue.cuh"
#include "gemm/load.cuh"
#include "gemm/mainloop.cuh"
+1 -1
View File
@@ -2,7 +2,7 @@
// Collective epilogue: fused bias, the bf16 scatter of the fp32 accumulators
// through the reclaimed operand shared memory, and the coalesced copy-out.
#include "../common.h"
#include "fp8/common.h"
#include "policy.cuh"
namespace astrai {
+2 -2
View File
@@ -5,8 +5,8 @@
// The staging invariants and the swizzle derivation live in
// docs/developer/cuda_kernels.md.
#include "../../common/cp_async.cuh"
#include "../common.h"
#include "common/cp_async.cuh"
#include "fp8/common.h"
#include "policy.cuh"
namespace astrai {
+2 -2
View File
@@ -7,8 +7,8 @@
#include <type_traits>
#include "../../common/mma.cuh"
#include "../common.h"
#include "common/mma.cuh"
#include "fp8/common.h"
#include "load.cuh"
#include "policy.cuh"
+1 -1
View File
@@ -5,7 +5,7 @@
#include <type_traits>
#include "../common.h"
#include "fp8/common.h"
namespace astrai {
namespace fp8 {
+1 -1
View File
@@ -8,7 +8,7 @@
#include <mutex>
#include <unordered_map>
#include "../common/device.cuh"
#include "common/device.cuh"
#include "gemm.cuh"
#include "quantize.cuh"
+1 -1
View File
@@ -10,7 +10,7 @@
#include <cstdint>
#include "common.h"
#include "../common/reduce.cuh"
#include "common/reduce.cuh"
namespace astrai {
namespace fp8 {
+2 -2
View File
@@ -1,5 +1,5 @@
// Compile:
// nvcc -I csrc -arch=sm_89 -O3 --use_fast_math --ptxas-options=-O3 \
// nvcc -I csrc/kernels -arch=sm_89 -O3 --use_fast_math --ptxas-options=-O3 \
// --extra-device-vectorization -Xcompiler -fopenmp \
// csrc/tests/attn_paged_test.cu \
// -o /tmp/test_paged && /tmp/test_paged
@@ -7,7 +7,7 @@
#include <cstring>
#include <vector>
#include "test_utils.cuh"
#include "../kernels/attention/dispatchers.cuh"
#include "attention/dispatchers.cuh"
using namespace astrai::attention;
+2 -2
View File
@@ -1,13 +1,13 @@
/*
Pure-C test — uses shared dispatcher. Combines the decode (split-KV) and
prefill (split-Q) correctness checks + benchmarks into one binary.
nvcc -I csrc -arch=sm_89 -O3 \
nvcc -I csrc/kernels -arch=sm_89 -O3 \
--use_fast_math --ptxas-options=-O3 --extra-device-vectorization \
-Xcompiler -fopenmp csrc/tests/attn_test.cu -o test && ./test
*/
#include "test_utils.cuh"
#include "../kernels/attention/dispatchers.cuh"
#include "attention/dispatchers.cuh"
using namespace astrai::attention;
+3 -3
View File
@@ -6,7 +6,7 @@ Part 1 exercises one bf16 -> fp8 -> mma.sync m16n8k32 instruction pair
Part 2 checks launch_fp8_gemm across all four operand layouts, both K
tiles, and ragged shapes against an fp32 CPU reference.
nvcc -I csrc -arch=sm_89 -std=c++17 -O3 csrc/tests/fp8_test.cu -o /tmp/fp8_test \
nvcc -I csrc/kernels -arch=sm_89 -std=c++17 -O3 csrc/tests/fp8_test.cu -o /tmp/fp8_test \
&& /tmp/fp8_test
*/
@@ -21,8 +21,8 @@ nvcc -I csrc -arch=sm_89 -std=c++17 -O3 csrc/tests/fp8_test.cu -o /tmp/fp8_test
#include <type_traits>
#include <vector>
#include "../kernels/common/mma.cuh"
#include "../kernels/fp8/gemm.cuh"
#include "common/mma.cuh"
#include "fp8/gemm.cuh"
using namespace astrai::fp8;