AMD-AGI / AMD-AGI/Infera

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

Open
#40 6 comments 0 reactions 0 assignees View on GitHub

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)

  • AttentionPlatform.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).
  • MoEFusedMoEModularKernel (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_clspass-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.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.