Skip to content

feat(cake_selective_state_update): promote the headdim-64 / dstate-128 single-token decode row (Nemotron-H, granite-4.0-h) #6040

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

Merged
yyihuang merged 1 commit into from
Oct 4, 2026
Merged

feat(cake_selective_state_update): promote the headdim-64 / dstate-128 single-token decode row (Nemotron-H, granite-4.0-h) #6040

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
207 changes: 207 additions & 0 deletions csrc/cake_selective_state_update/cake_selective_state_update_4570410fcf7e2771ab30_binding.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,207 @@
/*
* Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
* 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_selective_state_update_4570410fcf7e2771ab30'.
#include <cuda.h>
#include <cuda_bf16.h>
#include <cuda_fp16.h>
#include <cuda_runtime.h>

#include <cstdint>

#include "tvm_ffi_utils.h"

extern "C" __global__ void kernel_cake_selective_state_update_4570410fcf7e2771ab30(
float* __restrict__ state, __nv_bfloat16* __restrict__ x, unsigned long long dt_addr,
unsigned long long a_addr, __nv_bfloat16* __restrict__ B, __nv_bfloat16* __restrict__ C,
unsigned long long d_addr, __nv_bfloat16* __restrict__ z, unsigned long long dt_bias_addr,
__nv_bfloat16* __restrict__ output, unsigned long long state_batch_indices_addr,
unsigned long long dst_state_batch_indices_addr, int nheads, int ngroups, int dim_tiles,
unsigned long long state_stride_slot, long long x_batch_stride, long long b_batch_stride,
long long c_batch_stride, long long out_batch_stride, long long dt_batch_stride,
long long dt_head_stride, long long a_head_stride, long long d_head_stride,
long long dt_bias_head_stride, int dt_softplus, int has_z, int disable_state_update,
long long pad_slot_id);

namespace cake_host_shim_825081a78f52121f {

using tvm::ffi::TensorView;

void Run(TensorView arg_state, TensorView arg_x, int64_t arg_dt_addr, int64_t arg_a_addr,
TensorView arg_B, TensorView arg_C, int64_t arg_d_addr, TensorView arg_z,
int64_t arg_dt_bias_addr, TensorView arg_output, int64_t arg_state_batch_indices_addr,
int64_t arg_dst_state_batch_indices_addr, int64_t arg_nheads, int64_t arg_ngroups,
int64_t arg_dim_tiles, int64_t arg_state_stride_slot, int64_t arg_x_batch_stride,
int64_t arg_b_batch_stride, int64_t arg_c_batch_stride, int64_t arg_out_batch_stride,
int64_t arg_dt_batch_stride, int64_t arg_dt_head_stride, int64_t arg_a_head_stride,
int64_t arg_d_head_stride, int64_t arg_dt_bias_head_stride, int64_t arg_dt_softplus,
int64_t arg_has_z, int64_t arg_disable_state_update, int64_t arg_pad_slot_id,
int64_t grid_x, int64_t grid_y, int64_t grid_z) {
DLDevice dev = arg_state.device();
tvm::ffi::CUDADeviceGuard device_guard(dev.device_id);
check_cuda_tensor(arg_state, "state");
check_dtype(arg_state, DLDataType{kDLFloat, 32, 1}, "state");
check_contiguous(arg_state, "state");
check_cuda_tensor(arg_x, "x");
check_dtype(arg_x, DLDataType{kDLBfloat, 16, 1}, "x");
check_contiguous(arg_x, "x");
TVM_FFI_CHECK(arg_dt_addr >= 0LL && arg_dt_addr <= 9223372036854775807LL, ValueError)
<< "scalar 'dt_addr' value " << arg_dt_addr
<< " is outside u64 range [0, 9223372036854775807]";
TVM_FFI_CHECK(arg_a_addr >= 0LL && arg_a_addr <= 9223372036854775807LL, ValueError)
<< "scalar 'a_addr' value " << arg_a_addr << " is outside u64 range [0, 9223372036854775807]";
check_cuda_tensor(arg_B, "B");
check_dtype(arg_B, DLDataType{kDLBfloat, 16, 1}, "B");
check_contiguous(arg_B, "B");
check_cuda_tensor(arg_C, "C");
check_dtype(arg_C, DLDataType{kDLBfloat, 16, 1}, "C");
check_contiguous(arg_C, "C");
TVM_FFI_CHECK(arg_d_addr >= 0LL && arg_d_addr <= 9223372036854775807LL, ValueError)
<< "scalar 'd_addr' value " << arg_d_addr << " is outside u64 range [0, 9223372036854775807]";
check_cuda_tensor(arg_z, "z");
check_dtype(arg_z, DLDataType{kDLBfloat, 16, 1}, "z");
check_contiguous(arg_z, "z");
TVM_FFI_CHECK(arg_dt_bias_addr >= 0LL && arg_dt_bias_addr <= 9223372036854775807LL, ValueError)
<< "scalar 'dt_bias_addr' value " << arg_dt_bias_addr
<< " is outside u64 range [0, 9223372036854775807]";
check_cuda_tensor(arg_output, "output");
check_dtype(arg_output, DLDataType{kDLBfloat, 16, 1}, "output");
check_contiguous(arg_output, "output");
TVM_FFI_CHECK(
arg_state_batch_indices_addr >= 0LL && arg_state_batch_indices_addr <= 9223372036854775807LL,
ValueError)
<< "scalar 'state_batch_indices_addr' value " << arg_state_batch_indices_addr
<< " is outside u64 range [0, 9223372036854775807]";
TVM_FFI_CHECK(arg_dst_state_batch_indices_addr >= 0LL &&
arg_dst_state_batch_indices_addr <= 9223372036854775807LL,
ValueError)
<< "scalar 'dst_state_batch_indices_addr' value " << arg_dst_state_batch_indices_addr
<< " is outside u64 range [0, 9223372036854775807]";
TVM_FFI_CHECK(arg_nheads >= -2147483648LL && arg_nheads <= 2147483647LL, ValueError)
<< "scalar 'nheads' value " << arg_nheads
<< " is outside i32 range [-2147483648, 2147483647]";
TVM_FFI_CHECK(arg_ngroups >= -2147483648LL && arg_ngroups <= 2147483647LL, ValueError)
<< "scalar 'ngroups' value " << arg_ngroups
<< " is outside i32 range [-2147483648, 2147483647]";
TVM_FFI_CHECK(arg_dim_tiles >= -2147483648LL && arg_dim_tiles <= 2147483647LL, ValueError)
<< "scalar 'dim_tiles' value " << arg_dim_tiles
<< " is outside i32 range [-2147483648, 2147483647]";
TVM_FFI_CHECK(arg_state_stride_slot >= 0LL && arg_state_stride_slot <= 9223372036854775807LL,
ValueError)
<< "scalar 'state_stride_slot' value " << arg_state_stride_slot
<< " is outside u64 range [0, 9223372036854775807]";
TVM_FFI_CHECK(arg_dt_softplus >= -2147483648LL && arg_dt_softplus <= 2147483647LL, ValueError)
<< "scalar 'dt_softplus' value " << arg_dt_softplus
<< " is outside i32 range [-2147483648, 2147483647]";
TVM_FFI_CHECK(arg_has_z >= -2147483648LL && arg_has_z <= 2147483647LL, ValueError)
<< "scalar 'has_z' value " << arg_has_z << " is outside i32 range [-2147483648, 2147483647]";
TVM_FFI_CHECK(
arg_disable_state_update >= -2147483648LL && arg_disable_state_update <= 2147483647LL,
ValueError)
<< "scalar 'disable_state_update' value " << arg_disable_state_update
<< " is outside i32 range [-2147483648, 2147483647]";
check_same_device(arg_x, arg_state, "x", "state");
check_same_device(arg_B, arg_state, "B", "state");
check_same_device(arg_C, arg_state, "C", "state");
check_same_device(arg_z, arg_state, "z", "state");
check_same_device(arg_output, arg_state, "output", "state");
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);
float* p_state = static_cast<float*>(arg_state.data_ptr());
__nv_bfloat16* p_x = static_cast<__nv_bfloat16*>(arg_x.data_ptr());
uint64_t v_dt_addr = (uint64_t)arg_dt_addr;
uint64_t v_a_addr = (uint64_t)arg_a_addr;
__nv_bfloat16* p_B = static_cast<__nv_bfloat16*>(arg_B.data_ptr());
__nv_bfloat16* p_C = static_cast<__nv_bfloat16*>(arg_C.data_ptr());
uint64_t v_d_addr = (uint64_t)arg_d_addr;
__nv_bfloat16* p_z = static_cast<__nv_bfloat16*>(arg_z.data_ptr());
uint64_t v_dt_bias_addr = (uint64_t)arg_dt_bias_addr;
__nv_bfloat16* p_output = static_cast<__nv_bfloat16*>(arg_output.data_ptr());
uint64_t v_state_batch_indices_addr = (uint64_t)arg_state_batch_indices_addr;
uint64_t v_dst_state_batch_indices_addr = (uint64_t)arg_dst_state_batch_indices_addr;
int32_t v_nheads = (int32_t)arg_nheads;
int32_t v_ngroups = (int32_t)arg_ngroups;
int32_t v_dim_tiles = (int32_t)arg_dim_tiles;
uint64_t v_state_stride_slot = (uint64_t)arg_state_stride_slot;
int64_t v_x_batch_stride = (int64_t)arg_x_batch_stride;
int64_t v_b_batch_stride = (int64_t)arg_b_batch_stride;
int64_t v_c_batch_stride = (int64_t)arg_c_batch_stride;
int64_t v_out_batch_stride = (int64_t)arg_out_batch_stride;
int64_t v_dt_batch_stride = (int64_t)arg_dt_batch_stride;
int64_t v_dt_head_stride = (int64_t)arg_dt_head_stride;
int64_t v_a_head_stride = (int64_t)arg_a_head_stride;
int64_t v_d_head_stride = (int64_t)arg_d_head_stride;
int64_t v_dt_bias_head_stride = (int64_t)arg_dt_bias_head_stride;
int32_t v_dt_softplus = (int32_t)arg_dt_softplus;
int32_t v_has_z = (int32_t)arg_has_z;
int32_t v_disable_state_update = (int32_t)arg_disable_state_update;
int64_t v_pad_slot_id = (int64_t)arg_pad_slot_id;
void* kargs[] = {&p_state,
&p_x,
&v_dt_addr,
&v_a_addr,
&p_B,
&p_C,
&v_d_addr,
&p_z,
&v_dt_bias_addr,
&p_output,
&v_state_batch_indices_addr,
&v_dst_state_batch_indices_addr,
&v_nheads,
&v_ngroups,
&v_dim_tiles,
&v_state_stride_slot,
&v_x_batch_stride,
&v_b_batch_stride,
&v_c_batch_stride,
&v_out_batch_stride,
&v_dt_batch_stride,
&v_dt_head_stride,
&v_a_head_stride,
&v_d_head_stride,
&v_dt_bias_head_stride,
&v_dt_softplus,
&v_has_z,
&v_disable_state_update,
&v_pad_slot_id};

dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z);
dim3 block(128u, 1u, 1u);

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

} // namespace cake_host_shim_825081a78f52121f

TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_825081a78f52121f::Run);
Loading
Loading