[Bugfix] Bound accepted-token state lookups in GDN/KDA spec decode - #50021
[Bugfix] Bound accepted-token state lookups in GDN/KDA spec decode#50021amittell wants to merge 6 commits into
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: dd1fd0389c
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
Pull request overview
This PR hardens the hybrid GDN + MTP speculative decode path against CUDA illegal memory accesses by bounding two indices derived from num_accepted_tokens, ensuring out-of-range acceptance counts don’t turn into out-of-bounds reads and wild pointer dereferences in downstream Triton kernels.
Changes:
- Mask the
ssm_state_indicesload infused_recurrent_gated_delta_rule_fwd_kernelso invalidnum_accepted_tokensvalues fall into the existing “invalid state” early-return path. - Mask block-table column loads in
_copy_mamba_state_blockto prevent out-of-range accepted-token-derived columns from reading outside a request’s block-table row.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
vllm/v1/worker/mamba_utils.py |
Adds masked block-table column loads in _copy_mamba_state_block to prevent out-of-range accepted-token-derived indexing from producing illegal addresses during state copies. |
vllm/third_party/flash_linear_attention/ops/fused_recurrent.py |
Masks the initial-state index lookup derived from num_accepted_tokens - 1 to avoid out-of-bounds reads and subsequent invalid state dereferences. |
Comments suppressed due to low confidence (3)
vllm/v1/worker/mamba_utils.py:102
- Same sentinel issue as above: block ID 0 (
NULL_BLOCK_ID) should be treated as invalid. Otherwise, an in-range but unallocatedsrc_colcan cause reads from the reserved padding block 0.
src_block_id = tl.load(
block_table_base + src_col, mask=src_col_ok, other=-1
).to(tl.int64)
if src_block_id < 0:
return
vllm/v1/worker/mamba_utils.py:129
- Same sentinel issue as above: block ID 0 (
NULL_BLOCK_ID) should be treated as invalid to avoid copying from the reserved padding block.
src_block_id = tl.load(
block_table_base + src_col, mask=src_col_ok, other=-1
).to(tl.int64)
if src_block_id < 0:
return
vllm/v1/worker/mamba_utils.py:150
- Same sentinel issue as above in the temporal-state path: block ID 0 (
NULL_BLOCK_ID) should be treated as invalid. Otherwise an in-range but unallocatedtmp_colcan copy from the reserved padding block.
actual_src_block_id = tl.load(
block_table_base + tmp_col, mask=tmp_col_ok, other=-1
).to(tl.int64)
if actual_src_block_id < 0:
return
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
Thanks for the review — pushed Bound the sigmoid-gating kernel (Codex P1). Correct and important: the Qwen3.5/Qwen3.6 GDN decode path calls Reject Initialize outputs when rejecting (Codex P1). The The masking change was validated end-to-end on an RTX PRO 6000 (SM120, driver 610) with Qwen3.6-27B-NVFP4, MTP-3, prefix caching, fp8 KV, full PIECEWISE cudagraph: 0 GPU faults across a long-prompt sweep to 32k tokens, 0 cache poisoning, 149-158 t/s warm. |
Independent GPU red/green confirmationI independently exercised the exact current-source bounds changes on RTX 5090/SM120 with Qwen3.6 hybrid GDN + MTP n=3. Five deterministic GPU regressions failed before the bounded-index changes and passed with all three files mounted:
This is independent of the FlashInfer/XQA MTP routing issue in #49010: fixing the attention route restored long-context recall, but it does not remove this wild-address/adjacent-state class. I therefore consider this PR required for a production-safe hybrid-GDN MTP stack, not a workaround for the attention bug. |
|
The contributor-side checks and requested GPU evidence are complete. The remaining |
…bounded block-table load in the SD-conv branch)
|
This pull request has merge conflicts that must be resolved before it can be |
…ment stride The 71171a880 bound compared init_token_idx against stride(1) of the state-indices tensor -- 1 for a contiguous [batch, T] layout -- so every accepted count > 1 was wrongly failed closed (zero output, no state update), and a non-contiguous layout with stride(1) > stride(0) would have reopened the out-of-bounds lookup. The invalid-count regression could not catch this: an oversized count fails closed under both the wrong and the correct bound. Bound against stride(0) (elements per batch row; == T when contiguous), matching fused_recurrent.py's i_t < stride_indices_seq pattern, and add the discriminating valid-side regression: num_accepted=2 in a 3-column row must produce real output and a state update. RTX 5090 evidence: the new test fails on 71171a880 exactly as predicted (all-zero output) and passes here; the invalid-count test passes on both (demonstrating it was non-discriminating); the full tests/kernels/mamba/test_mamba_ssm.py file passes 288/288 (18 skipped). Found by the depthfirst review bot on the PR diff. Assisted-by: Claude (Anthropic) Signed-off-by: Alex Mittell <mittell@me.com>
a303cf1 to
9a198c0
Compare
…DA state indices Replaces the local vllm-project#48475-style clamp with the upstream bounds fix (PR vllm-project#50021, open): masked row-bounded loads for ssm_state_indices and block-table columns, fail-closed zero output for invalid counts, so stale/too-large accepted counts can no longer dereference wild state addresses (Xid 13/31). Clamps kept as defense in depth in FLA kernels.
…tale zero-accept rows Builder-level fix (open PR vllm-project#51508): rows whose async-scheduling step was discarded get their whole spec_state_indices row set to NULL_BLOCK_ID, so FLA/conv kernels skip them entirely (no initial-state read, no final-state write) instead of advancing state for a dead request. Kernel clamps retained as defense in depth (already applied via vllm-project#50021 port).
Drop the local vllm-project#48475-style clamps in fused_recurrent/fused_sigmoid_gating so a zero num_accepted_tokens falls through the row-masked load into the invalid-state path (zeroed output, state untouched) per PR vllm-project#50021, instead of silently reading slot 0. Adapt the PR vllm-project#51508 defense test that assumed clamp equivalence, and point the GDN builder fixture at the local Qwen3.5-0.8B checkpoint for offline runs.
|
We maintain a serving stack for hybrid-GDN Qwen models on consumer GPUs and have multiple field reports of this crash class. We built a fast reproducer and ran a controlled matrix against v0.27.1, including this PR's current head. Three results that seem important: 1. This PR's head does not stop the crash. We applied the five runtime-file changes at 2. The crash is gated on async scheduling — it looks like a cross-stream race, not an index bug. Matched single-variable A/B, identical client and config (MTP n=4,
Drafter-off and prefix-caching-off arms also never fault. Our reading: step N+1's metadata build consumes 3. Separately, we found a second, distinct fault while isolating this one — crossing sequence position 32,768 with the MTP drafter active permanently kills draft acceptance engine-wide (reproduces with async scheduling disabled, so it is not this PR's bug). Filed with full evidence and reproducer as #52873. Reproducer below (stdlib-only; point at an OpenAI-compatible endpoint; Reproducer#!/usr/bin/env python3
"""Fast reproducer: CUDA illegal memory access (Xid 31) on Qwen3-Next-family
hybrid GDN + MTP spec decode + prefix caching, vLLM v0.27.1.
Shape: a single growing multi-turn conversation (agent-style). Every turn is a
prefix-cache-hit resume with speculative tokens scheduled. On 2x RTX 3090 (TP=2,
fp8 KV, MTP n=4) the engine dies with `gdn_attn.py build -> illegal memory access`
within 6k-13k generated tokens (~5-10 min), reproduced 4/4 runs. With the drafter
off (SPEC_N=0) the same traffic runs indefinitely.
Usage: python3 gdn-mtp-apc-repro.py [base_url] [model]
"""
import json, sys, time, urllib.request
BASE = sys.argv[1] if len(sys.argv) > 1 else "http://localhost:8095"
MODEL = sys.argv[2] if len(sys.argv) > 2 else "qwen3.8-27b"
SYS = ("You are a senior systems engineer pair-programming with the user on a large "
"C++ and Python codebase for GPU inference. Be concrete and show code. ") * 40
TOPICS = [
"Refactor this CUDA kernel launcher to support streams. Write the full code.",
"Design a block-paged KV cache allocator; write the C++ header with comments.",
"Write a Python asyncio scheduler for batched LLM requests with timeouts.",
"Explain then implement speculative-decoding token verification in PyTorch.",
"Write unit tests for a ring-buffer class; include edge cases and comments.",
"Port this concept to Triton: fused RMSNorm + residual add. Full kernel please.",
"Debug: our TP=2 all-reduce hangs on PCIe. List hypotheses, then a bisect plan.",
"Write a bash script that soaks an OpenAI endpoint and logs TPS per request.",
]
def chat(messages, max_tokens=1000):
req = urllib.request.Request(BASE + "/v1/chat/completions",
data=json.dumps({"model": MODEL, "messages": messages, "max_tokens": max_tokens,
"temperature": 0.7, "top_p": 0.95,
"chat_template_kwargs": {"enable_thinking": False}}).encode(),
headers={"Content-Type": "application/json"})
with urllib.request.urlopen(req, timeout=600) as r:
out = json.loads(r.read().decode())
ch = out["choices"][0]["message"]["content"] or ""
return ch, out.get("usage", {}).get("completion_tokens", 0)
msgs = [{"role": "system", "content": SYS}]
total, turns = 0, 0
while total < 220_000:
msgs.append({"role": "user", "content": f"[turn {turns}] {TOPICS[turns % len(TOPICS)]} "
f"Focus on a module named engine_{turns} with a {5 + turns % 7}-stage pipeline."})
try:
content, comp = chat(msgs)
except Exception as e:
print(f"FAILURE after {total} generated tokens: {e}")
print("check: docker logs <container> | grep -m5 'illegal memory' ; dmesg | grep Xid")
sys.exit(2)
if len(content.strip()) < 20:
msgs.pop(); turns += 1; continue # keep the history well-formed
msgs.append({"role": "assistant", "content": content})
total += comp; turns += 1
if turns % 10 == 0:
print(f"generated={total} turns={turns}", flush=True)
print("no failure up to 220k generated tokens") |
|
Thank you — the matched single-variable A/B plus the The accepted-count sync is conditional, and prefix caching turns it off
# Sync num_accepted_tokens from CPU (set by
# _update_states_after_model_execute for hybrid models).
# Skipped under async scheduling (non-align): the CPU copy races with
# the in-flight D2H copy and with input-batch row moves.
needs_cpu_accepted_counts = self.num_accepted_tokens_event is not None and not (
self.use_async_scheduling and self.cache_config.mamba_cache_mode != "align"
)
if needs_cpu_accepted_counts:
assert self.num_accepted_tokens_event is not None
self.num_accepted_tokens_event.synchronize()
Falsifiable in one flag:
|
| config | result |
|---|---|
Qwen3.8-27B hybrid GDN, MTP-3, APC, --mamba-cache-mode align, async scheduling auto-on, TP=1, v0.26.0 |
16,918 generated tokens, 0 faults, engine healthy — straight through the 6.4k–12.9k window where you died 5/5 |
| Same recipe in production (NVFP4, RTX 6000 Blackwell, TP=1, align, async on) | no illegal memory access in the unit's journal history; no Xid in dmesg on that host |
Caveat I will flag myself: both of our arms are TP=1 on Blackwell, so this is consistent with the gate being the variable, not proof of it — TP=2 rank synchronization is its own timing regime, and your 2×3090 result may well have a second contributing factor. I have a non-align arm (mamba_cache_mode="all", everything else identical, same node class) staging now and will post it either way; if it faults where the align arm did not, the gate is confirmed independently of arch and TP.
Why this PR's head did not help you
The PR was developed and validated entirely in align mode with --enforce-eager — see the repro line and the test plan in the description. Both sidestep the async race: align keeps the sync, and eager mode removes the CUDA-graph stream structure. So it targets a different defect — accepted-count-derived indices that are correctly synchronized but still out of range (zero count → i_t == -1 reads before the row; stale or oversized count reads past it), which is deterministic and reproduces with async scheduling off. Its negative control (test_invalid_block_table_lookup_does_not_copy_state) still fails 4/4 against pristine main for us, so the defect it bounds is live independent of yours.
So I read your result as "there are at least two bugs here", not "the bounds are wrong". That said, your run shows how easily the PR's scope reads as broader than it is, so I will make the description state the align + eager validation scope explicitly rather than leaving align as a parenthetical, and reference this comment for the async/non-align path.
What a fix for your half probably looks like
Not more bounds: the GPU-side consumer needs to be ordered against the producer. Either the metadata build waits on num_accepted_tokens_event unconditionally (the current gate makes the wait conditional on a cache mode that has nothing to do with whether the GPU tensor is safe to read), or the L326/L460 reads move onto the stream that wrote them so the ordering is implicit rather than host-mediated. Happy to test a candidate patch of either shape.
On #52873
The position-32,768 acceptance collapse reproducing with async scheduling disabled does read as independent, and a hard boundary at exactly 2^15 is a strong hint on its own. We have hybrid-GDN capacity across Blackwell (RTX 6000, GB10, 5090) and can run that as a deterministic boundary sweep if a second data point on different silicon would help.
|
Correction to my previous comment, and it sharpens the predicate rather than weakening it. I wrote that if cache_config.enable_prefix_caching:
if cache_config.mamba_cache_mode == "none":
cache_config.mamba_cache_mode = (
"all" if model_config.supports_mamba_prefix_caching else "align"
)So with prefix caching on and no explicit flag, a model that supports mamba prefix caching gets How I caught it: I claimed our fleet was a clean data point at the same config as yours. It isn't. Our Qwen3.8-27B logs i.e. it auto-selects The upshot for your report is a tighter exposure predicate than I gave:
which fits Qwen3-Next-family + APC landing in
I have the genuine non-align arm running now (same node, same model, everything identical, but |
|
Non-align arm, as promised — and it did not reproduce, which is evidence against the gate being sufficient on its own. Reporting it because I said I would either way. Single variable against my earlier align run: same node, same model, same flags, but
So on Qwen3.8-27B, TP=1, single Blackwell (GB10), v0.26.0, MTP-3, APC, your reproducer verbatim: the missing synchronize is not sufficient to produce the fault. That is a real strike against my "this gate is the mechanism" framing, and you should weight it accordingly. What I think survives, and what does not:
Given that, the
Honest limits on our side: we cannot currently match your rig. Our Blackwell pair is both GPUs in production, and the GB10s are one GPU per node, so a TP=2 arm would need multi-node Ray rather than a same-box tensor-parallel split — not a like-for-like substitute for two 3090s on one host. If you can run the one-flag |
|
This is the missing piece, and it resolves your arm B cleanly: our config runs Qwen3.8-27B reports Which means, for our failure, Where "too late" isIn The fix we shipped, and what it buysTwo lines: hoist the wait to the top of @contextmanager
def synchronize_input_prep(self):
if self.prepare_inputs_event is None:
yield
return
self.prepare_inputs_event.synchronize()
# order this step's _update_states/_prepare_inputs after the previous
# step's spec-decode postprocess, which records this event
if self.num_accepted_tokens_event is not None:
self.num_accepted_tokens_event.synchronize()
try:
yield
finally:
self.prepare_inputs_event.record()This is your "wait unconditionally" shape, with the extra point that in Result on the rig you can't reproduce on (2×3090, TP=2, Ampere, v0.27.1, MTP n=4, fp8 KV, APC, async on): 3/3 boots clean to 33–36k generated tokens each, straight through the 6.4–13k window where we died 5/5 unpatched, across the ctx-32k crossing and a compaction; bench 79.1 narr / 108.5 code / 1756 prefill tok/s, i.e. ≥ baseline. We've shipped it as a downstream install-script overlay on our stack. Honest limitsWe have not isolated TP=2 as a necessary co-factor — the early-sync fix closes the window regardless of why the window is wide enough to fire, so "align-late-sync race" and "TP=2 timing widens it" are not separated by our data. Your arm B says the skip branch alone doesn't fault on TP=1/Blackwell; our result says the align branch does fault on TP=2/Ampere and the early hoist fixes it. Both can be true. Update — we tried the TP=1 arm to isolate the collective, and it's memory-blocked on our hardware. The 18.2 GiB AutoRound-INT4 weights plus fixed cudagraph/GDN-state scratch don't fit one 24 GB Ampere card: OOM at an identical 2.37 GiB shortfall across On #52873Yes please to the Blackwell boundary sweep — a second-silicon data point is exactly what that one needs. Note the boundary is not a fixed position: three async-on boots here collapsed at ctx ~21,025 while an async-off boot collapsed at ~32,570, reproducible within a config but shifting with it (details on #52873). So a sweep that brackets a fixed 2^15 may miss it; varying the config (async on/off, TP, depth) and watching where the per-window acceptance flips to 0 is the more informative shape. |
Follow-up stress validation with the accepted-token bounds appliedFollowing my earlier red/green confirmation, I kept the relevant #50021 hunks in the final MTP-3 integration stack and ran a materially larger async/concurrent workload. This is additional soak evidence, not a new isolated A/B. Environment: RTX 5090, Qwen3.8-27B NVFP4, NVFP4 KV, GDN/MTP-3, async scheduling, up to eight sequences.
The longer soak adds coverage for mixed prompt lengths, async structured output, prefix caching, and divergent agent histories on top of the original targeted red/green result. Disclosure: AI-assisted analysis and comment posting; the runs and measurements were produced and verified by me on the hardware described. |
…re-pads an already-padded layer to the shared page (boot AssertionError class); commensurate per-token-bytes zero guard; AR speculator draft-prefill dispatches on num_tokens_padded (restores upstream vllm-project#47352, reverted by carried vllm-project#48244 pick); gemma4_dspark + laguna_dflash fused-KV dtype derived from norm weight + bias cast (fp16 drafters); unpadded() propagates mm_req_doc_ranges; V1 runner zeroes padded-row num_prompt_tokens_cpu (stale-length dummy-row class); fused_recurrent + fused_sigmoid_gating INPLACE_FINAL_STATE load masked to the row (SM-fault class, mirrors vllm-project#50021's sibling bound); profile_cudagraph_memory empty-sample guard; TQ spec head_size_v; packed-codec (kvarn_/turboquant_) branch in _validate_cache_dtype logging; max_page_block_lcm includes padded-to-max layers (PR vllm-project#52804 intent) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Thanks - that closes the loop on arm B cleanly: we were never in the same branch, and "the sync runs too late" fits the code as far as I can verify it. Before upstreaming your hoist I ran the duplicate-work check, and there is an open PR already aimed at this hazard: #51599 (closes #51571, njhill reviewing). It takes the other route: Evidence from our side:
So I am not opening a hoist PR while #51599 is open - same defect, and the buffer-isolation shape is the better mechanism. The deciding experiment is yours if you are willing: your ~10-minute reproducer against #51599's head on the 2x3090 rig. Crash gone -> #51599 is the upstream fix and I will link your field data on that PR. Crash still there -> the hoist covers something isolation does not, and I will upstream it with you credited for the diagnosis, the fix, and the evidence. Dropping the planned Blackwell sweep for #52873 since you closed it as checkpoint-specific. @seanyourhighness thanks for the extended async soak on the bounds half. Determinized D2H repro (RTX 5090)"""Determinized demonstration of the align-mode accepted-count hazard.
Main's shape: postprocess_mamba_align_gpu D2H-copies accepted counts into
input_batch.num_accepted_tokens_cpu_tensor (pinned, condense()-managed).
Under async scheduling, _update_states/condense() compacts rows of that
same buffer BEFORE num_accepted_tokens_event is synchronized, so host
writes and the in-flight DMA interleave on one buffer. torch.cuda._sleep
delays the DMA so each ordering is observable deterministically instead of
depending on bus timing.
Scenario: prev step had 4 requests [counts 4,3,2,1]; request at row 0
finished; condense moves the last row (3) into row 0 and the batch shrinks
to 3. A new request is admitted at row 3 with init count 1.
"""
import torch
assert torch.cuda.is_available()
counts_gpu = torch.tensor([4, 3, 2, 1], dtype=torch.int32, device="cuda")
SLEEP_CYCLES = 400_000_000
def condense_and_admit(buf):
buf[0] = buf[3] # condense: move last row into the freed slot
buf[3] = 1 # _update_states: init a newly admitted request's row
def run(isolated: bool):
batch_cpu = torch.tensor([4, 3, 2, 1], dtype=torch.int32).pin_memory()
runner_cpu = torch.zeros(4, dtype=torch.int32).pin_memory()
ev = torch.cuda.Event()
torch.cuda._sleep(SLEEP_CYCLES) # step N still on the GPU
target = runner_cpu if isolated else batch_cpu
target.copy_(counts_gpu, non_blocking=True) # postprocess's D2H
ev.record()
condense_and_admit(batch_cpu) # host, BEFORE the sync
ev.synchronize() # _prepare_inputs' wait
return batch_cpu.tolist(), runner_cpu.tolist()
main_batch, _ = run(isolated=False)
iso_batch, iso_runner = run(isolated=True)
want_batch = [1, 3, 2, 1] # row0 = surviving req's count, row3 = new req init
print(f"host wrote after condense+admit: {want_batch}")
print(f"main shape -> input_batch after sync: {main_batch}"
f" {'HOST WRITES DESTROYED BY LATE DMA' if main_batch != want_batch else 'ok'}")
print(f"#51599 shape -> input_batch after sync: {iso_batch}"
f" runner buf: {iso_runner}"
f" {'ok - both coherent' if iso_batch == want_batch and iso_runner == [4, 3, 2, 1] else 'CORRUPTED'}")Output: |
|
Ran your deciding experiment on the 2×3090 (TP=2, Ampere sm_86, PCIe-only, vLLM v0.27.1, fp8 KV,
The crash site is not the buffer #51599 isolates. Both the stock and #51599 arms die here: #51599 retargets the One caveat on the checkpoint, since it matters for interpreting this. We can no longer reproduce this crash on the checkpoint we actually ship. Our fast tier moved from an early Intel-AutoRound INT4 export ("Avuja") to a re-export ("Frozenlock") — same quant method/bits/group-size, the one difference being that the earlier export excluded
Happy to share the full per-arm logs, or to re-run against a revised #51599 that also isolates the |
|
That settles it, and it goes against the call I made. #51599 does not cover your crash, so the hoist is upstreamed as #53613 with you as co-author, credited for the diagnosis, the fix and the field evidence. Thanks for running three arms instead of two - the #51599-only arm is the one that makes this decidable. One correction, because a reviewer will make it otherwise. An illegal memory access is asynchronous and sticky: the frame it surfaces in is the first CUDA call that checks the error, not the kernel that raised it. What the code does support is your earlier framing, and that is what the PR says:
On the revised-#51599 offer: I would not go that way, and you should not have to. Those input-prep copies are H2D out of buffers allocated per call, so there is no shared residency to move - the defect is ordering, not destination. #51599 should still land on its own merits: the host-side condense corruption it fixes (#51571) is real and independent, and my determinized Two asks. Attach your per-arm logs to the new PR so the field evidence sits with the change rather than three issues away. And when you get a window, run the branch head - it is the same two lines you shipped plus a comment and four regression tests, but I would rather cite a run of the exact tree than a reconstruction of it. One number from our rig you will want, because it reframes the perf argument for both of us. I timed both accepted-count waits directly. On main the existing wait at The retired-export caveat is in the PR as you stated it. It bears on how easy the race is to hit, not on whether it is one; an unordered write-after-read does not become ordered because a checkpoint stopped exposing it. |
|
Version applicability, since this has been open across two releases and it decides how anyone backports it. Blob shas per ref (
Four of the seven are identical at v0.27.1, v0.28.0 and Three do move: the AMD KDA kernel changed between v0.27.1 and v0.28.0, and Ready for review whenever a maintainer has bandwidth. Every label on this PR so far is from |
Model evaluationGSM8K, 1319 questions, 5-shot, temperature 0, seed 42, AI assistance was used for this evaluation. Every number below is copied from the run's JSON Why this configurationThis PR only changes GDN/KDA/Mamba spec-decode kernels, so a GSM8K run without speculative Both were verified from the server log of every run, not assumed: 83 Hardware and software
Serve command (identical for every run):
Deviation worth recording: flashinfer had to be moved from the branch's pinned MethodBoth arms ran on the same install, toggling only the patch with The server was restarted for every single run, so no run inherits another run's prefix cache. The Results - 1319 questions, client concurrency 32
Per-arm spread:
Do the arms differ beyond noise? No.Pooling the two runs per arm, patched is 2216/2638 = 0.840030 and stock is 2229/2638 = 0.844958.
For scale, the binomial standard error of a single 1319-question run at p = 0.84 is 1.009 pp, Throughput is likewise indistinguishable: patched is +0.25% on output tok/s, against a per-arm Conclusion: no accuracy regression, no throughput regression, no change in speculative Separate finding, present in BOTH arms, not caused or fixed by this PRWhile selecting the eval configuration I hit a hard output-quality cliff as client concurrency Sweep of 200 questions on one server boot per arm:
At concurrency 64 the "invalid" responses are not malformed answers - they are requests that A full 1319-question matrix at concurrency 128 reproduces it at scale, again equally in both arms:
patched mean 0.567854 (spread 0.009098), stock mean 0.569371 (spread 0.012130), difference The empty-output set is not deterministic. Comparing the captured raw completions of two What I could not do
Reproducegit diff $(git merge-base pr50021 main) pr50021 > /tmp/pr50021.patch # 693 lines
VLLM_USE_PRECOMPILED=1 uv pip install -e . --torch-backend=auto
uv pip install --extra-index-url https://flashinfer.ai/whl/ flashinfer-python==0.6.17 flashinfer-cubin==0.6.17
# boot the server with the flags above, then per arm:
.venv/bin/python -c "
import sys; sys.path.insert(0, 'tests/evals/gsm8k')
from gsm8k_eval import evaluate_gsm8k
print(evaluate_gsm8k(num_questions=1319, num_shots=5, max_tokens=256,
temperature=0.0, seed=42, max_concurrency=32,
request_timeout_seconds=7200))"
# toggle arms with: git apply -R /tmp/pr50021.patch (and git checkout -f pr50021 to restore)
|
Summary
On a hybrid GDN model (Qwen3.5 / Qwen3.6) with MTP speculative decoding and prefix caching (
--mamba-cache-mode align), the engine dies withCUDA error: unspecified launch failurewithin 7-10 requests of agent-shaped traffic. The GPU faults, not the runtime: the signature is an SM address exception (Xid 13,ESR 0x404000) or an MMU fault (Xid 31), depending on whether the wild address happens to be mapped.This is the crash half of the problems reported around hybrid-Mamba + MTP. It is distinct from the prefix-cache corruption in #43559 (the coordinator-level EAGLE cache-peek gating for Mamba), which is already handled on current
main. This PR does not touch that path and does not claim to fix #43559; it fixes a separate GPU fault in two kernels downstream.Root cause: unchecked accepted-count-derived indices
Speculative decoding produces a per-request accepted-token count. Multiple GPU state consumers turn that count into an array index without bounding it, so a count that is stale, zero, or too large indexes outside its tensor and yields a wild address the GPU then dereferences.
Site 1
fused_recurrent_gated_delta_rule_fwd_kernel(fused_recurrent.py):i_t(= count - 1) is unbounded against astride_indices_seq-column tensor. A zero accepted count givesi_t == -1, a read before this request's row (before the tensor fori_n == 0); a stale or too-large count reads past the row. Thestate_idx <= 0guard only rejects non-positive values, so an out-of-range read that returns a garbage positive int flows into the address math and faults the SM.Site 2
_copy_mamba_state_block(mamba_utils.py): the block-table columns derive from the same count and index the per-request block-table row unbounded; the loaded block id becomesstate_base_addr + block_id * state_block_stride, which is then read and written.The two sites fire under different loads. A light decode load exercises only Site 1; a heavy cache-transition load (long prefills, A->B->A prefix reuse) also drives Site 2. Both must be bounded.
The fix
Both loads are masked to the valid range, so an out-of-range index falls into the existing invalid-state path instead of producing an address. No stream-ordering change, no device sync, no measurable throughput cost.
Test plan
RTX 5090 (sm_120), Qwen3.6-27B NVFP4, TP=1,
--enforce-eager --max-model-len 16384, MTP 3 (qwen3_5_mtp), prefix caching on,alignmode.Xid 31)The A->B->A probe issues an agent-shaped sequence (long prefill, prefix reuse, 20 tool schemas) and reports whether a poisoned prefix reappears; 68/68 returned a clean verdict, which also confirms the out-of-range early-return does not drop a needed state copy. Warm throughput and MTP acceptance length (3.98-4.00 of 4) are unchanged.
Note: a separate, still-open livelock (#49203)
Independently of this crash, the same stack can occasionally hang: engine alive,
/v1/modelsanswering, but the in-flight request stuck at 0% GPU util with no Xid. It is rare and timing-variable (seen once, then not across the 68 probes here) and matches open issue #49203. This PR does not address it and does not claim to; the 68/68 clean-verdict result rules out this change as a cause.Follow-up bounds audit
A follow-up audit expanded the same fail-closed rule to the remaining consumers in this path:
Validation after this follow-up: full pre-commit passed; on an RTX 5090, 186 FLA/causal-conv kernel tests and all 21 fused Mamba postprocess tests passed (207 total).
Kimi K3 KDA expansion
The same accepted-count-derived state selection pattern also existed in Kimi K3 KDA fused recurrent decode, in both the NVIDIA and AMD vendored kernels. This PR now masks that initial state-index load to the request row, zeroes all invalid-count output tokens, preserves state on invalid counts, and releases NVIDIA PDL dependents before the new empty/invalid early returns.
Additional validation on an RTX 5090:
num_accepted0 and 4).e7f66b199with only the new test added: 4/4 failed on old KDA for NVIDIA and AMD, proving the regression is non-vacuous.AI assistance disclosure
OpenAI Codex assisted with the follow-up bounds audit, implementation, and regression-test generation. The submitter owns the conclusions and final review.
Mamba2 selective-state expansion
The Mamba2 selective-state-update kernel also derives its initial state-slot lookup from
num_accepted_tokens - 1. It already clamps the lower side, but had no upper row bound. This update preserves that zero-count behavior and makes an oversized count fail closed: it writes zero output and returns before reading or writing state.RTX 5090 red/green regression: a count one past a three-column state row made the prior source select an adjacent row and emit nonzero output; the new kernel passes by emitting zero output and preserving the complete state tensor. Pre-commit passes for both changed files.