From 27abb7c5e7f4670a8d005a97a3782cbc49ed81c6 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Thu, 3 Sep 2026 06:46:25 +0800 Subject: [PATCH] perf: drop swiglu warp-rows variant for an M=8 block-size rule - 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 --- csrc/kernels/bf16_swiglu.cu | 142 ++++++----------------------- docs/developer/cuda_kernels.md | 14 +-- docs/developer/swiglu_benchmark.md | 25 +++++ 3 files changed, 59 insertions(+), 122 deletions(-) diff --git a/csrc/kernels/bf16_swiglu.cu b/csrc/kernels/bf16_swiglu.cu index 99fe0b1..749d768 100644 --- a/csrc/kernels/bf16_swiglu.cu +++ b/csrc/kernels/bf16_swiglu.cu @@ -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 #include @@ -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 +template __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 +template 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<<>>( - x, up_weight, gate_weight, output, n, k - ); -} - -template -__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( - x + static_cast(row) * k - ); - const auto* up4 = reinterpret_cast( - up_weight + static_cast(output_index) * k - ); - const auto* gate4 = reinterpret_cast( - gate_weight + static_cast(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(&x_raw); - const auto* up_values = - reinterpret_cast(&up_raw); - const auto* gate_values = - reinterpret_cast(&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(row) * n + output_index] = - __float2bfloat16_rn(up_sum * silu); - } -} - -template -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<<>>( + bf16_swiglu_kernel<<>>( 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(n); const int k_int = static_cast(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()); diff --git a/docs/developer/cuda_kernels.md b/docs/developer/cuda_kernels.md index 0c3c132..9303b6f 100644 --- a/docs/developer/cuda_kernels.md +++ b/docs/developer/cuda_kernels.md @@ -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. diff --git a/docs/developer/swiglu_benchmark.md b/docs/developer/swiglu_benchmark.md index 54b37df..575ad81 100644 --- a/docs/developer/swiglu_benchmark.md +++ b/docs/developer/swiglu_benchmark.md @@ -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.