WIP: Add PPU-native DeepGEMM BF16 unquantized MoE

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,
+        )

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()

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:

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]]
 ]: