Add CDNA3 FP8 support

[Parxd]

This PR resolves #34 and adds:

• base fp8e4m3_fnuz type
  • utilities for packing and initialization
• fp8 support for rt_base, st, st_subtile
• fp8 support for MMA w/ 16x16x32 MFMA
• 8192x8192x8192 fp8 GEMM kernel example (non-scaled)
  • PyTorch reference (correctness-only, see hipBLASLt performance reference below)

The "tile" concepts (e.g. rt_base and st) currently use the generic TILE_*_DIM to determine base tile shapes. This PR adds a template specialization based on a tile's majorness (i.e. row v. col) for this value, since the use of the fp8 16x16x32 MFMA instruction introduces the need for a non-square base tile.

Modified st and st_subtile to store their majorness as well to accurately track width in subtiles & for swizzling, with a default of row-major (i.e. no change; same for TILE_*_DIM) as to not break any existing code.

All row-major st tiles remain bank conflict-free, but I noticed conflicts for some GEMM configs (both fp8 and bf16) that require col-major RTs; this change could help specialize the idx function for col-major st to setup conflict-free tiles later.

HipKittens benchmark

(.venv) root@6:/workdir/kernels/gemm/fp8fp32/mi300x/8192_256_256_64_32# make
# @mkdir -p 
/opt/rocm-7.2.1/bin/hipcc 256_256_64_32.cpp -DKITTENS_CDNA3 --offload-arch=gfx942 -std=c++20 -w -I/workdir/include -I/workdir/prototype -I/usr/include/python3.12 -I/opt/venv/lib/python3.12/site-packages/pybind11/include -L/usr/lib/python3.12/config-3.12-x86_64-linux-gnu -L/usr/lib/x86_64-linux-gnu  -ldl  -lm  -shared -fPIC -Rpass-analysis=kernel-resource-usage -I/workdir/include -I/opt/rocm-7.2.1/include -I/opt/rocm-7.2.1/include/hip  \
    -o tk_kernel.cpython-312-x86_64-linux-gnu.so 2>&1 # | tee 
256_256_64_32.cpp:32:1: remark: Function Name: _Z8micro_tk13micro_globals [-Rpass-analysis=kernel-resource-usage]
   32 | void micro_tk(const micro_globals g) {
      | ^
256_256_64_32.cpp:32:1: remark:     TotalSGPRs: 36 [-Rpass-analysis=kernel-resource-usage]
256_256_64_32.cpp:32:1: remark:     VGPRs: 250 [-Rpass-analysis=kernel-resource-usage]
256_256_64_32.cpp:32:1: remark:     AGPRs: 0 [-Rpass-analysis=kernel-resource-usage]
256_256_64_32.cpp:32:1: remark:     ScratchSize [bytes/lane]: 0 [-Rpass-analysis=kernel-resource-usage]
256_256_64_32.cpp:32:1: remark:     Dynamic Stack: False [-Rpass-analysis=kernel-resource-usage]
256_256_64_32.cpp:32:1: remark:     Occupancy [waves/SIMD]: 2 [-Rpass-analysis=kernel-resource-usage]
256_256_64_32.cpp:32:1: remark:     SGPRs Spill: 0 [-Rpass-analysis=kernel-resource-usage]
256_256_64_32.cpp:32:1: remark:     VGPRs Spill: 0 [-Rpass-analysis=kernel-resource-usage]
256_256_64_32.cpp:32:1: remark:     LDS Size [bytes/block]: 0 [-Rpass-analysis=kernel-resource-usage]

(.venv) root@6:/workdir/kernels/gemm/fp8fp32/mi300x/8192_256_256_64_32# python test_python.py 
src: 256_256_64_32.cpp
Warning: tk_kernel.cpython-313-x86_64-linux-gnu.so not found at /workdir/kernels/gemm/fp8fp32/mi300x/8192_256_256_64_32/tk_kernel.cpython-313-x86_64-linux-gnu.so, skipping.
C.dtype=torch.float32
Average execution time: 0.9551 ms
Performance: 1151.20 TFLOPS for 8192x8192 matrix multiplication.

Max error between kernel and reference: 0.0
Max error: 0.0
Mean error: 0.0
Number of large errors (>0.1): 0

TLDR:

• Shape: 8192x8192x8192
• Perf.: 1151 TFLOPs
• Warmup iters: 500
• Timed iters: 100

hipBLASLt benchmark

(.venv) root@6:/workdir/rocm-libraries/projects/hipblaslt/build/release/clients# ./hipblaslt-bench --api_method c --stride_a 0 --stride_b 0 --stride_c 0 --stride_d 0 --alpha 1.000000 --beta 0.000000 --transA T --transB N --batch_count 1 --scaleA 1 --scaleB 1 --a_type f8_r --b_type f8_r --c_type f32_r --d_type f32_r --scale_type f32_r --bias_type f32_r --compute_type f32_r --rotating 512 --iters 100 --cold_iters 500 -m 8192 -n 8192 -k 8192 --lda 8192 --ldb 8192 --ldc 8192 --ldd 8192

hipBLASLt version: 100202
hipBLASLt git version: dabb6df2b9
Query device success: there are 1 devices. (Target device ID is 0)
Device ID 0 : AMD Instinct MI300X VF gfx942:sramecc+:xnack-
with 205.8 GB memory, max. SCLK 2100 MHz, max. MCLK 1300 MHz, compute capability 9.4
maxGridDimX 2147483647, sharedMemPerBlock 65.5 KB, maxThreadsPerBlock 1024, warpSize 64

hipblaslt-bench INFO: Using f8_fnuz_r instead of f8_r
hipblaslt-bench INFO: Using f8_fnuz_r instead of f8_r
Rotating buffer 512 MiB. Needed Size: 384 MiB. Needed block count: 2 (Capped to max iters: 500)
Is supported 1 / Total solutions: 1
[0]:transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us
    T,N,0,1,8192,8192,8192,1,8192,67108864,0,8192,67108864,8192,67108864,8192,67108864,f8_fnuz_r,f8_fnuz_r,f32_r,f32_r,f32_r,1,1,0,0,0,0,0,none,0,f32_r,f32_r,512,0,0,1.12708e+06,384.402,975.54

TLDR:

• Shape: 8192x8192x8192
• Type: fp8fp32 (identity scale)
• Perf.: 1127 TFLOPs
• Warmup iters: 500
• Timed iters: 100

Details:

• Docker container: docker.io/rocm/pytorch
• ROCm version: 7.2.1
• Device: MI300X

[Parxd]

Hi @willhu-jpg just wondering if you had a chance to review this, no rush if CDNA3 stuff is lower priority at the moment + happy to make any changes