[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 |
|
Pushed RTX 5090 red/green proof: with a three-column state row and This intentionally does not duplicate #51508: that PR handles stale zero-count state semantics at the metadata layer; this is the missing upper-row bound in the Mamba2 kernel. |
![depthfirst-app[bot]](https://avatars.githubusercontent.com/in/1021002?s=60&v=4){kind=small}
|
This pull request has merge conflicts that must be resolved before it can be |
…e path On a hybrid GDN model (Qwen3.5/Qwen3.6) with MTP speculative decoding and prefix caching (--mamba-cache-mode align), the engine dies with "CUDA error: unspecified launch failure" within 7-10 requests of agent-shaped traffic. The GPU faults (Xid 13 SM address exception, or Xid 31 MMU fault), not the runtime. Speculative decoding produces a per-request accepted-token count. Two kernels turn that count into an array index without bounding it, so a stale, zero, or too-large count indexes outside its tensor and yields a wild address that the GPU then dereferences. Site 1, fused_recurrent_gated_delta_rule_fwd_kernel: i_t (num_accepted_tokens minus 1) is unbounded against a stride_indices_seq-column tensor; i_t == -1 reads before the row. The existing state_idx <= 0 guard still passes a garbage positive index into h0 + state_idx * stride_init_state_token. Site 2, _copy_mamba_state_block: the block-table columns derive from the same count and index the per-request block-table row unbounded; the loaded block id becomes state_base_addr + block_id * state_block_stride, then read and written. 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. Signed-off-by: Alex Mittell <mittell@me.com>
Address review feedback: - fused_sigmoid_gating.py: the Qwen3.5/Qwen3.6 GDN decode path calls fused_sigmoid_gating_delta_rule_update, not fused_recurrent_gated_delta_rule, and its kernel has the identical unbounded ssm_state_indices load. Apply the same row-bounded mask so the Qwen path is actually covered. - mamba_utils.py: NULL_BLOCK_ID is 0, not negative. The out-of-range early return used ``< 0`` (only the -1 sentinel), so a stale-but-in-range column landing on an unallocated (0) slot would still copy block 0. Change all four block-id guards to ``<= 0`` so both the sentinel and NULL_BLOCK_ID are rejected, matching the FLA kernels' own ``state_idx <= 0`` convention. Fix the comment that wrongly described 0 as negative. Signed-off-by: Alex Mittell <mittell@me.com> Signed-off-by: Alex Mittell <amittell@gmail.com>
Zero rejected kernel output deterministically, fail closed before causal-conv state address math, and exercise block-table row bounds with GPU regressions. Assisted-by: OpenAI Codex Signed-off-by: Alex Mittell <mittell@me.com>
Mask Kimi K3 KDA's accepted-token-derived state-index load to the request row for both vendored NVIDIA and AMD recurrent kernels. Invalid accepted counts now fail closed by zeroing output and preserving state instead of selecting an adjacent state row. Add a shared KDA regression that covers zero and too-large accepted counts for both implementations. The test keeps old out-of-row reads inside allocated storage so failures are deterministic semantic regressions. Assisted-by: OpenAI Codex Signed-off-by: Alex Mittell <mittell@me.com>
Assisted-by: OpenAI Codex Signed-off-by: Alex Mittell <mittell@me.com>
…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. |
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 mini-beast 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.