| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from abc import ABC, abstractmethod |
| 4 | from collections.abc import Callable |
| 5 | from dataclasses import dataclass |
| 6 | from enum import Enum |
| 7 | from math import prod |
| 8 | from typing import final |
| 9 | |
| 10 | import torch |
| 11 | |
| 12 | import vllm.envs as envs |
| 13 | from vllm.logger import init_logger |
| 14 | from vllm.model_executor.layers.fused_moe.activation import ( |
| 15 | ApplyMoEActivationConfig, |
| 16 | MoEActivation, |
| 17 | apply_moe_activation, |
| 18 | ) |
| 19 | from vllm.model_executor.layers.fused_moe.config import ( |
| 20 | FusedMoEConfig, |
| 21 | FusedMoEParallelConfig, |
| 22 | FusedMoEQuantConfig, |
| 23 | RoutingMethodType, |
| 24 | ) |
| 25 | from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput |
| 26 | from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( |
| 27 | SharedExperts, |
| 28 | SharedExpertsOrder, |
| 29 | ) |
| 30 | from vllm.model_executor.layers.fused_moe.utils import ( |
| 31 | _resize_cache, |
| 32 | ) |
| 33 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 34 | QuantKey, |
| 35 | ) |
| 36 | from vllm.platforms import current_platform |
| 37 | from vllm.utils.torch_utils import PIN_MEMORY |
| 38 | from vllm.v1.worker.ubatching import ( |
| 39 | dbo_enabled, |
| 40 | dbo_maybe_run_recv_hook, |
| 41 | dbo_register_recv_hook, |
| 42 | dbo_yield, |
| 43 | ) |
| 44 | from vllm.v1.worker.workspace import current_workspace_manager |
| 45 | |
| 46 | logger = init_logger(__name__) |
| 47 | |
| 48 | # |
| 49 | # This file defines a set of base classes used to make MoE kernels more modular. |
| 50 | # The goal is to be able to utilize different communication mechanisms with |
| 51 | # any fused MoE kernel without needing to have combinatoric implementations. |
| 52 | # |
| 53 | # The fused moe kernels are broken down into the following components: |
| 54 | # |
| 55 | # [Router] → [Quantize-Dispatch] → [Permute-Experts-Unpermute] → [Combine] |
| 56 | # |
| 57 | # Each component will be independent of (but may inform) the others except for |
| 58 | # [Quantize-Dispatch] and `[Combine] (see below). The components can then be |
| 59 | # mixed and matched with so that DP+EP can be supported easily for multiple |
| 60 | # MoE kernel implementations. |
| 61 | # |
| 62 | # The following main classes are defined: |
| 63 | # * FusedMoEPrepareAndFinalizeModular - an abstract base class for preparation of MoE |
| 64 | # inputs (e.g. quantization, distribution) and finalization of Moe outputs. |
| 65 | # The prepare method must take care of any needed quantization and the |
| 66 | # finalize method, informed by the FusedMoEExpertsModular method, |
| 67 | # may apply weights and/or do the final reduction of the output. |
| 68 | # * FusedMoEExpertsModular - an abstract base class for the main fused |
| 69 | # MoE operation, i.e matmul + act_mul + optionally quant + matmul. |
| 70 | # Some FusedMoEExpertsModular implementations may choose to do |
| 71 | # the weight application and/or reduction. The class communicates this |
| 72 | # to [Finalize] via a TopKWeightAndReduce object. |
| 73 | # * FusedMoEModularKernel - an interface class that combines a |
| 74 | # FusedMoEPrepareAndFinalizeModular and a FusedMoEExpertsModular to |
| 75 | # provide the standard fused MoE kernel interface. |
| 76 | # * TopKWeightAndReduce - A TopKWeightAndReduce implementation chosen |
| 77 | # by the FusedMoEExpertsModular implementation that is passed |
| 78 | # on to [Finalize]. |
| 79 | # |
| 80 | # [Quantize-Prepare] and [Finalize] functionality are bundled into a single |
| 81 | # class `FusedMoEPrepareAndFinalizeModular` since they could use collective |
| 82 | # communication mechanisms that need to be consistent. |
| 83 | # |
| 84 | |
| 85 | |
| 86 | class FusedMoEActivationFormat(Enum): |
| 87 | """The standard activation format (num_tokens, hidden dim).""" |
| 88 | |
| 89 | Standard = ("standard",) |
| 90 | """ |
| 91 | The batched experts format (num experts, max tokens per expert, hidden dim) |
| 92 | """ |
| 93 | BatchedExperts = ("batched_experts",) |
| 94 | |
| 95 | |
| 96 | @dataclass |
| 97 | class ExpertTokensMetadata: |
| 98 | """Metadata regarding expert-token routing.""" |
| 99 | |
| 100 | expert_num_tokens: torch.Tensor | None |
| 101 | expert_num_tokens_cpu: torch.Tensor | None |
| 102 | psum_recv_per_rank: torch.Tensor | None = None |
| 103 | |
| 104 | @staticmethod |
| 105 | def make_from_list( |
| 106 | expert_num_tokens_list: list[int], device: str |
| 107 | ) -> "ExpertTokensMetadata": |
| 108 | expert_num_tokens_cpu = torch.tensor( |
| 109 | expert_num_tokens_list, |
| 110 | device="cpu", |
| 111 | dtype=torch.int32, |
| 112 | pin_memory=PIN_MEMORY, |
| 113 | ) |
| 114 | return ExpertTokensMetadata( |
| 115 | expert_num_tokens=expert_num_tokens_cpu.to(device, non_blocking=True), |
| 116 | expert_num_tokens_cpu=expert_num_tokens_cpu, |
| 117 | ) |
| 118 | |
| 119 | |
| 120 | class TopKWeightAndReduce(ABC): |
| 121 | """An abstract base class for weight application and reduction implementations.""" |
| 122 | |
| 123 | @abstractmethod |
| 124 | def apply( |
| 125 | self, |
| 126 | output: torch.Tensor | None, |
| 127 | fused_expert_output: torch.Tensor, |
| 128 | topk_weights: torch.Tensor, |
| 129 | topk_ids: torch.Tensor, |
| 130 | apply_router_weight_on_input: bool, |
| 131 | ) -> torch.Tensor: |
| 132 | """Apply topk_weights to the fused_experts_outputs and/or reduce. |
| 133 | If an output tensor is not passed, it will be created in the |
| 134 | function. |
| 135 | """ |
| 136 | raise NotImplementedError |
| 137 | |
| 138 | |
| 139 | # |
| 140 | # PrepareResultType is a tuple of: |
| 141 | # - quantized + dispatched a. |
| 142 | # - quantized + dispatched a1_scales. |
| 143 | # - Optional ExpertTokensMetadata containing gpu/cpu tensors |
| 144 | # as big as the number of local experts with the information about the |
| 145 | # number of tokens assigned to each local expert. |
| 146 | # - Optional dispatched expert topk IDs |
| 147 | # - Optional dispatched expert topk weight |
| 148 | # |
| 149 | # See `prepare` method below. |
| 150 | # |
| 151 | PrepareResultType = tuple[ |
| 152 | torch.Tensor, |
| 153 | torch.Tensor | None, |
| 154 | ExpertTokensMetadata | None, |
| 155 | torch.Tensor | None, |
| 156 | torch.Tensor | None, |
| 157 | ] |
| 158 | |
| 159 | # |
| 160 | # PrepareResultType is a tuple of: |
| 161 | # - quantized + dispatched a. |
| 162 | # - quantized + dispatched a1_scales. |
| 163 | # - dispatched router logits. |
| 164 | # |
| 165 | # See `prepare_monolithic` method below. |
| 166 | # |
| 167 | PrepareMonolithicResultType = tuple[ |
| 168 | torch.Tensor, |
| 169 | torch.Tensor | None, |
| 170 | torch.Tensor, |
| 171 | ] |
| 172 | |
| 173 | ReceiverType = Callable[[], PrepareResultType] |
| 174 | |
| 175 | ################################################################################ |
| 176 | # Prepare/Finalize |
| 177 | ################################################################################ |
| 178 | |
| 179 | |
| 180 | class FusedMoEPrepareAndFinalize(ABC): |
| 181 | """An abstract base class for the [Quantize-Prepare] and [Finalize] steps |
| 182 | described above. |
| 183 | |
| 184 | There are two variants of this class: |
| 185 | * FusedMoEPrepareAndFinalizeModular - this operates on topk ids and weights |
| 186 | * FusedMoEPrepareAndFinalizeMonolithic - the operates on router_logits |
| 187 | """ |
| 188 | |
| 189 | def post_init_setup(self, fused_experts: "FusedMoEExperts"): |
| 190 | """Initialize FusedMoEPrepareAndFinalizeModular settings that depend on |
| 191 | FusedMoEExpertsModular experts object. |
| 192 | The FusedMoEPrepareAndFinalizeModular implementations that have such |
| 193 | dependencies may choose to override this function. |
| 194 | """ |
| 195 | return |
| 196 | |
| 197 | @property |
| 198 | @abstractmethod |
| 199 | def activation_format(self) -> FusedMoEActivationFormat: |
| 200 | """A property indicating the output format of the activations for the |
| 201 | 'prepare' method. |
| 202 | """ |
| 203 | raise NotImplementedError |
| 204 | |
| 205 | @abstractmethod |
| 206 | def topk_indices_dtype(self) -> torch.dtype | None: |
| 207 | """The PrepareFinalize All2All implementations generally constrain the |
| 208 | dtype of the topk_ids they support. This function returns the |
| 209 | required topk indices dtype so it can be respected. |
| 210 | Return None if there are no such restrictions. |
| 211 | """ |
| 212 | raise NotImplementedError |
| 213 | |
| 214 | @abstractmethod |
| 215 | def max_num_tokens_per_rank(self) -> int | None: |
| 216 | """Some PrepareFinalize All2All implementations are batched. Meaning, |
| 217 | they can process only as set of tokens at a time. This |
| 218 | function returns the batch size i.e the maximum number of tokens |
| 219 | the implementation can process at a time. |
| 220 | Return None if there are no such restrictions. |
| 221 | """ |
| 222 | raise NotImplementedError |
| 223 | |
| 224 | @abstractmethod |
| 225 | def num_dispatchers(self) -> int: |
| 226 | raise NotImplementedError |
| 227 | |
| 228 | @abstractmethod |
| 229 | def output_is_reduced(self) -> bool: |
| 230 | """Indicates whether or not the output of finalize is reduced across all |
| 231 | ranks. |
| 232 | """ |
| 233 | raise NotImplementedError |
| 234 | |
| 235 | def supports_async(self) -> bool: |
| 236 | """Indicates whether or not this class implements prepare_async and |
| 237 | finalize_async. |
| 238 | """ |
| 239 | return False |
| 240 | |
| 241 | def supports_deferred_moe_finalize(self) -> bool: |
| 242 | """Whether ``finalize`` can be skipped for a deferring consumer. |
| 243 | |
| 244 | An implementation opts in only if everything it does in ``finalize`` |
| 245 | -- the top-k reduction, and any combine or reduce-scatter -- is work |
| 246 | the consumer takes over. |
| 247 | """ |
| 248 | return False |
| 249 | |
| 250 | |
| 251 | # TODO: pass FusedMoEParallelConfig in as ctor parameter? |
| 252 | class FusedMoEPrepareAndFinalizeModular(FusedMoEPrepareAndFinalize): |
| 253 | """An abstract base class for the [Quantize-Prepare] and [Finalize] steps |
| 254 | described above for the Modular case. |
| 255 | """ |
| 256 | |
| 257 | @abstractmethod |
| 258 | def prepare( |
| 259 | self, |
| 260 | a1: torch.Tensor, |
| 261 | topk_weights: torch.Tensor, |
| 262 | topk_ids: torch.Tensor, |
| 263 | num_experts: int, |
| 264 | expert_map: torch.Tensor | None, |
| 265 | apply_router_weight_on_input: bool, |
| 266 | quant_config: FusedMoEQuantConfig, |
| 267 | defer_input_quant: bool, |
| 268 | ) -> PrepareResultType: |
| 269 | """Perform any quantization (and/or) dispatching needed for this kernel. |
| 270 | - a1: The (unquantized) input to the MoE layer. |
| 271 | - topk_ids: The topk ids. |
| 272 | - topk_weights: The topk weights. |
| 273 | - num_experts: The total number of experts in the global expert space. |
| 274 | - expert_map: A tensor mapping expert indices from the global expert |
| 275 | space to the local expert space of the expert parallel shard. |
| 276 | - apply_router_weight_on_input: When True, apply the weights to the |
| 277 | activations, before quantization + dispatching. |
| 278 | - quant_config: Quantization info provided by the fused experts. |
| 279 | - defer_input_quant: Runtime parameter indicating whether or not to |
| 280 | defer input quantization to the FusedMoEExpertsModular |
| 281 | in cases where the compute kernel expects unquantized inputs |
| 282 | |
| 283 | Returns a tuple of: |
| 284 | - quantized + dispatched a. |
| 285 | - Optional quantized + dispatched a1_scales. |
| 286 | - Optional ExpertTokensMetadata containing gpu/cpu tensors |
| 287 | as big as the number of local experts with the information about the |
| 288 | number of tokens assigned to each local expert. |
| 289 | - Optional dispatched expert topk IDs |
| 290 | - Optional dispatched expert topk weight |
| 291 | """ |
| 292 | raise NotImplementedError |
| 293 | |
| 294 | def prepare_async( |
| 295 | self, |
| 296 | a1: torch.Tensor, |
| 297 | topk_weights: torch.Tensor, |
| 298 | topk_ids: torch.Tensor, |
| 299 | num_experts: int, |
| 300 | expert_map: torch.Tensor | None, |
| 301 | apply_router_weight_on_input: bool, |
| 302 | quant_config: FusedMoEQuantConfig, |
| 303 | defer_input_quant: bool, |
| 304 | ) -> tuple[Callable, ReceiverType] | ReceiverType: |
| 305 | """Perform any quantization (and/or) dispatching needed for this kernel |
| 306 | but do not wait for results from other workers. |
| 307 | - a1: The (unquantized) input to the MoE layer. |
| 308 | - a1_scale: Optional scales for a1 |
| 309 | - a2_scale: Optional scales for the second MoE gemm. Required to make |
| 310 | sure the quantization is consistent for both gemms. |
| 311 | - topk_ids: The topk ids. |
| 312 | - topk_weights: The topk weights. |
| 313 | - num_experts: The total number of experts in the global expert space. |
| 314 | - expert_map: A tensor mapping expert indices from the global expert |
| 315 | space to the local expert space of the expert parallel shard. |
| 316 | - apply_router_weight_on_input: When True, apply the weights to the |
| 317 | activations, before quantization + dispatching. |
| 318 | - defer_input_quant: Runtime parameter indicating whether or not to |
| 319 | defer input quantization to the FusedMoEExpertsModular |
| 320 | in cases where the compute kernel expects unquantized inputs |
| 321 | |
| 322 | Returns a callback or a hook callback pair that when invoked waits for |
| 323 | results from other workers and has the same return signature as |
| 324 | `prepare`, if a hook is returned this is more lightweight check that |
| 325 | the recv is complete without doing extra work (used by DBO, will be |
| 326 | refactored in the very near future) |
| 327 | |
| 328 | e.g. |
| 329 | |
| 330 | ret = obj.prepare_async(...) |
| 331 | |
| 332 | if isinstance(ret, tuple): |
| 333 | hook, receiver = ret |
| 334 | hook() |
| 335 | |
| 336 | if hook is not None: |
| 337 | a, a_scales, expert_meta, topk_ids, topk_weights = receiver() |
| 338 | |
| 339 | is equivalent to: |
| 340 | |
| 341 | a, a_scales, expert_meta, topk_ids, topk_weights = obj.prepare(...) |
| 342 | """ |
| 343 | raise NotImplementedError |
| 344 | |
| 345 | @abstractmethod |
| 346 | def finalize( |
| 347 | self, |
| 348 | output: torch.Tensor, |
| 349 | fused_expert_output: torch.Tensor, |
| 350 | topk_weights: torch.Tensor, |
| 351 | topk_ids: torch.Tensor, |
| 352 | apply_router_weight_on_input: bool, |
| 353 | weight_and_reduce_impl: TopKWeightAndReduce, |
| 354 | ) -> None: |
| 355 | """Perform any combine plus apply weights and perform a reduction on the |
| 356 | fused experts output. |
| 357 | - output: The output tensor, written in place. Must be (M, K) shape. |
| 358 | - fused_expert_output: The unweighted, unreduced output of the fused |
| 359 | experts, it will have (M, topk, K) shape. |
| 360 | - topk_weights: The weights to be applied to the fused_experts_output. |
| 361 | - topk_ids: The topk_ids. |
| 362 | - apply_router_weight_on_input: When False, apply the weights to |
| 363 | fused_expert_output. |
| 364 | - weight_and_reduce_impl: An optional TopKWeightAndReduce |
| 365 | implementation. |
| 366 | """ |
| 367 | raise NotImplementedError |
| 368 | |
| 369 | def finalize_async( |
| 370 | self, |
| 371 | output: torch.Tensor, |
| 372 | fused_expert_output: torch.Tensor, |
| 373 | topk_weights: torch.Tensor, |
| 374 | topk_ids: torch.Tensor, |
| 375 | apply_router_weight_on_input: bool, |
| 376 | weight_and_reduce_impl: TopKWeightAndReduce, |
| 377 | ) -> tuple[Callable, Callable] | Callable: |
| 378 | """Perform any combine plus apply weights and perform a reduction on the |
| 379 | fused experts output but do not wait for results from other workers. |
| 380 | - output: The output tensor, written in place. Must be (M, K) shape. |
| 381 | - fused_expert_output: The unweighted, unreduced output of the fused |
| 382 | experts, it will have (M, topk, K) shape. |
| 383 | - topk_weights: The weights to be applied to the fused_experts_output. |
| 384 | - topk_ids: The topk_ids. |
| 385 | - apply_router_weight_on_input: When False, apply the weights to |
| 386 | fused_expert_output. |
| 387 | - weight_and_reduce_impl: An optional TopKWeightAndReduce |
| 388 | implementation. |
| 389 | |
| 390 | Returns a callback or a hook callback pair that when invoked waits for |
| 391 | results from other workers and has the same return signature as |
| 392 | `finalize`, if a hook is returned this is more lightweight check that |
| 393 | the recv is complete without doing extra work (used by DBO, will be |
| 394 | refactored in the very near future) |
| 395 | |
| 396 | ret = obj.finalize_async(output, ...) |
| 397 | ... output not valid yet ... |
| 398 | if isinstance(ret, tuple): |
| 399 | hook, receiver = ret |
| 400 | hook() |
| 401 | receiver() |
| 402 | ... output valid here ... |
| 403 | |
| 404 | is equivalent to: |
| 405 | |
| 406 | obj.finalize(output, ...) |
| 407 | """ |
| 408 | raise NotImplementedError |
| 409 | |
| 410 | |
| 411 | class FusedMoEPrepareAndFinalizeMonolithic(FusedMoEPrepareAndFinalize): |
| 412 | """An abstract base class for the [Quantize-Prepare] and [Finalize] steps |
| 413 | described above for the monolithic case. |
| 414 | """ |
| 415 | |
| 416 | @abstractmethod |
| 417 | def prepare( |
| 418 | self, |
| 419 | a1: torch.Tensor, |
| 420 | router_logits: torch.Tensor, |
| 421 | quant_config: FusedMoEQuantConfig, |
| 422 | defer_input_quant: bool = False, |
| 423 | ) -> PrepareMonolithicResultType: |
| 424 | """Optional method for subclasses compatible with monolithic |
| 425 | FusedMoEExpertsModular kernels. |
| 426 | |
| 427 | Perform any quantization (and/or) dispatching needed for this kernel. |
| 428 | - a1: The (unquantized) input to the MoE layer. |
| 429 | - quant_config: Quantization info provided by the fused experts. |
| 430 | - defer_input_quant: Runtime parameter indicating whether or not to |
| 431 | defer input quantization to the FusedMoEExpertsModular |
| 432 | |
| 433 | Returns a tuple of: |
| 434 | - quantized + dispatched a. |
| 435 | - Optional quantized + dispatched a1_scales. |
| 436 | """ |
| 437 | raise NotImplementedError |
| 438 | |
| 439 | @abstractmethod |
| 440 | def finalize(self, fused_expert_output: torch.Tensor) -> torch.Tensor: |
| 441 | """Optional method for subclasses compatible with monolithic |
| 442 | FusedMoEExpertsModular kernels. |
| 443 | |
| 444 | Perform any combine plus apply weights and perform a reduction on the |
| 445 | fused experts output. |
| 446 | - fused_expert_output: The unweighted, unreduced output of the fused |
| 447 | experts, it will have (M, topk, K) shape. |
| 448 | """ |
| 449 | raise NotImplementedError |
| 450 | |
| 451 | |
| 452 | ################################################################################ |
| 453 | # Experts |
| 454 | ################################################################################ |
| 455 | |
| 456 | |
| 457 | # TODO: add supported activations method (return string) |
| 458 | class FusedMoEExperts(ABC): |
| 459 | # ROCm AITER kernels consume a 0/1 local-expert mask (with a trailing |
| 460 | # sentinel slot); every other backend consumes the canonical -1/local-slot |
| 461 | # expert_map. RoutedExperts.expert_map reads this flag to pick which to hand |
| 462 | # the active experts kernel. |
| 463 | consumes_expert_mask: bool = False |
| 464 | |
| 465 | def __init__( |
| 466 | self, |
| 467 | moe_config: FusedMoEConfig, |
| 468 | quant_config: FusedMoEQuantConfig, |
| 469 | max_num_tokens: int | None = None, |
| 470 | num_dispatchers: int | None = None, |
| 471 | ): |
| 472 | """moe_config: MoE layer configuration. |
| 473 | quant_config: Quantization parameters for this experts instance. |
| 474 | """ |
| 475 | if self.activation_format() == FusedMoEActivationFormat.Standard and ( |
| 476 | max_num_tokens is not None or num_dispatchers is not None |
| 477 | ): |
| 478 | raise ValueError( |
| 479 | "max_num_tokens and num_dispatchers should only be set for " |
| 480 | "BatchedExperts activation format." |
| 481 | ) |
| 482 | elif self.activation_format() == FusedMoEActivationFormat.BatchedExperts and ( |
| 483 | max_num_tokens is None or num_dispatchers is None |
| 484 | ): |
| 485 | raise ValueError( |
| 486 | "max_num_tokens and num_dispatchers must be set for " |
| 487 | "BatchedExperts activation format." |
| 488 | ) |
| 489 | |
| 490 | self.moe_config = moe_config |
| 491 | self.quant_config = quant_config |
| 492 | self.activation_config = ApplyMoEActivationConfig.from_configs( |
| 493 | moe_config, quant_config |
| 494 | ) |
| 495 | self.max_num_tokens = max_num_tokens |
| 496 | self.num_dispatchers = num_dispatchers |
| 497 | |
| 498 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # noqa: B027 |
| 499 | pass |
| 500 | |
| 501 | @staticmethod |
| 502 | def is_monolithic() -> bool: |
| 503 | raise NotImplementedError("Implemented by subclasses.") |
| 504 | |
| 505 | @property |
| 506 | def expects_unquantized_inputs(self) -> bool: |
| 507 | """Whether or not the PrepareFinalize should defer input quantization |
| 508 | in the prepare step. If True, then the Experts kernel will |
| 509 | execute the input quantization itself. |
| 510 | |
| 511 | Sample subclasses that override are AITER and FlashInfer CUTLASS. |
| 512 | """ |
| 513 | return False |
| 514 | |
| 515 | @staticmethod |
| 516 | @abstractmethod |
| 517 | def activation_format() -> FusedMoEActivationFormat: |
| 518 | """A property which is a tuple of the input and output activation formats |
| 519 | for the 'apply' method. |
| 520 | """ |
| 521 | raise NotImplementedError |
| 522 | |
| 523 | # |
| 524 | # Various helpers for registering support for various features. |
| 525 | # Used by the oracle to select a particular kernel for a deployment. |
| 526 | # |
| 527 | |
| 528 | @staticmethod |
| 529 | def is_supported_config( |
| 530 | cls: type["FusedMoEExperts"], |
| 531 | moe_config: FusedMoEConfig, |
| 532 | weight_key: QuantKey | None, |
| 533 | activation_key: QuantKey | None, |
| 534 | activation_format: FusedMoEActivationFormat, |
| 535 | ) -> tuple[bool, str | None]: |
| 536 | def _make_reason(reason: str) -> str: |
| 537 | return f"kernel does not support {reason}" |
| 538 | |
| 539 | if not cls._supports_current_device(): |
| 540 | return False, _make_reason(f"current device {current_platform.device_name}") |
| 541 | elif not (moe_config.is_act_and_mul or cls._supports_no_act_and_mul()): |
| 542 | return False, _make_reason("no act_and_mul MLP layer") |
| 543 | elif not cls._supports_activation(moe_config.activation): |
| 544 | return False, _make_reason(f"{moe_config.activation} activation") |
| 545 | elif not cls._supports_quant_scheme(weight_key, activation_key): |
| 546 | return False, _make_reason( |
| 547 | f"quantization scheme {weight_key}x{activation_key}" |
| 548 | ) |
| 549 | elif not cls._supports_parallel_config(moe_config.moe_parallel_config): |
| 550 | return False, _make_reason( |
| 551 | f"parallel config {moe_config.moe_parallel_config}" |
| 552 | ) |
| 553 | elif moe_config.has_hash_routing and cls.is_monolithic(): |
| 554 | return False, _make_reason("hash routing") |
| 555 | elif not cls._supports_routing_method( |
| 556 | moe_config.routing_method, weight_key, activation_key |
| 557 | ): |
| 558 | return False, _make_reason(f"routing method {moe_config.routing_method}") |
| 559 | elif not cls._supports_router_logits_dtype( |
| 560 | moe_config.router_logits_dtype, |
| 561 | moe_config.routing_method, |
| 562 | ): |
| 563 | return False, _make_reason( |
| 564 | f"router logits dtype {moe_config.router_logits_dtype}" |
| 565 | ) |
| 566 | elif not cls._supports_shape(moe_config.hidden_dim): |
| 567 | return False, _make_reason( |
| 568 | f"{moe_config.hidden_dim} hidden dim is not supported" |
| 569 | ) |
| 570 | elif activation_format != cls.activation_format(): |
| 571 | return False, _make_reason(f"{activation_format.value} activation format") |
| 572 | elif envs.VLLM_BATCH_INVARIANT and not cls._supports_batch_invariance(): |
| 573 | return False, _make_reason("batch invariance") |
| 574 | elif moe_config.is_lora_enabled and not cls.supports_lora(): |
| 575 | return False, _make_reason("LoRA") |
| 576 | return True, None |
| 577 | |
| 578 | @staticmethod |
| 579 | @abstractmethod |
| 580 | def _supports_current_device() -> bool: |
| 581 | """Whether the kernel supports the current device type |
| 582 | (compute cability and current platform). |
| 583 | """ |
| 584 | raise NotImplementedError |
| 585 | |
| 586 | @staticmethod |
| 587 | @abstractmethod |
| 588 | def _supports_no_act_and_mul() -> bool: |
| 589 | """Whether the kernel supports act_and_mul=False, i.e. |
| 590 | non-gated MoE models like Nemotron-Nano. |
| 591 | """ |
| 592 | raise NotImplementedError |
| 593 | |
| 594 | @staticmethod |
| 595 | @abstractmethod |
| 596 | def _supports_quant_scheme( |
| 597 | weight_key: QuantKey | None, |
| 598 | activation_key: QuantKey | None, |
| 599 | ) -> bool: |
| 600 | raise NotImplementedError |
| 601 | |
| 602 | @staticmethod |
| 603 | @abstractmethod |
| 604 | def _supports_activation(activation: MoEActivation) -> bool: |
| 605 | """Whether the kernel supports a particular act function.""" |
| 606 | raise NotImplementedError |
| 607 | |
| 608 | @staticmethod |
| 609 | @abstractmethod |
| 610 | def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: |
| 611 | """Whether the kernel supports deployment in particular parallel config. |
| 612 | |
| 613 | Can be overridden if a kernel does not support EP, SP or some other |
| 614 | configuration. |
| 615 | """ |
| 616 | raise NotImplementedError |
| 617 | |
| 618 | @staticmethod |
| 619 | def _supports_routing_method( |
| 620 | routing_method: RoutingMethodType, |
| 621 | weight_key: QuantKey | None, |
| 622 | activation_key: QuantKey | None, |
| 623 | ) -> bool: |
| 624 | """Whether the kernel supports a routing method (e.g. GroupedTopK). |
| 625 | |
| 626 | Can be overridden by monolithic kernels that execute the router |
| 627 | in addition to the experts if certain routers are not supported. |
| 628 | """ |
| 629 | return True |
| 630 | |
| 631 | @staticmethod |
| 632 | def _supports_router_logits_dtype( |
| 633 | router_logits_dtype: torch.dtype | None, |
| 634 | routing_method: RoutingMethodType, |
| 635 | ) -> bool: |
| 636 | """Whether a kernel supports a particular dtype for router logits input. |
| 637 | |
| 638 | Can be overridden by monolithic kernels that execute the router |
| 639 | in addition to the experts if certain dtypes are not supported. |
| 640 | """ |
| 641 | return True |
| 642 | |
| 643 | @staticmethod |
| 644 | def _supports_shape(hidden_dim: int) -> bool: |
| 645 | """Whether a kernel supports a particular shape. Can be overridden if a kernel |
| 646 | has specific shape requirements. |
| 647 | """ |
| 648 | return True |
| 649 | |
| 650 | @staticmethod |
| 651 | def _supports_batch_invariance() -> bool: |
| 652 | """Whether the kernel supports batch invariance, i.e. the output does not |
| 653 | depend on the order of the tokens in the input batch. This is useful |
| 654 | for determining if the kernel can used with VLLM_BATCH_INVARIANT=1. |
| 655 | """ |
| 656 | return False |
| 657 | |
| 658 | # |
| 659 | # Various helpers for accessing quantization parameters from the |
| 660 | # quant_config. |
| 661 | # |
| 662 | |
| 663 | @property |
| 664 | def quant_dtype(self) -> torch.dtype | str | None: |
| 665 | return self.quant_config.quant_dtype |
| 666 | |
| 667 | @property |
| 668 | def weight_quant_dtype(self) -> torch.dtype | str | None: |
| 669 | return self.quant_config.weight_quant_dtype |
| 670 | |
| 671 | @property |
| 672 | def block_shape(self) -> list[int] | None: |
| 673 | return self.quant_config.block_shape |
| 674 | |
| 675 | @property |
| 676 | def per_act_token_quant(self) -> bool: |
| 677 | return self.quant_config.per_act_token_quant |
| 678 | |
| 679 | @property |
| 680 | def per_out_ch_quant(self) -> bool: |
| 681 | return self.quant_config.per_out_ch_quant |
| 682 | |
| 683 | @property |
| 684 | def a1_scale(self) -> torch.Tensor | None: |
| 685 | return self.quant_config.a1_scale |
| 686 | |
| 687 | @property |
| 688 | def a2_scale(self) -> torch.Tensor | None: |
| 689 | return self.quant_config.a2_scale |
| 690 | |
| 691 | @property |
| 692 | def a1_gscale(self) -> torch.Tensor | None: |
| 693 | return self.quant_config.a1_gscale |
| 694 | |
| 695 | @property |
| 696 | def a2_gscale(self) -> torch.Tensor | None: |
| 697 | return self.quant_config.a2_gscale |
| 698 | |
| 699 | @property |
| 700 | def w1_scale(self) -> torch.Tensor | None: |
| 701 | return self.quant_config.w1_scale |
| 702 | |
| 703 | @property |
| 704 | def w2_scale(self) -> torch.Tensor | None: |
| 705 | return self.quant_config.w2_scale |
| 706 | |
| 707 | @property |
| 708 | def w1_zp(self) -> torch.Tensor | None: |
| 709 | return self.quant_config.w1_zp |
| 710 | |
| 711 | @property |
| 712 | def w2_zp(self) -> torch.Tensor | None: |
| 713 | return self.quant_config.w2_zp |
| 714 | |
| 715 | @property |
| 716 | def w1_bias(self) -> torch.Tensor | None: |
| 717 | return self.quant_config.w1_bias |
| 718 | |
| 719 | @property |
| 720 | def w2_bias(self) -> torch.Tensor | None: |
| 721 | return self.quant_config.w2_bias |
| 722 | |
| 723 | @property |
| 724 | def g1_alphas(self) -> torch.Tensor | None: |
| 725 | return self.quant_config.g1_alphas |
| 726 | |
| 727 | @property |
| 728 | def g2_alphas(self) -> torch.Tensor | None: |
| 729 | return self.quant_config.g2_alphas |
| 730 | |
| 731 | @staticmethod |
| 732 | def supports_lora() -> bool: |
| 733 | """Return True if this expert impl natively handles LoRA. |
| 734 | |
| 735 | LoRA-aware experts should mix in LoRAExpertsMixin, which flips this |
| 736 | to True and provides the per-forward LoRA state plumbing. |
| 737 | """ |
| 738 | return False |
| 739 | |
| 740 | def supports_packed_ue8m0_act_scales(self) -> bool: |
| 741 | """A flag indicating whether or not this class can process packed ue8m0 |
| 742 | activation scales. |
| 743 | """ |
| 744 | return False |
| 745 | |
| 746 | |
| 747 | class FusedMoEExpertsModular(FusedMoEExperts): |
| 748 | """An abstract base class for the [Permute-Experts-Unpermute] step described |
| 749 | above. |
| 750 | """ |
| 751 | |
| 752 | @staticmethod |
| 753 | def is_monolithic() -> bool: |
| 754 | return False |
| 755 | |
| 756 | def moe_problem_size( |
| 757 | self, |
| 758 | a1: torch.Tensor, |
| 759 | w1: torch.Tensor, |
| 760 | w2: torch.Tensor, |
| 761 | topk_ids: torch.Tensor, |
| 762 | ) -> tuple[int, int, int, int, int]: |
| 763 | """Extract the MoE problem size from the given tensor arguments: |
| 764 | - a: The hidden states, input to the MoE layer. |
| 765 | - w1: The first set of expert weights. |
| 766 | - w2: The second set of expert weights. |
| 767 | - topk_ids: The topk ids. |
| 768 | |
| 769 | Note: extracting the problem shape from the weight and activation |
| 770 | tensors is not obvious. It needs to be done this way specifically |
| 771 | due to subtle issues with particular kernels, e.g. the int4 kernels |
| 772 | divide the trailing dimension by two, so it's not "correct" to |
| 773 | extract N or K from the trailing dimension of w1 or w2. Similarly, |
| 774 | some kernels transpose the weights, so this needs to be kept in mind. |
| 775 | |
| 776 | Note: This implementation covers most cases. However, if experts |
| 777 | require a specialized implementation, like MarlinExperts, they are free |
| 778 | to override this function. |
| 779 | """ |
| 780 | assert len(w1.shape) == 3 and len(w2.shape) == 3 |
| 781 | E, N, _ = w1.shape |
| 782 | K = a1.size(-1) |
| 783 | |
| 784 | if a1.dim() == 2: |
| 785 | # Make sure we are using the correct a1 (pre-permute). |
| 786 | assert topk_ids.size(0) == a1.size(0), f"{topk_ids.size(0)} != {a1.size(0)}" |
| 787 | M = a1.size(0) |
| 788 | else: |
| 789 | assert a1.dim() == 3 |
| 790 | assert a1.size(0) == E, f"{a1.size(0)} == {E}" |
| 791 | M = a1.size(1) # This is max_num_tokens |
| 792 | |
| 793 | assert topk_ids.dim() == 2 |
| 794 | topk = topk_ids.size(1) |
| 795 | |
| 796 | return E, M, N, K, topk |
| 797 | |
| 798 | def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype: |
| 799 | """Workspace type: The dtype to use for the workspace tensors.""" |
| 800 | return act_dtype |
| 801 | |
| 802 | @abstractmethod |
| 803 | def workspace_shapes( |
| 804 | self, |
| 805 | M: int, |
| 806 | N: int, |
| 807 | K: int, |
| 808 | topk: int, |
| 809 | global_num_experts: int, |
| 810 | local_num_experts: int, |
| 811 | expert_tokens_meta: ExpertTokensMetadata | None, |
| 812 | activation: MoEActivation, |
| 813 | ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: |
| 814 | """Compute the shapes for the temporary and final outputs of the two gemms |
| 815 | and activation in the fused expert function. Since the gemms are |
| 816 | independent, the workspace for the first gemm can be shared with the |
| 817 | workspace for the last gemm. |
| 818 | |
| 819 | Inputs: |
| 820 | - M: number of tokens. |
| 821 | - N: Row (or column) dimension of expert weights. |
| 822 | - K: hidden dimension |
| 823 | - topk: The number of top-k experts to select. |
| 824 | - global_num_experts: global number of experts. |
| 825 | - local_num_experts: local number of experts due to DP/EP. |
| 826 | - expert_tokens_meta: number of tokens per expert metadata for batched |
| 827 | format. |
| 828 | |
| 829 | Returns a tuple of: |
| 830 | - workspace13 shape tuple: must be large enough to hold the |
| 831 | result of either expert gemm. |
| 832 | - workspace2 shape tuple: must be large enough to hold the |
| 833 | result of the activation function. |
| 834 | - output shape tuple: must be exact size of the final gemm output. |
| 835 | - Note: workspace shapes can be 0 if the workspace is not needed. |
| 836 | But in order for activation chunking to work, the first dimension |
| 837 | of each tuple must be the number of tokens when the shape is |
| 838 | not 0. |
| 839 | """ |
| 840 | raise NotImplementedError |
| 841 | |
| 842 | @staticmethod |
| 843 | def adjust_N_for_activation(N: int, activation: MoEActivation) -> int: |
| 844 | """Calculate the output dimension for the activation function. |
| 845 | |
| 846 | For *_no_mul activations (e.g. relu2_no_mul), |
| 847 | there's no gate/up split, so output size equals input size (N). |
| 848 | |
| 849 | For regular gated activations (e.g., silu, gelu, swigluoai), |
| 850 | output size is N // 2 due to gate × activation(up) multiplication. |
| 851 | |
| 852 | Args: |
| 853 | N: The intermediate size (width of w1/w3 weights). |
| 854 | activation: The activation function enum. |
| 855 | |
| 856 | Returns: |
| 857 | The output dimension after activation. |
| 858 | |
| 859 | """ |
| 860 | return N if not activation.is_gated else N // 2 |
| 861 | |
| 862 | def activation( |
| 863 | self, |
| 864 | activation: MoEActivation, |
| 865 | output: torch.Tensor, |
| 866 | input: torch.Tensor, |
| 867 | *, |
| 868 | topk_ids: torch.Tensor | None = None, |
| 869 | expert_map: torch.Tensor | None = None, |
| 870 | valid_token_counts: torch.Tensor | None = None, |
| 871 | ) -> None: |
| 872 | apply_moe_activation( |
| 873 | activation, |
| 874 | output, |
| 875 | input, |
| 876 | activation_config=self.activation_config, |
| 877 | topk_ids=topk_ids, |
| 878 | expert_map=expert_map, |
| 879 | valid_token_counts=valid_token_counts, |
| 880 | ) |
| 881 | |
| 882 | @abstractmethod |
| 883 | def finalize_weight_and_reduce_impl(self) -> TopKWeightAndReduce: |
| 884 | raise NotImplementedError |
| 885 | |
| 886 | @abstractmethod |
| 887 | def apply( |
| 888 | self, |
| 889 | output: torch.Tensor, |
| 890 | hidden_states: torch.Tensor, |
| 891 | w1: torch.Tensor, |
| 892 | w2: torch.Tensor, |
| 893 | topk_weights: torch.Tensor, |
| 894 | topk_ids: torch.Tensor, |
| 895 | activation: MoEActivation, |
| 896 | global_num_experts: int, |
| 897 | expert_map: torch.Tensor | None, |
| 898 | a1q_scale: torch.Tensor | None, |
| 899 | a2_scale: torch.Tensor | None, |
| 900 | workspace13: torch.Tensor, |
| 901 | workspace2: torch.Tensor, |
| 902 | expert_tokens_meta: ExpertTokensMetadata | None, |
| 903 | apply_router_weight_on_input: bool, |
| 904 | ) -> UnfinalizedMoEOutput | None: |
| 905 | """This function computes the intermediate result of a Mixture of Experts |
| 906 | (MoE) layer using two sets of weights, w1 and w2. |
| 907 | |
| 908 | Writes into `output` and returns None, unless the implementation stopped |
| 909 | after GEMM2 and left the top-k reduction to a fused consumer, in which |
| 910 | case `output` is untouched and the unfinalized result is returned. |
| 911 | |
| 912 | Args: |
| 913 | output: (torch.Tensor): The unweighted, unreduced output tensor. |
| 914 | hidden_states: (torch.Tensor): The (quantized) input tensor to the MoE |
| 915 | layer. |
| 916 | w1 (torch.Tensor): The first set of expert weights. |
| 917 | w2 (torch.Tensor): The second set of expert weights. |
| 918 | topk_weights: A map of row to expert weights. Some implementations |
| 919 | choose to do weight application. |
| 920 | topk_ids (torch.Tensor): A map of row to expert id. |
| 921 | activation (str): The activation function to apply after the first |
| 922 | MoE layer. |
| 923 | global_num_experts (int): The total number of experts in the global |
| 924 | expert space. |
| 925 | expert_map (Optional[torch.Tensor]): A tensor mapping expert indices |
| 926 | from the global expert space to the local expert space of the |
| 927 | expert parallel shard. |
| 928 | a1q_scale (Optional[torch.Tensor]): Optional quantized scale to be |
| 929 | used for a1. Result of quantization from prepare/finalize and not |
| 930 | from the FusedMoEQuantConfig. |
| 931 | a2_scale (Optional[torch.Tensor]): Optional quantized scale to be |
| 932 | used for the second gemm's activations. |
| 933 | workspace13 (torch.Tensor): A scratch tensor used for gemm outputs |
| 934 | must be large enough to hold output of either MoE gemm. |
| 935 | workspace2 (torch.Tensor): A scratch tensor used for the activation |
| 936 | function. |
| 937 | expert_tokens_meta (Optional[ExpertTokensMetadata]): An optional |
| 938 | ExpertTokensMetadata object containing gpu/cpu tensors |
| 939 | as big as the number of local experts with the information about |
| 940 | the number of tokens assigned to each local expert. |
| 941 | apply_router_weight_on_input: True if router weights are already |
| 942 | applied on the input. This is relevant if the implementation |
| 943 | chooses to do weight application. |
| 944 | |
| 945 | """ |
| 946 | raise NotImplementedError |
| 947 | |
| 948 | |
| 949 | class FusedMoEExpertsMonolithic(FusedMoEExperts): |
| 950 | """An abstract base class for the [Permute-Experts-Unpermute] step described |
| 951 | above, but with the monolithic interface (accepts router logits |
| 952 | rather than topk ids and weights). |
| 953 | """ |
| 954 | |
| 955 | @staticmethod |
| 956 | def _supports_routing_method( |
| 957 | routing_method: RoutingMethodType, |
| 958 | weight_key: QuantKey | None, |
| 959 | activation_key: QuantKey | None, |
| 960 | ) -> bool: |
| 961 | """Whether the kernel supports a routing method (e.g. GroupedTopK). |
| 962 | |
| 963 | Monolithic kernels should explicitly opt-in to support. |
| 964 | """ |
| 965 | raise NotImplementedError |
| 966 | |
| 967 | @staticmethod |
| 968 | def _supports_router_logits_dtype( |
| 969 | router_logits_dtype: torch.dtype | None, |
| 970 | routing_method: RoutingMethodType, |
| 971 | ) -> bool: |
| 972 | """Whether the kernel supports a dtype for router logits. |
| 973 | |
| 974 | Modular kernels should opt-in to support. |
| 975 | """ |
| 976 | raise NotImplementedError |
| 977 | |
| 978 | @staticmethod |
| 979 | def is_monolithic() -> bool: |
| 980 | return True |
| 981 | |
| 982 | routing_replay_capture_fn: Callable[[torch.Tensor], None] | None = None |
| 983 | _routing_replay_buffer: torch.Tensor | None = None |
| 984 | |
| 985 | def supports_routing_replay_capture(self) -> bool: |
| 986 | """Whether this expert supports routing replay capture. |
| 987 | |
| 988 | Subclasses backed by a kernel that exposes routed expert IDs |
| 989 | (e.g. FlashInfer's ``routing_replay_out``) should override. |
| 990 | """ |
| 991 | return False |
| 992 | |
| 993 | def set_capture_fn( |
| 994 | self, |
| 995 | capture_fn: Callable[[torch.Tensor], None] | None, |
| 996 | ) -> None: |
| 997 | self.routing_replay_capture_fn = capture_fn |
| 998 | if capture_fn is None: |
| 999 | self._routing_replay_buffer = None |
| 1000 | return |
| 1001 | # Allocate for per-rank batches gathered across the DP or EP group. |
| 1002 | dispatch_group_size = ( |
| 1003 | self.moe_config.ep_size |
| 1004 | if self.moe_config.use_ep |
| 1005 | else self.moe_config.dp_size |
| 1006 | ) |
| 1007 | max_num_replay_tokens = self.moe_config.max_num_tokens * dispatch_group_size |
| 1008 | self._routing_replay_buffer = torch.empty( |
| 1009 | (max_num_replay_tokens, self.moe_config.experts_per_token), |
| 1010 | dtype=torch.int16, |
| 1011 | device=self.moe_config.device, |
| 1012 | ) |
| 1013 | |
| 1014 | def _maybe_make_routing_replay_buffer( |
| 1015 | self, |
| 1016 | num_tokens: int, |
| 1017 | device: torch.device, |
| 1018 | ) -> torch.Tensor | None: |
| 1019 | if self.routing_replay_capture_fn is None: |
| 1020 | return None |
| 1021 | buf = self._routing_replay_buffer |
| 1022 | assert buf is not None |
| 1023 | if buf.shape[0] < num_tokens or buf.device != device: |
| 1024 | raise ValueError( |
| 1025 | "Routing replay buffer was initialized for " |
| 1026 | f"{buf.shape[0]} tokens on {buf.device}, but the kernel " |
| 1027 | f"received {num_tokens} tokens on {device}." |
| 1028 | ) |
| 1029 | return buf |
| 1030 | |
| 1031 | def _maybe_dispatch_routing_replay( |
| 1032 | self, |
| 1033 | routing_replay_out: torch.Tensor | None, |
| 1034 | num_tokens: int, |
| 1035 | ) -> None: |
| 1036 | if routing_replay_out is None or self.routing_replay_capture_fn is None: |
| 1037 | return |
| 1038 | self.routing_replay_capture_fn(routing_replay_out[:num_tokens]) |
| 1039 | |
| 1040 | def apply( |
| 1041 | self, |
| 1042 | hidden_states: torch.Tensor, |
| 1043 | w1: torch.Tensor, |
| 1044 | w2: torch.Tensor, |
| 1045 | router_logits: torch.Tensor, |
| 1046 | activation: MoEActivation, |
| 1047 | global_num_experts: int, |
| 1048 | expert_map: torch.Tensor | None, |
| 1049 | a1q_scale: torch.Tensor | None, |
| 1050 | apply_router_weight_on_input: bool, |
| 1051 | # grouped topk + fused topk bias parameters |
| 1052 | num_expert_group: int | None = None, |
| 1053 | e_score_correction_bias: torch.Tensor | None = None, |
| 1054 | routed_scaling_factor: float | None = None, |
| 1055 | topk_group: int | None = None, |
| 1056 | ) -> torch.Tensor | UnfinalizedMoEOutput: |
| 1057 | """Same as ``FusedMoEExperts.apply``, except uses router_logits as opposed |
| 1058 | to the topk_ids and topk_weights. This is useful for kernels |
| 1059 | with fused router and fused_experts (e.g. FLASHINFER_TRTLLM). |
| 1060 | """ |
| 1061 | raise NotImplementedError |
| 1062 | |
| 1063 | |
| 1064 | ################################################################################ |
| 1065 | # Kernel |
| 1066 | ################################################################################ |
| 1067 | |
| 1068 | |
| 1069 | @final |
| 1070 | class FusedMoEKernelModularImpl: |
| 1071 | def __init__( |
| 1072 | self, |
| 1073 | prepare_finalize: FusedMoEPrepareAndFinalizeModular, |
| 1074 | fused_experts: FusedMoEExpertsModular, |
| 1075 | ): |
| 1076 | self.prepare_finalize = prepare_finalize |
| 1077 | self.fused_experts = fused_experts |
| 1078 | self.shared_experts: SharedExperts | None = None |
| 1079 | moe_parallel_config = fused_experts.moe_config.moe_parallel_config |
| 1080 | self.moe_parallel_config = moe_parallel_config |
| 1081 | self.is_dp_ep = ( |
| 1082 | moe_parallel_config is not None |
| 1083 | and moe_parallel_config.dp_size > 1 |
| 1084 | and moe_parallel_config.use_ep |
| 1085 | ) |
| 1086 | |
| 1087 | def _allocate_buffers( |
| 1088 | self, |
| 1089 | out_dtype: torch.dtype, |
| 1090 | device: torch.device, |
| 1091 | M_chunk: int, |
| 1092 | M_full: int, |
| 1093 | N: int, |
| 1094 | K: int, |
| 1095 | top_k: int, |
| 1096 | global_num_experts: int, |
| 1097 | local_num_experts: int, |
| 1098 | expert_tokens_meta: ExpertTokensMetadata | None, |
| 1099 | activation: MoEActivation, |
| 1100 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 1101 | """Allocate temporary and output buffers for the fused experts op. |
| 1102 | Inputs: |
| 1103 | - out_dtype: output type of workspace and output tensors. |
| 1104 | - device: the device of the workspace and output tensors. |
| 1105 | See `workspace_shapes` for a description of the remainder of arguments. |
| 1106 | Returns a tuple of (workspace13, workspace2, output) tensors. |
| 1107 | """ |
| 1108 | assert M_full > 0 and M_chunk > 0 |
| 1109 | |
| 1110 | workspace_dtype = self.fused_experts.workspace_dtype(out_dtype) |
| 1111 | |
| 1112 | # Get intermediate workspace shapes based off the chunked M size. |
| 1113 | workspace13_shape, workspace2_shape, _ = self.fused_experts.workspace_shapes( |
| 1114 | M_chunk, |
| 1115 | N, |
| 1116 | K, |
| 1117 | top_k, |
| 1118 | global_num_experts, |
| 1119 | local_num_experts, |
| 1120 | expert_tokens_meta, |
| 1121 | activation, |
| 1122 | ) |
| 1123 | |
| 1124 | # Get final output shape based on the full M size. |
| 1125 | _, _, fused_out_shape = self.fused_experts.workspace_shapes( |
| 1126 | M_full, |
| 1127 | N, |
| 1128 | K, |
| 1129 | top_k, |
| 1130 | global_num_experts, |
| 1131 | local_num_experts, |
| 1132 | expert_tokens_meta, |
| 1133 | activation, |
| 1134 | ) |
| 1135 | |
| 1136 | # We can reuse the memory between cache1 and cache3 because by the |
| 1137 | # time we need cache3, we're done with cache1. |
| 1138 | # Reuse workspace13 for the output since there is only one chunk. |
| 1139 | max_shape_size = max(prod(workspace13_shape), prod(fused_out_shape)) |
| 1140 | |
| 1141 | if current_platform.is_cpu(): |
| 1142 | # Every CPU FusedMoEExpertsModular kernel reports zero-sized |
| 1143 | # workspace13/workspace2 (it manages its own scratch space |
| 1144 | # internally) and just needs the output buffer, so skip the |
| 1145 | # workspace manager entirely here -- it's the one part of this |
| 1146 | # call graph Dynamo can't trace (ContextVar-based lane lookup), |
| 1147 | # and CPU never actually needs its cross-chunk memory reuse. |
| 1148 | common_workspace = torch.empty( |
| 1149 | (max_shape_size,), dtype=workspace_dtype, device=device |
| 1150 | ) |
| 1151 | workspace2 = torch.empty( |
| 1152 | workspace2_shape, dtype=workspace_dtype, device=device |
| 1153 | ) |
| 1154 | workspace13 = _resize_cache(common_workspace, workspace13_shape) |
| 1155 | fused_out = _resize_cache(common_workspace, fused_out_shape) |
| 1156 | return workspace13, workspace2, fused_out |
| 1157 | |
| 1158 | common_workspace, workspace2 = current_workspace_manager().get_simultaneous( |
| 1159 | ((max_shape_size,), workspace_dtype), |
| 1160 | (workspace2_shape, workspace_dtype), |
| 1161 | ) |
| 1162 | workspace13 = _resize_cache(common_workspace, workspace13_shape) |
| 1163 | fused_out = _resize_cache(common_workspace, fused_out_shape) |
| 1164 | |
| 1165 | return workspace13, workspace2, fused_out |
| 1166 | |
| 1167 | def _maybe_apply_shared_experts( |
| 1168 | self, |
| 1169 | shared_experts: SharedExperts | None, |
| 1170 | shared_experts_input: torch.Tensor | None, |
| 1171 | ): |
| 1172 | if shared_experts is not None: |
| 1173 | assert self.prepare_finalize.supports_async() |
| 1174 | assert shared_experts_input is not None |
| 1175 | shared_experts( |
| 1176 | shared_experts_input, |
| 1177 | SharedExpertsOrder.MK_INTERNAL_OVERLAPPED, |
| 1178 | ) |
| 1179 | |
| 1180 | def _prepare( |
| 1181 | self, |
| 1182 | hidden_states: torch.Tensor, |
| 1183 | topk_weights: torch.Tensor, |
| 1184 | topk_ids: torch.Tensor, |
| 1185 | global_num_experts: int, |
| 1186 | expert_map: torch.Tensor | None, |
| 1187 | apply_router_weight_on_input: bool, |
| 1188 | ) -> tuple[ |
| 1189 | torch.Tensor, |
| 1190 | torch.Tensor | None, |
| 1191 | ExpertTokensMetadata | None, |
| 1192 | torch.Tensor, |
| 1193 | torch.Tensor, |
| 1194 | ]: |
| 1195 | """The _prepare method is a wrapper around self.prepare_finalize.prepare |
| 1196 | that handles DBO and async. |
| 1197 | """ |
| 1198 | if not self.prepare_finalize.supports_async(): |
| 1199 | # We shouldn't be running an a2a kernel that doesn't |
| 1200 | # support async prepare/finalize |
| 1201 | # TODO(lucas): enable in follow-up |
| 1202 | assert not dbo_enabled() |
| 1203 | |
| 1204 | ( |
| 1205 | a1q, |
| 1206 | a1q_scale, |
| 1207 | expert_tokens_meta, |
| 1208 | _expert_topk_ids, |
| 1209 | _expert_topk_weights, |
| 1210 | ) = self.prepare_finalize.prepare( |
| 1211 | hidden_states, |
| 1212 | topk_weights, |
| 1213 | topk_ids, |
| 1214 | global_num_experts, |
| 1215 | expert_map, |
| 1216 | apply_router_weight_on_input, |
| 1217 | self.fused_experts.quant_config, |
| 1218 | defer_input_quant=self.fused_experts.expects_unquantized_inputs, |
| 1219 | ) |
| 1220 | else: |
| 1221 | # Overlap shared expert compute with all2all dispatch. |
| 1222 | dbo_maybe_run_recv_hook() |
| 1223 | prepare_ret = self.prepare_finalize.prepare_async( |
| 1224 | hidden_states, |
| 1225 | topk_weights, |
| 1226 | topk_ids, |
| 1227 | global_num_experts, |
| 1228 | expert_map, |
| 1229 | apply_router_weight_on_input, |
| 1230 | self.fused_experts.quant_config, |
| 1231 | defer_input_quant=self.fused_experts.expects_unquantized_inputs, |
| 1232 | ) |
| 1233 | |
| 1234 | # TODO(lucas): refactor this in the alternative schedules followup |
| 1235 | # currently unpack if we have hook + receiver pair or just |
| 1236 | # receiver (see finalize_async docstring) |
| 1237 | hook, receiver = ( |
| 1238 | prepare_ret if isinstance(prepare_ret, tuple) else (None, prepare_ret) |
| 1239 | ) |
| 1240 | |
| 1241 | if hook is not None: |
| 1242 | if dbo_enabled(): |
| 1243 | # If DBO is being used, register the hook with the ubatch |
| 1244 | # context and call it in dbo_maybe_run_recv_hook instead of |
| 1245 | # passing it to the receiver. |
| 1246 | dbo_register_recv_hook(hook) |
| 1247 | dbo_yield() |
| 1248 | else: |
| 1249 | hook() |
| 1250 | |
| 1251 | ( |
| 1252 | a1q, |
| 1253 | a1q_scale, |
| 1254 | expert_tokens_meta, |
| 1255 | _expert_topk_ids, |
| 1256 | _expert_topk_weights, |
| 1257 | ) = receiver() |
| 1258 | |
| 1259 | # Maybe prepare gathered topk_ids and topk_weights from other EP ranks. |
| 1260 | topk_ids = topk_ids if _expert_topk_ids is None else _expert_topk_ids |
| 1261 | topk_weights = ( |
| 1262 | topk_weights if _expert_topk_weights is None else _expert_topk_weights |
| 1263 | ) |
| 1264 | |
| 1265 | return a1q, a1q_scale, expert_tokens_meta, topk_ids, topk_weights |
| 1266 | |
| 1267 | def _fused_experts( |
| 1268 | self, |
| 1269 | in_dtype: torch.dtype, |
| 1270 | a1q: torch.Tensor, |
| 1271 | a1q_scale: torch.Tensor | None, |
| 1272 | w1: torch.Tensor, |
| 1273 | w2: torch.Tensor, |
| 1274 | topk_weights: torch.Tensor, |
| 1275 | topk_ids: torch.Tensor, |
| 1276 | activation: MoEActivation, |
| 1277 | global_num_experts: int, |
| 1278 | local_num_experts: int, |
| 1279 | expert_map: torch.Tensor | None, |
| 1280 | apply_router_weight_on_input: bool, |
| 1281 | expert_tokens_meta: ExpertTokensMetadata | None, |
| 1282 | output_alias: torch.Tensor | None = None, |
| 1283 | ) -> torch.Tensor | UnfinalizedMoEOutput: |
| 1284 | _, M_full, N, K, top_k = self.fused_experts.moe_problem_size( |
| 1285 | a1q, w1, w2, topk_ids |
| 1286 | ) |
| 1287 | |
| 1288 | # This happens when none of the tokens from the all2all reach this |
| 1289 | # EP rank. Also, note that this is only relevant for CUDAGraph |
| 1290 | # incompatible all2all kernels like the DeepEP high-throughput |
| 1291 | # kernels. CUDAGraph compatible all2all kernels like the DeepEP |
| 1292 | # low-latency kernels are always batched and can never run into |
| 1293 | # the tensor.numel() == 0 case. |
| 1294 | if M_full == 0: |
| 1295 | return torch.empty_like(a1q, dtype=in_dtype) |
| 1296 | |
| 1297 | workspace13, workspace2, fused_out = self._allocate_buffers( |
| 1298 | in_dtype, |
| 1299 | a1q.device, |
| 1300 | M_full, |
| 1301 | M_full, |
| 1302 | N, |
| 1303 | K, |
| 1304 | top_k, |
| 1305 | global_num_experts, |
| 1306 | local_num_experts, |
| 1307 | expert_tokens_meta, |
| 1308 | activation, |
| 1309 | ) |
| 1310 | |
| 1311 | use_output_alias = ( |
| 1312 | output_alias is not None |
| 1313 | and output_alias.shape == fused_out.shape |
| 1314 | and output_alias.dtype == fused_out.dtype |
| 1315 | and output_alias.device == fused_out.device |
| 1316 | and output_alias.is_contiguous() |
| 1317 | ) |
| 1318 | |
| 1319 | # If caller's output buffer already matches fused_out shape/dtype, alias |
| 1320 | # to skip the redundant copy in TopKWeightAndReduceNoOP.apply downstream. |
| 1321 | # This eliminates ~94% of __amd_rocclr_copyBuffer events (Copy 2 of the |
| 1322 | # double-copy MoE write-back path). |
| 1323 | if current_platform.is_rocm(): |
| 1324 | from vllm._aiter_ops import rocm_aiter_ops |
| 1325 | |
| 1326 | if use_output_alias and rocm_aiter_ops.is_fused_moe_enabled(): |
| 1327 | fused_out = output_alias |
| 1328 | elif use_output_alias: |
| 1329 | fused_out = output_alias |
| 1330 | |
| 1331 | unfinalized = self.fused_experts.apply( |
| 1332 | output=fused_out, |
| 1333 | hidden_states=a1q, |
| 1334 | w1=w1, |
| 1335 | w2=w2, |
| 1336 | topk_weights=topk_weights, |
| 1337 | topk_ids=topk_ids, |
| 1338 | activation=activation, |
| 1339 | global_num_experts=global_num_experts, |
| 1340 | expert_map=expert_map, |
| 1341 | a1q_scale=a1q_scale, |
| 1342 | a2_scale=self.fused_experts.a2_scale, |
| 1343 | workspace13=workspace13, |
| 1344 | workspace2=workspace2, |
| 1345 | expert_tokens_meta=expert_tokens_meta, |
| 1346 | apply_router_weight_on_input=apply_router_weight_on_input, |
| 1347 | ) |
| 1348 | |
| 1349 | # Experts that stopped after GEMM2 wrote nothing into `fused_out`; the |
| 1350 | # top-k reduction they left open belongs to whoever takes this. |
| 1351 | if isinstance(unfinalized, UnfinalizedMoEOutput): |
| 1352 | return unfinalized |
| 1353 | |
| 1354 | return fused_out |
| 1355 | |
| 1356 | def _finalize( |
| 1357 | self, |
| 1358 | output: torch.Tensor, |
| 1359 | fused_out: torch.Tensor, |
| 1360 | hidden_states: torch.Tensor, |
| 1361 | topk_weights: torch.Tensor, |
| 1362 | topk_ids: torch.Tensor, |
| 1363 | apply_router_weight_on_input: bool, |
| 1364 | shared_experts: SharedExperts | None, |
| 1365 | shared_experts_input: torch.Tensor | None, |
| 1366 | ) -> torch.Tensor: |
| 1367 | """The _finalize method is a wrapper around self.prepare_finalize.finalize |
| 1368 | that handles DBO, async and shared expert overlap. |
| 1369 | |
| 1370 | Args: |
| 1371 | output: Tensor the finalized result is written into. |
| 1372 | fused_out: The unweighted, unreduced output of the fused experts. |
| 1373 | hidden_states: The input tensor to the MoE layer. |
| 1374 | topk_weights: A map of row to expert weights. |
| 1375 | topk_ids: A map of row to expert id. |
| 1376 | apply_router_weight_on_input: True if the router weights were |
| 1377 | already applied on the input. |
| 1378 | shared_experts: SharedExperts | None. The shared experts if any. |
| 1379 | shared_experts_input: Optional separate input for shared experts. |
| 1380 | When latent MoE is used, hidden_states is the latent-projected |
| 1381 | tensor (smaller dimension) used by routed experts, while |
| 1382 | shared_experts_input is the original hidden_states (full |
| 1383 | dimension) needed by the shared expert MLP. |
| 1384 | |
| 1385 | """ |
| 1386 | if not self.prepare_finalize.supports_async(): |
| 1387 | assert not dbo_enabled() |
| 1388 | |
| 1389 | self.prepare_finalize.finalize( |
| 1390 | output, |
| 1391 | fused_out, |
| 1392 | topk_weights, |
| 1393 | topk_ids, |
| 1394 | apply_router_weight_on_input, |
| 1395 | self.fused_experts.finalize_weight_and_reduce_impl(), |
| 1396 | ) |
| 1397 | else: |
| 1398 | finalize_ret = self.prepare_finalize.finalize_async( |
| 1399 | output, |
| 1400 | fused_out, |
| 1401 | topk_weights, |
| 1402 | topk_ids, |
| 1403 | apply_router_weight_on_input, |
| 1404 | self.fused_experts.finalize_weight_and_reduce_impl(), |
| 1405 | ) |
| 1406 | self._maybe_apply_shared_experts(shared_experts, shared_experts_input) |
| 1407 | |
| 1408 | # TODO(lucas): refactor this in the alternative schedules followup |
| 1409 | # currently unpack if we have hook + receiver pair or just |
| 1410 | # receiver (see finalize_async docstring) |
| 1411 | hook, receiver = ( |
| 1412 | finalize_ret |
| 1413 | if isinstance(finalize_ret, tuple) |
| 1414 | else (None, finalize_ret) |
| 1415 | ) |
| 1416 | |
| 1417 | if hook is not None: |
| 1418 | if dbo_enabled(): |
| 1419 | # If DBO is being used, register the hook with the ubatch |
| 1420 | # context and call it in dbo_maybe_run_recv_hook instead of |
| 1421 | # passing it to the receiver. |
| 1422 | dbo_register_recv_hook(hook) |
| 1423 | dbo_yield() |
| 1424 | else: |
| 1425 | hook() |
| 1426 | |
| 1427 | receiver() |
| 1428 | |
| 1429 | return output |
| 1430 | |
| 1431 | def apply( |
| 1432 | self, |
| 1433 | hidden_states: torch.Tensor, |
| 1434 | w1: torch.Tensor, |
| 1435 | w2: torch.Tensor, |
| 1436 | topk_ids: torch.Tensor, |
| 1437 | topk_weights: torch.Tensor, |
| 1438 | activation: MoEActivation = MoEActivation.SILU, |
| 1439 | global_num_experts: int = -1, |
| 1440 | expert_map: torch.Tensor | None = None, |
| 1441 | apply_router_weight_on_input: bool = False, |
| 1442 | shared_experts: SharedExperts | None = None, |
| 1443 | shared_experts_input: torch.Tensor | None = None, |
| 1444 | ) -> torch.Tensor | UnfinalizedMoEOutput: |
| 1445 | """This function computes a Mixture of Experts (MoE) layer using two sets |
| 1446 | of weights, w1 and w2, and top-k gating mechanism. |
| 1447 | |
| 1448 | Args: |
| 1449 | hidden_states: (torch.Tensor): The input tensor to the MoE layer. |
| 1450 | w1 (torch.Tensor): The first set of expert weights. |
| 1451 | w2 (torch.Tensor): The second set of expert weights. |
| 1452 | topk_ids (torch.Tensor): A map of row to expert id. |
| 1453 | topk_weights (torch.Tensor): The topk weights applied at the end of |
| 1454 | the layer. |
| 1455 | activation (MoEActivation): The activation function to apply after |
| 1456 | the first MoE layer. |
| 1457 | global_num_experts (int): The total number of experts in the global |
| 1458 | expert space. |
| 1459 | expert_map (Optional[torch.Tensor]): A tensor mapping expert indices |
| 1460 | from the global expert space to the local expert space of the |
| 1461 | expert parallel shard. |
| 1462 | apply_router_weight_on_input (bool): When true, the topk weights are |
| 1463 | applied directly on the inputs. This is only applicable when |
| 1464 | topk is 1. |
| 1465 | shared_experts: SharedExperts | None. The shared experts if any. |
| 1466 | shared_experts_input (Optional[torch.Tensor]): Optional separate |
| 1467 | input for shared experts. For latent MoE, this is the original |
| 1468 | hidden_states before latent projection. |
| 1469 | |
| 1470 | Returns: |
| 1471 | torch.Tensor: The output tensor after applying the MoE layer, or |
| 1472 | the unfinalized output when the experts left the top-k reduction to |
| 1473 | a fused consumer. |
| 1474 | |
| 1475 | """ |
| 1476 | output = torch.empty_like(hidden_states) |
| 1477 | |
| 1478 | local_num_experts = w1.shape[0] |
| 1479 | if global_num_experts == -1: |
| 1480 | global_num_experts = local_num_experts |
| 1481 | |
| 1482 | a1q, a1q_scale, expert_tokens_meta, topk_ids, topk_weights = self._prepare( |
| 1483 | hidden_states, |
| 1484 | topk_weights, |
| 1485 | topk_ids, |
| 1486 | global_num_experts, |
| 1487 | expert_map, |
| 1488 | apply_router_weight_on_input, |
| 1489 | ) |
| 1490 | |
| 1491 | # Stash the original unquantized hidden states on the LoRA context |
| 1492 | # so apply_w13_lora sees correct-magnitude activations instead of |
| 1493 | # the potentially quantized values produced by _prepare(). |
| 1494 | lora_ctx = getattr(self.fused_experts, "_lora_context", None) |
| 1495 | if lora_ctx is not None: |
| 1496 | lora_ctx.original_hidden_states = hidden_states |
| 1497 | |
| 1498 | fused_out = self._fused_experts( |
| 1499 | in_dtype=hidden_states.dtype, |
| 1500 | a1q=a1q, |
| 1501 | a1q_scale=a1q_scale, |
| 1502 | w1=w1, |
| 1503 | w2=w2, |
| 1504 | topk_weights=topk_weights, |
| 1505 | topk_ids=topk_ids, |
| 1506 | activation=activation, |
| 1507 | global_num_experts=global_num_experts, |
| 1508 | local_num_experts=local_num_experts, |
| 1509 | expert_map=expert_map, |
| 1510 | apply_router_weight_on_input=apply_router_weight_on_input, |
| 1511 | expert_tokens_meta=expert_tokens_meta, |
| 1512 | output_alias=output, |
| 1513 | ) |
| 1514 | |
| 1515 | if lora_ctx is not None: |
| 1516 | lora_ctx.original_hidden_states = None |
| 1517 | |
| 1518 | if isinstance(fused_out, UnfinalizedMoEOutput): |
| 1519 | # Nothing below can run on an unfinalized output: finalize is the |
| 1520 | # local top-k reduction (and, for DP/EP, the combine) that the |
| 1521 | # consumer takes over. |
| 1522 | if not self.prepare_finalize.supports_deferred_moe_finalize(): |
| 1523 | raise RuntimeError( |
| 1524 | f"{type(self.prepare_finalize).__name__} cannot pass through " |
| 1525 | "a deferred MoE output." |
| 1526 | ) |
| 1527 | return fused_out |
| 1528 | |
| 1529 | return self._finalize( |
| 1530 | output, |
| 1531 | fused_out, |
| 1532 | hidden_states, |
| 1533 | topk_weights, |
| 1534 | topk_ids, |
| 1535 | apply_router_weight_on_input, |
| 1536 | shared_experts=shared_experts, |
| 1537 | shared_experts_input=shared_experts_input, |
| 1538 | ) |
| 1539 | |
| 1540 | |
| 1541 | @final |
| 1542 | class FusedMoEKernelMonolithicImpl: |
| 1543 | def __init__( |
| 1544 | self, |
| 1545 | prepare_finalize: FusedMoEPrepareAndFinalizeMonolithic, |
| 1546 | fused_experts: FusedMoEExpertsMonolithic, |
| 1547 | ): |
| 1548 | self.prepare_finalize = prepare_finalize |
| 1549 | self.fused_experts = fused_experts |
| 1550 | |
| 1551 | def apply( |
| 1552 | self, |
| 1553 | hidden_states: torch.Tensor, |
| 1554 | w1: torch.Tensor, |
| 1555 | w2: torch.Tensor, |
| 1556 | router_logits: torch.Tensor, |
| 1557 | activation: MoEActivation, |
| 1558 | global_num_experts: int, |
| 1559 | expert_map: torch.Tensor | None, |
| 1560 | apply_router_weight_on_input: bool, |
| 1561 | # grouped topk + fused topk bias parameters |
| 1562 | num_expert_group: int | None = None, |
| 1563 | e_score_correction_bias: torch.Tensor | None = None, |
| 1564 | routed_scaling_factor: float | None = None, |
| 1565 | topk_group: int | None = None, |
| 1566 | ) -> torch.Tensor | UnfinalizedMoEOutput: |
| 1567 | """Same as forward(), except uses router_logits as opposed |
| 1568 | to the topk_ids and topk_weights. This is used for kernels |
| 1569 | that have fused router + experts (e.g. FLASHINFER_TRTLLM). |
| 1570 | """ |
| 1571 | a1q, a1q_scale, router_logits = self.prepare_finalize.prepare( |
| 1572 | hidden_states, |
| 1573 | router_logits=router_logits, |
| 1574 | quant_config=self.fused_experts.quant_config, |
| 1575 | defer_input_quant=self.fused_experts.expects_unquantized_inputs, |
| 1576 | ) |
| 1577 | |
| 1578 | fused_out = self.fused_experts.apply( |
| 1579 | hidden_states=a1q, |
| 1580 | w1=w1, |
| 1581 | w2=w2, |
| 1582 | router_logits=router_logits, |
| 1583 | activation=activation, |
| 1584 | global_num_experts=global_num_experts, |
| 1585 | expert_map=expert_map, |
| 1586 | apply_router_weight_on_input=apply_router_weight_on_input, |
| 1587 | a1q_scale=a1q_scale, |
| 1588 | # grouped topk + fused topk bias parameters |
| 1589 | num_expert_group=num_expert_group, |
| 1590 | e_score_correction_bias=e_score_correction_bias, |
| 1591 | routed_scaling_factor=routed_scaling_factor, |
| 1592 | topk_group=topk_group, |
| 1593 | ) |
| 1594 | |
| 1595 | if isinstance(fused_out, UnfinalizedMoEOutput): |
| 1596 | if not self.prepare_finalize.supports_deferred_moe_finalize(): |
| 1597 | raise RuntimeError( |
| 1598 | f"{type(self.prepare_finalize).__name__} cannot pass through " |
| 1599 | "a deferred MoE output." |
| 1600 | ) |
| 1601 | return fused_out |
| 1602 | return self.prepare_finalize.finalize(fused_out) |
| 1603 | |
| 1604 | |
| 1605 | @final |
| 1606 | class FusedMoEKernel: |
| 1607 | def __init__( |
| 1608 | self, |
| 1609 | prepare_finalize: FusedMoEPrepareAndFinalize, |
| 1610 | fused_experts: FusedMoEExperts, |
| 1611 | ): |
| 1612 | super().__init__() |
| 1613 | |
| 1614 | # Initialize the implementation (monolithic or modular). |
| 1615 | self.impl: FusedMoEKernelModularImpl | FusedMoEKernelMonolithicImpl |
| 1616 | if isinstance( |
| 1617 | prepare_finalize, FusedMoEPrepareAndFinalizeModular |
| 1618 | ) and isinstance(fused_experts, FusedMoEExpertsModular): |
| 1619 | self.impl = FusedMoEKernelModularImpl( |
| 1620 | prepare_finalize, |
| 1621 | fused_experts, |
| 1622 | ) |
| 1623 | |
| 1624 | elif isinstance( |
| 1625 | prepare_finalize, FusedMoEPrepareAndFinalizeMonolithic |
| 1626 | ) and isinstance(fused_experts, FusedMoEExpertsMonolithic): |
| 1627 | self.impl = FusedMoEKernelMonolithicImpl( |
| 1628 | prepare_finalize, |
| 1629 | fused_experts, |
| 1630 | ) |
| 1631 | |
| 1632 | else: |
| 1633 | raise ValueError( |
| 1634 | "prepare_finalize and fused_experts must both be either monolithic " |
| 1635 | f"or non-monolithic but got {prepare_finalize.__class__.__name__} " |
| 1636 | f"and {fused_experts.__class__.__name__}" |
| 1637 | ) |
| 1638 | |
| 1639 | self._post_init_setup() |
| 1640 | |
| 1641 | @property |
| 1642 | def can_overlap_shared_experts(self) -> bool: |
| 1643 | if isinstance(self.impl, FusedMoEKernelModularImpl): |
| 1644 | return self.impl.prepare_finalize.supports_async() |
| 1645 | else: |
| 1646 | return False |
| 1647 | |
| 1648 | @property |
| 1649 | def is_monolithic(self) -> bool: |
| 1650 | return isinstance(self.impl, FusedMoEKernelMonolithicImpl) |
| 1651 | |
| 1652 | @property |
| 1653 | def prepare_finalize(self) -> FusedMoEPrepareAndFinalize: |
| 1654 | return self.impl.prepare_finalize |
| 1655 | |
| 1656 | @property |
| 1657 | def fused_experts(self) -> FusedMoEExperts: |
| 1658 | return self.impl.fused_experts |
| 1659 | |
| 1660 | @property |
| 1661 | def moe_config(self) -> FusedMoEConfig: |
| 1662 | return self.fused_experts.moe_config |
| 1663 | |
| 1664 | def supports_lora(self) -> bool: |
| 1665 | return self.fused_experts.supports_lora() |
| 1666 | |
| 1667 | def _post_init_setup(self): |
| 1668 | """Resolve any leftover setup dependencies between self.prepare_finalize |
| 1669 | and self.fused_experts here. |
| 1670 | """ |
| 1671 | self.prepare_finalize.post_init_setup(self.impl.fused_experts) |
| 1672 | assert ( |
| 1673 | self.prepare_finalize.activation_format |
| 1674 | == self.fused_experts.activation_format() |
| 1675 | ) |
| 1676 | |
| 1677 | def output_is_reduced(self) -> bool: |
| 1678 | """Indicates whether or not the output of fused MoE kernel |
| 1679 | is reduced across all ranks. |
| 1680 | """ |
| 1681 | return self.prepare_finalize.output_is_reduced() |
| 1682 | |
| 1683 | def supports_deferred_moe_finalize(self) -> bool: |
| 1684 | return self.prepare_finalize.supports_deferred_moe_finalize() |
| 1685 | |
| 1686 | def apply_monolithic( |
| 1687 | self, |
| 1688 | hidden_states: torch.Tensor, |
| 1689 | w1: torch.Tensor, |
| 1690 | w2: torch.Tensor, |
| 1691 | router_logits: torch.Tensor, |
| 1692 | activation: MoEActivation, |
| 1693 | global_num_experts: int, |
| 1694 | expert_map: torch.Tensor | None, |
| 1695 | apply_router_weight_on_input: bool, |
| 1696 | # grouped topk + fused topk bias parameters |
| 1697 | num_expert_group: int | None = None, |
| 1698 | e_score_correction_bias: torch.Tensor | None = None, |
| 1699 | routed_scaling_factor: float | None = None, |
| 1700 | topk_group: int | None = None, |
| 1701 | ) -> torch.Tensor | UnfinalizedMoEOutput: |
| 1702 | assert isinstance(self.impl, FusedMoEKernelMonolithicImpl) |
| 1703 | return self.impl.apply( |
| 1704 | hidden_states=hidden_states, |
| 1705 | w1=w1, |
| 1706 | w2=w2, |
| 1707 | router_logits=router_logits, |
| 1708 | activation=activation, |
| 1709 | global_num_experts=global_num_experts, |
| 1710 | expert_map=expert_map, |
| 1711 | apply_router_weight_on_input=apply_router_weight_on_input, |
| 1712 | num_expert_group=num_expert_group, |
| 1713 | e_score_correction_bias=e_score_correction_bias, |
| 1714 | routed_scaling_factor=routed_scaling_factor, |
| 1715 | topk_group=topk_group, |
| 1716 | ) |
| 1717 | |
| 1718 | def apply( |
| 1719 | self, |
| 1720 | hidden_states: torch.Tensor, |
| 1721 | w1: torch.Tensor, |
| 1722 | w2: torch.Tensor, |
| 1723 | topk_weights: torch.Tensor, |
| 1724 | topk_ids: torch.Tensor, |
| 1725 | activation: MoEActivation, |
| 1726 | global_num_experts: int, |
| 1727 | expert_map: torch.Tensor | None, |
| 1728 | apply_router_weight_on_input: bool, |
| 1729 | shared_experts: SharedExperts | None = None, |
| 1730 | shared_experts_input: torch.Tensor | None = None, |
| 1731 | ) -> torch.Tensor | UnfinalizedMoEOutput: |
| 1732 | assert isinstance(self.impl, FusedMoEKernelModularImpl) |
| 1733 | return self.impl.apply( |
| 1734 | hidden_states=hidden_states, |
| 1735 | w1=w1, |
| 1736 | w2=w2, |
| 1737 | topk_weights=topk_weights, |
| 1738 | topk_ids=topk_ids, |
| 1739 | activation=activation, |
| 1740 | global_num_experts=global_num_experts, |
| 1741 | expert_map=expert_map, |
| 1742 | apply_router_weight_on_input=apply_router_weight_on_input, |
| 1743 | shared_experts=shared_experts, |
| 1744 | shared_experts_input=shared_experts_input, |
| 1745 | ) |