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 14, 05:30 UTC
[Misc] Support routing logic simulation (#21990)
Signed-off-by: Ming Yang <minos.future@gmail.com> Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Ming Yang committed
Aug 6, 2025 at 23:06 UTC
82216dc21f777584bcf53ab1fe4936390c1737bf
4 files changed
+481
tests/test_routing_simulator.py
new
+171
@@ -0,0 +1,171 @@
1
+#!/usr/bin/env python3
2
+# SPDX-License-Identifier: Apache-2.0
3
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4
+"""
5
+Test script for the token-to-expert routing simulator.
6
+
7
+This script demonstrates how to use the routing simulator to test
8
+different routing strategies and analyze their performance, including
9
+integration tests with FusedMoE layer.
10
+"""
11
+
12
+import pytest
13
+import torch
14
+
15
+from vllm.model_executor.layers.fused_moe.routing_simulator import (
16
+ DistributionBasedRouting, RoutingSimulator)
17
+
18
+
19
+@pytest.fixture
20
+def device():
21
+ """Fixture to provide the appropriate device for testing."""
22
+ return torch.device("cuda" if torch.cuda.is_available() else "cpu")
23
+
24
+
25
+@pytest.mark.parametrize("num_tokens", [1, 16, 256])
26
+@pytest.mark.parametrize("hidden_size", [64, 1024])
27
+@pytest.mark.parametrize("num_experts", [16, 128])
28
+@pytest.mark.parametrize("top_k", [1, 4])
29
+def test_basic_functionality(
30
+ num_tokens: int,
31
+ hidden_size: int,
32
+ num_experts: int,
33
+ top_k: int,
34
+ device,
35
+):
36
+ """Test basic functionality of the routing simulator."""
37
+ # Test each routing strategy
38
+ strategies = RoutingSimulator.get_available_strategies()
39
+
40
+ hidden_states = torch.randn(num_tokens, hidden_size, device=device)
41
+ router_logits = torch.randn(num_tokens, num_experts, device=device)
42
+
43
+ for strategy in strategies:
44
+ # Simulate routing
45
+ topk_weights, topk_ids = RoutingSimulator.simulate_routing(
46
+ hidden_states=hidden_states,
47
+ router_logits=router_logits,
48
+ strategy_name=strategy,
49
+ top_k=top_k,
50
+ )
51
+
52
+ # Check output shapes
53
+ assert topk_weights.shape == (
54
+ num_tokens,
55
+ top_k,
56
+ ), f"Wrong weights shape for {strategy}"
57
+ assert topk_ids.shape == (
58
+ num_tokens,
59
+ top_k,
60
+ ), f"Wrong ids shape for {strategy}"
61
+
62
+ # Check that expert IDs are valid
63
+ assert (topk_ids.min()
64
+ >= 0), f"Invalid expert ID (negative) for {strategy}"
65
+ assert (topk_ids.max()
66
+ < num_experts), f"Invalid expert ID (too large) for {strategy}"
67
+
68
+
69
+def test_routing_strategy_integration(monkeypatch, device):
70
+ """Test that the routing strategy environment variable works with
71
+ FusedMoE."""
72
+ pytest.importorskip("vllm.model_executor.layers.fused_moe.layer")
73
+
74
+ import vllm.envs as envs
75
+ from vllm.model_executor.layers.fused_moe.layer import FusedMoE
76
+
77
+ # Test parameters
78
+ num_tokens = 32
79
+ hidden_size = 16
80
+ num_experts = 4
81
+ top_k = 2
82
+
83
+ # Create test data
84
+ hidden_states = torch.randn(num_tokens, hidden_size, device=device)
85
+ router_logits = torch.randn(num_tokens, num_experts, device=device)
86
+
87
+ # Test different routing strategies
88
+ strategies = RoutingSimulator.get_available_strategies()
89
+
90
+ for strategy in strategies:
91
+ # Set environment variable
92
+ env_name = "VLLM_MOE_ROUTING_SIMULATION_STRATEGY"
93
+ monkeypatch.setenv(env_name, strategy)
94
+
95
+ # Force reload of environment variable
96
+ envs.environment_variables[env_name] = lambda s=strategy: s
97
+
98
+ # Test the select_experts method
99
+ topk_weights, topk_ids = FusedMoE.select_experts(
100
+ hidden_states=hidden_states,
101
+ router_logits=router_logits,
102
+ top_k=top_k,
103
+ use_grouped_topk=False,
104
+ renormalize=True,
105
+ indices_type=torch.long)
106
+
107
+ # Verify output shapes
108
+ assert topk_weights.shape == (
109
+ num_tokens, top_k), f"Wrong weights shape for {strategy}"
110
+ assert topk_ids.shape == (num_tokens,
111
+ top_k), f"Wrong ids shape for {strategy}"
112
+
113
+ # Verify expert IDs are valid
114
+ assert topk_ids.min(
115
+ ) >= 0, f"Invalid expert ID (negative) for {strategy}"
116
+ assert topk_ids.max(
117
+ ) < num_experts, f"Invalid expert ID (too large) for {strategy}"
118
+
119
+
120
+def test_distribution_based_routing_with_custom_strategy():
121
+ """Test registering and using DistributionBasedRouting with custom
122
+ parameters."""
123
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
124
+
125
+ # Register custom distribution-based strategy
126
+ custom_strategy = DistributionBasedRouting(distribution="normal",
127
+ mean=2.0,
128
+ std=0.5)
129
+ RoutingSimulator.register_strategy("custom_normal", custom_strategy)
130
+
131
+ # Test data
132
+ num_tokens = 60
133
+ hidden_size = 48
134
+ num_experts = 6
135
+ top_k = 3
136
+
137
+ hidden_states = torch.randn(num_tokens, hidden_size, device=device)
138
+ router_logits = torch.randn(num_tokens, num_experts, device=device)
139
+
140
+ # Use the custom strategy
141
+ topk_weights, topk_ids = RoutingSimulator.simulate_routing(
142
+ hidden_states=hidden_states,
143
+ router_logits=router_logits,
144
+ strategy_name="custom_normal",
145
+ top_k=top_k)
146
+
147
+ # Check output shapes
148
+ assert topk_weights.shape == (num_tokens, top_k)
149
+ assert topk_ids.shape == (num_tokens, top_k)
150
+
151
+ # Check that expert IDs are valid
152
+ assert topk_ids.min() >= 0
153
+ assert topk_ids.max() < num_experts
154
+
155
+
156
+def test_instance_compatibility():
157
+ """Test that static methods work correctly."""
158
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
159
+
160
+ # Test static method directly
161
+ hidden_states = torch.randn(10, 8, device=device)
162
+ router_logits = torch.randn(10, 4, device=device)
163
+
164
+ topk_weights, topk_ids = RoutingSimulator.simulate_routing(
165
+ hidden_states=hidden_states,
166
+ router_logits=router_logits,
167
+ strategy_name="uniform_random",
168
+ top_k=2)
169
+
170
+ assert topk_weights.shape == (10, 2)
171
+ assert topk_ids.shape == (10, 2)
vllm/envs.py
+9
@@ -989,6 +989,15 @@ environment_variables: dict[str, Callable[[], Any]] = {
989
"VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE":
990
lambda: int(os.getenv("VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE", "163840")),
991
992
+ # MoE routing strategy selector.
993
+ # See `RoutingSimulator.get_available_strategies()` # for available
994
+ # strategies.
995
+ # Cutstom routing strategies can be registered by
996
+ # RoutingSimulator.register_strategy()
997
+ # Note: custom strategies may not produce correct model outputs
998
+ "VLLM_MOE_ROUTING_SIMULATION_STRATEGY":
999
+ lambda: os.environ.get("VLLM_MOE_ROUTING_SIMULATION_STRATEGY", "").lower(),
1000
+
1001
# Regex timeout for use by the vLLM tool parsing plugins.
1002
"VLLM_TOOL_PARSE_REGEX_TIMEOUT_SECONDS":
1003
lambda: int(os.getenv("VLLM_TOOL_PARSE_REGEX_TIMEOUT_SECONDS", "1")),
vllm/model_executor/layers/fused_moe/layer.py
+12
@@ -28,6 +28,8 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import (
28
FusedMoEPermuteExpertsUnpermute, FusedMoEPrepareAndFinalize)
29
from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
30
is_rocm_aiter_moe_enabled)
31
+from vllm.model_executor.layers.fused_moe.routing_simulator import (
32
+ RoutingSimulator)
33
from vllm.model_executor.layers.quantization.base_config import (
34
QuantizationConfig, QuantizeMethodBase)
35
from vllm.model_executor.utils import set_weight_attrs
@@ -1362,6 +1364,16 @@ class FusedMoE(torch.nn.Module):
1364
"""
1365
from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
1366
1367
+ # Check if we should use a routing simulation strategy
1368
+ routing_strategy = envs.VLLM_MOE_ROUTING_SIMULATION_STRATEGY
1369
+ if routing_strategy != "":
1370
+ return RoutingSimulator.simulate_routing(
1371
+ hidden_states=hidden_states,
1372
+ router_logits=router_logits,
1373
+ strategy_name=routing_strategy,
1374
+ top_k=top_k,
1375
+ indices_type=indices_type)
1376
+
1377
# DeepSeekv2 uses grouped_top_k
1378
if use_grouped_topk:
1379
assert topk_group is not None
vllm/model_executor/layers/fused_moe/routing_simulator.py
new
+289
@@ -0,0 +1,289 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Token-to-Expert Routing Simulator
5
+
6
+This module provides a framework for simulating and testing different
7
+token-to-expert routing strategies for Mixture of Experts (MoE) models.
8
+It supports routing logic customization and includes example implementations
9
+like uniform random routing.
10
+"""
11
+
12
+from abc import ABC, abstractmethod
13
+from typing import Optional
14
+
15
+import torch
16
+
17
+
18
+class RoutingStrategy(ABC):
19
+ """Base class for token-to-expert routing strategies."""
20
+
21
+ @abstractmethod
22
+ def route_tokens(
23
+ self,
24
+ hidden_states: torch.Tensor,
25
+ router_logits: torch.Tensor,
26
+ top_k: int,
27
+ indices_type: Optional[torch.dtype] = None,
28
+ ) -> tuple[torch.Tensor, torch.Tensor]:
29
+ """
30
+ Route tokens to experts.
31
+
32
+ Args:
33
+ hidden_states: Input hidden states [num_tokens, hidden_size]
34
+ router_logits: Router logits [num_tokens, num_experts]
35
+ top_k: Number of experts to select per token
36
+ indices_type: Data type for expert indices
37
+
38
+ Returns:
39
+ tuple of (topk_weights, topk_ids)
40
+ """
41
+ pass
42
+
43
+
44
+class DistributionBasedRouting(RoutingStrategy):
45
+ """
46
+ Distribution-based random routing strategy with configurable distributions.
47
+
48
+ This routing strategy randomly selects experts for each token based on
49
+ different probability distributions. Currently supports uniform and normal
50
+ distributions for testing different routing patterns.
51
+ """
52
+
53
+ def __init__(self, distribution: str = "uniform", **distribution_params):
54
+ """
55
+ Initialize distribution-based routing.
56
+
57
+ Args:
58
+ distribution: Type of distribution to use for sampling
59
+ - "uniform": Uniform distribution (default)
60
+ - "normal": Normal/Gaussian distribution
61
+ **distribution_params: Parameters specific to the
62
+ chosen distribution
63
+ For "uniform": No additional parameters needed
64
+ For "normal": mean (default: 0.0), std (default: 1.0)
65
+ """
66
+ self.distribution = distribution.lower()
67
+ self.distribution_params = distribution_params
68
+
69
+ # Validate distribution and parameters
70
+ self._validate_distribution_params()
71
+
72
+ def _validate_distribution_params(self):
73
+ """Validate distribution type and parameters."""
74
+ valid_distributions = ["uniform", "normal"]
75
+
76
+ if self.distribution not in valid_distributions:
77
+ raise ValueError(f"Unsupported distribution: {self.distribution}. "
78
+ f"Supported distributions: {valid_distributions}")
79
+
80
+ # Set default parameters if not provided
81
+ if self.distribution == "normal":
82
+ self.distribution_params.setdefault("mean", 0.0)
83
+ self.distribution_params.setdefault("std", 1.0)
84
+
85
+ def route_tokens(
86
+ self,
87
+ hidden_states: torch.Tensor,
88
+ router_logits: torch.Tensor,
89
+ top_k: int,
90
+ indices_type: Optional[torch.dtype] = None,
91
+ ) -> tuple[torch.Tensor, torch.Tensor]:
92
+ """
93
+ Randomly select experts for each token using the specified distribution.
94
+
95
+ Args:
96
+ hidden_states: Input hidden states [num_tokens, hidden_size]
97
+ router_logits: Router logits [num_tokens, num_experts]
98
+ top_k: Number of experts to select per token
99
+ indices_type: Data type for expert indices
100
+
101
+ Returns:
102
+ tuple of (topk_weights, topk_ids) where:
103
+ - topk_weights: Weights based on distribution sampling
104
+ - topk_ids: Expert indices sampled from the distribution
105
+ """
106
+ num_tokens = hidden_states.shape[0]
107
+ num_experts = router_logits.shape[-1]
108
+
109
+ if indices_type is None:
110
+ indices_type = torch.long
111
+
112
+ # Generate expert IDs based on the specified distribution
113
+ topk_ids = self._sample_expert_ids(num_tokens, num_experts, top_k,
114
+ hidden_states.device, indices_type)
115
+
116
+ # Generate weights based on the distribution
117
+ topk_weights = self._generate_weights(num_tokens, top_k,
118
+ hidden_states.device)
119
+
120
+ return topk_weights, topk_ids
121
+
122
+ def _sample_expert_ids(
123
+ self,
124
+ num_tokens: int,
125
+ num_experts: int,
126
+ top_k: int,
127
+ device: torch.device,
128
+ indices_type: torch.dtype,
129
+ ) -> torch.Tensor:
130
+ """Sample expert IDs based on the specified distribution."""
131
+
132
+ if self.distribution == "uniform":
133
+ # Uniform random sampling
134
+ return torch.randint(
135
+ low=0,
136
+ high=num_experts,
137
+ size=(num_tokens, top_k),
138
+ dtype=indices_type,
139
+ device=device,
140
+ )
141
+
142
+ elif self.distribution == "normal":
143
+ # For normal distribution, sample continuous values and map to
144
+ # expert IDs
145
+ continuous_samples = self._sample_continuous_distribution(
146
+ num_tokens, top_k, device)
147
+
148
+ # Map continuous samples to expert indices
149
+ # Normalize to [0, 1] range and scale to [0, num_experts)
150
+ normalized_samples = self._normalize_samples(continuous_samples)
151
+ expert_ids = (normalized_samples * num_experts).long()
152
+ expert_ids = torch.clamp(expert_ids, 0, num_experts - 1)
153
+
154
+ return expert_ids.to(dtype=indices_type)
155
+
156
+ else:
157
+ raise ValueError(f"Unsupported distribution: {self.distribution}")
158
+
159
+ def _sample_continuous_distribution(self, num_tokens: int, top_k: int,
160
+ device: torch.device) -> torch.Tensor:
161
+ """Sample from continuous distributions."""
162
+ shape = (num_tokens, top_k)
163
+
164
+ if self.distribution == "normal":
165
+ mean = self.distribution_params["mean"]
166
+ std = self.distribution_params["std"]
167
+ return torch.normal(mean, std, size=shape, device=device)
168
+
169
+ else:
170
+ raise ValueError(
171
+ f"Unsupported continuous distribution: {self.distribution}")
172
+
173
+ def _normalize_samples(self, samples: torch.Tensor) -> torch.Tensor:
174
+ """Normalize samples to [0, 1] range."""
175
+ if self.distribution == "normal":
176
+ # Use sigmoid to map normal distribution to [0, 1]
177
+ return torch.sigmoid(samples)
178
+
179
+ else:
180
+ raise ValueError(f"Unsupported distribution for normalization: "
181
+ f"{self.distribution}")
182
+
183
+ def _generate_weights(self, num_tokens: int, top_k: int,
184
+ device: torch.device) -> torch.Tensor:
185
+ """Generate weights based on the distribution."""
186
+ if self.distribution == "uniform":
187
+ # All-ones weights for uniform distribution
188
+ return torch.ones(
189
+ (num_tokens, top_k),
190
+ dtype=torch.float32,
191
+ device=device,
192
+ )
193
+
194
+ elif self.distribution == "normal":
195
+ # For normal distribution, generate weights from the same
196
+ # distribution
197
+ continuous_weights = self._sample_continuous_distribution(
198
+ num_tokens, top_k, device)
199
+ # Normalize to positive values and sum to 1
200
+ weights = torch.abs(continuous_weights)
201
+ weights = weights / weights.sum(dim=-1, keepdim=True)
202
+ return weights
203
+
204
+ else:
205
+ raise ValueError(
206
+ f"Unsupported distribution for weight generation: "
207
+ f"{self.distribution}")
208
+
209
+ def get_distribution_info(self) -> dict:
210
+ """Get information about the current distribution configuration."""
211
+ return {
212
+ "distribution": self.distribution,
213
+ "parameters": self.distribution_params.copy()
214
+ }
215
+
216
+
217
+class RoutingSimulator:
218
+ """
219
+ Token-to-Expert Routing Simulator.
220
+
221
+ This class provides a framework for testing and comparing different
222
+ routing strategies for MoE models. It can simulate routing behavior
223
+ and collect statistics for analysis.
224
+ """
225
+
226
+ # Class-level registry of routing strategies
227
+ _routing_strategies: dict[str, RoutingStrategy] = {
228
+ # Basic routing strategies
229
+ "uniform_random":
230
+ DistributionBasedRouting(distribution="uniform", mean=0.0, std=1.0),
231
+ "normal_routing":
232
+ DistributionBasedRouting(distribution="normal", mean=0.0, std=1.0),
233
+ }
234
+
235
+ @classmethod
236
+ def register_strategy(cls, name: str, strategy: RoutingStrategy):
237
+ """
238
+ Register a custom routing strategy.
239
+
240
+ Args:
241
+ name: Name of the strategy
242
+ strategy: RoutingStrategy instance
243
+ """
244
+ cls._routing_strategies[name] = strategy
245
+
246
+ @classmethod
247
+ def get_available_strategies(cls):
248
+ """
249
+ Get list of available routing strategy names.
250
+
251
+ Returns:
252
+ List of available strategy names
253
+ """
254
+ return list(cls._routing_strategies.keys())
255
+
256
+ @staticmethod
257
+ def simulate_routing(
258
+ hidden_states: torch.Tensor,
259
+ router_logits: torch.Tensor,
260
+ strategy_name: str,
261
+ top_k: int,
262
+ indices_type: Optional[torch.dtype] = None,
263
+ ) -> tuple[torch.Tensor, torch.Tensor]:
264
+ """
265
+ Simulate token-to-expert routing using the specified strategy.
266
+
267
+ Args:
268
+ hidden_states: Input hidden states [num_tokens, hidden_size]
269
+ router_logits: Router logits [num_tokens, num_experts]
270
+ strategy_name: Name of the routing strategy to use
271
+ top_k: Number of experts to select per token
272
+ indices_type: Data type for expert indices
273
+
274
+ Returns:
275
+ tuple of (topk_weights, topk_ids)
276
+ """
277
+ if strategy_name not in RoutingSimulator._routing_strategies:
278
+ raise ValueError(
279
+ f"Unknown routing strategy: {strategy_name}. "
280
+ f"Available strategies: "
281
+ f"{list(RoutingSimulator._routing_strategies.keys())}")
282
+
283
+ strategy = RoutingSimulator._routing_strategies[strategy_name]
284
+ return strategy.route_tokens(
285
+ hidden_states=hidden_states,
286
+ router_logits=router_logits,
287
+ top_k=top_k,
288
+ indices_type=indices_type,
289
+ )