feat: add fused moe shared-expert and add-rmsnorm optimization

[blueswhen]

No description provided.

[gemini-code-assist]

Code Review

This pull request introduces performance optimizations for Qwen3.5 and Qwen3Next models, including fused residual RMSNorm, shared expert gating, and specialized Triton kernels. The review feedback identifies a critical race condition in moe_align_fused_kernel requiring a tl.debug_barrier(), and a sequence parallelism bug in Qwen3Next where a hardcoded all_reduce should be replaced with _tpsp_reduce. Additionally, the feedback suggests using Triton's built-in tl.sigmoid for cleaner code, asserting tensor contiguity before calling .view() to prevent runtime crashes, and removing unused dead code in gdn_decode_pack.py.

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]

There is a critical Write-After-Write / Read-After-Write race condition here. Since moe_align_fused_kernel is launched with multiple warps (default num_warps=4), all warps in the thread block will execute the if ZERO_EXPERT_TOKEN_NUM: block and write 0 to expert_token_num_ptr concurrently. Because warps are scheduled independently, some warps might execute the tl.store (writing 0) after other warps have already executed tl.atomic_add (incrementing the count). This results in the incremented values being overwritten by 0, leading to incorrect token counts and silent data corruption or crashes in subsequent grouped GEMMs. Adding a barrier (tl.debug_barrier()) immediately after the tl.store ensures all zeroing writes are completed and visible before any warp proceeds to the atomic additions.

Suggested change

	if ZERO_EXPERT_TOKEN_NUM:
	expert_offs = tl.arange(0, BLOCK_EXPERT)
	tl.store(expert_token_num_ptr + expert_offs, 0, mask=expert_offs < expert_num)
	if ZERO_EXPERT_TOKEN_NUM:
	expert_offs = tl.arange(0, BLOCK_EXPERT)
	tl.store(expert_token_num_ptr + expert_offs, 0, mask=expert_offs < expert_num)
	tl.debug_barrier()

[gemini-code-assist]

Hardcoding all_reduce here breaks sequence parallelism (SP) support for full attention layers. When SP is enabled, the reduction should be a reduce_scatter (which is handled automatically by layer._tpsp_reduce), not an all_reduce. Hardcoding all_reduce causes shape/value mismatches and incorrect results because each rank will get the sum of the sequence-parallel slices instead of its own sequence slice. Please use layer._tpsp_reduce to dynamically handle both TP and SP reductions correctly.

                o = layer._get_o_local(o, infer_state, layer_weight)
                o = layer._tpsp_reduce(o, infer_state)

[gemini-code-assist]

Triton provides a built-in tl.sigmoid function which is cleaner, more readable, and potentially better optimized than the manual sigmoid formula.

Suggested change

	gate = 1.0 / (1.0 + tl.exp(-gate))
	gate = tl.sigmoid(gate)

[gemini-code-assist]

Calling .view() on potentially non-contiguous tensors can raise a RuntimeError. It is safer to assert contiguity before calling .view() to prevent unexpected runtime crashes.

Suggested change

	shared = shared.view(token_num, hidden_dim)
	gate = gate.view(token_num, gate.shape[-1])
	assert shared.is_contiguous()
	assert gate.is_contiguous()
	assert shared.is_contiguous()
	assert gate.is_contiguous()
	shared = shared.view(token_num, hidden_dim)
	gate = gate.view(token_num, gate.shape[-1])

[gemini-code-assist]

Triton provides a built-in tl.sigmoid function which is cleaner, more readable, and potentially better optimized than the manual sigmoid formula.

Suggested change

	gate_vals = 1.0 / (1.0 + tl.exp(-gate_vals))
	gate_vals = tl.sigmoid(gate_vals)

[gemini-code-assist]

Triton provides a built-in tl.sigmoid function which is cleaner, more readable, and potentially better optimized than the manual sigmoid formula.

Suggested change

	gate_vals = 1.0 / (1.0 + tl.exp(-gate_vals))
	gate_vals = tl.sigmoid(gate_vals)

[gemini-code-assist]

Calling .view() on potentially non-contiguous tensors can raise a RuntimeError. It is safer to assert contiguity before calling .view() to prevent unexpected runtime crashes.

Suggested change

	def add_shared_expert_gate_(hidden: torch.Tensor, shared: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
	hidden_arg = hidden.view(-1, hidden.shape[-1])
	shared_arg = shared.view(-1, hidden.shape[-1])
	gate_arg = gate.view(-1, gate.shape[-1])
	def add_shared_expert_gate_(hidden: torch.Tensor, shared: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
	assert hidden.is_contiguous()
	assert shared.is_contiguous()
	assert gate.is_contiguous()
	hidden_arg = hidden.view(-1, hidden.shape[-1])
	shared_arg = shared.view(-1, hidden.shape[-1])
	gate_arg = gate.view(-1, gate.shape[-1])

[gemini-code-assist]

Calling .view() on potentially non-contiguous tensors can raise a RuntimeError. It is safer to assert contiguity before calling .view() to prevent unexpected runtime crashes.

Suggested change

	def sigmoid_mul_(x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
	x_arg = x.view(-1, x.shape[-1])
	gate_arg = gate.view(-1, gate.shape[-1])
	assert gate_arg.shape[0] == x_arg.shape[0] and gate_arg.shape[1] in (1, x_arg.shape[1])
	def sigmoid_mul_(x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
	assert x.is_contiguous()
	assert gate.is_contiguous()
	x_arg = x.view(-1, x.shape[-1])
	gate_arg = gate.view(-1, gate.shape[-1])

[gemini-code-assist]

The function pack_gdn_decode_inputs and its corresponding kernel _pack_gdn_decode_kernel are defined but never imported or used anywhere in the codebase. Removing this dead code will keep the codebase clean and maintainable.