Initialize TP context for router profiling

runs/frontier-qwen30-vllm020-profile-v1/jobs_router_full.toml

@@ -1,18 +1,18 @@
 version = 1
 
 [[jobs]]
-name = "qwen30-vllm020-router-full-20260716-v2-disable-tp-init"
+name = "qwen30-vllm020-router-full-20260716-v3-tp-context"
 gpus = 1
 gpu_model = "H20"
 hosts = ["dash0"]
 command = "cd /home/admin/cpfs/wjh/aituner/aituner-qwen30-vllm020-profile-v1/runs/frontier-qwen30-vllm020-profile-v1 && timeout --signal=TERM --kill-after=30s 1020 bash run_router_full.sh"
-artifacts = ["artifacts/router-full-v2"]
+artifacts = ["artifacts/router-full-v3"]
 
 [jobs.env]
 HOME = "/tmp/wjh"
 XDG_CACHE_HOME = "/tmp/wjh/.cache"
 VLLM_CACHE_ROOT = "/tmp/wjh/.cache/vllm"
-OUTPUT_ROOT = "/home/admin/cpfs/wjh/aituner/aituner-qwen30-vllm020-profile-fleet/artifacts/router-full-v2"
+OUTPUT_ROOT = "/home/admin/cpfs/wjh/aituner/aituner-qwen30-vllm020-profile-fleet/artifacts/router-full-v3"
 VENV_ROOT = "/tmp/wjh/venvs/vllm-0.20.0-cu129-profiler-v1"
 VLLM_SOURCE = "/home/admin/cpfs/wjh/agentic-kv/third_party/vllm_v20_build"
 MODEL = "/home/admin/cpfs/wjh/models/Qwen/Qwen3-30B-A3B"

runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_router.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 
 import argparse
 import json
+import socket
 import statistics
 import subprocess
 from pathlib import Path
@@ -93,6 +94,12 @@ def main() -> None:
         raise SystemExit(f"model contract mismatch: expected {expected}, got {observed}")
 
     from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
+    from vllm.distributed import (
+        destroy_distributed_environment,
+        destroy_model_parallel,
+        init_distributed_environment,
+        initialize_model_parallel,
+    )
     from vllm.model_executor.layers.fused_moe import fused_topk
     from vllm.model_executor.layers.linear import ReplicatedLinear
 
@@ -108,65 +115,84 @@ def main() -> None:
     )
 
     rows: list[dict[str, Any]] = []
-    with set_current_vllm_config(VllmConfig(model_config=model_config)):
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener:
-        gate = ReplicatedLinear(
+        listener.bind(("127.0.0.1", 0))
-            HIDDEN_DIM,
+        distributed_init_method = f"tcp://127.0.0.1:{listener.getsockname()[1]}"
-            NUM_EXPERTS,
+    init_distributed_environment(
-            bias=False,
+        world_size=1,
-            quant_config=None,
+        rank=0,
-            prefix="model.layers.0.mlp.gate",
+        local_rank=0,
-            disable_tp=True,
+        distributed_init_method=distributed_init_method,
-        ).to(device=device, dtype=torch.bfloat16)
+    )
-        gate.weight.data.uniform_(-0.01, 0.01)
+    try:
+        with set_current_vllm_config(VllmConfig(model_config=model_config)):
+            initialize_model_parallel(tensor_model_parallel_size=1)
+            gate = ReplicatedLinear(
+                HIDDEN_DIM,
+                NUM_EXPERTS,
+                bias=False,
+                quant_config=None,
+                prefix="model.layers.0.mlp.gate",
+            ).to(device=device, dtype=torch.bfloat16)
+            gate.weight.data.uniform_(-0.01, 0.01)
 
-        for num_tokens in args.num_tokens:
+            for num_tokens in args.num_tokens:
-            hidden = torch.empty(
+                hidden = torch.empty(
-                (num_tokens, HIDDEN_DIM), device=device, dtype=torch.bfloat16
+                    (num_tokens, HIDDEN_DIM), device=device, dtype=torch.bfloat16
-            ).uniform_(-0.1, 0.1)
+                ).uniform_(-0.1, 0.1)
-            logits, gate_time = measure_ms(
+                logits, gate_time = measure_ms(
-                lambda: gate(hidden)[0], args.warmup_iters, args.repeats
+                    lambda: gate(hidden)[0], args.warmup_iters, args.repeats
-            )
+                )
-            topk_result, topk_time = measure_ms(
+                topk_result, topk_time = measure_ms(
-                lambda: fused_topk(hidden, logits, TOP_K, renormalize=True),
+                    lambda: fused_topk(hidden, logits, TOP_K, renormalize=True),
-                args.warmup_iters,
+                    args.warmup_iters,
-                args.repeats,
+                    args.repeats,
-            )
+                )
 
-            def gate_and_topk() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+                def gate_and_topk() -> tuple[
-                current_logits, _ = gate(hidden)
+                    torch.Tensor, torch.Tensor, torch.Tensor
-                return fused_topk(hidden, current_logits, TOP_K, renormalize=True)
+                ]:
+                    current_logits, _ = gate(hidden)
+                    return fused_topk(
+                        hidden, current_logits, TOP_K, renormalize=True
+                    )
 
-            combined_result, combined_time = measure_ms(
+                combined_result, combined_time = measure_ms(
-                gate_and_topk, args.warmup_iters, args.repeats
+                    gate_and_topk, args.warmup_iters, args.repeats
-            )
+                )
-            topk_weights, topk_ids, _ = topk_result
+                topk_weights, topk_ids, _ = topk_result
-            combined_weights, combined_ids, _ = combined_result
+                combined_weights, combined_ids, _ = combined_result
-            if logits.shape != (num_tokens, NUM_EXPERTS):
+                if logits.shape != (num_tokens, NUM_EXPERTS):
-                raise SystemExit(f"invalid gate output shape: {tuple(logits.shape)}")
+                    raise SystemExit(
-            if topk_ids.shape != (num_tokens, TOP_K):
+                        f"invalid gate output shape: {tuple(logits.shape)}"
-                raise SystemExit(f"invalid top-k shape: {tuple(topk_ids.shape)}")
+                    )
-            torch.testing.assert_close(
+                if topk_ids.shape != (num_tokens, TOP_K):
-                topk_weights.sum(dim=-1),
+                    raise SystemExit(f"invalid top-k shape: {tuple(topk_ids.shape)}")
-                torch.ones(num_tokens, device=device),
+                torch.testing.assert_close(
-                atol=1e-5,
+                    topk_weights.sum(dim=-1),
-                rtol=1e-5,
+                    torch.ones(num_tokens, device=device),
-            )
+                    atol=1e-5,
-            torch.testing.assert_close(combined_weights, topk_weights)
+                    rtol=1e-5,
-            torch.testing.assert_close(combined_ids, topk_ids)
+                )
-            additive_median = gate_time["median"] + topk_time["median"]
+                torch.testing.assert_close(combined_weights, topk_weights)
-            row = {
+                torch.testing.assert_close(combined_ids, topk_ids)
-                "num_tokens": num_tokens,
+                additive_median = gate_time["median"] + topk_time["median"]
-                "gate_linear_time_ms": gate_time,
+                row = {
-                "routing_topk_time_ms": topk_time,
+                    "num_tokens": num_tokens,
-                "gate_plus_topk_time_ms": combined_time,
+                    "gate_linear_time_ms": gate_time,
-                "median_nonadditivity_ratio": (
+                    "routing_topk_time_ms": topk_time,
-                    combined_time["median"] / additive_median
+                    "gate_plus_topk_time_ms": combined_time,
-                    if additive_median > 0
+                    "median_nonadditivity_ratio": (
-                    else 1.0
+                        combined_time["median"] / additive_median
-                ),
+                        if additive_median > 0
-            }
+                        else 1.0
-            rows.append(row)
+                    ),
-            print(json.dumps(row, sort_keys=True), flush=True)
+                }
+                rows.append(row)
+                print(json.dumps(row, sort_keys=True), flush=True)
+    finally:
+        destroy_model_parallel()
+        destroy_distributed_environment()
 
     payload = {
         "schema_version": "qwen30_vllm020_router_raw.v1",