Skip to content

fix(cake_rmsnorm): sm_100a and sm_107a dw-accuracy backward routes #8827

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
Nov 5, 2026
Merged

fix(cake_rmsnorm): sm_100a and sm_107a dw-accuracy backward routes #8827

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
1,081 changes: 782 additions & 299 deletions flashinfer/cake_rmsnorm_train_backend/build.json
View file
Open in desktop

Large diffs are not rendered by default.

2,534 changes: 1,689 additions & 845 deletions flashinfer/cake_rmsnorm_train_backend/cake_jit.py
View file
Open in desktop

Large diffs are not rendered by default.

696 changes: 443 additions & 253 deletions flashinfer/cake_rmsnorm_train_backend/cake_windows.py
View file
Open in desktop

Large diffs are not rendered by default.

5 changes: 0 additions & 5 deletions flashinfer/cake_rmsnorm_train_backend/csrc/cake_rmsnorm_train_07f9f2fbea7163639c27_kernel.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -116,12 +116,7 @@ kernel_cake_rmsnorm_train_07f9f2fbea7163639c27(__nv_bfloat16* __restrict__ g, __

extern __shared__ __align__(1024) char smem_raw[];
int smem;
#if __CUDA_ARCH__ == 1000
asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw));
smem = make_warp_uniform(smem);
#else
smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw);
#endif

const int mbar_base = smem;
#define bars_addr (mbar_base + 0)
Expand Down
16 changes: 8 additions & 8 deletions ...orm_train_ddedaa3be7466a4faa97_binding.cu → ...orm_train_083b5f2e2a0d1b4f1c82_binding.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,21 +14,21 @@
* limitations under the License.
*/

// Generated by Cake: the tvm-ffi host binding of kernel_cake_rmsnorm_train_ddedaa3be7466a4faa97, for the device source that defines it. Do not edit.
// Generated by Cake: the tvm-ffi host binding of kernel_cake_rmsnorm_train_083b5f2e2a0d1b4f1c82, for the device source that defines it. Do not edit.
#include <cuda_runtime.h>
#include <cuda_bf16.h>

#include "cake_ffi_checks.h"

#include <cstdint>

extern "C" __global__ void kernel_cake_rmsnorm_train_ddedaa3be7466a4faa97(__nv_bfloat16* __restrict__ g, __nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ w, float* __restrict__ r, __nv_bfloat16* __restrict__ dx, float* __restrict__ dw, double* __restrict__ partial, unsigned int* __restrict__ counters, __nv_bfloat16* __restrict__ g_h, int T, int n_chunks, int rows_per_chunk, long long g_stride_p0, long long x_stride_p0, long long g_h_stride_p0);
extern "C" __global__ void kernel_cake_rmsnorm_train_083b5f2e2a0d1b4f1c82(__nv_bfloat16* __restrict__ g, __nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ w, float* __restrict__ r, __nv_bfloat16* __restrict__ dx, float* __restrict__ dw, double* __restrict__ partial, unsigned int* __restrict__ counters, __nv_bfloat16* __restrict__ g_h, int T, int n_chunks, int rows_per_chunk, long long g_stride_p0, long long x_stride_p0, long long g_h_stride_p0);

namespace cake_binding_kernel_cake_rmsnorm_train_ddedaa3be7466a4faa97 {
namespace cake_binding_kernel_cake_rmsnorm_train_083b5f2e2a0d1b4f1c82 {

using tvm::ffi::TensorView;

constexpr const char* KERNEL_NAME = "kernel_cake_rmsnorm_train_ddedaa3be7466a4faa97";
constexpr const char* KERNEL_NAME = "kernel_cake_rmsnorm_train_083b5f2e2a0d1b4f1c82";

void CheckCuda(cudaError_t status, const char* what) {
TVM_FFI_CHECK(status == cudaSuccess, RuntimeError)
Expand Down Expand Up @@ -111,7 +111,7 @@ void Run(TensorView a_g, TensorView a_x, TensorView a_w, TensorView a_r, TensorV
static thread_local uint64_t prepared_devices = 0;
const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0};
if (device_bit == 0 || (prepared_devices & device_bit) == 0) {
CheckCuda(cudaFuncSetAttribute(reinterpret_cast<const void*>(kernel_cake_rmsnorm_train_ddedaa3be7466a4faa97), cudaFuncAttributeMaxDynamicSharedMemorySize, 148736), "cudaFuncSetAttribute");
CheckCuda(cudaFuncSetAttribute(reinterpret_cast<const void*>(kernel_cake_rmsnorm_train_083b5f2e2a0d1b4f1c82), cudaFuncAttributeMaxDynamicSharedMemorySize, 148736), "cudaFuncSetAttribute");
prepared_devices |= device_bit;
}
}
Expand All @@ -136,11 +136,11 @@ void Run(TensorView a_g, TensorView a_x, TensorView a_w, TensorView a_r, TensorV
config.blockDim = dim3(384, 1, 1);
config.dynamicSmemBytes = 148736;
config.stream = stream;
CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast<const void*>(kernel_cake_rmsnorm_train_ddedaa3be7466a4faa97), kernel_args),
CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast<const void*>(kernel_cake_rmsnorm_train_083b5f2e2a0d1b4f1c82), kernel_args),
"cudaLaunchKernelExC");
}

} // namespace cake_binding_kernel_cake_rmsnorm_train_ddedaa3be7466a4faa97
} // namespace cake_binding_kernel_cake_rmsnorm_train_083b5f2e2a0d1b4f1c82

TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_rmsnorm_train_ddedaa3be7466a4faa97::Run);
TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_rmsnorm_train_083b5f2e2a0d1b4f1c82::Run);
TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version);
693 changes: 417 additions & 276 deletions ...norm_train_ddedaa3be7466a4faa97_kernel.cu → ...norm_train_083b5f2e2a0d1b4f1c82_kernel.cu
View file
Open in desktop

Large diffs are not rendered by default.

20 changes: 10 additions & 10 deletions ...orm_train_0c84999caf2b86c2e09e_binding.cu → ...orm_train_0af20ef9ee6868c5569d_binding.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,21 +14,21 @@
* limitations under the License.
*/

// Generated by Cake: the tvm-ffi host binding of kernel_cake_rmsnorm_train_0c84999caf2b86c2e09e, for the device source that defines it. Do not edit.
// Generated by Cake: the tvm-ffi host binding of kernel_cake_rmsnorm_train_0af20ef9ee6868c5569d, for the device source that defines it. Do not edit.
#include <cuda_runtime.h>
#include <cuda_bf16.h>

#include "cake_ffi_checks.h"

#include <cstdint>

extern "C" __global__ void kernel_cake_rmsnorm_train_0c84999caf2b86c2e09e(__nv_bfloat16* __restrict__ g, __nv_bfloat16* __restrict__ x, float* __restrict__ w, float* __restrict__ r, __nv_bfloat16* __restrict__ dx, float* __restrict__ dw, float* __restrict__ partial, unsigned int* __restrict__ counters, __nv_bfloat16* __restrict__ g_h, int T, int n_chunks, int rows_per_chunk, long long g_stride_p0, long long x_stride_p0, long long g_h_stride_p0);
extern "C" __global__ void kernel_cake_rmsnorm_train_0af20ef9ee6868c5569d(__nv_bfloat16* __restrict__ g, __nv_bfloat16* __restrict__ x, float* __restrict__ w, float* __restrict__ r, __nv_bfloat16* __restrict__ dx, float* __restrict__ dw, double* __restrict__ partial, unsigned int* __restrict__ counters, __nv_bfloat16* __restrict__ g_h, int T, int n_chunks, int rows_per_chunk, long long g_stride_p0, long long x_stride_p0, long long g_h_stride_p0);

namespace cake_binding_kernel_cake_rmsnorm_train_0c84999caf2b86c2e09e {
namespace cake_binding_kernel_cake_rmsnorm_train_0af20ef9ee6868c5569d {

using tvm::ffi::TensorView;

constexpr const char* KERNEL_NAME = "kernel_cake_rmsnorm_train_0c84999caf2b86c2e09e";
constexpr const char* KERNEL_NAME = "kernel_cake_rmsnorm_train_0af20ef9ee6868c5569d";

void CheckCuda(cudaError_t status, const char* what) {
TVM_FFI_CHECK(status == cudaSuccess, RuntimeError)
Expand Down Expand Up @@ -63,7 +63,7 @@ void Run(TensorView a_g, TensorView a_x, TensorView a_w, TensorView a_r, TensorV
check_cuda_tensor(a_dw, "dw");
check_dtype(a_dw, DLDataType{kDLFloat, 32, 1}, "dw");
check_cuda_tensor(a_partial, "partial");
check_dtype(a_partial, DLDataType{kDLFloat, 32, 1}, "partial");
check_dtype(a_partial, DLDataType{kDLFloat, 64, 1}, "partial");
check_cuda_tensor(a_counters, "counters");
check_dtype(a_counters, DLDataType{kDLUInt, 32, 1}, "counters");
check_cuda_tensor(a_g_h, "g_h");
Expand Down Expand Up @@ -111,7 +111,7 @@ void Run(TensorView a_g, TensorView a_x, TensorView a_w, TensorView a_r, TensorV
static thread_local uint64_t prepared_devices = 0;
const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0};
if (device_bit == 0 || (prepared_devices & device_bit) == 0) {
CheckCuda(cudaFuncSetAttribute(reinterpret_cast<const void*>(kernel_cake_rmsnorm_train_0c84999caf2b86c2e09e), cudaFuncAttributeMaxDynamicSharedMemorySize, 99840), "cudaFuncSetAttribute");
CheckCuda(cudaFuncSetAttribute(reinterpret_cast<const void*>(kernel_cake_rmsnorm_train_0af20ef9ee6868c5569d), cudaFuncAttributeMaxDynamicSharedMemorySize, 66688), "cudaFuncSetAttribute");
prepared_devices |= device_bit;
}
}
Expand All @@ -134,13 +134,13 @@ void Run(TensorView a_g, TensorView a_x, TensorView a_w, TensorView a_r, TensorV
cudaLaunchConfig_t config{};
config.gridDim = dim3(static_cast<unsigned>(grid_x), static_cast<unsigned>(grid_y), static_cast<unsigned>(grid_z));
config.blockDim = dim3(256, 1, 1);
config.dynamicSmemBytes = 99840;
config.dynamicSmemBytes = 66688;
config.stream = stream;
CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast<const void*>(kernel_cake_rmsnorm_train_0c84999caf2b86c2e09e), kernel_args),
CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast<const void*>(kernel_cake_rmsnorm_train_0af20ef9ee6868c5569d), kernel_args),
"cudaLaunchKernelExC");
}

} // namespace cake_binding_kernel_cake_rmsnorm_train_0c84999caf2b86c2e09e
} // namespace cake_binding_kernel_cake_rmsnorm_train_0af20ef9ee6868c5569d

TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_rmsnorm_train_0c84999caf2b86c2e09e::Run);
TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_rmsnorm_train_0af20ef9ee6868c5569d::Run);
TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version);
Loading
Loading