perf: TRT-LLM MoE Block-FP8 activation optimization (#2063)

<!-- .github/pull_request_template.md -->

## 📌 Description

- Small optimization to the activation kernel for block-FP8 MoE for
large batch size.

| BS | Baseline, us | Optimized, us |
| ------------- | ------------- | ------------- |
| 1 | 2.4  | 2.1  |
| 32  | 3.5  | 2.6 |
| 256  | 21.7  | 8.7  |
| 1024 | 84.4  | 23.8  |
| 4096  | 333 | 87.0  |
| 16384  | 1330  | 365  |

- Adding micro-benchmark for DS FP8 implemented by @IwakuraRein.

<!-- What does this PR do? Briefly describe the changes and why they’re
needed. -->

## 🔍 Related Issues

<!-- Link any related issues here -->

## 🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull
request, please make sure the following items are complete.

### ✅ Pre-commit Checks

- [x] I have installed `pre-commit` by running `pip install pre-commit`
(or used your preferred method).
- [x] I have installed the hooks with `pre-commit install`.
- [x] I have run the hooks manually with `pre-commit run --all-files`
and fixed any reported issues.

> If you are unsure about how to set up `pre-commit`, see [the
pre-commit documentation](https://pre-commit.com/).

## 🧪 Tests

- [x] Tests have been added or updated as needed.
- [x] All tests are passing (`unittest`, etc.).

## Reviewer Notes

<!-- Optional: anything you'd like reviewers to focus on, concerns, etc.
-->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Improved Mixture-of-Experts inference with configurable multi-token
batching per GPU core for higher throughput.
* Expanded FP8 quantization with a new block-scale mode and dynamic,
hardware-aware kernel scheduling for better utilization and numerical
stability.
* Vectorized max-reduction and per-block scaling to accelerate
reductions and improve output scaling precision.
* Autotuner/CLI now exposes the FP8 block quantization option for
tuning.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Siyuan Fu <[email protected]>
Co-authored-by: Siyuan Fu <[email protected]>

Author: nekorobov

benchmarks/bench_trtllm_gen_fused_moe_autotuner.py

@@ -11,6 +11,8 @@
 from flashinfer.fused_moe import (
     trtllm_fp4_block_scale_moe,
     trtllm_fp8_per_tensor_scale_moe,
+    trtllm_fp8_block_scale_moe,
+    WeightLayout,
 )
 from flashinfer.autotuner import autotune
 from flashinfer.testing.utils import bench_gpu_time
@@ -21,15 +23,15 @@
 
 
 def fp8_quantize(x):
-    max = x.float().abs().nan_to_num().max()
+    max = x.abs().max().float()
     scale = FLOAT8_E4M3_MAX / max
     x = (x * scale).to(torch.float8_e4m3fn)
     return x, 1.0 / scale
 
 
 def bench_trtllm_gen_fused_moe_autotuner_fp8(
     tune_max_num_tokens: Optional[int],
-    quant_mode: Literal["Fp8-Per-Tensor"],
+    quant_mode: Literal["Fp8-Per-Tensor", "Fp8-Block"],
     num_tokens: int,
     num_experts: int,
     hidden_size: int,
@@ -41,55 +43,110 @@ def bench_trtllm_gen_fused_moe_autotuner_fp8(
     device = torch.device("cuda:0")
     enable_pdl = device_support_pdl(device)
     routing_logits = torch.rand(num_tokens, num_experts, device=device).to(
-        torch.bfloat16
+        torch.float32
     )
     hidden_states = torch.randn(num_tokens, hidden_size, device=device).to(
         torch.bfloat16
     )
+    routing_bias = torch.randn(num_experts, device="cuda", dtype=torch.bfloat16)
     w13 = torch.randn(
         num_experts, intermediate_size * 2, hidden_size, device=device
     ).to(torch.bfloat16)
     w2 = torch.randn(num_experts, hidden_size, intermediate_size, device=device).to(
         torch.bfloat16
     )
 
-    hidden_states, hidden_states_scale = fp8_quantize(hidden_states)
-    w13, w13_scale = fp8_quantize(w13)
-    w2, w2_scale = fp8_quantize(w2)
+    is_block_scale = quant_mode == "Fp8-Block"
+    if not is_block_scale:
+        hidden_states, hidden_states_scale = fp8_quantize(hidden_states)
+        w13, w13_scale = fp8_quantize(w13)
+        w2, w2_scale = fp8_quantize(w2)
+    else:
+        # block scale quantization is too slow, so we use per-tensor quantization for now
+        hidden_states, hidden_states_scale = fp8_quantize(hidden_states)
+        w13, w13_scale = fp8_quantize(w13)
+        w2, w2_scale = fp8_quantize(w2)
+        hidden_states_scale = torch.full(
+            (hidden_size // 128, num_tokens), hidden_states_scale.item(), device=device
+        )
+        w13_scale = torch.full(
+            (num_experts, intermediate_size * 2 // 128, hidden_size // 128),
+            w13_scale.item(),
+            device=device,
+        )
+        w2_scale = torch.full(
+            (num_experts, hidden_size // 128, intermediate_size // 128),
+            w2_scale.item(),
+            device=device,
+        )
 
-    output1_scale_scalar = torch.tensor(
-        [hidden_states_scale * w13_scale] * num_experts, device=device
+    output1_scale_scalar = (
+        torch.tensor([hidden_states_scale * w13_scale] * num_experts, device=device)
+        if not is_block_scale
+        else None
     )
-    output1_scales_gate_scalar = torch.ones(
-        num_experts, device=device, dtype=torch.float32
+    output1_scales_gate_scalar = (
+        torch.ones(num_experts, device=device, dtype=torch.float32)
+        if not is_block_scale
+        else None
     )
-    output2_scale_scalar = torch.tensor(
-        [hidden_states_scale * w2_scale] * num_experts, device=device
+    output2_scale_scalar = (
+        torch.tensor([hidden_states_scale * w2_scale] * num_experts, device=device)
+        if not is_block_scale
+        else None
     )
 
-    fn = lambda: trtllm_fp8_per_tensor_scale_moe(
-        routing_logits,
-        None,  # routing_bias
-        hidden_states,
-        w13,
-        output1_scale_scalar,
-        output1_scales_gate_scalar,
-        w2,
-        output2_scale_scalar,
-        num_experts,
-        top_k,
-        None,  # n_group
-        None,  # topk_group
-        intermediate_size,
-        0,  # local_expert_offset
-        num_experts,
-        1.0,  # routed_scaling_factor
-        False,  # use_routing_scales_on_input
-        None,
-        RoutingMethodType.TopK.value,
-        enable_pdl,
-        num_tokens if tune_max_num_tokens is None else tune_max_num_tokens,
-    )
+    if is_block_scale:
+        fn = lambda: trtllm_fp8_block_scale_moe(
+            routing_logits,
+            routing_bias,
+            hidden_states,
+            hidden_states_scale,
+            w13,
+            w13_scale,
+            w2,
+            w2_scale,
+            num_experts,
+            top_k,
+            8,  # n_group
+            4,  # topk_group
+            intermediate_size,
+            0,  # local_expert_offset
+            num_experts,
+            2.5,  # routed_scaling_factor
+            None,  # tile_tokens_dim
+            RoutingMethodType.DeepSeekV3.value,
+            True,  # use_shuffled_weight
+            WeightLayout.BlockMajorK.value,  # weight_layout
+            enable_pdl=enable_pdl,
+            tune_max_num_tokens=num_tokens
+            if tune_max_num_tokens is None
+            else tune_max_num_tokens,
+        )
+    else:
+        fn = lambda: trtllm_fp8_per_tensor_scale_moe(
+            routing_logits,
+            None,  # routing_bias
+            hidden_states,
+            w13,
+            output1_scale_scalar,
+            output1_scales_gate_scalar,
+            w2,
+            output2_scale_scalar,
+            num_experts,
+            top_k,
+            None,  # n_group
+            None,  # topk_group
+            intermediate_size,
+            0,  # local_expert_offset
+            num_experts,
+            1.0,  # routed_scaling_factor
+            False,  # use_routing_scales_on_input
+            None,  # tile_tokens_dim
+            RoutingMethodType.TopK.value,
+            enable_pdl,
+            num_tokens if tune_max_num_tokens is None else tune_max_num_tokens,
+        )
 
     def bench(do_autotune):
         with autotune(do_autotune):
@@ -135,6 +192,7 @@ def bench_trtllm_gen_fused_moe_autotuner_fp4(
             torch.tensor([448.0 * 6.0], device=device),
             sf_vec_size=16,
             sf_use_ue8m0=False,
+            is_sf_swizzled_layout=False,
         )
         hidden_states_scale = hidden_states_scale.view(torch.float8_e4m3fn).reshape(
             num_tokens, -1
@@ -263,7 +321,13 @@ def bench(do_autotune):
         "--quant-mode",
         type=str,
         default="MxFP4xMxFP8",
-        choices=["NvFP4xNvFP4", "MxFP4xMxFP8", "MxFP4xBf16", "Fp8-Per-Tensor"],
+        choices=[
+            "NvFP4xNvFP4",
+            "MxFP4xMxFP8",
+            "MxFP4xBf16",
+            "Fp8-Per-Tensor",
+            "Fp8-Block",
+        ],
         help="Quantization mode",
     )
     parser.add_argument("--num-tokens", type=int, default=512, help="Number of tokens")
@@ -288,7 +352,7 @@ def bench(do_autotune):
         "--iterations", type=int, default=100, help="Number of benchmark iterations"
     )
     args = parser.parse_args()
-    if args.quant_mode == "Fp8-Per-Tensor":
+    if args.quant_mode in ["Fp8-Per-Tensor", "Fp8-Block"]:
         bench_trtllm_gen_fused_moe_autotuner_fp8(
             args.tune_max_num_tokens,
             args.quant_mode,