Replace hardcoded FP8 e4m3fn max (448.0) with a parameterized FP8_MAX derived from torch.finfo(dtype).max in all 5 blockwise quantization Triton kernels, their Python wrapper functions, and the 3 PyTorch reference implementations. This allows the kernels to operate with both float8_e4m3fn (NVIDIA, max=448) and float8_e4m3fnuz (AMD MI300, max=240).

kernels.py:

• Add FP8_MAX as a tl.constexpr parameter to triton_fp8_blockwise_act_quant_lhs_kernel, triton_fp8_blockwise_act_quant_rhs_kernel, triton_fp8_blockwise_act_quant_transposed_lhs_kernel, triton_fp8_blockwise_weight_quant_rhs_kernel, and triton_fp8_blockwise_weight_quant_transposed_rhs_kernel. Each kernel previously had max_fp8_e4m3 = 448.0 / min_fp8_e4m3 = -448.0 inline; these are replaced with the passed-in FP8_MAX and -FP8_MAX.
• In the 5 wrapper functions, compute fp8_max = torch.finfo(dtype).max and forward it to the kernel call. Widen the dtype assertion from [torch.float8_e4m3fn] to {torch.float8_e4m3fn, torch.float8_e4m3fnuz}. Default dtype parameter changed from torch.float8_e4m3fn to the platform-aware e4m3_dtype (from torchao.float8.config).
• In the 3 reference implementations (torch_blockwise_scale_act_quant_lhs, torch_blockwise_scale_act_quant_rhs, torch_blockwise_scale_weight_quant), replace hardcoded torch.finfo(torch.float8_e4m3fn) with torch.finfo(dtype) and cast outputs to the passed dtype instead of torch.float8_e4m3fn.

test_blockwise_kernels.py:

• Replace is_sm_at_least_90() capability gate with is_sm_at_least_90() || is_MI300() || is_MI350() across all 7 tests.
• Replace hardcoded torch.float8_e4m3fn parametrize values with e4m3_dtype.
• Remove @skip_if_rocm decorators from the 5 quantization kernel tests.

Benchmark Results (AMD Instinct MI300X)

Environment: PyTorch 2.9.1+rocm7.2.0, Triton 3.5.1+rocm7.2.0, single MI300X GPU

Correctness

All 9 test configurations produce bit-identical results between old (per-expert loop) and new (grouped GEMM kernel) paths (max_diff = 0.0).

GEMM Kernel Only: per-expert loop (old) vs grouped kernel (new)

Full Forward: old Triton vs new Triton vs BF16 baseline

Forward+Backward (new Triton path, end-to-end)

Key takeaways:

• The new grouped GEMM kernel provides 1.2x-11x speedup over the old per-expert loop at the kernel level, with the largest gains on workloads with many experts and smaller per-expert M.
• End-to-end forward speedup is 1.05x-6.1x over the old path (quantization overhead is now the dominant cost at larger sizes).
• The Triton FP8 GEMM is currently slower than BF16 rocBLAS because it does not yet use hardware FP8 matrix cores; the benefit comes from reduced memory traffic which will be more impactful at scale.

cc: @BowenBao