From 2e2a479c08c54adf707f3b209695df480b9c8747 Mon Sep 17 00:00:00 2001 From: project6 Date: Fri, 7 Aug 2026 09:09:20 +0000 Subject: [PATCH] =?UTF-8?q?perf(moe):=20translate=20CCCL=20dispatch=5Fcopy?= =?UTF-8?q?=5Fmdspan.cuh=20=E2=80=94=20contiguous=20slice=20fast=20path?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- qwen3_6_scripts/qwen3_5.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/qwen3_6_scripts/qwen3_5.py b/qwen3_6_scripts/qwen3_5.py index e1adbddf..acf56c80 100644 --- a/qwen3_6_scripts/qwen3_5.py +++ b/qwen3_6_scripts/qwen3_5.py @@ -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)