feat(vllm): op-injection plugin + decode MoE experts kernel (1.28x at small M)

[jiejingzhangamd]

Closes the first loop of #40: inject custom Attention/MoE kernels into stock vLLM out-of-tree (no fork), plus a first real optimized kernel with a measured win.

Mechanism (no fork)

A single vllm.general_plugins entry point (register_ops) installs both seams after vllm.platforms initializes, so it's import-safe (a platform_plugins subclass circular-imports during vLLM's eager current_platform resolution). No-op unless ROCm + INFERA_VLLM_OPS_DISABLE != 1.

• Attention — patches get_attn_backend_cls on the resolved platform (pass-through; INFERA_ATTN_BACKEND="mod:Backend" injects a custom backend). ATOM's MLA-patch style, import-safe.
• MoE — infera_fused_experts (same signature as vLLM fused_experts) with a variant registry (@register_experts_variant, INFERA_MOE_EXPERTS=name); default is a bitwise pass-through to the builtin.

Optimize loop

bench/op_loop/moe_experts_loop.py measures builtin (baseline) vs plugin op vs a pure-torch reference (correctness oracle) at a model's real MoE dims. A new kernel is one @register_experts_variant + re-run.

First real kernel: infera_decode

A small-M SwiGLU experts kernel — two Triton GEMV kernels (fp32 accumulate, expert-id-indexed weight reads, fused gate/up/SiLU and down+router-combine, ~5.2 TB/s effective HBM). A dispatch guard delegates to the builtin where it'd lose or be unsafe (T>16 or 2·T·topk>E; non-bf16/fp16 weights ⇒ Kimi MXFP4 safely delegates; expert_map/non-SiLU/router-on-input).

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

tokens	baseline ms	infera_decode ms	speedup	rel err
1	0.173	0.135	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	4.351	4.349	1.00× (delegated)	0
512	5.815	5.810	1.00× (delegated)	0

A decode-step specialist (the latency-critical, token-by-token regime), not an aiter replacement — safe by construction everywhere else. rel err is below the fp32 torch oracle's own error vs the builtin.

Validation

import vllm clean; both seams install via load_general_plugins; results independently re-verified. ruff clean. Full progression + negative-result notes in #40.

[jiejingzhangamd]

Wired into the serve path + end-to-end decode win (09c4369)

The MoE seam now actually runs infera_decode in a real serve (it was a pass-through), self-gating so it's never a regression.

Weight-layout state is now on the interface (_weights_are_aiter_shuffled): aiter pre-shuffles MoE weights in-place with no per-tensor marker, so the seam reads the aiter master flag and delegates when weights are shuffled — this is what the earlier "garbage with aiter on" was; it now cleanly delegates to the builtin instead. INFERA_MOE_ASSUME_PLAIN=1 overrides. Verified: aiter-on → delegate (correct), aiter-off → kernel runs.

E2E — Qwen3.5-35B-A3B (bf16 MoE), MI355X, TP=1, batch-1 decode (independently reproduced, 5 reps, no overlap; a fire counter confirms the kernel ran):

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

infera_decode is the fastest config — +8.5% ITL / +9% tok/s over the aiter-on default. 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 (all self-delegating — never a regression): bf16/fp16 only (MXFP4/Kimi delegates), batch ≤ 16, plain weight layout (aiter-shuffled delegates). Mixing aiter-prefill + this decode kernel needs weight un-shuffling — future work.