[bugfix] fix linear_decoupled_in_proj CP

[Jintao-Huang]

No description provided.

[gemini-code-assist]

Code Review

This pull request updates the tensor reshaping logic in gated_delta_net.py to use num_key_heads_per_tp (based on tensor parallel size) instead of num_key_heads_per_device prior to the sequence parallel/context parallel all-to-all operation. The feedback suggests using PyTorch's automatic dimension inference (-1) in the view calls to simplify the code and improve robustness.

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]

For improved readability and robustness, you can use -1 in the view method to let PyTorch automatically infer the size of the last dimension. This avoids manual calculation and makes the code cleaner.

Suggested change

	qkvz = qkvz.view(qkvz.shape[:-1] + (num_key_heads_per_tp, qkvz.shape[-1] // num_key_heads_per_tp))
	ba = ba.view(ba.shape[:-1] + (num_key_heads_per_tp, ba.shape[-1] // num_key_heads_per_tp))
	qkvz = qkvz.view(*qkvz.shape[:-1], num_key_heads_per_tp, -1)
	ba = ba.view(*ba.shape[:-1], num_key_heads_per_tp, -1)