def fused_post_expert_finalize(
routed: torch.Tensor,
weight: torch.Tensor | None,
topk_weights: torch.Tensor,
eps: float,
gain_center: float,
weight_first: bool,
out: torch.Tensor,
*,
topk_ids: torch.Tensor | None = None,
expert_map: torch.Tensor | None = None,
permuted_idx: torch.Tensor | None = None,
) -> torch.Tensor:
"""``out[t] = sum_k tw[t, k] * norm(routed[t, k])`` written in ``out.dtype``.
Routes with a negative ``expert_map`` entry or ``permuted_idx == -1`` were
never written and are masked before the read.
"""
if expert_map is not None:
assert topk_ids is not None and topk_ids.stride(1) == 1
if permuted_idx is not None:
assert permuted_idx.dtype == torch.int32 and permuted_idx.stride(1) == 1
T, K = permuted_idx.shape
H = routed.shape[-1]
stride_xt, stride_xk = routed.stride(0), 0
else:
T, K, H = routed.shape
stride_xt, stride_xk = routed.stride(0), routed.stride(1)
assert routed.stride(-1) == 1 and out.stride(1) == 1
assert topk_weights.stride(1) == 1
assert out.shape == (T, H) and topk_weights.shape == (T, K)
if T == 0:
return out
block_h = triton.next_power_of_2(H)
assert block_h <= 16384, "row too wide for the single-tile finalize"
block_k = triton.next_power_of_2(K)
tiled = block_k * block_h <= 65536
kernel = (
_post_expert_finalize_kernel if tiled else _post_expert_finalize_loop_kernel
)
tile_kwargs = {"BLOCK_K": block_k} if tiled else {}
kernel[(T,)](
routed,
weight if weight is not None else routed,
topk_weights,
out,
topk_ids if topk_ids is not None else topk_weights,
expert_map if expert_map is not None else topk_weights,
expert_map is not None,
topk_ids.stride(0) if topk_ids is not None else 0,
permuted_idx if permuted_idx is not None else topk_weights,
permuted_idx is not None,
permuted_idx.stride(0) if permuted_idx is not None else 0,
K,
H,
eps,
gain_center,
stride_xt,
stride_xk,
topk_weights.stride(0),
out.stride(0),
BLOCK_H=block_h,
HAS_WEIGHT=weight is not None,
WEIGHT_FIRST=weight_first,
num_warps=8 if (block_k * block_h if tiled else block_h) >= 4096 else 4,
**tile_kwargs,
)
return out