[Qwen3Next] preserve linear-attn-mask optimization under torch.compile/export

[yuvrajsharma9981]

Hi,

`torch.export.export` fails on Qwen3Next-family models with `GuardOnDataDependentSymNode: Could not guard on data-dependent expression Eq(u0, 1)`. The crash traces to `Qwen3NextModel._update_linear_attn_mask`:

```python
if (past_key_values is not None and past_key_values.has_previous_state()) or (
attention_mask is not None and torch.all(attention_mask == 1)
):
linear_attn_mask = None
```

`torch.all(attention_mask == 1)` produces a 0-dim bool tensor, and Python's `if` does an implicit `.item()` on it — an unbacked symbolic int the exporter can't resolve. Net effect: any user wanting an AOT package (`torch._inductor.aoti_compile_and_package` → `.pt2`) for any model in this family is blocked at the export step.

I tripped on this trying to AOT compile Qwen3.5 for fast serving — the eager forward works, the export step crashes.

Scope

Fix lands at the modular source-of-truth, so the same patch propagates to all four models that inherit `Qwen3NextModel._update_linear_attn_mask`:

• Qwen3Next (direct)
• Qwen3.5 (`Qwen3_5TextModel(Qwen3NextModel)`)
• Qwen3.5-MoE (`Qwen3_5MoeTextModel` via the same lineage)
• OLMo Hybrid (`OlmoHybridModel(Qwen3NextModel)`)

Fix

Smallest behavior-preserving thing I could come up with: keep the eager-mode fast-path identical, and skip the data-dependent branch only when `torch.compiler.is_compiling()` is true. The downstream linear-attention layer treats an all-1s mask as a cheap no-op, so the exported graph runs correctly for the no-padding case that the eager path was short-circuiting.

```python
def _update_linear_attn_mask(self, attention_mask, past_key_values):
linear_attn_mask = attention_mask
if past_key_values is not None and past_key_values.has_previous_state():
return None
if torch.compiler.is_compiling():
return linear_attn_mask
if attention_mask is not None and torch.all(attention_mask == 1):
linear_attn_mask = None
return linear_attn_mask
```

Two notes on the ordering:

1. The cached-forward check stays first so users exporting a decode-step graph still get the cached-skip optimization baked into the resulting graph — that branch is already export-compatible (Python object state, not a tensor `.item()`).
2. `torch.compiler.is_compiling()` is the public PyTorch idiom for "behave differently under trace"; runtime behavior for everyone not exporting is byte-identical to before.

Reproducer

Fails on v5.9.0 + torch 2.11:

```python
import torch
from transformers import AutoModelForCausalLM

m = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3.5-4B", torch_dtype=torch.bfloat16)
m.eval()

class W(torch.nn.Module):
def init(s, m): super().init(); s.m = m
def forward(s, ids, mask):
return s.m(input_ids=ids, attention_mask=mask).logits

ids = torch.ones(2, 128, dtype=torch.long)
mask = torch.ones(2, 128, dtype=torch.long)

torch.export.export(W(m), (ids, mask), dynamic_shapes={
"ids": {0: torch.export.Dim.AUTO, 1: torch.export.Dim.AUTO},
"mask": {0: torch.export.Dim.AUTO, 1: torch.export.Dim.AUTO},
})
```

After the change, verified locally with a forked-source install:

• eager runtime with all-1s mask still returns `None` (existing optimization preserved)
• `torch.export.export(...)` succeeds and traces a clean graph

Commits

• First commit edited the generated `modeling_qwen3_5.py` directly — CI correctly flagged this via `check_repository_consistency`.
• Second commit moves the fix to `modular_qwen3_next.py` (the source-of-truth) and regenerates the affected `modeling_*.py` files via `make fix-repo`. Both checks pass locally now.

Happy to add tests under `tests/models/qwen3_next/` (and the inheriting models) if that's the preferred shape — held off pending guidance on existing export-compat coverage conventions.

Thanks!

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: olmo_hybrid, qwen3_5, qwen3_5_moe, qwen3_next

[Rocketknight1]

cc @ArthurZucker for text models!