[Bugfix] Fix DeepSeek V4 mHC broadcast buffer for weight sync

[HollowMan6]

Purpose

Fix Stale broadcast weights after refit. finalize_mhc_broadcast_weights() rebound layer.hc_attn_fn_broadcast to a freshly-allocated tensor on every call. CUDA graphs captured against the original tensor's address keep replaying with the old pointer, so a weight refit (repeat load_weights() after capture, e.g. RL weight sync) left captured graphs reading stale broadcast weights. Fixed by allocating only when the buffer is None and otherwise copy_()-ing in place, keeping the address stable.

Test Plan

Two unit tests added to tests/kernels/test_mhc_kernels.py:

1. test_deepseek_v4_mhc_broadcast_finalize_sums_hc_streams — first finalize allocates the buffer with the stream-summed values.
2. test_deepseek_v4_mhc_broadcast_refit_refreshes_in_place — after mutating hc_attn_fn and re-finalizing, the buffer is the same tensor object with updated values (the refit contract).
   pytest tests/kernels/test_mhc_kernels.py -k mhc_broadcast -v
   pytest tests/kernels/test_mhc_kernels.py -k "hc_head_tilelang or mhc_broadcast"

End to end tests on verl side with DSV4 training

Test Result

• -k mhc_broadcast: 2 passed.
• Combined with the tilelang hc_head kernel tests (which flip
  torch.set_default_device to CUDA earlier in the file): 10 passed.
• Full file collects all 53 tests cleanly.
• Mutation check: with the model.py fix reverted to HEAD, the refit test
  fails on the buffer-identity assertion and the sums test still passes — the
  pair pins both the pre-existing semantics and the new in-place contract.
• End to end tests on verl side with DSV4 training returned to normal training inference mismatch.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[aoshen02]

/ci run

[mergify]

Hi @HollowMan6, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

[github-actions]

✅ No earlier CI build exists for this PR. Use /ci run first.

[HollowMan6]

/ci run

[github-actions]

✅ Triggered Buildkite CI #84213 for commit e1060155aca7.