Add exact vLLM 0.20 MoE profile smoke

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

@@ -0,0 +1,18 @@
+version = 1
+
+[[jobs]]
+name = "qwen30-vllm020-moe-smoke-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_moe_smoke.sh"
+artifacts = ["artifacts/moe-smoke"]
+
+[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/moe-smoke"
+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_moe.py

@@ -0,0 +1,361 @@
+#!/usr/bin/env python3
+"""Profile vLLM 0.20's Qwen3-30B unquantized MoE kernel at TP-local shapes."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import statistics
+import subprocess
+from pathlib import Path
+from typing import Any
+
+import torch
+import vllm
+
+
+VLLM_VERSION = "0.20.0"
+VLLM_COMMIT = "88d34c6409e9fb3c7b8ca0c04756f061d2099eb1"
+HIDDEN_DIM = 2048
+INTERMEDIATE_DIM = 768
+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("--tp", type=int, choices=(1, 2, 4), nargs="+", default=[1, 2, 4])
+    parser.add_argument("--num-tokens", type=int, nargs="+", default=[8])
+    parser.add_argument(
+        "--routing-modes",
+        choices=("uniform_random_logits", "hotset8"),
+        nargs="+",
+        default=["uniform_random_logits"],
+    )
+    parser.add_argument("--warmup-iters", type=int, default=3)
+    parser.add_argument("--repeats", type=int, default=5)
+    parser.add_argument("--device", default="cuda:0")
+    parser.add_argument("--check-reference", action="store_true")
+    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 routing_inputs(
+    mode: str, num_tokens: int, device: torch.device
+) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]:
+    from vllm.model_executor.layers.fused_moe import fused_topk
+
+    if mode == "uniform_random_logits":
+        logits = torch.randn(
+            (num_tokens, NUM_EXPERTS), device=device, dtype=torch.bfloat16
+        )
+        hidden_for_topk = torch.empty(
+            (num_tokens, HIDDEN_DIM), device=device, dtype=torch.bfloat16
+        )
+        weights, ids, _ = fused_topk(
+            hidden_for_topk,
+            logits,
+            TOP_K,
+            renormalize=True,
+        )
+    elif mode == "hotset8":
+        ids = torch.arange(TOP_K, device=device, dtype=torch.int32).repeat(
+            num_tokens, 1
+        )
+        weights = torch.full(
+            (num_tokens, TOP_K),
+            1.0 / TOP_K,
+            device=device,
+            dtype=torch.float32,
+        )
+    else:
+        raise ValueError(mode)
+
+    counts = torch.bincount(ids.flatten().to(torch.int64), minlength=NUM_EXPERTS)
+    counts_cpu = counts.cpu().tolist()
+    mean_load = num_tokens * TOP_K / NUM_EXPERTS
+    variance = sum((count - mean_load) ** 2 for count in counts_cpu) / NUM_EXPERTS
+    return weights, ids, {
+        "active_experts": sum(count > 0 for count in counts_cpu),
+        "min_tokens_per_expert": min(counts_cpu),
+        "max_tokens_per_expert": max(counts_cpu),
+        "load_cv": math.sqrt(variance) / mean_load if mean_load else 0.0,
+        "counts": counts_cpu,
+    }
+
+
+def reference_partial_output(
+    hidden: torch.Tensor,
+    w13_original: torch.Tensor,
+    w2: torch.Tensor,
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+) -> torch.Tensor:
+    output = torch.zeros_like(hidden)
+    for token in range(hidden.shape[0]):
+        for route in range(TOP_K):
+            expert = int(topk_ids[token, route])
+            gate_up = torch.mv(w13_original[expert], hidden[token])
+            gate, up = gate_up.chunk(2)
+            activated = torch.nn.functional.silu(gate) * up
+            expert_output = torch.mv(w2[expert], activated)
+            output[token].add_(
+                expert_output * topk_weights[token, route].to(expert_output.dtype)
+            )
+    return output
+
+
+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}")
+    model_config = json.loads(args.model.joinpath("config.json").read_text())
+    expected_model = {
+        "hidden_size": HIDDEN_DIM,
+        "moe_intermediate_size": INTERMEDIATE_DIM,
+        "num_experts": NUM_EXPERTS,
+        "num_experts_per_tok": TOP_K,
+        "norm_topk_prob": True,
+        "torch_dtype": "bfloat16",
+    }
+    observed_model = {key: model_config.get(key) for key in expected_model}
+    if observed_model != expected_model:
+        raise SystemExit(
+            f"model contract mismatch: expected {expected_model}, got {observed_model}"
+        )
+
+    from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
+    from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+    from vllm.model_executor.layers.fused_moe.config import (
+        FUSED_MOE_UNQUANTIZED_CONFIG,
+        FusedMoEConfig,
+        FusedMoEParallelConfig,
+        RoutingMethodType,
+    )
+    from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
+        UnquantizedMoeBackend,
+        convert_to_unquantized_kernel_format,
+        make_unquantized_moe_kernel,
+        select_unquantized_moe_backend,
+    )
+    from vllm.utils.math_utils import next_power_of_2
+    from vllm.v1.worker.workspace import init_workspace_manager
+
+    device = torch.device(args.device)
+    torch.accelerator.set_device_index(device)
+    torch.manual_seed(20260716)
+    init_workspace_manager(args.device)
+    max_num_tokens = next_power_of_2(max(args.num_tokens))
+
+    rows: list[dict[str, Any]] = []
+    for tp in args.tp:
+        parallel = FusedMoEParallelConfig(
+            tp_size=tp,
+            tp_rank=0,
+            pcp_size=1,
+            pcp_rank=0,
+            dp_size=1,
+            dp_rank=0,
+            ep_size=1,
+            ep_rank=0,
+            sp_size=1,
+            use_ep=False,
+            all2all_backend="allgather_reducescatter",
+            enable_eplb=False,
+        )
+        moe_config = FusedMoEConfig(
+            num_experts=NUM_EXPERTS,
+            experts_per_token=TOP_K,
+            hidden_dim=HIDDEN_DIM,
+            intermediate_size_per_partition=INTERMEDIATE_DIM // tp,
+            num_local_experts=NUM_EXPERTS,
+            num_logical_experts=NUM_EXPERTS,
+            activation=MoEActivation.SILU,
+            device=device,
+            routing_method=RoutingMethodType.Renormalize,
+            moe_parallel_config=parallel,
+            in_dtype=torch.bfloat16,
+            max_num_tokens=max_num_tokens,
+        )
+        vllm_config = VllmConfig(
+            parallel_config=ParallelConfig(tensor_parallel_size=tp)
+        )
+        with set_current_vllm_config(vllm_config):
+            backend, experts_cls = select_unquantized_moe_backend(moe_config)
+            if backend != UnquantizedMoeBackend.FLASHINFER_CUTLASS:
+                raise SystemExit(
+                    "runtime backend mismatch: expected FlashInfer CUTLASS, "
+                    f"got {backend.value} at TP={tp}"
+                )
+            if experts_cls is None:
+                raise SystemExit(f"missing experts class for {backend.value}")
+
+            w13_original = torch.empty(
+                (NUM_EXPERTS, 2 * (INTERMEDIATE_DIM // tp), HIDDEN_DIM),
+                device=device,
+                dtype=torch.bfloat16,
+            ).uniform_(-0.01, 0.01)
+            w2 = torch.empty(
+                (NUM_EXPERTS, HIDDEN_DIM, INTERMEDIATE_DIM // tp),
+                device=device,
+                dtype=torch.bfloat16,
+            ).uniform_(-0.01, 0.01)
+
+            class Layer:
+                pass
+
+            layer = Layer()
+            layer.moe_config = moe_config
+            w13_kernel, w2_kernel = convert_to_unquantized_kernel_format(
+                backend,
+                layer=layer,
+                w13_weight=w13_original,
+                w2_weight=w2,
+            )
+            kernel = make_unquantized_moe_kernel(
+                quant_config=FUSED_MOE_UNQUANTIZED_CONFIG,
+                moe_config=moe_config,
+                backend=backend,
+                experts_cls=experts_cls,
+            )
+
+            reference_checked = False
+            for routing_mode in args.routing_modes:
+                for num_tokens in args.num_tokens:
+                    hidden = torch.empty(
+                        (num_tokens, HIDDEN_DIM),
+                        device=device,
+                        dtype=torch.bfloat16,
+                    ).uniform_(-0.1, 0.1)
+                    topk_weights, topk_ids, load = routing_inputs(
+                        routing_mode, num_tokens, device
+                    )
+
+                    for _ in range(args.warmup_iters):
+                        output = kernel.apply(
+                            hidden_states=hidden,
+                            w1=w13_kernel,
+                            w2=w2_kernel,
+                            topk_weights=topk_weights,
+                            topk_ids=topk_ids,
+                            activation=MoEActivation.SILU,
+                            global_num_experts=NUM_EXPERTS,
+                            expert_map=None,
+                            apply_router_weight_on_input=False,
+                        )
+                    torch.accelerator.synchronize()
+
+                    samples: list[float] = []
+                    for _ in range(args.repeats):
+                        start = torch.cuda.Event(enable_timing=True)
+                        end = torch.cuda.Event(enable_timing=True)
+                        start.record()
+                        output = kernel.apply(
+                            hidden_states=hidden,
+                            w1=w13_kernel,
+                            w2=w2_kernel,
+                            topk_weights=topk_weights,
+                            topk_ids=topk_ids,
+                            activation=MoEActivation.SILU,
+                            global_num_experts=NUM_EXPERTS,
+                            expert_map=None,
+                            apply_router_weight_on_input=False,
+                        )
+                        end.record()
+                        torch.accelerator.synchronize()
+                        samples.append(float(start.elapsed_time(end)))
+
+                    if output.shape != hidden.shape or not torch.isfinite(output).all():
+                        raise SystemExit(
+                            f"invalid MoE output TP={tp} M={num_tokens} mode={routing_mode}"
+                        )
+                    if args.check_reference and not reference_checked:
+                        check_tokens = min(2, num_tokens)
+                        reference = reference_partial_output(
+                            hidden[:check_tokens],
+                            w13_original,
+                            w2,
+                            topk_weights[:check_tokens],
+                            topk_ids[:check_tokens],
+                        )
+                        torch.testing.assert_close(
+                            output[:check_tokens], reference, atol=0.03, rtol=0.03
+                        )
+                        reference_checked = True
+
+                    row = {
+                        "tensor_parallel_size": tp,
+                        "num_tokens": num_tokens,
+                        "routing_mode": routing_mode,
+                        "backend": backend.value,
+                        "intermediate_size_per_partition": INTERMEDIATE_DIM // tp,
+                        "output_is_reduced": kernel.output_is_reduced(),
+                        "time_ms": stats_ms(samples),
+                        "routing_load": load,
+                    }
+                    rows.append(row)
+                    print(
+                        json.dumps(
+                            {
+                                "tp": tp,
+                                "num_tokens": num_tokens,
+                                "routing_mode": routing_mode,
+                                "backend": backend.value,
+                                "median_ms": row["time_ms"]["median"],
+                            },
+                            sort_keys=True,
+                        ),
+                        flush=True,
+                    )
+
+            del kernel, w13_kernel, w2_kernel, w13_original, w2
+            torch.accelerator.empty_cache()
+
+    payload = {
+        "schema_version": "qwen30_vllm020_moe_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",
+            "weight_quantization": "none",
+            "top_k": TOP_K,
+            "norm_topk_prob": True,
+        },
+        "measurement_scope": (
+            "vLLM modular MoE prepare+FlashInfer CUTLASS experts+finalize; "
+            "router linear/top-k and TP all-reduce excluded"
+        ),
+        "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()

runs/frontier-qwen30-vllm020-profile-v1/run_moe_smoke.sh

@@ -0,0 +1,47 @@
+#!/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}"
+mkdir -p "${OUTPUT_ROOT}/logs" "${OUTPUT_ROOT}/provenance" "${OUTPUT_ROOT}/raw"
+exec > >(tee -a "${OUTPUT_ROOT}/logs/smoke.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 operator=FlashInfer-CUTLASS-unquantized-MoE tp_local_shapes=1,2,4 tokens=8 routing=uniform_random_logits reference_check=on dtype=BF16 output=${OUTPUT_ROOT} expected_wall=3-8m 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
+
+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_moe.py run_moe_smoke.sh \
+  > "${OUTPUT_ROOT}/provenance/source.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"
+
+timeout --signal=TERM --kill-after=30s 780 \
+  "${VENV_ROOT}/bin/python" profile_vllm020_moe.py \
+  --vllm-source "${VLLM_SOURCE}" \
+  --model "${MODEL}" \
+  --output "${OUTPUT_ROOT}/raw/moe-smoke.json" \
+  --tp 1 2 4 \
+  --num-tokens 8 \
+  --routing-modes uniform_random_logits \
+  --warmup-iters 3 \
+  --repeats 5 \
+  --check-reference
+
+test -s "${OUTPUT_ROOT}/raw/moe-smoke.json"
+sha256sum "${OUTPUT_ROOT}/raw/moe-smoke.json" "${OUTPUT_ROOT}/provenance"/* \
+  > "${OUTPUT_ROOT}/artifacts.sha256"
+date -u +"END_UTC=%Y-%m-%dT%H:%M:%SZ"
+echo "MOE_SMOKE_COMPLETE"