fix(dsv4): zero inactive prefill MoE tails

[high-cloud]

Summary

• zero inactive HC tail rows before each prefill MoE call
• keep the fix scoped to prefill_fwd.py; no MoE, HC head, RoPE, or metadata changes
• prevents inactive attention tail rows from being consumed by MoE shared/gate compute when num_tokens < T

Validation

• python -m py_compile models/deepseek/v4/prefill_fwd.py
• git diff --check -- models/deepseek/v4/prefill_fwd.py
• compile-only: python models/deepseek/v4/prefill_fwd.py -p a2a3 --ep 2 -d 0,1 --num-tokens 4 --compile-only
• EP2 NPU: task_20260709_183916_13296387544, devices 6,7, runtime done (78.38s), PASS
• EP8 NPU: task_20260709_185059_17335683456, devices 8,9,10,11,12,13,14,15, runtime done (266.87s), PASS

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 4b383ec8-4259-4bd3-92de-9154a5260540

📥 Commits

Reviewing files that changed from the base of the PR and between f406b34 and 046e250.

📒 Files selected for processing (1)

• models/deepseek/v4/prefill_fwd.py

---

📝 Walkthrough

Walkthrough

DeepSeek V4 prefill now explicitly zeros inactive token tail blocks along the hidden-channel dimension after every attention stage and before the corresponding MoE stage. Exported JIT entrypoint signatures and attention/MoE arguments remain unchanged.

Changes

DeepSeek V4 prefill tail handling

Layer / File(s)	Summary
Tail-zeroing helper models/deepseek/v4/prefill_fwd.py	Adds _zero_hc_tail, constants for tail chunking, token-count clamping, and SPMD-based zeroing of inactive hidden-channel blocks.
Attention output integration models/deepseek/v4/prefill_fwd.py	Applies tail zeroing after layer-0, layer-1, CSA, HCA, and final attention outputs before MoE consumption.

Estimated code review effort: 2 (Simple) | ~10 minutes

Possibly related PRs

• hw-native-sys/pypto-lib#613: Both changes modify the DeepSeek V4 prefill path where attention outputs are passed into MoE.

Poem

A bunny found some tails astray,
And zeroed the unused bits away.
Each attention hop now leaves them neat,
Before MoE gets its tensor treat.
Hop, hop—clean data’s on its way!

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: zeroing inactive prefill MoE tails.
Description check	✅ Passed	The description accurately matches the scoped fix and validation steps in the changeset.
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.

---

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 introduces a _zero_hc_tail helper function in models/deepseek/v4/prefill_fwd.py to zero out the tail of the hidden states tensor for inactive tokens during prefill forward passes, applying it after various attention layers. The review feedback recommends adding an explanatory comment to a self-assignment statement within this function, which appears to be a compiler workaround to ensure the tensor is materialized.

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]

This no-op assignment appears to be a compiler workaround, likely to act as a keep-alive for x_hc and ensure it's materialized before the loop. For improved maintainability, please add a comment explaining its purpose.

Suggested change

	x_hc[0:1, 0:1, 0:16] = x_hc[0:1, 0:1, 0:16]
	# Keep-alive for x_hc to ensure it's materialized before use in the loop.
	x_hc[0:1, 0:1, 0:16] = x_hc[0:1, 0:1, 0:16]