| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from abc import abstractmethod |
| 4 | from collections.abc import Callable |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | from vllm.distributed.eplb.eplb_state import EplbLayerState |
| 9 | from vllm.model_executor.layers.fused_moe.router.fused_moe_router import ( |
| 10 | FusedMoERouter, |
| 11 | ) |
| 12 | from vllm.platforms import current_platform |
| 13 | from vllm.triton_utils import tl, triton |
| 14 | from vllm.v1.worker.ubatching import dbo_current_ubatch_id |
| 15 | |
| 16 | if current_platform.is_cuda_alike() or current_platform.is_xpu(): |
| 17 | |
| 18 | @triton.jit |
| 19 | def _eplb_map_and_record_i32_kernel( |
| 20 | topk_ids_ptr, |
| 21 | logical_replica_count_ptr, |
| 22 | logical_to_physical_ptr, |
| 23 | out_ids_ptr, |
| 24 | out_ptr, |
| 25 | record_enabled_ptr, |
| 26 | num_unpadded_tokens_ptr, |
| 27 | num_logical_experts, |
| 28 | map_slots, |
| 29 | out_size, |
| 30 | numel, |
| 31 | num_active_experts, |
| 32 | HAS_NUM_UNPADDED: tl.constexpr, |
| 33 | BLOCK_SIZE: tl.constexpr, |
| 34 | ): |
| 35 | pid = tl.program_id(0) |
| 36 | offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) |
| 37 | mask = offs < numel |
| 38 | |
| 39 | expert_id = tl.load(topk_ids_ptr + offs, mask=mask, other=0).to(tl.int64) |
| 40 | valid_expert = (expert_id >= 0) & (expert_id < num_logical_experts) |
| 41 | safe_expert_id = tl.where(valid_expert, expert_id, 0) |
| 42 | |
| 43 | # 1. Convert the logical expert ids to physical expert ids |
| 44 | replica_count = tl.load( |
| 45 | logical_replica_count_ptr + safe_expert_id, |
| 46 | mask=mask & valid_expert, |
| 47 | other=1, |
| 48 | ) |
| 49 | # Avoid invalid modulo/div by forcing at least 1. |
| 50 | replica_count = tl.maximum(replica_count, 1) |
| 51 | # floor(2^32 / phi), classic Knuth multiplicative hash multiplier. |
| 52 | KNUTH_MULTIPLIER = 2654435769 |
| 53 | token_idx = (offs // num_active_experts).to(tl.int64) |
| 54 | hashed = (token_idx * KNUTH_MULTIPLIER) & 0xFFFFFFFF |
| 55 | replica_idx = hashed % replica_count |
| 56 | map_index = safe_expert_id * map_slots + replica_idx |
| 57 | physical_id = tl.load( |
| 58 | logical_to_physical_ptr + map_index, |
| 59 | mask=mask & valid_expert, |
| 60 | other=-1, |
| 61 | ) |
| 62 | tl.store(out_ids_ptr + offs, physical_id, mask=mask) |
| 63 | |
| 64 | # 2. Record expert load metrics. |
| 65 | |
| 66 | # TODO(bowen): When using `FusedMoEModularKernel`, this |
| 67 | # can be done in a more unified way, since |
| 68 | # `FusedMoEPrepareAndFinalize` will return the expert |
| 69 | # token count, in some cases directly from the kernel. |
| 70 | # However, now there are many code paths not using |
| 71 | # the modular kernel, e.g. calling `fused_experts`, |
| 72 | # so we decide to keep the logic here. |
| 73 | # |
| 74 | # If later refactor moved all the MoE kernel calls |
| 75 | # to the modular kernel, we can move this logic there |
| 76 | # to achieve better efficiency. |
| 77 | |
| 78 | record_enabled = tl.load(record_enabled_ptr) != 0 |
| 79 | # Skip padded tokens when recording. |
| 80 | if HAS_NUM_UNPADDED: |
| 81 | num_unpadded_tokens = tl.load(num_unpadded_tokens_ptr) |
| 82 | is_unpadded = offs < num_unpadded_tokens * num_active_experts |
| 83 | else: |
| 84 | is_unpadded = True |
| 85 | valid = ( |
| 86 | mask |
| 87 | & record_enabled |
| 88 | & is_unpadded |
| 89 | & (physical_id >= 0) |
| 90 | & (physical_id < out_size) |
| 91 | ) |
| 92 | safe_physical_id = tl.where(physical_id >= 0, physical_id, 0) |
| 93 | tl.atomic_add(out_ptr + safe_physical_id, 1, mask=valid) |
| 94 | |
| 95 | def _eplb_map_and_record_triton( |
| 96 | topk_ids: torch.Tensor, |
| 97 | logical_to_physical_map: torch.Tensor, |
| 98 | logical_replica_count: torch.Tensor, |
| 99 | expert_load_view: torch.Tensor, |
| 100 | record_enabled: torch.Tensor, |
| 101 | num_unpadded_tokens: torch.Tensor | None, |
| 102 | ) -> torch.Tensor: |
| 103 | topk_ids_in = topk_ids.contiguous().to(dtype=torch.int32) |
| 104 | numel = topk_ids_in.numel() |
| 105 | if numel == 0: |
| 106 | return topk_ids |
| 107 | num_active_experts = topk_ids_in.shape[-1] |
| 108 | out_flat = torch.empty((numel,), device=topk_ids.device, dtype=topk_ids.dtype) |
| 109 | grid = lambda meta: (triton.cdiv(numel, meta["BLOCK_SIZE"]),) |
| 110 | assert expert_load_view.is_contiguous() |
| 111 | _eplb_map_and_record_i32_kernel[grid]( |
| 112 | topk_ids_in, |
| 113 | logical_replica_count.contiguous(), |
| 114 | logical_to_physical_map.contiguous(), |
| 115 | out_flat, |
| 116 | expert_load_view, |
| 117 | record_enabled, |
| 118 | num_unpadded_tokens, |
| 119 | logical_replica_count.shape[0], |
| 120 | logical_to_physical_map.shape[1], |
| 121 | expert_load_view.shape[0], |
| 122 | numel, |
| 123 | num_active_experts, |
| 124 | HAS_NUM_UNPADDED=num_unpadded_tokens is not None, |
| 125 | BLOCK_SIZE=256, |
| 126 | ) |
| 127 | return out_flat.reshape(topk_ids.shape) |
| 128 | |
| 129 | def eplb_map_to_physical_and_record( |
| 130 | topk_ids: torch.Tensor, |
| 131 | expert_load_view: torch.Tensor, |
| 132 | logical_to_physical_map: torch.Tensor, |
| 133 | logical_replica_count: torch.Tensor, |
| 134 | record_enabled: torch.Tensor, |
| 135 | num_unpadded_tokens: torch.Tensor | None = None, |
| 136 | ) -> torch.Tensor: |
| 137 | # Fused triton implementation: mapping + optional recording in one kernel. |
| 138 | return _eplb_map_and_record_triton( |
| 139 | topk_ids=topk_ids, |
| 140 | logical_to_physical_map=logical_to_physical_map, |
| 141 | logical_replica_count=logical_replica_count, |
| 142 | expert_load_view=expert_load_view, |
| 143 | record_enabled=record_enabled, |
| 144 | num_unpadded_tokens=num_unpadded_tokens, |
| 145 | ) |
| 146 | else: |
| 147 | |
| 148 | def eplb_map_to_physical_and_record( |
| 149 | topk_ids: torch.Tensor, |
| 150 | expert_load_view: torch.Tensor, |
| 151 | logical_to_physical_map: torch.Tensor, |
| 152 | logical_replica_count: torch.Tensor, |
| 153 | record_enabled: torch.Tensor, |
| 154 | num_unpadded_tokens: torch.Tensor | None = None, |
| 155 | ) -> torch.Tensor: |
| 156 | return topk_ids |
| 157 | |
| 158 | |
| 159 | class BaseRouter(FusedMoERouter): |
| 160 | """Base class providing common functionality for all router implementations. |
| 161 | |
| 162 | This class implements the template method pattern where select_experts() handles |
| 163 | common pre-processing and post-processing, delegating the actual routing logic |
| 164 | to the abstract _compute_routing() method. |
| 165 | """ |
| 166 | |
| 167 | def __init__( |
| 168 | self, |
| 169 | top_k: int, |
| 170 | global_num_experts: int, |
| 171 | eplb_state: EplbLayerState | None = None, |
| 172 | ): |
| 173 | """Args: |
| 174 | top_k: Number of experts to select per token |
| 175 | global_num_experts: Total number of experts |
| 176 | eplb_state: Optional EPLBLayerState for load balancing |
| 177 | |
| 178 | """ |
| 179 | super().__init__(eplb_state=eplb_state) |
| 180 | self.top_k = top_k |
| 181 | self.global_num_experts = global_num_experts |
| 182 | self.capture_fn: Callable[[torch.Tensor], None] | None = None |
| 183 | |
| 184 | def set_capture_fn(self, capture_fn: Callable[[torch.Tensor], None] | None) -> None: |
| 185 | """Set a capture callback for logical routed expert IDs.""" |
| 186 | self.capture_fn = capture_fn |
| 187 | |
| 188 | def _validate_eplb_state(self) -> None: |
| 189 | """Validate that EPLB state is properly initialized if EPLB is enabled.""" |
| 190 | if self.eplb_state is not None: |
| 191 | eplb_state = self.eplb_state |
| 192 | if eplb_state.expert_load_view is None: |
| 193 | raise ValueError("EPLB requires expert_load_view != None") |
| 194 | if eplb_state.logical_to_physical_map is None: |
| 195 | raise ValueError("EPLB requires logical_to_physical_map != None") |
| 196 | if eplb_state.logical_replica_count is None: |
| 197 | raise ValueError("EPLB requires logical_replica_count != None") |
| 198 | if eplb_state.should_record_tensor is None: |
| 199 | raise ValueError("EPLB requires should_record_tensor != None") |
| 200 | if eplb_state.num_unpadded_tokens_tensors is None: |
| 201 | raise ValueError("EPLB requires num_unpadded_tokens_tensors != None") |
| 202 | |
| 203 | def _apply_eplb_mapping(self, topk_ids: torch.Tensor) -> torch.Tensor: |
| 204 | """Apply EPLB mapping to convert logical expert IDs to physical expert IDs.""" |
| 205 | if self.eplb_state is not None: |
| 206 | eplb_state = self.eplb_state |
| 207 | assert eplb_state.expert_load_view is not None |
| 208 | assert eplb_state.logical_to_physical_map is not None |
| 209 | assert eplb_state.logical_replica_count is not None |
| 210 | assert eplb_state.should_record_tensor is not None |
| 211 | assert eplb_state.num_unpadded_tokens_tensors is not None |
| 212 | return eplb_map_to_physical_and_record( |
| 213 | topk_ids=topk_ids, |
| 214 | logical_to_physical_map=eplb_state.logical_to_physical_map, |
| 215 | logical_replica_count=eplb_state.logical_replica_count, |
| 216 | expert_load_view=eplb_state.expert_load_view, |
| 217 | record_enabled=eplb_state.should_record_tensor, |
| 218 | num_unpadded_tokens=eplb_state.num_unpadded_tokens_tensors[ |
| 219 | dbo_current_ubatch_id() |
| 220 | ], |
| 221 | ) |
| 222 | return topk_ids |
| 223 | |
| 224 | def _convert_indices_dtype( |
| 225 | self, topk_ids: torch.Tensor, indices_type: torch.dtype | None |
| 226 | ) -> torch.Tensor: |
| 227 | """Convert topk_ids to the desired dtype if needed.""" |
| 228 | if (indices_type is not None) and topk_ids.dtype != indices_type: |
| 229 | topk_ids = topk_ids.to(dtype=indices_type) |
| 230 | |
| 231 | assert topk_ids.dtype == indices_type or indices_type is None |
| 232 | return topk_ids |
| 233 | |
| 234 | @abstractmethod |
| 235 | def _compute_routing( |
| 236 | self, |
| 237 | hidden_states: torch.Tensor, |
| 238 | router_logits: torch.Tensor, |
| 239 | indices_type: torch.dtype | None, |
| 240 | *, |
| 241 | input_ids: torch.Tensor | None = None, |
| 242 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 243 | """Compute the actual routing logic. |
| 244 | |
| 245 | This method must be implemented by subclasses to provide the specific |
| 246 | routing algorithm (e.g., grouped_topk, fused_topk, custom routing, etc.). |
| 247 | |
| 248 | Args: |
| 249 | hidden_states: Input hidden states |
| 250 | router_logits: Router logits for expert selection |
| 251 | indices_type: Desired dtype for expert indices (may be None) |
| 252 | input_ids: Token ids, for routers that condition on them |
| 253 | |
| 254 | Returns: |
| 255 | tuple of (topk_weights, topk_ids) |
| 256 | |
| 257 | """ |
| 258 | raise NotImplementedError |
| 259 | |
| 260 | def _select_experts( |
| 261 | self, |
| 262 | hidden_states: torch.Tensor, |
| 263 | router_logits: torch.Tensor, |
| 264 | topk_indices_dtype: torch.dtype | None = None, |
| 265 | *, |
| 266 | input_ids: torch.Tensor | None = None, |
| 267 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 268 | """Route the input hidden states to the top-k experts based on the |
| 269 | router logits. |
| 270 | |
| 271 | This method implements the template method pattern: |
| 272 | 1. Validates EPLB state |
| 273 | 2. Calls _compute_routing() to get topk_weights and topk_ids |
| 274 | 3. Applies EPLB mapping if enabled |
| 275 | 4. Converts indices dtype if needed |
| 276 | |
| 277 | Returns: |
| 278 | (topk_weights, topk_ids) |
| 279 | (tuple[torch.Tensor, torch.Tensor]): |
| 280 | The weights and expert ids computation result. |
| 281 | |
| 282 | **Compatibility**: When EPLB is not enabled, the returned ids are |
| 283 | equivalent to global logical ids, so should be compatible with |
| 284 | plain MoE implementations without redundant experts. |
| 285 | |
| 286 | """ |
| 287 | # Step 1: Validate EPLB state |
| 288 | self._validate_eplb_state() |
| 289 | |
| 290 | # Step 2: Compute routing (delegated to subclass) |
| 291 | topk_weights, topk_ids = self._compute_routing( |
| 292 | hidden_states, router_logits, topk_indices_dtype, input_ids=input_ids |
| 293 | ) |
| 294 | |
| 295 | # Capture logical ids before EPLB mapping. |
| 296 | if self.capture_fn is not None: |
| 297 | self.capture_fn(topk_ids) |
| 298 | |
| 299 | # Step 3: Apply EPLB mapping |
| 300 | topk_ids = self._apply_eplb_mapping(topk_ids) |
| 301 | |
| 302 | # Step 4: Convert indices dtype |
| 303 | topk_ids = self._convert_indices_dtype(topk_ids, topk_indices_dtype) |
| 304 | |
| 305 | return topk_weights, topk_ids |