Skip to content

perf(cake_gated_act_mxfp8): band-path program for the SM107 forward both entry below 2^20 #8023

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 30, 2026
Merged

perf(cake_gated_act_mxfp8): band-path program for the SM107 forward both entry below 2^20 #8023

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
23 changes: 23 additions & 0 deletions flashinfer/cake_gated_act_mxfp8_backend/build.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1578,6 +1578,29 @@
"csrc/cake_gated_act_mxfp8_dd19cc9df9f6cf060a97_binding.cu"
]
},
{
"arch": "sm_107a",
"closure_sha256": "70a788a2fa82a216b95b2a657014b191869170665c385789d0104cff922ee185",
"entry": "run",
"gencode": "-gencode=arch=compute_107a,code=sm_107a",
"min_cuda": "13.4",
"name": "cake_gated_act_mxfp8_e62d5f6dda3083f31950_sm_107a",
"nvcc_flags": [
"--prec-div=true",
"--prec-sqrt=true",
"--ftz=false"
],
"nvcc_flags_if": [],
"path": "aot/cake_gated_act_mxfp8_e62d5f6dda3083f31950_sm_107a.so",
"sha256": {
"csrc/cake_gated_act_mxfp8_e62d5f6dda3083f31950_binding.cu": "b43d9d0b40c2f387f9cbb87700442e72351f90cdd81960b509d491e07d29fa7b",
"csrc/cake_gated_act_mxfp8_e62d5f6dda3083f31950_kernel.cu": "cee181a75c1de648acf411ddfea8f981086fae4972eedfa186000692463687c9"
},
"sources": [
"csrc/cake_gated_act_mxfp8_e62d5f6dda3083f31950_kernel.cu",
"csrc/cake_gated_act_mxfp8_e62d5f6dda3083f31950_binding.cu"
]
},
{
"arch": "sm_100a",
"closure_sha256": "868e35d8508dd8fb540376ef9f63d5165730e90306e1c231642357cff205a818",
Expand Down
77 changes: 63 additions & 14 deletions flashinfer/cake_gated_act_mxfp8_backend/cake_jit.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1314,12 +1314,10 @@
"route": "both",
"select_rule": None,
},
# The forward, rowwise and colwise: the route's fallback on sm_107a; on
# sm_103a it runs when M * K >= 2^21 and no record of a higher threshold
# admits the call.
"cake_gated_act_mxfp8_0520c14fde53ef7bfea7": {
# The forward, rowwise and colwise, with 4 CTAs per SM and the called band
# path: the route's fallback on sm_107a (no select_rule).
"cake_gated_act_mxfp8_e62d5f6dda3083f31950": {
"arches": [
"sm_103a",
"sm_107a",
],
"arg_plan": [
Expand All @@ -1340,16 +1338,16 @@
"tile_cols": 128,
"tile_rows": 32,
},
"kernel": "cake_gated_act_mxfp8_0520c14fde53ef7bfea7",
"kernel_symbol": "kernel_cake_gated_act_mxfp8_0520c14fde53ef7bfea7",
"kernel": "cake_gated_act_mxfp8_e62d5f6dda3083f31950",
"kernel_symbol": "kernel_cake_gated_act_mxfp8_e62d5f6dda3083f31950",
"launch": {
"block": (256, 1, 1),
"cluster": (1, 1, 1),
"cooperative": False,
"dynamic_smem_bytes": 8192,
"use_pdl": False,
},
"library": "cake_gated_act_mxfp8_0520c14fde53ef7bfea7",
"library": "cake_gated_act_mxfp8_e62d5f6dda3083f31950",
"logical": {
"K": "K",
"M": "M",
Expand All @@ -1361,12 +1359,7 @@
},
"role": "fwd",
"route": "both",
"select_rule": {
"elements": {
"sm_103a": 2097152,
},
"kind": "elements_at_least",
},
"select_rule": None,
},
# The forward, rowwise and colwise, with 8 CTAs per SM, the packed sigmoid
# and the called band path (on sm_100a only): it runs when M * K >= 2^22
Expand Down Expand Up @@ -1475,6 +1468,61 @@
"kind": "elements_at_least",
},
},
# The forward, rowwise and colwise (on sm_103a and sm_107a only): it runs
# when M * K >= 2^21 (2^20 on sm_107a) and no record of a higher threshold
# admits the call.
"cake_gated_act_mxfp8_0520c14fde53ef7bfea7": {
"arches": [
"sm_103a",
"sm_107a",
],
"arg_plan": [
["buffer", "gated_input"],
["buffer", "row_output"],
["buffer", "col_output"],
["buffer", "row_scales"],
["buffer", "col_scales"],
["parameter", "M"],
["parameter", "K"],
["grid", "grid_x"],
["grid", "grid_y"],
["grid", "grid_z"],
],
"ffi_entry": "run",
"grid_rule": {
"kind": "tiles",
"tile_cols": 128,
"tile_rows": 32,
},
"kernel": "cake_gated_act_mxfp8_0520c14fde53ef7bfea7",
"kernel_symbol": "kernel_cake_gated_act_mxfp8_0520c14fde53ef7bfea7",
"launch": {
"block": (256, 1, 1),
"cluster": (1, 1, 1),
"cooperative": False,
"dynamic_smem_bytes": 8192,
"use_pdl": False,
},
"library": "cake_gated_act_mxfp8_0520c14fde53ef7bfea7",
"logical": {
"K": "K",
"M": "M",
"col_output": "col_output",
"col_scales": "col_scales",
"gated_input": "gated_input",
"row_output": "row_output",
"row_scales": "row_scales",
},
"role": "fwd",
"route": "both",
"select_rule": {
"elements": {
"sm_103a": 2097152,
"sm_107a": 1048576,
},
"kind": "elements_at_least",
},
},
# The forward, rowwise and colwise, with 8 CTAs per SM and the packed
# sigmoid (on sm_103a only): it runs when M * K >= 2^22 and no record of a
# higher threshold admits the call.
Expand Down Expand Up @@ -3401,6 +3449,7 @@
"cake_gated_act_mxfp8_6902c8391e6f660586a4",
],
"fwd_both": [
"cake_gated_act_mxfp8_e62d5f6dda3083f31950",
"cake_gated_act_mxfp8_0520c14fde53ef7bfea7",
"cake_gated_act_mxfp8_3b4157aba492acf411db",
"cake_gated_act_mxfp8_bbe04581406928a16cbc",
Expand Down
93 changes: 93 additions & 0 deletions ...er/cake_gated_act_mxfp8_backend/csrc/cake_gated_act_mxfp8_e62d5f6dda3083f31950_binding.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,93 @@
/*
* 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_e62d5f6dda3083f31950, 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_e62d5f6dda3083f31950(__nv_bfloat16* __restrict__ gated_input, 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_e62d5f6dda3083f31950 {

using tvm::ffi::TensorView;

constexpr const char* KERNEL_NAME = "kernel_cake_gated_act_mxfp8_e62d5f6dda3083f31950";

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_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_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_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_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_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(256, 1, 1);
config.dynamicSmemBytes = 8192;
config.stream = stream;
CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast<const void*>(kernel_cake_gated_act_mxfp8_e62d5f6dda3083f31950), kernel_args),
"cudaLaunchKernelExC");
}

} // namespace cake_binding_kernel_cake_gated_act_mxfp8_e62d5f6dda3083f31950

TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_gated_act_mxfp8_e62d5f6dda3083f31950::Run);
TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version);
Loading
Loading