arch(moe): translate CCCL block_histogram.cuh — segment size histogram for expert load analysis

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)
This commit is contained in:
project6
2026-08-07 09:10:17 +00:00
parent 83192486d3
commit c1936a55cb

View File

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