P72 profile-run cap couples qwen3_next GDN buffer size → runtime batch bounded by GENESIS_PROFILE_RUN_CAP_M, not --max-num-batched-tokens (Ada SM89, k8v4)

[jsboige]

Summary. On v7.72.5 the P72 profile-run cap (GENESIS_PROFILE_RUN_CAP_M) lets the engine boot at --max-num-batched-tokens > 4096, but it also sizes the qwen3_next (GatedDeltaNet) compiled buffers at the cap value. So the runtime batch is bounded by GENESIS_PROFILE_RUN_CAP_M, not by --max-num-batched-tokens: any single combined forward whose token count lands in (cap, max_num_batched_tokens] overflows the GDN buffer and kills the worker. With the cap at 4096, 4096 is the effective runtime ceiling. First — thank you for the patch tree; it's the only thing that makes TurboQuant k8v4 + hybrid Qwen3.6 work on our hardware at all.

Environment

• Image: vllm-qwen36-genesis-tq:v7.72.5-vllm01d4d1ad3 (your tree on nightly 01d4d1ad3)
• Hardware: 2× RTX 4090 (Ada, SM89), TP=2 + EP=2
• Model: Qwen3.6-35B-A3B AWQ 4-bit, --kv-cache-dtype turboquant_k8v4, --max-model-len 262144, gpu-util 0.82
• KV cache ~2.0M tokens, full CUDA graphs, vision on. Stable for 7+ days at --max-num-batched-tokens 4096.

What we tried. To lift throughput we set --max-num-batched-tokens 8192 with P72 (GENESIS_PROFILE_RUN_CAP_M=4096) + P74 (chunk-clamp). Two distinct failures, depending on the cap:

(A) Boot crash — P72 OFF, batch ramped to 32768. Dynamo fake-tensor mismatch in determine_available_memory()'s profiling forward:

mul(FakeTensor(..., 65536, 128), FakeTensor(..., 16*s59, 128))   # 65536 = 2 × max_num_batched_tokens

The TQ-k8v4 GQA/multi-query path (P40 / P67) can't specialize that symbolic shape at large batch → profiling forward never completes. This is why P72 exists.

(B) Runtime crash — P72 ON at 4096, batch 8192. Engine boots clean, a short bench passes, then ~1h25 into the deploy under real mixed traffic the worker dies:

RuntimeError: setStorage: sizes [5536, 16, 128], strides [2048, 128, 1], storage offset 0,
and itemsize 2 requiring a storage size of 22675456 are out of bounds for storage of size 16777216
                                                                    # qwen3_next.py:495

The math is decisive:

• 16777216 = 4096 × 16 × 128 × 2 — GDN buffer sized for M = the P72 cap (4096)
• 22675456 = 5536 × 16 × 128 × 2 — a real combined forward batch of 5536 tokens

5536 > 4096 → overflow, via the inductor / cuda-graph as_strided aliasing path (gen_alias_from_base).

Why our pre-deploy bench missed it (a caution for others). An 80K single-prefill "crash test" is P74-chunked to ≤4096 per forward, and N=16 concurrent decode is ~16 tokens/step — neither ever produces a single forward in (4096, 8192]. It takes co-scheduled long-context prefill + decode in the same step to sum one forward past the cap. A valid test must force a single forward strictly in (cap, max_num_batched_tokens] (e.g. several concurrent ~2-3K-token prefills landing in one step).

The coupling, stated plainly. GENESIS_PROFILE_RUN_CAP_M does double duty: it bounds the memory-profiling shape (good — dodges crash A) and it sizes the qwen3_next GDN compiled buffers (bad — bounds the safe runtime batch). Raising both the cap and forcing a fresh compile to 8192 would re-expose crash A, so it's not a workaround we're willing to run on prod.

Question / ask. Is there a path to decouple the GDN buffer sizing from the profile cap — i.e. size the qwen3_next compiled buffers from --max-num-batched-tokens (or max_num_seqs × something) while keeping the determine_available_memory() profiling shape capped to dodge crash A? Or, alternatively, a large-batch profile route that lets the TQ-k8v4 path specialize at the real batch so P72 isn't needed at all? If 4096 is fundamental to the k8v4 GDN geometry on this model, that's a useful thing to document too.

Happy to test a candidate patch on Ada SM89 in a prod-adjacent setup and report back with the same crash-reproducing traffic shape. I can also provide fuller logs / the compile_ranges and boot summary if useful.

[Sandermage]

Thank you — this is one of the most precise bug reports we've received. The arithmetic (16777216 = 4096·16·128·2 vs 22675456 = 5536·16·128·2), the (cap, max_num_batched_tokens] band analysis, and the note about why a plain 80K prefill or N=16 decode never reproduces it are all exactly right. That analysis let us pin the root cause precisely.

What we found

Your diagnosis of the symptom is correct, but the buffer that overflows is not sized by P72, and it's not the pristine vLLM allocation either. It's one of our own patches: P28 (GDN core_attn_out prealloc).

• Pristine vLLM (then qwen3_next.py, now model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py) allocates the GDN output buffer per forward from the real token count: core_attn_out = torch.zeros((num_tokens, num_v_heads//tp, head_v_dim), ...) where num_tokens = hidden_states.size(0). That self-sizes and cannot overflow.
• P28 (a legacy, default-on Genesis patch) replaces that per-forward torch.zeros with a layer-persistent pre-allocated buffer sized at a resolved token budget, sliced [:num_tokens] for reuse — a throughput optimization (it removes a per-step zeros kernel on the GDN decode hot path).
• On v7.72.5, P28's budget was hardcoded to 4096 and never consulted the live scheduler config. So at --max-num-batched-tokens 8192, the buffer stayed 4096 rows; a co-scheduled prefill+decode forward of 5536 tokens did buf[:5536] on a 4096-row buffer, which silently returns 4096 rows, and Inductor's as_strided / gen_alias_from_base alias regeneration then died with your setStorage … out of bounds.

So GENESIS_PROFILE_RUN_CAP_M did indeed look like the culprit (both happened to be 4096), but the coupling was P28's budget defaulting to the same 4096, not P72 reaching into the GDN buffer. P72 only ever patches the determine_available_memory() profiling forward (_dummy_run(max_num_tokens, is_profile=True) in gpu_model_runner.py) — that's confirmed unchanged.

The fix (already in our current tree — vLLM 0.23.1rc1.dev424+g3f5a1e173)

The two roles are now cleanly decoupled, on two independent layers:

1. Budget decoupled from the profile cap. P28's prealloc budget now resolves via a central resolver (prealloc_budget.resolve_token_budget) whose priority chain is:
   explicit hint → GENESIS_PREALLOC_TOKEN_BUDGET → GENESIS_GDN_MAX_BATCHED_TOKENS → vllm scheduler_config.max_num_batched_tokens → 4096 default.
   So with --max-num-batched-tokens 8192 and no override, the GDN buffer is sized 8192, following your batch budget — not GENESIS_PROFILE_RUN_CAP_M. The profiling forward stays capped (P72 unchanged), so crash A is still dodged; the GDN buffer no longer inherits that cap.

2. In-graph capacity guard (belt-and-suspenders). The patched forward_cuda now only reuses the prealloc when it's actually big enough, and otherwise falls back to a fresh full-size torch.zeros:

   core_attn_out = (
       self._genesis_gdn_core_attn_buf[:num_tokens].zero_()
       if getattr(self, '_genesis_gdn_core_attn_buf', None) is not None
       and self._genesis_gdn_core_attn_buf.shape[0] >= num_tokens
       else torch.zeros(
           (num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim),
           dtype=hidden_states.dtype, device=hidden_states.device,
       )
   )

   Both branches are pure tensor ops (dynamo-safe; the is not None / >= resolve at trace time). Even if the budget were somehow undersized, an over-capacity forward can no longer alias past the prealloc storage — it just self-sizes, exactly like pristine vLLM. This is the direct answer to your "decouple the GDN buffer sizing from the profile cap" ask.

To be clear on your alternative question: 4096 is not fundamental to the k8v4 GDN geometry. It was only ever the conservative default of a budget resolver, not a kernel constraint.

What we verified (and the honest limits of it)

On our hardware — 2× A5000 (sm_86), not your Ada sm_89 — running the same vLLM dev424 pin and the same model family (Qwen3.6-35B-A3B, hybrid GDN+Mamba, TQ k8v4, MTP):

• The fix is live and applied on PROD (P28 v2 with the capacity guard confirmed in the running container; budget resolver routes through scheduler_config). Worker healthy 20h, zero setStorage / illegal-memory-access events, coherent generation.
• Red→green regression test (added as tests/legacy/test_issue33_gdn_buffer_decoupling.py): it fails against the v7.72.5 behavior (budget hardcoded 4096 ignoring scheduler_config; no capacity guard; [:5536] on a 4096-row buffer truncates) and passes against dev424. Includes the exact overflow band — a 5536-token forward with the buffer nominally at the 4096 cap returns a correct full-size (5536, 16, 128) tensor instead of a truncated 4096-row view.

What we could not test: your exact Ada sm_89 + TP=2/EP=2 + AWQ + vision + 262K-ctx co-scheduled traffic shape that summed a single forward into (4096, 8192]. We're sm_86 and don't run EP=2 or vision. Our verification proves the decoupling logic and no-regression; it does not replay your precise repro. If you can reproduce-then-confirm on Ada with your crash-reproducing traffic (several concurrent ~2-3K-token prefills landing in one step, batched=8192), that would close the loop on the hardware we can't cover.

Recommendation / next step

The cleanest path is to upgrade off v7.72.5 to a tree that carries P28 v2 + the budget resolver (current dev424 baseline). If upgrading the whole tree isn't immediately practical, the minimal mitigation on your existing image is to make the prealloc budget explicit so it can't fall back to 4096:

-e GENESIS_PREALLOC_TOKEN_BUDGET=8192   # = your --max-num-batched-tokens

(set it equal to --max-num-batched-tokens). That sizes the GDN buffer to the batch and sidesteps the truncation even on the older forward path — though the v2 capacity guard in the upgraded tree is the real fix and removes the foot-gun entirely.

Happy to share the branch / the two patch diffs (P28 v2 forward_cuda guard + resolve_token_budget scheduler probe) and the regression test if useful, and very glad to get your Ada confirmation.