@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 8189a15914
Fork 0

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 4, 03:45 UTC

[Core] Replace routing replay with device cache and async D2H pipeline (#39917)

Signed-off-by: Tomer Barnatan <tbarnatan@nvidia.com>

TomerBN-Nvidia committed May 7, 2026 at 21:24 UTC 8189a15914ca48461acf106f126c58ef7e41c9ee
15 files changed +1386 -630
docs/training/routed_experts_replay.md new
+289
@@ -0,0 +1,289 @@
1 +# Routed Experts Replay
2 +
3 +## Overview
4 +
5 +Routed experts replay captures which MoE (Mixture of Experts) experts process each token during inference and returns this information alongside the generated text. This is essential for **reinforcement learning (RL) training pipelines** (such as GRPO and RLHF) where the training step needs to reconstruct expert routing decisions from the inference pass.
6 +
7 +When enabled, each API response includes:
8 +
9 +- **`prompt_routed_experts`**: A `[prompt_len, num_moe_layers, top_k]` array of expert IDs for the prompt tokens (at the response level, shared across completions).
10 +- **`routed_experts`**: A `[gen_len, num_moe_layers, top_k]` array of expert IDs for the generated tokens (per completion).
11 +
12 +For example, a model with 40 MoE layers and top-22 routing that processes a 100-token prompt and generates 50 tokens would return:
13 +
14 +- `prompt_routed_experts`: shape `[100, 40, 22]`
15 +- `routed_experts`: shape `[50, 40, 22]`
16 +
17 +Each value is an int16 expert ID in the range `[0, num_experts)`.
18 +
19 +## Quickstart
20 +
21 +### OpenAI API Server
22 +
23 +```bash
24 +vllm serve <MODEL> \
25 + --enable-return-routed-experts \
26 + --tensor-parallel-size 4 \
27 + --enable-expert-parallel
28 +```
29 +
30 +Then query the `/v1/completions` endpoint as usual. The response includes routing data:
31 +
32 +```python
33 +import requests
34 +
35 +resp = requests.post("http://localhost:8000/v1/completions", json={
36 + "model": "<MODEL>",
37 + "prompt": "Explain quantum computing.",
38 + "max_tokens": 64,
39 + "temperature": 0.0,
40 +}).json()
41 +
42 +# Generation routing (per completion choice)
43 +gen_routing = resp["choices"][0]["routed_experts"] # [gen_len, layers, top_k]
44 +
45 +# Prompt routing (shared across all choices)
46 +prompt_routing = resp["prompt_routed_experts"] # [prompt_len, layers, top_k]
47 +
48 +print(f"Prompt routing shape: [{len(prompt_routing)}, "
49 + f"{len(prompt_routing[0])}, {len(prompt_routing[0][0])}]")
50 +print(f"Gen routing shape: [{len(gen_routing)}, "
51 + f"{len(gen_routing[0])}, {len(gen_routing[0][0])}]")
52 +```
53 +
54 +### Python SDK (Offline Inference)
55 +
56 +```python
57 +from vllm import LLM, SamplingParams
58 +
59 +llm = LLM(
60 + model="<MODEL>",
61 + enable_return_routed_experts=True,
62 + tensor_parallel_size=4,
63 + enable_expert_parallel=True,
64 +)
65 +
66 +outputs = llm.generate(
67 + ["Explain quantum computing."],
68 + SamplingParams(temperature=0, max_tokens=64),
69 +)
70 +
71 +result = outputs[0]
72 +
73 +# Prompt routing: numpy array, shape [prompt_len, num_moe_layers, top_k]
74 +prompt_routing = result.prompt_routed_experts
75 +print(f"Prompt routing: {prompt_routing.shape}, dtype={prompt_routing.dtype}")
76 +
77 +# Generation routing: numpy array, shape [gen_len, num_moe_layers, top_k]
78 +gen_routing = result.outputs[0].routed_experts
79 +print(f"Gen routing: {gen_routing.shape}, dtype={gen_routing.dtype}")
80 +```
81 +
82 +## Output Format
83 +
84 +### `CompletionOutput.routed_experts`
85 +
86 +- **Type**: `numpy.ndarray` (Python SDK) or `list[list[list[int]]]` (JSON API)
87 +- **Shape**: `[gen_len, num_moe_layers, top_k]`
88 +- **Dtype**: `int16`
89 +- **Content**: Expert IDs for **generated tokens only**. `gen_len` matches the number of generated tokens (i.e., `usage.completion_tokens` or fewer).
90 +
91 +### `RequestOutput.prompt_routed_experts`
92 +
93 +- **Type**: `numpy.ndarray` (Python SDK) or `list[list[list[int]]]` (JSON API)
94 +- **Shape**: `[prompt_len, num_moe_layers, top_k]`
95 +- **Dtype**: `int16`
96 +- **Content**: Expert IDs for **prompt tokens only**. `prompt_len` matches `usage.prompt_tokens`. This field lives on the request-level response (not per-choice), because prompt routing is shared across all completions when `n > 1`.
97 +
98 +### Why Separate Prompt and Generation Routing?
99 +
100 +When a request has multiple completions (`n > 1`), each completion shares the same prompt but produces different generated text. Storing prompt routing once on the `RequestOutput` (rather than duplicating it on every `CompletionOutput`) avoids redundant data. For RL training, the consumer typically needs:
101 +
102 +1. The prompt routing (once) to reconstruct the forward pass for the shared prefix.
103 +2. The per-completion generation routing to reconstruct each completion's forward pass.
104 +
105 +## Architecture
106 +
107 +### Data Flow
108 +
109 +```text
110 +Forward Pass Async D2H Pipeline Output
111 +───────────── ────────────────── ──────
112 +FusedMoE layer After forward pass: On request finish:
113 +writes topk_ids ──────► D2H copy to pinned ──────► Extract from host cache
114 +to device buffer staging buffer Split at prompt_len
115 +(L, N, K) int16 (via CUDA stream) Trim gen to output len
116 + Scatter to per-request Serialize to API response
117 + host cache (numpy)
118 +```
119 +
120 +### Device Cache
121 +
122 +A pre-allocated GPU buffer with layout `(L, N, K)` where:
123 +
124 +- `L` = number of MoE layers
125 +- `N` = `max_num_batched_tokens`
126 +- `K` = `num_experts_per_tok` (top-k)
127 +
128 +The `(L, N, K)` layout ensures that `buffer[layer_id]` gives a contiguous `(N, K)` view per layer. Each `FusedMoE` layer gets a persistent reference to its slice via `module._routing_replay_out = buffer[layer_id]`.
129 +
130 +**Dtype**: `int16` — sufficient for expert IDs (max ~512 experts in practice) and half the memory of `int32`.
131 +
132 +### Host Cache
133 +
134 +Per-request numpy arrays for accumulating routing data across decode steps. Each request gets a lazily allocated `(seq_len, L, K)` int16 buffer that grows as the sequence lengthens. Buffers are freed when a request completes.
135 +
136 +### Async D2H Pipeline
137 +
138 +After each forward pass, the model runner issues a non-blocking device-to-host copy on a dedicated CUDA stream:
139 +
140 +1. **Copy**: `pinned_staging[:, :total_tokens, :].copy_(device_buffer[:, :total_tokens, :])` on a separate stream, recorded with a CUDA event.
141 +2. **Scatter** (deferred to next step): On the *next* forward pass, synchronize the event (effectively free — an entire forward pass has elapsed) and scatter the staging data into per-request host cache buffers using the token positions.
142 +
143 +This design ensures the D2H copy overlaps with the next forward pass, minimizing GPU stall time.
144 +
145 +### CUDA Graph Compatibility
146 +
147 +CUDA graph compatibility requires two mechanisms:
148 +
149 +1. **Persistent tensor attribute**: Each `FusedMoE` layer stores a reference to its buffer slice as `module._routing_replay_out`. Because `torch.compile` captures module attributes by reference, graph replay always writes to the live buffer — not a stale snapshot.
150 +
151 +2. **Static marking**: Both the full `(L, N, K)` buffer and each per-layer `(N, K)` view are marked with `cudagraph_mark_tensor_static()`. This prevents CUDA graphs from snapshot/restore behavior that would zero the buffer on replay.
152 +
153 +### Multi-Node Support
154 +
155 +On multi-node tensor-parallel setups, all TP ranks allocate a device buffer (required for symmetric CUDA graph structure), but only TP rank 0 runs the D2H pipeline and host cache. Routing data flows from the model runner through `ModelRunnerOutput` via Ray DAG to the scheduler — no shared memory or file locks needed.
156 +
157 +### Routing Capture Path
158 +
159 +For the **non-monolithic (Triton) kernel path** (e.g., BF16 MoE), routing is captured after `select_experts()` in the MoE runner:
160 +
161 +```python
162 +routing_replay_out = getattr(layer, "_routing_replay_out", None)
163 +topk_weights, topk_ids = self.router.select_experts(...)
164 +
165 +if routing_replay_out is not None:
166 + routing_replay_out[:topk_ids.shape[0]].copy_(topk_ids.to(torch.int16))
167 +```
168 +
169 +For the **monolithic kernel path** (e.g., FP8/MXFP8 via FlashInfer), `routing_replay_out` is threaded through the `apply_monolithic()` call chain and FlashInfer writes expert IDs directly during routing inside the fused kernel.
170 +
171 +### MTP (Multi-Token Prediction) Handling
172 +
173 +With MTP speculative decoding, the model captures routing for all tokens including speculative ones that may later be rejected. When a request finishes, the generation routing is trimmed to match the actual number of accepted output tokens:
174 +
175 +```python
176 +num_gen = self.detokenizer.num_output_tokens()
177 +if gen_routed_experts.shape[0] > num_gen and num_gen > 0:
178 + gen_routed_experts = gen_routed_experts[:num_gen]
179 +```
180 +
181 +This ensures the routing array length always matches the token IDs in the response.
182 +
183 +## Design Decisions
184 +
185 +### Why Replace SharedMemory with Device Cache?
186 +
187 +The previous implementation used `multiprocessing.SharedMemory` with `fcntl` file locking to transfer routing data from GPU workers to the scheduler. This approach had fundamental problems:
188 +
189 +- **Multi-node**: `SharedMemory` is node-local. On multi-node TP setups (required for 400B+ parameter models), the scheduler on node 0 cannot read shared memory from workers on other nodes.
190 +- **Performance**: Synchronous `.cpu().numpy()` D2H transfers block the GPU. File-based locking adds further overhead.
191 +- **CUDA graphs**: The callback-based capture mechanism bakes tensor references at trace time, causing stale data on graph replay.
192 +
193 +The device cache approach solves all three: data flows through Ray DAG (works multi-node), D2H is async (non-blocking), and persistent tensor attributes work with CUDA graphs.
194 +
195 +### Why `(L, N, K)` Layout Instead of `(N, L, K)`?
196 +
197 +FlashInfer's `routing_replay_out` parameter expects a contiguous `(N, K)` tensor per layer. With `(L, N, K)` layout, `buffer[layer_id]` gives a contiguous `(N, K)` view with zero-copy slicing. The previous `(N, L, K)` layout would require non-contiguous indexing or an explicit copy.
198 +
199 +### Why int16 Instead of int32?
200 +
201 +Expert IDs are small integers (typically 0-255 for models with up to 256 experts). `int16` supports up to 32,767 experts — far more than any current model — while halving GPU memory usage and D2H bandwidth compared to `int32`.
202 +
203 +### Why Split Prompt and Generation Routing?
204 +
205 +RL training pipelines process prompt and generation routing separately:
206 +
207 +- Prompt routing reconstructs the shared forward pass for the input.
208 +- Generation routing reconstructs each sampled trajectory.
209 +
210 +With `n > 1` completions, all completions share the same prompt routing. Duplicating it per completion would waste memory proportional to `n * prompt_len * L * K`. Instead, `prompt_routed_experts` is stored once on `RequestOutput` and shared.
211 +
212 +### Why Async D2H Instead of Synchronous Copy?
213 +
214 +A synchronous `.cpu()` call forces the GPU to drain its command queue before the copy can begin, stalling the pipeline. The async approach:
215 +
216 +1. Issues the copy on a separate CUDA stream (non-blocking to the main compute stream).
217 +2. Defers the host-side scatter to the *next* step, by which time the copy has finished.
218 +
219 +This means the D2H transfer overlaps entirely with the next forward pass, adding near-zero latency to the critical path.
220 +
221 +### Why All TP Ranks Get a Device Buffer?
222 +
223 +CUDA graph capture records the exact sequence of kernel calls and their arguments. If only rank 0 had a device buffer, the `FusedMoE` layer would take a different code path on rank 0 vs. other ranks (one writes to a buffer, others don't). This asymmetry causes different CUDA graph structures across ranks, which can lead to NCCL deadlocks during collective operations inside the graph. Giving all ranks a real buffer ensures symmetric graph structure. Only rank 0 does the D2H copy and host cache management.
224 +
225 +## Performance
226 +
227 +Routing replay adds a small overhead from the device buffer writes and async D2H copies. On tested configurations:
228 +
229 +- **Throughput overhead** (random data, ISL=1024, OSL=1024): **~2%**
230 +- **Memory overhead** (int16 buffer, 40 layers, 8192 tokens, top-22): **~14 MB per GPU**
231 +- **Accuracy impact** (GSM8K): **Zero** (pass@1 identical with and without routing replay)
232 +
233 +The overhead is dominated by the per-layer `.copy_()` during the forward pass. The async D2H pipeline runs entirely in the background.
234 +
235 +## Supported Configurations
236 +
237 +| Configuration | Supported |
238 +|------------------------------------------|-----------------------------------------------------------|
239 +| BF16 Triton MoE (non-monolithic) | Yes |
240 +| FP8/MXFP8 FlashInfer MoE (monolithic) | Yes (requires FlashInfer with `routing_replay_out`) |
241 +| CUDA graphs | Yes |
242 +| Multi-node tensor parallelism | Yes |
243 +| Data parallelism (DP) | Yes |
244 +| Expert parallelism (EP) | Yes |
245 +| Prefix caching | Yes (cached positions marked with `-1` sentinel) |
246 +| MTP speculative decoding | Yes (gen routing trimmed to accepted tokens) |
247 +| `n > 1` (multiple completions) | Yes (prompt routing shared, gen routing per-completion) |
248 +
249 +## Limitations
250 +
251 +- **Streaming**: Routing data is only available when the request finishes (not streamed incrementally).
252 +- **V1 engine only**: Routing replay is implemented for the vLLM V1 engine.
253 +- **Preempted requests**: When a request is preempted by the scheduler (and later resumed via re-prefill), any routing already accumulated in the worker's host cache for that request is dropped without being emitted. The consumer sees `routed_experts=None` for the resumed request with no other signal. Partial-rollout and async-RL pipelines that rely on routing for preempted requests should either disable preemption (`--no-enable-chunked-prefill` / sufficient KV headroom) or reconstruct routing on the resumed prefill.
254 +- **Async scheduling**: Not supported; rejected at config time. The worker-side stop predicate reads `req_state.output_token_ids[-1]`, which under async scheduling is the placeholder `-1` until `AsyncGPUModelRunnerOutput` resolves the real sampled token, so EOS / stop-token finishes would silently drop routing. Use sync scheduling (the default when `--enable-return-routed-experts` is set, or set explicitly with the appropriate scheduler config).
255 +- **Sequence parallelism / naive DP MoE dispatch**: Not supported on the FusedMoE layer; rejected at bind time. SP shards `topk_ids` along dim 0 across the TP group so each rank only captures `1/sp_size` of the rows; naive DP dispatch all-gathers tokens across DP ranks before routing, so `topk_ids.shape[0]` exceeds the per-rank buffer size. Both raise `NotImplementedError` from `bind_routing_capture_to_model`.
256 +- **Pipeline / prefill-context / decode-context parallelism**: Not yet validated; rejected at config time.
257 +
258 +## CLI Reference
259 +
260 +| Flag | Description |
261 +|------------------------------------|------------------------------------------------------------------------|
262 +| `--enable-return-routed-experts` | Enable routing replay capture and return expert IDs in API responses. |
263 +
264 +## API Reference
265 +
266 +### Completions (`/v1/completions`)
267 +
268 +**Response-level field:**
269 +
270 +| Field | Type | Description |
271 +|---------------------------|-------------------------------------|-----------------------------------------------------------------------------|
272 +| `prompt_routed_experts` | `list[list[list[int]]]` or `null` | Expert IDs for prompt tokens. Shape: `[prompt_len, num_moe_layers, top_k]`. |
273 +
274 +**Choice-level field:**
275 +
276 +| Field | Type | Description |
277 +|--------------------|-------------------------------------|-------------------------------------------------------------------------------|
278 +| `routed_experts` | `list[list[list[int]]]` or `null` | Expert IDs for generated tokens. Shape: `[gen_len, num_moe_layers, top_k]`. |
279 +
280 +### Chat Completions (`/v1/chat/completions`)
281 +
282 +Same fields as above on `ChatCompletionResponse` and `ChatCompletionResponseChoice`.
283 +
284 +### Python SDK
285 +
286 +| Object | Field | Type | Description |
287 +|----------------------|---------------------------|--------------------------|-----------------------------|
288 +| `RequestOutput` | `prompt_routed_experts` | `np.ndarray` or `None` | `[prompt_len, L, K]` i16 |
289 +| `CompletionOutput` | `routed_experts` | `np.ndarray` or `None` | `[gen_len, L, K]` int16 |
tests/model_executor/test_routed_experts_capture.py
+132 -215
@@ -1,245 +1,162 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import types
4 -from types import SimpleNamespace
5 -from unittest.mock import patch
3
4 import pytest
5 import torch
6
10 -from vllm.distributed.eplb.eplb_state import EplbLayerState
11 -from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
12 -from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
13 - RoutedExpertsCapturer,
14 -)
15 -from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter
16 -
7 pytestmark = pytest.mark.cpu_test
8
19 -_REC_MODULE = "vllm.model_executor.layers.fused_moe.routed_experts_capturer"
20 -
21 -
22 -def _capturer_with_buffer(
23 - *,
24 - max_tokens: int = 8,
25 - num_layers: int = 4,
26 - num_experts_per_tok: int = 2,
27 - dp_rank: int = 0,
28 -) -> RoutedExpertsCapturer:
29 - c = RoutedExpertsCapturer()
30 - c.dp_rank = dp_rank
31 - c._device_buffer = torch.full(
32 - (max_tokens, num_layers, num_experts_per_tok),
33 - -1,
34 - dtype=torch.int32,
35 - )
36 - return c
37 -
38 -
39 -class DummyRouter(BaseRouter):
40 - @property
41 - def routing_method_type(self) -> RoutingMethodType:
42 - return RoutingMethodType.FUSED_TOPK
43 -
44 - def _compute_routing(
45 - self, hidden_states, router_logits, indices_type, *, input_ids=None
46 - ):
47 - topk_ids = torch.tensor([[1, 2], [3, 4]], dtype=torch.int64)
48 - topk_weights = torch.ones_like(topk_ids, dtype=torch.float32)
49 - return topk_weights, topk_ids
50 -
51 - def _apply_eplb_mapping(self, topk_ids: torch.Tensor) -> torch.Tensor:
52 - # Make mapping observable without requiring CUDA EPLB path.
53 - return topk_ids + 10
54 -
55 -
56 -def _make_router() -> DummyRouter:
57 - return DummyRouter(
58 - top_k=2,
59 - global_num_experts=16,
60 - eplb_state=EplbLayerState(),
61 - enable_eplb=False,
62 - indices_type_getter=None,
63 - )
64 -
65 -
66 -def test_base_router_capture_pre_eplb_mapping():
67 - router = _make_router()
68 - captured = []
69 -
70 - def capture_fn(ids):
71 - captured.append(ids.clone())
72 -
73 - router.set_capture_fn(capture_fn)
74 - topk_weights, topk_ids = router.select_experts(
75 - hidden_states=torch.empty(1),
76 - router_logits=torch.empty(1),
77 - )
78 -
79 - assert topk_weights.shape == topk_ids.shape
80 - assert len(captured) == 1
81 - assert torch.equal(captured[0], torch.tensor([[1, 2], [3, 4]]))
82 - assert torch.equal(topk_ids, torch.tensor([[11, 12], [13, 14]]))
83 -
84 -
85 -def test_base_router_capture_with_eplb_enabled():
86 - router = _make_router()
87 - router.enable_eplb = True
88 - router.eplb_state.expert_load_view = torch.zeros(32, dtype=torch.int64)
89 - router.eplb_state.logical_to_physical_map = torch.arange(32).view(32, 1)
90 - router.eplb_state.logical_replica_count = torch.ones(32, dtype=torch.int64)
91 - router.eplb_state.should_record_tensor = torch.ones((), dtype=torch.bool)
92 -
93 - captured = []
94 -
95 - def capture_fn(ids):
96 - captured.append(ids.clone())
97 -
98 - router.set_capture_fn(capture_fn)
99 - _, topk_ids = router.select_experts(
100 - hidden_states=torch.empty(1),
101 - router_logits=torch.empty(1),
102 - )
103 -
104 - assert len(captured) == 1
105 - # Capture should see logical ids pre-EPLB mapping.
106 - assert torch.equal(captured[0], torch.tensor([[1, 2], [3, 4]]))
107 - # Our DummyRouter mapping adds +10.
108 - assert torch.equal(topk_ids, torch.tensor([[11, 12], [13, 14]]))
109 -
110 -
111 -def test_gpu_model_runner_binds_router_capture(monkeypatch):
112 - from vllm.v1.worker import gpu_model_runner as gmr
9
114 - class DummyFusedMoE:
115 - def __init__(self):
116 - self.layer_id = 7
117 - self.router = _make_router()
10 +def test_bind_routing_capture_to_model_sets_layer_view(monkeypatch):
11 + import vllm.model_executor.layers.fused_moe.layer as fused_moe_layer
12 + import vllm.model_executor.layers.fused_moe.routed_experts_capturer as rec_mod
13
119 - class DummyCapturer:
120 - def __init__(self):
121 - self.calls = []
14 + class _DummyMoEConfig:
15 + is_sequence_parallel = False
16 + dp_size = 1
17
123 - def capture(self, layer_id, topk_ids):
124 - self.calls.append((layer_id, topk_ids))
18 + class _DummyQuantMethod:
19 + supports_internal_mk = True
20
126 - dummy_module = DummyFusedMoE()
21 + class DummyFusedMoE:
22 + _routing_replay_out: torch.Tensor
23
128 - # Patch the runtime import inside _bind_routed_experts_capturer.
129 - import vllm.model_executor.layers.fused_moe.layer as fused_moe_layer
24 + def __init__(self, moe_layer_id):
25 + self.moe_layer_id = moe_layer_id
26 + self.moe_config = _DummyMoEConfig()
27 + self.quant_method = _DummyQuantMethod()
28
29 monkeypatch.setattr(fused_moe_layer, "FusedMoE", DummyFusedMoE)
30
133 - dummy_self = types.SimpleNamespace(
134 - compilation_config=types.SimpleNamespace(
135 - static_forward_context={"dummy": dummy_module}
136 - )
137 - )
31 + num_layers, num_tokens, top_k = 4, 8, 2
32 + buffer = torch.zeros((num_layers, num_tokens, top_k), dtype=torch.int16)
33
139 - capturer = DummyCapturer()
140 - gmr.GPUModelRunner._bind_routed_experts_capturer(dummy_self, capturer)
34 + class DummyDeviceCache:
35 + def __init__(self, buf):
36 + self.buffer = buf
37
142 - assert dummy_module.router.capture_fn is not None
143 - dummy_module.router.capture_fn(torch.tensor([[5, 6]]))
38 + class DummyCapturer:
39 + def get_device_cache(self):
40 + return DummyDeviceCache(buffer)
41
145 - assert len(capturer.calls) == 1
146 - layer_id, topk_ids = capturer.calls[0]
147 - assert layer_id == 7
148 - assert torch.equal(topk_ids, torch.tensor([[5, 6]]))
42 + monkeypatch.setattr(rec_mod, "get_global_experts_capturer", lambda: DummyCapturer())
43
44 + m0 = DummyFusedMoE(moe_layer_id=0)
45 + m2 = DummyFusedMoE(moe_layer_id=2)
46
151 -def test_gpu_model_runner_binding_stage(monkeypatch):
152 - from vllm.v1.worker import gpu_model_runner as gmr
47 + class DummyModel:
48 + def modules(self):
49 + return iter([m0, m2])
50
154 - class DummyFusedMoE:
155 - def __init__(self):
156 - self.layer_id = 11
157 - self.router = _make_router()
51 + rec_mod.bind_routing_capture_to_model(DummyModel())
52 +
53 + assert torch.equal(m0._routing_replay_out, buffer[0])
54 + assert torch.equal(m2._routing_replay_out, buffer[2])
55 +
56 +
57 +def test_bind_routing_capture_to_model_noop_when_disabled(monkeypatch):
58 + import vllm.model_executor.layers.fused_moe.routed_experts_capturer as rec_mod
59
60 class DummyCapturer:
160 - def __init__(self):
161 - self.calls = []
61 + def get_device_cache(self):
62 + return None
63
163 - def capture(self, layer_id, topk_ids):
164 - self.calls.append((layer_id, topk_ids))
64 + monkeypatch.setattr(rec_mod, "get_global_experts_capturer", lambda: DummyCapturer())
65
166 - dummy_module = DummyFusedMoE()
66 + class DummyModel:
67 + def modules(self):
68 + return iter([])
69
168 - import vllm.model_executor.layers.fused_moe.layer as fused_moe_layer
70 + rec_mod.bind_routing_capture_to_model(DummyModel())
71
170 - monkeypatch.setattr(fused_moe_layer, "FusedMoE", DummyFusedMoE)
72
172 - dummy_self = types.SimpleNamespace(
173 - compilation_config=types.SimpleNamespace(
174 - static_forward_context={"dummy": dummy_module}
73 +# =========================================================================
74 +# Tests for device-cache routing replay architecture
75 +# =========================================================================
76 +
77 +
78 +class TestRoutedExpertsDeviceCache:
79 + """Tests for _RoutedExpertsDeviceCache (GPU buffer for routing data)."""
80 +
81 + def test_allocation_shape_and_dtype(self):
82 + """Device cache allocates (L, N, K) int16 buffer."""
83 + from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
84 + _RoutedExpertsDeviceCache,
85 + )
86 +
87 + cache = _RoutedExpertsDeviceCache(
88 + num_hidden_layers=40,
89 + max_num_batched_tokens=8192,
90 + num_experts_per_tok=8,
91 + device="cpu",
92 + )
93 + assert cache.buffer.shape == (40, 8192, 8)
94 + assert cache.buffer.dtype == torch.int16
95 +
96 + def test_per_layer_view_is_contiguous(self):
97 + """buffer[layer_id] gives contiguous (N, K) view for FlashInfer."""
98 + from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
99 + _RoutedExpertsDeviceCache,
100 + )
101 +
102 + cache = _RoutedExpertsDeviceCache(
103 + num_hidden_layers=40,
104 + max_num_batched_tokens=8192,
105 + num_experts_per_tok=8,
106 + device="cpu",
107 + )
108 + layer_view = cache.buffer[0]
109 + assert layer_view.is_contiguous()
110 + assert layer_view.shape == (8192, 8)
111 +
112 +
113 +class TestRoutedExpertsHostCache:
114 + """Tests for _RoutedExpertsHostCache (per-request numpy buffer)."""
115 +
116 + def test_sentinel_initialization(self):
117 + """Host cache initializes with zeros by default."""
118 + import numpy as np
119 +
120 + from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
121 + _RoutedExpertsHostCache,
122 + )
123 +
124 + cache = _RoutedExpertsHostCache(
125 + num_hidden_layers=40,
126 + num_experts_per_tok=8,
127 + max_model_len=1024,
128 + )
129 + buf = cache.get_or_grow_buffer("req1", max_pos=100)
130 + assert buf.dtype == np.int16
131 + assert (buf == 0).all(), "Host cache must initialize with zeros"
132 +
133 + def test_grow_preserves_existing_data(self):
134 + """Growing the buffer preserves previously written data."""
135 + from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
136 + _RoutedExpertsHostCache,
137 + )
138 +
139 + cache = _RoutedExpertsHostCache(
140 + num_hidden_layers=40,
141 + num_experts_per_tok=8,
142 + max_model_len=1024,
143 + )
144 + buf = cache.get_or_grow_buffer("req1", max_pos=50)
145 + buf[0, 0, 0] = 42
146 + buf2 = cache.get_or_grow_buffer("req1", max_pos=200)
147 + assert buf2[0, 0, 0] == 42, "Data lost during buffer grow"
148 +
149 + def test_free_request_removes_buffer(self):
150 + """Freeing a request removes its buffer."""
151 + from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
152 + _RoutedExpertsHostCache,
153 + )
154 +
155 + cache = _RoutedExpertsHostCache(
156 + num_hidden_layers=40,
157 + num_experts_per_tok=8,
158 + max_model_len=1024,
159 )
176 - )
177 -
178 - # Before binding, no capture hook.
179 - assert dummy_module.router.capture_fn is None
180 -
181 - capturer = DummyCapturer()
182 - gmr.GPUModelRunner._bind_routed_experts_capturer(dummy_self, capturer)
183 -
184 - # After binding, hook should exist and be callable.
185 - assert callable(dummy_module.router.capture_fn)
186 - dummy_module.router.capture_fn(torch.tensor([[9, 10]]))
187 - assert len(capturer.calls) == 1
188 -
189 -
190 -def test_routed_experts_capturer_single_dp_no_metadata():
191 - """dp_metadata is None: capture writes the full topk_ids rows."""
192 - capturer = _capturer_with_buffer(dp_rank=0)
193 - topk = torch.tensor([[1, 2], [3, 4], [5, 6]], dtype=torch.int32)
194 - ctx = SimpleNamespace(dp_metadata=None)
195 - with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx):
196 - capturer.capture(layer_id=0, topk_ids=topk)
197 - assert torch.equal(capturer._device_buffer[:3, 0, :], topk)
198 - assert capturer._device_buffer[3, 0, 0].item() == -1
199 -
200 -
201 -def test_routed_experts_capturer_dp_naive_concatenated_all_ranks():
202 - """n == sum(num_tokens_dp): slice this rank's segment from concatenated topk."""
203 - capturer = _capturer_with_buffer(dp_rank=1)
204 - num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32)
205 - ctx = SimpleNamespace(
206 - dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp)
207 - )
208 - # Concatenated order: rank0 rows then rank1 rows.
209 - topk = torch.tensor(
210 - [[0, 1], [2, 3], [10, 11], [12, 13], [14, 15]], dtype=torch.int32
211 - )
212 - with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx):
213 - capturer.capture(layer_id=0, topk_ids=topk)
214 - want = topk[2:5]
215 - assert torch.equal(capturer._device_buffer[:3, 0, :], want)
216 -
217 -
218 -def test_routed_experts_capturer_dp_modular_local_tokens():
219 - """n == token_num_per_dp: topk is already local to this DP rank."""
220 - capturer = _capturer_with_buffer(dp_rank=1)
221 - num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32)
222 - ctx = SimpleNamespace(
223 - dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp)
224 - )
225 - topk = torch.tensor([[10, 11], [12, 13], [14, 15]], dtype=torch.int32)
226 - with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx):
227 - capturer.capture(layer_id=0, topk_ids=topk)
228 - assert torch.equal(capturer._device_buffer[:3, 0, :], topk)
229 -
230 -
231 -def test_routed_experts_capturer_dp_unexpected_batch_raises():
232 - """Mismatch between topk batch dim and DP layout: fail fast."""
233 - capturer = _capturer_with_buffer(dp_rank=0)
234 - num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32)
235 - ctx = SimpleNamespace(
236 - dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp)
237 - )
238 - # total=5, local=2: n=1 matches neither naive (5) nor modular (2).
239 - topk = torch.tensor([[1, 2]], dtype=torch.int32)
240 - with (
241 - patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx),
242 - pytest.raises(AssertionError, match="unexpected topk_ids batch dim"),
243 - ):
244 - capturer.capture(layer_id=0, topk_ids=topk)
245 - assert capturer._device_buffer[0, 0, 0].item() == -1
160 + cache.get_or_grow_buffer("req1", max_pos=50)
161 + cache.free_request("req1")
162 + assert cache.get_buffer("req1") is None
vllm/config/vllm.py
+51
@@ -1157,6 +1157,12 @@ class VllmConfig:
1157 if envs.VLLM_USE_V2_MODEL_RUNNER:
1158 self._validate_v2_model_runner()
1159
1160 + if (
1161 + self.model_config is not None
1162 + and self.model_config.enable_return_routed_experts
1163 + ):
1164 + self._validate_return_routed_experts()
1165 +
1166 # Re-compute compile ranges after platform-specific config updates
1167 # (e.g., XPU may lower max_num_batched_tokens when MLA is enabled)
1168 self._set_compile_ranges()
@@ -1846,6 +1852,51 @@ class VllmConfig:
1852 + ", ".join(unsupported)
1853 )
1854
1855 + def _validate_return_routed_experts(self) -> None:
1856 + """Reject parallelism configurations not yet validated with
1857 + --enable-return-routed-experts.
1858 +
1859 + Validated scope (PR #39917): TP, EP, DP, single-node and multi-node,
1860 + prefix caching, and speculative decoding (MTP validated end-to-end;
1861 + Eagle/Eagle3/Ngram/Medusa supported by construction since the
1862 + routing buffer is bound only to the target model and verified-token
1863 + routing lands at the correct positions during the main forward).
1864 +
1865 + Out-of-scope (block until validated): PP > 1, prefill context
1866 + parallelism (PCP) > 1, decode context parallelism (DCP) > 1,
1867 + async scheduling.
1868 + """
1869 + unsupported: list[str] = []
1870 +
1871 + if self.parallel_config.pipeline_parallel_size > 1:
1872 + unsupported.append(
1873 + "pipeline parallelism "
1874 + f"(pipeline_parallel_size="
1875 + f"{self.parallel_config.pipeline_parallel_size})"
1876 + )
1877 + if self.parallel_config.prefill_context_parallel_size > 1:
1878 + unsupported.append(
1879 + "prefill context parallelism "
1880 + f"(prefill_context_parallel_size="
1881 + f"{self.parallel_config.prefill_context_parallel_size})"
1882 + )
1883 + if self.parallel_config.decode_context_parallel_size > 1:
1884 + unsupported.append(
1885 + "decode context parallelism "
1886 + f"(decode_context_parallel_size="
1887 + f"{self.parallel_config.decode_context_parallel_size})"
1888 + )
1889 + if self.scheduler_config.async_scheduling:
1890 + unsupported.append("async scheduling")
1891 +
1892 + if unsupported:
1893 + raise ValueError(
1894 + "--enable-return-routed-experts is not yet validated with: "
1895 + + ", ".join(unsupported)
1896 + + ". Disable these features or omit "
1897 + "--enable-return-routed-experts."
1898 + )
1899 +
1900 def validate_block_size(self) -> None:
1901 """Validate block_size against DCP and mamba constraints.
1902
vllm/entrypoints/openai/chat_completion/protocol.py
+4
@@ -92,12 +92,16 @@ class ChatCompletionResponseChoice(OpenAIBaseModel):
92 # not part of the OpenAI spec but is useful for tracing the tokens
93 # in agent scenarios
94 token_ids: list[int] | None = None
95 + routed_experts: list[list[list[int]]] | None = None # [gen_len, num_layers, top_k]
96
97
98 class ChatCompletionResponse(OpenAIBaseModel):
99 id: str = Field(default_factory=lambda: f"chatcmpl-{random_uuid()}")
100 object: Literal["chat.completion"] = "chat.completion"
101 created: int = Field(default_factory=lambda: int(time.time()))
102 + prompt_routed_experts: list[list[list[int]]] | None = (
103 + None # [prompt_len, num_layers, top_k]
104 + )
105 model: str
106 choices: list[ChatCompletionResponseChoice]
107 service_tier: Literal["auto", "default", "flex", "scale", "priority"] | None = None
vllm/entrypoints/openai/chat_completion/serving.py
+15
@@ -1088,6 +1088,11 @@ class OpenAIServingChat(OpenAIServing):
1088 token_ids=(
1089 as_list(output.token_ids) if request.return_token_ids else None
1090 ),
1091 + routed_experts=(
1092 + output.routed_experts.tolist()
1093 + if output.routed_experts is not None
1094 + else None
1095 + ),
1096 )
1097 choices.append(choice_data)
1098 continue
@@ -1309,6 +1314,11 @@ class OpenAIServingChat(OpenAIServing):
1314 token_ids=(
1315 as_list(output.token_ids) if request.return_token_ids else None
1316 ),
1317 + routed_experts=(
1318 + output.routed_experts.tolist()
1319 + if output.routed_experts is not None
1320 + else None
1321 + ),
1322 )
1323 choice_data = maybe_filter_parallel_tool_calls(choice_data, request)
1324
@@ -1348,6 +1358,10 @@ class OpenAIServingChat(OpenAIServing):
1358
1359 request_metadata.final_usage_info = usage
1360
1361 + prompt_routed_experts = None
1362 + if final_res.prompt_routed_experts is not None:
1363 + prompt_routed_experts = final_res.prompt_routed_experts.tolist()
1364 +
1365 response = ChatCompletionResponse(
1366 id=request_id,
1367 created=created_time,
@@ -1360,6 +1374,7 @@ class OpenAIServingChat(OpenAIServing):
1374 final_res.prompt_token_ids if request.return_token_ids else None
1375 ),
1376 kv_transfer_params=final_res.kv_transfer_params,
1377 + prompt_routed_experts=prompt_routed_experts,
1378 )
1379
1380 # Log complete response if output logging is enabled
vllm/entrypoints/openai/completion/protocol.py
+4
@@ -468,12 +468,16 @@ class CompletionResponseChoice(OpenAIBaseModel):
468 token_ids: list[int] | None = None # For response
469 prompt_logprobs: list[dict[int, Logprob] | None] | None = None
470 prompt_token_ids: list[int] | None = None # For prompt
471 + routed_experts: list[list[list[int]]] | None = None # [gen_len, num_layers, top_k]
472
473
474 class CompletionResponse(OpenAIBaseModel):
475 id: str = Field(default_factory=lambda: f"cmpl-{random_uuid()}")
476 object: Literal["text_completion"] = "text_completion"
477 created: int = Field(default_factory=lambda: int(time.time()))
478 + prompt_routed_experts: list[list[list[int]]] | None = (
479 + None # [prompt_len, num_layers, top_k]
480 + )
481 model: str
482 choices: list[CompletionResponseChoice]
483 service_tier: Literal["auto", "default", "flex", "scale", "priority"] | None = None
vllm/entrypoints/openai/completion/serving.py
+11
@@ -542,6 +542,11 @@ class OpenAIServingCompletion(OpenAIServing):
542 token_ids=(
543 as_list(output.token_ids) if request.return_token_ids else None
544 ),
545 + routed_experts=(
546 + output.routed_experts.tolist()
547 + if output.routed_experts is not None
548 + else None
549 + ),
550 )
551 choices.append(choice_data)
552
@@ -565,8 +570,13 @@ class OpenAIServingCompletion(OpenAIServing):
570 )
571
572 request_metadata.final_usage_info = usage
573 + prompt_routed_experts = None
574 if final_res_batch:
575 kv_transfer_params = final_res_batch[0].kv_transfer_params
576 + pre = final_res_batch[0].prompt_routed_experts
577 + if pre is not None:
578 + prompt_routed_experts = pre.tolist()
579 +
580 return CompletionResponse(
581 id=request_id,
582 created=created_time,
@@ -575,6 +585,7 @@ class OpenAIServingCompletion(OpenAIServing):
585 usage=usage,
586 system_fingerprint=self.system_fingerprint,
587 kv_transfer_params=kv_transfer_params,
588 + prompt_routed_experts=prompt_routed_experts,
589 )
590
591 def _create_completion_logprobs(
vllm/model_executor/layers/fused_moe/layer.py
+7
@@ -246,6 +246,9 @@ class FusedMoE(PluggableLayer):
246 not supported by the router (or the experts).
247 """
248
249 + # Auto-incrementing layer ID for routing replay buffer binding.
250 + _next_moe_layer_id: int = 0
251 +
252 # --8<-- [end:fused_moe]
253
254 def __init__(
@@ -290,6 +293,10 @@ class FusedMoE(PluggableLayer):
293 ):
294 super().__init__()
295
296 + # Assign unique layer ID for routing replay buffer binding.
297 + self.moe_layer_id = FusedMoE._next_moe_layer_id
298 + FusedMoE._next_moe_layer_id += 1
299 +
300 if params_dtype is None:
301 params_dtype = torch.get_default_dtype()
302 self.params_dtype = params_dtype
vllm/model_executor/layers/fused_moe/routed_experts_capturer.py
+760 -298
@@ -1,371 +1,833 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -# Adapted from
4 -# https://github.com/sgl-project/sglang/blob/bed301a5acaa9577c9aa706468bdf242f6a43051/python/sglang/srt/layers/moe/routed_experts_capturer.py
5 -
3 from __future__ import annotations
4
8 -import fcntl
5 +import contextlib
6 import logging
10 -import os
11 -import tempfile
12 -from collections.abc import Generator
13 -from contextlib import contextmanager
14 -from multiprocessing import shared_memory
15 -from unittest.mock import patch
7 +from abc import ABC, abstractmethod
8
9 import numpy as np
10 import torch
11 +import torch.distributed
12
20 -from vllm.config import VllmConfig
21 -from vllm.distributed import get_tensor_model_parallel_rank
22 -from vllm.forward_context import get_forward_context
23 -from vllm.platforms import current_platform
13 +from vllm.config.model import ModelConfig
14
15 logger = logging.getLogger(__name__)
16
17
28 -def _get_num_experts_per_tok(hf_config) -> int:
29 - """Resolve the per-token expert count from the HF config.
18 +# ---------------------------------------------------------------------------
19 +# Custom op for routing capture -- traceable by torch.compile / Dynamo.
20 +#
21 +# Registered as a formal custom op so that torch.compile traces through it
22 +# cleanly without graph breaks. ALL TP ranks call this op with a real
23 +# device buffer to ensure identical CUDA graph structure (symmetry).
24 +# Non-rank-0 buffers are written but never read for D2H.
25 +# ---------------------------------------------------------------------------
26
31 - Different model families store this under different attribute names
32 - (e.g. ``num_experts_per_tok`` for DeepSeek, ``top_k_experts`` for Gemma 4).
33 - """
34 - val = getattr(hf_config, "num_experts_per_tok", None)
35 - if val is None:
36 - val = getattr(hf_config, "top_k_experts", None)
37 - if val is None:
38 - raise ValueError(
39 - "Cannot determine num_experts_per_tok: HF config has neither "
40 - "'num_experts_per_tok' nor 'top_k_experts'"
41 - )
42 - return val
27
28 +@torch.library.custom_op("vllm::capture_routing", mutates_args={"buffer"})
29 +def capture_routing_op(
30 + buffer: torch.Tensor,
31 + topk_ids: torch.Tensor,
32 + layer_id: int,
33 + batch_size: int,
34 +) -> None:
35 + buffer[layer_id, :batch_size, :].copy_(
36 + topk_ids[:batch_size].to(buffer.dtype), non_blocking=True
37 + )
38
45 -# Constants
46 -_TMP_DIR = tempfile.gettempdir()
47 -_LOCK_FILE_PREFIX = os.path.join(_TMP_DIR, "vllm_routed_experts")
48 -_BUFFER_PREFIX = "vllm_routed_experts_buffer"
39
50 -# Global singleton instances
51 -_global_experts_capturer: RoutedExpertsCapturer | None = None
52 -_global_experts_reader: RoutedExpertsReader | None = None
40 +@capture_routing_op.register_fake
41 +def _capture_routing_op_fake(
42 + buffer: torch.Tensor,
43 + topk_ids: torch.Tensor,
44 + layer_id: int,
45 + batch_size: int,
46 +) -> None:
47 + pass
48
49
55 -@contextmanager
56 -def _file_lock(lock_file: str, mode: str = "wb+") -> Generator[None, None, None]:
57 - """Context manager for file-based locking."""
58 - with open(lock_file, mode) as fp:
59 - fcntl.flock(fp, fcntl.LOCK_EX)
60 - try:
61 - yield
62 - finally:
63 - fcntl.flock(fp, fcntl.LOCK_UN)
50 +_MB = 1024 * 1024
51
52
66 -def _create_or_attach_shared_memory(
67 - name: str, size: int, lock_file: str
68 -) -> shared_memory.SharedMemory:
69 - """Create or attach to shared memory with proper locking."""
70 - # Ensure lock file exists before acquiring lock
71 - with open(lock_file, "wb"):
72 - pass
53 +class _RoutedExpertsDeviceCache:
54 + """Per-device (GPU) cache for capturing routed expert IDs during forward
55 + pass. Always writes at row 0 so that CUDA graph replay sees the same
56 + addresses that were recorded at capture time.
57 + """
58
74 - with _file_lock(lock_file):
75 - try:
76 - shm = shared_memory.SharedMemory(name=name, create=True, size=size)
77 - except FileExistsError:
78 - shm = shared_memory.SharedMemory(name=name, create=False, size=size)
59 + DTYPE = torch.int16
60
80 - if shm.size != size:
81 - logger.warning(
82 - "Shared memory %s size mismatch; recreating",
83 - name,
84 - )
85 - shm.close()
86 - shm.unlink()
87 - try:
88 - shm = shared_memory.SharedMemory(name=name, create=True, size=size)
89 - logger.info("Created shared memory %s", name)
90 - except FileExistsError:
91 - shm = shared_memory.SharedMemory(name=name, create=False, size=size)
92 - logger.info("Linked to existing shared memory %s", name)
61 + def __init__(
62 + self,
63 + max_num_batched_tokens: int,
64 + num_hidden_layers: int,
65 + num_experts_per_tok: int,
66 + device: str,
67 + ) -> None:
68 + # Layout: (L, N, K) so that buffer[layer_id] is a contiguous (N, K)
69 + # view — required by the FlashInfer routing-replay kernel which
70 + # writes expert IDs assuming contiguous row-major memory.
71 + self.num_hidden_layers = num_hidden_layers
72 + self.buffer = torch.zeros(
73 + (num_hidden_layers, max_num_batched_tokens, num_experts_per_tok),
74 + dtype=self.DTYPE,
75 + device=device,
76 + )
77 + self._finalize_allocation_log()
78 +
79 + def get_buffer_size_bytes(self):
80 + return self.buffer.nbytes
81 +
82 + def capture_fwd_routed_experts(self, layer_id: int, topk_ids: torch.Tensor):
83 + assert layer_id is not None, "capturing routing experts but get layer_id None"
84 + batch, _ = topk_ids.shape
85 + self.buffer[layer_id, :batch, :].copy_(topk_ids, non_blocking=True)
86 +
87 + def _finalize_allocation_log(self):
88 + buf_mb = self.get_buffer_size_bytes() / _MB
89 + logger.info(
90 + "Routing experts device buffer allocated. shape=%s, size=%.2f MB",
91 + tuple(self.buffer.shape),
92 + buf_mb,
93 + )
94
94 - return shm
95
96 +class _RoutedExpertsHostCache:
97 + """Host (CPU) cache using numpy arrays for per-request routing data.
98
97 -class RoutedExpertsCapturer:
99 + Numpy arrays avoid torch dispatcher overhead for scatter operations.
100 + Lazy per-request allocation avoids a massive up-front buffer.
101 """
99 - Capturer for routed experts with device and optional shared memory buffer.
102
101 - This class captures expert routing decisions during model forward passes
102 - and optionally stores them in shared memory for cross-process access.
103 - """
103 + DTYPE = np.int16
104 +
105 + def __init__(
106 + self,
107 + num_hidden_layers: int,
108 + num_experts_per_tok: int,
109 + max_model_len: int,
110 + ) -> None:
111 + self.max_model_len = max_model_len
112 + self.num_hidden_layers = num_hidden_layers
113 + self.num_experts_per_tok = num_experts_per_tok
114 +
115 + self._req_buffers: dict[str, np.ndarray] = {}
116 + self._filled_len: dict[str, int] = {}
117 + self._total_allocated_bytes = 0
118
105 - _instance: RoutedExpertsCapturer | None = None
119 + self._finalize_allocation_log()
120
107 - def __init__(self) -> None:
108 - self._device_buffer: torch.Tensor | None = None
109 - self._shm: shared_memory.SharedMemory | None = None
110 - self._host_buffer_view: np.ndarray | None = None
111 - self._lock_file: str | None = None
121 + def get_buffer_size_bytes(self) -> int:
122 + return self._total_allocated_bytes
123
113 - @classmethod
114 - def create(cls) -> RoutedExpertsCapturer:
115 - """Create a global singleton instance."""
116 - global _global_experts_capturer
117 - if _global_experts_capturer is not None:
118 - raise RuntimeError("Experts capturer already created.")
124 + def get_or_grow_buffer(self, req_id: str, max_pos: int) -> np.ndarray:
125 + required_len = max_pos + 1
126
120 - _global_experts_capturer = cls()
121 - return _global_experts_capturer
127 + if req_id not in self._req_buffers:
128 + buf = np.zeros(
129 + (required_len, self.num_hidden_layers, self.num_experts_per_tok),
130 + dtype=self.DTYPE,
131 + )
132 + self._req_buffers[req_id] = buf
133 + self._total_allocated_bytes += buf.nbytes
134 + return buf
135 +
136 + buf = self._req_buffers[req_id]
137 + if buf.shape[0] >= required_len:
138 + return buf
139 +
140 + new_len = min(max(required_len, buf.shape[0] * 2), self.max_model_len)
141 + new_buf = np.zeros(
142 + (new_len, self.num_hidden_layers, self.num_experts_per_tok),
143 + dtype=self.DTYPE,
144 + )
145 + new_buf[: buf.shape[0]] = buf
146 + self._total_allocated_bytes += new_buf.nbytes - buf.nbytes
147 + self._req_buffers[req_id] = new_buf
148 + return new_buf
149 +
150 + def get_buffer(self, req_id: str) -> np.ndarray | None:
151 + return self._req_buffers.get(req_id)
152 +
153 + def update_filled_len(self, req_id: str, max_pos: int) -> None:
154 + new_len = max_pos + 1
155 + self._filled_len[req_id] = max(self._filled_len.get(req_id, 0), new_len)
156 +
157 + def get_filled_len(self, req_id: str) -> int:
158 + return self._filled_len.get(req_id, 0)
159 +
160 + def free_request(self, req_id: str) -> None:
161 + if req_id in self._req_buffers:
162 + self._total_allocated_bytes -= self._req_buffers.pop(req_id).nbytes
163 + self._filled_len.pop(req_id, None)
164 +
165 + def _finalize_allocation_log(self):
166 + logger.info(
167 + "Routing experts host cache initialized (lazy allocation). "
168 + "max_model_len=%s, layers=%s, experts_per_tok=%s",
169 + self.max_model_len,
170 + self.num_hidden_layers,
171 + self.num_experts_per_tok,
172 + )
173
174 +
175 +class RoutedExpertsCapturer(ABC):
176 @staticmethod
124 - def get_instance() -> RoutedExpertsCapturer | None:
125 - """Get the global singleton instance."""
126 - return _global_experts_capturer
177 + def create(
178 + enable: bool,
179 + model_config: ModelConfig,
180 + num_fused_shared_experts: int,
181 + max_num_batched_tokens: int,
182 + max_model_len: int,
183 + device: str,
184 + shared_host_cache: _RoutedExpertsHostCache | None = None,
185 + skip_host_cache: bool = False,
186 + ):
187 + if enable:
188 + return _RoutedExpertsCapturerReal(
189 + model_config,
190 + max_num_batched_tokens=max_num_batched_tokens,
191 + num_fused_shared_experts=num_fused_shared_experts,
192 + max_model_len=max_model_len,
193 + device=device,
194 + shared_host_cache=shared_host_cache,
195 + skip_host_cache=skip_host_cache,
196 + )
197 + return _RoutedExpertsCapturerNoop()
198 +
199 + @abstractmethod
200 + def capture(self, layer_id: int, topk_ids: torch.Tensor):
201 + raise NotImplementedError
202 +
203 + def get_routed_experts(
204 + self, req_id: str, seqlen: int | None = None, free_slot: bool = True
205 + ):
206 + raise NotImplementedError
207
128 - def init_buffer(
208 + def sync_fwd_experts_buffer_DtoH(
209 self,
130 - max_num_batched_tokens: int,
131 - max_num_kv_tokens: int,
132 - vllm_config: VllmConfig,
133 - ) -> None:
134 - """
135 - Initialize the device buffer and optionally shared memory buffer.
210 + positions: torch.Tensor,
211 + num_scheduled_tokens: dict[str, int],
212 + ):
213 + raise NotImplementedError
214
137 - Args:
138 - max_num_batched_tokens: Maximum number of tokens in a batch.
139 - max_num_kv_tokens: Maximum number of KV tokens for shared memory.
140 - vllm_config: vllm configuration containing layer and expert info.
141 - """
215 + def finalize_pending_copy(self):
216 + raise NotImplementedError
217
143 - if self._device_buffer is not None:
144 - raise RuntimeError("Device buffer has already been initialized")
218 + def get_host_cache(self):
219 + raise NotImplementedError
220
146 - hf_config = vllm_config.model_config.hf_text_config
147 - num_layers = hf_config.num_hidden_layers
148 - num_experts_per_tok = _get_num_experts_per_tok(hf_config)
221 + def get_device_cache(self):
222 + raise NotImplementedError
223
150 - # Initialize device buffer
151 - self._device_buffer = torch.zeros(
152 - (max_num_batched_tokens, num_layers, num_experts_per_tok),
153 - dtype=torch.int32,
154 - device=current_platform.device_type,
155 - )
156 - self.dp_rank = vllm_config.parallel_config.data_parallel_rank
224
158 - if get_tensor_model_parallel_rank() != 0:
159 - return
225 +def _count_moe_layers(hf_config) -> int:
226 + """Count the number of MoE layers in a model.
227
161 - # Initialize shared memory
162 - shape = (max_num_kv_tokens, num_layers, num_experts_per_tok)
163 - buffer_size = int(np.prod(shape)) * np.dtype(np.int32).itemsize
164 - instance_id = vllm_config.instance_id
165 - self._lock_file = f"{_LOCK_FILE_PREFIX}_{instance_id}_{self.dp_rank}.lock"
166 - shm_name = f"{_BUFFER_PREFIX}_{instance_id}_{self.dp_rank}"
228 + Resolves three known config shapes:
229 + - Nemotron-style: an explicit ``layers_block_type`` list with "moe" entries.
230 + - Qwen3MoE / DeepSeek-style sparse: ``decoder_sparse_step > 1`` with optional
231 + ``mlp_only_layers`` exclusions.
232 + - Default: every layer is MoE except those listed in ``mlp_only_layers``.
233 + """
234 + layers_block_type = getattr(hf_config, "layers_block_type", None)
235 + if layers_block_type is not None:
236 + return layers_block_type.count("moe")
237 + n = hf_config.num_hidden_layers
238 + mlp_only = getattr(hf_config, "mlp_only_layers", None) or []
239 + step = getattr(hf_config, "decoder_sparse_step", 1) or 1
240 + if step > 1:
241 + return sum(1 for i in range(n) if (i + 1) % step == 0 and i not in mlp_only)
242 + return n - sum(1 for i in mlp_only if 0 <= i < n)
243 +
244 +
245 +class _RoutedExpertsCapturerReal(RoutedExpertsCapturer):
246 + """Capturer with GPU device cache and CPU host cache.
247 +
248 + Performance strategy -- async D2H with optimized host-cache scatter:
249 +
250 + Every decode step we issue a non-blocking D2H copy on a dedicated
251 + CUDA stream. The scatter into per-request host-cache buffers is
252 + deferred to the start of the NEXT step (by which time the copy has
253 + finished). The scatter loop is optimized with direct scalar access
254 + to avoid numpy slice views, int() conversions, and .max() calls.
255 +
256 + At extraction time (when a request finishes), data is already in a
257 + contiguous host buffer -- just a numpy slice, no concatenation.
258 + """
259
168 - self._shm = _create_or_attach_shared_memory(
169 - shm_name, buffer_size, self._lock_file
170 - )
171 - self._host_buffer_view = np.ndarray(shape, dtype=np.int32, buffer=self._shm.buf)
172 - self._host_buffer_view.fill(0)
260 + def __init__(
261 + self,
262 + model_config: ModelConfig,
263 + max_num_batched_tokens: int,
264 + num_fused_shared_experts: int,
265 + max_model_len: int,
266 + device: str,
267 + shared_host_cache: _RoutedExpertsHostCache | None = None,
268 + skip_host_cache: bool = False,
269 + ):
270 + self.num_fused_shared_experts = num_fused_shared_experts
271 + self.num_hidden_layers = _count_moe_layers(model_config.hf_text_config)
272 + self.num_experts_per_tok = model_config.hf_text_config.num_experts_per_tok
273 + self.max_num_batched_tokens = max_num_batched_tokens
274 + self.max_model_len = max_model_len
275 + self._skip_host_cache = skip_host_cache
276 +
277 + if skip_host_cache:
278 + self.host_cache = None
279 + logger.info("Skipping host cache for device %s (non-rank-0)", device)
280 + elif shared_host_cache is not None:
281 + self.host_cache = shared_host_cache
282 + else:
283 + self.host_cache = _RoutedExpertsHostCache(
284 + num_hidden_layers=self.num_hidden_layers,
285 + num_experts_per_tok=self.num_experts_per_tok,
286 + max_model_len=self.max_model_len,
287 + )
288
174 - logger.debug(
175 - "Created shared memory buffer '%s' with shape %s",
176 - shm_name,
177 - shape,
289 + self.device_cache = _RoutedExpertsDeviceCache(
290 + max_num_batched_tokens=self.max_num_batched_tokens,
291 + num_hidden_layers=self.num_hidden_layers,
292 + num_experts_per_tok=self.num_experts_per_tok,
293 + device=device,
294 )
295
180 - def capture(self, layer_id: int, topk_ids: torch.Tensor) -> None:
181 - """
182 - Capture expert routing decisions for a specific layer.
296 + # ---- Async D2H pipeline (rank-0 only) ----
297 + # Non-rank-0 workers only need the device buffer for symmetric
298 + # CUDA graph capture; they skip the D2H pipeline entirely.
299 + self._has_pending_copy = False
300 + self._pending_positions: np.ndarray | None = None
301 + self._pending_num_scheduled: dict[str, int] | None = None
302 + self._pending_total_tokens: int = 0
303 +
304 + if not skip_host_cache:
305 + # Same (L, N, K) layout as device_cache.buffer.
306 + self._pinned_staging = torch.zeros(
307 + (
308 + self.num_hidden_layers,
309 + max_num_batched_tokens,
310 + self.num_experts_per_tok,
311 + ),
312 + dtype=_RoutedExpertsDeviceCache.DTYPE,
313 + pin_memory=True,
314 + )
315 + # Private device snapshot: source for the async D2H. Decouples
316 + # the in-flight copy from device_cache.buffer, which the next
317 + # step's MoE writes overwrite in place on main_stream.
318 + self._device_staging = torch.empty_like(self.device_cache.buffer)
319 + self._copy_stream = torch.cuda.Stream(device=device)
320 + self._copy_event = torch.cuda.Event()
321 +
322 + pinned_mb = self._pinned_staging.nbytes / _MB
323 + logger.info(
324 + "Routing experts pinned staging buffer allocated. "
325 + "shape=%s, size=%.2f MB",
326 + tuple(self._pinned_staging.shape),
327 + pinned_mb,
328 + )
329 + else:
330 + self._pinned_staging = None
331 + self._device_staging = None
332 + self._copy_stream = None
333 + self._copy_event = None
334 + logger.info(
335 + "Routing experts device-only capturer (rank != 0). "
336 + "Device buffer shape=%s",
337 + tuple(self.device_cache.buffer.shape),
338 + )
339
184 - Args:
185 - layer_id: The layer index.
186 - topk_ids: Tensor of shape (batch_size, num_routed_experts).
187 - """
188 - if self._device_buffer is None:
189 - raise RuntimeError("Buffer not initialized. Call init_buffer() first.")
190 -
191 - ctx = get_forward_context()
192 - if ctx.dp_metadata is None: # single dp
193 - start_loc = 0
194 - end_loc = topk_ids.shape[0]
195 - token_num_per_dp = topk_ids.shape[0]
196 - else: # multi dp
197 - num_tokens_dp = ctx.dp_metadata.num_tokens_across_dp_cpu
198 - token_num_per_dp = int(num_tokens_dp[self.dp_rank].item())
199 - total = int(num_tokens_dp.sum().item())
200 - n = topk_ids.shape[0]
201 -
202 - if n == total:
203 - # Naive dispatch: all DP ranks' tokens concatenated before routing.
204 - cumsum = torch.cumsum(num_tokens_dp, dim=0)
205 - end_loc = int(cumsum[self.dp_rank].item())
206 - start_loc = end_loc - token_num_per_dp
207 - elif n == token_num_per_dp:
208 - # Modular-kernel path: DP combine happens inside quant_method.apply;
209 - # select_experts only sees this rank's tokens.
210 - start_loc = 0
211 - end_loc = token_num_per_dp
212 - else:
213 - raise AssertionError(
214 - "RoutedExpertsCapturer: unexpected topk_ids batch dim "
215 - f"{n} (expected {total} or {token_num_per_dp} "
216 - f"for dp_rank={self.dp_rank})"
217 - )
340 + def capture(self, layer_id: int, topk_ids: torch.Tensor):
341 + self.device_cache.capture_fwd_routed_experts(layer_id, topk_ids)
342 +
343 + # ------------------------------------------------------------------
344 + # sync_fwd_experts_buffer_DtoH -- called AFTER the forward pass
345 + # ------------------------------------------------------------------
346
219 - if layer_id >= self._device_buffer.shape[1]:
347 + def sync_fwd_experts_buffer_DtoH(
348 + self,
349 + positions: torch.Tensor,
350 + num_scheduled_tokens: dict[str, int],
351 + ):
352 + if self.host_cache is None:
353 return
354
222 - self._device_buffer[:token_num_per_dp, layer_id, :] = topk_ids[
223 - start_loc:end_loc, :
224 - ]
355 + # 1. Finalize previous async copy -- the copy had an entire
356 + # forward pass to complete so event.synchronize() is ~free.
357 + if self._has_pending_copy:
358 + self._copy_event.synchronize()
359 + self._scatter_to_host()
360 + self._has_pending_copy = False
361
226 - def clear_buffer(self) -> None:
227 - """Clear the device buffer."""
228 - if self._device_buffer is not None:
229 - self._device_buffer.zero_()
362 + total_tokens = sum(num_scheduled_tokens.values())
363 + if total_tokens == 0:
364 + return
365
231 - def save_captured_experts(self, indices: np.ndarray) -> None:
232 - """
233 - Save captured experts from device buffer to shared memory.
366 + # 2. Snapshot the device buffer on main_stream into a private
367 + # staging buffer, then issue the D2H from the staging buffer
368 + # on a dedicated copy stream. The snapshot serializes after
369 + # the current step's MoE writes (same stream) and is private
370 + # from the next step's MoE writes, so the in-flight D2H is
371 + # not aliased by step N+1's forward under async scheduling.
372 + main_stream = torch.cuda.current_stream(self._copy_stream.device)
373 + self._device_staging[:, :total_tokens, :].copy_(
374 + self.device_cache.buffer[:, :total_tokens, :], non_blocking=True
375 + )
376 + with torch.cuda.stream(self._copy_stream):
377 + self._copy_stream.wait_stream(main_stream)
378 + self._pinned_staging[:, :total_tokens, :].copy_(
379 + self._device_staging[:, :total_tokens, :], non_blocking=True
380 + )
381 + self._copy_event.record()
382
235 - Args:
236 - indices: Array of indices indicating where to store the data.
237 - """
238 - if get_tensor_model_parallel_rank() != 0:
239 - return
240 - if self._lock_file is None:
241 - raise RuntimeError("Shared memory not initialized.")
242 - if self._host_buffer_view is None:
243 - return
244 - if self._device_buffer is None:
245 - raise RuntimeError("Device buffer not initialized.")
383 + # 3. Save metadata for deferred scatter.
384 + self._pending_positions = positions.numpy().copy()
385 + self._pending_num_scheduled = num_scheduled_tokens
386 + self._pending_total_tokens = total_tokens
387 + self._has_pending_copy = True
388
247 - num_tokens = len(indices)
248 - data = self._device_buffer[:num_tokens, :, :].cpu().numpy()
389 + # ------------------------------------------------------------------
390 + # Optimized scatter into pre-allocated host-cache buffers
391 + # ------------------------------------------------------------------
392
250 - with _file_lock(self._lock_file):
251 - self._host_buffer_view[indices, :, :] = data
393 + def _scatter_to_host(self):
394 + """Scatter D2H data into per-request host cache buffers.
395
253 - def cleanup(self) -> None:
254 - """Explicitly clean up shared memory resources."""
255 - if self._shm is not None:
256 - try:
257 - self._shm.close()
258 - self._shm.unlink()
259 - except Exception:
260 - logger.debug("Exception during cleanup for capturer", exc_info=True)
261 - finally:
262 - self._shm = None
396 + Staging layout is (L, N, K). Host cache layout is (seq_len, L, K).
397 + We transpose the staging slice to (N, L, K) before scattering so
398 + that indexing by token position naturally yields (L, K) rows.
399 + """
400 + # Transpose (L, N, K) -> (N, L, K) for the active token range.
401 + host_values = (
402 + self._pinned_staging[:, : self._pending_total_tokens, :]
403 + .numpy()
404 + .transpose(1, 0, 2)
405 + )
406 + positions_np = self._pending_positions
407 + host_cache = self.host_cache
408 + assert self._pending_num_scheduled is not None
409 + assert positions_np is not None
410 + assert host_cache is not None
411 +
412 + offset = 0
413 + for req_id, n_tokens in self._pending_num_scheduled.items():
414 + if n_tokens == 0:
415 + continue
416 +
417 + if n_tokens == 1:
418 + pos_val = int(positions_np[offset])
419 + buf = host_cache.get_or_grow_buffer(req_id, pos_val)
420 + buf[pos_val] = host_values[offset]
421 + host_cache.update_filled_len(req_id, pos_val)
422 + else:
423 + pos = positions_np[offset : offset + n_tokens]
424 + max_pos = int(pos[-1]) if n_tokens > 0 else 0
425 + if n_tokens > 1:
426 + max_pos = int(pos.max())
427 + buf = host_cache.get_or_grow_buffer(req_id, max_pos)
428 + buf[pos] = host_values[offset : offset + n_tokens]
429 + host_cache.update_filled_len(req_id, max_pos)
430 +
431 + offset += n_tokens
432 +
433 + self._pending_positions = None
434 + self._pending_num_scheduled = None
435 + self._pending_total_tokens = 0
436 +
437 + # ------------------------------------------------------------------
438 + # finalize_pending_copy -- call before reading host cache
439 + # ------------------------------------------------------------------
440 +
441 + def finalize_pending_copy(self):
442 + """Ensure the most recent async D2H copy has been scattered into
443 + host cache buffers. Call before get_routed_experts."""
444 + if self._has_pending_copy:
445 + self._copy_event.synchronize()
446 + self._scatter_to_host()
447 + self._has_pending_copy = False
448 +
449 + # ------------------------------------------------------------------
450 + # Extraction -- O(1), just a numpy slice
451 + # ------------------------------------------------------------------
452 +
453 + def get_routed_experts(
454 + self,
455 + req_id: str,
456 + seqlen: int | None = None,
457 + free_slot: bool = True,
458 + ):
459 + if self.host_cache is None:
460 + return None
461 + buf = self.host_cache.get_buffer(req_id)
462 + if buf is None:
463 + return None
464 + filled = self.host_cache.get_filled_len(req_id)
465 + if filled <= 0:
466 + return None
467 + effective_len = min(filled, seqlen) if seqlen is not None else filled
468 + result = buf[:effective_len].copy()
469 + if free_slot:
470 + self.host_cache.free_request(req_id)
471 + return result
472 +
473 + def get_host_cache(self):
474 + return self.host_cache
475 +
476 + def get_device_cache(self):
477 + return self.device_cache
478 +
479 +
480 +class _RoutedExpertsCapturerNoop(RoutedExpertsCapturer):
481 + def __init__(self):
482 + pass
483
264 - def __del__(self) -> None:
265 - """Clean up shared memory on destruction."""
266 - self.cleanup()
484 + def capture(self, layer_id: int, topk_ids: torch.Tensor):
485 + pass
486
487 + def get_routed_experts(self, req_id: str, seqlen=None, free_slot=True):
488 + return None
489
269 -class RoutedExpertsReader:
270 - """
271 - Reader for routed experts from shared memory.
490 + def sync_fwd_experts_buffer_DtoH(self, positions, num_scheduled_tokens):
491 + pass
492
273 - This class attaches to shared memory created by RoutedExpertsCapturer
274 - and reads expert routing decisions.
275 - """
493 + def finalize_pending_copy(self):
494 + pass
495
277 - _instance: RoutedExpertsReader | None = None
496 + def get_host_cache(self):
497 + return None
498
279 - def __init__(self) -> None:
280 - self._shm: shared_memory.SharedMemory | None = None
281 - self._host_buffer_view: np.ndarray | None = None
282 - self._lock_file: str | None = None
499 + def get_device_cache(self):
500 + pass
501
284 - @classmethod
285 - def create(cls) -> RoutedExpertsReader:
286 - """Create a global singleton instance."""
287 - global _global_experts_reader
288 - if _global_experts_reader is not None:
289 - raise RuntimeError("Experts reader already created.")
502
291 - _global_experts_reader = cls()
292 - return _global_experts_reader
503 +# Global capturer instance (per-process)
504 +_global_expert_capturer: RoutedExpertsCapturer | None = _RoutedExpertsCapturerNoop()
505 +_shared_host_cache: _RoutedExpertsHostCache | None = None
506
294 - @staticmethod
295 - def get_instance() -> RoutedExpertsReader | None:
296 - """Get the global singleton instance."""
297 - if _global_experts_reader is None:
298 - logger.info("Experts reader not initialized.")
299 - return _global_experts_reader
507
301 - def attach_buffer(
302 - self,
303 - max_num_kv_tokens: int,
304 - vllm_config: VllmConfig,
305 - ) -> None:
306 - """
307 - Attach to an existing shared memory buffer.
508 +def get_global_experts_capturer():
509 + return _global_expert_capturer
510
309 - Args:
310 - max_num_kv_tokens: Maximum number of KV tokens.
311 - vllm_config: vllm configuration.
312 - """
313 - if self._shm is not None:
314 - logger.warning("Already attached to shared memory buffer.")
315 - return # Already attached
316 -
317 - hf_config = vllm_config.model_config.hf_text_config
318 - shape = (
319 - max_num_kv_tokens,
320 - hf_config.num_hidden_layers,
321 - _get_num_experts_per_tok(hf_config),
322 - )
511
324 - self.dp_rank = vllm_config.parallel_config.data_parallel_rank
325 - instance_id = vllm_config.instance_id
326 - self._lock_file = f"{_LOCK_FILE_PREFIX}_{instance_id}_{self.dp_rank}.lock"
327 - shm_name = f"{_BUFFER_PREFIX}_{instance_id}_{self.dp_rank}"
512 +def set_global_experts_capturer(capturer: RoutedExpertsCapturer):
513 + global _global_expert_capturer
514 + _global_expert_capturer = capturer
515
329 - with _file_lock(self._lock_file, mode="rb+"):
330 - # Avoid resource_tracker registering the shared memory
331 - with patch(
332 - "multiprocessing.resource_tracker.register",
333 - lambda *args, **kwargs: None,
334 - ):
335 - self._shm = shared_memory.SharedMemory(name=shm_name)
516
337 - self._host_buffer_view = np.ndarray(
338 - shape, dtype=np.int32, buffer=self._shm.buf
339 - )
517 +def extract_routed_experts_for_current_batch(
518 + req_ids: list[str],
519 + requests: dict,
520 + req_id_to_index: dict[str, int],
521 + num_tokens_no_spec: np.ndarray,
522 + max_model_len: int,
523 +) -> dict[str, np.ndarray] | None:
524 + """Extract routed experts for requests predicted to finish this step.
525
341 - def get_routed_experts(self, indices: np.ndarray) -> np.ndarray:
342 - """
343 - Read routed expert data from shared memory.
526 + Checks all stop conditions the scheduler will check (max_tokens,
527 + EOS token, stop tokens, max_model_len) so that every finished
528 + request gets its routing data attached to the ModelRunnerOutput.
529
345 - Args:
346 - indices: Array of indices to read.
530 + Args:
531 + req_ids: Ordered request IDs for the current batch.
532 + requests: Map of req_id to CachedRequestState (read-only).
533 + req_id_to_index: Map of req_id to input batch index.
534 + num_tokens_no_spec: Array of total token counts per request index.
535 + max_model_len: Maximum model sequence length.
536 + """
537 + capturer = get_global_experts_capturer()
538 + if capturer is None:
539 + return None
540 + host_cache = capturer.get_host_cache()
541 + if host_cache is None:
542 + return None
543 +
544 + finishing_req_ids: list[str] = []
545 + for req_id in req_ids:
546 + req_state = requests.get(req_id)
547 + if req_state is None:
548 + continue
549 + sp = req_state.sampling_params
550 + if sp is None:
551 + continue
552 + output_ids = req_state.output_token_ids
553 + if not output_ids:
554 + continue
555 + if len(output_ids) < sp.min_tokens:
556 + continue
557 +
558 + finishing = False
559 + last_token = output_ids[-1]
560 +
561 + # EOS token (mirrors check_stop: eos_token_id is None
562 + # when ignore_eos=True, so this naturally respects that)
563 + if last_token == sp.eos_token_id:
564 + finishing = True
565 +
566 + # Explicit stop token IDs
567 + if not finishing and sp.stop_token_ids and last_token in sp.stop_token_ids:
568 + finishing = True
569 +
570 + # max_tokens / max_model_len length cap
571 + if not finishing:
572 + if sp.max_tokens is not None and len(output_ids) >= sp.max_tokens:
573 + finishing = True
574 + else:
575 + req_idx = req_id_to_index.get(req_id)
576 + if req_idx is not None:
577 + total = num_tokens_no_spec[req_idx]
578 + if total >= max_model_len:
579 + finishing = True
580 +
581 + if finishing:
582 + finishing_req_ids.append(req_id)
583 +
584 + if not finishing_req_ids:
585 + return None
586 +
587 + # At least one request is finishing: ensure the latest async D2H
588 + # copy has been scattered into the host cache.
589 + capturer.finalize_pending_copy()
590 +
591 + result: dict[str, np.ndarray] = {}
592 + for req_id in finishing_req_ids:
593 + seqlen = host_cache.get_filled_len(req_id)
594 + if seqlen <= 0:
595 + continue
596 + experts = capturer.get_routed_experts(req_id, seqlen=seqlen, free_slot=False)
597 + if experts is not None:
598 + result[req_id] = experts
599 +
600 + return result if result else None
601 +
602 +
603 +def free_routing_buffers(
604 + finished_req_ids: set[str],
605 + preempted_req_ids: set[str] | None = None,
606 +) -> None:
607 + """Free host cache buffers for finished and preempted requests.
608 +
609 + Finished requests had their routing data extracted in the previous
610 + step.
611 +
612 + Preempted requests are re-prefilled from scratch when they resume,
613 + so their host-cache buffer is freed here. This means any routing
614 + already accumulated in the host cache for the preempted request is
615 + dropped without being emitted on a ``ModelRunnerOutput`` --
616 + consumers see ``routed_experts=None`` for those requests with no
617 + other signal. Partial-rollout / async-RL pipelines that depend on
618 + receiving routing for preempted requests should treat preemption
619 + as a routing-data loss event and either keep preemption disabled
620 + or reconstruct routing on the resumed prefill.
621 + """
622 + capturer = get_global_experts_capturer()
623 + if capturer is None:
624 + return
625 + host_cache = capturer.get_host_cache()
626 + if host_cache is None:
627 + return
628 +
629 + for req_id in finished_req_ids:
630 + host_cache.free_request(req_id)
631 + if preempted_req_ids:
632 + for req_id in preempted_req_ids:
633 + host_cache.free_request(req_id)
634 +
635 +
636 +def issue_routing_d2h_copy(
637 + input_batch_req_ids: list[str],
638 + num_scheduled_tokens: dict[str, int],
639 + positions: torch.Tensor,
640 + positions_cpu: torch.Tensor,
641 +) -> None:
642 + """Issue async D2H copy of routed experts after the forward pass.
643 +
644 + Called EARLY in the execute_model epilogue so the copy overlaps with
645 + eplb, kv_connector finalization, and draft work.
646 + finalize_pending_copy() + get_routed_experts() happen later in
647 + extract_routed_experts_for_current_batch().
648 + """
649 + capturer = get_global_experts_capturer()
650 + if capturer is None:
651 + return
652 +
653 + ordered = {
654 + req_id: num_scheduled_tokens[req_id]
655 + for req_id in input_batch_req_ids
656 + if req_id in num_scheduled_tokens
657 + }
658 + n = sum(ordered.values())
659 + positions_cpu[:n].copy_(positions[:n])
660 + capturer.sync_fwd_experts_buffer_DtoH(
661 + positions=positions_cpu[:n],
662 + num_scheduled_tokens=ordered,
663 + )
664 +
665 +
666 +def split_routed_experts(
667 + routed_experts: np.ndarray,
668 + prompt_len: int,
669 + num_output_tokens: int | None = None,
670 +) -> tuple[np.ndarray | None, np.ndarray | None]:
671 + """Split routing data into prompt and generation portions.
672 +
673 + Args:
674 + routed_experts: Full routing array of shape (seq_len, L, K).
675 + prompt_len: Number of prompt tokens for the request.
676 + num_output_tokens: Actual number of generated tokens (from
677 + detokenizer). When provided, the generation portion is
678 + clipped to this length — necessary with MTP where the model
679 + runner may capture routing for more tokens than the final
680 + output contains.
681 +
682 + Returns:
683 + (prompt_routed_experts, gen_routed_experts) numpy arrays, either
684 + of which may be None if the corresponding portion is empty.
685 + """
686 + prompt_routed_experts = routed_experts[:prompt_len]
687 + gen_routed_experts = routed_experts[prompt_len:]
688 +
689 + # Clip generation routing to match actual output tokens.
690 + if (
691 + num_output_tokens is not None
692 + and gen_routed_experts.shape[0] > num_output_tokens
693 + and num_output_tokens > 0
694 + ):
695 + gen_routed_experts = gen_routed_experts[:num_output_tokens]
696 +
697 + if prompt_routed_experts.size == 0:
698 + prompt_routed_experts = None
699 + if gen_routed_experts.size == 0:
700 + gen_routed_experts = None
701 +
702 + return prompt_routed_experts, gen_routed_experts
703 +
704 +
705 +def get_shared_host_cache() -> _RoutedExpertsHostCache | None:
706 + return _shared_host_cache
707 +
708 +
709 +def create_shared_host_cache(
710 + model_config: ModelConfig,
711 + max_model_len: int,
712 +) -> _RoutedExpertsHostCache:
713 + global _shared_host_cache
714 + num_hidden_layers = _count_moe_layers(model_config.hf_text_config)
715 + num_experts_per_tok = model_config.hf_text_config.num_experts_per_tok
716 + _shared_host_cache = _RoutedExpertsHostCache(
717 + num_hidden_layers=num_hidden_layers,
718 + num_experts_per_tok=num_experts_per_tok,
719 + max_model_len=max_model_len,
720 + )
721 + return _shared_host_cache
722 +
723 +
724 +def init_routed_experts_capturer_with_shared_cache(
725 + enable: bool,
726 + model_config: ModelConfig,
727 + num_fused_shared_experts: int,
728 + max_num_batched_tokens: int,
729 + max_model_len: int,
730 + device: str,
731 + rank: int = 0,
732 + world_size: int = 1,
733 +) -> RoutedExpertsCapturer:
734 + """Initialize capturer with rank-aware handling (only rank 0 captures)."""
735 + if not enable:
736 + capturer = _RoutedExpertsCapturerNoop()
737 + set_global_experts_capturer(capturer)
738 + return capturer
739 +
740 + if world_size > 1 and rank != 0:
741 + # Non-rank-0 workers get a device-only capturer (no host cache,
742 + # no D2H pipeline) so that ALL ranks have a real device buffer.
743 + # This ensures the custom op call in every MoE layer produces
744 + # identical CUDA graph structure across TP ranks.
745 + logger.info("Creating device-only routed experts capturer for rank %s", rank)
746 + capturer = RoutedExpertsCapturer.create(
747 + enable=True,
748 + model_config=model_config,
749 + num_fused_shared_experts=num_fused_shared_experts,
750 + max_num_batched_tokens=max_num_batched_tokens,
751 + max_model_len=max_model_len,
752 + device=device,
753 + skip_host_cache=True,
754 + )
755 + set_global_experts_capturer(capturer)
756 + return capturer
757 +
758 + capturer = RoutedExpertsCapturer.create(
759 + enable=True,
760 + model_config=model_config,
761 + num_fused_shared_experts=num_fused_shared_experts,
762 + max_num_batched_tokens=max_num_batched_tokens,
763 + max_model_len=max_model_len,
764 + device=device,
765 + skip_host_cache=False,
766 + )
767 + set_global_experts_capturer(capturer)
768 + return capturer
769 +
770 +
771 +def bind_routing_capture_to_model(model) -> None:
772 + """Bind routing capture buffers to all FusedMoE layers in the model.
773 +
774 + Must be called AFTER init_routed_experts_capturer_with_shared_cache()
775 + and BEFORE CUDA graph capture. All TP ranks get a real buffer so
776 + that the custom op call produces identical graph structure.
777 + """
778 + from vllm.model_executor.layers.fused_moe.layer import FusedMoE
779 +
780 + capturer = get_global_experts_capturer()
781 + device_cache = capturer.get_device_cache()
782 + if device_cache is None:
783 + return # routing capture not enabled
784 +
785 + buffer = device_cache.buffer
786 +
787 + # Mark the buffer so CUDA graphs do NOT snapshot/restore its contents.
788 + if hasattr(torch.compiler, "cudagraph_mark_tensor_static"):
789 + torch.compiler.cudagraph_mark_tensor_static(buffer)
790 + elif hasattr(torch._C, "_set_static_address_tag"):
791 + torch._C._set_static_address_tag(buffer, True)
792 + with contextlib.suppress(Exception):
793 + torch._dynamo.mark_static_address(buffer)
794 +
795 + bound = 0
796 + for module in model.modules():
797 + if isinstance(module, FusedMoE) and hasattr(module, "moe_layer_id"):
798 + # Per-FusedMoE configurations not yet validated for routing
799 + # capture. These signals are only set after model init, so a
800 + # config-level guard cannot see them.
801 + if module.moe_config.is_sequence_parallel:
802 + raise NotImplementedError(
803 + "routed-experts capture is not yet validated with "
804 + "sequence parallelism on the FusedMoE layer "
805 + "(moe_config.is_sequence_parallel=True)."
806 + )
807 + if (
808 + module.moe_config.dp_size > 1
809 + and not module.quant_method.supports_internal_mk
810 + ):
811 + raise NotImplementedError(
812 + "routed-experts capture is not yet validated with "
813 + "naive DP dispatch (non-modular quant method "
814 + f"{type(module.quant_method).__name__}, "
815 + f"dp_size={module.moe_config.dp_size})."
816 + )
817
348 - Returns:
349 - Copy of the expert routing data for the given indices.
350 - """
351 - if self._host_buffer_view is None:
352 - raise RuntimeError("Buffer not attached. Call attach_buffer() first.")
353 - if self._lock_file is None:
354 - raise RuntimeError("Lock file not initialized.")
355 -
356 - with _file_lock(self._lock_file, mode="rb+"):
357 - return self._host_buffer_view[indices, :, :].copy()
358 -
359 - def cleanup(self) -> None:
360 - """Explicitly clean up resources (close without unlink)."""
361 - if self._shm is not None:
362 - try:
363 - self._shm.close()
364 - except Exception:
365 - logger.debug("Exception during cleanup for reader", exc_info=True)
366 - finally:
367 - self._shm = None
368 -
369 - def __del__(self) -> None:
370 - """Close shared memory on destruction (do not unlink)."""
371 - self.cleanup()
818 + layer_id = module.moe_layer_id
819 + layer_buf = buffer[layer_id] # (N_max, K)
820 + module._routing_replay_out = layer_buf
821 + # Mark each per-layer view as static so CUDA graphs don't
822 + # snapshot/restore or relocate the buffer during replay.
823 + if hasattr(torch.compiler, "cudagraph_mark_tensor_static"):
824 + torch.compiler.cudagraph_mark_tensor_static(layer_buf)
825 + with contextlib.suppress(Exception):
826 + torch._dynamo.mark_static_address(layer_buf)
827 + bound += 1
828 +
829 + logger.info(
830 + "Bound routing capture buffer to %s FusedMoE layers. Buffer shape=%s",
831 + bound,
832 + tuple(buffer.shape),
833 + )
vllm/model_executor/layers/fused_moe/runner/moe_runner.py
+8
@@ -451,6 +451,10 @@ class MoERunner(MoERunnerInterface):
451 shared_experts_input, SharedExpertsOrder.NO_OVERLAP
452 )
453
454 + # Get routing replay buffer from persistent layer attribute
455 + # (set by bind_routing_capture_to_model during capturer init)
456 + routing_replay_out = getattr(layer, "_routing_replay_out", None)
457 +
458 if self._quant_method.is_monolithic:
459 fused_out = self._quant_method.apply_monolithic(
460 layer=layer,
@@ -465,6 +469,10 @@ class MoERunner(MoERunnerInterface):
469 input_ids=input_ids,
470 )
471
472 + # Write routing data for non-monolithic path (Triton, etc.)
473 + if routing_replay_out is not None:
474 + routing_replay_out[: topk_ids.shape[0]].copy_(topk_ids.to(torch.int16))
475 +
476 # Passing shared_experts_input in case SharedExpertsOrder is
477 # MK_INTERNAL_OVERLAPPED.
478 fused_out = self._quant_method.apply(
vllm/outputs.py
+4
@@ -121,6 +121,7 @@ class RequestOutput:
121 num_cached_tokens: int | None = None,
122 *,
123 kv_transfer_params: dict[str, Any] | None = None,
124 + prompt_routed_experts: np.ndarray | None = None,
125 # Forward compatibility, code that uses args added in new release can
126 # still run with older versions of vLLM without breaking.
127 **kwargs: Any,
@@ -141,12 +142,15 @@ class RequestOutput:
142 self.encoder_prompt_token_ids = encoder_prompt_token_ids
143 self.num_cached_tokens = num_cached_tokens
144 self.kv_transfer_params = kv_transfer_params
145 + self.prompt_routed_experts = prompt_routed_experts
146
147 def add(self, next_output: "RequestOutput", aggregate: bool) -> None:
148 """Merge subsequent RequestOutput into this one"""
149
150 self.finished |= next_output.finished
151 self.kv_transfer_params = next_output.kv_transfer_params
152 + if next_output.prompt_routed_experts is not None:
153 + self.prompt_routed_experts = next_output.prompt_routed_experts
154
155 for next_completion in next_output.outputs:
156 for i, completion in enumerate(self.outputs):
vllm/v1/core/sched/scheduler.py
+7 -70
@@ -7,8 +7,6 @@ from collections.abc import Iterable
7 from dataclasses import replace
8 from typing import Any
9
10 -import numpy as np
11 -
10 from vllm import envs
11 from vllm.compilation.cuda_graph import CUDAGraphStat
12 from vllm.config import VllmConfig
@@ -27,9 +25,6 @@ from vllm.distributed.kv_transfer.kv_connector.v1 import (
25 from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata
26 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
27 from vllm.logger import init_logger
30 -from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
31 - RoutedExpertsReader,
32 -)
28 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
29 from vllm.multimodal.encoder_budget import MultiModalBudget
30 from vllm.v1.core.encoder_cache_manager import (
@@ -52,7 +47,7 @@ from vllm.v1.core.sched.request_queue import (
47 )
48 from vllm.v1.core.sched.utils import check_stop, remove_all
49 from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs
55 -from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig
50 +from vllm.v1.kv_cache_interface import KVCacheConfig
51 from vllm.v1.metrics.perf import ModelMetrics, PerfStats
52 from vllm.v1.metrics.stats import PrefixCacheStats, SchedulerStats
53 from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput
@@ -260,43 +255,6 @@ class Scheduler(SchedulerInterface):
255 if self.log_stats and vllm_config.observability_config.enable_mfu_metrics:
256 self.perf_metrics = ModelMetrics(vllm_config)
257
263 - if self.vllm_config.model_config.enable_return_routed_experts:
264 - assert self.dcp_world_size == 1 and self.pcp_world_size == 1, (
265 - "enable_return_routed_experts does not support context parallelism "
266 - "(dcp_world_size > 1 or pcp_world_size > 1)"
267 - )
268 -
269 - self.routed_experts_reader = RoutedExpertsReader.create()
270 -
271 - assert len(kv_cache_config.kv_cache_groups) > 0, (
272 - "enable_return_routed_experts requires at least one kv cache group"
273 - )
274 - # Find the attention group for routed experts indexing.
275 - self.routed_experts_attn_gid = 0
276 - for gid, group in enumerate(kv_cache_config.kv_cache_groups):
277 - if isinstance(group.kv_cache_spec, AttentionSpec):
278 - self.routed_experts_attn_gid = gid
279 - break
280 - min_block_size = min(
281 - [
282 - group.kv_cache_spec.block_size
283 - for group in kv_cache_config.kv_cache_groups
284 - ]
285 - )
286 - num_groups = len(kv_cache_config.kv_cache_groups)
287 - self.max_num_kv_tokens = (
288 - kv_cache_config.num_blocks // num_groups
289 - ) * min_block_size
290 - dcp_size = self.vllm_config.parallel_config.decode_context_parallel_size
291 - pcp_size = self.vllm_config.parallel_config.prefill_context_parallel_size
292 - if pcp_size * dcp_size > 1:
293 - self.max_num_kv_tokens *= pcp_size * dcp_size
294 -
295 - self.routed_experts_reader.attach_buffer(
296 - max_num_kv_tokens=self.max_num_kv_tokens,
297 - vllm_config=self.vllm_config,
298 - )
299 -
258 self._pause_state: PauseState = PauseState.UNPAUSED
259
260 def _mamba_block_aligned_split(
@@ -1415,11 +1373,15 @@ class Scheduler(SchedulerInterface):
1373 request.resumable = False
1374 stopped = True
1375
1376 + # Get routing data from ModelRunnerOutput (via worker D2H pipeline)
1377 routed_experts = None
1378 + if (
1379 + model_runner_output.routed_experts_dict is not None
1380 + and req_id in model_runner_output.routed_experts_dict
1381 + ):
1382 + routed_experts = model_runner_output.routed_experts_dict[req_id]
1383 finish_reason = None
1384 if stopped:
1421 - routed_experts = self._get_routed_experts(request)
1422 -
1385 # Capture finish_reason BEFORE _handle_stopped_request, which may
1386 # reset the status to WAITING for streaming requests that continue.
1387 finish_reason = request.get_finished_reason()
@@ -1594,31 +1556,6 @@ class Scheduler(SchedulerInterface):
1556 self._enqueue_waiting_request(request)
1557 return False
1558
1597 - def _get_routed_experts(self, request: Request) -> np.ndarray | None:
1598 - if not self.vllm_config.model_config.enable_return_routed_experts:
1599 - return None
1600 -
1601 - kv_blocks = self.kv_cache_manager.get_blocks(request.request_id)
1602 - block_ids = kv_blocks.get_block_ids()[self.routed_experts_attn_gid]
1603 - num_tokens = request.num_tokens - 1
1604 -
1605 - # compute slot mapping using attention group's block_size
1606 - block_ids_array = np.array(block_ids, dtype=np.int32)
1607 - num_blocks = len(block_ids)
1608 - attn_group = self.kv_cache_config.kv_cache_groups[self.routed_experts_attn_gid]
1609 - block_size = attn_group.kv_cache_spec.block_size
1610 -
1611 - # generate block offsets
1612 - block_offsets = np.arange(0, block_size)
1613 -
1614 - # compute slot mapping: slot = block_id * block_size + offset
1615 - slot_mapping = (
1616 - block_offsets.reshape((1, block_size))
1617 - + block_ids_array.reshape((num_blocks, 1)) * block_size
1618 - ).flatten()[:num_tokens]
1619 -
1620 - return self.routed_experts_reader.get_routed_experts(indices=slot_mapping)
1621 -
1559 def _update_request_with_output(
1560 self, request: Request, new_token_ids: list[int]
1561 ) -> tuple[list[int], bool]:
vllm/v1/engine/output_processor.py
+27 -2
@@ -11,6 +11,9 @@ import numpy as np
11 import torch
12
13 from vllm.lora.request import LoRARequest
14 +from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
15 + split_routed_experts,
16 +)
17 from vllm.outputs import (
18 STREAM_FINISHED,
19 CompletionOutput,
@@ -314,8 +317,24 @@ class RequestState:
317 finished,
318 )
319
320 + # Split routing data into prompt and generation portions.
321 + # Prompt routing lives on RequestOutput (shared across n>1
322 + # completions); generation routing lives on each CompletionOutput.
323 + prompt_routed_experts = None
324 + gen_routed_experts = None
325 + if routed_experts is not None:
326 + prompt_len = len(self.prompt_token_ids) if self.prompt_token_ids else 0
327 + num_gen = (
328 + self.detokenizer.num_output_tokens()
329 + if self.detokenizer is not None
330 + else None
331 + )
332 + prompt_routed_experts, gen_routed_experts = split_routed_experts(
333 + routed_experts, prompt_len, num_gen
334 + )
335 +
336 output = self._new_completion_output(
318 - new_token_ids, finish_reason, stop_reason, routed_experts
337 + new_token_ids, finish_reason, stop_reason, gen_routed_experts
338 )
339
340 if self.parent_req is None:
@@ -327,7 +346,11 @@ class RequestState:
346 external_req_id = self.parent_req.external_req_id
347
348 return self._new_request_output(
330 - external_req_id, outputs, finished, kv_transfer_params
349 + external_req_id,
350 + outputs,
351 + finished,
352 + kv_transfer_params,
353 + prompt_routed_experts,
354 )
355
356 def _new_request_output(
@@ -336,6 +359,7 @@ class RequestState:
359 outputs: list[CompletionOutput] | list[PoolingOutput],
360 finished: bool,
361 kv_transfer_params: dict[str, Any] | None = None,
362 + prompt_routed_experts: np.ndarray | None = None,
363 ) -> RequestOutput | PoolingRequestOutput:
364 # If prompt embeds were used, put placeholder prompt token ids
365 prompt_token_ids = self.prompt_token_ids
@@ -371,6 +395,7 @@ class RequestState:
395 kv_transfer_params=kv_transfer_params,
396 num_cached_tokens=self.num_cached_tokens,
397 metrics=self.stats,
398 + prompt_routed_experts=prompt_routed_experts,
399 )
400
401 def _new_completion_output(
vllm/v1/outputs.py
+3
@@ -198,6 +198,9 @@ class ModelRunnerOutput:
198 # req_id -> num_nans_in_logits
199 num_nans_in_logits: dict[str, int] | None = None
200
201 + # req_id -> routed experts ndarray of shape (seq_len, num_moe_layers, top_k)
202 + routed_experts_dict: dict[str, np.ndarray] | None = None
203 +
204 # information related to cudagraph execution
205 cudagraph_stats: CUDAGraphStat | None = None
206
vllm/v1/worker/gpu_model_runner.py
+64 -45
@@ -54,7 +54,11 @@ from vllm.lora.layers import LoRAMapping, LoRAMappingType
54 from vllm.model_executor.layers.attention import Attention, MLAAttention
55 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
56 from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
57 - RoutedExpertsCapturer,
57 + extract_routed_experts_for_current_batch,
58 + free_routing_buffers,
59 + get_global_experts_capturer,
60 + init_routed_experts_capturer_with_shared_cache,
61 + issue_routing_d2h_copy,
62 )
63 from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
64 initialize_mamba_ssu_backend,
@@ -1092,6 +1096,12 @@ class GPUModelRunner(
1096 for req_id in scheduler_output.finished_req_ids:
1097 self.input_batch.remove_request(req_id)
1098
1099 + if self.routed_experts_initialized:
1100 + free_routing_buffers(
1101 + scheduler_output.finished_req_ids,
1102 + scheduler_output.preempted_req_ids,
1103 + )
1104 +
1105 # Zero GPU memory for freshly allocated cache blocks to prevent
1106 # stale NaN/data from corrupting attention or SSM computation.
1107 if scheduler_output.new_block_ids_to_zero:
@@ -2175,10 +2185,6 @@ class GPUModelRunner(
2185 block_table_gid_0 = _get_block_table(0)
2186 slot_mapping_gid_0 = slot_mappings[0]
2187
2178 - if self.routed_experts_initialized:
2179 - attn_gid = self.routed_experts_attn_gid
2180 - slot_mapping_attn = slot_mappings[attn_gid]
2181 - self.slot_mapping = slot_mapping_attn[:num_tokens].cpu().numpy()
2188 num_computed_tokens_cpu = self.input_batch.num_computed_tokens_cpu_tensor[
2189 :num_reqs_padded
2190 ]
@@ -3859,11 +3865,9 @@ class GPUModelRunner(
3865 )
3866
3867 if self.routed_experts_initialized:
3862 - capturer = RoutedExpertsCapturer.get_instance()
3868 + capturer = get_global_experts_capturer()
3869 if capturer is not None:
3864 - capturer.clear_buffer() # noqa
3865 - else:
3866 - logger.error("RoutedExpertsCapturer not initialized.")
3870 + capturer.finalize_pending_copy()
3871
3872 # If ngram_gpu is used, we need to copy the scheduler_output to avoid
3873 # the modification has influence on the scheduler_output in engine core process.
@@ -4373,6 +4377,14 @@ class GPUModelRunner(
4377 scheduler_output.total_num_scheduled_tokens,
4378 )
4379
4380 + if self.routed_experts_initialized:
4381 + issue_routing_d2h_copy(
4382 + input_batch_req_ids=self.input_batch.req_ids,
4383 + num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
4384 + positions=self.positions,
4385 + positions_cpu=self._positions_cpu,
4386 + )
4387 +
4388 if propose_drafts_after_bookkeeping:
4389 # ngram and other speculative decoding methods use the sampled
4390 # tokens on the CPU, so they are run after bookkeeping.
@@ -4392,12 +4404,15 @@ class GPUModelRunner(
4404 self.kv_connector_output = None
4405
4406 with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"):
4407 + routed_experts_dict = None
4408 if self.routed_experts_initialized:
4396 - capturer = RoutedExpertsCapturer.get_instance()
4397 - if capturer is not None:
4398 - capturer.save_captured_experts(indices=self.slot_mapping) # noqa
4399 - else:
4400 - logger.error("RoutedExpertsCapturer not initialized.")
4409 + routed_experts_dict = extract_routed_experts_for_current_batch(
4410 + req_ids=req_ids_output_copy,
4411 + requests=self.requests,
4412 + req_id_to_index=self.input_batch.req_id_to_index,
4413 + num_tokens_no_spec=self.input_batch.num_tokens_no_spec,
4414 + max_model_len=self.max_model_len,
4415 + )
4416
4417 output = ModelRunnerOutput(
4418 req_ids=req_ids_output_copy,
@@ -4411,6 +4426,7 @@ class GPUModelRunner(
4426 else None,
4427 num_nans_in_logits=num_nans_in_logits,
4428 cudagraph_stats=cudagraph_stats,
4429 + routed_experts_dict=routed_experts_dict,
4430 )
4431
4432 if not self.use_async_scheduling:
@@ -6138,6 +6154,7 @@ class GPUModelRunner(
6154 "Skipping CUDA graph capture. To turn on CUDA graph capture, "
6155 "ensure `cudagraph_mode` was not manually set to `NONE`"
6156 )
6157 + self.init_routed_experts_capturer()
6158 return 0
6159
6160 # Initialize encoder CUDA graph manager if enabled.
@@ -6171,6 +6188,13 @@ class GPUModelRunner(
6188
6189 start_time = time.perf_counter()
6190
6191 + # Initialize the routed experts capturer once before any CUDA graph
6192 + # capture. Must happen before graphs are captured so the buffer
6193 + # address is baked into the graph. Do NOT call this inside
6194 + # _capture_cudagraphs() -- creating the capturer twice replaces
6195 + # the device buffer, causing graphs to write to a dead buffer.
6196 + self.init_routed_experts_capturer()
6197 +
6198 # Trigger CUDA graph capture for specific shapes.
6199 # Capture the large shapes first so that the smaller shapes
6200 # can reuse the memory pool allocated for the large shapes.
@@ -6965,45 +6989,40 @@ class GPUModelRunner(
6989 "Initializing routed experts capturer, enable_return_routed_experts: %s",
6990 self.model_config.enable_return_routed_experts,
6991 )
6968 - routed_experts_capturer = RoutedExpertsCapturer.create()
6969 - self.routed_experts_attn_gid = self._get_attention_kv_cache_gid()
6970 - min_block_size = min(
6971 - [
6972 - group.kv_cache_spec.block_size
6973 - for group in self.kv_cache_config.kv_cache_groups
6974 - ]
6975 - )
6976 - num_groups = len(self.kv_cache_config.kv_cache_groups)
6977 - self.max_num_kv_tokens = (
6978 - self.kv_cache_config.num_blocks // num_groups
6979 - ) * min_block_size
6980 - dcp_size = self.vllm_config.parallel_config.decode_context_parallel_size
6981 - pcp_size = self.vllm_config.parallel_config.prefill_context_parallel_size
6982 - if pcp_size * dcp_size > 1:
6983 - self.max_num_kv_tokens *= pcp_size * dcp_size
6992 + from vllm.distributed import get_tp_group
6993 +
6994 + if hasattr(self.model_config.hf_text_config, "n_shared_experts"):
6995 + num_fused_shared_experts = 1
6996 + else:
6997 + num_fused_shared_experts = 0
6998
6985 - routed_experts_capturer.init_buffer(
6999 + tp_group = get_tp_group()
7000 + init_routed_experts_capturer_with_shared_cache(
7001 + enable=self.model_config.enable_return_routed_experts,
7002 + model_config=self.model_config,
7003 + num_fused_shared_experts=num_fused_shared_experts,
7004 max_num_batched_tokens=self.scheduler_config.max_num_batched_tokens,
6987 - max_num_kv_tokens=self.max_num_kv_tokens,
6988 - vllm_config=self.vllm_config,
7005 + max_model_len=self.max_model_len,
7006 + device=self.device,
7007 + rank=tp_group.rank_in_group,
7008 + world_size=tp_group.world_size,
7009 )
6990 - self._bind_routed_experts_capturer(routed_experts_capturer)
7010 + self._bind_routed_experts_capturer()
7011 self.routed_experts_initialized = True
7012
6993 - def _bind_routed_experts_capturer(self, capturer: RoutedExpertsCapturer) -> None:
6994 - from vllm.model_executor.layers.fused_moe.layer import FusedMoE
6995 - from vllm.model_executor.layers.fused_moe.router.base_router import (
6996 - BaseRouter,
7013 + # Pinned CPU buffer for async positions D2H (avoids sync .cpu() call)
7014 + self._positions_cpu = torch.empty(
7015 + self.scheduler_config.max_num_batched_tokens,
7016 + dtype=torch.long,
7017 + pin_memory=True,
7018 )
7019
6999 - for module in self.compilation_config.static_forward_context.values():
7000 - if isinstance(module, FusedMoE) and isinstance(module.router, BaseRouter):
7001 - layer_id = module.layer_id
7002 -
7003 - def _capture_fn(topk_ids, _layer_id=layer_id, _capturer=capturer):
7004 - _capturer.capture(_layer_id, topk_ids)
7020 + def _bind_routed_experts_capturer(self) -> None:
7021 + from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
7022 + bind_routing_capture_to_model,
7023 + )
7024
7006 - module.router.set_capture_fn(_capture_fn)
7025 + bind_routing_capture_to_model(self.model)
7026
7027 def may_add_encoder_only_layers_to_kv_cache_config(self) -> None:
7028 """