From 4570376ad3532848b0df7a18225085d98e9890a2 Mon Sep 17 00:00:00 2001 From: wangys2008 Date: Tue, 28 Jul 2026 17:34:00 +0800 Subject: [PATCH] feat(metax): add MoE optimizations and contiguous prefill for MetaX backend - activation: fallback from SiluAndMul/GeluAndMul to torch.ops._C or native PyTorch ops - flash_attn: VLLM020_CONTIGUOUS_SINGLE_PREFILL for single-request prefill with contiguous KV cache - fused_moe: add moe_align_block_size, moe_sum, SGL MoE invoke, BF16 noscale fast path, even K fastpath, naive decode mode, stage trace - fused_moe_utils: split staged configs, runtime profile config selection (VLLM020_MOE_RUNTIME_PROFILE), stage2 realign, mcoplib stage1/stage2 token fast paths, small M block override - metax config: add vendor:metax backend for topk_softmax Co-Authored-By: Claude --- .gitignore | 1 + .../backends/vendor/metax/impl/activation.py | 59 +- .../vendor/metax/impl/attention/flash_attn.py | 49 +- .../backends/vendor/metax/impl/fused_moe.py | 372 ++++++++++ .../dispatch/backends/vendor/metax/metax.py | 27 + .../backends/vendor/metax/register_ops.py | 18 + vllm_fl/dispatch/config/metax.yaml | 1 + vllm_fl/ops/fused_moe/fused_moe_utils.py | 651 ++++++++++++++++-- 8 files changed, 1072 insertions(+), 106 deletions(-) diff --git a/.gitignore b/.gitignore index 059ddda3f..d5854158b 100644 --- a/.gitignore +++ b/.gitignore @@ -14,3 +14,4 @@ coverage.xml # Reports report*.json test-results-* +*.orig diff --git a/vllm_fl/dispatch/backends/vendor/metax/impl/activation.py b/vllm_fl/dispatch/backends/vendor/metax/impl/activation.py index 49cd735e4..bc3c28c94 100644 --- a/vllm_fl/dispatch/backends/vendor/metax/impl/activation.py +++ b/vllm_fl/dispatch/backends/vendor/metax/impl/activation.py @@ -14,42 +14,39 @@ # SPDX-License-Identifier: Apache-2.0 # 2026 - Modified by MetaX Integrated Circuits (Shanghai) Co., Ltd. All Rights Reserved. +from __future__ import annotations + import torch -from vllm.model_executor.layers.activation import ( - SiluAndMul, - GeluAndMul, -) +import torch.nn.functional as F def silu_and_mul_maca(obj, x: torch.Tensor) -> torch.Tensor: - """ - SiLU activation followed by element-wise multiplication using CUDA. - - Uses vLLM's optimized CUDA kernel when available. - - Args: - obj: The calling obj (for interface consistency) - x: Input tensor of shape [..., 2*d] + """SiLU activation followed by element-wise multiplication.""" + d = x.shape[-1] // 2 + op = getattr(torch.ops._C, "silu_and_mul", None) + if op is not None: + out = torch.empty(*x.shape[:-1], d, dtype=x.dtype, device=x.device) + op(out, x) + return out - Returns: - Output tensor of shape [..., d] - """ - act_fn = SiluAndMul() - return act_fn.forward_cuda(x) + x1, x2 = x[..., :d], x[..., d:] + return F.silu(x1) * x2 def gelu_and_mul_maca(obj, x: torch.Tensor) -> torch.Tensor: - """ - GELU activation followed by element-wise multiplication using CUDA. - - Uses vLLM's optimized CUDA kernel when available. - - Args: - obj: The calling obj (for interface consistency) - x: Input tensor of shape [..., 2*d] - - Returns: - Output tensor of shape [..., d] - """ - act_fn = GeluAndMul() - return act_fn.forward_cuda(x) + """GELU activation followed by element-wise multiplication.""" + d = x.shape[-1] // 2 + approximate = getattr(obj, "approximate", "none") if obj is not None else "none" + + if approximate == "tanh": + op = getattr(torch.ops._C, "gelu_tanh_and_mul", None) + else: + op = getattr(torch.ops._C, "gelu_and_mul", None) + + if op is not None: + out = torch.empty(*x.shape[:-1], d, dtype=x.dtype, device=x.device) + op(out, x) + return out + + x1, x2 = x[..., :d], x[..., d:] + return F.gelu(x1, approximate=approximate) * x2 diff --git a/vllm_fl/dispatch/backends/vendor/metax/impl/attention/flash_attn.py b/vllm_fl/dispatch/backends/vendor/metax/impl/attention/flash_attn.py index e031044a9..dfdb98721 100644 --- a/vllm_fl/dispatch/backends/vendor/metax/impl/attention/flash_attn.py +++ b/vllm_fl/dispatch/backends/vendor/metax/impl/attention/flash_attn.py @@ -4,6 +4,7 @@ """Attention layer with FlashAttention.""" from dataclasses import dataclass +import os from typing import ClassVar import numpy as np @@ -798,13 +799,53 @@ def forward( # For handling prefill decode split num_decode_tokens = attn_metadata.num_decode_tokens if attn_metadata.num_prefills > 0: + prefill_seq_lens = attn_metadata.prefill_seq_lens.tolist() cu_prefix_kv_lens = torch.tensor( - [0] + attn_metadata.prefill_seq_lens.tolist(), + [0] + prefill_seq_lens, device=attn_metadata.prefill_seq_lens.device, dtype=torch.int32, ).cumsum(dim=0, dtype=torch.int32) - output[num_decode_tokens:num_actual_tokens] = ( - flash_attn_varlen_func( + use_contiguous_prefill = ( + os.environ.get( + "VLLM020_CONTIGUOUS_SINGLE_PREFILL", "" + ).strip().lower() + in ("1", "true", "yes", "on") + and attn_metadata.num_prefills == 1 + and not self.kv_cache_dtype.startswith("fp8") + ) + if use_contiguous_prefill: + seq_len = prefill_seq_lens[0] + block_size = key_cache.shape[1] + num_kv_blocks = cdiv(seq_len, block_size) + block_ids = attn_metadata.prefill_block_table[ + 0, :num_kv_blocks + ] + contiguous_key = torch.index_select( + key_cache, 0, block_ids + ).reshape(-1, self.num_kv_heads, self.head_size)[:seq_len] + contiguous_value = torch.index_select( + value_cache, 0, block_ids + ).reshape(-1, self.num_kv_heads, self.head_size)[:seq_len] + logger.info_once( + "Using contiguous single-request prefill FlashAttention." + ) + prefill_output = flash_attn_varlen_func( + q=query[num_decode_tokens:num_actual_tokens], + k=contiguous_key, + v=contiguous_value, + cu_seqlens_q=attn_metadata.prefill_query_start_loc, + cu_seqlens_k=cu_prefix_kv_lens, + max_seqlen_q=attn_metadata.max_query_len, + max_seqlen_k=attn_metadata.prefill_max_seq_len, + softmax_scale=self.scale, + causal=attn_metadata.causal, + window_size=self.sliding_window, + alibi_slopes=self.alibi_slopes, + softcap=self.logits_soft_cap, + s_aux=self.sinks, + ) + else: + prefill_output = flash_attn_varlen_func( q=query[num_decode_tokens:num_actual_tokens], k=key_cache, v=value_cache, @@ -820,7 +861,7 @@ def forward( softcap=self.logits_soft_cap, s_aux=self.sinks, ) - ) + output[num_decode_tokens:num_actual_tokens] = prefill_output if attn_metadata.num_decodes > 0: # Use flash_attn_with_kvcache for normal decoding. decode_query = query[:num_decode_tokens] diff --git a/vllm_fl/dispatch/backends/vendor/metax/impl/fused_moe.py b/vllm_fl/dispatch/backends/vendor/metax/impl/fused_moe.py index 012dcafab..bbcf021f4 100644 --- a/vllm_fl/dispatch/backends/vendor/metax/impl/fused_moe.py +++ b/vllm_fl/dispatch/backends/vendor/metax/impl/fused_moe.py @@ -5,9 +5,60 @@ """ from typing import Any, Optional +import os +import time import torch +from vllm.logger import init_logger from vllm.triton_utils import tl, triton +from vllm.utils.math_utils import round_up + +logger = init_logger(__name__) + +_SGL_ZERO_BIAS_CACHE: dict[tuple[int, int, torch.dtype, torch.device], torch.Tensor] = {} +_MOE_STAGE_TRACE_COUNTS: dict[tuple[Any, ...], int] = {} + + +def moe_align_block_size_maca( + topk_ids: torch.Tensor, + block_size: int, + num_experts: int, + expert_map: torch.Tensor | None = None, + pad_sorted_ids: bool = False, + ignore_invalid_experts: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + from vllm._custom_ops import moe_align_block_size + + max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1) + if pad_sorted_ids: + max_num_tokens_padded = round_up(max_num_tokens_padded, block_size) + if topk_ids.numel() < num_experts: + max_num_tokens_padded = min( + topk_ids.numel() * block_size, max_num_tokens_padded + ) + sorted_ids = torch.empty( + (max_num_tokens_padded,), dtype=torch.int32, device=topk_ids.device + ) + max_num_m_blocks = triton.cdiv(max_num_tokens_padded, block_size) + expert_ids = torch.empty( + (max_num_m_blocks,), dtype=torch.int32, device=topk_ids.device + ) + num_tokens_post_pad = torch.empty((1), dtype=torch.int32, device=topk_ids.device) + + moe_align_block_size( + topk_ids, + num_experts, + block_size, + sorted_ids, + expert_ids, + num_tokens_post_pad, + expert_map if ignore_invalid_experts else None, + ) + + if expert_map is not None and not ignore_invalid_experts: + expert_ids = expert_map[expert_ids] + + return sorted_ids, expert_ids, num_tokens_post_pad def topk_softmax_maca( @@ -25,6 +76,147 @@ def topk_softmax_maca( return topk_weights, topk_indices +def moe_sum_maca(inp, out): + from vllm._custom_ops import moe_sum + + moe_sum(inp, out) + + +def _enable_sgl_zero_bias() -> bool: + value = os.getenv("VLLM020_SGL_ZERO_BIAS", "0").strip().lower() + return value in ("1", "true", "yes", "on") + + +def _enable_bf16_noscale_moe() -> bool: + value = os.getenv("VLLM020_MOE_BF16_NOSCALE", "0").strip().lower() + return value in ("1", "true", "yes", "on") + + +def _enable_even_k_fastpath() -> bool: + value = os.getenv("VLLM020_MOE_EVEN_K_FASTPATH", "0").strip().lower() + return value in ("1", "true", "yes", "on") + + +def _enable_stage2_full_n_fastpath() -> bool: + value = os.getenv("VLLM020_MOE_STAGE2_FULL_N_FASTPATH", "0").strip().lower() + return value in ("1", "true", "yes", "on") + + +def _enable_all_experts_local_fastpath() -> bool: + value = os.getenv("VLLM020_MOE_ALL_EXPERTS_LOCAL_FASTPATH", "0").strip().lower() + return value in ("1", "true", "yes", "on") + + +def _enable_stage_trace() -> bool: + value = os.getenv("VLLM020_MOE_STAGE_TRACE", "0").strip().lower() + return value in ("1", "true", "yes", "on") + + +def _naive_decode_max_m() -> int: + try: + return int(os.getenv("VLLM020_MOE_NAIVE_DECODE_MAX_M", "0")) + except ValueError: + return 0 + + +def _stage_trace_limit() -> int: + try: + return int(os.getenv("VLLM020_MOE_STAGE_TRACE_MAX", "256")) + except ValueError: + return 256 + + +def _stage_name(topk_weights: Optional[torch.Tensor], mul_routed_weight: bool, + top_k: int) -> str: + if topk_weights is None and not mul_routed_weight and top_k != 1: + return "stage1_w1" + if topk_weights is not None and top_k == 1: + return "stage2_w2" + return "unknown" + + +def _trace_moe_stage( + *, + stage: str, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + EM: int, + top_k: int, + mul_routed_weight: bool, + config: dict[str, Any], + block_size_k: int, + even_ks: bool, + sorted_token_ids: Optional[torch.Tensor], + naive_decode: bool, +) -> None: + if not _enable_stage_trace(): + return + + key = ( + stage, + int(A.size(0)), + int(B.size(1)), + int(B.size(2)), + int(EM), + int(top_k), + bool(mul_routed_weight), + int(config.get("BLOCK_SIZE_M", -1)), + int(config.get("BLOCK_SIZE_N", -1)), + int(block_size_k), + int(config.get("GROUP_SIZE_M", -1)), + int(config.get("SPLIT_K", 1)), + int(config.get("num_warps", -1)), + int(config.get("num_stages", -1)), + bool(even_ks), + bool(sorted_token_ids is None), + bool(naive_decode), + ) + count = _MOE_STAGE_TRACE_COUNTS.get(key, 0) + 1 + _MOE_STAGE_TRACE_COUNTS[key] = count + if count != 1: + return + if len(_MOE_STAGE_TRACE_COUNTS) > _stage_trace_limit(): + return + + logger.info( + "[VLLM020_MOE_STAGE_TRACE] ts=%.6f pid=%d stage=%s M=%d N=%d K=%d " + "C_shape=%s EM=%d top_k=%d mul_routed=%s sorted=%s cfg=%s " + "BLOCK_SIZE_K=%d even_ks=%s naive_decode=%s", + time.time(), + os.getpid(), + stage, + A.size(0), + B.size(1), + B.size(2), + tuple(C.shape), + EM, + top_k, + mul_routed_weight, + sorted_token_ids is not None, + { + "BLOCK_SIZE_M": config.get("BLOCK_SIZE_M"), + "BLOCK_SIZE_N": config.get("BLOCK_SIZE_N"), + "GROUP_SIZE_M": config.get("GROUP_SIZE_M"), + "SPLIT_K": config.get("SPLIT_K", 1), + "num_warps": config.get("num_warps"), + "num_stages": config.get("num_stages"), + }, + block_size_k, + even_ks, + naive_decode, + ) + + +def _get_sgl_zero_bias(B: torch.Tensor) -> torch.Tensor: + key = (B.size(0), B.size(1), B.dtype, B.device) + cached = _SGL_ZERO_BIAS_CACHE.get(key) + if cached is None: + cached = torch.zeros((B.size(0), B.size(1)), dtype=B.dtype, device=B.device) + _SGL_ZERO_BIAS_CACHE[key] = cached + return cached + + def invoke_fused_moe_triton_kernel_maca( A: torch.Tensor, B: torch.Tensor, @@ -46,6 +238,7 @@ def invoke_fused_moe_triton_kernel_maca( per_channel_quant: bool = False, block_shape: Optional[list[int]] = None, B_bias: Optional[torch.Tensor] = None, + _topk_ids: Optional[torch.Tensor] = None, ): """ MetaX implementation of invoke_fused_moe_triton_kernel using mcoplib's @@ -58,6 +251,81 @@ def invoke_fused_moe_triton_kernel_maca( assert topk_weights is None or topk_weights.stride(1) == 1 assert sorted_token_ids is None or sorted_token_ids.stride(0) == 1 + if ( + os.getenv("VLLM020_SGL_MOE", "0").lower() in ("1", "true", "yes", "on") + and _topk_ids is not None + ): + try: + from mcoplib.triton_fused_moe import sgl_invoke_fused_moe_kernel + + sgl_keys = { + "BLOCK_SIZE_M", + "BLOCK_SIZE_N", + "BLOCK_SIZE_K", + "GROUP_SIZE_M", + "num_warps", + "num_stages", + } + sgl_config = {key: value for key, value in config.items() if key in sgl_keys} + sgl_bias = ( + _get_sgl_zero_bias(B) + if B_bias is None and _enable_sgl_zero_bias() + else B_bias + ) + sgl_invoke_fused_moe_kernel( + A, + B, + sgl_bias, + C, + A_scale, + B_scale, + None, + topk_weights, + _topk_ids, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + mul_routed_weight, + top_k, + sgl_config, + compute_type, + use_fp8_w8a8, + use_int8_w8a8, + use_int8_w8a16, + use_int4_w4a16, + per_channel_quant, + block_shape=block_shape, + filter_expert=True, + ) + return + except Exception as err: + logger.warning_once( + "mcoplib sgl_invoke_fused_moe_kernel failed, fallback to " + "fused_moe_triton_kernel: %s", + err, + ) + + stage = _stage_name(topk_weights, mul_routed_weight, top_k) + request_tokens = C.size(0) if C.ndim >= 2 else A.size(0) + naive_decode = False + naive_decode_max_m = _naive_decode_max_m() + if ( + naive_decode_max_m > 0 + and _topk_ids is not None + and request_tokens <= naive_decode_max_m + and not use_fp8_w8a8 + and not use_int8_w8a8 + and not use_int8_w8a16 + and not use_int4_w4a16 + ): + sorted_token_ids = None + expert_ids = _topk_ids.reshape(-1) + config = config.copy() + config["BLOCK_SIZE_M"] = 1 + config["GROUP_SIZE_M"] = 1 + config["SPLIT_K"] = 1 + naive_decode = True + M = A.size(0) num_tokens = M * top_k if sorted_token_ids is not None: @@ -83,6 +351,107 @@ def invoke_fused_moe_triton_kernel_maca( BLOCK_SIZE_K = config.pop("BLOCK_SIZE_K") if block_shape is not None: BLOCK_SIZE_K = min(BLOCK_SIZE_K, min(block_shape[0], block_shape[1])) + even_ks = ( + _enable_even_k_fastpath() + and B.size(2) % (BLOCK_SIZE_K * config.get("SPLIT_K", 1)) == 0 + ) + stage2_full_n_fastpath = ( + _enable_stage2_full_n_fastpath() + and stage == "stage2_w2" + and A.dtype == torch.bfloat16 + and B.dtype == torch.bfloat16 + and C.dtype == torch.bfloat16 + and A_scale is None + and B_scale is None + and B_bias is None + and not use_fp8_w8a8 + and not use_int8_w8a8 + and not use_int8_w8a16 + and not use_int4_w4a16 + and config.get("SPLIT_K", 1) == 1 + and top_k == 1 + and even_ks + and B.size(1) % config["BLOCK_SIZE_N"] == 0 + ) + if stage2_full_n_fastpath: + logger.info_once( + "VLLM020_MOE_STAGE2_FULL_N_FASTPATH enabled: N=%d BLOCK_SIZE_N=%d", + B.size(1), + config["BLOCK_SIZE_N"], + ) + _trace_moe_stage( + stage=stage, + A=A, + B=B, + C=C, + EM=EM, + top_k=top_k, + mul_routed_weight=mul_routed_weight, + config=config, + block_size_k=BLOCK_SIZE_K, + even_ks=even_ks, + sorted_token_ids=sorted_token_ids, + naive_decode=naive_decode, + ) + + if ( + _enable_bf16_noscale_moe() + and A.dtype == torch.bfloat16 + and B.dtype == torch.bfloat16 + and C.dtype == torch.bfloat16 + and A_scale is None + and B_scale is None + and B_bias is None + and not use_fp8_w8a8 + and not use_int8_w8a8 + and not use_int8_w8a16 + and not use_int4_w4a16 + and config.get("SPLIT_K", 1) == 1 + ): + try: + from mcoplib.triton_fused_moe import fused_moe_triton_kernel_bf16_noscale + + noscale_config = config.copy() + noscale_config["BLOCK_SIZE_K"] = BLOCK_SIZE_K + noscale_config.pop("SPLIT_K", None) + noscale_grid = lambda META: ( + triton.cdiv(EM, META["BLOCK_SIZE_M"]) + * triton.cdiv(B.size(1), META["BLOCK_SIZE_N"]), + ) + fused_moe_triton_kernel_bf16_noscale( + noscale_grid, + A, + B, + C, + topk_weights, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + B.size(1), + B.size(2), + EM, + num_tokens, + A.stride(0), + A.stride(1), + B.stride(0), + B.stride(2), + B.stride(1), + C.stride(1), + C.stride(2), + naive_block_assignment=(sorted_token_ids is None), + MUL_ROUTED_WEIGHT=mul_routed_weight, + top_k=top_k, + compute_type=compute_type, + FAST_F32_TO_BF16=True, + **noscale_config, + ) + return + except Exception as err: + logger.warning_once( + "bf16 noscale fused MoE kernel failed, fallback to " + "fused_moe_triton_kernel: %s", + err, + ) fused_moe_triton_kernel( grid, @@ -128,5 +497,8 @@ def invoke_fused_moe_triton_kernel_maca( BLOCK_SIZE_K=BLOCK_SIZE_K, E=B.size(0), FAST_F32_TO_BF16=True, + EVEN_KS=even_ks, + STAGE2_FULL_N_FASTPATH=stage2_full_n_fastpath, + ALL_EXPERTS_LOCAL=_enable_all_experts_local_fastpath(), **config, ) diff --git a/vllm_fl/dispatch/backends/vendor/metax/metax.py b/vllm_fl/dispatch/backends/vendor/metax/metax.py index 4ec499fec..25a4e6480 100644 --- a/vllm_fl/dispatch/backends/vendor/metax/metax.py +++ b/vllm_fl/dispatch/backends/vendor/metax/metax.py @@ -160,6 +160,31 @@ def attention_backend(self, use_mla: bool = False, use_sparse: bool = False) -> # Default to FLASH_ATTN return AttentionBackendEnum.FLASH_ATTN.get_path() + def moe_align_block_size( + self, + topk_ids: torch.Tensor, + block_size: int, + num_experts: int, + expert_map: Optional[torch.Tensor] = None, + pad_sorted_ids: bool = False, + ignore_invalid_experts: bool = False, + ): + from .impl.fused_moe import moe_align_block_size_maca + + return moe_align_block_size_maca( + topk_ids, + block_size, + num_experts, + expert_map, + pad_sorted_ids, + ignore_invalid_experts, + ) + + def moe_sum(self, inp, out): + from .impl.fused_moe import moe_sum_maca + + moe_sum_maca(inp, out) + def topk_softmax( self, topk_weights, @@ -196,6 +221,7 @@ def invoke_fused_moe_triton_kernel( per_channel_quant=False, block_shape=None, B_bias=None, + _topk_ids=None, ): from .impl.fused_moe import invoke_fused_moe_triton_kernel_maca @@ -220,4 +246,5 @@ def invoke_fused_moe_triton_kernel( per_channel_quant=per_channel_quant, block_shape=block_shape, B_bias=B_bias, + _topk_ids=_topk_ids, ) diff --git a/vllm_fl/dispatch/backends/vendor/metax/register_ops.py b/vllm_fl/dispatch/backends/vendor/metax/register_ops.py index 268c5dea3..de2c59863 100644 --- a/vllm_fl/dispatch/backends/vendor/metax/register_ops.py +++ b/vllm_fl/dispatch/backends/vendor/metax/register_ops.py @@ -73,6 +73,24 @@ def register_builtins(registry) -> None: vendor="metax", priority=BackendPriority.VENDOR, ), + # MoE align + OpImpl( + op_name="moe_align_block_size", + impl_id="vendor.metax", + kind=BackendImplKind.VENDOR, + fn=_bind_is_available(backend.moe_align_block_size, is_avail), + vendor="metax", + priority=BackendPriority.VENDOR, + ), + # MoE sum + OpImpl( + op_name="moe_sum", + impl_id="vendor.metax", + kind=BackendImplKind.VENDOR, + fn=_bind_is_available(backend.moe_sum, is_avail), + vendor="metax", + priority=BackendPriority.VENDOR, + ), # topk softmax OpImpl( op_name="topk_softmax", diff --git a/vllm_fl/dispatch/config/metax.yaml b/vllm_fl/dispatch/config/metax.yaml index 0133b0ff1..56f6a3119 100644 --- a/vllm_fl/dispatch/config/metax.yaml +++ b/vllm_fl/dispatch/config/metax.yaml @@ -54,6 +54,7 @@ op_backends: - flagos - reference topk_softmax: + - vendor:metax - flagos - reference diff --git a/vllm_fl/ops/fused_moe/fused_moe_utils.py b/vllm_fl/ops/fused_moe/fused_moe_utils.py index 43d529fcb..cff99aef6 100644 --- a/vllm_fl/ops/fused_moe/fused_moe_utils.py +++ b/vllm_fl/ops/fused_moe/fused_moe_utils.py @@ -1,19 +1,8 @@ -# Copyright 2026 FlagOS Contributors -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - from enum import Enum +import json +import os +import time from typing import Any import torch @@ -37,16 +26,374 @@ get_flashinfer_moe_backend, ) from vllm.triton_utils import tl, triton -from vllm_fl.dispatch import CachedOp +from vllm_fl.dispatch import call_op from vllm_fl.ops.fused_moe.activation import apply_moe_activation from vllm_fl.utils import use_flaggems -_moe_align_block_size = CachedOp("moe_align_block_size") -_invoke_fused_moe_triton_kernel = CachedOp("invoke_fused_moe_triton_kernel") -_moe_sum = CachedOp("moe_sum") - logger = init_logger(__name__) +_MOE_CONFIG_KEYS = { + "BLOCK_SIZE_M", + "BLOCK_SIZE_N", + "BLOCK_SIZE_K", + "GROUP_SIZE_M", + "SPLIT_K", + "num_warps", + "num_stages", +} + +_MOE_RUNTIME_PROFILE_LOGGED: set[tuple[str, int]] = set() +_MOE_RUNTIME_PROFILE_TRACE_LOGGED: set[tuple[Any, ...]] = set() + +_MOE_SMALL_DECODE_CONFIGS: dict[int, dict[str, dict[str, int]]] = { + 1: { + "stage1": {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 32, "SPLIT_K": 1, "num_warps": 4, "num_stages": 3}, + "stage2": {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 1, "SPLIT_K": 1, "num_warps": 4, "num_stages": 3}, + }, + 2: { + "stage1": {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 16, "SPLIT_K": 1, "num_warps": 4, "num_stages": 3}, + "stage2": {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 64, "SPLIT_K": 1, "num_warps": 4, "num_stages": 3}, + }, + 4: { + "stage1": {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 8, "SPLIT_K": 1, "num_warps": 4, "num_stages": 3}, + "stage2": {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 8, "SPLIT_K": 1, "num_warps": 4, "num_stages": 3}, + }, + 8: { + "stage1": {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 32, "SPLIT_K": 1, "num_warps": 4, "num_stages": 3}, + "stage2": {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 1, "SPLIT_K": 1, "num_warps": 4, "num_stages": 3}, + }, + 16: { + "stage1": {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 8, "SPLIT_K": 1, "num_warps": 4, "num_stages": 3}, + "stage2": {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 1, "SPLIT_K": 1, "num_warps": 4, "num_stages": 3}, + }, + 24: { + "stage1": {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 16, "SPLIT_K": 1, "num_warps": 4, "num_stages": 3}, + "stage2": {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 1, "SPLIT_K": 1, "num_warps": 4, "num_stages": 3}, + }, + 32: { + "stage1": {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 1, "SPLIT_K": 1, "num_warps": 4, "num_stages": 3}, + "stage2": {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 1, "SPLIT_K": 1, "num_warps": 4, "num_stages": 3}, + }, + 40: { + "stage1": {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 1, "SPLIT_K": 1, "num_warps": 4, "num_stages": 3}, + "stage2": {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 16, "SPLIT_K": 1, "num_warps": 4, "num_stages": 3}, + }, + 48: { + "stage1": {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 1, "SPLIT_K": 1, "num_warps": 4, "num_stages": 3}, + "stage2": {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 1, "SPLIT_K": 1, "num_warps": 4, "num_stages": 3}, + }, + 56: { + "stage1": {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 16, "SPLIT_K": 1, "num_warps": 4, "num_stages": 3}, + "stage2": {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 1, "SPLIT_K": 1, "num_warps": 4, "num_stages": 3}, + }, + 64: { + "stage1": {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 1, "SPLIT_K": 1, "num_warps": 4, "num_stages": 3}, + "stage2": {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 1, "SPLIT_K": 1, "num_warps": 4, "num_stages": 3}, + }, +} + + +def _sanitize_moe_config(config: dict[str, Any]) -> dict[str, Any]: + sanitized = {key: value for key, value in config.items() if key in _MOE_CONFIG_KEYS} + sanitized.setdefault("SPLIT_K", 1) + return sanitized + + +def _limit_config_for_mcoplib(config: dict[str, Any]) -> dict[str, Any]: + limited = config.copy() + if limited.get("BLOCK_SIZE_K", 0) > 64: + limited["BLOCK_SIZE_K"] = 64 + if limited.get("num_stages", 0) > 3: + limited["num_stages"] = 3 + return limited + + +def _split_staged_moe_config( + config: dict[str, Any], +) -> tuple[dict[str, Any], dict[str, Any]]: + stage1 = config.get("stage1", config) + stage2 = config.get("stage2", config) + stage1_config = _limit_config_for_mcoplib(_sanitize_moe_config(stage1)) + stage2_config = _limit_config_for_mcoplib(_sanitize_moe_config(stage2)) + + # moe_align_block_size pads according to the stage1 BLOCK_SIZE_M. Keep w2 + # on the same M block contract unless the assignment path is redesigned. + if stage2_config.get("BLOCK_SIZE_M") != stage1_config.get("BLOCK_SIZE_M"): + stage2_config = stage2_config.copy() + stage2_config["BLOCK_SIZE_M"] = stage1_config["BLOCK_SIZE_M"] + return stage1_config, stage2_config + + +def _enable_moe_oot_without_flaggems() -> bool: + value = os.environ.get("VLLM_FL_ENABLE_MOE_METHOD_OOT", "").strip().lower() + return value in ("1", "true", "yes", "on") + + +def _enable_v626_moe_cfg() -> bool: + value = os.environ.get("VLLM_FL_MOE_V626_CFG", "").strip().lower() + return value in ("1", "true", "yes", "on") + + +def _enable_mcop_moe_sum() -> bool: + value = os.environ.get("VLLM020_MCOP_MOE_SUM", "").strip().lower() + return value in ("1", "true", "yes", "on") + + +def _moe_stage2_token_sum_max_m() -> int: + try: + return int(os.environ.get("VLLM020_MOE_STAGE2_TOKEN_SUM_MAX_M", "0")) + except ValueError: + return 0 + + +def _moe_stage2_token_sum_block_n() -> int: + try: + return int(os.environ.get("VLLM020_MOE_STAGE2_TOKEN_SUM_BLOCK_N", "64")) + except ValueError: + return 64 + + +def _moe_stage1_token_max_m() -> int: + try: + return int(os.environ.get("VLLM020_MOE_STAGE1_TOKEN_MAX_M", "0")) + except ValueError: + return 0 + + +def _moe_stage1_token_block_n() -> int: + try: + return int(os.environ.get("VLLM020_MOE_STAGE1_TOKEN_BLOCK_N", "64")) + except ValueError: + return 64 + + +def _enable_stage2_realign() -> bool: + value = os.environ.get("VLLM020_MOE_STAGE2_REALIGN", "0").strip().lower() + return value in ("1", "true", "yes", "on") + + +def _force_stage2_realign() -> bool: + value = os.environ.get("VLLM020_MOE_STAGE2_FORCE_REALIGN", "0").strip().lower() + return value in ("1", "true", "yes", "on") + + +def _moe_stage2_realign_min_m() -> int: + try: + return int(os.environ.get("VLLM020_MOE_STAGE2_REALIGN_MIN_M", "0")) + except ValueError: + return 0 + + +def _env_int(name: str) -> int | None: + value = os.environ.get(name, "").strip() + if not value: + return None + try: + return int(value) + except ValueError: + logger.warning_once("Ignore invalid %s=%s; expected int", name, value) + return None + + +def _apply_stage2_gemm_override( + stage2_config: dict[str, Any], +) -> tuple[dict[str, Any], bool]: + env_to_key = { + "VLLM020_MOE_STAGE2_BLOCK_M": "BLOCK_SIZE_M", + "VLLM020_MOE_STAGE2_BLOCK_N": "BLOCK_SIZE_N", + "VLLM020_MOE_STAGE2_BLOCK_K": "BLOCK_SIZE_K", + "VLLM020_MOE_STAGE2_GROUP_M": "GROUP_SIZE_M", + "VLLM020_MOE_STAGE2_SPLIT_K": "SPLIT_K", + "VLLM020_MOE_STAGE2_WARPS": "num_warps", + "VLLM020_MOE_STAGE2_STAGES": "num_stages", + } + updated = stage2_config.copy() + changed = False + for env_name, key in env_to_key.items(): + value = _env_int(env_name) + if value is None: + continue + if value <= 0: + logger.warning_once("Ignore invalid %s=%d; expected > 0", env_name, value) + continue + updated[key] = value + changed = True + if changed: + updated = _limit_config_for_mcoplib(_sanitize_moe_config(updated)) + logger.info_once( + "VLLM020 stage2 GEMM override enabled: %s", + _cfg_signature(updated), + ) + return updated, changed + + +def _moe_smallm_block_m_max_m() -> int: + try: + return int(os.environ.get("VLLM020_MOE_SMALLM_BLOCK_M_MAX_M", "0")) + except ValueError: + return 0 + + +def _moe_smallm_block_m() -> int: + try: + return int(os.environ.get("VLLM020_MOE_SMALLM_BLOCK_M", "0")) + except ValueError: + return 0 + + +def _apply_smallm_block_m_override( + num_tokens: int, + stage1_config: dict[str, Any], + stage2_config: dict[str, Any], +) -> tuple[dict[str, Any], dict[str, Any], bool]: + max_m = _moe_smallm_block_m_max_m() + block_m = _moe_smallm_block_m() + if max_m <= 0 or block_m <= 0 or int(num_tokens) > max_m: + return stage1_config, stage2_config, False + + if block_m not in (4, 8, 16, 32, 64, 128): + logger.warning_once( + "Ignore invalid VLLM020_MOE_SMALLM_BLOCK_M=%s; expected one of " + "4/8/16/32/64/128", + block_m, + ) + return stage1_config, stage2_config, False + + stage1_config = stage1_config.copy() + stage2_config = stage2_config.copy() + old_stage1_m = stage1_config.get("BLOCK_SIZE_M") + old_stage2_m = stage2_config.get("BLOCK_SIZE_M") + stage1_config["BLOCK_SIZE_M"] = block_m + stage2_config["BLOCK_SIZE_M"] = block_m + logger.info_once( + "VLLM020_MOE_SMALLM_BLOCK_M override enabled: max_m=%d block_m=%d " + "(stage1 %s -> %d, stage2 %s -> %d)", + max_m, + block_m, + old_stage1_m, + block_m, + old_stage2_m, + block_m, + ) + return stage1_config, stage2_config, True + + +def _cfg_signature(config: dict[str, Any]) -> str: + keys = ("BLOCK_SIZE_M", "BLOCK_SIZE_N", "BLOCK_SIZE_K", "GROUP_SIZE_M", + "num_warps", "num_stages") + return "/".join(str(config.get(key, "-")) for key in keys) + + +def _moe_runtime_profile() -> str: + return os.environ.get("VLLM020_MOE_RUNTIME_PROFILE", "").strip().lower() + + +def _trace_moe_runtime_config( + num_tokens: int, + stage1_config: dict[str, Any], + stage2_config: dict[str, Any], + selected_by_runtime_profile: bool, +) -> None: + path = os.environ.get("VLLM020_MOE_RUNTIME_PROFILE_TRACE", "").strip() + if not path: + return + + key = ( + int(num_tokens), + _cfg_signature(stage1_config), + _cfg_signature(stage2_config), + selected_by_runtime_profile, + ) + if key in _MOE_RUNTIME_PROFILE_TRACE_LOGGED: + return + _MOE_RUNTIME_PROFILE_TRACE_LOGGED.add(key) + + max_records = int(os.environ.get("VLLM020_MOE_RUNTIME_PROFILE_TRACE_MAX", "128")) + if len(_MOE_RUNTIME_PROFILE_TRACE_LOGGED) > max_records: + return + + record = { + "ts": time.time(), + "pid": os.getpid(), + "profile": _moe_runtime_profile(), + "num_tokens": int(num_tokens), + "selected_by_runtime_profile": selected_by_runtime_profile, + "stage1": {k: stage1_config.get(k) for k in sorted(_MOE_CONFIG_KEYS)}, + "stage2": {k: stage2_config.get(k) for k in sorted(_MOE_CONFIG_KEYS)}, + } + try: + with open(path, "a", encoding="utf-8") as f: + f.write(json.dumps(record, sort_keys=True) + "\n") + except OSError as exc: + logger.warning_once("Failed to write MoE runtime profile trace %s: %s", path, exc) + + +def _select_runtime_moe_profile_config( + num_tokens: int, + stage1_config: dict[str, Any], + stage2_config: dict[str, Any], +) -> tuple[dict[str, Any], dict[str, Any], bool]: + profile = _moe_runtime_profile() + if not profile: + return stage1_config, stage2_config, False + + if profile in ("1k", "short1k"): + enabled_buckets = {40, 48, 56, 64} + elif profile in ("4k", "short4k", "decode"): + enabled_buckets = set(_MOE_SMALL_DECODE_CONFIGS) + else: + logger.warning_once( + "Unknown VLLM020_MOE_RUNTIME_PROFILE=%s; using default MoE config", + profile, + ) + return stage1_config, stage2_config, False + + if int(num_tokens) not in enabled_buckets: + return stage1_config, stage2_config, False + + selected = _MOE_SMALL_DECODE_CONFIGS.get(int(num_tokens)) + if selected is None: + return stage1_config, stage2_config, False + + tuned_stage1, tuned_stage2 = _split_staged_moe_config(selected) + log_key = (profile, int(num_tokens)) + if log_key not in _MOE_RUNTIME_PROFILE_LOGGED: + _MOE_RUNTIME_PROFILE_LOGGED.add(log_key) + logger.info( + "VLLM020_MOE_RUNTIME_PROFILE=%s selected M=%d stage1=%s stage2=%s", + profile, + int(num_tokens), + _cfg_signature(tuned_stage1), + _cfg_signature(tuned_stage2), + ) + return tuned_stage1, tuned_stage2, True + + +def _apply_v626_moe_cfg(B: torch.Tensor, config: dict[str, Any]) -> dict[str, Any]: + try: + k_dim = B.shape[2] + except Exception: + return config + if k_dim < 1024: + return config + tuned = config.copy() + tuned["BLOCK_SIZE_N"] = 128 + tuned["BLOCK_SIZE_K"] = 64 + tuned["GROUP_SIZE_M"] = 1 + return tuned + + +def _v626_base_moe_cfg() -> dict[str, int]: + return { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_warps": 4, + "num_stages": 3, + } + def _get_priority_backends(moe_config: FusedMoEConfig) -> list[UnquantizedMoeBackend]: """ @@ -61,7 +408,12 @@ def _move_to_back( ) -> None: backends.append(backends.pop(backends.index(backend))) - if current_platform.is_rocm(): + if current_platform.is_out_of_tree(): + _AVAILABLE_BACKENDS = [ + UnquantizedMoeBackend.TRITON, + UnquantizedMoeBackend.BATCHED_TRITON, + ] + elif current_platform.is_rocm(): _AVAILABLE_BACKENDS = [ UnquantizedMoeBackend.AITER, UnquantizedMoeBackend.TRITON, @@ -102,7 +454,9 @@ def select_unquantized_moe_backend_oot(moe_config: FusedMoEConfig, if current_platform.is_tpu(): return UnquantizedMoeBackend.TPU, None - if current_platform.is_out_of_tree() and use_flaggems(): + if current_platform.is_out_of_tree() and ( + use_flaggems() or _enable_moe_oot_without_flaggems() + ): return UnquantizedMoeBackend.TRITON, TritonExpertsFL if moe_config.is_lora_enabled: @@ -274,7 +628,7 @@ def _prepare_expert_assignment( ), ) - return _moe_align_block_size( + return call_op("moe_align_block_size", topk_ids, config["BLOCK_SIZE_M"], global_num_experts, @@ -301,8 +655,7 @@ def apply( expert_tokens_meta: mk.ExpertTokensMetadata | None, apply_router_weight_on_input: bool, ): - # Fast path (no LoRA, NVIDIA only): single fused FlagGems call. This - # implementation supports both quantized and BF16 experts. + # Fast path (no LoRA, NVIDIA only): single fused FlagGems call. if self._lora_context is None and current_platform.is_cuda(): import flag_gems @@ -369,6 +722,33 @@ def apply( num_tokens, block_shape=self.block_shape, ) + if _enable_v626_moe_cfg(): + base_config = _v626_base_moe_cfg() + stage1_config = _apply_v626_moe_cfg(w1, base_config) + stage2_config = base_config + selected_by_runtime_profile = False + else: + stage1_config, stage2_config = _split_staged_moe_config(config) + stage1_config, stage2_config, selected_by_runtime_profile = ( + _select_runtime_moe_profile_config( + num_tokens, stage1_config, stage2_config + ) + ) + _trace_moe_runtime_config( + num_tokens, stage1_config, stage2_config, selected_by_runtime_profile + ) + stage1_config, stage2_config, smallm_block_m_override = ( + _apply_smallm_block_m_override(num_tokens, stage1_config, stage2_config) + ) + if smallm_block_m_override: + _trace_moe_runtime_config( + num_tokens, stage1_config, stage2_config, True + ) + stage2_config, stage2_override = _apply_stage2_gemm_override(stage2_config) + if stage2_override: + _trace_moe_runtime_config( + num_tokens, stage1_config, stage2_config, True + ) if hidden_states.dtype == torch.bfloat16: compute_type = tl.bfloat16 @@ -395,7 +775,7 @@ def apply( sorted_token_ids, expert_ids, num_tokens_post_padded = ( _prepare_expert_assignment( topk_ids, - config, + stage1_config, num_tokens, top_k_num, global_num_experts, @@ -405,29 +785,38 @@ def apply( block_shape=self.block_shape, ) ) - - _invoke_fused_moe_triton_kernel( - hidden_states, - w1, - intermediate_cache1, - a1q_scale, - self.w1_scale, - None, # topk_weights - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - False, # mul_routed_weights - top_k_num, - config, - compute_type=compute_type, - use_fp8_w8a8=self.quant_config.use_fp8_w8a8, - use_int8_w8a8=self.quant_config.use_int8_w8a8, - use_int8_w8a16=self.quant_config.use_int8_w8a16, - use_int4_w4a16=self.quant_config.use_int4_w4a16, - per_channel_quant=self.per_act_token_quant, - block_shape=self.block_shape, - B_bias=self.w1_bias, + sorted_token_ids_stage2 = sorted_token_ids + expert_ids_stage2 = expert_ids + num_tokens_post_padded_stage2 = num_tokens_post_padded + stage2_block_m_differs = ( + stage2_config.get("BLOCK_SIZE_M") != stage1_config.get("BLOCK_SIZE_M") + ) + stage2_realign_enabled = ( + _enable_stage2_realign() + and num_tokens >= _moe_stage2_realign_min_m() ) + if ( + stage2_realign_enabled + and (stage2_block_m_differs or _force_stage2_realign()) + ): + ( + sorted_token_ids_stage2, + expert_ids_stage2, + num_tokens_post_padded_stage2, + ) = _prepare_expert_assignment( + topk_ids, + stage2_config, + num_tokens, + top_k_num, + global_num_experts, + expert_map, + use_int8_w8a16=self.quant_config.use_int8_w8a16, + use_int4_w4a16=self.quant_config.use_int4_w4a16, + block_shape=self.block_shape, + ) + elif stage2_block_m_differs: + stage2_config = stage2_config.copy() + stage2_config["BLOCK_SIZE_M"] = stage1_config["BLOCK_SIZE_M"] # LoRA w13: applied to intermediate_cache1 before activation, using # hidden_states as the lora_a input. moe_lora_align_block_size is @@ -437,6 +826,70 @@ def apply( num_tokens_post_padded_lora = None token_lora_mapping = None lora_context = self._lora_context + + stage1_token_used = False + stage1_token_max_m = _moe_stage1_token_max_m() + if ( + stage1_token_max_m > 0 + and num_tokens <= stage1_token_max_m + and lora_context is None + and expert_map is None + and hidden_states.dtype == torch.bfloat16 + and w1.dtype == torch.bfloat16 + and intermediate_cache1.dtype == torch.bfloat16 + and a1q_scale is None + and self.w1_scale is None + and self.w1_bias is None + and not self.quant_config.use_fp8_w8a8 + and not self.quant_config.use_int8_w8a8 + and not self.quant_config.use_int8_w8a16 + and not self.quant_config.use_int4_w4a16 + ): + try: + from mcoplib.triton_fused_moe import moe_stage1_token_triton + + moe_stage1_token_triton( + hidden_states, + w1, + intermediate_cache1, + topk_ids, + compute_type, + block_size_n=_moe_stage1_token_block_n(), + block_size_k=stage1_config.get("BLOCK_SIZE_K", 64), + ) + stage1_token_used = True + except Exception as err: + logger.warning_once( + "mcoplib moe_stage1_token_triton failed, fallback to " + "stage1 fused_moe: %s", + err, + ) + + if not stage1_token_used: + call_op("invoke_fused_moe_triton_kernel", + hidden_states, + w1, + intermediate_cache1, + a1q_scale, + self.w1_scale, + None, # topk_weights + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + False, # mul_routed_weights + top_k_num, + stage1_config, + compute_type=compute_type, + use_fp8_w8a8=self.quant_config.use_fp8_w8a8, + use_int8_w8a8=self.quant_config.use_int8_w8a8, + use_int8_w8a16=self.quant_config.use_int8_w8a16, + use_int4_w4a16=self.quant_config.use_int4_w4a16, + per_channel_quant=self.per_act_token_quant, + block_shape=self.block_shape, + B_bias=self.w1_bias, + _topk_ids=topk_ids, + ) + if lora_context is not None: ( sorted_token_ids_lora, @@ -471,28 +924,71 @@ def apply( quantization_emulation=self.quantization_emulation, ) - _invoke_fused_moe_triton_kernel( - qintermediate_cache2, - w2, - intermediate_cache3, - a2q_scale, - self.w2_scale, - topk_weights, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - not apply_router_weight_on_input, - 1, - config, - compute_type=compute_type, - use_fp8_w8a8=self.quant_config.use_fp8_w8a8, - use_int8_w8a8=self.quant_config.use_int8_w8a8, - use_int8_w8a16=self.quant_config.use_int8_w8a16, - use_int4_w4a16=self.quant_config.use_int4_w4a16, - per_channel_quant=self.per_act_token_quant, - block_shape=self.block_shape, - B_bias=self.w2_bias, - ) + stage2_token_sum_used = False + stage2_token_sum_max_m = _moe_stage2_token_sum_max_m() + if ( + stage2_token_sum_max_m > 0 + and num_tokens <= stage2_token_sum_max_m + and lora_context is None + and topk_weights is not None + and not apply_router_weight_on_input + and qintermediate_cache2.dtype == torch.bfloat16 + and w2.dtype == torch.bfloat16 + and output.dtype == torch.bfloat16 + and a2q_scale is None + and self.w2_scale is None + and self.w2_bias is None + and not self.quant_config.use_fp8_w8a8 + and not self.quant_config.use_int8_w8a8 + and not self.quant_config.use_int8_w8a16 + and not self.quant_config.use_int4_w4a16 + ): + try: + from mcoplib.triton_fused_moe import moe_stage2_token_sum_triton + + moe_stage2_token_sum_triton( + qintermediate_cache2, + w2, + output, + topk_weights, + topk_ids, + True, + compute_type, + block_size_n=_moe_stage2_token_sum_block_n(), + block_size_k=stage2_config.get("BLOCK_SIZE_K", 64), + ) + stage2_token_sum_used = True + except Exception as err: + logger.warning_once( + "mcoplib moe_stage2_token_sum_triton failed, fallback to " + "stage2 fused_moe + moe_sum: %s", + err, + ) + + if not stage2_token_sum_used: + call_op("invoke_fused_moe_triton_kernel", + qintermediate_cache2, + w2, + intermediate_cache3, + a2q_scale, + self.w2_scale, + topk_weights, + sorted_token_ids_stage2, + expert_ids_stage2, + num_tokens_post_padded_stage2, + not apply_router_weight_on_input, + 1, + stage2_config, + compute_type=compute_type, + use_fp8_w8a8=self.quant_config.use_fp8_w8a8, + use_int8_w8a8=self.quant_config.use_int8_w8a8, + use_int8_w8a16=self.quant_config.use_int8_w8a16, + use_int4_w4a16=self.quant_config.use_int4_w4a16, + per_channel_quant=self.per_act_token_quant, + block_shape=self.block_shape, + B_bias=self.w2_bias, + _topk_ids=topk_ids, + ) # LoRA w2: applied to intermediate_cache3 before moe_sum, using the # unquantized intermediate_cache2 as the lora_a input. Reuses the @@ -514,7 +1010,20 @@ def apply( ) # separate function is required for MoE + LoRA - self.moe_sum(intermediate_cache3, output) + if not stage2_token_sum_used: + self.moe_sum(intermediate_cache3, output) def moe_sum(self, input: torch.Tensor, output: torch.Tensor) -> None: - _moe_sum(input, output) + if _enable_mcop_moe_sum(): + try: + from mcoplib.triton_fused_moe import moe_sum_reduce_triton + + moe_sum_reduce_triton(input, output, 1.0) + return + except Exception as err: + logger.warning_once( + "mcoplib moe_sum_reduce_triton failed, fallback to dispatch " + "moe_sum: %s", + err, + ) + call_op("moe_sum", input, output)