kanshan/TensorRT-LLMs
1
0
Fork 0
mirror of https://github.com/NVIDIA/TensorRT-LLM.git synced 2026-02-05 02:31:33 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

[None][fix] Fix CuteDSL MoE unittest (#10983)

Browse Source 
Signed-off-by: Enwei Zhu <21126786+syuoni@users.noreply.github.com>
This commit is contained in:
Enwei Zhu 2026-01-26 08:34:17 +08:00 committed by GitHub
parent 45d7022cc3
commit ffab217974
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 2 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py
View File

@ -740,6 +740,8 @@ def test_nvfp4_gather_grouped_gemm_swiglu_blackwell(
token_final_scales, token_selected_experts = routing_logits.topk(top_k, dim=-1)
token_selected_experts = token_selected_experts.to(torch.int32)
token_final_scales = token_final_scales.softmax(dim=-1).to(torch.float32)
# Ensure at least one valid token
token_selected_experts[0] = 0
(
tile_idx_to_group_idx,