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

fix: decode loop also needs int32 token_ids for hash router

Browse Source
This commit is contained in:
biondizzle
2026-06-01 01:58:45 +00:00
parent 905623793b
commit cfd2468c61
1 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
single_shot_inference.py
View File

@@ -715,8 +715,9 @@ def main():
print(f"\nDecoding (max {MAX_NEW_TOKENS} tokens)...")
for step in range(MAX_NEW_TOKENS):
t1 = time.time()
tid = torch.tensor([all_tokens[-1]], dtype=torch.long, device='cuda:0')
dec_pos = torch.tensor([len(all_tokens)-1], dtype=torch.long, device='cuda:0')
tid_int64 = torch.tensor([all_tokens[-1]], dtype=torch.long, device='cuda:0')
tid = tid_int64.to(torch.int32) # hash router needs int32
dec_pos = torch.tensor([len(all_tokens)-1], dtype=torch.long, device='cuda:0)
X = mHCLayer.init_state(embed(tid_int64))
for li in range(n_layers):
gpu = li % NUM_GPUS