A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 6, 21:45 UTC
Fix routed experts capture for hybrid models (Mamba + Attention) (#35744)
Signed-off-by: arlenxu <arlenxu@tencent.com> Signed-off-by: xhx1022 <1737006628@qq.com> Co-authored-by: arlenxu <arlenxu@tencent.com>
Hongxin Xu committed
Mar 11, 2026 at 23:53 UTC
bea02cdf93bcf9fe94a0efb3240f22facd5e1ac2
4 files changed
+442
-16
examples/offline_inference/routed_experts_e2e.py
new
+384
@@ -0,0 +1,384 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+End-to-end example for routed experts capture with hybrid models.
5
+
6
+Validates that:
7
+1. routed_experts is returned in CompletionOutput for MoE models.
8
+2. Expert IDs are within valid range.
9
+3. Results are deterministic across runs (baseline vs reference).
10
+
11
+Usage:
12
+ python examples/offline_inference/routed_experts_e2e.py \
13
+ --model Qwen/Qwen3-30B-A3B \
14
+ --tp 4 \
15
+ --max-model-len 4096 \
16
+ --num-prompts 20 \
17
+ --max-new-tokens 50
18
+"""
19
+
20
+from __future__ import annotations
21
+
22
+import argparse
23
+import asyncio
24
+import logging
25
+import os
26
+import uuid
27
+from dataclasses import dataclass, field
28
+
29
+import numpy as np
30
+
31
+from vllm.engine.arg_utils import AsyncEngineArgs
32
+
33
+logger = logging.getLogger(__name__)
34
+
35
+DEFAULT_MODEL = "Qwen/Qwen3-30B-A3B"
36
+
37
+TEST_PROMPTS = [
38
+ "Hello, my name is",
39
+ "The capital of France is",
40
+ "Explain quantum computing in simple terms:",
41
+ "Write a Python function that sorts a list:",
42
+ "The meaning of life is",
43
+ "In a distant galaxy, there was a",
44
+ "The best way to learn programming is",
45
+ "Once upon a time in a land far away,",
46
+ "The theory of relativity states that",
47
+ "How does photosynthesis work?",
48
+ "Describe the process of machine learning:",
49
+ "What are the benefits of exercise?",
50
+ "The history of artificial intelligence began",
51
+ "Translate the following to French: Hello world",
52
+ "Summarize the plot of Romeo and Juliet:",
53
+ "What is the difference between TCP and UDP?",
54
+ "The water cycle consists of",
55
+ "Explain how a neural network learns:",
56
+ "The periodic table organizes elements by",
57
+ "Write a haiku about the ocean:",
58
+]
59
+
60
+
61
+@dataclass
62
+class InferenceResult:
63
+ """Result from a single inference run."""
64
+
65
+ experts_list: list[np.ndarray] = field(default_factory=list)
66
+ token_ids_list: list[list[int]] = field(default_factory=list)
67
+ num_experts: int = 0
68
+
69
+
70
+# ---------------------------------------------------------------------------
71
+# Inference helpers
72
+# ---------------------------------------------------------------------------
73
+
74
+
75
+async def _run_async_inference(
76
+ engine_args: AsyncEngineArgs,
77
+ prompts: list[str],
78
+ max_new_tokens: int,
79
+) -> InferenceResult:
80
+ """Run inference using AsyncLLM."""
81
+ from vllm.sampling_params import SamplingParams
82
+ from vllm.v1.engine.async_llm import AsyncLLM
83
+
84
+ engine = AsyncLLM.from_engine_args(engine_args)
85
+
86
+ hf_config = engine.model_config.hf_text_config
87
+ num_experts: int = getattr(hf_config, "num_experts", 0) or getattr(
88
+ hf_config, "num_local_experts", 0
89
+ )
90
+ assert num_experts > 0, "Could not determine num_experts from model config"
91
+
92
+ sampling_params = SamplingParams(
93
+ temperature=0,
94
+ max_tokens=max_new_tokens,
95
+ )
96
+
97
+ async def _generate_one(prompt: str, idx: int):
98
+ request_id = str(uuid.uuid4())
99
+ final_output = None
100
+ async for output in engine.generate(prompt, sampling_params, request_id):
101
+ final_output = output
102
+ assert final_output is not None
103
+
104
+ completion = final_output.outputs[0]
105
+ routed = completion.routed_experts
106
+ num_prompt_tokens = len(final_output.prompt_token_ids)
107
+ num_generated_tokens = len(completion.token_ids)
108
+ expected_len = num_prompt_tokens + num_generated_tokens - 1
109
+ assert routed is not None, f"Prompt {idx}: routed_experts is None"
110
+ assert routed.shape[0] == expected_len, (
111
+ f"Prompt {idx}: routed_experts length {routed.shape[0]} != "
112
+ f"prompt ({num_prompt_tokens}) + generated ({num_generated_tokens})"
113
+ f" - 1 = {expected_len}"
114
+ )
115
+ return idx, routed, list(completion.token_ids)
116
+
117
+ tasks = [_generate_one(p, i) for i, p in enumerate(prompts)]
118
+ outputs = await asyncio.gather(*tasks)
119
+
120
+ # Sort by original index to maintain prompt order
121
+ outputs.sort(key=lambda x: x[0])
122
+
123
+ result = InferenceResult(num_experts=num_experts)
124
+ for _, routed, token_ids in outputs:
125
+ result.experts_list.append(routed)
126
+ result.token_ids_list.append(token_ids)
127
+
128
+ engine.shutdown()
129
+ return result
130
+
131
+
132
+def run_inference(
133
+ model: str,
134
+ prompts: list[str],
135
+ max_new_tokens: int = 50,
136
+ tp: int = 1,
137
+ max_model_len: int = 4096,
138
+) -> InferenceResult:
139
+ """Run inference with routed experts capture enabled via AsyncLLM."""
140
+ engine_args = AsyncEngineArgs(
141
+ model=model,
142
+ enable_return_routed_experts=True,
143
+ tensor_parallel_size=tp,
144
+ max_model_len=max_model_len,
145
+ disable_log_stats=True,
146
+ attention_backend="FLASH_ATTN",
147
+ )
148
+
149
+ result = asyncio.run(_run_async_inference(engine_args, prompts, max_new_tokens))
150
+
151
+ from vllm.platforms import current_platform
152
+
153
+ if current_platform.is_cuda_alike():
154
+ current_platform.empty_cache()
155
+
156
+ return result
157
+
158
+
159
+# ---------------------------------------------------------------------------
160
+# Validation helpers
161
+# ---------------------------------------------------------------------------
162
+
163
+
164
+def validate_expert_ids(
165
+ experts_list: list[np.ndarray],
166
+ num_experts: int,
167
+) -> None:
168
+ """Check that all expert IDs are within valid range [0, num_experts)."""
169
+ for i, experts in enumerate(experts_list):
170
+ assert np.all(experts >= 0), (
171
+ f"Prompt {i}: negative expert IDs found, min={experts.min()}"
172
+ )
173
+ assert np.all(experts < num_experts), (
174
+ f"Prompt {i}: expert ID out of range [0, {num_experts}), "
175
+ f"max={experts.max()}"
176
+ )
177
+
178
+
179
+def validate_shapes(experts_list: list[np.ndarray]) -> None:
180
+ """Check that all routed_experts arrays have at least 2 dimensions."""
181
+ for i, experts in enumerate(experts_list):
182
+ assert experts.ndim >= 2, (
183
+ f"Prompt {i}: expected at least 2D array, got shape {experts.shape}"
184
+ )
185
+ logger.info("Prompt %d: routed_experts shape = %s", i, experts.shape)
186
+
187
+
188
+# ---------------------------------------------------------------------------
189
+# Comparison helpers
190
+# ---------------------------------------------------------------------------
191
+
192
+
193
+def compare_token_ids(
194
+ baseline: list[list[int]],
195
+ reference: list[list[int]],
196
+) -> float:
197
+ """Compare token IDs from two runs. Returns mismatch ratio."""
198
+ assert len(baseline) == len(reference), (
199
+ f"Length mismatch: {len(baseline)} vs {len(reference)}"
200
+ )
201
+
202
+ total_tokens = 0
203
+ total_mismatches = 0
204
+
205
+ for i, (base, ref) in enumerate(zip(baseline, reference)):
206
+ min_len = min(len(base), len(ref))
207
+ max_len = max(len(base), len(ref))
208
+ matches = 0
209
+ for a, b in zip(base[:min_len], ref[:min_len]):
210
+ if a != b:
211
+ break
212
+ matches += 1
213
+
214
+ total_mismatches += max_len - matches
215
+ total_tokens += max_len
216
+
217
+ if matches < min_len or len(base) != len(ref):
218
+ print(
219
+ f" Prompt {i}: token_ids len={len(base)} vs {len(ref)}, "
220
+ f"mismatches={max_len - matches}/{max_len}"
221
+ )
222
+
223
+ if total_tokens == 0:
224
+ raise ValueError("No tokens to compare")
225
+
226
+ mismatch_ratio = total_mismatches / total_tokens
227
+ print(
228
+ f"Token ID mismatches: {total_mismatches}/{total_tokens} ({mismatch_ratio:.4%})"
229
+ )
230
+ return mismatch_ratio
231
+
232
+
233
+def compare_routed_experts(
234
+ baseline: list[np.ndarray],
235
+ reference: list[np.ndarray],
236
+ threshold: float = 0.05,
237
+) -> float:
238
+ """Compare two runs of routed experts. Returns mismatch ratio.
239
+
240
+ Raises AssertionError if ratio exceeds threshold.
241
+ """
242
+ assert len(baseline) == len(reference), (
243
+ f"Length mismatch: {len(baseline)} vs {len(reference)}"
244
+ )
245
+
246
+ total_elements = 0
247
+ total_mismatches = 0
248
+
249
+ for i, (base, ref) in enumerate(zip(baseline, reference)):
250
+ min_len = min(len(base), len(ref))
251
+ max_len = max(len(base), len(ref))
252
+ if min_len == 0:
253
+ continue
254
+
255
+ base_trimmed = base[:min_len]
256
+ ref_trimmed = ref[:min_len]
257
+
258
+ matches = 0
259
+ for a, b in zip(base_trimmed, ref_trimmed):
260
+ if a.sum() != b.sum():
261
+ break
262
+ matches += 1
263
+
264
+ total_mismatches += max_len - matches
265
+ total_elements += max_len
266
+
267
+ if matches < min_len or len(base) != len(ref):
268
+ print(
269
+ f" Prompt {i}: routed_experts len={len(base)} vs {len(ref)}, "
270
+ f"mismatches={max_len - matches}/{max_len}"
271
+ )
272
+
273
+ if total_elements == 0:
274
+ raise ValueError("No elements to compare")
275
+
276
+ mismatch_ratio = total_mismatches / total_elements
277
+ print(
278
+ f"Routed experts mismatches: {total_mismatches}/{total_elements} "
279
+ f"({mismatch_ratio:.4%})"
280
+ )
281
+
282
+ assert mismatch_ratio < threshold, (
283
+ f"Too many mismatches: {total_mismatches}/{total_elements} "
284
+ f"({mismatch_ratio:.4%}) exceeds threshold {threshold:.4%}"
285
+ )
286
+
287
+ return mismatch_ratio
288
+
289
+
290
+# ---------------------------------------------------------------------------
291
+# CLI entry point
292
+# ---------------------------------------------------------------------------
293
+
294
+
295
+def main():
296
+ os.environ.setdefault("VLLM_BATCH_INVARIANT", "1")
297
+
298
+ parser = argparse.ArgumentParser(
299
+ description="Test routed experts capture for MoE models"
300
+ )
301
+ parser.add_argument("--model", type=str, default=DEFAULT_MODEL)
302
+ parser.add_argument("--tp", type=int, default=1)
303
+ parser.add_argument("--max-model-len", type=int, default=4096)
304
+ parser.add_argument("--num-prompts", type=int, default=20)
305
+ parser.add_argument("--max-new-tokens", type=int, default=50)
306
+ parser.add_argument(
307
+ "--deterministic",
308
+ action="store_true",
309
+ help="Run twice and compare results for determinism check",
310
+ )
311
+ parser.add_argument(
312
+ "--threshold",
313
+ type=float,
314
+ default=0.05,
315
+ help="Maximum allowed mismatch ratio for determinism check",
316
+ )
317
+ args = parser.parse_args()
318
+
319
+ logging.basicConfig(level=logging.INFO)
320
+ prompts = TEST_PROMPTS[: args.num_prompts]
321
+
322
+ print(f"Model: {args.model}")
323
+ print(f"TP: {args.tp}")
324
+ print(f"Prompts: {len(prompts)}")
325
+ print(f"Max new tokens: {args.max_new_tokens}")
326
+ print()
327
+
328
+ print("=== Run 1 (baseline) ===")
329
+ baseline = run_inference(
330
+ model=args.model,
331
+ prompts=prompts,
332
+ max_new_tokens=args.max_new_tokens,
333
+ tp=args.tp,
334
+ max_model_len=args.max_model_len,
335
+ )
336
+ print(f"num_experts (from model config): {baseline.num_experts}")
337
+
338
+ print("\n=== Validation ===")
339
+ validate_shapes(baseline.experts_list)
340
+ validate_expert_ids(baseline.experts_list, num_experts=baseline.num_experts)
341
+ print(f"All {len(baseline.experts_list)} results passed validation.")
342
+
343
+ for i, experts in enumerate(baseline.experts_list):
344
+ print(
345
+ f" Prompt {i}: shape={experts.shape}, "
346
+ f"min={experts.min()}, max={experts.max()}"
347
+ )
348
+
349
+ if args.deterministic:
350
+ print("\n=== Run 2 (reference) ===")
351
+ reference = run_inference(
352
+ model=args.model,
353
+ prompts=prompts,
354
+ max_new_tokens=args.max_new_tokens,
355
+ tp=args.tp,
356
+ max_model_len=args.max_model_len,
357
+ )
358
+
359
+ print("\n=== Determinism Check ===")
360
+ validate_expert_ids(reference.experts_list, num_experts=baseline.num_experts)
361
+
362
+ print("\n--- Token IDs ---")
363
+ token_mismatch = compare_token_ids(
364
+ baseline.token_ids_list, reference.token_ids_list
365
+ )
366
+
367
+ print("\n--- Routed Experts ---")
368
+ expert_mismatch = compare_routed_experts(
369
+ baseline.experts_list,
370
+ reference.experts_list,
371
+ threshold=args.threshold,
372
+ )
373
+
374
+ print(
375
+ f"\nDeterminism check passed. "
376
+ f"Token mismatch: {token_mismatch:.4%}, "
377
+ f"Expert mismatch: {expert_mismatch:.4%}"
378
+ )
379
+
380
+ print("\nAll tests passed!")
381
+
382
+
383
+if __name__ == "__main__":
384
+ main()
vllm/v1/core/sched/scheduler.py
+24
-6
@@ -52,7 +52,7 @@ from vllm.v1.core.sched.request_queue import (
52
)
53
from vllm.v1.core.sched.utils import check_stop, remove_all
54
from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs
55
-from vllm.v1.kv_cache_interface import KVCacheConfig
55
+from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig
56
from vllm.v1.metrics.perf import ModelMetrics, PerfStats
57
from vllm.v1.metrics.stats import PrefixCacheStats, SchedulerStats
58
from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput
@@ -259,9 +259,26 @@ class Scheduler(SchedulerInterface):
259
assert len(kv_cache_config.kv_cache_groups) > 0, (
260
"enable_return_routed_experts requires at least one kv cache group"
261
)
262
+ # Find the attention group for routed experts indexing.
263
+ self.routed_experts_attn_gid = 0
264
+ for gid, group in enumerate(kv_cache_config.kv_cache_groups):
265
+ if isinstance(group.kv_cache_spec, AttentionSpec):
266
+ self.routed_experts_attn_gid = gid
267
+ break
268
+ min_block_size = min(
269
+ [
270
+ group.kv_cache_spec.block_size
271
+ for group in kv_cache_config.kv_cache_groups
272
+ ]
273
+ )
274
+ num_groups = len(kv_cache_config.kv_cache_groups)
275
self.max_num_kv_tokens = (
263
- kv_cache_config.num_blocks // len(kv_cache_config.kv_cache_groups) + 1
264
- ) * self.block_size
276
+ kv_cache_config.num_blocks // num_groups
277
+ ) * min_block_size
278
+ dcp_size = self.vllm_config.parallel_config.decode_context_parallel_size
279
+ pcp_size = self.vllm_config.parallel_config.prefill_context_parallel_size
280
+ if pcp_size * dcp_size > 1:
281
+ self.max_num_kv_tokens *= pcp_size * dcp_size
282
283
self.routed_experts_reader.attach_buffer(
284
max_num_kv_tokens=self.max_num_kv_tokens,
@@ -1561,13 +1578,14 @@ class Scheduler(SchedulerInterface):
1578
return None
1579
1580
kv_blocks = self.kv_cache_manager.get_blocks(request.request_id)
1564
- block_ids = kv_blocks.get_block_ids()[0]
1581
+ block_ids = kv_blocks.get_block_ids()[self.routed_experts_attn_gid]
1582
num_tokens = request.num_tokens - 1
1583
1567
- # compute slot mapping
1584
+ # compute slot mapping using attention group's block_size
1585
block_ids_array = np.array(block_ids, dtype=np.int32)
1586
num_blocks = len(block_ids)
1570
- block_size = self.block_size
1587
+ attn_group = self.kv_cache_config.kv_cache_groups[self.routed_experts_attn_gid]
1588
+ block_size = attn_group.kv_cache_spec.block_size
1589
1590
# generate block offsets
1591
block_offsets = np.arange(0, block_size)
vllm/v1/worker/gpu_model_runner.py
+31
-10
@@ -422,6 +422,9 @@ class GPUModelRunner(
422
)
423
# This will be overridden in load_model()
424
self.is_multimodal_pruning_enabled = False
425
+ # Set to True after init_routed_experts_capturer() completes.
426
+ # Prevents routed experts code from running during profiling/dummy run.
427
+ self.routed_experts_initialized = False
428
self.max_model_len = model_config.max_model_len
429
430
# Always set to false after the first forward pass
@@ -1951,8 +1954,10 @@ class GPUModelRunner(
1954
block_table_gid_0 = _get_block_table(0)
1955
slot_mapping_gid_0 = slot_mappings[0]
1956
1954
- if self.model_config.enable_return_routed_experts:
1955
- self.slot_mapping = slot_mapping_gid_0[:num_tokens].cpu().numpy()
1957
+ if self.routed_experts_initialized:
1958
+ attn_gid = self.routed_experts_attn_gid
1959
+ slot_mapping_attn = slot_mappings[attn_gid]
1960
+ self.slot_mapping = slot_mapping_attn[:num_tokens].cpu().numpy()
1961
cm_base = CommonAttentionMetadata(
1962
query_start_loc=self.query_start_loc.gpu[: num_reqs_padded + 1],
1963
query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs_padded + 1],
@@ -3540,7 +3545,7 @@ class GPUModelRunner(
3545
"after execute_model() returns None."
3546
)
3547
3543
- if self.vllm_config.model_config.enable_return_routed_experts:
3548
+ if self.routed_experts_initialized:
3549
capturer = RoutedExpertsCapturer.get_instance()
3550
if capturer is not None:
3551
capturer.clear_buffer() # noqa
@@ -4049,7 +4054,7 @@ class GPUModelRunner(
4054
self.kv_connector_output = None
4055
4056
with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"):
4052
- if self.model_config.enable_return_routed_experts:
4057
+ if self.routed_experts_initialized:
4058
capturer = RoutedExpertsCapturer.get_instance()
4059
if capturer is not None:
4060
capturer.save_captured_experts(indices=self.slot_mapping) # noqa
@@ -6531,8 +6536,12 @@ class GPUModelRunner(
6536
kv_transfer_group.register_kv_caches(kv_caches)
6537
kv_transfer_group.set_host_xfer_buffer_ops(copy_kv_blocks)
6538
6534
- if self.model_config.enable_return_routed_experts:
6535
- self.init_routed_experts_capturer()
6539
+ def _get_attention_kv_cache_gid(self) -> int:
6540
+ """Find the KV cache group index for attention layers."""
6541
+ for gid, group in enumerate(self.kv_cache_config.kv_cache_groups):
6542
+ if isinstance(group.kv_cache_spec, AttentionSpec):
6543
+ return gid
6544
+ return 0
6545
6546
def init_routed_experts_capturer(self):
6547
logger.info(
@@ -6540,17 +6549,29 @@ class GPUModelRunner(
6549
self.model_config.enable_return_routed_experts,
6550
)
6551
routed_experts_capturer = RoutedExpertsCapturer.create()
6543
- block_size = self.cache_config.block_size
6552
+ self.routed_experts_attn_gid = self._get_attention_kv_cache_gid()
6553
+ min_block_size = min(
6554
+ [
6555
+ group.kv_cache_spec.block_size
6556
+ for group in self.kv_cache_config.kv_cache_groups
6557
+ ]
6558
+ )
6559
+ num_groups = len(self.kv_cache_config.kv_cache_groups)
6560
self.max_num_kv_tokens = (
6545
- self.kv_cache_config.num_blocks // len(self.kv_cache_config.kv_cache_groups)
6546
- + 1
6547
- ) * block_size
6561
+ self.kv_cache_config.num_blocks // num_groups
6562
+ ) * min_block_size
6563
+ dcp_size = self.vllm_config.parallel_config.decode_context_parallel_size
6564
+ pcp_size = self.vllm_config.parallel_config.prefill_context_parallel_size
6565
+ if pcp_size * dcp_size > 1:
6566
+ self.max_num_kv_tokens *= pcp_size * dcp_size
6567
+
6568
routed_experts_capturer.init_buffer(
6569
max_num_batched_tokens=self.scheduler_config.max_num_batched_tokens,
6570
max_num_kv_tokens=self.max_num_kv_tokens,
6571
vllm_config=self.vllm_config,
6572
)
6573
self._bind_routed_experts_capturer(routed_experts_capturer)
6574
+ self.routed_experts_initialized = True
6575
6576
def _bind_routed_experts_capturer(self, capturer: RoutedExpertsCapturer) -> None:
6577
from vllm.model_executor.layers.fused_moe.layer import FusedMoE
vllm/v1/worker/gpu_worker.py
+3
@@ -552,6 +552,9 @@ class Worker(WorkerBase):
552
else:
553
self.model_runner.initialize_kv_cache(kv_cache_config)
554
555
+ if self.model_config.enable_return_routed_experts:
556
+ self.model_runner.init_routed_experts_capturer()
557
+
558
# Build KV-zero metadata outside the CuMem pool so the bookkeeping
559
# GPU tensors (seg_addrs, block-id buffers) use the standard PyTorch
560
# allocator and are not discarded during sleep/wake cycles.