| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Triton-based MoE expert implementations.""" |
| 4 | |
| 5 | import torch |
| 6 | |
| 7 | import vllm.model_executor.layers.fused_moe.modular_kernel as mk |
| 8 | from vllm import _custom_ops as ops |
| 9 | from vllm.model_executor.layers.fused_moe.activation import ( |
| 10 | MoEActivation, |
| 11 | apply_moe_activation_supported, |
| 12 | ) |
| 13 | from vllm.model_executor.layers.fused_moe.config import ( |
| 14 | FusedMoEConfig, |
| 15 | FusedMoEParallelConfig, |
| 16 | FusedMoEQuantConfig, |
| 17 | ) |
| 18 | from vllm.model_executor.layers.fused_moe.experts.lora_experts_mixin import ( |
| 19 | LoRAExpertsMixin, |
| 20 | ) |
| 21 | from vllm.model_executor.layers.fused_moe.fused_moe import ( |
| 22 | _prepare_expert_assignment, |
| 23 | invoke_fused_moe_triton_kernel, |
| 24 | invoke_fused_moe_wna16_triton_kernel, |
| 25 | try_get_optimal_moe_config, |
| 26 | ) |
| 27 | from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( |
| 28 | moe_align_block_size, |
| 29 | ) |
| 30 | from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( |
| 31 | TopKWeightAndReduceNoOP, |
| 32 | ) |
| 33 | from vllm.model_executor.layers.fused_moe.utils import ( |
| 34 | _resize_cache, |
| 35 | moe_kernel_quantize_input, |
| 36 | swiglu_limit_func, |
| 37 | ) |
| 38 | from vllm.model_executor.layers.quantization.utils.fp8_utils import ( |
| 39 | is_deep_gemm_e8m0_used, |
| 40 | ) |
| 41 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 42 | FP8_DTYPE, |
| 43 | QuantKey, |
| 44 | kFp8Dynamic128Sym, |
| 45 | kFp8DynamicTensorSym, |
| 46 | kFp8DynamicTokenSym, |
| 47 | kFp8Static128BlockSym, |
| 48 | kFp8StaticChannelSym, |
| 49 | kFp8StaticTensorSym, |
| 50 | kInt4Static, |
| 51 | kInt4Static32, |
| 52 | kInt4Static32Asym, |
| 53 | kInt4StaticAsym, |
| 54 | kInt8DynamicTensorSym, |
| 55 | kInt8DynamicTokenSym, |
| 56 | kInt8Static, |
| 57 | kInt8StaticChannelSym, |
| 58 | kInt8StaticTensorSym, |
| 59 | ) |
| 60 | from vllm.platforms import current_platform |
| 61 | from vllm.triton_utils import tl |
| 62 | from vllm.utils.multi_stream_utils import maybe_execute_in_parallel |
| 63 | |
| 64 | |
| 65 | class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular): |
| 66 | """Triton-based fused MoE expert implementation.""" |
| 67 | |
| 68 | @staticmethod |
| 69 | def is_supported_config( |
| 70 | cls: type[mk.FusedMoEExperts], |
| 71 | moe_config: FusedMoEConfig, |
| 72 | weight_key: QuantKey | None, |
| 73 | activation_key: QuantKey | None, |
| 74 | activation_format: mk.FusedMoEActivationFormat, |
| 75 | ) -> tuple[bool, str | None]: |
| 76 | supported, reason = mk.FusedMoEExperts.is_supported_config( |
| 77 | cls, |
| 78 | moe_config, |
| 79 | weight_key, |
| 80 | activation_key, |
| 81 | activation_format, |
| 82 | ) |
| 83 | if not supported or not current_platform.is_cuda_alike(): |
| 84 | return supported, reason |
| 85 | |
| 86 | padded_num_experts = (moe_config.num_experts + 31) // 32 * 32 |
| 87 | if padded_num_experts >= 1024: |
| 88 | return False, ( |
| 89 | "kernel's moe_align_block_size requires fewer than 1024 padded experts" |
| 90 | ) |
| 91 | return True, None |
| 92 | |
| 93 | def __init__( |
| 94 | self, |
| 95 | moe_config: FusedMoEConfig, |
| 96 | quant_config: FusedMoEQuantConfig, |
| 97 | ): |
| 98 | # Whether quantized MOE runs natively, or through |
| 99 | # higher-precision + activation QDQ. |
| 100 | self.quantization_emulation = False |
| 101 | super().__init__(moe_config, quant_config) |
| 102 | |
| 103 | @staticmethod |
| 104 | def activation_format() -> mk.FusedMoEActivationFormat: |
| 105 | return mk.FusedMoEActivationFormat.Standard |
| 106 | |
| 107 | @property |
| 108 | def expects_unquantized_inputs(self) -> bool: |
| 109 | # Defer activation quantization to apply() only when LoRA is active AND |
| 110 | # tokens are dispatched across ranks (DP+EP all2all). |
| 111 | return ( |
| 112 | self._lora_context is not None |
| 113 | and self.quant_dtype is not None |
| 114 | and self.moe_config.moe_parallel_config.use_all2all_kernels |
| 115 | ) |
| 116 | |
| 117 | @staticmethod |
| 118 | def _supports_current_device() -> bool: |
| 119 | return current_platform.is_cuda_alike() or current_platform.is_xpu() |
| 120 | |
| 121 | @staticmethod |
| 122 | def _supports_no_act_and_mul() -> bool: |
| 123 | return True |
| 124 | |
| 125 | @staticmethod |
| 126 | def _supports_quant_scheme( |
| 127 | weight_key: QuantKey | None, |
| 128 | activation_key: QuantKey | None, |
| 129 | ) -> bool: |
| 130 | # INT8 requires at least 7.5 (Turing) on CUDA. ROCm CDNA GPUs |
| 131 | # (e.g. MI2xx/MI3xx/gfx950) provide native INT8 matrix-core support and |
| 132 | # the Triton int8_w8a8 fused MoE kernel handles them. |
| 133 | # XPU reaches the same Triton kernel through the same launcher, and both |
| 134 | # int8 activation-quant paths work there: per-token uses the Triton |
| 135 | # `per_token_quant_int8`, per-tensor resolves to the XPU branch of |
| 136 | # `scaled_int8_quant`, which is plain elementwise arithmetic. |
| 137 | device_supports_int8 = ( |
| 138 | ( |
| 139 | current_platform.is_cuda() |
| 140 | and current_platform.has_device_capability((7, 5)) |
| 141 | ) |
| 142 | or current_platform.is_rocm() |
| 143 | or current_platform.is_xpu() |
| 144 | ) |
| 145 | |
| 146 | supported: list[tuple[QuantKey | None, QuantKey | None]] = [(None, None)] |
| 147 | if device_supports_int8: |
| 148 | # Activations are consumed as float and quantized to int8 |
| 149 | # dynamically inside the kernel, so only dynamic-activation int8 |
| 150 | # schemes are supported (static-activation int8 is not). |
| 151 | supported += [ |
| 152 | # per-channel weight + dynamic per-token activation |
| 153 | (kInt8StaticChannelSym, kInt8DynamicTokenSym), |
| 154 | # per-tensor weight + dynamic per-tensor activation |
| 155 | (kInt8StaticTensorSym, kInt8DynamicTensorSym), |
| 156 | ] |
| 157 | if current_platform.supports_fp8(): |
| 158 | supported += [ |
| 159 | (kFp8Static128BlockSym, kFp8Dynamic128Sym), |
| 160 | (kFp8StaticChannelSym, kFp8DynamicTokenSym), |
| 161 | (kFp8StaticTensorSym, kFp8DynamicTokenSym), |
| 162 | (kFp8StaticTensorSym, kFp8StaticTensorSym), |
| 163 | (kFp8StaticTensorSym, kFp8DynamicTensorSym), |
| 164 | ] |
| 165 | # Block-quantized FP8 with arbitrary block shape: the Triton |
| 166 | # kernels take the block shape as a runtime argument. |
| 167 | if ( |
| 168 | weight_key is not None |
| 169 | and activation_key == kFp8Dynamic128Sym |
| 170 | and weight_key.dtype == FP8_DTYPE |
| 171 | and weight_key.symmetric |
| 172 | and weight_key.scale.static |
| 173 | and weight_key.scale.dtype == torch.float32 |
| 174 | and weight_key.scale.group_shape.row > 1 |
| 175 | and weight_key.scale.group_shape.col > 1 |
| 176 | ): |
| 177 | return True |
| 178 | return (weight_key, activation_key) in supported |
| 179 | |
| 180 | @staticmethod |
| 181 | def _supports_activation(activation: MoEActivation) -> bool: |
| 182 | return apply_moe_activation_supported(activation) |
| 183 | |
| 184 | @staticmethod |
| 185 | def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: |
| 186 | return not ( |
| 187 | moe_parallel_config.use_fi_nvl_two_sided_kernels |
| 188 | or moe_parallel_config.use_fi_nvl_one_sided_kernels |
| 189 | ) |
| 190 | |
| 191 | @staticmethod |
| 192 | def _supports_batch_invariance(): |
| 193 | return True |
| 194 | |
| 195 | def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: |
| 196 | return TopKWeightAndReduceNoOP() |
| 197 | |
| 198 | def activation( |
| 199 | self, |
| 200 | activation: MoEActivation, |
| 201 | output: torch.Tensor, |
| 202 | input: torch.Tensor, |
| 203 | **kwargs, |
| 204 | ) -> None: |
| 205 | activation_config = self.activation_config |
| 206 | if ( |
| 207 | activation == MoEActivation.SILU |
| 208 | and activation_config.clamp_limit is not None |
| 209 | ): |
| 210 | swiglu_limit_func(output, input, activation_config.clamp_limit) |
| 211 | return |
| 212 | |
| 213 | # SWIGLUOAI_UNINTERLEAVE routes to torch.ops._C.silu_and_mul_with_clamp |
| 214 | # via apply_moe_activation() and requires clamp_limit to be set. |
| 215 | if activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE: |
| 216 | assert activation_config.clamp_limit is not None, ( |
| 217 | "SWIGLUOAI_UNINTERLEAVE requires gemm1_clamp_limit" |
| 218 | ) |
| 219 | |
| 220 | super().activation( |
| 221 | activation, |
| 222 | output, |
| 223 | input, |
| 224 | ) |
| 225 | |
| 226 | def workspace_shapes( |
| 227 | self, |
| 228 | M: int, |
| 229 | N: int, |
| 230 | K: int, |
| 231 | topk: int, |
| 232 | global_num_experts: int, |
| 233 | local_num_experts: int, |
| 234 | expert_tokens_meta: mk.ExpertTokensMetadata | None, |
| 235 | activation: MoEActivation, |
| 236 | ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: |
| 237 | activation_out_dim = self.adjust_N_for_activation(N, activation) |
| 238 | workspace1 = (M, topk, max(activation_out_dim, K)) |
| 239 | workspace2 = (M, topk, max(N, K)) |
| 240 | output = (M, K) |
| 241 | return (workspace1, workspace2, output) |
| 242 | |
| 243 | def apply( |
| 244 | self, |
| 245 | output: torch.Tensor, |
| 246 | hidden_states: torch.Tensor, |
| 247 | w1: torch.Tensor, |
| 248 | w2: torch.Tensor, |
| 249 | topk_weights: torch.Tensor, |
| 250 | topk_ids: torch.Tensor, |
| 251 | activation: MoEActivation, |
| 252 | global_num_experts: int, |
| 253 | expert_map: torch.Tensor | None, |
| 254 | a1q_scale: torch.Tensor | None, |
| 255 | a2_scale: torch.Tensor | None, |
| 256 | workspace13: torch.Tensor, |
| 257 | workspace2: torch.Tensor, |
| 258 | expert_tokens_meta: mk.ExpertTokensMetadata | None, |
| 259 | apply_router_weight_on_input: bool, |
| 260 | ): |
| 261 | # Check constraints. |
| 262 | if self.quant_config.use_int4_w4a16: |
| 263 | assert hidden_states.size(-1) // 2 == w1.size(2), "Hidden size mismatch" |
| 264 | else: |
| 265 | assert hidden_states.size(-1) == w1.size(2), ( |
| 266 | f"Hidden size mismatch {hidden_states.size(-1)} != {w1.size(2)}" |
| 267 | ) |
| 268 | |
| 269 | assert hidden_states.is_contiguous(), "Hidden_states must be contiguous" |
| 270 | assert hidden_states.dim() == 2 |
| 271 | assert w1.stride(-1) == 1, "Stride of last dimension must be 1" |
| 272 | assert w2.stride(-1) == 1, "Stride of last dimension must be 1" |
| 273 | assert hidden_states.dtype in [ |
| 274 | torch.float32, |
| 275 | torch.float16, |
| 276 | torch.bfloat16, |
| 277 | torch.float8_e4m3fn, |
| 278 | torch.float8_e4m3fnuz, |
| 279 | torch.int8, |
| 280 | ] |
| 281 | |
| 282 | # We declared expects_unquantized_inputs (LoRA + DP/EP all2all), so the |
| 283 | # prepare step deferred activation quantization to this kernel: |
| 284 | # `hidden_states` arrives unquantized. Keep the unquantized tensor for |
| 285 | # the LoRA shrink input and quantize a copy here for the base GEMM |
| 286 | # (mirrors what the prepare step would have done, but after the |
| 287 | # all-gather so the layout matches the gathered topk_ids / token map). |
| 288 | lora_unquantized_hidden_states: torch.Tensor | None = None |
| 289 | if self.expects_unquantized_inputs: |
| 290 | assert a1q_scale is None |
| 291 | lora_unquantized_hidden_states = hidden_states |
| 292 | hidden_states, a1q_scale = moe_kernel_quantize_input( |
| 293 | hidden_states, |
| 294 | self.a1_scale, |
| 295 | self.quant_dtype, |
| 296 | self.per_act_token_quant, |
| 297 | self.block_shape, |
| 298 | quantization_emulation=self.quantization_emulation, |
| 299 | ) |
| 300 | |
| 301 | E, num_tokens, N, K, top_k_num = self.moe_problem_size( |
| 302 | hidden_states, w1, w2, topk_ids |
| 303 | ) |
| 304 | |
| 305 | if global_num_experts == -1: |
| 306 | global_num_experts = E |
| 307 | |
| 308 | config = try_get_optimal_moe_config( |
| 309 | w1.size(), |
| 310 | w2.size(), |
| 311 | top_k_num, |
| 312 | self.quant_config.config_name(hidden_states.dtype), |
| 313 | num_tokens, |
| 314 | block_shape=self.block_shape, |
| 315 | ) |
| 316 | |
| 317 | if hidden_states.dtype == torch.bfloat16: |
| 318 | compute_type = tl.bfloat16 |
| 319 | elif hidden_states.dtype == torch.float16: |
| 320 | compute_type = tl.float16 |
| 321 | elif hidden_states.dtype == torch.float32: |
| 322 | compute_type = tl.float32 |
| 323 | elif ( |
| 324 | hidden_states.dtype == torch.float8_e4m3fn |
| 325 | or hidden_states.dtype == torch.float8_e4m3fnuz |
| 326 | or hidden_states.dtype == torch.int8 |
| 327 | ): |
| 328 | compute_type = tl.bfloat16 |
| 329 | else: |
| 330 | raise ValueError(f"Unsupported compute_type: {hidden_states.dtype}") |
| 331 | |
| 332 | # Note that the output tensor might be in workspace1 |
| 333 | intermediate_cache1 = _resize_cache(workspace2, (num_tokens, top_k_num, N)) |
| 334 | cache2_dim = self.adjust_N_for_activation(N, activation) |
| 335 | intermediate_cache2 = _resize_cache( |
| 336 | workspace13, (num_tokens * top_k_num, cache2_dim) |
| 337 | ) |
| 338 | intermediate_cache3 = _resize_cache(workspace2, (num_tokens, top_k_num, K)) |
| 339 | |
| 340 | # Under EP, drop the top-k slots routed to remote experts at alignment |
| 341 | # time instead of launching `off_experts == -1` GEMM blocks that only |
| 342 | # write zeros; the pad-aware `moe_sum` below skips the same slots. |
| 343 | # Their rows in the workspace caches are then never written, so every |
| 344 | # consumer of those caches must be row-local: a dynamic per-tensor |
| 345 | # scale for the second GEMM's activation would take its amax over the |
| 346 | # untouched rows, and LoRA aligns the full `topk_ids` itself. |
| 347 | a2_scale_is_global = ( |
| 348 | self.quant_dtype is not None |
| 349 | and a2_scale is None |
| 350 | and self.block_shape is None |
| 351 | and (not self.per_act_token_quant or self.quantization_emulation) |
| 352 | ) |
| 353 | skip_invalid = ( |
| 354 | expert_map is not None |
| 355 | and self._lora_context is None |
| 356 | and not a2_scale_is_global |
| 357 | ) |
| 358 | |
| 359 | # Include fused shared-expert rows while preserving EP remapping. |
| 360 | num_align_experts = w1.shape[0] if expert_map is None else global_num_experts |
| 361 | sorted_token_ids, expert_ids, num_tokens_post_padded = ( |
| 362 | _prepare_expert_assignment( |
| 363 | topk_ids, |
| 364 | config, |
| 365 | num_tokens, |
| 366 | top_k_num, |
| 367 | num_align_experts, |
| 368 | expert_map, |
| 369 | use_int8_w8a16=self.quant_config.use_int8_w8a16, |
| 370 | use_int4_w4a16=self.quant_config.use_int4_w4a16, |
| 371 | block_shape=self.block_shape, |
| 372 | ignore_invalid_experts=skip_invalid, |
| 373 | ) |
| 374 | ) |
| 375 | |
| 376 | # LoRA w13: applied to intermediate_cache1 before activation. When |
| 377 | # the LoRA layer requested a dual-stream schedule, we run base w13 |
| 378 | # GEMM on the default stream and the LoRA fast-path on aux_stream; |
| 379 | # the LoRA writes its delta into a fresh zero buffer (add_inputs= |
| 380 | # False) and we sum it into intermediate_cache1 after both finish. |
| 381 | # |
| 382 | # The LoRA shrink kernel needs unquantized, gathered-layout |
| 383 | # activations. When activation quant was deferred to this kernel |
| 384 | # (expects_unquantized_inputs), the input we quantized above is exactly |
| 385 | # that, so use it directly. Otherwise fall back to the context stash |
| 386 | # (e.g. weight-only quant), guarding on a row-count match so a |
| 387 | # DP-gathered layout never indexes a local stash out of bounds. |
| 388 | sorted_token_ids_lora = None |
| 389 | expert_ids_lora = None |
| 390 | num_tokens_post_padded_lora = None |
| 391 | token_lora_mapping = None |
| 392 | lora_context = self._lora_context |
| 393 | if lora_unquantized_hidden_states is not None: |
| 394 | lora_x = lora_unquantized_hidden_states |
| 395 | elif ( |
| 396 | lora_context is not None |
| 397 | and lora_context.original_hidden_states is not None |
| 398 | and lora_context.original_hidden_states.shape[0] == hidden_states.shape[0] |
| 399 | ): |
| 400 | lora_x = lora_context.original_hidden_states |
| 401 | else: |
| 402 | lora_x = hidden_states |
| 403 | |
| 404 | # TODO: The fallback to self.a1_scale was added for deferred static |
| 405 | # activation quantization in https://github.com/vllm-project/vllm/pull/40857. |
| 406 | # Activation emulation relies solely on `a1q_scale` output of |
| 407 | # `moe_kernel_quantize_input` - this should be adapted to |
| 408 | # always solely rely on `a1q_scale`. |
| 409 | input_scale = ( |
| 410 | a1q_scale |
| 411 | if self.quantization_emulation |
| 412 | else (a1q_scale if a1q_scale is not None else self.a1_scale) |
| 413 | ) |
| 414 | |
| 415 | def _base_w13_fn(): |
| 416 | invoke_fused_moe_triton_kernel( |
| 417 | hidden_states, |
| 418 | w1, |
| 419 | intermediate_cache1, |
| 420 | input_scale, |
| 421 | self.w1_scale, |
| 422 | None, # topk_weights |
| 423 | sorted_token_ids, |
| 424 | expert_ids, |
| 425 | num_tokens_post_padded, |
| 426 | False, # mul_routed_weights |
| 427 | top_k_num, |
| 428 | config, |
| 429 | compute_type=compute_type, |
| 430 | use_fp8_w8a8=self.quant_config.use_fp8_w8a8, |
| 431 | use_int8_w8a8=self.quant_config.use_int8_w8a8, |
| 432 | use_int8_w8a16=self.quant_config.use_int8_w8a16, |
| 433 | use_int4_w4a16=self.quant_config.use_int4_w4a16, |
| 434 | per_channel_quant=self.per_act_token_quant, |
| 435 | block_shape=self.block_shape, |
| 436 | B_bias=self.w1_bias, |
| 437 | ) |
| 438 | |
| 439 | if lora_context is not None and lora_context.aux_stream is not None: |
| 440 | # add_inputs=False: kernel overwrites lora_delta_w13. zeros (not |
| 441 | # empty) so untouched rows -- e.g. blocks where every program |
| 442 | # early-exits because lora_id<0 -- stay at zero and the trailing |
| 443 | # add_() is a no-op there. |
| 444 | lora_delta_w13 = torch.zeros_like(intermediate_cache1) |
| 445 | |
| 446 | def _lora_w13_fn(): |
| 447 | return self.apply_w13_lora( |
| 448 | lora_context, |
| 449 | y=lora_delta_w13, |
| 450 | x=lora_x, |
| 451 | topk_ids=topk_ids, |
| 452 | topk_weights=topk_weights, |
| 453 | expert_map=expert_map, |
| 454 | w1=w1, |
| 455 | w2=w2, |
| 456 | num_tokens=num_tokens, |
| 457 | top_k_num=top_k_num, |
| 458 | add_inputs=False, |
| 459 | ) |
| 460 | |
| 461 | assert lora_context.events is not None |
| 462 | _, lora_meta = maybe_execute_in_parallel( |
| 463 | _base_w13_fn, |
| 464 | _lora_w13_fn, |
| 465 | lora_context.events[0], |
| 466 | lora_context.events[1], |
| 467 | lora_context.aux_stream, |
| 468 | ) |
| 469 | ( |
| 470 | sorted_token_ids_lora, |
| 471 | expert_ids_lora, |
| 472 | num_tokens_post_padded_lora, |
| 473 | token_lora_mapping, |
| 474 | ) = lora_meta |
| 475 | intermediate_cache1.add_(lora_delta_w13) |
| 476 | else: |
| 477 | _base_w13_fn() |
| 478 | if lora_context is not None: |
| 479 | ( |
| 480 | sorted_token_ids_lora, |
| 481 | expert_ids_lora, |
| 482 | num_tokens_post_padded_lora, |
| 483 | token_lora_mapping, |
| 484 | ) = self.apply_w13_lora( |
| 485 | lora_context, |
| 486 | y=intermediate_cache1, |
| 487 | x=lora_x, |
| 488 | topk_ids=topk_ids, |
| 489 | topk_weights=topk_weights, |
| 490 | expert_map=expert_map, |
| 491 | w1=w1, |
| 492 | w2=w2, |
| 493 | num_tokens=num_tokens, |
| 494 | top_k_num=top_k_num, |
| 495 | ) |
| 496 | |
| 497 | a2q_scale: torch.Tensor | None = None |
| 498 | |
| 499 | # Fuse SiLU+Mul + FP8 block quantize into a single kernel |
| 500 | # when conditions permit (gated SiLU, fp8 block quant with |
| 501 | # group_size=128, no LoRA requiring the BF16 intermediate). |
| 502 | # The fused kernel has no clamp parameter, so a configured |
| 503 | # SwiGLU clamp limit falls through to the unfused path. |
| 504 | if ( |
| 505 | activation == MoEActivation.SILU |
| 506 | and self.activation_config.clamp_limit is None |
| 507 | and self.quant_config.use_fp8_w8a8 |
| 508 | and self.block_shape == [128, 128] |
| 509 | and lora_context is None |
| 510 | and not is_deep_gemm_e8m0_used() |
| 511 | ): |
| 512 | qintermediate_cache2, a2q_scale = ops.silu_and_mul_per_block_quant( |
| 513 | intermediate_cache1.view(-1, N), |
| 514 | group_size=128, |
| 515 | quant_dtype=current_platform.fp8_dtype(), |
| 516 | ) |
| 517 | else: |
| 518 | self.activation( |
| 519 | activation, intermediate_cache2, intermediate_cache1.view(-1, N) |
| 520 | ) |
| 521 | |
| 522 | qintermediate_cache2, a2q_scale = moe_kernel_quantize_input( |
| 523 | intermediate_cache2, |
| 524 | a2_scale, |
| 525 | self.quant_dtype, |
| 526 | self.per_act_token_quant, |
| 527 | self.block_shape, |
| 528 | quantization_emulation=self.quantization_emulation, |
| 529 | ) |
| 530 | |
| 531 | # LoRA w2: applied to intermediate_cache3 before moe_sum, using the |
| 532 | # unquantized intermediate_cache2 as the lora_a input. Reuses the |
| 533 | # sorted_token_ids_lora computed above. Same dual-stream pattern as |
| 534 | # the w13 pair: base GEMM on default stream, LoRA delta on aux, |
| 535 | # join via .add_() into intermediate_cache3. |
| 536 | def _base_w2_fn(): |
| 537 | invoke_fused_moe_triton_kernel( |
| 538 | qintermediate_cache2, |
| 539 | w2, |
| 540 | intermediate_cache3, |
| 541 | a2q_scale, |
| 542 | self.w2_scale, |
| 543 | topk_weights, |
| 544 | sorted_token_ids, |
| 545 | expert_ids, |
| 546 | num_tokens_post_padded, |
| 547 | not apply_router_weight_on_input, |
| 548 | 1, |
| 549 | config, |
| 550 | compute_type=compute_type, |
| 551 | use_fp8_w8a8=self.quant_config.use_fp8_w8a8, |
| 552 | use_int8_w8a8=self.quant_config.use_int8_w8a8, |
| 553 | use_int8_w8a16=self.quant_config.use_int8_w8a16, |
| 554 | use_int4_w4a16=self.quant_config.use_int4_w4a16, |
| 555 | per_channel_quant=self.per_act_token_quant, |
| 556 | block_shape=self.block_shape, |
| 557 | B_bias=self.w2_bias, |
| 558 | ) |
| 559 | |
| 560 | if lora_context is not None and lora_context.aux_stream is not None: |
| 561 | lora_delta_w2 = torch.zeros_like(intermediate_cache3) |
| 562 | |
| 563 | def _lora_w2_fn(): |
| 564 | self.apply_w2_lora( |
| 565 | lora_context, |
| 566 | y=lora_delta_w2, |
| 567 | x=intermediate_cache2, |
| 568 | topk_weights=topk_weights, |
| 569 | sorted_token_ids_lora=sorted_token_ids_lora, |
| 570 | expert_ids_lora=expert_ids_lora, |
| 571 | num_tokens_post_padded_lora=num_tokens_post_padded_lora, |
| 572 | token_lora_mapping=token_lora_mapping, |
| 573 | num_tokens=num_tokens, |
| 574 | w1=w1, |
| 575 | w2=w2, |
| 576 | top_k_num=top_k_num, |
| 577 | add_inputs=False, |
| 578 | ) |
| 579 | |
| 580 | assert lora_context.events is not None |
| 581 | maybe_execute_in_parallel( |
| 582 | _base_w2_fn, |
| 583 | _lora_w2_fn, |
| 584 | lora_context.events[2], |
| 585 | lora_context.events[3], |
| 586 | lora_context.aux_stream, |
| 587 | ) |
| 588 | intermediate_cache3.add_(lora_delta_w2) |
| 589 | else: |
| 590 | _base_w2_fn() |
| 591 | if lora_context is not None: |
| 592 | self.apply_w2_lora( |
| 593 | lora_context, |
| 594 | y=intermediate_cache3, |
| 595 | x=intermediate_cache2, |
| 596 | topk_weights=topk_weights, |
| 597 | sorted_token_ids_lora=sorted_token_ids_lora, |
| 598 | expert_ids_lora=expert_ids_lora, |
| 599 | num_tokens_post_padded_lora=num_tokens_post_padded_lora, |
| 600 | token_lora_mapping=token_lora_mapping, |
| 601 | num_tokens=num_tokens, |
| 602 | w1=w1, |
| 603 | w2=w2, |
| 604 | top_k_num=top_k_num, |
| 605 | ) |
| 606 | |
| 607 | # separate function is required for MoE + LoRA |
| 608 | self.moe_sum(intermediate_cache3, output, topk_ids, expert_map) |
| 609 | |
| 610 | def moe_sum( |
| 611 | self, |
| 612 | input: torch.Tensor, |
| 613 | output: torch.Tensor, |
| 614 | topk_ids: torch.Tensor | None = None, |
| 615 | expert_map: torch.Tensor | None = None, |
| 616 | ) -> None: |
| 617 | if expert_map is not None: |
| 618 | # Skip the slots whose expert is not on this rank: the rows the |
| 619 | # alignment dropped, or the zeros the `-1` blocks wrote. |
| 620 | ops.moe_sum(input, output, topk_ids, expert_map) |
| 621 | else: |
| 622 | ops.moe_sum(input, output) |
| 623 | |
| 624 | |
| 625 | class TritonWNA16Experts(TritonExperts): |
| 626 | @staticmethod |
| 627 | def _supports_current_device() -> bool: |
| 628 | return current_platform.is_cuda_alike() or current_platform.is_xpu() |
| 629 | |
| 630 | @staticmethod |
| 631 | def _supports_no_act_and_mul() -> bool: |
| 632 | return True |
| 633 | |
| 634 | @staticmethod |
| 635 | def _supports_quant_scheme( |
| 636 | weight_key: QuantKey | None, |
| 637 | activation_key: QuantKey | None, |
| 638 | ) -> bool: |
| 639 | SUPPORTED_W = [ |
| 640 | kInt4Static, |
| 641 | kInt4StaticAsym, |
| 642 | kInt8Static, |
| 643 | kInt4Static32, |
| 644 | kInt4Static32Asym, |
| 645 | # other group sizes? |
| 646 | ] |
| 647 | return weight_key in SUPPORTED_W |
| 648 | |
| 649 | @staticmethod |
| 650 | def _supports_activation(activation: MoEActivation) -> bool: |
| 651 | return activation in [ |
| 652 | MoEActivation.SILU, |
| 653 | MoEActivation.GELU, |
| 654 | MoEActivation.GELU_TANH, |
| 655 | MoEActivation.SWIGLUOAI, |
| 656 | MoEActivation.SWIGLUOAI_UNINTERLEAVE, |
| 657 | MoEActivation.SWIGLUSTEP, |
| 658 | MoEActivation.SILU_NO_MUL, |
| 659 | MoEActivation.GELU_NO_MUL, |
| 660 | MoEActivation.GELU_TANH_NO_MUL, |
| 661 | MoEActivation.RELU2_NO_MUL, |
| 662 | ] |
| 663 | |
| 664 | @staticmethod |
| 665 | def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: |
| 666 | return not ( |
| 667 | moe_parallel_config.use_fi_nvl_two_sided_kernels |
| 668 | or moe_parallel_config.use_fi_nvl_one_sided_kernels |
| 669 | ) |
| 670 | |
| 671 | def apply( |
| 672 | self, |
| 673 | output: torch.Tensor, |
| 674 | hidden_states: torch.Tensor, |
| 675 | w1: torch.Tensor, |
| 676 | w2: torch.Tensor, |
| 677 | topk_weights: torch.Tensor, |
| 678 | topk_ids: torch.Tensor, |
| 679 | activation: MoEActivation, |
| 680 | global_num_experts: int, |
| 681 | expert_map: torch.Tensor | None, |
| 682 | a1q_scale: torch.Tensor | None, |
| 683 | a2_scale: torch.Tensor | None, |
| 684 | workspace13: torch.Tensor, |
| 685 | workspace2: torch.Tensor, |
| 686 | expert_tokens_meta: mk.ExpertTokensMetadata | None, |
| 687 | apply_router_weight_on_input: bool, |
| 688 | ): |
| 689 | # Check constraints. |
| 690 | if self.quant_config.use_int4_w4a16: |
| 691 | assert hidden_states.size(-1) // 2 == w1.size(2), ( |
| 692 | f"Hidden size mismatch {hidden_states.size(-1) // 2} == {w1.size(2)}" |
| 693 | ) |
| 694 | else: |
| 695 | assert hidden_states.size(-1) == w1.size(2), ( |
| 696 | f"Hidden size mismatch {hidden_states.size(-1)} != {w1.size(2)}" |
| 697 | ) |
| 698 | |
| 699 | assert hidden_states.is_contiguous(), "Hidden_states must be contiguous" |
| 700 | assert hidden_states.dim() == 2 |
| 701 | assert w1.stride(-1) == 1, "Stride of last dimension must be 1" |
| 702 | assert w2.stride(-1) == 1, "Stride of last dimension must be 1" |
| 703 | assert hidden_states.dtype in [ |
| 704 | torch.float32, |
| 705 | torch.float16, |
| 706 | torch.bfloat16, |
| 707 | torch.float8_e4m3fn, |
| 708 | torch.float8_e4m3fnuz, |
| 709 | ] |
| 710 | |
| 711 | E, num_tokens, N, K, top_k_num = self.moe_problem_size( |
| 712 | hidden_states, w1, w2, topk_ids |
| 713 | ) |
| 714 | |
| 715 | if global_num_experts == -1: |
| 716 | global_num_experts = E |
| 717 | |
| 718 | config = try_get_optimal_moe_config( |
| 719 | w1.size(), |
| 720 | w2.size(), |
| 721 | top_k_num, |
| 722 | self.quant_config.config_name(hidden_states.dtype), |
| 723 | num_tokens, |
| 724 | block_shape=self.block_shape, |
| 725 | ) |
| 726 | |
| 727 | if hidden_states.dtype == torch.bfloat16: |
| 728 | compute_type = tl.bfloat16 |
| 729 | elif hidden_states.dtype == torch.float16: |
| 730 | compute_type = tl.float16 |
| 731 | elif hidden_states.dtype == torch.float32: |
| 732 | compute_type = tl.float32 |
| 733 | elif ( |
| 734 | hidden_states.dtype == torch.float8_e4m3fn |
| 735 | or hidden_states.dtype == torch.float8_e4m3fnuz |
| 736 | ): |
| 737 | compute_type = tl.bfloat16 |
| 738 | else: |
| 739 | raise ValueError(f"Unsupported compute_type: {hidden_states.dtype}") |
| 740 | |
| 741 | # Note that the output tensor might be in workspace1 |
| 742 | intermediate_cache1 = _resize_cache(workspace2, (num_tokens, top_k_num, N)) |
| 743 | activation_out_dim = self.adjust_N_for_activation(N, activation) |
| 744 | intermediate_cache2 = _resize_cache( |
| 745 | workspace13, (num_tokens * top_k_num, activation_out_dim) |
| 746 | ) |
| 747 | intermediate_cache3 = _resize_cache(workspace2, (num_tokens, top_k_num, K)) |
| 748 | |
| 749 | # Include fused shared-expert rows while preserving EP remapping. |
| 750 | num_align_experts = w1.shape[0] if expert_map is None else global_num_experts |
| 751 | sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( |
| 752 | topk_ids, config["BLOCK_SIZE_M"], num_align_experts, expert_map |
| 753 | ) |
| 754 | |
| 755 | invoke_fused_moe_wna16_triton_kernel( |
| 756 | hidden_states, |
| 757 | w1, |
| 758 | intermediate_cache1, |
| 759 | self.w1_scale, |
| 760 | self.quant_config.w1_zp, |
| 761 | None, # topk_weights |
| 762 | sorted_token_ids, |
| 763 | expert_ids, |
| 764 | num_tokens_post_padded, |
| 765 | False, # mul_routed_weights |
| 766 | top_k_num, |
| 767 | config, |
| 768 | compute_type=compute_type, |
| 769 | use_int8_w8a16=self.quant_config.use_int8_w8a16, |
| 770 | use_int4_w4a16=self.quant_config.use_int4_w4a16, |
| 771 | block_shape=self.block_shape, |
| 772 | ) |
| 773 | |
| 774 | self.activation( |
| 775 | activation, intermediate_cache2, intermediate_cache1.view(-1, N) |
| 776 | ) |
| 777 | |
| 778 | a2q_scale: torch.Tensor | None = None |
| 779 | |
| 780 | qintermediate_cache2, a2q_scale = moe_kernel_quantize_input( |
| 781 | intermediate_cache2, |
| 782 | a2_scale, |
| 783 | self.quant_dtype, |
| 784 | self.per_act_token_quant, |
| 785 | self.block_shape, |
| 786 | ) |
| 787 | |
| 788 | invoke_fused_moe_wna16_triton_kernel( |
| 789 | qintermediate_cache2, |
| 790 | w2, |
| 791 | intermediate_cache3, |
| 792 | self.w2_scale, |
| 793 | self.quant_config.w2_zp, |
| 794 | topk_weights, |
| 795 | sorted_token_ids, |
| 796 | expert_ids, |
| 797 | num_tokens_post_padded, |
| 798 | not apply_router_weight_on_input, |
| 799 | 1, |
| 800 | config, |
| 801 | compute_type=compute_type, |
| 802 | use_int8_w8a16=self.quant_config.use_int8_w8a16, |
| 803 | use_int4_w4a16=self.quant_config.use_int4_w4a16, |
| 804 | block_shape=self.block_shape, |
| 805 | ) |
| 806 | |
| 807 | # separate function is required for MoE + LoRA |
| 808 | self.moe_sum(intermediate_cache3, output) |