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

perf(moe): CCCL basic_vector pattern — batch GPU→CPU sync in segment detection

Browse Source 
thrust basic_vector.cu: device→host copy is batched (D = H, one memcpy).
Our MoE segment loop did int() per iteration — N separate GPU→CPU syncs.
Fix: .tolist() does ONE sync for all segment boundaries.

Maps to: qwen3_6_scripts/qwen3_5.py (_pure_pytorch_experts prefill path)
This commit is contained in:
project6
2026-08-07 09:03:06 +00:00
parent 57b83ed19e
commit 32fdae237a
1 changed files with 8 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
qwen3_6_scripts/qwen3_5.py
View File

@@ -895,9 +895,14 @@ class Qwen3_5MoeSparseBlock(nn.Module):
seg_eids = sorted_eids[seg_starts]
# Process each expert segment (contiguous tokens → single F.linear)
for seg_i in range(len(seg_starts)):
s, e = int(seg_starts[seg_i]), int(seg_ends[seg_i])
eid = int(seg_eids[seg_i])
# CCCL basic_vector.cu: device→host copy should be batched.
# .tolist() does ONE GPU→CPU sync vs int() doing one per element.
seg_starts_cpu = seg_starts.tolist()
seg_ends_cpu = seg_ends.tolist()
seg_eids_cpu = seg_eids.tolist()
for seg_i in range(len(seg_starts_cpu)):
s, e = seg_starts_cpu[seg_i], seg_ends_cpu[seg_i]
eid = seg_eids_cpu[seg_i]
tok_ids_seg = sorted_tok_ids[s:e]
topk_pos_seg = sorted_topk_pos[s:e]