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

refactor: drop gemv variant shape tables and flatten kernel dir

Browse Source 
- delete the warp-tiled kernel and both per-shape (N,K) selector tables; block size is 256 threads everywhere except M=8 with N*K <= 12 MiB, which keeps a 128-thread CTA
- HBM-streaming measurements (weight copies rotated through L2, the real decode regime) show the variants within ~3% on L20 because the kernel is bandwidth-bound; the retired tables were tuned against an L2-resident loop and sometimes picked the slowest variant ((2048,8192) M=8: coop128 6% slower than coop256)
- a shape no longer switches kernels (and accumulation order) with M, removing one shape-dependent nondeterminism source
- remove the stale split-K launcher comment
- move bf16_gemv.cu and bf16_swiglu.cu from csrc/kernels/gemv/ to csrc/kernels/ beside rotary_emb.cu; the family keeps no shared headers
- rename test_bf16_gemv_matches_half_cta_edge_bands to test_bf16_gemv_matches_m8_edge_bands and update docs/developer/cuda_kernels.md

Benchmark: L20 (sm_89), PyTorch 2.11.0+cu128, interleaved CUDA-event timing with rotated weight copies exceeding the 96MB L2; variant spread <=3% across 14 shapes x M in {1,2,4,8}, and the retained rule wins 5-9% at M=8 small weights ((512,3584), (1536,1536), (6912,1536))
This commit is contained in:
ViperEkura
2026-09-03 06:34:20 +08:00
parent d4a292b36b
commit d6f757dc13
5 changed files with 28 additions and 148 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/CMakeLists.txt
+2 -2
View File
@@ -70,8 +70,8 @@ set(KERNEL_SRCS
attention/prefill.cu
attention/paged_decode.cu
attention/paged_prefill.cu
gemv/bf16_gemv.cu
gemv/bf16_swiglu.cu
bf16_gemv.cu
bf16_swiglu.cu
rotary_emb.cu
)
csrc/kernels/gemv/bf16_gemv.cu → csrc/kernels/bf16_gemv.cu
+12 -133
View File
@@ -14,7 +14,6 @@ namespace {
constexpr int kThreads = 256;
constexpr int kHalfCtaThreads = 128;
constexpr int kWarpSize = 32;
constexpr int kWarpTiledThreads = 128;
__device__ __forceinline__ float warp_sum(float value) {
#pragma unroll
@@ -159,120 +158,6 @@ __global__ void bf16_gemv_kernel(
}
}
template <int Rows>
__global__ void bf16_gemv_aligned_warp_tiled_kernel(
const __nv_bfloat16* __restrict__ x,
const __nv_bfloat16* __restrict__ weight,
const __nv_bfloat16* __restrict__ bias,
__nv_bfloat16* __restrict__ output,
int n,
int k
) {
constexpr int kWarpsPerBlock = kWarpTiledThreads / kWarpSize;
const int lane = threadIdx.x & (kWarpSize - 1);
const int warp = threadIdx.x / kWarpSize;
const int output_index = blockIdx.x * kWarpsPerBlock + warp;
if (output_index >= n) {
return;
}
// The launcher selects this path only when each row is 16-byte aligned.
// Four independent output rows per CTA remove the block-wide reduction
// barrier and improve occupancy for the medium LLaMA projection bands.
const int vectors = k / 8;
const auto* x4 = reinterpret_cast<const uint4*>(x);
const auto* w4 = reinterpret_cast<const uint4*>(weight) +
static_cast<int64_t>(output_index) * vectors;
float sums[Rows] = {};
for (int vector = lane; vector < vectors; vector += kWarpSize) {
const uint4 wv_raw = w4[vector];
const auto* wv = reinterpret_cast<const __nv_bfloat162*>(&wv_raw);
#pragma unroll
for (int row = 0; row < Rows; ++row) {
const uint4 xv_raw =
x4[static_cast<int64_t>(row) * vectors + vector];
const auto* xv = reinterpret_cast<const __nv_bfloat162*>(&xv_raw);
#pragma unroll
for (int pair = 0; pair < 4; ++pair) {
sums[row] = fmaf(
__bfloat162float(__low2bfloat16(xv[pair])),
__bfloat162float(__low2bfloat16(wv[pair])),
sums[row]
);
sums[row] = fmaf(
__bfloat162float(__high2bfloat16(xv[pair])),
__bfloat162float(__high2bfloat16(wv[pair])),
sums[row]
);
}
}
}
#pragma unroll
for (int row = 0; row < Rows; ++row) {
sums[row] = warp_sum(sums[row]);
if (lane == 0) {
if (bias != nullptr) {
sums[row] += __bfloat162float(bias[output_index]);
}
output[row * n + output_index] = __float2bfloat16_rn(sums[row]);
}
}
}
template <int Rows>
constexpr bool use_warp_tiled_kernel(int n, int k) {
// These bands are intentionally narrow and are validated by the common
// transformer benchmark. The 256-thread cooperative kernel remains the
// fallback for arbitrary K, larger projections, and M=2 (where the
// single-warp reduction regresses the current vectorized kernel).
if constexpr (Rows == 4) {
return (n == 1024 && k == 4096) ||
(n == 4096 && k == 4096) ||
(n == 11008 && k == 4096) ||
(n == 4096 && k == 11008);
}
return false;
}
template <int Rows>
constexpr bool use_half_cta_kernel(int n, int k) {
// A 128-thread CTA reduces synchronization and scheduling overhead for
// selected medium decode projections. Keep the selector exact: long-K
// and bandwidth-saturated shapes regress, and the winning bands differ
// materially with the number of reused input rows.
if constexpr (Rows == 1) {
return n == 8192 && k == 2048;
}
if constexpr (Rows == 2) {
return (n == 4096 && k == 4096) ||
(n == 11008 && k == 4096) ||
(n == 3584 && k == 3584) ||
(n == 2048 && k == 2048) ||
(n == 8192 && k == 2048);
}
if constexpr (Rows == 4) {
return (n == 5120 && k == 5120) ||
(n == 3584 && k == 3584) ||
(n == 2048 && k == 2048) ||
(n == 8192 && k == 2048);
}
if constexpr (Rows == 8) {
return (n == 4096 && k == 4096) ||
(n == 11008 && k == 4096) ||
(n == 4096 && k == 11008) ||
(n == 1024 && k == 4096) ||
(n == 5120 && k == 5120) ||
(n == 512 && k == 3584) ||
(n == 3584 && k == 3584) ||
(n == 1024 && k == 8192) ||
(n == 2048 && k == 2048) ||
(n == 8192 && k == 2048) ||
(n == 2048 && k == 8192);
}
return false;
}
template <int Rows>
void launch_bf16_gemv(
const __nv_bfloat16* x,
@@ -283,27 +168,23 @@ void launch_bf16_gemv(
int k,
cudaStream_t stream
) {
const bool aligned_rows = k % 8 == 0 &&
(reinterpret_cast<uintptr_t>(x) & 15u) == 0u &&
(reinterpret_cast<uintptr_t>(weight) & 15u) == 0u;
if constexpr (Rows == 4) {
if (aligned_rows && use_warp_tiled_kernel<Rows>(n, k)) {
constexpr int kWarpsPerBlock = kWarpTiledThreads / kWarpSize;
const int blocks = (n + kWarpsPerBlock - 1) / kWarpsPerBlock;
bf16_gemv_aligned_warp_tiled_kernel<Rows>
<<<blocks, kWarpTiledThreads, 0, stream>>>(
// Decode is HBM weight-streaming bound: with weights rotated through L2,
// 128/256-thread CTAs measure within noise on L20 except for small
// weight matrices at the largest decode batch, where the smaller CTA
// wins 5-9% (see docs/developer/decode_linear_benchmark.md).
constexpr int64_t kSmallWeightLimit = int64_t{12} << 20;
if constexpr (Rows == 8) {
if (k % 8 == 0 &&
(reinterpret_cast<uintptr_t>(x) & 15u) == 0u &&
(reinterpret_cast<uintptr_t>(weight) & 15u) == 0u &&
static_cast<int64_t>(n) * k <= kSmallWeightLimit) {
bf16_gemv_kernel<Rows, kHalfCtaThreads>
<<<n, kHalfCtaThreads, 0, stream>>>(
x, weight, bias, output, n, k
);
return;
}
}
if (aligned_rows && use_half_cta_kernel<Rows>(n, k)) {
bf16_gemv_kernel<Rows, kHalfCtaThreads>
<<<n, kHalfCtaThreads, 0, stream>>>(
x, weight, bias, output, n, k
);
return;
}
bf16_gemv_kernel<Rows, kThreads><<<n, kThreads, 0, stream>>>(
x, weight, bias, output, n, k
);
@@ -366,8 +247,6 @@ torch::Tensor bf16_gemv(
auto output = x.dim() == 1 ? torch::empty({n}, x.options())
: torch::empty({m, n}, x.options());
// Small-N decode shapes (GQA k/v projections) cannot fill the GPU with
// one block per output row; split K across extra blocks and reduce.
const auto* x_ptr = reinterpret_cast<const __nv_bfloat16*>(x.data_ptr());
const auto* weight_ptr =
reinterpret_cast<const __nv_bfloat16*>(weight.data_ptr());
csrc/kernels/gemv/bf16_swiglu.cu → csrc/kernels/bf16_swiglu.cu
View File
docs/developer/cuda_kernels.md
+13 -12
View File
@@ -14,27 +14,28 @@ selected by guarded model dispatchers described below.
| `attn_paged_decode` | `attention/paged_decode.cu` | Paged KV cache decode attention |
| `attn_paged_prefill` | `attention/paged_prefill.cu` | Paged KV cache prefill attention (ragged batch) |
| `rotary_emb` | `rotary_emb.cu` | Fused rotary embedding (cos/sin lookup + rotation) |
| `bf16_gemv` | `gemv/bf16_gemv.cu` | M=1..8 BF16 linear with FP32 accumulation (sm_80+) |
| `bf16_swiglu` | `gemv/bf16_swiglu.cu` | Fused M=1..8 BF16 up/gate projections and SwiGLU epilogue (sm_80+) |
| `bf16_gemv` | `bf16_gemv.cu` | M=1..8 BF16 linear with FP32 accumulation (sm_80+) |
| `bf16_swiglu` | `bf16_swiglu.cu` | Fused M=1..8 BF16 up/gate projections and SwiGLU epilogue (sm_80+) |
| `fp8_ops` | `fp8/ops.cu` | FP8 quantization + tensor-core GEMM (sm_89+) |
### BF16 GEMV primitive
`astrai.extension.bf16_gemv(x, weight, bias=None)` accepts a contiguous BF16
input shaped `[K]` or `[M, K]`, with `M` in `[1, 8]` and any positive `K`, and
row-major weights `[N, K]`. The general path assigns one 256-thread CTA to an
output row and computes all M results together, reusing the weight row across
tokens. For measured aligned M=4 medium projections, a 128-thread CTA instead
assigns one output to each of four warps. That removes the CTA-wide reduction
barrier and exposes four neighboring outputs without changing accumulation.
row-major weights `[N, K]`. One CTA computes an output row for all M tokens
together, reusing the weight row across tokens. CTA size is 256 threads,
except for small weight matrices (`N*K <= 12 MiB`) at `M=8`, where a
128-thread CTA measured 5-9% faster on L20. Variant selection is otherwise
intentionally shape-free: under HBM-streaming conditions (weights rotated
through L2, as in real decode) the kernel is bandwidth-bound and block-size
choice measures within noise, so earlier per-shape variant tables were
removed along with the warp-tiled kernel.
The weight stream uses 128-bit vectorized loads anchored at each row's first
16-byte-aligned address with scalar head/tail sweeps for unaligned remainders,
so arbitrary `K` and storage offsets stay correct. The warp-tiled path is used
only when both tensors and every row are 16-byte aligned; all other calls keep
the general arbitrary-K path. Accumulation is FP32; the optional BF16 bias is
fused before the BF16 store. The launcher uses the current CUDA stream, is
CUDA Graph capture-safe, and requires sm_80 or newer.
so arbitrary `K` and storage offsets stay correct. Accumulation is FP32; the
optional BF16 bias is fused before the BF16 store. The launcher uses the
current CUDA stream, is CUDA Graph capture-safe, and requires sm_80 or newer.
Model `Linear` calls route through the lightweight linear backend. Set
`ASTRAI_GEMV=0` for an unconditional `F.linear` fallback, `1` to force the
tests/extension/test_gemv.py
+1 -1
View File
@@ -93,7 +93,7 @@ def test_bf16_gemv_matches_common_transformer_shapes(m, n, k):
(8, 2048, 8192),
],
)
def test_bf16_gemv_matches_half_cta_edge_bands(m, n, k):
def test_bf16_gemv_matches_m8_edge_bands(m, n, k):
torch.manual_seed(2026 + m + n + k)
x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
weight = torch.empty(n, k, device="cuda", dtype=torch.bfloat16)