Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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_0cd15705e6bea14f3eb6'.
#include <cuda.h>
#include <cuda_bf16.h>
#include <cuda_fp16.h>
Expand All @@ -31,10 +31,10 @@
#include <unordered_map>
#include <vector>

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_0cd15705e6bea14f3eb6(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_4d9d3a7fb10cffa8 {

using tvm::ffi::Optional;
using tvm::ffi::TensorView;
Expand Down Expand Up @@ -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<const void*>(kernel_cake_dsv4_57e7b93def70f2be356f),
reinterpret_cast<const void*>(kernel_cake_dsv4_0cd15705e6bea14f3eb6),
cudaFuncAttributeMaxDynamicSharedMemorySize,
198528);
}
Expand All @@ -367,23 +367,18 @@ 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_0cd15705e6bea14f3eb6 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 = 198528u;
config.stream = stream;
cudaError_t launch_status = cudaLaunchKernelExC(
&config, reinterpret_cast<const void*>(kernel_cake_dsv4_57e7b93def70f2be356f), kargs);
cudaError_t launch_status = cudaLaunchKernel(
reinterpret_cast<const void*>(kernel_cake_dsv4_0cd15705e6bea14f3eb6), grid, block, kargs, 198528u, stream);
TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError)
<< "cudaLaunchKernelExC for kernel_cake_dsv4_57e7b93def70f2be356f failed: "
<< "cudaLaunchKernel for kernel_cake_dsv4_0cd15705e6bea14f3eb6 failed: "
<< cudaGetErrorString(launch_status);
}

} // namespace cake_host_shim_16cf3a7d0d01e13d
} // namespace cake_host_shim_4d9d3a7fb10cffa8

TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_16cf3a7d0d01e13d::Run);
TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_4d9d3a7fb10cffa8::Run);

Large diffs are not rendered by default.

Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
*/

// Generated by the Cake source exporter — do not edit.
// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_bf2bffa3152d5dce396f'.
// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_0d2de23ce48a72102ba5'.
#include <cuda.h>
#include <cuda_bf16.h>
#include <cuda_fp16.h>
Expand All @@ -31,10 +31,10 @@
#include <unordered_map>
#include <vector>

extern "C" __global__ void kernel_cake_dsv4_bf2bffa3152d5dce396f(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_0d2de23ce48a72102ba5(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, int num_query_tokens);


namespace cake_host_shim_0100aa295c27e44a {
namespace cake_host_shim_cdf3191b47d44f36 {

using tvm::ffi::Optional;
using tvm::ffi::TensorView;
Expand Down Expand Up @@ -304,7 +304,7 @@ inline CUtensorMap EncodeTma_tmap_compressed_kv(const TensorView& t) {
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_seq_lens, TensorView arg_cum_seq_lens_q, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, 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_max_q_len, int64_t arg_batch_size, int64_t arg_has_sinks, int64_t grid_x, int64_t grid_y, int64_t grid_z) {
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_seq_lens, TensorView arg_cum_seq_lens_q, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, 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_max_q_len, int64_t arg_batch_size, int64_t arg_has_sinks, int64_t arg_num_query_tokens, 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");
Expand Down Expand Up @@ -362,6 +362,9 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_
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]";
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]";
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");
Expand Down Expand Up @@ -397,7 +400,8 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_
int32_t v_max_q_len = (int32_t)arg_max_q_len;
int32_t v_batch_size = (int32_t)arg_batch_size;
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_seq_lens, &p_cum_seq_lens_q, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_max_q_len, &v_batch_size, &v_has_sinks};
int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens;
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_seq_lens, &p_cum_seq_lens_q, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_max_q_len, &v_batch_size, &v_has_sinks, &v_num_query_tokens};

static std::mutex smem_mu;
static auto* smem_status_by_device =
Expand All @@ -414,7 +418,7 @@ void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_
}
if (smem_status == cudaSuccess) {
smem_status = cudaFuncSetAttribute(
reinterpret_cast<const void*>(kernel_cake_dsv4_bf2bffa3152d5dce396f),
reinterpret_cast<const void*>(kernel_cake_dsv4_0d2de23ce48a72102ba5),
cudaFuncAttributeMaxDynamicSharedMemorySize,
218752);
}
Expand All @@ -424,23 +428,18 @@ 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_bf2bffa3152d5dce396f failed: "
<< "cudaFuncSetAttribute for kernel_cake_dsv4_0d2de23ce48a72102ba5 failed: "
<< cudaGetErrorString(smem_status);
dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z);
dim3 block(640u, 1u, 1u);

cudaLaunchConfig_t config{};
config.gridDim = grid;
config.blockDim = block;
config.dynamicSmemBytes = 218752u;
config.stream = stream;
cudaError_t launch_status = cudaLaunchKernelExC(
&config, reinterpret_cast<const void*>(kernel_cake_dsv4_bf2bffa3152d5dce396f), kargs);
cudaError_t launch_status = cudaLaunchKernel(
reinterpret_cast<const void*>(kernel_cake_dsv4_0d2de23ce48a72102ba5), grid, block, kargs, 218752u, stream);
TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError)
<< "cudaLaunchKernelExC for kernel_cake_dsv4_bf2bffa3152d5dce396f failed: "
<< "cudaLaunchKernel for kernel_cake_dsv4_0d2de23ce48a72102ba5 failed: "
<< cudaGetErrorString(launch_status);
}

} // namespace cake_host_shim_0100aa295c27e44a
} // namespace cake_host_shim_cdf3191b47d44f36

TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_0100aa295c27e44a::Run);
TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_cdf3191b47d44f36::Run);
Original file line number Diff line number Diff line change
Expand Up @@ -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_bf2bffa3152d5dce396f(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_0d2de23ce48a72102ba5(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, int num_query_tokens)
{
const int tid = threadIdx.x;
const int warp = make_warp_uniform(tid / 32);
Expand Down Expand Up @@ -2189,7 +2189,39 @@ kernel_cake_dsv4_bf2bffa3152d5dce396f(const __grid_constant__ CUtensorMap tmap_q
int q_len_3 = q_end_3 - q_begin_3;
int query_idx_3 = (unsigned int)q_begin_3 + work_x_3;
int active_topk_3 = 0;
int pg0_rows[4];
int pg1_rows[4];
if (work_x_3 < (unsigned int)q_len_3) {
int pg_lane = lane * 4;
if (pg_lane < sparse_topk) {
{
const int4* _ivptr_0 = reinterpret_cast<const int4*>(swa_indices + (query_idx_3 * swa_index_stride) + pg_lane);
int4 _ivld_0;
_ivld_0 = *_ivptr_0;
pg0_rows[0 + 0] = _ivld_0.x;
pg0_rows[0 + 1] = _ivld_0.y;
pg0_rows[0 + 2] = _ivld_0.z;
pg0_rows[0 + 3] = _ivld_0.w;
}
}
if (128 + pg_lane < sparse_topk) {
{
const int4* _ivptr_1 = reinterpret_cast<const int4*>(compressed_indices + (query_idx_3 * compressed_index_stride) + pg_lane);
int4 _ivld_1;
_ivld_1 = *_ivptr_1;
pg1_rows[0 + 0] = _ivld_1.x;
pg1_rows[0 + 1] = _ivld_1.y;
pg1_rows[0 + 2] = _ivld_1.z;
pg1_rows[0 + 3] = _ivld_1.w;
}
}
int pf_lane = lane * 4;
#pragma unroll
for (int pf_tile = 0; pf_tile < 2; pf_tile++) {
int pf_col = pf_tile * 128 + pf_lane;
int* pf_src = ((pf_col < 128) ? (swa_indices + (query_idx_3 * swa_index_stride + pf_col)) : (compressed_indices + (query_idx_3 * compressed_index_stride + (pf_col - 128))));
if (pf_col < sparse_topk) { asm volatile("prefetch.global.L2 [%0];" :: "l"((uint64_t)(pf_src))); }
}
int _max_0 = ((sparse_topk_lens[query_idx_3] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_3] + sparse_topk_lens_offset) : (0));
int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk));
active_topk_3 = _min_0;
Expand All @@ -2200,18 +2232,55 @@ kernel_cake_dsv4_bf2bffa3152d5dce396f(const __grid_constant__ CUtensorMap tmap_q
int last_aligned = _max_1;
int* swa_row = swa_indices + (query_idx_3 * swa_index_stride);
int* compressed_row = compressed_indices + (query_idx_3 * compressed_index_stride);
if (num_pairs > 0) {
int pg_base = lane * 4;
#pragma unroll
for (int duplicate = 0; duplicate < 2; duplicate++) {
mbarrier_wait(page_empty_addr + (page_stage) * 8, page_phase);
#pragma unroll
for (int pg_tile = 0; pg_tile < 2; pg_tile++) {
int pg_page = pg_tile * 128 + pg_base;
int pg_safe = ((pg_page <= last_aligned) ? pg_page : last_aligned);
int pg_src = (pg_safe & 127) >> 2;
int _shfl_0 = __shfl_sync(0xFFFFFFFF, pg0_rows[0], pg_src);
int pg_a0 = _shfl_0;
int _shfl_1 = __shfl_sync(0xFFFFFFFF, pg0_rows[1], pg_src);
int pg_a1 = _shfl_1;
int _shfl_2 = __shfl_sync(0xFFFFFFFF, pg0_rows[2], pg_src);
int pg_a2 = _shfl_2;
int _shfl_3 = __shfl_sync(0xFFFFFFFF, pg0_rows[3], pg_src);
int pg_a3 = _shfl_3;
int _shfl_4 = __shfl_sync(0xFFFFFFFF, pg1_rows[0], pg_src);
int pg_b0 = _shfl_4;
int _shfl_5 = __shfl_sync(0xFFFFFFFF, pg1_rows[1], pg_src);
int pg_b1 = _shfl_5;
int _shfl_6 = __shfl_sync(0xFFFFFFFF, pg1_rows[2], pg_src);
int pg_b2 = _shfl_6;
int _shfl_7 = __shfl_sync(0xFFFFFFFF, pg1_rows[3], pg_src);
int pg_b3 = _shfl_7;
int pg_v0 = ((pg_safe >= 128) ? pg_b0 : pg_a0);
int pg_v1 = ((pg_safe >= 128) ? pg_b1 : pg_a1);
int pg_v2 = ((pg_safe >= 128) ? pg_b2 : pg_a2);
int pg_v3 = ((pg_safe >= 128) ? pg_b3 : pg_a3);
asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_page_offsets_addr + (page_stage * 1024 + (unsigned int)(pg_page * 4))), "r"(pg_v0), "r"(pg_v1), "r"(pg_v2), "r"(pg_v3) : "memory");
}
mbarrier_arrive(page_full_addr + (page_stage) * 8);
page_stage += 1;
if (page_stage == 6) { page_stage = 0; page_phase ^= 1; }
}
}
#pragma unroll 1
for (int pair_2 = 0; pair_2 < num_pairs; pair_2++) {
for (int pair_2 = 1; pair_2 < num_pairs; pair_2++) {
int lane_base = lane * 4;
#pragma unroll
for (int duplicate = 0; duplicate < 2; duplicate++) {
for (int duplicate_1 = 0; duplicate_1 < 2; duplicate_1++) {
mbarrier_wait(page_empty_addr + (page_stage) * 8, page_phase);
int page0 = pair_2 * 2 * 128 + lane_base;
int page1 = page0 + 128;
int safe0 = ((page0 <= last_aligned) ? page0 : last_aligned);
int safe1 = ((page1 <= last_aligned) ? page1 : last_aligned);
int* src0 = ((safe0 < 128) ? swa_row + safe0 : compressed_row + (safe0 - 128));
int* src1 = ((safe1 < 128) ? swa_row + safe1 : compressed_row + (safe1 - 128));
int* src0 = ((safe0 < 128) ? (swa_row + safe0) : (compressed_row + (safe0 - 128)));
int* src1 = ((safe1 < 128) ? (swa_row + safe1) : (compressed_row + (safe1 - 128)));
asm volatile("cp.async.ca.shared::cta.global [%0], [%1], 16;"
:: "r"(smem_page_offsets_addr + page_stage * 1024 + (unsigned int)(lane_base * 4)), "l"(src0));
asm volatile("cp.async.ca.shared::cta.global [%0], [%1], 16;"
Expand Down
Loading
Loading