Mixed Precision GEMM

Implement mixed-precision matrix multiplication with automatic type conversion and scaling

This guide demonstrates how to implement a mixed-precision GEMM kernel where input matrices have different data types. The kernel automatically handles type conversion and optional scaling operations.

Overview

Mixed-precision GEMM is essential for efficient AI inference workloads. This implementation supports:

• Narrow precision tensors: int4, int8, uint8, fp8
• Wide precision tensors: fp16, bf16
• Two transformation modes:
  • Convert-only: Direct type conversion
  • Convert-scale: Type conversion followed by element-wise scaling

Architecture

The mixed-input GEMM uses a warp-specialized persistent kernel design:

Input A (narrow)     Input B (wide)
     │                    │
     ├─ TMA Load          ├─ TMA Load
     ↓                    ↓
 Shared Memory        Shared Memory
     ↓                    │
 Transform Warp           │
(convert + scale)         │
     ↓                    ↓
Shared/Tensor Memory ──→ MMA Warps
     ↓                    ↓
 Accumulator (wide precision)
     ↓
 Epilogue Warps
     ↓
 Output C

Implementation

Running Examples

Convert-Only Mode (int8 → bf16)

python examples/python/CuTeDSL/blackwell/mixed_input_gemm/mixed_input_gemm.py \
  --a_dtype Int8 \
  --b_dtype BFloat16 \
  --scale_granularity_m 0 \
  --scale_granularity_k 0 \
  --c_dtype BFloat16 \
  --acc_dtype Float32 \
  --mma_tiler_mnk 128,128,64 \
  --cluster_shape_mn 1,1 \
  --mnkl 256,512,8192,1

python examples/python/CuTeDSL/blackwell/mixed_input_gemm/mixed_input_gemm.py \
  --a_dtype Int4 \
  --b_dtype BFloat16 \
  --scale_granularity_m 1 \
  --scale_granularity_k 256 \
  --c_dtype BFloat16 \
  --acc_dtype Float32 \
  --mma_tiler_mnk 256,128,128 \
  --cluster_shape_mn 2,1 \
  --use_2cta_instrs \
  --use_tma_store \
  --mnkl 1024,8192,6144,16

Performance Profiling

ncu python examples/python/CuTeDSL/blackwell/mixed_input_gemm/mixed_input_gemm.py \
  --a_dtype Int8 --b_dtype BFloat16 \
  --scale_granularity_m 0 --scale_granularity_k 0 \
  --c_dtype BFloat16 --acc_dtype Float32 \
  --mma_tiler_mnk 128,128,64 --cluster_shape_mn 1,1 \
  --mnkl 256,512,8192,1 \
  --warmup_iterations 1 --iterations 10 --skip_ref_check

Key Concepts

Scale Granularity

Scale granularity determines how many elements share the same scale factor:

• scale_granularity_m: Number of M-mode elements per scale
• scale_granularity_k: Number of K-mode elements per scale

For a GEMM with shape (M, N, K, L), the scale tensor shape is:

(M // scale_granularity_m, K // scale_granularity_k, L)

Example: For M=1024, K=6144, L=16 with granularity (1, 256):

• Scale tensor shape: (1024, 24, 16)
• Each element in M has its own scale
• Every 256 elements in K share a scale

Warp Specialization Benefits

Warp specialization improves performance by:

1. Reduced synchronization: Warps operate independently
2. Better register allocation: Each warp type has optimized register usage
3. Increased parallelism: Different operations happen concurrently
4. Cache efficiency: Specialized warps have more predictable access patterns

Transformation Storage Location

Transformed data can be stored in:

Shared Memory (SMEM):

• Used for K-major layouts
• Allows all MMA warps to access
• Higher SMEM usage

Tensor Memory (TMEM):

• Used for M-major layouts (Blackwell only)
• Direct accumulator access
• Saves SMEM bandwidth

Data Type Support

Narrow Precision (A)	Wide Precision (B)	Transform Mode	Accumulator
Int8, Uint8	Float16, BFloat16	Convert-Only	Float32
Int4	Float16, BFloat16	Convert-Scale	Float32
Float8E4M3FN	Float16, BFloat16	Convert-Only	Float32
Float8E5M2	Float16, BFloat16	Convert-Only	Float32

Constraints

• Scale granularity M must be 1 (current limitation)
• Scale granularity K must be a multiple of MMA tile K
• Scale tensor must be in M-major layout
• Narrow precision tensor A must use supported types
• Both input tensors require 16-byte alignment on contiguous dimension

Source Code

Mixed Input GEMM

Complete source: examples/python/CuTeDSL/blackwell/mixed_input_gemm/mixed_input_gemm.py