From 7969874cdebdf2b9a755f931e30e5f4c6caf1fca Mon Sep 17 00:00:00 2001 From: project6-dev Date: Wed, 12 Aug 2026 10:06:19 +0000 Subject: [PATCH] fix: cast router_logits to float32 before topk kernel router_logits from ReplicatedLinear is float16. PyTorch path does .float() but kernel path passed float16 directly. BI-V100 corex CUB half-precision softmax produces garbage. Cast to float32 first. --- qwen3_6_scripts/qwen3_5.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/qwen3_6_scripts/qwen3_5.py b/qwen3_6_scripts/qwen3_5.py index eeda4793..aee361aa 100644 --- a/qwen3_6_scripts/qwen3_5.py +++ b/qwen3_6_scripts/qwen3_5.py @@ -1611,7 +1611,7 @@ class Qwen3_5MoeSparseBlock(nn.Module): # Source: xllm/core/kernels/cuda/moe/moe_topk_softmax_kernels.cuh if _USE_COREX_MOE_TOPK_SOFTMAX: topk_weights, topk_ids = _corex_moe_topk_softmax.moe_topk_softmax( - router_logits, self.top_k, True) + router_logits.float(), self.top_k, True) topk_ids = topk_ids.to(torch.int64) topk_weights = topk_weights.to(hidden_states.dtype) else: