perf(moe): translate CCCL dispatch_copy_mdspan.cuh — contiguous slice fast path

dispatch_copy_mdspan.cuh entire design:
  1. Check is_exhaustive() + have_same_strides() (layout compatibility)
  2. Fast path: if contiguous, use DeviceTransform (1D memcpy-like kernel)
  3. Slow path: if non-contiguous, use DeviceFor::for_each_in_extents

Translation to MoE segment loop:
  After sorting tokens by expert_id, tokens routed to the same expert
  often have consecutive original indices. When they do, hidden_states
  slice is zero-copy (view) vs fancy indexing (allocates new tensor).

  Check: tok_ids_seg[-1] == tok_ids_seg[0] + n - 1 (contiguous range)
  Fast: hidden_states[first:first+n] (zero-copy slice)
  Slow: hidden_states[tok_ids_seg] (gather with copy)

CCCL source: cub/cub/device/dispatch/dispatch_copy_mdspan.cuh
Maps to: qwen3_6_scripts/qwen3_5.py (_pure_pytorch_experts)
This commit is contained in:
project6
2026-08-07 09:09:20 +00:00
parent be630106b2
commit 2e2a479c08

View File

@@ -1006,7 +1006,17 @@ class Qwen3_5MoeSparseBlock(nn.Module):
tok_ids_seg = sorted_tok_ids[s:e]
topk_pos_seg = sorted_topk_pos[s:e]
tokens = hidden_states[tok_ids_seg] # (n, H) — contiguous gather
# dispatch_copy_mdspan.cuh: check if data is exhaustive (contiguous).
# If token IDs form a contiguous range, use slice (zero-copy)
# instead of fancy indexing (allocates new tensor).
n_seg = e - s
first_tok = int(tok_ids_seg[0])
if n_seg > 1 and int(tok_ids_seg[-1]) == first_tok + n_seg - 1:
# Fast path: contiguous slice (no copy)
tokens = hidden_states[first_tok:first_tok + n_seg]
else:
# Slow path: gather by index
tokens = hidden_states[tok_ids_seg]
gate_up = F.linear(tokens, w13[eid]) # (n, 2*I)
gate, up = gate_up.chunk(2, dim=-1)
act = F.silu(gate) * up # (n, I)