[QST] How to optimize cutlass int8 fprop2d kernel with NCxHWx interleaved layout, which is ~12% slower than TensorRT on Jetson Orin

Summary

I'm observing a performance gap between CUTLASS INT8 NCxHWx interleaved Conv2d Fprop kernel and TensorRT's equivalent kernel on Jetson AGX Orin. For a specific convolution configuration, CUTLASS is approximately 12% slower than TensorRT.

Environment

• Hardware: NVIDIA Jetson AGX Orin
• GPU Architecture: SM87 (Ampere)
• CUTLASS Version: 4.2 (tested with latest)
• TensorRT Version: 8.6 (bundled with JetPack)

Problem Description

I benchmarked the INT8 Conv2d Fprop kernel with NCxHWx<32> interleaved layout using both CUTLASS templates and TensorRT, and found a noticeable performance difference.

Convolution Parameters

Performance Comparison

TensorRT Kernel Name (from nsys profiling)

sm80_xmma_fprop_implicit_gemm_interleaved_i8i8_i8i32_f32_nchw_vect_c_32kcrs_vect_c_32_nchw_vect_c_32_tilesize128x128x64_stage4_warpsize2x2x1_g1_tensor16x8x32_t1r3s3_execute_kernel_trt

From the kernel name, TensorRT appears to use:

• Tile size: 128x128x64
• 4 pipeline stages
• Warp configuration: 2x2x1
• Tensor instruction: 16x8x32
• Specialized for 3x3 filter (t1r3s3)

CUTLASS Kernel Configuration

I am using the following CUTLASS template configuration:

using Conv2dKernel = typename cutlass::conv::kernel::DefaultConv2dFprop<
  int8_t, cutlass::layout::TensorNCxHWx<32>,  // Input: NCxHWx<32>
  int8_t, cutlass::layout::TensorCxRSKx<32>,  // Filter: CxRSKx<32>
  int8_t, cutlass::layout::TensorNCxHWx<32>,  // Output: NCxHWx<32>
  int32_t,                                     // Accumulator
  cutlass::arch::OpClassTensorOp,
  cutlass::arch::Sm80,
  cutlass::gemm::GemmShape<128, 128, 64>,      // ThreadBlock shape
  cutlass::gemm::GemmShape<64, 64, 64>,        // Warp shape
  cutlass::gemm::GemmShape<16, 8, 32>,         // Instruction shape
  cutlass::epilogue::thread::LinearCombinationClamp<
    int8_t,
    64 / cutlass::sizeof_bits<int8_t>::value,  // 8 elements per access
    int32_t,
    float,
    cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling
  >,
  cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<1>,
  4,                                           // Stages
  cutlass::arch::OpMultiplyAddSaturate,
  cutlass::conv::IteratorAlgorithm::kOptimized,
  cutlass::conv::StrideSupport::kStrided,
  16, 16                                       // Alignment
>::Kernel;

Questions

How to modify and optimize cutlass int8 fprop2d kernel to meet with tensorrt's performance

Additional Information

I can provide additional nsys profiling reports if needed.

Thank you for your guidance on how to achieve TensorRT-level performance with CUTLASS templates!

Start by reviewing the reported DefaultConv2dFprop configuration and the TensorRT kernel parameters, then compare the supplied nsys profiling results for the Jetson AGX Orin case. The work would be done when the CUTLASS INT8 NCxHWx<32> fprop2d kernel configuration is identified and documented as matching TensorRT-level performance for the stated convolution.