vLLM op-injection plugin: custom Attention + MoE kernels via out-of-tree entry_points (no fork)

[jiejingzhangamd]

Goal

Ship a pip-installable vLLM plugin that injects Infera/HyperLoom-optimized Attention and MoE kernels into stock vLLM — no vLLM fork — using vLLM's out-of-tree plugin mechanism (vllm.platform_plugins + vllm.general_plugins). This is the delivery vehicle for op-injection: swap just the hot kernels while staying on upstream vLLM's model code, so we track vLLM releases instead of maintaining a fork.

Mechanism (verified against vLLM 0.23; cf. AMD's vLLM-ATOM)

• Attention → Platform.get_attn_backend_cls() returns a custom AttentionBackend / AttentionImpl. Clean, per-kernel seam (also settable via VLLM_ATTENTION_BACKEND). This is exactly what vLLM-ATOM uses (ATOMPlatform(RocmPlatform) → AiterBackend/AiterMLABackend).
• MoE → FusedMoEModularKernel (FusedMoEPrepareAndFinalize + FusedMoEPermuteExpertsUnpermute) or FusedMoEMethodBase.apply() — swap just the experts GEMM while keeping vLLM's routing / quant / EP-DP dispatch.
• vLLM-ATOM instead replaces the whole model via ModelRegistry.register_model; we prefer the less-invasive per-op seams above, and keep whole-model wrapping only as a fallback where a per-op seam is insufficient (e.g. deep MLA fusion).

Packaged as entry_points so it activates on install, gated by an env kill-switch (INFERA_VLLM_OPS_DISABLE=1), mirroring ATOM's ATOM_DISABLE_VLLM_PLUGIN.

Plan

• Plugin skeleton package + entry_points (register_platform, register_ops), env-gated.
• InferaPlatform(RocmPlatform) with get_attn_backend_cls → pass-through attention backend (delegates to the default impl) as the injection seam.
• MoE: a FusedMoEModularKernel experts pass-through stub as the injection seam.
• Validation: plugin loads under vLLM; a model runs unchanged with the plugin active vs INFERA_VLLM_OPS_DISABLE=1; numerics match (pass-through = bitwise identical).
• Wire a first real optimized kernel behind the seam (attention first).

Non-goals (for now)

Whole-model wrapper; quant-method rewrites; SGLang (separate mechanism).

Progress will be posted as comments on this issue.

[jiejingzhangamd]

Progress — plugin skeleton landed (branch feat/vllm-op-injection, 60f68b7)

Done

• Out-of-tree plugin package infera/engine/vllm/ops/ + entry_points in pyproject.toml (vllm.platform_plugins → register_platform, vllm.general_plugins → register_ops). No vLLM fork.
• Attention seam: InferaPlatform(RocmPlatform).get_attn_backend_cls (pass-through — delegates to the default ROCm backend; the point where a custom AttentionBackend is substituted).
• MoE seam: install_moe_ops() — where a FusedMoEModularKernel experts impl / FusedMoEMethodBase.apply() will replace the default while keeping vLLM's routing/quant/EP.
• Env gates: no-op unless ROCm present and INFERA_VLLM_OPS_DISABLE != 1.

Validated (vLLM 0.23.1rc1 ROCm image, MI355X)

• vLLM discovers infera in both plugin groups; import vllm is clean on the default install.
• With INFERA_VLLM_OPS_PLATFORM=1, vLLM logs Platform plugin infera is activated → the mechanism works end-to-end.

Blocker found + plan

• Activating the RocmPlatform subclass via platform_plugins currently circular-imports: vLLM resolves current_platform eagerly during import vllm (it's a lazy __getattr__ specifically to avoid this), and platform.py importing vllm.platforms.rocm at module top re-enters vllm.platforms before it finishes.
• Platform activation is therefore opt-in (INFERA_VLLM_OPS_PLATFORM=1) so the default install stays safe.
• Fix: move the attention seam to a register_ops monkey-patch of get_attn_backend_cls (the same approach vLLM-ATOM uses for its MLA forward_impl patch), which runs after vllm.platforms is initialized. Keep the platform subclass only if it can be made import-safe.

Next

• Attention seam via register_ops monkey-patch (import-safe); re-validate Platform plugin activated + a real serve.
• MoE: wire a concrete FusedMoEPermuteExpertsUnpermute pass-through, then a real experts kernel.
• Numerics parity check (pass-through = bitwise identical) with plugin on vs INFERA_VLLM_OPS_DISABLE=1.

[jiejingzhangamd]

Progress — attention seam is import-safe & validated (f3d98ac)

Fixed the circular-import blocker by dropping the platform-subclass activation and moving the attention seam to a single vllm.general_plugins hook that patches the already-resolved platform (runs after vllm.platforms inits, before the model builds — same style as vLLM-ATOM's MLA forward_impl patch, but import-safe).

• Removed the vllm.platform_plugins entry point + 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.
• register_ops now installs both seams.

Validated (vLLM 0.23 ROCm image, MI355X):

import vllm OK 0.23.1rc1...           # clean, no circular import
infera_ops discovered (general_plugins)
load_general_plugins() → attention seam patched on RocmPlatform ✓
INFERA_ATTN_BACKEND override → get_attn_backend_cls() returns the injected backend ✓

Status

• Plugin plumbing (entry_points, env gates, ROCm guard)
• Attention seam — import-safe, validated (pass-through + override)
• MoE experts seam — wire a concrete FusedMoEModularKernel / FusedMoEPermuteExpertsUnpermute pass-through (currently a logged no-op)
• Numerics parity (plugin on vs INFERA_VLLM_OPS_DISABLE=1) on a real serve
• Wire a first real optimized kernel behind the attention seam

Next up: the MoE experts pass-through.

[jiejingzhangamd]

Progress — MoE experts seam wired; both seams validated (ca937e1)

install_moe_ops() now patches the MoE modular-kernel compute chokepoints with a delegating pass-through. In vLLM 0.23 the MoE API was refactored (FusedMoEKernel.apply / apply_monolithic — the route → experts-GEMM → combine entry points), so the patch is defensive: it locates the kernel class and no-ops with a warning if the API has moved.

• Single point where an Infera/HyperLoom experts kernel replaces the default, keeping vLLM's routing / quant / EP-DP dispatch (INFERA_MOE_EXPERTS reserved).
• Less-invasive than vLLM-ATOM's whole-model register_model wrapper.

Validated (vLLM 0.23 ROCm image, MI355X):

import vllm OK                                   # clean
load_general_plugins():
  attention seam patched: True
  FusedMoEKernel.apply           wrapped: True
  FusedMoEKernel.apply_monolithic wrapped: True

Status

• Plugin plumbing (single general_plugins hook, env gates, ROCm guard)
• Attention seam — import-safe patch of get_attn_backend_cls (pass-through + INFERA_ATTN_BACKEND override)
• MoE experts seam — import-safe patch of FusedMoEKernel.apply/apply_monolithic (pass-through, INFERA_MOE_EXPERTS reserved)
• Numerics parity on a real serve (plugin on vs INFERA_VLLM_OPS_DISABLE=1) — pass-through is identical by construction; confirm end-to-end
• Wire a first real optimized kernel behind each seam (attention first)

The skeleton is complete: both OPs (Attention + MoE) have import-safe, observable, bitwise-identical injection seams in stock vLLM via one out-of-tree entry point — no fork. Ready to drop real kernels behind either seam.

[jiejingzhangamd]

Progress — MoE experts optimize-loop working on Kimi-2.6 dims (f1eb31c)

Turned the MoE seam into a runnable optimize loop and measured it at real Kimi-2.6 MoE dimensions.

Plugin op: infera_fused_experts() — same signature as vLLM fused_experts, with a variant registry (@register_experts_variant("name"), selected via INFERA_MOE_EXPERTS). Empty by default ⇒ bitwise pass-through to the built-in kernel (which on ROCm already dispatches to aiter). Falls back to built-in on any error.

Harness: bench/op_loop/moe_experts_loop.py — measures built-in (baseline), plugin op, and a pure-torch reference (correctness oracle) on identical inputs; reports median latency + max abs/rel error.

Measured (MI355X, vLLM 0.23 ROCm), Kimi-2.6 MoE dims — E=384, H=7168, I=2048, top_k=8, 512 tokens:

impl                                median ms   max|Δ| vs base      rel   speedup
reference (torch oracle)               11.658        3.89e-05   5.04e-03    0.10x   <- oracle, validates numerics
baseline (built-in kernel)              5.807        0.00e+00   0.00e+00    1.00x
plugin op (infera_fused_experts)        5.814        0.00e+00   0.00e+00    1.00x   <- verified pass-through

(64-expert quick run shown; full 384-expert baseline = 5.81 ms/call. The reference matches the built-in within bf16 tolerance and the harness cleanly detects the 10x latency gap — so it will show a real kernel's speedup and verify its correctness.)

The loop is closed: baseline (built-in) → swap the op in the plugin (@register_experts_variant + INFERA_MOE_EXPERTS=name) → re-run the bench → compare latency + numerics → iterate. A HyperLoom-tuned experts kernel is now a one-function drop-in.

Status

• Attention + MoE seams (import-safe, validated)
• MoE experts op + optimize-loop harness at Kimi-2.6 dims (baseline measured, pass-through + correctness verified)
• Drop a first real optimized experts kernel behind @register_experts_variant and show a measured speedup
• e2e serve parity (plugin on vs INFERA_VLLM_OPS_DISABLE=1)

[jiejingzhangamd]

Real kernel landed → PR #43

infera_decode — a small-M SwiGLU experts kernel (two Triton GEMV kernels), registered via @register_experts_variant. Independently re-verified at Kimi-2.6 dims (E=384):

tokens	baseline	infera_decode	speedup	rel err
1	0.173 ms	0.135 ms	1.28×	2.4e-3
8	0.959	0.878	1.09×	4.9e-3
16	1.767	1.657	1.07×	4.9e-3
64/512	—	—	1.00× (guard delegates)	0

Honest scope: a decode-step specialist — it beats the aiter-backed builtin only at small M (where the grouped-GEMM's sort/align/dedup overhead dominates), and self-delegates via a dispatch guard (T>16 or 2·T·topk>E, non-bf16 weights, etc.) everywhere else, so it's never a regression and Kimi MXFP4 safely delegates. Not an aiter replacement.

Loop closed end-to-end (baseline → variant → measure → compare). Opened as PR #43.

[jiejingzhangamd]

Optimize loop closed — profile + tune rings added (3d56a63, in PR #43)

The loop is now profile → tune → inject → re-profile, not just an A/B bench:

• profile_op.py — roofline: achieved HBM BW vs peak → bandwidth-bound (near optimal) vs launch/occupancy-bound (headroom); --kernels gives the per-kernel device-time split.
• tune_op.py — coordinate-descent autotune (correct configs only); --inject rewrites _TUNE_DEFAULTS in the plugin — the 植入 step.
• _decode_tunables() precedence: per-key env > INFERA_MOE_TUNE_FILE (JSON) > baked defaults.
• README with the loop diagram + commands.

One full cycle, MI355X, Kimi-2.6 dims, decode T=1:

profile  → 0.137 ms, 5.14 TB/s (64% of HBM peak) → bandwidth-bound; gate/up kernel = 64% of time
tune --inject → DN_WARPS 4→8 baked in (1.28x → 1.31x vs builtin)
re-profile → 0.135 ms, 5.23 TB/s → confirmed

The profile verdict (bandwidth-bound at ~65% of peak, reads each expert once = the traffic floor) says the decode kernel is near its roofline — further wins need less traffic (dtype / dedup), not more tuning. That's the loop doing its job: it tells you when to stop tuning and change strategy.