diff --git a/sparkinfer/moe/_shared/kernels/w4a16/kernel.py b/sparkinfer/moe/_shared/kernels/w4a16/kernel.py index bcd4cca4..c0cb0e9b 100644 --- a/sparkinfer/moe/_shared/kernels/w4a16/kernel.py +++ b/sparkinfer/moe/_shared/kernels/w4a16/kernel.py @@ -205,6 +205,11 @@ def _fake_m_for_specialization(size_m: int) -> int: (128, 3, 4, 8, False): 249, (128, 3, 8, 4, False): 250, (256, 4, 16, 4, False): 255, + # Measured by the mixed-Trellis block-64 qualification gate. This is the + # stock one-grid prefill tile geometry (N=128, K=128) with four 16-row + # route blocks per CTA. Keeping the measured entry permits block-64 + # routing without changing the K3/K4 accumulation geometry. + (256, 4, 8, 8, False): 255, (128, 4, 4, 8, False): 255, (128, 4, 8, 4, False): 255, } @@ -701,6 +706,8 @@ def __init__( fused_topk_sum: bool = False, fused_sum_topk: int = 1, schedule_whole_tiles: bool = False, + schedule_route_block_factor: int = 1, + paired_m8_routes: bool = False, ): if element_dtype not in {"bf16", "fp16"}: raise ValueError(f"unsupported element_dtype {element_dtype!r}") @@ -876,6 +883,18 @@ def __init__( # the split-K tail machinery entirely. Requires the host to bound the # wave count; used by the exact-geometry hybrid decode schedule. self.schedule_whole_tiles = bool(schedule_whole_tiles) + self.schedule_route_block_factor = int(schedule_route_block_factor) + if self.schedule_route_block_factor < 1: + raise ValueError("schedule_route_block_factor must be >= 1") + if self.schedule_route_block_factor != 1 and ( + not self.schedule_whole_tiles + or self.direct_topk_routes + or self.dense_route_fast_path + ): + raise ValueError( + "grouped route-block scheduling requires route-packed " + "whole-tile execution" + ) if ( self.schedule_whole_tiles and not self.direct_topk_routes @@ -890,6 +909,16 @@ def __init__( raise ValueError("fused_sum_topk must be >= 1") self.cta_m_blocks = int(_covering_count(moe_block_size, 16)) self.uses_m_block_8 = moe_block_size == 8 + self.paired_m8_routes = bool(paired_m8_routes) + if self.paired_m8_routes and ( + not self.uses_m_block_8 + or not self.schedule_whole_tiles + or self.schedule_route_block_factor != 2 + ): + raise ValueError( + "paired_m8_routes requires M8 whole-tile scheduling with " + "schedule_route_block_factor=2" + ) self.max_m_blocks = int(max_m_blocks) if torch.cuda.is_available(): props = torch.cuda.get_device_properties(torch.cuda.current_device()) @@ -921,6 +950,11 @@ def __init__( # W4A16 shared-memory geometry, in int4 units unless noted. self.a_sh_stride = 16 * self.cta_k_blocks // 8 self.a_sh_stage = self.a_sh_stride * (16 * self.cta_m_blocks) + if self.paired_m8_routes: + # The M8 ldmatrix mapping consumes a padded 16-row slab: rows 8-15 + # must remain zero. A paired tile therefore needs two independent + # 16-row slabs even though only eight rows in each slab are live. + self.a_sh_stage *= 2 self.a_gl_rd_delta_o = 16 * self.cta_k_blocks // 8 self.a_sh_wr_delta = self.a_sh_stride * ( self.cta_threads // self.a_gl_rd_delta_o @@ -962,9 +996,12 @@ def __init__( self.s_sh_stage = self.s_tb_groups * self.s_sh_stride self.tb_n_warps = self.cta_n_blocks // 4 - sh_block_route_indices = self.moe_block_size // 4 - sh_rd_block_route_indices = self.moe_block_size // 4 - sh_block_topk_weights = self.moe_block_size // 2 + route_metadata_rows = self.moe_block_size * ( + 2 if self.paired_m8_routes else 1 + ) + sh_block_route_indices = route_metadata_rows // 4 + sh_rd_block_route_indices = route_metadata_rows // 4 + sh_block_topk_weights = route_metadata_rows // 2 self.sh_valid_count_off = ( sh_block_route_indices + sh_rd_block_route_indices + sh_block_topk_weights ) @@ -1034,6 +1071,8 @@ def __cache_key__(self) -> tuple[object, ...]: # entry even when their arithmetic geometry otherwise matches. self.blocks_per_sm, self.schedule_whole_tiles, + self.schedule_route_block_factor, + self.paired_m8_routes, ) @cute.jit @@ -1289,7 +1328,7 @@ def _run_persistent_gemm( ) // Int32(self.moe_block_size) elif cutlass.const_expr(not self.direct_topk_routes): route_blocks = packed_route_count[Int32(0)].to(Int32) // Int32( - self.moe_block_size + self.moe_block_size * self.schedule_route_block_factor ) k_tiles = Int32(self.k_tiles) global_mn_tiles = route_blocks * n_tiles @@ -1613,6 +1652,85 @@ def _read_moe_block_data( cute.arch.sync_threads() return valid_count + @cute.jit + def _read_moe_block_data_pair( + self, + packed_route_indices: cute.Tensor, + topk_weights_flat: cute.Tensor, + smem_base: Int32, + tid: Int32, + route_block_idx: Int32, + global_scale_f32: cutlass.Float32, + active_size_m: Int32, + ): + """Load two adjacent M8 route blocks into one 16-row metadata slab.""" + route_indices_int4_addr = self._int4_addr( + smem_base, Int32(self.sh_route_off) + tid + ) + route_indices_gmem = get_ptr_as_int64( + packed_route_indices, + route_block_idx * Int32(self.moe_block_size) + tid * Int32(4), + ) + cp_async4_shared_global_pred( + route_indices_int4_addr, + route_indices_gmem, + (tid < Int32(2 * self.moe_block_size // 4)).to(Int32), + ) + cute.arch.cp_async_commit_group() + cute.arch.cp_async_wait_group(0) + cute.arch.sync_threads() + + if tid >= Int32(self.cta_threads - 32): + lane = tid - Int32(self.cta_threads - 32) + valid0 = Int32(0) + valid1 = Int32(0) + if lane < Int32(2 * self.moe_block_size): + idx = ld_shared_i32_relaxed( + smem_base + Int32(self.sh_route_off * 16) + lane * Int32(4) + ) + valid = (idx < active_size_m * Int32(self.top_k)).to(Int32) + if lane < Int32(self.moe_block_size): + valid0 = valid + else: + valid1 = valid + valid0 = cute.arch.warp_redux_sync(valid0, "add") + valid1 = cute.arch.warp_redux_sync(valid1, "add") + if lane == Int32(0): + valid_addr = smem_base + Int32(self.sh_valid_count_off * 16) + st_shared_i32(valid_addr, valid0) + st_shared_i32(valid_addr + Int32(4), valid1) + + if tid < Int32(2 * self.moe_block_size): + idx = ld_shared_i32_relaxed( + smem_base + Int32(self.sh_route_off * 16) + tid * Int32(4) + ) + rd_row = idx // Int32(self.top_k) + if cutlass.const_expr(self.route_major_a): + rd_row = idx + st_shared_i32( + smem_base + Int32(self.sh_rd_route_off * 16) + tid * Int32(4), + rd_row, + ) + if cutlass.const_expr(self.mul_topk_weights): + safe_idx = idx + if idx >= active_size_m * Int32(self.top_k): + safe_idx = Int32(0) + topk = ( + topk_weights_flat[safe_idx].to(cutlass.Float32) + * global_scale_f32 + ) + st_shared_u32( + smem_base + Int32(self.sh_topk_off * 16) + tid * Int32(4), + self._broadcast_f32_to_elem2(topk), + ) + + cute.arch.sync_threads() + valid_addr = smem_base + Int32(self.sh_valid_count_off * 16) + block_valid_rows0 = ld_shared_i32_relaxed(valid_addr) + block_valid_rows1 = ld_shared_i32_relaxed(valid_addr + Int32(4)) + cute.arch.sync_threads() + return block_valid_rows0, block_valid_rows1 + @cute.jit def _run_tile( self, @@ -1744,6 +1862,42 @@ def _tile_common_prologue( s_sh_rd, ) + @cute.jit + def _tile_common_prologue_pair( + self, + global_scale: cute.Tensor, + packed_route_indices: cute.Tensor, + topk_weights_flat: cute.Tensor, + smem_base: Int32, + tid: Int32, + route_block_idx: Int32, + expert_idx: Int32, + output_n_tile: Int32, + active_size_m: Int32, + ): + global_scale_f32 = global_scale[expert_idx].to(cutlass.Float32) + if cutlass.const_expr(self.scale_format_e8m0_k32): + if cutlass.const_expr(self.is_fp16): + global_scale_f32 *= cutlass.Float32(_E8M0_K32_FP16_GLOBAL_COMPENSATION) + else: + global_scale_f32 *= cutlass.Float32(_E8M0_K32_BF16_GLOBAL_COMPENSATION) + block_valid_rows0, block_valid_rows1 = self._read_moe_block_data_pair( + packed_route_indices, + topk_weights_flat, + smem_base, + tid, + route_block_idx, + global_scale_f32, + active_size_m, + ) + offsets = self._tile_stream_offsets(tid, expert_idx, output_n_tile) + return ( + global_scale_f32, + block_valid_rows0, + block_valid_rows1, + *offsets, + ) + @cute.jit def _tile_stream_offsets(self, tid: Int32, expert_idx: Int32, output_n_tile: Int32): a_gl_stride = Int32(self.size_k // 8) @@ -1883,6 +2037,8 @@ def _run_tile_m8( k_tiles, reduce_k_tile, block_valid_rows, + Int32(0), + False, a_gl_stride, b_gl_stride, s_gl_stride, @@ -1964,6 +2120,7 @@ def _run_tile_m8( tid, output_n_tile, block_valid_rows, + Int32(0), global_scale_f32, reduce_slice_count, reduce_slice_idx, @@ -1971,6 +2128,194 @@ def _run_tile_m8( True, ) + @cute.jit + def _run_tile_m8_pair( + self, + a_bf16_flat: cute.Tensor, + a_alt_bf16_flat: cute.Tensor, + b_i32_flat: cute.Tensor, + c_bf16_flat: cute.Tensor, + scales_i32_flat: cute.Tensor, + global_scale: cute.Tensor, + packed_route_indices: cute.Tensor, + topk_weights_flat: cute.Tensor, + c_tmp_f32_flat: cute.Tensor, + locks_i32_flat: cute.Tensor, + smem_base: Int32, + tid: Int32, + route_block_idx: Int32, + expert_idx: Int32, + output_n_tile: Int32, + reduce_k_tile: Int32, + reduce_tile_count: Int32, + reduce_slice_count: Int32, + reduce_slice_idx: Int32, + lock_slot: Int32, + active_size_m: Int32, + ): + ( + global_scale_f32, + block_valid_rows0, + block_valid_rows1, + a_gl_stride, + b_gl_stride, + s_gl_stride, + scales_expert_off, + b_gl_rd_base, + a_gl_rd_row, + a_gl_rd_col0, + a_sh_wr, + a_rows_per_iter, + b_sh_rd, + s_sh_rd, + ) = self._tile_common_prologue_pair( + global_scale, + packed_route_indices, + topk_weights_flat, + smem_base, + tid, + route_block_idx, + expert_idx, + output_n_tile, + active_size_m, + ) + a0_sh_rd = self._a_shared_read_offset(tid, 8) + a1_sh_rd = a0_sh_rd + Int32(self.a_sh_rd_delta_i) + + acc0 = [ + cute.make_rmem_tensor((_SCALAR_ACC_FRAGMENT_WIDTH,), cutlass.Float32) + for _ in range(16 // _SCALAR_ACC_FRAGMENT_WIDTH) + ] + acc1 = [ + cute.make_rmem_tensor((_SCALAR_ACC_FRAGMENT_WIDTH,), cutlass.Float32) + for _ in range(16 // _SCALAR_ACC_FRAGMENT_WIDTH) + ] + for frag in cutlass.range_constexpr(16 // _SCALAR_ACC_FRAGMENT_WIDTH): + acc0[frag].fill(0.0) + acc1[frag].fill(0.0) + + k_tiles = reduce_tile_count + self._prefetch_initial_tiles( + a_bf16_flat, + a_alt_bf16_flat, + b_i32_flat, + scales_i32_flat, + smem_base, + tid, + k_tiles, + reduce_k_tile, + block_valid_rows0, + block_valid_rows1, + True, + a_gl_stride, + b_gl_stride, + s_gl_stride, + scales_expert_off, + b_gl_rd_base, + a_gl_rd_row, + a_gl_rd_col0, + a_sh_wr, + a_rows_per_iter, + output_n_tile, + expert_idx, + ) + + b_scale_cur = cute.make_rmem_tensor((2, 4), Uint32) + b_scale_next = cute.make_rmem_tensor((2, 4), Uint32) + self._load_b_scale_register_bundle( + b_scale_cur, + smem_base, + tid, + b_sh_rd, + s_sh_rd, + Int32(0), + Int32(0), + ) + a0_regs_cur = cute.make_rmem_tensor((2,), Uint32) + a0_regs_next = cute.make_rmem_tensor((2,), Uint32) + a1_regs_cur = cute.make_rmem_tensor((2,), Uint32) + a1_regs_next = cute.make_rmem_tensor((2,), Uint32) + self._load_a_registers_m8_bundle( + a0_regs_cur, smem_base, a0_sh_rd, Int32(0), Int32(0) + ) + self._load_a_registers_m8_bundle( + a1_regs_cur, smem_base, a1_sh_rd, Int32(0), Int32(0) + ) + self._run_mma_pipeline_m8_pair( + a_bf16_flat, + a_alt_bf16_flat, + b_i32_flat, + scales_i32_flat, + smem_base, + tid, + acc0, + acc1, + b_scale_cur, + b_scale_next, + a0_regs_cur, + a0_regs_next, + a1_regs_cur, + a1_regs_next, + b_sh_rd, + s_sh_rd, + a0_sh_rd, + a1_sh_rd, + k_tiles, + reduce_k_tile, + block_valid_rows0, + block_valid_rows1, + a_gl_stride, + b_gl_stride, + s_gl_stride, + scales_expert_off, + b_gl_rd_base, + a_gl_rd_row, + a_gl_rd_col0, + a_sh_wr, + a_rows_per_iter, + output_n_tile, + expert_idx, + ) + + self._finish_tile( + acc0, + acc0, + acc0, + acc0, + c_bf16_flat, + c_tmp_f32_flat, + locks_i32_flat, + smem_base, + tid, + output_n_tile, + block_valid_rows0, + Int32(0), + global_scale_f32, + reduce_slice_count, + reduce_slice_idx, + lock_slot, + True, + ) + self._finish_tile( + acc1, + acc1, + acc1, + acc1, + c_bf16_flat, + c_tmp_f32_flat, + locks_i32_flat, + smem_base, + tid, + output_n_tile, + block_valid_rows1, + Int32(self.moe_block_size), + global_scale_f32, + reduce_slice_count, + reduce_slice_idx, + lock_slot + Int32(1), + True, + ) + @cute.jit def _run_tile_large_m( self, @@ -2068,6 +2413,8 @@ def _run_tile_large_m( k_tiles, reduce_k_tile, block_valid_rows, + Int32(0), + False, a_gl_stride, b_gl_stride, s_gl_stride, @@ -2149,6 +2496,7 @@ def _run_tile_large_m( tid, output_n_tile, block_valid_rows, + Int32(0), global_scale_f32, reduce_slice_count, reduce_slice_idx, @@ -2226,6 +2574,8 @@ def _run_mma_pipeline( k_tiles, reduce_k_tile, block_valid_rows, + Int32(0), + False, a_gl_stride, b_gl_stride, s_gl_stride, @@ -2327,6 +2677,168 @@ def _run_mma_pipeline( uses_m_block_8, ) + @cute.jit + def _run_mma_pipeline_m8_pair( + self, + a_bf16_flat: cute.Tensor, + a_alt_bf16_flat: cute.Tensor, + b_i32_flat: cute.Tensor, + scales_i32_flat: cute.Tensor, + smem_base: Int32, + tid: Int32, + acc0, + acc1, + b_scale_cur: cute.Tensor, + b_scale_next: cute.Tensor, + a0_regs_cur: cute.Tensor, + a0_regs_next: cute.Tensor, + a1_regs_cur: cute.Tensor, + a1_regs_next: cute.Tensor, + b_sh_rd: Int32, + s_sh_rd: Int32, + a0_sh_rd: Int32, + a1_sh_rd: Int32, + k_tiles: Int32, + reduce_k_tile: Int32, + block_valid_rows0: Int32, + block_valid_rows1: Int32, + a_gl_stride: Int32, + b_gl_stride: Int32, + s_gl_stride: Int32, + scales_expert_off: Int32, + b_gl_rd_base: Int32, + a_gl_rd_row: Int32, + a_gl_rd_col0: Int32, + a_sh_wr: Int32, + a_rows_per_iter: Int32, + output_n_tile: Int32, + expert_idx: Int32, + ): + b_frag = cute.make_rmem_tensor((2, 2), Uint32) + tile_idx = Int32(0) + while tile_idx < k_tiles: + for pipe in cutlass.range_constexpr(_STAGES): + if tile_idx < k_tiles: + for kk in cutlass.range_constexpr(self.b_sh_wr_iters): + self._load_next_fragment_bundle_m8_pair( + b_scale_next, + a0_regs_next, + a1_regs_next, + smem_base, + tid, + b_sh_rd, + s_sh_rd, + a0_sh_rd, + a1_sh_rd, + pipe, + kk, + tile_idx, + k_tiles, + ) + + self._prefetch_pipeline_step( + a_bf16_flat, + a_alt_bf16_flat, + b_i32_flat, + scales_i32_flat, + smem_base, + tid, + pipe, + kk, + tile_idx, + k_tiles, + reduce_k_tile, + block_valid_rows0, + block_valid_rows1, + True, + a_gl_stride, + b_gl_stride, + s_gl_stride, + scales_expert_off, + b_gl_rd_base, + a_gl_rd_row, + a_gl_rd_col0, + a_sh_wr, + a_rows_per_iter, + output_n_tile, + expert_idx, + ) + + for jj in cutlass.range_constexpr(4): + if cutlass.const_expr(self.weight_layout_trellis256): + self._scaled_dequant_b_fragment_trellis256( + b_frag, + b_scale_cur[0, jj], + b_scale_cur[1, jj], + ) + elif cutlass.const_expr(self.weight_layout_nf3): + lo_w = b_scale_cur[0, jj // 2] + lo16 = (lo_w >> Uint32(16 * (jj % 2))) & Uint32( + 0xFFFF + ) + hi8 = ( + b_scale_cur[0, 2] >> Uint32(8 * jj) + ) & Uint32(0xFF) + self._scaled_dequant_b_fragment_nf3( + b_frag, + lo16, + hi8, + b_scale_cur[1, jj], + ) + else: + q, s = self._select_b_scale_register( + jj, b_scale_cur + ) + self._scaled_dequant_b_fragment(b_frag, q, s) + self._mma_accumulate_m8( + acc0, + jj, + a0_regs_cur, + b_frag, + ) + self._mma_accumulate_m8( + acc1, + jj, + a1_regs_cur, + b_frag, + ) + + self._copy_a_register_bundle_m8( + a0_regs_cur, a0_regs_next + ) + self._copy_a_register_bundle_m8( + a1_regs_cur, a1_regs_next + ) + self._copy_b_scale_register_bundle( + b_scale_cur, b_scale_next + ) + tile_idx += Int32(1) + cute.arch.sync_threads() + if tile_idx < k_tiles: + self._load_b_scale_register_bundle( + b_scale_cur, + smem_base, + tid, + b_sh_rd, + s_sh_rd, + Int32(0), + Int32(0), + ) + self._load_a_registers_m8_bundle( + a0_regs_cur, + smem_base, + a0_sh_rd, + Int32(0), + Int32(0), + ) + self._load_a_registers_m8_bundle( + a1_regs_cur, + smem_base, + a1_sh_rd, + Int32(0), + Int32(0), + ) + @cute.jit def _finish_tile( self, @@ -2341,6 +2853,7 @@ def _finish_tile( tid: Int32, output_n_tile: Int32, block_valid_rows: Int32, + metadata_row_base: Int32, global_scale_f32: cutlass.Float32, reduce_slice_count: Int32, reduce_slice_idx: Int32, @@ -2385,6 +2898,7 @@ def _finish_tile( tid, output_n_tile, block_valid_rows, + metadata_row_base, global_scale_f32, ) else: @@ -2957,6 +3471,80 @@ def _load_next_fragment_bundle( uses_m_block_8, ) + @cute.jit + def _load_next_fragment_bundle_m8_pair( + self, + b_scale_next: cute.Tensor, + a0_regs_next: cute.Tensor, + a1_regs_next: cute.Tensor, + smem_base: Int32, + tid: Int32, + b_sh_rd: Int32, + s_sh_rd: Int32, + a0_sh_rd: Int32, + a1_sh_rd: Int32, + pipe: cutlass.Constexpr[int], + kk: cutlass.Constexpr[int], + tile_idx: Int32, + k_tiles: Int32, + ): + self._clear_b_scale_register_bundle(b_scale_next) + self._clear_a_register_bundle_m8(a0_regs_next) + self._clear_a_register_bundle_m8(a1_regs_next) + + if cutlass.const_expr(kk + 1 < self.b_sh_wr_iters): + if tile_idx < k_tiles: + self._load_b_scale_register_bundle( + b_scale_next, + smem_base, + tid, + b_sh_rd, + s_sh_rd, + Int32(pipe), + Int32(kk + 1), + ) + self._load_a_registers_m8_bundle( + a0_regs_next, + smem_base, + a0_sh_rd, + Int32(pipe), + Int32(kk + 1), + ) + self._load_a_registers_m8_bundle( + a1_regs_next, + smem_base, + a1_sh_rd, + Int32(pipe), + Int32(kk + 1), + ) + else: + next_tile = tile_idx + Int32(1) + if next_tile < k_tiles: + next_pipe = Int32((pipe + 1) % _STAGES) + self._load_b_scale_register_bundle( + b_scale_next, + smem_base, + tid, + b_sh_rd, + s_sh_rd, + next_pipe, + Int32(0), + ) + self._load_a_registers_m8_bundle( + a0_regs_next, + smem_base, + a0_sh_rd, + next_pipe, + Int32(0), + ) + self._load_a_registers_m8_bundle( + a1_regs_next, + smem_base, + a1_sh_rd, + next_pipe, + Int32(0), + ) + @cute.jit def _scaled_dequant_b_fragment(self, frag: cute.Tensor, q: Uint32, s: Uint32): bq1 = q @@ -3409,6 +3997,8 @@ def _stage_k_tile_async( pipe: Int32, tile_idx: Int32, block_valid_rows: Int32, + block_valid_rows1: Int32, + paired_m8: cutlass.Constexpr[bool], a_gl_stride: Int32, b_gl_stride: Int32, s_gl_stride: Int32, @@ -3423,10 +4013,24 @@ def _stage_k_tile_async( ): for i in cutlass.range_constexpr(self.a_sh_wr_iters): row = a_rows_per_iter * Int32(i) + a_gl_rd_row + metadata_row = row + route_rows = Int32(self.moe_block_size) + if cutlass.const_expr(paired_m8): + route_rows = Int32(2 * self.moe_block_size) + metadata_row = Int32(-1) + if row < Int32(self.moe_block_size): + metadata_row = row + elif ( + row >= Int32(2 * self.moe_block_size) + and row < Int32(3 * self.moe_block_size) + ): + metadata_row = row - Int32(self.moe_block_size) route_index = Int32(0) - if row < Int32(self.moe_block_size): + if metadata_row >= Int32(0) and metadata_row < route_rows: route_index = ld_shared_i32_relaxed( - smem_base + Int32(self.sh_rd_route_off * 16) + row * Int32(4) + smem_base + + Int32(self.sh_rd_route_off * 16) + + metadata_row * Int32(4) ) a_int4 = ( Int64(route_index) * Int64(a_gl_stride) @@ -3450,9 +4054,22 @@ def _stage_k_tile_async( # stage, so this adds neither shared memory nor MMA work. if output_n_tile >= Int32(self.n_tiles // 2): a_src = get_ptr_as_int64(a_alt_bf16_flat, a_int4 * Int32(8)) + row_valid = row < block_valid_rows + if cutlass.const_expr(paired_m8): + if row < Int32(self.moe_block_size): + row_valid = row < block_valid_rows + elif ( + row >= Int32(2 * self.moe_block_size) + and row < Int32(3 * self.moe_block_size) + ): + row_valid = ( + row - Int32(2 * self.moe_block_size) < block_valid_rows1 + ) + else: + row_valid = row < Int32(0) if cutlass.const_expr(self.has_k_tile_tail): a_k_int4 = tile_idx * Int32(self.a_gl_rd_delta_o) + a_gl_rd_col0 - if row < block_valid_rows and a_k_int4 < a_gl_stride: + if row_valid and a_k_int4 < a_gl_stride: cp_async4_shared_global( a_dst, a_src, @@ -3463,7 +4080,7 @@ def _stage_k_tile_async( cp_async4_shared_global_pred( a_dst, a_src, - (row < block_valid_rows).to(Int32), + row_valid.to(Int32), ) if cutlass.const_expr(self.weight_layout_trellis256): @@ -3626,6 +4243,8 @@ def _prefetch_pipeline_step( k_tiles: Int32, reduce_k_tile: Int32, block_valid_rows: Int32, + block_valid_rows1: Int32, + paired_m8: cutlass.Constexpr[bool], a_gl_stride: Int32, b_gl_stride: Int32, s_gl_stride: Int32, @@ -3651,6 +4270,8 @@ def _prefetch_pipeline_step( k_tiles, reduce_k_tile, block_valid_rows, + block_valid_rows1, + paired_m8, a_gl_stride, b_gl_stride, s_gl_stride, @@ -3676,6 +4297,8 @@ def _prefetch_initial_tiles( k_tiles: Int32, reduce_k_tile: Int32, block_valid_rows: Int32, + block_valid_rows1: Int32, + paired_m8: cutlass.Constexpr[bool], a_gl_stride: Int32, b_gl_stride: Int32, s_gl_stride: Int32, @@ -3700,6 +4323,8 @@ def _prefetch_initial_tiles( Int32(pipe), reduce_k_tile + Int32(pipe), block_valid_rows, + block_valid_rows1, + paired_m8, a_gl_stride, b_gl_stride, s_gl_stride, @@ -3731,6 +4356,8 @@ def _prefetch_lookahead_tile( k_tiles: Int32, reduce_k_tile: Int32, block_valid_rows: Int32, + block_valid_rows1: Int32, + paired_m8: cutlass.Constexpr[bool], a_gl_stride: Int32, b_gl_stride: Int32, s_gl_stride: Int32, @@ -3755,6 +4382,8 @@ def _prefetch_lookahead_tile( Int32((pipe + _STAGES - 1) % _STAGES), reduce_k_tile + fetch_tile, block_valid_rows, + block_valid_rows1, + paired_m8, a_gl_stride, b_gl_stride, s_gl_stride, @@ -3981,13 +4610,17 @@ def _drain_output_smem( c_sh_rd: Int32, c_sh_rd_delta: Int32, block_valid_rows: Int32, + metadata_row_base: Int32, store_iters: cutlass.Constexpr[int], ): for _ in cutlass.range_constexpr(store_iters): row = c_gl_wr // c_gl_stride if row < block_valid_rows: + metadata_row = metadata_row_base + row route_index = ld_shared_i32_relaxed( - smem_base + Int32(self.sh_route_off * 16) + row * Int32(4) + smem_base + + Int32(self.sh_route_off * 16) + + metadata_row * Int32(4) ) true_idx = Int64(route_index) * Int64(c_gl_stride) + Int64( c_gl_wr % c_gl_stride @@ -3997,7 +4630,9 @@ def _drain_output_smem( ) if cutlass.const_expr(self.mul_topk_weights): scale_bf2 = ld_shared_u32( - smem_base + Int32(self.sh_topk_off * 16) + row * Int32(4) + smem_base + + Int32(self.sh_topk_off * 16) + + metadata_row * Int32(4) ) q0 = self._elem2_mul(q0, scale_bf2) q1 = self._elem2_mul(q1, scale_bf2) @@ -4047,14 +4682,18 @@ def _drain_output_smem_tail( c_sh_rd: Int32, c_sh_rd_delta: Int32, block_valid_rows: Int32, + metadata_row_base: Int32, store_iters: cutlass.Constexpr[int], ): for _ in cutlass.range_constexpr(store_iters): row = c_gl_wr // c_gl_stride_covered col_word = c_gl_wr - row * c_gl_stride_covered if row < block_valid_rows and col_word < c_gl_stride: + metadata_row = metadata_row_base + row route_index = ld_shared_i32_relaxed( - smem_base + Int32(self.sh_route_off * 16) + row * Int32(4) + smem_base + + Int32(self.sh_route_off * 16) + + metadata_row * Int32(4) ) true_idx = Int64(route_index) * Int64(c_gl_stride) + Int64(col_word) q0, q1, q2, q3 = ld_shared_v4_u32( @@ -4062,7 +4701,9 @@ def _drain_output_smem_tail( ) if cutlass.const_expr(self.mul_topk_weights): scale_bf2 = ld_shared_u32( - smem_base + Int32(self.sh_topk_off * 16) + row * Int32(4) + smem_base + + Int32(self.sh_topk_off * 16) + + metadata_row * Int32(4) ) q0 = self._elem2_mul(q0, scale_bf2) q1 = self._elem2_mul(q1, scale_bf2) @@ -4102,6 +4743,7 @@ def _store_tile_m8( tid: Int32, output_n_tile: Int32, block_valid_rows: Int32, + metadata_row_base: Int32, global_scale_f32: cutlass.Float32, ): if cutlass.const_expr(self.has_n_tile_tail): @@ -4184,6 +4826,7 @@ def _store_tile_m8( c_sh_rd, c_sh_rd_delta, block_valid_rows, + metadata_row_base, store_iters, ) else: @@ -4196,6 +4839,7 @@ def _store_tile_m8( c_sh_rd, c_sh_rd_delta, block_valid_rows, + metadata_row_base, store_iters, ) @@ -4495,6 +5139,7 @@ def _store_tile_large_m( c_sh_rd, c_sh_rd_delta, block_valid_rows, + Int32(0), store_iters, ) else: @@ -4507,6 +5152,7 @@ def _store_tile_large_m( c_sh_rd, c_sh_rd_delta, block_valid_rows, + Int32(0), store_iters, ) @@ -4585,6 +5231,8 @@ def __init__( fc2_tile_k: int, moe_block_size: int, max_m_blocks: int, + fc2_moe_block_size: int | None = None, + fc2_schedule_route_block_factor: int = 1, element_dtype: str = "bf16", fast_math: bool = True, swiglu_limit: float | None = None, @@ -4648,6 +5296,34 @@ def __init__( self.num_experts = int(num_experts) self.top_k = int(top_k) self.moe_block_size = int(moe_block_size) + self.fc2_moe_block_size = int( + moe_block_size if fc2_moe_block_size is None else fc2_moe_block_size + ) + self.fc2_schedule_route_block_factor = int( + fc2_schedule_route_block_factor + ) + if ( + self.fc2_moe_block_size not in _ALLOWED_ROUTED_SIZES + or self.moe_block_size % self.fc2_moe_block_size != 0 + ): + raise ValueError( + "FC2 route subtile must be an allowed divisor of the packed " + f"route block: packed={self.moe_block_size}, " + f"fc2={self.fc2_moe_block_size}" + ) + expected_fc2_schedule_factor = ( + self.moe_block_size // self.fc2_moe_block_size + ) + if ( + self.fc2_schedule_route_block_factor < 1 + or expected_fc2_schedule_factor % self.fc2_schedule_route_block_factor + != 0 + ): + raise ValueError( + "FC2 schedule factor must divide one packed route block: " + f"factor={self.fc2_schedule_route_block_factor}, " + f"maximum={expected_fc2_schedule_factor}" + ) self.activation = activation self.activation_is_gated = is_gated self.activation_is_situ = activation == SITU @@ -4752,8 +5428,10 @@ def __init__( ), tile_n=fc2_tile_n, tile_k=fc2_tile_k, - moe_block_size=moe_block_size, - max_m_blocks=max_m_blocks, + moe_block_size=self.fc2_moe_block_size, + max_m_blocks=( + max_m_blocks * self.moe_block_size // self.fc2_moe_block_size + ), element_dtype=element_dtype, weight_layout=weight_layout, scale_format=scale_format, @@ -4764,6 +5442,11 @@ def __init__( fused_topk_sum=self.tc_decode_fused_sum, fused_sum_topk=int(top_k), schedule_whole_tiles=self.schedule_whole_tiles, + schedule_route_block_factor=self.fc2_schedule_route_block_factor, + paired_m8_routes=( + self.fc2_moe_block_size == 8 + and self.fc2_schedule_route_block_factor == 2 + ), ) self.cta_threads = max(self.fc1.cta_threads, self.fc2.cta_threads) if self.fc1.cta_threads != self.fc2.cta_threads: @@ -5067,7 +5750,6 @@ def _moe_body( active_m, ) self._grid_barrier(locks_i32_flat, tid, grid_x) - if cutlass.const_expr(self.tc_decode_fused_sum): # The TC-decode FC2 epilogue atomically accumulates per-route # partials directly into the per-token output, so the output must be @@ -5202,7 +5884,6 @@ def _moe_body( active_m * Int32(self.top_k), fc2_emit_tile, ) - @cute.jit def _grid_barrier( self, diff --git a/sparkinfer/moe/_shared/kernels/w4a16/mixed_trellis.py b/sparkinfer/moe/_shared/kernels/w4a16/mixed_trellis.py index 71fce755..8d8c6a51 100644 --- a/sparkinfer/moe/_shared/kernels/w4a16/mixed_trellis.py +++ b/sparkinfer/moe/_shared/kernels/w4a16/mixed_trellis.py @@ -55,6 +55,9 @@ class MixedTrellisCompileResult: fc2_tile_k: int fc2_tile_n: int moe_block_size: int + fc2_moe_block_size: int + fc2_schedule_route_block_factor: int + fc2_paired_m8_routes: bool max_m_blocks: int blocks_per_sm: int sms: int @@ -109,7 +112,7 @@ class MixedTrellisTier(Protocol): class W4A16MixedTrellisKernel: """One cooperative grid over two native Trellis bitrates.""" - ABI_VERSION = 1 + ABI_VERSION = 4 def __init__( self, @@ -145,12 +148,36 @@ def __init__( for phase in ("fc1", "fc2"): gemms = tuple(getattr(moe, phase) for moe in (driver, tier0, tier1)) geometry = tuple( - (g.n_tiles, g.k_tiles, g.tile_n, g.tile_k, g.cta_threads) for g in gemms + ( + g.n_tiles, + g.k_tiles, + g.tile_n, + g.tile_k, + g.cta_threads, + g.moe_block_size, + g.schedule_route_block_factor, + g.paired_m8_routes, + ) + for g in gemms ) if geometry[1:] != geometry[:-1]: raise ValueError( f"mixed Trellis kernels disagree on {phase} geometry: {geometry}" ) + fc2_factor = int(driver.fc2.schedule_route_block_factor) + expected_factor = int(driver.moe_block_size // driver.fc2.moe_block_size) + if fc2_factor < 1 or expected_factor % fc2_factor != 0: + raise ValueError( + "mixed Trellis FC2 schedule factor must divide one packed " + f"route block: factor={fc2_factor}, maximum={expected_factor}" + ) + expected_pair = fc2_factor == 2 and driver.fc2.moe_block_size == 8 + if bool(driver.fc2.paired_m8_routes) != expected_pair: + raise ValueError( + "mixed Trellis FC2 pair contract mismatch: " + f"factor={fc2_factor}, m={driver.fc2.moe_block_size}, " + f"paired={driver.fc2.paired_m8_routes}" + ) if tier0.num_experts > 256 or tier1.num_experts > 256: raise ValueError("tier-local expert ids must fit in eight bits") if driver.num_experts != tier0.num_experts + tier1.num_experts: @@ -214,7 +241,16 @@ def _emit_tier_tile( reduce_slice_idx: Int32, lock_slot: Int32, ): - combined_expert = block_expert_ids[route_block_idx].to(Int32) + metadata_block_idx = route_block_idx + if cutlass.const_expr(not is_fc1): + metadata_block_idx = route_block_idx // Int32( + self.driver.moe_block_size + // ( + self.driver.fc2.moe_block_size + * self.driver.fc2.schedule_route_block_factor + ) + ) + combined_expert = block_expert_ids[metadata_block_idx].to(Int32) if combined_expert >= Int32(0) and combined_expert < Int32(self.total_experts): descriptor = descriptor_map[combined_expert].to(Int32) if descriptor >= Int32(0): @@ -225,57 +261,135 @@ def _emit_tier_tile( gemm = self.tier0.fc1 else: gemm = self.tier0.fc2 - gemm._run_tile( - a_flat, - a_alt_flat, - t0_b_flat, - c_flat, - t0_scales_flat, - t0_global_scale, - packed_route_indices, - topk_weights, - c_tmp, - locks, - smem_base, - tid, - route_block_idx, - local_expert, - output_n_tile, - reduce_k_tile, - reduce_tile_count, - reduce_slice_count, - reduce_slice_idx, - lock_slot, - active_size_m, - ) + if cutlass.const_expr(gemm.paired_m8_routes): + tile_route_block_idx = ( + route_block_idx + * Int32(gemm.schedule_route_block_factor) + ) + gemm._run_tile_m8_pair( + a_flat, + a_alt_flat, + t0_b_flat, + c_flat, + t0_scales_flat, + t0_global_scale, + packed_route_indices, + topk_weights, + c_tmp, + locks, + smem_base, + tid, + tile_route_block_idx, + local_expert, + output_n_tile, + reduce_k_tile, + reduce_tile_count, + reduce_slice_count, + reduce_slice_idx, + lock_slot * Int32(gemm.schedule_route_block_factor), + active_size_m, + ) + else: + for subtile in cutlass.range_constexpr( + gemm.schedule_route_block_factor + ): + tile_route_block_idx = ( + route_block_idx + * Int32(gemm.schedule_route_block_factor) + + Int32(subtile) + ) + gemm._run_tile( + a_flat, + a_alt_flat, + t0_b_flat, + c_flat, + t0_scales_flat, + t0_global_scale, + packed_route_indices, + topk_weights, + c_tmp, + locks, + smem_base, + tid, + tile_route_block_idx, + local_expert, + output_n_tile, + reduce_k_tile, + reduce_tile_count, + reduce_slice_count, + reduce_slice_idx, + lock_slot + * Int32(gemm.schedule_route_block_factor) + + Int32(subtile), + active_size_m, + ) elif tier == Int32(1) and local_expert < Int32(self.tier1.num_experts): if cutlass.const_expr(is_fc1): gemm = self.tier1.fc1 else: gemm = self.tier1.fc2 - gemm._run_tile( - a_flat, - a_alt_flat, - t1_b_flat, - c_flat, - t1_scales_flat, - t1_global_scale, - packed_route_indices, - topk_weights, - c_tmp, - locks, - smem_base, - tid, - route_block_idx, - local_expert, - output_n_tile, - reduce_k_tile, - reduce_tile_count, - reduce_slice_count, - reduce_slice_idx, - lock_slot, - active_size_m, - ) + if cutlass.const_expr(gemm.paired_m8_routes): + tile_route_block_idx = ( + route_block_idx + * Int32(gemm.schedule_route_block_factor) + ) + gemm._run_tile_m8_pair( + a_flat, + a_alt_flat, + t1_b_flat, + c_flat, + t1_scales_flat, + t1_global_scale, + packed_route_indices, + topk_weights, + c_tmp, + locks, + smem_base, + tid, + tile_route_block_idx, + local_expert, + output_n_tile, + reduce_k_tile, + reduce_tile_count, + reduce_slice_count, + reduce_slice_idx, + lock_slot * Int32(gemm.schedule_route_block_factor), + active_size_m, + ) + else: + for subtile in cutlass.range_constexpr( + gemm.schedule_route_block_factor + ): + tile_route_block_idx = ( + route_block_idx + * Int32(gemm.schedule_route_block_factor) + + Int32(subtile) + ) + gemm._run_tile( + a_flat, + a_alt_flat, + t1_b_flat, + c_flat, + t1_scales_flat, + t1_global_scale, + packed_route_indices, + topk_weights, + c_tmp, + locks, + smem_base, + tid, + tile_route_block_idx, + local_expert, + output_n_tile, + reduce_k_tile, + reduce_tile_count, + reduce_slice_count, + reduce_slice_idx, + lock_slot + * Int32(gemm.schedule_route_block_factor) + + Int32(subtile), + active_size_m, + ) @cute.jit def __call__( @@ -557,6 +671,8 @@ def make_kernel(num_experts: int, bits: int) -> W4A16FusedMoeKernel: fc2_tile_k=fc2_tile_k, moe_block_size=moe_block_size, max_m_blocks=max_m_blocks, + fc2_moe_block_size=(8 if int(moe_block_size) == 64 else moe_block_size), + fc2_schedule_route_block_factor=(2 if int(moe_block_size) == 64 else 1), element_dtype="fp16", weight_layout="trellis3_t256", scale_format="e4m3_k32", @@ -573,9 +689,15 @@ def make_kernel(num_experts: int, bits: int) -> W4A16FusedMoeKernel: tier0=make_kernel(int(tier0_num_experts), int(tier0_bits)), tier1=make_kernel(int(tier1_num_experts), int(tier1_bits)), ) - if kernel.shared_words * 4 > int(max_shared_mem) - 512: + # shared_words is the complete dynamically allocated MemRange used by the + # cooperative kernel. CUDA permits a launch exactly at the device's + # opt-in shared-memory limit; rejecting an additional 512 bytes here + # unnecessarily excludes the stock mixed-K tile geometry at block-64. + if kernel.shared_words * 4 > int(max_shared_mem): raise ValueError( - "mixed Trellis shared-memory requirement exceeds the device limit" + "mixed Trellis shared-memory requirement exceeds the device limit: " + f"required={kernel.shared_words * 4} " + f"limit={int(max_shared_mem)}" ) device = int(torch.cuda.current_device()) cache_key = ( @@ -698,6 +820,11 @@ def tier_args(experts: int, bits: int): fc2_tile_k=fc2_tile_k, fc2_tile_n=fc2_tile_n, moe_block_size=int(moe_block_size), + fc2_moe_block_size=int(kernel.driver.fc2.moe_block_size), + fc2_schedule_route_block_factor=int( + kernel.driver.fc2.schedule_route_block_factor + ), + fc2_paired_m8_routes=bool(kernel.driver.fc2.paired_m8_routes), max_m_blocks=int(max_m_blocks), blocks_per_sm=int(kernel.blocks_per_sm), sms=int(sms), diff --git a/tests/moe/test_w4a16_mixed_trellis.py b/tests/moe/test_w4a16_mixed_trellis.py index 3949583a..32810e62 100644 --- a/tests/moe/test_w4a16_mixed_trellis.py +++ b/tests/moe/test_w4a16_mixed_trellis.py @@ -72,6 +72,7 @@ def _serial_tier( topk_weights: torch.Tensor, topk_ids: torch.Tensor, expert_map: torch.Tensor, + block_size_m: int = 8, ) -> torch.Tensor: m, topk = int(topk_ids.shape[0]), int(topk_ids.shape[1]) buffers = make_w4a16_packed_buffers( @@ -82,7 +83,7 @@ def _serial_tier( device=x.device, route_num_experts=int(expert_map.numel()), full_rotation=True, - block_size_m=8, + block_size_m=block_size_m, ) assert buffers.rotation_a_gate is not None assert buffers.rotation_a_up is not None @@ -104,7 +105,7 @@ def _serial_tier( expert_counts=buffers.expert_counts, expert_map=expert_map, output_expert_map=expert_map, - route_block_size_m=8, + route_block_size_m=block_size_m, intermediate_rotation_scales=prepared.intermediate_rotations, full_rotation=True, suh_gate_table=prepared.gate_suh, @@ -260,6 +261,118 @@ def test_build_tiered_maps_rejects_invalid_partitions() -> None: build_tiered_maps((0, 4), (1, 2), device=torch.device("cpu")) +@pytest.mark.skipif(not _sm12x_available(), reason="requires an SM120/SM121 GPU") +def test_one_grid_block64_avoids_serial_prefill_drift() -> None: + """Block-64 packing plus FC2 subtiles preserves stock one-grid arithmetic.""" + + torch.manual_seed(20260801) + device = torch.device("cuda", torch.cuda.current_device()) + m, hidden, intermediate, topk = 64, 512, 256, 8 + tile_config = (128, 128, 32, 512) + tier0_experts, tier1_experts = 6, 2 + + prepared_tiers = tuple( + _prepared( + experts=experts, + hidden=hidden, + intermediate=intermediate, + bits=bits, + seed=seed, + device=device, + tile_config=tile_config, + ) + for experts, bits, seed in ( + (tier0_experts, 3, 301), + (tier1_experts, 4, 401), + ) + ) + + x = (torch.randn((m, hidden), device=device) * 1.0e-3).to(torch.bfloat16) + topk_ids = torch.tensor( + [0, 6, 1, 7, 2, 3, 4, 5], dtype=torch.int32, device=device + ).expand(m, -1).contiguous() + topk_weights = torch.softmax( + torch.randn((m, topk), dtype=torch.float32, device=device), dim=-1 + ) + props = torch.cuda.get_device_properties(device) + global_to_combined, descriptor = build_tiered_maps( + range(tier0_experts), range(tier0_experts, 8), device=device + ) + + def one_grid( + block_size_m: int, + ) -> tuple[torch.Tensor, object, tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + route_slots = max_packed_route_slots(m * topk, block_size_m, 8) + launch = compile_mixed_trellis( + size_m=m, + hidden_size=hidden, + intermediate_size=intermediate, + tier0_num_experts=tier0_experts, + tier1_num_experts=tier1_experts, + top_k=topk, + max_m_blocks=(route_slots + block_size_m - 1) // block_size_m, + moe_block_size=block_size_m, + sms=int(props.multi_processor_count), + max_shared_mem=int(props.shared_memory_per_block_optin), + force_tile_config=tile_config, + ) + buffers = make_mixed_trellis_buffers( + launch, device=device, sms=int(props.multi_processor_count) + ) + output = run_mixed_trellis( + x, + prepared_tiers[0], + prepared_tiers[1], + topk_weights, + topk_ids, + global_to_combined, + descriptor, + combine_trellis_rotations(*prepared_tiers), + launch, + buffers, + ).clone() + phase_outputs = ( + buffers.fc1.clone(), + buffers.activated.clone(), + buffers.fc2.clone(), + ) + return output, launch, phase_outputs + + reference, reference_launch, reference_phases = one_grid(8) + candidate, candidate_launch, candidate_phases = one_grid(64) + torch.cuda.synchronize(device) + + phase_equal = tuple( + torch.equal(candidate_phase, reference_phase) + for candidate_phase, reference_phase in zip( + candidate_phases, reference_phases, strict=True + ) + ) + print( + "mixed_trellis_block64_fc2_subtile_parity " + f"phases={phase_equal} final={torch.equal(candidate, reference)} " + f"packed_block={candidate_launch.moe_block_size} " + f"fc2_subtile={candidate_launch.fc2_moe_block_size} " + f"fc2_schedule_factor={candidate_launch.fc2_schedule_route_block_factor} " + f"regs={candidate_launch.registers_per_thread} " + f"local={candidate_launch.local_memory_bytes} " + f"smem={candidate_launch.shared_memory_bytes}" + ) + assert reference_launch.moe_block_size == 8 + assert reference_launch.fc2_moe_block_size == 8 + assert reference_launch.fc2_schedule_route_block_factor == 1 + assert candidate_launch.moe_block_size == 64 + assert candidate_launch.fc2_moe_block_size == 8 + assert candidate_launch.fc2_schedule_route_block_factor == 2 + assert candidate_launch.fc2_paired_m8_routes is True + assert phase_equal == (True, True, True) + assert torch.equal(candidate, reference) + assert candidate_launch.local_memory_bytes == 0 + assert candidate_launch.shared_memory_bytes <= int( + props.shared_memory_per_block_optin + ) + + @pytest.mark.skipif(not _sm12x_available(), reason="requires an SM120/SM121 GPU") def test_glm52_large_m_mixed_k3_k4_matches_serial() -> None: """Cover the production prefill shape that exposed lost FC1 reductions."""