| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Output contract between a MoE layer and a consumer that fuses its tail.""" |
| 4 | |
| 5 | from dataclasses import dataclass |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | |
| 10 | @dataclass |
| 11 | class UnfinalizedMoEOutput: |
| 12 | """Unfinalized output of a MoE kernel. |
| 13 | |
| 14 | Kernels that can stop after GEMM2 (the TRTLLM-Gen ``do_finalize=False`` path) |
| 15 | hand back their permuted, unweighted output plus the routing weights and the |
| 16 | permute map, so that the top-k reduction can be fused with whatever follows |
| 17 | -- the shared-expert add and the tensor-parallel all-reduce -- instead of |
| 18 | running as its own kernel. |
| 19 | |
| 20 | The buffers are consumed as-is by the fused kernels, which index |
| 21 | ``gemm2_permuted`` by row: it must be densely packed at ``hidden_dim``, and |
| 22 | its row count (an autotuner-dependent padded value) is never referenced. |
| 23 | """ |
| 24 | |
| 25 | # [num_permuted_rows, hidden_dim] permuted, unweighted GEMM2 output. |
| 26 | gemm2_permuted: torch.Tensor |
| 27 | # [num_tokens, top_k] routing weights, in the activation dtype when the |
| 28 | # kernel routed itself and in the router's dtype for modular experts -- |
| 29 | # consumers index them as such, and a buffer typed wider than its contents |
| 30 | # reads as garbage. These already carry routed_scaling_factor for routing |
| 31 | # methods that fold it in. |
| 32 | expert_weights: torch.Tensor |
| 33 | # [num_tokens, top_k] int32 permute map; -1 = expert not local to this rank. |
| 34 | expanded_idx_to_permuted_idx: torch.Tensor |
| 35 | |
| 36 | |
| 37 | def convert_flashinfer_moe_output( |
| 38 | flashinfer_output: torch.Tensor | list[torch.Tensor], |
| 39 | *, |
| 40 | do_finalize: bool, |
| 41 | num_tokens: int, |
| 42 | top_k: int, |
| 43 | finalized_output: torch.Tensor | None = None, |
| 44 | ) -> torch.Tensor | UnfinalizedMoEOutput: |
| 45 | """Normalize the two FlashInfer TRTLLM MoE return layouts. |
| 46 | |
| 47 | Args: |
| 48 | flashinfer_output: Tensor returned by the FlashInfer BF16 wrapper's |
| 49 | legacy finalized path, or its mode-dependent tensor list. |
| 50 | do_finalize: Whether FlashInfer ran its top-k finalize step. |
| 51 | num_tokens: Number of input tokens. |
| 52 | top_k: Number of routed experts per token. |
| 53 | finalized_output: Optional destination passed to FlashInfer's ``output`` |
| 54 | argument. |
| 55 | |
| 56 | Returns: |
| 57 | A finalized tensor or the structured deferred-finalize output. |
| 58 | |
| 59 | Raises: |
| 60 | ValueError: If FlashInfer returns an unexpected layout. |
| 61 | |
| 62 | """ |
| 63 | if num_tokens < 0: |
| 64 | raise ValueError("num_tokens must be non-negative.") |
| 65 | if top_k <= 0: |
| 66 | raise ValueError("top_k must be positive.") |
| 67 | |
| 68 | if do_finalize: |
| 69 | if isinstance(flashinfer_output, torch.Tensor): |
| 70 | returned_output = flashinfer_output |
| 71 | elif len(flashinfer_output) == 1 and isinstance( |
| 72 | flashinfer_output[0], torch.Tensor |
| 73 | ): |
| 74 | returned_output = flashinfer_output[0] |
| 75 | else: |
| 76 | raise ValueError( |
| 77 | "Finalized FlashInfer MoE output must contain exactly one tensor." |
| 78 | ) |
| 79 | if finalized_output is not None and ( |
| 80 | returned_output.shape != finalized_output.shape |
| 81 | or returned_output.dtype != finalized_output.dtype |
| 82 | or returned_output.device != finalized_output.device |
| 83 | or returned_output.data_ptr() != finalized_output.data_ptr() |
| 84 | ): |
| 85 | raise ValueError( |
| 86 | "Finalized FlashInfer MoE output must alias the provided destination." |
| 87 | ) |
| 88 | return returned_output |
| 89 | |
| 90 | if finalized_output is not None: |
| 91 | raise ValueError("Deferred FlashInfer MoE output cannot have a destination.") |
| 92 | if isinstance(flashinfer_output, torch.Tensor) or len(flashinfer_output) != 3: |
| 93 | raise ValueError( |
| 94 | "Deferred FlashInfer MoE output must contain GEMM2 output, routing " |
| 95 | "weights, and a permutation map." |
| 96 | ) |
| 97 | gemm2_permuted, expert_weights, expanded_idx = flashinfer_output |
| 98 | if gemm2_permuted.ndim != 2: |
| 99 | raise ValueError("Deferred FlashInfer GEMM2 output must be rank-2.") |
| 100 | if expanded_idx.dtype != torch.int32: |
| 101 | raise ValueError("Deferred FlashInfer permutation map must use int32.") |
| 102 | expected_routes = num_tokens * top_k |
| 103 | if expert_weights.numel() != expected_routes: |
| 104 | raise ValueError( |
| 105 | "Deferred FlashInfer routing weights must contain " |
| 106 | f"num_tokens * top_k ({expected_routes}) elements." |
| 107 | ) |
| 108 | if expanded_idx.numel() != expected_routes: |
| 109 | raise ValueError( |
| 110 | "Deferred FlashInfer permutation map must contain " |
| 111 | f"num_tokens * top_k ({expected_routes}) elements." |
| 112 | ) |
| 113 | if not (gemm2_permuted.device == expert_weights.device == expanded_idx.device): |
| 114 | raise ValueError("Deferred FlashInfer outputs must share a device.") |
| 115 | if not all( |
| 116 | output.is_contiguous() |
| 117 | for output in (gemm2_permuted, expert_weights, expanded_idx) |
| 118 | ): |
| 119 | raise ValueError("Deferred FlashInfer outputs must be contiguous.") |
| 120 | return UnfinalizedMoEOutput( |
| 121 | gemm2_permuted=gemm2_permuted, |
| 122 | expert_weights=expert_weights.view(num_tokens, top_k), |
| 123 | expanded_idx_to_permuted_idx=expanded_idx.view(num_tokens, top_k), |
| 124 | ) |
| 125 | |
| 126 | |
| 127 | @dataclass |
| 128 | class MoEOutput: |
| 129 | """A MoE layer's output with its final reduction still open. |
| 130 | |
| 131 | Returned by layers whose MoE runs un-reduced (``reduce_results=False``) so |
| 132 | that the consumer -- typically the next layer's RMSNorm -- can fuse the |
| 133 | tensor-parallel all-reduce into itself instead of paying for a standalone |
| 134 | one. Keeping the shared-expert output and the routed scale separate leaves |
| 135 | that reduction the consumer's to schedule; when the routed output is still |
| 136 | unfinalized, the top-k reduction is open too and can fold into the same |
| 137 | kernel. |
| 138 | |
| 139 | Producers only leave the routed output unfinalized when a fused consumer can |
| 140 | actually take that form -- the token ceiling and topology support are theirs |
| 141 | to check -- so an ``UnfinalizedMoEOutput`` here means the fused path applies, |
| 142 | and a consumer need not re-derive that. |
| 143 | """ |
| 144 | |
| 145 | # Un-reduced routed output, either finalized or not. |
| 146 | routed: torch.Tensor | UnfinalizedMoEOutput |
| 147 | # Un-reduced shared-expert output, folded in by whoever finalizes. |
| 148 | shared_output: torch.Tensor | None = None |
| 149 | # Applied to the routed sum before the shared add. Already folded into the |
| 150 | # routing weights for methods that do so, in which case this is 1.0. |
| 151 | routed_scaling_factor: float = 1.0 |