From 533089c9d2dd19045188ae4a51fb75f19de9ad89 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 16 May 2026 18:21:27 +0000 Subject: [PATCH] fix: token_indices slice bug + torch.zeros for float4/float8 dtypes --- cutedsl/bridge.py | 8 ++++---- vllm/nvfp4_cutedsl.py | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/cutedsl/bridge.py b/cutedsl/bridge.py index ed6932bc..2822d8c5 100644 --- a/cutedsl/bridge.py +++ b/cutedsl/bridge.py @@ -303,10 +303,10 @@ def _get_compiled_kernel(num_experts, device, mma_tiler_mn, cluster_shape_mn): K_packed = 256 # minimal N_packed = 256 tokens = 1 - dummy_a = torch.zeros(tokens, K_packed, dtype=torch.float4_e2m1fn_x2, device=device) - dummy_b = torch.zeros(num_experts, K_packed, N_packed, dtype=torch.float4_e2m1fn_x2, device=device) - dummy_sfa = torch.zeros(1, 1, dtype=torch.float8_e4m3fn, device=device) - dummy_sfb = torch.zeros(1, 1, dtype=torch.float8_e4m3fn, device=device) + dummy_a = torch.zeros(tokens, K_packed, dtype=torch.uint8, device=device).view(torch.float4_e2m1fn_x2) + dummy_b = torch.zeros(num_experts, K_packed, N_packed, dtype=torch.uint8, device=device).view(torch.float4_e2m1fn_x2) + dummy_sfa = torch.zeros(1, 1, dtype=torch.float16, device=device).to(torch.float8_e4m3fn) + dummy_sfb = torch.zeros(1, 1, dtype=torch.float16, device=device).to(torch.float8_e4m3fn) dummy_c = torch.zeros(tokens, N_packed, dtype=torch.bfloat16, device=device) dummy_offs = torch.zeros(num_experts, dtype=torch.int32, device=device) ws_size = kernel.get_workspace_size(num_experts) diff --git a/vllm/nvfp4_cutedsl.py b/vllm/nvfp4_cutedsl.py index 3796c2ea..aad6791e 100644 --- a/vllm/nvfp4_cutedsl.py +++ b/vllm/nvfp4_cutedsl.py @@ -199,7 +199,7 @@ class CuTeDSLMoERunner: # ── Build slot mapping ── flat_ids = topk_ids.reshape(-1) flat_weights = topk_weights.reshape(-1) - token_indices = self._token_indices[:num_tokens].reshape(-1) + token_indices = self._token_indices[:num_tokens * top_k] sort_idx = flat_ids.argsort(stable=True) sorted_ids = flat_ids[sort_idx]