Add: support 1M context in DeepSeek V4 SWA decode

[sjduan]

• Pin MAX_SEQ_LEN to 1,048,576 as the SWA-local context ceiling
  instead of M.max_position_embeddings, leaving the global model
  configuration untouched.
• Take freqs_cos and freqs_sin as token-local [T_DYN, ROPE_HEAD_DIM]
  tensors on decode_swa, decode_swa_tp1 and l3_decode_swa, binding
  their token axis dynamic on every test and L3 entry.
• Drop the on-device swa_rope_step gather from both decode paths and
  feed freqs_cos and freqs_sin straight into qkv_proj_rope and
  sparse_attn_swa; position_ids stays in the ABI for host admission
  and golden semantics.
• Build the RoPE rows for the active positions inside
  build_tensor_specs from rope_theta, replacing the full-context
  build_rope_tables table, and cache them across the cos and sin
  initializers.
• Size the logical block table to the 1M ceiling so high positions
  wrap into the fixed physical pool.
• Accept a scalar or comma-separated --start-pos on the full and tp1
  entries, deriving the batch from the list length and rejecting
  negative positions or ones past MAX_SEQ_LEN; the output-only entry
  rejects the flag.
• Thread start_pos through build_distributed_tensor_specs so the
  distributed specs share the single-rank fixture path, and size the
  kv_cache comparator mapping to the actual token count.
• Drop the per-position gather from golden_decode_swa_tp1, which now
  consumes the token-local RoPE tensors directly.
• Add --enable-l2-swimlane to the full, output and tp1 entries.

Validated on A2/A3 devices at TP1 and TP2 for 128, 16K and 1M local
contexts plus a mixed (16K, 16K, 128K, 512K) workload at S=8, with
both kv_cache and x_out passing. Related to #962.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 31b9311c-baab-4e65-abf1-53d748b9ae84

📥 Commits

Reviewing files that changed from the base of the PR and between 1525587 and da1dd43.

📒 Files selected for processing (1)

• models/deepseek_v4_flash_dspark/decode_swa.py

Included review availability: Your plan provides up to 1 included review per hour; 0 remain after this review.

---

📝 Walkthrough

Walkthrough

SWA decode now uses token-local dynamic RoPE tensors. The harness computes active-position rows, supports mixed start positions up to 1M positions, validates bounds, and forwards runtime options through TP1 and distributed execution paths.

Changes

SWA token-local RoPE flow

Layer / File(s)	Summary
Token-local RoPE device ABI models/deepseek_v4_flash_dspark/decode_swa.py	Decode, distributed, and TP1 functions now accept dynamic RoPE tensors and pass them directly to projection and attention operations.
RoPE fixture and reference generation models/deepseek_v4_flash_dspark/decode_swa.py	Fixtures compute and cache cosine and sine rows for active positions. Golden references consume these rows directly.
Dynamic tensor specifications and positions models/deepseek_v4_flash_dspark/decode_swa.py	Tensor specifications use token-shaped RoPE dimensions and propagate start positions across distributed ranks.
Position-aware harness execution models/deepseek_v4_flash_dspark/decode_swa.py	The harness parses start positions and L2 swimlane settings, validates entry-point combinations, and forwards runtime configuration.
Estimated code review effort: 4 (Complex)	~45 minutes

Merge Risk: ⚪ Minimal · up to da1dd

The PR is merge-ready after normal checks and review; no actionable merge-blocking risk remains.

Possibly related issues

• [Feature][P2] Enable elastic 1M-context decode attention with block size 32 #962 — Directly covers token-local RoPE inputs and 1M-context SWA position support.

Possibly related PRs

• hw-native-sys/pypto-lib#903 — Updates the same SWA decode RoPE handling and tensor-spec fixtures.
• hw-native-sys/pypto-lib#910 — Adds related dynamic tensor dimensions and batch handling in decode_swa.py.
• hw-native-sys/pypto-lib#987 — Modifies the same decode functions and tensor specifications refined here.

Sequence Diagram(s)

sequenceDiagram
  participant CLI
  participant build_tensor_specs
  participant init_rope_rows
  participant decode_swa_tp1
  CLI->>build_tensor_specs: pass start_pos and batch
  build_tensor_specs->>init_rope_rows: compute active-position RoPE rows
  init_rope_rows-->>build_tensor_specs: return token-local cosine and sine rows
  build_tensor_specs-->>decode_swa_tp1: provide dynamic tensor specifications
  CLI->>decode_swa_tp1: pass start_pos and enable_l2_swimlane

Loading

Poem

A rabbit hops through rows of light,
With RoPE tokens trimmed just right.
Start positions lead the way,
Through million-step blocks they play.
Dynamic tensors bloom and flow—
Swiмlane winds make kernels glow.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly summarizes the primary change: adding 1M-token context support to DeepSeek V4 SWA decode.
Description check	✅ Passed	The description directly explains the 1M-context support, token-local RoPE changes, start-position handling, and validation results.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}