Optimize W8A8 Triton backward matmul

[zzlol63]

Summary

• fuse row-wise activation and weight dequantization scales into the W8A8 Triton matmul epilogue
• store directly in the requested output dtype, avoiding the full INT32/FP32 intermediate and a second pointwise CUDA kernel
• add a faster 3-stage autotune candidate for the row-major backward layout
• retain a functionally equivalent non-Triton fallback

Performance

Measured on an NVIDIA RTX PRO 6000 Blackwell with the repository-root venv (torch 2.12.0+cu130, Triton 3.7.0), using the benchmark shape M=2098, K=3072, N=3088.

Baseline and final implementations ran in the same process against identical resident tensors. Both were precompiled/autotuned, measurement order alternated over 7 rounds, and each sample contained 1,000 launches timed with CUDA events. Values below are medians.

Path	Baseline	Final	Improvement
INT8 raw forward	82.474 us	82.419 us	+0.07%
INT8 raw backward	106.072 us	99.274 us	+6.85%
INT8 complete backward	121.638 us	108.499 us	+12.11%
FP8 raw forward	81.987 us	82.285 us	-0.36%
FP8 raw backward	107.642 us	100.085 us	+7.55%
FP8 complete backward	123.958 us	114.634 us	+8.13%

complete backward includes activation quantization, matmul, scale/dequantization, and output conversion under torch.compile(fullgraph=True).

A CUDA profiler check confirmed that the previous compiled path launches _mm_kernel followed by an Inductor pointwise scale/cast kernel. The fused path launches only _mm_kernel.

Validation

• ruff check on all three changed Python files
• in-memory syntax compilation
• INT8 exact-reference checks
• FP8 reference checks within expected tolerance
• irregular matrix shape (73 x 96 @ 96 x 80) covering masked edge tiles

Disclaimer

AI assistance was used to author the code (both Claude Code and Codex).