diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_3b3789d6356f15f9cbe3_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_09e14f1d9cbd6739db61_binding.cu similarity index 97% rename from csrc/cake_dsv4/sm_100a/cake_dsv4_3b3789d6356f15f9cbe3_binding.cu rename to csrc/cake_dsv4/sm_100a/cake_dsv4_09e14f1d9cbd6739db61_binding.cu index 08603a80ed1..20cc47f3686 100644 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_3b3789d6356f15f9cbe3_binding.cu +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_09e14f1d9cbd6739db61_binding.cu @@ -15,7 +15,7 @@ */ // Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_3b3789d6356f15f9cbe3'. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_09e14f1d9cbd6739db61'. #include #include #include @@ -31,10 +31,10 @@ #include #include -extern "C" __global__ void kernel_cake_dsv4_3b3789d6356f15f9cbe3(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int batch_size, int max_q_len, int ragged_query, int has_sinks); +extern "C" __global__ void kernel_cake_dsv4_09e14f1d9cbd6739db61(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int batch_size, int max_q_len, int ragged_query, int has_sinks); -namespace cake_host_shim_d2767e2d94919bb7 { +namespace cake_host_shim_db4b534bf8fd17b7 { using tvm::ffi::Optional; using tvm::ffi::TensorView; @@ -422,7 +422,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ } if (smem_status == cudaSuccess) { smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_3b3789d6356f15f9cbe3), + reinterpret_cast(kernel_cake_dsv4_09e14f1d9cbd6739db61), cudaFuncAttributeMaxDynamicSharedMemorySize, 198144); } @@ -432,7 +432,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ } } TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_3b3789d6356f15f9cbe3 failed: " + << "cudaFuncSetAttribute for kernel_cake_dsv4_09e14f1d9cbd6739db61 failed: " << cudaGetErrorString(smem_status); dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); dim3 block(256u, 1u, 1u); @@ -443,12 +443,12 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ config.dynamicSmemBytes = 198144u; config.stream = stream; cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_3b3789d6356f15f9cbe3), kargs); + &config, reinterpret_cast(kernel_cake_dsv4_09e14f1d9cbd6739db61), kargs); TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_3b3789d6356f15f9cbe3 failed: " + << "cudaLaunchKernelExC for kernel_cake_dsv4_09e14f1d9cbd6739db61 failed: " << cudaGetErrorString(launch_status); } -} // namespace cake_host_shim_d2767e2d94919bb7 +} // namespace cake_host_shim_db4b534bf8fd17b7 -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_d2767e2d94919bb7::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_db4b534bf8fd17b7::Run); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_3b3789d6356f15f9cbe3_kernel.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_09e14f1d9cbd6739db61_kernel.cu similarity index 99% rename from csrc/cake_dsv4/sm_100a/cake_dsv4_3b3789d6356f15f9cbe3_kernel.cu rename to csrc/cake_dsv4/sm_100a/cake_dsv4_09e14f1d9cbd6739db61_kernel.cu index 3a0cb8057d1..bdf35286adc 100644 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_3b3789d6356f15f9cbe3_kernel.cu +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_09e14f1d9cbd6739db61_kernel.cu @@ -1060,7 +1060,7 @@ __device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { extern "C" { __global__ __launch_bounds__(256, 1) void -kernel_cake_dsv4_3b3789d6356f15f9cbe3(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int batch_size, int max_q_len, int ragged_query, int has_sinks) +kernel_cake_dsv4_09e14f1d9cbd6739db61(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int batch_size, int max_q_len, int ragged_query, int has_sinks) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); @@ -1595,15 +1595,6 @@ kernel_cake_dsv4_3b3789d6356f15f9cbe3(const __grid_constant__ CUtensorMap tmap_q int tile_count_2 = (active_topk_2 + 128 - 1) / 128; unsigned int load_kv_stage = 0; unsigned int load_kv_phase = 1; - if (load_warp_rank == 0) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(q_full_addr, 65536); - #pragma unroll - for (int q_stage = 0; q_stage < 8; q_stage += 8) { - tma_4d_gmem2smem(smem_q_addr + (unsigned int)(q_stage * 8192), (&tmap_q), 0, 0, q_stage, query_idx_2, q_full_addr); - } - } - } int group = load_warp_rank + lane * 3; int row0 = 0; int row1 = 0; @@ -1642,6 +1633,15 @@ kernel_cake_dsv4_3b3789d6356f15f9cbe3(const __grid_constant__ CUtensorMap tmap_q row2 = ((raw_rows[2] >= 0) ? raw_rows[2] : 0); row3 = ((raw_rows[3] >= 0) ? raw_rows[3] : 0); } + if (load_warp_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(q_full_addr, 65536); + #pragma unroll + for (int q_stage = 0; q_stage < 8; q_stage += 8) { + tma_4d_gmem2smem(smem_q_addr + (unsigned int)(q_stage * 8192), (&tmap_q), 0, 0, q_stage, query_idx_2, q_full_addr); + } + } + } #pragma unroll for (int tile_2 = 0; tile_2 < 1; tile_2++) { if (tile_count_2 > tile_2) { diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_57e7b93def70f2be356f_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_0c8ab955ac7b62e7bfee_binding.cu similarity index 97% rename from csrc/cake_dsv4/sm_100a/cake_dsv4_57e7b93def70f2be356f_binding.cu rename to csrc/cake_dsv4/sm_100a/cake_dsv4_0c8ab955ac7b62e7bfee_binding.cu index da00a1e33ac..2f4858f771a 100644 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_57e7b93def70f2be356f_binding.cu +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_0c8ab955ac7b62e7bfee_binding.cu @@ -15,7 +15,7 @@ */ // Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_57e7b93def70f2be356f'. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_0c8ab955ac7b62e7bfee'. #include #include #include @@ -31,10 +31,10 @@ #include #include -extern "C" __global__ void kernel_cake_dsv4_57e7b93def70f2be356f(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int num_head_tiles, int has_sinks); +extern "C" __global__ void kernel_cake_dsv4_0c8ab955ac7b62e7bfee(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int num_head_tiles, int has_sinks); -namespace cake_host_shim_16cf3a7d0d01e13d { +namespace cake_host_shim_900ed6f94bca9ff0 { using tvm::ffi::Optional; using tvm::ffi::TensorView; @@ -357,7 +357,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_O, Te } if (smem_status == cudaSuccess) { smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_57e7b93def70f2be356f), + reinterpret_cast(kernel_cake_dsv4_0c8ab955ac7b62e7bfee), cudaFuncAttributeMaxDynamicSharedMemorySize, 198528); } @@ -367,7 +367,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_O, Te } } TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_57e7b93def70f2be356f failed: " + << "cudaFuncSetAttribute for kernel_cake_dsv4_0c8ab955ac7b62e7bfee failed: " << cudaGetErrorString(smem_status); dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); dim3 block(512u, 1u, 1u); @@ -378,12 +378,12 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_O, Te config.dynamicSmemBytes = 198528u; config.stream = stream; cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_57e7b93def70f2be356f), kargs); + &config, reinterpret_cast(kernel_cake_dsv4_0c8ab955ac7b62e7bfee), kargs); TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_57e7b93def70f2be356f failed: " + << "cudaLaunchKernelExC for kernel_cake_dsv4_0c8ab955ac7b62e7bfee failed: " << cudaGetErrorString(launch_status); } -} // namespace cake_host_shim_16cf3a7d0d01e13d +} // namespace cake_host_shim_900ed6f94bca9ff0 -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_16cf3a7d0d01e13d::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_900ed6f94bca9ff0::Run); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_57e7b93def70f2be356f_kernel.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_0c8ab955ac7b62e7bfee_kernel.cu similarity index 89% rename from csrc/cake_dsv4/sm_100a/cake_dsv4_57e7b93def70f2be356f_kernel.cu rename to csrc/cake_dsv4/sm_100a/cake_dsv4_0c8ab955ac7b62e7bfee_kernel.cu index bb28b1bf0d1..c36df67ae57 100644 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_57e7b93def70f2be356f_kernel.cu +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_0c8ab955ac7b62e7bfee_kernel.cu @@ -999,7 +999,7 @@ __device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { extern "C" { __global__ __launch_bounds__(512, 1) void -kernel_cake_dsv4_57e7b93def70f2be356f(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int num_head_tiles, int has_sinks) +kernel_cake_dsv4_0c8ab955ac7b62e7bfee(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int num_head_tiles, int has_sinks) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); @@ -1370,29 +1370,129 @@ kernel_cake_dsv4_57e7b93def70f2be356f(const __grid_constant__ CUtensorMap tmap_q " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" : "=r"(*reinterpret_cast(&_tmem_load_1[32])), "=r"(*reinterpret_cast(&_tmem_load_1[33])), "=r"(*reinterpret_cast(&_tmem_load_1[34])), "=r"(*reinterpret_cast(&_tmem_load_1[35])), "=r"(*reinterpret_cast(&_tmem_load_1[36])), "=r"(*reinterpret_cast(&_tmem_load_1[37])), "=r"(*reinterpret_cast(&_tmem_load_1[38])), "=r"(*reinterpret_cast(&_tmem_load_1[39])), "=r"(*reinterpret_cast(&_tmem_load_1[40])), "=r"(*reinterpret_cast(&_tmem_load_1[41])), "=r"(*reinterpret_cast(&_tmem_load_1[42])), "=r"(*reinterpret_cast(&_tmem_load_1[43])), "=r"(*reinterpret_cast(&_tmem_load_1[44])), "=r"(*reinterpret_cast(&_tmem_load_1[45])), "=r"(*reinterpret_cast(&_tmem_load_1[46])), "=r"(*reinterpret_cast(&_tmem_load_1[47])), "=r"(*reinterpret_cast(&_tmem_load_1[48])), "=r"(*reinterpret_cast(&_tmem_load_1[49])), "=r"(*reinterpret_cast(&_tmem_load_1[50])), "=r"(*reinterpret_cast(&_tmem_load_1[51])), "=r"(*reinterpret_cast(&_tmem_load_1[52])), "=r"(*reinterpret_cast(&_tmem_load_1[53])), "=r"(*reinterpret_cast(&_tmem_load_1[54])), "=r"(*reinterpret_cast(&_tmem_load_1[55])), "=r"(*reinterpret_cast(&_tmem_load_1[56])), "=r"(*reinterpret_cast(&_tmem_load_1[57])), "=r"(*reinterpret_cast(&_tmem_load_1[58])), "=r"(*reinterpret_cast(&_tmem_load_1[59])), "=r"(*reinterpret_cast(&_tmem_load_1[60])), "=r"(*reinterpret_cast(&_tmem_load_1[61])), "=r"(*reinterpret_cast(&_tmem_load_1[62])), "=r"(*reinterpret_cast(&_tmem_load_1[63])) : "r"(taddr + 256 + (unsigned int)row_addr + 32)); - int col_base = col_half_1 * 64; - if (head_idx_1 < num_heads) { - #pragma unroll - for (int offset = 0; offset < 64; offset += 8) { - { - const float2 _prescale2_0 = {inv_sum * output_scale, inv_sum * output_scale}; - #if __CUDA_ARCH__ >= 1000 - #pragma unroll - for (int _ps = 0; _ps < 4; _ps++) - mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_1[offset])[_ps], _prescale2_0); - #else - #pragma unroll - for (int _ps = 0; _ps < 8; _ps++) - _tmem_load_1[offset + _ps] *= inv_sum * output_scale; - #endif - __nv_bfloat162 _pk[4]; - _pk[0] = __floats2bfloat162_rn(_tmem_load_1[offset + 0], _tmem_load_1[offset + 1]); - _pk[1] = __floats2bfloat162_rn(_tmem_load_1[offset + 2], _tmem_load_1[offset + 3]); - _pk[2] = __floats2bfloat162_rn(_tmem_load_1[offset + 4], _tmem_load_1[offset + 5]); - _pk[3] = __floats2bfloat162_rn(_tmem_load_1[offset + 6], _tmem_load_1[offset + 7]); - *reinterpret_cast(&((__nv_bfloat16*)(O + (output_base + col_base + offset)))[0]) = *reinterpret_cast(&_pk[0]); - } - } + float2 _f2_14 = make_float2(inv_sum * output_scale, inv_sum * output_scale); + #pragma unroll + for (int i_1 = 0; i_1 < 64; i_1 += 2) { + float2 _f2_15 = make_float2(_tmem_load_1[i_1], _tmem_load_1[i_1 + 1]); + float2 _mul_f32x2_0; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_0) : "l"(*(const unsigned long long*)&_f2_15), "l"(*(const unsigned long long*)&_f2_14)); + _tmem_load_1[i_1] = _mul_f32x2_0.x; + _tmem_load_1[i_1 + 1] = _mul_f32x2_0.y; + } + unsigned int packed_o[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_1[_lp*2 + 0], _tmem_load_1[_lp*2+1 + 0])); + packed_o[_lp] = *(uint32_t*)&_bf2; + } + bool stage_hi_1 = (lane & 1) != 0; + bool stage_hi_2 = (lane & 2) != 0; + bool stage_hi_4 = (lane & 4) != 0; + unsigned int _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[0] : packed_o[4]), 1); + packed_o[0] = ((stage_hi_1) ? _shfl_xor_2 : packed_o[0]); + packed_o[4] = ((stage_hi_1) ? packed_o[4] : _shfl_xor_2); + unsigned int _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[1] : packed_o[5]), 1); + packed_o[1] = ((stage_hi_1) ? _shfl_xor_3 : packed_o[1]); + packed_o[5] = ((stage_hi_1) ? packed_o[5] : _shfl_xor_3); + unsigned int _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[2] : packed_o[6]), 1); + packed_o[2] = ((stage_hi_1) ? _shfl_xor_4 : packed_o[2]); + packed_o[6] = ((stage_hi_1) ? packed_o[6] : _shfl_xor_4); + unsigned int _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[3] : packed_o[7]), 1); + packed_o[3] = ((stage_hi_1) ? _shfl_xor_5 : packed_o[3]); + packed_o[7] = ((stage_hi_1) ? packed_o[7] : _shfl_xor_5); + unsigned int _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[8] : packed_o[12]), 1); + packed_o[8] = ((stage_hi_1) ? _shfl_xor_6 : packed_o[8]); + packed_o[12] = ((stage_hi_1) ? packed_o[12] : _shfl_xor_6); + unsigned int _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[9] : packed_o[13]), 1); + packed_o[9] = ((stage_hi_1) ? _shfl_xor_7 : packed_o[9]); + packed_o[13] = ((stage_hi_1) ? packed_o[13] : _shfl_xor_7); + unsigned int _shfl_xor_8 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[10] : packed_o[14]), 1); + packed_o[10] = ((stage_hi_1) ? _shfl_xor_8 : packed_o[10]); + packed_o[14] = ((stage_hi_1) ? packed_o[14] : _shfl_xor_8); + unsigned int _shfl_xor_9 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[11] : packed_o[15]), 1); + packed_o[11] = ((stage_hi_1) ? _shfl_xor_9 : packed_o[11]); + packed_o[15] = ((stage_hi_1) ? packed_o[15] : _shfl_xor_9); + unsigned int _shfl_xor_10 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[16] : packed_o[20]), 1); + packed_o[16] = ((stage_hi_1) ? _shfl_xor_10 : packed_o[16]); + packed_o[20] = ((stage_hi_1) ? packed_o[20] : _shfl_xor_10); + unsigned int _shfl_xor_11 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[17] : packed_o[21]), 1); + packed_o[17] = ((stage_hi_1) ? _shfl_xor_11 : packed_o[17]); + packed_o[21] = ((stage_hi_1) ? packed_o[21] : _shfl_xor_11); + unsigned int _shfl_xor_12 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[18] : packed_o[22]), 1); + packed_o[18] = ((stage_hi_1) ? _shfl_xor_12 : packed_o[18]); + packed_o[22] = ((stage_hi_1) ? packed_o[22] : _shfl_xor_12); + unsigned int _shfl_xor_13 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[19] : packed_o[23]), 1); + packed_o[19] = ((stage_hi_1) ? _shfl_xor_13 : packed_o[19]); + packed_o[23] = ((stage_hi_1) ? packed_o[23] : _shfl_xor_13); + unsigned int _shfl_xor_14 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[24] : packed_o[28]), 1); + packed_o[24] = ((stage_hi_1) ? _shfl_xor_14 : packed_o[24]); + packed_o[28] = ((stage_hi_1) ? packed_o[28] : _shfl_xor_14); + unsigned int _shfl_xor_15 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[25] : packed_o[29]), 1); + packed_o[25] = ((stage_hi_1) ? _shfl_xor_15 : packed_o[25]); + packed_o[29] = ((stage_hi_1) ? packed_o[29] : _shfl_xor_15); + unsigned int _shfl_xor_16 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[26] : packed_o[30]), 1); + packed_o[26] = ((stage_hi_1) ? _shfl_xor_16 : packed_o[26]); + packed_o[30] = ((stage_hi_1) ? packed_o[30] : _shfl_xor_16); + unsigned int _shfl_xor_17 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[27] : packed_o[31]), 1); + packed_o[27] = ((stage_hi_1) ? _shfl_xor_17 : packed_o[27]); + packed_o[31] = ((stage_hi_1) ? packed_o[31] : _shfl_xor_17); + int lane_t = lane & 1; + unsigned int chunk_words[4]; + chunk_words[0] = packed_o[0]; + chunk_words[1] = packed_o[1]; + chunk_words[2] = packed_o[2]; + chunk_words[3] = packed_o[3]; + if (head_tile_1 * 64 + (my_row_1 & -2 | 0) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + head_tile_1 * 64 + (my_row_1 & -2 | 0)) * 512 + v_chunk * 128 + col_half_1 * 64 + (0 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[4]; + chunk_words[1] = packed_o[5]; + chunk_words[2] = packed_o[6]; + chunk_words[3] = packed_o[7]; + if (head_tile_1 * 64 + (my_row_1 & -2 | 1) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + head_tile_1 * 64 + (my_row_1 & -2 | 1)) * 512 + v_chunk * 128 + col_half_1 * 64 + (0 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[8]; + chunk_words[1] = packed_o[9]; + chunk_words[2] = packed_o[10]; + chunk_words[3] = packed_o[11]; + if (head_tile_1 * 64 + (my_row_1 & -2 | 0) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + head_tile_1 * 64 + (my_row_1 & -2 | 0)) * 512 + v_chunk * 128 + col_half_1 * 64 + (2 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[12]; + chunk_words[1] = packed_o[13]; + chunk_words[2] = packed_o[14]; + chunk_words[3] = packed_o[15]; + if (head_tile_1 * 64 + (my_row_1 & -2 | 1) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + head_tile_1 * 64 + (my_row_1 & -2 | 1)) * 512 + v_chunk * 128 + col_half_1 * 64 + (2 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[16]; + chunk_words[1] = packed_o[17]; + chunk_words[2] = packed_o[18]; + chunk_words[3] = packed_o[19]; + if (head_tile_1 * 64 + (my_row_1 & -2 | 0) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + head_tile_1 * 64 + (my_row_1 & -2 | 0)) * 512 + v_chunk * 128 + col_half_1 * 64 + (4 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[20]; + chunk_words[1] = packed_o[21]; + chunk_words[2] = packed_o[22]; + chunk_words[3] = packed_o[23]; + if (head_tile_1 * 64 + (my_row_1 & -2 | 1) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + head_tile_1 * 64 + (my_row_1 & -2 | 1)) * 512 + v_chunk * 128 + col_half_1 * 64 + (4 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[24]; + chunk_words[1] = packed_o[25]; + chunk_words[2] = packed_o[26]; + chunk_words[3] = packed_o[27]; + if (head_tile_1 * 64 + (my_row_1 & -2 | 0) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + head_tile_1 * 64 + (my_row_1 & -2 | 0)) * 512 + v_chunk * 128 + col_half_1 * 64 + (6 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[28]; + chunk_words[1] = packed_o[29]; + chunk_words[2] = packed_o[30]; + chunk_words[3] = packed_o[31]; + if (head_tile_1 * 64 + (my_row_1 & -2 | 1) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + head_tile_1 * 64 + (my_row_1 & -2 | 1)) * 512 + v_chunk * 128 + col_half_1 * 64 + (6 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; } asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); asm volatile("tcgen05.fence::before_thread_sync;"); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_9cf30aacb43b71034660_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_12ac507a3c73f2784621_binding.cu similarity index 96% rename from csrc/cake_dsv4/sm_103a/cake_dsv4_9cf30aacb43b71034660_binding.cu rename to csrc/cake_dsv4/sm_100a/cake_dsv4_12ac507a3c73f2784621_binding.cu index b5754808d14..579d030d38f 100644 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_9cf30aacb43b71034660_binding.cu +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_12ac507a3c73f2784621_binding.cu @@ -15,7 +15,7 @@ */ // Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_9cf30aacb43b71034660'. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_12ac507a3c73f2784621'. #include #include #include @@ -31,10 +31,10 @@ #include #include -extern "C" __global__ void kernel_cake_dsv4_9cf30aacb43b71034660(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks); +extern "C" __global__ void kernel_cake_dsv4_12ac507a3c73f2784621(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks); -namespace cake_host_shim_8b6a5b581ef12c6c { +namespace cake_host_shim_96d072de9e3f7f61 { using tvm::ffi::Optional; using tvm::ffi::TensorView; @@ -353,7 +353,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_O, Te } if (smem_status == cudaSuccess) { smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_9cf30aacb43b71034660), + reinterpret_cast(kernel_cake_dsv4_12ac507a3c73f2784621), cudaFuncAttributeMaxDynamicSharedMemorySize, 198528); } @@ -363,7 +363,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_O, Te } } TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_9cf30aacb43b71034660 failed: " + << "cudaFuncSetAttribute for kernel_cake_dsv4_12ac507a3c73f2784621 failed: " << cudaGetErrorString(smem_status); dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); dim3 block(512u, 1u, 1u); @@ -374,12 +374,12 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_O, Te config.dynamicSmemBytes = 198528u; config.stream = stream; cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_9cf30aacb43b71034660), kargs); + &config, reinterpret_cast(kernel_cake_dsv4_12ac507a3c73f2784621), kargs); TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_9cf30aacb43b71034660 failed: " + << "cudaLaunchKernelExC for kernel_cake_dsv4_12ac507a3c73f2784621 failed: " << cudaGetErrorString(launch_status); } -} // namespace cake_host_shim_8b6a5b581ef12c6c +} // namespace cake_host_shim_96d072de9e3f7f61 -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_8b6a5b581ef12c6c::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_96d072de9e3f7f61::Run); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_4ac65f00a4a032c10ecb_kernel.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_12ac507a3c73f2784621_kernel.cu similarity index 89% rename from csrc/cake_dsv4/sm_100a/cake_dsv4_4ac65f00a4a032c10ecb_kernel.cu rename to csrc/cake_dsv4/sm_100a/cake_dsv4_12ac507a3c73f2784621_kernel.cu index 1e58d437826..0fa750be2db 100644 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_4ac65f00a4a032c10ecb_kernel.cu +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_12ac507a3c73f2784621_kernel.cu @@ -999,7 +999,7 @@ __device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { extern "C" { __global__ __launch_bounds__(512, 1) void -kernel_cake_dsv4_4ac65f00a4a032c10ecb(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks) +kernel_cake_dsv4_12ac507a3c73f2784621(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); @@ -1364,29 +1364,129 @@ kernel_cake_dsv4_4ac65f00a4a032c10ecb(const __grid_constant__ CUtensorMap tmap_q " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" : "=r"(*reinterpret_cast(&_tmem_load_1[32])), "=r"(*reinterpret_cast(&_tmem_load_1[33])), "=r"(*reinterpret_cast(&_tmem_load_1[34])), "=r"(*reinterpret_cast(&_tmem_load_1[35])), "=r"(*reinterpret_cast(&_tmem_load_1[36])), "=r"(*reinterpret_cast(&_tmem_load_1[37])), "=r"(*reinterpret_cast(&_tmem_load_1[38])), "=r"(*reinterpret_cast(&_tmem_load_1[39])), "=r"(*reinterpret_cast(&_tmem_load_1[40])), "=r"(*reinterpret_cast(&_tmem_load_1[41])), "=r"(*reinterpret_cast(&_tmem_load_1[42])), "=r"(*reinterpret_cast(&_tmem_load_1[43])), "=r"(*reinterpret_cast(&_tmem_load_1[44])), "=r"(*reinterpret_cast(&_tmem_load_1[45])), "=r"(*reinterpret_cast(&_tmem_load_1[46])), "=r"(*reinterpret_cast(&_tmem_load_1[47])), "=r"(*reinterpret_cast(&_tmem_load_1[48])), "=r"(*reinterpret_cast(&_tmem_load_1[49])), "=r"(*reinterpret_cast(&_tmem_load_1[50])), "=r"(*reinterpret_cast(&_tmem_load_1[51])), "=r"(*reinterpret_cast(&_tmem_load_1[52])), "=r"(*reinterpret_cast(&_tmem_load_1[53])), "=r"(*reinterpret_cast(&_tmem_load_1[54])), "=r"(*reinterpret_cast(&_tmem_load_1[55])), "=r"(*reinterpret_cast(&_tmem_load_1[56])), "=r"(*reinterpret_cast(&_tmem_load_1[57])), "=r"(*reinterpret_cast(&_tmem_load_1[58])), "=r"(*reinterpret_cast(&_tmem_load_1[59])), "=r"(*reinterpret_cast(&_tmem_load_1[60])), "=r"(*reinterpret_cast(&_tmem_load_1[61])), "=r"(*reinterpret_cast(&_tmem_load_1[62])), "=r"(*reinterpret_cast(&_tmem_load_1[63])) : "r"(taddr + 256 + (unsigned int)row_addr + 32)); - int col_base = col_half_1 * 64; - if (my_row_1 < num_heads) { - #pragma unroll - for (int offset = 0; offset < 64; offset += 8) { - { - const float2 _prescale2_0 = {inv_sum * output_scale, inv_sum * output_scale}; - #if __CUDA_ARCH__ >= 1000 - #pragma unroll - for (int _ps = 0; _ps < 4; _ps++) - mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_1[offset])[_ps], _prescale2_0); - #else - #pragma unroll - for (int _ps = 0; _ps < 8; _ps++) - _tmem_load_1[offset + _ps] *= inv_sum * output_scale; - #endif - __nv_bfloat162 _pk[4]; - _pk[0] = __floats2bfloat162_rn(_tmem_load_1[offset + 0], _tmem_load_1[offset + 1]); - _pk[1] = __floats2bfloat162_rn(_tmem_load_1[offset + 2], _tmem_load_1[offset + 3]); - _pk[2] = __floats2bfloat162_rn(_tmem_load_1[offset + 4], _tmem_load_1[offset + 5]); - _pk[3] = __floats2bfloat162_rn(_tmem_load_1[offset + 6], _tmem_load_1[offset + 7]); - *reinterpret_cast(&((__nv_bfloat16*)(O + (output_base + col_base + offset)))[0]) = *reinterpret_cast(&_pk[0]); - } - } + float2 _f2_14 = make_float2(inv_sum * output_scale, inv_sum * output_scale); + #pragma unroll + for (int i_1 = 0; i_1 < 64; i_1 += 2) { + float2 _f2_15 = make_float2(_tmem_load_1[i_1], _tmem_load_1[i_1 + 1]); + float2 _mul_f32x2_0; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_0) : "l"(*(const unsigned long long*)&_f2_15), "l"(*(const unsigned long long*)&_f2_14)); + _tmem_load_1[i_1] = _mul_f32x2_0.x; + _tmem_load_1[i_1 + 1] = _mul_f32x2_0.y; + } + unsigned int packed_o[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_1[_lp*2 + 0], _tmem_load_1[_lp*2+1 + 0])); + packed_o[_lp] = *(uint32_t*)&_bf2; + } + bool stage_hi_1 = (lane & 1) != 0; + bool stage_hi_2 = (lane & 2) != 0; + bool stage_hi_4 = (lane & 4) != 0; + unsigned int _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[0] : packed_o[4]), 1); + packed_o[0] = ((stage_hi_1) ? _shfl_xor_2 : packed_o[0]); + packed_o[4] = ((stage_hi_1) ? packed_o[4] : _shfl_xor_2); + unsigned int _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[1] : packed_o[5]), 1); + packed_o[1] = ((stage_hi_1) ? _shfl_xor_3 : packed_o[1]); + packed_o[5] = ((stage_hi_1) ? packed_o[5] : _shfl_xor_3); + unsigned int _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[2] : packed_o[6]), 1); + packed_o[2] = ((stage_hi_1) ? _shfl_xor_4 : packed_o[2]); + packed_o[6] = ((stage_hi_1) ? packed_o[6] : _shfl_xor_4); + unsigned int _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[3] : packed_o[7]), 1); + packed_o[3] = ((stage_hi_1) ? _shfl_xor_5 : packed_o[3]); + packed_o[7] = ((stage_hi_1) ? packed_o[7] : _shfl_xor_5); + unsigned int _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[8] : packed_o[12]), 1); + packed_o[8] = ((stage_hi_1) ? _shfl_xor_6 : packed_o[8]); + packed_o[12] = ((stage_hi_1) ? packed_o[12] : _shfl_xor_6); + unsigned int _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[9] : packed_o[13]), 1); + packed_o[9] = ((stage_hi_1) ? _shfl_xor_7 : packed_o[9]); + packed_o[13] = ((stage_hi_1) ? packed_o[13] : _shfl_xor_7); + unsigned int _shfl_xor_8 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[10] : packed_o[14]), 1); + packed_o[10] = ((stage_hi_1) ? _shfl_xor_8 : packed_o[10]); + packed_o[14] = ((stage_hi_1) ? packed_o[14] : _shfl_xor_8); + unsigned int _shfl_xor_9 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[11] : packed_o[15]), 1); + packed_o[11] = ((stage_hi_1) ? _shfl_xor_9 : packed_o[11]); + packed_o[15] = ((stage_hi_1) ? packed_o[15] : _shfl_xor_9); + unsigned int _shfl_xor_10 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[16] : packed_o[20]), 1); + packed_o[16] = ((stage_hi_1) ? _shfl_xor_10 : packed_o[16]); + packed_o[20] = ((stage_hi_1) ? packed_o[20] : _shfl_xor_10); + unsigned int _shfl_xor_11 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[17] : packed_o[21]), 1); + packed_o[17] = ((stage_hi_1) ? _shfl_xor_11 : packed_o[17]); + packed_o[21] = ((stage_hi_1) ? packed_o[21] : _shfl_xor_11); + unsigned int _shfl_xor_12 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[18] : packed_o[22]), 1); + packed_o[18] = ((stage_hi_1) ? _shfl_xor_12 : packed_o[18]); + packed_o[22] = ((stage_hi_1) ? packed_o[22] : _shfl_xor_12); + unsigned int _shfl_xor_13 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[19] : packed_o[23]), 1); + packed_o[19] = ((stage_hi_1) ? _shfl_xor_13 : packed_o[19]); + packed_o[23] = ((stage_hi_1) ? packed_o[23] : _shfl_xor_13); + unsigned int _shfl_xor_14 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[24] : packed_o[28]), 1); + packed_o[24] = ((stage_hi_1) ? _shfl_xor_14 : packed_o[24]); + packed_o[28] = ((stage_hi_1) ? packed_o[28] : _shfl_xor_14); + unsigned int _shfl_xor_15 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[25] : packed_o[29]), 1); + packed_o[25] = ((stage_hi_1) ? _shfl_xor_15 : packed_o[25]); + packed_o[29] = ((stage_hi_1) ? packed_o[29] : _shfl_xor_15); + unsigned int _shfl_xor_16 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[26] : packed_o[30]), 1); + packed_o[26] = ((stage_hi_1) ? _shfl_xor_16 : packed_o[26]); + packed_o[30] = ((stage_hi_1) ? packed_o[30] : _shfl_xor_16); + unsigned int _shfl_xor_17 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[27] : packed_o[31]), 1); + packed_o[27] = ((stage_hi_1) ? _shfl_xor_17 : packed_o[27]); + packed_o[31] = ((stage_hi_1) ? packed_o[31] : _shfl_xor_17); + int lane_t = lane & 1; + unsigned int chunk_words[4]; + chunk_words[0] = packed_o[0]; + chunk_words[1] = packed_o[1]; + chunk_words[2] = packed_o[2]; + chunk_words[3] = packed_o[3]; + if ((my_row_1 & -2 | 0) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + (my_row_1 & -2 | 0)) * 512 + v_chunk * 128 + col_half_1 * 64 + (0 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[4]; + chunk_words[1] = packed_o[5]; + chunk_words[2] = packed_o[6]; + chunk_words[3] = packed_o[7]; + if ((my_row_1 & -2 | 1) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + (my_row_1 & -2 | 1)) * 512 + v_chunk * 128 + col_half_1 * 64 + (0 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[8]; + chunk_words[1] = packed_o[9]; + chunk_words[2] = packed_o[10]; + chunk_words[3] = packed_o[11]; + if ((my_row_1 & -2 | 0) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + (my_row_1 & -2 | 0)) * 512 + v_chunk * 128 + col_half_1 * 64 + (2 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[12]; + chunk_words[1] = packed_o[13]; + chunk_words[2] = packed_o[14]; + chunk_words[3] = packed_o[15]; + if ((my_row_1 & -2 | 1) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + (my_row_1 & -2 | 1)) * 512 + v_chunk * 128 + col_half_1 * 64 + (2 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[16]; + chunk_words[1] = packed_o[17]; + chunk_words[2] = packed_o[18]; + chunk_words[3] = packed_o[19]; + if ((my_row_1 & -2 | 0) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + (my_row_1 & -2 | 0)) * 512 + v_chunk * 128 + col_half_1 * 64 + (4 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[20]; + chunk_words[1] = packed_o[21]; + chunk_words[2] = packed_o[22]; + chunk_words[3] = packed_o[23]; + if ((my_row_1 & -2 | 1) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + (my_row_1 & -2 | 1)) * 512 + v_chunk * 128 + col_half_1 * 64 + (4 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[24]; + chunk_words[1] = packed_o[25]; + chunk_words[2] = packed_o[26]; + chunk_words[3] = packed_o[27]; + if ((my_row_1 & -2 | 0) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + (my_row_1 & -2 | 0)) * 512 + v_chunk * 128 + col_half_1 * 64 + (6 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[28]; + chunk_words[1] = packed_o[29]; + chunk_words[2] = packed_o[30]; + chunk_words[3] = packed_o[31]; + if ((my_row_1 & -2 | 1) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + (my_row_1 & -2 | 1)) * 512 + v_chunk * 128 + col_half_1 * 64 + (6 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; } asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); asm volatile("tcgen05.fence::before_thread_sync;"); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_9526f80c89f2b5ef4c6c_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_15eabf5d2a5f50db201b_binding.cu similarity index 97% rename from csrc/cake_dsv4/sm_100a/cake_dsv4_9526f80c89f2b5ef4c6c_binding.cu rename to csrc/cake_dsv4/sm_100a/cake_dsv4_15eabf5d2a5f50db201b_binding.cu index 931e515eb10..d58a59b6a60 100644 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_9526f80c89f2b5ef4c6c_binding.cu +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_15eabf5d2a5f50db201b_binding.cu @@ -15,7 +15,7 @@ */ // Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_9526f80c89f2b5ef4c6c'. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_15eabf5d2a5f50db201b'. #include #include #include @@ -31,10 +31,10 @@ #include #include -extern "C" __global__ void kernel_cake_dsv4_9526f80c89f2b5ef4c6c(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int num_query_tokens, int has_sinks); +extern "C" __global__ void kernel_cake_dsv4_15eabf5d2a5f50db201b(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int num_query_tokens, int has_sinks); -namespace cake_host_shim_638caaf740c22a7d { +namespace cake_host_shim_6c292f2d365fae68 { using tvm::ffi::Optional; using tvm::ffi::TensorView; @@ -404,7 +404,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ } if (smem_status == cudaSuccess) { smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_9526f80c89f2b5ef4c6c), + reinterpret_cast(kernel_cake_dsv4_15eabf5d2a5f50db201b), cudaFuncAttributeMaxDynamicSharedMemorySize, 203904); } @@ -414,7 +414,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ } } TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_9526f80c89f2b5ef4c6c failed: " + << "cudaFuncSetAttribute for kernel_cake_dsv4_15eabf5d2a5f50db201b failed: " << cudaGetErrorString(smem_status); dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); dim3 block(512u, 1u, 1u); @@ -425,12 +425,12 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ config.dynamicSmemBytes = 203904u; config.stream = stream; cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_9526f80c89f2b5ef4c6c), kargs); + &config, reinterpret_cast(kernel_cake_dsv4_15eabf5d2a5f50db201b), kargs); TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_9526f80c89f2b5ef4c6c failed: " + << "cudaLaunchKernelExC for kernel_cake_dsv4_15eabf5d2a5f50db201b failed: " << cudaGetErrorString(launch_status); } -} // namespace cake_host_shim_638caaf740c22a7d +} // namespace cake_host_shim_6c292f2d365fae68 -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_638caaf740c22a7d::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_6c292f2d365fae68::Run); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_9526f80c89f2b5ef4c6c_kernel.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_15eabf5d2a5f50db201b_kernel.cu similarity index 98% rename from csrc/cake_dsv4/sm_100a/cake_dsv4_9526f80c89f2b5ef4c6c_kernel.cu rename to csrc/cake_dsv4/sm_100a/cake_dsv4_15eabf5d2a5f50db201b_kernel.cu index b17e9aa149f..265088b7d34 100644 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_9526f80c89f2b5ef4c6c_kernel.cu +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_15eabf5d2a5f50db201b_kernel.cu @@ -1082,7 +1082,7 @@ __device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { extern "C" { __global__ __launch_bounds__(512, 1) void -kernel_cake_dsv4_9526f80c89f2b5ef4c6c(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int num_query_tokens, int has_sinks) +kernel_cake_dsv4_15eabf5d2a5f50db201b(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int num_query_tokens, int has_sinks) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); @@ -1640,24 +1640,39 @@ kernel_cake_dsv4_9526f80c89f2b5ef4c6c(const __grid_constant__ CUtensorMap tmap_q int col_base = v_stage_1 * 128 + col_half_1 * 64; if (my_row_1 < num_heads) { #pragma unroll - for (int offset = 0; offset < 64; offset += 8) { + for (int offset = 0; offset < 64; offset += 16) { { const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; #if __CUDA_ARCH__ >= 1000 #pragma unroll - for (int _ps = 0; _ps < 4; _ps++) + for (int _ps = 0; _ps < 8; _ps++) mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_4[offset])[_ps], _prescale2_1); #else #pragma unroll - for (int _ps = 0; _ps < 8; _ps++) + for (int _ps = 0; _ps < 16; _ps++) _tmem_load_4[offset + _ps] *= inv_sum * output_scale; #endif - __nv_bfloat162 _pk[4]; - _pk[0] = __floats2bfloat162_rn(_tmem_load_4[offset + 0], _tmem_load_4[offset + 1]); - _pk[1] = __floats2bfloat162_rn(_tmem_load_4[offset + 2], _tmem_load_4[offset + 3]); - _pk[2] = __floats2bfloat162_rn(_tmem_load_4[offset + 4], _tmem_load_4[offset + 5]); - _pk[3] = __floats2bfloat162_rn(_tmem_load_4[offset + 6], _tmem_load_4[offset + 7]); - *reinterpret_cast(&((__nv_bfloat16*)(O + (output_base + col_base + offset)))[0]) = *reinterpret_cast(&_pk[0]); + { + __nv_bfloat162 _pk0 = __floats2bfloat162_rn(_tmem_load_4[offset + 0], _tmem_load_4[offset + 1]); + unsigned _pk_u0 = *reinterpret_cast(&_pk0); + __nv_bfloat162 _pk1 = __floats2bfloat162_rn(_tmem_load_4[offset + 2], _tmem_load_4[offset + 3]); + unsigned _pk_u1 = *reinterpret_cast(&_pk1); + __nv_bfloat162 _pk2 = __floats2bfloat162_rn(_tmem_load_4[offset + 4], _tmem_load_4[offset + 5]); + unsigned _pk_u2 = *reinterpret_cast(&_pk2); + __nv_bfloat162 _pk3 = __floats2bfloat162_rn(_tmem_load_4[offset + 6], _tmem_load_4[offset + 7]); + unsigned _pk_u3 = *reinterpret_cast(&_pk3); + __nv_bfloat162 _pk4 = __floats2bfloat162_rn(_tmem_load_4[offset + 8], _tmem_load_4[offset + 9]); + unsigned _pk_u4 = *reinterpret_cast(&_pk4); + __nv_bfloat162 _pk5 = __floats2bfloat162_rn(_tmem_load_4[offset + 10], _tmem_load_4[offset + 11]); + unsigned _pk_u5 = *reinterpret_cast(&_pk5); + __nv_bfloat162 _pk6 = __floats2bfloat162_rn(_tmem_load_4[offset + 12], _tmem_load_4[offset + 13]); + unsigned _pk_u6 = *reinterpret_cast(&_pk6); + __nv_bfloat162 _pk7 = __floats2bfloat162_rn(_tmem_load_4[offset + 14], _tmem_load_4[offset + 15]); + unsigned _pk_u7 = *reinterpret_cast(&_pk7); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(&((__nv_bfloat16*)(O + (output_base + col_base + offset)))[0])), "r"(_pk_u0), "r"(_pk_u1), "r"(_pk_u2), "r"(_pk_u3), "r"(_pk_u4), "r"(_pk_u5), "r"(_pk_u6), "r"(_pk_u7) : "memory"); + } } } } diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_fbf138c989dfde9f764f_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_21a93d40caefcf5b252d_binding.cu similarity index 97% rename from csrc/cake_dsv4/sm_100a/cake_dsv4_fbf138c989dfde9f764f_binding.cu rename to csrc/cake_dsv4/sm_100a/cake_dsv4_21a93d40caefcf5b252d_binding.cu index e832234c20a..5accf360b11 100644 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_fbf138c989dfde9f764f_binding.cu +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_21a93d40caefcf5b252d_binding.cu @@ -15,7 +15,7 @@ */ // Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_fbf138c989dfde9f764f'. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_21a93d40caefcf5b252d'. #include #include #include @@ -31,10 +31,10 @@ #include #include -extern "C" __global__ void kernel_cake_dsv4_fbf138c989dfde9f764f(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int max_q_len, int batch_size, int has_sinks); +extern "C" __global__ void kernel_cake_dsv4_21a93d40caefcf5b252d(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int max_q_len, int batch_size, int has_sinks); -namespace cake_host_shim_e92a6da8d9e51169 { +namespace cake_host_shim_443c43e36fe3b626 { using tvm::ffi::Optional; using tvm::ffi::TensorView; @@ -414,7 +414,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ } if (smem_status == cudaSuccess) { smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_fbf138c989dfde9f764f), + reinterpret_cast(kernel_cake_dsv4_21a93d40caefcf5b252d), cudaFuncAttributeMaxDynamicSharedMemorySize, 147072); } @@ -424,7 +424,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ } } TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_fbf138c989dfde9f764f failed: " + << "cudaFuncSetAttribute for kernel_cake_dsv4_21a93d40caefcf5b252d failed: " << cudaGetErrorString(smem_status); dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); dim3 block(640u, 1u, 1u); @@ -435,12 +435,12 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ config.dynamicSmemBytes = 147072u; config.stream = stream; cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_fbf138c989dfde9f764f), kargs); + &config, reinterpret_cast(kernel_cake_dsv4_21a93d40caefcf5b252d), kargs); TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_fbf138c989dfde9f764f failed: " + << "cudaLaunchKernelExC for kernel_cake_dsv4_21a93d40caefcf5b252d failed: " << cudaGetErrorString(launch_status); } -} // namespace cake_host_shim_e92a6da8d9e51169 +} // namespace cake_host_shim_443c43e36fe3b626 -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_e92a6da8d9e51169::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_443c43e36fe3b626::Run); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_fbf138c989dfde9f764f_kernel.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_21a93d40caefcf5b252d_kernel.cu similarity index 95% rename from csrc/cake_dsv4/sm_100a/cake_dsv4_fbf138c989dfde9f764f_kernel.cu rename to csrc/cake_dsv4/sm_100a/cake_dsv4_21a93d40caefcf5b252d_kernel.cu index f09dcbb858b..bc691ca97b1 100644 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_fbf138c989dfde9f764f_kernel.cu +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_21a93d40caefcf5b252d_kernel.cu @@ -998,7 +998,7 @@ __device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { extern "C" { __global__ __launch_bounds__(640, 1) void -kernel_cake_dsv4_fbf138c989dfde9f764f(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int max_q_len, int batch_size, int has_sinks) +kernel_cake_dsv4_21a93d40caefcf5b252d(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int max_q_len, int batch_size, int has_sinks) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); @@ -1206,8 +1206,8 @@ kernel_cake_dsv4_fbf138c989dfde9f764f(const __grid_constant__ CUtensorMap tmap_q int active_topk = 0; if (work_x < (unsigned int)q_len) { int _max_4 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); - int _min_3 = ((_max_4) < (sparse_topk) ? (_max_4) : (sparse_topk)); - active_topk = _min_3; + int _min_7 = ((_max_4) < (sparse_topk) ? (_max_4) : (sparse_topk)); + active_topk = _min_7; } int num_steps = (active_topk + 128 - 1) / 128; float row_max_pair[2]; @@ -1534,8 +1534,8 @@ kernel_cake_dsv4_fbf138c989dfde9f764f(const __grid_constant__ CUtensorMap tmap_q int active_topk_1 = 0; if (work_x_1 < (unsigned int)q_len_1) { int _max_14 = ((sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) : (0)); - int _min_4 = ((_max_14) < (sparse_topk) ? (_max_14) : (sparse_topk)); - active_topk_1 = _min_4; + int _min_8 = ((_max_14) < (sparse_topk) ? (_max_14) : (sparse_topk)); + active_topk_1 = _min_8; } int num_steps_1 = (active_topk_1 + 128 - 1) / 128; int head_group = work_y_1 / 4; @@ -1798,8 +1798,8 @@ kernel_cake_dsv4_fbf138c989dfde9f764f(const __grid_constant__ CUtensorMap tmap_q int active_topk_2 = 0; if (work_x_2 < (unsigned int)q_len_2) { int _max_3 = ((sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) : (0)); - int _min_2 = ((_max_3) < (sparse_topk) ? (_max_3) : (sparse_topk)); - active_topk_2 = _min_2; + int _min_6 = ((_max_3) < (sparse_topk) ? (_max_3) : (sparse_topk)); + active_topk_2 = _min_6; } int num_steps_2 = (active_topk_2 + 128 - 1) / 128; int value_quarter_1 = work_y_2 % 4; @@ -2434,30 +2434,45 @@ kernel_cake_dsv4_fbf138c989dfde9f764f(const __grid_constant__ CUtensorMap tmap_q } mbarrier_wait(page_full_addr + (page_k_stage) * 8, page_k_phase); int page_base = 0; + int warps_tile_idx = 0; + int _min_2 = ((active_topk_4) < (128) ? (active_topk_4) : (128)); + int warps_valid_rows = _min_2; + int _min_3 = (((warps_valid_rows + 3) / 4) < (8) ? ((warps_valid_rows + 3) / 4) : (8)); + int warps_tile = _min_3; + int warps_ge1 = 8; + int warps_ge2 = 8; + if (warps_tile_idx >= 1) { + warps_ge1 = warps_tile; + } + if (warps_tile_idx >= 2) { + warps_ge2 = warps_tile; + } #pragma unroll for (int hs_3 = 0; hs_3 < 4; hs_3++) { int hs_sel_2 = ((hs_3 == 3) ? value_quarter_2 : ((value_quarter_2 > hs_3) ? hs_3 : hs_3 + 1)); mbarrier_wait(kv_empty_addr + (kv_stage_1) * 8, kv_phase_1); if (load_rank == 0) { if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_1) * 8, 16384); + mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_1) * 8, ((hs_3 >= 4) ? warps_ge1 : warps_ge2) * 2048); } } asm volatile("barrier.sync 9, 256;" ::: "memory"); - int rows_all[16]; - #pragma unroll - for (int group = 0; group < 4; group++) { - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&rows_all[group * 4])), "=r"(*reinterpret_cast(&rows_all[(group * 4) + 1])), "=r"(*reinterpret_cast(&rows_all[(group * 4) + 2])), "=r"(*reinterpret_cast(&rows_all[(group * 4) + 3])) - : "r"(smem_page_offsets_addr + page_k_stage * 1024 + (unsigned int)((page_base + load_rank * 4 + group * 32) * 4))); - } - if (elect_sync()) { - { - #pragma unroll - for (int group_1 = 0; group_1 < 4; group_1++) { - const int row_base = load_rank * 4 + group_1 * 32; - int dst_k = smem_kv_addr + kv_stage_1 * 16384 + (unsigned int)(row_base * 128); - tma_gather4_gmem2smem(dst_k, (&tmap_swa_kv), hs_sel_2 * 128, rows_all[group_1 * 4], rows_all[group_1 * 4 + 1], rows_all[group_1 * 4 + 2], rows_all[group_1 * 4 + 3], kv_full_addr + (kv_stage_1) * 8); + if (load_rank < ((hs_3 >= 4) ? warps_ge1 : warps_ge2)) { + int rows_all[16]; + #pragma unroll + for (int group = 0; group < 4; group++) { + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&rows_all[group * 4])), "=r"(*reinterpret_cast(&rows_all[(group * 4) + 1])), "=r"(*reinterpret_cast(&rows_all[(group * 4) + 2])), "=r"(*reinterpret_cast(&rows_all[(group * 4) + 3])) + : "r"(smem_page_offsets_addr + page_k_stage * 1024 + (unsigned int)((page_base + load_rank * 4 + group * 32) * 4))); + } + if (elect_sync()) { + { + #pragma unroll + for (int group_1 = 0; group_1 < 4; group_1++) { + const int row_base = load_rank * 4 + group_1 * 32; + int dst_k = smem_kv_addr + kv_stage_1 * 16384 + (unsigned int)(row_base * 128); + tma_gather4_gmem2smem(dst_k, (&tmap_swa_kv), hs_sel_2 * 128, rows_all[group_1 * 4], rows_all[group_1 * 4 + 1], rows_all[group_1 * 4 + 2], rows_all[group_1 * 4 + 3], kv_full_addr + (kv_stage_1) * 8); + } } } } @@ -2471,37 +2486,52 @@ kernel_cake_dsv4_fbf138c989dfde9f764f(const __grid_constant__ CUtensorMap tmap_q mbarrier_wait(page_full_addr + (page_k_stage) * 8, page_k_phase); } int page_base_0 = (next_tile & 1) * 128; + int warps_tile_idx_1 = next_tile; + int _min_4 = ((active_topk_4 - next_tile * 128) < (128) ? (active_topk_4 - next_tile * 128) : (128)); + int warps_valid_rows_2 = _min_4; + int _min_5 = (((warps_valid_rows_2 + 3) / 4) < (8) ? ((warps_valid_rows_2 + 3) / 4) : (8)); + int warps_tile_3 = _min_5; + int warps_ge1_4 = 8; + int warps_ge2_5 = 8; + if (warps_tile_idx_1 >= 1) { + warps_ge1_4 = warps_tile_3; + } + if (warps_tile_idx_1 >= 2) { + warps_ge2_5 = warps_tile_3; + } #pragma unroll for (int hs_4 = 0; hs_4 < 4; hs_4++) { int hs_sel_3 = ((hs_4 == 3) ? value_quarter_2 : ((value_quarter_2 > hs_4) ? hs_4 : hs_4 + 1)); mbarrier_wait(kv_empty_addr + (kv_stage_1) * 8, kv_phase_1); if (load_rank == 0) { if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_1) * 8, 16384); + mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_1) * 8, ((hs_4 >= 4) ? warps_ge1_4 : warps_ge2_5) * 2048); } } asm volatile("barrier.sync 9, 256;" ::: "memory"); - int rows_all_1[16]; - #pragma unroll - for (int group_2 = 0; group_2 < 4; group_2++) { - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&rows_all_1[group_2 * 4])), "=r"(*reinterpret_cast(&rows_all_1[(group_2 * 4) + 1])), "=r"(*reinterpret_cast(&rows_all_1[(group_2 * 4) + 2])), "=r"(*reinterpret_cast(&rows_all_1[(group_2 * 4) + 3])) - : "r"(smem_page_offsets_addr + page_k_stage * 1024 + (unsigned int)((page_base_0 + load_rank * 4 + group_2 * 32) * 4))); - } - if (elect_sync()) { - if (next_tile == 0) { - #pragma unroll - for (int group_3 = 0; group_3 < 4; group_3++) { - const int row_base_1 = load_rank * 4 + group_3 * 32; - int dst_k_1 = smem_kv_addr + kv_stage_1 * 16384 + (unsigned int)(row_base_1 * 128); - tma_gather4_gmem2smem(dst_k_1, (&tmap_swa_kv), hs_sel_3 * 128, rows_all_1[group_3 * 4], rows_all_1[group_3 * 4 + 1], rows_all_1[group_3 * 4 + 2], rows_all_1[group_3 * 4 + 3], kv_full_addr + (kv_stage_1) * 8); - } - } else { - #pragma unroll - for (int group_4 = 0; group_4 < 4; group_4++) { - const int row_base_2 = load_rank * 4 + group_4 * 32; - int dst_k_2 = smem_kv_addr + kv_stage_1 * 16384 + (unsigned int)(row_base_2 * 128); - tma_gather4_gmem2smem(dst_k_2, (&tmap_compressed_kv), hs_sel_3 * 128, rows_all_1[group_4 * 4], rows_all_1[group_4 * 4 + 1], rows_all_1[group_4 * 4 + 2], rows_all_1[group_4 * 4 + 3], kv_full_addr + (kv_stage_1) * 8); + if (load_rank < ((hs_4 >= 4) ? warps_ge1_4 : warps_ge2_5)) { + int rows_all_1[16]; + #pragma unroll + for (int group_2 = 0; group_2 < 4; group_2++) { + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&rows_all_1[group_2 * 4])), "=r"(*reinterpret_cast(&rows_all_1[(group_2 * 4) + 1])), "=r"(*reinterpret_cast(&rows_all_1[(group_2 * 4) + 2])), "=r"(*reinterpret_cast(&rows_all_1[(group_2 * 4) + 3])) + : "r"(smem_page_offsets_addr + page_k_stage * 1024 + (unsigned int)((page_base_0 + load_rank * 4 + group_2 * 32) * 4))); + } + if (elect_sync()) { + if (next_tile == 0) { + #pragma unroll + for (int group_3 = 0; group_3 < 4; group_3++) { + const int row_base_1 = load_rank * 4 + group_3 * 32; + int dst_k_1 = smem_kv_addr + kv_stage_1 * 16384 + (unsigned int)(row_base_1 * 128); + tma_gather4_gmem2smem(dst_k_1, (&tmap_swa_kv), hs_sel_3 * 128, rows_all_1[group_3 * 4], rows_all_1[group_3 * 4 + 1], rows_all_1[group_3 * 4 + 2], rows_all_1[group_3 * 4 + 3], kv_full_addr + (kv_stage_1) * 8); + } + } else { + #pragma unroll + for (int group_4 = 0; group_4 < 4; group_4++) { + const int row_base_2 = load_rank * 4 + group_4 * 32; + int dst_k_2 = smem_kv_addr + kv_stage_1 * 16384 + (unsigned int)(row_base_2 * 128); + tma_gather4_gmem2smem(dst_k_2, (&tmap_compressed_kv), hs_sel_3 * 128, rows_all_1[group_4 * 4], rows_all_1[group_4 * 4 + 1], rows_all_1[group_4 * 4 + 2], rows_all_1[group_4 * 4 + 3], kv_full_addr + (kv_stage_1) * 8); + } } } } diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_e1660a7f44d156286fea_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_42fa8e98da5fa46d1a94_binding.cu similarity index 86% rename from csrc/cake_dsv4/sm_100a/cake_dsv4_e1660a7f44d156286fea_binding.cu rename to csrc/cake_dsv4/sm_100a/cake_dsv4_42fa8e98da5fa46d1a94_binding.cu index d3bccd7047a..214a8b9199b 100644 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_e1660a7f44d156286fea_binding.cu +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_42fa8e98da5fa46d1a94_binding.cu @@ -15,7 +15,7 @@ */ // Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_e1660a7f44d156286fea'. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_42fa8e98da5fa46d1a94'. #include #include #include @@ -31,10 +31,10 @@ #include #include -extern "C" __global__ void kernel_cake_dsv4_e1660a7f44d156286fea(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); +extern "C" __global__ void kernel_cake_dsv4_42fa8e98da5fa46d1a94(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, const __grid_constant__ CUtensorMap O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); -namespace cake_host_shim_244c86188c16b5f1 { +namespace cake_host_shim_c0e0335c8d570682 { using tvm::ffi::Optional; using tvm::ffi::TensorView; @@ -390,6 +390,57 @@ inline CUtensorMap EncodeTma_tmap_compressed_v(const TensorView& t) { return tm; } +// 4D TMA descriptor for buffer 'O' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_O(const TensorView& t, int64_t tp_num_heads, int64_t tp_num_query_tokens) { + TVM_FFI_CHECK(t.ndim() >= 1, ValueError) + << "TMA source 'O' must have at least 1 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'O' must have unit innermost stride, got " << t.stride(-1); + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(8), (uint64_t)(4), (uint64_t)((tp_num_query_tokens * tp_num_heads))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'O' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1] && 1u <= global_dim[2] && 32u <= global_dim[3], ValueError) + << "TMA box (64, 1, 1, 32) exceeds resolved global dims for 'O'"; + int64_t carrier_stride_0 = 64; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 512; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = 2048; + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 1u, 1u, 32u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'O') failed: CUresult=" << (int)r; + return tm; +} + void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_compressed_k, TensorView arg_tmap_swa_v, TensorView arg_tmap_compressed_v, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_num_query_tokens, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z) { DLDevice dev = arg_tmap_q.device(); ScopedCudaDevice device_guard(dev.device_id); @@ -477,7 +528,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_c CUtensorMap p_tmap_compressed_k = EncodeTma_tmap_compressed_k(arg_tmap_compressed_k); CUtensorMap p_tmap_swa_v = EncodeTma_tmap_swa_v(arg_tmap_swa_v); CUtensorMap p_tmap_compressed_v = EncodeTma_tmap_compressed_v(arg_tmap_compressed_v); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); + CUtensorMap p_O = EncodeTma_O(arg_O, (int64_t)arg_num_heads, (int64_t)arg_num_query_tokens); float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); @@ -510,7 +561,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_c } if (smem_status == cudaSuccess) { smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_e1660a7f44d156286fea), + reinterpret_cast(kernel_cake_dsv4_42fa8e98da5fa46d1a94), cudaFuncAttributeMaxDynamicSharedMemorySize, 220032); } @@ -520,7 +571,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_c } } TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_e1660a7f44d156286fea failed: " + << "cudaFuncSetAttribute for kernel_cake_dsv4_42fa8e98da5fa46d1a94 failed: " << cudaGetErrorString(smem_status); dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); dim3 block(512u, 1u, 1u); @@ -544,12 +595,12 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_c config.attrs = attrs; config.numAttrs = n; cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_e1660a7f44d156286fea), kargs); + &config, reinterpret_cast(kernel_cake_dsv4_42fa8e98da5fa46d1a94), kargs); TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_e1660a7f44d156286fea failed: " + << "cudaLaunchKernelExC for kernel_cake_dsv4_42fa8e98da5fa46d1a94 failed: " << cudaGetErrorString(launch_status); } -} // namespace cake_host_shim_244c86188c16b5f1 +} // namespace cake_host_shim_c0e0335c8d570682 -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_244c86188c16b5f1::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_c0e0335c8d570682::Run); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_e1660a7f44d156286fea_kernel.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_42fa8e98da5fa46d1a94_kernel.cu similarity index 96% rename from csrc/cake_dsv4/sm_100a/cake_dsv4_e1660a7f44d156286fea_kernel.cu rename to csrc/cake_dsv4/sm_100a/cake_dsv4_42fa8e98da5fa46d1a94_kernel.cu index 5c3bf235129..a1879057902 100644 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_e1660a7f44d156286fea_kernel.cu +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_42fa8e98da5fa46d1a94_kernel.cu @@ -1181,6 +1181,15 @@ __device__ __forceinline__ void tma_gather4_gmem2smem_mc_cta2( } +__device__ __forceinline__ void tma_store_4d( + const void *tmap, int x, int y, int z, int w, unsigned smem_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.global.shared::cta.tile.bulk_group" + " [%0, {%1, %2, %3, %4}], [%5];" + :: "l"(tmap), "r"(x), "r"(y), "r"(z), "r"(w), "r"(smem_addr) : "memory"); +} + + __device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { asm volatile( "{\n\t" @@ -1195,7 +1204,7 @@ __device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint extern "C" { __global__ __launch_bounds__(512, 1) __cluster_dims__(2,1,1) void -kernel_cake_dsv4_e1660a7f44d156286fea(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items) +kernel_cake_dsv4_42fa8e98da5fa46d1a94(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, const __grid_constant__ CUtensorMap O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); @@ -1221,10 +1230,11 @@ kernel_cake_dsv4_e1660a7f44d156286fea(const __grid_constant__ CUtensorMap tmap_q #define stats_addr (mbar_base + 168) #define stats_empty_addr (mbar_base + 184) #define o_full_addr (mbar_base + 200) - #define o_first_slice_seeded_addr (mbar_base + 208) - #define tmem_dealloc_peer_addr (mbar_base + 216) - #define index_valid_full_addr (mbar_base + 224) - #define index_valid_empty_addr (mbar_base + 240) + #define epi_q_free_addr (mbar_base + 208) + #define o_first_slice_seeded_addr (mbar_base + 216) + #define tmem_dealloc_peer_addr (mbar_base + 224) + #define index_valid_full_addr (mbar_base + 232) + #define index_valid_empty_addr (mbar_base + 248) const int bid = blockIdx.x; const int num_bids = gridDim.x; @@ -1257,8 +1267,8 @@ kernel_cake_dsv4_e1660a7f44d156286fea(const __grid_constant__ CUtensorMap tmap_q asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); - // Mbarrier init (18 pipeline groups, 0 ordered-sequence groups, 32 barriers) - // Mbarriers at smem_raw[0..256) + // Mbarrier init (19 pipeline groups, 0 ordered-sequence groups, 33 barriers) + // Mbarriers at smem_raw[0..264) if (warp == 0) { uint32_t leader = elect_sync(); @@ -1305,16 +1315,18 @@ kernel_cake_dsv4_e1660a7f44d156286fea(const __grid_constant__ CUtensorMap tmap_q mbarrier_init(smem + 192, 128); // o_full: 1 barriers, init_count=1 mbarrier_init(smem + 200, 1); + // epi_q_free: 1 barriers, init_count=4 + mbarrier_init(smem + 208, 4); // o_first_slice_seeded: 1 barriers, init_count=256 - mbarrier_init(smem + 208, 256); + mbarrier_init(smem + 216, 256); // tmem_dealloc_peer: 1 barriers, init_count=32 - mbarrier_init(smem + 216, 32); - // index_valid_full: 2 barriers, init_count=32 mbarrier_init(smem + 224, 32); + // index_valid_full: 2 barriers, init_count=32 mbarrier_init(smem + 232, 32); + mbarrier_init(smem + 240, 32); // index_valid_empty: 2 barriers, init_count=128 - mbarrier_init(smem + 240, 128); mbarrier_init(smem + 248, 128); + mbarrier_init(smem + 256, 128); asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); } } @@ -1323,9 +1335,9 @@ kernel_cake_dsv4_e1660a7f44d156286fea(const __grid_constant__ CUtensorMap tmap_q asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); // TMEM alloc (512 columns, 512 used) - volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 256); + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 264); if (warp == 0) { - int _tmem_hold = smem + 256; + int _tmem_hold = smem + 264; asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); __syncwarp(); asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); @@ -1697,11 +1709,15 @@ kernel_cake_dsv4_e1660a7f44d156286fea(const __grid_constant__ CUtensorMap tmap_q asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum)); partial_lse[stat_offset] = ((final_sum > 0.0f) ? final_max * softmax_scale_log2_1 + _log2_0 : -CAKE_INF); } + float epi_scale = inv_sum * output_scale; + int epi_box_base = smem_q_addr + (unsigned int)(warp % 4 * 16384); + int epi_row0 = query_idx_1 * num_heads + cta_rank * 64 + head_row_base_1; #pragma unroll for (int acc_stage_1 = 0; acc_stage_1 < 2; acc_stage_1++) { #pragma unroll for (int local_slice_1 = 0; local_slice_1 < 2; local_slice_1++) { int logical_slice = acc_stage_1 * 2 * 2 + n_half_1 * 2 + local_slice_1; + int epi_box = epi_box_base + (acc_stage_1 * 2 + local_slice_1) * 4096; int o_base_1 = taddr + 128 + (unsigned int)(acc_stage_1 * 128) + (unsigned int)(local_slice_1 * 64) + (unsigned int)(output_row_base << 16); #pragma unroll for (int c_1 = 0; c_1 < 64; c_1 += 32) { @@ -1711,46 +1727,35 @@ kernel_cake_dsv4_e1660a7f44d156286fea(const __grid_constant__ CUtensorMap tmap_q " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]) : "r"(o_base_1 + c_1)); - int gmem_base = output_offset + logical_slice * 64 + c_1; + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + float epi_scaled[32]; #pragma unroll - for (int j = 0; j < 32; j += 16) { - { - const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; - #if __CUDA_ARCH__ >= 1000 - #pragma unroll - for (int _ps = 0; _ps < 8; _ps++) - mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_3[j])[_ps], _prescale2_1); - #else - #pragma unroll - for (int _ps = 0; _ps < 16; _ps++) - _tmem_load_3[j + _ps] *= inv_sum * output_scale; - #endif - { - __nv_bfloat162 _pk0 = __floats2bfloat162_rn(_tmem_load_3[j + 0], _tmem_load_3[j + 1]); - unsigned _pk_u0 = *reinterpret_cast(&_pk0); - __nv_bfloat162 _pk1 = __floats2bfloat162_rn(_tmem_load_3[j + 2], _tmem_load_3[j + 3]); - unsigned _pk_u1 = *reinterpret_cast(&_pk1); - __nv_bfloat162 _pk2 = __floats2bfloat162_rn(_tmem_load_3[j + 4], _tmem_load_3[j + 5]); - unsigned _pk_u2 = *reinterpret_cast(&_pk2); - __nv_bfloat162 _pk3 = __floats2bfloat162_rn(_tmem_load_3[j + 6], _tmem_load_3[j + 7]); - unsigned _pk_u3 = *reinterpret_cast(&_pk3); - __nv_bfloat162 _pk4 = __floats2bfloat162_rn(_tmem_load_3[j + 8], _tmem_load_3[j + 9]); - unsigned _pk_u4 = *reinterpret_cast(&_pk4); - __nv_bfloat162 _pk5 = __floats2bfloat162_rn(_tmem_load_3[j + 10], _tmem_load_3[j + 11]); - unsigned _pk_u5 = *reinterpret_cast(&_pk5); - __nv_bfloat162 _pk6 = __floats2bfloat162_rn(_tmem_load_3[j + 12], _tmem_load_3[j + 13]); - unsigned _pk_u6 = *reinterpret_cast(&_pk6); - __nv_bfloat162 _pk7 = __floats2bfloat162_rn(_tmem_load_3[j + 14], _tmem_load_3[j + 15]); - unsigned _pk_u7 = *reinterpret_cast(&_pk7); - asm volatile( - "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" - :: "l"((void*)(&((__nv_bfloat16*)(O + (gmem_base + j)))[0])), "r"(_pk_u0), "r"(_pk_u1), "r"(_pk_u2), "r"(_pk_u3), "r"(_pk_u4), "r"(_pk_u5), "r"(_pk_u6), "r"(_pk_u7) : "memory"); - } - } + for (int e = 0; e < 32; e++) { + epi_scaled[e] = _tmem_load_3[e] * epi_scale; + } + unsigned int epi_packed[16]; + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(epi_scaled[_lp*2 + 0], epi_scaled[_lp*2+1 + 0])); + epi_packed[_lp] = *(uint32_t*)&_bf2; } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((epi_box + (lane * 128 + c_1 * 2 ^ (lane * 128 + c_1 * 2 >> 7 & 7) << 4))), "r"(epi_packed[0]), "r"(epi_packed[1]), "r"(epi_packed[2]), "r"(epi_packed[3]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((epi_box + (lane * 128 + (c_1 * 2 + 16) ^ (lane * 128 + (c_1 * 2 + 16) >> 7 & 7) << 4))), "r"(epi_packed[4]), "r"(epi_packed[5]), "r"(epi_packed[6]), "r"(epi_packed[7]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((epi_box + (lane * 128 + (c_1 * 2 + 32) ^ (lane * 128 + (c_1 * 2 + 32) >> 7 & 7) << 4))), "r"(epi_packed[8]), "r"(epi_packed[9]), "r"(epi_packed[10]), "r"(epi_packed[11]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((epi_box + (lane * 128 + (c_1 * 2 + 48) ^ (lane * 128 + (c_1 * 2 + 48) >> 7 & 7) << 4))), "r"(epi_packed[12]), "r"(epi_packed[13]), "r"(epi_packed[14]), "r"(epi_packed[15]) : "memory"); + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + __syncwarp(); + if (elect_sync()) { + tma_store_4d((&O), 0, logical_slice, split_idx_1, epi_row0, epi_box); + asm volatile("cp.async.bulk.commit_group;"); } } } + if (elect_sync()) { + asm volatile("cp.async.bulk.wait_group.read 0;"); + mbarrier_arrive(epi_q_free_addr); + } asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); asm volatile( "{\n\t" @@ -1761,6 +1766,9 @@ kernel_cake_dsv4_e1660a7f44d156286fea(const __grid_constant__ CUtensorMap tmap_q :: "r"(o_empty_addr), "r"(0) : "memory"); tile_cursor_1 = tile_cursor_1 + num_kv_tiles_1; } + if (elect_sync()) { + asm volatile("cp.async.bulk.wait_group 0;"); + } asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); asm volatile("tcgen05.fence::before_thread_sync;"); } @@ -2107,6 +2115,7 @@ kernel_cake_dsv4_e1660a7f44d156286fea(const __grid_constant__ CUtensorMap tmap_q unsigned int k_stage_1 = 0; unsigned int v_stage_1 = 0; unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_epi_q_free_0 = 1; unsigned int _phase_k_empty = 1; unsigned int _phase_v_empty = 1; #pragma unroll 1 @@ -2117,6 +2126,8 @@ kernel_cake_dsv4_e1660a7f44d156286fea(const __grid_constant__ CUtensorMap tmap_q int first_kv_tile_1 = split_idx_2 * num_kv_tiles_3; mbarrier_wait(q_empty_addr, _phase_q_empty_0); _phase_q_empty_0 ^= 1; + mbarrier_wait(epi_q_free_addr, _phase_epi_q_free_0); + _phase_epi_q_free_0 ^= 1; if (load_warp_rank == 0) { if (cta_rank == 0) { if (elect_sync()) { diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_75827bd38b443de66f67_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_530fe7e6cb8c7d09744d_binding.cu similarity index 97% rename from csrc/cake_dsv4/sm_100a/cake_dsv4_75827bd38b443de66f67_binding.cu rename to csrc/cake_dsv4/sm_100a/cake_dsv4_530fe7e6cb8c7d09744d_binding.cu index 7cd6edd23c7..51916b01598 100644 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_75827bd38b443de66f67_binding.cu +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_530fe7e6cb8c7d09744d_binding.cu @@ -15,7 +15,7 @@ */ // Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_75827bd38b443de66f67'. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_530fe7e6cb8c7d09744d'. #include #include #include @@ -31,10 +31,10 @@ #include #include -extern "C" __global__ void kernel_cake_dsv4_75827bd38b443de66f67(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items); +extern "C" __global__ void kernel_cake_dsv4_530fe7e6cb8c7d09744d(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items); -namespace cake_host_shim_fefc853e0bc1ea14 { +namespace cake_host_shim_ed72cc11693a1ade { using tvm::ffi::Optional; using tvm::ffi::TensorView; @@ -413,7 +413,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ } if (smem_status == cudaSuccess) { smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_75827bd38b443de66f67), + reinterpret_cast(kernel_cake_dsv4_530fe7e6cb8c7d09744d), cudaFuncAttributeMaxDynamicSharedMemorySize, 205824); } @@ -423,7 +423,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ } } TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_75827bd38b443de66f67 failed: " + << "cudaFuncSetAttribute for kernel_cake_dsv4_530fe7e6cb8c7d09744d failed: " << cudaGetErrorString(smem_status); dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); dim3 block(512u, 1u, 1u); @@ -434,12 +434,12 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ config.dynamicSmemBytes = 205824u; config.stream = stream; cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_75827bd38b443de66f67), kargs); + &config, reinterpret_cast(kernel_cake_dsv4_530fe7e6cb8c7d09744d), kargs); TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_75827bd38b443de66f67 failed: " + << "cudaLaunchKernelExC for kernel_cake_dsv4_530fe7e6cb8c7d09744d failed: " << cudaGetErrorString(launch_status); } -} // namespace cake_host_shim_fefc853e0bc1ea14 +} // namespace cake_host_shim_ed72cc11693a1ade -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_fefc853e0bc1ea14::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_ed72cc11693a1ade::Run); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_75827bd38b443de66f67_kernel.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_530fe7e6cb8c7d09744d_kernel.cu similarity index 54% rename from csrc/cake_dsv4/sm_100a/cake_dsv4_75827bd38b443de66f67_kernel.cu rename to csrc/cake_dsv4/sm_100a/cake_dsv4_530fe7e6cb8c7d09744d_kernel.cu index fb67b80f0b2..37358d2befd 100644 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_75827bd38b443de66f67_kernel.cu +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_530fe7e6cb8c7d09744d_kernel.cu @@ -1192,7 +1192,7 @@ __device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { extern "C" { __global__ __launch_bounds__(512, 1) void -kernel_cake_dsv4_75827bd38b443de66f67(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items) +kernel_cake_dsv4_530fe7e6cb8c7d09744d(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); @@ -1464,7 +1464,7 @@ kernel_cake_dsv4_75827bd38b443de66f67(const __grid_constant__ CUtensorMap tmap_q asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_sparse_indices_addr + (unsigned int)(index_offset * 4)), "r"(sparse_rows[0]), "r"(sparse_rows[1]), "r"(sparse_rows[2]), "r"(sparse_rows[3]) : "memory"); } asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile("barrier.sync 8, 128;" ::: "memory"); + asm volatile("barrier.sync 8, 256;" ::: "memory"); } int k_index_stage_base = smem_sparse_indices_addr; { @@ -1477,7 +1477,8 @@ kernel_cake_dsv4_75827bd38b443de66f67(const __grid_constant__ CUtensorMap tmap_q } int k_dst = smem_kv_addr + load_kv_stage * 16384; int group = lane; - if (group < 32) { + int k_group_limit = ((work_idx_1 == (unsigned int)bid) ? 16 : 32); + if (group < k_group_limit) { int group_offset = k_cta_offset + group * 4; int raw_rows[4]; asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" @@ -1655,6 +1656,29 @@ kernel_cake_dsv4_75827bd38b443de66f67(const __grid_constant__ CUtensorMap tmap_q asm volatile("tcgen05.fence::before_thread_sync;"); mbarrier_arrive(s_seeded_addr); } + { + asm volatile("barrier.sync 8, 256;" ::: "memory"); + if (work_idx_2 == (unsigned int)bid) { + const int helper_stage = warp % 4; + int helper_dst = smem_kv_addr + (unsigned int)(helper_stage * 16384); + int helper_group = lane; + if (helper_group >= 16) { + int helper_rows[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&helper_rows[0])), "=r"(*reinterpret_cast(&helper_rows[(0) + 1])), "=r"(*reinterpret_cast(&helper_rows[(0) + 2])), "=r"(*reinterpret_cast(&helper_rows[(0) + 3])) + : "r"(smem_sparse_indices_addr + (unsigned int)(helper_group * 16))); + int hrow0 = ((helper_rows[0] >= 0) ? helper_rows[0] : 0); + int hrow1 = ((helper_rows[1] >= 0) ? helper_rows[1] : 0); + int hrow2 = ((helper_rows[2] >= 0) ? helper_rows[2] : 0); + int hrow3 = ((helper_rows[3] >= 0) ? helper_rows[3] : 0); + if (first_tile_1 == 0) { + tma_gather4_gmem2smem(helper_dst + helper_group * 512, (&tmap_swa_kv), helper_stage * 128, hrow0, hrow1, hrow2, hrow3, kv_full_addr + (helper_stage) * 8); + } else { + tma_gather4_gmem2smem(helper_dst + helper_group * 512, (&tmap_compressed_kv), helper_stage * 128, hrow0, hrow1, hrow2, hrow3, kv_full_addr + (helper_stage) * 8); + } + } + } + } { mbarrier_wait(q_full_addr, _phase_q_full_0_1); _phase_q_full_0_1 ^= 1; @@ -1696,121 +1720,437 @@ kernel_cake_dsv4_75827bd38b443de66f67(const __grid_constant__ CUtensorMap tmap_q invalid_cols[w] = _vote_0; any_invalid = any_invalid | invalid_cols[w]; } - float _tmem_load_0[4]; - tmem_ld_x4(&_tmem_load_0[0], s_base); + float _tmem_load_0[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]) + : "r"(s_base)); float _tmem_load_1[4]; tmem_ld_x4(&_tmem_load_1[0], s_base); + float _tmem_load_2[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]) + : "r"(s_base + 32)); + float _tmem_load_3[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]) + : "r"(s_base + 64)); + float _tmem_load_4[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_4[0]), "=f"(_tmem_load_4[1]), "=f"(_tmem_load_4[2]), "=f"(_tmem_load_4[3]), "=f"(_tmem_load_4[4]), "=f"(_tmem_load_4[5]), "=f"(_tmem_load_4[6]), "=f"(_tmem_load_4[7]), "=f"(_tmem_load_4[8]), "=f"(_tmem_load_4[9]), "=f"(_tmem_load_4[10]), "=f"(_tmem_load_4[11]), "=f"(_tmem_load_4[12]), "=f"(_tmem_load_4[13]), "=f"(_tmem_load_4[14]), "=f"(_tmem_load_4[15]), "=f"(_tmem_load_4[16]), "=f"(_tmem_load_4[17]), "=f"(_tmem_load_4[18]), "=f"(_tmem_load_4[19]), "=f"(_tmem_load_4[20]), "=f"(_tmem_load_4[21]), "=f"(_tmem_load_4[22]), "=f"(_tmem_load_4[23]), "=f"(_tmem_load_4[24]), "=f"(_tmem_load_4[25]), "=f"(_tmem_load_4[26]), "=f"(_tmem_load_4[27]), "=f"(_tmem_load_4[28]), "=f"(_tmem_load_4[29]), "=f"(_tmem_load_4[30]), "=f"(_tmem_load_4[31]) + : "r"(s_base + 96)); { - #pragma unroll - for (int frag = 0; frag < 4; frag++) { - float _tmem_load_2[32]; - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]) - : "r"(s_base + frag * 32)); - int frag_valid = valid_sparse_cols - frag * 32; - uint32_t _slice_lo_mask_1; - { - int _lim_0 = frag_valid; - if (_lim_0 <= 0) { _slice_lo_mask_1 = 0u; } - else if (_lim_0 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_0)); - } + int frag_valid_0 = valid_sparse_cols; + uint32_t _slice_lo_mask_1; + { + int _lim_0 = frag_valid_0; + if (_lim_0 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_0)); } - if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_2[0] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_2[1] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_2[2] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_2[3] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_2[4] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_2[5] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_2[6] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_2[7] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_2[8] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_2[9] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_2[10] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_2[11] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_2[12] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_2[13] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_2[14] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_2[15] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_2[16] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_2[17] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_2[18] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_2[19] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_2[20] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_2[21] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_2[22] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_2[23] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_2[24] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_2[25] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_2[26] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_2[27] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_2[28] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_2[29] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_2[30] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_2[31] = -CAKE_INF; - if (any_invalid != 0) { - unsigned int frag_mask_word = invalid_cols[frag]; - _tmem_load_2[0] = (((frag_mask_word & 1) != 0) ? -CAKE_INF : _tmem_load_2[0]); - _tmem_load_2[1] = (((frag_mask_word >> 1 & 1) != 0) ? -CAKE_INF : _tmem_load_2[1]); - _tmem_load_2[2] = (((frag_mask_word >> 2 & 1) != 0) ? -CAKE_INF : _tmem_load_2[2]); - _tmem_load_2[3] = (((frag_mask_word >> 3 & 1) != 0) ? -CAKE_INF : _tmem_load_2[3]); - _tmem_load_2[4] = (((frag_mask_word >> 4 & 1) != 0) ? -CAKE_INF : _tmem_load_2[4]); - _tmem_load_2[5] = (((frag_mask_word >> 5 & 1) != 0) ? -CAKE_INF : _tmem_load_2[5]); - _tmem_load_2[6] = (((frag_mask_word >> 6 & 1) != 0) ? -CAKE_INF : _tmem_load_2[6]); - _tmem_load_2[7] = (((frag_mask_word >> 7 & 1) != 0) ? -CAKE_INF : _tmem_load_2[7]); - _tmem_load_2[8] = (((frag_mask_word >> 8 & 1) != 0) ? -CAKE_INF : _tmem_load_2[8]); - _tmem_load_2[9] = (((frag_mask_word >> 9 & 1) != 0) ? -CAKE_INF : _tmem_load_2[9]); - _tmem_load_2[10] = (((frag_mask_word >> 10 & 1) != 0) ? -CAKE_INF : _tmem_load_2[10]); - _tmem_load_2[11] = (((frag_mask_word >> 11 & 1) != 0) ? -CAKE_INF : _tmem_load_2[11]); - _tmem_load_2[12] = (((frag_mask_word >> 12 & 1) != 0) ? -CAKE_INF : _tmem_load_2[12]); - _tmem_load_2[13] = (((frag_mask_word >> 13 & 1) != 0) ? -CAKE_INF : _tmem_load_2[13]); - _tmem_load_2[14] = (((frag_mask_word >> 14 & 1) != 0) ? -CAKE_INF : _tmem_load_2[14]); - _tmem_load_2[15] = (((frag_mask_word >> 15 & 1) != 0) ? -CAKE_INF : _tmem_load_2[15]); - _tmem_load_2[16] = (((frag_mask_word >> 16 & 1) != 0) ? -CAKE_INF : _tmem_load_2[16]); - _tmem_load_2[17] = (((frag_mask_word >> 17 & 1) != 0) ? -CAKE_INF : _tmem_load_2[17]); - _tmem_load_2[18] = (((frag_mask_word >> 18 & 1) != 0) ? -CAKE_INF : _tmem_load_2[18]); - _tmem_load_2[19] = (((frag_mask_word >> 19 & 1) != 0) ? -CAKE_INF : _tmem_load_2[19]); - _tmem_load_2[20] = (((frag_mask_word >> 20 & 1) != 0) ? -CAKE_INF : _tmem_load_2[20]); - _tmem_load_2[21] = (((frag_mask_word >> 21 & 1) != 0) ? -CAKE_INF : _tmem_load_2[21]); - _tmem_load_2[22] = (((frag_mask_word >> 22 & 1) != 0) ? -CAKE_INF : _tmem_load_2[22]); - _tmem_load_2[23] = (((frag_mask_word >> 23 & 1) != 0) ? -CAKE_INF : _tmem_load_2[23]); - _tmem_load_2[24] = (((frag_mask_word >> 24 & 1) != 0) ? -CAKE_INF : _tmem_load_2[24]); - _tmem_load_2[25] = (((frag_mask_word >> 25 & 1) != 0) ? -CAKE_INF : _tmem_load_2[25]); - _tmem_load_2[26] = (((frag_mask_word >> 26 & 1) != 0) ? -CAKE_INF : _tmem_load_2[26]); - _tmem_load_2[27] = (((frag_mask_word >> 27 & 1) != 0) ? -CAKE_INF : _tmem_load_2[27]); - _tmem_load_2[28] = (((frag_mask_word >> 28 & 1) != 0) ? -CAKE_INF : _tmem_load_2[28]); - _tmem_load_2[29] = (((frag_mask_word >> 29 & 1) != 0) ? -CAKE_INF : _tmem_load_2[29]); - _tmem_load_2[30] = (((frag_mask_word >> 30 & 1) != 0) ? -CAKE_INF : _tmem_load_2[30]); - _tmem_load_2[31] = (((frag_mask_word >> 31 & 1) != 0) ? -CAKE_INF : _tmem_load_2[31]); + } + if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_0[0] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_0[1] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_0[2] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_0[3] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_0[4] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_0[5] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_0[6] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_0[7] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_0[8] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_0[9] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_0[10] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_0[11] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_0[12] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_0[13] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_0[14] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_0[15] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_0[16] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_0[17] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_0[18] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_0[19] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_0[20] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_0[21] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_0[22] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_0[23] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_0[24] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_0[25] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_0[26] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_0[27] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_0[28] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_0[29] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_0[30] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_0[31] = -CAKE_INF; + if (any_invalid != 0) { + unsigned int frag_mask_word_0 = invalid_cols[0]; + _tmem_load_0[0] = (((frag_mask_word_0 & 1) != 0) ? -CAKE_INF : _tmem_load_0[0]); + _tmem_load_0[1] = (((frag_mask_word_0 >> 1 & 1) != 0) ? -CAKE_INF : _tmem_load_0[1]); + _tmem_load_0[2] = (((frag_mask_word_0 >> 2 & 1) != 0) ? -CAKE_INF : _tmem_load_0[2]); + _tmem_load_0[3] = (((frag_mask_word_0 >> 3 & 1) != 0) ? -CAKE_INF : _tmem_load_0[3]); + _tmem_load_0[4] = (((frag_mask_word_0 >> 4 & 1) != 0) ? -CAKE_INF : _tmem_load_0[4]); + _tmem_load_0[5] = (((frag_mask_word_0 >> 5 & 1) != 0) ? -CAKE_INF : _tmem_load_0[5]); + _tmem_load_0[6] = (((frag_mask_word_0 >> 6 & 1) != 0) ? -CAKE_INF : _tmem_load_0[6]); + _tmem_load_0[7] = (((frag_mask_word_0 >> 7 & 1) != 0) ? -CAKE_INF : _tmem_load_0[7]); + _tmem_load_0[8] = (((frag_mask_word_0 >> 8 & 1) != 0) ? -CAKE_INF : _tmem_load_0[8]); + _tmem_load_0[9] = (((frag_mask_word_0 >> 9 & 1) != 0) ? -CAKE_INF : _tmem_load_0[9]); + _tmem_load_0[10] = (((frag_mask_word_0 >> 10 & 1) != 0) ? -CAKE_INF : _tmem_load_0[10]); + _tmem_load_0[11] = (((frag_mask_word_0 >> 11 & 1) != 0) ? -CAKE_INF : _tmem_load_0[11]); + _tmem_load_0[12] = (((frag_mask_word_0 >> 12 & 1) != 0) ? -CAKE_INF : _tmem_load_0[12]); + _tmem_load_0[13] = (((frag_mask_word_0 >> 13 & 1) != 0) ? -CAKE_INF : _tmem_load_0[13]); + _tmem_load_0[14] = (((frag_mask_word_0 >> 14 & 1) != 0) ? -CAKE_INF : _tmem_load_0[14]); + _tmem_load_0[15] = (((frag_mask_word_0 >> 15 & 1) != 0) ? -CAKE_INF : _tmem_load_0[15]); + _tmem_load_0[16] = (((frag_mask_word_0 >> 16 & 1) != 0) ? -CAKE_INF : _tmem_load_0[16]); + _tmem_load_0[17] = (((frag_mask_word_0 >> 17 & 1) != 0) ? -CAKE_INF : _tmem_load_0[17]); + _tmem_load_0[18] = (((frag_mask_word_0 >> 18 & 1) != 0) ? -CAKE_INF : _tmem_load_0[18]); + _tmem_load_0[19] = (((frag_mask_word_0 >> 19 & 1) != 0) ? -CAKE_INF : _tmem_load_0[19]); + _tmem_load_0[20] = (((frag_mask_word_0 >> 20 & 1) != 0) ? -CAKE_INF : _tmem_load_0[20]); + _tmem_load_0[21] = (((frag_mask_word_0 >> 21 & 1) != 0) ? -CAKE_INF : _tmem_load_0[21]); + _tmem_load_0[22] = (((frag_mask_word_0 >> 22 & 1) != 0) ? -CAKE_INF : _tmem_load_0[22]); + _tmem_load_0[23] = (((frag_mask_word_0 >> 23 & 1) != 0) ? -CAKE_INF : _tmem_load_0[23]); + _tmem_load_0[24] = (((frag_mask_word_0 >> 24 & 1) != 0) ? -CAKE_INF : _tmem_load_0[24]); + _tmem_load_0[25] = (((frag_mask_word_0 >> 25 & 1) != 0) ? -CAKE_INF : _tmem_load_0[25]); + _tmem_load_0[26] = (((frag_mask_word_0 >> 26 & 1) != 0) ? -CAKE_INF : _tmem_load_0[26]); + _tmem_load_0[27] = (((frag_mask_word_0 >> 27 & 1) != 0) ? -CAKE_INF : _tmem_load_0[27]); + _tmem_load_0[28] = (((frag_mask_word_0 >> 28 & 1) != 0) ? -CAKE_INF : _tmem_load_0[28]); + _tmem_load_0[29] = (((frag_mask_word_0 >> 29 & 1) != 0) ? -CAKE_INF : _tmem_load_0[29]); + _tmem_load_0[30] = (((frag_mask_word_0 >> 30 & 1) != 0) ? -CAKE_INF : _tmem_load_0[30]); + _tmem_load_0[31] = (((frag_mask_word_0 >> 31 & 1) != 0) ? -CAKE_INF : _tmem_load_0[31]); + } + float2 _reg_reduce_max2_1 = {-CAKE_INF, -CAKE_INF}; + _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_0[0], _tmem_load_0[1])); + _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_0[2], _tmem_load_0[3])); + _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_0[4], _tmem_load_0[5])); + _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_0[6], _tmem_load_0[7])); + _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_0[8], _tmem_load_0[9])); + _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_0[10], _tmem_load_0[11])); + _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_0[12], _tmem_load_0[13])); + _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_0[14], _tmem_load_0[15])); + _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_0[16], _tmem_load_0[17])); + _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_0[18], _tmem_load_0[19])); + _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_0[20], _tmem_load_0[21])); + _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_0[22], _tmem_load_0[23])); + _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_0[24], _tmem_load_0[25])); + _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_0[26], _tmem_load_0[27])); + _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_0[28], _tmem_load_0[29])); + _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_0[30], _tmem_load_0[31])); + float _tmem_load_0_max = row_max_reduce(_reg_reduce_max2_1); + float _max_3 = max_noftz(new_max, _tmem_load_0_max); + new_max = _max_3; + int frag_valid_1 = valid_sparse_cols - 32; + uint32_t _slice_lo_mask_2; + { + int _lim_2 = frag_valid_1; + if (_lim_2 <= 0) { _slice_lo_mask_2 = 0u; } + else if (_lim_2 >= 32) { _slice_lo_mask_2 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_2) : "r"(_lim_2)); + } + } + if (!(_slice_lo_mask_2 & (1u << 0))) _tmem_load_2[0] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 1))) _tmem_load_2[1] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 2))) _tmem_load_2[2] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 3))) _tmem_load_2[3] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 4))) _tmem_load_2[4] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 5))) _tmem_load_2[5] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 6))) _tmem_load_2[6] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 7))) _tmem_load_2[7] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 8))) _tmem_load_2[8] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 9))) _tmem_load_2[9] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 10))) _tmem_load_2[10] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 11))) _tmem_load_2[11] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 12))) _tmem_load_2[12] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 13))) _tmem_load_2[13] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 14))) _tmem_load_2[14] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 15))) _tmem_load_2[15] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 16))) _tmem_load_2[16] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 17))) _tmem_load_2[17] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 18))) _tmem_load_2[18] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 19))) _tmem_load_2[19] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 20))) _tmem_load_2[20] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 21))) _tmem_load_2[21] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 22))) _tmem_load_2[22] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 23))) _tmem_load_2[23] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 24))) _tmem_load_2[24] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 25))) _tmem_load_2[25] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 26))) _tmem_load_2[26] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 27))) _tmem_load_2[27] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 28))) _tmem_load_2[28] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 29))) _tmem_load_2[29] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 30))) _tmem_load_2[30] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 31))) _tmem_load_2[31] = -CAKE_INF; + if (any_invalid != 0) { + unsigned int frag_mask_word_1 = invalid_cols[1]; + _tmem_load_2[0] = (((frag_mask_word_1 & 1) != 0) ? -CAKE_INF : _tmem_load_2[0]); + _tmem_load_2[1] = (((frag_mask_word_1 >> 1 & 1) != 0) ? -CAKE_INF : _tmem_load_2[1]); + _tmem_load_2[2] = (((frag_mask_word_1 >> 2 & 1) != 0) ? -CAKE_INF : _tmem_load_2[2]); + _tmem_load_2[3] = (((frag_mask_word_1 >> 3 & 1) != 0) ? -CAKE_INF : _tmem_load_2[3]); + _tmem_load_2[4] = (((frag_mask_word_1 >> 4 & 1) != 0) ? -CAKE_INF : _tmem_load_2[4]); + _tmem_load_2[5] = (((frag_mask_word_1 >> 5 & 1) != 0) ? -CAKE_INF : _tmem_load_2[5]); + _tmem_load_2[6] = (((frag_mask_word_1 >> 6 & 1) != 0) ? -CAKE_INF : _tmem_load_2[6]); + _tmem_load_2[7] = (((frag_mask_word_1 >> 7 & 1) != 0) ? -CAKE_INF : _tmem_load_2[7]); + _tmem_load_2[8] = (((frag_mask_word_1 >> 8 & 1) != 0) ? -CAKE_INF : _tmem_load_2[8]); + _tmem_load_2[9] = (((frag_mask_word_1 >> 9 & 1) != 0) ? -CAKE_INF : _tmem_load_2[9]); + _tmem_load_2[10] = (((frag_mask_word_1 >> 10 & 1) != 0) ? -CAKE_INF : _tmem_load_2[10]); + _tmem_load_2[11] = (((frag_mask_word_1 >> 11 & 1) != 0) ? -CAKE_INF : _tmem_load_2[11]); + _tmem_load_2[12] = (((frag_mask_word_1 >> 12 & 1) != 0) ? -CAKE_INF : _tmem_load_2[12]); + _tmem_load_2[13] = (((frag_mask_word_1 >> 13 & 1) != 0) ? -CAKE_INF : _tmem_load_2[13]); + _tmem_load_2[14] = (((frag_mask_word_1 >> 14 & 1) != 0) ? -CAKE_INF : _tmem_load_2[14]); + _tmem_load_2[15] = (((frag_mask_word_1 >> 15 & 1) != 0) ? -CAKE_INF : _tmem_load_2[15]); + _tmem_load_2[16] = (((frag_mask_word_1 >> 16 & 1) != 0) ? -CAKE_INF : _tmem_load_2[16]); + _tmem_load_2[17] = (((frag_mask_word_1 >> 17 & 1) != 0) ? -CAKE_INF : _tmem_load_2[17]); + _tmem_load_2[18] = (((frag_mask_word_1 >> 18 & 1) != 0) ? -CAKE_INF : _tmem_load_2[18]); + _tmem_load_2[19] = (((frag_mask_word_1 >> 19 & 1) != 0) ? -CAKE_INF : _tmem_load_2[19]); + _tmem_load_2[20] = (((frag_mask_word_1 >> 20 & 1) != 0) ? -CAKE_INF : _tmem_load_2[20]); + _tmem_load_2[21] = (((frag_mask_word_1 >> 21 & 1) != 0) ? -CAKE_INF : _tmem_load_2[21]); + _tmem_load_2[22] = (((frag_mask_word_1 >> 22 & 1) != 0) ? -CAKE_INF : _tmem_load_2[22]); + _tmem_load_2[23] = (((frag_mask_word_1 >> 23 & 1) != 0) ? -CAKE_INF : _tmem_load_2[23]); + _tmem_load_2[24] = (((frag_mask_word_1 >> 24 & 1) != 0) ? -CAKE_INF : _tmem_load_2[24]); + _tmem_load_2[25] = (((frag_mask_word_1 >> 25 & 1) != 0) ? -CAKE_INF : _tmem_load_2[25]); + _tmem_load_2[26] = (((frag_mask_word_1 >> 26 & 1) != 0) ? -CAKE_INF : _tmem_load_2[26]); + _tmem_load_2[27] = (((frag_mask_word_1 >> 27 & 1) != 0) ? -CAKE_INF : _tmem_load_2[27]); + _tmem_load_2[28] = (((frag_mask_word_1 >> 28 & 1) != 0) ? -CAKE_INF : _tmem_load_2[28]); + _tmem_load_2[29] = (((frag_mask_word_1 >> 29 & 1) != 0) ? -CAKE_INF : _tmem_load_2[29]); + _tmem_load_2[30] = (((frag_mask_word_1 >> 30 & 1) != 0) ? -CAKE_INF : _tmem_load_2[30]); + _tmem_load_2[31] = (((frag_mask_word_1 >> 31 & 1) != 0) ? -CAKE_INF : _tmem_load_2[31]); + } + float2 _reg_reduce_max2_3 = {-CAKE_INF, -CAKE_INF}; + _reg_reduce_max2_3.x = max_noftz(_reg_reduce_max2_3.x, max_noftz(_tmem_load_2[0], _tmem_load_2[1])); + _reg_reduce_max2_3.y = max_noftz(_reg_reduce_max2_3.y, max_noftz(_tmem_load_2[2], _tmem_load_2[3])); + _reg_reduce_max2_3.x = max_noftz(_reg_reduce_max2_3.x, max_noftz(_tmem_load_2[4], _tmem_load_2[5])); + _reg_reduce_max2_3.y = max_noftz(_reg_reduce_max2_3.y, max_noftz(_tmem_load_2[6], _tmem_load_2[7])); + _reg_reduce_max2_3.x = max_noftz(_reg_reduce_max2_3.x, max_noftz(_tmem_load_2[8], _tmem_load_2[9])); + _reg_reduce_max2_3.y = max_noftz(_reg_reduce_max2_3.y, max_noftz(_tmem_load_2[10], _tmem_load_2[11])); + _reg_reduce_max2_3.x = max_noftz(_reg_reduce_max2_3.x, max_noftz(_tmem_load_2[12], _tmem_load_2[13])); + _reg_reduce_max2_3.y = max_noftz(_reg_reduce_max2_3.y, max_noftz(_tmem_load_2[14], _tmem_load_2[15])); + _reg_reduce_max2_3.x = max_noftz(_reg_reduce_max2_3.x, max_noftz(_tmem_load_2[16], _tmem_load_2[17])); + _reg_reduce_max2_3.y = max_noftz(_reg_reduce_max2_3.y, max_noftz(_tmem_load_2[18], _tmem_load_2[19])); + _reg_reduce_max2_3.x = max_noftz(_reg_reduce_max2_3.x, max_noftz(_tmem_load_2[20], _tmem_load_2[21])); + _reg_reduce_max2_3.y = max_noftz(_reg_reduce_max2_3.y, max_noftz(_tmem_load_2[22], _tmem_load_2[23])); + _reg_reduce_max2_3.x = max_noftz(_reg_reduce_max2_3.x, max_noftz(_tmem_load_2[24], _tmem_load_2[25])); + _reg_reduce_max2_3.y = max_noftz(_reg_reduce_max2_3.y, max_noftz(_tmem_load_2[26], _tmem_load_2[27])); + _reg_reduce_max2_3.x = max_noftz(_reg_reduce_max2_3.x, max_noftz(_tmem_load_2[28], _tmem_load_2[29])); + _reg_reduce_max2_3.y = max_noftz(_reg_reduce_max2_3.y, max_noftz(_tmem_load_2[30], _tmem_load_2[31])); + float _tmem_load_2_max = row_max_reduce(_reg_reduce_max2_3); + float _max_4 = max_noftz(new_max, _tmem_load_2_max); + new_max = _max_4; + int frag_valid_2 = valid_sparse_cols - 64; + uint32_t _slice_lo_mask_3; + { + int _lim_4 = frag_valid_2; + if (_lim_4 <= 0) { _slice_lo_mask_3 = 0u; } + else if (_lim_4 >= 32) { _slice_lo_mask_3 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_3) : "r"(_lim_4)); + } + } + if (!(_slice_lo_mask_3 & (1u << 0))) _tmem_load_3[0] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 1))) _tmem_load_3[1] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 2))) _tmem_load_3[2] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 3))) _tmem_load_3[3] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 4))) _tmem_load_3[4] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 5))) _tmem_load_3[5] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 6))) _tmem_load_3[6] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 7))) _tmem_load_3[7] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 8))) _tmem_load_3[8] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 9))) _tmem_load_3[9] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 10))) _tmem_load_3[10] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 11))) _tmem_load_3[11] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 12))) _tmem_load_3[12] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 13))) _tmem_load_3[13] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 14))) _tmem_load_3[14] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 15))) _tmem_load_3[15] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 16))) _tmem_load_3[16] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 17))) _tmem_load_3[17] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 18))) _tmem_load_3[18] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 19))) _tmem_load_3[19] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 20))) _tmem_load_3[20] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 21))) _tmem_load_3[21] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 22))) _tmem_load_3[22] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 23))) _tmem_load_3[23] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 24))) _tmem_load_3[24] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 25))) _tmem_load_3[25] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 26))) _tmem_load_3[26] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 27))) _tmem_load_3[27] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 28))) _tmem_load_3[28] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 29))) _tmem_load_3[29] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 30))) _tmem_load_3[30] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 31))) _tmem_load_3[31] = -CAKE_INF; + if (any_invalid != 0) { + unsigned int frag_mask_word_2 = invalid_cols[2]; + _tmem_load_3[0] = (((frag_mask_word_2 & 1) != 0) ? -CAKE_INF : _tmem_load_3[0]); + _tmem_load_3[1] = (((frag_mask_word_2 >> 1 & 1) != 0) ? -CAKE_INF : _tmem_load_3[1]); + _tmem_load_3[2] = (((frag_mask_word_2 >> 2 & 1) != 0) ? -CAKE_INF : _tmem_load_3[2]); + _tmem_load_3[3] = (((frag_mask_word_2 >> 3 & 1) != 0) ? -CAKE_INF : _tmem_load_3[3]); + _tmem_load_3[4] = (((frag_mask_word_2 >> 4 & 1) != 0) ? -CAKE_INF : _tmem_load_3[4]); + _tmem_load_3[5] = (((frag_mask_word_2 >> 5 & 1) != 0) ? -CAKE_INF : _tmem_load_3[5]); + _tmem_load_3[6] = (((frag_mask_word_2 >> 6 & 1) != 0) ? -CAKE_INF : _tmem_load_3[6]); + _tmem_load_3[7] = (((frag_mask_word_2 >> 7 & 1) != 0) ? -CAKE_INF : _tmem_load_3[7]); + _tmem_load_3[8] = (((frag_mask_word_2 >> 8 & 1) != 0) ? -CAKE_INF : _tmem_load_3[8]); + _tmem_load_3[9] = (((frag_mask_word_2 >> 9 & 1) != 0) ? -CAKE_INF : _tmem_load_3[9]); + _tmem_load_3[10] = (((frag_mask_word_2 >> 10 & 1) != 0) ? -CAKE_INF : _tmem_load_3[10]); + _tmem_load_3[11] = (((frag_mask_word_2 >> 11 & 1) != 0) ? -CAKE_INF : _tmem_load_3[11]); + _tmem_load_3[12] = (((frag_mask_word_2 >> 12 & 1) != 0) ? -CAKE_INF : _tmem_load_3[12]); + _tmem_load_3[13] = (((frag_mask_word_2 >> 13 & 1) != 0) ? -CAKE_INF : _tmem_load_3[13]); + _tmem_load_3[14] = (((frag_mask_word_2 >> 14 & 1) != 0) ? -CAKE_INF : _tmem_load_3[14]); + _tmem_load_3[15] = (((frag_mask_word_2 >> 15 & 1) != 0) ? -CAKE_INF : _tmem_load_3[15]); + _tmem_load_3[16] = (((frag_mask_word_2 >> 16 & 1) != 0) ? -CAKE_INF : _tmem_load_3[16]); + _tmem_load_3[17] = (((frag_mask_word_2 >> 17 & 1) != 0) ? -CAKE_INF : _tmem_load_3[17]); + _tmem_load_3[18] = (((frag_mask_word_2 >> 18 & 1) != 0) ? -CAKE_INF : _tmem_load_3[18]); + _tmem_load_3[19] = (((frag_mask_word_2 >> 19 & 1) != 0) ? -CAKE_INF : _tmem_load_3[19]); + _tmem_load_3[20] = (((frag_mask_word_2 >> 20 & 1) != 0) ? -CAKE_INF : _tmem_load_3[20]); + _tmem_load_3[21] = (((frag_mask_word_2 >> 21 & 1) != 0) ? -CAKE_INF : _tmem_load_3[21]); + _tmem_load_3[22] = (((frag_mask_word_2 >> 22 & 1) != 0) ? -CAKE_INF : _tmem_load_3[22]); + _tmem_load_3[23] = (((frag_mask_word_2 >> 23 & 1) != 0) ? -CAKE_INF : _tmem_load_3[23]); + _tmem_load_3[24] = (((frag_mask_word_2 >> 24 & 1) != 0) ? -CAKE_INF : _tmem_load_3[24]); + _tmem_load_3[25] = (((frag_mask_word_2 >> 25 & 1) != 0) ? -CAKE_INF : _tmem_load_3[25]); + _tmem_load_3[26] = (((frag_mask_word_2 >> 26 & 1) != 0) ? -CAKE_INF : _tmem_load_3[26]); + _tmem_load_3[27] = (((frag_mask_word_2 >> 27 & 1) != 0) ? -CAKE_INF : _tmem_load_3[27]); + _tmem_load_3[28] = (((frag_mask_word_2 >> 28 & 1) != 0) ? -CAKE_INF : _tmem_load_3[28]); + _tmem_load_3[29] = (((frag_mask_word_2 >> 29 & 1) != 0) ? -CAKE_INF : _tmem_load_3[29]); + _tmem_load_3[30] = (((frag_mask_word_2 >> 30 & 1) != 0) ? -CAKE_INF : _tmem_load_3[30]); + _tmem_load_3[31] = (((frag_mask_word_2 >> 31 & 1) != 0) ? -CAKE_INF : _tmem_load_3[31]); + } + float2 _reg_reduce_max2_5 = {-CAKE_INF, -CAKE_INF}; + _reg_reduce_max2_5.x = max_noftz(_reg_reduce_max2_5.x, max_noftz(_tmem_load_3[0], _tmem_load_3[1])); + _reg_reduce_max2_5.y = max_noftz(_reg_reduce_max2_5.y, max_noftz(_tmem_load_3[2], _tmem_load_3[3])); + _reg_reduce_max2_5.x = max_noftz(_reg_reduce_max2_5.x, max_noftz(_tmem_load_3[4], _tmem_load_3[5])); + _reg_reduce_max2_5.y = max_noftz(_reg_reduce_max2_5.y, max_noftz(_tmem_load_3[6], _tmem_load_3[7])); + _reg_reduce_max2_5.x = max_noftz(_reg_reduce_max2_5.x, max_noftz(_tmem_load_3[8], _tmem_load_3[9])); + _reg_reduce_max2_5.y = max_noftz(_reg_reduce_max2_5.y, max_noftz(_tmem_load_3[10], _tmem_load_3[11])); + _reg_reduce_max2_5.x = max_noftz(_reg_reduce_max2_5.x, max_noftz(_tmem_load_3[12], _tmem_load_3[13])); + _reg_reduce_max2_5.y = max_noftz(_reg_reduce_max2_5.y, max_noftz(_tmem_load_3[14], _tmem_load_3[15])); + _reg_reduce_max2_5.x = max_noftz(_reg_reduce_max2_5.x, max_noftz(_tmem_load_3[16], _tmem_load_3[17])); + _reg_reduce_max2_5.y = max_noftz(_reg_reduce_max2_5.y, max_noftz(_tmem_load_3[18], _tmem_load_3[19])); + _reg_reduce_max2_5.x = max_noftz(_reg_reduce_max2_5.x, max_noftz(_tmem_load_3[20], _tmem_load_3[21])); + _reg_reduce_max2_5.y = max_noftz(_reg_reduce_max2_5.y, max_noftz(_tmem_load_3[22], _tmem_load_3[23])); + _reg_reduce_max2_5.x = max_noftz(_reg_reduce_max2_5.x, max_noftz(_tmem_load_3[24], _tmem_load_3[25])); + _reg_reduce_max2_5.y = max_noftz(_reg_reduce_max2_5.y, max_noftz(_tmem_load_3[26], _tmem_load_3[27])); + _reg_reduce_max2_5.x = max_noftz(_reg_reduce_max2_5.x, max_noftz(_tmem_load_3[28], _tmem_load_3[29])); + _reg_reduce_max2_5.y = max_noftz(_reg_reduce_max2_5.y, max_noftz(_tmem_load_3[30], _tmem_load_3[31])); + float _tmem_load_3_max = row_max_reduce(_reg_reduce_max2_5); + float _max_5 = max_noftz(new_max, _tmem_load_3_max); + new_max = _max_5; + int frag_valid_3 = valid_sparse_cols - 96; + uint32_t _slice_lo_mask_4; + { + int _lim_6 = frag_valid_3; + if (_lim_6 <= 0) { _slice_lo_mask_4 = 0u; } + else if (_lim_6 >= 32) { _slice_lo_mask_4 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_4) : "r"(_lim_6)); } - float2 _reg_reduce_max2_1 = {-CAKE_INF, -CAKE_INF}; - _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_2[0], _tmem_load_2[1])); - _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_2[2], _tmem_load_2[3])); - _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_2[4], _tmem_load_2[5])); - _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_2[6], _tmem_load_2[7])); - _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_2[8], _tmem_load_2[9])); - _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_2[10], _tmem_load_2[11])); - _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_2[12], _tmem_load_2[13])); - _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_2[14], _tmem_load_2[15])); - _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_2[16], _tmem_load_2[17])); - _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_2[18], _tmem_load_2[19])); - _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_2[20], _tmem_load_2[21])); - _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_2[22], _tmem_load_2[23])); - _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_2[24], _tmem_load_2[25])); - _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_2[26], _tmem_load_2[27])); - _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_2[28], _tmem_load_2[29])); - _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_2[30], _tmem_load_2[31])); - float _tmem_load_2_max = row_max_reduce(_reg_reduce_max2_1); - float _max_3 = max_noftz(new_max, _tmem_load_2_max); - new_max = _max_3; } + if (!(_slice_lo_mask_4 & (1u << 0))) _tmem_load_4[0] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 1))) _tmem_load_4[1] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 2))) _tmem_load_4[2] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 3))) _tmem_load_4[3] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 4))) _tmem_load_4[4] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 5))) _tmem_load_4[5] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 6))) _tmem_load_4[6] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 7))) _tmem_load_4[7] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 8))) _tmem_load_4[8] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 9))) _tmem_load_4[9] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 10))) _tmem_load_4[10] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 11))) _tmem_load_4[11] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 12))) _tmem_load_4[12] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 13))) _tmem_load_4[13] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 14))) _tmem_load_4[14] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 15))) _tmem_load_4[15] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 16))) _tmem_load_4[16] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 17))) _tmem_load_4[17] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 18))) _tmem_load_4[18] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 19))) _tmem_load_4[19] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 20))) _tmem_load_4[20] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 21))) _tmem_load_4[21] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 22))) _tmem_load_4[22] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 23))) _tmem_load_4[23] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 24))) _tmem_load_4[24] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 25))) _tmem_load_4[25] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 26))) _tmem_load_4[26] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 27))) _tmem_load_4[27] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 28))) _tmem_load_4[28] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 29))) _tmem_load_4[29] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 30))) _tmem_load_4[30] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 31))) _tmem_load_4[31] = -CAKE_INF; + if (any_invalid != 0) { + unsigned int frag_mask_word_3 = invalid_cols[3]; + _tmem_load_4[0] = (((frag_mask_word_3 & 1) != 0) ? -CAKE_INF : _tmem_load_4[0]); + _tmem_load_4[1] = (((frag_mask_word_3 >> 1 & 1) != 0) ? -CAKE_INF : _tmem_load_4[1]); + _tmem_load_4[2] = (((frag_mask_word_3 >> 2 & 1) != 0) ? -CAKE_INF : _tmem_load_4[2]); + _tmem_load_4[3] = (((frag_mask_word_3 >> 3 & 1) != 0) ? -CAKE_INF : _tmem_load_4[3]); + _tmem_load_4[4] = (((frag_mask_word_3 >> 4 & 1) != 0) ? -CAKE_INF : _tmem_load_4[4]); + _tmem_load_4[5] = (((frag_mask_word_3 >> 5 & 1) != 0) ? -CAKE_INF : _tmem_load_4[5]); + _tmem_load_4[6] = (((frag_mask_word_3 >> 6 & 1) != 0) ? -CAKE_INF : _tmem_load_4[6]); + _tmem_load_4[7] = (((frag_mask_word_3 >> 7 & 1) != 0) ? -CAKE_INF : _tmem_load_4[7]); + _tmem_load_4[8] = (((frag_mask_word_3 >> 8 & 1) != 0) ? -CAKE_INF : _tmem_load_4[8]); + _tmem_load_4[9] = (((frag_mask_word_3 >> 9 & 1) != 0) ? -CAKE_INF : _tmem_load_4[9]); + _tmem_load_4[10] = (((frag_mask_word_3 >> 10 & 1) != 0) ? -CAKE_INF : _tmem_load_4[10]); + _tmem_load_4[11] = (((frag_mask_word_3 >> 11 & 1) != 0) ? -CAKE_INF : _tmem_load_4[11]); + _tmem_load_4[12] = (((frag_mask_word_3 >> 12 & 1) != 0) ? -CAKE_INF : _tmem_load_4[12]); + _tmem_load_4[13] = (((frag_mask_word_3 >> 13 & 1) != 0) ? -CAKE_INF : _tmem_load_4[13]); + _tmem_load_4[14] = (((frag_mask_word_3 >> 14 & 1) != 0) ? -CAKE_INF : _tmem_load_4[14]); + _tmem_load_4[15] = (((frag_mask_word_3 >> 15 & 1) != 0) ? -CAKE_INF : _tmem_load_4[15]); + _tmem_load_4[16] = (((frag_mask_word_3 >> 16 & 1) != 0) ? -CAKE_INF : _tmem_load_4[16]); + _tmem_load_4[17] = (((frag_mask_word_3 >> 17 & 1) != 0) ? -CAKE_INF : _tmem_load_4[17]); + _tmem_load_4[18] = (((frag_mask_word_3 >> 18 & 1) != 0) ? -CAKE_INF : _tmem_load_4[18]); + _tmem_load_4[19] = (((frag_mask_word_3 >> 19 & 1) != 0) ? -CAKE_INF : _tmem_load_4[19]); + _tmem_load_4[20] = (((frag_mask_word_3 >> 20 & 1) != 0) ? -CAKE_INF : _tmem_load_4[20]); + _tmem_load_4[21] = (((frag_mask_word_3 >> 21 & 1) != 0) ? -CAKE_INF : _tmem_load_4[21]); + _tmem_load_4[22] = (((frag_mask_word_3 >> 22 & 1) != 0) ? -CAKE_INF : _tmem_load_4[22]); + _tmem_load_4[23] = (((frag_mask_word_3 >> 23 & 1) != 0) ? -CAKE_INF : _tmem_load_4[23]); + _tmem_load_4[24] = (((frag_mask_word_3 >> 24 & 1) != 0) ? -CAKE_INF : _tmem_load_4[24]); + _tmem_load_4[25] = (((frag_mask_word_3 >> 25 & 1) != 0) ? -CAKE_INF : _tmem_load_4[25]); + _tmem_load_4[26] = (((frag_mask_word_3 >> 26 & 1) != 0) ? -CAKE_INF : _tmem_load_4[26]); + _tmem_load_4[27] = (((frag_mask_word_3 >> 27 & 1) != 0) ? -CAKE_INF : _tmem_load_4[27]); + _tmem_load_4[28] = (((frag_mask_word_3 >> 28 & 1) != 0) ? -CAKE_INF : _tmem_load_4[28]); + _tmem_load_4[29] = (((frag_mask_word_3 >> 29 & 1) != 0) ? -CAKE_INF : _tmem_load_4[29]); + _tmem_load_4[30] = (((frag_mask_word_3 >> 30 & 1) != 0) ? -CAKE_INF : _tmem_load_4[30]); + _tmem_load_4[31] = (((frag_mask_word_3 >> 31 & 1) != 0) ? -CAKE_INF : _tmem_load_4[31]); + } + float2 _reg_reduce_max2_7 = {-CAKE_INF, -CAKE_INF}; + _reg_reduce_max2_7.x = max_noftz(_reg_reduce_max2_7.x, max_noftz(_tmem_load_4[0], _tmem_load_4[1])); + _reg_reduce_max2_7.y = max_noftz(_reg_reduce_max2_7.y, max_noftz(_tmem_load_4[2], _tmem_load_4[3])); + _reg_reduce_max2_7.x = max_noftz(_reg_reduce_max2_7.x, max_noftz(_tmem_load_4[4], _tmem_load_4[5])); + _reg_reduce_max2_7.y = max_noftz(_reg_reduce_max2_7.y, max_noftz(_tmem_load_4[6], _tmem_load_4[7])); + _reg_reduce_max2_7.x = max_noftz(_reg_reduce_max2_7.x, max_noftz(_tmem_load_4[8], _tmem_load_4[9])); + _reg_reduce_max2_7.y = max_noftz(_reg_reduce_max2_7.y, max_noftz(_tmem_load_4[10], _tmem_load_4[11])); + _reg_reduce_max2_7.x = max_noftz(_reg_reduce_max2_7.x, max_noftz(_tmem_load_4[12], _tmem_load_4[13])); + _reg_reduce_max2_7.y = max_noftz(_reg_reduce_max2_7.y, max_noftz(_tmem_load_4[14], _tmem_load_4[15])); + _reg_reduce_max2_7.x = max_noftz(_reg_reduce_max2_7.x, max_noftz(_tmem_load_4[16], _tmem_load_4[17])); + _reg_reduce_max2_7.y = max_noftz(_reg_reduce_max2_7.y, max_noftz(_tmem_load_4[18], _tmem_load_4[19])); + _reg_reduce_max2_7.x = max_noftz(_reg_reduce_max2_7.x, max_noftz(_tmem_load_4[20], _tmem_load_4[21])); + _reg_reduce_max2_7.y = max_noftz(_reg_reduce_max2_7.y, max_noftz(_tmem_load_4[22], _tmem_load_4[23])); + _reg_reduce_max2_7.x = max_noftz(_reg_reduce_max2_7.x, max_noftz(_tmem_load_4[24], _tmem_load_4[25])); + _reg_reduce_max2_7.y = max_noftz(_reg_reduce_max2_7.y, max_noftz(_tmem_load_4[26], _tmem_load_4[27])); + _reg_reduce_max2_7.x = max_noftz(_reg_reduce_max2_7.x, max_noftz(_tmem_load_4[28], _tmem_load_4[29])); + _reg_reduce_max2_7.y = max_noftz(_reg_reduce_max2_7.y, max_noftz(_tmem_load_4[30], _tmem_load_4[31])); + float _tmem_load_4_max = row_max_reduce(_reg_reduce_max2_7); + float _max_6 = max_noftz(new_max, _tmem_load_4_max); + new_max = _max_6; } float _fma_0 = __fmaf_rn(row_max_val, softmax_scale_log2, (-new_max) * softmax_scale_log2); float delta = _fma_0; @@ -1825,187 +2165,367 @@ kernel_cake_dsv4_75827bd38b443de66f67(const __grid_constant__ CUtensorMap tmap_q float max_scaled = safe_max * softmax_scale_log2; float block_sum = 0.0f; { - int p_off = ((phase != 0) ? 224 : 96); - int p_base = taddr + (unsigned int)p_off + (unsigned int)(tmem_row_base << 16); + int p_off_all = ((phase != 0) ? 224 : 96); + int p_base_all = taddr + (unsigned int)p_off_all + (unsigned int)(tmem_row_base << 16); + const float2 _fma_b2_8 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_9 = {-max_scaled, -max_scaled}; + float2 _fma_pair_10 = fma_f32x2(make_float2(_tmem_load_0[0], _tmem_load_0[1]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[0] = _fma_pair_10.x; + _tmem_load_0[1] = _fma_pair_10.y; + float2 _fma_pair_11 = fma_f32x2(make_float2(_tmem_load_0[2], _tmem_load_0[3]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[2] = _fma_pair_11.x; + _tmem_load_0[3] = _fma_pair_11.y; + float2 _fma_pair_12 = fma_f32x2(make_float2(_tmem_load_0[4], _tmem_load_0[5]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[4] = _fma_pair_12.x; + _tmem_load_0[5] = _fma_pair_12.y; + float2 _fma_pair_13 = fma_f32x2(make_float2(_tmem_load_0[6], _tmem_load_0[7]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[6] = _fma_pair_13.x; + _tmem_load_0[7] = _fma_pair_13.y; + float2 _fma_pair_14 = fma_f32x2(make_float2(_tmem_load_0[8], _tmem_load_0[9]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[8] = _fma_pair_14.x; + _tmem_load_0[9] = _fma_pair_14.y; + float2 _fma_pair_15 = fma_f32x2(make_float2(_tmem_load_0[10], _tmem_load_0[11]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[10] = _fma_pair_15.x; + _tmem_load_0[11] = _fma_pair_15.y; + float2 _fma_pair_16 = fma_f32x2(make_float2(_tmem_load_0[12], _tmem_load_0[13]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[12] = _fma_pair_16.x; + _tmem_load_0[13] = _fma_pair_16.y; + float2 _fma_pair_17 = fma_f32x2(make_float2(_tmem_load_0[14], _tmem_load_0[15]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[14] = _fma_pair_17.x; + _tmem_load_0[15] = _fma_pair_17.y; + float2 _fma_pair_18 = fma_f32x2(make_float2(_tmem_load_0[16], _tmem_load_0[17]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[16] = _fma_pair_18.x; + _tmem_load_0[17] = _fma_pair_18.y; + float2 _fma_pair_19 = fma_f32x2(make_float2(_tmem_load_0[18], _tmem_load_0[19]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[18] = _fma_pair_19.x; + _tmem_load_0[19] = _fma_pair_19.y; + float2 _fma_pair_20 = fma_f32x2(make_float2(_tmem_load_0[20], _tmem_load_0[21]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[20] = _fma_pair_20.x; + _tmem_load_0[21] = _fma_pair_20.y; + float2 _fma_pair_21 = fma_f32x2(make_float2(_tmem_load_0[22], _tmem_load_0[23]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[22] = _fma_pair_21.x; + _tmem_load_0[23] = _fma_pair_21.y; + float2 _fma_pair_22 = fma_f32x2(make_float2(_tmem_load_0[24], _tmem_load_0[25]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[24] = _fma_pair_22.x; + _tmem_load_0[25] = _fma_pair_22.y; + float2 _fma_pair_23 = fma_f32x2(make_float2(_tmem_load_0[26], _tmem_load_0[27]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[26] = _fma_pair_23.x; + _tmem_load_0[27] = _fma_pair_23.y; + float2 _fma_pair_24 = fma_f32x2(make_float2(_tmem_load_0[28], _tmem_load_0[29]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[28] = _fma_pair_24.x; + _tmem_load_0[29] = _fma_pair_24.y; + float2 _fma_pair_25 = fma_f32x2(make_float2(_tmem_load_0[30], _tmem_load_0[31]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[30] = _fma_pair_25.x; + _tmem_load_0[31] = _fma_pair_25.y; #pragma unroll - for (int frag_1 = 0; frag_1 < 4; frag_1++) { - float _tmem_load_3[32]; - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]) - : "r"(s_base + frag_1 * 32)); - int frag_valid_1 = valid_sparse_cols - frag_1 * 32; - uint32_t _slice_lo_mask_3; - { - int _lim_2 = frag_valid_1; - if (_lim_2 <= 0) { _slice_lo_mask_3 = 0u; } - else if (_lim_2 >= 32) { _slice_lo_mask_3 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_3) : "r"(_lim_2)); - } + for (int _le = 0; _le < 32; _le++) { + _tmem_load_0[_le] = approx_exp2(_tmem_load_0[_le]); + } + float2 _reg_reduce_sum2_26 = make_float2(0.0f, 0.0f); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[0], _tmem_load_0[1])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[2], _tmem_load_0[3])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[4], _tmem_load_0[5])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[6], _tmem_load_0[7])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[8], _tmem_load_0[9])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[10], _tmem_load_0[11])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[12], _tmem_load_0[13])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[14], _tmem_load_0[15])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[16], _tmem_load_0[17])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[18], _tmem_load_0[19])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[20], _tmem_load_0[21])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[22], _tmem_load_0[23])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[24], _tmem_load_0[25])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[26], _tmem_load_0[27])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[28], _tmem_load_0[29])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[30], _tmem_load_0[31])); + float _tmem_load_0_sum = _reg_reduce_sum2_26.x + _reg_reduce_sum2_26.y; + block_sum = block_sum + _tmem_load_0_sum; + { + uint32_t _pv_packed[8]; + #pragma unroll + for (int _j = 0; _j < 8; _j++) { + uint32_t _pk; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_pk) : "f"(_tmem_load_0[0 + _j * 4]), "f"(_tmem_load_0[0 + _j * 4 + 1]), + "f"(_tmem_load_0[0 + _j * 4 + 2]), "f"(_tmem_load_0[0 + _j * 4 + 3])); + _pv_packed[_j] = _pk; } - if (!(_slice_lo_mask_3 & (1u << 0))) _tmem_load_3[0] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 1))) _tmem_load_3[1] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 2))) _tmem_load_3[2] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 3))) _tmem_load_3[3] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 4))) _tmem_load_3[4] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 5))) _tmem_load_3[5] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 6))) _tmem_load_3[6] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 7))) _tmem_load_3[7] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 8))) _tmem_load_3[8] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 9))) _tmem_load_3[9] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 10))) _tmem_load_3[10] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 11))) _tmem_load_3[11] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 12))) _tmem_load_3[12] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 13))) _tmem_load_3[13] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 14))) _tmem_load_3[14] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 15))) _tmem_load_3[15] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 16))) _tmem_load_3[16] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 17))) _tmem_load_3[17] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 18))) _tmem_load_3[18] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 19))) _tmem_load_3[19] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 20))) _tmem_load_3[20] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 21))) _tmem_load_3[21] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 22))) _tmem_load_3[22] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 23))) _tmem_load_3[23] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 24))) _tmem_load_3[24] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 25))) _tmem_load_3[25] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 26))) _tmem_load_3[26] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 27))) _tmem_load_3[27] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 28))) _tmem_load_3[28] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 29))) _tmem_load_3[29] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 30))) _tmem_load_3[30] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 31))) _tmem_load_3[31] = -CAKE_INF; - if (any_invalid != 0) { - unsigned int frag_mask_word_1 = invalid_cols[frag_1]; - _tmem_load_3[0] = (((frag_mask_word_1 & 1) != 0) ? -CAKE_INF : _tmem_load_3[0]); - _tmem_load_3[1] = (((frag_mask_word_1 >> 1 & 1) != 0) ? -CAKE_INF : _tmem_load_3[1]); - _tmem_load_3[2] = (((frag_mask_word_1 >> 2 & 1) != 0) ? -CAKE_INF : _tmem_load_3[2]); - _tmem_load_3[3] = (((frag_mask_word_1 >> 3 & 1) != 0) ? -CAKE_INF : _tmem_load_3[3]); - _tmem_load_3[4] = (((frag_mask_word_1 >> 4 & 1) != 0) ? -CAKE_INF : _tmem_load_3[4]); - _tmem_load_3[5] = (((frag_mask_word_1 >> 5 & 1) != 0) ? -CAKE_INF : _tmem_load_3[5]); - _tmem_load_3[6] = (((frag_mask_word_1 >> 6 & 1) != 0) ? -CAKE_INF : _tmem_load_3[6]); - _tmem_load_3[7] = (((frag_mask_word_1 >> 7 & 1) != 0) ? -CAKE_INF : _tmem_load_3[7]); - _tmem_load_3[8] = (((frag_mask_word_1 >> 8 & 1) != 0) ? -CAKE_INF : _tmem_load_3[8]); - _tmem_load_3[9] = (((frag_mask_word_1 >> 9 & 1) != 0) ? -CAKE_INF : _tmem_load_3[9]); - _tmem_load_3[10] = (((frag_mask_word_1 >> 10 & 1) != 0) ? -CAKE_INF : _tmem_load_3[10]); - _tmem_load_3[11] = (((frag_mask_word_1 >> 11 & 1) != 0) ? -CAKE_INF : _tmem_load_3[11]); - _tmem_load_3[12] = (((frag_mask_word_1 >> 12 & 1) != 0) ? -CAKE_INF : _tmem_load_3[12]); - _tmem_load_3[13] = (((frag_mask_word_1 >> 13 & 1) != 0) ? -CAKE_INF : _tmem_load_3[13]); - _tmem_load_3[14] = (((frag_mask_word_1 >> 14 & 1) != 0) ? -CAKE_INF : _tmem_load_3[14]); - _tmem_load_3[15] = (((frag_mask_word_1 >> 15 & 1) != 0) ? -CAKE_INF : _tmem_load_3[15]); - _tmem_load_3[16] = (((frag_mask_word_1 >> 16 & 1) != 0) ? -CAKE_INF : _tmem_load_3[16]); - _tmem_load_3[17] = (((frag_mask_word_1 >> 17 & 1) != 0) ? -CAKE_INF : _tmem_load_3[17]); - _tmem_load_3[18] = (((frag_mask_word_1 >> 18 & 1) != 0) ? -CAKE_INF : _tmem_load_3[18]); - _tmem_load_3[19] = (((frag_mask_word_1 >> 19 & 1) != 0) ? -CAKE_INF : _tmem_load_3[19]); - _tmem_load_3[20] = (((frag_mask_word_1 >> 20 & 1) != 0) ? -CAKE_INF : _tmem_load_3[20]); - _tmem_load_3[21] = (((frag_mask_word_1 >> 21 & 1) != 0) ? -CAKE_INF : _tmem_load_3[21]); - _tmem_load_3[22] = (((frag_mask_word_1 >> 22 & 1) != 0) ? -CAKE_INF : _tmem_load_3[22]); - _tmem_load_3[23] = (((frag_mask_word_1 >> 23 & 1) != 0) ? -CAKE_INF : _tmem_load_3[23]); - _tmem_load_3[24] = (((frag_mask_word_1 >> 24 & 1) != 0) ? -CAKE_INF : _tmem_load_3[24]); - _tmem_load_3[25] = (((frag_mask_word_1 >> 25 & 1) != 0) ? -CAKE_INF : _tmem_load_3[25]); - _tmem_load_3[26] = (((frag_mask_word_1 >> 26 & 1) != 0) ? -CAKE_INF : _tmem_load_3[26]); - _tmem_load_3[27] = (((frag_mask_word_1 >> 27 & 1) != 0) ? -CAKE_INF : _tmem_load_3[27]); - _tmem_load_3[28] = (((frag_mask_word_1 >> 28 & 1) != 0) ? -CAKE_INF : _tmem_load_3[28]); - _tmem_load_3[29] = (((frag_mask_word_1 >> 29 & 1) != 0) ? -CAKE_INF : _tmem_load_3[29]); - _tmem_load_3[30] = (((frag_mask_word_1 >> 30 & 1) != 0) ? -CAKE_INF : _tmem_load_3[30]); - _tmem_load_3[31] = (((frag_mask_word_1 >> 31 & 1) != 0) ? -CAKE_INF : _tmem_load_3[31]); + tmem_st_x8_u32(p_base_all, _pv_packed); + } + const float2 _fma_b2_27 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_28 = {-max_scaled, -max_scaled}; + float2 _fma_pair_29 = fma_f32x2(make_float2(_tmem_load_2[0], _tmem_load_2[1]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[0] = _fma_pair_29.x; + _tmem_load_2[1] = _fma_pair_29.y; + float2 _fma_pair_30 = fma_f32x2(make_float2(_tmem_load_2[2], _tmem_load_2[3]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[2] = _fma_pair_30.x; + _tmem_load_2[3] = _fma_pair_30.y; + float2 _fma_pair_31 = fma_f32x2(make_float2(_tmem_load_2[4], _tmem_load_2[5]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[4] = _fma_pair_31.x; + _tmem_load_2[5] = _fma_pair_31.y; + float2 _fma_pair_32 = fma_f32x2(make_float2(_tmem_load_2[6], _tmem_load_2[7]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[6] = _fma_pair_32.x; + _tmem_load_2[7] = _fma_pair_32.y; + float2 _fma_pair_33 = fma_f32x2(make_float2(_tmem_load_2[8], _tmem_load_2[9]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[8] = _fma_pair_33.x; + _tmem_load_2[9] = _fma_pair_33.y; + float2 _fma_pair_34 = fma_f32x2(make_float2(_tmem_load_2[10], _tmem_load_2[11]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[10] = _fma_pair_34.x; + _tmem_load_2[11] = _fma_pair_34.y; + float2 _fma_pair_35 = fma_f32x2(make_float2(_tmem_load_2[12], _tmem_load_2[13]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[12] = _fma_pair_35.x; + _tmem_load_2[13] = _fma_pair_35.y; + float2 _fma_pair_36 = fma_f32x2(make_float2(_tmem_load_2[14], _tmem_load_2[15]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[14] = _fma_pair_36.x; + _tmem_load_2[15] = _fma_pair_36.y; + float2 _fma_pair_37 = fma_f32x2(make_float2(_tmem_load_2[16], _tmem_load_2[17]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[16] = _fma_pair_37.x; + _tmem_load_2[17] = _fma_pair_37.y; + float2 _fma_pair_38 = fma_f32x2(make_float2(_tmem_load_2[18], _tmem_load_2[19]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[18] = _fma_pair_38.x; + _tmem_load_2[19] = _fma_pair_38.y; + float2 _fma_pair_39 = fma_f32x2(make_float2(_tmem_load_2[20], _tmem_load_2[21]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[20] = _fma_pair_39.x; + _tmem_load_2[21] = _fma_pair_39.y; + float2 _fma_pair_40 = fma_f32x2(make_float2(_tmem_load_2[22], _tmem_load_2[23]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[22] = _fma_pair_40.x; + _tmem_load_2[23] = _fma_pair_40.y; + float2 _fma_pair_41 = fma_f32x2(make_float2(_tmem_load_2[24], _tmem_load_2[25]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[24] = _fma_pair_41.x; + _tmem_load_2[25] = _fma_pair_41.y; + float2 _fma_pair_42 = fma_f32x2(make_float2(_tmem_load_2[26], _tmem_load_2[27]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[26] = _fma_pair_42.x; + _tmem_load_2[27] = _fma_pair_42.y; + float2 _fma_pair_43 = fma_f32x2(make_float2(_tmem_load_2[28], _tmem_load_2[29]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[28] = _fma_pair_43.x; + _tmem_load_2[29] = _fma_pair_43.y; + float2 _fma_pair_44 = fma_f32x2(make_float2(_tmem_load_2[30], _tmem_load_2[31]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[30] = _fma_pair_44.x; + _tmem_load_2[31] = _fma_pair_44.y; + #pragma unroll + for (int _le = 0; _le < 32; _le++) { + _tmem_load_2[_le] = approx_exp2(_tmem_load_2[_le]); + } + float2 _reg_reduce_sum2_45 = make_float2(0.0f, 0.0f); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[0], _tmem_load_2[1])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[2], _tmem_load_2[3])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[4], _tmem_load_2[5])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[6], _tmem_load_2[7])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[8], _tmem_load_2[9])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[10], _tmem_load_2[11])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[12], _tmem_load_2[13])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[14], _tmem_load_2[15])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[16], _tmem_load_2[17])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[18], _tmem_load_2[19])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[20], _tmem_load_2[21])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[22], _tmem_load_2[23])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[24], _tmem_load_2[25])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[26], _tmem_load_2[27])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[28], _tmem_load_2[29])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[30], _tmem_load_2[31])); + float _tmem_load_2_sum = _reg_reduce_sum2_45.x + _reg_reduce_sum2_45.y; + block_sum = block_sum + _tmem_load_2_sum; + { + uint32_t _pv_packed[8]; + #pragma unroll + for (int _j = 0; _j < 8; _j++) { + uint32_t _pk; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_pk) : "f"(_tmem_load_2[0 + _j * 4]), "f"(_tmem_load_2[0 + _j * 4 + 1]), + "f"(_tmem_load_2[0 + _j * 4 + 2]), "f"(_tmem_load_2[0 + _j * 4 + 3])); + _pv_packed[_j] = _pk; } - const float2 _fma_b2_3 = {softmax_scale_log2, softmax_scale_log2}; - const float2 _fma_c2_4 = {-max_scaled, -max_scaled}; - float2 _fma_pair_5 = fma_f32x2(make_float2(_tmem_load_3[0], _tmem_load_3[1]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[0] = _fma_pair_5.x; - _tmem_load_3[1] = _fma_pair_5.y; - float2 _fma_pair_6 = fma_f32x2(make_float2(_tmem_load_3[2], _tmem_load_3[3]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[2] = _fma_pair_6.x; - _tmem_load_3[3] = _fma_pair_6.y; - float2 _fma_pair_7 = fma_f32x2(make_float2(_tmem_load_3[4], _tmem_load_3[5]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[4] = _fma_pair_7.x; - _tmem_load_3[5] = _fma_pair_7.y; - float2 _fma_pair_8 = fma_f32x2(make_float2(_tmem_load_3[6], _tmem_load_3[7]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[6] = _fma_pair_8.x; - _tmem_load_3[7] = _fma_pair_8.y; - float2 _fma_pair_9 = fma_f32x2(make_float2(_tmem_load_3[8], _tmem_load_3[9]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[8] = _fma_pair_9.x; - _tmem_load_3[9] = _fma_pair_9.y; - float2 _fma_pair_10 = fma_f32x2(make_float2(_tmem_load_3[10], _tmem_load_3[11]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[10] = _fma_pair_10.x; - _tmem_load_3[11] = _fma_pair_10.y; - float2 _fma_pair_11 = fma_f32x2(make_float2(_tmem_load_3[12], _tmem_load_3[13]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[12] = _fma_pair_11.x; - _tmem_load_3[13] = _fma_pair_11.y; - float2 _fma_pair_12 = fma_f32x2(make_float2(_tmem_load_3[14], _tmem_load_3[15]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[14] = _fma_pair_12.x; - _tmem_load_3[15] = _fma_pair_12.y; - float2 _fma_pair_13 = fma_f32x2(make_float2(_tmem_load_3[16], _tmem_load_3[17]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[16] = _fma_pair_13.x; - _tmem_load_3[17] = _fma_pair_13.y; - float2 _fma_pair_14 = fma_f32x2(make_float2(_tmem_load_3[18], _tmem_load_3[19]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[18] = _fma_pair_14.x; - _tmem_load_3[19] = _fma_pair_14.y; - float2 _fma_pair_15 = fma_f32x2(make_float2(_tmem_load_3[20], _tmem_load_3[21]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[20] = _fma_pair_15.x; - _tmem_load_3[21] = _fma_pair_15.y; - float2 _fma_pair_16 = fma_f32x2(make_float2(_tmem_load_3[22], _tmem_load_3[23]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[22] = _fma_pair_16.x; - _tmem_load_3[23] = _fma_pair_16.y; - float2 _fma_pair_17 = fma_f32x2(make_float2(_tmem_load_3[24], _tmem_load_3[25]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[24] = _fma_pair_17.x; - _tmem_load_3[25] = _fma_pair_17.y; - float2 _fma_pair_18 = fma_f32x2(make_float2(_tmem_load_3[26], _tmem_load_3[27]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[26] = _fma_pair_18.x; - _tmem_load_3[27] = _fma_pair_18.y; - float2 _fma_pair_19 = fma_f32x2(make_float2(_tmem_load_3[28], _tmem_load_3[29]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[28] = _fma_pair_19.x; - _tmem_load_3[29] = _fma_pair_19.y; - float2 _fma_pair_20 = fma_f32x2(make_float2(_tmem_load_3[30], _tmem_load_3[31]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[30] = _fma_pair_20.x; - _tmem_load_3[31] = _fma_pair_20.y; + tmem_st_x8_u32(p_base_all + 8, _pv_packed); + } + const float2 _fma_b2_46 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_47 = {-max_scaled, -max_scaled}; + float2 _fma_pair_48 = fma_f32x2(make_float2(_tmem_load_3[0], _tmem_load_3[1]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[0] = _fma_pair_48.x; + _tmem_load_3[1] = _fma_pair_48.y; + float2 _fma_pair_49 = fma_f32x2(make_float2(_tmem_load_3[2], _tmem_load_3[3]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[2] = _fma_pair_49.x; + _tmem_load_3[3] = _fma_pair_49.y; + float2 _fma_pair_50 = fma_f32x2(make_float2(_tmem_load_3[4], _tmem_load_3[5]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[4] = _fma_pair_50.x; + _tmem_load_3[5] = _fma_pair_50.y; + float2 _fma_pair_51 = fma_f32x2(make_float2(_tmem_load_3[6], _tmem_load_3[7]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[6] = _fma_pair_51.x; + _tmem_load_3[7] = _fma_pair_51.y; + float2 _fma_pair_52 = fma_f32x2(make_float2(_tmem_load_3[8], _tmem_load_3[9]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[8] = _fma_pair_52.x; + _tmem_load_3[9] = _fma_pair_52.y; + float2 _fma_pair_53 = fma_f32x2(make_float2(_tmem_load_3[10], _tmem_load_3[11]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[10] = _fma_pair_53.x; + _tmem_load_3[11] = _fma_pair_53.y; + float2 _fma_pair_54 = fma_f32x2(make_float2(_tmem_load_3[12], _tmem_load_3[13]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[12] = _fma_pair_54.x; + _tmem_load_3[13] = _fma_pair_54.y; + float2 _fma_pair_55 = fma_f32x2(make_float2(_tmem_load_3[14], _tmem_load_3[15]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[14] = _fma_pair_55.x; + _tmem_load_3[15] = _fma_pair_55.y; + float2 _fma_pair_56 = fma_f32x2(make_float2(_tmem_load_3[16], _tmem_load_3[17]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[16] = _fma_pair_56.x; + _tmem_load_3[17] = _fma_pair_56.y; + float2 _fma_pair_57 = fma_f32x2(make_float2(_tmem_load_3[18], _tmem_load_3[19]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[18] = _fma_pair_57.x; + _tmem_load_3[19] = _fma_pair_57.y; + float2 _fma_pair_58 = fma_f32x2(make_float2(_tmem_load_3[20], _tmem_load_3[21]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[20] = _fma_pair_58.x; + _tmem_load_3[21] = _fma_pair_58.y; + float2 _fma_pair_59 = fma_f32x2(make_float2(_tmem_load_3[22], _tmem_load_3[23]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[22] = _fma_pair_59.x; + _tmem_load_3[23] = _fma_pair_59.y; + float2 _fma_pair_60 = fma_f32x2(make_float2(_tmem_load_3[24], _tmem_load_3[25]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[24] = _fma_pair_60.x; + _tmem_load_3[25] = _fma_pair_60.y; + float2 _fma_pair_61 = fma_f32x2(make_float2(_tmem_load_3[26], _tmem_load_3[27]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[26] = _fma_pair_61.x; + _tmem_load_3[27] = _fma_pair_61.y; + float2 _fma_pair_62 = fma_f32x2(make_float2(_tmem_load_3[28], _tmem_load_3[29]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[28] = _fma_pair_62.x; + _tmem_load_3[29] = _fma_pair_62.y; + float2 _fma_pair_63 = fma_f32x2(make_float2(_tmem_load_3[30], _tmem_load_3[31]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[30] = _fma_pair_63.x; + _tmem_load_3[31] = _fma_pair_63.y; + #pragma unroll + for (int _le = 0; _le < 32; _le++) { + _tmem_load_3[_le] = approx_exp2(_tmem_load_3[_le]); + } + float2 _reg_reduce_sum2_64 = make_float2(0.0f, 0.0f); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[0], _tmem_load_3[1])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[2], _tmem_load_3[3])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[4], _tmem_load_3[5])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[6], _tmem_load_3[7])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[8], _tmem_load_3[9])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[10], _tmem_load_3[11])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[12], _tmem_load_3[13])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[14], _tmem_load_3[15])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[16], _tmem_load_3[17])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[18], _tmem_load_3[19])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[20], _tmem_load_3[21])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[22], _tmem_load_3[23])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[24], _tmem_load_3[25])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[26], _tmem_load_3[27])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[28], _tmem_load_3[29])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[30], _tmem_load_3[31])); + float _tmem_load_3_sum = _reg_reduce_sum2_64.x + _reg_reduce_sum2_64.y; + block_sum = block_sum + _tmem_load_3_sum; + { + uint32_t _pv_packed[8]; #pragma unroll - for (int _le = 0; _le < 32; _le++) { - _tmem_load_3[_le] = approx_exp2(_tmem_load_3[_le]); + for (int _j = 0; _j < 8; _j++) { + uint32_t _pk; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_pk) : "f"(_tmem_load_3[0 + _j * 4]), "f"(_tmem_load_3[0 + _j * 4 + 1]), + "f"(_tmem_load_3[0 + _j * 4 + 2]), "f"(_tmem_load_3[0 + _j * 4 + 3])); + _pv_packed[_j] = _pk; } - float2 _reg_reduce_sum2_21 = make_float2(0.0f, 0.0f); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[0], _tmem_load_3[1])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[2], _tmem_load_3[3])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[4], _tmem_load_3[5])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[6], _tmem_load_3[7])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[8], _tmem_load_3[9])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[10], _tmem_load_3[11])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[12], _tmem_load_3[13])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[14], _tmem_load_3[15])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[16], _tmem_load_3[17])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[18], _tmem_load_3[19])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[20], _tmem_load_3[21])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[22], _tmem_load_3[23])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[24], _tmem_load_3[25])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[26], _tmem_load_3[27])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[28], _tmem_load_3[29])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[30], _tmem_load_3[31])); - float _tmem_load_3_sum = _reg_reduce_sum2_21.x + _reg_reduce_sum2_21.y; - block_sum = block_sum + _tmem_load_3_sum; - { - uint32_t _pv_packed[8]; - #pragma unroll - for (int _j = 0; _j < 8; _j++) { - uint32_t _pk; - asm("{\n\t" - ".reg .b16 _lo, _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}\n" - : "=r"(_pk) : "f"(_tmem_load_3[0 + _j * 4]), "f"(_tmem_load_3[0 + _j * 4 + 1]), - "f"(_tmem_load_3[0 + _j * 4 + 2]), "f"(_tmem_load_3[0 + _j * 4 + 3])); - _pv_packed[_j] = _pk; - } - tmem_st_x8_u32(p_base + frag_1 * 8, _pv_packed); + tmem_st_x8_u32(p_base_all + 16, _pv_packed); + } + const float2 _fma_b2_65 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_66 = {-max_scaled, -max_scaled}; + float2 _fma_pair_67 = fma_f32x2(make_float2(_tmem_load_4[0], _tmem_load_4[1]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[0] = _fma_pair_67.x; + _tmem_load_4[1] = _fma_pair_67.y; + float2 _fma_pair_68 = fma_f32x2(make_float2(_tmem_load_4[2], _tmem_load_4[3]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[2] = _fma_pair_68.x; + _tmem_load_4[3] = _fma_pair_68.y; + float2 _fma_pair_69 = fma_f32x2(make_float2(_tmem_load_4[4], _tmem_load_4[5]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[4] = _fma_pair_69.x; + _tmem_load_4[5] = _fma_pair_69.y; + float2 _fma_pair_70 = fma_f32x2(make_float2(_tmem_load_4[6], _tmem_load_4[7]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[6] = _fma_pair_70.x; + _tmem_load_4[7] = _fma_pair_70.y; + float2 _fma_pair_71 = fma_f32x2(make_float2(_tmem_load_4[8], _tmem_load_4[9]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[8] = _fma_pair_71.x; + _tmem_load_4[9] = _fma_pair_71.y; + float2 _fma_pair_72 = fma_f32x2(make_float2(_tmem_load_4[10], _tmem_load_4[11]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[10] = _fma_pair_72.x; + _tmem_load_4[11] = _fma_pair_72.y; + float2 _fma_pair_73 = fma_f32x2(make_float2(_tmem_load_4[12], _tmem_load_4[13]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[12] = _fma_pair_73.x; + _tmem_load_4[13] = _fma_pair_73.y; + float2 _fma_pair_74 = fma_f32x2(make_float2(_tmem_load_4[14], _tmem_load_4[15]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[14] = _fma_pair_74.x; + _tmem_load_4[15] = _fma_pair_74.y; + float2 _fma_pair_75 = fma_f32x2(make_float2(_tmem_load_4[16], _tmem_load_4[17]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[16] = _fma_pair_75.x; + _tmem_load_4[17] = _fma_pair_75.y; + float2 _fma_pair_76 = fma_f32x2(make_float2(_tmem_load_4[18], _tmem_load_4[19]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[18] = _fma_pair_76.x; + _tmem_load_4[19] = _fma_pair_76.y; + float2 _fma_pair_77 = fma_f32x2(make_float2(_tmem_load_4[20], _tmem_load_4[21]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[20] = _fma_pair_77.x; + _tmem_load_4[21] = _fma_pair_77.y; + float2 _fma_pair_78 = fma_f32x2(make_float2(_tmem_load_4[22], _tmem_load_4[23]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[22] = _fma_pair_78.x; + _tmem_load_4[23] = _fma_pair_78.y; + float2 _fma_pair_79 = fma_f32x2(make_float2(_tmem_load_4[24], _tmem_load_4[25]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[24] = _fma_pair_79.x; + _tmem_load_4[25] = _fma_pair_79.y; + float2 _fma_pair_80 = fma_f32x2(make_float2(_tmem_load_4[26], _tmem_load_4[27]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[26] = _fma_pair_80.x; + _tmem_load_4[27] = _fma_pair_80.y; + float2 _fma_pair_81 = fma_f32x2(make_float2(_tmem_load_4[28], _tmem_load_4[29]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[28] = _fma_pair_81.x; + _tmem_load_4[29] = _fma_pair_81.y; + float2 _fma_pair_82 = fma_f32x2(make_float2(_tmem_load_4[30], _tmem_load_4[31]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[30] = _fma_pair_82.x; + _tmem_load_4[31] = _fma_pair_82.y; + #pragma unroll + for (int _le = 0; _le < 32; _le++) { + _tmem_load_4[_le] = approx_exp2(_tmem_load_4[_le]); + } + float2 _reg_reduce_sum2_83 = make_float2(0.0f, 0.0f); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[0], _tmem_load_4[1])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[2], _tmem_load_4[3])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[4], _tmem_load_4[5])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[6], _tmem_load_4[7])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[8], _tmem_load_4[9])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[10], _tmem_load_4[11])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[12], _tmem_load_4[13])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[14], _tmem_load_4[15])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[16], _tmem_load_4[17])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[18], _tmem_load_4[19])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[20], _tmem_load_4[21])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[22], _tmem_load_4[23])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[24], _tmem_load_4[25])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[26], _tmem_load_4[27])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[28], _tmem_load_4[29])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[30], _tmem_load_4[31])); + float _tmem_load_4_sum = _reg_reduce_sum2_83.x + _reg_reduce_sum2_83.y; + block_sum = block_sum + _tmem_load_4_sum; + { + uint32_t _pv_packed[8]; + #pragma unroll + for (int _j = 0; _j < 8; _j++) { + uint32_t _pk; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_pk) : "f"(_tmem_load_4[0 + _j * 4]), "f"(_tmem_load_4[0 + _j * 4 + 1]), + "f"(_tmem_load_4[0 + _j * 4 + 2]), "f"(_tmem_load_4[0 + _j * 4 + 3])); + _pv_packed[_j] = _pk; } + tmem_st_x8_u32(p_base_all + 24, _pv_packed); } } { @@ -2079,13 +2599,13 @@ kernel_cake_dsv4_75827bd38b443de66f67(const __grid_constant__ CUtensorMap tmap_q int o_base = taddr + 256 + (unsigned int)(vs * 128) + (unsigned int)corr_row; #pragma unroll for (int c = 0; c < 128; c += 16) { - float _tmem_load_4[16]; - tmem_ld_x16(&_tmem_load_4[0], o_base + c); + float _tmem_load_5[16]; + tmem_ld_x16(&_tmem_load_5[0], o_base + c); const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; #pragma unroll for (int _ls = 0; _ls < 8; _ls++) - mul_f32x2_inplace(&reinterpret_cast(_tmem_load_4)[_ls], _scale2_0); - tmem_st_x16_f32(o_base + c, _tmem_load_4); + mul_f32x2_inplace(&reinterpret_cast(_tmem_load_5)[_ls], _scale2_0); + tmem_st_x16_f32(o_base + c, _tmem_load_5); asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); } } @@ -2118,11 +2638,11 @@ kernel_cake_dsv4_75827bd38b443de66f67(const __grid_constant__ CUtensorMap tmap_q int o_base_epi = taddr + 256 + (unsigned int)(vs_1 * 128) + (unsigned int)corr_row; #pragma unroll for (int c_1 = 0; c_1 < 128; c_1 += 32) { - float _tmem_load_5[32]; + float _tmem_load_6[32]; asm volatile( "tcgen05.ld.sync.aligned.32x32b.x32.b32" " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(_tmem_load_5[0]), "=f"(_tmem_load_5[1]), "=f"(_tmem_load_5[2]), "=f"(_tmem_load_5[3]), "=f"(_tmem_load_5[4]), "=f"(_tmem_load_5[5]), "=f"(_tmem_load_5[6]), "=f"(_tmem_load_5[7]), "=f"(_tmem_load_5[8]), "=f"(_tmem_load_5[9]), "=f"(_tmem_load_5[10]), "=f"(_tmem_load_5[11]), "=f"(_tmem_load_5[12]), "=f"(_tmem_load_5[13]), "=f"(_tmem_load_5[14]), "=f"(_tmem_load_5[15]), "=f"(_tmem_load_5[16]), "=f"(_tmem_load_5[17]), "=f"(_tmem_load_5[18]), "=f"(_tmem_load_5[19]), "=f"(_tmem_load_5[20]), "=f"(_tmem_load_5[21]), "=f"(_tmem_load_5[22]), "=f"(_tmem_load_5[23]), "=f"(_tmem_load_5[24]), "=f"(_tmem_load_5[25]), "=f"(_tmem_load_5[26]), "=f"(_tmem_load_5[27]), "=f"(_tmem_load_5[28]), "=f"(_tmem_load_5[29]), "=f"(_tmem_load_5[30]), "=f"(_tmem_load_5[31]) + : "=f"(_tmem_load_6[0]), "=f"(_tmem_load_6[1]), "=f"(_tmem_load_6[2]), "=f"(_tmem_load_6[3]), "=f"(_tmem_load_6[4]), "=f"(_tmem_load_6[5]), "=f"(_tmem_load_6[6]), "=f"(_tmem_load_6[7]), "=f"(_tmem_load_6[8]), "=f"(_tmem_load_6[9]), "=f"(_tmem_load_6[10]), "=f"(_tmem_load_6[11]), "=f"(_tmem_load_6[12]), "=f"(_tmem_load_6[13]), "=f"(_tmem_load_6[14]), "=f"(_tmem_load_6[15]), "=f"(_tmem_load_6[16]), "=f"(_tmem_load_6[17]), "=f"(_tmem_load_6[18]), "=f"(_tmem_load_6[19]), "=f"(_tmem_load_6[20]), "=f"(_tmem_load_6[21]), "=f"(_tmem_load_6[22]), "=f"(_tmem_load_6[23]), "=f"(_tmem_load_6[24]), "=f"(_tmem_load_6[25]), "=f"(_tmem_load_6[26]), "=f"(_tmem_load_6[27]), "=f"(_tmem_load_6[28]), "=f"(_tmem_load_6[29]), "=f"(_tmem_load_6[30]), "=f"(_tmem_load_6[31]) : "r"(o_base_epi + c_1)); int gmem_base = o_offset + vs_1 * 128 + c_1; if (head_idx < num_heads) { @@ -2133,28 +2653,28 @@ kernel_cake_dsv4_75827bd38b443de66f67(const __grid_constant__ CUtensorMap tmap_q #if __CUDA_ARCH__ >= 1000 #pragma unroll for (int _ps = 0; _ps < 8; _ps++) - mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_5[j])[_ps], _prescale2_1); + mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_6[j])[_ps], _prescale2_1); #else #pragma unroll for (int _ps = 0; _ps < 16; _ps++) - _tmem_load_5[j + _ps] *= inv_sum * output_scale; + _tmem_load_6[j + _ps] *= inv_sum * output_scale; #endif { - __nv_bfloat162 _pk0 = __floats2bfloat162_rn(_tmem_load_5[j + 0], _tmem_load_5[j + 1]); + __nv_bfloat162 _pk0 = __floats2bfloat162_rn(_tmem_load_6[j + 0], _tmem_load_6[j + 1]); unsigned _pk_u0 = *reinterpret_cast(&_pk0); - __nv_bfloat162 _pk1 = __floats2bfloat162_rn(_tmem_load_5[j + 2], _tmem_load_5[j + 3]); + __nv_bfloat162 _pk1 = __floats2bfloat162_rn(_tmem_load_6[j + 2], _tmem_load_6[j + 3]); unsigned _pk_u1 = *reinterpret_cast(&_pk1); - __nv_bfloat162 _pk2 = __floats2bfloat162_rn(_tmem_load_5[j + 4], _tmem_load_5[j + 5]); + __nv_bfloat162 _pk2 = __floats2bfloat162_rn(_tmem_load_6[j + 4], _tmem_load_6[j + 5]); unsigned _pk_u2 = *reinterpret_cast(&_pk2); - __nv_bfloat162 _pk3 = __floats2bfloat162_rn(_tmem_load_5[j + 6], _tmem_load_5[j + 7]); + __nv_bfloat162 _pk3 = __floats2bfloat162_rn(_tmem_load_6[j + 6], _tmem_load_6[j + 7]); unsigned _pk_u3 = *reinterpret_cast(&_pk3); - __nv_bfloat162 _pk4 = __floats2bfloat162_rn(_tmem_load_5[j + 8], _tmem_load_5[j + 9]); + __nv_bfloat162 _pk4 = __floats2bfloat162_rn(_tmem_load_6[j + 8], _tmem_load_6[j + 9]); unsigned _pk_u4 = *reinterpret_cast(&_pk4); - __nv_bfloat162 _pk5 = __floats2bfloat162_rn(_tmem_load_5[j + 10], _tmem_load_5[j + 11]); + __nv_bfloat162 _pk5 = __floats2bfloat162_rn(_tmem_load_6[j + 10], _tmem_load_6[j + 11]); unsigned _pk_u5 = *reinterpret_cast(&_pk5); - __nv_bfloat162 _pk6 = __floats2bfloat162_rn(_tmem_load_5[j + 12], _tmem_load_5[j + 13]); + __nv_bfloat162 _pk6 = __floats2bfloat162_rn(_tmem_load_6[j + 12], _tmem_load_6[j + 13]); unsigned _pk_u6 = *reinterpret_cast(&_pk6); - __nv_bfloat162 _pk7 = __floats2bfloat162_rn(_tmem_load_5[j + 14], _tmem_load_5[j + 15]); + __nv_bfloat162 _pk7 = __floats2bfloat162_rn(_tmem_load_6[j + 14], _tmem_load_6[j + 15]); unsigned _pk_u7 = *reinterpret_cast(&_pk7); asm volatile( "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_7771bcc5c13a5ff83890_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_5703afd5b398f2b6c980_binding.cu similarity index 97% rename from csrc/cake_dsv4/sm_100a/cake_dsv4_7771bcc5c13a5ff83890_binding.cu rename to csrc/cake_dsv4/sm_100a/cake_dsv4_5703afd5b398f2b6c980_binding.cu index 2df363fef17..e9e716552b4 100644 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_7771bcc5c13a5ff83890_binding.cu +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_5703afd5b398f2b6c980_binding.cu @@ -15,7 +15,7 @@ */ // Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_7771bcc5c13a5ff83890'. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_5703afd5b398f2b6c980'. #include #include #include @@ -31,10 +31,10 @@ #include #include -extern "C" __global__ void kernel_cake_dsv4_7771bcc5c13a5ff83890(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items); +extern "C" __global__ void kernel_cake_dsv4_5703afd5b398f2b6c980(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items); -namespace cake_host_shim_392ec7fc32b7041f { +namespace cake_host_shim_5678cf46270e8859 { using tvm::ffi::Optional; using tvm::ffi::TensorView; @@ -416,7 +416,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ } if (smem_status == cudaSuccess) { smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_7771bcc5c13a5ff83890), + reinterpret_cast(kernel_cake_dsv4_5703afd5b398f2b6c980), cudaFuncAttributeMaxDynamicSharedMemorySize, 223744); } @@ -426,7 +426,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ } } TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_7771bcc5c13a5ff83890 failed: " + << "cudaFuncSetAttribute for kernel_cake_dsv4_5703afd5b398f2b6c980 failed: " << cudaGetErrorString(smem_status); dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); dim3 block(512u, 1u, 1u); @@ -450,12 +450,12 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ config.attrs = attrs; config.numAttrs = n; cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_7771bcc5c13a5ff83890), kargs); + &config, reinterpret_cast(kernel_cake_dsv4_5703afd5b398f2b6c980), kargs); TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_7771bcc5c13a5ff83890 failed: " + << "cudaLaunchKernelExC for kernel_cake_dsv4_5703afd5b398f2b6c980 failed: " << cudaGetErrorString(launch_status); } -} // namespace cake_host_shim_392ec7fc32b7041f +} // namespace cake_host_shim_5678cf46270e8859 -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_392ec7fc32b7041f::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_5678cf46270e8859::Run); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_7771bcc5c13a5ff83890_kernel.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_5703afd5b398f2b6c980_kernel.cu similarity index 99% rename from csrc/cake_dsv4/sm_100a/cake_dsv4_7771bcc5c13a5ff83890_kernel.cu rename to csrc/cake_dsv4/sm_100a/cake_dsv4_5703afd5b398f2b6c980_kernel.cu index 048415f9878..e481fa741c8 100644 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_7771bcc5c13a5ff83890_kernel.cu +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_5703afd5b398f2b6c980_kernel.cu @@ -1220,7 +1220,7 @@ __device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { extern "C" { __global__ __launch_bounds__(512, 1) __cluster_dims__(2,1,1) void -kernel_cake_dsv4_7771bcc5c13a5ff83890(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items) +kernel_cake_dsv4_5703afd5b398f2b6c980(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); @@ -1626,10 +1626,10 @@ kernel_cake_dsv4_7771bcc5c13a5ff83890(const __grid_constant__ CUtensorMap tmap_q :: "r"((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); } } - if (load_warp_rank < 2) { - int v_dst = smem_v_full_addr + load_v_stage * 32768 + (unsigned int)(load_warp_rank * 16384); - int v_col = v_chunk * 256 + load_warp_rank * 128; - int group_2 = lane; + if (lane < 16) { + int v_dst = smem_v_full_addr + load_v_stage * 32768 + (unsigned int)(load_warp_rank / 2 * 16384); + int v_col = v_chunk * 256 + load_warp_rank / 2 * 128; + int group_2 = (load_warp_rank & 1) * 16 + lane; int group_offset_2 = (prev_tile & 1) * 128 + group_2 * 4; int raw_rows_2[4]; asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" @@ -1668,10 +1668,10 @@ kernel_cake_dsv4_7771bcc5c13a5ff83890(const __grid_constant__ CUtensorMap tmap_q :: "r"((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); } } - if (load_warp_rank < 2) { - int v_dst_1 = smem_v_full_addr + load_v_stage * 32768 + (unsigned int)(load_warp_rank * 16384); - int v_col_1 = v_chunk * 256 + load_warp_rank * 128; - int group_3 = lane; + if (lane < 16) { + int v_dst_1 = smem_v_full_addr + load_v_stage * 32768 + (unsigned int)(load_warp_rank / 2 * 16384); + int v_col_1 = v_chunk * 256 + load_warp_rank / 2 * 128; + int group_3 = (load_warp_rank & 1) * 16 + lane; int group_offset_3 = (last_tile & 1) * 128 + group_3 * 4; int raw_rows_3[4]; asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_ba9ed6a6233761682cc3_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_714df747e199d0adbe08_binding.cu similarity index 86% rename from csrc/cake_dsv4/sm_100a/cake_dsv4_ba9ed6a6233761682cc3_binding.cu rename to csrc/cake_dsv4/sm_100a/cake_dsv4_714df747e199d0adbe08_binding.cu index 5883041ded4..d7c15323daa 100644 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_ba9ed6a6233761682cc3_binding.cu +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_714df747e199d0adbe08_binding.cu @@ -15,7 +15,7 @@ */ // Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_ba9ed6a6233761682cc3'. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_714df747e199d0adbe08'. #include #include #include @@ -31,10 +31,10 @@ #include #include -extern "C" __global__ void kernel_cake_dsv4_ba9ed6a6233761682cc3(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); +extern "C" __global__ void kernel_cake_dsv4_714df747e199d0adbe08(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, const __grid_constant__ CUtensorMap O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); -namespace cake_host_shim_37348cf6921b17a8 { +namespace cake_host_shim_dab7193014e792b3 { using tvm::ffi::Optional; using tvm::ffi::TensorView; @@ -390,6 +390,57 @@ inline CUtensorMap EncodeTma_tmap_compressed_v(const TensorView& t) { return tm; } +// 4D TMA descriptor for buffer 'O' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_O(const TensorView& t, int64_t tp_num_heads, int64_t tp_num_query_tokens) { + TVM_FFI_CHECK(t.ndim() >= 1, ValueError) + << "TMA source 'O' must have at least 1 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'O' must have unit innermost stride, got " << t.stride(-1); + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(8), (uint64_t)(5), (uint64_t)((tp_num_query_tokens * tp_num_heads))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'O' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1] && 1u <= global_dim[2] && 32u <= global_dim[3], ValueError) + << "TMA box (64, 1, 1, 32) exceeds resolved global dims for 'O'"; + int64_t carrier_stride_0 = 64; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 512; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = 2560; + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 1u, 1u, 32u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'O') failed: CUresult=" << (int)r; + return tm; +} + void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_compressed_k, TensorView arg_tmap_swa_v, TensorView arg_tmap_compressed_v, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_num_query_tokens, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z) { DLDevice dev = arg_tmap_q.device(); ScopedCudaDevice device_guard(dev.device_id); @@ -477,7 +528,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_c CUtensorMap p_tmap_compressed_k = EncodeTma_tmap_compressed_k(arg_tmap_compressed_k); CUtensorMap p_tmap_swa_v = EncodeTma_tmap_swa_v(arg_tmap_swa_v); CUtensorMap p_tmap_compressed_v = EncodeTma_tmap_compressed_v(arg_tmap_compressed_v); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); + CUtensorMap p_O = EncodeTma_O(arg_O, (int64_t)arg_num_heads, (int64_t)arg_num_query_tokens); float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); @@ -510,7 +561,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_c } if (smem_status == cudaSuccess) { smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_ba9ed6a6233761682cc3), + reinterpret_cast(kernel_cake_dsv4_714df747e199d0adbe08), cudaFuncAttributeMaxDynamicSharedMemorySize, 231808); } @@ -520,7 +571,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_c } } TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_ba9ed6a6233761682cc3 failed: " + << "cudaFuncSetAttribute for kernel_cake_dsv4_714df747e199d0adbe08 failed: " << cudaGetErrorString(smem_status); dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); dim3 block(512u, 1u, 1u); @@ -544,12 +595,12 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_c config.attrs = attrs; config.numAttrs = n; cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_ba9ed6a6233761682cc3), kargs); + &config, reinterpret_cast(kernel_cake_dsv4_714df747e199d0adbe08), kargs); TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_ba9ed6a6233761682cc3 failed: " + << "cudaLaunchKernelExC for kernel_cake_dsv4_714df747e199d0adbe08 failed: " << cudaGetErrorString(launch_status); } -} // namespace cake_host_shim_37348cf6921b17a8 +} // namespace cake_host_shim_dab7193014e792b3 -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_37348cf6921b17a8::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_dab7193014e792b3::Run); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_ba9ed6a6233761682cc3_kernel.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_714df747e199d0adbe08_kernel.cu similarity index 96% rename from csrc/cake_dsv4/sm_100a/cake_dsv4_ba9ed6a6233761682cc3_kernel.cu rename to csrc/cake_dsv4/sm_100a/cake_dsv4_714df747e199d0adbe08_kernel.cu index 5677c4f50cd..cec2500ed5f 100644 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_ba9ed6a6233761682cc3_kernel.cu +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_714df747e199d0adbe08_kernel.cu @@ -1181,6 +1181,15 @@ __device__ __forceinline__ void tma_gather4_gmem2smem_mc_cta2( } +__device__ __forceinline__ void tma_store_4d( + const void *tmap, int x, int y, int z, int w, unsigned smem_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.global.shared::cta.tile.bulk_group" + " [%0, {%1, %2, %3, %4}], [%5];" + :: "l"(tmap), "r"(x), "r"(y), "r"(z), "r"(w), "r"(smem_addr) : "memory"); +} + + __device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { asm volatile( "{\n\t" @@ -1195,7 +1204,7 @@ __device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint extern "C" { __global__ __launch_bounds__(512, 1) __cluster_dims__(2,1,1) void -kernel_cake_dsv4_ba9ed6a6233761682cc3(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items) +kernel_cake_dsv4_714df747e199d0adbe08(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, const __grid_constant__ CUtensorMap O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); @@ -1221,10 +1230,11 @@ kernel_cake_dsv4_ba9ed6a6233761682cc3(const __grid_constant__ CUtensorMap tmap_q #define stats_addr (mbar_base + 184) #define stats_empty_addr (mbar_base + 200) #define o_full_addr (mbar_base + 216) - #define o_first_slice_seeded_addr (mbar_base + 224) - #define tmem_dealloc_peer_addr (mbar_base + 232) - #define index_valid_full_addr (mbar_base + 240) - #define index_valid_empty_addr (mbar_base + 256) + #define epi_q_free_addr (mbar_base + 224) + #define o_first_slice_seeded_addr (mbar_base + 232) + #define tmem_dealloc_peer_addr (mbar_base + 240) + #define index_valid_full_addr (mbar_base + 248) + #define index_valid_empty_addr (mbar_base + 264) const int bid = blockIdx.x; const int num_bids = gridDim.x; @@ -1257,8 +1267,8 @@ kernel_cake_dsv4_ba9ed6a6233761682cc3(const __grid_constant__ CUtensorMap tmap_q asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); - // Mbarrier init (18 pipeline groups, 0 ordered-sequence groups, 34 barriers) - // Mbarriers at smem_raw[231488..231760) + // Mbarrier init (19 pipeline groups, 0 ordered-sequence groups, 35 barriers) + // Mbarriers at smem_raw[231488..231768) if (warp == 0) { uint32_t leader = elect_sync(); @@ -1307,16 +1317,18 @@ kernel_cake_dsv4_ba9ed6a6233761682cc3(const __grid_constant__ CUtensorMap tmap_q mbarrier_init(smem + 231696, 128); // o_full: 1 barriers, init_count=1 mbarrier_init(smem + 231704, 1); + // epi_q_free: 1 barriers, init_count=4 + mbarrier_init(smem + 231712, 4); // o_first_slice_seeded: 1 barriers, init_count=256 - mbarrier_init(smem + 231712, 256); + mbarrier_init(smem + 231720, 256); // tmem_dealloc_peer: 1 barriers, init_count=32 - mbarrier_init(smem + 231720, 32); - // index_valid_full: 2 barriers, init_count=32 mbarrier_init(smem + 231728, 32); + // index_valid_full: 2 barriers, init_count=32 mbarrier_init(smem + 231736, 32); + mbarrier_init(smem + 231744, 32); // index_valid_empty: 2 barriers, init_count=128 - mbarrier_init(smem + 231744, 128); mbarrier_init(smem + 231752, 128); + mbarrier_init(smem + 231760, 128); asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); } } @@ -1325,9 +1337,9 @@ kernel_cake_dsv4_ba9ed6a6233761682cc3(const __grid_constant__ CUtensorMap tmap_q asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); // TMEM alloc (512 columns, 512 used) - volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 231760); + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 231768); if (warp == 0) { - int _tmem_hold = smem + 231760; + int _tmem_hold = smem + 231768; asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); __syncwarp(); asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); @@ -1705,11 +1717,15 @@ kernel_cake_dsv4_ba9ed6a6233761682cc3(const __grid_constant__ CUtensorMap tmap_q asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum)); partial_lse[stat_offset] = ((final_sum > 0.0f) ? final_max * softmax_scale_log2_1 + _log2_0 : -CAKE_INF); } + float epi_scale = inv_sum * output_scale; + int epi_box_base = smem_q_addr + (unsigned int)(warp % 4 * 16384); + int epi_row0 = query_idx_1 * num_heads + cta_rank * 64 + head_row_base_1; #pragma unroll for (int acc_stage_1 = 0; acc_stage_1 < 2; acc_stage_1++) { #pragma unroll for (int local_slice_1 = 0; local_slice_1 < 2; local_slice_1++) { int logical_slice = acc_stage_1 * 2 * 2 + n_half_1 * 2 + local_slice_1; + int epi_box = epi_box_base + (acc_stage_1 * 2 + local_slice_1) * 4096; int o_base_1 = taddr + 128 + (unsigned int)(acc_stage_1 * 128) + (unsigned int)(local_slice_1 * 64) + (unsigned int)(output_row_base << 16); #pragma unroll for (int c_1 = 0; c_1 < 64; c_1 += 32) { @@ -1719,46 +1735,35 @@ kernel_cake_dsv4_ba9ed6a6233761682cc3(const __grid_constant__ CUtensorMap tmap_q " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]) : "r"(o_base_1 + c_1)); - int gmem_base = output_offset + logical_slice * 64 + c_1; + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + float epi_scaled[32]; #pragma unroll - for (int j = 0; j < 32; j += 16) { - { - const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; - #if __CUDA_ARCH__ >= 1000 - #pragma unroll - for (int _ps = 0; _ps < 8; _ps++) - mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_3[j])[_ps], _prescale2_1); - #else - #pragma unroll - for (int _ps = 0; _ps < 16; _ps++) - _tmem_load_3[j + _ps] *= inv_sum * output_scale; - #endif - { - __nv_bfloat162 _pk0 = __floats2bfloat162_rn(_tmem_load_3[j + 0], _tmem_load_3[j + 1]); - unsigned _pk_u0 = *reinterpret_cast(&_pk0); - __nv_bfloat162 _pk1 = __floats2bfloat162_rn(_tmem_load_3[j + 2], _tmem_load_3[j + 3]); - unsigned _pk_u1 = *reinterpret_cast(&_pk1); - __nv_bfloat162 _pk2 = __floats2bfloat162_rn(_tmem_load_3[j + 4], _tmem_load_3[j + 5]); - unsigned _pk_u2 = *reinterpret_cast(&_pk2); - __nv_bfloat162 _pk3 = __floats2bfloat162_rn(_tmem_load_3[j + 6], _tmem_load_3[j + 7]); - unsigned _pk_u3 = *reinterpret_cast(&_pk3); - __nv_bfloat162 _pk4 = __floats2bfloat162_rn(_tmem_load_3[j + 8], _tmem_load_3[j + 9]); - unsigned _pk_u4 = *reinterpret_cast(&_pk4); - __nv_bfloat162 _pk5 = __floats2bfloat162_rn(_tmem_load_3[j + 10], _tmem_load_3[j + 11]); - unsigned _pk_u5 = *reinterpret_cast(&_pk5); - __nv_bfloat162 _pk6 = __floats2bfloat162_rn(_tmem_load_3[j + 12], _tmem_load_3[j + 13]); - unsigned _pk_u6 = *reinterpret_cast(&_pk6); - __nv_bfloat162 _pk7 = __floats2bfloat162_rn(_tmem_load_3[j + 14], _tmem_load_3[j + 15]); - unsigned _pk_u7 = *reinterpret_cast(&_pk7); - asm volatile( - "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" - :: "l"((void*)(&((__nv_bfloat16*)(O + (gmem_base + j)))[0])), "r"(_pk_u0), "r"(_pk_u1), "r"(_pk_u2), "r"(_pk_u3), "r"(_pk_u4), "r"(_pk_u5), "r"(_pk_u6), "r"(_pk_u7) : "memory"); - } - } + for (int e = 0; e < 32; e++) { + epi_scaled[e] = _tmem_load_3[e] * epi_scale; + } + unsigned int epi_packed[16]; + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(epi_scaled[_lp*2 + 0], epi_scaled[_lp*2+1 + 0])); + epi_packed[_lp] = *(uint32_t*)&_bf2; } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((epi_box + (lane * 128 + c_1 * 2 ^ (lane * 128 + c_1 * 2 >> 7 & 7) << 4))), "r"(epi_packed[0]), "r"(epi_packed[1]), "r"(epi_packed[2]), "r"(epi_packed[3]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((epi_box + (lane * 128 + (c_1 * 2 + 16) ^ (lane * 128 + (c_1 * 2 + 16) >> 7 & 7) << 4))), "r"(epi_packed[4]), "r"(epi_packed[5]), "r"(epi_packed[6]), "r"(epi_packed[7]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((epi_box + (lane * 128 + (c_1 * 2 + 32) ^ (lane * 128 + (c_1 * 2 + 32) >> 7 & 7) << 4))), "r"(epi_packed[8]), "r"(epi_packed[9]), "r"(epi_packed[10]), "r"(epi_packed[11]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((epi_box + (lane * 128 + (c_1 * 2 + 48) ^ (lane * 128 + (c_1 * 2 + 48) >> 7 & 7) << 4))), "r"(epi_packed[12]), "r"(epi_packed[13]), "r"(epi_packed[14]), "r"(epi_packed[15]) : "memory"); + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + __syncwarp(); + if (elect_sync()) { + tma_store_4d((&O), 0, logical_slice, split_idx_1, epi_row0, epi_box); + asm volatile("cp.async.bulk.commit_group;"); } } } + if (elect_sync()) { + asm volatile("cp.async.bulk.wait_group.read 0;"); + mbarrier_arrive(epi_q_free_addr); + } asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); asm volatile( "{\n\t" @@ -1769,6 +1774,9 @@ kernel_cake_dsv4_ba9ed6a6233761682cc3(const __grid_constant__ CUtensorMap tmap_q :: "r"(o_empty_addr), "r"(0) : "memory"); tile_cursor_1 = tile_cursor_1 + num_kv_tiles_1; } + if (elect_sync()) { + asm volatile("cp.async.bulk.wait_group 0;"); + } asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); asm volatile("tcgen05.fence::before_thread_sync;"); } @@ -2118,6 +2126,7 @@ kernel_cake_dsv4_ba9ed6a6233761682cc3(const __grid_constant__ CUtensorMap tmap_q unsigned int k_stage_1 = 0; unsigned int v_stage_1 = 0; unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_epi_q_free_0 = 1; unsigned int _phase_k_empty = 1; unsigned int _phase_v_empty = 1; #pragma unroll 1 @@ -2131,6 +2140,8 @@ kernel_cake_dsv4_ba9ed6a6233761682cc3(const __grid_constant__ CUtensorMap tmap_q int first_kv_tile_1 = first_owner_tile_3; mbarrier_wait(q_empty_addr, _phase_q_empty_0); _phase_q_empty_0 ^= 1; + mbarrier_wait(epi_q_free_addr, _phase_epi_q_free_0); + _phase_epi_q_free_0 ^= 1; if (load_warp_rank == 0) { if (cta_rank == 0) { if (elect_sync()) { diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_fb6409904b35ed0f7a1a_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_9c3b9b4fa7cba6e44ffc_binding.cu similarity index 97% rename from csrc/cake_dsv4/sm_100a/cake_dsv4_fb6409904b35ed0f7a1a_binding.cu rename to csrc/cake_dsv4/sm_100a/cake_dsv4_9c3b9b4fa7cba6e44ffc_binding.cu index 45f5eb50419..e0f1513b40a 100644 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_fb6409904b35ed0f7a1a_binding.cu +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_9c3b9b4fa7cba6e44ffc_binding.cu @@ -15,7 +15,7 @@ */ // Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_fb6409904b35ed0f7a1a'. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_9c3b9b4fa7cba6e44ffc'. #include #include #include @@ -31,10 +31,10 @@ #include #include -extern "C" __global__ void kernel_cake_dsv4_fb6409904b35ed0f7a1a(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items); +extern "C" __global__ void kernel_cake_dsv4_9c3b9b4fa7cba6e44ffc(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items); -namespace cake_host_shim_c54faed3b4eb93e3 { +namespace cake_host_shim_42a0596262e73abf { using tvm::ffi::Optional; using tvm::ffi::TensorView; @@ -416,7 +416,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ } if (smem_status == cudaSuccess) { smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_fb6409904b35ed0f7a1a), + reinterpret_cast(kernel_cake_dsv4_9c3b9b4fa7cba6e44ffc), cudaFuncAttributeMaxDynamicSharedMemorySize, 190976); } @@ -426,7 +426,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ } } TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_fb6409904b35ed0f7a1a failed: " + << "cudaFuncSetAttribute for kernel_cake_dsv4_9c3b9b4fa7cba6e44ffc failed: " << cudaGetErrorString(smem_status); dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); dim3 block(512u, 1u, 1u); @@ -450,12 +450,12 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ config.attrs = attrs; config.numAttrs = n; cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_fb6409904b35ed0f7a1a), kargs); + &config, reinterpret_cast(kernel_cake_dsv4_9c3b9b4fa7cba6e44ffc), kargs); TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_fb6409904b35ed0f7a1a failed: " + << "cudaLaunchKernelExC for kernel_cake_dsv4_9c3b9b4fa7cba6e44ffc failed: " << cudaGetErrorString(launch_status); } -} // namespace cake_host_shim_c54faed3b4eb93e3 +} // namespace cake_host_shim_42a0596262e73abf -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_c54faed3b4eb93e3::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_42a0596262e73abf::Run); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_fb6409904b35ed0f7a1a_kernel.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_9c3b9b4fa7cba6e44ffc_kernel.cu similarity index 77% rename from csrc/cake_dsv4/sm_100a/cake_dsv4_fb6409904b35ed0f7a1a_kernel.cu rename to csrc/cake_dsv4/sm_100a/cake_dsv4_9c3b9b4fa7cba6e44ffc_kernel.cu index bcc67b72a4c..a623a6872fa 100644 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_fb6409904b35ed0f7a1a_kernel.cu +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_9c3b9b4fa7cba6e44ffc_kernel.cu @@ -62,11 +62,11 @@ static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap ali #define SMEM_SMEM_Q_FULL_STAGE_BYTES 32768 #define SMEM_SMEM_Q_FULL_STRIDE 32768 #define SMEM_SMEM_K_FULL_OFF 33792 -#define SMEM_SMEM_K_FULL_STAGE_BYTES 32768 -#define SMEM_SMEM_K_FULL_STRIDE 32768 -#define SMEM_SMEM_V_FULL_OFF 99328 -#define SMEM_SMEM_V_FULL_STAGE_BYTES 32768 -#define SMEM_SMEM_V_FULL_STRIDE 32768 +#define SMEM_SMEM_K_FULL_STAGE_BYTES 65536 +#define SMEM_SMEM_K_FULL_STRIDE 65536 +#define SMEM_SMEM_V_FULL_OFF 33792 +#define SMEM_SMEM_V_FULL_STAGE_BYTES 65536 +#define SMEM_SMEM_V_FULL_STRIDE 65536 #define SMEM_SMEM_STATS_MAX_OFF 164864 #define SMEM_SMEM_STATS_MAX_STAGE_BYTES 1024 #define SMEM_SMEM_STATS_MAX_STRIDE 1024 @@ -88,6 +88,18 @@ static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap ali #define SMEM_SMEM_P_FP8_OFF 174592 #define SMEM_SMEM_P_FP8_STAGE_BYTES 8192 #define SMEM_SMEM_P_FP8_STRIDE 8192 +#define SMEM_SMEM_MERGE_RECV_OFF 33792 +#define SMEM_SMEM_MERGE_RECV_STAGE_BYTES 65536 +#define SMEM_SMEM_MERGE_RECV_STRIDE 65536 +#define SMEM_SMEM_MERGE_SEND_OFF 99328 +#define SMEM_SMEM_MERGE_SEND_STAGE_BYTES 65536 +#define SMEM_SMEM_MERGE_SEND_STRIDE 65536 +#define SMEM_SMEM_MERGE_STATS_RECV_OFF 174592 +#define SMEM_SMEM_MERGE_STATS_RECV_STAGE_BYTES 512 +#define SMEM_SMEM_MERGE_STATS_RECV_STRIDE 512 +#define SMEM_SMEM_MERGE_STATS_SEND_OFF 182784 +#define SMEM_SMEM_MERGE_STATS_SEND_STAGE_BYTES 512 +#define SMEM_SMEM_MERGE_STATS_SEND_STRIDE 512 #define SMEM_TOTAL 190976 #define THREADS 512 @@ -308,17 +320,14 @@ __device__ __forceinline__ void mbarrier_wait_token_cluster_hint( } -__device__ __forceinline__ void tcgen05_mma_f8f6f4_cta2( +__device__ __forceinline__ void tcgen05_mma_f8f6f4( int taddr, uint64_t a_desc, uint64_t b_desc, uint32_t i_desc, int enable_input_d) { asm volatile( "{\n\t" ".reg .pred p;\n\t" - ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" "setp.ne.b32 p, %4, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" - "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "tcgen05.mma.cta_group::2.kind::f8f6f4 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" + "tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], %1, %2, %3, p;\n\t" "}\n" :: "r"(taddr), "l"(a_desc), "l"(b_desc), "r"(i_desc), "r"(enable_input_d) @@ -344,15 +353,15 @@ __device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t } -__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { +__device__ __forceinline__ void elect_commit(int mbar_addr) { asm volatile( "{\n\t" ".reg .pred leader;\n\t" "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" - ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" "}\n" - :: "r"(mbar_addr), "h"(cta_mask) : "memory"); + :: "r"(mbar_addr)); } @@ -1112,11 +1121,11 @@ __device__ __forceinline__ void tma_2d_gmem2smem( } -__device__ __forceinline__ void tma_4d_gmem2smem_cta2( +__device__ __forceinline__ void tma_4d_gmem2smem( int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { asm volatile( - "cp.async.bulk.tensor.4d.shared::cluster.global" - ".mbarrier::complete_tx::bytes.cta_group::2" + "cp.async.bulk.tensor.4d.shared::cta.global" + ".mbarrier::complete_tx::bytes" " [%0], [%1, {%2, %3, %4, %5}], [%6];" :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), "r"(mbar_addr) : "memory"); @@ -1141,62 +1150,11 @@ __device__ __forceinline__ void tma_gather4_gmem2smem( } -__device__ __forceinline__ void tma_gather4_gmem2smem_cta2( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form; see tma_gather4_gmem2smem above. - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem_mc( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr, unsigned short cta_mask) { - // Multicast variant: the PTX grammar ties the .shared::cluster - // destination to .multicast::cluster + ctaMask (cf. cuda_ptx / - // SM100_TMA_LOAD_MULTICAST_2D_GATHER4 in CUTLASS). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" - ".mbarrier::complete_tx::bytes.multicast::cluster" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr), "h"(cta_mask) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem_mc_cta2( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr, unsigned short cta_mask) { - // Multicast + cta_group::2 variant; see tma_gather4_gmem2smem_mc. - asm volatile( - "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" - ".mbarrier::complete_tx::bytes.multicast::cluster.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr), "h"(cta_mask) : "memory"); -} - - -__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { asm volatile( - "{\n\t" - ".reg .b16 lo, hi;\n\t" - "mov.b32 {lo, hi}, %1;\n\t" - "tcgen05.commit.cta_group::2.mbarrier::arrive::one" - ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" - "}\n" - :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); } @@ -1220,7 +1178,7 @@ __device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { extern "C" { __global__ __launch_bounds__(512, 1) __cluster_dims__(2,1,1) void -kernel_cake_dsv4_fb6409904b35ed0f7a1a(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items) +kernel_cake_dsv4_9c3b9b4fa7cba6e44ffc(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); @@ -1253,6 +1211,8 @@ kernel_cake_dsv4_fb6409904b35ed0f7a1a(const __grid_constant__ CUtensorMap tmap_q #define pv_pair_ready_addr (mbar_base + 288) #define tmem_dealloc_addr (mbar_base + 304) #define tmem_dealloc_peer_addr (mbar_base + 312) + #define merge_ready_addr (mbar_base + 320) + #define merge_full_addr (mbar_base + 328) const int bid = blockIdx.x; const int num_bids = gridDim.x; @@ -1274,8 +1234,8 @@ kernel_cake_dsv4_fb6409904b35ed0f7a1a(const __grid_constant__ CUtensorMap tmap_q const int smem_q_full_addr = smem + 1024; uint8_t* smem_k_full = reinterpret_cast(smem_raw + 33792); const int smem_k_full_addr = smem + 33792; - uint8_t* smem_v_full = reinterpret_cast(smem_raw + 99328); - const int smem_v_full_addr = smem + 99328; + uint8_t* smem_v_full = reinterpret_cast(smem_raw + 33792); + const int smem_v_full_addr = smem + 33792; float* smem_stats_max = reinterpret_cast(smem_raw + 164864); const int smem_stats_max_addr = smem + 164864; float* smem_stats_sum = reinterpret_cast(smem_raw + 165888); @@ -1290,25 +1250,33 @@ kernel_cake_dsv4_fb6409904b35ed0f7a1a(const __grid_constant__ CUtensorMap tmap_q const int smem_sparse_indices_addr = smem + 168448; uint8_t* smem_p_fp8 = reinterpret_cast(smem_raw + 174592); const int smem_p_fp8_addr = smem + 174592; + float* smem_merge_recv = reinterpret_cast(smem_raw + 33792); + const int smem_merge_recv_addr = smem + 33792; + float* smem_merge_send = reinterpret_cast(smem_raw + 99328); + const int smem_merge_send_addr = smem + 99328; + float* smem_merge_stats_recv = reinterpret_cast(smem_raw + 174592); + const int smem_merge_stats_recv_addr = smem + 174592; + float* smem_merge_stats_send = reinterpret_cast(smem_raw + 182784); + const int smem_merge_stats_send_addr = smem + 182784; { asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); } - // Mbarrier init (21 pipeline groups, 0 ordered-sequence groups, 40 barriers) - // Mbarriers at smem_raw[0..320) + // Mbarrier init (23 pipeline groups, 0 ordered-sequence groups, 42 barriers) + // Mbarriers at smem_raw[0..336) if (warp == 12) { uint32_t leader = elect_sync(); if (leader) { - // q_full: 1 barriers, init_count=2 - mbarrier_init(smem + 0, 2); + // q_full: 1 barriers, init_count=1 + mbarrier_init(smem + 0, 1); // q_empty: 1 barriers, init_count=1 mbarrier_init(smem + 8, 1); // --- pipeline 'k_pipe' --- - // k_full: 2 barriers, init_count=2 - mbarrier_init(smem + 16, 2); - mbarrier_init(smem + 24, 2); + // k_full: 2 barriers, init_count=1 + mbarrier_init(smem + 16, 1); + mbarrier_init(smem + 24, 1); // k_empty: 2 barriers, init_count=1 mbarrier_init(smem + 32, 1); mbarrier_init(smem + 40, 1); @@ -1395,6 +1363,10 @@ kernel_cake_dsv4_fb6409904b35ed0f7a1a(const __grid_constant__ CUtensorMap tmap_q mbarrier_init(smem + 296, 64); // tmem_dealloc: 1 barriers, init_count=288 mbarrier_init(smem + 304, 288); + // merge_ready: 1 barriers, init_count=64 + mbarrier_init(smem + 320, 64); + // merge_full: 1 barriers, init_count=1 + mbarrier_init(smem + 328, 1); asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); } } @@ -1419,16 +1391,13 @@ kernel_cake_dsv4_fb6409904b35ed0f7a1a(const __grid_constant__ CUtensorMap tmap_q __syncwarp(); - asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); - asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); - // TMEM alloc (512 columns, 512 used) - volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 320); + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 336); if (warp == 0) { - int _tmem_hold = smem + 320; - asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + int _tmem_hold = smem + 336; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); __syncwarp(); - asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); } __syncthreads(); @@ -1522,8 +1491,6 @@ kernel_cake_dsv4_fb6409904b35ed0f7a1a(const __grid_constant__ CUtensorMap tmap_q unsigned int _phase_q_empty_0 = 1; unsigned int _phase_index_full = 0; unsigned int _phase_k_empty = 1; - unsigned int _phase_k_empty_1 = 1; - unsigned int _phase_v_empty = 1; #pragma unroll for (unsigned int work_idx = 0; work_idx < 1; work_idx++) { int split_idx_1 = 0; @@ -1538,176 +1505,52 @@ kernel_cake_dsv4_fb6409904b35ed0f7a1a(const __grid_constant__ CUtensorMap tmap_q int first_tile_1 = split_idx_1 * tiles_per_split_1; int num_kv_tiles = tiles_per_split_1; int sparse_extent = num_kv_tiles * 128; + int wt_num_local_tiles = (num_kv_tiles - bid % 2 + 1) / 2; mbarrier_wait(q_empty_addr, _phase_q_empty_0); _phase_q_empty_0 ^= 1; if (load_warp_rank == 0) { if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + mbarrier_arrive_expect_tx(q_full_addr, 32768); #pragma unroll for (int q_stage = 0; q_stage < 4; q_stage++) { - asm volatile( - "cp.async.bulk.tensor.4d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5}], [%6], %7;" - :: "r"(smem_q_full_addr + (unsigned int)(q_stage * 8192)), "l"((&tmap_q)), "r"(0), "r"(bid % 2 * 64), "r"(q_stage), "r"(query_idx_1), - "r"(((q_full_addr) & 0xFEFFFFFF)), "h"((uint16_t)(1 << bid % 2)) : "memory"); + tma_4d_gmem2smem(smem_q_full_addr + (unsigned int)(q_stage * 8192), (&tmap_q), 0, 0, q_stage, query_idx_1, q_full_addr); } } } - { + #pragma unroll 1 + for (int wt_tile = 0; wt_tile < wt_num_local_tiles; wt_tile++) { + int wt_global_tile = first_tile_1 + wt_tile * 2 + bid % 2; mbarrier_wait(index_full_addr + (load_k_index_stage) * 8, _phase_index_full); asm volatile("tcgen05.fence::after_thread_sync;"); - } - int k_index_stage_base = smem_sparse_indices_addr + load_k_index_stage * 1024; - { + int wt_index_base = smem_sparse_indices_addr + load_k_index_stage * 1024; mbarrier_wait(k_empty_addr + (load_k_stage) * 8, _phase_k_empty); if (load_warp_rank == 0) { if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + mbarrier_arrive_expect_tx(k_full_addr + (load_k_stage) * 8, 65536); } } - int k_dst = smem_k_full_addr + load_k_stage * 32768 + (unsigned int)(load_warp_rank * 8192); - int group = lane; - if (group < 16) { - int group_offset = k_cta_offset + group * 4; - int raw_rows[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows[0])), "=r"(*reinterpret_cast(&raw_rows[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows[(0) + 3])) - : "r"(k_index_stage_base + group_offset * 4)); - int row0 = ((raw_rows[0] >= 0) ? raw_rows[0] : 0); - int row1 = ((raw_rows[1] >= 0) ? raw_rows[1] : 0); - int row2 = ((raw_rows[2] >= 0) ? raw_rows[2] : 0); - int row3 = ((raw_rows[3] >= 0) ? raw_rows[3] : 0); - if (first_tile_1 == 0) { - tma_gather4_gmem2smem_mc_cta2(k_dst + group * 512, (&tmap_swa_kv), load_warp_rank * 128, row0, row1, row2, row3, ((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } else { - tma_gather4_gmem2smem_mc_cta2(k_dst + group * 512, (&tmap_compressed_kv), load_warp_rank * 128, row0, row1, row2, row3, ((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } + int wt_dst = smem_k_full_addr + load_k_stage * 65536 + (unsigned int)(load_warp_rank * 16384); + int wt_group = lane; + int wt_group_offset = bid % 2 * 128 + wt_group * 4; + int wt_rows[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&wt_rows[0])), "=r"(*reinterpret_cast(&wt_rows[(0) + 1])), "=r"(*reinterpret_cast(&wt_rows[(0) + 2])), "=r"(*reinterpret_cast(&wt_rows[(0) + 3])) + : "r"(wt_index_base + wt_group_offset * 4)); + int wt_row0 = ((wt_rows[0] >= 0) ? wt_rows[0] : 0); + int wt_row1 = ((wt_rows[1] >= 0) ? wt_rows[1] : 0); + int wt_row2 = ((wt_rows[2] >= 0) ? wt_rows[2] : 0); + int wt_row3 = ((wt_rows[3] >= 0) ? wt_rows[3] : 0); + if (wt_global_tile == 0) { + tma_gather4_gmem2smem(wt_dst + wt_group * 512, (&tmap_swa_kv), load_warp_rank * 128, wt_row0, wt_row1, wt_row2, wt_row3, k_full_addr + (load_k_stage) * 8); + } else { + tma_gather4_gmem2smem(wt_dst + wt_group * 512, (&tmap_compressed_kv), load_warp_rank * 128, wt_row0, wt_row1, wt_row2, wt_row3, k_full_addr + (load_k_stage) * 8); } load_k_stage += 1; if (load_k_stage == 2) { load_k_stage = 0; _phase_k_empty ^= 1; } - } - #pragma unroll 1 - for (int tile = 1; tile < num_kv_tiles; tile++) { - if ((1 & (int)((tile & 1) == 0)) != 0) { - mbarrier_wait(index_full_addr + (load_k_index_stage) * 8, _phase_index_full); - asm volatile("tcgen05.fence::after_thread_sync;"); - } - k_index_stage_base = smem_sparse_indices_addr + load_k_index_stage * 1024; - { - mbarrier_wait(k_empty_addr + (load_k_stage) * 8, _phase_k_empty); - if (load_warp_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - int k_dst_1 = smem_k_full_addr + load_k_stage * 32768 + (unsigned int)(load_warp_rank * 8192); - int group_1 = lane; - if (group_1 < 16) { - int group_offset_1 = (tile & 1) * 128 + k_cta_offset + group_1 * 4; - int raw_rows_1[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_1[0])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 3])) - : "r"(k_index_stage_base + group_offset_1 * 4)); - int row0_1 = ((raw_rows_1[0] >= 0) ? raw_rows_1[0] : 0); - int row1_1 = ((raw_rows_1[1] >= 0) ? raw_rows_1[1] : 0); - int row2_1 = ((raw_rows_1[2] >= 0) ? raw_rows_1[2] : 0); - int row3_1 = ((raw_rows_1[3] >= 0) ? raw_rows_1[3] : 0); - tma_gather4_gmem2smem_mc_cta2(k_dst_1 + group_1 * 512, (&tmap_compressed_kv), load_warp_rank * 128, row0_1, row1_1, row2_1, row3_1, ((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } - load_k_stage += 1; - if (load_k_stage == 2) { load_k_stage = 0; _phase_k_empty ^= 1; } - if ((tile & 1) != 0) { - load_k_index_stage += 1; - if (load_k_index_stage == 6) { load_k_index_stage = 0; _phase_index_full ^= 1; } - } - } - int prev_tile = tile - 1; - int global_prev_tile = first_tile_1 + prev_tile; - int v_index_stage_base = smem_sparse_indices_addr + load_v_index_stage * 1024; - { - mbarrier_wait(v_empty_addr + (load_v_stage) * 8, _phase_v_empty); - if (load_warp_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - if (lane < 16) { - int v_dst = smem_v_full_addr + load_v_stage * 32768 + (unsigned int)(load_warp_rank / 2 * 16384); - int v_col = v_chunk * 256 + load_warp_rank / 2 * 128; - int group_2 = (load_warp_rank & 1) * 16 + lane; - int group_offset_2 = (prev_tile & 1) * 128 + group_2 * 4; - int raw_rows_2[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_2[0])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 3])) - : "r"(v_index_stage_base + group_offset_2 * 4)); - int row0_2 = ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : 0); - int row1_2 = ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : 0); - int row2_2 = ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : 0); - int row3_2 = ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : 0); - if (global_prev_tile == 0) { - tma_gather4_gmem2smem_mc_cta2(v_dst + group_2 * 512, (&tmap_swa_kv), v_col, row0_2, row1_2, row2_2, row3_2, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } else { - tma_gather4_gmem2smem_mc_cta2(v_dst + group_2 * 512, (&tmap_compressed_kv), v_col, row0_2, row1_2, row2_2, row3_2, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } - } - load_v_stage += 1; - if (load_v_stage == 2) { load_v_stage = 0; _phase_v_empty ^= 1; } - if ((prev_tile & 1) != 0) { - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(index_empty_addr + (load_v_index_stage) * 8); - load_v_index_stage = load_v_index_stage + 1; - if (load_v_index_stage == 6) { - load_v_index_stage = 0; - } - } - } - } - int last_tile = num_kv_tiles - 1; - int global_last_tile = first_tile_1 + last_tile; - { - mbarrier_wait(v_empty_addr + (load_v_stage) * 8, _phase_v_empty); - if (load_warp_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - if (lane < 16) { - int v_dst_1 = smem_v_full_addr + load_v_stage * 32768 + (unsigned int)(load_warp_rank / 2 * 16384); - int v_col_1 = v_chunk * 256 + load_warp_rank / 2 * 128; - int group_3 = (load_warp_rank & 1) * 16 + lane; - int group_offset_3 = (last_tile & 1) * 128 + group_3 * 4; - int raw_rows_3[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_3[0])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 3])) - : "r"(smem_sparse_indices_addr + load_v_index_stage * 1024 + (unsigned int)(group_offset_3 * 4))); - int row0_3 = ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : 0); - int row1_3 = ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : 0); - int row2_3 = ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : 0); - int row3_3 = ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : 0); - if (global_last_tile == 0) { - tma_gather4_gmem2smem_mc_cta2(v_dst_1 + group_3 * 512, (&tmap_swa_kv), v_col_1, row0_3, row1_3, row2_3, row3_3, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } else { - tma_gather4_gmem2smem_mc_cta2(v_dst_1 + group_3 * 512, (&tmap_compressed_kv), v_col_1, row0_3, row1_3, row2_3, row3_3, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } - } - load_v_stage += 1; - if (load_v_stage == 2) { load_v_stage = 0; _phase_v_empty ^= 1; } asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(index_empty_addr + (load_v_index_stage) * 8); - load_v_index_stage = load_v_index_stage + 1; - if (load_v_index_stage == 6) { - load_v_index_stage = 0; - } + mbarrier_arrive(index_empty_addr + (load_k_index_stage) * 8); + load_k_index_stage += 1; + if (load_k_index_stage == 6) { load_k_index_stage = 0; _phase_index_full ^= 1; } } } mbarrier_arrive(tmem_dealloc_addr); @@ -1743,14 +1586,14 @@ kernel_cake_dsv4_fb6409904b35ed0f7a1a(const __grid_constant__ CUtensorMap tmap_q int active_topk = _min_0; float row_max_val = -CAKE_INF; float row_sum_val = 0.0f; - int sink_head = ((1) ? bid % 2 * 64 + my_row : my_row); - if (has_sinks != 0 && sink_head < num_heads && split_idx_2 == 0) { + int sink_head = ((1) ? my_row : my_row); + if (has_sinks != 0 && sink_head < num_heads && split_idx_2 + bid % 2 == 0) { row_max_val = sinks[sink_head] * 1.4426950408889634f / softmax_scale_log2; row_sum_val = 1.0f; } #pragma unroll 1 - for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { - int pipeline_tile = softmax_tile_cursor + tile_1; + for (int tile = 0; tile < (num_kv_tiles_1 - bid % 2 + 1) / 2; tile++) { + int pipeline_tile = softmax_tile_cursor + tile; int phase = pipeline_tile & 1; int s_wait_phase = pipeline_tile >> 1 & 1; mbarrier_wait(s_full_addr + (phase) * 8, s_wait_phase); @@ -1759,14 +1602,14 @@ kernel_cake_dsv4_fb6409904b35ed0f7a1a(const __grid_constant__ CUtensorMap tmap_q int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_score_row_base << 16); float new_max = row_max_val; int valid_sparse_cols = ((active_topk < sparse_topk) ? active_topk : sparse_topk); - valid_sparse_cols = valid_sparse_cols - (first_tile_2 + tile_1) * 128 - n_half * 64; + valid_sparse_cols = valid_sparse_cols - (first_tile_2 + tile * 2 + bid % 2) * 128 - n_half * 64; if (valid_sparse_cols < 0) { valid_sparse_cols = 0; } if (valid_sparse_cols > 64) { valid_sparse_cols = 64; } - int index_tile_addr = smem_sparse_indices_addr + softmax_index_stage * 1024 + (unsigned int)((tile_1 & 1) * 512); + int index_tile_addr = smem_sparse_indices_addr + softmax_index_stage * 1024 + (unsigned int)(bid % 2 * 512); int index_lane_addr = index_tile_addr + (n_half * 64 + lane) * 4; int staged_index[2]; unsigned int invalid_cols[2]; @@ -1778,7 +1621,7 @@ kernel_cake_dsv4_fb6409904b35ed0f7a1a(const __grid_constant__ CUtensorMap tmap_q invalid_cols[w] = _vote_0; any_invalid = any_invalid | invalid_cols[w]; } - if ((tile_1 & 1) != 0 || tile_1 + 1 == num_kv_tiles_1) { + if (tile >= 0) { asm volatile("tcgen05.fence::before_thread_sync;"); mbarrier_arrive(index_empty_addr + (softmax_index_stage) * 8); softmax_index_stage = softmax_index_stage + 1; @@ -2216,7 +2059,7 @@ kernel_cake_dsv4_fb6409904b35ed0f7a1a(const __grid_constant__ CUtensorMap tmap_q smem_stats_final_max[stats_row] = row_max_val; asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); mbarrier_arrive(sum_ready_addr); - softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles_1; + softmax_tile_cursor = softmax_tile_cursor + (num_kv_tiles_1 - bid % 2 + 1) / 2; } asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); asm volatile("tcgen05.fence::before_thread_sync;"); @@ -2246,6 +2089,8 @@ kernel_cake_dsv4_fb6409904b35ed0f7a1a(const __grid_constant__ CUtensorMap tmap_q unsigned int _phase_pv_done_0 = 0; unsigned int _phase_o_done_0 = 0; unsigned int _phase_sum_ready_0 = 0; + unsigned int _phase_merge_ready_0 = 0; + unsigned int _phase_merge_full_0 = 0; #pragma unroll for (unsigned int work_idx_2 = 0; work_idx_2 < 1; work_idx_2++) { int split_idx_3 = 0; @@ -2260,19 +2105,19 @@ kernel_cake_dsv4_fb6409904b35ed0f7a1a(const __grid_constant__ CUtensorMap tmap_q int first_tile_3 = split_idx_3 * tiles_per_split_3; int num_kv_tiles_2 = tiles_per_split_3; #pragma unroll 1 - for (int tile_2 = 0; tile_2 < num_kv_tiles_2; tile_2++) { - int pipeline_tile_1 = correction_tile_cursor + tile_2; + for (int tile_1 = 0; tile_1 < (num_kv_tiles_2 - bid % 2 + 1) / 2; tile_1++) { + int pipeline_tile_1 = correction_tile_cursor + tile_1; int phase_1 = pipeline_tile_1 & 1; int stats_wait_phase = pipeline_tile_1 >> 1 & 1; mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); asm volatile("tcgen05.fence::after_thread_sync;"); float acc_scale_1 = smem_stats_max[phase_1 * 128 + stats_row_1]; - if (tile_2 > 0) { + if (tile_1 > 0) { mbarrier_wait(pv_done_addr, _phase_pv_done_0); _phase_pv_done_0 ^= 1; asm volatile("tcgen05.fence::after_thread_sync;"); } - if (tile_2 > 0) { + if (tile_1 > 0) { int _vote_1 = __any_sync(0xFFFFFFFF, acc_scale_1 < 1.0f); int any_rescale = _vote_1; if (any_rescale != 0) { @@ -2301,60 +2146,157 @@ kernel_cake_dsv4_fb6409904b35ed0f7a1a(const __grid_constant__ CUtensorMap tmap_q mbarrier_wait(sum_ready_addr, _phase_sum_ready_0); _phase_sum_ready_0 ^= 1; asm volatile("tcgen05.fence::after_thread_sync;"); - float total_sum = smem_stats_sum[stats_row_1]; - float final_max = smem_stats_final_max[stats_row_1]; - { - smem_corr_warp_pair_exchange[stats_row_1] = total_sum; - asm volatile("barrier.sync %0, 64;" :: "r"(4 + warp % 2) : "memory"); - total_sum = total_sum + smem_corr_warp_pair_exchange[stats_row_1 ^ 64]; + float wt_sum = smem_stats_sum[stats_row_1]; + float wt_max = smem_stats_final_max[stats_row_1]; + smem_corr_warp_pair_exchange[stats_row_1] = wt_sum; + asm volatile("barrier.sync %0, 64;" :: "r"(4 + warp % 2) : "memory"); + wt_sum = wt_sum + smem_corr_warp_pair_exchange[stats_row_1 ^ 64]; + int wt_peer_rank = bid % 2 ^ 1; + int wt_quadrant = warp - 4; + if (warp == 4) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(merge_full_addr, 66048); + } } - float _rcp_0 = approx_rcp(total_sum); - float inv_sum = ((total_sum > 0.0f) ? _rcp_0 : 0.0f); - int head_idx = my_row_1; - { - head_idx = bid % 2 * 64 + my_row_1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(merge_ready_addr), "r"(wt_peer_rank) : "memory"); + if (n_half_1 == 0) { + smem_merge_stats_send[my_row_1] = wt_max; + smem_merge_stats_send[64 + my_row_1] = wt_sum; } - int direct_o_offset = (query_idx_3 * num_heads + head_idx) * 512 + v_chunk_1 * 256; - int partial_o_offset = (query_idx_3 * num_heads + head_idx + split_idx_3) * 512 + ((0) ? v_chunk_1 : n_half_1) * 256; - int o_offset = partial_o_offset; + int wt_send_block0 = (bid % 2 ^ 1) * 2; #pragma unroll - for (int vs_1 = 0; vs_1 < 2; vs_1++) { - int o_base_epi = taddr + 256 + (unsigned int)(vs_1 * 128) + (unsigned int)corr_row; + for (int wt_b = 0; wt_b < 2; wt_b++) { #pragma unroll - for (int c_1 = 0; c_1 < 128; c_1 += 32) { + for (int wt_c = 0; wt_c < 64; wt_c += 32) { float _tmem_load_5[32]; asm volatile( "tcgen05.ld.sync.aligned.32x32b.x32.b32" " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" : "=f"(_tmem_load_5[0]), "=f"(_tmem_load_5[1]), "=f"(_tmem_load_5[2]), "=f"(_tmem_load_5[3]), "=f"(_tmem_load_5[4]), "=f"(_tmem_load_5[5]), "=f"(_tmem_load_5[6]), "=f"(_tmem_load_5[7]), "=f"(_tmem_load_5[8]), "=f"(_tmem_load_5[9]), "=f"(_tmem_load_5[10]), "=f"(_tmem_load_5[11]), "=f"(_tmem_load_5[12]), "=f"(_tmem_load_5[13]), "=f"(_tmem_load_5[14]), "=f"(_tmem_load_5[15]), "=f"(_tmem_load_5[16]), "=f"(_tmem_load_5[17]), "=f"(_tmem_load_5[18]), "=f"(_tmem_load_5[19]), "=f"(_tmem_load_5[20]), "=f"(_tmem_load_5[21]), "=f"(_tmem_load_5[22]), "=f"(_tmem_load_5[23]), "=f"(_tmem_load_5[24]), "=f"(_tmem_load_5[25]), "=f"(_tmem_load_5[26]), "=f"(_tmem_load_5[27]), "=f"(_tmem_load_5[28]), "=f"(_tmem_load_5[29]), "=f"(_tmem_load_5[30]), "=f"(_tmem_load_5[31]) - : "r"(o_base_epi + c_1)); - int gmem_base = o_offset + vs_1 * 128 + c_1; - if (head_idx < num_heads) { + : "r"(taddr + 256 + (unsigned int)((wt_send_block0 + wt_b) * 64) + (unsigned int)wt_c + (unsigned int)corr_row)); + #pragma unroll + for (int wt_k = 0; wt_k < 32; wt_k += 4) { + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(smem_merge_send_addr + (unsigned int)((((wt_b * 64 + wt_c + wt_k) / 4 * 4 + wt_quadrant) * 32 + lane) * 16)), "f"(_tmem_load_5[wt_k]), "f"(_tmem_load_5[wt_k + 1]), "f"(_tmem_load_5[wt_k + 2]), "f"(_tmem_load_5[wt_k + 3]) : "memory"); + } + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 6, 64;" ::: "memory"); + if (warp == 4) { + if (elect_sync()) { + mbarrier_wait(merge_ready_addr, _phase_merge_ready_0); + _phase_merge_ready_0 ^= 1; + uint32_t _mapa_0; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_0) : "r"(smem_merge_recv_addr), "r"(wt_peer_rank)); + uint32_t _mapa_1; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_1) : "r"(merge_full_addr), "r"(wt_peer_rank)); + asm volatile( + "cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes" + " [%0], [%1], %2, [%3];" + :: "r"(_mapa_0), "r"(smem_merge_send_addr), "r"((uint32_t)(65536)), "r"(_mapa_1) + : "memory"); + uint32_t _mapa_2; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_2) : "r"(smem_merge_stats_recv_addr), "r"(wt_peer_rank)); + uint32_t _mapa_3; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_3) : "r"(merge_full_addr), "r"(wt_peer_rank)); + asm volatile( + "cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes" + " [%0], [%1], %2, [%3];" + :: "r"(_mapa_2), "r"(smem_merge_stats_send_addr), "r"((uint32_t)(512)), "r"(_mapa_3) + : "memory"); + } + } + mbarrier_wait(merge_full_addr, _phase_merge_full_0); + _phase_merge_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + float wt_peer_max = smem_merge_stats_recv[my_row_1]; + float wt_peer_sum = smem_merge_stats_recv[64 + my_row_1]; + float _max_5 = max_noftz(wt_max, wt_peer_max); + float wt_merged_max = _max_5; + float _exp2_1 = approx_exp2((wt_max - wt_merged_max) * softmax_scale_log2_1); + float wt_s_own = ((wt_max > -CAKE_INF) ? _exp2_1 : 0.0f); + float _exp2_2 = approx_exp2((wt_peer_max - wt_merged_max) * softmax_scale_log2_1); + float wt_s_peer = ((wt_peer_max > -CAKE_INF) ? _exp2_2 : 0.0f); + float wt_merged_sum = wt_sum * wt_s_own + wt_peer_sum * wt_s_peer; + float _rcp_0 = approx_rcp(wt_merged_sum); + float wt_inv_sum = ((wt_merged_sum > 0.0f) ? _rcp_0 : 0.0f); + int wt_head_idx = my_row_1; + int wt_o_base = (query_idx_3 * num_heads + wt_head_idx) * 512; + if (bid % 2 == 0 && n_half_1 == 0 && wt_head_idx < num_heads) { + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(wt_merged_sum)); + partial_lse[query_idx_3 * num_heads + wt_head_idx + split_idx_3] = ((wt_merged_sum > 0.0f) ? wt_merged_max * softmax_scale_log2_1 + _log2_0 : -CAKE_INF); + } + int wt_keep_block0 = bid % 2 * 2; + #pragma unroll + for (int wt_b_1 = 0; wt_b_1 < 2; wt_b_1++) { + #pragma unroll + for (int wt_c_1 = 0; wt_c_1 < 64; wt_c_1 += 32) { + float _tmem_load_6[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_6[0]), "=f"(_tmem_load_6[1]), "=f"(_tmem_load_6[2]), "=f"(_tmem_load_6[3]), "=f"(_tmem_load_6[4]), "=f"(_tmem_load_6[5]), "=f"(_tmem_load_6[6]), "=f"(_tmem_load_6[7]), "=f"(_tmem_load_6[8]), "=f"(_tmem_load_6[9]), "=f"(_tmem_load_6[10]), "=f"(_tmem_load_6[11]), "=f"(_tmem_load_6[12]), "=f"(_tmem_load_6[13]), "=f"(_tmem_load_6[14]), "=f"(_tmem_load_6[15]), "=f"(_tmem_load_6[16]), "=f"(_tmem_load_6[17]), "=f"(_tmem_load_6[18]), "=f"(_tmem_load_6[19]), "=f"(_tmem_load_6[20]), "=f"(_tmem_load_6[21]), "=f"(_tmem_load_6[22]), "=f"(_tmem_load_6[23]), "=f"(_tmem_load_6[24]), "=f"(_tmem_load_6[25]), "=f"(_tmem_load_6[26]), "=f"(_tmem_load_6[27]), "=f"(_tmem_load_6[28]), "=f"(_tmem_load_6[29]), "=f"(_tmem_load_6[30]), "=f"(_tmem_load_6[31]) + : "r"(taddr + 256 + (unsigned int)((wt_keep_block0 + wt_b_1) * 64) + (unsigned int)wt_c_1 + (unsigned int)corr_row)); + float wt_peer[32]; + #pragma unroll + for (int wt_k_1 = 0; wt_k_1 < 32; wt_k_1 += 4) { + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&wt_peer[wt_k_1])), "=r"(*reinterpret_cast(&wt_peer[(wt_k_1) + 1])), "=r"(*reinterpret_cast(&wt_peer[(wt_k_1) + 2])), "=r"(*reinterpret_cast(&wt_peer[(wt_k_1) + 3])) + : "r"(smem_merge_recv_addr + (unsigned int)((((wt_b_1 * 64 + wt_c_1 + wt_k_1) / 4 * 4 + wt_quadrant) * 32 + lane) * 16))); + } + const float2 _scale2_1 = {wt_s_own, wt_s_own}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + mul_f32x2_inplace(&reinterpret_cast(_tmem_load_6)[_ls], _scale2_1); + const float2 _scale2_2 = {wt_s_peer, wt_s_peer}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + mul_f32x2_inplace(&reinterpret_cast(wt_peer)[_ls], _scale2_2); + #pragma unroll + for (int _la = 0; _la < 16; _la++) + add_f32x2_inplace(&reinterpret_cast(_tmem_load_6)[_la], reinterpret_cast(wt_peer)[_la]); + int wt_gmem = wt_o_base + (wt_keep_block0 + wt_b_1) * 128 + n_half_1 * 64 + wt_c_1; + if (wt_head_idx < num_heads) { #pragma unroll - for (int j = 0; j < 32; j += 8) { + for (int wt_j = 0; wt_j < 32; wt_j += 8) { { - const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; + const float2 _prescale2_3 = {wt_inv_sum * output_scale, wt_inv_sum * output_scale}; #if __CUDA_ARCH__ >= 1000 #pragma unroll for (int _ps = 0; _ps < 4; _ps++) - mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_5[j])[_ps], _prescale2_1); + mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_6[wt_j])[_ps], _prescale2_3); #else #pragma unroll for (int _ps = 0; _ps < 8; _ps++) - _tmem_load_5[j + _ps] *= inv_sum * output_scale; + _tmem_load_6[wt_j + _ps] *= wt_inv_sum * output_scale; #endif __nv_bfloat162 _pk[4]; - _pk[0] = __floats2bfloat162_rn(_tmem_load_5[j + 0], _tmem_load_5[j + 1]); - _pk[1] = __floats2bfloat162_rn(_tmem_load_5[j + 2], _tmem_load_5[j + 3]); - _pk[2] = __floats2bfloat162_rn(_tmem_load_5[j + 4], _tmem_load_5[j + 5]); - _pk[3] = __floats2bfloat162_rn(_tmem_load_5[j + 6], _tmem_load_5[j + 7]); - *reinterpret_cast(&((__nv_bfloat16*)(O + (gmem_base + j)))[0]) = *reinterpret_cast(&_pk[0]); + _pk[0] = __floats2bfloat162_rn(_tmem_load_6[wt_j + 0], _tmem_load_6[wt_j + 1]); + _pk[1] = __floats2bfloat162_rn(_tmem_load_6[wt_j + 2], _tmem_load_6[wt_j + 3]); + _pk[2] = __floats2bfloat162_rn(_tmem_load_6[wt_j + 4], _tmem_load_6[wt_j + 5]); + _pk[3] = __floats2bfloat162_rn(_tmem_load_6[wt_j + 6], _tmem_load_6[wt_j + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(O + (wt_gmem + wt_j)))[0]) = *reinterpret_cast(&_pk[0]); } } } } } - correction_tile_cursor = correction_tile_cursor + num_kv_tiles_2; + correction_tile_cursor = correction_tile_cursor + (num_kv_tiles_2 - bid % 2 + 1) / 2; } asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); asm volatile("tcgen05.fence::before_thread_sync;"); @@ -2376,13 +2318,8 @@ kernel_cake_dsv4_fb6409904b35ed0f7a1a(const __grid_constant__ CUtensorMap tmap_q unsigned int mma_v_stage = 0; unsigned int mma_kv_stage = 0; int mma_tile_cursor = 0; - unsigned int _phase_s_seeded_0 = 0; unsigned int _phase_q_full_0_3 = 0; - unsigned int _phase_q_pair_ready_0 = 0; unsigned int _phase_k_full = 0; - unsigned int _phase_k_full_1 = 0; - unsigned int _phase_kv_pair_ready = 0; - unsigned int _phase_v_full = 0; #pragma unroll for (unsigned int work_idx_3 = 0; work_idx_3 < 1; work_idx_3++) { int split_idx_4 = 0; @@ -2392,410 +2329,487 @@ kernel_cake_dsv4_fb6409904b35ed0f7a1a(const __grid_constant__ CUtensorMap tmap_q int tiles_per_split_4 = all_num_kv_tiles_4 + 1 - 1; int first_tile_4 = split_idx_4 * tiles_per_split_4; int num_kv_tiles_3 = tiles_per_split_4; - { - if (cta_rank == 0) { - mbarrier_wait(q_full_addr, _phase_q_full_0_3); - _phase_q_full_0_3 ^= 1; - } - } - int first_pv = 1; + int wt_num_local_tiles_1 = (num_kv_tiles_3 - bid % 2 + 1) / 2; + mbarrier_wait(q_full_addr, _phase_q_full_0_3); + _phase_q_full_0_3 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int wt_first_pv = 1; #pragma unroll 1 - for (int tile_3 = 0; tile_3 < num_kv_tiles_3; tile_3++) { - int pipeline_tile_2 = mma_tile_cursor + tile_3; - int phase_2 = pipeline_tile_2 & 1; - int score_col = ((phase_2 != 0) ? 128 : 0); - { - if (cta_rank == 0) { - mbarrier_wait(k_full_addr + (mma_k_stage) * 8, _phase_k_full); - int _mma_a_lo_0 = (((smem_q_full_addr) >> 4) & 0x3FFF) + (0) * 2048; - int _mma_b_lo_0 = (((smem_k_full_addr) >> 4) & 0x3FFF) + (mma_k_stage) * 2048; - asm volatile( + for (int wt_tile_1 = 0; wt_tile_1 < wt_num_local_tiles_1; wt_tile_1++) { + int wt_pipeline_tile = mma_tile_cursor + wt_tile_1; + int wt_phase = wt_pipeline_tile & 1; + int wt_score_col = ((wt_phase != 0) ? 128 : 0); + mbarrier_wait(k_full_addr + (mma_k_stage) * 8, _phase_k_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_0 = make_warp_uniform((((smem_q_full_addr) >> 4) & 0x3FFF) + (0) * 2048); + int _mma_b_lo_0 = make_warp_uniform((((smem_k_full_addr) >> 4) & 0x3FFF) + (mma_k_stage) * 4096); + asm volatile( "{\n\t" ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" ".reg .b64 da, db;\n\t" "elect.sync _|leader, 0xFFFFFFFF;\n\t" "setp.ne.b32 p0, %3, 0;\n\t" "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "" "mov.b32 adhi, 0x40004040;\n\t" "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 136314896;\n\t" + "mov.b32 id, 69206032;\n\t" "mov.b32 alo, %0;\n\t" "mov.b32 blo, %1;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 2;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 2;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 2;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "}\n" - :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(0)); - int _mma_a_lo_1 = (((smem_q_full_addr + 8192) >> 4) & 0x3FFF) + (0) * 2048; - int _mma_b_lo_1 = (((smem_k_full_addr + 8192) >> 4) & 0x3FFF) + (mma_k_stage) * 2048; - asm volatile( + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (wt_score_col))), "r"(0)); + int _mma_a_lo_1 = make_warp_uniform((((smem_q_full_addr + 8192) >> 4) & 0x3FFF) + (0) * 2048); + int _mma_b_lo_1 = make_warp_uniform((((smem_k_full_addr + 16384) >> 4) & 0x3FFF) + (mma_k_stage) * 4096); + asm volatile( "{\n\t" ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" ".reg .b64 da, db;\n\t" "elect.sync _|leader, 0xFFFFFFFF;\n\t" "setp.ne.b32 p0, %3, 0;\n\t" "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "" "mov.b32 adhi, 0x40004040;\n\t" "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 136314896;\n\t" + "mov.b32 id, 69206032;\n\t" "mov.b32 alo, %0;\n\t" "mov.b32 blo, %1;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 2;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 2;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 2;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "}\n" - :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); - int _mma_a_lo_2 = (((smem_q_full_addr + 16384) >> 4) & 0x3FFF) + (0) * 2048; - int _mma_b_lo_2 = (((smem_k_full_addr + 16384) >> 4) & 0x3FFF) + (mma_k_stage) * 2048; - asm volatile( + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_tmem_scratch + (wt_score_col))), "r"(1)); + int _mma_a_lo_2 = make_warp_uniform((((smem_q_full_addr + 16384) >> 4) & 0x3FFF) + (0) * 2048); + int _mma_b_lo_2 = make_warp_uniform((((smem_k_full_addr + 32768) >> 4) & 0x3FFF) + (mma_k_stage) * 4096); + asm volatile( "{\n\t" ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" ".reg .b64 da, db;\n\t" "elect.sync _|leader, 0xFFFFFFFF;\n\t" "setp.ne.b32 p0, %3, 0;\n\t" "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "" "mov.b32 adhi, 0x40004040;\n\t" "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 136314896;\n\t" + "mov.b32 id, 69206032;\n\t" "mov.b32 alo, %0;\n\t" "mov.b32 blo, %1;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 2;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 2;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 2;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "}\n" - :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); - int _mma_a_lo_3 = (((smem_q_full_addr + 24576) >> 4) & 0x3FFF) + (0) * 2048; - int _mma_b_lo_3 = (((smem_k_full_addr + 24576) >> 4) & 0x3FFF) + (mma_k_stage) * 2048; - asm volatile( + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (wt_score_col))), "r"(1)); + int _mma_a_lo_3 = make_warp_uniform((((smem_q_full_addr + 24576) >> 4) & 0x3FFF) + (0) * 2048); + int _mma_b_lo_3 = make_warp_uniform((((smem_k_full_addr + 49152) >> 4) & 0x3FFF) + (mma_k_stage) * 4096); + asm volatile( "{\n\t" ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" ".reg .b64 da, db;\n\t" "elect.sync _|leader, 0xFFFFFFFF;\n\t" "setp.ne.b32 p0, %3, 0;\n\t" "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "" "mov.b32 adhi, 0x40004040;\n\t" "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 136314896;\n\t" + "mov.b32 id, 69206032;\n\t" "mov.b32 alo, %0;\n\t" "mov.b32 blo, %1;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 2;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 2;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 2;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "}\n" - :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); - if (cta_rank == 0) { - elect_commit_cg2_multicast(s_full_addr + (phase_2) * 8, (uint16_t)(3)); - } - if (cta_rank == 0) { - elect_commit_cg2_multicast(k_empty_addr + (mma_k_stage) * 8, (uint16_t)(3)); - } - } - mma_k_stage += 1; - if (mma_k_stage == 2) { mma_k_stage = 0; _phase_k_full ^= 1; } - } - if (tile_3 > 0) { - int prev_pipeline_tile = pipeline_tile_2 - 1; - int prev_phase = prev_pipeline_tile & 1; - int pv_wait_phase = prev_pipeline_tile >> 1 & 1; - mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); - mbarrier_wait(corr_done_addr + (prev_phase) * 8, pv_wait_phase); + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"((tmem_tmem_scratch + (wt_score_col))), "r"(1)); + elect_commit(s_full_addr + (wt_phase) * 8); + if (wt_tile_1 > 0) { + int wt_prev_pipeline_tile = wt_pipeline_tile - 1; + int wt_prev_phase = wt_prev_pipeline_tile & 1; + int wt_pv_wait_phase = wt_prev_pipeline_tile >> 1 & 1; + mbarrier_wait(p_full_addr + (wt_prev_phase) * 8, wt_pv_wait_phase); + mbarrier_wait(corr_done_addr + (wt_prev_phase) * 8, wt_pv_wait_phase); asm volatile("tcgen05.fence::after_thread_sync;"); - { - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(pv_pair_ready_addr + (unsigned int)(prev_phase * 8)), "r"(0) : "memory"); - if (cta_rank == 0) { - mbarrier_wait_cluster_hint(pv_pair_ready_addr + (prev_phase) * 8, pv_wait_phase, 10000000); - } - } - int p_col = ((prev_phase != 0) ? 224 : 96); - { - if (cta_rank == 0) { - mbarrier_wait(v_full_addr + (mma_v_stage) * 8, _phase_v_full); - int _mma_a_lo_6 = (((smem_p_fp8_addr) >> 4) & 0x3FFF) + (prev_phase) * 512; - int _mma_b_lo_6 = ((((smem_v_full_addr) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; - asm volatile( + unsigned int wt_prev_kv_stage = mma_k_stage ^ 1; + int _mma_a_lo_4 = make_warp_uniform((((smem_p_fp8_addr) >> 4) & 0x3FFF) + (wt_prev_phase) * 512); + int _mma_b_lo_4 = make_warp_uniform(((((smem_v_full_addr) >> 4) & 0x3FFF) | 0x4000000) + (wt_prev_kv_stage) * 4096); + asm volatile( "{\n\t" ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" ".reg .b64 da, db;\n\t" "elect.sync _|leader, 0xFFFFFFFF;\n\t" "setp.ne.b32 p0, %3, 0;\n\t" "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "" "mov.b32 adhi, 0x40004040;\n\t" "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138477584;\n\t" + "mov.b32 id, 69271568;\n\t" "mov.b32 alo, %0;\n\t" "mov.b32 blo, %1;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 256;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 256;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 256;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "}\n" - :: "r"(_mma_a_lo_6), "r"(_mma_b_lo_6), "r"((tmem_tmem_scratch + (256))), "r"(((first_pv) ? 0 : 1))); - int _mma_a_lo_7 = (((smem_p_fp8_addr) >> 4) & 0x3FFF) + (prev_phase) * 512; - int _mma_b_lo_7 = ((((smem_v_full_addr + 16384) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; - asm volatile( + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"((tmem_tmem_scratch + (256))), "r"(((wt_first_pv) ? 0 : 1))); + int _mma_a_lo_5 = make_warp_uniform((((smem_p_fp8_addr) >> 4) & 0x3FFF) + (wt_prev_phase) * 512); + int _mma_b_lo_5 = make_warp_uniform(((((smem_v_full_addr + 16384) >> 4) & 0x3FFF) | 0x4000000) + (wt_prev_kv_stage) * 4096); + asm volatile( "{\n\t" ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" ".reg .b64 da, db;\n\t" "elect.sync _|leader, 0xFFFFFFFF;\n\t" "setp.ne.b32 p0, %3, 0;\n\t" "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "" "mov.b32 adhi, 0x40004040;\n\t" "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138477584;\n\t" + "mov.b32 id, 69271568;\n\t" "mov.b32 alo, %0;\n\t" "mov.b32 blo, %1;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 256;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 256;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 256;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "}\n" - :: "r"(_mma_a_lo_7), "r"(_mma_b_lo_7), "r"((tmem_tmem_scratch + (384))), "r"(((first_pv) ? 0 : 1))); - if (cta_rank == 0) { - elect_commit_cg2_multicast(v_empty_addr + (mma_v_stage) * 8, (uint16_t)(3)); - } - } - mma_v_stage += 1; - if (mma_v_stage == 2) { mma_v_stage = 0; _phase_v_full ^= 1; } - } - first_pv = 0; - if (cta_rank == 0) { - elect_commit_cg2_multicast(pv_done_addr, (uint16_t)(3)); - } - } - } - int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_3 - 1; - int last_phase = last_pipeline_tile & 1; - int drain_wait_phase = last_pipeline_tile >> 1 & 1; - mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); - mbarrier_wait(corr_done_addr + (last_phase) * 8, drain_wait_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - { - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(pv_pair_ready_addr + (unsigned int)(last_phase * 8)), "r"(0) : "memory"); - if (cta_rank == 0) { - mbarrier_wait_cluster_hint(pv_pair_ready_addr + (last_phase) * 8, drain_wait_phase, 10000000); - } - } - int p_col_last = ((last_phase != 0) ? 224 : 96); - { - if (cta_rank == 0) { - mbarrier_wait(v_full_addr + (mma_v_stage) * 8, _phase_v_full); - int _mma_a_lo_10 = (((smem_p_fp8_addr) >> 4) & 0x3FFF) + (last_phase) * 512; - int _mma_b_lo_10 = ((((smem_v_full_addr) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; + :: "r"(_mma_a_lo_5), "r"(_mma_b_lo_5), "r"((tmem_tmem_scratch + (320))), "r"(((wt_first_pv) ? 0 : 1))); + int _mma_a_lo_6 = make_warp_uniform((((smem_p_fp8_addr) >> 4) & 0x3FFF) + (wt_prev_phase) * 512); + int _mma_b_lo_6 = make_warp_uniform(((((smem_v_full_addr + 32768) >> 4) & 0x3FFF) | 0x4000000) + (wt_prev_kv_stage) * 4096); asm volatile( "{\n\t" ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" ".reg .b64 da, db;\n\t" "elect.sync _|leader, 0xFFFFFFFF;\n\t" "setp.ne.b32 p0, %3, 0;\n\t" "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "" "mov.b32 adhi, 0x40004040;\n\t" "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138477584;\n\t" + "mov.b32 id, 69271568;\n\t" "mov.b32 alo, %0;\n\t" "mov.b32 blo, %1;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 256;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 256;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 256;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "}\n" - :: "r"(_mma_a_lo_10), "r"(_mma_b_lo_10), "r"((tmem_tmem_scratch + (256))), "r"(((first_pv) ? 0 : 1))); - int _mma_a_lo_11 = (((smem_p_fp8_addr) >> 4) & 0x3FFF) + (last_phase) * 512; - int _mma_b_lo_11 = ((((smem_v_full_addr + 16384) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; + :: "r"(_mma_a_lo_6), "r"(_mma_b_lo_6), "r"((tmem_tmem_scratch + (384))), "r"(((wt_first_pv) ? 0 : 1))); + int _mma_a_lo_7 = make_warp_uniform((((smem_p_fp8_addr) >> 4) & 0x3FFF) + (wt_prev_phase) * 512); + int _mma_b_lo_7 = make_warp_uniform(((((smem_v_full_addr + 49152) >> 4) & 0x3FFF) | 0x4000000) + (wt_prev_kv_stage) * 4096); asm volatile( "{\n\t" ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" ".reg .b64 da, db;\n\t" "elect.sync _|leader, 0xFFFFFFFF;\n\t" "setp.ne.b32 p0, %3, 0;\n\t" "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "" "mov.b32 adhi, 0x40004040;\n\t" "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138477584;\n\t" + "mov.b32 id, 69271568;\n\t" "mov.b32 alo, %0;\n\t" "mov.b32 blo, %1;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 256;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 256;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 256;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "}\n" - :: "r"(_mma_a_lo_11), "r"(_mma_b_lo_11), "r"((tmem_tmem_scratch + (384))), "r"(((first_pv) ? 0 : 1))); - if (cta_rank == 0) { - elect_commit_cg2_multicast(v_empty_addr + (mma_v_stage) * 8, (uint16_t)(3)); - } + :: "r"(_mma_a_lo_7), "r"(_mma_b_lo_7), "r"((tmem_tmem_scratch + (448))), "r"(((wt_first_pv) ? 0 : 1))); + wt_first_pv = 0; + elect_commit(k_empty_addr + (wt_prev_kv_stage) * 8); + elect_commit(pv_done_addr); } - mma_v_stage += 1; - if (mma_v_stage == 2) { mma_v_stage = 0; _phase_v_full ^= 1; } - } - if (cta_rank == 0) { - elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + mma_k_stage += 1; + if (mma_k_stage == 2) { mma_k_stage = 0; _phase_k_full ^= 1; } } - if (cta_rank == 0) { - elect_commit_cg2_multicast(o_done_addr, (uint16_t)(3)); + if (wt_num_local_tiles_1 > 0) { + int wt_last_pipeline_tile = mma_tile_cursor + wt_num_local_tiles_1 - 1; + int wt_last_phase = wt_last_pipeline_tile & 1; + int wt_drain_wait_phase = wt_last_pipeline_tile >> 1 & 1; + mbarrier_wait(p_full_addr + (wt_last_phase) * 8, wt_drain_wait_phase); + mbarrier_wait(corr_done_addr + (wt_last_phase) * 8, wt_drain_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int wt_last_kv_stage = mma_k_stage ^ 1; + int _mma_a_lo_8 = make_warp_uniform((((smem_p_fp8_addr) >> 4) & 0x3FFF) + (wt_last_phase) * 512); + int _mma_b_lo_8 = make_warp_uniform(((((smem_v_full_addr) >> 4) & 0x3FFF) | 0x4000000) + (wt_last_kv_stage) * 4096); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69271568;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_8), "r"(_mma_b_lo_8), "r"((tmem_tmem_scratch + (256))), "r"(((wt_first_pv) ? 0 : 1))); + int _mma_a_lo_9 = make_warp_uniform((((smem_p_fp8_addr) >> 4) & 0x3FFF) + (wt_last_phase) * 512); + int _mma_b_lo_9 = make_warp_uniform(((((smem_v_full_addr + 16384) >> 4) & 0x3FFF) | 0x4000000) + (wt_last_kv_stage) * 4096); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69271568;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_9), "r"(_mma_b_lo_9), "r"((tmem_tmem_scratch + (320))), "r"(((wt_first_pv) ? 0 : 1))); + int _mma_a_lo_10 = make_warp_uniform((((smem_p_fp8_addr) >> 4) & 0x3FFF) + (wt_last_phase) * 512); + int _mma_b_lo_10 = make_warp_uniform(((((smem_v_full_addr + 32768) >> 4) & 0x3FFF) | 0x4000000) + (wt_last_kv_stage) * 4096); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69271568;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_10), "r"(_mma_b_lo_10), "r"((tmem_tmem_scratch + (384))), "r"(((wt_first_pv) ? 0 : 1))); + int _mma_a_lo_11 = make_warp_uniform((((smem_p_fp8_addr) >> 4) & 0x3FFF) + (wt_last_phase) * 512); + int _mma_b_lo_11 = make_warp_uniform(((((smem_v_full_addr + 49152) >> 4) & 0x3FFF) | 0x4000000) + (wt_last_kv_stage) * 4096); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69271568;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_11), "r"(_mma_b_lo_11), "r"((tmem_tmem_scratch + (448))), "r"(((wt_first_pv) ? 0 : 1))); + elect_commit(k_empty_addr + (wt_last_kv_stage) * 8); } - mma_tile_cursor = mma_tile_cursor + num_kv_tiles_3; + elect_commit(q_empty_addr); + elect_commit(o_done_addr); + mma_tile_cursor = mma_tile_cursor + wt_num_local_tiles_1; } mbarrier_arrive(tmem_dealloc_addr); unsigned int _phase_tmem_dealloc_0 = 0; mbarrier_wait(tmem_dealloc_addr, _phase_tmem_dealloc_0); _phase_tmem_dealloc_0 ^= 1; asm volatile("tcgen05.fence::after_thread_sync;"); - { - int tmem_dealloc_peer_rank = bid % 2 ^ 1; - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(tmem_dealloc_peer_addr), "r"(tmem_dealloc_peer_rank) : "memory"); - mbarrier_wait(tmem_dealloc_peer_addr, 0); - int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); - asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); - } + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); } } // ---- Role: empty1 ---- diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_f55f69afe4ac3ec96817_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_bab2852cc0ea4a6d4860_binding.cu similarity index 97% rename from csrc/cake_dsv4/sm_103a/cake_dsv4_f55f69afe4ac3ec96817_binding.cu rename to csrc/cake_dsv4/sm_100a/cake_dsv4_bab2852cc0ea4a6d4860_binding.cu index d0f749f0def..a3213ad3a76 100644 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_f55f69afe4ac3ec96817_binding.cu +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_bab2852cc0ea4a6d4860_binding.cu @@ -15,7 +15,7 @@ */ // Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_f55f69afe4ac3ec96817'. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_bab2852cc0ea4a6d4860'. #include #include #include @@ -31,10 +31,10 @@ #include #include -extern "C" __global__ void kernel_cake_dsv4_f55f69afe4ac3ec96817(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ cum_seq_lens_q, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); +extern "C" __global__ void kernel_cake_dsv4_bab2852cc0ea4a6d4860(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ cum_seq_lens_q, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); -namespace cake_host_shim_90ab13be7c9c8334 { +namespace cake_host_shim_66b5cf782c79bc55 { using tvm::ffi::Optional; using tvm::ffi::TensorView; @@ -409,7 +409,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ } if (smem_status == cudaSuccess) { smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_f55f69afe4ac3ec96817), + reinterpret_cast(kernel_cake_dsv4_bab2852cc0ea4a6d4860), cudaFuncAttributeMaxDynamicSharedMemorySize, 203136); } @@ -419,7 +419,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ } } TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_f55f69afe4ac3ec96817 failed: " + << "cudaFuncSetAttribute for kernel_cake_dsv4_bab2852cc0ea4a6d4860 failed: " << cudaGetErrorString(smem_status); dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); dim3 block(512u, 1u, 1u); @@ -430,12 +430,12 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ config.dynamicSmemBytes = 203136u; config.stream = stream; cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_f55f69afe4ac3ec96817), kargs); + &config, reinterpret_cast(kernel_cake_dsv4_bab2852cc0ea4a6d4860), kargs); TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_f55f69afe4ac3ec96817 failed: " + << "cudaLaunchKernelExC for kernel_cake_dsv4_bab2852cc0ea4a6d4860 failed: " << cudaGetErrorString(launch_status); } -} // namespace cake_host_shim_90ab13be7c9c8334 +} // namespace cake_host_shim_66b5cf782c79bc55 -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_90ab13be7c9c8334::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_66b5cf782c79bc55::Run); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_5e5d6648b7bd42bcd0d6_kernel.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_bab2852cc0ea4a6d4860_kernel.cu similarity index 99% rename from csrc/cake_dsv4/sm_100a/cake_dsv4_5e5d6648b7bd42bcd0d6_kernel.cu rename to csrc/cake_dsv4/sm_100a/cake_dsv4_bab2852cc0ea4a6d4860_kernel.cu index 2a64d3f6d3e..cb8f8e0f09e 100644 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_5e5d6648b7bd42bcd0d6_kernel.cu +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_bab2852cc0ea4a6d4860_kernel.cu @@ -365,6 +365,19 @@ __device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { } +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + __device__ __forceinline__ float approx_exp2(float x) { float y; asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); @@ -1062,35 +1075,12 @@ __device__ __forceinline__ void tcgen05_commit(int mbar_addr) { } -__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x16.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) - : "r"(tmem_addr)); -} - - __device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { tmem_ld_x16(dst, addr); asm volatile("tcgen05.wait::ld.sync.aligned;"); } -__device__ __forceinline__ void tmem_ld_x8(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x8.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]) - : "r"(tmem_addr)); -} - - __device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { uint32_t result; asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" @@ -1132,7 +1122,7 @@ cake_pipeline_state_at(CakePipelineState state) { extern "C" { __global__ __launch_bounds__(512, 1) void -kernel_cake_dsv4_5e5d6648b7bd42bcd0d6(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ cum_seq_lens_q, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items) +kernel_cake_dsv4_bab2852cc0ea4a6d4860(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ cum_seq_lens_q, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); @@ -1926,25 +1916,30 @@ kernel_cake_dsv4_5e5d6648b7bd42bcd0d6(const __grid_constant__ CUtensorMap tmap_q int row_bank_1 = ((output_half_1 != 0) ? 1048576 : 0); int col_base = col_half_1 * 128; #pragma unroll 2 - for (int col_1 = 0; col_1 < 128; col_1 += 8) { - float _tmem_load_4[8]; + for (int wcol = 0; wcol < 128; wcol += 16) { + float _tmem_load_4[16]; asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x8.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8], 128;" - : "=r"(*reinterpret_cast(&_tmem_load_4[0])), "=r"(*reinterpret_cast(&_tmem_load_4[1])), "=r"(*reinterpret_cast(&_tmem_load_4[2])), "=r"(*reinterpret_cast(&_tmem_load_4[3])), "=r"(*reinterpret_cast(&_tmem_load_4[4])), "=r"(*reinterpret_cast(&_tmem_load_4[5])), "=r"(*reinterpret_cast(&_tmem_load_4[6])), "=r"(*reinterpret_cast(&_tmem_load_4[7])) - : "r"(taddr + 256 + (unsigned int)row_bank_1 + (unsigned int)col_1)); + "tcgen05.ld.sync.aligned.16x32bx2.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16], 128;" + : "=r"(*reinterpret_cast(&_tmem_load_4[0])), "=r"(*reinterpret_cast(&_tmem_load_4[1])), "=r"(*reinterpret_cast(&_tmem_load_4[2])), "=r"(*reinterpret_cast(&_tmem_load_4[3])), "=r"(*reinterpret_cast(&_tmem_load_4[4])), "=r"(*reinterpret_cast(&_tmem_load_4[5])), "=r"(*reinterpret_cast(&_tmem_load_4[6])), "=r"(*reinterpret_cast(&_tmem_load_4[7])), "=r"(*reinterpret_cast(&_tmem_load_4[8])), "=r"(*reinterpret_cast(&_tmem_load_4[9])), "=r"(*reinterpret_cast(&_tmem_load_4[10])), "=r"(*reinterpret_cast(&_tmem_load_4[11])), "=r"(*reinterpret_cast(&_tmem_load_4[12])), "=r"(*reinterpret_cast(&_tmem_load_4[13])), "=r"(*reinterpret_cast(&_tmem_load_4[14])), "=r"(*reinterpret_cast(&_tmem_load_4[15])) + : "r"(taddr + 256 + (unsigned int)row_bank_1 + (unsigned int)wcol)); const float2 _scale2_1 = {inv_sum, inv_sum}; #pragma unroll - for (int _ls = 0; _ls < 4; _ls++) + for (int _ls = 0; _ls < 8; _ls++) mul_f32x2_inplace(&reinterpret_cast(_tmem_load_4)[_ls], _scale2_1); - uint32_t _tmem_load_4_bf16[4]; + uint32_t _tmem_load_4_bf16[8]; #pragma unroll - for (int _lp = 0; _lp < 4; _lp++) { + for (int _lp = 0; _lp < 8; _lp++) { __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_4[_lp*2 + 0], _tmem_load_4[_lp*2+1 + 0])); _tmem_load_4_bf16[_lp] = *(uint32_t*)&_bf2; } if (output_in_bounds != 0) { - reinterpret_cast(O + (output_base + output_half_1 * 256 + col_base + col_1))[0] = reinterpret_cast(_tmem_load_4_bf16)[0]; + { + const unsigned* _raw_stv8_2 = reinterpret_cast(_tmem_load_4_bf16); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(O + (output_base + output_half_1 * 256 + col_base + wcol))), "r"(_raw_stv8_2[0]), "r"(_raw_stv8_2[1]), "r"(_raw_stv8_2[2]), "r"(_raw_stv8_2[3]), "r"(_raw_stv8_2[4]), "r"(_raw_stv8_2[5]), "r"(_raw_stv8_2[6]), "r"(_raw_stv8_2[7]) : "memory"); + } } } } diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_family_068c73b21d7cbcd958b8_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_family_5383a615a7e502f6fdac_binding.cu similarity index 87% rename from csrc/cake_dsv4/sm_100a/cake_dsv4_family_068c73b21d7cbcd958b8_binding.cu rename to csrc/cake_dsv4/sm_100a/cake_dsv4_family_5383a615a7e502f6fdac_binding.cu index 4515647e55f..3047ed0720b 100644 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_family_068c73b21d7cbcd958b8_binding.cu +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_family_5383a615a7e502f6fdac_binding.cu @@ -15,7 +15,7 @@ */ // Generated by cake.runtime.ffi_dispatch — do not edit. -// tvm-ffi dispatch module for kernel family 'cake_dsv4_family_068c73b21d7cbcd958b8'. +// tvm-ffi dispatch module for kernel family 'cake_dsv4_family_5383a615a7e502f6fdac'. // dispatch-policy-revision: unpinned #include #include @@ -34,9 +34,9 @@ #include extern "C" __global__ void kernel_cake_dsv4_d906a03bd66a7488fc29(__nv_bfloat16* __restrict__ partial_O, float* __restrict__ partial_lse, __nv_bfloat16* __restrict__ O, int num_q_heads, int num_split); -extern "C" __global__ void kernel_cake_dsv4_ba9ed6a6233761682cc3(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); +extern "C" __global__ void kernel_cake_dsv4_714df747e199d0adbe08(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, const __grid_constant__ CUtensorMap O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); -namespace cake_host_shim_3cfd00f6b2ce7022 { +namespace cake_host_shim_3ce874e06ea0fc76 { using tvm::ffi::Optional; using tvm::ffi::TensorView; @@ -161,7 +161,7 @@ inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* } } -namespace variant_cake_dsv4_ba9ed6a6233761682cc3_sm_100a_f23ac2823b80 { +namespace variant_cake_dsv4_714df747e199d0adbe08_sm_100a_591fb3561ebd { // 4D TMA descriptor for buffer 'tmap_q' — compiled from the // descriptor's std.Expr global_dim/global_strides/checks record. @@ -394,6 +394,57 @@ inline CUtensorMap EncodeTma_tmap_compressed_v(const TensorView& t) { return tm; } +// 4D TMA descriptor for buffer 'O' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_O(const TensorView& t, int64_t tp_num_heads, int64_t tp_num_query_tokens) { + TVM_FFI_CHECK(t.ndim() >= 1, ValueError) + << "TMA source 'O' must have at least 1 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'O' must have unit innermost stride, got " << t.stride(-1); + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(8), (uint64_t)(5), (uint64_t)((tp_num_query_tokens * tp_num_heads))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'O' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1] && 1u <= global_dim[2] && 32u <= global_dim[3], ValueError) + << "TMA box (64, 1, 1, 32) exceeds resolved global dims for 'O'"; + int64_t carrier_stride_0 = 64; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 512; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = 2560; + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 1u, 1u, 32u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'O') failed: CUresult=" << (int)r; + return tm; +} + void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_compressed_k, TensorView arg_tmap_swa_v, TensorView arg_tmap_compressed_v, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_num_query_tokens, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { DLDevice dev = arg_tmap_q.device(); ScopedCudaDevice device_guard(dev.device_id); @@ -481,7 +532,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_c CUtensorMap p_tmap_compressed_k = EncodeTma_tmap_compressed_k(arg_tmap_compressed_k); CUtensorMap p_tmap_swa_v = EncodeTma_tmap_swa_v(arg_tmap_swa_v); CUtensorMap p_tmap_compressed_v = EncodeTma_tmap_compressed_v(arg_tmap_compressed_v); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); + CUtensorMap p_O = EncodeTma_O(arg_O, (int64_t)arg_num_heads, (int64_t)arg_num_query_tokens); float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); @@ -514,7 +565,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_c } if (smem_status == cudaSuccess) { smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_ba9ed6a6233761682cc3), + reinterpret_cast(kernel_cake_dsv4_714df747e199d0adbe08), cudaFuncAttributeMaxDynamicSharedMemorySize, 231808); } @@ -524,7 +575,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_c } } TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_ba9ed6a6233761682cc3 failed: " + << "cudaFuncSetAttribute for kernel_cake_dsv4_714df747e199d0adbe08 failed: " << cudaGetErrorString(smem_status); dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); dim3 block(512u, 1u, 1u); @@ -545,12 +596,12 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_c config.attrs = attrs; config.numAttrs = n; cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_ba9ed6a6233761682cc3), kargs); + &config, reinterpret_cast(kernel_cake_dsv4_714df747e199d0adbe08), kargs); TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_ba9ed6a6233761682cc3 failed: " + << "cudaLaunchKernelExC for kernel_cake_dsv4_714df747e199d0adbe08 failed: " << cudaGetErrorString(launch_status); } -} // namespace variant_cake_dsv4_ba9ed6a6233761682cc3_sm_100a_f23ac2823b80 +} // namespace variant_cake_dsv4_714df747e199d0adbe08_sm_100a_591fb3561ebd namespace variant_cake_dsv4_d906a03bd66a7488fc29_sm_100a_cd6fb31da15f { @@ -629,24 +680,24 @@ void Dispatch(tvm::ffi::TensorView arg_Q, tvm::ffi::TensorView arg_SWA_cache, tv int64_t gx = (total_work_items * 2); int64_t gy = 1; int64_t gz = 1; - cake_host_shim_3cfd00f6b2ce7022::variant_cake_dsv4_ba9ed6a6233761682cc3_sm_100a_f23ac2823b80::Run(arg_Q, arg_SWA_cache, arg_compressed_KV_cache, arg_SWA_cache, arg_compressed_KV_cache, arg_partial_O, arg_partial_lse, arg_swa_indices, arg_compressed_indices, arg_sparse_topk_lens, arg_sinks, arg_bmm1_scale, arg_bmm2_scale, num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, gx, gy, gz, stream); + cake_host_shim_3ce874e06ea0fc76::variant_cake_dsv4_714df747e199d0adbe08_sm_100a_591fb3561ebd::Run(arg_Q, arg_SWA_cache, arg_compressed_KV_cache, arg_SWA_cache, arg_compressed_KV_cache, arg_partial_O, arg_partial_lse, arg_swa_indices, arg_compressed_indices, arg_sparse_topk_lens, arg_sinks, arg_bmm1_scale, arg_bmm2_scale, num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, gx, gy, gz, stream); } { int64_t gx = num_query_tokens; int64_t gy = num_heads; int64_t gz = 1; - cake_host_shim_3cfd00f6b2ce7022::variant_cake_dsv4_d906a03bd66a7488fc29_sm_100a_cd6fb31da15f::Run(arg_partial_O, arg_partial_lse, arg_O, num_heads, num_split, gx, gy, gz, stream); + cake_host_shim_3ce874e06ea0fc76::variant_cake_dsv4_d906a03bd66a7488fc29_sm_100a_cd6fb31da15f::Run(arg_partial_O, arg_partial_lse, arg_O, num_heads, num_split, gx, gy, gz, stream); } return; } default: - TVM_FFI_CHECK(false, ValueError) << "no dispatch route for cake_dsv4_family_068c73b21d7cbcd958b8 shape (num_heads=" << num_heads << ", num_query_tokens=" << num_query_tokens << ", swa_index_stride=" << swa_index_stride << ", compressed_index_stride=" << compressed_index_stride << ", sparse_topk_lens_offset=" << sparse_topk_lens_offset << ", sparse_topk=" << sparse_topk << ", has_sinks=" << has_sinks << ", total_work_items=" << total_work_items << ", num_split=" << num_split << ")"; + TVM_FFI_CHECK(false, ValueError) << "no dispatch route for cake_dsv4_family_5383a615a7e502f6fdac shape (num_heads=" << num_heads << ", num_query_tokens=" << num_query_tokens << ", swa_index_stride=" << swa_index_stride << ", compressed_index_stride=" << compressed_index_stride << ", sparse_topk_lens_offset=" << sparse_topk_lens_offset << ", sparse_topk=" << sparse_topk << ", has_sinks=" << has_sinks << ", total_work_items=" << total_work_items << ", num_split=" << num_split << ")"; } } int64_t RouteOf(int32_t num_heads, int32_t num_query_tokens, int32_t swa_index_stride, int32_t compressed_index_stride, int32_t sparse_topk_lens_offset, int32_t sparse_topk, int32_t has_sinks, int32_t total_work_items, int32_t num_split) { return SelectVariant(num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, num_split); } -} // namespace cake_host_shim_3cfd00f6b2ce7022 +} // namespace cake_host_shim_3ce874e06ea0fc76 -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_3cfd00f6b2ce7022::Dispatch); -TVM_FFI_DLL_EXPORT_TYPED_FUNC(route_of, cake_host_shim_3cfd00f6b2ce7022::RouteOf); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_3ce874e06ea0fc76::Dispatch); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(route_of, cake_host_shim_3ce874e06ea0fc76::RouteOf); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_family_a84dbf9944ef3cfb088b_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_family_9c96528827a52c42bbd7_binding.cu similarity index 87% rename from csrc/cake_dsv4/sm_100a/cake_dsv4_family_a84dbf9944ef3cfb088b_binding.cu rename to csrc/cake_dsv4/sm_100a/cake_dsv4_family_9c96528827a52c42bbd7_binding.cu index c24f61f5b90..d224bcffd84 100644 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_family_a84dbf9944ef3cfb088b_binding.cu +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_family_9c96528827a52c42bbd7_binding.cu @@ -15,7 +15,7 @@ */ // Generated by cake.runtime.ffi_dispatch — do not edit. -// tvm-ffi dispatch module for kernel family 'cake_dsv4_family_a84dbf9944ef3cfb088b'. +// tvm-ffi dispatch module for kernel family 'cake_dsv4_family_9c96528827a52c42bbd7'. // dispatch-policy-revision: unpinned #include #include @@ -34,9 +34,9 @@ #include extern "C" __global__ void kernel_cake_dsv4_d50dd820272ad24f7027(__nv_bfloat16* __restrict__ partial_O, float* __restrict__ partial_lse, __nv_bfloat16* __restrict__ O, int num_q_heads, int num_split); -extern "C" __global__ void kernel_cake_dsv4_e1660a7f44d156286fea(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); +extern "C" __global__ void kernel_cake_dsv4_42fa8e98da5fa46d1a94(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, const __grid_constant__ CUtensorMap O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); -namespace cake_host_shim_c06b4c8dc8a15cae { +namespace cake_host_shim_6534755d2bf16a2e { using tvm::ffi::Optional; using tvm::ffi::TensorView; @@ -161,7 +161,7 @@ inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* } } -namespace variant_cake_dsv4_e1660a7f44d156286fea_sm_100a_994fe0d2939f { +namespace variant_cake_dsv4_42fa8e98da5fa46d1a94_sm_100a_6276393afd19 { // 4D TMA descriptor for buffer 'tmap_q' — compiled from the // descriptor's std.Expr global_dim/global_strides/checks record. @@ -394,6 +394,57 @@ inline CUtensorMap EncodeTma_tmap_compressed_v(const TensorView& t) { return tm; } +// 4D TMA descriptor for buffer 'O' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_O(const TensorView& t, int64_t tp_num_heads, int64_t tp_num_query_tokens) { + TVM_FFI_CHECK(t.ndim() >= 1, ValueError) + << "TMA source 'O' must have at least 1 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'O' must have unit innermost stride, got " << t.stride(-1); + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(8), (uint64_t)(4), (uint64_t)((tp_num_query_tokens * tp_num_heads))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'O' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1] && 1u <= global_dim[2] && 32u <= global_dim[3], ValueError) + << "TMA box (64, 1, 1, 32) exceeds resolved global dims for 'O'"; + int64_t carrier_stride_0 = 64; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 512; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = 2048; + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 1u, 1u, 32u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'O') failed: CUresult=" << (int)r; + return tm; +} + void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_compressed_k, TensorView arg_tmap_swa_v, TensorView arg_tmap_compressed_v, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_num_query_tokens, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { DLDevice dev = arg_tmap_q.device(); ScopedCudaDevice device_guard(dev.device_id); @@ -481,7 +532,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_c CUtensorMap p_tmap_compressed_k = EncodeTma_tmap_compressed_k(arg_tmap_compressed_k); CUtensorMap p_tmap_swa_v = EncodeTma_tmap_swa_v(arg_tmap_swa_v); CUtensorMap p_tmap_compressed_v = EncodeTma_tmap_compressed_v(arg_tmap_compressed_v); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); + CUtensorMap p_O = EncodeTma_O(arg_O, (int64_t)arg_num_heads, (int64_t)arg_num_query_tokens); float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); @@ -514,7 +565,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_c } if (smem_status == cudaSuccess) { smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_e1660a7f44d156286fea), + reinterpret_cast(kernel_cake_dsv4_42fa8e98da5fa46d1a94), cudaFuncAttributeMaxDynamicSharedMemorySize, 220032); } @@ -524,7 +575,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_c } } TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_e1660a7f44d156286fea failed: " + << "cudaFuncSetAttribute for kernel_cake_dsv4_42fa8e98da5fa46d1a94 failed: " << cudaGetErrorString(smem_status); dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); dim3 block(512u, 1u, 1u); @@ -545,12 +596,12 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_c config.attrs = attrs; config.numAttrs = n; cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_e1660a7f44d156286fea), kargs); + &config, reinterpret_cast(kernel_cake_dsv4_42fa8e98da5fa46d1a94), kargs); TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_e1660a7f44d156286fea failed: " + << "cudaLaunchKernelExC for kernel_cake_dsv4_42fa8e98da5fa46d1a94 failed: " << cudaGetErrorString(launch_status); } -} // namespace variant_cake_dsv4_e1660a7f44d156286fea_sm_100a_994fe0d2939f +} // namespace variant_cake_dsv4_42fa8e98da5fa46d1a94_sm_100a_6276393afd19 namespace variant_cake_dsv4_d50dd820272ad24f7027_sm_100a_aa3105f3b9de { @@ -629,24 +680,24 @@ void Dispatch(tvm::ffi::TensorView arg_Q, tvm::ffi::TensorView arg_SWA_cache, tv int64_t gx = (total_work_items * 2); int64_t gy = 1; int64_t gz = 1; - cake_host_shim_c06b4c8dc8a15cae::variant_cake_dsv4_e1660a7f44d156286fea_sm_100a_994fe0d2939f::Run(arg_Q, arg_SWA_cache, arg_compressed_KV_cache, arg_SWA_cache, arg_compressed_KV_cache, arg_partial_O, arg_partial_lse, arg_swa_indices, arg_compressed_indices, arg_sparse_topk_lens, arg_sinks, arg_bmm1_scale, arg_bmm2_scale, num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, gx, gy, gz, stream); + cake_host_shim_6534755d2bf16a2e::variant_cake_dsv4_42fa8e98da5fa46d1a94_sm_100a_6276393afd19::Run(arg_Q, arg_SWA_cache, arg_compressed_KV_cache, arg_SWA_cache, arg_compressed_KV_cache, arg_partial_O, arg_partial_lse, arg_swa_indices, arg_compressed_indices, arg_sparse_topk_lens, arg_sinks, arg_bmm1_scale, arg_bmm2_scale, num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, gx, gy, gz, stream); } { int64_t gx = num_query_tokens; int64_t gy = num_heads; int64_t gz = 1; - cake_host_shim_c06b4c8dc8a15cae::variant_cake_dsv4_d50dd820272ad24f7027_sm_100a_aa3105f3b9de::Run(arg_partial_O, arg_partial_lse, arg_O, num_heads, num_split, gx, gy, gz, stream); + cake_host_shim_6534755d2bf16a2e::variant_cake_dsv4_d50dd820272ad24f7027_sm_100a_aa3105f3b9de::Run(arg_partial_O, arg_partial_lse, arg_O, num_heads, num_split, gx, gy, gz, stream); } return; } default: - TVM_FFI_CHECK(false, ValueError) << "no dispatch route for cake_dsv4_family_a84dbf9944ef3cfb088b shape (num_heads=" << num_heads << ", num_query_tokens=" << num_query_tokens << ", swa_index_stride=" << swa_index_stride << ", compressed_index_stride=" << compressed_index_stride << ", sparse_topk_lens_offset=" << sparse_topk_lens_offset << ", sparse_topk=" << sparse_topk << ", has_sinks=" << has_sinks << ", total_work_items=" << total_work_items << ", num_split=" << num_split << ")"; + TVM_FFI_CHECK(false, ValueError) << "no dispatch route for cake_dsv4_family_9c96528827a52c42bbd7 shape (num_heads=" << num_heads << ", num_query_tokens=" << num_query_tokens << ", swa_index_stride=" << swa_index_stride << ", compressed_index_stride=" << compressed_index_stride << ", sparse_topk_lens_offset=" << sparse_topk_lens_offset << ", sparse_topk=" << sparse_topk << ", has_sinks=" << has_sinks << ", total_work_items=" << total_work_items << ", num_split=" << num_split << ")"; } } int64_t RouteOf(int32_t num_heads, int32_t num_query_tokens, int32_t swa_index_stride, int32_t compressed_index_stride, int32_t sparse_topk_lens_offset, int32_t sparse_topk, int32_t has_sinks, int32_t total_work_items, int32_t num_split) { return SelectVariant(num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, num_split); } -} // namespace cake_host_shim_c06b4c8dc8a15cae +} // namespace cake_host_shim_6534755d2bf16a2e -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_c06b4c8dc8a15cae::Dispatch); -TVM_FFI_DLL_EXPORT_TYPED_FUNC(route_of, cake_host_shim_c06b4c8dc8a15cae::RouteOf); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_6534755d2bf16a2e::Dispatch); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(route_of, cake_host_shim_6534755d2bf16a2e::RouteOf); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_4c41fb8347f137d4d46f_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_064850af5858dd03c104_binding.cu similarity index 86% rename from csrc/cake_dsv4/sm_103a/cake_dsv4_4c41fb8347f137d4d46f_binding.cu rename to csrc/cake_dsv4/sm_103a/cake_dsv4_064850af5858dd03c104_binding.cu index 9ea7ff58e8d..0392d451107 100644 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_4c41fb8347f137d4d46f_binding.cu +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_064850af5858dd03c104_binding.cu @@ -15,7 +15,7 @@ */ // Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_4c41fb8347f137d4d46f'. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_064850af5858dd03c104'. #include #include #include @@ -31,10 +31,10 @@ #include #include -extern "C" __global__ void kernel_cake_dsv4_4c41fb8347f137d4d46f(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); +extern "C" __global__ void kernel_cake_dsv4_064850af5858dd03c104(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, const __grid_constant__ CUtensorMap O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); -namespace cake_host_shim_22e55bd245aced16 { +namespace cake_host_shim_66599af918c96e07 { using tvm::ffi::Optional; using tvm::ffi::TensorView; @@ -390,6 +390,57 @@ inline CUtensorMap EncodeTma_tmap_compressed_v(const TensorView& t) { return tm; } +// 4D TMA descriptor for buffer 'O' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_O(const TensorView& t, int64_t tp_num_heads, int64_t tp_num_query_tokens) { + TVM_FFI_CHECK(t.ndim() >= 1, ValueError) + << "TMA source 'O' must have at least 1 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'O' must have unit innermost stride, got " << t.stride(-1); + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(8), (uint64_t)(4), (uint64_t)((tp_num_query_tokens * tp_num_heads))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'O' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1] && 1u <= global_dim[2] && 32u <= global_dim[3], ValueError) + << "TMA box (64, 1, 1, 32) exceeds resolved global dims for 'O'"; + int64_t carrier_stride_0 = 64; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 512; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = 2048; + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 1u, 1u, 32u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'O') failed: CUresult=" << (int)r; + return tm; +} + void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_compressed_k, TensorView arg_tmap_swa_v, TensorView arg_tmap_compressed_v, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_num_query_tokens, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z) { DLDevice dev = arg_tmap_q.device(); ScopedCudaDevice device_guard(dev.device_id); @@ -477,7 +528,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_c CUtensorMap p_tmap_compressed_k = EncodeTma_tmap_compressed_k(arg_tmap_compressed_k); CUtensorMap p_tmap_swa_v = EncodeTma_tmap_swa_v(arg_tmap_swa_v); CUtensorMap p_tmap_compressed_v = EncodeTma_tmap_compressed_v(arg_tmap_compressed_v); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); + CUtensorMap p_O = EncodeTma_O(arg_O, (int64_t)arg_num_heads, (int64_t)arg_num_query_tokens); float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); @@ -510,7 +561,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_c } if (smem_status == cudaSuccess) { smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_4c41fb8347f137d4d46f), + reinterpret_cast(kernel_cake_dsv4_064850af5858dd03c104), cudaFuncAttributeMaxDynamicSharedMemorySize, 220032); } @@ -520,7 +571,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_c } } TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_4c41fb8347f137d4d46f failed: " + << "cudaFuncSetAttribute for kernel_cake_dsv4_064850af5858dd03c104 failed: " << cudaGetErrorString(smem_status); dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); dim3 block(512u, 1u, 1u); @@ -544,12 +595,12 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_c config.attrs = attrs; config.numAttrs = n; cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_4c41fb8347f137d4d46f), kargs); + &config, reinterpret_cast(kernel_cake_dsv4_064850af5858dd03c104), kargs); TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_4c41fb8347f137d4d46f failed: " + << "cudaLaunchKernelExC for kernel_cake_dsv4_064850af5858dd03c104 failed: " << cudaGetErrorString(launch_status); } -} // namespace cake_host_shim_22e55bd245aced16 +} // namespace cake_host_shim_66599af918c96e07 -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_22e55bd245aced16::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_66599af918c96e07::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_4c41fb8347f137d4d46f_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_064850af5858dd03c104_kernel.cu similarity index 96% rename from csrc/cake_dsv4/sm_103a/cake_dsv4_4c41fb8347f137d4d46f_kernel.cu rename to csrc/cake_dsv4/sm_103a/cake_dsv4_064850af5858dd03c104_kernel.cu index d55adfe2ced..1016950b3e6 100644 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_4c41fb8347f137d4d46f_kernel.cu +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_064850af5858dd03c104_kernel.cu @@ -1181,6 +1181,15 @@ __device__ __forceinline__ void tma_gather4_gmem2smem_mc_cta2( } +__device__ __forceinline__ void tma_store_4d( + const void *tmap, int x, int y, int z, int w, unsigned smem_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.global.shared::cta.tile.bulk_group" + " [%0, {%1, %2, %3, %4}], [%5];" + :: "l"(tmap), "r"(x), "r"(y), "r"(z), "r"(w), "r"(smem_addr) : "memory"); +} + + __device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { asm volatile( "{\n\t" @@ -1195,7 +1204,7 @@ __device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint extern "C" { __global__ __launch_bounds__(512, 1) __cluster_dims__(2,1,1) void -kernel_cake_dsv4_4c41fb8347f137d4d46f(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items) +kernel_cake_dsv4_064850af5858dd03c104(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, const __grid_constant__ CUtensorMap O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); @@ -1220,10 +1229,11 @@ kernel_cake_dsv4_4c41fb8347f137d4d46f(const __grid_constant__ CUtensorMap tmap_q #define stats_addr (mbar_base + 168) #define stats_empty_addr (mbar_base + 184) #define o_full_addr (mbar_base + 200) - #define o_first_slice_seeded_addr (mbar_base + 208) - #define tmem_dealloc_peer_addr (mbar_base + 216) - #define index_valid_full_addr (mbar_base + 224) - #define index_valid_empty_addr (mbar_base + 240) + #define epi_q_free_addr (mbar_base + 208) + #define o_first_slice_seeded_addr (mbar_base + 216) + #define tmem_dealloc_peer_addr (mbar_base + 224) + #define index_valid_full_addr (mbar_base + 232) + #define index_valid_empty_addr (mbar_base + 248) const int bid = blockIdx.x; const int num_bids = gridDim.x; @@ -1256,8 +1266,8 @@ kernel_cake_dsv4_4c41fb8347f137d4d46f(const __grid_constant__ CUtensorMap tmap_q asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); - // Mbarrier init (18 pipeline groups, 0 ordered-sequence groups, 32 barriers) - // Mbarriers at smem_raw[0..256) + // Mbarrier init (19 pipeline groups, 0 ordered-sequence groups, 33 barriers) + // Mbarriers at smem_raw[0..264) if (warp == 0) { uint32_t leader = elect_sync(); @@ -1304,16 +1314,18 @@ kernel_cake_dsv4_4c41fb8347f137d4d46f(const __grid_constant__ CUtensorMap tmap_q mbarrier_init(smem + 192, 128); // o_full: 1 barriers, init_count=1 mbarrier_init(smem + 200, 1); + // epi_q_free: 1 barriers, init_count=4 + mbarrier_init(smem + 208, 4); // o_first_slice_seeded: 1 barriers, init_count=256 - mbarrier_init(smem + 208, 256); + mbarrier_init(smem + 216, 256); // tmem_dealloc_peer: 1 barriers, init_count=32 - mbarrier_init(smem + 216, 32); - // index_valid_full: 2 barriers, init_count=32 mbarrier_init(smem + 224, 32); + // index_valid_full: 2 barriers, init_count=32 mbarrier_init(smem + 232, 32); + mbarrier_init(smem + 240, 32); // index_valid_empty: 2 barriers, init_count=128 - mbarrier_init(smem + 240, 128); mbarrier_init(smem + 248, 128); + mbarrier_init(smem + 256, 128); asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); } } @@ -1322,9 +1334,9 @@ kernel_cake_dsv4_4c41fb8347f137d4d46f(const __grid_constant__ CUtensorMap tmap_q asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); // TMEM alloc (512 columns, 512 used) - volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 256); + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 264); if (warp == 0) { - int _tmem_hold = smem + 256; + int _tmem_hold = smem + 264; asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); __syncwarp(); asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); @@ -1696,11 +1708,15 @@ kernel_cake_dsv4_4c41fb8347f137d4d46f(const __grid_constant__ CUtensorMap tmap_q asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum)); partial_lse[stat_offset] = ((final_sum > 0.0f) ? final_max * softmax_scale_log2_1 + _log2_0 : -CAKE_INF); } + float epi_scale = inv_sum * output_scale; + int epi_box_base = smem_q_addr + (unsigned int)(warp % 4 * 16384); + int epi_row0 = query_idx_1 * num_heads + cta_rank * 64 + head_row_base_1; #pragma unroll for (int acc_stage_1 = 0; acc_stage_1 < 2; acc_stage_1++) { #pragma unroll for (int local_slice_1 = 0; local_slice_1 < 2; local_slice_1++) { int logical_slice = acc_stage_1 * 2 * 2 + n_half_1 * 2 + local_slice_1; + int epi_box = epi_box_base + (acc_stage_1 * 2 + local_slice_1) * 4096; int o_base_1 = taddr + 128 + (unsigned int)(acc_stage_1 * 128) + (unsigned int)(local_slice_1 * 64) + (unsigned int)(output_row_base << 16); #pragma unroll for (int c_1 = 0; c_1 < 64; c_1 += 32) { @@ -1710,46 +1726,35 @@ kernel_cake_dsv4_4c41fb8347f137d4d46f(const __grid_constant__ CUtensorMap tmap_q " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]) : "r"(o_base_1 + c_1)); - int gmem_base = output_offset + logical_slice * 64 + c_1; + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + float epi_scaled[32]; #pragma unroll - for (int j = 0; j < 32; j += 16) { - { - const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; - #if __CUDA_ARCH__ >= 1000 - #pragma unroll - for (int _ps = 0; _ps < 8; _ps++) - mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_3[j])[_ps], _prescale2_1); - #else - #pragma unroll - for (int _ps = 0; _ps < 16; _ps++) - _tmem_load_3[j + _ps] *= inv_sum * output_scale; - #endif - { - __nv_bfloat162 _pk0 = __floats2bfloat162_rn(_tmem_load_3[j + 0], _tmem_load_3[j + 1]); - unsigned _pk_u0 = *reinterpret_cast(&_pk0); - __nv_bfloat162 _pk1 = __floats2bfloat162_rn(_tmem_load_3[j + 2], _tmem_load_3[j + 3]); - unsigned _pk_u1 = *reinterpret_cast(&_pk1); - __nv_bfloat162 _pk2 = __floats2bfloat162_rn(_tmem_load_3[j + 4], _tmem_load_3[j + 5]); - unsigned _pk_u2 = *reinterpret_cast(&_pk2); - __nv_bfloat162 _pk3 = __floats2bfloat162_rn(_tmem_load_3[j + 6], _tmem_load_3[j + 7]); - unsigned _pk_u3 = *reinterpret_cast(&_pk3); - __nv_bfloat162 _pk4 = __floats2bfloat162_rn(_tmem_load_3[j + 8], _tmem_load_3[j + 9]); - unsigned _pk_u4 = *reinterpret_cast(&_pk4); - __nv_bfloat162 _pk5 = __floats2bfloat162_rn(_tmem_load_3[j + 10], _tmem_load_3[j + 11]); - unsigned _pk_u5 = *reinterpret_cast(&_pk5); - __nv_bfloat162 _pk6 = __floats2bfloat162_rn(_tmem_load_3[j + 12], _tmem_load_3[j + 13]); - unsigned _pk_u6 = *reinterpret_cast(&_pk6); - __nv_bfloat162 _pk7 = __floats2bfloat162_rn(_tmem_load_3[j + 14], _tmem_load_3[j + 15]); - unsigned _pk_u7 = *reinterpret_cast(&_pk7); - asm volatile( - "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" - :: "l"((void*)(&((__nv_bfloat16*)(O + (gmem_base + j)))[0])), "r"(_pk_u0), "r"(_pk_u1), "r"(_pk_u2), "r"(_pk_u3), "r"(_pk_u4), "r"(_pk_u5), "r"(_pk_u6), "r"(_pk_u7) : "memory"); - } - } + for (int e = 0; e < 32; e++) { + epi_scaled[e] = _tmem_load_3[e] * epi_scale; + } + unsigned int epi_packed[16]; + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(epi_scaled[_lp*2 + 0], epi_scaled[_lp*2+1 + 0])); + epi_packed[_lp] = *(uint32_t*)&_bf2; } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((epi_box + (lane * 128 + c_1 * 2 ^ (lane * 128 + c_1 * 2 >> 7 & 7) << 4))), "r"(epi_packed[0]), "r"(epi_packed[1]), "r"(epi_packed[2]), "r"(epi_packed[3]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((epi_box + (lane * 128 + (c_1 * 2 + 16) ^ (lane * 128 + (c_1 * 2 + 16) >> 7 & 7) << 4))), "r"(epi_packed[4]), "r"(epi_packed[5]), "r"(epi_packed[6]), "r"(epi_packed[7]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((epi_box + (lane * 128 + (c_1 * 2 + 32) ^ (lane * 128 + (c_1 * 2 + 32) >> 7 & 7) << 4))), "r"(epi_packed[8]), "r"(epi_packed[9]), "r"(epi_packed[10]), "r"(epi_packed[11]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((epi_box + (lane * 128 + (c_1 * 2 + 48) ^ (lane * 128 + (c_1 * 2 + 48) >> 7 & 7) << 4))), "r"(epi_packed[12]), "r"(epi_packed[13]), "r"(epi_packed[14]), "r"(epi_packed[15]) : "memory"); + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + __syncwarp(); + if (elect_sync()) { + tma_store_4d((&O), 0, logical_slice, split_idx_1, epi_row0, epi_box); + asm volatile("cp.async.bulk.commit_group;"); } } } + if (elect_sync()) { + asm volatile("cp.async.bulk.wait_group.read 0;"); + mbarrier_arrive(epi_q_free_addr); + } asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); asm volatile( "{\n\t" @@ -1760,6 +1765,9 @@ kernel_cake_dsv4_4c41fb8347f137d4d46f(const __grid_constant__ CUtensorMap tmap_q :: "r"(o_empty_addr), "r"(0) : "memory"); tile_cursor_1 = tile_cursor_1 + num_kv_tiles_1; } + if (elect_sync()) { + asm volatile("cp.async.bulk.wait_group 0;"); + } asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); asm volatile("tcgen05.fence::before_thread_sync;"); } @@ -2106,6 +2114,7 @@ kernel_cake_dsv4_4c41fb8347f137d4d46f(const __grid_constant__ CUtensorMap tmap_q unsigned int k_stage_1 = 0; unsigned int v_stage_1 = 0; unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_epi_q_free_0 = 1; unsigned int _phase_k_empty = 1; unsigned int _phase_v_empty = 1; #pragma unroll 1 @@ -2116,6 +2125,8 @@ kernel_cake_dsv4_4c41fb8347f137d4d46f(const __grid_constant__ CUtensorMap tmap_q int first_kv_tile_1 = split_idx_2 * num_kv_tiles_3; mbarrier_wait(q_empty_addr, _phase_q_empty_0); _phase_q_empty_0 ^= 1; + mbarrier_wait(epi_q_free_addr, _phase_epi_q_free_0); + _phase_epi_q_free_0 ^= 1; if (load_warp_rank == 0) { if (cta_rank == 0) { if (elect_sync()) { diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_5e5d6648b7bd42bcd0d6_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_0a925b9bcb3044777c76_binding.cu similarity index 97% rename from csrc/cake_dsv4/sm_100a/cake_dsv4_5e5d6648b7bd42bcd0d6_binding.cu rename to csrc/cake_dsv4/sm_103a/cake_dsv4_0a925b9bcb3044777c76_binding.cu index b4398bf2025..f5e9697f131 100644 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_5e5d6648b7bd42bcd0d6_binding.cu +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_0a925b9bcb3044777c76_binding.cu @@ -15,7 +15,7 @@ */ // Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_5e5d6648b7bd42bcd0d6'. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_0a925b9bcb3044777c76'. #include #include #include @@ -31,10 +31,10 @@ #include #include -extern "C" __global__ void kernel_cake_dsv4_5e5d6648b7bd42bcd0d6(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ cum_seq_lens_q, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); +extern "C" __global__ void kernel_cake_dsv4_0a925b9bcb3044777c76(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ cum_seq_lens_q, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); -namespace cake_host_shim_1d4218665228ec14 { +namespace cake_host_shim_52ed16aa7da8966c { using tvm::ffi::Optional; using tvm::ffi::TensorView; @@ -409,7 +409,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ } if (smem_status == cudaSuccess) { smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_5e5d6648b7bd42bcd0d6), + reinterpret_cast(kernel_cake_dsv4_0a925b9bcb3044777c76), cudaFuncAttributeMaxDynamicSharedMemorySize, 203136); } @@ -419,7 +419,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ } } TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_5e5d6648b7bd42bcd0d6 failed: " + << "cudaFuncSetAttribute for kernel_cake_dsv4_0a925b9bcb3044777c76 failed: " << cudaGetErrorString(smem_status); dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); dim3 block(512u, 1u, 1u); @@ -430,12 +430,12 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ config.dynamicSmemBytes = 203136u; config.stream = stream; cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_5e5d6648b7bd42bcd0d6), kargs); + &config, reinterpret_cast(kernel_cake_dsv4_0a925b9bcb3044777c76), kargs); TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_5e5d6648b7bd42bcd0d6 failed: " + << "cudaLaunchKernelExC for kernel_cake_dsv4_0a925b9bcb3044777c76 failed: " << cudaGetErrorString(launch_status); } -} // namespace cake_host_shim_1d4218665228ec14 +} // namespace cake_host_shim_52ed16aa7da8966c -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_1d4218665228ec14::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_52ed16aa7da8966c::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_f55f69afe4ac3ec96817_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_0a925b9bcb3044777c76_kernel.cu similarity index 99% rename from csrc/cake_dsv4/sm_103a/cake_dsv4_f55f69afe4ac3ec96817_kernel.cu rename to csrc/cake_dsv4/sm_103a/cake_dsv4_0a925b9bcb3044777c76_kernel.cu index 53eb116e5d9..6b6f53ff26e 100644 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_f55f69afe4ac3ec96817_kernel.cu +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_0a925b9bcb3044777c76_kernel.cu @@ -365,6 +365,19 @@ __device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { } +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + __device__ __forceinline__ float approx_exp2(float x) { float y; asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); @@ -1062,35 +1075,12 @@ __device__ __forceinline__ void tcgen05_commit(int mbar_addr) { } -__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x16.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) - : "r"(tmem_addr)); -} - - __device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { tmem_ld_x16(dst, addr); asm volatile("tcgen05.wait::ld.sync.aligned;"); } -__device__ __forceinline__ void tmem_ld_x8(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x8.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]) - : "r"(tmem_addr)); -} - - __device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { uint32_t result; asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" @@ -1132,7 +1122,7 @@ cake_pipeline_state_at(CakePipelineState state) { extern "C" { __global__ __launch_bounds__(512, 1) void -kernel_cake_dsv4_f55f69afe4ac3ec96817(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ cum_seq_lens_q, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items) +kernel_cake_dsv4_0a925b9bcb3044777c76(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ cum_seq_lens_q, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); @@ -1925,25 +1915,30 @@ kernel_cake_dsv4_f55f69afe4ac3ec96817(const __grid_constant__ CUtensorMap tmap_q int row_bank_1 = ((output_half_1 != 0) ? 1048576 : 0); int col_base = col_half_1 * 128; #pragma unroll 2 - for (int col_1 = 0; col_1 < 128; col_1 += 8) { - float _tmem_load_4[8]; + for (int wcol = 0; wcol < 128; wcol += 16) { + float _tmem_load_4[16]; asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x8.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8], 128;" - : "=r"(*reinterpret_cast(&_tmem_load_4[0])), "=r"(*reinterpret_cast(&_tmem_load_4[1])), "=r"(*reinterpret_cast(&_tmem_load_4[2])), "=r"(*reinterpret_cast(&_tmem_load_4[3])), "=r"(*reinterpret_cast(&_tmem_load_4[4])), "=r"(*reinterpret_cast(&_tmem_load_4[5])), "=r"(*reinterpret_cast(&_tmem_load_4[6])), "=r"(*reinterpret_cast(&_tmem_load_4[7])) - : "r"(taddr + 256 + (unsigned int)row_bank_1 + (unsigned int)col_1)); + "tcgen05.ld.sync.aligned.16x32bx2.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16], 128;" + : "=r"(*reinterpret_cast(&_tmem_load_4[0])), "=r"(*reinterpret_cast(&_tmem_load_4[1])), "=r"(*reinterpret_cast(&_tmem_load_4[2])), "=r"(*reinterpret_cast(&_tmem_load_4[3])), "=r"(*reinterpret_cast(&_tmem_load_4[4])), "=r"(*reinterpret_cast(&_tmem_load_4[5])), "=r"(*reinterpret_cast(&_tmem_load_4[6])), "=r"(*reinterpret_cast(&_tmem_load_4[7])), "=r"(*reinterpret_cast(&_tmem_load_4[8])), "=r"(*reinterpret_cast(&_tmem_load_4[9])), "=r"(*reinterpret_cast(&_tmem_load_4[10])), "=r"(*reinterpret_cast(&_tmem_load_4[11])), "=r"(*reinterpret_cast(&_tmem_load_4[12])), "=r"(*reinterpret_cast(&_tmem_load_4[13])), "=r"(*reinterpret_cast(&_tmem_load_4[14])), "=r"(*reinterpret_cast(&_tmem_load_4[15])) + : "r"(taddr + 256 + (unsigned int)row_bank_1 + (unsigned int)wcol)); const float2 _scale2_1 = {inv_sum, inv_sum}; #pragma unroll - for (int _ls = 0; _ls < 4; _ls++) + for (int _ls = 0; _ls < 8; _ls++) mul_f32x2_inplace(&reinterpret_cast(_tmem_load_4)[_ls], _scale2_1); - uint32_t _tmem_load_4_bf16[4]; + uint32_t _tmem_load_4_bf16[8]; #pragma unroll - for (int _lp = 0; _lp < 4; _lp++) { + for (int _lp = 0; _lp < 8; _lp++) { __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_4[_lp*2 + 0], _tmem_load_4[_lp*2+1 + 0])); _tmem_load_4_bf16[_lp] = *(uint32_t*)&_bf2; } if (output_in_bounds != 0) { - reinterpret_cast(O + (output_base + output_half_1 * 256 + col_base + col_1))[0] = reinterpret_cast(_tmem_load_4_bf16)[0]; + { + const unsigned* _raw_stv8_2 = reinterpret_cast(_tmem_load_4_bf16); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(O + (output_base + output_half_1 * 256 + col_base + wcol))), "r"(_raw_stv8_2[0]), "r"(_raw_stv8_2[1]), "r"(_raw_stv8_2[2]), "r"(_raw_stv8_2[3]), "r"(_raw_stv8_2[4]), "r"(_raw_stv8_2[5]), "r"(_raw_stv8_2[6]), "r"(_raw_stv8_2[7]) : "memory"); + } } } } diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_f89bc6561b38462e9538_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_1bad879f597b2e2adf3e_binding.cu similarity index 97% rename from csrc/cake_dsv4/sm_103a/cake_dsv4_f89bc6561b38462e9538_binding.cu rename to csrc/cake_dsv4/sm_103a/cake_dsv4_1bad879f597b2e2adf3e_binding.cu index 02f9f0f2ff8..8ef2566dad1 100644 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_f89bc6561b38462e9538_binding.cu +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_1bad879f597b2e2adf3e_binding.cu @@ -15,7 +15,7 @@ */ // Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_f89bc6561b38462e9538'. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_1bad879f597b2e2adf3e'. #include #include #include @@ -31,10 +31,10 @@ #include #include -extern "C" __global__ void kernel_cake_dsv4_f89bc6561b38462e9538(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int num_head_tiles, int has_sinks); +extern "C" __global__ void kernel_cake_dsv4_1bad879f597b2e2adf3e(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int num_head_tiles, int has_sinks); -namespace cake_host_shim_13f3ad5a6154b3a4 { +namespace cake_host_shim_406af89a732a5e2c { using tvm::ffi::Optional; using tvm::ffi::TensorView; @@ -357,7 +357,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_O, Te } if (smem_status == cudaSuccess) { smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_f89bc6561b38462e9538), + reinterpret_cast(kernel_cake_dsv4_1bad879f597b2e2adf3e), cudaFuncAttributeMaxDynamicSharedMemorySize, 198528); } @@ -367,7 +367,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_O, Te } } TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_f89bc6561b38462e9538 failed: " + << "cudaFuncSetAttribute for kernel_cake_dsv4_1bad879f597b2e2adf3e failed: " << cudaGetErrorString(smem_status); dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); dim3 block(512u, 1u, 1u); @@ -378,12 +378,12 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_O, Te config.dynamicSmemBytes = 198528u; config.stream = stream; cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_f89bc6561b38462e9538), kargs); + &config, reinterpret_cast(kernel_cake_dsv4_1bad879f597b2e2adf3e), kargs); TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_f89bc6561b38462e9538 failed: " + << "cudaLaunchKernelExC for kernel_cake_dsv4_1bad879f597b2e2adf3e failed: " << cudaGetErrorString(launch_status); } -} // namespace cake_host_shim_13f3ad5a6154b3a4 +} // namespace cake_host_shim_406af89a732a5e2c -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_13f3ad5a6154b3a4::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_406af89a732a5e2c::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_f89bc6561b38462e9538_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_1bad879f597b2e2adf3e_kernel.cu similarity index 89% rename from csrc/cake_dsv4/sm_103a/cake_dsv4_f89bc6561b38462e9538_kernel.cu rename to csrc/cake_dsv4/sm_103a/cake_dsv4_1bad879f597b2e2adf3e_kernel.cu index d9d9646824e..0a4d274ffc7 100644 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_f89bc6561b38462e9538_kernel.cu +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_1bad879f597b2e2adf3e_kernel.cu @@ -999,7 +999,7 @@ __device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { extern "C" { __global__ __launch_bounds__(512, 1) void -kernel_cake_dsv4_f89bc6561b38462e9538(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int num_head_tiles, int has_sinks) +kernel_cake_dsv4_1bad879f597b2e2adf3e(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int num_head_tiles, int has_sinks) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); @@ -1369,29 +1369,129 @@ kernel_cake_dsv4_f89bc6561b38462e9538(const __grid_constant__ CUtensorMap tmap_q " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" : "=r"(*reinterpret_cast(&_tmem_load_1[32])), "=r"(*reinterpret_cast(&_tmem_load_1[33])), "=r"(*reinterpret_cast(&_tmem_load_1[34])), "=r"(*reinterpret_cast(&_tmem_load_1[35])), "=r"(*reinterpret_cast(&_tmem_load_1[36])), "=r"(*reinterpret_cast(&_tmem_load_1[37])), "=r"(*reinterpret_cast(&_tmem_load_1[38])), "=r"(*reinterpret_cast(&_tmem_load_1[39])), "=r"(*reinterpret_cast(&_tmem_load_1[40])), "=r"(*reinterpret_cast(&_tmem_load_1[41])), "=r"(*reinterpret_cast(&_tmem_load_1[42])), "=r"(*reinterpret_cast(&_tmem_load_1[43])), "=r"(*reinterpret_cast(&_tmem_load_1[44])), "=r"(*reinterpret_cast(&_tmem_load_1[45])), "=r"(*reinterpret_cast(&_tmem_load_1[46])), "=r"(*reinterpret_cast(&_tmem_load_1[47])), "=r"(*reinterpret_cast(&_tmem_load_1[48])), "=r"(*reinterpret_cast(&_tmem_load_1[49])), "=r"(*reinterpret_cast(&_tmem_load_1[50])), "=r"(*reinterpret_cast(&_tmem_load_1[51])), "=r"(*reinterpret_cast(&_tmem_load_1[52])), "=r"(*reinterpret_cast(&_tmem_load_1[53])), "=r"(*reinterpret_cast(&_tmem_load_1[54])), "=r"(*reinterpret_cast(&_tmem_load_1[55])), "=r"(*reinterpret_cast(&_tmem_load_1[56])), "=r"(*reinterpret_cast(&_tmem_load_1[57])), "=r"(*reinterpret_cast(&_tmem_load_1[58])), "=r"(*reinterpret_cast(&_tmem_load_1[59])), "=r"(*reinterpret_cast(&_tmem_load_1[60])), "=r"(*reinterpret_cast(&_tmem_load_1[61])), "=r"(*reinterpret_cast(&_tmem_load_1[62])), "=r"(*reinterpret_cast(&_tmem_load_1[63])) : "r"(taddr + 256 + (unsigned int)row_addr + 32)); - int col_base = col_half_1 * 64; - if (head_idx_1 < num_heads) { - #pragma unroll - for (int offset = 0; offset < 64; offset += 8) { - { - const float2 _prescale2_0 = {inv_sum * output_scale, inv_sum * output_scale}; - #if __CUDA_ARCH__ >= 1000 - #pragma unroll - for (int _ps = 0; _ps < 4; _ps++) - mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_1[offset])[_ps], _prescale2_0); - #else - #pragma unroll - for (int _ps = 0; _ps < 8; _ps++) - _tmem_load_1[offset + _ps] *= inv_sum * output_scale; - #endif - __nv_bfloat162 _pk[4]; - _pk[0] = __floats2bfloat162_rn(_tmem_load_1[offset + 0], _tmem_load_1[offset + 1]); - _pk[1] = __floats2bfloat162_rn(_tmem_load_1[offset + 2], _tmem_load_1[offset + 3]); - _pk[2] = __floats2bfloat162_rn(_tmem_load_1[offset + 4], _tmem_load_1[offset + 5]); - _pk[3] = __floats2bfloat162_rn(_tmem_load_1[offset + 6], _tmem_load_1[offset + 7]); - *reinterpret_cast(&((__nv_bfloat16*)(O + (output_base + col_base + offset)))[0]) = *reinterpret_cast(&_pk[0]); - } - } + float2 _f2_14 = make_float2(inv_sum * output_scale, inv_sum * output_scale); + #pragma unroll + for (int i_1 = 0; i_1 < 64; i_1 += 2) { + float2 _f2_15 = make_float2(_tmem_load_1[i_1], _tmem_load_1[i_1 + 1]); + float2 _mul_f32x2_0; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_0) : "l"(*(const unsigned long long*)&_f2_15), "l"(*(const unsigned long long*)&_f2_14)); + _tmem_load_1[i_1] = _mul_f32x2_0.x; + _tmem_load_1[i_1 + 1] = _mul_f32x2_0.y; + } + unsigned int packed_o[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_1[_lp*2 + 0], _tmem_load_1[_lp*2+1 + 0])); + packed_o[_lp] = *(uint32_t*)&_bf2; + } + bool stage_hi_1 = (lane & 1) != 0; + bool stage_hi_2 = (lane & 2) != 0; + bool stage_hi_4 = (lane & 4) != 0; + unsigned int _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[0] : packed_o[4]), 1); + packed_o[0] = ((stage_hi_1) ? _shfl_xor_2 : packed_o[0]); + packed_o[4] = ((stage_hi_1) ? packed_o[4] : _shfl_xor_2); + unsigned int _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[1] : packed_o[5]), 1); + packed_o[1] = ((stage_hi_1) ? _shfl_xor_3 : packed_o[1]); + packed_o[5] = ((stage_hi_1) ? packed_o[5] : _shfl_xor_3); + unsigned int _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[2] : packed_o[6]), 1); + packed_o[2] = ((stage_hi_1) ? _shfl_xor_4 : packed_o[2]); + packed_o[6] = ((stage_hi_1) ? packed_o[6] : _shfl_xor_4); + unsigned int _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[3] : packed_o[7]), 1); + packed_o[3] = ((stage_hi_1) ? _shfl_xor_5 : packed_o[3]); + packed_o[7] = ((stage_hi_1) ? packed_o[7] : _shfl_xor_5); + unsigned int _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[8] : packed_o[12]), 1); + packed_o[8] = ((stage_hi_1) ? _shfl_xor_6 : packed_o[8]); + packed_o[12] = ((stage_hi_1) ? packed_o[12] : _shfl_xor_6); + unsigned int _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[9] : packed_o[13]), 1); + packed_o[9] = ((stage_hi_1) ? _shfl_xor_7 : packed_o[9]); + packed_o[13] = ((stage_hi_1) ? packed_o[13] : _shfl_xor_7); + unsigned int _shfl_xor_8 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[10] : packed_o[14]), 1); + packed_o[10] = ((stage_hi_1) ? _shfl_xor_8 : packed_o[10]); + packed_o[14] = ((stage_hi_1) ? packed_o[14] : _shfl_xor_8); + unsigned int _shfl_xor_9 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[11] : packed_o[15]), 1); + packed_o[11] = ((stage_hi_1) ? _shfl_xor_9 : packed_o[11]); + packed_o[15] = ((stage_hi_1) ? packed_o[15] : _shfl_xor_9); + unsigned int _shfl_xor_10 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[16] : packed_o[20]), 1); + packed_o[16] = ((stage_hi_1) ? _shfl_xor_10 : packed_o[16]); + packed_o[20] = ((stage_hi_1) ? packed_o[20] : _shfl_xor_10); + unsigned int _shfl_xor_11 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[17] : packed_o[21]), 1); + packed_o[17] = ((stage_hi_1) ? _shfl_xor_11 : packed_o[17]); + packed_o[21] = ((stage_hi_1) ? packed_o[21] : _shfl_xor_11); + unsigned int _shfl_xor_12 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[18] : packed_o[22]), 1); + packed_o[18] = ((stage_hi_1) ? _shfl_xor_12 : packed_o[18]); + packed_o[22] = ((stage_hi_1) ? packed_o[22] : _shfl_xor_12); + unsigned int _shfl_xor_13 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[19] : packed_o[23]), 1); + packed_o[19] = ((stage_hi_1) ? _shfl_xor_13 : packed_o[19]); + packed_o[23] = ((stage_hi_1) ? packed_o[23] : _shfl_xor_13); + unsigned int _shfl_xor_14 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[24] : packed_o[28]), 1); + packed_o[24] = ((stage_hi_1) ? _shfl_xor_14 : packed_o[24]); + packed_o[28] = ((stage_hi_1) ? packed_o[28] : _shfl_xor_14); + unsigned int _shfl_xor_15 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[25] : packed_o[29]), 1); + packed_o[25] = ((stage_hi_1) ? _shfl_xor_15 : packed_o[25]); + packed_o[29] = ((stage_hi_1) ? packed_o[29] : _shfl_xor_15); + unsigned int _shfl_xor_16 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[26] : packed_o[30]), 1); + packed_o[26] = ((stage_hi_1) ? _shfl_xor_16 : packed_o[26]); + packed_o[30] = ((stage_hi_1) ? packed_o[30] : _shfl_xor_16); + unsigned int _shfl_xor_17 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[27] : packed_o[31]), 1); + packed_o[27] = ((stage_hi_1) ? _shfl_xor_17 : packed_o[27]); + packed_o[31] = ((stage_hi_1) ? packed_o[31] : _shfl_xor_17); + int lane_t = lane & 1; + unsigned int chunk_words[4]; + chunk_words[0] = packed_o[0]; + chunk_words[1] = packed_o[1]; + chunk_words[2] = packed_o[2]; + chunk_words[3] = packed_o[3]; + if (head_tile_1 * 64 + (my_row_1 & -2 | 0) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + head_tile_1 * 64 + (my_row_1 & -2 | 0)) * 512 + v_chunk * 128 + col_half_1 * 64 + (0 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[4]; + chunk_words[1] = packed_o[5]; + chunk_words[2] = packed_o[6]; + chunk_words[3] = packed_o[7]; + if (head_tile_1 * 64 + (my_row_1 & -2 | 1) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + head_tile_1 * 64 + (my_row_1 & -2 | 1)) * 512 + v_chunk * 128 + col_half_1 * 64 + (0 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[8]; + chunk_words[1] = packed_o[9]; + chunk_words[2] = packed_o[10]; + chunk_words[3] = packed_o[11]; + if (head_tile_1 * 64 + (my_row_1 & -2 | 0) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + head_tile_1 * 64 + (my_row_1 & -2 | 0)) * 512 + v_chunk * 128 + col_half_1 * 64 + (2 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[12]; + chunk_words[1] = packed_o[13]; + chunk_words[2] = packed_o[14]; + chunk_words[3] = packed_o[15]; + if (head_tile_1 * 64 + (my_row_1 & -2 | 1) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + head_tile_1 * 64 + (my_row_1 & -2 | 1)) * 512 + v_chunk * 128 + col_half_1 * 64 + (2 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[16]; + chunk_words[1] = packed_o[17]; + chunk_words[2] = packed_o[18]; + chunk_words[3] = packed_o[19]; + if (head_tile_1 * 64 + (my_row_1 & -2 | 0) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + head_tile_1 * 64 + (my_row_1 & -2 | 0)) * 512 + v_chunk * 128 + col_half_1 * 64 + (4 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[20]; + chunk_words[1] = packed_o[21]; + chunk_words[2] = packed_o[22]; + chunk_words[3] = packed_o[23]; + if (head_tile_1 * 64 + (my_row_1 & -2 | 1) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + head_tile_1 * 64 + (my_row_1 & -2 | 1)) * 512 + v_chunk * 128 + col_half_1 * 64 + (4 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[24]; + chunk_words[1] = packed_o[25]; + chunk_words[2] = packed_o[26]; + chunk_words[3] = packed_o[27]; + if (head_tile_1 * 64 + (my_row_1 & -2 | 0) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + head_tile_1 * 64 + (my_row_1 & -2 | 0)) * 512 + v_chunk * 128 + col_half_1 * 64 + (6 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[28]; + chunk_words[1] = packed_o[29]; + chunk_words[2] = packed_o[30]; + chunk_words[3] = packed_o[31]; + if (head_tile_1 * 64 + (my_row_1 & -2 | 1) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + head_tile_1 * 64 + (my_row_1 & -2 | 1)) * 512 + v_chunk * 128 + col_half_1 * 64 + (6 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; } asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); asm volatile("tcgen05.fence::before_thread_sync;"); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_1e3cf4361de31c58b69b_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_1e3cf4361de31c58b69b_binding.cu new file mode 100644 index 00000000000..28868ff1e78 --- /dev/null +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_1e3cf4361de31c58b69b_binding.cu @@ -0,0 +1,436 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Generated by the Cake source exporter — do not edit. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_1e3cf4361de31c58b69b'. +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#include +#include +#include +#include +#include +#include +#include + +extern "C" __global__ void kernel_cake_dsv4_1e3cf4361de31c58b69b(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int num_query_tokens, int has_sinks); + + +namespace cake_host_shim_4f4d09577aaed6ed { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t CakeDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +// 4D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + int64_t d3 = t.size(t.ndim() - 3); + TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 64u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_compressed_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_compressed_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_compressed_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_compressed_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_compressed_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_compressed_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_compressed_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_kv') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_compressed_kv, TensorView arg_O, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_num_query_tokens, int64_t arg_has_sinks, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckContiguous(arg_tmap_q, "tmap_q"); + CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); + CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 4, 16, 1); + CheckCudaTensor(arg_tmap_compressed_kv, "tmap_compressed_kv"); + CheckDtype(arg_tmap_compressed_kv, "tmap_compressed_kv", 4, 16, 1); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_swa_indices, "swa_indices"); + CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); + CheckContiguous(arg_swa_indices, "swa_indices"); + CheckCudaTensor(arg_compressed_indices, "compressed_indices"); + CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); + CheckContiguous(arg_compressed_indices, "compressed_indices"); + CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); + CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); + CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); + CheckContiguous(arg_bmm1_scale, "bmm1_scale"); + CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); + CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); + CheckContiguous(arg_bmm2_scale, "bmm2_scale"); + TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) + << "scalar 'num_heads' value " << arg_num_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) + << "scalar 'swa_index_stride' value " << arg_swa_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) + << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) + << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) + << "scalar 'sparse_topk' value " << arg_sparse_topk + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_query_tokens >= -2147483648LL && arg_num_query_tokens <= 2147483647LL, ValueError) + << "scalar 'num_query_tokens' value " << arg_num_query_tokens + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) + << "scalar 'has_sinks' value " << arg_has_sinks + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); + CheckSameCudaDevice(arg_tmap_compressed_kv, arg_tmap_q, "tmap_compressed_kv", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); + CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); + CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); + CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + + cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); + CUtensorMap p_tmap_compressed_kv = EncodeTma_tmap_compressed_kv(arg_tmap_compressed_kv); + __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); + int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); + int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); + int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); + float* p_sinks = static_cast(arg_sinks.data_ptr()); + float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); + float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); + int32_t v_num_heads = (int32_t)arg_num_heads; + int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; + int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; + int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; + int32_t v_sparse_topk = (int32_t)arg_sparse_topk; + int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens; + int32_t v_has_sinks = (int32_t)arg_has_sinks; + void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_tmap_compressed_kv, &p_O, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_num_query_tokens, &v_has_sinks}; + + static std::mutex smem_mu; + static auto* smem_status_by_device = + new std::unordered_map(); + cudaError_t smem_status = cudaSuccess; + { + std::lock_guard lock(smem_mu); + auto it = smem_status_by_device->find(dev.device_id); + if (it == smem_status_by_device->end()) { + int current_device = -1; + smem_status = cudaGetDevice(¤t_device); + if (smem_status == cudaSuccess && current_device != dev.device_id) { + smem_status = cudaErrorInvalidDevice; + } + if (smem_status == cudaSuccess) { + smem_status = cudaFuncSetAttribute( + reinterpret_cast(kernel_cake_dsv4_1e3cf4361de31c58b69b), + cudaFuncAttributeMaxDynamicSharedMemorySize, + 203904); + } + smem_status_by_device->emplace(dev.device_id, smem_status); + } else { + smem_status = it->second; + } + } + TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) + << "cudaFuncSetAttribute for kernel_cake_dsv4_1e3cf4361de31c58b69b failed: " + << cudaGetErrorString(smem_status); + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(512u, 1u, 1u); + + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 203904u; + config.stream = stream; + cudaError_t launch_status = cudaLaunchKernelExC( + &config, reinterpret_cast(kernel_cake_dsv4_1e3cf4361de31c58b69b), kargs); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernelExC for kernel_cake_dsv4_1e3cf4361de31c58b69b failed: " + << cudaGetErrorString(launch_status); +} + +} // namespace cake_host_shim_4f4d09577aaed6ed + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_4f4d09577aaed6ed::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_3e5a644d61420c9c7ea1_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_1e3cf4361de31c58b69b_kernel.cu similarity index 96% rename from csrc/cake_dsv4/sm_103a/cake_dsv4_3e5a644d61420c9c7ea1_kernel.cu rename to csrc/cake_dsv4/sm_103a/cake_dsv4_1e3cf4361de31c58b69b_kernel.cu index d4ab4dbfb83..8aebdd9da7c 100644 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_3e5a644d61420c9c7ea1_kernel.cu +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_1e3cf4361de31c58b69b_kernel.cu @@ -30,8 +30,6 @@ struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); -template -struct __align__(128) CakeTensorMapPack { CakeTensorMap maps[N]; }; #if defined(__CUDACC_RTC__) typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; @@ -1084,7 +1082,7 @@ __device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { extern "C" { __global__ __launch_bounds__(512, 1) void -kernel_cake_dsv4_3e5a644d61420c9c7ea1(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int num_query_tokens, int has_sinks) +kernel_cake_dsv4_1e3cf4361de31c58b69b(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int num_query_tokens, int has_sinks) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); @@ -1116,13 +1114,6 @@ kernel_cake_dsv4_3e5a644d61420c9c7ea1(CakeTensorMap const* tmap_q, CakeTensorMap const int num_bids = gridDim.x; const int cta_rank = 0; - if (tid == 0) { - asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_q)) : "memory"); - asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_swa_kv)) : "memory"); - asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_compressed_kv)) : "memory"); - } - __syncthreads(); - // Kernel setup ops __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); @@ -1648,24 +1639,39 @@ kernel_cake_dsv4_3e5a644d61420c9c7ea1(CakeTensorMap const* tmap_q, CakeTensorMap int col_base = v_stage_1 * 128 + col_half_1 * 64; if (my_row_1 < num_heads) { #pragma unroll - for (int offset = 0; offset < 64; offset += 8) { + for (int offset = 0; offset < 64; offset += 16) { { const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; #if __CUDA_ARCH__ >= 1000 #pragma unroll - for (int _ps = 0; _ps < 4; _ps++) + for (int _ps = 0; _ps < 8; _ps++) mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_4[offset])[_ps], _prescale2_1); #else #pragma unroll - for (int _ps = 0; _ps < 8; _ps++) + for (int _ps = 0; _ps < 16; _ps++) _tmem_load_4[offset + _ps] *= inv_sum * output_scale; #endif - __nv_bfloat162 _pk[4]; - _pk[0] = __floats2bfloat162_rn(_tmem_load_4[offset + 0], _tmem_load_4[offset + 1]); - _pk[1] = __floats2bfloat162_rn(_tmem_load_4[offset + 2], _tmem_load_4[offset + 3]); - _pk[2] = __floats2bfloat162_rn(_tmem_load_4[offset + 4], _tmem_load_4[offset + 5]); - _pk[3] = __floats2bfloat162_rn(_tmem_load_4[offset + 6], _tmem_load_4[offset + 7]); - *reinterpret_cast(&((__nv_bfloat16*)(O + (output_base + col_base + offset)))[0]) = *reinterpret_cast(&_pk[0]); + { + __nv_bfloat162 _pk0 = __floats2bfloat162_rn(_tmem_load_4[offset + 0], _tmem_load_4[offset + 1]); + unsigned _pk_u0 = *reinterpret_cast(&_pk0); + __nv_bfloat162 _pk1 = __floats2bfloat162_rn(_tmem_load_4[offset + 2], _tmem_load_4[offset + 3]); + unsigned _pk_u1 = *reinterpret_cast(&_pk1); + __nv_bfloat162 _pk2 = __floats2bfloat162_rn(_tmem_load_4[offset + 4], _tmem_load_4[offset + 5]); + unsigned _pk_u2 = *reinterpret_cast(&_pk2); + __nv_bfloat162 _pk3 = __floats2bfloat162_rn(_tmem_load_4[offset + 6], _tmem_load_4[offset + 7]); + unsigned _pk_u3 = *reinterpret_cast(&_pk3); + __nv_bfloat162 _pk4 = __floats2bfloat162_rn(_tmem_load_4[offset + 8], _tmem_load_4[offset + 9]); + unsigned _pk_u4 = *reinterpret_cast(&_pk4); + __nv_bfloat162 _pk5 = __floats2bfloat162_rn(_tmem_load_4[offset + 10], _tmem_load_4[offset + 11]); + unsigned _pk_u5 = *reinterpret_cast(&_pk5); + __nv_bfloat162 _pk6 = __floats2bfloat162_rn(_tmem_load_4[offset + 12], _tmem_load_4[offset + 13]); + unsigned _pk_u6 = *reinterpret_cast(&_pk6); + __nv_bfloat162 _pk7 = __floats2bfloat162_rn(_tmem_load_4[offset + 14], _tmem_load_4[offset + 15]); + unsigned _pk_u7 = *reinterpret_cast(&_pk7); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(&((__nv_bfloat16*)(O + (output_base + col_base + offset)))[0])), "r"(_pk_u0), "r"(_pk_u1), "r"(_pk_u2), "r"(_pk_u3), "r"(_pk_u4), "r"(_pk_u5), "r"(_pk_u6), "r"(_pk_u7) : "memory"); + } } } } @@ -2252,7 +2258,7 @@ kernel_cake_dsv4_3e5a644d61420c9c7ea1(CakeTensorMap const* tmap_q, CakeTensorMap mbarrier_arrive_expect_tx(q_full_addr, 65536); #pragma unroll for (int q_subload = 0; q_subload < 8; q_subload++) { - tma_4d_gmem2smem(smem_q_addr + (unsigned int)(q_subload * 8192), tmap_q, 0, 0, q_subload, query_idx_6, q_full_addr); + tma_4d_gmem2smem(smem_q_addr + (unsigned int)(q_subload * 8192), (&tmap_q), 0, 0, q_subload, query_idx_6, q_full_addr); } } } @@ -2284,8 +2290,8 @@ kernel_cake_dsv4_3e5a644d61420c9c7ea1(CakeTensorMap const* tmap_q, CakeTensorMap int dst = smem_kv_addr + (unsigned int)(kv_stage * 32768) + (unsigned int)(row_offset * 64 * 2); if (elect_sync()) { { - tma_gather4_gmem2smem(dst, tmap_swa_kv, kv_stage * 128, row0, row1, row2, row3, kv_full_addr + (kv_stage) * 8); - tma_gather4_gmem2smem(dst + 16384, tmap_swa_kv, kv_stage * 128 + 64, row0, row1, row2, row3, kv_full_addr + (kv_stage) * 8); + tma_gather4_gmem2smem(dst, (&tmap_swa_kv), kv_stage * 128, row0, row1, row2, row3, kv_full_addr + (kv_stage) * 8); + tma_gather4_gmem2smem(dst + 16384, (&tmap_swa_kv), kv_stage * 128 + 64, row0, row1, row2, row3, kv_full_addr + (kv_stage) * 8); } } } @@ -2313,8 +2319,8 @@ kernel_cake_dsv4_3e5a644d61420c9c7ea1(CakeTensorMap const* tmap_q, CakeTensorMap int dst_1 = smem_kv_addr + (unsigned int)(kv_stage_1 * 32768) + (unsigned int)(row_offset_1 * 64 * 2); if (elect_sync()) { { - tma_gather4_gmem2smem(dst_1, tmap_compressed_kv, kv_stage_1 * 128, row0_1, row1_1, row2_1, row3_1, kv_full_addr + (kv_stage_1) * 8); - tma_gather4_gmem2smem(dst_1 + 16384, tmap_compressed_kv, kv_stage_1 * 128 + 64, row0_1, row1_1, row2_1, row3_1, kv_full_addr + (kv_stage_1) * 8); + tma_gather4_gmem2smem(dst_1, (&tmap_compressed_kv), kv_stage_1 * 128, row0_1, row1_1, row2_1, row3_1, kv_full_addr + (kv_stage_1) * 8); + tma_gather4_gmem2smem(dst_1 + 16384, (&tmap_compressed_kv), kv_stage_1 * 128 + 64, row0_1, row1_1, row2_1, row3_1, kv_full_addr + (kv_stage_1) * 8); } } } @@ -2349,8 +2355,8 @@ kernel_cake_dsv4_3e5a644d61420c9c7ea1(CakeTensorMap const* tmap_q, CakeTensorMap int dst_2 = smem_v_addr + (unsigned int)(kv_stage_2 * 32768) + (unsigned int)(row_offset_2 * 64 * 2); if (elect_sync()) { { - tma_gather4_gmem2smem(dst_2, tmap_swa_kv, kv_stage_2 * 128, row0_2, row1_2, row2_2, row3_2, kv_full_addr + (kv_stage_2) * 8); - tma_gather4_gmem2smem(dst_2 + 16384, tmap_swa_kv, kv_stage_2 * 128 + 64, row0_2, row1_2, row2_2, row3_2, kv_full_addr + (kv_stage_2) * 8); + tma_gather4_gmem2smem(dst_2, (&tmap_swa_kv), kv_stage_2 * 128, row0_2, row1_2, row2_2, row3_2, kv_full_addr + (kv_stage_2) * 8); + tma_gather4_gmem2smem(dst_2 + 16384, (&tmap_swa_kv), kv_stage_2 * 128 + 64, row0_2, row1_2, row2_2, row3_2, kv_full_addr + (kv_stage_2) * 8); } } } @@ -2388,8 +2394,8 @@ kernel_cake_dsv4_3e5a644d61420c9c7ea1(CakeTensorMap const* tmap_q, CakeTensorMap int dst_3 = smem_kv_addr + (unsigned int)(kv_stage_3 * 32768) + (unsigned int)(row_offset_3 * 64 * 2); if (elect_sync()) { { - tma_gather4_gmem2smem(dst_3, tmap_compressed_kv, kv_stage_3 * 128, row0_3, row1_3, row2_3, row3_3, kv_full_addr + (kv_stage_3) * 8); - tma_gather4_gmem2smem(dst_3 + 16384, tmap_compressed_kv, kv_stage_3 * 128 + 64, row0_3, row1_3, row2_3, row3_3, kv_full_addr + (kv_stage_3) * 8); + tma_gather4_gmem2smem(dst_3, (&tmap_compressed_kv), kv_stage_3 * 128, row0_3, row1_3, row2_3, row3_3, kv_full_addr + (kv_stage_3) * 8); + tma_gather4_gmem2smem(dst_3 + 16384, (&tmap_compressed_kv), kv_stage_3 * 128 + 64, row0_3, row1_3, row2_3, row3_3, kv_full_addr + (kv_stage_3) * 8); } } } @@ -2426,8 +2432,8 @@ kernel_cake_dsv4_3e5a644d61420c9c7ea1(CakeTensorMap const* tmap_q, CakeTensorMap int dst_4 = smem_v_addr + (unsigned int)(kv_stage_4 * 32768) + (unsigned int)(row_offset_4 * 64 * 2); if (elect_sync()) { { - tma_gather4_gmem2smem(dst_4, tmap_compressed_kv, kv_stage_4 * 128, row0_4, row1_4, row2_4, row3_4, kv_full_addr + (kv_stage_4) * 8); - tma_gather4_gmem2smem(dst_4 + 16384, tmap_compressed_kv, kv_stage_4 * 128 + 64, row0_4, row1_4, row2_4, row3_4, kv_full_addr + (kv_stage_4) * 8); + tma_gather4_gmem2smem(dst_4, (&tmap_compressed_kv), kv_stage_4 * 128, row0_4, row1_4, row2_4, row3_4, kv_full_addr + (kv_stage_4) * 8); + tma_gather4_gmem2smem(dst_4 + 16384, (&tmap_compressed_kv), kv_stage_4 * 128 + 64, row0_4, row1_4, row2_4, row3_4, kv_full_addr + (kv_stage_4) * 8); } } } @@ -2468,8 +2474,8 @@ kernel_cake_dsv4_3e5a644d61420c9c7ea1(CakeTensorMap const* tmap_q, CakeTensorMap int dst_5 = smem_v_addr + (unsigned int)(kv_stage_5 * 32768) + (unsigned int)(row_offset_5 * 64 * 2); if (elect_sync()) { { - tma_gather4_gmem2smem(dst_5, tmap_swa_kv, kv_stage_5 * 128, row0_5, row1_5, row2_5, row3_5, kv_full_addr + (kv_stage_5) * 8); - tma_gather4_gmem2smem(dst_5 + 16384, tmap_swa_kv, kv_stage_5 * 128 + 64, row0_5, row1_5, row2_5, row3_5, kv_full_addr + (kv_stage_5) * 8); + tma_gather4_gmem2smem(dst_5, (&tmap_swa_kv), kv_stage_5 * 128, row0_5, row1_5, row2_5, row3_5, kv_full_addr + (kv_stage_5) * 8); + tma_gather4_gmem2smem(dst_5 + 16384, (&tmap_swa_kv), kv_stage_5 * 128 + 64, row0_5, row1_5, row2_5, row3_5, kv_full_addr + (kv_stage_5) * 8); } } } @@ -2497,8 +2503,8 @@ kernel_cake_dsv4_3e5a644d61420c9c7ea1(CakeTensorMap const* tmap_q, CakeTensorMap int dst_6 = smem_v_addr + (unsigned int)(kv_stage_6 * 32768) + (unsigned int)(row_offset_6 * 64 * 2); if (elect_sync()) { { - tma_gather4_gmem2smem(dst_6, tmap_compressed_kv, kv_stage_6 * 128, row0_6, row1_6, row2_6, row3_6, kv_full_addr + (kv_stage_6) * 8); - tma_gather4_gmem2smem(dst_6 + 16384, tmap_compressed_kv, kv_stage_6 * 128 + 64, row0_6, row1_6, row2_6, row3_6, kv_full_addr + (kv_stage_6) * 8); + tma_gather4_gmem2smem(dst_6, (&tmap_compressed_kv), kv_stage_6 * 128, row0_6, row1_6, row2_6, row3_6, kv_full_addr + (kv_stage_6) * 8); + tma_gather4_gmem2smem(dst_6 + 16384, (&tmap_compressed_kv), kv_stage_6 * 128 + 64, row0_6, row1_6, row2_6, row3_6, kv_full_addr + (kv_stage_6) * 8); } } } diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_c967e7efe220160e28d3_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_1ff52fece183b851b441_binding.cu similarity index 97% rename from csrc/cake_dsv4/sm_103a/cake_dsv4_c967e7efe220160e28d3_binding.cu rename to csrc/cake_dsv4/sm_103a/cake_dsv4_1ff52fece183b851b441_binding.cu index eeec0e7ee24..d59944bacd5 100644 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_c967e7efe220160e28d3_binding.cu +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_1ff52fece183b851b441_binding.cu @@ -15,7 +15,7 @@ */ // Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_c967e7efe220160e28d3'. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_1ff52fece183b851b441'. #include #include #include @@ -32,10 +32,10 @@ #include struct CakeTensorMap; -extern "C" __global__ void kernel_cake_dsv4_c967e7efe220160e28d3(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int batch_size, int max_q_len, int ragged_query, int has_sinks); +extern "C" __global__ void kernel_cake_dsv4_1ff52fece183b851b441(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int batch_size, int max_q_len, int ragged_query, int has_sinks); -namespace cake_host_shim_2f26cde03513c4ba { +namespace cake_host_shim_831b6e81c0a7ccc0 { using tvm::ffi::Optional; using tvm::ffi::TensorView; @@ -500,7 +500,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ } if (smem_status == cudaSuccess) { smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_c967e7efe220160e28d3), + reinterpret_cast(kernel_cake_dsv4_1ff52fece183b851b441), cudaFuncAttributeMaxDynamicSharedMemorySize, 198144); } @@ -510,7 +510,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ } } TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_c967e7efe220160e28d3 failed: " + << "cudaFuncSetAttribute for kernel_cake_dsv4_1ff52fece183b851b441 failed: " << cudaGetErrorString(smem_status); dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); dim3 block(256u, 1u, 1u); @@ -523,12 +523,12 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ config.dynamicSmemBytes = 198144u; config.stream = stream; cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_c967e7efe220160e28d3), kargs); + &config, reinterpret_cast(kernel_cake_dsv4_1ff52fece183b851b441), kargs); TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_c967e7efe220160e28d3 failed: " + << "cudaLaunchKernelExC for kernel_cake_dsv4_1ff52fece183b851b441 failed: " << cudaGetErrorString(launch_status); } -} // namespace cake_host_shim_2f26cde03513c4ba +} // namespace cake_host_shim_831b6e81c0a7ccc0 -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_2f26cde03513c4ba::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_831b6e81c0a7ccc0::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_c967e7efe220160e28d3_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_1ff52fece183b851b441_kernel.cu similarity index 99% rename from csrc/cake_dsv4/sm_103a/cake_dsv4_c967e7efe220160e28d3_kernel.cu rename to csrc/cake_dsv4/sm_103a/cake_dsv4_1ff52fece183b851b441_kernel.cu index b4dce1a94ee..47cf58b610b 100644 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_c967e7efe220160e28d3_kernel.cu +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_1ff52fece183b851b441_kernel.cu @@ -1062,7 +1062,7 @@ __device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { extern "C" { __global__ __launch_bounds__(256, 1) void -kernel_cake_dsv4_c967e7efe220160e28d3(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int batch_size, int max_q_len, int ragged_query, int has_sinks) +kernel_cake_dsv4_1ff52fece183b851b441(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int batch_size, int max_q_len, int ragged_query, int has_sinks) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); @@ -1603,15 +1603,6 @@ kernel_cake_dsv4_c967e7efe220160e28d3(CakeTensorMap const* tmap_q, CakeTensorMap int tile_count_2 = (active_topk_2 + 128 - 1) / 128; unsigned int load_kv_stage = 0; unsigned int load_kv_phase = 1; - if (load_warp_rank == 0) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(q_full_addr, 65536); - #pragma unroll - for (int q_stage = 0; q_stage < 8; q_stage += 8) { - tma_4d_gmem2smem(smem_q_addr + (unsigned int)(q_stage * 8192), tmap_q, 0, 0, q_stage, query_idx_2, q_full_addr); - } - } - } int group = load_warp_rank + lane * 3; int row0 = 0; int row1 = 0; @@ -1650,6 +1641,15 @@ kernel_cake_dsv4_c967e7efe220160e28d3(CakeTensorMap const* tmap_q, CakeTensorMap row2 = ((raw_rows[2] >= 0) ? raw_rows[2] : 0); row3 = ((raw_rows[3] >= 0) ? raw_rows[3] : 0); } + if (load_warp_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(q_full_addr, 65536); + #pragma unroll + for (int q_stage = 0; q_stage < 8; q_stage += 8) { + tma_4d_gmem2smem(smem_q_addr + (unsigned int)(q_stage * 8192), tmap_q, 0, 0, q_stage, query_idx_2, q_full_addr); + } + } + } #pragma unroll for (int tile_2 = 0; tile_2 < 1; tile_2++) { if (tile_count_2 > tile_2) { diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_3e5a644d61420c9c7ea1_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_3e5a644d61420c9c7ea1_binding.cu deleted file mode 100644 index 6d2febf4ef8..00000000000 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_3e5a644d61420c9c7ea1_binding.cu +++ /dev/null @@ -1,516 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer team. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -// Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_3e5a644d61420c9c7ea1'. -#include -#include -#include -#include - -#include "tvm_ffi_utils.h" - -#include -#include -#include -#include -#include -#include -#include - -struct CakeTensorMap; -extern "C" __global__ void kernel_cake_dsv4_3e5a644d61420c9c7ea1(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int num_query_tokens, int has_sinks); - - -namespace cake_host_shim_6afd242c7dfa9fb2 { - -using tvm::ffi::Optional; -using tvm::ffi::TensorView; - -class ScopedCudaDevice { - public: - explicit ScopedCudaDevice(int device_id) { - cudaError_t error = cudaGetDevice(&previous_device_); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed before host-shim launch: cudaError=" - << static_cast(error); - if (previous_device_ != device_id) { - error = cudaSetDevice(device_id); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaSetDevice failed before host-shim launch for cuda:" - << device_id << ": cudaError=" << static_cast(error); - restore_ = true; - } - } - - ScopedCudaDevice(const ScopedCudaDevice&) = delete; - ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; - - ~ScopedCudaDevice() noexcept { - if (restore_) { - (void)cudaSetDevice(previous_device_); - } - } - - private: - int previous_device_ = -1; - bool restore_ = false; -}; - -inline int64_t CakeDeviceMultiprocessorCount(int device_id) { - constexpr int kMaxCachedCudaDevices = 64; - TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) - << "physical-SM-count cache does not cover cuda:" << device_id; - static std::atomic count_by_device[kMaxCachedCudaDevices]{}; - int cached = count_by_device[device_id].load(std::memory_order_acquire); - if (cached > 0) return cached; - - int count = 0; - cudaError_t error = cudaDeviceGetAttribute( - &count, cudaDevAttrMultiProcessorCount, device_id); - TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) - << "querying multiProcessorCount failed for cuda:" << device_id - << ": cudaError=" << static_cast(error) << ", count=" << count; - // Concurrent first launches may repeat the immutable device query, but all - // publication is atomic and every later hot-path lookup is one acquire load. - count_by_device[device_id].store(count, std::memory_order_release); - return count; -} - -inline void CheckCudaTensor(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) - << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; -} - -inline void CheckSameCudaDevice( - const TensorView& t, - const TensorView& reference, - const char* name, - const char* reference_name) { - TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) - << name << " must be on the same CUDA device as " << reference_name - << ": got cuda:" << t.device().device_id - << " versus cuda:" << reference.device().device_id; -} - -inline void CheckCurrentCudaDevice( - const TensorView& reference, - const char* reference_name) { - int current_device = -1; - cudaError_t error = cudaGetDevice(¤t_device); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed while validating " << reference_name - << ": cudaError=" << static_cast(error); - TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) - << "current CUDA device must match " << reference_name - << ": current=cuda:" << current_device - << ", tensor=cuda:" << reference.device().device_id; -} - -inline void CheckContiguous(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; -} - -inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { - DLDataType d = t.dtype(); - TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) - << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits - << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits - << ", lanes=" << (int)d.lanes << ")"; -} - -// A logical axis.outer(trailing) folds every source dim above the trailing -// dimensions. Shape products are independent of physical strides, so verify -// the leading dimensions form one dense row-major chain instead of inventing -// a "folded stride". The descriptor reads its exact adjacent physical step -// separately through stride[-(trailing + 1)]. -inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { - TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) - << name << " cannot fold leading dimensions above " << trailing - << " trailing dims from ndim=" << t.ndim(); - int outer_last = t.ndim() - trailing - 1; - if (outer_last <= 0) { - return; - } - int64_t step = t.stride(outer_last); - TVM_FFI_CHECK(step > 0, ValueError) - << name << " physical strides must be positive"; - int64_t expected = step; - for (int axis = outer_last - 1; axis >= 0; --axis) { - expected *= t.size(axis + 1); - if (t.size(axis) > 1) { - TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) - << name << " leading dims are not physically foldable above " << trailing - << " trailing dims: stride(" << axis << ")=" << t.stride(axis) - << ", expected " << expected; - } - } -} - -// Explicit immutable caller storage: initialize once, outside graph capture, -// exactly as the production TmaDeviceSlot path. The owner must retain the -// allocation and must never expose it as mutable scratch or rewrite its bytes. -template -static inline void PrepareImmutableCallerTmaWorkspace( - const tvm::ffi::Tensor& retained_workspace, - const CUtensorMap (&maps)[N], cudaStream_t stream) { - void* workspace = retained_workspace.data_ptr(); - CUcontext context = nullptr; - CUresult result = cuCtxGetCurrent(&context); - TVM_FFI_CHECK(result == CUDA_SUCCESS && context != nullptr, RuntimeError) - << "immutable TMA workspace requires an active CUDA context"; - CUcontext allocation_context = nullptr; - result = cuPointerGetAttribute( - &allocation_context, CU_POINTER_ATTRIBUTE_CONTEXT, - reinterpret_cast(workspace)); - TVM_FFI_CHECK(result == CUDA_SUCCESS && allocation_context == context, RuntimeError) - << "immutable TMA workspace must belong to the active CUDA context"; - unsigned long long buffer_id = 0; - result = cuPointerGetAttribute( - &buffer_id, CU_POINTER_ATTRIBUTE_BUFFER_ID, - reinterpret_cast(workspace)); - TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) - << "querying immutable TMA allocation identity failed"; - std::string key = std::to_string(reinterpret_cast(context)); - key += ":" + std::to_string(reinterpret_cast(workspace)); - key += ":" + std::to_string(buffer_id); - std::string descriptor_bytes(reinterpret_cast(maps), sizeof(maps)); - struct InitializedWorkspace { - tvm::ffi::Tensor owner; - std::string descriptor_bytes; - }; - static std::mutex mu; - // Keep the actual Tensor allocation alive: a caching allocator may reuse a - // freed suballocation without changing its enclosing CUDA buffer identity. - static auto* initialized = new std::unordered_map(); - std::lock_guard lock(mu); - auto it = initialized->find(key); - if (it != initialized->end()) { - TVM_FFI_CHECK(it->second.descriptor_bytes == descriptor_bytes, ValueError) - << "immutable TMA descriptors changed; provide new private descriptor storage"; - return; - } - CUstreamCaptureStatus capture_status = CU_STREAM_CAPTURE_STATUS_NONE; - result = cuStreamIsCapturing(reinterpret_cast(stream), &capture_status); - TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) - << "querying immutable TMA initialization capture state failed"; - TVM_FFI_CHECK(capture_status == CU_STREAM_CAPTURE_STATUS_NONE, RuntimeError) - << "immutable TMA storage must be initialized before CUDA Graph capture"; - // Synchronous host-to-device initialization completes before publishing the - // slot to any thread/stream, matching TmaDeviceSlot. The consumer retains its - // generated tensor-map acquire; the mutable upload's release fence is unchanged. - result = cuMemcpyHtoD(reinterpret_cast(workspace), maps, sizeof(maps)); - TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) - << "initializing immutable caller TMA descriptors failed"; - initialized->emplace( - std::move(key), InitializedWorkspace{retained_workspace, std::move(descriptor_bytes)}); -} - -static inline void* CallerTmaWorkspaceSlot( - const tvm::ffi::TensorView& workspace, size_t slot) { - return static_cast(workspace.data_ptr()) + slot * sizeof(CUtensorMap); -} - -// 4D TMA descriptor for buffer 'tmap_q' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 3, ValueError) - << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - int64_t d2 = t.size(t.ndim() - 2); - int64_t d3 = t.size(t.ndim() - 3); - TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) - << "TMA source 'tmap_q' trailing dims must be positive"; - TVM_FFI_CHECK(d1 % 64 == 0, ValueError) - << "TMA source 'tmap_q' extent " << d1 - << " must divide exactly by " << 64; - uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) - << "TMA box (64, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; - int64_t carrier_stride_0 = d1; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; - int64_t carrier_stride_1 = 64; - TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; - TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; - TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; - int64_t carrier_stride_2 = (d2 * d1); - TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; - TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; - TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; - uint64_t global_strides[3] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - (uint64_t)((carrier_stride_1 * 16) / 8), - (uint64_t)((carrier_stride_2 * 16) / 8), - }; - uint32_t box_dim[4] = {64u, 64u, 1u, 1u}; - uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_compressed_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_compressed_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_compressed_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_compressed_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_compressed_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_compressed_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_compressed_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_kv') failed: CUresult=" << (int)r; - return tm; -} - -void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_compressed_kv, TensorView arg_O, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_num_query_tokens, int64_t arg_has_sinks, tvm::ffi::Tensor arg_tma_descriptor_workspace, int64_t grid_x, int64_t grid_y, int64_t grid_z) { - DLDevice dev = arg_tmap_q.device(); - ScopedCudaDevice device_guard(dev.device_id); - CheckCudaTensor(arg_tmap_q, "tmap_q"); - CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); - CheckContiguous(arg_tmap_q, "tmap_q"); - CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); - CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 4, 16, 1); - CheckCudaTensor(arg_tmap_compressed_kv, "tmap_compressed_kv"); - CheckDtype(arg_tmap_compressed_kv, "tmap_compressed_kv", 4, 16, 1); - CheckCudaTensor(arg_O, "O"); - CheckDtype(arg_O, "O", 4, 16, 1); - CheckContiguous(arg_O, "O"); - CheckCudaTensor(arg_swa_indices, "swa_indices"); - CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); - CheckContiguous(arg_swa_indices, "swa_indices"); - CheckCudaTensor(arg_compressed_indices, "compressed_indices"); - CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); - CheckContiguous(arg_compressed_indices, "compressed_indices"); - CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); - CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckCudaTensor(arg_sinks, "sinks"); - CheckDtype(arg_sinks, "sinks", 2, 32, 1); - CheckContiguous(arg_sinks, "sinks"); - CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); - CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); - CheckContiguous(arg_bmm1_scale, "bmm1_scale"); - CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); - CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); - CheckContiguous(arg_bmm2_scale, "bmm2_scale"); - TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) - << "scalar 'num_heads' value " << arg_num_heads - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) - << "scalar 'swa_index_stride' value " << arg_swa_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) - << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) - << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) - << "scalar 'sparse_topk' value " << arg_sparse_topk - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_num_query_tokens >= -2147483648LL && arg_num_query_tokens <= 2147483647LL, ValueError) - << "scalar 'num_query_tokens' value " << arg_num_query_tokens - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) - << "scalar 'has_sinks' value " << arg_has_sinks - << " is outside i32 range [-2147483648, 2147483647]"; - CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); - CheckSameCudaDevice(arg_tmap_compressed_kv, arg_tmap_q, "tmap_compressed_kv", "tmap_q"); - CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); - CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); - CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); - CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); - CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); - CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); - CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); - CheckCudaTensor(arg_tma_descriptor_workspace, "tma_descriptor_workspace"); - CheckSameCudaDevice(arg_tma_descriptor_workspace, arg_tmap_q, "tma_descriptor_workspace", "tmap_q"); - CheckContiguous(arg_tma_descriptor_workspace, "tma_descriptor_workspace"); - CheckDtype(arg_tma_descriptor_workspace, "tma_descriptor_workspace", 1, 8, 1); - TVM_FFI_CHECK(arg_tma_descriptor_workspace.numel() >= 384, ValueError) - << "tma_descriptor_workspace requires at least 384 bytes"; - TVM_FFI_CHECK( - reinterpret_cast(arg_tma_descriptor_workspace.data_ptr()) % 128 == 0, - ValueError) - << "tma_descriptor_workspace must be 128-byte aligned"; - TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) - << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y - << ", " << grid_z << ")"; - - cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); - CUtensorMap h_tmap_q = EncodeTma_tmap_q(arg_tmap_q); - CakeTensorMap const* p_tmap_q = static_cast(CallerTmaWorkspaceSlot(arg_tma_descriptor_workspace, 0u)); - CUtensorMap h_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); - CakeTensorMap const* p_tmap_swa_kv = static_cast(CallerTmaWorkspaceSlot(arg_tma_descriptor_workspace, 1u)); - CUtensorMap h_tmap_compressed_kv = EncodeTma_tmap_compressed_kv(arg_tmap_compressed_kv); - CakeTensorMap const* p_tmap_compressed_kv = static_cast(CallerTmaWorkspaceSlot(arg_tma_descriptor_workspace, 2u)); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); - int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); - int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); - int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); - float* p_sinks = static_cast(arg_sinks.data_ptr()); - float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); - float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); - int32_t v_num_heads = (int32_t)arg_num_heads; - int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; - int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; - int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; - int32_t v_sparse_topk = (int32_t)arg_sparse_topk; - int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens; - int32_t v_has_sinks = (int32_t)arg_has_sinks; - void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_tmap_compressed_kv, &p_O, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_num_query_tokens, &v_has_sinks}; - - static std::mutex smem_mu; - static auto* smem_status_by_device = - new std::unordered_map(); - cudaError_t smem_status = cudaSuccess; - { - std::lock_guard lock(smem_mu); - auto it = smem_status_by_device->find(dev.device_id); - if (it == smem_status_by_device->end()) { - int current_device = -1; - smem_status = cudaGetDevice(¤t_device); - if (smem_status == cudaSuccess && current_device != dev.device_id) { - smem_status = cudaErrorInvalidDevice; - } - if (smem_status == cudaSuccess) { - smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_3e5a644d61420c9c7ea1), - cudaFuncAttributeMaxDynamicSharedMemorySize, - 203904); - } - smem_status_by_device->emplace(dev.device_id, smem_status); - } else { - smem_status = it->second; - } - } - TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_3e5a644d61420c9c7ea1 failed: " - << cudaGetErrorString(smem_status); - dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); - dim3 block(512u, 1u, 1u); - const CUtensorMap workspace_maps[] = {h_tmap_q, h_tmap_swa_kv, h_tmap_compressed_kv}; - PrepareImmutableCallerTmaWorkspace(arg_tma_descriptor_workspace, workspace_maps, stream); - - cudaLaunchConfig_t config{}; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = 203904u; - config.stream = stream; - cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_3e5a644d61420c9c7ea1), kargs); - TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_3e5a644d61420c9c7ea1 failed: " - << cudaGetErrorString(launch_status); -} - -} // namespace cake_host_shim_6afd242c7dfa9fb2 - -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_6afd242c7dfa9fb2::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_125332c2066d89bd7610_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_4ab0487d5fde70d4569e_binding.cu similarity index 86% rename from csrc/cake_dsv4/sm_103a/cake_dsv4_125332c2066d89bd7610_binding.cu rename to csrc/cake_dsv4/sm_103a/cake_dsv4_4ab0487d5fde70d4569e_binding.cu index e463bc58de8..9e594768617 100644 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_125332c2066d89bd7610_binding.cu +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_4ab0487d5fde70d4569e_binding.cu @@ -15,7 +15,7 @@ */ // Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_125332c2066d89bd7610'. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_4ab0487d5fde70d4569e'. #include #include #include @@ -31,10 +31,10 @@ #include #include -extern "C" __global__ void kernel_cake_dsv4_125332c2066d89bd7610(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); +extern "C" __global__ void kernel_cake_dsv4_4ab0487d5fde70d4569e(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, const __grid_constant__ CUtensorMap O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); -namespace cake_host_shim_0f71e276da509b5c { +namespace cake_host_shim_df4c5e4bdb73dbc3 { using tvm::ffi::Optional; using tvm::ffi::TensorView; @@ -390,6 +390,57 @@ inline CUtensorMap EncodeTma_tmap_compressed_v(const TensorView& t) { return tm; } +// 4D TMA descriptor for buffer 'O' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_O(const TensorView& t, int64_t tp_num_heads, int64_t tp_num_query_tokens) { + TVM_FFI_CHECK(t.ndim() >= 1, ValueError) + << "TMA source 'O' must have at least 1 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'O' must have unit innermost stride, got " << t.stride(-1); + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(8), (uint64_t)(5), (uint64_t)((tp_num_query_tokens * tp_num_heads))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'O' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1] && 1u <= global_dim[2] && 32u <= global_dim[3], ValueError) + << "TMA box (64, 1, 1, 32) exceeds resolved global dims for 'O'"; + int64_t carrier_stride_0 = 64; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 512; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = 2560; + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 1u, 1u, 32u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'O') failed: CUresult=" << (int)r; + return tm; +} + void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_compressed_k, TensorView arg_tmap_swa_v, TensorView arg_tmap_compressed_v, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_num_query_tokens, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z) { DLDevice dev = arg_tmap_q.device(); ScopedCudaDevice device_guard(dev.device_id); @@ -477,7 +528,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_c CUtensorMap p_tmap_compressed_k = EncodeTma_tmap_compressed_k(arg_tmap_compressed_k); CUtensorMap p_tmap_swa_v = EncodeTma_tmap_swa_v(arg_tmap_swa_v); CUtensorMap p_tmap_compressed_v = EncodeTma_tmap_compressed_v(arg_tmap_compressed_v); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); + CUtensorMap p_O = EncodeTma_O(arg_O, (int64_t)arg_num_heads, (int64_t)arg_num_query_tokens); float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); @@ -510,7 +561,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_c } if (smem_status == cudaSuccess) { smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_125332c2066d89bd7610), + reinterpret_cast(kernel_cake_dsv4_4ab0487d5fde70d4569e), cudaFuncAttributeMaxDynamicSharedMemorySize, 231936); } @@ -520,7 +571,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_c } } TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_125332c2066d89bd7610 failed: " + << "cudaFuncSetAttribute for kernel_cake_dsv4_4ab0487d5fde70d4569e failed: " << cudaGetErrorString(smem_status); dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); dim3 block(512u, 1u, 1u); @@ -544,12 +595,12 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_c config.attrs = attrs; config.numAttrs = n; cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_125332c2066d89bd7610), kargs); + &config, reinterpret_cast(kernel_cake_dsv4_4ab0487d5fde70d4569e), kargs); TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_125332c2066d89bd7610 failed: " + << "cudaLaunchKernelExC for kernel_cake_dsv4_4ab0487d5fde70d4569e failed: " << cudaGetErrorString(launch_status); } -} // namespace cake_host_shim_0f71e276da509b5c +} // namespace cake_host_shim_df4c5e4bdb73dbc3 -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_0f71e276da509b5c::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_df4c5e4bdb73dbc3::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_125332c2066d89bd7610_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_4ab0487d5fde70d4569e_kernel.cu similarity index 96% rename from csrc/cake_dsv4/sm_103a/cake_dsv4_125332c2066d89bd7610_kernel.cu rename to csrc/cake_dsv4/sm_103a/cake_dsv4_4ab0487d5fde70d4569e_kernel.cu index d4c22d86df8..17aca3f4231 100644 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_125332c2066d89bd7610_kernel.cu +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_4ab0487d5fde70d4569e_kernel.cu @@ -1178,6 +1178,15 @@ __device__ __forceinline__ void tma_gather4_gmem2smem_mc_cta2( } +__device__ __forceinline__ void tma_store_4d( + const void *tmap, int x, int y, int z, int w, unsigned smem_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.global.shared::cta.tile.bulk_group" + " [%0, {%1, %2, %3, %4}], [%5];" + :: "l"(tmap), "r"(x), "r"(y), "r"(z), "r"(w), "r"(smem_addr) : "memory"); +} + + __device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { asm volatile( "{\n\t" @@ -1192,7 +1201,7 @@ __device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint extern "C" { __global__ __launch_bounds__(512, 1) __cluster_dims__(2,1,1) void -kernel_cake_dsv4_125332c2066d89bd7610(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items) +kernel_cake_dsv4_4ab0487d5fde70d4569e(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, const __grid_constant__ CUtensorMap O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); @@ -1218,10 +1227,11 @@ kernel_cake_dsv4_125332c2066d89bd7610(const __grid_constant__ CUtensorMap tmap_q #define stats_addr (mbar_base + 312) #define stats_empty_addr (mbar_base + 328) #define o_full_addr (mbar_base + 344) - #define o_first_slice_seeded_addr (mbar_base + 352) - #define tmem_dealloc_peer_addr (mbar_base + 360) - #define index_valid_full_addr (mbar_base + 368) - #define index_valid_empty_addr (mbar_base + 384) + #define epi_q_free_addr (mbar_base + 352) + #define o_first_slice_seeded_addr (mbar_base + 360) + #define tmem_dealloc_peer_addr (mbar_base + 368) + #define index_valid_full_addr (mbar_base + 376) + #define index_valid_empty_addr (mbar_base + 392) const int bid = blockIdx.x; const int num_bids = gridDim.x; @@ -1252,8 +1262,8 @@ kernel_cake_dsv4_125332c2066d89bd7610(const __grid_constant__ CUtensorMap tmap_q asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); - // Mbarrier init (19 pipeline groups, 0 ordered-sequence groups, 50 barriers) - // Mbarriers at smem_raw[231488..231888) + // Mbarrier init (20 pipeline groups, 0 ordered-sequence groups, 51 barriers) + // Mbarriers at smem_raw[231488..231896) if (warp == 0) { uint32_t leader = elect_sync(); @@ -1319,16 +1329,18 @@ kernel_cake_dsv4_125332c2066d89bd7610(const __grid_constant__ CUtensorMap tmap_q mbarrier_init(smem + 231824, 128); // o_full: 1 barriers, init_count=1 mbarrier_init(smem + 231832, 1); + // epi_q_free: 1 barriers, init_count=4 + mbarrier_init(smem + 231840, 4); // o_first_slice_seeded: 1 barriers, init_count=256 - mbarrier_init(smem + 231840, 256); + mbarrier_init(smem + 231848, 256); // tmem_dealloc_peer: 1 barriers, init_count=32 - mbarrier_init(smem + 231848, 32); - // index_valid_full: 2 barriers, init_count=32 mbarrier_init(smem + 231856, 32); + // index_valid_full: 2 barriers, init_count=32 mbarrier_init(smem + 231864, 32); + mbarrier_init(smem + 231872, 32); // index_valid_empty: 2 barriers, init_count=128 - mbarrier_init(smem + 231872, 128); mbarrier_init(smem + 231880, 128); + mbarrier_init(smem + 231888, 128); asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); } } @@ -1337,9 +1349,9 @@ kernel_cake_dsv4_125332c2066d89bd7610(const __grid_constant__ CUtensorMap tmap_q asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); // TMEM alloc (512 columns, 512 used) - volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 231888); + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 231896); if (warp == 0) { - int _tmem_hold = smem + 231888; + int _tmem_hold = smem + 231896; asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); __syncwarp(); asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); @@ -1717,11 +1729,15 @@ kernel_cake_dsv4_125332c2066d89bd7610(const __grid_constant__ CUtensorMap tmap_q asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum)); partial_lse[stat_offset] = ((final_sum > 0.0f) ? final_max * softmax_scale_log2_1 + _log2_0 : -CAKE_INF); } + float epi_scale = inv_sum * output_scale; + int epi_box_base = smem_q_addr + (unsigned int)(warp % 4 * 16384); + int epi_row0 = query_idx_1 * num_heads + cta_rank * 64 + head_row_base_1; #pragma unroll for (int acc_stage_1 = 0; acc_stage_1 < 2; acc_stage_1++) { #pragma unroll for (int local_slice_1 = 0; local_slice_1 < 2; local_slice_1++) { int logical_slice = acc_stage_1 * 2 * 2 + n_half_1 * 2 + local_slice_1; + int epi_box = epi_box_base + (acc_stage_1 * 2 + local_slice_1) * 4096; int o_base_1 = taddr + 128 + (unsigned int)(acc_stage_1 * 128) + (unsigned int)(local_slice_1 * 64) + (unsigned int)(output_row_base << 16); #pragma unroll for (int c_1 = 0; c_1 < 64; c_1 += 32) { @@ -1731,46 +1747,35 @@ kernel_cake_dsv4_125332c2066d89bd7610(const __grid_constant__ CUtensorMap tmap_q " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]) : "r"(o_base_1 + c_1)); - int gmem_base = output_offset + logical_slice * 64 + c_1; + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + float epi_scaled[32]; #pragma unroll - for (int j = 0; j < 32; j += 16) { - { - const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; - #if __CUDA_ARCH__ >= 1000 - #pragma unroll - for (int _ps = 0; _ps < 8; _ps++) - mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_3[j])[_ps], _prescale2_1); - #else - #pragma unroll - for (int _ps = 0; _ps < 16; _ps++) - _tmem_load_3[j + _ps] *= inv_sum * output_scale; - #endif - { - __nv_bfloat162 _pk0 = __floats2bfloat162_rn(_tmem_load_3[j + 0], _tmem_load_3[j + 1]); - unsigned _pk_u0 = *reinterpret_cast(&_pk0); - __nv_bfloat162 _pk1 = __floats2bfloat162_rn(_tmem_load_3[j + 2], _tmem_load_3[j + 3]); - unsigned _pk_u1 = *reinterpret_cast(&_pk1); - __nv_bfloat162 _pk2 = __floats2bfloat162_rn(_tmem_load_3[j + 4], _tmem_load_3[j + 5]); - unsigned _pk_u2 = *reinterpret_cast(&_pk2); - __nv_bfloat162 _pk3 = __floats2bfloat162_rn(_tmem_load_3[j + 6], _tmem_load_3[j + 7]); - unsigned _pk_u3 = *reinterpret_cast(&_pk3); - __nv_bfloat162 _pk4 = __floats2bfloat162_rn(_tmem_load_3[j + 8], _tmem_load_3[j + 9]); - unsigned _pk_u4 = *reinterpret_cast(&_pk4); - __nv_bfloat162 _pk5 = __floats2bfloat162_rn(_tmem_load_3[j + 10], _tmem_load_3[j + 11]); - unsigned _pk_u5 = *reinterpret_cast(&_pk5); - __nv_bfloat162 _pk6 = __floats2bfloat162_rn(_tmem_load_3[j + 12], _tmem_load_3[j + 13]); - unsigned _pk_u6 = *reinterpret_cast(&_pk6); - __nv_bfloat162 _pk7 = __floats2bfloat162_rn(_tmem_load_3[j + 14], _tmem_load_3[j + 15]); - unsigned _pk_u7 = *reinterpret_cast(&_pk7); - asm volatile( - "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" - :: "l"((void*)(&((__nv_bfloat16*)(O + (gmem_base + j)))[0])), "r"(_pk_u0), "r"(_pk_u1), "r"(_pk_u2), "r"(_pk_u3), "r"(_pk_u4), "r"(_pk_u5), "r"(_pk_u6), "r"(_pk_u7) : "memory"); - } - } + for (int e = 0; e < 32; e++) { + epi_scaled[e] = _tmem_load_3[e] * epi_scale; + } + unsigned int epi_packed[16]; + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(epi_scaled[_lp*2 + 0], epi_scaled[_lp*2+1 + 0])); + epi_packed[_lp] = *(uint32_t*)&_bf2; } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((epi_box + (lane * 128 + c_1 * 2 ^ (lane * 128 + c_1 * 2 >> 7 & 7) << 4))), "r"(epi_packed[0]), "r"(epi_packed[1]), "r"(epi_packed[2]), "r"(epi_packed[3]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((epi_box + (lane * 128 + (c_1 * 2 + 16) ^ (lane * 128 + (c_1 * 2 + 16) >> 7 & 7) << 4))), "r"(epi_packed[4]), "r"(epi_packed[5]), "r"(epi_packed[6]), "r"(epi_packed[7]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((epi_box + (lane * 128 + (c_1 * 2 + 32) ^ (lane * 128 + (c_1 * 2 + 32) >> 7 & 7) << 4))), "r"(epi_packed[8]), "r"(epi_packed[9]), "r"(epi_packed[10]), "r"(epi_packed[11]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((epi_box + (lane * 128 + (c_1 * 2 + 48) ^ (lane * 128 + (c_1 * 2 + 48) >> 7 & 7) << 4))), "r"(epi_packed[12]), "r"(epi_packed[13]), "r"(epi_packed[14]), "r"(epi_packed[15]) : "memory"); + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + __syncwarp(); + if (elect_sync()) { + tma_store_4d((&O), 0, logical_slice, split_idx_1, epi_row0, epi_box); + asm volatile("cp.async.bulk.commit_group;"); } } } + if (elect_sync()) { + asm volatile("cp.async.bulk.wait_group.read 0;"); + mbarrier_arrive(epi_q_free_addr); + } asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); asm volatile( "{\n\t" @@ -1781,6 +1786,9 @@ kernel_cake_dsv4_125332c2066d89bd7610(const __grid_constant__ CUtensorMap tmap_q :: "r"(o_empty_addr), "r"(0) : "memory"); tile_cursor_1 = tile_cursor_1 + num_kv_tiles_1; } + if (elect_sync()) { + asm volatile("cp.async.bulk.wait_group 0;"); + } asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); asm volatile("tcgen05.fence::before_thread_sync;"); } @@ -2063,6 +2071,7 @@ kernel_cake_dsv4_125332c2066d89bd7610(const __grid_constant__ CUtensorMap tmap_q unsigned int v_stage_1 = 0; int load_tile_cursor = 0; unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_epi_q_free_0 = 1; unsigned int _phase_slot_empty = 1; unsigned int _phase_v_empty = 1; #pragma unroll 1 @@ -2077,6 +2086,8 @@ kernel_cake_dsv4_125332c2066d89bd7610(const __grid_constant__ CUtensorMap tmap_q int group = load_warp_rank * 4 + lane; mbarrier_wait(q_empty_addr, _phase_q_empty_0); _phase_q_empty_0 ^= 1; + mbarrier_wait(epi_q_free_addr, _phase_epi_q_free_0); + _phase_epi_q_free_0 ^= 1; if (load_warp_rank == 0) { if (cta_rank == 0) { if (elect_sync()) { diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_a2be51a1fd34ec0d9a3b_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_868bb2291a1464f14e05_binding.cu similarity index 97% rename from csrc/cake_dsv4/sm_103a/cake_dsv4_a2be51a1fd34ec0d9a3b_binding.cu rename to csrc/cake_dsv4/sm_103a/cake_dsv4_868bb2291a1464f14e05_binding.cu index 035179e18f8..fda3c5b89ff 100644 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_a2be51a1fd34ec0d9a3b_binding.cu +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_868bb2291a1464f14e05_binding.cu @@ -15,7 +15,7 @@ */ // Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_a2be51a1fd34ec0d9a3b'. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_868bb2291a1464f14e05'. #include #include #include @@ -32,10 +32,10 @@ #include struct CakeTensorMap; -extern "C" __global__ void kernel_cake_dsv4_a2be51a1fd34ec0d9a3b(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items); +extern "C" __global__ void kernel_cake_dsv4_868bb2291a1464f14e05(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items); -namespace cake_host_shim_5841b4c9e46c06fc { +namespace cake_host_shim_262e38a78b1a7fce { using tvm::ffi::Optional; using tvm::ffi::TensorView; @@ -491,7 +491,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ } if (smem_status == cudaSuccess) { smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_a2be51a1fd34ec0d9a3b), + reinterpret_cast(kernel_cake_dsv4_868bb2291a1464f14e05), cudaFuncAttributeMaxDynamicSharedMemorySize, 205824); } @@ -501,7 +501,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ } } TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_a2be51a1fd34ec0d9a3b failed: " + << "cudaFuncSetAttribute for kernel_cake_dsv4_868bb2291a1464f14e05 failed: " << cudaGetErrorString(smem_status); dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); dim3 block(512u, 1u, 1u); @@ -514,12 +514,12 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ config.dynamicSmemBytes = 205824u; config.stream = stream; cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_a2be51a1fd34ec0d9a3b), kargs); + &config, reinterpret_cast(kernel_cake_dsv4_868bb2291a1464f14e05), kargs); TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_a2be51a1fd34ec0d9a3b failed: " + << "cudaLaunchKernelExC for kernel_cake_dsv4_868bb2291a1464f14e05 failed: " << cudaGetErrorString(launch_status); } -} // namespace cake_host_shim_5841b4c9e46c06fc +} // namespace cake_host_shim_262e38a78b1a7fce -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_5841b4c9e46c06fc::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_262e38a78b1a7fce::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_a2be51a1fd34ec0d9a3b_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_868bb2291a1464f14e05_kernel.cu similarity index 54% rename from csrc/cake_dsv4/sm_103a/cake_dsv4_a2be51a1fd34ec0d9a3b_kernel.cu rename to csrc/cake_dsv4/sm_103a/cake_dsv4_868bb2291a1464f14e05_kernel.cu index 35598530ac2..47307e2798a 100644 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_a2be51a1fd34ec0d9a3b_kernel.cu +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_868bb2291a1464f14e05_kernel.cu @@ -1194,7 +1194,7 @@ __device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { extern "C" { __global__ __launch_bounds__(512, 1) void -kernel_cake_dsv4_a2be51a1fd34ec0d9a3b(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items) +kernel_cake_dsv4_868bb2291a1464f14e05(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); @@ -1472,7 +1472,7 @@ kernel_cake_dsv4_a2be51a1fd34ec0d9a3b(CakeTensorMap const* tmap_q, CakeTensorMap asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_sparse_indices_addr + (unsigned int)(index_offset * 4)), "r"(sparse_rows[0]), "r"(sparse_rows[1]), "r"(sparse_rows[2]), "r"(sparse_rows[3]) : "memory"); } asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile("barrier.sync 8, 128;" ::: "memory"); + asm volatile("barrier.sync 8, 256;" ::: "memory"); } int k_index_stage_base = smem_sparse_indices_addr; { @@ -1485,7 +1485,8 @@ kernel_cake_dsv4_a2be51a1fd34ec0d9a3b(CakeTensorMap const* tmap_q, CakeTensorMap } int k_dst = smem_kv_addr + load_kv_stage * 16384; int group = lane; - if (group < 32) { + int k_group_limit = ((work_idx_1 == (unsigned int)bid) ? 16 : 32); + if (group < k_group_limit) { int group_offset = k_cta_offset + group * 4; int raw_rows[4]; asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" @@ -1663,6 +1664,29 @@ kernel_cake_dsv4_a2be51a1fd34ec0d9a3b(CakeTensorMap const* tmap_q, CakeTensorMap asm volatile("tcgen05.fence::before_thread_sync;"); mbarrier_arrive(s_seeded_addr); } + { + asm volatile("barrier.sync 8, 256;" ::: "memory"); + if (work_idx_2 == (unsigned int)bid) { + const int helper_stage = warp % 4; + int helper_dst = smem_kv_addr + (unsigned int)(helper_stage * 16384); + int helper_group = lane; + if (helper_group >= 16) { + int helper_rows[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&helper_rows[0])), "=r"(*reinterpret_cast(&helper_rows[(0) + 1])), "=r"(*reinterpret_cast(&helper_rows[(0) + 2])), "=r"(*reinterpret_cast(&helper_rows[(0) + 3])) + : "r"(smem_sparse_indices_addr + (unsigned int)(helper_group * 16))); + int hrow0 = ((helper_rows[0] >= 0) ? helper_rows[0] : 0); + int hrow1 = ((helper_rows[1] >= 0) ? helper_rows[1] : 0); + int hrow2 = ((helper_rows[2] >= 0) ? helper_rows[2] : 0); + int hrow3 = ((helper_rows[3] >= 0) ? helper_rows[3] : 0); + if (first_tile_1 == 0) { + tma_gather4_gmem2smem(helper_dst + helper_group * 512, tmap_swa_kv, helper_stage * 128, hrow0, hrow1, hrow2, hrow3, kv_full_addr + (helper_stage) * 8); + } else { + tma_gather4_gmem2smem(helper_dst + helper_group * 512, tmap_compressed_kv, helper_stage * 128, hrow0, hrow1, hrow2, hrow3, kv_full_addr + (helper_stage) * 8); + } + } + } + } { mbarrier_wait(q_full_addr, _phase_q_full_0_1); _phase_q_full_0_1 ^= 1; @@ -1704,121 +1728,437 @@ kernel_cake_dsv4_a2be51a1fd34ec0d9a3b(CakeTensorMap const* tmap_q, CakeTensorMap invalid_cols[w] = _vote_0; any_invalid = any_invalid | invalid_cols[w]; } - float _tmem_load_0[4]; - tmem_ld_x4(&_tmem_load_0[0], s_base); + float _tmem_load_0[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]) + : "r"(s_base)); float _tmem_load_1[4]; tmem_ld_x4(&_tmem_load_1[0], s_base); + float _tmem_load_2[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]) + : "r"(s_base + 32)); + float _tmem_load_3[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]) + : "r"(s_base + 64)); + float _tmem_load_4[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_4[0]), "=f"(_tmem_load_4[1]), "=f"(_tmem_load_4[2]), "=f"(_tmem_load_4[3]), "=f"(_tmem_load_4[4]), "=f"(_tmem_load_4[5]), "=f"(_tmem_load_4[6]), "=f"(_tmem_load_4[7]), "=f"(_tmem_load_4[8]), "=f"(_tmem_load_4[9]), "=f"(_tmem_load_4[10]), "=f"(_tmem_load_4[11]), "=f"(_tmem_load_4[12]), "=f"(_tmem_load_4[13]), "=f"(_tmem_load_4[14]), "=f"(_tmem_load_4[15]), "=f"(_tmem_load_4[16]), "=f"(_tmem_load_4[17]), "=f"(_tmem_load_4[18]), "=f"(_tmem_load_4[19]), "=f"(_tmem_load_4[20]), "=f"(_tmem_load_4[21]), "=f"(_tmem_load_4[22]), "=f"(_tmem_load_4[23]), "=f"(_tmem_load_4[24]), "=f"(_tmem_load_4[25]), "=f"(_tmem_load_4[26]), "=f"(_tmem_load_4[27]), "=f"(_tmem_load_4[28]), "=f"(_tmem_load_4[29]), "=f"(_tmem_load_4[30]), "=f"(_tmem_load_4[31]) + : "r"(s_base + 96)); { - #pragma unroll - for (int frag = 0; frag < 4; frag++) { - float _tmem_load_2[32]; - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]) - : "r"(s_base + frag * 32)); - int frag_valid = valid_sparse_cols - frag * 32; - uint32_t _slice_lo_mask_1; - { - int _lim_0 = frag_valid; - if (_lim_0 <= 0) { _slice_lo_mask_1 = 0u; } - else if (_lim_0 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_0)); - } + int frag_valid_0 = valid_sparse_cols; + uint32_t _slice_lo_mask_1; + { + int _lim_0 = frag_valid_0; + if (_lim_0 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_0)); } - if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_2[0] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_2[1] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_2[2] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_2[3] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_2[4] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_2[5] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_2[6] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_2[7] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_2[8] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_2[9] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_2[10] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_2[11] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_2[12] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_2[13] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_2[14] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_2[15] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_2[16] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_2[17] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_2[18] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_2[19] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_2[20] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_2[21] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_2[22] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_2[23] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_2[24] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_2[25] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_2[26] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_2[27] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_2[28] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_2[29] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_2[30] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_2[31] = -CAKE_INF; - if (any_invalid != 0) { - unsigned int frag_mask_word = invalid_cols[frag]; - _tmem_load_2[0] = (((frag_mask_word & 1) != 0) ? -CAKE_INF : _tmem_load_2[0]); - _tmem_load_2[1] = (((frag_mask_word >> 1 & 1) != 0) ? -CAKE_INF : _tmem_load_2[1]); - _tmem_load_2[2] = (((frag_mask_word >> 2 & 1) != 0) ? -CAKE_INF : _tmem_load_2[2]); - _tmem_load_2[3] = (((frag_mask_word >> 3 & 1) != 0) ? -CAKE_INF : _tmem_load_2[3]); - _tmem_load_2[4] = (((frag_mask_word >> 4 & 1) != 0) ? -CAKE_INF : _tmem_load_2[4]); - _tmem_load_2[5] = (((frag_mask_word >> 5 & 1) != 0) ? -CAKE_INF : _tmem_load_2[5]); - _tmem_load_2[6] = (((frag_mask_word >> 6 & 1) != 0) ? -CAKE_INF : _tmem_load_2[6]); - _tmem_load_2[7] = (((frag_mask_word >> 7 & 1) != 0) ? -CAKE_INF : _tmem_load_2[7]); - _tmem_load_2[8] = (((frag_mask_word >> 8 & 1) != 0) ? -CAKE_INF : _tmem_load_2[8]); - _tmem_load_2[9] = (((frag_mask_word >> 9 & 1) != 0) ? -CAKE_INF : _tmem_load_2[9]); - _tmem_load_2[10] = (((frag_mask_word >> 10 & 1) != 0) ? -CAKE_INF : _tmem_load_2[10]); - _tmem_load_2[11] = (((frag_mask_word >> 11 & 1) != 0) ? -CAKE_INF : _tmem_load_2[11]); - _tmem_load_2[12] = (((frag_mask_word >> 12 & 1) != 0) ? -CAKE_INF : _tmem_load_2[12]); - _tmem_load_2[13] = (((frag_mask_word >> 13 & 1) != 0) ? -CAKE_INF : _tmem_load_2[13]); - _tmem_load_2[14] = (((frag_mask_word >> 14 & 1) != 0) ? -CAKE_INF : _tmem_load_2[14]); - _tmem_load_2[15] = (((frag_mask_word >> 15 & 1) != 0) ? -CAKE_INF : _tmem_load_2[15]); - _tmem_load_2[16] = (((frag_mask_word >> 16 & 1) != 0) ? -CAKE_INF : _tmem_load_2[16]); - _tmem_load_2[17] = (((frag_mask_word >> 17 & 1) != 0) ? -CAKE_INF : _tmem_load_2[17]); - _tmem_load_2[18] = (((frag_mask_word >> 18 & 1) != 0) ? -CAKE_INF : _tmem_load_2[18]); - _tmem_load_2[19] = (((frag_mask_word >> 19 & 1) != 0) ? -CAKE_INF : _tmem_load_2[19]); - _tmem_load_2[20] = (((frag_mask_word >> 20 & 1) != 0) ? -CAKE_INF : _tmem_load_2[20]); - _tmem_load_2[21] = (((frag_mask_word >> 21 & 1) != 0) ? -CAKE_INF : _tmem_load_2[21]); - _tmem_load_2[22] = (((frag_mask_word >> 22 & 1) != 0) ? -CAKE_INF : _tmem_load_2[22]); - _tmem_load_2[23] = (((frag_mask_word >> 23 & 1) != 0) ? -CAKE_INF : _tmem_load_2[23]); - _tmem_load_2[24] = (((frag_mask_word >> 24 & 1) != 0) ? -CAKE_INF : _tmem_load_2[24]); - _tmem_load_2[25] = (((frag_mask_word >> 25 & 1) != 0) ? -CAKE_INF : _tmem_load_2[25]); - _tmem_load_2[26] = (((frag_mask_word >> 26 & 1) != 0) ? -CAKE_INF : _tmem_load_2[26]); - _tmem_load_2[27] = (((frag_mask_word >> 27 & 1) != 0) ? -CAKE_INF : _tmem_load_2[27]); - _tmem_load_2[28] = (((frag_mask_word >> 28 & 1) != 0) ? -CAKE_INF : _tmem_load_2[28]); - _tmem_load_2[29] = (((frag_mask_word >> 29 & 1) != 0) ? -CAKE_INF : _tmem_load_2[29]); - _tmem_load_2[30] = (((frag_mask_word >> 30 & 1) != 0) ? -CAKE_INF : _tmem_load_2[30]); - _tmem_load_2[31] = (((frag_mask_word >> 31 & 1) != 0) ? -CAKE_INF : _tmem_load_2[31]); + } + if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_0[0] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_0[1] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_0[2] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_0[3] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_0[4] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_0[5] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_0[6] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_0[7] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_0[8] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_0[9] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_0[10] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_0[11] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_0[12] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_0[13] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_0[14] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_0[15] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_0[16] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_0[17] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_0[18] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_0[19] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_0[20] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_0[21] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_0[22] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_0[23] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_0[24] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_0[25] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_0[26] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_0[27] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_0[28] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_0[29] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_0[30] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_0[31] = -CAKE_INF; + if (any_invalid != 0) { + unsigned int frag_mask_word_0 = invalid_cols[0]; + _tmem_load_0[0] = (((frag_mask_word_0 & 1) != 0) ? -CAKE_INF : _tmem_load_0[0]); + _tmem_load_0[1] = (((frag_mask_word_0 >> 1 & 1) != 0) ? -CAKE_INF : _tmem_load_0[1]); + _tmem_load_0[2] = (((frag_mask_word_0 >> 2 & 1) != 0) ? -CAKE_INF : _tmem_load_0[2]); + _tmem_load_0[3] = (((frag_mask_word_0 >> 3 & 1) != 0) ? -CAKE_INF : _tmem_load_0[3]); + _tmem_load_0[4] = (((frag_mask_word_0 >> 4 & 1) != 0) ? -CAKE_INF : _tmem_load_0[4]); + _tmem_load_0[5] = (((frag_mask_word_0 >> 5 & 1) != 0) ? -CAKE_INF : _tmem_load_0[5]); + _tmem_load_0[6] = (((frag_mask_word_0 >> 6 & 1) != 0) ? -CAKE_INF : _tmem_load_0[6]); + _tmem_load_0[7] = (((frag_mask_word_0 >> 7 & 1) != 0) ? -CAKE_INF : _tmem_load_0[7]); + _tmem_load_0[8] = (((frag_mask_word_0 >> 8 & 1) != 0) ? -CAKE_INF : _tmem_load_0[8]); + _tmem_load_0[9] = (((frag_mask_word_0 >> 9 & 1) != 0) ? -CAKE_INF : _tmem_load_0[9]); + _tmem_load_0[10] = (((frag_mask_word_0 >> 10 & 1) != 0) ? -CAKE_INF : _tmem_load_0[10]); + _tmem_load_0[11] = (((frag_mask_word_0 >> 11 & 1) != 0) ? -CAKE_INF : _tmem_load_0[11]); + _tmem_load_0[12] = (((frag_mask_word_0 >> 12 & 1) != 0) ? -CAKE_INF : _tmem_load_0[12]); + _tmem_load_0[13] = (((frag_mask_word_0 >> 13 & 1) != 0) ? -CAKE_INF : _tmem_load_0[13]); + _tmem_load_0[14] = (((frag_mask_word_0 >> 14 & 1) != 0) ? -CAKE_INF : _tmem_load_0[14]); + _tmem_load_0[15] = (((frag_mask_word_0 >> 15 & 1) != 0) ? -CAKE_INF : _tmem_load_0[15]); + _tmem_load_0[16] = (((frag_mask_word_0 >> 16 & 1) != 0) ? -CAKE_INF : _tmem_load_0[16]); + _tmem_load_0[17] = (((frag_mask_word_0 >> 17 & 1) != 0) ? -CAKE_INF : _tmem_load_0[17]); + _tmem_load_0[18] = (((frag_mask_word_0 >> 18 & 1) != 0) ? -CAKE_INF : _tmem_load_0[18]); + _tmem_load_0[19] = (((frag_mask_word_0 >> 19 & 1) != 0) ? -CAKE_INF : _tmem_load_0[19]); + _tmem_load_0[20] = (((frag_mask_word_0 >> 20 & 1) != 0) ? -CAKE_INF : _tmem_load_0[20]); + _tmem_load_0[21] = (((frag_mask_word_0 >> 21 & 1) != 0) ? -CAKE_INF : _tmem_load_0[21]); + _tmem_load_0[22] = (((frag_mask_word_0 >> 22 & 1) != 0) ? -CAKE_INF : _tmem_load_0[22]); + _tmem_load_0[23] = (((frag_mask_word_0 >> 23 & 1) != 0) ? -CAKE_INF : _tmem_load_0[23]); + _tmem_load_0[24] = (((frag_mask_word_0 >> 24 & 1) != 0) ? -CAKE_INF : _tmem_load_0[24]); + _tmem_load_0[25] = (((frag_mask_word_0 >> 25 & 1) != 0) ? -CAKE_INF : _tmem_load_0[25]); + _tmem_load_0[26] = (((frag_mask_word_0 >> 26 & 1) != 0) ? -CAKE_INF : _tmem_load_0[26]); + _tmem_load_0[27] = (((frag_mask_word_0 >> 27 & 1) != 0) ? -CAKE_INF : _tmem_load_0[27]); + _tmem_load_0[28] = (((frag_mask_word_0 >> 28 & 1) != 0) ? -CAKE_INF : _tmem_load_0[28]); + _tmem_load_0[29] = (((frag_mask_word_0 >> 29 & 1) != 0) ? -CAKE_INF : _tmem_load_0[29]); + _tmem_load_0[30] = (((frag_mask_word_0 >> 30 & 1) != 0) ? -CAKE_INF : _tmem_load_0[30]); + _tmem_load_0[31] = (((frag_mask_word_0 >> 31 & 1) != 0) ? -CAKE_INF : _tmem_load_0[31]); + } + float2 _reg_reduce_max2_1 = {-CAKE_INF, -CAKE_INF}; + _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_0[0], _tmem_load_0[1])); + _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_0[2], _tmem_load_0[3])); + _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_0[4], _tmem_load_0[5])); + _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_0[6], _tmem_load_0[7])); + _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_0[8], _tmem_load_0[9])); + _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_0[10], _tmem_load_0[11])); + _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_0[12], _tmem_load_0[13])); + _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_0[14], _tmem_load_0[15])); + _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_0[16], _tmem_load_0[17])); + _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_0[18], _tmem_load_0[19])); + _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_0[20], _tmem_load_0[21])); + _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_0[22], _tmem_load_0[23])); + _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_0[24], _tmem_load_0[25])); + _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_0[26], _tmem_load_0[27])); + _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_0[28], _tmem_load_0[29])); + _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_0[30], _tmem_load_0[31])); + float _tmem_load_0_max = row_max_reduce(_reg_reduce_max2_1); + float _max_3 = max_noftz(new_max, _tmem_load_0_max); + new_max = _max_3; + int frag_valid_1 = valid_sparse_cols - 32; + uint32_t _slice_lo_mask_2; + { + int _lim_2 = frag_valid_1; + if (_lim_2 <= 0) { _slice_lo_mask_2 = 0u; } + else if (_lim_2 >= 32) { _slice_lo_mask_2 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_2) : "r"(_lim_2)); + } + } + if (!(_slice_lo_mask_2 & (1u << 0))) _tmem_load_2[0] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 1))) _tmem_load_2[1] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 2))) _tmem_load_2[2] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 3))) _tmem_load_2[3] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 4))) _tmem_load_2[4] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 5))) _tmem_load_2[5] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 6))) _tmem_load_2[6] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 7))) _tmem_load_2[7] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 8))) _tmem_load_2[8] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 9))) _tmem_load_2[9] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 10))) _tmem_load_2[10] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 11))) _tmem_load_2[11] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 12))) _tmem_load_2[12] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 13))) _tmem_load_2[13] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 14))) _tmem_load_2[14] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 15))) _tmem_load_2[15] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 16))) _tmem_load_2[16] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 17))) _tmem_load_2[17] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 18))) _tmem_load_2[18] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 19))) _tmem_load_2[19] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 20))) _tmem_load_2[20] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 21))) _tmem_load_2[21] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 22))) _tmem_load_2[22] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 23))) _tmem_load_2[23] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 24))) _tmem_load_2[24] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 25))) _tmem_load_2[25] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 26))) _tmem_load_2[26] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 27))) _tmem_load_2[27] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 28))) _tmem_load_2[28] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 29))) _tmem_load_2[29] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 30))) _tmem_load_2[30] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 31))) _tmem_load_2[31] = -CAKE_INF; + if (any_invalid != 0) { + unsigned int frag_mask_word_1 = invalid_cols[1]; + _tmem_load_2[0] = (((frag_mask_word_1 & 1) != 0) ? -CAKE_INF : _tmem_load_2[0]); + _tmem_load_2[1] = (((frag_mask_word_1 >> 1 & 1) != 0) ? -CAKE_INF : _tmem_load_2[1]); + _tmem_load_2[2] = (((frag_mask_word_1 >> 2 & 1) != 0) ? -CAKE_INF : _tmem_load_2[2]); + _tmem_load_2[3] = (((frag_mask_word_1 >> 3 & 1) != 0) ? -CAKE_INF : _tmem_load_2[3]); + _tmem_load_2[4] = (((frag_mask_word_1 >> 4 & 1) != 0) ? -CAKE_INF : _tmem_load_2[4]); + _tmem_load_2[5] = (((frag_mask_word_1 >> 5 & 1) != 0) ? -CAKE_INF : _tmem_load_2[5]); + _tmem_load_2[6] = (((frag_mask_word_1 >> 6 & 1) != 0) ? -CAKE_INF : _tmem_load_2[6]); + _tmem_load_2[7] = (((frag_mask_word_1 >> 7 & 1) != 0) ? -CAKE_INF : _tmem_load_2[7]); + _tmem_load_2[8] = (((frag_mask_word_1 >> 8 & 1) != 0) ? -CAKE_INF : _tmem_load_2[8]); + _tmem_load_2[9] = (((frag_mask_word_1 >> 9 & 1) != 0) ? -CAKE_INF : _tmem_load_2[9]); + _tmem_load_2[10] = (((frag_mask_word_1 >> 10 & 1) != 0) ? -CAKE_INF : _tmem_load_2[10]); + _tmem_load_2[11] = (((frag_mask_word_1 >> 11 & 1) != 0) ? -CAKE_INF : _tmem_load_2[11]); + _tmem_load_2[12] = (((frag_mask_word_1 >> 12 & 1) != 0) ? -CAKE_INF : _tmem_load_2[12]); + _tmem_load_2[13] = (((frag_mask_word_1 >> 13 & 1) != 0) ? -CAKE_INF : _tmem_load_2[13]); + _tmem_load_2[14] = (((frag_mask_word_1 >> 14 & 1) != 0) ? -CAKE_INF : _tmem_load_2[14]); + _tmem_load_2[15] = (((frag_mask_word_1 >> 15 & 1) != 0) ? -CAKE_INF : _tmem_load_2[15]); + _tmem_load_2[16] = (((frag_mask_word_1 >> 16 & 1) != 0) ? -CAKE_INF : _tmem_load_2[16]); + _tmem_load_2[17] = (((frag_mask_word_1 >> 17 & 1) != 0) ? -CAKE_INF : _tmem_load_2[17]); + _tmem_load_2[18] = (((frag_mask_word_1 >> 18 & 1) != 0) ? -CAKE_INF : _tmem_load_2[18]); + _tmem_load_2[19] = (((frag_mask_word_1 >> 19 & 1) != 0) ? -CAKE_INF : _tmem_load_2[19]); + _tmem_load_2[20] = (((frag_mask_word_1 >> 20 & 1) != 0) ? -CAKE_INF : _tmem_load_2[20]); + _tmem_load_2[21] = (((frag_mask_word_1 >> 21 & 1) != 0) ? -CAKE_INF : _tmem_load_2[21]); + _tmem_load_2[22] = (((frag_mask_word_1 >> 22 & 1) != 0) ? -CAKE_INF : _tmem_load_2[22]); + _tmem_load_2[23] = (((frag_mask_word_1 >> 23 & 1) != 0) ? -CAKE_INF : _tmem_load_2[23]); + _tmem_load_2[24] = (((frag_mask_word_1 >> 24 & 1) != 0) ? -CAKE_INF : _tmem_load_2[24]); + _tmem_load_2[25] = (((frag_mask_word_1 >> 25 & 1) != 0) ? -CAKE_INF : _tmem_load_2[25]); + _tmem_load_2[26] = (((frag_mask_word_1 >> 26 & 1) != 0) ? -CAKE_INF : _tmem_load_2[26]); + _tmem_load_2[27] = (((frag_mask_word_1 >> 27 & 1) != 0) ? -CAKE_INF : _tmem_load_2[27]); + _tmem_load_2[28] = (((frag_mask_word_1 >> 28 & 1) != 0) ? -CAKE_INF : _tmem_load_2[28]); + _tmem_load_2[29] = (((frag_mask_word_1 >> 29 & 1) != 0) ? -CAKE_INF : _tmem_load_2[29]); + _tmem_load_2[30] = (((frag_mask_word_1 >> 30 & 1) != 0) ? -CAKE_INF : _tmem_load_2[30]); + _tmem_load_2[31] = (((frag_mask_word_1 >> 31 & 1) != 0) ? -CAKE_INF : _tmem_load_2[31]); + } + float2 _reg_reduce_max2_3 = {-CAKE_INF, -CAKE_INF}; + _reg_reduce_max2_3.x = max_noftz(_reg_reduce_max2_3.x, max_noftz(_tmem_load_2[0], _tmem_load_2[1])); + _reg_reduce_max2_3.y = max_noftz(_reg_reduce_max2_3.y, max_noftz(_tmem_load_2[2], _tmem_load_2[3])); + _reg_reduce_max2_3.x = max_noftz(_reg_reduce_max2_3.x, max_noftz(_tmem_load_2[4], _tmem_load_2[5])); + _reg_reduce_max2_3.y = max_noftz(_reg_reduce_max2_3.y, max_noftz(_tmem_load_2[6], _tmem_load_2[7])); + _reg_reduce_max2_3.x = max_noftz(_reg_reduce_max2_3.x, max_noftz(_tmem_load_2[8], _tmem_load_2[9])); + _reg_reduce_max2_3.y = max_noftz(_reg_reduce_max2_3.y, max_noftz(_tmem_load_2[10], _tmem_load_2[11])); + _reg_reduce_max2_3.x = max_noftz(_reg_reduce_max2_3.x, max_noftz(_tmem_load_2[12], _tmem_load_2[13])); + _reg_reduce_max2_3.y = max_noftz(_reg_reduce_max2_3.y, max_noftz(_tmem_load_2[14], _tmem_load_2[15])); + _reg_reduce_max2_3.x = max_noftz(_reg_reduce_max2_3.x, max_noftz(_tmem_load_2[16], _tmem_load_2[17])); + _reg_reduce_max2_3.y = max_noftz(_reg_reduce_max2_3.y, max_noftz(_tmem_load_2[18], _tmem_load_2[19])); + _reg_reduce_max2_3.x = max_noftz(_reg_reduce_max2_3.x, max_noftz(_tmem_load_2[20], _tmem_load_2[21])); + _reg_reduce_max2_3.y = max_noftz(_reg_reduce_max2_3.y, max_noftz(_tmem_load_2[22], _tmem_load_2[23])); + _reg_reduce_max2_3.x = max_noftz(_reg_reduce_max2_3.x, max_noftz(_tmem_load_2[24], _tmem_load_2[25])); + _reg_reduce_max2_3.y = max_noftz(_reg_reduce_max2_3.y, max_noftz(_tmem_load_2[26], _tmem_load_2[27])); + _reg_reduce_max2_3.x = max_noftz(_reg_reduce_max2_3.x, max_noftz(_tmem_load_2[28], _tmem_load_2[29])); + _reg_reduce_max2_3.y = max_noftz(_reg_reduce_max2_3.y, max_noftz(_tmem_load_2[30], _tmem_load_2[31])); + float _tmem_load_2_max = row_max_reduce(_reg_reduce_max2_3); + float _max_4 = max_noftz(new_max, _tmem_load_2_max); + new_max = _max_4; + int frag_valid_2 = valid_sparse_cols - 64; + uint32_t _slice_lo_mask_3; + { + int _lim_4 = frag_valid_2; + if (_lim_4 <= 0) { _slice_lo_mask_3 = 0u; } + else if (_lim_4 >= 32) { _slice_lo_mask_3 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_3) : "r"(_lim_4)); + } + } + if (!(_slice_lo_mask_3 & (1u << 0))) _tmem_load_3[0] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 1))) _tmem_load_3[1] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 2))) _tmem_load_3[2] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 3))) _tmem_load_3[3] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 4))) _tmem_load_3[4] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 5))) _tmem_load_3[5] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 6))) _tmem_load_3[6] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 7))) _tmem_load_3[7] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 8))) _tmem_load_3[8] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 9))) _tmem_load_3[9] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 10))) _tmem_load_3[10] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 11))) _tmem_load_3[11] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 12))) _tmem_load_3[12] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 13))) _tmem_load_3[13] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 14))) _tmem_load_3[14] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 15))) _tmem_load_3[15] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 16))) _tmem_load_3[16] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 17))) _tmem_load_3[17] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 18))) _tmem_load_3[18] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 19))) _tmem_load_3[19] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 20))) _tmem_load_3[20] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 21))) _tmem_load_3[21] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 22))) _tmem_load_3[22] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 23))) _tmem_load_3[23] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 24))) _tmem_load_3[24] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 25))) _tmem_load_3[25] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 26))) _tmem_load_3[26] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 27))) _tmem_load_3[27] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 28))) _tmem_load_3[28] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 29))) _tmem_load_3[29] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 30))) _tmem_load_3[30] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 31))) _tmem_load_3[31] = -CAKE_INF; + if (any_invalid != 0) { + unsigned int frag_mask_word_2 = invalid_cols[2]; + _tmem_load_3[0] = (((frag_mask_word_2 & 1) != 0) ? -CAKE_INF : _tmem_load_3[0]); + _tmem_load_3[1] = (((frag_mask_word_2 >> 1 & 1) != 0) ? -CAKE_INF : _tmem_load_3[1]); + _tmem_load_3[2] = (((frag_mask_word_2 >> 2 & 1) != 0) ? -CAKE_INF : _tmem_load_3[2]); + _tmem_load_3[3] = (((frag_mask_word_2 >> 3 & 1) != 0) ? -CAKE_INF : _tmem_load_3[3]); + _tmem_load_3[4] = (((frag_mask_word_2 >> 4 & 1) != 0) ? -CAKE_INF : _tmem_load_3[4]); + _tmem_load_3[5] = (((frag_mask_word_2 >> 5 & 1) != 0) ? -CAKE_INF : _tmem_load_3[5]); + _tmem_load_3[6] = (((frag_mask_word_2 >> 6 & 1) != 0) ? -CAKE_INF : _tmem_load_3[6]); + _tmem_load_3[7] = (((frag_mask_word_2 >> 7 & 1) != 0) ? -CAKE_INF : _tmem_load_3[7]); + _tmem_load_3[8] = (((frag_mask_word_2 >> 8 & 1) != 0) ? -CAKE_INF : _tmem_load_3[8]); + _tmem_load_3[9] = (((frag_mask_word_2 >> 9 & 1) != 0) ? -CAKE_INF : _tmem_load_3[9]); + _tmem_load_3[10] = (((frag_mask_word_2 >> 10 & 1) != 0) ? -CAKE_INF : _tmem_load_3[10]); + _tmem_load_3[11] = (((frag_mask_word_2 >> 11 & 1) != 0) ? -CAKE_INF : _tmem_load_3[11]); + _tmem_load_3[12] = (((frag_mask_word_2 >> 12 & 1) != 0) ? -CAKE_INF : _tmem_load_3[12]); + _tmem_load_3[13] = (((frag_mask_word_2 >> 13 & 1) != 0) ? -CAKE_INF : _tmem_load_3[13]); + _tmem_load_3[14] = (((frag_mask_word_2 >> 14 & 1) != 0) ? -CAKE_INF : _tmem_load_3[14]); + _tmem_load_3[15] = (((frag_mask_word_2 >> 15 & 1) != 0) ? -CAKE_INF : _tmem_load_3[15]); + _tmem_load_3[16] = (((frag_mask_word_2 >> 16 & 1) != 0) ? -CAKE_INF : _tmem_load_3[16]); + _tmem_load_3[17] = (((frag_mask_word_2 >> 17 & 1) != 0) ? -CAKE_INF : _tmem_load_3[17]); + _tmem_load_3[18] = (((frag_mask_word_2 >> 18 & 1) != 0) ? -CAKE_INF : _tmem_load_3[18]); + _tmem_load_3[19] = (((frag_mask_word_2 >> 19 & 1) != 0) ? -CAKE_INF : _tmem_load_3[19]); + _tmem_load_3[20] = (((frag_mask_word_2 >> 20 & 1) != 0) ? -CAKE_INF : _tmem_load_3[20]); + _tmem_load_3[21] = (((frag_mask_word_2 >> 21 & 1) != 0) ? -CAKE_INF : _tmem_load_3[21]); + _tmem_load_3[22] = (((frag_mask_word_2 >> 22 & 1) != 0) ? -CAKE_INF : _tmem_load_3[22]); + _tmem_load_3[23] = (((frag_mask_word_2 >> 23 & 1) != 0) ? -CAKE_INF : _tmem_load_3[23]); + _tmem_load_3[24] = (((frag_mask_word_2 >> 24 & 1) != 0) ? -CAKE_INF : _tmem_load_3[24]); + _tmem_load_3[25] = (((frag_mask_word_2 >> 25 & 1) != 0) ? -CAKE_INF : _tmem_load_3[25]); + _tmem_load_3[26] = (((frag_mask_word_2 >> 26 & 1) != 0) ? -CAKE_INF : _tmem_load_3[26]); + _tmem_load_3[27] = (((frag_mask_word_2 >> 27 & 1) != 0) ? -CAKE_INF : _tmem_load_3[27]); + _tmem_load_3[28] = (((frag_mask_word_2 >> 28 & 1) != 0) ? -CAKE_INF : _tmem_load_3[28]); + _tmem_load_3[29] = (((frag_mask_word_2 >> 29 & 1) != 0) ? -CAKE_INF : _tmem_load_3[29]); + _tmem_load_3[30] = (((frag_mask_word_2 >> 30 & 1) != 0) ? -CAKE_INF : _tmem_load_3[30]); + _tmem_load_3[31] = (((frag_mask_word_2 >> 31 & 1) != 0) ? -CAKE_INF : _tmem_load_3[31]); + } + float2 _reg_reduce_max2_5 = {-CAKE_INF, -CAKE_INF}; + _reg_reduce_max2_5.x = max_noftz(_reg_reduce_max2_5.x, max_noftz(_tmem_load_3[0], _tmem_load_3[1])); + _reg_reduce_max2_5.y = max_noftz(_reg_reduce_max2_5.y, max_noftz(_tmem_load_3[2], _tmem_load_3[3])); + _reg_reduce_max2_5.x = max_noftz(_reg_reduce_max2_5.x, max_noftz(_tmem_load_3[4], _tmem_load_3[5])); + _reg_reduce_max2_5.y = max_noftz(_reg_reduce_max2_5.y, max_noftz(_tmem_load_3[6], _tmem_load_3[7])); + _reg_reduce_max2_5.x = max_noftz(_reg_reduce_max2_5.x, max_noftz(_tmem_load_3[8], _tmem_load_3[9])); + _reg_reduce_max2_5.y = max_noftz(_reg_reduce_max2_5.y, max_noftz(_tmem_load_3[10], _tmem_load_3[11])); + _reg_reduce_max2_5.x = max_noftz(_reg_reduce_max2_5.x, max_noftz(_tmem_load_3[12], _tmem_load_3[13])); + _reg_reduce_max2_5.y = max_noftz(_reg_reduce_max2_5.y, max_noftz(_tmem_load_3[14], _tmem_load_3[15])); + _reg_reduce_max2_5.x = max_noftz(_reg_reduce_max2_5.x, max_noftz(_tmem_load_3[16], _tmem_load_3[17])); + _reg_reduce_max2_5.y = max_noftz(_reg_reduce_max2_5.y, max_noftz(_tmem_load_3[18], _tmem_load_3[19])); + _reg_reduce_max2_5.x = max_noftz(_reg_reduce_max2_5.x, max_noftz(_tmem_load_3[20], _tmem_load_3[21])); + _reg_reduce_max2_5.y = max_noftz(_reg_reduce_max2_5.y, max_noftz(_tmem_load_3[22], _tmem_load_3[23])); + _reg_reduce_max2_5.x = max_noftz(_reg_reduce_max2_5.x, max_noftz(_tmem_load_3[24], _tmem_load_3[25])); + _reg_reduce_max2_5.y = max_noftz(_reg_reduce_max2_5.y, max_noftz(_tmem_load_3[26], _tmem_load_3[27])); + _reg_reduce_max2_5.x = max_noftz(_reg_reduce_max2_5.x, max_noftz(_tmem_load_3[28], _tmem_load_3[29])); + _reg_reduce_max2_5.y = max_noftz(_reg_reduce_max2_5.y, max_noftz(_tmem_load_3[30], _tmem_load_3[31])); + float _tmem_load_3_max = row_max_reduce(_reg_reduce_max2_5); + float _max_5 = max_noftz(new_max, _tmem_load_3_max); + new_max = _max_5; + int frag_valid_3 = valid_sparse_cols - 96; + uint32_t _slice_lo_mask_4; + { + int _lim_6 = frag_valid_3; + if (_lim_6 <= 0) { _slice_lo_mask_4 = 0u; } + else if (_lim_6 >= 32) { _slice_lo_mask_4 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_4) : "r"(_lim_6)); } - float2 _reg_reduce_max2_1 = {-CAKE_INF, -CAKE_INF}; - _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_2[0], _tmem_load_2[1])); - _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_2[2], _tmem_load_2[3])); - _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_2[4], _tmem_load_2[5])); - _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_2[6], _tmem_load_2[7])); - _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_2[8], _tmem_load_2[9])); - _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_2[10], _tmem_load_2[11])); - _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_2[12], _tmem_load_2[13])); - _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_2[14], _tmem_load_2[15])); - _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_2[16], _tmem_load_2[17])); - _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_2[18], _tmem_load_2[19])); - _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_2[20], _tmem_load_2[21])); - _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_2[22], _tmem_load_2[23])); - _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_2[24], _tmem_load_2[25])); - _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_2[26], _tmem_load_2[27])); - _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_2[28], _tmem_load_2[29])); - _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_2[30], _tmem_load_2[31])); - float _tmem_load_2_max = row_max_reduce(_reg_reduce_max2_1); - float _max_3 = max_noftz(new_max, _tmem_load_2_max); - new_max = _max_3; } + if (!(_slice_lo_mask_4 & (1u << 0))) _tmem_load_4[0] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 1))) _tmem_load_4[1] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 2))) _tmem_load_4[2] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 3))) _tmem_load_4[3] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 4))) _tmem_load_4[4] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 5))) _tmem_load_4[5] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 6))) _tmem_load_4[6] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 7))) _tmem_load_4[7] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 8))) _tmem_load_4[8] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 9))) _tmem_load_4[9] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 10))) _tmem_load_4[10] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 11))) _tmem_load_4[11] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 12))) _tmem_load_4[12] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 13))) _tmem_load_4[13] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 14))) _tmem_load_4[14] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 15))) _tmem_load_4[15] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 16))) _tmem_load_4[16] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 17))) _tmem_load_4[17] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 18))) _tmem_load_4[18] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 19))) _tmem_load_4[19] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 20))) _tmem_load_4[20] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 21))) _tmem_load_4[21] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 22))) _tmem_load_4[22] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 23))) _tmem_load_4[23] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 24))) _tmem_load_4[24] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 25))) _tmem_load_4[25] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 26))) _tmem_load_4[26] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 27))) _tmem_load_4[27] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 28))) _tmem_load_4[28] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 29))) _tmem_load_4[29] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 30))) _tmem_load_4[30] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 31))) _tmem_load_4[31] = -CAKE_INF; + if (any_invalid != 0) { + unsigned int frag_mask_word_3 = invalid_cols[3]; + _tmem_load_4[0] = (((frag_mask_word_3 & 1) != 0) ? -CAKE_INF : _tmem_load_4[0]); + _tmem_load_4[1] = (((frag_mask_word_3 >> 1 & 1) != 0) ? -CAKE_INF : _tmem_load_4[1]); + _tmem_load_4[2] = (((frag_mask_word_3 >> 2 & 1) != 0) ? -CAKE_INF : _tmem_load_4[2]); + _tmem_load_4[3] = (((frag_mask_word_3 >> 3 & 1) != 0) ? -CAKE_INF : _tmem_load_4[3]); + _tmem_load_4[4] = (((frag_mask_word_3 >> 4 & 1) != 0) ? -CAKE_INF : _tmem_load_4[4]); + _tmem_load_4[5] = (((frag_mask_word_3 >> 5 & 1) != 0) ? -CAKE_INF : _tmem_load_4[5]); + _tmem_load_4[6] = (((frag_mask_word_3 >> 6 & 1) != 0) ? -CAKE_INF : _tmem_load_4[6]); + _tmem_load_4[7] = (((frag_mask_word_3 >> 7 & 1) != 0) ? -CAKE_INF : _tmem_load_4[7]); + _tmem_load_4[8] = (((frag_mask_word_3 >> 8 & 1) != 0) ? -CAKE_INF : _tmem_load_4[8]); + _tmem_load_4[9] = (((frag_mask_word_3 >> 9 & 1) != 0) ? -CAKE_INF : _tmem_load_4[9]); + _tmem_load_4[10] = (((frag_mask_word_3 >> 10 & 1) != 0) ? -CAKE_INF : _tmem_load_4[10]); + _tmem_load_4[11] = (((frag_mask_word_3 >> 11 & 1) != 0) ? -CAKE_INF : _tmem_load_4[11]); + _tmem_load_4[12] = (((frag_mask_word_3 >> 12 & 1) != 0) ? -CAKE_INF : _tmem_load_4[12]); + _tmem_load_4[13] = (((frag_mask_word_3 >> 13 & 1) != 0) ? -CAKE_INF : _tmem_load_4[13]); + _tmem_load_4[14] = (((frag_mask_word_3 >> 14 & 1) != 0) ? -CAKE_INF : _tmem_load_4[14]); + _tmem_load_4[15] = (((frag_mask_word_3 >> 15 & 1) != 0) ? -CAKE_INF : _tmem_load_4[15]); + _tmem_load_4[16] = (((frag_mask_word_3 >> 16 & 1) != 0) ? -CAKE_INF : _tmem_load_4[16]); + _tmem_load_4[17] = (((frag_mask_word_3 >> 17 & 1) != 0) ? -CAKE_INF : _tmem_load_4[17]); + _tmem_load_4[18] = (((frag_mask_word_3 >> 18 & 1) != 0) ? -CAKE_INF : _tmem_load_4[18]); + _tmem_load_4[19] = (((frag_mask_word_3 >> 19 & 1) != 0) ? -CAKE_INF : _tmem_load_4[19]); + _tmem_load_4[20] = (((frag_mask_word_3 >> 20 & 1) != 0) ? -CAKE_INF : _tmem_load_4[20]); + _tmem_load_4[21] = (((frag_mask_word_3 >> 21 & 1) != 0) ? -CAKE_INF : _tmem_load_4[21]); + _tmem_load_4[22] = (((frag_mask_word_3 >> 22 & 1) != 0) ? -CAKE_INF : _tmem_load_4[22]); + _tmem_load_4[23] = (((frag_mask_word_3 >> 23 & 1) != 0) ? -CAKE_INF : _tmem_load_4[23]); + _tmem_load_4[24] = (((frag_mask_word_3 >> 24 & 1) != 0) ? -CAKE_INF : _tmem_load_4[24]); + _tmem_load_4[25] = (((frag_mask_word_3 >> 25 & 1) != 0) ? -CAKE_INF : _tmem_load_4[25]); + _tmem_load_4[26] = (((frag_mask_word_3 >> 26 & 1) != 0) ? -CAKE_INF : _tmem_load_4[26]); + _tmem_load_4[27] = (((frag_mask_word_3 >> 27 & 1) != 0) ? -CAKE_INF : _tmem_load_4[27]); + _tmem_load_4[28] = (((frag_mask_word_3 >> 28 & 1) != 0) ? -CAKE_INF : _tmem_load_4[28]); + _tmem_load_4[29] = (((frag_mask_word_3 >> 29 & 1) != 0) ? -CAKE_INF : _tmem_load_4[29]); + _tmem_load_4[30] = (((frag_mask_word_3 >> 30 & 1) != 0) ? -CAKE_INF : _tmem_load_4[30]); + _tmem_load_4[31] = (((frag_mask_word_3 >> 31 & 1) != 0) ? -CAKE_INF : _tmem_load_4[31]); + } + float2 _reg_reduce_max2_7 = {-CAKE_INF, -CAKE_INF}; + _reg_reduce_max2_7.x = max_noftz(_reg_reduce_max2_7.x, max_noftz(_tmem_load_4[0], _tmem_load_4[1])); + _reg_reduce_max2_7.y = max_noftz(_reg_reduce_max2_7.y, max_noftz(_tmem_load_4[2], _tmem_load_4[3])); + _reg_reduce_max2_7.x = max_noftz(_reg_reduce_max2_7.x, max_noftz(_tmem_load_4[4], _tmem_load_4[5])); + _reg_reduce_max2_7.y = max_noftz(_reg_reduce_max2_7.y, max_noftz(_tmem_load_4[6], _tmem_load_4[7])); + _reg_reduce_max2_7.x = max_noftz(_reg_reduce_max2_7.x, max_noftz(_tmem_load_4[8], _tmem_load_4[9])); + _reg_reduce_max2_7.y = max_noftz(_reg_reduce_max2_7.y, max_noftz(_tmem_load_4[10], _tmem_load_4[11])); + _reg_reduce_max2_7.x = max_noftz(_reg_reduce_max2_7.x, max_noftz(_tmem_load_4[12], _tmem_load_4[13])); + _reg_reduce_max2_7.y = max_noftz(_reg_reduce_max2_7.y, max_noftz(_tmem_load_4[14], _tmem_load_4[15])); + _reg_reduce_max2_7.x = max_noftz(_reg_reduce_max2_7.x, max_noftz(_tmem_load_4[16], _tmem_load_4[17])); + _reg_reduce_max2_7.y = max_noftz(_reg_reduce_max2_7.y, max_noftz(_tmem_load_4[18], _tmem_load_4[19])); + _reg_reduce_max2_7.x = max_noftz(_reg_reduce_max2_7.x, max_noftz(_tmem_load_4[20], _tmem_load_4[21])); + _reg_reduce_max2_7.y = max_noftz(_reg_reduce_max2_7.y, max_noftz(_tmem_load_4[22], _tmem_load_4[23])); + _reg_reduce_max2_7.x = max_noftz(_reg_reduce_max2_7.x, max_noftz(_tmem_load_4[24], _tmem_load_4[25])); + _reg_reduce_max2_7.y = max_noftz(_reg_reduce_max2_7.y, max_noftz(_tmem_load_4[26], _tmem_load_4[27])); + _reg_reduce_max2_7.x = max_noftz(_reg_reduce_max2_7.x, max_noftz(_tmem_load_4[28], _tmem_load_4[29])); + _reg_reduce_max2_7.y = max_noftz(_reg_reduce_max2_7.y, max_noftz(_tmem_load_4[30], _tmem_load_4[31])); + float _tmem_load_4_max = row_max_reduce(_reg_reduce_max2_7); + float _max_6 = max_noftz(new_max, _tmem_load_4_max); + new_max = _max_6; } float _fma_0 = __fmaf_rn(row_max_val, softmax_scale_log2, (-new_max) * softmax_scale_log2); float delta = _fma_0; @@ -1833,187 +2173,367 @@ kernel_cake_dsv4_a2be51a1fd34ec0d9a3b(CakeTensorMap const* tmap_q, CakeTensorMap float max_scaled = safe_max * softmax_scale_log2; float block_sum = 0.0f; { - int p_off = ((phase != 0) ? 224 : 96); - int p_base = taddr + (unsigned int)p_off + (unsigned int)(tmem_row_base << 16); + int p_off_all = ((phase != 0) ? 224 : 96); + int p_base_all = taddr + (unsigned int)p_off_all + (unsigned int)(tmem_row_base << 16); + const float2 _fma_b2_8 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_9 = {-max_scaled, -max_scaled}; + float2 _fma_pair_10 = fma_f32x2(make_float2(_tmem_load_0[0], _tmem_load_0[1]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[0] = _fma_pair_10.x; + _tmem_load_0[1] = _fma_pair_10.y; + float2 _fma_pair_11 = fma_f32x2(make_float2(_tmem_load_0[2], _tmem_load_0[3]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[2] = _fma_pair_11.x; + _tmem_load_0[3] = _fma_pair_11.y; + float2 _fma_pair_12 = fma_f32x2(make_float2(_tmem_load_0[4], _tmem_load_0[5]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[4] = _fma_pair_12.x; + _tmem_load_0[5] = _fma_pair_12.y; + float2 _fma_pair_13 = fma_f32x2(make_float2(_tmem_load_0[6], _tmem_load_0[7]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[6] = _fma_pair_13.x; + _tmem_load_0[7] = _fma_pair_13.y; + float2 _fma_pair_14 = fma_f32x2(make_float2(_tmem_load_0[8], _tmem_load_0[9]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[8] = _fma_pair_14.x; + _tmem_load_0[9] = _fma_pair_14.y; + float2 _fma_pair_15 = fma_f32x2(make_float2(_tmem_load_0[10], _tmem_load_0[11]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[10] = _fma_pair_15.x; + _tmem_load_0[11] = _fma_pair_15.y; + float2 _fma_pair_16 = fma_f32x2(make_float2(_tmem_load_0[12], _tmem_load_0[13]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[12] = _fma_pair_16.x; + _tmem_load_0[13] = _fma_pair_16.y; + float2 _fma_pair_17 = fma_f32x2(make_float2(_tmem_load_0[14], _tmem_load_0[15]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[14] = _fma_pair_17.x; + _tmem_load_0[15] = _fma_pair_17.y; + float2 _fma_pair_18 = fma_f32x2(make_float2(_tmem_load_0[16], _tmem_load_0[17]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[16] = _fma_pair_18.x; + _tmem_load_0[17] = _fma_pair_18.y; + float2 _fma_pair_19 = fma_f32x2(make_float2(_tmem_load_0[18], _tmem_load_0[19]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[18] = _fma_pair_19.x; + _tmem_load_0[19] = _fma_pair_19.y; + float2 _fma_pair_20 = fma_f32x2(make_float2(_tmem_load_0[20], _tmem_load_0[21]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[20] = _fma_pair_20.x; + _tmem_load_0[21] = _fma_pair_20.y; + float2 _fma_pair_21 = fma_f32x2(make_float2(_tmem_load_0[22], _tmem_load_0[23]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[22] = _fma_pair_21.x; + _tmem_load_0[23] = _fma_pair_21.y; + float2 _fma_pair_22 = fma_f32x2(make_float2(_tmem_load_0[24], _tmem_load_0[25]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[24] = _fma_pair_22.x; + _tmem_load_0[25] = _fma_pair_22.y; + float2 _fma_pair_23 = fma_f32x2(make_float2(_tmem_load_0[26], _tmem_load_0[27]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[26] = _fma_pair_23.x; + _tmem_load_0[27] = _fma_pair_23.y; + float2 _fma_pair_24 = fma_f32x2(make_float2(_tmem_load_0[28], _tmem_load_0[29]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[28] = _fma_pair_24.x; + _tmem_load_0[29] = _fma_pair_24.y; + float2 _fma_pair_25 = fma_f32x2(make_float2(_tmem_load_0[30], _tmem_load_0[31]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[30] = _fma_pair_25.x; + _tmem_load_0[31] = _fma_pair_25.y; #pragma unroll - for (int frag_1 = 0; frag_1 < 4; frag_1++) { - float _tmem_load_3[32]; - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]) - : "r"(s_base + frag_1 * 32)); - int frag_valid_1 = valid_sparse_cols - frag_1 * 32; - uint32_t _slice_lo_mask_3; - { - int _lim_2 = frag_valid_1; - if (_lim_2 <= 0) { _slice_lo_mask_3 = 0u; } - else if (_lim_2 >= 32) { _slice_lo_mask_3 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_3) : "r"(_lim_2)); - } + for (int _le = 0; _le < 32; _le++) { + _tmem_load_0[_le] = approx_exp2(_tmem_load_0[_le]); + } + float2 _reg_reduce_sum2_26 = make_float2(0.0f, 0.0f); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[0], _tmem_load_0[1])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[2], _tmem_load_0[3])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[4], _tmem_load_0[5])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[6], _tmem_load_0[7])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[8], _tmem_load_0[9])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[10], _tmem_load_0[11])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[12], _tmem_load_0[13])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[14], _tmem_load_0[15])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[16], _tmem_load_0[17])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[18], _tmem_load_0[19])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[20], _tmem_load_0[21])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[22], _tmem_load_0[23])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[24], _tmem_load_0[25])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[26], _tmem_load_0[27])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[28], _tmem_load_0[29])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[30], _tmem_load_0[31])); + float _tmem_load_0_sum = _reg_reduce_sum2_26.x + _reg_reduce_sum2_26.y; + block_sum = block_sum + _tmem_load_0_sum; + { + uint32_t _pv_packed[8]; + #pragma unroll + for (int _j = 0; _j < 8; _j++) { + uint32_t _pk; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_pk) : "f"(_tmem_load_0[0 + _j * 4]), "f"(_tmem_load_0[0 + _j * 4 + 1]), + "f"(_tmem_load_0[0 + _j * 4 + 2]), "f"(_tmem_load_0[0 + _j * 4 + 3])); + _pv_packed[_j] = _pk; } - if (!(_slice_lo_mask_3 & (1u << 0))) _tmem_load_3[0] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 1))) _tmem_load_3[1] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 2))) _tmem_load_3[2] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 3))) _tmem_load_3[3] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 4))) _tmem_load_3[4] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 5))) _tmem_load_3[5] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 6))) _tmem_load_3[6] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 7))) _tmem_load_3[7] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 8))) _tmem_load_3[8] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 9))) _tmem_load_3[9] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 10))) _tmem_load_3[10] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 11))) _tmem_load_3[11] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 12))) _tmem_load_3[12] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 13))) _tmem_load_3[13] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 14))) _tmem_load_3[14] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 15))) _tmem_load_3[15] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 16))) _tmem_load_3[16] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 17))) _tmem_load_3[17] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 18))) _tmem_load_3[18] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 19))) _tmem_load_3[19] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 20))) _tmem_load_3[20] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 21))) _tmem_load_3[21] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 22))) _tmem_load_3[22] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 23))) _tmem_load_3[23] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 24))) _tmem_load_3[24] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 25))) _tmem_load_3[25] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 26))) _tmem_load_3[26] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 27))) _tmem_load_3[27] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 28))) _tmem_load_3[28] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 29))) _tmem_load_3[29] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 30))) _tmem_load_3[30] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 31))) _tmem_load_3[31] = -CAKE_INF; - if (any_invalid != 0) { - unsigned int frag_mask_word_1 = invalid_cols[frag_1]; - _tmem_load_3[0] = (((frag_mask_word_1 & 1) != 0) ? -CAKE_INF : _tmem_load_3[0]); - _tmem_load_3[1] = (((frag_mask_word_1 >> 1 & 1) != 0) ? -CAKE_INF : _tmem_load_3[1]); - _tmem_load_3[2] = (((frag_mask_word_1 >> 2 & 1) != 0) ? -CAKE_INF : _tmem_load_3[2]); - _tmem_load_3[3] = (((frag_mask_word_1 >> 3 & 1) != 0) ? -CAKE_INF : _tmem_load_3[3]); - _tmem_load_3[4] = (((frag_mask_word_1 >> 4 & 1) != 0) ? -CAKE_INF : _tmem_load_3[4]); - _tmem_load_3[5] = (((frag_mask_word_1 >> 5 & 1) != 0) ? -CAKE_INF : _tmem_load_3[5]); - _tmem_load_3[6] = (((frag_mask_word_1 >> 6 & 1) != 0) ? -CAKE_INF : _tmem_load_3[6]); - _tmem_load_3[7] = (((frag_mask_word_1 >> 7 & 1) != 0) ? -CAKE_INF : _tmem_load_3[7]); - _tmem_load_3[8] = (((frag_mask_word_1 >> 8 & 1) != 0) ? -CAKE_INF : _tmem_load_3[8]); - _tmem_load_3[9] = (((frag_mask_word_1 >> 9 & 1) != 0) ? -CAKE_INF : _tmem_load_3[9]); - _tmem_load_3[10] = (((frag_mask_word_1 >> 10 & 1) != 0) ? -CAKE_INF : _tmem_load_3[10]); - _tmem_load_3[11] = (((frag_mask_word_1 >> 11 & 1) != 0) ? -CAKE_INF : _tmem_load_3[11]); - _tmem_load_3[12] = (((frag_mask_word_1 >> 12 & 1) != 0) ? -CAKE_INF : _tmem_load_3[12]); - _tmem_load_3[13] = (((frag_mask_word_1 >> 13 & 1) != 0) ? -CAKE_INF : _tmem_load_3[13]); - _tmem_load_3[14] = (((frag_mask_word_1 >> 14 & 1) != 0) ? -CAKE_INF : _tmem_load_3[14]); - _tmem_load_3[15] = (((frag_mask_word_1 >> 15 & 1) != 0) ? -CAKE_INF : _tmem_load_3[15]); - _tmem_load_3[16] = (((frag_mask_word_1 >> 16 & 1) != 0) ? -CAKE_INF : _tmem_load_3[16]); - _tmem_load_3[17] = (((frag_mask_word_1 >> 17 & 1) != 0) ? -CAKE_INF : _tmem_load_3[17]); - _tmem_load_3[18] = (((frag_mask_word_1 >> 18 & 1) != 0) ? -CAKE_INF : _tmem_load_3[18]); - _tmem_load_3[19] = (((frag_mask_word_1 >> 19 & 1) != 0) ? -CAKE_INF : _tmem_load_3[19]); - _tmem_load_3[20] = (((frag_mask_word_1 >> 20 & 1) != 0) ? -CAKE_INF : _tmem_load_3[20]); - _tmem_load_3[21] = (((frag_mask_word_1 >> 21 & 1) != 0) ? -CAKE_INF : _tmem_load_3[21]); - _tmem_load_3[22] = (((frag_mask_word_1 >> 22 & 1) != 0) ? -CAKE_INF : _tmem_load_3[22]); - _tmem_load_3[23] = (((frag_mask_word_1 >> 23 & 1) != 0) ? -CAKE_INF : _tmem_load_3[23]); - _tmem_load_3[24] = (((frag_mask_word_1 >> 24 & 1) != 0) ? -CAKE_INF : _tmem_load_3[24]); - _tmem_load_3[25] = (((frag_mask_word_1 >> 25 & 1) != 0) ? -CAKE_INF : _tmem_load_3[25]); - _tmem_load_3[26] = (((frag_mask_word_1 >> 26 & 1) != 0) ? -CAKE_INF : _tmem_load_3[26]); - _tmem_load_3[27] = (((frag_mask_word_1 >> 27 & 1) != 0) ? -CAKE_INF : _tmem_load_3[27]); - _tmem_load_3[28] = (((frag_mask_word_1 >> 28 & 1) != 0) ? -CAKE_INF : _tmem_load_3[28]); - _tmem_load_3[29] = (((frag_mask_word_1 >> 29 & 1) != 0) ? -CAKE_INF : _tmem_load_3[29]); - _tmem_load_3[30] = (((frag_mask_word_1 >> 30 & 1) != 0) ? -CAKE_INF : _tmem_load_3[30]); - _tmem_load_3[31] = (((frag_mask_word_1 >> 31 & 1) != 0) ? -CAKE_INF : _tmem_load_3[31]); + tmem_st_x8_u32(p_base_all, _pv_packed); + } + const float2 _fma_b2_27 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_28 = {-max_scaled, -max_scaled}; + float2 _fma_pair_29 = fma_f32x2(make_float2(_tmem_load_2[0], _tmem_load_2[1]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[0] = _fma_pair_29.x; + _tmem_load_2[1] = _fma_pair_29.y; + float2 _fma_pair_30 = fma_f32x2(make_float2(_tmem_load_2[2], _tmem_load_2[3]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[2] = _fma_pair_30.x; + _tmem_load_2[3] = _fma_pair_30.y; + float2 _fma_pair_31 = fma_f32x2(make_float2(_tmem_load_2[4], _tmem_load_2[5]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[4] = _fma_pair_31.x; + _tmem_load_2[5] = _fma_pair_31.y; + float2 _fma_pair_32 = fma_f32x2(make_float2(_tmem_load_2[6], _tmem_load_2[7]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[6] = _fma_pair_32.x; + _tmem_load_2[7] = _fma_pair_32.y; + float2 _fma_pair_33 = fma_f32x2(make_float2(_tmem_load_2[8], _tmem_load_2[9]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[8] = _fma_pair_33.x; + _tmem_load_2[9] = _fma_pair_33.y; + float2 _fma_pair_34 = fma_f32x2(make_float2(_tmem_load_2[10], _tmem_load_2[11]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[10] = _fma_pair_34.x; + _tmem_load_2[11] = _fma_pair_34.y; + float2 _fma_pair_35 = fma_f32x2(make_float2(_tmem_load_2[12], _tmem_load_2[13]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[12] = _fma_pair_35.x; + _tmem_load_2[13] = _fma_pair_35.y; + float2 _fma_pair_36 = fma_f32x2(make_float2(_tmem_load_2[14], _tmem_load_2[15]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[14] = _fma_pair_36.x; + _tmem_load_2[15] = _fma_pair_36.y; + float2 _fma_pair_37 = fma_f32x2(make_float2(_tmem_load_2[16], _tmem_load_2[17]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[16] = _fma_pair_37.x; + _tmem_load_2[17] = _fma_pair_37.y; + float2 _fma_pair_38 = fma_f32x2(make_float2(_tmem_load_2[18], _tmem_load_2[19]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[18] = _fma_pair_38.x; + _tmem_load_2[19] = _fma_pair_38.y; + float2 _fma_pair_39 = fma_f32x2(make_float2(_tmem_load_2[20], _tmem_load_2[21]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[20] = _fma_pair_39.x; + _tmem_load_2[21] = _fma_pair_39.y; + float2 _fma_pair_40 = fma_f32x2(make_float2(_tmem_load_2[22], _tmem_load_2[23]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[22] = _fma_pair_40.x; + _tmem_load_2[23] = _fma_pair_40.y; + float2 _fma_pair_41 = fma_f32x2(make_float2(_tmem_load_2[24], _tmem_load_2[25]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[24] = _fma_pair_41.x; + _tmem_load_2[25] = _fma_pair_41.y; + float2 _fma_pair_42 = fma_f32x2(make_float2(_tmem_load_2[26], _tmem_load_2[27]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[26] = _fma_pair_42.x; + _tmem_load_2[27] = _fma_pair_42.y; + float2 _fma_pair_43 = fma_f32x2(make_float2(_tmem_load_2[28], _tmem_load_2[29]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[28] = _fma_pair_43.x; + _tmem_load_2[29] = _fma_pair_43.y; + float2 _fma_pair_44 = fma_f32x2(make_float2(_tmem_load_2[30], _tmem_load_2[31]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[30] = _fma_pair_44.x; + _tmem_load_2[31] = _fma_pair_44.y; + #pragma unroll + for (int _le = 0; _le < 32; _le++) { + _tmem_load_2[_le] = approx_exp2(_tmem_load_2[_le]); + } + float2 _reg_reduce_sum2_45 = make_float2(0.0f, 0.0f); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[0], _tmem_load_2[1])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[2], _tmem_load_2[3])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[4], _tmem_load_2[5])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[6], _tmem_load_2[7])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[8], _tmem_load_2[9])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[10], _tmem_load_2[11])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[12], _tmem_load_2[13])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[14], _tmem_load_2[15])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[16], _tmem_load_2[17])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[18], _tmem_load_2[19])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[20], _tmem_load_2[21])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[22], _tmem_load_2[23])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[24], _tmem_load_2[25])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[26], _tmem_load_2[27])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[28], _tmem_load_2[29])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[30], _tmem_load_2[31])); + float _tmem_load_2_sum = _reg_reduce_sum2_45.x + _reg_reduce_sum2_45.y; + block_sum = block_sum + _tmem_load_2_sum; + { + uint32_t _pv_packed[8]; + #pragma unroll + for (int _j = 0; _j < 8; _j++) { + uint32_t _pk; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_pk) : "f"(_tmem_load_2[0 + _j * 4]), "f"(_tmem_load_2[0 + _j * 4 + 1]), + "f"(_tmem_load_2[0 + _j * 4 + 2]), "f"(_tmem_load_2[0 + _j * 4 + 3])); + _pv_packed[_j] = _pk; } - const float2 _fma_b2_3 = {softmax_scale_log2, softmax_scale_log2}; - const float2 _fma_c2_4 = {-max_scaled, -max_scaled}; - float2 _fma_pair_5 = fma_f32x2(make_float2(_tmem_load_3[0], _tmem_load_3[1]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[0] = _fma_pair_5.x; - _tmem_load_3[1] = _fma_pair_5.y; - float2 _fma_pair_6 = fma_f32x2(make_float2(_tmem_load_3[2], _tmem_load_3[3]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[2] = _fma_pair_6.x; - _tmem_load_3[3] = _fma_pair_6.y; - float2 _fma_pair_7 = fma_f32x2(make_float2(_tmem_load_3[4], _tmem_load_3[5]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[4] = _fma_pair_7.x; - _tmem_load_3[5] = _fma_pair_7.y; - float2 _fma_pair_8 = fma_f32x2(make_float2(_tmem_load_3[6], _tmem_load_3[7]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[6] = _fma_pair_8.x; - _tmem_load_3[7] = _fma_pair_8.y; - float2 _fma_pair_9 = fma_f32x2(make_float2(_tmem_load_3[8], _tmem_load_3[9]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[8] = _fma_pair_9.x; - _tmem_load_3[9] = _fma_pair_9.y; - float2 _fma_pair_10 = fma_f32x2(make_float2(_tmem_load_3[10], _tmem_load_3[11]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[10] = _fma_pair_10.x; - _tmem_load_3[11] = _fma_pair_10.y; - float2 _fma_pair_11 = fma_f32x2(make_float2(_tmem_load_3[12], _tmem_load_3[13]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[12] = _fma_pair_11.x; - _tmem_load_3[13] = _fma_pair_11.y; - float2 _fma_pair_12 = fma_f32x2(make_float2(_tmem_load_3[14], _tmem_load_3[15]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[14] = _fma_pair_12.x; - _tmem_load_3[15] = _fma_pair_12.y; - float2 _fma_pair_13 = fma_f32x2(make_float2(_tmem_load_3[16], _tmem_load_3[17]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[16] = _fma_pair_13.x; - _tmem_load_3[17] = _fma_pair_13.y; - float2 _fma_pair_14 = fma_f32x2(make_float2(_tmem_load_3[18], _tmem_load_3[19]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[18] = _fma_pair_14.x; - _tmem_load_3[19] = _fma_pair_14.y; - float2 _fma_pair_15 = fma_f32x2(make_float2(_tmem_load_3[20], _tmem_load_3[21]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[20] = _fma_pair_15.x; - _tmem_load_3[21] = _fma_pair_15.y; - float2 _fma_pair_16 = fma_f32x2(make_float2(_tmem_load_3[22], _tmem_load_3[23]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[22] = _fma_pair_16.x; - _tmem_load_3[23] = _fma_pair_16.y; - float2 _fma_pair_17 = fma_f32x2(make_float2(_tmem_load_3[24], _tmem_load_3[25]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[24] = _fma_pair_17.x; - _tmem_load_3[25] = _fma_pair_17.y; - float2 _fma_pair_18 = fma_f32x2(make_float2(_tmem_load_3[26], _tmem_load_3[27]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[26] = _fma_pair_18.x; - _tmem_load_3[27] = _fma_pair_18.y; - float2 _fma_pair_19 = fma_f32x2(make_float2(_tmem_load_3[28], _tmem_load_3[29]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[28] = _fma_pair_19.x; - _tmem_load_3[29] = _fma_pair_19.y; - float2 _fma_pair_20 = fma_f32x2(make_float2(_tmem_load_3[30], _tmem_load_3[31]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[30] = _fma_pair_20.x; - _tmem_load_3[31] = _fma_pair_20.y; + tmem_st_x8_u32(p_base_all + 8, _pv_packed); + } + const float2 _fma_b2_46 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_47 = {-max_scaled, -max_scaled}; + float2 _fma_pair_48 = fma_f32x2(make_float2(_tmem_load_3[0], _tmem_load_3[1]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[0] = _fma_pair_48.x; + _tmem_load_3[1] = _fma_pair_48.y; + float2 _fma_pair_49 = fma_f32x2(make_float2(_tmem_load_3[2], _tmem_load_3[3]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[2] = _fma_pair_49.x; + _tmem_load_3[3] = _fma_pair_49.y; + float2 _fma_pair_50 = fma_f32x2(make_float2(_tmem_load_3[4], _tmem_load_3[5]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[4] = _fma_pair_50.x; + _tmem_load_3[5] = _fma_pair_50.y; + float2 _fma_pair_51 = fma_f32x2(make_float2(_tmem_load_3[6], _tmem_load_3[7]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[6] = _fma_pair_51.x; + _tmem_load_3[7] = _fma_pair_51.y; + float2 _fma_pair_52 = fma_f32x2(make_float2(_tmem_load_3[8], _tmem_load_3[9]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[8] = _fma_pair_52.x; + _tmem_load_3[9] = _fma_pair_52.y; + float2 _fma_pair_53 = fma_f32x2(make_float2(_tmem_load_3[10], _tmem_load_3[11]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[10] = _fma_pair_53.x; + _tmem_load_3[11] = _fma_pair_53.y; + float2 _fma_pair_54 = fma_f32x2(make_float2(_tmem_load_3[12], _tmem_load_3[13]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[12] = _fma_pair_54.x; + _tmem_load_3[13] = _fma_pair_54.y; + float2 _fma_pair_55 = fma_f32x2(make_float2(_tmem_load_3[14], _tmem_load_3[15]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[14] = _fma_pair_55.x; + _tmem_load_3[15] = _fma_pair_55.y; + float2 _fma_pair_56 = fma_f32x2(make_float2(_tmem_load_3[16], _tmem_load_3[17]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[16] = _fma_pair_56.x; + _tmem_load_3[17] = _fma_pair_56.y; + float2 _fma_pair_57 = fma_f32x2(make_float2(_tmem_load_3[18], _tmem_load_3[19]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[18] = _fma_pair_57.x; + _tmem_load_3[19] = _fma_pair_57.y; + float2 _fma_pair_58 = fma_f32x2(make_float2(_tmem_load_3[20], _tmem_load_3[21]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[20] = _fma_pair_58.x; + _tmem_load_3[21] = _fma_pair_58.y; + float2 _fma_pair_59 = fma_f32x2(make_float2(_tmem_load_3[22], _tmem_load_3[23]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[22] = _fma_pair_59.x; + _tmem_load_3[23] = _fma_pair_59.y; + float2 _fma_pair_60 = fma_f32x2(make_float2(_tmem_load_3[24], _tmem_load_3[25]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[24] = _fma_pair_60.x; + _tmem_load_3[25] = _fma_pair_60.y; + float2 _fma_pair_61 = fma_f32x2(make_float2(_tmem_load_3[26], _tmem_load_3[27]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[26] = _fma_pair_61.x; + _tmem_load_3[27] = _fma_pair_61.y; + float2 _fma_pair_62 = fma_f32x2(make_float2(_tmem_load_3[28], _tmem_load_3[29]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[28] = _fma_pair_62.x; + _tmem_load_3[29] = _fma_pair_62.y; + float2 _fma_pair_63 = fma_f32x2(make_float2(_tmem_load_3[30], _tmem_load_3[31]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[30] = _fma_pair_63.x; + _tmem_load_3[31] = _fma_pair_63.y; + #pragma unroll + for (int _le = 0; _le < 32; _le++) { + _tmem_load_3[_le] = approx_exp2(_tmem_load_3[_le]); + } + float2 _reg_reduce_sum2_64 = make_float2(0.0f, 0.0f); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[0], _tmem_load_3[1])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[2], _tmem_load_3[3])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[4], _tmem_load_3[5])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[6], _tmem_load_3[7])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[8], _tmem_load_3[9])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[10], _tmem_load_3[11])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[12], _tmem_load_3[13])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[14], _tmem_load_3[15])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[16], _tmem_load_3[17])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[18], _tmem_load_3[19])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[20], _tmem_load_3[21])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[22], _tmem_load_3[23])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[24], _tmem_load_3[25])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[26], _tmem_load_3[27])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[28], _tmem_load_3[29])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[30], _tmem_load_3[31])); + float _tmem_load_3_sum = _reg_reduce_sum2_64.x + _reg_reduce_sum2_64.y; + block_sum = block_sum + _tmem_load_3_sum; + { + uint32_t _pv_packed[8]; #pragma unroll - for (int _le = 0; _le < 32; _le++) { - _tmem_load_3[_le] = approx_exp2(_tmem_load_3[_le]); + for (int _j = 0; _j < 8; _j++) { + uint32_t _pk; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_pk) : "f"(_tmem_load_3[0 + _j * 4]), "f"(_tmem_load_3[0 + _j * 4 + 1]), + "f"(_tmem_load_3[0 + _j * 4 + 2]), "f"(_tmem_load_3[0 + _j * 4 + 3])); + _pv_packed[_j] = _pk; } - float2 _reg_reduce_sum2_21 = make_float2(0.0f, 0.0f); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[0], _tmem_load_3[1])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[2], _tmem_load_3[3])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[4], _tmem_load_3[5])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[6], _tmem_load_3[7])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[8], _tmem_load_3[9])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[10], _tmem_load_3[11])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[12], _tmem_load_3[13])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[14], _tmem_load_3[15])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[16], _tmem_load_3[17])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[18], _tmem_load_3[19])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[20], _tmem_load_3[21])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[22], _tmem_load_3[23])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[24], _tmem_load_3[25])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[26], _tmem_load_3[27])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[28], _tmem_load_3[29])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[30], _tmem_load_3[31])); - float _tmem_load_3_sum = _reg_reduce_sum2_21.x + _reg_reduce_sum2_21.y; - block_sum = block_sum + _tmem_load_3_sum; - { - uint32_t _pv_packed[8]; - #pragma unroll - for (int _j = 0; _j < 8; _j++) { - uint32_t _pk; - asm("{\n\t" - ".reg .b16 _lo, _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}\n" - : "=r"(_pk) : "f"(_tmem_load_3[0 + _j * 4]), "f"(_tmem_load_3[0 + _j * 4 + 1]), - "f"(_tmem_load_3[0 + _j * 4 + 2]), "f"(_tmem_load_3[0 + _j * 4 + 3])); - _pv_packed[_j] = _pk; - } - tmem_st_x8_u32(p_base + frag_1 * 8, _pv_packed); + tmem_st_x8_u32(p_base_all + 16, _pv_packed); + } + const float2 _fma_b2_65 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_66 = {-max_scaled, -max_scaled}; + float2 _fma_pair_67 = fma_f32x2(make_float2(_tmem_load_4[0], _tmem_load_4[1]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[0] = _fma_pair_67.x; + _tmem_load_4[1] = _fma_pair_67.y; + float2 _fma_pair_68 = fma_f32x2(make_float2(_tmem_load_4[2], _tmem_load_4[3]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[2] = _fma_pair_68.x; + _tmem_load_4[3] = _fma_pair_68.y; + float2 _fma_pair_69 = fma_f32x2(make_float2(_tmem_load_4[4], _tmem_load_4[5]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[4] = _fma_pair_69.x; + _tmem_load_4[5] = _fma_pair_69.y; + float2 _fma_pair_70 = fma_f32x2(make_float2(_tmem_load_4[6], _tmem_load_4[7]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[6] = _fma_pair_70.x; + _tmem_load_4[7] = _fma_pair_70.y; + float2 _fma_pair_71 = fma_f32x2(make_float2(_tmem_load_4[8], _tmem_load_4[9]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[8] = _fma_pair_71.x; + _tmem_load_4[9] = _fma_pair_71.y; + float2 _fma_pair_72 = fma_f32x2(make_float2(_tmem_load_4[10], _tmem_load_4[11]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[10] = _fma_pair_72.x; + _tmem_load_4[11] = _fma_pair_72.y; + float2 _fma_pair_73 = fma_f32x2(make_float2(_tmem_load_4[12], _tmem_load_4[13]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[12] = _fma_pair_73.x; + _tmem_load_4[13] = _fma_pair_73.y; + float2 _fma_pair_74 = fma_f32x2(make_float2(_tmem_load_4[14], _tmem_load_4[15]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[14] = _fma_pair_74.x; + _tmem_load_4[15] = _fma_pair_74.y; + float2 _fma_pair_75 = fma_f32x2(make_float2(_tmem_load_4[16], _tmem_load_4[17]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[16] = _fma_pair_75.x; + _tmem_load_4[17] = _fma_pair_75.y; + float2 _fma_pair_76 = fma_f32x2(make_float2(_tmem_load_4[18], _tmem_load_4[19]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[18] = _fma_pair_76.x; + _tmem_load_4[19] = _fma_pair_76.y; + float2 _fma_pair_77 = fma_f32x2(make_float2(_tmem_load_4[20], _tmem_load_4[21]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[20] = _fma_pair_77.x; + _tmem_load_4[21] = _fma_pair_77.y; + float2 _fma_pair_78 = fma_f32x2(make_float2(_tmem_load_4[22], _tmem_load_4[23]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[22] = _fma_pair_78.x; + _tmem_load_4[23] = _fma_pair_78.y; + float2 _fma_pair_79 = fma_f32x2(make_float2(_tmem_load_4[24], _tmem_load_4[25]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[24] = _fma_pair_79.x; + _tmem_load_4[25] = _fma_pair_79.y; + float2 _fma_pair_80 = fma_f32x2(make_float2(_tmem_load_4[26], _tmem_load_4[27]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[26] = _fma_pair_80.x; + _tmem_load_4[27] = _fma_pair_80.y; + float2 _fma_pair_81 = fma_f32x2(make_float2(_tmem_load_4[28], _tmem_load_4[29]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[28] = _fma_pair_81.x; + _tmem_load_4[29] = _fma_pair_81.y; + float2 _fma_pair_82 = fma_f32x2(make_float2(_tmem_load_4[30], _tmem_load_4[31]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[30] = _fma_pair_82.x; + _tmem_load_4[31] = _fma_pair_82.y; + #pragma unroll + for (int _le = 0; _le < 32; _le++) { + _tmem_load_4[_le] = approx_exp2(_tmem_load_4[_le]); + } + float2 _reg_reduce_sum2_83 = make_float2(0.0f, 0.0f); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[0], _tmem_load_4[1])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[2], _tmem_load_4[3])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[4], _tmem_load_4[5])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[6], _tmem_load_4[7])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[8], _tmem_load_4[9])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[10], _tmem_load_4[11])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[12], _tmem_load_4[13])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[14], _tmem_load_4[15])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[16], _tmem_load_4[17])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[18], _tmem_load_4[19])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[20], _tmem_load_4[21])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[22], _tmem_load_4[23])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[24], _tmem_load_4[25])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[26], _tmem_load_4[27])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[28], _tmem_load_4[29])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[30], _tmem_load_4[31])); + float _tmem_load_4_sum = _reg_reduce_sum2_83.x + _reg_reduce_sum2_83.y; + block_sum = block_sum + _tmem_load_4_sum; + { + uint32_t _pv_packed[8]; + #pragma unroll + for (int _j = 0; _j < 8; _j++) { + uint32_t _pk; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_pk) : "f"(_tmem_load_4[0 + _j * 4]), "f"(_tmem_load_4[0 + _j * 4 + 1]), + "f"(_tmem_load_4[0 + _j * 4 + 2]), "f"(_tmem_load_4[0 + _j * 4 + 3])); + _pv_packed[_j] = _pk; } + tmem_st_x8_u32(p_base_all + 24, _pv_packed); } } { @@ -2087,13 +2607,13 @@ kernel_cake_dsv4_a2be51a1fd34ec0d9a3b(CakeTensorMap const* tmap_q, CakeTensorMap int o_base = taddr + 256 + (unsigned int)(vs * 128) + (unsigned int)corr_row; #pragma unroll for (int c = 0; c < 128; c += 16) { - float _tmem_load_4[16]; - tmem_ld_x16(&_tmem_load_4[0], o_base + c); + float _tmem_load_5[16]; + tmem_ld_x16(&_tmem_load_5[0], o_base + c); const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; #pragma unroll for (int _ls = 0; _ls < 8; _ls++) - mul_f32x2_inplace(&reinterpret_cast(_tmem_load_4)[_ls], _scale2_0); - tmem_st_x16_f32(o_base + c, _tmem_load_4); + mul_f32x2_inplace(&reinterpret_cast(_tmem_load_5)[_ls], _scale2_0); + tmem_st_x16_f32(o_base + c, _tmem_load_5); asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); } } @@ -2126,11 +2646,11 @@ kernel_cake_dsv4_a2be51a1fd34ec0d9a3b(CakeTensorMap const* tmap_q, CakeTensorMap int o_base_epi = taddr + 256 + (unsigned int)(vs_1 * 128) + (unsigned int)corr_row; #pragma unroll for (int c_1 = 0; c_1 < 128; c_1 += 32) { - float _tmem_load_5[32]; + float _tmem_load_6[32]; asm volatile( "tcgen05.ld.sync.aligned.32x32b.x32.b32" " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(_tmem_load_5[0]), "=f"(_tmem_load_5[1]), "=f"(_tmem_load_5[2]), "=f"(_tmem_load_5[3]), "=f"(_tmem_load_5[4]), "=f"(_tmem_load_5[5]), "=f"(_tmem_load_5[6]), "=f"(_tmem_load_5[7]), "=f"(_tmem_load_5[8]), "=f"(_tmem_load_5[9]), "=f"(_tmem_load_5[10]), "=f"(_tmem_load_5[11]), "=f"(_tmem_load_5[12]), "=f"(_tmem_load_5[13]), "=f"(_tmem_load_5[14]), "=f"(_tmem_load_5[15]), "=f"(_tmem_load_5[16]), "=f"(_tmem_load_5[17]), "=f"(_tmem_load_5[18]), "=f"(_tmem_load_5[19]), "=f"(_tmem_load_5[20]), "=f"(_tmem_load_5[21]), "=f"(_tmem_load_5[22]), "=f"(_tmem_load_5[23]), "=f"(_tmem_load_5[24]), "=f"(_tmem_load_5[25]), "=f"(_tmem_load_5[26]), "=f"(_tmem_load_5[27]), "=f"(_tmem_load_5[28]), "=f"(_tmem_load_5[29]), "=f"(_tmem_load_5[30]), "=f"(_tmem_load_5[31]) + : "=f"(_tmem_load_6[0]), "=f"(_tmem_load_6[1]), "=f"(_tmem_load_6[2]), "=f"(_tmem_load_6[3]), "=f"(_tmem_load_6[4]), "=f"(_tmem_load_6[5]), "=f"(_tmem_load_6[6]), "=f"(_tmem_load_6[7]), "=f"(_tmem_load_6[8]), "=f"(_tmem_load_6[9]), "=f"(_tmem_load_6[10]), "=f"(_tmem_load_6[11]), "=f"(_tmem_load_6[12]), "=f"(_tmem_load_6[13]), "=f"(_tmem_load_6[14]), "=f"(_tmem_load_6[15]), "=f"(_tmem_load_6[16]), "=f"(_tmem_load_6[17]), "=f"(_tmem_load_6[18]), "=f"(_tmem_load_6[19]), "=f"(_tmem_load_6[20]), "=f"(_tmem_load_6[21]), "=f"(_tmem_load_6[22]), "=f"(_tmem_load_6[23]), "=f"(_tmem_load_6[24]), "=f"(_tmem_load_6[25]), "=f"(_tmem_load_6[26]), "=f"(_tmem_load_6[27]), "=f"(_tmem_load_6[28]), "=f"(_tmem_load_6[29]), "=f"(_tmem_load_6[30]), "=f"(_tmem_load_6[31]) : "r"(o_base_epi + c_1)); int gmem_base = o_offset + vs_1 * 128 + c_1; if (head_idx < num_heads) { @@ -2141,28 +2661,28 @@ kernel_cake_dsv4_a2be51a1fd34ec0d9a3b(CakeTensorMap const* tmap_q, CakeTensorMap #if __CUDA_ARCH__ >= 1000 #pragma unroll for (int _ps = 0; _ps < 8; _ps++) - mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_5[j])[_ps], _prescale2_1); + mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_6[j])[_ps], _prescale2_1); #else #pragma unroll for (int _ps = 0; _ps < 16; _ps++) - _tmem_load_5[j + _ps] *= inv_sum * output_scale; + _tmem_load_6[j + _ps] *= inv_sum * output_scale; #endif { - __nv_bfloat162 _pk0 = __floats2bfloat162_rn(_tmem_load_5[j + 0], _tmem_load_5[j + 1]); + __nv_bfloat162 _pk0 = __floats2bfloat162_rn(_tmem_load_6[j + 0], _tmem_load_6[j + 1]); unsigned _pk_u0 = *reinterpret_cast(&_pk0); - __nv_bfloat162 _pk1 = __floats2bfloat162_rn(_tmem_load_5[j + 2], _tmem_load_5[j + 3]); + __nv_bfloat162 _pk1 = __floats2bfloat162_rn(_tmem_load_6[j + 2], _tmem_load_6[j + 3]); unsigned _pk_u1 = *reinterpret_cast(&_pk1); - __nv_bfloat162 _pk2 = __floats2bfloat162_rn(_tmem_load_5[j + 4], _tmem_load_5[j + 5]); + __nv_bfloat162 _pk2 = __floats2bfloat162_rn(_tmem_load_6[j + 4], _tmem_load_6[j + 5]); unsigned _pk_u2 = *reinterpret_cast(&_pk2); - __nv_bfloat162 _pk3 = __floats2bfloat162_rn(_tmem_load_5[j + 6], _tmem_load_5[j + 7]); + __nv_bfloat162 _pk3 = __floats2bfloat162_rn(_tmem_load_6[j + 6], _tmem_load_6[j + 7]); unsigned _pk_u3 = *reinterpret_cast(&_pk3); - __nv_bfloat162 _pk4 = __floats2bfloat162_rn(_tmem_load_5[j + 8], _tmem_load_5[j + 9]); + __nv_bfloat162 _pk4 = __floats2bfloat162_rn(_tmem_load_6[j + 8], _tmem_load_6[j + 9]); unsigned _pk_u4 = *reinterpret_cast(&_pk4); - __nv_bfloat162 _pk5 = __floats2bfloat162_rn(_tmem_load_5[j + 10], _tmem_load_5[j + 11]); + __nv_bfloat162 _pk5 = __floats2bfloat162_rn(_tmem_load_6[j + 10], _tmem_load_6[j + 11]); unsigned _pk_u5 = *reinterpret_cast(&_pk5); - __nv_bfloat162 _pk6 = __floats2bfloat162_rn(_tmem_load_5[j + 12], _tmem_load_5[j + 13]); + __nv_bfloat162 _pk6 = __floats2bfloat162_rn(_tmem_load_6[j + 12], _tmem_load_6[j + 13]); unsigned _pk_u6 = *reinterpret_cast(&_pk6); - __nv_bfloat162 _pk7 = __floats2bfloat162_rn(_tmem_load_5[j + 14], _tmem_load_5[j + 15]); + __nv_bfloat162 _pk7 = __floats2bfloat162_rn(_tmem_load_6[j + 14], _tmem_load_6[j + 15]); unsigned _pk_u7 = *reinterpret_cast(&_pk7); asm volatile( "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_9e1188b192c9c806d7c6_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_983db4c490ceed3c455a_binding.cu similarity index 97% rename from csrc/cake_dsv4/sm_103a/cake_dsv4_9e1188b192c9c806d7c6_binding.cu rename to csrc/cake_dsv4/sm_103a/cake_dsv4_983db4c490ceed3c455a_binding.cu index b3d17f4a22c..b5c6878ec55 100644 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_9e1188b192c9c806d7c6_binding.cu +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_983db4c490ceed3c455a_binding.cu @@ -15,7 +15,7 @@ */ // Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_9e1188b192c9c806d7c6'. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_983db4c490ceed3c455a'. #include #include #include @@ -31,10 +31,10 @@ #include #include -extern "C" __global__ void kernel_cake_dsv4_9e1188b192c9c806d7c6(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items); +extern "C" __global__ void kernel_cake_dsv4_983db4c490ceed3c455a(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items); -namespace cake_host_shim_4aa7ddc0eb16149e { +namespace cake_host_shim_8e7f30cc5fb7222a { using tvm::ffi::Optional; using tvm::ffi::TensorView; @@ -416,7 +416,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ } if (smem_status == cudaSuccess) { smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_9e1188b192c9c806d7c6), + reinterpret_cast(kernel_cake_dsv4_983db4c490ceed3c455a), cudaFuncAttributeMaxDynamicSharedMemorySize, 223744); } @@ -426,7 +426,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ } } TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_9e1188b192c9c806d7c6 failed: " + << "cudaFuncSetAttribute for kernel_cake_dsv4_983db4c490ceed3c455a failed: " << cudaGetErrorString(smem_status); dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); dim3 block(512u, 1u, 1u); @@ -450,12 +450,12 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ config.attrs = attrs; config.numAttrs = n; cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_9e1188b192c9c806d7c6), kargs); + &config, reinterpret_cast(kernel_cake_dsv4_983db4c490ceed3c455a), kargs); TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_9e1188b192c9c806d7c6 failed: " + << "cudaLaunchKernelExC for kernel_cake_dsv4_983db4c490ceed3c455a failed: " << cudaGetErrorString(launch_status); } -} // namespace cake_host_shim_4aa7ddc0eb16149e +} // namespace cake_host_shim_8e7f30cc5fb7222a -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_4aa7ddc0eb16149e::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_8e7f30cc5fb7222a::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_9e1188b192c9c806d7c6_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_983db4c490ceed3c455a_kernel.cu similarity index 99% rename from csrc/cake_dsv4/sm_103a/cake_dsv4_9e1188b192c9c806d7c6_kernel.cu rename to csrc/cake_dsv4/sm_103a/cake_dsv4_983db4c490ceed3c455a_kernel.cu index b26d3631ccb..a01f2cc29ce 100644 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_9e1188b192c9c806d7c6_kernel.cu +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_983db4c490ceed3c455a_kernel.cu @@ -1220,7 +1220,7 @@ __device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { extern "C" { __global__ __launch_bounds__(512, 1) __cluster_dims__(2,1,1) void -kernel_cake_dsv4_9e1188b192c9c806d7c6(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items) +kernel_cake_dsv4_983db4c490ceed3c455a(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); @@ -1625,10 +1625,10 @@ kernel_cake_dsv4_9e1188b192c9c806d7c6(const __grid_constant__ CUtensorMap tmap_q :: "r"((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); } } - if (load_warp_rank < 2) { - int v_dst = smem_v_full_addr + load_v_stage * 32768 + (unsigned int)(load_warp_rank * 16384); - int v_col = v_chunk * 256 + load_warp_rank * 128; - int group_2 = lane; + if (lane < 16) { + int v_dst = smem_v_full_addr + load_v_stage * 32768 + (unsigned int)(load_warp_rank / 2 * 16384); + int v_col = v_chunk * 256 + load_warp_rank / 2 * 128; + int group_2 = (load_warp_rank & 1) * 16 + lane; int group_offset_2 = (prev_tile & 1) * 128 + group_2 * 4; int raw_rows_2[4]; asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" @@ -1667,10 +1667,10 @@ kernel_cake_dsv4_9e1188b192c9c806d7c6(const __grid_constant__ CUtensorMap tmap_q :: "r"((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); } } - if (load_warp_rank < 2) { - int v_dst_1 = smem_v_full_addr + load_v_stage * 32768 + (unsigned int)(load_warp_rank * 16384); - int v_col_1 = v_chunk * 256 + load_warp_rank * 128; - int group_3 = lane; + if (lane < 16) { + int v_dst_1 = smem_v_full_addr + load_v_stage * 32768 + (unsigned int)(load_warp_rank / 2 * 16384); + int v_col_1 = v_chunk * 256 + load_warp_rank / 2 * 128; + int group_3 = (load_warp_rank & 1) * 16 + lane; int group_offset_3 = (last_tile & 1) * 128 + group_3 * 4; int raw_rows_3[4]; asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_c20607c57cc11784015e_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_98db1f55c8c129e16005_binding.cu similarity index 97% rename from csrc/cake_dsv4/sm_103a/cake_dsv4_c20607c57cc11784015e_binding.cu rename to csrc/cake_dsv4/sm_103a/cake_dsv4_98db1f55c8c129e16005_binding.cu index a287b064d89..0f76ba73e68 100644 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_c20607c57cc11784015e_binding.cu +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_98db1f55c8c129e16005_binding.cu @@ -15,7 +15,7 @@ */ // Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_c20607c57cc11784015e'. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_98db1f55c8c129e16005'. #include #include #include @@ -31,10 +31,10 @@ #include #include -extern "C" __global__ void kernel_cake_dsv4_c20607c57cc11784015e(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items); +extern "C" __global__ void kernel_cake_dsv4_98db1f55c8c129e16005(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items); -namespace cake_host_shim_9514fa50724352e6 { +namespace cake_host_shim_9d011e971631673b { using tvm::ffi::Optional; using tvm::ffi::TensorView; @@ -416,7 +416,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ } if (smem_status == cudaSuccess) { smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_c20607c57cc11784015e), + reinterpret_cast(kernel_cake_dsv4_98db1f55c8c129e16005), cudaFuncAttributeMaxDynamicSharedMemorySize, 190976); } @@ -426,7 +426,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ } } TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_c20607c57cc11784015e failed: " + << "cudaFuncSetAttribute for kernel_cake_dsv4_98db1f55c8c129e16005 failed: " << cudaGetErrorString(smem_status); dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); dim3 block(512u, 1u, 1u); @@ -450,12 +450,12 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ config.attrs = attrs; config.numAttrs = n; cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_c20607c57cc11784015e), kargs); + &config, reinterpret_cast(kernel_cake_dsv4_98db1f55c8c129e16005), kargs); TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_c20607c57cc11784015e failed: " + << "cudaLaunchKernelExC for kernel_cake_dsv4_98db1f55c8c129e16005 failed: " << cudaGetErrorString(launch_status); } -} // namespace cake_host_shim_9514fa50724352e6 +} // namespace cake_host_shim_9d011e971631673b -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_9514fa50724352e6::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_9d011e971631673b::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_c20607c57cc11784015e_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_98db1f55c8c129e16005_kernel.cu similarity index 77% rename from csrc/cake_dsv4/sm_103a/cake_dsv4_c20607c57cc11784015e_kernel.cu rename to csrc/cake_dsv4/sm_103a/cake_dsv4_98db1f55c8c129e16005_kernel.cu index e88df9fd3d5..6364f2a635d 100644 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_c20607c57cc11784015e_kernel.cu +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_98db1f55c8c129e16005_kernel.cu @@ -62,11 +62,11 @@ static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap ali #define SMEM_SMEM_Q_FULL_STAGE_BYTES 32768 #define SMEM_SMEM_Q_FULL_STRIDE 32768 #define SMEM_SMEM_K_FULL_OFF 33792 -#define SMEM_SMEM_K_FULL_STAGE_BYTES 32768 -#define SMEM_SMEM_K_FULL_STRIDE 32768 -#define SMEM_SMEM_V_FULL_OFF 99328 -#define SMEM_SMEM_V_FULL_STAGE_BYTES 32768 -#define SMEM_SMEM_V_FULL_STRIDE 32768 +#define SMEM_SMEM_K_FULL_STAGE_BYTES 65536 +#define SMEM_SMEM_K_FULL_STRIDE 65536 +#define SMEM_SMEM_V_FULL_OFF 33792 +#define SMEM_SMEM_V_FULL_STAGE_BYTES 65536 +#define SMEM_SMEM_V_FULL_STRIDE 65536 #define SMEM_SMEM_STATS_MAX_OFF 164864 #define SMEM_SMEM_STATS_MAX_STAGE_BYTES 1024 #define SMEM_SMEM_STATS_MAX_STRIDE 1024 @@ -88,6 +88,18 @@ static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap ali #define SMEM_SMEM_P_FP8_OFF 174592 #define SMEM_SMEM_P_FP8_STAGE_BYTES 8192 #define SMEM_SMEM_P_FP8_STRIDE 8192 +#define SMEM_SMEM_MERGE_RECV_OFF 33792 +#define SMEM_SMEM_MERGE_RECV_STAGE_BYTES 65536 +#define SMEM_SMEM_MERGE_RECV_STRIDE 65536 +#define SMEM_SMEM_MERGE_SEND_OFF 99328 +#define SMEM_SMEM_MERGE_SEND_STAGE_BYTES 65536 +#define SMEM_SMEM_MERGE_SEND_STRIDE 65536 +#define SMEM_SMEM_MERGE_STATS_RECV_OFF 174592 +#define SMEM_SMEM_MERGE_STATS_RECV_STAGE_BYTES 512 +#define SMEM_SMEM_MERGE_STATS_RECV_STRIDE 512 +#define SMEM_SMEM_MERGE_STATS_SEND_OFF 182784 +#define SMEM_SMEM_MERGE_STATS_SEND_STAGE_BYTES 512 +#define SMEM_SMEM_MERGE_STATS_SEND_STRIDE 512 #define SMEM_TOTAL 190976 #define THREADS 512 @@ -308,17 +320,14 @@ __device__ __forceinline__ void mbarrier_wait_token_cluster_hint( } -__device__ __forceinline__ void tcgen05_mma_f8f6f4_cta2( +__device__ __forceinline__ void tcgen05_mma_f8f6f4( int taddr, uint64_t a_desc, uint64_t b_desc, uint32_t i_desc, int enable_input_d) { asm volatile( "{\n\t" ".reg .pred p;\n\t" - ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" "setp.ne.b32 p, %4, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" - "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "tcgen05.mma.cta_group::2.kind::f8f6f4 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" + "tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], %1, %2, %3, p;\n\t" "}\n" :: "r"(taddr), "l"(a_desc), "l"(b_desc), "r"(i_desc), "r"(enable_input_d) @@ -344,15 +353,15 @@ __device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t } -__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { +__device__ __forceinline__ void elect_commit(int mbar_addr) { asm volatile( "{\n\t" ".reg .pred leader;\n\t" "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" - ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" "}\n" - :: "r"(mbar_addr), "h"(cta_mask) : "memory"); + :: "r"(mbar_addr)); } @@ -1112,11 +1121,11 @@ __device__ __forceinline__ void tma_2d_gmem2smem( } -__device__ __forceinline__ void tma_4d_gmem2smem_cta2( +__device__ __forceinline__ void tma_4d_gmem2smem( int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { asm volatile( - "cp.async.bulk.tensor.4d.shared::cluster.global" - ".mbarrier::complete_tx::bytes.cta_group::2" + "cp.async.bulk.tensor.4d.shared::cta.global" + ".mbarrier::complete_tx::bytes" " [%0], [%1, {%2, %3, %4, %5}], [%6];" :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), "r"(mbar_addr) : "memory"); @@ -1141,62 +1150,11 @@ __device__ __forceinline__ void tma_gather4_gmem2smem( } -__device__ __forceinline__ void tma_gather4_gmem2smem_cta2( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form; see tma_gather4_gmem2smem above. - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem_mc( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr, unsigned short cta_mask) { - // Multicast variant: the PTX grammar ties the .shared::cluster - // destination to .multicast::cluster + ctaMask (cf. cuda_ptx / - // SM100_TMA_LOAD_MULTICAST_2D_GATHER4 in CUTLASS). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" - ".mbarrier::complete_tx::bytes.multicast::cluster" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr), "h"(cta_mask) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem_mc_cta2( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr, unsigned short cta_mask) { - // Multicast + cta_group::2 variant; see tma_gather4_gmem2smem_mc. - asm volatile( - "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" - ".mbarrier::complete_tx::bytes.multicast::cluster.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr), "h"(cta_mask) : "memory"); -} - - -__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { asm volatile( - "{\n\t" - ".reg .b16 lo, hi;\n\t" - "mov.b32 {lo, hi}, %1;\n\t" - "tcgen05.commit.cta_group::2.mbarrier::arrive::one" - ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" - "}\n" - :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); } @@ -1220,7 +1178,7 @@ __device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { extern "C" { __global__ __launch_bounds__(512, 1) __cluster_dims__(2,1,1) void -kernel_cake_dsv4_c20607c57cc11784015e(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items) +kernel_cake_dsv4_98db1f55c8c129e16005(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); @@ -1252,6 +1210,8 @@ kernel_cake_dsv4_c20607c57cc11784015e(const __grid_constant__ CUtensorMap tmap_q #define pv_pair_ready_addr (mbar_base + 288) #define tmem_dealloc_addr (mbar_base + 304) #define tmem_dealloc_peer_addr (mbar_base + 312) + #define merge_ready_addr (mbar_base + 320) + #define merge_full_addr (mbar_base + 328) const int bid = blockIdx.x; const int num_bids = gridDim.x; @@ -1273,8 +1233,8 @@ kernel_cake_dsv4_c20607c57cc11784015e(const __grid_constant__ CUtensorMap tmap_q const int smem_q_full_addr = smem + 1024; uint8_t* smem_k_full = reinterpret_cast(smem_raw + 33792); const int smem_k_full_addr = smem + 33792; - uint8_t* smem_v_full = reinterpret_cast(smem_raw + 99328); - const int smem_v_full_addr = smem + 99328; + uint8_t* smem_v_full = reinterpret_cast(smem_raw + 33792); + const int smem_v_full_addr = smem + 33792; float* smem_stats_max = reinterpret_cast(smem_raw + 164864); const int smem_stats_max_addr = smem + 164864; float* smem_stats_sum = reinterpret_cast(smem_raw + 165888); @@ -1289,25 +1249,33 @@ kernel_cake_dsv4_c20607c57cc11784015e(const __grid_constant__ CUtensorMap tmap_q const int smem_sparse_indices_addr = smem + 168448; uint8_t* smem_p_fp8 = reinterpret_cast(smem_raw + 174592); const int smem_p_fp8_addr = smem + 174592; + float* smem_merge_recv = reinterpret_cast(smem_raw + 33792); + const int smem_merge_recv_addr = smem + 33792; + float* smem_merge_send = reinterpret_cast(smem_raw + 99328); + const int smem_merge_send_addr = smem + 99328; + float* smem_merge_stats_recv = reinterpret_cast(smem_raw + 174592); + const int smem_merge_stats_recv_addr = smem + 174592; + float* smem_merge_stats_send = reinterpret_cast(smem_raw + 182784); + const int smem_merge_stats_send_addr = smem + 182784; { asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); } - // Mbarrier init (21 pipeline groups, 0 ordered-sequence groups, 40 barriers) - // Mbarriers at smem_raw[0..320) + // Mbarrier init (23 pipeline groups, 0 ordered-sequence groups, 42 barriers) + // Mbarriers at smem_raw[0..336) if (warp == 12) { uint32_t leader = elect_sync(); if (leader) { - // q_full: 1 barriers, init_count=2 - mbarrier_init(smem + 0, 2); + // q_full: 1 barriers, init_count=1 + mbarrier_init(smem + 0, 1); // q_empty: 1 barriers, init_count=1 mbarrier_init(smem + 8, 1); // --- pipeline 'k_pipe' --- - // k_full: 2 barriers, init_count=2 - mbarrier_init(smem + 16, 2); - mbarrier_init(smem + 24, 2); + // k_full: 2 barriers, init_count=1 + mbarrier_init(smem + 16, 1); + mbarrier_init(smem + 24, 1); // k_empty: 2 barriers, init_count=1 mbarrier_init(smem + 32, 1); mbarrier_init(smem + 40, 1); @@ -1394,6 +1362,10 @@ kernel_cake_dsv4_c20607c57cc11784015e(const __grid_constant__ CUtensorMap tmap_q mbarrier_init(smem + 296, 64); // tmem_dealloc: 1 barriers, init_count=288 mbarrier_init(smem + 304, 288); + // merge_ready: 1 barriers, init_count=64 + mbarrier_init(smem + 320, 64); + // merge_full: 1 barriers, init_count=1 + mbarrier_init(smem + 328, 1); asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); } } @@ -1418,16 +1390,13 @@ kernel_cake_dsv4_c20607c57cc11784015e(const __grid_constant__ CUtensorMap tmap_q __syncwarp(); - asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); - asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); - // TMEM alloc (512 columns, 512 used) - volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 320); + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 336); if (warp == 0) { - int _tmem_hold = smem + 320; - asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + int _tmem_hold = smem + 336; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); __syncwarp(); - asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); } __syncthreads(); @@ -1517,8 +1486,6 @@ kernel_cake_dsv4_c20607c57cc11784015e(const __grid_constant__ CUtensorMap tmap_q unsigned int _phase_q_empty_0 = 1; unsigned int _phase_index_full = 0; unsigned int _phase_k_empty = 1; - unsigned int _phase_k_empty_1 = 1; - unsigned int _phase_v_empty = 1; #pragma unroll for (unsigned int work_idx = 0; work_idx < 1; work_idx++) { int split_idx_1 = 0; @@ -1533,176 +1500,52 @@ kernel_cake_dsv4_c20607c57cc11784015e(const __grid_constant__ CUtensorMap tmap_q int first_tile_1 = split_idx_1 * tiles_per_split_1; int num_kv_tiles = tiles_per_split_1; int sparse_extent = num_kv_tiles * 128; + int wt_num_local_tiles = (num_kv_tiles - bid % 2 + 1) / 2; mbarrier_wait(q_empty_addr, _phase_q_empty_0); _phase_q_empty_0 ^= 1; if (load_warp_rank == 0) { if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + mbarrier_arrive_expect_tx(q_full_addr, 32768); #pragma unroll for (int q_stage = 0; q_stage < 4; q_stage++) { - asm volatile( - "cp.async.bulk.tensor.4d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5}], [%6], %7;" - :: "r"(smem_q_full_addr + (unsigned int)(q_stage * 8192)), "l"((&tmap_q)), "r"(0), "r"(bid % 2 * 64), "r"(q_stage), "r"(query_idx_1), - "r"(((q_full_addr) & 0xFEFFFFFF)), "h"((uint16_t)(1 << bid % 2)) : "memory"); + tma_4d_gmem2smem(smem_q_full_addr + (unsigned int)(q_stage * 8192), (&tmap_q), 0, 0, q_stage, query_idx_1, q_full_addr); } } } - { + #pragma unroll 1 + for (int wt_tile = 0; wt_tile < wt_num_local_tiles; wt_tile++) { + int wt_global_tile = first_tile_1 + wt_tile * 2 + bid % 2; mbarrier_wait(index_full_addr + (load_k_index_stage) * 8, _phase_index_full); asm volatile("tcgen05.fence::after_thread_sync;"); - } - int k_index_stage_base = smem_sparse_indices_addr + load_k_index_stage * 1024; - { + int wt_index_base = smem_sparse_indices_addr + load_k_index_stage * 1024; mbarrier_wait(k_empty_addr + (load_k_stage) * 8, _phase_k_empty); if (load_warp_rank == 0) { if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + mbarrier_arrive_expect_tx(k_full_addr + (load_k_stage) * 8, 65536); } } - int k_dst = smem_k_full_addr + load_k_stage * 32768 + (unsigned int)(load_warp_rank * 8192); - int group = lane; - if (group < 16) { - int group_offset = k_cta_offset + group * 4; - int raw_rows[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows[0])), "=r"(*reinterpret_cast(&raw_rows[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows[(0) + 3])) - : "r"(k_index_stage_base + group_offset * 4)); - int row0 = ((raw_rows[0] >= 0) ? raw_rows[0] : 0); - int row1 = ((raw_rows[1] >= 0) ? raw_rows[1] : 0); - int row2 = ((raw_rows[2] >= 0) ? raw_rows[2] : 0); - int row3 = ((raw_rows[3] >= 0) ? raw_rows[3] : 0); - if (first_tile_1 == 0) { - tma_gather4_gmem2smem_mc_cta2(k_dst + group * 512, (&tmap_swa_kv), load_warp_rank * 128, row0, row1, row2, row3, ((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } else { - tma_gather4_gmem2smem_mc_cta2(k_dst + group * 512, (&tmap_compressed_kv), load_warp_rank * 128, row0, row1, row2, row3, ((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } + int wt_dst = smem_k_full_addr + load_k_stage * 65536 + (unsigned int)(load_warp_rank * 16384); + int wt_group = lane; + int wt_group_offset = bid % 2 * 128 + wt_group * 4; + int wt_rows[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&wt_rows[0])), "=r"(*reinterpret_cast(&wt_rows[(0) + 1])), "=r"(*reinterpret_cast(&wt_rows[(0) + 2])), "=r"(*reinterpret_cast(&wt_rows[(0) + 3])) + : "r"(wt_index_base + wt_group_offset * 4)); + int wt_row0 = ((wt_rows[0] >= 0) ? wt_rows[0] : 0); + int wt_row1 = ((wt_rows[1] >= 0) ? wt_rows[1] : 0); + int wt_row2 = ((wt_rows[2] >= 0) ? wt_rows[2] : 0); + int wt_row3 = ((wt_rows[3] >= 0) ? wt_rows[3] : 0); + if (wt_global_tile == 0) { + tma_gather4_gmem2smem(wt_dst + wt_group * 512, (&tmap_swa_kv), load_warp_rank * 128, wt_row0, wt_row1, wt_row2, wt_row3, k_full_addr + (load_k_stage) * 8); + } else { + tma_gather4_gmem2smem(wt_dst + wt_group * 512, (&tmap_compressed_kv), load_warp_rank * 128, wt_row0, wt_row1, wt_row2, wt_row3, k_full_addr + (load_k_stage) * 8); } load_k_stage += 1; if (load_k_stage == 2) { load_k_stage = 0; _phase_k_empty ^= 1; } - } - #pragma unroll 1 - for (int tile = 1; tile < num_kv_tiles; tile++) { - if ((1 & (int)((tile & 1) == 0)) != 0) { - mbarrier_wait(index_full_addr + (load_k_index_stage) * 8, _phase_index_full); - asm volatile("tcgen05.fence::after_thread_sync;"); - } - k_index_stage_base = smem_sparse_indices_addr + load_k_index_stage * 1024; - { - mbarrier_wait(k_empty_addr + (load_k_stage) * 8, _phase_k_empty); - if (load_warp_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - int k_dst_1 = smem_k_full_addr + load_k_stage * 32768 + (unsigned int)(load_warp_rank * 8192); - int group_1 = lane; - if (group_1 < 16) { - int group_offset_1 = (tile & 1) * 128 + k_cta_offset + group_1 * 4; - int raw_rows_1[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_1[0])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 3])) - : "r"(k_index_stage_base + group_offset_1 * 4)); - int row0_1 = ((raw_rows_1[0] >= 0) ? raw_rows_1[0] : 0); - int row1_1 = ((raw_rows_1[1] >= 0) ? raw_rows_1[1] : 0); - int row2_1 = ((raw_rows_1[2] >= 0) ? raw_rows_1[2] : 0); - int row3_1 = ((raw_rows_1[3] >= 0) ? raw_rows_1[3] : 0); - tma_gather4_gmem2smem_mc_cta2(k_dst_1 + group_1 * 512, (&tmap_compressed_kv), load_warp_rank * 128, row0_1, row1_1, row2_1, row3_1, ((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } - load_k_stage += 1; - if (load_k_stage == 2) { load_k_stage = 0; _phase_k_empty ^= 1; } - if ((tile & 1) != 0) { - load_k_index_stage += 1; - if (load_k_index_stage == 6) { load_k_index_stage = 0; _phase_index_full ^= 1; } - } - } - int prev_tile = tile - 1; - int global_prev_tile = first_tile_1 + prev_tile; - int v_index_stage_base = smem_sparse_indices_addr + load_v_index_stage * 1024; - { - mbarrier_wait(v_empty_addr + (load_v_stage) * 8, _phase_v_empty); - if (load_warp_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - if (lane < 16) { - int v_dst = smem_v_full_addr + load_v_stage * 32768 + (unsigned int)(load_warp_rank / 2 * 16384); - int v_col = v_chunk * 256 + load_warp_rank / 2 * 128; - int group_2 = (load_warp_rank & 1) * 16 + lane; - int group_offset_2 = (prev_tile & 1) * 128 + group_2 * 4; - int raw_rows_2[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_2[0])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 3])) - : "r"(v_index_stage_base + group_offset_2 * 4)); - int row0_2 = ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : 0); - int row1_2 = ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : 0); - int row2_2 = ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : 0); - int row3_2 = ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : 0); - if (global_prev_tile == 0) { - tma_gather4_gmem2smem_mc_cta2(v_dst + group_2 * 512, (&tmap_swa_kv), v_col, row0_2, row1_2, row2_2, row3_2, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } else { - tma_gather4_gmem2smem_mc_cta2(v_dst + group_2 * 512, (&tmap_compressed_kv), v_col, row0_2, row1_2, row2_2, row3_2, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } - } - load_v_stage += 1; - if (load_v_stage == 2) { load_v_stage = 0; _phase_v_empty ^= 1; } - if ((prev_tile & 1) != 0) { - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(index_empty_addr + (load_v_index_stage) * 8); - load_v_index_stage = load_v_index_stage + 1; - if (load_v_index_stage == 6) { - load_v_index_stage = 0; - } - } - } - } - int last_tile = num_kv_tiles - 1; - int global_last_tile = first_tile_1 + last_tile; - { - mbarrier_wait(v_empty_addr + (load_v_stage) * 8, _phase_v_empty); - if (load_warp_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - if (lane < 16) { - int v_dst_1 = smem_v_full_addr + load_v_stage * 32768 + (unsigned int)(load_warp_rank / 2 * 16384); - int v_col_1 = v_chunk * 256 + load_warp_rank / 2 * 128; - int group_3 = (load_warp_rank & 1) * 16 + lane; - int group_offset_3 = (last_tile & 1) * 128 + group_3 * 4; - int raw_rows_3[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_3[0])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 3])) - : "r"(smem_sparse_indices_addr + load_v_index_stage * 1024 + (unsigned int)(group_offset_3 * 4))); - int row0_3 = ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : 0); - int row1_3 = ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : 0); - int row2_3 = ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : 0); - int row3_3 = ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : 0); - if (global_last_tile == 0) { - tma_gather4_gmem2smem_mc_cta2(v_dst_1 + group_3 * 512, (&tmap_swa_kv), v_col_1, row0_3, row1_3, row2_3, row3_3, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } else { - tma_gather4_gmem2smem_mc_cta2(v_dst_1 + group_3 * 512, (&tmap_compressed_kv), v_col_1, row0_3, row1_3, row2_3, row3_3, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } - } - load_v_stage += 1; - if (load_v_stage == 2) { load_v_stage = 0; _phase_v_empty ^= 1; } asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(index_empty_addr + (load_v_index_stage) * 8); - load_v_index_stage = load_v_index_stage + 1; - if (load_v_index_stage == 6) { - load_v_index_stage = 0; - } + mbarrier_arrive(index_empty_addr + (load_k_index_stage) * 8); + load_k_index_stage += 1; + if (load_k_index_stage == 6) { load_k_index_stage = 0; _phase_index_full ^= 1; } } } mbarrier_arrive(tmem_dealloc_addr); @@ -1738,14 +1581,14 @@ kernel_cake_dsv4_c20607c57cc11784015e(const __grid_constant__ CUtensorMap tmap_q int active_topk = _min_0; float row_max_val = -CAKE_INF; float row_sum_val = 0.0f; - int sink_head = ((1) ? bid % 2 * 64 + my_row : my_row); - if (has_sinks != 0 && sink_head < num_heads && split_idx_2 == 0) { + int sink_head = ((1) ? my_row : my_row); + if (has_sinks != 0 && sink_head < num_heads && split_idx_2 + bid % 2 == 0) { row_max_val = sinks[sink_head] * 1.4426950408889634f / softmax_scale_log2; row_sum_val = 1.0f; } #pragma unroll 1 - for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { - int pipeline_tile = softmax_tile_cursor + tile_1; + for (int tile = 0; tile < (num_kv_tiles_1 - bid % 2 + 1) / 2; tile++) { + int pipeline_tile = softmax_tile_cursor + tile; int phase = pipeline_tile & 1; int s_wait_phase = pipeline_tile >> 1 & 1; mbarrier_wait(s_full_addr + (phase) * 8, s_wait_phase); @@ -1754,14 +1597,14 @@ kernel_cake_dsv4_c20607c57cc11784015e(const __grid_constant__ CUtensorMap tmap_q int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_score_row_base << 16); float new_max = row_max_val; int valid_sparse_cols = ((active_topk < sparse_topk) ? active_topk : sparse_topk); - valid_sparse_cols = valid_sparse_cols - (first_tile_2 + tile_1) * 128 - n_half * 64; + valid_sparse_cols = valid_sparse_cols - (first_tile_2 + tile * 2 + bid % 2) * 128 - n_half * 64; if (valid_sparse_cols < 0) { valid_sparse_cols = 0; } if (valid_sparse_cols > 64) { valid_sparse_cols = 64; } - int index_tile_addr = smem_sparse_indices_addr + softmax_index_stage * 1024 + (unsigned int)((tile_1 & 1) * 512); + int index_tile_addr = smem_sparse_indices_addr + softmax_index_stage * 1024 + (unsigned int)(bid % 2 * 512); int index_lane_addr = index_tile_addr + (n_half * 64 + lane) * 4; int staged_index[2]; unsigned int invalid_cols[2]; @@ -1773,7 +1616,7 @@ kernel_cake_dsv4_c20607c57cc11784015e(const __grid_constant__ CUtensorMap tmap_q invalid_cols[w] = _vote_0; any_invalid = any_invalid | invalid_cols[w]; } - if ((tile_1 & 1) != 0 || tile_1 + 1 == num_kv_tiles_1) { + if (tile >= 0) { asm volatile("tcgen05.fence::before_thread_sync;"); mbarrier_arrive(index_empty_addr + (softmax_index_stage) * 8); softmax_index_stage = softmax_index_stage + 1; @@ -2225,7 +2068,7 @@ kernel_cake_dsv4_c20607c57cc11784015e(const __grid_constant__ CUtensorMap tmap_q smem_stats_final_max[stats_row] = row_max_val; asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); mbarrier_arrive(sum_ready_addr); - softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles_1; + softmax_tile_cursor = softmax_tile_cursor + (num_kv_tiles_1 - bid % 2 + 1) / 2; } asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); asm volatile("tcgen05.fence::before_thread_sync;"); @@ -2255,6 +2098,8 @@ kernel_cake_dsv4_c20607c57cc11784015e(const __grid_constant__ CUtensorMap tmap_q unsigned int _phase_pv_done_0 = 0; unsigned int _phase_o_done_0 = 0; unsigned int _phase_sum_ready_0 = 0; + unsigned int _phase_merge_ready_0 = 0; + unsigned int _phase_merge_full_0 = 0; #pragma unroll for (unsigned int work_idx_2 = 0; work_idx_2 < 1; work_idx_2++) { int split_idx_3 = 0; @@ -2269,19 +2114,19 @@ kernel_cake_dsv4_c20607c57cc11784015e(const __grid_constant__ CUtensorMap tmap_q int first_tile_3 = split_idx_3 * tiles_per_split_3; int num_kv_tiles_2 = tiles_per_split_3; #pragma unroll 1 - for (int tile_2 = 0; tile_2 < num_kv_tiles_2; tile_2++) { - int pipeline_tile_1 = correction_tile_cursor + tile_2; + for (int tile_1 = 0; tile_1 < (num_kv_tiles_2 - bid % 2 + 1) / 2; tile_1++) { + int pipeline_tile_1 = correction_tile_cursor + tile_1; int phase_1 = pipeline_tile_1 & 1; int stats_wait_phase = pipeline_tile_1 >> 1 & 1; mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); asm volatile("tcgen05.fence::after_thread_sync;"); float acc_scale_1 = smem_stats_max[phase_1 * 128 + stats_row_1]; - if (tile_2 > 0) { + if (tile_1 > 0) { mbarrier_wait(pv_done_addr, _phase_pv_done_0); _phase_pv_done_0 ^= 1; asm volatile("tcgen05.fence::after_thread_sync;"); } - if (tile_2 > 0) { + if (tile_1 > 0) { int _vote_1 = __any_sync(0xFFFFFFFF, acc_scale_1 < 1.0f); int any_rescale = _vote_1; if (any_rescale != 0) { @@ -2310,60 +2155,157 @@ kernel_cake_dsv4_c20607c57cc11784015e(const __grid_constant__ CUtensorMap tmap_q mbarrier_wait(sum_ready_addr, _phase_sum_ready_0); _phase_sum_ready_0 ^= 1; asm volatile("tcgen05.fence::after_thread_sync;"); - float total_sum = smem_stats_sum[stats_row_1]; - float final_max = smem_stats_final_max[stats_row_1]; - { - smem_corr_warp_pair_exchange[stats_row_1] = total_sum; - asm volatile("barrier.sync %0, 64;" :: "r"(4 + warp % 2) : "memory"); - total_sum = total_sum + smem_corr_warp_pair_exchange[stats_row_1 ^ 64]; + float wt_sum = smem_stats_sum[stats_row_1]; + float wt_max = smem_stats_final_max[stats_row_1]; + smem_corr_warp_pair_exchange[stats_row_1] = wt_sum; + asm volatile("barrier.sync %0, 64;" :: "r"(4 + warp % 2) : "memory"); + wt_sum = wt_sum + smem_corr_warp_pair_exchange[stats_row_1 ^ 64]; + int wt_peer_rank = bid % 2 ^ 1; + int wt_quadrant = warp - 4; + if (warp == 4) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(merge_full_addr, 66048); + } } - float _rcp_0 = approx_rcp(total_sum); - float inv_sum = ((total_sum > 0.0f) ? _rcp_0 : 0.0f); - int head_idx = my_row_1; - { - head_idx = bid % 2 * 64 + my_row_1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(merge_ready_addr), "r"(wt_peer_rank) : "memory"); + if (n_half_1 == 0) { + smem_merge_stats_send[my_row_1] = wt_max; + smem_merge_stats_send[64 + my_row_1] = wt_sum; } - int direct_o_offset = (query_idx_3 * num_heads + head_idx) * 512 + v_chunk_1 * 256; - int partial_o_offset = (query_idx_3 * num_heads + head_idx + split_idx_3) * 512 + ((0) ? v_chunk_1 : n_half_1) * 256; - int o_offset = partial_o_offset; + int wt_send_block0 = (bid % 2 ^ 1) * 2; #pragma unroll - for (int vs_1 = 0; vs_1 < 2; vs_1++) { - int o_base_epi = taddr + 256 + (unsigned int)(vs_1 * 128) + (unsigned int)corr_row; + for (int wt_b = 0; wt_b < 2; wt_b++) { #pragma unroll - for (int c_1 = 0; c_1 < 128; c_1 += 32) { + for (int wt_c = 0; wt_c < 64; wt_c += 32) { float _tmem_load_4[32]; asm volatile( "tcgen05.ld.sync.aligned.32x32b.x32.b32" " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" : "=f"(_tmem_load_4[0]), "=f"(_tmem_load_4[1]), "=f"(_tmem_load_4[2]), "=f"(_tmem_load_4[3]), "=f"(_tmem_load_4[4]), "=f"(_tmem_load_4[5]), "=f"(_tmem_load_4[6]), "=f"(_tmem_load_4[7]), "=f"(_tmem_load_4[8]), "=f"(_tmem_load_4[9]), "=f"(_tmem_load_4[10]), "=f"(_tmem_load_4[11]), "=f"(_tmem_load_4[12]), "=f"(_tmem_load_4[13]), "=f"(_tmem_load_4[14]), "=f"(_tmem_load_4[15]), "=f"(_tmem_load_4[16]), "=f"(_tmem_load_4[17]), "=f"(_tmem_load_4[18]), "=f"(_tmem_load_4[19]), "=f"(_tmem_load_4[20]), "=f"(_tmem_load_4[21]), "=f"(_tmem_load_4[22]), "=f"(_tmem_load_4[23]), "=f"(_tmem_load_4[24]), "=f"(_tmem_load_4[25]), "=f"(_tmem_load_4[26]), "=f"(_tmem_load_4[27]), "=f"(_tmem_load_4[28]), "=f"(_tmem_load_4[29]), "=f"(_tmem_load_4[30]), "=f"(_tmem_load_4[31]) - : "r"(o_base_epi + c_1)); - int gmem_base = o_offset + vs_1 * 128 + c_1; - if (head_idx < num_heads) { + : "r"(taddr + 256 + (unsigned int)((wt_send_block0 + wt_b) * 64) + (unsigned int)wt_c + (unsigned int)corr_row)); + #pragma unroll + for (int wt_k = 0; wt_k < 32; wt_k += 4) { + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(smem_merge_send_addr + (unsigned int)((((wt_b * 64 + wt_c + wt_k) / 4 * 4 + wt_quadrant) * 32 + lane) * 16)), "f"(_tmem_load_4[wt_k]), "f"(_tmem_load_4[wt_k + 1]), "f"(_tmem_load_4[wt_k + 2]), "f"(_tmem_load_4[wt_k + 3]) : "memory"); + } + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 6, 64;" ::: "memory"); + if (warp == 4) { + if (elect_sync()) { + mbarrier_wait(merge_ready_addr, _phase_merge_ready_0); + _phase_merge_ready_0 ^= 1; + uint32_t _mapa_0; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_0) : "r"(smem_merge_recv_addr), "r"(wt_peer_rank)); + uint32_t _mapa_1; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_1) : "r"(merge_full_addr), "r"(wt_peer_rank)); + asm volatile( + "cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes" + " [%0], [%1], %2, [%3];" + :: "r"(_mapa_0), "r"(smem_merge_send_addr), "r"((uint32_t)(65536)), "r"(_mapa_1) + : "memory"); + uint32_t _mapa_2; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_2) : "r"(smem_merge_stats_recv_addr), "r"(wt_peer_rank)); + uint32_t _mapa_3; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_3) : "r"(merge_full_addr), "r"(wt_peer_rank)); + asm volatile( + "cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes" + " [%0], [%1], %2, [%3];" + :: "r"(_mapa_2), "r"(smem_merge_stats_send_addr), "r"((uint32_t)(512)), "r"(_mapa_3) + : "memory"); + } + } + mbarrier_wait(merge_full_addr, _phase_merge_full_0); + _phase_merge_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + float wt_peer_max = smem_merge_stats_recv[my_row_1]; + float wt_peer_sum = smem_merge_stats_recv[64 + my_row_1]; + float _max_5 = max_noftz(wt_max, wt_peer_max); + float wt_merged_max = _max_5; + float _exp2_1 = approx_exp2((wt_max - wt_merged_max) * softmax_scale_log2_1); + float wt_s_own = ((wt_max > -CAKE_INF) ? _exp2_1 : 0.0f); + float _exp2_2 = approx_exp2((wt_peer_max - wt_merged_max) * softmax_scale_log2_1); + float wt_s_peer = ((wt_peer_max > -CAKE_INF) ? _exp2_2 : 0.0f); + float wt_merged_sum = wt_sum * wt_s_own + wt_peer_sum * wt_s_peer; + float _rcp_0 = approx_rcp(wt_merged_sum); + float wt_inv_sum = ((wt_merged_sum > 0.0f) ? _rcp_0 : 0.0f); + int wt_head_idx = my_row_1; + int wt_o_base = (query_idx_3 * num_heads + wt_head_idx) * 512; + if (bid % 2 == 0 && n_half_1 == 0 && wt_head_idx < num_heads) { + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(wt_merged_sum)); + partial_lse[query_idx_3 * num_heads + wt_head_idx + split_idx_3] = ((wt_merged_sum > 0.0f) ? wt_merged_max * softmax_scale_log2_1 + _log2_0 : -CAKE_INF); + } + int wt_keep_block0 = bid % 2 * 2; + #pragma unroll + for (int wt_b_1 = 0; wt_b_1 < 2; wt_b_1++) { + #pragma unroll + for (int wt_c_1 = 0; wt_c_1 < 64; wt_c_1 += 32) { + float _tmem_load_5[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_5[0]), "=f"(_tmem_load_5[1]), "=f"(_tmem_load_5[2]), "=f"(_tmem_load_5[3]), "=f"(_tmem_load_5[4]), "=f"(_tmem_load_5[5]), "=f"(_tmem_load_5[6]), "=f"(_tmem_load_5[7]), "=f"(_tmem_load_5[8]), "=f"(_tmem_load_5[9]), "=f"(_tmem_load_5[10]), "=f"(_tmem_load_5[11]), "=f"(_tmem_load_5[12]), "=f"(_tmem_load_5[13]), "=f"(_tmem_load_5[14]), "=f"(_tmem_load_5[15]), "=f"(_tmem_load_5[16]), "=f"(_tmem_load_5[17]), "=f"(_tmem_load_5[18]), "=f"(_tmem_load_5[19]), "=f"(_tmem_load_5[20]), "=f"(_tmem_load_5[21]), "=f"(_tmem_load_5[22]), "=f"(_tmem_load_5[23]), "=f"(_tmem_load_5[24]), "=f"(_tmem_load_5[25]), "=f"(_tmem_load_5[26]), "=f"(_tmem_load_5[27]), "=f"(_tmem_load_5[28]), "=f"(_tmem_load_5[29]), "=f"(_tmem_load_5[30]), "=f"(_tmem_load_5[31]) + : "r"(taddr + 256 + (unsigned int)((wt_keep_block0 + wt_b_1) * 64) + (unsigned int)wt_c_1 + (unsigned int)corr_row)); + float wt_peer[32]; + #pragma unroll + for (int wt_k_1 = 0; wt_k_1 < 32; wt_k_1 += 4) { + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&wt_peer[wt_k_1])), "=r"(*reinterpret_cast(&wt_peer[(wt_k_1) + 1])), "=r"(*reinterpret_cast(&wt_peer[(wt_k_1) + 2])), "=r"(*reinterpret_cast(&wt_peer[(wt_k_1) + 3])) + : "r"(smem_merge_recv_addr + (unsigned int)((((wt_b_1 * 64 + wt_c_1 + wt_k_1) / 4 * 4 + wt_quadrant) * 32 + lane) * 16))); + } + const float2 _scale2_1 = {wt_s_own, wt_s_own}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + mul_f32x2_inplace(&reinterpret_cast(_tmem_load_5)[_ls], _scale2_1); + const float2 _scale2_2 = {wt_s_peer, wt_s_peer}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + mul_f32x2_inplace(&reinterpret_cast(wt_peer)[_ls], _scale2_2); + #pragma unroll + for (int _la = 0; _la < 16; _la++) + add_f32x2_inplace(&reinterpret_cast(_tmem_load_5)[_la], reinterpret_cast(wt_peer)[_la]); + int wt_gmem = wt_o_base + (wt_keep_block0 + wt_b_1) * 128 + n_half_1 * 64 + wt_c_1; + if (wt_head_idx < num_heads) { #pragma unroll - for (int j = 0; j < 32; j += 8) { + for (int wt_j = 0; wt_j < 32; wt_j += 8) { { - const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; + const float2 _prescale2_3 = {wt_inv_sum * output_scale, wt_inv_sum * output_scale}; #if __CUDA_ARCH__ >= 1000 #pragma unroll for (int _ps = 0; _ps < 4; _ps++) - mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_4[j])[_ps], _prescale2_1); + mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_5[wt_j])[_ps], _prescale2_3); #else #pragma unroll for (int _ps = 0; _ps < 8; _ps++) - _tmem_load_4[j + _ps] *= inv_sum * output_scale; + _tmem_load_5[wt_j + _ps] *= wt_inv_sum * output_scale; #endif __nv_bfloat162 _pk[4]; - _pk[0] = __floats2bfloat162_rn(_tmem_load_4[j + 0], _tmem_load_4[j + 1]); - _pk[1] = __floats2bfloat162_rn(_tmem_load_4[j + 2], _tmem_load_4[j + 3]); - _pk[2] = __floats2bfloat162_rn(_tmem_load_4[j + 4], _tmem_load_4[j + 5]); - _pk[3] = __floats2bfloat162_rn(_tmem_load_4[j + 6], _tmem_load_4[j + 7]); - *reinterpret_cast(&((__nv_bfloat16*)(O + (gmem_base + j)))[0]) = *reinterpret_cast(&_pk[0]); + _pk[0] = __floats2bfloat162_rn(_tmem_load_5[wt_j + 0], _tmem_load_5[wt_j + 1]); + _pk[1] = __floats2bfloat162_rn(_tmem_load_5[wt_j + 2], _tmem_load_5[wt_j + 3]); + _pk[2] = __floats2bfloat162_rn(_tmem_load_5[wt_j + 4], _tmem_load_5[wt_j + 5]); + _pk[3] = __floats2bfloat162_rn(_tmem_load_5[wt_j + 6], _tmem_load_5[wt_j + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(O + (wt_gmem + wt_j)))[0]) = *reinterpret_cast(&_pk[0]); } } } } } - correction_tile_cursor = correction_tile_cursor + num_kv_tiles_2; + correction_tile_cursor = correction_tile_cursor + (num_kv_tiles_2 - bid % 2 + 1) / 2; } asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); asm volatile("tcgen05.fence::before_thread_sync;"); @@ -2385,13 +2327,8 @@ kernel_cake_dsv4_c20607c57cc11784015e(const __grid_constant__ CUtensorMap tmap_q unsigned int mma_v_stage = 0; unsigned int mma_kv_stage = 0; int mma_tile_cursor = 0; - unsigned int _phase_s_seeded_0 = 0; unsigned int _phase_q_full_0_3 = 0; - unsigned int _phase_q_pair_ready_0 = 0; unsigned int _phase_k_full = 0; - unsigned int _phase_k_full_1 = 0; - unsigned int _phase_kv_pair_ready = 0; - unsigned int _phase_v_full = 0; #pragma unroll for (unsigned int work_idx_3 = 0; work_idx_3 < 1; work_idx_3++) { int split_idx_4 = 0; @@ -2401,410 +2338,487 @@ kernel_cake_dsv4_c20607c57cc11784015e(const __grid_constant__ CUtensorMap tmap_q int tiles_per_split_4 = all_num_kv_tiles_4 + 1 - 1; int first_tile_4 = split_idx_4 * tiles_per_split_4; int num_kv_tiles_3 = tiles_per_split_4; - { - if (cta_rank == 0) { - mbarrier_wait(q_full_addr, _phase_q_full_0_3); - _phase_q_full_0_3 ^= 1; - } - } - int first_pv = 1; + int wt_num_local_tiles_1 = (num_kv_tiles_3 - bid % 2 + 1) / 2; + mbarrier_wait(q_full_addr, _phase_q_full_0_3); + _phase_q_full_0_3 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int wt_first_pv = 1; #pragma unroll 1 - for (int tile_3 = 0; tile_3 < num_kv_tiles_3; tile_3++) { - int pipeline_tile_2 = mma_tile_cursor + tile_3; - int phase_2 = pipeline_tile_2 & 1; - int score_col = ((phase_2 != 0) ? 128 : 0); - { - if (cta_rank == 0) { - mbarrier_wait(k_full_addr + (mma_k_stage) * 8, _phase_k_full); - int _mma_a_lo_0 = (((smem_q_full_addr) >> 4) & 0x3FFF) + (0) * 2048; - int _mma_b_lo_0 = (((smem_k_full_addr) >> 4) & 0x3FFF) + (mma_k_stage) * 2048; - asm volatile( + for (int wt_tile_1 = 0; wt_tile_1 < wt_num_local_tiles_1; wt_tile_1++) { + int wt_pipeline_tile = mma_tile_cursor + wt_tile_1; + int wt_phase = wt_pipeline_tile & 1; + int wt_score_col = ((wt_phase != 0) ? 128 : 0); + mbarrier_wait(k_full_addr + (mma_k_stage) * 8, _phase_k_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_0 = make_warp_uniform((((smem_q_full_addr) >> 4) & 0x3FFF) + (0) * 2048); + int _mma_b_lo_0 = make_warp_uniform((((smem_k_full_addr) >> 4) & 0x3FFF) + (mma_k_stage) * 4096); + asm volatile( "{\n\t" ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" ".reg .b64 da, db;\n\t" "elect.sync _|leader, 0xFFFFFFFF;\n\t" "setp.ne.b32 p0, %3, 0;\n\t" "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "" "mov.b32 adhi, 0x40004040;\n\t" "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 136314896;\n\t" + "mov.b32 id, 69206032;\n\t" "mov.b32 alo, %0;\n\t" "mov.b32 blo, %1;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 2;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 2;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 2;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "}\n" - :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(0)); - int _mma_a_lo_1 = (((smem_q_full_addr + 8192) >> 4) & 0x3FFF) + (0) * 2048; - int _mma_b_lo_1 = (((smem_k_full_addr + 8192) >> 4) & 0x3FFF) + (mma_k_stage) * 2048; - asm volatile( + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (wt_score_col))), "r"(0)); + int _mma_a_lo_1 = make_warp_uniform((((smem_q_full_addr + 8192) >> 4) & 0x3FFF) + (0) * 2048); + int _mma_b_lo_1 = make_warp_uniform((((smem_k_full_addr + 16384) >> 4) & 0x3FFF) + (mma_k_stage) * 4096); + asm volatile( "{\n\t" ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" ".reg .b64 da, db;\n\t" "elect.sync _|leader, 0xFFFFFFFF;\n\t" "setp.ne.b32 p0, %3, 0;\n\t" "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "" "mov.b32 adhi, 0x40004040;\n\t" "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 136314896;\n\t" + "mov.b32 id, 69206032;\n\t" "mov.b32 alo, %0;\n\t" "mov.b32 blo, %1;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 2;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 2;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 2;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "}\n" - :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); - int _mma_a_lo_2 = (((smem_q_full_addr + 16384) >> 4) & 0x3FFF) + (0) * 2048; - int _mma_b_lo_2 = (((smem_k_full_addr + 16384) >> 4) & 0x3FFF) + (mma_k_stage) * 2048; - asm volatile( + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_tmem_scratch + (wt_score_col))), "r"(1)); + int _mma_a_lo_2 = make_warp_uniform((((smem_q_full_addr + 16384) >> 4) & 0x3FFF) + (0) * 2048); + int _mma_b_lo_2 = make_warp_uniform((((smem_k_full_addr + 32768) >> 4) & 0x3FFF) + (mma_k_stage) * 4096); + asm volatile( "{\n\t" ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" ".reg .b64 da, db;\n\t" "elect.sync _|leader, 0xFFFFFFFF;\n\t" "setp.ne.b32 p0, %3, 0;\n\t" "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "" "mov.b32 adhi, 0x40004040;\n\t" "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 136314896;\n\t" + "mov.b32 id, 69206032;\n\t" "mov.b32 alo, %0;\n\t" "mov.b32 blo, %1;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 2;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 2;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 2;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "}\n" - :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); - int _mma_a_lo_3 = (((smem_q_full_addr + 24576) >> 4) & 0x3FFF) + (0) * 2048; - int _mma_b_lo_3 = (((smem_k_full_addr + 24576) >> 4) & 0x3FFF) + (mma_k_stage) * 2048; - asm volatile( + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (wt_score_col))), "r"(1)); + int _mma_a_lo_3 = make_warp_uniform((((smem_q_full_addr + 24576) >> 4) & 0x3FFF) + (0) * 2048); + int _mma_b_lo_3 = make_warp_uniform((((smem_k_full_addr + 49152) >> 4) & 0x3FFF) + (mma_k_stage) * 4096); + asm volatile( "{\n\t" ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" ".reg .b64 da, db;\n\t" "elect.sync _|leader, 0xFFFFFFFF;\n\t" "setp.ne.b32 p0, %3, 0;\n\t" "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "" "mov.b32 adhi, 0x40004040;\n\t" "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 136314896;\n\t" + "mov.b32 id, 69206032;\n\t" "mov.b32 alo, %0;\n\t" "mov.b32 blo, %1;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 2;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 2;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 2;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "}\n" - :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); - if (cta_rank == 0) { - elect_commit_cg2_multicast(s_full_addr + (phase_2) * 8, (uint16_t)(3)); - } - if (cta_rank == 0) { - elect_commit_cg2_multicast(k_empty_addr + (mma_k_stage) * 8, (uint16_t)(3)); - } - } - mma_k_stage += 1; - if (mma_k_stage == 2) { mma_k_stage = 0; _phase_k_full ^= 1; } - } - if (tile_3 > 0) { - int prev_pipeline_tile = pipeline_tile_2 - 1; - int prev_phase = prev_pipeline_tile & 1; - int pv_wait_phase = prev_pipeline_tile >> 1 & 1; - mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); - mbarrier_wait(corr_done_addr + (prev_phase) * 8, pv_wait_phase); + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"((tmem_tmem_scratch + (wt_score_col))), "r"(1)); + elect_commit(s_full_addr + (wt_phase) * 8); + if (wt_tile_1 > 0) { + int wt_prev_pipeline_tile = wt_pipeline_tile - 1; + int wt_prev_phase = wt_prev_pipeline_tile & 1; + int wt_pv_wait_phase = wt_prev_pipeline_tile >> 1 & 1; + mbarrier_wait(p_full_addr + (wt_prev_phase) * 8, wt_pv_wait_phase); + mbarrier_wait(corr_done_addr + (wt_prev_phase) * 8, wt_pv_wait_phase); asm volatile("tcgen05.fence::after_thread_sync;"); - { - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(pv_pair_ready_addr + (unsigned int)(prev_phase * 8)), "r"(0) : "memory"); - if (cta_rank == 0) { - mbarrier_wait_cluster_hint(pv_pair_ready_addr + (prev_phase) * 8, pv_wait_phase, 10000000); - } - } - int p_col = ((prev_phase != 0) ? 224 : 96); - { - if (cta_rank == 0) { - mbarrier_wait(v_full_addr + (mma_v_stage) * 8, _phase_v_full); - int _mma_a_lo_6 = (((smem_p_fp8_addr) >> 4) & 0x3FFF) + (prev_phase) * 512; - int _mma_b_lo_6 = ((((smem_v_full_addr) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; - asm volatile( + unsigned int wt_prev_kv_stage = mma_k_stage ^ 1; + int _mma_a_lo_4 = make_warp_uniform((((smem_p_fp8_addr) >> 4) & 0x3FFF) + (wt_prev_phase) * 512); + int _mma_b_lo_4 = make_warp_uniform(((((smem_v_full_addr) >> 4) & 0x3FFF) | 0x4000000) + (wt_prev_kv_stage) * 4096); + asm volatile( "{\n\t" ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" ".reg .b64 da, db;\n\t" "elect.sync _|leader, 0xFFFFFFFF;\n\t" "setp.ne.b32 p0, %3, 0;\n\t" "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "" "mov.b32 adhi, 0x40004040;\n\t" "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138477584;\n\t" + "mov.b32 id, 69271568;\n\t" "mov.b32 alo, %0;\n\t" "mov.b32 blo, %1;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 256;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 256;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 256;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "}\n" - :: "r"(_mma_a_lo_6), "r"(_mma_b_lo_6), "r"((tmem_tmem_scratch + (256))), "r"(((first_pv) ? 0 : 1))); - int _mma_a_lo_7 = (((smem_p_fp8_addr) >> 4) & 0x3FFF) + (prev_phase) * 512; - int _mma_b_lo_7 = ((((smem_v_full_addr + 16384) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; - asm volatile( + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"((tmem_tmem_scratch + (256))), "r"(((wt_first_pv) ? 0 : 1))); + int _mma_a_lo_5 = make_warp_uniform((((smem_p_fp8_addr) >> 4) & 0x3FFF) + (wt_prev_phase) * 512); + int _mma_b_lo_5 = make_warp_uniform(((((smem_v_full_addr + 16384) >> 4) & 0x3FFF) | 0x4000000) + (wt_prev_kv_stage) * 4096); + asm volatile( "{\n\t" ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" ".reg .b64 da, db;\n\t" "elect.sync _|leader, 0xFFFFFFFF;\n\t" "setp.ne.b32 p0, %3, 0;\n\t" "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "" "mov.b32 adhi, 0x40004040;\n\t" "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138477584;\n\t" + "mov.b32 id, 69271568;\n\t" "mov.b32 alo, %0;\n\t" "mov.b32 blo, %1;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 256;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 256;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 256;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "}\n" - :: "r"(_mma_a_lo_7), "r"(_mma_b_lo_7), "r"((tmem_tmem_scratch + (384))), "r"(((first_pv) ? 0 : 1))); - if (cta_rank == 0) { - elect_commit_cg2_multicast(v_empty_addr + (mma_v_stage) * 8, (uint16_t)(3)); - } - } - mma_v_stage += 1; - if (mma_v_stage == 2) { mma_v_stage = 0; _phase_v_full ^= 1; } - } - first_pv = 0; - if (cta_rank == 0) { - elect_commit_cg2_multicast(pv_done_addr, (uint16_t)(3)); - } - } - } - int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_3 - 1; - int last_phase = last_pipeline_tile & 1; - int drain_wait_phase = last_pipeline_tile >> 1 & 1; - mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); - mbarrier_wait(corr_done_addr + (last_phase) * 8, drain_wait_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - { - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(pv_pair_ready_addr + (unsigned int)(last_phase * 8)), "r"(0) : "memory"); - if (cta_rank == 0) { - mbarrier_wait_cluster_hint(pv_pair_ready_addr + (last_phase) * 8, drain_wait_phase, 10000000); - } - } - int p_col_last = ((last_phase != 0) ? 224 : 96); - { - if (cta_rank == 0) { - mbarrier_wait(v_full_addr + (mma_v_stage) * 8, _phase_v_full); - int _mma_a_lo_10 = (((smem_p_fp8_addr) >> 4) & 0x3FFF) + (last_phase) * 512; - int _mma_b_lo_10 = ((((smem_v_full_addr) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; + :: "r"(_mma_a_lo_5), "r"(_mma_b_lo_5), "r"((tmem_tmem_scratch + (320))), "r"(((wt_first_pv) ? 0 : 1))); + int _mma_a_lo_6 = make_warp_uniform((((smem_p_fp8_addr) >> 4) & 0x3FFF) + (wt_prev_phase) * 512); + int _mma_b_lo_6 = make_warp_uniform(((((smem_v_full_addr + 32768) >> 4) & 0x3FFF) | 0x4000000) + (wt_prev_kv_stage) * 4096); asm volatile( "{\n\t" ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" ".reg .b64 da, db;\n\t" "elect.sync _|leader, 0xFFFFFFFF;\n\t" "setp.ne.b32 p0, %3, 0;\n\t" "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "" "mov.b32 adhi, 0x40004040;\n\t" "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138477584;\n\t" + "mov.b32 id, 69271568;\n\t" "mov.b32 alo, %0;\n\t" "mov.b32 blo, %1;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 256;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 256;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 256;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "}\n" - :: "r"(_mma_a_lo_10), "r"(_mma_b_lo_10), "r"((tmem_tmem_scratch + (256))), "r"(((first_pv) ? 0 : 1))); - int _mma_a_lo_11 = (((smem_p_fp8_addr) >> 4) & 0x3FFF) + (last_phase) * 512; - int _mma_b_lo_11 = ((((smem_v_full_addr + 16384) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; + :: "r"(_mma_a_lo_6), "r"(_mma_b_lo_6), "r"((tmem_tmem_scratch + (384))), "r"(((wt_first_pv) ? 0 : 1))); + int _mma_a_lo_7 = make_warp_uniform((((smem_p_fp8_addr) >> 4) & 0x3FFF) + (wt_prev_phase) * 512); + int _mma_b_lo_7 = make_warp_uniform(((((smem_v_full_addr + 49152) >> 4) & 0x3FFF) | 0x4000000) + (wt_prev_kv_stage) * 4096); asm volatile( "{\n\t" ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" ".reg .b64 da, db;\n\t" "elect.sync _|leader, 0xFFFFFFFF;\n\t" "setp.ne.b32 p0, %3, 0;\n\t" "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "" "mov.b32 adhi, 0x40004040;\n\t" "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138477584;\n\t" + "mov.b32 id, 69271568;\n\t" "mov.b32 alo, %0;\n\t" "mov.b32 blo, %1;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 256;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 256;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "add.u32 alo, alo, 2;\n\t" "add.u32 blo, blo, 256;\n\t" "mov.b64 da, {alo, adhi};\n\t" "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" "}\n" - :: "r"(_mma_a_lo_11), "r"(_mma_b_lo_11), "r"((tmem_tmem_scratch + (384))), "r"(((first_pv) ? 0 : 1))); - if (cta_rank == 0) { - elect_commit_cg2_multicast(v_empty_addr + (mma_v_stage) * 8, (uint16_t)(3)); - } + :: "r"(_mma_a_lo_7), "r"(_mma_b_lo_7), "r"((tmem_tmem_scratch + (448))), "r"(((wt_first_pv) ? 0 : 1))); + wt_first_pv = 0; + elect_commit(k_empty_addr + (wt_prev_kv_stage) * 8); + elect_commit(pv_done_addr); } - mma_v_stage += 1; - if (mma_v_stage == 2) { mma_v_stage = 0; _phase_v_full ^= 1; } + mma_k_stage += 1; + if (mma_k_stage == 2) { mma_k_stage = 0; _phase_k_full ^= 1; } } - if (cta_rank == 0) { - elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); - } - if (cta_rank == 0) { - elect_commit_cg2_multicast(o_done_addr, (uint16_t)(3)); + if (wt_num_local_tiles_1 > 0) { + int wt_last_pipeline_tile = mma_tile_cursor + wt_num_local_tiles_1 - 1; + int wt_last_phase = wt_last_pipeline_tile & 1; + int wt_drain_wait_phase = wt_last_pipeline_tile >> 1 & 1; + mbarrier_wait(p_full_addr + (wt_last_phase) * 8, wt_drain_wait_phase); + mbarrier_wait(corr_done_addr + (wt_last_phase) * 8, wt_drain_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int wt_last_kv_stage = mma_k_stage ^ 1; + int _mma_a_lo_8 = make_warp_uniform((((smem_p_fp8_addr) >> 4) & 0x3FFF) + (wt_last_phase) * 512); + int _mma_b_lo_8 = make_warp_uniform(((((smem_v_full_addr) >> 4) & 0x3FFF) | 0x4000000) + (wt_last_kv_stage) * 4096); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69271568;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_8), "r"(_mma_b_lo_8), "r"((tmem_tmem_scratch + (256))), "r"(((wt_first_pv) ? 0 : 1))); + int _mma_a_lo_9 = make_warp_uniform((((smem_p_fp8_addr) >> 4) & 0x3FFF) + (wt_last_phase) * 512); + int _mma_b_lo_9 = make_warp_uniform(((((smem_v_full_addr + 16384) >> 4) & 0x3FFF) | 0x4000000) + (wt_last_kv_stage) * 4096); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69271568;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_9), "r"(_mma_b_lo_9), "r"((tmem_tmem_scratch + (320))), "r"(((wt_first_pv) ? 0 : 1))); + int _mma_a_lo_10 = make_warp_uniform((((smem_p_fp8_addr) >> 4) & 0x3FFF) + (wt_last_phase) * 512); + int _mma_b_lo_10 = make_warp_uniform(((((smem_v_full_addr + 32768) >> 4) & 0x3FFF) | 0x4000000) + (wt_last_kv_stage) * 4096); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69271568;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_10), "r"(_mma_b_lo_10), "r"((tmem_tmem_scratch + (384))), "r"(((wt_first_pv) ? 0 : 1))); + int _mma_a_lo_11 = make_warp_uniform((((smem_p_fp8_addr) >> 4) & 0x3FFF) + (wt_last_phase) * 512); + int _mma_b_lo_11 = make_warp_uniform(((((smem_v_full_addr + 49152) >> 4) & 0x3FFF) | 0x4000000) + (wt_last_kv_stage) * 4096); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69271568;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_11), "r"(_mma_b_lo_11), "r"((tmem_tmem_scratch + (448))), "r"(((wt_first_pv) ? 0 : 1))); + elect_commit(k_empty_addr + (wt_last_kv_stage) * 8); } - mma_tile_cursor = mma_tile_cursor + num_kv_tiles_3; + elect_commit(q_empty_addr); + elect_commit(o_done_addr); + mma_tile_cursor = mma_tile_cursor + wt_num_local_tiles_1; } mbarrier_arrive(tmem_dealloc_addr); unsigned int _phase_tmem_dealloc_0 = 0; mbarrier_wait(tmem_dealloc_addr, _phase_tmem_dealloc_0); _phase_tmem_dealloc_0 ^= 1; asm volatile("tcgen05.fence::after_thread_sync;"); - { - int tmem_dealloc_peer_rank = bid % 2 ^ 1; - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(tmem_dealloc_peer_addr), "r"(tmem_dealloc_peer_rank) : "memory"); - mbarrier_wait(tmem_dealloc_peer_addr, 0); - int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); - asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); - } + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); } } // ---- Role: empty1 ---- diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_7a4b9556b3c84ebff6ce_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_c1d444bba4cdb8273c7c_binding.cu similarity index 97% rename from csrc/cake_dsv4/sm_103a/cake_dsv4_7a4b9556b3c84ebff6ce_binding.cu rename to csrc/cake_dsv4/sm_103a/cake_dsv4_c1d444bba4cdb8273c7c_binding.cu index b7c794c30a4..b03ef9f4476 100644 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_7a4b9556b3c84ebff6ce_binding.cu +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_c1d444bba4cdb8273c7c_binding.cu @@ -15,7 +15,7 @@ */ // Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_7a4b9556b3c84ebff6ce'. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_c1d444bba4cdb8273c7c'. #include #include #include @@ -32,10 +32,10 @@ #include struct CakeTensorMap; -extern "C" __global__ void kernel_cake_dsv4_7a4b9556b3c84ebff6ce(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int max_q_len, int batch_size, int has_sinks); +extern "C" __global__ void kernel_cake_dsv4_c1d444bba4cdb8273c7c(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int max_q_len, int batch_size, int has_sinks); -namespace cake_host_shim_1779b169b6ce3d43 { +namespace cake_host_shim_1a24ae3a58f0bd58 { using tvm::ffi::Optional; using tvm::ffi::TensorView; @@ -492,7 +492,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ } if (smem_status == cudaSuccess) { smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_7a4b9556b3c84ebff6ce), + reinterpret_cast(kernel_cake_dsv4_c1d444bba4cdb8273c7c), cudaFuncAttributeMaxDynamicSharedMemorySize, 147072); } @@ -502,7 +502,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ } } TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_7a4b9556b3c84ebff6ce failed: " + << "cudaFuncSetAttribute for kernel_cake_dsv4_c1d444bba4cdb8273c7c failed: " << cudaGetErrorString(smem_status); dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); dim3 block(512u, 1u, 1u); @@ -515,12 +515,12 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ config.dynamicSmemBytes = 147072u; config.stream = stream; cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_7a4b9556b3c84ebff6ce), kargs); + &config, reinterpret_cast(kernel_cake_dsv4_c1d444bba4cdb8273c7c), kargs); TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_7a4b9556b3c84ebff6ce failed: " + << "cudaLaunchKernelExC for kernel_cake_dsv4_c1d444bba4cdb8273c7c failed: " << cudaGetErrorString(launch_status); } -} // namespace cake_host_shim_1779b169b6ce3d43 +} // namespace cake_host_shim_1a24ae3a58f0bd58 -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_1779b169b6ce3d43::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_1a24ae3a58f0bd58::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_7a4b9556b3c84ebff6ce_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_c1d444bba4cdb8273c7c_kernel.cu similarity index 95% rename from csrc/cake_dsv4/sm_103a/cake_dsv4_7a4b9556b3c84ebff6ce_kernel.cu rename to csrc/cake_dsv4/sm_103a/cake_dsv4_c1d444bba4cdb8273c7c_kernel.cu index ca25505d652..f2cfc0f9c4f 100644 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_7a4b9556b3c84ebff6ce_kernel.cu +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_c1d444bba4cdb8273c7c_kernel.cu @@ -1000,7 +1000,7 @@ __device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { extern "C" { __global__ __launch_bounds__(512, 1) void -kernel_cake_dsv4_7a4b9556b3c84ebff6ce(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int max_q_len, int batch_size, int has_sinks) +kernel_cake_dsv4_c1d444bba4cdb8273c7c(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int max_q_len, int batch_size, int has_sinks) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); @@ -1214,8 +1214,8 @@ kernel_cake_dsv4_7a4b9556b3c84ebff6ce(CakeTensorMap const* tmap_q, CakeTensorMap int active_topk = 0; if (work_x < (unsigned int)q_len) { int _max_4 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); - int _min_3 = ((_max_4) < (sparse_topk) ? (_max_4) : (sparse_topk)); - active_topk = _min_3; + int _min_7 = ((_max_4) < (sparse_topk) ? (_max_4) : (sparse_topk)); + active_topk = _min_7; } int num_steps = (active_topk + 128 - 1) / 128; float row_max_pair[2]; @@ -1542,8 +1542,8 @@ kernel_cake_dsv4_7a4b9556b3c84ebff6ce(CakeTensorMap const* tmap_q, CakeTensorMap int active_topk_1 = 0; if (work_x_1 < (unsigned int)q_len_1) { int _max_14 = ((sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) : (0)); - int _min_4 = ((_max_14) < (sparse_topk) ? (_max_14) : (sparse_topk)); - active_topk_1 = _min_4; + int _min_8 = ((_max_14) < (sparse_topk) ? (_max_14) : (sparse_topk)); + active_topk_1 = _min_8; } int num_steps_1 = (active_topk_1 + 128 - 1) / 128; int head_group = work_y_1 / 4; @@ -1806,8 +1806,8 @@ kernel_cake_dsv4_7a4b9556b3c84ebff6ce(CakeTensorMap const* tmap_q, CakeTensorMap int active_topk_2 = 0; if (work_x_2 < (unsigned int)q_len_2) { int _max_3 = ((sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) : (0)); - int _min_2 = ((_max_3) < (sparse_topk) ? (_max_3) : (sparse_topk)); - active_topk_2 = _min_2; + int _min_6 = ((_max_3) < (sparse_topk) ? (_max_3) : (sparse_topk)); + active_topk_2 = _min_6; } int num_steps_2 = (active_topk_2 + 128 - 1) / 128; int value_quarter_1 = work_y_2 % 4; @@ -2443,30 +2443,45 @@ kernel_cake_dsv4_7a4b9556b3c84ebff6ce(CakeTensorMap const* tmap_q, CakeTensorMap } mbarrier_wait(page_full_addr + (page_k_stage) * 8, page_k_phase); int page_base = 0; + int warps_tile_idx = 0; + int _min_2 = ((active_topk_4) < (128) ? (active_topk_4) : (128)); + int warps_valid_rows = _min_2; + int _min_3 = (((warps_valid_rows + 3) / 4) < (4) ? ((warps_valid_rows + 3) / 4) : (4)); + int warps_tile = _min_3; + int warps_ge1 = 4; + int warps_ge2 = 4; + if (warps_tile_idx >= 1) { + warps_ge1 = warps_tile; + } + if (warps_tile_idx >= 2) { + warps_ge2 = warps_tile; + } #pragma unroll for (int hs_3 = 0; hs_3 < 4; hs_3++) { int hs_sel_2 = ((hs_3 == 3) ? value_quarter_2 : ((value_quarter_2 > hs_3) ? hs_3 : hs_3 + 1)); mbarrier_wait(kv_empty_addr + (kv_stage_1) * 8, kv_phase_1); if (load_rank == 0) { if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_1) * 8, 16384); + mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_1) * 8, ((hs_3 >= 4) ? warps_ge1 : warps_ge2) * 4096); } } asm volatile("barrier.sync 9, 128;" ::: "memory"); - int rows_all[32]; - #pragma unroll - for (int group = 0; group < 8; group++) { - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&rows_all[group * 4])), "=r"(*reinterpret_cast(&rows_all[(group * 4) + 1])), "=r"(*reinterpret_cast(&rows_all[(group * 4) + 2])), "=r"(*reinterpret_cast(&rows_all[(group * 4) + 3])) - : "r"(smem_page_offsets_addr + page_k_stage * 1024 + (unsigned int)((page_base + load_rank * 4 + group * 16) * 4))); - } - if (elect_sync()) { - { - #pragma unroll - for (int group_1 = 0; group_1 < 8; group_1++) { - const int row_base = load_rank * 4 + group_1 * 16; - int dst_k = smem_kv_addr + kv_stage_1 * 16384 + (unsigned int)(row_base * 128); - tma_gather4_gmem2smem(dst_k, tmap_swa_kv, hs_sel_2 * 128, rows_all[group_1 * 4], rows_all[group_1 * 4 + 1], rows_all[group_1 * 4 + 2], rows_all[group_1 * 4 + 3], kv_full_addr + (kv_stage_1) * 8); + if (load_rank < ((hs_3 >= 4) ? warps_ge1 : warps_ge2)) { + int rows_all[32]; + #pragma unroll + for (int group = 0; group < 8; group++) { + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&rows_all[group * 4])), "=r"(*reinterpret_cast(&rows_all[(group * 4) + 1])), "=r"(*reinterpret_cast(&rows_all[(group * 4) + 2])), "=r"(*reinterpret_cast(&rows_all[(group * 4) + 3])) + : "r"(smem_page_offsets_addr + page_k_stage * 1024 + (unsigned int)((page_base + load_rank * 4 + group * 16) * 4))); + } + if (elect_sync()) { + { + #pragma unroll + for (int group_1 = 0; group_1 < 8; group_1++) { + const int row_base = load_rank * 4 + group_1 * 16; + int dst_k = smem_kv_addr + kv_stage_1 * 16384 + (unsigned int)(row_base * 128); + tma_gather4_gmem2smem(dst_k, tmap_swa_kv, hs_sel_2 * 128, rows_all[group_1 * 4], rows_all[group_1 * 4 + 1], rows_all[group_1 * 4 + 2], rows_all[group_1 * 4 + 3], kv_full_addr + (kv_stage_1) * 8); + } } } } @@ -2480,37 +2495,52 @@ kernel_cake_dsv4_7a4b9556b3c84ebff6ce(CakeTensorMap const* tmap_q, CakeTensorMap mbarrier_wait(page_full_addr + (page_k_stage) * 8, page_k_phase); } int page_base_0 = (next_tile & 1) * 128; + int warps_tile_idx_1 = next_tile; + int _min_4 = ((active_topk_4 - next_tile * 128) < (128) ? (active_topk_4 - next_tile * 128) : (128)); + int warps_valid_rows_2 = _min_4; + int _min_5 = (((warps_valid_rows_2 + 3) / 4) < (4) ? ((warps_valid_rows_2 + 3) / 4) : (4)); + int warps_tile_3 = _min_5; + int warps_ge1_4 = 4; + int warps_ge2_5 = 4; + if (warps_tile_idx_1 >= 1) { + warps_ge1_4 = warps_tile_3; + } + if (warps_tile_idx_1 >= 2) { + warps_ge2_5 = warps_tile_3; + } #pragma unroll for (int hs_4 = 0; hs_4 < 4; hs_4++) { int hs_sel_3 = ((hs_4 == 3) ? value_quarter_2 : ((value_quarter_2 > hs_4) ? hs_4 : hs_4 + 1)); mbarrier_wait(kv_empty_addr + (kv_stage_1) * 8, kv_phase_1); if (load_rank == 0) { if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_1) * 8, 16384); + mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_1) * 8, ((hs_4 >= 4) ? warps_ge1_4 : warps_ge2_5) * 4096); } } asm volatile("barrier.sync 9, 128;" ::: "memory"); - int rows_all_1[32]; - #pragma unroll - for (int group_2 = 0; group_2 < 8; group_2++) { - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&rows_all_1[group_2 * 4])), "=r"(*reinterpret_cast(&rows_all_1[(group_2 * 4) + 1])), "=r"(*reinterpret_cast(&rows_all_1[(group_2 * 4) + 2])), "=r"(*reinterpret_cast(&rows_all_1[(group_2 * 4) + 3])) - : "r"(smem_page_offsets_addr + page_k_stage * 1024 + (unsigned int)((page_base_0 + load_rank * 4 + group_2 * 16) * 4))); - } - if (elect_sync()) { - if (next_tile == 0) { - #pragma unroll - for (int group_3 = 0; group_3 < 8; group_3++) { - const int row_base_1 = load_rank * 4 + group_3 * 16; - int dst_k_1 = smem_kv_addr + kv_stage_1 * 16384 + (unsigned int)(row_base_1 * 128); - tma_gather4_gmem2smem(dst_k_1, tmap_swa_kv, hs_sel_3 * 128, rows_all_1[group_3 * 4], rows_all_1[group_3 * 4 + 1], rows_all_1[group_3 * 4 + 2], rows_all_1[group_3 * 4 + 3], kv_full_addr + (kv_stage_1) * 8); - } - } else { - #pragma unroll - for (int group_4 = 0; group_4 < 8; group_4++) { - const int row_base_2 = load_rank * 4 + group_4 * 16; - int dst_k_2 = smem_kv_addr + kv_stage_1 * 16384 + (unsigned int)(row_base_2 * 128); - tma_gather4_gmem2smem(dst_k_2, tmap_compressed_kv, hs_sel_3 * 128, rows_all_1[group_4 * 4], rows_all_1[group_4 * 4 + 1], rows_all_1[group_4 * 4 + 2], rows_all_1[group_4 * 4 + 3], kv_full_addr + (kv_stage_1) * 8); + if (load_rank < ((hs_4 >= 4) ? warps_ge1_4 : warps_ge2_5)) { + int rows_all_1[32]; + #pragma unroll + for (int group_2 = 0; group_2 < 8; group_2++) { + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&rows_all_1[group_2 * 4])), "=r"(*reinterpret_cast(&rows_all_1[(group_2 * 4) + 1])), "=r"(*reinterpret_cast(&rows_all_1[(group_2 * 4) + 2])), "=r"(*reinterpret_cast(&rows_all_1[(group_2 * 4) + 3])) + : "r"(smem_page_offsets_addr + page_k_stage * 1024 + (unsigned int)((page_base_0 + load_rank * 4 + group_2 * 16) * 4))); + } + if (elect_sync()) { + if (next_tile == 0) { + #pragma unroll + for (int group_3 = 0; group_3 < 8; group_3++) { + const int row_base_1 = load_rank * 4 + group_3 * 16; + int dst_k_1 = smem_kv_addr + kv_stage_1 * 16384 + (unsigned int)(row_base_1 * 128); + tma_gather4_gmem2smem(dst_k_1, tmap_swa_kv, hs_sel_3 * 128, rows_all_1[group_3 * 4], rows_all_1[group_3 * 4 + 1], rows_all_1[group_3 * 4 + 2], rows_all_1[group_3 * 4 + 3], kv_full_addr + (kv_stage_1) * 8); + } + } else { + #pragma unroll + for (int group_4 = 0; group_4 < 8; group_4++) { + const int row_base_2 = load_rank * 4 + group_4 * 16; + int dst_k_2 = smem_kv_addr + kv_stage_1 * 16384 + (unsigned int)(row_base_2 * 128); + tma_gather4_gmem2smem(dst_k_2, tmap_compressed_kv, hs_sel_3 * 128, rows_all_1[group_4 * 4], rows_all_1[group_4 * 4 + 1], rows_all_1[group_4 * 4 + 2], rows_all_1[group_4 * 4 + 3], kv_full_addr + (kv_stage_1) * 8); + } } } } diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_4ac65f00a4a032c10ecb_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_d9754db4fe394a927e67_binding.cu similarity index 96% rename from csrc/cake_dsv4/sm_100a/cake_dsv4_4ac65f00a4a032c10ecb_binding.cu rename to csrc/cake_dsv4/sm_103a/cake_dsv4_d9754db4fe394a927e67_binding.cu index 245c914ddcf..171a6c68149 100644 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_4ac65f00a4a032c10ecb_binding.cu +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_d9754db4fe394a927e67_binding.cu @@ -15,7 +15,7 @@ */ // Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_4ac65f00a4a032c10ecb'. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_d9754db4fe394a927e67'. #include #include #include @@ -31,10 +31,10 @@ #include #include -extern "C" __global__ void kernel_cake_dsv4_4ac65f00a4a032c10ecb(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks); +extern "C" __global__ void kernel_cake_dsv4_d9754db4fe394a927e67(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks); -namespace cake_host_shim_71507c1f0a7f757c { +namespace cake_host_shim_01a51b3a9acbf4d0 { using tvm::ffi::Optional; using tvm::ffi::TensorView; @@ -353,7 +353,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_O, Te } if (smem_status == cudaSuccess) { smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_4ac65f00a4a032c10ecb), + reinterpret_cast(kernel_cake_dsv4_d9754db4fe394a927e67), cudaFuncAttributeMaxDynamicSharedMemorySize, 198528); } @@ -363,7 +363,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_O, Te } } TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_4ac65f00a4a032c10ecb failed: " + << "cudaFuncSetAttribute for kernel_cake_dsv4_d9754db4fe394a927e67 failed: " << cudaGetErrorString(smem_status); dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); dim3 block(512u, 1u, 1u); @@ -374,12 +374,12 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_O, Te config.dynamicSmemBytes = 198528u; config.stream = stream; cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_4ac65f00a4a032c10ecb), kargs); + &config, reinterpret_cast(kernel_cake_dsv4_d9754db4fe394a927e67), kargs); TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_4ac65f00a4a032c10ecb failed: " + << "cudaLaunchKernelExC for kernel_cake_dsv4_d9754db4fe394a927e67 failed: " << cudaGetErrorString(launch_status); } -} // namespace cake_host_shim_71507c1f0a7f757c +} // namespace cake_host_shim_01a51b3a9acbf4d0 -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_71507c1f0a7f757c::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_01a51b3a9acbf4d0::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_9cf30aacb43b71034660_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_d9754db4fe394a927e67_kernel.cu similarity index 89% rename from csrc/cake_dsv4/sm_103a/cake_dsv4_9cf30aacb43b71034660_kernel.cu rename to csrc/cake_dsv4/sm_103a/cake_dsv4_d9754db4fe394a927e67_kernel.cu index 3398f54dc49..5de2a0df80a 100644 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_9cf30aacb43b71034660_kernel.cu +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_d9754db4fe394a927e67_kernel.cu @@ -999,7 +999,7 @@ __device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { extern "C" { __global__ __launch_bounds__(512, 1) void -kernel_cake_dsv4_9cf30aacb43b71034660(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks) +kernel_cake_dsv4_d9754db4fe394a927e67(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); @@ -1363,29 +1363,129 @@ kernel_cake_dsv4_9cf30aacb43b71034660(const __grid_constant__ CUtensorMap tmap_q " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" : "=r"(*reinterpret_cast(&_tmem_load_1[32])), "=r"(*reinterpret_cast(&_tmem_load_1[33])), "=r"(*reinterpret_cast(&_tmem_load_1[34])), "=r"(*reinterpret_cast(&_tmem_load_1[35])), "=r"(*reinterpret_cast(&_tmem_load_1[36])), "=r"(*reinterpret_cast(&_tmem_load_1[37])), "=r"(*reinterpret_cast(&_tmem_load_1[38])), "=r"(*reinterpret_cast(&_tmem_load_1[39])), "=r"(*reinterpret_cast(&_tmem_load_1[40])), "=r"(*reinterpret_cast(&_tmem_load_1[41])), "=r"(*reinterpret_cast(&_tmem_load_1[42])), "=r"(*reinterpret_cast(&_tmem_load_1[43])), "=r"(*reinterpret_cast(&_tmem_load_1[44])), "=r"(*reinterpret_cast(&_tmem_load_1[45])), "=r"(*reinterpret_cast(&_tmem_load_1[46])), "=r"(*reinterpret_cast(&_tmem_load_1[47])), "=r"(*reinterpret_cast(&_tmem_load_1[48])), "=r"(*reinterpret_cast(&_tmem_load_1[49])), "=r"(*reinterpret_cast(&_tmem_load_1[50])), "=r"(*reinterpret_cast(&_tmem_load_1[51])), "=r"(*reinterpret_cast(&_tmem_load_1[52])), "=r"(*reinterpret_cast(&_tmem_load_1[53])), "=r"(*reinterpret_cast(&_tmem_load_1[54])), "=r"(*reinterpret_cast(&_tmem_load_1[55])), "=r"(*reinterpret_cast(&_tmem_load_1[56])), "=r"(*reinterpret_cast(&_tmem_load_1[57])), "=r"(*reinterpret_cast(&_tmem_load_1[58])), "=r"(*reinterpret_cast(&_tmem_load_1[59])), "=r"(*reinterpret_cast(&_tmem_load_1[60])), "=r"(*reinterpret_cast(&_tmem_load_1[61])), "=r"(*reinterpret_cast(&_tmem_load_1[62])), "=r"(*reinterpret_cast(&_tmem_load_1[63])) : "r"(taddr + 256 + (unsigned int)row_addr + 32)); - int col_base = col_half_1 * 64; - if (my_row_1 < num_heads) { - #pragma unroll - for (int offset = 0; offset < 64; offset += 8) { - { - const float2 _prescale2_0 = {inv_sum * output_scale, inv_sum * output_scale}; - #if __CUDA_ARCH__ >= 1000 - #pragma unroll - for (int _ps = 0; _ps < 4; _ps++) - mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_1[offset])[_ps], _prescale2_0); - #else - #pragma unroll - for (int _ps = 0; _ps < 8; _ps++) - _tmem_load_1[offset + _ps] *= inv_sum * output_scale; - #endif - __nv_bfloat162 _pk[4]; - _pk[0] = __floats2bfloat162_rn(_tmem_load_1[offset + 0], _tmem_load_1[offset + 1]); - _pk[1] = __floats2bfloat162_rn(_tmem_load_1[offset + 2], _tmem_load_1[offset + 3]); - _pk[2] = __floats2bfloat162_rn(_tmem_load_1[offset + 4], _tmem_load_1[offset + 5]); - _pk[3] = __floats2bfloat162_rn(_tmem_load_1[offset + 6], _tmem_load_1[offset + 7]); - *reinterpret_cast(&((__nv_bfloat16*)(O + (output_base + col_base + offset)))[0]) = *reinterpret_cast(&_pk[0]); - } - } + float2 _f2_14 = make_float2(inv_sum * output_scale, inv_sum * output_scale); + #pragma unroll + for (int i_1 = 0; i_1 < 64; i_1 += 2) { + float2 _f2_15 = make_float2(_tmem_load_1[i_1], _tmem_load_1[i_1 + 1]); + float2 _mul_f32x2_0; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_0) : "l"(*(const unsigned long long*)&_f2_15), "l"(*(const unsigned long long*)&_f2_14)); + _tmem_load_1[i_1] = _mul_f32x2_0.x; + _tmem_load_1[i_1 + 1] = _mul_f32x2_0.y; + } + unsigned int packed_o[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_1[_lp*2 + 0], _tmem_load_1[_lp*2+1 + 0])); + packed_o[_lp] = *(uint32_t*)&_bf2; + } + bool stage_hi_1 = (lane & 1) != 0; + bool stage_hi_2 = (lane & 2) != 0; + bool stage_hi_4 = (lane & 4) != 0; + unsigned int _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[0] : packed_o[4]), 1); + packed_o[0] = ((stage_hi_1) ? _shfl_xor_2 : packed_o[0]); + packed_o[4] = ((stage_hi_1) ? packed_o[4] : _shfl_xor_2); + unsigned int _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[1] : packed_o[5]), 1); + packed_o[1] = ((stage_hi_1) ? _shfl_xor_3 : packed_o[1]); + packed_o[5] = ((stage_hi_1) ? packed_o[5] : _shfl_xor_3); + unsigned int _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[2] : packed_o[6]), 1); + packed_o[2] = ((stage_hi_1) ? _shfl_xor_4 : packed_o[2]); + packed_o[6] = ((stage_hi_1) ? packed_o[6] : _shfl_xor_4); + unsigned int _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[3] : packed_o[7]), 1); + packed_o[3] = ((stage_hi_1) ? _shfl_xor_5 : packed_o[3]); + packed_o[7] = ((stage_hi_1) ? packed_o[7] : _shfl_xor_5); + unsigned int _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[8] : packed_o[12]), 1); + packed_o[8] = ((stage_hi_1) ? _shfl_xor_6 : packed_o[8]); + packed_o[12] = ((stage_hi_1) ? packed_o[12] : _shfl_xor_6); + unsigned int _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[9] : packed_o[13]), 1); + packed_o[9] = ((stage_hi_1) ? _shfl_xor_7 : packed_o[9]); + packed_o[13] = ((stage_hi_1) ? packed_o[13] : _shfl_xor_7); + unsigned int _shfl_xor_8 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[10] : packed_o[14]), 1); + packed_o[10] = ((stage_hi_1) ? _shfl_xor_8 : packed_o[10]); + packed_o[14] = ((stage_hi_1) ? packed_o[14] : _shfl_xor_8); + unsigned int _shfl_xor_9 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[11] : packed_o[15]), 1); + packed_o[11] = ((stage_hi_1) ? _shfl_xor_9 : packed_o[11]); + packed_o[15] = ((stage_hi_1) ? packed_o[15] : _shfl_xor_9); + unsigned int _shfl_xor_10 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[16] : packed_o[20]), 1); + packed_o[16] = ((stage_hi_1) ? _shfl_xor_10 : packed_o[16]); + packed_o[20] = ((stage_hi_1) ? packed_o[20] : _shfl_xor_10); + unsigned int _shfl_xor_11 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[17] : packed_o[21]), 1); + packed_o[17] = ((stage_hi_1) ? _shfl_xor_11 : packed_o[17]); + packed_o[21] = ((stage_hi_1) ? packed_o[21] : _shfl_xor_11); + unsigned int _shfl_xor_12 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[18] : packed_o[22]), 1); + packed_o[18] = ((stage_hi_1) ? _shfl_xor_12 : packed_o[18]); + packed_o[22] = ((stage_hi_1) ? packed_o[22] : _shfl_xor_12); + unsigned int _shfl_xor_13 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[19] : packed_o[23]), 1); + packed_o[19] = ((stage_hi_1) ? _shfl_xor_13 : packed_o[19]); + packed_o[23] = ((stage_hi_1) ? packed_o[23] : _shfl_xor_13); + unsigned int _shfl_xor_14 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[24] : packed_o[28]), 1); + packed_o[24] = ((stage_hi_1) ? _shfl_xor_14 : packed_o[24]); + packed_o[28] = ((stage_hi_1) ? packed_o[28] : _shfl_xor_14); + unsigned int _shfl_xor_15 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[25] : packed_o[29]), 1); + packed_o[25] = ((stage_hi_1) ? _shfl_xor_15 : packed_o[25]); + packed_o[29] = ((stage_hi_1) ? packed_o[29] : _shfl_xor_15); + unsigned int _shfl_xor_16 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[26] : packed_o[30]), 1); + packed_o[26] = ((stage_hi_1) ? _shfl_xor_16 : packed_o[26]); + packed_o[30] = ((stage_hi_1) ? packed_o[30] : _shfl_xor_16); + unsigned int _shfl_xor_17 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[27] : packed_o[31]), 1); + packed_o[27] = ((stage_hi_1) ? _shfl_xor_17 : packed_o[27]); + packed_o[31] = ((stage_hi_1) ? packed_o[31] : _shfl_xor_17); + int lane_t = lane & 1; + unsigned int chunk_words[4]; + chunk_words[0] = packed_o[0]; + chunk_words[1] = packed_o[1]; + chunk_words[2] = packed_o[2]; + chunk_words[3] = packed_o[3]; + if ((my_row_1 & -2 | 0) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + (my_row_1 & -2 | 0)) * 512 + v_chunk * 128 + col_half_1 * 64 + (0 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[4]; + chunk_words[1] = packed_o[5]; + chunk_words[2] = packed_o[6]; + chunk_words[3] = packed_o[7]; + if ((my_row_1 & -2 | 1) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + (my_row_1 & -2 | 1)) * 512 + v_chunk * 128 + col_half_1 * 64 + (0 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[8]; + chunk_words[1] = packed_o[9]; + chunk_words[2] = packed_o[10]; + chunk_words[3] = packed_o[11]; + if ((my_row_1 & -2 | 0) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + (my_row_1 & -2 | 0)) * 512 + v_chunk * 128 + col_half_1 * 64 + (2 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[12]; + chunk_words[1] = packed_o[13]; + chunk_words[2] = packed_o[14]; + chunk_words[3] = packed_o[15]; + if ((my_row_1 & -2 | 1) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + (my_row_1 & -2 | 1)) * 512 + v_chunk * 128 + col_half_1 * 64 + (2 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[16]; + chunk_words[1] = packed_o[17]; + chunk_words[2] = packed_o[18]; + chunk_words[3] = packed_o[19]; + if ((my_row_1 & -2 | 0) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + (my_row_1 & -2 | 0)) * 512 + v_chunk * 128 + col_half_1 * 64 + (4 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[20]; + chunk_words[1] = packed_o[21]; + chunk_words[2] = packed_o[22]; + chunk_words[3] = packed_o[23]; + if ((my_row_1 & -2 | 1) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + (my_row_1 & -2 | 1)) * 512 + v_chunk * 128 + col_half_1 * 64 + (4 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[24]; + chunk_words[1] = packed_o[25]; + chunk_words[2] = packed_o[26]; + chunk_words[3] = packed_o[27]; + if ((my_row_1 & -2 | 0) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + (my_row_1 & -2 | 0)) * 512 + v_chunk * 128 + col_half_1 * 64 + (6 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[28]; + chunk_words[1] = packed_o[29]; + chunk_words[2] = packed_o[30]; + chunk_words[3] = packed_o[31]; + if ((my_row_1 & -2 | 1) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + (my_row_1 & -2 | 1)) * 512 + v_chunk * 128 + col_half_1 * 64 + (6 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; } asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); asm volatile("tcgen05.fence::before_thread_sync;"); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_d23652aae3370c30d978_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_e48c5e0a89a41eb1e57a_binding.cu similarity index 97% rename from csrc/cake_dsv4/sm_103a/cake_dsv4_d23652aae3370c30d978_binding.cu rename to csrc/cake_dsv4/sm_103a/cake_dsv4_e48c5e0a89a41eb1e57a_binding.cu index 51c14700a2a..0b839f4f0ba 100644 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_d23652aae3370c30d978_binding.cu +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_e48c5e0a89a41eb1e57a_binding.cu @@ -15,7 +15,7 @@ */ // Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_d23652aae3370c30d978'. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_e48c5e0a89a41eb1e57a'. #include #include #include @@ -32,10 +32,10 @@ #include struct CakeTensorMap; -extern "C" __global__ void kernel_cake_dsv4_d23652aae3370c30d978(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int max_q_len, int batch_size, int has_sinks); +extern "C" __global__ void kernel_cake_dsv4_e48c5e0a89a41eb1e57a(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int max_q_len, int batch_size, int has_sinks); -namespace cake_host_shim_9b1b8d760815189c { +namespace cake_host_shim_b67714d8806118f0 { using tvm::ffi::Optional; using tvm::ffi::TensorView; @@ -492,7 +492,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ } if (smem_status == cudaSuccess) { smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_d23652aae3370c30d978), + reinterpret_cast(kernel_cake_dsv4_e48c5e0a89a41eb1e57a), cudaFuncAttributeMaxDynamicSharedMemorySize, 218752); } @@ -502,7 +502,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ } } TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_d23652aae3370c30d978 failed: " + << "cudaFuncSetAttribute for kernel_cake_dsv4_e48c5e0a89a41eb1e57a failed: " << cudaGetErrorString(smem_status); dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); dim3 block(640u, 1u, 1u); @@ -515,12 +515,12 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_ config.dynamicSmemBytes = 218752u; config.stream = stream; cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_d23652aae3370c30d978), kargs); + &config, reinterpret_cast(kernel_cake_dsv4_e48c5e0a89a41eb1e57a), kargs); TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_d23652aae3370c30d978 failed: " + << "cudaLaunchKernelExC for kernel_cake_dsv4_e48c5e0a89a41eb1e57a failed: " << cudaGetErrorString(launch_status); } -} // namespace cake_host_shim_9b1b8d760815189c +} // namespace cake_host_shim_b67714d8806118f0 -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_9b1b8d760815189c::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_b67714d8806118f0::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_d23652aae3370c30d978_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_e48c5e0a89a41eb1e57a_kernel.cu similarity index 96% rename from csrc/cake_dsv4/sm_103a/cake_dsv4_d23652aae3370c30d978_kernel.cu rename to csrc/cake_dsv4/sm_103a/cake_dsv4_e48c5e0a89a41eb1e57a_kernel.cu index f3f50174875..f841e132e77 100644 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_d23652aae3370c30d978_kernel.cu +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_e48c5e0a89a41eb1e57a_kernel.cu @@ -1000,7 +1000,7 @@ __device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { extern "C" { __global__ __launch_bounds__(640, 1) void -kernel_cake_dsv4_d23652aae3370c30d978(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int max_q_len, int batch_size, int has_sinks) +kernel_cake_dsv4_e48c5e0a89a41eb1e57a(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int max_q_len, int batch_size, int has_sinks) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); @@ -1210,8 +1210,8 @@ kernel_cake_dsv4_d23652aae3370c30d978(CakeTensorMap const* tmap_q, CakeTensorMap int active_topk = 0; if (work_x < (unsigned int)q_len) { int _max_4 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); - int _min_3 = ((_max_4) < (sparse_topk) ? (_max_4) : (sparse_topk)); - active_topk = _min_3; + int _min_7 = ((_max_4) < (sparse_topk) ? (_max_4) : (sparse_topk)); + active_topk = _min_7; } int num_steps = (active_topk + 128 - 1) / 128; float row_max_pair[2]; @@ -1509,8 +1509,8 @@ kernel_cake_dsv4_d23652aae3370c30d978(CakeTensorMap const* tmap_q, CakeTensorMap int active_topk_1 = 0; if (work_x_1 < (unsigned int)q_len_1) { int _max_14 = ((sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) : (0)); - int _min_4 = ((_max_14) < (sparse_topk) ? (_max_14) : (sparse_topk)); - active_topk_1 = _min_4; + int _min_8 = ((_max_14) < (sparse_topk) ? (_max_14) : (sparse_topk)); + active_topk_1 = _min_8; } int num_steps_1 = (active_topk_1 + 128 - 1) / 128; int head_group = work_y_1 / 4; @@ -1773,8 +1773,8 @@ kernel_cake_dsv4_d23652aae3370c30d978(CakeTensorMap const* tmap_q, CakeTensorMap int active_topk_2 = 0; if (work_x_2 < (unsigned int)q_len_2) { int _max_3 = ((sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) : (0)); - int _min_2 = ((_max_3) < (sparse_topk) ? (_max_3) : (sparse_topk)); - active_topk_2 = _min_2; + int _min_6 = ((_max_3) < (sparse_topk) ? (_max_3) : (sparse_topk)); + active_topk_2 = _min_6; } int num_steps_2 = (active_topk_2 + 128 - 1) / 128; int value_quarter_1 = work_y_2 % 4; @@ -2530,28 +2530,43 @@ kernel_cake_dsv4_d23652aae3370c30d978(CakeTensorMap const* tmap_q, CakeTensorMap } mbarrier_wait(page_full_addr + (page_k_stage) * 8, page_k_phase); int page_base = 0; + int warps_tile_idx = 0; + int _min_2 = ((active_topk_4) < (128) ? (active_topk_4) : (128)); + int warps_valid_rows = _min_2; + int _min_3 = (((warps_valid_rows + 3) / 4) < (8) ? ((warps_valid_rows + 3) / 4) : (8)); + int warps_tile = _min_3; + int warps_ge1 = 8; + int warps_ge2 = 8; + if (warps_tile_idx >= 1) { + warps_ge1 = warps_tile; + } + if (warps_tile_idx >= 2) { + warps_ge2 = warps_tile; + } #pragma unroll for (int hs_3 = 0; hs_3 < 4; hs_3++) { int hs_sel_2 = ((hs_3 == 3) ? value_quarter_2 : ((value_quarter_2 > hs_3) ? hs_3 : hs_3 + 1)); mbarrier_wait(kv_empty_addr + (kv_stage_1) * 8, kv_phase_1); if (load_rank == 0) { if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_1) * 8, 32768); + mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_1) * 8, ((hs_3 >= 2) ? warps_ge1 : warps_ge2) * 4096); } } asm volatile("barrier.sync 9, 256;" ::: "memory"); - #pragma unroll - for (int group = 0; group < 4; group++) { - const int row_base = load_rank * 4 + group * 32; - int rows[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&rows[0])), "=r"(*reinterpret_cast(&rows[(0) + 1])), "=r"(*reinterpret_cast(&rows[(0) + 2])), "=r"(*reinterpret_cast(&rows[(0) + 3])) - : "r"(smem_page_offsets_addr + page_k_stage * 1024 + (unsigned int)((page_base + row_base) * 4))); - int dst_k = smem_kv_addr + kv_stage_1 * 32768 + (unsigned int)(row_base * 64 * 2); - if (elect_sync()) { - { - tma_gather4_gmem2smem(dst_k, tmap_swa_kv, hs_sel_2 * 128, rows[0], rows[1], rows[2], rows[3], kv_full_addr + (kv_stage_1) * 8); - tma_gather4_gmem2smem(dst_k + 16384, tmap_swa_kv, hs_sel_2 * 128 + 64, rows[0], rows[1], rows[2], rows[3], kv_full_addr + (kv_stage_1) * 8); + if (load_rank < ((hs_3 >= 2) ? warps_ge1 : warps_ge2)) { + #pragma unroll + for (int group = 0; group < 4; group++) { + const int row_base = load_rank * 4 + group * 32; + int rows[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&rows[0])), "=r"(*reinterpret_cast(&rows[(0) + 1])), "=r"(*reinterpret_cast(&rows[(0) + 2])), "=r"(*reinterpret_cast(&rows[(0) + 3])) + : "r"(smem_page_offsets_addr + page_k_stage * 1024 + (unsigned int)((page_base + row_base) * 4))); + int dst_k = smem_kv_addr + kv_stage_1 * 32768 + (unsigned int)(row_base * 64 * 2); + if (elect_sync()) { + { + tma_gather4_gmem2smem(dst_k, tmap_swa_kv, hs_sel_2 * 128, rows[0], rows[1], rows[2], rows[3], kv_full_addr + (kv_stage_1) * 8); + tma_gather4_gmem2smem(dst_k + 16384, tmap_swa_kv, hs_sel_2 * 128 + 64, rows[0], rows[1], rows[2], rows[3], kv_full_addr + (kv_stage_1) * 8); + } } } } @@ -2565,31 +2580,46 @@ kernel_cake_dsv4_d23652aae3370c30d978(CakeTensorMap const* tmap_q, CakeTensorMap mbarrier_wait(page_full_addr + (page_k_stage) * 8, page_k_phase); } int page_base_0 = (next_tile & 1) * 128; + int warps_tile_idx_1 = next_tile; + int _min_4 = ((active_topk_4 - next_tile * 128) < (128) ? (active_topk_4 - next_tile * 128) : (128)); + int warps_valid_rows_2 = _min_4; + int _min_5 = (((warps_valid_rows_2 + 3) / 4) < (8) ? ((warps_valid_rows_2 + 3) / 4) : (8)); + int warps_tile_3 = _min_5; + int warps_ge1_4 = 8; + int warps_ge2_5 = 8; + if (warps_tile_idx_1 >= 1) { + warps_ge1_4 = warps_tile_3; + } + if (warps_tile_idx_1 >= 2) { + warps_ge2_5 = warps_tile_3; + } #pragma unroll for (int hs_4 = 0; hs_4 < 4; hs_4++) { int hs_sel_3 = ((hs_4 == 3) ? value_quarter_2 : ((value_quarter_2 > hs_4) ? hs_4 : hs_4 + 1)); mbarrier_wait(kv_empty_addr + (kv_stage_1) * 8, kv_phase_1); if (load_rank == 0) { if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_1) * 8, 32768); + mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_1) * 8, ((hs_4 >= 2) ? warps_ge1_4 : warps_ge2_5) * 4096); } } asm volatile("barrier.sync 9, 256;" ::: "memory"); - #pragma unroll - for (int group_1 = 0; group_1 < 4; group_1++) { - const int row_base_1 = load_rank * 4 + group_1 * 32; - int rows_1[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&rows_1[0])), "=r"(*reinterpret_cast(&rows_1[(0) + 1])), "=r"(*reinterpret_cast(&rows_1[(0) + 2])), "=r"(*reinterpret_cast(&rows_1[(0) + 3])) - : "r"(smem_page_offsets_addr + page_k_stage * 1024 + (unsigned int)((page_base_0 + row_base_1) * 4))); - int dst_k_1 = smem_kv_addr + kv_stage_1 * 32768 + (unsigned int)(row_base_1 * 64 * 2); - if (elect_sync()) { - if (next_tile == 0) { - tma_gather4_gmem2smem(dst_k_1, tmap_swa_kv, hs_sel_3 * 128, rows_1[0], rows_1[1], rows_1[2], rows_1[3], kv_full_addr + (kv_stage_1) * 8); - tma_gather4_gmem2smem(dst_k_1 + 16384, tmap_swa_kv, hs_sel_3 * 128 + 64, rows_1[0], rows_1[1], rows_1[2], rows_1[3], kv_full_addr + (kv_stage_1) * 8); - } else { - tma_gather4_gmem2smem(dst_k_1, tmap_compressed_kv, hs_sel_3 * 128, rows_1[0], rows_1[1], rows_1[2], rows_1[3], kv_full_addr + (kv_stage_1) * 8); - tma_gather4_gmem2smem(dst_k_1 + 16384, tmap_compressed_kv, hs_sel_3 * 128 + 64, rows_1[0], rows_1[1], rows_1[2], rows_1[3], kv_full_addr + (kv_stage_1) * 8); + if (load_rank < ((hs_4 >= 2) ? warps_ge1_4 : warps_ge2_5)) { + #pragma unroll + for (int group_1 = 0; group_1 < 4; group_1++) { + const int row_base_1 = load_rank * 4 + group_1 * 32; + int rows_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&rows_1[0])), "=r"(*reinterpret_cast(&rows_1[(0) + 1])), "=r"(*reinterpret_cast(&rows_1[(0) + 2])), "=r"(*reinterpret_cast(&rows_1[(0) + 3])) + : "r"(smem_page_offsets_addr + page_k_stage * 1024 + (unsigned int)((page_base_0 + row_base_1) * 4))); + int dst_k_1 = smem_kv_addr + kv_stage_1 * 32768 + (unsigned int)(row_base_1 * 64 * 2); + if (elect_sync()) { + if (next_tile == 0) { + tma_gather4_gmem2smem(dst_k_1, tmap_swa_kv, hs_sel_3 * 128, rows_1[0], rows_1[1], rows_1[2], rows_1[3], kv_full_addr + (kv_stage_1) * 8); + tma_gather4_gmem2smem(dst_k_1 + 16384, tmap_swa_kv, hs_sel_3 * 128 + 64, rows_1[0], rows_1[1], rows_1[2], rows_1[3], kv_full_addr + (kv_stage_1) * 8); + } else { + tma_gather4_gmem2smem(dst_k_1, tmap_compressed_kv, hs_sel_3 * 128, rows_1[0], rows_1[1], rows_1[2], rows_1[3], kv_full_addr + (kv_stage_1) * 8); + tma_gather4_gmem2smem(dst_k_1 + 16384, tmap_compressed_kv, hs_sel_3 * 128 + 64, rows_1[0], rows_1[1], rows_1[2], rows_1[3], kv_full_addr + (kv_stage_1) * 8); + } } } } diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_family_956aa0f74f4a82042f2c_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_family_0d09e8b242706c4e05a9_binding.cu similarity index 87% rename from csrc/cake_dsv4/sm_103a/cake_dsv4_family_956aa0f74f4a82042f2c_binding.cu rename to csrc/cake_dsv4/sm_103a/cake_dsv4_family_0d09e8b242706c4e05a9_binding.cu index b5b1f4e43a8..bc529abed8b 100644 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_family_956aa0f74f4a82042f2c_binding.cu +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_family_0d09e8b242706c4e05a9_binding.cu @@ -15,7 +15,7 @@ */ // Generated by cake.runtime.ffi_dispatch — do not edit. -// tvm-ffi dispatch module for kernel family 'cake_dsv4_family_956aa0f74f4a82042f2c'. +// tvm-ffi dispatch module for kernel family 'cake_dsv4_family_0d09e8b242706c4e05a9'. // dispatch-policy-revision: unpinned #include #include @@ -34,9 +34,9 @@ #include extern "C" __global__ void kernel_cake_dsv4_230143c89d42c6c89387(__nv_bfloat16* __restrict__ partial_O, float* __restrict__ partial_lse, __nv_bfloat16* __restrict__ O, int num_q_heads, int num_split); -extern "C" __global__ void kernel_cake_dsv4_125332c2066d89bd7610(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); +extern "C" __global__ void kernel_cake_dsv4_4ab0487d5fde70d4569e(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, const __grid_constant__ CUtensorMap O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); -namespace cake_host_shim_3743d63f7b758ba5 { +namespace cake_host_shim_4bb3ed291da4a466 { using tvm::ffi::Optional; using tvm::ffi::TensorView; @@ -161,7 +161,7 @@ inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* } } -namespace variant_cake_dsv4_125332c2066d89bd7610_sm_103a_aaa5b590c060 { +namespace variant_cake_dsv4_4ab0487d5fde70d4569e_sm_103a_7e73b6214907 { // 4D TMA descriptor for buffer 'tmap_q' — compiled from the // descriptor's std.Expr global_dim/global_strides/checks record. @@ -394,6 +394,57 @@ inline CUtensorMap EncodeTma_tmap_compressed_v(const TensorView& t) { return tm; } +// 4D TMA descriptor for buffer 'O' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_O(const TensorView& t, int64_t tp_num_heads, int64_t tp_num_query_tokens) { + TVM_FFI_CHECK(t.ndim() >= 1, ValueError) + << "TMA source 'O' must have at least 1 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'O' must have unit innermost stride, got " << t.stride(-1); + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(8), (uint64_t)(5), (uint64_t)((tp_num_query_tokens * tp_num_heads))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'O' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1] && 1u <= global_dim[2] && 32u <= global_dim[3], ValueError) + << "TMA box (64, 1, 1, 32) exceeds resolved global dims for 'O'"; + int64_t carrier_stride_0 = 64; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 512; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = 2560; + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 1u, 1u, 32u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'O') failed: CUresult=" << (int)r; + return tm; +} + void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_compressed_k, TensorView arg_tmap_swa_v, TensorView arg_tmap_compressed_v, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_num_query_tokens, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { DLDevice dev = arg_tmap_q.device(); ScopedCudaDevice device_guard(dev.device_id); @@ -481,7 +532,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_c CUtensorMap p_tmap_compressed_k = EncodeTma_tmap_compressed_k(arg_tmap_compressed_k); CUtensorMap p_tmap_swa_v = EncodeTma_tmap_swa_v(arg_tmap_swa_v); CUtensorMap p_tmap_compressed_v = EncodeTma_tmap_compressed_v(arg_tmap_compressed_v); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); + CUtensorMap p_O = EncodeTma_O(arg_O, (int64_t)arg_num_heads, (int64_t)arg_num_query_tokens); float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); @@ -514,7 +565,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_c } if (smem_status == cudaSuccess) { smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_125332c2066d89bd7610), + reinterpret_cast(kernel_cake_dsv4_4ab0487d5fde70d4569e), cudaFuncAttributeMaxDynamicSharedMemorySize, 231936); } @@ -524,7 +575,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_c } } TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_125332c2066d89bd7610 failed: " + << "cudaFuncSetAttribute for kernel_cake_dsv4_4ab0487d5fde70d4569e failed: " << cudaGetErrorString(smem_status); dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); dim3 block(512u, 1u, 1u); @@ -545,12 +596,12 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_c config.attrs = attrs; config.numAttrs = n; cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_125332c2066d89bd7610), kargs); + &config, reinterpret_cast(kernel_cake_dsv4_4ab0487d5fde70d4569e), kargs); TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_125332c2066d89bd7610 failed: " + << "cudaLaunchKernelExC for kernel_cake_dsv4_4ab0487d5fde70d4569e failed: " << cudaGetErrorString(launch_status); } -} // namespace variant_cake_dsv4_125332c2066d89bd7610_sm_103a_aaa5b590c060 +} // namespace variant_cake_dsv4_4ab0487d5fde70d4569e_sm_103a_7e73b6214907 namespace variant_cake_dsv4_230143c89d42c6c89387_sm_103a_c6d3bf2db9b6 { @@ -629,24 +680,24 @@ void Dispatch(tvm::ffi::TensorView arg_Q, tvm::ffi::TensorView arg_SWA_cache, tv int64_t gx = (total_work_items * 2); int64_t gy = 1; int64_t gz = 1; - cake_host_shim_3743d63f7b758ba5::variant_cake_dsv4_125332c2066d89bd7610_sm_103a_aaa5b590c060::Run(arg_Q, arg_SWA_cache, arg_compressed_KV_cache, arg_SWA_cache, arg_compressed_KV_cache, arg_partial_O, arg_partial_lse, arg_swa_indices, arg_compressed_indices, arg_sparse_topk_lens, arg_sinks, arg_bmm1_scale, arg_bmm2_scale, num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, gx, gy, gz, stream); + cake_host_shim_4bb3ed291da4a466::variant_cake_dsv4_4ab0487d5fde70d4569e_sm_103a_7e73b6214907::Run(arg_Q, arg_SWA_cache, arg_compressed_KV_cache, arg_SWA_cache, arg_compressed_KV_cache, arg_partial_O, arg_partial_lse, arg_swa_indices, arg_compressed_indices, arg_sparse_topk_lens, arg_sinks, arg_bmm1_scale, arg_bmm2_scale, num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, gx, gy, gz, stream); } { int64_t gx = num_query_tokens; int64_t gy = (num_heads / 4); int64_t gz = 1; - cake_host_shim_3743d63f7b758ba5::variant_cake_dsv4_230143c89d42c6c89387_sm_103a_c6d3bf2db9b6::Run(arg_partial_O, arg_partial_lse, arg_O, num_heads, num_split, gx, gy, gz, stream); + cake_host_shim_4bb3ed291da4a466::variant_cake_dsv4_230143c89d42c6c89387_sm_103a_c6d3bf2db9b6::Run(arg_partial_O, arg_partial_lse, arg_O, num_heads, num_split, gx, gy, gz, stream); } return; } default: - TVM_FFI_CHECK(false, ValueError) << "no dispatch route for cake_dsv4_family_956aa0f74f4a82042f2c shape (num_heads=" << num_heads << ", num_query_tokens=" << num_query_tokens << ", swa_index_stride=" << swa_index_stride << ", compressed_index_stride=" << compressed_index_stride << ", sparse_topk_lens_offset=" << sparse_topk_lens_offset << ", sparse_topk=" << sparse_topk << ", has_sinks=" << has_sinks << ", total_work_items=" << total_work_items << ", num_split=" << num_split << ")"; + TVM_FFI_CHECK(false, ValueError) << "no dispatch route for cake_dsv4_family_0d09e8b242706c4e05a9 shape (num_heads=" << num_heads << ", num_query_tokens=" << num_query_tokens << ", swa_index_stride=" << swa_index_stride << ", compressed_index_stride=" << compressed_index_stride << ", sparse_topk_lens_offset=" << sparse_topk_lens_offset << ", sparse_topk=" << sparse_topk << ", has_sinks=" << has_sinks << ", total_work_items=" << total_work_items << ", num_split=" << num_split << ")"; } } int64_t RouteOf(int32_t num_heads, int32_t num_query_tokens, int32_t swa_index_stride, int32_t compressed_index_stride, int32_t sparse_topk_lens_offset, int32_t sparse_topk, int32_t has_sinks, int32_t total_work_items, int32_t num_split) { return SelectVariant(num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, num_split); } -} // namespace cake_host_shim_3743d63f7b758ba5 +} // namespace cake_host_shim_4bb3ed291da4a466 -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_3743d63f7b758ba5::Dispatch); -TVM_FFI_DLL_EXPORT_TYPED_FUNC(route_of, cake_host_shim_3743d63f7b758ba5::RouteOf); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_4bb3ed291da4a466::Dispatch); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(route_of, cake_host_shim_4bb3ed291da4a466::RouteOf); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_family_ec44ddd94e3be88cfc83_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_family_ebb85adb872529ca9e83_binding.cu similarity index 87% rename from csrc/cake_dsv4/sm_103a/cake_dsv4_family_ec44ddd94e3be88cfc83_binding.cu rename to csrc/cake_dsv4/sm_103a/cake_dsv4_family_ebb85adb872529ca9e83_binding.cu index f7090b2f731..aa2fd81aeb1 100644 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_family_ec44ddd94e3be88cfc83_binding.cu +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_family_ebb85adb872529ca9e83_binding.cu @@ -15,7 +15,7 @@ */ // Generated by cake.runtime.ffi_dispatch — do not edit. -// tvm-ffi dispatch module for kernel family 'cake_dsv4_family_ec44ddd94e3be88cfc83'. +// tvm-ffi dispatch module for kernel family 'cake_dsv4_family_ebb85adb872529ca9e83'. // dispatch-policy-revision: unpinned #include #include @@ -34,9 +34,9 @@ #include extern "C" __global__ void kernel_cake_dsv4_397f817e2931a6cc7c92(__nv_bfloat16* __restrict__ partial_O, float* __restrict__ partial_lse, __nv_bfloat16* __restrict__ O, int num_q_heads, int num_split); -extern "C" __global__ void kernel_cake_dsv4_4c41fb8347f137d4d46f(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); +extern "C" __global__ void kernel_cake_dsv4_064850af5858dd03c104(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, const __grid_constant__ CUtensorMap O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); -namespace cake_host_shim_9b7f43581bd6022f { +namespace cake_host_shim_71d3307de7b4df2d { using tvm::ffi::Optional; using tvm::ffi::TensorView; @@ -161,7 +161,7 @@ inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* } } -namespace variant_cake_dsv4_4c41fb8347f137d4d46f_sm_103a_994fe0d2939f { +namespace variant_cake_dsv4_064850af5858dd03c104_sm_103a_6276393afd19 { // 4D TMA descriptor for buffer 'tmap_q' — compiled from the // descriptor's std.Expr global_dim/global_strides/checks record. @@ -394,6 +394,57 @@ inline CUtensorMap EncodeTma_tmap_compressed_v(const TensorView& t) { return tm; } +// 4D TMA descriptor for buffer 'O' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_O(const TensorView& t, int64_t tp_num_heads, int64_t tp_num_query_tokens) { + TVM_FFI_CHECK(t.ndim() >= 1, ValueError) + << "TMA source 'O' must have at least 1 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'O' must have unit innermost stride, got " << t.stride(-1); + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(8), (uint64_t)(4), (uint64_t)((tp_num_query_tokens * tp_num_heads))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'O' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1] && 1u <= global_dim[2] && 32u <= global_dim[3], ValueError) + << "TMA box (64, 1, 1, 32) exceeds resolved global dims for 'O'"; + int64_t carrier_stride_0 = 64; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 512; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = 2048; + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 1u, 1u, 32u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'O') failed: CUresult=" << (int)r; + return tm; +} + void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_compressed_k, TensorView arg_tmap_swa_v, TensorView arg_tmap_compressed_v, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_num_query_tokens, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { DLDevice dev = arg_tmap_q.device(); ScopedCudaDevice device_guard(dev.device_id); @@ -481,7 +532,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_c CUtensorMap p_tmap_compressed_k = EncodeTma_tmap_compressed_k(arg_tmap_compressed_k); CUtensorMap p_tmap_swa_v = EncodeTma_tmap_swa_v(arg_tmap_swa_v); CUtensorMap p_tmap_compressed_v = EncodeTma_tmap_compressed_v(arg_tmap_compressed_v); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); + CUtensorMap p_O = EncodeTma_O(arg_O, (int64_t)arg_num_heads, (int64_t)arg_num_query_tokens); float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); @@ -514,7 +565,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_c } if (smem_status == cudaSuccess) { smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_4c41fb8347f137d4d46f), + reinterpret_cast(kernel_cake_dsv4_064850af5858dd03c104), cudaFuncAttributeMaxDynamicSharedMemorySize, 220032); } @@ -524,7 +575,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_c } } TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_4c41fb8347f137d4d46f failed: " + << "cudaFuncSetAttribute for kernel_cake_dsv4_064850af5858dd03c104 failed: " << cudaGetErrorString(smem_status); dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); dim3 block(512u, 1u, 1u); @@ -545,12 +596,12 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_c config.attrs = attrs; config.numAttrs = n; cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_4c41fb8347f137d4d46f), kargs); + &config, reinterpret_cast(kernel_cake_dsv4_064850af5858dd03c104), kargs); TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_4c41fb8347f137d4d46f failed: " + << "cudaLaunchKernelExC for kernel_cake_dsv4_064850af5858dd03c104 failed: " << cudaGetErrorString(launch_status); } -} // namespace variant_cake_dsv4_4c41fb8347f137d4d46f_sm_103a_994fe0d2939f +} // namespace variant_cake_dsv4_064850af5858dd03c104_sm_103a_6276393afd19 namespace variant_cake_dsv4_397f817e2931a6cc7c92_sm_103a_aa3105f3b9de { @@ -629,24 +680,24 @@ void Dispatch(tvm::ffi::TensorView arg_Q, tvm::ffi::TensorView arg_SWA_cache, tv int64_t gx = (total_work_items * 2); int64_t gy = 1; int64_t gz = 1; - cake_host_shim_9b7f43581bd6022f::variant_cake_dsv4_4c41fb8347f137d4d46f_sm_103a_994fe0d2939f::Run(arg_Q, arg_SWA_cache, arg_compressed_KV_cache, arg_SWA_cache, arg_compressed_KV_cache, arg_partial_O, arg_partial_lse, arg_swa_indices, arg_compressed_indices, arg_sparse_topk_lens, arg_sinks, arg_bmm1_scale, arg_bmm2_scale, num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, gx, gy, gz, stream); + cake_host_shim_71d3307de7b4df2d::variant_cake_dsv4_064850af5858dd03c104_sm_103a_6276393afd19::Run(arg_Q, arg_SWA_cache, arg_compressed_KV_cache, arg_SWA_cache, arg_compressed_KV_cache, arg_partial_O, arg_partial_lse, arg_swa_indices, arg_compressed_indices, arg_sparse_topk_lens, arg_sinks, arg_bmm1_scale, arg_bmm2_scale, num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, gx, gy, gz, stream); } { int64_t gx = num_query_tokens; int64_t gy = num_heads; int64_t gz = 1; - cake_host_shim_9b7f43581bd6022f::variant_cake_dsv4_397f817e2931a6cc7c92_sm_103a_aa3105f3b9de::Run(arg_partial_O, arg_partial_lse, arg_O, num_heads, num_split, gx, gy, gz, stream); + cake_host_shim_71d3307de7b4df2d::variant_cake_dsv4_397f817e2931a6cc7c92_sm_103a_aa3105f3b9de::Run(arg_partial_O, arg_partial_lse, arg_O, num_heads, num_split, gx, gy, gz, stream); } return; } default: - TVM_FFI_CHECK(false, ValueError) << "no dispatch route for cake_dsv4_family_ec44ddd94e3be88cfc83 shape (num_heads=" << num_heads << ", num_query_tokens=" << num_query_tokens << ", swa_index_stride=" << swa_index_stride << ", compressed_index_stride=" << compressed_index_stride << ", sparse_topk_lens_offset=" << sparse_topk_lens_offset << ", sparse_topk=" << sparse_topk << ", has_sinks=" << has_sinks << ", total_work_items=" << total_work_items << ", num_split=" << num_split << ")"; + TVM_FFI_CHECK(false, ValueError) << "no dispatch route for cake_dsv4_family_ebb85adb872529ca9e83 shape (num_heads=" << num_heads << ", num_query_tokens=" << num_query_tokens << ", swa_index_stride=" << swa_index_stride << ", compressed_index_stride=" << compressed_index_stride << ", sparse_topk_lens_offset=" << sparse_topk_lens_offset << ", sparse_topk=" << sparse_topk << ", has_sinks=" << has_sinks << ", total_work_items=" << total_work_items << ", num_split=" << num_split << ")"; } } int64_t RouteOf(int32_t num_heads, int32_t num_query_tokens, int32_t swa_index_stride, int32_t compressed_index_stride, int32_t sparse_topk_lens_offset, int32_t sparse_topk, int32_t has_sinks, int32_t total_work_items, int32_t num_split) { return SelectVariant(num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, num_split); } -} // namespace cake_host_shim_9b7f43581bd6022f +} // namespace cake_host_shim_71d3307de7b4df2d -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_9b7f43581bd6022f::Dispatch); -TVM_FFI_DLL_EXPORT_TYPED_FUNC(route_of, cake_host_shim_9b7f43581bd6022f::RouteOf); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_71d3307de7b4df2d::Dispatch); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(route_of, cake_host_shim_71d3307de7b4df2d::RouteOf); diff --git a/flashinfer/jit/cake_dsv4.py b/flashinfer/jit/cake_dsv4.py index 4381bdac79c..10b77c359f2 100644 --- a/flashinfer/jit/cake_dsv4.py +++ b/flashinfer/jit/cake_dsv4.py @@ -85,7 +85,7 @@ }, "sources": ["cake_dsv4_family_06111202fc1bd2f72060_binding.cu"], }, - "cake_dsv4_family_068c73b21d7cbcd958b8": { + "cake_dsv4_family_5383a615a7e502f6fdac": { "arch": "sm_100a", "compile_flags": ["-std=c++17"], "dependencies": [ @@ -94,7 +94,7 @@ ], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "567dd4bd885f3cb2e808cfe4aa2012ef8bbc25e3698b5a7febf987f09df63eb9", + "identity": "78bbedd13dc7248f5985a1e029ffb2bd8e7be9402c95fda93a698cb73e806ace", "route_entry": "route_of", "routes": [ { @@ -159,15 +159,15 @@ "workspace_dtypes": {}, "workspace_keys": [], }, - "sources": ["cake_dsv4_family_068c73b21d7cbcd958b8_binding.cu"], + "sources": ["cake_dsv4_family_5383a615a7e502f6fdac_binding.cu"], }, - "cake_dsv4_family_a84dbf9944ef3cfb088b": { + "cake_dsv4_family_9c96528827a52c42bbd7": { "arch": "sm_100a", "compile_flags": ["-std=c++17"], "dependencies": ["bf16_h128_topk128x_split4_sm100", "split_reduce"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "7a1c48125738408e45898acfc24926b83b872852cfb31556b4a0dc5cf6b30b8c", + "identity": "048575f321b03f7bd49bbb1f639d8f3790b65d4044648b653e005f9ba51de109", "route_entry": "route_of", "routes": [ { @@ -232,7 +232,7 @@ "workspace_dtypes": {}, "workspace_keys": [], }, - "sources": ["cake_dsv4_family_a84dbf9944ef3cfb088b_binding.cu"], + "sources": ["cake_dsv4_family_9c96528827a52c42bbd7_binding.cu"], }, "cake_dsv4_family_c6fc162d8e0b1752b4f7": { "arch": "sm_100a", @@ -391,7 +391,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "b3715376b8c6e27f09d941cd1d259130d7a54170640d87e0b534d83708ead0fd", + "identity": "bddbb08633a8c728dfaafad7d7d2bebf224379cf83d5b0ba3637eb6520917fb7", "sources": [ "cake_dsv4_314d943a39ab7eade2c8_kernel.cu", "cake_dsv4_314d943a39ab7eade2c8_binding.cu", @@ -429,7 +429,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "24344094f11b37ec4b591eed3a095f38440fc3d66ab7d03357d47093633b9cbb", + "identity": "5f0ec3e7b912f6be4f30e9622a5f2ac879da9c5f09963122bcb02004f32e0f3e", "sources": [ "cake_dsv4_d994492f2dfb135a7dbf_kernel.cu", "cake_dsv4_d994492f2dfb135a7dbf_binding.cu", @@ -451,7 +451,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "1b1cc19373e52f67981ef2cff298918c5579ee69afa319f8c1ed8c1c5cf4ee7c", + "identity": "056b13094b2c64be626f377c14517e7bf67014b79476caae010d9e927d8ce4c7", "sources": [ "cake_dsv4_d906a03bd66a7488fc29_kernel.cu", "cake_dsv4_d906a03bd66a7488fc29_binding.cu", @@ -484,10 +484,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "82976f52e6a753f3e2fe9a76552492077e07345543038a013d52d92ae7d037a2", + "identity": "4c903082c53d323a454998a3b5305b7bcf38cc671c2661ab6426801494b3731d", "sources": [ - "cake_dsv4_57e7b93def70f2be356f_kernel.cu", - "cake_dsv4_57e7b93def70f2be356f_binding.cu", + "cake_dsv4_0c8ab955ac7b62e7bfee_kernel.cu", + "cake_dsv4_0c8ab955ac7b62e7bfee_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -522,7 +522,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "f2200c038a3de20ae07fe52f1a3db30d71941feff6d47fe4db8cd1a45e9d9d44", + "identity": "1bbf6b3a268b2084c0832b70c721618ad84d91c6feb75adba665af2238a7da5a", "sources": [ "cake_dsv4_e6d521df78ecb146fd49_kernel.cu", "cake_dsv4_e6d521df78ecb146fd49_binding.cu", @@ -537,7 +537,7 @@ ["tma_buffer", "tmap_compressed_k"], ["tma_buffer", "tmap_swa_v"], ["tma_buffer", "tmap_compressed_v"], - ["buffer", "O"], + ["tma_buffer", "O"], ["buffer", "partial_lse"], ["buffer", "swa_indices"], ["buffer", "compressed_indices"], @@ -560,10 +560,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "5609e5563448b459fc62af27735acd289478d73d4dfa97cc276009a78b472705", + "identity": "175b037118c3f8d346bc89148513a38f643f541121ed3485f18b9d8c58694b44", "sources": [ - "cake_dsv4_e1660a7f44d156286fea_kernel.cu", - "cake_dsv4_e1660a7f44d156286fea_binding.cu", + "cake_dsv4_42fa8e98da5fa46d1a94_kernel.cu", + "cake_dsv4_42fa8e98da5fa46d1a94_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -575,7 +575,7 @@ ["tma_buffer", "tmap_compressed_k"], ["tma_buffer", "tmap_swa_v"], ["tma_buffer", "tmap_compressed_v"], - ["buffer", "O"], + ["tma_buffer", "O"], ["buffer", "partial_lse"], ["buffer", "swa_indices"], ["buffer", "compressed_indices"], @@ -598,10 +598,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "11243ea83b1c9138a3751dda4de6de1d9ba24712eddce8a30c46fa7027f09e4f", + "identity": "6af64a72ba546c5afa01185851bbd377cb0ee1b41e85a30292124934a9518da3", "sources": [ - "cake_dsv4_ba9ed6a6233761682cc3_kernel.cu", - "cake_dsv4_ba9ed6a6233761682cc3_binding.cu", + "cake_dsv4_714df747e199d0adbe08_kernel.cu", + "cake_dsv4_714df747e199d0adbe08_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -630,7 +630,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "dce42787024ee775d92b9917c6aef6c75769aa31d20cc420a5986533fae9dc29", + "identity": "d68dac00e6e7cf66a1d5ada72032d13c8f4d11349ccbba93b4dc9d69bcbedb6c", "sources": [ "cake_dsv4_a96a3de25becc6c0a422_kernel.cu", "cake_dsv4_a96a3de25becc6c0a422_binding.cu", @@ -668,7 +668,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "7a08a9d6d26f32e66a62402142d1fdb47a1358542fe79be3c76445e4826c6824", + "identity": "ac092d7073e4b4a91c2d48a58a124dfb9e682342c050bac9688017681469e0ee", "sources": [ "cake_dsv4_c0e3d0af9a266850de48_kernel.cu", "cake_dsv4_c0e3d0af9a266850de48_binding.cu", @@ -703,7 +703,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "fe3d4502bb2dfe17a4a4d94fe8f02215e9ee7ae48fd4c70ef3e29326fe7e3640", + "identity": "a910d48878523c7afc2a69bda42a34e7b36b2731d2c4a633592bad70e22551fd", "sources": [ "cake_dsv4_422eafcc723b03b8e044_kernel.cu", "cake_dsv4_422eafcc723b03b8e044_binding.cu", @@ -725,7 +725,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "2e14912ba8d257e634816d04ff5cce6abc380b94e8b771da784bb8e3d97e2b53", + "identity": "1c7462a3ce0fef598330a20566a419bf922d70f655b6493501807217ffac4232", "sources": [ "cake_dsv4_697c8662eb01c8d6bd50_kernel.cu", "cake_dsv4_697c8662eb01c8d6bd50_binding.cu", @@ -763,10 +763,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "e1b220a2d15a79b981a5e717f4d92d99f416852bcea16cb3c6d377e39c9d9671", + "identity": "427e28dadac3a6dfe75d9810c7687601813da18412b5afe0fc94a074b9f09ca9", "sources": [ - "cake_dsv4_3b3789d6356f15f9cbe3_kernel.cu", - "cake_dsv4_3b3789d6356f15f9cbe3_binding.cu", + "cake_dsv4_09e14f1d9cbd6739db61_kernel.cu", + "cake_dsv4_09e14f1d9cbd6739db61_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -797,10 +797,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "53b1e34a7eaa5c285ff343627132fc0cf9a943a2f7860e2bc10cb45797d989f6", + "identity": "cf8478ea922b580842f74f3ea666062aa418264360e9d65577a18fdd4f8bdedc", "sources": [ - "cake_dsv4_9526f80c89f2b5ef4c6c_kernel.cu", - "cake_dsv4_9526f80c89f2b5ef4c6c_binding.cu", + "cake_dsv4_15eabf5d2a5f50db201b_kernel.cu", + "cake_dsv4_15eabf5d2a5f50db201b_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -833,7 +833,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "14152b76d08f7df3fc1e57ac1a9c48ee800539333b92cd59c65ec148f1025558", + "identity": "f923f90544f3007aa8e2db40fa8998257b8dc655c2fe04fe04b4d86a2bcc7853", "sources": [ "cake_dsv4_bf2bffa3152d5dce396f_kernel.cu", "cake_dsv4_bf2bffa3152d5dce396f_binding.cu", @@ -865,7 +865,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "d9b38ca3f8c4bf622b0e37322c11c66c33889b7444792e64a4fbb398f0671e71", + "identity": "cf95eea83e8bdabd58514f45be73fa45bb47398dc513f838d56a6dd5baf873c4", "sources": [ "cake_dsv4_88d5662b6147fe74ed0c_kernel.cu", "cake_dsv4_88d5662b6147fe74ed0c_binding.cu", @@ -897,10 +897,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "db3a18b88f9a22de1bbaf8e18f45b1281ce7ed5f69a7a102ea55af5e8ca5cf24", + "identity": "b6d2b855936ced61ba3dfa126235e6c1fde09b383afa90c887e6c26b7f036592", "sources": [ - "cake_dsv4_4ac65f00a4a032c10ecb_kernel.cu", - "cake_dsv4_4ac65f00a4a032c10ecb_binding.cu", + "cake_dsv4_12ac507a3c73f2784621_kernel.cu", + "cake_dsv4_12ac507a3c73f2784621_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -941,7 +941,7 @@ ], "entry": "run", "host_linkage_flags": [], - "identity": "2c813be0fac81518383683a7faa658c9e452fb1f1966745bdbc4306327d3dee1", + "identity": "5001ca9df604adc9b0728459de36a9f4fd2f544d95d5b6363eb8dd740e4ce4c2", "sources": [ "cake_dsv4_2dd31513440dbd1a0ef2_kernel.cu", "cake_dsv4_2dd31513440dbd1a0ef2_binding.cu", @@ -982,7 +982,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "4dce24ad4f95460dfa2ecc045bf899c67c4e7927a1bb362b2f3cf1d5d24ceb26", + "identity": "8968f0f75b48b5e140c407a3610819fcf62395d53fb2d97d25c896f52b60dfd8", "sources": [ "cake_dsv4_77c162685d381298b786_kernel.cu", "cake_dsv4_77c162685d381298b786_binding.cu", @@ -1022,7 +1022,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "ca35d52a3fbfe936fa85d8f8e4c62b2f624c9b1043519640803738047d5c17c3", + "identity": "17011dd54760db735a19a3b3f6ac9312dd9ab0baa8828fd98ce06cd66deaa64b", "sources": [ "cake_dsv4_119728d98fe411f18a22_kernel.cu", "cake_dsv4_119728d98fe411f18a22_binding.cu", @@ -1057,10 +1057,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "1dc5dc743aca6bb61cb8956a6906a9cb03997d40adfc53f5eb5308dc16fd279a", + "identity": "ec7155bae85cda09aa4c0a15c0b804ffbe91284f388c51a7e60971bae78adde3", "sources": [ - "cake_dsv4_5e5d6648b7bd42bcd0d6_kernel.cu", - "cake_dsv4_5e5d6648b7bd42bcd0d6_binding.cu", + "cake_dsv4_bab2852cc0ea4a6d4860_kernel.cu", + "cake_dsv4_bab2852cc0ea4a6d4860_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -1093,10 +1093,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "7777f8fda7eda670a2733c023b36fd5c05a0cef563e051a8a715ed22d5cc8b79", + "identity": "b4d0530f81c68e3ce8c5d71100116cd46594116ec0c3baa5a71435b99ef03033", "sources": [ - "cake_dsv4_fbf138c989dfde9f764f_kernel.cu", - "cake_dsv4_fbf138c989dfde9f764f_binding.cu", + "cake_dsv4_21a93d40caefcf5b252d_kernel.cu", + "cake_dsv4_21a93d40caefcf5b252d_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -1129,10 +1129,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "3640ee1967a674a977c4ce19a91cfc1c0095d43d59c6ea1fea8cd6c9157331cc", + "identity": "b49f25ffcfc867c0feb10b2184dc867b5e9538760d711e94b08ac337f028099f", "sources": [ - "cake_dsv4_7771bcc5c13a5ff83890_kernel.cu", - "cake_dsv4_7771bcc5c13a5ff83890_binding.cu", + "cake_dsv4_5703afd5b398f2b6c980_kernel.cu", + "cake_dsv4_5703afd5b398f2b6c980_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -1168,10 +1168,10 @@ ], "entry": "run", "host_linkage_flags": [], - "identity": "ef787d9bc5f14ef9eab0e77e05e3e8a3a350341db6b57f23816515a16a92b7d0", + "identity": "c9c07f67f19855f4b65189ccf12c9f9154c53c4fe8b771f198aa2b6ebe9b3948", "sources": [ - "cake_dsv4_fb6409904b35ed0f7a1a_kernel.cu", - "cake_dsv4_fb6409904b35ed0f7a1a_binding.cu", + "cake_dsv4_9c3b9b4fa7cba6e44ffc_kernel.cu", + "cake_dsv4_9c3b9b4fa7cba6e44ffc_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -1204,10 +1204,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "2a1fc4a24d65057c1db942da73a0d62ca1b02b8a1a1002b05532990638dd2aca", + "identity": "7f336bc7088175237e4e51cbd4b8e212eb7e9024aaa080e7d655d611bbd6679e", "sources": [ - "cake_dsv4_75827bd38b443de66f67_kernel.cu", - "cake_dsv4_75827bd38b443de66f67_binding.cu", + "cake_dsv4_530fe7e6cb8c7d09744d_kernel.cu", + "cake_dsv4_530fe7e6cb8c7d09744d_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -1226,7 +1226,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "014512652fd83c055323b34e9c4807829313cd35fbcde3a1d44cd87e89404265", + "identity": "f3b805d486f48af9a2e1e8bc5d443a18f1509b90ce0326a93cb92e64a8fa2589", "sources": [ "cake_dsv4_d50dd820272ad24f7027_kernel.cu", "cake_dsv4_d50dd820272ad24f7027_binding.cu", @@ -1237,21 +1237,33 @@ }, "sm_103a": { "programs": { - "cake_dsv4_family_1b5ee6b39dec3c3538c6": { + "cake_dsv4_family_0d09e8b242706c4e05a9": { "arch": "sm_103a", "compile_flags": ["-std=c++17"], - "dependencies": ["bf16_h128_prefill_v42"], + "dependencies": [ + "bf16_h128_topk4x_v52", + "split_reduce_cake_dsv4_230143c89d42c6c89387", + ], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "860901dbab47cf43c94ed40b20ba0847a3137cf7deb3f0a88e699c67963ca890", + "identity": "897382fde021af3afe3cbf39d8771c6f45a849efd11f723e57ca4fe800f91f6a", "route_entry": "route_of", - "routes": [{"index": 0, "stage_variants": ["bf16_h128_prefill_v42"]}], + "routes": [ + { + "index": 0, + "stage_variants": [ + "bf16_h128_topk4x_v52", + "split_reduce_cake_dsv4_230143c89d42c6c89387", + ], + } + ], "signature": { "scalar_dtypes": { "compressed_index_stride": "int32", "has_sinks": "int32", "num_heads": "int32", "num_query_tokens": "int32", + "num_split": "int32", "sparse_topk": "int32", "sparse_topk_lens_offset": "int32", "swa_index_stride": "int32", @@ -1266,6 +1278,7 @@ "sparse_topk", "has_sinks", "total_work_items", + "num_split", ], "tensor_dtypes": { "O": "bfloat16", @@ -1275,6 +1288,7 @@ "bmm2_scale": "float32", "compressed_KV_cache": "bfloat16", "compressed_indices": "int32", + "partial_O": "bfloat16", "partial_lse": "float32", "sinks": "float32", "sparse_topk_lens": "int32", @@ -1284,6 +1298,7 @@ "Q", "SWA_cache", "compressed_KV_cache", + "partial_O", "partial_lse", "swa_indices", "compressed_indices", @@ -1296,19 +1311,17 @@ "workspace_dtypes": {}, "workspace_keys": [], }, - "sources": ["cake_dsv4_family_1b5ee6b39dec3c3538c6_binding.cu"], + "sources": ["cake_dsv4_family_0d09e8b242706c4e05a9_binding.cu"], }, - "cake_dsv4_family_42ae3e93a01e7a3bf20d": { + "cake_dsv4_family_1b5ee6b39dec3c3538c6": { "arch": "sm_103a", "compile_flags": ["-std=c++17"], - "dependencies": ["bf16_h128_prefill_v42_snake"], + "dependencies": ["bf16_h128_prefill_v42"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "285cf074aae8f3086ca239ade301c7eb313c22c15286c2f26f0386be98ebb4f5", + "identity": "860901dbab47cf43c94ed40b20ba0847a3137cf7deb3f0a88e699c67963ca890", "route_entry": "route_of", - "routes": [ - {"index": 0, "stage_variants": ["bf16_h128_prefill_v42_snake"]} - ], + "routes": [{"index": 0, "stage_variants": ["bf16_h128_prefill_v42"]}], "signature": { "scalar_dtypes": { "compressed_index_stride": "int32", @@ -1359,18 +1372,18 @@ "workspace_dtypes": {}, "workspace_keys": [], }, - "sources": ["cake_dsv4_family_42ae3e93a01e7a3bf20d_binding.cu"], + "sources": ["cake_dsv4_family_1b5ee6b39dec3c3538c6_binding.cu"], }, - "cake_dsv4_family_585100536b80a2cec8ba": { + "cake_dsv4_family_42ae3e93a01e7a3bf20d": { "arch": "sm_103a", "compile_flags": ["-std=c++17"], - "dependencies": ["bf16_h128_topk128x_row_first"], + "dependencies": ["bf16_h128_prefill_v42_snake"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "72b48d77a88ee5d7e37313209035e21d2037d2b74d4793d7660367a6c9003750", + "identity": "285cf074aae8f3086ca239ade301c7eb313c22c15286c2f26f0386be98ebb4f5", "route_entry": "route_of", "routes": [ - {"index": 0, "stage_variants": ["bf16_h128_topk128x_row_first"]} + {"index": 0, "stage_variants": ["bf16_h128_prefill_v42_snake"]} ], "signature": { "scalar_dtypes": { @@ -1422,27 +1435,18 @@ "workspace_dtypes": {}, "workspace_keys": [], }, - "sources": ["cake_dsv4_family_585100536b80a2cec8ba_binding.cu"], + "sources": ["cake_dsv4_family_42ae3e93a01e7a3bf20d_binding.cu"], }, - "cake_dsv4_family_956aa0f74f4a82042f2c": { + "cake_dsv4_family_585100536b80a2cec8ba": { "arch": "sm_103a", "compile_flags": ["-std=c++17"], - "dependencies": [ - "bf16_h128_topk4x_v52", - "split_reduce_cake_dsv4_230143c89d42c6c89387", - ], + "dependencies": ["bf16_h128_topk128x_row_first"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "53306debf3bddc3348c65222e0ed7b69fc316a69af2ce6d7314061f98104b233", + "identity": "72b48d77a88ee5d7e37313209035e21d2037d2b74d4793d7660367a6c9003750", "route_entry": "route_of", "routes": [ - { - "index": 0, - "stage_variants": [ - "bf16_h128_topk4x_v52", - "split_reduce_cake_dsv4_230143c89d42c6c89387", - ], - } + {"index": 0, "stage_variants": ["bf16_h128_topk128x_row_first"]} ], "signature": { "scalar_dtypes": { @@ -1450,7 +1454,6 @@ "has_sinks": "int32", "num_heads": "int32", "num_query_tokens": "int32", - "num_split": "int32", "sparse_topk": "int32", "sparse_topk_lens_offset": "int32", "swa_index_stride": "int32", @@ -1465,7 +1468,6 @@ "sparse_topk", "has_sinks", "total_work_items", - "num_split", ], "tensor_dtypes": { "O": "bfloat16", @@ -1475,7 +1477,6 @@ "bmm2_scale": "float32", "compressed_KV_cache": "bfloat16", "compressed_indices": "int32", - "partial_O": "bfloat16", "partial_lse": "float32", "sinks": "float32", "sparse_topk_lens": "int32", @@ -1485,7 +1486,6 @@ "Q", "SWA_cache", "compressed_KV_cache", - "partial_O", "partial_lse", "swa_indices", "compressed_indices", @@ -1498,9 +1498,9 @@ "workspace_dtypes": {}, "workspace_keys": [], }, - "sources": ["cake_dsv4_family_956aa0f74f4a82042f2c_binding.cu"], + "sources": ["cake_dsv4_family_585100536b80a2cec8ba_binding.cu"], }, - "cake_dsv4_family_ec44ddd94e3be88cfc83": { + "cake_dsv4_family_ebb85adb872529ca9e83": { "arch": "sm_103a", "compile_flags": ["-std=c++17"], "dependencies": [ @@ -1509,7 +1509,7 @@ ], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "7c7b9440d3cf7f9ad7a4726319dfcba580a58bf74449e03f048096044be3bee5", + "identity": "42d62a52aae085402e0ffb43549c81c591c5c7bd90b08b4d71ecdf5e0c44fb1f", "route_entry": "route_of", "routes": [ { @@ -1574,7 +1574,7 @@ "workspace_dtypes": {}, "workspace_keys": [], }, - "sources": ["cake_dsv4_family_ec44ddd94e3be88cfc83_binding.cu"], + "sources": ["cake_dsv4_family_ebb85adb872529ca9e83_binding.cu"], }, }, "variants": { @@ -1609,7 +1609,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "3c18335a89ae686e86a9d8f7e02bedfceda722cf8eca63b648d0bf1871fe4cf5", + "identity": "89495d4d9f2fe8ee1b03873aa440cff30e4d413047e7aae57b372c68b430df7f", "sources": [ "cake_dsv4_a40bd49fee9655d8b1c6_kernel.cu", "cake_dsv4_a40bd49fee9655d8b1c6_binding.cu", @@ -1647,7 +1647,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "4ca772248ca31846c587fafbac94c2da2efdf58a844e78832268403caca66abb", + "identity": "525777264ca547d1e7d9e732ff609890eda28ca396da31184dce78c7c23253bb", "sources": [ "cake_dsv4_d338c32a69477808a1ef_kernel.cu", "cake_dsv4_d338c32a69477808a1ef_binding.cu", @@ -1680,10 +1680,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "55e52ad1bad7c5f6d9f0287db61de94bd0c95041a6cb2933c495784d302f9d68", + "identity": "669f6c1916289193a683ac74b7d3c5967c0a8a5ea1e2e652f0f99a9bdcbb40cd", "sources": [ - "cake_dsv4_f89bc6561b38462e9538_kernel.cu", - "cake_dsv4_f89bc6561b38462e9538_binding.cu", + "cake_dsv4_1bad879f597b2e2adf3e_kernel.cu", + "cake_dsv4_1bad879f597b2e2adf3e_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -1718,7 +1718,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "8fdcc0ef01be0f9758a29d5556b2f7bae6864aa2ed8a79e56200ea722bb60c78", + "identity": "f4a7ad606f7e1cbc294f4021c2f443da82e6735e6118ec12af544bb948f798bb", "sources": [ "cake_dsv4_af6f9d240794a942a124_kernel.cu", "cake_dsv4_af6f9d240794a942a124_binding.cu", @@ -1733,7 +1733,7 @@ ["tma_buffer", "tmap_compressed_k"], ["tma_buffer", "tmap_swa_v"], ["tma_buffer", "tmap_compressed_v"], - ["buffer", "O"], + ["tma_buffer", "O"], ["buffer", "partial_lse"], ["buffer", "swa_indices"], ["buffer", "compressed_indices"], @@ -1756,10 +1756,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "5bed0c801c9f03c5a4157722eca0768a8457156f1aab8d80ee1d002b00029ca5", + "identity": "6008a89e4a7f9470bbd0b258bad5d335dc84a0027763f5c00745d93787c5efdf", "sources": [ - "cake_dsv4_4c41fb8347f137d4d46f_kernel.cu", - "cake_dsv4_4c41fb8347f137d4d46f_binding.cu", + "cake_dsv4_064850af5858dd03c104_kernel.cu", + "cake_dsv4_064850af5858dd03c104_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -1771,7 +1771,7 @@ ["tma_buffer", "tmap_compressed_k"], ["tma_buffer", "tmap_swa_v"], ["tma_buffer", "tmap_compressed_v"], - ["buffer", "O"], + ["tma_buffer", "O"], ["buffer", "partial_lse"], ["buffer", "swa_indices"], ["buffer", "compressed_indices"], @@ -1794,10 +1794,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "1ca21a0f7b4c6edfddb5340b989d3cc2964c772e7ecd16fc74bcea8088fd5eec", + "identity": "2adf6df71253d0032819574f14607f44a8c86dae4c489e3d77dd50806316dcf6", "sources": [ - "cake_dsv4_125332c2066d89bd7610_kernel.cu", - "cake_dsv4_125332c2066d89bd7610_binding.cu", + "cake_dsv4_4ab0487d5fde70d4569e_kernel.cu", + "cake_dsv4_4ab0487d5fde70d4569e_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -1826,7 +1826,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "181de2841cecdd15e69e910e6afec211743131dcf5fb7fed2023027d0570b314", + "identity": "19423e1c286750ba3f6842d6809313687520c8ae382c1e5b3f14297c75155a96", "sources": [ "cake_dsv4_ae27071dabfc85fb47c1_kernel.cu", "cake_dsv4_ae27071dabfc85fb47c1_binding.cu", @@ -1864,7 +1864,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "456fc2e0220f173e148d85b2fd7b3d73c245f83a64a29847de39a0a29b8a183b", + "identity": "2780b283e3889a568e1f4bce3009a3a3d71b4c076930bc8481381308d5379d67", "sources": [ "cake_dsv4_6b7225f9e37a48354992_kernel.cu", "cake_dsv4_6b7225f9e37a48354992_binding.cu", @@ -1899,7 +1899,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "9f675bcf369d90feaf54659043953dd0734c40b14e65b5cc2e8db3ad756b6b3f", + "identity": "3c9bff552cbbe14c33c1d07469dad72b836d2c3816b91dcf6150b1affc5f9aed", "sources": [ "cake_dsv4_fb5ca25f5af355ef89b5_kernel.cu", "cake_dsv4_fb5ca25f5af355ef89b5_binding.cu", @@ -1921,7 +1921,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "e88d9e3ab9ee0f46c4d50a03e328e447396b8ab121c01f134d447065b10394f0", + "identity": "8b08b4fd9e1a83550104302232433821bb3a83d4ce9762fcfee3fd540bc2af59", "sources": [ "cake_dsv4_0590fb03afd3391febef_kernel.cu", "cake_dsv4_0590fb03afd3391febef_binding.cu", @@ -1960,10 +1960,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "6991e3e7c4210570953f9cdc15c6293f15bf9a0e1798b739abee8df34c6d3551", + "identity": "7cbab296d40363c89a8116b277938459a7a15cbd3ea2bad99ac36d24f0f92166", "sources": [ - "cake_dsv4_c967e7efe220160e28d3_kernel.cu", - "cake_dsv4_c967e7efe220160e28d3_binding.cu", + "cake_dsv4_1ff52fece183b851b441_kernel.cu", + "cake_dsv4_1ff52fece183b851b441_binding.cu", ], "tma_workspace_bytes": 384, }, @@ -1987,7 +1987,6 @@ ["parameter", "sparse_topk"], ["parameter", "num_query_tokens"], ["parameter", "has_sinks"], - ["workspace", "tma_descriptor_workspace"], ["grid", "grid_x"], ["grid", "grid_y"], ["grid", "grid_z"], @@ -1995,12 +1994,12 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "aae71071ae4aa66b4bdd00a852c061efda61475dc4a69cdf4b2bd2705f4a4661", + "identity": "6f16d1265dc27464472cc2464c8faa84a788e83345d6c5fa698c57897dcad362", "sources": [ - "cake_dsv4_3e5a644d61420c9c7ea1_kernel.cu", - "cake_dsv4_3e5a644d61420c9c7ea1_binding.cu", + "cake_dsv4_1e3cf4361de31c58b69b_kernel.cu", + "cake_dsv4_1e3cf4361de31c58b69b_binding.cu", ], - "tma_workspace_bytes": 384, + "tma_workspace_bytes": 0, }, "bf16_h8_h16_source_exact": { "arch": "sm_103a", @@ -2032,10 +2031,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "ad819a642ec063a95836b08029677d93d1c6516c4b096e3106f9032257a12bd0", + "identity": "a08e408c6cdc5bbb63728eae56cb3c556c8b15513eb1a23d68d25712751ba903", "sources": [ - "cake_dsv4_d23652aae3370c30d978_kernel.cu", - "cake_dsv4_d23652aae3370c30d978_binding.cu", + "cake_dsv4_e48c5e0a89a41eb1e57a_kernel.cu", + "cake_dsv4_e48c5e0a89a41eb1e57a_binding.cu", ], "tma_workspace_bytes": 384, }, @@ -2064,7 +2063,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "12c3d8e97fc54d26fe53f365da8b548e3cfe388c84a0101316c88a55a840ab0e", + "identity": "a837022a12132a14e91375f54c3ff7fcabea99a3cae8d5583d13070284df984f", "sources": [ "cake_dsv4_fdbe4c4737d7f6dd1d9e_kernel.cu", "cake_dsv4_fdbe4c4737d7f6dd1d9e_binding.cu", @@ -2096,10 +2095,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "5b806976ca4dcb0b14105e2f1b06df0630452af2283a61256c8df96b0ec8b7e4", + "identity": "9503eaa1245660fa320fbbd82dfec0253c49d57bdd1b2208f3dcc687d8cb1587", "sources": [ - "cake_dsv4_9cf30aacb43b71034660_kernel.cu", - "cake_dsv4_9cf30aacb43b71034660_binding.cu", + "cake_dsv4_d9754db4fe394a927e67_kernel.cu", + "cake_dsv4_d9754db4fe394a927e67_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -2140,7 +2139,7 @@ ], "entry": "run", "host_linkage_flags": [], - "identity": "72c980c50469dd2c8c22342b2384ccd6d6770535e77343c9aea15c9f1f99860f", + "identity": "8a292d1b18d0e15c944d598de09148bc80b5d5bd09a855539482791253d4ff15", "sources": [ "cake_dsv4_85ef6cd9a950ca02c484_kernel.cu", "cake_dsv4_85ef6cd9a950ca02c484_binding.cu", @@ -2181,7 +2180,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "4cbb2c362b859ca355d32bb1c802aabf58d88fd8f9ee36140273165754e8820c", + "identity": "cfef2a7086ba2562eeebc50d196c8486ab4522e3b059ce3a0952915c03d8dad0", "sources": [ "cake_dsv4_32e01d0e8961dfa26bf3_kernel.cu", "cake_dsv4_32e01d0e8961dfa26bf3_binding.cu", @@ -2221,7 +2220,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "9499459d9e38d660258217d59ae733e72c9421c68b93e53534379a4353cf4145", + "identity": "587da4e34f457b41b1b540e71a2a841fe03c0a5e0f056e846da82eeafe874185", "sources": [ "cake_dsv4_bd3f29f66e3196dc9451_kernel.cu", "cake_dsv4_bd3f29f66e3196dc9451_binding.cu", @@ -2256,10 +2255,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "a0c09ef4273aa62ee82bc9f0f221af205a14fbbc44509bbdf674ff321f5f66b6", + "identity": "910f0fea7bd4f6be5f5b62ddb6ca05b4365799b0e8fe122b4a02018ee0da8ae4", "sources": [ - "cake_dsv4_f55f69afe4ac3ec96817_kernel.cu", - "cake_dsv4_f55f69afe4ac3ec96817_binding.cu", + "cake_dsv4_0a925b9bcb3044777c76_kernel.cu", + "cake_dsv4_0a925b9bcb3044777c76_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -2293,10 +2292,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "d29ebd3f5584b9189fb165159f8117f93dcf919fd9b35bf6ddfccdbbb3f05593", + "identity": "30120c1636cb83664be716d2b2df67d2f2dd5f8faaa9b7b2e05237fa9a402557", "sources": [ - "cake_dsv4_7a4b9556b3c84ebff6ce_kernel.cu", - "cake_dsv4_7a4b9556b3c84ebff6ce_binding.cu", + "cake_dsv4_c1d444bba4cdb8273c7c_kernel.cu", + "cake_dsv4_c1d444bba4cdb8273c7c_binding.cu", ], "tma_workspace_bytes": 384, }, @@ -2329,10 +2328,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "7ff0f4e03616e8a841b5726c9bae67e2ec29b56d21c7da962980cf1d0fba8d57", + "identity": "22eff1cf99b136db0c474ddb719cbca1a10a199b8d7a1639c83182dad1c6ac15", "sources": [ - "cake_dsv4_9e1188b192c9c806d7c6_kernel.cu", - "cake_dsv4_9e1188b192c9c806d7c6_binding.cu", + "cake_dsv4_983db4c490ceed3c455a_kernel.cu", + "cake_dsv4_983db4c490ceed3c455a_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -2365,10 +2364,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "936e66b0b3c4a12668dcca14881ec70097823c247ca505c455fd0db5da4c99b5", + "identity": "4e9cdeed8ed71d70dc7b1d02f7a417e006a6ff14789f8af8779658b85fa60c5d", "sources": [ - "cake_dsv4_c20607c57cc11784015e_kernel.cu", - "cake_dsv4_c20607c57cc11784015e_binding.cu", + "cake_dsv4_98db1f55c8c129e16005_kernel.cu", + "cake_dsv4_98db1f55c8c129e16005_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -2402,10 +2401,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "61781095b5e9377262e277a85c730dcd0bb3588b3633cbe248765b64980e8e23", + "identity": "325bbfbb67496e600ba75e025966b45f9df4db915c0e75e7d27f1438a6ad5036", "sources": [ - "cake_dsv4_a2be51a1fd34ec0d9a3b_kernel.cu", - "cake_dsv4_a2be51a1fd34ec0d9a3b_binding.cu", + "cake_dsv4_868bb2291a1464f14e05_kernel.cu", + "cake_dsv4_868bb2291a1464f14e05_binding.cu", ], "tma_workspace_bytes": 384, }, @@ -2424,7 +2423,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "2a8255d11ab33428fbaa167fb8b2a10bf0267969da6ddfb81d9974cd7b62296f", + "identity": "a1b1f15c277a8778c33e5288abfbbb1dbacf1e2c6625e3e1fe39277c65023443", "sources": [ "cake_dsv4_230143c89d42c6c89387_kernel.cu", "cake_dsv4_230143c89d42c6c89387_binding.cu", @@ -2446,7 +2445,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "3f909a9786eea6a97a37b095198212ee64c0fddc5e3d9e362fa48443f946c528", + "identity": "ebb9afd30b8660b2b019a33b9d32b1468d7dd2f86458522b4b720506188ed321", "sources": [ "cake_dsv4_397f817e2931a6cc7c92_kernel.cu", "cake_dsv4_397f817e2931a6cc7c92_binding.cu",