From 2102146c01c4c86eec61b7a4e26453bf2504ef30 Mon Sep 17 00:00:00 2001 From: project6 Date: Fri, 7 Aug 2026 09:17:23 +0000 Subject: [PATCH] =?UTF-8?q?refactor(moe):=20translate=20thrust=20mode.cu?= =?UTF-8?q?=20pipeline=20=E2=80=94=20unique=5Fconsecutive=20replaces=20man?= =?UTF-8?q?ual=20boundary=20detect?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit thrust/examples/mode.cu entire design (80 lines): Complete GPU pipeline: sort → unique_count → reduce_by_key → max_element Key operations: 1. thrust::sort — bring equal keys together (we already do: argsort) 2. thrust::unique_count — precompute number of unique keys for allocation 3. thrust::reduce_by_key(data, constant_iterator<1>) — count per key 4. thrust::max_element — find the mode (highest count) Design principle: every step is a GPU primitive, no CPU round-trips. constant_iterator<1> trick: turns reduce_by_key into count_by_key. Translation to MoE segment detection: Previous (4 GPU ops + CPU tensors): changes = cat([True, sorted[1:] != sorted[:-1]]) seg_starts = changes.nonzero() seg_ends = cat([seg_starts[1:], tensor([len])]) seg_eids = sorted[seg_starts] Now (1 fused GPU op): seg_eids, _, seg_counts = torch.unique_consecutive(sorted, return_counts=True) seg_ends = seg_counts.cumsum(0) seg_starts = cat([0, seg_ends[:-1]]) unique_consecutive IS mode.cu's sort+reduce_by_key fused: it returns (unique_keys, inverse, counts) — exactly the data mode.cu builds from reduce_by_key(data, constant_iterator<1>, keys_out, counts_out). 3 fewer GPU kernel launches per MoE forward. CCCL source: thrust/examples/mode.cu Maps to: qwen3_6_scripts/qwen3_5.py (_pure_pytorch_experts) --- qwen3_6_scripts/qwen3_5.py | 31 +++++++++++++------------------ 1 file changed, 13 insertions(+), 18 deletions(-) diff --git a/qwen3_6_scripts/qwen3_5.py b/qwen3_6_scripts/qwen3_5.py index c6694fc6..4f7ae2f4 100644 --- a/qwen3_6_scripts/qwen3_5.py +++ b/qwen3_6_scripts/qwen3_5.py @@ -1050,25 +1050,20 @@ class Qwen3_5MoeSparseBlock(nn.Module): sorted_tok_ids = flat_tok_ids[sort_idx] sorted_topk_pos = flat_topk_pos[sort_idx] - # Find segment boundaries — CCCL reduce_by_key: identify contiguous runs - # This replaces the unique().tolist() + per-expert mask.nonzero() pattern - changes = torch.cat([ - torch.tensor([True], device=sorted_eids.device), - sorted_eids[1:] != sorted_eids[:-1], - ]) - seg_starts = changes.nonzero(as_tuple=True)[0] - seg_ends = torch.cat([seg_starts[1:], - torch.tensor([len(sorted_eids)], device=seg_starts.device)]) - seg_eids = sorted_eids[seg_starts] + # thrust/examples/mode.cu complete pipeline translation: + # sort → unique_count → reduce_by_key(data, constant_iterator<1>) → max_element + # torch.unique_consecutive = sort's reduce_by_key in one fused call. + # Returns (unique_keys, inverse, counts) — mode.cu builds the same from + # sort + reduce_by_key(data, constant_iterator<1>, keys, counts). + # Replaces: changes detection → nonzero → concat → 3 separate GPU ops. + seg_eids, _inv, seg_counts = torch.unique_consecutive( + sorted_eids, return_inverse=True, return_counts=True) + seg_ends = seg_counts.cumsum(0) + seg_starts = torch.cat([ + torch.zeros(1, dtype=seg_ends.dtype, device=seg_ends.device), + seg_ends[:-1]]) - # Process each expert segment (contiguous tokens → single F.linear) - # CCCL basic_vector.cu: device→host copy should be batched. - # .tolist() does ONE GPU→CPU sync vs int() doing one per element. - # block_histogram.cuh: compute segment size histogram to understand - # expert load distribution. This enables: - # 1. Logging: understand if MoE routing is balanced or skewed - # 2. Future: batch small segments into padded GEMM (HISTO_SORT vs HISTO_ATOMIC) - seg_sizes = seg_ends - seg_starts # GPU tensor + # Process each expert segment seg_starts_cpu = seg_starts.tolist() seg_ends_cpu = seg_ends.tolist() seg_eids_cpu = seg_eids.tolist()