From 60f68b711be7ac5dd4535856205f3fda12474849 Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Thu, 30 Jul 2026 04:29:13 +0000 Subject: [PATCH 1/8] =?UTF-8?q?feat(vllm):=20op-injection=20plugin=20skele?= =?UTF-8?q?ton=20=E2=80=94=20Attention=20+=20MoE=20seams=20via=20entry=5Fp?= =?UTF-8?q?oints=20(#40)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Out-of-tree vLLM plugin (no fork) that will inject Infera/HyperLoom kernels: - entry_points: vllm.platform_plugins (register_platform) + vllm.general_plugins (register_ops); no-op unless ROCm is present and INFERA_VLLM_OPS_DISABLE != 1. - Attention seam: InferaPlatform(RocmPlatform).get_attn_backend_cls (pass-through). - MoE seam: install_moe_ops() — FusedMoEModularKernel experts, pass-through stub. Validated on the vLLM ROCm image: vLLM discovers `infera` in both plugin groups and `import vllm` is clean with the default install. Platform activation is opt-in (INFERA_VLLM_OPS_PLATFORM=1) pending an import-safety fix: activating the RocmPlatform subclass here re-enters vllm.platforms during its eager current_platform resolution and circular-imports. Planned fix (#40): move the attention seam to a register_ops monkey-patch of get_attn_backend_cls (ATOM's MLA approach), keeping the platform subclass only if made import-safe. Signed-off-by: Zhang, Jiejing --- infera/engine/vllm/ops/__init__.py | 17 ++++++++ infera/engine/vllm/ops/moe.py | 37 ++++++++++++++++ infera/engine/vllm/ops/platform.py | 41 ++++++++++++++++++ infera/engine/vllm/ops/register.py | 69 ++++++++++++++++++++++++++++++ pyproject.toml | 9 ++++ 5 files changed, 173 insertions(+) create mode 100644 infera/engine/vllm/ops/__init__.py create mode 100644 infera/engine/vllm/ops/moe.py create mode 100644 infera/engine/vllm/ops/platform.py create mode 100644 infera/engine/vllm/ops/register.py diff --git a/infera/engine/vllm/ops/__init__.py b/infera/engine/vllm/ops/__init__.py new file mode 100644 index 00000000..eb7cccd2 --- /dev/null +++ b/infera/engine/vllm/ops/__init__.py @@ -0,0 +1,17 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""Infera vLLM op-injection plugin (issue #40). + +Injects Infera/HyperLoom-optimized **Attention** and **MoE** kernels into stock +vLLM via vLLM's out-of-tree plugin mechanism (``vllm.platform_plugins`` + +``vllm.general_plugins``) — no vLLM fork. Both hooks are no-ops when +``INFERA_VLLM_OPS_DISABLE=1``. + +Seams: + * Attention → :class:`infera.engine.vllm.ops.platform.InferaPlatform` + (``get_attn_backend_cls``). + * MoE experts → :func:`infera.engine.vllm.ops.moe.install_moe_ops`. +""" diff --git a/infera/engine/vllm/ops/moe.py b/infera/engine/vllm/ops/moe.py new file mode 100644 index 00000000..8b998968 --- /dev/null +++ b/infera/engine/vllm/ops/moe.py @@ -0,0 +1,37 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""MoE experts injection seam (issue #40). + +vLLM runs MoE through a ``FusedMoE`` layer whose experts kernel is a +``FusedMoEModularKernel`` (``FusedMoEPrepareAndFinalize`` + +``FusedMoEPermuteExpertsUnpermute``) or a ``FusedMoEMethodBase.apply()``. +:func:`install_moe_ops` is where an Infera/HyperLoom experts kernel replaces the +default while keeping vLLM's routing / quant / EP-DP dispatch — the less-invasive +alternative to vLLM-ATOM's whole-model ``register_model`` wrapper. + +Today it is a logged pass-through so the seam is wired and observable without +changing numerics; the real kernel lands behind this function. +""" + +from __future__ import annotations + +import logging + +logger = logging.getLogger(__name__) + +_INSTALLED = False + + +def install_moe_ops() -> None: + """Install the Infera MoE experts kernel (pass-through stub for now).""" + global _INSTALLED + if _INSTALLED: + return + _INSTALLED = True + # TODO(#40): register a FusedMoEPermuteExpertsUnpermute / FusedMoEMethodBase + # implementation dispatching to the Infera experts-GEMM kernel. Until then + # this is a no-op so behaviour is unchanged while the seam is in place. + logger.info("infera-vllm-ops: MoE experts seam active (pass-through)") diff --git a/infera/engine/vllm/ops/platform.py b/infera/engine/vllm/ops/platform.py new file mode 100644 index 00000000..e4c7edb2 --- /dev/null +++ b/infera/engine/vllm/ops/platform.py @@ -0,0 +1,41 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""InferaPlatform — the **attention** injection seam (issue #40). + +Subclasses vLLM's ``RocmPlatform`` and overrides ``get_attn_backend_cls``. Today +it is a pass-through (delegates to ``RocmPlatform``, so numerics are bitwise +identical), but it is the single supported point where an Infera/HyperLoom +``AttentionBackend`` is substituted — exactly the seam vLLM-ATOM uses +(``ATOMPlatform`` → ``AiterBackend`` / ``AiterMLABackend``). + +Imported by vLLM only via the ``platform_plugins`` qualname string, so the vLLM +import below never runs outside a vLLM process. +""" + +from __future__ import annotations + +import logging + +from vllm.platforms.rocm import RocmPlatform + +logger = logging.getLogger(__name__) + +# Set to a ``"module.path:AttentionBackend"`` string to inject a custom attention +# backend; ``None`` = pass-through (vLLM's default ROCm backend selection). The +# custom class must implement vLLM's ``AttentionBackend`` / ``AttentionImpl``. +INFERA_ATTN_BACKEND: str | None = None + + +class InferaPlatform(RocmPlatform): + """ROCm platform with an attention-backend injection seam.""" + + @classmethod + def get_attn_backend_cls(cls, *args, **kwargs): + if INFERA_ATTN_BACKEND is not None: + logger.info("infera-vllm-ops: attention backend → %s", INFERA_ATTN_BACKEND) + return INFERA_ATTN_BACKEND + # Pass-through: vLLM's default ROCm attention backend, unchanged. + return super().get_attn_backend_cls(*args, **kwargs) diff --git a/infera/engine/vllm/ops/register.py b/infera/engine/vllm/ops/register.py new file mode 100644 index 00000000..7ff89e7c --- /dev/null +++ b/infera/engine/vllm/ops/register.py @@ -0,0 +1,69 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""Entry points for the Infera vLLM op-injection plugin (issue #40). + +vLLM discovers these via ``entry_points`` and calls them at startup: + + ``vllm.platform_plugins`` → :func:`register_platform` returns the qualname of + :class:`InferaPlatform` (or ``None`` to skip), whose ``get_attn_backend_cls`` + is the **attention** injection seam. + ``vllm.general_plugins`` → :func:`register_ops` installs the **MoE** experts + injection seam. + +Both are no-ops when ``INFERA_VLLM_OPS_DISABLE=1``. Kept free of top-level vLLM / +torch imports so this module is safe to import in any environment; the ROCm-only +platform is loaded lazily by vLLM via the qualname string. +""" + +from __future__ import annotations + +import logging +import os + +logger = logging.getLogger(__name__) + +_PLATFORM_QUALNAME = "infera.engine.vllm.ops.platform:InferaPlatform" + + +def _disabled() -> bool: + return os.environ.get("INFERA_VLLM_OPS_DISABLE", "0") == "1" + + +def _is_rocm() -> bool: + """Cheap ROCm probe (no torch import): the platform subclasses RocmPlatform, + so only activate where ROCm is present — leave CUDA/CPU vLLM untouched.""" + if os.environ.get("ROCM_PATH") or os.environ.get("HIP_VISIBLE_DEVICES"): + return True + import glob + + return bool(glob.glob("/opt/rocm*")) + + +def register_platform() -> str | None: + """vLLM ``platform_plugins`` hook: return InferaPlatform's qualname, else None. + + Opt-in via ``INFERA_VLLM_OPS_PLATFORM=1``. Default off (returns None) because + activating a ``RocmPlatform`` subclass here re-enters ``vllm.platforms`` while + it is still initializing (vLLM resolves ``current_platform`` eagerly during + ``import vllm``), so importing ``platform.py`` at that moment circular-imports. + The attention seam will move to an import-safe path (a ``register_ops`` + monkey-patch of ``get_attn_backend_cls``, ATOM's MLA approach) — tracked in #40. + """ + if _disabled() or not _is_rocm(): + return None + if os.environ.get("INFERA_VLLM_OPS_PLATFORM", "0") != "1": + return None + logger.info("infera-vllm-ops: activating platform %s", _PLATFORM_QUALNAME) + return _PLATFORM_QUALNAME + + +def register_ops() -> None: + """vLLM ``general_plugins`` hook: install the MoE experts injection seam.""" + if _disabled() or not _is_rocm(): + return + from infera.engine.vllm.ops.moe import install_moe_ops + + install_moe_ops() diff --git a/pyproject.toml b/pyproject.toml index 6ac21f91..9d766243 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,6 +41,15 @@ infera-kvd-l3-bench = "infera.kvd.bench.l3_bench:main" # node + PD preflight suite (gpu / network / storage / firmware / host probes) infera-preflight = "infera.tools.preflight.cli:main" +# vLLM op-injection plugin (issue #40): inject custom Attention/MoE kernels into +# stock vLLM out-of-tree — no fork. Both hooks no-op unless ROCm is present and +# INFERA_VLLM_OPS_DISABLE != 1. +[project.entry-points."vllm.platform_plugins"] +infera = "infera.engine.vllm.ops.register:register_platform" + +[project.entry-points."vllm.general_plugins"] +infera_ops = "infera.engine.vllm.ops.register:register_ops" + [build-system] requires = ["setuptools>=69", "setuptools_scm[toml]>=8", "wheel"] build-backend = "setuptools.build_meta" From f3d98ac0ffe2246112bf6a1f71a4f29470224dfd Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Thu, 30 Jul 2026 04:58:16 +0000 Subject: [PATCH 2/8] =?UTF-8?q?feat(vllm):=20import-safe=20op-injection=20?= =?UTF-8?q?=E2=80=94=20attention=20seam=20via=20general-plugin=20patch=20(?= =?UTF-8?q?#40)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the platform-subclass activation (which circular-imported during vLLM's eager current_platform resolution) with a single vllm.general_plugins hook that patches the resolved platform after vllm.platforms is initialized: - register_ops installs both seams; drop the vllm.platform_plugins entry point and InferaPlatform subclass. - attention.py: install_attention_ops() monkey-patches get_attn_backend_cls on the current platform — pass-through by default, returns INFERA_ATTN_BACKEND ("module:Backend") when set. ATOM's MLA-patch style, but import-safe. - moe.py: MoE experts seam unchanged (pass-through stub). Validated on the vLLM 0.23 ROCm image: `import vllm` is clean; the general plugin is discovered and run; the attention seam patches RocmPlatform and the override path returns the injected backend. No circular import, no fork. Signed-off-by: Zhang, Jiejing --- infera/engine/vllm/ops/__init__.py | 11 +++--- infera/engine/vllm/ops/attention.py | 58 +++++++++++++++++++++++++++++ infera/engine/vllm/ops/platform.py | 41 -------------------- infera/engine/vllm/ops/register.py | 56 +++++++++++----------------- pyproject.toml | 8 ++-- 5 files changed, 89 insertions(+), 85 deletions(-) create mode 100644 infera/engine/vllm/ops/attention.py delete mode 100644 infera/engine/vllm/ops/platform.py diff --git a/infera/engine/vllm/ops/__init__.py b/infera/engine/vllm/ops/__init__.py index eb7cccd2..68c519e3 100644 --- a/infera/engine/vllm/ops/__init__.py +++ b/infera/engine/vllm/ops/__init__.py @@ -6,12 +6,13 @@ """Infera vLLM op-injection plugin (issue #40). Injects Infera/HyperLoom-optimized **Attention** and **MoE** kernels into stock -vLLM via vLLM's out-of-tree plugin mechanism (``vllm.platform_plugins`` + -``vllm.general_plugins``) — no vLLM fork. Both hooks are no-ops when -``INFERA_VLLM_OPS_DISABLE=1``. +vLLM via a single out-of-tree ``vllm.general_plugins`` hook — no vLLM fork. The +hook (:func:`infera.engine.vllm.ops.register.register_ops`) runs after +``vllm.platforms`` is initialized and patches the resolved platform / MoE layer, +so it is import-safe. No-op when ``INFERA_VLLM_OPS_DISABLE=1`` or off-ROCm. Seams: - * Attention → :class:`infera.engine.vllm.ops.platform.InferaPlatform` - (``get_attn_backend_cls``). + * Attention → :func:`infera.engine.vllm.ops.attention.install_attention_ops` + (patches ``get_attn_backend_cls``; ``INFERA_ATTN_BACKEND`` selects a backend). * MoE experts → :func:`infera.engine.vllm.ops.moe.install_moe_ops`. """ diff --git a/infera/engine/vllm/ops/attention.py b/infera/engine/vllm/ops/attention.py new file mode 100644 index 00000000..559b255d --- /dev/null +++ b/infera/engine/vllm/ops/attention.py @@ -0,0 +1,58 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""Attention injection seam (issue #40) — import-safe monkey-patch. + +Rather than activating a ``RocmPlatform`` *subclass* through ``platform_plugins`` +(which circular-imports: vLLM resolves ``current_platform`` eagerly during +``import vllm``, and a platform module that imports ``vllm.platforms.rocm`` at top +level re-enters ``vllm.platforms`` before it finishes), we patch +``get_attn_backend_cls`` on the already-resolved platform from ``register_ops`` — +a ``general_plugins`` hook that runs *after* ``vllm.platforms`` is initialized and +*before* the model (and thus attention-backend selection) is built. Same style as +vLLM-ATOM's MLA ``forward_impl`` patch. + +Set ``INFERA_ATTN_BACKEND="module.path:AttentionBackend"`` to substitute a custom +attention backend; unset ⇒ pass-through (delegate to vLLM's default selection). +The custom class must implement vLLM's ``AttentionBackend`` / ``AttentionImpl``. +""" + +from __future__ import annotations + +import logging +import os + +logger = logging.getLogger(__name__) + +_PATCHED = False + + +def install_attention_ops() -> None: + """Patch the current platform's ``get_attn_backend_cls`` (idempotent).""" + global _PATCHED + if _PATCHED: + return + + override = os.environ.get("INFERA_ATTN_BACKEND") or None + # Safe here: general plugins load after vllm.platforms is initialized. + from vllm.platforms import current_platform + + platform_cls = type(current_platform) + original = platform_cls.get_attn_backend_cls # bound classmethod (pre-patch) + + def _patched(cls, *args, **kwargs): + if override is not None: + logger.info("infera-vllm-ops: attention backend → %s", override) + return override + # Pass-through: vLLM's default selection, unchanged. + return original(*args, **kwargs) + + platform_cls.get_attn_backend_cls = classmethod(_patched) + _PATCHED = True + logger.info( + "infera-vllm-ops: attention seam installed on %s (backend=%s)", + platform_cls.__name__, + override or "pass-through", + ) diff --git a/infera/engine/vllm/ops/platform.py b/infera/engine/vllm/ops/platform.py deleted file mode 100644 index e4c7edb2..00000000 --- a/infera/engine/vllm/ops/platform.py +++ /dev/null @@ -1,41 +0,0 @@ -############################################################################### -# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. -# -# SPDX-License-Identifier: MIT -############################################################################### -"""InferaPlatform — the **attention** injection seam (issue #40). - -Subclasses vLLM's ``RocmPlatform`` and overrides ``get_attn_backend_cls``. Today -it is a pass-through (delegates to ``RocmPlatform``, so numerics are bitwise -identical), but it is the single supported point where an Infera/HyperLoom -``AttentionBackend`` is substituted — exactly the seam vLLM-ATOM uses -(``ATOMPlatform`` → ``AiterBackend`` / ``AiterMLABackend``). - -Imported by vLLM only via the ``platform_plugins`` qualname string, so the vLLM -import below never runs outside a vLLM process. -""" - -from __future__ import annotations - -import logging - -from vllm.platforms.rocm import RocmPlatform - -logger = logging.getLogger(__name__) - -# Set to a ``"module.path:AttentionBackend"`` string to inject a custom attention -# backend; ``None`` = pass-through (vLLM's default ROCm backend selection). The -# custom class must implement vLLM's ``AttentionBackend`` / ``AttentionImpl``. -INFERA_ATTN_BACKEND: str | None = None - - -class InferaPlatform(RocmPlatform): - """ROCm platform with an attention-backend injection seam.""" - - @classmethod - def get_attn_backend_cls(cls, *args, **kwargs): - if INFERA_ATTN_BACKEND is not None: - logger.info("infera-vllm-ops: attention backend → %s", INFERA_ATTN_BACKEND) - return INFERA_ATTN_BACKEND - # Pass-through: vLLM's default ROCm attention backend, unchanged. - return super().get_attn_backend_cls(*args, **kwargs) diff --git a/infera/engine/vllm/ops/register.py b/infera/engine/vllm/ops/register.py index 7ff89e7c..5526d5e0 100644 --- a/infera/engine/vllm/ops/register.py +++ b/infera/engine/vllm/ops/register.py @@ -3,19 +3,20 @@ # # SPDX-License-Identifier: MIT ############################################################################### -"""Entry points for the Infera vLLM op-injection plugin (issue #40). +"""Entry point for the Infera vLLM op-injection plugin (issue #40). -vLLM discovers these via ``entry_points`` and calls them at startup: +vLLM discovers this via ``entry_points`` and calls it once at startup: - ``vllm.platform_plugins`` → :func:`register_platform` returns the qualname of - :class:`InferaPlatform` (or ``None`` to skip), whose ``get_attn_backend_cls`` - is the **attention** injection seam. - ``vllm.general_plugins`` → :func:`register_ops` installs the **MoE** experts - injection seam. + ``vllm.general_plugins`` → :func:`register_ops` installs both injection seams + by patching the already-resolved platform / MoE layer — a general plugin + runs *after* ``vllm.platforms`` is initialized and *before* the model is + built, so it is import-safe (unlike a ``platform_plugins`` subclass, which + re-enters ``vllm.platforms`` during vLLM's eager ``current_platform`` + resolution and circular-imports). -Both are no-ops when ``INFERA_VLLM_OPS_DISABLE=1``. Kept free of top-level vLLM / -torch imports so this module is safe to import in any environment; the ROCm-only -platform is loaded lazily by vLLM via the qualname string. +No-op unless ROCm is present and ``INFERA_VLLM_OPS_DISABLE != 1``. Kept free of +top-level vLLM / torch imports so this module is safe to import anywhere; the +seams import vLLM internals lazily, inside :func:`register_ops`. """ from __future__ import annotations @@ -25,16 +26,14 @@ logger = logging.getLogger(__name__) -_PLATFORM_QUALNAME = "infera.engine.vllm.ops.platform:InferaPlatform" - def _disabled() -> bool: return os.environ.get("INFERA_VLLM_OPS_DISABLE", "0") == "1" def _is_rocm() -> bool: - """Cheap ROCm probe (no torch import): the platform subclasses RocmPlatform, - so only activate where ROCm is present — leave CUDA/CPU vLLM untouched.""" + """Cheap ROCm probe (no torch import) — the seams target ROCm, so leave + CUDA/CPU vLLM untouched.""" if os.environ.get("ROCM_PATH") or os.environ.get("HIP_VISIBLE_DEVICES"): return True import glob @@ -42,28 +41,17 @@ def _is_rocm() -> bool: return bool(glob.glob("/opt/rocm*")) -def register_platform() -> str | None: - """vLLM ``platform_plugins`` hook: return InferaPlatform's qualname, else None. - - Opt-in via ``INFERA_VLLM_OPS_PLATFORM=1``. Default off (returns None) because - activating a ``RocmPlatform`` subclass here re-enters ``vllm.platforms`` while - it is still initializing (vLLM resolves ``current_platform`` eagerly during - ``import vllm``), so importing ``platform.py`` at that moment circular-imports. - The attention seam will move to an import-safe path (a ``register_ops`` - monkey-patch of ``get_attn_backend_cls``, ATOM's MLA approach) — tracked in #40. - """ - if _disabled() or not _is_rocm(): - return None - if os.environ.get("INFERA_VLLM_OPS_PLATFORM", "0") != "1": - return None - logger.info("infera-vllm-ops: activating platform %s", _PLATFORM_QUALNAME) - return _PLATFORM_QUALNAME - - def register_ops() -> None: - """vLLM ``general_plugins`` hook: install the MoE experts injection seam.""" - if _disabled() or not _is_rocm(): + """vLLM ``general_plugins`` hook: install the Attention + MoE injection seams.""" + if _disabled(): + logger.info("infera-vllm-ops: disabled (INFERA_VLLM_OPS_DISABLE=1)") + return + if not _is_rocm(): + logger.info("infera-vllm-ops: no ROCm detected — seams not installed") return + + from infera.engine.vllm.ops.attention import install_attention_ops from infera.engine.vllm.ops.moe import install_moe_ops + install_attention_ops() install_moe_ops() diff --git a/pyproject.toml b/pyproject.toml index 9d766243..44fe816c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,11 +42,9 @@ infera-kvd-l3-bench = "infera.kvd.bench.l3_bench:main" infera-preflight = "infera.tools.preflight.cli:main" # vLLM op-injection plugin (issue #40): inject custom Attention/MoE kernels into -# stock vLLM out-of-tree — no fork. Both hooks no-op unless ROCm is present and -# INFERA_VLLM_OPS_DISABLE != 1. -[project.entry-points."vllm.platform_plugins"] -infera = "infera.engine.vllm.ops.register:register_platform" - +# stock vLLM out-of-tree — no fork. No-op unless ROCm is present and +# INFERA_VLLM_OPS_DISABLE != 1. A general plugin (not a platform plugin) so it is +# import-safe — it patches the resolved platform/MoE after vllm.platforms inits. [project.entry-points."vllm.general_plugins"] infera_ops = "infera.engine.vllm.ops.register:register_ops" From ca937e1d10cf962bee3c450860ce25c1221195d2 Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Thu, 30 Jul 2026 05:03:43 +0000 Subject: [PATCH 3/8] =?UTF-8?q?feat(vllm):=20MoE=20experts=20seam=20?= =?UTF-8?q?=E2=80=94=20wrap=20FusedMoEKernel=20compute=20chokepoints=20(#4?= =?UTF-8?q?0)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit install_moe_ops() now patches the MoE modular-kernel compute methods (FusedMoEKernel.apply / apply_monolithic in vLLM 0.23) with a delegating pass-through — the single point where an Infera/HyperLoom experts kernel replaces the default while keeping vLLM's routing / quant / EP-DP dispatch. INFERA_MOE_EXPERTS is reserved for selecting a custom impl. Defensive against MoE-API churn: no-ops with a warning if the kernel class/methods aren't found. Validated on the vLLM 0.23 ROCm image: load_general_plugins() installs both seams; FusedMoEKernel.apply and apply_monolithic are wrapped; import stays clean. Signed-off-by: Zhang, Jiejing --- infera/engine/vllm/ops/moe.py | 75 ++++++++++++++++++++++++++++------- 1 file changed, 61 insertions(+), 14 deletions(-) diff --git a/infera/engine/vllm/ops/moe.py b/infera/engine/vllm/ops/moe.py index 8b998968..61ea2383 100644 --- a/infera/engine/vllm/ops/moe.py +++ b/infera/engine/vllm/ops/moe.py @@ -5,33 +5,80 @@ ############################################################################### """MoE experts injection seam (issue #40). -vLLM runs MoE through a ``FusedMoE`` layer whose experts kernel is a -``FusedMoEModularKernel`` (``FusedMoEPrepareAndFinalize`` + -``FusedMoEPermuteExpertsUnpermute``) or a ``FusedMoEMethodBase.apply()``. -:func:`install_moe_ops` is where an Infera/HyperLoom experts kernel replaces the -default while keeping vLLM's routing / quant / EP-DP dispatch — the less-invasive -alternative to vLLM-ATOM's whole-model ``register_model`` wrapper. - -Today it is a logged pass-through so the seam is wired and observable without -changing numerics; the real kernel lands behind this function. +vLLM runs every MoE block through a modular kernel (``FusedMoEKernel`` in vLLM +0.23; its ``apply`` / ``apply_monolithic`` route → experts-GEMM → combine). We +patch those methods from ``register_ops`` (a general plugin, so ``vllm`` is +initialized) with a pass-through wrapper — the single point where an +Infera/HyperLoom experts kernel replaces the default while keeping vLLM's +routing / quant / EP-DP dispatch. This is the less-invasive alternative to +vLLM-ATOM's whole-model ``register_model`` wrapper. + +Today it is a delegating pass-through (bitwise identical). ``INFERA_MOE_EXPERTS`` +is reserved for selecting a custom experts implementation. The MoE modular-kernel +API moves between vLLM versions, so the patch is defensive: if the class/methods +aren't found it logs and no-ops rather than raising. """ from __future__ import annotations import logging +import os logger = logging.getLogger(__name__) _INSTALLED = False +# MoE compute chokepoints to wrap, newest vLLM first. Extend as the API moves. +_KERNEL_CANDIDATES = (("vllm.model_executor.layers.fused_moe.modular_kernel", "FusedMoEKernel"),) +_METHODS = ("apply", "apply_monolithic") + def install_moe_ops() -> None: - """Install the Infera MoE experts kernel (pass-through stub for now).""" + """Wrap the MoE modular-kernel compute methods with the Infera seam (idempotent).""" global _INSTALLED if _INSTALLED: return _INSTALLED = True - # TODO(#40): register a FusedMoEPermuteExpertsUnpermute / FusedMoEMethodBase - # implementation dispatching to the Infera experts-GEMM kernel. Until then - # this is a no-op so behaviour is unchanged while the seam is in place. - logger.info("infera-vllm-ops: MoE experts seam active (pass-through)") + + import importlib + + kernel = None + for mod_name, cls_name in _KERNEL_CANDIDATES: + try: + kernel = getattr(importlib.import_module(mod_name), cls_name) + break + except (ImportError, AttributeError): + continue + if kernel is None: + logger.warning( + "infera-vllm-ops: MoE seam not wired — no known modular-kernel class on this vLLM" + ) + return + + experts = os.environ.get("INFERA_MOE_EXPERTS") or None + wrapped = [] + for name in _METHODS: + original = getattr(kernel, name, None) + if original is None or getattr(original, "_infera_wrapped", False): + continue + setattr(kernel, name, _make_seam(original)) + wrapped.append(name) + + logger.info( + "infera-vllm-ops: MoE experts seam on %s.{%s} (experts=%s)", + kernel.__name__, + ",".join(wrapped) or "", + experts or "pass-through", + ) + + +def _make_seam(original): + """Delegating pass-through — the point a custom experts kernel is dropped in.""" + + def _seam(self, *args, **kwargs): + # TODO(#40): dispatch to the Infera experts-GEMM kernel when + # INFERA_MOE_EXPERTS is set. Until then, delegate unchanged. + return original(self, *args, **kwargs) + + _seam._infera_wrapped = True + return _seam From f1eb31cd1578eb7e960129c6be30f9f6d0db6f7a Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Thu, 30 Jul 2026 07:43:05 +0000 Subject: [PATCH 4/8] feat(vllm): MoE experts op + optimize-loop microbench at Kimi-2.6 dims (#40) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Turn the MoE seam into a measurable optimize loop: - infera_fused_experts(): the plugin's swappable experts op, same signature as vLLM fused_experts. A variant registry (@register_experts_variant) lets a custom/HyperLoom kernel be selected via INFERA_MOE_EXPERTS; empty by default, so the op is a bitwise pass-through to the built-in (which on ROCm already dispatches to aiter). Falls back to built-in on any failure. - bench/op_loop/moe_experts_loop.py: measures built-in (baseline), plugin op, and a pure-torch reference (correctness oracle) on identical inputs at Kimi-2.6 MoE dims (E=384, H=7168, I=2048, top_k=8) — median latency + max abs/rel error. Measured on MI355X (vLLM 0.23 ROCm): built-in experts = 5.81 ms/call at Kimi-2.6 dims; plugin op is a verified pass-through (0 diff); the torch reference matches the built-in within bf16 tolerance (rel 5e-3) and the harness detects a 10x latency gap vs the reference. The loop is ready: a real kernel is one @register_experts_variant + re-run. Signed-off-by: Zhang, Jiejing --- bench/op_loop/moe_experts_loop.py | 129 ++++++++++++++++++++++++++++++ infera/engine/vllm/ops/moe.py | 41 ++++++++++ 2 files changed, 170 insertions(+) create mode 100644 bench/op_loop/moe_experts_loop.py diff --git a/bench/op_loop/moe_experts_loop.py b/bench/op_loop/moe_experts_loop.py new file mode 100644 index 00000000..482df710 --- /dev/null +++ b/bench/op_loop/moe_experts_loop.py @@ -0,0 +1,129 @@ +#!/usr/bin/env python3 +"""MoE-experts op optimization loop (issue #40). + +The fast inner loop for iterating custom MoE kernels behind the Infera vLLM +op-injection plugin: measure the **built-in** experts kernel (baseline), the +**plugin op** (``infera.engine.vllm.ops.moe.infera_fused_experts`` — swap the +variant with ``INFERA_MOE_EXPERTS``), and a **pure-torch reference** (correctness +oracle), on identical inputs at a real model's MoE dimensions (Kimi-2.6 by +default). Reports median latency + max abs/rel error, so a new kernel is a +one-line swap → re-run → compare. + + python moe_experts_loop.py # Kimi-2.6 dims, builtin + INFERA_MOE_EXPERTS=aiter python moe_experts_loop.py # optimized variant + python moe_experts_loop.py --experts 64 --tokens 256 # quick +""" + +import argparse +import os + +import torch + + +def build_inputs(E, H, Dm, top_k, T, dtype, device, seed=0): + g = torch.Generator(device=device).manual_seed(seed) + x = torch.randn(T, H, dtype=dtype, device=device, generator=g) * 0.1 + # w1 = [gate; up] stacked → [E, 2I, H]; w2 = down → [E, H, Dm] + w1 = ( + torch.randn(E, 2 * Dm, H, dtype=dtype, device=device, generator=g) * (H**-0.5) + ).contiguous() + w2 = (torch.randn(E, H, Dm, dtype=dtype, device=device, generator=g) * (Dm**-0.5)).contiguous() + logits = torch.randn(T, E, dtype=torch.float32, device=device, generator=g) + weights, ids = torch.topk(torch.softmax(logits, dim=-1), top_k, dim=-1) + return x, w1, w2, weights.contiguous(), ids.to(torch.int32).contiguous() + + +def reference(x, w1, w2, topk_weights, topk_ids): + """Pure-torch SwiGLU MoE (router weight applied on output) — the oracle.""" + T, H = x.shape + E, twoI, _ = w1.shape + Dm = twoI // 2 + out = torch.zeros(T, H, dtype=torch.float32, device=x.device) + xf = x.float() + for e in range(E): + sel = topk_ids == e + if not sel.any(): + continue + tok, slot = sel.nonzero(as_tuple=True) + gu = xf[tok] @ w1[e].float().t() + g, u = gu[:, :Dm], gu[:, Dm:] + o = (torch.nn.functional.silu(g) * u) @ w2[e].float().t() + out.index_add_(0, tok, o * topk_weights[tok, slot].unsqueeze(1)) + return out + + +def timed(fn, iters=30, warmup=8): + for _ in range(warmup): + fn() + torch.cuda.synchronize() + ts = [] + for _ in range(iters): + s, e = torch.cuda.Event(True), torch.cuda.Event(True) + s.record() + fn() + e.record() + torch.cuda.synchronize() + ts.append(s.elapsed_time(e)) + ts.sort() + return ts[len(ts) // 2] + + +def err(a, b): + a, b = a.float(), b.float() + denom = b.abs().max().clamp_min(1e-6) + return (a - b).abs().max().item(), ((a - b).abs().max() / denom).item() + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--experts", type=int, default=384) # Kimi-2.6 + ap.add_argument("--hidden", type=int, default=7168) + ap.add_argument("--inter", type=int, default=2048) + ap.add_argument("--topk", type=int, default=8) + ap.add_argument("--tokens", type=int, default=512) + ap.add_argument("--dtype", default="bfloat16") + ap.add_argument("--skip-ref", action="store_true") + args = ap.parse_args() + + dev = "cuda" + dt = getattr(torch, args.dtype) + E, H, Dm, top_k, T = args.experts, args.hidden, args.inter, args.topk, args.tokens + print( + f"MoE experts op @ Kimi-2.6 dims: E={E} H={H} Dm={Dm} top_k={top_k} tokens={T} dtype={args.dtype}\n" + f"variant (INFERA_MOE_EXPERTS)={os.environ.get('INFERA_MOE_EXPERTS', 'builtin')}\n" + ) + x, w1, w2, tw, ti = build_inputs(E, H, Dm, top_k, T, dt, dev) + + from vllm.model_executor.layers.fused_moe import fused_experts as builtin + + from infera.engine.vllm.ops.moe import infera_fused_experts as plugin_op + + def call(fn): + return fn(x, w1, w2, tw, ti, global_num_experts=E) + + base_out = call(builtin) + plug_out = call(plugin_op) + + rows = [] + if not args.skip_ref: + ref_out = reference(x, w1, w2, tw, ti) + rows.append( + ( + "reference (torch oracle)", + timed(lambda: reference(x, w1, w2, tw, ti), iters=5, warmup=1), + err(ref_out, base_out), + ) + ) + base_t = timed(lambda: call(builtin)) + plug_t = timed(lambda: call(plugin_op)) + rows.append(("baseline (built-in kernel)", base_t, (0.0, 0.0))) + rows.append(("plugin op (infera_fused_experts)", plug_t, err(plug_out, base_out))) + + print(f"{'impl':34s} {'median ms':>10s} {'max|Δ| vs base':>16s} {'rel':>10s} {'speedup':>9s}") + for name, ms, (ad, rd) in rows: + sp = f"{base_t / ms:.2f}x" if ms > 0 else "-" + print(f"{name:34s} {ms:10.3f} {ad:16.2e} {rd:10.2e} {sp:>9s}") + + +if __name__ == "__main__": + main() diff --git a/infera/engine/vllm/ops/moe.py b/infera/engine/vllm/ops/moe.py index 61ea2383..b05b1afa 100644 --- a/infera/engine/vllm/ops/moe.py +++ b/infera/engine/vllm/ops/moe.py @@ -82,3 +82,44 @@ def _seam(self, *args, **kwargs): _seam._infera_wrapped = True return _seam + + +# Custom experts-kernel variants register here (name -> callable with the same +# signature as vLLM's ``fused_experts``). Empty by default: the built-in kernel +# is the baseline (and on ROCm already dispatches to aiter), so a novel +# HyperLoom-tuned kernel is added with @register_experts_variant("name") and +# selected at runtime via ``INFERA_MOE_EXPERTS=name``. +_EXPERTS_VARIANTS: dict[str, object] = {} + + +def register_experts_variant(name: str): + """Decorator: register a custom MoE experts kernel under ``name``.""" + + def deco(fn): + _EXPERTS_VARIANTS[name.lower()] = fn + return fn + + return deco + + +def infera_fused_experts(*args, **kwargs): + """The plugin's swappable MoE experts op (issue #40) — the unit the + optimize-loop measures. ``INFERA_MOE_EXPERTS`` selects a registered variant; + unset / ``builtin`` (or an unknown/failed variant) uses vLLM's built-in + kernel, so the plugin op is a bitwise pass-through until a real kernel lands. + """ + from vllm.model_executor.layers.fused_moe import fused_experts as _builtin + + variant = (os.environ.get("INFERA_MOE_EXPERTS") or "builtin").lower() + fn = _EXPERTS_VARIANTS.get(variant) + if fn is None: + if variant not in ("builtin", "off", "0", ""): + logger.warning("infera-vllm-ops: no MoE experts variant %r — using builtin", variant) + return _builtin(*args, **kwargs) + try: + return fn(*args, **kwargs) + except Exception as exc: # noqa: BLE001 + logger.warning( + "infera-vllm-ops: experts variant %r failed (%s) — using builtin", variant, exc + ) + return _builtin(*args, **kwargs) From 84bdd2ee2e88cce87b450ef744e4bc615785980b Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Thu, 30 Jul 2026 12:58:17 +0000 Subject: [PATCH 5/8] =?UTF-8?q?feat(vllm):=20decode-regime=20MoE=20experts?= =?UTF-8?q?=20kernel=20=E2=80=94=201.07-1.28x=20over=20aiter=20at=20small?= =?UTF-8?q?=20M=20(#40)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit infera_decode: a small-M SwiGLU experts variant (@register_experts_variant, INFERA_MOE_EXPERTS=infera_decode) — two Triton GEMV kernels (fp32 accumulate, no tl.dot): - gate/up + SiLU fused, expert-id-indexed weight reads (no gather/scatter), - down-projection fused with the router-weighted top-k combine (no atomics). Reads each selected expert's weights once per (token, slot) pair — the traffic floor for decode — hitting ~5.2 TB/s effective HBM after tuning (block sizes env-tunable via INFERA_MOE_*). A dispatch guard delegates to the builtin where the kernel would lose or be unsafe: T > 16 or 2*T*topk > E (collision-heavy / large batch), non-bf16/fp16 weights (so Kimi MXFP4 fp4-packed weights delegate), or expert_map / non-SiLU / apply_router_weight_on_input. So it never runs where it loses. Measured (MI355X, vLLM 0.23 ROCm, Kimi-2.6 dims E=384 H=7168 I=2048 top_k=8, bf16): T=1 1.28x (0.173->0.135 ms) | T=8 1.09x | T=16 1.07x | T>16 delegates (1.00x, 0 diff) rel err vs builtin <= 5e-3 (below the fp32 torch oracle's own error vs builtin). Independently re-verified via bench/op_loop/moe_experts_loop.py. Signed-off-by: Zhang, Jiejing --- infera/engine/vllm/ops/moe.py | 226 +++++++++++++++++++++++++++++++++- 1 file changed, 222 insertions(+), 4 deletions(-) diff --git a/infera/engine/vllm/ops/moe.py b/infera/engine/vllm/ops/moe.py index b05b1afa..c01f2e16 100644 --- a/infera/engine/vllm/ops/moe.py +++ b/infera/engine/vllm/ops/moe.py @@ -85,10 +85,9 @@ def _seam(self, *args, **kwargs): # Custom experts-kernel variants register here (name -> callable with the same -# signature as vLLM's ``fused_experts``). Empty by default: the built-in kernel -# is the baseline (and on ROCm already dispatches to aiter), so a novel -# HyperLoom-tuned kernel is added with @register_experts_variant("name") and -# selected at runtime via ``INFERA_MOE_EXPERTS=name``. +# signature as vLLM's ``fused_experts``). The built-in kernel is the baseline +# (and on ROCm already dispatches to aiter); a variant is added with +# @register_experts_variant("name") and selected via ``INFERA_MOE_EXPERTS=name``. _EXPERTS_VARIANTS: dict[str, object] = {} @@ -123,3 +122,222 @@ def infera_fused_experts(*args, **kwargs): "infera-vllm-ops: experts variant %r failed (%s) — using builtin", variant, exc ) return _builtin(*args, **kwargs) + + +# --------------------------------------------------------------------------- +# "infera_decode": decode-regime (small-M) MoE experts kernel (issue #40). +# +# At M = a few tokens x top_k, the experts op is a handful of skinny GEMVs and +# is HBM-bandwidth / launch-overhead bound; the general grouped-GEMM path +# (sort/align/scatter + tl.dot tiles sized for large M) leaves a lot on the +# table. This variant runs exactly two Triton launches and reads each selected +# expert's weights exactly once — the traffic lower bound: +# +# kernel 1: per (token, slot) pair, expert e = topk_ids[t, s]: +# h[t,s,:] = SiLU(x[t] @ w1[e, :I].T) * (x[t] @ w1[e, I:].T) +# (gate/up read in one pass, fp32 accumulate, fp32 intermediate) +# kernel 2: per token, out[t] = sum_s w2[e_ts] @ h[t,s] * topk_weight[t,s] +# (down-proj + weighted combine fused, no atomics) +# +# Router weight on the *output* (apply_router_weight_on_input=False), matching +# vLLM's fused_experts / the harness oracle. bf16/fp16 weights, no quant. +# Above INFERA_MOE_DECODE_MAX_TOKENS tokens (default 16) it delegates to the +# built-in kernel: re-reading weights per pair cannot win once M is large. +# --------------------------------------------------------------------------- + +try: # Triton is present on the ROCm image; degrade gracefully elsewhere. + import triton + import triton.language as tl + + _HAS_TRITON = True +except ImportError: # pragma: no cover + _HAS_TRITON = False + + +if _HAS_TRITON: + + @triton.jit + def _infera_moe_gateup_silu( + x_ptr, # [T, H] activations + w1_ptr, # [E, 2I, H] gate;up stacked + ids_ptr, # [T*K] int expert ids + h_ptr, # [T*K, I] fp32 intermediate out + H: tl.constexpr, + I: tl.constexpr, # noqa: E741 + stride_xt, + stride_w1e, + stride_w1r, + TOPK: tl.constexpr, + BLOCK_I: tl.constexpr, + BLOCK_H: tl.constexpr, + ): + pair = tl.program_id(0) # token*TOPK + slot + pid_i = tl.program_id(1) + tok = pair // TOPK + e = tl.load(ids_ptr + pair).to(tl.int64) + ri = pid_i * BLOCK_I + tl.arange(0, BLOCK_I) # h-neuron rows + wg_ptrs = w1_ptr + e * stride_w1e + ri[:, None] * stride_w1r + wu_ptrs = wg_ptrs + I * stride_w1r + acc_g = tl.zeros((BLOCK_I,), dtype=tl.float32) + acc_u = tl.zeros((BLOCK_I,), dtype=tl.float32) + for h0 in range(0, H, BLOCK_H): + rh = h0 + tl.arange(0, BLOCK_H) + xv = tl.load(x_ptr + tok * stride_xt + rh).to(tl.float32) + wg = tl.load(wg_ptrs + rh[None, :]).to(tl.float32) + wu = tl.load(wu_ptrs + rh[None, :]).to(tl.float32) + acc_g += tl.sum(wg * xv[None, :], axis=1) + acc_u += tl.sum(wu * xv[None, :], axis=1) + h = acc_g * tl.sigmoid(acc_g) * acc_u # SiLU(gate) * up + tl.store(h_ptr + pair * I + ri, h) + + @triton.jit + def _infera_moe_down_combine( + h_ptr, # [T*K, I] fp32 intermediate + w2_ptr, # [E, H, I] down + ids_ptr, # [T*K] + tw_ptr, # [T*K] fp32 router weights + out_ptr, # [T, H] + H: tl.constexpr, + I: tl.constexpr, # noqa: E741 + stride_w2e, + stride_w2r, + stride_ot, + TOPK: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_I: tl.constexpr, + ): + tok = tl.program_id(0) + pid_h = tl.program_id(1) + rh = pid_h * BLOCK_H + tl.arange(0, BLOCK_H) # output rows + acc = tl.zeros((BLOCK_H,), dtype=tl.float32) + for slot in range(TOPK): + pair = tok * TOPK + slot + e = tl.load(ids_ptr + pair).to(tl.int64) + rw = tl.load(tw_ptr + pair).to(tl.float32) + w2_ptrs = w2_ptr + e * stride_w2e + rh[:, None] * stride_w2r + acc_e = tl.zeros((BLOCK_H,), dtype=tl.float32) + for i0 in range(0, I, BLOCK_I): + ri = i0 + tl.arange(0, BLOCK_I) + hv = tl.load(h_ptr + pair * I + ri) + wv = tl.load(w2_ptrs + ri[None, :]).to(tl.float32) + acc_e += tl.sum(wv * hv[None, :], axis=1) + acc += acc_e * rw + tl.store(out_ptr + tok * stride_ot + rh, acc.to(out_ptr.dtype.element_ty)) + + +def _decode_tunables(): + """Block sizes / warps, env-overridable for the optimize loop.""" + + def _i(name, default): + return int(os.environ.get(name, default)) + + # Defaults tuned on MI300-class (gfx942) at Kimi-2.6 dims, T=1..8: both + # kernels hit ~5.2-5.3 TB/s effective on the selected-expert weight traffic. + return ( + _i("INFERA_MOE_GU_BLOCK_I", 32), + _i("INFERA_MOE_GU_BLOCK_H", 512), + _i("INFERA_MOE_GU_WARPS", 8), + _i("INFERA_MOE_DN_BLOCK_H", 8), + _i("INFERA_MOE_DN_BLOCK_I", 512), + _i("INFERA_MOE_DN_WARPS", 4), + ) + + +@register_experts_variant("infera_decode") +def _infera_decode_experts( + hidden_states, + w1, + w2, + topk_weights, + topk_ids, + activation=None, + apply_router_weight_on_input: bool = False, + global_num_experts: int = -1, + expert_map=None, + quant_config=None, + **kwargs, +): + """Small-M fused SwiGLU experts kernel (see module comment above).""" + import torch + + if not _HAS_TRITON: + raise RuntimeError("triton not available") + # bf16/fp16 GEMV only — delegate on any quantized/packed weights (e.g. Kimi + # MXFP4 w1/w2 are fp4-packed, not float), which the builtin path handles. + if w1.dtype not in (torch.bfloat16, torch.float16) or w2.dtype != w1.dtype: + raise RuntimeError(f"unsupported weight dtype {w1.dtype}/{w2.dtype}") + if activation is not None and "silu" not in str(activation).lower(): + raise RuntimeError(f"activation {activation!r} unsupported") + if apply_router_weight_on_input: + raise RuntimeError("apply_router_weight_on_input unsupported") + if expert_map is not None: + raise RuntimeError("expert_map unsupported") + + T, H = hidden_states.shape + E, twoI, _ = w1.shape + I = twoI // 2 # noqa: E741 + K = topk_ids.shape[1] + + # This kernel reads the expert weights once per (token, slot) pair, so it + # wins while pairs are few relative to E (expert collisions across tokens + # are rare and the grouped-GEMM's dedup buys nothing). Measured on gfx942: + # win at T<=8 for E=384 (1.09-1.28x), win at T<=4 for E=64, lose beyond. + max_tokens = int(os.environ.get("INFERA_MOE_DECODE_MAX_TOKENS", "16")) + if T > max_tokens or 2 * T * K > E: # large-M / collision-heavy: delegate. + from vllm.model_executor.layers.fused_moe import fused_experts as _builtin + + return _builtin( + hidden_states, + w1, + w2, + topk_weights, + topk_ids, + apply_router_weight_on_input=apply_router_weight_on_input, + global_num_experts=global_num_experts, + expert_map=expert_map, + quant_config=quant_config, + **kwargs, + ) + + gu_bi, gu_bh, gu_w, dn_bh, dn_bi, dn_w = _decode_tunables() + if I % gu_bi or H % gu_bh or H % dn_bh or I % dn_bi: + raise RuntimeError(f"dims H={H} I={I} not divisible by block sizes") + + x = hidden_states.contiguous() + ids = topk_ids.reshape(-1).contiguous() + tw = topk_weights.reshape(-1).to(torch.float32).contiguous() + hbuf = torch.empty((T * K, I), dtype=torch.float32, device=x.device) + out = torch.empty_like(x) + + _infera_moe_gateup_silu[(T * K, I // gu_bi)]( + x, + w1, + ids, + hbuf, + H, + I, + x.stride(0), + w1.stride(0), + w1.stride(1), + TOPK=K, + BLOCK_I=gu_bi, + BLOCK_H=gu_bh, + num_warps=gu_w, + ) + _infera_moe_down_combine[(T, H // dn_bh)]( + hbuf, + w2, + ids, + tw, + out, + H, + I, + w2.stride(0), + w2.stride(1), + x.stride(0), + TOPK=K, + BLOCK_H=dn_bh, + BLOCK_I=dn_bi, + num_warps=dn_w, + ) + return out From 3d56a63998e84bbc3c25edb7e20bcd3f12b36e08 Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Thu, 30 Jul 2026 13:11:03 +0000 Subject: [PATCH 6/8] =?UTF-8?q?feat(vllm):=20close=20the=20MoE-experts=20o?= =?UTF-8?q?ptimize=20loop=20=E2=80=94=20profile=20+=20tune=20rings=20(#40)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add the two missing rings so it's a real profile → tune → inject → re-profile loop: - bench/op_loop/profile_op.py: roofline profile of an experts impl — achieved HBM BW vs peak → bandwidth-bound (near optimal) vs launch/occupancy-bound (headroom); --kernels adds the per-kernel device-time split (torch profiler). - bench/op_loop/tune_op.py: coordinate-descent autotune of the infera_decode block/warp config (keeps only correct configs, rel < 1e-2); --inject rewrites _TUNE_DEFAULTS in the plugin (the inject step) so re-profile picks it up. - moe.py: _decode_tunables() precedence per-key env > INFERA_MOE_TUNE_FILE (JSON) > baked _TUNE_DEFAULTS, so the loop drives config without code edits. - bench/op_loop/README.md: the loop, with commands. Demonstrated one full cycle on MI355X (Kimi-2.6 dims, decode T=1): profile → bandwidth-bound at 5.1 TB/s (64% of HBM peak), gate/up kernel dominates (64%); tune --inject found DN_WARPS 4->8 and baked it in (1.28x -> 1.31x); re-profile confirmed 5.23 TB/s. _TUNE_DEFAULTS updated to the tuned (32,512,8,8,512,8). Signed-off-by: Zhang, Jiejing --- bench/op_loop/README.md | 56 +++++++++++++++++ bench/op_loop/profile_op.py | 110 ++++++++++++++++++++++++++++++++++ bench/op_loop/tune_op.py | 106 ++++++++++++++++++++++++++++++++ infera/engine/vllm/ops/moe.py | 46 +++++++++----- 4 files changed, 304 insertions(+), 14 deletions(-) create mode 100644 bench/op_loop/README.md create mode 100644 bench/op_loop/profile_op.py create mode 100644 bench/op_loop/tune_op.py diff --git a/bench/op_loop/README.md b/bench/op_loop/README.md new file mode 100644 index 00000000..9f2acdbb --- /dev/null +++ b/bench/op_loop/README.md @@ -0,0 +1,56 @@ +# MoE-experts op optimize loop (issue #40) + +The closed loop for iterating custom kernels behind the [Infera vLLM op-injection +plugin](../../infera/engine/vllm/ops/): **profile → tune → inject → re-profile**, +on a real model's op dimensions (Kimi-2.6 by default). A new kernel is a +`@register_experts_variant` in `infera/engine/vllm/ops/moe.py`; these scripts +measure, diagnose, and tune it. + +Run inside the vLLM ROCm image with the repo mounted and `PYTHONPATH` set, e.g. +`docker run --device=/dev/kfd --device=/dev/dri -v :/work -w /work +-e PYTHONPATH=/work bash -lc ""`. + +## The loop + +``` + ┌────────────────────────────────────────────────┐ + ▼ │ + profile_op.py ──► tune_op.py --inject ──► (plugin updated) ┘ + where's the time? search configs, _TUNE_DEFAULTS + BW / roofline / keep correct ones, rewritten → + which kernel bake the winner in re-profile confirms +``` + +| Ring | Script | What it does | +| --- | --- | --- | +| **measure** | `moe_experts_loop.py` | built-in vs plugin-op vs torch-reference: latency + max abs/rel error. The A/B + correctness gate. | +| **profile** | `profile_op.py` | roofline: achieved HBM BW vs peak → *bandwidth-bound* (near optimal) vs *launch/occupancy-bound* (headroom); `--kernels` adds the per-kernel device-time split. | +| **tune** | `tune_op.py` | coordinate-descent over block/warp configs, keeping only correct ones; `--inject` rewrites `_TUNE_DEFAULTS` in the plugin (the 植入 step). | + +## Example cycle + +```bash +# 1. measure the current op vs the built-in (baseline) +INFERA_MOE_EXPERTS=infera_decode python moe_experts_loop.py --tokens 1 + +# 2. profile it — is there headroom, and which kernel to attack? +python profile_op.py --impl infera_decode --tokens 1 --kernels + +# 3. tune and inject the winner into the plugin +python tune_op.py --tokens 1 --inject + +# 4. re-profile to confirm (new default is now baked in) +python profile_op.py --impl infera_decode --tokens 1 +``` + +Measured on MI355X (vLLM 0.23 ROCm), Kimi-2.6 dims, decode `T=1`: the built-in +experts kernel = 0.171 ms; `infera_decode` after this loop = **0.135 ms +(1.31× )** at **~5.2 TB/s (65% of HBM peak)** — bandwidth-bound, so further wins +need less traffic (dtype / dedup), not more tuning. + +## Config precedence + +`_decode_tunables()` reads, in order: per-key env (`INFERA_MOE_GU_BLOCK_I`, …) → +`INFERA_MOE_TUNE_FILE` (JSON) → the baked `_TUNE_DEFAULTS`. So `--inject` +(rewrites `_TUNE_DEFAULTS`) makes a tuned config permanent, while env / tune-file +let you A/B without editing source. diff --git a/bench/op_loop/profile_op.py b/bench/op_loop/profile_op.py new file mode 100644 index 00000000..17149e72 --- /dev/null +++ b/bench/op_loop/profile_op.py @@ -0,0 +1,110 @@ +#!/usr/bin/env python3 +"""Profile ring of the MoE-experts optimize loop (issue #40). + +Roofline profile of an experts impl at a given shape: median latency, the +selected-expert weight traffic it must move, achieved effective HBM bandwidth, +%-of-peak, and a bottleneck verdict (bandwidth-bound ⇒ near optimal, cut traffic; +launch/occupancy-bound ⇒ headroom, tune tiling/warps). ``--kernels`` adds the +per-kernel device-time split (torch profiler) so you see which kernel to tune. + + python profile_op.py --impl infera_decode --tokens 1 + python profile_op.py --impl infera_decode --tokens 1 --kernels + python profile_op.py --impl builtin --tokens 1 # profile the baseline +""" + +import argparse +import os +import sys + +import torch + +sys.path.insert(0, os.path.dirname(__file__)) +from moe_experts_loop import build_inputs, timed # noqa: E402 + + +def weight_bytes(T, K, H, Dm, dbytes): + # The decode kernel reads each (token, slot) expert's gate/up (2·Dm·H) + down + # (H·Dm) weights once — the traffic floor for this regime. + return T * K * (2 * Dm * H + H * Dm) * dbytes + + +def peak_gbps(override): + if override: + return override + try: + p = torch.cuda.get_device_properties(0) + bw = p.memory_clock_rate * 1e3 * (p.memory_bus_width / 8) * 2 / 1e9 + if bw > 6000: # trust only a plausibly-HBM3e-class figure + return bw + except Exception: # noqa: BLE001 + pass + return 8000.0 # MI355X HBM3e ~8 TB/s; override with --peak-gbps + + +def get_impl(name): + from vllm.model_executor.layers.fused_moe import fused_experts as builtin + + if name in ("builtin", "baseline"): + return builtin + os.environ["INFERA_MOE_EXPERTS"] = name + from infera.engine.vllm.ops.moe import infera_fused_experts + + return infera_fused_experts + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--impl", default="infera_decode") + ap.add_argument("--experts", type=int, default=384) + ap.add_argument("--hidden", type=int, default=7168) + ap.add_argument("--inter", type=int, default=2048) + ap.add_argument("--topk", type=int, default=8) + ap.add_argument("--tokens", type=int, default=1) + ap.add_argument("--dtype", default="bfloat16") + ap.add_argument("--peak-gbps", type=float, default=0.0) + ap.add_argument("--kernels", action="store_true") + args = ap.parse_args() + + dt = getattr(torch, args.dtype) + E, H, Dm, K, T = args.experts, args.hidden, args.inter, args.topk, args.tokens + x, w1, w2, tw, ti = build_inputs(E, H, Dm, K, T, dt, "cuda") + impl = get_impl(args.impl) + + def call(): + return impl(x, w1, w2, tw, ti, global_num_experts=E) + + call() # warm / JIT + ms = timed(call) + wb = weight_bytes(T, K, H, Dm, torch.finfo(dt).bits // 8) + bw_gbps = (wb / 1e9) / (ms / 1e3) # GB/s + peak = peak_gbps(args.peak_gbps) # GB/s + pct = 100 * bw_gbps / peak + verdict = ( + "bandwidth-bound — near peak; win by moving less traffic (dtype / dedup)" + if pct >= 55 + else "launch/occupancy-bound — headroom; tune tiling/warps, cut launches" + ) + print(f"impl={args.impl} T={T} E={E} H={H} I={Dm} top_k={K} dtype={args.dtype}") + print(f" latency : {ms:.4f} ms") + print(f" weight traffic : {wb / 1e6:.1f} MB (each selected expert read once)") + print( + f" achieved BW : {bw_gbps / 1e3:.2f} TB/s ({pct:.0f}% of {peak / 1e3:.1f} TB/s peak)" + ) + print(f" bottleneck : {verdict}") + + if args.kernels: + from torch.profiler import ProfilerActivity, profile + + for _ in range(5): + call() + torch.cuda.synchronize() + with profile(activities=[ProfilerActivity.CUDA]) as prof: + for _ in range(20): + call() + torch.cuda.synchronize() + print("\n per-kernel device time (top 6):") + print(prof.key_averages().table(sort_by="self_device_time_total", row_limit=6)) + + +if __name__ == "__main__": + main() diff --git a/bench/op_loop/tune_op.py b/bench/op_loop/tune_op.py new file mode 100644 index 00000000..a15ec583 --- /dev/null +++ b/bench/op_loop/tune_op.py @@ -0,0 +1,106 @@ +#!/usr/bin/env python3 +"""Tune ring of the MoE-experts optimize loop (issue #40). + +Coordinate-descent autotune of the ``infera_decode`` kernel's block/warp config +at a given shape: sweeps the gate/up kernel, then the down/combine kernel, +keeping only numerically-correct configs (rel < 1e-2 vs builtin), reports the +best, and ``--inject`` bakes it into the plugin (rewrites ``_TUNE_DEFAULTS`` in +infera/engine/vllm/ops/moe.py) — the "植入" step, so re-profiling picks it up. + + python tune_op.py --tokens 1 # tune, print winner + python tune_op.py --tokens 1 --inject # tune + write into the plugin +""" + +import argparse +import os +import re +import sys + +import torch + +sys.path.insert(0, os.path.dirname(__file__)) +from moe_experts_loop import build_inputs, timed # noqa: E402 + +KEYS = ( + "INFERA_MOE_GU_BLOCK_I", + "INFERA_MOE_GU_BLOCK_H", + "INFERA_MOE_GU_WARPS", + "INFERA_MOE_DN_BLOCK_H", + "INFERA_MOE_DN_BLOCK_I", + "INFERA_MOE_DN_WARPS", +) +GU_GRID = [(bi, bh, w) for bi in (16, 32, 64) for bh in (256, 512, 1024) for w in (4, 8)] +DN_GRID = [(bh, bi, w) for bh in (8, 16, 32) for bi in (256, 512, 1024) for w in (4, 8)] + + +def measure(fn, x, w1, w2, tw, ti, E, ref): + try: + out = fn(x, w1, w2, tw, ti, global_num_experts=E) + rel = ((out.float() - ref).abs().max() / ref.abs().max().clamp_min(1e-6)).item() + if rel > 1e-2: + return None, rel + return timed(lambda: fn(x, w1, w2, tw, ti, global_num_experts=E)), rel + except Exception: # noqa: BLE001 + return None, float("inf") + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--experts", type=int, default=384) + ap.add_argument("--hidden", type=int, default=7168) + ap.add_argument("--inter", type=int, default=2048) + ap.add_argument("--topk", type=int, default=8) + ap.add_argument("--tokens", type=int, default=1) + ap.add_argument("--dtype", default="bfloat16") + ap.add_argument("--inject", action="store_true") + args = ap.parse_args() + + dt = getattr(torch, args.dtype) + E, H, Dm, K, T = args.experts, args.hidden, args.inter, args.topk, args.tokens + x, w1, w2, tw, ti = build_inputs(E, H, Dm, K, T, dt, "cuda") + from vllm.model_executor.layers.fused_moe import fused_experts as builtin + + from infera.engine.vllm.ops.moe import _EXPERTS_VARIANTS, _TUNE_DEFAULTS + + fn = _EXPERTS_VARIANTS["infera_decode"] + ref = builtin(x, w1, w2, tw, ti, global_num_experts=E).float() + base_t = timed(lambda: builtin(x, w1, w2, tw, ti, global_num_experts=E)) + + cfg = list(_TUNE_DEFAULTS) # start from current baked defaults + + def setenv(c): + for k, v in zip(KEYS, c): + os.environ[k] = str(v) + + def best_over(grid, slot): # slot: 0 for GU (keys 0:3), 3 for DN (keys 3:6) + best, best_ms = None, float("inf") + for combo in grid: + trial = cfg.copy() + trial[slot : slot + 3] = list(combo) + setenv(trial) + ms, rel = measure(fn, x, w1, w2, tw, ti, E, ref) + if ms is not None and ms < best_ms: + best, best_ms = combo, ms + return best, best_ms + + print(f"tuning infera_decode @ T={T} E={E} H={H} I={Dm} top_k={K}; baseline {base_t:.4f} ms") + gu, _ = best_over(GU_GRID, 0) + cfg[0:3] = list(gu) + dn, ms = best_over(DN_GRID, 3) + cfg[3:6] = list(dn) + setenv(cfg) + ms, rel = measure(fn, x, w1, w2, tw, ti, E, ref) + print(f" default {tuple(_TUNE_DEFAULTS)} ") + print(f" best {tuple(cfg)} -> {ms:.4f} ms ({base_t / ms:.2f}x vs builtin, rel {rel:.1e})") + + if args.inject: + path = os.path.join(os.path.dirname(__file__), "..", "..", "infera/engine/vllm/ops/moe.py") + path = os.path.abspath(path) + src = open(path).read() + new = re.sub(r"_TUNE_DEFAULTS = \([^)]*\)", f"_TUNE_DEFAULTS = {tuple(cfg)}", src, count=1) + open(path, "w").write(new) + print(f" injected _TUNE_DEFAULTS = {tuple(cfg)} into {path}") + + +if __name__ == "__main__": + main() diff --git a/infera/engine/vllm/ops/moe.py b/infera/engine/vllm/ops/moe.py index c01f2e16..7d677a06 100644 --- a/infera/engine/vllm/ops/moe.py +++ b/infera/engine/vllm/ops/moe.py @@ -225,21 +225,39 @@ def _infera_moe_down_combine( tl.store(out_ptr + tok * stride_ot + rh, acc.to(out_ptr.dtype.element_ty)) +_TUNE_KEYS = ( + "INFERA_MOE_GU_BLOCK_I", + "INFERA_MOE_GU_BLOCK_H", + "INFERA_MOE_GU_WARPS", + "INFERA_MOE_DN_BLOCK_H", + "INFERA_MOE_DN_BLOCK_I", + "INFERA_MOE_DN_WARPS", +) +# Tuned on gfx942/gfx950 at Kimi-2.6 dims, T=1..8 (~5.2-5.3 TB/s effective on the +# selected-expert weight traffic). The tune ring rewrites this tuple in place +# (tune_op.py --inject), or supply INFERA_MOE_TUNE_FILE (JSON) / per-key env. +_TUNE_DEFAULTS = (32, 512, 8, 8, 512, 8) + + +def _load_tune_file() -> dict: + path = os.environ.get("INFERA_MOE_TUNE_FILE") + if not path: + return {} + try: + import json + + with open(path) as f: + return json.load(f) + except Exception: # noqa: BLE001 + return {} + + def _decode_tunables(): - """Block sizes / warps, env-overridable for the optimize loop.""" - - def _i(name, default): - return int(os.environ.get(name, default)) - - # Defaults tuned on MI300-class (gfx942) at Kimi-2.6 dims, T=1..8: both - # kernels hit ~5.2-5.3 TB/s effective on the selected-expert weight traffic. - return ( - _i("INFERA_MOE_GU_BLOCK_I", 32), - _i("INFERA_MOE_GU_BLOCK_H", 512), - _i("INFERA_MOE_GU_WARPS", 8), - _i("INFERA_MOE_DN_BLOCK_H", 8), - _i("INFERA_MOE_DN_BLOCK_I", 512), - _i("INFERA_MOE_DN_WARPS", 4), + """Block sizes / warps. Precedence: per-key env > tune file > baked defaults, + so profile→tune→inject→profile can drive the config without code edits.""" + fromfile = _load_tune_file() + return tuple( + int(os.environ.get(k, fromfile.get(k, d))) for k, d in zip(_TUNE_KEYS, _TUNE_DEFAULTS) ) From 09c4369c7ced60c9b58a603f4b55559c4d4b69e7 Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Thu, 30 Jul 2026 15:29:31 +0000 Subject: [PATCH 7/8] feat(vllm): wire infera_decode into the serve path + e2e decode win (#40) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Make the MoE seam actually run the selected variant in a real serve (it was a pass-through), and add a weight-layout gate so it's safe regardless of the aiter flag: - _make_seam intercepts the modular FusedMoEKernel.apply: when a variant is selected and the request is one it supports, run it and return the combined [T,H] output; otherwise delegate to the untouched original. - Self-gating (never a regression, never garbage): delegates unless plain (non-aiter-shuffled) bf16/fp16 weights + small M + unquantized + no expert_map/EP + no shared-expert overlap + modular apply. The key gate is _weights_are_aiter_shuffled() — aiter pre-shuffles MoE weights in-place with no per-tensor marker, so we read the aiter master flag (INFERA_MOE_ASSUME_PLAIN=1 overrides). This is what the earlier "garbage with aiter on" was; now it cleanly delegates. - A fire counter (INFERA_MOE_DECODE_DEBUG) proves the kernel actually runs in a serve, not silently delegates. - bench/_moe_decode_e2e_{serve.sh,client.py}: single-stream batch-1 decode ITL. E2E (Qwen3.5-35B-A3B bf16 MoE, MI355X, TP=1, batch-1 decode), independently reproduced (5 reps, no overlap): aiter-off + infera_decode = 5.26 ms/tok vs the aiter-on production default 5.75 ms/tok — the fastest config, +8.5% ITL / +9% tok/s (aiter gives no batch-1 decode benefit here, so disabling it is free and the MoE kernel then wins). Confined to bf16 / batch<=16 / plain weights, all self-delegating. Details in bench/op_loop/README.md. Signed-off-by: Zhang, Jiejing --- bench/_moe_decode_e2e_client.py | 89 ++++++++++++++++++ bench/_moe_decode_e2e_serve.sh | 34 +++++++ bench/op_loop/README.md | 27 ++++++ infera/engine/vllm/ops/moe.py | 157 ++++++++++++++++++++++++++++---- 4 files changed, 291 insertions(+), 16 deletions(-) create mode 100644 bench/_moe_decode_e2e_client.py create mode 100644 bench/_moe_decode_e2e_serve.sh diff --git a/bench/_moe_decode_e2e_client.py b/bench/_moe_decode_e2e_client.py new file mode 100644 index 00000000..257323e6 --- /dev/null +++ b/bench/_moe_decode_e2e_client.py @@ -0,0 +1,89 @@ +#!/usr/bin/env python3 +"""Single-stream decode ITL benchmark for the infera_decode MoE experiment (#40). + +Batch-1, temperature 0, stream=True, measure inter-token latency (ITL). Reports +median/mean ITL (ms/tok) and decode throughput (tok/s) over the streamed tokens, +discarding the first token (TTFT / prefill) so we measure decode steps only. +""" + +import argparse +import json +import statistics +import time +import urllib.request + + +def run(port, prompt, max_tokens, warmup): + url = f"http://127.0.0.1:{port}/v1/completions" + body = { + "model": "qwen35", + "prompt": prompt, + "max_tokens": max_tokens, + "temperature": 0.0, + "stream": True, + } + data = json.dumps(body).encode() + req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"}) + t_prev = None + itls = [] + n = 0 + t0 = time.perf_counter() + with urllib.request.urlopen(req) as resp: + for raw in resp: + line = raw.decode().strip() + if not line or not line.startswith("data:"): + continue + payload = line[len("data:") :].strip() + if payload == "[DONE]": + break + obj = json.loads(payload) + txt = obj["choices"][0].get("text", "") + if txt == "": + continue + now = time.perf_counter() + if t_prev is not None: + itls.append((now - t_prev) * 1000.0) + t_prev = now + n += 1 + wall = time.perf_counter() - t0 + # discard warmup decode steps + body_itls = itls[warmup:] if len(itls) > warmup else itls + return { + "tokens": n, + "wall_s": wall, + "median_itl_ms": statistics.median(body_itls) if body_itls else float("nan"), + "mean_itl_ms": statistics.mean(body_itls) if body_itls else float("nan"), + "p10_itl_ms": statistics.quantiles(body_itls, n=10)[0] + if len(body_itls) > 10 + else float("nan"), + "decode_toks": len(body_itls), + "decode_tok_s": (len(body_itls) / (sum(body_itls) / 1000.0)) if body_itls else float("nan"), + } + + +if __name__ == "__main__": + ap = argparse.ArgumentParser() + ap.add_argument("--port", type=int, default=8012) + ap.add_argument("--max-tokens", type=int, default=256) + ap.add_argument("--warmup", type=int, default=8, help="decode steps to discard") + ap.add_argument("--reps", type=int, default=3) + ap.add_argument("--label", default="") + args = ap.parse_args() + prompt = ( + "You are a helpful assistant. Write a detailed, step-by-step explanation " + "of how a modern GPU executes a matrix multiplication, covering memory " + "hierarchy, tiling, and warp scheduling. Be thorough and precise.\n\nAnswer:" + ) + results = [] + for r in range(args.reps): + res = run(args.port, prompt, args.max_tokens, args.warmup) + results.append(res) + print( + f"[{args.label}] rep{r}: tokens={res['tokens']} " + f"median_itl={res['median_itl_ms']:.3f}ms mean_itl={res['mean_itl_ms']:.3f}ms " + f"decode_tok/s={res['decode_tok_s']:.2f} (n_decode={res['decode_toks']})" + ) + # aggregate across reps on the per-rep medians + med = statistics.median([r["median_itl_ms"] for r in results]) + tps = statistics.median([r["decode_tok_s"] for r in results]) + print(f"[{args.label}] AGG median_itl={med:.3f}ms decode_tok/s={tps:.2f}") diff --git a/bench/_moe_decode_e2e_serve.sh b/bench/_moe_decode_e2e_serve.sh new file mode 100644 index 00000000..cdc280f7 --- /dev/null +++ b/bench/_moe_decode_e2e_serve.sh @@ -0,0 +1,34 @@ +#!/usr/bin/env bash +# Dedicated serve driver for the infera_decode MoE-experts e2e experiment (issue #40). +# Usage: _moe_decode_e2e_serve.sh +set -euo pipefail +MODE="${1:?kernel|baseline}" +PORT="${2:-8012}" +MODEL=/mnt/vast/john/huggingface/Qwen3.5-35B-A3B + +export PYTHONPATH=/mnt/vast/jiejing/workspace/Optimus +export INFERA_MOE_DECODE_DEBUG=1 +export INFERA_MOE_FIRE_FILE=/tmp/infera_moe_fires.txt +export VLLM_ROCM_USE_AITER=0 # plain (unshuffled) MoE weight layout +export INFERA_MOE_DECODE_MAX_TOKENS=16 + +if [[ "$MODE" == "kernel" ]]; then + export INFERA_MOE_EXPERTS=infera_decode + export INFERA_VLLM_OPS_DISABLE=0 +else + # baseline: builtin experts (aiter). Seam installs but selects no variant. + export INFERA_MOE_EXPERTS=builtin + export INFERA_VLLM_OPS_DISABLE=0 +fi + +cd /root +exec python -m vllm.entrypoints.openai.api_server \ + --model "$MODEL" \ + --served-model-name qwen35 \ + --tensor-parallel-size 1 \ + --trust-remote-code \ + --gpu-memory-utilization 0.85 \ + --max-model-len 8192 \ + --max-num-seqs 8 \ + --no-enable-log-requests \ + --port "$PORT" diff --git a/bench/op_loop/README.md b/bench/op_loop/README.md index 9f2acdbb..b0d14086 100644 --- a/bench/op_loop/README.md +++ b/bench/op_loop/README.md @@ -48,6 +48,33 @@ experts kernel = 0.171 ms; `infera_decode` after this loop = **0.135 ms (1.31× )** at **~5.2 TB/s (65% of HBM peak)** — bandwidth-bound, so further wins need less traffic (dtype / dedup), not more tuning. +## End-to-end (does the op win show up in a serve?) + +`../_moe_decode_e2e_serve.sh` + `../_moe_decode_e2e_client.py` measure single-stream +(batch-1) decode ITL with the kernel wired into the serving path (the plugin's MoE +seam runs `infera_decode` on genuine small-M bf16 decode steps, verified via a +fire counter). Measured on **Qwen3.5-35B-A3B (bf16 MoE)**, MI355X, TP=1: + +| config | ITL (ms/tok) | decode tok/s | +| --- | --- | --- | +| aiter **on** — production default | 5.754 | 173.3 | +| aiter off — builtin (triton) | 5.650 | 176.7 | +| **aiter off — `infera_decode`** | **5.264** | **188.8** | + +So `infera_decode` is the **fastest config, ~8.5% better ITL than the aiter-on +default** — because at batch-1 aiter gives no decode benefit here (slightly +negative), so turning it off is free and the MoE kernel then wins. The op-level +~1.3× dilutes to single-digit e2e because routed experts are ~⅓ of the +bandwidth-bound batch-1 decode step. + +**Limitations (by design, all self-delegating — never a regression):** +- **bf16/fp16 only** — MXFP4/quantized weights delegate (so Kimi-2.6 MXFP4 sees no effect). +- **batch ≤ 16** — larger batches delegate and aiter's large-M wins take over. +- **plain weight layout** — aiter pre-shuffles MoE weights, so the seam delegates when + `VLLM_ROCM_USE_AITER` is set; the kernel runs with aiter's MoE path off (or + `INFERA_MOE_ASSUME_PLAIN=1`). Mixing aiter-prefill with this decode kernel needs + weight un-shuffling — future work. + ## Config precedence `_decode_tunables()` reads, in order: per-key env (`INFERA_MOE_GU_BLOCK_I`, …) → diff --git a/infera/engine/vllm/ops/moe.py b/infera/engine/vllm/ops/moe.py index 7d677a06..00f40e48 100644 --- a/infera/engine/vllm/ops/moe.py +++ b/infera/engine/vllm/ops/moe.py @@ -7,16 +7,24 @@ vLLM runs every MoE block through a modular kernel (``FusedMoEKernel`` in vLLM 0.23; its ``apply`` / ``apply_monolithic`` route → experts-GEMM → combine). We -patch those methods from ``register_ops`` (a general plugin, so ``vllm`` is -initialized) with a pass-through wrapper — the single point where an -Infera/HyperLoom experts kernel replaces the default while keeping vLLM's -routing / quant / EP-DP dispatch. This is the less-invasive alternative to -vLLM-ATOM's whole-model ``register_model`` wrapper. - -Today it is a delegating pass-through (bitwise identical). ``INFERA_MOE_EXPERTS`` -is reserved for selecting a custom experts implementation. The MoE modular-kernel -API moves between vLLM versions, so the patch is defensive: if the class/methods -aren't found it logs and no-ops rather than raising. +patch ``apply`` from ``register_ops`` (a general plugin, so ``vllm`` is +initialized) — the single point where an Infera/HyperLoom experts kernel replaces +the default while keeping vLLM's routing / quant / EP-DP dispatch. Less invasive +than vLLM-ATOM's whole-model ``register_model`` wrapper. + +``INFERA_MOE_EXPERTS=`` selects a registered variant; when the request is +one the variant supports the seam runs it and returns the combined ``[T, H]`` +output, otherwise it delegates to the untouched original. The seam is safe and +self-gating — it delegates (never garbage, never a regression) unless ALL hold: +a variant is selected, plain (non-aiter-shuffled) bf16/fp16 weights +(see ``_weights_are_aiter_shuffled``), small M (the ``infera_decode`` guard), +unquantized, no expert_map/EP, no shared-expert overlap, modular (topk) ``apply``. +The modular-kernel API moves between vLLM versions, so patching is defensive. + +Scope of the first kernel (``infera_decode``): a **decode-step specialist** — it +helps single-stream / low-concurrency **bf16** MoE decode (measured e2e below), +and delegates everywhere else. Not an aiter replacement, and (until weight +un-shuffling is added) it runs only with aiter's MoE path off / plain weights. """ from __future__ import annotations @@ -61,7 +69,7 @@ def install_moe_ops() -> None: original = getattr(kernel, name, None) if original is None or getattr(original, "_infera_wrapped", False): continue - setattr(kernel, name, _make_seam(original)) + setattr(kernel, name, _make_seam(original, name)) wrapped.append(name) logger.info( @@ -72,12 +80,99 @@ def install_moe_ops() -> None: ) -def _make_seam(original): - """Delegating pass-through — the point a custom experts kernel is dropped in.""" +# Count of genuine infera_decode Triton executions (set inside the variant, past +# every guard/delegate). Verifies the kernel actually fires during a serve — a +# silent delegate would otherwise give a false "no change". When +# INFERA_MOE_DECODE_DEBUG=1: the first fire prints to stderr and, if +# INFERA_MOE_FIRE_FILE is set, the running count is written there periodically. +_KERNEL_FIRE_COUNT = 0 + + +def _selected_variant(): + """The custom experts callable selected by ``INFERA_MOE_EXPERTS`` (or None for + the built-in aiter/modular path).""" + name = (os.environ.get("INFERA_MOE_EXPERTS") or "builtin").lower() + if name in ("builtin", "off", "0", ""): + return None + return _EXPERTS_VARIANTS.get(name) + + +def _weights_are_aiter_shuffled() -> bool: + """Weight-layout state on the interface: when vLLM's aiter path is active it + pre-shuffles the MoE weights into aiter's private layout at load + (``rocm_aiter_ops.shuffle_weights`` in fused_moe/oracle/*), an in-place + ``.data`` swap with no per-tensor marker — so a plain-layout kernel would + silently misread them (garbage, no exception). The aiter master switch is the + reliable signal; when set, the seam MUST delegate. `INFERA_MOE_ASSUME_PLAIN=1` + overrides (you asserted plain weights, e.g. a custom load path).""" + if os.environ.get("INFERA_MOE_ASSUME_PLAIN") == "1": + return False + try: + import vllm.envs as envs + + return bool(getattr(envs, "VLLM_ROCM_USE_AITER", False)) + except Exception: # noqa: BLE001 + return False + + +def _make_seam(original, method_name): + """MoE-experts injection seam. + + For the modular ``apply`` (topk_weights/topk_ids based) we intercept genuine + small-M **decode** steps and run the selected custom experts variant, which + returns the fully combined ``[T, H]`` MoE output (router weights applied on + output) — exactly what the modular flow returns before the downstream TP + all-reduce. prepare/finalize is a no-op identity for bf16-unquantized + single-node TP, so bypassing it is equivalent. Everything the variant does + not support (large M, quantized weights, expert_map/EP, shared-expert + overlap, non-SiLU, monolithic router+experts) delegates to the untouched + original so prefill / large batches keep the aiter path. + """ + + if method_name != "apply": + # apply_monolithic fuses routing (router_logits, no topk) — not something + # the decode experts kernel handles. Leave untouched. + def _passthrough(self, *args, **kwargs): + return original(self, *args, **kwargs) + + _passthrough._infera_wrapped = True + return _passthrough def _seam(self, *args, **kwargs): - # TODO(#40): dispatch to the Infera experts-GEMM kernel when - # INFERA_MOE_EXPERTS is set. Until then, delegate unchanged. + variant = _selected_variant() + if variant is not None: + # FusedMoEKernel.apply is always called by keyword (see + # unquantized_fused_moe_method / fused_moe_modular_method). + hidden_states = kwargs.get("hidden_states") + # When can_overlap_shared_experts is False (non-async prepare/finalize, + # e.g. single-node TP / NoEP), the modular _finalize does NOT run the + # shared experts — the runner computes and combines them OUTSIDE the + # kernel (SharedExpertsOrder.NO_OVERLAP, see moe_runner._apply_quant_ + # method). So the shared_experts handed to apply here are inert and we + # may replace the routed-experts compute. If overlap is active, the + # kernel owns the shared-expert compute — delegate to stay correct. + can_overlap = bool(getattr(self, "can_overlap_shared_experts", False)) + # Weight-layout gate: the kernel reads plain [E, 2I, H]/[E, H, I] + # tensors; if aiter pre-shuffled them at load, delegate (else garbage). + if hidden_states is not None and not can_overlap and not _weights_are_aiter_shuffled(): + try: + out = variant( + hidden_states, + kwargs["w1"], + kwargs["w2"], + kwargs["topk_weights"], + kwargs["topk_ids"], + activation=kwargs.get("activation"), + apply_router_weight_on_input=kwargs.get( + "apply_router_weight_on_input", False + ), + global_num_experts=kwargs.get("global_num_experts", -1), + expert_map=kwargs.get("expert_map"), + _infera_delegate=lambda: original(self, *args, **kwargs), + ) + return out + except Exception as exc: # noqa: BLE001 + logger.warning("infera-vllm-ops: decode experts seam fell back (%s)", exc) return original(self, *args, **kwargs) _seam._infera_wrapped = True @@ -273,9 +368,16 @@ def _infera_decode_experts( global_num_experts: int = -1, expert_map=None, quant_config=None, + _infera_delegate=None, **kwargs, ): - """Small-M fused SwiGLU experts kernel (see module comment above).""" + """Small-M fused SwiGLU experts kernel (see module comment above). + + ``_infera_delegate`` (when supplied by the serve seam) re-runs the original + modular MoE flow — i.e. the untouched aiter path — for anything this kernel + does not handle, instead of the standalone ``fused_experts`` fallback used by + the offline microbench. + """ import torch if not _HAS_TRITON: @@ -302,6 +404,8 @@ def _infera_decode_experts( # win at T<=8 for E=384 (1.09-1.28x), win at T<=4 for E=64, lose beyond. max_tokens = int(os.environ.get("INFERA_MOE_DECODE_MAX_TOKENS", "16")) if T > max_tokens or 2 * T * K > E: # large-M / collision-heavy: delegate. + if _infera_delegate is not None: + return _infera_delegate() from vllm.model_executor.layers.fused_moe import fused_experts as _builtin return _builtin( @@ -321,6 +425,27 @@ def _infera_decode_experts( if I % gu_bi or H % gu_bh or H % dn_bh or I % dn_bi: raise RuntimeError(f"dims H={H} I={I} not divisible by block sizes") + # Verification counter: incremented only here, where the Triton kernels + # actually launch (past every delegate/guard) — so it counts genuine decode + # kernel executions, not seam entries or delegated prefill steps. + global _KERNEL_FIRE_COUNT + _KERNEL_FIRE_COUNT += 1 + if os.environ.get("INFERA_MOE_DECODE_DEBUG") == "1": + if _KERNEL_FIRE_COUNT == 1: + import sys + + sys.stderr.write( + f"[infera-moe] infera_decode kernel FIRST FIRE (T={T} K={K} E={E} H={H} I={I})\n" + ) + sys.stderr.flush() + fire_file = os.environ.get("INFERA_MOE_FIRE_FILE") + if fire_file and _KERNEL_FIRE_COUNT % 500 == 0: + try: + with open(fire_file, "w") as _f: + _f.write(str(_KERNEL_FIRE_COUNT)) + except Exception: # noqa: BLE001 + pass + x = hidden_states.contiguous() ids = topk_ids.reshape(-1).contiguous() tw = topk_weights.reshape(-1).to(torch.float32).contiguous() From 3c3557848720880d44326d29e7ba824ee7bb1738 Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Thu, 30 Jul 2026 15:52:48 +0000 Subject: [PATCH 8/8] refactor(op-loop): op-agnostic optimize-loop scaffold (#40) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Generalize the loop from three MoE-specific scripts into a reusable scaffold — the point is the framework, not any one kernel: - framework.py: OpSpec + registry + generic measure/profile/tune. An op declares make_inputs / baseline / candidate / reference / traffic_bytes / tune_*; the loop drives all three rings from that, skipping whatever an op doesn't provide. - loop.py: one CLI over any registered op — `loop.py {measure,profile,tune,list} --op [-d k=v] [--kernels] [--inject]`. - ops/moe_experts.py: the MoE spec (the worked example / template). Adding an op is now writing one OpSpec file + register_op — no new script. - Remove moe_experts_loop.py / profile_op.py / tune_op.py (folded in). - README rewritten around the scaffold + an "adding an op" table. Re-validated via the new CLI on MI355X (Kimi-2.6 dims, T=1): measure = 1.29x candidate vs builtin (oracle rel 3.9e-3), profile = bandwidth-bound at 5.1 TB/s (64% peak), tune = 1.33x. infera_decode is unchanged — just the candidate the moe_experts spec points at. Signed-off-by: Zhang, Jiejing --- bench/op_loop/README.md | 112 +++++++++--------- bench/op_loop/framework.py | 187 ++++++++++++++++++++++++++++++ bench/op_loop/loop.py | 57 +++++++++ bench/op_loop/moe_experts_loop.py | 129 --------------------- bench/op_loop/ops/__init__.py | 1 + bench/op_loop/ops/moe_experts.py | 116 ++++++++++++++++++ bench/op_loop/profile_op.py | 110 ------------------ bench/op_loop/tune_op.py | 106 ----------------- 8 files changed, 413 insertions(+), 405 deletions(-) create mode 100644 bench/op_loop/framework.py create mode 100644 bench/op_loop/loop.py delete mode 100644 bench/op_loop/moe_experts_loop.py create mode 100644 bench/op_loop/ops/__init__.py create mode 100644 bench/op_loop/ops/moe_experts.py delete mode 100644 bench/op_loop/profile_op.py delete mode 100644 bench/op_loop/tune_op.py diff --git a/bench/op_loop/README.md b/bench/op_loop/README.md index b0d14086..e8732d14 100644 --- a/bench/op_loop/README.md +++ b/bench/op_loop/README.md @@ -1,59 +1,67 @@ -# MoE-experts op optimize loop (issue #40) +# Op optimize-loop scaffold (issue #40) -The closed loop for iterating custom kernels behind the [Infera vLLM op-injection -plugin](../../infera/engine/vllm/ops/): **profile → tune → inject → re-profile**, -on a real model's op dimensions (Kimi-2.6 by default). A new kernel is a -`@register_experts_variant` in `infera/engine/vllm/ops/moe.py`; these scripts -measure, diagnose, and tune it. +An **op-agnostic** `measure → profile → tune → inject` loop for iterating custom +kernels behind the [Infera vLLM op-injection plugin](../../infera/engine/vllm/ops/). +The scaffold is the deliverable; a kernel (e.g. `infera_decode`) is just a +candidate plugged in. **Adding an op is writing one `OpSpec` — no new script.** -Run inside the vLLM ROCm image with the repo mounted and `PYTHONPATH` set, e.g. +``` +framework.py OpSpec + registry + generic measure/profile/tune +loop.py one CLI over any registered op +ops/.py one OpSpec per op (self-registers). moe_experts.py = the template. +``` + +Run inside the vLLM ROCm image with the repo mounted and `PYTHONPATH` set (so the +edited plugin is used): `docker run --device=/dev/kfd --device=/dev/dri -v :/work -w /work --e PYTHONPATH=/work bash -lc ""`. +-e PYTHONPATH=/work bash -lc "cd bench/op_loop && "`. ## The loop -``` - ┌────────────────────────────────────────────────┐ - ▼ │ - profile_op.py ──► tune_op.py --inject ──► (plugin updated) ┘ - where's the time? search configs, _TUNE_DEFAULTS - BW / roofline / keep correct ones, rewritten → - which kernel bake the winner in re-profile confirms +```bash +python loop.py list # registered ops +python loop.py measure --op moe_experts # baseline vs candidate vs oracle +python loop.py profile --op moe_experts --kernels # roofline + per-kernel split +python loop.py tune --op moe_experts --inject # autotune, bake winner into plugin +python loop.py measure --op moe_experts -d tokens=1 -d experts=64 # override dims ``` -| Ring | Script | What it does | -| --- | --- | --- | -| **measure** | `moe_experts_loop.py` | built-in vs plugin-op vs torch-reference: latency + max abs/rel error. The A/B + correctness gate. | -| **profile** | `profile_op.py` | roofline: achieved HBM BW vs peak → *bandwidth-bound* (near optimal) vs *launch/occupancy-bound* (headroom); `--kernels` adds the per-kernel device-time split. | -| **tune** | `tune_op.py` | coordinate-descent over block/warp configs, keeping only correct ones; `--inject` rewrites `_TUNE_DEFAULTS` in the plugin (the 植入 step). | +| Ring | What it does | +| --- | --- | +| **measure** | baseline (built-in) vs candidate (plugin op) vs reference (oracle): latency + rel error. The A/B + correctness gate. | +| **profile** | roofline from `traffic_bytes`: achieved HBM BW vs peak → bandwidth-bound (near optimal) vs launch/occupancy-bound (headroom); `--kernels` adds the per-kernel device-time split. | +| **tune** | sweep the op's `tune_env` configs, keep only correct ones; `--inject` calls the op's `inject` to bake the winner into the plugin. | -## Example cycle +## Adding an op -```bash -# 1. measure the current op vs the built-in (baseline) -INFERA_MOE_EXPERTS=infera_decode python moe_experts_loop.py --tokens 1 +Write `ops/.py` with an `OpSpec` and `register_op` it — the CLI picks it up +by name. Provide what applies (the loop skips the rest): -# 2. profile it — is there headroom, and which kernel to attack? -python profile_op.py --impl infera_decode --tokens 1 --kernels +| Field | For | +| --- | --- | +| `make_inputs(dims, dev)` | build the op's tensors at a model's dims | +| `baseline(*inputs)` | the engine's built-in op (A) | +| `candidate(*inputs)` | the plugin's op — the selected variant (B) | +| `reference(*inputs)` | correctness oracle (optional) | +| `traffic_bytes(dims)` | bytes moved, for the roofline (optional) | +| `tune_env` / `tune_grid` / `inject` | the tune ring (optional) | -# 3. tune and inject the winner into the plugin -python tune_op.py --tokens 1 --inject +`ops/moe_experts.py` is the reference implementation (Kimi-2.6 dims, the +`infera_fused_experts` candidate, a torch-SwiGLU oracle, block/warp tuning). -# 4. re-profile to confirm (new default is now baked in) -python profile_op.py --impl infera_decode --tokens 1 -``` +## Worked example: `moe_experts` / `infera_decode` -Measured on MI355X (vLLM 0.23 ROCm), Kimi-2.6 dims, decode `T=1`: the built-in -experts kernel = 0.171 ms; `infera_decode` after this loop = **0.135 ms -(1.31× )** at **~5.2 TB/s (65% of HBM peak)** — bandwidth-bound, so further wins -need less traffic (dtype / dedup), not more tuning. +Measured on MI355X (vLLM 0.23 ROCm), Kimi-2.6 dims, decode `T=1`: built-in +experts = 0.171 ms; `infera_decode` after this loop = **0.135 ms (1.31× )** at +**~5.2 TB/s (65% of HBM peak)** — the profile verdict (bandwidth-bound, reads each +expert once = traffic floor) says it's near its roofline, so further wins need +*less traffic* (dtype), not more tuning. -## End-to-end (does the op win show up in a serve?) +### End-to-end (does the op win reach a serve?) -`../_moe_decode_e2e_serve.sh` + `../_moe_decode_e2e_client.py` measure single-stream -(batch-1) decode ITL with the kernel wired into the serving path (the plugin's MoE -seam runs `infera_decode` on genuine small-M bf16 decode steps, verified via a -fire counter). Measured on **Qwen3.5-35B-A3B (bf16 MoE)**, MI355X, TP=1: +`../_moe_decode_e2e_serve.sh` + `../_moe_decode_e2e_client.py` measure batch-1 +decode ITL with the kernel wired into the serving path (fire-counter verified). +**Qwen3.5-35B-A3B (bf16 MoE)**, MI355X, TP=1: | config | ITL (ms/tok) | decode tok/s | | --- | --- | --- | @@ -61,23 +69,7 @@ fire counter). Measured on **Qwen3.5-35B-A3B (bf16 MoE)**, MI355X, TP=1: | aiter off — builtin (triton) | 5.650 | 176.7 | | **aiter off — `infera_decode`** | **5.264** | **188.8** | -So `infera_decode` is the **fastest config, ~8.5% better ITL than the aiter-on -default** — because at batch-1 aiter gives no decode benefit here (slightly -negative), so turning it off is free and the MoE kernel then wins. The op-level -~1.3× dilutes to single-digit e2e because routed experts are ~⅓ of the -bandwidth-bound batch-1 decode step. - -**Limitations (by design, all self-delegating — never a regression):** -- **bf16/fp16 only** — MXFP4/quantized weights delegate (so Kimi-2.6 MXFP4 sees no effect). -- **batch ≤ 16** — larger batches delegate and aiter's large-M wins take over. -- **plain weight layout** — aiter pre-shuffles MoE weights, so the seam delegates when - `VLLM_ROCM_USE_AITER` is set; the kernel runs with aiter's MoE path off (or - `INFERA_MOE_ASSUME_PLAIN=1`). Mixing aiter-prefill with this decode kernel needs - weight un-shuffling — future work. - -## Config precedence - -`_decode_tunables()` reads, in order: per-key env (`INFERA_MOE_GU_BLOCK_I`, …) → -`INFERA_MOE_TUNE_FILE` (JSON) → the baked `_TUNE_DEFAULTS`. So `--inject` -(rewrites `_TUNE_DEFAULTS`) makes a tuned config permanent, while env / tune-file -let you A/B without editing source. +Fastest config, **+8.5% ITL over the aiter-on default** (at batch-1 aiter gives no +decode benefit, so disabling it is free and the MoE kernel wins). Self-delegating +limitations: bf16 only, batch ≤ 16, plain (non-aiter-shuffled) weights — never a +regression. diff --git a/bench/op_loop/framework.py b/bench/op_loop/framework.py new file mode 100644 index 00000000..5d0ae060 --- /dev/null +++ b/bench/op_loop/framework.py @@ -0,0 +1,187 @@ +#!/usr/bin/env python3 +"""Op-agnostic optimize-loop scaffold (issue #40). + +The point of this directory is the *scaffold*, not any one kernel: a uniform +**measure → profile → tune → inject** loop that works for ANY op behind the +Infera vLLM op-injection plugin. Adding an op is writing one :class:`OpSpec` and +registering it (see ``ops/``); `loop.py` then drives measure/profile/tune for it +by name. The kernels (e.g. ``infera_decode``) are just candidates plugged in. + +An OpSpec tells the loop how to exercise one op: + * ``make_inputs(dims, device)`` → the tensors the op takes, + * ``baseline(*inputs)`` → the engine's built-in op (the A in A/B), + * ``candidate(*inputs)`` → the plugin's op (selected variant; the B), + * ``reference(*inputs)`` (optional) → a correctness oracle, + * ``traffic_bytes(dims)`` (optional)→ bytes moved, for the roofline profile, + * ``tune_env`` / ``tune_grid`` / ``inject`` (optional) → the tune ring. +Anything an op doesn't provide, the loop simply skips for that op. +""" + +from __future__ import annotations + +import importlib +import os +from collections.abc import Callable +from dataclasses import dataclass + +import torch + +# --- registry -------------------------------------------------------------- + +_REGISTRY: dict[str, OpSpec] = {} + + +def register_op(spec: OpSpec) -> OpSpec: + _REGISTRY[spec.name] = spec + return spec + + +def get_op(name: str) -> OpSpec: + if name not in _REGISTRY: + # ops self-register on import; import the module named after the op. + importlib.import_module(f"ops.{name}") + return _REGISTRY[name] + + +def list_ops() -> list[str]: + return sorted(_REGISTRY) + + +@dataclass +class OpSpec: + name: str + default_dims: dict + make_inputs: Callable # (dims, device) -> tuple of tensors + baseline: Callable # (*inputs) -> Tensor + candidate: Callable # (*inputs) -> Tensor + reference: Callable | None = None # (*inputs) -> Tensor + traffic_bytes: Callable | None = None # (dims) -> int + tune_env: tuple[str, ...] = () # env keys the tuner sweeps + tune_grid: Callable | None = None # (dims) -> list[tuple] of full configs + inject: Callable | None = None # (config) -> None (bake into the plugin) + peak_gbps: float = 8000.0 # HBM peak for the roofline (MI355X ~8 TB/s) + + +# --- shared helpers -------------------------------------------------------- + + +def timed(fn, iters=30, warmup=8) -> float: + for _ in range(warmup): + fn() + torch.cuda.synchronize() + ts = [] + for _ in range(iters): + s, e = torch.cuda.Event(True), torch.cuda.Event(True) + s.record() + fn() + e.record() + torch.cuda.synchronize() + ts.append(s.elapsed_time(e)) + ts.sort() + return ts[len(ts) // 2] + + +def rel_err(a, b) -> float: + a, b = a.float(), b.float() + return ((a - b).abs().max() / b.abs().max().clamp_min(1e-6)).item() + + +def _cast(v): + try: + return int(v) + except (ValueError, TypeError): + return v + + +def _dims(spec: OpSpec, overrides: dict) -> dict: + d = dict(spec.default_dims) + d.update({k: _cast(v) for k, v in overrides.items()}) + return d + + +# --- the three rings, op-agnostic ----------------------------------------- + + +def measure(spec: OpSpec, overrides: dict, dev="cuda"): + dims = _dims(spec, overrides) + inp = spec.make_inputs(dims, dev) + base = spec.baseline(*inp) + cand = spec.candidate(*inp) + print(f"op={spec.name} dims={dims}") + rows = [("baseline (built-in)", timed(lambda: spec.baseline(*inp)), 0.0)] + if spec.reference is not None: + ref = spec.reference(*inp) + rows.insert( + 0, ("reference (oracle)", timed(lambda: spec.reference(*inp), 5, 1), rel_err(ref, base)) + ) + rows.append(("candidate (plugin op)", timed(lambda: spec.candidate(*inp)), rel_err(cand, base))) + base_t = next(t for n, t, _ in rows if n.startswith("baseline")) + print(f" {'impl':26s} {'median ms':>10s} {'rel vs base':>12s} {'speedup':>9s}") + for n, t, r in rows: + print(f" {n:26s} {t:10.4f} {r:12.2e} {base_t / t:8.2f}x") + + +def profile(spec: OpSpec, overrides: dict, kernels=False, dev="cuda"): + dims = _dims(spec, overrides) + inp = spec.make_inputs(dims, dev) + spec.candidate(*inp) # warm + ms = timed(lambda: spec.candidate(*inp)) + print(f"op={spec.name} dims={dims}\n latency : {ms:.4f} ms") + if spec.traffic_bytes is not None: + wb = spec.traffic_bytes(dims) + bw = (wb / 1e9) / (ms / 1e3) # GB/s + pct = 100 * bw / spec.peak_gbps + verdict = ( + "bandwidth-bound — near peak; win by moving less traffic" + if pct >= 55 + else "launch/occupancy-bound — headroom; tune tiling/warps" + ) + print(f" traffic : {wb / 1e6:.1f} MB") + print( + f" achieved BW : {bw / 1e3:.2f} TB/s ({pct:.0f}% of {spec.peak_gbps / 1e3:.1f} TB/s peak)" + ) + print(f" bottleneck : {verdict}") + if kernels: + from torch.profiler import ProfilerActivity + from torch.profiler import profile as tprofile + + for _ in range(5): + spec.candidate(*inp) + torch.cuda.synchronize() + with tprofile(activities=[ProfilerActivity.CUDA]) as prof: + for _ in range(20): + spec.candidate(*inp) + torch.cuda.synchronize() + print("\n per-kernel device time (top 6):") + print(prof.key_averages().table(sort_by="self_device_time_total", row_limit=6)) + + +def tune(spec: OpSpec, overrides: dict, inject=False, dev="cuda"): + if spec.tune_grid is None or not spec.tune_env: + print(f"op={spec.name}: not tunable (no tune_grid/tune_env)") + return + dims = _dims(spec, overrides) + inp = spec.make_inputs(dims, dev) + ref = spec.baseline(*inp).float() + base_t = timed(lambda: spec.baseline(*inp)) + best, best_t = None, float("inf") + for cfg in spec.tune_grid(dims): + for k, v in zip(spec.tune_env, cfg): + os.environ[k] = str(v) + try: + out = spec.candidate(*inp) + if rel_err(out, ref) > 1e-2: + continue + t = timed(lambda: spec.candidate(*inp)) + except Exception: # noqa: BLE001 + continue + if t < best_t: + best, best_t = cfg, t + print(f"op={spec.name} dims={dims} baseline {base_t:.4f} ms") + if best is None: + print(" no correct config found") + return + print(f" best {best} -> {best_t:.4f} ms ({base_t / best_t:.2f}x vs baseline)") + if inject and spec.inject is not None: + spec.inject(best) + print(f" injected {best} into the plugin") diff --git a/bench/op_loop/loop.py b/bench/op_loop/loop.py new file mode 100644 index 00000000..3ef8f23d --- /dev/null +++ b/bench/op_loop/loop.py @@ -0,0 +1,57 @@ +#!/usr/bin/env python3 +"""Uniform CLI for the op optimize-loop scaffold (issue #40). + +One driver, any registered op — measure / profile / tune by name: + + python loop.py measure --op moe_experts + python loop.py profile --op moe_experts --kernels + python loop.py tune --op moe_experts --inject + python loop.py measure --op moe_experts -d tokens=1 -d experts=64 # override dims + python loop.py list + +Adding an op is writing one ``OpSpec`` in ``ops/.py`` (see +``ops/moe_experts.py`` as the template) — no new script. The plugin kernel is +just the candidate the op's spec points at. +""" + +import argparse +import os +import sys + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import framework as fw # noqa: E402 + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("stage", choices=["measure", "profile", "tune", "list"]) + ap.add_argument("--op", help="registered op name (see `loop.py list`)") + ap.add_argument("-d", "--dim", action="append", default=[], help="dims override k=v") + ap.add_argument("--kernels", action="store_true", help="profile: per-kernel split") + ap.add_argument("--inject", action="store_true", help="tune: bake winner into the plugin") + args = ap.parse_args() + + if args.stage == "list": + import glob + import importlib + + for f in glob.glob(os.path.join(os.path.dirname(__file__), "ops", "*.py")): + name = os.path.splitext(os.path.basename(f))[0] + if not name.startswith("_"): + importlib.import_module(f"ops.{name}") + print("registered ops:", ", ".join(fw.list_ops()) or "(none)") + return + if not args.op: + ap.error("--op is required") + spec = fw.get_op(args.op) + overrides = dict(kv.split("=", 1) for kv in args.dim) + if args.stage == "measure": + fw.measure(spec, overrides) + elif args.stage == "profile": + fw.profile(spec, overrides, kernels=args.kernels) + elif args.stage == "tune": + fw.tune(spec, overrides, inject=args.inject) + + +if __name__ == "__main__": + main() diff --git a/bench/op_loop/moe_experts_loop.py b/bench/op_loop/moe_experts_loop.py deleted file mode 100644 index 482df710..00000000 --- a/bench/op_loop/moe_experts_loop.py +++ /dev/null @@ -1,129 +0,0 @@ -#!/usr/bin/env python3 -"""MoE-experts op optimization loop (issue #40). - -The fast inner loop for iterating custom MoE kernels behind the Infera vLLM -op-injection plugin: measure the **built-in** experts kernel (baseline), the -**plugin op** (``infera.engine.vllm.ops.moe.infera_fused_experts`` — swap the -variant with ``INFERA_MOE_EXPERTS``), and a **pure-torch reference** (correctness -oracle), on identical inputs at a real model's MoE dimensions (Kimi-2.6 by -default). Reports median latency + max abs/rel error, so a new kernel is a -one-line swap → re-run → compare. - - python moe_experts_loop.py # Kimi-2.6 dims, builtin - INFERA_MOE_EXPERTS=aiter python moe_experts_loop.py # optimized variant - python moe_experts_loop.py --experts 64 --tokens 256 # quick -""" - -import argparse -import os - -import torch - - -def build_inputs(E, H, Dm, top_k, T, dtype, device, seed=0): - g = torch.Generator(device=device).manual_seed(seed) - x = torch.randn(T, H, dtype=dtype, device=device, generator=g) * 0.1 - # w1 = [gate; up] stacked → [E, 2I, H]; w2 = down → [E, H, Dm] - w1 = ( - torch.randn(E, 2 * Dm, H, dtype=dtype, device=device, generator=g) * (H**-0.5) - ).contiguous() - w2 = (torch.randn(E, H, Dm, dtype=dtype, device=device, generator=g) * (Dm**-0.5)).contiguous() - logits = torch.randn(T, E, dtype=torch.float32, device=device, generator=g) - weights, ids = torch.topk(torch.softmax(logits, dim=-1), top_k, dim=-1) - return x, w1, w2, weights.contiguous(), ids.to(torch.int32).contiguous() - - -def reference(x, w1, w2, topk_weights, topk_ids): - """Pure-torch SwiGLU MoE (router weight applied on output) — the oracle.""" - T, H = x.shape - E, twoI, _ = w1.shape - Dm = twoI // 2 - out = torch.zeros(T, H, dtype=torch.float32, device=x.device) - xf = x.float() - for e in range(E): - sel = topk_ids == e - if not sel.any(): - continue - tok, slot = sel.nonzero(as_tuple=True) - gu = xf[tok] @ w1[e].float().t() - g, u = gu[:, :Dm], gu[:, Dm:] - o = (torch.nn.functional.silu(g) * u) @ w2[e].float().t() - out.index_add_(0, tok, o * topk_weights[tok, slot].unsqueeze(1)) - return out - - -def timed(fn, iters=30, warmup=8): - for _ in range(warmup): - fn() - torch.cuda.synchronize() - ts = [] - for _ in range(iters): - s, e = torch.cuda.Event(True), torch.cuda.Event(True) - s.record() - fn() - e.record() - torch.cuda.synchronize() - ts.append(s.elapsed_time(e)) - ts.sort() - return ts[len(ts) // 2] - - -def err(a, b): - a, b = a.float(), b.float() - denom = b.abs().max().clamp_min(1e-6) - return (a - b).abs().max().item(), ((a - b).abs().max() / denom).item() - - -def main(): - ap = argparse.ArgumentParser() - ap.add_argument("--experts", type=int, default=384) # Kimi-2.6 - ap.add_argument("--hidden", type=int, default=7168) - ap.add_argument("--inter", type=int, default=2048) - ap.add_argument("--topk", type=int, default=8) - ap.add_argument("--tokens", type=int, default=512) - ap.add_argument("--dtype", default="bfloat16") - ap.add_argument("--skip-ref", action="store_true") - args = ap.parse_args() - - dev = "cuda" - dt = getattr(torch, args.dtype) - E, H, Dm, top_k, T = args.experts, args.hidden, args.inter, args.topk, args.tokens - print( - f"MoE experts op @ Kimi-2.6 dims: E={E} H={H} Dm={Dm} top_k={top_k} tokens={T} dtype={args.dtype}\n" - f"variant (INFERA_MOE_EXPERTS)={os.environ.get('INFERA_MOE_EXPERTS', 'builtin')}\n" - ) - x, w1, w2, tw, ti = build_inputs(E, H, Dm, top_k, T, dt, dev) - - from vllm.model_executor.layers.fused_moe import fused_experts as builtin - - from infera.engine.vllm.ops.moe import infera_fused_experts as plugin_op - - def call(fn): - return fn(x, w1, w2, tw, ti, global_num_experts=E) - - base_out = call(builtin) - plug_out = call(plugin_op) - - rows = [] - if not args.skip_ref: - ref_out = reference(x, w1, w2, tw, ti) - rows.append( - ( - "reference (torch oracle)", - timed(lambda: reference(x, w1, w2, tw, ti), iters=5, warmup=1), - err(ref_out, base_out), - ) - ) - base_t = timed(lambda: call(builtin)) - plug_t = timed(lambda: call(plugin_op)) - rows.append(("baseline (built-in kernel)", base_t, (0.0, 0.0))) - rows.append(("plugin op (infera_fused_experts)", plug_t, err(plug_out, base_out))) - - print(f"{'impl':34s} {'median ms':>10s} {'max|Δ| vs base':>16s} {'rel':>10s} {'speedup':>9s}") - for name, ms, (ad, rd) in rows: - sp = f"{base_t / ms:.2f}x" if ms > 0 else "-" - print(f"{name:34s} {ms:10.3f} {ad:16.2e} {rd:10.2e} {sp:>9s}") - - -if __name__ == "__main__": - main() diff --git a/bench/op_loop/ops/__init__.py b/bench/op_loop/ops/__init__.py new file mode 100644 index 00000000..78f2c04a --- /dev/null +++ b/bench/op_loop/ops/__init__.py @@ -0,0 +1 @@ +"""Op specs for the optimize-loop scaffold — one file per op, self-registering.""" diff --git a/bench/op_loop/ops/moe_experts.py b/bench/op_loop/ops/moe_experts.py new file mode 100644 index 00000000..13e5eda1 --- /dev/null +++ b/bench/op_loop/ops/moe_experts.py @@ -0,0 +1,116 @@ +"""MoE-experts op spec (issue #40) — the worked example for the scaffold. + +Baseline = vLLM's built-in ``fused_experts`` (aiter on ROCm); candidate = the +plugin's ``infera_fused_experts`` (whichever variant ``INFERA_MOE_EXPERTS`` +selects); reference = a pure-torch SwiGLU oracle. Default dims are Kimi-2.6's MoE +block. This is the template: a new op is one file like this + ``register_op``. +""" + +import os + +import framework as fw +import torch + + +def make_inputs(dims, dev): + E, H, Dm, K, T = (dims["experts"], dims["hidden"], dims["inter"], dims["topk"], dims["tokens"]) + dt = getattr(torch, dims["dtype"]) + g = torch.Generator(device=dev).manual_seed(0) + x = torch.randn(T, H, dtype=dt, device=dev, generator=g) * 0.1 + w1 = (torch.randn(E, 2 * Dm, H, dtype=dt, device=dev, generator=g) * (H**-0.5)).contiguous() + w2 = (torch.randn(E, H, Dm, dtype=dt, device=dev, generator=g) * (Dm**-0.5)).contiguous() + logits = torch.randn(T, E, dtype=torch.float32, device=dev, generator=g) + tw, ti = torch.topk(torch.softmax(logits, dim=-1), K, dim=-1) + return x, w1, w2, tw.contiguous(), ti.to(torch.int32).contiguous() + + +def baseline(x, w1, w2, tw, ti): + from vllm.model_executor.layers.fused_moe import fused_experts + + return fused_experts(x, w1, w2, tw, ti, global_num_experts=w1.shape[0]) + + +def candidate(x, w1, w2, tw, ti): + from infera.engine.vllm.ops.moe import infera_fused_experts + + return infera_fused_experts(x, w1, w2, tw, ti, global_num_experts=w1.shape[0]) + + +def reference(x, w1, w2, tw, ti): + """Pure-torch SwiGLU MoE (router weight on output) — correctness oracle.""" + T, H = x.shape + E, twoDm, _ = w1.shape + Dm = twoDm // 2 + out = torch.zeros(T, H, dtype=torch.float32, device=x.device) + xf = x.float() + for e in range(E): + sel = ti == e + if not sel.any(): + continue + tok, slot = sel.nonzero(as_tuple=True) + gu = xf[tok] @ w1[e].float().t() + g, u = gu[:, :Dm], gu[:, Dm:] + o = (torch.nn.functional.silu(g) * u) @ w2[e].float().t() + out.index_add_(0, tok, o * tw[tok, slot].unsqueeze(1)) + return out + + +def traffic_bytes(dims): + db = torch.finfo(getattr(torch, dims["dtype"])).bits // 8 + return dims["tokens"] * dims["topk"] * 3 * dims["inter"] * dims["hidden"] * db + + +_TUNE_ENV = ( + "INFERA_MOE_GU_BLOCK_I", + "INFERA_MOE_GU_BLOCK_H", + "INFERA_MOE_GU_WARPS", + "INFERA_MOE_DN_BLOCK_H", + "INFERA_MOE_DN_BLOCK_I", + "INFERA_MOE_DN_WARPS", +) + + +def tune_grid(dims): + # A small neighbourhood of the impactful knobs (warps + one block dim each). + return [ + (32, gu_bh, gu_w, 8, dn_bi, dn_w) + for gu_bh in (256, 512) + for gu_w in (4, 8) + for dn_bi in (256, 512) + for dn_w in (4, 8) + ] + + +def inject(cfg): + import re + + p = os.path.abspath( + os.path.join(os.path.dirname(__file__), "../../../infera/engine/vllm/ops/moe.py") + ) + src = open(p).read() + open(p, "w").write( + re.sub(r"_TUNE_DEFAULTS = \([^)]*\)", f"_TUNE_DEFAULTS = {tuple(cfg)}", src, count=1) + ) + + +fw.register_op( + fw.OpSpec( + name="moe_experts", + default_dims={ + "experts": 384, + "hidden": 7168, + "inter": 2048, + "topk": 8, + "tokens": 1, + "dtype": "bfloat16", + }, + make_inputs=make_inputs, + baseline=baseline, + candidate=candidate, + reference=reference, + traffic_bytes=traffic_bytes, + tune_env=_TUNE_ENV, + tune_grid=tune_grid, + inject=inject, + ) +) diff --git a/bench/op_loop/profile_op.py b/bench/op_loop/profile_op.py deleted file mode 100644 index 17149e72..00000000 --- a/bench/op_loop/profile_op.py +++ /dev/null @@ -1,110 +0,0 @@ -#!/usr/bin/env python3 -"""Profile ring of the MoE-experts optimize loop (issue #40). - -Roofline profile of an experts impl at a given shape: median latency, the -selected-expert weight traffic it must move, achieved effective HBM bandwidth, -%-of-peak, and a bottleneck verdict (bandwidth-bound ⇒ near optimal, cut traffic; -launch/occupancy-bound ⇒ headroom, tune tiling/warps). ``--kernels`` adds the -per-kernel device-time split (torch profiler) so you see which kernel to tune. - - python profile_op.py --impl infera_decode --tokens 1 - python profile_op.py --impl infera_decode --tokens 1 --kernels - python profile_op.py --impl builtin --tokens 1 # profile the baseline -""" - -import argparse -import os -import sys - -import torch - -sys.path.insert(0, os.path.dirname(__file__)) -from moe_experts_loop import build_inputs, timed # noqa: E402 - - -def weight_bytes(T, K, H, Dm, dbytes): - # The decode kernel reads each (token, slot) expert's gate/up (2·Dm·H) + down - # (H·Dm) weights once — the traffic floor for this regime. - return T * K * (2 * Dm * H + H * Dm) * dbytes - - -def peak_gbps(override): - if override: - return override - try: - p = torch.cuda.get_device_properties(0) - bw = p.memory_clock_rate * 1e3 * (p.memory_bus_width / 8) * 2 / 1e9 - if bw > 6000: # trust only a plausibly-HBM3e-class figure - return bw - except Exception: # noqa: BLE001 - pass - return 8000.0 # MI355X HBM3e ~8 TB/s; override with --peak-gbps - - -def get_impl(name): - from vllm.model_executor.layers.fused_moe import fused_experts as builtin - - if name in ("builtin", "baseline"): - return builtin - os.environ["INFERA_MOE_EXPERTS"] = name - from infera.engine.vllm.ops.moe import infera_fused_experts - - return infera_fused_experts - - -def main(): - ap = argparse.ArgumentParser() - ap.add_argument("--impl", default="infera_decode") - ap.add_argument("--experts", type=int, default=384) - ap.add_argument("--hidden", type=int, default=7168) - ap.add_argument("--inter", type=int, default=2048) - ap.add_argument("--topk", type=int, default=8) - ap.add_argument("--tokens", type=int, default=1) - ap.add_argument("--dtype", default="bfloat16") - ap.add_argument("--peak-gbps", type=float, default=0.0) - ap.add_argument("--kernels", action="store_true") - args = ap.parse_args() - - dt = getattr(torch, args.dtype) - E, H, Dm, K, T = args.experts, args.hidden, args.inter, args.topk, args.tokens - x, w1, w2, tw, ti = build_inputs(E, H, Dm, K, T, dt, "cuda") - impl = get_impl(args.impl) - - def call(): - return impl(x, w1, w2, tw, ti, global_num_experts=E) - - call() # warm / JIT - ms = timed(call) - wb = weight_bytes(T, K, H, Dm, torch.finfo(dt).bits // 8) - bw_gbps = (wb / 1e9) / (ms / 1e3) # GB/s - peak = peak_gbps(args.peak_gbps) # GB/s - pct = 100 * bw_gbps / peak - verdict = ( - "bandwidth-bound — near peak; win by moving less traffic (dtype / dedup)" - if pct >= 55 - else "launch/occupancy-bound — headroom; tune tiling/warps, cut launches" - ) - print(f"impl={args.impl} T={T} E={E} H={H} I={Dm} top_k={K} dtype={args.dtype}") - print(f" latency : {ms:.4f} ms") - print(f" weight traffic : {wb / 1e6:.1f} MB (each selected expert read once)") - print( - f" achieved BW : {bw_gbps / 1e3:.2f} TB/s ({pct:.0f}% of {peak / 1e3:.1f} TB/s peak)" - ) - print(f" bottleneck : {verdict}") - - if args.kernels: - from torch.profiler import ProfilerActivity, profile - - for _ in range(5): - call() - torch.cuda.synchronize() - with profile(activities=[ProfilerActivity.CUDA]) as prof: - for _ in range(20): - call() - torch.cuda.synchronize() - print("\n per-kernel device time (top 6):") - print(prof.key_averages().table(sort_by="self_device_time_total", row_limit=6)) - - -if __name__ == "__main__": - main() diff --git a/bench/op_loop/tune_op.py b/bench/op_loop/tune_op.py deleted file mode 100644 index a15ec583..00000000 --- a/bench/op_loop/tune_op.py +++ /dev/null @@ -1,106 +0,0 @@ -#!/usr/bin/env python3 -"""Tune ring of the MoE-experts optimize loop (issue #40). - -Coordinate-descent autotune of the ``infera_decode`` kernel's block/warp config -at a given shape: sweeps the gate/up kernel, then the down/combine kernel, -keeping only numerically-correct configs (rel < 1e-2 vs builtin), reports the -best, and ``--inject`` bakes it into the plugin (rewrites ``_TUNE_DEFAULTS`` in -infera/engine/vllm/ops/moe.py) — the "植入" step, so re-profiling picks it up. - - python tune_op.py --tokens 1 # tune, print winner - python tune_op.py --tokens 1 --inject # tune + write into the plugin -""" - -import argparse -import os -import re -import sys - -import torch - -sys.path.insert(0, os.path.dirname(__file__)) -from moe_experts_loop import build_inputs, timed # noqa: E402 - -KEYS = ( - "INFERA_MOE_GU_BLOCK_I", - "INFERA_MOE_GU_BLOCK_H", - "INFERA_MOE_GU_WARPS", - "INFERA_MOE_DN_BLOCK_H", - "INFERA_MOE_DN_BLOCK_I", - "INFERA_MOE_DN_WARPS", -) -GU_GRID = [(bi, bh, w) for bi in (16, 32, 64) for bh in (256, 512, 1024) for w in (4, 8)] -DN_GRID = [(bh, bi, w) for bh in (8, 16, 32) for bi in (256, 512, 1024) for w in (4, 8)] - - -def measure(fn, x, w1, w2, tw, ti, E, ref): - try: - out = fn(x, w1, w2, tw, ti, global_num_experts=E) - rel = ((out.float() - ref).abs().max() / ref.abs().max().clamp_min(1e-6)).item() - if rel > 1e-2: - return None, rel - return timed(lambda: fn(x, w1, w2, tw, ti, global_num_experts=E)), rel - except Exception: # noqa: BLE001 - return None, float("inf") - - -def main(): - ap = argparse.ArgumentParser() - ap.add_argument("--experts", type=int, default=384) - ap.add_argument("--hidden", type=int, default=7168) - ap.add_argument("--inter", type=int, default=2048) - ap.add_argument("--topk", type=int, default=8) - ap.add_argument("--tokens", type=int, default=1) - ap.add_argument("--dtype", default="bfloat16") - ap.add_argument("--inject", action="store_true") - args = ap.parse_args() - - dt = getattr(torch, args.dtype) - E, H, Dm, K, T = args.experts, args.hidden, args.inter, args.topk, args.tokens - x, w1, w2, tw, ti = build_inputs(E, H, Dm, K, T, dt, "cuda") - from vllm.model_executor.layers.fused_moe import fused_experts as builtin - - from infera.engine.vllm.ops.moe import _EXPERTS_VARIANTS, _TUNE_DEFAULTS - - fn = _EXPERTS_VARIANTS["infera_decode"] - ref = builtin(x, w1, w2, tw, ti, global_num_experts=E).float() - base_t = timed(lambda: builtin(x, w1, w2, tw, ti, global_num_experts=E)) - - cfg = list(_TUNE_DEFAULTS) # start from current baked defaults - - def setenv(c): - for k, v in zip(KEYS, c): - os.environ[k] = str(v) - - def best_over(grid, slot): # slot: 0 for GU (keys 0:3), 3 for DN (keys 3:6) - best, best_ms = None, float("inf") - for combo in grid: - trial = cfg.copy() - trial[slot : slot + 3] = list(combo) - setenv(trial) - ms, rel = measure(fn, x, w1, w2, tw, ti, E, ref) - if ms is not None and ms < best_ms: - best, best_ms = combo, ms - return best, best_ms - - print(f"tuning infera_decode @ T={T} E={E} H={H} I={Dm} top_k={K}; baseline {base_t:.4f} ms") - gu, _ = best_over(GU_GRID, 0) - cfg[0:3] = list(gu) - dn, ms = best_over(DN_GRID, 3) - cfg[3:6] = list(dn) - setenv(cfg) - ms, rel = measure(fn, x, w1, w2, tw, ti, E, ref) - print(f" default {tuple(_TUNE_DEFAULTS)} ") - print(f" best {tuple(cfg)} -> {ms:.4f} ms ({base_t / ms:.2f}x vs builtin, rel {rel:.1e})") - - if args.inject: - path = os.path.join(os.path.dirname(__file__), "..", "..", "infera/engine/vllm/ops/moe.py") - path = os.path.abspath(path) - src = open(path).read() - new = re.sub(r"_TUNE_DEFAULTS = \([^)]*\)", f"_TUNE_DEFAULTS = {tuple(cfg)}", src, count=1) - open(path, "w").write(new) - print(f" injected _TUNE_DEFAULTS = {tuple(cfg)} into {path}") - - -if __name__ == "__main__": - main()