chore: update vlm moe config and tune scripts (#30866)

benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py

@@ -140,13 +140,14 @@ class FlashInferFusedAllReduceParams:
         world_size: int,
         use_fp32_lamport: bool = False,
         max_token_num: int = 1024,
+        fp32_acc: bool = True,
     ):
         self.rank = rank
         self.world_size = world_size
         self.use_fp32_lamport = use_fp32_lamport
         self.trigger_completion_at_end = True
         self.launch_with_pdl = True
-        self.fp32_acc = True
+        self.fp32_acc = fp32_acc
         self.max_token_num = max_token_num
 
     def get_trtllm_fused_allreduce_kwargs(self):
@@ -1134,6 +1135,12 @@ def main():
         action="store_true",
         help="Disable oneshot mode for FlashInfer operations",
     )
+    parser.add_argument(
+        "--fp32-acc",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="Use FP32 accumulation in FlashInfer fused all-reduce (default: enabled).",
+    )
     parser.add_argument(
         "--warmup", type=int, default=5, help="Number of warmup iterations"
     )
@@ -1238,6 +1245,7 @@ def main():
                 rank=rank,
                 world_size=world_size,
                 max_token_num=max_num_token,
+                fp32_acc=args.fp32_acc,
             )
 
     # Collect all results for markdown export

benchmark/kernels/fused_moe_triton/common_utils.py

@@ -95,6 +95,7 @@ def get_model_config(
         "DeepseekV4ForCausalLM",
         "Glm4MoeForCausalLM",
         "GlmMoeDsaForCausalLM",
+        "KimiVLForConditionalGeneration",
         "MistralLarge3ForCausalLM",
     ]:
         E = (config.n_routed_experts // ep_size) + (

benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 # Adapted from https://github.com/vllm-project/vllm/blob/main/benchmarks/kernels/benchmark_moe.py
 import argparse
+import json
 import time
 from contextlib import nullcontext
 from datetime import datetime
@@ -242,7 +243,6 @@ def benchmark_config(
 
 @ray.remote(num_gpus=1)
 class BenchmarkWorker:
-
     def __init__(self, seed: int, server_args: ServerArgs) -> None:
         torch.set_default_device(get_device())
         torch.get_device_module().manual_seed_all(0)
@@ -398,10 +398,12 @@ def main(args: argparse.Namespace):
     use_int4_w4a16 = args.dtype == "int4_w4a16"
     per_channel_quant = args.per_channel_quant
 
-    if args.batch_size is None:
-        batch_sizes = get_default_batch_sizes()
-    else:
+    if args.batch_sizes is not None:
+        batch_sizes = args.batch_sizes
+    elif args.batch_size is not None:
         batch_sizes = [args.batch_size]
+    else:
+        batch_sizes = get_default_batch_sizes()
 
     ray.init()
     num_gpus = int(ray.available_resources()["GPU"])
@@ -419,9 +421,19 @@ def main(args: argparse.Namespace):
         return ray.get(outputs)
 
     if args.tune:
-        search_space = get_configs_compute_bound()
+        if args.search_space_file:
+            with open(args.search_space_file) as f:
+                search_space = json.load(f)
+            if not isinstance(search_space, list) or not all(
+                isinstance(config, dict) for config in search_space
+            ):
+                raise ValueError(
+                    "--search-space-file must contain a JSON list of configs"
+                )
+        else:
+            search_space = get_configs_compute_bound()
         if block_shape is not None:
-            block_n, block_k = block_shape[0], block_shape[1]
+            block_k = block_shape[1]
             search_space = [
                 config
                 for config in search_space
@@ -522,7 +534,18 @@ if __name__ == "__main__":
     )
     parser.add_argument("--seed", type=int, default=0)
     parser.add_argument("--batch-size", type=int, required=False)
+    parser.add_argument(
+        "--batch-sizes",
+        type=int,
+        nargs="+",
+        help="Tune or benchmark an explicit set of token counts in parallel.",
+    )
     parser.add_argument("--tune", action="store_true")
+    parser.add_argument(
+        "--search-space-file",
+        type=str,
+        help="JSON file containing an explicit list of Triton configs to evaluate with --tune.",
+    )
     parser.add_argument("--disable-shared-experts-fusion", action="store_true")
     args = parser.parse_args()