Fix: clear LM-head signals after forward

[high-cloud]

• Clear local hidden and logits completion counters after final gather
• Let retained CommDomains safely reuse the fixed done epoch

[coderabbitai]

📝 Walkthrough

Walkthrough

The lm_head function now clears per-peer hidden and logits completion counters after gathering logits, allowing the communication domain to reuse the same completion epoch on the next forward pass.

Changes

LM Head Signal Reset

Layer / File(s)	Summary
Clear completion counters models/deepseek/v4-flash/lm_head.py	Adds a core-group scope that writes zero to hidden_done and logits_done for each src_tp after logits gathering.

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

Poem

I’m a bunny with counters to clear,
Zeroing signals so paths reappear.
Logits hop through the gate,
Then epochs reset their state—
A tidy new forward draws near.

🚥 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 change: clearing LM-head signals after forward.
Description check	✅ Passed	The description directly describes clearing completion counters and reusing the done epoch.

---

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.}