feat: trtrllm-gen global scaled FP8 GEMMs (#1829)

In low latency context, it is not uncommon to encounter memory bandwidth
bound GEMMs with a tiny leading dimension M. These cases are currently
not addressed as efficiently as they could by library implementations.
To fill this gap, I propose to expose generated GEMM kernels optimized
for small batch sizes, which saturate memory bandwidth to a higher
degree.

The main challenge in doing so is that these GEMMs expect the weight
tensor (second operand) to be pre-processed into a layout more amenable
to maximizing memory bandwidth saturation. As such it is not practical
to expose them under the same API as the other GEMMs, as they are not
interchangeable without changing the caller's implementation. I have
tentatively exposed these GEMMs as "flavored" GEMMs, by contrast with
the more "vanilla" GEMMs currently available.

Summary of the changes:

- Added cpp runner to be jitted for these new GEMMs:
`csrc/trtllm_flavored_gemm_runner.cu`
- A separate `flashinfer/trtllm_flavored_gemm.py` file containing the
Python interface of the new GEMMs
- Some stylistic refactoring of the autotuner done while understanding
how it works
- Tests
- Benchmarks
- Some other minor cleanups along the way. Note I will undo the
extraction of `fp8_utils.py` as the implementations of `to_fp8` differ
between the places I extracted it for

Next step:

I will add more kernels for larger batch sizes. This is required because
the weight matrix shuffling commits the user to this interface.
Therefore, they also need efficient kernels for larger batches, which
they will encounter for prefills for example, when not doing disagg.

Benchmarking results on GB200:
```
m=1 n=2560 k=16384 9.65 TFLOPs/s over 0.008694 ms, 4.83 TB/s
m=1 n=2560 k=32768 11.34 TFLOPs/s over 0.014797 ms, 5.67 TB/s
m=1 n=5120 k=16384 15.10 TFLOPs/s over 0.011110 ms, 7.55 TB/s
m=1 n=5120 k=32768 12.21 TFLOPs/s over 0.027491 ms, 6.10 TB/s
m=1 n=8192 k=16384 11.75 TFLOPs/s over 0.022851 ms, 5.87 TB/s
m=1 n=8192 k=32768 13.06 TFLOPs/s over 0.041114 ms, 6.53 TB/s
m=2 n=2560 k=16384 18.38 TFLOPs/s over 0.009130 ms, 4.60 TB/s
m=2 n=2560 k=32768 21.21 TFLOPs/s over 0.015821 ms, 5.31 TB/s
m=2 n=5120 k=16384 30.21 TFLOPs/s over 0.011107 ms, 7.56 TB/s
m=2 n=5120 k=32768 24.41 TFLOPs/s over 0.027491 ms, 6.11 TB/s
m=2 n=8192 k=16384 23.43 TFLOPs/s over 0.022912 ms, 5.86 TB/s
m=2 n=8192 k=32768 26.15 TFLOPs/s over 0.041056 ms, 6.54 TB/s
m=4 n=2560 k=16384 36.22 TFLOPs/s over 0.009264 ms, 4.54 TB/s
m=4 n=2560 k=32768 43.55 TFLOPs/s over 0.015408 ms, 5.45 TB/s
m=4 n=5120 k=16384 60.40 TFLOPs/s over 0.011110 ms, 7.56 TB/s
m=4 n=5120 k=32768 48.82 TFLOPs/s over 0.027494 ms, 6.11 TB/s
m=4 n=8192 k=16384 46.71 TFLOPs/s over 0.022989 ms, 5.84 TB/s
m=4 n=8192 k=32768 52.10 TFLOPs/s over 0.041216 ms, 6.52 TB/s
m=8 n=2560 k=16384 72.47 TFLOPs/s over 0.009261 ms, 4.55 TB/s
m=8 n=2560 k=32768 84.84 TFLOPs/s over 0.015821 ms, 5.32 TB/s
m=8 n=5120 k=16384 120.84 TFLOPs/s over 0.011107 ms, 7.57 TB/s
m=8 n=5120 k=32768 97.37 TFLOPs/s over 0.027568 ms, 6.10 TB/s
m=8 n=8192 k=16384 93.41 TFLOPs/s over 0.022989 ms, 5.85 TB/s
m=8 n=8192 k=32768 104.21 TFLOPs/s over 0.041216 ms, 6.52 TB/s
m=16 n=2560 k=16384 138.70 TFLOPs/s over 0.009677 ms, 4.37 TB/s
m=16 n=2560 k=32768 174.22 TFLOPs/s over 0.015408 ms, 5.48 TB/s
m=16 n=5120 k=16384 231.03 TFLOPs/s over 0.011619 ms, 7.26 TB/s
m=16 n=5120 k=32768 190.13 TFLOPs/s over 0.028237 ms, 5.97 TB/s
m=16 n=8192 k=16384 180.96 TFLOPs/s over 0.023734 ms, 5.68 TB/s
m=16 n=8192 k=32768 205.52 TFLOPs/s over 0.041795 ms, 6.44 TB/s
m=32 n=2560 k=16384 260.92 TFLOPs/s over 0.010288 ms, 4.14 TB/s
m=32 n=2560 k=32768 322.64 TFLOPs/s over 0.016640 ms, 5.11 TB/s
m=32 n=5120 k=16384 421.01 TFLOPs/s over 0.012752 ms, 6.65 TB/s
m=32 n=5120 k=32768 371.18 TFLOPs/s over 0.028928 ms, 5.85 TB/s
m=32 n=8192 k=16384 348.80 TFLOPs/s over 0.024627 ms, 5.49 TB/s
m=32 n=8192 k=32768 400.89 TFLOPs/s over 0.042854 ms, 6.30 TB/s
m=64 n=2560 k=16384 466.29 TFLOPs/s over 0.011514 ms, 3.76 TB/s
m=64 n=2560 k=32768 458.96 TFLOPs/s over 0.023395 ms, 3.69 TB/s
m=64 n=5120 k=16384 673.11 TFLOPs/s over 0.015952 ms, 5.37 TB/s
m=64 n=5120 k=32768 679.79 TFLOPs/s over 0.031590 ms, 5.40 TB/s
m=64 n=8192 k=16384 648.00 TFLOPs/s over 0.026512 ms, 5.14 TB/s
m=64 n=8192 k=32768 766.41 TFLOPs/s over 0.044832 ms, 6.06 TB/s
```

Author: hypdeb

benchmarks/bench_mm_fp8.py

@@ -0,0 +1,98 @@
+"""
+Copyright (c) 2025 by FlashInfer team.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+  http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+from typing import Dict
+from flashinfer.autotuner import autotune
+from flashinfer.trtllm_low_latency_gemm import prepare_low_latency_gemm_weights
+import numpy as np
+import torch
+
+from flashinfer import mm_fp8
+from flashinfer.testing.utils import bench_gpu_time
+
+_cache_permute_indices: Dict[torch.Size, torch.Tensor] = {}
+
+
+def to_float8(
+    x: torch.Tensor, dtype=torch.float8_e4m3fn
+) -> tuple[torch.Tensor, torch.Tensor]:
+    finfo = torch.finfo(dtype)
+    min_val, max_val = x.aminmax()
+    amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12)
+    scale = finfo.max / amax
+    x_scl_sat = (x * scale).clamp(min=finfo.min, max=finfo.max)
+    return x_scl_sat.to(dtype), scale.float().reciprocal()
+
+
+def bench_mm_fp8(m, n, k, in_dtype, out_dtype):
+    torch.manual_seed(123)
+    input_tensor = torch.randn([m, k], device="cuda", dtype=torch.bfloat16)
+    input_fp8, input_inv_s = to_float8(input_tensor, dtype=in_dtype)
+
+    # mat2 row  major -> column major
+    mat2 = torch.randn([n, k], device="cuda", dtype=torch.bfloat16)
+    mat2_fp8, mat2_inv_s = to_float8(mat2, dtype=in_dtype)
+
+    res = torch.zeros([m, n], device="cuda", dtype=out_dtype)
+    global_scale = input_inv_s * mat2_inv_s
+
+    # Do row shuffling.
+    prepared_weights = prepare_low_latency_gemm_weights(
+        mat2_fp8, _cache_permute_indices
+    )
+
+    with autotune(True):
+        mm_fp8(
+            input_fp8,
+            prepared_weights,
+            global_scale,
+            out=res,
+        )
+
+    measurements = bench_gpu_time(
+        lambda: mm_fp8(
+            input_fp8,
+            prepared_weights,
+            global_scale,
+            res,
+        ),
+        dry_run_time_ms=500,
+        repeat_time_ms=2500,
+        use_cuda_graph=True,
+    )
+    ms = np.median(measurements)
+    tflops_per_second = 2 * m * n * k * 1e-9 / ms
+
+    bandwidth = (
+        (
+            input_fp8.numel() * input_fp8.element_size()
+            + prepared_weights.numel() * prepared_weights.element_size()
+            + res.numel() * res.element_size()
+        )
+        / ms
+        / 1e9
+    )
+
+    print(
+        f"mm_fp8 m={m} n={n} k={k} in_dtype={in_dtype} out_dtype={out_dtype}: {tflops_per_second:.2f} TFLOPs/s over {ms:.6f} ms, {bandwidth:.2f} TB/s"
+    )
+
+
+if __name__ == "__main__":
+    for m in [1, 2, 4, 8, 16, 32, 64]:
+        for n in [2560, 5120, 8192]:
+            for k in [16384, 32768]:
+                bench_mm_fp8(m, n, k, torch.float8_e4m3fn, torch.bfloat16)

csrc/trtllm_gemm_runner.cu

@@ -43,19 +43,20 @@ struct TrtllmGenGemmRunnerOptions {
 int64_t select_kernel_fp8(int32_t M, int32_t N, int32_t K,
                           const gemm::gemm::GemmInterface& interface) {
   static constexpr const char* KERNEL_NAME_HIGH_N_K_RATIO =
-      "gemm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s6_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_"
+      "gemm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s6_et64x8_m64x8x32_cga1x1x1_16dp256b_rM_TN_"
+      "transOut_"
       "noShflA_dsFp8_schedP2x2x1x3_sm100f";
 
   static constexpr const char* KERNEL_NAME_LOW_N_K_RATIO =
-      "gemm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s6_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_"
+      "gemm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s6_et64x32_m64x32x32_cga1x1x1_16dp256b_rM_TN_"
       "transOut_noShflA_dsFp8_schedS_sm100f";
 
   static constexpr const char* KERNEL_NAME_LARGE_N =
-      "gemm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s6_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_"
+      "gemm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s6_et64x32_m64x32x32_cga1x1x1_16dp256b_rM_TN_"
       "transOut_noShflA_dsFp8_schedP2x2x1x3_sm100f";
 
   static constexpr const char* KERNEL_NAME_DEFAULT =
-      "gemm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_"
+      "gemm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_rM_TN_"
       "transOut_noShflA_dsFp8_schedS_sm100f";
 
   double const n_k_ratio = static_cast<double>(N) / static_cast<double>(K);