Skip to content

perf(cake_sparse_mla): round-7 DSv4 sparse-MLA Cake programs for SM100/SM103 (numerics-unlocked levers; stacked on #5700) #5728

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
yyihuang wants to merge 3 commits into from
+4,994 −3,054
Closed

perf(cake_sparse_mla): round-7 DSv4 sparse-MLA Cake programs for SM100/SM103 (numerics-unlocked levers; stacked on #5700) #5728

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
07fed14
perf(cake_sparse_mla): regenerate sm_103a programs from Cake 56f72c78…
averyhNV Sep 30, 2026
1042f9c
perf(cake_sparse_mla): regenerate sm_100a programs from Cake 56f72c78…
averyhNV Sep 30, 2026
5ee5417
perf(cake_sparse_mla): drop the round-6 program sources superseded by…
averyhNV Sep 30, 2026
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
18 changes: 9 additions & 9 deletions ...cake_dsv4_3b3789d6356f15f9cbe3_binding.cu → ...cake_dsv4_09e14f1d9cbd6739db61_binding.cu
View file
Open in desktop
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_3b3789d6356f15f9cbe3'.
// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_09e14f1d9cbd6739db61'.
#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_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;
Expand Down Expand Up @@ -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<const void*>(kernel_cake_dsv4_3b3789d6356f15f9cbe3),
reinterpret_cast<const void*>(kernel_cake_dsv4_09e14f1d9cbd6739db61),
cudaFuncAttributeMaxDynamicSharedMemorySize,
198144);
}
Expand All @@ -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);
Expand All @@ -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<const void*>(kernel_cake_dsv4_3b3789d6356f15f9cbe3), kargs);
&config, reinterpret_cast<const void*>(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);
20 changes: 10 additions & 10 deletions .../cake_dsv4_3b3789d6356f15f9cbe3_kernel.cu → .../cake_dsv4_09e14f1d9cbd6739db61_kernel.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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) {
Expand Down
18 changes: 9 additions & 9 deletions ...cake_dsv4_57e7b93def70f2be356f_binding.cu → ...cake_dsv4_0c8ab955ac7b62e7bfee_binding.cu
View file
Open in desktop
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_0c8ab955ac7b62e7bfee'.
#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_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;
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_0c8ab955ac7b62e7bfee),
cudaFuncAttributeMaxDynamicSharedMemorySize,
198528);
}
Expand All @@ -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);
Expand All @@ -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<const void*>(kernel_cake_dsv4_57e7b93def70f2be356f), kargs);
&config, reinterpret_cast<const void*>(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);
Loading
Loading