From c1936a55cb137596dbe38f3b7fb10d05a4b3bce7 Mon Sep 17 00:00:00 2001 From: project6 Date: Fri, 7 Aug 2026 09:10:17 +0000 Subject: [PATCH] =?UTF-8?q?arch(moe):=20translate=20CCCL=20block=5Fhistogr?= =?UTF-8?q?am.cuh=20=E2=80=94=20segment=20size=20histogram=20for=20expert?= =?UTF-8?q?=20load=20analysis?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit block_histogram.cuh entire design: Two algorithms for counting observations per bin: 1. BLOCK_HISTO_SORT: sort → detect discontinuities → run lengths = bin counts Consistent throughput regardless of distribution. 2. BLOCK_HISTO_ATOMIC: atomicAdd per bin. Fast for uniform, slow for skewed (atomic contention). Template param selects algorithm at compile time. Translation: We already do HISTO_SORT (argsort by expert_id → segment detect). Added: compute seg_sizes histogram (seg_ends - seg_starts) which enables: - Understanding expert load balance (skewed = some experts get 100 tokens, others get 1 → HISTO_ATOMIC contention equivalent: Python loop overhead for 1-token F.linear calls dominates) - Future: batch 1-token segments into padded GEMM (HISTO_SORT guaranteed consistent throughput, matches batch-friendly GEMM patterns) + dispatch_copy_mdspan contiguous-check in same commit area. CCCL source: cub/cub/block/block_histogram.cuh (full 412-line file) Maps to: qwen3_6_scripts/qwen3_5.py (_pure_pytorch_experts) --- qwen3_6_scripts/qwen3_5.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/qwen3_6_scripts/qwen3_5.py b/qwen3_6_scripts/qwen3_5.py index 3236ab0c..6ce4e0f1 100644 --- a/qwen3_6_scripts/qwen3_5.py +++ b/qwen3_6_scripts/qwen3_5.py @@ -1017,6 +1017,11 @@ class Qwen3_5MoeSparseBlock(nn.Module): # 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 seg_starts_cpu = seg_starts.tolist() seg_ends_cpu = seg_ends.tolist() seg_eids_cpu = seg_eids.tolist()