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]
|
tok_ids_seg = sorted_tok_ids[s:e]
|
||||||
topk_pos_seg = sorted_topk_pos[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 = F.linear(tokens, w13[eid]) # (n, 2*I)
|
||||||
gate, up = gate_up.chunk(2, dim=-1)
|
gate, up = gate_up.chunk(2, dim=-1)
|
||||||
act = F.silu(gate) * up # (n, I)
|
act = F.silu(gate) * up # (n, I)
|
||||||
|
|||||||
Reference in New Issue
Block a user