[Rotary] Separate varlen lengths from cache capacity

[taking-lying-flat]

Summary

Packed varlen currently overloads one max_seqlen value with three different meanings: the RoPE cache capacity, the batch's logical maximum sequence length, and the configured context limit. This can make a 4K/32K packed batch allocate and launch as if it were 256K.

This PR:

• validates scalar varlen rotary offsets against each segment length instead of the packed token total;
• adds a cached get_max_seqlen(cu_seqlens, cu_seqlens_cpu) helper;
• separates batch_max_seqlen from rope_cache_length across rotary-using attention layers;
• passes the logical batch maximum to FlashAttention, DeltaFormer, NSA, and MoBA varlen paths;
• removes per-layer Python scalar synchronization from left-padded decode by retaining the safe pre-adjustment upper bound;
• removes the redundant physical batch axis from CUDA and Ascend varlen rotary grids;
• covers the scalar validation, physical B > 1 grid, logical cache capacity, and FlashAttention launch arguments with regression tests.

Implementation checklist

• Scalar packed-varlen validation uses max(diff(cu_seqlens)), not packed T.
• RoPE cache capacity is independent from FlashAttention's logical batch maximum.
• Left-padding decode no longer evaluates max() / .item() on per-sample CUDA offsets in each layer.
• Varlen rotary launches one physical batch program per (token chunk, head).
• Existing dense behavior and configured-context decode preallocation are preserved.

Test plan

Hardware and environment: NVIDIA RTX A1000 Laptop GPU (4 GiB), PyTorch 2.13.0+cu130, CUDA 13.x, locally rebuilt flash-attn==2.8.3.post1.

• pytest -q tests/modules/test_rotary.py tests/ops/utils/test_index.py tests/layers/test_attn_varlen_pack_layout.py
  • 174 passed
• Cache/decode coverage across Attention, BitAttention, MLA, MoBA, NSA, DeltaFormer, MultiScaleRetention, Parallax, Raven, and Rodimus
  • 40 passed
• tests/ops/test_attn.py and tests/ops/test_moba.py
  • 36 passed
• Affected model suites run separately to fit the 4 GiB GPU
  • 56 passed, 14 skipped
• Real flash_attn_varlen_func causal forward/backward smoke test
  • finite output and finite Q/K/V gradients
• pre-commit run --files <all changed files>
  • all hooks passed

The full large naive-reference configurations in tests/ops/test_deltaformer.py and tests/ops/test_nsa.py exceed this GPU's 4 GiB capacity. Their smaller related packed-varlen forward/backward cases pass. Ascend changes were statically reviewed and linted but could not be executed without NPU hardware.

Benchmark / NCU (kernel changes only)

Same RTX A1000 Laptop GPU before/after. NCU was not collected.

Packed prefill/cache capacity

BF16, rotary dim 128, eight 4K segments, 32K packed tokens, configured context 256K:

Metric	Baseline	This PR	Change
RoPE cache	64.0 MiB	1.0 MiB	-98.4%
Peak allocated memory	233.251 MiB	57.253 MiB	-75.5%
FlashAttention max length	262144	4096	64x smaller
Cold elapsed time	430.68 ms	442.69 ms	within compilation/startup noise

Left-padded decode synchronization

Real FlashAttention, FP16, 8 Attention layers, hidden size 128, 128-token prefill, 64 measured decode steps:

Batch	Baseline median	This PR median	Latency reduction	Throughput increase
1	4.2032 ms	3.9178 ms	6.8%	7.3%
8	4.7711 ms	3.8538 ms	19.2%	23.8%
32	6.4738 ms	4.0056 ms	38.1%	61.6%

This intentionally small decode workload makes host synchronization visible. Larger models will spend a greater fraction in GEMMs, so the percentage improvement should not be extrapolated directly. The removed scalar synchronization itself measured approximately 13.8 / 76.2 / 286.7 microseconds per affected layer at batch sizes 1 / 8 / 32.

Dense rotary math and its grid are unchanged; the physical-grid change applies only when cu_seqlens is present.

Breaking changes

None. The new DeltaFormer maximum-length argument is optional, and the MLA attention_mask default is backward-compatible.

Checklist

• I have read CONTRIBUTING.md and follow its conventions (code style, docstrings, commit prefixes).
• I have read AGENTS.md and, where my change matches its scope, the relevant skill under .agents/skills.
• Dependent tests pass locally or in CI, and new behavior is covered by tests where applicable (tick as N/A for changes with no testable code, e.g. docs-only).
• Kernel changes include same-hardware before/after benchmark numbers, dense + varlen where applicable (tick as N/A when no kernel code changed).
• This PR is minor/cosmetic-only (typo, formatting, style-only tweaks) — tick only if it is, and justify below.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-08-30T15:06:05.905787Z	4ba9219	PR opened

ℹ️ 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" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[zhiyuan1i]

Style issues:

1. P1: 10 layers have the dead branch rope_cache_length = (q_len + seqlen_offset if cu_seqlens is None else rope_cache_length + seqlen_offset) — both arms equal rope_cache_length + seqlen_offset since initialization already branches on cu_seqlens. All 10 should be rope_cache_length += seqlen_offset.

2. P2: rotary.py:188 and NPU rotary.py:163 hand-write cu_seqlens[1:] - cu_seqlens[:-1] — this PR just unified prepare_lens(cu_seqlens), reuse it.

3. P2: nsa.py:97 — batch_max_seqlen is single-use (nsa doesn't go through flash attn), inline it like parallax/yoco.

4. P2: moba.py:177-180 — rotary_cu_seqlens uses ternary, cu_seqlens_cpu uses if — pick one style.

5. P2: multiscale_retention.py:233/239, yoco.py:180/186 — adjacent lines use q.shape[1] and q_len interchangeably, unify to q_len.