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

perf: drop swiglu warp-rows variant for an M=8 block-size rule

Browse Source 
- delete the warp-per-row kernel and the (6912,1536) M=2/4/8 dispatch table; under rotated cold weights the warp path is 2-6% slower than CTA reuse at M=2/4, and the table had been tuned against L2-resident timing
- a single CTA-reuse kernel now serves all M in [1, 8]; block size is 256 threads for M in [1, 7] and 128 for M=8, where the shorter shared-memory reduction tree wins
- document in docs/developer/swiglu_benchmark.md that the earlier operator numbers were L2-resident: the fused kernel sits at the dual-stream cold-read floor (702 vs 699 GB/s at (6912,1536); 369 vs 370 GB/s at (11008,4096)) and wide matrices cap at ~370-400 GB/s even for pure reads, so the reported M=8 -23% regression does not survive the cold regime
- update docs/developer/cuda_kernels.md accordingly

Benchmark: L20 (sm_89), PyTorch 2.11.0+cu128, rotated weight copies >= 240 MB to defeat the 96 MB L2; end-to-end through the built module at (6912,1536) reaches 738-752 GB/s for M in [1, 4] and 702 GB/s at M=8, about +8% at M=2/4 and +6% at M=8 over the removed warp path
This commit is contained in:
ViperEkura
2026-09-03 06:46:25 +08:00
parent d6f757dc13
commit 27abb7c5e7
3 changed files with 59 additions and 122 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/kernels/bf16_swiglu.cu
+26 -116
View File
@@ -1,4 +1,8 @@
// Fused small-M BF16 SwiGLU primitive for decode-time dense MLP layers.
// One CTA per output column; each weight pair is read once and reused across
// all decode rows. Bandwidth-bound in the cold-HBM decode regime, so variant
// selection beyond the M=8 block-size rule is noise (see
// docs/developer/swiglu_benchmark.md).
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
@@ -11,9 +15,7 @@
namespace {
constexpr int kThreads = 256;
constexpr int kWarpSize = 32;
constexpr int kWarps = kThreads / kWarpSize;
__device__ __forceinline__ float warp_sum(float value) {
#pragma unroll
@@ -27,7 +29,7 @@ __device__ __forceinline__ float round_bf16(float value) {
return __bfloat162float(__float2bfloat16_rn(value));
}
template <int Rows>
template <int Threads, int Rows>
__global__ void bf16_swiglu_kernel(
const __nv_bfloat16* __restrict__ x,
const __nv_bfloat16* __restrict__ up_weight,
@@ -36,6 +38,7 @@ __global__ void bf16_swiglu_kernel(
int n,
int k
) {
constexpr int kWarps = Threads / kWarpSize;
const int output_index = blockIdx.x;
const int lane = threadIdx.x & (kWarpSize - 1);
const int warp = threadIdx.x / kWarpSize;
@@ -120,7 +123,7 @@ __global__ void bf16_swiglu_kernel(
}
}
template <int Rows>
template <int Threads, int Rows>
void launch_bf16_swiglu(
const __nv_bfloat16* x,
const __nv_bfloat16* up_weight,
@@ -130,85 +133,7 @@ void launch_bf16_swiglu(
int k,
cudaStream_t stream
) {
bf16_swiglu_kernel<Rows><<<n, kThreads, 0, stream>>>(
x, up_weight, gate_weight, output, n, k
);
}
template <int Rows>
__global__ void bf16_swiglu_warp_rows_kernel(
const __nv_bfloat16* __restrict__ x,
const __nv_bfloat16* __restrict__ up_weight,
const __nv_bfloat16* __restrict__ gate_weight,
__nv_bfloat16* __restrict__ output,
int n,
int k
) {
const int output_index = blockIdx.x;
const int row = threadIdx.x / kWarpSize;
const int lane = threadIdx.x & (kWarpSize - 1);
const int vector_count = k / 8;
float up_sum = 0.0f;
float gate_sum = 0.0f;
const auto* x4 = reinterpret_cast<const uint4*>(
x + static_cast<int64_t>(row) * k
);
const auto* up4 = reinterpret_cast<const uint4*>(
up_weight + static_cast<int64_t>(output_index) * k
);
const auto* gate4 = reinterpret_cast<const uint4*>(
gate_weight + static_cast<int64_t>(output_index) * k
);
// A warp owns one decode row. Same-address weight reads from sibling
// warps are served through the read-only/L1 path, while each row avoids
// CTA-wide shared-memory reductions and synchronization.
for (int vector_index = lane;
vector_index < vector_count;
vector_index += kWarpSize) {
const uint4 x_raw = x4[vector_index];
const uint4 up_raw = up4[vector_index];
const uint4 gate_raw = gate4[vector_index];
const auto* x_values = reinterpret_cast<const __nv_bfloat162*>(&x_raw);
const auto* up_values =
reinterpret_cast<const __nv_bfloat162*>(&up_raw);
const auto* gate_values =
reinterpret_cast<const __nv_bfloat162*>(&gate_raw);
#pragma unroll
for (int pair = 0; pair < 4; ++pair) {
const float2 xv = __bfloat1622float2(x_values[pair]);
const float2 uv = __bfloat1622float2(up_values[pair]);
const float2 gv = __bfloat1622float2(gate_values[pair]);
up_sum = fmaf(xv.x, uv.x, up_sum);
up_sum = fmaf(xv.y, uv.y, up_sum);
gate_sum = fmaf(xv.x, gv.x, gate_sum);
gate_sum = fmaf(xv.y, gv.y, gate_sum);
}
}
up_sum = warp_sum(up_sum);
gate_sum = warp_sum(gate_sum);
if (lane == 0) {
up_sum = round_bf16(up_sum);
gate_sum = round_bf16(gate_sum);
const float silu =
round_bf16(gate_sum / (1.0f + expf(-gate_sum)));
output[static_cast<int64_t>(row) * n + output_index] =
__float2bfloat16_rn(up_sum * silu);
}
}
template <int Rows>
void launch_bf16_swiglu_warp_rows(
const __nv_bfloat16* x,
const __nv_bfloat16* up_weight,
const __nv_bfloat16* gate_weight,
__nv_bfloat16* output,
int n,
int k,
cudaStream_t stream
) {
bf16_swiglu_warp_rows_kernel<Rows><<<n, Rows * kWarpSize, 0, stream>>>(
bf16_swiglu_kernel<Threads, Rows><<<n, Threads, 0, stream>>>(
x, up_weight, gate_weight, output, n, k
);
}
@@ -288,66 +213,51 @@ torch::Tensor bf16_swiglu(
reinterpret_cast<__nv_bfloat16*>(output.data_ptr());
const int n_int = static_cast<int>(n);
const int k_int = static_cast<int>(k);
const bool use_warp_rows = n_int == 6912 && k_int == 1536;
// Block size 256 keeps the weight streams at the HBM bandwidth floor for
// M in [1, 7]; M=8 halves the CTA so each thread owns more of the row
// and the shared-memory reduction tree shrinks (measured on L20 with
// rotated cold weights; larger CTAs only add idle warps).
switch (m) {
case 1:
launch_bf16_swiglu<1>(
launch_bf16_swiglu<256, 1>(
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
);
break;
case 2:
if (use_warp_rows) {
launch_bf16_swiglu_warp_rows<2>(
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
);
} else {
launch_bf16_swiglu<2>(
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
);
}
launch_bf16_swiglu<256, 2>(
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
);
break;
case 3:
launch_bf16_swiglu<3>(
launch_bf16_swiglu<256, 3>(
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
);
break;
case 4:
if (use_warp_rows) {
launch_bf16_swiglu_warp_rows<4>(
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
);
} else {
launch_bf16_swiglu<4>(
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
);
}
launch_bf16_swiglu<256, 4>(
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
);
break;
case 5:
launch_bf16_swiglu<5>(
launch_bf16_swiglu<256, 5>(
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
);
break;
case 6:
launch_bf16_swiglu<6>(
launch_bf16_swiglu<256, 6>(
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
);
break;
case 7:
launch_bf16_swiglu<7>(
launch_bf16_swiglu<256, 7>(
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
);
break;
case 8:
if (use_warp_rows) {
launch_bf16_swiglu_warp_rows<8>(
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
);
} else {
launch_bf16_swiglu<8>(
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
);
}
launch_bf16_swiglu<128, 8>(
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
);
break;
}
C10_CUDA_CHECK(cudaGetLastError());
docs/developer/cuda_kernels.md
+8 -6
View File
@@ -125,12 +125,14 @@ BF16 inputs with `M` in `[1, 8]` and K divisible by 8. It preserves the BF16
rounding boundaries of the two projection outputs, SiLU output, and final
product while accumulating dot products in FP32.
The kernel contains two output-row tilings. A CTA-reuse path reads each up/gate
weight chunk once and applies it to all M rows. The native AstrAI 1B shape
`(N,K)=(6912,1536)` uses one warp per decode row for M=2/4/8; on L20 this
removes the shared reductions and barrier and reduces M=4 CUDA-Graph latency
from 0.0324 ms to 0.0181 ms. Wider LLaMA/GPT-NeoX matrices keep CTA reuse,
because duplicating their weight reads across row warps regressed 1.3-4.2%.
The kernel is a single CTA-reuse tiling: one CTA per output column reads each
up/gate weight chunk once and applies it to all M rows. Block size is 256
threads for M in [1, 7] and 128 for M=8, where the shorter shared-memory
reduction tree wins under cold-HBM decode traffic. An earlier per-shape
`(6912,1536)` warp-per-row variant and its dispatch table were removed: HBM
measurements with rotated weights showed the table was tuned against L2-cache
regime timing and was up to 6% slower than CTA reuse at M=2/4; the kernel is
bandwidth-bound, so finer variant selection is noise.
Dense `MLP` modules route through the SwiGLU backend. `ASTRAI_SWIGLU=0` keeps
the unfused linear backend, and `1` explicitly forces the fused primitive.
docs/developer/swiglu_benchmark.md
+25
View File
@@ -56,3 +56,28 @@ the existing unfused linear backend, including any independently qualified
GEMV dispatch. `ASTRAI_SWIGLU=1` remains an explicit benchmark/experimentation
switch for callers that accept normal BF16 reduction-order variation. A future
automatic band must repeat both the performance and checkpoint-output gates.
## HBM re-measurement and kernel simplification
The operator numbers above are L2-resident: the AstrAI pair is 40.5 MB,
smaller than the 96 MB L2, so a tight timing loop re-reads warm weights
(13.75 us implies ~3.1 TB/s, far above the 864 GB/s spec). Real decode rotates
~1 GB of per-layer weights through L2 every step, so every call is cold.
Re-measuring with rotated weight copies (>= 240 MB working set) on the same
L20 showed:
- The fused CTA-reuse kernel sits at the dual-stream cold-read floor
(702 vs 699 GB/s at (6912,1536); 369 vs 370 GB/s at (11008,4096)). Wide
LLaMA matrices cap at ~370-400 GB/s regardless of kernel, even for a
pure-read loop, so the old per-variant gaps there were noise.
- The `(6912,1536)` warp-per-row variant (formerly M=2/4/8) is 2-6% slower
than CTA reuse at M=2/4 under cold weights and no longer wins at M=8 once
the CTA drops to 128 threads. It and its dispatch table were deleted.
- New rule: 256 threads for M in [1, 7], 128 threads for M=8. End-to-end
through the built module at (6912,1536): 738-752 GB/s for M in [1, 4] and
702 GB/s at M=8 (+6% over the removed warp path).
The M=8 CUDA-Graph regression reported above (`-23.24%`) does not survive the
cold-weight regime: cuBLAS reaches L2 bandwidth in the warm loop while both
fused paths converge to the same HBM floor.