Profile Qwen MoE router on vLLM 0.20

2026-07-16 22:43:55 +08:00
parent b8f5f1dcf3
commit fcca259475
5 changed files with 266 additions and 3 deletions
--- a/runs/frontier-qwen30-vllm020-profile-v1/jobs_frontier_linear_full.toml
+++ b/runs/frontier-qwen30-vllm020-profile-v1/jobs_frontier_linear_full.toml
@@ -1,18 +1,18 @@
 version = 1

 [[jobs]]
-name = "qwen30-vllm020-frontier-linear-full-20260716-v1"
+name = "qwen30-vllm020-frontier-linear-full-20260716-v2-max-tokens"
 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 1620 bash run_frontier_linear_full.sh"
-artifacts = ["artifacts/frontier-linear-full-v1"]
+artifacts = ["artifacts/frontier-linear-full-v2"]

 [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/frontier-linear-full-v1"
+OUTPUT_ROOT = "/home/admin/cpfs/wjh/aituner/aituner-qwen30-vllm020-profile-fleet/artifacts/frontier-linear-full-v2"
 VENV_ROOT = "/tmp/wjh/venvs/vllm-0.20.0-cu129-profiler-v1"
 FRONTIER_ROOT = "/home/admin/cpfs/wjh/frontier-qwen30-vllm020-profile-v1/Frontier"
 VLLM_SOURCE = "/home/admin/cpfs/wjh/agentic-kv/third_party/vllm_v20_build"
--- a/runs/frontier-qwen30-vllm020-profile-v1/jobs_router_full.toml
+++ b/runs/frontier-qwen30-vllm020-profile-v1/jobs_router_full.toml
@@ -0,0 +1,18 @@
+version = 1
+
+[[jobs]]
+name = "qwen30-vllm020-router-full-20260716-v1"
+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-v1"]
+
+[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-v1"
+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"
--- a/runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_router.py
+++ b/runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_router.py
@@ -0,0 +1,195 @@
+#!/usr/bin/env python3
+"""Profile Qwen3's replicated MoE gate and fused top-k in vLLM 0.20."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import statistics
+import subprocess
+from pathlib import Path
+from typing import Any, Callable
+
+import torch
+import vllm
+
+
+VLLM_VERSION = "0.20.0"
+VLLM_COMMIT = "88d34c6409e9fb3c7b8ca0c04756f061d2099eb1"
+HIDDEN_DIM = 2048
+NUM_EXPERTS = 128
+TOP_K = 8
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--vllm-source", type=Path, required=True)
+    parser.add_argument("--model", type=Path, required=True)
+    parser.add_argument("--output", type=Path, required=True)
+    parser.add_argument("--num-tokens", type=int, nargs="+", required=True)
+    parser.add_argument("--warmup-iters", type=int, default=5)
+    parser.add_argument("--repeats", type=int, default=20)
+    parser.add_argument("--device", default="cuda:0")
+    return parser.parse_args()
+
+
+def git_head(repo: Path) -> str:
+    return subprocess.check_output(
+        ["git", "-C", str(repo), "rev-parse", "HEAD"], text=True
+    ).strip()
+
+
+def stats_ms(samples: list[float]) -> dict[str, float]:
+    return {
+        "min": min(samples),
+        "max": max(samples),
+        "mean": statistics.fmean(samples),
+        "median": statistics.median(samples),
+        "std": statistics.pstdev(samples),
+    }
+
+
+def measure_ms(
+    fn: Callable[[], Any], warmup_iters: int, repeats: int
+) -> tuple[Any, dict[str, float]]:
+    result = None
+    for _ in range(warmup_iters):
+        result = fn()
+    torch.accelerator.synchronize()
+    samples: list[float] = []
+    for _ in range(repeats):
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+        result = fn()
+        end.record()
+        torch.accelerator.synchronize()
+        samples.append(float(start.elapsed_time(end)))
+    return result, stats_ms(samples)
+
+
+def main() -> None:
+    args = parse_args()
+    if vllm.__version__ != VLLM_VERSION:
+        raise SystemExit(f"expected vLLM {VLLM_VERSION}, got {vllm.__version__}")
+    source_head = git_head(args.vllm_source)
+    if source_head != VLLM_COMMIT:
+        raise SystemExit(f"expected vLLM source {VLLM_COMMIT}, got {source_head}")
+
+    raw_model_config = json.loads(args.model.joinpath("config.json").read_text())
+    observed = {
+        "hidden_size": raw_model_config.get("hidden_size"),
+        "num_experts": raw_model_config.get("num_experts"),
+        "num_experts_per_tok": raw_model_config.get("num_experts_per_tok"),
+        "norm_topk_prob": raw_model_config.get("norm_topk_prob"),
+    }
+    expected = {
+        "hidden_size": HIDDEN_DIM,
+        "num_experts": NUM_EXPERTS,
+        "num_experts_per_tok": TOP_K,
+        "norm_topk_prob": True,
+    }
+    if observed != expected:
+        raise SystemExit(f"model contract mismatch: expected {expected}, got {observed}")
+
+    from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
+    from vllm.model_executor.layers.fused_moe import fused_topk
+    from vllm.model_executor.layers.linear import ReplicatedLinear
+
+    device = torch.device(args.device)
+    torch.accelerator.set_device_index(device)
+    torch.manual_seed(20260716)
+    model_config = ModelConfig(
+        model=str(args.model),
+        dtype="bfloat16",
+        max_model_len=8192,
+        skip_tokenizer_init=True,
+        generation_config="vllm",
+    )
+
+    rows: list[dict[str, Any]] = []
+    with set_current_vllm_config(VllmConfig(model_config=model_config)):
+        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:
+            hidden = torch.empty(
+                (num_tokens, HIDDEN_DIM), device=device, dtype=torch.bfloat16
+            ).uniform_(-0.1, 0.1)
+            logits, gate_time = measure_ms(
+                lambda: gate(hidden)[0], args.warmup_iters, args.repeats
+            )
+            topk_result, topk_time = measure_ms(
+                lambda: fused_topk(hidden, logits, TOP_K, renormalize=True),
+                args.warmup_iters,
+                args.repeats,
+            )
+
+            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)
+
+            combined_result, combined_time = measure_ms(
+                gate_and_topk, args.warmup_iters, args.repeats
+            )
+            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)}")
+            if topk_ids.shape != (num_tokens, TOP_K):
+                raise SystemExit(f"invalid top-k shape: {tuple(topk_ids.shape)}")
+            torch.testing.assert_close(
+                topk_weights.sum(dim=-1),
+                torch.ones(num_tokens, device=device),
+                atol=1e-5,
+                rtol=1e-5,
+            )
+            torch.testing.assert_close(combined_weights, topk_weights)
+            torch.testing.assert_close(combined_ids, topk_ids)
+            additive_median = gate_time["median"] + topk_time["median"]
+            row = {
+                "num_tokens": num_tokens,
+                "gate_linear_time_ms": gate_time,
+                "routing_topk_time_ms": topk_time,
+                "gate_plus_topk_time_ms": combined_time,
+                "median_nonadditivity_ratio": (
+                    combined_time["median"] / additive_median
+                    if additive_median > 0
+                    else 1.0
+                ),
+            }
+            rows.append(row)
+            print(json.dumps(row, sort_keys=True), flush=True)
+
+    payload = {
+        "schema_version": "qwen30_vllm020_router_raw.v1",
+        "environment": {
+            "vllm_version": vllm.__version__,
+            "vllm_source_commit": source_head,
+            "torch_version": torch.__version__,
+            "torch_cuda": torch.version.cuda,
+            "gpu": torch.cuda.get_device_name(device),
+            "model": str(args.model),
+            "dtype": "bfloat16",
+            "gate_replication": "replicated_across_tp",
+            "top_k": TOP_K,
+            "norm_topk_prob": True,
+        },
+        "measurement_scope": (
+            "vLLM ReplicatedLinear gate and fused_topk; measured separately and "
+            "as the actual sequential router path"
+        ),
+        "rows": rows,
+    }
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
+
+
+if __name__ == "__main__":
+    main()
--- a/runs/frontier-qwen30-vllm020-profile-v1/run_frontier_linear_full.sh
+++ b/runs/frontier-qwen30-vllm020-profile-v1/run_frontier_linear_full.sh
@@ -47,6 +47,7 @@ timeout --signal=TERM --kill-after=30s 1380 \
  --device h20 \
  --models qwen3-a3b-30b-moe \
  --num_tensor_parallel_workers 1 2 4 \
+  --max_tokens 8192 \
  --num_tokens_list "${TOKENS[@]}" \
  --profile_method cuda_event \
  --precision BF16 \
--- a/runs/frontier-qwen30-vllm020-profile-v1/run_router_full.sh
+++ b/runs/frontier-qwen30-vllm020-profile-v1/run_router_full.sh
@@ -0,0 +1,49 @@
+#!/usr/bin/env bash
+
+set -euo pipefail
+
+OUTPUT_ROOT="${OUTPUT_ROOT:?OUTPUT_ROOT must be set}"
+VENV_ROOT="${VENV_ROOT:-/tmp/wjh/venvs/vllm-0.20.0-cu129-profiler-v1}"
+VLLM_SOURCE="${VLLM_SOURCE:-/home/admin/cpfs/wjh/agentic-kv/third_party/vllm_v20_build}"
+MODEL="${MODEL:-/home/admin/cpfs/wjh/models/Qwen/Qwen3-30B-A3B}"
+TOKENS=(1 8 16 32 64 128 256 512 1024 2048 4096 8192)
+mkdir -p "${OUTPUT_ROOT}/logs" "${OUTPUT_ROOT}/provenance" "${OUTPUT_ROOT}/raw"
+exec > >(tee -a "${OUTPUT_ROOT}/logs/full.log") 2>&1
+
+IFS=',' read -r -a GPU_IDS <<< "${CUDA_VISIBLE_DEVICES:?fleet GPU is required}"
+if [[ "${#GPU_IDS[@]}" -ne 1 ]]; then
+  echo "ERROR: expected exactly one GPU, got ${CUDA_VISIBLE_DEVICES}" >&2
+  exit 1
+fi
+
+echo "PROFILE_LAUNCH_ECHO host=$(hostname) gpu=${CUDA_VISIBLE_DEVICES} model=${MODEL} runtime=vLLM-0.20.0+cu129 operators=ReplicatedLinear,fused_topk tokens=${TOKENS[*]} dtype=BF16 output=${OUTPUT_ROOT} expected_wall=2-6m hard_wall=900s hard_gpu_cap=0.25_H20h"
+date -u +"START_UTC=%Y-%m-%dT%H:%M:%SZ"
+nvidia-smi --query-gpu=index,name,driver_version,memory.used,utilization.gpu --format=csv,noheader
+
+test "$(git -C "${VLLM_SOURCE}" rev-parse HEAD)" = "88d34c6409e9fb3c7b8ca0c04756f061d2099eb1"
+test -s "${MODEL}/config.json"
+git rev-parse HEAD > "${OUTPUT_ROOT}/provenance/aituner.commit"
+git -C "${VLLM_SOURCE}" rev-parse HEAD > "${OUTPUT_ROOT}/provenance/vllm-source.commit"
+sha256sum profile_vllm020_router.py run_router_full.sh \
+  > "${OUTPUT_ROOT}/provenance/source.sha256"
+sha256sum "${MODEL}/config.json" > "${OUTPUT_ROOT}/provenance/model-config.sha256"
+uv pip freeze --python "${VENV_ROOT}/bin/python" \
+  > "${OUTPUT_ROOT}/provenance/pip-freeze.txt"
+nvidia-smi --query-gpu=index,uuid,name,driver_version,memory.total \
+  --format=csv,noheader > "${OUTPUT_ROOT}/provenance/gpus.csv"
+printf '%s\n' "${TOKENS[@]}" > "${OUTPUT_ROOT}/provenance/tokens.txt"
+
+timeout --signal=TERM --kill-after=30s 780 \
+  "${VENV_ROOT}/bin/python" profile_vllm020_router.py \
+  --vllm-source "${VLLM_SOURCE}" \
+  --model "${MODEL}" \
+  --output "${OUTPUT_ROOT}/raw/router.json" \
+  --num-tokens "${TOKENS[@]}" \
+  --warmup-iters 5 \
+  --repeats 20
+
+test -s "${OUTPUT_ROOT}/raw/router.json"
+sha256sum "${OUTPUT_ROOT}/raw/router.json" "${OUTPUT_ROOT}/provenance"/* \
+  > "${OUTPUT_ROOT}/artifacts.sha256"
+date -u +"END_UTC=%Y-%m-%dT%H:%M:%SZ"
+echo "ROUTER_FULL_COMPLETE cases=${#TOKENS[@]}"