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

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}")

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 <kernel|baseline> <port>
+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"

bench/op_loop/README.md

@@ -0,0 +1,75 @@
+# Op optimize-loop scaffold (issue #40)
+
+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.**
+
+```
+framework.py   OpSpec + registry + generic measure/profile/tune
+loop.py        one CLI over any registered op
+ops/<name>.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 <repo>:/work -w /work
+-e PYTHONPATH=/work <vllm-rocm-image> bash -lc "cd bench/op_loop && <cmd>"`.
+
+## The loop
+
+```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 | 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. |
+
+## Adding an op
+
+Write `ops/<name>.py` with an `OpSpec` and `register_op` it — the CLI picks it up
+by name. Provide what applies (the loop skips the rest):
+
+| 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) |
+
+`ops/moe_experts.py` is the reference implementation (Kimi-2.6 dims, the
+`infera_fused_experts` candidate, a torch-SwiGLU oracle, block/warp tuning).
+
+## Worked example: `moe_experts` / `infera_decode`
+
+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 reach a serve?)
+
+`../_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 |
+| --- | --- | --- |
+| aiter **on** — production default | 5.754 | 173.3 |
+| aiter off — builtin (triton) | 5.650 | 176.7 |
+| **aiter off — `infera_decode`** | **5.264** | **188.8** |
+
+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.

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")