From 31ca357c61a39c6606d39fc265ccc98ad58a63c3 Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Mon, 24 Aug 2026 14:49:55 +0800
Subject: [PATCH] refactor: namespace csrc kernels and extract common helpers

- 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
---
 csrc/kernels/attention/common.h               | 12 +++-
 csrc/kernels/attention/decode.cu              |  2 +
 csrc/kernels/attention/decode_split_kv.cuh    |  9 ++-
 .../kernels/attention/decode_split_kv_mma.cuh | 19 +++--
 csrc/kernels/attention/dispatchers.cuh        |  7 +-
 csrc/kernels/attention/entry_utils.cuh        | 11 ++-
 csrc/kernels/attention/layout_policies.cuh    |  6 ++
 csrc/kernels/attention/mma_utils.cuh          | 40 +++--------
 csrc/kernels/attention/paged_decode.cu        |  2 +
 csrc/kernels/attention/paged_prefill.cu       |  2 +
 csrc/kernels/attention/prefill.cu             |  2 +
 csrc/kernels/attention/prefill_split_q.cuh    | 16 ++---
 .../kernels/attention/prefill_split_q_mma.cuh | 16 +++--
 csrc/kernels/attention/warp_utils.cuh         | 13 ----
 csrc/kernels/common/cp_async.cuh              | 69 ++++++++++++++++++
 csrc/kernels/common/reduce.cuh                | 48 +++++++++++++
 csrc/kernels/fp8/common.h                     |  9 +++
 csrc/kernels/fp8/gemm.cuh                     | 70 +++----------------
 csrc/kernels/fp8/ops.cu                       | 25 ++++---
 csrc/tests/attn_paged_test.cu                 |  2 +
 csrc/tests/attn_test.cu                       |  2 +
 21 files changed, 243 insertions(+), 139 deletions(-)
 delete mode 100644 csrc/kernels/attention/warp_utils.cuh
 create mode 100644 csrc/kernels/common/cp_async.cuh
 create mode 100644 csrc/kernels/common/reduce.cuh

diff --git a/csrc/kernels/attention/common.h b/csrc/kernels/attention/common.h
index fe773fe..6fe3ab0 100644
--- a/csrc/kernels/attention/common.h
+++ b/csrc/kernels/attention/common.h
@@ -1,5 +1,10 @@
 #pragma once
 
+// Pure POD header
+
+namespace astrai {
+namespace attention {
+
 // Tensor layout for Q/K/V tensors passed to attention kernels.
 // Internally, kernels always operate on BHLD [batch, n_heads, seq_len, head_dim].
 // When the caller passes BLHD, dims 1 and 2 are transposed at entry.
@@ -8,6 +13,9 @@ enum TensorLayout : int {
     BLHD = 1,  // [batch, seq_len, n_heads, head_dim]
 };
 
+// Split-KV workspace cap: max decode splits per (batch, q_head).
+constexpr int MAX_SPLITS = 32;
+
 
 // Unified attention params covering BOTH addressing modes:
 //   - Contiguous K/V: dense [batch, kv_head, kv_len, head_dim] tensors (k/v).
@@ -73,5 +81,7 @@ struct AttentionParams {
     int num_splits;
     AT* __restrict__ o_part;
     AT* __restrict__ ml_part;
-
 };
+
+}  // namespace attention
+}  // namespace astrai
diff --git a/csrc/kernels/attention/decode.cu b/csrc/kernels/attention/decode.cu
index 6d48f26..e373820 100644
--- a/csrc/kernels/attention/decode.cu
+++ b/csrc/kernels/attention/decode.cu
@@ -1,6 +1,8 @@
 #include "dispatchers.cuh"
 #include "entry_utils.cuh"
 
+using namespace astrai::attention;
+
 torch::Tensor attn_decode(
     torch::Tensor q,
     torch::Tensor k,
diff --git a/csrc/kernels/attention/decode_split_kv.cuh b/csrc/kernels/attention/decode_split_kv.cuh
index 564fe30..56c7f19 100644
--- a/csrc/kernels/attention/decode_split_kv.cuh
+++ b/csrc/kernels/attention/decode_split_kv.cuh
@@ -3,7 +3,11 @@
 #include <float.h>
 #include "common.h"
 #include "layout_policies.cuh"
-#include "warp_utils.cuh"
+#include "../common/reduce.cuh"
+
+namespace astrai {
+namespace attention {
+
 constexpr int DC_CHUNK = 64;
 
 // Scalar split-KV decode (fallback for sm < 80, no tensor cores), unified
@@ -142,3 +146,6 @@ __global__ void attn_decode_combine_kernel(AttentionParams<bf16> p) {
     int o_off = KV::q_decode_base(p, batch, q_head) + d * p.q_d_stride;
     p.o_ptr[o_off] = __float2bfloat16(acc * inv);
 }
+
+}  // namespace attention
+}  // namespace astrai
diff --git a/csrc/kernels/attention/decode_split_kv_mma.cuh b/csrc/kernels/attention/decode_split_kv_mma.cuh
index a454e7c..6026938 100644
--- a/csrc/kernels/attention/decode_split_kv_mma.cuh
+++ b/csrc/kernels/attention/decode_split_kv_mma.cuh
@@ -4,7 +4,9 @@
 #include "common.h"
 #include "layout_policies.cuh"
 #include "mma_utils.cuh"
-#include "warp_utils.cuh"
+
+namespace astrai {
+namespace attention {
 
 // Split-K (FlashDecoding) tensor-core decode via GQA head-packing, unified
 // across contiguous and paged (SGLang flat-pool) K/V via the KV template
@@ -78,10 +80,10 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
             KVAddr a = KV::template decode_addr<Traits::VEC>(
                 p, kctx, batch, kv_head, kc, d, valid, pass == 0);
             int off = r * Traits::LD + swiz_col(d, r, Traits::SWIZ_MASK);
-            cp_async_16_pred(&dK[off], a.k, a.valid);
-            cp_async_16_pred(&dV[off], a.v, a.valid);
+            astrai::cp_async_16(&dK[off], a.k, a.valid);
+            astrai::cp_async_16(&dV[off], a.v, a.valid);
         }
-        cp_async_commit();
+        astrai::cp_async_commit_group();
     };
 
     // ---- Multi-stage cp.async pipeline ----
@@ -126,9 +128,9 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
 
         for (int it = 0; it < ntiles; it++) {
             if (it + 1 == ntiles)
-                cp_async_wait_group<0>();
+                astrai::cp_async_wait_group<0>();
             else
-                cp_async_wait_group<STAGES - 1>();
+                astrai::cp_async_wait_group<STAGES - 1>();
             __syncwarp();
             process_tile(it, it & (STAGES - 1));
             __syncwarp();
@@ -139,7 +141,7 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
         // Fewer tiles than stages: load all, wait for all, process.
         for (int i = 0; i < ntiles; i++)
             load_tile(ti_begin + i, i);
-        cp_async_wait_group<0>();
+        astrai::cp_async_wait_all();
         __syncwarp();
         for (int it = 0; it < ntiles; it++)
             process_tile(it, it);
@@ -181,3 +183,6 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
         }
     }
 }
+
+}  // namespace attention
+}  // namespace astrai
diff --git a/csrc/kernels/attention/dispatchers.cuh b/csrc/kernels/attention/dispatchers.cuh
index 1fc412f..f6b78a7 100644
--- a/csrc/kernels/attention/dispatchers.cuh
+++ b/csrc/kernels/attention/dispatchers.cuh
@@ -10,7 +10,6 @@
 
 #include <cuda_runtime.h>
 #include <algorithm>
-#include "warp_utils.cuh"
 #include "layout_policies.cuh"
 #include "prefill_split_q.cuh"
 #include "decode_split_kv.cuh"
@@ -19,6 +18,9 @@
 #include "decode_split_kv_mma.cuh"
 #endif
 
+namespace astrai {
+namespace attention {
+
 // Split-KV: compute number of splits to fill all SMs for small-batch decode.
 // Caps splits so each split processes at least `min_tiles_per_split` tiles,
 // avoiding excessive loop/prologue overhead when tiles are small.
@@ -231,3 +233,6 @@ static inline void dispatch_paged_decode(AttentionParams<bf16>& p, cudaStream_t
 
     attn_decode_combine_kernel<PagedKV><<<p.batch * p.q_head, p.head_dim, 0, stream>>>(p);
 }
+
+}  // namespace attention
+}  // namespace astrai
diff --git a/csrc/kernels/attention/entry_utils.cuh b/csrc/kernels/attention/entry_utils.cuh
index 7f55705..c4d15ac 100644
--- a/csrc/kernels/attention/entry_utils.cuh
+++ b/csrc/kernels/attention/entry_utils.cuh
@@ -3,9 +3,6 @@
 #include <torch/extension.h>
 #include <c10/cuda/CUDAGuard.h>
 #include "common.h"
-#include "warp_utils.cuh"
-
-using bf16 = __nv_bfloat16;
 
 // Dispatch head_dim: shared macro — avoids C++20 lambda template syntax.
 // Usage: DISPATCH_HEAD_DIM(hd, fn, args...)
@@ -21,6 +18,11 @@ using bf16 = __nv_bfloat16;
                          " (supported: 32, 64, 128, 256)"); \
     }
 
+namespace astrai {
+namespace attention {
+
+using bf16 = __nv_bfloat16;
+
 // The split kernel unconditionally writes every (batch, q_head, split) slot it
 // owns — including empty split ranges, which store m = -FLT_MAX so the combine
 // skips them. Allocators are therefore left uninitialized (torch::empty); the
@@ -356,3 +358,6 @@ inline void attn_pack_paged_prefill_params(
     p.o_part = nullptr;
     p.ml_part = nullptr;
 }
+
+}  // namespace attention
+}  // namespace astrai
diff --git a/csrc/kernels/attention/layout_policies.cuh b/csrc/kernels/attention/layout_policies.cuh
index c8ec715..0354c1d 100644
--- a/csrc/kernels/attention/layout_policies.cuh
+++ b/csrc/kernels/attention/layout_policies.cuh
@@ -26,6 +26,9 @@
 #define DEVICE_FORCEINLINE static __device__ __forceinline__
 #define HOST_DEV_FORCEINLINE static __host__ __device__ __forceinline__
 
+namespace astrai {
+namespace attention {
+
 using bf16 = __nv_bfloat16;
 
 // ============================================================================
@@ -253,3 +256,6 @@ struct PagedKV {
         return kv_addr_from_token(p, c, token, d);
     }
 };
+
+}  // namespace attention
+}  // namespace astrai
diff --git a/csrc/kernels/attention/mma_utils.cuh b/csrc/kernels/attention/mma_utils.cuh
index 06aaebe..7cceb33 100644
--- a/csrc/kernels/attention/mma_utils.cuh
+++ b/csrc/kernels/attention/mma_utils.cuh
@@ -3,6 +3,7 @@
 #include <cuda_fp16.h>
 #include <cuda_runtime.h>
 
+#include "../common/cp_async.cuh"
 #include "../common/mma.cuh"
 
 // Predicated cp.async (4-operand form) requires CUDA 11.2+.
@@ -11,6 +12,9 @@
 #error "AstrAI CUDA kernels require CUDA 11.2 or later (CUDART_VERSION >= 11020)."
 #endif
 
+namespace astrai {
+namespace attention {
+
 // ============================================================================
 // KernelTraits — FlashAttention-v2 style compile-time configuration bundle.
 //
@@ -75,36 +79,9 @@ __device__ __forceinline__ int swiz_col(int d, int r, int mask = 7) {
     return ((d >> 3) ^ (r & mask)) << 3 | (d & 7);
 }
 
-// Predicated cp.async: copy 16 bytes when `pred`, otherwise zero-fill.
-// BypassL1 defaults to .cg (L2 only); false selects .ca (L1 + L2).
-// src_size=0 means no bytes are read, so an out-of-bounds address is safe.
-template <bool BypassL1 = true>
-__device__ __forceinline__ void cp_async_16_pred(bf16* smem_ptr,
-                                                  const void* gmem_ptr,
-                                                  bool pred) {
-    unsigned smem_addr = __cvta_generic_to_shared(smem_ptr);
-    int src_size = pred ? 16 : 0;
-    if constexpr (BypassL1) {
-        asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;"
-                     :: "r"(smem_addr), "l"(gmem_ptr), "r"(src_size));
-    } else {
-        asm volatile("cp.async.ca.shared.global [%0], [%1], 16, %2;"
-                     :: "r"(smem_addr), "l"(gmem_ptr), "r"(src_size));
-    }
-}
-
-__device__ __forceinline__ void cp_async_commit() {
-    asm volatile("cp.async.commit_group;");
-}
-
-__device__ __forceinline__ void cp_async_wait_all() {
-    asm volatile("cp.async.wait_all;");
-}
-
-template <int N>
-__device__ __forceinline__ void cp_async_wait_group() {
-    asm volatile("cp.async.wait_group %0;" :: "n"(N));
-}
+// cp.async primitives live in the shared template (common/cp_async.cuh):
+// `astrai::cp_async_16` (predicated), `astrai::cp_async_commit_group`,
+// `astrai::cp_async_wait_group<N>` / `_wait_all` stage the K/V tiles.
 
 // ---------------------------------------------------------------------------
 // Q-load: load query rows directly from global memory into mma A-operand
@@ -272,3 +249,6 @@ __device__ inline void mma_pv_accumulate(
         }
     }
 }
+
+}  // namespace attention
+}  // namespace astrai
diff --git a/csrc/kernels/attention/paged_decode.cu b/csrc/kernels/attention/paged_decode.cu
index bff3662..93dcaef 100644
--- a/csrc/kernels/attention/paged_decode.cu
+++ b/csrc/kernels/attention/paged_decode.cu
@@ -1,6 +1,8 @@
 #include "dispatchers.cuh"
 #include "entry_utils.cuh"
 
+using namespace astrai::attention;
+
 torch::Tensor attn_paged_decode(
     torch::Tensor q,
     torch::Tensor k_cache,
diff --git a/csrc/kernels/attention/paged_prefill.cu b/csrc/kernels/attention/paged_prefill.cu
index 157401f..812aad7 100644
--- a/csrc/kernels/attention/paged_prefill.cu
+++ b/csrc/kernels/attention/paged_prefill.cu
@@ -1,6 +1,8 @@
 #include "dispatchers.cuh"
 #include "entry_utils.cuh"
 
+using namespace astrai::attention;
+
 torch::Tensor attn_paged_prefill(
     torch::Tensor q,
     torch::Tensor k_cache,
diff --git a/csrc/kernels/attention/prefill.cu b/csrc/kernels/attention/prefill.cu
index 47a8e8a..9bd9399 100644
--- a/csrc/kernels/attention/prefill.cu
+++ b/csrc/kernels/attention/prefill.cu
@@ -1,6 +1,8 @@
 #include "dispatchers.cuh"
 #include "entry_utils.cuh"
 
+using namespace astrai::attention;
+
 torch::Tensor attn_prefill(
     torch::Tensor q,
     torch::Tensor k,
diff --git a/csrc/kernels/attention/prefill_split_q.cuh b/csrc/kernels/attention/prefill_split_q.cuh
index acdc2b1..2717a38 100644
--- a/csrc/kernels/attention/prefill_split_q.cuh
+++ b/csrc/kernels/attention/prefill_split_q.cuh
@@ -3,6 +3,10 @@
 #include <cuda_bf16.h>
 #include "common.h"
 #include "layout_policies.cuh"
+#include "../common/reduce.cuh"
+
+namespace astrai {
+namespace attention {
 
 using bf16 = __nv_bfloat16;
 
@@ -11,14 +15,7 @@ using bf16 = __nv_bfloat16;
 // compile-time bools — the compiler eliminates dead branches.
 // Unified across contiguous and paged (SGLang flat-pool) K/V via KV.
 // Templated on <HEAD_DIM, KV, G, ROWS, P_BC, IsCausal, HasMask>.
-
-template <int G>
-__device__ __forceinline__ float group_reduce_sum(float v, unsigned mask) {
-#pragma unroll
-    for (int o = G / 2; o > 0; o >>= 1)
-        v += __shfl_xor_sync(mask, v, o);
-    return v;
-}
+// group_reduce_sum<G> lives in common/reduce.cuh (astrai::).
 
 // load 8 contiguous bf16 from (16-byte aligned) smem as one float4
 __device__ __forceinline__ void ld8(const bf16* p, float* o) {
@@ -155,3 +152,6 @@ __global__ void attn_prefill_split_q_kernel_t(AttentionParams<bf16> p) {
             p.o_ptr[o_off + i * p.q_d_stride] = __float2bfloat16(acc[i] * rl);
     }
 }
+
+}  // namespace attention
+}  // namespace astrai
diff --git a/csrc/kernels/attention/prefill_split_q_mma.cuh b/csrc/kernels/attention/prefill_split_q_mma.cuh
index 2dc616b..4692529 100644
--- a/csrc/kernels/attention/prefill_split_q_mma.cuh
+++ b/csrc/kernels/attention/prefill_split_q_mma.cuh
@@ -5,6 +5,9 @@
 #include "layout_policies.cuh"
 #include "mma_utils.cuh"
 
+namespace astrai {
+namespace attention {
+
 // Tensor-core prefill flash attention (raw mma.sync PTX), unified across
 // contiguous and paged (SGLang flat-pool) K/V via the KV template parameter.
 // One warp owns BR=16 query rows. S = Q@K^T and O = P@V run on bf16 tensor
@@ -85,10 +88,10 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
             int token = KV::resolve_token(p, kctx, kc, valid);
             KVAddr a = KV::kv_addr_from_token(p, kctx, token, d);
             int off = r * Traits::LD + swiz_col(d, r, Traits::SWIZ_MASK);
-            cp_async_16_pred(&dK[off], a.k, a.valid);
-            cp_async_16_pred(&dV[off], a.v, a.valid);
+            astrai::cp_async_16(&dK[off], a.k, a.valid);
+            astrai::cp_async_16(&dV[off], a.v, a.valid);
         }
-        cp_async_commit();
+        astrai::cp_async_commit_group();
     };
 
     // ---- Prologue: issue first tile load ----
@@ -98,7 +101,7 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
         int buf = ti & 1;
 
         // Wait for current tile, then publish cross-warp + guard buffer reuse.
-        cp_async_wait_group<0>();
+        astrai::cp_async_wait_group<0>();
         __syncthreads();
         if (ti < t_end) load_tile(ti + 1, (ti + 1) & 1);
 
@@ -149,9 +152,12 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
         }
         if (qr1 < q_len) {
             __nv_bfloat162 v = __floats2bfloat162_rn(Oacc[dn8][2] * rl1,
-                                                      Oacc[dn8][3] * rl1);
+                                                     Oacc[dn8][3] * rl1);
             *reinterpret_cast<__nv_bfloat162*>(
                 &p.o_ptr[o_base + qr1 * p.q_l_stride + d * p.q_d_stride]) = v;
         }
     }
 }
+
+}  // namespace attention
+}  // namespace astrai
diff --git a/csrc/kernels/attention/warp_utils.cuh b/csrc/kernels/attention/warp_utils.cuh
deleted file mode 100644
index 241626e..0000000
--- a/csrc/kernels/attention/warp_utils.cuh
+++ /dev/null
@@ -1,13 +0,0 @@
-#pragma once
-#include <cuda_bf16.h>
-
-using bf16 = __nv_bfloat16;
-
-static constexpr int MAX_SPLITS = 32;
-
-__device__ inline float warp_reduce_sum(float val) {
-    #pragma unroll
-    for (int offset = 16; offset > 0; offset >>= 1)
-        val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
-    return val;
-}
diff --git a/csrc/kernels/common/cp_async.cuh b/csrc/kernels/common/cp_async.cuh
new file mode 100644
index 0000000..add6c9e
--- /dev/null
+++ b/csrc/kernels/common/cp_async.cuh
@@ -0,0 +1,69 @@
+// Shared cp.async primitives — pure CUDA, no torch.
+//
+// One header for the async-copy pipeline used by both the attention kernels
+// (predicated 16-byte K/V tile staging) and the fp8 GEMM (predicated operand
+// staging + wait_group dispatch). PTX requires wait_group's operand to be an
+// immediate, hence the template forms.
+
+#pragma once
+
+#include <cuda_runtime.h>
+
+namespace astrai {
+
+// Predicated cp.async: copy 16 bytes when `pred`, otherwise zero-fill.
+// src_size=0 means no bytes are read, so an out-of-bounds address is safe.
+// BypassL1 defaults to .cg (L2 only); false selects .ca (L1 + L2).
+// `T` is the smem element type; only the destination pointer's type matters.
+template <typename T, bool BypassL1 = true>
+__device__ __forceinline__ void cp_async_16(T* smem_ptr, const void* gmem_ptr,
+                                            bool pred) {
+    const unsigned smem_addr = __cvta_generic_to_shared(smem_ptr);
+    const int src_size = pred ? 16 : 0;
+    if constexpr (BypassL1) {
+        asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;"
+                     :: "r"(smem_addr), "l"(gmem_ptr), "r"(src_size));
+    } else {
+        asm volatile("cp.async.ca.shared.global [%0], [%1], 16, %2;"
+                     :: "r"(smem_addr), "l"(gmem_ptr), "r"(src_size));
+    }
+}
+
+// Commit all outstanding cp.async ops of this thread as one group.
+__device__ __forceinline__ void cp_async_commit_group() {
+    asm volatile("cp.async.commit_group;");
+}
+
+// Wait for every committed group (pipeline drain).
+__device__ __forceinline__ void cp_async_wait_all() {
+    asm volatile("cp.async.wait_all;");
+}
+
+// Wait until at most KeepGroups committed groups are still in flight.
+// PTX requires an immediate operand; keep it as a template argument so the
+// stage policy stays 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));
+}
+
+// Runtime dispatch over cp_async_wait_group<N>: unrolls into a compare
+// ladder over [0, MaxKeepGroups] so the immediate-only PTX constraint is
+// hidden behind a runtime `keep_groups` (used by the fp8 GEMM pipeline,
+// whose remaining-tile count is dynamic).
+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>();
+    }
+}
+
+}  // namespace astrai
diff --git a/csrc/kernels/common/reduce.cuh b/csrc/kernels/common/reduce.cuh
new file mode 100644
index 0000000..fb8a890
--- /dev/null
+++ b/csrc/kernels/common/reduce.cuh
@@ -0,0 +1,48 @@
+// Shared warp/block reduction + atomic helpers — pure CUDA, no torch.
+//
+// Extracted from the attention and fp8 families so both share one
+// implementation: warp_reduce_sum (decode scalar kernel), warp_reduce_max +
+// atomic_max_float (fp8 quantize amax), group_reduce_sum<G> (prefill scalar
+// kernel).
+
+#pragma once
+
+namespace astrai {
+
+// Full-warp butterfly sum reduction (32 lanes).
+__device__ __forceinline__ float warp_reduce_sum(float val) {
+#pragma unroll
+    for (int offset = 16; offset > 0; offset >>= 1)
+        val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
+    return val;
+}
+
+// Full-warp butterfly max reduction (32 lanes).
+__device__ __forceinline__ float warp_reduce_max(float value) {
+#pragma unroll
+    for (int offset = 16; offset > 0; offset >>= 1)
+        value = fmaxf(value, __shfl_xor_sync(0xffffffffu, value, offset));
+    return value;
+}
+
+// Sub-warp group reduction over G consecutive lanes (G a power of two).
+// `mask` is the full participating-lane mask of the group (see the
+// prefill scalar kernel's gmask computation).
+template <int G>
+__device__ __forceinline__ float group_reduce_sum(float v, unsigned mask) {
+#pragma unroll
+    for (int o = G / 2; o > 0; o >>= 1)
+        v += __shfl_xor_sync(mask, v, o);
+    return v;
+}
+
+// Unsigned-bit-pattern atomicMax for non-negative floats; a null
+// destination disables the update (kernels with optional amax slots).
+__device__ __forceinline__ void atomic_max_float(float* destination,
+                                                 float value) {
+    if (destination)
+        atomicMax(reinterpret_cast<unsigned*>(destination),
+                  __float_as_uint(value));
+}
+
+}  // namespace astrai
diff --git a/csrc/kernels/fp8/common.h b/csrc/kernels/fp8/common.h
index a1eefa4..85aaabf 100644
--- a/csrc/kernels/fp8/common.h
+++ b/csrc/kernels/fp8/common.h
@@ -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
diff --git a/csrc/kernels/fp8/gemm.cuh b/csrc/kernels/fp8/gemm.cuh
index 6d60e44..d970bae 100644
--- a/csrc/kernels/fp8/gemm.cuh
+++ b/csrc/kernels/fp8/gemm.cuh
@@ -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
diff --git a/csrc/kernels/fp8/ops.cu b/csrc/kernels/fp8/ops.cu
index f021134..c54c587 100644
--- a/csrc/kernels/fp8/ops.cu
+++ b/csrc/kernels/fp8/ops.cu
@@ -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
diff --git a/csrc/tests/attn_paged_test.cu b/csrc/tests/attn_paged_test.cu
index 9010526..be11b06 100644
--- a/csrc/tests/attn_paged_test.cu
+++ b/csrc/tests/attn_paged_test.cu
@@ -9,6 +9,8 @@
 #include "test_utils.cuh"
 #include "../kernels/attention/dispatchers.cuh"
 
+using namespace astrai::attention;
+
 struct PagedDecodeDispatch { AttentionParams<bf16>& p; template<int H> void operator()() { dispatch_paged_decode<H>(p, 0); } };
 struct PagedPrefillDispatch { AttentionParams<bf16>& p; template<int H> void operator()() { dispatch_paged_prefill<H>(p, 0); } };
 
diff --git a/csrc/tests/attn_test.cu b/csrc/tests/attn_test.cu
index b42c2ec..b7328e4 100644
--- a/csrc/tests/attn_test.cu
+++ b/csrc/tests/attn_test.cu
@@ -9,6 +9,8 @@ nvcc -I csrc -arch=sm_89 -O3 \
 #include "test_utils.cuh"
 #include "../kernels/attention/dispatchers.cuh"
 
+using namespace astrai::attention;
+
 struct DecodeDispatch { AttentionParams<bf16>& p; template<int H> void operator()() { dispatch_decode<H>(p, 0); } };
 struct PrefillDispatch { AttentionParams<bf16>& p; template<int H> void operator()() { dispatch_prefill<H>(p, 0); } };