@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 8c79ad6580

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 3, 16:30 UTC

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

Signed-off-by: aoshen02 <aoshen@inferact.ai>

aoshen02 committed May 14, 2026 at 14:49 UTC 8c79ad65804d042dfaf7ce0e9fd28f503f9266a9

15 files changed +614 -1377

docs/training/routed_experts_replay.md deleted

-289

@@ -1,289 +0,0 @@

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 -

287 -|----------------------|---------------------------|--------------------------|-----------------------------|

tests/model_executor/test_routed_experts_capture.py

+199 -117

@@ -1,162 +1,244 @@

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

7 import pytest

8 import torch

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 +

17 pytestmark = pytest.mark.cpu_test

19 +_REC_MODULE = "vllm.model_executor.layers.fused_moe.routed_experts_capturer"

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

14 - class _DummyMoEConfig:

15 - is_sequence_parallel = False

16 - dp_size = 1

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

18 - class _DummyQuantMethod:

19 - supports_internal_mk = True

21 - class DummyFusedMoE:

22 - _routing_replay_out: torch.Tensor

39 +class DummyRouter(BaseRouter):

40 + @property

41 + def routing_method_type(self) -> RoutingMethodType:

42 + return RoutingMethodType.FUSED_TOPK

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()

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

29 - monkeypatch.setattr(fused_moe_layer, "FusedMoE", DummyFusedMoE)

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

31 - num_layers, num_tokens, top_k = 4, 8, 2

32 - buffer = torch.zeros((num_layers, num_tokens, top_k), dtype=torch.int16)

34 - class DummyDeviceCache:

35 - def __init__(self, buf):

36 - self.buffer = buf

56 +def _make_router(eplb_state: EplbLayerState | None = None) -> DummyRouter:

57 + return DummyRouter(

58 + top_k=2,

59 + global_num_experts=16,

60 + eplb_state=eplb_state,

61 + indices_type_getter=None,

62 + )

38 - class DummyCapturer:

39 - def get_device_cache(self):

40 - return DummyDeviceCache(buffer)

42 - monkeypatch.setattr(rec_mod, "get_global_experts_capturer", lambda: DummyCapturer())

65 +def test_base_router_capture_pre_eplb_mapping():

66 + router = _make_router()

67 + captured = []

44 - m0 = DummyFusedMoE(moe_layer_id=0)

45 - m2 = DummyFusedMoE(moe_layer_id=2)

69 + def capture_fn(ids):

70 + captured.append(ids.clone())

47 - class DummyModel:

48 - def modules(self):

49 - return iter([m0, m2])

72 + router.set_capture_fn(capture_fn)

73 + topk_weights, topk_ids = router.select_experts(

74 + hidden_states=torch.empty(1),

75 + router_logits=torch.empty(1),

76 + )

51 - rec_mod.bind_routing_capture_to_model(DummyModel())

78 + assert topk_weights.shape == topk_ids.shape

79 + assert len(captured) == 1

80 + assert torch.equal(captured[0], torch.tensor([[1, 2], [3, 4]]))

81 + assert torch.equal(topk_ids, torch.tensor([[11, 12], [13, 14]]))

53 - assert torch.equal(m0._routing_replay_out, buffer[0])

54 - assert torch.equal(m2._routing_replay_out, buffer[2])

84 +def test_base_router_capture_with_eplb_enabled():

85 + eplb_state = EplbLayerState()

86 + eplb_state.expert_load_view = torch.zeros(32, dtype=torch.int64)

87 + eplb_state.logical_to_physical_map = torch.arange(32).view(32, 1)

88 + eplb_state.logical_replica_count = torch.ones(32, dtype=torch.int64)

89 + eplb_state.should_record_tensor = torch.ones((), dtype=torch.bool)

90 + router = _make_router(eplb_state=eplb_state)

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

92 + captured = []

60 - class DummyCapturer:

61 - def get_device_cache(self):

62 - return None

94 + def capture_fn(ids):

95 + captured.append(ids.clone())

64 - monkeypatch.setattr(rec_mod, "get_global_experts_capturer", lambda: DummyCapturer())

97 + router.set_capture_fn(capture_fn)

98 + _, topk_ids = router.select_experts(

99 + hidden_states=torch.empty(1),

100 + router_logits=torch.empty(1),

101 + )

102

66 - class DummyModel:

67 - def modules(self):

68 - return iter([])

103 + assert len(captured) == 1

104 + # Capture should see logical ids pre-EPLB mapping.

105 + assert torch.equal(captured[0], torch.tensor([[1, 2], [3, 4]]))

106 + # Our DummyRouter mapping adds +10.

107 + assert torch.equal(topk_ids, torch.tensor([[11, 12], [13, 14]]))

108

70 - rec_mod.bind_routing_capture_to_model(DummyModel())

109

110 +def test_gpu_model_runner_binds_router_capture(monkeypatch):

111 + from vllm.v1.worker import gpu_model_runner as gmr

112

73 -# =========================================================================

74 -# Tests for device-cache routing replay architecture

75 -# =========================================================================

113 + class DummyFusedMoE:

114 + def __init__(self):

115 + self.layer_id = 7

116 + self.router = _make_router()

117

118 + class DummyCapturer:

119 + def __init__(self):

120 + self.calls = []

121

78 -class TestRoutedExpertsDeviceCache:

79 - """Tests for _RoutedExpertsDeviceCache (GPU buffer for routing data)."""

122 + def capture(self, layer_id, topk_ids):

123 + self.calls.append((layer_id, topk_ids))

124

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 - )

125 + dummy_module = DummyFusedMoE()

126

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

127 + # Patch the runtime import inside _bind_routed_experts_capturer.

128 + import vllm.model_executor.layers.fused_moe.layer as fused_moe_layer

129

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 - )

130 + monkeypatch.setattr(fused_moe_layer, "FusedMoE", DummyFusedMoE)

131

102 - cache = _RoutedExpertsDeviceCache(

103 - num_hidden_layers=40,

104 - max_num_batched_tokens=8192,

105 - num_experts_per_tok=8,

106 - device="cpu",

132 + dummy_self = types.SimpleNamespace(

133 + compilation_config=types.SimpleNamespace(

134 + static_forward_context={"dummy": dummy_module}

135 )

108 - layer_view = cache.buffer[0]

109 - assert layer_view.is_contiguous()

110 - assert layer_view.shape == (8192, 8)

136 + )

137

138 + capturer = DummyCapturer()

139 + gmr.GPUModelRunner._bind_routed_experts_capturer(dummy_self, capturer)

140

113 -class TestRoutedExpertsHostCache:

114 - """Tests for _RoutedExpertsHostCache (per-request numpy buffer)."""

141 + assert dummy_module.router.capture_fn is not None

142 + dummy_module.router.capture_fn(torch.tensor([[5, 6]]))

143

116 - def test_sentinel_initialization(self):

117 - """Host cache initializes with zeros by default."""

118 - import numpy as np

144 + assert len(capturer.calls) == 1

145 + layer_id, topk_ids = capturer.calls[0]

146 + assert layer_id == 7

147 + assert torch.equal(topk_ids, torch.tensor([[5, 6]]))

148

120 - from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (

121 - _RoutedExpertsHostCache,

122 - )

149

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 - )

150 +def test_gpu_model_runner_binding_stage(monkeypatch):

151 + from vllm.v1.worker import gpu_model_runner as gmr

152

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 - )

153 + class DummyFusedMoE:

154 + def __init__(self):

155 + self.layer_id = 11

156 + self.router = _make_router()

157 +

158 + class DummyCapturer:

159 + def __init__(self):

160 + self.calls = []

161 +

162 + def capture(self, layer_id, topk_ids):

163 + self.calls.append((layer_id, topk_ids))

164 +

165 + dummy_module = DummyFusedMoE()

166 +

167 + import vllm.model_executor.layers.fused_moe.layer as fused_moe_layer

168 +

169 + monkeypatch.setattr(fused_moe_layer, "FusedMoE", DummyFusedMoE)

170

155 - cache = _RoutedExpertsHostCache(

156 - num_hidden_layers=40,

157 - num_experts_per_tok=8,

158 - max_model_len=1024,

171 + dummy_self = types.SimpleNamespace(

172 + compilation_config=types.SimpleNamespace(

173 + static_forward_context={"dummy": dummy_module}

174 )

160 - cache.get_or_grow_buffer("req1", max_pos=50)

161 - cache.free_request("req1")

162 - assert cache.get_buffer("req1") is None

175 + )

176 +

177 + # Before binding, no capture hook.

178 + assert dummy_module.router.capture_fn is None

179 +

180 + capturer = DummyCapturer()

181 + gmr.GPUModelRunner._bind_routed_experts_capturer(dummy_self, capturer)

182 +

183 + # After binding, hook should exist and be callable.

184 + assert callable(dummy_module.router.capture_fn)

185 + dummy_module.router.capture_fn(torch.tensor([[9, 10]]))

186 + assert len(capturer.calls) == 1

187 +

188 +

189 +def test_routed_experts_capturer_single_dp_no_metadata():

190 + """dp_metadata is None: capture writes the full topk_ids rows."""

191 + capturer = _capturer_with_buffer(dp_rank=0)

192 + topk = torch.tensor([[1, 2], [3, 4], [5, 6]], dtype=torch.int32)

193 + ctx = SimpleNamespace(dp_metadata=None)

194 + with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx):

195 + capturer.capture(layer_id=0, topk_ids=topk)

196 + assert torch.equal(capturer._device_buffer[:3, 0, :], topk)

197 + assert capturer._device_buffer[3, 0, 0].item() == -1

198 +

199 +

200 +def test_routed_experts_capturer_dp_naive_concatenated_all_ranks():

201 + """n == sum(num_tokens_dp): slice this rank's segment from concatenated topk."""

202 + capturer = _capturer_with_buffer(dp_rank=1)

203 + num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32)

204 + ctx = SimpleNamespace(

205 + dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp)

206 + )

207 + # Concatenated order: rank0 rows then rank1 rows.

208 + topk = torch.tensor(

209 + [[0, 1], [2, 3], [10, 11], [12, 13], [14, 15]], dtype=torch.int32

210 + )

211 + with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx):

212 + capturer.capture(layer_id=0, topk_ids=topk)

213 + want = topk[2:5]

214 + assert torch.equal(capturer._device_buffer[:3, 0, :], want)

215 +

216 +

217 +def test_routed_experts_capturer_dp_modular_local_tokens():

218 + """n == token_num_per_dp: topk is already local to this DP rank."""

219 + capturer = _capturer_with_buffer(dp_rank=1)

220 + num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32)

221 + ctx = SimpleNamespace(

222 + dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp)

223 + )

224 + topk = torch.tensor([[10, 11], [12, 13], [14, 15]], dtype=torch.int32)

225 + with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx):

226 + capturer.capture(layer_id=0, topk_ids=topk)

227 + assert torch.equal(capturer._device_buffer[:3, 0, :], topk)

228 +

229 +

230 +def test_routed_experts_capturer_dp_unexpected_batch_raises():

231 + """Mismatch between topk batch dim and DP layout: fail fast."""

232 + capturer = _capturer_with_buffer(dp_rank=0)

233 + num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32)

234 + ctx = SimpleNamespace(

235 + dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp)

236 + )

237 + # total=5, local=2: n=1 matches neither naive (5) nor modular (2).

238 + topk = torch.tensor([[1, 2]], dtype=torch.int32)

239 + with (

240 + patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx),

241 + pytest.raises(AssertionError, match="unexpected topk_ids batch dim"),

242 + ):

243 + capturer.capture(layer_id=0, topk_ids=topk)

244 + assert capturer._device_buffer[0, 0, 0].item() == -1

vllm/config/vllm.py

-51

@@ -1212,12 +1212,6 @@ class VllmConfig:

1212 if envs.VLLM_USE_V2_MODEL_RUNNER:

1213 self._validate_v2_model_runner()

1214

1215 - if (

1216 - self.model_config is not None

1217 - and self.model_config.enable_return_routed_experts

1218 - ):

1219 - self._validate_return_routed_experts()

1220 -

1215 # Re-compute compile ranges after platform-specific config updates

1216 # (e.g., XPU may lower max_num_batched_tokens when MLA is enabled)

1217 self._set_compile_ranges()

@@ -1933,51 +1927,6 @@ class VllmConfig:

1927 + ", ".join(unsupported)

1928 )

1929

1936 - def _validate_return_routed_experts(self) -> None:

1937 - """Reject parallelism configurations not yet validated with

1938 - --enable-return-routed-experts.

1939 -

1940 - Validated scope (PR #39917): TP, EP, DP, single-node and multi-node,

1941 - prefix caching, and speculative decoding (MTP validated end-to-end;

1942 - Eagle/Eagle3/Ngram/Medusa supported by construction since the

1943 - routing buffer is bound only to the target model and verified-token

1944 - routing lands at the correct positions during the main forward).

1945 -

1946 - Out-of-scope (block until validated): PP > 1, prefill context

1947 - parallelism (PCP) > 1, decode context parallelism (DCP) > 1,

1948 - async scheduling.

1949 - """

1950 - unsupported: list[str] = []

1951 -

1952 - if self.parallel_config.pipeline_parallel_size > 1:

1953 - unsupported.append(

1954 - "pipeline parallelism "

1955 - f"(pipeline_parallel_size="

1956 - f"{self.parallel_config.pipeline_parallel_size})"

1957 - )

1958 - if self.parallel_config.prefill_context_parallel_size > 1:

1959 - unsupported.append(

1960 - "prefill context parallelism "

1961 - f"(prefill_context_parallel_size="

1962 - f"{self.parallel_config.prefill_context_parallel_size})"

1963 - )

1964 - if self.parallel_config.decode_context_parallel_size > 1:

1965 - unsupported.append(

1966 - "decode context parallelism "

1967 - f"(decode_context_parallel_size="

1968 - f"{self.parallel_config.decode_context_parallel_size})"

1969 - )

1970 - if self.scheduler_config.async_scheduling:

1971 - unsupported.append("async scheduling")

1972 -

1973 - if unsupported:

1974 - raise ValueError(

1975 - "--enable-return-routed-experts is not yet validated with: "

1976 - + ", ".join(unsupported)

1977 - + ". Disable these features or omit "

1978 - "--enable-return-routed-experts."

1979 - )

1980 -

1930 def validate_block_size(self) -> None:

1931 """Validate block_size against DCP and mamba constraints.

1932

vllm/entrypoints/openai/chat_completion/protocol.py

-4

@@ -92,16 +92,12 @@ 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]

97 class ChatCompletionResponse(OpenAIBaseModel):

98 id: str = Field(default_factory=lambda: f"chatcmpl-{random_uuid()}")

99 object: Literal["chat.completion"] = "chat.completion"

100 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 - )

101 model: str

102 choices: list[ChatCompletionResponseChoice]

103 service_tier: Literal["auto", "default", "flex", "scale", "priority"] | None = None

vllm/entrypoints/openai/chat_completion/serving.py

-15

@@ -1101,11 +1101,6 @@ class OpenAIServingChat(OpenAIServing):

1101 token_ids=(

1102 as_list(output.token_ids) if request.return_token_ids else None

1103 ),

1104 - routed_experts=(

1105 - output.routed_experts.tolist()

1106 - if output.routed_experts is not None

1107 - else None

1108 - ),

1104 )

1105 choices.append(choice_data)

1106 continue

@@ -1327,11 +1322,6 @@ class OpenAIServingChat(OpenAIServing):

1322 token_ids=(

1323 as_list(output.token_ids) if request.return_token_ids else None

1324 ),

1330 - routed_experts=(

1331 - output.routed_experts.tolist()

1332 - if output.routed_experts is not None

1333 - else None

1334 - ),

1325 )

1326 choice_data = maybe_filter_parallel_tool_calls(choice_data, request)

1327

@@ -1371,10 +1361,6 @@ class OpenAIServingChat(OpenAIServing):

1361

1362 request_metadata.final_usage_info = usage

1363

1374 - prompt_routed_experts = None

1375 - if final_res.prompt_routed_experts is not None:

1376 - prompt_routed_experts = final_res.prompt_routed_experts.tolist()

1377 -

1364 # ``final_res.prompt`` is the rendered chat-templated prompt text

1365 prompt_text = final_res.prompt if request.return_prompt_text else None

1366

@@ -1391,7 +1377,6 @@ class OpenAIServingChat(OpenAIServing):

1377 ),

1378 prompt_text=prompt_text,

1379 kv_transfer_params=final_res.kv_transfer_params,

1394 - prompt_routed_experts=prompt_routed_experts,

1380 )

1381

1382 # Log complete response if output logging is enabled

vllm/entrypoints/openai/completion/protocol.py

-4

@@ -477,16 +477,12 @@ class CompletionResponseChoice(OpenAIBaseModel):

477 token_ids: list[int] | None = None # For response

478 prompt_logprobs: list[dict[int, Logprob] | None] | None = None

479 prompt_token_ids: list[int] | None = None # For prompt

480 - routed_experts: list[list[list[int]]] | None = None # [gen_len, num_layers, top_k]

480

481

482 class CompletionResponse(OpenAIBaseModel):

483 id: str = Field(default_factory=lambda: f"cmpl-{random_uuid()}")

484 object: Literal["text_completion"] = "text_completion"

485 created: int = Field(default_factory=lambda: int(time.time()))

487 - prompt_routed_experts: list[list[list[int]]] | None = (

488 - None # [prompt_len, num_layers, top_k]

489 - )

486 model: str

487 choices: list[CompletionResponseChoice]

488 service_tier: Literal["auto", "default", "flex", "scale", "priority"] | None = None

vllm/entrypoints/openai/completion/serving.py

-11

@@ -551,11 +551,6 @@ class OpenAIServingCompletion(OpenAIServing):

551 token_ids=(

552 as_list(output.token_ids) if request.return_token_ids else None

553 ),

554 - routed_experts=(

555 - output.routed_experts.tolist()

556 - if output.routed_experts is not None

557 - else None

558 - ),

554 )

555 choices.append(choice_data)

556

@@ -579,13 +574,8 @@ class OpenAIServingCompletion(OpenAIServing):

574 )

575

576 request_metadata.final_usage_info = usage

582 - prompt_routed_experts = None

577 if final_res_batch:

578 kv_transfer_params = final_res_batch[0].kv_transfer_params

585 - pre = final_res_batch[0].prompt_routed_experts

586 - if pre is not None:

587 - prompt_routed_experts = pre.tolist()

588 -

579 return CompletionResponse(

580 id=request_id,

581 created=created_time,

@@ -594,7 +584,6 @@ class OpenAIServingCompletion(OpenAIServing):

584 usage=usage,

585 system_fingerprint=self.system_fingerprint,

586 kv_transfer_params=kv_transfer_params,

597 - prompt_routed_experts=prompt_routed_experts,

587 )

588

589 def _create_completion_logprobs(

vllm/model_executor/layers/fused_moe/layer.py

-7

@@ -100,9 +100,6 @@ class FusedMoE(PluggableLayer):

100 not supported by the router (or the experts).

101 """

102

103 - # Auto-incrementing layer ID for routing replay buffer binding.

104 - _next_moe_layer_id: int = 0

105 -

103 # --8<-- [end:fused_moe]

104

105 def __init__(

@@ -148,10 +145,6 @@ class FusedMoE(PluggableLayer):

145 ):

146 super().__init__()

147

151 - # Assign unique layer ID for routing replay buffer binding.

152 - self.moe_layer_id = FusedMoE._next_moe_layer_id

153 - FusedMoE._next_moe_layer_id += 1

154 -

148 if params_dtype is None:

149 params_dtype = torch.get_default_dtype()

150 self.params_dtype = params_dtype

vllm/model_executor/layers/fused_moe/routed_experts_capturer.py

+298 -760

@@ -1,833 +1,371 @@

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 +

6 from __future__ import annotations

5 -import contextlib

8 +import fcntl

9 import logging

7 -from abc import ABC, abstractmethod

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

17 import numpy as np

18 import torch

11 -import torch.distributed

13 -from vllm.config.model import ModelConfig

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

25 logger = logging.getLogger(__name__)

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 -# ---------------------------------------------------------------------------

28 +def _get_num_experts_per_tok(hf_config) -> int:

29 + """Resolve the per-token expert count from the HF config.

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

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 - )

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"

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

50 +# Global singleton instances

51 +_global_experts_capturer: RoutedExpertsCapturer | None = None

52 +_global_experts_reader: RoutedExpertsReader | None = None

50 -_MB = 1024 * 1024

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)

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 - """

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

59 - DTYPE = torch.int16

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)

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 - )

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)

94 + return shm

96 -class _RoutedExpertsHostCache:

97 - """Host (CPU) cache using numpy arrays for per-request routing data.

99 - Numpy arrays avoid torch dispatcher overhead for scatter operations.

100 - Lazy per-request allocation avoids a massive up-front buffer.

97 +class RoutedExpertsCapturer:

98 """

99 + Capturer for routed experts with device and optional shared memory buffer.

100

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 -

119 - self._finalize_allocation_log()

101 + This class captures expert routing decisions during model forward passes

102 + and optionally stores them in shared memory for cross-process access.

103 + """

104

121 - def get_buffer_size_bytes(self) -> int:

122 - return self._total_allocated_bytes

105 + _instance: RoutedExpertsCapturer | None = None

106

124 - def get_or_grow_buffer(self, req_id: str, max_pos: int) -> np.ndarray:

125 - required_len = max_pos + 1

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

112

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 - )

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.")

119

120 + _global_experts_capturer = cls()

121 + return _global_experts_capturer

122

175 -class RoutedExpertsCapturer(ABC):

123 @staticmethod

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

124 + def get_instance() -> RoutedExpertsCapturer | None:

125 + """Get the global singleton instance."""

126 + return _global_experts_capturer

127

208 - def sync_fwd_experts_buffer_DtoH(

128 + def init_buffer(

129 self,

210 - positions: torch.Tensor,

211 - num_scheduled_tokens: dict[str, int],

212 - ):

213 - raise NotImplementedError

214 -

215 - def finalize_pending_copy(self):

216 - raise NotImplementedError

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.

136

218 - def get_host_cache(self):

219 - raise NotImplementedError

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 + """

142

221 - def get_device_cache(self):

222 - raise NotImplementedError

143 + if self._device_buffer is not None:

144 + raise RuntimeError("Device buffer has already been initialized")

145

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)

149

225 -def _count_moe_layers(hf_config) -> int:

226 - """Count the number of MoE layers in a model.

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

157

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 - """

158 + if get_tensor_model_parallel_rank() != 0:

159 + return

160

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 - )

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}"

167

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,

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)

173

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 - )

174 + logger.debug(

175 + "Created shared memory buffer '%s' with shape %s",

176 + shm_name,

177 + shape,

178 + )

179

340 - def capture(self, layer_id: int, topk_ids: torch.Tensor):

341 - self.device_cache.capture_fwd_routed_experts(layer_id, topk_ids)

180 + def capture(self, layer_id: int, topk_ids: torch.Tensor) -> None:

181 + """

182 + Capture expert routing decisions for a specific layer.

183

343 - # ------------------------------------------------------------------

344 - # sync_fwd_experts_buffer_DtoH -- called AFTER the forward pass

345 - # ------------------------------------------------------------------

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 + )

218

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:

219 + if layer_id >= self._device_buffer.shape[1]:

220 return

221

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

222 + self._device_buffer[:token_num_per_dp, layer_id, :] = topk_ids[

223 + start_loc:end_loc, :

224 + ]

225

362 - total_tokens = sum(num_scheduled_tokens.values())

363 - if total_tokens == 0:

364 - return

226 + def clear_buffer(self) -> None:

227 + """Clear the device buffer."""

228 + if self._device_buffer is not None:

229 + self._device_buffer.zero_()

230

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()

231 + def save_captured_experts(self, indices: np.ndarray) -> None:

232 + """

233 + Save captured experts from device buffer to shared memory.

234

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

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.")

246

389 - # ------------------------------------------------------------------

390 - # Optimized scatter into pre-allocated host-cache buffers

391 - # ------------------------------------------------------------------

247 + num_tokens = len(indices)

248 + data = self._device_buffer[:num_tokens, :, :].cpu().numpy()

249

393 - def _scatter_to_host(self):

394 - """Scatter D2H data into per-request host cache buffers.

250 + with _file_lock(self._lock_file):

251 + self._host_buffer_view[indices, :, :] = data

252

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

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

263

484 - def capture(self, layer_id: int, topk_ids: torch.Tensor):

485 - pass

264 + def __del__(self) -> None:

265 + """Clean up shared memory on destruction."""

266 + self.cleanup()

267

487 - def get_routed_experts(self, req_id: str, seqlen=None, free_slot=True):

488 - return None

268

490 - def sync_fwd_experts_buffer_DtoH(self, positions, num_scheduled_tokens):

491 - pass

269 +class RoutedExpertsReader:

270 + """

271 + Reader for routed experts from shared memory.

272

493 - def finalize_pending_copy(self):

494 - pass

273 + This class attaches to shared memory created by RoutedExpertsCapturer

274 + and reads expert routing decisions.

275 + """

276

496 - def get_host_cache(self):

497 - return None

277 + _instance: RoutedExpertsReader | None = None

278

499 - def get_device_cache(self):

500 - pass

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

283

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.")

290

503 -# Global capturer instance (per-process)

504 -_global_expert_capturer: RoutedExpertsCapturer | None = _RoutedExpertsCapturerNoop()

505 -_shared_host_cache: _RoutedExpertsHostCache | None = None

291 + _global_experts_reader = cls()

292 + return _global_experts_reader

293

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

300

508 -def get_global_experts_capturer():

509 - return _global_expert_capturer

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.

308

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 + )

323

512 -def set_global_experts_capturer(capturer: RoutedExpertsCapturer):

513 - global _global_expert_capturer

514 - _global_expert_capturer = capturer

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}"

328

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)

336

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.

337 + self._host_buffer_view = np.ndarray(

338 + shape, dtype=np.int32, buffer=self._shm.buf

339 + )

340

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.

341 + def get_routed_experts(self, indices: np.ndarray) -> np.ndarray:

342 + """

343 + Read routed expert data from shared memory.

344

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 - )

345 + Args:

346 + indices: Array of indices to read.

347

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 - )

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()

vllm/model_executor/layers/fused_moe/runner/moe_runner.py

-8

@@ -511,10 +511,6 @@ class MoERunner(MoERunnerInterface):

511 shared_experts_input, SharedExpertsOrder.NO_OVERLAP

512 )

513

514 - # Get routing replay buffer from persistent layer attribute

515 - # (set by bind_routing_capture_to_model during capturer init)

516 - routing_replay_out = getattr(layer, "_routing_replay_out", None)

517 -

514 if self._quant_method.is_monolithic:

515 fused_out = self._quant_method.apply_monolithic(

516 layer=layer,

@@ -529,10 +525,6 @@ class MoERunner(MoERunnerInterface):

525 input_ids=input_ids,

526 )

527

532 - # Write routing data for non-monolithic path (Triton, etc.)

533 - if routing_replay_out is not None:

534 - routing_replay_out[: topk_ids.shape[0]].copy_(topk_ids.to(torch.int16))

535 -

528 # Passing shared_experts_input in case SharedExpertsOrder is

529 # MK_INTERNAL_OVERLAPPED.

530 fused_out = self._quant_method.apply(

vllm/outputs.py

-4

@@ -121,7 +121,6 @@ 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,

124 # Forward compatibility, code that uses args added in new release can

125 # still run with older versions of vLLM without breaking.

126 **kwargs: Any,

@@ -142,15 +141,12 @@ class RequestOutput:

141 self.encoder_prompt_token_ids = encoder_prompt_token_ids

142 self.num_cached_tokens = num_cached_tokens

143 self.kv_transfer_params = kv_transfer_params

145 - self.prompt_routed_experts = prompt_routed_experts

144

145 def add(self, next_output: "RequestOutput", aggregate: bool) -> None:

146 """Merge subsequent RequestOutput into this one"""

147

148 self.finished |= next_output.finished

149 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

150

151 for next_completion in next_output.outputs:

152 for i, completion in enumerate(self.outputs):

vllm/v1/core/sched/scheduler.py

+70 -7

@@ -7,6 +7,8 @@ from collections.abc import Iterable

7 from dataclasses import replace

8 from typing import Any

10 +import numpy as np

11 +

12 from vllm import envs

13 from vllm.compilation.cuda_graph import CUDAGraphStat

14 from vllm.config import VllmConfig

@@ -25,6 +27,9 @@ from vllm.distributed.kv_transfer.kv_connector.v1 import (

27 from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata

28 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats

29 from vllm.logger import init_logger

30 +from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (

31 + RoutedExpertsReader,

32 +)

33 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry

34 from vllm.multimodal.encoder_budget import MultiModalBudget

35 from vllm.v1.core.encoder_cache_manager import (

@@ -47,7 +52,7 @@ from vllm.v1.core.sched.request_queue import (

52 )

53 from vllm.v1.core.sched.utils import check_stop, remove_all

54 from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs

50 -from vllm.v1.kv_cache_interface import KVCacheConfig

55 +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig

56 from vllm.v1.metrics.perf import ModelMetrics, PerfStats

57 from vllm.v1.metrics.stats import PrefixCacheStats, SchedulerStats

58 from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput

@@ -253,6 +258,43 @@ class Scheduler(SchedulerInterface):

258 if self.log_stats and vllm_config.observability_config.enable_mfu_metrics:

259 self.perf_metrics = ModelMetrics(vllm_config)

260

261 + if self.vllm_config.model_config.enable_return_routed_experts:

262 + assert self.dcp_world_size == 1 and self.pcp_world_size == 1, (

263 + "enable_return_routed_experts does not support context parallelism "

264 + "(dcp_world_size > 1 or pcp_world_size > 1)"

265 + )

266 +

267 + self.routed_experts_reader = RoutedExpertsReader.create()

268 +

269 + assert len(kv_cache_config.kv_cache_groups) > 0, (

270 + "enable_return_routed_experts requires at least one kv cache group"

271 + )

272 + # Find the attention group for routed experts indexing.

273 + self.routed_experts_attn_gid = 0

274 + for gid, group in enumerate(kv_cache_config.kv_cache_groups):

275 + if isinstance(group.kv_cache_spec, AttentionSpec):

276 + self.routed_experts_attn_gid = gid

277 + break

278 + min_block_size = min(

279 + [

280 + group.kv_cache_spec.block_size

281 + for group in kv_cache_config.kv_cache_groups

282 + ]

283 + )

284 + num_groups = len(kv_cache_config.kv_cache_groups)

285 + self.max_num_kv_tokens = (

286 + kv_cache_config.num_blocks // num_groups

287 + ) * min_block_size

288 + dcp_size = self.vllm_config.parallel_config.decode_context_parallel_size

289 + pcp_size = self.vllm_config.parallel_config.prefill_context_parallel_size

290 + if pcp_size * dcp_size > 1:

291 + self.max_num_kv_tokens *= pcp_size * dcp_size

292 +

293 + self.routed_experts_reader.attach_buffer(

294 + max_num_kv_tokens=self.max_num_kv_tokens,

295 + vllm_config=self.vllm_config,

296 + )

297 +

298 self._pause_state: PauseState = PauseState.UNPAUSED

299

300 def _mamba_block_aligned_split(

@@ -1371,15 +1413,11 @@ class Scheduler(SchedulerInterface):

1413 request.resumable = False

1414 stopped = True

1415

1374 - # Get routing data from ModelRunnerOutput (via worker D2H pipeline)

1416 routed_experts = None

1376 - if (

1377 - model_runner_output.routed_experts_dict is not None

1378 - and req_id in model_runner_output.routed_experts_dict

1379 - ):

1380 - routed_experts = model_runner_output.routed_experts_dict[req_id]

1417 finish_reason = None

1418 if stopped:

1419 + routed_experts = self._get_routed_experts(request)

1420 +

1421 # Capture finish_reason BEFORE _handle_stopped_request, which may

1422 # reset the status to WAITING for streaming requests that continue.

1423 finish_reason = request.get_finished_reason()

@@ -1554,6 +1592,31 @@ class Scheduler(SchedulerInterface):

1592 self._enqueue_waiting_request(request)

1593 return False

1594

1595 + def _get_routed_experts(self, request: Request) -> np.ndarray | None:

1596 + if not self.vllm_config.model_config.enable_return_routed_experts:

1597 + return None

1598 +

1599 + kv_blocks = self.kv_cache_manager.get_blocks(request.request_id)

1600 + block_ids = kv_blocks.get_block_ids()[self.routed_experts_attn_gid]

1601 + num_tokens = request.num_tokens - 1

1602 +

1603 + # compute slot mapping using attention group's block_size

1604 + block_ids_array = np.array(block_ids, dtype=np.int32)

1605 + num_blocks = len(block_ids)

1606 + attn_group = self.kv_cache_config.kv_cache_groups[self.routed_experts_attn_gid]

1607 + block_size = attn_group.kv_cache_spec.block_size

1608 +

1609 + # generate block offsets

1610 + block_offsets = np.arange(0, block_size)

1611 +

1612 + # compute slot mapping: slot = block_id * block_size + offset

1613 + slot_mapping = (

1614 + block_offsets.reshape((1, block_size))

1615 + + block_ids_array.reshape((num_blocks, 1)) * block_size

1616 + ).flatten()[:num_tokens]

1617 +

1618 + return self.routed_experts_reader.get_routed_experts(indices=slot_mapping)

1619 +

1620 def _update_request_with_output(

1621 self, request: Request, new_token_ids: list[int]

1622 ) -> tuple[list[int], bool]:

vllm/v1/engine/output_processor.py

+2 -27

@@ -11,9 +11,6 @@ import numpy as np

11 import torch

13 from vllm.lora.request import LoRARequest

14 -from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (

15 - split_routed_experts,

16 -)

14 from vllm.outputs import (

15 STREAM_FINISHED,

16 CompletionOutput,

@@ -317,24 +314,8 @@ class RequestState:

314 finished,

315 )

316

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 -

317 output = self._new_completion_output(

337 - new_token_ids, finish_reason, stop_reason, gen_routed_experts

318 + new_token_ids, finish_reason, stop_reason, routed_experts

319 )

320

321 if self.parent_req is None:

@@ -346,11 +327,7 @@ class RequestState:

327 external_req_id = self.parent_req.external_req_id

328

329 return self._new_request_output(

349 - external_req_id,

350 - outputs,

351 - finished,

352 - kv_transfer_params,

353 - prompt_routed_experts,

330 + external_req_id, outputs, finished, kv_transfer_params

331 )

332

333 def _new_request_output(

@@ -359,7 +336,6 @@ class RequestState:

336 outputs: list[CompletionOutput] | list[PoolingOutput],

337 finished: bool,

338 kv_transfer_params: dict[str, Any] | None = None,

362 - prompt_routed_experts: np.ndarray | None = None,

339 ) -> RequestOutput | PoolingRequestOutput:

340 # If prompt embeds were used, put placeholder prompt token ids

341 prompt_token_ids = self.prompt_token_ids

@@ -395,7 +371,6 @@ class RequestState:

371 kv_transfer_params=kv_transfer_params,

372 num_cached_tokens=self.num_cached_tokens,

373 metrics=self.stats,

398 - prompt_routed_experts=prompt_routed_experts,

374 )

375

376 def _new_completion_output(

vllm/v1/outputs.py

-3

@@ -198,9 +198,6 @@ 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 -

201 # information related to cudagraph execution

202 cudagraph_stats: CUDAGraphStat | None = None

203

vllm/v1/worker/gpu_model_runner.py

+45 -70

@@ -54,11 +54,7 @@ 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 - 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,

57 + RoutedExpertsCapturer,

58 )

59 from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (

60 initialize_mamba_ssu_backend,

@@ -1112,12 +1108,6 @@ class GPUModelRunner(

1108 for req_id in scheduler_output.finished_req_ids:

1109 self.input_batch.remove_request(req_id)

1110

1115 - if self.routed_experts_initialized:

1116 - free_routing_buffers(

1117 - scheduler_output.finished_req_ids,

1118 - scheduler_output.preempted_req_ids,

1119 - )

1120 -

1111 # Zero GPU memory for freshly allocated cache blocks to prevent

1112 # stale NaN/data from corrupting attention or SSM computation.

1113 if scheduler_output.new_block_ids_to_zero:

@@ -2201,6 +2191,10 @@ class GPUModelRunner(

2191 block_table_gid_0 = _get_block_table(0)

2192 slot_mapping_gid_0 = slot_mappings[0]

2193

2194 + if self.routed_experts_initialized:

2195 + attn_gid = self.routed_experts_attn_gid

2196 + slot_mapping_attn = slot_mappings[attn_gid]

2197 + self.slot_mapping = slot_mapping_attn[:num_tokens].cpu().numpy()

2198 num_computed_tokens_cpu = self.input_batch.num_computed_tokens_cpu_tensor[

2199 :num_reqs_padded

2200 ]

@@ -3884,9 +3878,11 @@ class GPUModelRunner(

3878 )

3879

3880 if self.routed_experts_initialized:

3887 - capturer = get_global_experts_capturer()

3881 + capturer = RoutedExpertsCapturer.get_instance()

3882 if capturer is not None:

3889 - capturer.finalize_pending_copy()

3883 + capturer.clear_buffer() # noqa

3884 + else:

3885 + logger.error("RoutedExpertsCapturer not initialized.")

3886

3887 # If ngram_gpu is used, we need to copy the scheduler_output to avoid

3888 # the modification has influence on the scheduler_output in engine core process.

@@ -4400,14 +4396,6 @@ class GPUModelRunner(

4396 scheduler_output.total_num_scheduled_tokens,

4397 )

4398

4403 - if self.routed_experts_initialized:

4404 - issue_routing_d2h_copy(

4405 - input_batch_req_ids=self.input_batch.req_ids,

4406 - num_scheduled_tokens=scheduler_output.num_scheduled_tokens,

4407 - positions=self.positions,

4408 - positions_cpu=self._positions_cpu,

4409 - )

4410 -

4399 if propose_drafts_after_bookkeeping:

4400 # ngram and other speculative decoding methods use the sampled

4401 # tokens on the CPU, so they are run after bookkeeping.

@@ -4427,15 +4415,12 @@ class GPUModelRunner(

4415 self.kv_connector_output = None

4416

4417 with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"):

4430 - routed_experts_dict = None

4418 if self.routed_experts_initialized:

4432 - routed_experts_dict = extract_routed_experts_for_current_batch(

4433 - req_ids=req_ids_output_copy,

4434 - requests=self.requests,

4435 - req_id_to_index=self.input_batch.req_id_to_index,

4436 - num_tokens_no_spec=self.input_batch.num_tokens_no_spec,

4437 - max_model_len=self.max_model_len,

4438 - )

4419 + capturer = RoutedExpertsCapturer.get_instance()

4420 + if capturer is not None:

4421 + capturer.save_captured_experts(indices=self.slot_mapping) # noqa

4422 + else:

4423 + logger.error("RoutedExpertsCapturer not initialized.")

4424

4425 output = ModelRunnerOutput(

4426 req_ids=req_ids_output_copy,

@@ -4449,7 +4434,6 @@ class GPUModelRunner(

4434 else None,

4435 num_nans_in_logits=num_nans_in_logits,

4436 cudagraph_stats=cudagraph_stats,

4452 - routed_experts_dict=routed_experts_dict,

4437 )

4438

4439 if not self.use_async_scheduling:

@@ -6242,8 +6226,6 @@ class GPUModelRunner(

6226 "Skipping CUDA graph capture. To turn on CUDA graph capture, "

6227 "ensure `cudagraph_mode` was not manually set to `NONE`"

6228 )

6245 - if self.model_config.enable_return_routed_experts:

6246 - self.init_routed_experts_capturer()

6229 return 0

6230

6231 # Initialize encoder CUDA graph manager if enabled.

@@ -6277,14 +6259,6 @@ class GPUModelRunner(

6259

6260 start_time = time.perf_counter()

6261

6280 - # Initialize the routed experts capturer once before any CUDA graph

6281 - # capture. Must happen before graphs are captured so the buffer

6282 - # address is baked into the graph. Do NOT call this inside

6283 - # _capture_cudagraphs() -- creating the capturer twice replaces

6284 - # the device buffer, causing graphs to write to a dead buffer.

6285 - if self.model_config.enable_return_routed_experts:

6286 - self.init_routed_experts_capturer()

6287 -

6262 # Trigger CUDA graph capture for specific shapes.

6263 # Capture the large shapes first so that the smaller shapes

6264 # can reuse the memory pool allocated for the large shapes.

@@ -7079,44 +7053,45 @@ class GPUModelRunner(

7053 "Initializing routed experts capturer, enable_return_routed_experts: %s",

7054 self.model_config.enable_return_routed_experts,

7055 )

7082 - if not self.model_config.enable_return_routed_experts:

7083 - self.routed_experts_initialized = False

7084 - return

7085 -

7086 - from vllm.distributed import get_tp_group

7087 -

7088 - if hasattr(self.model_config.hf_text_config, "n_shared_experts"):

7089 - num_fused_shared_experts = 1

7090 - else:

7091 - num_fused_shared_experts = 0

7056 + routed_experts_capturer = RoutedExpertsCapturer.create()

7057 + self.routed_experts_attn_gid = self._get_attention_kv_cache_gid()

7058 + min_block_size = min(

7059 + [

7060 + group.kv_cache_spec.block_size

7061 + for group in self.kv_cache_config.kv_cache_groups

7062 + ]

7063 + )

7064 + num_groups = len(self.kv_cache_config.kv_cache_groups)

7065 + self.max_num_kv_tokens = (

7066 + self.kv_cache_config.num_blocks // num_groups

7067 + ) * min_block_size

7068 + dcp_size = self.vllm_config.parallel_config.decode_context_parallel_size

7069 + pcp_size = self.vllm_config.parallel_config.prefill_context_parallel_size

7070 + if pcp_size * dcp_size > 1:

7071 + self.max_num_kv_tokens *= pcp_size * dcp_size

7072

7093 - tp_group = get_tp_group()

7094 - init_routed_experts_capturer_with_shared_cache(

7095 - enable=self.model_config.enable_return_routed_experts,

7096 - model_config=self.model_config,

7097 - num_fused_shared_experts=num_fused_shared_experts,

7073 + routed_experts_capturer.init_buffer(

7074 max_num_batched_tokens=self.scheduler_config.max_num_batched_tokens,

7099 - max_model_len=self.max_model_len,

7100 - device=self.device,

7101 - rank=tp_group.rank_in_group,

7102 - world_size=tp_group.world_size,

7075 + max_num_kv_tokens=self.max_num_kv_tokens,

7076 + vllm_config=self.vllm_config,

7077 )

7104 - self._bind_routed_experts_capturer()

7078 + self._bind_routed_experts_capturer(routed_experts_capturer)

7079 self.routed_experts_initialized = True

7080

7107 - # Pinned CPU buffer for async positions D2H (avoids sync .cpu() call)

7108 - self._positions_cpu = torch.empty(

7109 - self.scheduler_config.max_num_batched_tokens,

7110 - dtype=torch.long,

7111 - pin_memory=True,

7081 + def _bind_routed_experts_capturer(self, capturer: RoutedExpertsCapturer) -> None:

7082 + from vllm.model_executor.layers.fused_moe.layer import FusedMoE

7083 + from vllm.model_executor.layers.fused_moe.router.base_router import (

7084 + BaseRouter,

7085 )

7086

7114 - def _bind_routed_experts_capturer(self) -> None:

7115 - from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (

7116 - bind_routing_capture_to_model,

7117 - )

7087 + for module in self.compilation_config.static_forward_context.values():

7088 + if isinstance(module, FusedMoE) and isinstance(module.router, BaseRouter):

7089 + layer_id = module.layer_id

7090 +

7091 + def _capture_fn(topk_ids, _layer_id=layer_id, _capturer=capturer):

7092 + _capturer.capture(_layer_id, topk_ids)

7093

7119 - bind_routing_capture_to_model(self.model)

7094 + module.router.set_capture_fn(_capture_fn)

7095

7096 def may_add_encoder_only_layers_to_kv_cache_config(self) -> None:

7097 """