dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

fix: corex_batched_gemm_kernel — use OpClassTensorOp + arch::Cu10 + FP32 accumulator

Browse Source 
Root cause of 25ms (vs expected 2.5ms):
1. ElementAccumulator was half_t → now float (FP32 accumulation)
2. Missing OpClassTensorOp → was defaulting to OpClassSimt (CUDA cores only)
3. Missing arch::Cu10 → was defaulting to arch::Sm61

With these fixes it should use __ivcorex_matrix_mad_f32x4_f16x4 (TCU)
same as moe_cutlass_batched.cu which benchmarked at 2.462ms.
This commit is contained in:
Claude
2026-08-15 12:34:01 +00:00
parent 3481f2903f
commit d8d241bf9f
1 changed files with 11 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
ex_engine/xllm_kernels/cuda/corex_batched_gemm_kernel.cu
View File

@@ -35,17 +35,24 @@ cudaError_t cutlass_batched_hgemm(
using ElementA = cutlass::half_t; using ElementA = cutlass::half_t;
using ElementB = cutlass::half_t; using ElementB = cutlass::half_t;
using ElementC = cutlass::half_t; using ElementC = cutlass::half_t;
using ElementAccumulator = cutlass::half_t; using ElementAccumulator = float;
using Gemm = cutlass::gemm::device::GemmBatched< using Gemm = cutlass::gemm::device::GemmBatched<
ElementA, cutlass::layout::ColumnMajor, // A ElementA, cutlass::layout::ColumnMajor, // A
ElementB, cutlass::layout::ColumnMajor, // B ElementB, cutlass::layout::ColumnMajor, // B
ElementC, cutlass::layout::ColumnMajor, // C ElementC, cutlass::layout::ColumnMajor, // C
ElementAccumulator // accumulator ElementAccumulator, // accumulator = FP32
cutlass::arch::OpClassTensorOp, // use TCU (not SIMT)
cutlass::arch::Cu10 // BI-V100 arch
// Defaults from DefaultGemmConfiguration:
// ThreadblockShape = <128, 128, 32>
// WarpShape = <32, 32, 32>
// InstructionShape = <16, 16, 16>
// Stages = 2
>; >;
ElementAccumulator alpha_val(1.0f); float alpha_val = 1.0f;
ElementAccumulator beta_val(0.0f); float beta_val = 0.0f;
Gemm gemm_op; Gemm gemm_op;