class ROCmLatentMoERunner(MoERunner):
"""MoE runner for latent MoE with a replicated routed up-projection.
Mirrors CUDA's LatentMoERunner, but currently only the up projection
-sharded path is implemented. (Tier 2)
Native path: the replicated up-proj produces the full hidden dim on every
rank, so the base runner combines routed + shared correctly at any TP size.
"""
def __init__(
self,
*args,
**kwargs,
) -> None:
super().__init__(*args, **kwargs)
transform = self.routed_output_transform
up_proj = getattr(transform, "up_proj", None)
tp_size = get_tensor_model_parallel_world_size()
self._up_proj_preshard = bool(getattr(transform, "row_sharded", False))
self._up_proj_shard_size = 0
self._tail_shardable = (
up_proj is not None
and tp_size > 1
and (self._up_proj_preshard or up_proj.weight.shape[0] % tp_size == 0)
and self._shared_experts is not None
and not self.moe_config.is_sequence_parallel
and self.routed_scaling_factor == 1.0
)
if self._tail_shardable:
assert up_proj is not None
self._up_proj_shard_size = (
up_proj.weight.shape[0]
if self._up_proj_preshard
else up_proj.weight.shape[0] // tp_size
)
else:
logger.warning_once(
"K3 latent-MoE tail is not shardable under this config, "
"falling back to the replicated up-projection.",
scope="global",
)
self._logged_sharded_tail = False
self._quantized_routed_input: bool | None = None
def _shard_up_proj_tail(
self,
fused_output: torch.Tensor,
shared_output: torch.Tensor,
trunc_size: int | None,
) -> torch.Tensor:
"""Tier 2: column-parallel up-projection folded into the final reduce."""
if not self._logged_sharded_tail:
self._logged_sharded_tail = True
logger.info_once(
"Kimi-K3 latent-MoE tail: up-projecting only this rank's "
"hidden shard into the shared output.",
scope="global",
)
transform = self.routed_output_transform
assert transform is not None
latent = tensor_model_parallel_all_reduce(fused_output)
if transform.norm is not None:
latent = transform.norm(latent)
shard_size = self._up_proj_shard_size
shard_start = get_tensor_model_parallel_rank() * shard_size
weight = transform.up_proj.weight
up_proj_shard = (
weight
if self._up_proj_preshard
else weight.narrow(0, shard_start, shard_size)
)
hidden_shard = shared_output.narrow(-1, shard_start, shard_size)
# hidden_shard += latent @ up_proj_shard.T, accumulated in the GEMM's
# beta-add epilogue so folding in the shared partial costs no kernel.
hidden_shard.addmm_(latent, up_proj_shard.t())
return self._maybe_reduce_final_output(
shared_output, trunc_size, output_is_reduced=False
)
def forward(
self,
hidden_states: torch.Tensor,
router_logits: torch.Tensor,
input_ids: torch.Tensor | None = None,
shared_experts_input: torch.Tensor | None = None,
) -> torch.Tensor:
if self._tail_shardable and not self._fused_output_is_reduced:
return self._fused_forward(
hidden_states, router_logits, input_ids, shared_experts_input
)
return super().forward(
hidden_states, router_logits, input_ids, shared_experts_input
)
def _fused_forward(
self,
hidden_states: torch.Tensor,
router_logits: torch.Tensor,
input_ids: torch.Tensor | None,
shared_experts_input: torch.Tensor | None,
) -> torch.Tensor:
# The modular kernel is selected after weight loading. Resolve its
# compatible prepare adapter once, during warmup before graph capture.
if self._quantized_routed_input is None:
self._quantized_routed_input = enable_quantized_routed_input(self)
transform = self.routed_input_transform
defer_projection = (
self._quantized_routed_input
and shared_experts_input is None
and transform is not None
and transform._token_group is not None
and hidden_states.shape[0] >= transform.min_sharded_tokens
and hidden_states.dtype == torch.bfloat16
)
# When the caller pre-applies the routed input transform outside the
# runner (e.g. to overlap it on a separate stream), it passes the
# already-transformed routed input as ``hidden_states`` and the original
# hidden states as ``shared_experts_input``; skip the transform then.
if defer_projection:
shared_experts_input = hidden_states
og_hidden_dim_pre_xform: int | None = self.moe_config.hidden_dim
og_hidden_dim_post_xform: int | None = hidden_states.shape[-1]
elif shared_experts_input is None:
hidden_states, shared_experts_input = self.apply_routed_input_transform(
hidden_states
)
if not defer_projection:
hidden_states, og_hidden_dim_pre_xform, og_hidden_dim_post_xform = (
self._maybe_pad_hidden_states(
shared_experts_input,
hidden_states,
)
)
result = self._forward_entry(
hidden_states,
router_logits,
shared_experts_input,
input_ids,
self._encode_layer_name(),
self.moe_config.hidden_dim_unpadded
if defer_projection or self._quant_method.has_unpadded_output
else 0,
)
shared_output, fused_output = cast(tuple[torch.Tensor, torch.Tensor], result)
if og_hidden_dim_pre_xform is not None:
fused_output = fused_output[..., :og_hidden_dim_pre_xform]
result = self._shard_up_proj_tail(
fused_output, shared_output, og_hidden_dim_post_xform
)
return self._maybe_add_zero_expert_output(result)