From 76eea36e3853a080ad859d957cf255dcade38597 Mon Sep 17 00:00:00 2001
From: Alison Shao <54658187+alisonshao@users.noreply.github.com>
Date: Tue, 8 Sep 2026 19:48:21 -0700
Subject: [PATCH] Cast fp32 routing weights to bf16 in the Kimi-K3 fused
 finalize (#38588)

---
 python/sglang/srt/layers/k3_ar_fusion.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/python/sglang/srt/layers/k3_ar_fusion.py b/python/sglang/srt/layers/k3_ar_fusion.py
index b91774a5c..d5e0a116e 100644
--- a/python/sglang/srt/layers/k3_ar_fusion.py
+++ b/python/sglang/srt/layers/k3_ar_fusion.py
@@ -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,