| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """End-to-end example for routed experts capture with hybrid models. |
| 4 | |
| 5 | Validates that: |
| 6 | 1. routed_experts is returned in CompletionOutput for MoE models. |
| 7 | 2. Expert IDs are within valid range. |
| 8 | 3. Results are deterministic across runs (baseline vs reference). |
| 9 | |
| 10 | Usage: |
| 11 | python examples/rl/routed_experts_e2e.py \ |
| 12 | --model Qwen/Qwen3-30B-A3B \ |
| 13 | --tp 4 \ |
| 14 | --max-model-len 4096 \ |
| 15 | --num-prompts 20 \ |
| 16 | --max-new-tokens 50 |
| 17 | """ |
| 18 | |
| 19 | from __future__ import annotations |
| 20 | |
| 21 | import argparse |
| 22 | import asyncio |
| 23 | import logging |
| 24 | import os |
| 25 | import uuid |
| 26 | from dataclasses import dataclass, field |
| 27 | |
| 28 | import numpy as np |
| 29 | |
| 30 | from vllm.engine.arg_utils import AsyncEngineArgs |
| 31 | |
| 32 | logger = logging.getLogger(__name__) |
| 33 | |
| 34 | DEFAULT_MODEL = "Qwen/Qwen3-30B-A3B" |
| 35 | |
| 36 | TEST_PROMPTS = [ |
| 37 | "Hello, my name is", |
| 38 | "The capital of France is", |
| 39 | "Explain quantum computing in simple terms:", |
| 40 | "Write a Python function that sorts a list:", |
| 41 | "The meaning of life is", |
| 42 | "In a distant galaxy, there was a", |
| 43 | "The best way to learn programming is", |
| 44 | "Once upon a time in a land far away,", |
| 45 | "The theory of relativity states that", |
| 46 | "How does photosynthesis work?", |
| 47 | "Describe the process of machine learning:", |
| 48 | "What are the benefits of exercise?", |
| 49 | "The history of artificial intelligence began", |
| 50 | "Translate the following to French: Hello world", |
| 51 | "Summarize the plot of Romeo and Juliet:", |
| 52 | "What is the difference between TCP and UDP?", |
| 53 | "The water cycle consists of", |
| 54 | "Explain how a neural network learns:", |
| 55 | "The periodic table organizes elements by", |
| 56 | "Write a haiku about the ocean:", |
| 57 | ] |
| 58 | |
| 59 | |
| 60 | @dataclass |
| 61 | class InferenceResult: |
| 62 | """Result from a single inference run.""" |
| 63 | |
| 64 | experts_list: list[np.ndarray] = field(default_factory=list) |
| 65 | token_ids_list: list[list[int]] = field(default_factory=list) |
| 66 | num_experts: int = 0 |
| 67 | |
| 68 | |
| 69 | # --------------------------------------------------------------------------- |
| 70 | # Inference helpers |
| 71 | # --------------------------------------------------------------------------- |
| 72 | |
| 73 | |
| 74 | async def _run_async_inference( |
| 75 | engine_args: AsyncEngineArgs, |
| 76 | prompts: list[str], |
| 77 | max_new_tokens: int, |
| 78 | ) -> InferenceResult: |
| 79 | """Run inference using AsyncLLM.""" |
| 80 | from vllm.sampling_params import SamplingParams |
| 81 | from vllm.v1.engine.async_llm import AsyncLLM |
| 82 | |
| 83 | engine = AsyncLLM.from_engine_args(engine_args) |
| 84 | |
| 85 | hf_config = engine.model_config.hf_text_config |
| 86 | num_experts: int = getattr(hf_config, "num_experts", 0) or getattr( |
| 87 | hf_config, "num_local_experts", 0 |
| 88 | ) |
| 89 | assert num_experts > 0, "Could not determine num_experts from model config" |
| 90 | |
| 91 | sampling_params = SamplingParams( |
| 92 | temperature=0, |
| 93 | max_tokens=max_new_tokens, |
| 94 | ) |
| 95 | |
| 96 | async def _generate_one(prompt: str, idx: int): |
| 97 | request_id = str(uuid.uuid4()) |
| 98 | final_output = None |
| 99 | async for output in engine.generate(prompt, sampling_params, request_id): |
| 100 | final_output = output |
| 101 | assert final_output is not None |
| 102 | |
| 103 | completion = final_output.outputs[0] |
| 104 | routed = completion.routed_experts |
| 105 | num_prompt_tokens = len(final_output.prompt_token_ids) |
| 106 | num_generated_tokens = len(completion.token_ids) |
| 107 | expected_len = num_prompt_tokens + num_generated_tokens - 1 |
| 108 | assert routed is not None, f"Prompt {idx}: routed_experts is None" |
| 109 | assert routed.shape[0] == expected_len, ( |
| 110 | f"Prompt {idx}: routed_experts length {routed.shape[0]} != " |
| 111 | f"prompt ({num_prompt_tokens}) + generated ({num_generated_tokens})" |
| 112 | f" - 1 = {expected_len}" |
| 113 | ) |
| 114 | return idx, routed, list(completion.token_ids) |
| 115 | |
| 116 | tasks = [_generate_one(p, i) for i, p in enumerate(prompts)] |
| 117 | outputs = await asyncio.gather(*tasks) |
| 118 | |
| 119 | # Sort by original index to maintain prompt order |
| 120 | outputs.sort(key=lambda x: x[0]) |
| 121 | |
| 122 | result = InferenceResult(num_experts=num_experts) |
| 123 | for _, routed, token_ids in outputs: |
| 124 | result.experts_list.append(routed) |
| 125 | result.token_ids_list.append(token_ids) |
| 126 | |
| 127 | engine.shutdown() |
| 128 | return result |
| 129 | |
| 130 | |
| 131 | def run_inference( |
| 132 | model: str, |
| 133 | prompts: list[str], |
| 134 | max_new_tokens: int = 50, |
| 135 | tp: int = 1, |
| 136 | max_model_len: int = 4096, |
| 137 | ) -> InferenceResult: |
| 138 | """Run inference with routed experts capture enabled via AsyncLLM.""" |
| 139 | engine_args = AsyncEngineArgs( |
| 140 | model=model, |
| 141 | enable_return_routed_experts=True, |
| 142 | tensor_parallel_size=tp, |
| 143 | max_model_len=max_model_len, |
| 144 | disable_log_stats=True, |
| 145 | attention_backend="FLASH_ATTN", |
| 146 | ) |
| 147 | |
| 148 | result = asyncio.run(_run_async_inference(engine_args, prompts, max_new_tokens)) |
| 149 | |
| 150 | from vllm.platforms import current_platform |
| 151 | |
| 152 | if current_platform.is_cuda_alike(): |
| 153 | current_platform.empty_cache() |
| 154 | |
| 155 | return result |
| 156 | |
| 157 | |
| 158 | # --------------------------------------------------------------------------- |
| 159 | # Validation helpers |
| 160 | # --------------------------------------------------------------------------- |
| 161 | |
| 162 | |
| 163 | def validate_expert_ids( |
| 164 | experts_list: list[np.ndarray], |
| 165 | num_experts: int, |
| 166 | ) -> None: |
| 167 | """Check that all expert IDs are within valid range [0, num_experts).""" |
| 168 | for i, experts in enumerate(experts_list): |
| 169 | assert np.all(experts >= 0), ( |
| 170 | f"Prompt {i}: negative expert IDs found, min={experts.min()}" |
| 171 | ) |
| 172 | assert np.all(experts < num_experts), ( |
| 173 | f"Prompt {i}: expert ID out of range [0, {num_experts}), " |
| 174 | f"max={experts.max()}" |
| 175 | ) |
| 176 | |
| 177 | |
| 178 | def validate_shapes(experts_list: list[np.ndarray]) -> None: |
| 179 | """Check that all routed_experts arrays have at least 2 dimensions.""" |
| 180 | for i, experts in enumerate(experts_list): |
| 181 | assert experts.ndim >= 2, ( |
| 182 | f"Prompt {i}: expected at least 2D array, got shape {experts.shape}" |
| 183 | ) |
| 184 | logger.info("Prompt %d: routed_experts shape = %s", i, experts.shape) |
| 185 | |
| 186 | |
| 187 | # --------------------------------------------------------------------------- |
| 188 | # Comparison helpers |
| 189 | # --------------------------------------------------------------------------- |
| 190 | |
| 191 | |
| 192 | def compare_token_ids( |
| 193 | baseline: list[list[int]], |
| 194 | reference: list[list[int]], |
| 195 | ) -> float: |
| 196 | """Compare token IDs from two runs. Returns mismatch ratio.""" |
| 197 | assert len(baseline) == len(reference), ( |
| 198 | f"Length mismatch: {len(baseline)} vs {len(reference)}" |
| 199 | ) |
| 200 | |
| 201 | total_tokens = 0 |
| 202 | total_mismatches = 0 |
| 203 | |
| 204 | for i, (base, ref) in enumerate(zip(baseline, reference)): |
| 205 | min_len = min(len(base), len(ref)) |
| 206 | max_len = max(len(base), len(ref)) |
| 207 | matches = 0 |
| 208 | for a, b in zip(base[:min_len], ref[:min_len]): |
| 209 | if a != b: |
| 210 | break |
| 211 | matches += 1 |
| 212 | |
| 213 | total_mismatches += max_len - matches |
| 214 | total_tokens += max_len |
| 215 | |
| 216 | if matches < min_len or len(base) != len(ref): |
| 217 | print( |
| 218 | f" Prompt {i}: token_ids len={len(base)} vs {len(ref)}, " |
| 219 | f"mismatches={max_len - matches}/{max_len}" |
| 220 | ) |
| 221 | |
| 222 | if total_tokens == 0: |
| 223 | raise ValueError("No tokens to compare") |
| 224 | |
| 225 | mismatch_ratio = total_mismatches / total_tokens |
| 226 | print( |
| 227 | f"Token ID mismatches: {total_mismatches}/{total_tokens} ({mismatch_ratio:.4%})" |
| 228 | ) |
| 229 | return mismatch_ratio |
| 230 | |
| 231 | |
| 232 | def compare_routed_experts( |
| 233 | baseline: list[np.ndarray], |
| 234 | reference: list[np.ndarray], |
| 235 | threshold: float = 0.05, |
| 236 | ) -> float: |
| 237 | """Compare two runs of routed experts. Returns mismatch ratio. |
| 238 | |
| 239 | Raises AssertionError if ratio exceeds threshold. |
| 240 | """ |
| 241 | assert len(baseline) == len(reference), ( |
| 242 | f"Length mismatch: {len(baseline)} vs {len(reference)}" |
| 243 | ) |
| 244 | |
| 245 | total_elements = 0 |
| 246 | total_mismatches = 0 |
| 247 | |
| 248 | for i, (base, ref) in enumerate(zip(baseline, reference)): |
| 249 | min_len = min(len(base), len(ref)) |
| 250 | max_len = max(len(base), len(ref)) |
| 251 | if min_len == 0: |
| 252 | continue |
| 253 | |
| 254 | base_trimmed = base[:min_len] |
| 255 | ref_trimmed = ref[:min_len] |
| 256 | |
| 257 | matches = 0 |
| 258 | for a, b in zip(base_trimmed, ref_trimmed): |
| 259 | if a.sum() != b.sum(): |
| 260 | break |
| 261 | matches += 1 |
| 262 | |
| 263 | total_mismatches += max_len - matches |
| 264 | total_elements += max_len |
| 265 | |
| 266 | if matches < min_len or len(base) != len(ref): |
| 267 | print( |
| 268 | f" Prompt {i}: routed_experts len={len(base)} vs {len(ref)}, " |
| 269 | f"mismatches={max_len - matches}/{max_len}" |
| 270 | ) |
| 271 | |
| 272 | if total_elements == 0: |
| 273 | raise ValueError("No elements to compare") |
| 274 | |
| 275 | mismatch_ratio = total_mismatches / total_elements |
| 276 | print( |
| 277 | f"Routed experts mismatches: {total_mismatches}/{total_elements} " |
| 278 | f"({mismatch_ratio:.4%})" |
| 279 | ) |
| 280 | |
| 281 | assert mismatch_ratio < threshold, ( |
| 282 | f"Too many mismatches: {total_mismatches}/{total_elements} " |
| 283 | f"({mismatch_ratio:.4%}) exceeds threshold {threshold:.4%}" |
| 284 | ) |
| 285 | |
| 286 | return mismatch_ratio |
| 287 | |
| 288 | |
| 289 | # --------------------------------------------------------------------------- |
| 290 | # CLI entry point |
| 291 | # --------------------------------------------------------------------------- |
| 292 | |
| 293 | |
| 294 | def main(): |
| 295 | os.environ.setdefault("VLLM_BATCH_INVARIANT", "1") |
| 296 | |
| 297 | parser = argparse.ArgumentParser( |
| 298 | description="Test routed experts capture for MoE models" |
| 299 | ) |
| 300 | parser.add_argument("--model", type=str, default=DEFAULT_MODEL) |
| 301 | parser.add_argument("--tp", type=int, default=1) |
| 302 | parser.add_argument("--max-model-len", type=int, default=4096) |
| 303 | parser.add_argument("--num-prompts", type=int, default=20) |
| 304 | parser.add_argument("--max-new-tokens", type=int, default=50) |
| 305 | parser.add_argument( |
| 306 | "--deterministic", |
| 307 | action="store_true", |
| 308 | help="Run twice and compare results for determinism check", |
| 309 | ) |
| 310 | parser.add_argument( |
| 311 | "--threshold", |
| 312 | type=float, |
| 313 | default=0.05, |
| 314 | help="Maximum allowed mismatch ratio for determinism check", |
| 315 | ) |
| 316 | args = parser.parse_args() |
| 317 | |
| 318 | logging.basicConfig(level=logging.INFO) |
| 319 | prompts = TEST_PROMPTS[: args.num_prompts] |
| 320 | |
| 321 | print(f"Model: {args.model}") |
| 322 | print(f"TP: {args.tp}") |
| 323 | print(f"Prompts: {len(prompts)}") |
| 324 | print(f"Max new tokens: {args.max_new_tokens}") |
| 325 | print() |
| 326 | |
| 327 | print("=== Run 1 (baseline) ===") |
| 328 | baseline = run_inference( |
| 329 | model=args.model, |
| 330 | prompts=prompts, |
| 331 | max_new_tokens=args.max_new_tokens, |
| 332 | tp=args.tp, |
| 333 | max_model_len=args.max_model_len, |
| 334 | ) |
| 335 | print(f"num_experts (from model config): {baseline.num_experts}") |
| 336 | |
| 337 | print("\n=== Validation ===") |
| 338 | validate_shapes(baseline.experts_list) |
| 339 | validate_expert_ids(baseline.experts_list, num_experts=baseline.num_experts) |
| 340 | print(f"All {len(baseline.experts_list)} results passed validation.") |
| 341 | |
| 342 | for i, experts in enumerate(baseline.experts_list): |
| 343 | print( |
| 344 | f" Prompt {i}: shape={experts.shape}, " |
| 345 | f"min={experts.min()}, max={experts.max()}" |
| 346 | ) |
| 347 | |
| 348 | if args.deterministic: |
| 349 | print("\n=== Run 2 (reference) ===") |
| 350 | reference = run_inference( |
| 351 | model=args.model, |
| 352 | prompts=prompts, |
| 353 | max_new_tokens=args.max_new_tokens, |
| 354 | tp=args.tp, |
| 355 | max_model_len=args.max_model_len, |
| 356 | ) |
| 357 | |
| 358 | print("\n=== Determinism Check ===") |
| 359 | validate_expert_ids(reference.experts_list, num_experts=baseline.num_experts) |
| 360 | |
| 361 | print("\n--- Token IDs ---") |
| 362 | token_mismatch = compare_token_ids( |
| 363 | baseline.token_ids_list, reference.token_ids_list |
| 364 | ) |
| 365 | |
| 366 | print("\n--- Routed Experts ---") |
| 367 | expert_mismatch = compare_routed_experts( |
| 368 | baseline.experts_list, |
| 369 | reference.experts_list, |
| 370 | threshold=args.threshold, |
| 371 | ) |
| 372 | |
| 373 | print( |
| 374 | f"\nDeterminism check passed. " |
| 375 | f"Token mismatch: {token_mismatch:.4%}, " |
| 376 | f"Expert mismatch: {expert_mismatch:.4%}" |
| 377 | ) |
| 378 | |
| 379 | print("\nAll tests passed!") |
| 380 | |
| 381 | |
| 382 | if __name__ == "__main__": |
| 383 | main() |