perf(cake_gated_act_mxfp8): narrow tiles for small problems and four rows per thread for large ones

flashinfer/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_2513ecc12f03a6081c2e_binding.cu

@@ -0,0 +1,98 @@
+/*
+ * Copyright (c) 2026 by FlashInfer team.
+ *
+ * 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 Cake: the tvm-ffi host binding of kernel_cake_gated_act_mxfp8_2513ecc12f03a6081c2e, 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_gated_act_mxfp8_2513ecc12f03a6081c2e(__nv_bfloat16* __restrict__ gated_input, __nv_bfloat16* __restrict__ grad, uint8_t* __restrict__ row_output, uint8_t* __restrict__ col_output, uint8_t* __restrict__ row_scales, uint8_t* __restrict__ col_scales, int M, int K);
+
+namespace cake_binding_kernel_cake_gated_act_mxfp8_2513ecc12f03a6081c2e {
+
+using tvm::ffi::TensorView;
+
+constexpr const char* KERNEL_NAME = "kernel_cake_gated_act_mxfp8_2513ecc12f03a6081c2e";
+
+void CheckCuda(cudaError_t status, const char* what) {
+  TVM_FFI_CHECK(status == cudaSuccess, RuntimeError)
+      << KERNEL_NAME << ": " << what << " failed: " << cudaGetErrorString(status);
+}
+
+void* DataPtr(const TensorView& t) { return static_cast<char*>(t.data_ptr()) + t.byte_offset(); }
+
+void Run(TensorView a_gated_input, TensorView a_grad, TensorView a_row_output, TensorView a_col_output, TensorView a_row_scales, TensorView a_col_scales, int64_t a_M, int64_t a_K, int64_t grid_x, int64_t grid_y, int64_t grid_z) {
+  DLDevice dev = a_gated_input.device();
+  tvm::ffi::CUDADeviceGuard device_guard(dev.device_id);
+  check_cuda_tensor(a_gated_input, "gated_input");
+  check_dtype(a_gated_input, DLDataType{kDLBfloat, 16, 1}, "gated_input");
+  check_cuda_tensor(a_grad, "grad");
+  check_dtype(a_grad, DLDataType{kDLBfloat, 16, 1}, "grad");
+  check_cuda_tensor(a_row_output, "row_output");
+  check_dtype(a_row_output, DLDataType{kDLUInt, 8, 1}, "row_output");
+  check_cuda_tensor(a_col_output, "col_output");
+  check_dtype(a_col_output, DLDataType{kDLUInt, 8, 1}, "col_output");
+  check_cuda_tensor(a_row_scales, "row_scales");
+  check_dtype(a_row_scales, DLDataType{kDLUInt, 8, 1}, "row_scales");
+  check_cuda_tensor(a_col_scales, "col_scales");
+  check_dtype(a_col_scales, DLDataType{kDLUInt, 8, 1}, "col_scales");
+  TVM_FFI_CHECK(a_M >= -2147483648LL && a_M <= 2147483647LL, ValueError)
+      << KERNEL_NAME << ": 'M' = " << a_M
+      << " is outside the i32 range [-2147483648, 2147483647]";
+  TVM_FFI_CHECK(a_K >= -2147483648LL && a_K <= 2147483647LL, ValueError)
+      << KERNEL_NAME << ": 'K' = " << a_K
+      << " is outside the i32 range [-2147483648, 2147483647]";
+  check_contiguous(a_gated_input, "gated_input");
+  check_contiguous(a_grad, "grad");
+  check_same_device(a_grad, a_gated_input, "grad", "gated_input");
+  check_contiguous(a_row_output, "row_output");
+  check_same_device(a_row_output, a_gated_input, "row_output", "gated_input");
+  check_contiguous(a_col_output, "col_output");
+  check_same_device(a_col_output, a_gated_input, "col_output", "gated_input");
+  check_contiguous(a_row_scales, "row_scales");
+  check_same_device(a_row_scales, a_gated_input, "row_scales", "gated_input");
+  check_contiguous(a_col_scales, "col_scales");
+  check_same_device(a_col_scales, a_gated_input, "col_scales", "gated_input");
+  TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError)
+      << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") must be positive";
+  TVM_FFI_CHECK(grid_x <= 2147483647 && grid_y <= 65535 && grid_z <= 65535, ValueError)
+      << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") exceeds CUDA's limits";
+  cudaStream_t stream = static_cast<cudaStream_t>(TVMFFIEnvGetStream(dev.device_type, dev.device_id));
+  void* v_gated_input = DataPtr(a_gated_input);
+  void* v_grad = DataPtr(a_grad);
+  void* v_row_output = DataPtr(a_row_output);
+  void* v_col_output = DataPtr(a_col_output);
+  void* v_row_scales = DataPtr(a_row_scales);
+  void* v_col_scales = DataPtr(a_col_scales);
+  int32_t v_M = static_cast<int32_t>(a_M);
+  int32_t v_K = static_cast<int32_t>(a_K);
+  void* kernel_args[] = {&v_gated_input, &v_grad, &v_row_output, &v_col_output, &v_row_scales, &v_col_scales, &v_M, &v_K};
+  cudaLaunchConfig_t config{};
+  config.gridDim = dim3(static_cast<unsigned>(grid_x), static_cast<unsigned>(grid_y), static_cast<unsigned>(grid_z));
+  config.blockDim = dim3(128, 1, 1);
+  config.dynamicSmemBytes = 4096;
+  config.stream = stream;
+  CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast<const void*>(kernel_cake_gated_act_mxfp8_2513ecc12f03a6081c2e), kernel_args),
+            "cudaLaunchKernelExC");
+}
+
+}  // namespace cake_binding_kernel_cake_gated_act_mxfp8_2513ecc12f03a6081c2e
+
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_gated_act_mxfp8_2513ecc12f03a6081c2e::Run);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version);