[Kimi-K3] Accept fp32 routing weights in the fused MoE finalize (#38612)

Co-authored-by: Mohammad Angkad <mohammad.angkad@radixark.ai>

test/registered/kernels/ops/kimi_k3/test_ar_fusion.py

@@ -248,10 +248,13 @@ def test_ar_fusion_push_norm(num_tokens: int, rows_per_token: int):
 FIN_TOPK = 16
 
 
-def _build_permuted_layout(num_tokens: int, seed: int):
+def _build_permuted_layout(
+    num_tokens: int, seed: int, w_dtype: torch.dtype = torch.float32
+):
     """trtllm-gen permuted gemm2 layout (rows grouped by expert, per-expert
     tile padding). Deterministic on CPU: idx/weights are identical on every
-    rank (TP semantics — same routing), gemm2 values are per-rank."""
+    rank (TP semantics — same routing), gemm2 values are per-rank.
+    ``w_dtype`` is the routing-weight dtype the deferred finalize hands back."""
     num_experts, tile = 896, 8
     gen = torch.Generator(device="cpu").manual_seed(seed)
     topk_ids = torch.stack(
@@ -268,7 +271,7 @@ def _build_permuted_layout(num_tokens: int, seed: int):
     for i, e in enumerate(topk_ids.flatten().tolist()):
         idx[i] = bases[e] + fill[e]
         fill[e] += 1
-    weights = torch.rand(num_tokens, FIN_TOPK, generator=gen).to(torch.bfloat16)
+    weights = torch.rand(num_tokens, FIN_TOPK, generator=gen).to(w_dtype)
     num_rows = int(padded.sum())
     g = torch.Generator(device="cpu").manual_seed(seed * 31 + dist.get_rank())
     gemm2 = (torch.randn(num_rows, NORM_DIM, generator=g) * 2).to(torch.bfloat16)
@@ -296,13 +299,14 @@ def _finalize_norm_ref(gemm2, idx, weights, norm_w, eps: float) -> torch.Tensor:
     return (total * factor * norm_w.float()).to(torch.bfloat16)
 
 
+@pytest.mark.parametrize("w_dtype", [torch.float32, torch.bfloat16])
 @pytest.mark.parametrize("bs", PUSH_BS)
 @torch.inference_mode()
-def test_ar_fusion_finalize_push_norm(bs: int):
+def test_ar_fusion_finalize_push_norm(bs: int, w_dtype: torch.dtype):
     comm = _init_comm()
     world = comm.world_size
     eps = 1e-6
-    gemm2, idx, weights = _build_permuted_layout(bs, seed=bs + 23)
+    gemm2, idx, weights = _build_permuted_layout(bs, seed=bs + 23, w_dtype=w_dtype)
     g = torch.Generator(device="cpu").manual_seed(77)
     norm_w = (torch.rand(NORM_DIM, generator=g) + 0.5).to(torch.bfloat16).to(_device())
     ref = _finalize_norm_ref(gemm2, idx, weights, norm_w, eps)
@@ -325,7 +329,12 @@ def test_ar_fusion_finalize_push_norm_stress():
     norm_w = (torch.rand(NORM_DIM, generator=g) + 0.5).to(torch.bfloat16).to(_device())
     for it in range(12):
         bs = (1, 8, 32)[it % 3]
-        gemm2, idx, weights = _build_permuted_layout(bs, seed=9000 + it)
+        # alternate the routing-weight precision: both kernel instantiations
+        # share the one push workspace
+        w_dtype = (torch.float32, torch.bfloat16)[it % 2]
+        gemm2, idx, weights = _build_permuted_layout(
+            bs, seed=9000 + it, w_dtype=w_dtype
+        )
         ref = _finalize_norm_ref(gemm2, idx, weights, norm_w, eps)
         out = torch.empty(bs, NORM_DIM, dtype=torch.bfloat16, device=_device())
         all_reduce.finalize_all_reduce_push_norm(