From 57b404d79dbc6325d8fa1f9c697c31fffb08916a Mon Sep 17 00:00:00 2001 From: ppu-dev Date: Mon, 20 Jul 2026 13:48:48 +0800 Subject: [PATCH] Add PPU-native DeepGEMM BF16 unquantized MoE (VLLM_FL_MOE=deepgemm) Route unquantized BF16 MoE experts through deep_gemm's nopad grouped GEMM (m_grouped_gemm_bf16_bf16_bf16_nt_nopad + deepgemm_moe_permute, small-M decode auto-dispatches to the GEMV kernel) instead of FlagGems Triton fused_moe, on out-of-tree (PPU/thead) platforms. Opt-in via VLLM_FL_MOE=deepgemm; unset = unchanged. Also fixes an UnboundLocalError in _get_priority_backends for out-of-tree platforms when FlagGems is disabled. Co-Authored-By: Claude Opus 4.8 --- vllm_fl/ops/fused_moe/deepgemm_experts.py | 138 +++++++++++++++++++ vllm_fl/ops/fused_moe/deepgemm_microbench.py | 122 ++++++++++++++++ vllm_fl/ops/fused_moe/fused_moe_utils.py | 15 ++ vllm_fl/utils.py | 10 ++ 4 files changed, 285 insertions(+) create mode 100644 vllm_fl/ops/fused_moe/deepgemm_experts.py create mode 100644 vllm_fl/ops/fused_moe/deepgemm_microbench.py diff --git a/vllm_fl/ops/fused_moe/deepgemm_experts.py b/vllm_fl/ops/fused_moe/deepgemm_experts.py new file mode 100644 index 000000000..fe5fffb74 --- /dev/null +++ b/vllm_fl/ops/fused_moe/deepgemm_experts.py @@ -0,0 +1,138 @@ +# Copyright (c) 2025 BAAI. All rights reserved. +"""PPU-native DeepGEMM BF16 unquantized MoE experts (env-gated). + +Opt-in via ``VLLM_FL_MOE=deepgemm`` (see ``vllm_fl.utils.use_deepgemm_moe``). +Replaces the default FlagGems Triton ``fused_moe`` expert compute with +``deep_gemm``'s grouped BF16 GEMM — the same kernels the vendor's native vLLM +0.19 build used (``m_grouped_gemm_bf16_bf16_bf16_nt`` / ``..._gemv``). + +Uses the **nopad** grouped GEMM with a compact (block_align=1) permute: each +expert's rows are packed with NO 128-row padding, and small-M decode auto- +dispatches to the GEMV kernel. (The contiguous/128-aligned layout wastes ~128x +compute per active expert on sparse decode — do NOT use it here.) + +Pipeline (BF16, no FP8 scales): + deepgemm_moe_permute(block_align=1) → nopad GEMM1 → silu_and_mul + → nopad GEMM2 → weighted unpermute+reduce (ep_gather) + +Permute/gather are vendor/vLLM Triton kernels (CUDA-graph safe); ``m_rows`` +(exact per-expert token counts) is fed to the nopad kernel so no host sync / +internal bincount is needed. +""" + +import torch + +import deep_gemm +from deep_gemm.deep_gemm_tuner.deepgemm_tools import deepgemm_moe_permute +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.deep_gemm_utils import ( + compute_aligned_M, + ep_gather, +) +from vllm.model_executor.layers.fused_moe.fused_moe import TritonExperts +from vllm.model_executor.layers.fused_moe.utils import _resize_cache + +from vllm_fl.ops.fused_moe.activation import apply_moe_activation + + +class DeepGemmExpertsFL(TritonExperts): + """OOT unquantized BF16 MoE experts backed by deep_gemm nopad grouped GEMM. + + Subclasses ``TritonExperts`` to inherit ``moe_problem_size``, + ``adjust_N_for_activation`` and the ``TopKWeightAndReduceNoOP`` finalize + contract; overrides ``workspace_shapes`` (compact M_sum = M*topk) and + ``apply``. + """ + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: "mk.ExpertTokensMetadata | None", + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + # Compact layout (block_align=1): no per-expert 128-row padding. + M_sum = compute_aligned_M(M, topk, local_num_experts, 1, expert_tokens_meta) + activation_out_dim = self.adjust_N_for_activation(N, activation) + workspace1 = (M_sum, max(activation_out_dim, K)) + workspace2 = (M_sum, max(N, K)) + output = (M, K) + return (workspace1, workspace2, output) + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: "mk.ExpertTokensMetadata | None", + apply_router_weight_on_input: bool, + ): + assert hidden_states.dtype == torch.bfloat16, ( + "DeepGemmExpertsFL only supports bf16 unquantized MoE" + ) + assert hidden_states.is_contiguous() + assert expert_map is None, ( + "DeepGemmExpertsFL does not support expert parallelism (expert_map)" + ) + + a1 = hidden_states # [M, K] + M, K = a1.shape + local_num_experts, N, K_w = w1.shape # w1: [E, 2I, K] + assert K_w == K + topk = topk_ids.size(1) + + # Kernels use -1 for invalid ids -> topk_ids must be signed (router: int32). + if not topk_ids.dtype.is_signed: + topk_ids = topk_ids.to(torch.int32) + + # ---- compact permute: pack tokens per-expert (no 128 padding) ---- + # returns: a1_perm [M_sum, K], m_indices [M_sum], inv_perm [M, topk], + # m_rows (expert_num_tokens) [E]. M_sum == M * topk. + a1_perm, _scale_out, m_indices, inv_perm, m_rows = deepgemm_moe_permute( + a1, None, topk_ids, local_num_experts, block_align=1, block_k=K + ) + M_sum = a1_perm.size(0) + + # ---- grouped GEMM 1 (nopad): [M_sum, K] x [E, 2I, K]^T -> [M_sum, 2I] ---- + mm1_out = _resize_cache(workspace2, (M_sum, N)) + deep_gemm.m_grouped_gemm_bf16_bf16_bf16_nt_nopad( + a1_perm, w1, mm1_out, m_indices, m_rows + ) + + # ---- activation: silu_and_mul -> [M_sum, I] ---- + activation_out_dim = self.adjust_N_for_activation(N, activation) + act_out = _resize_cache(workspace13, (M_sum, activation_out_dim)) + apply_moe_activation(activation, act_out, mm1_out.view(-1, N)) + + # ---- grouped GEMM 2 (nopad): [M_sum, I] x [E, K, I]^T -> [M_sum, K] ---- + mm2_out = _resize_cache(workspace2, (M_sum, K)) + deep_gemm.m_grouped_gemm_bf16_bf16_bf16_nt_nopad( + act_out, w2, mm2_out, m_indices, m_rows + ) + + # ---- weighted unpermute + reduce over topk -> output [M, K] ---- + if apply_router_weight_on_input: + topk_weights = torch.ones_like(topk_weights) + ep_gather( + input_tensor=mm2_out, + recv_topk_ids=topk_ids, + recv_topk_weight=topk_weights, + input_index=inv_perm, + expert_map=None, + output_tensor=output, + ) diff --git a/vllm_fl/ops/fused_moe/deepgemm_microbench.py b/vllm_fl/ops/fused_moe/deepgemm_microbench.py new file mode 100644 index 000000000..c00af50a2 --- /dev/null +++ b/vllm_fl/ops/fused_moe/deepgemm_microbench.py @@ -0,0 +1,122 @@ +# Copyright (c) 2025 BAAI. All rights reserved. +"""Step-0 gate microbench for the PPU-native DeepGEMM BF16 MoE path. + +Validates, on the real PPU device and Qwen3.5-MoE shapes, that +``deep_gemm.m_grouped_gemm_bf16_bf16_bf16_nt_nopad`` can drive the two expert +GEMMs of an unquantized BF16 MoE, and pins down the permute/unpermute approach +(plan risks R1/R3/R5) before touching serving. + +Uses a self-contained pure-torch argsort permute (dtype-agnostic, no dependency +on vLLM's fp8-oriented ep_scatter nor the missing +get_mk_alignment_for_contiguous_layout). The real DeepGemmExpertsFL may switch +to a Triton permute later; here we only need to confirm the GEMM math + speed. + +Run: python -m vllm_fl.ops.fused_moe.deepgemm_microbench +""" + +import torch +import torch.nn.functional as F + +import deep_gemm + + +E = 256 # experts (TP, no EP -> all experts local) +I = 256 # intermediate size per partition (moe_intermediate_size / TP2) +H = 2048 # hidden size +TOPK = 8 +DTYPE = torch.bfloat16 +DEVICE = "cuda" + + +def _grouped_gemm_nopad(lhs, rhs, m_indices, m_rows): + """lhs [m,k] @ rhs[G,n,k]^T grouped by m_indices -> out [m,n].""" + m, k = lhs.shape + G, n, k2 = rhs.shape + assert k == k2 + out = torch.empty(m, n, dtype=DTYPE, device=DEVICE) + deep_gemm.m_grouped_gemm_bf16_bf16_bf16_nt_nopad(lhs, rhs, out, m_indices, m_rows) + return out + + +def _permute(hidden, topk_ids): + """Pure-torch permute: group (token,expert) rows contiguously by expert. + + Returns a1[m,H], m_indices[m] int32, m_rows[E] int32, order, tok_of_row. + """ + M = hidden.shape[0] + flat_expert = topk_ids.reshape(-1).to(torch.int32) # [M*topk] + tok_idx = torch.arange(M, device=DEVICE).repeat_interleave(TOPK) # [M*topk] + order = torch.argsort(flat_expert) # group by expert + m_indices = flat_expert[order].contiguous() + tok_of_row = tok_idx[order].contiguous() + a1 = hidden.index_select(0, tok_of_row).contiguous() # [m,H] + m_rows = torch.bincount(flat_expert, minlength=E).to(torch.int32) + return a1, m_indices, m_rows, order, tok_of_row + + +def _run(M, w1, w2, gen, check_numeric): + hidden = torch.randn(M, H, dtype=DTYPE, device=DEVICE, generator=gen) + topk_ids = torch.randint(0, E, (M, TOPK), device=DEVICE, generator=gen, dtype=torch.int64) + topk_w = torch.rand(M, TOPK, dtype=torch.float32, device=DEVICE, generator=gen) + + def pipeline(): + a1, m_indices, m_rows, order, tok_of_row = _permute(hidden, topk_ids) + mm1 = _grouped_gemm_nopad(a1, w1, m_indices, m_rows) # [m, 2I] + act = F.silu(mm1[:, :I].float()).to(DTYPE) * mm1[:, I:] # silu_and_mul -> [m, I] + act = act.contiguous() + mm3 = _grouped_gemm_nopad(act, w2, m_indices, m_rows) # [m, H] + # unpermute + weighted reduce + tmp = torch.empty_like(mm3) + tmp[order] = mm3 + tmp = tmp.view(M, TOPK, H) + out = (tmp.float() * topk_w.unsqueeze(-1)).sum(1).to(DTYPE) + return out + + out = pipeline() + torch.cuda.synchronize() + ok = bool(torch.isfinite(out).all()) + print(f"\n=== M={M} (m={M*TOPK}) === out {tuple(out.shape)} finite={ok}") + + if check_numeric: + # reference via gathered per-row bmm (bounded memory: only small M) + a1, m_indices, m_rows, order, tok_of_row = _permute(hidden, topk_ids) + w1g = w1.index_select(0, m_indices) # [m,2I,H] + gate_up = torch.bmm(w1g.float(), a1.float().unsqueeze(-1)).squeeze(-1) # [m,2I] + ref_act = F.silu(gate_up[:, :I]) * gate_up[:, I:] # [m,I] + w2g = w2.index_select(0, m_indices) # [m,H,I] + ref_row = torch.bmm(w2g.float(), ref_act.unsqueeze(-1)).squeeze(-1) # [m,H] + tmp = torch.empty(M * TOPK, H, device=DEVICE) + tmp[order] = ref_row + ref = (tmp.view(M, TOPK, H) * topk_w.unsqueeze(-1)).sum(1) + diff = (out.float() - ref).abs() + rel = diff.max().item() / (ref.abs().max().item() + 1e-6) + print(f" numeric vs torch ref: max_abs={diff.max().item():.3f} " + f"max_rel={rel:.4f} -> {'PASS' if rel < 2e-2 else 'CHECK'}") + + # timing + for _ in range(5): + pipeline() + torch.cuda.synchronize() + s = torch.cuda.Event(enable_timing=True); e = torch.cuda.Event(enable_timing=True) + s.record() + for _ in range(20): + pipeline() + e.record(); torch.cuda.synchronize() + print(f" full pipeline: {s.elapsed_time(e)/20*1000:.1f} us/iter") + + +def main(): + print("device:", torch.cuda.get_device_name(0)) + gen = torch.Generator(device=DEVICE).manual_seed(0) + w1 = torch.randn(E, 2 * I, H, dtype=DTYPE, device=DEVICE, generator=gen) * 0.02 + w2 = torch.randn(E, H, I, dtype=DTYPE, device=DEVICE, generator=gen) * 0.02 + + # decode-ish (small M -> should hit gemv path) with numeric check + _run(16, w1, w2, gen, check_numeric=True) + # prefill-ish (large M -> gemm path), finite + latency only + _run(2048, w1, w2, gen, check_numeric=False) + print("\nGATE: see finite/PASS above") + + +if __name__ == "__main__": + main() diff --git a/vllm_fl/ops/fused_moe/fused_moe_utils.py b/vllm_fl/ops/fused_moe/fused_moe_utils.py index eb60b11f9..bc77dd420 100644 --- a/vllm_fl/ops/fused_moe/fused_moe_utils.py +++ b/vllm_fl/ops/fused_moe/fused_moe_utils.py @@ -66,6 +66,15 @@ def _move_to_back( _AVAILABLE_BACKENDS = [UnquantizedMoeBackend.XPU] elif current_platform.is_cpu(): _AVAILABLE_BACKENDS = [UnquantizedMoeBackend.CPU] + else: + # Out-of-tree platforms (e.g. the FL plugin backends). When FlagGems is + # disabled (USE_FLAGGEMS=0) we still need a native fallback, otherwise + # _AVAILABLE_BACKENDS would be unbound. Native Triton MoE kernels work + # on these CUDA-like devices. + _AVAILABLE_BACKENDS = [ + UnquantizedMoeBackend.TRITON, + UnquantizedMoeBackend.BATCHED_TRITON, + ] return _AVAILABLE_BACKENDS ## Adopt from select_unquantized_moe_backend @@ -84,6 +93,12 @@ def select_unquantized_moe_backend_oot(moe_config: FusedMoEConfig, return UnquantizedMoeBackend.TPU, None if current_platform.is_out_of_tree() and use_flaggems(): + # Opt-in (VLLM_FL_MOE=deepgemm): route BF16 experts through the + # PPU-native DeepGEMM grouped GEMM instead of FlagGems Triton fused_moe. + from vllm_fl.utils import use_deepgemm_moe + if use_deepgemm_moe(): + from vllm_fl.ops.fused_moe.deepgemm_experts import DeepGemmExpertsFL + return UnquantizedMoeBackend.TRITON, DeepGemmExpertsFL return UnquantizedMoeBackend.TRITON, TritonExpertsFL if moe_config.is_lora_enabled: diff --git a/vllm_fl/utils.py b/vllm_fl/utils.py index 3f30bcbc8..36f3017c1 100644 --- a/vllm_fl/utils.py +++ b/vllm_fl/utils.py @@ -88,6 +88,16 @@ def use_flaggems(default: bool = True) -> bool: return value.lower() in ("true", "1") +def use_deepgemm_moe() -> bool: + """Opt-in switch to route the unquantized BF16 MoE experts through the + PPU-native DeepGEMM grouped GEMM (``deep_gemm``) instead of the default + FlagGems/Triton fused_moe. Enabled with ``VLLM_FL_MOE=deepgemm``. + + When unset, behavior is identical to the previous default path. + """ + return os.environ.get("VLLM_FL_MOE", "").strip().lower() == "deepgemm" + + def get_flag_gems_whitelist_blacklist() -> Tuple[ Optional[list[str]], Optional[list[str]] ]: