attn_res kernel latency improvements (#50185)
gnovack committed
Aug 6, 2026 at 08:53 UTC
7b4ed49628abd7860a435d6798feef76a944cb02
1 file changed
+84
-39
csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu
+84
-39
@@ -91,7 +91,8 @@ __device__ __forceinline__ void mbarrier_wait(uint64_t& barrier, int phase) {
91
"bra WAIT;\n"
92
"DONE:\n"
93
"}\n" ::"r"(barrier_addr),
94
- "r"(phase));
94
+ "r"(phase)
95
+ : "memory");
96
}
97
98
__device__ __forceinline__ void mbarrier_arrive(uint64_t& barrier) {
@@ -101,7 +102,8 @@ __device__ __forceinline__ void mbarrier_arrive(uint64_t& barrier) {
102
"{\n"
103
".reg .b64 state;\n"
104
"mbarrier.arrive.shared::cta.b64 state, [%0];\n"
104
- "}\n" ::"r"(barrier_addr));
105
+ "}\n" ::"r"(barrier_addr)
106
+ : "memory");
107
}
108
109
__device__ __forceinline__ void fence_mbarrier_init() {
@@ -227,14 +229,14 @@ __device__ __forceinline__ void cp_async_bulk(void* smem_dst,
229
: "memory");
230
}
231
230
-template <int H, int NC = N_CHUNK_DEFAULT, bool RELEASE_TMEM = false,
231
- bool HAS_DELTA = false, bool HAS_OUTPUT_NORM = false,
232
- bool OUTPUT_NORM_IN_SMEM = false>
232
+template <int H, int N, int NC = N_CHUNK_DEFAULT, int B = 1,
233
+ bool RELEASE_TMEM = false, bool HAS_DELTA = false,
234
+ bool HAS_OUTPUT_NORM = false, bool OUTPUT_NORM_IN_SMEM = false>
235
__global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(
236
const bf16_t* __restrict__ block_res, bf16_t* __restrict__ layer_res,
237
const bf16_t* __restrict__ delta, const bf16_t* __restrict__ res_w,
236
- const bf16_t* __restrict__ rms_w, bf16_t* __restrict__ output, int N, int T,
237
- int B, int block_stride_m, int block_stride_r, float rms_eps,
238
+ const bf16_t* __restrict__ rms_w, bf16_t* __restrict__ output, int T,
239
+ int block_stride_m, int block_stride_r, float rms_eps,
240
const bf16_t* __restrict__ output_norm_weight, float output_norm_eps) {
241
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && __CUDA_ARCH__ < 1100
242
constexpr float LOG2_E = 1.4426950408889634f;
@@ -258,7 +260,7 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(
260
const int lane = tid & 31;
261
const int TB = T * B;
262
const int num_ctas = gridDim.x;
261
- const int num_chunks = (N + N_CHUNK - 1) / N_CHUNK;
263
+ constexpr int num_chunks = (N + N_CHUNK - 1) / N_CHUNK;
264
265
const int comp_wid = wid - 1;
266
const int comp_tid = tid - 32;
@@ -329,11 +331,16 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(
331
if constexpr (H == 7168) {
332
if (si == SLICES_PER_GROUP - 1) {
333
int h_base = 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4;
332
- #pragma unroll
333
- for (int j = 0; j < 4; j++) {
334
- int h = h_base + j;
335
- q_cache[si * VEC + j] =
336
- __bfloat162float(rms_w[h]) * __bfloat162float(res_w[h]);
334
+ int2 rms_v = *reinterpret_cast<const int2*>(rms_w + h_base);
335
+ int2 res_v = *reinterpret_cast<const int2*>(res_w + h_base);
336
+ auto* rms2 = reinterpret_cast<__nv_bfloat162*>(&rms_v);
337
+ auto* res2 = reinterpret_cast<__nv_bfloat162*>(&res_v);
338
+ #pragma unroll
339
+ for (int k = 0; k < 2; k++) {
340
+ float2 rf = __bfloat1622float2(rms2[k]);
341
+ float2 sf = __bfloat1622float2(res2[k]);
342
+ q_cache[si * VEC + 2 * k] = rf.x * sf.x;
343
+ q_cache[si * VEC + 2 * k + 1] = rf.y * sf.y;
344
}
345
continue;
346
}
@@ -341,11 +348,16 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(
348
int dt = si * CONSUMER_GROUPS + group;
349
if (dt >= NHT) continue;
350
int h_base = dt * K_TILE + k_local;
344
- #pragma unroll
345
- for (int j = 0; j < VEC; j++) {
346
- int h = h_base + j;
347
- q_cache[si * VEC + j] =
348
- __bfloat162float(rms_w[h]) * __bfloat162float(res_w[h]);
351
+ int4 rms_v = *reinterpret_cast<const int4*>(rms_w + h_base);
352
+ int4 res_v = *reinterpret_cast<const int4*>(res_w + h_base);
353
+ auto* rms2 = reinterpret_cast<__nv_bfloat162*>(&rms_v);
354
+ auto* res2 = reinterpret_cast<__nv_bfloat162*>(&res_v);
355
+ #pragma unroll
356
+ for (int k = 0; k < 4; k++) {
357
+ float2 rf = __bfloat1622float2(rms2[k]);
358
+ float2 sf = __bfloat1622float2(res2[k]);
359
+ q_cache[si * VEC + 2 * k] = rf.x * sf.x;
360
+ q_cache[si * VEC + 2 * k + 1] = rf.y * sf.y;
361
}
362
}
363
}
@@ -404,6 +416,7 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(
416
acc32[i] = 0.f;
417
}
418
419
+ #pragma unroll
420
for (int ci = 0; ci < num_chunks; ci++, gci++) {
421
int ns = ci * N_CHUNK;
422
int an = min(N_CHUNK, N - ns);
@@ -854,13 +867,13 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(
867
#endif
868
}
869
857
-template <int H, int NC = N_CHUNK_DEFAULT, bool RELEASE_TMEM = false,
870
+template <int H, int N, int NC = N_CHUNK_DEFAULT, bool RELEASE_TMEM = false,
871
bool HAS_DELTA = false, bool HAS_OUTPUT_NORM = false,
872
bool OUTPUT_NORM_IN_SMEM = false>
873
static void launch_fwd(const bf16_t* block_residual, bf16_t* layer_residual,
874
const bf16_t* delta, const bf16_t* res_weight,
862
- const bf16_t* rms_weight, bf16_t* output, int N, int T,
863
- int B, float rms_eps, int num_sm, cudaStream_t stream,
875
+ const bf16_t* rms_weight, bf16_t* output, int T,
876
+ float rms_eps, int num_sm, cudaStream_t stream,
877
const bf16_t* output_norm_weight = nullptr,
878
float output_norm_eps = 0.f, int block_stride_m = 0,
879
int block_stride_r = 0) {
@@ -870,7 +883,7 @@ static void launch_fwd(const bf16_t* block_residual, bf16_t* layer_residual,
883
sizeof(FwdSmemPlan<NC>) + 15) &
884
~size_t(15);
885
auto kernel =
873
- &attn_res_fwd_online_v2_kernel<H, NC, RELEASE_TMEM, HAS_DELTA,
886
+ &attn_res_fwd_online_v2_kernel<H, N, NC, 1, RELEASE_TMEM, HAS_DELTA,
887
HAS_OUTPUT_NORM, OUTPUT_NORM_IN_SMEM>;
888
static bool attrs_set = false;
889
if (!attrs_set) {
@@ -892,7 +905,7 @@ static void launch_fwd(const bf16_t* block_residual, bf16_t* layer_residual,
905
config.attrs = attrs;
906
config.numAttrs = 1;
907
cudaLaunchKernelEx(&config, kernel, block_residual, layer_residual, delta,
895
- res_weight, rms_weight, output, N, T, B, block_stride_m,
908
+ res_weight, rms_weight, output, T, block_stride_m,
909
block_stride_r, rms_eps, output_norm_weight,
910
output_norm_eps);
911
}
@@ -926,33 +939,65 @@ void kimi_k3_attn_res(torch::stable::Tensor& prefix,
939
// Two-source chunks and two resident CTAs are beneficial once setup is
940
// amortized by the long, full eight-block prefill workload.
941
if (num_blocks == 8 && num_tokens >= 4096) {
929
- launch_fwd<7168, 2, true, true, true, true>(
942
+ launch_fwd<7168, 9, 2, true, true, true, true>(
943
static_cast<bf16_t const*>(blocks.data_ptr()),
944
static_cast<bf16_t*>(prefix.data_ptr()),
945
static_cast<bf16_t const*>(delta.data_ptr()),
946
static_cast<bf16_t const*>(qk_weight.data_ptr()),
947
static_cast<bf16_t const*>(norm_weight.data_ptr()),
935
- static_cast<bf16_t*>(output.data_ptr()),
936
- static_cast<int>(num_blocks) + 1, num_tokens, 1,
948
+ static_cast<bf16_t*>(output.data_ptr()), num_tokens,
949
static_cast<float>(eps), properties->multiProcessorCount,
950
get_current_cuda_stream(device),
951
static_cast<bf16_t const*>(output_norm_weight.data_ptr()),
952
static_cast<float>(output_norm_eps), static_cast<int>(blocks.stride(0)),
953
static_cast<int>(blocks.stride(1)));
954
} else {
943
- launch_fwd<7168, 4, false, true, true, true>(
944
- static_cast<bf16_t const*>(blocks.data_ptr()),
945
- static_cast<bf16_t*>(prefix.data_ptr()),
946
- static_cast<bf16_t const*>(delta.data_ptr()),
947
- static_cast<bf16_t const*>(qk_weight.data_ptr()),
948
- static_cast<bf16_t const*>(norm_weight.data_ptr()),
949
- static_cast<bf16_t*>(output.data_ptr()),
950
- static_cast<int>(num_blocks) + 1, num_tokens, 1,
951
- static_cast<float>(eps), properties->multiProcessorCount,
952
- get_current_cuda_stream(device),
953
- static_cast<bf16_t const*>(output_norm_weight.data_ptr()),
954
- static_cast<float>(output_norm_eps), static_cast<int>(blocks.stride(0)),
955
- static_cast<int>(blocks.stride(1)));
955
+ auto dispatch = [&](auto nsrc_tok) {
956
+ constexpr int NSRC = decltype(nsrc_tok)::value;
957
+ launch_fwd<7168, NSRC, 3, false, true, true, true>(
958
+ static_cast<bf16_t const*>(blocks.data_ptr()),
959
+ static_cast<bf16_t*>(prefix.data_ptr()),
960
+ static_cast<bf16_t const*>(delta.data_ptr()),
961
+ static_cast<bf16_t const*>(qk_weight.data_ptr()),
962
+ static_cast<bf16_t const*>(norm_weight.data_ptr()),
963
+ static_cast<bf16_t*>(output.data_ptr()), num_tokens,
964
+ static_cast<float>(eps), properties->multiProcessorCount,
965
+ get_current_cuda_stream(device),
966
+ static_cast<bf16_t const*>(output_norm_weight.data_ptr()),
967
+ static_cast<float>(output_norm_eps),
968
+ static_cast<int>(blocks.stride(0)),
969
+ static_cast<int>(blocks.stride(1)));
970
+ };
971
+ // The source count is num_blocks + 1; specialising on it makes the chunk
972
+ // count, per-chunk source count and prefix-chunk test compile-time.
973
+ switch (num_blocks) {
974
+ case 1:
975
+ dispatch(std::integral_constant<int, 2>{});
976
+ break;
977
+ case 2:
978
+ dispatch(std::integral_constant<int, 3>{});
979
+ break;
980
+ case 3:
981
+ dispatch(std::integral_constant<int, 4>{});
982
+ break;
983
+ case 4:
984
+ dispatch(std::integral_constant<int, 5>{});
985
+ break;
986
+ case 5:
987
+ dispatch(std::integral_constant<int, 6>{});
988
+ break;
989
+ case 6:
990
+ dispatch(std::integral_constant<int, 7>{});
991
+ break;
992
+ case 7:
993
+ dispatch(std::integral_constant<int, 8>{});
994
+ break;
995
+ case 8:
996
+ dispatch(std::integral_constant<int, 9>{});
997
+ break;
998
+ default:
999
+ STD_TORCH_CHECK(false, "Kimi K3 AttnRes: num_blocks must be 1..8");
1000
+ }
1001
}
1002
cudaError_t const error = cudaGetLastError();
1003
STD_TORCH_CHECK(