Skip to content

perf(cake_grouped_fp8_gemm): SM107a UE8M0 scale-word reuse for BLOCK_N 128 block-scaled programs #8972

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 6, 2026
+337 −126
Merged

perf(cake_grouped_fp8_gemm): SM107a UE8M0 scale-word reuse for BLOCK_N 128 block-scaled programs #8972

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
16 changes: 8 additions & 8 deletions ..._fp8_gemm_28834b76bb7d1f408674_binding.cu → ..._fp8_gemm_1c465144240239ebb5ca_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_grouped_fp8_gemm_28834b76bb7d1f408674, for the device source that defines it. Do not edit.
// Generated by Cake: the tvm-ffi host binding of kernel_cake_grouped_fp8_gemm_1c465144240239ebb5ca, for the device source that defines it. Do not edit.
#include <cuda_runtime.h>
#include <cuda.h>

#include "cake_ffi_checks.h"

#include <cstdint>

extern "C" __global__ void kernel_cake_grouped_fp8_gemm_28834b76bb7d1f408674(const __grid_constant__ CUtensorMap A, const __grid_constant__ CUtensorMap A64, const __grid_constant__ CUtensorMap A32, const __grid_constant__ CUtensorMap B, int* __restrict__ SFA, int* __restrict__ SFB, int* __restrict__ m_indices, const __grid_constant__ CUtensorMap C_tma, unsigned int shape_m, unsigned int shape_n, unsigned int grid_n, unsigned int k_tiles, unsigned int sfa_row_stride, unsigned int sfa_col_stride, unsigned int sfb_group_stride, unsigned int sfb_row_stride, unsigned int sfb_col_stride);
extern "C" __global__ void kernel_cake_grouped_fp8_gemm_1c465144240239ebb5ca(const __grid_constant__ CUtensorMap A, const __grid_constant__ CUtensorMap A64, const __grid_constant__ CUtensorMap A32, const __grid_constant__ CUtensorMap B, int* __restrict__ SFA, int* __restrict__ SFB, int* __restrict__ m_indices, const __grid_constant__ CUtensorMap C_tma, unsigned int shape_m, unsigned int shape_n, unsigned int grid_n, unsigned int k_tiles, unsigned int sfa_row_stride, unsigned int sfa_col_stride, unsigned int sfb_group_stride, unsigned int sfb_row_stride, unsigned int sfb_col_stride);

namespace cake_binding_kernel_cake_grouped_fp8_gemm_28834b76bb7d1f408674 {
namespace cake_binding_kernel_cake_grouped_fp8_gemm_1c465144240239ebb5ca {

using tvm::ffi::TensorView;

constexpr const char* KERNEL_NAME = "kernel_cake_grouped_fp8_gemm_28834b76bb7d1f408674";
constexpr const char* KERNEL_NAME = "kernel_cake_grouped_fp8_gemm_1c465144240239ebb5ca";

void CheckCuda(cudaError_t status, const char* what) {
TVM_FFI_CHECK(status == cudaSuccess, RuntimeError)
Expand Down Expand Up @@ -326,7 +326,7 @@ void Run(TensorView a_A, TensorView a_A64, TensorView a_A32, TensorView a_B, Ten
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_grouped_fp8_gemm_28834b76bb7d1f408674), cudaFuncAttributeMaxDynamicSharedMemorySize, 220160), "cudaFuncSetAttribute");
CheckCuda(cudaFuncSetAttribute(reinterpret_cast<const void*>(kernel_cake_grouped_fp8_gemm_1c465144240239ebb5ca), cudaFuncAttributeMaxDynamicSharedMemorySize, 220160), "cudaFuncSetAttribute");
prepared_devices |= device_bit;
}
}
Expand All @@ -353,11 +353,11 @@ void Run(TensorView a_A, TensorView a_A64, TensorView a_A32, TensorView a_B, Ten
config.blockDim = dim3(224, 1, 1);
config.dynamicSmemBytes = 220160;
config.stream = stream;
CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast<const void*>(kernel_cake_grouped_fp8_gemm_28834b76bb7d1f408674), kernel_args),
CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast<const void*>(kernel_cake_grouped_fp8_gemm_1c465144240239ebb5ca), kernel_args),
"cudaLaunchKernelExC");
}

} // namespace cake_binding_kernel_cake_grouped_fp8_gemm_28834b76bb7d1f408674
} // namespace cake_binding_kernel_cake_grouped_fp8_gemm_1c465144240239ebb5ca

TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_grouped_fp8_gemm_28834b76bb7d1f408674::Run);
TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_grouped_fp8_gemm_1c465144240239ebb5ca::Run);
TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version);
2 changes: 1 addition & 1 deletion ...d_fp8_gemm_28834b76bb7d1f408674_kernel.cu → ...d_fp8_gemm_1c465144240239ebb5ca_kernel.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -215,7 +215,7 @@ __device__ __forceinline__ void tmem_ld_x8(float* dst, int tmem_addr) {
extern "C" {

__global__ __launch_bounds__(224, LAUNCH_MIN_BLOCKS) void
kernel_cake_grouped_fp8_gemm_28834b76bb7d1f408674(const __grid_constant__ CUtensorMap A, const __grid_constant__ CUtensorMap A64, const __grid_constant__ CUtensorMap A32, const __grid_constant__ CUtensorMap B, int* __restrict__ SFA, int* __restrict__ SFB, int* __restrict__ m_indices, const __grid_constant__ CUtensorMap C_tma, unsigned int shape_m, unsigned int shape_n, unsigned int grid_n, unsigned int k_tiles, unsigned int sfa_row_stride, unsigned int sfa_col_stride, unsigned int sfb_group_stride, unsigned int sfb_row_stride, unsigned int sfb_col_stride)
kernel_cake_grouped_fp8_gemm_1c465144240239ebb5ca(const __grid_constant__ CUtensorMap A, const __grid_constant__ CUtensorMap A64, const __grid_constant__ CUtensorMap A32, const __grid_constant__ CUtensorMap B, int* __restrict__ SFA, int* __restrict__ SFB, int* __restrict__ m_indices, const __grid_constant__ CUtensorMap C_tma, unsigned int shape_m, unsigned int shape_n, unsigned int grid_n, unsigned int k_tiles, unsigned int sfa_row_stride, unsigned int sfa_col_stride, unsigned int sfb_group_stride, unsigned int sfb_row_stride, unsigned int sfb_col_stride)
{
const int tid = threadIdx.x;
const int warp = make_warp_uniform(tid / 32);
Expand Down
16 changes: 8 additions & 8 deletions ..._fp8_gemm_b63c85680050204fa4a2_binding.cu → ..._fp8_gemm_348a937af226cd4d36f9_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_grouped_fp8_gemm_b63c85680050204fa4a2, for the device source that defines it. Do not edit.
// Generated by Cake: the tvm-ffi host binding of kernel_cake_grouped_fp8_gemm_348a937af226cd4d36f9, for the device source that defines it. Do not edit.
#include <cuda_runtime.h>
#include <cuda.h>

#include "cake_ffi_checks.h"

#include <cstdint>

extern "C" __global__ void kernel_cake_grouped_fp8_gemm_b63c85680050204fa4a2(const __grid_constant__ CUtensorMap A, const __grid_constant__ CUtensorMap A64, const __grid_constant__ CUtensorMap A32, const __grid_constant__ CUtensorMap B, int* __restrict__ SFA, int* __restrict__ SFB, int* __restrict__ m_indices, const __grid_constant__ CUtensorMap C_tma, unsigned int shape_m, unsigned int shape_n, unsigned int grid_n, unsigned int k_tiles, unsigned int sfa_row_stride, unsigned int sfa_col_stride, unsigned int sfb_group_stride, unsigned int sfb_row_stride, unsigned int sfb_col_stride);
extern "C" __global__ void kernel_cake_grouped_fp8_gemm_348a937af226cd4d36f9(const __grid_constant__ CUtensorMap A, const __grid_constant__ CUtensorMap A64, const __grid_constant__ CUtensorMap A32, const __grid_constant__ CUtensorMap B, int* __restrict__ SFA, int* __restrict__ SFB, int* __restrict__ m_indices, const __grid_constant__ CUtensorMap C_tma, unsigned int shape_m, unsigned int shape_n, unsigned int grid_n, unsigned int k_tiles, unsigned int sfa_row_stride, unsigned int sfa_col_stride, unsigned int sfb_group_stride, unsigned int sfb_row_stride, unsigned int sfb_col_stride);

namespace cake_binding_kernel_cake_grouped_fp8_gemm_b63c85680050204fa4a2 {
namespace cake_binding_kernel_cake_grouped_fp8_gemm_348a937af226cd4d36f9 {

using tvm::ffi::TensorView;

constexpr const char* KERNEL_NAME = "kernel_cake_grouped_fp8_gemm_b63c85680050204fa4a2";
constexpr const char* KERNEL_NAME = "kernel_cake_grouped_fp8_gemm_348a937af226cd4d36f9";

void CheckCuda(cudaError_t status, const char* what) {
TVM_FFI_CHECK(status == cudaSuccess, RuntimeError)
Expand Down Expand Up @@ -326,7 +326,7 @@ void Run(TensorView a_A, TensorView a_A64, TensorView a_A32, TensorView a_B, Ten
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_grouped_fp8_gemm_b63c85680050204fa4a2), cudaFuncAttributeMaxDynamicSharedMemorySize, 220160), "cudaFuncSetAttribute");
CheckCuda(cudaFuncSetAttribute(reinterpret_cast<const void*>(kernel_cake_grouped_fp8_gemm_348a937af226cd4d36f9), cudaFuncAttributeMaxDynamicSharedMemorySize, 220160), "cudaFuncSetAttribute");
prepared_devices |= device_bit;
}
}
Expand All @@ -353,11 +353,11 @@ void Run(TensorView a_A, TensorView a_A64, TensorView a_A32, TensorView a_B, Ten
config.blockDim = dim3(224, 1, 1);
config.dynamicSmemBytes = 220160;
config.stream = stream;
CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast<const void*>(kernel_cake_grouped_fp8_gemm_b63c85680050204fa4a2), kernel_args),
CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast<const void*>(kernel_cake_grouped_fp8_gemm_348a937af226cd4d36f9), kernel_args),
"cudaLaunchKernelExC");
}

} // namespace cake_binding_kernel_cake_grouped_fp8_gemm_b63c85680050204fa4a2
} // namespace cake_binding_kernel_cake_grouped_fp8_gemm_348a937af226cd4d36f9

TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_grouped_fp8_gemm_b63c85680050204fa4a2::Run);
TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_grouped_fp8_gemm_348a937af226cd4d36f9::Run);
TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version);
Loading
Loading