| 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 Mock, patch |
| 6 | |
| 7 | import pytest |
| 8 | import torch |
| 9 | |
| 10 | from vllm.config import VllmConfig |
| 11 | from vllm.config.compilation import CompilationMode |
| 12 | from vllm.distributed.eplb.eplb_state import EplbLayerState |
| 13 | from vllm.model_executor.layers.fused_moe.config import RoutingMethodType |
| 14 | from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( |
| 15 | RoutedExpertsCapturer, |
| 16 | bind_routed_experts_capturer, |
| 17 | ) |
| 18 | from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter |
| 19 | from vllm.v1.kv_cache_interface import ( |
| 20 | FullAttentionSpec, |
| 21 | KVCacheGroupSpec, |
| 22 | MLAAttentionSpec, |
| 23 | UniformTypeKVCacheSpecs, |
| 24 | ) |
| 25 | |
| 26 | pytestmark = pytest.mark.cpu_test |
| 27 | |
| 28 | _REC_MODULE = "vllm.model_executor.layers.fused_moe.routed_experts_capturer" |
| 29 | |
| 30 | |
| 31 | def _capturer_with_buffer( |
| 32 | *, |
| 33 | max_tokens: int = 8, |
| 34 | num_layers: int = 4, |
| 35 | num_experts_per_tok: int = 2, |
| 36 | dp_rank: int = 0, |
| 37 | tp_size: int = 1, |
| 38 | dtype: torch.dtype = torch.int32, |
| 39 | ) -> RoutedExpertsCapturer: |
| 40 | # Bypass __init__ so the test can use a CPU buffer and skip the |
| 41 | # VllmConfig dependency. The CUDA device-tensor allocation in the |
| 42 | # real constructor is not what we are exercising here. |
| 43 | c = RoutedExpertsCapturer.__new__(RoutedExpertsCapturer) |
| 44 | c.dp_rank = dp_rank |
| 45 | c.tp_size = tp_size |
| 46 | c.output_dtype = torch.uint8 |
| 47 | c.device_buffer = torch.full( |
| 48 | (max_tokens, num_layers, num_experts_per_tok), |
| 49 | -1, |
| 50 | dtype=dtype, |
| 51 | ) |
| 52 | return c |
| 53 | |
| 54 | |
| 55 | class DummyRouter(BaseRouter): |
| 56 | @property |
| 57 | def routing_method_type(self) -> RoutingMethodType: |
| 58 | return RoutingMethodType.FUSED_TOPK |
| 59 | |
| 60 | def _compute_routing( |
| 61 | self, hidden_states, router_logits, indices_type, *, input_ids=None |
| 62 | ): |
| 63 | topk_ids = torch.tensor([[1, 2], [3, 4]], dtype=torch.int64) |
| 64 | topk_weights = torch.ones_like(topk_ids, dtype=torch.float32) |
| 65 | return topk_weights, topk_ids |
| 66 | |
| 67 | def _apply_eplb_mapping(self, topk_ids: torch.Tensor) -> torch.Tensor: |
| 68 | # Make mapping observable without requiring CUDA EPLB path. |
| 69 | return topk_ids + 10 |
| 70 | |
| 71 | |
| 72 | def _make_router(eplb_state: EplbLayerState | None = None) -> DummyRouter: |
| 73 | return DummyRouter( |
| 74 | top_k=2, |
| 75 | global_num_experts=16, |
| 76 | eplb_state=eplb_state, |
| 77 | ) |
| 78 | |
| 79 | |
| 80 | def _make_modular_routed_experts(): |
| 81 | return types.SimpleNamespace( |
| 82 | quant_method=types.SimpleNamespace(is_monolithic=False), |
| 83 | ) |
| 84 | |
| 85 | |
| 86 | def _full_attention_kv_group( |
| 87 | spec_type: type[FullAttentionSpec] = FullAttentionSpec, |
| 88 | ) -> KVCacheGroupSpec: |
| 89 | full_attention = spec_type( |
| 90 | block_size=16, |
| 91 | num_kv_heads=1, |
| 92 | head_size=1, |
| 93 | dtype=torch.float32, |
| 94 | ) |
| 95 | return KVCacheGroupSpec( |
| 96 | ["layer"], |
| 97 | UniformTypeKVCacheSpecs( |
| 98 | block_size=16, |
| 99 | kv_cache_specs={"layer": full_attention}, |
| 100 | ), |
| 101 | ) |
| 102 | |
| 103 | |
| 104 | @pytest.mark.parametrize("eplb_enabled", [False, True]) |
| 105 | def test_base_router_capture_pre_eplb_mapping(eplb_enabled): |
| 106 | eplb_state = None |
| 107 | if eplb_enabled: |
| 108 | eplb_state = EplbLayerState() |
| 109 | eplb_state.expert_load_view = torch.zeros(32, dtype=torch.int64) |
| 110 | eplb_state.logical_to_physical_map = torch.arange(32).view(32, 1) |
| 111 | eplb_state.logical_replica_count = torch.ones(32, dtype=torch.int64) |
| 112 | eplb_state.should_record_tensor = torch.ones((), dtype=torch.bool) |
| 113 | eplb_state.num_unpadded_tokens_tensors = [torch.tensor(0, dtype=torch.int32)] |
| 114 | router = _make_router(eplb_state) |
| 115 | |
| 116 | captured = [] |
| 117 | |
| 118 | def capture_fn(ids): |
| 119 | captured.append(ids.clone()) |
| 120 | |
| 121 | router.set_capture_fn(capture_fn) |
| 122 | topk_weights, topk_ids = router.select_experts( |
| 123 | hidden_states=torch.empty(1), |
| 124 | router_logits=torch.empty(1), |
| 125 | ) |
| 126 | |
| 127 | assert topk_weights.shape == topk_ids.shape |
| 128 | assert len(captured) == 1 |
| 129 | assert torch.equal(captured[0], torch.tensor([[1, 2], [3, 4]])) |
| 130 | assert torch.equal(topk_ids, torch.tensor([[11, 12], [13, 14]])) |
| 131 | |
| 132 | |
| 133 | def test_public_binding_binds_target_model_router(monkeypatch): |
| 134 | class DummyFusedMoE: |
| 135 | def __init__(self, layer_id): |
| 136 | self.layer_id = layer_id |
| 137 | self.router = _make_router() |
| 138 | self._quant_method = _make_modular_routed_experts().quant_method |
| 139 | |
| 140 | target_module = DummyFusedMoE(layer_id=7) |
| 141 | |
| 142 | import vllm.model_executor.layers.fused_moe.layer as fused_moe_layer |
| 143 | |
| 144 | monkeypatch.setattr(fused_moe_layer, "MoERunner", DummyFusedMoE) |
| 145 | calls = [] |
| 146 | capturer = types.SimpleNamespace(capture=lambda *args: calls.append(args)) |
| 147 | |
| 148 | bind_routed_experts_capturer( |
| 149 | types.SimpleNamespace(modules=lambda: [target_module]), capturer |
| 150 | ) |
| 151 | |
| 152 | assert target_module.router.capture_fn is not None |
| 153 | topk_ids = torch.tensor([[5, 6]]) |
| 154 | target_module.router.capture_fn(topk_ids) |
| 155 | assert calls == [(7, topk_ids)] |
| 156 | |
| 157 | |
| 158 | def test_public_binding_supports_direct_capture_source(): |
| 159 | class DirectCaptureSource(torch.nn.Module): |
| 160 | def __init__(self): |
| 161 | super().__init__() |
| 162 | self.layer_id = 3 |
| 163 | self.capture_fn = None |
| 164 | |
| 165 | source = DirectCaptureSource() |
| 166 | calls = [] |
| 167 | capturer = types.SimpleNamespace(capture=lambda *args: calls.append(args)) |
| 168 | |
| 169 | bind_routed_experts_capturer(source, capturer) |
| 170 | |
| 171 | assert source.capture_fn is not None |
| 172 | topk_ids = torch.tensor([[1, 2]]) |
| 173 | source.capture_fn(topk_ids) |
| 174 | assert calls == [(3, topk_ids)] |
| 175 | |
| 176 | |
| 177 | def test_public_binding_rejects_monolithic_without_replay_support(monkeypatch): |
| 178 | class DummyFusedMoE: |
| 179 | def __init__(self): |
| 180 | self.layer_id = 3 |
| 181 | self.router = _make_router() |
| 182 | # Use a concrete monolithic expert and override its capability |
| 183 | # instead of instantiating the abstract base class directly. |
| 184 | from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( |
| 185 | CPUExpertsFp8, |
| 186 | ) |
| 187 | |
| 188 | fused_experts = CPUExpertsFp8.__new__(CPUExpertsFp8) |
| 189 | self.routed_experts = types.SimpleNamespace( |
| 190 | quant_method=types.SimpleNamespace( |
| 191 | is_monolithic=True, |
| 192 | moe_kernel=types.SimpleNamespace( |
| 193 | impl=types.SimpleNamespace(fused_experts=fused_experts) |
| 194 | ), |
| 195 | ) |
| 196 | ) |
| 197 | self._quant_method = self.routed_experts.quant_method |
| 198 | self._quant_method.moe_kernel.impl.fused_experts = fused_experts |
| 199 | fused_experts.supports_routing_replay_capture = lambda: False |
| 200 | |
| 201 | class DummyCapturer: |
| 202 | def capture(self, layer_id, topk_ids): |
| 203 | pass |
| 204 | |
| 205 | dummy_module = DummyFusedMoE() |
| 206 | import vllm.model_executor.layers.fused_moe.layer as fused_moe_layer |
| 207 | |
| 208 | monkeypatch.setattr(fused_moe_layer, "MoERunner", DummyFusedMoE) |
| 209 | |
| 210 | with pytest.raises(ValueError, match="monolithic MoE kernel"): |
| 211 | bind_routed_experts_capturer( |
| 212 | types.SimpleNamespace(modules=lambda: [dummy_module]), DummyCapturer() |
| 213 | ) |
| 214 | |
| 215 | |
| 216 | def test_routed_experts_capturer_single_dp_no_metadata(): |
| 217 | """dp_metadata is None: capture writes the full topk_ids rows.""" |
| 218 | capturer = _capturer_with_buffer(dp_rank=0) |
| 219 | topk = torch.tensor([[1, 2], [3, 4], [5, 6]], dtype=torch.int32) |
| 220 | ctx = SimpleNamespace(dp_metadata=None) |
| 221 | with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx): |
| 222 | capturer.capture(layer_id=0, topk_ids=topk) |
| 223 | assert torch.equal(capturer.device_buffer[:3, 0, :], topk) |
| 224 | assert capturer.device_buffer[3, 0, 0].item() == -1 |
| 225 | |
| 226 | |
| 227 | @pytest.mark.parametrize( |
| 228 | ("num_experts", "dtype_name", "torch_dtype"), |
| 229 | [(256, "uint8", torch.uint8), (257, "uint16", torch.uint16)], |
| 230 | ) |
| 231 | def test_routed_experts_capturer_exposes_output_profile( |
| 232 | monkeypatch, num_experts, dtype_name, torch_dtype |
| 233 | ): |
| 234 | import vllm.model_executor.layers.fused_moe.routed_experts_capturer as module |
| 235 | |
| 236 | monkeypatch.setattr(module, "current_platform", SimpleNamespace(device_type="cpu")) |
| 237 | config = SimpleNamespace( |
| 238 | model_config=SimpleNamespace( |
| 239 | get_total_num_hidden_layers=lambda: 3, |
| 240 | get_num_experts=lambda: num_experts, |
| 241 | get_num_experts_per_tok=lambda: 2, |
| 242 | hf_text_config=SimpleNamespace(model_type="test"), |
| 243 | ), |
| 244 | parallel_config=SimpleNamespace(data_parallel_rank=0, tensor_parallel_size=1), |
| 245 | ) |
| 246 | |
| 247 | capturer = RoutedExpertsCapturer(8, config) |
| 248 | |
| 249 | assert capturer.shape_per_token == (3, 2) |
| 250 | assert capturer.output_dtype_name == dtype_name |
| 251 | assert capturer.output_dtype == torch_dtype |
| 252 | assert capturer.device_buffer.shape == (8, 3, 2) |
| 253 | |
| 254 | |
| 255 | @pytest.mark.parametrize("output_dtype", [torch.uint8, torch.uint16]) |
| 256 | def test_routed_experts_capturer_narrows_snapshot(output_dtype): |
| 257 | capturer = _capturer_with_buffer(dtype=torch.int32) |
| 258 | capturer.output_dtype = output_dtype |
| 259 | topk = torch.tensor([[1, 2], [254, 255]], dtype=torch.int64) |
| 260 | ctx = SimpleNamespace(dp_metadata=None) |
| 261 | with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx): |
| 262 | capturer.capture(layer_id=0, topk_ids=topk) |
| 263 | |
| 264 | output = capturer.snapshot_routing_data(2) |
| 265 | assert capturer.device_buffer.dtype == torch.int32 |
| 266 | assert capturer.device_buffer[:2, 0, :].tolist() == topk.tolist() |
| 267 | assert output.dtype == output_dtype |
| 268 | assert output[:, 0, :].tolist() == topk.tolist() |
| 269 | assert output.data_ptr() != capturer.device_buffer.data_ptr() |
| 270 | |
| 271 | |
| 272 | def test_routed_experts_capturer_dp_naive_concatenated_all_ranks(): |
| 273 | """N == sum(num_tokens_dp): slice this rank's segment from concatenated topk.""" |
| 274 | capturer = _capturer_with_buffer(dp_rank=1) |
| 275 | num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32) |
| 276 | ctx = SimpleNamespace( |
| 277 | dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp) |
| 278 | ) |
| 279 | # Concatenated order: rank0 rows then rank1 rows. |
| 280 | topk = torch.tensor( |
| 281 | [[0, 1], [2, 3], [10, 11], [12, 13], [14, 15]], dtype=torch.int32 |
| 282 | ) |
| 283 | with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx): |
| 284 | capturer.capture(layer_id=0, topk_ids=topk) |
| 285 | want = topk[2:5] |
| 286 | assert torch.equal(capturer.device_buffer[:3, 0, :], want) |
| 287 | |
| 288 | |
| 289 | def test_routed_experts_capturer_dp_modular_local_tokens(): |
| 290 | """N == token_num_per_dp: topk is already local to this DP rank.""" |
| 291 | capturer = _capturer_with_buffer(dp_rank=1) |
| 292 | num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32) |
| 293 | ctx = SimpleNamespace( |
| 294 | dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp) |
| 295 | ) |
| 296 | topk = torch.tensor([[10, 11], [12, 13], [14, 15]], dtype=torch.int32) |
| 297 | with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx): |
| 298 | capturer.capture(layer_id=0, topk_ids=topk) |
| 299 | assert torch.equal(capturer.device_buffer[:3, 0, :], topk) |
| 300 | |
| 301 | |
| 302 | def test_routed_experts_capturer_dp_ep_gathered_shards(): |
| 303 | capturer = _capturer_with_buffer(dp_rank=1, tp_size=2) |
| 304 | num_tokens_dp = torch.tensor([3, 2], dtype=torch.int32) |
| 305 | ctx = SimpleNamespace( |
| 306 | dp_metadata=SimpleNamespace( |
| 307 | num_tokens_across_dp_cpu=num_tokens_dp, |
| 308 | local_sizes=[2, 2, 1, 1], |
| 309 | ) |
| 310 | ) |
| 311 | topk = torch.tensor( |
| 312 | [[0, 1], [2, 3], [4, 5], [-1, -1], [10, 11], [12, 13]], |
| 313 | dtype=torch.int32, |
| 314 | ) |
| 315 | |
| 316 | with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx): |
| 317 | capturer.capture(layer_id=0, topk_ids=topk) |
| 318 | |
| 319 | assert torch.equal(capturer.device_buffer[:2, 0], topk[4:]) |
| 320 | |
| 321 | |
| 322 | def test_routed_experts_capturer_sp_modular_gathers_tp_shards(): |
| 323 | capturer = _capturer_with_buffer(dp_rank=1, tp_size=2) |
| 324 | num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32) |
| 325 | ctx = SimpleNamespace( |
| 326 | dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp) |
| 327 | ) |
| 328 | local_shard = torch.tensor([[10, 11], [12, 13]], dtype=torch.int32) |
| 329 | gathered = torch.tensor([[10, 11], [12, 13], [14, 15], [-1, -1]], dtype=torch.int32) |
| 330 | tp_group = Mock() |
| 331 | tp_group.all_gather.return_value = gathered |
| 332 | |
| 333 | with ( |
| 334 | patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx), |
| 335 | patch(f"{_REC_MODULE}.get_tp_group", return_value=tp_group), |
| 336 | ): |
| 337 | capturer.capture(layer_id=0, topk_ids=local_shard) |
| 338 | |
| 339 | tp_group.all_gather.assert_called_once_with(local_shard, dim=0) |
| 340 | assert torch.equal(capturer.device_buffer[:3, 0], gathered[:3]) |
| 341 | |
| 342 | |
| 343 | def test_routed_experts_capturer_dp_unexpected_batch_raises(): |
| 344 | """Mismatch between topk batch dim and DP layout: fail fast.""" |
| 345 | capturer = _capturer_with_buffer(dp_rank=0) |
| 346 | num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32) |
| 347 | ctx = SimpleNamespace( |
| 348 | dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp) |
| 349 | ) |
| 350 | # total=5, local=2: n=1 matches neither naive (5) nor modular (2). |
| 351 | topk = torch.tensor([[1, 2]], dtype=torch.int32) |
| 352 | with ( |
| 353 | patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx), |
| 354 | pytest.raises(AssertionError, match="unexpected topk_ids batch dim"), |
| 355 | ): |
| 356 | capturer.capture(layer_id=0, topk_ids=topk) |
| 357 | assert capturer.device_buffer[0, 0, 0].item() == -1 |
| 358 | |
| 359 | |
| 360 | def test_get_aux_output_connector_passes_config(monkeypatch): |
| 361 | import vllm.distributed.aux_output_connector.worker as aux_output_worker |
| 362 | |
| 363 | connector = Mock() |
| 364 | constructor = Mock(return_value=connector) |
| 365 | monkeypatch.setattr(aux_output_worker, "AuxOutputWorkerConnector", constructor) |
| 366 | config = SimpleNamespace( |
| 367 | scheduler_config=SimpleNamespace(max_num_batched_tokens=32) |
| 368 | ) |
| 369 | model = Mock() |
| 370 | kv_cache_config = Mock() |
| 371 | |
| 372 | result = aux_output_worker.get_aux_output_connector(model, config, kv_cache_config) |
| 373 | |
| 374 | constructor.assert_called_once_with( |
| 375 | model=model, |
| 376 | kv_cache_config=kv_cache_config, |
| 377 | vllm_config=config, |
| 378 | ) |
| 379 | assert result is connector |
| 380 | |
| 381 | |
| 382 | def test_aux_output_worker_connector_binds_capture_on_non_output_rank(monkeypatch): |
| 383 | import vllm.distributed.aux_output_connector.worker as aux_output_worker |
| 384 | |
| 385 | capturer = Mock() |
| 386 | constructor = Mock(return_value=capturer) |
| 387 | bind = Mock() |
| 388 | monkeypatch.setattr(aux_output_worker, "RoutedExpertsCapturer", constructor) |
| 389 | monkeypatch.setattr(aux_output_worker, "bind_routed_experts_capturer", bind) |
| 390 | monkeypatch.setattr( |
| 391 | aux_output_worker, |
| 392 | "get_tp_group", |
| 393 | lambda: SimpleNamespace(is_first_rank=False, world_size=1), |
| 394 | ) |
| 395 | |
| 396 | config = SimpleNamespace( |
| 397 | aux_output_config=SimpleNamespace(enable_return_routed_experts=True), |
| 398 | kv_transfer_config=None, |
| 399 | max_concurrent_batches=2, |
| 400 | scheduler_config=SimpleNamespace(max_num_batched_tokens=32), |
| 401 | ) |
| 402 | model = Mock() |
| 403 | connector = aux_output_worker.AuxOutputWorkerConnector( |
| 404 | vllm_config=config, |
| 405 | model=model, |
| 406 | kv_cache_config=SimpleNamespace(kv_cache_groups=[_full_attention_kv_group()]), |
| 407 | ) |
| 408 | |
| 409 | constructor.assert_called_once_with( |
| 410 | max_num_batched_tokens=32, |
| 411 | vllm_config=config, |
| 412 | ) |
| 413 | bind.assert_called_once_with(model, capturer) |
| 414 | connector.begin_step(Mock()) |
| 415 | assert connector.prepare_output(Mock()) is None |
| 416 | capturer.snapshot_routing_data.assert_not_called() |
| 417 | |
| 418 | |
| 419 | def test_aux_output_worker_connector_default_capacity(monkeypatch): |
| 420 | import vllm.distributed.aux_output_connector.worker as aux_output_worker |
| 421 | |
| 422 | tp_group = SimpleNamespace(is_first_rank=True, world_size=1) |
| 423 | store_constructor = Mock() |
| 424 | background_store_constructor = Mock(side_effect=lambda store, **_: store) |
| 425 | capturer = SimpleNamespace(shape_per_token=(2,), output_dtype_name="int32") |
| 426 | monkeypatch.setattr(aux_output_worker, "get_tp_group", lambda: tp_group) |
| 427 | monkeypatch.setattr( |
| 428 | aux_output_worker, "RoutedExpertsCapturer", Mock(return_value=capturer) |
| 429 | ) |
| 430 | monkeypatch.setattr(aux_output_worker, "bind_routed_experts_capturer", Mock()) |
| 431 | monkeypatch.setattr( |
| 432 | aux_output_worker, |
| 433 | "resolve_kv_cache_block_sizes", |
| 434 | lambda *_: (32, 16), |
| 435 | ) |
| 436 | monkeypatch.setattr(aux_output_worker, "BlockObjectStore", store_constructor) |
| 437 | monkeypatch.setattr( |
| 438 | aux_output_worker, "BackgroundBlockObjectStore", background_store_constructor |
| 439 | ) |
| 440 | monkeypatch.setattr(aux_output_worker, "RoutedExpertsBuffer", Mock()) |
| 441 | |
| 442 | config = SimpleNamespace( |
| 443 | aux_output_config=SimpleNamespace(max_bytes=None), |
| 444 | kv_transfer_config=None, |
| 445 | cache_config=SimpleNamespace(enable_prefix_caching=True), |
| 446 | scheduler_config=SimpleNamespace(max_num_seqs=8, max_num_batched_tokens=32), |
| 447 | max_concurrent_batches=2, |
| 448 | ) |
| 449 | kwargs = dict( |
| 450 | vllm_config=config, |
| 451 | model=Mock(), |
| 452 | kv_cache_config=SimpleNamespace( |
| 453 | num_blocks=10, |
| 454 | kv_cache_groups=[_full_attention_kv_group(MLAAttentionSpec)], |
| 455 | ), |
| 456 | ) |
| 457 | |
| 458 | aux_output_worker.AuxOutputWorkerConnector(**kwargs) |
| 459 | assert store_constructor.call_args.kwargs["max_bytes"] == 2560 |
| 460 | assert store_constructor.call_args.kwargs["object_nbytes"] == 128 |
| 461 | assert background_store_constructor.call_args.kwargs["max_pending_batches"] == 16 |
| 462 | |
| 463 | |
| 464 | def test_v2_model_runner_accepts_routed_experts(monkeypatch): |
| 465 | monkeypatch.setattr("importlib.metadata.entry_points", lambda **_: ()) |
| 466 | config = SimpleNamespace( |
| 467 | model_config=SimpleNamespace( |
| 468 | use_mla=False, |
| 469 | logits_processors=None, |
| 470 | enable_prompt_embeds=False, |
| 471 | ), |
| 472 | aux_output_config=SimpleNamespace(enable_return_routed_experts=True), |
| 473 | speculative_config=None, |
| 474 | parallel_config=SimpleNamespace( |
| 475 | prefill_context_parallel_size=1, |
| 476 | tensor_parallel_size=1, |
| 477 | distributed_executor_backend=None, |
| 478 | pipeline_parallel_size=1, |
| 479 | enable_dbo=False, |
| 480 | use_ubatching=False, |
| 481 | enable_elastic_ep=False, |
| 482 | ), |
| 483 | compilation_config=SimpleNamespace( |
| 484 | mode=CompilationMode.NONE, |
| 485 | pass_config=SimpleNamespace(enable_sp=False), |
| 486 | ), |
| 487 | cache_config=SimpleNamespace( |
| 488 | kv_sharing_fast_prefill=False, |
| 489 | mamba_cache_mode="none", |
| 490 | ), |
| 491 | ec_transfer_config=None, |
| 492 | ) |
| 493 | |
| 494 | unsupported = VllmConfig._get_v2_model_runner_unsupported_features(config) |
| 495 | |
| 496 | assert "routed experts capture" not in unsupported |