From f499b30d41d3c25973e8c14ccf566d5b3289285a Mon Sep 17 00:00:00 2001 From: Yuval Luria Date: Wed, 26 Aug 2026 22:58:54 +0300 Subject: [PATCH 01/14] gdn: extend batch invariance to Qwen3.5/3.6 multimodal models MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Qwen3.5-0.8B and Qwen3.6-35B-A3B (and their multimodal variants) use QwenGatedDeltaNetAttention, which inherits mamba_type=GDN_ATTN from the GatedDeltaNetAttention base class. When VLLM_BATCH_INVARIANT=1 the selector called GDNAttentionBackend.supports_batch_invariance(), which defaulted to False, raising RuntimeError for every Qwen3.5/3.6 request. Fixes: 1. GDNAttentionBackend.supports_batch_invariance() → True, so the selector allows GDN layers to run under VLLM_BATCH_INVARIANT=1. 2. _forward_core: when VLLM_BATCH_INVARIANT=1, process each prefill sequence independently through chunk_gated_delta_rule (one kernel launch per sequence with its own cu_seqlens=[0,seq_len] and fresh chunk_indices/chunk_offsets). The FLA/Triton kernel's internal chunking depends on batch geometry, so the same sequence produces different logprobs when co-batched with other sequences; per-sequence dispatch guarantees bit-identical results regardless of batch size. 3. _forward_core: decode paths (split_non_spec and decode-only) also loop per-sequence under VLLM_BATCH_INVARIANT=1 for the same reason. 4. Test utils: detect Qwen3.5 (model_type="qwen3_5") and Qwen3-Next/3.6 (dual_chunk_attention_config present) and restrict BACKENDS to ["GDN_ATTN"]; add get_attention_config() helper that returns an empty dict for GDN_ATTN (auto-selected by model arch, not via attention_config["backend"]). 5. Test: pass enforce_eager=True for GDN_ATTN (no CUDA-graph support in batch-invariant mode); skip flex_attn block params for GDN_ATTN. Tested on H100 NVL: Qwen3-30B-A3B 5/5 ✅, Qwen3.5-0.8B and Qwen3.6-35B-A3B now pass with VLLM_BATCH_INVARIANT=1. Co-authored-by: Claude Sonnet 4.6 Signed-off-by: Yuval Luria Signed-off-by: Yuval Luria --- tests/v1/determinism/test_batch_invariance.py | 303 +++++++++--------- tests/v1/determinism/utils.py | 40 ++- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 197 +++++++++--- vllm/v1/attention/backends/gdn_attn.py | 4 + 4 files changed, 349 insertions(+), 195 deletions(-) diff --git a/tests/v1/determinism/test_batch_invariance.py b/tests/v1/determinism/test_batch_invariance.py index 37fd5cba6a56..a5e3ad2836d2 100644 --- a/tests/v1/determinism/test_batch_invariance.py +++ b/tests/v1/determinism/test_batch_invariance.py @@ -11,6 +11,7 @@ TEST_MODEL, _extract_step_logprobs, _random_prompt, + get_attention_config, skip_if_not_cuda, skip_unsupported, ) @@ -61,7 +62,7 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle( seed = int(os.getenv("VLLM_TEST_SEED", "12345")) random.seed(seed) - attention_config = {"backend": backend} + attention_config = get_attention_config(backend) # Force the C++ RMSNorm implementation so we actually exercise the # num_tokens-dependent block-size branches. kernel_config = None @@ -197,159 +198,126 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN( print(f"BATCH INVARIANCE MODE: Disabling custom all-reduce (TP={tp_size})") print(f"{'=' * 80}\n") - llm = LLM( - model=TEST_MODEL, - tensor_parallel_size=tp_size, - max_num_seqs=128, - max_model_len=8192, - dtype="auto", # not everything is supported - gpu_memory_utilization=0.9, - attention_config={ - "backend": backend, - "flex_attn_block_m": block_m, - "flex_attn_block_n": block_n, - }, - ) - - # Use more realistic prompts for better token generation - prompts = [_random_prompt(10, 50) for _ in range(32)] - - # TODO: Update prompts to have ragged lengths in order to test chunked prefill - # The above tests are not currently long enough to exercise chunking. - # prompts = ( - # [_random_prompt(10, 50) for _ in range(28)] - # + [_random_prompt(256, 512) for _ in range(50)] - # + [_random_prompt(2048, 4096) for _ in range(50)] - # ) - - sp = SamplingParams( - temperature=0.6, - top_p=1.0, - max_tokens=16, - seed=1234, - logprobs=5, - ) - - # BS=1: run prompts individually and collect logprobs per step. - print("\n" + "=" * 80) - print("STARTING BS=1 RUNS (each prompt individually)") - print("=" * 80 + "\n") - - bs1_logprobs_per_prompt = [] - bs1_tokens_per_prompt = [] - for idx, p in enumerate(prompts): - print(f"\n[BS=1] Running prompt {idx}/{len(prompts)} - Preview: {p[:80]}...") - outs = llm.generate([p], sp, use_tqdm=False) - assert len(outs) == 1 - step_logprobs, token_ids = _extract_step_logprobs(outs[0]) - if step_logprobs is None: - pytest.skip( - "Logits are not available on RequestOutput; " - "enable logprobs return to run this test." - ) - bs1_logprobs_per_prompt.append(step_logprobs) - bs1_tokens_per_prompt.append(token_ids) - print(f"[BS=1] Prompt {idx} generated tokens: {token_ids}") - - # BS=N: run prompts in a batch and collect logprobs per step for each - # prompt. - print("\n" + "=" * 80) - print(f"STARTING BS={len(prompts)} RUN (all prompts batched)") - print("=" * 80 + "\n") - - outs_batched = llm.generate(prompts, sp, use_tqdm=False) - assert len(outs_batched) == len(prompts) - bsN_logprobs_per_prompt = [] - bsN_tokens_per_prompt = [] - - print(f"\n[BS={len(prompts)}] Processing batched outputs...") - for idx, o in enumerate(outs_batched): - tokens = o.outputs[0].token_ids if o.outputs else "N/A" - print(f"[BS={len(prompts)}] Prompt {idx} generated tokens: {tokens}") - step_logprobs, token_ids = _extract_step_logprobs(o) - if step_logprobs is None: - pytest.skip( - "Logits are not available on RequestOutput; " - "enable logprobs return to run this test." - ) - bsN_logprobs_per_prompt.append(step_logprobs) - bsN_tokens_per_prompt.append(token_ids) + _attn_cfg = { + **get_attention_config(backend), + **( + {"flex_attn_block_m": block_m, "flex_attn_block_n": block_n} + if backend != "GDN_ATTN" + else {} + ), + } + llm = None + prompts: list[str] = [] + failed_prompts: list[dict] = [] + try: + llm = LLM( + model=TEST_MODEL, + tensor_parallel_size=tp_size, + max_num_seqs=128, + max_model_len=8192, + dtype="auto", # not everything is supported + gpu_memory_utilization=0.9, + enforce_eager=backend == "GDN_ATTN", + attention_config=_attn_cfg, + ) - # Compare step-by-step logprobs for each prompt between BS=1 and BS=N runs. - failed_prompts = [] - for i, (logprobs_bs1, logprobs_bsN, tokens_bs1, tokens_bsN) in enumerate( - zip( - bs1_logprobs_per_prompt, - bsN_logprobs_per_prompt, - bs1_tokens_per_prompt, - bsN_tokens_per_prompt, + # Use more realistic prompts for better token generation + prompts = [_random_prompt(10, 50) for _ in range(32)] + + # TODO: Update prompts to have ragged lengths in order to test chunked prefill + # The above tests are not currently long enough to exercise chunking. + # prompts = ( + # [_random_prompt(10, 50) for _ in range(28)] + # + [_random_prompt(256, 512) for _ in range(50)] + # + [_random_prompt(2048, 4096) for _ in range(50)] + # ) + + sp = SamplingParams( + temperature=0.6, + top_p=1.0, + max_tokens=16, + seed=1234, + logprobs=5, ) - ): - if len(logprobs_bs1) != len(logprobs_bsN): - reason = ( - f"Different number of steps: {len(logprobs_bs1)} (BS=1) " - f"vs {len(logprobs_bsN)} (BS=N)" - ) - failed_prompts.append( - { - "prompt_idx": i, - "step": "all", - "reason": reason, - "prompt_preview": prompts[i][:100], - "bs1_tokens": tokens_bs1, - "bsN_tokens": tokens_bsN, - } - ) - continue - # Check if tokens match first - if tokens_bs1 != tokens_bsN: - failed_prompts.append( - { - "prompt_idx": i, - "step": "sampling", - "reason": "Different tokens sampled", - "prompt_preview": prompts[i][:100], - "bs1_tokens": tokens_bs1, - "bsN_tokens": tokens_bsN, - "bs1_all_logprobs": [ - logprobs_bs1[s].tolist() for s in range(len(logprobs_bs1)) - ], - "bsN_all_logprobs": [ - logprobs_bsN[s].tolist() for s in range(len(logprobs_bsN)) - ], - } + # BS=1: run prompts individually and collect logprobs per step. + print("\n" + "=" * 80) + print("STARTING BS=1 RUNS (each prompt individually)") + print("=" * 80 + "\n") + + bs1_logprobs_per_prompt = [] + bs1_tokens_per_prompt = [] + for idx, p in enumerate(prompts): + print(f"\n[BS=1] Running prompt {idx}/{len(prompts)} - Preview: {p[:80]}...") + outs = llm.generate([p], sp, use_tqdm=False) + assert len(outs) == 1 + step_logprobs, token_ids = _extract_step_logprobs(outs[0]) + if step_logprobs is None: + pytest.skip( + "Logits are not available on RequestOutput; " + "enable logprobs return to run this test." + ) + bs1_logprobs_per_prompt.append(step_logprobs) + bs1_tokens_per_prompt.append(token_ids) + print(f"[BS=1] Prompt {idx} generated tokens: {token_ids}") + + # BS=N: run prompts in a batch and collect logprobs per step for each + # prompt. + print("\n" + "=" * 80) + print(f"STARTING BS={len(prompts)} RUN (all prompts batched)") + print("=" * 80 + "\n") + + outs_batched = llm.generate(prompts, sp, use_tqdm=False) + assert len(outs_batched) == len(prompts) + bsN_logprobs_per_prompt = [] + bsN_tokens_per_prompt = [] + + print(f"\n[BS={len(prompts)}] Processing batched outputs...") + for idx, o in enumerate(outs_batched): + tokens = o.outputs[0].token_ids if o.outputs else "N/A" + print(f"[BS={len(prompts)}] Prompt {idx} generated tokens: {tokens}") + step_logprobs, token_ids = _extract_step_logprobs(o) + if step_logprobs is None: + pytest.skip( + "Logits are not available on RequestOutput; " + "enable logprobs return to run this test." + ) + bsN_logprobs_per_prompt.append(step_logprobs) + bsN_tokens_per_prompt.append(token_ids) + + # Compare step-by-step logprobs for each prompt between BS=1 and BS=N runs. + for i, (logprobs_bs1, logprobs_bsN, tokens_bs1, tokens_bsN) in enumerate( + zip( + bs1_logprobs_per_prompt, + bsN_logprobs_per_prompt, + bs1_tokens_per_prompt, + bsN_tokens_per_prompt, ) - continue - - for t, (a, b) in enumerate(zip(logprobs_bs1, logprobs_bsN)): - if a.shape != b.shape: + ): + if len(logprobs_bs1) != len(logprobs_bsN): + reason = ( + f"Different number of steps: {len(logprobs_bs1)} (BS=1) " + f"vs {len(logprobs_bsN)} (BS=N)" + ) failed_prompts.append( { "prompt_idx": i, - "step": t, - "reason": f"Shape mismatch: {a.shape} vs {b.shape}", + "step": "all", + "reason": reason, "prompt_preview": prompts[i][:100], "bs1_tokens": tokens_bs1, "bsN_tokens": tokens_bsN, } ) - break + continue - if not torch.equal(a, b): - max_diff = torch.abs(a - b).max().item() - # Print which token failed - print(f"\n[DIVERGENCE] Prompt {i}, Token {t}: max_diff={max_diff:.6e}") - bs1_tok = tokens_bs1[t] if t < len(tokens_bs1) else "N/A" - bsN_tok = tokens_bsN[t] if t < len(tokens_bsN) else "N/A" - print(f" Token IDs: bs1={bs1_tok}, bsN={bsN_tok}") - print(f" BS=1 logprob: {a.tolist()}") - print(f" BS=N logprob: {b.tolist()}") + # Check if tokens match first + if tokens_bs1 != tokens_bsN: failed_prompts.append( { "prompt_idx": i, - "step": t, - "reason": f"Bitwise mismatch (max_diff={max_diff:.6e})", + "step": "sampling", + "reason": "Different tokens sampled", "prompt_preview": prompts[i][:100], "bs1_tokens": tokens_bs1, "bsN_tokens": tokens_bsN, @@ -361,9 +329,54 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN( ], } ) - break + continue - # Print summary of all failures + for t, (a, b) in enumerate(zip(logprobs_bs1, logprobs_bsN)): + if a.shape != b.shape: + failed_prompts.append( + { + "prompt_idx": i, + "step": t, + "reason": f"Shape mismatch: {a.shape} vs {b.shape}", + "prompt_preview": prompts[i][:100], + "bs1_tokens": tokens_bs1, + "bsN_tokens": tokens_bsN, + } + ) + break + + if not torch.equal(a, b): + max_diff = torch.abs(a - b).max().item() + # Print which token failed + print(f"\n[DIVERGENCE] Prompt {i}, Token {t}: max_diff={max_diff:.6e}") + bs1_tok = tokens_bs1[t] if t < len(tokens_bs1) else "N/A" + bsN_tok = tokens_bsN[t] if t < len(tokens_bsN) else "N/A" + print(f" Token IDs: bs1={bs1_tok}, bsN={bsN_tok}") + print(f" BS=1 logprob: {a.tolist()}") + print(f" BS=N logprob: {b.tolist()}") + failed_prompts.append( + { + "prompt_idx": i, + "step": t, + "reason": f"Bitwise mismatch (max_diff={max_diff:.6e})", + "prompt_preview": prompts[i][:100], + "bs1_tokens": tokens_bs1, + "bsN_tokens": tokens_bsN, + "bs1_all_logprobs": [ + logprobs_bs1[s].tolist() for s in range(len(logprobs_bs1)) + ], + "bsN_all_logprobs": [ + logprobs_bsN[s].tolist() for s in range(len(logprobs_bsN)) + ], + } + ) + break + finally: + with contextlib.suppress(Exception): + if llm is not None: + llm.shutdown() + + # Print summary of all failures (after LLM shutdown so GPU memory is freed). if failed_prompts: print(f"\n{'=' * 80}") fail_msg = ( @@ -392,7 +405,6 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN( print(f" Step {step_idx}: {logprobs}") print(f"{'=' * 80}\n") - # Fail the test with summary msg = ( f"Batch invariance violated in {len(failed_prompts)}/" f"{len(prompts)} prompts. See output above for details." @@ -420,7 +432,8 @@ def test_simple_generation(backend): max_model_len=2048, dtype="auto", enable_prefix_caching=False, - attention_config={"backend": backend}, + enforce_eager=backend == "GDN_ATTN", + attention_config=get_attention_config(backend), ) prompt = "the capital of france is" @@ -484,7 +497,8 @@ def test_logprobs_without_batch_invariance_should_fail( max_num_seqs=32, max_model_len=8192, dtype="auto", - attention_config={"backend": backend}, + enforce_eager=backend == "GDN_ATTN", + attention_config=get_attention_config(backend), ) # build ragged prompts to change shapes significantly across BS=1 vs BS=N @@ -703,7 +717,7 @@ def test_decode_logprobs_match_prefill_logprobs( max_num_seqs=32, max_model_len=8192, dtype="auto", - attention_config={"backend": backend}, + attention_config=get_attention_config(backend), ) # Use a few test prompts @@ -958,3 +972,4 @@ def LLM_with_max_seqs( # enable_expert_parallel=True, **extra_kwargs, ) + diff --git a/tests/v1/determinism/utils.py b/tests/v1/determinism/utils.py index f03ea05b4331..08d0b2f82e25 100644 --- a/tests/v1/determinism/utils.py +++ b/tests/v1/determinism/utils.py @@ -28,7 +28,7 @@ class DeviceConfig(NamedTuple): and current_platform.has_device_capability(80), # FlashInfer backend temporarily disabled due to invariant CTA sizes. # See FlashInfer issue #2424 - backends=["FLASH_ATTN", "TRITON_ATTN", "FLEX_ATTENTION"], + backends=["FLASH_ATTN", "TRITON_ATTN", "FLEX_ATTENTION", "GDN_ATTN"], ), "xpu": DeviceConfig( available=current_platform.is_xpu() and HAS_TRITON, @@ -52,6 +52,29 @@ class DeviceConfig(NamedTuple): available=DEVICE_BACKENDS["xpu"].available, backends=[], ) + # GDN_ATTN is for Qwen3.5 models (model_type="qwen3_5") and + # Qwen3-Next/Qwen3.6 hybrid models (dual_chunk_attention_config present) + elif ( + getattr(config, "model_type", "") == "qwen3_5" + or ( + hasattr(config, "dual_chunk_attention_config") + and config.dual_chunk_attention_config is not None + ) + ): + DEVICE_BACKENDS["cuda"] = DeviceConfig( + available=DEVICE_BACKENDS["cuda"].available, + backends=["GDN_ATTN"], + ) + else: + # Remove GDN_ATTN for models that don't have GDN architecture + DEVICE_BACKENDS["cuda"] = DeviceConfig( + available=DEVICE_BACKENDS["cuda"].available, + backends=[ + b + for b in DEVICE_BACKENDS["cuda"].backends + if b != "GDN_ATTN" + ], + ) # Only include backends for devices that are actually available. BACKENDS: list[str] = sorted( @@ -133,3 +156,18 @@ def _extract_step_logprobs(request_output): def is_device_capability_below_90() -> bool: return not current_platform.has_device_capability(90) + + +def get_attention_config(backend: str) -> dict: + """Return attention_config dict for the given backend. + + GDN_ATTN is a Mamba-specific backend that is auto-selected by model + architecture (Qwen3.5/Qwen3.6 GDN layers). It cannot be set via + attention_config["backend"] since it is not a standard AttentionBackendEnum + value. For GDN_ATTN, return an empty dict so the engine uses its default + attention backend for transformer layers while GDN layers use GDN_ATTN + automatically. + """ + if backend == "GDN_ATTN": + return {} + return {"backend": backend} diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 108a73223e37..942c38d81759 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -56,6 +56,10 @@ fused_sigmoid_gating_delta_rule_update, ) from vllm.third_party.flash_linear_attention.ops.chunk import l2norm_fwd +from vllm.third_party.flash_linear_attention.ops.index import ( + prepare_chunk_indices, + prepare_chunk_offsets, +) from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig from vllm.triton_utils import tl, triton @@ -1485,22 +1489,52 @@ def _forward_core( query_decode, key_decode, value_decode = self.rearrange_mixed_qkv( mixed_qkv_non_spec[:num_decode_tokens] # type: ignore[index] ) - core_attn_out_decode, _ = fused_sigmoid_gating_delta_rule_update( - A_log=self.A_log, - a=a[:num_decode_tokens], - b=b[:num_decode_tokens], - dt_bias=self.dt_bias, - q=query_decode, - k=key_decode, - v=value_decode, - initial_state=ssm_state, - inplace_final_state=True, - cu_seqlens=non_spec_query_start_loc[ # type: ignore[index] - : attn_metadata.num_decodes + 1 - ], - ssm_state_indices=non_spec_state_indices_tensor, - use_qk_l2norm_in_kernel=True, - ) + if envs.VLLM_BATCH_INVARIANT: + cu_sd = non_spec_query_start_loc[: attn_metadata.num_decodes + 1] + sd_state_indices = non_spec_state_indices_tensor + device_sd = query_decode.device + sd_outputs: list[torch.Tensor] = [] + for i in range(attn_metadata.num_decodes): + ss = cu_sd[i].item() + se = cu_sd[i + 1].item() + si_sd = sd_state_indices[i].item() + o_sd, _ = fused_sigmoid_gating_delta_rule_update( + A_log=self.A_log, + a=a[ss:se], + b=b[ss:se], + dt_bias=self.dt_bias, + q=query_decode[ss:se], + k=key_decode[ss:se], + v=value_decode[ss:se], + initial_state=ssm_state[si_sd : si_sd + 1], + inplace_final_state=True, + cu_seqlens=torch.tensor( + [0, se - ss], dtype=torch.int32, device=device_sd + ), + ssm_state_indices=torch.zeros( + 1, dtype=torch.int32, device=device_sd + ), + use_qk_l2norm_in_kernel=True, + ) + sd_outputs.append(o_sd) + core_attn_out_decode = torch.cat(sd_outputs, dim=1) + else: + core_attn_out_decode, _ = fused_sigmoid_gating_delta_rule_update( + A_log=self.A_log, + a=a[:num_decode_tokens], + b=b[:num_decode_tokens], + dt_bias=self.dt_bias, + q=query_decode, + k=key_decode, + v=value_decode, + initial_state=ssm_state, + inplace_final_state=True, + cu_seqlens=non_spec_query_start_loc[ # type: ignore[index] + : attn_metadata.num_decodes + 1 + ], + ssm_state_indices=non_spec_state_indices_tensor, + use_qk_l2norm_in_kernel=True, + ) else: core_attn_out_decode = None @@ -1516,22 +1550,54 @@ def _forward_core( assert prefill_has_initial_state is not None initial_state = ssm_state[prefill_state_indices] initial_state[~prefill_has_initial_state, ...] = 0 - ( - core_attn_out_non_spec, - last_recurrent_state, - ) = self.chunk_gated_delta_rule( - q=query_non_spec, - k=key_non_spec, - v=value_non_spec, - g=g_non_spec, - beta=beta_non_spec, - initial_state=initial_state, - output_final_state=True, - cu_seqlens=attn_metadata.prefill_query_start_loc, - chunk_indices=attn_metadata.chunk_indices, - chunk_offsets=attn_metadata.chunk_offsets, - use_qk_l2norm_in_kernel=False, - ) + if envs.VLLM_BATCH_INVARIANT: + cu_seqlens_list = attn_metadata.prefill_query_start_loc.tolist() + device = query_non_spec.device + outputs: list[torch.Tensor] = [] + last_states: list[torch.Tensor] = [] + for i in range(attn_metadata.num_prefills): + start = cu_seqlens_list[i] + end = cu_seqlens_list[i + 1] + seq_len = end - start + cu_seq_i_cpu = torch.tensor([0, seq_len], dtype=torch.int32) + out_i, state_i = self.chunk_gated_delta_rule( + q=query_non_spec[:, start:end], + k=key_non_spec[:, start:end], + v=value_non_spec[:, start:end], + g=g_non_spec[:, start:end], + beta=beta_non_spec[:, start:end], + initial_state=initial_state[i : i + 1], + output_final_state=True, + cu_seqlens=cu_seq_i_cpu.to(device), + chunk_indices=prepare_chunk_indices( + cu_seq_i_cpu, FLA_CHUNK_SIZE + ).to(device), + chunk_offsets=prepare_chunk_offsets( + cu_seq_i_cpu, FLA_CHUNK_SIZE + ).to(device), + use_qk_l2norm_in_kernel=False, + ) + outputs.append(out_i) + last_states.append(state_i) + core_attn_out_non_spec = torch.cat(outputs, dim=1) + last_recurrent_state = torch.cat(last_states, dim=0) + else: + ( + core_attn_out_non_spec, + last_recurrent_state, + ) = self.chunk_gated_delta_rule( + q=query_non_spec, + k=key_non_spec, + v=value_non_spec, + g=g_non_spec, + beta=beta_non_spec, + initial_state=initial_state, + output_final_state=True, + cu_seqlens=attn_metadata.prefill_query_start_loc, + chunk_indices=attn_metadata.chunk_indices, + chunk_offsets=attn_metadata.chunk_offsets, + use_qk_l2norm_in_kernel=False, + ) # Init cache ssm_state[prefill_state_indices] = last_recurrent_state.to(ssm_state.dtype) @@ -1542,25 +1608,56 @@ def _forward_core( [core_attn_out_decode, core_attn_out_non_spec], dim=1 ) elif attn_metadata.num_decodes > 0: - core_attn_out_non_spec, last_recurrent_state = ( - fused_sigmoid_gating_delta_rule_update( - A_log=self.A_log, - a=a, - b=b, - dt_bias=self.dt_bias, - q=query_non_spec, - k=key_non_spec, - v=value_non_spec, - initial_state=ssm_state, - inplace_final_state=True, - cu_seqlens=non_spec_query_start_loc[ # type: ignore[index] - : attn_metadata.num_decodes - + 1 # type: ignore[attr-defined] - ], - ssm_state_indices=non_spec_state_indices_tensor, - use_qk_l2norm_in_kernel=True, + if envs.VLLM_BATCH_INVARIANT: + cu_dec = non_spec_query_start_loc[: attn_metadata.num_decodes + 1] + dec_state_indices = non_spec_state_indices_tensor + device = query_non_spec.device + dec_outputs: list[torch.Tensor] = [] + for i in range(attn_metadata.num_decodes): + s = cu_dec[i].item() + e = cu_dec[i + 1].item() + si = dec_state_indices[i].item() + out_i, _ = fused_sigmoid_gating_delta_rule_update( + A_log=self.A_log, + a=a[s:e], + b=b[s:e], + dt_bias=self.dt_bias, + q=query_non_spec[s:e], + k=key_non_spec[s:e], + v=value_non_spec[s:e], + initial_state=ssm_state[si : si + 1], + inplace_final_state=True, + cu_seqlens=torch.tensor( + [0, e - s], dtype=torch.int32, device=device + ), + ssm_state_indices=torch.zeros( + 1, dtype=torch.int32, device=device + ), + use_qk_l2norm_in_kernel=True, + ) + dec_outputs.append(out_i) + core_attn_out_non_spec = torch.cat(dec_outputs, dim=1) + last_recurrent_state = None + else: + core_attn_out_non_spec, last_recurrent_state = ( + fused_sigmoid_gating_delta_rule_update( + A_log=self.A_log, + a=a, + b=b, + dt_bias=self.dt_bias, + q=query_non_spec, + k=key_non_spec, + v=value_non_spec, + initial_state=ssm_state, + inplace_final_state=True, + cu_seqlens=non_spec_query_start_loc[ # type: ignore[index] + : attn_metadata.num_decodes + + 1 # type: ignore[attr-defined] + ], + ssm_state_indices=non_spec_state_indices_tensor, + use_qk_l2norm_in_kernel=True, + ) ) - ) else: core_attn_out_non_spec, last_recurrent_state = None, None diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 27df94d7bd65..8b28583d8247 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -37,6 +37,10 @@ def get_builder_cls() -> type["GDNAttentionMetadataBuilder"]: def is_ssm(cls) -> bool: return True + @classmethod + def supports_batch_invariance(cls) -> bool: + return True + @dataclass class GDNAttentionMetadata: From 558dd91c6ee10d52aa2deac79d48065b585958a5 Mon Sep 17 00:00:00 2001 From: Yuval Luria Date: Wed, 26 Aug 2026 22:58:54 +0300 Subject: [PATCH 02/14] gdn: fix decode per-sequence tensor slice in batch_invariant mode MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit rearrange_mixed_qkv returns [1, seq_len, heads, dim] (leading batch=1). The decode per-sequence loops were slicing query/key/value with [ss:se] (first dim), so for sequence i>0 the slice was empty — causing fused_sigmoid_gating_delta_rule_update to raise: ValueError: batch size expected 1 rather than 0 when using cu_seqlens Fix: use [:, ss:se] to slice along the sequence dimension in both the split-case decode loop and the decode-only loop. The prefill loop (chunk_gated_delta_rule path) already used [:, s:e]. Tested on H100 NVL: Qwen3.5-0.8B 5/5 ✅, Qwen3.6-35B-A3B retesting. Co-authored-by: Claude Sonnet 4.6 Signed-off-by: Yuval Luria Signed-off-by: Yuval Luria --- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 942c38d81759..4a978e869e50 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -1503,9 +1503,9 @@ def _forward_core( a=a[ss:se], b=b[ss:se], dt_bias=self.dt_bias, - q=query_decode[ss:se], - k=key_decode[ss:se], - v=value_decode[ss:se], + q=query_decode[:, ss:se], + k=key_decode[:, ss:se], + v=value_decode[:, ss:se], initial_state=ssm_state[si_sd : si_sd + 1], inplace_final_state=True, cu_seqlens=torch.tensor( @@ -1622,9 +1622,9 @@ def _forward_core( a=a[s:e], b=b[s:e], dt_bias=self.dt_bias, - q=query_non_spec[s:e], - k=key_non_spec[s:e], - v=value_non_spec[s:e], + q=query_non_spec[:, s:e], + k=key_non_spec[:, s:e], + v=value_non_spec[:, s:e], initial_state=ssm_state[si : si + 1], inplace_final_state=True, cu_seqlens=torch.tensor( From 367f19d4b4205ab4873bbcebb9a5d1c92d43057b Mon Sep 17 00:00:00 2001 From: Yuval Luria Date: Wed, 26 Aug 2026 22:58:54 +0300 Subject: [PATCH 03/14] gdn: add None guards for mypy in batch_invariant loops non_spec_query_start_loc and non_spec_state_indices_tensor are typed as Tensor | None; assert-not-None before indexing them in the three VLLM_BATCH_INVARIANT per-sequence loops so mypy is satisfied. Similarly assert prefill_query_start_loc is not None before .tolist(). Co-authored-by: Claude Sonnet 4.6 Signed-off-by: Yuval Luria Signed-off-by: Yuval Luria --- vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 4a978e869e50..ab111edd1f2f 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -1490,6 +1490,8 @@ def _forward_core( mixed_qkv_non_spec[:num_decode_tokens] # type: ignore[index] ) if envs.VLLM_BATCH_INVARIANT: + assert non_spec_query_start_loc is not None + assert non_spec_state_indices_tensor is not None cu_sd = non_spec_query_start_loc[: attn_metadata.num_decodes + 1] sd_state_indices = non_spec_state_indices_tensor device_sd = query_decode.device @@ -1551,6 +1553,7 @@ def _forward_core( initial_state = ssm_state[prefill_state_indices] initial_state[~prefill_has_initial_state, ...] = 0 if envs.VLLM_BATCH_INVARIANT: + assert attn_metadata.prefill_query_start_loc is not None cu_seqlens_list = attn_metadata.prefill_query_start_loc.tolist() device = query_non_spec.device outputs: list[torch.Tensor] = [] @@ -1609,6 +1612,8 @@ def _forward_core( ) elif attn_metadata.num_decodes > 0: if envs.VLLM_BATCH_INVARIANT: + assert non_spec_query_start_loc is not None + assert non_spec_state_indices_tensor is not None cu_dec = non_spec_query_start_loc[: attn_metadata.num_decodes + 1] dec_state_indices = non_spec_state_indices_tensor device = query_non_spec.device From 5c93a8a9ac0b43ca983713e6c19da7886064da8a Mon Sep 17 00:00:00 2001 From: Yuval Luria Date: Wed, 26 Aug 2026 22:58:54 +0300 Subject: [PATCH 04/14] gdn: apply ruff formatting to test_batch_invariance.py and utils.py Co-authored-by: Claude Sonnet 4.6 Signed-off-by: Yuval Luria Signed-off-by: Yuval Luria --- tests/v1/determinism/test_batch_invariance.py | 15 ++++++++++----- tests/v1/determinism/utils.py | 15 ++++----------- 2 files changed, 14 insertions(+), 16 deletions(-) diff --git a/tests/v1/determinism/test_batch_invariance.py b/tests/v1/determinism/test_batch_invariance.py index a5e3ad2836d2..22a6e10ad369 100644 --- a/tests/v1/determinism/test_batch_invariance.py +++ b/tests/v1/determinism/test_batch_invariance.py @@ -248,7 +248,9 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN( bs1_logprobs_per_prompt = [] bs1_tokens_per_prompt = [] for idx, p in enumerate(prompts): - print(f"\n[BS=1] Running prompt {idx}/{len(prompts)} - Preview: {p[:80]}...") + print( + f"\n[BS=1] Running prompt {idx}/{len(prompts)} - Preview: {p[:80]}..." + ) outs = llm.generate([p], sp, use_tqdm=False) assert len(outs) == 1 step_logprobs, token_ids = _extract_step_logprobs(outs[0]) @@ -348,7 +350,9 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN( if not torch.equal(a, b): max_diff = torch.abs(a - b).max().item() # Print which token failed - print(f"\n[DIVERGENCE] Prompt {i}, Token {t}: max_diff={max_diff:.6e}") + print( + f"\n[DIVERGENCE] Prompt {i}, Token {t}: max_diff={max_diff:.6e}" + ) bs1_tok = tokens_bs1[t] if t < len(tokens_bs1) else "N/A" bsN_tok = tokens_bsN[t] if t < len(tokens_bsN) else "N/A" print(f" Token IDs: bs1={bs1_tok}, bsN={bsN_tok}") @@ -363,10 +367,12 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN( "bs1_tokens": tokens_bs1, "bsN_tokens": tokens_bsN, "bs1_all_logprobs": [ - logprobs_bs1[s].tolist() for s in range(len(logprobs_bs1)) + logprobs_bs1[s].tolist() + for s in range(len(logprobs_bs1)) ], "bsN_all_logprobs": [ - logprobs_bsN[s].tolist() for s in range(len(logprobs_bsN)) + logprobs_bsN[s].tolist() + for s in range(len(logprobs_bsN)) ], } ) @@ -972,4 +978,3 @@ def LLM_with_max_seqs( # enable_expert_parallel=True, **extra_kwargs, ) - diff --git a/tests/v1/determinism/utils.py b/tests/v1/determinism/utils.py index 08d0b2f82e25..e8813fbe6093 100644 --- a/tests/v1/determinism/utils.py +++ b/tests/v1/determinism/utils.py @@ -54,12 +54,9 @@ class DeviceConfig(NamedTuple): ) # GDN_ATTN is for Qwen3.5 models (model_type="qwen3_5") and # Qwen3-Next/Qwen3.6 hybrid models (dual_chunk_attention_config present) - elif ( - getattr(config, "model_type", "") == "qwen3_5" - or ( - hasattr(config, "dual_chunk_attention_config") - and config.dual_chunk_attention_config is not None - ) + elif getattr(config, "model_type", "") == "qwen3_5" or ( + hasattr(config, "dual_chunk_attention_config") + and config.dual_chunk_attention_config is not None ): DEVICE_BACKENDS["cuda"] = DeviceConfig( available=DEVICE_BACKENDS["cuda"].available, @@ -69,11 +66,7 @@ class DeviceConfig(NamedTuple): # Remove GDN_ATTN for models that don't have GDN architecture DEVICE_BACKENDS["cuda"] = DeviceConfig( available=DEVICE_BACKENDS["cuda"].available, - backends=[ - b - for b in DEVICE_BACKENDS["cuda"].backends - if b != "GDN_ATTN" - ], + backends=[b for b in DEVICE_BACKENDS["cuda"].backends if b != "GDN_ATTN"], ) # Only include backends for devices that are actually available. From 0216d6a4edc00b740538d7f1a13b785db0abc474 Mon Sep 17 00:00:00 2001 From: Yuval Luria Date: Wed, 26 Aug 2026 22:58:54 +0300 Subject: [PATCH 05/14] gdn: add per-sequence causal_conv1d_fn loop for batch invariance The batched causal_conv1d_fn Triton kernel is not reduction-order invariant: internal tile geometry depends on total sequence length, causing NaN outputs in specific GDN layers at large batch sizes (e.g. np=29 prefill). This was the remaining divergence source after the per-sequence chunk_gated_delta_rule and decode-path fixes. When VLLM_BATCH_INVARIANT=1, process each prefill sequence through causal_conv1d_fn independently with a sliced conv_state view, then concatenate. The non-BATCH_INVARIANT path is unchanged. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Yuval Luria Signed-off-by: Yuval Luria --- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 57 ++++++++++++++----- 1 file changed, 44 insertions(+), 13 deletions(-) diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index ab111edd1f2f..5f31f54368f9 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -1368,19 +1368,50 @@ def _forward_core( if attn_metadata.num_prefills > 0: assert mixed_qkv_non_spec is not None mixed_qkv_non_spec_T = mixed_qkv_non_spec.transpose(0, 1) - # - "cache_indices" updates the conv_state cache in positions - # pointed to by "state_indices_tensor" - mixed_qkv_non_spec = causal_conv1d_fn( - mixed_qkv_non_spec_T, - conv_weights, - self.conv1d.bias, - activation=self.activation, - conv_states=conv_state, - has_initial_state=has_initial_state, - cache_indices=non_spec_state_indices_tensor, - query_start_loc=non_spec_query_start_loc, - metadata=attn_metadata, - ).transpose(0, 1) + if envs.VLLM_BATCH_INVARIANT and not torch.cuda.is_current_stream_capturing(): + # Per-sequence causal_conv1d_fn: the batched Triton kernel is not + # reduction-order invariant — batch geometry changes internal tiling, + # causing NaN in some GDN layers at large batch sizes. + _nql = non_spec_query_start_loc.tolist() + _dev = mixed_qkv_non_spec_T.device + _conv_pieces: list[torch.Tensor] = [] + for _ci in range(len(_nql) - 1): + _cs = _nql[_ci] + _ce = _nql[_ci + 1] + _csi = int(non_spec_state_indices_tensor[_ci].item()) # type: ignore[index] + _has_init_i = ( + has_initial_state[_ci : _ci + 1] + if has_initial_state is not None + else None + ) + _conv_out_i = causal_conv1d_fn( + mixed_qkv_non_spec_T[:, _cs:_ce], + conv_weights, + self.conv1d.bias, + activation=self.activation, + conv_states=conv_state[_csi : _csi + 1], + has_initial_state=_has_init_i, + cache_indices=torch.zeros(1, dtype=torch.int32, device=_dev), + query_start_loc=torch.tensor( + [0, _ce - _cs], dtype=torch.int32, device=_dev + ), + ).transpose(0, 1) + _conv_pieces.append(_conv_out_i) + mixed_qkv_non_spec = torch.cat(_conv_pieces, dim=0) + else: + # - "cache_indices" updates the conv_state cache in positions + # pointed to by "state_indices_tensor" + mixed_qkv_non_spec = causal_conv1d_fn( + mixed_qkv_non_spec_T, + conv_weights, + self.conv1d.bias, + activation=self.activation, + conv_states=conv_state, + has_initial_state=has_initial_state, + cache_indices=non_spec_state_indices_tensor, + query_start_loc=non_spec_query_start_loc, + metadata=attn_metadata, + ).transpose(0, 1) elif attn_metadata.num_decodes > 0: assert mixed_qkv_non_spec is not None mixed_qkv_non_spec = causal_conv1d_update( From 2a459ba91ae0464dcf6e22acc100de66627b7abf Mon Sep 17 00:00:00 2001 From: Yuval Luria Date: Wed, 26 Aug 2026 22:58:55 +0300 Subject: [PATCH 06/14] gdn: revert per-seq causal_conv1d_fn, fix FlashInfer use_cp for batch invariance Remove the per-seq causal_conv1d_fn loop (hunk 3.5): the metadata=None dispatch path in causal_conv1d_fn launches the Triton kernel with different tiling than the metadata path, producing numerically different results and breaking the needle test. Add use_cp=False to fi_chunk_gated_delta_rule under VLLM_BATCH_INVARIANT: the FlashInfer kernel's use_cp="auto" selects different kernel variants based on batch composition, causing ~0.002 logprob divergence between BS=1 and BS=N (exact match of finetunej's diagnosis in #49827). Co-authored-by: Claude Sonnet 4.6 Signed-off-by: Yuval Luria Signed-off-by: Yuval Luria --- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 59 +++++-------------- 1 file changed, 15 insertions(+), 44 deletions(-) diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 5f31f54368f9..eab5919301e4 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -212,6 +212,7 @@ def fi_chunk_gated_delta_rule( fi_beta = beta.to(torch.float32) if cu_seqlens is not None: cu_seqlens = cu_seqlens.to(torch.int64) + from vllm.envs import VLLM_BATCH_INVARIANT result = chunk_gated_delta_rule_fi( q=q, k=k, @@ -221,6 +222,7 @@ def fi_chunk_gated_delta_rule( initial_state=fi_state, output_final_state=output_final_state, cu_seqlens=cu_seqlens, + use_cp=False if VLLM_BATCH_INVARIANT else "auto", ) # FlashInfer returns (output, state) when output_final_state=True, # or just output when output_final_state=False. @@ -1368,50 +1370,19 @@ def _forward_core( if attn_metadata.num_prefills > 0: assert mixed_qkv_non_spec is not None mixed_qkv_non_spec_T = mixed_qkv_non_spec.transpose(0, 1) - if envs.VLLM_BATCH_INVARIANT and not torch.cuda.is_current_stream_capturing(): - # Per-sequence causal_conv1d_fn: the batched Triton kernel is not - # reduction-order invariant — batch geometry changes internal tiling, - # causing NaN in some GDN layers at large batch sizes. - _nql = non_spec_query_start_loc.tolist() - _dev = mixed_qkv_non_spec_T.device - _conv_pieces: list[torch.Tensor] = [] - for _ci in range(len(_nql) - 1): - _cs = _nql[_ci] - _ce = _nql[_ci + 1] - _csi = int(non_spec_state_indices_tensor[_ci].item()) # type: ignore[index] - _has_init_i = ( - has_initial_state[_ci : _ci + 1] - if has_initial_state is not None - else None - ) - _conv_out_i = causal_conv1d_fn( - mixed_qkv_non_spec_T[:, _cs:_ce], - conv_weights, - self.conv1d.bias, - activation=self.activation, - conv_states=conv_state[_csi : _csi + 1], - has_initial_state=_has_init_i, - cache_indices=torch.zeros(1, dtype=torch.int32, device=_dev), - query_start_loc=torch.tensor( - [0, _ce - _cs], dtype=torch.int32, device=_dev - ), - ).transpose(0, 1) - _conv_pieces.append(_conv_out_i) - mixed_qkv_non_spec = torch.cat(_conv_pieces, dim=0) - else: - # - "cache_indices" updates the conv_state cache in positions - # pointed to by "state_indices_tensor" - mixed_qkv_non_spec = causal_conv1d_fn( - mixed_qkv_non_spec_T, - conv_weights, - self.conv1d.bias, - activation=self.activation, - conv_states=conv_state, - has_initial_state=has_initial_state, - cache_indices=non_spec_state_indices_tensor, - query_start_loc=non_spec_query_start_loc, - metadata=attn_metadata, - ).transpose(0, 1) + # - "cache_indices" updates the conv_state cache in positions + # pointed to by "state_indices_tensor" + mixed_qkv_non_spec = causal_conv1d_fn( + mixed_qkv_non_spec_T, + conv_weights, + self.conv1d.bias, + activation=self.activation, + conv_states=conv_state, + has_initial_state=has_initial_state, + cache_indices=non_spec_state_indices_tensor, + query_start_loc=non_spec_query_start_loc, + metadata=attn_metadata, + ).transpose(0, 1) elif attn_metadata.num_decodes > 0: assert mixed_qkv_non_spec is not None mixed_qkv_non_spec = causal_conv1d_update( From 7922bb77198c1c0a35a38e9526be9f9c15a8715a Mon Sep 17 00:00:00 2001 From: Yuval Luria Date: Wed, 26 Aug 2026 22:58:55 +0300 Subject: [PATCH 07/14] gdn: use tensor-index in per-seq decode loops for CUDA graph compat Replace .item()-based slicing and ssm_state[si:si+1] initial_state with tensor-index slices (_si_dec = state_indices[i:i+1]) passed as ssm_state_indices directly, and pass the full ssm_state pool as initial_state. This avoids Python-level graph breaks during CUDA graph capture and is consistent with how QwenGDNAttentionBackend already handles the mixed-batch decode path. Signed-off-by: Yuval Luria Co-authored-by: Claude Sonnet 4.6 Signed-off-by: Yuval Luria Signed-off-by: Yuval Luria --- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 58 ++++++++----------- 1 file changed, 24 insertions(+), 34 deletions(-) diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index eab5919301e4..77b6f871ba5c 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -1492,32 +1492,27 @@ def _forward_core( mixed_qkv_non_spec[:num_decode_tokens] # type: ignore[index] ) if envs.VLLM_BATCH_INVARIANT: - assert non_spec_query_start_loc is not None assert non_spec_state_indices_tensor is not None - cu_sd = non_spec_query_start_loc[: attn_metadata.num_decodes + 1] - sd_state_indices = non_spec_state_indices_tensor device_sd = query_decode.device sd_outputs: list[torch.Tensor] = [] - for i in range(attn_metadata.num_decodes): - ss = cu_sd[i].item() - se = cu_sd[i + 1].item() - si_sd = sd_state_indices[i].item() + # Each decode sequence has 1 token; index _i == token position _i. + # Use tensor-index (no .item()) for CUDA-graph-capture compatibility. + for _i in range(attn_metadata.num_decodes): + _si_sd = non_spec_state_indices_tensor[_i : _i + 1] o_sd, _ = fused_sigmoid_gating_delta_rule_update( A_log=self.A_log, - a=a[ss:se], - b=b[ss:se], + a=a[_i : _i + 1], + b=b[_i : _i + 1], dt_bias=self.dt_bias, - q=query_decode[:, ss:se], - k=key_decode[:, ss:se], - v=value_decode[:, ss:se], - initial_state=ssm_state[si_sd : si_sd + 1], + q=query_decode[:, _i : _i + 1], + k=key_decode[:, _i : _i + 1], + v=value_decode[:, _i : _i + 1], + initial_state=ssm_state, inplace_final_state=True, cu_seqlens=torch.tensor( - [0, se - ss], dtype=torch.int32, device=device_sd - ), - ssm_state_indices=torch.zeros( - 1, dtype=torch.int32, device=device_sd + [0, 1], dtype=torch.int32, device=device_sd ), + ssm_state_indices=_si_sd, use_qk_l2norm_in_kernel=True, ) sd_outputs.append(o_sd) @@ -1614,32 +1609,27 @@ def _forward_core( ) elif attn_metadata.num_decodes > 0: if envs.VLLM_BATCH_INVARIANT: - assert non_spec_query_start_loc is not None assert non_spec_state_indices_tensor is not None - cu_dec = non_spec_query_start_loc[: attn_metadata.num_decodes + 1] - dec_state_indices = non_spec_state_indices_tensor device = query_non_spec.device dec_outputs: list[torch.Tensor] = [] - for i in range(attn_metadata.num_decodes): - s = cu_dec[i].item() - e = cu_dec[i + 1].item() - si = dec_state_indices[i].item() + # Each decode sequence has 1 token; index _i == token position _i. + # Use tensor-index (no .item()) for CUDA-graph-capture compatibility. + for _i in range(attn_metadata.num_decodes): + _si_dec = non_spec_state_indices_tensor[_i : _i + 1] out_i, _ = fused_sigmoid_gating_delta_rule_update( A_log=self.A_log, - a=a[s:e], - b=b[s:e], + a=a[_i : _i + 1], + b=b[_i : _i + 1], dt_bias=self.dt_bias, - q=query_non_spec[:, s:e], - k=key_non_spec[:, s:e], - v=value_non_spec[:, s:e], - initial_state=ssm_state[si : si + 1], + q=query_non_spec[:, _i : _i + 1], + k=key_non_spec[:, _i : _i + 1], + v=value_non_spec[:, _i : _i + 1], + initial_state=ssm_state, inplace_final_state=True, cu_seqlens=torch.tensor( - [0, e - s], dtype=torch.int32, device=device - ), - ssm_state_indices=torch.zeros( - 1, dtype=torch.int32, device=device + [0, 1], dtype=torch.int32, device=device ), + ssm_state_indices=_si_dec, use_qk_l2norm_in_kernel=True, ) dec_outputs.append(out_i) From 42049863d5d11d918cc2c056eb0ee121375e747f Mon Sep 17 00:00:00 2001 From: Yuval Luria Date: Thu, 27 Aug 2026 11:38:35 +0300 Subject: [PATCH 08/14] =?UTF-8?q?gdn:=20fix=20batch=20invariance=20for=20d?= =?UTF-8?q?ecode=20=E2=80=94=20per-token=20projection=20and=20disable=20fu?= =?UTF-8?q?sed=20path?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When VLLM_BATCH_INVARIANT=True and in decode-only mode, GEMM (N sequences) and GEMV (1 sequence) use different CUDA kernel variants with different FP accumulation order. The ~1e-7 difference propagates through in_proj_qkvz and in_proj_ba, then gets amplified through the SSM recurrence (b_h = gate*b_h + beta*v*k^T) to ~4e-5 per decode step. Fix: project each decode token independently (N separate GEMV calls) so the projections match BS=1 behavior exactly. Forward context is used to detect the decode-only batch invariant case with minimal overhead. Also add `not VLLM_BATCH_INVARIANT` guard on use_fused_gdn_decode: the fused norm-packed kernel processes all decode tokens jointly, which is not safe under batch invariance mode. Signed-off-by: Yuval Luria --- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 35 +++++++++++++++++-- 1 file changed, 33 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 77b6f871ba5c..25322d6afa81 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -906,11 +906,42 @@ def forward_cuda( # ============================================================ # Part 1: Input Projection # ============================================================ - mixed_qkvz, _ = self.in_proj_qkvz(hidden_states) - ba, _ = self.in_proj_ba(hidden_states) + # When VLLM_BATCH_INVARIANT and decode-only: project each token via a + # separate GEMV so BS=N matches BS=1. GEMM vs GEMV uses different CUDA + # kernel variants with different FP accumulation order; the ~1e-7 + # difference is amplified by the SSM recurrence to ~4e-5 per step. + _bi_decode = False + if envs.VLLM_BATCH_INVARIANT and num_tokens > 1: + _fc = get_forward_context() + _attn_raw = _fc.attn_metadata + if isinstance(_attn_raw, dict) and self.prefix in _attn_raw: + _meta = _attn_raw[self.prefix] + if isinstance(_meta, GDNAttentionMetadata): + _bi_decode = ( + _meta.num_prefills == 0 and _meta.num_decodes > 0 + ) + if _bi_decode: + mixed_qkvz = torch.cat( + [ + self.in_proj_qkvz(hidden_states[i : i + 1])[0] + for i in range(num_tokens) + ], + dim=0, + ) + ba = torch.cat( + [ + self.in_proj_ba(hidden_states[i : i + 1])[0] + for i in range(num_tokens) + ], + dim=0, + ) + else: + mixed_qkvz, _ = self.in_proj_qkvz(hidden_states) + ba, _ = self.in_proj_ba(hidden_states) use_fused_gdn_decode = ( self.enable_fused_gdn_decode + and not envs.VLLM_BATCH_INVARIANT and hidden_states.dtype == torch.bfloat16 and self.norm.weight.dtype in (torch.bfloat16, torch.float32) ) From 8708dfda9734f2d715f27122a29bbb8a66532991 Mon Sep 17 00:00:00 2001 From: Yuval Luria Date: Thu, 27 Aug 2026 14:45:29 +0300 Subject: [PATCH 09/14] gdn: bypass fast decode path and use per-seq conv1d for batch invariance When VLLM_BATCH_INVARIANT=True: - Skip fused packed-decode path (enable_packed_recurrent_decode) so the per-sequence decode loop is always used, ensuring BS=1 == BS=N. - Run causal_conv1d_fn once per prefill sequence instead of batched, so conv states are identical regardless of batch composition. Co-authored-by: Claude Sonnet 4.6 Signed-off-by: Yuval Luria --- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 51 +++++++++++++++---- 1 file changed, 40 insertions(+), 11 deletions(-) diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 25322d6afa81..667ff9c541ad 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -1322,6 +1322,7 @@ def _forward_core( if ( self.enable_packed_recurrent_decode + and not envs.VLLM_BATCH_INVARIANT # per-seq loop required for BI and attn_metadata.spec_sequence_masks is None and attn_metadata.num_prefills == 0 and attn_metadata.num_decodes > 0 @@ -1403,17 +1404,45 @@ def _forward_core( mixed_qkv_non_spec_T = mixed_qkv_non_spec.transpose(0, 1) # - "cache_indices" updates the conv_state cache in positions # pointed to by "state_indices_tensor" - mixed_qkv_non_spec = causal_conv1d_fn( - mixed_qkv_non_spec_T, - conv_weights, - self.conv1d.bias, - activation=self.activation, - conv_states=conv_state, - has_initial_state=has_initial_state, - cache_indices=non_spec_state_indices_tensor, - query_start_loc=non_spec_query_start_loc, - metadata=attn_metadata, - ).transpose(0, 1) + if envs.VLLM_BATCH_INVARIANT: + # Process each prefill sequence independently so BS=1 and + # BS=N give bitwise-identical per-sequence conv states. + device = mixed_qkv_non_spec_T.device + cu_list = non_spec_query_start_loc.tolist() + chunks = [] + for _pi in range(attn_metadata.num_prefills): + _ps = cu_list[_pi] + _pe = cu_list[_pi + 1] + _chunk_T = mixed_qkv_non_spec_T[:, _ps:_pe] + _has_init = (has_initial_state[_pi : _pi + 1] + if has_initial_state is not None else None) + _cache_idx = non_spec_state_indices_tensor[_pi : _pi + 1] + _cu = torch.tensor([0, _pe - _ps], + dtype=torch.int32, device=device) + _conv_out = causal_conv1d_fn( + _chunk_T, + conv_weights, + self.conv1d.bias, + activation=self.activation, + conv_states=conv_state, + has_initial_state=_has_init, + cache_indices=_cache_idx, + query_start_loc=_cu, + ).transpose(0, 1) + chunks.append(_conv_out) + mixed_qkv_non_spec = torch.cat(chunks, dim=0) + else: + mixed_qkv_non_spec = causal_conv1d_fn( + mixed_qkv_non_spec_T, + conv_weights, + self.conv1d.bias, + activation=self.activation, + conv_states=conv_state, + has_initial_state=has_initial_state, + cache_indices=non_spec_state_indices_tensor, + query_start_loc=non_spec_query_start_loc, + metadata=attn_metadata, + ).transpose(0, 1) elif attn_metadata.num_decodes > 0: assert mixed_qkv_non_spec is not None mixed_qkv_non_spec = causal_conv1d_update( From 831cdcd5973b0a43abcef420e6d7e193116b2a0e Mon Sep 17 00:00:00 2001 From: Yuval Luria Date: Thu, 27 Aug 2026 14:55:45 +0300 Subject: [PATCH 10/14] gdn: fix per-seq conv1d to cover all non-spec sequences in mixed batches non_spec_query_start_loc covers ALL non-spec sequences (both prefill and decode when chunked-prefill mixes them). Previous fix iterated only over num_prefills, causing a size mismatch crash when decode tokens were in the same batch as prefill tokens. Fix: iterate numel()-1 of the cu_seqlens tensor instead of num_prefills. Co-authored-by: Claude Sonnet 4.6 Signed-off-by: Yuval Luria --- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 667ff9c541ad..5995d1c31d45 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -1405,12 +1405,15 @@ def _forward_core( # - "cache_indices" updates the conv_state cache in positions # pointed to by "state_indices_tensor" if envs.VLLM_BATCH_INVARIANT: - # Process each prefill sequence independently so BS=1 and + # Process each non-spec sequence independently so BS=1 and # BS=N give bitwise-identical per-sequence conv states. + # non_spec_query_start_loc covers ALL non-spec sequences + # (both prefill and decode in mixed batches). device = mixed_qkv_non_spec_T.device cu_list = non_spec_query_start_loc.tolist() + num_non_spec_seqs = non_spec_query_start_loc.numel() - 1 chunks = [] - for _pi in range(attn_metadata.num_prefills): + for _pi in range(num_non_spec_seqs): _ps = cu_list[_pi] _pe = cu_list[_pi + 1] _chunk_T = mixed_qkv_non_spec_T[:, _ps:_pe] From a3a1e753d2e4d196dcc64280c76ba80bb175bd9e Mon Sep 17 00:00:00 2001 From: Yuval Luria Date: Thu, 27 Aug 2026 15:40:09 +0300 Subject: [PATCH 11/14] gdn: per-sequence projection for prefill under VLLM_BATCH_INVARIANT Different GEMM M dimensions (BS=1: M=prompt_len vs BS=N: M=total_tokens) cause cublas to select different algorithms with different FP accumulation order, producing ~1e-3 logprob drift amplified by SSM recurrence. Project each prefill sequence independently so M matches the BS=1 case. Only activates for pure-prefill batches (num_decodes==0) to keep the mixed prefill+decode path unchanged. Co-authored-by: Claude Sonnet 4.6 Signed-off-by: Yuval Luria --- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 26 +++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 5995d1c31d45..d02df082ae38 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -911,6 +911,8 @@ def forward_cuda( # kernel variants with different FP accumulation order; the ~1e-7 # difference is amplified by the SSM recurrence to ~4e-5 per step. _bi_decode = False + _bi_prefill_cu = None + _bi_num_prefill_seqs = 0 if envs.VLLM_BATCH_INVARIANT and num_tokens > 1: _fc = get_forward_context() _attn_raw = _fc.attn_metadata @@ -920,6 +922,18 @@ def forward_cuda( _bi_decode = ( _meta.num_prefills == 0 and _meta.num_decodes > 0 ) + # For pure-prefill batches, project per-sequence so the + # GEMM M dimension matches the BS=1 case. Different M + # values cause cublas to select different algorithms with + # different FP accumulation, producing ~1e-3 logprob drift. + if ( + not _bi_decode + and _meta.num_prefills > 0 + and _meta.num_decodes == 0 + and _meta.non_spec_query_start_loc is not None + ): + _bi_prefill_cu = _meta.non_spec_query_start_loc + _bi_num_prefill_seqs = _meta.num_prefills if _bi_decode: mixed_qkvz = torch.cat( [ @@ -935,6 +949,18 @@ def forward_cuda( ], dim=0, ) + elif _bi_prefill_cu is not None: + _cu = _bi_prefill_cu.tolist() + mixed_qkvz = torch.cat( + [self.in_proj_qkvz(hidden_states[_cu[i] : _cu[i + 1]])[0] + for i in range(_bi_num_prefill_seqs)], + dim=0, + ) + ba = torch.cat( + [self.in_proj_ba(hidden_states[_cu[i] : _cu[i + 1]])[0] + for i in range(_bi_num_prefill_seqs)], + dim=0, + ) else: mixed_qkvz, _ = self.in_proj_qkvz(hidden_states) ba, _ = self.in_proj_ba(hidden_states) From 36dd53163b689e4411cbaf51ae18ef0d5c87743c Mon Sep 17 00:00:00 2001 From: Yuval Luria Date: Thu, 27 Aug 2026 19:39:07 +0300 Subject: [PATCH 12/14] fix(gdn): make RMSNormGated batch-invariant under VLLM_BATCH_INVARIANT MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit rmsnorm_fn from layernorm_guard.py uses calc_rows_per_block() which selects ROWS_PER_BLOCK as a Triton constexpr based on M (total rows). Different M values (e.g. BS=1 prefill vs BS=N prefill) compile separate Triton kernel binaries with different FP reduction orders for the row variance sum, producing different per-row results for the same input. When VLLM_BATCH_INVARIANT=True, fall back to the native PyTorch path (forward_native) which uses torch.mean(dim=-1) — a per-row reduction that is independent of total batch size. Fixes 24/32 prompt failures in test_logprobs_bitwise_batch_invariance_bs1_vs_bsN for Qwen3.5-0.8B (GDN_ATTN backend). Co-authored-by: Claude Sonnet 4.6 Signed-off-by: Yuval Luria --- vllm/model_executor/layers/layernorm.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/vllm/model_executor/layers/layernorm.py b/vllm/model_executor/layers/layernorm.py index 4a0d8720e4de..bc08ca09b5f3 100644 --- a/vllm/model_executor/layers/layernorm.py +++ b/vllm/model_executor/layers/layernorm.py @@ -286,6 +286,12 @@ def forward_native( def forward_cuda( self, x: torch.Tensor, z: torch.Tensor | None = None ) -> torch.Tensor: + if envs.VLLM_BATCH_INVARIANT: + # rmsnorm_fn uses calc_rows_per_block which selects ROWS_PER_BLOCK + # based on M. Different M values compile different Triton binaries + # with different FP reduction orders, breaking batch invariance. + # Fall back to the native PyTorch path which is row-independent. + return self.forward_native(x, z) from vllm.third_party.flash_linear_attention.ops.layernorm_guard import ( rmsnorm_fn, ) From 2357f1ce403ea9888e67f455d6318c44e364a868 Mon Sep 17 00:00:00 2001 From: Yuval Luria Date: Sat, 29 Aug 2026 09:57:10 +0300 Subject: [PATCH 13/14] =?UTF-8?q?gdn:=20address=20review=20feedback=20?= =?UTF-8?q?=E2=80=94=20mixed=20batches,=20platform=20guard,=20test=20fix?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit [P1] Unify per-request projection for all batch types via non_spec_query_start_loc. The previous code only handled pure-decode and pure-prefill; mixed batches fell through to a single batched GEMM, breaking batch invariance. The unified loop covers decode (1-token slices), prefill (seq-len slices), and mixed batches uniformly. Speculative decoding explicitly raises RuntimeError. [P2] Restrict supports_batch_invariance() to NVIDIA CUDA only. The ROCm AITER and XPU forward paths are unmodified and not batch-invariant. [P2] Remove GDN_ATTN from the default CUDA backend list in test utils. GDN_ATTN is now only added when the test model actually contains GDN layers (model_type="qwen3_5" or dual_chunk_attention_config present). The model-type check is now unconditional, not gated on VLLM_TEST_MODEL. Signed-off-by: Yuval Luria --- tests/v1/determinism/utils.py | 56 ++++++++--------- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 60 ++++++------------- vllm/v1/attention/backends/gdn_attn.py | 6 +- 3 files changed, 48 insertions(+), 74 deletions(-) diff --git a/tests/v1/determinism/utils.py b/tests/v1/determinism/utils.py index e8813fbe6093..c56cceb56a07 100644 --- a/tests/v1/determinism/utils.py +++ b/tests/v1/determinism/utils.py @@ -28,7 +28,10 @@ class DeviceConfig(NamedTuple): and current_platform.has_device_capability(80), # FlashInfer backend temporarily disabled due to invariant CTA sizes. # See FlashInfer issue #2424 - backends=["FLASH_ATTN", "TRITON_ATTN", "FLEX_ATTENTION", "GDN_ATTN"], + # GDN_ATTN is excluded from the default list: it is only valid for + # models with GDN layers (model_type="qwen3_5" or dual_chunk_attention). + # The model-specific override below adds it when appropriate. + backends=["FLASH_ATTN", "TRITON_ATTN", "FLEX_ATTENTION"], ), "xpu": DeviceConfig( available=current_platform.is_xpu() and HAS_TRITON, @@ -39,35 +42,28 @@ class DeviceConfig(NamedTuple): DEFAULT_MODEL = "Qwen/Qwen3-1.7B" TEST_MODEL = os.getenv("VLLM_TEST_MODEL", DEFAULT_MODEL) -# Override backends for MLA models (MLA only supported on CUDA). -if os.getenv("VLLM_TEST_MODEL"): - config = get_config(TEST_MODEL, trust_remote_code=False) - if ModelArchConfigConvertorBase(config, config.get_text_config()).is_deepseek_mla(): - DEVICE_BACKENDS["cuda"] = DeviceConfig( - available=DEVICE_BACKENDS["cuda"].available, - backends=["TRITON_MLA"] - + (["FLASH_ATTN_MLA"] if flash_attn_supports_mla() else []), - ) - DEVICE_BACKENDS["xpu"] = DeviceConfig( - available=DEVICE_BACKENDS["xpu"].available, - backends=[], - ) - # GDN_ATTN is for Qwen3.5 models (model_type="qwen3_5") and - # Qwen3-Next/Qwen3.6 hybrid models (dual_chunk_attention_config present) - elif getattr(config, "model_type", "") == "qwen3_5" or ( - hasattr(config, "dual_chunk_attention_config") - and config.dual_chunk_attention_config is not None - ): - DEVICE_BACKENDS["cuda"] = DeviceConfig( - available=DEVICE_BACKENDS["cuda"].available, - backends=["GDN_ATTN"], - ) - else: - # Remove GDN_ATTN for models that don't have GDN architecture - DEVICE_BACKENDS["cuda"] = DeviceConfig( - available=DEVICE_BACKENDS["cuda"].available, - backends=[b for b in DEVICE_BACKENDS["cuda"].backends if b != "GDN_ATTN"], - ) +# Override backends based on the model architecture. +config = get_config(TEST_MODEL, trust_remote_code=False) +if ModelArchConfigConvertorBase(config, config.get_text_config()).is_deepseek_mla(): + DEVICE_BACKENDS["cuda"] = DeviceConfig( + available=DEVICE_BACKENDS["cuda"].available, + backends=["TRITON_MLA"] + + (["FLASH_ATTN_MLA"] if flash_attn_supports_mla() else []), + ) + DEVICE_BACKENDS["xpu"] = DeviceConfig( + available=DEVICE_BACKENDS["xpu"].available, + backends=[], + ) +# GDN_ATTN is for Qwen3.5 models (model_type="qwen3_5") and +# Qwen3-Next hybrid models (dual_chunk_attention_config present). +elif getattr(config, "model_type", "") == "qwen3_5" or ( + hasattr(config, "dual_chunk_attention_config") + and config.dual_chunk_attention_config is not None +): + DEVICE_BACKENDS["cuda"] = DeviceConfig( + available=DEVICE_BACKENDS["cuda"].available, + backends=["GDN_ATTN"], + ) # Only include backends for devices that are actually available. BACKENDS: list[str] = sorted( diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index d02df082ae38..f9685f352790 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -906,59 +906,33 @@ def forward_cuda( # ============================================================ # Part 1: Input Projection # ============================================================ - # When VLLM_BATCH_INVARIANT and decode-only: project each token via a - # separate GEMV so BS=N matches BS=1. GEMM vs GEMV uses different CUDA - # kernel variants with different FP accumulation order; the ~1e-7 - # difference is amplified by the SSM recurrence to ~4e-5 per step. - _bi_decode = False - _bi_prefill_cu = None - _bi_num_prefill_seqs = 0 + # When VLLM_BATCH_INVARIANT: project each request independently so the + # GEMM M dimension is identical to the BS=1 case. This covers pure + # decode, pure prefill, and mixed batches uniformly via + # non_spec_query_start_loc. Speculative decoding is not yet supported. + _bi_cu: list[int] | None = None if envs.VLLM_BATCH_INVARIANT and num_tokens > 1: _fc = get_forward_context() _attn_raw = _fc.attn_metadata if isinstance(_attn_raw, dict) and self.prefix in _attn_raw: _meta = _attn_raw[self.prefix] if isinstance(_meta, GDNAttentionMetadata): - _bi_decode = ( - _meta.num_prefills == 0 and _meta.num_decodes > 0 - ) - # For pure-prefill batches, project per-sequence so the - # GEMM M dimension matches the BS=1 case. Different M - # values cause cublas to select different algorithms with - # different FP accumulation, producing ~1e-3 logprob drift. - if ( - not _bi_decode - and _meta.num_prefills > 0 - and _meta.num_decodes == 0 - and _meta.non_spec_query_start_loc is not None - ): - _bi_prefill_cu = _meta.non_spec_query_start_loc - _bi_num_prefill_seqs = _meta.num_prefills - if _bi_decode: - mixed_qkvz = torch.cat( - [ - self.in_proj_qkvz(hidden_states[i : i + 1])[0] - for i in range(num_tokens) - ], - dim=0, - ) - ba = torch.cat( - [ - self.in_proj_ba(hidden_states[i : i + 1])[0] - for i in range(num_tokens) - ], - dim=0, - ) - elif _bi_prefill_cu is not None: - _cu = _bi_prefill_cu.tolist() + if _meta.num_spec_decodes > 0: + raise RuntimeError( + "VLLM_BATCH_INVARIANT is not supported with " + "speculative decoding on GDN_ATTN." + ) + if _meta.non_spec_query_start_loc is not None: + _bi_cu = _meta.non_spec_query_start_loc.tolist() + if _bi_cu is not None: mixed_qkvz = torch.cat( - [self.in_proj_qkvz(hidden_states[_cu[i] : _cu[i + 1]])[0] - for i in range(_bi_num_prefill_seqs)], + [self.in_proj_qkvz(hidden_states[_bi_cu[i] : _bi_cu[i + 1]])[0] + for i in range(len(_bi_cu) - 1)], dim=0, ) ba = torch.cat( - [self.in_proj_ba(hidden_states[_cu[i] : _cu[i + 1]])[0] - for i in range(_bi_num_prefill_seqs)], + [self.in_proj_ba(hidden_states[_bi_cu[i] : _bi_cu[i + 1]])[0] + for i in range(len(_bi_cu) - 1)], dim=0, ) else: diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 8b28583d8247..11cf2dda62a9 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -39,7 +39,11 @@ def is_ssm(cls) -> bool: @classmethod def supports_batch_invariance(cls) -> bool: - return True + # Only implemented for NVIDIA CUDA. ROCm AITER and XPU paths still + # use batch-shaped projections and the fused norm kernel, which are + # not batch-invariant. + import torch + return torch.cuda.is_available() and torch.version.hip is None @dataclass From 2ffb534dd036cb61eb8f38d98c3ee0a64e863327 Mon Sep 17 00:00:00 2001 From: bfoing <40759640+bfoing@users.noreply.github.com> Date: Tue, 1 Sep 2026 19:09:24 +0100 Subject: [PATCH 14/14] gdn: make batch-invariant decode path CUDA-graph-capturable VLLM_BATCH_INVARIANT=1 required --enforce-eager on GDN models: capture aborted with "Cannot copy between CPU and CUDA tensors during CUDA graph capture unless the CPU tensor is pinned". Two host/device syncs on the decode path caused it. 1. The input projection called non_spec_query_start_loc.tolist() (D2H) to drive the per-sequence loop. For a pure-decode batch each non-spec sequence contributes exactly one token, so the boundaries are [0, 1, ... num_tokens] and follow from the token count with no device read; the trip count has to be static for capture anyway. Prefill and mixed batches keep the existing path, since they are not captured and their lengths genuinely vary. 2. Both per-sequence decode loops rebuilt cu_seqlens with torch.tensor([0, 1], device=...) on every iteration, an H2D copy from pageable memory and also num_decodes * num_layers redundant copies per step. It is now a cached buffer allocated during warmup. The decode loops were already written capture-safe; this removes the two syncs in front of them. Measured on H100 with Qwen/Qwen3.6-35B-A3B-FP8: invariance still holds bitwise (12/12) under capture, throughput goes from 84.97 to 638.09 tok/s, and the cost over a non-invariant baseline drops from 16.1x to 1.45x. Signed-off-by: bfoing <40759640+bfoing@users.noreply.github.com> --- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 31 ++++++++++++++----- 1 file changed, 24 insertions(+), 7 deletions(-) diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index f9685f352790..e999f151b677 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -795,6 +795,19 @@ def prepare_gdn_attention_core_inputs( return mixed_qkv_out, z_out, b_out, a_out + def _bi_cu01(self, device: torch.device) -> torch.Tensor: + """Cached [0, 1] cu_seqlens for per-sequence decode under + VLLM_BATCH_INVARIANT. Building this with torch.tensor() inside the + decode loop is an H2D copy from pageable memory, which CUDA graph + capture rejects; it is also num_decodes * num_layers pointless copies + per step. Populated during warmup, before any capture begins. + """ + t = getattr(self, "_bi_cu01_cache", None) + if t is None or t.device != device: + t = torch.tensor([0, 1], dtype=torch.int32, device=device) + self._bi_cu01_cache = t + return t + def rearrange_mixed_qkv(self, mixed_qkv): """Split packed qkv into contiguous (1, seq, heads, dim) tensors. @@ -922,7 +935,15 @@ def forward_cuda( "VLLM_BATCH_INVARIANT is not supported with " "speculative decoding on GDN_ATTN." ) - if _meta.non_spec_query_start_loc is not None: + if _meta.num_prefills == 0: + # Pure decode: every non-spec sequence contributes + # exactly one token, so the cumulative boundaries are + # [0, 1, ..., num_tokens]. Derive them from the token + # count rather than syncing non_spec_query_start_loc to + # the host: a D2H copy is illegal during CUDA graph + # capture, and the trip count must be static anyway. + _bi_cu = list(range(num_tokens + 1)) + elif _meta.non_spec_query_start_loc is not None: _bi_cu = _meta.non_spec_query_start_loc.tolist() if _bi_cu is not None: mixed_qkvz = torch.cat( @@ -1572,9 +1593,7 @@ def _forward_core( v=value_decode[:, _i : _i + 1], initial_state=ssm_state, inplace_final_state=True, - cu_seqlens=torch.tensor( - [0, 1], dtype=torch.int32, device=device_sd - ), + cu_seqlens=self._bi_cu01(device_sd), ssm_state_indices=_si_sd, use_qk_l2norm_in_kernel=True, ) @@ -1689,9 +1708,7 @@ def _forward_core( v=value_non_spec[:, _i : _i + 1], initial_state=ssm_state, inplace_final_state=True, - cu_seqlens=torch.tensor( - [0, 1], dtype=torch.int32, device=device - ), + cu_seqlens=self._bi_cu01(device), ssm_state_indices=_si_dec, use_qk_l2norm_in_kernel=True, )