Skip to content

Commit 970315b

Browse files
committed
[None][perf] optimize DSV4 FP8 o_a decode epilogue
Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>
1 parent 6ce2609 commit 970315b

3 files changed

Lines changed: 327 additions & 14 deletions

File tree

tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3821,6 +3821,10 @@ def forward_fp8out(
38213821
sf_n = ceil_div(n, 128)
38223822
aligned_mn = pad_up(m, 4)
38233823
n_tiles_per_group = ceil_div(n, 128)
3824+
# The cooperative SMEM epilogue wins through M=32. Larger shapes
3825+
# retain the register epilogue and share one dynamic-M kernel.
3826+
use_fp8_smem_epilogue = m <= 32
3827+
fp8_smem_row_iters = ceil_div(m, 16) if use_fp8_smem_epilogue else 1
38243828

38253829
if c_tensor.dtype != torch.float8_e4m3fn or c_tensor.shape != (
38263830
batch_size, m, n):
@@ -3840,6 +3844,8 @@ def forward_fp8out(
38403844
mma_tiler_mn,
38413845
cluster_shape_mn,
38423846
self.use_tvm_ffi,
3847+
fp8_smem_row_iters,
3848+
use_fp8_smem_epilogue,
38433849
)
38443850
if cache_key not in self.__class__.kernel_cache:
38453851
a_ptr = make_ptr(cutlass.Float8E4M3FN,
@@ -3894,6 +3900,8 @@ def forward_fp8out(
38943900
sf_out_ptr,
38953901
aligned_mn,
38963902
n_tiles_per_group,
3903+
fp8_smem_row_iters,
3904+
use_fp8_smem_epilogue,
38973905
options="--opt-level 2 --enable-tvm-ffi",
38983906
)
38993907
self.__class__.kernel_cache[cache_key] = compiled_gemm

0 commit comments

Comments
 (0)