gahow/aituner
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Initialize TP context for router profiling

Browse Source
This commit is contained in:
Gahow Wang
2026-07-16 22:54:17 +08:00
parent 414838a799
commit 07b1eb4b75
2 changed files with 85 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
runs/frontier-qwen30-vllm020-profile-v1/jobs_router_full.toml
View File

@@ -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"

34
runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_router.py
View File

@@ -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,14 +115,24 @@ def main() -> None:
)
rows: list[dict[str, Any]] = []
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener:
listener.bind(("127.0.0.1", 0))
distributed_init_method = f"tcp://127.0.0.1:{listener.getsockname()[1]}"
init_distributed_environment(
world_size=1,
rank=0,
local_rank=0,
distributed_init_method=distributed_init_method,
)
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",
disable_tp=True,
).to(device=device, dtype=torch.bfloat16)
gate.weight.data.uniform_(-0.01, 0.01)
@@ -132,9 +149,13 @@ def main() -> None:
args.repeats,
)
def gate_and_topk() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
def gate_and_topk() -> tuple[
torch.Tensor, torch.Tensor, torch.Tensor
]:
current_logits, _ = gate(hidden)
return fused_topk(hidden, current_logits, TOP_K, renormalize=True)
return fused_topk(
hidden, current_logits, TOP_K, renormalize=True
)
combined_result, combined_time = measure_ms(
gate_and_topk, args.warmup_iters, args.repeats
@@ -142,7 +163,9 @@ def main() -> None:
topk_weights, topk_ids, _ = topk_result
combined_weights, combined_ids, _ = combined_result
if logits.shape != (num_tokens, NUM_EXPERTS):
raise SystemExit(f"invalid gate output shape: {tuple(logits.shape)}")
raise SystemExit(
f"invalid gate output shape: {tuple(logits.shape)}"
)
if topk_ids.shape != (num_tokens, TOP_K):
raise SystemExit(f"invalid top-k shape: {tuple(topk_ids.shape)}")
torch.testing.assert_close(
@@ -167,6 +190,9 @@ def main() -> None:
}
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",