better profiling for trtllm-gen moe benchmark

Signed-off-by: Siyuan Fu <siyuanf@nvidia.com>

Author: IwakuraRein

benchmarks/bench_trtllm_gen_fused_moe_autotuner.py

@@ -2,6 +2,7 @@
 from typing import Optional, Literal
 import torch
 import numpy as np
+from functools import partial
 from flashinfer import (
     RoutingMethodType,
     GatedActType,
@@ -143,27 +144,26 @@ def bench_trtllm_gen_fused_moe_autotuner_fp8(
             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
-            RoutingMethodType.TopK.value,
-            enable_pdl,
-            num_tokens if tune_max_num_tokens is None else tune_max_num_tokens,
+        fn = partial(
+            trtllm_fp8_per_tensor_scale_moe,
+            routing_logits=routing_logits,
+            routing_bias=None,  # routing_bias
+            output1_scale_scalar=output1_scale_scalar,
+            output1_scales_gate_scalar=output1_scales_gate_scalar,
+            output2_scale_scalar=output2_scale_scalar,
+            num_experts=num_experts,
+            top_k=top_k,
+            n_group=None,  # n_group
+            topk_group=None,  # topk_group
+            intermediate_size=intermediate_size,
+            local_expert_offset=0,  # local_expert_offset
+            routed_scaling_factor=1.0,  # routed_scaling_factor
+            use_routing_scales_on_input=False,  # use_routing_scales_on_input
+            routing_method_type=RoutingMethodType.TopK.value,
+            enable_pdl=enable_pdl,
+            tune_max_num_tokens=num_tokens
+            if tune_max_num_tokens is None
+            else tune_max_num_tokens,
         )
 
     def bench(do_autotune):
@@ -173,6 +173,14 @@ def bench(do_autotune):
             fn,
             dry_run_iters=warmups,
             repeat_iters=iterations,
+            enable_cupti=True,
+            use_cuda_graph=True,
+            input_kwargs={
+                "hidden_states": hidden_states,
+                "gemm1_weights": w13,
+                "gemm2_weights": w2,
+            },
+            cold_l2_cache=True,
         )
         median_ms = np.median(ms_list)
         return median_ms
@@ -280,37 +288,31 @@ def bench_trtllm_gen_fused_moe_autotuner_fp4(
     output2_scale_scalar = torch.tensor(
         [hidden_states_global_scale * w2_global_scale] * num_experts, device=device
     )
-    fn = lambda: trtllm_fp4_block_scale_moe(
-        routing_logits,
-        None,  # routing_bias
-        hidden_states,
-        hidden_states_scale,
-        w13,
-        w13_scale,
-        bias13,
-        None,  # gemm1_alpha
-        None,  # gemm1_beta
-        None,  # gemm1_clamp_limit
-        w2,
-        w2_scale,
-        bias2,
-        output1_scale_scalar,
-        output1_scale_gate_scalar,
-        output2_scale_scalar,
-        num_experts,
-        top_k,
-        None,  # n_group
-        None,  # topk_group
-        intermediate_size,
-        0,  # local_expert_offset
-        num_experts,
-        None,  # routed_scaling_factor
-        RoutingMethodType.Renormalize.value,
-        True,
-        enable_pdl,
-        GatedActType.SwiGlu.value,  # gated_act_type
-        None,
-        num_tokens if tune_max_num_tokens is None else tune_max_num_tokens,
+    fn = partial(
+        trtllm_fp4_block_scale_moe,
+        routing_logits=routing_logits,
+        routing_bias=None,  # routing_bias
+        gemm1_alpha=None,  # gemm1_alpha
+        gemm1_beta=None,  # gemm1_beta
+        gemm1_clamp_limit=None,  # gemm1_clamp_limit
+        output1_scale_scalar=output1_scale_scalar,
+        output1_scale_gate_scalar=output1_scale_gate_scalar,
+        output2_scale_scalar=output2_scale_scalar,
+        num_experts=num_experts,
+        top_k=top_k,
+        n_group=None,  # n_group
+        topk_group=None,  # topk_group
+        intermediate_size=intermediate_size,
+        local_expert_offset=0,  # local_expert_offset
+        routed_scaling_factor=None,  # routed_scaling_factor
+        routing_method_type=RoutingMethodType.Renormalize.value,
+        do_finalize=True,
+        enable_pdl=enable_pdl,
+        gated_act_type=GatedActType.SwiGlu.value,  # gated_act_type
+        output=None,
+        tune_max_num_tokens=num_tokens
+        if tune_max_num_tokens is None
+        else tune_max_num_tokens,
     )
 
     def bench(do_autotune):
@@ -320,6 +322,18 @@ def bench(do_autotune):
             fn,
             dry_run_iters=warmups,
             repeat_iters=iterations,
+            enable_cupti=True,
+            use_cuda_graph=True,
+            input_kwargs={
+                "hidden_states": hidden_states,
+                "gemm1_weights": w13,
+                "gemm1_weights_scale": w13_scale,
+                "gemm2_weights": w2,
+                "gemm2_weights_scale": w2_scale,
+                "gemm1_bias": bias13,
+                "gemm2_bias": bias2,
+            },
+            cold_l2_cache=True,
         )
         median_ms = np.median(ms_list)
         return median_ms
@@ -370,29 +384,27 @@ def bench_trtllm_gen_fused_moe_autotuner_mxint4(
         intermediate_size // 32,
     )
 
-    fn = lambda: trtllm_mxint4_block_scale_moe(
-        routing_logits,
-        routing_bias,
-        hidden_states,
-        w13,
-        w13_scale,
-        None,  # gemm1_alpha
-        None,  # gemm1_beta
-        None,  # gemm1_clamp_limit
-        w2,
-        w2_scale,
-        num_experts,
-        top_k,
-        1,  # n_group
-        1,  # topk_group
-        intermediate_size,
-        0,  # local_expert_offset
-        num_experts,
-        None,  # routed_scaling_factor
-        RoutingMethodType.DeepSeekV3.value,
-        enable_pdl,
-        None,
-        num_tokens if tune_max_num_tokens is None else tune_max_num_tokens,
+    fn = partial(
+        trtllm_mxint4_block_scale_moe,
+        routing_logits=routing_logits,
+        routing_bias=routing_bias,
+        hidden_states=hidden_states,
+        gemm1_alpha=None,  # gemm1_alpha
+        gemm1_beta=None,  # gemm1_beta
+        gemm1_clamp_limit=None,  # gemm1_clamp_limit
+        num_experts=num_experts,
+        top_k=top_k,
+        n_group=1,  # n_group
+        topk_group=1,  # topk_group
+        intermediate_size=intermediate_size,
+        local_expert_offset=0,  # local_expert_offset
+        routed_scaling_factor=None,  # routed_scaling_factor
+        routing_method_type=RoutingMethodType.DeepSeekV3.value,
+        enable_pdl=enable_pdl,
+        output=None,
+        tune_max_num_tokens=num_tokens
+        if tune_max_num_tokens is None
+        else tune_max_num_tokens,
     )
 
     def bench(do_autotune):
@@ -402,6 +414,16 @@ def bench(do_autotune):
             fn,
             dry_run_iters=warmups,
             repeat_iters=iterations,
+            enable_cupti=True,
+            use_cuda_graph=True,
+            input_kwargs={
+                "hidden_states": hidden_states,
+                "gemm1_weights": w13,
+                "gemm1_weights_scale": w13_scale,
+                "gemm2_weights": w2,
+                "gemm2_weights_scale": w2_scale,
+            },
+            cold_l2_cache=True,
         )
         median_ms = np.median(ms_list)
         return median_ms