[Performance]: Shape-gated backend selection for FP8 block-scaling GEMM

[georgeliu95]

[Performance]: Shape-gated backend selection for FP8 block-scaling GEMM

Proposal to improve performance

Summary

I would like to discuss adding shape-gated backend selection for FP8 block-scaling GEMM in TensorRT-LLM, instead of relying on a single backend choice for all shapes within an architecture path.

The goal is not to make any alternative backend a global default. The proposal is to keep the current TensorRT-LLM path as the conservative default, while allowing a guarded dispatch mechanism to select an alternative backend only for shapes where it is known to be beneficial and safe.

Current upstream status checked before filing

As of TensorRT-LLM v1.3.0rc20, the FP8 block-scaling GEMM Python wrapper appears to expose a single Fp8BlockScalingGemmRunner tactic (get_valid_tactics() returns [0]) and calls torch.ops.trtllm.fp8_block_scaling_gemm_impl(...) through AutoTuner.

The C++ thop implementation then selects the implementation by SM version (SM103/SM100 -> Blackwell, SM90 -> Hopper) rather than by shape or measured backend performance. I did not find a current shape-gated backend-selection policy for this op.

Relevant public source locations:

• Python wrapper: tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
• C++ thop dispatch: cpp/tensorrt_llm/thop/fp8BlockScalingGemm.cpp

Motivation

The shape distribution can differ substantially between workloads:

• LLM decode and small-prefill paths often use small or moderate M, where a low-overhead default path is usually preferable.
• Diffusion or video-diffusion transformer blocks can exercise much larger M shapes, where a different backend may be more appropriate.

A global backend switch would risk regressing latency-sensitive or small-shape paths. A shape-gated policy can preserve the current safe default while creating room for workload-specific optimization.

Suggested direction

A safe implementation could look like:

1. Expose separate deterministic internal tactics for the current TensorRT-LLM implementation and any alternative FP8 block-scaling GEMM backend.
2. Add an availability check for each alternative backend, scoped to supported architectures and layouts.
3. Use a profiling/autotune cache or exact-shape policy table to select a backend only after validation.
4. Hard-fallback to TensorRT-LLM default for:
   • CUDA graph capture;
   • unsupported architecture/layout;
   • small M;
   • cache miss;
   • backend identity mismatch;
   • correctness validation failure.
5. Keep the policy conservative until there is end-to-end evidence on target workloads.

Non-goals

• This is not a request to make any alternative backend the global default.
• This is not a performance-regression report.
• This is not an end-to-end model speedup claim.
• This is not a request to change Blackwell behavior based on Hopper-only assumptions.

Report of performance regression

This is not a regression report.

This is a performance-design discussion: FP8 block-scaling GEMM may benefit from shape-aware backend selection, but the current default path should remain the fallback.

Misc discussion on performance

The main design concern is avoiding policy over-generalization.

The likely policy buckets are:

1. Small-M LLM decode/prefill: keep TensorRT-LLM default.
2. Medium or mixed shapes: use profiling/cache or a conservative allowlist.
3. Large-M diffusion/video-diffusion-style shapes: consider an alternative backend, but keep explicit counterexample handling.

The most useful next step may be to add a TensorRT-LLM-owned benchmark/validation path for FP8 block-scaling GEMM backend selection, then decide whether the dispatch policy should live in the existing AutoTuner/cache infrastructure or closer to the C++ thop dispatch layer.

Your current environment

I am filing this as a design/performance discussion based on public source inspection and local experimentation. Detailed benchmark data can be shared separately through an appropriate channel if maintainers want to evaluate the policy further.