biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: token_indices slice bug + torch.zeros for float4/float8 dtypes

Browse Source
This commit is contained in:
biondizzle
2026-05-16 18:21:27 +00:00
parent 54c470e535
commit 533089c9d2
2 changed files with 5 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
cutedsl/bridge.py
View File

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

2
vllm/nvfp4_cutedsl.py
View File

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