vLLM hybrid Mamba + MTP corruption reproducer (reorder threshold)

README.md

Reproducer: hybrid Mamba + MTP corruption from a lowered batch-reorder threshold (vLLM 0.27.1)

Files

• probe_sitecustomize.py — put its directory on PYTHONPATH of vllm serve and set RACE_PROBE=log,geom RACE_PROBE_LOG=/path/probe.jsonl. Logs the runner's reorder_batch_threshold, every builder's threshold, the rows of every Mamba2 prefill step ([request, num_computed_before, tokens_this_step]) and every batch swap. No behaviour change. RACE_PROBE=log,geom,reorderfix additionally forces the runner threshold to 1 + k (the monkeypatch used as the positive control).
• burst.py — load generator. --mode pairs --pair-long-reps 1200 --pair-delay 1.5 --pair-delay-jitter 1.0 --concurrency 60 --n 400 --max-tokens 3000 sends ~30 concurrent pairs of one ~60k-token prompt followed 1.5 ± 1 s later by a short one; flags responses as corrupted by out-of-vocabulary rate against the batch's own vocabulary, and records each request's engine id.
• reanalyze.py <burst.jsonl> <probe.jsonl> — adds the "thousands of tokens rendering to no text" signature and cross-tabulates corruption with whether a request's first decode step sat in the Mamba2 prefill region.

Server (production flags)

vllm serve nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 --trust-remote-code \
  --max-model-len 262144 --max-num-seqs 96 --gpu-memory-utilization 0.9 \
  --attention-backend TRITON_ATTN --mamba-backend triton --mamba-cache-mode align \
  --kv-cache-dtype fp8 --async-scheduling --enable-prefix-caching \
  --reasoning-parser nemotron_v3 --enable-auto-tool-choice --tool-call-parser qwen3_coder \
  --speculative-config '{"method":"mtp","num_speculative_tokens":3}'

Workaround: add "attention_backend":"TRITON_ATTN" to --speculative-config.

burst.py

"""Sustained-load harness: keep `--concurrency` chat requests in flight until `--n` are done,
then flag corrupted outputs. Each prompt carries a unique marker; a response containing another
request's marker is direct evidence of state shared across sequences. Word-salad is flagged by
the out-of-vocabulary rate against the batch's own vocabulary (words used by >=3% of responses),
which is ~0.03 on normal reasoning and >0.3 on the corrupted production output, plus the
'th'/'ch' fragment signature of that output. Full text is kept for re-analysis.
"""
import argparse, asyncio, json, random, re, string, time, zlib
from collections import Counter
import aiohttp

QUESTIONS = [
    "Let N be the number of ordered pairs (a, b) of positive integers with a, b <= 60 such that a*b is a perfect square and a + b is divisible by 7. Find N.",
    "A uniform solid sphere of mass 2.0 kg and radius 0.10 m rolls without slipping down a 30 degree incline from rest. What is its translational speed after its centre has descended 1.5 m vertically, and what fraction of its kinetic energy is rotational? (A) 4.6 m/s, 2/7 (B) 5.4 m/s, 2/7 (C) 4.6 m/s, 2/5 (D) 5.4 m/s, 2/5",
    "Compound X (C8H10O) shows a broad IR band at 3300 cm^-1, and its 1H NMR shows a 5H multiplet at 7.3 ppm, a 1H quartet at 4.9 ppm, a 1H broad singlet at 2.0 ppm and a 3H doublet at 1.5 ppm. Treatment with PCC gives Y. Identify X and Y and give the product of Y with excess CH3MgBr followed by aqueous workup. (A) 1-phenylethanol; acetophenone; 2-phenylpropan-2-ol (B) 2-phenylethanol; phenylacetaldehyde; 1-phenylpropan-2-ol (C) benzyl methyl ether; benzaldehyde; 1-phenylethanol (D) 4-methylphenol; no reaction; no reaction",
    "In a population in Hardy-Weinberg equilibrium, an X-linked recessive condition affects 1 in 2500 males. What fraction of females are carriers, and what fraction of females are affected? (A) 1/1250 carriers, 1/6,250,000 affected (B) 1/2500 carriers, 1/2500 affected (C) 2/2500 carriers, 1/6,250,000 affected (D) 1/50 carriers, 1/2500 affected",
    "Find the remainder when 7^2026 + 2026^7 is divided by 1000.",
    "Two events in frame S occur at the same place, 3.0 microseconds apart. In frame S' they are 5.0 microseconds apart. What is the spatial separation of the events in S', and the relative speed of the frames? (A) 1200 m, 0.8c (B) 1200 m, 0.6c (C) 900 m, 0.8c (D) 1500 m, 0.6c",
    "A fair coin is flipped 12 times. What is the probability that no two consecutive flips are both heads? Give the answer as a reduced fraction.",
    "A protein-coding gene has a point mutation in the third position of a codon that changes UAC to UAA. Which statement about the effect is correct, and why? (A) Silent, because third-position changes are always synonymous (B) Nonsense, producing a truncated protein (C) Missense, substituting one amino acid (D) Frameshift, altering all downstream codons",
]
SIG = {"th", "ch", "chch", "thch", "thth"}
FILLER = ("Context notes (may be ignored if not needed): the following background describes standard conventions used in "
          "this problem set. Quantities are in SI units unless stated otherwise, ideal conditions are assumed, and answers "
          "should be justified from first principles rather than recalled. ")


import contextlib
_nullcm = contextlib.nullcontext


async def pair(session, url, model, k, markers, max_tokens, sem, out, delay, jitter=0.0, long_reps=480):
    """A 24k-token-prompt request, then `delay` seconds later a short one: the short one's
    prefill lands in a step where the long one is still on a continuing chunk."""
    i_long, i_short = 2 * k, 2 * k + 1
    async with sem:
        t_long = asyncio.create_task(one(session, url, model, i_long, markers[i_long], True, max_tokens, None, out, filler_reps=long_reps, kind="long"))
        await asyncio.sleep(max(0.0, delay + random.uniform(-jitter, jitter)))
        t_short = asyncio.create_task(one(session, url, model, i_short, markers[i_short], False, max_tokens, None, out, filler_reps=0, kind="follower"))
        await asyncio.gather(t_long, t_short)


def analyse(rows):
    ok = [r for r in rows if "text" in r]
    cnt = Counter()
    for r in ok:
        cnt.update({w.lower() for w in re.findall(r"[A-Za-z]{3,}", r["text"])})
    vocab = {w for w, c in cnt.items() if c >= max(3, 0.03 * len(ok))}
    for r in ok:
        ws = [w.lower() for w in re.findall(r"[A-Za-z]+", r["text"])]
        long_ws = [w for w in ws if len(w) >= 3]
        r["oov"] = round(sum(1 for w in long_ws if w not in vocab) / max(1, len(long_ws)), 4)
        r["sig"] = round(sum(1 for w in ws if w in SIG) / max(1, len(ws)), 4)
        r["zlib"] = round(len(zlib.compress(r["text"].encode())) / max(1, len(r["text"])), 3)
        # production word-salad: oov 0.65 and 'th'/'ch' at 7% of words; chemistry traces reach sig 0.08 but oov < 0.08
        r["garbage"] = r["oov"] > 0.30 or (r["sig"] > 0.04 and r["oov"] > 0.10)
    return ok


async def one(session, url, model, i, marker, long_prompt, max_tokens, sem, out, filler_reps=None, kind=None):
    q = QUESTIONS[i % len(QUESTIONS)]
    # mixed mode: 1 in 7 prompts is ~24k tokens (tau2-conversation scale), the other "long" ones
    # ~2.3k (livecodebench scale), the rest a bare question. pairs mode passes filler_reps explicitly.
    if filler_reps is None:
        filler_reps = 480 if i % 7 == 0 else 45 if long_prompt else 0
    filler = FILLER * filler_reps
    body = {"model": model, "temperature": 1.0, "top_p": 0.95, "seed": i, "max_tokens": max_tokens,
            "chat_template_kwargs": {"enable_thinking": True},
            "messages": [{"role": "user", "content": f"{filler}Reference code {marker}.\n\n{q}\n\nThink step by step, then state the final answer."}]}
    async with (sem if sem is not None else _nullcm()):
        t0 = time.time()
        try:
            async with session.post(f"{url}/v1/chat/completions", json=body, timeout=aiohttp.ClientTimeout(total=7200)) as resp:
                d = await resp.json()
        except Exception as e:
            out.append({"i": i, "error": str(e)[:200], "t_start": t0}); return
    ch = d["choices"][0]; m = ch["message"]
    text = (m.get("reasoning_content") or m.get("reasoning") or "") + "</think>" + (m.get("content") or "")
    out.append({"i": i, "req_id": d.get("id"), "marker": marker, "long": long_prompt, "kind": kind or ("24k" if filler_reps >= 400 else "2k" if filler_reps else "short"), "finish": ch["finish_reason"],
                "completion_tokens": d.get("usage", {}).get("completion_tokens"), "prompt_tokens": d.get("usage", {}).get("prompt_tokens"),
                "t_start": t0, "t_end": time.time(), "text": text})


async 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("--n", type=int, default=600); ap.add_argument("--max-tokens", type=int, default=16000)
    ap.add_argument("--concurrency", type=int, default=110); ap.add_argument("--out", required=True)
    ap.add_argument("--mode", choices=["mixed", "pairs"], default="mixed"); ap.add_argument("--pair-delay", type=float, default=1.0); ap.add_argument("--pair-delay-jitter", type=float, default=0.0); ap.add_argument("--pair-long-reps", type=int, default=480, help="FILLER repetitions for the long prompt (~50 tokens each)")
    a = ap.parse_args()
    rnd = random.Random(0)
    markers = ["".join(rnd.choices(string.ascii_lowercase, k=9)) for _ in range(a.n)]
    out = []; sem = asyncio.Semaphore(a.concurrency)
    async with aiohttp.ClientSession() as s:
        if a.mode == "pairs":
            psem = asyncio.Semaphore(max(1, a.concurrency // 2))
            await asyncio.gather(*(pair(s, a.url, a.model, k, markers, a.max_tokens, psem, out, a.pair_delay, a.pair_delay_jitter, a.pair_long_reps) for k in range(a.n // 2)))
        else:
            await asyncio.gather(*(one(s, a.url, a.model, i, markers[i], i % 2 == 0, a.max_tokens, sem, out) for i in range(a.n)))
    ok = analyse(out)
    mset = {m: i for i, m in enumerate(markers)}
    for r in ok:
        r["foreign_markers"] = sorted({mset[m] for m in mset if m in r["text"] and mset[m] != r["i"]})
    with open(a.out, "w") as f:
        for r in sorted(out, key=lambda r: r["i"]): f.write(json.dumps(r) + "\n")
    errs = len(out) - len(ok)
    g = [r for r in ok if r["garbage"]]; leak = [r for r in ok if r["foreign_markers"]]
    import statistics as st
    print(f"RESULT n={a.n} errors={errs} garbage={len(g)} ({100*len(g)/max(1,len(ok)):.1f}%) foreign_marker_hits={len(leak)} "
          f"finish_length={sum(r['finish']=='length' for r in ok)} median_tokens={st.median(r['completion_tokens'] or 0 for r in ok):.0f} "
          f"median_oov={st.median(r['oov'] for r in ok):.3f} wall={max(r['t_end'] for r in ok)-min(r['t_start'] for r in ok):.0f}s")
    from collections import defaultdict
    byk = defaultdict(lambda: [0, 0])
    for r in ok: byk[r["kind"]][0] += 1; byk[r["kind"]][1] += r["garbage"]
    print("  by prompt kind: " + "; ".join(f"{k}: {gg}/{n} garbage" for k, (n, gg) in sorted(byk.items())))
    for r in g[:10]: print("  GARBAGE i=%d long=%s oov=%.3f sig=%.4f tokens=%s finish=%s req=%s head=%r" % (r["i"], r["long"], r["oov"], r["sig"], r["completion_tokens"], r["finish"], r["req_id"], r["text"][:90]))
    for r in leak[:5]: print("  LEAK i=%d contains markers of %s" % (r["i"], r["foreign_markers"]))

asyncio.run(main())

probe_sitecustomize.py

"""Race probe for vLLM's hybrid-Mamba accepted-token bookkeeping (loaded via PYTHONPATH).

Inert unless RACE_PROBE is set. RACE_PROBE=log records, per engine step that reorders
the batch, whether the previous step's non-blocking D2H copy of num_accepted_tokens had
already landed before the CPU row moves (early=True is the racy case) and which requests
were swapped. RACE_PROBE=fix additionally keeps num_accepted_tokens_cpu in the previous
step's row layout across swaps under async scheduling, which is the layout _prepare_inputs
indexes it by (prev_positions). Patches land lazily when gpu_model_runner is imported.
"""
import importlib.abc, importlib.util, json, os, sys, time

_MODE = os.environ.get("RACE_PROBE", "")
_TARGETS = {"vllm.v1.worker.gpu_model_runner", "vllm.v1.attention.backends.mamba2_attn"}
_state = {"step": 0, "swaps": [], "early": 0, "with_event": 0, "async": None, "log": None, "done": set()}


def _log(rec):
    if _state["log"] is None:
        _state["log"] = open(os.environ.get("RACE_PROBE_LOG", "/tmp/race_probe.log"), "a", buffering=1)
    rec["t"] = time.time(); rec["pid"] = os.getpid()
    _state["log"].write(json.dumps(rec) + "\n")


def _patch_batch(IB):
    orig_swap = IB.swap_states

    def swap_states(self, i1, i2):
        a = int(self.num_accepted_tokens_cpu[i1]); b = int(self.num_accepted_tokens_cpu[i2])
        _state["swaps"].append([i1, i2, self._req_ids[i1], self._req_ids[i2], a, b])
        orig_swap(self, i1, i2)
        if "fix" in _MODE and _state["async"]:
            self.num_accepted_tokens_cpu[i1] = a
            self.num_accepted_tokens_cpu[i2] = b

    IB.swap_states = swap_states


def _graft_pr_fix(mod):
    """RACE_PROBE=prfix: install the PR's calculate_reorder_batch_threshold and the
    requires_decode_ordering flags on the builder classes, exactly as the patch does."""
    import importlib, os.path
    ns = {}
    exec(open(os.path.join(os.path.dirname(__file__), "pr_runner_fix.py")).read(), ns)
    mod.GPUModelRunner.calculate_reorder_batch_threshold = ns["calculate_reorder_batch_threshold"]
    importlib.import_module("vllm.v1.attention.backend").AttentionMetadataBuilder.requires_decode_ordering = False
    for m, cls in (("vllm.v1.attention.backends.mamba_attn", "BaseMambaAttentionMetadataBuilder"),
                   ("vllm.v1.attention.backends.gdn_attn", "GDNAttentionMetadataBuilder"),
                   ("vllm.v1.attention.backends.linear_attn", "LinearAttentionMetadataBuilder")):
        try:
            setattr(getattr(importlib.import_module(m), cls), "requires_decode_ordering", True)
        except Exception as e:
            _log({"event": "prfix_flag_failed", "module": m, "error": repr(e)})
    _log({"event": "prfix_installed"})


def _patch_runner(mod):
    R = mod.GPUModelRunner
    if "prfix" in _MODE:
        _graft_pr_fix(mod)
    orig_us = R._update_states

    def _update_states(self, scheduler_output):
        if _state["async"] is None:
            _state["async"] = bool(getattr(self, "use_async_scheduling", False))
            try:
                _patch_batch(type(self.input_batch))  # InputBatch is fully loaded by now
                _log({"event": "patched_batch", "cls": type(self.input_batch).__name__})
            except Exception as e:
                _log({"event": "patch_failed", "module": "input_batch", "error": repr(e)})
            _log({"event": "start", "mode": _MODE, "async": _state["async"],
                  "mamba_cache_mode": getattr(self.cache_config, "mamba_cache_mode", None),
                  "runner_reorder_batch_threshold": getattr(self, "reorder_batch_threshold", None),
                  "builder_thresholds": {type(g.get_metadata_builder()).__name__: g.get_metadata_builder().reorder_batch_threshold for g in self._attn_group_iterator()} if hasattr(self, "_attn_group_iterator") else None})
        ev = getattr(self, "num_accepted_tokens_event", None)
        early = bool(ev.query()) if ev is not None else None
        _state["swaps"] = []
        out = orig_us(self, scheduler_output)
        _state["step"] += 1
        _state["req_ids"] = list(self.input_batch.req_ids)  # post-reorder batch order, for the prefill-geometry log
        if early is not None:
            _state["with_event"] += 1; _state["early"] += early
        if _state["swaps"]:
            _log({"event": "swap", "step": _state["step"], "early": early,
                  "num_reqs": self.input_batch.num_reqs, "swaps": _state["swaps"]})
        if _state["step"] % 1000 == 0:
            _log({"event": "tick", "step": _state["step"], "steps_with_event": _state["with_event"], "early_steps": _state["early"]})
        return out

    R._update_states = _update_states

    if "reorderfix" in _MODE and hasattr(R, "calculate_reorder_batch_threshold"):
        orig_calc = R.calculate_reorder_batch_threshold

        def calculate_reorder_batch_threshold(self):
            orig_calc(self)
            sc = getattr(self.vllm_config, "speculative_config", None)
            k = (sc.num_speculative_tokens or 0) if sc is not None else 0
            before = self.reorder_batch_threshold
            # the Mamba2 builders split decodes/prefills at 1 + k; make the batch reorder use the same threshold
            self.reorder_batch_threshold = max(before or 1, 1 + k)
            _log({"event": "reorderfix", "before": before, "after": self.reorder_batch_threshold})

        R.calculate_reorder_batch_threshold = calculate_reorder_batch_threshold


def _patch_mamba2_builder(mod):
    """Log the composition of every Mamba2 prefill step: which requests are in the prefill
    region, their already-computed tokens (0 = cold) and this step's chunk length."""
    B = mod.Mamba2AttentionMetadataBuilder
    orig_build = B.build

    def build(self, common_prefix_len, common_attn_metadata, fast_build=False, **kwargs):
        m = orig_build(self, common_prefix_len, common_attn_metadata, fast_build, **kwargs)
        try:
            if "geom" in _MODE and m.num_prefills > 0:
                cam = common_attn_metadata
                nr, npf = m.num_reqs, m.num_prefills
                qsl = cam.query_start_loc_cpu.tolist()
                seq = cam.seq_lens_cpu_upper_bound.tolist()
                rows = []
                ids = _state.get("req_ids") or []
                for r in range(nr - npf, nr):
                    qlen = qsl[r + 1] - qsl[r]
                    rows.append([ids[r] if r < len(ids) else None, seq[r] - qlen, qlen])
                _log({"event": "prefill", "step": _state["step"], "num_reqs": nr, "num_decode_tokens": m.num_decode_tokens,
                      "rows": rows})  # rows: [req_id, num_computed_before_this_step, tokens_this_step]
        except Exception as e:
            _log({"event": "geom_failed", "error": repr(e)})
        return m

    B.build = build


class _Finder(importlib.abc.MetaPathFinder):
    def find_spec(self, name, path, target=None):
        if name not in _TARGETS or name in _state["done"]:
            return None
        _state["done"].add(name)
        spec = importlib.util.find_spec(name)
        if spec is None or spec.loader is None:
            return None
        loader = spec.loader
        orig_exec = loader.exec_module

        def exec_module(module):
            orig_exec(module)
            try:
                (_patch_runner if name.endswith("gpu_model_runner") else _patch_mamba2_builder)(module)
                _log({"event": "patched", "module": name})
            except Exception as e:  # never take the engine down for a probe
                _log({"event": "patch_failed", "module": name, "error": repr(e)})

        loader.exec_module = exec_module
        return spec


if _MODE:
    sys.meta_path.insert(0, _Finder())

reanalyze.py

"""Augmented corruption detector + geometry correlation.
usage: reanalyze.py <burst.jsonl> [probe.jsonl]"""
import json, re, sys
from collections import defaultdict
rows = [json.loads(l) for l in open(sys.argv[1])]; ok = [r for r in rows if "text" in r]
def corrupt(r):
    words = len(re.findall(r"[A-Za-z]{2,}", r["text"])); tok = r["completion_tokens"] or 0
    empty = tok > 200 and words < 0.02 * tok          # thousands of tokens that render to (almost) no text
    return r["garbage"] or empty, ("salad" if r["garbage"] else "empty" if empty else "")
kinds = defaultdict(lambda: [0, 0, 0])
for r in ok:
    c, k = corrupt(r); kinds[r.get("kind", "?")][0] += 1; kinds[r.get("kind", "?")][1] += (k == "salad"); kinds[r.get("kind", "?")][2] += (k == "empty")
print("by kind (n, salad, empty):", dict(kinds))
if len(sys.argv) > 2:
    P = [json.loads(l) for l in open(sys.argv[2]) if '"prefill"' in l]
    first = {}
    for e in P: first.setdefault(e["step"], e)
    byreq = defaultdict(list)
    for e in first.values():
        for pos, (rid, c, n) in enumerate(e["rows"]):
            byreq[rid.rsplit("-", 1)[0]].append((e["step"], c, n, pos, len(e["rows"]), e["rows"][0][1], e["rows"][0][2]))
    stats = defaultdict(lambda: [0, 0, 0])
    detail = []
    for r in ok:
        evs = sorted(byreq.get(r["req_id"], [])); p = r["prompt_tokens"]
        misfirst = [e for e in evs if p <= e[1] <= p + 4 and e[2] <= 4]
        mislater = [e for e in evs if e[1] > p + 4 and e[2] <= 4]
        key = ("FIRST decode step in Mamba prefill region" if misfirst else "first decode step clean") + (" | later decode steps in prefill region" if mislater else "")
        c, k = corrupt(r); stats[key][0] += 1; stats[key][1] += (k == "salad"); stats[key][2] += (k == "empty")
        if misfirst or c: detail.append((r["i"], r["kind"], p, r["completion_tokens"], k or "ok", [(e[0], e[1], e[2], f"{e[3]}/{e[4]}", f"lead=({e[5]},{e[6]})") for e in misfirst[:2]]))
    for k, (n, s, e) in sorted(stats.items()): print(f"  {k:<75} n={n:4d} salad={s} empty={e}")
    print("requests with a misclassified first decode step, or corrupted:")
    for d in detail: print("  ", d)