fix(cake_kda_det): gate-gradient accuracy of the GLM intra stage (diagonal TF32 rounding)

flashinfer/experimental/cake_kda_det_train/cake_jit.py

@@ -7366,7 +7366,7 @@
         ],
         "ffi_entry": "run",
         "grid_rule": {"kind": "chunks", "x_factor": 1},
-        "kernel_symbol": "kernel_cake_kda_det_train_328afe43951979e798d1",
+        "kernel_symbol": "kernel_cake_kda_det_train_ad4c93dbb0435933c8a1",
         "l2n": ["fused"],
         "launch": {
             "block": (288, 1, 1),
@@ -7375,7 +7375,7 @@
             "dynamic_smem_bytes": 115200,
             "use_pdl": False,
         },
-        "library": "cake_kda_det_train_328afe43951979e798d1",
+        "library": "cake_kda_det_train_ad4c93dbb0435933c8a1",
         "logical": {
             "beta": "beta_int",
             "chunk_bos": "chunk_bos",
@@ -7457,7 +7457,7 @@
         ],
         "ffi_entry": "run",
         "grid_rule": {"kind": "chunks", "x_factor": 1},
-        "kernel_symbol": "kernel_cake_kda_det_train_9ad80f98f22e0b3bce08",
+        "kernel_symbol": "kernel_cake_kda_det_train_a135a53a639f33f921ab",
         "l2n": ["fused"],
         "launch": {
             "block": (288, 1, 1),
@@ -7466,7 +7466,7 @@
             "dynamic_smem_bytes": 115200,
             "use_pdl": False,
         },
-        "library": "cake_kda_det_train_9ad80f98f22e0b3bce08",
+        "library": "cake_kda_det_train_a135a53a639f33f921ab",
         "logical": {
             "beta": "beta_int",
             "chunk_bos": "chunk_bos",

flashinfer/experimental/cake_kda_det_train/csrc/cake_kda_det_train_328afe43951979e798d1_binding.cu → flashinfer/experimental/cake_kda_det_train/csrc/cake_kda_det_train_a135a53a639f33f921ab_binding.cu

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-// Generated by Cake: the tvm-ffi host binding of kernel_cake_kda_det_train_328afe43951979e798d1, for the device source that defines it. Do not edit.
+// Generated by Cake: the tvm-ffi host binding of kernel_cake_kda_det_train_a135a53a639f33f921ab, for the device source that defines it. Do not edit.
 #include <cuda_runtime.h>
 #include <cuda.h>
 #include <cuda_bf16.h>
@@ -23,13 +23,13 @@
 
 #include <cstdint>
 
-extern "C" __global__ void kernel_cake_kda_det_train_328afe43951979e798d1(const __grid_constant__ CUtensorMap dAqk, const __grid_constant__ CUtensorMap dAkk, const __grid_constant__ CUtensorMap gk, const __grid_constant__ CUtensorMap k_e, const __grid_constant__ CUtensorMap q_e, const __grid_constant__ CUtensorMap dq_f_t, const __grid_constant__ CUtensorMap dk_f_t, const __grid_constant__ CUtensorMap dg_f_t, float* __restrict__ beta, float* __restrict__ dq_f, float* __restrict__ dk_f, float* __restrict__ dg_f, float* __restrict__ db_f, int* __restrict__ chunk_bos, int* __restrict__ chunk_len, __nv_bfloat16* __restrict__ q_raw, __nv_bfloat16* __restrict__ k_raw, float* __restrict__ qk_rstd, __nv_bfloat16* __restrict__ dq_ext, __nv_bfloat16* __restrict__ dk_ext, float* __restrict__ dg_ext, __nv_bfloat16* __restrict__ dbeta_ext, int num_heads, int num_qk_heads, int group, int num_tokens, int q_raw_row_stride, int q_raw_head_stride, int k_raw_row_stride, int k_raw_head_stride);
+extern "C" __global__ void kernel_cake_kda_det_train_a135a53a639f33f921ab(const __grid_constant__ CUtensorMap dAqk, const __grid_constant__ CUtensorMap dAkk, const __grid_constant__ CUtensorMap gk, const __grid_constant__ CUtensorMap k_e, const __grid_constant__ CUtensorMap q_e, const __grid_constant__ CUtensorMap dq_f_t, const __grid_constant__ CUtensorMap dk_f_t, const __grid_constant__ CUtensorMap dg_f_t, float* __restrict__ beta, float* __restrict__ dq_f, float* __restrict__ dk_f, float* __restrict__ dg_f, float* __restrict__ db_f, int* __restrict__ chunk_bos, int* __restrict__ chunk_len, __nv_bfloat16* __restrict__ q_raw, __nv_bfloat16* __restrict__ k_raw, float* __restrict__ qk_rstd, __nv_bfloat16* __restrict__ dq_ext, __nv_bfloat16* __restrict__ dk_ext, float* __restrict__ dg_ext, __nv_bfloat16* __restrict__ dbeta_ext, int num_heads, int num_qk_heads, int group, int num_tokens, int q_raw_row_stride, int q_raw_head_stride, int k_raw_row_stride, int k_raw_head_stride);
 
-namespace cake_binding_kernel_cake_kda_det_train_328afe43951979e798d1 {
+namespace cake_binding_kernel_cake_kda_det_train_a135a53a639f33f921ab {
 
 using tvm::ffi::TensorView;
 
-constexpr const char* KERNEL_NAME = "kernel_cake_kda_det_train_328afe43951979e798d1";
+constexpr const char* KERNEL_NAME = "kernel_cake_kda_det_train_a135a53a639f33f921ab";
 
 void CheckCuda(cudaError_t status, const char* what) {
   TVM_FFI_CHECK(status == cudaSuccess, RuntimeError)
@@ -487,7 +487,7 @@ void Run(TensorView a_dAqk, TensorView a_dAkk, TensorView a_gk, TensorView a_k_e
     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(cudaSetDevice(dev.device_id), "cudaSetDevice");
-      CheckCuda(cudaFuncSetAttribute(reinterpret_cast<const void*>(kernel_cake_kda_det_train_328afe43951979e798d1), cudaFuncAttributeMaxDynamicSharedMemorySize, 115200), "cudaFuncSetAttribute");
+      CheckCuda(cudaFuncSetAttribute(reinterpret_cast<const void*>(kernel_cake_kda_det_train_a135a53a639f33f921ab), cudaFuncAttributeMaxDynamicSharedMemorySize, 115200), "cudaFuncSetAttribute");
       prepared_devices |= device_bit;
     }
   }
@@ -527,11 +527,11 @@ void Run(TensorView a_dAqk, TensorView a_dAkk, TensorView a_gk, TensorView a_k_e
   config.blockDim = dim3(288, 1, 1);
   config.dynamicSmemBytes = 115200;
   config.stream = stream;
-  CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast<const void*>(kernel_cake_kda_det_train_328afe43951979e798d1), kernel_args),
+  CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast<const void*>(kernel_cake_kda_det_train_a135a53a639f33f921ab), kernel_args),
             "cudaLaunchKernelExC");
 }
 
-}  // namespace cake_binding_kernel_cake_kda_det_train_328afe43951979e798d1
+}  // namespace cake_binding_kernel_cake_kda_det_train_a135a53a639f33f921ab
 
-TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_kda_det_train_328afe43951979e798d1::Run);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_kda_det_train_a135a53a639f33f921ab::Run);
 TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version);