refactor(dsv4): port #734 rope_cs collapse to HCA/SWA decode sparse kernels

[wangqin1723-max]

Summary

Ports the small sparse-kernel piece of CSA's #734 refactor to the HCA & SWA decode paths: collapse rope_cs from a 2-task spmd fan-out into a single pl.at, matching what #734 did in CSA's decode_sparse_attn.py. Structural parity only — bit-exact, wall-neutral.

Changes

• decode_sparse_attn_hca.py / decode_sparse_attn_swa.py: rope_cs pl.spmd(HALF_ROPE // ROPE_TILE) (2 tasks) → single pl.at(CORE_GROUP) operating on full ranges; drop the ROPE_TILE / ROPE_INTERLEAVE_TILE constants. Bit-identical math, just un-tiled.
• decode_sparse_attn_swa.py: add allow_early_resolve=True to swa_valid_bias (parity with HCA build_valid and CSA's fused plan).

Validation (a2a3, real device)

Both bit-exact PASS:

• decode_attention_hca.py e2e: x_out / kv_cache PASS.
• decode_attention_swa.py e2e: x_out / kv_cache PASS.

Scope notes

• Intentionally not included: the compressor_ratio128 task-count refactor (Refactor: cut runtime task count in DeepSeek-V4 decode CSA flow #734's scatter/pool/rmsnorm fusion + RMS_TILE drop). A same-session A/B showed it is wall-neutral at standalone decode resolution, so it was dropped to keep this PR minimal. The real HCA/SWA decode long poles (qk_pv, swa_gather_kv) are compute/bandwidth-bound and are being pursued separately.
• This PR is structural parity (rope_cs form), wall-neutral; no perf claim.

[coderabbitai]

📝 Walkthrough

Walkthrough

The changes simplify compressor pooling and cache finalization, and replace tiled inverse-RoPE precomputation with full-tensor core-group operations in HCA and SWA decode kernels. SWA valid-bias token processing also enables early resolution.

Changes

DeepSeek decode kernel updates

Layer / File(s)	Summary
Compressor pooling and fused finalization models/deepseek/v4/decode_compressor_ratio128.py	Pooling uses a fixed 16-row padded tile, while RMSNorm, RoPE, and conditional cache writes are fused into one core-group stage.
HCA inverse-RoPE precompute models/deepseek/v4/decode_sparse_attn_hca.py	Full inverse-RoPE cosine and signed-sine tensors are computed in one core-group region instead of tiled SPMD loops.
SWA validation and inverse-RoPE precompute models/deepseek/v4/decode_sparse_attn_swa.py	Valid-bias token processing allows early resolution, and inverse-RoPE tensors are generated in one full-tensor core-group region.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Poem

I’m a bunny with tiles in a row,
Watching pooled values neatly flow.
RoPE gathers as one,
Cache writes are done,
And early resolves make kernels go!

🚥 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 matches the main rope_cs refactor, though it omits the compressor and SWA bias changes.
Description check	✅ Passed	The description is directly about the same HCA/SWA sparse-kernel refactor 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.}

[gemini-code-assist]

Code Review

This pull request refactors the RoPE and RMSNorm computation blocks across several DeepSeek v4 decode models, fusing tasks into core-group level blocks and simplifying indexing and padding logic. However, in decode_compressor_ratio128.py, the changes incorrectly replace the dynamic batch size b_dim with the static constant B when slicing cos/sin tensors and when looping for cache writes. This breaks dynamic batch support and can lead to runtime out-of-bounds errors or incorrect behavior for varying batch sizes.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

The compressor function compressor_ratio128 is designed to support dynamic batch sizes via b_dim (bound to B_DYN). Hardcoding the static constant B (which is 8) when slicing cos and sin will cause runtime out-of-bounds errors if the actual batch size at runtime is smaller than B, or ignore active rows if the batch size is larger. Use b_dim instead of B to preserve dynamic batch support.

Suggested change

	cos_b[0:B, 0 : ROPE_HEAD_DIM // 2] = cos[0:B, 0 : ROPE_HEAD_DIM // 2]
	sin_b[0:B, 0 : ROPE_HEAD_DIM // 2] = sin[0:B, 0 : ROPE_HEAD_DIM // 2]
	cos_b[0:b_dim, 0 : ROPE_HEAD_DIM // 2] = cos[0:b_dim, 0 : ROPE_HEAD_DIM // 2]
	sin_b[0:b_dim, 0 : ROPE_HEAD_DIM // 2] = sin[0:b_dim, 0 : ROPE_HEAD_DIM // 2]

References

1. In dynamic kernels, do not replace dynamic loop bounds with static padding constants, as it can cause over-iteration past the dynamic batch size.

[gemini-code-assist]

Similarly, the loop for writing to the cache should iterate over the dynamic batch size b_dim instead of the static constant B to ensure correct behavior for dynamic batch sizes.

Suggested change

	for inner in pl.range(B):
	for inner in pl.range(b_dim):

References

1. In dynamic kernels, do not replace dynamic loop bounds with static padding constants, as it can cause over-iteration past the dynamic batch size.