docker-images: glm53-kindling

docker-images/glm53-kindling/.dockerignore

@@ -0,0 +1,16 @@
+# The build context is the repo root, so the vllm symlink resolves inside it.
+# Denylist rather than allowlist: the Dockerfile's COPY lines are the list of
+# what goes in, and this only keeps the upload small. The status server under
+# .submodules/spark-agent/vllm/ must stay.
+.git
+.submodules/*/.git
+.submodules/spark-agent/agent
+dev/
+smoketest/
+compose/
+.env
+.env.*
+**/__pycache__
+**/*.pyc
+.DS_Store
+*.md

docker-images/glm53-kindling/.env.example

@@ -0,0 +1,34 @@
+# Optional. Copy to compose/.env (not the repo root) only to override
+# something: docker compose reads .env from the compose file's directory.
+# Every tuned knob has its default in image/entrypoint.sh, and a copy here
+# silently wins over the measured value. KNOBS.md lists every knob.
+
+# --- image and host paths ------------------------------------------------------
+# Set when you built under your own tag (image/build.sh's default otherwise).
+#IMAGE=spark-glm53:mine
+# Set when the checkpoint and drafter are not under /srv/models on this box.
+#MODEL_HOST_DIR=/srv/models/glm-5.3-flash-nvfp4
+#DFLASH_HOST_DIR=/srv/models/glm-5.3-flash-dflash2
+# Set to keep JIT caches and logs in a host directory instead of a volume.
+#CACHE_HOME=/srv/container-cache
+#LOG_DIR=/srv/logs
+
+# --- roles: set both, on every box, to pin the head instead of electing it -------
+# The head's LAN address, the same on every box, the head included.
+#HEAD_HOST=192.0.2.10
+# head on the box at HEAD_HOST, worker on the rest.
+#ROLE=worker
+
+# --- networking: set when mentatd's tags do not describe this box ---------------
+# This box's LAN address, when mentatd tags nothing lan.
+#VLLM_HOST_IP=192.0.2.10
+# Fabric prefixes with the trailing dot, one per ConnectX PCIe root, when
+# mentatd tags no ConnectX interface rdma.
+#FABRIC_SUBNETS="198.18.0. 198.19.0."
+
+# --- other ----------------------------------------------------------------------
+# Per box. Set only when the interface holding VLLM_HOST_IP is the wrong one.
+#NCCL_SOCKET_IFNAME=
+# Set a URL only to pin one address (port form, 8002/v1 and 8082/mcp, otherwise).
+#MENTAT_OPENAI_API=
+#MENTAT_MCP_API=

docker-images/glm53-kindling/.gitignore

@@ -0,0 +1,7 @@
+__pycache__/
+*.pyc
+.DS_Store
+.env
+*.log
+quality-data/
+gate-results/

docker-images/glm53-kindling/.gitmodules

@@ -0,0 +1,3 @@
+[submodule ".submodules/spark-agent"]
+	path = .submodules/spark-agent
+	url = https://github.com/mmastrac/spark-agent.git

docker-images/glm53-kindling/.submodules/spark-agent/.gitignore

@@ -0,0 +1,3 @@
+__pycache__/
+*.pyc
+.DS_Store

docker-images/glm53-kindling/.submodules/spark-agent/KNOBS.md

@@ -0,0 +1,17 @@
+# Knobs
+
+The agent has no entrypoint script. It reads these directly
+(`agent/spark-agent.py`). The compose file requires `LOG_DIR` (the host side
+of `/logs`), `MENTAT_ROUTER_URL` and `ALLOWED_SOURCES` in the node's `.env`.
+
+| variable | default | meaning |
+|---|---|---|
+| `AGENT_LOG_DIR` | `/logs` | log and snapshot directory, mounted read-only |
+| `MENTAT_ROUTER_URL` | (none) | mentat-serve, e.g. `http://router:6381`; source of the model table |
+| `MENTAT_GROUP` | `agent-<hostname>` | group the agent registers its MCP under |
+| `MENTAT_DAEMON` | `127.0.0.1:6379` | local mentatd control address |
+| `ALLOWED_SOURCES` | `127.0.0.1` | address prefixes let in; the router's host is always let in |
+| `LOAD_INTERVAL_S` | `5` | how often the router, peers and engines are read |
+| `AGENT_PORT` | `8090` | port peers are linked on when `--port` is not given |
+| `PY_SPY` | `py-spy` | py-spy binary for `engine_stacks` |
+| `SPARK_MEMORY_PY` | `/usr/local/bin/spark-memory.py` | helper behind `memory_accounting` |

docker-images/glm53-kindling/.submodules/spark-agent/MODELS.md

@@ -0,0 +1,121 @@
+# Model recipe checklist
+
+Every model repo (`spark-<name>`) is checked against this list.
+
+## Layout
+
+```
+model.yaml
+.env.example
+README.md
+KNOBS.md
+.gitmodules
+.submodules/spark-agent/
+compose/
+image/
+smoketest/
+```
+
+## model.yaml
+
+- [ ] `name` matches the served model name and the mentat group
+- [ ] `huggingface` id of the exact checkpoint served
+- [ ] `description`: what it is and why this model was chosen
+- [ ] `vram.budget_gib`, `vram.resident_gib` (measured, with the date), `vram.kv_pin_bytes`
+- [ ] `tp`, `pp`
+- [ ] `api` and `endpoints`
+- [ ] Records measured figures; the entrypoint sets them. Where both carry a value (`tp`, `kv_pin_bytes`), they agree, and the entrypoint wins
+
+## .env.example
+
+- [ ] Every variable in `compose/` that has no default, listed
+- [ ] Per-node values called out as per-node
+- [ ] Placeholders only: no fleet addresses, box names or secrets
+- [ ] Live copy is `compose/.env`: compose reads `.env` beside the compose file
+- [ ] `.env` and `compose/.env` are in `.gitignore`
+
+## mentat integration
+
+- [ ] TP=1 / PP=1: `python -m ray.register &` beside the engine, no `ray` executable
+- [ ] Multi-node: `ray start` from the mentatd binary, plus `mentatd-probe-machine`
+- [ ] `MENTAT_GROUP` defaults to the served name
+- [ ] `MENTAT_OPENAI_API` in port form (`8000/v1`), on the API rank only
+- [ ] `MENTAT_MCP_API` in port form, on every rank
+- [ ] `MENTAT_MODEL_PROVIDER` names the real engine (`vllm`, `sglang`, `llamacpp`). The router counts tokens for `vllm` only; never claim `vllm` to get it
+- [ ] `MENTAT_NODE_IP` unset unless the stack is single-node on loopback
+- [ ] Shim version matches the fleet daemons (`MENTAT_VERSION`)
+- [ ] Shows as healthy in mentat-serve's `/`, and a request through the router answers
+- [ ] Status-server tools show in the router's `/mcp`
+- [ ] Status port from the table below, so tenants on one box stay disjoint
+
+## .submodules/spark-agent
+
+vLLM engines only. For SGLang or llama.cpp, mark these N/A: the status server
+reads vLLM's metrics.
+
+- [ ] Submodule at `.submodules/spark-agent`, HTTPS URL
+- [ ] `vllm -> .submodules/spark-agent/vllm` symlink
+- [ ] No local copy of `status-server.py`
+- [ ] No `AGENT_URL` or `/register`: the agent has no registry
+- [ ] Pin bumped on purpose, never left behind on a stale commit
+- [ ] Every status-server knob the recipe relies on is set explicitly: `SERVICE_NAME`, `STAGE_FILE`, `STAGES`, `PORT`, `STATUS_PORT`, `ROLE`, `PEERS`
+- [ ] TP=1 with no self-test: `SERVING_WHEN_READY=1`, not a watcher loop
+
+## compose/
+
+- [ ] `<name>.yaml`, run with `docker compose -f compose/<name>.yaml up -d`
+- [ ] Top-level `name: <model>`: from `compose/` the project would otherwise be called `compose`, and every recipe would share it
+- [ ] Every variable is `${VAR:-default}` or `${VAR:?why}`
+- [ ] No fleet addresses, box names or host paths as defaults
+- [ ] `network_mode: host`; ports disjoint from the other tenants on the box
+- [ ] No added capabilities: host tools belong to the agent
+- [ ] KV pinned (`--kv-cache-memory` or equivalent), never derived from profiling
+- [ ] Model weights mounted read-only
+- [ ] JIT caches (FlashInfer, TileLang, Triton, torch) on a persistent mount
+- [ ] Except the FlashInfer autotune cache when TP > 1: ephemeral, or ranks desync and deadlock
+- [ ] Restart policy set
+
+## image/
+
+- [ ] Everything the Dockerfile copies lives here: Dockerfile, entrypoint, patches, verify script
+- [ ] Build context is the repo root (`docker build -f image/Dockerfile .`), so the `vllm` symlink stays inside it
+- [ ] Base pinned by digest or nightly SHA, never `latest` or `nightly`
+- [ ] `ARG BASE` and `ARG MENTAT_VERSION` above the first `FROM`
+- [ ] `mentat-artifacts` pulled registry-qualified (`mmastrac/...`)
+- [ ] Shim wheel installed `--no-deps`
+- [ ] `verify-base.py` fails the build if the model arch, the sm_120 cubins or `ray.register` are missing
+- [ ] Entrypoint baked in, not bind-mounted
+- [ ] Entrypoint reads its knobs from the environment; `KNOBS.md` regenerated
+- [ ] `build.sh` with a `TAG`; builds from a plain `rsync -a` copy
+
+## smoketest/
+
+- [ ] `run.sh <base> [served-name]`, exit code is the failure count
+- [ ] `lib.sh` byte-identical to the other repos
+- [ ] One `t_*` per case, every check a jq expression
+- [ ] Fails against the wrong model, not only passes against the right one
+
+## Docs
+
+- [ ] README: what it serves, where, how to run it, how to roll back
+- [ ] Measured numbers carry a date
+- [ ] No fleet addresses or box names
+
+## Status ports
+
+Taken, so a new recipe picks a free one. API ports are in each README.
+
+| port | recipe |
+|---|---|
+| 8022 | dgemma |
+| 8081 | ds4-flash |
+| 8082 | glm53 |
+| 8083 | glm53-exl3 |
+| 8084 | qwen38fn (llama.cpp's own status page) |
+| 8181 | qwen36-a3b |
+| 8182 | dots-ocr |
+| 8183 | whisper |
+| 8184 | qwen3-embedding |
+| 8185 | qwen38-flashnext |
+| 8186 | qwen38 |
+| 8090 | the host agent |

docker-images/glm53-kindling/.submodules/spark-agent/README.md

@@ -0,0 +1,40 @@
+# spark-agent
+
+The MCP side of the GB10 / ASUS GX10 Spark fleet, in two parts:
+
+| path | what |
+|---|---|
+| `agent/` | The host agent image. One per node, privileged, serves no model. |
+| `vllm/` | The status server that runs inside a vLLM model container. |
+
+Both register with mentat. mentat-serve merges every group's MCP into one
+endpoint, so clients reach every node and every model container through the
+router, and `__group` picks which one answers.
+
+## agent/
+
+Holds the tools that need the host: every process, ptrace for py-spy, the
+GPU, logs, and the host's snapshots. The model containers then need no added
+capabilities.
+
+    cd agent && TAG=spark-agent:v6 ./build.sh
+    docker compose -f spark-agent.yaml up -d
+
+It registers its MCP endpoint with the local mentatd as group
+`agent-<hostname>`. Its page on `:8090` shows every node and every model,
+read from mentat-serve (`MENTAT_ROUTER_URL`).
+
+## vllm/
+
+`status-server.py` serves a model container's stage page and its MCP tools.
+A recipe takes it as a submodule rather than a copy:
+
+    git submodule add https://github.com/mmastrac/spark-agent.git .submodules/spark-agent
+    ln -s .submodules/spark-agent/vllm vllm
+
+and its Dockerfile copies `vllm/status-server.py`. Its knobs are environment
+variables, and an empty one counts as unset. Set `SERVING_WHEN_READY=1` on a
+single-node engine with no self-test, so it reads as serving once the API
+answers. [MODELS.md](MODELS.md) is
+the checklist every recipe is held to. BuildKit follows a symlink
+that stays inside the build context, so the image gets the file itself.

docker-images/glm53-kindling/.submodules/spark-agent/agent/Dockerfile

@@ -0,0 +1,49 @@
+# The host agent for a GB10 node. Holds the privileged tools so the container
+# running the model does not have to.
+#
+# Small on purpose: stdlib Python plus py-spy. nvidia-smi is not installed
+# here, the NVIDIA container runtime injects it along with the driver
+# libraries when the compose file reserves a GPU.
+# Pulled registry-qualified: a bare `mentat-artifacts:<ver>` tag resolves to a
+# local image, and those have been mislabelled before.
+ARG MENTAT_VERSION=0.12.0
+FROM mmastrac/mentat-artifacts:${MENTAT_VERSION} AS mentat
+FROM python:3.12-slim
+
+# pciutils/usbutils carry pci.ids and usb.ids, without which lspci and lsusb
+# report bare hex vendor:device pairs. iproute2 backs the network tool.
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        iproute2 pciutils usbutils procps \
+    && rm -rf /var/lib/apt/lists/* \
+    && command -v ip && command -v lspci && command -v lsusb && command -v ps
+
+# py-spy reads another process's Python stack without stopping it. Container
+# processes are visible through the host's PID namespace, which the compose
+# file shares, so this reaches the engine without the model container holding
+# any tracing capability itself.
+RUN pip install --no-cache-dir py-spy && py-spy --version
+
+# The accounting helper runs live rather than reading the periodic snapshot,
+# because pid: host lets this container see every process and nvidia-smi gives
+# the GPU figure that process RSS never shows. The host's own memory timer
+# runs the same script.
+COPY spark-memory.py /usr/local/bin/spark-memory.py
+
+# The mentat shim, for ray.register only: the agent announces its MCP endpoint
+# to the local mentatd, and mentat-serve merges it. --no-deps, because the
+# shim installs as `ray` and must not pull real ray in.
+COPY --from=mentat /out/ /tmp/mentat-out/
+RUN pip install --no-cache-dir --no-deps /tmp/mentat-out/mentatd-*-py3-none-any.whl \
+    && rm -rf /tmp/mentat-out \
+    && python3 -c "from ray import register; register.connect, register.parse_args"
+
+COPY spark-agent.py /usr/local/bin/spark-agent.py
+RUN chmod +x /usr/local/bin/spark-agent.py /usr/local/bin/spark-memory.py \
+    && python3 -c "import ast; ast.parse(open('/usr/local/bin/spark-agent.py').read()); \
+ast.parse(open('/usr/local/bin/spark-memory.py').read())"
+
+# 8090 = MCP and the node status endpoint. The compose file decides which
+# addresses it is reachable on.
+EXPOSE 8090
+ENTRYPOINT ["python3", "/usr/local/bin/spark-agent.py"]
+CMD ["--port", "8090"]

docker-images/glm53-kindling/.submodules/spark-agent/agent/build.sh

@@ -0,0 +1,7 @@
+#!/bin/bash
+# Build the host agent image. Run from this directory.
+set -euo pipefail
+cd "$(dirname "$0")"
+TAG="${TAG:-spark-agent:v6}"
+sudo docker build -t "$TAG" -f Dockerfile .
+echo "BUILD_EXIT=$? TAG=$TAG"

docker-images/glm53-kindling/.submodules/spark-agent/agent/spark-agent.py

@@ -0,0 +1,873 @@
+#!/usr/bin/env python3
+"""Host agent: general MCP tools for one node, served through mentat.
+
+Runs beside the model containers rather than inside one, so the capabilities
+that need a host (reading every process, ptrace, kernel messages) leave the
+process that runs the model.
+
+Tools return what the underlying command printed. A parsed view discards the
+fields nobody thought to keep, and on 2026-08-25 the answer to where 100 GiB
+of memory went was in `nvidia-smi --query-compute-apps`, which the curated GPU
+tool did not carry. Correlating several sources is a separate tool, named for
+the question it answers.
+
+The agent registers its MCP endpoint with the local mentatd, and mentat-serve
+merges it with every other node's, so clients reach every node's tools
+through the router. Stdlib only, apart from the mentat shim's ray.register,
+and it runs without that. Prototype on a spare port with:
+    python3 spark-agent.py --port 8090
+"""
+import argparse
+import json
+import os
+import shutil
+import socket
+import subprocess
+import threading
+import time
+import urllib.parse
+import urllib.request
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+
+
+class Server(ThreadingHTTPServer):
+    """The HTTP listener.
+
+    socketserver defaults `request_queue_size` to 5, and that is the listen()
+    backlog. Past five connections the kernel has not yet handed over, it
+    refuses new ones, and the refusal reaches the client as a transport error
+    naming nothing -- it read as model failures through an afternoon of
+    benchmarking before the listener was suspected. It must be a class
+    attribute: TCPServer.__init__ calls listen() before the instance exists to
+    set it on. The kernel clamps to net.core.somaxconn.
+    """
+
+    request_queue_size = 1024
+    daemon_threads = True
+
+HOSTNAME = socket.gethostname()
+STARTED = time.time()
+LOG_DIR = os.environ.get("AGENT_LOG_DIR", "/logs")
+AGENT_PORT = int(os.environ.get("AGENT_PORT", "8090"))
+# Set from --port in main(); the page links peers on whatever port this
+# agent is actually served on rather than assuming the default.
+LISTEN_PORT = AGENT_PORT
+MAX_BYTES = 64 * 1024
+
+# mentat owns the model list: containers register with mentatd, and the
+# mentat-serve router decides what is routable. The agent reads the router's
+# table rather than keeping a registry of its own, so the two cannot disagree.
+MENTAT_ROUTER_URL = os.environ.get("MENTAT_ROUTER_URL", "")
+# How often the router and the engines are read. Token rates are differenced
+# across this interval, so it also sets how twitchy the numbers look.
+LOAD_INTERVAL_S = float(os.environ.get("LOAD_INTERVAL_S", "5"))
+
+
+def run(cmd: list[str], timeout: float = 15) -> str:
+    """Combined stdout and stderr, or a line saying why there is none."""
+    if not shutil.which(cmd[0]):
+        return f"({cmd[0]} is not installed on {HOSTNAME})"
+    try:
+        r = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
+    except subprocess.TimeoutExpired:
+        return f"({' '.join(cmd)} timed out after {timeout}s)"
+    except OSError as e:
+        return f"({' '.join(cmd)} failed: {type(e).__name__}: {e})"
+    return (r.stdout or "") + (r.stderr or "")
+
+
+def read_file(path: str) -> str:
+    try:
+        with open(path) as f:
+            return f.read()
+    except OSError as e:
+        return f"({path}: {type(e).__name__}: {e})"
+
+
+def cap(text: str, limit: int = MAX_BYTES) -> str:
+    b = text.encode("utf-8", "replace")
+    if len(b) <= limit:
+        return text
+    head = b[:limit].decode("utf-8", "ignore")
+    return f"{head}\n\n[truncated at {limit} bytes of {len(b)}]"
+
+
+def stamped(body: str, source: str) -> str:
+    """Every answer says where it came from and when."""
+    return (f"[{HOSTNAME} {time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())}] "
+            f"{source}\n{body}")
+
+
+# --- raw tools -------------------------------------------------------------
+# One command in, its output out. Nothing between.
+
+def t_meminfo(**_):
+    return stamped(read_file("/proc/meminfo"), "cat /proc/meminfo")
+
+
+def t_gpu(args: str = "", **_):
+    """nvidia-smi with the caller's own flags, from a fixed set.
+
+    --query-gpu reports N/A for memory on GB10 while --query-compute-apps
+    still gives a per-pid figure, so the caller needs to pick.
+    """
+    allowed = {
+        "": ["nvidia-smi"],
+        "query": ["nvidia-smi", "--query-gpu=name,utilization.gpu,clocks.sm,"
+                  "temperature.gpu,power.draw,memory.total,memory.used",
+                  "--format=csv"],
+        "apps": ["nvidia-smi", "--query-compute-apps=pid,process_name,used_memory",
+                 "--format=csv"],
+        "topo": ["nvidia-smi", "topo", "-m"],
+        "clocks": ["nvidia-smi", "-q", "-d", "CLOCK,PERFORMANCE"],
+    }
+    if args not in allowed:
+        return f"unknown selection {args!r}. Available: {', '.join(sorted(allowed))}"
+    return stamped(cap(run(allowed[args])), " ".join(allowed[args]))
+
+
+def t_processes(**_):
+    return stamped(cap(run(["ps", "-eo",
+                            "pid,ppid,rss,pcpu,etime,nlwp,comm,args", "--sort=-rss"])),
+                   "ps -eo pid,ppid,rss,pcpu,etime,nlwp,comm,args --sort=-rss")
+
+
+def t_network(**_):
+    parts = [("ip -o -4 addr", run(["ip", "-o", "-4", "addr"])),
+             ("ip -o link", run(["ip", "-o", "link"])),
+             ("rdma link show", run(["rdma", "link", "show"]))]
+    return stamped("\n".join(f"$ {c}\n{o}" for c, o in parts), "network")
+
+
+def t_rdma_counters(device: str = "", **_):
+    base = "/sys/class/infiniband"
+    out = []
+    try:
+        devs = [device] if device else sorted(os.listdir(base))
+    except OSError as e:
+        return stamped(f"({base}: {e})", "rdma counters")
+    for d in devs:
+        for group in ("counters", "hw_counters"):
+            p = f"{base}/{d}/ports/1/{group}"
+            try:
+                names = sorted(os.listdir(p))
+            except OSError:
+                continue
+            out.append(f"# {d}/{group}")
+            for n in names:
+                out.append(f"{n} = {read_file(f'{p}/{n}').strip()}")
+    return stamped(cap("\n".join(out) or "(no infiniband devices)"), "rdma counters")
+
+
+def t_pci(**_):
+    return stamped(cap(run(["lspci", "-nnk"])), "lspci -nnk")
+
+
+def t_usb(**_):
+    return stamped(cap(run(["lsusb", "-t"])), "lsusb -t")
+
+
+def t_docker_ps(**_):
+    """Container inventory, from the host's snapshot.
+
+    Reading it from a file keeps the Docker socket out of this process. The
+    socket is root-equivalent on the host, and this endpoint answers a model.
+    """
+    p = os.path.join(LOG_DIR, "docker-ps.log")
+    try:
+        age = time.time() - os.stat(p).st_mtime
+    except OSError as e:
+        return stamped(f"({p}: {type(e).__name__}: {e})", "docker ps snapshot")
+    return stamped(cap(read_file(p)), f"docker-ps.log, {age:.0f}s old")
+
+
+def t_systemd(**_):
+    """Running services, from the host's snapshot.
+
+    Same reason as t_docker_ps: this image has no systemctl, and the way to
+    give it one is to mount the host's /run/systemd, which grants control and
+    not only a read. So the host writes the list and this reads a file.
+    """
+    p = os.path.join(LOG_DIR, "systemd-units.log")
+    try:
+        age = time.time() - os.stat(p).st_mtime
+    except OSError as e:
+        return stamped(f"({p}: {type(e).__name__}: {e})", "systemd snapshot")
+    return stamped(cap(read_file(p)), f"systemd-units.log, {age:.0f}s old")
+
+
+def t_list_logs(**_):
+    try:
+        names = sorted(os.listdir(LOG_DIR))
+    except OSError as e:
+        return stamped(f"({LOG_DIR}: {e})", "list logs")
+    rows = []
+    for n in names:
+        p = os.path.join(LOG_DIR, n)
+        try:
+            st = os.stat(p)
+        except OSError:
+            continue
+        rows.append(f"{st.st_size:>12}  {time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime(st.st_mtime))}  {n}")
+    return stamped("\n".join(rows) or "(empty)", f"ls -l {LOG_DIR}")
+
+
+def t_read_log(file: str = "", lines: int = 200, grep: str = "", **_):
+    """Tail of one file under the log directory, optionally filtered."""
+    if not file:
+        return "Name a file. Use list_logs to see them."
+    root = os.path.realpath(LOG_DIR)
+    target = os.path.realpath(os.path.join(root, file))
+    if target != root and not target.startswith(root + os.sep):
+        return "That path leaves the log directory."
+    body = read_file(target)
+    out = body.splitlines()
+    if grep:
+        try:
+            import re
+            rx = re.compile(grep)
+            out = [l for l in out if rx.search(l)]
+        except re.error as e:
+            return f"(bad pattern: {e})"
+    tail = out[-max(1, min(int(lines or 200), 5000)):]
+    return stamped(cap("\n".join(tail)), f"{file} ({len(out)} matching lines)")
+
+
+def _engine_pids() -> list[tuple[int, str]]:
+    """(host pid, label) for the engine processes, engine core first.
+
+    These are container processes seen through the host's /proc, which is why
+    the agent can dump them without the container holding CAP_SYS_PTRACE.
+    """
+    out = []
+    try:
+        entries = os.listdir("/proc")
+    except OSError:
+        return out
+    for d in entries:
+        if not d.isdigit():
+            continue
+        try:
+            with open(f"/proc/{d}/cmdline", "rb") as f:
+                cmd = f.read().replace(b"\0", b" ").decode("utf-8", "replace")
+        except OSError:
+            continue
+        if "EngineCore" in cmd:
+            out.insert(0, (int(d), "EngineCore"))
+        elif "RayWorkerProc" in cmd:
+            out.append((int(d), "RayWorkerProc"))
+        elif "/bin/vllm serve" in cmd:
+            out.append((int(d), "APIServer"))
+    return out[:4]
+
+
+PY_SPY = os.environ.get("PY_SPY", "py-spy")
+
+
+def t_engine_stacks(rounds: int = 3, gap_s: float = 5, **_):
+    """Python stacks of the engine processes, sampled several times.
+
+    One sample cannot tell a slow kernel from queue backpressure. CUDA
+    launches are async, so when the GPU stream backs up the launch queue fills
+    and whatever the CPU tries to launch next is what blocks. A frame that
+    stays put across rounds is the kernel actually running. A frame that moves
+    means the cause is upstream of where the sample landed.
+    """
+    rounds = max(1, min(int(rounds or 3), 10))
+    gap_s = max(0.5, min(float(gap_s or 5), 30))
+    pids = _engine_pids()
+    if not pids:
+        return stamped("(no engine processes found)", "engine stacks")
+    out = []
+    for i in range(rounds):
+        if i:
+            time.sleep(gap_s)
+        out.append(f"----- round {i+1}/{rounds} at "
+                   f"{time.strftime('%H:%M:%SZ', time.gmtime())} -----")
+        for pid, label in pids:
+            if label == "APIServer" and i:
+                continue          # waits on the engine by construction
+            out.append(f"===== {label} (host pid {pid}) =====")
+            out.append(run([PY_SPY, "dump", "--pid", str(pid)], timeout=30))
+    return stamped(cap("\n".join(out)), f"{PY_SPY} dump x{rounds}")
+
+
+# --- tools that correlate ---------------------------------------------------
+
+def t_memory_accounting(**_):
+    """Every byte of host memory, by owner.
+
+    Separate from meminfo because it needs three sources: /proc/meminfo,
+    per-process RSS, and the GPU's per-pid figure. GPU allocations skip
+    process RSS entirely, so ps accounts for about 6 GiB of 121 on a serving
+    node and the total only closes once the GPU number is added.
+    """
+    helper = os.environ.get("SPARK_MEMORY_PY", "/usr/local/bin/spark-memory.py")
+    if os.path.exists(helper):
+        return stamped(cap(run(["python3", helper, "-v"], timeout=30)), helper)
+    return stamped(f"({helper} not present)", "memory accounting")
+
+
+def t_models(**_):
+    """The router's view of every model, with the load this agent scraped."""
+    cluster = router_snapshot()
+    with _load_lock:
+        loads = dict(_load)
+    return json.dumps({"router": MENTAT_ROUTER_URL, **cluster,
+                       "load": loads}, indent=2)
+
+
+
+# --- load sampling ----------------------------------------------------------
+# The agent scrapes each engine the router lists rather than the engine
+# reporting its load. That keeps every model image out of this: a container
+# announces where it serves to mentat, and what it is doing is derived here.
+
+_load: dict[str, dict] = {}
+_load_lock = threading.Lock()
+_load_prev: dict[str, tuple[float, float, float]] = {}
+
+# The names are exact and were read off a running engine. kv_cache_usage_perc
+# is NOT gpu_cache_usage_perc, which is what the older status page scraped and
+# which silently reports nothing on this vLLM.
+_GAUGES = {
+    "vllm:num_requests_running": "running",
+    "vllm:num_requests_waiting": "waiting",
+    "vllm:kv_cache_usage_perc": "kv",
+}
+_COUNTERS = {
+    "vllm:prompt_tokens_total": "prompt",
+    "vllm:generation_tokens_total": "gen",
+}
+
+
+def _parse_metrics(text: str) -> dict:
+    """Pull the handful of series worth showing out of Prometheus text.
+
+    Labels are ignored beyond stripping them: one engine per container here,
+    so a bare sum over duplicate series is the right reading.
+    """
+    out: dict[str, float] = {}
+    for line in text.splitlines():
+        if not line or line[0] == "#":
+            continue
+        name, _, rest = line.partition("{")
+        if rest:
+            _, _, value = rest.partition("}")
+        else:
+            name, _, value = line.partition(" ")
+        key = _GAUGES.get(name.strip()) or _COUNTERS.get(name.strip())
+        if not key:
+            continue
+        try:
+            out[key] = out.get(key, 0.0) + float(value.strip())
+        except ValueError:
+            continue
+    return out
+
+
+
+
+def read_interconnect() -> dict:
+    """This node's RoCE rates, from the host snapshot.
+
+    The snapshot samples two counter reads a second apart, which is the only
+    way to say whether the link is busy now rather than how much it has ever
+    carried. Absent on a node with no cable plugged in, which is not an error.
+    """
+    p = os.path.join(LOG_DIR, "interconnect.log")
+    try:
+        age = time.time() - os.stat(p).st_mtime
+        text = read_file(p)
+    except OSError:
+        return {}
+    if age > 300:                       # a snapshot this old describes nothing
+        return {}
+    # The snapshot writes: "<dev> tx <n> MB/s  rx <n> MB/s", so the units are
+    # fields of their own and the rx label lands at index 4, not 3.
+    for line in text.splitlines():
+        parts = line.split()
+        if len(parts) >= 7 and parts[1] == "tx" and parts[4] == "rx":
+            try:
+                return {"dev": parts[0], "tx_mbs": float(parts[2]),
+                        "rx_mbs": float(parts[5]), "age_s": round(age)}
+            except ValueError:
+                continue
+    return {}
+
+
+# How long a driver fault stays interesting. Anything older is history: these
+# lines are never cleared, so without a window a single fault at boot would
+# read as a permanently broken node.
+DRIVER_ERR_WINDOW_S = 900
+
+
+def read_driver_errors() -> dict:
+    """NVRM and Xid lines from the host snapshot, counted and dated.
+
+    Returns {} when the snapshot is missing, stale or clean. `recent` counts
+    the window; `total` is the whole boot, which can be far larger.
+
+    This exists because nothing watched it. On 2026-08-27 the driver logged
+    NV_ERR_NO_MEMORY every few minutes for 23 minutes while the model answered
+    requests at normal speed, and the first visible symptom was a box that had
+    to be power cycled.
+    """
+    p = os.path.join(LOG_DIR, "driver-errors.log")
+    try:
+        age = time.time() - os.stat(p).st_mtime
+        text = read_file(p)
+    except OSError:
+        return {}
+    if age > 300:
+        return {}
+    total, recent, newest, sample = 0, 0, None, ""
+    now = time.time()
+    for line in text.splitlines():
+        if line.startswith("# matched "):
+            try:
+                total = int(line.split()[2])
+            except (IndexError, ValueError):
+                pass
+            continue
+        if not line.startswith("["):
+            continue
+        stamp, _, msg = line[1:].partition("] ")
+        try:
+            # dmesg -T renders "Thu Aug 27 22:13:26 2026" in local time.
+            when = time.mktime(time.strptime(stamp, "%a %b %d %H:%M:%S %Y"))
+        except ValueError:
+            continue
+        if newest is None or when > newest:
+            newest, sample = when, msg.strip()[:200]
+        if now - when <= DRIVER_ERR_WINDOW_S:
+            recent += 1
+    if not total and newest is None:
+        return {}
+    return {"total": total, "recent": recent, "sample": sample,
+            "last_s": round(now - newest) if newest else None}
+
+
+
+
+# --- mentat ------------------------------------------------------------------
+# mentat does the cluster work this agent used to do itself. mentatd knows the
+# nodes, and mentat-serve routes the models and merges every group's MCP. The
+# agent registers its own MCP endpoint as a group of one named for the node, so
+# its host tools reach clients through the router, where `__group` picks the
+# node.
+MENTAT_GROUP = os.environ.get("MENTAT_GROUP", f"agent-{HOSTNAME}")
+MENTAT_DAEMON = os.environ.get("MENTAT_DAEMON", "127.0.0.1:6379")
+# Short, because one loop reads the router, every peer agent and every engine
+# in turn, and a dead node should cost a round seconds, not minutes.
+FETCH_TIMEOUT_S = 2.0
+
+
+def register_loop():
+    """Hold this agent's registration with the local daemon.
+
+    ray.register's main() refuses a registration with no OpenAI endpoint,
+    because one usually means a model that never came up. This agent has no
+    model on purpose, so it calls connect() and keeps main()'s retry shape.
+    """
+    try:
+        from ray import register
+    except ImportError as e:
+        print(f"mentat registration off: {e}", flush=True)
+        return
+    args = register.parse_args([
+        "--address", MENTAT_DAEMON, "--group", MENTAT_GROUP,
+        "--mcp", f"{LISTEN_PORT}/mcp", "--container", "spark-agent"])
+    delay = 1.0
+    while True:
+        try:
+            if register.connect(args) > register.SETTLED_S:
+                delay = 1.0
+        except (OSError, RuntimeError, ValueError) as e:
+            print(f"mentat registration: {type(e).__name__}: {e}", flush=True)
+        time.sleep(delay)
+        delay = min(delay * 2, register.RETRY_MAX_S)
+
+
+def _get_json(url: str) -> dict:
+    with urllib.request.urlopen(url, timeout=FETCH_TIMEOUT_S) as r:
+        return json.loads(r.read().decode())
+
+
+# The last complete read of the cluster: the router's table and what each
+# node's agent reported about itself.
+_cluster: dict = {"router": None, "router_error": "not read yet", "nodes": {}}
+_cluster_lock = threading.Lock()
+
+
+def router_snapshot() -> dict:
+    with _cluster_lock:
+        return dict(_cluster)
+
+
+def read_nodes(daemons: dict) -> dict[str, dict]:
+    """daemon address -> that node's agent status, or why there is none.
+
+    The router's watch set is the node list. mentat knows nothing of the
+    interlink or the driver, so each node's agent is asked for its own.
+    """
+    out = {}
+    for addr, d in sorted(daemons.items()):
+        ip = addr.rsplit(":", 1)[0]
+        entry = {"ip": ip, "addrs": [ip] + [a.rsplit(":", 1)[0]
+                                            for a in d.get("alternates") or []],
+                 "connected": bool(d.get("connected")),
+                 "daemon_error": d.get("error")}
+        try:
+            entry["agent"] = _get_json(f"http://{ip}:{LISTEN_PORT}/status.json")
+        except Exception as e:
+            entry["agent_error"] = f"{type(e).__name__}: {e}"
+        out[addr] = entry
+    return out
+
+
+def watch_cluster():
+    """Read the router and every node's agent, then scrape each engine.
+
+    Token rates come from differencing the totals, because vLLM publishes
+    counters and a counter says nothing about now.
+    """
+    while True:
+        router, err = None, "MENTAT_ROUTER_URL is not set"
+        if MENTAT_ROUTER_URL:
+            try:
+                router, err = _get_json(MENTAT_ROUTER_URL.rstrip("/") + "/"), None
+            except Exception as e:
+                err = f"{type(e).__name__}: {e}"
+        nodes = read_nodes((router or {}).get("daemons") or {})
+        with _cluster_lock:
+            _cluster.update(router=router, router_error=err, nodes=nodes)
+        groups = (router or {}).get("groups") or {}
+        for name, g in groups.items():
+            if g.get("openai"):
+                sample_load(name, g["openai"])
+        with _load_lock:
+            for gone in set(_load) - set(groups):
+                _load.pop(gone, None)
+        time.sleep(LOAD_INTERVAL_S)
+
+
+def sample_load(name: str, url: str):
+    parts = urllib.parse.urlsplit(url)
+    try:
+        with urllib.request.urlopen(
+                f"{parts.scheme}://{parts.netloc}/metrics",
+                timeout=FETCH_TIMEOUT_S) as r:
+            m = _parse_metrics(r.read().decode("utf-8", "replace"))
+    except Exception:
+        with _load_lock:
+            _load.pop(name, None)
+        _load_prev.pop(name, None)
+        return
+    now = time.time()
+    entry = {"running": int(m.get("running", 0)),
+             "waiting": int(m.get("waiting", 0)),
+             "kv": round(m.get("kv", 0.0) * 100, 1)}
+    prev = _load_prev.get(name)
+    if prev:
+        dt = now - prev[0]
+        # A restarted engine resets its counters, so a negative delta is a
+        # new engine rather than a rate. Drop it and re-baseline.
+        if dt >= 1.0:
+            dp, dg = m.get("prompt", 0.0) - prev[1], m.get("gen", 0.0) - prev[2]
+            if dp >= 0 and dg >= 0:
+                entry["prompt_tps"] = round(dp / dt, 1)
+                entry["gen_tps"] = round(dg / dt, 1)
+    _load_prev[name] = (now, m.get("prompt", 0.0), m.get("gen", 0.0))
+    with _load_lock:
+        _load[name] = entry
+
+
+TOOLS = {
+    "meminfo": (t_meminfo, "Raw /proc/meminfo.", {}),
+    "gpu": (t_gpu, "nvidia-smi output. Pick a view with args: '' for the "
+            "default table, 'query' for a csv of utilisation/clocks/temp/power, "
+            "'apps' for per-pid GPU memory, 'topo' for the topology matrix, "
+            "'clocks' for clock and performance detail. On GB10 --query-gpu "
+            "reports N/A for memory while 'apps' still answers.",
+            {"args": {"type": "string"}}),
+    "processes": (t_processes, "Raw ps, sorted by RSS. GPU allocations do not "
+                  "appear here.", {}),
+    "network": (t_network, "Raw ip addr, ip link and rdma link output.", {}),
+    "rdma_counters": (t_rdma_counters, "Every RoCE counter for a port, raw. "
+                      "Rates need two reads: these are lifetime totals.",
+                      {"device": {"type": "string"}}),
+    "pci": (t_pci, "Raw lspci -nnk, including the bound driver.", {}),
+    "usb": (t_usb, "Raw lsusb -t.", {}),
+    "docker_ps": (t_docker_ps, "Raw docker ps -a.", {}),
+    "systemd_units": (t_systemd, "Running systemd services.", {}),
+    "list_logs": (t_list_logs, "Files in the log directory with size and mtime.", {}),
+    "read_log": (t_read_log, "Tail of one log file, optionally filtered by a "
+                 "regex. Capped at 64 KB.",
+                 {"file": {"type": "string"}, "lines": {"type": "integer"},
+                  "grep": {"type": "string"}}),
+    "memory_accounting": (t_memory_accounting, "Host memory by owner, "
+                          "reconciled against MemTotal. Needs /proc/meminfo, "
+                          "per-process RSS and the GPU's per-pid figure "
+                          "together, which is why it is not raw meminfo.", {}),
+    "engine_stacks": (t_engine_stacks, "Python stacks of the engine processes, "
+                      "sampled over several rounds. A frame that stays put is "
+                      "the kernel running; one that moves is queue "
+                      "backpressure, since an async CUDA launch blocks on "
+                      "whatever it tries to launch next. Reads container "
+                      "processes through the host's /proc.",
+                      {"rounds": {"type": "integer"},
+                       "gap_s": {"type": "number"}}),
+    "models": (t_models, "mentat-serve's table of every group and model, "
+               "with the token load this agent scraped from each engine.", {}),
+}
+
+
+
+
+# --- cluster page -----------------------------------------------------------
+# Every agent reads the same router, so any node can draw the cluster and
+# there is no page that only one node can serve.
+_CSS = """body{font:14px/1.5 ui-monospace,SFMono-Regular,Menlo,monospace;
+max-width:60rem;margin:2rem auto;padding:0 1rem;background:#111;color:#ddd}
+h1{font-size:1.2rem}h2{font-size:1rem;margin-top:2rem;color:#9ad}
+table{border-collapse:collapse;width:100%;margin:.5rem 0}
+th,td{text-align:left;padding:.3rem .6rem;border-bottom:1px solid #333}
+th{color:#888;font-weight:normal}.ok{color:#7c7}.bad{color:#e77}.dim{color:#777}
+.self{color:#fc6}"""
+_DIM_DASH = "<span class=dim>-</span>"
+
+
+def _esc(v) -> str:
+    return (str(v).replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;"))
+
+
+def _driver_cell(dv: dict) -> str:
+    # A fault inside the window is the alarm. Older ones stay visible but dim,
+    # because the count never resets short of a reboot.
+    if dv.get("recent"):
+        return (f"<span class=bad title=\"{_esc(dv.get('sample') or '')}\">"
+                f"{dv['recent']} in {DRIVER_ERR_WINDOW_S // 60}m</span>")
+    if dv.get("total"):
+        return f"<span class=dim>{dv['total']} earlier</span>"
+    return _DIM_DASH
+
+
+def cluster_page() -> bytes:
+    snap = router_snapshot()
+    router = snap.get("router") or {}
+    groups = router.get("groups") or {}
+    with _load_lock:
+        loads = dict(_load)
+
+    # An engine URL names an address. Map each address back to the node that
+    # holds it, so the models table can say where a model runs.
+    by_ip: dict[str, str] = {}
+    for n in snap.get("nodes", {}).values():
+        name = (n.get("agent") or {}).get("node") or n["ip"]
+        for a in n["addrs"]:
+            by_ip[a] = name
+
+    def node_of(url: str | None) -> str:
+        host = urllib.parse.urlsplit(url or "").hostname or ""
+        return by_ip.get(host, host or "?")
+
+    served = {}
+    for g in groups.values():
+        if g.get("openai"):
+            n = node_of(g["openai"])
+            served[n] = served.get(n, 0) + 1
+
+    rows = []
+    for n in snap.get("nodes", {}).values():
+        agent = n.get("agent") or {}
+        name = agent.get("node") or n["ip"]
+        link = (f"<a href=\"http://{_esc(n['ip'])}:{LISTEN_PORT}/\" "
+                f"style=\"color:#9ad\">{_esc(name)}</a>")
+        if name == HOSTNAME:
+            link += " <span class=self>this node</span>"
+        mentat = ("<span class=ok>connected</span>" if n["connected"] else
+                  f"<span class=bad>{_esc(n.get('daemon_error') or 'down')}</span>")
+        if "agent" in n:
+            ic = agent.get("interconnect") or {}
+            ics = (f"{ic['tx_mbs']:.0f}/{ic['rx_mbs']:.0f} MB/s"
+                   if ic else _DIM_DASH)
+            drv = _driver_cell(agent.get("driver") or {})
+        else:
+            ics = f"<span class=bad title=\"{_esc(n.get('agent_error'))}\">no agent</span>"
+            drv = _DIM_DASH
+        rows.append(
+            f"<tr><td>{link}</td><td>{_esc(', '.join(n['addrs']))}</td>"
+            f"<td>{mentat}</td><td>{ics}</td><td>{drv}</td>"
+            f"<td>{served.get(name, 0)}</td></tr>")
+
+    models = []
+    for name, g in sorted(groups.items()):
+        if not g.get("openai"):
+            continue                    # agents and other MCP-only groups
+        healthy = bool(g.get("healthy"))
+        ld = loads.get(name) or {}
+        if "gen_tps" in ld:
+            # Prompt and generation are separate on purpose: a long prefill
+            # shows as prompt tokens with generation at zero, which is the
+            # difference between a busy engine and a stalled one.
+            tps = f"{ld['prompt_tps']:.0f} in / {ld['gen_tps']:.0f} out"
+        elif ld:
+            tps = "<span class=dim>warming</span>"
+        else:
+            tps = _DIM_DASH
+        q = f"{ld['running']}/{ld['waiting']}" if ld else _DIM_DASH
+        kv = f"{ld['kv']:.0f}%" if ld else _DIM_DASH
+        models.append(
+            f"<tr><td>{_esc(', '.join(g.get('models') or [name]))}</td>"
+            f"<td>{_esc(name)}</td><td>{_esc(node_of(g['openai']))}</td>"
+            f"<td class={'ok' if healthy else 'bad'}>"
+            f"{'healthy' if healthy else _esc(g.get('why_not') or 'not ready')}</td>"
+            f"<td>{tps}</td><td>{q}</td><td>{kv}</td>"
+            f"<td class=dim>{_esc(g['openai'])}</td></tr>")
+
+    err = snap.get("router_error")
+    router_note = (f"<p class=bad>router unreachable: {_esc(err)}</p>" if err else "")
+    empty = "router unreachable" if err else "none routed"
+    html = f"""<!doctype html><meta charset=utf-8>
+<meta http-equiv=refresh content=10>
+<title>spark cluster</title><style>{_CSS}</style>
+<h1>spark cluster &middot; seen from {_esc(HOSTNAME)}</h1>
+{router_note}
+<h2>nodes</h2>
+<table><tr><th>node</th><th>addresses</th><th>mentatd</th>
+<th>interlink tx/rx</th><th>driver faults</th><th>models</th></tr>
+{''.join(rows) or '<tr><td colspan=6 class=dim>no daemons known</td></tr>'}</table>
+<h2>models</h2>
+<table><tr><th>model</th><th>group</th><th>node</th><th>state</th>
+<th>tokens/s</th><th>run/wait</th><th>kv</th>
+<th>serving</th></tr>{''.join(models) or f'<tr><td colspan=8 class=dim>{empty}</td></tr>'}</table>
+<p class=dim>from <a href="{_esc(MENTAT_ROUTER_URL)}/" style="color:#9ad">{_esc(MENTAT_ROUTER_URL)}</a>
+&middot; refreshes every 10s
+&middot; <a href="/status.json" style="color:#9ad">status.json</a></p>"""
+    return html.encode()
+
+
+# --- who may ask ------------------------------------------------------------
+# The network is the only control: the endpoint answers our own subnets and
+# the router. Docker bridge ranges are absent on purpose, so a container on
+# this host cannot read host state just by being here. A bridge-networked
+# client still reaches the agent when it dials a real host address, because
+# Docker masquerades the source to the host's own IP on that route.
+# Address prefixes, not CIDR blocks. Loopback alone by default: the compose
+# file takes the real subnets from the node's .env.
+ALLOWED_SOURCES = [p.strip() for p in os.environ.get(
+    "ALLOWED_SOURCES", "127.0.0.1").split(",") if p.strip()]
+# The router forwards every merged tool call, so its address is always let in.
+ROUTER_HOST = urllib.parse.urlsplit(MENTAT_ROUTER_URL).hostname or ""
+
+
+def source_allowed(addr: str) -> bool:
+    return addr == ROUTER_HOST or any(addr.startswith(p) for p in ALLOWED_SOURCES)
+
+
+class Handler(BaseHTTPRequestHandler):
+    protocol_version = "HTTP/1.1"
+
+    def log_message(self, *a):
+        pass
+
+    def _deny_source(self) -> bool:
+        addr = self.client_address[0]
+        if source_allowed(addr):
+            return False
+        self._send(403, b'{"error":"source not permitted"}\n', "application/json")
+        return True
+
+    def _send(self, code: int, body: bytes, ctype: str):
+        self.send_response(code)
+        self.send_header("Content-Type", ctype)
+        self.send_header("Content-Length", str(len(body)))
+        self.end_headers()
+        self.wfile.write(body)
+
+    def do_GET(self):
+        if self._deny_source():
+            return
+        path = self.path.rstrip("/") or "/"
+        if path == "/":
+            self._send(200, cluster_page(), "text/html; charset=utf-8")
+            return
+        if path in ("/healthz", "/status.json"):
+            # Peer agents read this for the node table, so it carries this
+            # node's interlink and driver figures.
+            self._send(200, json.dumps({
+                "node": HOSTNAME, "uptime_s": round(time.time() - STARTED, 1),
+                "group": MENTAT_GROUP, "tools": sorted(TOOLS),
+                "interconnect": read_interconnect(),
+                "driver": read_driver_errors(),
+            }, indent=2).encode(), "application/json")
+            return
+        self._send(404, b"not found\n", "text/plain")
+
+    def do_POST(self):
+        if self._deny_source():
+            return
+        path = self.path.rstrip("/") or "/"
+        length = int(self.headers.get("Content-Length") or 0)
+        raw = self.rfile.read(length) if length else b"{}"
+        try:
+            payload = json.loads(raw.decode() or "{}")
+        except ValueError:
+            self._send(400, b'{"error":"bad json"}', "application/json")
+            return
+        if path == "/mcp":
+            self._send(200, json.dumps(mcp(payload)).encode(), "application/json")
+            return
+        self._send(404, b'{"error":"not found"}', "application/json")
+
+
+def mcp(req: dict) -> dict:
+    rid, method = req.get("id"), req.get("method")
+    def ok(result): return {"jsonrpc": "2.0", "id": rid, "result": result}
+
+    if method == "initialize":
+        return ok({"protocolVersion": "2024-11-05",
+                   "capabilities": {"tools": {"listChanged": False}},
+                   "serverInfo": {"name": f"spark-agent@{HOSTNAME}",
+                                  "version": "0.2"}})
+    if method == "tools/list":
+        return ok({"tools": [{"name": n, "description": d,
+                              "inputSchema": {"type": "object",
+                                              "properties": dict(p)}}
+                             for n, (_, d, p) in sorted(TOOLS.items())]})
+    if method == "tools/call":
+        params = req.get("params") or {}
+        name = params.get("name")
+        if name not in TOOLS:
+            return ok({"content": [{"type": "text",
+                                    "text": f"no tool {name!r} on {HOSTNAME}"}],
+                       "isError": True})
+        try:
+            text = TOOLS[name][0](**(params.get("arguments") or {}))
+        except Exception as e:
+            return ok({"content": [{"type": "text",
+                                    "text": f"{type(e).__name__}: {e}"}],
+                       "isError": True})
+        return ok({"content": [{"type": "text", "text": text}], "isError": False})
+    return {"jsonrpc": "2.0", "id": rid,
+            "error": {"code": -32601, "message": f"no method {method!r}"}}
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--port", type=int, default=8090)
+    ap.add_argument("--bind", default="0.0.0.0")
+    a = ap.parse_args()
+    global LISTEN_PORT
+    LISTEN_PORT = a.port
+    threading.Thread(target=watch_cluster, daemon=True).start()
+    threading.Thread(target=register_loop, daemon=True).start()
+    srv = Server((a.bind, a.port), Handler)
+    print(f"spark-agent on {a.bind}:{a.port} as {HOSTNAME}, group "
+          f"{MENTAT_GROUP}, {len(TOOLS)} tools, logs at {LOG_DIR}", flush=True)
+    srv.serve_forever()
+
+
+if __name__ == "__main__":
+    main()

docker-images/glm53-kindling/.submodules/spark-agent/agent/spark-agent.yaml

@@ -0,0 +1,58 @@
+# The agent runs on every node, including ones with no model on them.
+#
+#   docker compose -f spark-agent.yaml up -d
+#
+# It holds what the model container gave up. Keep the two asymmetric: this one
+# is privileged and serves no model, that one serves a model and is not.
+services:
+  spark-agent:
+    image: ${IMAGE:-spark-agent:v6}
+    container_name: spark-agent
+    restart: unless-stopped
+
+    # mentatd's control port is on 127.0.0.1, and mentat-serve dials the
+    # agent's MCP port on the host's own addresses. Both want the host's
+    # network.
+    network_mode: host
+
+    # Container processes appear in the host's PID namespace, which is what
+    # lets py-spy here reach an engine running in another container. This is
+    # the whole reason the model container can drop CAP_SYS_PTRACE.
+    pid: host
+    # SYS_PTRACE and nothing else. Sharing the host PID namespace does not
+    # widen the capability set on its own, so dropping the default 14 first
+    # takes this from 15 capabilities to 1. Read /proc/self/status to check:
+    # under pid: host, /proc/1 is the host's init rather than this process.
+    cap_drop:
+      - ALL
+    cap_add:
+      - SYS_PTRACE
+
+    # Read-only throughout: the agent observes and never changes the host.
+    # /proc needs no mount, because pid: host already makes the container's
+    # own /proc the host's. No Docker socket either: it is root-equivalent,
+    # and this endpoint answers a model. The host snapshot timer writes
+    # docker-ps.log for the container inventory instead.
+    volumes:
+      - /sys:/sys:ro
+      - ${LOG_DIR:?set LOG_DIR in .env to the host log directory}:/logs:ro
+
+    environment:
+      - AGENT_LOG_DIR=/logs
+      # nvidia-smi and the driver libraries arrive through the container
+      # runtime rather than the image, so the image stays small and the
+      # versions cannot drift from the host's.
+      - NVIDIA_VISIBLE_DEVICES=all
+      # The router the page reads models from. Its address is also let in
+      # to the MCP port, because it forwards every merged tool call.
+      - MENTAT_ROUTER_URL=${MENTAT_ROUTER_URL:?set MENTAT_ROUTER_URL in .env}
+      # Address prefixes allowed to reach the page and MCP: the node subnets.
+      - ALLOWED_SOURCES=${ALLOWED_SOURCES:?set ALLOWED_SOURCES in .env}
+
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: all
+              capabilities: [gpu]

docker-images/glm53-kindling/.submodules/spark-agent/agent/spark-memory.py

@@ -0,0 +1,149 @@
+#!/usr/bin/env python3
+"""Account for every byte of host memory on a GB10 node.
+
+The GPU shares the host's memory, and its allocations do not appear in process
+RSS: on one node the engine holds about 100 GiB that `ps` reports as 6 GiB
+across every process on the box. `nvidia-smi --query-gpu=memory.used` answers
+N/A here, but `--query-compute-apps` still gives a per-pid figure, which is
+what makes the total reconcile.
+
+Prints one table that sums to MemTotal, with the residual named rather than
+hidden, so a number that moves after a config change has somewhere to land.
+"""
+import os
+import re
+import subprocess
+import sys
+from collections import defaultdict
+
+GIB = 1024 ** 3
+
+
+def meminfo() -> dict[str, int]:
+    out = {}
+    with open("/proc/meminfo") as f:
+        for line in f:
+            k, _, v = line.partition(":")
+            out[k] = int(v.split()[0]) * 1024
+    return out
+
+
+def gpu_by_pid() -> dict[int, int]:
+    """Bytes of GPU memory per pid, empty when the driver declines to say."""
+    try:
+        r = subprocess.run(
+            ["nvidia-smi", "--query-compute-apps=pid,used_memory",
+             "--format=csv,noheader,nounits"],
+            capture_output=True, text=True, timeout=15)
+    except (OSError, subprocess.SubprocessError):
+        return {}
+    out = {}
+    for line in r.stdout.splitlines():
+        parts = [p.strip() for p in line.split(",")]
+        if len(parts) == 2 and parts[0].isdigit() and parts[1].isdigit():
+            out[int(parts[0])] = int(parts[1]) * 1024 * 1024
+    return out
+
+
+def classify(pid: int, cmd: str) -> str:
+    try:
+        with open(f"/proc/{pid}/cgroup") as f:
+            cg = f.read()
+    except OSError:
+        cg = ""
+    in_docker = "docker" in cg or "containerd" in cg
+    if "EngineCore" in cmd or re.search(r"\bvllm\b", cmd):
+        return "vllm engine"
+    if "RayWorkerProc" in cmd or "ray::" in cmd:
+        return "ray workers"
+    if "raylet" in cmd or "gcs_server" in cmd or "/ray/" in cmd:
+        return "ray system"
+    if "status-server.py" in cmd:
+        return "status server"
+    if in_docker:
+        return "other container processes"
+    if "dockerd" in cmd or "containerd" in cmd:
+        return "docker daemons"
+    if pid == 1 or "/systemd" in cmd or "systemd-" in cmd:
+        return "systemd and units"
+    if "sshd" in cmd:
+        return "ssh sessions"
+    return "other host processes"
+
+
+def processes() -> tuple[dict[str, int], dict[str, int], dict[str, list]]:
+    """(rss by role, gpu by role, notable pids by role)."""
+    gpu = gpu_by_pid()
+    rss_by, gpu_by = defaultdict(int), defaultdict(int)
+    top = defaultdict(list)
+    for entry in os.listdir("/proc"):
+        if not entry.isdigit():
+            continue
+        pid = int(entry)
+        try:
+            with open(f"/proc/{pid}/cmdline", "rb") as f:
+                cmd = f.read().replace(b"\0", b" ").decode("utf-8", "replace")
+            with open(f"/proc/{pid}/statm") as f:
+                rss = int(f.read().split()[1]) * os.sysconf("SC_PAGE_SIZE")
+        except (OSError, ValueError, IndexError):
+            continue
+        role = classify(pid, cmd)
+        rss_by[role] += rss
+        g = gpu.get(pid, 0)
+        gpu_by[role] += g
+        if g or rss > 256 * 1024 * 1024:
+            top[role].append((g or rss, pid, cmd.strip()[:60], bool(g)))
+    return rss_by, gpu_by, top
+
+
+def main() -> int:
+    m = meminfo()
+    total = m["MemTotal"]
+    rss_by, gpu_by, top = processes()
+
+    gpu_total = sum(gpu_by.values())
+    # Process RSS on this box excludes the GPU allocation, so the two add
+    # rather than overlap. Shmem is already inside Cached.
+    rss_total = sum(rss_by.values())
+    cache = m["Cached"] + m["Buffers"] - m["Shmem"]
+    kernel = (m["SUnreclaim"] + m["KernelStack"] + m["PageTables"]
+              + m.get("VmallocUsed", 0) + m.get("Percpu", 0))
+    reclaimable = m["SReclaimable"]
+    free = m["MemFree"]
+
+    rows = [("GPU allocations", gpu_total), ("process memory", rss_total),
+            ("page cache", cache), ("tmpfs and shared", m["Shmem"]),
+            ("kernel", kernel), ("kernel, reclaimable", reclaimable),
+            ("free", free)]
+    named = sum(v for _, v in rows)
+    rows.append(("unaccounted", total - named))
+
+    print(f"{'MemTotal':<28}{total / GIB:>9.2f} GiB\n")
+    for label, v in rows:
+        print(f"  {label:<26}{v / GIB:>9.2f} GiB{v / total:>8.1%}")
+
+    print(f"\n{'GPU allocations by role':<28}")
+    for role, v in sorted(gpu_by.items(), key=lambda kv: -kv[1]):
+        if v:
+            print(f"  {role:<26}{v / GIB:>9.2f} GiB")
+    print(f"\n{'process memory by role':<28}")
+    for role, v in sorted(rss_by.items(), key=lambda kv: -kv[1]):
+        if v > 32 * 1024 * 1024:
+            print(f"  {role:<26}{v / GIB:>9.2f} GiB")
+
+    if "-v" in sys.argv:
+        print("\nlargest single processes")
+        flat = [x for lst in top.values() for x in lst]
+        for size, pid, cmd, is_gpu in sorted(flat, reverse=True)[:12]:
+            kind = "GPU" if is_gpu else "RSS"
+            print(f"  {size / GIB:>7.2f} GiB {kind}  pid {pid:<8} {cmd}")
+
+    if m.get("SwapTotal"):
+        used = m["SwapTotal"] - m["SwapFree"]
+        print(f"\nswap  {used / GIB:.2f} of {m['SwapTotal'] / GIB:.2f} GiB used, "
+              f"{m.get('SwapCached', 0) / GIB:.2f} GiB cached")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

docker-images/glm53-kindling/.submodules/spark-agent/vllm/status-server.py

@@ -0,0 +1,1522 @@
+#!/usr/bin/env python3
+"""Tiny status page and engine MCP tools for a vLLM container.
+
+Runs on its own port from container start, so there is something to look at
+while the model loads. It does NOT proxy inference traffic -- putting a Python hop in front of vLLM would add
+a failure mode and latency to every request for no benefit.
+
+Stage is read from a file the entrypoint appends to, so the page survives the
+entrypoint being mid-exec.
+
+On a multi-node engine every rank polls its PEERS and renders the whole
+cluster, so any box answers "is the cluster up?" without checking the others by
+hand. A TP rank that never joined is the common failure, and it is only visible
+by comparing the nodes. A single-node engine sets no PEERS.
+
+Every setting is an environment variable, and an empty one counts as unset:
+compose passes knobs through as `${VAR:-}`, so an empty string means "use the
+default", never "the value is empty".
+"""
+import glob
+import http.server
+import importlib.metadata
+import importlib.util
+import json
+import os
+import re
+import shutil
+import socket
+import socketserver
+import subprocess
+import threading
+import time
+import urllib.error
+import urllib.request
+
+def _env(name: str, default: str = "") -> str:
+    """The variable's value, or the default when it is unset OR empty."""
+    v = os.environ.get(name, "").strip()
+    return v if v else default
+
+
+SERVICE = _env("SERVICE_NAME", _env("SERVED_NAME", "vllm"))
+STAGE_FILE = _env("STAGE_FILE", f"/tmp/{SERVICE}-stage")
+PORT = int(_env("STATUS_PORT", "8081"))
+# vLLM's own default port. Every recipe sets PORT; a default that named another
+# model's port would scrape that engine on a shared box without any error.
+VLLM_PORT = int(_env("PORT", "8000"))
+START = time.time()
+HOSTNAME = socket.gethostname()
+
+ROLE = _env("ROLE", "head")
+
+# Other ranks to poll, "host:port,host:port", for a multi-node engine. Empty
+# for a single-node one: this node then renders alone.
+PEERS = [p.strip() for p in _env("PEERS").split(",") if p.strip()]
+PEER_POLL_S = float(_env("PEER_POLL_S", "4"))
+PEER_TIMEOUT_S = float(_env("PEER_TIMEOUT_S", "2"))
+
+# The worker holds a TP rank and never serves an API, so it needs its own
+# vocabulary -- otherwise it sits forever on a head-oriented stage and looks
+# stuck when it is healthy.
+WORKER_STAGES = [
+    ("starting", "Container started"),
+    ("joining", "Joining the Ray cluster"),
+    ("worker-ready", "Holding TP rank -- ready (no API on this node)"),
+]
+
+# One label covers everything between exec and the API binding: weight
+# loading, KV cache setup, cudagraph capture and warmup. Weights are a small
+# part of that now (~75s of a ~700s boot), so naming only them made the page
+# look stuck on a step that had long finished.
+_DEFAULT_STAGES = [
+    ("starting", "Container started"),
+    ("waiting-workers", "Waiting for TP workers to join"),
+    ("loading", "Loading model weights and initializing engine"),
+    ("self-test", "Self-test: checking the model answers correctly"),
+    ("serving", "Serving"),
+]
+
+
+def _parse_stages(spec: str) -> list[tuple[str, str]]:
+    """`key:label` pairs, comma separated, in the order the entrypoint writes
+    them. A model with a different boot sequence sets STAGES rather than
+    carrying its own copy of this file. A label cannot contain a comma."""
+    out = []
+    for part in spec.split(","):
+        key, _, label = part.partition(":")
+        if key.strip():
+            out.append((key.strip(), label.strip() or key.strip()))
+    return out
+
+
+STAGES = _parse_stages(_env("STAGES")) if _env("STAGES") else _DEFAULT_STAGES
+
+# 1 marks the node serving as soon as the engine answers /v1/models. For an
+# entrypoint that runs no self-test and execs into vLLM, which leaves nothing
+# behind to write "serving". Off by default: a model with a self-test must not
+# be called healthy before it passes.
+SERVING_WHEN_READY = _env("SERVING_WHEN_READY", "0") == "1"
+
+# Terminal failure states. Kept out of STAGES so they never render as a step to
+# be reached -- they replace the current step and turn the page red.
+FAILED_STAGES = {"self-test-failed"}
+
+
+def stages():
+    return WORKER_STAGES if ROLE == "worker" else STAGES
+
+
+def read_stage():
+    try:
+        with open(STAGE_FILE) as f:
+            lines = [l.strip() for l in f if l.strip()]
+        return lines[-1] if lines else "starting", lines
+    except OSError:
+        return "starting", []
+
+
+# Cache directories to report. These are the things that populate slowly during
+# a cold start -- FlashInfer JIT is ~170 MB and dominates the wait -- so showing
+# their size and growth turns an opaque "loading" into visible progress.
+# Resolve FlashInfer's cache the way FlashInfer itself does, rather than from
+# FLASHINFER_CACHE_DIR -- that name is a module constant, not an env var, so
+# reading it here reported a directory nothing ever wrote to and made a cold
+# cache look permanently warm.
+_FI_BASE = _env("FLASHINFER_WORKSPACE_BASE", os.path.expanduser("~"))
+FLASHINFER_DIR = os.path.join(_FI_BASE, ".cache", "flashinfer")
+
+CACHE_DIRS = [
+    ("flashinfer JIT", FLASHINFER_DIR),
+    ("flashinfer autotune", _env("VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR")),
+    ("torch.compile", _env("VLLM_CACHE_ROOT")),
+    ("triton", _env("TRITON_CACHE_DIR")),
+    ("tilelang", _env("TILELANG_CACHE_DIR")),
+]
+_size_cache: dict[str, tuple[float, int, float]] = {}
+_baseline: dict[str, int] = {}
+
+# A cache directory written to within this many seconds is "active". Size alone
+# is a poor progress signal -- a cache can be re-writing entries at the same
+# total size, which reads as "idle" when it is the thing holding up the boot.
+# The newest mtime says which stage is doing work RIGHT NOW.
+ACTIVE_WINDOW_S = float(_env("CACHE_ACTIVE_WINDOW_S", "25"))
+CACHE_ROOT_PATH = _env("CACHE_ROOT", "/root/.cache")
+
+
+def dir_stat(path: str, ttl: float = 4.0) -> tuple[int, float]:
+    """(bytes, newest mtime). One walk for both; memoised so a 5s refresh is cheap."""
+    if not path or not os.path.isdir(path):
+        return -1, 0.0
+    now = time.time()
+    hit = _size_cache.get(path)
+    if hit and now - hit[0] < ttl:
+        return hit[1], hit[2]
+    total, newest = 0, 0.0
+    for root, _, files in os.walk(path):
+        for f in files:
+            try:
+                st = os.stat(os.path.join(root, f))
+            except OSError:
+                continue
+            total += st.st_size
+            if st.st_mtime > newest:
+                newest = st.st_mtime
+    _size_cache[path] = (now, total, newest)
+    _baseline.setdefault(path, total)
+    return total, newest
+
+
+def human(n: int) -> str:
+    if n < 0:
+        return "-"
+    for unit in ("B", "KB", "MB", "GB"):
+        if n < 1024 or unit == "GB":
+            return f"{n:.0f} {unit}" if unit == "B" else f"{n:.1f} {unit}"
+        n /= 1024.0
+    return f"{n:.1f} GB"
+
+
+def cache_rows():
+    out = []
+    now = time.time()
+    for label, path in CACHE_DIRS:
+        sz, newest = dir_stat(path)
+        grown = sz - _baseline.get(path, sz) if sz >= 0 else 0
+        idle = (now - newest) if newest else None
+        out.append({"name": label, "path": path, "bytes": sz,
+                    "size": human(sz), "grown": human(grown) if grown > 0 else "",
+                    "active": idle is not None and idle < ACTIVE_WINDOW_S,
+                    "idle_s": None if idle is None else int(idle)})
+    return out
+
+
+def vllm_up():
+    try:
+        with urllib.request.urlopen(
+            f"http://127.0.0.1:{VLLM_PORT}/v1/models", timeout=2
+        ) as r:
+            return r.status == 200
+    except Exception:
+        return False
+
+
+THROUGHPUT_INTERVAL_S = float(_env("THROUGHPUT_INTERVAL_S", "60"))
+_tp_lock = threading.Lock()
+_tp: list[dict] = []          # newest last, bounded
+_TP_KEEP = 15
+
+
+def scrape_metrics() -> dict | None:
+    """Counters we need for a rate, or None if the API is not up yet."""
+    try:
+        with urllib.request.urlopen(
+            f"http://127.0.0.1:{VLLM_PORT}/metrics", timeout=5
+        ) as r:
+            text = r.read().decode()
+    except Exception:
+        return None
+    out = {"t": time.time()}
+    wanted = {
+        "generation_tokens_total": "gen",
+        "prompt_tokens_total": "prompt",
+        "num_requests_running": "running",
+        "num_requests_waiting": "waiting",
+        "spec_decode_num_draft_tokens_total": "draft",
+        "spec_decode_num_accepted_tokens_total": "accepted",
+        "kv_cache_usage_perc": "kv_usage",
+        "prefix_cache_queries_total": "pc_q",
+        "prefix_cache_hits_total": "pc_hit",
+        # Histograms: _sum and _count together give a mean, and differencing
+        # them between samples gives a mean over the WINDOW rather than over
+        # the process lifetime -- which is what "is it slow right now" needs.
+        "time_to_first_token_seconds_sum": "ttft_s",
+        "time_to_first_token_seconds_count": "ttft_n",
+        "inter_token_latency_seconds_sum": "itl_s",
+        "inter_token_latency_seconds_count": "itl_n",
+        "request_time_per_output_token_seconds_sum": "tpot_s",
+        "request_time_per_output_token_seconds_count": "tpot_n",
+    }
+    finished = {}
+    waiting_why = {}
+    for line in text.splitlines():
+        if not line.startswith("vllm:"):
+            continue
+        # request_success_total is labelled by finished_reason (stop, length,
+        # abort). Aborts are the interesting one -- they are how a failure
+        # actually shows up here.
+        # WHY a request is queued decides which knob can help:
+        #   capacity -> blocked on KV cache space; capping prefill chunk size
+        #               does nothing, you need KV or shorter contexts
+        #   deferred -> held back by scheduling; long_prefill_token_threshold
+        #               and the batched-token budget are the levers
+        # Without this the two are indistinguishable from "waiting: N".
+        if line.startswith("vllm:num_requests_waiting_by_reason{"):
+            m = re.search(r'reason="([^"]+)"', line)
+            if m:
+                try:
+                    v = float(line.rsplit(" ", 1)[1])
+                    if v > 0:
+                        waiting_why[m.group(1)] = v
+                except (IndexError, ValueError):
+                    pass
+            continue
+        if line.startswith("vllm:request_success_total{"):
+            m = re.search(r'finished_reason="([^"]+)"', line)
+            if m:
+                try:
+                    finished[m.group(1)] = float(line.rsplit(" ", 1)[1])
+                except (IndexError, ValueError):
+                    pass
+            continue
+        for key, short in wanted.items():
+            if line.startswith(f"vllm:{key}{{"):
+                try:
+                    out[short] = float(line.rsplit(" ", 1)[1])
+                except (IndexError, ValueError):
+                    pass
+    if finished:
+        out["finished"] = finished
+    out["waiting_why"] = waiting_why
+    return out if "gen" in out else None
+
+
+def throughput_loop():
+    prev = None
+    while True:
+        cur = scrape_metrics()
+        if cur and prev:
+            dt = cur["t"] - prev["t"]
+            if dt > 0:
+                d_draft = cur.get("draft", 0) - prev.get("draft", 0)
+                d_acc = cur.get("accepted", 0) - prev.get("accepted", 0)
+
+                def win_mean(sum_k, cnt_k):
+                    """Mean over this window from two cumulative histogram fields."""
+                    ds = cur.get(sum_k, 0) - prev.get(sum_k, 0)
+                    dn = cur.get(cnt_k, 0) - prev.get(cnt_k, 0)
+                    return round(ds / dn, 4) if dn > 0 else None
+
+                d_q = cur.get("pc_q", 0) - prev.get("pc_q", 0)
+                d_h = cur.get("pc_hit", 0) - prev.get("pc_hit", 0)
+                d_kv = (cur.get("kv_usage", 0) or 0) - (prev.get("kv_usage", 0) or 0)
+                fin_now = cur.get("finished") or {}
+                fin_prev = prev.get("finished") or {}
+                d_fin = {k: fin_now[k] - fin_prev.get(k, 0) for k in fin_now
+                         if fin_now[k] - fin_prev.get(k, 0) > 0}
+
+                with _tp_lock:
+                    _tp.append({
+                        "t": cur["t"],
+                        "decode": round((cur["gen"] - prev["gen"]) / dt, 1),
+                        "prefill": round((cur.get("prompt", 0) - prev.get("prompt", 0)) / dt, 1),
+                        "running": int(cur.get("running", 0)),
+                        "waiting": int(cur.get("waiting", 0)),
+                        # Only meaningful when drafts were actually produced in
+                        # this window; a lifetime ratio would mask a live drop.
+                        "accept": round(d_acc / d_draft, 3) if d_draft > 0 else None,
+                        "draft_tok_s": round(d_draft / dt, 1) if d_draft > 0 else None,
+                        "ttft_s": win_mean("ttft_s", "ttft_n"),
+                        "inter_token_s": win_mean("itl_s", "itl_n"),
+                        "tpot_s": win_mean("tpot_s", "tpot_n"),
+                        "kv_cache_used_pct": (round(cur["kv_usage"] * 100, 2)
+                                              if "kv_usage" in cur else None),
+                        "prefix_hit_rate": round(d_h / d_q, 3) if d_q > 0 else None,
+                        "finished": d_fin or None,
+                        "waiting_why": cur.get("waiting_why") or None,
+                        # Prefill that vLLM has not credited yet. Chunked
+                        # prefill emits no IterationStats until its last chunk,
+                        # so decode and prefill both read 0 through a prompt
+                        # that may take minutes. KV comes from the block
+                        # manager instead, so it keeps moving -- and it is the
+                        # only thing on this page that shows the engine is
+                        # busy. Without this the page says "idle" while a 200K
+                        # prompt is 90 seconds into its prefill.
+                        "uncounted_prefill": bool(
+                            d_kv > 0
+                            and not (cur["gen"] - prev["gen"])
+                            and not (cur.get("prompt", 0) - prev.get("prompt", 0))
+                        ),
+                    })
+                    del _tp[:-_TP_KEEP]
+        if cur:
+            prev = cur
+        time.sleep(THROUGHPUT_INTERVAL_S)
+
+
+def throughput_rows():
+    with _tp_lock:
+        return list(_tp)
+
+
+# --- host load ------------------------------------------------------------
+# Sampled on BOTH roles, unlike throughput. Tonight a boot sat on "loading"
+# with every cache flat while three compilers ran at 100% -- the page looked
+# idle when the box was saturated. CPU busy needs two /proc/stat samples, so it
+# lives on a timer rather than being computed per request.
+HOST_SAMPLE_S = float(_env("HOST_SAMPLE_S", "15"))
+_host_lock = threading.Lock()
+_host: dict = {}
+
+
+def _cpu_times():
+    try:
+        with open("/proc/stat") as f:
+            parts = f.readline().split()
+        vals = [float(x) for x in parts[1:8]]
+        idle = vals[3] + (vals[4] if len(vals) > 4 else 0.0)
+        return sum(vals), idle
+    except (OSError, ValueError, IndexError):
+        return None
+
+
+def _gpu_sample():
+    try:
+        r = subprocess.run(
+            ["nvidia-smi",
+             "--query-gpu=utilization.gpu,temperature.gpu,power.draw,clocks.sm,memory.used,memory.total",
+             "--format=csv,noheader,nounits"],
+            capture_output=True, text=True, timeout=6)
+        line = (r.stdout or "").strip().splitlines()
+        if not line:
+            return {}
+        f = [x.strip() for x in line[0].split(",")]
+        def num(i):
+            try:
+                return float(f[i])
+            except (ValueError, IndexError):
+                return None
+        return {"gpu_util": num(0), "gpu_temp": num(1), "gpu_power": num(2),
+                "gpu_clock": num(3), "gpu_mem_used": num(4), "gpu_mem_total": num(5)}
+    except Exception:
+        return {}
+
+
+def _disk_sectors() -> dict:
+    """{device: (sectors_read, sectors_written)} for whole devices only.
+
+    Partitions are skipped so a device is not double-counted against itself.
+    """
+    out = {}
+    try:
+        with open("/proc/diskstats") as f:
+            for line in f:
+                p = line.split()
+                if len(p) < 10:
+                    continue
+                name = p[2]
+                if not re.fullmatch(r"(nvme\d+n\d+|sd[a-z]+|vd[a-z]+|mmcblk\d+)", name):
+                    continue
+                try:
+                    out[name] = (int(p[5]), int(p[9]))
+                except ValueError:
+                    continue
+    except OSError:
+        pass
+    return out
+
+
+def host_loop():
+    prev = _cpu_times()
+    prev_disk = _disk_sectors()
+    while True:
+        time.sleep(HOST_SAMPLE_S)
+        cur = _cpu_times()
+        busy = None
+        if prev and cur:
+            dt, di = cur[0] - prev[0], cur[1] - prev[1]
+            if dt > 0:
+                busy = max(0.0, min(100.0, 100.0 * (dt - di) / dt))
+        prev = cur or prev
+        sample = {"cpu_busy": None if busy is None else round(busy, 1),
+                  "ncpu": os.cpu_count()}
+        try:
+            with open("/proc/uptime") as f:
+                sample["host_uptime_s"] = int(float(f.read().split()[0]))
+        except (OSError, ValueError, IndexError):
+            pass
+        try:
+            l1, l5, l15 = os.getloadavg()
+            sample.update({"load1": round(l1, 2), "load5": round(l5, 2),
+                           "load15": round(l15, 2)})
+        except OSError:
+            pass
+        space = {}
+        for label, path in (("root", "/"), ("cache", CACHE_ROOT_PATH),
+                            ("logs", MCP_LOG_DIR)):
+            try:
+                st = os.statvfs(path)
+            except OSError:
+                continue
+            total = st.f_blocks * st.f_frsize
+            free = st.f_bavail * st.f_frsize
+            if total:
+                space[label] = {
+                    "path": path,
+                    "total_gb": round(total / 2**30, 1),
+                    "free_gb": round(free / 2**30, 1),
+                    "used_pct": round(100 * (1 - free / total), 1),
+                }
+        if space:
+            sample["disk_space"] = space
+
+        cur_disk = _disk_sectors()
+        disks = {}
+        for dev, (r, w) in cur_disk.items():
+            pr = prev_disk.get(dev)
+            if pr:
+                # 512-byte sectors -> MB/s over the sample interval.
+                disks[dev] = {
+                    "read_mb_s": round((r - pr[0]) * 512 / 1048576 / HOST_SAMPLE_S, 1),
+                    "write_mb_s": round((w - pr[1]) * 512 / 1048576 / HOST_SAMPLE_S, 1),
+                }
+        prev_disk = cur_disk
+        if disks:
+            sample["disk"] = disks
+        sample.update(_gpu_sample())
+        with _host_lock:
+            _host.clear()
+            _host.update(sample)
+
+
+def host_stats() -> dict:
+    with _host_lock:
+        return dict(_host)
+
+
+def local_status(caches: bool = True):
+    cur, history = read_stage()
+    if ROLE != "worker" and cur in ("starting", "loading") and vllm_up():
+        if SERVING_WHEN_READY:
+            cur = "serving"
+        elif "self-test" in dict(STAGES):
+            # API is bound but the entrypoint has not declared the self-test
+            # passed. Never claim serving on its behalf.
+            cur = "self-test"
+    healthy = (cur == "worker-ready") if ROLE == "worker" else (cur == "serving")
+    failed = cur in FAILED_STAGES
+    return {
+        "stage": cur,
+        "role": ROLE,
+        "healthy": healthy,
+        "failed": failed,
+        "serving": ROLE != "worker" and cur == "serving",
+        "elapsed_s": int(time.time() - START),
+        "history": history,
+        # Walks every cache directory, so a health probe does not ask for it.
+        "caches": cache_rows() if caches else [],
+        "throughput": throughput_rows()[-1] if throughput_rows() else None,
+        "throughput_history": throughput_rows(),
+        "load": host_stats(),
+        "host": HOSTNAME,
+    }
+
+
+# --- peer polling -----------------------------------------------------------
+# Polled on a timer into a snapshot rather than fetched while rendering, so a
+# hung or unplugged peer costs a stale row instead of a page that never loads.
+# Peers are asked for /healthz, which is deliberately LOCAL-ONLY -- if it
+# included peer data the nodes would poll each other forever. It is /status.json
+# without the cache walk.
+_peers_lock = threading.Lock()
+_peers: dict[str, dict] = {p: {"addr": p, "reachable": None} for p in PEERS}
+
+
+def poll_peer(addr: str) -> dict:
+    try:
+        req = urllib.request.Request(f"http://{addr}/healthz")
+        with urllib.request.urlopen(req, timeout=PEER_TIMEOUT_S) as r:
+            data = json.loads(r.read().decode())
+        # /healthz answers 503 until healthy; that is a valid answer, not
+        # an error, and urllib only raises on it via HTTPError (handled below).
+        data.update({"addr": addr, "reachable": True, "seen": time.time()})
+        return data
+    except urllib.error.HTTPError as e:
+        try:
+            data = json.loads(e.read().decode())
+            data.update({"addr": addr, "reachable": True, "seen": time.time()})
+            return data
+        except Exception:
+            return {"addr": addr, "reachable": False,
+                    "error": f"HTTP {e.code}", "seen": time.time()}
+    except Exception as e:
+        return {"addr": addr, "reachable": False,
+                "error": type(e).__name__, "seen": time.time()}
+
+
+def peer_loop():
+    while True:
+        for addr in PEERS:
+            got = poll_peer(addr)
+            with _peers_lock:
+                _peers[addr] = got
+        time.sleep(PEER_POLL_S)
+
+
+def cluster_rows():
+    """One row per node, this one included, with self marked."""
+    with _peers_lock:
+        snap = dict(_peers)
+    rows = []
+    seen_self = False
+    for addr in PEERS:
+        p = snap.get(addr, {"addr": addr, "reachable": None})
+        is_self = p.get("host") == HOSTNAME
+        seen_self = seen_self or is_self
+        rows.append({
+            "addr": addr,
+            "self": is_self,
+            # Two different names, for two different jobs:
+            #   addr  -- how THIS node polls the peer, often an address on the
+            #            fabric between the ranks. Not resolvable from anywhere
+            #            else, so it must never be a link target.
+            #   name  -- how a READER reaches the peer: the box's real hostname,
+            #            which is what resolves on the LAN.
+            # Falls back to the poll address only when the peer is down and has
+            # therefore not told us its hostname.
+            "name": p.get("host") or addr.rsplit(":", 1)[0],
+            "port": addr.rsplit(":", 1)[-1] if ":" in addr else str(PORT),
+            "role": p.get("role", "?"),
+            "stage": p.get("stage", "unreachable" if p.get("reachable") is False else "..."),
+            "healthy": bool(p.get("healthy")),
+            "reachable": p.get("reachable"),
+            "elapsed_s": p.get("elapsed_s"),
+            "error": p.get("error", ""),
+            "throughput": p.get("throughput"),
+            "throughput_history": p.get("throughput_history") or [],
+            "load": p.get("load") or {},
+        })
+    # If this node is not in PEERS (misconfigured, or a name that resolves
+    # elsewhere) show it anyway rather than rendering a cluster it is absent from.
+    if not seen_self:
+        me = local_status()
+        rows.insert(0, {
+            "addr": f"{HOSTNAME}:{PORT}", "self": True,
+            "name": HOSTNAME, "port": str(PORT),
+            "role": me["role"], "stage": me["stage"], "healthy": me["healthy"],
+            "reachable": True, "elapsed_s": me["elapsed_s"], "error": "",
+            "load": host_stats(),
+        })
+    return rows
+
+
+# --- MCP ------------------------------------------------------------------
+# OpenWebUI 0.11 speaks MCP natively over streamable HTTP, so the model can ask
+# this node about itself instead of a human reading the page. Deliberately
+# constrained: every tool is READ-ONLY, takes no arguments, and runs a fixed
+# argv (never a shell). There is no auth on this port, so the whitelist IS the
+# security boundary -- keep it boring. Notably absent: anything needing the
+# Docker socket, which would be root-equivalent on the host.
+MCP_PROTOCOL_VERSION = "2025-06-18"
+
+
+def _run(argv: list[str], timeout: float = 8.0) -> str:
+    try:
+        r = subprocess.run(argv, capture_output=True, text=True, timeout=timeout)
+        out = (r.stdout or "") + (r.stderr or "")
+        return out.strip() or "(no output)"
+    except FileNotFoundError:
+        raise RuntimeError(f"{argv[0]} is not available in this container")
+    except subprocess.TimeoutExpired:
+        raise RuntimeError(f"{argv[0]} timed out after {timeout}s")
+
+
+def _read(path: str, default: str = "") -> str:
+    try:
+        with open(path) as f:
+            return f.read().strip()
+    except OSError:
+        return default
+
+
+def _run_or_note(argv: list[str]) -> str:
+    """_run for a report that should name a missing binary and carry on,
+    rather than lose every other line to it."""
+    try:
+        return _run(argv)
+    except RuntimeError as e:
+        return f"({e})"
+
+
+def _t_node_status():
+    me = local_status()
+    return json.dumps({k: me[k] for k in
+                       ("host", "role", "stage", "healthy", "serving", "elapsed_s")},
+                      indent=2)
+
+
+def _t_cluster_status():
+    rows = cluster_rows()
+    return json.dumps({"cluster_healthy": all(r["healthy"] for r in rows),
+                       "nodes": rows}, indent=2, default=str)
+
+
+def _t_throughput():
+    hist = throughput_rows()
+    if not hist:
+        for n in cluster_rows():
+            if not n["self"] and n.get("throughput"):
+                return json.dumps({
+                    "source": n["name"],
+                    "interval_s": THROUGHPUT_INTERVAL_S,
+                    "samples_newest_last": n.get("throughput_history")
+                                           or [n["throughput"]],
+                }, indent=2)
+        return "No throughput samples yet (the API may still be loading)."
+    return json.dumps({"source": HOSTNAME, "interval_s": THROUGHPUT_INTERVAL_S,
+                       "samples_newest_last": hist}, indent=2)
+
+
+_LATENCY_HISTS = (
+    "time_to_first_token_seconds", "inter_token_latency_seconds",
+    "request_queue_time_seconds", "request_prefill_time_seconds",
+    "request_decode_time_seconds",
+)
+
+
+def _hist_buckets() -> dict[str, dict[float, float]]:
+    """{histogram: {le: cumulative count}} from the engine's metrics endpoint."""
+    try:
+        with urllib.request.urlopen(
+            f"http://127.0.0.1:{VLLM_PORT}/metrics", timeout=5) as r:
+            body = r.read().decode()
+    except Exception:
+        return {}
+    out: dict[str, dict[float, float]] = {}
+    for line in body.splitlines():
+        if not line.startswith("vllm:") or "_bucket" not in line:
+            continue
+        name = line.split("{", 1)[0].split("_bucket")[0].split(":", 1)[-1]
+        if name not in _LATENCY_HISTS:
+            continue
+        try:
+            le = float(line.split('le="', 1)[1].split('"', 1)[0])
+            out.setdefault(name, {})[le] = float(line.rsplit(" ", 1)[-1])
+        except (IndexError, ValueError):
+            pass
+    return out
+
+
+def _t_latency_percentiles(window_s: int = 20, **_):
+    """Percentiles over a window, rather than over the process lifetime."""
+    window_s = max(5, min(int(window_s or 20), 120))
+    before = _hist_buckets()
+    if not before:
+        return "The engine is not serving metrics yet."
+    time.sleep(window_s)
+    after = _hist_buckets()
+
+    out: dict[str, object] = {
+        "window_s": window_s,
+        "now": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+        "note": "Empty means nothing of that kind happened in the window. "
+                "inter_token_latency needs a request decoding, and one token "
+                "arrives per scheduler step, so it doubles as step time.",
+    }
+    for name in _LATENCY_HISTS:
+        a, b = before.get(name, {}), after.get(name, {})
+        if not b:
+            continue
+        delta = {le: b[le] - a.get(le, 0.0) for le in b}
+        les = sorted(delta)
+        total = delta[les[-1]] if les else 0.0          # +Inf holds the count
+        if total <= 0:
+            out[name] = {"observations": 0}
+            continue
+        qs = {}
+        for q in (0.5, 0.9, 0.99):
+            target = q * total
+            qs[f"p{int(q * 100)}"] = next(
+                (le for le in les if delta[le] >= target), les[-1])
+        out[name] = {"observations": int(total), "seconds_at_most": qs}
+    return json.dumps(out, indent=2)
+
+
+def _t_metrics(grep: str = "", **_):
+    """The engine's Prometheus text, filtered by a regex.
+
+    Raw because which counter matters is decided by the question being asked.
+    The summary tools carry a fixed set and this carries everything else.
+    """
+    try:
+        with urllib.request.urlopen(
+                f"http://127.0.0.1:{VLLM_PORT}/metrics", timeout=5) as r:
+            body = r.read().decode()
+    except Exception as e:
+        return f"({type(e).__name__}: {e})"
+    if grep:
+        try:
+            rx = re.compile(grep)
+        except re.error as e:
+            return f"(bad pattern: {e})"
+        body = "\n".join(l for l in body.splitlines()
+                          if not l.startswith("#") and rx.search(l))
+    b = body.encode()
+    if len(b) > MCP_LOG_MAX_BYTES:
+        body = (b[:MCP_LOG_MAX_BYTES].decode("utf-8", "ignore")
+                + f"\n\n[truncated at {MCP_LOG_MAX_BYTES} bytes of {len(b)}]")
+    return body or "(no lines matched)"
+
+
+def _t_serve_args(**_):
+    """The command line the engine actually runs, from /proc/1/cmdline."""
+    try:
+        with open("/proc/1/cmdline", "rb") as f:
+            return f.read().replace(b"\0", b" ").decode(errors="replace")
+    except OSError as e:
+        return f"({type(e).__name__}: {e})"
+
+
+def _t_cache_sizes():
+    return json.dumps(cache_rows(), indent=2)
+
+
+# The log mount, reported in the host sample's free-space figures. Reading logs
+# belongs to the host agent.
+MCP_LOG_DIR = _env("MCP_LOG_DIR", "/logs")
+# Cap on one tool answer (the metrics tool). 16 KB, not 64: this lands directly
+# in a model's context, and more crowds out the reasoning it is meant to inform.
+MCP_LOG_MAX_BYTES = int(_env("MCP_LOG_MAX_BYTES", "16384"))
+
+
+def _now_header() -> str:
+    lt = time.strftime("%Y-%m-%d %H:%M:%S %Z", time.localtime())
+    ut = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
+    return f"[now: {lt} / {ut} on {HOSTNAME}]"
+
+
+def _t_versions():
+    """Driver, CUDA, firmware and Python stack versions."""
+    out = {}
+    # /proc/driver/nvidia/version is the kernel module's own claim, which is
+    # what actually matters after a driver update -- nvidia-smi can report a
+    # different userspace version if the two drift.
+    out["nvidia_kernel_module"] = _read("/proc/driver/nvidia/version", "(absent)")
+    out["nvidia_smi"] = _run_or_note(
+        ["nvidia-smi",
+         "--query-gpu=name,driver_version,vbios_version,compute_cap,serial",
+         "--format=csv"])
+    nvcc = _run_or_note(["nvcc", "--version"])
+    out["cuda_toolkit"] = nvcc.splitlines()[-1] if nvcc else ""
+    for label, path in (("bios_version", "/sys/class/dmi/id/bios_version"),
+                        ("bios_date", "/sys/class/dmi/id/bios_date"),
+                        ("board_name", "/sys/class/dmi/id/board_name"),
+                        ("board_vendor", "/sys/class/dmi/id/board_vendor"),
+                        ("product_name", "/sys/class/dmi/id/product_name")):
+        v = _read(path)
+        if v:
+            out[label] = v
+    out["kernel"] = _read("/proc/sys/kernel/osrelease")
+    # Read from package metadata, never by importing: importing torch and vLLM
+    # here would keep over a gigabyte resident in this process for the rest of
+    # the container's life. `ray` is mentat's shim, distributed as mentatd.
+    for label, dists in (("torch", ["torch"]), ("vllm", ["vllm"]),
+                         ("flashinfer", ["flashinfer-python", "flashinfer"]),
+                         ("tilelang", ["tilelang"]), ("triton", ["triton"]),
+                         ("ray", ["mentatd", "ray"])):
+        for d in dists:
+            try:
+                out[f"py_{label}"] = f"{d} {importlib.metadata.version(d)}"
+                break
+            except importlib.metadata.PackageNotFoundError:
+                continue
+        else:
+            out[f"py_{label}"] = "(not installed)"
+    return json.dumps(out, indent=2)
+
+
+RAY_DASHBOARD = _env("RAY_DASHBOARD", "http://127.0.0.1:8265")
+
+
+def _t_ray_status():
+    """Ray cluster resources plus per-actor state.
+
+    The dashboard is read over LOCALHOST and never exposed: it carries the Job
+    Submission API, which runs arbitrary code, so binding it to 0.0.0.0 would
+    put remote code execution on the LAN. This tool returns a curated read-only
+    view instead.
+    """
+    out = {"resources": _run(["ray", "status"], timeout=15)}
+    try:
+        req = urllib.request.Request(f"{RAY_DASHBOARD}/logical/actors")
+        with urllib.request.urlopen(req, timeout=8) as r:
+            data = json.loads(r.read().decode())
+        acts = (data.get("data") or {}).get("actors") or {}
+        out["actors"] = [
+            {"class": a.get("actorClass"), "state": a.get("state"),
+             "ip": (a.get("address") or {}).get("ipAddress"), "pid": a.get("pid")}
+            for a in acts.values()
+        ]
+        # A TP rank whose actor is missing or not ALIVE is the signature of a
+        # cluster that loads on one node and then waits forever on the other.
+        alive = sum(1 for a in out["actors"] if a["state"] == "ALIVE")
+        out["actors_alive"] = f"{alive} of {len(out['actors'])}"
+    except Exception as e:
+        out["actors"] = f"dashboard unreachable: {type(e).__name__}: {e}"
+    return json.dumps(out, indent=2)
+
+# --- filesystem introspection ----------------------------------------------
+# Three read-only tools so the model can answer questions about the deployment
+# from the files themselves instead of guessing. Tonight's parser-name hunt and
+# "which cache was written during the stall" both took a human many round trips
+# and are one tool call each.
+#
+# Every one of these takes ARGUMENTS, so the containment is the security
+# boundary, not the tool whitelist. Rules, all enforced below:
+#   - paths must resolve inside SEARCH_ROOTS (realpath, so symlinks cannot escape)
+#   - arguments are passed as argv, never through a shell
+#   - output is capped so one call cannot flood a context window
+def _default_search_roots() -> str:
+    """vLLM's source wherever this image has it, then the caches and logs.
+
+    A source build keeps it in /src/vllm; a stock image has only the installed
+    package. find_spec locates the package without importing it.
+    """
+    roots = []
+    if os.path.isdir("/src/vllm"):
+        roots.append("/src/vllm")
+    try:
+        spec = importlib.util.find_spec("vllm")
+        if spec and spec.origin:
+            roots.append(os.path.dirname(spec.origin))
+    except (ImportError, ValueError):
+        pass
+    return ",".join(roots + ["/root/.cache", "/logs", "/cache"])
+
+
+SEARCH_ROOTS = [r.strip() for r in _env(
+    "SEARCH_ROOTS", _default_search_roots()).split(",") if r.strip()]
+SEARCH_MAX_BYTES = int(_env("SEARCH_MAX_BYTES", "16384"))
+
+
+def _resolve_under_roots(path: str) -> str:
+    """Resolve a caller path inside one of SEARCH_ROOTS, or raise."""
+    if not path:
+        raise RuntimeError(f"path is required; allowed roots: {', '.join(SEARCH_ROOTS)}")
+    target = os.path.realpath(path)
+    for root in SEARCH_ROOTS:
+        r = os.path.realpath(root)
+        if target == r or target.startswith(r + os.sep):
+            if os.path.exists(target):
+                return target
+            raise RuntimeError(f"no such path: {path}")
+    raise RuntimeError(f"path {path!r} is outside the allowed roots: "
+                       f"{', '.join(SEARCH_ROOTS)}")
+
+
+def _cap(text: str) -> str:
+    if len(text) > SEARCH_MAX_BYTES:
+        return (text[:SEARCH_MAX_BYTES]
+                + f"\n... (truncated at {SEARCH_MAX_BYTES} bytes)")
+    return text
+
+
+def _t_find_files(path: str = "", name: str = "", newer_than_minutes: int = 0,
+                  min_size_mb: float = 0, limit: int = 200, **_):
+    root = _resolve_under_roots(path)
+    argv = ["find", root]
+    if name:
+        if "/" in name or len(name) > 120:
+            raise RuntimeError("name must be a bare glob, e.g. '*.safetensors'")
+        argv += ["-name", name]
+    try:
+        mins = int(newer_than_minutes)
+    except (TypeError, ValueError):
+        mins = 0
+    if mins > 0:
+        argv += ["-mmin", f"-{mins}"]
+    try:
+        mb = float(min_size_mb)
+    except (TypeError, ValueError):
+        mb = 0
+    if mb > 0:
+        argv += ["-size", f"+{int(mb * 1024)}k"]
+    argv += ["-type", "f", "-printf", "%TY-%Tm-%Td %TH:%TM  %10s  %p\\n"]
+    out = _run(argv, timeout=30)
+    lines = out.splitlines()
+    try:
+        n = max(1, min(int(limit), 2000))
+    except (TypeError, ValueError):
+        n = 200
+    head = "\n".join(lines[:n])
+    extra = f"\n... ({len(lines) - n} more matches not shown)" if len(lines) > n else ""
+    return _now_header() + f"  {len(lines)} matches\n" + _cap(head + extra)
+
+
+def _t_search_files(pattern: str = "", path: str = "", glob: str = "",
+                    limit: int = 100, **_):
+    root = _resolve_under_roots(path)
+    if not pattern or len(pattern) > 200:
+        raise RuntimeError("pattern is required and must be 200 characters or fewer")
+    try:
+        n = max(1, min(int(limit), 1000))
+    except (TypeError, ValueError):
+        n = 100
+    argv = ["rg", "--line-number", "--no-heading", "--color", "never",
+            "--max-count", "20", "--max-filesize", "8M", "-m", str(n)]
+    if glob:
+        if len(glob) > 120:
+            raise RuntimeError("glob must be 120 characters or fewer")
+        argv += ["--glob", glob]
+    argv += ["--", pattern, root]
+    out = _run(argv, timeout=45)
+    return _now_header() + f"  pattern={pattern!r} under {root}\n" + _cap(out)
+
+
+NO_ARGS = {"type": "object", "properties": {}, "required": []}
+
+MCP_TOOLS = [
+    ("node_status", "Stage, role, health and uptime of THIS node's model "
+     "container. The worker node holds a TP rank and never serves an API, so "
+     "'worker-ready' is its healthy state, not an error.", _t_node_status, NO_ARGS),
+    ("cluster_status", "Stage and health of every node in the serving cluster, "
+     "polled from each rank's status server. Use this to answer 'is the "
+     "cluster up?' -- one node serving while another is not is the common "
+     "failure. A single-node engine shows only itself.",
+     _t_cluster_status, NO_ARGS),
+    ("metrics", "The engine's raw Prometheus text, filtered by a regex. "
+     "Everything the engine reports is here, including counters the summary "
+     "tools do not carry.", _t_metrics, {
+         "type": "object",
+         "properties": {"grep": {
+             "type": "string",
+             "description": "Regex; only matching lines are returned.",
+         }},
+         "required": [],
+     }),
+    ("serve_args", "The exact command line the engine runs. Says which flags "
+     "are live rather than what a config file asks for.", _t_serve_args, NO_ARGS),
+    ("throughput", "Recent decode and prefill tokens/sec, running and queued "
+     "request counts, and speculative-decoding acceptance rate. Rates are "
+     "deltas between samples, not lifetime averages. Acceptance is blank "
+     "without speculative decoding; far below the model's usual rate it "
+     "suggests mismatched or corrupted weights.", _t_throughput, NO_ARGS),
+    ("latency_percentiles", "TTFT, inter-token latency, queue time, prefill and "
+     "decode durations as p50/p90/p99 over a window you choose, taken as the "
+     "difference between two scrapes so the numbers describe now rather than "
+     "the process lifetime. Queue time is what request starvation looks like. "
+     "Inter-token latency is one token per scheduler step, so it also reads as "
+     "step time. Blocks for the window.", _t_latency_percentiles, {
+         "type": "object",
+         "properties": {"window_s": {
+             "type": "integer",
+             "description": "Seconds to sample, 5 to 120. Default 20.",
+         }},
+         "required": [],
+     }),
+    ("cache_sizes", "Sizes of the JIT/compile caches (FlashInfer, torch.compile, "
+     "Triton, TileLang) and which is being written to right now. Useful for "
+     "seeing what a slow startup is actually doing.", _t_cache_sizes, NO_ARGS),
+    # (conditionally appended below when Ray is present)
+]
+
+RAY_TOOL = [
+    ("ray_status", "Ray cluster resources and the state of every actor, "
+     "including one RayWorkerProc per tensor-parallel rank. If a rank's actor "
+     "is missing or not ALIVE, that is why the head loaded and then waited "
+     "forever. Read from the dashboard over localhost; the dashboard is not "
+     "exposed because it carries a job API that executes arbitrary code.",
+     _t_ray_status, NO_ARGS),
+]
+
+if _env("ENABLE_RAY_TOOL", "auto") == "1" or (
+    _env("ENABLE_RAY_TOOL", "auto") == "auto"
+    and os.path.isdir("/tmp/ray")
+):
+    MCP_TOOLS += RAY_TOOL
+
+MCP_TOOLS += [
+    ("versions", "Firmware and software versions: NVIDIA kernel module, driver, "
+     "VBIOS, compute capability, CUDA toolkit, system BIOS/board, kernel, and "
+     "the Python stack (torch, vllm, flashinfer, tilelang, triton, ray). The "
+     "kernel-module version and nvidia-smi's driver version can disagree after "
+     "an update, which is worth noticing.", _t_versions, NO_ARGS),
+    ("find_files", "Find files under an allowed root, filtered by name glob, "
+     "age and size. Answers questions like 'which cache was written during the "
+     "stall' or 'how big are the shards' directly, instead of inferring them. "
+     f"Allowed roots: {', '.join(SEARCH_ROOTS)}.", _t_find_files, {
+         "type": "object",
+         "properties": {
+             "path": {"type": "string",
+                      "description": f"Directory to search. Must be inside "
+                                     f"{', '.join(SEARCH_ROOTS)}."},
+             "name": {"type": "string",
+                      "description": "Bare filename glob, e.g. '*.safetensors'. "
+                                     "No slashes."},
+             "newer_than_minutes": {"type": "integer",
+                                    "description": "Only files modified within "
+                                                   "this many minutes."},
+             "min_size_mb": {"type": "number",
+                             "description": "Only files at least this large."},
+             "limit": {"type": "integer", "description": "Max results (default 200)."},
+         },
+         "required": ["path"],
+     }),
+    ("search_files", "Search file CONTENTS with ripgrep under an allowed root. "
+     "Use this to read how the running vLLM actually behaves -- which env var a "
+     "cache honours, what a parser registers as, what a flag defaults to -- "
+     f"rather than guessing from the name. Allowed roots: {', '.join(SEARCH_ROOTS)}.",
+     _t_search_files, {
+         "type": "object",
+         "properties": {
+             "pattern": {"type": "string",
+                         "description": "Regular expression. Max 200 characters."},
+             "path": {"type": "string",
+                      "description": f"Directory to search, inside "
+                                     f"{', '.join(SEARCH_ROOTS)}."},
+             "glob": {"type": "string",
+                      "description": "Restrict to matching files, e.g. '*.py'."},
+             "limit": {"type": "integer", "description": "Max matches (default 100)."},
+         },
+         "required": ["pattern", "path"],
+     }),
+]
+
+# search_files runs ripgrep, which a stock vLLM image does not ship. Listing it
+# anyway would advertise a tool that can only fail.
+if shutil.which("rg") is None:
+    MCP_TOOLS = [t for t in MCP_TOOLS if t[0] != "search_files"]
+
+MCP_BY_NAME = {n: (d, f, sc) for n, d, f, sc in MCP_TOOLS}
+
+
+# --- per-node dispatch ------------------------------------------------------
+# Tools whose answer is a property of ONE machine. Aggregate tools
+# (cluster_status, host_load, throughput) already span the cluster and take no
+# node argument. Omitting `node` answers locally, which means whichever node the
+# MCP client is pointed at -- so the default needs no special case.
+NODE_LOCAL = {
+    "node_status", "cache_sizes", "versions", "find_files", "search_files",
+    # Only the head runs a Ray dashboard, but keep it node-addressable so a
+    # client pointed at the worker can still ask the head.
+    "ray_status",
+}
+NODE_ARG = {
+    "node": {
+        "type": "string",
+        "description": "Which node to answer for, e.g. a hostname from "
+                       "cluster_status. Omit for the node serving this request.",
+    }
+}
+
+
+def _peer_addr_for(node: str) -> str | None:
+    """Match a caller-supplied name against hostname or configured address."""
+    want = node.strip().lower()
+    for n in cluster_rows():
+        if n["self"]:
+            continue
+        if want in (str(n.get("host", "")).lower(),
+                    str(n.get("name", "")).lower(),
+                    str(n.get("addr", "")).lower(),
+                    str(n.get("addr", "")).split(":")[0].lower()):
+            return n["addr"]
+    return None
+
+
+def _is_self(node: str) -> bool:
+    want = node.strip().lower()
+    return want in (HOSTNAME.lower(), "self", "local", "this")
+
+
+def proxy_tool_call(node: str, name: str, args: dict) -> dict:
+    """Forward one tools/call to a peer and return its result verbatim.
+
+    The forwarded request carries NO node argument, so the peer answers locally
+    and cannot bounce it onward -- that is what stops two nodes ping-ponging a
+    request between them.
+    """
+    addr = _peer_addr_for(node)
+    if addr is None:
+        known = [HOSTNAME] + [n["host"] or n["name"] for n in cluster_rows()
+                              if not n["self"]]
+        return {"content": [{"type": "text",
+                             "text": f"unknown node {node!r}; known nodes: "
+                                     f"{', '.join(x for x in known if x)}"}],
+                "isError": True}
+    body = json.dumps({"jsonrpc": "2.0", "id": 1, "method": "tools/call",
+                       "params": {"name": name, "arguments": args}}).encode()
+    req = urllib.request.Request(f"http://{addr}/mcp", data=body,
+                                 headers={"Content-Type": "application/json"})
+    try:
+        with urllib.request.urlopen(req, timeout=30) as r:
+            resp = json.loads(r.read().decode())
+    except Exception as e:
+        return {"content": [{"type": "text",
+                             "text": f"could not reach {node} at {addr}: "
+                                     f"{type(e).__name__}: {e}"}],
+                "isError": True}
+    if "result" in resp:
+        return resp["result"]
+    return {"content": [{"type": "text", "text": json.dumps(resp)}],
+            "isError": True}
+
+def mcp_handle(req: dict) -> dict | None:
+    """One JSON-RPC request -> one response, or None for a notification."""
+    rid = req.get("id")
+    method = req.get("method", "")
+
+    def ok(result):
+        return {"jsonrpc": "2.0", "id": rid, "result": result}
+
+    def err(code, msg):
+        return {"jsonrpc": "2.0", "id": rid, "error": {"code": code, "message": msg}}
+
+    if method == "initialize":
+        want = (req.get("params") or {}).get("protocolVersion")
+        return ok({
+            "protocolVersion": want if isinstance(want, str) else MCP_PROTOCOL_VERSION,
+            "capabilities": {"tools": {"listChanged": False}},
+            "serverInfo": {"name": f"{SERVICE}-status ({HOSTNAME})", "version": "1"},
+        })
+    if method.startswith("notifications/"):
+        return None                      # nothing to acknowledge
+    if method == "ping":
+        return ok({})
+    if method == "tools/list":
+        tools = []
+        for n, d, _, sc in MCP_TOOLS:
+            if n in NODE_LOCAL:
+                sc = dict(sc)
+                sc["properties"] = {**sc.get("properties", {}), **NODE_ARG}
+            tools.append({"name": n, "description": d, "inputSchema": sc})
+        return ok({"tools": tools})
+    if method == "tools/call":
+        params = req.get("params") or {}
+        name = params.get("name")
+        entry = MCP_BY_NAME.get(name)
+        if entry is None:
+            # A tool error is reported in the RESULT, not as a protocol error,
+            # so the model can see it and recover rather than the call failing.
+            return ok({"content": [{"type": "text",
+                                    "text": f"unknown tool {name!r}; available: "
+                                            f"{', '.join(MCP_BY_NAME)}"}],
+                       "isError": True})
+        args = params.get("arguments") or {}
+        if not isinstance(args, dict):
+            args = {}
+        node = args.pop("node", None)
+        if node and name in NODE_LOCAL and not _is_self(str(node)):
+            return ok(proxy_tool_call(str(node), name, args))
+        # Ignore arguments for tools that declare none, rather than raising a
+        # TypeError at a client that sent a harmless empty/extra field.
+        if not (entry[2].get("properties") or {}):
+            args = {}
+        try:
+            text = entry[1](**args)
+        except Exception as e:
+            return ok({"content": [{"type": "text",
+                                    "text": f"{type(e).__name__}: {e}"}],
+                       "isError": True})
+        return ok({"content": [{"type": "text", "text": text}], "isError": False})
+    return err(-32601, f"method not found: {method}")
+
+
+def memory_report():
+    """Host memory beside what the workers say they hold.
+
+    Every allocator counter is per process, so the workers publish theirs and
+    this only collects them. `unaccounted` is the interesting column: host
+    memory that is neither free nor reserved by torch is NCCL, cuBLAS or the
+    JIT caches, and on unified memory it competes with the same pool that
+    decides whether this node keeps forking.
+    """
+    meminfo = {}
+    try:
+        with open("/proc/meminfo") as fh:
+            for line in fh:
+                key, _, rest = line.partition(":")
+                meminfo[key] = int(rest.split()[0]) * 1024
+    except OSError:
+        pass
+
+    workers = []
+    stats_dir = _env("TORCH_MEM_STATS_DIR")
+    if stats_dir:
+        for name in sorted(glob.glob(os.path.join(stats_dir, "mem-stats-pid*.json"))):
+            try:
+                with open(name) as fh:
+                    workers.append(json.load(fh))
+            except (OSError, ValueError):
+                continue
+
+    total = meminfo.get("MemTotal", 0)
+    reserved = sum(w.get("reserved", 0) for w in workers)
+    out = {
+        "host": {
+            "total": total,
+            "free": meminfo.get("MemFree", 0),
+            "available": meminfo.get("MemAvailable", 0),
+            "cached": meminfo.get("Cached", 0),
+        },
+        "workers": workers,
+        "torch_reserved": reserved,
+        "torch_allocated": sum(w.get("allocated", 0) for w in workers),
+    }
+    if total and workers:
+        out["unaccounted"] = total - reserved - meminfo.get("MemAvailable", 0)
+    return out
+
+
+class Handler(http.server.BaseHTTPRequestHandler):
+    # Default is HTTP/1.0, which closes the connection after every response --
+    # an MCP client makes a request per tool call and would reconnect each time.
+    # Safe here because every response path sets an explicit Content-Length.
+    protocol_version = "HTTP/1.1"
+
+    def log_message(self, *a):
+        pass
+
+    def _json(self, obj, code):
+        body = json.dumps(obj).encode()
+        self.send_response(code)
+        self.send_header("Content-Type", "application/json")
+        self.send_header("Content-Length", str(len(body)))
+        self.end_headers()
+        self.wfile.write(body)
+
+    def do_POST(self):
+        if (self.path.rstrip("/") or "/") != "/mcp":
+            self.send_response(404); self.send_header("Content-Length", "0")
+            self.end_headers(); return
+        try:
+            n = int(self.headers.get("Content-Length") or 0)
+            payload = json.loads(self.rfile.read(n).decode() or "{}")
+        except Exception as e:
+            self._json({"jsonrpc": "2.0", "id": None,
+                        "error": {"code": -32700, "message": f"parse error: {e}"}}, 400)
+            return
+        # A client may batch requests in a list.
+        if isinstance(payload, list):
+            out = [r for r in (mcp_handle(m) for m in payload) if r is not None]
+            if not out:
+                self.send_response(202); self.send_header("Content-Length", "0")
+                self.end_headers(); return
+            self._json(out, 200)
+            return
+        resp = mcp_handle(payload)
+        if resp is None:                      # notification: no body
+            self.send_response(202); self.send_header("Content-Length", "0")
+            self.end_headers(); return
+        self._json(resp, 200)
+
+    def do_GET(self):
+        path = self.path.rstrip("/") or "/"
+
+        # The transport allows a GET for server-pushed events. Nothing here
+        # pushes, so decline rather than hold a connection open forever.
+        if path == "/mcp":
+            self.send_response(405)
+            self.send_header("Allow", "POST")
+            self.send_header("Content-Length", "0")
+            self.end_headers()
+            return
+
+        # LOCAL ONLY -- see the note on _peers above. Peers poll this endpoint.
+        if path in ("/healthz", "/status.json"):
+            me = local_status(caches=(path == "/status.json"))
+            self._json(me, 200 if me["healthy"] else 503)
+            return
+
+        if path == "/memory":
+            self._json(memory_report(), 200)
+            return
+
+        me = local_status()
+        if path == "/cluster.json":
+            rows = cluster_rows()
+            allup = all(r["healthy"] for r in rows) if rows else False
+            self._json({"cluster_healthy": allup, "nodes": rows, "local": me},
+                       200 if allup else 503)
+            return
+
+        cur, healthy, elapsed = me["stage"], me["healthy"], me["elapsed_s"]
+        failed = me.get("failed")
+        # A failure is not a step in the sequence, so anchor the walk on the last
+        # real stage reached; otherwise every step renders as done (ticked) and
+        # the page reads as success.
+        anchor = cur
+        if failed:
+            anchor = next((h for h in reversed(me["history"])
+                           if h in dict(stages())), "loading")
+        rows = []
+        reached = False
+        for key, label in stages():
+            if key == anchor:
+                mark, cls, reached = ("&#9679;", "bad", True) if failed else ("&#9679;", "now", True)
+            elif reached:
+                mark, cls = "&#9675;", "todo"
+            else:
+                mark, cls = "&#10003;", "done"
+            rows.append(f'<li class="{cls}"><span>{mark}</span> {label}</li>')
+        if failed:
+            rows.append('<li class="bad"><span>&#10007;</span> '
+                        'SELF-TEST FAILED &mdash; the model loaded but answered '
+                        'incorrectly. Check the container log.</li>')
+
+        # The worker scrapes nothing; show whichever node published numbers.
+        tp_hist = throughput_rows()
+        tp_now = tp_hist[-1] if tp_hist else None
+        tp_src = HOSTNAME
+        if tp_now is None:
+            for n in cluster_rows():
+                if not n["self"] and n.get("throughput"):
+                    tp_now, tp_src = n["throughput"], n["name"]
+                    break
+
+        def _tp_cells(r):
+            acc = "-" if r.get("accept") is None else f'{r["accept"]*100:.1f}%'
+            # A zero here is ambiguous, and the ambiguity cost days: vLLM does
+            # not credit prompt tokens until a chunked prefill emits its first
+            # token, so a 200K prompt reads 0/0 for minutes while the GPU is
+            # saturated. Say which zero this is.
+            pre = (f'{r["prefill"]:.0f}' if not r.get("uncounted_prefill")
+                   else '<span class=now>prefilling</span>')
+            return (f'<td class=n>{r["decode"]:.1f}</td>'
+                    f'<td class=n>{pre}</td>'
+                    f'<td class=n>{r["running"]}</td>'
+                    f'<td class=n>{r["waiting"]}</td>'
+                    f'<td class=n>{acc}</td>')
+
+        if tp_now:
+            hist = "".join(
+                f'<tr><td>{time.strftime("%H:%M:%S", time.localtime(r["t"]))}</td>'
+                f'{_tp_cells(r)}</tr>' for r in reversed(tp_hist[-8:]))
+            if not hist:
+                hist = f'<tr><td>now</td>{_tp_cells(tp_now)}</tr>'
+            tp_html = (
+                f'<h2>throughput &middot; {tp_src}</h2><table>'
+                f'<tr><th>time</th><th>decode tok/s</th><th>prefill tok/s</th>'
+                f'<th>running</th><th>queued</th><th>spec accept</th></tr>'
+                f'{hist}</table>'
+                f'<div class=foot style="margin-top:.5rem">'
+                f'sampled every {int(THROUGHPUT_INTERVAL_S)}s &middot; '
+                f'rates are deltas between samples, not lifetime averages'
+                f'<br>&ldquo;prefilling&rdquo; means KV is growing but vLLM has '
+                f'not credited the tokens yet &mdash; a long prompt reports '
+                f'nothing until its last chunk</div>')
+        else:
+            tp_html = ""
+
+        cache_html = []
+        for c in cache_rows():
+            if c["active"]:
+                cls, when = "now", "writing now"
+            elif c["idle_s"] is None:
+                cls, when = "", "-"
+            else:
+                cls, when = "", f'{c["idle_s"]}s ago'
+            dot = " &#9679;" if c["active"] else ""
+            cache_html.append(
+                f'<tr><td class="{cls}">{c["name"]}{dot}</td>'
+                f'<td class=n>{c["size"]}</td>'
+                f'<td class="n g">{c["grown"]}</td>'
+                f'<td class="n {cls}">{when}</td></tr>')
+
+        crows = []
+        for n in cluster_rows():
+            if n["reachable"] is False:
+                dot, cls, detail = "&#9679;", "bad", "unreachable"
+            elif n["healthy"]:
+                dot, cls, detail = "&#9679;", "ok", n["stage"]
+            else:
+                dot, cls, detail = "&#9679;", "wait", n["stage"]
+            el = "" if n["elapsed_s"] is None else f"{n['elapsed_s']//60}m{n['elapsed_s']%60:02d}s"
+            # Link the reachable hostname, NOT n["addr"] -- see cluster_rows.
+            link = (n["name"] if n["self"]
+                    else f'<a href="http://{n["name"]}:{n["port"]}/">{n["name"]}</a>')
+            tag = ' <span class="me">this node</span>' if n["self"] else ""
+            ld = n.get("load") or {}
+            def _pct(v):
+                return "-" if v is None else f"{float(v):.0f}%"
+            cpu = _pct(ld.get("cpu_busy"))
+            gpu = _pct(ld.get("gpu_util"))
+            if ld.get("gpu_temp") is not None:
+                gpu += f" {ld['gpu_temp']:.0f}&deg;C"
+            dsk = ld.get("disk") or {}
+            if dsk:
+                rd = sum(d["read_mb_s"] for d in dsk.values())
+                wr = sum(d["write_mb_s"] for d in dsk.values())
+                io = f"{rd:.0f}/{wr:.0f}"
+            else:
+                io = "-"
+            up = ld.get("host_uptime_s")
+            if up is None:
+                upt = "-"
+            elif up >= 86400:
+                upt = f"{up // 86400}d{(up % 86400) // 3600}h"
+            elif up >= 3600:
+                upt = f"{up // 3600}h{(up % 3600) // 60}m"
+            else:
+                upt = f"{up // 60}m"
+            crows.append(
+                f'<tr><td class="{cls}">{dot}</td><td>{link}{tag}</td>'
+                f'<td>{n["role"]}</td><td class="{cls}">{detail}</td>'
+                f'<td class=n>{el}</td><td class=n>{upt}</td>'
+                f'<td class=n>{cpu}</td><td class=n>{gpu}</td><td class=n>{io}</td></tr>')
+
+        body = f"""<!doctype html><meta charset=utf-8>
+<meta http-equiv=refresh content=5>
+<title>{SERVICE} &middot; {cur}</title>
+<style>
+ body{{font:14px/1.6 ui-monospace,SFMono-Regular,Menlo,monospace;
+      background:#111;color:#ddd;margin:0;padding:2.5rem;}}
+ h1{{font-size:1.1rem;margin:0 0 .25rem;color:#fff}}
+ h2{{font-size:.8rem;margin:1.75rem 0 .5rem;color:#888;font-weight:400;
+     text-transform:uppercase;letter-spacing:.08em}}
+ .sub{{color:#888;margin-bottom:1.5rem}}
+ ul{{list-style:none;padding:0;max-width:46rem}}
+ li{{padding:.35rem 0}} li span{{display:inline-block;width:1.5rem}}
+ .done{{color:#4a8}} .now{{color:#fc6;font-weight:600}} .todo{{color:#555}}
+ table{{border-collapse:collapse;font-size:13px}}
+ th{{text-align:left;color:#888;font-weight:400;padding:.2rem 1.5rem .2rem 0}}
+ td{{padding:.2rem 1.5rem .2rem 0;color:#bbb}}
+ td.n{{text-align:right;font-variant-numeric:tabular-nums}}
+ td.g{{color:#4a8}}
+ td.ok{{color:#4a8}} td.wait{{color:#fc6}} td.bad{{color:#e66}}
+ .me{{color:#666;font-size:11px}}
+ .foot{{margin-top:1.75rem;color:#666;font-size:12px}}
+ a{{color:#6af}}
+</style>
+<h1>{SERVICE} <span style="color:#888;font-weight:400">&middot; {ROLE}</span></h1>
+<div class=sub>{HOSTNAME} &middot; {elapsed//60}m{elapsed%60:02d}s elapsed</div>
+<ul>{''.join(rows)}</ul>
+<h2>cluster</h2>
+<table>
+<tr><th></th><th>node</th><th>role</th><th>stage</th><th>container</th><th>host up</th><th>cpu</th><th>gpu</th><th>disk r/w</th></tr>
+{''.join(crows)}
+</table>
+{tp_html}
+<h2>caches</h2>
+<table>
+<tr><th>cache</th><th>size</th><th>grown</th><th>last write</th></tr>
+{''.join(cache_html)}
+</table>
+<div class=foot>
+ {("Holding a TP rank; the API lives on the head node." if ROLE == "worker"
+    else "API ready at <a href='/v1/models'>:%d</a>" % VLLM_PORT) if healthy
+  else "Starting up; each stage above ticks as it completes."}
+ <br>refreshes every 5s &middot; <a href="/status.json">status.json</a>
+ &middot; <a href="/cluster.json">cluster.json</a>
+ &middot; MCP at <code>/mcp</code> ({len(MCP_TOOLS)} read-only tools)
+</div>"""
+        b = body.encode()
+        self.send_response(200)
+        self.send_header("Content-Type", "text/html; charset=utf-8")
+        self.send_header("Content-Length", str(len(b)))
+        self.end_headers()
+        self.wfile.write(b)
+
+
+class Server(socketserver.ThreadingTCPServer):
+    allow_reuse_address = True
+    daemon_threads = True
+
+
+if __name__ == "__main__":
+    threading.Thread(target=peer_loop, daemon=True).start()
+    threading.Thread(target=host_loop, daemon=True).start()
+    if ROLE != "worker":
+        # Only the head serves an API, so only it has metrics to scrape.
+        threading.Thread(target=throughput_loop, daemon=True).start()
+    Server(("0.0.0.0", PORT), Handler).serve_forever()

docker-images/glm53-kindling/AGENTS.md

@@ -0,0 +1,197 @@
+# Working on this repo
+
+This is the recipe I run GLM-5.3-Flash with on four GB10 boxes. People copy it,
+so a change that makes it slower or subtly wrong costs more than the change
+was worth. Most of what follows is about proving a change is neither before it
+lands.
+
+If you're an agent working here, read this whole file first. If you're a
+person, the same rules apply to you.
+
+## How changes land
+
+Every change goes up as a pull request against `main`. One change per PR,
+with the gate results in the description. I merge.
+
+Before anything leaves your machine, scrub it. No IP addresses, hostnames, box
+names, usernames or paths from your own setup. The compose files use relative
+paths and `.env` values for a reason.
+
+Don't push to a PR branch until the change passes the gate below. Pushing
+something "to see what CI thinks" doesn't work here: there's no CI that can run
+this model, and the PR is the record of what was measured.
+
+A docs-only change doesn't need the gate. A setup fix that helps someone
+starting fresh (an entrypoint guard, a `.env.example` comment, a
+troubleshooting entry) needs its own logic tested, not the speed suite.
+
+## The gate
+
+A change passes when:
+
+1. every quality check below is inside its bar, and
+2. prefill and decode are at par or faster than the last full gate, on
+   repeated samples, not one run.
+
+Run it on the head node, from the repo root, with the stack up and nothing
+else using it:
+
+```
+gate/run.sh full my-change
+```
+
+`gate/run.sh quick my-change` is the reduced gate. It skips HumanEval, the
+count probe, RigMark and half the needle and prefill samples. It's fine for
+checking your own work as you go. It's only enough for a PR when I've said
+so for that change.
+
+The full gate takes a couple of hours. GSM8K and HumanEval need their data in
+`QUALITY_DIR` (see `experimental/quality/quality.py`). RigMark needs a checkout
+of [RigMark](https://github.com/alexellis/rigmark) at c5a0db0 in
+`RIGMARK_DIR`. The gate skips it without one, and then the PR isn't done.
+
+### Quality
+
+| Check | Script | Bar |
+|---|---|---|
+| Smoketest | `smoketest/run.sh` | 8/8 |
+| NLL on three fixed texts | `gate/nll.py compare` | notes +0.008 to +0.025, the others within ±0.007 |
+| GSM8K, first 250 | `experimental/quality/quality.py` | 95.5 to 98.5% |
+| HumanEval | `quality.py`, then `run_he.py` with no network | 154 to 156 of 164 |
+| Count to 200, 5 runs | `gate/count.py` | 0 corrupt |
+| Tool call at 42k context, 40 runs | `dev/repro/toolcall_corruption.py` | 0 diverging on the last gate (stock: 10) |
+| Agent tools | `experimental/quality/agent_tools.py` | 10/10, 0 bad args |
+| Needle, 32k to 480k | `dev/repro/needle.py` | 12/12 |
+| RigMark output gates | RigMark `bench.py` | 6/6 |
+
+The bars are wide because greedy decoding here isn't reproducible across
+boots. FlashInfer's autotuner picks different kernels on each boot, the
+rounding changes, and borderline answers flip. HumanEval has about a dozen
+problems that pass or fail from boot to boot, and
+`experimental/quality/he_detail.py` shows which ones flipped between runs. GSM8K runs 8 requests at once, so
+batching changes run to run too: the same config on the same boot has scored
+96.4 and then 98.4.
+
+So one number near an edge proves nothing. A real regression either falls well
+outside a bar or moves several checks together. The last one I caught (NVFP4
+for the attention projections) took GSM8K to 95.6 and HumanEval to 154 at
+once. When you're unsure, compare the two paths inside one boot, where
+generated text is byte for byte deterministic.
+
+NLL isn't. `nll-ref.json` was recorded with the dense layers left in bf16
+(`fp8.yaml` left out), and FP8 dense moved the notes text up by about 0.015.
+Since then it has read anywhere from +0.008 to +0.025, sometimes 0.006 apart
+on the same boot. Watch the other two texts, and watch for all three moving
+the same way. To add a text, drop it in `gate/texts/`, boot without
+`fp8.yaml`, and record a new reference with `gate/nll.py record`.
+
+### Speed
+
+| Measure | Script | Samples |
+|---|---|---|
+| Prefill, 32k and 128k, cold | `gate/prefill.py` | 4 (mean) |
+| Decode, structured / code / prose | `gate/decode.py` | median of 3 |
+| Decode, older prompt set | `gate/decode_stream.py` | median of 3 |
+| Concurrent streams 1 to 16, mixed prompts | `gate/conc.py` | 1 |
+| Preemptions | `/metrics` | must stay 0 |
+
+Par means within the noise of the last full gate: about 0.5% for prefill and
+1 to 2% for decode. Anything lower is a loss, even if it's small. I would
+rather not have a change than carry a 1% regression forever.
+
+The last full gate (2026-09-28):
+
+| | |
+|---|---|
+| Prefill 32k / 128k | 4,946 / 4,750 tok/s |
+| Decode structured / code / prose | 170.3 / 120.8 / 65.8 tok/s |
+| Decode, older prompts | 162.2 / 127.6 / 89.2 tok/s |
+| Concurrent 1 / 2 / 4 / 8 / 16 | 125 / 103 / 146 / 192 / 251 tok/s |
+| RigMark code / prose / structured | 107.9 / 61.7 / 157.1 tok/s |
+
+When your PR lands with better numbers, update this table in the same PR.
+
+## kindling.json
+
+`kindling.json` at the repo root is the README's results table in a form a
+program can read. The Kindling AI site reads it from `main`, so whatever is
+merged there is what people see.
+
+Keep it in step with the README:
+
+- When a PR changes the README results table, change `kindling.json` in the
+  same PR. Same numbers, same rounding. A PR that changes one and not the
+  other isn't done.
+- It mirrors the README table, not the last-full-gate table above. The gate
+  table moves with every PR; the README is what we publish.
+- One entry in `configs` per TP column. Leave out a number that wasn't
+  measured. Don't write 0, and don't fill a gap with an estimate.
+- Check it parses before you push: `python3 -m json.tool kindling.json`.
+
+The fields:
+
+| field | what |
+|---|---|
+| `title`, `variant`, `summary` | the card heading, the quant tag, one or two sentences |
+| `model`, `hardware`, `stack` | the checkpoint, the boxes, a few short tags |
+| `default_config` | the `id` of the config shown first |
+| `configs[].id`, `label`, `nodes` | e.g. `tp4`, `TP=4`, `4` |
+| `configs[].decode` | single-stream decode tok/s by workload: `code`, `prose`, `structured` |
+| `configs[].prefill` | cold prefill, `[{ "context": tokens, "tps": tok/s }]` |
+| `configs[].concurrency` | aggregate tok/s, `[{ "streams": n, "aggregate_tps": tok/s }]` |
+| `configs[].max_context`, `kv_pool_tokens`, `max_concurrent` | integers: tokens, tokens, requests |
+| `configs[].quality` | `[{ "name": ..., "value": ... }]`, shown as short text |
+| `configs[].notes` | one line of caveats |
+
+Context lengths are in tokens: 32k is `32768`, 128k is `131072`.
+
+## Extra checks for some changes
+
+The gate covers the common cases. Some changes need more.
+
+Kernels and anything else that changes numerics get a bit-exact or
+tolerance test in `dev/patch-tests/` or `dev/kernel-tests/` first, against the
+path they replace. Say in the PR which one it is and why a tolerance is safe,
+if you needed one.
+
+Anything that touches scheduling or runs differently with more than one
+request gets the concurrency sweep too. `gate/conc_workload.py` runs one
+workload (code, prose, structured or mixed) at the concurrency levels you give
+it, each stream with its own prompt:
+
+```
+for w in code prose mixed; do python3 gate/conc_workload.py $w 1 2 4 8 16 32 64; done
+```
+
+Report it against the same sweep on `main`. The targets I care about are 155
+tok/s at 4 streams and 350 at 16, on the mixed set.
+
+Memory and KV changes need the boot log's `GPU KV cache size` line before and
+after, and a run at `MAX_NUM_SEQS` streams with vLLM's `Running:` count, to
+show they all fit.
+
+Anything that changes how weights are processed needs a bump of
+`SNAPSHOT_VERSION` in `experimental/snapshot/weight_snapshot.py` (code) or a
+new `VLLM_WEIGHT_SNAPSHOT_TAG` (environment knobs), or a boot without
+restoring. A shape change fails the restore loudly. A same-shape change
+restores stale weights and says nothing.
+
+## Things that have cost a day
+
+Measure prefill on cold prompts. A repeated prompt measures the prefix cache:
+a 200k probe that took 198 s cold came back in 6 s warm. `gate/prefill.py`
+draws fresh text for each run.
+
+Count tokens, not stream events. With speculative decoding one event carries
+several tokens, and events per second undercount decode badly.
+
+Keep the FlashInfer autotune cache ephemeral. If the four ranks read caches
+that differ, TP=4 deadlocks.
+
+Don't benchmark on a box with a desktop session running. It shares the
+unified memory and the power budget with the model, and the entrypoint warns
+about it.
+
+A worker at `ROLE=head` crash-loops on mentat's "group already has an active
+driver session" while the real head waits for GPUs. The entrypoint refuses to
+start that way now, but it's the first thing to check when a boot hangs.

docker-images/glm53-kindling/KNOBS.md

@@ -0,0 +1,89 @@
+# Knobs
+
+Environment variables read by `entrypoint.sh` and `discover.sh`, extracted mechanically.
+A default of _(none)_ means the variable is referenced without one --
+check the entrypoint for whether it is required or merely optional.
+Meaning and rationale stay in the entrypoint and compose comments.
+
+| variable | default |
+|---|---|
+| `ALLOW_STOCK` | _(empty)_ |
+| `API_PORT` | `8002` |
+| `AUTOTUNE_CACHE` | _(empty)_ |
+| `BLOCK_SIZE` | `2304` |
+| `CACHE_ROOT` | `/root/.cache` |
+| `CACHE_TAG` | `${SHARED_TAG}-${_opthash}` _(derived)_ |
+| `CHAT_TEMPLATE` | `/usr/local/share/glm53-chat-template.jinja` |
+| `CLUSTER_SUBNET` | `${FABRIC_SUBNETS%% *}` _(derived)_ |
+| `CONTAINER_NAME` | `glm53` |
+| `CUDAGRAPH_CAPTURE_SIZES` | `8 16 32 64 96 128 192 256` |
+| `CUDAGRAPH_MAX` | `MAX_NUM_SEQS` x (1 + `SPEC_TOKENS`), at most 512 _(derived)_; unset, vLLM picks it when the draft width varies with batch size |
+| `CUDAGRAPH_MODE` | `FULL_AND_PIECEWISE` |
+| `CUDA_GRAPHS` | `1` |
+| `DECODE_RESERVE_TOKENS` | _(empty)_ |
+| `DFLASH_MODEL` | `/models/glm-5.3-flash-dflash2` |
+| `EXTRA_ARGS` | _(empty)_ |
+| `FABRIC_CHECK` | `1` |
+| `FABRIC_CHECK_GBPS_PER_DEVICE` | `95` |
+| `FABRIC_CHECK_PORT` | `29511` |
+| `FABRIC_CHECK_TIMEOUT_S` | `120` |
+| `FABRIC_SUBNETS` | `$CLUSTER_SUBNET`, else one prefix per `rdma`-tagged address from mentatd _(derived)_ |
+| `FABRIC_WAIT_S` | `0` |
+| `GLOO_SOCKET_IFNAME` | _(empty)_ |
+| `GPU_MEM_UTIL` | `0.88` |
+| `HEAD_HOST` | elected when `ROLE` is also unset _(derived)_ |
+| `ITERATION_DETAILS` | _(empty)_ |
+| `KV_CACHE_DTYPE` | `fp8_e4m3` |
+| `KV_CACHE_MEMORY` | `27917287424` |
+| `LIMIT_MM` | `{\"image\":16,\"video\":0\}` |
+| `LONG_PREFILL_TOKEN_THRESHOLD` | `2304` |
+| `MAX_JOBS` | `4` |
+| `MAX_MODEL_LEN` | `524288` |
+| `MAX_NUM_BATCHED_TOKENS` | `16384` |
+| `MAX_NUM_SEQS` | `32` |
+| `MCP_LOG_DIR` | `/logs` |
+| `MENTAT_GROUP` | `${SERVICE_NAME:-glm53}` _(derived)_ |
+| `MENTAT_MCP_API` | `${STATUS_PORT:-8082}/mcp` _(derived)_ |
+| `MENTAT_MODEL_PROVIDER` | `vllm` |
+| `MENTAT_NODE_IP` | mentatd's `node_ip` when electing _(derived)_ |
+| `MENTAT_OPENAI_API` | `${API_PORT:-8002}/v1` _(derived)_ |
+| `MODEL_DIR` | `/models/glm-5.3-flash-nvfp4` |
+| `MOE_BACKEND` | `flashinfer_cutlass` |
+| `MTP` | `1` |
+| `NCCL_DEBUG` | `INFO` |
+| `NCCL_IB_GID_INDEX` | _(empty)_ |
+| `NCCL_IB_HCA` | _(empty)_ |
+| `NCCL_MAX_NCHANNELS` | `8` |
+| `NCCL_SOCKET_IFNAME` | `$GLOO_SOCKET_IFNAME` _(derived)_ |
+| `PREFLIGHT` | `1` |
+| `PYTORCH_CUDA_ALLOC_CONF` | `expandable_segments:True` |
+| `RAY_ADDRESS` | `127.0.0.1:6379` when electing, else `${HEAD_HOST:?set HEAD_HOST to the head node address}:6379` _(derived)_ |
+| `RAY_MEMORY_MONITOR_REFRESH_MS` | `0` |
+| `RAY_OBJECT_STORE_MEMORY` | `4294967296` |
+| `ROCE_SETTLE_S` | `60` |
+| `ROLE` | elected when `HEAD_HOST` is also unset, else `head` |
+| `SAFETENSORS_LOAD_STRATEGY` | `eager` |
+| `SELF_TEST` | `1` |
+| `SERVED_NAME` | `glm53` |
+| `SERVICE_NAME` | `glm53` |
+| `SHARED_TAG` | `glm53-${_arch}-${_kv}` _(derived)_ |
+| `SKIP_MM_PROFILING` | `0` |
+| `SPEC_EXTRA` | _(empty)_ |
+| `SPEC_METHOD` | `dflash` |
+| `SPEC_TOKENS` | `7` |
+| `STAGE_FILE` | `/tmp/glm53-stage` |
+| `STATUS_PORT` | `8082` |
+| `TILELANG_CACHE_DIR` | `${CACHE_ROOT}/${SHARED_TAG}/tilelang` _(derived)_ |
+| `TOOL_PARSER` | `glm47` |
+| `TOPK_BACKEND` | `per_row` |
+| `TORCH_MEM_FRACTION` | `0.92` |
+| `TP` | `4` (or `2`, `3`, `RING4`) |
+| `TRITON_CACHE_DIR` | `${CACHE_ROOT}/${SHARED_TAG}/triton` _(derived)_ |
+| `VLLM_ARX_TWO_SHOT_MAX_KB` | `2048` (with arx.yaml; `0` sends all-reduces over 256 KB to NCCL; NCCL is faster at 4 MB) |
+| `VLLM_DRAFT_TRUNC_EXPLORE` | `4` (with adaptive-k.yaml; every Nth step skips the draft cut, and only those steps train the estimator) |
+| `VLLM_DRAFT_TRUNC_TAU` | `0.3` (with adaptive-k.yaml; `0` turns the draft cut off) |
+| `VLLM_ENGINE_READY_TIMEOUT_S` | `3600` |
+| `VLLM_GLM5NEXT_RECOVERSSM` | `0` (1 with recoverssm.yaml) |
+| `VLLM_HOST_IP` | the `lan`-tagged address from mentatd, else its `node_ip` _(derived)_ |
+| `VLLM_WEIGHT_SNAPSHOT_DIR` | _(empty)_ |
+| `WORKER_WAIT_S` | `0` |

docker-images/glm53-kindling/NOTES.md

@@ -0,0 +1,250 @@
+# Notes
+
+The measurements and diagnosis behind the recipe in [README.md](README.md).
+Dates are when each was measured; image tags (v4, v6, sm90-v21) are the
+builds of this recipe at the time.
+
+## Moving to a vLLM nightly, 2026-09-23
+
+Before the nightly, this ran a per-model base image with 16 vLLM source files
+bind-mounted from five compose overrides (in git history, up to 6524882). Every one of those
+files is now upstream, a flag, or a small anchored patch baked into the image,
+and nothing is mounted over it. Images v4 to v7 used `nightly-0961bbae`; v8
+moved to `nightly-ddd6fbca` on 2026-09-26 and has not been measured.
+
+Measured on the same four nodes, thinking off, median of 3, 2026-09-23:
+
+| | old sm90-v21 + DFlash2 | v4 + DFlash2 |
+|---|---|---|
+| decode structured / code / prose | 105.8 / 71.9 / 50.7 tok/s | 104.4 / 79.1 / 53.4 |
+| acceptance | 89.0 / 58.3 / 36.0 % | 88.5 / 65.0 / 38.2 % |
+| prefill 31.6k / 126k | 2,229 / 2,176 tok/s | 2,389 / 2,421 |
+| KV pool at the 26 GiB pin | 2.02M tokens | 2.63M |
+| needle 33k+136k | 6/6 | 6/6 |
+| greedy @42k, 8 runs | 7 distinct | 4 distinct |
+
+v6 (thinking off as low effort, below) decodes 109.8 / 88.8 / 52.6 on the same
+three, because DFlash2's acceptance rose. MTP (the checkpoint's own head,
+`SPEC_METHOD=mtp`) also works, and decoded 57.2 / 54.4 / 45.6 on v4: correct,
+but DFlash2 is what carries structured and code.
+
+### What it took, in the order it broke
+
+1. mentat 0.7.0 refuses a bare `ray start --head` while `RAY_ADDRESS` is set.
+   The entrypoint passes `--address`.
+2. The MTP layer of nvidia/GLM-5.3-Flash-NVFP4 is BF16 but `config.json` says
+   NVFP4 (`image/patches/glm53_mtp_bf16.py`).
+3. The SM90 sparse-MLA builder planned an fp8 KV cache as uint8
+   (`image/patches/sm90_fp8_kv_dtype.py`).
+4. The GLM-5-Next KV grouper drops out on the drafter's sliding-window layers
+   (`image/patches/glm53_dflash2_kv_groups.py`, tested by
+   `dev/patch-tests/_glm53_dflash2_kv_groups_test.py`).
+5. A persisted FlashInfer autotune cache deadlocks the next TP=4 boot: rank 0
+   saves per-rank MoE entries the others then miss. The entrypoint keeps it
+   ephemeral and wipes it on every start.
+
+mentat drops a remote rank's traceback, so a worker that raises during setup
+shows only the message on the head. To see the stack, bind-mount a copy of
+`ray/_host.py` with `traceback.print_exc()` added to both `except BaseException`
+handlers; that is how 2 and 3 were found.
+
+## The model
+
+320B total / 18B active, natively multimodal MoE, `Glm5NextForConditionalGeneration`.
+45 layers split **34 KDA linear-attention + 11 DeepSeek-sparse-attention**, 288
+routed experts (8 active + 1 shared), an MTP head, 1M context.
+
+Only the 11 sparse layers carry a KV cache and they use `kv_lora_rank: 512`, so
+KV is cheap and the memory pressure is essentially all weights. The 26 GiB pin
+holds 2,632,595 tokens with DFlash2 at a 524288 window, 5.02 full-length
+requests (2026-09-23).
+
+The KDA layers keep one recurrent state per sequence, allocated after weights,
+and exactly 32 fit (2026-08-26). That, not throughput, caps `MAX_NUM_SEQS`.
+
+## Why NVFP4 and not the official checkpoint
+
+Measured when the model was released (2026-08-26), at TP=2:
+
+| checkpoint | size | per node at TP=2 | fits? |
+|---|---|---|---|
+| BF16 | 598.5 GiB | 299 GiB | no |
+| FP8 e4m3 (the default) | 305.8 GiB | 153 GiB | **no** — 121.6 GiB per node |
+| NVFP4 | **181.3 GiB** | **90.7 GiB** | yes |
+
+Sizes are from the safetensors indexes, not estimates. The routed experts are
+NVFP4; both attention flavours, the vision tower, shared experts, routers,
+embeddings and `lm_head` stay BF16. The deployed checkpoint is nvidia's
+(`nvidia/GLM-5.3-Flash-NVFP4`); the LibertAIDAI and RedHat quants were served
+before it.
+
+**There is no usable GGUF.** llama.cpp has no `glm5_next` support, so the GGUF
+repos on HF are not loadable by anything.
+
+## What the image is
+
+A pinned vLLM nightly plus anchored patches, each described in the Dockerfile
+header, and:
+
+- **mentat** — the daemon binary as `ray`, `mentatd-probe-machine`, and the
+  shim wheel. TP=4 across four boxes needs real placement, which a TP=1 tenant
+  gets from `ray.register` alone.
+- **iproute2 and ibverbs** — the entrypoint finds this node's fabric interface
+  with `ip`, and NCCL needs ibverbs for RoCE (without it NCCL falls back to TCP
+  over the LAN and every step crawls). Not replaceable by a hardcoded interface
+  name: the nodes carry the link on different ports.
+- **spark-agent's status server**, plus ripgrep, pciutils and usbutils for its
+  tools.
+
+GB10 is sm_121 and the nightly's torch reports up to sm_120. That reads like a
+blocker and is not one: sm_120 cubins run on sm_121. `verify-base.py` asserts
+it, with every patch marker, the `ray` binary and `ray.register`, so a rebased
+base fails the build instead of producing an image that loads ~90 GiB and then
+dies.
+
+## Why it loads plain safetensors
+
+`--load-format auto`, and deliberately no sharded-state cache. The fast-boot
+cache is dumped *after* the NVFP4 kernel-format transform, so loading one
+re-runs that transform, permutes the fused gate/up halves and serves fluent
+nonsense with no error. A guard in the loader that skips an already
+transformed tensor prevents this. This image has no such guard, so the only
+safe load is the one that runs the transform exactly once.
+
+Adopt the sharded-state cache only together with that guard. The cost of not
+having it is boot time, which is CPU-bound on weight processing.
+
+## PP is off, and turning it on needs files this repo no longer has
+
+`pp-mtp-override.yaml` was in the old deployment until 2026-09-10. Every file
+it mounted was PP-only or inert at `pipeline_parallel_size=1`: two activate
+solely when `VLLM_NCCL_NET_*` is set, `mtp.py` is gated on
+`get_pp_group().world_size > 1`, `speculative.py` computed the value PP=1
+already yields, and its `model.py` and `model_runner.py` were shadowed by the
+dflash copies at the same destinations. Removing it left draft acceptance at
+1.79 of 7 against a control of 1.75, and decode medians of 41.1 against 38.9
+with spreads of 9.4 and 6.0. Its files are in this repo's history
+(`git log --all -- '*pp-mtp-override*'`); port them before raising PP above 1.
+
+## The RoCE GID index is per node and per boot
+
+The entrypoint derives `NCCL_IB_GID_INDEX` at start; do not pin it. The RoCE
+GID table is indexed by (address, RoCE version) in the order addresses
+appeared, so nodes do not agree and a removed address leaves a hole. Read on
+2026-08-26 from two nodes:
+
+| node | RoCE device | table | value |
+|---|---|---|---|
+| A | `rocep1s0f1` | 0,1,2,**hole**,4,5,6 | **6** |
+| B | `rocep1s0f0` | 0,1,2,3,4,5 | **5** |
+
+Reboot A with the static profile in place and its table comes up dense, moving
+6 to 5, so a pinned 6 names an empty slot and NCCL fails every TP init. To read
+a table:
+
+    for i in $(seq 0 9); do p=/sys/class/infiniband/<dev>/ports/1; \
+      echo "$i $(cat $p/gid_attrs/types/$i) $(cat $p/gids/$i)"; done
+
+The right entry is the RoCE v2 one for the node's own static fabric address,
+unlike the IPv6 link-locals, which regenerate.
+
+## Thinking off is outside what the model was trained on
+
+The official template (zai-org FP8, and the RedHat checkpoint) has no
+non-thinking mode at all: every kwarg renders
+`<|system|>Reasoning Effort: Max ... <|assistant|><think>`. Our first fix
+closed the block empty (`<think></think>`), and long structured output then
+degraded. Measured 2026-09-23, "count from 1 to 200", temperature 0:
+
+| mode | nvidia NVFP4, cutlass, MTP | official FP8, triton |
+|---|---|---|
+| thinking on (default) | clean 8/8 | clean 6/6 |
+| thinking off (`<think></think>`) | corrupt 3-8 of 8 | corrupt 8/8 |
+
+The failures are repeats (`35 36 37 37`), skips, or a jump to copying the prompt
+text. Every checkpoint (nvidia, RedHat, official FP8), both MoE backends,
+weight-only marlin everywhere, fp8 or bf16 KV, TP=2 or 4, graphs or eager, and
+the old sm90-v21 image all fail the thinking-off probe the same way. It is not
+a serving bug: the per-layer references in `dev/kernel-tests/*_ref.py` match
+the decode path. Short answers and the Korean corruption probe are fine with
+thinking off; long repetitive output is not.
+
+**Since v6, thinking off means low effort.** The official template takes
+`reasoning_effort` (`low`, `high`, anything else is `max`), passed at the top
+level of the request or in `chat_template_kwargs`. `image/chat-template.jinja`
+maps `thinking: false` / `enable_thinking: false` to `Reasoning Effort: Low`
+(an explicit `reasoning_effort` wins) and always opens `<think>`.
+`image/patches/glm53_reasoning_always_parsed.py` makes the reasoning parser
+track `<think>` on every request, since stock glm47_moe stops parsing it when
+either kwarg is false and the short trace would land in `content`. Measured on
+v6 (2026-09-23): thinking off gives ~67 reasoning chars and a clean count 6/6.
+
+## v8 against v7, 2026-09-26
+
+Same four boxes and scripts, v7 measured just before the switch. v7 ran with
+two overrides that v8 has built in: the SM90 kpool plan (vllm#58704) and fp32
+FlashKDA state (vllm#58846).
+
+| | v7 | v8 |
+|---|---|---|
+| decode counting / code / prose | 101.4 / 75.6 / 34.1 tok/s | 109.9 / 84.4 / 38.4 |
+| acceptance | 96.1 / 72.2 / 23.5 % | 97.2 / 74.7 / 24.5 |
+| prefill @31k / @117-124k | 2,699 / 2,662 tok/s | 2,730 / 2,679 |
+| tool-call probe, 40 runs at 42k | 19 agree with the majority | 22 |
+| boot to API | ~10 min | 7.5 min (weights 238 s) |
+
+Prefill is at parity. The first long prompt after a boot measured 2,309 tok/s
+and is left out: it pays for JIT compiles. The tool-call probe still diverges,
+with 14 and 16 distinct completions, so read 19 against 22 as noise.
+
+## Both ConnectX roots, 2026-09-26
+
+The 3,365 tok/s at 126k recorded on 2026-09-23 was a partial prefix-cache hit.
+The benchmark seeded its random words by seed alone and sized the prompt from a
+calibration run, so two runs with the same seed and slightly different
+calibrations shared every word up to the shorter length. Reproduced: the same
+seed with a different calibration measured 3,391 and 3,396 tok/s, with ~30k
+tokens served from the cache. Fresh seeds measure ~2,680 on v7 and v8 alike.
+
+Measured with the model stopped, all-reduce bus bandwidth across the four boxes
+(bf16, 134 MB to 1 GB, `NCCL_MAX_NCHANNELS=8` unless noted):
+
+| | Gb/s |
+|---|---|
+| both roots, NCCL's choice | 175-191 |
+| both roots, Ring | 180-191 |
+| both roots, Tree | 55-93 |
+| one root | 95-111 |
+| both roots, 16 channels | 156-160 |
+
+Ring far ahead of Tree is the healthy ordering; the latched slow state after
+a DAC hot-plug reverses it. First-touch prefill, fresh random words, after a
+warm-up request:
+
+| | 32k | 126k |
+|---|---|---|
+| one root | 2,470 tok/s | 2,412 |
+| both roots | 2,750 | 2,680 |
+
+So the second root is worth 11% of prefill, not 40%. Collectives are a
+minority of prefill time, which is why doubling their bandwidth moves it this
+little.
+
+`NCCL_MAX_NCHANNELS=8` was chosen on 2026-09-06 while the fabric ran at 12
+Gb/s. Remeasured on the healthy fabric (2026-09-29, first-touch prefill, two
+alternating passes of three warmed runs):
+
+| channel cap | 32k | 128k |
+|---|---|---|
+| 8 (the entrypoint default) | 4,905 tok/s | 4,701 |
+| 32 (NCCL's own choice) | 4,769 | 4,595 |
+
+So 8 stays. Forcing `NCCL_PROTO=Simple` was also 0.5-1.5% slower: NCCL picks
+LL even for 50 MB messages here, but prefill's NCCL time is ranks waiting for
+the slowest one, not the protocol.
+
+## Not yet measured
+
+- Why production moved from marlin to `flashinfer_cutlass` on 2026-09-21 is not
+  recorded. It passes the thinking-on corruption probe; the greedy
+  determinism repros in `dev/repro/` were measured on marlin.

docker-images/glm53-kindling/README.md

@@ -0,0 +1,89 @@
+# glm53-kindling — vendored kindlingai/glm-5.3-flash-gx10 (B-mode)
+
+Upstream: <https://github.com/kindlingai/glm-5.3-flash-gx10> @ `c748079d45e6`
+(re-vendored from `27cf839c2133`; carries the TP3/TP6 DFlash block-drop fix
++ the optional dispram display-RAM overlay).
+(our fork name for their tree; original mmastrac/glm-5.3-flash-4x-gx10).
+GLM-5.3-Flash NVFP4 on 2/4/6 GB10 boxes: stock vLLM nightly + 15 anchored
+patches + the `experimental/` overlay stack (arx RDMA collectives, weight
+snapshots, adaptive-k scheduling, fp8 dense, megamoe, profiling fixes,
+sequence parallelism, RecoverSSM), DFlash2 drafter, mentat instead of Ray.
+
+## Status: integration mode B (2026-10-01)
+
+Operator decision: **image + patches under sm121 rendezvous first**, mode A
+(their full entrypoint + resident `mentatd` daemons) only if B proves out.
+
+What B keeps vs skips:
+
+| kept (baked/transcribed) | skipped (mode A territory) |
+|---|---|
+| their image build (all 15 patches, FlashKDA rebuild, chat template) | `/entrypoint.sh` orchestration + `discover.sh` |
+| all nine overlay file sets — **baked** by `build.sh`, not mounted | resident `mentatd` host daemons (`ray` shim, 6379-6382) |
+| every measured knob: args + env transcribed 1:1 per TP | preflight / fabric-check / status server / self-test |
+| DFlash2 spec, AdaptiveK scheduler, per-TP KV pin + limits | `MODEL_HOST_DIR` local-disk rule (NAS, operator call) |
+
+Recipes: `recipes/glm53-kindling-{2x,4x,6x}/recipe.toml`. The overlay
+`environment:` blocks from `experimental/compose/*.yaml` live in recipe
+`[env]`, the `entrypoint.sh` `exec vllm serve …` construction in `[args]`.
+`glm53-kindling` == `entrypoint` defaults == "the order the numbers at the
+top [of the upstream README] were measured".
+
+## Deviations (each in the recipe headers too)
+
+- **Executor**: `mp` + sm121 rank flags instead of mentat's Ray surface.
+- **Weights on NFS**: upstream requires each box's local disk (boot-time NFS
+  races — upstream README). We serve from the NAS HF cache (operator call;
+  download `nvidia/GLM-5.3-Flash-NVFP4` + `incoai/GLM-5.3-Flash-DFlash2`
+  land there). Freeing ~181 GiB local per node ×4 was not an option.
+- **Snapshot tags are TP-scoped** (`fp8w4c-tp4-…`): upstream snapshots are
+  per TP rank; its single shared tag is safe only on a fixed-topology box
+  set. TP2/TP4 recipes must never share the dir.
+- **Port 30000** (`API_PORT` upstream default 8002), served name `glm53`
+  (upstream default, already simple).
+- Chat template stays upstream's (`thinking:false` → low reasoning effort;
+  GLM-5.3-Flash has no true non-thinking mode).
+
+Full upstream documentation (their README for mode A reference) is kept
+verbatim as `README.upstream.md`. KNOBS.md / NOTES.md / AGENTS.md are
+theirs and unmodified.
+
+Note on `27cf839c2133` (#64): upstream moved spec-config assembly into
+the entrypoint (overlays now pass `SPEC_EXTRA`, a second
+`--speculative-config` is refused — adaptive-k's EXTRA_ARGS used to
+silently win the same way our first transcription predated the fix).
+B-mode recipes keep the effective merged config inline in `[args]`,
+matching the pre/post-fix TP4 result upstream live-checked — no change
+needed, but re-transcribe from the new EXTRA/SPEC_EXTRA shape if a
+recipe ever grows a mode-A switch.
+
+## Build
+
+```
+sm121 build                        # build.sh, tag spark-glm53:v9 (default)
+TAG=spark-glm53:v10 ./build.sh      # or directly on a node
+```
+
+Context = this directory (repo root upstream; `.dockerignore` from upstream
+applies). `Dockerfile.bake`/`Dockerfile.sm121` are generated at build time
+(node-side, not committed).
+
+## Vendoring an upstream update
+
+```
+git clone --recursive --depth 1 https://github.com/kindlingai/glm-5.3-flash-gx10 /tmp/new
+rsync -a --delete --exclude build.sh --exclude README.md … /tmp/new/ ./
+```
+
+then re-check: `image/entrypoint.sh` (knob construction — recipes mirror
+it), `experimental/compose/*.yaml` (overlays/env), `KNOBS.md`.
+
+## Model contract
+
+- `nvidia/GLM-5.3-Flash-NVFP4` rev `da920bb0b9f…` (181 GiB, ModelOpt NVFP4
+  — NOT the local-inference-lab/-Spark cuts; their MoE kernels read the
+  checkpoint's own per-projection input scales).
+- `incoai/GLM-5.3-Flash-DFlash2` rev `bf582e4eacc…` (2.2 GiB, **CC-BY-NC-ND
+  4.0 non-commercial** — upstream flags this; same family we already serve).
+- TP3/TP6 only: a zero-padded checkpoint from
+  `experimental/tp3/make_padded.py` (66/66/2304 heads don't split by 3/6).

docker-images/glm53-kindling/build.sh

@@ -0,0 +1,59 @@
+#!/usr/bin/env bash
+# [sm121 vendor] Build the vendored kindlingai/glm-5.3-flash-gx10 image.
+#
+# Upstream builds `image/Dockerfile` against the repo root (the `vllm` symlink
+# into `.submodules/spark-agent` must stay inside the context). At runtime the
+# `glm53` script stacks eight `experimental/compose/*.yaml` overlays that
+# mount replacement Python/CUDA files over site-packages — the measured
+# "arx snapshot adaptive-k fp8 megamoe fixes sp recoverssm" stack.
+#
+# sm121 recipes mount no file trees from the node, so this build BAKES those
+# overlay files into the image: a generated Dockerfile.bake COPY stage
+# appended after upstream's own Dockerfile (used unmodified — the bake step
+# runs after verify-base.py). The per-overlay env/args live in the recipes
+# (recipes/glm53-kindling-*), mirroring the compose environment blocks.
+#
+# Deltas vs upstream c748079d45e6: only this file and README.md; every
+# upstream file is byte-identical. To vendor a new upstream: clone fresh
+# (with submodules), rsync -a over this dir keeping build.sh + README.md.
+#
+# c748079d carries the TP3/TP6 DFlash block-drop fix (drop the restated
+# SPEC_EXTRA so adaptive-k's copy, with disable_eagle_block_drop, applies)
+# and the optional dispram overlay (display-RAM KV carveout, no-op without
+# kindling-spark-os's dispramd — we don't run it, so it stays out of the
+# bake stack).
+set -euo pipefail
+cd "$(dirname "$0")"
+TAG="${TAG:-spark-glm53:v10}"
+
+# Overlay stack (bake order irrelevant: identical target paths). tp3.yaml
+# carries the only files TP6 needs — the pads activate only when the config
+# holds tp_pad_orig, a no-op on TP4.
+OVERLAYS=(arx snapshot adaptive-k fp8 megamoe fixes sp recoverssm tp3)
+
+python3 - "${OVERLAYS[@]}" > Dockerfile.bake <<'PY'
+import re, sys
+
+print("# [sm121] generated by build.sh: experimental/ overlays baked in")
+print("# (content parity with experimental/compose/*.yaml runtime mounts).")
+seen = set()
+for name in sys.argv[1:]:
+    for line in open(f"experimental/compose/{name}.yaml"):
+        if not re.match(r"\s*-\s+\.\./experimental/", line):
+            continue
+        parts = line.strip()[len("- ") :].split(":")
+        if len(parts) < 2 or "${" in parts[0]:
+            continue
+        src, dst = parts[0][len("../") :], parts[1]
+        if (src, dst) in seen:
+            continue
+        seen.add((src, dst))
+        print(f"COPY {src} {dst}")
+PY
+
+cat image/Dockerfile Dockerfile.bake > Dockerfile.sm121
+echo "baked $(grep -c '^COPY experimental' Dockerfile.bake) overlay files"
+
+# shellcheck disable=SC2086
+docker buildx build --progress=plain --load -t "$TAG" -f Dockerfile.sm121 . 2>&1
+echo "=== Done: $TAG ==="

docker-images/glm53-kindling/compose/glm53.yaml

@@ -0,0 +1,276 @@
+# GLM-5.3-Flash NVFP4, TP=4 across four GB10 nodes (one head, three workers),
+# with the DFlash2 drafter.
+#
+# Run the SAME file on every node:
+#   docker compose -f compose/glm53.yaml up -d
+# compose/.env is optional. Addresses and roles come from mentat, and paths
+# and the image have defaults (.env.example lists the overrides). Compose
+# reads .env from this file's directory, whatever the working directory.
+#
+# Ordering: from cold, any rank may start first -- registration retries
+# forever, and the host-level mentatd (mentatd.yaml from
+# https://github.com/mmastrac/mentat) just has to be running on every box. Replacing a running stack is different: take all four
+# ranks down first (README, "Replacing a running stack").
+#
+# Nothing else fits beside it. GLM's weights alone are ~90.7 GiB per node out
+# of 121.6 GiB, so stop any other model on these boxes first.
+#
+# network_mode: host is required -- RoCE/NCCL needs the real interfaces, and
+# Ray's actor ports are not practically mappable. It is also how the
+# container reaches the box's mentatd, on 127.0.0.1:6379 and :6380.
+#
+# Pinned so the project is not named after this directory ("compose"). Before
+# the first `up`, take down a stack started under another project name: this
+# one reuses container_name glm53 and would collide with it.
+name: glm53
+services:
+  glm53:
+    # image/build.sh's default tag. Keep the two in step: the entrypoint's
+    # defaults live in the image, and an older tag boots silently mistuned.
+    image: ${IMAGE:-spark-glm53:v9}
+    container_name: glm53
+    network_mode: host
+    ipc: host
+    ulimits:
+      memlock: -1
+      stack: 67108864
+      # A crash must not dump core. The engine maps ~90 GiB, and the dumping
+      # thread holds the NVIDIA RM lock the whole time it writes: on
+      # 2026-08-27 that turned one CUDA OOM on a node into a box with no
+      # ssh, no userspace and no way back except a power cycle, while the
+      # kernel sat in nvidia_close for 860+ seconds. core_pattern on these
+      # hosts is a plain file, so nothing rate-limits or discards it either.
+      core: 0
+    devices:
+      - /dev/infiniband:/dev/infiniband
+    volumes:
+      # The full 181 GiB checkpoint, read-only. There is no transcoded cache:
+      # this image loads plain safetensors, so EVERY rank reads the whole
+      # checkpoint and it must be present on EACH node's local disk.
+      # Local rather than NFS on purpose -- an NFS mount races the network at
+      # boot and fails (seen on one node on 2026-08-26), and a model that only
+      # starts when the NAS is awake will eventually fail to start.
+      - ${MODEL_HOST_DIR:-/srv/models/glm-5.3-flash-nvfp4}:/models/glm-5.3-flash-nvfp4:ro
+      # The host's list of running systemd units, read by the preflight to warn
+      # about a desktop session (README, "What you need").
+      - /run/systemd/units:/host/systemd-units:ro
+      # The DFlash2 drafter, also on every node: speculation runs on every TP
+      # rank. A separate Qwen3-based model that reads auxiliary hidden states
+      # from target layers 5, 14, 24, 33 and 42 (image/patches/glm53_eagle3_aux.py
+      # supplies them; upstream GLM5next does not).
+      - ${DFLASH_HOST_DIR:-/srv/models/glm-5.3-flash-dflash2}:/models/glm-5.3-flash-dflash2:ro
+      # JIT caches (FlashInfer, torch.compile, Triton, TileLang), and the
+      # experimental weight snapshots (~48 GB a node). WITHOUT this they live
+      # inside the container and are discarded on every `compose down`, so
+      # each restart recompiles. The FlashInfer autotune cache is the
+      # exception: at TP>1 the entrypoint keeps it in /tmp and wipes it.
+      # Unset, a named volume. `down -v` deletes it.
+      - ${CACHE_HOME:-cache}:/root/.cache
+      # The engine log and the allocator counters behind /memory, read through
+      # the status server's MCP tools. Unset, a named volume.
+      - ${LOG_DIR:-logs}:/logs
+    # RULE: any knob that also has a default in entrypoint.sh is left EMPTY
+    # here. A default in both places means two places to change and one
+    # silently winning -- which happened repeatedly on 2026-08-27, with the
+    # compose defaults from the pre-SM90 attempts quietly overriding tuned
+    # values (MTP off, GPU_MEM_UTIL 0.86, MAX_MODEL_LEN 131072). The exceptions
+    # are the knobs the status server reads itself (SERVICE_NAME, SERVED_NAME,
+    # PEERS, MCP_LOG_DIR, THROUGHPUT_INTERVAL_S). Its own defaults are
+    # generic, not this model's, so they are set here explicitly.
+    environment:
+      # Passed through only when .env (or the shell running compose) sets
+      # them. Unset, the entrypoint takes them from the box's mentatd.
+      #
+      # ROLE and HEAD_HOST, both unset: the boxes elect the live agent with
+      # the lowest address as head. Set, the roles are fixed: ROLE is head on
+      # one box and worker on the rest, and HEAD_HOST is the head's LAN
+      # address on every box, the head included.
+      - ROLE
+      - HEAD_HOST
+      # This node's LAN address. mentat identifies a node by it, and the agent
+      # and daemon must agree on one string. Unset = the address mentatd tags
+      # lan.
+      - VLLM_HOST_IP
+      # Address prefix of the fabric, with its trailing dot. Unset = the first
+      # of FABRIC_SUBNETS.
+      - CLUSTER_SUBNET
+      # The status server reads these two itself, and its default name is
+      # "vllm", so they carry a default here too.
+      - SERVICE_NAME=${SERVICE_NAME:-glm53}
+      - SERVED_NAME=${SERVED_NAME:-glm53}
+      # mentat subcluster identity; empty = SERVICE_NAME (the entrypoint's
+      # default). Set it only to run the same model as a second parallel
+      # stack, which needs a distinct group per stack.
+      - MENTAT_GROUP=${MENTAT_GROUP:-}
+      # Empty = 127.0.0.1:6379 when the boxes elect a head, else HEAD_HOST:6379.
+      - RAY_ADDRESS=${RAY_ADDRESS:-}
+      # One subnet per ConnectX PCIe root NCCL should use, space-separated. Each
+      # names an interface; the entrypoint maps it to its RDMA device and GID.
+      # Both roots made a 126k prefill 11% faster on TP=4 (2026-09-26). Empty =
+      # CLUSTER_SUBNET when that is set, else one subnet per address mentatd
+      # tags rdma.
+      - FABRIC_SUBNETS=${FABRIC_SUBNETS:-}
+      # How long to wait for the fabric and for the peer's GPUs before giving
+      # up. Both default to waiting forever, because the usual reason either is
+      # missing is a switch reboot, and a container that exits then just races
+      # it. Set a bound only where a stuck stack must surface as a failure.
+      - FABRIC_WAIT_S=${FABRIC_WAIT_S:-}
+      - WORKER_WAIT_S=${WORKER_WAIT_S:-}
+      # How long the RoCE device and its GID table get to appear once the
+      # address is already up. Local state, so it settles in seconds or is
+      # broken: this one is deliberately bounded, and the container exits so
+      # the failure is visible rather than a container that waits forever.
+      - ROCE_SETTLE_S=${ROCE_SETTLE_S:-}
+      # LEFT EMPTY ON PURPOSE so the entrypoint derives it at start. There is
+      # no correct constant: the GID table is indexed by (address, RoCE version)
+      # with first-free allocation, so the index differs per node AND per boot.
+      # On 2026-08-26 the right entry was 6 on one node and 5 on its peer, and the
+      # first node's table had an empty slot 3 left by four failed DHCP cycles plus
+      # a profile replacement during an earlier boot. Reboot that node with the static
+      # profile in place and its correct index moves 6 -> 5, so a pinned value
+      # is a latent NCCL failure rather than a fix. Set it only to override.
+      - NCCL_IB_GID_INDEX=${NCCL_IB_GID_INDEX:-}
+      # Empty = all four RDMA devices, which is what reaches 196 Gb/s rather
+      # than 112. Override only to restrict NCCL to one PCIe root.
+      - NCCL_IB_HCA=${NCCL_IB_HCA:-}
+      # Empty = the interface holding VLLM_HOST_IP, derived in the entrypoint.
+      - NCCL_SOCKET_IFNAME=${NCCL_SOCKET_IFNAME:-}
+      # Channel count and NCCL logging; defaults and why in the entrypoint.
+      - NCCL_MIN_NCHANNELS=${NCCL_MIN_NCHANNELS:-}
+      - NCCL_MAX_NCHANNELS=${NCCL_MAX_NCHANNELS:-}
+      - NCCL_DEBUG=${NCCL_DEBUG:-}
+      - TP=${TP:-}
+      # 1 lets the stack start without any experimental/ overlay (stock
+      # kernels, about half the README's speed); otherwise the entrypoint refuses.
+      - ALLOW_STOCK=${ALLOW_STOCK:-}
+      - PREFLIGHT=${PREFLIGHT:-}
+      - FABRIC_CHECK=${FABRIC_CHECK:-}
+      # Speculation. SPEC_METHOD is dflash, mtp or none (entrypoint default
+      # dflash, k=7); MTP=0 turns it off whichever is named. An MTP drafter
+      # costs ~3.4 GiB and a rejection-sampler warmup, which were what killed
+      # rank 0 on 2026-08-27 while this comment wrongly said MTP was off.
+      - MTP=${MTP:-}
+      - SPEC_METHOD=${SPEC_METHOD:-}
+      - SPEC_TOKENS=${SPEC_TOKENS:-}
+      - SPEC_EXTRA=${SPEC_EXTRA:-}
+      - DFLASH_MODEL=${DFLASH_MODEL:-}
+      - MODEL_DIR=${MODEL_DIR:-}
+      # Default and KV arithmetic in the entrypoint. Raise it against the
+      # "GPU KV cache size" line at boot.
+      - MAX_MODEL_LEN=${MAX_MODEL_LEN:-}
+      # Per-prompt multimodal caps. Raising the image cap raises the peak a
+      # single request can demand. With SKIP_MM_PROFILING=1 the vision encoder
+      # is never budgeted at startup, so those bytes come out of the same
+      # headroom a long prefill needs.
+      - LIMIT_MM=${LIMIT_MM:-}
+      - SKIP_MM_PROFILING=${SKIP_MM_PROFILING:-}
+      # lazy|eager|prefetch|torchao; eager in the entrypoint, which says why.
+      - SAFETENSORS_LOAD_STRATEGY=${SAFETENSORS_LOAD_STRATEGY:-}
+      # Which NVFP4 MoE kernel: flashinfer_cutlass (entrypoint default) or
+      # marlin. The entrypoint has the history.
+      - MOE_BACKEND=${MOE_BACKEND:-}
+      # 0.88 is the measured ceiling under mentat (2026-08-27): CUDA-visible
+      # free memory on the head is 108.0 GiB regardless of page cache, 0.90
+      # wants 109.46 and failed the startup check four times, 0.88 wants
+      # 107.03 and served with 0.63 GiB to spare -- KV 1,842,951 tokens vs
+      # ~1.07M at the old 0.84. Under real Ray the ceiling was lower still
+      # (its object store and daemons sat in this pool).
+      #
+      # 0.90 starts on a freshly booted box, which is the trap. It is worth
+      # real KV -- 1,541,820 tokens against 1,200,343 here, +28%, because KV
+      # gets only what the weights leave -- and it still must not be used: it
+      # ran the driver out of memory 90 seconds into serving
+      # (NV_ERR_NO_MEMORY in dmesg, hours before anything visible went wrong)
+      # and both ranks died with the boxes needing a power cycle. A driver OOM
+      # during startup PROFILING is normal at any setting; one during steady
+      # serving is the signal.
+      - GPU_MEM_UTIL=${GPU_MEM_UTIL:-}
+      # Allocator and per-worker memory ceiling; defaults and why in the
+      # entrypoint.
+      - PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF:-}
+      - TORCH_MEM_FRACTION=${TORCH_MEM_FRACTION:-}
+      # Set to a writable directory to record an allocation stack per block
+      # and dump a snapshot the first time the ceiling above is crossed. The
+      # stack is taken in the allocating thread, so it names the caller. It
+      # costs time and memory to record, so leave it empty outside an
+      # investigation. /logs is already mounted and MCP-readable.
+      - TORCH_MEM_TRACE_DIR=${TORCH_MEM_TRACE_DIR:-}
+      # Where each worker publishes its allocator counters, which is how the
+      # status page's /memory answers whether torch or something else is
+      # holding the node's memory. One counter read per worker per interval,
+      # so unlike the trace above this is cheap enough to leave on.
+      - TORCH_MEM_STATS_DIR=${TORCH_MEM_STATS_DIR:-/logs}
+      - TORCH_MEM_STATS_INTERVAL=${TORCH_MEM_STATS_INTERVAL:-}
+      # Tool-call parser. Unset = the image's default: the stock glm47, or
+      # glm47_failclosed on an image built before that became the default
+      # (the default lives in the image's entrypoint). Set glm47 or
+      # glm47_failclosed (image/patches/glm47_failclosed.py) to choose.
+      - TOOL_PARSER=${TOOL_PARSER:-}
+      # Appended to the serve command verbatim. Nothing the entrypoint selects
+      # belongs here: EXTRA_ARGS is one string, and an override restating it
+      # drops whatever else it carried -- how the fail-closed tool parser once
+      # vanished from production.
+      - EXTRA_ARGS=${EXTRA_ARGS:-}
+      # KV pool in bytes; the entrypoint pins 26 GiB. Pinned, vLLM "skips
+      # memory profiling" AND "does not respect the gpu_memory_utilization
+      # config", so the pin is not a fix for a warmup OOM -- lower GPU_MEM_UTIL
+      # or the pin for that.
+      - KV_CACHE_MEMORY=${KV_CACHE_MEMORY:-}
+      # Empty = fp8_e4m3, the entrypoint's default and what every measurement
+      # here used. Never plain `fp8`: that selects the fp8_ds_mla packed format,
+      # which hardcodes pe_dim == 64. This model is NoPE sparse MLA
+      # (qk_rope_head_dim is 0), so that kernel refuses at KV cache init --
+      # after a full 10-minute weight load.
+      - KV_CACHE_DTYPE=${KV_CACHE_DTYPE:-}
+      # Left EMPTY on purpose so the entrypoint owns the defaults. A default
+      # here as well as there means two places to change and one silently
+      # winning.
+      - MAX_NUM_BATCHED_TOKENS=${MAX_NUM_BATCHED_TOKENS:-}
+      # Capped by the KDA recurrent state, not by throughput. 34 of 45 layers
+      # are linear attention and each keeps one fixed-size state per sequence
+      # ("Mamba cache block"); exactly 32 fit after weights at GPU_MEM_UTIL
+      # 0.86. vLLM's default of 256 aborts at CUDA graph capture AFTER the full
+      # weight load. Empty here so the entrypoint owns the default.
+      - MAX_NUM_SEQS=${MAX_NUM_SEQS:-}
+      - LONG_PREFILL_TOKEN_THRESHOLD=${LONG_PREFILL_TOKEN_THRESHOLD:-}
+      - DECODE_RESERVE_TOKENS=${DECODE_RESERVE_TOKENS:-}
+      - CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-}
+      - CUDAGRAPH_METRICS=${CUDAGRAPH_METRICS:-}
+      - ITERATION_DETAILS=${ITERATION_DETAILS:-}
+      # Empty = 8002 (API) and 8082 (status page and MCP), from the
+      # entrypoint.
+      - API_PORT=${API_PORT:-}
+      - STATUS_PORT=${STATUS_PORT:-}
+      # Status-page peers, host:port, comma-separated. Empty shows this node
+      # alone; mentatd-serve gives the cluster view.
+      - PEERS=${PEERS:-}
+      - MCP_LOG_DIR=${MCP_LOG_DIR:-/logs}
+      # mentat's agent id is group@container@node_ip. Without this the agent
+      # falls back to the hostname, which under network_mode: host is the
+      # HOST's name, so every container on a node would claim the same id.
+      - CONTAINER_NAME=${CONTAINER_NAME:-glm53}
+      # What the agent announces to mentatd-serve. Empty = port form (8002/v1,
+      # 8082/mcp), which the router resolves against every address the node
+      # announces. A URL pins one address. A worker holds a TP rank and
+      # answers no requests, so the entrypoint drops its OpenAI endpoint
+      # however this is set.
+      - MENTAT_OPENAI_API=${MENTAT_OPENAI_API:-}
+      - MENTAT_MCP_API=${MENTAT_MCP_API:-}
+      - MENTAT_MODEL_PROVIDER=${MENTAT_MODEL_PROVIDER:-}
+      - MAX_JOBS=${MAX_JOBS:-}
+      - SELF_TEST=${SELF_TEST:-}
+      - THROUGHPUT_INTERVAL_S=${THROUGHPUT_INTERVAL_S:-60}
+      - NVIDIA_VISIBLE_DEVICES=all
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: all
+              capabilities: [gpu]
+    restart: unless-stopped
+
+# Used only when CACHE_HOME or LOG_DIR is unset.
+volumes:
+  cache:
+  logs:

docker-images/glm53-kindling/dev/README.md

@@ -0,0 +1 @@
+Dev-only material that is not in the image: the token-corruption diagnosis and its repros, kernel tests, patch tests, the step tap, and `entrypoint-spec-test.sh` (the entrypoint's speculative config under each setting, no GPU needed).

docker-images/glm53-kindling/dev/TOPK-CORRUPTION.md

@@ -0,0 +1,120 @@
+# Greedy decoding is not deterministic at long context
+
+GLM-5.3-Flash-NVFP4, vLLM `0.1.dev20051+g487ecf187`, 4x GB10 (sm121) at TP=4.
+
+## What happens
+
+At `temperature=0` with a fixed seed, identical requests sometimes return
+different completions. Most differences are invisible. The one that reaches a
+user is a corrupted **tool name**: `glm47_moe` sets `validate_tool_names=True`,
+so a call whose name is not in the request's tool list emits **zero deltas** and
+finishes `stop` with no content and no tool calls. Clients report it as "model
+returned a completed response with no content" and retry it as a dead endpoint.
+
+## Cause
+
+**The marlin NVFP4 MoE is not bit-deterministic at M >> 1**, and the model
+amplifies that into whole-token differences.
+
+Measured, in order:
+
+1. Among runs that answer *correctly* at the failing position, P(correct token)
+   ranges 1.00 down to 0.48 -- 35 distinct values in 38 runs. A wrong token is
+   the tail of that spread, not a separate failure mode.
+2. Per-layer tap across identical runs: layers 0-2 and layer 3's attention input
+   are bit-identical; layer 3's attention output takes 12 distinct values.
+3. Fixing the FA2 input order makes layer 3 bit-identical, and the first
+   divergence moves to layer 4 with identical input -- so it arrives in the KDA
+   recurrent state, which prefill produced.
+4. Prefill tap: layer 3's `mlp_out`, the first marlin MoE, differs in 8 of 8
+   runs at M=800.
+5. In-engine, evaluating that MoE block twice on the same input differs at
+   M=800 and M=2304, and is bit-stable at M=1.
+6. Padding the prompt so the uncached recompute is a single row makes decode
+   bit-identical across 11 runs, spread 0.000000.
+
+Amplification is what turns ulp noise into a different token: 288-expert
+routing plus eleven sparse selections per token. With layer 3 pinned, the
+indexer's 512th logit still spreads by 1-30 units by layer 39.
+
+## Do not use the prefill re-score as an oracle
+
+An earlier version of this document treated a token's logprob under
+`prompt_logprobs` as ground truth and read a 13-21 logprob gap as proof that
+decode was wrong. It is not: the re-score is its own multi-row prefill, so it is
+another draw of the same nondeterministic computation, taken at n=1. It shows
+that two draws disagreed, not which one is right.
+
+## Ruled out, by measurement
+
+Top-k selection, both kernels: values match `torch.topk` exactly at every
+coarse-bin population up to a whole row in one bin, and real indexer logits hold
+exact boundary ties in 3 of 34,320 layer-rows, none at the layer where
+divergence enters. The FA2 sparse-MLA kernel (exact, bit-repeatable). DeepGEMM
+MQA logits (bit-deterministic; paged equals non-paged). CUDA graphs (eager
+reproduces). Speculative decoding (off reproduces). KV dtype, memory pressure,
+prefix caching, temperature, conversation history, the request payload, the
+uninitialised `topk_indices_buffer`, the MoE router's Python fallback (the fused
+kernel runs here), and TP/NCCL (ranks agree bitwise).
+
+Consequently neither `hpc.topk_filtered` (Tencent/hpc-ops#93) nor the upstream
+sm121 determinism set -- vllm#55122, the MoE finalize knob, vllm#53899 -- touches
+this. `use_fused_finalize` belongs to the FlashInfer CUTLASS path; the backend
+here is marlin.
+
+## Reading the rates
+
+Per-condition rates ranged 2.5%-22.5%, including the same configuration twice at
+13% and 22.5%. At n=40 nothing in that spread is distinguishable, and comparing
+conditions on those numbers cost a night. Size the sample from the expected rate
+first, or measure determinism directly with a per-layer tap, which decides on a
+single pair of runs.
+
+## Writing a tap
+
+Three ways it silently produces nothing, all of them paid for:
+
+- `GLM53_TAP` in the container does not reach the worker: mentat spawns actors
+  without forwarding it. Gate on a file under `/logs` instead.
+- Gating on `t.shape[0] == 1` to exclude prefill excludes all of decode, because
+  speculative decode carries 1 + draft tokens per step.
+- `.item()` inside the forward invalidates CUDA graph capture and crash-loops the
+  engine. Accumulate into a preallocated device tensor, and keep the step counter
+  on device too: host code does not run on graph replay, so a Python counter
+  freezes at its capture value.
+
+## The MoE kernel is not the whole story
+
+2026-09-13, NVIDIA's ModelOpt NVFP4 checkpoint with `MOE_BACKEND=flashinfer_cutlass`
+(which does run on sm_121 -- the entrypoint's `cudaErrorNoKernelImageForDevice`
+warning and the upstream issues are about auto-selection skipping CUTLASS, not
+about the kernel being absent):
+
+| condition | distinct completions from 16 |
+|---|---|
+| marlin, RedHat weights, content only, 48 tokens | 3 |
+| CUTLASS, NVIDIA weights, content only, 48 tokens | 9, sizes [6,2,2,1,1,1,1,1,1] |
+| CUTLASS, NVIDIA weights, reasoning included, 400 tokens | 16 |
+
+Replacing the kernel this document blames did not restore reproducibility, and
+under matched conditions it is worse. Two variables differ between the first two
+rows, kernel and checkpoint, so the attribution needs marlin on the NVIDIA
+weights -- one flag, the same boot -- before the kernel is cleared or convicted.
+One sample of 16 each, and the sizing warning above applies.
+
+The third row is not comparable to the others: it hashes the reasoning trace as
+well as the content, which finds divergence a 48-token content hash cannot. That
+is the better test of the forward pass and the worse test for comparing against
+the baseline.
+
+## Open
+
+A deterministic NVFP4 MoE at M >> 1 on sm121, in any backend. The
+single-row-recompute trick removes the variance but only for prompts whose
+uncached tail is one token.
+
+## Reproducing
+
+`repro/` holds three scripts and what each one measures. `repro/greedy_nondet.py`
+is the small one: standard library, one repeated paragraph, 3 distinct
+completions from 16 identical greedy requests at 42k and 1 from 16 at 2k.

docker-images/glm53-kindling/dev/entrypoint-spec-test.sh

@@ -0,0 +1,41 @@
+#!/usr/bin/env bash
+# Tests the entrypoint's speculative config (image/entrypoint.sh, from `SPEC=()` to its "speculative:" line) without a GPU:
+# runs those exact lines under each setting and checks the one --speculative-config vLLM gets. The reference for the default
+# stack is what vLLM used before SPEC_EXTRA: the second --speculative-config adaptive-k.yaml put in EXTRA_ARGS (vLLM keeps
+# the last one). Usage: dev/entrypoint-spec-test.sh   (prints PASS or the failing case; exit status = failures)
+cd "$(dirname "$0")/.."
+block=$(sed -n '/^SPEC=()$/,/^\[\[ \${#SPEC\[@\]} -gt 0 \]\] && echo "speculative:/p' image/entrypoint.sh)
+extra=$(sed -n 's/^ *- SPEC_EXTRA=//p' experimental/compose/adaptive-k.yaml)
+d1=$(mktemp -d); d2=$(mktemp -d); touch "$d1/config.json" "$d2/config.json"; trap 'rm -rf "$d1" "$d2"' EXIT
+fails=0
+run() {  # run <env assignments...>: prints the final SPEC[1] (or nothing), exit status of the block
+  env -i PATH="$PATH" BLOCK="$BLOCK" "$@" bash -c 'eval "$BLOCK" >/dev/null; printf %s "${SPEC[1]:-}"' 2>/dev/null
+}
+export BLOCK="$block"
+check() {  # check <name> <python condition on c (the parsed config, or None) and rc>
+  local name=$1 cond=$2; shift 2
+  out=$(run "$@"); rc=$?
+  python3 -c 'import json, sys; out, rc = sys.argv[1], int(sys.argv[2]); c = json.loads(out) if out else None
+d1, d2 = sys.argv[4], sys.argv[5]; sys.exit(0 if eval(sys.argv[3]) else 1)' "$out" "$rc" "$cond" "$d1" "$d2" \
+    || { echo "FAIL $name: rc=$rc out=$out"; fails=$((fails + 1)); }
+}
+ref='{"method": "dflash", "disable_eagle_block_drop": True, "model": d1, "num_speculative_tokens": 7, "attention_backend": "TRITON_ATTN", "num_speculative_tokens_per_batch_size": [[1,1,7],[2,2,5],[3,3,4],[4,4,3],[5,5,2],[6,32,7]]}'
+check "default stack = the old overlay's config" "rc == 0 and c == $ref" SPEC_METHOD=dflash DFLASH_MODEL="$d1" SPEC_EXTRA="$extra"
+for f in tp3 tp6; do  # these restate EXTRA_ARGS after adaptive-k.yaml and keep its SPEC_EXTRA (#66)
+  ! grep -q '^ *- SPEC_EXTRA=' "experimental/compose/$f.yaml" || { echo "FAIL $f.yaml sets its own SPEC_EXTRA"; fails=$((fails + 1)); }
+  check "$f.yaml = the default stack's config" "rc == 0 and c == $ref" SPEC_METHOD=dflash DFLASH_MODEL="$d1" \
+    SPEC_EXTRA="$extra" EXTRA_ARGS="$(sed -n 's/^ *- EXTRA_ARGS=//p' experimental/compose/$f.yaml)"
+done
+check "stock (no overlay) unchanged" "rc == 0 and out == json.dumps({'method': 'dflash', 'model': d1, 'num_speculative_tokens': 7}, separators=(',', ':'))" SPEC_METHOD=dflash DFLASH_MODEL="$d1"
+check "SPEC_TOKENS takes effect, widths capped" "rc == 0 and c['num_speculative_tokens'] == 5 and [r[2] for r in c['num_speculative_tokens_per_batch_size']] == [5,5,4,3,2,5]" SPEC_METHOD=dflash SPEC_TOKENS=5 DFLASH_MODEL="$d1" SPEC_EXTRA="$extra"
+check "DFLASH_MODEL takes effect" "rc == 0 and c['model'] == d2" SPEC_METHOD=dflash DFLASH_MODEL="$d2" SPEC_EXTRA="$extra"
+check "mtp with the DFlash overlay is refused" "rc != 0" SPEC_METHOD=mtp SPEC_EXTRA="$extra"
+check "none with the DFlash overlay is refused" "rc != 0" SPEC_METHOD=none SPEC_EXTRA="$extra"
+check "mtp without the overlay still works" "rc == 0 and c == {'method': 'mtp', 'num_speculative_tokens': 4}" SPEC_METHOD=mtp
+check "a second --speculative-config in EXTRA_ARGS is refused" "rc != 0" SPEC_METHOD=dflash DFLASH_MODEL="$d1" EXTRA_ARGS='--speculative-config {"method":"dflash"}'
+check "SPEC_EXTRA cannot set the method, model or width" "rc != 0" SPEC_METHOD=dflash DFLASH_MODEL="$d1" SPEC_EXTRA='{"model":"/x"}'
+check "SPEC_EXTRA must be a JSON object" "rc != 0" SPEC_METHOD=dflash DFLASH_MODEL="$d1" SPEC_EXTRA='[1]'
+grep -q '"${SPEC\[\*\]:-}" == \*num_speculative_tokens_per_batch_size\*' image/entrypoint.sh \
+  || { echo "FAIL CUDA graph ceiling still reads the draft table from EXTRA_ARGS"; fails=$((fails + 1)); }
+[ $fails = 0 ] && echo PASS
+exit $fails

docker-images/glm53-kindling/dev/kernel-tests/README.md

@@ -0,0 +1,76 @@
+# Decode-path kernel tests (GLM-5.3-Flash on GB10, spark-glm53:v4)
+
+Written 2026-09-23 while chasing the count-to-200 token corruption. Each file is
+self-contained and runs inside the image on one GB10; none needs the weights.
+
+    scp <file> admin@<box>:/tmp/kvtest/
+    sudo docker run --rm --gpus all --ipc=host -v /tmp/kvtest:/kv \
+        --entrypoint python3 -w / localhost:5000/spark-glm53:v4 /kv/<file>
+
+| file | what it proves | result on v4 |
+|---|---|---|
+| `_topk_mhcpost_fp8mla_test.py` | T1 per_row top-k + expand_pools_and_append_tail return exactly the causal set for short rows under any garbage past the row; T3 mhc_post_tilelang is pure and matches the torch reference; T5 SM90 fp8 MLA wrapper planned as float8_e4m3fn (patches/sm90_fp8_kv_dtype.py) matches a bf16 reference at <=0.5% on 8-row verify shapes and across 2048 | PASS |
+| `_kda_spec_vs_plain_test.py` | T2 causal_conv1d_update + fused_recurrent_kda spec path (k=3 and k=7, random acceptance) is bit-identical to the one-token-per-step plain path; run at H=64 and H=16 | PASS |
+| `_cudagraph_replay_test.py` | T6a SM90 wrapper captured at the dummy lens vLLM uses (1..8), replayed at real lens after replan: exact. T6b DeepGEMM paged MQA logits + top-k under replay still select the full causal pool set (its logits reference is wrong: my test lays scales per row, DeepGEMM keeps them per page; ignore rel_err) | T6a PASS, T6b selection PASS |
+| `_cutlass_nvfp4_moe_rows_test.py` | T7 FlashInfer CUTLASS NVFP4 fused MoE at the per-rank shape (E=288,K=4096,N=512,topk 8): every batched row (M up to 64) is bit-identical to the same row run alone; repeat calls are deterministic | PASS |
+| `_cutlass_nvfp4_moe_sanity_test.py` | T7b outputs are finite, non-zero, input-dependent and row-local (my torch dequant reference disagrees by ~70%; that is the reference's scale convention, not the kernel) | as described |
+
+`_cudagraph_replay_test.py` needs the import fix already applied (vllm.utils.platform_utils).
+
+Added after the spec-off result (plain decode corrupts too):
+
+| file | what it proves | result on v4 |
+|---|---|---|
+| `_tail_compaction_test.py` | T8 production-shaped top-k rows (pools, -1 gap, tail at columns 2048+) through the SM90 index remap: every tail slot is inside the first `ctx` compacted entries the plan reads, eager and under graph replay, positions 99..5003 | PASS |
+| `_kda_numerics_test.py` | T9 FlashKDA prefill (the nightly's new SM12x default), Triton chunk prefill (old image path) and fused_recurrent_kda plain decode vs an fp32 torch reference of the kernel math; state seeded from prefill; L=20/300/2500 | PASS (<=1.1% bf16 error, decode state rel err ~1e-6) |
+| `_cudagraph_replay_shapes_test.py` | T6a at the spec-off shape: one request captured with dummy seq_len 1, replayed at 100/700/2051 and back down to 37 | PASS (the R=4 case aborts on a test-side cache size, not a kernel error) |
+
+## Collective-corruption test on the four GLM boxes (no model loaded)
+
+`nccl_collective_test.py` + `nccl_test_launch.sh`. The launcher derives the
+per-box network exactly as entrypoint.sh does (LAN address for the bootstrap
+socket, RoCE devices and GID index for the data plane, NCCL_MAX_NCHANNELS=8)
+and runs the test, which emulates a TP=4 decode step's collectives (90 bf16
+all-reduces of [M,4096] per step, one bf16 logits all-gather [M,38720],
+a [2304,4096] "prefill" all-reduce every 200 steps) with deterministic inputs
+and checks every result (bit-exact for integer data and all-gathers, 2 bf16
+ulps for float data, all ranks identical, repeat identical).
+
+Copy both files to the same path on all four boxes (e.g. /home/admin/kernel-tests).
+Ranks: node0 (192.0.2.1) = 0, 192.0.2.2 = 1, 192.0.2.3 = 2, 192.0.2.4 = 3. Run the four
+commands within ~60 s of each other (rank 0 first). VLLM_HOST_IP is the box's
+LAN address (what production passes); the launcher derives the RoCE side.
+
+    # on each box, with RANK and VLLM_HOST_IP set for that box:
+    sudo docker run --rm --gpus all --network host --ipc host --ulimit memlock=-1 \
+      --device /dev/infiniband:/dev/infiniband --cap-add SYS_PTRACE \
+      -v /home/admin/kernel-tests:/kt \
+      -e RANK=<0..3> -e WORLD_SIZE=4 -e MASTER_ADDR=192.0.2.1 -e MASTER_PORT=29555 \
+      -e VLLM_HOST_IP=<this box's 192.0.2.x> -e CLUSTER_SUBNET=198.51.100. \
+      $VARIANT --entrypoint bash spark-glm53:v5 /kt/nccl_test_launch.sh --mode vllm --iters 2000
+
+Variants (VARIANT):
+  default              (nothing)
+  simple protocol      -e NCCL_PROTO=Simple
+  ring algorithm       -e NCCL_ALGO=Ring
+  no RDMA (TCP)        -e NCCL_IB_DISABLE=1
+  both HCAs            -e FABRIC_SUBNETS="198.51.100. 203.0.113."   (only if production uses both)
+  communicator         --mode pynccl   |   --mode torch          (instead of --mode vllm)
+
+Each rank prints its NCCL_* env, a progress line every 100 steps, every
+failure with step/collective/M/first bad index, and rank 0 prints
+`RESULT: failures per rank [...], collectives per rank N, CLEAN|CORRUPT`.
+2000 steps x 90 all-reduces x 2 = 360k all-reduces plus 2000 all-gathers;
+expect a few minutes. NCCL_DEBUG=INFO on top shows the devices NCCL chose.
+
+## Real-weight MoE value test (T10)
+
+`_cutlass_nvfp4_moe_realweights_test.py [nvidia|rh]` needs `/kv/moe-l5-<name>.safetensors`
+(16 experts of layer 5 plus the gate, extracted read-only from the checkpoint; the
+extractor is scratch code, ~220 MiB each). Result on v4, both checkpoints: the
+FlashInfer CUTLASS kernel matches the w4a4 reference (dequantised weights, activations
+quantised by vLLM's own scaled_fp4_quant with the checkpoint's global input scale)
+within 2-4%; the w4a4 result differs from the w4a16 quantity (what marlin computes)
+by 13-20% at |x|<~2, 25-30% at 2x that, 62% once inputs exceed the calibrated amax
+(6 * 448 * input_scale = 1.88 for gate/up on both checkpoints), because the static
+global activation scale clips every element above it.

docker-images/glm53-kindling/dev/kernel-tests/_cudagraph_replay_shapes_test.py

@@ -0,0 +1,235 @@
+#!/usr/bin/env python3
+"""T6: FULL-cudagraph replay checks for the two decode kernels whose plan is
+built host-side: the SM90 FlashInfer MLA wrapper (_SM90State) and the
+DeepGEMM paged MQA logits + per_row top-k of the sparse indexer.
+
+vLLM captures decode graphs with a dummy batch whose seq_len equals the query
+length (InputBatch.make_dummy: seq_lens = ceil(num_tokens/num_reqs), positions
+zero), then replans every step and replays. This mirrors that: capture with
+8 rows of lens 1..8, replay with real lens, compare against a reference.
+"""
+import sys
+import traceback
+
+import torch
+
+torch.manual_seed(0)
+dev = torch.device("cuda")
+results = {}
+
+
+def section(name):
+    def deco(fn):
+        def run():
+            print(f"\n===== {name} =====", flush=True)
+            try:
+                fn()
+                results[name] = "PASS"
+            except Exception as e:  # noqa: BLE001
+                results[name] = f"FAIL: {e}"
+                traceback.print_exc()
+        return run
+    return deco
+
+
+def t6a_impl(R, cap_lens, replays):
+    from vllm import _custom_ops as ops
+    from vllm.v1.attention.backends.mla.flashinfer_mla_sparse_sm90 import _SM90State
+
+    heads, d = 16, 512
+    sm_scale = 1.0 / 16.0
+    topk_width = 2176
+    max_tokens = 256
+    N = 2300
+    bs = 64
+    num_blocks = N // bs + 3
+    kv_c = torch.randn(N, d, device=dev, dtype=torch.bfloat16)
+    k_pe = torch.empty(N, 0, device=dev, dtype=torch.bfloat16)
+    kv_cache = torch.zeros(num_blocks, bs, d, device=dev, dtype=torch.uint8)
+    slots = (torch.arange(N, device=dev) + 37).to(torch.int64)
+    ops.concat_and_cache_mla(kv_c, k_pe, kv_cache, slots, "fp8_e4m3", torch.tensor(1.0, device=dev))
+    deq = kv_cache.view(torch.float8_e4m3fn).reshape(-1, d).float()
+    flat = kv_cache.view(torch.float8_e4m3fn).reshape(-1, 1, d)
+    ckv, kpe = flat[..., :d], flat[..., d:]
+
+    state = _SM90State(dev, heads, torch.float8_e4m3fn, max_tokens, topk_width,
+                       kv_lora_rank=d, qk_rope_head_dim=0, sm_scale=sm_scale)
+    # static input buffers (what the captured graph reads)
+    q_nope = torch.randn(R, heads, d, device=dev, dtype=torch.bfloat16)
+    q_pe = torch.empty(R, heads, 0, device=dev, dtype=torch.bfloat16)
+    out_buf = torch.empty(R, heads, d, device=dev, dtype=torch.bfloat16)
+
+    def set_rows(lens):
+        state.kv_indices.fill_(0)
+        for r, L in enumerate(lens):
+            state.kv_indices[r * topk_width: r * topk_width + L] = slots[:L].to(torch.int32)
+        state.plan(R, torch.tensor(lens, dtype=torch.int32))
+
+    def reference(lens):
+        ref = torch.empty(R, heads, d, device=dev)
+        for r, L in enumerate(lens):
+            K = deq[slots[:L]]
+            s = (q_nope[r].float() @ K.T) * sm_scale
+            ref[r] = torch.softmax(s, dim=-1) @ K
+        return ref
+
+    def run_once():
+        out = state.wrapper.run(q_nope, q_pe, ckv, kpe, ckv_scale=1.0, kpe_scale=1.0)
+        out_buf.copy_(out)
+
+    # --- capture with the dummy lens vLLM uses ---
+    set_rows(cap_lens)
+    s = torch.cuda.Stream()
+    s.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(s):
+        for _ in range(2):
+            run_once()
+    torch.cuda.current_stream().wait_stream(s)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        run_once()
+    torch.cuda.synchronize()
+
+    worst = 0.0
+    for lens in replays:
+        q_nope.copy_(torch.randn_like(q_nope))
+        set_rows(lens)
+        graph.replay()
+        torch.cuda.synchronize()
+        ref = reference(lens)
+        err = (out_buf.float() - ref).abs()
+        rel = err.max().item() / max(ref.abs().max().item(), 1e-6)
+        row_err = [round(v, 4) for v in err.flatten(1).max(dim=1).values.tolist()]
+        nan = torch.isnan(out_buf).any().item()
+        print(f"  replay lens={lens[0]}..{lens[-1]}: rel err {rel:.4g} nan={nan} per-row {row_err}")
+        worst = max(worst, rel)
+        assert not nan
+    print(f"T6a R={R} cap={cap_lens} worst replay rel err {worst:.4g}")
+    assert worst < 0.1
+
+
+@section("T6a SM90 MLA wrapper: capture at dummy lens, replay at real lens")
+def t6a():
+    t6a_impl(8, list(range(1, 9)), [list(range(1, 9)), list(range(100, 108)), list(range(693, 701)), list(range(2044, 2052))])
+    # spec-off production shape: one request, one token, dummy seq_len 1
+    t6a_impl(1, [1], [[1], [100], [700], [2051], [100], [37], [2600 // 4 * 4 + 3], [5]])
+    # MTP k=3 shape: one request, 4 tokens, dummy seq_lens 1..4
+    t6a_impl(4, [1, 2, 3, 4], [[1, 2, 3, 4], [100, 101, 102, 103], [700, 701, 702, 703], [2049, 2050, 2051, 2052], [40, 41, 42, 43]])
+
+
+@section("T6b indexer paged MQA logits + per_row topk: capture short, replay long")
+def t6b():
+    from vllm.utils.deep_gemm import fp8_fp4_paged_mqa_logits, get_paged_mqa_logits_metadata
+    from vllm.models.glm5next.nvidia.ops import kpool_compress as kpool_ops
+    from vllm import _custom_ops as ops
+    from vllm.utils.platform_utils import num_compute_units
+
+    # Indexer cache: [num_blocks, pool_block_size, 1, head_dim+4] uint8 (fp8 + fp32 scale)
+    head_dim = 128
+    kpool = 4
+    index_heads = 32  # index_n_heads
+    pool_bs = 64  # production indexer kernel page: 64 pools = 256 tokens
+    num_blocks = 80
+    max_pool_len = 524288 // kpool
+    fp8 = torch.float8_e4m3fn
+    k_pool = torch.randn(num_blocks * pool_bs, head_dim, device=dev).clamp(-4, 4)
+    kv_cache = torch.zeros(num_blocks, pool_bs, 1, head_dim + 4, device=dev, dtype=torch.uint8)
+    kv_cache[..., :head_dim] = k_pool.to(fp8).view(torch.uint8).reshape(num_blocks, pool_bs, 1, head_dim)
+    kv_cache[..., head_dim:] = torch.ones(num_blocks * pool_bs, 1, 1, device=dev).view(torch.uint8).reshape(num_blocks, pool_bs, 1, 4)
+    num_sms = num_compute_units(dev.index if dev.index is not None else 0)
+
+    R = 8
+    # per-token expanded block table (flattened path): every row -> request 0's blocks [1..num_blocks-1]
+    block_table = torch.zeros(R, num_blocks, device=dev, dtype=torch.int32)
+    block_table[:, : num_blocks - 1] = torch.arange(1, num_blocks, device=dev, dtype=torch.int32)
+    q = torch.randn(R, 1, index_heads, head_dim, device=dev).to(fp8)
+    weights = torch.rand(R, index_heads, device=dev, dtype=torch.float32)
+    seq_lens = torch.zeros(R, 1, device=dev, dtype=torch.int32)  # pool-granular, refreshed
+    sched_buf = torch.empty((num_sms + 1, 2), dtype=torch.int32, device=dev)
+    select_k = 2048 // kpool
+    pool_topk = torch.full((R, select_k), -1, dtype=torch.int32, device=dev)
+    logits_holder = {}
+
+    def set_lens(pool_lens):
+        seq_lens[:, 0] = torch.tensor(pool_lens, device=dev, dtype=torch.int32)
+        meta = get_paged_mqa_logits_metadata(seq_lens, pool_bs, num_sms)
+        sched = sched_buf[: meta.shape[0]]
+        sched[:] = meta
+        return sched
+
+    def run_once(sched):
+        logits = fp8_fp4_paged_mqa_logits((q, None), kv_cache, weights, seq_lens, block_table,
+                                          sched, max_model_len=max_pool_len, clean_logits=False)
+        logits_holder["l"] = logits
+        pool_topk.fill_(-1)
+        ops.top_k_per_row_decode(logits, 1, seq_lens, pool_topk, R, logits.stride(0), logits.stride(1), select_k)
+
+    def reference(pool_lens):
+        # logits[r, p] = sum_h w[r,h] * relu(q[r,h] . k[p]) (DeepSeek indexer semantics)
+        out = []
+        for r, L in enumerate(pool_lens):
+            blocks = block_table[r, : (L + pool_bs - 1) // pool_bs]
+            K = k_pool.view(num_blocks, pool_bs, head_dim)[blocks.long()].reshape(-1, head_dim)[:L]
+            Kq = K.to(fp8).float()
+            s = torch.relu(q[r, 0].float() @ Kq.T)  # [heads, L]
+            out.append((weights[r][:, None] * s).sum(0))
+        return out
+
+    cap_lens = [max(1, (i + 1) // kpool) for i in range(1, R + 1)]  # dummy: seq_len = query_len
+    sched = set_lens(cap_lens)
+    s = torch.cuda.Stream()
+    s.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(s):
+        for _ in range(2):
+            run_once(sched)
+    torch.cuda.current_stream().wait_stream(s)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        run_once(sched)
+    torch.cuda.synchronize()
+
+    bad = 0
+    for tok_lens in (list(range(100, 108)), list(range(693, 701)), list(range(1500, 1508)), list(range(2600, 2608)), list(range(4600, 4608))):
+        pool_lens = [t // kpool for t in tok_lens]
+        # the schedule buffer is refreshed in place; its slot count must be stable
+        meta = get_paged_mqa_logits_metadata(seq_lens.new_tensor([[p] for p in pool_lens]), pool_bs, num_sms)
+        assert meta.shape[0] == sched.shape[0], f"schedule slot count changed {meta.shape} vs {sched.shape}"
+        set_lens(pool_lens)
+        graph.replay()
+        torch.cuda.synchronize()
+        logits = logits_holder["l"]
+        refs = reference(pool_lens)
+        for r, L in enumerate(pool_lens):
+            got = logits[r, :L].float()
+            ref = refs[r]
+            nan = torch.isnan(got).any().item()
+            err = (got - ref).abs().max().item() / max(ref.abs().max().item(), 1e-6)
+            sel = pool_topk[r][pool_topk[r] >= 0].tolist()
+            if L <= select_k:
+                ok_sel = set(sel) == set(range(L)) and len(sel) == L
+            else:
+                ok_sel = set(sel) == set(ref.topk(select_k).indices.tolist())
+            if nan or err > 0.05 or not ok_sel:
+                bad += 1
+                miss = sorted(set(range(min(L, select_k))) - set(sel))[:8] if L <= select_k else []
+                print(f"  BAD replay tok_len={tok_lens[r]} pools={L}: nan={nan} rel_err={err:.3g} "
+                      f"sel_ok={ok_sel} nsel={len(sel)} missing_pools={miss}")
+        print(f"  replay tok_lens {tok_lens[0]}..{tok_lens[-1]}: checked")
+    # eager comparison
+    pool_lens = [t // kpool for t in range(693, 701)]
+    sched = set_lens(pool_lens)
+    run_once(sched)
+    torch.cuda.synchronize()
+    refs = reference(pool_lens)
+    e = max(((logits_holder["l"][r, :L].float() - refs[r]).abs().max() / refs[r].abs().max()).item() for r, L in enumerate(pool_lens))
+    print(f"  eager 693..700 rel err {e:.3g}")
+    print(f"T6b bad rows: {bad}")
+    assert bad == 0
+
+
+for fn in (t6a,):
+    fn()
+print("\n===== SUMMARY =====")
+for k, v in results.items():
+    print(f"{k}: {v}")
+sys.exit(0 if all(v == "PASS" for v in results.values()) else 1)

docker-images/glm53-kindling/dev/kernel-tests/_cudagraph_replay_test.py

@@ -0,0 +1,242 @@
+#!/usr/bin/env python3
+"""T6: FULL-cudagraph replay checks for the two decode kernels whose plan is
+built host-side: the SM90 FlashInfer MLA wrapper (_SM90State) and the
+DeepGEMM paged MQA logits + per_row top-k of the sparse indexer.
+
+vLLM captures decode graphs with a dummy batch whose seq_len equals the query
+length (InputBatch.make_dummy: seq_lens = ceil(num_tokens/num_reqs), positions
+zero), then replans every step and replays. This mirrors that: capture with
+8 rows of lens 1..8, replay with real lens, compare against a reference.
+"""
+import sys
+import traceback
+
+import torch
+
+torch.manual_seed(0)
+dev = torch.device("cuda")
+results = {}
+
+
+def section(name):
+    def deco(fn):
+        def run():
+            print(f"\n===== {name} =====", flush=True)
+            try:
+                fn()
+                results[name] = "PASS"
+            except Exception as e:  # noqa: BLE001
+                results[name] = f"FAIL: {e}"
+                traceback.print_exc()
+        return run
+    return deco
+
+
+@section("T6a SM90 MLA wrapper: capture at dummy lens, replay at real lens")
+def t6a():
+    from vllm import _custom_ops as ops
+    from vllm.v1.attention.backends.mla.flashinfer_mla_sparse_sm90 import _SM90State
+
+    heads, d = 16, 512
+    sm_scale = 1.0 / 16.0
+    topk_width = 2176
+    max_tokens = 256
+    R = 8
+    N = 2300
+    bs = 64
+    num_blocks = N // bs + 3
+    kv_c = torch.randn(N, d, device=dev, dtype=torch.bfloat16)
+    k_pe = torch.empty(N, 0, device=dev, dtype=torch.bfloat16)
+    kv_cache = torch.zeros(num_blocks, bs, d, device=dev, dtype=torch.uint8)
+    slots = (torch.arange(N, device=dev) + 37).to(torch.int64)
+    ops.concat_and_cache_mla(kv_c, k_pe, kv_cache, slots, "fp8_e4m3", torch.tensor(1.0, device=dev))
+    deq = kv_cache.view(torch.float8_e4m3fn).reshape(-1, d).float()
+    flat = kv_cache.view(torch.float8_e4m3fn).reshape(-1, 1, d)
+    ckv, kpe = flat[..., :d], flat[..., d:]
+
+    state = _SM90State(dev, heads, torch.float8_e4m3fn, max_tokens, topk_width,
+                       kv_lora_rank=d, qk_rope_head_dim=0, sm_scale=sm_scale)
+    # static input buffers (what the captured graph reads)
+    q_nope = torch.randn(R, heads, d, device=dev, dtype=torch.bfloat16)
+    q_pe = torch.empty(R, heads, 0, device=dev, dtype=torch.bfloat16)
+    out_buf = torch.empty(R, heads, d, device=dev, dtype=torch.bfloat16)
+
+    def set_rows(lens):
+        state.kv_indices.fill_(0)
+        for r, L in enumerate(lens):
+            state.kv_indices[r * topk_width: r * topk_width + L] = slots[:L].to(torch.int32)
+        state.plan(R, torch.tensor(lens, dtype=torch.int32))
+
+    def reference(lens):
+        ref = torch.empty(R, heads, d, device=dev)
+        for r, L in enumerate(lens):
+            K = deq[slots[:L]]
+            s = (q_nope[r].float() @ K.T) * sm_scale
+            ref[r] = torch.softmax(s, dim=-1) @ K
+        return ref
+
+    def run_once():
+        out = state.wrapper.run(q_nope, q_pe, ckv, kpe, ckv_scale=1.0, kpe_scale=1.0)
+        out_buf.copy_(out)
+
+    # --- capture with the dummy lens vLLM uses ---
+    cap_lens = list(range(1, R + 1))
+    set_rows(cap_lens)
+    s = torch.cuda.Stream()
+    s.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(s):
+        for _ in range(2):
+            run_once()
+    torch.cuda.current_stream().wait_stream(s)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        run_once()
+    torch.cuda.synchronize()
+
+    worst = 0.0
+    for lens in (
+        cap_lens,
+        list(range(100, 108)),
+        list(range(693, 701)),
+        list(range(400, 408)),
+        [1, 5, 63, 64, 65, 2047, 2048, 2051],
+        list(range(2044, 2052)),
+    ):
+        q_nope.copy_(torch.randn_like(q_nope))
+        set_rows(lens)
+        graph.replay()
+        torch.cuda.synchronize()
+        ref = reference(lens)
+        err = (out_buf.float() - ref).abs()
+        rel = err.max().item() / max(ref.abs().max().item(), 1e-6)
+        row_err = [round(v, 4) for v in err.flatten(1).max(dim=1).values.tolist()]
+        nan = torch.isnan(out_buf).any().item()
+        print(f"  replay lens={lens[0]}..{lens[-1]}: rel err {rel:.4g} nan={nan} per-row {row_err}")
+        worst = max(worst, rel)
+        assert not nan
+    # eager run at the same lens for comparison
+    set_rows(list(range(693, 701)))
+    run_once()
+    torch.cuda.synchronize()
+    ref = reference(list(range(693, 701)))
+    print(f"  eager lens=693..700: rel err {((out_buf.float() - ref).abs().max() / ref.abs().max()).item():.4g}")
+    print(f"T6a worst replay rel err {worst:.4g}")
+    assert worst < 0.1
+
+
+@section("T6b indexer paged MQA logits + per_row topk: capture short, replay long")
+def t6b():
+    from vllm.utils.deep_gemm import fp8_fp4_paged_mqa_logits, get_paged_mqa_logits_metadata
+    from vllm.models.glm5next.nvidia.ops import kpool_compress as kpool_ops
+    from vllm import _custom_ops as ops
+    from vllm.utils.platform_utils import num_compute_units
+
+    # Indexer cache: [num_blocks, pool_block_size, 1, head_dim+4] uint8 (fp8 + fp32 scale)
+    head_dim = 128
+    kpool = 4
+    index_heads = 32  # index_n_heads
+    pool_bs = 64  # production indexer kernel page: 64 pools = 256 tokens
+    num_blocks = 80
+    max_pool_len = 524288 // kpool
+    fp8 = torch.float8_e4m3fn
+    k_pool = torch.randn(num_blocks * pool_bs, head_dim, device=dev).clamp(-4, 4)
+    kv_cache = torch.zeros(num_blocks, pool_bs, 1, head_dim + 4, device=dev, dtype=torch.uint8)
+    kv_cache[..., :head_dim] = k_pool.to(fp8).view(torch.uint8).reshape(num_blocks, pool_bs, 1, head_dim)
+    kv_cache[..., head_dim:] = torch.ones(num_blocks * pool_bs, 1, 1, device=dev).view(torch.uint8).reshape(num_blocks, pool_bs, 1, 4)
+    num_sms = num_compute_units(dev.index if dev.index is not None else 0)
+
+    R = 8
+    # per-token expanded block table (flattened path): every row -> request 0's blocks [1..num_blocks-1]
+    block_table = torch.zeros(R, num_blocks, device=dev, dtype=torch.int32)
+    block_table[:, : num_blocks - 1] = torch.arange(1, num_blocks, device=dev, dtype=torch.int32)
+    q = torch.randn(R, 1, index_heads, head_dim, device=dev).to(fp8)
+    weights = torch.rand(R, index_heads, device=dev, dtype=torch.float32)
+    seq_lens = torch.zeros(R, 1, device=dev, dtype=torch.int32)  # pool-granular, refreshed
+    sched_buf = torch.empty((num_sms + 1, 2), dtype=torch.int32, device=dev)
+    select_k = 2048 // kpool
+    pool_topk = torch.full((R, select_k), -1, dtype=torch.int32, device=dev)
+    logits_holder = {}
+
+    def set_lens(pool_lens):
+        seq_lens[:, 0] = torch.tensor(pool_lens, device=dev, dtype=torch.int32)
+        meta = get_paged_mqa_logits_metadata(seq_lens, pool_bs, num_sms)
+        sched = sched_buf[: meta.shape[0]]
+        sched[:] = meta
+        return sched
+
+    def run_once(sched):
+        logits = fp8_fp4_paged_mqa_logits((q, None), kv_cache, weights, seq_lens, block_table,
+                                          sched, max_model_len=max_pool_len, clean_logits=False)
+        logits_holder["l"] = logits
+        pool_topk.fill_(-1)
+        ops.top_k_per_row_decode(logits, 1, seq_lens, pool_topk, R, logits.stride(0), logits.stride(1), select_k)
+
+    def reference(pool_lens):
+        # logits[r, p] = sum_h w[r,h] * relu(q[r,h] . k[p]) (DeepSeek indexer semantics)
+        out = []
+        for r, L in enumerate(pool_lens):
+            blocks = block_table[r, : (L + pool_bs - 1) // pool_bs]
+            K = k_pool.view(num_blocks, pool_bs, head_dim)[blocks.long()].reshape(-1, head_dim)[:L]
+            Kq = K.to(fp8).float()
+            s = torch.relu(q[r, 0].float() @ Kq.T)  # [heads, L]
+            out.append((weights[r][:, None] * s).sum(0))
+        return out
+
+    cap_lens = [max(1, (i + 1) // kpool) for i in range(1, R + 1)]  # dummy: seq_len = query_len
+    sched = set_lens(cap_lens)
+    s = torch.cuda.Stream()
+    s.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(s):
+        for _ in range(2):
+            run_once(sched)
+    torch.cuda.current_stream().wait_stream(s)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        run_once(sched)
+    torch.cuda.synchronize()
+
+    bad = 0
+    for tok_lens in (list(range(100, 108)), list(range(693, 701)), list(range(1500, 1508)), list(range(2600, 2608)), list(range(4600, 4608))):
+        pool_lens = [t // kpool for t in tok_lens]
+        # the schedule buffer is refreshed in place; its slot count must be stable
+        meta = get_paged_mqa_logits_metadata(seq_lens.new_tensor([[p] for p in pool_lens]), pool_bs, num_sms)
+        assert meta.shape[0] == sched.shape[0], f"schedule slot count changed {meta.shape} vs {sched.shape}"
+        set_lens(pool_lens)
+        graph.replay()
+        torch.cuda.synchronize()
+        logits = logits_holder["l"]
+        refs = reference(pool_lens)
+        for r, L in enumerate(pool_lens):
+            got = logits[r, :L].float()
+            ref = refs[r]
+            nan = torch.isnan(got).any().item()
+            err = (got - ref).abs().max().item() / max(ref.abs().max().item(), 1e-6)
+            sel = pool_topk[r][pool_topk[r] >= 0].tolist()
+            if L <= select_k:
+                ok_sel = set(sel) == set(range(L)) and len(sel) == L
+            else:
+                ok_sel = set(sel) == set(ref.topk(select_k).indices.tolist())
+            if nan or err > 0.05 or not ok_sel:
+                bad += 1
+                miss = sorted(set(range(min(L, select_k))) - set(sel))[:8] if L <= select_k else []
+                print(f"  BAD replay tok_len={tok_lens[r]} pools={L}: nan={nan} rel_err={err:.3g} "
+                      f"sel_ok={ok_sel} nsel={len(sel)} missing_pools={miss}")
+        print(f"  replay tok_lens {tok_lens[0]}..{tok_lens[-1]}: checked")
+    # eager comparison
+    pool_lens = [t // kpool for t in range(693, 701)]
+    sched = set_lens(pool_lens)
+    run_once(sched)
+    torch.cuda.synchronize()
+    refs = reference(pool_lens)
+    e = max(((logits_holder["l"][r, :L].float() - refs[r]).abs().max() / refs[r].abs().max()).item() for r, L in enumerate(pool_lens))
+    print(f"  eager 693..700 rel err {e:.3g}")
+    print(f"T6b bad rows: {bad}")
+    assert bad == 0
+
+
+for fn in (t6a, t6b):
+    fn()
+print("\n===== SUMMARY =====")
+for k, v in results.items():
+    print(f"{k}: {v}")
+sys.exit(0 if all(v == "PASS" for v in results.values()) else 1)

docker-images/glm53-kindling/dev/kernel-tests/_cutlass_nvfp4_moe_realweights_test.py

@@ -0,0 +1,162 @@
+#!/usr/bin/env python3
+"""T10: FlashInfer CUTLASS NVFP4 fused MoE VALUES on real GLM-5.3-Flash experts.
+
+Loads 16 experts of layer 5 (extracted read-only from a checkpoint into
+/kv/moe-l5-<name>.safetensors), builds the kernel-format weights exactly as
+vLLM's FLASHINFER_CUTLASS path does (reorder [w1;w3]->[w3;w1], swizzled block
+scales, g1_alphas = weight_scale_2 * input_scale, a1_gscale = 1/input_scale),
+routes random tokens among those experts, and compares the kernel against:
+
+  ref_a4 : dequantised weights (vLLM's dequantize_to_dtype, linear scales),
+           activations quantised by vLLM's own scaled_fp4_quant with the
+           checkpoint's global input scales and dequantised -- what the
+           kernel is supposed to compute (w4a4)
+  ref_a16: dequantised weights, bf16 activations -- the w4a16 quantity
+           (what the marlin backend computes)
+
+Both checkpoints are supported: ModelOpt naming (weight/weight_scale/
+weight_scale_2/input_scale) and compressed-tensors naming (weight_packed/
+weight_scale/weight_global_scale/input_global_scale, reciprocal convention).
+usage: python3 this.py [nvidia|rh]
+"""
+import sys
+
+import torch
+from safetensors.torch import load_file
+
+from vllm import _custom_ops as ops
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
+    activation_to_flashinfer_type,
+)
+from vllm.model_executor.layers.quantization.utils.nvfp4_utils import swizzle_blockscale
+from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
+    reorder_w1w3_to_w3w1,
+)
+from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
+    dequantize_to_dtype,
+)
+from vllm.utils.flashinfer import flashinfer_cutlass_fused_moe
+
+name = sys.argv[1] if len(sys.argv) > 1 else "nvidia"
+dev = torch.device("cuda")
+torch.manual_seed(0)
+E_REAL, E, K, N, TOPK = 16, 288, 4096, 2048, 8
+LIMIT = 10.0
+fp8 = torch.float8_e4m3fn
+sd = load_file(f"/kv/moe-l5-{name}.safetensors", device="cuda")
+pre = next(k for k in sd if ".experts.0." in k).split(".experts.0.")[0] + ".experts."
+
+
+def get(e, proj, kind):
+    """kind in weight, scale, gscale (weight_scale_2 convention), iscale (input_scale convention)."""
+    base = f"{pre}{e}.{proj}."
+    if kind == "weight":
+        return sd.get(base + "weight", sd.get(base + "weight_packed"))
+    if kind == "scale":
+        return sd[base + "weight_scale"]
+    if kind == "gscale":
+        return sd[base + "weight_scale_2"].reshape(()) if base + "weight_scale_2" in sd else 1.0 / sd[base + "weight_global_scale"].reshape(())
+    if kind == "iscale":
+        return sd[base + "input_scale"].reshape(()) if base + "input_scale" in sd else 1.0 / sd[base + "input_global_scale"].reshape(())
+    raise KeyError(kind)
+
+
+# checkpoint-layout tensors, padded to E experts (zeros beyond the real ones)
+w13 = torch.zeros(E, 2 * N, K // 2, dtype=torch.uint8, device=dev)
+w13_scale = torch.ones(E, 2 * N, K // 16, device=dev).to(fp8)
+w2 = torch.zeros(E, K, N // 2, dtype=torch.uint8, device=dev)
+w2_scale = torch.ones(E, K, N // 16, device=dev).to(fp8)
+w13_gs = torch.ones(E, device=dev)
+w2_gs = torch.ones(E, device=dev)
+a13 = torch.zeros(E_REAL, device=dev)
+a2 = torch.zeros(E_REAL, device=dev)
+for e in range(E_REAL):
+    w13[e, :N] = get(e, "gate_proj", "weight"); w13[e, N:] = get(e, "up_proj", "weight")
+    w13_scale[e, :N] = get(e, "gate_proj", "scale"); w13_scale[e, N:] = get(e, "up_proj", "scale")
+    g1, g3 = float(get(e, "gate_proj", "gscale")), float(get(e, "up_proj", "gscale"))
+    assert abs(g1 - g3) <= 1e-9 * max(abs(g1), 1e-30), f"gate/up weight_scale_2 differ for expert {e}: {g1} vs {g3}"
+    w13_gs[e] = g1
+    w2[e] = get(e, "down_proj", "weight"); w2_scale[e] = get(e, "down_proj", "scale"); w2_gs[e] = float(get(e, "down_proj", "gscale"))
+    a13[e] = max(float(get(e, "gate_proj", "iscale")), float(get(e, "up_proj", "iscale")))
+    a2[e] = float(get(e, "down_proj", "iscale"))
+print(f"[{name}] input_scale gate/up per expert: min {a13.min():.6g} max {a13.max():.6g}; down: min {a2.min():.6g} max {a2.max():.6g}")
+print(f"[{name}] weight_scale_2 w13: {w13_gs[:E_REAL].min():.4g}..{w13_gs[:E_REAL].max():.4g}; w2: {w2_gs[:E_REAL].min():.4g}..{w2_gs[:E_REAL].max():.4g}")
+# FLASHINFER_CUTLASS uses one global activation scale = max over experts (prepare_nvfp4_moe_layer_for_fi_or_cutlass)
+a13_g = a13.max().item(); a2_g = a2.max().item()
+w13_input_scale = torch.full((E,), a13_g, device=dev); w2_input_scale = torch.full((E,), a2_g, device=dev)
+
+# kernel format (prepare_nvfp4_moe_layer_for_fi_or_cutlass + FlashInferExperts.process_weights_after_loading)
+w13_k, w13_scale_k = reorder_w1w3_to_w3w1(w13.clone(), w13_scale.clone())
+w13_scale_k = swizzle_blockscale(w13_scale_k); w2_scale_k = swizzle_blockscale(w2_scale.clone())
+g1_alphas = (w13_gs * w13_input_scale).contiguous(); g2_alphas = (w2_gs * w2_input_scale).contiguous()
+a1_gscale = (1.0 / w13_input_scale).contiguous(); a2_gscale = (1.0 / w2_input_scale).contiguous()
+quant_scales = [a1_gscale, w13_scale_k.view(torch.int32), g1_alphas, a2_gscale, w2_scale_k.view(torch.int32), g2_alphas]
+swiglu_limit = torch.full((E,), LIMIT, device=dev)
+act = activation_to_flashinfer_type(MoEActivation.SILU)
+
+
+def kernel(x, ids, w):
+    M = x.shape[0]
+    out = torch.empty(M, K, device=dev, dtype=torch.bfloat16)
+    xq, xsf = ops.scaled_fp4_quant(x, a1_gscale[:1].reshape(()).clone(), is_sf_swizzled_layout=True)
+    flashinfer_cutlass_fused_moe(
+        input=xq, token_selected_experts=ids.to(torch.int), token_final_scales=w,
+        fc1_expert_weights=w13_k.view(torch.long), fc2_expert_weights=w2.view(torch.long),
+        fc1_expert_biases=None, fc2_expert_biases=None, swiglu_alpha=None, swiglu_beta=None,
+        swiglu_limit=swiglu_limit, output=out, output_dtype=torch.bfloat16, quant_scales=quant_scales,
+        input_sf=xsf, tp_size=1, tp_rank=0, ep_size=1, ep_rank=0, activation_type=act,
+        use_deepseek_fp8_block_scale=False, use_mxfp8_act_scaling=False, use_w4_group_scaling=False)
+    return out, xq, xsf
+
+
+# dequantised weights for the real experts
+W13 = dequantize_to_dtype(w13[:E_REAL], w13_scale[:E_REAL], w13_gs[:E_REAL], torch.float32, 16, swizzle=False)  # [E, 2N, K]
+W2 = dequantize_to_dtype(w2[:E_REAL], w2_scale[:E_REAL], w2_gs[:E_REAL], torch.float32, 16, swizzle=False)  # [E, K, N]
+
+
+def act_q(h, gscale_inv):
+    """quantise like the kernel's a2 path: scaled_fp4_quant with global scale 1/input_scale, then dequantise."""
+    q, sf = ops.scaled_fp4_quant(h.to(torch.bfloat16), torch.tensor(1.0 / gscale_inv, device=dev), is_sf_swizzled_layout=True)
+    return dequantize_to_dtype(q, sf, torch.tensor(gscale_inv, device=dev), torch.float32, 16, swizzle=True)
+
+
+def reference(x, ids, w, xq, xsf, quant_act):
+    M = x.shape[0]
+    if quant_act:
+        xd = dequantize_to_dtype(xq, xsf, torch.tensor(a13_g, device=dev), torch.float32, 16, swizzle=True)
+    else:
+        xd = x.float()
+    out = torch.zeros(M, K, device=dev)
+    for r in range(M):
+        for j in range(TOPK):
+            e = int(ids[r, j])
+            gate = xd[r] @ W13[e, :N].T
+            up = xd[r] @ W13[e, N:].T
+            h = torch.nn.functional.silu(gate.clamp(max=LIMIT)) * up.clamp(-LIMIT, LIMIT)
+            hd = act_q(h[None], a2_g)[0] if quant_act else h
+            out[r] += w[r, j] * (hd @ W2[e].T)
+    return out
+
+
+gen = torch.Generator(device=dev).manual_seed(1)
+worst_a4 = worst_a16 = 0.0
+for M, amp in ((1, 0.5), (4, 0.5), (8, 0.5), (8, 1.0), (16, 0.5), (1, 2.0)):
+    x = (torch.randn(M, K, device=dev, generator=gen) * amp).to(torch.bfloat16)
+    ids = torch.stack([torch.randperm(E_REAL, device=dev, generator=gen)[:TOPK] for _ in range(M)]).to(torch.int32)
+    w = torch.rand(M, TOPK, device=dev, generator=gen) + 0.1
+    w = (w / w.sum(-1, keepdim=True)) * 2.5
+    out, xq, xsf = kernel(x, ids, w)
+    torch.cuda.synchronize()
+    ref4 = reference(x, ids, w, xq, xsf, True)
+    ref16 = reference(x, ids, w, xq, xsf, False)
+    e4 = (out.float() - ref4).abs().amax(1) / (ref4.abs().amax(1) + 1e-6)
+    e16 = (out.float() - ref16).abs().amax(1) / (ref16.abs().amax(1) + 1e-6)
+    q_loss = (ref4 - ref16).abs().amax(1) / (ref16.abs().amax(1) + 1e-6)
+    cos = torch.nn.functional.cosine_similarity(out.float(), ref16, dim=1)
+    print(f"[{name}] M={M} amp={amp}: kernel vs ref_a4 rel err per row {[round(v, 3) for v in e4.tolist()]}; "
+          f"kernel vs ref_a16 {[round(v, 3) for v in e16.tolist()]}; a4-vs-a16 quant loss {[round(v, 3) for v in q_loss.tolist()]}; "
+          f"cos(kernel, ref_a16) min {cos.min():.4f}; |out| mean {out.float().abs().mean():.3g} nan={torch.isnan(out).any().item()}")
+    worst_a4 = max(worst_a4, e4.max().item()); worst_a16 = max(worst_a16, e16.max().item())
+print(f"[{name}] T10 worst: kernel vs w4a4 reference {worst_a4:.3f}, kernel vs w4a16 reference {worst_a16:.3f}")
+print("T10:", "PASS" if worst_a4 < 0.1 else "FAIL")

docker-images/glm53-kindling/dev/kernel-tests/_cutlass_nvfp4_moe_rows_test.py

@@ -0,0 +1,143 @@
+#!/usr/bin/env python3
+"""T7: FlashInfer CUTLASS NVFP4 fused MoE on GB10 -- row consistency vs M.
+
+Mirrors FlashInferExperts.apply (vllm/model_executor/layers/fused_moe/experts/
+flashinfer_cutlass_moe.py) for the nvfp4 path at GLM-5.3-Flash's per-rank
+shape (E=288, K=4096, N=512, top-k 8, swiglu_limit 10): pre-quantized
+activation via ops.scaled_fp4_quant (swizzled sf), [w3;w1] packed weights,
+swizzle_blockscale'd block scales viewed as int32, per-expert global scales,
+output= preallocated.
+
+For each M, every batched row is compared with the same row run alone (M=1,
+same routing). A row that disagrees is checked against every other row's
+solo output to detect row mixing.
+"""
+import sys
+
+import torch
+
+from vllm import _custom_ops as ops
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
+    activation_to_flashinfer_type,
+)
+from vllm.model_executor.layers.quantization.utils.nvfp4_utils import swizzle_blockscale
+from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
+    reorder_w1w3_to_w3w1,
+)
+from vllm.utils.flashinfer import flashinfer_cutlass_fused_moe
+
+torch.manual_seed(0)
+dev = torch.device("cuda")
+E, K, N, TOPK = 288, 4096, 512, 8
+fp8 = torch.float8_e4m3fn
+
+# ---- synthetic weights in checkpoint layout, then kernel layout ----
+g = torch.Generator(device=dev).manual_seed(7)
+w13 = torch.randint(0, 256, (E, 2 * N, K // 2), device=dev, dtype=torch.uint8, generator=g)
+w2 = torch.randint(0, 256, (E, K, N // 2), device=dev, dtype=torch.uint8, generator=g)
+w13_scale = (torch.rand(E, 2 * N, K // 16, device=dev, generator=g) * 1.5 + 0.5).to(fp8)
+w2_scale = (torch.rand(E, K, N // 16, device=dev, generator=g) * 1.5 + 0.5).to(fp8)
+w13_scale_2 = torch.full((E,), 1.0 / 96, device=dev)  # weight global scale
+w2_scale_2 = torch.full((E,), 1.0 / 96, device=dev)
+w13_input_scale = torch.full((E,), 8.0 / (448 * 6), device=dev)  # amax 8 -> ModelOpt convention
+w2_input_scale = torch.full((E,), 8.0 / (448 * 6), device=dev)
+
+w13_k, w13_scale_k = reorder_w1w3_to_w3w1(w13.clone(), w13_scale.clone())
+w13_scale_k = swizzle_blockscale(w13_scale_k)
+w2_scale_k = swizzle_blockscale(w2_scale.clone())
+g1_alphas = (w13_scale_2 * w13_input_scale).contiguous()
+g2_alphas = (w2_scale_2 * w2_input_scale).contiguous()
+a1_gscale = (1.0 / w13_input_scale).contiguous()
+a2_gscale = (1.0 / w2_input_scale).contiguous()
+swiglu_limit = torch.full((E,), 10.0, device=dev, dtype=torch.float32)
+act = activation_to_flashinfer_type(MoEActivation.SILU)
+quant_scales = [a1_gscale, w13_scale_k.view(torch.int32), g1_alphas,
+                a2_gscale, w2_scale_k.view(torch.int32), g2_alphas]
+fc1 = w13_k.view(torch.long)
+fc2 = w2.view(torch.long)
+
+
+def moe(x_bf16, topk_ids, topk_w, prequant=True, tp_size=1, tp_rank=0):
+    M = x_bf16.shape[0]
+    out = torch.empty(M, K, device=dev, dtype=torch.bfloat16)
+    if prequant:
+        xq, xsf = ops.scaled_fp4_quant(x_bf16, a1_gscale[:1].reshape(()).clone(), is_sf_swizzled_layout=True)
+        inp, sf = xq, xsf
+    else:
+        inp, sf = x_bf16, None
+    flashinfer_cutlass_fused_moe(
+        input=inp,
+        token_selected_experts=topk_ids.to(torch.int),
+        token_final_scales=topk_w,
+        fc1_expert_weights=fc1,
+        fc2_expert_weights=fc2,
+        fc1_expert_biases=None,
+        fc2_expert_biases=None,
+        swiglu_alpha=None,
+        swiglu_beta=None,
+        swiglu_limit=swiglu_limit,
+        output=out,
+        output_dtype=torch.bfloat16,
+        quant_scales=quant_scales,
+        input_sf=sf,
+        tp_size=tp_size,
+        tp_rank=tp_rank,
+        ep_size=1,
+        ep_rank=0,
+        activation_type=act,
+        use_deepseek_fp8_block_scale=False,
+        use_mxfp8_act_scaling=False,
+        use_w4_group_scaling=False,
+    )
+    return out
+
+
+def routing(M, gen):
+    ids = torch.stack([torch.randperm(E, device=dev, generator=gen)[:TOPK] for _ in range(M)])
+    w = torch.rand(M, TOPK, device=dev, generator=gen) + 0.1
+    w = (w / w.sum(-1, keepdim=True)) * 2.5  # routed_scaling_factor
+    return ids.to(torch.int32), w.to(torch.float32)
+
+
+def check(prequant, tp_size=1, Ms=(1, 2, 4, 8, 16, 32, 64), trials=12, tag=""):
+    gen = torch.Generator(device=dev).manual_seed(11)
+    total_bad = 0
+    for M in Ms:
+        worst = 0.0
+        bad_rows = 0
+        mixes = []
+        nondet = 0.0
+        for t in range(trials):
+            x = torch.randn(M, K, device=dev, generator=gen).to(torch.bfloat16) * 1.5
+            ids, w = routing(M, gen)
+            outB = moe(x, ids, w, prequant, tp_size)
+            outB2 = moe(x, ids, w, prequant, tp_size)
+            nondet = max(nondet, (outB.float() - outB2.float()).abs().max().item())
+            solo = torch.stack([moe(x[r : r + 1], ids[r : r + 1], w[r : r + 1], prequant, tp_size)[0] for r in range(M)])
+            torch.cuda.synchronize()
+            for r in range(M):
+                scale = solo[r].float().abs().max().item() + 1e-6
+                err = (outB[r].float() - solo[r].float()).abs().max().item() / scale
+                worst = max(worst, err)
+                if err > 0.05:
+                    bad_rows += 1
+                    # does the batched row match some other row's solo output?
+                    d = (solo.float() - outB[r].float()[None]).abs().amax(dim=1) / scale
+                    j = int(d.argmin())
+                    mixes.append((t, r, round(err, 3), j, round(float(d[j]), 3)))
+        total_bad += bad_rows
+        print(f"  [{tag}] M={M:3d}: worst row err vs solo {worst:.4f}, bad rows {bad_rows}/{M * trials}, "
+              f"repeat-call nondeterminism max|d| {nondet:.4g}" + (f", mixes(trial,row,err,best_row,best_err)={mixes[:6]}" if mixes else ""))
+    return total_bad
+
+
+bad = 0
+print("== production path: pre-quantized nvfp4 input + swizzled input_sf, tp_size=1 ==")
+bad += check(True, 1, tag="prequant")
+print("== bf16 input (kernel-internal quant), tp_size=1 ==")
+bad += check(False, 1, tag="bf16in")
+print("== production tp flags: prequant, tp_size=4 tp_rank=0 ==")
+bad += check(True, 4, Ms=(1, 4, 8, 16), trials=8, tag="prequant-tp4")
+print("T7:", "PASS" if bad == 0 else f"FAIL ({bad} bad rows)")
+sys.exit(0 if bad == 0 else 1)

docker-images/glm53-kindling/dev/kernel-tests/_cutlass_nvfp4_moe_sanity_test.py

@@ -0,0 +1,78 @@
+#!/usr/bin/env python3
+"""T7b: sanity for T7 -- are the CUTLASS NVFP4 MoE outputs real (finite,
+non-zero, input-dependent) and do they agree with a torch dequantised
+reference? Same synthetic layer as kt7.py."""
+import sys
+
+import torch
+
+sys.argv = [sys.argv[0]]
+exec(open("/kv/kt7.py").read().split("def routing(")[0])  # reuse weights + moe()
+
+FP4_LUT = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6], device=dev)
+
+
+def deq_fp4(packed: torch.Tensor, scale_fp8: torch.Tensor, gscale: torch.Tensor) -> torch.Tensor:
+    """packed uint8 [..., K/2], scale [..., K/16] fp8, gscale scalar -> float [..., K]."""
+    lo = packed & 0xF
+    hi = packed >> 4
+    nib = torch.stack([lo, hi], dim=-1).reshape(*packed.shape[:-1], packed.shape[-1] * 2)
+    mag = FP4_LUT[(nib & 7).long()]
+    sign = torch.where((nib & 8) != 0, -1.0, 1.0)
+    vals = mag * sign
+    sc = scale_fp8.float().repeat_interleave(16, dim=-1)
+    return vals * sc * gscale
+
+
+def quant_deq_act(x: torch.Tensor, gscale_inv: float) -> torch.Tensor:
+    """Emulate nvfp4 activation quant: per-16 block fp8 scale, global scale."""
+    M, Kd = x.shape
+    xb = x.float().view(M, Kd // 16, 16)
+    amax = xb.abs().amax(-1, keepdim=True)
+    # block scale (fp8) = amax / 6 * global; ModelOpt: sf = amax/6 / input_scale
+    sf = (amax / 6.0 * gscale_inv).to(fp8).float()
+    sf = torch.where(sf == 0, torch.ones_like(sf), sf)
+    q = (xb / (sf / gscale_inv)).clamp(-6, 6)
+    # round to e2m1 grid
+    grid = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6], device=x.device)
+    idx = (q.abs().unsqueeze(-1) - grid).abs().argmin(-1)
+    qv = grid[idx] * torch.sign(q)
+    return (qv * (sf / gscale_inv)).view(M, Kd)
+
+
+gen = torch.Generator(device=dev).manual_seed(3)
+M = 8
+x = torch.randn(M, K, device=dev, generator=gen).to(torch.bfloat16) * 1.5
+ids = torch.stack([torch.randperm(E, device=dev, generator=gen)[:TOPK] for _ in range(M)]).to(torch.int32)
+w = torch.rand(M, TOPK, device=dev, generator=gen) + 0.1
+w = (w / w.sum(-1, keepdim=True)) * 2.5
+out = moe(x, ids, w, True, 1)
+torch.cuda.synchronize()
+print("out finite:", torch.isfinite(out).all().item(), "abs mean:", out.float().abs().mean().item(),
+      "abs max:", out.float().abs().max().item(), "zero rows:", int((out.float().abs().amax(1) == 0).sum()))
+out_b = moe(x, ids, w, False, 1)
+print("bf16-input variant vs prequant: max |d| =", (out.float() - out_b.float()).abs().max().item())
+x2 = x.clone(); x2[0] = torch.randn(K, device=dev, generator=gen).to(torch.bfloat16)
+out2 = moe(x2, ids, w, True, 1)
+print("changing row 0 input changes row 0 output:", (out2[0].float() - out[0].float()).abs().max().item() > 1e-3,
+      "and leaves other rows unchanged:", (out2[1:].float() - out[1:].float()).abs().max().item())
+
+# torch reference for a few rows
+a1_inv = float(1.0 / a1_gscale[0])  # = w13_input_scale
+xd = quant_deq_act(x, a1_inv)
+ref = torch.zeros(M, K, device=dev)
+for r in range(M):
+    for j in range(TOPK):
+        e = int(ids[r, j])
+        W13 = deq_fp4(w13[e], w13_scale[e], float(w13_scale_2[e]))  # [2N, K] rows: [w1(gate); w3(up)]
+        W2 = deq_fp4(w2[e], w2_scale[e], float(w2_scale_2[e]))  # [K, N]
+        gate = xd[r] @ W13[:N].T
+        up = xd[r] @ W13[N:].T
+        gate = gate.clamp(max=10.0)
+        up = up.clamp(-10.0, 10.0)
+        h = torch.nn.functional.silu(gate) * up
+        hd = quant_deq_act(h[None], float(w2_input_scale[e]))[0]
+        ref[r] += w[r, j] * (hd @ W2.T)
+err = (out.float() - ref).abs().amax(1) / (ref.abs().amax(1) + 1e-6)
+print("per-row rel err vs torch reference:", [round(v, 3) for v in err.tolist()])
+print("ref abs mean:", ref.abs().mean().item())

docker-images/glm53-kindling/dev/kernel-tests/_kda_numerics_test.py

@@ -0,0 +1,110 @@
+#!/usr/bin/env python3
+"""T9: GLM-5.3-Flash KDA numerics against an fp32 torch reference.
+
+Covers the three kernels the nightly uses for the recurrent state (per rank:
+16 heads x 128, safe gate, lower_bound -5): FlashKDA chunked prefill (the
+nightly's new default on SM12x), the Triton chunk_kda_with_fused_gate prefill
+(what the old image ran), and fused_recurrent_kda plain decode, with the
+decode state seeded from the prefill's final state. The reference replicates
+the recurrent kernel's math: l2norm(q,k), q*=K^-0.5, g = lb/(1+exp(-exp(A)(g1+b))),
+h*=exp(g), v-=h@k, v*=sigmoid(beta), h+=v k^T, o=h@q.
+"""
+import sys
+
+import torch
+
+import vllm._flashkda_C  # noqa: F401
+from vllm.models.glm5next.nvidia.ops.third_party.kda import (
+    chunk_kda_with_fused_gate,
+    fused_recurrent_kda,
+)
+
+dev = torch.device("cuda")
+H, K, V = 16, 128, 128
+LB = -5.0
+torch.manual_seed(5)
+
+
+def ref_run(q, k, v, g1, beta, a_log, dt_bias, h0):
+    """q,k,v,g1: [T,H,K] bf16; beta [T,H]; h0 [H,V,K] fp32 -> (o [T,H,V] fp32, h)"""
+    T = q.shape[0]
+    h = h0.clone().float()
+    o = torch.zeros(T, H, V, device=dev)
+    A = torch.exp(a_log.view(H).float())
+    for t in range(T):
+        qf, kf, vf = q[t].float(), k[t].float(), v[t].float()
+        qn = qf / torch.sqrt((qf * qf).sum(-1, keepdim=True) + 1e-6) * (K ** -0.5)
+        kn = kf / torch.sqrt((kf * kf).sum(-1, keepdim=True) + 1e-6)
+        gate = LB / (1.0 + torch.exp(-(A[:, None] * (g1[t].float() + dt_bias.float()))))  # [H,K]
+        h = h * torch.exp(gate)[:, None, :]
+        vp = vf - torch.einsum("hvk,hk->hv", h, kn)
+        vp = vp * torch.sigmoid(beta[t].float())[:, None]
+        h = h + vp[:, :, None] * kn[:, None, :]
+        o[t] = torch.einsum("hvk,hk->hv", h, qn)
+    return o, h
+
+
+def rel(a, b):
+    return ((a.float() - b.float()).abs().max() / (b.float().abs().max() + 1e-6)).item()
+
+
+a_log = torch.log(torch.rand(1, 1, H, 1, device=dev) * 8 + 1).float().contiguous()
+dt_bias = (torch.randn(H * K, device=dev) * 0.5).float().contiguous()
+ok = True
+for L in (20, 300, 2500):
+    T = L + 40
+    q = torch.randn(T, H, K, device=dev).to(torch.bfloat16)
+    k = torch.randn(T, H, K, device=dev).to(torch.bfloat16)
+    v = torch.randn(T, H, V, device=dev).to(torch.bfloat16)
+    g1 = torch.randn(T, H, K, device=dev).to(torch.bfloat16)
+    beta = torch.randn(T, H, device=dev).to(torch.bfloat16)
+    h0 = torch.zeros(H, V, K, device=dev)
+    o_ref, h_ref_pref = ref_run(q[:L], k[:L], v[:L], g1[:L], beta[:L], a_log, dt_bias.view(H, K), h0)
+
+    # --- FlashKDA prefill (nightly default) ---
+    ws = torch.ops._flashkda_C.get_workspace_size(T, H, 4)
+    workspace = torch.empty(ws, dtype=torch.uint8, device=dev)
+    out = torch.empty(1, L, H, V, device=dev, dtype=torch.bfloat16)
+    init = torch.zeros(1, H, V, K, device=dev)
+    final = torch.zeros(1, H, V, K, device=dev)
+    cu = torch.tensor([0, L], device=dev, dtype=torch.int32)
+    torch.ops._flashkda_C.fwd(q[:L][None].contiguous(), k[:L][None].contiguous(), v[:L][None].contiguous(),
+                              g1[:L][None].contiguous(), beta[:L][None].contiguous(), K ** -0.5, out, workspace,
+                              a_log.view(-1), dt_bias.view(-1, K), LB, init.contiguous(), final, cu, None, None)
+    torch.cuda.synchronize()
+    e_o = rel(out[0], o_ref); e_h = rel(final[0], h_ref_pref)
+    print(f"L={L} flashkda prefill: out rel err {e_o:.4g}, final state rel err {e_h:.4g}, "
+          f"nan={torch.isnan(out).any().item() or torch.isnan(final).any().item()}")
+    ok &= e_o < 0.05 and e_h < 0.05
+
+    # --- Triton chunk prefill (old image path) ---
+    init_t = torch.zeros(1, H, V, K, device=dev)
+    o_t, h_t = chunk_kda_with_fused_gate(
+        q=q[:L][None], k=k[:L][None], v=v[:L][None], raw_g=g1[:L][None],
+        beta=torch.sigmoid(beta[:L].float())[None], A_log=a_log, g_bias=dt_bias,
+        initial_state=init_t, output_final_state=True, use_qk_l2norm_in_kernel=True,
+        cu_seqlens=cu, safe_gate=True, lower_bound=LB)
+    torch.cuda.synchronize()
+    e_ot = rel(o_t[0], o_ref); e_ht = rel(h_t[0], h_ref_pref)
+    print(f"L={L} triton chunk prefill: out rel err {e_ot:.4g}, final state rel err {e_ht:.4g}")
+    ok &= e_ot < 0.05 and e_ht < 0.05
+
+    # --- plain decode from the flashkda state ---
+    pool = torch.zeros(4, H, V, K, device=dev)
+    pool[1] = final[0]
+    o_dec_ref, h_dec_ref = ref_run(q[L:], k[L:], v[L:], g1[L:], beta[L:], a_log, dt_bias.view(H, K), h_ref_pref)
+    worst = 0.0
+    for i in range(L, T):
+        o, _ = fused_recurrent_kda(q=q[i:i+1][None], k=k[i:i+1][None], v=v[i:i+1][None], g=g1[i:i+1][None],
+                                   beta=beta[i:i+1][None], initial_state=pool, use_qk_l2norm_in_kernel=True,
+                                   cu_seqlens=torch.tensor([0, 1], device=dev, dtype=torch.int32),
+                                   ssm_state_indices=torch.tensor([1], device=dev, dtype=torch.int32),
+                                   sigmoid_beta=True, a_log=a_log, g_bias=dt_bias, compute_gate=True, lower_bound=LB)
+        torch.cuda.synchronize()
+        worst = max(worst, rel(o[0, 0], o_dec_ref[i - L]))
+    e_hd = rel(pool[1], h_dec_ref)
+    print(f"L={L} plain decode x40: worst out rel err {worst:.4g}, state after decode rel err {e_hd:.4g}, "
+          f"|state| {pool[1].abs().mean().item():.4g}")
+    ok &= worst < 0.05 and e_hd < 0.05
+print("T9:", "PASS" if ok else "FAIL")
+sys.exit(0 if ok else 1)

docker-images/glm53-kindling/dev/kernel-tests/_kda_spec_vs_plain_test.py

@@ -0,0 +1,151 @@
+#!/usr/bin/env python3
+"""T2: GLM-5.3-Flash KDA spec-decode path vs the plain one-token-per-step path.
+
+Drives the exact kernels kda.py uses (causal_conv1d_update with
+num_accepted_tokens/query_start_loc, fused_recurrent_kda with 2D
+ssm_state_indices + num_accepted_tokens) through a simulated speculative
+decode with random acceptance, and compares every ACCEPTED row against the
+same token processed one at a time by the non-spec decode kernels (what a
+spec-off boot runs). GLM-5.3-Flash dims at TP=1: 64 heads x 128, conv width 4.
+"""
+import sys
+
+import torch
+
+from vllm.model_executor.layers.mamba.ops.causal_conv1d import causal_conv1d_update
+from vllm.models.glm5next.nvidia.ops.third_party.kda import fused_recurrent_kda
+
+torch.manual_seed(1234)
+dev = torch.device("cuda")
+H, D = 64, 128
+PROJ = H * D
+CONV_DIM = 3 * PROJ
+WIDTH = 4
+LOWER_BOUND = -5.0
+LINES = 64
+
+
+def torch_conv_ref(xs: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
+    """xs [T, dim] bf16, w [dim, WIDTH] fp32 -> silu(causal conv) [T, dim] fp32."""
+    T = xs.shape[0]
+    x = xs.float()
+    pad = torch.zeros(WIDTH - 1, x.shape[1], device=dev)
+    xp = torch.cat([pad, x], 0)
+    out = torch.zeros_like(x)
+    for j in range(WIDTH):
+        out += w[:, j][None, :] * xp[j : j + T]
+    return torch.nn.functional.silu(out)
+
+
+def run(k: int, T: int = 260, warm: int = 20, seed: int = 0, verbose: bool = True):
+    g = torch.Generator(device=dev).manual_seed(seed)
+    state_len = WIDTH - 1 + k
+    conv_w = (torch.randn(CONV_DIM, WIDTH, device=dev, generator=g) * 0.3).float()
+    a_log = torch.log(torch.rand(1, 1, H, 1, device=dev, generator=g) * 8 + 1).float().contiguous()
+    dt_bias = (torch.randn(PROJ, device=dev, generator=g) * 0.5).float().contiguous()
+
+    xs = torch.randn(T, CONV_DIM, device=dev, generator=g).to(torch.bfloat16)
+    g1s = torch.randn(T, H, D, device=dev, generator=g).to(torch.bfloat16)
+    betas = torch.randn(T, H, device=dev, generator=g).to(torch.bfloat16)
+    garbage = torch.randn(T, CONV_DIM, device=dev, generator=g).to(torch.bfloat16)
+    ggarb = torch.randn(T, H, D, device=dev, generator=g).to(torch.bfloat16)
+    bgarb = torch.randn(T, H, device=dev, generator=g).to(torch.bfloat16)
+
+    conv_pool = torch.zeros(LINES, CONV_DIM, state_len, device=dev, dtype=torch.bfloat16)
+    rec_pool = torch.zeros(LINES, H, D, D, device=dev, dtype=torch.float32)
+
+    ref_line = 1
+    spec_cols = torch.tensor([[2 + i for i in range(k + 1)]], device=dev, dtype=torch.int32)
+
+    def plain_step(line: int, t: int):
+        x = xs[t : t + 1].clone()
+        y = causal_conv1d_update(
+            x, conv_pool, conv_w, None, activation="silu",
+            conv_state_indices=torch.tensor([line], device=dev, dtype=torch.int32),
+        )
+        q, kk, v = y.split(PROJ, dim=-1)
+        o, _ = fused_recurrent_kda(
+            q=q.reshape(1, -1, H, D), k=kk.reshape(1, -1, H, D), v=v.reshape(1, -1, H, D),
+            g=g1s[t : t + 1].reshape(1, -1, H, D), beta=betas[t : t + 1].reshape(1, -1, H),
+            initial_state=rec_pool, use_qk_l2norm_in_kernel=True,
+            cu_seqlens=torch.tensor([0, 1], device=dev, dtype=torch.int32),
+            ssm_state_indices=torch.tensor([line], device=dev, dtype=torch.int32),
+            sigmoid_beta=True, a_log=a_log, g_bias=dt_bias, compute_gate=True,
+            lower_bound=LOWER_BOUND,
+        )
+        return y[0].float(), o[0, 0].float()
+
+    # Warm both from the same state using the plain path on the ref line, then copy.
+    for t in range(warm):
+        plain_step(ref_line, t)
+    conv_pool[int(spec_cols[0, 0])] = conv_pool[ref_line]
+    rec_pool[int(spec_cols[0, 0])] = rec_pool[ref_line]
+
+    # Reference outputs for every remaining token via the plain path.
+    o_ref = torch.zeros(T, H, D, device=dev)
+    y_ref = torch.zeros(T, CONV_DIM, device=dev)
+    for t in range(warm, T):
+        y_ref[t], o_ref[t] = plain_step(ref_line, t)
+    y_torch = torch_conv_ref(xs, conv_w)
+
+    # Spec path.
+    p = warm
+    na_prev = 1
+    step = 0
+    worst_o = worst_y = worst_yt = 0.0
+    hist = []
+    while p + k + 1 <= T:
+        na = int(torch.randint(1, k + 2, (1,), generator=g, device=dev))
+        xq = torch.cat([xs[p : p + na], garbage[p + na : p + k + 1]], 0).clone()
+        gq = torch.cat([g1s[p : p + na], ggarb[p + na : p + k + 1]], 0)
+        bq = torch.cat([betas[p : p + na], bgarb[p + na : p + k + 1]], 0)
+        assert xq.shape[0] == k + 1
+        nacc = torch.tensor([na_prev], device=dev, dtype=torch.int32)
+        qsl = torch.tensor([0, k + 1], device=dev, dtype=torch.int32)
+        y = causal_conv1d_update(
+            xq, conv_pool, conv_w, None, activation="silu",
+            conv_state_indices=spec_cols[:, 0], num_accepted_tokens=nacc,
+            query_start_loc=qsl, max_query_len=k + 1,
+        )
+        q, kk, v = y.split(PROJ, dim=-1)
+        o, _ = fused_recurrent_kda(
+            q=q.reshape(1, -1, H, D), k=kk.reshape(1, -1, H, D), v=v.reshape(1, -1, H, D),
+            g=gq.reshape(1, -1, H, D), beta=bq.reshape(1, -1, H),
+            initial_state=rec_pool, use_qk_l2norm_in_kernel=True,
+            cu_seqlens=qsl, ssm_state_indices=spec_cols, num_accepted_tokens=nacc,
+            sigmoid_beta=True, a_log=a_log, g_bias=dt_bias, compute_gate=True,
+            lower_bound=LOWER_BOUND,
+        )
+        torch.cuda.synchronize()
+        assert torch.isfinite(o[0, :na].float()).all(), f"spec recurrent output has NaN/Inf at step {step}"
+        assert torch.isfinite(o_ref[p : p + na]).all(), f"plain recurrent output has NaN/Inf at p={p}"
+        assert o_ref[p : p + na].abs().max().item() > 1e-3, f"plain recurrent output is ~zero at p={p}"
+        ey = (y[:na].float() - y_ref[p : p + na]).abs().max().item()
+        eyt = (y[:na].float() - y_torch[p : p + na]).abs().max().item()
+        eo = (o[0, :na].float() - o_ref[p : p + na]).abs().max().item()
+        oscale = o_ref[p : p + na].abs().max().item()
+        hist.append((step, p, na_prev, na, ey, eo, oscale))
+        worst_y = max(worst_y, ey)
+        worst_yt = max(worst_yt, eyt)
+        worst_o = max(worst_o, eo / max(oscale, 1e-6))
+        if verbose and (eo / max(oscale, 1e-6) > 0.02 or ey > 0.05):
+            per_row = [(o[0, i].float() - o_ref[p + i]).abs().max().item() for i in range(na)]
+            print(f"  k={k} step={step} p={p} na_prev={na_prev} na={na}: conv err {ey:.3g} "
+                  f"rec err {eo:.3g} (scale {oscale:.3g}) per-row {[round(e, 3) for e in per_row]}")
+        na_prev = na
+        p += na
+        step += 1
+    print(f"k={k}: recurrent out |mean| {o_ref[warm:].abs().mean().item():.4g} max {o_ref[warm:].abs().max().item():.4g}; "
+          f"conv out |mean| {y_ref[warm:].abs().mean().item():.4g}")
+    print(f"k={k}: {step} spec steps, {p - warm} accepted tokens; worst conv err vs plain {worst_y:.3g}, "
+          f"vs torch {worst_yt:.3g}; worst recurrent rel err {worst_o:.3g}")
+    return worst_y, worst_o
+
+
+ok = True
+for k in (3, 7):
+    for seed in (0, 1):
+        wy, wo = run(k, seed=seed)
+        ok &= wy < 0.05 and wo < 0.02
+print("T2:", "PASS" if ok else "FAIL")
+sys.exit(0 if ok else 1)

docker-images/glm53-kindling/dev/kernel-tests/_sm90_single_row_sweep_test.py

@@ -0,0 +1,94 @@
+#!/usr/bin/env python3
+"""T16: the SM90 sparse-MLA wrapper at ONE query row, every kv_len 1..MAXL.
+
+T5 only ran the wrapper at 8 query rows. Every decode step is a single row,
+and FlashInfer's MLA plan splits a lone row's KV range across CTAs and merges
+partial results, so the single-row case has its own code path. This sweeps
+kv_len 1..MAXL at R=1 (and R=2 pairs, R=93 prefill-like control) for H=16
+(TP=4) and H=32 (TP=2), fp8 and bf16 caches, with the plan sized both like the
+old tests (64 rows) and like production (max_num_batched_tokens rows: the
+padded rows keep the full-width lens), against an fp32 torch reference on the
+dequantised cache. Flags any row above REL_THR.
+
+    python3 _sm90_single_row_sweep_test.py [MAXL] [max_tokens_prod]
+"""
+import sys
+
+import torch
+
+from vllm import _custom_ops as ops
+from vllm.v1.attention.backends.mla.flashinfer_mla_sparse_sm90 import _SM90State
+
+dev = torch.device("cuda")
+torch.manual_seed(0)
+MAXL = int(sys.argv[1]) if len(sys.argv) > 1 else 400
+PROD_MAX_TOKENS = int(sys.argv[2]) if len(sys.argv) > 2 else 8192
+d = 512
+sm_scale = 1.0 / (256**0.5)
+topk_width = 2176
+bs = 64
+REL_THR = 2e-2
+
+
+def make_cache(N, fp8):
+    num_blocks = (N + 200) // bs + 2
+    kv_c = torch.randn(N, d, device=dev, dtype=torch.bfloat16)
+    k_pe = torch.empty(N, 0, device=dev, dtype=torch.bfloat16)
+    kv_cache = torch.zeros(num_blocks, bs, d, device=dev, dtype=torch.uint8 if fp8 else torch.bfloat16)
+    slots = (torch.arange(N, device=dev) + 37).to(torch.int64)
+    ops.concat_and_cache_mla(kv_c, k_pe, kv_cache, slots, "fp8_e4m3" if fp8 else "auto", torch.tensor(1.0, device=dev))
+    deq = (kv_cache.view(torch.float8_e4m3fn) if fp8 else kv_cache).reshape(-1, d).float()
+    return kv_cache, slots, deq
+
+
+def run(state, kv_cache, slots, deq, lens, heads, fp8, q_nope=None):
+    R = len(lens)
+    if q_nope is None:
+        q_nope = torch.randn(R, heads, d, device=dev, dtype=torch.bfloat16)
+    q_pe = torch.empty(R, heads, 0, device=dev, dtype=torch.bfloat16)
+    state.kv_indices.fill_(0)
+    for r, L in enumerate(lens):
+        state.kv_indices[r * topk_width : r * topk_width + L] = slots[:L].to(torch.int32)
+    state.plan(R, torch.tensor(lens, dtype=torch.int32))
+    flat = (kv_cache.view(torch.float8_e4m3fn) if fp8 else kv_cache).reshape(-1, 1, d)
+    kw = {"ckv_scale": 1.0, "kpe_scale": 1.0} if fp8 else {}
+    out = state.wrapper.run(q_nope, q_pe, flat[..., :d], flat[..., d:], **kw)
+    torch.cuda.synchronize()
+    rels = []
+    for r, L in enumerate(lens):
+        K = deq[slots[:L]]
+        p = torch.softmax((q_nope[r].float() @ K.T) * sm_scale, dim=-1)
+        ref = p @ K
+        o = out[r].float()
+        rels.append(float((o - ref).abs().max() / (ref.abs().max() + 1e-6)))
+    return rels
+
+
+worst_all = 0.0
+bad_total = 0
+for heads in (16, 32):
+    for fp8 in (True, False):
+        kv_cache, slots, deq = make_cache(MAXL + 100, fp8)
+        for max_tokens in (64, PROD_MAX_TOKENS):
+            state = _SM90State(dev, heads, torch.float8_e4m3fn if fp8 else torch.bfloat16, max_tokens, topk_width,
+                               kv_lora_rank=d, qk_rope_head_dim=0, sm_scale=sm_scale)
+            bad = []
+            worst = 0.0
+            for L in range(1, MAXL + 1):
+                r = run(state, kv_cache, slots, deq, [L], heads, fp8)[0]
+                worst = max(worst, r)
+                if r > REL_THR:
+                    bad.append((L, round(r, 4)))
+            # pairs (S2) and a prefill-like 93-row batch
+            pair_worst = 0.0
+            for L in (93, 176, 274, 302):
+                pair_worst = max(pair_worst, max(run(state, kv_cache, slots, deq, [L, L], heads, fp8)))
+            pre_rows = min(93, max_tokens)
+            pre = max(run(state, kv_cache, slots, deq, list(range(1, pre_rows + 1)), heads, fp8))
+            tag = f"H={heads} kv={'fp8' if fp8 else 'bf16'} max_tokens={max_tokens}"
+            print(f"{tag:36s} R=1 worst rel {worst:.4g} over L=1..{MAXL}, {len(bad)} rows > {REL_THR}: {bad[:12]}{' ...' if len(bad) > 12 else ''}")
+            print(f"{'':36s} R=2 pairs worst {pair_worst:.4g}; R={pre_rows} prefill-like worst {pre:.4g}")
+            worst_all = max(worst_all, worst, pair_worst, pre)
+            bad_total += len(bad)
+print(f"\nT16 worst rel {worst_all:.4g}; {bad_total} single-row cases above {REL_THR}")
+print("T16: PASS" if bad_total == 0 and worst_all < REL_THR else "T16: FAIL")

docker-images/glm53-kindling/dev/kernel-tests/_tail_compaction_test.py

@@ -0,0 +1,92 @@
+#!/usr/bin/env python3
+"""T8: does the SM90 path see the tail tokens?
+
+Builds indexer rows exactly as the GLM kpool decode path leaves them in
+topk_indices_buffer (width 2176): pool tokens at columns [0, 4*pool_len),
+-1 up to column 2048, the 0-3 tail tokens at columns 2048.., -1 after.
+Runs triton_convert_req_index_to_global_index(return_valid_counts=True) as
+FlashInferMLASparseSM90Impl.forward_mqa does and checks that the first `ctx`
+compacted entries (what _SM90State plans for) contain every causal slot,
+in particular the tail. Also runs it inside a CUDA graph.
+"""
+import sys
+
+import torch
+
+from vllm.v1.attention.backends.mla.sparse_utils import (
+    triton_convert_req_index_to_global_index,
+)
+
+dev = torch.device("cuda")
+WIDTH = 2176
+TOPK = 2048
+KPOOL = 4
+BLOCK = 2304
+
+
+def build_rows(positions, block_table_row):
+    R = len(positions)
+    rows = torch.full((R, WIDTH), -1, dtype=torch.int32, device=dev)
+    for r, pos in enumerate(positions):
+        seq = pos + 1
+        pool_len = seq // KPOOL
+        n_hist = min(pool_len, TOPK // KPOOL) * KPOOL
+        # pools are selected in an arbitrary order by top-k; use a permutation
+        pools = torch.randperm(pool_len, device=dev)[: TOPK // KPOOL]
+        toks = (pools[:, None] * KPOOL + torch.arange(KPOOL, device=dev)[None, :]).reshape(-1)
+        rows[r, : toks.numel()] = toks.to(torch.int32)
+        tail = torch.arange(pool_len * KPOOL, seq, device=dev, dtype=torch.int32)
+        rows[r, TOPK : TOPK + tail.numel()] = tail
+    return rows
+
+
+bad = 0
+for positions in ([99, 100, 101, 102, 103, 104, 105, 106], [600, 601, 602, 603], [2047, 2048, 2049, 2050, 2051, 2052, 2053, 2054], [5000, 5001, 5002, 5003]):
+    R = len(positions)
+    req_id = torch.zeros(R, dtype=torch.int32, device=dev)
+    block_table = torch.tensor([[7, 3, 11] + [0] * 5], dtype=torch.int32, device=dev)
+    rows = build_rows(positions, block_table[0])
+    for graph in (False, True):
+        out = torch.empty_like(rows)
+        counts = torch.empty(R, dtype=torch.int32, device=dev)
+        def run():
+            triton_convert_req_index_to_global_index(
+                req_id, block_table, rows, BLOCK_SIZE=BLOCK, NUM_TOPK_TOKENS=WIDTH,
+                return_valid_counts=True, out=out, valid_counts_out=counts)
+        if graph:
+            s = torch.cuda.Stream(); s.wait_stream(torch.cuda.current_stream())
+            with torch.cuda.stream(s):
+                run(); run()
+            torch.cuda.current_stream().wait_stream(s)
+            g = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(g):
+                run()
+            rows2 = build_rows(positions, block_table[0]); rows.copy_(rows2)
+            g.replay()
+        else:
+            run()
+        torch.cuda.synchronize()
+        for r, pos in enumerate(positions):
+            ctx = pos + 1
+            lens = ctx if ctx <= TOPK else TOPK + ctx % KPOOL
+            got = out[r, :lens].tolist()
+            # expected slots: block_table[tok // BLOCK]*BLOCK + tok % BLOCK for all causal toks
+            toks = torch.arange(ctx, device=dev)
+            if ctx > TOPK:
+                exp_toks = set((toks[-(ctx % KPOOL):] if ctx % KPOOL else toks[:0]).tolist())
+                need_tail_only = True
+            else:
+                exp_toks = set(toks.tolist()); need_tail_only = False
+            exp_slots = {int(block_table[0, t // BLOCK]) * BLOCK + t % BLOCK for t in exp_toks}
+            gotset = set(got)
+            missing = exp_slots - gotset
+            tail_slots = {int(block_table[0, t // BLOCK]) * BLOCK + t % BLOCK for t in range((ctx // KPOOL) * KPOOL, ctx)}
+            tail_missing = tail_slots - gotset
+            neg = sum(1 for v in got if v < 0)
+            ok = (not tail_missing) and (need_tail_only or not missing) and int(counts[r]) == lens and neg == 0
+            if not ok:
+                bad += 1
+            print(f"  graph={graph} pos={pos} ctx={ctx} planned_lens={lens} valid_count={int(counts[r])} "
+                  f"neg_in_prefix={neg} tail_missing={sorted(tail_missing)} missing={len(missing)} -> {'ok' if ok else 'BAD'}")
+print("T8:", "PASS" if bad == 0 else f"FAIL ({bad})")
+sys.exit(0 if bad == 0 else 1)

docker-images/glm53-kindling/dev/kernel-tests/_topk_mhcpost_fp8mla_test.py

@@ -0,0 +1,224 @@
+#!/usr/bin/env python3
+"""Kernel-level checks for the GLM-5.3-Flash decode path on GB10 (spark-glm53:v4).
+
+T1: sparse-indexer decode top-k (per_row) + expand_pools_and_append_tail on
+    short rows: must return exactly the causal set for every row.
+T3: mhc_post_tilelang must not write into any of its inputs, and must match
+    the torch reference.
+T5: FlashInfer SM90 MLA wrapper with an fp8 KV cache (planned as
+    float8_e4m3fn, as patches/sm90_fp8_kv_dtype.py does) against a bf16
+    reference, through vLLM's own _SM90State and concat_and_cache_mla.
+"""
+import sys
+import traceback
+
+import torch
+
+torch.manual_seed(0)
+dev = torch.device("cuda")
+results = {}
+
+
+def section(name):
+    def deco(fn):
+        def run():
+            print(f"\n===== {name} =====", flush=True)
+            try:
+                fn()
+                results[name] = "PASS"
+            except Exception as e:  # noqa: BLE001
+                results[name] = f"FAIL: {e}"
+                traceback.print_exc()
+        return run
+    return deco
+
+
+@section("T1 per_row topk + expand (short rows)")
+def t1():
+    from vllm import _custom_ops as ops
+    from vllm.models.glm5next.nvidia.ops import kpool_compress as kpool_ops
+    from vllm.model_executor.layers.indexer_topk import SparseIndexerTopk
+
+    max_pool_len = 524288 // 4
+    kpool = 4
+    select_k = 2048 // kpool
+    bad = 0
+    for garbage in ("neg", "pos", "nan", "inf"):
+        for trial in range(3):
+            rows = 8
+            base = int(torch.randint(20, 3000, (1,)))
+            lens_tok = torch.arange(base, base + rows, dtype=torch.int32)
+            pool_lens = lens_tok // kpool
+            fill = {"neg": -1e4, "pos": 1e4, "nan": float("nan"), "inf": float("inf")}[garbage]
+            logits = torch.full((rows, max_pool_len), fill, device=dev, dtype=torch.float32)
+            for r in range(rows):
+                n = int(pool_lens[r])
+                logits[r, :n] = torch.randn(n, device=dev)
+            seq_lens = pool_lens.to(torch.int32).to(dev).view(rows, 1)
+            out = torch.full((rows, select_k), -1, dtype=torch.int32, device=dev)
+            ops.top_k_per_row_decode(
+                logits, 1, seq_lens, out, rows, logits.stride(0), logits.stride(1), select_k
+            )
+            torch.cuda.synchronize()
+            # reference
+            for r in range(rows):
+                n = int(pool_lens[r])
+                got = out[r][out[r] >= 0].tolist()
+                if n <= select_k:
+                    exp = set(range(n))
+                else:
+                    exp = set(logits[r, :n].topk(select_k).indices.tolist())
+                if set(got) != exp or len(got) != len(set(got)):
+                    bad += 1
+                    print(f"  topk MISMATCH garbage={garbage} row={r} len_tok={int(lens_tok[r])} pools={n}: "
+                          f"got {len(got)} ({len(set(got))} uniq), missing={sorted(exp - set(got))[:10]}, "
+                          f"extra={sorted(set(got) - exp)[:10]}")
+            expanded = kpool_ops.expand_pools_and_append_tail(out.to(torch.int64), lens_tok.to(dev), kpool)
+            torch.cuda.synchronize()
+            for r in range(rows):
+                L = int(lens_tok[r])
+                got = expanded[r][expanded[r] >= 0].tolist()
+                if L <= 2048:
+                    exp = set(range(L))
+                    if set(got) != exp or len(got) != L:
+                        bad += 1
+                        print(f"  expand MISMATCH garbage={garbage} row={r} L={L}: got {len(got)} "
+                              f"missing={sorted(exp - set(got))[:10]} extra={sorted(set(got) - exp)[:10]}")
+                else:
+                    if len(got) != 2048 + (L % kpool) or max(got) >= L or len(set(got)) != len(got):
+                        bad += 1
+                        print(f"  expand MISMATCH(long) row={r} L={L}: n={len(got)} max={max(got)}")
+    # Also drive the module path production uses.
+    mod = SparseIndexerTopk("per_row")
+    rows = 16
+    lens_tok = torch.randint(5, 900, (rows,), dtype=torch.int32)
+    pool_lens = (lens_tok // kpool)
+    logits = torch.full((rows, max_pool_len), 1e4, device=dev)
+    for r in range(rows):
+        logits[r, : int(pool_lens[r])] = torch.randn(int(pool_lens[r]), device=dev)
+    out = torch.full((rows, select_k), -1, dtype=torch.int32, device=dev)
+    mod(logits, pool_lens.to(dev).view(rows, 1), 1, out, select_k, int(lens_tok.max()))
+    torch.cuda.synchronize()
+    for r in range(rows):
+        n = int(pool_lens[r])
+        got = out[r][out[r] >= 0].tolist()
+        if set(got) != set(range(n)) or len(got) != n:
+            bad += 1
+            print(f"  module MISMATCH row={r} pools={n}: got {len(got)}")
+    print(f"T1 mismatches: {bad}")
+    assert bad == 0
+
+
+@section("T3 mhc_post_tilelang purity + reference")
+def t3():
+    from vllm.model_executor.kernels.mhc.torch import mhc_post_torch
+
+    T, n, H = 24, 4, 4096
+    worst = 0.0
+    for pdtype in (torch.float32, torch.bfloat16):
+        x = torch.randn(T, H, device=dev, dtype=torch.bfloat16)
+        residual = torch.randn(T, n, H, device=dev, dtype=torch.bfloat16)
+        post = torch.randn(T, n, 1, device=dev, dtype=pdtype)
+        comb = torch.randn(T, n, n, device=dev, dtype=pdtype)
+        keep = [t.clone() for t in (x, residual, post, comb)]
+        try:
+            out = torch.ops.vllm.mhc_post_tilelang(x, residual, post, comb)
+        except Exception as e:  # noqa: BLE001
+            print(f"  post/comb dtype {pdtype}: op rejected ({type(e).__name__}: {str(e)[:120]})")
+            continue
+        torch.cuda.synchronize()
+        for name, a, b in zip(("x", "residual", "post", "comb"), (x, residual, post, comb), keep):
+            assert torch.equal(a, b), f"mhc_post_tilelang MUTATED input {name} (post dtype {pdtype})"
+        ref = mhc_post_torch(x, residual, post, comb)
+        err = (out.float() - ref.float()).abs().max().item()
+        scale = ref.float().abs().max().item()
+        worst = max(worst, err / max(scale, 1e-6))
+        print(f"  post/comb dtype {pdtype}: out {tuple(out.shape)} {out.dtype}, max abs err {err:.4g} (ref max {scale:.4g})")
+        assert not out.data_ptr() in {t.data_ptr() for t in (x, residual, post, comb)}
+    print(f"T3 worst relative error {worst:.3g}")
+    assert worst < 0.05
+
+
+@section("T5 SM90 fp8 MLA wrapper vs reference")
+def t5():
+    from vllm import _custom_ops as ops
+    from vllm.v1.attention.backends.mla.flashinfer_mla_sparse_sm90 import _SM90State
+
+    heads = 16  # 64 heads / TP4
+    d = 512
+    sm_scale = 1.0 / (256 ** 0.5)
+    topk_width = 2176  # (2048 + 3) rounded up to 128
+    bs = 64
+    max_tokens = 64
+
+    def run_case(kv_dtype_name, N, lens, q_mag=1.0, kv_mag=1.0, label=""):
+        fp8 = kv_dtype_name == "fp8_e4m3"
+        num_blocks = (N + 100) // bs + 2
+        kv_c = torch.randn(N, d, device=dev, dtype=torch.bfloat16) * kv_mag
+        k_pe = torch.empty(N, 0, device=dev, dtype=torch.bfloat16)
+        if fp8:
+            kv_cache = torch.zeros(num_blocks, bs, d, device=dev, dtype=torch.uint8)
+        else:
+            kv_cache = torch.zeros(num_blocks, bs, d, device=dev, dtype=torch.bfloat16)
+        slots = (torch.arange(N, device=dev) + 37).to(torch.int64)
+        scale = torch.tensor(1.0, device=dev, dtype=torch.float32)
+        ops.concat_and_cache_mla(kv_c, k_pe, kv_cache, slots, kv_dtype_name if fp8 else "auto", scale)
+        torch.cuda.synchronize()
+        if fp8:
+            deq = kv_cache.view(torch.float8_e4m3fn).reshape(-1, d).float()
+        else:
+            deq = kv_cache.reshape(-1, d).float()
+        # cache write check
+        werr = (deq[slots] - kv_c.float()).abs().max().item()
+        R = len(lens)
+        q_nope = torch.randn(R, heads, d, device=dev, dtype=torch.bfloat16) * q_mag
+        q_pe = torch.empty(R, heads, 0, device=dev, dtype=torch.bfloat16)
+        state = _SM90State(dev, heads, torch.float8_e4m3fn if fp8 else torch.bfloat16,
+                           max_tokens, topk_width, kv_lora_rank=d, qk_rope_head_dim=0, sm_scale=sm_scale)
+        state.kv_indices.fill_(0)
+        for r, L in enumerate(lens):
+            state.kv_indices[r * topk_width: r * topk_width + L] = slots[:L].to(torch.int32)
+        state.plan(R, torch.tensor(lens, dtype=torch.int32))
+        flat = (kv_cache.view(torch.float8_e4m3fn) if fp8 else kv_cache).reshape(-1, 1, d)
+        ckv, kpe = flat[..., :d], flat[..., d:]
+        kw = {"ckv_scale": 1.0, "kpe_scale": 1.0} if fp8 else {}
+        out = state.wrapper.run(q_nope, q_pe, ckv, kpe, **kw)
+        torch.cuda.synchronize()
+        # reference in fp32 on the dequantised cache
+        ref = torch.empty(R, heads, d, device=dev, dtype=torch.float32)
+        for r, L in enumerate(lens):
+            K = deq[slots[:L]]  # [L, d]
+            s = (q_nope[r].float() @ K.T) * sm_scale  # [heads, L]
+            p = torch.softmax(s, dim=-1)
+            ref[r] = p @ K
+        err = (out.float() - ref).abs()
+        rel = err.max().item() / max(ref.abs().max().item(), 1e-6)
+        # per-row error
+        row_err = err.flatten(1).max(dim=1).values.tolist()
+        print(f"  [{label}] kv={kv_dtype_name} N={N} lens={lens[:3]}..{lens[-1]} write_err={werr:.3g} "
+              f"max_abs_err={err.max().item():.4g} rel={rel:.4g} out_shape={tuple(out.shape)} {out.dtype}")
+        print(f"     per-row max err: {[round(v, 4) for v in row_err]}")
+        return rel
+
+    worst = 0.0
+    # 8-row spec-verify step at short context (the count task regime)
+    for N, base in ((110, 103), (300, 293), (700, 693)):
+        lens = list(range(base, base + 8))
+        worst = max(worst, run_case("fp8_e4m3", N, lens, label="short8"))
+        run_case("auto", N, lens, label="short8-bf16")
+    # single row, mixed lengths, and lengths crossing 2048 (topk + tail)
+    N = 2100
+    worst = max(worst, run_case("fp8_e4m3", N, [1, 5, 63, 64, 65, 2047, 2048, 2051], label="mixed"))
+    # larger magnitudes (fp8 saturation regime)
+    worst = max(worst, run_case("fp8_e4m3", 300, list(range(293, 301)), q_mag=4.0, kv_mag=8.0, label="big"))
+    print(f"T5 worst relative error (fp8): {worst:.4g}")
+    assert worst < 0.1
+
+
+for fn in (t1, t3, t5):
+    fn()
+
+print("\n===== SUMMARY =====")
+for k, v in results.items():
+    print(f"{k}: {v}")
+sys.exit(0 if all(v == "PASS" for v in results.values()) else 1)

docker-images/glm53-kindling/dev/kernel-tests/hc_ref.py

@@ -0,0 +1,123 @@
+#!/usr/bin/env python3
+"""Offline reference for the mHC ops in a DEEP glm53_step_tap dump.
+
+The dump holds every mHC op call's inputs (by keyword: x, residual,
+post_layer_mix, comb_res_mix) and outputs (residual, post, comb, layer input)
+per layer: `hc_pre` (layer 0), `hc_fused` (attn pre, fused with the previous
+layer's post), `hc_fused@1` (ffn pre) and `hc_post` (last layer). This replays
+each call with vLLM's own torch implementation (mhc_pre_torch / mhc_post_torch)
+on the checkpoint's hc_* parameters and reports the relative difference of
+every output. The CUDA op fuses the RMSNorm into the layer input, which the
+torch reference does not, so the norm is applied here before comparing out3.
+
+    python3 hc_ref.py --ckpt /ckpt /hit/vec-rank0-step3257.pt:92 /hit/vec-rank0-step3283.pt:92
+"""
+
+import argparse
+import json
+import os
+
+import torch
+from safetensors import safe_open
+
+import vllm.model_executor.kernels.mhc as mk
+
+dev = torch.device("cuda")
+PREFIX = "model.language_model.layers."
+
+
+class Ckpt:
+    def __init__(self, d):
+        self.d = d
+        cfg = json.load(open(os.path.join(d, "config.json")))
+        self.cfg = cfg.get("text_config", cfg)
+        self.map = json.load(open(os.path.join(d, "model.safetensors.index.json")))["weight_map"]
+        self._open = {}
+
+    def get(self, name):
+        f = self._open.get(self.map[name])
+        if f is None:
+            f = safe_open(os.path.join(self.d, self.map[name]), "pt", device="cuda")
+            self._open[self.map[name]] = f
+        return f.get_tensor(name)
+
+
+def rel(a, b):
+    a, b = a.float().reshape(-1), b.float().reshape(-1)
+    return float((a - b).norm() / (b.norm() + 1e-12))
+
+
+def rmsnorm(x, w, eps):
+    xf = x.float()
+    return (xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps) * w.float()).to(x.dtype)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--ckpt", required=True)
+    ap.add_argument("--thr", type=float, default=1e-2)
+    ap.add_argument("specs", nargs="+")
+    args = ap.parse_args()
+    ck = Ckpt(args.ckpt)
+    cfg = ck.cfg
+    n = int(cfg.get("mhc_num_residual_streams") or 4)
+    hid = int(cfg["hidden_size"])
+    rms_eps = float(cfg.get("rms_norm_eps", 1e-5))
+    hc_eps = float(cfg.get("hc_eps") or 1e-6)
+    post_mult = float(cfg.get("mhc_post_mult_value") or 2.0)
+    sink = int(cfg.get("mhc_sinkhorn_iterations") or 20)
+    for spec in args.specs:
+        path, _, pos = spec.rpartition(":")
+        b = torch.load(path, map_location="cpu")
+        rows = [i for i, r in enumerate(b["rows"]) if r["pos"] == int(pos)]
+        if not rows:
+            print(f"{spec}: no row at pos {pos}")
+            continue
+        i = rows[-1]
+        cps = b["cps"]
+        g = lambda k: cps[k][i].to(dev)
+        layers = sorted({int(k[1:].split(".")[0]) for k in cps if ".hc_" in k})
+        print(f"\n== {os.path.basename(path)} pos {pos} (rows in dump: {len(b['rows'])}) ==")
+        print(f"{'layer':6s} {'call':10s} {'residual':>10s} {'post':>10s} {'comb':>10s} {'x(normed)':>10s}")
+        worst = 0.0
+        for L in layers:
+            P = f"{PREFIX}{L}."
+            params = {
+                "attn": (ck.get(P + "hc_attn_fn").float(), ck.get(P + "hc_attn_scale").float(), ck.get(P + "hc_attn_base").float(), ck.get(P + "input_layernorm.weight")),
+                "ffn": (ck.get(P + "hc_ffn_fn").float(), ck.get(P + "hc_ffn_scale").float(), ck.get(P + "hc_ffn_base").float(), ck.get(P + "post_attention_layernorm.weight")),
+            }
+            for call, kind in (("hc_pre", "attn"), ("hc_fused", "attn"), ("hc_fused@1", "ffn")):
+                key = f"L{L}.{call}.out0"
+                if key not in cps:
+                    continue
+                fn, sc, base, nw = params[kind]
+                if call == "hc_pre":
+                    residual_cur = g(f"L{L}.hc_pre.in_residual").view(n, hid)
+                    res_out = None
+                    post_d, comb_d, x_d = g(key), g(f"L{L}.{call}.out1"), g(f"L{L}.{call}.out2")
+                else:
+                    x = g(f"L{L}.{call}.in_x")
+                    residual = g(f"L{L}.{call}.in_residual").view(n, hid)
+                    post = g(f"L{L}.{call}.in_post_layer_mix").view(n, 1)
+                    comb = g(f"L{L}.{call}.in_comb_res_mix").view(n, n)
+                    residual_cur = mk.mhc_post_torch(x, residual, post, comb)
+                    res_out = g(key).view(n, hid)
+                    post_d, comb_d, x_d = g(f"L{L}.{call}.out1"), g(f"L{L}.{call}.out2"), g(f"L{L}.{call}.out3")
+                post_r, comb_r, xin_r = mk.mhc_pre_torch(residual_cur, fn, sc, base, rms_eps, hc_eps, hc_eps, post_mult, sink)
+                x_r = rmsnorm(xin_r.reshape(-1, hid), nw, rms_eps)
+                e = [rel(res_out, residual_cur) if res_out is not None else float("nan"), rel(post_d, post_r), rel(comb_d, comb_r), rel(x_d, x_r)]
+                worst = max(worst, *[v for v in e if v == v])
+                flag = "  <--" if any(v == v and v > args.thr for v in e) else ""
+                print(f"L{L:<5d} {call:10s} {e[0]:10.2e} {e[1]:10.2e} {e[2]:10.2e} {e[3]:10.2e}{flag}")
+            key = f"L{L}.hc_post.out0"
+            if key in cps:
+                x, residual = g(f"L{L}.hc_post.in0"), g(f"L{L}.hc_post.in1").view(n, hid)
+                post, comb = g(f"L{L}.hc_post.in2").view(n, 1), g(f"L{L}.hc_post.in3").view(n, n)
+                r = rel(g(key).view(n, hid), mk.mhc_post_torch(x, residual, post, comb))
+                worst = max(worst, r)
+                print(f"L{L:<5d} {'hc_post':10s} {r:10.2e}{'  <--' if r > args.thr else ''}")
+        print(f"worst relative difference: {worst:.3e} (threshold {args.thr})")
+
+
+if __name__ == "__main__":
+    main()

docker-images/glm53-kindling/dev/kernel-tests/hf_compare.py

@@ -0,0 +1,83 @@
+#!/usr/bin/env python3
+"""Compare HF's per-layer residual streams (hf_ref.py output) with vLLM tap dumps
+at one position, layer by layer.
+
+HF layer L output = the 4 residual streams after layer L's ffn post-mix. In
+vLLM's DEEP dump the same tensor is the NEXT layer's fused op output
+`L{L+1}.hc_fused.out0` (the deferred post applied), and for the last layer
+`L44.hc_post.out0`. The final normed hidden is compared by applying the norm
+weight to the stream mean of vLLM's last output.
+
+    python3 hf_compare.py --hf /out/hf_layers.pt --ckpt /ckpt --pos 338 /hit/vec-rank0-step464.pt /hit/vec-rank0-step470.pt
+"""
+
+import argparse
+import json
+import os
+
+import torch
+from safetensors import safe_open
+
+
+def rel(a, b):
+    a, b = a.float().reshape(-1), b.float().reshape(-1)
+    return float((a - b).norm() / (b.norm() + 1e-12))
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--hf", required=True)
+    ap.add_argument("--ckpt", required=True)
+    ap.add_argument("--pos", type=int, required=True)
+    ap.add_argument("dumps", nargs="+")
+    args = ap.parse_args()
+    hf = torch.load(args.hf, map_location="cpu")
+    kp = hf["keep_pos"]
+    if args.pos not in kp:
+        print("HF kept positions", kp)
+        return 1
+    hi = kp.index(args.pos)
+    nlayers = max(hf["layers"]) + 1
+    idx = json.load(open(os.path.join(args.ckpt, "model.safetensors.index.json")))["weight_map"]
+    nm = "model.language_model.norm.weight"
+    normw = safe_open(os.path.join(args.ckpt, idx[nm]), "pt").get_tensor(nm).float()
+    cfg = json.load(open(os.path.join(args.ckpt, "config.json")))
+    tc = cfg.get("text_config", cfg)
+    eps = float(tc.get("rms_norm_eps", 1e-5))
+    dumps = []
+    for p in args.dumps:
+        b = torch.load(p, map_location="cpu")
+        rows = [i for i, r in enumerate(b["rows"]) if r["pos"] == args.pos]
+        if not rows:
+            print(f"{p}: no row at pos {args.pos}")
+            return 1
+        dumps.append((os.path.basename(p), b["cps"], rows[-1]))
+    print(f"pos {args.pos}: relative l2 difference of the residual streams vs HF, per layer")
+    print("layer  " + "  ".join(f"{n:>28s}" for n, _, _ in dumps) + "   (vLLM dump vs vLLM dump)")
+    for L in range(nlayers):
+        h = hf["layers"][L][hi].float()  # [4, 4096]
+        key = f"L{L + 1}.hc_fused.out0" if L + 1 < nlayers else f"L{L}.hc_post.out0"
+        vals, vecs = [], []
+        for n, cps, i in dumps:
+            if key not in cps:
+                vals.append(float("nan"))
+                vecs.append(None)
+                continue
+            v = cps[key][i].float().view(4, -1)
+            vals.append(rel(v, h))
+            vecs.append(v)
+        cross = rel(vecs[0], vecs[1]) if len(vecs) > 1 and vecs[0] is not None and vecs[1] is not None else float("nan")
+        print(f"L{L:<5d}" + "  ".join(f"{x:28.4f}" for x in vals) + f"   {cross:.4f}")
+    # final hidden: HF normed vs norm(mean of vLLM last streams)
+    hfin = hf["final_hidden"][hi].float()
+    print("final normed hidden vs HF:")
+    for n, cps, i in dumps:
+        key = f"L{nlayers - 1}.hc_post.out0"
+        if key in cps:
+            m = cps[key][i].float().view(4, -1).mean(0)
+            x = m * torch.rsqrt(m.pow(2).mean() + eps) * normw
+            print(f"  {n}: rel {rel(x, hfin):.4f}")
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

docker-images/glm53-kindling/dev/kernel-tests/hf_ref.py

@@ -0,0 +1,228 @@
+#!/usr/bin/env python3
+"""Independent ground truth for GLM-5.3-Flash: HF transformers' own glm5_next
+modeling code (torch-only paths: no fla, no causal_conv1d, no kernels hub, eager
+attention), fed the official FP8 checkpoint dequantised layer by layer.
+
+The text model is built on the meta device; each decoder layer's parameters are
+materialised on the GPU (FP8 e4m3 x 128x128 weight_scale_inv -> bf16) by a
+forward pre-hook right before the layer runs and released after it, so a
+339-token forward of the 320 GB checkpoint fits in GB10 memory. One forward
+gives every position at once.
+
+Outputs (in OUT dir):
+  hf_top5.json      top-5 (token, logprob) at the requested positions
+  hf_top5_all.json  top-5 at every position, plus the logprob of the token the
+                    vLLM run actually produced at each generated index
+  hf_layers.pt      per-layer residual streams [4, 4096] at the watched
+                    positions, the final normed hidden, and the input ids
+
+    python3 hf_ref.py --ckpt /ckpt --tapshot /hit/tapshot.json --out /out \
+        --gen 310 --watch-gen 74,142,148,176,310
+"""
+
+import argparse
+import gc
+import json
+import os
+import resource
+import time
+
+import torch
+import torch.nn.functional as F
+from safetensors import safe_open
+
+os.environ.setdefault("HF_HUB_OFFLINE", "1")
+os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
+
+dev = torch.device("cuda")
+BS = 128  # FP8 block size
+DTYPE = torch.bfloat16
+
+
+class Ckpt:
+    def __init__(self, d):
+        self.d = d
+        self.map = json.load(open(os.path.join(d, "model.safetensors.index.json")))["weight_map"]
+        self._open = {}
+
+    def raw(self, name):
+        f = self._open.get(self.map[name])
+        if f is None:
+            f = safe_open(os.path.join(self.d, self.map[name]), "pt", device="cpu")
+            self._open[self.map[name]] = f
+        return f.get_tensor(name)
+
+    def has(self, name):
+        return name in self.map
+
+    def close(self):
+        self._open.clear()
+
+    def load(self, name):
+        """bf16 tensor on the GPU; FP8 block-quantised weights are dequantised."""
+        t = self.raw(name)
+        if t.dtype == torch.float8_e4m3fn:
+            s = self.raw(name[: -len(".weight")] + ".weight_scale_inv").to(dev).float()
+            t = t.to(dev).float()
+            out, inn = t.shape
+            s = s.repeat_interleave(BS, 0)[:out].repeat_interleave(BS, 1)[:, :inn]
+            return (t * s).to(DTYPE)
+        return t.to(dev, DTYPE)
+
+
+def layer_tensor(ck, L, hf, E):
+    P = f"model.language_model.layers.{L}."
+    if hf.startswith("self_attn.forget_gate."):
+        return ck.load(P + "self_attn." + hf.split("forget_gate.", 1)[1])
+    if hf == "self_attn.conv1d.weight":
+        return torch.cat([ck.load(P + f"self_attn.{c}_conv1d.weight") for c in "qkv"], 0)
+    hc = {"attn_hc.fn": "hc_attn_fn", "attn_hc.base": "hc_attn_base", "attn_hc.scale": "hc_attn_scale",
+          "ffn_hc.fn": "hc_ffn_fn", "ffn_hc.base": "hc_ffn_base", "ffn_hc.scale": "hc_ffn_scale"}
+    if hf in hc:
+        return ck.load(P + hc[hf])
+    if hf == "mlp.experts.gate_up_proj":
+        return torch.stack([torch.cat([ck.load(P + f"mlp.experts.{e}.gate_proj.weight"),
+                                       ck.load(P + f"mlp.experts.{e}.up_proj.weight")], 0) for e in range(E)])
+    if hf == "mlp.experts.down_proj":
+        return torch.stack([ck.load(P + f"mlp.experts.{e}.down_proj.weight") for e in range(E)])
+    return ck.load(P + hf)
+
+
+def set_param(root, name, tensor, buffer=False):
+    mod = root
+    parts = name.split(".")
+    for p in parts[:-1]:
+        mod = getattr(mod, p)
+    if buffer:
+        mod._buffers[parts[-1]] = tensor
+    else:
+        mod._parameters[parts[-1]] = torch.nn.Parameter(tensor, requires_grad=False)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--ckpt", required=True)
+    ap.add_argument("--tapshot", required=True)
+    ap.add_argument("--out", required=True)
+    ap.add_argument("--gen", type=int, default=310, help="generated tokens to append to the prompt")
+    ap.add_argument("--watch-gen", default="74,142,148,176,310", help="generated indices whose predicting position is reported")
+    ap.add_argument("--dtype", default="bf16", choices=["bf16", "fp32"], help="model/weight dtype (fp32 = perturbation/truth run)")
+    ap.add_argument("--bias-dtype", default="model", choices=["model", "fp32"],
+                    help="dtype of the router e_score_correction_bias buffer: 'model' (bf16 when --dtype bf16, "
+                         "as from_pretrained does) or fp32 (what vLLM uses)")
+    args = ap.parse_args()
+    global DTYPE
+    DTYPE = torch.float32 if args.dtype == "fp32" else torch.bfloat16
+    os.makedirs(args.out, exist_ok=True)
+    from transformers import AutoConfig, AutoTokenizer
+    from transformers.models.glm5_next.modeling_glm5_next import Glm5NextTextModel
+
+    snap = json.load(open(args.tapshot))
+    prompt = list(snap["prompt"])
+    tok = AutoTokenizer.from_pretrained(args.ckpt)
+    gen_ids = tok("".join(snap["strs"][: args.gen]), add_special_tokens=False)["input_ids"]
+    assert len(gen_ids) == args.gen, f"re-tokenised {len(gen_ids)} ids, expected {args.gen}"
+    ids = prompt + gen_ids
+    T = len(ids)
+    P0 = len(prompt)
+    watch_gen = [int(x) for x in args.watch_gen.split(",") if x]
+    watch_pos = [P0 + g - 1 for g in watch_gen]  # the row whose logits predict generated index g
+    print(f"prompt {P0} + gen {args.gen} = {T} tokens; watch positions {dict(zip(watch_gen, watch_pos))}", flush=True)
+
+    cfg = AutoConfig.from_pretrained(args.ckpt)
+    tc = cfg.text_config
+    tc._attn_implementation = "eager"
+    E = int(tc.num_local_experts)
+    torch.set_default_dtype(DTYPE)
+    with torch.device("meta"):
+        model = Glm5NextTextModel(tc)
+    torch.set_default_dtype(torch.float32)
+    model.eval()
+    ck = Ckpt(args.ckpt)
+    nb = sum(1 for _ in model.buffers())
+    print(f"model built on meta: {len(list(model.parameters()))} params, {nb} buffers", flush=True)
+    for name, buf in model.named_buffers():
+        print("  buffer", name, tuple(buf.shape), buf.dtype)
+
+    # non-layer weights
+    set_param(model, "embed_tokens.weight", ck.load("model.language_model.embed_tokens.weight"))
+    set_param(model, "norm.weight", ck.load("model.language_model.norm.weight"))
+    lm_head = ck.load("lm_head.weight")
+
+    saved = {"ids": ids, "watch_pos": watch_pos, "layers": {}}
+    keep_pos = sorted(set(watch_pos + [T - 2, T - 1]))
+    timing = {}
+
+    def pre_hook(L):
+        def f(module, args_, kwargs_):
+            t0 = time.time()
+            for name, p in list(module.named_parameters()):
+                if p.device.type == "meta":
+                    set_param(module, name, layer_tensor(ck, L, name, E))
+            for name, b in list(module.named_buffers()):
+                if b.device.type == "meta":
+                    t = layer_tensor(ck, L, name, E)
+                    if args.bias_dtype == "fp32" and name.endswith("e_score_correction_bias"):
+                        t = ck.raw(f"model.language_model.layers.{L}.mlp.gate.e_score_correction_bias").to(dev).float()
+                    set_param(module, name, t, buffer=True)
+            torch.cuda.synchronize()
+            timing[L] = time.time() - t0
+        return f
+
+    def post_hook(L):
+        def f(module, args_, kwargs_, out):
+            hs = out[0] if isinstance(out, tuple) else out  # [1, T, hc, hidden]
+            saved["layers"][L] = hs[0, keep_pos].detach().to("cpu")
+            for name, p in list(module.named_parameters()):
+                set_param(module, name, torch.empty(p.shape, device="meta", dtype=p.dtype))
+            for name, b in list(module.named_buffers()):
+                set_param(module, name, torch.empty(b.shape, device="meta", dtype=b.dtype), buffer=True)
+            ck.close()
+            gc.collect()
+            torch.cuda.empty_cache()
+            rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1e6
+            print(f"layer {L:2d} done: load {timing.get(L, 0):.1f}s  cuda alloc {torch.cuda.memory_allocated() / 1e9:.1f} GB "
+                  f"reserved {torch.cuda.memory_reserved() / 1e9:.1f} GB  host maxrss {rss:.1f} GB", flush=True)
+        return f
+
+    for L, layer in enumerate(model.layers):
+        layer.register_forward_pre_hook(pre_hook(L), with_kwargs=True)
+        layer.register_forward_hook(post_hook(L), with_kwargs=True)
+
+    t0 = time.time()
+    with torch.no_grad():
+        out = model(input_ids=torch.tensor([ids], device=dev), use_cache=False)
+        h = out.last_hidden_state[0].float()  # [T, hidden], normed
+        logits = F.linear(h, lm_head.float())
+        logp = torch.log_softmax(logits, dim=-1)
+    print(f"forward done in {time.time() - t0:.0f}s", flush=True)
+    saved["final_hidden"] = h[keep_pos].to("cpu")
+    saved["keep_pos"] = keep_pos
+    torch.save(saved, os.path.join(args.out, "hf_layers.pt"))
+
+    top = torch.topk(logp, 5, dim=-1)
+    res = {}
+    for g, p in zip(watch_gen, watch_pos):
+        row = [(tok.decode([int(i)]), int(i), round(float(v), 4)) for v, i in zip(top.values[p], top.indices[p])]
+        res[str(g)] = {"pos": p, "top5": row, "margin": round(float(top.values[p, 0] - top.values[p, 1]), 4)}
+        print(f"gen {g} (pos {p}): " + "  ".join(f"{s!r} {v}" for s, i, v in row) + f"   margin {res[str(g)]['margin']}", flush=True)
+    json.dump(res, open(os.path.join(args.out, "hf_top5.json"), "w"), indent=1)
+
+    allrows = []
+    for g in range(1, args.gen + 1):
+        p = P0 + g - 1
+        actual = ids[p + 1] if p + 1 < T else None
+        row = {"gen": g, "pos": p,
+               "top5": [(tok.decode([int(i)]), round(float(v), 4)) for v, i in zip(top.values[p], top.indices[p])],
+               "actual": (tok.decode([actual]), round(float(logp[p, actual]), 4)) if actual is not None else None,
+               "margin": round(float(top.values[p, 0] - top.values[p, 1]), 4)}
+        allrows.append(row)
+    json.dump(allrows, open(os.path.join(args.out, "hf_top5_all.json"), "w"))
+    thin = [(r["gen"], r["margin"], r["top5"][0][0]) for r in allrows if r["margin"] < 2.0]
+    wrong = [(r["gen"], r["top5"][0][0], r["actual"]) for r in allrows if r["actual"] and r["top5"][0][0] != r["actual"][0]]
+    print(f"HF margins over {len(allrows)} generated positions: thin(<2 nats) {thin[:30]}", flush=True)
+    print(f"HF argmax != token vLLM produced at: {wrong[:30]}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

docker-images/glm53-kindling/dev/kernel-tests/hf_sublayer_ref.py

@@ -0,0 +1,110 @@
+#!/usr/bin/env python3
+"""HF's own sublayer modules on vLLM's dumped inputs: an implementation-
+independent check of the token-local sublayers (mHC hyper-connection, dense
+MLP, MoE with router) against what vLLM produced for the same rows.
+
+For each layer in the dump:
+  hc_pre   HF Glm5NextTextHyperConnection(layer params) on vLLM's residual
+           streams (`L<i>.hc_fused.out0`, i.e. after the deferred post) ->
+           (post, comb, collapsed); vs vLLM `hc_fused.out1/out2` and, after
+           HF's input_layernorm, `hc_fused.out3`.
+  hc_post  HF's combine `post*x + comb^T residual` on vLLM's fused-op inputs
+           vs vLLM `hc_fused.out0` (and the ffn-site call `hc_fused@1`).
+  mlp      HF Glm5NextTextMLP / Glm5NextTextMoE (dequantised bf16 weights)
+           on vLLM `mlp.in0` vs vLLM `mlp.out0`.
+
+    python3 hf_sublayer_ref.py --ckpt /ckpt --layers 0,1,2,3 /hit/vec-rank0-step470.pt
+"""
+
+import argparse
+import os
+import sys
+
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import hf_ref  # noqa: E402
+
+os.environ.setdefault("HF_HUB_OFFLINE", "1")
+os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
+dev = torch.device("cuda")
+
+
+def rel(a, b):
+    a, b = a.float().reshape(-1), b.float().reshape(-1)
+    return float((a - b).norm() / (b.norm() + 1e-12))
+
+
+def fill(module, ck, L, prefix, E):
+    for name, p in list(module.named_parameters()):
+        hf_ref.set_param(module, name, hf_ref.layer_tensor(ck, L, prefix + name, E))
+    for name, b in list(module.named_buffers()):
+        hf_ref.set_param(module, name, hf_ref.layer_tensor(ck, L, prefix + name, E), buffer=True)
+    return module
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--ckpt", required=True)
+    ap.add_argument("--layers", default="0,1,2,3")
+    ap.add_argument("dump")
+    args = ap.parse_args()
+    from transformers import AutoConfig
+    from transformers.models.glm5_next import modeling_glm5_next as M
+
+    tc = AutoConfig.from_pretrained(args.ckpt).text_config
+    E = int(tc.num_local_experts)
+    n, hid = int(tc.hc_mult), int(tc.hidden_size)
+    ck = hf_ref.Ckpt(args.ckpt)
+    b = torch.load(args.dump, map_location="cpu")
+    cps = b["cps"]
+    nrows = len(b["rows"])
+    print(f"{os.path.basename(args.dump)}: {nrows} rows, positions {b['rows'][0]['pos']}..{b['rows'][-1]['pos']}")
+    torch.set_default_dtype(torch.bfloat16)
+    for L in [int(x) for x in args.layers.split(",")]:
+        with torch.no_grad():
+            # --- hyper-connection, attention site: input = streams after the previous post ---
+            for site, key, hcname, normname in (("attn", f"L{L}.hc_fused", "attn_hc.", "input_layernorm."),
+                                                 ("ffn", f"L{L}.hc_fused@1", "ffn_hc.", "post_attention_layernorm.")):
+                if f"{key}.out0" not in cps:
+                    continue
+                with torch.device("meta"):
+                    hc = M.Glm5NextTextHyperConnection(tc)
+                    ln = M.Glm5NextTextRMSNorm(hid, tc.rms_norm_eps)
+                fill(hc, ck, L, hcname, E)
+                fill(ln, ck, L, normname, E)
+                # HF post-combine on vLLM's fused-op inputs vs vLLM out0 (residual streams)
+                x = cps[f"{key}.in_x"].to(dev)
+                res = cps[f"{key}.in_residual"].to(dev).view(nrows, n, hid)
+                post_in = cps[f"{key}.in_post_layer_mix"].to(dev).view(nrows, n)
+                comb_in = cps[f"{key}.in_comb_res_mix"].to(dev).view(nrows, n, n)
+                streams = post_in.to(x.dtype).unsqueeze(-1) * x.unsqueeze(-2) + torch.matmul(comb_in.to(x.dtype).transpose(-1, -2), res)
+                e_post = rel(streams, cps[f"{key}.out0"].to(dev).view(nrows, n, hid))
+                # HF hyper-connection on vLLM's resulting streams vs vLLM's post/comb/x
+                v_streams = cps[f"{key}.out0"].to(dev).view(1, nrows, n, hid)
+                post, comb, collapsed = hc(v_streams)
+                e_p = rel(post[0], cps[f"{key}.out1"].to(dev).view(nrows, n))
+                e_c = rel(comb[0], cps[f"{key}.out2"].to(dev).view(nrows, n, n))
+                e_x = rel(ln(collapsed[0]), cps[f"{key}.out3"].to(dev))
+                print(f"L{L:<3d} hc {site:4s}: post-combine {e_post:.2e} | pre: post {e_p:.2e} comb {e_c:.2e} x(normed) {e_x:.2e}")
+            # --- MLP / MoE on vLLM's mlp input ---
+            key = f"L{L}.mlp"
+            if f"{key}.in0" in cps:
+                dense = L < int(tc.first_k_dense_replace)
+                with torch.device("meta"):
+                    mlp = M.Glm5NextTextMLP(tc) if dense else M.Glm5NextTextMoE(tc)
+                fill(mlp, ck, L, "mlp.", E)
+                xin = cps[f"{key}.in0"].to(dev)
+                out = mlp(xin.view(1, nrows, hid) if not dense else xin)
+                out = out.view(nrows, hid)
+                e = rel(out, cps[f"{key}.out0"].to(dev))
+                per_row = [(r["pos"], round(rel(out[k], cps[f"{key}.out0"][k].to(dev)), 3)) for k, r in enumerate(b["rows"])]
+                worst = sorted(per_row, key=lambda t: -t[1])[:4]
+                print(f"L{L:<3d} mlp {'dense' if dense else 'moe  '}: vLLM out vs HF {e:.2e} (worst rows {worst})")
+                del mlp
+                torch.cuda.empty_cache()
+                ck.close()
+
+
+if __name__ == "__main__":
+    main()

docker-images/glm53-kindling/dev/kernel-tests/kda_ref.py

@@ -0,0 +1,153 @@
+#!/usr/bin/env python3
+"""Offline decode-step reference for every KDA layer in a glm53_step_tap dump.
+
+A decode dump holds, per KDA layer and tapped row, everything the layer's
+decode path consumed and produced: `kda.qkv` (merged q|k|v projection, before
+the short conv), `kda.g1`, `kda.beta` (raw), the request's `conv_before` and
+`rec_before` state blocks, and `kda.out`, `conv_after`, `rec_after`. This
+tool replays the step in fp32 from those inputs with the checkpoint's conv
+weights, A_log and dt_bias (sliced to the dump's TP rank), using the same
+math as T9's reference (l2-normed q/k, bounded sigmoid gate, sigmoid beta),
+and reports how far the dumped output and states are from it. Both conv
+state layouts (dim-first / state-first) and both recurrent state layouts
+([H,V,K] / [H,K,V]) are tried; the matching one is reported.
+
+    python3 kda_ref.py --ckpt /ckpt /hit/vec-rank0-step3257.pt:92 [more specs]
+"""
+
+import argparse
+import json
+import os
+import re
+
+import torch
+from safetensors import safe_open
+
+dev = torch.device("cuda")
+PREFIX = "model.language_model.layers."
+
+
+class Ckpt:
+    def __init__(self, d):
+        self.d = d
+        cfg = json.load(open(os.path.join(d, "config.json")))
+        self.cfg = cfg.get("text_config", cfg)
+        self.map = json.load(open(os.path.join(d, "model.safetensors.index.json")))["weight_map"]
+        self._open = {}
+
+    def get(self, name):
+        f = self._open.get(self.map[name])
+        if f is None:
+            f = safe_open(os.path.join(self.d, self.map[name]), "pt", device="cuda")
+            self._open[self.map[name]] = f
+        return f.get_tensor(name)
+
+    def has(self, name):
+        return name in self.map
+
+
+def rel(a, b):
+    a, b = a.float().reshape(-1), b.float().reshape(-1)
+    return f"l2 {float((a - b).norm() / (b.norm() + 1e-12)):.2e} max {float((a - b).abs().max() / (b.abs().max() + 1e-12)):.2e}"
+
+
+def step_ref(q, k, v, g1, beta, a_log, dt_bias, h, lb, D):
+    """One recurrent step. q,k,v,g1 [H,D]; beta [H]; h [H,V,K]; returns (o [H,V], h)."""
+    A = torch.exp(a_log.float()).view(-1)
+    qn = q / torch.sqrt((q * q).sum(-1, keepdim=True) + 1e-6) * (D**-0.5)
+    kn = k / torch.sqrt((k * k).sum(-1, keepdim=True) + 1e-6)
+    gate = lb / (1.0 + torch.exp(-(A[:, None] * (g1 + dt_bias))))
+    h = h * torch.exp(gate)[:, None, :]
+    vp = v - torch.einsum("hvk,hk->hv", h, kn)
+    vp = vp * torch.sigmoid(beta)[:, None]
+    h = h + vp[:, :, None] * kn[:, None, :]
+    return torch.einsum("hvk,hk->hv", h, qn), h
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--ckpt", required=True)
+    ap.add_argument("--layers", default="", help="comma list; default every KDA layer in the dump")
+    ap.add_argument("--state-vs", default="", help="a prefill dump (extended tap): compare this step's rec/conv_after "
+                    "with the state that run left after its last position (extra kda.state_after_req)")
+    ap.add_argument("specs", nargs="+", help="vec.pt:pos")
+    args = ap.parse_args()
+    ck = Ckpt(args.ckpt)
+    cfg = ck.cfg
+    lac = cfg.get("linear_attn_config") or {}
+    H_all = int(cfg.get("linear_num_heads") or lac["num_heads"])
+    D = int(cfg.get("linear_head_dim") or lac["head_dim"])
+    W = int(cfg.get("linear_conv_kernel_dim") or lac["short_conv_kernel_size"])
+    lb = float(cfg["linear_lower_bound"] if "linear_lower_bound" in cfg else lac["gate_lower_bound"])
+    for spec in args.specs:
+        path, _, pos = spec.rpartition(":")
+        m = re.search(r"rank(-?\d+)", os.path.basename(path))
+        rank = int(m.group(1)) if m else 0
+        b = torch.load(path, map_location="cpu")
+        rows = [i for i, r in enumerate(b["rows"]) if r["pos"] == int(pos)]
+        if not rows:
+            print(f"{spec}: no row at pos {pos}")
+            continue
+        i = rows[-1]
+        cps = b["cps"]
+        layers = [int(x) for x in args.layers.split(",") if x] or sorted(
+            int(k[1:].split(".")[0]) for k in cps if k.endswith(".kda.out")
+        )
+        other = torch.load(args.state_vs, map_location="cpu") if args.state_vs else None
+        print(f"\n== {os.path.basename(path)} pos {pos} rank {rank} (row {b['rows'][i]}) ==")
+        print(f"{'layer':6s} {'conv':22s} {'out vs ref':24s} {'rec_after vs ref':24s} layouts   |out|max")
+        for L in layers:
+            g = lambda n: cps[f"L{L}.kda.{n}"][i].to(dev)
+            qkv, g1, beta, out = g("qkv").float(), g("g1").float(), g("beta").float(), g("out").float()
+            cb, ca, rb, ra = g("conv_before").float(), g("conv_after").float(), g("rec_before").float(), g("rec_after").float()
+            H = beta.numel()
+            tp = H_all // H
+            P = H * D  # local projection size
+            assert qkv.numel() == 3 * P, (qkv.shape, P)
+            base = f"{PREFIX}{L}.self_attn."
+            if not ck.has(base + "q_conv1d.weight"):
+                base = f"{PREFIX}{L}.self_attn.linear_attn."
+            ws = []
+            for p in ("q", "k", "v"):
+                w = ck.get(f"{base}{p}_conv1d.weight").float().reshape(H_all * D, W)
+                ws.append(w[rank * P : (rank + 1) * P])
+            cw = torch.cat(ws, 0)  # [3P, W]
+            a_log = ck.get(base + "A_log").float().reshape(-1)[rank * H : (rank + 1) * H]
+            dt_bias = ck.get(base + "dt_bias").float().reshape(-1)[rank * P : (rank + 1) * P].view(H, D)
+            # --- short conv, both state layouts ---
+            best = None
+            for layout in ("DS", "SD"):
+                st = cb.view(3 * P, W - 1) if layout == "DS" else cb.view(W - 1, 3 * P).T
+                win = torch.cat([st, qkv[:, None]], 1)  # [3P, W]
+                y = torch.nn.functional.silu((win * cw).sum(1))
+                new_st = win[:, 1:]
+                ca_v = ca.view(3 * P, W - 1) if layout == "DS" else ca.view(W - 1, 3 * P).T
+                e = float((new_st - ca_v).abs().max())
+                if best is None or e < best[1]:
+                    best = (layout, e, y)
+            conv_layout, conv_err, y = best
+            q, k, v = y.split(P)
+            q, k, v, g1r = q.view(H, D), k.view(H, D), v.view(H, D), g1.view(H, D)
+            # --- recurrent step, both state layouts ---
+            res = None
+            for hl in ("HVK", "HKV"):
+                h0 = rb.view(H, D, D)
+                if hl == "HKV":
+                    h0 = h0.transpose(1, 2)
+                o, h1 = step_ref(q, k, v, g1r, beta, a_log, dt_bias, h0, lb, D)
+                h1_cmp = h1 if hl == "HVK" else h1.transpose(1, 2)
+                e_o = float((o.reshape(-1) - out).norm() / (out.norm() + 1e-12))
+                if res is None or e_o < res[0]:
+                    res = (e_o, hl, o, h1_cmp)
+            e_o, hl, o, h1 = res
+            sv = ""
+            if other is not None:
+                st = other.get("extra", {}).get(f"L{L}.kda.state_after_req")
+                if st:
+                    oc, orc = next(iter(st.values()))
+                    sv = f"  | vs {os.path.basename(args.state_vs)} state: conv {rel(ca, oc.reshape(-1).to(dev))} rec {rel(ra, orc.reshape(-1).to(dev))}"
+            print(f"L{L:<5d} {conv_layout} state maxdiff {conv_err:.1e}   {rel(out, o.reshape(-1)):24s} {rel(ra, h1.reshape(-1)):24s} {conv_layout}/{hl}  {float(out.abs().max()):.3f}{sv}")
+
+
+if __name__ == "__main__":
+    main()

docker-images/glm53-kindling/dev/kernel-tests/mla_ref.py

@@ -0,0 +1,159 @@
+#!/usr/bin/env python3
+"""Offline decode-step reference for the MLA layers in a glm53_step_tap dump.
+
+Two halves, per MLA layer and tapped row:
+
+  up-projection  `mla.out` (the SM90 kernel's latent output, [H, kv_lora_rank])
+                 times the checkpoint's W_UV (v half of kv_b_proj, per head)
+                 against the dumped `mla.attn_out`. Needs only today's dumps.
+  kernel         softmax(scale * q_nope . K_j) V_j over the latent rows the
+                 kernel attended (K_j = V_j = fp8 cache row * k_scale), against
+                 the dumped `mla.out`. Needs `mla.q_nope`, `mla.k_scale` and
+                 `mla.kv_rows` from the extended tap (extra dict in the blob).
+
+    python3 mla_ref.py --ckpt /ckpt /hit/vec-rank0-step3257.pt:92
+"""
+
+import argparse
+import json
+import os
+import re
+
+import torch
+from safetensors import safe_open
+
+from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
+    dequantize_to_dtype,
+)
+
+dev = torch.device("cuda")
+PREFIX = "model.language_model.layers."
+
+
+class Ckpt:
+    def __init__(self, d):
+        self.d = d
+        cfg = json.load(open(os.path.join(d, "config.json")))
+        self.cfg = cfg.get("text_config", cfg)
+        self.map = json.load(open(os.path.join(d, "model.safetensors.index.json")))["weight_map"]
+        self._open = {}
+
+    def get(self, name):
+        f = self._open.get(self.map[name])
+        if f is None:
+            f = safe_open(os.path.join(self.d, self.map[name]), "pt", device="cuda")
+            self._open[self.map[name]] = f
+        return f.get_tensor(name)
+
+    def has(self, name):
+        return name in self.map
+
+    def linear(self, base):
+        """fp32 [out, in] whether the layer is NVFP4 or BF16 in the checkpoint."""
+        if self.has(base + ".weight_scale"):
+            w = self.get(base + ".weight")
+            s = self.get(base + ".weight_scale")
+            g = self.get(base + ".weight_scale_2").reshape(()).float()
+            return dequantize_to_dtype(w[None], s[None], g[None], torch.float32, 16, swizzle=False)[0], "nvfp4"
+        return self.get(base + ".weight").float(), "bf16"
+
+
+def rel(a, b):
+    a, b = a.float().reshape(-1), b.float().reshape(-1)
+    return f"l2 {float((a - b).norm() / (b.norm() + 1e-12)):.2e} max {float((a - b).abs().max() / (b.abs().max() + 1e-12)):.2e}"
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--ckpt", required=True)
+    ap.add_argument("--layers", default="")
+    ap.add_argument("--kv-vs", default="", help="another dump: compare the rows this step READ (extra kv_rows) "
+                    "against the cache_row that dump WROTE at the same positions (e.g. the prefill replay)")
+    ap.add_argument("specs", nargs="+")
+    args = ap.parse_args()
+    ck = Ckpt(args.ckpt)
+    cfg = ck.cfg
+    H_all, nope, vd, lora = (int(cfg[k]) for k in ("num_attention_heads", "qk_nope_head_dim", "v_head_dim", "kv_lora_rank"))
+    rope = int(cfg.get("qk_rope_head_dim") or 0)
+    scale = (nope + rope) ** -0.5
+    for spec in args.specs:
+        path, _, pos = spec.rpartition(":")
+        m = re.search(r"rank(-?\d+)", os.path.basename(path))
+        rank = int(m.group(1)) if m else 0
+        b = torch.load(path, map_location="cpu")
+        rows = [i for i, r in enumerate(b["rows"]) if r["pos"] == int(pos)]
+        if not rows:
+            print(f"{spec}: no row at pos {pos}")
+            continue
+        i = rows[-1]
+        cps, extra = b["cps"], b.get("extra", {})
+        layers = [int(x) for x in args.layers.split(",") if x] or sorted(
+            int(k[1:].split(".")[0]) for k in cps if k.endswith(".mla.out")
+        )
+        print(f"\n== {os.path.basename(path)} pos {pos} rank {rank} ==")
+        print(f"{'layer':6s} {'kv_b':6s} {'up-proj: attn_out vs ref':28s} {'kernel: out vs ref':28s} rows  |out|max")
+        for L in layers:
+            out = cps[f"L{L}.mla.out"][i].to(dev).float()
+            H = out.numel() // lora
+            out = out.view(H, lora)
+            attn = cps[f"L{L}.mla.attn_out"][i].to(dev).float().view(H, vd)
+            W, kind = ck.linear(f"{PREFIX}{L}.self_attn.kv_b_proj")
+            W = W.view(H_all, nope + vd, lora)[rank * H : (rank + 1) * H]
+            W_UV = W[:, nope:, :]  # [H, vd, lora]
+            attn_ref = torch.einsum("hl,hvl->hv", out, W_UV)
+            up = rel(attn, attn_ref)
+            kern = "(needs q_nope/kv_rows from the extended tap)"
+            nrows = "-"
+            qn_key = f"L{L}.mla.q_nope"
+            kv_key = f"L{L}.mla.kv_rows"
+            if qn_key in cps and kv_key in extra:
+                q_nope = cps[qn_key][i].to(dev).float().view(H, lora)
+                kv = extra[kv_key][i].to(dev)
+                ks = float(extra.get(f"L{L}.mla.k_scale", 1.0))
+                if kv.dtype == torch.uint8:
+                    kv = kv.view(torch.float8_e4m3fn)
+                latent = kv.float() * ks  # [N, lora]
+                logits = torch.einsum("hl,nl->hn", q_nope, latent) * scale
+                p = torch.softmax(logits, dim=-1)
+                out_ref = p @ latent
+                kern = rel(out, out_ref)
+                nrows = str(latent.shape[0])
+            # fp8 KV write side: the cache row read back this step vs the bf16 latent written
+            wr = ""
+            cr_key, kvc_key = f"L{L}.mla.cache_row", f"L{L}.mla.kvc"
+            if cr_key in cps and kvc_key in cps:
+                cr = cps[cr_key][i].to(dev)
+                kvc = cps[kvc_key][i].to(dev).float()
+                if cr.dtype == torch.uint8:
+                    cr = cr.view(torch.float8_e4m3fn)
+                crf = cr.float()
+                ks = float(extra.get(f"L{L}.mla.k_scale", 1.0))
+                fit = float((crf * kvc).sum() / ((crf * crf).sum() + 1e-12))  # least-squares scale
+                wr = f"kv write: {rel(crf * ks, kvc)} (k_scale {ks:g}, ls-fit {fit:.4g})"
+            print(f"L{L:<5d} {kind:6s} {up:28s} {kern:28s} {nrows:5s} {float(out.abs().max()):.3f}  {wr}")
+            # rows this step read vs rows another run wrote at the same positions
+            if args.kv_vs and kv_key in extra:
+                other = torch.load(args.kv_vs, map_location="cpu")
+                ocr = other["cps"].get(cr_key)
+                topk = cps[f"L{L}.mla.topk"][i]
+                kv = extra[kv_key][i]
+                if kv.dtype == torch.uint8:
+                    kv = kv.view(torch.float8_e4m3fn)
+                kvf = kv.float()
+                worst, bad = 0.0, []
+                for j in range(kvf.shape[0]):
+                    p = int(topk[j])
+                    orow = [k for k, r in enumerate(other["rows"]) if r["pos"] == p]
+                    if not orow or ocr is None:
+                        continue
+                    o = ocr[orow[-1]]
+                    o = (o.view(torch.float8_e4m3fn) if o.dtype == torch.uint8 else o).float()
+                    d = float((kvf[j] - o).norm() / (o.norm() + 1e-12))
+                    worst = max(worst, d)
+                    if d > 0.5:
+                        bad.append((p, round(d, 3)))
+                print(f"        kv rows read vs {os.path.basename(args.kv_vs)} rows written: worst rel {worst:.3f}; rows > 0.5: {bad[:10]}")
+
+
+if __name__ == "__main__":
+    main()

docker-images/glm53-kindling/dev/kernel-tests/mlp_ref.py

@@ -0,0 +1,198 @@
+#!/usr/bin/env python3
+"""Offline MLP/MoE reference for rows dumped by glm53_step_tap (DEEP mode).
+
+For each `file:pos` spec the tool takes the row's `L<L>.mlp.in0` (the MLP
+input, replicated) and `L<L>.mlp.out0` (the MLP output after the TP
+all-reduce) from the vector dump, recomputes the layer from the nvidia NVFP4
+checkpoint two ways, and reports how far the dumped output is from each:
+
+  ref16  dequantised fp32 weights, fp32 activations, no activation quant
+         (what a W4A16 path such as marlin approximates)
+  ref4   same, but the activations are quantised to NVFP4 with the layer's
+         static input_scale before each GEMM, exactly like the W4A4 cutlass
+         path (scaled_fp4_quant + dequant)
+
+Layers < first_k_dense_replace (0..2) are dense Glm5NextMLP; the others are
+MoE: sigmoid scores, noaux_tc top-8 on scores + e_score_correction_bias,
+weights renormalised and scaled by routed_scaling_factor, plus the BF16
+shared expert. Between every pair of specs it also prints the input
+difference, the dumped output difference, and what ref16/ref4 make of the
+same two inputs, which separates legitimate sensitivity from a wrong path.
+
+    python3 mlp_ref.py --ckpt /ckpt --layer 0 /hit/vec-rank0-step3257.pt:92 /hit/vec-rank0-step3283.pt:92
+"""
+
+import argparse
+import json
+import os
+
+import torch
+from safetensors import safe_open
+
+from vllm import _custom_ops as ops
+from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
+    dequantize_to_dtype,
+)
+
+dev = torch.device("cuda")
+PREFIX = "model.language_model.layers."
+
+
+class Ckpt:
+    def __init__(self, d: str):
+        self.d = d
+        cfg = json.load(open(os.path.join(d, "config.json")))
+        self.cfg = cfg.get("text_config", cfg)
+        self.map = json.load(open(os.path.join(d, "model.safetensors.index.json")))["weight_map"]
+        self._open = {}
+
+    def get(self, name: str) -> torch.Tensor:
+        shard = self.map[name]
+        f = self._open.get(shard)
+        if f is None:
+            f = safe_open(os.path.join(self.d, shard), "pt", device="cuda")
+            self._open[shard] = f
+        return f.get_tensor(name)
+
+    def has(self, name: str) -> bool:
+        return name in self.map
+
+    def fp4_linear(self, base: str):
+        """Dequantised fp32 weight [out, in], and (weight_scale_2, input_scale)."""
+        w = self.get(base + ".weight")
+        s = self.get(base + ".weight_scale")
+        g = self.get(base + ".weight_scale_2").reshape(()).float()
+        a = self.get(base + ".input_scale").reshape(()).float()
+        W = dequantize_to_dtype(w[None], s[None], g[None], torch.float32, 16, swizzle=False)[0]
+        return W, float(g), float(a)
+
+
+def quant_act(x: torch.Tensor, input_scale: float) -> torch.Tensor:
+    """NVFP4-quantise a [M, K] activation with a static global scale, then dequantise (fp32)."""
+    q, sf = ops.scaled_fp4_quant(
+        x.to(torch.bfloat16), torch.tensor(1.0 / input_scale, device=dev), is_sf_swizzled_layout=True
+    )
+    return dequantize_to_dtype(q, sf, torch.tensor(input_scale, device=dev), torch.float32, 16, swizzle=True)
+
+
+def swiglu(gate: torch.Tensor, up: torch.Tensor, limit: float) -> torch.Tensor:
+    return torch.nn.functional.silu(gate.clamp(max=limit)) * up.clamp(-limit, limit)
+
+
+def mlp3(x: torch.Tensor, Wg, Wu, Wd, a_gu: float | None, a_d: float | None, limit: float) -> torch.Tensor:
+    """gate/up/down with optional activation quant (a_* = input_scale, None = no quant)."""
+    xq = quant_act(x, a_gu) if a_gu else x
+    h = swiglu(xq @ Wg.T, xq @ Wu.T, limit)
+    hq = quant_act(h, a_d) if a_d else h
+    return hq @ Wd.T
+
+
+def dense_refs(ck: Ckpt, L: int, X: torch.Tensor, limit: float):
+    base = f"{PREFIX}{L}.mlp."
+    Wg, gg, ag = ck.fp4_linear(base + "gate_proj")
+    Wu, gu, au = ck.fp4_linear(base + "up_proj")
+    Wd, gd, ad = ck.fp4_linear(base + "down_proj")
+    a_gu = max(ag, au)  # vLLM's merged gate_up takes the max shard input_scale
+    print(f"dense L{L}: input_scale gate {ag:.5g} up {au:.5g} (merged {a_gu:.5g}) down {ad:.5g}; "
+          f"weight_scale_2 {gg:.4g}/{gu:.4g}/{gd:.4g}")
+    r16 = mlp3(X, Wg, Wu, Wd, None, None, limit)
+    r4 = mlp3(X, Wg, Wu, Wd, a_gu, ad, limit)
+    return r16, r4, {}
+
+
+def moe_refs(ck: Ckpt, L: int, X: torch.Tensor, limit: float):
+    cfg = ck.cfg
+    E, K8, rsf = int(cfg["n_routed_experts"]), int(cfg["num_experts_per_tok"]), float(cfg["routed_scaling_factor"])
+    base = f"{PREFIX}{L}.mlp."
+    gw = ck.get(base + "gate.weight").float()  # [E, hidden]
+    bias = ck.get(base + "gate.e_score_correction_bias").float() if ck.has(base + "gate.e_score_correction_bias") else torch.zeros(E, device=dev)
+    logits = X @ gw.T
+    scores = torch.sigmoid(logits)
+    choice = scores + bias
+    top = torch.topk(choice, K8 + 1, dim=-1)
+    ids = top.indices[:, :K8]
+    margin = (top.values[:, K8 - 1] - top.values[:, K8]).tolist()
+    w = torch.gather(scores, 1, ids)
+    if cfg.get("norm_topk_prob", True):
+        w = w / w.sum(-1, keepdim=True)
+    w = w * rsf
+    # global activation scales as the cutlass path uses them: max over experts
+    a13 = max(float(ck.get(f"{base}experts.{e}.{p}.input_scale").reshape(())) for e in range(E) for p in ("gate_proj", "up_proj"))
+    a2 = max(float(ck.get(f"{base}experts.{e}.down_proj.input_scale").reshape(())) for e in range(E))
+    print(f"moe L{L}: global input_scale w13 {a13:.5g} w2 {a2:.5g}; routed_scaling {rsf}; renorm {cfg.get('norm_topk_prob')}")
+    # shared expert: BF16 (in the quantisation ignore list)
+    sg = ck.get(base + "shared_experts.gate_proj.weight").float()
+    su = ck.get(base + "shared_experts.up_proj.weight").float()
+    sd = ck.get(base + "shared_experts.down_proj.weight").float()
+    shared = mlp3(X, sg, su, sd, None, None, limit)
+    r16 = shared.clone()
+    r4 = shared.clone()
+    cache = {}
+    for r in range(X.shape[0]):
+        for j in range(K8):
+            e = int(ids[r, j])
+            if e not in cache:
+                Wg, _, _ = ck.fp4_linear(f"{base}experts.{e}.gate_proj")
+                Wu, _, _ = ck.fp4_linear(f"{base}experts.{e}.up_proj")
+                Wd, _, _ = ck.fp4_linear(f"{base}experts.{e}.down_proj")
+                cache[e] = (Wg, Wu, Wd)
+            Wg, Wu, Wd = cache[e]
+            xr = X[r : r + 1]
+            r16[r] += w[r, j] * mlp3(xr, Wg, Wu, Wd, None, None, limit)[0]
+            r4[r] += w[r, j] * mlp3(xr, Wg, Wu, Wd, a13, a2, limit)[0]
+    info = {"ids": ids.tolist(), "w": [[round(v, 4) for v in row] for row in w.tolist()], "margin_8_vs_9": margin}
+    return r16, r4, info
+
+
+def rel(a: torch.Tensor, b: torch.Tensor) -> str:
+    l2 = float((a - b).norm() / (b.norm() + 1e-12))
+    mx = float((a - b).abs().max() / (b.abs().max() + 1e-12))
+    return f"l2 {l2:.4f} max {mx:.4f}"
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--ckpt", required=True)
+    ap.add_argument("--layer", type=int, required=True)
+    ap.add_argument("specs", nargs="+", help="vec.pt:pos")
+    args = ap.parse_args()
+    ck = Ckpt(args.ckpt)
+    L = args.layer
+    limit = float(ck.cfg.get("swiglu_limit") or 1e30)
+    dense = L < int(ck.cfg.get("first_k_dense_replace", 0))
+    names, X, Y = [], [], []
+    for spec in args.specs:
+        path, _, pos = spec.rpartition(":")
+        b = torch.load(path, map_location="cpu")
+        rows = [i for i, r in enumerate(b["rows"]) if r["pos"] == int(pos)]
+        if not rows:
+            print(f"{spec}: no row at pos {pos}")
+            return 1
+        i = rows[-1]
+        names.append(f"{os.path.basename(path)}:{pos}")
+        X.append(b["cps"][f"L{L}.mlp.in0"][i].float())
+        Y.append(b["cps"][f"L{L}.mlp.out0"][i].float())
+    X = torch.stack(X).to(dev)
+    Y = torch.stack(Y).to(dev)
+    with torch.no_grad():
+        r16, r4, info = (dense_refs if dense else moe_refs)(ck, L, X, limit)
+    print(f"\nlayer L{L} ({'dense' if dense else 'moe'}), swiglu_limit {limit}")
+    for i, n in enumerate(names):
+        extra = ""
+        if info:
+            extra = f"  experts {info['ids'][i]} w {info['w'][i]} margin8/9 {info['margin_8_vs_9'][i]:.4f}"
+        print(f"{n:34s} |in|max {float(X[i].abs().max()):.3f} |out|max {float(Y[i].abs().max()):.3f}{extra}")
+        print(f"   dumped vs ref16: {rel(Y[i], r16[i])}   dumped vs ref4: {rel(Y[i], r4[i])}   ref4 vs ref16 (quant loss): {rel(r4[i], r16[i])}")
+    print("\npairs (a vs b):")
+    for i in range(len(names)):
+        for j in range(i + 1, len(names)):
+            same = ""
+            if info and info["ids"][i] != info["ids"][j]:
+                same = f"  EXPERT SETS DIFFER: {sorted(set(info['ids'][i]) ^ set(info['ids'][j]))}"
+            print(f"{names[i]} vs {names[j]}{same}")
+            print(f"   input {rel(X[i], X[j])} | dumped out {rel(Y[i], Y[j])} | ref16 out {rel(r16[i], r16[j])} | ref4 out {rel(r4[i], r4[j])}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

docker-images/glm53-kindling/dev/kernel-tests/nccl_collective_test.py

@@ -0,0 +1,208 @@
+#!/usr/bin/env python3
+"""Collective-corruption test for the four GLM boxes, no model loaded.
+
+Repeats the collectives a GLM-5.3-Flash TP=4 decode step issues -- bf16
+all-reduce of [M, 4096] (o_proj, KDA o_proj, MoE, embedding: ~90 per step)
+and the bf16 logits all-gather of [M, 38720] -> [M, 154880] -- with inputs
+every rank can regenerate deterministically, and checks each result:
+
+  all-gather : bit-exact against the locally rebuilt concatenation
+  all-reduce : "int" data (bf16 integers in [-32, 32], exact in bf16) must be
+               bit-exact; "float" data must be within 2 bf16 ulps of an fp32
+               reference; in both cases every rank must hold the identical
+               result (checked by a gloo all_gather of a hash) and a repeated
+               all-reduce of the same input must be bit-identical.
+
+Launch one process per box with RANK/WORLD_SIZE/MASTER_ADDR/MASTER_PORT set
+(see nccl_test_launch.sh). --mode picks the communicator:
+  vllm   : vLLM's own TP group (init_distributed_environment +
+           initialize_model_parallel), i.e. the same dispatcher and PyNCCL
+           communicator the model runner uses
+  pynccl : vllm.distributed.device_communicators.pynccl.PyNcclCommunicator
+           on a gloo process group
+  torch  : torch.distributed NCCL backend
+"""
+import argparse
+import os
+import sys
+import time
+
+import torch
+import torch.distributed as dist
+
+HIDDEN = 4096
+VOCAB = 154880
+
+
+def parse():
+    p = argparse.ArgumentParser()
+    p.add_argument("--mode", default="vllm", choices=["vllm", "pynccl", "torch"])
+    p.add_argument("--iters", type=int, default=2000, help="decode steps to emulate")
+    p.add_argument("--ar-per-step", type=int, default=90)
+    p.add_argument("--tokens", default="1,2,4,8,16,64")
+    p.add_argument("--prefill-tokens", type=int, default=2304, help="0 to skip")
+    p.add_argument("--data", default="both", choices=["int", "float", "both"])
+    return p.parse_args()
+
+
+def main():
+    args = parse()
+    rank = int(os.environ["RANK"])
+    world = int(os.environ["WORLD_SIZE"])
+    master = os.environ["MASTER_ADDR"]
+    port = int(os.environ["MASTER_PORT"])
+    torch.cuda.set_device(0)
+    dev = torch.device("cuda:0")
+    init = f"tcp://{master}:{port}"
+
+    for k in sorted(os.environ):
+        if k.startswith(("NCCL_", "GLOO_", "VLLM_HOST_IP")):
+            print(f"[rank {rank}] {k}={os.environ[k]}", flush=True)
+
+    if args.mode == "vllm":
+        from vllm.distributed import (
+            get_tp_group,
+            init_distributed_environment,
+            initialize_model_parallel,
+        )
+
+        from vllm.config import VllmConfig, set_current_vllm_config
+
+        _cfg_ctx = set_current_vllm_config(VllmConfig())
+        _cfg_ctx.__enter__()  # the TP group and its communicator read the config
+        init_distributed_environment(world, rank, init, local_rank=0, backend="nccl")
+        initialize_model_parallel(tensor_model_parallel_size=world)
+        tp = get_tp_group()
+        cpu_group = tp.cpu_group
+
+        def all_reduce(x):
+            return tp.all_reduce(x)
+
+        def all_gather(x):
+            return tp.all_gather(x, dim=-1)
+
+        names = getattr(getattr(tp, "device_communicator", None), "__class__", type(None)).__name__
+        print(f"[rank {rank}] vllm tp group world={tp.world_size} device_communicator={names}", flush=True)
+    else:
+        dist.init_process_group("gloo" if args.mode == "pynccl" else "nccl",
+                                init_method=init, rank=rank, world_size=world)
+        cpu_group = dist.new_group(backend="gloo") if args.mode == "torch" else dist.group.WORLD
+        if args.mode == "pynccl":
+            from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
+
+            comm = PyNcclCommunicator(group=dist.group.WORLD, device=dev)
+            assert not comm.disabled
+
+            def all_reduce(x):
+                return comm.all_reduce(x)
+
+            def all_gather(x):
+                out = torch.empty(x.shape[0], x.shape[1] * world, dtype=x.dtype, device=dev)
+                # nccl all_gather is rank-major over the flat buffer
+                tmp = torch.empty(world, x.shape[0], x.shape[1], dtype=x.dtype, device=dev)
+                comm.all_gather(tmp, x.contiguous())
+                out.copy_(tmp.permute(1, 0, 2).reshape(x.shape[0], -1))
+                return out
+        else:
+            def all_reduce(x):
+                y = x.clone()
+                dist.all_reduce(y)
+                return y
+
+            def all_gather(x):
+                parts = [torch.empty_like(x) for _ in range(world)]
+                dist.all_gather(parts, x.contiguous())
+                return torch.cat(parts, dim=-1)
+
+    def gen(seed, shape, data):
+        g = torch.Generator(device=dev).manual_seed(seed)
+        if data == "int":
+            return torch.randint(-32, 33, shape, device=dev, generator=g).to(torch.bfloat16)
+        return (torch.randn(shape, device=dev, generator=g) * 3).to(torch.bfloat16)
+
+    def rank_hash(t):
+        v = t.contiguous().view(torch.int16).reshape(-1).to(torch.int64)
+        w = torch.arange(1, v.numel() + 1, device=dev, dtype=torch.int64)
+        return int((v * w).sum().item() & 0xFFFFFFFFFFFF)
+
+    def check_consistent(y, tag):
+        h = rank_hash(y)
+        objs = [None] * world
+        dist.all_gather_object(objs, h, group=cpu_group)
+        if len(set(objs)) != 1:
+            return f"{tag}: ranks disagree, hashes={objs}"
+        return None
+
+    tokens = [int(t) for t in args.tokens.split(",")]
+    datas = ["int", "float"] if args.data == "both" else [args.data]
+    failures = 0
+    t0 = time.time()
+    total = 0
+    for it in range(args.iters):
+        M = tokens[it % len(tokens)]
+        data = datas[it % len(datas)]
+        for j in range(args.ar_per_step):
+            seed = it * 1_000_003 + j * 7919
+            xs = [gen(seed + r * 31, (M, HIDDEN), data) for r in range(world)]
+            ref32 = torch.zeros(M, HIDDEN, device=dev)
+            for r in range(world):
+                ref32 += xs[r].float()
+            y = all_reduce(xs[rank])
+            y2 = all_reduce(xs[rank])
+            total += 2
+            torch.cuda.synchronize()
+            msg = None
+            if not torch.isfinite(y).all():
+                msg = "non-finite output"
+            elif data == "int":
+                if not torch.equal(y.float(), ref32):
+                    bad = (y.float() != ref32).nonzero()
+                    msg = f"int all-reduce mismatch at {bad[:4].tolist()} got {y.float()[tuple(bad[0])].item()} want {ref32[tuple(bad[0])].item()}"
+            else:
+                err = (y.float() - ref32).abs()
+                # bound relative to the partial sums NCCL may round, not the
+                # (possibly cancelled) final value; the int rows are the exact check
+                ulp = sum(x.float().abs() for x in xs).clamp(min=1e-2) * 2 ** -7
+                if (err > 2 * ulp).any():
+                    bad = (err > 2 * ulp).nonzero()
+                    msg = f"float all-reduce error at {bad[:4].tolist()} got {y.float()[tuple(bad[0])].item()} want {ref32[tuple(bad[0])].item()}"
+            if msg is None and not torch.equal(y, y2):
+                msg = "repeat all-reduce differs (nondeterministic)"
+            if msg is None and (j % 15 == 0):
+                msg = check_consistent(y, "all-reduce")
+            if msg:
+                failures += 1
+                print(f"[rank {rank}] FAIL step={it} ar={j} M={M} data={data}: {msg}", flush=True)
+        # logits all-gather
+        seed = it * 1_000_003 + 999_331
+        xs = [gen(seed + r * 31, (M, VOCAB // world), "float") for r in range(world)]
+        y = all_gather(xs[rank])
+        total += 1
+        torch.cuda.synchronize()
+        exp = torch.cat(xs, dim=-1)
+        if not torch.equal(y, exp):
+            bad = (y != exp).nonzero()
+            failures += 1
+            print(f"[rank {rank}] FAIL step={it} all-gather M={M}: mismatch at {bad[:4].tolist()} (first bad column {int(bad[0,1])} = rank {int(bad[0,1]) // (VOCAB // world)}'s slice)", flush=True)
+        if args.prefill_tokens and it % 200 == 0:
+            xs = [gen(seed + 5 + r * 31, (args.prefill_tokens, HIDDEN), "int") for r in range(world)]
+            ref32 = sum(x.float() for x in xs)
+            y = all_reduce(xs[rank])
+            total += 1
+            torch.cuda.synchronize()
+            if not torch.equal(y.float(), ref32):
+                failures += 1
+                print(f"[rank {rank}] FAIL step={it} prefill all-reduce [{args.prefill_tokens},4096] mismatch", flush=True)
+        if it % 100 == 0 and rank == 0:
+            print(f"[rank 0] step {it}/{args.iters} M={M} {data} collectives={total} failures={failures} {time.time() - t0:.0f}s", flush=True)
+    # global verdict
+    objs = [None] * world
+    dist.all_gather_object(objs, failures, group=cpu_group)
+    if rank == 0:
+        print(f"RESULT: failures per rank {objs}, collectives per rank {total}, {'CLEAN' if sum(objs) == 0 else 'CORRUPT'}", flush=True)
+    dist.barrier(group=cpu_group)
+    sys.exit(0 if sum(objs) == 0 else 1)
+
+
+if __name__ == "__main__":
+    main()

docker-images/glm53-kindling/dev/kernel-tests/nccl_test_launch.sh

@@ -0,0 +1,71 @@
+#!/usr/bin/env bash
+# Launch nccl_collective_test.py inside the spark-glm53 image with the SAME
+# per-box network derivation entrypoint.sh performs (LAN address for the
+# bootstrap, RoCE devices + GID index for the data plane). Run one container
+# per box; see README.md for the docker command.
+#
+# Required env: RANK (0 on node0), WORLD_SIZE=4, MASTER_ADDR=192.0.2.1,
+# MASTER_PORT. Optional: CLUSTER_SUBNET (198.51.100.), FABRIC_SUBNETS,
+# VLLM_HOST_IP (the LAN address; derived from the default route if unset),
+# any NCCL_* override (NCCL_PROTO, NCCL_ALGO, NCCL_IB_DISABLE, ...).
+# Remaining arguments go to the python test.
+set -euo pipefail
+
+CLUSTER_SUBNET="${CLUSTER_SUBNET:-198.51.100.}"
+FABRIC_SUBNETS="${FABRIC_SUBNETS:-$CLUSTER_SUBNET}"
+
+if [[ -z "${VLLM_HOST_IP:-}" ]]; then
+  VLLM_HOST_IP=$(ip -o -4 route get 1.1.1.1 2>/dev/null | awk '{for(i=1;i<=NF;i++) if($i=="src"){print $(i+1); exit}}')
+fi
+export VLLM_HOST_IP
+GLOO_SOCKET_IFNAME="${GLOO_SOCKET_IFNAME:-$(ip -o -4 addr show | awk -v ip="$VLLM_HOST_IP" '$4 ~ "^"ip"/" {print $2; exit}')}"
+export GLOO_SOCKET_IFNAME
+export NCCL_SOCKET_IFNAME="${NCCL_SOCKET_IFNAME:-$GLOO_SOCKET_IFNAME}"
+
+# Same as entrypoint.sh fabric_port(): "<rdma-device> <gid-index>" for the
+# port holding an address in $1.
+fabric_port() {
+  local prefix="$1" found addr ifname dev="" hex i t g d n
+  found=$(ip -o -4 addr show 2>/dev/null \
+      | awk -v p="$prefix" '$4 ~ "^"p {split($4,a,"/"); print $2, a[1]; exit}')
+  [[ -n "$found" ]] || return 1
+  ifname="${found%% *}"; addr="${found##* }"
+  for d in /sys/class/infiniband/*; do
+    for n in "$d"/ports/1/gid_attrs/ndevs/*; do
+      [[ -f "$n" ]] || continue
+      [[ "$(cat "$n" 2>/dev/null)" == "$ifname" ]] || continue
+      dev="$(basename "$d")"; break 2
+    done
+  done
+  [[ -n "$dev" ]] || return 1
+  hex=$(printf '%s' "$addr" | awk -F. '{printf "%02x%02x:%02x%02x", $1,$2,$3,$4}')
+  for i in $(seq 0 15); do
+    t=$(cat "/sys/class/infiniband/$dev/ports/1/gid_attrs/types/$i" 2>/dev/null) || continue
+    g=$(cat "/sys/class/infiniband/$dev/ports/1/gids/$i" 2>/dev/null) || continue
+    [[ "$t" == "RoCE v2" && "$g" == *"ffff:$hex" ]] || continue
+    printf '%s %s\n' "$dev" "$i"; return 0
+  done
+  return 1
+}
+
+if [[ "${NCCL_IB_DISABLE:-0}" != "1" && ( -z "${NCCL_IB_HCA:-}" || -z "${NCCL_IB_GID_INDEX:-}" ) ]]; then
+  _hcas=""; _gid=""
+  for _p in $FABRIC_SUBNETS; do
+    _r=$(fabric_port "$_p") || continue
+    _d="${_r%% *}"; _i="${_r##* }"
+    if [[ -n "$_gid" && "$_i" != "$_gid" ]]; then
+      echo "FATAL: fabric ports disagree on GID index ($_d at $_i, expected $_gid)" >&2; exit 1
+    fi
+    _gid="${_gid:-$_i}"; _hcas="${_hcas:+$_hcas,}$_d"
+  done
+  if [[ -z "$_hcas" ]]; then
+    echo "FATAL: no RoCE v2 GID for any of: $FABRIC_SUBNETS" >&2; ip -br addr show >&2; exit 1
+  fi
+  export NCCL_IB_HCA="${NCCL_IB_HCA:-$_hcas}" NCCL_IB_GID_INDEX="${NCCL_IB_GID_INDEX:-$_gid}"
+fi
+export NCCL_MAX_NCHANNELS="${NCCL_MAX_NCHANNELS:-8}"
+export NCCL_DEBUG="${NCCL_DEBUG:-WARN}"
+
+echo "rank=${RANK} host_ip=${VLLM_HOST_IP} ifname=${GLOO_SOCKET_IFNAME} NCCL_IB_HCA=${NCCL_IB_HCA:-} gid=${NCCL_IB_GID_INDEX:-} NCCL_IB_DISABLE=${NCCL_IB_DISABLE:-0} PROTO=${NCCL_PROTO:-} ALGO=${NCCL_ALGO:-}"
+cd /
+exec python3 "$(dirname "$0")/nccl_collective_test.py" "$@"

docker-images/glm53-kindling/dev/patch-tests/_entrypoint_discovery_test.sh

@@ -0,0 +1,144 @@
+#!/usr/bin/env bash
+# Runs image/discover.sh against mentatd status taken from the four GLM boxes
+# (fixtures/mentat, addresses moved to the documentation ranges), with
+# mentat_query answering from the fixtures and sleep stubbed out. Needs only
+# bash and python3. Exit status is the number of failures.
+set -euo pipefail
+here=$(cd "$(dirname "$0")" && pwd)
+disc="$here/../../image/discover.sh"
+tmp=$(mktemp -d)
+trap 'rm -rf "$tmp"' EXIT
+fails=0
+
+check() {
+  if [[ "$2" == "$3" ]]; then echo "  ok   $1"
+  else echo "  FAIL $1: got '$2', want '$3'"; fails=$(( fails + 1 )); fi
+}
+
+# One box's boot as far as the entrypoint takes discovery, in a clean
+# environment. $1 is the fixture directory, $2 the box (the last octet of its
+# address), the rest VAR=value settings from .env. Both waits are bounded, so
+# a broken parse ends in a failure while sleep is stubbed. Prints
+# VLLM_HOST_IP|FABRIC_SUBNETS|CLUSTER_SUBNET|MENTAT_NODE_IP|HEAD_HOST|ROLE|queries
+# on stdout, and the entrypoint's own messages on stderr.
+boot() {
+  local fx=$1 node=$2; shift 2
+  : > "$tmp/queries"
+  env -i PATH="$PATH" DISC="$disc" FX="$fx" NODE="$node" QLOG="$tmp/queries" \
+      TP=4 MENTAT_GROUP=glm53 RAY_ADDRESS=127.0.0.1:6379 \
+      FABRIC_WAIT_S=600 WORKER_WAIT_S=600 "$@" bash -c '
+    set -euo pipefail
+    . "$DISC"
+    mentat_query() {
+      echo "$1" >> "$QLOG"
+      case "$1" in
+        local) cat "$FX/local-$NODE.json" 2>/dev/null ;;
+        group) cat "$FX/group-glm53.json" 2>/dev/null ;;
+      esac
+    }
+    sleep() { :; }
+    # As entrypoint.sh does, in order.
+    _elect=""
+    [[ -z "${ROLE:-}" && -z "${HEAD_HOST:-}" ]] && _elect=1
+    discover_network >&2
+    CLUSTER_SUBNET="${CLUSTER_SUBNET:-${FABRIC_SUBNETS%% *}}"
+    if [[ -n "$_elect" ]]; then elect_head >&2; else fixed_role; fi
+    echo "${VLLM_HOST_IP:-}|${FABRIC_SUBNETS:-}|${CLUSTER_SUBNET:-}|${MENTAT_NODE_IP:-}|${HEAD_HOST:-}|${ROLE:-}|$(sort -u "$QLOG" | tr "\n" " ")"
+  '
+}
+
+# A copy of the fixtures with $2 (python, `local` and `group` as dicts keyed
+# by box) applied, in directory $1.
+variant() {
+  mkdir -p "$1"
+  python3 - "$here/fixtures/mentat" "$1" "$2" <<'PY'
+import json, os, sys
+src, dst, edit = sys.argv[1:]
+local = {n: json.load(open(f"{src}/local-{n}.json")) for n in ("36", "70", "77", "93")}
+group = json.load(open(f"{src}/group-glm53.json"))
+agents = group["groups"]["glm53"]["agents"]
+exec(edit)
+for n, d in local.items():
+    json.dump(d, open(f"{dst}/local-{n}.json", "w"))
+json.dump(group, open(f"{dst}/group-glm53.json", "w"))
+PY
+}
+
+echo "== all four boxes, as the fleet tags them today (one rdma root)"
+for n in 36 70 77 93; do
+  role=worker; [[ $n == 36 ]] && role=head
+  check "box .$n" "$(boot "$here/fixtures/mentat" $n 2>/dev/null)" \
+    "192.0.2.$n|198.18.0.|198.18.0.|192.0.2.$n|192.0.2.36|$role|group local "
+done
+boot "$here/fixtures/mentat" 70 2>&1 >/dev/null | grep '^election:' | sed 's/^/       /'
+
+echo "== both PCIe roots tagged rdma, listed out of order"
+variant "$tmp/roots" '
+for n, d in local.items():
+    a = f"198.19.0.{n}"
+    d["addrs"].insert(1, a)
+    d["addr_tags"][a] = ["connectx", "rdma"]
+    d["addr_ifaces"][a] = "enP2p1s0f1np1"'
+for n in 36 93; do
+  check "box .$n prefixes" "$(boot "$tmp/roots" $n 2>/dev/null | cut -d'|' -f2,3)" \
+    "198.18.0. 198.19.0.|198.18.0."
+done
+
+echo "== nothing tagged lan: VLLM_HOST_IP falls back to the daemon's node_ip"
+variant "$tmp/nolan" 'local["70"]["addr_tags"]["192.0.2.70"] = []'
+check "box .70" "$(boot "$tmp/nolan" 70 2>/dev/null | cut -d'|' -f1)" "192.0.2.70"
+
+echo "== a gone agent and a register-only agent below the lowest box"
+variant "$tmp/stale" '
+agents["glm53@glm53@192.0.2.5"] = dict(agents["glm53@glm53@192.0.2.36"],
+    node_ip="192.0.2.5", alive=False, gone_since_ms=1790624000000)
+agents["glm53@glm53-api@192.0.2.4"] = dict(agents["glm53@glm53@192.0.2.36"],
+    node_ip="192.0.2.4", machine={"cpus": 20, "gpus": [], "memory": 0})'
+for n in 36 77; do
+  role=worker; [[ $n == 36 ]] && role=head
+  check "box .$n" "$(boot "$tmp/stale" $n 2>/dev/null | cut -d'|' -f5,6)" "192.0.2.36|$role"
+done
+
+echo "== fewer than TP live agents: keeps waiting, then WORKER_WAIT_S ends it"
+variant "$tmp/three" 'agents["glm53@glm53@192.0.2.77"]["alive"] = False'
+set +e
+out=$(boot "$tmp/three" 70 WORKER_WAIT_S=65 2>&1); rc=$?
+set -e
+check "exit status" "$rc" 1
+check "waited" "$(grep -c '^electing: waiting for 4 agents in group glm53, have 3' <<< "$out")" 2
+check "reason" "$(grep -o 'FATAL: only 3 of 4 agents.*' <<< "$out")" \
+  "FATAL: only 3 of 4 agents in group glm53 after 65s (192.0.2.36 192.0.2.70 192.0.2.93)"
+
+echo "== mentatd not answering, or no rdma address: FABRIC_WAIT_S ends the wait"
+set +e
+out=$(boot "$tmp/none" 70 FABRIC_WAIT_S=30 2>&1); rc=$?
+set -e
+check "no daemon" "$rc $(grep -c '^FATAL: no answer from mentatd' <<< "$out")" "1 1"
+variant "$tmp/nordma" 'local["70"]["addr_tags"]["198.18.0.1"] = ["connectx"]'
+set +e
+out=$(boot "$tmp/nordma" 70 FABRIC_WAIT_S=30 2>&1); rc=$?
+set -e
+check "no rdma" "$rc $(grep -c '^FATAL: mentatd lists no address tagged rdma' <<< "$out")" "1 1"
+
+echo "== explicit settings skip discovery"
+check "all four set" \
+  "$(boot /nonexistent 70 ROLE=worker HEAD_HOST=192.0.2.93 VLLM_HOST_IP=192.0.2.70 CLUSTER_SUBNET=10.0.0. 2>/dev/null)" \
+  "192.0.2.70||10.0.0.||192.0.2.93|worker|"
+check "HEAD_HOST alone: a box that is not HEAD_HOST works" \
+  "$(boot /nonexistent 70 HEAD_HOST=192.0.2.93 VLLM_HOST_IP=192.0.2.70 CLUSTER_SUBNET=10.0.0. 2>/dev/null | cut -d'|' -f5-)" \
+  "192.0.2.93|worker|"
+check "HEAD_HOST alone: the box at HEAD_HOST leads" \
+  "$(boot /nonexistent 93 HEAD_HOST=192.0.2.93 VLLM_HOST_IP=192.0.2.93 CLUSTER_SUBNET=10.0.0. 2>/dev/null | cut -d'|' -f5-)" \
+  "192.0.2.93|head|"
+check "roles set, networking from mentat" \
+  "$(boot "$here/fixtures/mentat" 70 ROLE=worker HEAD_HOST=192.0.2.93 2>/dev/null)" \
+  "192.0.2.70|198.18.0.|198.18.0.||192.0.2.93|worker|local "
+check "networking set, roles elected" \
+  "$(boot "$here/fixtures/mentat" 93 VLLM_HOST_IP=192.0.2.93 CLUSTER_SUBNET=10.0.0. 2>/dev/null)" \
+  "192.0.2.93||10.0.0.|192.0.2.93|192.0.2.36|worker|group local "
+check "FABRIC_SUBNETS alone names CLUSTER_SUBNET" \
+  "$(boot /nonexistent 70 ROLE=worker HEAD_HOST=192.0.2.93 VLLM_HOST_IP=192.0.2.70 "FABRIC_SUBNETS=10.0.0. 10.0.1." 2>/dev/null)" \
+  "192.0.2.70|10.0.0. 10.0.1.|10.0.0.||192.0.2.93|worker|"
+
+echo "$fails failure(s)"
+exit "$fails"

docker-images/glm53-kindling/dev/patch-tests/_glm47_failclosed_stream_test.py

@@ -0,0 +1,74 @@
+"""End-to-end checks of the fail-closed parser on real GLM tool-call text, fed
+once whole (a non-streaming response) and once token by token (streaming).
+
+The model writes arguments as <arg_key>/<arg_value> pairs, which the engine
+turns into JSON; validating the raw text refused every streamed call. Needs the
+checkpoint's tokenizer at /model and the plugin at /tmp/plugin.py.
+"""
+import json
+
+from transformers import AutoTokenizer
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
+from vllm.tool_parsers.abstract_tool_parser import ToolParserManager as M
+
+M.import_tool_parser("/tmp/plugin.py")
+Parser = M.get_tool_parser("glm47_failclosed")
+tok = AutoTokenizer.from_pretrained("/model", trust_remote_code=True)
+tools = [
+    {"type": "function", "function": {"name": "bash", "parameters": {"type": "object", "properties": {
+        "command": {"type": "string"}, "timeout": {"type": "integer"}}, "required": ["command"]}}},
+    {"type": "function", "function": {"name": "read_file", "parameters": {"type": "object", "properties": {
+        "path": {"type": "string"}}}}},
+]
+req = ChatCompletionRequest(model="glm53", messages=[{"role": "user", "content": "x"}], tools=tools)
+R = "_rejected_by_server"
+CASES = [  # (label, model output, [(name, argument keys or R)])
+    ("good", "</think>Listing.<tool_call>bash<arg_key>command</arg_key><arg_value>ls /tmp</arg_value></tool_call>",
+     [("bash", ["command"])]),
+    ("two args", "</think><tool_call>bash<arg_key>command</arg_key><arg_value>sleep 1</arg_value>"
+     "<arg_key>timeout</arg_key><arg_value>30</arg_value></tool_call>", [("bash", ["command", "timeout"])]),
+    ("key not in schema", "</think><tool_call>bash<arg_key>cmd</arg_key><arg_value>ls</arg_value></tool_call>",
+     [("bash", R)]),
+    ("good then bad", "</think><tool_call>read_file<arg_key>path</arg_key><arg_value>/etc/hosts</arg_value>"
+     "</tool_call><tool_call>bash<arg_key>cmd</arg_key><arg_value>ls</arg_value></tool_call>",
+     [("read_file", ["path"]), ("bash", R)]),
+    ("markup after name", "</think><tool_call>bash</arg_key><arg_key>command</arg_key><arg_value>pwd"
+     "</arg_value></tool_call>", [("bash", ["command"])]),
+    ("tool not offered", "</think><tool_call>curl<arg_key>url</arg_key><arg_value>x</arg_value></tool_call>",
+     [("curl", R)]),
+]
+
+
+def whole(text):
+    info = Parser(tok, req.tools).extract_tool_calls(text, req)
+    return [(c.function.name, c.function.arguments) for c in info.tool_calls]
+
+
+def streamed(text):
+    p = Parser(tok, req.tools)
+    ids = tok.encode(text, add_special_tokens=False)
+    names, args, prev_text = {}, {}, ""
+    for i in range(len(ids)):
+        cur = tok.decode(ids[: i + 1])
+        d = p.extract_tool_calls_streaming(prev_text, cur, cur[len(prev_text):], ids[:i], ids[: i + 1], [ids[i]], req)
+        for tc in (d.tool_calls or []) if d else []:
+            if tc.function and tc.function.name:
+                names[tc.index] = tc.function.name
+            if tc.function and tc.function.arguments:
+                args[tc.index] = args.get(tc.index, "") + tc.function.arguments
+        prev_text = cur
+    return [(names.get(k), args.get(k, "")) for k in sorted(set(names) | set(args))]
+
+
+def shape(calls):
+    out = []
+    for name, arguments in calls:
+        parsed = json.loads(arguments or "{}")
+        out.append((name, R if R in parsed else sorted(parsed)))
+    return out
+
+
+for label, text, want in CASES:
+    for mode, run in (("whole", whole), ("streamed", streamed)):
+        got = shape(run(text))
+        print(f"  {label:<18} {mode:<9} {'ok' if got == want else 'WANTED ' + str(want) + ' GOT'} {got}")

docker-images/glm53-kindling/dev/patch-tests/_glm47_failclosed_test.py

@@ -0,0 +1,43 @@
+"""Model-free checks of the fail-closed validator and name resolution."""
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionToolsParam
+from vllm.tool_parsers.abstract_tool_parser import ToolParserManager as M
+
+M.import_tool_parser("/tmp/plugin.py")
+P = M.get_tool_parser("glm47_failclosed")._parser_engine_cls
+
+def tool(name, props):
+    return ChatCompletionToolsParam.model_validate({
+        "type": "function",
+        "function": {"name": name, "parameters": {"type": "object", "properties": props}},
+    })
+
+parser = P.__new__(P)  # only _tools is needed
+parser._tools = [
+    tool("bash", {"command": {}, "timeout": {}, "workdir": {}}),
+    tool("run_code", {"code": {}, "description": {}}),
+]
+
+print("== name resolution")
+for emitted, want in [("bash", "bash"), ("bash</arg_key>x", "bash"),
+                      ("bash1635", None), ("run_code", "run_code"),
+                      ("nope", None), ("run_code</arg_value>", "run_code")]:
+    got = parser._resolve_name(emitted)
+    print(f"  {emitted:<24} -> {str(got):<10} {'ok' if got == want else 'WANTED ' + str(want)}")
+
+print("== validation")
+cases = [
+    ("good call", "bash", '{"command": "ls", "timeout": 5}', True),
+    ("no args", "bash", "", True),
+    ("unknown tool", "list_files", '{"path": "/"}', False),
+    ("corrupted name", "bash</arg_value><arg_key>x", '{"command": "ls"}', False),
+    ("key not in schema", "bash", '{"prefix": "x"}', False),
+    ("corrupt key", "run_code",
+     '{"print_code_snapshot</arg_value><arg_key>description": "x"}', False),
+    ("not json", "bash", '{"command": ', False),
+    ("args not an object", "bash", '["command"]', False),
+]
+for label, name, args, want_ok in cases:
+    reason = parser._reject_reason(name, args)
+    ok = reason is None
+    mark = "ok" if ok == want_ok else "UNEXPECTED"
+    print(f"  {label:<20} {'accept' if ok else 'refuse':<7} {mark:<11} {(reason or '')[:60]}")

docker-images/glm53-kindling/dev/patch-tests/_glm53_dflash2_kv_groups_test.py

@@ -0,0 +1,354 @@
+#!/usr/bin/env python3
+"""Rebuild GLM-5.3-Flash TP=4 + DFlash2's KV specs on the nightly and drive
+get_kv_cache_groups / get_kv_cache_configs. Mode 'before' expects the stock
+tree's NotImplementedError; mode 'after' checks the patched tree end to end
+(groups, configs, block sizes, coordinator, pickle, a small GPU allocation).
+"""
+import os
+import pickle
+import sys
+from dataclasses import replace
+
+import torch
+
+MODE = sys.argv[1]
+MODEL = "/models/glm-5.3-flash-nvfp4-nvidia"
+DRAFT = "/models/glm-5.3-flash-dflash2"
+KV_PIN = 27917287424
+TP = 4
+
+from vllm.engine.arg_utils import EngineArgs  # noqa: E402
+
+args = EngineArgs(
+    model=MODEL,
+    tensor_parallel_size=TP,
+    distributed_executor_backend="ray",
+    max_model_len=524288,
+    kv_cache_dtype="fp8_e4m3",
+    block_size=2304,
+    kv_cache_memory_bytes=KV_PIN,
+    enable_prefix_caching=True,
+    max_num_batched_tokens=16384,
+    max_num_seqs=32,
+    long_prefill_token_threshold=2304,
+    trust_remote_code=True,
+    gpu_memory_utilization=0.88,
+    speculative_config={
+        "method": "dflash",
+        "model": DRAFT,
+        "num_speculative_tokens": 7,
+    },
+)
+vllm_config = args.create_engine_config()
+cc = vllm_config.cache_config
+mc = vllm_config.model_config
+# The engine core resolves the layout before grouping (resolve_kv_cache_layout):
+# candidates[0] of the workers' supported list. FlashInferMLASparseSM90Backend,
+# DeepseekV32IndexerBackend (GLM indexer) and KpoolTailBackend each declare only
+# LBHNC; TritonAttn (the drafter) declares no preference. So: LBHNC.
+cc.kv_cache_layout = "LBHNC"
+sc = vllm_config.speculative_config
+print(f"== vllm {__import__('vllm').__version__}")
+print(
+    f"cache: block_size={cc.block_size} cache_dtype={cc.cache_dtype} "
+    f"mamba_block_size={cc.mamba_block_size} mamba_page_size_padded={cc.mamba_page_size_padded} "
+    f"mamba_cache_mode={cc.mamba_cache_mode} mamba_cache_dtype={cc.mamba_cache_dtype} "
+    f"mamba_ssm_cache_dtype={cc.mamba_ssm_cache_dtype} use_kda_recoverssm={cc.use_kda_recoverssm} "
+    f"kv_cache_memory_bytes={cc.kv_cache_memory_bytes} kv_cache_layout={cc.kv_cache_layout} "
+    f"skip_page_size_padded={cc.skip_page_size_padded}"
+)
+print(
+    f"spec: method={sc.method} k={sc.num_speculative_tokens} draft_arch={sc.draft_model_config.architectures} "
+    f"use_eagle={sc.use_eagle()} block_drop={sc.use_eagle_block_drop()} "
+    f"lookahead={vllm_config.num_prefill_lookahead_tokens} max_in_flight={vllm_config.max_in_flight_tokens} "
+    f"max_model_len={mc.max_model_len}"
+)
+
+from vllm.config import set_current_vllm_config  # noqa: E402
+from vllm.model_executor.layers.mamba.mamba_utils import (  # noqa: E402
+    MambaStateDtypeCalculator,
+    MambaStateShapeCalculator,
+)
+from vllm.models.glm5next.common.attention import (  # noqa: E402
+    Glm5NextIndexerCache,
+    Glm5NextTailCache,
+)
+from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype  # noqa: E402
+from vllm.v1.kv_cache_interface import (  # noqa: E402
+    MambaSpec,
+    MLAAttentionSpec,
+    SlidingWindowSpec,
+    UniformTypeKVCacheSpecs,
+    get_kv_quant_mode,
+)
+
+hf = mc.hf_text_config
+mla_layers = list(hf.linear_attn_config["full_attn_layers"])
+kda_layers = list(hf.linear_attn_config["kda_layers"])
+assert hf.kv_lora_rank == 512 and hf.index_kpool == 4 and hf.index_head_dim == 128
+P = "language_model.model.layers"
+
+fp8 = kv_cache_dtype_str_to_dtype("fp8_e4m3", mc)
+qm = get_kv_quant_mode("fp8_e4m3")
+from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention  # noqa: E402
+
+kda_type = GatedDeltaNetAttention.mamba_type.fget(None)
+
+specs = {}
+with set_current_vllm_config(vllm_config):
+    for i in range(hf.num_hidden_layers):
+        if i in mla_layers:
+            kc = Glm5NextIndexerCache(
+                head_dim=hf.index_head_dim + hf.index_head_dim // 128 * 4,
+                dtype=torch.uint8,
+                prefix=f"{P}.{i}.self_attn.indexer.k_cache",
+                cache_config=cc,
+                index_kpool=hf.index_kpool,
+            )
+            specs[kc.prefix] = kc.get_kv_cache_spec(vllm_config)
+            tc = Glm5NextTailCache(
+                head_dim=hf.index_head_dim,
+                dtype=torch.bfloat16,
+                prefix=f"{P}.{i}.self_attn.indexer.tail_cache",
+                cache_config=cc,
+                index_kpool=hf.index_kpool,
+            )
+            specs[tc.prefix] = tc.get_kv_cache_spec(vllm_config)
+            # mla_attention.MLAAttention.get_kv_cache_spec, head_size = kv_lora_rank + qk_rope(0)
+            specs[f"{P}.{i}.self_attn.mla_attn"] = MLAAttentionSpec(
+                block_size=cc.block_size,
+                num_kv_heads=1,
+                head_size=hf.kv_lora_rank + hf.qk_rope_head_dim,
+                dtype=fp8,
+                cache_dtype_str="fp8_e4m3",
+                kv_quant_mode=qm,
+                state_content_bytes=None,
+                is_index_group_leader=True,
+                non_causal_multi_token_decode=False,
+            )
+        else:
+            assert i in kda_layers
+            shapes = MambaStateShapeCalculator.kda_state_shape(
+                TP,
+                hf.linear_num_heads,
+                hf.linear_head_dim,
+                conv_kernel_size=hf.linear_conv_kernel_dim,
+                num_spec=vllm_config.num_speculative_tokens,
+            )
+            dtypes = MambaStateDtypeCalculator.kda_state_dtype(
+                mc.dtype, cc.mamba_cache_dtype, cc.mamba_ssm_cache_dtype
+            )
+            specs[f"{P}.{i}.self_attn"] = MambaSpec(
+                shapes=tuple(shapes),
+                dtypes=dtypes,
+                block_size=cc.mamba_block_size,
+                page_size_padded=cc.mamba_page_size_padded,
+                mamba_type=kda_type,
+                tp_replicated=False,
+                mamba_cache_mode=cc.mamba_cache_mode,
+                num_speculative_blocks=(
+                    0 if cc.use_kda_recoverssm else vllm_config.num_speculative_tokens
+                ),
+            )
+
+# The drafter: build a real Attention layer under the draft config so the
+# backend selection and Attention.get_kv_cache_spec are the nightly's own.
+dhf = sc.draft_model_config.hf_config
+draft_cfg = replace(
+    vllm_config,
+    model_config=sc.draft_model_config,
+    attention_config=replace(
+        vllm_config.attention_config, backend=None, use_non_causal=True
+    ),
+)
+draft_names = [f"model.layers.{i}.self_attn.attn" for i in range(dhf.num_hidden_layers)]
+draft_backend = None
+try:
+    from vllm.model_executor.layers.attention.attention import Attention
+
+    with set_current_vllm_config(draft_cfg):
+        attn = Attention(
+            dhf.num_attention_heads // TP,
+            dhf.head_dim,
+            dhf.head_dim**-0.5,
+            num_kv_heads=max(1, dhf.num_key_value_heads // TP),
+            cache_config=cc,
+            quant_config=None,
+            per_layer_sliding_window=dhf.sliding_window,
+            prefix=draft_names[0],
+        )
+        draft_spec = attn.get_kv_cache_spec(draft_cfg)
+        draft_backend = attn.get_attn_backend()
+    print(
+        f"draft backend: {draft_backend.get_name()} kernel blocks "
+        f"{draft_backend.get_supported_kernel_block_sizes()}"
+    )
+except Exception as e:  # noqa: BLE001
+    print(f"!! real drafter Attention construction failed: {type(e).__name__}: {e}")
+    print("   falling back to a hand-built SlidingWindowSpec at the primary block")
+    draft_spec = SlidingWindowSpec(
+        block_size=cc.block_size,
+        num_kv_heads=max(1, dhf.num_key_value_heads // TP),
+        head_size=dhf.head_dim,
+        dtype=fp8,
+        kv_quant_mode=qm,
+        sliding_window=dhf.sliding_window,
+    )
+for name in draft_names:
+    specs[name] = draft_spec
+print(f"draft spec: {draft_spec}")
+
+print(f"== {len(specs)} layer specs")
+seen = set()
+for name, spec in specs.items():
+    key = (type(spec).__name__, spec.block_size, spec.page_size_bytes)
+    if key in seen:
+        continue
+    seen.add(key)
+    print(f"  {name}: {type(spec).__name__} block={spec.block_size} page={spec.page_size_bytes}")
+
+from vllm.v1.core import kv_cache_utils as U  # noqa: E402
+
+specs_list = [dict(specs) for _ in range(TP)]
+if MODE == "before":
+    try:
+        U.get_kv_cache_configs(vllm_config, specs_list, [KV_PIN] * TP)
+    except NotImplementedError as e:
+        print(f"== BEFORE: NotImplementedError as expected:\n   {e}")
+        sys.exit(0)
+    print("!! BEFORE: no error raised; the stock tree changed")
+    sys.exit(1)
+
+assert MODE == "after"
+groups = U.get_kv_cache_groups(vllm_config, dict(specs))
+print(f"== AFTER: {len(groups)} groups")
+for gid, g in enumerate(groups):
+    s = g.kv_cache_spec
+    inner = (
+        f"{len(s.kv_cache_specs)} inner specs "
+        f"{sorted({type(x).__name__ for x in s.kv_cache_specs.values()})}"
+        if isinstance(s, UniformTypeKVCacheSpecs)
+        else ""
+    )
+    print(
+        f"  [{gid}] {type(s).__name__} block={s.block_size} page={s.page_size_bytes} "
+        f"cacheable={s.prefix_cacheable} eagle={g.is_eagle_group} layers={len(g.layer_names)} "
+        f"{g.layer_names[0]}..{g.layer_names[-1]} {inner}"
+    )
+
+configs = U.get_kv_cache_configs(vllm_config, specs_list, [KV_PIN] * TP)
+cfg = configs[0]
+assert all(c.num_blocks == cfg.num_blocks for c in configs)
+bpb = U._pool_bytes_per_block(cfg.kv_cache_groups)
+print(f"== config: num_blocks={cfg.num_blocks} bytes/block={bpb} ({bpb/2**20:.2f} MiB) "
+      f"pool={cfg.num_blocks*bpb/2**30:.2f} GiB of {KV_PIN/2**30:.2f} GiB pinned")
+layout = U._glm5_next_tensor_layout(cfg.kv_cache_groups)
+assert layout is not None
+attn_group, mamba_groups, mla_names, idx_names, mla_page, idx_page, tail_names, tail_page, draft_groups = layout
+print(f"   layout: {len(mla_names)} MLA x {mla_page} + {len(idx_names)} idx x {idx_page}; tail {len(tail_names)} x {tail_page}; "
+      f"{len(mamba_groups)} mamba groups; {len(draft_groups)} draft groups "
+      f"({sum(len(g.layer_names) for g in draft_groups)} layers x {draft_groups[0].kv_cache_spec.page_size_bytes} B)")
+size = cfg.kv_cache_tensors[0].size
+assert size == bpb * cfg.num_blocks
+mla_end = len(mla_names) * mla_page * cfg.num_blocks
+idx_end = mla_end + len(idx_names) * idx_page * cfg.num_blocks
+regions = []
+for t in cfg.kv_cache_tensors:
+    assert t.size == size
+    if t.layers[0] in draft_names:
+        span = t.layer_stride * len(t.layers)
+        assert t.offset >= idx_end and t.offset + span <= size, (t, idx_end, size)
+        regions.append((t.offset, t.offset + span, t.layers[0], t.block_stride))
+regions.sort()
+assert len(regions) == len(draft_names)
+for a, b in zip(regions, regions[1:]):
+    assert a[1] <= b[0], "draft regions overlap"
+assert regions[-1][1] == size and regions[0][0] == idx_end, "draft regions do not tile the tail of the pool"
+print(f"   draft regions: {len(regions)} x {regions[0][1]-regions[0][0]} B, block_stride={regions[0][3]}, "
+      f"from {idx_end} to {size} (pool end)")
+
+sched_bs, hash_bs = U.resolve_kv_cache_block_sizes(cfg, vllm_config)
+participating = [g.kv_cache_spec.block_size for g in cfg.kv_cache_groups if g.kv_cache_spec.prefix_cacheable]
+print(f"== block sizes: scheduler={sched_bs} hash={hash_bs} engine-core min(participating)={min(participating)} "
+      f"participating={participating}")
+assert sched_bs == cc.block_size and min(participating) == cc.block_size
+
+sched_cfg = U.generate_scheduler_kv_cache_config(configs)
+tokens, conc = U.get_kv_cache_capacity(vllm_config, sched_cfg)
+print(f"== capacity: {tokens:,} tokens, {conc:.2f}x at {mc.max_model_len}")
+for gid, g in enumerate(sched_cfg.kv_cache_groups):
+    s = g.kv_cache_spec
+    need = U.cdiv(s.max_memory_usage_bytes(vllm_config), s.page_size_bytes)
+    print(f"   group {gid} {type(s).__name__}: {need} blocks per max-len request")
+
+from vllm.v1.core.kv_cache_coordinator import get_kv_cache_coordinator  # noqa: E402
+
+coord = get_kv_cache_coordinator(
+    kv_cache_config=sched_cfg,
+    max_model_len=mc.max_model_len,
+    max_in_flight_tokens=vllm_config.max_in_flight_tokens,
+    use_eagle=sc.use_eagle_block_drop(),
+    enable_caching=True,
+    enable_kv_cache_events=False,
+    dcp_world_size=1,
+    pcp_world_size=1,
+    scheduler_block_size=sched_bs,
+    hash_block_size=hash_bs,
+    num_prefill_lookahead=vllm_config.num_prefill_lookahead_tokens,
+)
+print(f"== coordinator: {type(coord).__name__}; managers "
+      f"{[type(m).__name__ for m in coord.single_type_managers]}")
+print(f"   hit-lookup groups: {[(type(g.spec).__name__, g.group_ids, g.use_eagle) for g in coord.attention_groups]}")
+print(f"   eagle_group_ids={sorted(coord.eagle_group_ids)} partial_hash_hits={coord.enable_partial_hash_hits}")
+
+rt = pickle.loads(pickle.dumps(cfg))
+assert rt.kv_cache_groups[-1].kv_cache_spec == cfg.kv_cache_groups[-1].kv_cache_spec
+assert type(rt.kv_cache_groups[-1].kv_cache_spec).__module__ == "vllm.v1.core.kv_cache_utils"
+assert not rt.kv_cache_groups[-1].kv_cache_spec.prefix_cacheable
+print("== pickle round trip ok (draft spec by module path, still opted out)")
+
+# PP projection: a stage that holds no drafter layer must still take the GLM
+# layout (empty draft group), count no draft bytes and emit no draft tensor.
+stage_specs = {k: v for k, v in specs.items() if k not in draft_names}
+proj = U._project_kv_cache_groups_to_worker(groups, stage_specs)
+assert proj[-1].layer_names == [] and isinstance(proj[-1].kv_cache_spec, U.DraftSlidingWindowSpec)
+pl = U._glm5_next_tensor_layout(proj)
+assert pl is not None and pl[-1][0].layer_names == []
+stage_cfg = U.get_kv_cache_config_from_groups(vllm_config, proj, KV_PIN)
+stage_bpb = U._pool_bytes_per_block(proj)
+assert stage_bpb == len(mla_names) * mla_page + len(idx_names) * idx_page
+assert not any(t.layers[0] in draft_names for t in stage_cfg.kv_cache_tensors)
+assert U._max_memory_usage_bytes_from_groups(vllm_config, proj) < U._max_memory_usage_bytes_from_groups(vllm_config, groups)
+print(f"== PP projection without the drafter: {len(proj)} groups (last empty), bytes/block={stage_bpb}, "
+      f"num_blocks={stage_cfg.num_blocks}, no draft tensors")
+
+# A small GPU allocation with the real view builder and a kernel split for the drafter.
+if torch.cuda.is_available():
+    from vllm.v1.worker.utils import allocate_kv_cache, select_common_block_size
+
+    small = U.get_kv_cache_config_from_groups(vllm_config, cfg.kv_cache_groups, 2 * 2**30)
+    kernel = []
+    for g in small.kv_cache_groups:
+        s = g.kv_cache_spec
+        if g.layer_names and g.layer_names[0] in draft_names and draft_backend is not None:
+            kernel.append(select_common_block_size(s.block_size, [draft_backend]))
+        else:
+            kernel.append(s.block_size)
+    lay = cc.get_resolved_kv_cache_layout()
+    layout_name = lay.name
+    caches = allocate_kv_cache(small, torch.device("cuda"), lay, kernel)
+    dv = caches[draft_names[0]]
+    print(f"== GPU alloc ({small.num_blocks} blocks, layout {layout_name}, kernel blocks {kernel}): "
+          f"draft view shape={tuple(dv.shape)} dtype={dv.dtype} strides={dv.stride()}")
+    # Draft regions must not alias anything else: poison one and check the rest.
+    dv.view(torch.uint8).fill_(0xAB)
+    for name, view in caches.items():
+        if name in draft_names:
+            continue
+        v = view.view(torch.uint8) if view.dtype != torch.uint8 else view
+        assert int(v.max()) == 0, f"{name} aliases the draft region"
+    other = caches[draft_names[1]].view(torch.uint8)
+    assert int(other.max()) == 0, "draft layers alias each other"
+    print("   draft region is unshared (poison did not leak into any other view)")
+    del caches
+print("== ALL CHECKS PASSED")

docker-images/glm53-kindling/dev/patch-tests/_glm53_recoverssm_boundary_test.py

@@ -0,0 +1,289 @@
+#!/usr/bin/env python3
+"""KDA RecoverSSM align-mode commit with a block table filled only as far as
+align mode allocates it.
+
+Align mode allocates cdiv(num_computed + query_len, block_size) KDA blocks per
+request. The runner's gathered block table keeps whatever an earlier request
+left past that. A commit whose accepted tokens end exactly on a block boundary
+used to store the request's state in column n // block_size, one past the
+allocation, and the state was lost (2026-09-29).
+
+One commit over rows that each end in a different place relative to a
+boundary (block_size 16, up to 7 drafts):
+
+  exact, all accepted      nc 24 + 8 of 8   = 32   (the bug)
+  exact, all accepted, k=2 nc 29 + 3 of 3   = 32   (the bug)
+  exact, partial           nc 26 + 6 of 8   = 32
+  exact, partial, k=2      nc 30 + 2 of 3   = 32
+  past a boundary          nc 28 + 8 of 8   = 36
+  mid-block                nc 17 + 3 of 8   = 20
+  from a boundary          nc 32 + 8 of 8   = 40
+
+Each row's columns past cdiv(nc + q, 16) point at unallocated blocks, or at
+the null block (0) in the second pass. Checks, bit for bit against the stock
+fused_recurrent_kda:
+  - the state after the last accepted token is in the allocated block holding
+    that token, with the conv window at column 0
+  - a row that reaches a boundary has the state after the boundary token in
+    the boundary block
+  - no other block changed
+  - record_state_blocks points the running state column at that block
+
+Fails on the old column math (n // block_size) and passes on the new one.
+
+Needs a CUDA device and the glm53 image's vLLM, no weights. The module is
+found as in _glm53_recoverssm_test.py. On a GB10 that is not serving:
+
+    mkdir -p /tmp/rs && cp dev/patch-tests/_glm53_recoverssm_boundary_test.py \\
+        experimental/fixes/recoverssm.py /tmp/rs/
+    sudo docker run --rm --gpus all --ipc=host -v /tmp/rs:/rs -w / \\
+        --entrypoint python3 <glm53 image> /rs/_glm53_recoverssm_boundary_test.py
+
+It allocates about 250 MB of GPU memory.
+"""
+import importlib.util
+import os
+import sys
+from pathlib import Path
+from types import SimpleNamespace
+
+import torch
+
+from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first
+from vllm.models.glm5next.nvidia.ops.third_party.kda import fused_recurrent_kda
+from vllm.v1.attention.backends.recoverssm_metadata import (
+    RecoverSSMPostprocessMetadata,
+)
+
+
+def _load_module():
+    name = "vllm.models.glm5next.common.recoverssm"
+    try:
+        return importlib.import_module(name)
+    except ImportError:
+        pass
+    here = Path(__file__).resolve().parent
+    candidates = [
+        os.environ.get("RECOVERSSM_PY"),
+        here / "recoverssm.py",
+        here.parent.parent / "experimental" / "fixes" / "recoverssm.py",
+    ]
+    for path in candidates:
+        if path and Path(path).is_file():
+            spec = importlib.util.spec_from_file_location(name, path)
+            module = importlib.util.module_from_spec(spec)
+            sys.modules[name] = module
+            spec.loader.exec_module(module)
+            print(f"recoverssm module: {path}")
+            return module
+    raise SystemExit("recoverssm.py not found; set RECOVERSSM_PY")
+
+
+rs = _load_module()
+
+dev = torch.device("cuda")
+H, D = 16, 128  # per rank at TP=4
+PROJ = H * D
+CONV_DIM = 3 * PROJ
+WIDTH = 4
+HIST = WIDTH - 1
+LOWER_BOUND = -5.0
+DTYPE = torch.bfloat16
+K = 7
+Q = K + 1
+BS = 16  # mamba block size
+W = 6  # block table width
+NUM_LAYERS = 2
+DS = is_conv_state_dim_first()
+
+# (name, num_computed, query_len, num_sampled)
+CASES = [
+    ("exact, all accepted", 24, 8, 8),
+    ("exact, all accepted, k=2", 29, 3, 3),
+    ("exact, partial", 26, 8, 6),
+    ("exact, partial, k=2", 30, 3, 2),
+    ("past a boundary", 28, 8, 8),
+    ("mid-block", 17, 8, 3),
+    ("from a boundary", 32, 8, 8),
+]
+ROWS = len(CASES)
+LINES = 1 + ROWS * W
+
+
+def cdiv(a: int, b: int) -> int:
+    return -(-a // b)
+
+
+def block_id(row: int, col: int) -> int:
+    return 1 + row * W + col
+
+
+def conv_pool(lines: int) -> tuple[torch.Tensor, torch.Tensor]:
+    """(raw kv_cache[0], its (lines, dim, state_len) view)."""
+    state_len = HIST + K
+    if DS:
+        raw = torch.zeros(lines, CONV_DIM, state_len, device=dev, dtype=DTYPE)
+        return raw, raw
+    raw = torch.zeros(lines, state_len, CONV_DIM, device=dev, dtype=DTYPE)
+    return raw, raw.transpose(-1, -2)
+
+
+def make_layer(g: torch.Generator):
+    layer = SimpleNamespace()
+    layer.A_log = torch.log(
+        torch.rand(1, 1, H, 1, device=dev, generator=g) * 8 + 1
+    ).float()
+    layer.dt_bias = (torch.randn(PROJ, device=dev, generator=g) * 0.5).float()
+    layer.kda_lower_bound = LOWER_BOUND
+    conv_raw, layer.conv = conv_pool(LINES)
+    layer.conv.copy_(torch.randn(layer.conv.shape, device=dev, generator=g).to(DTYPE))
+    layer.rec = torch.randn(LINES, H, D, D, device=dev, generator=g) * 0.1
+    layer.kv_cache = (conv_raw, layer.rec)
+    layer.recoverssm_records = rs.allocate_records(ROWS, H, D, Q, DTYPE, dev)
+    return layer
+
+
+def run(unallocated_fill: str, seed: int) -> list[str]:
+    g = torch.Generator(device=dev).manual_seed(seed)
+    layers = [make_layer(g) for _ in range(NUM_LAYERS)]
+    ctx = rs.Glm5NextRecoverSSMCommitContext.create(
+        layers, spec_query_len=Q, max_num_reqs=ROWS
+    )
+
+    block_table = torch.zeros(ROWS, W, dtype=torch.int32)
+    source = []
+    for i, (_, nc, q, _) in enumerate(CASES):
+        allocated = cdiv(nc + q, BS)
+        assert allocated < W
+        for c in range(W):
+            if c < allocated or unallocated_fill == "stale":
+                block_table[i, c] = block_id(i, c)
+        # The verify reads the block the runner's pre-step copy moved the
+        # state to: (seq_len - 1) // block_size.
+        source.append(block_id(i, (nc + q - 1) // BS))
+    block_table = block_table.to(dev)
+    idx = torch.tensor(source, dtype=torch.int32, device=dev)
+    qlens = [q for _, _, q, _ in CASES]
+    qsl_cpu = torch.zeros(ROWS + 1, dtype=torch.int32)
+    qsl_cpu[1:] = torch.cumsum(torch.tensor(qlens), 0)
+    qsl = qsl_cpu.to(dev)
+    T = int(qsl_cpu[-1])
+    num_computed = torch.tensor([nc for _, nc, _, _ in CASES], dtype=torch.int32, device=dev)
+    num_sampled = torch.tensor([n for _, _, _, n in CASES], dtype=torch.int32, device=dev)
+
+    # Stock reference: row i's state after token t goes to scratch slot
+    # 1 + i * Q + t, starting from its checkpoint in slot 1 + i * Q.
+    ref_cols = (
+        1 + torch.arange(ROWS * Q, dtype=torch.int32, device=dev)
+    ).view(ROWS, Q)
+    ones = torch.ones(ROWS, dtype=torch.int32, device=dev)
+
+    refs, before = [], []
+    for layer in layers:
+        qkv = [torch.randn(1, T, H, D, device=dev, generator=g).to(DTYPE) for _ in range(3)]
+        g1 = torch.randn(1, T, H, D, device=dev, generator=g).to(DTYPE)
+        beta = torch.randn(1, T, H, device=dev, generator=g).to(DTYPE)
+        scratch = torch.zeros(1 + ROWS * Q, H, D, D, device=dev, dtype=torch.float32)
+        scratch[ref_cols[:, 0].long()] = layer.rec[idx.long()]
+        fused_recurrent_kda(
+            q=qkv[0], k=qkv[1], v=qkv[2], g=g1, beta=beta, initial_state=scratch,
+            use_qk_l2norm_in_kernel=True, cu_seqlens=qsl, ssm_state_indices=ref_cols,
+            num_accepted_tokens=ones, sigmoid_beta=True, a_log=layer.A_log,
+            g_bias=layer.dt_bias, compute_gate=True, lower_bound=LOWER_BOUND,
+        )
+        refs.append(scratch)
+        rs.kda_recoverssm_verify(
+            q=qkv[0], k=qkv[1], v=qkv[2], g=g1, beta=beta, a_log=layer.A_log,
+            g_bias=layer.dt_bias, lower_bound=LOWER_BOUND,
+            checkpoint_state=layer.rec, records=layer.recoverssm_records,
+            query_start_loc=qsl, state_indices=idx,
+        )
+        # The conv state as the verify's conv update leaves it is random here:
+        # the commit only moves windows, and the checks read them from this copy.
+        before.append((layer.rec.clone(), layer.conv.clone()))
+
+    ctx.commit(
+        num_sampled, idx, qsl,
+        block_table=block_table, num_computed_tokens=num_computed, mamba_block_size=BS,
+    )
+
+    state_idx = torch.zeros(ROWS, dtype=torch.int32, device=dev)
+    num_accepted = torch.full((ROWS,), 9, dtype=torch.int32, device=dev)
+    rs.record_state_blocks(
+        RecoverSSMPostprocessMetadata(
+            num_spec_decodes=ROWS,
+            request_indices=None,
+            block_table=block_table,
+            num_computed_tokens=num_computed,
+            block_size=BS,
+        ),
+        torch.arange(ROWS, dtype=torch.int32, device=dev),
+        num_sampled,
+        state_idx,
+        num_accepted,
+    )
+    torch.cuda.synchronize()
+
+    failures = []
+
+    def fail(msg: str) -> None:
+        failures.append(f"[{unallocated_fill} seed={seed}] {msg}")
+
+    written = set()
+    for i, (name, nc, q, n) in enumerate(CASES):
+        final_col = (nc + n - 1) // BS
+        if final_col >= cdiv(nc + q, BS):
+            fail(f"{name}: final column {final_col} is not allocated")
+            continue
+        final_b = block_id(i, final_col)
+        written.add(final_b)
+        next_boundary = (nc // BS + 1) * BS
+        boundary_b = None
+        if nc + n >= next_boundary:
+            boundary_b = block_id(i, next_boundary // BS - 1)
+            written.add(boundary_b)
+        src = source[i]
+        for li, layer in enumerate(layers):
+            ref = refs[li]
+            _, conv0 = before[li]
+            if not torch.equal(layer.rec[final_b], ref[int(ref_cols[i, n - 1])]):
+                fail(f"{name} layer {li}: state after token {nc + n} not in block "
+                     f"{final_b} (column {final_col})")
+            if not torch.equal(layer.conv[final_b, :, :HIST], conv0[src, :, n - 1 : n - 1 + HIST]):
+                fail(f"{name} layer {li}: conv window not in block {final_b}")
+            if boundary_b is not None:
+                blen = next_boundary - nc
+                if not torch.equal(layer.rec[boundary_b], ref[int(ref_cols[i, blen - 1])]):
+                    fail(f"{name} layer {li}: boundary state after token "
+                         f"{next_boundary} not in block {boundary_b}")
+                if not torch.equal(
+                    layer.conv[boundary_b, :, :HIST], conv0[src, :, blen - 1 : blen - 1 + HIST]
+                ):
+                    fail(f"{name} layer {li}: boundary conv window not in block {boundary_b}")
+        if int(state_idx[i]) != final_col:
+            fail(f"{name}: recorded state column {int(state_idx[i])}, want {final_col}")
+        if int(num_accepted[i]) != 1:
+            fail(f"{name}: num_accepted {int(num_accepted[i])}, want 1")
+
+    for li, layer in enumerate(layers):
+        rec0, conv0 = before[li]
+        for b in range(1, LINES):
+            if b in written:
+                continue
+            if not torch.equal(layer.rec[b], rec0[b]):
+                fail(f"layer {li}: recurrent state of block {b} changed")
+            if not torch.equal(layer.conv[b], conv0[b]):
+                fail(f"layer {li}: conv state of block {b} changed")
+    return failures
+
+
+failures = []
+for fill in ("stale", "null"):
+    for seed in (0, 1):
+        failures += run(fill, seed)
+for msg in failures[:20]:
+    print("  FAIL", msg)
+if failures:
+    print(f"RECOVERSSM BOUNDARY: FAIL ({len(failures)} failures)")
+    sys.exit(1)
+print(f"RECOVERSSM BOUNDARY: PASS ({len(CASES)} cases x 2 fills x 2 seeds)")

docker-images/glm53-kindling/dev/patch-tests/_glm53_recoverssm_test.py

@@ -0,0 +1,474 @@
+#!/usr/bin/env python3
+"""KDA RecoverSSM (experimental/fixes/recoverssm.py) against the stock spec path.
+
+Drives both through the same simulated speculative decode at GLM-5.3-Flash's
+per-rank KDA shapes (TP=4: 16 heads x 128, conv width 4), several layers per
+KV group, with batch-level k varying step to step (2, 3, 4, 5, 7) and each
+row accepting 0..k drafts.
+
+  stock:      causal_conv1d_update(num_accepted_tokens = last step's count)
+              + fused_recurrent_kda with 1 + k state slots per request
+  recoverssm: the same conv call with num_accepted_tokens = 1
+              + kda_recoverssm_verify from one checkpoint block
+              + Glm5NextRecoverSSMCommitContext.commit(num_sampled)
+
+Every step, per layer, bit for bit:
+  - conv output and recurrent output for every verified token
+  - the checkpoint is unchanged by the verify
+  - after the commit: the recurrent state equals the stock slot n - 1, and the
+    conv window at column 0 equals the stock window at column n - 1
+  - align mode (prefix caching): the state is committed to the block holding
+    the last committed token, and a step that crosses a block boundary also
+    leaves the stock state after the boundary token in the boundary block
+Runs with and without request_indices (spec rows interleaved with other rows)
+and with a trailing cudagraph padding row.
+
+Pass: everything bit-identical. The summary counts checks and failures per
+check class, with the first failing step (later failures in a class can be
+fallout: the two paths then carry different states). RECOVERSSM_TEST_ATOL=<x>
+accepts differences up to x and reports them (diagnosis only).
+--drift N runs one chained config for N steps and prints how far the
+committed state and the outputs move from the stock path's.
+
+No weights needed. The module is found as vllm.models.glm5next.common.recoverssm
+when recoverssm.yaml mounts it, else RECOVERSSM_PY, else next to this file,
+else ../../experimental/fixes/recoverssm.py. On one GB10:
+
+    mkdir -p /tmp/rs && cp dev/patch-tests/_glm53_recoverssm_test.py \\
+        experimental/fixes/recoverssm.py /tmp/rs/
+    sudo docker run --rm --gpus all --ipc=host -v /tmp/rs:/rs -w / \\
+        --entrypoint python3 <glm53 image> /rs/_glm53_recoverssm_test.py [--bench]
+
+or `docker cp` both files into a glm53 container and `docker exec` python3 on
+the test. It allocates up to ~0.7 GB of GPU memory: run it on a box that is
+not serving, since a serving GB10 has 1.5-3 GiB free.
+--bench adds a timing of the stock spec kernel against verify + commit at
+34 layers.
+"""
+import importlib.util
+import os
+import sys
+from pathlib import Path
+from types import SimpleNamespace
+
+import torch
+
+from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first
+from vllm.model_executor.layers.mamba.ops.causal_conv1d import causal_conv1d_update
+from vllm.models.glm5next.nvidia.ops.third_party.kda import fused_recurrent_kda
+
+
+def _load_module():
+    name = "vllm.models.glm5next.common.recoverssm"
+    try:
+        return importlib.import_module(name)
+    except ImportError:
+        pass
+    here = Path(__file__).resolve().parent
+    candidates = [
+        os.environ.get("RECOVERSSM_PY"),
+        here / "recoverssm.py",
+        here.parent.parent / "experimental" / "fixes" / "recoverssm.py",
+    ]
+    for path in candidates:
+        if path and Path(path).is_file():
+            spec = importlib.util.spec_from_file_location(name, path)
+            module = importlib.util.module_from_spec(spec)
+            sys.modules[name] = module
+            spec.loader.exec_module(module)
+            print(f"recoverssm module: {path}")
+            return module
+    raise SystemExit("recoverssm.py not found; set RECOVERSSM_PY")
+
+
+rs = _load_module()
+
+dev = torch.device("cuda")
+H, D = 16, 128  # per rank at TP=4 (64 heads)
+PROJ = H * D
+CONV_DIM = 3 * PROJ
+WIDTH = 4
+LOWER_BOUND = -5.0
+DTYPE = torch.bfloat16
+LEVELS = (2, 3, 4, 5, 7)
+ATOL = float(os.environ.get("RECOVERSSM_TEST_ATOL", "0"))
+DS = is_conv_state_dim_first()
+
+
+CLASSES = (
+    "conv output",
+    "recurrent output",
+    "checkpoint untouched",
+    "committed state",
+    "committed conv",
+    "boundary state",
+    "boundary conv",
+)
+
+
+class Checker:
+    def __init__(self) -> None:
+        self.checks = dict.fromkeys(CLASSES, 0)
+        self.failures = dict.fromkeys(CLASSES, 0)
+        self.inexact = dict.fromkeys(CLASSES, 0)
+        self.worst = dict.fromkeys(CLASSES, 0.0)
+        self.first: dict[str, str] = {}
+        self.printed = 0
+
+    def eq(self, cls: str, where: str, got: torch.Tensor, want: torch.Tensor) -> None:
+        self.checks[cls] += 1
+        if torch.equal(got, want):
+            return
+        diff = (got.float() - want.float()).abs()
+        err = float(diff.nan_to_num(float("inf")).max())
+        self.worst[cls] = max(self.worst[cls], err)
+        if err <= ATOL:
+            self.inexact[cls] += 1
+            return
+        self.failures[cls] += 1
+        self.first.setdefault(cls, where)
+        if self.printed < 10:
+            self.printed += 1
+            scale = float(want.float().abs().max())
+            print(f"  MISMATCH {cls} at {where}: max abs {err:.3g} (scale "
+                  f"{scale:.3g}), {int((diff > 0).sum())} of {diff.numel()} differ")
+
+    def report(self) -> None:
+        print(f"{'class':<22}{'checks':>8}{'fail':>8}{'inexact':>9}{'worst':>11}  first failure")
+        for cls in CLASSES:
+            print(f"{cls:<22}{self.checks[cls]:>8}{self.failures[cls]:>8}"
+                  f"{self.inexact[cls]:>9}{self.worst[cls]:>11.3g}  {self.first.get(cls, '-')}")
+
+
+def conv_pool(lines: int, state_len: int) -> tuple[torch.Tensor, torch.Tensor]:
+    """(raw kv_cache[0] tensor, its (lines, dim, state_len) view for the kernels)."""
+    if DS:
+        raw = torch.zeros(lines, CONV_DIM, state_len, device=dev, dtype=DTYPE)
+        return raw, raw
+    raw = torch.zeros(lines, state_len, CONV_DIM, device=dev, dtype=DTYPE)
+    return raw, raw.transpose(-1, -2)
+
+
+def make_layer(g: torch.Generator, lines_stock: int, lines_rs: int, k: int, rows: int):
+    state_len = WIDTH - 1 + k
+    layer = SimpleNamespace()
+    layer.conv_w = (torch.randn(CONV_DIM, WIDTH, device=dev, generator=g) * 0.3).float()
+    layer.A_log = torch.log(
+        torch.rand(1, 1, H, 1, device=dev, generator=g) * 8 + 1
+    ).float()
+    layer.dt_bias = (torch.randn(PROJ, device=dev, generator=g) * 0.5).float()
+    layer.kda_lower_bound = LOWER_BOUND
+    layer.conv_s_raw, layer.conv_s = conv_pool(lines_stock, state_len)
+    layer.rec_s = torch.zeros(lines_stock, H, D, D, device=dev, dtype=torch.float32)
+    conv_r_raw, layer.conv_r = conv_pool(lines_rs, state_len)
+    layer.rec_r = torch.zeros(lines_rs, H, D, D, device=dev, dtype=torch.float32)
+    layer.kv_cache = (conv_r_raw, layer.rec_r)
+    layer.recoverssm_records = rs.allocate_records(rows, H, D, k + 1, DTYPE, dev)
+    return layer
+
+
+def stock_step(layer, x, g1, beta, qsl, cols, na_prev, k):
+    y = causal_conv1d_update(
+        x.clone(), layer.conv_s, layer.conv_w, None, activation="silu",
+        conv_state_indices=cols[:, 0], num_accepted_tokens=na_prev,
+        query_start_loc=qsl, max_query_len=k + 1,
+    )
+    q, kk, v = y.split(PROJ, dim=-1)
+    o, _ = fused_recurrent_kda(
+        q=q.reshape(1, -1, H, D), k=kk.reshape(1, -1, H, D), v=v.reshape(1, -1, H, D),
+        g=g1, beta=beta, initial_state=layer.rec_s, use_qk_l2norm_in_kernel=True,
+        cu_seqlens=qsl, ssm_state_indices=cols, num_accepted_tokens=na_prev,
+        sigmoid_beta=True, a_log=layer.A_log, g_bias=layer.dt_bias,
+        compute_gate=True, lower_bound=LOWER_BOUND,
+    )
+    return y, o
+
+
+def rs_step(layer, x, g1, beta, qsl, idx, ones, k):
+    y = causal_conv1d_update(
+        x.clone(), layer.conv_r, layer.conv_w, None, activation="silu",
+        conv_state_indices=idx, num_accepted_tokens=ones,
+        query_start_loc=qsl, max_query_len=k + 1,
+    )
+    q, kk, v = y.split(PROJ, dim=-1)
+    o = rs.kda_recoverssm_verify(
+        q=q.reshape(1, -1, H, D), k=kk.reshape(1, -1, H, D), v=v.reshape(1, -1, H, D),
+        g=g1, beta=beta, a_log=layer.A_log, g_bias=layer.dt_bias,
+        lower_bound=LOWER_BOUND, checkpoint_state=layer.rec_r,
+        records=layer.recoverssm_records, query_start_loc=qsl, state_indices=idx,
+    )
+    return y, o
+
+
+def run(k: int, align: bool, seed: int, chk: Checker, steps: int = 40,
+        num_reqs: int = 6, num_layers: int = 3, block_size: int = 16,
+        drift: list | None = None) -> None:
+    g = torch.Generator(device=dev).manual_seed(seed)
+    cpu = torch.Generator().manual_seed(seed)
+    B = num_reqs
+    rows = B + 1  # one trailing padding row, as a FULL-graph batch has
+    width = 2 + (steps * (k + 1)) // block_size + 2
+    lines_stock = 1 + B * (k + 1)
+    lines_rs = 1 + B * (width if align else 1)
+    layers = [make_layer(g, lines_stock, lines_rs, k, rows) for _ in range(num_layers)]
+    ctx = rs.Glm5NextRecoverSSMCommitContext.create(
+        layers, spec_query_len=k + 1, max_num_reqs=rows
+    )
+
+    stock_cols = torch.zeros(rows, k + 1, dtype=torch.int32)
+    for i in range(B):
+        stock_cols[i] = torch.arange(1 + i * (k + 1), 1 + (i + 1) * (k + 1))
+    stock_cols = stock_cols.to(dev)
+    if align:
+        block_table = torch.zeros(rows, width, dtype=torch.int32)
+        for i in range(B):
+            block_table[i] = torch.arange(1 + i * width, 1 + (i + 1) * width)
+        block_table = block_table.to(dev)
+        num_computed = torch.randint(1, 2 * block_size, (B,), generator=cpu)
+        cur_col = ((num_computed - 1) // block_size).tolist()
+    else:
+        block_table = None
+        num_computed = torch.zeros(B, dtype=torch.int64)
+        cur_col = [0] * B
+
+    def rs_block(i: int, col: int) -> int:
+        return int(block_table[i, col]) if align else 1 + i
+
+    # Same random checkpoint in both paths: stock slot 0, read with na = 1.
+    for layer in layers:
+        for i in range(B):
+            state = torch.randn(H, D, D, device=dev, generator=g) * 0.1
+            hist = torch.randn(CONV_DIM, WIDTH - 1, device=dev, generator=g).to(DTYPE)
+            layer.rec_s[int(stock_cols[i, 0])] = state
+            layer.conv_s[int(stock_cols[i, 0]), :, : WIDTH - 1] = hist
+            b = rs_block(i, cur_col[i])
+            layer.rec_r[b] = state
+            layer.conv_r[b, :, : WIDTH - 1] = hist
+    na_prev = torch.ones(rows, dtype=torch.int32, device=dev)
+    ones = torch.ones(rows, dtype=torch.int32, device=dev)
+
+    for step in range(steps):
+        k_step = min(LEVELS[int(torch.randint(len(LEVELS), (1,), generator=cpu))], k)
+        # Most rows verify k_step drafts. About one in five has no drafts.
+        qlens = [
+            1 if torch.rand(1, generator=cpu) < 0.2 else k_step + 1 for _ in range(B)
+        ] + [0]
+        qsl_cpu = torch.zeros(rows + 1, dtype=torch.int32)
+        qsl_cpu[1:] = torch.cumsum(torch.tensor(qlens), 0)
+        qsl = qsl_cpu.to(dev)
+        T = int(qsl_cpu[-1])
+        x = torch.randn(T, CONV_DIM, device=dev, generator=g).to(DTYPE)
+        g1 = torch.randn(1, T, H, D, device=dev, generator=g).to(DTYPE)
+        # beta as a column slice of a wider projection, like kda.py's (its row
+        # is 6416 wide: a multiple of 16, as here, which Triton specializes on).
+        beta = torch.randn(1, T, H + 16, device=dev, generator=g).to(DTYPE)[..., :H]
+
+        # Runner align preprocess: the running block moves to the column the
+        # verify reads, (seq_len - 1) // block_size.
+        if align:
+            for i in range(B):
+                dst = (int(num_computed[i]) + qlens[i] - 1) // block_size
+                if dst != cur_col[i]:
+                    src_b, dst_b = rs_block(i, cur_col[i]), rs_block(i, dst)
+                    for layer in layers:
+                        layer.rec_r[dst_b] = layer.rec_r[src_b]
+                        layer.conv_r[dst_b] = layer.conv_r[src_b]
+                    cur_col[i] = dst
+        idx = torch.tensor(
+            [rs_block(i, cur_col[i]) for i in range(B)] + [0],
+            dtype=torch.int32, device=dev,
+        )
+
+        for li, layer in enumerate(layers):
+            before = layer.rec_r[idx[:B].long()].clone()
+            y_s, o_s = stock_step(layer, x, g1, beta, qsl, stock_cols, na_prev, k)
+            y_r, o_r = rs_step(layer, x, g1, beta, qsl, idx, ones, k)
+            tag = f"k={k} align={align} seed={seed} step={step} layer={li}"
+            chk.eq("conv output", tag, y_r[:T], y_s[:T])
+            chk.eq("recurrent output", tag, o_r[:, :T], o_s[:, :T])
+            chk.eq("checkpoint untouched", tag, layer.rec_r[idx[:B].long()], before)
+            if drift is not None:
+                out_err = float((o_r[:, :T].float() - o_s[:, :T].float()).abs().max())
+
+        accepted = [int(torch.randint(0, q, (1,), generator=cpu)) for q in qlens[:B]]
+        n_sampled = [a + 1 for a in accepted]
+
+        # Odd steps: spec rows interleaved with other batch rows.
+        interleave = step % 2 == 1
+        if interleave:
+            batch_rows = [2 * i + 1 for i in range(B)]
+            nb = 2 * B + 1
+            request_indices = torch.tensor(batch_rows, dtype=torch.int32, device=dev)
+        else:
+            batch_rows = list(range(B))
+            nb = B
+            request_indices = None
+        num_sampled = torch.randint(0, 9, (nb,), generator=cpu, dtype=torch.int32)
+        nc_batch = torch.randint(0, 99, (nb,), generator=cpu, dtype=torch.int32)
+        bt_batch = torch.zeros(nb, width if align else 1, dtype=torch.int32)
+        for i, r in enumerate(batch_rows):
+            num_sampled[r] = n_sampled[i]
+            nc_batch[r] = int(num_computed[i])
+            if align:
+                bt_batch[r] = block_table[i].cpu()
+        align_kwargs = {}
+        if align:
+            align_kwargs = dict(
+                block_table=bt_batch.to(dev),
+                num_computed_tokens=nc_batch.to(dev),
+                mamba_block_size=block_size,
+            )
+        ctx.commit(
+            num_sampled.to(dev), idx[:B], qsl[: B + 1],
+            request_indices=request_indices, **align_kwargs,
+        )
+
+        state_err = 0.0
+        for i in range(B):
+            n = n_sampled[i]
+            nc = int(num_computed[i])
+            # The block holding the last committed token.
+            final_col = (nc + n - 1) // block_size if align else 0
+            final_b = rs_block(i, final_col)
+            for li, layer in enumerate(layers):
+                tag = f"k={k} align={align} seed={seed} step={step} layer={li} req={i} n={n}"
+                s_col = int(stock_cols[i, n - 1])
+                c_col = int(stock_cols[i, 0])
+                chk.eq("committed state", tag, layer.rec_r[final_b], layer.rec_s[s_col])
+                chk.eq("committed conv", tag,
+                       layer.conv_r[final_b, :, : WIDTH - 1],
+                       layer.conv_s[c_col, :, n - 1 : n - 1 + WIDTH - 1])
+                if drift is not None:
+                    want = layer.rec_s[s_col]
+                    state_err = max(state_err, float(
+                        (layer.rec_r[final_b] - want).abs().max() / want.abs().max()))
+                next_boundary = (nc // block_size + 1) * block_size
+                if align and nc + n >= next_boundary:
+                    blen = next_boundary - nc
+                    bb = rs_block(i, next_boundary // block_size - 1)
+                    chk.eq("boundary state", tag,
+                           layer.rec_r[bb], layer.rec_s[int(stock_cols[i, blen - 1])])
+                    chk.eq("boundary conv", tag,
+                           layer.conv_r[bb, :, : WIDTH - 1],
+                           layer.conv_s[c_col, :, blen - 1 : blen - 1 + WIDTH - 1])
+            if align:
+                num_computed[i] = nc + n
+                cur_col[i] = final_col
+        na_prev = torch.tensor(n_sampled + [1], dtype=torch.int32, device=dev)
+        if drift is not None:
+            drift.append((step, state_err, out_err))
+    torch.cuda.synchronize()
+
+
+def drift_run(steps: int) -> None:
+    """Chained, never resynced: how far RecoverSSM moves from the stock path."""
+    samples: list = []
+    run(7, False, 3, Checker(), steps=steps, num_layers=1, drift=samples)
+    every = max(1, steps // 10)
+    print(f"drift over {steps} steps (k=7, 6 requests; state error relative to "
+          "the state's max, output error absolute):")
+    worst_state = worst_out = 0.0
+    for step, state_err, out_err in samples:
+        worst_state = max(worst_state, state_err)
+        worst_out = max(worst_out, out_err)
+        if (step + 1) % every == 0:
+            print(f"  step {step + 1:>6}: state {state_err:.3g} (worst so far "
+                  f"{worst_state:.3g}), last-layer output {out_err:.3g} "
+                  f"(worst {worst_out:.3g})")
+
+
+def bench(num_reqs: int, k: int = 7, num_layers: int = 34, iters: int = 20) -> None:
+    """Stock spec kernel vs verify + commit, all rows verifying k drafts.
+
+    Layers share one pool here (timing only), so the numbers are per step for
+    num_layers layers of one KV group.
+    """
+    g = torch.Generator(device=dev).manual_seed(0)
+    q_len = k + 1
+    rows = num_reqs
+    layer = make_layer(g, 1 + rows * q_len, 1 + rows, k, rows)
+    layer.rec_s.normal_(generator=g).mul_(0.1)
+    layer.rec_r.normal_(generator=g).mul_(0.1)
+    layers = [layer] * num_layers
+    ctx = rs.Glm5NextRecoverSSMCommitContext.create(
+        layers, spec_query_len=q_len, max_num_reqs=rows
+    )
+    T = rows * q_len
+    qsl = torch.arange(0, T + 1, q_len, dtype=torch.int32, device=dev)
+    cols = (1 + torch.arange(rows * q_len, dtype=torch.int32, device=dev)).view(rows, q_len)
+    idx = 1 + torch.arange(rows, dtype=torch.int32, device=dev)
+    na = torch.full((rows,), q_len, dtype=torch.int32, device=dev)
+    y = torch.randn(T, CONV_DIM, device=dev, generator=g).to(DTYPE)
+    q, kk, v = (t.reshape(1, -1, H, D) for t in y.split(PROJ, dim=-1))
+    g1 = torch.randn(1, T, H, D, device=dev, generator=g).to(DTYPE)
+    beta = torch.randn(1, T, H, device=dev, generator=g).to(DTYPE)
+    out = torch.empty(1, T, H, D, device=dev, dtype=DTYPE)
+    sampled = torch.full((rows,), q_len, dtype=torch.int32, device=dev)
+
+    def stock():
+        for _ in range(num_layers):
+            fused_recurrent_kda(
+                q=q, k=kk, v=v, g=g1, beta=beta, initial_state=layer.rec_s,
+                use_qk_l2norm_in_kernel=True, cu_seqlens=qsl, ssm_state_indices=cols,
+                num_accepted_tokens=na, out=out, sigmoid_beta=True, a_log=layer.A_log,
+                g_bias=layer.dt_bias, compute_gate=True, lower_bound=LOWER_BOUND,
+            )
+
+    def recover():
+        for _ in range(num_layers):
+            rs.kda_recoverssm_verify(
+                q=q, k=kk, v=v, g=g1, beta=beta, a_log=layer.A_log,
+                g_bias=layer.dt_bias, lower_bound=LOWER_BOUND,
+                checkpoint_state=layer.rec_r, records=layer.recoverssm_records,
+                query_start_loc=qsl, state_indices=idx, out=out,
+            )
+        ctx.commit(sampled, idx, qsl)
+
+    def verify_only():
+        for _ in range(num_layers):
+            rs.kda_recoverssm_verify(
+                q=q, k=kk, v=v, g=g1, beta=beta, a_log=layer.A_log,
+                g_bias=layer.dt_bias, lower_bound=LOWER_BOUND,
+                checkpoint_state=layer.rec_r, records=layer.recoverssm_records,
+                query_start_loc=qsl, state_indices=idx, out=out,
+            )
+
+    def timed(fn) -> float:
+        fn()
+        torch.cuda.synchronize()
+        start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
+        start.record()
+        for _ in range(iters):
+            fn()
+        end.record()
+        torch.cuda.synchronize()
+        return start.elapsed_time(end) / iters
+
+    t_stock, t_rec, t_ver = timed(stock), timed(recover), timed(verify_only)
+    print(f"bench rows={rows} k={k} layers={num_layers}: stock {t_stock:.2f} ms, "
+          f"verify {t_ver:.2f} ms + commit {t_rec - t_ver:.2f} ms = {t_rec:.2f} ms")
+
+
+chk = Checker()
+for k in (7, 3):
+    for align in (False, True):
+        for seed in (0, 1):
+            run(k, align, seed, chk)
+            print(f"k={k} align={align} seed={seed}: {sum(chk.checks.values())} "
+                  f"checks so far, {sum(chk.failures.values())} failures")
+chk.report()
+if "--bench" in sys.argv:
+    for n in (16, 32, 64):
+        bench(n)
+if "--drift" in sys.argv:
+    drift_run(int(sys.argv[sys.argv.index("--drift") + 1]))
+checks, failures = sum(chk.checks.values()), sum(chk.failures.values())
+inexact = sum(chk.inexact.values())
+if failures:
+    print(f"RECOVERSSM: FAIL ({failures} of {checks} checks)")
+    sys.exit(1)
+if inexact:
+    print(f"RECOVERSSM: PASS within atol {ATOL} ({inexact} of {checks} "
+          "not bit-identical)")
+else:
+    print(f"RECOVERSSM: PASS, bit-identical ({checks} checks)")

docker-images/glm53-kindling/dev/patch-tests/_glm53_tp3_sparse_mla_test.py

@@ -0,0 +1,70 @@
+#!/usr/bin/env python3
+"""experimental/fixes/gb10_sparse_mla.py at head counts that are not powers of two.
+
+TP=3 leaves 22 MLA heads per rank (66 padded / 3). The kernel used to take
+tl.arange(0, H), which needs a power of two; it now tiles at the next power of
+two (>= 16) with a head mask. This checks both launch paths (one program per
+token, and the split + combine path decode takes) against a float32 torch
+reference, at H = 16 and 32 (unchanged tiles) and 22 and 24 (masked), and
+H = 11 (TP=6: 66 / 6, a 16-row tile with 5 rows masked).
+
+Pass: every case within 5e-3 relative error of the reference, and the masked
+cases no worse than twice the H=16 error. On one GB10:
+
+    mkdir -p /tmp/smla && cp dev/patch-tests/_glm53_tp3_sparse_mla_test.py \\
+        experimental/fixes/gb10_sparse_mla.py /tmp/smla/
+    sudo docker run --rm --gpus all -v /tmp/smla:/t -w /t \\
+        --entrypoint python3 <glm53 image> /t/_glm53_tp3_sparse_mla_test.py
+"""
+import importlib.util
+import os
+import sys
+
+import torch
+
+here = os.path.dirname(os.path.abspath(__file__))
+path = next(p for p in (os.path.join(here, "gb10_sparse_mla.py"),
+                        os.path.join(here, "../../experimental/fixes/gb10_sparse_mla.py")) if os.path.exists(p))
+spec = importlib.util.spec_from_file_location("gb10_sparse_mla", path)
+m = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(m)
+
+dev = torch.device("cuda")
+D, S, TOPK = 512, 8192, 512
+
+
+def reference(q, kv, idx, lens, sm_scale, kv_scale):
+    kvf = kv.float() * kv_scale
+    out = torch.empty(q.shape, dtype=torch.float32, device=q.device)
+    for t in range(q.shape[0]):
+        rows = kvf[idx[t, :lens[t]].long()]                      # [n, D]
+        s = (q[t].float() @ rows.T) * sm_scale                    # [H, n]
+        out[t] = torch.softmax(s, -1) @ rows
+    return out
+
+
+def case(H, T, splits, g):
+    q = torch.randn(T, H, D, generator=g, device=dev).to(torch.bfloat16)
+    kv = (torch.randn(S, D, generator=g, device=dev) * 0.5).to(torch.float8_e4m3fn)
+    lens = torch.randint(TOPK // 2, TOPK + 1, (T,), generator=g, device=dev).int()
+    idx = torch.stack([torch.randperm(S, generator=g, device=dev)[:TOPK] for _ in range(T)]).int()
+    for t in range(T):
+        idx[t, lens[t]:] = -1
+    sm_scale, kv_scale = D ** -0.5, 0.75
+    got = m.sparse_mla(q, kv, idx, lens, sm_scale, kv_scale, splits=splits).float()
+    ref = reference(q, kv, idx, lens, sm_scale, kv_scale)
+    return ((got - ref).norm() / ref.norm()).item()
+
+
+g = torch.Generator(device=dev).manual_seed(0)
+errs, fails = {}, 0
+for T, splits in ((64, 1), (4, 8)):
+    for H in (11, 16, 22, 24, 32):
+        e = case(H, T, splits, g)
+        errs[(H, T)] = e
+        print(f"H={H:2d} T={T:2d} splits={splits}: rel err {e:.2e}", flush=True)
+for (H, T), e in errs.items():
+    if e > 5e-3 or e > 2 * max(errs[(16, T)], 1e-4):
+        print(f"FAIL H={H} T={T}: {e:.2e}"); fails += 1
+print("PASS" if not fails else f"{fails} FAILED")
+sys.exit(1 if fails else 0)

docker-images/glm53-kindling/dev/patch-tests/_glm53_tp3pad_test.py

@@ -0,0 +1,103 @@
+#!/usr/bin/env python3
+"""experimental/tp3/tp3pad.py: the load-time zero padding behind TP=3.
+
+Builds a small GLM-5.3-Flash-shaped set of tensors (one MLA layer, one KDA
+layer, routed experts in bf16 and NVFP4 layout, a shared expert) and runs
+pad_weights with a config that pads heads 4->6, KDA heads 4->6 and the MoE
+intermediate 32->48. Checks:
+  - every sized tensor reaches the padded shape, the original block is
+    byte-identical and the padding is zero (uint8 NVFP4 and scales included)
+  - tensors that are not sized by a padded dim pass through untouched
+  - a bf16 expert forward, silu(x Wg) * (x Wu) Wd, is bit-identical padded
+  - with no tp_pad_orig in the config, weights pass through as-is
+CPU only, no vLLM needed (vllm.logger is stubbed when absent):
+
+    python3 dev/patch-tests/_glm53_tp3pad_test.py
+"""
+import importlib.util
+import logging
+import os
+import sys
+import types
+
+import torch
+
+try:
+    import vllm.logger  # noqa: F401
+except Exception:
+    stub = types.ModuleType("vllm.logger")
+    stub.init_logger = logging.getLogger
+    sys.modules.setdefault("vllm", types.ModuleType("vllm"))
+    sys.modules["vllm.logger"] = stub
+
+here = os.path.dirname(os.path.abspath(__file__))
+path = next(p for p in (os.path.join(here, "tp3pad.py"),
+                        os.path.join(here, "../../experimental/tp3/tp3pad.py")) if os.path.exists(p))
+spec = importlib.util.spec_from_file_location("tp3pad", path)
+m = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(m)
+
+H0, H1, K0, K1, I0, I1 = 4, 6, 4, 6, 32, 48
+nope, rope, vd, hd, hid = 8, 4, 8, 16, 64
+cfg = types.SimpleNamespace(num_attention_heads=H1, linear_num_heads=K1, moe_intermediate_size=I1,
+                            qk_nope_head_dim=nope, qk_rope_head_dim=rope, v_head_dim=vd,
+                            linear_head_dim=hd, n_shared_experts=1,
+                            tp_pad_orig={"num_attention_heads": H0, "linear_num_heads": K0,
+                                         "moe_intermediate_size": I0})
+g = torch.Generator().manual_seed(0)
+r = lambda *s: torch.randn(*s, generator=g).to(torch.bfloat16)
+u8 = lambda *s: torch.randint(1, 255, s, generator=g, dtype=torch.uint8)
+L = "model.layers.0"
+W = {  # name: (tensor, padded shape or None if untouched)
+    f"{L}.self_attn.q_b_proj.weight": (r(H0 * (nope + rope), 32), (H1 * (nope + rope), 32)),
+    f"{L}.self_attn.kv_b_proj.weight": (r(H0 * (nope + vd), 16), (H1 * (nope + vd), 16)),
+    f"{L}.self_attn.o_proj.weight": (r(hid, H0 * vd), (hid, H1 * vd)),
+    f"{L}.self_attn.kv_a_proj_with_mqa.weight": (r(16 + rope, hid), None),
+    f"{L}.self_attn.indexer.wq_b.weight": (r(H0 * (nope + rope), 32), None),
+    "model.layers.1.self_attn.q_proj.weight": (r(K0 * hd, hid), (K1 * hd, hid)),
+    "model.layers.1.self_attn.q_conv1d.weight": (r(K0 * hd, 1, 4), (K1 * hd, 1, 4)),
+    "model.layers.1.self_attn.b_proj.weight": (r(K0, hid), (K1, hid)),
+    "model.layers.1.self_attn.A_log": (r(K0), (K1,)),
+    "model.layers.1.self_attn.dt_bias": (r(K0 * hd), (K1 * hd,)),
+    "model.layers.1.self_attn.o_proj.weight": (r(hid, K0 * hd), (hid, K1 * hd)),
+    f"{L}.mlp.experts.0.gate_proj.weight": (r(I0, hid), (I1, hid)),
+    f"{L}.mlp.experts.0.up_proj.weight": (r(I0, hid), (I1, hid)),
+    f"{L}.mlp.experts.0.down_proj.weight": (r(hid, I0), (hid, I1)),
+    f"{L}.mlp.experts.1.gate_proj.weight": (u8(I0, hid // 2), (I1, hid // 2)),
+    f"{L}.mlp.experts.1.gate_proj.weight_scale": (u8(I0, hid // 16), (I1, hid // 16)),
+    f"{L}.mlp.experts.1.down_proj.weight": (u8(hid, I0 // 2), (hid, I1 // 2)),
+    f"{L}.mlp.experts.1.down_proj.weight_scale": (u8(hid, I0 // 16), (hid, I1 // 16)),
+    f"{L}.mlp.experts.1.down_proj.weight_scale_2": (r(1), None),
+    f"{L}.mlp.shared_experts.gate_proj.weight": (r(I0, hid), (I1, hid)),
+    f"{L}.mlp.shared_experts.down_proj.weight": (r(hid, I0), (hid, I1)),
+    f"{L}.mlp.gate.weight": (r(8, hid), None),
+}
+fails = 0
+out = dict(m.pad_weights(((k, v[0]) for k, v in W.items()), cfg))
+for name, (t, shape) in W.items():
+    p = out[name]
+    if shape is None:
+        ok = p is t
+    else:
+        a, b = (x.view(torch.uint8) if x.element_size() == 1 else x for x in (t, p))
+        sl = tuple(slice(0, n) for n in t.shape)
+        pad = b.clone(); pad[sl] = 0
+        ok = tuple(p.shape) == shape and p.dtype == t.dtype and torch.equal(b[sl], a) and not pad.any()
+    if not ok:
+        print(f"FAIL {name}: {tuple(t.shape)} -> {tuple(p.shape)}"); fails += 1
+
+x = r(3, hid).float()
+e = lambda gw, uw, dw: (torch.nn.functional.silu(x @ gw.float().T) * (x @ uw.float().T)) @ dw.float().T
+E = f"{L}.mlp.experts.0."
+before = e(W[E + "gate_proj.weight"][0], W[E + "up_proj.weight"][0], W[E + "down_proj.weight"][0])
+after = e(out[E + "gate_proj.weight"], out[E + "up_proj.weight"], out[E + "down_proj.weight"])
+if not torch.equal(before, after):
+    print("FAIL expert forward differs after padding"); fails += 1
+
+plain = types.SimpleNamespace(**{k: v for k, v in vars(cfg).items() if k != "tp_pad_orig"})
+if any(a is not b for (_, a), b in zip(m.pad_weights(((k, v[0]) for k, v in W.items()), plain),
+                                        (v[0] for v in W.values()))):
+    print("FAIL without tp_pad_orig the weights were changed"); fails += 1
+
+print(f"{len(W)} tensors checked, expert forward, pass-through: " + ("PASS" if not fails else f"{fails} FAILED"))
+sys.exit(1 if fails else 0)

docker-images/glm53-kindling/dev/patch-tests/fixtures/mentat/group-glm53.json

@@ -0,0 +1,151 @@
+{
+  "control_port": 6379,
+  "groups": {
+    "glm53": {
+      "agents": {
+        "glm53@glm53@": {
+          "alive": true,
+          "container": "glm53",
+          "degraded": false,
+          "gone_since_ms": null,
+          "gpus_free": [],
+          "machine": {
+            "cpus": 20,
+            "gpus": [
+              {
+                "index": 0,
+                "memory": 130596028416,
+                "name": "NVIDIA GB10",
+                "uma": true,
+                "vendor": "nvidia"
+              }
+            ],
+            "memory": 130596028416
+          },
+          "node_id": "6d656e7461743a3139322e302e322e39330000000000000000000000",
+          "node_ip": "192.0.2.93",
+          "pid": 2315,
+          "services": {
+            "mcp": {
+              "host": "",
+              "note": "",
+              "path": "/mcp",
+              "port": 8082,
+              "provider": ""
+            },
+            "openai": {
+              "host": "",
+              "note": "",
+              "path": "/v1",
+              "port": 8002,
+              "provider": "vllm"
+            }
+          }
+        },
+        "glm53@glm53@192.0.2.36": {
+          "alive": true,
+          "container": "glm53",
+          "degraded": false,
+          "gone_since_ms": null,
+          "gpus_free": [],
+          "machine": {
+            "cpus": 20,
+            "gpus": [
+              {
+                "index": 0,
+                "memory": 130596028416,
+                "name": "NVIDIA GB10",
+                "uma": true,
+                "vendor": "nvidia"
+              }
+            ],
+            "memory": 130596028416
+          },
+          "node_id": "6d656e7461743a3139322e302e322e33360000000000000000000000",
+          "node_ip": "192.0.2.36",
+          "pid": 1,
+          "services": {
+            "mcp": {
+              "host": "",
+              "note": "",
+              "path": "/mcp",
+              "port": 8082,
+              "provider": ""
+            }
+          }
+        },
+        "glm53@glm53@192.0.2.70": {
+          "alive": true,
+          "container": "glm53",
+          "degraded": false,
+          "gone_since_ms": null,
+          "gpus_free": [],
+          "machine": {
+            "cpus": 20,
+            "gpus": [
+              {
+                "index": 0,
+                "memory": 130596188160,
+                "name": "NVIDIA GB10",
+                "uma": true,
+                "vendor": "nvidia"
+              }
+            ],
+            "memory": 130596188160
+          },
+          "node_id": "6d656e7461743a3139322e302e322e37300000000000000000000000",
+          "node_ip": "192.0.2.70",
+          "pid": 1,
+          "services": {
+            "mcp": {
+              "host": "",
+              "note": "",
+              "path": "/mcp",
+              "port": 8082,
+              "provider": ""
+            }
+          }
+        },
+        "glm53@glm53@192.0.2.77": {
+          "alive": true,
+          "container": "glm53",
+          "degraded": false,
+          "gone_since_ms": null,
+          "gpus_free": [],
+          "machine": {
+            "cpus": 20,
+            "gpus": [
+              {
+                "index": 0,
+                "memory": 130596028416,
+                "name": "NVIDIA GB10",
+                "uma": true,
+                "vendor": "nvidia"
+              }
+            ],
+            "memory": 130596028416
+          },
+          "node_id": "6d656e7461743a3139322e302e322e37370000000000000000000000",
+          "node_ip": "192.0.2.77",
+          "pid": 1,
+          "services": {
+            "mcp": {
+              "host": "",
+              "note": "",
+              "path": "/mcp",
+              "port": 8082,
+              "provider": ""
+            }
+          }
+        }
+      },
+      "gpus_total": 4,
+      "gpus_used": 4
+    }
+  },
+  "head_generation": 1,
+  "head_node_id": "6d656e7461743a3139322e302e322e31323200000000000000000000",
+  "node_id": "6d656e7461743a3139322e302e322e31323200000000000000000000",
+  "node_ip": "192.0.2.122",
+  "proto": "0.99"
+}

docker-images/glm53-kindling/dev/patch-tests/fixtures/mentat/local-36.json

@@ -0,0 +1,26 @@
+{
+  "addr_ifaces": {
+    "192.0.2.36": "enP7s7",
+    "198.18.0.36": "enp1s0f1np1"
+  },
+  "addr_tags": {
+    "192.0.2.36": [
+      "lan"
+    ],
+    "198.18.0.36": [
+      "connectx",
+      "rdma"
+    ]
+  },
+  "addrs": [
+    "192.0.2.36",
+    "198.18.0.36"
+  ],
+  "control_port": 6379,
+  "groups": {},
+  "head_generation": 1,
+  "head_node_id": "6d656e7461743a3139322e302e322e31323200000000000000000000",
+  "node_id": "6d656e7461743a3139322e302e322e33360000000000000000000000",
+  "node_ip": "192.0.2.36",
+  "proto": "0.99"
+}

docker-images/glm53-kindling/dev/patch-tests/fixtures/mentat/local-70.json

@@ -0,0 +1,26 @@
+{
+  "addr_ifaces": {
+    "192.0.2.70": "enP7s7",
+    "198.18.0.1": "enp1s0f0np0"
+  },
+  "addr_tags": {
+    "192.0.2.70": [
+      "lan"
+    ],
+    "198.18.0.1": [
+      "connectx",
+      "rdma"
+    ]
+  },
+  "addrs": [
+    "192.0.2.70",
+    "198.18.0.1"
+  ],
+  "control_port": 6379,
+  "groups": {},
+  "head_generation": 1,
+  "head_node_id": "6d656e7461743a3139322e302e322e31323200000000000000000000",
+  "node_id": "6d656e7461743a3139322e302e322e37300000000000000000000000",
+  "node_ip": "192.0.2.70",
+  "proto": "0.99"
+}

docker-images/glm53-kindling/dev/patch-tests/fixtures/mentat/local-77.json

@@ -0,0 +1,26 @@
+{
+  "addr_ifaces": {
+    "192.0.2.77": "enP7s7",
+    "198.18.0.2": "enp1s0f0np0"
+  },
+  "addr_tags": {
+    "192.0.2.77": [
+      "lan"
+    ],
+    "198.18.0.2": [
+      "connectx",
+      "rdma"
+    ]
+  },
+  "addrs": [
+    "192.0.2.77",
+    "198.18.0.2"
+  ],
+  "control_port": 6379,
+  "groups": {},
+  "head_generation": 1,
+  "head_node_id": "6d656e7461743a3139322e302e322e31323200000000000000000000",
+  "node_id": "6d656e7461743a3139322e302e322e37370000000000000000000000",
+  "node_ip": "192.0.2.77",
+  "proto": "0.99"
+}

docker-images/glm53-kindling/dev/patch-tests/fixtures/mentat/local-93.json

@@ -0,0 +1,26 @@
+{
+  "addr_ifaces": {
+    "192.0.2.93": "enP7s7",
+    "198.18.0.93": "enp1s0f1np1"
+  },
+  "addr_tags": {
+    "192.0.2.93": [
+      "lan"
+    ],
+    "198.18.0.93": [
+      "connectx",
+      "rdma"
+    ]
+  },
+  "addrs": [
+    "192.0.2.93",
+    "198.18.0.93"
+  ],
+  "control_port": 6379,
+  "groups": {},
+  "head_generation": 1,
+  "head_node_id": "6d656e7461743a3139322e302e322e31323200000000000000000000",
+  "node_id": "6d656e7461743a3139322e302e322e39330000000000000000000000",
+  "node_ip": "192.0.2.93",
+  "proto": "0.99"
+}

docker-images/glm53-kindling/dev/repro/README.md

@@ -0,0 +1,130 @@
+# Reproductions for the greedy-decoding corruption
+
+Three scripts, smallest first. See `../TOPK-CORRUPTION.md` for the diagnosis
+these came out of. The two HTTP scripts use the standard library only and take
+`--url` and `--model`, so they can be attached to an issue as they are.
+
+## greedy_nondet.py
+
+Identical greedy requests over a ~42k-token prompt return different
+completions. The prompt is one paragraph repeated to length, so there is
+nothing model-specific or agent-specific in it.
+
+Measured against GLM-5.3-Flash-NVFP4, 4x GB10 (sm121) at TP=4:
+
+    42034 tokens, 16 runs   3 distinct completions
+     2104 tokens, 16 runs   1 distinct completion
+
+Use this one for a determinism issue. It shows the effect and nothing else.
+
+## toolcall_corruption.py
+
+The same nondeterminism landing inside a tool call, which is what reaches a
+user. It rebuilds a 42272-token coding-agent transcript (47 tools, 24
+tool-call/result pairs, generated from fixed seeds) and asks for a Bash call
+repeating a command from earlier in the conversation, so the correct output is
+fixed and every difference is legible.
+
+One run of 40 at temperature 0, same seed, 5 distinct completions:
+
+    30 runs  correct
+     4 runs  pos 23  ' identical' for '>&'       -> pytest -q 2 identical 2>&1
+     3 runs  pos 26  ' pytest'    for ' tail'    -> ... 2>&1 | pytest -q 2>&1 | tail -40
+     2 runs  pos 23  ' Bash'      for '>&'       -> pytest -q 2 Bash 2>&1
+     1 run   pos 31  ' Read'      for 'description'
+
+The last three corrupt a tool name rather than an argument. That is the path to
+the symptom clients actually report: `glm47_moe.py` runs with
+`validate_tool_names=True`, so an unknown name emits zero deltas and the request
+finishes `stop` with no content and no tool calls.
+
+It is 340 lines because shrinking it stops it reproducing. Two reductions were
+measured and both went bit-stable over 40 runs: replacing the generated file
+listings with the command's own output repeated, and dropping to 3 tools with
+one Read/Edit chain. Both make every token of the command high-margin. The
+length matters too, and 2k tokens is stable over 16 runs.
+
+## marlin_moe_nondet.py
+
+`fused_marlin_moe` returns different bits for identical inputs at some M. No
+checkpoint and no server: random NVFP4 weights on one CUDA device, in a
+container that has vLLM importable.
+
+Measured earlier at the GLM-5.3-Flash geometry (288 experts, top-8, hidden
+4096, intermediate 2048), 24 repetitions per M:
+
+    M=3104   23/23 runs differ
+    M=2304    9/39 runs differ
+    M=1536    6/11 runs differ
+    M=3072, 2816, 2272, 2048, 1024, 800   bit-identical
+
+It needs roughly 4 GB free, so it cannot run while the engine holds
+`gpu_memory_utilization=0.88`. It has not been re-run since the numbers above
+were taken, and no measurement yet connects those M values to the ones the
+engine actually issues, so treat it as a separate finding rather than the cause
+of the two scripts above.
+
+# Determinism and prefix-cache scans
+
+Three scripts for checking a server, independent of the corruption above.
+`determinism.py` and `prefix_cache.py` use the standard library only and take
+`--url` and `--model`. Run them with no other traffic on the server: both flag
+requests that finished alongside theirs.
+
+## determinism.py
+
+`prompt` sends the same prompt three times at each of 3k, 6k, 10k, 14k and 40k
+tokens with `prompt_logprobs`, each under a fresh `cache_salt`, and reports per
+pair of runs how many logprobs differ, where the first difference is and the
+largest. `generate` does the same for 48 greedy tokens of one 5208-token prompt,
+six times.
+
+On stock kernels the runs differ from the first token (the marlin MoE above).
+With the MoE overrides the remaining differences came from the sparse
+indexer's top-k: it returned the selected pools in a run-dependent order, and
+where two pools tie at the k-th score it kept a run-dependent one of them. Both
+start at a fixed position and change every token after it. Ties are rare, one
+or two rows per layer in a 10k prompt, which is why they were not the entry
+point on the stock stack.
+
+## prefix_cache.py
+
+Fills the cache with one prompt, then sends a second prompt that shares its
+first L tokens three times: once able to hit the cache and twice under fresh
+salts. It compares the cache hit with a cold run, and the two cold runs with
+each other as the noise floor. A length where the hit differs earlier or by more
+than twice the floor runs twice more, and it fails only if all three runs do. The lengths default to multiples and half
+multiples of the served `block_size` and `mamba_block_size`.
+
+- A difference from the first generated token, or a large one early, at a
+  length that ends partway through a mamba block, is the pattern of
+  vllm-project/vllm#54076: a hit restores the KDA state from before the hit
+  point while the attention KV is complete.
+- Small differences that start 20 or more tokens in turn up now and then at a
+  random length with speculative decoding on, and do not repeat on a rerun.
+  They are not the cache.
+- Hits one block short of L are vLLM dropping the last block for speculative
+  decoding, on purpose.
+
+Run `determinism.py` first; this means nothing unless two cold runs agree.
+
+## topk_ties.py
+
+`top_k_per_row_prefill` on synthetic scores with an exact tie at the k-th value,
+200 calls on the same input. Needs one CUDA device and vLLM, no checkpoint.
+
+    1 row                    1 distinct output
+    64 rows x 1300           156-166 distinct outputs (two sessions)
+    16128 rows x 1300/8000   200 distinct outputs
+
+## boundary_parity.py
+
+RecoverSSM's KDA state at a block boundary. The model repeats a passage from a
+prompt that ends ~150 tokens short of a mamba block boundary, in 16 variants
+that shift where the verify steps land, with k forced to 1 through adaptive-k's
+control file (the docstring has the commands). Each variant passes when its
+output reaches the passage's last line.
+
+Before the fix, with RecoverSSM on, 8 of 16 variants looped or turned to noise
+right after the boundary (TP=4, 2026-09-29). With `VLLM_GLM5NEXT_RECOVERSSM=0`,
+0 of 16. Needs `/usr/share/common-licenses/GPL-3` on the host that runs it.

docker-images/glm53-kindling/dev/repro/boundary_parity.py

@@ -0,0 +1,70 @@
+#!/usr/bin/env python3
+"""Does a verify step that commits exactly to a KDA block boundary break the reply?
+
+With prefix caching on, vLLM keeps each request's KDA state in the block of its
+current position (align mode). RecoverSSM used to store the state after a
+verify in column committed // block_size: when every draft was accepted and the
+commit ended exactly on a boundary, that is the next block, which is not
+allocated yet, so the state was lost and the rest of the reply turned to noise.
+
+The model repeats a passage word for word (drafts almost all accepted) from a
+prompt that ends ~150 tokens short of a boundary; 16 variants add 0-15 filler
+tokens so the verify steps land at different offsets. Force k=1 first, so every
+fully accepted step commits exactly 2 tokens:
+
+    echo "force 1" | sudo tee /var/lib/docker/volumes/glm53_cache/_data/adaptive-k
+    python3 dev/repro/boundary_parity.py http://127.0.0.1:8002
+    sudo rm /var/lib/docker/volumes/glm53_cache/_data/adaptive-k
+
+A variant passes when its output reaches the passage's last line. On the
+broken code about half the variants loop or turn to noise right after the
+boundary; with the fix, or with VLLM_GLM5NEXT_RECOVERSSM=0, none should.
+Reads the passage from /usr/share/common-licenses/GPL-3.
+
+usage: boundary_parity.py BASE_URL [VARIANTS]
+"""
+import json, os, sys, urllib.request
+
+BASE = sys.argv[1]
+N = int(sys.argv[2]) if len(sys.argv) > 2 else 16
+MODEL = os.environ.get("GATE_MODEL", "glm53")
+_m = urllib.request.urlopen(BASE + "/metrics", timeout=30).read().decode()
+MAMBA = int(_m.split('mamba_block_size="', 1)[1].split('"', 1)[0])
+passage = " ".join(open("/usr/share/common-licenses/GPL-3").read().split()[400:1000])
+last_line = " ".join(passage.split()[-12:])
+
+
+def post(path, body):
+    r = urllib.request.Request(BASE + path, json.dumps(body).encode(), {"Content-Type": "application/json"})
+    return json.loads(urllib.request.urlopen(r, timeout=1800).read())
+
+
+def messages(filler):
+    return [{"role": "user", "content": "Notes (ignore these): " + filler + "\n\nRepeat the following passage "
+             "exactly, word for word, with nothing before or after it:\n\n" + passage}]
+
+
+def prompt_len(filler):
+    return post("/tokenize", {"model": MODEL, "messages": messages(filler), "add_generation_prompt": True,
+                              "chat_template_kwargs": {"enable_thinking": False}})["count"]
+
+
+filler = ""
+for _ in range(8):  # one-token words until the prompt ends ~160 tokens short of a boundary
+    short = (MAMBA - 160 - prompt_len(filler) % MAMBA) % MAMBA
+    if short == 0:
+        break
+    filler += " the" * short
+broken = 0
+for v in range(N):
+    f = filler + " the" * v
+    P = prompt_len(f)
+    r = post("/v1/chat/completions", {"model": MODEL, "messages": messages(f), "max_tokens": 1200,
+                                      "temperature": 0, "cache_salt": os.urandom(8).hex(),
+                                      "chat_template_kwargs": {"enable_thinking": False}})
+    out = r["choices"][0]["message"].get("content") or ""
+    ok = last_line in " ".join(out.split())
+    broken += not ok
+    print(f"variant {v:2d}  prompt {P}  boundary {-P % MAMBA:3d} tokens into the reply  "
+          f"out {r['usage']['completion_tokens']:4d} tokens  {'ok' if ok else 'BROKEN'}  tail {out[-70:]!r}", flush=True)
+print(f"{'FAIL' if broken else 'PASS'}: {broken} of {N} broken")

docker-images/glm53-kindling/dev/repro/decode.py

@@ -0,0 +1,26 @@
+# Per-workload decode with DFlash2 acceptance: thinking off, temperature 0,
+# 512 tokens, median of 3 after one discarded warm-up. Same prompts as the
+# 2026-09-06..26 measurements. usage: decode.py [base-url]
+import json, re, statistics, sys, time, urllib.request
+BASE = sys.argv[1] if len(sys.argv) > 1 else "http://127.0.0.1:8002"  # no /v1
+W = [("structured", "Count from 1 to 200, comma separated. No commentary."),
+     ("code", "Write a red-black tree in Python with insert, delete and rebalancing. Code only."),
+     ("prose", "Explain how a hash map works, in flowing prose. No code, no lists.")]
+def metrics():
+    raw = urllib.request.urlopen(BASE + "/metrics", timeout=30).read().decode()
+    get = lambda k: sum(float(v) for v in re.findall(rf"^vllm:{k}\S*\s+([0-9.e+]+)$", raw, re.M))
+    return get("spec_decode_num_accepted_tokens_total"), get("spec_decode_num_draft_tokens_total")
+def run(prompt, n=512):
+    body = json.dumps({"model": "glm53", "messages": [{"role": "user", "content": prompt}],
+        "max_tokens": n, "temperature": 0, "chat_template_kwargs": {"thinking": False}}).encode()
+    r = urllib.request.Request(BASE + "/v1/chat/completions", data=body, headers={"Content-Type": "application/json"})
+    t = time.time(); d = json.loads(urllib.request.urlopen(r, timeout=600).read()); e = time.time() - t
+    return d["usage"]["completion_tokens"] / e
+print("%-12s %10s %8s %s" % ("workload", "tok/s", "accept", "runs"))
+for name, p in W:
+    run(p, 16)
+    a0, d0 = metrics()
+    rates = [run(p) for _ in range(3)]
+    a1, d1 = metrics()
+    acc = (a1 - a0) / (d1 - d0) if d1 > d0 else float("nan")
+    print("%-12s %10.1f %7.1f%% %s" % (name, statistics.median(rates), acc * 100, " ".join("%.1f" % r for r in rates)), flush=True)

docker-images/glm53-kindling/dev/repro/determinism.py

@@ -0,0 +1,118 @@
+#!/usr/bin/env python3
+"""Do identical requests give identical logprobs?
+
+Standard library only. Point it at a running vLLM OpenAI server with no other
+traffic:
+
+    python determinism.py --url http://127.0.0.1:8000 --model glm53 prompt
+    python determinism.py --url http://127.0.0.1:8000 --model glm53 generate
+
+prompt: for each length, sends one prompt RUNS times with prompt_logprobs and
+reports, per pair of runs, how many prompt logprobs differ, the first differing
+position and the largest difference.
+
+generate: sends one prompt RUNS times and compares the generated tokens and
+their logprobs.
+
+Every request carries a fresh cache_salt, so none of them hits the prefix
+cache. A request that ran alongside other traffic is flagged, because a
+different batch shape changes the rounding.
+"""
+import argparse
+import itertools
+import json
+import os
+import random
+import string
+import urllib.request
+
+
+def post(base, path, payload, timeout=1800):
+    request = urllib.request.Request(
+        base + path, data=json.dumps(payload).encode(), headers={"content-type": "application/json"})
+    with urllib.request.urlopen(request, timeout=timeout) as response:
+        return json.loads(response.read())
+
+
+def finished(base):
+    """Requests the server has completed so far."""
+    text = urllib.request.urlopen(base + "/metrics", timeout=30).read().decode()
+    return sum(float(line.split()[-1]) for line in text.splitlines()
+               if line.startswith("vllm:request_success_total"))
+
+
+def tokens(base, model, n, seed):
+    """n token ids of pseudo-words from a fixed seed: the same ids on any server with this tokenizer."""
+    rng = random.Random(seed)
+    ids = []
+    while len(ids) < n:
+        text = " ".join("".join(rng.choice(string.ascii_lowercase) for _ in range(rng.randint(2, 9)))
+                        for _ in range(3000)) + ".\n"
+        ids += post(base, "/tokenize", {"model": model, "prompt": text, "add_special_tokens": False})["tokens"]
+    return ids[:n]
+
+
+def complete(base, model, prompt, **extra):
+    before = finished(base)
+    out = post(base, "/v1/completions", {"model": model, "prompt": prompt, "temperature": 0,
+                                         "cache_salt": os.urandom(8).hex(), **extra})
+    others = finished(base) - before - 1
+    return out["choices"][0], int(others)
+
+
+def compare(a, b):
+    diffs = [abs(x - y) for x, y in zip(a, b)]
+    first = next((i for i, d in enumerate(diffs) if d), None)
+    return sum(d > 0 for d in diffs), first, max(diffs, default=0.0)
+
+
+def prompt_mode(args):
+    for n in args.lengths:
+        prompt = tokens(args.url, args.model, n, seed=n)
+        runs = []
+        for _ in range(args.runs):
+            choice, others = complete(args.url, args.model, prompt, max_tokens=1, prompt_logprobs=0)
+            runs.append([next(iter(e.values()))["logprob"] for e in choice["prompt_logprobs"][1:]])
+            if others:
+                print(f"  warning: {others} other requests finished during a {n}-token run")
+        for i, j in itertools.combinations(range(len(runs)), 2):
+            count, first, largest = compare(runs[i], runs[j])
+            print(f"{n:6d} tokens  run {j} vs {i}: {count} of {len(runs[i])} differ, "
+                  f"first at {first}, max {largest:.4f}")
+
+
+def generate_mode(args):
+    prompt = tokens(args.url, args.model, args.length, seed=args.length)
+    runs = []
+    for i in range(args.runs):
+        choice, others = complete(args.url, args.model, prompt, max_tokens=args.max_tokens, logprobs=1)
+        runs.append((choice["logprobs"]["tokens"], choice["logprobs"]["token_logprobs"]))
+        print(f"run {i}: first logprobs {[round(x, 4) for x in runs[-1][1][:3]]}"
+              + (f"  ({others} other requests during the run)" if others else ""))
+    for i in range(1, len(runs)):
+        count, first, largest = compare(runs[0][1], runs[i][1])
+        print(f"run {i} vs 0: tokens {'same' if runs[i][0] == runs[0][0] else 'DIFFER'}, "
+              f"{count} logprobs differ, first at {first}, max {largest:.4f}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument("--url", default="http://127.0.0.1:8000")
+    parser.add_argument("--model", default="glm53")
+    sub = parser.add_subparsers(dest="mode", required=True)
+    p = sub.add_parser("prompt")
+    p.add_argument("--lengths", type=lambda s: [int(x) for x in s.split(",")],
+                   default=[3000, 6000, 10000, 14000, 40000])
+    p.add_argument("--runs", type=int, default=3)
+    g = sub.add_parser("generate")
+    g.add_argument("--length", type=int, default=5208)
+    g.add_argument("--runs", type=int, default=6)
+    g.add_argument("--max-tokens", type=int, default=48)
+    args = parser.parse_args()
+    # The first long request after a boot can read differently; spend it here.
+    complete(args.url, args.model, tokens(args.url, args.model, 8000, seed=0), max_tokens=1)
+    (prompt_mode if args.mode == "prompt" else generate_mode)(args)
+
+
+if __name__ == "__main__":
+    main()

docker-images/glm53-kindling/dev/repro/greedy_nondet.py

@@ -0,0 +1,140 @@
+#!/usr/bin/env python3
+"""Identical greedy requests over a long prompt return different completions.
+
+Standard library only. Point it at a running vLLM OpenAI server:
+
+    python greedy_nondet.py --url http://127.0.0.1:8000 --model glm53
+
+Sends the same ~42k-token prompt N times at temperature 0 with a fixed seed and
+groups the completions. More than one group means the forward pass is not
+reproducible.
+"""
+import argparse
+import collections
+import hashlib
+import json
+import queue
+import threading
+import urllib.request
+
+# Repeated verbatim to reach the target length. Content does not matter: any
+# filler of the same size behaves the same.
+PARAGRAPH = """The job runner reads a JSON manifest, then runs each job in order.
+Every job carries a name, a command, a retry count and an environment overlay.
+The overlay is applied on top of the inherited environment on each attempt, so a
+key set to null must remove the inherited value rather than set it to the empty
+string. Failures are retried up to the configured count, and the runner returns
+the worst exit code it saw. The exporter reads the same manifest and publishes
+one gauge per job plus a counter for retries.
+"""
+
+QUESTION = "\nIn one sentence: what does the runner return when every job fails?"
+
+
+def post(url, payload, timeout=1800):
+    req = urllib.request.Request(
+        url, data=json.dumps(payload).encode(), headers={"content-type": "application/json"}
+    )
+    with urllib.request.urlopen(req, timeout=timeout) as r:
+        return json.loads(r.read())
+
+
+def build_prompt(base, model, target):
+    """Repeat PARAGRAPH until the tokenizer reports at least `target` tokens."""
+    low, high = 1, 16
+    while True:
+        text = PARAGRAPH * high + QUESTION
+        if count_tokens(base, model, text) >= target:
+            break
+        low, high = high, high * 2
+    while low < high:
+        mid = (low + high) // 2
+        if count_tokens(base, model, PARAGRAPH * mid + QUESTION) >= target:
+            high = mid
+        else:
+            low = mid + 1
+    text = PARAGRAPH * low + QUESTION
+    return text, count_tokens(base, model, text)
+
+
+def count_tokens(base, model, text):
+    return post(base + "/tokenize", {"model": model, "prompt": text}, timeout=300)["count"]
+
+
+def complete(base, model, prompt, max_tokens):
+    body = {
+        "model": model,
+        "messages": [{"role": "user", "content": prompt}],
+        "temperature": 0,
+        "seed": 1234,
+        "max_tokens": max_tokens,
+    }
+    choice = post(base + "/v1/chat/completions", body)["choices"][0]
+    message = choice["message"]
+    # Compare the reasoning too, under either spelling: `reasoning_content` from
+    # SGLang, `reasoning` from vLLM. A thinking model whose trace does not close
+    # inside max_tokens returns empty content, and hashing content alone then
+    # groups every run together and reports perfect agreement -- the divergence
+    # this script exists to find happens in the trace as readily as in the answer.
+    text = "".join(
+        message.get(key) or ""
+        for key in ("reasoning_content", "reasoning", "content")
+    )
+    if not text:
+        raise SystemExit(
+            f"empty completion and empty reasoning ({choice.get('finish_reason')}, "
+            f"{max_tokens} max_tokens): nothing to compare. Raise --max-tokens so "
+            "the trace can close, or send reasoning_effort=none."
+        )
+    return text
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--url", default="http://127.0.0.1:8000")
+    ap.add_argument("--model", required=True)
+    ap.add_argument("--prompt-tokens", type=int, default=42000)
+    ap.add_argument("--max-tokens", type=int, default=48)
+    ap.add_argument("--runs", type=int, default=16)
+    ap.add_argument("--concurrency", type=int, default=3)
+    a = ap.parse_args()
+    base = a.url.rstrip("/")
+
+    prompt, n_tokens = build_prompt(base, a.model, a.prompt_tokens)
+    print(f"prompt: {n_tokens} tokens, {a.runs} runs at temperature 0, seed 1234")
+
+    results, lock = [], threading.Lock()
+    work = queue.Queue()
+    for i in range(a.runs):
+        work.put(i)
+
+    def worker():
+        while True:
+            try:
+                work.get_nowait()
+            except queue.Empty:
+                return
+            try:
+                out = complete(base, a.model, prompt, a.max_tokens)
+            except Exception as exc:
+                out = f"<error {exc!r}>"
+            with lock:
+                results.append(out)
+                print(f"  {len(results):>3}/{a.runs} {hashlib.sha1(out.encode()).hexdigest()[:8]}",
+                      flush=True)
+
+    threads = [threading.Thread(target=worker) for _ in range(a.concurrency)]
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join()
+
+    groups = collections.Counter(results)
+    print(f"\n{len(groups)} distinct completion(s) from {a.runs} identical requests")
+    for text, n in groups.most_common():
+        print(f"\n--- {n} run(s) ---\n{text}")
+    raise SystemExit(0 if len(groups) == 1 else 1)
+
+
+if __name__ == "__main__":
+    main()

docker-images/glm53-kindling/dev/repro/marlin_moe_nondet.py

@@ -0,0 +1,77 @@
+#!/usr/bin/env python3
+"""fused_marlin_moe returns different bits for the same inputs at some M.
+
+No checkpoint needed: random NVFP4 weights, one CUDA device. Runs the same
+call repeatedly and compares the outputs bit for bit.
+
+    python marlin_moe_nondet.py
+"""
+import argparse
+
+import torch
+from vllm.model_executor.layers.fused_moe.experts.marlin_moe import fused_marlin_moe
+from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
+    rand_marlin_weight_nvfp4_like,
+)
+from vllm.scalar_type import scalar_types
+
+DEV = "cuda"
+
+
+def quant_stack(experts, n, k, group):
+    q, s, g = [], [], []
+    for _ in range(experts):
+        w = torch.randn(n, k, device=DEV, dtype=torch.bfloat16) * 0.05
+        _, qw, sw, gw = rand_marlin_weight_nvfp4_like(w, group)
+        q.append(qw), s.append(sw), g.append(gw)
+    return torch.stack(q), torch.stack(s), torch.stack(g)
+
+
+def trial(w1, w2, m, reps, experts, hidden, topk):
+    gen = torch.Generator(device=DEV).manual_seed(1)
+    x = (torch.randn(m, hidden, device=DEV, generator=gen) * 0.5).to(torch.bfloat16)
+    weights, ids = torch.topk(
+        torch.softmax(torch.randn(m, experts, device=DEV, generator=gen), -1), topk, dim=-1
+    )
+    weights = (weights / weights.sum(-1, keepdim=True)).float()
+    ids = ids.to(torch.int32)
+
+    outs = []
+    for _ in range(reps):
+        out = fused_marlin_moe(
+            x, w1[0], w2[0], None, None, w1[1], w2[1], weights, ids,
+            scalar_types.float4_e2m1f.id, global_num_experts=experts,
+            global_scale1=w1[2], global_scale2=w2[2],
+        )
+        torch.cuda.synchronize()
+        outs.append(out.clone())
+
+    differing = sum(1 for o in outs[1:] if not torch.equal(o, outs[0]))
+    worst = max(float((o.float() - outs[0].float()).abs().max()) for o in outs[1:])
+    return differing, worst
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--experts", type=int, default=288)
+    ap.add_argument("--topk", type=int, default=8)
+    ap.add_argument("--hidden", type=int, default=4096)
+    ap.add_argument("--inter", type=int, default=2048)
+    ap.add_argument("--group", type=int, default=16)
+    ap.add_argument("--reps", type=int, default=24)
+    ap.add_argument("--m", type=int, nargs="+", default=[1, 1024, 2048, 2816, 3072, 3104, 3136])
+    a = ap.parse_args()
+
+    torch.manual_seed(0)
+    w1 = quant_stack(a.experts, 2 * a.inter, a.hidden, a.group)  # gate and up
+    w2 = quant_stack(a.experts, a.hidden, a.inter, a.group)      # down
+    print(f"{torch.cuda.get_device_name(0)}  experts={a.experts} topk={a.topk} "
+          f"hidden={a.hidden} inter={a.inter} reps={a.reps}")
+    for m in a.m:
+        differing, worst = trial(w1, w2, m, a.reps, a.experts, a.hidden, a.topk)
+        print(f"M={m:<6} runs differing from the first: {differing}/{a.reps - 1}   "
+              f"max abs diff {worst:.3e}")
+
+
+if __name__ == "__main__":
+    main()

docker-images/glm53-kindling/dev/repro/needle.py

@@ -0,0 +1,94 @@
+#!/usr/bin/env python3
+"""Long-context retrieval: can the model find one fact buried in a long prompt?
+
+Standard library only.
+
+    python needle.py --url http://127.0.0.1:8000 --model glm53
+
+Places a distinctive sentence at several depths in filler of several lengths and
+asks for it back. Reports one line per (length, depth) with the exact-match
+verdict, so a partial credit answer still reads as a miss.
+"""
+import argparse
+import json
+import time
+import urllib.request
+
+NEEDLE = "The access code for the Vermilion vault is 7741-KESTREL."
+ANSWER = "7741-KESTREL"
+QUESTION = ("\n\nQuestion: what is the access code for the Vermilion vault? "
+            "Reply with the code alone.")
+
+# Numbered so the filler is not a degenerate repeat of one block.
+FILLER = ("Maintenance note {n}: the west pump was serviced and the intake "
+          "filter replaced; readings held within tolerance for the shift.\n")
+
+
+def post(url, payload, timeout=3600):
+    request = urllib.request.Request(
+        url, data=json.dumps(payload).encode(), headers={"content-type": "application/json"})
+    with urllib.request.urlopen(request, timeout=timeout) as response:
+        return json.loads(response.read())
+
+
+def count(base, model, text):
+    return post(base + "/tokenize", {"model": model, "prompt": text}, timeout=600)["count"]
+
+
+def build(base, model, target, depth):
+    """Filler of ~target tokens with the needle at `depth` through it."""
+    per = count(base, model, FILLER.format(n=1))
+    lines = max(1, target // max(per, 1))
+    body = [FILLER.format(n=i) for i in range(lines)]
+    at = min(len(body), max(0, int(len(body) * depth)))
+    body.insert(at, NEEDLE + "\n")
+    text = "".join(body)
+    return text + QUESTION
+
+
+def ask(base, model, prompt, max_tokens):
+    body = {"model": model, "messages": [{"role": "user", "content": prompt}],
+            "max_tokens": max_tokens, "temperature": 0,
+            "chat_template_kwargs": {"enable_thinking": False}}
+    started = time.time()
+    result = post(base + "/v1/chat/completions", body)
+    message = result["choices"][0]["message"]
+    text = (message.get("content") or "") + " " + (message.get("reasoning") or "")
+    return text.strip(), result.get("usage", {}), time.time() - started
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--url", default="http://127.0.0.1:8000")
+    ap.add_argument("--model", required=True)
+    ap.add_argument("--lengths", type=int, nargs="+", default=[32000, 128000, 256000])
+    ap.add_argument("--depths", type=float, nargs="+", default=[0.1, 0.5, 0.9])
+    ap.add_argument("--max-tokens", type=int, default=200)
+    a = ap.parse_args()
+    base = a.url.rstrip("/")
+
+    passed = total = 0
+    for length in a.lengths:
+        for depth in a.depths:
+            prompt = build(base, a.model, length, depth)
+            try:
+                text, usage, secs = ask(base, a.model, prompt, a.max_tokens)
+            except Exception as exc:
+                print(f"  {length:>7} tok  depth {depth:>4.0%}  ERROR {exc!r}", flush=True)
+                total += 1
+                continue
+            hit = ANSWER in text
+            passed += hit
+            total += 1
+            finish = usage.get("completion_tokens")
+            capped = finish is not None and finish >= a.max_tokens
+            note = " (hit max_tokens)" if capped and not hit else ""
+            print(f"  {usage.get('prompt_tokens', '?'):>7} tok  depth {depth:>4.0%}  "
+                  f"{'FOUND' if hit else 'MISS ':5}  {secs:6.1f}s{note}  {text[:60]!r}",
+                  flush=True)
+    print(f"\n  retrieval: {passed}/{total}")
+    raise SystemExit(0 if passed == total else 1)
+
+
+if __name__ == "__main__":
+    main()

docker-images/glm53-kindling/dev/repro/perf-repeat.py

@@ -0,0 +1,50 @@
+"""Decode rate over N short-prompt runs, plus DFlash2 draft acceptance.
+
+Acceptance length, not rate, is what predicts throughput for a k-token drafter.
+"""
+import json, re, statistics, sys, time, urllib.request
+BASE, MODEL, N = sys.argv[1], sys.argv[2], int(sys.argv[3])
+
+def metrics():
+    raw = urllib.request.urlopen(BASE.replace("/v1", "") + "/metrics", timeout=30).read().decode()
+    out = {}
+    for key in ("spec_decode_num_accepted_tokens_total", "spec_decode_num_draft_tokens_total"):
+        m = re.search(rf"^vllm:{key}\S*\s+([0-9.e+]+)$", raw, re.M)
+        if m:
+            out[key] = float(m.group(1))
+    return out
+
+def once(prompt, max_tokens):
+    body = {"model": MODEL, "messages": [{"role": "user", "content": prompt}],
+            "max_tokens": max_tokens, "temperature": 0, "stream": True,
+            "stream_options": {"include_usage": True}}
+    req = urllib.request.Request(BASE + "/chat/completions", json.dumps(body).encode(),
+                                {"Content-Type": "application/json"})
+    t0 = time.time(); ttft = None; usage = None
+    with urllib.request.urlopen(req, timeout=1800) as r:
+        for raw in r:
+            line = raw.decode("utf-8", "replace").strip()
+            if not line.startswith("data:"): continue
+            d = line[5:].strip()
+            if d == "[DONE]": break
+            ev = json.loads(d)
+            if ev.get("usage"): usage = ev["usage"]
+            for ch in ev.get("choices", []):
+                delta = ch.get("delta", {})
+                if (delta.get("content") or delta.get("reasoning")) and ttft is None:
+                    ttft = time.time() - t0
+    total = time.time() - t0
+    completion = (usage or {}).get("completion_tokens", 0)
+    return (completion - 1) / (total - ttft) if ttft and total > ttft else 0
+
+before = metrics()
+rates = [once("Explain unified memory on GB10 in two sentences.", 200) for _ in range(N)]
+after = metrics()
+acc = after.get("spec_decode_num_accepted_tokens_total", 0) - before.get("spec_decode_num_accepted_tokens_total", 0)
+draft = after.get("spec_decode_num_draft_tokens_total", 0) - before.get("spec_decode_num_draft_tokens_total", 0)
+print(f"  decode tok/s over {N} runs: " + " ".join(f"{r:.1f}" for r in sorted(rates)))
+print(f"  median {statistics.median(rates):.1f}  mean {statistics.mean(rates):.1f}  "
+      f"spread {max(rates)-min(rates):.1f}")
+if draft:
+    print(f"  draft: accepted {acc:.0f} of {draft:.0f} = {acc/draft:.3f}; "
+          f"accepted per step {acc/(draft/7):.2f} of 7")

docker-images/glm53-kindling/dev/repro/prefill.py

@@ -0,0 +1,32 @@
+"""Cold-cache prefill throughput: unique filler each time so nothing is cached."""
+import json, sys, time, urllib.request, uuid
+BASE, MODEL = sys.argv[1], sys.argv[2]
+def once(target):
+    nonce = uuid.uuid4().hex
+    unit = f"Reading {nonce} note {{n}}: the pump was serviced and the filter replaced.\n"
+    text = "".join(unit.format(n=i) for i in range(target // 16))
+    body = {"model": MODEL, "messages": [{"role": "user", "content": text + "\nReply with OK."}],
+            "max_tokens": 4, "temperature": 0, "stream": True,
+            "stream_options": {"include_usage": True}}
+    req = urllib.request.Request(BASE + "/chat/completions", json.dumps(body).encode(),
+                                {"Content-Type": "application/json"})
+    t0 = time.time(); ttft = None; usage = None
+    with urllib.request.urlopen(req, timeout=3600) as r:
+        for raw in r:
+            line = raw.decode("utf-8", "replace").strip()
+            if not line.startswith("data:"): continue
+            d = line[5:].strip()
+            if d == "[DONE]": break
+            ev = json.loads(d)
+            if ev.get("usage"): usage = ev["usage"]
+            for ch in ev.get("choices", []):
+                d = ch.get("delta") or {}
+                if (d.get("content") or d.get("reasoning") or d.get("reasoning_content")) and ttft is None:
+                    ttft = time.time() - t0
+    n = (usage or {}).get("prompt_tokens", 0)
+    if ttft:
+        print(f"  prompt={n:>7} ttft={ttft:6.1f}s  prefill={n/ttft:8.0f} tok/s", flush=True)
+    else:
+        print(f"  prompt={n:>7} no first token seen (usage={usage})", flush=True)
+for t in ([int(a) for a in sys.argv[3:]] or [8000, 42000, 128000]):
+    once(t)

docker-images/glm53-kindling/dev/repro/prefix_cache.py

@@ -0,0 +1,134 @@
+#!/usr/bin/env python3
+"""Does a prefix-cache hit change what the model generates?
+
+Standard library only. Point it at a running vLLM OpenAI server with prefix
+caching on and no other traffic:
+
+    python prefix_cache.py --url http://127.0.0.1:8000 --model glm53
+
+For each shared-prefix length L: request A (the first L + 700 tokens of a
+document) fills the cache under a salt. Request B (the same first L tokens,
+then 600 tokens of other text) runs twice: under A's salt, so it can hit the
+cache, and twice under fresh salts, so it prefills cold. All three generate
+greedily. The warm run is compared with the first cold run, and the two cold
+runs with each other: that second comparison is the noise floor, the difference
+two runs show without the cache involved. A length whose cache hit differs by
+more than the floor runs twice more with new salts: corruption repeats at the
+same length, while occasional decode noise (the verify length picked from step
+timing, for example) does not. The hit size comes from the
+server's prefix-cache hit counter, since usage.cached_tokens can read 0 on
+hybrid models.
+
+Run determinism.py first: this compares two runs, so it means something only
+when two cold runs already agree. Generated logprobs are used because
+prompt_logprobs requests skip the prefix cache.
+
+The default lengths are multiples and half multiples of the attention block
+and the mamba block, read from the server's cache_config_info.
+"""
+import argparse
+import json
+import os
+import random
+import string
+import urllib.request
+
+
+def post(base, path, payload, timeout=1800):
+    request = urllib.request.Request(
+        base + path, data=json.dumps(payload).encode(), headers={"content-type": "application/json"})
+    with urllib.request.urlopen(request, timeout=timeout) as response:
+        return json.loads(response.read())
+
+
+def metrics(base):
+    return urllib.request.urlopen(base + "/metrics", timeout=30).read().decode()
+
+
+def counter(text, name):
+    return sum(float(line.split()[-1]) for line in text.splitlines() if line.startswith(name))
+
+
+def label(text, name):
+    return text.split(f'{name}="', 1)[1].split('"', 1)[0]
+
+
+def tokens(base, model, n, seed):
+    """n token ids of pseudo-words from a fixed seed."""
+    rng = random.Random(seed)
+    ids = []
+    while len(ids) < n:
+        text = " ".join("".join(rng.choice(string.ascii_lowercase) for _ in range(rng.randint(2, 9)))
+                        for _ in range(3000)) + ".\n"
+        ids += post(base, "/tokenize", {"model": model, "prompt": text, "add_special_tokens": False})["tokens"]
+    return ids[:n]
+
+
+def generate(base, model, prompt, salt, max_tokens):
+    before = metrics(base)
+    out = post(base, "/v1/completions", {"model": model, "prompt": prompt, "temperature": 0,
+                                         "max_tokens": max_tokens, "logprobs": 1, "cache_salt": salt})
+    after = metrics(base)
+    hit = counter(after, "vllm:prefix_cache_hits_total") - counter(before, "vllm:prefix_cache_hits_total")
+    others = counter(after, "vllm:request_success_total") - counter(before, "vllm:request_success_total") - 1
+    lp = out["choices"][0]["logprobs"]
+    return lp["tokens"], lp["token_logprobs"], int(hit), int(others)
+
+
+def compare(a, b):
+    """(first differing generated token or None, largest logprob difference, same tokens)."""
+    diffs = [abs(x - y) for x, y in zip(a[1], b[1])]
+    return next((i for i, d in enumerate(diffs) if d), None), max(diffs, default=0.0), a[0] == b[0]
+
+
+def describe(first, largest, same_tokens):
+    if first is None and same_tokens:
+        return "identical"
+    return f"from token {first}, max {largest:.4f}" + ("" if same_tokens else ", tokens differ")
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument("--url", default="http://127.0.0.1:8000")
+    parser.add_argument("--model", default="glm53")
+    parser.add_argument("--lengths", type=lambda s: [int(x) for x in s.split(",")])
+    parser.add_argument("--max-tokens", type=int, default=48)
+    args = parser.parse_args()
+
+    text = metrics(args.url)
+    block, mamba = int(label(text, "block_size")), int(label(text, "mamba_block_size"))
+    print(f"block_size {block}, mamba_block_size {mamba}, mamba_cache_mode {label(text, 'mamba_cache_mode')}")
+    lengths = args.lengths or sorted({b * h // 2 for b in (block, mamba) for h in (3, 4, 5, 6, 8, 9, 14, 15)})
+
+    doc = tokens(args.url, args.model, max(lengths) + 800, seed=1)
+    tail = tokens(args.url, args.model, 600, seed=2)
+    def trial(n):
+        """(hit, warm vs cold, cold vs cold, other requests, worse than the floor) for length n."""
+        salt = os.urandom(8).hex()
+        generate(args.url, args.model, doc[:n + 700], salt, args.max_tokens)
+        prompt = doc[:n] + tail
+        warm = generate(args.url, args.model, prompt, salt, args.max_tokens)
+        cold = generate(args.url, args.model, prompt, os.urandom(8).hex(), args.max_tokens)
+        cold2 = generate(args.url, args.model, prompt, os.urandom(8).hex(), args.max_tokens)
+        wf, wl, ws = compare(warm, cold)
+        cf, cl, cs = compare(cold, cold2)
+        # Worse than the floor: an earlier first difference, or a larger one.
+        worse = (wf is not None or not ws) and (cf is None or (wf or 0) < cf or wl > 2 * cl)
+        return warm[2], (wf, wl, ws), (cf, cl, cs), warm[3] + cold[3] + cold2[3], worse
+
+    failed = 0
+    for n in lengths:
+        hit, w, c, others, worse = trial(n)
+        verdict = ""
+        if worse:
+            repeats = sum(trial(n)[4] for _ in range(2))
+            verdict = f"  <- worse than noise, and again in {repeats} of 2 reruns"
+            failed += repeats == 2
+        print(f"L {n:6d} (mod {block} = {n % block:5d}, mod {mamba} = {n % mamba:5d})  hit {hit:6d}  "
+              f"warm vs cold: {describe(*w)};  cold vs cold: {describe(*c)}{verdict}"
+              + (f"  ({others} other requests)" if others else ""))
+    print("PASS" if not failed else f"{failed} of {len(lengths)} lengths repeatedly worse than the noise floor")
+
+
+if __name__ == "__main__":
+    main()

docker-images/glm53-kindling/dev/repro/toolcall_corruption.py

@@ -0,0 +1,341 @@
+#!/usr/bin/env python3
+"""Greedy decoding corrupts tokens inside a tool call at long context.
+
+Standard library only. Point it at a running vLLM OpenAI server:
+
+    python toolcall_corruption.py --url http://127.0.0.1:8000 --model glm53
+
+Builds a 42272-token coding-agent transcript (47 tools, 24 tool-call/result
+pairs, all generated from fixed seeds so every run sends identical bytes), then
+sends it N times at temperature 0 with a fixed seed. The correct continuation is
+a Bash call repeating a command from earlier in the conversation, so any
+difference between runs is visible token by token.
+
+Observed on 4x GB10 (sm121) at TP=4, GLM-5.3-Flash-NVFP4. One run of 40, all at
+temperature 0 with the same seed, gave 5 distinct completions:
+
+    30 runs  correct
+     4 runs  pos 23  ' identical' for '>&'  -> pytest -q 2 identical 2>&1
+     3 runs  pos 26  ' pytest'    for ' tail' -> ... 2>&1 | pytest -q 2>&1 | tail -40
+     2 runs  pos 23  ' Bash'      for '>&'  -> pytest -q 2 Bash 2>&1
+     1 run   pos 31  ' Read'      for 'description' -> <arg_key> Read more: https://...
+
+A corrupted tool NAME is worse than a corrupted argument: the GLM parser runs
+with validate_tool_names=True, so an unknown name emits zero deltas and the
+request finishes `stop` with no content and no tool calls. Clients report that
+as "completed response with no content".
+
+Reducing the transcript weakens it. A short prompt (2k tokens) is bit-stable
+over 16 runs, and a transcript whose padding repeats the command's own output
+is stable over 40, because every token of the command is then high-margin.
+"""
+import argparse
+import collections
+import json
+import queue
+import random
+import threading
+import time
+import urllib.request
+
+SYSTEM = """You are a coding agent operating inside a developer's repository. You help by reading files, searching the codebase, running commands, and editing code. Always investigate before you change anything. Prefer small, verifiable steps. When you need information from the repository, call a tool rather than guessing. When a task is complete, summarize what you did in plain language.
+
+Guidelines:
+- Never invent file contents. Read them.
+- Use Grep to locate symbols and Glob to locate files by name.
+- Use Bash only for commands that cannot be done with the other tools.
+- Keep edits minimal and targeted; do not reformat unrelated code.
+- If a command fails, read the error and decide the next step; do not repeat the same command blindly.
+- Report absolute paths.
+"""
+
+def tool(name, desc, props, req):
+    return {"type": "function", "function": {"name": name, "description": desc,
+            "parameters": {"type": "object", "properties": props, "required": req}}}
+
+TOOLS = [
+    tool("Read", "Read a file from the filesystem. Returns the content with line numbers.",
+         {"file_path": {"type": "string", "description": "Absolute path"},
+          "offset": {"type": "integer"}, "limit": {"type": "integer"}}, ["file_path"]),
+    tool("Write", "Write a file, overwriting it.",
+         {"file_path": {"type": "string"}, "content": {"type": "string"}}, ["file_path", "content"]),
+    tool("Edit", "Exact string replacement in a file.",
+         {"file_path": {"type": "string"}, "old_string": {"type": "string"},
+          "new_string": {"type": "string"}, "replace_all": {"type": "boolean"}},
+         ["file_path", "old_string", "new_string"]),
+    tool("Bash", "Run a shell command and return stdout/stderr.",
+         {"command": {"type": "string"}, "timeout": {"type": "integer"},
+          "description": {"type": "string"}}, ["command"]),
+    tool("Grep", "Search file contents with a regex (ripgrep).",
+         {"pattern": {"type": "string"}, "path": {"type": "string"}, "glob": {"type": "string"},
+          "output_mode": {"type": "string", "enum": ["content", "files_with_matches", "count"]},
+          "-n": {"type": "boolean"}, "-i": {"type": "boolean"}}, ["pattern"]),
+    tool("Glob", "Find files by glob pattern.",
+         {"pattern": {"type": "string"}, "path": {"type": "string"}}, ["pattern"]),
+    tool("WebFetch", "Fetch a URL and return its text.",
+         {"url": {"type": "string"}, "prompt": {"type": "string"}}, ["url", "prompt"]),
+    tool("WebSearch", "Search the web.", {"query": {"type": "string"}}, ["query"]),
+    tool("TodoWrite", "Replace the todo list.",
+         {"todos": {"type": "array", "items": {"type": "object", "properties": {
+             "content": {"type": "string"}, "status": {"type": "string",
+             "enum": ["pending", "in_progress", "completed"]}, "activeForm": {"type": "string"}},
+             "required": ["content", "status", "activeForm"]}}}, ["todos"]),
+    tool("Agent", "Launch a sub-agent for a multi-step task.",
+         {"description": {"type": "string"}, "prompt": {"type": "string"},
+          "subagent_type": {"type": "string"}}, ["description", "prompt"]),
+    tool("NotebookEdit", "Edit a Jupyter notebook cell.",
+         {"notebook_path": {"type": "string"}, "cell_id": {"type": "string"},
+          "new_source": {"type": "string"}, "cell_type": {"type": "string"},
+          "edit_mode": {"type": "string"}}, ["notebook_path", "new_source"]),
+    tool("AskUserQuestion", "Ask the user a clarifying question.",
+         {"questions": {"type": "array", "items": {"type": "object"}}}, ["questions"]),
+    tool("KillShell", "Kill a background shell.", {"shell_id": {"type": "string"}}, ["shell_id"]),
+    tool("TaskOutput", "Read output of a background task.",
+         {"task_id": {"type": "string"}, "block": {"type": "boolean"}}, ["task_id"]),
+    tool("EnterPlanMode", "Switch to planning mode.", {}, []),
+    tool("ExitPlanMode", "Leave planning mode with a plan.", {"plan": {"type": "string"}}, ["plan"]),
+    tool("LSP", "Language server query.",
+         {"operation": {"type": "string"}, "filePath": {"type": "string"},
+          "line": {"type": "integer"}, "character": {"type": "integer"}}, ["operation", "filePath"]),
+    tool("SlashCommand", "Run a slash command.", {"command": {"type": "string"}}, ["command"]),
+]
+
+FILE_A = "\n".join(f"{i:>6}\t" + line for i, line in enumerate("""#!/usr/bin/env python3
+import json
+import os
+import sys
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+
+
+@dataclass
+class Job:
+    name: str
+    command: list[str]
+    retries: int = 0
+    env: dict = field(default_factory=dict)
+
+
+def load_jobs(path: Path) -> list[Job]:
+    with path.open() as f:
+        raw = json.load(f)
+    jobs = []
+    for item in raw["jobs"]:
+        jobs.append(Job(name=item["name"], command=item["command"],
+                        retries=item.get("retries", 0), env=item.get("env", {})))
+    return jobs
+
+
+def run_job(job: Job) -> int:
+    import subprocess
+    env = dict(os.environ)
+    env.update(job.env)
+    attempt = 0
+    while True:
+        started = time.monotonic()
+        proc = subprocess.run(job.command, env=env, capture_output=True, text=True)
+        elapsed = time.monotonic() - started
+        if proc.returncode == 0:
+            print(f"ok {job.name} in {elapsed:.1f}s")
+            return 0
+        attempt += 1
+        print(f"fail {job.name} rc={proc.returncode} attempt={attempt}", file=sys.stderr)
+        print(proc.stderr, file=sys.stderr)
+        if attempt > job.retries:
+            return proc.returncode
+
+
+def main() -> int:
+    if len(sys.argv) != 2:
+        print("usage: runner.py JOBS.json", file=sys.stderr)
+        return 2
+    jobs = load_jobs(Path(sys.argv[1]))
+    worst = 0
+    for job in jobs:
+        rc = run_job(job)
+        worst = max(worst, rc)
+    return worst
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
+""".splitlines(), 1))
+
+TEST_OUT = """============================= test session starts ==============================
+platform linux -- Python 3.12.3, pytest-8.3.2
+collected 14 items
+
+tests/test_runner.py ..........F...                                      [100%]
+
+=================================== FAILURES ===================================
+____________________________ test_retry_env_isolated ___________________________
+
+    def test_retry_env_isolated(tmp_path, monkeypatch):
+        monkeypatch.setenv("SHARED", "outer")
+        job = Job(name="x", command=["sh", "-c", "test \\"$SHARED\\" = inner"], retries=1, env={"SHARED": "inner"})
+>       assert run_job(job) == 0
+E       assert 1 == 0
+
+tests/test_runner.py:61: AssertionError
+----------------------------- Captured stderr call -----------------------------
+fail x rc=1 attempt=1
+fail x rc=1 attempt=2
+=========================== short test summary info ============================
+FAILED tests/test_runner.py::test_retry_env_isolated - assert 1 == 0
+========================= 1 failed, 13 passed in 0.41s =========================
+"""
+
+GREP_OUT = """/home/dev/proj/runner.py:32:    env = dict(os.environ)
+/home/dev/proj/runner.py:33:    env.update(job.env)
+/home/dev/proj/tests/test_runner.py:55:def test_retry_env_isolated(tmp_path, monkeypatch):
+/home/dev/proj/tests/test_runner.py:57:    job = Job(name="x", command=["sh", "-c", "test \\"$SHARED\\" = inner"], retries=1, env={"SHARED": "inner"})
+/home/dev/proj/docs/design.md:14:Each job's `env` overlays the process environment; keys set to null must unset."""
+
+def tc(name, args, cid):
+    return {"id": cid, "type": "function", "function": {"name": name, "arguments": json.dumps(args)}}
+
+def extra_tools(n):
+    out = []
+    for i in range(n):
+        out.append(tool(
+            f"mcp__server{i%5}__op{i}",
+            f"Operation {i} on server {i%5}: queries the {['inventory','billing','metrics','tickets','deploys'][i%5]} system and returns a JSON document. Use when the user asks about {['stock levels','invoices','latency','support cases','releases'][i%5]}.",
+            {"query": {"type": "string", "description": "Free-text query"},
+             "limit": {"type": "integer", "description": "Max rows", "default": 50},
+             "since": {"type": "string", "description": "ISO-8601 timestamp"},
+             "fields": {"type": "array", "items": {"type": "string"}}},
+            ["query"]))
+    return out
+
+TOOLS47 = TOOLS + extra_tools(47 - len(TOOLS))
+
+def big_file(seed, lines=220):
+    rnd = random.Random(seed)
+    words = "config loader parse validate schema emit retry backoff socket frame header payload checksum rotate flush queue worker shard replica leader follower epoch term".split()
+    out = []
+    for i in range(1, lines + 1):
+        k = rnd.randint(2, 7)
+        ident = "_".join(rnd.choice(words) for _ in range(2))
+        out.append(f"{i:>6}\t" + ("    " * rnd.randint(0, 3)) + rnd.choice([
+            f"def {ident}(self, {', '.join(rnd.choice(words) for _ in range(k))}):",
+            f"{ident} = {rnd.choice(words)}.{rnd.choice(words)}({rnd.randint(0, 999)})",
+            f"if {ident} is None or {rnd.choice(words)} > {rnd.randint(1, 64)}:",
+            f"return {ident}",
+            f"# {rnd.choice(words)} {rnd.choice(words)} {rnd.choice(words)} handled below",
+            f"raise ValueError(f\"{ident} out of range: {{{rnd.choice(words)}}}\")",
+            f"for {rnd.choice(words)} in {ident}:",
+            f"self._{ident}[{rnd.choice(words)}] = {rnd.choice(words)}",
+            "",
+        ]))
+    return "\n".join(out)
+
+def build(nonce, pairs=24):
+    msgs = [{"role": "system", "content": SYSTEM + f"\n\nSession: {nonce}\n"}]
+    msgs.append({"role": "user", "content": "The nightly job runner in /home/dev/proj is flaky: some jobs report success but their env overrides are dropped on retry, and the metrics exporter shows gaps. Investigate the whole flow (runner.py, the exporter under svc/, and the tests), find the root causes, and fix them. Run the tests when you think you are done."})
+    files = [f"/home/dev/proj/{p}" for p in ["runner.py", "svc/exporter.py", "svc/metrics.py", "svc/queue.py", "svc/config.py", "tests/test_runner.py", "tests/test_exporter.py", "svc/retry.py", "svc/shard.py", "svc/leader.py"]]
+    rnd = random.Random(nonce)
+    for i in range(pairs):
+        cid = f"call_{i+1}"
+        kind = i % 6
+        if kind == 0:
+            call = tc("Bash", {"command": "cd /home/dev/proj && python -m pytest -q 2>&1 | tail -40", "description": "Run tests"}, cid)
+            result = TEST_OUT
+        elif kind in (1, 2, 4):
+            f = files[(i * 3) % len(files)]
+            call = tc("Read", {"file_path": f}, cid)
+            result = FILE_A if f.endswith("runner.py") else big_file(f + nonce[:2])
+        elif kind == 3:
+            call = tc("Grep", {"pattern": rnd.choice(["env", "retry", "flush", "shard", "epoch"]), "path": "/home/dev/proj", "output_mode": "content", "-n": True}, cid)
+            result = GREP_OUT + "\n" + "\n".join(f"/home/dev/proj/svc/{rnd.choice(['exporter','queue','retry'])}.py:{rnd.randint(1,300)}:    {rnd.choice(['env','retry','flush'])} = {rnd.choice(['self','cfg','job'])}.{rnd.choice(['get','pop','update'])}(\"{rnd.choice(['SHARED','TIMEOUT','SHARD'])}\")" for _ in range(25))
+        else:
+            call = tc(f"mcp__server{i%5}__op{i%47}", {"query": "job runner gaps last 24h", "limit": 20}, cid)
+            result = json.dumps({"rows": [{"ts": f"2026-09-0{1+(j%6)}T0{j%9}:1{j%5}:00Z", "job": f"nightly-{j}", "status": rnd.choice(["ok", "retry", "fail"]), "env_keys": rnd.randint(0, 5), "latency_ms": rnd.randint(50, 9000)} for j in range(20)], "total": 20})
+        msgs.append({"role": "assistant", "content": "", "tool_calls": [call]})
+        msgs.append({"role": "tool", "tool_call_id": cid, "content": result})
+    return msgs
+
+def post(base, path, payload, timeout=1800):
+    request = urllib.request.Request(
+        base + path, data=json.dumps(payload).encode(),
+        headers={"content-type": "application/json"})
+    with urllib.request.urlopen(request, timeout=timeout) as response:
+        return json.loads(response.read())
+
+
+def run_once(base, model, messages, max_tokens):
+    """Return the generated tokens, including the tool-call control tokens."""
+    body = {"model": model, "messages": messages, "tools": TOOLS47,
+            "temperature": 0, "seed": 1234, "max_tokens": max_tokens,
+            "logprobs": True, "top_logprobs": 1}
+    choice = post(base, "/v1/chat/completions", body)["choices"][0]
+    return [t["token"] for t in (choice.get("logprobs") or {}).get("content") or []]
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--url", default="http://127.0.0.1:8000")
+    ap.add_argument("--model", required=True)
+    ap.add_argument("--max-tokens", type=int, default=40)
+    ap.add_argument("--runs", type=int, default=40)
+    ap.add_argument("--concurrency", type=int, default=1)
+    a = ap.parse_args()
+    base = a.url.rstrip("/")
+
+    messages = build("presstest")
+    n_tokens = post(base, "/tokenize",
+                    {"model": a.model, "messages": messages, "tools": TOOLS47,
+                     "add_generation_prompt": True}, timeout=300)["count"]
+    print(f"prompt: {n_tokens} tokens, {a.runs} runs, temperature 0, seed 1234",
+          flush=True)
+
+    runs, lock = [], threading.Lock()
+    work = queue.Queue()
+    for i in range(a.runs):
+        work.put(i)
+
+    def worker():
+        while True:
+            try:
+                work.get_nowait()
+            except queue.Empty:
+                return
+            started = time.time()
+            try:
+                tokens = run_once(base, a.model, messages, a.max_tokens)
+            except Exception as exc:
+                tokens = [f"<error {exc!r}>"]
+            with lock:
+                runs.append(tokens)
+                print(f"  {len(runs):>3}/{a.runs} {time.time() - started:5.1f}s "
+                      f"{''.join(tokens)[:80]!r}", flush=True)
+
+    threads = [threading.Thread(target=worker) for _ in range(a.concurrency)]
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join()
+
+    texts = ["".join(t) for t in runs]
+    groups = collections.Counter(texts)
+    majority = groups.most_common(1)[0][0]
+    reference = runs[texts.index(majority)]
+    divergent = sum(n for text, n in groups.items() if text != majority)
+
+    print(f"\n{divergent} of {a.runs} runs diverge from the majority "
+          f"({len(groups)} distinct completions)")
+    print(f"\nmajority ({groups[majority]} runs):\n  {majority!r}")
+    for text, n in groups.most_common()[1:]:
+        tokens = runs[texts.index(text)]
+        pos = next((i for i in range(min(len(tokens), len(reference)))
+                    if tokens[i] != reference[i]), min(len(tokens), len(reference)))
+        got = tokens[pos] if pos < len(tokens) else "<end>"
+        want = reference[pos] if pos < len(reference) else "<end>"
+        print(f"\n{n} run(s) diverge at token {pos}: got {got!r}, expected {want!r}"
+              f"\n  {text!r}")
+
+    raise SystemExit(0 if divergent == 0 else 1)
+
+
+if __name__ == "__main__":
+    main()

docker-images/glm53-kindling/dev/repro/topk_ties.py

@@ -0,0 +1,43 @@
+#!/usr/bin/env python3
+"""top_k_per_row_prefill keeps a run-dependent subset of pools tied at the k-th score.
+
+No checkpoint and no server: one CUDA device in a container that has vLLM
+importable.
+
+    python topk_ties.py
+
+Each row gets random scores plus two extra copies of its k-th largest score,
+so the cut falls inside an exact tie. The same input then goes through the
+kernel 200 times, and the distinct outputs are counted (each row's ids sorted,
+so only the chosen set counts). One row alone comes back the same every time.
+A launch with as many rows as a real prefill chunk comes back different on
+almost every call.
+"""
+import torch
+
+import vllm._custom_ops as ops
+
+K = 512
+CALLS = 200
+
+
+def distinct_outputs(rows, cols, extra_ties):
+    torch.manual_seed(0)
+    scores = torch.randn(rows, cols, device="cuda")
+    kth = scores.topk(K, dim=1).values[:, -1:]
+    # Put the copies on entries below the cut, chosen at random per row.
+    below = torch.rand(rows, cols, device="cuda") + (scores >= kth).float() * 2
+    scores.scatter_(1, below.argsort(dim=1)[:, :extra_ties], kth.expand(-1, extra_ties))
+    starts = torch.zeros(rows, dtype=torch.int32, device="cuda")
+    ends = torch.full((rows,), cols, dtype=torch.int32, device="cuda")
+    seen = set()
+    for _ in range(CALLS):
+        out = torch.full((rows, K), -1, dtype=torch.int32, device="cuda")
+        ops.top_k_per_row_prefill(scores, starts, ends, out, rows, scores.stride(0), scores.stride(1), K)
+        seen.add(out.sort(dim=1).values.cpu().numpy().tobytes())
+    return len(seen)
+
+
+for rows, cols in [(1, 1300), (1, 8000), (64, 1300), (16128, 1300), (16128, 8000)]:
+    print(f"{rows:5d} rows x {cols} pools, 2 extra tied at the cut: "
+          f"{distinct_outputs(rows, cols, 2)} distinct outputs in {CALLS} calls")

docker-images/glm53-kindling/dev/tap/README.md

@@ -0,0 +1,100 @@
+# glm53 step tap
+
+Per-step, per-rank instrumentation for the GLM-5.3-Flash corruption hunt.
+Opt-in through `GLM53_TAP=1`; costs one env lookup when off.
+
+Files:
+
+- `glm53_step_tap.py`  the tap (runtime module, no rebuild needed)
+- `glm53_step_tap.pth` one-line loader; Python imports it at interpreter start
+  from any site directory
+- `glm53_tap_report.py` offline report over the per-rank JSONL files
+
+## What it records (per step, per rank)
+
+- rows: request id, state index, position, input token, device `seq_len`,
+  host `seq_lens_cpu_upper_bound`, host and device `num_computed_tokens`
+- pairs: rows in the same batch with the same position and the same token
+  prefix (the S2 setup). Every checkpoint is compared row A vs row B on the
+  device, in execution order; `first_diff` names the first checkpoint where
+  the rows differ, `diffs` lists them all.
+- checkpoints without `GLM53_TAP_DEEP` (work under torch.compile and
+  piecewise/breakable graphs, since they sit at custom-op boundaries):
+  - `L<i>.kda.{qkv,g1,beta}` inputs, `L<i>.kda.out` output,
+    `L<i>.kda.{conv,rec}_{before,after}` the request's own state block,
+    `L<i>.kda.state_idx`
+  - `L<i>.mla.{q,kvc,kpe}` inputs, `L<i>.mla.attn_out`, `L<i>.mla.slot`,
+    `L<i>.mla.cache_row` (the cache row just written, read back),
+    `L<i>.mla.cache_reread` (the row written last step, read again)
+  - `L<i>.mla.{planned,valid}` SM90 planned kv_len vs the indexer's valid
+    count for the row, `L<i>.mla.topk` logical indices, `L<i>.mla.kvidx`
+    physical indices, `L<i>.mla.out` the SM90 output
+  - `L<i>.idx.{hidden,q,k,w,gate,topk,slot,tail_slot}`
+- checkpoints with `GLM53_TAP_DEEP=1` (pure eager only): forward hooks on
+  every decoder layer and its `mhc_pre_op`, `mhc_fused_post_pre_op`,
+  `self_attn`, `mhc_post_op`, `mlp`, `input_layernorm`,
+  `post_attention_layernorm`, as `L<i>.<sub>.in<k>` / `.out<k>`. A second
+  call of the same module in a layer gets `@1` (the fused hc op runs twice:
+  `hc_fused` = attn pre, `hc_fused@1` = ffn pre).
+- events: `host_dev_mismatch`, `plan_mismatch`, `state_changed` (a KDA
+  state block read at step t differs from what the same request wrote at
+  t-1), `kv_changed` (an MLA cache row changed between the write and the
+  next step's re-read), sampler `argmax`/`top2`/`gap`/`sampled` for the rows
+  this rank samples, `fed` = `last_sampled_tokens` per request after
+  post_update (what this rank feeds next), `nc_after`.
+
+## Running it
+
+Bind-mount the module and loader into the image's site dir on every box,
+and set the env. No image rebuild. With the compose layering, put the env in
+`.env` (it beats compose and the entrypoint).
+
+```
+# on each GLM box, e.g. as an extra compose volume / env, or on a throwaway run:
+-v /home/admin/tap/glm53_step_tap.py:/usr/local/lib/python3.12/dist-packages/glm53_step_tap.py:ro
+-v /home/admin/tap/glm53_step_tap.pth:/usr/local/lib/python3.12/dist-packages/glm53_step_tap.pth:ro
+-v /home/admin/tap/out:/tmp/glm53-tap
+-e GLM53_TAP=1 -e GLM53_TAP_DIR=/tmp/glm53-tap -e GLM53_TAP_VECTORS=1
+# for sublayer-level localisation:
+-e GLM53_TAP_DEEP=1   plus  --enforce-eager  (EXTRA_ARGS; eager still reproduces per EXP2)
+```
+
+Worker processes inherit the env and run site at start, so every rank taps
+itself; the stderr line `[glm53-tap] rank N armed, writing ...` confirms it.
+The vLLM worker log also shows `[glm53-tap] runner patched` / `KDA patched`
+/ `MLAAttention patched` / `SM90 forward_mqa patched` / `indexer patched` /
+`sampler patched` as the modules import.
+
+Rebuild variant: `COPY patches/tap/glm53_step_tap.py patches/tap/glm53_step_tap.pth /usr/local/lib/python3.12/dist-packages/`.
+
+Then send the S2 pair (two identical raw-token-id completions in flight
+together), collect `tap-rank*.jsonl` from all four boxes into one directory,
+and run:
+
+```
+python3 glm53_tap_report.py DIR                 # pairs, events, cross-rank, tokens
+python3 glm53_tap_report.py DIR --req cmpl-abc  # one request
+python3 glm53_tap_report.py DIR --steps 200-320 --all-pairs
+```
+
+Reading the pair section: `FIRST SPLIT at L7.mlp.out0` with `L7.mlp.in0`
+absent from `differing` means the MoE of layer 7 split two identical rows;
+`L12.kda.out` first with `L12.kda.qkv` and `L12.kda.{conv,rec}_before`
+equal means the KDA kernels did; `L12.kda.rec_before` first means the two
+requests' state blocks already differed when the step began (look at the
+earlier steps' `state_changed`). `L3.mla.out` first with `L3.mla.q`,
+`L3.mla.topk`, `L3.mla.cache_row` equal points at the SM90 kernel or its
+plan (`plan_mismatch` says whether the planned length was wrong).
+
+Vectors: `GLM53_TAP_VECTORS=1` dumps every checkpoint's tapped rows to
+`vec-rank<r>-step<N>.pt` at the step a pair first splits (2 = every step;
+about 140 MB per step per rank with DEEP on). Load with `torch.load`; the
+`cps` dict maps checkpoint name to `[rows, bytes-as-int32]` tensors in the
+same row order as `rows`.
+
+Self-test (no vLLM):
+
+```
+sudo docker run --rm --gpus all -v /home/admin/tap:/tap -e GLM53_TAP=1 \
+  --entrypoint python3 localhost:5000/spark-glm53:v4 /tap/glm53_step_tap.py --selftest
+```

docker-images/glm53-kindling/dev/tap/glm53_step_tap.pth

@@ -0,0 +1 @@
+import glm53_step_tap

docker-images/glm53-kindling/dev/tap/glm53_step_tap.py

@@ -0,0 +1,1210 @@
+"""Per-step tap for GLM-5.3-Flash under vLLM's V2 GPU runner.
+
+What it answers, per decode step and per TP rank:
+
+  * Pairs: two rows in the same batch with the same position and the same
+    token prefix (the S2 setup) are compared sublayer by sublayer, in
+    execution order, and the FIRST checkpoint where the rows differ is
+    logged with the layer and sublayer name. Rows that differ at a
+    sublayer's output but not at its input name the component.
+  * Placement: KV slot, mamba state block, indexer pool slot and tail slot
+    for every tapped row, plus the SM90 planned kv_len against the indexer's
+    valid count for that row.
+  * Continuity: the KDA conv/recurrent state of a request is hashed before
+    and after every step; a "before" that differs from the previous step's
+    "after" for the same block means something else wrote it in between.
+    The MLA cache row written at step t is re-read at step t+1 the same way.
+  * Feed: the token each rank feeds next (last_sampled_tokens after
+    post_update) and the host/device num_computed_tokens, so cross-rank
+    disagreement shows up in the per-rank files.
+
+Everything is off unless GLM53_TAP=1; the import then costs one env lookup.
+When armed, the hooks install lazily through a post-import hook on the vLLM
+modules, so importing this file never imports torch or vLLM by itself.
+
+Env:
+  GLM53_TAP=1             arm
+  GLM53_TAP_DIR           output dir, default /tmp/glm53-tap (tap-rank<r>.jsonl)
+  GLM53_TAP_REQ           substring filter on request id ("" = every request)
+  GLM53_TAP_DEEP=1        forward hooks on every decoder sublayer. Needs pure
+                          eager (--enforce-eager); refused under torch.compile
+                          unless GLM53_TAP_DEEP=force.
+  GLM53_TAP_VECTORS=0|1|2 0 hashes only; 1 also dump the tapped rows of every
+                          checkpoint at the step a pair first splits; 2 every step
+  GLM53_TAP_VEC_STEPS=A-B dump vectors for steps A..B regardless of pairs
+  GLM53_TAP_MAX_ROWS      rows recorded per step (default 64; the last N
+                          matching tokens of the batch)
+  GLM53_TAP_PAIR_ALL=1    pair rows by position only, skip the prefix check
+  GLM53_TAP_SELFTEST=1    (with `python3 -m glm53_step_tap`) exercise the
+                          machinery on a toy module, no vLLM needed
+
+Loading without a rebuild: bind-mount this file and glm53_step_tap.pth into
+/usr/local/lib/python3.12/dist-packages/ (see README.md).
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+
+_ENABLED = os.environ.get("GLM53_TAP", "0").strip().lower() not in ("", "0", "no", "false")
+
+
+def _env_int(name: str, default: int) -> int:
+    try:
+        return int(os.environ.get(name, "") or default)
+    except ValueError:
+        return default
+
+
+# --------------------------------------------------------------------------
+# Post-import hook: run a callback right after a named module finishes import.
+# --------------------------------------------------------------------------
+
+
+def _install_post_import_hooks(targets: dict) -> None:
+    import importlib.abc
+    import importlib.util
+
+    class _Finder(importlib.abc.MetaPathFinder):
+        def __init__(self) -> None:
+            self._busy: set[str] = set()
+
+        def find_spec(self, fullname, path=None, target=None):
+            if fullname not in targets or fullname in self._busy:
+                return None
+            self._busy.add(fullname)
+            try:
+                spec = importlib.util.find_spec(fullname)
+            finally:
+                self._busy.discard(fullname)
+            if spec is None or spec.loader is None:
+                return None
+            orig = spec.loader
+            cb = targets[fullname]
+
+            class _Loader(importlib.abc.Loader):
+                def create_module(self, spec_):
+                    if hasattr(orig, "create_module"):
+                        return orig.create_module(spec_)
+                    return None
+
+                def exec_module(self, module):
+                    orig.exec_module(module)
+                    try:
+                        cb(module)
+                    except Exception:  # noqa: BLE001
+                        _log_exc(f"post-import hook for {fullname}")
+
+            spec.loader = _Loader()
+            return spec
+
+    sys.meta_path.insert(0, _Finder())
+
+
+def _log_exc(where: str) -> None:
+    import traceback
+
+    sys.stderr.write(f"[glm53-tap] error in {where}:\n{traceback.format_exc()}\n")
+    sys.stderr.flush()
+
+
+# --------------------------------------------------------------------------
+# The tap itself (imports torch lazily; only constructed when armed).
+# --------------------------------------------------------------------------
+
+_SHARDED = ("kda.", "mla.q", "mla.out", "mla.kvidx")
+
+
+class Tap:
+    def __init__(self) -> None:
+        import torch
+
+        self.torch = torch
+        self.dir = os.environ.get("GLM53_TAP_DIR", "/tmp/glm53-tap")
+        self.req_filter = os.environ.get("GLM53_TAP_REQ", "")
+        deep = os.environ.get("GLM53_TAP_DEEP", "0").strip().lower()
+        self.deep = deep not in ("", "0", "no", "false")
+        self.deep_force = deep == "force"
+        self.vectors = _env_int("GLM53_TAP_VECTORS", 0)
+        vs = os.environ.get("GLM53_TAP_VEC_STEPS", "")
+        self.vec_steps = None
+        if vs:
+            a, _, b = vs.partition("-")
+            self.vec_steps = (int(a or 0), int(b or 1 << 40))
+        self.max_rows = max(1, _env_int("GLM53_TAP_MAX_ROWS", 64))
+        self.pair_all = os.environ.get("GLM53_TAP_PAIR_ALL", "0") not in ("", "0")
+        self.rank: int | None = None
+        self.fh = None
+        self.step = 0
+        self.header_written = False
+        # per-step state
+        self.cur: dict | None = None
+        self.n_tokens = 0
+        self.n_pad = 0
+        self.tapped: list[int] = []
+        self.tapped_t = None
+        self.req_of_tok = None  # CPU LongTensor [n_tokens]
+        self.row_req: list[str] = []
+        self.row_pos: list[int] = []
+        self.batch_req_ids: list[str] = []
+        self.pairs: list[tuple[int, int]] = []
+        self.cps: list[tuple[str, object]] = []  # (name, rows tensor [k, m] int view)
+        self.ints: dict[str, object] = {}
+        self.ints_cpu: dict[str, list] = {}
+        self.sample_rows: list[dict] = []
+        self.extra: dict[str, object] = {}  # per-step, saved only in the vector blob
+        self.err: list[str] = []
+        self.call_counts: dict[str, int] = {}
+        # cross-step memory
+        self.last_state: dict[tuple[str, int], int] = {}  # (cpname, state_idx) -> hash
+        self.kv_prev: dict[tuple[str, str], tuple[int, int]] = {}  # (layer, req) -> (slot, hash)
+        self.pair_split: dict[tuple[str, str], int] = {}  # first step a pair split
+        self.runner = None
+        self._w = None
+
+    # ---- utilities -------------------------------------------------------
+
+    def note_err(self, where: str) -> None:
+        import traceback
+
+        msg = f"{where}: {traceback.format_exc(limit=3)}"
+        if self.cur is not None:
+            self.cur.setdefault("err", []).append(msg)
+        sys.stderr.write(f"[glm53-tap] {msg}\n")
+
+    def _rank(self) -> int:
+        if self.rank is not None:
+            return self.rank
+        r = None
+        try:
+            from vllm.distributed.parallel_state import get_tensor_model_parallel_rank
+
+            r = int(get_tensor_model_parallel_rank())
+        except Exception:  # noqa: BLE001
+            r = None
+        if r is None:
+            try:
+                r = int(os.environ.get("RANK", ""))
+            except ValueError:
+                r = -1
+        self.rank = r
+        return r
+
+    def _open(self):
+        if self.fh is None:
+            os.makedirs(self.dir, exist_ok=True)
+            r = self._rank()
+            path = os.path.join(self.dir, f"tap-rank{r}.jsonl")
+            self.fh = open(path, "a", buffering=1)
+            sys.stderr.write(f"[glm53-tap] rank {r} armed, writing {path}\n")
+        return self.fh
+
+    def _weights(self, m: int, device):
+        torch = self.torch
+        if self._w is None or self._w.numel() < m or self._w.device != device:
+            n = max(m, 1 << 16)
+            w = torch.arange(1, n + 1, dtype=torch.int64, device=device)
+            self._w = w * 0x9E3779B1 | 1
+        return self._w[:m]
+
+    def _row_hash(self, iv):
+        """[k, m] int32/uint8 -> [k] int64 order-sensitive hash (device)."""
+        w = self._weights(iv.shape[1], iv.device)
+        return (iv.to(self.torch.int64) * w).sum(dim=1)
+
+    def _rows_view(self, t):
+        """Return [n_tokens_view, -1] view if t has a token dim, else None."""
+        torch = self.torch
+        if not torch.is_tensor(t) or t.dim() == 0 or t.numel() == 0:
+            return None
+        if isinstance(t, torch.nn.Parameter) or t.requires_grad:
+            return None
+        n, npad = self.n_tokens, self.n_pad
+        s = t.shape
+        if s[0] in (n, npad):
+            return t.reshape(s[0], -1)
+        if t.dim() >= 2 and s[0] == 1 and s[1] in (n, npad):
+            return t.reshape(s[1], -1)
+        return None
+
+    def _as_int_rows(self, rows):
+        torch = self.torch
+        rows = rows.contiguous()
+        if rows.dtype == torch.bool:
+            rows = rows.to(torch.uint8)
+        k, m = rows.shape
+        es = rows.element_size()
+        if (m * es) % 4 == 0 and rows.dtype != torch.uint8:
+            return rows.view(torch.int32) if es != 4 else rows.view(torch.int32)
+        return rows.view(torch.uint8)
+
+    def capturing(self) -> bool:
+        try:
+            return bool(self.torch.cuda.is_current_stream_capturing())
+        except Exception:  # noqa: BLE001
+            return False
+
+    # ---- checkpoints ------------------------------------------------------
+
+    def cp(self, name: str, t) -> None:
+        """Record tensor rows of the tapped tokens under a checkpoint name."""
+        if self.cur is None or self.tapped_t is None:
+            return
+        try:
+            v = self._rows_view(t)
+            if v is None:
+                return
+            if self.capturing():
+                return
+            rows = v.index_select(0, self.tapped_t)
+            self.cps.append((name, self._as_int_rows(rows), rows))
+        except Exception:  # noqa: BLE001
+            self.note_err(f"cp {name}")
+
+    def cp_rows(self, name: str, rows) -> None:
+        """Record an already row-selected tensor [k, ...] (k = tapped rows)."""
+        if self.cur is None or self.tapped_t is None:
+            return
+        try:
+            if self.capturing():
+                return
+            k = len(self.tapped)
+            rows = rows.reshape(k, -1)
+            self.cps.append((name, self._as_int_rows(rows), rows))
+        except Exception:  # noqa: BLE001
+            self.note_err(f"cp_rows {name}")
+
+    def ival(self, name: str, t) -> None:
+        """Small per-tapped-row integer values (device or CPU tensor / list)."""
+        if self.cur is None:
+            return
+        try:
+            torch = self.torch
+            if torch.is_tensor(t):
+                if t.device.type != "cpu":
+                    if self.capturing():
+                        return
+                    self.ints[name] = t.detach().reshape(-1).to(torch.int64)
+                else:
+                    self.ints_cpu[name] = [int(x) for x in t.reshape(-1).tolist()]
+            else:
+                self.ints_cpu[name] = [int(x) for x in t]
+        except Exception:  # noqa: BLE001
+            self.note_err(f"ival {name}")
+
+    def ival_tok(self, name: str, t) -> None:
+        """Per-token device tensor [n_tokens] -> keep only tapped rows."""
+        if self.cur is None or self.tapped_t is None:
+            return
+        try:
+            if self.capturing():
+                return
+            self.ints[name] = t.reshape(-1)[: self.n_pad].index_select(0, self.tapped_t).to(self.torch.int64)
+        except Exception:  # noqa: BLE001
+            self.note_err(f"ival_tok {name}")
+
+    def layer_of(self, prefix: str) -> str:
+        # "model.layers.12.self_attn..." -> "L12"
+        parts = prefix.split(".")
+        for i, p in enumerate(parts):
+            if p == "layers" and i + 1 < len(parts) and parts[i + 1].isdigit():
+                return f"L{parts[i + 1]}"
+        return prefix
+
+    def count(self, key: str) -> int:
+        c = self.call_counts.get(key, 0)
+        self.call_counts[key] = c + 1
+        return c
+
+    # ---- step lifecycle ---------------------------------------------------
+
+    def begin_step(self, runner, input_batch) -> None:
+        torch = self.torch
+        self.cur = None
+        self.cps = []
+        self.ints = {}
+        self.ints_cpu = {}
+        self.sample_rows = []
+        self.extra = {}
+        self.call_counts = {}
+        self.pairs = []
+        self.tapped = []
+        self.tapped_t = None
+        self.runner = runner
+        try:
+            req_ids = list(input_batch.req_ids)
+            n = int(input_batch.num_tokens)
+            npad = int(input_batch.num_tokens_after_padding)
+            nreq = int(input_batch.num_reqs)
+            qsl = [int(x) for x in input_batch.query_start_loc_np[: nreq + 1]]
+            match = [i for i, r in enumerate(req_ids) if (not self.req_filter) or (self.req_filter in r)]
+            if not match:
+                return
+            req_of_tok = torch.empty(n, dtype=torch.int64)
+            for r in range(nreq):
+                req_of_tok[qsl[r] : qsl[r + 1]] = r
+            toks = [t for t in range(n) if int(req_of_tok[t]) in set(match)]
+            toks = toks[-self.max_rows :]
+            if not toks:
+                return
+            pos = input_batch.positions[:n].detach().to("cpu")
+            ids = input_batch.input_ids[:n].detach().to("cpu")
+            seq_dev = input_batch.seq_lens[:nreq].detach().to("cpu")
+            seq_ub = input_batch.seq_lens_cpu_upper_bound[:nreq]
+            nc_host = input_batch.num_computed_tokens_np[:nreq]
+            idx_np = input_batch.idx_mapping_np[:nreq]
+            nc_dev = runner.req_states.num_computed_tokens.gpu[input_batch.idx_mapping[:nreq]].detach().to("cpu")
+            self.step += 1
+            self.n_tokens, self.n_pad = n, npad
+            self.req_of_tok = req_of_tok
+            self.batch_req_ids = req_ids
+            self.tapped = toks
+            self.tapped_t = torch.tensor(toks, dtype=torch.int64, device=input_batch.positions.device)
+            rows = []
+            self.row_req, self.row_pos = [], []
+            for t in toks:
+                r = int(req_of_tok[t])
+                rows.append(
+                    {
+                        "i": t,
+                        "req": req_ids[r],
+                        "state": int(idx_np[r]),
+                        "pos": int(pos[t]),
+                        "tok": int(ids[t]),
+                        "seq_dev": int(seq_dev[r]),
+                        "seq_ub": int(seq_ub[r]),
+                        "nc_host": int(nc_host[r]),
+                        "nc_dev": int(nc_dev[r]),
+                    }
+                )
+                self.row_req.append(req_ids[r])
+                self.row_pos.append(int(pos[t]))
+            self.cur = {
+                "step": self.step,
+                "rank": self._rank(),
+                "t": __import__("time").time(),
+                "n_tokens": n,
+                "n_reqs": nreq,
+                "rows": rows,
+            }
+            hd = [x["req"] for x in rows if x["nc_host"] != x["nc_dev"] or x["seq_dev"] != x["seq_ub"]]
+            if hd:
+                self.cur["host_dev_mismatch"] = sorted(set(hd))
+            # pairs: distinct requests, same position (+ same prefix unless PAIR_ALL)
+            all_ids = runner.req_states.all_token_ids.gpu
+            for a in range(len(toks)):
+                for b in range(a + 1, len(toks)):
+                    if len(self.pairs) >= 8:
+                        break
+                    if self.row_req[a] == self.row_req[b] or self.row_pos[a] != self.row_pos[b]:
+                        continue
+                    if rows[a]["tok"] != rows[b]["tok"]:
+                        continue
+                    if not self.pair_all:
+                        p = self.row_pos[a]
+                        sa, sb = rows[a]["state"], rows[b]["state"]
+                        if p > 0 and not bool(torch.equal(all_ids[sa, :p], all_ids[sb, :p])):
+                            continue
+                    self.pairs.append((a, b))
+            self.cur["pairs"] = [
+                {"a": self.row_req[a], "b": self.row_req[b], "pos": self.row_pos[a]} for a, b in self.pairs
+            ]
+        except Exception:  # noqa: BLE001
+            self.cur = None
+            self.note_err("begin_step")
+
+    def flush(self) -> None:
+        if self.cur is None:
+            return
+        cur = self.cur
+        torch = self.torch
+        try:
+            k = len(self.tapped)
+            names = [n for n, _, _ in self.cps]
+            # hashes: one stack, one D2H
+            if self.cps:
+                hs = torch.stack([self._row_hash(iv) for _, iv, _ in self.cps])  # [C, k]
+                hs_cpu = hs.to("cpu").tolist()
+            else:
+                hs_cpu = []
+            cps_out = {}
+            for i, n in enumerate(names):
+                nm = n if n not in cps_out else f"{n}#{i}"
+                cps_out[nm] = hs_cpu[i]
+            cur["cps"] = cps_out
+            # pair diffs (physical kv_indices differ between rows by construction)
+            if self.pairs and self.cps:
+                cmp_idx = [c for c, n in enumerate(names) if not n.endswith(".mla.kvidx")]
+                flags = []
+                for a, b in self.pairs:
+                    flags.append(torch.stack([(self.cps[c][1][a] != self.cps[c][1][b]).any() for c in cmp_idx]))
+                flags = torch.stack(flags).to("cpu")  # [P, C']
+                for pi, (a, b) in enumerate(self.pairs):
+                    diff = [names[c] for j, c in enumerate(cmp_idx) if bool(flags[pi, j])]
+                    first = diff[0] if diff else None
+                    key = (self.row_req[a], self.row_req[b])
+                    rec = cur["pairs"][pi]
+                    rec["first_diff"] = first
+                    rec["ndiff"] = len(diff)
+                    rec["diffs"] = diff[:60]
+                    rec["ncp"] = len(names)
+                    if first is not None and key not in self.pair_split:
+                        self.pair_split[key] = self.step
+                        rec["first_split_step"] = True
+                        # max-abs at the first diverging checkpoint, in the cp's own dtype view
+                        try:
+                            c = names.index(first)
+                            iv = self.cps[c][1]
+                            rec["first_diff_nbytes_ne"] = int(
+                                (iv[a].view(torch.uint8) != iv[b].view(torch.uint8)).sum().item()
+                            )
+                        except Exception:  # noqa: BLE001
+                            pass
+            # ints
+            ints_out = dict(self.ints_cpu)
+            for n, t in self.ints.items():
+                ints_out[n] = t.to("cpu").tolist()
+            cur["ints"] = ints_out
+            # continuity: KDA state before(t) vs after(t-1); MLA row re-read
+            changed = []
+            for n, h in cps_out.items():
+                if not (n.endswith(".conv_before") or n.endswith(".rec_before")):
+                    continue
+                sidx_name = n.rsplit(".", 1)[0] + ".state_idx"
+                sidx = ints_out.get(sidx_name)
+                if sidx is None:
+                    continue
+                for r in range(k):
+                    key = (n, int(sidx[r]))
+                    prev = self.last_state.get(key)
+                    if prev is not None and prev != h[r]:
+                        changed.append({"cp": n, "req": self.row_req[r], "state": int(sidx[r]), "pos": self.row_pos[r]})
+            for n, h in cps_out.items():
+                if n.endswith(".conv_after") or n.endswith(".rec_after"):
+                    sidx = ints_out.get(n.rsplit(".", 1)[0] + ".state_idx")
+                    if sidx is None:
+                        continue
+                    for r in range(k):
+                        self.last_state[(n.replace("_after", "_before"), int(sidx[r]))] = h[r]
+            if changed:
+                cur["state_changed"] = changed
+            kvch = []
+            for n, h in cps_out.items():
+                if not n.endswith(".mla.cache_reread"):
+                    continue
+                layer = n.split(".")[0]
+                for r in range(k):
+                    prev = self.kv_prev.get((layer, self.row_req[r]))
+                    if prev is not None and prev[1] != h[r] and h[r] != 0:
+                        kvch.append({"layer": layer, "req": self.row_req[r], "slot": prev[0], "pos": self.row_pos[r]})
+            for n, h in cps_out.items():
+                if n.endswith(".mla.cache_row"):
+                    layer = n.split(".")[0]
+                    slots = ints_out.get(f"{layer}.mla.slot")
+                    if slots is None:
+                        continue
+                    for r in range(k):
+                        self.kv_prev[(layer, self.row_req[r])] = (int(slots[r]), h[r])
+            if kvch:
+                cur["kv_changed"] = kvch
+            # SM90 plan vs valid
+            pm = []
+            for n, v in ints_out.items():
+                if n.endswith(".mla.planned"):
+                    valid = ints_out.get(n.replace(".planned", ".valid"))
+                    if valid is None:
+                        continue
+                    for r in range(k):
+                        if v[r] != valid[r]:
+                            pm.append({"cp": n, "req": self.row_req[r], "planned": v[r], "valid": valid[r], "pos": self.row_pos[r]})
+            if pm:
+                cur["plan_mismatch"] = pm
+            if self.sample_rows:
+                cur["sample"] = self.sample_rows
+            # vectors
+            want_vec = self.vectors >= 2 or (
+                self.vectors == 1 and any(p.get("first_split_step") for p in cur.get("pairs", []))
+            )
+            # A window file (GLM53_TAP_DIR/vec_steps, "A-B") overrides the env var and is
+            # re-read every step, so the window can be set after the boot's self-test
+            # requests have shifted the step count.
+            try:
+                with open(os.path.join(self.dir, "vec_steps")) as f:
+                    a, _, b = f.read().strip().partition("-")
+                    self.vec_steps = (int(a or 0), int(b or 1 << 40))
+            except (OSError, ValueError):
+                pass
+            if self.vec_steps is not None:
+                want_vec = self.vec_steps[0] <= self.step <= self.vec_steps[1]
+            if want_vec and self.cps:
+                blob = {"rows": cur["rows"], "cps": {}, "extra": {}}
+                def _cpu(v):
+                    if torch.is_tensor(v):
+                        return v.to("cpu")
+                    if isinstance(v, (list, tuple)):
+                        return type(v)(_cpu(x) for x in v)
+                    if isinstance(v, dict):
+                        return {kk: _cpu(vv) for kk, vv in v.items()}
+                    return v
+
+                for en, ev in self.extra.items():
+                    try:
+                        blob["extra"][en] = _cpu(ev)
+                    except Exception:  # noqa: BLE001
+                        pass
+                for i, (n, _iv, rows) in enumerate(self.cps):
+                    nm = n if n not in blob["cps"] else f"{n}#{i}"
+                    blob["cps"][nm] = rows.to("cpu")
+                os.makedirs(self.dir, exist_ok=True)
+                p = os.path.join(self.dir, f"vec-rank{self._rank()}-step{self.step}.pt")
+                torch.save(blob, p)
+                cur["vectors"] = p
+        except Exception:  # noqa: BLE001
+            self.note_err("flush")
+        try:
+            import json
+
+            fh = self._open()
+            if not self.header_written:
+                self.header_written = True
+                fh.write(json.dumps({"header": True, "rank": self._rank(), "pid": os.getpid(), **self._config_summary()}) + "\n")
+            fh.write(json.dumps(cur, default=str) + "\n")
+        except Exception:  # noqa: BLE001
+            _log_exc("flush write")
+        self.cur = None
+        self.cps = []
+        self.ints = {}
+        self.ints_cpu = {}
+        self.sample_rows = []
+
+    def _config_summary(self) -> dict:
+        out = {"deep": self.deep, "vectors": self.vectors, "req_filter": self.req_filter}
+        try:
+            cfg = self.runner.vllm_config
+            cc = cfg.compilation_config
+            out.update(
+                {
+                    "compile_mode": str(cc.mode),
+                    "cudagraph_mode": str(cc.cudagraph_mode),
+                    "tp": int(cfg.parallel_config.tensor_parallel_size),
+                    "async_scheduling": bool(cfg.scheduler_config.async_scheduling),
+                    "max_concurrent_batches": int(getattr(cfg, "max_concurrent_batches", 0) or 0),
+                    "block_size": int(cfg.cache_config.block_size),
+                    "kv_cache_dtype": str(cfg.cache_config.cache_dtype),
+                    "enforce_eager": bool(cfg.model_config.enforce_eager),
+                }
+            )
+        except Exception:  # noqa: BLE001
+            pass
+        return out
+
+
+TAP: Tap | None = None
+
+
+def _tap() -> Tap:
+    global TAP
+    if TAP is None:
+        TAP = Tap()
+    return TAP
+
+
+# --------------------------------------------------------------------------
+# Patches per module
+# --------------------------------------------------------------------------
+
+
+def _patch_runner(mod) -> None:
+    import functools
+
+    R = mod.GPUModelRunner
+    tap = _tap()
+
+    o_prepare = R.prepare_inputs
+
+    @functools.wraps(o_prepare)
+    def prepare_inputs(self, *a, **k):
+        ib = o_prepare(self, *a, **k)
+        try:
+            tap.begin_step(self, ib)
+        except Exception:  # noqa: BLE001
+            tap.note_err("prepare_inputs")
+        return ib
+
+    R.prepare_inputs = prepare_inputs
+
+    o_exec = R.execute_model
+
+    @functools.wraps(o_exec)
+    def execute_model(self, *a, **k):
+        # a step whose sample_tokens never ran (or a dummy run) must not leak
+        if tap.cur is not None:
+            try:
+                tap.cur.setdefault("err", []).append("flushed at next execute_model (sample_tokens not seen)")
+                tap.flush()
+            except Exception:  # noqa: BLE001
+                tap.note_err("late flush")
+        tap.cur = None
+        if k.get("dummy_run") or k.get("is_profile") or (len(a) > 2 and a[2]):
+            tap.tapped_t = None
+        return o_exec(self, *a, **k)
+
+    R.execute_model = execute_model
+
+    o_post = R.postprocess_sampled
+
+    @functools.wraps(o_post)
+    def postprocess_sampled(self, idx_mapping, sampled_tokens, num_sampled, num_rejected, *a, **k):
+        out = o_post(self, idx_mapping, sampled_tokens, num_sampled, num_rejected, *a, **k)
+        try:
+            if tap.cur is not None and not tap.capturing():
+                fed = self.req_states.last_sampled_tokens[idx_mapping].detach().to("cpu").tolist()
+                nc = self.req_states.num_computed_tokens.gpu[idx_mapping].detach().to("cpu").tolist()
+                st = sampled_tokens[:, 0].detach().to("cpu").tolist() if sampled_tokens.dim() == 2 else []
+                reqs = tap.batch_req_ids or [str(i) for i in range(len(fed))]
+                tap.cur["fed"] = {r: (v if isinstance(v, list) else int(v)) for r, v in zip(reqs, fed)}  # [n, 1] in this nightly
+                tap.cur["nc_after"] = {r: int(v) for r, v in zip(reqs, nc)}
+                if st:
+                    tap.cur["sampled0"] = {r: int(v) for r, v in zip(reqs, st)}
+        except Exception:  # noqa: BLE001
+            tap.note_err("postprocess_sampled")
+        return out
+
+    R.postprocess_sampled = postprocess_sampled
+
+    o_sample_tokens = R.sample_tokens
+
+    @functools.wraps(o_sample_tokens)
+    def sample_tokens(self, *a, **k):
+        try:
+            out = o_sample_tokens(self, *a, **k)
+        finally:
+            try:
+                if tap.cur is not None:
+                    tap.flush()
+            except Exception:  # noqa: BLE001
+                tap.note_err("sample_tokens flush")
+        return out
+
+    R.sample_tokens = sample_tokens
+
+    o_load = R.load_model
+
+    @functools.wraps(o_load)
+    def load_model(self, *a, **k):
+        out = o_load(self, *a, **k)
+        try:
+            tap.runner = self
+            if tap.deep:
+                _install_deep_hooks(self, tap)
+        except Exception:  # noqa: BLE001
+            tap.note_err("load_model hooks")
+        return out
+
+    R.load_model = load_model
+    sys.stderr.write("[glm53-tap] runner patched\n")
+
+
+def _install_deep_hooks(runner, tap: Tap) -> None:
+    torch = tap.torch
+    try:
+        mode = runner.vllm_config.compilation_config.mode
+        eager = str(mode).endswith("NONE") or int(getattr(mode, "value", 0) or 0) == 0
+    except Exception:  # noqa: BLE001
+        eager = False
+    if not eager and not tap.deep_force:
+        sys.stderr.write(
+            "[glm53-tap] DEEP hooks refused: torch.compile is on (use --enforce-eager, or GLM53_TAP_DEEP=force)\n"
+        )
+        return
+    model = runner.model
+    layers = [(n, m) for n, m in model.named_modules() if type(m).__name__ == "Glm5NextDecoderLayer"]
+    if not layers:
+        sys.stderr.write("[glm53-tap] DEEP: no Glm5NextDecoderLayer modules found\n")
+        return
+
+    def mk(name: str, in_names: bool = True):
+        def hook(mod, args, kwargs, out):
+            if tap.cur is None:
+                return
+            try:
+                c = tap.count(name)
+                tag = name if c == 0 else f"{name}@{c}"
+                if in_names:
+                    for i, t in enumerate(args):
+                        tap.cp(f"{tag}.in{i}", t)
+                    for kname, t in kwargs.items():
+                        tap.cp(f"{tag}.in_{kname}", t)
+                outs = out if isinstance(out, (tuple, list)) else (out,)
+                for i, t in enumerate(outs):
+                    tap.cp(f"{tag}.out{i}", t)
+            except Exception:  # noqa: BLE001
+                tap.note_err(f"hook {name}")
+
+        return hook
+
+    n_hooks = 0
+    for lname, layer in layers:
+        L = tap.layer_of(lname + ".x")
+        kind = type(getattr(layer, "self_attn", None)).__name__
+        attn_tag = "kda" if "Linear" in kind else "mla"
+        layer.register_forward_hook(mk(f"{L}.layer"), with_kwargs=True)
+        n_hooks += 1
+        for sub, tag in (
+            ("mhc_pre_op", "hc_pre"),
+            ("mhc_fused_post_pre_op", "hc_fused"),
+            ("self_attn", f"attn_{attn_tag}"),
+            ("mhc_post_op", "hc_post"),
+            ("mlp", "mlp"),
+            ("input_layernorm", "ln_in"),
+            ("post_attention_layernorm", "ln_post"),
+        ):
+            m = getattr(layer, sub, None)
+            if isinstance(m, torch.nn.Module):
+                m.register_forward_hook(mk(f"{L}.{tag}"), with_kwargs=True)
+                n_hooks += 1
+        # Inside the MoE: the router (gate logits, top-k derivable offline),
+        # the fused experts (their output before the TP all-reduce) and the
+        # shared experts. mlp.in/out alone cannot show an expert-choice split.
+        mlp = getattr(layer, "mlp", None)
+        for sub, tag in (("gate", "moe_gate"), ("experts", "moe_experts"), ("shared_experts", "moe_shared")):
+            m = getattr(mlp, sub, None)
+            if isinstance(m, torch.nn.Module):
+                m.register_forward_hook(mk(f"{L}.{tag}"), with_kwargs=True)
+                n_hooks += 1
+    sys.stderr.write(f"[glm53-tap] DEEP: {n_hooks} forward hooks on {len(layers)} decoder layers\n")
+
+
+def _patch_kda(mod) -> None:
+    import functools
+
+    tap = _tap()
+    C = mod.Glm5NextLinearAttention
+    o_forward = C._forward
+
+    def _state_rows(self, md):
+        """Gather the conv/recurrent state rows of the tapped tokens' requests."""
+        torch = tap.torch
+        idx = md.non_spec_state_indices_tensor
+        if idx is None:
+            idx = md.spec_state_indices_tensor[:, 0] if md.spec_state_indices_tensor is not None else None
+        if idx is None:
+            return None, None, None
+        if len(set(tap.row_req)) < len(tap.tapped):
+            return None, None, None  # prefill rows: one state per request, skip
+        req_rows = tap.req_of_tok[tap.tapped].to(idx.device)
+        sidx = idx.reshape(-1)[req_rows].to(torch.int64)
+        conv_state, rec_state = self.kv_cache
+        return sidx, conv_state.index_select(0, sidx), rec_state.index_select(0, sidx)
+
+    @functools.wraps(o_forward)
+    def _forward(self, qkv_proj_states, g1, beta, core_attn_out, *a, **k):
+        L = tap.layer_of(self.prefix)
+        md = None
+        sidx = None
+        if tap.cur is not None:
+            try:
+                from vllm.forward_context import get_forward_context
+
+                raw = get_forward_context().attn_metadata
+                md = raw.get(self.prefix) if isinstance(raw, dict) else None
+                if md is not None:
+                    tap.cp(f"{L}.kda.qkv", qkv_proj_states)
+                    tap.cp(f"{L}.kda.g1", g1)
+                    tap.cp(f"{L}.kda.beta", beta)
+                    sidx, conv_rows, rec_rows = _state_rows(self, md)
+                    if sidx is not None:
+                        tap.ival(f"{L}.kda.state_idx", sidx)
+                        tap.cp_rows(f"{L}.kda.conv_before", conv_rows)
+                        tap.cp_rows(f"{L}.kda.rec_before", rec_rows)
+                    if L == "L1" or L == "L0":
+                        tap.ival("kda.num_prefills", [int(md.num_prefills)])
+                        tap.ival("kda.num_decodes", [int(md.num_decodes)])
+                        if md.has_initial_state is not None:
+                            tap.ival("kda.has_initial_state", md.has_initial_state.to(tap.torch.int64))
+            except Exception:  # noqa: BLE001
+                tap.note_err(f"kda pre {L}")
+        out = o_forward(self, qkv_proj_states, g1, beta, core_attn_out, *a, **k)
+        if tap.cur is not None and md is not None:
+            try:
+                tap.cp(f"{L}.kda.out", core_attn_out)
+                if sidx is not None:
+                    _, conv_rows, rec_rows = _state_rows(self, md)
+                    tap.cp_rows(f"{L}.kda.conv_after", conv_rows)
+                    tap.cp_rows(f"{L}.kda.rec_after", rec_rows)
+                elif tap.vectors or tap.vec_steps is not None:
+                    # prefill rows (several per request): keep ONE state per request
+                    # after the step, in the blob's extra, keyed by request id
+                    idx = md.non_spec_state_indices_tensor
+                    if idx is not None:
+                        conv_state, rec_state = self.kv_cache
+                        per_req = {}
+                        for t, r in zip(tap.tapped, tap.row_req):
+                            if r not in per_req:
+                                per_req[r] = int(idx.reshape(-1)[int(tap.req_of_tok[t])])
+                        tap.extra[f"{L}.kda.state_after_req"] = {
+                            r: (conv_state[si].clone(), rec_state[si].clone()) for r, si in per_req.items()
+                        }
+            except Exception:  # noqa: BLE001
+                tap.note_err(f"kda post {L}")
+        return out
+
+    C._forward = _forward
+    sys.stderr.write("[glm53-tap] KDA patched\n")
+
+
+def _patch_mla_layer(mod) -> None:
+    import functools
+
+    tap = _tap()
+    C = mod.MLAAttention
+    o_forward = C.forward
+
+    def _cache(self):
+        kv = self.kv_cache
+        if isinstance(kv, (list, tuple)):
+            kv = kv[0]
+        return kv
+
+    @functools.wraps(o_forward)
+    def forward(self, q, kv_c_normed, k_pe, *a, **k):
+        L = tap.layer_of(self.layer_name)
+        slot_rows = None
+        if tap.cur is not None:
+            try:
+                from vllm.forward_context import get_forward_context
+
+                fc = get_forward_context()
+                tap.cp(f"{L}.mla.q", q)
+                tap.cp(f"{L}.mla.kvc", kv_c_normed)
+                tap.cp(f"{L}.mla.kpe", k_pe)
+                sm = fc.slot_mapping.get(self.layer_name) if isinstance(fc.slot_mapping, dict) else None
+                kv = _cache(self)
+                if sm is not None and tap.torch.is_tensor(kv) and kv.numel() > 0:
+                    slot_rows = sm.reshape(-1)[: tap.n_pad].index_select(0, tap.tapped_t).to(tap.torch.int64)
+                    tap.ival(f"{L}.mla.slot", slot_rows)
+                    flat = kv.reshape(-1, kv.shape[-1])
+                    # re-read the rows this request wrote last step (continuity)
+                    prev = [tap.kv_prev.get((L, r), (0, 0))[0] for r in tap.row_req]
+                    if any(p for p in prev):
+                        pv = tap.torch.tensor(prev, dtype=tap.torch.int64, device=flat.device)
+                        tap.cp_rows(f"{L}.mla.cache_reread", flat.index_select(0, pv.clamp_(min=0)))
+            except Exception:  # noqa: BLE001
+                tap.note_err(f"mla pre {L}")
+        out = o_forward(self, q, kv_c_normed, k_pe, *a, **k)
+        if tap.cur is not None:
+            try:
+                tap.cp(f"{L}.mla.attn_out", out)
+                if slot_rows is not None:
+                    kv = _cache(self)
+                    flat = kv.reshape(-1, kv.shape[-1])
+                    tap.cp_rows(f"{L}.mla.cache_row", flat.index_select(0, slot_rows.clamp(min=0)))
+            except Exception:  # noqa: BLE001
+                tap.note_err(f"mla post {L}")
+        return out
+
+    C.forward = forward
+    sys.stderr.write("[glm53-tap] MLAAttention patched\n")
+
+
+def _patch_sm90(mod) -> None:
+    import functools
+
+    tap = _tap()
+    C = mod.FlashInferMLASparseSM90Impl
+    o = C.forward_mqa
+
+    @functools.wraps(o)
+    def forward_mqa(self, q, kv_c_and_k_pe_cache, attn_metadata, layer):
+        L = tap.layer_of(getattr(layer, "layer_name", ""))
+        n = None
+        if tap.cur is not None:
+            try:
+                torch = tap.torch
+                qn = q[0] if isinstance(q, tuple) else q
+                n = qn.shape[0]
+                st = attn_metadata.state
+                planned = st._lens_cpu[:n]
+                tap.ival_tok(f"{L}.mla.planned", planned.to(qn.device))
+                topk = self.topk_indices_buffer[:n]
+                tap.ival_tok(f"{L}.mla.valid", (topk >= 0).sum(dim=1))
+                tap.cp(f"{L}.mla.topk", topk)
+                # the kernel's real query inputs (absorbed latent q) for an offline reference
+                if isinstance(q, tuple):
+                    tap.cp(f"{L}.mla.q_nope", q[0])
+                    tap.cp(f"{L}.mla.q_pe", q[1])
+                tap.extra[f"{L}.mla.k_scale"] = float(getattr(layer, "_k_scale_float", 1.0) or 1.0)
+            except Exception:  # noqa: BLE001
+                tap.note_err(f"sm90 pre {L}")
+        out = o(self, q, kv_c_and_k_pe_cache, attn_metadata, layer)
+        if tap.cur is not None and n is not None:
+            try:
+                st = attn_metadata.state
+                w = st.topk_width
+                kvi = st.kv_indices[: n * w].reshape(n, w)
+                tap.cp(f"{L}.mla.kvidx", kvi)
+                tap.ival_tok(f"{L}.mla.kvidx0", kvi[:, 0])
+                # the latent rows the kernel attended, per tapped row (variable length):
+                # only when vectors are being kept this step, since this is the bulk
+                if tap.vectors or tap.vec_steps is not None:
+                    flat = kv_c_and_k_pe_cache.reshape(-1, kv_c_and_k_pe_cache.shape[-1])
+                    valid = (self.topk_indices_buffer[:n] >= 0).sum(dim=1)
+                    rows = []
+                    for t in tap.tapped:
+                        nv = int(valid[t])
+                        rows.append(flat.index_select(0, kvi[t, :nv].to(torch.int64)).clone())
+                    tap.extra[f"{L}.mla.kv_rows"] = rows
+                res = out[0] if isinstance(out, tuple) else out
+                tap.cp(f"{L}.mla.out", res)
+            except Exception:  # noqa: BLE001
+                tap.note_err(f"sm90 post {L}")
+        return out
+
+    C.forward_mqa = forward_mqa
+    sys.stderr.write("[glm53-tap] SM90 forward_mqa patched\n")
+
+
+def _patch_indexer(mod) -> None:
+    import functools
+
+    tap = _tap()
+    o = mod.sparse_attn_indexer_kpool
+
+    # positional layout after `weights` (forward_cuda passes everything
+    # positionally): quant_block_size, scale_fmt, topk_tokens, head_dim,
+    # max_pool_len, total_seq_lens, topk_indices_buffer, skip_k_cache_insert,
+    # use_fp4_cache, gate_score, compress_ape, index_kpool, positions,
+    # tail_kv_cache, tail_prefix, topk_backend
+    def _arg(a, k_, i, name):
+        if name in k_:
+            return k_[name]
+        return a[i] if i < len(a) else None
+
+    @functools.wraps(o)
+    def sparse_attn_indexer_kpool(hidden_states, k_cache_prefix, kv_cache, q_quant, q_scale, k, weights, *a, **k_):
+        L = None
+        if tap.cur is not None:
+            try:
+                from vllm.forward_context import get_forward_context
+                from vllm.utils.torch_utils import _resolve_layer_name
+
+                name = _resolve_layer_name(k_cache_prefix)
+                L = tap.layer_of(name)
+                tap.cp(f"{L}.idx.hidden", hidden_states)
+                tap.cp(f"{L}.idx.q", q_quant)
+                tap.cp(f"{L}.idx.k", k)
+                tap.cp(f"{L}.idx.w", weights)
+                gs = _arg(a, k_, 9, "gate_score")
+                if gs is not None:
+                    tap.cp(f"{L}.idx.gate", gs)
+                raw = get_forward_context().attn_metadata
+                if isinstance(raw, dict):
+                    md = raw.get(name)
+                    if md is not None and md.slot_mapping is not None:
+                        tap.ival_tok(f"{L}.idx.slot", md.slot_mapping)
+                    tp = _arg(a, k_, 14, "tail_prefix")
+                    if tp is not None:
+                        tm = raw.get(_resolve_layer_name(tp))
+                        if tm is not None and tm.slot_mapping is not None:
+                            tap.ival_tok(f"{L}.idx.tail_slot", tm.slot_mapping)
+            except Exception:  # noqa: BLE001
+                tap.note_err("indexer pre")
+        out = o(hidden_states, k_cache_prefix, kv_cache, q_quant, q_scale, k, weights, *a, **k_)
+        if tap.cur is not None and L is not None:
+            try:
+                tap.cp(f"{L}.idx.topk", out)
+            except Exception:  # noqa: BLE001
+                tap.note_err("indexer post")
+        return out
+
+    mod.sparse_attn_indexer_kpool = sparse_attn_indexer_kpool
+    sys.stderr.write("[glm53-tap] indexer patched\n")
+
+
+def _patch_sampler(mod) -> None:
+    import functools
+
+    tap = _tap()
+    C = mod.Sampler
+    o = C.__call__
+
+    @functools.wraps(o)
+    def __call__(self, logits, input_batch):
+        pre = None
+        if tap.cur is not None:
+            try:
+                if logits.shape[0] > 0 and not tap.capturing():
+                    top = logits.float().topk(2, dim=-1)
+                    pre = (top.values.to("cpu"), top.indices.to("cpu"), list(input_batch.req_ids))
+            except Exception:  # noqa: BLE001
+                tap.note_err("sampler pre")
+        out = o(self, logits, input_batch)
+        if tap.cur is not None and pre is not None:
+            try:
+                vals, ids, reqs = pre
+                st = out.sampled_token_ids
+                st = st[:, 0].to("cpu").tolist() if st.dim() == 2 else st.reshape(-1).to("cpu").tolist()
+                for i, r in enumerate(reqs[: vals.shape[0]]):
+                    if tap.req_filter and tap.req_filter not in r:
+                        continue
+                    tap.sample_rows.append(
+                        {
+                            "req": r,
+                            "argmax": int(ids[i, 0]),
+                            "top2": int(ids[i, 1]),
+                            "gap": float(vals[i, 0] - vals[i, 1]),
+                            "sampled": int(st[i]) if i < len(st) else None,
+                        }
+                    )
+            except Exception:  # noqa: BLE001
+                tap.note_err("sampler post")
+        return out
+
+    C.__call__ = __call__
+    sys.stderr.write("[glm53-tap] sampler patched\n")
+
+
+_TARGETS = {
+    "vllm.v1.worker.gpu.model_runner": _patch_runner,
+    "vllm.models.glm5next.common.kda": _patch_kda,
+    "vllm.model_executor.layers.attention.mla_attention": _patch_mla_layer,
+    "vllm.v1.attention.backends.mla.flashinfer_mla_sparse_sm90": _patch_sm90,
+    "vllm.models.glm5next.nvidia.sparse_indexer": _patch_indexer,
+    "vllm.v1.worker.gpu.sample.sampler": _patch_sampler,
+}
+
+
+def _arm() -> None:
+    # If a target is already imported (e.g. armed late), patch it now.
+    for name, cb in _TARGETS.items():
+        m = sys.modules.get(name)
+        if m is not None:
+            try:
+                cb(m)
+            except Exception:  # noqa: BLE001
+                _log_exc(f"late patch {name}")
+    _install_post_import_hooks({n: cb for n, cb in _TARGETS.items() if n not in sys.modules})
+
+
+if _ENABLED and os.environ.get("GLM53_TAP_SELFTEST", "0") in ("", "0"):
+    _arm()
+
+
+# --------------------------------------------------------------------------
+# Self-test: no vLLM. Exercises row views, pairing, hashing, diffs, hooks.
+# --------------------------------------------------------------------------
+
+
+def _selftest() -> int:
+    import json
+    import tempfile
+    import types
+
+    import torch
+
+    os.environ["GLM53_TAP"] = "1"
+    os.environ.setdefault("GLM53_TAP_DIR", tempfile.mkdtemp(prefix="glm53-tap-"))
+    os.environ["GLM53_TAP_DEEP"] = "force"
+    tap = _tap()
+    dev = "cuda" if torch.cuda.is_available() else "cpu"
+
+    # fake runner / input batch with 3 requests: A and B identical prefixes at pos 5, C at pos 9
+    class RS:
+        pass
+
+    n = 3
+    rs = RS()
+    rs.num_computed_tokens = types.SimpleNamespace(gpu=torch.tensor([5, 5, 9, 0], device=dev))
+    ids = torch.zeros(4, 32, dtype=torch.int64, device=dev)
+    ids[0, :6] = torch.arange(6, device=dev)
+    ids[1, :6] = torch.arange(6, device=dev)
+    ids[2, :10] = torch.arange(10, device=dev)
+    rs.all_token_ids = types.SimpleNamespace(gpu=ids)
+    rs.last_sampled_tokens = torch.tensor([7, 7, 3, 0], device=dev)
+    runner = types.SimpleNamespace(req_states=rs, execute_model_state=None, vllm_config=None)
+    ib = types.SimpleNamespace(
+        req_ids=["reqA", "reqB", "reqC"],
+        num_tokens=n,
+        num_tokens_after_padding=4,
+        num_reqs=n,
+        query_start_loc_np=[0, 1, 2, 3],
+        positions=torch.tensor([5, 5, 9, 0], device=dev),
+        input_ids=torch.tensor([5, 5, 9, 0], device=dev),
+        seq_lens=torch.tensor([6, 6, 10, 0], device=dev),
+        seq_lens_cpu_upper_bound=torch.tensor([6, 6, 10, 0]),
+        num_computed_tokens_np=[5, 5, 9],
+        idx_mapping_np=[0, 1, 2],
+        idx_mapping=torch.tensor([0, 1, 2], device=dev),
+    )
+    tap.begin_step(runner, ib)
+    assert tap.cur is not None, "begin_step produced no record"
+    assert tap.pairs == [(0, 1)], f"pairs {tap.pairs}"
+    # hooks on a toy stack: layer0 row-local, layer1 leaks row 1's value into row 0
+    class Sub(torch.nn.Module):
+        def __init__(self, leak):
+            super().__init__()
+            self.leak = leak
+
+        def forward(self, x):
+            y = x * 2
+            if self.leak:
+                y = y.clone()
+                y[0] += 1.0
+            return y
+
+    class Layer(torch.nn.Module):
+        def __init__(self, leak):
+            super().__init__()
+            self.self_attn = Sub(False)
+            self.mlp = Sub(leak)
+
+        def forward(self, x):
+            return self.mlp(self.self_attn(x)), None
+
+    Layer.__name__ = "Glm5NextDecoderLayer"
+    Sub.__name__ = "Glm5NextLinearAttention"
+    model = torch.nn.Module()
+    model.layers = torch.nn.ModuleList([Layer(False), Layer(True)]).to(dev)
+    runner.model = model
+    runner.vllm_config = types.SimpleNamespace(compilation_config=types.SimpleNamespace(mode="CompilationMode.NONE"))
+    _install_deep_hooks(runner, tap)
+    x = torch.randn(4, 16, device=dev)
+    x[1] = x[0]
+    for layer in model.layers:
+        x, _ = layer(x)
+    # a fake KDA state continuity + plan mismatch + cache reread
+    tap.ival("L3.kda.state_idx", torch.tensor([11, 12, 13], device=dev))
+    st = torch.randn(3, 8, device=dev)
+    tap.cp_rows("L3.kda.rec_before", st)
+    tap.cp_rows("L3.kda.rec_after", st + 1)
+    tap.ival_tok("L2.mla.planned", torch.tensor([6, 6, 10, 0], device=dev))
+    tap.ival_tok("L2.mla.valid", torch.tensor([6, 5, 10, 0], device=dev))
+    tap.flush()
+    # step 2: state_before differs from step 1's after for state 12 -> state_changed
+    tap.begin_step(runner, ib)
+    tap.ival("L3.kda.state_idx", torch.tensor([11, 12, 13], device=dev))
+    st2 = st + 1
+    st2[1] += 5
+    tap.cp_rows("L3.kda.rec_before", st2)
+    tap.flush()
+    path = os.path.join(tap.dir, f"tap-rank{tap._rank()}.jsonl")
+    with open(path) as f:
+        recs = [json.loads(l) for l in f]
+    s1, s2 = recs[1], recs[2]
+    p = s1["pairs"][0]
+    assert p["first_diff"] == "L1.mlp.out0", f"first_diff {p['first_diff']} diffs {p['diffs']}"
+    assert "L1.mlp.in0" not in p["diffs"], p["diffs"]
+    assert s1["plan_mismatch"][0]["req"] == "reqB", s1.get("plan_mismatch")
+    assert s2["state_changed"][0]["req"] == "reqB", s2.get("state_changed")
+    print("selftest OK:", path)
+    print(" first_diff:", p["first_diff"], "| ndiff", p["ndiff"], "| ncp", p["ncp"])
+    print(" plan_mismatch:", s1["plan_mismatch"])
+    print(" state_changed:", s2["state_changed"])
+    return 0
+
+
+if __name__ == "__main__":
+    if "--selftest" in sys.argv or os.environ.get("GLM53_TAP_SELFTEST", "0") not in ("", "0"):
+        sys.exit(_selftest())
+    print(__doc__)

docker-images/glm53-kindling/dev/tap/glm53_tap_report.py

@@ -0,0 +1,243 @@
+#!/usr/bin/env python3
+"""Offline report over glm53_step_tap output.
+
+    python3 glm53_tap_report.py DIR [--req SUBSTR] [--steps A-B] [--all-pairs]
+    python3 glm53_tap_report.py --compare DECODE.pt PREFILL.pt --pos N [--rel 0.05]
+
+--compare diffs the row at absolute position N of two vector dumps (one from
+a decode step, one from a prefill replay of the same ids) checkpoint by
+checkpoint, in float, in the decode file's execution order, and flags the
+first checkpoint whose relative difference exceeds --rel. KDA state rows are
+absent from prefill dumps (one state per request), so the first divergence
+there is read from kda.qkv/g1/beta (layer input) vs kda.out (kernel output).
+
+DIR holds tap-rank<r>.jsonl for one or more ranks (copy each box's file into
+the same directory). Steps are aligned by their step number: every rank runs
+the same scheduler output sequence, so step N is the same batch on every rank.
+
+Sections:
+  pairs        per step, per pair: first checkpoint where the two rows differ
+               (execution order), and the step a pair first split
+  events       host/device num_computed mismatch, SM90 planned!=valid,
+               KDA state changed between steps, MLA cache row changed
+               between steps, sampler argmax!=sampled
+  ranks        cross-rank disagreement on fed tokens, input ids, positions,
+               slots, and every replicated checkpoint hash
+  tokens       per request: the token fed at each step on rank 0 (find the
+               step where the wrong token appears, then look up the sections
+               above at that step and the one before it)
+"""
+
+from __future__ import annotations
+
+import argparse
+import glob
+import json
+import os
+import sys
+from collections import defaultdict
+
+SHARDED = ("kda.", "mla.q", "mla.out", "mla.kvidx", "mla.attn_out", "attn_mla.out", "attn_kda.out")
+
+
+def load(d: str):
+    ranks = {}
+    for p in sorted(glob.glob(os.path.join(d, "tap-rank*.jsonl"))):
+        r = int(os.path.basename(p)[len("tap-rank") : -len(".jsonl")])
+        recs = []
+        hdr = None
+        with open(p) as f:
+            for line in f:
+                line = line.strip()
+                if not line:
+                    continue
+                try:
+                    o = json.loads(line)
+                except json.JSONDecodeError:
+                    continue
+                if o.get("header"):
+                    hdr = o
+                else:
+                    recs.append(o)
+        ranks[r] = (hdr, {o["step"]: o for o in recs})
+    return ranks
+
+
+def replicated(name: str) -> bool:
+    return not any(s in name for s in SHARDED)
+
+
+def compare(a_path: str, b_path: str, pos: int, rel_thr: float) -> int:
+    import torch
+
+    A = torch.load(a_path, map_location="cpu")
+    B = torch.load(b_path, map_location="cpu")
+
+    def row_of(blob, name):
+        cands = [i for i, r in enumerate(blob["rows"]) if r["pos"] == pos]
+        if not cands:
+            print(f"{name}: no row at pos {pos}; positions {sorted({r['pos'] for r in blob['rows']})[:12]}...")
+            return None
+        return cands[-1]
+
+    ia, ib = row_of(A, a_path), row_of(B, b_path)
+    if ia is None or ib is None:
+        return 1
+    print(f"row A {A['rows'][ia]}\nrow B {B['rows'][ib]}")
+    first = None
+    for name, ta in A["cps"].items():
+        tb = B["cps"].get(name)
+        if tb is None:
+            continue
+        va, vb = ta[ia].float().reshape(-1), tb[ib].float().reshape(-1)
+        if va.numel() != vb.numel():
+            print(f"{name:32s} shape mismatch {tuple(ta.shape)} vs {tuple(tb.shape)}")
+            continue
+        if ta.dtype in (torch.int32, torch.int64, torch.uint8, torch.int8):
+            neq = int((va != vb).sum())
+            mark = "  <-- differs" if neq else ""
+            print(f"{name:32s} int  {neq} of {va.numel()} entries differ{mark}")
+            if neq and first is None:
+                first = name
+            continue
+        d = (va - vb).abs()
+        rel = float(d.norm() / (vb.norm() + 1e-12))
+        mark = ""
+        if rel > rel_thr and first is None:
+            first = name
+            mark = "  <-- FIRST above threshold"
+        print(f"{name:32s} maxabs {float(d.max()):.4g} rel {rel:.4g} |B|max {float(vb.abs().max()):.4g}{mark}")
+    print("\nfirst checkpoint above threshold:", first)
+    return 0
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("dir", nargs="?")
+    ap.add_argument("--compare", nargs=2, metavar=("A.pt", "B.pt"))
+    ap.add_argument("--pos", type=int)
+    ap.add_argument("--rel", type=float, default=0.05)
+    ap.add_argument("--req", default="")
+    ap.add_argument("--steps", default="")
+    ap.add_argument("--all-pairs", action="store_true", help="print every pair every step, not only splits")
+    args = ap.parse_args()
+    if args.compare:
+        if args.pos is None:
+            print("--compare needs --pos N (absolute position of the row)")
+            return 2
+        return compare(args.compare[0], args.compare[1], args.pos, args.rel)
+    if not args.dir:
+        ap.print_help()
+        return 2
+    ranks = load(args.dir)
+    if not ranks:
+        print("no tap-rank*.jsonl in", args.dir)
+        return 1
+    lo, hi = 0, 1 << 40
+    if args.steps:
+        a, _, b = args.steps.partition("-")
+        lo, hi = int(a or 0), int(b or hi)
+
+    def want(o):
+        return lo <= o["step"] <= hi and (not args.req or any(args.req in x["req"] for x in o["rows"]))
+
+    for r, (hdr, _) in sorted(ranks.items()):
+        print(f"rank {r}: {json.dumps(hdr) if hdr else 'no header'}")
+    r0 = min(ranks)
+    steps0 = ranks[r0][1]
+
+    print("\n== pairs ==")
+    seen_split = set()
+    for s in sorted(steps0):
+        o = steps0[s]
+        if not want(o):
+            continue
+        for p in o.get("pairs", []):
+            key = (p["a"], p["b"])
+            split = p.get("first_diff")
+            if split and key not in seen_split:
+                seen_split.add(key)
+                print(f"step {s} pos {p['pos']} pair {p['a']} / {p['b']}: FIRST SPLIT at {split} "
+                      f"({p.get('ndiff')} of {p.get('ncp')} checkpoints differ; {p.get('first_diff_nbytes_ne')} bytes differ there)")
+                print("   differing:", ", ".join(p.get("diffs", [])[:40]))
+            elif args.all_pairs or (split and args.req):
+                print(f"step {s} pos {p['pos']} pair {p['a']} / {p['b']}: {split or 'identical'} ({p.get('ndiff')}/{p.get('ncp')})")
+    if not seen_split:
+        print("(no pair split in the selected steps)")
+
+    print("\n== events ==")
+    n_ev = 0
+    for r, (_, steps) in sorted(ranks.items()):
+        for s in sorted(steps):
+            o = steps[s]
+            if not want(o):
+                continue
+            for k in ("host_dev_mismatch", "plan_mismatch", "state_changed", "kv_changed", "err"):
+                v = o.get(k)
+                if v:
+                    n_ev += 1
+                    print(f"rank {r} step {s} {k}: {json.dumps(v)[:600]}")
+            for sm in o.get("sample", []):
+                if sm.get("sampled") is not None and sm["sampled"] != sm["argmax"]:
+                    n_ev += 1
+                    print(f"rank {r} step {s} sampled!=argmax {sm}")
+    if not n_ev:
+        print("(none)")
+
+    print("\n== ranks ==")
+    if len(ranks) < 2:
+        print("(single rank file; nothing to compare)")
+    else:
+        n_bad = 0
+        for s in sorted(steps0):
+            o0 = steps0[s]
+            if not want(o0):
+                continue
+            others = [(r, ranks[r][1].get(s)) for r in sorted(ranks) if r != r0]
+            for r, o in others:
+                if o is None:
+                    print(f"step {s}: rank {r} has no record")
+                    n_bad += 1
+                    continue
+                bad = []
+                for f in ("fed", "nc_after", "sampled0"):
+                    if o0.get(f) != o.get(f):
+                        bad.append(f"{f} r{r0}={o0.get(f)} r{r}={o.get(f)}")
+                rows0 = [(x["req"], x["pos"], x["tok"], x["seq_dev"], x["nc_dev"]) for x in o0["rows"]]
+                rows1 = [(x["req"], x["pos"], x["tok"], x["seq_dev"], x["nc_dev"]) for x in o["rows"]]
+                if rows0 != rows1:
+                    bad.append(f"rows r{r0}={rows0} r{r}={rows1}")
+                for n, h in o0.get("cps", {}).items():
+                    if not replicated(n):
+                        continue
+                    h1 = o.get("cps", {}).get(n)
+                    if h1 is not None and h1 != h:
+                        bad.append(f"cp {n}")
+                for n, v in o0.get("ints", {}).items():
+                    if n.endswith((".mla.slot", ".idx.slot", ".idx.tail_slot", ".kda.state_idx", ".mla.planned", ".mla.valid")):
+                        v1 = o.get("ints", {}).get(n)
+                        if v1 is not None and v1 != v:
+                            bad.append(f"int {n} r{r0}={v} r{r}={v1}")
+                if bad:
+                    n_bad += 1
+                    print(f"step {s} rank {r} vs {r0}: " + "; ".join(bad[:12]) + (" ..." if len(bad) > 12 else ""))
+        if not n_bad:
+            print("(all ranks agree on fed tokens, rows, placement and replicated checkpoints)")
+
+    print("\n== tokens (rank %d) ==" % r0)
+    per_req = defaultdict(list)
+    for s in sorted(steps0):
+        o = steps0[s]
+        if not want(o):
+            continue
+        for req, tok in (o.get("fed") or {}).items():
+            if not args.req or args.req in req:
+                pos = next((x["pos"] for x in o["rows"] if x["req"] == req), None)
+                per_req[req].append((s, pos, tok))
+    for req, seq in per_req.items():
+        print(f"{req}: " + " ".join(f"{s}:{p}>{t}" for s, p, t in seq))
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

docker-images/glm53-kindling/experimental/README.md

@@ -0,0 +1,297 @@
+# Experimental speedups
+
+Compose overrides that stack on `compose/glm53.yaml`. Each one mounts files
+over the image this repo builds (vLLM nightly ddd6fbca); several replace whole vLLM or
+FlashInfer files, so they only match that image.
+
+## How to use it
+
+Requirements: the image built from this tree (`image/build.sh`) and the normal four-node setup from the main README (the same `compose/glm53.yaml` and a per-node `compose/.env` on every box). The overrides replace files inside that image, so they do not apply to other images.
+
+1. Put both ConnectX PCIe roots on the fabric: tag both interfaces `rdma` in mentatd's `MENTAT_ANNOUNCE_IFACES` (main README, step 4), or set `FABRIC_SUBNETS` to both subnets in each node's `compose/.env`. The entrypoint derives `NCCL_IB_HCA` and `NCCL_IB_GID_INDEX` from it, and arx and arxbig use the same two devices. Each rank's log shows them (`arx all-reduce: rank r/4 on [dev0, dev1]`). With only one device they log a warning and fall back to NCCL.
+2. Start the stack with every override, on every node: `./glm53 up -d` from the repo root. The script holds the list of compose files, in order (later files win).
+
+3. The first boot loads the checkpoint normally (~8 min), compiles the CUDA extensions, and writes a processed-weight snapshot per rank under `weight-snapshots` in the cache mount (`CACHE_HOME`, or the `glm53_cache` volume), ~48 GB per node. Later boots restore it (~3.5 min).
+4. Check: `bash smoketest/run.sh http://<head>:8002` should pass 8/8, and each rank's log should show `arx all-reduce: rank r/4` and `arxbig all-gather: rank r/4`.
+
+### Turning pieces off
+
+Leave a compose file out to drop that piece, or set its switch in `.env` (every node must use the same values):
+
+| Switch | Default | What it does |
+|---|---|---|
+| `VLLM_ARXBIG` | 1 | arxbig RDMA collectives for prefill (arx.yaml) |
+| `VLLM_ARXBIG_AG` | 0 | route prefill all-gathers through arxbig; off because its output sits in pinned memory, where GEMMs reading it run 3.7x slower |
+| `VLLM_ARXBIG_RS` | 1 | RDMA reduce-scatter buffers (~0.5 GB pinned per rank); needed by `VLLM_GLM_SP_MOE_FUSED` |
+| `VLLM_ARX_RING` | 0 (1 with `TP=RING4`) | for boxes cabled in a ring with no switch: arx and arxbig open QPs only to rank r-1 and r+1, and at four ranks the rank between two others relays their data, with the same results as over the switch. arx takes rings of 2, 3 or 4; arxbig takes 2 or 4, so a ring of 3 needs `VLLM_ARXBIG=0` |
+| `ARX_RING_PREV_HCAS` | from mentat with `TP=RING4` | with `VLLM_ARX_RING=1`: the two RDMA devices of the port cabled to rank r-1. Each device uses the GID of its own IPv4 address. Device r must share its subnet with device r of the neighbour's list for that cable: root 0 then root 1 when every box is cabled the same way, root 1 first on a box whose cabling is mirrored |
+| `ARX_RING_NEXT_HCAS` | unset | the same for the port cabled to rank r+1 |
+| `VLLM_GLM_SP_TP` | 1 | sequence parallelism for forwards of `VLLM_GLM_SP_MIN_TOKENS` (1024) or more (sp.yaml) |
+| `VLLM_GLM_SP_FP8_GATHER` | 1 | gather KDA attention inputs as FP8 |
+| `VLLM_GLM_SP_MOE_FUSED` | 1 | MoE combine feeding the RDMA reduce-scatter; needs arx.yaml with `VLLM_ARXBIG_RS=1` and megamoe.yaml with `VLLM_MOE_PREFILL=1` |
+| `VLLM_GLM_SP_MOE_QUANT_GATHER` | 1 | with `VLLM_GLM_SP_MOE_FUSED`: route each rank's own rows and gather the MoE input as NVFP4 + FP8 instead of bf16 (same results) |
+| `VLLM_GLM_ARX_PREFETCH` | 1 | L2 prefetch during decode all-reduces; needs arx.yaml |
+| `VLLM_MEGAMOE` | 1 | decode MoE kernel for batches of up to `VLLM_MEGAMOE_MAX_TOKENS` (8) |
+| `VLLM_MOE_PREFILL` | 1 | prefill MoE kernel for batches of `VLLM_MOE_PREFILL_MIN_TOKENS` (1024) or more |
+| `VLLM_MOE_PREFILL_Y8` | 1 | FP8 per-expert rows in the fused prefill MoE |
+| `VLLM_TRITON_SPARSE_MLA` | 1 | Triton sparse MLA instead of FlashInfer's (fixes.yaml) |
+| `MAX_NUM_SEQS` | 50 | requests decoding at once (fixes.yaml); about 50 KDA states fit in the pinned KV pool, against the stock 32 |
+| `VLLM_GLM5NEXT_RECOVERSSM` | 1 | with recoverssm.yaml: KDA drafts verified from one saved state per request, bit-exact; raises `MAX_NUM_SEQS` to 64 at TP=4 and 16 at TP=2 |
+| `VLLM_GLM5NEXT_DRAFT_POOL` | 1 | the drafter's KV in its own small pool instead of a page in every KV block: ~40% more KV tokens (fixes.yaml) |
+| `VLLM_DENSE_W4` | in_proj, o_proj, shared experts, drafter | regex of dense layers stored as NVFP4 (the rest are FP8); add `\|lm_head$` for the opt-in NVFP4 lm_head |
+| `VLLM_DENSE_FP8_LM_HEAD` | 1 | FP8 lm_head |
+| `VLLM_ADAPTIVE_K_ONLINE` | 1 | adaptive-k learns the step cost while serving; 0 uses the fixed table below |
+| `VLLM_ADAPTIVE_K_MODEL` | fitted on this stack | starting step cost: ms, ms per expert touched, ms per verified token |
+| `VLLM_ADAPTIVE_K_COST_MS` | measured on this stack | with `VLLM_ADAPTIVE_K_ONLINE=0`: step cost per verify length |
+
+Debug switches: `VLLM_MOE_PREFILL_CHECK=N` and `VLLM_GLM_SP_MOE_CHECK=N` also run the stock path on the first N prefill batches and log the difference.
+
+### Snapshot version and tag
+
+A snapshot holds processed weights. Its key covers the checkpoint, the TP layout, the MoE backend and `max_num_batched_tokens`. Two kinds of change need more:
+
+- A code change to how weights are processed (`weight_snapshot.py`, `dense_fp8.py`, the loader patches): bump `SNAPSHOT_VERSION` in `weight_snapshot.py`. Snapshot names start with `v<version>-`, and each boot deletes the snapshots of every other version from its snapshot directory, for all TP sizes.
+- A knob that changes processing (`VLLM_DENSE_W4`, `VLLM_DENSE_FP8*`): a new `VLLM_WEIGHT_SNAPSHOT_TAG` (set in fp8.yaml). Snapshots under the old tag stay on disk until the next version bump.
+
+A shape change fails the restore with an error. A same-shape change (such as which layers are NVFP4) would restore stale weights without one.
+
+### Quality check
+
+`experimental/quality/quality.py` (GSM8K and HumanEval against a running server; see its docstring for the data files and the no-network HumanEval run).
+`experimental/quality/prefill_block.py` measures what a long prefill does to other requests: gaps in a stream that is already generating, and time to first token for short requests arriving meanwhile.
+`experimental/quality/agent_tools.py` runs five small agent tasks (list, read, search and write files in an in-memory tree) streamed and not streamed, and checks the result; `python3 agent_tools.py http://<head>:8002`.
+
+## Results
+
+Single stream, thinking off, 512 tokens (`dev/repro/decode.py` prompts):
+
+| | structured | code | prose |
+|---|---|---|---|
+| stock, DFlash2 k=7 | 121.9 | 91.3 | 38.7 |
+| all overrides | 170.3 | 120.8 | 65.8 |
+
+Concurrent streams, aggregate tok/s at 1/2/4/8 streams: stock 85.5/65.4/99.5/147.8,
+all overrides 126/102/146/198 (16 streams: 271). With a different prompt per
+stream, code runs 317 tok/s at 16 streams and 451 at 32, and 50 mixed streams
+run 404. RigMark (reasoning=low) code / prose / structured 107.9 / 61.7 / 157.1
+tok/s. Needle retrieval 12/12 up to 507k tokens. KV cache 3.63M tokens per
+rank at the 26 GiB pin and 50 sequences (2.63M before the drafter got its own
+pool).
+
+Cold prefill, tok/s (random-word prompts, nothing cached):
+
+| | 32k | 128k |
+|---|---|---|
+| stock | 2,730 | 2,679 |
+| all overrides | 4,946 | 4,750 |
+
+Boot goes from about 8 minutes to about 3.5 once snapshots exist.
+
+## The pieces
+
+- **arx** (`arx/`): all-reduce over RDMA for tensors up to 256 KB, without
+  NCCL. It takes 13-27 us where NCCL takes 80-90. The GPU publishes a
+  sequence number, and a CPU thread posts the RDMA writes over both ConnectX
+  roots.
+- **snapshot** (`snapshot/weight_snapshot.py`, `base_loader.py`): the first
+  boot writes each rank's processed weights to disk. Later boots build the
+  model on dummy weights through the normal path, then overwrite every tensor
+  in place (about 15 s for 46 GiB). `modelopt.py` and
+  `flashinfer_cutlass_moe.py` are left from an earlier design and may no
+  longer be needed.
+- **fp8** (`snapshot/dense_fp8.py`): the NVFP4 checkpoint leaves about 4 GB
+  per rank of dense linears in bf16, and decode reads them every step. This
+  converts them to FP8 after load (per-channel weight scales, per-token
+  activation scales, CUTLASS scaled_mm). Prefill-sized batches go through in
+  2048-row pieces when the weight is over 16 MiB: CUTLASS rereads the whole
+  weight for every row of tiles, and past L2 that comes from DRAM, so one call
+  on a 16k x 6416 x 4096 GEMM runs at 68 TFLOPS and the pieces at 169. +14-20% decode, and about +0.5% NLL
+  on prose. Layers matching `VLLM_DENSE_W4` (by default the KDA in_proj and
+  every drafter layer, attention o_proj and the shared experts) go to NVFP4 instead, through `megamoe/megadense4.cu`
+  (W4A16 at the 4-bit roofline, about twice as fast as the FP8 GEMM). They
+  keep an FP8 copy for batches of more than 32 tokens. That is +6-10% decode.
+  The in_proj costs about as much NLL again as FP8 did. The drafter layers only
+  change acceptance. NVFP4 on every dense layer cost 2-3x more NLL, so it is
+  not the default.
+- **megamoe** (`megamoe/`): an NVFP4 MoE kernel for batches of up to 8 tokens.
+  It reads the CUTLASS backend's own tensors, so CUTLASS still handles
+  prefill. Activations stay 16-bit, which is exact in the weights, where
+  CUTLASS rounds them to FP4. With `VLLM_MOE_PREFILL=1` (the default here),
+  batches of 1024+ tokens take `megamoe/moe_prefill.cu`: the same W4A4 math as
+  CUTLASS on the same tensors, with the token gather, SwiGLU and FP4
+  requantization folded into the fc1 GEMM (hand-written block-scaled
+  `mma.sync` NVFP4). About 18 ms per layer at 16k tokens against CUTLASS's
+  22-26, equally close to an fp32 reference. +7-8% prefill.
+- **adaptive-k** (`adaptive-k/adaptive_k.py`): a scheduler that picks how many
+  of the 7 DFlash2 drafts each step verifies (2, 3, 4, 5 or 7). It uses recent
+  per-position acceptance, per request and averaged over all requests, and a
+  step cost it learns from measured step times: a fixed part, the MoE experts
+  the step's tokens touch, and a per-token part. Weight reads follow the
+  experts touched, so a 128-token step costs far less than a straight line
+  through small ones, and at 16+ streams the scheduler verifies more drafts
+  than a fixed table would allow. The drafter uses Triton attention, which
+  avoids a mid-step host sync.
+  The scheduler picks the step's draft count before the drafts exist. Once
+  they do, `draft-trunc/` cuts each request's drafts where their predicted
+  survival falls below `VLLM_DRAFT_TRUNC_TAU` (0.3). The prediction is vLLM's
+  acceptance estimator, which scores each draft from the logit of its most
+  likely token (AUC 0.91-0.93 per position on logged steps). Dead drafts keep
+  the verify shape, route to expert -1 (the MoE kernels skip them) and are
+  rejected, so greedy output is unchanged, token for token and logprob for
+  logprob. A cut step grades only the drafts the estimator expected to
+  survive, and fitting on those steps drifted it low over a long mixed
+  workload. So every 4th step (`VLLM_DRAFT_TRUNC_EXPLORE`) skips the cut, and
+  only those steps train the estimator. On TP=4, paired with
+  `gate/bench/bench.py` after a long mixed warm-up and net of drift: code
+  +3.3%, prose +3.2%, mixed streams +3.5 to 6.9%, structured -1.1%.
+  `VLLM_DRAFT_TRUNC_TAU=0` turns it off. `draft-trunc/` holds whole copies of
+  five vLLM files, so it checks the image's vLLM commit and stops the engine
+  on any other.
+- **recoverssm** (`fixes/recoverssm.py`, `compose/recoverssm.yaml`;
+  `VLLM_GLM5NEXT_RECOVERSSM=0` turns it off): the KDA layers keep one recurrent state
+  per request instead of one per draft position, and after sampling replay
+  the accepted tokens through the stock kernel, so outputs are bit-identical.
+  Ported from vLLM's Kimi-K3 RecoverSSM. At TP=2, 16 requests run at once
+  instead of 3, and 8 streams decode 44% faster. At TP=4, 61 run at 64
+  streams instead of 49 (+28% aggregate).
+- **fixes** (`fixes/`): things found by profiling.
+  - The fused mHC post + pre-norm decode kernel read its projection weights
+    in fp32, 1.5 MB a call, though the checkpoint stores them in bf16. It
+    now reads a cached bf16 copy (mhc_tilelang.py, tilelang_kernels.py):
+    8.2 us a call instead of 11.9, ~1% faster decode at one stream, and
+    greedy logprobs identical to the fp32 path.
+  - FlashInfer's MLA planner cloned a 136 MB metadata buffer on every step,
+    only to roll it back on error.
+  - The GLM indexer's head gate ran a fp32 GEMM that took 69 us per layer.
+  - FlashInfer's sparse MLA runs at about 10 TFLOPS on GB10 whatever the
+    indices. `gb10_sparse_mla.py` is a Triton kernel for GLM's NoPE MLA (one
+    program per query, all 16 heads, 32 gathered keys per block) that runs
+    3.3x faster. `VLLM_TRITON_SPARSE_MLA=0` goes back to FlashInfer.
+    At decode sizes (a request's 1 + 7 draft tokens) one program per query
+    left most of the 48 SMs idle, and FlashInfer decoded 3-6% faster; below
+    96 query tokens the kernel now splits each token's rows across programs
+    and combines them (4.4x faster at 8 tokens), which recovers most of that
+    without FlashInfer's per-step host sync and keeps Triton's prefill.
+  - FlashKDA needs dense q, k and v, and the KDA short conv wrote them as one
+    [tokens, 3 x 2048] tensor, so each KDA layer copied three 56 MB tensors
+    per 16k-token step. `causal_conv1d.py` gains an `out_group` argument that
+    writes each group of features as its own dense tensor, and `kda.py` uses
+    it. The outputs are bit-identical; prefill +2-3%.
+  - The DFlash2 drafter's KV rode the target's pool: every 2304-token block
+    carried the drafter's pages for its five layers, 30% of the KV cache,
+    though a sliding-window layer only ever holds its 2048-token window.
+    `kv_cache_utils.py`, `kv_cache_coordinator.py`, `kv_cache_interface.py`,
+    `worker_utils.py` and `warmup.py` give the drafter's group its own pool
+    (sized for `max_num_seqs` windows plus one prefill chunk, 105 blocks,
+    0.58 GiB) and keep its block ids out of the target's zeroing and warmup.
+    At the same 26 GiB pin: 1,416 -> 1,976 blocks, 2.63M -> 3.67M KV tokens
+    at 32 sequences (3.63M at the 50 fixes.yaml sets).
+    Draft acceptance and speed are unchanged. `VLLM_GLM5NEXT_DRAFT_POOL=0`
+    goes back; `VLLM_GLM5NEXT_DRAFT_POOL_BLOCKS` sets the pool's size.
+- **sp** (`fixes/model.py`): sequence parallelism for prefill. A forward of
+  at least 1024 tokens keeps the residual stream split across the TP ranks, so
+  mHC and the norms run on a quarter of the tokens, with an all-gather before
+  and a reduce-scatter after attention and the MLP. +17% prefill. Decode stays
+  plain TP: applied to every batch, SP cost decode 10-15%. The KDA layers'
+  attention inputs are gathered as per-token FP8, half the bytes, since
+  in_proj's FP8 GEMM would quantize the same rows the same way (+2%).
+- **arxbig** (`arx/arxbig.cu`, `VLLM_ARXBIG`, `VLLM_ARXBIG_RS`): all-gather and
+  reduce-scatter over RoCE for prefill-sized SP collectives, from pinned
+  buffers the ConnectX writes directly. The all-gather is ~10% faster than
+  NCCL's (187 vs 165 Gb/s), but its result lands in pinned memory, and on
+  GB10 a GEMM that reads its operand from pinned memory runs 3.7x slower, so
+  prefill gathers stay on NCCL (`VLLM_ARXBIG_AG=0`). The point is the
+  reduce-scatter. With `VLLM_GLM_SP_MOE_FUSED`, MoE layers under SP run the
+  router, the shared expert and moe_prefill's fc1/fc2 in model.py, and one
+  kernel writes shared + scaled routed sum (fp32, rounded once) straight into
+  the send buffer, publishing rows as it goes, so the network runs under the
+  finalize. It replaces the runner's scale and add passes and NCCL's
+  reduce-scatter: ~12.5 -> ~8.4 ms per MoE layer at 16k tokens, prefill +2.5-3%.
+  With `VLLM_MOE_PREFILL_Y8` (default) fc2 writes its per-expert rows as e4m3
+  with a scale per 128 columns, halving the 1.07 GB each layer writes and reads
+  back: prefill +5%, GSM8K/HumanEval/NLL/tool-call results unchanged.
+  With `VLLM_GLM_SP_MOE_QUANT_GATHER` (default) each rank routes its own rows
+  and quantizes them the way the experts would (NVFP4 for the routed experts,
+  per-token FP8 for the shared expert's GEMM), and the gather moves those
+  instead of bf16: 0.56 + 1 bytes per value instead of 2, and a quarter of the
+  fp32 router GEMM per rank. The layer output matches the bf16 gather to
+  within 1e-4 relative; prefill +3-4%.
+  All ranks write to all peers at once and this fabric has no PFC, so
+  incast drops packets and go-back-N retransmits make individual calls vary
+  (4-9 ms); per-destination serialization was worse (two QPs cannot fill a
+  link).
+- **ring** (`TP=RING4`; `image/fabric_ring.py`, ring mode in `arx/`): TP=4 on
+  four boxes cabled in a loop. A ConnectX-7 can't forward RoCE for another
+  box, because RoCE addressed to its MAC goes to its own RDMA engine, so every
+  collective uses only neighbour links. mentat 0.17's ring claim places rank i
+  on member i in cable order and gives each rank its interface toward each
+  neighbour. `fabric_ring.py` runs at each rank's Python start and turns that
+  into NCCL's ring algorithm with a channel graph that receives on the port
+  toward the previous rank and sends on the port toward the next, and into
+  arx's devices. The graph carries the speeds NCCL computes for its own ring:
+  with others it picked other protocols, the reductions summed in another
+  order, and a greedy text diverged. arx and arxbig relay data for the
+  opposite box through a neighbour, raw, so every rank sums the same values
+  in rank order and results stay bit-identical. Each cable needs its own
+  subnet for each PCIe root. On a switch, with both neighbours on the one
+  port, `arx/test_ring.py` passed bit for bit, decode was ~1% slower from the
+  relay and prefill 3-4% slower (2026-09-28). Not yet measured on cables.
+  Three boxes, one cable per pair, also form a ring, and there both other
+  ranks are neighbours, so arx sends every partial direct and relays nothing.
+  arx's ring mode takes 3 ranks; arxbig's does not, because its ring
+  reduce-scatter splits the rank two away across both neighbours, so a ring
+  of 3 runs with `VLLM_ARXBIG=0`. The entrypoint has no three-box ring mode
+  (only `TP=RING4`), so the NCCL ring and `ARX_RING_*_HCAS` are set by hand.
+  `fabric_ring.py` lists each port's devices in PCI order, root 0 first,
+  which fits boxes that are all cabled the same way; on a box cabled
+  mirrored, root 0 of a port faces the neighbour's root 1, and its lists
+  must put root 1 first. On cables, `arx/test_ring.py --mode ring` at
+  `WORLD_SIZE=3`, cut down to the arx all-reduce checks, the all-reduce
+  burst and the latency bench, passed 85 checks per rank with
+  12.2 / 19.7 / 42.8 / 69.1 us at 8 / 64 / 256 / 512 KB
+  ([#52](https://github.com/kindlingai/glm-5.3-flash-gx10/issues/52)).
+- **arx prefetch** (`VLLM_GLM_ARX_PREFETCH`): while a decode all-reduce waits
+  for its peers, its threads ask L2 for the weights the next kernels read (the
+  router gate and shared expert after attention, the next in_proj after the
+  MoE). About -0.6 ms per 32 ms step.
+
+`dense_fp8.simulate_nvfp4` is a diagnostic. It rounds the dense weights
+through NVFP4 to measure the quality cost: about +1% NLL on prose.
+
+## Sources
+
+Most of `fixes/`, and a few files elsewhere, are vLLM or FlashInfer files with
+our changes, bind-mounted over the originals in the image. They keep their
+Apache-2.0 headers. The rest is ours.
+
+| Origin | Files |
+|---|---|
+| vLLM, changed | `arx/cuda_communicator.py`; `fixes/attention.py`, `causal_conv1d.py`, `flashinfer_mla_sparse_sm90.py`, `gdn_attn.py`, `kda.py`, `kv_cache_coordinator.py`, `kv_cache_interface.py`, `kv_cache_utils.py`, `model.py`, `tilelang_kernels.py`, `warmup.py`, `worker_utils.py`; `snapshot/base_loader.py`, `flashinfer_cutlass_moe.py`, `modelopt.py` |
+| FlashInfer, changed | `fixes/_fa_common.py` |
+| ours | `arx/arx.py`, `arx_vllm.cu`, `arxbig.cu`; `adaptive-k/adaptive_k.py`; `fixes/gb10_sparse_mla.py`, `recoverssm.py` (after vLLM's Kimi-K3 RecoverSSM); `megamoe/`; `snapshot/weight_snapshot.py`, `dense_fp8.py`; `quality/` |
+
+Other people's work in here:
+
+- arx and arxbig take the RDMA path MTU from the ports rather than assuming
+  4096, from [Chuck](https://github.com/chuck-ads)'s
+  [#6](https://github.com/mmastrac/glm-5.3-flash-4x-gx10/pull/6).
+- `tilelang_kernels.py` carries vLLM's copy of SGLang's mHC kernel, and
+  `causal_conv1d.py` vLLM's adaptation of Tri Dao's causal-conv1d, both
+  credited in the file.
+- The decode split in `gb10_sparse_mla.py` is the Flash-Decoding scheme (Dao
+  et al., 2023): split each query's keys across programs, then merge the
+  partial softmax results.
+
+## Quality
+
+`quality/quality.py` (GSM8K first 250 and HumanEval, greedy, thinking off;
+HumanEval programs run by `quality/run_he.py` in a no-network container),
+against stock on the same boxes:
+
+| | stock | all overrides |
+|---|---|---|
+| GSM8K (250) | 97.2% | 97.2% |
+| HumanEval pass@1 | 156/164 | 156/164 |
+| count to 200, thinking off, 5 runs | 0 corrupt | 0 corrupt |
+| tool call at 42k context, 40 greedy runs | 10 diverge | 0 diverge |
+| NLL vs bf16 dense layers | within ±0.006 | prose +0.008 to +0.025, code within ±0.007 |

docker-images/glm53-kindling/experimental/adaptive-k/adaptive_k.py

@@ -0,0 +1,374 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Choose how many drafts each step verifies from recent draft acceptance.
+
+The drafter proposes num_speculative_tokens drafts every step. Verifying fewer
+costs less (each extra verified token adds MoE experts to the step), so when
+acceptance falls off early, as it does in free prose, verifying all of them
+wastes time. This scheduler picks k, the number of drafts the next step
+verifies, from the levels that the dynamic speculative decoding schedule
+(num_speculative_tokens_per_batch_size) captured CUDA graphs for.
+
+Per request it keeps c[i], an average of the rate at which draft i is accepted
+given that draft i - 1 was. A step that verified d drafts and accepted a < d of
+them observes accepts at 0..a-1 and a reject at a. A step that accepted all d
+says nothing about positions d and beyond, so those drift towards g, the same
+rates averaged over every request, raised or lowered by how far the request's
+c[d - 1] is from g[d - 1]. A request that beats the average where it was last
+tested is expected to beat it deeper too, so k can climb for it after a run of
+poorer requests. g starts each new request's c, and its own unverified
+positions drift slowly towards its deepest verified one.
+
+Expected tokens for a step verifying k drafts is 1 + sum over i < k of
+c[0] * ... * c[i]. The next k maximises the batch's expected tokens over the
+step's cost, and changes only for a gain of at least VLLM_ADAPTIVE_K_MARGIN.
+
+The step's cost comes from a model learned while serving (VLLM_ADAPTIVE_K_ONLINE,
+default on): step_ms = base + per_expert * experts(d * t) + per_token * t for t
+verified tokens, where experts(n) is how many of the MoE's experts n independent
+tokens touch and d the batch's routing diversity (repetitive output touches
+fewer experts than its token count suggests). A step's weight reads scale with
+the experts it touches, so large steps cost far less than a straight line
+through small ones. The three coefficients start from VLLM_ADAPTIVE_K_MODEL,
+and a small Kalman filter moves them and d to fit each decode step's measured
+time, d much faster than the coefficients. VLLM_ADAPTIVE_K_ONLINE=0 uses the
+fixed table VLLM_ADAPTIVE_K_COST_MS ("tokens:ms,...") instead.
+
+With VLLM_ADAPTIVE_K_PER_REQUEST=1, a batch of several requests gets a k per
+request instead: every (request, draft) slot is scored by its survival
+probability and the best slots are admitted up to the budget that maximises
+expected tokens over step cost. Unequal k takes the batch off the full CUDA
+graph for attention (piecewise instead).
+
+If the file named by VLLM_ADAPTIVE_K_CONTROL holds "force N", every step
+verifies N drafts, which is how the cost table is measured.
+
+VLLM_ADAPTIVE_K_CODE_MODE (Chuck's idea) also keeps rates for two kinds of
+output that accept far longer draft runs than prose: code inside ``` fences,
+and tool calls between <tool_call> and </tool_call>. Each request keeps them
+beside its usual rates, updated only while it is in that mode, and each mode
+has its own g. "gated" uses a mode's rates only when every decoding request is
+in that mode, since the batch shares one k, and is the default; "1" uses them
+per request, and "0" turns code mode off. A fence is a run of three backticks,
+which the tokenizer can split across tokens (a closing fence becomes "``" then
+"`\n"), so runs are counted across tokens.
+"""
+
+import os
+import time
+from collections import Counter
+
+import numpy as np
+
+from vllm.logger import init_logger
+from vllm.v1.core.sched.async_scheduler import AsyncScheduler
+from vllm.v1.core.sched.output import SchedulerOutput
+
+logger = init_logger(__name__)
+
+_ALPHA = float(os.environ.get("VLLM_ADAPTIVE_K_ALPHA", "0.25"))
+_MARGIN = float(os.environ.get("VLLM_ADAPTIVE_K_MARGIN", "0.03"))
+_PRIOR = 0.8
+# g's rate for verified positions, and its drift for unverified ones.
+_GLOBAL_ALPHA = 0.05
+_GLOBAL_DRIFT = 0.01
+_PER_REQUEST = os.environ.get("VLLM_ADAPTIVE_K_PER_REQUEST") == "1"
+_CODE_MODE = os.environ.get("VLLM_ADAPTIVE_K_CODE_MODE", "gated")  # "gated", "1" or "0" (off)
+_PROSE, _CODE, _TOOL = 0, 1, 2
+_ONLINE = os.environ.get("VLLM_ADAPTIVE_K_ONLINE", "1") == "1"
+# Fitted on GLM-5.3-Flash at TP=4 on GB10 (k 2, 3, 5 and 7 forced, 1-32 streams, code,
+# prose and structured prompts): ms, ms per expert touched, ms per token.
+_MODEL = tuple(float(v) for v in (os.environ.get("VLLM_ADAPTIVE_K_MODEL") or "27.0,0.635,0.542").split(","))
+# Routed experts and experts per token.
+_EXPERTS, _TOPK = 288, 8
+
+
+class _CostModel:
+    """step_ms = base + per_expert * experts(d * t) + per_token * t, learned online.
+
+    An extended Kalman filter over (base, per_expert, per_token, d). Only d
+    gets process noise: the workload changes while serving and the hardware
+    does not. A run of steps with one size, where only one combination is
+    observable, then moves d and leaves the coefficients alone.
+    """
+
+    _DRIFT = np.diag([0.0, 0.0, 0.0, 2e-3])
+
+    def __init__(self, theta) -> None:
+        self.x = np.array([*theta, 1.0])
+        # Confidence in the seed: base to ~5 ms, the slopes to a few percent.
+        self.P = np.diag([25.0, 4e-4, 4e-3, 0.1])
+        self.n = 0
+        self.err = 0.0
+
+    @property
+    def d(self) -> float:
+        return float(self.x[3])
+
+    @staticmethod
+    def experts(n: float) -> float:
+        return _EXPERTS * (1.0 - (1.0 - _TOPK / _EXPERTS) ** n)
+
+    def cost(self, tokens: int) -> float:
+        base, per_expert, per_token, d = self.x
+        return max(base + per_expert * self.experts(d * tokens) + per_token * tokens, 1.0)
+
+    def observe(self, tokens: int, ms: float) -> None:
+        predicted = self.cost(tokens)
+        self.err += 0.05 * (abs(ms - predicted) / predicted - self.err)
+        self.n += 1
+        base, per_expert, per_token, d = self.x
+        touched = self.experts(d * tokens)
+        # d(experts)/dd = (288 - experts) * t * -ln(1 - 8/288)
+        slope = (_EXPERTS - touched) * tokens * -np.log1p(-_TOPK / _EXPERTS)
+        H = np.array([1.0, touched, float(tokens), per_expert * slope])
+        P = self.P + self._DRIFT
+        Ph = P @ H
+        gain = Ph / (H @ Ph + (0.05 * predicted) ** 2)
+        self.x = self.x + gain * (ms - predicted)
+        self.x[3] = np.clip(self.x[3], 0.05, 2.0)
+        self.P = P - np.outer(gain, Ph)
+
+
+def _parse_costs(spec: str) -> tuple[np.ndarray, np.ndarray]:
+    points = sorted((int(t), float(ms)) for t, ms in (p.split(":") for p in spec.split(",")))
+    return np.array([p[0] for p in points], float), np.array([p[1] for p in points], float)
+
+
+class AdaptiveKScheduler(AsyncScheduler):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        lookup = self.dynamic_sd_lookup or [self.num_spec_tokens]
+        self._levels = sorted({k for k in lookup[1:] if 0 < k <= self.num_spec_tokens}) or [self.num_spec_tokens]
+        self._cost_x, self._cost_y = _parse_costs(
+            os.environ.get("VLLM_ADAPTIVE_K_COST_MS", "2:47,4:53,8:65"))
+        self._control = os.environ.get("VLLM_ADAPTIVE_K_CONTROL", "")
+        self._control_mtime = 0.0
+        self._force: int | None = None
+        self._rates: dict[str, np.ndarray] = {}
+        self._global = np.full(self.num_spec_tokens, _PRIOR)
+        # Code mode: per (request, mode) rates, per-mode g, and each request's mode.
+        self._mode_rates: dict[tuple[str, int], np.ndarray] = {}
+        self._mode_global = {m: np.full(self.num_spec_tokens, _PRIOR) for m in (_CODE, _TOOL)}
+        self._mode: dict[str, int] = {}
+        self._in_fence: dict[str, bool] = {}
+        self._in_tool: dict[str, bool] = {}
+        self._scanned: dict[str, int] = {}
+        self._run: dict[str, int] = {}  # backticks at the end of each request's output
+        self._switches = 0
+        self._backticks: dict[int, str] = {}
+        self._tool_open = self._tool_close = -1
+        if _CODE_MODE != "0":
+            self._load_mode_tokens()
+        self._k = self.num_spec_tokens
+        self._last_log = 0.0
+        self._steps: Counter[int] = Counter()
+        self._model = _CostModel(_MODEL) if _ONLINE else None
+        self._last_output = 0.0
+        self._last_model_log = 0.0
+        if self._model:
+            logger.info("adaptive k: levels %s, online cost model from %s", self._levels, _MODEL)
+        else:
+            logger.info("adaptive k: levels %s, cost %s", self._levels,
+                        dict(zip(self._cost_x.astype(int).tolist(), self._cost_y.tolist())))
+
+    def _load_mode_tokens(self) -> None:
+        """The text of every token holding a backtick, and the tool-call markers' ids."""
+        from vllm.tokenizers.registry import cached_tokenizer_from_config
+
+        tok = cached_tokenizer_from_config(self.vllm_config.model_config)
+        self._backticks = {i: t for i in range(len(tok)) if "`" in (t := tok.decode([i]))}
+        self._tool_open = tok.convert_tokens_to_ids("<tool_call>")
+        self._tool_close = tok.convert_tokens_to_ids("</tool_call>")
+        logger.info("adaptive k: code mode %s, %d tokens hold a backtick, tool call markers %d/%d",
+                    _CODE_MODE, len(self._backticks), self._tool_open, self._tool_close)
+
+    def _track_modes(self, req_ids: list[str]) -> None:
+        """Advance each request's fence and tool-call state over its new output."""
+        for r in req_ids:
+            request = self.requests.get(r)
+            if request is None:
+                continue
+            out = request.output_token_ids
+            run, fence, tool = self._run.get(r, 0), self._in_fence.get(r, False), self._in_tool.get(r, False)
+            for t in out[self._scanned.get(r, 0):]:
+                if t == self._tool_open or t == self._tool_close:
+                    tool, run = t == self._tool_open, 0
+                    continue
+                text = self._backticks.get(t)
+                if text is None:
+                    run = 0
+                    continue
+                for ch in text:
+                    if ch != "`":
+                        run = 0
+                        continue
+                    run += 1
+                    fence ^= run == 3
+            self._scanned[r], self._run[r], self._in_fence[r], self._in_tool[r] = len(out), run, fence, tool
+            mode = _TOOL if tool else _CODE if fence else _PROSE
+            self._switches += mode != self._mode.get(r, _PROSE)
+            self._mode[r] = mode
+
+    def _active_rates(self, decoding: list[str]) -> list[np.ndarray]:
+        """Each request's rates for choosing k: its mode's where code mode says so."""
+        if _CODE_MODE == "0":
+            return [self._rates[r] for r in decoding]
+        modes = [self._mode.get(r, _PROSE) for r in decoding]
+        if _CODE_MODE == "gated" and (len(set(modes)) > 1 or modes[0] == _PROSE):
+            return [self._rates[r] for r in decoding]
+        return [self._rates[r] if m == _PROSE else
+                self._mode_rates.setdefault((r, m), self._mode_global[m].copy())
+                for r, m in zip(decoding, modes)]
+
+    def update_from_output(self, scheduler_output, model_runner_output):
+        if self._model is not None:
+            self._time_step(scheduler_output)
+        return super().update_from_output(scheduler_output, model_runner_output)
+
+    def _time_step(self, out: SchedulerOutput) -> None:
+        """Teach the cost model with this step, if it was a steady decode step.
+
+        Outputs arrive one step apart while the GPU stays busy, so the gap
+        between them is the step's time. A step with a prefill chunk is not a
+        decode step, a gap far past the prediction includes idle time, and one far
+        under it is a timing artifact.
+        """
+        now = time.monotonic()
+        last, self._last_output = self._last_output, now
+        tokens = out.total_num_scheduled_tokens
+        if not last or not tokens:
+            return
+        spec = out.scheduled_spec_decode_tokens
+        if any(n > 1 + len(spec.get(r, ())) for r, n in out.num_scheduled_tokens.items()):
+            return
+        ms = (now - last) * 1000.0
+        predicted = self._model.cost(tokens)
+        if not 0.5 * predicted < ms < 2.0 * predicted + 20.0:
+            return
+        self._model.observe(tokens, ms)
+        if now - self._last_model_log > 30:
+            self._last_model_log = now
+            m = self._model
+            logger.info("adaptive k: cost model base %.1f ms, %.3f ms/expert, %.3f ms/token, diversity %.2f, "
+                        "%d steps, mean error %.0f%%, %d code-mode switches", *m.x[:3], m.d, m.n,
+                        100 * m.err, self._switches)
+
+    def make_spec_decoding_stats(self, spec_decoding_stats, num_draft_tokens, num_accepted_tokens,
+                                 num_invalid_spec_tokens, request_id):
+        if num_draft_tokens:
+            self._observe(request_id, num_draft_tokens, num_accepted_tokens)
+        return super().make_spec_decoding_stats(spec_decoding_stats, num_draft_tokens, num_accepted_tokens,
+                                                num_invalid_spec_tokens, request_id)
+
+    def _observe(self, req_id: str, drafted: int, accepted: int) -> None:
+        g = self._global
+        c = self._rates.get(req_id)
+        if c is None:
+            c = self._rates[req_id] = g.copy()
+        self._update(c, g, drafted, accepted)
+        mode = self._mode.get(req_id, _PROSE)
+        if mode != _PROSE:
+            gm = self._mode_global[mode]
+            self._update(self._mode_rates.setdefault((req_id, mode), gm.copy()), gm, drafted, accepted)
+
+    @staticmethod
+    def _update(c: np.ndarray, g: np.ndarray, drafted: int, accepted: int) -> None:
+        for rates, alpha in ((c, _ALPHA), (g, _GLOBAL_ALPHA)):
+            rates[:accepted] += alpha * (1.0 - rates[:accepted])
+            if accepted < drafted:
+                rates[accepted] -= alpha * rates[accepted]
+        if accepted == drafted:
+            lead = c[drafted - 1] - g[drafted - 1]
+            c[drafted:] += _ALPHA * (np.clip(g[drafted:] + lead, 0.0, 1.0) - c[drafted:])
+            g[drafted:] += _GLOBAL_DRIFT * (g[drafted - 1] - g[drafted:])
+
+    def _step_cost(self, tokens: int) -> float:
+        if self._model is not None:
+            return self._model.cost(tokens)
+        x, y = self._cost_x, self._cost_y
+        if tokens <= x[-1]:
+            return float(np.interp(tokens, x, y))
+        return float(y[-1] + (y[-1] - y[-2]) / (x[-1] - x[-2]) * (tokens - x[-1]))
+
+    def _read_control(self) -> None:
+        if not self._control:
+            return
+        try:
+            mtime = os.stat(self._control).st_mtime
+        except FileNotFoundError:
+            self._force = None
+            return
+        if mtime != self._control_mtime:
+            self._control_mtime = mtime
+            words = open(self._control).read().split()
+            self._force = int(words[1]) if len(words) == 2 and words[0] == "force" else None
+            logger.info("adaptive k: control file says %s", self._force or "adapt")
+
+    def _choose(self, req_ids: list[str]) -> int:
+        self._read_control()
+        if self._force is not None:
+            return min(max(self._force, 1), self.num_spec_tokens)
+        decoding = [r for r in req_ids if r in self._rates]
+        if not decoding:
+            return self._k
+        survival = np.cumprod(np.stack(self._active_rates(decoding)), axis=1).sum(axis=0)
+        expected = {k: len(decoding) + survival[:k].sum() for k in self._levels}
+        score = {k: expected[k] / self._step_cost(len(decoding) * (k + 1)) for k in self._levels}
+        best = max(score, key=score.get)
+        current = self._k if self._k in score else best
+        return best if score[best] > score[current] * (1 + _MARGIN) else current
+
+    def _choose_per_request(self, req_ids: list[str]) -> dict[str, int]:
+        """A k for each request from one draft budget: every (request, draft)
+        slot is scored by its survival probability, and the budget is the
+        number of best slots that maximises expected tokens over step cost."""
+        decoding = [r for r in req_ids if r in self._rates]
+        if not decoding:
+            return {}
+        survival = np.cumprod(np.stack([self._rates[r] for r in decoding]), axis=1)
+        order = np.argsort(-survival, axis=None, kind="stable")
+        best_b, best_score, n = 0, len(decoding) / self._step_cost(len(decoding)), len(decoding)
+        gained = 0.0
+        for b, flat in enumerate(order, 1):
+            gained += survival.flat[flat]
+            score = (n + gained) / self._step_cost(n + b)
+            if score > best_score:
+                best_b, best_score = b, score
+        counts = np.bincount(order[:best_b] // survival.shape[1], minlength=n)
+        return {r: max(int(c), 1) for r, c in zip(decoding, counts)}
+
+    def _update_after_schedule(self, scheduler_output: SchedulerOutput) -> None:
+        for req_id in [r for r in self._rates if r not in self.requests]:
+            del self._rates[req_id]
+            for d in (self._mode, self._in_fence, self._in_tool, self._scanned, self._run):
+                d.pop(req_id, None)
+            for m in (_CODE, _TOOL):
+                self._mode_rates.pop((req_id, m), None)
+        if _CODE_MODE != "0":
+            self._track_modes(list(scheduler_output.num_scheduled_tokens))
+        self._read_control()
+        if self.num_spec_tokens and _PER_REQUEST and self._force is None:
+            per = self._choose_per_request(list(scheduler_output.num_scheduled_tokens))
+            if len(per) > 1:
+                scheduler_output.num_spec_tokens_to_schedule = max(per.values())
+                super()._update_after_schedule(scheduler_output)
+                for req_id, k in per.items():
+                    request = self.requests.get(req_id)
+                    if request is not None and request.spec_token_ids:
+                        request.spec_token_ids = [-1] * k
+                for k in per.values():
+                    self._steps[k] += 1
+                return
+        self._update_batch_k(scheduler_output)
+
+    def _update_batch_k(self, scheduler_output: SchedulerOutput) -> None:
+        if self.num_spec_tokens:
+            k = self._choose(list(scheduler_output.num_scheduled_tokens))
+            self._steps[k] += 1
+            now = time.monotonic()
+            if k != self._k and now - self._last_log > 5:
+                self._last_log = now
+                logger.info("adaptive k: %d -> %d; steps per k so far %s", self._k, k, dict(sorted(self._steps.items())))
+            self._k = k
+            scheduler_output.num_spec_tokens_to_schedule = self._k
+        super()._update_after_schedule(scheduler_output)

docker-images/glm53-kindling/experimental/arx/arx.py

@@ -0,0 +1,232 @@
+"""All-reduce over RoCE for GB10 tensor-parallel groups, without NCCL.
+
+NCCL costs about 80-90 us per decode-sized all-reduce on these boxes; this
+costs 13 us at 8 KB and 27 us at 64 KB. Each rank writes its partial to its
+peers with RDMA over both ConnectX roots (see arx_vllm.cu). Larger all-reduces,
+which are prefill-sized, stay on NCCL, whose bandwidth wins there.
+
+Enabled with VLLM_ARX_ALLREDUCE=1. The RDMA devices and GID index come from
+NCCL_IB_HCA (exactly two, one per root, in the same subnet order on every
+rank) and NCCL_IB_GID_INDEX.
+
+Ring mode (VLLM_ARX_RING=1) is for groups of 2, 3 or 4 cabled in a ring with
+no switch, with rank r's "next" port cabled to rank r+1. arx and arxbig open
+QPs only to r-1 and r+1, and at 4 the rank in between relays raw data for the
+rank beyond, so results match mesh mode bit for bit (see the .cu files). At 3
+both other ranks are neighbours and nothing is relayed. arxbig's ring
+reduce-scatter needs 2 or 4 ranks, so a ring of 3 runs arx alone
+(VLLM_ARXBIG=0). ARX_RING_PREV_HCAS and ARX_RING_NEXT_HCAS name the two RDMA
+devices of the port facing each neighbour, in the order whose root r shares a
+subnet with root r at that neighbour: root 0 then root 1 when every box is
+cabled the same way, the other order on a box whose cabling is mirrored. Each
+port has its own subnet, so each device uses the RoCE v2 GID of the IPv4
+address on its own netdev.
+"""
+import os
+
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
+
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+_ext = None
+_taken = False  # the extension holds one connection per process
+
+
+def _hcas(var: str) -> list[str]:
+    return [h.strip("=^").split(":")[0] for h in os.environ.get(var, "").split(",") if h]
+
+
+def roce_v2_gid(dev: str) -> int:
+    """GID index of dev's RoCE v2 GID for an IPv4 address on dev's own netdev.
+
+    The kernel fills the GID table from the netdev's addresses, and ndevs
+    names the netdev an entry came from, so a VLAN or other upper device's
+    entries are skipped.
+    """
+    base = f"/sys/class/infiniband/{dev}"
+    netdevs = set(os.listdir(f"{base}/device/net"))
+    port = f"{base}/ports/1"
+    found = []
+    for name in sorted(os.listdir(f"{port}/gids"), key=int):
+        try:
+            with open(f"{port}/gids/{name}") as f:
+                gid = f.read().strip()
+            with open(f"{port}/gid_attrs/types/{name}") as f:
+                kind = f.read().strip()
+            with open(f"{port}/gid_attrs/ndevs/{name}") as f:
+                ndev = f.read().strip()
+        except OSError:  # empty entries refuse to report their attributes
+            continue
+        if kind == "RoCE v2" and ndev in netdevs and gid.startswith("0000:0000:0000:0000:0000:ffff:"):
+            found.append(int(name))
+    if not found:
+        raise RuntimeError(f"arx: {dev} has no RoCE v2 GID for an IPv4 address on {sorted(netdevs)}")
+    if len(found) > 1:
+        logger.warning("arx: %s has IPv4 GIDs at %s; using %d", dev, found, found[0])
+    return found[0]
+
+
+def fabric(world: int, who: str):
+    """(devices, GID index per device, ring) for this rank, or None to use NCCL."""
+    if os.environ.get("VLLM_ARX_RING") == "1":
+        prev, nxt = _hcas("ARX_RING_PREV_HCAS"), _hcas("ARX_RING_NEXT_HCAS")
+        if len(prev) != 2 or len(nxt) != 2:
+            logger.warning("%s ring mode needs two devices each in ARX_RING_PREV_HCAS and ARX_RING_NEXT_HCAS, "
+                           "got %r and %r; using NCCL", who, prev, nxt)
+            return None
+        if world not in (2, 3, 4):
+            raise ValueError(f"{who} ring mode supports groups of 2 to 4 ranks, not {world}")
+        devs = prev + nxt
+        return devs, [roce_v2_gid(d) for d in devs], True
+    hcas = _hcas("NCCL_IB_HCA")
+    if len(hcas) != 2:
+        logger.warning("%s needs two RDMA devices in NCCL_IB_HCA, got %r; using NCCL", who, hcas)
+        return None
+    gid = int(os.environ.get("NCCL_IB_GID_INDEX", "5"))
+    return hcas, [gid, gid], False
+
+
+def _jit_load(name, *args, **kwargs):
+    """torch's cpp_extension.load(), safe after a build was killed midway (#53).
+
+    torch marks a build with a `lock` file and waits forever on one it finds, so a
+    build killed midway hangs every later load with nothing in the log. Builds here
+    also hold an flock, which the kernel drops when its holder dies: whoever holds
+    it is the only builder, and any `lock` file it finds is stale."""
+    import fcntl
+
+    from torch.utils.cpp_extension import _get_build_directory, load
+
+    build = _get_build_directory(name, False)
+    with open(os.path.join(build, "lock.flock"), "w") as f:
+        fcntl.flock(f, fcntl.LOCK_EX)
+        try:
+            os.remove(os.path.join(build, "lock"))
+        except FileNotFoundError:
+            pass
+        return load(name, *args, **kwargs)
+
+
+def _load():
+    global _ext
+    if _ext is None:
+        _ext = _jit_load(
+            "arx_vllm",
+            [os.path.join(os.path.dirname(__file__), "arx_vllm.cu")],
+            extra_cuda_cflags=["-O3", "-std=c++17", "-gencode=arch=compute_121a,code=sm_121a"],
+            extra_ldflags=["-libverbs"],
+        )
+    return _ext
+
+
+# (address, bytes) ranges for the next all-reduce to prefetch into L2 while it waits.
+_prefetch: tuple = ()
+
+
+def set_prefetch(tensors) -> None:
+    """The next arx all-reduce asks L2 for these tensors' bytes while it waits for peers."""
+    global _prefetch
+    _prefetch = tuple((t.data_ptr(), t.numel() * t.element_size()) for t in tensors if t is not None)
+
+
+class ArxCommunicator:
+    def __init__(self, group: ProcessGroup, device: torch.device):
+        global _taken
+        self.disabled = True
+        if _taken:
+            logger.warning("arx already serves another group in this process; using NCCL here")
+            return
+        rank, world = dist.get_rank(group), dist.get_world_size(group)
+        fab = fabric(world, "arx")
+        if fab is None:
+            return
+        devs, gids, ring = fab
+        with torch.cuda.device(device):
+            ext = _load()
+            info = ext.prepare(rank, world, devs, gids, ring)
+            infos: list[bytes] = [b""] * world
+            dist.all_gather_object(infos, info, group=group)
+            ext.connect(infos)
+        dist.barrier(group=group)
+        _taken = True
+        self.ext = ext
+        self.max_bytes = min(int(os.environ.get("VLLM_ARX_MAX_BYTES", 256 << 10)), ext.max_bytes())
+        # Two-shot (reduce-scatter + all-gather) for all-reduces above max_bytes,
+        # up to this size: 2-2.7x faster than NCCL from 512 KB. Mesh only; 0 is off.
+        self.two_shot_max = min(int(os.environ.get("VLLM_ARX_TWO_SHOT_MAX_KB", "0")) << 10, ext.max_bytes2())
+        self.world = world
+        self.disabled = False
+        logger.info("arx all-reduce: rank %d/%d%s on %s, gids %s, up to %d bytes%s", rank, world,
+                    " ring (prev, next)" if ring else "", devs, gids, self.max_bytes,
+                    f", two-shot up to {self.two_shot_max} bytes" if self.two_shot_max else "")
+
+    def should_use(self, t: torch.Tensor, two_shot: bool = False) -> bool:
+        """Whether arx takes t: up to max_bytes, or up to two_shot_max when the caller allows it."""
+        if not (t.dtype == torch.bfloat16 and t.is_cuda and t.is_contiguous() and t.numel() % 8 == 0):
+            return False
+        nbytes = t.numel() * 2
+        return nbytes <= self.max_bytes or (
+            two_shot and nbytes <= self.two_shot_max and t.numel() % (8 * self.world) == 0)
+
+    def all_reduce(self, t: torch.Tensor) -> torch.Tensor:
+        global _prefetch
+        out = torch.empty_like(t)
+        if t.numel() * 2 > self.max_bytes:
+            self.ext.allreduce2(t, out)
+            return out
+        pf, _prefetch = _prefetch, ()
+        self.ext.allreduce(t, out, [p for p, _ in pf], [b for _, b in pf])
+        return out
+
+
+_big_ext = None
+
+
+class ArxBig:
+    """Prefill-sized all-gathers over RoCE (arxbig.cu): ~10% faster than NCCL's.
+
+    Enabled with VLLM_ARXBIG=1 next to arx. all_gather returns a view of a
+    pinned buffer that the fourth all-gather after it overwrites, so a caller
+    that keeps the result longer must copy it.
+    """
+
+    def __init__(self, group: ProcessGroup, device: torch.device):
+        global _big_ext
+        self.disabled = True
+        self.rs = False
+        self.gather = False
+        rank, world = dist.get_rank(group), dist.get_world_size(group)
+        fab = fabric(world, "arxbig")
+        if fab is None:
+            return
+        devs, gids, ring = fab
+        slot = int(os.environ.get("VLLM_ARXBIG_SLOT_MB", "132")) << 20
+        slot -= slot % (world * 16)
+        with torch.cuda.device(device):
+            if _big_ext is None:
+                _big_ext = _jit_load("arxbig", [os.path.join(os.path.dirname(__file__), "arxbig.cu")],
+                                     extra_cuda_cflags=["-O3", "-std=c++17", "-gencode=arch=compute_121a,code=sm_121a"],
+                                     extra_ldflags=["-libverbs"])
+            self.rs = os.environ.get("VLLM_ARXBIG_RS") == "1"
+            info = _big_ext.prepare(rank, world, devs, gids, ring, slot, slot if self.rs else 0)
+            infos: list[bytes] = [b""] * world
+            dist.all_gather_object(infos, info, group=group)
+            _big_ext.connect(infos)
+            dist.barrier(group=group)
+        self.ext, self.world, self.slot = _big_ext, world, slot
+        self.min_bytes = int(os.environ.get("VLLM_ARXBIG_MIN_KB", "1024")) << 10
+        self.gather = os.environ.get("VLLM_ARXBIG_AG") == "1"
+        self.disabled = False
+        logger.info("arxbig all-gather: rank %d/%d%s, %d MiB slots", rank, world, " ring" if ring else "", slot >> 20)
+
+    def should_gather(self, t: torch.Tensor) -> bool:
+        total = t.numel() * t.element_size() * self.world
+        return (t.is_cuda and t.dim() >= 1 and (t.numel() * t.element_size()) % 16 == 0
+                and self.min_bytes <= total <= self.slot)
+
+    def all_gather(self, t: torch.Tensor) -> torch.Tensor:
+        return self.ext.all_gather(t.contiguous())

docker-images/glm53-kindling/experimental/arx/arx_vllm.cu

@@ -0,0 +1,620 @@
+// arx: one-shot all-reduce over RoCE for GB10, no NCCL.
+//
+// GB10's CPU and GPU share one coherent memory, so pinned host buffers are
+// GPU buffers at full speed and the ConnectX can DMA into them without
+// GPUDirect. Per all-reduce:
+//   1. the GPU writes its partial into send[seq & 1] and publishes seq;
+//   2. a CPU proxy sees seq, RDMA-writes the partial into every peer's
+//      recv[seq & 1][rank], then writes seq into the peer's flag[rank]
+//      (same QP, so the data is placed first);
+//   3. the GPU waits for every peer's flag to reach seq and sums.
+// Two buffers by parity are enough: a rank reaches seq + 2 only after every
+// peer consumed seq + 1, which needed this rank's seq + 1 flag, which the
+// proxy posts after this rank's seq data was read.
+//
+// The proxy posts every seq in order. The GPU can publish seq + 1 before the
+// proxy has seen seq (finishing seq needs only the peers' data), and skipping
+// seq would let peers accept seq + 1's flag over stale seq data.
+//
+// Every partial is split in half, one half per ConnectX root, each followed by
+// its own flag on that root's QP: flag[src * 2 + root]. Rank r sends to
+// r+1, r+2, ... in turn, so no receiver takes every sender at once.
+//
+// Ring mode is for boxes cabled in a ring with no switch. The ConnectX cannot
+// forward RoCE for other boxes, so QPs go only to prev = r-1 and next = r+1,
+// each over the two functions of the port facing that neighbour. A rank sends
+// its partial to next, and at world 3 and 4 to prev as well. At world 4 the
+// proxy forwards the partial from prev on to next. The relayed data lands in
+// the same recv[seq & 1][src] and flag[src * 2 + root] as a direct write, so
+// the GPU side is unchanged. At world 3 prev and next are the two other ranks,
+// so every partial goes direct and nothing is relayed. World 2 sends to next
+// only.
+//
+// The relay adds one reuse hazard. Rank m reads recv[par][m-1] to forward it
+// to m+1, and m-1 writes seq + 2 there next. m-1 publishes seq + 2 after
+// finishing seq + 1. Finishing needs m+1's seq + 1 partial. m+1 publishes that
+// after finishing seq, and finishing seq needed the relayed data. The relay
+// takes the partial's size from this rank's own publish of the same seq. The
+// same chain keeps this rank from publishing seq + 2 before the relay.
+//
+// Two-shot (mesh only), for partials too big to send whole to every peer:
+// allreduce2() runs two collectives in a row, each its own seq. Reduce-scatter:
+// the proxy writes chunk j of this rank's partial into rank j's recv slot, and
+// the GPU sums its own chunk from every rank (fp32, rank order, one rounding,
+// so the bits match the one-shot sum). All-gather: the proxy writes the summed
+// chunk into every peer's slot, and the GPU assembles the chunks in rank order.
+// Each rank sends 2 (world - 1) / world partials' worth instead of world - 1.
+// Ctl::mode tells the proxy which pattern a seq uses.
+//
+// vLLM form: prepare() returns this rank's connection details, the caller
+// all-gathers them over its own group, connect() wires the QPs and starts the
+// proxy, then allreduce(in, out) runs on the current stream and is
+// CUDA-graph capturable. in/out are bf16 device tensors.
+#include <arpa/inet.h>
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <infiniband/verbs.h>
+#include <netinet/in.h>
+#include <netinet/tcp.h>
+#include <sys/socket.h>
+#include <unistd.h>
+
+#include <algorithm>
+#include <atomic>
+#include <chrono>
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <thread>
+#include <vector>
+
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/extension.h>
+
+#define CK(x) do { cudaError_t e_ = (x); if (e_ != cudaSuccess) { fprintf(stderr, "%s:%d %s\n", __FILE__, __LINE__, cudaGetErrorString(e_)); exit(1); } } while (0)
+#define IBCK(x) do { if (!(x)) { fprintf(stderr, "%s:%d %s failed: %s\n", __FILE__, __LINE__, #x, strerror(errno)); exit(1); } } while (0)
+
+constexpr int kMaxWorld = 8;
+constexpr int kMaxDev = 4;  // ring mode: 2 ports x 2 roots
+constexpr size_t kMaxBytes = 512 << 10;  // one-shot partial: bf16 [32, 4096] is 256 KB
+constexpr size_t kSendBytes = 4 << 20;   // per parity: a two-shot partial
+constexpr size_t kSlotBytes = 1 << 20;   // per peer per parity: a one-shot partial or a two-shot chunk
+enum : uint64_t { kFull = 0, kScatter = 1, kGather = 2 };
+
+struct Ctl {                   // pinned; written by the GPU, read by the proxy
+  volatile uint64_t seq;       // last published partial
+  volatile uint64_t bytes[2];  // size of the partial in send[parity]
+  volatile uint64_t mode[2];   // kFull, kScatter or kGather for send[parity]
+  volatile uint64_t stop;
+  volatile uint64_t done;      // last seq the GPU finished summing (for the watchdog)
+};
+
+struct PeerInfo {  // what rank r tells everyone about its QPs
+  uint32_t qpn[kMaxWorld][2];  // per channel (see State), per root
+  uint8_t gid[kMaxDev][16];    // per RDMA device
+  uint32_t mtu[kMaxDev];       // per device: the port's active MTU (enum ibv_mtu)
+  uint64_t recv_addr, flag_addr;
+  uint32_t recv_rkey[kMaxDev], flag_rkey[kMaxDev];
+};
+
+// ---- GPU side ---------------------------------------------------------------
+
+__device__ __forceinline__ uint64_t ld_acquire_sys(const volatile uint64_t* p) {
+  uint64_t v;
+  asm volatile("ld.acquire.sys.global.u64 %0, [%1];" : "=l"(v) : "l"(p) : "memory");
+  return v;
+}
+
+__device__ __forceinline__ uint4 ld_cv(const void* p) {  // bypass caches: the NIC wrote it
+  uint4 v;
+  asm volatile("ld.global.cv.v4.u32 {%0,%1,%2,%3}, [%4];" : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "l"(p));
+  return v;
+}
+
+
+struct Dev {
+  __nv_bfloat16* send;           // [2][kSendBytes / 2]
+  const __nv_bfloat16* recv;     // [2][world][kSlotBytes / 2]
+  const volatile uint64_t* flag; // [2 * world]
+  Ctl* ctl;
+  unsigned* seq_dev;             // device: last completed seq
+  unsigned* blocks;         // device [2]: blocks past the copy, blocks finished
+  unsigned* go;                  // device: seq whose peer data has all arrived
+  int rank, world;
+};
+
+// Weight ranges the next kernels read; the kernel asks L2 for them before it waits.
+constexpr int kMaxPrefetch = 4;
+struct Prefetch {
+  const uint8_t* ptr[kMaxPrefetch];
+  long long bytes[kMaxPrefetch];
+  int count;
+};
+
+// out = sum over ranks of in. in and out may alias.
+__global__ void arx_kernel(Dev d, const __nv_bfloat16* __restrict__ in, __nv_bfloat16* out, int n, Prefetch pf) {
+  const uint64_t seq = *d.seq_dev + 1;
+  const int par = seq & 1;
+  __nv_bfloat16* mine = d.send + (size_t)par * (kSendBytes / 2);
+  const int stride = gridDim.x * blockDim.x * 8;
+  for (int i = (blockIdx.x * blockDim.x + threadIdx.x) * 8; i < n; i += stride)
+    *reinterpret_cast<uint4*>(mine + i) = *reinterpret_cast<const uint4*>(in + i);
+  __threadfence_system();
+  __syncthreads();
+  // Separate counters per phase: blocks that start late would otherwise add
+  // their first count after early blocks added their second, and the
+  // publish would be skipped or would go out before every block copied.
+  if (threadIdx.x == 0 && atomicAdd(&d.blocks[0], 1) == gridDim.x - 1) {  // last block publishes
+    d.blocks[0] = 0;
+    d.ctl->mode[par] = kFull;
+    d.ctl->bytes[par] = (uint64_t)n * 2;
+    __threadfence_system();
+    d.ctl->seq = seq;
+  }
+  // DRAM is idle while the peers' data is in flight: spread one bulk L2
+  // prefetch per thread over each range.
+  if (pf.count > 0) {
+    const long long nthr = (long long)gridDim.x * blockDim.x, t = (long long)blockIdx.x * blockDim.x + threadIdx.x;
+    for (int r = 0; r < pf.count; ++r) {
+      const long long chunk = ((pf.bytes[r] + nthr - 1) / nthr + 15) & ~15LL, off = t * chunk;
+      const long long len = min(chunk, pf.bytes[r] - off) & ~15LL;
+      if (len > 0)
+        asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;\n" ::"l"(pf.ptr[r] + off), "r"((unsigned)len));
+    }
+  }
+  // Block 0 polls the pinned flags the NIC writes; the others wait on a word in
+  // device memory, so only a few threads touch lines the NIC is filling.
+  if (blockIdx.x == 0) {
+    if (threadIdx.x < 2 * d.world && threadIdx.x / 2 != d.rank)
+      while (ld_acquire_sys(d.flag + threadIdx.x) < seq) {}
+    __syncthreads();
+    if (threadIdx.x == 0) { __threadfence(); atomicExch(d.go, (unsigned)seq); }
+  } else {
+    if (threadIdx.x == 0)
+      while (true) {
+        unsigned g;
+        asm volatile("ld.acquire.gpu.global.u32 %0, [%1];" : "=r"(g) : "l"(d.go) : "memory");
+        if (g >= (unsigned)seq) break;
+      }
+    __syncthreads();
+  }
+  for (int i = (blockIdx.x * blockDim.x + threadIdx.x) * 8; i < n; i += stride) {
+    uint4 v[kMaxWorld];
+#pragma unroll
+    for (int src = 0; src < kMaxWorld; ++src)
+      if (src < d.world)
+        v[src] = src == d.rank ? *reinterpret_cast<const uint4*>(mine + i)
+                               : ld_cv(d.recv + ((size_t)par * d.world + src) * (kSlotBytes / 2) + i);
+    float acc[8] = {};
+#pragma unroll
+    for (int src = 0; src < kMaxWorld; ++src)  // rank order: every rank gets the same bits
+      if (src < d.world) {
+        const __nv_bfloat16* b = reinterpret_cast<const __nv_bfloat16*>(&v[src]);
+#pragma unroll
+        for (int k = 0; k < 8; ++k) acc[k] += __bfloat162float(b[k]);
+      }
+    __align__(16) __nv_bfloat16 o[8];
+#pragma unroll
+    for (int k = 0; k < 8; ++k) o[k] = __float2bfloat16(acc[k]);
+    *reinterpret_cast<uint4*>(out + i) = *reinterpret_cast<const uint4*>(o);
+  }
+  __syncthreads();
+  if (threadIdx.x == 0 && atomicAdd(&d.blocks[1], 1) == gridDim.x - 1) {  // last block out
+    d.blocks[1] = 0;
+    *d.seq_dev = (unsigned)seq;
+    d.ctl->done = seq;
+  }
+}
+
+// The two-shot kernels: copy n bf16 into send[par], publish them as `mode`,
+// wait for every peer's flag, then leave the rest to the caller.
+__device__ __forceinline__ uint64_t publish_and_wait(Dev d, const __nv_bfloat16* __restrict__ in, int n,
+                                                     uint64_t mode) {
+  const uint64_t seq = *d.seq_dev + 1;
+  const int par = seq & 1;
+  __nv_bfloat16* mine = d.send + (size_t)par * (kSendBytes / 2);
+  const int stride = gridDim.x * blockDim.x * 8;
+  for (int i = (blockIdx.x * blockDim.x + threadIdx.x) * 8; i < n; i += stride)
+    *reinterpret_cast<uint4*>(mine + i) = *reinterpret_cast<const uint4*>(in + i);
+  __threadfence_system();
+  __syncthreads();
+  if (threadIdx.x == 0 && atomicAdd(&d.blocks[0], 1) == gridDim.x - 1) {
+    d.blocks[0] = 0;
+    d.ctl->mode[par] = mode;
+    d.ctl->bytes[par] = (uint64_t)n * 2;
+    __threadfence_system();
+    d.ctl->seq = seq;
+  }
+  if (blockIdx.x == 0) {
+    if (threadIdx.x < 2 * d.world && threadIdx.x / 2 != d.rank)
+      while (ld_acquire_sys(d.flag + threadIdx.x) < seq) {}
+    __syncthreads();
+    if (threadIdx.x == 0) { __threadfence(); atomicExch(d.go, (unsigned)seq); }
+  } else {
+    if (threadIdx.x == 0)
+      while (true) {
+        unsigned g;
+        asm volatile("ld.acquire.gpu.global.u32 %0, [%1];" : "=r"(g) : "l"(d.go) : "memory");
+        if (g >= (unsigned)seq) break;
+      }
+    __syncthreads();
+  }
+  return seq;
+}
+
+__device__ __forceinline__ void finish(Dev d, uint64_t seq) {
+  __syncthreads();
+  if (threadIdx.x == 0 && atomicAdd(&d.blocks[1], 1) == gridDim.x - 1) {
+    d.blocks[1] = 0;
+    *d.seq_dev = (unsigned)seq;
+    d.ctl->done = seq;
+  }
+}
+
+// out[chunk] = this rank's chunk of the sum of every rank's in[n].
+__global__ void arx_scatter_kernel(Dev d, const __nv_bfloat16* __restrict__ in, __nv_bfloat16* out, int n) {
+  const uint64_t seq = publish_and_wait(d, in, n, kScatter);
+  const int par = seq & 1, chunk = n / d.world;
+  const __nv_bfloat16* mine = d.send + (size_t)par * (kSendBytes / 2) + (size_t)d.rank * chunk;
+  const int stride = gridDim.x * blockDim.x * 8;
+  for (int i = (blockIdx.x * blockDim.x + threadIdx.x) * 8; i < chunk; i += stride) {
+    uint4 v[kMaxWorld];
+#pragma unroll
+    for (int src = 0; src < kMaxWorld; ++src)
+      if (src < d.world)
+        v[src] = src == d.rank ? *reinterpret_cast<const uint4*>(mine + i)
+                               : ld_cv(d.recv + ((size_t)par * d.world + src) * (kSlotBytes / 2) + i);
+    float acc[8] = {};
+#pragma unroll
+    for (int src = 0; src < kMaxWorld; ++src)  // rank order, as the one-shot sum
+      if (src < d.world) {
+        const __nv_bfloat16* b = reinterpret_cast<const __nv_bfloat16*>(&v[src]);
+#pragma unroll
+        for (int k = 0; k < 8; ++k) acc[k] += __bfloat162float(b[k]);
+      }
+    __align__(16) __nv_bfloat16 o[8];
+#pragma unroll
+    for (int k = 0; k < 8; ++k) o[k] = __float2bfloat16(acc[k]);
+    *reinterpret_cast<uint4*>(out + i) = *reinterpret_cast<const uint4*>(o);
+  }
+  finish(d, seq);
+}
+
+// out[world * chunk] = every rank's chunk, in rank order.
+__global__ void arx_gather_kernel(Dev d, const __nv_bfloat16* __restrict__ chunk_in, __nv_bfloat16* out, int chunk) {
+  const uint64_t seq = publish_and_wait(d, chunk_in, chunk, kGather);
+  const int par = seq & 1, n = chunk * d.world;
+  const int stride = gridDim.x * blockDim.x * 8;
+  for (int i = (blockIdx.x * blockDim.x + threadIdx.x) * 8; i < n; i += stride) {
+    const int src = i / chunk, j = i - src * chunk;
+    *reinterpret_cast<uint4*>(out + i) =
+        src == d.rank ? *reinterpret_cast<const uint4*>(chunk_in + j)
+                      : ld_cv(d.recv + ((size_t)par * d.world + src) * (kSlotBytes / 2) + j);
+  }
+  finish(d, seq);
+}
+
+// ---- host side ---------------------------------------------------------------
+
+namespace {
+// A channel is one QP per root to one peer. Mesh: channel j goes to rank j
+// over device r. Ring: channel 0 goes to prev over devices 0 and 1, channel 1
+// to next over devices 2 and 3, and meets the peer's channel on the other side.
+struct State {
+  int rank = -1, world = 0, ndev = 0;
+  bool ring = false;
+  int prev = 0, next = 0;
+  int gid_idx[kMaxDev] = {};
+  uint8_t *send_h = nullptr, *recv_h = nullptr;
+  uint64_t* flag_h = nullptr;
+  Ctl* ctl = nullptr;
+  ibv_context* ctx[kMaxDev] = {};
+  ibv_pd* pd[kMaxDev] = {};
+  ibv_mr *send_mr[kMaxDev] = {}, *recv_mr[kMaxDev] = {}, *flag_mr[kMaxDev] = {};
+  ibv_cq* cq[kMaxDev] = {};
+  ibv_qp* qp[kMaxWorld][2] = {};
+  PeerInfo mine{};
+  std::vector<PeerInfo> all;
+  Dev dev{};
+  bool connected = false;
+};
+State S;
+std::atomic<bool> g_proxy_err{false};
+
+// Idle backoff (#55): keep the exact busy spin while there is work and for ARX_IDLE_SPIN_MS after it (default
+// 200 ms), then nap ARX_IDLE_NAP_US (default 200 us) per loop turn. ARX_IDLE_NAP_US=0 keeps the old spin.
+struct IdleNap {
+  std::chrono::steady_clock::duration spin;
+  std::chrono::microseconds nap;
+  std::chrono::steady_clock::time_point last;
+  IdleNap() {
+    const char* s = std::getenv("ARX_IDLE_SPIN_MS");
+    const char* n = std::getenv("ARX_IDLE_NAP_US");
+    spin = std::chrono::milliseconds(s ? std::atol(s) : 200);
+    nap = std::chrono::microseconds(n ? std::atol(n) : 200);
+    last = std::chrono::steady_clock::now();
+    fprintf(stderr, "arx idle nap: spin %ld ms, nap %ld us%s\n", (long)(s ? std::atol(s) : 200), (long)nap.count(),
+            nap.count() <= 0 ? " (disabled = old spin)" : "");
+  }
+  void busy() { last = std::chrono::steady_clock::now(); }
+  void idle() {
+    if (nap.count() <= 0) return;
+    if (std::chrono::steady_clock::now() - last > spin) std::this_thread::sleep_for(nap);
+  }
+};
+
+int chan_peer(int c) { return S.ring ? (c ? S.next : S.prev) : c; }
+int chan_dev(int c, int r) { return S.ring ? c * 2 + r : r; }         // device here
+int chan_rdev(int c, int r) { return S.ring ? (1 - c) * 2 + r : r; }  // device at the peer
+
+// One write of [laddr, laddr + len) to the peer's recv at roff, then seq into
+// its flag slot, on channel c's root r QP.
+bool post(int c, int r, uint64_t laddr, uint32_t lkey, uint32_t len, uint64_t roff, int slot, uint64_t seq,
+          uint64_t& posted) {
+  const PeerInfo& p = S.all[chan_peer(c)];
+  const int rd = chan_rdev(c, r);
+  ibv_sge sg{laddr, len, lkey};
+  ibv_send_wr w{}, f{}, *bad;
+  w.opcode = IBV_WR_RDMA_WRITE; w.sg_list = &sg; w.num_sge = 1;
+  w.wr.rdma.remote_addr = p.recv_addr + roff;
+  w.wr.rdma.rkey = p.recv_rkey[rd];
+  ibv_sge fs{(uint64_t)&seq, 8, 0};
+  f.opcode = IBV_WR_RDMA_WRITE; f.sg_list = &fs; f.num_sge = 1; f.send_flags = IBV_SEND_INLINE;
+  f.wr.rdma.remote_addr = p.flag_addr + slot * 8; f.wr.rdma.rkey = p.flag_rkey[rd];
+  // Unsignalled WRs are reclaimed only by a later signalled one on the
+  // same QP, so the count is per QP.
+  if ((++posted & 31) == 0) f.send_flags |= IBV_SEND_SIGNALED;
+  w.next = &f;
+  if (int e = ibv_post_send(S.qp[c][r], &w, &bad)) {
+    fprintf(stderr, "arx rank %d: post to %d root %d failed: %s\n", S.rank, chan_peer(c), r, strerror(e));
+    return false;
+  }
+  return true;
+}
+
+void proxy_loop() {
+  uint64_t seq = 0, posted[kMaxWorld][2] = {};
+  const int rank = S.rank, world = S.world;
+  auto last_move = std::chrono::steady_clock::now();
+  uint64_t last_done = 0;
+  bool reported = false;
+  // Ring at world 4: the last seq relayed from prev to next, per root, and
+  // the size of this rank's own partial for each parity.
+  uint64_t relayed[2] = {}, own_bytes[2] = {};
+  const bool relay = S.ring && world == 4;
+  const volatile uint64_t* flags = S.flag_h;
+  IdleNap nap;
+  while (!S.ctl->stop) {
+    for (int r = 0; relay && r < 2; ++r)
+      while (relayed[r] < seq && flags[S.prev * 2 + r] > relayed[r]) {
+        std::atomic_thread_fence(std::memory_order_acquire);
+        const uint64_t s = ++relayed[r];
+        const int par = s & 1;
+        const uint32_t bytes = (uint32_t)own_bytes[par], half = (bytes / 2 + 15) & ~15u;
+        const uint32_t off = r ? half : 0, len = r ? bytes - half : half;
+        const uint64_t at = ((uint64_t)par * world + S.prev) * kSlotBytes + off;
+        if (!post(1, r, (uint64_t)(S.recv_h + at), S.recv_mr[chan_dev(1, r)]->lkey, len, at, S.prev * 2 + r, s,
+                  posted[1][r])) {
+          g_proxy_err = true;
+          return;
+        }
+      }
+    if (S.ctl->seq <= seq) {
+      // Watchdog: published work that has not completed for 2 s is a stall,
+      // and so is a peer whose flags are ahead of what this rank posted.
+      // Print what this rank sent and what it has heard from each peer.
+      const uint64_t done = S.ctl->done;
+      const auto now = std::chrono::steady_clock::now();
+      bool behind = false;
+      for (int j = 0; j < world; ++j)
+        if (j != rank) behind |= S.flag_h[j * 2] > seq;
+      if (done != last_done || (done >= seq && !behind)) {
+        last_done = done; last_move = now; reported = false;
+      }
+      else if (!reported && now - last_move > std::chrono::seconds(2)) {
+        reported = true;
+        char buf[512];
+        int o = snprintf(buf, sizeof buf, "arx rank %d STALL: published %lu posted %lu done %lu; flags from peers:",
+                         rank, (unsigned long)S.ctl->seq, (unsigned long)seq, (unsigned long)done);
+        for (int j = 0; j < world; ++j)
+          if (j != rank)
+            o += snprintf(buf + o, sizeof buf - o, " r%d=%lu/%lu", j, (unsigned long)S.flag_h[j * 2],
+                          (unsigned long)S.flag_h[j * 2 + 1]);
+        if (relay)
+          o += snprintf(buf + o, sizeof buf - o, "; relayed %lu/%lu", (unsigned long)relayed[0],
+                        (unsigned long)relayed[1]);
+        fprintf(stderr, "%s\n", buf);
+      }
+      if (relay && (relayed[0] < seq || relayed[1] < seq)) nap.busy(); else nap.idle();
+      continue;
+    }
+    nap.busy();
+    ++seq;  // every seq, in order
+    std::atomic_thread_fence(std::memory_order_acquire);
+    const int par = seq & 1;
+    const uint64_t mode = S.ctl->mode[par];
+    const uint32_t total = (uint32_t)S.ctl->bytes[par];
+    // A scatter sends each peer its own chunk; full and gather send everything.
+    const uint32_t bytes = mode == kScatter ? total / world : total;
+    const uint32_t half = (bytes / 2 + 15) & ~15u;  // root 0 takes the first half
+    own_bytes[par] = bytes;
+    // Mesh: channel j is rank j. Ring: next, then prev unless prev is next.
+    const int nchan = S.ring ? (world == 2 ? 1 : 2) : world - 1;
+    for (int step = 1; step <= nchan; ++step) {
+      const int c = S.ring ? 2 - step : (rank + step) % world;
+      const uint64_t base = mode == kScatter ? (uint64_t)chan_peer(c) * bytes : 0;
+      for (int r = 0; r < 2; ++r) {
+        const uint32_t off = r ? half : 0, len = r ? bytes - half : half;
+        if (!post(c, r, (uint64_t)(S.send_h + par * kSendBytes + base + off), S.send_mr[chan_dev(c, r)]->lkey, len,
+                  ((uint64_t)par * world + rank) * kSlotBytes + off, rank * 2 + r, seq, posted[c][r])) {
+          g_proxy_err = true;
+          return;
+        }
+      }
+    }
+    for (int r = 0; r < S.ndev; ++r) {
+      ibv_wc wc[16];
+      const int n = ibv_poll_cq(S.cq[r], 16, wc);
+      for (int i = 0; i < n; ++i)
+        if (wc[i].status != IBV_WC_SUCCESS) {
+          fprintf(stderr, "arx rank %d: completion error %d\n", rank, wc[i].status);
+          g_proxy_err = true;
+          return;
+        }
+    }
+  }
+}
+}  // namespace
+
+// Allocates buffers, opens the RDMA devices and creates the QPs. Returns this
+// rank's PeerInfo for the caller to all-gather. Mesh: devs are root 0 and root
+// 1. Ring: the port facing prev (root 0, root 1), then the port facing next.
+// gids: each device's GID index.
+py::bytes arx_prepare(int64_t rank, int64_t world, std::vector<std::string> devs, std::vector<int64_t> gids,
+                      bool ring) {
+  TORCH_CHECK(S.rank < 0, "arx is already prepared in this process");
+  TORCH_CHECK(world >= 2 && world <= kMaxWorld);
+  TORCH_CHECK(!ring || (world >= 2 && world <= 4), "arx ring mode supports world 2 to 4, not ", world);
+  S.ndev = ring ? 4 : 2;
+  TORCH_CHECK((int)devs.size() == S.ndev && (int)gids.size() == S.ndev, "arx: expected ", S.ndev, " devices");
+  S.rank = rank; S.world = world; S.ring = ring;
+  S.prev = (rank + world - 1) % world; S.next = (rank + 1) % world;
+  CK(cudaHostAlloc(&S.send_h, 2 * kSendBytes, cudaHostAllocMapped));
+  CK(cudaHostAlloc(&S.recv_h, 2 * world * kSlotBytes, cudaHostAllocMapped));
+  CK(cudaHostAlloc(&S.flag_h, 2 * kMaxWorld * sizeof(uint64_t), cudaHostAllocMapped));
+  CK(cudaHostAlloc(&S.ctl, sizeof(Ctl), cudaHostAllocMapped));
+  memset(S.flag_h, 0, 2 * kMaxWorld * sizeof(uint64_t));
+  memset((void*)S.ctl, 0, sizeof(Ctl));
+  int ndev;
+  ibv_device** list = ibv_get_device_list(&ndev);
+  const int acc = IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_WRITE;
+  for (int d = 0; d < S.ndev; ++d) {
+    S.gid_idx[d] = gids[d];
+    for (int i = 0; i < ndev; ++i)
+      if (devs[d] == ibv_get_device_name(list[i])) S.ctx[d] = ibv_open_device(list[i]);
+    TORCH_CHECK(S.ctx[d], "arx: no RDMA device ", devs[d]);
+    S.pd[d] = ibv_alloc_pd(S.ctx[d]); IBCK(S.pd[d]);
+    S.send_mr[d] = ibv_reg_mr(S.pd[d], S.send_h, 2 * kSendBytes, acc); IBCK(S.send_mr[d]);
+    S.recv_mr[d] = ibv_reg_mr(S.pd[d], S.recv_h, 2 * world * kSlotBytes, acc); IBCK(S.recv_mr[d]);
+    S.flag_mr[d] = ibv_reg_mr(S.pd[d], S.flag_h, 2 * kMaxWorld * sizeof(uint64_t), acc); IBCK(S.flag_mr[d]);
+    S.cq[d] = ibv_create_cq(S.ctx[d], 4096, nullptr, nullptr, 0); IBCK(S.cq[d]);
+    ibv_gid gid;
+    IBCK(ibv_query_gid(S.ctx[d], 1, S.gid_idx[d], &gid) == 0);
+    memcpy(S.mine.gid[d], gid.raw, 16);
+    ibv_port_attr port{};
+    IBCK(ibv_query_port(S.ctx[d], 1, &port) == 0);
+    S.mine.mtu[d] = port.active_mtu;
+    S.mine.recv_rkey[d] = S.recv_mr[d]->rkey;
+    S.mine.flag_rkey[d] = S.flag_mr[d]->rkey;
+  }
+  ibv_free_device_list(list);
+  for (int c = 0; c < (ring ? 2 : world); ++c)
+    for (int r = 0; r < 2; ++r) {
+      if (!ring && c == rank) continue;
+      const int d = chan_dev(c, r);
+      ibv_qp_init_attr ia{};
+      ia.send_cq = S.cq[d]; ia.recv_cq = S.cq[d]; ia.qp_type = IBV_QPT_RC;
+      ia.cap.max_send_wr = 1024; ia.cap.max_recv_wr = 1; ia.cap.max_send_sge = 1; ia.cap.max_recv_sge = 1;
+      ia.cap.max_inline_data = 64;
+      S.qp[c][r] = ibv_create_qp(S.pd[d], &ia); IBCK(S.qp[c][r]);
+      ibv_qp_attr a{};
+      a.qp_state = IBV_QPS_INIT; a.pkey_index = 0; a.port_num = 1; a.qp_access_flags = acc;
+      IBCK(ibv_modify_qp(S.qp[c][r], &a, IBV_QP_STATE | IBV_QP_PKEY_INDEX | IBV_QP_PORT | IBV_QP_ACCESS_FLAGS) == 0);
+      S.mine.qpn[c][r] = S.qp[c][r]->qp_num;
+    }
+  S.mine.recv_addr = (uint64_t)S.recv_h;
+  S.mine.flag_addr = (uint64_t)S.flag_h;
+  return py::bytes(reinterpret_cast<const char*>(&S.mine), sizeof(PeerInfo));
+}
+
+// Takes every rank's PeerInfo in rank order, brings the QPs up, starts the
+// proxy. The caller must barrier its group before the first allreduce.
+void arx_connect(std::vector<std::string> infos) {
+  TORCH_CHECK(S.rank >= 0 && !S.connected && (int)infos.size() == S.world);
+  S.all.resize(S.world);
+  for (int j = 0; j < S.world; ++j) {
+    TORCH_CHECK(infos[j].size() == sizeof(PeerInfo), "arx: PeerInfo size mismatch");
+    memcpy(&S.all[j], infos[j].data(), sizeof(PeerInfo));
+  }
+  for (int c = 0; c < (S.ring ? 2 : S.world); ++c)
+    for (int r = 0; r < 2; ++r) {
+      if (!S.ring && c == S.rank) continue;
+      const PeerInfo& p = S.all[chan_peer(c)];
+      const int d = chan_dev(c, r), rd = chan_rdev(c, r);
+      ibv_qp_attr a{};
+      a.qp_state = IBV_QPS_RTR; a.path_mtu = (ibv_mtu)std::min(S.mine.mtu[d], p.mtu[rd]);
+      a.dest_qp_num = S.ring ? p.qpn[1 - c][r] : p.qpn[S.rank][r]; a.rq_psn = 0;
+      a.max_dest_rd_atomic = 1; a.min_rnr_timer = 12;
+      a.ah_attr.is_global = 1; a.ah_attr.port_num = 1; a.ah_attr.grh.hop_limit = 1;
+      a.ah_attr.grh.sgid_index = S.gid_idx[d];
+      memcpy(a.ah_attr.grh.dgid.raw, p.gid[rd], 16);
+      IBCK(ibv_modify_qp(S.qp[c][r], &a, IBV_QP_STATE | IBV_QP_AV | IBV_QP_PATH_MTU | IBV_QP_DEST_QPN |
+                                             IBV_QP_RQ_PSN | IBV_QP_MAX_DEST_RD_ATOMIC | IBV_QP_MIN_RNR_TIMER) == 0);
+      if (S.ring || c == (S.rank + 1) % S.world)
+        fprintf(stderr, "arx rank %d%s root %d: path MTU %d bytes (ours %d, peer %d)\n", S.rank,
+                S.ring ? (c ? " next" : " prev") : "", r, 128 << a.path_mtu, 128 << S.mine.mtu[d], 128 << p.mtu[rd]);
+      a = {};
+      a.qp_state = IBV_QPS_RTS; a.timeout = 14; a.retry_cnt = 7; a.rnr_retry = 7; a.sq_psn = 0; a.max_rd_atomic = 1;
+      IBCK(ibv_modify_qp(S.qp[c][r], &a, IBV_QP_STATE | IBV_QP_TIMEOUT | IBV_QP_RETRY_CNT | IBV_QP_RNR_RETRY |
+                                             IBV_QP_SQ_PSN | IBV_QP_MAX_QP_RD_ATOMIC) == 0);
+    }
+  Dev& d = S.dev;
+  d.send = (__nv_bfloat16*)S.send_h; d.recv = (const __nv_bfloat16*)S.recv_h; d.flag = S.flag_h; d.ctl = S.ctl;
+  d.rank = S.rank; d.world = S.world;
+  CK(cudaMalloc(&d.seq_dev, 4)); CK(cudaMalloc(&d.blocks, 8)); CK(cudaMalloc(&d.go, 4));
+  CK(cudaMemset(d.seq_dev, 0, 4)); CK(cudaMemset(d.blocks, 0, 8)); CK(cudaMemset(d.go, 0, 4));
+  CK(cudaDeviceSynchronize());
+  std::thread(proxy_loop).detach();
+  S.connected = true;
+}
+
+void arx_allreduce(torch::Tensor in, torch::Tensor out, std::vector<int64_t> pf_ptrs, std::vector<int64_t> pf_bytes) {
+  TORCH_CHECK(S.connected, "arx is not connected");
+  TORCH_CHECK(!g_proxy_err, "arx proxy failed");
+  TORCH_CHECK(in.scalar_type() == at::kBFloat16 && out.scalar_type() == at::kBFloat16 && in.is_contiguous() &&
+              out.is_contiguous() && in.numel() == out.numel());
+  const int n = in.numel();
+  TORCH_CHECK(n % 8 == 0 && (size_t)n * 2 <= kMaxBytes);
+  const int threads = 256, blocks = std::max(1, std::min(8, n / (threads * 8)));
+  Prefetch pf{};
+  pf.count = (int)std::min<size_t>(pf_ptrs.size(), kMaxPrefetch);
+  for (int r = 0; r < pf.count; ++r) {
+    pf.ptr[r] = reinterpret_cast<const uint8_t*>(pf_ptrs[r]);
+    pf.bytes[r] = pf_bytes[r];
+  }
+  arx_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
+      S.dev, (const __nv_bfloat16*)in.data_ptr(), (__nv_bfloat16*)out.data_ptr(), n, pf);
+}
+
+int64_t arx_max_bytes() { return kMaxBytes; }
+
+// Two-shot all-reduce (mesh only): a reduce-scatter then an all-gather. n must
+// split into world chunks of whole 16-byte vectors.
+void arx_allreduce2(torch::Tensor in, torch::Tensor out) {
+  TORCH_CHECK(S.connected, "arx is not connected");
+  TORCH_CHECK(!g_proxy_err, "arx proxy failed");
+  TORCH_CHECK(!S.ring, "arx two-shot is mesh only");
+  TORCH_CHECK(in.scalar_type() == at::kBFloat16 && out.scalar_type() == at::kBFloat16 && in.is_contiguous() &&
+              out.is_contiguous() && in.numel() == out.numel());
+  const int n = in.numel(), chunk = n / S.world;
+  TORCH_CHECK(n % (8 * S.world) == 0 && (size_t)n * 2 <= kSendBytes && (size_t)chunk * 2 <= kSlotBytes);
+  auto mid = torch::empty({chunk}, in.options());
+  const int threads = 256;
+  auto stream = at::cuda::getCurrentCUDAStream();
+  arx_scatter_kernel<<<std::max(1, std::min(8, chunk / (threads * 8))), threads, 0, stream>>>(
+      S.dev, (const __nv_bfloat16*)in.data_ptr(), (__nv_bfloat16*)mid.data_ptr(), n);
+  arx_gather_kernel<<<std::max(1, std::min(8, n / (threads * 8))), threads, 0, stream>>>(
+      S.dev, (const __nv_bfloat16*)mid.data_ptr(), (__nv_bfloat16*)out.data_ptr(), chunk);
+}
+
+int64_t arx_max_bytes2() { return S.ring ? 0 : (int64_t)std::min(kSendBytes, kSlotBytes * S.world); }
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("prepare", &arx_prepare);
+  m.def("connect", &arx_connect);
+  m.def("allreduce", &arx_allreduce);
+  m.def("max_bytes", &arx_max_bytes);
+  m.def("allreduce2", &arx_allreduce2);
+  m.def("max_bytes2", &arx_max_bytes2);
+}

docker-images/glm53-kindling/experimental/arx/arxbig.cu

@@ -0,0 +1,848 @@
+// arxbig: prefill-sized all-gather and reduce-scatter over RoCE for GB10, no NCCL.
+//
+// Same transport as arx_vllm.cu (pinned host buffers are GPU buffers at full
+// speed on GB10, and the ConnectX DMAs them without GPUDirect; a CPU proxy
+// posts RDMA writes when the GPU publishes a sequence number), sized for
+// sequence-parallel prefill: tens of MB per rank per call.
+//
+// all_gather: the GPU copies this rank's slice into its slot of out[seq % 3];
+// the proxy writes that slot into the same place in every peer's out[seq % 3],
+// then a flag; the GPU waits for every peer's flags. The result is a view of
+// out[seq % 3], valid until three more all-gathers have run.
+// reduce_scatter: the GPU copies its full input into send[seq & 1]; the proxy
+// writes chunk j into peer j's recv[seq & 1][rank], then a flag; the GPU waits
+// and sums the world chunks for this rank in rank order (fp32).
+//
+// Every write is split in half, one half per ConnectX root, each followed by
+// its own flag on that root's QP. Rank r sends to r+1, r+2, ... in turn.
+// Reuse is safe for the same reason as in arx: a peer can only reach seq + 2
+// (or + 3) after this rank posted seq + 1, which it does after its stream
+// finished with seq.
+//
+// Ring mode (see arx_vllm.cu): QPs go only to prev and next, one channel
+// each. Data for the rank two away goes half each way round the ring. Of each
+// root's share of a piece, the first half travels next-ward and the second
+// prev-ward. The rank in between forwards each piece as soon as its flag
+// arrives.
+// all_gather: a rank sends its whole slice to prev and to next. The rank
+// forwards the next-ward half of prev's slice to next and the prev-ward half
+// of next's to prev, from and to the same place in out[seq % 3]. Each port
+// carries 1.5 slices out per call. Relaying whole slices one way would put 2
+// on one port.
+// reduce_scatter: chunk r+1 goes to next, chunk r-1 to prev, and chunk r+2
+// half to each. Each neighbour stages its half in its own recv[seq & 1][self],
+// a slot a rank never otherwise uses, and forwards it into recv[seq & 1][r]
+// at r+2. Each port carries 2 slices out, the least possible when relays carry
+// raw data. The two relays through a rank carry opposite halves of each piece,
+// so they share the staging slot.
+// World 2 sends half of everything each way and relays nothing.
+// Ring flags are per (source, arrival channel, root), because data from rank
+// r+2 arrives on both channels. Dev::mask names the flags the GPU waits on.
+// Staged chunks have their own flags for the proxy. A relay takes the piece's
+// bounds from this rank's own publish of the same seq.
+// Reuse with relays: rank m reads out[seq % 3] or its staging slot to forward
+// seq. The sender writes there again at seq + 3 or seq + 2, after finishing
+// seq + 1. Finishing seq + 1 needs data from m's other neighbour. That
+// neighbour published it after finishing seq, and finishing seq needed m's
+// relay. The same chain holds m's own seq + 2 publish, and its new piece
+// bounds, behind the relay. At the final destination the mesh argument holds
+// as is. Relayed data for seq + 2 left its source after the source finished
+// seq + 1, and that needed the destination's seq + 1 data.
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+#include <cuda_runtime.h>
+#include <infiniband/verbs.h>
+
+#include <algorithm>
+#include <atomic>
+#include <chrono>
+#include <cstdint>
+#include <cstdlib>
+#include <cstdio>
+#include <cstring>
+#include <thread>
+#include <vector>
+
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/extension.h>
+
+#define CK(x) do { cudaError_t e_ = (x); if (e_ != cudaSuccess) { fprintf(stderr, "%s:%d %s\n", __FILE__, __LINE__, cudaGetErrorString(e_)); exit(1); } } while (0)
+#define IBCK(x) do { if (!(x)) { fprintf(stderr, "%s:%d %s failed: %s\n", __FILE__, __LINE__, #x, strerror(errno)); exit(1); } } while (0)
+
+namespace {
+
+constexpr int kMaxWorld = 8, kAgSlots = 3, kChunks = 8, kRing = 256;
+constexpr int kMaxDev = 4;  // ring mode: 2 ports x 2 roots
+// Mesh: [op][src][root]. Ring: [op][src][channel][root], then [channel][root]
+// for staged reduce-scatter chunks.
+constexpr int kFlagWords = 2 * kMaxWorld * 4 + 4;
+// kRsPiece: one row range of one destination's chunk, published by a producer
+// kernel as soon as it has written those rows (see finalize_rs_kernel).
+enum Op : uint64_t { kAllGather = 1, kReduceScatter = 2, kRsPiece = 3 };
+
+// Each op goes out as kChunks pieces so the network starts on piece 0 while
+// the GPU still copies piece 1. A peer's flag for op seq, piece c, is
+// seq * 16 + c + 1: monotonic, and it covers every earlier piece on that QP.
+// Producers publish in any order, so each ring entry carries its own index + 1
+// (tag), written last; the proxy consumes entries strictly in index order.
+struct Work {
+  uint64_t op, seq, chunk, slice, lo, hi;
+  volatile uint64_t tag;
+};
+struct Ctl {                    // pinned; written by the GPU, read by the proxy
+  Work ring[kRing];
+  volatile uint64_t stop;
+};
+
+struct PeerInfo {
+  uint32_t qpn[kMaxWorld][2];  // per channel (see State), per root
+  uint8_t gid[kMaxDev][16];
+  uint32_t mtu[kMaxDev];  // per device: the port's active MTU (enum ibv_mtu)
+  uint64_t ag_addr, rs_addr, flag_addr;
+  uint32_t ag_rkey[kMaxDev], rs_rkey[kMaxDev], flag_rkey[kMaxDev];
+};
+
+struct Dev {
+  uint8_t* ag;                    // [kAgSlots][slot_bytes]
+  uint8_t* rs_send;               // [2][slot_bytes]
+  uint8_t* rs_recv;               // [2][world][slot_bytes / world]
+  const volatile uint64_t* flag;  // [2 ops][nslot], see kFlagWords
+  Ctl* ctl;
+  unsigned* blocks;               // [kChunks * kMaxWorld] blocks done with each piece
+  unsigned long long* work_next;  // next ring index to hand out
+  unsigned long long* ready;      // highest flag_of(seq, piece) whose data has fully arrived
+  int rank, world;
+  int nslot;      // flag words per op
+  uint32_t mask;  // the words of an op that carry data for this rank
+  size_t slot_bytes;
+};
+
+__device__ __forceinline__ uint64_t ld_acquire_sys(const volatile uint64_t* p) {
+  uint64_t v;
+  asm volatile("ld.acquire.sys.global.u64 %0, [%1];" : "=l"(v) : "l"(p) : "memory");
+  return v;
+}
+
+__device__ __forceinline__ uint4 ld_cv(const void* p) {
+  uint4 v;
+  asm volatile("ld.global.cv.v4.u32 {%0,%1,%2,%3}, [%4];" : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "l"(p));
+  return v;
+}
+
+#ifndef RS_CACHED
+#define RS_CACHED 1
+#endif
+__device__ __forceinline__ uint4 ld_plain(const void* p) { return *reinterpret_cast<const uint4*>(p); }
+#if RS_CACHED
+#define RS_LOAD ld_plain
+#else
+#define RS_LOAD ld_cv
+#endif
+
+__device__ __forceinline__ uint64_t flag_of(uint64_t seq, int c) { return seq * 16 + c + 1; }
+
+// Piece c of one rank's slice: [lo, hi) bytes, 16-aligned.
+__device__ __forceinline__ void piece(size_t slice, int c, size_t& lo, size_t& hi) {
+  const size_t step = (slice / kChunks + 15) & ~size_t(15);
+  lo = min(slice, step * c);
+  hi = c == kChunks - 1 ? slice : min(slice, step * (c + 1));
+}
+
+__device__ void publish(const Dev& d, uint64_t op, uint64_t seq, uint64_t chunk, uint64_t slice, uint64_t lo,
+                        uint64_t hi) {
+  const unsigned long long idx = atomicAdd(d.work_next, 1ull);
+  Work& it = d.ctl->ring[idx % kRing];
+  it.op = op; it.seq = seq; it.chunk = chunk; it.slice = slice; it.lo = lo; it.hi = hi;
+  __threadfence_system();
+  it.tag = idx + 1;
+}
+
+// Copy piece c of every slice in `slices` (count of them, stride `slice`) from
+// src to dst; the last block to finish a piece publishes it.
+__device__ void copy_publish(const Dev& d, const uint8_t* src, uint8_t* dst, size_t slice, int slices, uint64_t op,
+                             uint64_t seq) {
+  for (int c = 0; c < kChunks; ++c) {
+    size_t lo, hi;
+    piece(slice, c, lo, hi);
+    const size_t len = hi - lo, total = len * slices;
+    const size_t stride = (size_t)gridDim.x * blockDim.x * 16;
+    if (src != dst)
+      for (size_t i = ((size_t)blockIdx.x * blockDim.x + threadIdx.x) * 16; i < total; i += stride) {
+        const size_t s_ = i / len, o = s_ * slice + lo + i % len;
+        *reinterpret_cast<uint4*>(dst + o) = *reinterpret_cast<const uint4*>(src + o);
+      }
+    __threadfence_system();
+    __syncthreads();
+    // One counter per piece: blocks do not wait for each other between pieces.
+    if (threadIdx.x == 0 && atomicAdd(&d.blocks[c], 1) == gridDim.x - 1) {
+      d.blocks[c] = 0;
+      publish(d, op, seq, c, slice, lo, hi);
+    }
+  }
+}
+
+// Wait until every flag in d.mask for op reaches seq, piece c. Only
+// block 0 reads the pinned flags the NIC writes; it relays progress through a
+// device word, since many blocks polling those lines slows the NIC's writes.
+__device__ void wait_piece(const Dev& d, uint64_t op, uint64_t seq, int c) {
+  const uint64_t want = flag_of(seq, c);
+  if (blockIdx.x == 0) {
+    const volatile uint64_t* f = d.flag + (op == kAllGather ? 0 : d.nslot);
+    if (threadIdx.x < d.nslot && (d.mask >> threadIdx.x & 1))
+      while (ld_acquire_sys(f + threadIdx.x) < want) {}
+    __syncthreads();
+    if (threadIdx.x == 0) { __threadfence(); atomicMax(d.ready, (unsigned long long)want); }
+  } else {
+    if (threadIdx.x == 0)
+      while (true) {
+        unsigned long long g;
+        asm volatile("ld.acquire.gpu.global.u64 %0, [%1];" : "=l"(g) : "l"(d.ready) : "memory");
+        if (g >= want) break;
+      }
+    __syncthreads();
+  }
+}
+
+__global__ void ag_kernel(Dev d, const uint8_t* __restrict__ in, size_t slice, uint64_t seq) {
+  uint8_t* out = d.ag + (seq % kAgSlots) * d.slot_bytes;
+  copy_publish(d, in, out + d.rank * slice, slice, 1, kAllGather, seq);
+  wait_piece(d, kAllGather, seq, kChunks - 1);
+}
+
+// out [slice / 2] bf16 = sum over ranks of their chunk `rank` of in.
+__global__ void rs_kernel(Dev d, const uint8_t* __restrict__ in, size_t slice, __nv_bfloat16* __restrict__ out,
+                          uint64_t seq) {
+  const int par = seq & 1;
+  uint8_t* send = d.rs_send + par * d.slot_bytes;
+  copy_publish(d, in, send, slice, d.world, kReduceScatter, seq);
+  const uint8_t* recv = d.rs_recv + (size_t)par * d.slot_bytes;
+  const size_t stride = (size_t)gridDim.x * blockDim.x * 16;
+  for (int c = 0; c < kChunks; ++c) {
+    size_t lo, hi;
+    piece(slice, c, lo, hi);
+    wait_piece(d, kReduceScatter, seq, c);
+    for (size_t i = lo + ((size_t)blockIdx.x * blockDim.x + threadIdx.x) * 16; i < hi; i += stride) {
+      float acc[8] = {};
+      for (int src = 0; src < d.world; ++src) {  // rank order: every rank rounds the same way
+        const uint4 v = src == d.rank ? *reinterpret_cast<const uint4*>(send + (size_t)d.rank * slice + i)
+                                      : RS_LOAD(recv + (size_t)src * slice + i);
+        const __nv_bfloat16* b = reinterpret_cast<const __nv_bfloat16*>(&v);
+#pragma unroll
+        for (int k = 0; k < 8; ++k) acc[k] += __bfloat162float(b[k]);
+      }
+      __align__(16) __nv_bfloat16 o[8];
+#pragma unroll
+      for (int k = 0; k < 8; ++k) o[k] = __float2bfloat16(acc[k]);
+      *reinterpret_cast<uint4*>(reinterpret_cast<uint8_t*>(out) + i) = *reinterpret_cast<const uint4*>(o);
+    }
+  }
+}
+
+
+// Rows [lo, hi) of piece k of a destination's Rr-row chunk.
+__device__ __host__ __forceinline__ void row_piece(int64_t Rr, int k, int64_t& lo, int64_t& hi) {
+  lo = Rr * k / kChunks;
+  hi = Rr * (k + 1) / kChunks;
+}
+
+// MoE output straight into the reduce-scatter send buffer:
+//   row t = shared[t] + sum_k w[t, k] * y[pos[t, k]]  (fp32, rounded once),
+// zero for padding rows t >= T. Blocks take rows piece-major (piece 0 of
+// every destination first) and the last block of a piece publishes it, so
+// the network starts while later rows are still being computed.
+// y8s != nullptr: y holds e4m3 bytes with one fp32 scale per (row, 128 columns).
+__global__ void finalize_rs_kernel(Dev d, const void* __restrict__ yv, const float* __restrict__ y8s,
+                                   const int* __restrict__ pos, const float* __restrict__ w,
+                                   const __nv_bfloat16* __restrict__ shared, int64_t T, int64_t Tpad, int H, int topk,
+                                   uint64_t seq) {
+  const __nv_bfloat16* y = reinterpret_cast<const __nv_bfloat16*>(yv);
+  const uint8_t* y8 = reinterpret_cast<const uint8_t*>(yv);
+  const int64_t Rr = Tpad / d.world;
+  int64_t b = blockIdx.x, lo = 0, hi = 0;
+  int k = 0;
+  for (; k < kChunks; ++k) {
+    row_piece(Rr, k, lo, hi);
+    const int64_t n = (hi - lo) * d.world;
+    if (b < n) break;
+    b -= n;
+  }
+  const int j = (int)(b / (hi - lo));
+  const int64_t t = j * Rr + lo + b % (hi - lo);
+  __nv_bfloat16* dst = reinterpret_cast<__nv_bfloat16*>(d.rs_send + (seq & 1) * d.slot_bytes) + t * H;
+  __shared__ int ps[16];
+  __shared__ float ws[16];
+  if (t < T && threadIdx.x < topk) {
+    ps[threadIdx.x] = pos[t * topk + threadIdx.x];
+    ws[threadIdx.x] = w[t * topk + threadIdx.x];
+  }
+  __syncthreads();
+  for (int c = threadIdx.x * 8; c < H; c += blockDim.x * 8) {
+    float acc[8] = {};
+    if (t < T) {
+      const uint4 sv = *reinterpret_cast<const uint4*>(shared + t * H + c);
+      const __nv_bfloat162* sb = reinterpret_cast<const __nv_bfloat162*>(&sv);
+#pragma unroll
+      for (int i = 0; i < 4; ++i) {
+        const float2 f = __bfloat1622float2(sb[i]);
+        acc[2 * i] = f.x; acc[2 * i + 1] = f.y;
+      }
+      for (int q = 0; q < topk; ++q) {
+        if (y8s != nullptr) {
+          const uint2 v = *reinterpret_cast<const uint2*>(y8 + (size_t)ps[q] * H + c);
+          const float f = ws[q] * y8s[(size_t)ps[q] * (H / 128) + c / 128];
+          const __nv_fp8x4_e4m3* q4 = reinterpret_cast<const __nv_fp8x4_e4m3*>(&v);
+#pragma unroll
+          for (int i = 0; i < 2; ++i) {
+            const float4 x4 = static_cast<float4>(q4[i]);
+            acc[4 * i] += f * x4.x; acc[4 * i + 1] += f * x4.y; acc[4 * i + 2] += f * x4.z; acc[4 * i + 3] += f * x4.w;
+          }
+          continue;
+        }
+        const uint4 v = *reinterpret_cast<const uint4*>(y + (size_t)ps[q] * H + c);
+        const __nv_bfloat162* b2 = reinterpret_cast<const __nv_bfloat162*>(&v);
+#pragma unroll
+        for (int i = 0; i < 4; ++i) {
+          const float2 f = __bfloat1622float2(b2[i]);
+          acc[2 * i] += ws[q] * f.x;
+          acc[2 * i + 1] += ws[q] * f.y;
+        }
+      }
+    }
+    uint4 o;
+    __nv_bfloat162* ob = reinterpret_cast<__nv_bfloat162*>(&o);
+#pragma unroll
+    for (int i = 0; i < 4; ++i) ob[i] = __floats2bfloat162_rn(acc[2 * i], acc[2 * i + 1]);
+    *reinterpret_cast<uint4*>(dst + c) = o;
+  }
+  __threadfence_system();
+  __syncthreads();
+  const int idx = k * d.world + j;
+  if (threadIdx.x == 0 && atomicAdd(&d.blocks[idx], 1) == (unsigned)(hi - lo) - 1) {
+    d.blocks[idx] = 0;
+    const uint64_t row = (uint64_t)H * 2;
+    publish(d, kRsPiece, seq, (uint64_t)j * kChunks + k, Rr * row, lo * row, hi * row);
+  }
+}
+
+// out [Rr, H] = sum over ranks, in rank order, of their rows for this rank,
+// piece by piece as the pieces arrive.
+__global__ void rs_finish_kernel(Dev d, __nv_bfloat16* __restrict__ out, int64_t Rr, int H, uint64_t seq) {
+  const int par = seq & 1;
+  const size_t row = (size_t)H * 2, slice = Rr * row;
+  const uint8_t* send = d.rs_send + par * d.slot_bytes + d.rank * slice;
+  const uint8_t* recv = d.rs_recv + (size_t)par * d.slot_bytes;
+  const size_t stride = (size_t)gridDim.x * blockDim.x * 16;
+  for (int k = 0; k < kChunks; ++k) {
+    int64_t lo, hi;
+    row_piece(Rr, k, lo, hi);
+    wait_piece(d, kReduceScatter, seq, k);
+    for (size_t i = lo * row + ((size_t)blockIdx.x * blockDim.x + threadIdx.x) * 16; i < hi * row; i += stride) {
+      float acc[8] = {};
+      for (int src = 0; src < d.world; ++src) {
+        const uint4 v = src == d.rank ? *reinterpret_cast<const uint4*>(send + i)
+                                      : RS_LOAD(recv + (size_t)src * slice + i);
+        const __nv_bfloat16* b = reinterpret_cast<const __nv_bfloat16*>(&v);
+#pragma unroll
+        for (int q = 0; q < 8; ++q) acc[q] += __bfloat162float(b[q]);
+      }
+      __align__(16) __nv_bfloat16 o[8];
+#pragma unroll
+      for (int q = 0; q < 8; ++q) o[q] = __float2bfloat16(acc[q]);
+      *reinterpret_cast<uint4*>(reinterpret_cast<uint8_t*>(out) + i) = *reinterpret_cast<const uint4*>(o);
+    }
+  }
+}
+
+// A channel is one QP per root to one peer. Mesh: channel j goes to rank j
+// over device r. Ring: channel 0 goes to prev over devices 0 and 1, channel 1
+// to next over devices 2 and 3, and meets the peer's channel on the other side.
+struct State {
+  int rank = -1, world = 0, ndev = 0;
+  bool ring = false;
+  int prev = 0, next = 0, nslot = 0;
+  int gid_idx[kMaxDev] = {};
+  size_t slot_bytes = 0, rs_slot_bytes = 0;
+  uint8_t *ag = nullptr, *rs_send = nullptr, *rs_recv = nullptr;
+  uint64_t* flag = nullptr;
+  Ctl* ctl = nullptr;
+  ibv_context* ctx[kMaxDev] = {};
+  ibv_pd* pd[kMaxDev] = {};
+  ibv_mr *ag_mr[kMaxDev] = {}, *rs_send_mr[kMaxDev] = {}, *rs_recv_mr[kMaxDev] = {}, *flag_mr[kMaxDev] = {};
+  ibv_cq* cq[kMaxDev] = {};
+  ibv_qp* qp[kMaxWorld][2] = {};
+  PeerInfo mine{};
+  std::vector<PeerInfo> all;
+  Dev dev{};
+  bool connected = false;
+};
+State S;
+std::atomic<bool> g_err{false};
+
+int chan_peer(int c) { return S.ring ? (c ? S.next : S.prev) : c; }
+int chan_dev(int c, int r) { return S.ring ? c * 2 + r : r; }         // device here
+int chan_rdev(int c, int r) { return S.ring ? (1 - c) * 2 + r : r; }  // device at the peer
+
+bool post(int c, int r, uint64_t laddr, uint32_t lkey, uint64_t len, uint64_t raddr, uint32_t rkey, uint64_t flag_off,
+          uint64_t* seq_word) {
+  const int j = chan_peer(c);
+  ibv_sge sg{laddr, (uint32_t)len, lkey};
+  ibv_sge fs{(uint64_t)seq_word, 8, 0};
+  ibv_send_wr w{}, f{}, *bad;
+  w.opcode = IBV_WR_RDMA_WRITE; w.sg_list = &sg; w.num_sge = 1;
+  w.wr.rdma.remote_addr = raddr; w.wr.rdma.rkey = rkey;
+  f.opcode = IBV_WR_RDMA_WRITE; f.sg_list = &fs; f.num_sge = 1;
+  f.send_flags = IBV_SEND_INLINE | IBV_SEND_SIGNALED;
+  f.wr_id = (uint64_t)j;
+  f.wr.rdma.remote_addr = S.all[j].flag_addr + flag_off; f.wr.rdma.rkey = S.all[j].flag_rkey[chan_rdev(c, r)];
+  w.next = len ? &f : nullptr;
+  ibv_send_wr* head = len ? &w : &f;
+  if (int e = ibv_post_send(S.qp[c][r], head, &bad)) {
+    fprintf(stderr, "arxbig rank %d: post to %d root %d failed: %s\n", S.rank, j, r, strerror(e));
+    return false;
+  }
+  return true;
+}
+
+// Sends one piece [lo, hi) of op `it` to peer j, split over both roots, each
+// half followed by a flag write of `flag` on that root.
+bool send_piece(const Work& it, int j, uint64_t* flag) {
+  const int rank = S.rank, world = S.world;
+  const uint64_t len = it.hi - it.lo, half = (len / 2 + 15) & ~15ull;
+  for (int r = 0; r < 2; ++r) {
+    const uint64_t off = it.lo + (r ? half : 0), n = r ? len - half : half;
+    bool ok;
+    if (it.op == kAllGather) {
+      const uint64_t at = (it.seq % kAgSlots) * S.slot_bytes + rank * it.slice + off;
+      ok = post(j, r, (uint64_t)S.ag + at, S.ag_mr[r]->lkey, n, S.all[j].ag_addr + at, S.all[j].ag_rkey[r],
+                (rank * 2 + r) * 8, flag);
+    } else {
+      const uint64_t par = it.seq & 1;
+      ok = post(j, r, (uint64_t)S.rs_send + par * S.rs_slot_bytes + j * it.slice + off, S.rs_send_mr[r]->lkey, n,
+                S.all[j].rs_addr + par * S.rs_slot_bytes + rank * it.slice + off, S.all[j].rs_rkey[r],
+                (world * 2 + rank * 2 + r) * 8, flag);
+    }
+    if (!ok) return false;
+  }
+  return true;
+}
+
+// ---- ring mode ----------------------------------------------------------------
+
+// Byte offsets of a ring flag: data from src that arrived on channel ch, and
+// a staged reduce-scatter chunk that arrived on channel ch.
+uint64_t ring_slot(int op, int src, int ch, int r) { return ((uint64_t)op * S.nslot + (src * 2 + ch) * 2 + r) * 8; }
+uint64_t stage_slot(int ch, int r) { return ((uint64_t)2 * S.nslot + ch * 2 + r) * 8; }
+
+// Narrows [lo, hi) to half q. send_piece splits roots at the same point.
+void halve(uint64_t& lo, uint64_t& hi, int q) {
+  const uint64_t h = ((hi - lo) / 2 + 15) & ~15ull;
+  if (q) lo += h; else hi = lo + h;
+}
+
+// Piece [lo, hi) of this rank's all-gather slice, to both neighbours: whole at
+// world 4, half each way at world 2.
+bool ring_send_ag(const Work& it, uint64_t* flag) {
+  const uint64_t at = (it.seq % kAgSlots) * S.slot_bytes + S.rank * it.slice;
+  for (int c = 1; c >= 0; --c)
+    for (int r = 0; r < 2; ++r) {
+      uint64_t lo = it.lo, hi = it.hi;
+      halve(lo, hi, r);
+      if (S.world == 2) halve(lo, hi, 1 - c);
+      const PeerInfo& p = S.all[chan_peer(c)];
+      if (!post(c, r, (uint64_t)S.ag + at + lo, S.ag_mr[chan_dev(c, r)]->lkey, hi - lo, p.ag_addr + at + lo,
+                p.ag_rkey[chan_rdev(c, r)], ring_slot(0, S.rank, 1 - c, r), flag))
+        return false;
+    }
+  return true;
+}
+
+// Piece [lo, hi) of chunk j of this rank's reduce-scatter input. The chunk for
+// the rank two away goes half each way into the neighbours' staging slots.
+bool ring_send_chunk(const Work& it, int j, uint64_t* flag) {
+  const int world = S.world, f = (j - S.rank + world) % world;
+  const bool split = world == 2 || f == 2, stage = world == 4 && f == 2;
+  const uint64_t base = (it.seq & 1) * S.rs_slot_bytes;
+  for (int c = 1; c >= 0; --c) {
+    if (!split && (c == 1) != (f == 1)) continue;
+    const int peer = chan_peer(c);
+    const PeerInfo& p = S.all[peer];
+    for (int r = 0; r < 2; ++r) {
+      uint64_t lo = it.lo, hi = it.hi;
+      halve(lo, hi, r);
+      if (split) halve(lo, hi, 1 - c);
+      if (!post(c, r, (uint64_t)S.rs_send + base + j * it.slice + lo, S.rs_send_mr[chan_dev(c, r)]->lkey, hi - lo,
+                p.rs_addr + base + (stage ? peer : S.rank) * it.slice + lo, p.rs_rkey[chan_rdev(c, r)],
+                stage ? stage_slot(1 - c, r) : ring_slot(1, S.rank, 1 - c, r), flag))
+        return false;
+    }
+  }
+  return true;
+}
+
+// The proxy's relay state at world 4. Every rank runs the same
+// collectives in the same order, so this rank's own publishes say which seqs
+// are all-gathers or reduce-scatters and where each piece starts and ends.
+constexpr int kQ = 8;
+struct Bound { uint64_t seq, slice, lo, hi; };
+struct Relays {
+  Bound bnd[2][kChunks] = {};              // by seq parity, piece
+  uint64_t seqs[2][kQ] = {}, nseq[2] = {};  // per kind (0 all-gather, 1 reduce-scatter), in order
+  struct { uint64_t i; int c; } at[2][2][2] = {};  // per kind, out channel, root: next seq index and piece
+};
+
+bool relay_note(Relays& R, const Work& it) {
+  const int kind = it.op == kAllGather ? 0 : 1;
+  const int c = (int)(it.op == kRsPiece ? it.chunk % kChunks : it.chunk);
+  if (R.nseq[kind] == 0 || R.seqs[kind][(R.nseq[kind] - 1) % kQ] != it.seq) {
+    for (int o = 0; o < 2; ++o)
+      for (int r = 0; r < 2; ++r)
+        if (R.nseq[kind] - R.at[kind][o][r].i >= kQ) {
+          fprintf(stderr, "arxbig rank %d: relays fell %d collectives behind\n", S.rank, kQ);
+          return false;
+        }
+    R.seqs[kind][R.nseq[kind]++ % kQ] = it.seq;
+  }
+  R.bnd[it.seq & 1][c] = {it.seq, it.slice, it.lo, it.hi};
+  return true;
+}
+
+
+// Idle backoff (#55): keep the exact busy spin while there is work and for ARX_IDLE_SPIN_MS after it (default
+// 200 ms), then nap ARX_IDLE_NAP_US (default 200 us) per loop turn. ARX_IDLE_NAP_US=0 keeps the old spin.
+struct IdleNap {
+  std::chrono::steady_clock::duration spin;
+  std::chrono::microseconds nap;
+  std::chrono::steady_clock::time_point last;
+  IdleNap() {
+    const char* s = std::getenv("ARX_IDLE_SPIN_MS");
+    const char* n = std::getenv("ARX_IDLE_NAP_US");
+    spin = std::chrono::milliseconds(s ? std::atol(s) : 200);
+    nap = std::chrono::microseconds(n ? std::atol(n) : 200);
+    last = std::chrono::steady_clock::now();
+    fprintf(stderr, "arx idle nap: spin %ld ms, nap %ld us%s\n", (long)(s ? std::atol(s) : 200), (long)nap.count(),
+            nap.count() <= 0 ? " (disabled = old spin)" : "");
+  }
+  void busy() { last = std::chrono::steady_clock::now(); }
+  void idle() {
+    if (nap.count() <= 0) return;
+    if (std::chrono::steady_clock::now() - last > spin) std::this_thread::sleep_for(nap);
+  }
+};
+
+bool relay_pending(const Relays& R) {  // a relay this rank noted but has not forwarded yet
+  for (int kind = 0; kind < 2; ++kind)
+    for (int o = 0; o < 2; ++o)
+      for (int r = 0; r < 2; ++r)
+        if (R.at[kind][o][r].i < R.nseq[kind]) return true;
+  return false;
+}
+
+// Forwards every piece whose flag has arrived, in order per stream. Out
+// channel 1 carries prev's data on to next (first halves), channel 0 next's
+// on to prev (second halves).
+bool relay_step(Relays& R) {
+  const volatile uint64_t* flags = S.flag;
+  for (int kind = 0; kind < 2; ++kind)
+    for (int o = 0; o < 2; ++o)
+      for (int r = 0; r < 2; ++r) {
+        auto& st = R.at[kind][o][r];
+        const int src = o ? S.prev : S.next, peer = chan_peer(o);
+        const uint64_t in = (kind ? stage_slot(1 - o, r) : ring_slot(0, src, 1 - o, r)) / 8;
+        while (st.i < R.nseq[kind]) {
+          const uint64_t s = R.seqs[kind][st.i % kQ];
+          const Bound& B = R.bnd[s & 1][st.c];
+          uint64_t flag = s * 16 + st.c + 1;
+          if (B.seq != s || flags[in] < flag) break;
+          std::atomic_thread_fence(std::memory_order_acquire);
+          uint64_t lo = B.lo, hi = B.hi;
+          halve(lo, hi, r);
+          halve(lo, hi, 1 - o);
+          const PeerInfo& p = S.all[peer];
+          const int d = chan_dev(o, r), rd = chan_rdev(o, r);
+          bool ok;
+          if (kind == 0) {
+            const uint64_t at = (s % kAgSlots) * S.slot_bytes + src * B.slice;
+            ok = post(o, r, (uint64_t)S.ag + at + lo, S.ag_mr[d]->lkey, hi - lo, p.ag_addr + at + lo, p.ag_rkey[rd],
+                      ring_slot(0, src, 1 - o, r), &flag);
+          } else {
+            const uint64_t base = (s & 1) * S.rs_slot_bytes;
+            ok = post(o, r, (uint64_t)S.rs_recv + base + S.rank * B.slice + lo, S.rs_recv_mr[d]->lkey, hi - lo,
+                      p.rs_addr + base + src * B.slice + lo, p.rs_rkey[rd], ring_slot(1, src, 1 - o, r), &flag);
+          }
+          if (!ok) return false;
+          if (++st.c == kChunks) { st.c = 0; ++st.i; }
+        }
+      }
+  return true;
+}
+
+void proxy_loop() {
+  uint64_t done = 0;
+  static Work pend[kMaxWorld][kChunks];  // kRsPiece items that arrived early
+  uint32_t pend_mask[kMaxWorld] = {}, next_k[kMaxWorld] = {};
+  uint64_t cur_seq[kMaxWorld] = {};
+  const int rank = S.rank, world = S.world;
+  const bool relay = S.ring && world == 4;
+  static Relays R;
+  auto send = S.ring ? ring_send_chunk : send_piece;
+  IdleNap nap;
+  while (!S.ctl->stop) {
+    bool moved = false;
+    for (int d = 0; d < S.ndev; ++d) {
+      ibv_wc wc[32];
+      const int n = ibv_poll_cq(S.cq[d], 32, wc);
+      if (n > 0) moved = true;
+      for (int i = 0; i < n; ++i)
+        if (wc[i].status != IBV_WC_SUCCESS) {
+          fprintf(stderr, "arxbig rank %d: completion error %d\n", rank, wc[i].status);
+          g_err = true;
+          return;
+        }
+    }
+    if (relay && !relay_step(R)) { g_err = true; return; }
+    Work& slot = S.ctl->ring[done % kRing];
+    if (slot.tag != done + 1) {
+      if (moved || (relay && relay_pending(R))) nap.busy(); else nap.idle();
+      continue;
+    }
+    nap.busy();
+    std::atomic_thread_fence(std::memory_order_acquire);
+    Work it;
+    it.op = slot.op; it.seq = slot.seq; it.chunk = slot.chunk; it.slice = slot.slice; it.lo = slot.lo; it.hi = slot.hi;
+    ++done;
+    if (relay && !relay_note(R, it)) { g_err = true; return; }
+    if (it.op == kAllGather && S.ring) {
+      uint64_t flag = it.seq * 16 + it.chunk + 1;
+      if (!ring_send_ag(it, &flag)) { g_err = true; return; }
+      continue;
+    }
+    if (it.op != kRsPiece) {  // one piece for every peer
+      uint64_t flag = it.seq * 16 + it.chunk + 1;  // inline: copied at post time
+      for (int k = 1; k < world; ++k)
+        if (!send(it, (rank + k) % world, &flag)) { g_err = true; return; }
+      continue;
+    }
+    // Row piece `sub` of destination j's chunk. Peers read the flag as "all
+    // pieces up to sub arrived", so pieces go out in order per destination.
+    const int j = (int)(it.chunk / kChunks), sub = (int)(it.chunk % kChunks);
+    if (j == rank) continue;  // this rank's own rows never leave
+    if (cur_seq[j] != it.seq) { cur_seq[j] = it.seq; next_k[j] = 0; pend_mask[j] = 0; }
+    pend[j][sub] = it;
+    pend_mask[j] |= 1u << sub;
+    while (next_k[j] < kChunks && (pend_mask[j] >> next_k[j] & 1)) {
+      const int k = next_k[j]++;
+      uint64_t flag = it.seq * 16 + k + 1;
+      if (!send(pend[j][k], j, &flag)) { g_err = true; return; }
+    }
+  }
+}
+
+Dev make_dev() {
+  Dev d{};
+  d.ag = S.ag; d.rs_send = S.rs_send; d.rs_recv = S.rs_recv; d.flag = S.flag; d.ctl = S.ctl;
+  CK(cudaMalloc(&d.blocks, kChunks * kMaxWorld * 4)); CK(cudaMalloc(&d.ready, 8)); CK(cudaMalloc(&d.work_next, 8));
+  CK(cudaMemset(d.blocks, 0, kChunks * kMaxWorld * 4)); CK(cudaMemset(d.ready, 0, 8)); CK(cudaMemset(d.work_next, 0, 8));
+  d.rank = S.rank; d.world = S.world; d.slot_bytes = S.slot_bytes;
+  d.nslot = S.nslot;
+  for (int src = 0; src < S.world; ++src) {
+    if (src == S.rank) continue;
+    if (!S.ring) { d.mask |= 3u << (src * 2); continue; }
+    // f hops from src to this rank going next-ward, which arrives on channel 0.
+    const int f = (S.rank - src + S.world) % S.world;
+    if (2 * f <= S.world) d.mask |= 3u << ((src * 2 + 0) * 2);
+    if (2 * f >= S.world) d.mask |= 3u << ((src * 2 + 1) * 2);
+  }
+  return d;
+}
+
+}  // namespace
+
+// Mesh: devs are root 0 and root 1. Ring: the port facing prev (root 0, root
+// 1), then the port facing next. gids: each device's GID index.
+// rs_slot_bytes: 0 leaves reduce_scatter unavailable and saves its buffers.
+py::bytes prepare(int64_t rank, int64_t world, std::vector<std::string> devs, std::vector<int64_t> gids, bool ring,
+                  int64_t slot_bytes, int64_t rs_slot_bytes) {
+  TORCH_CHECK(S.rank < 0 && world >= 2 && world <= kMaxWorld && slot_bytes % (world * 16) == 0);
+  TORCH_CHECK(rs_slot_bytes == 0 || rs_slot_bytes == slot_bytes, "arxbig: reduce_scatter uses the same slot size");
+  TORCH_CHECK(!ring || world == 2 || world == 4, "arxbig ring mode supports world 2 or 4, not ", world,
+              ": the ring reduce-scatter sends the rank two away half each way. Set VLLM_ARXBIG=0 to run arx alone");
+  S.ndev = ring ? 4 : 2;
+  TORCH_CHECK((int)devs.size() == S.ndev && (int)gids.size() == S.ndev, "arxbig: expected ", S.ndev, " devices");
+  S.rank = rank; S.world = world; S.ring = ring; S.slot_bytes = slot_bytes; S.rs_slot_bytes = rs_slot_bytes;
+  S.prev = (rank + world - 1) % world; S.next = (rank + 1) % world;
+  S.nslot = (ring ? 4 : 2) * world;
+  const size_t rs_alloc = rs_slot_bytes ? 2 * rs_slot_bytes : 4096;
+  CK(cudaHostAlloc(&S.ag, kAgSlots * slot_bytes, cudaHostAllocMapped));
+  CK(cudaHostAlloc(&S.rs_send, rs_alloc, cudaHostAllocMapped));
+  CK(cudaHostAlloc(&S.rs_recv, rs_alloc, cudaHostAllocMapped));
+  CK(cudaHostAlloc(&S.flag, kFlagWords * sizeof(uint64_t), cudaHostAllocMapped));
+  CK(cudaHostAlloc(&S.ctl, sizeof(Ctl), cudaHostAllocMapped));
+  memset(S.flag, 0, kFlagWords * sizeof(uint64_t));
+  memset((void*)S.ctl, 0, sizeof(Ctl));
+  int ndev;
+  ibv_device** list = ibv_get_device_list(&ndev);
+  const int acc = IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_WRITE;
+  for (int d = 0; d < S.ndev; ++d) {
+    S.gid_idx[d] = gids[d];
+    for (int i = 0; i < ndev; ++i)
+      if (devs[d] == ibv_get_device_name(list[i])) S.ctx[d] = ibv_open_device(list[i]);
+    TORCH_CHECK(S.ctx[d], "arxbig: no RDMA device ", devs[d]);
+    S.pd[d] = ibv_alloc_pd(S.ctx[d]); IBCK(S.pd[d]);
+    S.ag_mr[d] = ibv_reg_mr(S.pd[d], S.ag, kAgSlots * slot_bytes, acc); IBCK(S.ag_mr[d]);
+    S.rs_send_mr[d] = ibv_reg_mr(S.pd[d], S.rs_send, rs_alloc, acc); IBCK(S.rs_send_mr[d]);
+    S.rs_recv_mr[d] = ibv_reg_mr(S.pd[d], S.rs_recv, rs_alloc, acc); IBCK(S.rs_recv_mr[d]);
+    S.flag_mr[d] = ibv_reg_mr(S.pd[d], S.flag, kFlagWords * sizeof(uint64_t), acc); IBCK(S.flag_mr[d]);
+    S.cq[d] = ibv_create_cq(S.ctx[d], 4096, nullptr, nullptr, 0); IBCK(S.cq[d]);
+    ibv_gid gid;
+    IBCK(ibv_query_gid(S.ctx[d], 1, S.gid_idx[d], &gid) == 0);
+    memcpy(S.mine.gid[d], gid.raw, 16);
+    ibv_port_attr port{};
+    IBCK(ibv_query_port(S.ctx[d], 1, &port) == 0);
+    S.mine.mtu[d] = port.active_mtu;
+    S.mine.ag_rkey[d] = S.ag_mr[d]->rkey;
+    S.mine.rs_rkey[d] = S.rs_recv_mr[d]->rkey;
+    S.mine.flag_rkey[d] = S.flag_mr[d]->rkey;
+  }
+  ibv_free_device_list(list);
+  for (int c = 0; c < (ring ? 2 : world); ++c)
+    for (int r = 0; r < 2; ++r) {
+      if (!ring && c == rank) continue;
+      const int d = chan_dev(c, r);
+      ibv_qp_init_attr ia{};
+      ia.send_cq = S.cq[d]; ia.recv_cq = S.cq[d]; ia.qp_type = IBV_QPT_RC;
+      ia.cap.max_send_wr = 1024; ia.cap.max_recv_wr = 1; ia.cap.max_send_sge = 1; ia.cap.max_recv_sge = 1;
+      ia.cap.max_inline_data = 64;
+      S.qp[c][r] = ibv_create_qp(S.pd[d], &ia); IBCK(S.qp[c][r]);
+      ibv_qp_attr a{};
+      a.qp_state = IBV_QPS_INIT; a.pkey_index = 0; a.port_num = 1; a.qp_access_flags = acc;
+      IBCK(ibv_modify_qp(S.qp[c][r], &a, IBV_QP_STATE | IBV_QP_PKEY_INDEX | IBV_QP_PORT | IBV_QP_ACCESS_FLAGS) == 0);
+      S.mine.qpn[c][r] = S.qp[c][r]->qp_num;
+    }
+  S.mine.ag_addr = (uint64_t)S.ag;
+  S.mine.rs_addr = (uint64_t)S.rs_recv;
+  S.mine.flag_addr = (uint64_t)S.flag;
+  return py::bytes(reinterpret_cast<const char*>(&S.mine), sizeof(PeerInfo));
+}
+
+void connect(std::vector<std::string> infos) {
+  TORCH_CHECK(S.rank >= 0 && !S.connected && (int)infos.size() == S.world);
+  S.all.resize(S.world);
+  for (int j = 0; j < S.world; ++j) {
+    TORCH_CHECK(infos[j].size() == sizeof(PeerInfo), "arxbig: PeerInfo size mismatch");
+    memcpy(&S.all[j], infos[j].data(), sizeof(PeerInfo));
+  }
+  for (int c = 0; c < (S.ring ? 2 : S.world); ++c)
+    for (int r = 0; r < 2; ++r) {
+      if (!S.ring && c == S.rank) continue;
+      const PeerInfo& p = S.all[chan_peer(c)];
+      const int d = chan_dev(c, r), rd = chan_rdev(c, r);
+      ibv_qp_attr a{};
+      a.qp_state = IBV_QPS_RTR; a.path_mtu = (ibv_mtu)std::min(S.mine.mtu[d], p.mtu[rd]);
+      a.dest_qp_num = S.ring ? p.qpn[1 - c][r] : p.qpn[S.rank][r]; a.rq_psn = 0;
+      a.max_dest_rd_atomic = 1; a.min_rnr_timer = 12;
+      a.ah_attr.is_global = 1; a.ah_attr.port_num = 1; a.ah_attr.grh.hop_limit = 1;
+      a.ah_attr.grh.sgid_index = S.gid_idx[d];
+      memcpy(a.ah_attr.grh.dgid.raw, p.gid[rd], 16);
+      IBCK(ibv_modify_qp(S.qp[c][r], &a, IBV_QP_STATE | IBV_QP_AV | IBV_QP_PATH_MTU | IBV_QP_DEST_QPN |
+                                              IBV_QP_RQ_PSN | IBV_QP_MAX_DEST_RD_ATOMIC | IBV_QP_MIN_RNR_TIMER) == 0);
+      if (S.ring || c == (S.rank + 1) % S.world)
+        fprintf(stderr, "arxbig rank %d%s root %d: path MTU %d bytes (ours %d, peer %d)\n", S.rank,
+                S.ring ? (c ? " next" : " prev") : "", r, 128 << a.path_mtu, 128 << S.mine.mtu[d], 128 << p.mtu[rd]);
+      a.qp_state = IBV_QPS_RTS; a.timeout = 14; a.retry_cnt = 7; a.rnr_retry = 7; a.sq_psn = 0; a.max_rd_atomic = 1;
+      IBCK(ibv_modify_qp(S.qp[c][r], &a, IBV_QP_STATE | IBV_QP_TIMEOUT | IBV_QP_RETRY_CNT | IBV_QP_RNR_RETRY |
+                                              IBV_QP_SQ_PSN | IBV_QP_MAX_QP_RD_ATOMIC) == 0);
+    }
+  S.dev = make_dev();
+  CK(cudaDeviceSynchronize());
+  std::thread(proxy_loop).detach();
+  S.connected = true;
+}
+
+static int blocks_for(size_t bytes) { return (int)std::max<size_t>(1, std::min<size_t>(48, bytes / (256 * 16 * 8))); }
+static uint64_t g_seq = 0;  // collectives run in stream order, so the host numbers them
+
+// in: this rank's slice (any dtype, contiguous). Returns the gathered data as a
+// view of the pinned output slot, shaped [world * in.size(0), ...]; it stays
+// valid until three more all-gathers have run.
+torch::Tensor all_gather(torch::Tensor in) {
+  TORCH_CHECK(S.connected && !g_err, "arxbig is not connected or its proxy failed");
+  TORCH_CHECK(in.is_contiguous() && in.is_cuda() && in.dim() >= 1);
+  const size_t slice = in.numel() * in.element_size();
+  TORCH_CHECK(slice % 16 == 0 && slice * S.world <= S.slot_bytes, "arxbig: all_gather slice size");
+  const uint64_t seq = ++g_seq;
+  ag_kernel<<<blocks_for(slice), 256, 0, at::cuda::getCurrentCUDAStream()>>>(S.dev, (const uint8_t*)in.data_ptr(),
+                                                                              slice, seq);
+  std::vector<int64_t> shape(in.sizes().begin(), in.sizes().end());
+  shape[0] *= S.world;
+  return torch::from_blob(S.ag + (seq % kAgSlots) * S.slot_bytes, shape, in.options());
+}
+
+// in: bf16 [world * m, ...] contiguous. Returns bf16 [m, ...]: the sum over
+// ranks of each rank's chunk `rank`.
+torch::Tensor reduce_scatter(torch::Tensor in) {
+  TORCH_CHECK(S.connected && !g_err, "arxbig is not connected or its proxy failed");
+  TORCH_CHECK(S.rs_slot_bytes, "arxbig: prepared without reduce_scatter buffers");
+  TORCH_CHECK(in.is_contiguous() && in.is_cuda() && in.scalar_type() == at::kBFloat16 && in.size(0) % S.world == 0);
+  const size_t slice = in.numel() * 2 / S.world;
+  TORCH_CHECK(slice % 16 == 0 && slice * S.world <= S.slot_bytes, "arxbig: reduce_scatter size");
+  std::vector<int64_t> shape(in.sizes().begin(), in.sizes().end());
+  shape[0] /= S.world;
+  auto out = torch::empty(shape, in.options());
+  const uint64_t seq = ++g_seq;
+  rs_kernel<<<blocks_for(slice * S.world), 256, 0, at::cuda::getCurrentCUDAStream()>>>(
+      S.dev, (const uint8_t*)in.data_ptr(), slice, (__nv_bfloat16*)out.data_ptr(), seq);
+  return out;
+}
+
+// The pinned buffer the next reduce_scatter sends from, as a bf16 tensor of
+// `shape`. A producer that writes its output here saves reduce_scatter's copy;
+// it stays valid until that reduce_scatter.
+torch::Tensor rs_input(std::vector<int64_t> shape) {
+  TORCH_CHECK(S.connected && S.rs_slot_bytes);
+  int64_t n = 1;
+  for (auto v : shape) n *= v;
+  TORCH_CHECK(n * 2 <= (int64_t)S.slot_bytes, "arxbig: rs_input too large");
+  const uint64_t par = (g_seq + 1) & 1;
+  return torch::from_blob(S.rs_send + par * S.slot_bytes, shape,
+                          torch::TensorOptions().dtype(at::kBFloat16).device(torch::kCUDA));
+}
+
+// The MoE's final combine, reduce-scattered: finalize_rs_kernel writes and
+// publishes rows; rs_finish (called next on the same stream) returns this
+// rank's [Tpad / world, H] sum. y bf16 [R, H], pos int32 [T * topk],
+// w fp32 [T, topk] (routing weights, scale folded in), shared bf16 [T, H].
+int64_t moe_finalize_rs(torch::Tensor y, torch::Tensor pos, torch::Tensor w, torch::Tensor shared, int64_t T,
+                        int64_t Tpad, c10::optional<torch::Tensor> y8s) {
+  TORCH_CHECK(S.connected && !g_err && S.rs_slot_bytes, "arxbig: reduce_scatter unavailable");
+  const int H = y.size(1), topk = w.size(1);
+  TORCH_CHECK(Tpad % S.world == 0 && Tpad / S.world >= kChunks && T <= Tpad && topk <= 16 && H % 8 == 0);
+  TORCH_CHECK((size_t)Tpad * H * 2 <= S.slot_bytes, "arxbig: MoE output larger than a slot");
+  TORCH_CHECK(y.is_contiguous() && pos.is_contiguous() && w.is_contiguous() && shared.is_contiguous());
+  TORCH_CHECK(!y8s.has_value() || (y.scalar_type() == at::kByte && H % 128 == 0), "arxbig: y8s needs e4m3 y");
+  const uint64_t seq = ++g_seq;
+  finalize_rs_kernel<<<Tpad, 128, 0, at::cuda::getCurrentCUDAStream()>>>(
+      S.dev, y.data_ptr(), y8s.has_value() ? y8s->data_ptr<float>() : nullptr, pos.data_ptr<int>(),
+      w.data_ptr<float>(), (const __nv_bfloat16*)shared.data_ptr(), T, Tpad, H, topk, seq);
+  return (int64_t)seq;
+}
+
+torch::Tensor rs_finish(int64_t seq, int64_t rows, int64_t H) {
+  TORCH_CHECK((uint64_t)seq == g_seq, "arxbig: rs_finish must follow its producer directly");
+  auto out = torch::empty({rows, H}, torch::TensorOptions().dtype(at::kBFloat16).device(torch::kCUDA));
+  rs_finish_kernel<<<blocks_for(rows * H * 2 * S.world), 256, 0, at::cuda::getCurrentCUDAStream()>>>(
+      S.dev, (__nv_bfloat16*)out.data_ptr(), rows, (int)H, (uint64_t)seq);
+  return out;
+}
+
+bool failed() { return g_err; }
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("prepare", &prepare);
+  m.def("connect", &connect);
+  m.def("all_gather", &all_gather);
+  m.def("reduce_scatter", &reduce_scatter);
+  m.def("rs_input", &rs_input);
+  m.def("moe_finalize_rs", &moe_finalize_rs, py::arg("y"), py::arg("pos"), py::arg("w"), py::arg("shared"), py::arg("T"), py::arg("Tpad"), py::arg("y8s") = py::none());
+  m.def("rs_finish", &rs_finish);
+  m.def("failed", &failed);
+}

docker-images/glm53-kindling/experimental/arx/cuda_communicator.py

@@ -0,0 +1,970 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+import os
+
+import torch
+from torch.distributed import ProcessGroup
+
+import vllm.envs as envs
+from vllm._aiter_ops import rocm_aiter_ops
+from vllm.distributed.device_communicators.all_reduce_utils import (
+    NCCL_SYMM_MEM_ALL_REDUCE_CONFIG,
+    should_nccl_symm_mem_ag_rs,
+    should_nccl_symm_mem_allreduce,
+)
+from vllm.distributed.device_communicators.pynccl import register_nccl_symmetric_ops
+from vllm.distributed.device_communicators.pynccl_allocator import (
+    is_symmetric_memory_enabled,
+)
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+
+from ..utils import StatelessProcessGroup
+from .aiter_custom_all_reduce import AiterCustomAllreduce
+from .base_device_communicator import DeviceCommunicatorBase
+
+logger = init_logger(__name__)
+
+
+class CudaCommunicator(DeviceCommunicatorBase):
+    def __init__(
+        self,
+        cpu_group: ProcessGroup,
+        device: torch.device | None = None,
+        device_group: ProcessGroup | None = None,
+        unique_name: str = "",
+        global_ranks: list[int] | None = None,
+        global_world_size: int | None = None,
+        tcp_store_group: StatelessProcessGroup | None = None,
+        use_all2all: bool = False,
+    ):
+        super().__init__(
+            cpu_group,
+            device,
+            device_group,
+            unique_name,
+            global_ranks,
+            global_world_size,
+            use_all2all=use_all2all,
+        )
+        # Match the group name exactly so ETP does not enable TP-only backends.
+        if unique_name.split(":")[0] != "tp":
+            # custom allreduce or torch symm mem can be used only by tp
+            use_custom_allreduce = False
+            use_torch_symm_mem = False
+            use_flashinfer_allreduce = False
+            use_flashinfer_pcie_ipc_allreduce = False
+            use_aiter_allreduce = False
+        else:
+            from vllm.distributed.parallel_state import _ENABLE_CUSTOM_ALL_REDUCE
+
+            use_custom_allreduce = _ENABLE_CUSTOM_ALL_REDUCE
+            use_torch_symm_mem = envs.VLLM_ALLREDUCE_USE_SYMM_MEM
+            # FlashInfer all-reduce does not provide a fixed reduction order.
+            use_flashinfer_allreduce = (
+                envs.VLLM_ALLREDUCE_USE_FLASHINFER and not envs.VLLM_BATCH_INVARIANT
+            )
+            use_flashinfer_pcie_ipc_allreduce = (
+                envs.VLLM_ALLREDUCE_USE_FLASHINFER_PCIE_IPC
+                and not envs.VLLM_BATCH_INVARIANT
+            )
+            use_aiter_allreduce = use_custom_allreduce and bool(
+                rocm_aiter_ops.is_custom_all_reduce_enabled()
+            )
+
+        self.use_custom_allreduce = use_custom_allreduce
+        self.use_torch_symm_mem = use_torch_symm_mem
+        self.use_flashinfer_allreduce = use_flashinfer_allreduce
+        self.use_flashinfer_pcie_ipc_allreduce = use_flashinfer_pcie_ipc_allreduce
+        self.use_aiter_allreduce = use_aiter_allreduce
+
+        # lazy import to avoid documentation build error
+        from vllm.distributed.device_communicators.custom_all_reduce import (
+            CustomAllreduce,
+        )
+        from vllm.distributed.device_communicators.flashinfer_all_reduce import (
+            FlashInferAllReduce,
+        )
+        from vllm.distributed.device_communicators.flashinfer_pcie_ipc_all_reduce import (  # noqa: E501
+            FlashInferPcieIpcAllReduce,
+        )
+        from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
+        from vllm.distributed.device_communicators.quick_all_reduce import (
+            QuickAllReduce,
+        )
+        from vllm.distributed.device_communicators.symm_mem import SymmMemCommunicator
+
+        self.pynccl_comm: PyNcclCommunicator | None = None
+        if self.world_size > 1:
+            self.pynccl_comm = PyNcclCommunicator(
+                group=self.cpu_group if tcp_store_group is None else tcp_store_group,
+                device=self.device,
+            )
+            if is_symmetric_memory_enabled():
+                register_nccl_symmetric_ops(self.pynccl_comm)
+
+        # GB10: RDMA all-reduce for decode-sized tensors (arx.py).
+        self.arx_comm = None
+        if (
+            unique_name.split(":")[0] == "tp"
+            and self.world_size > 1
+            and os.environ.get("VLLM_ARX_ALLREDUCE") == "1"
+        ):
+            from vllm.distributed.device_communicators.arx import ArxCommunicator
+
+            self.arx_comm = ArxCommunicator(self.cpu_group, self.device)
+        self.arxbig = None
+        if self.arx_comm is not None and os.environ.get("VLLM_ARXBIG") == "1":
+            from vllm.distributed.device_communicators.arx import ArxBig
+
+            self.arxbig = ArxBig(self.cpu_group, self.device)
+
+        self.ca_comm: CustomAllreduce | None = None
+        self.qr_comm: QuickAllReduce | None = None
+        self.symm_mem_comm: SymmMemCommunicator | None = None
+        self.fi_ar_comm: FlashInferAllReduce | None = None
+        self.fi_pcie_ipc_ar_comm: FlashInferPcieIpcAllReduce | None = None
+        self.aiter_ar_comm: AiterCustomAllreduce | None = None
+        self.use_aiter_ag_rs: bool = False
+
+        if use_torch_symm_mem and current_platform.is_cuda():
+            self.symm_mem_comm = SymmMemCommunicator(
+                group=self.cpu_group,
+                device=self.device,
+            )
+
+        if self.use_flashinfer_allreduce and self.world_size > 1:
+            self.fi_ar_comm = FlashInferAllReduce(
+                group=self.cpu_group,
+                device=self.device,
+            )
+
+        if (
+            self.use_flashinfer_pcie_ipc_allreduce
+            and self.world_size > 1
+            and self.device_group is not None
+        ):
+            self.fi_pcie_ipc_ar_comm = FlashInferPcieIpcAllReduce(
+                group=self.device_group,
+                tune_group=self.cpu_group,
+                device=self.device,
+            )
+
+        if self.use_aiter_allreduce and self.world_size > 1:
+            self.aiter_ar_comm = AiterCustomAllreduce(
+                group=self.cpu_group,
+                device=self.device,
+            )
+
+        if use_custom_allreduce and self.aiter_ar_comm is None and self.world_size > 1:
+            # Initialize a custom fast all-reduce implementation.
+            self.ca_comm = CustomAllreduce(
+                group=self.cpu_group,
+                device=self.device,
+                symm_mem_enabled=(
+                    self.symm_mem_comm is not None and not self.symm_mem_comm.disabled
+                ),
+            )
+
+        # AITER custom all-gather/reduce-scatter DP-attention dispatch/combine
+        if (
+            "dp" in unique_name
+            and self.world_size in (2, 4, 8)
+            and current_platform.is_rocm()
+            and rocm_aiter_ops.is_custom_all_reduce_enabled()
+        ):
+            self.aiter_ar_comm = AiterCustomAllreduce(
+                group=self.cpu_group,
+                device=self.device,
+            )
+            if self.aiter_ar_comm.disabled:
+                self.aiter_ar_comm = None
+            else:
+                self.use_aiter_ag_rs = True
+
+        if use_custom_allreduce and self.world_size > 1 and current_platform.is_rocm():
+            # Initialize a custom quick all-reduce implementation for AMD.
+            # Quick reduce is designed as a complement to custom allreduce
+            # (vLLM's or AITER's), so it is initialized for either backend.
+            # Based on quickreduce (https://github.com/mk1-project/quickreduce).
+            # On ROCm, 'use_custom_allreduce==True' means it must currently be
+            # an MI300 series.
+            self.qr_comm = QuickAllReduce(group=self.cpu_group, device=self.device)
+
+        if self.world_size > 1:
+            self._log_all_reduce_backend_selection()
+
+        if self.use_all2all:
+            if self.all2all_backend in ("naive", "allgather_reducescatter"):
+                from .all2all import AgRsAll2AllManager
+
+                self.all2all_manager = AgRsAll2AllManager(
+                    self.cpu_group, tcp_store_group
+                )
+            elif self.all2all_backend == "deepep_high_throughput":
+                from .all2all import DeepEPHTAll2AllManager
+
+                self.all2all_manager = DeepEPHTAll2AllManager(
+                    self.cpu_group, tcp_store_group
+                )
+            elif self.all2all_backend == "deepep_low_latency":
+                from .all2all import DeepEPLLAll2AllManager
+
+                self.all2all_manager = DeepEPLLAll2AllManager(
+                    self.cpu_group, tcp_store_group
+                )
+            elif self.all2all_backend in (
+                "mori_high_throughput",
+                "mori_low_latency",
+            ):
+                from .all2all import MoriAll2AllManager
+
+                self.all2all_manager = MoriAll2AllManager(
+                    self.cpu_group, self.all2all_backend
+                )
+            elif self.all2all_backend == "deepep_v2":
+                from .all2all import DeepEPV2All2AllManager
+
+                self.all2all_manager = DeepEPV2All2AllManager(
+                    self.cpu_group,
+                    tcp_store_group,
+                    device_group=self.device_group,
+                )
+            elif self.all2all_backend == "moonep":
+                from .all2all import MoonEPAll2AllManager
+
+                self.all2all_manager = MoonEPAll2AllManager(
+                    self.cpu_group,
+                    tcp_store_group,
+                    device_group=self.device_group,
+                )
+            elif self.all2all_backend == "nixl_ep":
+                from .all2all import NixlEPAll2AllManager
+
+                self.all2all_manager = NixlEPAll2AllManager(
+                    self.cpu_group, tcp_store_group
+                )
+            elif (
+                self.all2all_backend == "flashinfer_all2allv"
+                or self.all2all_backend == "flashinfer_nvlink_two_sided"
+            ):
+                if self.all2all_backend == "flashinfer_all2allv":
+                    logger.warning_once(
+                        "'flashinfer_all2allv' is deprecated and has been renamed to"
+                        "'flashinfer_nvlink_two_sided'. It will be removed in a future"
+                        "release."
+                    )
+                from .all2all import FlashInferNVLinkTwoSidedManager
+
+                self.all2all_manager = FlashInferNVLinkTwoSidedManager(
+                    self.cpu_group, tcp_store_group
+                )
+            elif self.all2all_backend == "flashinfer_nvlink_one_sided":
+                from .all2all import FlashInferNVLinkOneSidedManager
+
+                self.all2all_manager = FlashInferNVLinkOneSidedManager(self.cpu_group)
+            else:
+                raise ValueError(f"Unknown all2all backend: {self.all2all_backend}")
+
+            logger.info_once(
+                "Using %s all2all manager.",
+                self.all2all_manager.__class__.__name__,
+                scope="global",
+            )
+
+    def _log_all_reduce_backend_selection(self) -> None:
+        """Log the all-reduce backends that are active for this group.
+
+        The dispatch chain in ``all_reduce`` tries backends in this order and
+        falls through to the next one if the current backend rejects the
+        input (size/dtype gates) or is disabled. The list of "enabled"
+        backends below is the subset of potential backends that may be
+        chosen at dispatch time for this group; the actual per-call choice
+        depends on the input tensor.
+        """
+        all_potential_ar_backends = [
+            "FLASHINFER_PCIE_IPC",
+            "FLASHINFER",
+            "NCCL_SYMM_MEM",
+            "QUICK_REDUCE",
+            "AITER_CUSTOM",
+            "CUSTOM",
+            "SYMM_MEM",
+            "PYNCCL",
+        ]
+        enabled_ar_backends: list[str] = []
+        if (
+            self.fi_pcie_ipc_ar_comm is not None
+            and not self.fi_pcie_ipc_ar_comm.disabled
+        ):
+            enabled_ar_backends.append("FLASHINFER_PCIE_IPC")
+        if self.fi_ar_comm is not None and not self.fi_ar_comm.disabled:
+            enabled_ar_backends.append("FLASHINFER")
+        # Mirror the static preconditions of `should_nccl_symm_mem_allreduce`:
+        # VLLM_BATCH_INVARIANT off, NCCL symm mem enabled, world_size meets
+        # min_world_size, and world_size either has a tuned entry in
+        # `custom_ar_preferred_ranges` or is greater than
+        # `always_use_above_world_size`. World sizes that fail the latter (e.g.
+        # 5/6/7 with the default config) never dispatch NCCL symm mem
+        # regardless of input. The per-tensor-size check inside the function
+        # stays as a runtime decision.
+        nccl_symm_ws_ok = self.world_size >= NCCL_SYMM_MEM_ALL_REDUCE_CONFIG[
+            "min_world_size"
+        ] and (
+            self.world_size
+            in NCCL_SYMM_MEM_ALL_REDUCE_CONFIG["custom_ar_preferred_ranges"]
+            or self.world_size
+            > NCCL_SYMM_MEM_ALL_REDUCE_CONFIG["always_use_above_world_size"]
+        )
+        if (
+            self.pynccl_comm is not None
+            and not self.pynccl_comm.disabled
+            and is_symmetric_memory_enabled()
+            and not envs.VLLM_BATCH_INVARIANT
+            and nccl_symm_ws_ok
+        ):
+            enabled_ar_backends.append("NCCL_SYMM_MEM")
+        if self.qr_comm is not None and not self.qr_comm.disabled:
+            enabled_ar_backends.append("QUICK_REDUCE")
+        if (
+            self.use_aiter_allreduce
+            and self.aiter_ar_comm is not None
+            and not self.aiter_ar_comm.disabled
+        ):
+            enabled_ar_backends.append("AITER_CUSTOM")
+        if self.ca_comm is not None and not self.ca_comm.disabled:
+            enabled_ar_backends.append("CUSTOM")
+        if self.symm_mem_comm is not None and not self.symm_mem_comm.disabled:
+            enabled_ar_backends.append("SYMM_MEM")
+        if self.pynccl_comm is not None and not self.pynccl_comm.disabled:
+            enabled_ar_backends.append("PYNCCL")
+
+        logger.info_once(
+            "Using %s all-reduce backends (in dispatch order) for group "
+            "'%s' out of potential backends: %s.",
+            "[" + ", ".join(f"'{b}'" for b in enabled_ar_backends) + "]",
+            self.unique_name or "<unnamed>",
+            "[" + ", ".join(f"'{b}'" for b in all_potential_ar_backends) + "]",
+            scope="global",
+        )
+
+    # GB10: sequence-parallel collectives over RDMA, tried first by
+    # custom_all_gather / custom_reduce_scatter below: arx all-reduces for
+    # decode sizes (~13-27 us against ~80 for NCCL), arxbig for prefill-sized
+    # gathers. Returning None hands the call on.
+    def _arx_reduce_scatter(self, input_):
+        arx_comm = self.arx_comm
+        if arx_comm is None or arx_comm.disabled or not arx_comm.should_use(input_):
+            return None
+        chunk = input_.shape[0] // self.world_size
+        return arx_comm.all_reduce(input_)[self.rank_in_group * chunk:(self.rank_in_group + 1) * chunk]
+
+    def _arx_all_gather(self, input_):
+        # arxbig's gather lands in pinned memory, and the GEMMs that read it
+        # there ran 16% slower overall than NCCL's gather into device memory.
+        big = self.arxbig
+        if big is not None and not big.disabled and big.gather and big.should_gather(input_):
+            return big.all_gather(input_)
+        # Each rank's rows in their own slot of a zero tensor: the sum is the
+        # concatenation, exactly (x + 0 is x).
+        arx_comm = self.arx_comm
+        if (arx_comm is None or arx_comm.disabled or input_.dtype != torch.bfloat16
+                or input_.numel() * 2 * self.world_size > arx_comm.max_bytes):
+            return None
+        full = input_.new_zeros(input_.shape[0] * self.world_size, *input_.shape[1:])
+        if not arx_comm.should_use(full):
+            return None
+        chunk = input_.shape[0]
+        full[self.rank_in_group * chunk:(self.rank_in_group + 1) * chunk] = input_
+        return arx_comm.all_reduce(full)
+
+    def all_reduce(self, input_):
+        arx_comm = self.arx_comm
+        if arx_comm is not None and not arx_comm.disabled and arx_comm.should_use(input_, two_shot=True):
+            return arx_comm.all_reduce(input_)
+        fi_ar_comm = self.fi_ar_comm
+        use_fi_ar = (
+            fi_ar_comm is not None
+            and not fi_ar_comm.disabled
+            and fi_ar_comm.should_use_fi_ar(input_)
+        )
+
+        # since currently we perform copy input -> symm_input -> out-of-place AR
+        # return symm_output, we don't need to check if input is symmetric
+        if (
+            self.pynccl_comm is not None
+            and not use_fi_ar
+            and should_nccl_symm_mem_allreduce(self.pynccl_comm.world_size, input_)
+        ):
+            out = torch.ops.vllm.all_reduce_symmetric_with_copy(input_)
+            if out is not None:
+                return out
+        qr_comm = self.qr_comm
+        if (
+            qr_comm is not None
+            and not qr_comm.disabled
+            and qr_comm.should_quick_allreduce(input_)
+        ):
+            out = qr_comm.quick_all_reduce(input_)
+            assert out is not None
+            return out
+        fi_pcie_ipc_ar_comm = self.fi_pcie_ipc_ar_comm
+        if fi_pcie_ipc_ar_comm is not None and fi_pcie_ipc_ar_comm.should_use(input_):
+            return fi_pcie_ipc_ar_comm.all_reduce(input_)
+        if use_fi_ar:
+            assert fi_ar_comm is not None
+            out = fi_ar_comm.all_reduce(input_)
+            assert out is not None
+            return out
+        aiter_ar_comm = self.aiter_ar_comm
+        if (
+            self.use_aiter_allreduce
+            and aiter_ar_comm is not None
+            and not aiter_ar_comm.disabled
+            and aiter_ar_comm.should_custom_ar(input_)
+        ):
+            out = aiter_ar_comm.custom_all_reduce(input_)
+            assert out is not None
+            return out
+        ca_comm = self.ca_comm
+        if (
+            ca_comm is not None
+            and not ca_comm.disabled
+            and ca_comm.should_custom_ar(input_)
+        ):
+            out = ca_comm.custom_all_reduce(input_)
+            assert out is not None
+            return out
+        symm_mem_comm = self.symm_mem_comm
+        if symm_mem_comm is not None and symm_mem_comm.should_use_symm_mem(input_):
+            out = symm_mem_comm.all_reduce(input_)
+            assert out is not None
+            return out
+        pynccl_comm = self.pynccl_comm
+        if pynccl_comm is None or pynccl_comm.disabled:
+            out = input_.clone()
+            torch.distributed.all_reduce(out, group=self.device_group)
+            return out
+        assert pynccl_comm is not None
+        out = pynccl_comm.all_reduce(input_)
+        if out is None:
+            # fall back to the default all-reduce using PyTorch.
+            # this usually happens during testing.
+            # when we run the model, allreduce only happens for the TP
+            # group, where we always have either custom allreduce or pynccl.
+            out = input_.clone()
+            torch.distributed.all_reduce(out, group=self.device_group)
+        return out
+
+    def custom_all_gather(self, input_: torch.Tensor) -> torch.Tensor | None:
+        out = self._arx_all_gather(input_)
+        if out is not None:
+            return out
+        ca_comm = self.ca_comm
+        if ca_comm is None:
+            return None
+        return ca_comm.custom_all_gather(input_.contiguous())
+
+    def custom_reduce_scatter(self, input_: torch.Tensor) -> torch.Tensor | None:
+        out = self._arx_reduce_scatter(input_)
+        if out is not None:
+            return out
+        ca_comm = self.ca_comm
+        if ca_comm is None:
+            return None
+        return ca_comm.custom_reduce_scatter(input_.contiguous())
+
+    def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
+        # Route uniform dim-0 all-gathers through NVLS symmetric memory when
+        # enabled (mirrors reduce_scatter); otherwise fall back to the
+        # PyNccl/base-class all-gather. Sequence parallelism's
+        # gather-before-GEMM uses dim=0 with tp-aligned (uniform) shards.
+        if dim < 0:
+            dim += input_.dim()
+        if dim == 0 and should_nccl_symm_mem_ag_rs():
+            return self._all_gather_symm_mem(input_.contiguous())
+
+        pynccl_comm = self.pynccl_comm
+        if pynccl_comm is None or pynccl_comm.disabled:
+            return super().all_gather(input_, dim)
+
+        # On ROCm, the base-class all_gather (all_gather_into_tensor) is faster
+        # than the manual pynccl + torch.empty + movedim + reshape path below,
+        # which adds a per-call output allocation and (for dim != 0) an extra
+        # copy on every step. This is on the hot path for TP forward passes, so
+        # keep ROCm on the base-class collective to avoid a decode regression.
+        if current_platform.is_rocm():
+            return super().all_gather(input_, dim)
+
+        input_size = input_.size()
+        output_size = (input_size[0] * self.world_size,) + input_size[1:]
+        output_tensor = torch.empty(
+            output_size, dtype=input_.dtype, device=input_.device
+        )
+        pynccl_comm.all_gather(output_tensor, input_.contiguous())
+        output_tensor = output_tensor.reshape((self.world_size,) + input_size)
+        output_tensor = output_tensor.movedim(0, dim)
+        return output_tensor.reshape(
+            input_size[:dim]
+            + (self.world_size * input_size[dim],)
+            + input_size[dim + 1 :]
+        )
+
+    def reduce_scatter(self, input_: torch.Tensor, dim: int = -1):
+        world_size = self.world_size
+        pynccl_comm = self.pynccl_comm
+        assert pynccl_comm is not None
+        if dim < 0:
+            # Convert negative dim to positive.
+            dim += input_.dim()
+
+        # Note: This will produce an incorrect answer if we don't make
+        # the input_tensor contiguous. Possible bug in reduce_scatter_tensor?
+        input_tensor = input_.movedim(0, dim).contiguous()
+
+        assert input_tensor.shape[0] % world_size == 0
+        chunk_size = input_tensor.shape[0] // world_size
+        output_shape = (chunk_size,) + input_tensor.shape[1:]
+
+        if should_nccl_symm_mem_ag_rs():
+            output = self._reduce_scatter_symm_mem(input_tensor)
+        else:
+            output = torch.empty(
+                output_shape, dtype=input_tensor.dtype, device=input_tensor.device
+            )
+            pynccl_comm.reduce_scatter(output, input_tensor)
+
+        # Reshape before returning
+        return output.movedim(0, dim).contiguous()
+
+    def reduce_scatterv(
+        self, input_: torch.Tensor, dim: int = -1, sizes: list[int] | None = None
+    ):
+        world_size = self.world_size
+        pynccl_comm = self.pynccl_comm
+        assert pynccl_comm is not None
+        if dim < 0:
+            # Convert negative dim to positive.
+            dim += input_.dim()
+
+        # 'sizes' is not needed if all inputs in the same group have the same
+        # shape
+        if sizes is not None and all(s == sizes[0] for s in sizes):
+            sizes = None
+
+        # Note: This will produce an incorrect answer if we don't make
+        # the input_tensor contiguous. Possible bug in reduce_scatter_tensor?
+        input_tensor = input_.movedim(0, dim).contiguous()
+
+        if sizes is not None:
+            assert len(sizes) == world_size, f"{len(sizes)} == {world_size}"
+            assert input_tensor.shape[0] == sum(sizes)
+            chunk_size = sizes[self.rank_in_group]
+        else:
+            assert input_tensor.shape[0] % world_size == 0
+            chunk_size = input_tensor.shape[0] // world_size
+        output_shape = (chunk_size,) + input_tensor.shape[1:]
+
+        if self._can_use_aiter_ag_rs(sizes):
+            aiter_comm = self.aiter_ar_comm
+            assert aiter_comm is not None
+            if aiter_comm.should_custom_rs(input_tensor, dim=0):
+                output = torch.empty(
+                    output_shape, dtype=input_tensor.dtype, device=input_tensor.device
+                )
+                aiter_comm.custom_reduce_scatter(input_tensor, output, dim=0)
+                return output.movedim(0, dim).contiguous()
+
+        # Symmetric memory is only used when all ranks have uniform sizes.
+        # ncclCommWindowRegister is collective: asymmetric pool allocations
+        # from variable per-rank sizes cause deadlocks.
+        use_symm_mem = sizes is None and should_nccl_symm_mem_ag_rs()
+        if use_symm_mem:
+            output = self._reduce_scatter_symm_mem(input_tensor)
+        else:
+            output = torch.empty(
+                output_shape, dtype=input_tensor.dtype, device=input_tensor.device
+            )
+            use_deterministic_rs = envs.VLLM_BATCH_INVARIANT and world_size > 2
+            if use_deterministic_rs:
+                # Reduce to a fixed root (0) for determinism
+                reduced = torch.empty_like(input_tensor)
+                sizes = sizes if sizes else [chunk_size] * world_size
+                pynccl_comm.reduce(reduced, input_tensor, root=0)
+                pynccl_comm.scatter(output, reduced, sizes, root=0)
+            elif sizes is not None and sizes.count(sizes[0]) != len(sizes):
+                pynccl_comm.reduce_scatterv(output, input_tensor, sizes=sizes)
+            else:
+                pynccl_comm.reduce_scatter(output, input_tensor)
+
+        # Reshape before returning
+        return output.movedim(0, dim).contiguous()
+
+    def _get_symm_scratch(
+        self,
+        role: str,
+        shape: tuple[int, ...],
+        dtype: torch.dtype,
+        device: torch.device,
+    ) -> torch.Tensor:
+        """Persistent, pre-registered NCCL symmetric-memory scratch buffer.
+
+        Allocating a fresh symm tensor per collective pays the
+        ``nccl_symm_mem_context`` snapshot + window-registration scan on every
+        call (~0.5 ms/RS+AG pair, dwarfing the NVLS transfer itself). Instead we
+        allocate once per ``(role, shape, dtype)``, register once, and reuse.
+
+        Safe for serial (eager) sequence parallelism: each collective's result
+        is consumed on the same stream before the next same-role collective
+        reuses the buffer. Distinct roles (e.g. ``rs_in`` vs ``ag_out``, both
+        full-size) get distinct buffers so a reduce-scatter input copy never
+        clobbers a still-live all-gather output.
+        """
+        from vllm.distributed.device_communicators.pynccl_allocator import (
+            nccl_symm_mem_context,
+        )
+
+        pynccl_comm = self.pynccl_comm
+        assert pynccl_comm is not None
+        cache = self.__dict__.setdefault("_symm_scratch_bufs", {})
+        key = (role, tuple(shape), dtype)
+        buf = cache.get(key)
+        if buf is None:
+            with nccl_symm_mem_context(pynccl_comm):
+                buf = torch.empty(shape, dtype=dtype, device=device)
+            cache[key] = buf
+        return buf
+
+    def _reduce_scatter_symm_mem(
+        self,
+        input_tensor: torch.Tensor,
+    ) -> torch.Tensor:
+        """ReduceScatter using NCCL symmetric memory (NVLS).
+
+        Only called for uniform-size reduce_scatter (variable sizes are
+        guarded out by the caller to avoid asymmetric ncclCommWindowRegister).
+        Uses persistent pre-registered scratch (see _get_symm_scratch).
+        """
+        from vllm.distributed.device_communicators.pynccl_allocator import (
+            is_symmetric_memory_tensor,
+        )
+
+        pynccl_comm = self.pynccl_comm
+        assert pynccl_comm is not None
+
+        chunk = input_tensor.shape[0] // self.world_size
+        output_shape = (chunk,) + tuple(input_tensor.shape[1:])
+
+        symm_output = self._get_symm_scratch(
+            "rs_out", output_shape, input_tensor.dtype, input_tensor.device
+        )
+        # NVLS reduce-scatter (LDMC) requires the input in symmetric memory.
+        if is_symmetric_memory_tensor(input_tensor):
+            symm_input = input_tensor
+        else:
+            symm_input = self._get_symm_scratch(
+                "rs_in",
+                tuple(input_tensor.shape),
+                input_tensor.dtype,
+                input_tensor.device,
+            )
+            symm_input.copy_(input_tensor)
+
+        pynccl_comm.reduce_scatter(symm_output, symm_input)
+        return symm_output
+
+    def send(self, tensor: torch.Tensor, dst: int | None = None) -> None:
+        """Sends a tensor to the destination rank in a blocking way."""
+        """NOTE: `dst` is the local rank of the destination rank."""
+        if dst is None:
+            dst = (self.rank_in_group + 1) % self.world_size
+
+        pynccl_comm = self.pynccl_comm
+        if pynccl_comm is not None and not pynccl_comm.disabled:
+            pynccl_comm.send(tensor, dst)
+        else:
+            torch.distributed.send(tensor, self.ranks[dst], self.device_group)
+
+    def recv(
+        self, size: torch.Size, dtype: torch.dtype, src: int | None = None
+    ) -> torch.Tensor:
+        """Receives a tensor from the source rank."""
+        """NOTE: `src` is the local rank of the source rank."""
+        if src is None:
+            src = (self.rank_in_group - 1) % self.world_size
+
+        tensor = torch.empty(size, dtype=dtype, device=self.device)
+        pynccl_comm = self.pynccl_comm
+        if pynccl_comm is not None and not pynccl_comm.disabled:
+            pynccl_comm.recv(tensor, src)
+        else:
+            torch.distributed.recv(tensor, self.ranks[src], self.device_group)
+        return tensor
+
+    def broadcast(self, tensor: torch.Tensor, src: int = 0) -> torch.Tensor:
+        """Broadcast a tensor from source rank to all ranks."""
+        if self.world_size == 1:
+            return tensor
+
+        pynccl_comm = self.pynccl_comm
+        if pynccl_comm is not None and not pynccl_comm.disabled:
+            pynccl_comm.broadcast(tensor, src)
+            return tensor
+        else:
+            raise ValueError("No PyNCCL communicator found")
+
+    def destroy(self):
+        if self.pynccl_comm is not None:
+            self.pynccl_comm.destroy()
+            self.pynccl_comm = None
+        if self.ca_comm is not None:
+            self.ca_comm = None
+        if self.aiter_ar_comm is not None:
+            self.aiter_ar_comm.close()
+            self.aiter_ar_comm = None
+        if self.fi_ar_comm is not None:
+            self.fi_ar_comm.destroy()
+            self.fi_ar_comm = None
+        if self.fi_pcie_ipc_ar_comm is not None:
+            self.fi_pcie_ipc_ar_comm.destroy()
+            self.fi_pcie_ipc_ar_comm = None
+        if self.all2all_manager is not None:
+            self.all2all_manager.destroy()
+            self.all2all_manager = None  # type: ignore[assignment]
+
+    def _can_use_aiter_ag_rs(self, sizes: list[int] | None) -> bool:
+        """Whether the AITER custom AG/RS fast path may run for this collective.
+
+        Requires:
+        - uniform batches
+        - FULL CUDAgraphs
+        """
+        if (
+            not self.use_aiter_ag_rs
+            or self.aiter_ar_comm is None
+            or self.aiter_ar_comm.disabled
+        ):
+            return False
+        if sizes is not None:
+            return False
+
+        from vllm.config.compilation import CUDAGraphMode
+        from vllm.forward_context import get_forward_context
+
+        try:
+            ctx = get_forward_context()
+        except AssertionError:
+            return False
+        if ctx.cudagraph_runtime_mode != CUDAGraphMode.FULL:
+            return False
+        bd = ctx.batch_descriptor
+        return bd is not None and bd.uniform
+
+    def suspend(self) -> None:
+        if self.pynccl_comm is not None:
+            self.pynccl_comm.suspend()
+
+    def resume(self) -> None:
+        if self.pynccl_comm is not None:
+            self.pynccl_comm.resume()
+
+    def checkpoint_prepare(self) -> None:
+        # Only FlashInfer all-reduce and FlashInfer all2all are supported for now.
+        from .flashinfer_all_reduce import checkpoint_prepare_fi_ar_workspaces
+
+        checkpoint_prepare_fi_ar_workspaces(self.cpu_group)
+        if self.all2all_manager is not None:
+            self.all2all_manager.checkpoint_prepare()
+
+    def checkpoint_restore(self) -> None:
+        # Only FlashInfer all-reduce and FlashInfer all2all are supported for now.
+        from .flashinfer_all_reduce import checkpoint_restore_fi_ar_workspaces
+
+        checkpoint_restore_fi_ar_workspaces(self.cpu_group)
+        if self.all2all_manager is not None:
+            self.all2all_manager.checkpoint_restore()
+
+    def all_gatherv(
+        self,
+        input_: torch.Tensor | list[torch.Tensor],
+        dim: int = 0,
+        sizes: list[int] | None = None,
+    ):
+        if dim != 0:
+            raise NotImplementedError("only dim 0 all-gatherv is supported")
+        world_size = self.world_size
+
+        # 'sizes' is not needed if all inputs in the same group have the same
+        # shape
+        if sizes is not None and all(s == sizes[0] for s in sizes):
+            sizes = None
+
+        if self._can_use_aiter_ag_rs(sizes):
+            aiter_comm = self.aiter_ar_comm
+            assert aiter_comm is not None
+            if isinstance(input_, torch.Tensor):
+                if aiter_comm.should_custom_ag(input_):
+                    out = aiter_comm.custom_all_gather(input_, dim=0)
+                    if out is not None:
+                        return out
+            elif all(aiter_comm.should_custom_ag(inp) for inp in input_):
+                outs = [aiter_comm.custom_all_gather(inp, dim=0) for inp in input_]
+                if all(o is not None for o in outs):
+                    return outs
+
+        pynccl_comm = self.pynccl_comm
+        assert pynccl_comm is not None and not pynccl_comm.disabled
+
+        # Symmetric memory is only used when all ranks have uniform sizes.
+        # ncclCommWindowRegister is collective: asymmetric pool allocations
+        # from variable per-rank sizes cause deadlocks.
+        if sizes is None and should_nccl_symm_mem_ag_rs():
+            if isinstance(input_, torch.Tensor):
+                return self._all_gather_symm_mem(input_)
+            return self._all_gather_batched_symm_mem(input_)
+
+        def _all_gather_single(input_: torch.Tensor, sizes: list[int] | None = None):
+            input_size = input_.size()
+            if sizes is not None:
+                assert len(sizes) == world_size
+                assert input_.shape[dim] == sizes[self.rank_in_group], (
+                    f"{input_.shape[dim]} != {sizes[self.rank_in_group]}"
+                )
+                output_size = (sum(sizes),) + input_size[1:]
+            else:
+                output_size = (input_size[0] * world_size,) + input_size[1:]
+            # Allocate output tensor.
+            output_tensor = torch.empty(
+                output_size, dtype=input_.dtype, device=input_.device
+            )
+            if sizes is not None:
+                pynccl_comm.all_gatherv(output_tensor, input_, sizes=sizes)
+            else:
+                pynccl_comm.all_gather(output_tensor, input_)
+            return output_tensor
+
+        if isinstance(input_, torch.Tensor):
+            return _all_gather_single(input_, sizes)
+
+        output_list = []
+        pynccl_comm.group_start()
+        for inp in input_:
+            output_list.append(_all_gather_single(inp, sizes=sizes))
+        pynccl_comm.group_end()
+
+        return output_list
+
+    def _all_gather_symm_mem(self, input_: torch.Tensor) -> torch.Tensor:
+        """AllGather a single tensor using NCCL symmetric memory (NVLS).
+
+        Only the output needs to be in symmetric memory; NCCL does not
+        require the AG input to be symmetrically allocated.
+        """
+        pynccl_comm = self.pynccl_comm
+        assert pynccl_comm is not None
+
+        out_size = (input_.size(0) * self.world_size,) + tuple(input_.size()[1:])
+        # Persistent pre-registered scratch avoids the per-call symm-mem context
+        # snapshot/registration overhead (see _get_symm_scratch).
+        symm_output = self._get_symm_scratch(
+            "ag_out", out_size, input_.dtype, input_.device
+        )
+        pynccl_comm.all_gather(symm_output, input_)
+        return symm_output
+
+    def _all_gather_batched_symm_mem(
+        self, inputs: list[torch.Tensor]
+    ) -> list[torch.Tensor]:
+        """AllGather a list of tensors using NCCL symmetric memory (NVLS).
+
+        Uses group_start/group_end to batch the collectives.
+        Only the output needs to be in symmetric memory (see
+        _all_gather_symm_mem).
+        """
+        from vllm.distributed.device_communicators.pynccl_allocator import (
+            nccl_symm_mem_context,
+        )
+
+        pynccl_comm = self.pynccl_comm
+        assert pynccl_comm is not None
+        world_size = self.world_size
+
+        symm_outputs = []
+        with nccl_symm_mem_context(pynccl_comm):
+            for inp in inputs:
+                out_size = (inp.size(0) * world_size,) + inp.size()[1:]
+                symm_outputs.append(
+                    torch.empty(out_size, dtype=inp.dtype, device=inp.device)
+                )
+
+        pynccl_comm.group_start()
+        for symm_out, inp in zip(symm_outputs, inputs):
+            pynccl_comm.all_gather(symm_out, inp)
+        pynccl_comm.group_end()
+
+        return symm_outputs
+
+    def dispatch_router_logits(
+        self,
+        hidden_states: torch.Tensor,
+        router_logits: torch.Tensor,
+        is_sequence_parallel: bool = False,
+        extra_tensors: list[torch.Tensor] | None = None,
+    ) -> (
+        tuple[torch.Tensor, torch.Tensor]
+        | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]
+    ):
+        """Dispatch the hidden states and router logits to the appropriate device.
+        This is a no-op in the base class.
+        """
+        assert self.all2all_manager is not None
+        return self.all2all_manager.dispatch_router_logits(
+            hidden_states,
+            router_logits,
+            is_sequence_parallel,
+            extra_tensors,
+        )
+
+    def dispatch(
+        self,
+        hidden_states: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        is_sequence_parallel: bool = False,
+        extra_tensors: list[torch.Tensor] | None = None,
+    ) -> (
+        tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+        | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
+    ):
+        """Dispatch the hidden states and topk weights/ids to the appropriate device.
+        This is a no-op in the base class.
+        """
+        assert self.all2all_manager is not None
+        return self.all2all_manager.dispatch(
+            hidden_states,
+            topk_weights,
+            topk_ids,
+            is_sequence_parallel,
+            extra_tensors=extra_tensors,
+        )
+
+    def combine(
+        self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
+    ) -> torch.Tensor:
+        """Combine the hidden states and router logits from the appropriate device.
+        This is a no-op in the base class.
+        """
+        assert self.all2all_manager is not None
+        return self.all2all_manager.combine(
+            hidden_states,
+            is_sequence_parallel,
+        )
+
+    def batch_isend_irecv(self, p2p_ops: list):
+        pynccl_comm = self.pynccl_comm
+        if pynccl_comm is not None and not pynccl_comm.disabled:
+            pynccl_comm.batch_isend_irecv(p2p_ops)
+        else:
+            raise ValueError("No PyNCCL communicator found")

docker-images/glm53-kindling/experimental/arx/test_ring.py

@@ -0,0 +1,328 @@
+#!/usr/bin/env python3
+"""Bit-exact and speed test for arx and arxbig, in mesh or ring mode, without vLLM.
+
+Every rank builds every rank's inputs from fixed seeds and computes the
+result the way the kernels do (bf16 values summed in fp32 in rank order,
+rounded once). Each collective's output must match that reference bit for
+bit, and each rank's output must match what the other ranks expect it to
+hold. Covered: the arx all-reduce, arxbig's all_gather and reduce_scatter,
+and the fused MoE finalize + rs_finish path (with routing weights that are
+powers of two, so the reference can round the same way the kernel's FMAs
+do). Calls also go out in back-to-back bursts of mixed collectives, checked
+only afterwards, so buffer reuse and relays run at full speed. It then
+prints all-reduce latency and all-gather / reduce-scatter bandwidth.
+
+Mesh mode needs every box on the switch; ring mode needs the ring cabling,
+or the switch with both ARX_RING_*_HCAS set to the NCCL_IB_HCA devices (that
+checks the relay logic, not the speed). The two modes can't share one
+cabling, so each is its own run. Since both must match the same reference,
+passing both means they agree.
+
+Run one container per box from the repo root, rank 0 first, all within a
+minute. The serving stack must be down. Rank r must be ring member r in
+cable order (rank r's next port cabled to rank r+1's prev port).
+
+    sudo docker run --rm --gpus all --network host --ipc host --ulimit memlock=-1 \\
+      --device /dev/infiniband:/dev/infiniband -v "$PWD/experimental/arx:/arx:ro" \\
+      -v /tmp/arx-ext:/root/.cache/torch_extensions \\
+      -e RANK=<0..3> -e WORLD_SIZE=4 -e MASTER_ADDR=<rank 0's LAN address> -e MASTER_PORT=29610 \\
+      -e VLLM_HOST_IP=<this box's LAN address> $FABRIC \\
+      --entrypoint python3 <serving image> /arx/test_ring.py --mode <mesh|ring>
+
+with FABRIC, for mesh, the entrypoint's values (its log prints them):
+    -e NCCL_IB_HCA=<root 0 device>,<root 1 device> -e NCCL_IB_GID_INDEX=<index>
+and for ring, the two functions of the port facing each neighbour, e.g.
+    -e ARX_RING_PREV_HCAS=rocep1s0f0,roceP2p1s0f0 -e ARX_RING_NEXT_HCAS=rocep1s0f1,roceP2p1s0f1
+
+The /tmp/arx-ext mount keeps the compiled extensions between runs.
+WORLD_SIZE=2 works the same way. At WORLD_SIZE=3 ring mode is arx only:
+arxbig refuses a ring of 3, so this script stops at its setup. Every rank
+prints its failures and rank 0 ends with "RESULT <mode>: PASS" or "FAIL";
+the exit code is nonzero on a failure. A rank that makes no progress for
+120 s says where it was and exits.
+"""
+import argparse
+import hashlib
+import os
+import random
+import subprocess
+import sys
+import threading
+import time
+import zlib
+
+import torch
+import torch.distributed as dist
+
+BF16 = torch.bfloat16
+
+
+def lan_iface(ip: str) -> str:
+    out = subprocess.run(["ip", "-o", "-4", "addr", "show"], capture_output=True, text=True).stdout
+    for line in out.splitlines():
+        f = line.split()
+        if len(f) > 3 and f[3].split("/")[0] == ip:
+            return f[1]
+    sys.exit(f"no interface holds VLLM_HOST_IP={ip}")
+
+
+def gen(shape, dtype, *key) -> torch.Tensor:
+    g = torch.Generator(device="cuda")
+    g.manual_seed(zlib.crc32(repr(key).encode()))
+    if dtype == torch.uint8:
+        return torch.randint(0, 256, shape, generator=g, device="cuda", dtype=torch.uint8)
+    return (torch.randn(shape, generator=g, device="cuda") * 4).to(dtype)
+
+
+def rank_order_sum(xs) -> torch.Tensor:
+    acc = torch.zeros(xs[0].shape, dtype=torch.float32, device="cuda")
+    for x in xs:
+        acc += x.float()
+    return acc.to(BF16)
+
+
+def raw(t: torch.Tensor) -> torch.Tensor:
+    return t.contiguous().reshape(-1).view(torch.uint8)
+
+
+def digest(t: torch.Tensor) -> str:
+    return hashlib.sha256(raw(t).cpu().numpy().tobytes()).hexdigest()[:16]
+
+
+class Watchdog:
+    def __init__(self, rank: int, secs: float):
+        self.rank, self.secs, self.what, self.t = rank, secs, "setup", time.time()
+        threading.Thread(target=self._run, daemon=True).start()
+
+    def at(self, what: str) -> None:
+        self.what, self.t = what, time.time()
+
+    def _run(self) -> None:
+        while True:
+            time.sleep(5)
+            if time.time() - self.t > self.secs:
+                print(f"rank {self.rank}: no progress for {self.secs:.0f} s in {self.what}", flush=True)
+                os._exit(3)
+
+
+class Test:
+    def __init__(self, mode: str):
+        self.mode = mode
+        self.rank, self.world = dist.get_rank(), dist.get_world_size()
+        self.checks = self.failures = 0
+        sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+        import arx
+
+        dev = torch.device("cuda", 0)
+        self.ar = arx.ArxCommunicator(dist.group.WORLD, dev)
+        self.big = arx.ArxBig(dist.group.WORLD, dev)
+        if self.ar.disabled or self.big.disabled or not self.big.rs:
+            sys.exit("arx or arxbig did not come up; see the warnings above")
+        self.dog = Watchdog(self.rank, 120)  # after the extensions compile
+
+    # Each op issues one collective and returns (output, reference function,
+    # kind): "same" outputs are the whole reference on every rank, "chunk"
+    # outputs are this rank's chunk of it.
+
+    def allreduce(self, n: int, key):
+        ins = [gen((n,), BF16, "ar", key, s) for s in range(self.world)]
+        out = torch.empty_like(ins[self.rank])
+        self.ar.ext.allreduce(ins[self.rank], out, [], [])
+        return out, lambda: rank_order_sum(ins), "same"
+
+    def all_gather(self, nbytes: int, dtype, key):
+        n = nbytes // torch.empty((), dtype=dtype).element_size()
+        ins = [gen((n,), dtype, "ag", key, s) for s in range(self.world)]
+        out = self.big.all_gather(ins[self.rank]).clone()  # later gathers overwrite the returned slot
+        return out, lambda: torch.cat(ins), "same"
+
+    def reduce_scatter(self, chunk_bytes: int, key):
+        n = chunk_bytes // 2 * self.world
+        ins = [gen((n,), BF16, "rs", key, s) for s in range(self.world)]
+        out = self.big.ext.reduce_scatter(ins[self.rank])
+        return out, lambda: rank_order_sum(ins), "chunk"
+
+    def finalize_rs(self, T: int, Tpad: int, H: int, key, topk: int = 2):
+        def inputs(s):
+            R = T * topk
+            y = gen((R, H), BF16, "fy", key, s)
+            shared = gen((T, H), BF16, "fs", key, s)
+            g = torch.Generator(device="cuda")
+            g.manual_seed(zlib.crc32(repr(("fp", key, s)).encode()))
+            pos = torch.randint(0, R, (T * topk,), generator=g, device="cuda", dtype=torch.int32)
+            w = 2.0 ** torch.randint(-2, 2, (T, topk), generator=g, device="cuda").float()
+            return y, pos, w, shared
+
+        mine = inputs(self.rank)
+        seq = self.big.ext.moe_finalize_rs(*mine, T, Tpad, None)
+        out = self.big.ext.rs_finish(seq, Tpad // self.world, H)
+
+        def ref():
+            rows = []
+            for s in range(self.world):
+                y, pos, w, shared = inputs(s)
+                acc = shared.float()
+                p = pos.view(T, topk).long()
+                for q in range(topk):
+                    acc = acc + w[:, q:q + 1] * y[p[:, q]].float()
+                r = torch.zeros((Tpad, H), dtype=BF16, device="cuda")
+                r[:T] = acc.to(BF16)
+                rows.append(r)
+            return rank_order_sum(rows)
+
+        return out, ref, "chunk"
+
+    def check(self, name: str, out, ref_fn, kind: str) -> None:
+        self.dog.at(f"check {name}")
+        torch.cuda.synchronize()
+        ref = ref_fn()
+        if kind == "chunk":
+            chunks = ref.chunk(self.world)
+            want = [digest(c) for c in chunks]
+            ok = torch.equal(raw(out), raw(chunks[self.rank]))
+        else:
+            want = [digest(ref)] * self.world
+            ok = torch.equal(raw(out), raw(ref))
+        got: list = [None] * self.world
+        dist.all_gather_object(got, digest(out))
+        self.checks += 1
+        bad = [j for j in range(self.world) if got[j] != want[j]]
+        if not ok or bad:
+            self.failures += 1
+            first = ""
+            if not ok:
+                exp = chunks[self.rank] if kind == "chunk" else ref
+                diff = (raw(out) != raw(exp)).nonzero()
+                first = f", first differing byte {diff[0].item()} of {raw(exp).numel()}"
+            print(f"[{self.mode}] rank {self.rank} FAIL {name}: matches reference {ok}{first}; "
+                  f"ranks whose output differs from what this rank expects: {bad}", flush=True)
+
+    def issue(self, op: str, size: int, key):
+        if op == "ar":
+            return self.allreduce(size // 2, key)
+        if op == "ag":
+            return self.all_gather(size, BF16, key)
+        if op == "ag8":
+            return self.all_gather(size, torch.uint8, key)
+        if op == "rs":
+            return self.reduce_scatter(size, key)
+        Tpad = max(size // (2 * 1024), 8 * self.world)
+        Tpad -= Tpad % (8 * self.world)
+        return self.finalize_rs(Tpad - Tpad // 7, Tpad, 1024, key)
+
+    def burst(self, name: str, ops) -> None:
+        """Issues every op back to back, then checks them all."""
+        self.dog.at(f"burst {name}")
+        issued = [(f"{name}[{i}] {op} {size}", *self.issue(op, size, (name, i))) for i, (op, size) in enumerate(ops)]
+        for label, out, ref_fn, kind in issued:
+            self.check(label, out, ref_fn, kind)
+
+    def bench(self, label: str, fn, iters: int) -> float:
+        """Seconds per call, the slowest rank's."""
+        self.dog.at(f"bench {label}")
+        for _ in range(3):
+            fn()
+        torch.cuda.synchronize()
+        dist.barrier()
+        a, b = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
+        a.record()
+        for _ in range(iters):
+            fn()
+        b.record()
+        torch.cuda.synchronize()
+        t: list = [None] * self.world
+        dist.all_gather_object(t, a.elapsed_time(b) / 1e3 / iters)
+        return max(t)
+
+    def say(self, msg: str) -> None:
+        if self.rank == 0:
+            print(f"[{self.mode}] {msg}", flush=True)
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--mode", choices=["mesh", "ring"], required=True)
+    p.add_argument("--repeat", type=int, default=3, help="checks per size")
+    p.add_argument("--bursts", type=int, default=20, help="bursts of 16 mixed collectives")
+    args = p.parse_args()
+
+    os.environ["VLLM_ARX_RING"] = "1" if args.mode == "ring" else "0"
+    os.environ["VLLM_ARXBIG_RS"] = "1"
+    if "GLOO_SOCKET_IFNAME" not in os.environ and "VLLM_HOST_IP" in os.environ:
+        os.environ["GLOO_SOCKET_IFNAME"] = lan_iface(os.environ["VLLM_HOST_IP"])
+    rank, world = int(os.environ["RANK"]), int(os.environ["WORLD_SIZE"])
+    torch.cuda.set_device(0)
+    dist.init_process_group("gloo", init_method=f"tcp://{os.environ['MASTER_ADDR']}:{os.environ['MASTER_PORT']}",
+                            rank=rank, world_size=world)
+    for k in sorted(os.environ):
+        if k.startswith(("NCCL_IB", "ARX_", "VLLM_ARX", "GLOO_")):
+            print(f"[{args.mode}] rank {rank} {k}={os.environ[k]}", flush=True)
+    T = Test(args.mode)
+    KB, MB = 1 << 10, 1 << 20
+
+    T.say("checking sizes")
+    for n in (8, 64, 1024, 4096, 32768, 131072, 262144):
+        for i in range(args.repeat):
+            T.check(f"allreduce {n * 2} B", *T.allreduce(n, ("size", i)))
+    for nbytes in (16, 48, 4 * KB, 1 * MB + 16, 8 * MB, 32 * MB):
+        for i in range(args.repeat):
+            T.check(f"all_gather {nbytes} B", *T.all_gather(nbytes, BF16, ("size", i)))
+            T.check(f"all_gather u8 {nbytes} B", *T.all_gather(nbytes, torch.uint8, ("size", i)))
+    for chunk in (16, 48, 4 * KB, 1 * MB + 16, 8 * MB, 32 * MB):
+        for i in range(args.repeat):
+            T.check(f"reduce_scatter {chunk} B/chunk", *T.reduce_scatter(chunk, ("size", i)))
+    for Tn, Tpad, H in ((100, 64 * world, 1024), (1000, 1024, 4096), (16000, 16384, 4096)):
+        for i in range(args.repeat):
+            T.check(f"finalize_rs T={Tn} Tpad={Tpad} H={H}", *T.finalize_rs(Tn, Tpad, H, ("size", i)))
+
+    T.say("checking bursts of mixed collectives")
+    rng = random.Random(1234)  # the same sequence on every rank
+    sizes = {"ar": [16, 8 * KB, 64 * KB, 512 * KB], "ag": [16, 64 * KB, 1 * MB, 4 * MB],
+             "ag8": [48, 1 * MB], "rs": [16, 64 * KB, 1 * MB, 4 * MB], "fin": [1 * MB, 8 * MB]}
+    for b in range(args.bursts):
+        ops = []
+        for _ in range(16):
+            op = rng.choice(["ar", "ar", "ar", "ag", "ag8", "rs", "rs", "fin"])
+            ops.append((op, rng.choice(sizes[op])))
+        T.burst(f"burst{b}", ops)
+    T.burst("allreduce x64", [("ar", 8 * KB)] * 64)
+    T.burst("reduce_scatter x8", [("rs", 8 * MB)] * 8)
+    T.burst("finalize_rs x4", [("fin", 8 * MB)] * 4)
+
+    T.say("all-reduce latency (arx)")
+    for n in (4096, 32768, 131072, 262144):
+        x = gen((n,), BF16, "bench")
+        out = torch.empty_like(x)
+        t = T.bench(f"allreduce {n}", lambda: T.ar.ext.allreduce(x, out, [], []), 2000)
+        T.say(f"  allreduce {n * 2 // KB:4d} KB  {t * 1e6:7.1f} us")
+    T.say("all-gather / reduce-scatter bandwidth (arxbig), data received per rank")
+    for nbytes in (1 * MB, 8 * MB, 32 * MB):
+        x = gen((nbytes // 2,), BF16, "bench")
+        t = T.bench(f"all_gather {nbytes}", lambda: T.big.all_gather(x), 50)
+        T.say(f"  all_gather     {nbytes // MB:3d} MB/rank  {t * 1e3:7.3f} ms  "
+              f"{(world - 1) * nbytes * 8 / t / 1e9:6.1f} Gb/s")
+    for chunk in (1 * MB, 8 * MB, 32 * MB):
+        x = gen((chunk // 2 * world,), BF16, "bench")
+        t = T.bench(f"reduce_scatter {chunk}", lambda: T.big.ext.reduce_scatter(x), 50)
+        T.say(f"  reduce_scatter {chunk // MB:3d} MB/chunk {t * 1e3:7.3f} ms  "
+              f"{(world - 1) * chunk * 8 / t / 1e9:6.1f} Gb/s")
+    y, pos, w, shared = (gen((32768, 4096), BF16, "bench"), torch.randint(0, 32768, (32768,), device="cuda",
+                         dtype=torch.int32), torch.ones((16384, 2), device="cuda"), gen((16384, 4096), BF16, "bench"))
+
+    def fin():
+        T.big.ext.rs_finish(T.big.ext.moe_finalize_rs(y, pos, w, shared, 16384, 16384, None), 16384 // world, 4096)
+
+    t = T.bench("finalize_rs", fin, 20)
+    T.say(f"  finalize_rs + rs_finish, 16k x 4096 bf16: {t * 1e3:.3f} ms")
+
+    T.dog.at("result")
+    fails: list = [None] * world
+    dist.all_gather_object(fails, T.failures)
+    T.say(f"RESULT {args.mode}: {'PASS' if sum(fails) == 0 else 'FAIL'}: {T.checks} checks per rank, "
+          f"failures per rank {fails}")
+    dist.barrier()
+    os._exit(0 if sum(fails) == 0 else 1)  # the proxy threads never exit
+
+
+if __name__ == "__main__":
+    main()

docker-images/glm53-kindling/experimental/compose/adaptive-k.yaml

@@ -0,0 +1,37 @@
+# Adaptive speculation length: the scheduler picks how many of the 7 DFlash2
+# drafts each step verifies (adaptive_k.py). The per-batch-size schedule only
+# makes vLLM capture CUDA graphs for every level; the scheduler ignores it.
+# draft-trunc/: after the drafts exist, each request drops its drafts whose
+# predicted survival (vLLM's acceptance estimator) is below
+# VLLM_DRAFT_TRUNC_TAU. Dead drafts keep the verify shape, skip their MoE
+# experts and are rejected, so outputs are unchanged. Every
+# VLLM_DRAFT_TRUNC_EXPLORE-th step skips the cut and trains the estimator.
+# disable_eagle_block_drop: vLLM drops the last matching prefix-cache block for
+# EAGLE-style drafters, 2,304 tokens here at TP=4, and every cache hit then
+# recomputes it. With the drop off, an agent turn's time to first token falls by
+# about a third and the output is identical (#54).
+services:
+  glm53:
+    volumes:
+      - ../experimental/adaptive-k/adaptive_k.py:/usr/local/lib/python3.12/dist-packages/vllm/v1/core/sched/adaptive_k.py:ro
+      - ../experimental/draft-trunc/draft_trunc.py:/usr/local/lib/python3.12/dist-packages/vllm/v1/worker/gpu/spec_decode/draft_trunc.py:ro
+      - ../experimental/draft-trunc/model_runner.py:/usr/local/lib/python3.12/dist-packages/vllm/v1/worker/gpu/model_runner.py:ro
+      - ../experimental/draft-trunc/rejection_sampler.py:/usr/local/lib/python3.12/dist-packages/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py:ro
+      - ../experimental/draft-trunc/grouped_topk_router.py:/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py:ro
+      - ../experimental/draft-trunc/speculator.py:/usr/local/lib/python3.12/dist-packages/vllm/v1/worker/gpu/spec_decode/speculator.py:ro
+      - ../experimental/draft-trunc/dflash2_speculator.py:/usr/local/lib/python3.12/dist-packages/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py:ro
+    environment:
+      - EXTRA_ARGS=--scheduler-cls vllm.v1.core.sched.adaptive_k.AdaptiveKScheduler
+      # Merged into the entrypoint's --speculative-config (method, drafter and
+      # width stay SPEC_METHOD, DFLASH_MODEL and SPEC_TOKENS).
+      - SPEC_EXTRA={"disable_eagle_block_drop":true,"attention_backend":"TRITON_ATTN","num_speculative_tokens_per_batch_size":[[1,1,7],[2,2,5],[3,3,4],[4,4,3],[5,5,2],[6,32,7]]}
+      - VLLM_ADAPTIVE_K_CONTROL=/root/.cache/adaptive-k
+      - VLLM_DRAFT_TRUNC_TAU=${VLLM_DRAFT_TRUNC_TAU:-0.3}
+      - VLLM_DRAFT_TRUNC_EXPLORE=${VLLM_DRAFT_TRUNC_EXPLORE:-4}
+      - VLLM_DRAFT_TRUNC_CONTROL=/root/.cache/draft-trunc
+      - VLLM_DRAFT_TRUNC_LOG=${VLLM_DRAFT_TRUNC_LOG:-}
+      - VLLM_ADAPTIVE_K_PER_REQUEST=${VLLM_ADAPTIVE_K_PER_REQUEST:-0}
+      - VLLM_ADAPTIVE_K_CODE_MODE=${VLLM_ADAPTIVE_K_CODE_MODE:-gated}
+      - VLLM_ADAPTIVE_K_ONLINE=${VLLM_ADAPTIVE_K_ONLINE:-1}
+      - VLLM_ADAPTIVE_K_MODEL=${VLLM_ADAPTIVE_K_MODEL:-}
+      - VLLM_ADAPTIVE_K_COST_MS=${VLLM_ADAPTIVE_K_COST_MS:-2:29.7,3:33.1,4:36.5,5:39.8,6:43.1,7:46.0,8:48.9}

docker-images/glm53-kindling/experimental/compose/arx.yaml

@@ -0,0 +1,18 @@
+# EXPERIMENT: arx RDMA all-reduce for decode-sized TP all-reduces (see arx.py),
+# and with VLLM_ARXBIG=1 arxbig.cu's all-gather for prefill-sized SP gathers.
+#   docker compose -f compose/glm53.yaml -f compose/arx.yaml up -d
+services:
+  glm53:
+    volumes:
+      - ../experimental/arx/arx.py:/usr/local/lib/python3.12/dist-packages/vllm/distributed/device_communicators/arx.py:ro
+      - ../experimental/arx/arx_vllm.cu:/usr/local/lib/python3.12/dist-packages/vllm/distributed/device_communicators/arx_vllm.cu:ro
+      - ../experimental/arx/arxbig.cu:/usr/local/lib/python3.12/dist-packages/vllm/distributed/device_communicators/arxbig.cu:ro
+      - ../experimental/arx/cuda_communicator.py:/usr/local/lib/python3.12/dist-packages/vllm/distributed/device_communicators/cuda_communicator.py:ro
+    environment:
+      - VLLM_ARX_ALLREDUCE=1
+      - VLLM_ARXBIG=${VLLM_ARXBIG:-1}
+      - VLLM_ARXBIG_RS=${VLLM_ARXBIG_RS:-1}
+      - VLLM_ARXBIG_AG=${VLLM_ARXBIG_AG:-0}
+      # Decode all-reduces above 256 KB and up to this size run arx's two-shot mode
+      # (reduce-scatter + all-gather over RDMA); 0 sends them to NCCL.
+      - VLLM_ARX_TWO_SHOT_MAX_KB=${VLLM_ARX_TWO_SHOT_MAX_KB:-2048}

docker-images/glm53-kindling/experimental/compose/dispram.yaml

@@ -0,0 +1,23 @@
+# The GB10's display carveout as KV cache. The firmware reserves 2 GiB for a
+# display that the driver never uses on this chip. kindling-spark-os
+# (github.com/kindlingai/kindling-spark-os) runs dispramd, which lends that
+# range to CUDA processes. This file mounts dispramd's socket and the host's
+# copy of its vLLM plugin, so the client always matches the daemon. The plugin
+# adds what dispramd can lend (1.99 GiB) to the pinned KV pool and maps the
+# pool's tail onto it.
+#
+# kindling-spark-os's 64 KiB kernel also returns ~2 GiB of RAM, so at TP=4
+# this raises the pinned pool by 2 GiB (KV_EXTRA_MEMORY). TP=2 has no room for
+# it, and TP=3 and 6 are unmeasured.
+#
+# ./glm53 adds this file by itself when /run/dispram/dispram.sock exists.
+# DISPRAM=0 leaves it out. Every rank should have it: vLLM sizes the pool by
+# the smallest rank, so one box without dispramd cancels the gain.
+services:
+  glm53:
+    volumes:
+      - /run/dispram:/run/dispram
+      - /opt/kindling/dispram/python:/opt/dispram:ro
+    environment:
+      - PYTHONPATH=/opt/dispram
+      - KV_EXTRA_MEMORY=${KV_EXTRA_MEMORY:-2147483648}

docker-images/glm53-kindling/experimental/compose/fixes.yaml

@@ -0,0 +1,35 @@
+# Fixes found by profiling (2026-09-27): FlashInfer's per-step rollback clone,
+# the indexer's fp32 head-gate GEMM, a switch for the fused mHC kernel's token
+# limit, and a Triton sparse MLA kernel in place of FlashInfer's FA2 path
+# (VLLM_TRITON_SPARSE_MLA=1, about 3.3x faster on GB10). kda.py and
+# causal_conv1d.py: the KDA short conv writes q, k and v as separate dense
+# tensors, so FlashKDA no longer copies them (~51 ms per 16k-token step).
+# kv_cache_*.py, worker_utils.py, warmup.py: the drafter's KV gets its own
+# small pool, paged 256 tokens at a time, instead of a page in every KV block
+# (VLLM_GLM5NEXT_DRAFT_POOL=1, VLLM_GLM5NEXT_DRAFT_BLOCK_SIZE).
+# tilelang_kernels.py and mhc_tilelang.py: the fused mHC decode kernel reads
+# its projection weights (hc_*_fn, stored bf16 in the checkpoint) as a cached
+# bf16 copy instead of the model's fp32 one, which halves its weight reads and
+# leaves every product unchanged.
+services:
+  glm53:
+    volumes:
+      - ../experimental/fixes/_fa_common.py:/usr/local/lib/python3.12/dist-packages/flashinfer/mla/_batch_mla/_backends/_fa_common.py:ro
+      - ../experimental/fixes/attention.py:/usr/local/lib/python3.12/dist-packages/vllm/models/glm5next/common/attention.py:ro
+      - ../experimental/fixes/tilelang_kernels.py:/usr/local/lib/python3.12/dist-packages/vllm/model_executor/kernels/mhc/tilelang_kernels.py:ro
+      - ../experimental/fixes/mhc_tilelang.py:/usr/local/lib/python3.12/dist-packages/vllm/model_executor/kernels/mhc/tilelang.py:ro
+      - ../experimental/fixes/flashinfer_mla_sparse_sm90.py:/usr/local/lib/python3.12/dist-packages/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py:ro
+      - ../experimental/fixes/gb10_sparse_mla.py:/usr/local/lib/python3.12/dist-packages/vllm/v1/attention/backends/mla/gb10_sparse_mla.py:ro
+      - ../experimental/fixes/gdn_attn.py:/usr/local/lib/python3.12/dist-packages/vllm/v1/attention/backends/gdn_attn.py:ro
+      - ../experimental/fixes/kda.py:/usr/local/lib/python3.12/dist-packages/vllm/models/glm5next/common/kda.py:ro
+      - ../experimental/fixes/causal_conv1d.py:/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/mamba/ops/causal_conv1d.py:ro
+      - ../experimental/fixes/kv_cache_utils.py:/usr/local/lib/python3.12/dist-packages/vllm/v1/core/kv_cache_utils.py:ro
+      - ../experimental/fixes/kv_cache_coordinator.py:/usr/local/lib/python3.12/dist-packages/vllm/v1/core/kv_cache_coordinator.py:ro
+      - ../experimental/fixes/kv_cache_interface.py:/usr/local/lib/python3.12/dist-packages/vllm/v1/kv_cache_interface.py:ro
+      - ../experimental/fixes/worker_utils.py:/usr/local/lib/python3.12/dist-packages/vllm/v1/worker/utils.py:ro
+      - ../experimental/fixes/warmup.py:/usr/local/lib/python3.12/dist-packages/vllm/v1/worker/gpu/warmup.py:ro
+    environment:
+      - VLLM_MHC_FUSED_MAX_TOKENS=${VLLM_MHC_FUSED_MAX_TOKENS:-32}
+      - VLLM_TRITON_SPARSE_MLA=${VLLM_TRITON_SPARSE_MLA:-1}
+      - VLLM_GLM5NEXT_DRAFT_POOL=${VLLM_GLM5NEXT_DRAFT_POOL:-1}
+      - VLLM_GLM5NEXT_DRAFT_BLOCK_SIZE=${VLLM_GLM5NEXT_DRAFT_BLOCK_SIZE:-}

docker-images/glm53-kindling/experimental/compose/fp8.yaml

@@ -0,0 +1,15 @@
+# Dense linears the checkpoint left in bf16 (dense_fp8.py): FP8, except the
+# layers matching VLLM_DENSE_W4, which go to NVFP4 (megadense4.cu): by default
+# the KDA in_proj, attention o_proj, the shared experts and every DFlash
+# drafter layer. Needs snapshot.yaml's loader;
+# keeps its own snapshots.
+services:
+  glm53:
+    volumes:
+      - ../experimental/snapshot/dense_fp8.py:/usr/local/lib/python3.12/dist-packages/vllm/model_executor/model_loader/dense_fp8.py:ro
+      - ../experimental/megamoe/megadense4.cu:/opt/megamoe/megadense4.cu:ro
+    environment:
+      - VLLM_DENSE_FP8=1
+      - VLLM_DENSE_FP8_LM_HEAD=${VLLM_DENSE_FP8_LM_HEAD:-1}
+      - VLLM_DENSE_W4=${VLLM_DENSE_W4:-in_proj_qkvbfg_a$$|^model\.|self_attn\.o_proj$$|shared_experts\.(gate_up_proj|down_proj)$$}
+      - VLLM_WEIGHT_SNAPSHOT_TAG=fp8w4c-2026-09-28

docker-images/glm53-kindling/experimental/compose/megamoe.yaml

@@ -0,0 +1,19 @@
+# megamoe for decode-sized MoE batches (W4A16, CUTLASS's own tensors), and with
+# VLLM_MOE_PREFILL=1 moe_prefill.cu for batches of VLLM_MOE_PREFILL_MIN_TOKENS or
+# more; CUTLASS keeps the sizes in between. Needs snapshot.yaml's loader, which
+# calls megamoe_vllm.install() after loading. VLLM_MOE_PREFILL_Y8: fc2's per-expert
+# rows as e4m3 (scale per 128 columns) for the fused SP finalize: half the bytes.
+services:
+  glm53:
+    volumes:
+      - ../experimental/megamoe/megamoe_rm.cu:/opt/megamoe/megamoe_rm.cu:ro
+      - ../experimental/megamoe/moe_prefill.cu:/opt/megamoe/moe_prefill.cu:ro
+      - ../experimental/megamoe/megamoe_vllm.py:/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/fused_moe/megamoe_vllm.py:ro
+    environment:
+      - VLLM_MEGAMOE=${VLLM_MEGAMOE:-1}
+      - VLLM_MEGAMOE_MAX_TOKENS=${VLLM_MEGAMOE_MAX_TOKENS:-8}
+      - VLLM_MEGAMOE_VARIANT=${VLLM_MEGAMOE_VARIANT:-20}
+      - VLLM_MOE_PREFILL=${VLLM_MOE_PREFILL:-1}
+      - VLLM_MOE_PREFILL_MIN_TOKENS=${VLLM_MOE_PREFILL_MIN_TOKENS:-1024}
+      - VLLM_MOE_PREFILL_CHECK=${VLLM_MOE_PREFILL_CHECK:-0}
+      - VLLM_MOE_PREFILL_Y8=${VLLM_MOE_PREFILL_Y8:-1}

docker-images/glm53-kindling/experimental/compose/recoverssm.yaml

@@ -0,0 +1,17 @@
+# KDA spec verify from one checkpoint state per request, ported from Kimi-K3's
+# RecoverSSM (fixes/recoverssm.py). On with this file;
+# VLLM_GLM5NEXT_RECOVERSSM=0 turns it off. The stock spec path keeps 1 + k recurrent states
+# per request and writes one per verified token. This verifies from the
+# checkpoint, records each token's inputs, and after sampling replays only
+# the accepted tokens into it. Each request then holds one KDA block per
+# group, where the stock path holds 1 + k (8 at k=7). The records cost
+# ~3.2 MiB per max_num_seqs row per rank, outside the pinned pool.
+# Stack after fixes.yaml (kda.py, causal_conv1d.py). model.py is sp.yaml's
+# file: it installs the runner's commit hook, and SP stays off without sp.yaml.
+services:
+  glm53:
+    volumes:
+      - ../experimental/fixes/recoverssm.py:/usr/local/lib/python3.12/dist-packages/vllm/models/glm5next/common/recoverssm.py:ro
+      - ../experimental/fixes/model.py:/usr/local/lib/python3.12/dist-packages/vllm/models/glm5next/common/model.py:ro
+    environment:
+      - VLLM_GLM5NEXT_RECOVERSSM=${VLLM_GLM5NEXT_RECOVERSSM:-1}

docker-images/glm53-kindling/experimental/compose/snapshot.yaml

@@ -0,0 +1,14 @@
+# EXPERIMENT: local weight snapshots (weight_snapshot.py). The first boot loads
+# the checkpoint and writes one snapshot per rank; later boots load from it.
+services:
+  glm53:
+    volumes:
+      - ../experimental/snapshot/weight_snapshot.py:/usr/local/lib/python3.12/dist-packages/vllm/model_executor/model_loader/weight_snapshot.py:ro
+      - ../experimental/snapshot/base_loader.py:/usr/local/lib/python3.12/dist-packages/vllm/model_executor/model_loader/base_loader.py:ro
+      # CUTLASS NVFP4 MoE under weights_already_processed: allow it, and skip the
+      # in-place scale fold the snapshot already holds.
+      - ../experimental/snapshot/modelopt.py:/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/quantization/modelopt.py:ro
+      - ../experimental/snapshot/flashinfer_cutlass_moe.py:/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py:ro
+    environment:
+      - VLLM_WEIGHT_SNAPSHOT_DIR=/root/.cache/weight-snapshots
+      - VLLM_WEIGHT_SNAPSHOT_TAG=2026-09-28

docker-images/glm53-kindling/experimental/compose/sp.yaml

@@ -0,0 +1,27 @@
+# Sequence parallelism over TP for prefill (fixes/model.py): a forward with at
+# least VLLM_GLM_SP_MIN_TOKENS tokens runs mHC, norms and residuals on a quarter
+# of the tokens per rank, with a gather before and a reduce-scatter after
+# attention and the MLP. Decode batches stay plain TP: at that size the extra
+# collectives cost more than the replicated work they save.
+# VLLM_GLM_SP_FP8_GATHER: KDA attention inputs are gathered as per-token FP8
+# (in_proj's FP8 GEMM quantizes the same rows the same way).
+# VLLM_GLM_ARX_PREFETCH: decode all-reduces prefetch the next weights into L2
+# while they wait (needs arx.yaml).
+# VLLM_GLM_SP_MOE_FUSED: MoE layers under SP compute shared + scaled routed
+# sum in one kernel that feeds arxbig's RDMA reduce-scatter row by row (needs
+# arx.yaml with VLLM_ARXBIG_RS=1 and megamoe.yaml with VLLM_MOE_PREFILL=1).
+# VLLM_GLM_SP_MOE_QUANT_GATHER: those layers route each rank's own rows and
+# gather the input already quantized (NVFP4 for the routed experts, FP8 for the
+# shared expert) instead of as bf16, matching it to within 1e-4 relative.
+services:
+  glm53:
+    volumes:
+      - ../experimental/fixes/model.py:/usr/local/lib/python3.12/dist-packages/vllm/models/glm5next/common/model.py:ro
+    environment:
+      - VLLM_GLM_SP_TP=${VLLM_GLM_SP_TP:-1}
+      - VLLM_GLM_SP_MIN_TOKENS=${VLLM_GLM_SP_MIN_TOKENS:-1024}
+      - VLLM_GLM_SP_FP8_GATHER=${VLLM_GLM_SP_FP8_GATHER:-1}
+      - VLLM_GLM_ARX_PREFETCH=${VLLM_GLM_ARX_PREFETCH:-1}
+      - VLLM_GLM_SP_MOE_FUSED=${VLLM_GLM_SP_MOE_FUSED:-1}
+      - VLLM_GLM_SP_MOE_QUANT_GATHER=${VLLM_GLM_SP_MOE_QUANT_GATHER:-1}
+      - VLLM_GLM_SP_MOE_CHECK=${VLLM_GLM_SP_MOE_CHECK:-0}

docker-images/glm53-kindling/experimental/compose/tp3.yaml

@@ -0,0 +1,23 @@
+# TP=3: three boxes. GLM-5.3-Flash's 64 MLA heads, 64 KDA heads, 2048 MoE
+# intermediate and 154880 vocab do not split by 3, so experimental/tp3/
+# make_padded.py writes a config with padded sizes (66 / 66 / 2304) beside
+# hardlinks of the checkpoint, and tp3pad.py zero-pads the tensors as they
+# load. Zero heads and zero expert columns contribute nothing. Set TP=3 and
+# point MODEL_HOST_DIR / DFLASH_HOST_DIR at the padded dirs.
+#
+# Stack LAST, after adaptive-k.yaml (it restates that file's EXTRA_ARGS, and keeps
+# its SPEC_EXTRA) and
+# after fixes.yaml / sp.yaml / recoverssm.yaml (model.py there carries the
+# load hook, a no-op without tp_pad_orig in the config).
+services:
+  glm53:
+    volumes:
+      - ../experimental/tp3/tp3pad.py:/usr/local/lib/python3.12/dist-packages/vllm/models/glm5next/common/tp3pad.py:ro
+      - ../experimental/tp3/vocab_parallel_embedding.py:/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/vocab_parallel_embedding.py:ro
+    environment:
+      - TP=3
+      # adaptive-k.yaml's scheduler plus a data-parallel vision tower: its 16
+      # ViT heads and 10240-wide merger do not split by 3.
+      - EXTRA_ARGS=--scheduler-cls vllm.v1.core.sched.adaptive_k.AdaptiveKScheduler --mm-encoder-tp-mode data
+      # Padded weights process differently: never restore a TP=4 snapshot.
+      - VLLM_WEIGHT_SNAPSHOT_TAG=tp3pad-2026-09-28

docker-images/glm53-kindling/experimental/compose/tp6.yaml

@@ -0,0 +1,14 @@
+# TP=6: six boxes on the TP=3 zero-padded checkpoint (experimental/tp3):
+# 66 MLA / 66 KDA heads and 2304 MoE columns split as 11 / 11 / 384 per rank.
+# The DFlash2 drafter's 8 kv heads need 12 (48 q) at TP=6:
+#   make_padded.py ... --draft-heads 48 --draft-kv-heads 12 --tp 6
+# Stack LAST, as tp3.yaml.
+services:
+  glm53:
+    volumes:
+      - ../experimental/tp3/tp3pad.py:/usr/local/lib/python3.12/dist-packages/vllm/models/glm5next/common/tp3pad.py:ro
+      - ../experimental/tp3/vocab_parallel_embedding.py:/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/vocab_parallel_embedding.py:ro
+    environment:
+      - TP=6
+      - EXTRA_ARGS=--scheduler-cls vllm.v1.core.sched.adaptive_k.AdaptiveKScheduler --mm-encoder-tp-mode data
+      - VLLM_WEIGHT_SNAPSHOT_TAG=tp6pad-2026-09-29

docker-images/glm53-kindling/experimental/draft-trunc/dflash2_speculator.py

@@ -0,0 +1,265 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import Any
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.config.compilation import CUDAGraphMode
+from vllm.triton_utils import tl, triton
+from vllm.v1.worker.gpu.sample.gumbel import gumbel_noised_argmax
+from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator
+
+
+@triton.jit
+def _selector_walk_kernel(
+    scores_ptr,
+    candidate_ptr,
+    sample_pos_ptr,
+    req_state_ptr,
+    temperature_ptr,
+    seeds_ptr,
+    tokens_ptr,
+    realized_scores_ptr,
+    num_steps: tl.constexpr,
+    top_k: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    SAMPLE_PROBABILISTIC: tl.constexpr,
+    USE_FP64: tl.constexpr,
+):
+    row = tl.program_id(0)
+    offsets = tl.arange(0, BLOCK_K)
+    mask = offsets < top_k
+    req_state = tl.load(req_state_ptr + row * num_steps)
+    valid = req_state >= 0
+    temperature = tl.load(temperature_ptr + req_state, mask=valid, other=0.0)
+    seed = tl.load(seeds_ptr + req_state, mask=valid, other=0)
+    previous = 0
+    for step in range(num_steps):
+        flat = row * num_steps + step
+        score_base = (flat * top_k + previous) * top_k
+        scores = tl.load(
+            scores_ptr + score_base + offsets,
+            mask=mask & valid,
+            other=float("-inf"),
+        ).to(tl.float64 if USE_FP64 else tl.float32)
+        candidate_base = flat * top_k
+        candidates = tl.load(
+            candidate_ptr + candidate_base + offsets,
+            mask=mask & valid,
+            other=0,
+        )
+
+        # sample_pos is the predicted token's position P. Sampling keys a draw
+        # by the position before the sampled token, P-1.
+        sample_pos = tl.load(sample_pos_ptr + flat) - 1
+        _, index = gumbel_noised_argmax(
+            scores,
+            candidates,
+            mask & valid,
+            seed,
+            sample_pos,
+            temperature if SAMPLE_PROBABILISTIC else 0.0,
+            IS_DRAFTING=True,
+            USE_FP64=USE_FP64,
+        )
+
+        tl.store(
+            realized_scores_ptr + candidate_base + offsets,
+            scores,
+            mask=mask & valid,
+        )
+        token = tl.load(candidate_ptr + candidate_base + index, mask=valid, other=0)
+        tl.store(tokens_ptr + flat, token, mask=valid)
+        previous = index
+
+
+@triton.jit
+def _cache_draft_logits_kernel(
+    draft_logits_ptr,
+    cached_candidate_ptr,
+    candidate_ptr,
+    scores_ptr,
+    req_state_ptr,
+    draft_logits_stride_0,
+    draft_logits_stride_1,
+    num_steps: tl.constexpr,
+    top_k: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    flat = tl.program_id(0)
+    req_state = tl.load(req_state_ptr + flat)
+    step = flat % num_steps
+    offsets = tl.arange(0, BLOCK_K)
+    mask = (req_state >= 0) & (offsets < top_k)
+    candidate_base = flat * top_k
+    cache_base = (req_state * num_steps + step) * top_k
+    old_token_ids = tl.load(cached_candidate_ptr + cache_base + offsets, mask=mask)
+    logits_base = (
+        draft_logits_ptr
+        + req_state * draft_logits_stride_0
+        + step * draft_logits_stride_1
+    )
+    tl.store(logits_base + old_token_ids, -float("inf"), mask=mask)
+    token_ids = tl.load(candidate_ptr + candidate_base + offsets, mask=mask)
+    scores = tl.load(scores_ptr + candidate_base + offsets, mask=mask)
+    tl.store(logits_base + token_ids, scores, mask=mask)
+    tl.store(cached_candidate_ptr + cache_base + offsets, token_ids, mask=mask)
+
+
+class CandidateSampler:
+    """The shared DFlash2/LiLiCorr candidate walk and realized proposal cache."""
+
+    def __init__(
+        self, max_num_reqs: int, num_steps: int, top_k: int, device: torch.device
+    ):
+        self.num_steps = num_steps
+        self.top_k = top_k
+        self.scores = torch.empty(
+            max_num_reqs, num_steps, top_k, dtype=torch.float32, device=device
+        )
+        self.cached_candidate_ids = torch.zeros(
+            self.scores.shape, dtype=torch.int64, device=device
+        )
+
+    def sample(
+        self,
+        candidate_ids: torch.Tensor,
+        scores: torch.Tensor,
+        num_reqs: int,
+        sample_pos: torch.Tensor,
+        idx_mapping: torch.Tensor,
+        temperature: torch.Tensor,
+        seeds: torch.Tensor,
+        draft_tokens: torch.Tensor,
+        draft_logits: torch.Tensor | None,
+        use_fp64: bool,
+    ) -> None:
+        candidate_ids = candidate_ids.contiguous()
+        sample_pos = sample_pos.contiguous()
+        idx_mapping = idx_mapping.contiguous()
+        block_k = triton.next_power_of_2(self.top_k)
+        _selector_walk_kernel[(num_reqs,)](
+            scores.contiguous(),
+            candidate_ids.contiguous(),
+            sample_pos,
+            idx_mapping,
+            temperature,
+            seeds,
+            draft_tokens,
+            self.scores,
+            num_steps=self.num_steps,
+            top_k=self.top_k,
+            BLOCK_K=block_k,
+            SAMPLE_PROBABILISTIC=draft_logits is not None,
+            USE_FP64=use_fp64,
+            num_warps=1,
+        )
+
+        if draft_logits is not None:
+            self._cache_draft_logits(
+                candidate_ids, num_reqs * self.num_steps, idx_mapping, draft_logits
+            )
+
+    def _cache_draft_logits(
+        self,
+        candidate_ids: torch.Tensor,
+        num_sample: int,
+        idx_mapping: torch.Tensor,
+        draft_logits: torch.Tensor,
+    ) -> None:
+        block_k = triton.next_power_of_2(self.top_k)
+        _cache_draft_logits_kernel[(num_sample,)](
+            draft_logits,
+            self.cached_candidate_ids,
+            candidate_ids,
+            self.scores,
+            idx_mapping,
+            draft_logits.stride(0),
+            draft_logits.stride(1),
+            num_steps=self.num_steps,
+            top_k=self.top_k,
+            BLOCK_K=block_k,
+            num_warps=1,
+        )
+
+
+class DFlash2Speculator(DFlashSpeculator):
+    _speculator_name = "DFlash2"
+    _candidate_top_k_key = "selector_top_k"
+
+    def __init__(self, vllm_config: VllmConfig, device: torch.device):
+        super().__init__(vllm_config, device)
+        draft_config = self.draft_model_config.hf_config.dflash_config
+        self.top_k = int(draft_config[self._candidate_top_k_key])
+        self.candidate_sampler = CandidateSampler(
+            self.max_num_reqs, self.num_speculative_steps, self.top_k, device
+        )
+
+    def draft_logits_spec(self, vllm_config: VllmConfig) -> tuple[torch.dtype, float]:
+        # fp32 so the walk and the rejection that checks it read the same
+        # distribution; -inf because the cache kernel writes only the K
+        # candidates.
+        return torch.float32, -float("inf")
+
+    def _generate_draft(
+        self,
+        num_reqs: int,
+        num_tokens_padded: int,
+        attn_metadata: dict[str, Any] | None,
+        slot_mappings: dict[str, torch.Tensor] | None,
+        num_tokens_across_dp: torch.Tensor | None,
+        cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
+    ) -> None:
+        last_hidden_states = self._run_model(
+            num_tokens_padded,
+            attn_metadata,
+            slot_mappings,
+            num_tokens_across_dp,
+            cudagraph_runtime_mode,
+        )
+        num_sample = num_reqs * self.num_speculative_steps
+        hidden_states = last_hidden_states[self.sample_indices[:num_sample]].view(
+            num_reqs, self.num_speculative_steps, -1
+        )
+        candidate_ids, unary_logits = self.model.compute_candidates(
+            hidden_states.flatten(0, 1)
+        )
+        candidate_ids = candidate_ids.view(
+            num_reqs, self.num_speculative_steps, self.top_k
+        )
+        unary_logits = unary_logits.view_as(candidate_ids)
+        scores = self._score_candidates(candidate_ids, unary_logits, hidden_states)
+        self.candidate_sampler.sample(
+            candidate_ids,
+            scores,
+            num_reqs,
+            self.sample_pos,
+            self.sample_idx_mapping,
+            self.temperature,
+            self.seeds,
+            self.draft_tokens,
+            self.draft_logits,
+            self.use_fp64_gumbel,
+        )
+        if self.acceptance_estimator is not None:
+            self._maybe_predict_acceptance(
+                self.candidate_sampler.scores[:num_reqs].flatten(0, 1),
+                self.sample_idx_mapping[:num_sample],
+                self.sample_col[:num_sample],
+            )
+
+    def _score_candidates(
+        self,
+        candidate_ids: torch.Tensor,
+        unary_logits: torch.Tensor,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        num_reqs = candidate_ids.shape[0]
+        anchor_token_ids = self.input_buffers.input_ids[
+            : num_reqs * self.num_query_per_req : self.num_query_per_req
+        ]
+        return self.model.model.candidate_selector(
+            candidate_ids, unary_logits, hidden_states, anchor_token_ids
+        )

docker-images/glm53-kindling/experimental/draft-trunc/draft_trunc.py

@@ -0,0 +1,159 @@
+"""Per-step draft truncation.
+
+The scheduler picks how many drafts each request verifies before the drafts
+exist, so the verify shape is fixed. Truncation leaves the shape alone and
+marks a request's later drafts dead on the device: their MoE routes go to
+expert -1, which the MoE kernels skip without reading those experts' weights,
+and the rejection sampler sees them as -1, so acceptance stops before them.
+They come after the live tokens, so causal attention leaves the live tokens
+exact.
+
+By default (VLLM_DRAFT_TRUNC_TAU, 0.3 in adaptive-k.yaml) each request keeps
+its leading drafts whose predicted survival, the running product of the
+acceptance estimator's per-draft probabilities, is at least tau. A cut step
+grades only the drafts the estimator already expected to survive, so fitting
+on it would bias the estimator, and a position it always kills could never
+recover. Every VLLM_DRAFT_TRUNC_EXPLORE-th step (4 by default) therefore skips
+the cut, and only those steps are folded into the estimator. The file named by
+VLLM_DRAFT_TRUNC_CONTROL overrides that while it exists: "tau X" sets the
+threshold, "live N" keeps each request's first N drafts, and anything else
+truncates nothing.
+"""
+import os
+
+import torch
+import vllm
+
+from vllm.triton_utils import tl, triton
+
+# The other files in this directory are whole copies of vLLM's, with the cut added. They
+# match one vLLM commit, the image's VLLM_REF. On any other they would put back old code
+# over new, so the engine stops instead.
+VLLM_COMMIT = "ddd6fbca"
+if f"+g{VLLM_COMMIT}" not in vllm.__version__:
+    raise RuntimeError(f"experimental/draft-trunc copies vLLM {VLLM_COMMIT}'s files, but this image runs "
+                       f"vLLM {vllm.__version__}. Refresh the copies from this image, or remove the "
+                       "draft-trunc mounts from experimental/compose/adaptive-k.yaml.")
+
+_TAU = float(os.environ.get("VLLM_DRAFT_TRUNC_TAU") or 0)
+ENABLED = bool(os.environ.get("VLLM_DRAFT_TRUNC_CONTROL")) or _TAU > 0
+# Measurement: per verify step, each request's draft confidences (the
+# acceptance estimator's, computed when the drafts were made) and how many
+# tokens it sampled, one JSON line per step, on tensor-parallel rank 0.
+LOG = os.environ.get("VLLM_DRAFT_TRUNC_LOG", "")
+_LOG_MAX = int(os.environ.get("VLLM_DRAFT_TRUNC_LOG_STEPS", "20000"))
+_logged = 0
+_control = os.environ.get("VLLM_DRAFT_TRUNC_CONTROL", "")
+_mtime = 0.0
+_mode: tuple[str, float] | None = None
+_EXPLORE = int(os.environ.get("VLLM_DRAFT_TRUNC_EXPLORE") or 4)
+_steps = 0
+_cut = False
+
+
+def _read_control() -> tuple[str, float] | None:
+    """("live", N), ("tau", X), or None to truncate nothing."""
+    global _mtime, _mode
+    default = ("tau", _TAU) if _TAU > 0 else None
+    try:
+        mtime = os.stat(_control).st_mtime if _control else None
+    except OSError:
+        mtime = None
+    if mtime is None:
+        _mtime, _mode = 0.0, default
+        return default
+    if mtime != _mtime:
+        _mtime = mtime
+        words = open(_control).read().split()
+        _mode = (words[0], float(words[1])) if len(words) == 2 and words[0] in ("live", "tau") else None
+    return _mode
+
+
+def live_drafts() -> int | None:
+    """The control file's N, or None."""
+    mode = _read_control()
+    return int(mode[1]) if mode and mode[0] == "live" else None
+
+
+@triton.jit
+def _mark_dead_kernel(is_padding_ptr, query_start_loc_ptr, cu_num_logits_ptr, live, BLOCK: tl.constexpr):
+    b = tl.program_id(0)
+    num_drafts = tl.load(cu_num_logits_ptr + b + 1) - tl.load(cu_num_logits_ptr + b) - 1
+    end = tl.load(query_start_loc_ptr + b + 1)
+    j = tl.arange(0, BLOCK)
+    # A request's drafts are its last num_drafts tokens.
+    tl.store(is_padding_ptr + end - num_drafts + j, tl.full((BLOCK,), 1, tl.int1),
+             mask=(j >= live) & (j < num_drafts))
+
+
+@triton.jit
+def _mark_cut_kernel(is_padding_ptr, query_start_loc_ptr, cu_num_logits_ptr, pred_ptr, pred_stride,
+                     idx_mapping_ptr, tau, BLOCK: tl.constexpr):
+    b = tl.program_id(0)
+    num_drafts = tl.load(cu_num_logits_ptr + b + 1) - tl.load(cu_num_logits_ptr + b) - 1
+    end = tl.load(query_start_loc_ptr + b + 1)
+    state = tl.load(idx_mapping_ptr + b).to(tl.int64)
+    j = tl.arange(0, BLOCK)
+    in_range = j < num_drafts
+    p = tl.load(pred_ptr + state * pred_stride + j, mask=in_range & (state >= 0), other=1.0)
+    # Survival only falls along a request, so the dead drafts are a suffix.
+    dead = (tl.cumprod(p, axis=0) < tau) & in_range
+    tl.store(is_padding_ptr + end - num_drafts + j, tl.full((BLOCK,), 1, tl.int1), mask=dead)
+
+
+def mark_dead(is_padding: torch.Tensor, query_start_loc: torch.Tensor, cu_num_logits: torch.Tensor,
+              num_reqs: int, max_drafts: int, predictions: torch.Tensor | None = None,
+              idx_mapping: torch.Tensor | None = None) -> None:
+    """Set is_padding on each request's dead drafts, per the control file."""
+    global _steps, _cut
+    _cut = False
+    mode = _read_control()
+    if mode is None or num_reqs == 0 or max_drafts == 0:
+        return
+    block = triton.next_power_of_2(max_drafts)
+    if mode[0] == "live":
+        _mark_dead_kernel[(num_reqs,)](is_padding, query_start_loc, cu_num_logits, int(mode[1]), BLOCK=block)
+        _cut = True
+        return
+    if predictions is None or idx_mapping is None:
+        return
+    _steps += 1
+    if _EXPLORE > 0 and _steps % _EXPLORE == 0:
+        return
+    _mark_cut_kernel[(num_reqs,)](is_padding, query_start_loc, cu_num_logits, predictions, predictions.stride(0),
+                                  idx_mapping, mode[1], BLOCK=block)
+    _cut = True
+
+
+def graded() -> bool:
+    """Whether the last step verified every draft, so the estimator may learn from it."""
+    return not _cut
+
+
+def log_step(confidence: torch.Tensor, cu_num_logits: torch.Tensor, num_sampled: torch.Tensor,
+             estimator=None, idx_mapping: torch.Tensor | None = None) -> None:
+    """Append this step's (drafts verified, sampled, confidences) per request to LOG,
+    with the estimator's features and coefficients when it has one."""
+    global _logged
+    if _logged >= _LOG_MAX:
+        return
+    from vllm.distributed import get_tensor_model_parallel_rank
+
+    if get_tensor_model_parallel_rank() != 0:
+        return
+    n = num_sampled.shape[0]
+    drafts = (cu_num_logits[1:n + 1] - cu_num_logits[:n] - 1).tolist()
+    feats = (estimator.features[idx_mapping].tolist() if estimator is not None and idx_mapping is not None
+             else [None] * n)
+    rows = [{"k": k, "sampled": s, "conf": [round(c, 4) for c in conf[:k]],
+             "feat": None if f is None else [round(x, 3) for x in f[:k]]}
+            for k, s, conf, f in zip(drafts, num_sampled.tolist(), confidence.tolist(), feats) if k > 0]
+    if rows and estimator is not None and _logged % 200 == 0:
+        rows.append({"slope": estimator.slope.tolist(), "intercepts": estimator.intercepts.tolist(),
+                     "refits": estimator._refits})
+    if rows:
+        import json
+
+        with open(LOG, "a") as f:
+            f.write(json.dumps(rows) + "\n")
+        _logged += 1

docker-images/glm53-kindling/experimental/draft-trunc/grouped_topk_router.py

@@ -0,0 +1,367 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import os
+from functools import partial
+
+import torch
+
+from vllm import _custom_ops as ops
+from vllm import envs as envs
+from vllm._aiter_ops import rocm_aiter_ops
+from vllm.distributed.eplb.eplb_state import EplbLayerState
+from vllm.forward_context import get_forward_context, is_forward_context_available
+from vllm.model_executor.custom_op import CustomOp
+from vllm.model_executor.layers.fused_moe.config import (
+    RoutingMethodType,
+    get_routing_method_type,
+)
+from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
+    rocm_aiter_grouped_topk,
+)
+from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter
+from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (
+    fused_topk_bias,
+)
+from vllm.model_executor.layers.fused_moe.router.fused_topk_router import fused_topk
+from vllm.model_executor.utils import maybe_disable_graph_partition
+from vllm.platforms import current_platform
+
+
+def fused_grouped_topk(
+    hidden_states: torch.Tensor,
+    gating_output: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+    e_score_correction_bias: torch.Tensor,
+    num_expert_group: int = 0,
+    topk_group: int = 0,
+    scoring_func: str = "softmax",
+    routed_scaling_factor: float = 1.0,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    assert hidden_states.size(0) == gating_output.size(0), "Number of tokens mismatch"
+
+    if scoring_func == "sigmoid":
+        # Fully fused kernel path for sigmoid
+        topk_values, topk_indices = ops.grouped_topk(
+            gating_output,  # raw logits
+            num_expert_group,
+            topk_group,
+            topk,
+            renormalize,
+            routed_scaling_factor,
+            e_score_correction_bias,
+            1,  # scoring_func=1 for sigmoid
+        )
+    elif scoring_func == "softmax":
+        # Apply softmax in Python, then use fused kernel
+        # TODO: Add support for softmax in kernel
+        scores = torch.softmax(gating_output, dim=-1)
+        topk_values, topk_indices = ops.grouped_topk(
+            scores,  # pre-computed scores
+            num_expert_group,
+            topk_group,
+            topk,
+            renormalize,
+            routed_scaling_factor,
+            e_score_correction_bias,
+            0,  # scoring_func=0 (no activation, scores already computed)
+        )
+    else:
+        raise ValueError(f"Unsupported scoring function: {scoring_func}")
+
+    # Fused kernel outputs float32 values and int32 indices directly
+    return topk_values, topk_indices
+
+
+# This is used by the Deepseek-V2 and Deepseek-V3 model
+@torch.compile(
+    dynamic=True,
+    backend=current_platform.simple_compile_backend,
+    options=maybe_disable_graph_partition(current_platform.simple_compile_backend),
+)
+def grouped_topk(
+    hidden_states: torch.Tensor,
+    gating_output: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+    num_expert_group: int = 0,
+    topk_group: int = 0,
+    scoring_func: str = "softmax",
+    routed_scaling_factor: float = 1.0,
+    e_score_correction_bias: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if (
+        envs.VLLM_USE_FUSED_MOE_GROUPED_TOPK
+        and (current_platform.is_cuda() or current_platform.is_xpu())
+        and num_expert_group <= 32
+        and topk <= 32
+        and e_score_correction_bias is not None
+    ):
+        return fused_grouped_topk(
+            hidden_states=hidden_states,
+            gating_output=gating_output,
+            topk=topk,
+            renormalize=renormalize,
+            e_score_correction_bias=e_score_correction_bias,
+            num_expert_group=num_expert_group,
+            topk_group=topk_group,
+            scoring_func=scoring_func,
+            routed_scaling_factor=routed_scaling_factor,
+        )
+
+    assert hidden_states.size(0) == gating_output.size(0), "Number of tokens mismatch"
+
+    if scoring_func == "softmax":
+        scores = torch.softmax(gating_output, dim=-1)
+    elif scoring_func == "sigmoid":
+        scores = gating_output.sigmoid()
+    else:
+        raise ValueError(f"Unsupported scoring function: {scoring_func}")
+
+    num_token = scores.size(0)
+    if e_score_correction_bias is not None:
+        # Store original scores before applying correction bias. We use biased
+        # scores for expert selection but original scores for routing weights
+        original_scores = scores
+        scores = scores + e_score_correction_bias.unsqueeze(0)
+        group_scores = (
+            scores.view(num_token, num_expert_group, -1).topk(2, dim=-1)[0].sum(dim=-1)
+        )
+    else:
+        group_scores = (
+            scores.view(num_token, num_expert_group, -1).max(dim=-1).values
+        )  # [n, n_group]
+
+    # For batch invariance, use sorted=True to ensure deterministic expert selection
+    use_sorted = envs.VLLM_BATCH_INVARIANT
+    group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=use_sorted)[
+        1
+    ]  # [n, top_k_group]
+    group_mask = torch.zeros_like(group_scores)  # [n, n_group]
+    group_mask.scatter_(1, group_idx, 1)  # [n, n_group]
+    score_mask = (
+        group_mask.unsqueeze(-1)
+        .expand(num_token, num_expert_group, scores.size(-1) // num_expert_group)
+        .reshape(num_token, -1)
+    )  # [n, e]
+    tmp_scores = scores.masked_fill(~score_mask.bool(), float("-inf"))  # [n, e]
+
+    if e_score_correction_bias is not None:
+        topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=use_sorted)[1]
+        # Use original unbiased scores for the routing weights
+        topk_weights = original_scores.gather(1, topk_ids)
+    else:
+        topk_weights, topk_ids = torch.topk(
+            tmp_scores, k=topk, dim=-1, sorted=use_sorted
+        )
+
+    if renormalize:
+        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
+
+    if routed_scaling_factor != 1.0:
+        topk_weights = topk_weights * routed_scaling_factor
+    return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
+
+
+# --8<-- [start:grouped_topk]
+@CustomOp.register("grouped_topk")
+class GroupedTopk(CustomOp):
+    """GroupedTopk used by the Deepseek-V2 and Deepseek-V3 model."""
+
+    # --8<-- [end:grouped_topk]
+
+    def __init__(
+        self,
+        topk: int,
+        renormalize: bool,
+        num_expert_group: int = 0,
+        topk_group: int = 0,
+        scoring_func: str = "softmax",
+        routed_scaling_factor: float = 1.0,
+        num_fused_shared_experts: int = 0,
+    ) -> None:
+        super().__init__()
+        self.native_impl = grouped_topk
+        self.topk = topk
+        self.renormalize = renormalize
+        self.num_expert_group = num_expert_group
+        self.topk_group = topk_group
+        self.scoring_func = scoring_func
+        self.routed_scaling_factor = routed_scaling_factor
+        self.num_fused_shared_experts = num_fused_shared_experts
+
+    def forward_native(
+        self,
+        hidden_states: torch.Tensor,
+        gating_output: torch.Tensor,
+        e_score_correction_bias: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return self.native_impl(
+            hidden_states,
+            gating_output,
+            self.topk,
+            self.renormalize,
+            self.num_expert_group,
+            self.topk_group,
+            self.scoring_func,
+            self.routed_scaling_factor,
+            e_score_correction_bias,
+        )
+
+    def forward_cuda(
+        self,
+        hidden_states: torch.Tensor,
+        gating_output: torch.Tensor,
+        e_score_correction_bias: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return self.forward_native(
+            hidden_states, gating_output, e_score_correction_bias
+        )
+
+    def forward_hip(
+        self,
+        hidden_states: torch.Tensor,
+        gating_output: torch.Tensor,
+        e_score_correction_bias: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if rocm_aiter_ops.is_fused_moe_enabled():
+            if not rocm_aiter_ops.is_fusion_moe_shared_experts_enabled():
+                assert self.num_fused_shared_experts == 0
+            return rocm_aiter_grouped_topk(
+                hidden_states,
+                gating_output,
+                self.topk,
+                self.renormalize,
+                self.num_expert_group,
+                self.topk_group,
+                self.scoring_func,
+                self.routed_scaling_factor,
+                e_score_correction_bias,
+                self.num_fused_shared_experts,
+            )
+        else:
+            return self.forward_native(
+                hidden_states, gating_output, e_score_correction_bias
+            )
+
+
+class GroupedTopKRouter(BaseRouter):
+    """Router using grouped top-k routing (e.g., DeepSeekV2/V3)."""
+
+    def __init__(
+        self,
+        top_k: int,
+        global_num_experts: int,
+        num_expert_group: int,
+        topk_group: int,
+        renormalize: bool = True,
+        scoring_func: str = "softmax",
+        routed_scaling_factor: float = 1.0,
+        e_score_correction_bias: torch.Tensor | None = None,
+        num_fused_shared_experts: int = 0,
+        eplb_state: EplbLayerState | None = None,
+        skip_padding: bool = False,
+    ):
+        super().__init__(
+            top_k=top_k,
+            global_num_experts=global_num_experts,
+            eplb_state=eplb_state,
+        )
+        self.num_expert_group = num_expert_group
+        self.topk_group = topk_group
+        self.renormalize = renormalize
+        self.scoring_func = scoring_func
+        self.routed_scaling_factor = routed_scaling_factor
+        self.e_score_correction_bias = e_score_correction_bias
+        self.num_fused_shared_experts = num_fused_shared_experts
+        # Draft truncation (draft_trunc) marks dead verify tokens as padding.
+        self.skip_padding = skip_padding or bool(os.environ.get("VLLM_DRAFT_TRUNC_CONTROL"))
+
+    @property
+    def routing_method_type(self) -> RoutingMethodType:
+        return get_routing_method_type(
+            scoring_func=self.scoring_func,
+            top_k=self.top_k,
+            renormalize=self.renormalize,
+            num_expert_group=self.num_expert_group,
+            has_e_score_bias=self.e_score_correction_bias is not None,
+            routed_scaling_factor=self.routed_scaling_factor,
+        )
+
+    def _compute_routing(
+        self,
+        hidden_states: torch.Tensor,
+        router_logits: torch.Tensor,
+        indices_type: torch.dtype | None,
+        *,
+        input_ids: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Compute routing using grouped top-k."""
+
+        def valid_grouping() -> bool:
+            # Check if num_experts is greater than num_expert_group
+            # and is divisible by num_expert_group
+            num_experts = router_logits.shape[-1]
+            if num_experts <= self.num_expert_group:
+                return False
+            return num_experts % self.num_expert_group == 0
+
+        if not valid_grouping():
+            if self.e_score_correction_bias is not None:
+                topk_weights, topk_ids = fused_topk_bias(
+                    hidden_states=hidden_states,
+                    gating_output=router_logits,
+                    scoring_func=self.scoring_func,
+                    e_score_correction_bias=self.e_score_correction_bias.data,
+                    topk=self.top_k,
+                    renormalize=self.renormalize,
+                )
+                if self.routed_scaling_factor != 1.0:
+                    topk_weights *= self.routed_scaling_factor
+            else:
+                topk_weights, topk_ids, token_expert_indices = fused_topk(
+                    hidden_states=hidden_states,
+                    gating_output=router_logits,
+                    topk=self.top_k,
+                    renormalize=self.renormalize,
+                    indices_type=indices_type,
+                )
+            return topk_weights, topk_ids
+
+        # Select grouped_topk implementation
+        if rocm_aiter_ops.is_fused_moe_enabled():
+            if not rocm_aiter_ops.is_fusion_moe_shared_experts_enabled():
+                assert self.num_fused_shared_experts == 0
+            grouped_topk_impl = partial(
+                rocm_aiter_grouped_topk,
+                num_fused_shared_experts=self.num_fused_shared_experts,
+            )
+        else:
+            grouped_topk_impl = grouped_topk
+
+        topk_weights, topk_ids = grouped_topk_impl(
+            hidden_states=hidden_states,
+            gating_output=router_logits,
+            topk=self.top_k,
+            renormalize=self.renormalize,
+            num_expert_group=self.num_expert_group,
+            topk_group=self.topk_group,
+            scoring_func=self.scoring_func,
+            routed_scaling_factor=self.routed_scaling_factor,
+            e_score_correction_bias=self.e_score_correction_bias,
+        )
+
+        if (
+            self.skip_padding
+            and envs.VLLM_MOE_SKIP_PADDING
+            and is_forward_context_available()
+        ):
+            is_padding = get_forward_context().is_padding
+            # A sequence-parallel shard routes a slice of the batch, which the
+            # first rows of is_padding do not describe.
+            if is_padding is not None and is_padding.shape[0] == topk_ids.shape[0]:
+                is_padding = is_padding[: topk_ids.shape[0]].unsqueeze(1)
+                topk_weights = topk_weights.masked_fill(is_padding, 0)
+                topk_ids = topk_ids.masked_fill(is_padding, -1)
+
+        return topk_weights, topk_ids

docker-images/glm53-kindling/experimental/draft-trunc/model_runner.py

@@ -0,0 +1,2342 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""NOTE: Coding style guide for this file:
+This model runner is shared by all models: text and multimodal, generative
+and embedding, public and private. As a result, this file must only contain
+code that is common to every model. Model-specific behavior belongs in the
+appropriate model-specific files.
+
+In other words:
+* Be paranoid about changing this file. It should remain stable.
+* Be even more paranoid about adding new lines. It should remain minimal.
+
+Even for shared features (for example, different parallelism modes), keep the
+complexity out of this path. The less common the feature, the more it should be
+hidden. Prefer utility functions defined elsewhere and call them from here,
+instead of embedding feature-specific logic directly.
+"""
+
+import functools
+import gc
+import time
+from contextlib import AbstractContextManager
+from copy import deepcopy
+from typing import Any, NamedTuple
+
+import numpy as np
+import torch
+
+from vllm.v1.worker.gpu.spec_decode import draft_trunc
+import torch.nn as nn
+
+import vllm.envs as envs
+from vllm.compilation.counter import compilation_counter
+from vllm.compilation.cuda_graph import CUDAGraphStat
+from vllm.config import VllmConfig
+from vllm.config.compilation import CUDAGraphMode
+from vllm.distributed.aux_output_connector.worker import (
+    AuxOutputWorkerConnector,
+    get_aux_output_connector,
+)
+from vllm.distributed.parallel_state import get_dcp_group, get_pp_group
+from vllm.forward_context import BatchDescriptor, set_forward_context
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager
+from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
+    initialize_mamba_ssu_backend,
+)
+from vllm.model_executor.model_loader import get_model_loader
+from vllm.model_executor.models.interfaces import requires_raw_input_tokens
+from vllm.model_executor.offloader import (
+    create_offloader,
+    get_offloader,
+    set_offloader,
+)
+from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.encoder_budget import (
+    MultiModalBudget,
+)
+from vllm.platforms import current_platform
+from vllm.sequence import IntermediateTensors
+from vllm.tasks import SupportedTask
+from vllm.utils.gc_utils import freeze_gc_for_cudagraph_capture
+from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
+from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, async_tensor_h2d
+from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
+from vllm.v1.kv_cache_interface import (
+    KVCacheConfig,
+    MambaSpec,
+    UniformTypeKVCacheSpecs,
+)
+from vllm.v1.outputs import (
+    DraftTokenIds,
+    ECConnectorOutput,
+    ModelRunnerOutput,
+)
+from vllm.v1.watermarking import create_watermarker
+from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler
+from vllm.v1.watermarking.spec_decode import (
+    create_speculative_target_watermarker,
+    speculative_target_watermark_key,
+)
+from vllm.v1.worker.block_table import get_block_table_width
+from vllm.v1.worker.cp_utils import check_attention_cp_compatibility
+from vllm.v1.worker.gpu import pcp_manager as pcp
+from vllm.v1.worker.gpu.async_utils import (
+    AsyncOutput,
+    AsyncPoolingOutput,
+    StepTimingCollector,
+)
+from vllm.v1.worker.gpu.attn_utils import (
+    FastPrefillHelper,
+    build_slot_mappings_by_layer,
+    get_kv_cache_spec,
+    init_attn_backend,
+    init_kv_cache,
+)
+from vllm.v1.worker.gpu.block_table import BlockTables
+from vllm.v1.worker.gpu.buffer_utils import set_default_max_concurrency
+from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
+from vllm.v1.worker.gpu.cudagraph_utils import (
+    BatchExecutionDescriptor,
+    ModelCudaGraphManager,
+    has_compiled_submodule,
+    make_cudagraph_stats,
+)
+from vllm.v1.worker.gpu.cudagraph_utils import (
+    profile_cudagraph_memory as _profile_cudagraph_memory,
+)
+from vllm.v1.worker.gpu.dp_utils import DPSyncState, dispatch_cg_and_sync_dp
+from vllm.v1.worker.gpu.ec_connector import get_ec_connector
+from vllm.v1.worker.gpu.eplb_utils import EPLBController, step_eplb_after
+from vllm.v1.worker.gpu.input_batch import (
+    InputBatch,
+    InputBuffers,
+    combine_sampled_and_draft_tokens,
+    expand_idx_mapping,
+    post_update,
+    post_update_num_computed_tokens,
+    prepare_pos_seq_lens,
+    prepare_prefill_inputs,
+    set_dummy_context,
+)
+from vllm.v1.worker.gpu.kv_connector import (
+    NO_OP_KV_CONNECTOR,
+    KVConnector,
+    get_kv_connector,
+)
+from vllm.v1.worker.gpu.lora_utils import (
+    LoraState,
+    create_lora_capture_hook,
+    get_lora_capture_cases,
+    get_num_active_loras_for_dispatch,
+)
+from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
+from vllm.v1.worker.gpu.mm.lora import set_active_mm_loras
+from vllm.v1.worker.gpu.model_states import init_model_state
+from vllm.v1.worker.gpu.pool.pooling_runner import PoolingRunner
+from vllm.v1.worker.gpu.pp_utils import PPHandler
+from vllm.v1.worker.gpu.sample.batch_shard import (
+    BatchSharder,
+    all_to_all_logits,
+    gather_sampler_output,
+)
+from vllm.v1.worker.gpu.sample.logits_processor import build_custom_logits_processors
+from vllm.v1.worker.gpu.sample.output import SamplerOutput
+from vllm.v1.worker.gpu.sample.prompt_logprob import PromptLogprobsWorker
+from vllm.v1.worker.gpu.sample.sampler import Sampler
+from vllm.v1.worker.gpu.shutdown import free_before_shutdown
+from vllm.v1.worker.gpu.spec_decode import init_speculator
+from vllm.v1.worker.gpu.spec_decode.adaptive_verification import (
+    AdaptiveVerificationManager,
+    maybe_create_adaptive_verification_manager,
+    resolve_adaptive_cudagraph_mode,
+)
+from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import (
+    set_eagle3_aux_hidden_state_layers,
+    verify_supports_aux_hidden_states_over_pp,
+)
+from vllm.v1.worker.gpu.spec_decode.rejection_sampler import (
+    RejectionSampler,
+    get_max_chunk_logits,
+)
+from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator
+from vllm.v1.worker.gpu.spec_decode.utils import DraftTokensHandler
+from vllm.v1.worker.gpu.states import RequestState
+from vllm.v1.worker.gpu.structured_outputs import StructuredOutputsWorker
+from vllm.v1.worker.gpu.ubatch_utils import (
+    UBatchRunner,
+    UBatchState,
+    maybe_build_ubatch_runner,
+)
+from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
+from vllm.v1.worker.utils import (
+    KVBlockZeroer,
+    clear_layer_kv_caches,
+    copy_kv_cache_blocks_inplace,
+    get_uniform_decode_token_count,
+)
+from vllm.v1.worker.workspace import lock_workspace, use_workspace_lane
+
+logger = init_logger(__name__)
+
+
+class GPUModelRunner(LoRAModelRunnerMixin):
+    def __init__(self, vllm_config: VllmConfig, device: torch.device):
+        self.vllm_config = vllm_config
+        self.model_config = vllm_config.model_config
+        self.cache_config = vllm_config.cache_config
+        self.compilation_config = vllm_config.compilation_config
+        self.lora_config = vllm_config.lora_config
+        self.load_config = vllm_config.load_config
+        self.parallel_config = vllm_config.parallel_config
+        self.scheduler_config = vllm_config.scheduler_config
+        self.speculative_config = vllm_config.speculative_config
+        self._draft_workspace_lane = int(
+            self.speculative_config is not None and self.speculative_config.use_dspark()
+        )
+        self.observability_config = vllm_config.observability_config
+        self.jit_warmup_registry = JitWarmupRegistry(vllm_config)
+
+        self.device = device
+        self.dtype = self.model_config.dtype
+        self.kv_cache_dtype = self.dtype
+        if self.cache_config.cache_dtype != "auto":
+            # Quantized KV cache.
+            self.kv_cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[
+                self.cache_config.cache_dtype
+            ]
+
+        # Lazily initialized in _init_kv_zero_meta() when the KV cache needs
+        # zeroing (e.g. hybrid models with fp8 KV cache).
+        self.kv_block_zeroer: KVBlockZeroer | None = None
+
+        self.vocab_size = self.model_config.get_vocab_size()
+        self.max_model_len = self.model_config.max_model_len
+        self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
+        self.max_num_reqs = self.scheduler_config.max_num_seqs
+        self.is_encoder_decoder = self.model_config.is_encoder_decoder
+
+        self.output_copy_stream = torch.cuda.Stream(self.device)
+
+        # Pipeline parallelism.
+        self.use_pp = self.parallel_config.pipeline_parallel_size > 1
+        self.is_first_pp_rank = get_pp_group().is_first_rank
+        self.is_last_pp_rank = get_pp_group().is_last_rank
+
+        # Size the UVA buffer pools to the max number of concurrent in-flight
+        # steps. Must run before any pooled buffer is constructed
+        set_default_max_concurrency(vllm_config.max_concurrent_batches)
+
+        # PP broadcast/recv helper. Runs the collective on a side stream.
+        self.pp_handler: PPHandler | None = None
+
+        # Persistent buffer for intermediate tensors (non-first PP ranks).
+        self.intermediate_tensors: IntermediateTensors | None = None
+
+        # Data parallelism.
+        self.dp_size = self.parallel_config.data_parallel_size
+        self.dp_rank = self.parallel_config.data_parallel_rank
+
+        # Dual batch overlap. Created in initialize_kv_cache(), once everything
+        # it runs the microbatched forward with exists.
+        self.ubatch_runner: UBatchRunner | None = None
+
+        # Detect EP all2all peer faults to prevent emitting corrupted output.
+        # Only meaningful for MoE + DP with an FT-capable all2all backend.
+        self.check_ep_fault = False
+        if self.dp_size > 1 and self.model_config.is_moe:
+            self.check_ep_fault = get_ep_all2all_manager().support_fault_tolerance
+
+        # Decode context parallelism.
+        self.dcp_size = self.parallel_config.decode_context_parallel_size
+        self.use_dcp = self.dcp_size > 1
+        self.dcp_rank = get_dcp_group().rank_in_group if self.use_dcp else 0
+        self.cp_interleave = self.parallel_config.cp_kv_cache_interleave_size
+
+        # Multimodal
+        self.mm_registry = MULTIMODAL_REGISTRY
+        self.supports_mm_inputs = self.model_config.supports_multimodal_inputs
+        self.uses_inputs_embeds = (
+            self.supports_mm_inputs or self.model_config.enable_prompt_embeds
+        )
+        self.encoder_cache = None
+        if self.supports_mm_inputs and self.is_first_pp_rank:
+            self.encoder_cache = EncoderCache()
+        self.ec_connector = get_ec_connector(vllm_config, self.encoder_cache)
+
+        # Speculative decoding.
+        self.speculator = None
+        self.use_aux_hidden_state_outputs = False
+        self.num_speculative_steps = vllm_config.num_speculative_tokens
+        if self.speculative_config is not None:
+            if self.is_last_pp_rank:
+                self.speculator = init_speculator(self.vllm_config, self.device)
+
+            if self.speculative_config.method in (
+                "eagle3",
+                "dflash",
+                "dspark",
+                "extract_hidden_states",
+            ):
+                # Drafting may require auxiliary hidden states from target model outputs
+                self.use_aux_hidden_state_outputs = True
+
+        # Draft tokens propagation - for spec-dec + struct outputs.
+        self.draft_tokens_handler = DraftTokensHandler(self.device)
+
+        self.pcp_manager: pcp.PCPManager | None = None
+
+        # Pooling models.
+        self.is_pooling_model = self.model_config.runner_type == "pooling"
+        self.pooling_runner: PoolingRunner | None = None
+
+        num_prefill_lookahead = max(1, vllm_config.num_prefill_lookahead_tokens)
+
+        self.step_timing = StepTimingCollector()
+
+        # General request states.
+        self.req_states = RequestState(
+            max_num_reqs=self.max_num_reqs,
+            max_model_len=self.max_model_len,
+            max_num_batched_tokens=self.max_num_tokens,
+            num_speculative_steps=self.num_speculative_steps,
+            vocab_size=self.vocab_size,
+            device=self.device,
+            num_prefill_lookahead=num_prefill_lookahead,
+        )
+        self.adaptive_verification: AdaptiveVerificationManager | None = None
+        self.input_buffers = InputBuffers(
+            max_num_reqs=self.max_num_reqs,
+            max_num_tokens=self.max_num_tokens,
+            device=self.device,
+        )
+        self.fast_prefill: FastPrefillHelper | None = None
+        if self.use_pp:
+            self.pp_handler = PPHandler(
+                max_num_reqs=self.max_num_reqs,
+                num_speculative_steps=self.num_speculative_steps,
+                device=self.device,
+            )
+
+        # Samplers and decode_query_len created in load_model() after
+        # model_state exists (num_new_sampled_tokens_per_step from ModelState).
+        self.sampler: Sampler | None = None
+        self.rejection_sampler: RejectionSampler | None = None
+        self.batch_sharder: BatchSharder | None = None
+        self.prompt_logprobs_worker: PromptLogprobsWorker | None = None
+        self.structured_outputs_worker: StructuredOutputsWorker | None = None
+        self.cudagraph_manager: ModelCudaGraphManager | None = None
+
+        # LoRA-related workers.
+        self.lora_state = LoraState(max_num_reqs=self.max_num_reqs)
+        self.lora_capture_cases = [0]
+        if self.lora_config:
+            self.lora_capture_cases = get_lora_capture_cases(
+                self.lora_config, self.compilation_config
+            )
+
+        # KV Connector if configured.
+        self.kv_connector: KVConnector = NO_OP_KV_CONNECTOR
+
+        # For transferring state from execute_model to subsequent sample_tokens call.
+        self.execute_model_state: ExecuteModelState | None = None
+
+        # Expert parallelism load balancer.
+        self.eplb = EPLBController(self.parallel_config, self.device)
+        # The AuxOutput Connector owns R3 capture, copying, and storage.
+        self.aux_output_connector: AuxOutputWorkerConnector | None = None
+
+        set_offloader(create_offloader(self.vllm_config.offload_config))
+
+    def update_max_model_len(self, max_model_len: int) -> None:
+        self.max_model_len = max_model_len
+        self.req_states.max_model_len = max_model_len
+
+    def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
+        tasks: list[SupportedTask] = []
+        if self.model_config.runner_type == "generate":
+            tasks.extend(self.model_state.get_supported_generation_tasks())
+        if self.is_pooling_model:
+            # Do not rely on pooling_runner here, since this information is needed
+            # on the first PP rank, while pooling_runner is only initialized
+            # on the last PP rank.
+            tasks.extend(PoolingRunner.get_supported_tasks(self.model))
+        return tuple(tasks)
+
+    def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None:
+        time_before_load = time.perf_counter()
+        if load_dummy_weights:
+            self.load_config.load_format = "dummy"
+        self.eplb.prepare_load()
+        eplb_models_added = False
+        with DeviceMemoryProfiler() as m:
+            model_loader = get_model_loader(self.vllm_config.load_config)
+            logger.info_once("Loading model from scratch...")
+
+            # Capture warmup providers selected while constructing the model.
+            with self.jit_warmup_registry.activate():
+                self.model = model_loader.load_model(
+                    vllm_config=self.vllm_config,
+                    model_config=self.vllm_config.model_config,
+                )
+            if self.lora_config:
+                self.model = self.load_lora_model(
+                    self.model, self.vllm_config, self.device
+                )
+
+            if self.use_aux_hidden_state_outputs:
+                assert self.speculative_config is not None
+                set_eagle3_aux_hidden_state_layers(self.model, self.speculative_config)
+                if self.use_pp:
+                    assert self.speculative_config.method is not None
+                    verify_supports_aux_hidden_states_over_pp(
+                        self.model, self.speculative_config.method
+                    )
+                    assert self.pp_handler is not None
+                    self.pp_handler.configure_aux_hidden_state_relay(self.model)
+            if isinstance(self.speculator, DraftModelSpeculator):
+                with use_workspace_lane(self._draft_workspace_lane):
+                    self.speculator.load_model(self.model)
+                    eplb_models_added = self.eplb.maybe_register_speculator(
+                        self.speculator, self.speculative_config, load_dummy_weights
+                    )
+        time_after_load = time.perf_counter()
+
+        self.model_memory_usage = m.consumed_memory
+        logger.info(
+            "Model loading took %s GiB memory and %.6f seconds",
+            format_gib(m.consumed_memory),
+            time_after_load - time_before_load,
+        )
+
+        # Initialize the components that require the model.
+        self.model_state = init_model_state(
+            self.vllm_config, self.model, self.encoder_cache, self.device
+        )
+
+        self.decode_query_len = (
+            self.num_speculative_steps
+            + self.model_state.num_new_sampled_tokens_per_step
+        )
+
+        if self.parallel_config.enable_batch_sharded_sampling:
+            if hasattr(self.model, "compute_logits_local"):
+                self.batch_sharder = BatchSharder(
+                    max_num_reqs=self.max_num_reqs,
+                    max_num_logits_per_req=self.decode_query_len,
+                    device=self.device,
+                )
+                logger.info("Batch-sharded sampling enabled.")
+            else:
+                logger.warning_once(
+                    "Disabling batch-sharded sampling: %s does not implement "
+                    "compute_logits_local",
+                    type(self.model).__name__,
+                )
+
+        # Initialize samplers. Model states may override via custom_sampler().
+        config_processors = self.model_config.logits_processors or ()
+        custom_logits_processors = build_custom_logits_processors(
+            self.vllm_config, self.req_states, self.is_pooling_model, config_processors
+        )
+        if self.is_last_pp_rank and not self.is_pooling_model:
+            # V2 bypasses TopKTopPSampler, which registers native warmups.
+            # ROCm only: on CUDA these warmups add ~2 min to every engine start.
+            if current_platform.is_rocm():
+                from vllm.v1.sample.ops.topk_topp_sampler import (
+                    register_top_k_top_p_warmups,
+                )
+
+                with self.jit_warmup_registry.activate():
+                    register_top_k_top_p_warmups()
+
+            sampler_kwargs: dict[str, Any] = {
+                "vllm_config": self.vllm_config,
+                "max_num_reqs": self.max_num_reqs,
+                "vocab_size": self.vocab_size,
+                "device": self.device,
+                "req_states": self.req_states,
+                "logprobs_mode": self.model_config.logprobs_mode,
+                "num_speculative_tokens": self.decode_query_len,
+                "use_fp64_gumbel": self.model_config.use_fp64_gumbel,
+                "enable_trace_replay": self.model_config.enable_trace_replay,
+                "return_sampling_mask": self.model_config.return_sampling_mask,
+                "custom_logits_processors": custom_logits_processors,
+            }
+            if self.vllm_config.watermark_config is None:
+                self.sampler = Sampler(**sampler_kwargs)
+            else:
+                wm_config = self.vllm_config.watermark_config
+                watermarker = create_watermarker(wm_config)
+                if self.speculative_config is not None:
+                    watermarker = create_speculative_target_watermarker(watermarker)
+                self.sampler = GPUWatermarkSampler(
+                    watermarker,
+                    deduplicate_contexts=wm_config.deduplicate_contexts,
+                    deduplicate_contexts_max_history=(
+                        wm_config.deduplicate_contexts_max_history
+                    ),
+                    **sampler_kwargs,
+                )
+            custom = self.model_state.custom_sampler(self.sampler)
+
+            if custom:
+                self.vllm_config._check_watermarking_unsupported(custom_sampler=True)
+                self.sampler, self.rejection_sampler = custom
+            elif self.speculative_config is not None:
+                self.rejection_sampler = RejectionSampler(
+                    self.sampler,
+                    self.speculative_config,
+                    self.device,
+                    watermark_key=speculative_target_watermark_key(
+                        self.vllm_config.watermark_config
+                    ),
+                )
+            self.prompt_logprobs_worker = PromptLogprobsWorker(
+                self.max_num_reqs, logprobs_mode=self.model_config.logprobs_mode
+            )
+            self.structured_outputs_worker = StructuredOutputsWorker(
+                max_num_logits=self.max_num_reqs * self.decode_query_len,
+                vocab_size=self.vocab_size,
+                device=self.device,
+                mask_stride=self.decode_query_len,
+                num_bonus_tokens=self.model_state.num_new_sampled_tokens_per_step,
+            )
+
+        if self.is_pooling_model and self.is_last_pp_rank:
+            self.pooling_runner = PoolingRunner(self.model, self.vllm_config)
+        eplb_models_added |= self.eplb.maybe_register_model(
+            self.model, self.model_config, load_dummy_weights
+        )
+        self.eplb.maybe_start_async_loop(eplb_models_added)
+
+        if not self.is_first_pp_rank:
+            # For non-first PP ranks, create intermediate tensors sized
+            # for the max capture size so they can be sliced per batch.
+            # Save as persistent member so runtime can copy received data
+            # into the same addresses that the CUDA graphs captured.
+            self.intermediate_tensors = self.model.make_empty_intermediate_tensors(
+                batch_size=self.max_num_tokens,
+                dtype=self.model_config.dtype,
+                device=self.device,
+            )
+
+        get_offloader().post_init()
+
+    def get_model(self) -> nn.Module:
+        return self.model
+
+    def get_draft_model(self) -> nn.Module | None:
+        speculator = self.speculator
+        if not isinstance(speculator, DraftModelSpeculator):
+            return None
+        return speculator.model
+
+    def reload_weights(self, *args, **kwargs) -> None:
+        # TODO(Wentao): Use full version instead of import when fully migrated to v2
+        from vllm.v1.worker.gpu_model_runner import GPUModelRunner as GPUModelRunnerV1
+
+        GPUModelRunnerV1.reload_weights(self, *args, **kwargs)  # type: ignore[arg-type]
+
+    def update_config(self, *args, **kwargs) -> None:
+        # TODO(Wentao): Use full version instead of import when fully migrated to v2
+        from vllm.v1.worker.gpu_model_runner import GPUModelRunner as GPUModelRunnerV1
+
+        GPUModelRunnerV1.update_config(self, *args, **kwargs)  # type: ignore[arg-type]
+
+        # v2 reads config via self.vllm_config (e.g. in load_model), so keep it
+        # in sync with the attributes the v1 helper just replaced.
+        self.vllm_config.model_config = self.model_config
+        self.vllm_config.load_config = self.load_config
+
+    @functools.cached_property
+    def main_stream(self) -> torch.cuda.Stream:
+        # Cache the default CUDA stream to avoid lookup overhead.
+        return torch.cuda.current_stream(self.device)
+
+    def get_encoder_timing_stats(self) -> dict[str, dict[str, float | int]]:
+        encoder_runner = getattr(self.model_state, "encoder_runner", None)
+        if encoder_runner is None:
+            return {}
+        return encoder_runner.get_encoder_timing_stats()
+
+    def get_kv_cache_spec(self):
+        return get_kv_cache_spec(self.vllm_config)
+
+    def initialize_kv_cache(
+        self,
+        kv_cache_config: KVCacheConfig,
+        is_profiling: bool = False,
+        kv_cache_allocation_context: AbstractContextManager | None = None,
+    ) -> None:
+        # GPUWorker finalizes the PD interleave before KV cache initialization.
+        self.cp_interleave = self.parallel_config.cp_kv_cache_interleave_size
+        kv_cache_config = deepcopy(kv_cache_config)
+        self.kv_cache_config = kv_cache_config
+
+        block_table_max_model_len = self.max_model_len
+        if self.is_encoder_decoder:
+            # Cross-attention block tables need to index encoder tokens, which
+            # can exceed the decoder's max_model_len.
+            block_table_max_model_len = max(
+                block_table_max_model_len,
+                self.scheduler_config.max_num_encoder_input_tokens,
+                getattr(self.model_config.hf_config, "max_source_positions", 0),
+            )
+
+        block_sizes = []
+        max_num_blocks_per_group = []
+        slot_mapping_enabled = []
+        for kv_cache_group in kv_cache_config.kv_cache_groups:
+            spec = kv_cache_group.kv_cache_spec
+            block_sizes.append(spec.block_size)
+            layer_spec = (
+                spec.first_spec if isinstance(spec, UniformTypeKVCacheSpecs) else spec
+            )
+            slot_mapping_enabled.append(layer_spec.uses_slot_mapping)
+            # Let each cache type account for CP. Attention KV is DCP-sharded,
+            # while Mamba/GDN recurrent state is replicated across DCP ranks.
+            max_num_blocks = spec.max_num_blocks_per_req(
+                self.vllm_config, block_table_max_model_len
+            )
+            # Preserve each cache type's alignment requirements after applying
+            # its topology-aware block-table width.
+            max_num_blocks = get_block_table_width(
+                max_num_blocks,
+                spec.block_size,
+                token_alignment=spec.block_table_token_alignment,
+            )
+            max_num_blocks_per_group.append(max_num_blocks)
+
+        target_attn_layer_names = None
+        if isinstance(self.speculator, DraftModelSpeculator):
+            # Adaptive verification validates target attention separately.
+            target_attn_layer_names = {
+                layer_name
+                for group in self.kv_cache_config.kv_cache_groups
+                for layer_name in group.layer_names
+            } - self.speculator.draft_attn_layer_names
+        draft_attn_layer_names = None
+        if isinstance(self.speculator, DraftModelSpeculator):
+            draft_attn_layer_names = self.speculator.draft_attn_layer_names
+        # Metadata builders select attention kernels that need JIT warmup.
+        with self.jit_warmup_registry.activate():
+            (
+                self.attn_groups,
+                attn_cg_support,
+                self.kernel_block_sizes,
+            ) = init_attn_backend(
+                self.kv_cache_config,
+                self.vllm_config,
+                self.device,
+                draft_layer_names=draft_attn_layer_names,
+            )
+        additional_attn_cg_support = self.model_state.get_additional_cg_support()
+        attn_cg_support = attn_cg_support.narrow(*additional_attn_cg_support)
+        # The speculator clears the flag at load time when the checkpoint has
+        # no confidence head, so it holds the effective value.
+        self.adaptive_verification = maybe_create_adaptive_verification_manager(
+            enable_adaptive_verification=getattr(
+                self.speculator, "enable_adaptive_verification", False
+            ),
+            attn_groups=self.attn_groups,
+            attn_cg_support=attn_cg_support,
+            req_states=self.req_states,
+            query_start_loc=self.input_buffers.query_start_loc,
+            num_bonus_tokens=self.model_state.num_new_sampled_tokens_per_step,
+            max_total_logits=get_max_chunk_logits(self.vocab_size),
+            vllm_config=self.vllm_config,
+            target_layer_names=target_attn_layer_names,
+            additional_attn_cg_support=additional_attn_cg_support,
+        )
+
+        self.block_tables = BlockTables(
+            block_sizes=block_sizes,
+            max_num_reqs=self.max_num_reqs,
+            max_num_batched_tokens=self.max_num_tokens,
+            max_num_blocks_per_group=max_num_blocks_per_group,
+            device=self.device,
+            kernel_block_sizes=self.kernel_block_sizes,
+            slot_mapping_enabled=slot_mapping_enabled,
+            cp_size=self.dcp_size,
+            cp_rank=self.dcp_rank,
+            cp_interleave=self.cp_interleave,
+        )
+        self.pcp_manager = pcp.maybe_build_pcp_manager(
+            self.vllm_config,
+            self.device,
+            self.supports_mm_inputs,
+            self.block_tables,
+            cls=self.pcp_manager_cls,
+        )
+        self.ubatch_runner = maybe_build_ubatch_runner(
+            self.vllm_config,
+            self.device,
+            self.model_state,
+            self.attn_groups,
+            self.kv_cache_config,
+            self.max_num_reqs,
+        )
+        if self.speculator is not None:
+            self.speculator.pcp_manager = self.pcp_manager
+        initialize_mamba_ssu_backend(
+            self.vllm_config.mamba_config,
+            self.kv_cache_config,
+            use_replayssm=self.cache_config.use_replayssm,
+        )
+        piecewise_capture_available = bool(
+            envs.VLLM_USE_BREAKABLE_CUDAGRAPH or has_compiled_submodule(self.model)
+        )
+        if self.adaptive_verification is not None:
+            self.compilation_config.cudagraph_mode = resolve_adaptive_cudagraph_mode(
+                self.compilation_config.cudagraph_mode,
+                piecewise_capture_available=piecewise_capture_available,
+            )
+        cudagraph_mode = self.compilation_config.resolve_cudagraph_mode_and_sizes(
+            attn_cg_support.min_cg_support,
+            attn_cg_support.min_cg_attn_backend,
+            self.decode_query_len,
+            use_v2_model_runner=True,
+            tensor_parallel_size=self.parallel_config.tensor_parallel_size,
+            kv_cache_config=self.kv_cache_config,
+            max_num_reqs=self.max_num_reqs,
+            is_profiling=is_profiling,
+            piecewise_capture_available=piecewise_capture_available,
+        )
+        self.cudagraph_manager = ModelCudaGraphManager(
+            self.vllm_config,
+            self.device,
+            cudagraph_mode,
+            decode_query_len=self.decode_query_len,
+            lora_capture_cases=self.lora_capture_cases,
+            varlen_decode=self.adaptive_verification is not None,
+            ubatch_runner=self.ubatch_runner,
+        )
+        if self.cache_config.kv_sharing_fast_prefill and self.pcp_manager is None:
+            self.fast_prefill = FastPrefillHelper(
+                self.cudagraph_manager, self.max_num_tokens
+            )
+        check_attention_cp_compatibility(self.vllm_config, target_attn_layer_names)
+        if isinstance(self.speculator, DraftModelSpeculator):
+            # HACK(woosuk)
+            self.speculator.set_attn(
+                self.model_state,
+                self.kv_cache_config,
+                self.block_tables,
+                self.input_buffers,
+                self.attn_groups,
+            )
+        if self.speculator is not None:
+            # After set_attn, so the speculator can size its cudagraph mode
+            # to its own attention support.
+            self.speculator.init_cudagraph_manager(cudagraph_mode)
+
+        # Capture warmup providers that depend on allocated KV-cache strides.
+        with self.jit_warmup_registry.activate():
+            kv_caches_dict = init_kv_cache(
+                self.compilation_config.static_forward_context,
+                self.kv_cache_config,
+                self.device,
+                self.kernel_block_sizes,
+                self.vllm_config,
+                kv_cache_allocation_context=kv_cache_allocation_context,
+                block_tables=self.block_tables,
+            )
+        self.kv_caches = [
+            cache for cache in kv_caches_dict.values() if cache.device == self.device
+        ]
+        if is_profiling:
+            self.kv_connector = NO_OP_KV_CONNECTOR
+        else:
+            self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict)
+
+            # AuxOutput connector requires resolved kv_cache_config.
+            if self.vllm_config.aux_output_config.enabled:
+                self.aux_output_connector = get_aux_output_connector(
+                    self.model, self.vllm_config, kv_cache_config
+                )
+
+    def _init_kv_zero_meta(self) -> None:
+        """Build KV-block zeroing metadata; invoked from gpu_worker."""
+        self.kv_block_zeroer = KVBlockZeroer(
+            self.device,
+            attn_groups_iter=(g for groups in self.attn_groups for g in groups),
+            kernel_block_sizes=self.kernel_block_sizes,
+            static_forward_context=self.compilation_config.static_forward_context,
+            num_blocks=self.kv_cache_config.num_blocks,
+        )
+
+    @torch.inference_mode()
+    @step_eplb_after(is_dummy=True)
+    def _dummy_run(
+        self,
+        num_tokens: int,
+        *args,
+        skip_attn: bool = False,
+        uniform_decode: bool = False,
+        context_len: int = 0,
+        skip_eplb: bool = False,
+        is_profile: bool = False,
+        valid_dummy_state_slots: bool = False,
+        **kwargs,
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+        if skip_attn and not is_profile:
+            raise ValueError(
+                "skip_attn must only be True for initial memory profiling."
+            )
+
+        # Create a dummy scheduler output.
+        num_reqs = min(num_tokens, self.max_num_reqs)
+        if uniform_decode:
+            # HACK(lucas): for now since the worker is shared between MRV1 and MRV2,
+            # and for spec-decode with MTP we want to make sure the dummy runs use
+            # 1+num_speculative_tokens we use max here, this will likely be eventually
+            # changed in the worker: https://github.com/vllm-project/vllm/pull/35243
+            num_tokens = max(num_tokens, self.decode_query_len)
+            num_reqs = num_tokens // self.decode_query_len
+            assert num_tokens % self.decode_query_len == 0
+        elif self.speculator is not None:
+            num_reqs = min(
+                num_reqs, self.max_num_tokens // self.speculator.num_query_per_req
+            )
+        # Distribute the remainder evenly so no dummy request exceeds
+        # ceil(num_tokens / num_reqs) <= max_model_len tokens.
+        num_tokens_per_request = [
+            num_tokens // num_reqs + (i >= num_reqs - num_tokens % num_reqs)
+            for i in range(num_reqs)
+        ]
+
+        assert sum(num_tokens_per_request) == num_tokens
+        num_scheduled_tokens = {
+            f"_dummy_req_{i}": n for i, n in enumerate(num_tokens_per_request)
+        }
+        dummy_scheduler_output = SchedulerOutput.make_empty()
+        dummy_scheduler_output.total_num_scheduled_tokens = num_tokens
+        dummy_scheduler_output.num_scheduled_tokens = num_scheduled_tokens
+
+        # Disable any use of KVConnector for dummy runs.
+        self.kv_connector.set_disabled(True)
+
+        # Get the intermediate tensors for the dummy run.
+        intermediate_tensors = None
+        if not self.is_first_pp_rank:
+            assert self.intermediate_tensors is not None
+            intermediate_tensors = self.intermediate_tensors[:num_tokens]
+
+        max_loras = self.lora_config.max_loras if self.lora_config is not None else 0
+        with self.maybe_dummy_run_with_lora(
+            self.lora_config,
+            num_scheduled_tokens=np.array(num_tokens_per_request, dtype=np.int32),
+            num_sampled_tokens=None,
+            remove_lora=True,
+            num_active_loras=max_loras,
+        ):
+            # Execute the model.
+            self.execute_model(
+                dummy_scheduler_output,
+                intermediate_tensors=intermediate_tensors,
+                dummy_run=True,
+                skip_attn_for_dummy_run=skip_attn,
+                is_profile=is_profile,
+                context_len=context_len,
+                valid_dummy_state_slots=valid_dummy_state_slots,
+            )
+        self.kv_connector.set_disabled(False)
+
+        # Non-last PP ranks don't produce output for sampling.
+        if not self.is_last_pp_rank:
+            return None, None
+
+        assert self.execute_model_state is not None
+        input_batch = self.execute_model_state.input_batch
+        attn_metadata = self.execute_model_state.attn_metadata
+        slot_mappings_by_layer = self.execute_model_state.slot_mappings_by_layer
+        hidden_states = self.execute_model_state.hidden_states
+        aux_hidden_states = self.execute_model_state.aux_hidden_states
+        dp_sync = self.execute_model_state.dp_sync
+        self.execute_model_state = None
+
+        self.step_timing.forward_end()
+
+        # dummy run the eagle speculator's propose to ensure DP/EP sync.
+        if self.speculator is not None:
+            assert self.sampler is not None
+            self.step_timing.drafter_start()
+            mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None
+            if self.speculator.supports_mm_inputs:
+                all_false = torch.zeros(
+                    input_batch.num_tokens, dtype=torch.bool, device="cpu"
+                )
+                mm_inputs = [], all_false
+
+            # Let the target override the hidden state fed to the drafter
+            # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The
+            # target returns a persistent buffer sized at max_num_batched_tokens;
+            # slice to the active token count that propose() expects.
+            spec_hidden_states = hidden_states
+            if hasattr(self.model, "get_mtp_target_hidden_states"):
+                pre_hc_hidden_states = self.model.get_mtp_target_hidden_states()
+                spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]]  # type: ignore[union-attr]
+            if isinstance(self.sampler, GPUWatermarkSampler):
+                self.speculator.prepare_watermarking(
+                    self.sampler._get_contexts(input_batch.idx_mapping),
+                    self.sampler.watermarking.gpu[input_batch.idx_mapping],
+                )
+            with use_workspace_lane(self._draft_workspace_lane):
+                self.speculator.propose(
+                    input_batch=input_batch,
+                    attn_metadata=attn_metadata,
+                    slot_mappings=slot_mappings_by_layer,
+                    last_hidden_states=spec_hidden_states,
+                    aux_hidden_states=aux_hidden_states,
+                    num_sampled=torch.ones(
+                        input_batch.num_reqs, dtype=torch.int32, device=self.device
+                    ),
+                    num_rejected=torch.zeros(
+                        input_batch.num_reqs, dtype=torch.int32, device=self.device
+                    ),
+                    last_sampled=self.req_states.last_sampled_tokens,
+                    next_prefill_tokens=self.req_states.next_prefill_tokens,
+                    temperature=self.sampler.sampling_states.temperature.gpu,
+                    seeds=self.sampler.sampling_states.seeds.gpu,
+                    dp_sync=dp_sync,
+                    dummy_run=True,
+                    skip_attn_for_dummy_run=skip_attn,
+                    mm_inputs=mm_inputs,
+                    is_profile=is_profile,
+                )
+            self.step_timing.drafter_end()
+
+        assert hidden_states is not None  # Last PP rank always has hidden_states
+        sample_hidden_states = hidden_states[input_batch.logits_indices]
+        return hidden_states, sample_hidden_states
+
+    @torch.inference_mode()
+    def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None:
+        num_reqs = hidden_states.shape[0]
+        logits = self.model.compute_logits(hidden_states)
+        dummy_input_batch = InputBatch.make_dummy(
+            num_reqs, num_reqs, self.input_buffers
+        )
+
+        # NOTE(woosuk): During the initial memory profiling, the sampler may skip
+        # top_k, top_p, and logprobs, using less GPU memory than what is possible
+        # during actual execution.
+        assert self.sampler is not None
+        self.sampler(logits, dummy_input_batch)
+
+    @torch.inference_mode()
+    def _dummy_pooler_run(self, hidden_states: torch.Tensor) -> None:
+        assert self.pooling_runner is not None
+        self.pooling_runner.dummy_pooler_run(hidden_states)
+
+    @torch.inference_mode()
+    def profile_run(self) -> None:
+        if self.supports_mm_inputs and self.is_first_pp_rank:
+            mm_config = self.model_config.multimodal_config
+            if mm_config is not None and not mm_config.skip_mm_profiling:
+                mm_budget = MultiModalBudget(
+                    self.vllm_config, self.mm_registry, enable_cache=False
+                )
+                dummy_mm_inputs = mm_budget.get_dummy_encoder_profile_inputs()
+                self.model_state.encoder_runner.profile_encoder_cache(
+                    dummy_mm_inputs, mm_budget
+                )
+
+        hidden_states, sample_hidden_states = self._dummy_run(
+            self.max_num_tokens, skip_attn=True, is_profile=True
+        )
+
+        # Only run sampler/pooler on last PP rank (non-last ranks return None).
+        if self.is_last_pp_rank:
+            assert sample_hidden_states is not None
+            if self.pooling_runner is None:
+                self._dummy_sampler_run(sample_hidden_states)
+            else:
+                self._dummy_pooler_run(hidden_states)
+
+        torch.accelerator.synchronize()
+        del hidden_states, sample_hidden_states
+        self.reset_encoder_cache()
+        gc.collect()
+
+    def reset_mm_cache(self) -> None:
+        if self.encoder_cache is not None:
+            self.encoder_cache.reset_mm_cache()
+
+    def reset_encoder_cache(self) -> None:
+        if self.encoder_cache is not None:
+            self.encoder_cache.reset_encoder_cache()
+        if self.pooling_runner is not None:
+            self.pooling_runner.clear()
+
+    @torch.inference_mode()
+    def profile_cudagraph_memory(self) -> int:
+        """Estimate the GPU memory required to capture CUDA graphs."""
+        return _profile_cudagraph_memory(self)
+
+    def needs_cudagraph_capture(self) -> bool:
+        """Whether capture_model() has any CUDA graphs left to capture."""
+        return (
+            self.model_state.supports_mm_inputs
+            and self.model_state.encoder_runner.has_cudagraph()
+        ) or (
+            self.cudagraph_manager is not None
+            and self.cudagraph_manager.needs_capture()
+        )
+
+    @torch.inference_mode()
+    def capture_model(self, *, profile_only: bool = False) -> int:
+        assert self.cudagraph_manager is not None
+        capture_encoder = (
+            self.model_state.supports_mm_inputs
+            and self.model_state.encoder_runner.has_cudagraph()
+        )
+        capture_decoder = self.cudagraph_manager.needs_capture()
+        if not capture_encoder and not capture_decoder:
+            logger.warning(
+                "Skipping encoder and decoder CUDA graph capture. To enable "
+                "encoder capture, ensure `cudagraph_mm_encoder` is enabled; "
+                "to enable decoder capture, ensure `cudagraph_mode` is not `NONE`."
+            )
+            return 0
+
+        compilation_counter.num_gpu_runner_capture_triggers += 1
+
+        start_time = time.perf_counter()
+        with freeze_gc_for_cudagraph_capture():
+            torch.accelerator.empty_cache()
+            start_free_gpu_memory = torch.accelerator.get_memory_info()[0]
+
+            with self.maybe_setup_dummy_loras(self.lora_config):
+                if capture_encoder:
+                    self.model_state.encoder_runner.capture()
+
+                if capture_decoder:
+                    input_buffers = self.input_buffers
+                    if self.pcp_manager is not None:
+                        input_buffers = self.pcp_manager.input_buffers
+                    self.cudagraph_manager.capture(
+                        self.model,
+                        self.model_state,
+                        input_buffers,
+                        self.intermediate_tensors,
+                        self.block_tables,
+                        self.attn_groups,
+                        self.kv_cache_config,
+                        pcp_manager=self.pcp_manager,
+                        has_lora=self.lora_config is not None,
+                        use_aux_hidden_state_outputs=self.use_aux_hidden_state_outputs,
+                        lora_capture_hook=create_lora_capture_hook(
+                            self.lora_config, self
+                        ),
+                    )
+                    if self.speculator is not None:
+                        with use_workspace_lane(self._draft_workspace_lane):
+                            self.speculator.capture()
+                    if self.adaptive_verification is not None:
+                        with self.step_timing.collect() as timings:
+                            for batch in self.adaptive_verification.batches_to_profile(
+                                self.cudagraph_manager.captured_token_counts()
+                            ):
+                                self._dummy_run(**batch)
+                        self.adaptive_verification.set_initial_cost_curves(timings)
+                    self.kv_connector.reset_capture_state()
+
+            end_free_gpu_memory = torch.accelerator.get_memory_info()[0]
+
+        if not profile_only:
+            # Lock workspace to prevent resizing during execution. A resize after
+            # capture frees the static cuda graph buffer.
+            lock_workspace()
+
+        end_time = time.perf_counter()
+        elapsed_time = end_time - start_time
+        cuda_graph_size = start_free_gpu_memory - end_free_gpu_memory
+        # This usually takes 5~20 seconds.
+        logger.info(
+            "Graph capturing finished in %.0f secs, took %.2f GiB",
+            elapsed_time,
+            cuda_graph_size / (1 << 30),
+        )
+        return cuda_graph_size
+
+    def _remove_request(self, req_id: str) -> bool:
+        # Call model_state.remove_request *before* req_states.remove_request
+        # so the model_state can still look up the slot index.
+        self.model_state.remove_request(req_id)
+        req_idx = self.req_states.remove_request(req_id)
+        if req_idx is None:
+            return False
+        if self.pooling_runner is not None:
+            self.pooling_runner.remove_request(req_idx)
+        if self.pp_handler is not None:
+            self.pp_handler.on_req_idx_freed(req_idx)
+        if self.encoder_cache is not None:
+            self.encoder_cache.remove_request(req_id)
+        if self.prompt_logprobs_worker is not None:
+            self.prompt_logprobs_worker.remove_request(req_id)
+        self.lora_state.remove_request(req_id)
+        return True
+
+    def finish_requests(self, scheduler_output: SchedulerOutput) -> None:
+        finished_req_ids = scheduler_output.finished_req_ids
+        if self.pooling_runner is not None:
+            # Preempted docs keep their query-use reservation until rescheduled.
+            self.pooling_runner.on_requests_finished(finished_req_ids)
+        preempted_req_ids = scheduler_output.preempted_req_ids
+        if preempted_req_ids:
+            finished_req_ids = finished_req_ids.union(preempted_req_ids)
+        # Sorted so every TP rank frees request slots in the same order.
+        # Features like batch-sharded sampling derive rank request ownership
+        # from the slot index.
+        for req_id in sorted(finished_req_ids):
+            self._remove_request(req_id)
+
+    def free_states(self, scheduler_output: SchedulerOutput) -> None:
+        if self.encoder_cache is not None:
+            for mm_hash in scheduler_output.free_encoder_mm_hashes:
+                self.encoder_cache.free_encoder_cache(mm_hash)
+
+    def update_pp_decode_requests(self):
+        # For non-last PP ranks, update decode requests with sampler output from
+        # the prior step in which they were scheduled (pp_size steps ago).
+        if self.pp_handler is not None:
+            outputs = self.pp_handler.get_prev_sampled_outputs(
+                self.req_states.draft_tokens
+            )
+            if outputs is not None:
+                self.postprocess_sampled(**outputs)
+
+    def add_requests(self, scheduler_output: SchedulerOutput) -> None:
+        for new_req_data in scheduler_output.scheduled_new_reqs:
+            assert new_req_data.prefill_token_ids is not None
+            req_id = new_req_data.req_id
+
+            # Streaming input update: request already exists from a prior
+            # chunk. Remove old state so it can be cleanly re-added below
+            # with the updated prompt_token_ids and mm_features.
+            self._remove_request(req_id)
+
+            prompt_len = new_req_data.prompt_len
+            sampling_params = new_req_data.sampling_params
+            self.req_states.add_request(
+                req_id=req_id,
+                prompt_len=prompt_len,
+                all_token_ids=new_req_data.prefill_token_ids,
+                num_computed_tokens=new_req_data.num_computed_tokens,
+                max_tokens=sampling_params.max_tokens if sampling_params else 1,  # type: ignore[arg-type]
+            )
+            req_index = self.req_states.req_id_to_index[req_id]
+            if self.adaptive_verification is not None:
+                self.adaptive_verification.add_request(req_index)
+
+            if self.pooling_runner is not None:
+                assert new_req_data.pooling_params is not None
+                self.pooling_runner.add_request(
+                    req_id,
+                    req_index,
+                    new_req_data.pooling_params,
+                    new_req_data.prompt_token_ids,
+                )
+
+            if self.encoder_cache is not None:
+                self.encoder_cache.add_request(req_id, new_req_data.mm_features)
+
+            self.model_state.add_request(req_index, new_req_data)
+            self.block_tables.append_block_ids(
+                req_index, new_req_data.block_ids, overwrite=True
+            )
+            self.lora_state.add_request(req_id, req_index, new_req_data.lora_request)
+
+            if self.is_last_pp_rank and new_req_data.sampling_params is not None:
+                assert self.sampler is not None
+                self.sampler.add_request(req_index, new_req_data.sampling_params)
+                assert self.prompt_logprobs_worker is not None
+                self.prompt_logprobs_worker.add_request(
+                    req_id, req_index, new_req_data.sampling_params
+                )
+
+        if scheduler_output.scheduled_new_reqs:
+            self.req_states.apply_staged_writes()
+            self.model_state.apply_staged_writes()
+        if self.sampler is not None:
+            self.sampler.apply_staged_writes()
+
+    def update_requests(self, scheduler_output: SchedulerOutput) -> None:
+        # Add new blocks and update num_computed_tokens for the existing requests.
+        reqs = scheduler_output.scheduled_cached_reqs
+        table_updates = scheduler_output.block_table_updates or {}
+        for req_id, block_ids in table_updates.items():
+            if (idx := self.req_states.req_id_to_index.get(req_id)) is not None:
+                self.block_tables.append_block_ids(idx, block_ids, overwrite=True)
+        num_computed_tokens_np = self.req_states.num_computed_tokens_np
+        for req_id, num_computed_tokens, req_new_block_ids in zip(
+            reqs.req_ids, reqs.num_computed_tokens, reqs.new_block_ids
+        ):
+            req_index = self.req_states.req_id_to_index[req_id]
+            num_computed_tokens_np[req_index] = num_computed_tokens
+            if req_new_block_ids is not None and req_id not in table_updates:
+                self.block_tables.append_block_ids(
+                    req_index, req_new_block_ids, overwrite=False
+                )
+
+        # Update CPU num_computed_prefill_tokens.
+        np.minimum(
+            self.req_states.num_computed_tokens_np,
+            self.req_states.prefill_len.np,
+            out=self.req_states.num_computed_prefill_tokens,
+        )
+
+        # Zero GPU memory for freshly allocated cache blocks to prevent
+        # stale NaN/data from corrupting attention or SSM computation.
+        if scheduler_output.new_block_ids_to_zero:
+            assert self.kv_block_zeroer is not None
+            self.kv_block_zeroer.zero_block_ids(scheduler_output.new_block_ids_to_zero)
+
+        # Apply copy-on-write block copies for partial prefix-cache hits, after
+        # zeroing new blocks and before the forward pass reads them.
+        if scheduler_output.kv_cache_block_copies:
+            copy_kv_cache_blocks_inplace(
+                self.kv_caches,
+                self.kv_cache_config.num_blocks,
+                scheduler_output.kv_cache_block_copies,
+            )
+
+    def gather_batch_req_state(
+        self, scheduler_output: SchedulerOutput, dummy_run: bool
+    ) -> tuple["BatchReqState | None", int | None]:
+        """Gather CPU request state for the scheduled batch, in batch order.
+        Returns (batch_state, uniform_decode_token_count)
+        """
+        num_tokens_per_req = scheduler_output.num_scheduled_tokens
+        num_reqs = len(num_tokens_per_req)
+        num_toks = scheduler_output.total_num_scheduled_tokens
+        max_query_len = max(scheduler_output.num_scheduled_tokens.values())
+
+        if dummy_run:
+            # Dummy batches are uniform by construction.
+            return None, get_uniform_decode_token_count(
+                num_reqs, num_toks, max_query_len, has_prefill=False
+            )
+
+        draft_tokens = scheduler_output.scheduled_spec_decode_tokens
+        # batch_idx -> req_id
+        req_ids = sort_batch_req_ids(
+            num_tokens_per_req, draft_tokens, self.decode_query_len
+        )
+
+        numtoks_iter = map(num_tokens_per_req.__getitem__, req_ids)
+        num_scheduled_tokens = np.fromiter(numtoks_iter, dtype=np.int32, count=num_reqs)
+
+        idx_mapping_iter = map(self.req_states.req_id_to_index.__getitem__, req_ids)
+        idx_mapping_np = np.fromiter(idx_mapping_iter, dtype=np.intp, count=num_reqs)
+        prefill_len_np = self.req_states.prefill_len.np[idx_mapping_np]
+        num_computed_prefill_tokens_np = self.req_states.num_computed_prefill_tokens[
+            idx_mapping_np
+        ]
+        is_prefilling_np = num_computed_prefill_tokens_np < prefill_len_np
+
+        if self.adaptive_verification is not None and draft_tokens:
+            num_toks = self.adaptive_verification.get_num_tokens(
+                num_tokens_per_req, draft_tokens
+            )
+
+        batch_state = BatchReqState(
+            req_ids=req_ids,
+            num_scheduled_tokens=num_scheduled_tokens,
+            num_tokens=num_toks,
+            idx_mapping_np=idx_mapping_np,
+            prefill_len_np=prefill_len_np,
+            num_computed_prefill_tokens_np=num_computed_prefill_tokens_np,
+            is_prefilling_np=is_prefilling_np,
+            has_prefill=bool(is_prefilling_np.any()),
+        )
+        return batch_state, get_uniform_decode_token_count(
+            num_reqs, num_toks, max_query_len, batch_state.has_prefill
+        )
+
+    def prepare_inputs(
+        self,
+        scheduler_output: SchedulerOutput,
+        batch_req_state: "BatchReqState",
+        batch_desc: BatchExecutionDescriptor,
+        num_active_loras: int,
+    ) -> InputBatch:
+        num_tokens = batch_req_state.num_tokens
+        num_tokens_after_padding = max(num_tokens, batch_desc.num_tokens)
+        assert num_tokens > 0
+        is_padding = self.input_buffers.is_padding[:num_tokens_after_padding]
+        if envs.VLLM_MOE_SKIP_PADDING:
+            is_padding[:num_tokens].fill_(False)
+            is_padding[num_tokens:].fill_(True)
+
+        req_ids = batch_req_state.req_ids
+        num_scheduled_tokens_np = batch_req_state.num_scheduled_tokens
+        idx_mapping_np = batch_req_state.idx_mapping_np
+        idx_mapping = async_tensor_h2d(
+            idx_mapping_np, device=self.device, dtype=torch.int32
+        )
+        num_reqs = len(req_ids)
+
+        # Get the number of draft tokens for each request.
+        draft_tokens = scheduler_output.scheduled_spec_decode_tokens
+        num_draft_tokens_per_req = None
+        if not draft_tokens and self.model_state.num_new_sampled_tokens_per_step == 1:
+            # No draft token scheduled (common case).
+            total_num_draft_tokens = 0
+            total_num_logits = num_reqs
+            cu_num_logits_np = np.arange(num_reqs + 1, dtype=np.int32)
+            cu_num_logits = torch.arange(
+                num_reqs + 1, device=self.device, dtype=torch.int32
+            )
+            expanded_idx_mapping = idx_mapping
+            expanded_local_pos = torch.zeros(
+                num_reqs, dtype=torch.int32, device=self.device
+            )
+        else:
+            num_draft_tokens_per_req = np.fromiter(
+                (len(draft_tokens.get(req_id, ())) for req_id in req_ids),
+                dtype=np.int32,
+                count=num_reqs,
+            )
+            num_bonus_tokens = self.model_state.num_new_sampled_tokens_per_step
+            total_num_draft_tokens = int(num_draft_tokens_per_req.sum())
+            total_num_logits = num_reqs * num_bonus_tokens + total_num_draft_tokens
+            num_logits = num_draft_tokens_per_req + num_bonus_tokens
+            # combine_sampled_and_draft_tokens places a request's logits rows
+            # at [query_end - num_logits, query_end). Fewer query rows than
+            # that would silently select the preceding request's hidden states.
+            assert (num_scheduled_tokens_np >= num_logits).all()
+            cu_num_logits_np = np.empty(num_reqs + 1, dtype=np.int32)
+            cu_num_logits_np[0] = 0
+            np.cumsum(num_logits, out=cu_num_logits_np[1:])
+            cu_num_logits = async_tensor_h2d(cu_num_logits_np, device=self.device)
+
+        # The general branch also serves a no-draft batch with k != 1, and
+        # compact_batch's budget is only primed when drafts are scheduled.
+        adaptive_verification = self.adaptive_verification if draft_tokens else None
+        num_scheduled_tokens_upper_bound = num_scheduled_tokens_np
+        if adaptive_verification is not None:
+            # num_scheduled_tokens represents the draft budget evenly distributed across
+            # all verification requests, `reallocate_drafts` will unevenly assign the
+            # draft budget to requests on the GPU side only.
+            num_scheduled_tokens_np, cu_num_logits_np = (
+                adaptive_verification.compact_batch(
+                    num_draft_tokens_per_req,
+                    num_scheduled_tokens_np,
+                    cu_num_logits_np,
+                )
+            )
+
+        # Get query_start_loc.
+        # num_reqs_padded is None for PIECEWISE graphs (no request padding needed)
+        num_reqs_padded = batch_desc.num_reqs or num_reqs
+        query_start_loc_np = np.empty(self.max_num_reqs + 1, dtype=np.int32)
+        query_start_loc_np[0] = 0
+        np.cumsum(num_scheduled_tokens_np, out=query_start_loc_np[1 : num_reqs + 1])
+        # Pad for full CUDA graph mode.
+        # Some attention backends like FA3 require query_start_loc to be non-decreasing.
+        query_start_loc_np[num_reqs + 1 :] = num_tokens
+        query_start_loc = self.input_buffers.query_start_loc
+        async_tensor_h2d(query_start_loc_np, out=query_start_loc)
+        if adaptive_verification is not None:
+            cu_num_logits, query_start_loc, total_num_draft_tokens = (
+                adaptive_verification.reallocate_drafts(req_ids, idx_mapping)
+            )
+            total_num_logits = num_reqs * num_bonus_tokens + total_num_draft_tokens
+        if num_draft_tokens_per_req is not None:
+            expanded_idx_mapping, expanded_local_pos = expand_idx_mapping(
+                idx_mapping, total_num_logits, cu_num_logits, self.decode_query_len
+            )
+        query_start_loc_np = query_start_loc_np[: num_reqs_padded + 1]
+        query_start_loc = query_start_loc[: num_reqs_padded + 1]
+
+        # Get prefill tokens if any.
+        if batch_req_state.has_prefill:
+            prepare_prefill_inputs(
+                self.input_buffers.input_ids,
+                self.req_states.next_prefill_tokens,
+                idx_mapping,
+                query_start_loc,
+                self.req_states.all_token_ids.gpu,
+                self.req_states.prefill_len.gpu,
+                self.req_states.num_computed_tokens.gpu,
+            )
+
+        # Prepare positions and seq_lens.
+        prepare_pos_seq_lens(
+            idx_mapping,
+            query_start_loc,
+            self.req_states.num_computed_tokens.gpu,
+            self.input_buffers.positions,
+            self.input_buffers.seq_lens,
+        )
+        seq_lens = self.input_buffers.seq_lens[:num_reqs_padded]
+
+        # Some input token ids are directly read from the last sampled tokens
+        # and draft tokens. Also, get the logits indices to sample tokens from.
+        logits_indices = combine_sampled_and_draft_tokens(
+            self.input_buffers.input_ids,
+            idx_mapping,
+            self.req_states.last_sampled_tokens,
+            query_start_loc,
+            seq_lens,
+            self.req_states.prefill_len.gpu,
+            self.req_states.draft_tokens,
+            cu_num_logits,
+            total_num_logits,
+            self.model_state.num_new_sampled_tokens_per_step,
+        )
+        if draft_trunc.ENABLED and envs.VLLM_MOE_SKIP_PADDING:
+            est = getattr(self.speculator, "acceptance_estimator", None)
+            draft_trunc.mark_dead(is_padding, query_start_loc, cu_num_logits, num_reqs,
+                                  self.req_states.draft_tokens.shape[-1],
+                                  est.predictions if est is not None else None, idx_mapping)
+
+        fast_prefill = None
+        if self.fast_prefill is not None:
+            fast_prefill = self.fast_prefill.prepare(
+                logits_indices,
+                num_reqs,
+                cu_num_logits_np,
+                batch_req_state.has_prefill,
+                batch_desc,
+                num_active_loras=num_active_loras,
+            )
+
+        # CPU upper bound on seq_lens; padded entries left at zero.
+        num_computed_tokens_np = self.req_states.num_computed_tokens_np[idx_mapping_np]
+        seq_lens_cpu_upper_bound_np = np.zeros(num_reqs_padded, dtype=np.int32)
+        np.add(
+            num_computed_tokens_np,
+            num_scheduled_tokens_upper_bound,
+            out=seq_lens_cpu_upper_bound_np[:num_reqs],
+        )
+        seq_lens_cpu_upper_bound = torch.from_numpy(seq_lens_cpu_upper_bound_np)
+
+        prompt_lens = None
+        if self.model_config.rswa_window is not None:
+            # prompt_lens is only used in R-SWA case.
+            prompt_lens = self.req_states.prompt_len.gpu[idx_mapping]
+
+        input_batch = InputBatch(
+            req_ids=req_ids,
+            num_reqs=num_reqs,
+            num_reqs_after_padding=num_reqs_padded,
+            idx_mapping=idx_mapping,
+            idx_mapping_np=idx_mapping_np,
+            expanded_idx_mapping=expanded_idx_mapping,
+            expanded_local_pos=expanded_local_pos,
+            num_scheduled_tokens=num_scheduled_tokens_upper_bound,
+            num_tokens=num_tokens,
+            num_tokens_after_padding=num_tokens_after_padding,
+            num_draft_tokens=total_num_draft_tokens,
+            num_draft_tokens_per_req=num_draft_tokens_per_req,
+            query_start_loc=query_start_loc,
+            query_start_loc_np=query_start_loc_np,
+            seq_lens=seq_lens,
+            seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
+            dcp_local_seq_lens=None,
+            num_computed_tokens_np=num_computed_tokens_np,
+            prefill_len_np=batch_req_state.prefill_len_np,
+            num_computed_prefill_tokens_np=batch_req_state.num_computed_prefill_tokens_np,
+            is_prefilling_np=batch_req_state.is_prefilling_np,
+            has_prefill=batch_req_state.has_prefill,
+            input_ids=self.input_buffers.input_ids[:num_tokens_after_padding],
+            positions=self.input_buffers.positions[:num_tokens_after_padding],
+            is_padding=is_padding,
+            logits_indices=logits_indices,
+            cu_num_logits=cu_num_logits,
+            cu_num_logits_np=cu_num_logits_np,
+            has_structured_output_reqs=scheduler_output.has_structured_output_requests,
+            prompt_lens=prompt_lens,
+            fast_prefill=fast_prefill,
+            max_query_len=(
+                int(num_scheduled_tokens_upper_bound.max())
+                if adaptive_verification is not None
+                else None
+            ),
+        )
+        if self.pcp_manager is not None:
+            input_batch = self.pcp_manager.partition_batch(input_batch, batch_desc)
+        return input_batch
+
+    def prepare_attn(
+        self, input_batch: InputBatch
+    ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]:
+        if self.pcp_manager is not None:
+            return self.pcp_manager.prepare_attn(input_batch)
+
+        # Block tables: num_kv_cache_groups x [num_reqs_padded, max_num_blocks].
+        block_tables = self.block_tables.gather_block_tables(
+            input_batch.idx_mapping, num_reqs_padded=input_batch.num_reqs_after_padding
+        )
+        # Slot mappings: [num_kv_cache_groups, num_tokens_padded].
+        # Kernel pads beyond num_tokens with PAD_SLOT_ID.
+        slot_mappings = self.block_tables.compute_slot_mappings(
+            input_batch.idx_mapping,
+            input_batch.query_start_loc,
+            input_batch.positions,
+            num_tokens_padded=input_batch.num_tokens_after_padding,
+        )
+        return block_tables, slot_mappings
+
+    def prepare_dummy_attn(
+        self, input_batch: InputBatch, valid_state_slots: bool = False
+    ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]:
+        block_table_provider = self.pcp_manager or self.block_tables
+        block_tables = block_table_provider.get_dummy_block_tables(input_batch.num_reqs)
+        if valid_state_slots:
+            state_slots = torch.arange(
+                1, input_batch.num_reqs + 1, dtype=torch.int32, device=self.device
+            )
+            for block_table in block_tables:
+                block_table[:, 0].copy_(state_slots)
+        slot_mappings = block_table_provider.get_dummy_slot_mappings(
+            input_batch.num_tokens
+        )
+        return block_tables, slot_mappings
+
+    def sample(
+        self,
+        hidden_states: torch.Tensor,
+        input_batch: InputBatch,
+        grammar_output: GrammarOutput | None,
+    ) -> tuple[SamplerOutput, torch.Tensor, torch.Tensor]:
+        shard_metadata = None
+        global_input_batch = input_batch
+        if self.batch_sharder is not None:
+            # Shard the inputs along the batch dimension to sample in parallel
+            # across TP ranks.
+            input_batch, sorted_logits_indices, grammar_output, shard_metadata = (
+                self.batch_sharder.shard_sampler_inputs(input_batch, grammar_output)
+            )
+            # The hidden states must be gathered in rank-owner-sorted order
+            # before computing the partial-vocab logits, so that the all-to-all
+            # produces full-vocab logits for just the locally-owned requests.
+            sample_hidden_states = hidden_states[sorted_logits_indices]
+            local_logits = self.model.compute_logits_local(sample_hidden_states)
+            logits = all_to_all_logits(local_logits, shard_metadata)
+            logits = logits[:, : self.vocab_size]
+        else:
+            sample_hidden_states = hidden_states[input_batch.logits_indices]
+            logits = self.model.compute_logits(sample_hidden_states)
+
+        # A diffusion prefill has no logit rows even when a bitmask row
+        # arrived for it.
+        if grammar_output is not None and logits.shape[0] > 0:
+            # Apply grammar bitmask to the logits in-place.
+            assert self.structured_outputs_worker is not None
+            self.structured_outputs_worker.apply_grammar_bitmask(
+                logits,
+                input_batch,
+                grammar_output.structured_output_request_ids,
+                grammar_output.grammar_bitmask,
+            )
+
+        sampler_output: SamplerOutput | None
+        if input_batch.num_reqs == 0:
+            # This rank owns no requests this step. It contributes an
+            # all-padding block to the gather below.
+            sampler_output = None
+        elif input_batch.num_draft_tokens == 0 or self.rejection_sampler is None:
+            assert self.sampler is not None
+            sampler_output = self.sampler(logits, input_batch)
+        else:
+            # Rejection sampling for spec decoding.
+            assert self.rejection_sampler is not None
+            assert self.speculator is not None
+            sampler_output = self.rejection_sampler(
+                logits,
+                input_batch,
+                # Draft logits are needed for probabilistic rejection sampling.
+                self.speculator.draft_logits,
+            )
+
+        if shard_metadata is not None:
+            # Gather the sharded sampler outputs from the TP ranks into a single
+            # sampler output.
+            assert self.sampler is not None
+            sampler_output = gather_sampler_output(
+                sampler_output,
+                shard_metadata,
+                device=self.device,
+                global_batch=global_input_batch,
+                local_batch=input_batch,
+                gather_num_nans=self.sampler.compute_nans,
+                logprobs_dims=self.sampler.get_logprobs_dims(
+                    global_input_batch.idx_mapping_np,
+                    # Rejection sampler does not return logprob token ids.
+                    include_token_ids=(
+                        global_input_batch.num_draft_tokens == 0
+                        or self.rejection_sampler is None
+                    ),
+                ),
+            )
+
+        assert sampler_output is not None
+        return sampler_output, sampler_output.num_sampled, sampler_output.num_rejected
+
+    def postprocess_sampled(
+        self,
+        idx_mapping: torch.Tensor,  # May include -1 for masked entries
+        sampled_tokens: torch.Tensor,
+        num_sampled: torch.Tensor,
+        num_rejected: torch.Tensor,
+        query_start_loc: torch.Tensor | None = None,
+    ) -> None:
+        # Update the number of computed tokens.
+        output_bin_counts = None
+        if self.is_last_pp_rank:
+            assert self.sampler is not None
+            output_bin_counts = self.sampler.penalties_state.output_bin_counts
+        post_update(
+            idx_mapping,
+            self.req_states.num_computed_tokens.gpu,
+            self.req_states.last_sampled_tokens,
+            output_bin_counts,
+            sampled_tokens,
+            num_sampled,
+            num_rejected,
+            query_start_loc,
+            self.req_states.all_token_ids.gpu,
+            self.req_states.total_len.gpu,
+        )
+
+        self.model_state.postprocess_state(
+            idx_mapping, num_sampled, self.req_states.num_computed_tokens.gpu
+        )
+
+    def _merge_ec_connector_no_forward(
+        self, scheduler_output: SchedulerOutput, output: ModelRunnerOutput
+    ) -> ModelRunnerOutput:
+        """Let the EC connector send/recv on a step with no work to run."""
+        return ModelRunnerOutput.with_ec_conn_output(
+            output,
+            self.ec_connector.no_forward(scheduler_output).ec_connector_output,
+        )
+
+    @torch.inference_mode()
+    def execute_model(
+        self,
+        scheduler_output: SchedulerOutput,
+        intermediate_tensors: IntermediateTensors | None = None,
+        dummy_run: bool = False,
+        skip_attn_for_dummy_run: bool = False,
+        is_profile: bool = False,
+        context_len: int = 0,
+        valid_dummy_state_slots: bool = False,
+    ) -> ModelRunnerOutput | IntermediateTensors | None:
+        if not dummy_run:
+            # Update the request states.
+            self.update_pp_decode_requests()
+            self.finish_requests(scheduler_output)
+            self.free_states(scheduler_output)
+            self.add_requests(scheduler_output)
+            self.update_requests(scheduler_output)
+            self.block_tables.apply_staged_writes()
+            if self.aux_output_connector is not None:
+                # Register this step before the GPU forward.
+                self.aux_output_connector.begin_step(
+                    scheduler_output.aux_output_connector_metadata
+                )
+            if scheduler_output.total_num_scheduled_tokens == 0:
+                # No need to run the model.
+                empty_output = self.kv_connector.no_forward(scheduler_output)
+                return self._merge_ec_connector_no_forward(
+                    scheduler_output, empty_output
+                )
+
+        # Get batch descriptor and sync across DP ranks.
+        num_reqs = len(scheduler_output.num_scheduled_tokens)
+        num_toks = scheduler_output.total_num_scheduled_tokens
+        max_query_len = max(scheduler_output.num_scheduled_tokens.values())
+        batch_req_state, uniform_tok_count = self.gather_batch_req_state(
+            scheduler_output, dummy_run
+        )
+        if batch_req_state is not None:
+            num_toks = batch_req_state.num_tokens
+            if self.pcp_manager is not None:
+                num_toks = self.pcp_manager.get_num_tokens_for_dispatch(
+                    batch_req_state.num_scheduled_tokens,
+                    batch_req_state.is_prefilling_np,
+                )
+
+        num_active_loras = 0
+        if self.lora_config:
+            req_ids = list(scheduler_output.num_scheduled_tokens.keys())
+            num_active_loras = get_num_active_loras_for_dispatch(
+                self.lora_config, self.lora_state, req_ids, dummy_run
+            )
+
+        skip_compiled = False
+        if self.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs:
+            # Encoder-decoder models such as Whisper should run eager/non-compiled
+            # when encoder inputs are scheduled, because this step updates
+            # cross-attention cache with dynamic encoder outputs.
+            skip_compiled = True
+
+        batch_desc, dp_sync = dispatch_cg_and_sync_dp(
+            self.cudagraph_manager,
+            num_reqs,
+            num_toks,
+            uniform_tok_count,
+            self.dp_size,
+            self.dp_rank,
+            max_query_len=max_query_len,
+            need_eager=is_profile or skip_compiled,
+            num_active_loras=num_active_loras,
+            parallel_config=self.parallel_config,
+            allow_ubatching=(
+                self.ubatch_runner is not None and not skip_attn_for_dummy_run
+            ),
+            uniform_decode=uniform_tok_count == self.decode_query_len,
+        )
+
+        if batch_desc.num_tokens == 0:
+            # All DP ranks have zero tokens to run.
+            empty_output = self.kv_connector.no_forward(scheduler_output)
+            return self._merge_ec_connector_no_forward(scheduler_output, empty_output)
+
+        cudagraph_stats = None
+        if not dummy_run:
+            # Common case.
+            # Prepare all the inputs and copy to the input buffers.
+            if self.observability_config.cudagraph_metrics:
+                cudagraph_stats = make_cudagraph_stats(batch_desc, num_toks)
+            assert batch_req_state is not None
+            input_batch = self.prepare_inputs(
+                scheduler_output, batch_req_state, batch_desc, num_active_loras
+            )
+            block_tables, slot_mappings = self.prepare_attn(input_batch)
+            # Mamba "align" pre-copy: migrate recurrent state across block
+            # boundaries before the forward. Runs only on real batches, and
+            # before model_state.prepare_attn gathers num_accepted_tokens so the
+            # boundary reset is visible to the attention metadata.
+            self.model_state.preprocess_state(
+                input_batch,
+                block_tables,
+                self.kv_cache_config,
+                self.req_states.num_computed_tokens.gpu,
+            )
+
+            if self.lora_config:
+                # Activate LoRA adapters.
+                lora_inputs = self.lora_state.make_lora_inputs(
+                    input_batch.req_ids,
+                    input_batch.idx_mapping_np,
+                    input_batch.num_scheduled_tokens,
+                )
+                self._set_active_loras(*lora_inputs)
+        else:
+            # No actual tokens to run. A dummy run for DP or memory profiling.
+            dummy_num_reqs = batch_desc.num_reqs or num_reqs
+            input_batch = InputBatch.make_dummy(
+                dummy_num_reqs,
+                batch_desc.num_tokens,
+                self.input_buffers,
+                max_query_len=batch_desc.max_query_len,
+                # Profiling and warmup must route the dummy tokens to experts
+                # so MoE memory is measured and MoE kernels are exercised.
+                is_padding=not is_profile,
+            )
+            if self.pcp_manager is not None:
+                input_batch = self.pcp_manager.prepare_inputs_to_capture(input_batch)
+            if skip_attn_for_dummy_run:
+                assert batch_desc.cg_mode != CUDAGraphMode.FULL, (
+                    "Attention metadata must be prepared for dummy runs when using "
+                    "FULL cudagraph mode."
+                )
+                block_tables, slot_mappings = None, None
+            else:
+                block_tables, slot_mappings = self.prepare_dummy_attn(
+                    input_batch, valid_dummy_state_slots
+                )
+                if context_len:
+                    set_dummy_context(
+                        input_batch,
+                        self.block_tables,
+                        context_len,
+                        self.kv_cache_config.num_blocks,
+                        self.max_model_len,
+                        input_block_tables=block_tables,
+                    )
+
+        if self.dcp_size > 1 and not (dummy_run and skip_attn_for_dummy_run):
+            input_buffers = self.input_buffers
+            if self.pcp_manager is not None:
+                input_buffers = self.pcp_manager.input_buffers
+            input_batch.dcp_local_seq_lens = prepare_dcp_local_seq_lens(
+                input_buffers.dcp_local_seq_lens,
+                input_batch.seq_lens,
+                input_batch.num_reqs,
+                self.dcp_size,
+                self.dcp_rank,
+                self.cp_interleave,
+                num_reqs_padded=input_batch.num_reqs_after_padding,
+            )
+
+        attn_metadata = None
+        slot_mappings_by_layer = None
+        ubatch_state: UBatchState | None = None
+        if batch_desc.num_ubatches > 1:
+            assert self.ubatch_runner is not None
+            assert block_tables is not None and slot_mappings is not None
+            ubatch_state = self.ubatch_runner.prepare(
+                input_batch,
+                block_tables,
+                slot_mappings,
+                cg_mode=batch_desc.cg_mode,
+                for_capture=dummy_run and batch_desc.cg_mode == CUDAGraphMode.FULL,
+            )
+        elif not (dummy_run and skip_attn_for_dummy_run):
+            assert slot_mappings is not None
+            slot_mappings_by_layer = build_slot_mappings_by_layer(
+                slot_mappings, self.kv_cache_config
+            )
+            assert block_tables is not None
+            attn_groups = self.attn_groups
+            if dummy_run and is_profile:
+                # Mamba layers take a cheap warmup path with no metadata;
+                # attention metadata is still built so those kernels tune.
+                attn_groups = [
+                    [g for g in groups if not isinstance(g.kv_cache_spec, MambaSpec)]
+                    for groups in attn_groups
+                ]
+            attn_metadata = self.model_state.prepare_attn(
+                input_batch,
+                batch_desc.cg_mode,
+                block_tables,
+                slot_mappings,
+                attn_groups,
+                self.kv_cache_config,
+                # FULL replay reads capture-time metadata buffers. Re-stage them
+                # from the zeroed dummy block tables instead of retaining state
+                # indices from the previous real batch.
+                for_capture=dummy_run and batch_desc.cg_mode == CUDAGraphMode.FULL,
+            )
+
+        input_ids = input_batch.input_ids
+        inputs_embeds = None
+        ec_connector_output = None
+        if self.uses_inputs_embeds and self.is_first_pp_rank:
+            # Prepare inputs_embeds (MM encoder outputs and/or prompt_embeds
+            # overlay). Only first PP rank prepares them.
+            if dummy_run:
+                # Obtain embeddings of correct shape for compiled model.
+                inputs_embeds = self.model_state.dummy_inputs_embeds(
+                    input_batch.num_tokens_after_padding
+                )
+            else:
+                scheduled_encoder_inputs = scheduler_output.scheduled_encoder_inputs
+                if self.supports_mm_inputs and self.lora_config is not None:
+                    set_active_mm_loras(
+                        model=self.model,
+                        lora_manager=self.lora_manager,
+                        encoder_cache=self.encoder_cache,
+                        req_id_to_index=self.req_states.req_id_to_index,
+                        lora_state=self.lora_state,
+                        scheduled_encoder_inputs=scheduled_encoder_inputs,
+                    )
+                with self.ec_connector.maybe_get_output(
+                    scheduler_output
+                ) as ec_connector_output:
+                    inputs_embeds = self.model_state.prepare_inputs_embeds(
+                        scheduled_encoder_inputs, input_batch, self.req_states
+                    )
+            if inputs_embeds is not None and not requires_raw_input_tokens(self.model):
+                input_ids = None
+
+        model_inputs = {
+            "input_ids": input_ids,
+            "positions": input_batch.positions,
+            "inputs_embeds": inputs_embeds,
+            "intermediate_tensors": None,
+            # NOTE: Values returned by `prepare_inputs` will override the default
+            # values above.
+            **self.model_state.prepare_inputs(input_batch, self.req_states),
+        }
+        if not self.is_first_pp_rank:
+            # Update for non-first PP ranks.
+            model_inputs["input_ids"] = None
+            model_inputs["inputs_embeds"] = None
+
+            # Prepare the intermediate tensors.
+            assert intermediate_tensors is not None
+            assert self.intermediate_tensors is not None
+            n = input_batch.num_tokens_after_padding
+            new_tensors = {
+                k: v[:n]
+                if dummy_run
+                else v[:n].copy_(intermediate_tensors.tensors[k][:n])
+                for k, v in self.intermediate_tensors.tensors.items()
+            }
+            model_inputs["intermediate_tensors"] = IntermediateTensors(new_tensors)
+            del intermediate_tensors
+
+        # Update the EPLB meta.
+        ubatch_slices = ubatch_state.slices if ubatch_state is not None else None
+        self.eplb.prepare_forward(
+            self.model_config, input_batch.num_tokens, ubatch_slices
+        )
+
+        self.step_timing.record_batch(
+            input_batch, batch_desc.cg_mode == CUDAGraphMode.FULL
+        )
+        self.step_timing.forward_start()
+
+        connector_kwargs = dict(
+            scheduler_output=scheduler_output,
+            request_state_indices=input_batch.idx_mapping,
+            request_ids=input_batch.req_ids,
+            num_tokens=input_batch.num_tokens,
+        )
+
+        # Run model.
+        if batch_desc.cg_mode == CUDAGraphMode.FULL:
+            # Use explicit cudagraph replay for FULL mode.
+            # NOTE(woosuk): Here, we don't need to pass the input tensors,
+            # because they are already copied to the CUDA graph input buffers.
+            assert self.cudagraph_manager is not None
+            self.kv_connector.pre_forward(
+                **connector_kwargs, attn_metadata=attn_metadata
+            )
+            model_output = self.cudagraph_manager.run_fullgraph(batch_desc)
+        else:
+            # For piecewise and eager mode, just call model().
+            batch_descriptor = BatchDescriptor(
+                num_tokens=input_batch.num_tokens_after_padding,
+                has_lora=self.lora_config is not None,
+                num_active_loras=batch_desc.num_active_loras,
+            )
+
+            with set_forward_context(
+                attn_metadata,
+                self.vllm_config,
+                num_tokens=input_batch.num_tokens_after_padding,
+                cudagraph_runtime_mode=batch_desc.cg_mode,
+                num_tokens_across_dp=(
+                    dp_sync.num_tokens_across_dp if dp_sync is not None else None
+                ),
+                batch_descriptor=batch_descriptor,
+                ubatch_slices=ubatch_slices,
+                slot_mapping=slot_mappings_by_layer,
+                skip_compiled=skip_compiled,
+                is_padding=input_batch.is_padding,
+            ):
+                self.kv_connector.pre_forward(**connector_kwargs)
+                if ubatch_state is not None:
+                    assert self.ubatch_runner is not None
+                    model_output = self.ubatch_runner.run(
+                        self.model, model_inputs, ubatch_state
+                    )
+                elif batch_desc.cg_mode == CUDAGraphMode.PIECEWISE:
+                    # Run the PIECEWISE graph (compiled PW cudagraph or breakable
+                    # cudagraph, chosen inside run_pw_graph). cg_mode is only
+                    # PIECEWISE after the cudagraph manager exists.
+                    assert self.cudagraph_manager is not None
+                    model_output = self.cudagraph_manager.run_pw_graph(
+                        self.model, model_inputs
+                    )
+                else:
+                    # Eager (NONE): call the raw model directly.
+                    model_output = self.model(**model_inputs)
+
+        self.kv_connector.finish_forward()
+
+        if self.is_last_pp_rank:
+            if self.use_aux_hidden_state_outputs:
+                assert isinstance(model_output, tuple)
+                hidden_states, aux_hidden_states = model_output
+            else:
+                assert isinstance(model_output, torch.Tensor)
+                hidden_states = model_output
+                aux_hidden_states = None
+            output_intermediate_tensors = None
+        else:
+            assert isinstance(model_output, IntermediateTensors)
+            hidden_states = None
+            aux_hidden_states = None
+            output_intermediate_tensors = model_output
+
+        finished_req_ids = scheduler_output.finished_req_ids
+        self.execute_model_state = ExecuteModelState(
+            input_batch=input_batch,
+            attn_metadata=attn_metadata,
+            slot_mappings_by_layer=slot_mappings_by_layer,
+            hidden_states=hidden_states,
+            aux_hidden_states=aux_hidden_states,
+            dp_sync=dp_sync,
+            finished_req_ids=finished_req_ids,
+            ec_connector_output=ec_connector_output,
+            cudagraph_stats=cudagraph_stats,
+        )
+
+        if not self.is_last_pp_rank:
+            # Non-last PP rank: return IntermediateTensors for sending.
+            assert output_intermediate_tensors is not None
+            assert self.pp_handler is not None
+            return self.pp_handler.relay_aux_hidden_states(
+                model_inputs["intermediate_tensors"], output_intermediate_tensors
+            )
+        return None
+
+    @torch.inference_mode()
+    @step_eplb_after()
+    def sample_tokens(
+        self, grammar_output: GrammarOutput | None
+    ) -> AsyncOutput | ModelRunnerOutput | None:
+        if self.execute_model_state is None:
+            # The prior execute_model call must have failed.
+            return None
+
+        input_batch = self.execute_model_state.input_batch
+        attn_metadata = self.execute_model_state.attn_metadata
+        slot_mappings_by_layer = self.execute_model_state.slot_mappings_by_layer
+        hidden_states = self.execute_model_state.hidden_states
+        aux_hidden_states = self.execute_model_state.aux_hidden_states
+        dp_sync = self.execute_model_state.dp_sync
+        finished_req_ids = self.execute_model_state.finished_req_ids
+        ec_connector_output = self.execute_model_state.ec_connector_output
+        cudagraph_stats = self.execute_model_state.cudagraph_stats
+        self.execute_model_state = None
+
+        if not self.is_last_pp_rank:
+            # Non-last PP rank: hidden_states is None because this rank produced
+            # IntermediateTensors instead of final hidden states. Receive the
+            # sampled tokens broadcast from the last rank and update local state.
+            assert self.pp_handler is not None
+            all_decode_next = self.pp_handler.receive(input_batch)
+            # Optimistically update num_computed_tokens for entire batch here.
+            # Will be adjusted for rejections if necessary in update_requests.
+            self.postprocess_num_computed_tokens(input_batch)
+            if not all_decode_next:
+                # Might contain non-final prefill chunks, which will be scheduled
+                # in the immediate next step (rather than in pp_size steps).
+                self.model_state.postprocess_state(input_batch.idx_mapping, 0)
+
+            # Post-step KV connector related operations.
+            kv_connector_output = self.kv_connector.post_forward(finished_req_ids)
+            # The first PP rank holds the encoder cache, so pass its EC output on.
+            output = ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)
+            return ModelRunnerOutput.with_ec_conn_output(output, ec_connector_output)
+
+        # Last rank: sample tokens
+        assert hidden_states is not None
+        draft_hidden_states = hidden_states
+        if self.pcp_manager is not None:
+            hidden_states, aux_hidden_states, input_batch = (
+                self.pcp_manager.restore_for_sampling(hidden_states, aux_hidden_states)
+            )
+
+        sampler_output, num_sampled, num_rejected = self.sample(
+            hidden_states, input_batch, grammar_output
+        )
+
+        if self.pp_handler is not None:
+            # Broadcast to non-last PP ranks (handles spec decode multi-token).
+            self.pp_handler.broadcast(
+                sampler_output.sampled_token_ids,
+                num_sampled,
+                num_rejected,
+                input_batch,
+            )
+
+        assert self.prompt_logprobs_worker is not None
+        prompt_logprobs_dict = self.prompt_logprobs_worker.compute_prompt_logprobs(
+            self.model.compute_logits,
+            hidden_states,
+            input_batch,
+            self.req_states.all_token_ids.gpu,
+            self.req_states.num_computed_tokens.gpu,
+            self.req_states.prompt_len.np,
+        )
+
+        # Prepare the model runner output.
+        model_runner_output = ModelRunnerOutput(
+            req_ids=input_batch.req_ids,
+            # NOTE(woosuk): req_id_to_index is unused in this model runner.
+            # Only for compatibility with the existing model runner and scheduler.
+            req_id_to_index={req_id: i for i, req_id in enumerate(input_batch.req_ids)},
+            sampled_token_ids=None,  # type: ignore
+            prompt_logprobs_dict=prompt_logprobs_dict,  # type: ignore[arg-type]
+            cudagraph_stats=cudagraph_stats,
+        )
+        pending_aux_output = None
+        if self.aux_output_connector is not None:
+            pending_aux_output = self.aux_output_connector.prepare_output(input_batch)
+
+        # Start async output copy here so that it can overlap with speculator proposal.
+        async_output = AsyncOutput(
+            model_runner_output=model_runner_output,
+            sampler_output=sampler_output,
+            num_sampled_tokens=num_sampled,
+            main_stream=self.main_stream,
+            copy_stream=self.output_copy_stream,
+            check_ep_fault=self.check_ep_fault,
+            pending_aux_output=pending_aux_output,
+        )
+
+        mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None
+        # The encoder runner exists only on the first PP rank, so later ranks
+        # have no cached embeddings to gather.
+        if (
+            self.speculator is not None
+            and self.speculator.supports_mm_inputs
+            and self.model_state.supports_mm_inputs
+        ):
+            # Get cached multimodal embeddings for draft forward.
+            # NOTE: This is done here because postprocess updates
+            # num_computed_prefill_tokens.
+            # The EAGLE/MTP drafter reads one position ahead of the target.
+            # TODO(TheEpicDolphin): Gather MM embeddings for all speculative
+            # steps during multi-module MTP.
+            mm_inputs = self.model_state.gather_mm_embeddings(
+                input_batch, draft_lookahead=1
+            )
+
+        # Postprocess results and update request states.
+        # NOTE: This is intentionally done after creating the AsyncOutput,
+        # ensuring that `copy_event` is recorded before calling postprocess.
+        # This sequencing may slightly reduce latency as async D2H copy does not
+        # need to wait for the postprocess to finish.
+        self.postprocess_sampled(
+            input_batch.idx_mapping,
+            sampler_output.sampled_token_ids,
+            num_sampled,
+            num_rejected,
+            input_batch.query_start_loc,
+        )
+
+        if self.speculator is not None:
+            assert self.sampler is not None
+            if isinstance(self.speculator, DraftModelSpeculator) and draft_trunc.graded():
+                self.speculator.observe_verification(
+                    input_batch.idx_mapping, num_sampled, num_rejected
+                )
+            if draft_trunc.LOG:
+                draft_trunc.log_step(
+                    self.speculator.draft_token_confidence_probs[input_batch.idx_mapping],
+                    input_batch.cu_num_logits, num_sampled,
+                    getattr(self.speculator, "acceptance_estimator", None), input_batch.idx_mapping,
+                )
+            # Let the target override the hidden state fed to the drafter
+            # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The
+            # target returns a persistent buffer sized at max_num_batched_tokens;
+            # slice to the active token count that propose() expects.
+            spec_hidden_states = draft_hidden_states
+            if hasattr(self.model, "get_mtp_target_hidden_states"):
+                pre_hc_hidden_states = self.model.get_mtp_target_hidden_states()
+                spec_hidden_states = pre_hc_hidden_states[: draft_hidden_states.size(0)]
+            if isinstance(self.sampler, GPUWatermarkSampler):
+                self.speculator.prepare_watermarking(
+                    self.sampler._get_contexts(input_batch.idx_mapping),
+                    self.sampler.watermarking.gpu[input_batch.idx_mapping],
+                )
+            with use_workspace_lane(self._draft_workspace_lane):
+                draft_tokens = self.speculator.propose(
+                    input_batch,
+                    attn_metadata,
+                    slot_mappings_by_layer,
+                    spec_hidden_states,
+                    aux_hidden_states,
+                    num_sampled,
+                    num_rejected,
+                    self.req_states.last_sampled_tokens,
+                    self.req_states.next_prefill_tokens,
+                    self.sampler.sampling_states.temperature.gpu,
+                    self.sampler.sampling_states.seeds.gpu,
+                    dp_sync=dp_sync,
+                    mm_inputs=mm_inputs,
+                )
+            self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens
+            if self.adaptive_verification is not None:
+                self.adaptive_verification.record_confidences(
+                    self.speculator.draft_token_confidence_probs, input_batch
+                )
+
+        if self.num_speculative_steps > 0:
+            # Spec-decode and diffusion LLMs both use draft tokens but the latter does
+            # not have a speculator (i.e. self.speculator is None)
+            self.draft_tokens_handler.set_draft_tokens(
+                input_batch,
+                self.req_states.draft_tokens[input_batch.idx_mapping],
+            )
+            if self.pp_handler is not None:
+                self.pp_handler.broadcast_drafts(
+                    self.req_states.draft_tokens, input_batch
+                )
+
+        # Post-step KV connector related operations.
+        kv_connector_output = self.kv_connector.post_forward(finished_req_ids)
+        model_runner_output.kv_connector_output = kv_connector_output
+        model_runner_output.ec_connector_output = ec_connector_output
+
+        return async_output
+
+    def take_draft_token_ids(self) -> DraftTokenIds | None:
+        return self.draft_tokens_handler.get_draft_tokens()
+
+    @torch.inference_mode()
+    @step_eplb_after()
+    def pool(self) -> AsyncPoolingOutput | ModelRunnerOutput | None:
+        if self.execute_model_state is None:
+            # The prior execute_model call must have failed.
+            return None
+
+        input_batch = self.execute_model_state.input_batch
+        hidden_states = self.execute_model_state.hidden_states
+        finished_req_ids = self.execute_model_state.finished_req_ids
+        ec_connector_output = self.execute_model_state.ec_connector_output
+        self.execute_model_state = None
+
+        # Post-step KV connector related operations.
+        kv_connector_output = self.kv_connector.post_forward(finished_req_ids)
+
+        if not self.is_last_pp_rank:
+            self.postprocess_num_computed_tokens(input_batch)
+            output = ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)
+            return ModelRunnerOutput.with_ec_conn_output(output, ec_connector_output)
+
+        assert self.pooling_runner is not None
+        pooler_output, finished_mask = self.pooling_runner.pool(
+            hidden_states, input_batch, self.req_states
+        )
+
+        # Build the model runner output.
+        model_runner_output = ModelRunnerOutput(
+            req_ids=input_batch.req_ids,
+            req_id_to_index={req_id: i for i, req_id in enumerate(input_batch.req_ids)},
+            kv_connector_output=kv_connector_output,
+            ec_connector_output=ec_connector_output,
+        )
+        async_output = AsyncPoolingOutput(
+            model_runner_output=model_runner_output,
+            pooler_output=pooler_output,
+            finished_mask=finished_mask,
+            main_stream=self.main_stream,
+            copy_stream=self.output_copy_stream,
+        )
+
+        self.postprocess_num_computed_tokens(input_batch)
+        return async_output
+
+    def postprocess_num_computed_tokens(self, input_batch: InputBatch) -> None:
+        # Update the number of computed tokens.
+        post_update_num_computed_tokens(
+            input_batch.idx_mapping,
+            self.req_states.num_computed_tokens.gpu,
+            input_batch.query_start_loc,
+        )
+
+    def shutdown(self) -> None:
+        """Release GPU tensors (model weights, KV caches, workspace) so that
+        memory is reclaimable when running in the same process."""
+        torch.accelerator.synchronize()
+        if self.aux_output_connector is not None:
+            self.aux_output_connector.close()
+        set_offloader(None)
+        self.cudagraph_manager = None
+        self.fast_prefill = None
+        self.pooling_runner = None
+        if hasattr(self, "kv_caches"):
+            self.kv_caches.clear()
+        if hasattr(self, "attn_groups"):
+            self.attn_groups.clear()
+        if hasattr(self, "kv_cache_config"):
+            del self.kv_cache_config
+        if hasattr(self, "model_state") and self.model_state.supports_mm_inputs:
+            self.model_state.encoder_runner.clear()
+        free_before_shutdown(self.vllm_config)
+        if hasattr(self, "model_state"):
+            del self.model_state
+        # Detach the layer-level KV/state cache tensors before dropping the
+        # models; the model objects can outlive this runner.
+        speculator = getattr(self, "speculator", None)
+        if speculator is not None:
+            if draft_model := getattr(speculator, "model", None):
+                clear_layer_kv_caches(draft_model.modules())
+            self.speculator = None
+        if hasattr(self, "model"):
+            clear_layer_kv_caches(self.model.modules())
+            del self.model
+
+        gc.collect()
+        torch.accelerator.empty_cache()
+        logger.debug("Cleaned up model weights, KV caches, and workspace")
+
+    ########### EPLB methods start ###########
+    @property
+    def eplb_state(self):
+        return self.eplb.state
+
+    @eplb_state.setter
+    def eplb_state(self, state) -> None:
+        self.eplb.state = state
+
+    @property
+    def eep_eplb_suppressed(self) -> bool:
+        return self.eplb.suppressed
+
+    @eep_eplb_suppressed.setter
+    def eep_eplb_suppressed(self, suppressed: bool) -> None:
+        self.eplb.suppressed = suppressed
+
+    def setup_eplb_from_mapping(
+        self, expanded_physical_to_logical: torch.Tensor
+    ) -> None:
+        self.eplb.setup_from_mapping(
+            self.model, self.model_config, expanded_physical_to_logical
+        )
+
+    ########### EPLB methods end ###########
+
+    # Out-of-tree hardware runners can select a PCP manager class.
+    @property
+    def pcp_manager_cls(self) -> type[pcp.PCPManager]:
+        return pcp.PCPManager
+
+
+class ExecuteModelState(NamedTuple):
+    input_batch: InputBatch
+    attn_metadata: dict[str, Any] | None
+    slot_mappings_by_layer: dict[str, torch.Tensor] | None
+    hidden_states: torch.Tensor | None
+    aux_hidden_states: list[torch.Tensor] | None
+    dp_sync: DPSyncState | None
+    finished_req_ids: set[str]
+    ec_connector_output: ECConnectorOutput | None
+    cudagraph_stats: CUDAGraphStat | None
+
+
+class BatchReqState(NamedTuple):
+    """CPU request state for a scheduled batch, in batch (sorted) order."""
+
+    req_ids: list[str]
+    num_scheduled_tokens: np.ndarray  # [num_reqs]
+    # May be less than scheduler_output.total_num_scheduled_tokens:
+    # adaptive verification trims the draft budget before running.
+    num_tokens: int
+    idx_mapping_np: np.ndarray  # [num_reqs]
+    prefill_len_np: np.ndarray  # [num_reqs]
+    num_computed_prefill_tokens_np: np.ndarray  # [num_reqs]
+    is_prefilling_np: np.ndarray  # [num_reqs]
+    has_prefill: bool
+
+
+def sort_batch_req_ids(
+    num_tokens_per_req: dict[str, int],
+    draft_tokens: dict[str, list[int]],
+    decode_query_len: int,
+) -> list[str]:
+    # Order verification/decode -> short_extend -> prefill;
+    # split_decodes_and_prefills relies on decode-like requests leading.
+    key = lambda r: (
+        not draft_tokens.get(r),
+        (num := num_tokens_per_req[r]) != decode_query_len,
+        num,
+    )
+    return sorted(num_tokens_per_req, key=key)

docker-images/glm53-kindling/experimental/draft-trunc/rejection_sampler.py

@@ -0,0 +1,340 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections.abc import Iterable, Iterator
+from typing import Any
+
+import numpy as np
+import torch
+
+from vllm.v1.worker.gpu.spec_decode import draft_trunc
+
+from vllm.config import SpeculativeConfig
+from vllm.config.model import PROCESSED_LOGPROBS_MODES
+from vllm.triton_utils import tl, triton
+from vllm.v1.outputs import LogprobsTensors
+from vllm.v1.spec_decode.utils import unconditional_to_conditional_rates
+from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler
+from vllm.v1.worker.gpu.input_batch import (
+    InputBatch,
+    get_num_sampled_and_rejected,
+)
+from vllm.v1.worker.gpu.metrics.logits import get_num_nans
+from vllm.v1.worker.gpu.sample.logprob import compute_topk_scores
+from vllm.v1.worker.gpu.sample.output import SamplerOutput
+from vllm.v1.worker.gpu.sample.sampler import Sampler
+from vllm.v1.worker.gpu.sample.states import NO_LOGPROBS
+from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import (
+    rejection_sample,
+)
+
+# Cap on the FP32 target-logits buffer materialized by apply_sampling_params.
+# TODO(mgoin): Chunking is a workaround. The rejection kernels already upcast
+# per vocab block on load and apply ops like temperature and gumbel, so folding
+# sampling-param application into those kernels would remove this buffer and
+# its traffic entirely.
+MAX_CHUNK_BYTES = 2**30  # 1GB
+_FP32_BYTES = 4
+
+
+def get_max_chunk_logits(vocab_size: int) -> int:
+    """Largest number of logits rows one verification chunk may hold."""
+    return max(1, MAX_CHUNK_BYTES // (vocab_size * _FP32_BYTES))
+
+
+def _iter_request_chunks(
+    cu_num_logits: np.ndarray, max_chunk_logits: int
+) -> Iterator[tuple[int, int]]:
+    """Yield maximally packed request ranges without splitting requests."""
+    assert max_chunk_logits > 0
+    num_reqs = cu_num_logits.size - 1
+    start = 0
+    while start < num_reqs:
+        max_logit = int(cu_num_logits[start]) + max_chunk_logits
+        end = int(np.searchsorted(cu_num_logits, max_logit, side="right") - 1)
+        end = min(num_reqs, max(start + 1, end))
+        yield start, end
+        start = end
+
+
+@triton.jit
+def _flatten_sampled_kernel(
+    # [num_logits]
+    flat_sampled_ptr,
+    # [num_reqs, num_speculative_steps + 1]
+    sampled_ptr,
+    sampled_stride,
+    # [num_reqs]
+    num_sampled_ptr,
+    # [num_reqs + 1]
+    cu_num_logits_ptr,
+):
+    req_idx = tl.program_id(0)
+    start_idx = tl.load(cu_num_logits_ptr + req_idx)
+    num_sampled = tl.load(num_sampled_ptr + req_idx)
+    for i in range(num_sampled):
+        token_id = tl.load(sampled_ptr + req_idx * sampled_stride + i)
+        tl.store(flat_sampled_ptr + start_idx + i, token_id)
+
+
+class RejectionSampler:
+    def __init__(
+        self,
+        sampler: Sampler,
+        spec_config: SpeculativeConfig,
+        device: torch.device,
+        *,
+        watermark_key: int | None = None,
+    ):
+        self.sampler = sampler
+        self.watermark_key = watermark_key
+        if watermark_key is not None:
+            assert isinstance(sampler, GPUWatermarkSampler)
+        self.num_speculative_steps = spec_config.num_speculative_tokens
+        self.enable_adaptive_verification = spec_config.enable_adaptive_verification
+        rejection_sample_method = spec_config.rejection_sample_method
+        self.use_block_verification: bool = False
+        self.synthetic_conditional_rates: torch.Tensor | None = None
+        if rejection_sample_method == "synthetic":
+            assert spec_config.synthetic_acceptance_rates is not None
+            self.synthetic_conditional_rates = torch.tensor(
+                unconditional_to_conditional_rates(
+                    spec_config.synthetic_acceptance_rates
+                ),
+                dtype=torch.float32,
+                device=device,
+            )
+        elif rejection_sample_method == "block":
+            self.use_block_verification = True
+
+    def _get_logprobs_tensors(
+        self,
+        sampled: torch.Tensor,
+        num_sampled: torch.Tensor,
+        logits: torch.Tensor,
+        cu_num_logits: torch.Tensor,
+        cu_num_logits_np: np.ndarray,
+        max_num_logprobs: int,
+    ) -> LogprobsTensors | None:
+        if max_num_logprobs == NO_LOGPROBS:
+            return None
+
+        num_reqs = cu_num_logits.shape[0] - 1
+        num_logits = logits.shape[0]
+        flat_sampled = torch.zeros(
+            num_logits, dtype=sampled.dtype, device=sampled.device
+        )
+        _flatten_sampled_kernel[(num_reqs,)](
+            flat_sampled,
+            sampled,
+            sampled.stride(0),
+            num_sampled,
+            cu_num_logits,
+            num_warps=1,
+        )
+        expanded_logits = num_logits != num_reqs
+        cu_num_generated_tokens: list[int] | torch.Tensor | None = None
+        if expanded_logits:
+            if self.enable_adaptive_verification:
+                # Adaptive verification keeps the true per-request boundaries
+                # on device only; cu_num_logits_np holds the pre-compacted layout.
+                cu_num_generated_tokens = cu_num_logits.clone()
+            else:
+                cu_num_generated_tokens = cu_num_logits_np.tolist()
+        return compute_topk_scores(
+            logits,
+            max_num_logprobs,
+            flat_sampled,
+            cu_num_generated_tokens,
+            logits_mode=self.sampler.logprobs_mode
+            in ("raw_logits", "processed_logits"),
+        )
+
+    def _watermarking_kwargs(
+        self,
+        draft_sampled: torch.Tensor,
+        expanded_idx_mapping: torch.Tensor,
+        expanded_local_pos: torch.Tensor,
+    ) -> dict[str, Any]:
+        if self.watermark_key is None:
+            return {}
+        assert isinstance(self.sampler, GPUWatermarkSampler)
+        return {
+            "contexts": self.sampler._get_contexts(
+                expanded_idx_mapping, expanded_local_pos, draft_sampled
+            ),
+            "watermarking": self.sampler.watermarking.gpu,
+            "watermark_key": self.watermark_key,
+        }
+
+    def _verify(
+        self,
+        logits: torch.Tensor,
+        draft_logits: torch.Tensor | None,
+        draft_sampled: torch.Tensor,
+        accept_drafts: torch.Tensor,
+        pos: torch.Tensor,
+        cu_num_logits: torch.Tensor,
+        idx_mapping: torch.Tensor,
+        idx_mapping_np: np.ndarray,
+        expanded_idx_mapping: torch.Tensor,
+        expanded_local_pos: torch.Tensor,
+        seq_lens_upper_bound_np: np.ndarray,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        processed_logits = self.sampler.apply_sampling_params(
+            logits,
+            expanded_idx_mapping,
+            idx_mapping,
+            idx_mapping_np,
+            pos,
+            draft_sampled,
+            expanded_local_pos,
+            seq_lens_upper_bound_np,
+        )
+        sampled, num_sampled = rejection_sample(
+            processed_logits,
+            draft_logits,
+            accept_drafts,
+            cu_num_logits,
+            pos,
+            idx_mapping,
+            expanded_idx_mapping,
+            expanded_local_pos,
+            self.sampler.sampling_states.temperature.gpu,
+            self.sampler.sampling_states.seeds.gpu,
+            self.num_speculative_steps,
+            self.synthetic_conditional_rates,
+            use_fp64=self.sampler.use_fp64_gumbel,
+            use_block_verification=self.use_block_verification,
+            **self._watermarking_kwargs(
+                draft_sampled, expanded_idx_mapping, expanded_local_pos
+            ),
+        )
+        return processed_logits, sampled, num_sampled
+
+    def _verify_in_chunks(
+        self,
+        logits: torch.Tensor,
+        input_batch: InputBatch,
+        draft_logits: torch.Tensor | None,
+        draft_sampled: torch.Tensor,
+        pos: torch.Tensor,
+        max_chunk_logits: int,
+        max_num_logprobs: int,
+        accept_drafts: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, LogprobsTensors | None]:
+        cu_num_logits_np = input_batch.cu_num_logits_np
+        use_processed_logits = self.sampler.logprobs_mode in PROCESSED_LOGPROBS_MODES
+        num_reqs = input_batch.num_reqs
+
+        if logits.shape[0] <= max_chunk_logits:
+            # One chunk covers the batch. Adaptive verification compacts the logits
+            # without updating cu_num_logits_np (it keeps the pre-compacted layout),
+            # so the stale sums must not pick chunk boundaries; its budget cap
+            # guarantees the compacted batch always lands here.
+            request_chunks: Iterable[tuple[int, int]] = ((0, num_reqs),)
+        else:
+            assert not self.enable_adaptive_verification
+            request_chunks = _iter_request_chunks(cu_num_logits_np, max_chunk_logits)
+
+        sampled_chunks: list[torch.Tensor] = []
+        num_sampled_chunks: list[torch.Tensor] = []
+        logprobs_chunks: list[LogprobsTensors] = []
+
+        for start, end in request_chunks:
+            lo = int(cu_num_logits_np[start])
+            hi = int(cu_num_logits_np[end])
+            chunk_cu_num_logits_np = cu_num_logits_np[start : end + 1] - lo
+            chunk_cu_num_logits = input_batch.cu_num_logits[start : end + 1] - lo
+            # draft_logits uses persistent request-state indices and stays global.
+            processed_logits, sampled, num_sampled = self._verify(
+                logits[lo:hi],
+                draft_logits,
+                draft_sampled[lo:hi],
+                (draft_sampled if accept_drafts is None else accept_drafts)[lo:hi],
+                pos[lo:hi],
+                chunk_cu_num_logits,
+                input_batch.idx_mapping[start:end],
+                input_batch.idx_mapping_np[start:end],
+                input_batch.expanded_idx_mapping[lo:hi],
+                input_batch.expanded_local_pos[lo:hi],
+                input_batch.seq_lens_cpu_upper_bound.numpy()[start:end],
+            )
+            chunk_logprobs = self._get_logprobs_tensors(
+                sampled,
+                num_sampled,
+                processed_logits if use_processed_logits else logits[lo:hi],
+                chunk_cu_num_logits,
+                chunk_cu_num_logits_np,
+                max_num_logprobs,
+            )
+            if chunk_logprobs is not None:
+                logprobs_chunks.append(chunk_logprobs)
+            del processed_logits
+            sampled_chunks.append(sampled)
+            num_sampled_chunks.append(num_sampled)
+
+        if len(sampled_chunks) == 1:
+            logprobs_tensors = logprobs_chunks[0] if logprobs_chunks else None
+            return sampled_chunks[0], num_sampled_chunks[0], logprobs_tensors
+
+        logprobs_tensors = None
+        if logprobs_chunks:
+            expanded_logits = logits.shape[0] != input_batch.num_reqs
+            logprobs_tensors = LogprobsTensors.cat(
+                logprobs_chunks,
+                cu_num_generated_tokens=(
+                    cu_num_logits_np.tolist() if expanded_logits else None
+                ),
+            )
+
+        sampled = torch.cat(sampled_chunks)
+        num_sampled = torch.cat(num_sampled_chunks)
+        return sampled, num_sampled, logprobs_tensors
+
+    def __call__(
+        self,
+        logits: torch.Tensor,
+        input_batch: InputBatch,
+        draft_logits: torch.Tensor | None = None,
+    ) -> SamplerOutput:
+        # NOTE(woosuk): We intentionally compute num_nans before sampling to make clear
+        # that num_nans is computed before applying penalties and temperature.
+        num_nans = get_num_nans(logits) if self.sampler.compute_nans else None
+
+        draft_sampled = input_batch.input_ids[input_batch.logits_indices]
+        # A dead draft (draft_trunc) is -1 for acceptance only, so it never matches.
+        accept_drafts = draft_sampled
+        if draft_trunc.ENABLED:
+            accept_drafts = draft_sampled.masked_fill(input_batch.is_padding[input_batch.logits_indices], -1)
+        pos = input_batch.positions[input_batch.logits_indices]
+
+        max_num_logprobs = self.sampler.sampling_states.max_num_logprobs(
+            input_batch.idx_mapping_np
+        )
+        chunk_logit_limit = get_max_chunk_logits(logits.shape[1])
+        sampled, num_sampled, logprobs_tensors = self._verify_in_chunks(
+            logits,
+            input_batch,
+            draft_logits,
+            draft_sampled,
+            pos,
+            chunk_logit_limit,
+            max_num_logprobs,
+            accept_drafts=accept_drafts,
+        )
+
+        num_sampled, num_rejected = get_num_sampled_and_rejected(
+            num_sampled,
+            input_batch.seq_lens,
+            input_batch.cu_num_logits,
+            input_batch.idx_mapping,
+            self.sampler.req_states.prefill_len.gpu,
+        )
+
+        return SamplerOutput(
+            sampled_token_ids=sampled,
+            logprobs_tensors=logprobs_tensors,
+            num_nans=num_nans,
+            num_sampled=num_sampled,
+            num_rejected=num_rejected,
+        )

docker-images/glm53-kindling/experimental/draft-trunc/speculator.py

@@ -0,0 +1,539 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from abc import ABC, abstractmethod
+from collections.abc import Mapping
+from dataclasses import replace
+from typing import TYPE_CHECKING, Any
+
+import os
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from vllm.config import VllmConfig, get_layers_from_vllm_config
+from vllm.config.compilation import CUDAGraphMode
+from vllm.distributed.eplb.eplb_state import EplbState
+from vllm.logger import init_logger
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
+from vllm.model_executor.models import supports_multimodal_embeddings
+from vllm.v1.kv_cache_interface import KVCacheConfig
+from vllm.v1.watermarking import create_watermarker
+from vllm.v1.watermarking.spec_decode import (
+    DraftWatermarker,
+    create_speculative_draft_watermarker,
+)
+from vllm.v1.worker.gpu.attn_utils import (
+    build_attn_metadata,
+    init_attn_backend,
+)
+from vllm.v1.worker.gpu.block_table import BlockTables
+from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
+from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
+from vllm.v1.worker.gpu.dp_utils import DPSyncState
+from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
+from vllm.v1.worker.gpu.model_states.interface import ModelState
+from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
+from vllm.v1.worker.gpu.spec_decode.acceptance_estimator import (
+    OnlineAcceptanceEstimator,
+)
+from vllm.v1.worker.utils import AttentionGroup
+
+if TYPE_CHECKING:
+    from vllm.v1.worker.gpu.pcp_manager import PCPManager
+
+logger = init_logger(__name__)
+
+
+def _target_feeds_hc_residual(vllm_config: VllmConfig) -> bool:
+    """Whether the target replaces the drafter's input with its HC residual.
+
+    Keyed on the same hook the model runner calls to perform that swap. It is
+    resolved from the target model class because speculators are built before
+    the target model is instantiated.
+    """
+    from vllm.model_executor.model_loader.utils import get_model_cls
+
+    target_cls = get_model_cls(vllm_config.model_config)
+    return hasattr(target_cls, "get_mtp_target_hidden_states")
+
+
+class BaseSpeculator(ABC):
+    num_query_per_req: int = 1
+
+    @abstractmethod
+    def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None:
+        pass
+
+    @abstractmethod
+    def capture(self) -> None:
+        pass
+
+    @abstractmethod
+    def propose(
+        self,
+        input_batch: InputBatch,
+        attn_metadata: dict[str, Any],
+        slot_mappings: dict[str, torch.Tensor],
+        # [num_tokens, hidden_size]
+        last_hidden_states: torch.Tensor,
+        # num_layers x [num_tokens, hidden_size]
+        aux_hidden_states: list[torch.Tensor] | None,
+        # [num_reqs]
+        num_sampled: torch.Tensor,
+        # [num_reqs]
+        num_rejected: torch.Tensor,
+        # [max_num_reqs]
+        last_sampled: torch.Tensor,
+        # [num_prefill_lookahead, max_num_reqs]
+        next_prefill_tokens: torch.Tensor,
+        # [max_num_reqs]
+        temperature: torch.Tensor,
+        # [max_num_reqs]
+        seeds: torch.Tensor,
+        dp_sync: DPSyncState | None = None,
+        dummy_run: bool = False,
+        skip_attn_for_dummy_run: bool = False,
+        mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
+        is_profile: bool = False,
+    ) -> torch.Tensor:
+        pass
+
+
+class DraftModelSpeculator(BaseSpeculator):
+    def __init__(self, vllm_config: VllmConfig, device: torch.device):
+        self.vllm_config = vllm_config
+        self.device = device
+
+        assert vllm_config.speculative_config is not None
+        self.speculative_config = vllm_config.speculative_config
+        self.method = self.speculative_config.method
+        self.num_speculative_steps = self.speculative_config.num_speculative_tokens
+        self.draft_model_config = self.speculative_config.draft_model_config
+
+        self.scheduler_config = vllm_config.scheduler_config
+        self.max_num_reqs = self.scheduler_config.max_num_seqs
+        self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
+        self.max_model_len = vllm_config.model_config.max_model_len
+        self.draft_max_seq_len = self.max_model_len
+        # We need to get the hidden size from the draft model config because
+        # the draft model's hidden size can be different from the target model's
+        # hidden size (e.g., Llama 3.3 70B).
+        self.hidden_size = self.draft_model_config.get_hidden_size()
+        # Widen for HC-multiplexed residuals: a target that implements
+        # get_mtp_target_hidden_states() (e.g. DeepSeek V4) hands the drafter
+        # its pre-hc_head (T, hc_mult * hidden_size) residual in place of the
+        # collapsed hidden states, so the drafter's buffers must match. Key off
+        # that hook rather than hc_mult alone -- HY V4 runs iHC in its backbone
+        # (hc_mult=4) but its MTP head consumes the collapsed states, so
+        # widening it feeds propose() a 4x-too-wide buffer.
+        if _target_feeds_hc_residual(vllm_config):
+            hc_mult = getattr(self.draft_model_config.hf_config, "hc_mult", 1)
+            self.hidden_size = self.hidden_size * hc_mult
+        self.vocab_size = self.draft_model_config.get_vocab_size()
+        self.dtype = vllm_config.model_config.dtype
+        self.use_fp64_gumbel = vllm_config.model_config.use_fp64_gumbel
+        self.use_local_argmax_reduction = (
+            self.speculative_config.use_local_argmax_reduction
+        )
+
+        # DP configuration
+        self.dp_size = vllm_config.parallel_config.data_parallel_size
+        self.dp_rank = vllm_config.parallel_config.data_parallel_rank
+
+        self.eplb_state: EplbState | None = None
+
+        self.input_buffers = InputBuffers(
+            max_num_reqs=self.max_num_reqs,
+            max_num_tokens=self.max_num_tokens,
+            device=device,
+        )
+        self.idx_mapping = torch.zeros(
+            self.max_num_reqs, dtype=torch.int32, device=device
+        )
+        self.temperature = torch.zeros(
+            self.max_num_reqs, dtype=torch.float32, device=device
+        )
+        self.seeds = torch.zeros(self.max_num_reqs, dtype=torch.int64, device=device)
+        self.draft_tokens = torch.zeros(
+            self.max_num_reqs,
+            self.num_speculative_steps,
+            dtype=torch.int64,
+            device=device,
+        )
+        self.arange_np = np.arange(self.max_num_reqs + 1, dtype=np.int32)
+        self.draft_is_prefilling = torch.zeros(self.max_num_reqs, dtype=torch.bool)
+
+        self.draft_token_confidence_probs = torch.empty_like(
+            self.draft_tokens, dtype=torch.float32
+        )
+        self.enable_adaptive_verification = (
+            self.speculative_config.enable_adaptive_verification
+        )
+        # Draft truncation's log wants the confidences without the trimming.
+        self.use_acceptance_estimator = self.enable_adaptive_verification or bool(
+            os.environ.get("VLLM_DRAFT_TRUNC_LOG") or os.environ.get("VLLM_DRAFT_TRUNC_CONTROL")
+            or float(os.environ.get("VLLM_DRAFT_TRUNC_TAU") or 0) > 0
+        )
+        self.acceptance_estimator: OnlineAcceptanceEstimator | None = None
+
+        self.draft_logits: torch.Tensor | None = None
+        if self.speculative_config.draft_sample_method == "probabilistic":
+            # Pre-temperature logits, cached from the previous decode step.
+            dtype, fill = self.draft_logits_spec(vllm_config)
+            self.draft_logits = torch.full(
+                (
+                    self.max_num_reqs,
+                    self.num_speculative_steps,
+                    self.vocab_size,
+                ),
+                fill,
+                dtype=dtype,
+                device=device,
+            )
+
+        self.draft_watermarker: DraftWatermarker | None = None
+        watermark_config = getattr(vllm_config, "watermark_config", None)
+        if watermark_config is not None:
+            watermarker = create_watermarker(watermark_config)
+            self.draft_watermarker = create_speculative_draft_watermarker(
+                watermarker,
+                self.max_num_reqs,
+                device,
+                watermark_config.allow_target_only_watermarking,
+            )
+
+        self.supports_mm_inputs = False
+        self.pcp_manager: PCPManager | None = None
+
+    @abstractmethod
+    def load_draft_model(
+        self, target_model: nn.Module, target_attn_layer_names: set[str]
+    ) -> nn.Module:
+        pass
+
+    def load_model(self, target_model: nn.Module) -> None:
+        target_attn_layer_names = set(
+            get_layers_from_vllm_config(
+                self.vllm_config,
+                AttentionLayerBase,  # type: ignore[type-abstract]
+            ).keys()
+        )
+
+        self.model = self.load_draft_model(target_model, target_attn_layer_names)
+        self._validate_local_argmax_reduction()
+
+        all_attn_layers = set[str](
+            get_layers_from_vllm_config(
+                self.vllm_config,
+                AttentionLayerBase,  # type: ignore[type-abstract]
+            ).keys()
+        )
+        self.draft_attn_layer_names = all_attn_layers - target_attn_layer_names
+
+        target_supports_mm = self.vllm_config.model_config.supports_multimodal_inputs
+        draft_supports_mm = supports_multimodal_embeddings(self.model)
+        self.supports_mm_inputs = target_supports_mm and draft_supports_mm
+        if target_supports_mm and not draft_supports_mm:
+            logger.warning_once(
+                "Draft model %s does not support external multimodal embeddings. "
+                "Embeddings from the target model will not be passed to the "
+                "drafter; using text-only draft inputs instead.",
+                type(self.model).__name__,
+            )
+
+        if self.use_acceptance_estimator:
+            if self.use_local_argmax_reduction:
+                raise ValueError(
+                    "Adaptive verification without a confidence head estimates "
+                    "per-position acceptance from the draft logits, which "
+                    "use_local_argmax_reduction never materializes. Disable one "
+                    "of them."
+                )
+            self.acceptance_estimator = OnlineAcceptanceEstimator(
+                self.max_num_reqs,
+                self.num_speculative_steps,
+                self.device,
+            )
+
+    def set_eplb_state(self, eplb_state: EplbState) -> None:
+        """Inject EPLB state after construction."""
+        self.eplb_state = eplb_state
+
+    def _prepare_eplb_forward(self, num_unpadded_tokens: int) -> None:
+        """Call EPLB prepare_forward if EPLB is active for the draft model."""
+        if self.eplb_state is not None:
+            self.eplb_state.prepare_forward(
+                self.speculative_config.draft_model_config,
+                num_unpadded_tokens,
+            )
+
+    @property
+    def attn_vllm_config(self) -> VllmConfig:
+        """Config for the draft's attention metadata builders. Overridden by
+        speculators whose attention mode differs from the target's."""
+        return self.vllm_config
+
+    def set_attn(
+        self,
+        model_state: ModelState,
+        kv_cache_config: KVCacheConfig,
+        block_tables: BlockTables,
+        target_input_buffers: InputBuffers,
+        target_attn_groups: list[list[AttentionGroup]],
+    ) -> None:
+        self.model_state = model_state
+        self.kv_cache_config = kv_cache_config
+        self.attn_groups, self.attn_cg_support, _ = init_attn_backend(
+            kv_cache_config,
+            self.attn_vllm_config,
+            self.device,
+            active_layer_names=self.draft_attn_layer_names,
+        )
+        self.block_tables = block_tables
+        # The target model runner's buffers and attention groups. Draft
+        # prefill reuses the target model's attention metadata, so its
+        # cudagraph capture must build dummy metadata through the same
+        # builders and buffers.
+        self.target_input_buffers = target_input_buffers
+        self.target_attn_groups = target_attn_groups
+
+    def _build_attn_metadata(
+        self,
+        num_reqs: int,
+        batch_desc: BatchExecutionDescriptor,
+        query_start_loc_np: np.ndarray,
+        seq_lens_cpu_upper_bound: torch.Tensor,
+        step: int,
+        causal: bool | Mapping[int, bool] = True,
+        dcp_local_seq_lens: torch.Tensor | None = None,
+    ) -> dict[str, Any] | None:
+        num_reqs_padded = batch_desc.num_reqs or num_reqs
+        # A FULL graph replays a captured shape whose padded requests each hold
+        # a full query width, so attention must see the padded token count.
+        # PIECEWISE/eager needs the actual token count, because batch_desc may
+        # carry graph or DP padding that no request owns, which would desync it
+        # from query_start_loc.
+        num_tokens = (
+            batch_desc.num_tokens
+            if batch_desc.cg_mode == CUDAGraphMode.FULL
+            else int(query_start_loc_np[-1])
+        )
+        query_start_loc_cpu = torch.empty(num_reqs_padded + 1, dtype=torch.int32)
+        query_start_loc_cpu[: num_reqs + 1] = torch.from_numpy(
+            query_start_loc_np[: num_reqs + 1]
+        )
+        query_start_loc_cpu[num_reqs:] = query_start_loc_cpu[num_reqs]
+        max_query_len = int((query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]).max())
+        block_tables = [
+            x[:num_reqs_padded] for x in self.block_tables.input_block_tables
+        ]
+        slot_mappings = self.block_tables.slot_mappings[:, :num_tokens]
+        draft_seq_lens_cpu_upper_bound = torch.zeros(
+            num_reqs_padded, dtype=torch.int32, device="cpu"
+        )
+        torch.add(
+            seq_lens_cpu_upper_bound[:num_reqs],
+            step,
+            out=draft_seq_lens_cpu_upper_bound[:num_reqs],
+        )
+        draft_seq_lens_cpu_upper_bound[:num_reqs].clamp_(max=self.max_model_len)
+        if dcp_local_seq_lens is None and self.block_tables.cp_size > 1:
+            # Draft steps advance and rewind their own global sequence lengths,
+            # so the target model's DCP-local lengths may already be stale.
+            dcp_local_seq_lens = prepare_dcp_local_seq_lens(
+                self.input_buffers.dcp_local_seq_lens,
+                self.input_buffers.seq_lens,
+                num_reqs,
+                self.block_tables.cp_size,
+                self.block_tables.cp_rank,
+                self.block_tables.cp_interleave,
+            )
+        attn_metadata = build_attn_metadata(
+            attn_groups=self.attn_groups,
+            num_reqs=num_reqs_padded,
+            num_tokens=num_tokens,
+            query_start_loc_gpu=self.input_buffers.query_start_loc[
+                : num_reqs_padded + 1
+            ],
+            query_start_loc_cpu=query_start_loc_cpu,
+            max_query_len=max_query_len,
+            seq_lens=self.input_buffers.seq_lens[:num_reqs_padded],
+            dcp_local_seq_lens=(
+                None
+                if dcp_local_seq_lens is None
+                else dcp_local_seq_lens[:num_reqs_padded]
+            ),
+            max_seq_len=self.draft_max_seq_len,
+            block_tables=block_tables,
+            slot_mappings=slot_mappings,
+            kv_cache_config=self.kv_cache_config,
+            causal=causal,
+            seq_lens_cpu_upper_bound=draft_seq_lens_cpu_upper_bound,
+            is_prefilling=self.draft_is_prefilling[:num_reqs_padded],
+        )
+        return attn_metadata
+
+    def draft_logits_spec(self, vllm_config: VllmConfig) -> tuple[torch.dtype, float]:
+        """Dtype and fill for the cached proposal distribution.
+
+        Speculators that write only a subset of columns each step override this.
+        """
+        return vllm_config.model_config.head_dtype, 0.0
+
+    def _validate_local_argmax_reduction(self) -> None:
+        if not self.use_local_argmax_reduction:
+            return
+        if self.speculative_config.draft_sample_method == "probabilistic":
+            raise ValueError(
+                "use_local_argmax_reduction is not compatible with "
+                "draft_sample_method='probabilistic'."
+            )
+        if not hasattr(self.model, "get_top_tokens"):
+            raise ValueError(
+                "use_local_argmax_reduction is enabled but draft model "
+                f"{self.model.__class__.__name__} does not implement "
+                "get_top_tokens()."
+            )
+        logger.info(
+            "Using local argmax reduction for draft token generation "
+            "(communication: O(2*tp_size) vs O(vocab_size))."
+        )
+
+    def sample_draft(
+        self,
+        hidden_states: torch.Tensor,
+        sample_src_positions: torch.Tensor,
+        idx_mapping: torch.Tensor,
+        temperature: torch.Tensor,
+        seeds: torch.Tensor,
+        draft_step: torch.Tensor,
+        draft_logits: torch.Tensor | None,
+    ) -> torch.Tensor:
+        if draft_logits is not None:
+            logits = self.model.compute_logits(hidden_states)
+            sampled = gumbel_sample(
+                logits,
+                idx_mapping,
+                temperature,
+                seeds,
+                sample_src_positions,
+                apply_temperature=True,
+                is_drafting=True,
+                logits_cache=draft_logits,
+                logits_cache_col=draft_step,
+                use_fp64=self.use_fp64_gumbel,
+            )
+            if self.draft_watermarker is not None:
+                sampled = self.draft_watermarker.sample(
+                    logits, sampled, idx_mapping, temperature
+                )
+        elif self.use_local_argmax_reduction:
+            return self.model.get_top_tokens(hidden_states)
+        else:
+            logits = self.model.compute_logits(hidden_states)
+            sampled = logits.argmax(dim=-1)
+        self._maybe_predict_acceptance(logits, idx_mapping, draft_step)
+        return sampled
+
+    def _maybe_predict_acceptance(
+        self, logits: torch.Tensor, idx_mapping: torch.Tensor, draft_step: torch.Tensor
+    ) -> None:
+        if self.acceptance_estimator is not None:
+            self.acceptance_estimator.predict(
+                logits,
+                idx_mapping,
+                draft_step,
+                self.draft_token_confidence_probs,
+                self.temperature,
+            )
+
+    def observe_verification(
+        self,
+        idx_mapping: torch.Tensor,
+        num_sampled: torch.Tensor,
+        num_rejected: torch.Tensor,
+    ) -> None:
+        """Fold the target's verdict on the last drafts into the estimator.
+
+        Must run before the next `propose`, which overwrites the per-slot
+        features the verdict grades.
+        """
+        if self.acceptance_estimator is not None:
+            self.acceptance_estimator.step(idx_mapping, num_sampled, num_rejected)
+
+    def prepare_watermarking(
+        self, contexts: torch.Tensor, watermarking: torch.Tensor
+    ) -> None:
+        if self.draft_watermarker is None:
+            return
+        self.draft_watermarker.prepare(contexts, watermarking)
+
+    def _copy_request_inputs(
+        self,
+        num_reqs: int,
+        # [num_reqs]
+        idx_mapping: torch.Tensor,
+        # [max_num_reqs]
+        temperature: torch.Tensor,
+        # [max_num_reqs]
+        seeds: torch.Tensor,
+        dummy_run: bool = False,
+    ) -> None:
+        # Copy temperature, seeds, and idx mapping to the pre-allocated buffers.
+        # NOTE(woosuk): For draft sampling, we only consider the temperature
+        # and ignore the other sampling parameters such as top_k and top_p,
+        # for simplicity and performance.
+        # While this may slightly degrade the acceptance rate, it does not
+        # affect the output distribution after rejection sampling.
+        self.temperature.copy_(temperature)
+        self.seeds.copy_(seeds)
+        # idx_mapping == -1 marks a row the drafter must not act on: sampling
+        # skips it and compute_slot_mappings emits PAD. CUDA-graph padded rows
+        # always get it; a dummy batch gets it for every row, since its arange
+        # idx_mapping names request-state slots with stale block tables.
+        if dummy_run:
+            self.idx_mapping.fill_(-1)
+        else:
+            self.idx_mapping[:num_reqs].copy_(idx_mapping)
+            self.idx_mapping[num_reqs:].fill_(-1)
+
+    def _build_uniform_batch_dp_sync(
+        self,
+        target_dp_sync: DPSyncState,
+        num_reqs: int,
+        num_query_per_req: int = 1,
+    ) -> tuple[DPSyncState, int]:
+        num_batch_tokens = target_dp_sync.num_reqs * num_query_per_req
+        assert num_reqs * num_query_per_req <= num_batch_tokens, (
+            "reusing a DP sync that does not cover this batch's requests"
+        )
+        return replace(
+            target_dp_sync,
+            num_tokens_across_dp=torch.full_like(
+                target_dp_sync.num_tokens_across_dp, num_batch_tokens
+            ),
+            uniform_token_count=num_query_per_req,
+            eager=False,
+        ), num_batch_tokens
+
+    def _build_uniform_attn_metadata(
+        self,
+        batch_desc: BatchExecutionDescriptor,
+        num_reqs: int,
+        num_query_per_req: int,
+        seq_lens_cpu_upper_bound: torch.Tensor,
+        step: int,
+        causal: bool | Mapping[int, bool] = True,
+        dcp_local_seq_lens: torch.Tensor | None = None,
+    ) -> dict[str, Any] | None:
+        query_start_loc_np = self.arange_np[: num_reqs + 1] * num_query_per_req
+        return self._build_attn_metadata(
+            num_reqs=num_reqs,
+            batch_desc=batch_desc,
+            query_start_loc_np=query_start_loc_np,
+            seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
+            step=step,
+            causal=causal,
+            dcp_local_seq_lens=dcp_local_seq_lens,
+        )

docker-images/glm53-kindling/experimental/fixes/_fa_common.py

@@ -0,0 +1,638 @@
+"""
+Copyright (c) 2026 by FlashInfer team.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+"""
+
+import functools
+import math
+from typing import Callable, ClassVar, Optional, Protocol, Tuple, TypeVar, Union, cast
+
+import torch
+
+from ....jit import gen_batch_mla_module
+from ....utils import MaskMode, check_shape_dtype_device, get_compute_capability
+from ._capabilities import MLAPlanCapabilities, plan_capability_rejection_reason
+from .._planning import _MLAPlanArguments
+
+
+class _GeneratedBatchMLAModule(Protocol):
+    def plan(self, *args: object) -> object: ...
+
+    def plan_with_staged_workspace_bytes(self, *args: object) -> tuple[object, int]: ...
+
+    def run(self, *args: object) -> object: ...
+
+
+@functools.cache
+def get_batch_mla_module(
+    backend: str,
+    dtype_q: torch.dtype,
+    dtype_kv: torch.dtype,
+    dtype_o: torch.dtype,
+    dtype_idx: torch.dtype,
+    head_dim_ckv: int,
+    head_dim_kpe: int,
+    use_profiler: bool,
+) -> _GeneratedBatchMLAModule:
+    return gen_batch_mla_module(
+        backend,
+        dtype_q,
+        dtype_kv,
+        dtype_o,
+        dtype_idx,
+        head_dim_ckv,
+        head_dim_kpe,
+        use_profiler,
+    ).build_and_load()
+
+
+def _validate_generated_fa_plan(
+    *,
+    backend: str,
+    device: torch.device,
+    qo_indptr: torch.Tensor,
+    kv_indptr: torch.Tensor,
+    kv_indices: torch.Tensor,
+    kv_len_arr: torch.Tensor,
+    head_dim_ckv: int,
+    head_dim_kpe: int,
+    q_data_type: torch.dtype,
+    kv_data_type: torch.dtype,
+    output_dtype: torch.dtype,
+    scale_mode: str,
+) -> None:
+    for name, tensor in (
+        ("qo_indptr", qo_indptr),
+        ("kv_indptr", kv_indptr),
+        ("kv_indices", kv_indices),
+        ("kv_len_arr", kv_len_arr),
+    ):
+        if tensor.dtype != torch.int32:
+            raise ValueError(f"{name} must have dtype torch.int32, got {tensor.dtype}.")
+    if q_data_type not in (torch.float16, torch.bfloat16):
+        raise ValueError(
+            f"MLA q_data_type {q_data_type} is not supported by the {backend} backend."
+        )
+    supported_kv_dtypes = (torch.float16, torch.bfloat16, torch.float8_e4m3fn)
+    if kv_data_type not in supported_kv_dtypes:
+        raise ValueError(
+            f"MLA kv_data_type {kv_data_type} is not supported by the {backend} "
+            f"backend. Supported dtypes: {list(supported_kv_dtypes)}."
+        )
+    if output_dtype != q_data_type:
+        raise ValueError(
+            f"{backend} backend output_dtype must match q_data_type, got "
+            f"{output_dtype} and {q_data_type}."
+        )
+    if head_dim_kpe < 0:
+        raise ValueError(f"head_dim_kpe must be >= 0, got {head_dim_kpe}.")
+    if kv_data_type == torch.float8_e4m3fn:
+        major, minor = get_compute_capability(device)
+        if major not in (9, 12):
+            raise ValueError(
+                "FP8 kv_data_type for MLA requires an SM90 (Hopper) device, "
+                f"got SM{major}{minor}."
+            )
+        if q_data_type != torch.bfloat16:
+            raise ValueError(
+                "FP8 kv_data_type for MLA currently only supports "
+                f"q_data_type=torch.bfloat16, got {q_data_type}."
+            )
+        if head_dim_ckv != 512 or head_dim_kpe not in (0, 64):
+            raise ValueError(
+                "FP8 kv_data_type for MLA currently only supports "
+                "head_dim_ckv=512 and head_dim_kpe in (0, 64), got "
+                f"head_dim_ckv={head_dim_ckv}, head_dim_kpe={head_dim_kpe}."
+            )
+        if scale_mode != "kv-per-tensor":
+            raise ValueError("FP8 MLA plans require scale_mode='kv-per-tensor'.")
+    elif scale_mode != "default":
+        raise ValueError("non-FP8 MLA plans require scale_mode='default'.")
+
+
+class _BatchMLAGeneratedFaMechanics:
+    def __init__(
+        self,
+        *,
+        backend: Optional[str] = None,
+        float_workspace_buffer: torch.Tensor,
+        use_cuda_graph: bool,
+        qo_indptr_buf: Optional[torch.Tensor],
+        kv_indptr_buf: Optional[torch.Tensor],
+        kv_indices_buf: Optional[torch.Tensor],
+        kv_len_arr_buf: Optional[torch.Tensor],
+        int_workspace_buffer: Optional[torch.Tensor] = None,
+        pin_memory_int_workspace_buffer: Optional[torch.Tensor] = None,
+    ) -> None:
+        self._backend = backend
+        self._float_workspace_buffer = float_workspace_buffer
+        self.device = float_workspace_buffer.device
+        if int_workspace_buffer is None:
+            int_workspace_buffer = torch.empty(
+                (8 * 1024 * 1024,), dtype=torch.uint8, device=self.device
+            )
+        if pin_memory_int_workspace_buffer is None:
+            pin_memory_int_workspace_buffer = torch.empty(
+                int_workspace_buffer.shape,
+                dtype=int_workspace_buffer.dtype,
+                pin_memory=True,
+                device="cpu",
+            )
+        self._int_workspace_buffer = int_workspace_buffer
+        self._pin_memory_int_workspace_buffer = pin_memory_int_workspace_buffer
+        self._use_cuda_graph = use_cuda_graph
+        self._qo_indptr_buf = qo_indptr_buf
+        self._kv_indptr_buf = kv_indptr_buf
+        self._kv_indices_buf = kv_indices_buf
+        self._kv_len_arr_buf = kv_len_arr_buf
+
+    @staticmethod
+    def _storage_interval(tensor: torch.Tensor) -> tuple[int, int, int]:
+        start = tensor.storage_offset() * tensor.element_size()
+        return (
+            tensor.untyped_storage().data_ptr(),
+            start,
+            start + tensor.numel() * tensor.element_size(),
+        )
+
+    def _preflight_graph_metadata_buffers(
+        self,
+        *,
+        qo_indptr: torch.Tensor,
+        kv_indptr: torch.Tensor,
+        kv_indices: torch.Tensor,
+        kv_len_arr: torch.Tensor,
+    ) -> None:
+        named = (
+            ("qo_indptr", self._qo_indptr_buf, qo_indptr, False),
+            ("kv_indptr", self._kv_indptr_buf, kv_indptr, False),
+            ("kv_indices", self._kv_indices_buf, kv_indices, True),
+            ("kv_len_arr", self._kv_len_arr_buf, kv_len_arr, False),
+        )
+        for name, reserved, source, allow_larger in named:
+            if reserved is None:
+                raise ValueError(
+                    "CUDA graph mode requires reserved qo_indptr, kv_indptr, "
+                    "kv_indices, and kv_len_arr buffers."
+                )
+            if reserved.dtype != torch.int32:
+                raise ValueError(
+                    f"CUDA graph reserved {name} buffer must have dtype torch.int32."
+                )
+            if reserved.device != self.device:
+                raise ValueError(
+                    f"CUDA graph reserved {name} buffer must be on {self.device}."
+                )
+            if not reserved.is_contiguous():
+                raise ValueError(
+                    f"CUDA graph reserved {name} buffer must be contiguous."
+                )
+            if name == "kv_indices" and reserved.ndim != 1:
+                raise ValueError(
+                    "CUDA graph reserved kv_indices buffer must have rank 1."
+                )
+            if (allow_larger and reserved.shape[0] < source.shape[0]) or (
+                not allow_larger and reserved.shape != source.shape
+            ):
+                raise ValueError(
+                    f"CUDA graph reserved {name} buffer has insufficient or "
+                    "incompatible capacity."
+                )
+        intervals = [
+            (name, self._storage_interval(reserved))
+            for name, reserved, _, _ in named
+            if reserved is not None
+        ]
+        for index, (left_name, (left_ptr, left_start, left_end)) in enumerate(
+            intervals
+        ):
+            for right_name, (right_ptr, right_start, right_end) in intervals[
+                index + 1 :
+            ]:
+                if left_ptr == right_ptr and max(left_start, right_start) < min(
+                    left_end, right_end
+                ):
+                    raise ValueError(
+                        "CUDA graph reserved metadata buffers must not overlap: "
+                        f"{left_name} overlaps {right_name}."
+                    )
+        copy_targets = {
+            name: reserved[: source.shape[0]] if allow_larger else reserved
+            for name, reserved, source, allow_larger in named
+            if reserved is not None
+        }
+        for source_name, _, source, _ in named:
+            source_interval = self._storage_interval(source)
+            for target_name, target in copy_targets.items():
+                target_interval = self._storage_interval(target)
+                overlaps = source_interval[0] == target_interval[0] and max(
+                    source_interval[1], target_interval[1]
+                ) < min(source_interval[2], target_interval[2])
+                exact_corresponding_reuse = (
+                    source_name == target_name and source_interval == target_interval
+                )
+                if overlaps and not exact_corresponding_reuse:
+                    raise ValueError(
+                        f"CUDA graph {source_name} source overlaps reserved "
+                        f"{target_name} target; only exact same-source/same-target "
+                        "reuse is allowed."
+                    )
+
+    def _stage_metadata(
+        self,
+        *,
+        qo_indptr: torch.Tensor,
+        kv_indptr: torch.Tensor,
+        kv_indices: torch.Tensor,
+        kv_len_arr: torch.Tensor,
+    ) -> None:
+        if self._use_cuda_graph:
+            self._preflight_graph_metadata_buffers(
+                qo_indptr=qo_indptr,
+                kv_indptr=kv_indptr,
+                kv_indices=kv_indices,
+                kv_len_arr=kv_len_arr,
+            )
+            assert self._qo_indptr_buf is not None
+            assert self._kv_indptr_buf is not None
+            assert self._kv_indices_buf is not None
+            assert self._kv_len_arr_buf is not None
+            targets = (
+                (self._qo_indptr_buf, qo_indptr),
+                (self._kv_indptr_buf, kv_indptr),
+                (self._kv_indices_buf[: len(kv_indices)], kv_indices),
+                (self._kv_len_arr_buf, kv_len_arr),
+            )
+            # GB10 fix: no rollback snapshot. It cloned kv_indices_buf on every
+            # plan(), 136 MB per decode step for GLM's sparse MLA, to undo a
+            # copy that the preflight above has already made safe.
+            for target, source in targets:
+                target.copy_(source, non_blocking=True)
+        else:
+            self._qo_indptr_buf = qo_indptr.to(self.device, non_blocking=True)
+            self._kv_indptr_buf = kv_indptr.to(self.device, non_blocking=True)
+            self._kv_indices_buf = kv_indices.to(self.device, non_blocking=True)
+            self._kv_len_arr_buf = kv_len_arr.to(self.device, non_blocking=True)
+
+    def _plan_generated_fa(
+        self,
+        *,
+        module_loader: Callable[[], _GeneratedBatchMLAModule],
+        qo_indptr: torch.Tensor,
+        kv_indptr: torch.Tensor,
+        kv_indices: torch.Tensor,
+        kv_len_arr: torch.Tensor,
+        num_heads: int,
+        head_dim_ckv: int,
+        page_size: int,
+        causal: bool,
+        sm_scale: float,
+        q_data_type: torch.dtype,
+        kv_data_type: torch.dtype,
+        use_profiler: bool,
+    ) -> None:
+        # ---------------------------------------------------------------------------
+        # Build the generated backend plan
+        # ---------------------------------------------------------------------------
+        cached_module = module_loader()
+        qo_indptr_host = qo_indptr.to("cpu")
+        kv_indptr_host = kv_indptr.to("cpu")
+        kv_len_arr_host = kv_len_arr.to("cpu")
+        plan_args = (
+            self._float_workspace_buffer,
+            self._int_workspace_buffer,
+            self._pin_memory_int_workspace_buffer,
+            qo_indptr_host,
+            kv_indptr_host,
+            kv_len_arr_host,
+            num_heads,
+            head_dim_ckv,
+            causal,
+        )
+        if hasattr(cached_module, "plan_with_staged_workspace_bytes"):
+            plan_info, staged_int_workspace_bytes = (
+                cached_module.plan_with_staged_workspace_bytes(*plan_args)
+            )
+        else:
+            # Compatibility for externally cached legacy modules; generated modules
+            # are always adapted at load time.
+            plan_info = cached_module.plan(*plan_args)
+            staged_int_workspace_bytes = 0
+
+        # ---------------------------------------------------------------------------
+        # Stage metadata and publish backend state
+        # ---------------------------------------------------------------------------
+        self._stage_metadata(
+            qo_indptr=qo_indptr,
+            kv_indptr=kv_indptr,
+            kv_indices=kv_indices,
+            kv_len_arr=kv_len_arr,
+        )
+        self._cached_module = cached_module
+        self._causal = causal
+        self._page_size = page_size
+        self._sm_scale = sm_scale
+        self._head_dim_ckv = head_dim_ckv
+        self._q_data_type = q_data_type
+        self._kv_data_type = kv_data_type
+        self._use_profiler = use_profiler
+        self._plan_info = plan_info
+        self._staged_int_workspace_bytes = staged_int_workspace_bytes
+
+    def _validate_run_input_dtypes(
+        self,
+        *,
+        q_nope: torch.Tensor,
+        q_pe: torch.Tensor,
+        ckv_cache: torch.Tensor,
+        kpe_cache: torch.Tensor,
+    ) -> None:
+        if q_nope.dtype != self._q_data_type:
+            raise ValueError(
+                f"q_nope.dtype={q_nope.dtype} does not match the planned "
+                f"q_data_type={self._q_data_type}."
+            )
+        if q_pe.dtype != self._q_data_type:
+            raise ValueError(
+                f"q_pe.dtype={q_pe.dtype} does not match the planned "
+                f"q_data_type={self._q_data_type}."
+            )
+        if ckv_cache.dtype != self._kv_data_type:
+            raise ValueError(
+                f"ckv_cache.dtype={ckv_cache.dtype} does not match the planned "
+                f"kv_data_type={self._kv_data_type}."
+            )
+        if kpe_cache.dtype != self._kv_data_type:
+            raise ValueError(
+                f"kpe_cache.dtype={kpe_cache.dtype} does not match the planned "
+                f"kv_data_type={self._kv_data_type}."
+            )
+
+    def _run_generated_fa(
+        self,
+        *,
+        q_nope: torch.Tensor,
+        q_pe: torch.Tensor,
+        ckv_cache: torch.Tensor,
+        kpe_cache: torch.Tensor,
+        out: Optional[torch.Tensor],
+        lse: Optional[torch.Tensor],
+        return_lse: bool,
+        profiler_buffer: Optional[torch.Tensor],
+        return_lse_base_on_e: bool,
+        ckv_scale: Optional[float],
+        ckv_scale_arr: Optional[torch.Tensor],
+        kpe_scale: Optional[float],
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        # ---------------------------------------------------------------------------
+        # Validate inputs and scale arguments
+        # ---------------------------------------------------------------------------
+        self._validate_run_input_dtypes(
+            q_nope=q_nope,
+            q_pe=q_pe,
+            ckv_cache=ckv_cache,
+            kpe_cache=kpe_cache,
+        )
+        kv_is_fp8 = self._kv_data_type == torch.float8_e4m3fn
+        if kv_is_fp8:
+            if (ckv_scale is None) == (ckv_scale_arr is None):
+                raise ValueError(
+                    "Exactly one of ckv_scale or ckv_scale_arr is required when "
+                    "kv_data_type is FP8."
+                )
+            if kpe_scale is None:
+                raise ValueError("kpe_scale is required when kv_data_type is FP8.")
+            ckv_scale_f = 1.0 if ckv_scale is None else float(ckv_scale)
+            kpe_scale_f = float(kpe_scale)
+            if ckv_scale is not None and (
+                not math.isfinite(ckv_scale_f) or ckv_scale_f <= 0.0
+            ):
+                raise ValueError(
+                    f"ckv_scale must be a finite positive value, got {ckv_scale}"
+                )
+            if not math.isfinite(kpe_scale_f) or kpe_scale_f <= 0.0:
+                raise ValueError(
+                    f"kpe_scale must be a finite positive value, got {kpe_scale}"
+                )
+        else:
+            if (
+                ckv_scale is not None
+                or ckv_scale_arr is not None
+                or kpe_scale is not None
+            ):
+                raise ValueError(
+                    "ckv_scale / ckv_scale_arr / kpe_scale are only valid when "
+                    "kv_data_type is FP8."
+                )
+            ckv_scale_f = 1.0
+            kpe_scale_f = 1.0
+
+        # ---------------------------------------------------------------------------
+        # Prepare output and auxiliary buffers
+        # ---------------------------------------------------------------------------
+        if profiler_buffer is None and self._use_profiler:
+            raise ValueError("Profiler is enabled, profiler_buffer must be provided")
+        if out is None:
+            out = torch.empty_like(q_nope)
+        else:
+            check_shape_dtype_device(
+                out, q_nope.shape, q_nope.dtype, q_nope.device, "out"
+            )
+        if return_lse:
+            if lse is None:
+                lse = torch.empty(
+                    q_nope.shape[:2], dtype=torch.float32, device=self.device
+                )
+            else:
+                check_shape_dtype_device(
+                    lse, q_nope.shape[:2], torch.float32, q_nope.device, "lse"
+                )
+        if ckv_scale_arr is not None:
+            expected_scale_shape = (*ckv_cache.shape[:-1], self._head_dim_ckv // 128)
+            check_shape_dtype_device(
+                ckv_scale_arr,
+                expected_scale_shape,
+                torch.float32,
+                ckv_cache.device,
+                "ckv_scale_arr",
+            )
+            if not ckv_scale_arr.is_contiguous():
+                raise ValueError("ckv_scale_arr must be contiguous.")
+
+        # ---------------------------------------------------------------------------
+        # Launch the generated backend
+        # ---------------------------------------------------------------------------
+        mask_mode = MaskMode.CAUSAL.value if self._causal else MaskMode.NON_CAUSAL.value
+        profiler_args = (profiler_buffer,) if self._use_profiler else ()
+        self._cached_module.run(
+            self._float_workspace_buffer,
+            self._int_workspace_buffer,
+            self._plan_info,
+            q_nope,
+            q_pe,
+            ckv_cache,
+            kpe_cache,
+            self._kv_indices_buf,
+            out,
+            lse,
+            mask_mode,
+            q_nope.shape[1],
+            self._page_size,
+            self._sm_scale,
+            return_lse_base_on_e,
+            ckv_scale_f,
+            kpe_scale_f,
+            ckv_scale_arr,
+            *profiler_args,
+        )
+        return (out, lse) if return_lse else out
+
+
+_FaBackendT = TypeVar("_FaBackendT", bound="_BatchMLAPagedAttentionFaBackendBase")
+
+
+class _BatchMLAPagedAttentionFaBackendBase(_BatchMLAGeneratedFaMechanics):
+    _plan_capabilities: ClassVar[Optional[MLAPlanCapabilities]] = None
+
+    def __init__(
+        self,
+        *,
+        backend: Optional[str] = None,
+        float_workspace_buffer: torch.Tensor,
+        use_cuda_graph: bool,
+        qo_indptr_buf: Optional[torch.Tensor],
+        kv_indptr_buf: Optional[torch.Tensor],
+        kv_indices_buf: Optional[torch.Tensor],
+        kv_len_arr_buf: Optional[torch.Tensor],
+        query_split_widths: tuple[int, int],
+        kv_split_widths: tuple[int, int],
+        int_workspace_buffer: Optional[torch.Tensor] = None,
+        pin_memory_int_workspace_buffer: Optional[torch.Tensor] = None,
+    ) -> None:
+        if backend is None:
+            if self._plan_capabilities is None:
+                raise TypeError("generated-FA backend capabilities are required.")
+            backend = self._plan_capabilities.backend_name
+        super().__init__(
+            backend=backend,
+            float_workspace_buffer=float_workspace_buffer,
+            use_cuda_graph=use_cuda_graph,
+            qo_indptr_buf=qo_indptr_buf,
+            kv_indptr_buf=kv_indptr_buf,
+            kv_indices_buf=kv_indices_buf,
+            kv_len_arr_buf=kv_len_arr_buf,
+            int_workspace_buffer=int_workspace_buffer,
+            pin_memory_int_workspace_buffer=pin_memory_int_workspace_buffer,
+        )
+        self._query_split_widths = query_split_widths
+        self._kv_split_widths = kv_split_widths
+
+    def plan(
+        self,
+        *,
+        qo_indptr: torch.Tensor,
+        kv_indptr: torch.Tensor,
+        kv_indices: torch.Tensor,
+        kv_len_arr: torch.Tensor,
+        num_heads: int,
+        head_dim_ckv: int,
+        head_dim_kpe: int,
+        page_size: int,
+        causal: bool,
+        sm_scale: float,
+        q_data_type: torch.dtype,
+        kv_data_type: torch.dtype,
+        output_dtype: torch.dtype,
+        scale_mode: str,
+        use_profiler: bool,
+    ) -> None:
+        raise NotImplementedError
+
+    @classmethod
+    def plan_from_wrapper(
+        cls: type[_FaBackendT], args: _MLAPlanArguments
+    ) -> _FaBackendT:
+        assert cls._plan_capabilities is not None
+        if reason := plan_capability_rejection_reason(args, cls._plan_capabilities):
+            raise ValueError(reason)
+        csr = args.csr()
+        backend = cls(
+            float_workspace_buffer=args._float_workspace_buffer,
+            use_cuda_graph=args._use_cuda_graph,
+            qo_indptr_buf=args._qo_indptr_buf,
+            kv_indptr_buf=args._kv_indptr_buf,
+            kv_indices_buf=args._kv_indices_buf,
+            kv_len_arr_buf=args._kv_len_arr_buf,
+            query_split_widths=(args.head_dim_ckv, args.head_dim_kpe),
+            kv_split_widths=(args.head_dim_ckv, args.head_dim_kpe),
+            int_workspace_buffer=args._graph_plan_int_workspace_buffer,
+        )
+        if args._use_cuda_graph:
+            backend._preflight_graph_metadata_buffers(
+                qo_indptr=csr.qo_indptr,
+                kv_indptr=csr.kv_indptr,
+                kv_indices=csr.kv_indices,
+                kv_len_arr=csr.kv_len_arr,
+            )
+        backend.plan(
+            qo_indptr=csr.qo_indptr,
+            kv_indptr=csr.kv_indptr,
+            kv_indices=csr.kv_indices,
+            kv_len_arr=csr.kv_len_arr,
+            num_heads=args.num_heads,
+            head_dim_ckv=args.head_dim_ckv,
+            head_dim_kpe=args.head_dim_kpe,
+            page_size=args.page_size,
+            causal=args.causal,
+            sm_scale=args.sm_scale,
+            q_data_type=args.q_data_type,
+            kv_data_type=args.kv_data_type,
+            output_dtype=args.output_dtype,
+            scale_mode=args.scale_mode,
+            use_profiler=args.use_profiler,
+        )
+        return backend
+
+    def run_from_wrapper(
+        self,
+        *,
+        query: object,
+        kv_cache: object,
+        out: Optional[torch.Tensor],
+        lse: Optional[torch.Tensor],
+        return_lse: bool,
+        profiler_buffer: Optional[torch.Tensor],
+        kv_len: Optional[torch.Tensor],
+        page_table: Optional[torch.Tensor],
+        return_lse_base_on_e: bool,
+        o_scale: Optional[float],
+        ckv_scale: Optional[float],
+        ckv_scale_arr: Optional[torch.Tensor],
+        kpe_scale: Optional[float],
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        if kv_len is not None:
+            raise ValueError("kv_len is only supported with cutlass backend.")
+        if page_table is not None:
+            raise ValueError("page_table is only supported with cutlass backend.")
+        if o_scale is not None:
+            raise ValueError(
+                "o_scale is only supported with the cutlass backend for now."
+            )
+        q_nope, q_pe = cast(tuple[torch.Tensor, torch.Tensor], query)
+        ckv_cache, kpe_cache = cast(tuple[torch.Tensor, torch.Tensor], kv_cache)
+        return self._run_generated_fa(
+            q_nope=q_nope,
+            q_pe=q_pe,
+            ckv_cache=ckv_cache,
+            kpe_cache=kpe_cache,
+            out=out,
+            lse=lse,
+            return_lse=return_lse,
+            profiler_buffer=profiler_buffer,
+            return_lse_base_on_e=return_lse_base_on_e,
+            ckv_scale=ckv_scale,
+            ckv_scale_arr=ckv_scale_arr,
+            kpe_scale=kpe_scale,
+        )

docker-images/glm53-kindling/experimental/fixes/attention.py

@@ -0,0 +1,600 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from vllm.config import (
+    CacheConfig,
+    VllmConfig,
+)
+from vllm.distributed import (
+    get_tensor_model_parallel_world_size,
+)
+from vllm.logger import init_logger
+from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    MergedColumnParallelLinear,
+    ReplicatedLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.mla import MLAModules, MultiHeadLatentAttentionWrapper
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding, get_rope
+from vllm.model_executor.models.deepseek_v2 import (
+    DeepSeekV2FusedQkvAProjLinear,
+    DeepseekV32IndexerCache,
+    yarn_get_mscale,
+)
+from vllm.model_executor.utils import maybe_disable_graph_partition
+from vllm.models.glm5next.nvidia.ops.kpool_compress import fwht128_quant_fp8
+from vllm.models.glm5next.sparse_indexer import SparseAttnIndexerKpool
+from vllm.platforms import current_platform
+from vllm.transformers_utils.configs.glm5_next import Glm5NextConfig
+from vllm.utils.deep_gemm import PAGED_MQA_PAGE_SIZES
+from vllm.utils.math_utils import cdiv, next_power_of_2
+from vllm.v1.kv_cache_interface import KpoolTailSpec, MLAAttentionSpec
+
+logger = init_logger(__name__)
+
+# Shared torch.compile config for the indexer's small-kernel leaves. The MLA
+# indexer runs under breakable-CG (CompilationMode.NONE), which blocks FX-graph
+# fusion of the surrounding eager ops; carving each cluster into its own
+# @torch.compile leaf (backend==inductor) still fuses them. Matches the
+# grouped_topk / _cast_sigmoid leaf pattern.
+_INDEXER_COMPILE = dict(
+    dynamic=True,
+    backend=current_platform.simple_compile_backend,
+    options=maybe_disable_graph_partition(current_platform.simple_compile_backend),
+)
+
+
+@torch.compile(**_INDEXER_COMPILE)
+def _fused_indexer_k_norm(
+    x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, dim: int, eps: float
+) -> torch.Tensor:
+    # Fuse fp32 cast + layer_norm + cast-back (was 3 kernels) into one.
+    return F.layer_norm(x.float(), (dim,), weight, bias, eps).type_as(x)
+
+
+@torch.compile(**_INDEXER_COMPILE)
+def _fused_indexer_weight_scale(
+    weights: torch.Tensor, q_scale: torch.Tensor, scale: float
+) -> torch.Tensor:
+    # Fuse the weight-scaling muls (was 2 kernels) into one. `scale` folds
+    # softmax_scale (head_dim**-0.5) and n_head**-0.5 into a single constant.
+    return (weights.unsqueeze(-1) * q_scale * scale).squeeze(-1)
+
+
+@torch.compile(**_INDEXER_COMPILE)
+def _pad_indexer_heads(x: torch.Tensor, pad: int) -> torch.Tensor:
+    # DeepGEMM MQA-logits needs num_heads in {32,64}; zero-pad the head dim.
+    # Fuse new_zeros + cat (was 2 kernels) into one. Pad values are zero (exact
+    # in fp8 e4m3 and zero-weight in the logits sum), so numerically a no-op.
+    return torch.cat([x, x.new_zeros(x.shape[0], pad, *x.shape[2:])], dim=1)
+
+
+class Glm5NextIndexerCache(DeepseekV32IndexerCache):
+    """Indexer K cache that stores kpool-compressed entries.
+
+    Setting ``tokens_per_state = index_kpool`` on the KV cache spec makes vLLM's
+    indexer metadata builder emit pool-granular ``slot_mapping`` /
+    ``seq_lens`` / ``cu_seq_lens`` / ``page_table`` for free, and shrinks the
+    cache allocation store one state per ``index_kpool`` tokens. The pool
+    *content* (softmax-weighted sum vs keep-every-Nth) is computed by the
+    kpool compress kernel inside the indexer op — the cache only provides the
+    addressing, which is identical for both schemes.
+
+    The indexer shares one block with the co-located MLA (a single
+    ``MLAAttentionSpec`` / block_table), so ``block_size`` is the model-wide
+    ``cache_config.block_size``. DeepGEMM's paged-MQA kernel
+    (``csrc/apis/attention.hpp``) requires ``block_kv`` to be exactly 32 or
+    64, so the storage block is virtually split into pool pages of the
+    largest such size that tiles it (``storage_kernel_block_size``); this
+    needs ``block_size`` to be a multiple of ``index_kpool * 32`` (512 for
+    ``index_kpool = 16``). A smaller block (e.g. the default 64) silently
+    collapses ``storage_block_size`` (64 // 16 = 4) and only fails later at
+    the opaque C++ assert; ``get_kv_cache_spec`` guards this up front
+    instead.
+    """
+
+    def __init__(
+        self,
+        *,
+        head_dim: int,
+        dtype: torch.dtype,
+        prefix: str,
+        cache_config,
+        index_kpool: int,
+    ):
+        super().__init__(
+            head_dim=head_dim, dtype=dtype, prefix=prefix, cache_config=cache_config
+        )
+        assert index_kpool > 1, "Glm5NextIndexerCache expects index_kpool > 1"
+        # Keep chunked-prefill boundaries aligned to complete pools.
+        assert cache_config.block_size % index_kpool == 0, (
+            "Glm5NextIndexerCache: cache_config.block_size "
+            f"({cache_config.block_size}) must be a multiple of index_kpool "
+            f"({index_kpool}) so chunked-prefill boundaries stay pool-aligned."
+        )
+        self._index_kpool = index_kpool
+
+    def get_kv_cache_spec(self, vllm_config: VllmConfig):
+        from dataclasses import replace
+
+        spec = super().get_kv_cache_spec(vllm_config)
+        # ``tokens_per_state`` is the KV-spec representation of kpool
+        # compression in the current cache-layout API.
+        assert isinstance(spec, MLAAttentionSpec)
+        spec = replace(spec, tokens_per_state=self._index_kpool)
+
+        # DeepGEMM paged-MQA takes block_kv in {32, 64}; the storage block
+        # (= block_size // index_kpool) is virtually split into pool pages of
+        # the largest such size that tiles it, so it must be a multiple of 32.
+        storage_block_size = spec.block_size // self._index_kpool
+        assert (
+            spec.block_size % self._index_kpool == 0 and storage_block_size % 32 == 0
+        ), (
+            "Glm5NextIndexerCache: kpool indexer requires cache block_size to "
+            f"be a multiple of index_kpool * 32 ({self._index_kpool * 32}) so "
+            "that DeepGEMM paged-MQA pool pages (32 or 64 entries) tile the "
+            f"storage block, got block_size={spec.block_size} -> "
+            f"storage_block_size={storage_block_size}."
+        )
+        max_page_size = max(PAGED_MQA_PAGE_SIZES)
+        min_page_size = min(PAGED_MQA_PAGE_SIZES)
+        if storage_block_size <= max_page_size:
+            page_size = storage_block_size
+        elif storage_block_size % max_page_size == 0:
+            page_size = max_page_size
+        else:
+            page_size = min_page_size
+        return replace(
+            spec,
+            storage_block_size=page_size * self._index_kpool,
+        )
+
+
+class Glm5NextTailCache(DeepseekV32IndexerCache):
+    """Paged circular buffer for the kpool indexer's in-progress (tail) pool.
+
+    Holds the trailing incomplete pool's raw K + gate score: one block of
+    ``ring`` slots per request, overwritten in place by ``pos % ring``
+    as decode/spec-decode advances. Prefill seeds it (instead of discarding the
+    tail raw K+gate); the connector transfers it across PD; decode reads it to
+    compress the boundary pool correctly. ``KpoolTailSpec`` /
+    ``KpoolTailManager`` provide the no-prune, 1-block/req allocation that lets
+    the in-progress pool survive across steps and across transfer.
+
+    Stores raw bf16 K (``head_dim``) as the "K" half of each block and the
+    bf16 gate score (``head_dim``) as the "V" half -- not the fp8-compressed
+    entry, which lives in ``Glm5NextIndexerCache``.
+    """
+
+    def __init__(
+        self,
+        *,
+        head_dim: int,
+        dtype: torch.dtype,
+        prefix: str,
+        cache_config,
+        index_kpool: int,
+    ):
+        super().__init__(
+            head_dim=head_dim, dtype=dtype, prefix=prefix, cache_config=cache_config
+        )
+        assert index_kpool > 1, "Glm5NextTailCache expects index_kpool > 1"
+        self._index_kpool = index_kpool
+
+    def get_kv_cache_spec(self, vllm_config: VllmConfig):
+        # The two head slots form [K, gate score] in the generic
+        # [block, head, state, content] cache view.
+        # Drafts are stashed before acceptance. With a one-pool ring, the
+        # drafts behind a rejected pool-completing draft overwrite the keys
+        # read by its redo.
+        span = self._index_kpool + vllm_config.num_speculative_tokens
+        ring = self._index_kpool * next_power_of_2(cdiv(span, self._index_kpool))
+        # ring must divide the attention block size (a multiple of 128).
+        assert self.cache_config.block_size % ring == 0, (
+            f"Glm5NextTailCache: cache_config.block_size "
+            f"({self.cache_config.block_size}) must be a multiple of the "
+            f"tail ring ({ring})"
+        )
+        return KpoolTailSpec(
+            block_size=ring,
+            num_kv_heads=2,
+            head_size=self.head_dim,
+            head_size_v=0,
+            dtype=torch.bfloat16,
+            sliding_window=ring,
+        )
+
+    def get_attn_backend(self):
+        from vllm.v1.attention.backends.mla.indexer import KpoolTailBackend
+
+        return KpoolTailBackend
+
+
+class Indexer(nn.Module):
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        config: Glm5NextConfig,
+        hidden_size: int,
+        q_lora_rank: int,
+        quant_config: QuantizationConfig | None,
+        cache_config: CacheConfig | None,
+        topk_indices_buffer: torch.Tensor | None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.vllm_config = vllm_config
+        self.config = config
+        self.quant_config = quant_config
+        # self.indexer_cfg = config.attn_module_list_cfg[0]["attn_index"]
+        # Indexer is only constructed for v32 configs, where these sparse-indexer
+        # fields are guaranteed populated; narrow away the `int | None` declared
+        # on Glm5NextConfig for the optional-indexer case.
+        assert config.index_topk is not None
+        assert config.index_n_heads is not None
+        assert config.index_head_dim is not None
+        assert config.index_kpool is not None
+        self.topk_tokens = config.index_topk
+        self.n_head = config.index_n_heads  # 64
+        self.head_dim = config.index_head_dim  # 128
+        self.rope_dim = config.qk_rope_head_dim  # 64
+        self.index_kpool = config.index_kpool
+        self.q_lora_rank = q_lora_rank  # 1536
+
+        # kpool
+        self.index_kpool_compress_ape = nn.Parameter(
+            torch.zeros(self.index_kpool, self.head_dim, dtype=torch.float32)
+        )
+        # Keep the checkpoint name ``index_kpool_compress_gate`` without a
+        # ``.weight`` suffix. F.linear consumes its [head_dim, hidden_size] shape.
+        self.index_kpool_compress_gate = nn.Parameter(
+            torch.empty(self.head_dim, hidden_size, dtype=torch.bfloat16)
+        )
+
+        # no tensor parallel, just replicated
+        self.wq_b = ReplicatedLinear(
+            self.q_lora_rank,
+            self.head_dim * self.n_head,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.wq_b",
+        )
+        # Fused wk + weights_proj: single GEMM producing [head_dim + n_head].
+        # FP8 wk weights are upcasted to BF16 during loading to maintain fusion.
+        self.wk_weights_proj = MergedColumnParallelLinear(
+            hidden_size,
+            [self.head_dim, self.n_head],
+            bias=False,
+            quant_config=None,
+            disable_tp=True,
+            prefix=f"{prefix}.wk_weights_proj",
+        )
+        self.k_norm = LayerNorm(self.head_dim, eps=1e-6)
+        self.softmax_scale = self.head_dim**-0.5
+
+        # Hadamard-128 rotation of the indexer query is fused with the FP8
+        # quant (see forward: fwht128_quant_fp8) -- no precomputed matrix.
+
+        self.scale_fmt = "ue8m0"
+        self.quant_block_size = 128  # TODO: get from config
+        self.topk_indices_buffer = topk_indices_buffer
+        self._wp_fp32: torch.Tensor | None = None
+
+        # NOTE: (zyongye) we use fp8 naive cache,
+        #       where we store value in fp8 and scale in fp32
+        #       per self.quant_block_size element
+        self.k_cache = Glm5NextIndexerCache(
+            head_dim=self.head_dim + self.head_dim // self.quant_block_size * 4,
+            dtype=torch.uint8,
+            prefix=f"{prefix}.k_cache",
+            cache_config=cache_config,
+            index_kpool=self.index_kpool,
+        )
+        # Paged tail cache (in-progress pool's raw K + gate score). Written by
+        # prefill (seeds the boundary pool) and decode (per-step stash); read by
+        # the decode kernel to compress the boundary pool. Transferred across PD
+        # so the decode side sees the prefill tail. See KpoolTailSpec/Manager.
+        self.tail_cache = Glm5NextTailCache(
+            head_dim=self.head_dim,
+            dtype=torch.bfloat16,
+            prefix=f"{prefix}.tail_cache",
+            cache_config=cache_config,
+            index_kpool=self.index_kpool,
+        )
+        self.max_pool_len = vllm_config.model_config.max_model_len // self.index_kpool
+        self.prefix = prefix
+        from vllm.v1.attention.backends.mla.indexer import get_max_prefill_buffer_size
+
+        self.max_total_seq_len = get_max_prefill_buffer_size(vllm_config)
+        self.indexer_op = SparseAttnIndexerKpool(
+            self.k_cache,
+            self.quant_block_size,
+            self.scale_fmt,
+            self.topk_tokens,
+            self.head_dim,
+            self.max_pool_len,
+            self.max_total_seq_len,
+            self.topk_indices_buffer,
+            tail_cache=self.tail_cache,
+        )
+
+    def forward(
+        self, hidden_states: torch.Tensor, qr: torch.Tensor, positions, rotary_emb
+    ) -> torch.Tensor:
+        q, _ = self.wq_b(qr)
+        q = q.view(-1, self.n_head, self.head_dim)
+
+        # Compute the head gate in fp32; bf16 error can change near-tie pool
+        # rankings on long-context tasks. Cache it after weights are loaded.
+        kw, _ = self.wk_weights_proj(hidden_states)
+        k = kw[:, : self.head_dim]
+        # GB10 fix: bf16 x bf16 with fp32 accumulation and output. The inputs
+        # are exact in bf16, so this matches the fp32 GEMM up to summation
+        # order, and cuBLAS's fp32 small-N kernel took ~69 us here per layer.
+        if self._wp_fp32 is None:
+            self._wp_fp32 = self.wk_weights_proj.weight.data[self.head_dim :, :].t()
+        weights = torch.mm(hidden_states, self._wp_fp32, out_dtype=torch.float32)
+
+        k = _fused_indexer_k_norm(
+            k, self.k_norm.weight, self.k_norm.bias, self.head_dim, self.k_norm.eps
+        )
+
+        if self.rope_dim > 0:
+            q_pe, q_nope = torch.split(
+                q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
+            )
+            k_pe, k_nope = torch.split(
+                k, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
+            )
+
+            q_pe, k_pe = rotary_emb(positions, q_pe, k_pe.unsqueeze(1))
+            # Note: RoPE (NeoX) can introduce extra leading dimensions during
+            # compilation so we need to reshape back to token-flattened shapes
+            q_pe = q_pe.reshape(-1, self.n_head, self.rope_dim)
+            k_pe = k_pe.reshape(-1, 1, self.rope_dim)
+
+            # `rotary_emb` is shape-preserving; `q_pe` is already
+            # [num_tokens, n_head, rope_dim].
+            q = torch.cat([q_pe, q_nope], dim=-1)
+            # `k_pe` is [num_tokens, 1, rope_dim] (MQA).
+            k = torch.cat([k_pe.squeeze(-2), k_nope], dim=-1)
+        # else: qk_rope_head_dim=0 — no rope component. q is already
+        # [num_tokens, n_head, head_dim] and k is [num_tokens, head_dim] (all
+        # nope), so skip the rope split / rotary / cat entirely; otherwise the
+        # split/reshape would build 0-element tensors (breaks dynamo tracing).
+
+        # Rotate Q into the cached K basis before computing fp8 MQA logits.
+        # Fusing the fp32 FWHT and quantization avoids an intermediate HBM
+        # round-trip and bf16 matrix-rounding bias.
+        assert self.head_dim == 128 and self.quant_block_size == 128
+        assert self.scale_fmt == "ue8m0"
+        q = q.view(-1, self.head_dim)
+        q_fp8, q_scale = fwht128_quant_fp8(q)
+        q_fp8 = q_fp8.view(-1, self.n_head, self.head_dim)
+        q_scale = q_scale.view(-1, self.n_head, 1)
+
+        weights = _fused_indexer_weight_scale(
+            weights, q_scale, self.softmax_scale * self.n_head**-0.5
+        )
+
+        # kpool: per-token gate score driving the softmax-weighted pool. Computed
+        # from the same hidden_states that produced `k`, so it stays token-aligned.
+        # F.linear(x, gate) = x @ gate.T  with gate [head_dim, hidden_size].
+        gate_score = F.linear(hidden_states, self.index_kpool_compress_gate)
+
+        # DeepGEMM's MQA-logits kernels (fp8_mqa_logits /
+        # fp8_fp4_paged_mqa_logits) require num_heads in {32, 64}; this
+        # checkpoint uses index_n_heads=16. Zero-pad q and the per-head
+        # weights: logits are a weights-weighted sum over heads, so
+        # zero-weight padded heads contribute exactly nothing.
+        if self.n_head < 32:
+            pad = 32 - self.n_head
+            q_fp8 = _pad_indexer_heads(q_fp8, pad)
+            weights = _pad_indexer_heads(weights, pad)
+
+        return self.indexer_op(
+            hidden_states,
+            q_fp8,
+            k,
+            weights,
+            gate_score=gate_score,
+            compress_ape=self.index_kpool_compress_ape,
+            index_kpool=self.index_kpool,
+            positions=positions,
+        )
+
+
+class Glm5NextMLAAttention(nn.Module):
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        config: Glm5NextConfig,
+        hidden_size: int,
+        num_heads: int,
+        qk_nope_head_dim: int,
+        qk_rope_head_dim: int,
+        v_head_dim: int,
+        q_lora_rank: int | None,
+        kv_lora_rank: int,
+        max_position_embeddings: int = 8192,
+        cache_config: CacheConfig | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+        topk_indices_buffer: torch.Tensor | None = None,
+        input_size: int | None = None,
+        skip_rope: bool | None = False,
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.qk_nope_head_dim = qk_nope_head_dim
+        self.qk_rope_head_dim = qk_rope_head_dim
+        self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
+        self.v_head_dim = v_head_dim
+
+        self.q_lora_rank = q_lora_rank
+        self.kv_lora_rank = kv_lora_rank
+
+        self.num_heads = num_heads
+        tp_size = get_tensor_model_parallel_world_size()
+        assert num_heads % tp_size == 0
+        self.num_local_heads = num_heads // tp_size
+
+        self.scaling = self.qk_head_dim**-0.5
+        self.max_position_embeddings = max_position_embeddings
+
+        # Use input_size for projection input dimensions if provided,
+        # otherwise default to hidden_size (used in Eagle3 Deepseek with MLA)
+        proj_input_size = input_size if input_size is not None else self.hidden_size
+
+        if self.q_lora_rank is not None:
+            self.fused_qkv_a_proj = DeepSeekV2FusedQkvAProjLinear(
+                proj_input_size,
+                [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
+                quant_config=quant_config,
+                prefix=f"{prefix}.fused_qkv_a_proj",
+            )
+        else:
+            self.kv_a_proj_with_mqa = ReplicatedLinear(
+                proj_input_size,
+                self.kv_lora_rank + self.qk_rope_head_dim,
+                bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.kv_a_proj_with_mqa",
+            )
+
+        if self.q_lora_rank is not None:
+            self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps)
+            self.q_b_proj = ColumnParallelLinear(
+                self.q_lora_rank,
+                self.num_heads * self.qk_head_dim,
+                bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.q_b_proj",
+            )
+        else:
+            self.q_proj = ColumnParallelLinear(
+                proj_input_size,
+                self.num_heads * self.qk_head_dim,
+                bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.q_proj",
+            )
+        self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps)
+        self.kv_b_proj = ColumnParallelLinear(
+            self.kv_lora_rank,
+            self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.kv_b_proj",
+        )
+        self.o_proj = RowParallelLinear(
+            self.num_heads * self.v_head_dim,
+            self.hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+
+        if not skip_rope:
+            assert config.rope_parameters is not None
+            if config.rope_parameters["rope_type"] != "default":
+                config.rope_parameters["rope_type"] = (
+                    "deepseek_llama_scaling"
+                    if config.rope_parameters.get("attention_factor") == 1.0
+                    else "deepseek_yarn"
+                )
+
+            self.rotary_emb: RotaryEmbedding | None = get_rope(
+                qk_rope_head_dim,
+                max_position=max_position_embeddings,
+                rope_parameters=config.rope_parameters,
+                is_neox_style=False,
+            )
+
+            if (
+                config.rope_parameters["rope_type"] != "default"
+                and config.rope_parameters["rope_type"] == "deepseek_yarn"
+            ):
+                mscale_all_dim = config.rope_parameters.get("mscale_all_dim", False)
+                scaling_factor = config.rope_parameters["factor"]
+                mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
+                self.scaling = self.scaling * mscale * mscale
+        else:
+            self.rotary_emb = None
+
+        self.is_v32 = config.index_topk is not None
+
+        if self.is_v32:
+            self.indexer_rope_emb: RotaryEmbedding | None = get_rope(
+                qk_rope_head_dim,
+                max_position=max_position_embeddings,
+                rope_parameters=config.rope_parameters,
+                is_neox_style=not config.indexer_rope_interleave,
+            )
+            # The sparse indexer projects from the MLA q-lora rank, which is
+            # always set for v32 MLA configs; narrow away the `int | None`.
+            assert q_lora_rank is not None
+            self.indexer: Indexer | None = Indexer(
+                vllm_config,
+                config,
+                hidden_size,
+                q_lora_rank,
+                quant_config,
+                cache_config,
+                topk_indices_buffer,
+                f"{prefix}.indexer",
+            )
+
+        else:
+            self.indexer_rope_emb = None
+            self.indexer = None
+
+        mla_modules = MLAModules(
+            kv_a_layernorm=self.kv_a_layernorm,
+            kv_b_proj=self.kv_b_proj,
+            rotary_emb=self.rotary_emb,
+            o_proj=self.o_proj,
+            fused_qkv_a_proj=self.fused_qkv_a_proj
+            if self.q_lora_rank is not None
+            else None,
+            kv_a_proj_with_mqa=self.kv_a_proj_with_mqa
+            if self.q_lora_rank is None
+            else None,
+            q_a_layernorm=self.q_a_layernorm if self.q_lora_rank is not None else None,
+            q_b_proj=self.q_b_proj if self.q_lora_rank is not None else None,
+            q_proj=self.q_proj if self.q_lora_rank is None else None,
+            indexer=self.indexer,
+            indexer_rotary_emb=self.indexer_rope_emb,
+            is_sparse=self.is_v32,
+            topk_indices_buffer=topk_indices_buffer,
+        )
+
+        self.mla_attn = MultiHeadLatentAttentionWrapper(
+            self.hidden_size,
+            self.num_local_heads,
+            self.scaling,
+            self.qk_nope_head_dim,
+            self.qk_rope_head_dim,
+            self.v_head_dim,
+            self.q_lora_rank,
+            self.kv_lora_rank,
+            mla_modules,
+            cache_config,
+            quant_config,
+            prefix,
+            skip_topk=False,
+            fuse_qkv_rmsnorm=True,
+        )
+
+    def forward(
+        self, hidden_states: torch.Tensor, positions: torch.Tensor
+    ) -> torch.Tensor:
+        # The wrapper also runs the sparse indexer before MLA attention.
+        return self.mla_attn(positions, hidden_states)

docker-images/glm53-kindling/experimental/fixes/causal_conv1d.py

@@ -0,0 +1,1306 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Copyright (c) 2024, Tri Dao.
+# Adapted from https://github.com/Dao-AILab/causal-conv1d/blob/main/causal_conv1d/causal_conv1d_interface.py
+
+
+import numpy as np
+import torch
+
+from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
+from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID
+
+
+@triton.jit(do_not_specialize_on_alignment=["num_cache_lines"])
+def _causal_conv1d_fwd_kernel(  # continuous batching
+    # Pointers to matrices
+    x_ptr,  # (dim, cu_seqlen) holding `batch` of actual sequences + padded sequences
+    w_ptr,  # (dim, width)
+    bias_ptr,
+    initial_states_ptr,  # conv_states_ptr
+    cache_indices_ptr,  # (batch, n_blocks + padding) The second dimension contains
+    # the block indices relevant for each sequence
+    # plus potential 0-padding at the beginning and at the end
+    has_initial_states_ptr,
+    query_start_loc_ptr,
+    batch_ptr,
+    token_chunk_offset_ptr,
+    block_idx_first_scheduled_token,  # (batch,)
+    block_idx_last_scheduled_token,  # (batch,)
+    initial_state_idx,  # (batch,)
+    num_computed_tokens,  # (batch,)
+    o_ptr,  # (dim, seqlen) - actually pointing to x_ptr
+    # Matrix dimensions
+    dim: tl.constexpr,
+    num_cache_lines,  # added to support vLLM larger cache lines
+    # Strides
+    stride_x_dim: tl.constexpr,  # stride to get to next feature-value,
+    stride_x_token: tl.int64,  # stride to get to next token (same feature-index, same sequence-index)
+    stride_w_dim: tl.constexpr,  # stride to get to next dim-axis value
+    stride_w_width: tl.constexpr,  # stride to get to next width-axis value
+    stride_istate_seq: tl.constexpr,
+    stride_istate_dim: tl.constexpr,
+    stride_istate_token: tl.constexpr,
+    stride_cache_indices: tl.constexpr,
+    stride_o_dim: tl.constexpr,
+    stride_o_token: tl.int64,
+    stride_o_group: tl.int64,  # with OUT_GROUP: stride between groups of OUT_GROUP features
+    stride_block_m: tl.constexpr,  # Stride block to align divided by BLOCK_M
+    # others
+    pad_slot_id: tl.constexpr,
+    null_block_id: tl.constexpr,
+    # Meta-parameters
+    HAS_BIAS: tl.constexpr,
+    KERNEL_WIDTH: tl.constexpr,
+    SILU_ACTIVATION: tl.constexpr,
+    IS_APC_ENABLED: tl.constexpr,
+    HAS_NULL_BLOCK: tl.constexpr,
+    NP2_STATELEN: tl.constexpr,
+    OUT_GROUP: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    launch_pdl: tl.constexpr,
+):
+    conv_states_ptr = initial_states_ptr
+    conv_state_indices_ptr = cache_indices_ptr
+    stride_conv_state_seq = stride_istate_seq
+    stride_conv_state_dim = stride_istate_dim
+    stride_conv_state_tok = stride_istate_token
+    state_len = (
+        KERNEL_WIDTH - 1
+    )  # can be passed via argument if it's not the same as this value
+
+    if launch_pdl:
+        tl.extra.cuda.gdc_wait()
+
+    # one program handles one chunk in a single sequence
+    # rather than mixing sequences - to make updating initial_states across sequences efficiently
+
+    # single-sequence id
+    idx_seq = tl.load(batch_ptr + tl.program_id(0)).to(tl.int64)
+    chunk_offset = tl.load(token_chunk_offset_ptr + tl.program_id(0))
+
+    # BLOCK_N elements along the feature-dimension (channel)
+    idx_feats = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N)
+
+    if idx_seq == pad_slot_id:
+        if launch_pdl:
+            tl.extra.cuda.gdc_launch_dependents()
+        return
+
+    sequence_start_index = tl.load(query_start_loc_ptr + idx_seq)
+    sequence_end_index = tl.load(query_start_loc_ptr + idx_seq + 1)
+    # find the actual sequence length
+    seqlen = sequence_end_index - sequence_start_index
+
+    B_size: tl.constexpr = stride_block_m * BLOCK_M
+
+    if IS_APC_ENABLED:
+        # Handle the case if prefix caching is enabled.
+        # In particular, if prefix caching is enabled, the program write additional cache states to "cache_indices_ptr"
+
+        # Get the length of the completed sequence so far and compute the offset.
+        current_first_index = tl.load(block_idx_first_scheduled_token + idx_seq)
+        current_last_index = tl.load(block_idx_last_scheduled_token + idx_seq)
+        sequence_completed_index = tl.load(num_computed_tokens + idx_seq)
+
+        # Compute the offset where the first stride_block_m-aligned first full block is
+        # Value in "token-space"
+        sequence_completed_offset_token = sequence_completed_index % B_size
+        seq_completed_offset = B_size - sequence_completed_offset_token
+        seq_end_offset = (seqlen - seq_completed_offset) % B_size
+        last_full_block_token_index = sequence_end_index - seq_end_offset
+        # If the sequence without the sequence_offset_index is stride_cache_chunk-aligned, then the last full chunk is the second-to-last one
+        if seq_end_offset == 0:
+            last_full_block_token_index = last_full_block_token_index - B_size
+
+        # Get the number of blocks to be filled for the current sequence
+        # If n_block_to_fill = 0, then only the state at the sequence end is stored
+        n_block_to_fill = current_last_index - current_first_index
+
+        # Get the index of the init block
+        conv_state_init_index = tl.load(initial_state_idx + idx_seq)
+    else:
+        n_block_to_fill = 0
+        current_last_index = 0
+        conv_state_init_index = 0
+        current_first_index = 0
+        last_full_block_token_index = 0
+
+    token_offset = BLOCK_M * chunk_offset
+    segment_len = min(BLOCK_M, seqlen - token_offset)
+
+    # base of the sequence
+    x_base = (
+        x_ptr + sequence_start_index * stride_x_token + idx_feats * stride_x_dim
+    )  # [BLOCK_N,]
+
+    # cache_idx
+    conv_states_input_coord = tl.load(
+        conv_state_indices_ptr + idx_seq * stride_cache_indices + conv_state_init_index
+    ).to(tl.int64)
+
+    if HAS_NULL_BLOCK:  # noqa
+        if conv_states_input_coord == null_block_id:
+            # not processing as this is a null block (padding)
+            if launch_pdl:
+                tl.extra.cuda.gdc_launch_dependents()
+            return
+    conv_states_base = (
+        conv_states_ptr
+        + (conv_states_input_coord * stride_conv_state_seq)
+        + (idx_feats * stride_conv_state_dim)
+    )  # [BLOCK_N,]
+
+    w_base = w_ptr + (idx_feats * stride_w_dim)  # [BLOCK_N,]
+
+    # Does 2 things:
+    # 1. READ prior-block init-state data - [done by every Triton programs]
+    # 2. update conv_state with new data [only by the Triton program handles chunk_offset=0]
+    if chunk_offset == 0:
+        # read from conv_states
+        load_init_state = tl.load(has_initial_states_ptr + idx_seq).to(tl.int1)
+        if load_init_state:
+            # load from conv_states
+            prior_tokens = conv_states_base + (state_len - 1) * stride_conv_state_tok
+            mask_w = idx_feats < dim
+            if KERNEL_WIDTH == 2:
+                conv_states_ptrs = prior_tokens  # [BLOCK_N]
+                col0 = tl.load(conv_states_ptrs, mask_w, 0.0)
+            if KERNEL_WIDTH == 3:
+                conv_states_ptrs = prior_tokens  # [BLOCK_N]
+                col1 = tl.load(conv_states_ptrs, mask_w, 0.0)
+                conv_states_ptrs = prior_tokens - 1 * stride_conv_state_tok  # [BLOCK_N]
+                col0 = tl.load(conv_states_ptrs, mask_w, 0.0)
+            if KERNEL_WIDTH == 4:
+                conv_states_ptrs = prior_tokens  # [BLOCK_N]
+                col2 = tl.load(conv_states_ptrs, mask_w, 0.0)
+                conv_states_ptrs = prior_tokens - 1 * stride_conv_state_tok  # [BLOCK_N]
+                col1 = tl.load(conv_states_ptrs, mask_w, 0.0)
+                conv_states_ptrs = prior_tokens - 2 * stride_conv_state_tok  # [BLOCK_N]
+                col0 = tl.load(conv_states_ptrs, mask_w, 0.0)
+            if KERNEL_WIDTH == 5:
+                conv_states_ptrs = prior_tokens  # [BLOCK_N]
+                col3 = tl.load(conv_states_ptrs, mask_w, 0.0)
+                conv_states_ptrs = prior_tokens - 1 * stride_conv_state_tok  # [BLOCK_N]
+                col2 = tl.load(conv_states_ptrs, mask_w, 0.0)
+                conv_states_ptrs = prior_tokens - 2 * stride_conv_state_tok  # [BLOCK_N]
+                col1 = tl.load(conv_states_ptrs, mask_w, 0.0)
+                conv_states_ptrs = prior_tokens - 3 * stride_conv_state_tok  # [BLOCK_N]
+                col0 = tl.load(conv_states_ptrs, mask_w, 0.0)
+        else:
+            # prior-tokens are zeros
+            if KERNEL_WIDTH >= 2:  # STRATEGY1
+                # first chunk and does not have prior-token, so just set to 0
+                col0 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty)
+            if KERNEL_WIDTH >= 3:  # STRATEGY1
+                col1 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty)
+            if KERNEL_WIDTH >= 4:  # STRATEGY1
+                col2 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty)
+            if KERNEL_WIDTH >= 5:  # STRATEGY1
+                col3 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty)
+
+        # STEP 2:
+        # here prepare data for updating conv_state
+        if (
+            state_len <= seqlen
+        ):  # SMALL_CACHE=True (only move part of 'x' into conv_state cache)
+            # just read from 'x'
+            # copy 'x' data to conv_state
+            # load only 'x' data (and set 0 before 'x' if seqlen < state_len)
+            idx_tokens_last = (seqlen - state_len) + tl.arange(
+                0, NP2_STATELEN
+            )  # [BLOCK_M]
+            x_ptrs = (
+                x_ptr
+                + ((sequence_start_index + idx_tokens_last) * stride_x_token)[:, None]
+                + (idx_feats * stride_x_dim)[None, :]
+            )  # [BLOCK_M,BLOCK_N,]
+            mask_x = (
+                (idx_tokens_last >= 0)[:, None]
+                & (idx_tokens_last < seqlen)[:, None]
+                & (idx_feats < dim)[None, :]
+            )  # token-index  # token-index  # feature-index
+            loaded_x = tl.load(x_ptrs, mask_x, 0.0)
+            idx_tokens_conv = tl.arange(0, NP2_STATELEN)  # [BLOCK_M]
+
+            # Compute the offset where the last block should be written in the conv_states
+            conv_states_output_coord = tl.load(
+                conv_state_indices_ptr
+                + idx_seq * stride_cache_indices
+                + current_last_index
+            ).to(tl.int64)
+
+            conv_states_ptrs_target = (
+                conv_states_ptr
+                + (conv_states_output_coord * stride_conv_state_seq)  # Offset from seq
+                + (idx_feats * stride_conv_state_dim)
+            )[None, :] + (  # [BLOCK_N,]
+                idx_tokens_conv * stride_conv_state_tok
+            )[:, None]
+
+            mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :]
+            tl.debug_barrier()  #  NOTE: use this due to bug in Triton compiler
+            tl.store(conv_states_ptrs_target, loaded_x, mask)
+
+        else:
+            if load_init_state:
+                # update conv_state by shifting left, i.e. take last few cols from conv_state + cols from 'x'
+                idx_tokens_conv = tl.arange(0, NP2_STATELEN)  # [BLOCK_M]
+
+                conv_states_ptrs_source = (
+                    conv_states_ptr
+                    + (conv_states_input_coord * stride_conv_state_seq)
+                    + (idx_feats * stride_conv_state_dim)[None, :]
+                    + ((idx_tokens_conv + seqlen) * stride_conv_state_tok)[:, None]
+                )  # [BLOCK_M, BLOCK_N]
+                mask = (
+                    (conv_states_input_coord < num_cache_lines)
+                    & ((idx_tokens_conv + seqlen) < state_len)[:, None]
+                    & (idx_feats < dim)[None, :]
+                )
+                conv_state = tl.load(conv_states_ptrs_source, mask, other=0.0)
+
+                VAL = state_len - seqlen
+
+                x_ptrs = (
+                    x_base[None, :]
+                    + ((idx_tokens_conv - VAL) * stride_x_token)[:, None]
+                )  # [BLOCK_M, BLOCK_N]
+
+                mask_x = (
+                    (idx_tokens_conv - VAL >= 0)[:, None]
+                    & (idx_tokens_conv - VAL < seqlen)[:, None]
+                    & (idx_feats < dim)[None, :]
+                )  # token-index  # token-index  # feature-index
+                loaded_x = tl.load(x_ptrs, mask_x, 0.0)
+
+                tl.debug_barrier()  # need this due to the bug in tl.where not enforcing this when data is the result of another tl.load
+                new_conv_state = tl.where(
+                    mask, conv_state, loaded_x
+                )  # BUG in 'tl.where'  which requires a barrier before this
+                conv_states_ptrs_target = (
+                    conv_states_base
+                    + (idx_tokens_conv * stride_conv_state_tok)[:, None]
+                )  # [BLOCK_M, BLOCK_N]
+                mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[
+                    None, :
+                ]
+                tl.store(conv_states_ptrs_target, new_conv_state, mask)
+            else:  # load_init_state == False
+                # update conv_state by shifting left, BUT
+                # set cols prior to 'x' as zeros + cols from 'x'
+                idx_tokens_conv = tl.arange(0, NP2_STATELEN)  # [BLOCK_M]
+
+                VAL = state_len - seqlen
+
+                x_ptrs = (
+                    x_base[None, :]
+                    + ((idx_tokens_conv - VAL) * stride_x_token)[:, None]
+                )  # [BLOCK_M, BLOCK_N]
+
+                mask_x = (
+                    (idx_tokens_conv - VAL >= 0)[:, None]
+                    & (idx_tokens_conv - VAL < seqlen)[:, None]
+                    & (idx_feats < dim)[None, :]
+                )  # token-index  # token-index  # feature-index
+                new_conv_state = tl.load(x_ptrs, mask_x, 0.0)
+
+                conv_states_ptrs_target = (
+                    conv_states_base
+                    + (idx_tokens_conv * stride_conv_state_tok)[:, None]
+                )  # [BLOCK_M, BLOCK_N]
+                mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[
+                    None, :
+                ]
+                tl.store(conv_states_ptrs_target, new_conv_state, mask)
+
+    else:  # chunk_offset > 0
+        # read prior-token data from `x`
+        load_init_state = True
+        prior_tokens = x_base + (token_offset - 1) * stride_x_token
+        mask_w = idx_feats < dim
+        if KERNEL_WIDTH == 2:
+            conv_states_ptrs = prior_tokens  # [BLOCK_N]
+            col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca")
+        if KERNEL_WIDTH == 3:
+            conv_states_ptrs = prior_tokens  # [BLOCK_N]
+            col1 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca")
+            conv_states_ptrs = prior_tokens - 1 * stride_x_token  # [BLOCK_N]
+            col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca")
+        if KERNEL_WIDTH == 4:
+            conv_states_ptrs = prior_tokens  # [BLOCK_N]
+            col2 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca")
+            conv_states_ptrs = prior_tokens - 1 * stride_x_token  # [BLOCK_N]
+            col1 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca")
+            conv_states_ptrs = prior_tokens - 2 * stride_x_token  # [BLOCK_N]
+            col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca")
+        if KERNEL_WIDTH == 5:
+            # ruff: noqa: F841
+            conv_states_ptrs = prior_tokens  # [BLOCK_N]
+            col3 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca")
+            conv_states_ptrs = prior_tokens - 1 * stride_x_token  # [BLOCK_N]
+            col2 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca")
+            conv_states_ptrs = prior_tokens - 2 * stride_x_token  # [BLOCK_N]
+            col1 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca")
+            conv_states_ptrs = prior_tokens - 3 * stride_x_token  # [BLOCK_N]
+            col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca")
+
+        # Store intermediate states aligned with stride_block_m
+        # The additional states are cached starting from the last stride_block_m.
+        # For example:
+        # If n_block_to_fill = 0, then only the state at the sequence end is cached and the process below is not involved.
+        # If n_block_to_fill > 0, then the states at the sequence end and at the n_block_to_fill-last
+        # stride_block_m are cached.
+        # For example chunk_offset = n_block_to_fill stores the state at last_full_block
+        if (chunk_offset - 1) < n_block_to_fill:
+            # Store the states at the chunk boundaries from the start of the sequence
+            idx_tokens_last = (
+                last_full_block_token_index
+                - (n_block_to_fill - chunk_offset) * B_size
+                - state_len
+            ) + tl.arange(0, NP2_STATELEN)  # [BLOCK_M]
+            x_ptrs = (
+                x_ptr
+                + (idx_tokens_last * stride_x_token)[:, None]
+                + (idx_feats * stride_x_dim)[None, :]
+            )  # [BLOCK_M,BLOCK_N,]
+
+            mask_x = (idx_tokens_last >= 0)[:, None] & (idx_feats < dim)[
+                None, :
+            ]  # token-index  # token-index  # feature-index
+            loaded_x = tl.load(x_ptrs, mask_x, 0.0)
+            idx_tokens_conv = tl.arange(0, NP2_STATELEN)  # [BLOCK_M]
+
+            # cache_idx
+            conv_states_output_coord = tl.load(
+                conv_state_indices_ptr
+                + idx_seq * stride_cache_indices
+                + current_first_index
+                + (chunk_offset - 1)
+            ).to(tl.int64)
+
+            conv_states_ptrs_target = (
+                conv_states_ptr
+                + (conv_states_output_coord * stride_conv_state_seq)  # Offset from seq
+                + (idx_feats * stride_conv_state_dim)
+            )[None, :] + (  # [BLOCK_N,]
+                idx_tokens_conv * stride_conv_state_tok
+            )[:, None]
+
+            mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :]
+            tl.debug_barrier()  #  NOTE: use this due to bug in Triton compiler
+            tl.store(conv_states_ptrs_target, loaded_x, mask)
+
+    if HAS_BIAS:
+        bias = bias_ptr + idx_feats
+        mask_bias = idx_feats < dim
+        acc_preload = tl.load(bias, mask=mask_bias, other=0.0).to(
+            tl.float32
+        )  # [BLOCK_N]
+    else:
+        acc_preload = tl.zeros((BLOCK_N,), dtype=tl.float32)
+
+    x_base_1d = x_base + token_offset * stride_x_token  # starting of chunk
+
+    # PRE-LOAD WEIGHTS
+    mask_w = idx_feats < dim
+    if KERNEL_WIDTH >= 2:
+        w_ptrs = w_base + (0 * stride_w_width)  # [BLOCK_N] tensor
+        w_col0 = tl.load(w_ptrs, mask_w, other=0.0)
+        w_ptrs = w_base + (1 * stride_w_width)  # [BLOCK_N] tensor
+        w_col1 = tl.load(w_ptrs, mask_w, other=0.0)
+    if KERNEL_WIDTH >= 3:
+        w_ptrs = w_base + (2 * stride_w_width)  # [BLOCK_N] tensor
+        w_col2 = tl.load(w_ptrs, mask_w, other=0.0)
+    if KERNEL_WIDTH >= 4:
+        w_ptrs = w_base + (3 * stride_w_width)  # [BLOCK_N] tensor
+        w_col3 = tl.load(w_ptrs, mask_w, other=0.0)
+    mask_x_1d = idx_feats < dim
+
+    if launch_pdl:
+        tl.extra.cuda.gdc_launch_dependents()
+
+    for idx_token in range(segment_len):
+        acc = acc_preload
+
+        matrix_w = w_col0
+        matrix_x = col0
+        for j in tl.static_range(KERNEL_WIDTH):
+            if KERNEL_WIDTH == 2:
+                if j == 1:  # KERNEL_WIDTH-1:
+                    matrix_w = w_col1
+                    x_ptrs_1d = x_base_1d + idx_token * stride_x_token  # [BLOCK_N]
+                    matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d)
+            elif KERNEL_WIDTH == 3:
+                if j == 1:
+                    matrix_w = w_col1
+                    matrix_x = col1
+                elif j == 2:
+                    matrix_w = w_col2
+                    x_ptrs_1d = x_base_1d + idx_token * stride_x_token  # [BLOCK_N]
+                    matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d)
+            elif KERNEL_WIDTH == 4:
+                if j == 1:
+                    matrix_w = w_col1
+                    matrix_x = col1
+                elif j == 2:
+                    matrix_w = w_col2
+                    matrix_x = col2
+                elif j == 3:
+                    matrix_w = w_col3
+                    x_ptrs_1d = x_base_1d + idx_token * stride_x_token  # [BLOCK_N]
+                    matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d)
+
+            acc += matrix_x * matrix_w  # [BLOCK_N]
+
+        if KERNEL_WIDTH == 2:
+            col0 = matrix_x
+        elif KERNEL_WIDTH == 3:
+            col0 = col1
+            col1 = matrix_x
+        elif KERNEL_WIDTH == 4:
+            col0 = col1
+            col1 = col2
+            col2 = matrix_x
+
+        if SILU_ACTIVATION:
+            acc = acc / (1 + tl.exp(-acc))
+        mask_1d = (idx_token < segment_len) & (
+            idx_feats < dim
+        )  # token-index  # feature-index
+        if OUT_GROUP > 0:
+            o_feats = (idx_feats // OUT_GROUP) * stride_o_group + (idx_feats % OUT_GROUP) * stride_o_dim
+        else:
+            o_feats = idx_feats * stride_o_dim
+        o_ptrs = (
+            o_ptr
+            + (sequence_start_index + token_offset + idx_token) * stride_o_token
+            + o_feats
+        )
+
+        tl.store(o_ptrs, acc, mask=mask_1d)
+
+
+def causal_conv1d_fn(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None,
+    conv_states: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    cache_indices: torch.Tensor | None = None,
+    has_initial_state: torch.Tensor | None = None,
+    activation: str | None = "silu",
+    pad_slot_id: int = PAD_SLOT_ID,
+    null_block_id: int = NULL_BLOCK_ID,
+    block_idx_first_scheduled_token: torch.Tensor | None = None,
+    block_idx_last_scheduled_token: torch.Tensor | None = None,
+    initial_state_idx: torch.Tensor | None = None,
+    num_computed_tokens: torch.Tensor | None = None,
+    block_size_to_align=0,
+    metadata=None,
+    validate_data=False,
+    out_group: int = 0,
+):
+    """Support varlen + continuous batching when x is 2D tensor.
+
+    x: (dim,cu_seq_len)
+        cu_seq_len = total tokens of all seqs in that batch
+        sequences are concatenated from left to right for varlen
+    weight: (dim, width)
+    conv_states: (...,dim,width - 1) itype
+        updated inplace if cache_indices are not provided
+        [it use `cache_indices` to get the index to the cache of conv_state for that sequence
+
+        conv_state[cache_indices[i]] for seq-i - to be used as initial_state when has_initial_state[i] = True
+             and after that conv_state[cache_indices[i]] need to be shift-left and updated with values from 'x'
+        ]
+    query_start_loc: (batch + 1) int32
+        The cumulative sequence lengths of the sequences in
+        the batch, used to index into sequence. prepended by 0.
+        if
+        x = [5, 1, 1, 1] <- continuous batching (batch=4)
+        then
+        query_start_loc = [0, 5, 6, 7, 8] <- the starting index of the next sequence; while the last value is
+           the ending index of the last sequence
+        [length(query_start_loc)-1 == batch]
+        for example: query_start_loc = torch.Tensor([0,10,16,17]),
+        x.shape=(dim,17)
+    cache_indices: (batch)  int32
+        indicates the corresponding state index,
+        like so: conv_state = conv_states[cache_indices[batch_id]]
+    has_initial_state: (batch) bool
+        indicates whether should the kernel take the current state as initial
+        state for the calculations
+        [single boolean for each sequence in the batch: True or False]
+    bias: (dim,)
+    activation: either None or "silu" or "swish" or True
+    pad_slot_id: int
+        if cache_indices is passed, lets the kernel identify padded
+        entries that will not be processed,
+        for example: cache_indices = [pad_slot_id, 1, 20, pad_slot_id]
+        in this case, the kernel will not process entries at
+        indices 0 and 3
+    block_idx_first_scheduled_token: (batch,), dtype int32
+        The pointer into cache_indices, where the first cache block to be filled is located.
+    block_idx_last_scheduled_token: (batch,), dtype int32
+        The pointer into cache_indices, where the last cache block to be filled is located.
+    initial_state_idx: (batch,), dtype int32
+        The pointer into cache_indices, where the cache block containing the initial state is located.
+    num_computed_tokens: (batch,), dtype int32
+        The number of tokens already completed for each sequence
+    block_size_to_align: int
+        The block size to align the cached states to
+    out: same shape as `x`
+    """
+    if isinstance(activation, bool) and activation:
+        activation = "silu"
+
+    args = None
+    # Store original dtype to cast back at the end
+    original_x_dtype = x.dtype
+    x = x.to(conv_states.dtype)
+    if out_group:
+        # Output [dim // out_group, seqlen, out_group]: each group of features
+        # (q, k, v of a merged conv) comes out as its own dense tensor.
+        assert x.dim() == 2 and x.shape[0] % out_group == 0
+        out = torch.empty(x.shape[0] // out_group, x.shape[1], out_group, dtype=x.dtype, device=x.device)
+    else:
+        out = torch.empty_like(x)
+    if metadata is not None:
+        nums_dict = metadata.nums_dict
+        args = nums_dict
+        batch_ptr = metadata.batch_ptr
+        token_chunk_offset_ptr = metadata.token_chunk_offset_ptr
+    else:
+        seqlens = query_start_loc.diff().to("cpu")
+        args = seqlens
+        MAX_NUM_PROGRAMS = 1024
+
+        batch_ptr = torch.full(
+            (MAX_NUM_PROGRAMS,), PAD_SLOT_ID, dtype=torch.int32, device=x.device
+        )  # tracking which seq-idx the Triton program is handling
+        token_chunk_offset_ptr = torch.full(
+            (MAX_NUM_PROGRAMS,), PAD_SLOT_ID, dtype=torch.int32, device=x.device
+        )  # tracking BLOCK_M-based index in the sequence the Triton program is handling
+
+    is_channel_last = (x.stride(0) == 1) & (x.stride(1) > 1)
+    dim, cu_seqlen = x.shape
+    _, width = weight.shape
+    state_len = width - 1
+    np2_statelen = triton.next_power_of_2(state_len)
+
+    padded_batch = query_start_loc.size(0) - 1
+    stride_x_dim = x.stride(0)
+    stride_x_token = x.stride(1)
+    stride_w_dim = weight.stride(0)
+    stride_w_width = weight.stride(1)
+    stride_istate_seq = 0
+    stride_istate_dim = 0
+    stride_istate_token = 0
+    num_cache_lines = 0
+    BLOCK_M = 8
+    if conv_states is not None:
+        # extensions to support vLLM:
+        # 1. conv_states is used to replaced initial_states
+        # 2. conv_states serve as a cache with num cache lines can be larger than batch size
+        # 3. mapping from sequence x[idx] to a cache line at index as specified via cache_indices[idx]
+        # 4. computation can be skipped if cache_indices[idx] == pad_slot_id
+        num_cache_lines = conv_states.size(0)
+        assert (
+            num_cache_lines == conv_states.shape[0]
+            and dim == conv_states.shape[1]
+            and width - 1 <= conv_states.shape[2]
+        )
+        stride_istate_seq = conv_states.stride(0)
+        stride_istate_dim = conv_states.stride(1)
+        stride_istate_token = conv_states.stride(2)
+    stride_o_group = 0
+    if out_group:
+        stride_o_group, stride_o_token, stride_o_dim = out.stride()
+    elif out.dim() == 2:
+        stride_o_dim = out.stride(0)
+        stride_o_token = out.stride(1)
+    else:
+        stride_o_dim = out.stride(1)
+        stride_o_token = out.stride(2)
+    stride_cache_indices = cache_indices.stride(0) if cache_indices is not None else 0
+
+    if validate_data:
+        assert x.dim() == 2
+        assert query_start_loc is not None
+        assert query_start_loc.dim() == 1
+        assert x.stride(0) == 1 or x.stride(1) == 1
+        if bias is not None:
+            assert bias.dim() == 1
+            assert dim == bias.size(0)
+        if cache_indices is not None:
+            assert cache_indices.dim() == 1
+            assert padded_batch == cache_indices.size(0)
+        if has_initial_state is not None:
+            assert has_initial_state.size() == (padded_batch,)
+            assert conv_states is not None, (
+                "ERROR: `has_initial_state` is used, which needs also `conv_states`"
+            )
+        assert weight.stride(1) == 1
+        assert (dim, width) == weight.shape
+        assert is_channel_last, "Need to run in channel-last layout"
+        if block_size_to_align is not None and block_size_to_align > 0:
+            assert (block_size_to_align % BLOCK_M) == 0, (
+                "The mamba block size needs to be divisible by the BLOCK_M"
+            )
+        else:
+            block_size_to_align = BLOCK_M
+
+    if metadata is None:
+
+        def num_program(META, seqlens):
+            tot = 0
+
+            mlist = []
+            offsetlist = []  # type: ignore
+
+            nums = -(-seqlens // META["BLOCK_M"])
+
+            tot = nums.sum().item()
+            mlist = np.repeat(np.arange(len(nums)), nums)
+            for idx, num in enumerate(nums):
+                offsetlist.extend(
+                    range(num)
+                )  # chunk-idx if a sequence is split into multiple chunks
+
+            if META["batch_ptr"].nelement() < len(mlist):
+                newlen = len(mlist) + 1
+                META["batch_ptr"].resize_(newlen).fill_(PAD_SLOT_ID)
+                META["token_chunk_offset_ptr"].resize_(newlen).fill_(PAD_SLOT_ID)
+
+            if META["batch_ptr"].nelement() >= len(mlist):
+                META["batch_ptr"][0 : len(mlist)].copy_(
+                    torch.from_numpy(np.array(mlist))
+                )
+                META["token_chunk_offset_ptr"][0 : len(mlist)].copy_(
+                    torch.from_numpy(np.array(offsetlist))
+                )
+
+            META["batch_ptr"] = META["batch_ptr"].to(META["x_ptr"].device)
+            META["token_chunk_offset_ptr"] = META["token_chunk_offset_ptr"].to(
+                META["x_ptr"].device
+            )
+            return tot
+    else:
+
+        def num_program(META, nums_dict):
+            tot = nums_dict[META["BLOCK_M"]]["tot"]
+
+            mlist = nums_dict[META["BLOCK_M"]]["mlist"]
+            mlist_len = nums_dict[META["BLOCK_M"]]["mlist_len"]
+
+            offsetlist = nums_dict[META["BLOCK_M"]]["offsetlist"]
+
+            if nums_dict[META["BLOCK_M"]]["batch_ptr"] is not None:
+                META["batch_ptr"] = nums_dict[META["BLOCK_M"]]["batch_ptr"]
+                META["token_chunk_offset_ptr"] = nums_dict[META["BLOCK_M"]][
+                    "token_chunk_offset_ptr"
+                ]
+            else:
+                if META["batch_ptr"].nelement() < mlist_len:
+                    newlen = mlist_len + 1
+                    META["batch_ptr"].resize_(newlen).fill_(PAD_SLOT_ID)
+                    META["token_chunk_offset_ptr"].resize_(newlen).fill_(PAD_SLOT_ID)
+
+                if META["batch_ptr"].nelement() >= mlist_len:
+                    META["batch_ptr"][0:mlist_len].copy_(mlist)
+                    META["token_chunk_offset_ptr"][0:mlist_len].copy_(offsetlist)
+            return tot
+
+    def grid(META):
+        return (
+            num_program(META, args),
+            triton.cdiv(dim, META["BLOCK_N"]),
+        )
+
+    if batch_ptr.device != x.device:
+        batch_ptr = batch_ptr.to(x.device)
+        token_chunk_offset_ptr = token_chunk_offset_ptr.to(x.device)
+
+    _causal_conv1d_fwd_kernel[grid](
+        # Pointers to matrices
+        x,
+        weight,
+        bias,
+        conv_states,
+        cache_indices,
+        has_initial_state,
+        query_start_loc,
+        batch_ptr,
+        token_chunk_offset_ptr,
+        block_idx_first_scheduled_token,
+        block_idx_last_scheduled_token,
+        initial_state_idx,
+        num_computed_tokens,
+        out,
+        # Matrix dimensions
+        dim,
+        num_cache_lines,
+        # stride
+        stride_x_dim,
+        stride_x_token,
+        stride_w_dim,
+        stride_w_width,
+        stride_istate_seq,
+        stride_istate_dim,
+        stride_istate_token,
+        stride_cache_indices,
+        stride_o_dim,
+        stride_o_token,
+        stride_o_group,
+        block_size_to_align // BLOCK_M,
+        # others
+        pad_slot_id,
+        null_block_id,
+        # META
+        HAS_BIAS=bias is not None,
+        KERNEL_WIDTH=width,
+        SILU_ACTIVATION=activation in ["silu", "swish"],
+        IS_APC_ENABLED=block_idx_last_scheduled_token is not None,
+        HAS_NULL_BLOCK=null_block_id is not None,
+        NP2_STATELEN=np2_statelen,
+        OUT_GROUP=out_group,
+        # launch_cooperative_grid=True
+        BLOCK_M=BLOCK_M,
+        BLOCK_N=256,
+        num_stages=2,
+        launch_pdl=current_platform.is_arch_support_pdl(),
+    )
+    return out.to(original_x_dtype)
+
+
+@triton.jit(do_not_specialize_on_alignment=["num_cache_lines"])
+def _causal_conv1d_update_kernel(
+    # Pointers to matrices
+    x_ptr,  # (batch, dim, seqlen)
+    w_ptr,  # (dim, width)
+    bias_ptr,
+    conv_state_ptr,
+    conv_state_indices_ptr,
+    num_accepted_tokens_ptr,
+    query_start_loc_ptr,  # (batch + 1)
+    block_idx_last_scheduled_token,  # (batch,)
+    initial_state_idx,  # (batch,)
+    o_ptr,  # (batch, dim, seqlen)
+    # Matrix dimensions
+    batch: int,
+    dim: tl.constexpr,
+    seqlen: tl.constexpr,
+    state_len: tl.constexpr,
+    num_cache_lines,  # added to support vLLM larger cache lines
+    # Strides
+    stride_x_seq: tl.constexpr,
+    stride_x_dim: tl.constexpr,
+    stride_x_token: tl.int64,
+    stride_w_dim: tl.constexpr,
+    stride_w_width: tl.constexpr,
+    stride_conv_state_seq: tl.constexpr,
+    stride_conv_state_dim: tl.constexpr,
+    stride_conv_state_tok: tl.constexpr,
+    stride_state_indices: tl.constexpr,
+    stride_o_seq: tl.constexpr,
+    stride_o_dim: tl.constexpr,
+    stride_o_token: tl.int64,
+    # others
+    null_block_id: tl.constexpr,
+    # Meta-parameters
+    HAS_BIAS: tl.constexpr,
+    KERNEL_WIDTH: tl.constexpr,
+    SILU_ACTIVATION: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    IS_APC_ENABLED: tl.constexpr,
+    IS_SPEC_DECODING: tl.constexpr,
+    NP2_STATELEN: tl.constexpr,
+    HAS_NULL_BLOCK: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    launch_pdl: tl.constexpr,
+):
+    if launch_pdl:
+        tl.extra.cuda.gdc_wait()
+
+    # ruff: noqa: E501
+    idx_seq = tl.program_id(0)
+    if idx_seq >= batch:
+        if launch_pdl:
+            tl.extra.cuda.gdc_launch_dependents()
+        return
+
+    # [BLOCK_N,] elements along the feature-dimension (channel)
+    idx_feats = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N)
+
+    if IS_APC_ENABLED:
+        # Get the state from the initial_state_idx
+        conv_state_init = tl.load(initial_state_idx + idx_seq)
+        current_last_index = tl.load(block_idx_last_scheduled_token + idx_seq)
+    else:
+        conv_state_init = 0
+        current_last_index = 0
+
+    # cache_idx
+    conv_states_input_coord = tl.load(
+        conv_state_indices_ptr + idx_seq * stride_state_indices + conv_state_init
+    ).to(tl.int64)
+
+    if HAS_NULL_BLOCK:  # noqa
+        if conv_states_input_coord == null_block_id:
+            # not processing as this is not the actual sequence
+            if launch_pdl:
+                tl.extra.cuda.gdc_launch_dependents()
+            return
+
+    if IS_VARLEN:
+        query_start_index = tl.load(query_start_loc_ptr + idx_seq).to(tl.int64)
+        query_end_index = tl.load(query_start_loc_ptr + (idx_seq + 1)).to(tl.int64)
+        # revise state_len and seqlen
+        state_len = state_len - (seqlen - (query_end_index - query_start_index))
+        seqlen = query_end_index - query_start_index
+        x_offset = query_start_index * stride_x_token
+        o_offset = query_start_index * stride_o_token
+    else:
+        query_start_index = idx_seq * seqlen
+        query_end_index = query_start_index + seqlen
+        x_offset = idx_seq * stride_x_seq
+        o_offset = idx_seq * stride_o_seq
+
+    if query_start_index == query_end_index:
+        if launch_pdl:
+            tl.extra.cuda.gdc_launch_dependents()
+        return
+
+    if IS_SPEC_DECODING:
+        # The rolling of conv state:
+        #
+        # Before forward, the conv_state is:
+        # [history1, history2, ..., historyM].
+        #
+        # After forward, the conv_state becomes:
+        # [history2, ..., historyM, draft1, draft2, ..., draftN].
+        #
+        # After acceptance, it becomes:
+        #
+        # - accept 1 tokens: [history2, ..., historyM, draft1]
+        # - accept 2 tokens: [history3, ..., historyM, draft1, draft2]
+        # - and so on.
+        conv_state_token_offset = (
+            tl.load(num_accepted_tokens_ptr + idx_seq).to(tl.int64) - 1
+        )
+    else:
+        conv_state_token_offset = 0
+
+    # STEP 1: READ init_state data
+    conv_states_base = (
+        conv_state_ptr
+        + (conv_states_input_coord * stride_conv_state_seq)
+        + (idx_feats * stride_conv_state_dim)
+    )
+    mask_w = idx_feats < dim
+
+    prior_tokens = conv_states_base + conv_state_token_offset * stride_conv_state_tok
+    if KERNEL_WIDTH >= 2:
+        conv_states_ptrs = prior_tokens  # [BLOCK_N]
+        col0 = tl.load(conv_states_ptrs, mask_w, 0.0)
+    if KERNEL_WIDTH >= 3:
+        conv_states_ptrs = prior_tokens + 1 * stride_conv_state_tok  # [BLOCK_N]
+        col1 = tl.load(conv_states_ptrs, mask_w, 0.0)
+    if KERNEL_WIDTH >= 4:
+        conv_states_ptrs = prior_tokens + 2 * stride_conv_state_tok  # [BLOCK_N]
+        col2 = tl.load(conv_states_ptrs, mask_w, 0.0)
+    if KERNEL_WIDTH >= 5:
+        conv_states_ptrs = prior_tokens + 3 * stride_conv_state_tok  # [BLOCK_N]
+        col3 = tl.load(conv_states_ptrs, mask_w, 0.0)
+    if KERNEL_WIDTH >= 6:
+        conv_states_ptrs = prior_tokens + 4 * stride_conv_state_tok  # [BLOCK_N]
+        col4 = tl.load(conv_states_ptrs, mask_w, 0.0)
+
+    # STEP 2: assume state_len > seqlen
+    idx_tokens = tl.arange(0, NP2_STATELEN)  # [BLOCK_M]
+
+    # With speculative decoding, the conv_state updates works in a sliding
+    # window manner, at each forward pass, the tokens are shift by 1, so we
+    # load since idx_tokens + 1.
+    conv_state_ptrs_source = (
+        conv_state_ptr
+        + (conv_states_input_coord * stride_conv_state_seq)
+        + conv_state_token_offset * stride_conv_state_tok
+        + (idx_feats * stride_conv_state_dim)[None, :]
+        + ((idx_tokens + (1 if IS_SPEC_DECODING else seqlen)) * stride_conv_state_tok)[
+            :, None
+        ]
+    )  # [BLOCK_M, BLOCK_N]
+    mask = (
+        (conv_states_input_coord < num_cache_lines)
+        & ((idx_tokens + seqlen) < state_len)[:, None]
+        & (idx_feats < dim)[None, :]
+    )
+    conv_state = tl.load(conv_state_ptrs_source, mask, other=0.0)
+
+    VAL = state_len - seqlen
+    x_base = x_ptr + x_offset + (idx_feats * stride_x_dim)  # [BLOCK_N]
+
+    x_ptrs = (
+        x_base[None, :] + ((idx_tokens - VAL) * stride_x_token)[:, None]
+    )  # [BLOCK_M, BLOCK_N]
+
+    mask_x = (
+        (idx_tokens - VAL >= 0)[:, None]
+        & (idx_tokens - VAL < seqlen)[:, None]
+        & (idx_feats < dim)[None, :]
+    )  # token-index  # token-index  # feature-index
+    loaded_x = tl.load(x_ptrs, mask_x, 0.0)
+    tl.debug_barrier()
+
+    new_conv_state = tl.where(mask, conv_state, loaded_x)
+
+    # Get the state from the initial_state_idx
+    # cache_idx
+    conv_states_offset = tl.load(
+        conv_state_indices_ptr + idx_seq * stride_state_indices + current_last_index
+    ).to(tl.int64)
+    conv_state_ptrs_target = (
+        conv_state_ptr
+        + (conv_states_offset * stride_conv_state_seq)  # Offset from seq
+        + (idx_feats * stride_conv_state_dim)
+    )[None, :] + (  # [BLOCK_N,]
+        idx_tokens * stride_conv_state_tok
+    )[:, None]
+    mask = (idx_tokens < state_len)[:, None] & (idx_feats < dim)[None, :]
+    tl.store(conv_state_ptrs_target, new_conv_state, mask)
+
+    # STEP 3: init accumulator
+    if HAS_BIAS:
+        bias = bias_ptr + idx_feats
+        mask_bias = idx_feats < dim
+        acc_preload = tl.load(bias, mask=mask_bias, other=0.0).to(
+            tl.float32
+        )  # [BLOCK_N]
+    else:
+        acc_preload = tl.zeros((BLOCK_N,), dtype=tl.float32)
+
+    # STEP 4:
+    # PRE-LOAD WEIGHTS
+    # first kernel column, configured for weights to handle BLOCK_N features in range
+    w_base = w_ptr + (idx_feats * stride_w_dim)  # [BLOCK_N,]
+    mask_w = idx_feats < dim
+    if KERNEL_WIDTH >= 2:
+        w_ptrs = w_base + (0 * stride_w_width)  # [BLOCK_N] tensor
+        w_col0 = tl.load(w_ptrs, mask_w, other=0.0)
+        w_ptrs = w_base + (1 * stride_w_width)  # [BLOCK_N] tensor
+        w_col1 = tl.load(w_ptrs, mask_w, other=0.0)
+    if KERNEL_WIDTH >= 3:
+        w_ptrs = w_base + (2 * stride_w_width)  # [BLOCK_N] tensor
+        w_col2 = tl.load(w_ptrs, mask_w, other=0.0)
+    if KERNEL_WIDTH >= 4:
+        w_ptrs = w_base + (3 * stride_w_width)  # [BLOCK_N] tensor
+        w_col3 = tl.load(w_ptrs, mask_w, other=0.0)
+    if KERNEL_WIDTH >= 5:
+        w_ptrs = w_base + (4 * stride_w_width)  # [BLOCK_N] tensor
+        w_col4 = tl.load(w_ptrs, mask_w, other=0.0)
+    if KERNEL_WIDTH >= 6:
+        w_ptrs = w_base + (5 * stride_w_width)  # [BLOCK_N] tensor
+        w_col5 = tl.load(w_ptrs, mask_w, other=0.0)
+
+    x_base_1d = x_base  # starting of chunk [BLOCK_N]
+    mask_x_1d = idx_feats < dim
+
+    # STEP 5: compute each token
+    if launch_pdl:
+        tl.extra.cuda.gdc_launch_dependents()
+
+    for idx_token in tl.range(seqlen):
+        acc = acc_preload
+
+        matrix_w = w_col0
+        matrix_x = col0
+        for j in tl.static_range(KERNEL_WIDTH):
+            if KERNEL_WIDTH == 2:
+                if j == 1:  # KERNEL_WIDTH-1:
+                    matrix_w = w_col1
+                    x_ptrs_1d = x_base_1d + idx_token * stride_x_token  # [BLOCK_N]
+                    matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d)
+            elif KERNEL_WIDTH == 3:
+                if j == 1:
+                    matrix_w = w_col1
+                    matrix_x = col1
+                elif j == 2:
+                    matrix_w = w_col2
+                    x_ptrs_1d = x_base_1d + idx_token * stride_x_token  # [BLOCK_N]
+                    matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d)
+            elif KERNEL_WIDTH == 4:
+                if j == 1:
+                    matrix_w = w_col1
+                    matrix_x = col1
+                elif j == 2:
+                    matrix_w = w_col2
+                    matrix_x = col2
+                elif j == 3:
+                    matrix_w = w_col3
+                    x_ptrs_1d = x_base_1d + idx_token * stride_x_token  # [BLOCK_N]
+                    matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d)
+            elif KERNEL_WIDTH == 5:
+                if j == 1:
+                    matrix_w = w_col1
+                    matrix_x = col1
+                elif j == 2:
+                    matrix_w = w_col2
+                    matrix_x = col2
+                elif j == 3:
+                    matrix_w = w_col3
+                    matrix_x = col3
+                elif j == 4:
+                    matrix_w = w_col4
+                    x_ptrs_1d = x_base_1d + idx_token * stride_x_token  # [BLOCK_N]
+                    matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d)
+            elif KERNEL_WIDTH == 6:
+                if j == 1:
+                    matrix_w = w_col1
+                    matrix_x = col1
+                elif j == 2:
+                    matrix_w = w_col2
+                    matrix_x = col2
+                elif j == 3:
+                    matrix_w = w_col3
+                    matrix_x = col3
+                elif j == 4:
+                    matrix_w = w_col4
+                    matrix_x = col4
+                elif j == 5:
+                    matrix_w = w_col5
+                    x_ptrs_1d = x_base_1d + idx_token * stride_x_token  # [BLOCK_N]
+                    matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d)
+
+            acc += matrix_x * matrix_w  # [BLOCK_N]
+
+        if KERNEL_WIDTH == 2:
+            col0 = matrix_x
+        elif KERNEL_WIDTH == 3:
+            col0 = col1
+            col1 = matrix_x
+        elif KERNEL_WIDTH == 4:
+            col0 = col1
+            col1 = col2
+            col2 = matrix_x
+        elif KERNEL_WIDTH == 5:
+            col0 = col1
+            col1 = col2
+            col2 = col3
+            col3 = matrix_x
+        elif KERNEL_WIDTH == 6:
+            col0 = col1
+            col1 = col2
+            col2 = col3
+            col3 = col4
+            col4 = matrix_x
+
+        if SILU_ACTIVATION:
+            acc = acc / (1 + tl.exp(-acc))
+        mask_1d = (idx_token < seqlen) & (
+            idx_feats < dim
+        )  # token-index  # feature-index
+        o_ptrs = (
+            o_ptr + o_offset + idx_token * stride_o_token + (idx_feats * stride_o_dim)
+        )
+
+        tl.store(o_ptrs, acc, mask=mask_1d)
+
+
+def causal_conv1d_update(
+    x: torch.Tensor,
+    conv_state: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None = None,
+    activation: bool | str | None = None,
+    conv_state_indices: torch.Tensor | None = None,
+    num_accepted_tokens: torch.Tensor | None = None,
+    query_start_loc: torch.Tensor | None = None,
+    max_query_len: int = -1,
+    null_block_id: int = NULL_BLOCK_ID,
+    block_idx_last_scheduled_token: torch.Tensor | None = None,
+    initial_state_idx: torch.Tensor | None = None,
+    validate_data=False,
+    out: torch.Tensor | None = None,
+):
+    """x: Input tensor which can take the following shapes:
+
+    - `[batch, dim]` - single token prediction
+    - `[batch, dim, seqlen]` - single or multiple tokens prediction
+    - `[num_tokens, dim]` - continuous batching, where num_tokens is
+        the total tokens of all sequences in that batch
+
+    conv_state: (..., dim, state_len), where state_len >= width - 1
+    weight: (dim, width)
+    bias: (dim,)
+    conv_state_indices: (batch,), dtype int32
+        If not None, the conv_state is a larger tensor along the batch dim,
+        and we are selecting the batch coords specified by conv_state_indices.
+        Useful for a continuous batching scenario.
+    block_idx_last_scheduled_token: (batch,), dtype int32
+        The pointer into conv_state_indices, where the last cache block to be filled is located.
+    initial_state_idx: (batch,), dtype int32
+        The pointer into conv_state_indices, where the cache block containing the initial state is located.
+    num_accepted_tokens: (batch,), dtype int32
+        If not None, it indicates the number of accepted tokens for each
+        sequence in the batch.
+        This is used in speculative decoding, where the conv_state is updated
+        in a sliding window manner.
+    query_start_loc: (batch + 1,) int32
+        If not None, the inputs is given in a varlen fashion and this indicates
+        the starting index of each sequence in the batch.
+    max_query_len: int
+        If query_start_loc is not None, this indicates the maximum query
+        length in the batch.
+    null_block_id: int
+            Block ID used to identify padded entries in
+            conv_state_indices. Block 0 is the null block.
+            for example: conv_state_indices = [null_block_id, 1, 20, null_block_id]
+            in this case, the kernel will not process entries at
+            indices 0 and 3
+    out: optional output tensor with the same shape as `x`. When omitted,
+        the input is overwritten.
+    """
+    if validate_data:
+        assert null_block_id is not None
+        assert x.stride(1) == 1
+    if isinstance(activation, bool):
+        activation = "silu" if activation is True else None
+    elif activation is not None:
+        assert activation in ["silu", "swish"]
+
+    original_x_dtype = x.dtype
+    x = x.to(conv_state.dtype)
+    if out is None:
+        out = x
+    else:
+        if out.shape != x.shape:
+            raise ValueError(
+                f"`out` shape {tuple(out.shape)} must match `x` shape {tuple(x.shape)}."
+            )
+        if out.dtype != original_x_dtype or out.device != x.device:
+            raise ValueError(
+                "`out` must have the same dtype and device as the input `x`."
+            )
+    unsqueeze = query_start_loc is None and x.dim() == 2
+    if unsqueeze:
+        # make it (batch, dim, seqlen) with seqlen == 1
+        x = x.unsqueeze(-1)
+        out = out.unsqueeze(-1)
+    if query_start_loc is None:
+        batch, dim, seqlen = x.shape
+    else:
+        assert conv_state_indices is not None
+        batch = conv_state_indices.size(0)
+        dim = x.size(1)
+        seqlen = max_query_len
+    _, width = weight.shape
+    # conv_state: (..., dim, state_len), where state_len >= width - 1
+    num_cache_lines, _, state_len = conv_state.size()
+
+    if validate_data:
+        assert dim == weight.size(0)
+        assert state_len >= width - 1
+        # when above happens, we don't shift-left to keep any records in conv_state
+        assert dim == conv_state.size(1)
+        if conv_state_indices is None:
+            assert conv_state.size(0) >= batch
+        else:
+            assert batch == conv_state_indices.shape[0], (
+                f"ERROR: conv_state_indices should have shape ({batch},*) but got {conv_state_indices.shape}"
+            )
+
+        assert num_cache_lines >= batch
+        assert weight.stride(1) == 1  # Need this
+
+    stride_w_dim, stride_w_width = weight.stride()
+
+    if query_start_loc is None:
+        # X (batch, dim, seqlen)
+        stride_x_seq, stride_x_dim, stride_x_token = x.stride()
+        stride_o_seq, stride_o_dim, stride_o_token = out.stride()
+    else:
+        # X (dim, cu_seqlen)
+        stride_x_token, stride_x_dim = x.stride()
+        stride_x_seq = 0
+        stride_o_token, stride_o_dim = out.stride()
+        stride_o_seq = 0
+
+    stride_istate_seq, stride_istate_dim, stride_istate_token = conv_state.stride()
+    stride_state_indices = (
+        conv_state_indices.stride(0) if conv_state_indices is not None else 0
+    )
+    if num_accepted_tokens is not None:
+        state_len = width - 1 + (seqlen - 1)  # effective state_len needed
+    else:
+        state_len = width - 1
+    np2_statelen = triton.next_power_of_2(state_len)
+
+    def grid(META):
+        return (
+            batch,
+            triton.cdiv(dim, META["BLOCK_N"]),
+        )
+
+    _causal_conv1d_update_kernel[grid](
+        # Pointers to matrices
+        x,
+        weight,
+        bias,
+        conv_state,
+        conv_state_indices,
+        num_accepted_tokens,
+        query_start_loc,
+        block_idx_last_scheduled_token,
+        initial_state_idx,
+        out,
+        # Matrix dimensions
+        batch,
+        dim,
+        seqlen,
+        state_len,
+        num_cache_lines,
+        # stride
+        stride_x_seq,
+        stride_x_dim,
+        stride_x_token,
+        stride_w_dim,
+        stride_w_width,
+        stride_istate_seq,
+        stride_istate_dim,
+        stride_istate_token,
+        stride_state_indices,
+        stride_o_seq,
+        stride_o_dim,
+        stride_o_token,
+        # others
+        null_block_id,
+        # META
+        HAS_BIAS=bias is not None,
+        KERNEL_WIDTH=width,
+        SILU_ACTIVATION=activation in ["silu", "swish"],
+        IS_VARLEN=query_start_loc is not None,
+        IS_APC_ENABLED=block_idx_last_scheduled_token is not None,
+        IS_SPEC_DECODING=num_accepted_tokens is not None,
+        NP2_STATELEN=np2_statelen,
+        HAS_NULL_BLOCK=null_block_id is not None,
+        BLOCK_N=256,
+        launch_pdl=current_platform.is_arch_support_pdl(),
+    )
+    if unsqueeze:
+        out = out.squeeze(-1)
+    return out.to(original_x_dtype)
+
+
+if current_platform.is_cpu():
+    from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (
+        causal_conv1d_fn_cpu,
+        causal_conv1d_update_cpu,
+    )
+
+    causal_conv1d_fn = causal_conv1d_fn_cpu  # type: ignore
+    causal_conv1d_update = causal_conv1d_update_cpu  # type: ignore

docker-images/glm53-kindling/experimental/fixes/flashinfer_mla_sparse_sm90.py

@@ -0,0 +1,526 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""FlashInfer sparse MLA backend for SM90 (Hopper) NoPE models.
+
+Wraps FlashInfer's ``BatchMLAPagedAttentionWrapper`` (FA2/FA3 paths), which
+as of FlashInfer 0.6.18 supports ``head_dim_kpe=0`` (GLM-5.3-Flash NoPE MLA)
+and FP8 E4M3 KV caches on SM90 with in-kernel dequantization: the FP8 cache
+is read directly (half the bf16 HBM traffic) and converted to BF16 in shared
+memory, while queries stay BF16 (no query quantization).
+
+Sparsity rides the same trick the FA-based sparse backend uses: with
+``page_size=1`` the per-token top-k slot indices ARE the page table, so each
+query token becomes one varlen batch row whose ``kv_indices`` slice is its
+top-k row and whose ``kv_len`` is its valid count. Causality is already
+encoded by the indexer's selection, so ``causal=False``.
+
+CUDA-graph handling: ``plan()`` copies its inputs to host unconditionally,
+so it must stay outside graph capture. Each metadata builder owns a wrapper,
+reserved capture-stable device buffers, and the plan parameters. The wrapper
+bakes the per-row ``kv_len`` into its int schedule
+at plan() time — ``run()`` never reads the device-side buffer — so the
+metadata builder replans every step (outside capture) with exact host-side
+lengths derived from the batch's sequence lengths; a full-width schedule
+would send the kernel past each row's valid count into the -1 tail of the
+converted index buffer (illegal address). Per-step content (top-k slots)
+is written into the reserved buffers by kernels inside the captured
+forward, and captured runs read the refreshed plan buffers on replay.
+
+KV cache format: plain contiguous E4M3 ``[num_blocks, block_size, 512]``
+(uint8 storage) with a per-tensor ``k_scale``; BF16 caches also work. The
+per-token x 128-channel-group ``ckv_scale_arr`` layout is supported by the
+kernel but not wired yet (it needs a group-quantizing cache-write op).
+"""
+
+import os
+from dataclasses import dataclass
+from typing import Any, ClassVar
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.config.cache import CacheDType
+from vllm.model_executor.layers.attention.sparse_mla_attention import (
+    SparseMLACommonImpl,
+)
+from vllm.platforms.interface import DeviceCapability
+from vllm.triton_utils import tl, triton
+from vllm.utils.flashinfer import has_flashinfer_sm90_nope_mla
+from vllm.v1.attention.backend import (
+    AttentionBackend,
+    AttentionLayer,
+    CommonAttentionMetadata,
+    MLAAttentionImpl,
+    MultipleOf,
+)
+from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
+    FlashInferMLASparseMetadata,
+    FlashInferMLASparseMetadataBuilder,
+)
+from vllm.v1.attention.backends.mla.sparse_utils import (
+    triton_convert_req_index_to_global_index,
+)
+from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout
+
+_FP8_KV_DTYPES = ("fp8", "fp8_e4m3")
+_WORKSPACE_BYTES = 128 * 1024 * 1024
+
+# GB10: a Triton sparse MLA kernel in place of FlashInfer's FA2 path, which
+# runs at ~10 TFLOPS on SM121 (gb10_sparse_mla.py; about 3.3x faster).
+_TRITON = os.environ.get("VLLM_TRITON_SPARSE_MLA") == "1"
+
+
+@triton.jit
+def _pack_topk_indices(
+    slots,
+    indptr,
+    indices,
+    WIDTH: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    cols = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    start = tl.load(indptr + row)
+    end = tl.load(indptr + row + 1)
+    mask = (cols < end - start) & (cols < WIDTH)
+    values = tl.load(slots + row * WIDTH + cols, mask=mask, other=0)
+    tl.store(indices + start + cols, tl.maximum(values, 0), mask=mask)
+
+
+class FlashInferMLASparseSM90Backend(AttentionBackend):
+    supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
+    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
+        "auto",
+        "bfloat16",
+        "fp8",
+        "fp8_e4m3",
+    ]
+
+    @staticmethod
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
+        return [MultipleOf(64)]
+
+    @staticmethod
+    def get_name() -> str:
+        return "FLASHINFER_MLA_SPARSE_SM90"
+
+    @staticmethod
+    def get_builder_cls() -> type["FlashInferMLASparseSM90Builder"]:
+        return FlashInferMLASparseSM90Builder
+
+    @staticmethod
+    def get_impl_cls() -> type[MLAAttentionImpl]:
+        return FlashInferMLASparseSM90Impl
+
+    @classmethod
+    def get_supported_head_sizes(cls) -> list[int]:
+        # 512 = ckv 512 + kpe 0 (NoPE); 576 = ckv 512 + kpe 64.
+        return [512, 576]
+
+    @classmethod
+    def is_mla(cls) -> bool:
+        return True
+
+    @classmethod
+    def is_sparse(cls) -> bool:
+        return True
+
+    @classmethod
+    def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
+        return capability.major in (9, 12)
+
+    @classmethod
+    def supports_combination(
+        cls,
+        head_size: int,
+        dtype: torch.dtype,
+        kv_cache_dtype: CacheDType | None,
+        block_size: int | None,
+        use_mla: bool,
+        has_sink: bool,
+        use_sparse: bool,
+        use_mm_prefix: bool,
+        device_capability: DeviceCapability,
+    ) -> str | None:
+        if kv_cache_dtype in ("fp8", "fp8_e4m3") and not has_flashinfer_sm90_nope_mla():
+            return (
+                "FLASHINFER_MLA_SPARSE_SM90 fp8 KV requires FlashInfer with "
+                "SM90 MLA support (ckv_scale_arr in "
+                "BatchMLAPagedAttentionWrapper.run, FlashInfer >= 0.6.18)"
+            )
+        if not use_sparse:
+            return "FLASHINFER_MLA_SPARSE_SM90 requires sparse MLA"
+        from vllm.config import get_current_vllm_config
+
+        vllm_config = get_current_vllm_config()
+        if vllm_config.model_config is not None:
+            hf = vllm_config.model_config.hf_text_config
+            # The SM90 FA2/FA3 kernel covers ckv=512 with kpe in {0, 64}
+            # (NoPE models and DeepSeek-style rope MLA alike).
+            if hf.kv_lora_rank != 512:
+                return "FLASHINFER_MLA_SPARSE_SM90 requires kv_lora_rank=512"
+            if hf.qk_rope_head_dim not in (0, 64):
+                return "FLASHINFER_MLA_SPARSE_SM90 requires qk_rope_head_dim in (0, 64)"
+        return None
+
+    @staticmethod
+    def get_kv_cache_shape(
+        num_blocks: int,
+        block_size: int,
+        num_kv_heads: int,
+        head_size: int,
+        cache_dtype_str: str = "auto",
+    ) -> tuple[int, ...]:
+        return (num_blocks, block_size, head_size)
+
+    @classmethod
+    def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]:
+        return (KVCacheLayout.LBHNC,)
+
+
+class _SM90State:
+    """Builder-owned wrapper, capture-stable buffers, and plan parameters.
+
+    One instance serves every MLA layer in an attention group because the plan
+    depends only on the batch shape, not the layer.
+    """
+
+    def __init__(
+        self,
+        device: torch.device,
+        num_heads: int,
+        kv_dtype: torch.dtype,
+        max_tokens: int,
+        topk_width: int,
+        kv_lora_rank: int,
+        qk_rope_head_dim: int,
+        sm_scale: float,
+    ) -> None:
+        from flashinfer.mla import BatchMLAPagedAttentionWrapper
+
+        self.workspace = torch.empty(_WORKSPACE_BYTES, dtype=torch.uint8, device=device)
+        self.device = device
+        self.num_heads = num_heads
+        self.kv_dtype = kv_dtype
+        self.max_tokens = max_tokens
+        self.topk_width = topk_width
+        self.kv_lora_rank = kv_lora_rank
+        self.qk_rope_head_dim = qk_rope_head_dim
+        self.sm_scale = sm_scale
+        # User-reserved buffers: with use_cuda_graph=True plan() refreshes
+        # these in place, so run()'s captured kernels always read them.
+        self.kv_indices = torch.zeros(
+            max_tokens * topk_width, dtype=torch.int32, device=device
+        )
+        self.kv_len_arr = torch.full(
+            (max_tokens,), topk_width, dtype=torch.int32, device=device
+        )
+        self.kv_indptr = torch.zeros(max_tokens + 1, dtype=torch.int32, device=device)
+        self.wrapper = BatchMLAPagedAttentionWrapper(
+            self.workspace,
+            qo_indptr=torch.zeros(max_tokens + 1, dtype=torch.int32, device=device),
+            kv_indptr=self.kv_indptr,
+            kv_indices=self.kv_indices,
+            kv_len_arr=self.kv_len_arr,
+            use_cuda_graph=True,
+            backend=("fa3" if torch.cuda.get_device_capability()[0] == 9 else "fa2"),
+        )
+        self._arange_cpu = torch.arange(self.max_tokens + 1, dtype=torch.int32)
+        self._qo_cpu = torch.empty(self.max_tokens + 1, dtype=torch.int32)
+        self._kv_cpu = torch.empty(self.max_tokens + 1, dtype=torch.int32)
+        self._lens_cpu = torch.full(
+            (self.max_tokens,), self.topk_width, dtype=torch.int32
+        )
+
+    def plan(self, num_tokens: int, kv_lens: torch.Tensor) -> None:
+        """Replan with exact per-row KV lengths (CPU int32, ``[num_tokens]``).
+
+        The wrapper bakes kv_len into its int schedule from host values;
+        run() never reads the device kv_len_arr buffer. Scheduling with the
+        full buffer width while kv_indices rows carry a -1 tail past each
+        row's valid count makes the kernel compute ``-1 * ckv_stride_page``
+        (illegal address), so the lengths must be exact at plan time and
+        replanned every step as contexts grow. Must run outside CUDA graph
+        capture: the in-place refreshed plan_info/indptr buffers are what
+        captured runs read.
+        """
+        if torch.cuda.is_current_stream_capturing():
+            raise RuntimeError(
+                "FlashInferMLASparseSM90 plan() called inside CUDA graph "
+                "capture; lengths must be planned host-side before capture."
+            )
+        # CPU staging buffers are filled in place: plan() runs per step
+        # (once per draft/verify metadata build), so per-call allocations and
+        # device round trips are on the hot path. Passing CPU tensors lets
+        # the wrapper's internal .to("cpu") no-op; its reserved-buffer
+        # copy_ then performs the single H2D transfer per tensor.
+        # use_cuda_graph=True makes the wrapper copy qo/kv indptr into its
+        # fixed (max_tokens+1)-sized buffers with exact-size copy_, so the
+        # indptr must always be full-size. Rows past num_tokens are padded
+        # empty (qo_indptr flat at num_tokens) — zero-query rows read no q
+        # and schedule no work. FlashInfer 0.7 requires packed CSR offsets
+        # matching the exact lengths, including zero-length padded rows.
+        torch.clamp(self._arange_cpu, max=num_tokens, out=self._qo_cpu)
+        self._lens_cpu.zero_()
+        self._lens_cpu[:num_tokens] = kv_lens.to(torch.int32)
+        self._kv_cpu[0] = 0
+        torch.cumsum(self._lens_cpu, dim=0, out=self._kv_cpu[1:])
+        self.wrapper.plan(
+            self._qo_cpu,
+            self._kv_cpu,
+            self.kv_indices,
+            self._lens_cpu,
+            self.num_heads,
+            self.kv_lora_rank,  # head_dim_ckv
+            self.qk_rope_head_dim,  # 0 (NoPE) or 64 (rope MLA)
+            1,  # page_size: top-k slots are the page table
+            False,  # causal: encoded by the indexer's selection
+            self.sm_scale,
+            q_data_type=torch.bfloat16,
+            kv_data_type=self.kv_dtype,
+        )
+
+    def pack_indices(self, topk_slots: torch.Tensor) -> None:
+        num_tokens, width = topk_slots.shape
+        _pack_topk_indices[(num_tokens, triton.cdiv(width, 256))](
+            topk_slots, self.kv_indptr, self.kv_indices, width, 256
+        )
+
+
+@dataclass
+class FlashInferMLASparseSM90Metadata(FlashInferMLASparseMetadata):
+    state: _SM90State | None = None
+
+
+class FlashInferMLASparseSM90Builder(FlashInferMLASparseMetadataBuilder):
+    """Reuse the common sparse metadata (req ids, topk buffer access)."""
+
+    metadata_cls = FlashInferMLASparseSM90Metadata
+
+    def __init__(
+        self,
+        kv_cache_spec: "AttentionSpec",
+        layer_names: list[str],
+        vllm_config: "VllmConfig",
+        device: torch.device,
+    ) -> None:
+        super().__init__(kv_cache_spec, layer_names, vllm_config, device)
+        attention_layer = vllm_config.compilation_config.static_forward_context[
+            layer_names[0]
+        ]
+        self._attention_layer = attention_layer
+        impl = attention_layer.impl
+        if not isinstance(impl, FlashInferMLASparseSM90Impl):
+            raise TypeError(
+                "FlashInferMLASparseSM90Builder requires an SM90 FlashInfer "
+                f"implementation, got {type(impl).__name__}."
+            )
+        topk_indices_buffer = impl.topk_indices_buffer
+        assert topk_indices_buffer is not None
+        self.state = _SM90State(
+            device,
+            impl.num_heads,
+            # SM90-FP8-KV-DTYPE: fp8 caches are stored as uint8; plan with fp8.
+            (torch.float8_e4m3fn if impl.use_fp8_kv_cache else kv_cache_spec.dtype),
+            vllm_config.scheduler_config.max_num_batched_tokens,
+            topk_indices_buffer.shape[1],
+            kv_lora_rank=impl.kv_lora_rank,
+            qk_rope_head_dim=impl.qk_rope_head_dim,
+            sm_scale=impl.scale,
+        )
+        # seq_lens_cpu_upper_bound is optimistic on decode rows under async
+        # spec decode, so the fast sync-free path is only safe without it;
+        # under async scheduling the exact (device) positions are used at
+        # the cost of one D2H sync per metadata build.
+        self._async_scheduling = bool(vllm_config.scheduler_config.async_scheduling)
+        hf_config = vllm_config.model_config.hf_text_config
+        assert hf_config.index_topk is not None
+        self._index_topk = int(hf_config.index_topk)
+        self._index_kpool = int(getattr(hf_config, "index_kpool", 1) or 1)
+
+    def _kv_lens_host(self, cam: CommonAttentionMetadata) -> tuple[int, torch.Tensor]:
+        """Exact per-row KV lengths, host-side (the flashinfer wrapper bakes
+        them into its schedule at plan time; there is no device-side path).
+
+        A row for the j-th query token of request i attends
+        ``seq_lens[i] - q_len[i] + j + 1`` tokens. The indexer's selection
+        then bounds the valid count: contexts up to ``index_topk`` select
+        everything (valid == context); longer contexts keep the top
+        ``index_topk`` pool-expanded tokens plus the trailing incomplete
+        pool (valid == ``index_topk + context % index_kpool``). Both match
+        the count of non -1 entries the convert kernel produces.
+        """
+        num_reqs = cam.num_reqs
+        qsl = cam.query_start_loc_cpu[: num_reqs + 1]
+        num_rows = int(qsl[-1])
+        if num_rows == 0:
+            return 0, torch.zeros(0, dtype=torch.int32)
+        # Row context == position + 1. Without async scheduling the
+        # maintained host upper bound is exact, giving a sync-free path;
+        # under async scheduling it is optimistic on decode rows, so fall
+        # back to the exact device positions (one D2H sync per build). The
+        # seq_lens derivation equals position + 1 because positions are
+        # contiguous per request.
+        sl_host = cam.seq_lens_cpu_upper_bound
+        positions = cam.positions
+        if not self._async_scheduling and sl_host is not None:
+            seq_lens = sl_host[:num_reqs].to(torch.int32)
+            q_lens = qsl[1:] - qsl[:-1]
+            first_pos = seq_lens - q_lens
+            req_of_row = torch.repeat_interleave(
+                torch.arange(num_reqs, dtype=torch.int64), q_lens.to(torch.int64)
+            )
+            rows = torch.arange(num_rows, dtype=torch.int32)
+            ctx = (
+                first_pos.to(torch.int64)[req_of_row]
+                + rows.to(torch.int64)
+                - qsl.to(torch.int64)[req_of_row]
+                + 1
+            )
+        elif positions is not None and num_rows <= positions.shape[0]:
+            ctx = positions[:num_rows].cpu().to(torch.int64) + 1
+        else:
+            seq_lens = cam.seq_lens[:num_reqs].cpu().to(torch.int32)
+            q_lens = qsl[1:] - qsl[:-1]
+            first_pos = seq_lens - q_lens
+            req_of_row = torch.repeat_interleave(
+                torch.arange(num_reqs, dtype=torch.int64), q_lens.to(torch.int64)
+            )
+            rows = torch.arange(num_rows, dtype=torch.int32)
+            ctx = (
+                first_pos.to(torch.int64)[req_of_row]
+                + rows.to(torch.int64)
+                - qsl.to(torch.int64)[req_of_row]
+                + 1
+            )
+        topk = self._index_topk
+        lens = torch.where(ctx <= topk, ctx, topk + ctx % self._index_kpool)
+        return num_rows, lens.to(torch.int32)
+
+    def build(
+        self,
+        common_prefix_len: int,
+        common_attn_metadata: CommonAttentionMetadata,
+        fast_build: bool = False,
+    ) -> FlashInferMLASparseSM90Metadata:
+        metadata = super().build(common_prefix_len, common_attn_metadata, fast_build)
+        assert isinstance(metadata, FlashInferMLASparseSM90Metadata)
+        if _TRITON:
+            # The Triton kernel needs no host-side schedule, and the host KV
+            # lengths cost a device-to-host sync per step under async
+            # scheduling (it stalls the CPU until the previous step ends).
+            metadata.state = self.state
+            return metadata
+        # Replan every step outside any CUDA graph capture with this step's
+        # exact per-row lengths; captured runs read the refreshed buffers.
+        num_rows, kv_lens = self._kv_lens_host(common_attn_metadata)
+        # MHA prefills are removed from q before forward_mqa. Plan only the
+        # decode prefix in that case, using the same routing decision as MLA.
+        if metadata.num_prefills > 0 and self._attention_layer._use_sparse_mha(
+            metadata
+        ):
+            num_rows = metadata.num_decode_tokens
+            kv_lens = kv_lens[:num_rows]
+        self.state.plan(num_rows, kv_lens)
+        metadata.state = self.state
+        return metadata
+
+
+class FlashInferMLASparseSM90Impl(SparseMLACommonImpl[FlashInferMLASparseSM90Metadata]):
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        scale: float,
+        num_kv_heads: int,
+        alibi_slopes: list[float] | None,
+        sliding_window: int | None,
+        kv_cache_dtype: str,
+        logits_soft_cap: float | None,
+        attn_type: str,
+        kv_sharing_target_layer_name: str | None,
+        topk_indices_buffer: torch.Tensor | None = None,
+        indexer: Any | None = None,
+        **mla_args: Any,
+    ) -> None:
+        if any([alibi_slopes, sliding_window, logits_soft_cap]):
+            raise NotImplementedError(
+                "FlashInferMLASparseSM90Impl does not support alibi, sliding "
+                "window, or logits soft cap."
+            )
+        super().__init__(
+            num_heads,
+            head_size,
+            scale,
+            num_kv_heads,
+            alibi_slopes,
+            sliding_window,
+            kv_cache_dtype,
+            logits_soft_cap,
+            attn_type,
+            kv_sharing_target_layer_name,
+            indexer=indexer,
+            topk_indices_buffer=topk_indices_buffer,
+            **mla_args,
+        )
+        assert self.topk_indices_buffer is not None
+        self.supports_quant_query_input = False
+        self.use_fp8_kv_cache = self.kv_cache_dtype in _FP8_KV_DTYPES
+
+    def forward_mqa(
+        self,
+        q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+        kv_c_and_k_pe_cache: torch.Tensor,
+        attn_metadata: FlashInferMLASparseSM90Metadata,
+        layer: AttentionLayer,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        if not isinstance(q, tuple):
+            raise NotImplementedError(
+                "FlashInferMLASparseSM90Impl expects split (q_nope, q_rope)."
+            )
+        q_nope, q_rope = q
+        num_tokens = q_rope.shape[0]
+        # NoPE models hand a zero-width rope tensor through; rope MLA hands
+        # the real 64-dim part. The kernel takes both as-is.
+        q_pe = q_rope.reshape(num_tokens, self.num_heads, self.qk_rope_head_dim)
+
+        assert self.topk_indices_buffer is not None
+        topk_indices = self.topk_indices_buffer[:num_tokens]
+        # return_valid_counts=True keeps the compacted-prefix layout: valid
+        # entries at [0, valid_count), -1 past it — exactly the prefix the
+        # planned per-row lengths address.
+        topk_slots, valid_counts = triton_convert_req_index_to_global_index(
+            attn_metadata.req_id_per_token[:num_tokens],
+            attn_metadata.block_table,
+            topk_indices,
+            BLOCK_SIZE=attn_metadata.block_size,
+            NUM_TOPK_TOKENS=topk_indices.shape[1],
+            return_valid_counts=True,
+        )
+        if _TRITON:
+            from vllm.v1.attention.backends.mla.gb10_sparse_mla import sparse_mla
+
+            kv = kv_c_and_k_pe_cache.view(torch.float8_e4m3fn) if self.use_fp8_kv_cache else kv_c_and_k_pe_cache
+            out = sparse_mla(q_nope, kv.view(-1, self.head_size), topk_slots, valid_counts, self.scale,
+                             float(layer._k_scale_float or 1.0) if self.use_fp8_kv_cache else 1.0)
+            return out, None
+        state = attn_metadata.state
+        assert state is not None
+        # Pack valid prefixes at the offsets refreshed by plan() before replay.
+        state.pack_indices(topk_slots)
+
+        flat = (
+            kv_c_and_k_pe_cache.view(torch.float8_e4m3fn)
+            if self.use_fp8_kv_cache
+            else kv_c_and_k_pe_cache
+        ).reshape(-1, 1, self.head_size)
+        ckv = flat[..., : self.kv_lora_rank]
+        kpe = flat[..., self.kv_lora_rank :]
+
+        scale_kwargs = (
+            {"ckv_scale": float(layer._k_scale_float or 1.0), "kpe_scale": 1.0}
+            if self.use_fp8_kv_cache
+            else {}
+        )
+        out = state.wrapper.run(q_nope, q_pe, ckv, kpe, **scale_kwargs)
+        return out, None