vendor vllm fused_moe

Signed-off-by: Ceng23333 <441651826@qq.com>

Author: Ceng23333

pyproject.toml

@@ -11,10 +11,10 @@ dependencies = []
 requires-python = ">=3.8"
 
 [project.optional-dependencies]
-# Same interpreter as InfiniCore built with ``--aten=y`` (see vllm_kernel_reuse_evaluation.md).
+# Optional vLLM pip stack for benchmark/compare scripts only (MoE fused path is vendored in-tree).
 vllm = ["vllm==0.19.0"]
-# MiniCPM5 fused stub + jiuge: vLLM fused MoE + helper deps (flash-attn install separately; see InfiniLM/examples/FLASH_ATTN_AND_VLLM_FUSED_MOE.md).
-vllm-fused-moe = ["vllm==0.19.0", "janus>=1.0.0", "xxhash>=3.0.0"]
+# MiniCPM5 fused MoE + jiuge: Triton + helpers (flash-attn install separately; see InfiniLM/examples/FLASH_ATTN_AND_VLLM_FUSED_MOE.md).
+vllm-fused-moe = ["triton>=3.0.0", "janus>=1.0.0", "xxhash>=3.0.0"]
 
 classifiers = [
     "Programming Language :: Python :: 3",

python/infinicore/vendor/__init__.py

@@ -0,0 +1 @@
+# Namespace for vendored third-party slices (see subpackages).

python/infinicore/vendor/vllm_fused_moe/NOTICE

@@ -0,0 +1,16 @@
+This directory vendors vLLM fused MoE Triton kernels for MiniCPM / InfiniLM.
+
+Upstream: https://github.com/vllm-project/vllm (release 0.19.x)
+License: Apache-2.0 (see SPDX headers in individual files)
+
+InfiniLM changes:
+- Register torch.ops.infinilm.outplace_fused_experts (no vllm PyPI import).
+- Replace vLLM-only CUDA helpers (moe_align_block_size, moe_sum) with PyTorch.
+- Replace silu_and_mul custom op with pure torch SwiGLU.
+- Config search honors INFINILM_TUNED_CONFIG_FOLDER (and VLLM_TUNED_CONFIG_FOLDER).
+
+Tuned configs:
+``configs/`` is intentionally empty in-repo (kernel falls back to defaults). To use
+vLLM-style JSON tuning, copy ``E=...,N=...,device_name=....json`` files from upstream
+``vllm/model_executor/layers/fused_moe/configs/`` into ``configs/``, or point
+``INFINILM_TUNED_CONFIG_FOLDER`` at such a directory.

python/infinicore/vendor/vllm_fused_moe/__init__.py

@@ -0,0 +1,12 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""vLLM-derived fused MoE (Triton) registered as ``torch.ops.infinilm.*``."""
+
+from __future__ import annotations
+
+# Import for side effects: registers torch.library fragments on ``infinilm``.
+from . import fused_moe as _fused_moe  # noqa: F401
+from .activation import MoEActivation
+from .fused_moe import fused_experts, get_config_file_name
+
+__all__ = ["MoEActivation", "fused_experts", "get_config_file_name"]

python/infinicore/vendor/vllm_fused_moe/_generate_vendor_fused_moe.py

@@ -0,0 +1,165 @@
+# Internal one-shot helper (not imported at runtime): regenerate fused_moe.py from vLLM.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from pathlib import Path
+
+
+def main() -> None:
+    repo = Path(__file__).resolve().parents[5]
+    src = (
+        repo
+        / ".venv-vllm/lib/python3.12/site-packages/vllm/model_executor/layers/fused_moe/fused_moe.py"
+    )
+    dst = Path(__file__).resolve().parent / "fused_moe.py"
+    lines = src.read_text(encoding="utf-8").splitlines(keepends=True)
+    out = "".join(lines[:1896])
+
+    reps: list[tuple[str, str]] = [
+        ("import vllm.envs as envs", "from . import envs"),
+        (
+            "import vllm.model_executor.layers.fused_moe.modular_kernel as mk\n",
+            "",
+        ),
+        ("from vllm import _custom_ops as ops\n", "from . import ops_shim as ops\n"),
+        ("from vllm.logger import init_logger", "from .logging_utils import init_logger"),
+        (
+            "from vllm.model_executor.layers.fused_moe.activation import",
+            "from .activation import",
+        ),
+        (
+            "from vllm.model_executor.layers.fused_moe.config import (\n"
+            "    FUSED_MOE_UNQUANTIZED_CONFIG,\n"
+            "    FusedMoEConfig,\n"
+            "    FusedMoEParallelConfig,\n"
+            "    FusedMoEQuantConfig,\n"
+            "    _get_config_dtype_str,\n)",
+            "from .config_light import (\n"
+            "    FUSED_MOE_UNQUANTIZED_CONFIG,\n"
+            "    FusedMoEQuantConfig,\n"
+            "    _get_config_dtype_str,\n)",
+        ),
+        (
+            "from vllm.model_executor.layers.fused_moe.moe_align_block_size import",
+            "from .moe_align_block_size import",
+        ),
+        (
+            "from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (\n"
+            "    TopKWeightAndReduceNoOP,\n)\n",
+            "",
+        ),
+        (
+            "from vllm.model_executor.layers.fused_moe.utils import",
+            "from .utils_moe import",
+        ),
+        (
+            "from vllm.model_executor.layers.quantization.utils.mxfp4_utils import dequant_mxfp4\n",
+            "",
+        ),
+        (
+            "from vllm.model_executor.layers.quantization.utils.mxfp6_utils import dequant_mxfp6\n",
+            "",
+        ),
+        (
+            "from vllm.model_executor.layers.quantization.utils.quant_utils import (\n"
+            "    QuantKey,\n"
+            "    kFp8Dynamic128Sym,\n"
+            "    kFp8DynamicTensorSym,\n"
+            "    kFp8DynamicTokenSym,\n"
+            "    kFp8Static128BlockSym,\n"
+            "    kFp8StaticChannelSym,\n"
+            "    kFp8StaticTensorSym,\n)\n",
+            "",
+        ),
+        ("from vllm.platforms import current_platform", "from .platform_ctx import current_platform"),
+        ("from vllm.triton_utils import tl, triton", "import triton\nimport triton.language as tl"),
+        (
+            "from vllm.utils.torch_utils import direct_register_custom_op",
+            "from .torch_register import direct_register_custom_op, infinilm_fused_lib",
+        ),
+    ]
+    for a, b in reps:
+        if a not in out:
+            raise SystemExit(f"missing pattern fragment:\n{a[:200]}")
+        out = out.replace(a, b)
+
+    out = out.replace(
+        "activation_out_dim = mk.FusedMoEExpertsModular.adjust_N_for_activation(\n"
+        "        N, activation_enum\n"
+        "    )",
+        "activation_out_dim = N if not activation_enum.is_gated else N // 2",
+    )
+
+    out = out.replace(
+        "    from vllm.model_executor.layers.fused_moe import get_config\n\n"
+        "    override_config = get_config()",
+        "    override_config = None",
+    )
+
+    out = out.replace("torch.ops.vllm.", "torch.ops.infinilm.")
+
+    out = out.replace(
+        "direct_register_custom_op(\n    op_name=\"inplace_fused_experts\",",
+        "direct_register_custom_op(\n    op_name=\"inplace_fused_experts\",\n"
+        "    target_lib=infinilm_fused_lib,",
+    )
+    out = out.replace(
+        "direct_register_custom_op(\n    op_name=\"outplace_fused_experts\",",
+        "direct_register_custom_op(\n    op_name=\"outplace_fused_experts\",\n"
+        "    target_lib=infinilm_fused_lib,",
+    )
+
+    needle = "def dispatch_fused_moe_kernel(\n"
+    idx = out.find(needle)
+    if idx == -1:
+        raise SystemExit("dispatch_fused_moe_kernel not found")
+    insert_at = idx + len(needle)
+    # After opening def line + newline, insert guard at body start
+    guard = (
+        "    if (use_int8_w8a16 or use_int4_w4a16) and (\n"
+        "        block_shape is not None and block_shape[1] > 0\n"
+        "    ):\n"
+        "        raise NotImplementedError(\n"
+        "            \"InfiniLM vendor fused_moe: INT4/INT8 WNA16 CUDA path requires vLLM native ops.\"\n"
+        "        )\n\n"
+    )
+    # Find first line after def that's already indented (skip docstring? none)
+    # Insert right after `) -> None:` line end - actually after `):`
+    line_end = out.find(") -> None:", idx)
+    if line_end == -1:
+        line_end = out.find("):", idx)
+    body_start = out.find("\n", line_end) + 1
+    out = out[:body_start] + guard + out[body_start:]
+
+    ocp_start = out.find("    if ocp_mx_scheme is not None:\n        # TODO: On platforms")
+    if ocp_start == -1:
+        raise SystemExit("ocp_mx block not found")
+    ocp_end = out.find("\n    qhidden_states, a1q_scale = moe_kernel_quantize_input(", ocp_start)
+    if ocp_end == -1:
+        raise SystemExit("ocp_mx block end not found")
+    out = (
+        out[:ocp_start]
+        + "    if ocp_mx_scheme is not None:\n"
+        "        raise NotImplementedError(\n"
+        "            \"InfiniLM vendor fused_moe: ocp_mx_scheme is unsupported.\"\n"
+        "        )\n\n"
+        + out[ocp_end + 1 :]
+    )
+
+    hdr = (
+        '# SPDX-License-Identifier: Apache-2.0\n# SPDX-FileCopyrightText: Copyright contributors to the vLLM project\n"""Fused MoE Triton kernels (InfiniLM vendor; see NOTICE)."""\n\n'
+    )
+    # Docstring already in file - replace opening docstring line
+    out = out.replace(
+        '"""Fused MoE Triton kernels."""\n\n',
+        '"""Fused MoE Triton kernels (InfiniLM vendor; see NOTICE)."""\n\n',
+        1,
+    )
+
+    dst.write_text(out, encoding="utf-8")
+    print("wrote", dst)
+
+
+if __name__ == "__main__":
+    main()

python/infinicore/vendor/vllm_fused_moe/activation.py

@@ -0,0 +1,71 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from enum import Enum
+
+import torch
+import torch.nn.functional as F
+
+
+class MoEActivation(Enum):
+    SILU = "silu"
+    GELU = "gelu"
+    RELU2 = "relu2"
+    SWIGLUOAI = "swigluoai"
+    SWIGLUSTEP = "swiglustep"
+    SILU_NO_MUL = "silu_no_mul"
+    GELU_NO_MUL = "gelu_no_mul"
+    RELU2_NO_MUL = "relu2_no_mul"
+
+    @property
+    def is_gated(self) -> bool:
+        return not self.value.endswith("_no_mul")
+
+    @classmethod
+    def from_str(cls, s: str) -> MoEActivation:
+        for member in cls:
+            if member.value == s:
+                return member
+        valid = [m.value for m in cls]
+        raise ValueError(f"Unknown MoE activation: {s!r}. Valid activations: {valid}")
+
+
+def apply_moe_activation(
+    activation: MoEActivation,
+    output: torch.Tensor,
+    input: torch.Tensor,
+) -> torch.Tensor:
+    assert input.dim() == 2, "Input must be 2D"
+    assert output.dim() == 2, "Output must be 2D"
+    if activation.is_gated:
+        assert output.size(-1) * 2 == input.size(-1), (
+            f"{activation.value} expects 2x ratio: "
+            f"{output.size(-1) * 2} vs {input.size(-1)}"
+        )
+    else:
+        assert output.size(-1) == input.size(-1), (
+            f"{activation.value} expects equal sizes: "
+            f"{output.size(-1)} vs {input.size(-1)}"
+        )
+
+    if activation == MoEActivation.SILU:
+        gate, up = input.chunk(2, dim=-1)
+        torch.mul(F.silu(gate), up, out=output)
+    elif activation == MoEActivation.GELU:
+        gate, up = input.chunk(2, dim=-1)
+        torch.mul(F.gelu(gate), up, out=output)
+    elif activation in (MoEActivation.SWIGLUOAI, MoEActivation.SWIGLUSTEP):
+        raise NotImplementedError(f"{activation} requires vLLM Triton ops in upstream.")
+    elif activation == MoEActivation.SILU_NO_MUL:
+        output.copy_(F.silu(input))
+    elif activation == MoEActivation.GELU_NO_MUL:
+        output.copy_(F.gelu(input))
+    elif activation == MoEActivation.RELU2_NO_MUL:
+        tmp = F.relu(input)
+        torch.square(tmp, out=output)
+    else:
+        raise ValueError(f"Unsupported FusedMoe activation: {activation}")
+
+    return output

python/infinicore/vendor/vllm_fused_moe/config_light.py

@@ -0,0 +1,65 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+
+
+def _get_config_dtype_str(
+    dtype: torch.dtype,
+    use_fp8_w8a8: bool = False,
+    use_fp8_w8a16: bool = False,
+    use_int8_w8a16: bool = False,
+    use_int4_w4a16: bool = False,
+    ocp_mx_scheme: str | None = None,
+) -> str | None:
+    if use_fp8_w8a8:
+        return "fp8_w8a8"
+    elif use_fp8_w8a16:
+        return "fp8_w8a16"
+    elif use_int8_w8a16:
+        return "int8_w8a16"
+    elif use_int4_w4a16:
+        return "int4_w4a16"
+    elif ocp_mx_scheme is not None:
+        return None
+    elif dtype == torch.float:
+        return "float32"
+    return None
+
+
+@dataclass
+class FusedMoEQuantConfig:
+    """Minimal quant bundle for fused_experts (unquantized MoE only in this vendor)."""
+
+    use_fp8_w8a8: bool = False
+    use_int8_w8a8: bool = False
+    use_int8_w8a16: bool = False
+    use_int4_w4a16: bool = False
+    ocp_mx_scheme: str | None = None
+    per_act_token_quant: bool = False
+    block_shape: list[int] | None = None
+    w1_scale: torch.Tensor | None = None
+    w2_scale: torch.Tensor | None = None
+    w1_zp: torch.Tensor | None = None
+    w2_zp: torch.Tensor | None = None
+    a1_scale: torch.Tensor | None = None
+    a2_scale: torch.Tensor | None = None
+    w1_bias: torch.Tensor | None = None
+    w2_bias: torch.Tensor | None = None
+
+    def config_name(self, dtype: torch.dtype) -> str | None:
+        return _get_config_dtype_str(
+            use_fp8_w8a8=self.use_fp8_w8a8,
+            use_fp8_w8a16=False,
+            use_int8_w8a16=self.use_int8_w8a16,
+            use_int4_w4a16=self.use_int4_w4a16,
+            ocp_mx_scheme=self.ocp_mx_scheme,
+            dtype=dtype,
+        )
+
+
+FUSED_MOE_UNQUANTIZED_CONFIG = FusedMoEQuantConfig()

python/infinicore/vendor/vllm_fused_moe/envs.py

@@ -0,0 +1,14 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import os
+
+# Prefer InfiniLM env; keep VLLM_* as fallback for existing workflows.
+INFINILM_TUNED_CONFIG_FOLDER = os.environ.get("INFINILM_TUNED_CONFIG_FOLDER")
+VLLM_TUNED_CONFIG_FOLDER = INFINILM_TUNED_CONFIG_FOLDER or os.environ.get(
+    "VLLM_TUNED_CONFIG_FOLDER"
+)
+
+VLLM_BATCH_INVARIANT = os.environ.get("VLLM_BATCH_INVARIANT", "0") == "1"