Cast fp32 routing weights to bf16 in the Kimi-K3 fused finalize (#38588)

2026-09-09 10:48:21 +08:00
parent 7a464a7014
commit 76eea36e38
1 changed files with 2 additions and 0 deletions
@@ -371,6 +371,8 @@ def finalize_all_reduce_push_norm(
    state = _get_state()
    assert state is not None
    if expert_weights.dtype != torch.bfloat16:
        expert_weights = expert_weights.to(torch.bfloat16)
    return mod.finalize_all_reduce_push_norm(
        state.world_size,
        out,