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:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user