vLLM op-injection plugin: custom Attention + MoE kernels via out-of-tree entry_points (no fork)
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 21
- Forks
- 4
- Avg merge
- 2d 7h
- Merged PRs (30d)
- 28
Description
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 customAttentionBackend/AttentionImpl. Clean, per-kernel seam (also settable viaVLLM_ATTENTION_BACKEND). This is exactly what vLLM-ATOM uses (ATOMPlatform(RocmPlatform)→AiterBackend/AiterMLABackend). - MoE →
FusedMoEModularKernel(FusedMoEPrepareAndFinalize+FusedMoEPermuteExpertsUnpermute) orFusedMoEMethodBase.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)withget_attn_backend_cls→ pass-through attention backend (delegates to the default impl) as the injection seam. - MoE: a
FusedMoEModularKernelexperts 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.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with vLLM 0.23's vllm.platform_plugins and vllm.general_plugins entry points, then trace the Platform.get_attn_backend_cls and FusedMoEModularKernel seams described here. Validate that the packaged plugin loads, the model runs unchanged with and without INFERA_VLLM_OPS_DISABLE=1, and pass-through numerics are bitwise identical before wiring the first optimized kernel.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- ai-infra-agents, backend, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 42/100