dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

fix: cast router_logits to float32 before topk kernel

Browse Source 
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.
This commit is contained in:
project6-dev
2026-08-12 10:06:19 +00:00
parent 05c09a2801
commit 7969874cde
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
qwen3_6_scripts/qwen3_5.py
View File

@@ -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: