diff --git a/aiter/ops/flydsl/kernels/moe_sorting_kernel.py b/aiter/ops/flydsl/kernels/moe_sorting_kernel.py index 39be3eac0e9..273fa4d738d 100644 --- a/aiter/ops/flydsl/kernels/moe_sorting_kernel.py +++ b/aiter/ops/flydsl/kernels/moe_sorting_kernel.py @@ -180,6 +180,7 @@ def _lds_store_raw(raw_ptr, val, idx): _dummy_mask_cache = {} # device -> torch.Tensor(1, dtype=i32, value=1) +_dummy_local_tokens_cache = {} # dummy placeholder when has_local_tokens=False # --------------------------------------------------------------------------- @@ -193,6 +194,7 @@ def _compile_moe_sorting_oneshot( max_tokens: int = 128, unit_size: int = UNIT_SIZE, has_mask: bool = False, + has_local_tokens: bool = False, ): """Compile the oneshot MoE sorting kernel (single kernel, all phases in LDS). @@ -203,9 +205,18 @@ def _compile_moe_sorting_oneshot( topk : int Experts per token (e.g. 8 for DeepSeek R1). max_tokens : int - Upper bound on T for LDS sizing. Actual T is passed at runtime. + Upper bound on T for LDS sizing (always the static padded capacity, + never the dynamic per-call count -- keeps LDS size, grid/workspace + sizing, and kernel-variant selection identical across CUDA graph + capture and replay). unit_size : int GEMM tile-M for padding alignment (default 32). + has_local_tokens : bool + When True, an extra ``p_local_tokens`` [1]-element i32 tensor arg is + read on-device for the *exact* per-call token count, used for data-dependent loop bounds and the + ``num_valid_ids[1]`` output. Never synced to host -- CUDA-graph-safe. + The sentinel value written into padding rows always uses the static + ``max_tokens``-scale ``i32_tokens`` arg. """ arch = get_hip_arch() E = num_experts @@ -260,6 +271,7 @@ def moe_sorting_oneshot_kernel( num_valid_ids: fx.Tensor, moe_buf: fx.Tensor, expert_mask_tensor: fx.Tensor, + local_tokens_tensor: fx.Tensor, i32_tokens: fx.Int32, i32_moe_buf_elems: fx.Int32, ): @@ -267,7 +279,15 @@ def moe_sorting_oneshot_kernel( tid = gpu.thread_idx.x lane = tid % WARP_SIZE wave = tid // WARP_SIZE + tokens = i32_tokens + if has_local_tokens: + ltok_rsrc = buffer_ops.create_buffer_resource( + local_tokens_tensor, max_size=True + ) + tokens = buffer_ops.buffer_load( + ltok_rsrc, fx.Int32(0), vec_width=1, dtype=T.i32 + ) c_zero_i32 = fx.Int32(0) c_one_i32 = fx.Int32(1) c_oob_idx = fx.Int32(0x7FFFFFFF) @@ -612,7 +632,7 @@ def moe_sorting_oneshot_kernel( sorted_w_rsrc, c_zero_i32, total_padded_pre, - c_sentinel | tokens, + c_sentinel | i32_tokens, ONESHOT_BLOCK, tid, c_oob_idx, @@ -746,11 +766,13 @@ def launch_moe_sorting_oneshot( num_valid_ids_out: fx.Tensor, moe_buf: fx.Tensor, expert_mask_tensor: fx.Tensor, + local_tokens_tensor: fx.Tensor, i32_tokens: fx.Int32, i32_moe_buf_elems: fx.Int32, n_grid_blocks: fx.Int32, - stream: fx.Stream = fx.Stream(None), + stream: fx.Stream = None, ): + stream = stream if stream is not None else fx.Stream(None) launcher = moe_sorting_oneshot_kernel( topk_ids_tensor, topk_weights_tensor, @@ -760,6 +782,7 @@ def launch_moe_sorting_oneshot( num_valid_ids_out, moe_buf, expert_mask_tensor, + local_tokens_tensor, i32_tokens, i32_moe_buf_elems, ) @@ -794,6 +817,7 @@ def _compile_moe_sorting_multiphase( topk: int, unit_size: int = UNIT_SIZE, has_mask: bool = False, + has_local_tokens: bool = False, k4_block: int = 256, ): """Compile the multiphase MoE sorting kernels (2 or 4 kernels via HBM workspace). @@ -855,11 +879,12 @@ def _p23_scatter_mesh( my_start, my_end, i32_mesh_stride, + i32_scan_words_per_row, c_topk, K4_BLOCK, has_mask, ): - """P23 Step 4: EP mask check, read uint8 mesh, DPP prefix sum, scatter tokens.""" + lane = tid % WARP_SIZE wave = tid // WARP_SIZE K4_NUM_WAVES = K4_BLOCK // WARP_SIZE @@ -871,7 +896,7 @@ def _p23_scatter_mesh( mask_rsrc, my_expert, vec_width=1, dtype=T.i32 ) p23_bid_enabled = p23_bid_mask != c_zero - i32_words_per_row = i32_mesh_stride >> fx.Int32(2) + i32_words_per_row = i32_scan_words_per_row n_mesh_iters = (my_start != my_end).select( (i32_words_per_row + fx.Int32(K4_BLOCK - 1)) // fx.Int32(K4_BLOCK), c_zero ) @@ -997,8 +1022,9 @@ def launch_clear_ws( workspace: fx.Tensor, i32_total_elems: fx.Int32, n_grid: fx.Int32, - stream: fx.Stream = fx.Stream(None), + stream: fx.Stream = None, ): + stream = stream if stream is not None else fx.Stream(None) launcher = clear_workspace_kernel(workspace, i32_total_elems) launcher.launch(grid=(n_grid, 1, 1), block=(K1_BLOCK, 1, 1), stream=stream) @@ -1012,6 +1038,7 @@ def launch_clear_ws( def p0_scatter_kernel( topk_ids: fx.Tensor, workspace: fx.Tensor, + local_tokens_tensor: fx.Tensor, i32_tokens: fx.Int32, i32_mesh_stride: fx.Int32, i32_niters: fx.Int32, @@ -1024,7 +1051,15 @@ def p0_scatter_kernel( c_topk = fx.Int32(topk) c_one = fx.Int32(1) - total = i32_tokens * c_topk + tokens_ = i32_tokens + if has_local_tokens: + ltok_rsrc = buffer_ops.create_buffer_resource( + local_tokens_tensor, max_size=True + ) + tokens_ = buffer_ops.buffer_load( + ltok_rsrc, fx.Int32(0), vec_width=1, dtype=T.i32 + ) + total = tokens_ * c_topk _s = fx.Index(0) _e = ArithValue(i32_niters).index_cast(T.index) @@ -1047,14 +1082,21 @@ def p0_scatter_kernel( def launch_p0( topk_ids: fx.Tensor, workspace: fx.Tensor, + local_tokens_tensor: fx.Tensor, i32_tokens: fx.Int32, i32_mesh_stride: fx.Int32, i32_niters: fx.Int32, n_grid: fx.Int32, - stream: fx.Stream = fx.Stream(None), + stream: fx.Stream = None, ): + stream = stream if stream is not None else fx.Stream(None) launcher = p0_scatter_kernel( - topk_ids, workspace, i32_tokens, i32_mesh_stride, i32_niters + topk_ids, + workspace, + local_tokens_tensor, + i32_tokens, + i32_mesh_stride, + i32_niters, ) launcher.launch(grid=(n_grid, 1, 1), block=(K2_BLOCK, 1, 1), stream=stream) @@ -1077,6 +1119,8 @@ class P1SharedStorage: def p1_count_kernel( workspace: fx.Tensor, expert_mask_tensor: fx.Tensor, + local_tokens_tensor: fx.Tensor, + i32_tokens: fx.Int32, i32_mesh_stride: fx.Int32, i32_mesh_size: fx.Int32, ): @@ -1092,11 +1136,21 @@ def p1_count_kernel( reduce_mr = fx.SharedAllocator().allocate(P1SharedStorage).peek().reduce.ptr + # Data-dependent scan + tokens_ = i32_tokens + if has_local_tokens: + ltok_rsrc = buffer_ops.create_buffer_resource( + local_tokens_tensor, max_size=True + ) + tokens_ = buffer_ops.buffer_load( + ltok_rsrc, fx.Int32(0), vec_width=1, dtype=T.i32 + ) + mesh_row_i32_base = (eid * i32_mesh_stride) >> fx.Int32(2) - i32_words_per_row = i32_mesh_stride >> fx.Int32(2) - n_iters = (i32_words_per_row + fx.Int32(K3_WORDS_PER_ITER - 1)) >> fx.Int32( - K3_WORDS_PER_ITER_LOG2 - ) + i32_scan_words_per_row = (tokens_ + fx.Int32(3)) >> fx.Int32(2) + n_iters = ( + i32_scan_words_per_row + fx.Int32(K3_WORDS_PER_ITER - 1) + ) >> fx.Int32(K3_WORDS_PER_ITER_LOG2) if has_mask: mask_rsrc = buffer_ops.create_buffer_resource( @@ -1123,14 +1177,16 @@ def p1_count_kernel( word_base = fx.Int32(_i) * fx.Int32(K3_WORDS_PER_ITER) + tid * fx.Int32( K3_VEC_WIDTH ) - valid = word_base < i32_words_per_row + valid = word_base < i32_scan_words_per_row safe_addr = mesh_row_i32_base + valid.select(word_base, c_zero) vec4 = buffer_ops.buffer_load(ws_rsrc, safe_addr, vec_width=4, dtype=T.i32) iter_cnt = c_zero for _wi in range_constexpr(K3_VEC_WIDTH): word = Vec(vec4)[_wi] - word_valid = valid & ((word_base + fx.Int32(_wi)) < i32_words_per_row) + word_valid = valid & ( + (word_base + fx.Int32(_wi)) < i32_scan_words_per_row + ) b0 = word & c_ff b1 = (word >> fx.Int32(8)) & c_ff b2 = (word >> fx.Int32(16)) & c_ff @@ -1174,12 +1230,20 @@ def p1_count_kernel( def launch_p1( workspace: fx.Tensor, expert_mask_tensor: fx.Tensor, + local_tokens_tensor: fx.Tensor, + i32_tokens: fx.Int32, i32_mesh_stride: fx.Int32, i32_mesh_size: fx.Int32, - stream: fx.Stream = fx.Stream(None), + stream: fx.Stream = None, ): + stream = stream if stream is not None else fx.Stream(None) launcher = p1_count_kernel( - workspace, expert_mask_tensor, i32_mesh_stride, i32_mesh_size + workspace, + expert_mask_tensor, + local_tokens_tensor, + i32_tokens, + i32_mesh_stride, + i32_mesh_size, ) launcher.launch(grid=(E, 1, 1), block=(K3_BLOCK, 1, 1), stream=stream) @@ -1206,6 +1270,7 @@ def p0v2_kernel( topk_ids: fx.Tensor, workspace: fx.Tensor, expert_mask_tensor: fx.Tensor, + local_tokens_tensor: fx.Tensor, i32_tokens: fx.Int32, i32_mesh_stride: fx.Int32, i32_mesh_size: fx.Int32, @@ -1227,14 +1292,34 @@ def p0v2_kernel( reduce_mr = fx.SharedAllocator().allocate(P0V2SharedStorage).peek().reduce.ptr - # Precompute mesh row base (in i32 words) and words per row mesh_row_i32_base = (eid * i32_mesh_stride) >> fx.Int32(2) + i32_words_per_row = i32_mesh_stride >> fx.Int32(2) + tokens_ = i32_tokens + if has_local_tokens: + ltok_rsrc = buffer_ops.create_buffer_resource( + local_tokens_tensor, max_size=True + ) + tokens_ = buffer_ops.buffer_load( + ltok_rsrc, fx.Int32(0), vec_width=1, dtype=T.i32 + ) + + # Phase 3 (count) only needs to scan words that can hold real mesh + # bytes -- the first tokens_ columns (dynamic per-call count), never + # the full static row width used by Phase 1's clear. + i32_scan_words_per_row = (tokens_ + fx.Int32(3)) >> fx.Int32(2) + clear_niters = (i32_words_per_row + fx.Int32(P0V2_BLOCK - 1)) >> fx.Int32(9) - total_assignments = i32_tokens * c_topk + total_assignments = tokens_ * c_topk scatter_niters = (total_assignments + fx.Int32(P0V2_BLOCK - 1)) >> fx.Int32(9) + count_niters = (i32_scan_words_per_row + fx.Int32(P0V2_BLOCK - 1)) >> fx.Int32( + 9 + ) + # Hoist before if/else: AST rewriter extracts branches into separate + # functions, so variables must be defined in outer scope first. + is_local_expert = c_one != c_zero # EP: load mask, write cumsum=0 for masked experts, set loop bounds to 0 if has_mask: m_val = buffer_ops.buffer_load(mask_rsrc, eid, vec_width=1, dtype=T.i32) @@ -1245,6 +1330,7 @@ def p0v2_kernel( ) clear_niters = is_local_expert.select(clear_niters, c_zero) scatter_niters = is_local_expert.select(scatter_niters, c_zero) + count_niters = is_local_expert.select(count_niters, c_zero) # ---- Phase 1: Clear this expert's mesh row ---- for _ci in range( @@ -1293,7 +1379,6 @@ def p0v2_kernel( gpu.barrier() # ---- Phase 3: Count non-zero bytes + warp/cross-wave reduce ---- - count_niters = clear_niters # same loop structure, reuse (already EP-gated) for _ki, state in range( fx.Index(0), ArithValue(count_niters).index_cast(T.index), @@ -1303,7 +1388,7 @@ def p0v2_kernel( cnt_so_far = state[0] word_base = fx.Int32(_ki) * c_block + tid - valid = word_base < i32_words_per_row + valid = word_base < i32_scan_words_per_row safe_addr = mesh_row_i32_base + valid.select(word_base, c_zero) word = buffer_ops.buffer_load(ws_rsrc, safe_addr, vec_width=1, dtype=T.i32) @@ -1351,15 +1436,18 @@ def launch_p0v2( topk_ids: fx.Tensor, workspace: fx.Tensor, expert_mask_tensor: fx.Tensor, + local_tokens_tensor: fx.Tensor, i32_tokens: fx.Int32, i32_mesh_stride: fx.Int32, i32_mesh_size: fx.Int32, - stream: fx.Stream = fx.Stream(None), + stream: fx.Stream = None, ): + stream = stream if stream is not None else fx.Stream(None) launcher = p0v2_kernel( topk_ids, workspace, expert_mask_tensor, + local_tokens_tensor, i32_tokens, i32_mesh_stride, i32_mesh_size, @@ -1392,6 +1480,7 @@ def p23_kernel( num_valid_ids: fx.Tensor, moe_buf: fx.Tensor, expert_mask_tensor: fx.Tensor, + local_tokens_tensor: fx.Tensor, i32_tokens: fx.Int32, i32_mesh_stride: fx.Int32, i32_mesh_size: fx.Int32, @@ -1448,6 +1537,15 @@ def p23_kernel( if is_sort_block: my_expert = bid + tokens_ = i32_tokens + if has_local_tokens: + ltok_rsrc = buffer_ops.create_buffer_resource( + local_tokens_tensor, max_size=True + ) + tokens_ = buffer_ops.buffer_load( + ltok_rsrc, fx.Int32(0), vec_width=1, dtype=T.i32 + ) + # Step 1: Load expert counts from workspace -> pad to unit_size -> LDS cumsum # Process E experts in chunks of K4_BLOCK (256). Most models have # E <= 256, so the extra chunk is only needed for E > 256 @@ -1537,7 +1635,7 @@ def p23_kernel( num_valid_ids, max_size=True ) buffer_ops.buffer_store(total_padded, nvalid_rsrc, c_zero) - buffer_ops.buffer_store(i32_tokens, nvalid_rsrc, c_one) + buffer_ops.buffer_store(tokens_, nvalid_rsrc, c_one) # Step 3: Write sorted_expert_ids for THIS expert (using local_idx_p23 for EP) # Store local_idx to LDS cumsum[tid], barrier, read cumsum[my_expert] @@ -1563,6 +1661,8 @@ def p23_kernel( ) # Step 4: Mesh-based scatter (EP mask + uint8 mesh read + DPP prefix sum + scatter) + + i32_scan_words_per_row = (tokens_ + fx.Int32(3)) >> fx.Int32(2) scatter_end_pos_t0 = _p23_scatter_mesh( tid, scatter_mr, @@ -1575,6 +1675,7 @@ def p23_kernel( my_start, my_end, i32_mesh_stride, + i32_scan_words_per_row, c_topk, K4_BLOCK, has_mask, @@ -1602,13 +1703,15 @@ def launch_p23( num_valid_ids_out: fx.Tensor, moe_buf: fx.Tensor, expert_mask_tensor: fx.Tensor, + local_tokens_tensor: fx.Tensor, i32_tokens: fx.Int32, i32_mesh_stride: fx.Int32, i32_mesh_size: fx.Int32, i32_moe_buf_elems: fx.Int32, n_grid: fx.Int32, - stream: fx.Stream = fx.Stream(None), + stream: fx.Stream = None, ): + stream = stream if stream is not None else fx.Stream(None) launcher = p23_kernel( workspace, topk_weights_tensor, @@ -1618,6 +1721,7 @@ def launch_p23( num_valid_ids_out, moe_buf, expert_mask_tensor, + local_tokens_tensor, i32_tokens, i32_mesh_stride, i32_mesh_size, @@ -1636,17 +1740,20 @@ def launch_p0v2_p23( num_valid_ids_out: fx.Tensor, moe_buf: fx.Tensor, expert_mask_tensor: fx.Tensor, + local_tokens_tensor: fx.Tensor, i32_tokens: fx.Int32, i32_mesh_stride: fx.Int32, i32_mesh_size: fx.Int32, i32_moe_buf_elems: fx.Int32, n_grid_p23: fx.Int32, - stream: fx.Stream = fx.Stream(None), + stream: fx.Stream = None, ): + stream = stream if stream is not None else fx.Stream(None) l1 = p0v2_kernel( topk_ids, workspace, expert_mask_tensor, + local_tokens_tensor, i32_tokens, i32_mesh_stride, i32_mesh_size, @@ -1662,6 +1769,7 @@ def launch_p0v2_p23( num_valid_ids_out, moe_buf, expert_mask_tensor, + local_tokens_tensor, i32_tokens, i32_mesh_stride, i32_mesh_size, @@ -1680,6 +1788,7 @@ def launch_4k_fused( num_valid_ids_out: fx.Tensor, moe_buf: fx.Tensor, expert_mask_tensor: fx.Tensor, + local_tokens_tensor: fx.Tensor, i32_tokens: fx.Int32, i32_mesh_stride: fx.Int32, i32_mesh_size: fx.Int32, @@ -1689,18 +1798,29 @@ def launch_4k_fused( n_grid_k1: fx.Int32, n_grid_k2: fx.Int32, n_grid_p23: fx.Int32, - stream: fx.Stream = fx.Stream(None), + stream: fx.Stream = None, ): + stream = stream if stream is not None else fx.Stream(None) l1 = clear_workspace_kernel(workspace, i32_ws_total) l1.launch(grid=(n_grid_k1, 1, 1), block=(K1_BLOCK, 1, 1), stream=stream) l2 = p0_scatter_kernel( - topk_ids, workspace, i32_tokens, i32_mesh_stride, i32_p0_niters + topk_ids, + workspace, + local_tokens_tensor, + i32_tokens, + i32_mesh_stride, + i32_p0_niters, ) l2.launch(grid=(n_grid_k2, 1, 1), block=(K2_BLOCK, 1, 1), stream=stream) l3 = p1_count_kernel( - workspace, expert_mask_tensor, i32_mesh_stride, i32_mesh_size + workspace, + expert_mask_tensor, + local_tokens_tensor, + i32_tokens, + i32_mesh_stride, + i32_mesh_size, ) l3.launch(grid=(E, 1, 1), block=(K3_BLOCK, 1, 1), stream=stream) @@ -1713,6 +1833,7 @@ def launch_4k_fused( num_valid_ids_out, moe_buf, expert_mask_tensor, + local_tokens_tensor, i32_tokens, i32_mesh_stride, i32_mesh_size, @@ -1781,6 +1902,7 @@ def compile_moe_sorting( max_tokens=128, unit_size=UNIT_SIZE, has_mask=False, + has_local_tokens=False, k4_block=256, ): """Compile MoE sorting kernels for all paths (oneshot + multiphase). @@ -1794,12 +1916,14 @@ def compile_moe_sorting( max_tokens=max_tokens, unit_size=unit_size, has_mask=has_mask, + has_local_tokens=has_local_tokens, ) _, _, _, _, _, launch_p0v2_p23, launch_4k_fused = _compile_moe_sorting_multiphase( num_experts=num_experts, topk=topk, unit_size=unit_size, has_mask=has_mask, + has_local_tokens=has_local_tokens, k4_block=k4_block, ) return launch_oneshot, launch_p0v2_p23, launch_4k_fused @@ -1831,19 +1955,45 @@ def moe_sorting_flydsl( All output tensors (sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, moe_buf) must be pre-allocated by the caller. + ``num_local_tokens``, when a CUDA tensor, is never read on the host (no + ``.item()``/sync -- this is what makes the whole path safe to call from + inside ``torch.cuda.graph()`` capture, e.g. under DP-attention + EP). Host + -side sizing (LDS, workspace, grid dims, which kernel variant to launch) + always uses the static padded capacity ``topk_ids.shape[0]`` instead, so + kernel selection is identical across capture and replay regardless of the + real per-call token count. The exact count is read **on-device** inside + the compiled kernel (mirroring Opus's ``p_local_tokens`` pattern) and used + only for data-dependent loop bounds / the ``num_valid_ids[1]`` output. + Returns ------- sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, moe_buf """ topk = topk_ids.shape[1] - if num_local_tokens is not None: - M = ( - num_local_tokens.item() - if isinstance(num_local_tokens, torch.Tensor) - else int(num_local_tokens) - ) - else: + has_local_tokens = ( + isinstance(num_local_tokens, torch.Tensor) and num_local_tokens.is_cuda + ) + if has_local_tokens: + local_tokens_tensor = num_local_tokens M = topk_ids.shape[0] + else: + local_tokens_tensor = None + if isinstance(num_local_tokens, torch.Tensor): + M = int(num_local_tokens.item()) + else: + M = ( + int(num_local_tokens) + if num_local_tokens is not None + else topk_ids.shape[0] + ) + + if local_tokens_tensor is None: + local_tokens_tensor = _dummy_local_tokens_cache.get(topk_ids.device) + if local_tokens_tensor is None: + local_tokens_tensor = torch.zeros( + 1, dtype=torch.int32, device=topk_ids.device + ) + _dummy_local_tokens_cache[topk_ids.device] = local_tokens_tensor sub_tokens = _compute_sub_tokens(num_experts) @@ -1890,6 +2040,7 @@ def moe_sorting_flydsl( max_tokens=max_tokens, unit_size=unit_size, has_mask=has_mask, + has_local_tokens=has_local_tokens, ) oneshot_args = ( topk_ids, @@ -1900,6 +2051,7 @@ def moe_sorting_flydsl( num_valid_ids, moe_buf_i32, mask_tensor, + local_tokens_tensor, M, moe_buf_elems, n_grid_blocks, @@ -1927,6 +2079,7 @@ def moe_sorting_flydsl( topk=topk, unit_size=unit_size, has_mask=has_mask, + has_local_tokens=has_local_tokens, k4_block=k4_block, ) stream = torch.cuda.current_stream(device) @@ -1945,6 +2098,7 @@ def moe_sorting_flydsl( num_valid_ids, moe_buf_i32, mask_tensor, + local_tokens_tensor, M, mesh_stride, ws_mesh_i32, @@ -1968,6 +2122,7 @@ def moe_sorting_flydsl( num_valid_ids, moe_buf_i32, mask_tensor, + local_tokens_tensor, M, mesh_stride, ws_mesh_i32, diff --git a/op_tests/test_moe_sorting.py b/op_tests/test_moe_sorting.py index 08e364a5771..248f875b3a6 100644 --- a/op_tests/test_moe_sorting.py +++ b/op_tests/test_moe_sorting.py @@ -201,7 +201,14 @@ def test_moe_sorting( has_expert_mask=False, padding_extra=0, dispatch_policy=0, + accumulate=True, ): + """accumulate=False (with has_expert_mask=False) makes moe_sorting allocate + a (0,0) moe_buf placeholder instead of the full [M, model_dim] buffer -- + useful for isolating moe_buf-zeroing cost (which scales with the static + padded capacity M) from the mesh-scan work (which should NOT scale with M + once tokens_ correctly bounds it) when comparing padding_extra=0 vs >0. + """ topk_ids, topk_weights, expert_mask, num_local_tokens = _build_moe_sorting_inputs( token, model_dim, @@ -232,6 +239,7 @@ def test_moe_sorting( expert_mask, num_local_tokens, dispatch_policy, + accumulate=accumulate, ), "ck": lambda: moe_sorting( topk_ids, @@ -243,6 +251,7 @@ def test_moe_sorting( expert_mask, num_local_tokens, dispatch_policy, + accumulate=accumulate, ), } # FlyDSL kernel only supports dispatch_policy=0 today. @@ -257,6 +266,7 @@ def test_moe_sorting( expert_mask, num_local_tokens, dispatch_policy, + accumulate=accumulate, ) flops, nbytes = _moe_sorting_roofline(token, topk, E, model_dim, dtype) @@ -291,6 +301,197 @@ def test_moe_sorting( return ret +def test_moe_sorting_flydsl_cuda_graph_capture( + dtype, + token, + model_dim, + E, + topk, + has_expert_mask=False, +): + """Capture moe_sorting (FlyDSL backend) under a real torch.cuda.graph(), + then replay with a DIFFERENT num_local_tokens tensor value than was + captured. Regression test for the .item() capture-safety fix: reading + num_local_tokens on the host during capture used to raise + 'operation not permitted when stream is capturing'; and even a naive + host-side guard would still bake a stale capture-time count into the + graph. This asserts real dynamic behavior survives across replay. + """ + if not is_flydsl_available(): + aiter.logger.warning("flydsl unavailable; skipping cuda-graph capture test") + return + + set_moe_sorting_backend("flydsl") + + capture_tokens = token + replay_tokens = max(1, token // 2) # a different dynamic count at replay time + padding_extra = token # topk_ids capacity must cover the larger of the two + + topk_ids, topk_weights, expert_mask, num_local_tokens = _build_moe_sorting_inputs( + capture_tokens, + model_dim, + E, + topk, + dtype, + has_expert_mask, + padding_extra, + ) + assert isinstance(num_local_tokens, torch.Tensor) + + def run(): + return moe_sorting( + topk_ids, + topk_weights, + E, + model_dim, + dtype, + BLOCK_SIZE_M, + expert_mask, + num_local_tokens, + ) + + # Warm-up on a side stream (standard torch.cuda.graph() prerequisite): + # populates the FlyDSL JIT cache before capture, since compiling a new + # kernel variant mid-capture is itself capture-unsafe. + warmup_stream = torch.cuda.Stream() + warmup_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup_stream): + for _ in range(3): + run() + torch.cuda.current_stream().wait_stream(warmup_stream) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + num_local_tokens.fill_(capture_tokens) + with torch.cuda.graph(graph): + out = run() + torch.cuda.synchronize() + + # Replay with a different dynamic token count -- the num_local_tokens + # buffer is captured by reference, so mutating its contents (not the + # tensor object) is what a real decode step would do between replays. + num_local_tokens.fill_(replay_tokens) + graph.replay() + torch.cuda.synchronize() + + ref = run_torch_moe_sorting( + topk_ids, + topk_weights, + E, + BLOCK_SIZE_M, + expert_mask, + num_local_tokens, + ) + # num_rows here must be the static padded capacity (topk_ids.shape[0]), + # matching run_torch_moe_sorting's sentinel (topk << 24 | m) -- NOT the + # dynamic replay_tokens count. + errs = _compare_moe_sorting_outputs(ref, out, topk, topk_ids.shape[0]) + bad = {k: v for k, v in errs.items() if v} + assert not bad, ( + f"moe_sorting cuda-graph replay mismatch (captured token={capture_tokens}, " + f"replayed token={replay_tokens}, E={E}, topk={topk}, " + f"has_expert_mask={has_expert_mask}): {bad}" + ) + + +def test_moe_sorting_decode_graph_perf( + dtype, + capture_capacity, + model_dim, + E, + topk, + real_tokens, + has_expert_mask=False, + num_iters=50, +): + """Benchmark steady-state torch.cuda.graph() REPLAY latency + + This is what actual decode looks like: the graph is captured once per + bucket, then replayed every step with whatever the real dynamic count + happens to be that step + """ + real_tokens = min(real_tokens, capture_capacity) + topk_ids, topk_weights, expert_mask, num_local_tokens = _build_moe_sorting_inputs( + capture_capacity, + model_dim, + E, + topk, + dtype, + has_expert_mask, + padding_extra=0, # capture_capacity IS the static M; no extra slack + ) + if num_local_tokens is None: + num_local_tokens = torch.tensor( + [capture_capacity], dtype=topk_ids.dtype, device="cuda" + ) + + ret = { + "gfx": get_gfx(), + "E": E, + "topk": topk, + "has_expert_mask": has_expert_mask, + "capture_capacity": capture_capacity, + "real_tokens": real_tokens, + } + backends = ["opus", "ck"] + (["flydsl"] if is_flydsl_available() else []) + for name in backends: + set_moe_sorting_backend(name) + + def run(): + return moe_sorting( + topk_ids, + topk_weights, + E, + model_dim, + dtype, + BLOCK_SIZE_M, + expert_mask, + num_local_tokens, + ) + + # Warm-up on a side stream: populates any JIT cache before capture, + # since compiling mid-capture is itself capture-unsafe. + warmup_stream = torch.cuda.Stream() + warmup_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup_stream): + for _ in range(3): + run() + torch.cuda.current_stream().wait_stream(warmup_stream) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + # moe_sorting() allocates fresh output tensors on every call + # (torch.empty(...) inside _flydsl_moe_sorting / _moe_sorting_impl). + # Keep a live Python reference to what's captured -- CUDA graph + # semantics require capture-time allocations to stay referenced + # for the graph's lifetime, or the allocator can reclaim that + # memory (e.g. for a later unrelated allocation) while replay + # still writes into it. + captured_out = run() + torch.cuda.synchronize() + + # Set the REAL per-step token count, then time steady-state replay. + num_local_tokens.fill_(real_tokens) + for _ in range(5): + graph.replay() + torch.cuda.synchronize() + + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + start_event.record() + for _ in range(num_iters): + graph.replay() + end_event.record() + end_event.synchronize() + us = start_event.elapsed_time(end_event) * 1000 / num_iters + del captured_out + + ret[f"{name} us"] = us + + return ret + + def main(): if get_gfx() not in SUPPORTED_GFX: aiter.logger.warning("moe_sorting unsupported on %s; skipping", get_gfx()) @@ -371,6 +572,28 @@ def main(): default=[True, False], help="Expert mask.\n e.g.: -em f", ) + parser.add_argument( + "-accum", + "--accumulate", + type=dtypes.str2bool, + nargs="*", + default=[True], + help="Stage2 accumulate mode (False + no expert_mask allocates a (0,0) " + "moe_buf placeholder instead of the full [M, model_dim] buffer -- " + "isolates moe_buf-zeroing cost from mesh-scan work).\n e.g.: -accum f", + ) + parser.add_argument( + "-dg", + "--decode_graph", + type=int, + nargs="*", + default=None, + help="Real per-step token counts for a cuda-graph capture/replay decode " + "benchmark (typical decode batch sizes: 1 2 4 8 16 32 64 128). Each is " + "run against every -m value as the captured static capacity (real " + "tokens capped to the capacity). Omit to skip this benchmark.\n" + " e.g.: -dg 1 2 4 8 16 32 64 128", + ) args = parser.parse_args() if len(args.expert) != len(args.topk): @@ -380,10 +603,17 @@ def main(): for dtype in args.dtype: df = [] - for padding_extra, expert_mask, dispatch_policy, m in itertools.product( + for ( + padding_extra, + expert_mask, + dispatch_policy, + accumulate, + m, + ) in itertools.product( args.padding, args.expert_mask, args.dispatch_policy, + args.accumulate, args.m, ): for E, topk in model_configs: @@ -398,6 +628,7 @@ def main(): has_expert_mask=expert_mask, padding_extra=padding_extra, dispatch_policy=dispatch_policy, + accumulate=accumulate, ) ) df = pd.DataFrame(df) @@ -405,6 +636,47 @@ def main(): "moe_sorting summary (markdown):\n%s", df.to_markdown(index=False) ) + if is_flydsl_available(): + for expert_mask, m in itertools.product(args.expert_mask, args.m): + if m < 2: + continue # need capture/replay token counts to differ + for E, topk in model_configs: + test_moe_sorting_flydsl_cuda_graph_capture( + dtype, + m, + args.model_dim, + E, + topk, + has_expert_mask=expert_mask, + ) + aiter.logger.info( + "moe_sorting FlyDSL cuda-graph capture/replay: all passed" + ) + + if args.decode_graph: + decode_rows = [] + for expert_mask, capacity, real_tokens in itertools.product( + args.expert_mask, args.m, args.decode_graph + ): + for E, topk in model_configs: + decode_rows.append( + test_moe_sorting_decode_graph_perf( + dtype, + capacity, + args.model_dim, + E, + topk, + real_tokens, + has_expert_mask=expert_mask, + ) + ) + decode_df = pd.DataFrame(decode_rows) + aiter.logger.info( + "moe_sorting decode graph-replay summary (us, capture_capacity=static M, " + "real_tokens=actual per-step count):\n%s", + decode_df.to_markdown(index=False), + ) + if __name__ == "__main__": main()