From 2c552b49e9fd28dfe9e95736b1351ec367927a39 Mon Sep 17 00:00:00 2001 From: derek Date: Sat, 1 Aug 2026 10:26:59 -0400 Subject: [PATCH 1/4] test: prove one-grid block64 prefill parity Assisted-by: OpenAI Codex --- tests/moe/test_w4a16_mixed_trellis.py | 135 +++++++++++++++++++++++++- 1 file changed, 133 insertions(+), 2 deletions(-) diff --git a/tests/moe/test_w4a16_mixed_trellis.py b/tests/moe/test_w4a16_mixed_trellis.py index 3949583a..58178a9f 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,136 @@ 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: + """Keep block-64 routing without grouping K3 and K4 reductions.""" + + torch.manual_seed(20260801) + device = torch.device("cuda", torch.cuda.current_device()) + m, hidden, intermediate, topk = 64, 512, 256, 8 + reference_tiles = (128, 128, 32, 512) + block64_tiles = (128, 64, 64, 128) + tier0_experts, tier1_experts = 6, 2 + + def tiers(tile_config: tuple[int, int, int, int]): + return 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), + ) + ) + + reference_tiers = tiers(reference_tiles) + block64_tiers = tiers(block64_tiles) + for reference, candidate in zip(reference_tiers, block64_tiers, strict=True): + assert torch.equal(reference.w13, candidate.w13) + assert torch.equal(reference.w2, candidate.w2) + + 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 + ) + map0 = torch.cat( + ( + torch.arange(tier0_experts, dtype=torch.int32, device=device), + torch.full((tier1_experts,), -1, dtype=torch.int32, device=device), + ) + ) + map1 = torch.cat( + ( + torch.full((tier0_experts,), -1, dtype=torch.int32, device=device), + torch.arange(tier1_experts, dtype=torch.int32, device=device), + ) + ) + + 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( + prepared_tiers, + tile_config: tuple[int, int, int, int], + block_size_m: int, + ) -> 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) + ) + return 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() + + reference = one_grid(reference_tiers, reference_tiles, 8) + candidate = one_grid(block64_tiers, block64_tiles, 64) + serial_block8 = _serial_tier( + x, reference_tiers[0], topk_weights, topk_ids, map0, block_size_m=8 + ).clone() + serial_block8.add_( + _serial_tier( + x, reference_tiers[1], topk_weights, topk_ids, map1, block_size_m=8 + ) + ) + serial_block64 = _serial_tier( + x, block64_tiers[0], topk_weights, topk_ids, map0, block_size_m=64 + ).clone() + serial_block64.add_( + _serial_tier( + x, block64_tiers[1], topk_weights, topk_ids, map1, block_size_m=64 + ) + ) + torch.cuda.synchronize(device) + + denominator = reference.norm().clamp_min(1.0e-12) + serial8_error = float((serial_block8 - reference).norm() / denominator) + serial64_error = float((serial_block64 - reference).norm() / denominator) + candidate_error = float((candidate - reference).norm() / denominator) + print( + "mixed_trellis_block64_numerics " + f"serial_block8={serial8_error:.9e} " + f"serial_block64={serial64_error:.9e} " + f"one_grid_block64={candidate_error:.9e}" + ) + assert serial64_error > 5.0e-5 + assert candidate_error < 2.0e-5 + assert candidate_error < serial64_error * 0.2 + assert abs(candidate_error - serial8_error) < 1.0e-6 + + @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.""" From fab05be829edec3c560cc92e2427667ab9335c66 Mon Sep 17 00:00:00 2001 From: derek Date: Sat, 1 Aug 2026 11:30:18 -0400 Subject: [PATCH 2/4] moe: preserve mixed prefill arithmetic with FC2 subtiles --- .../moe/_shared/kernels/w4a16/kernel.py | 24 +++- .../_shared/kernels/w4a16/mixed_trellis.py | 32 ++++- tests/moe/test_w4a16_mixed_trellis.py | 120 +++++++----------- 3 files changed, 98 insertions(+), 78 deletions(-) diff --git a/sparkinfer/moe/_shared/kernels/w4a16/kernel.py b/sparkinfer/moe/_shared/kernels/w4a16/kernel.py index bcd4cca4..35c97645 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, } @@ -4585,6 +4590,7 @@ def __init__( fc2_tile_k: int, moe_block_size: int, max_m_blocks: int, + fc2_moe_block_size: int | None = None, element_dtype: str = "bf16", fast_math: bool = True, swiglu_limit: float | None = None, @@ -4648,6 +4654,18 @@ 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 + ) + 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}" + ) self.activation = activation self.activation_is_gated = is_gated self.activation_is_situ = activation == SITU @@ -4752,8 +4770,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, diff --git a/sparkinfer/moe/_shared/kernels/w4a16/mixed_trellis.py b/sparkinfer/moe/_shared/kernels/w4a16/mixed_trellis.py index 71fce755..8bdc5753 100644 --- a/sparkinfer/moe/_shared/kernels/w4a16/mixed_trellis.py +++ b/sparkinfer/moe/_shared/kernels/w4a16/mixed_trellis.py @@ -55,6 +55,7 @@ class MixedTrellisCompileResult: fc2_tile_k: int fc2_tile_n: int moe_block_size: int + fc2_moe_block_size: int max_m_blocks: int blocks_per_sm: int sms: int @@ -109,7 +110,7 @@ class MixedTrellisTier(Protocol): class W4A16MixedTrellisKernel: """One cooperative grid over two native Trellis bitrates.""" - ABI_VERSION = 1 + ABI_VERSION = 2 def __init__( self, @@ -145,7 +146,15 @@ 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, + ) + for g in gemms ) if geometry[1:] != geometry[:-1]: raise ValueError( @@ -214,7 +223,12 @@ 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 + ) + 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): @@ -557,6 +571,7 @@ 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), element_dtype="fp16", weight_layout="trellis3_t256", scale_format="e4m3_k32", @@ -573,9 +588,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 +719,7 @@ 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), 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 58178a9f..1f461a30 100644 --- a/tests/moe/test_w4a16_mixed_trellis.py +++ b/tests/moe/test_w4a16_mixed_trellis.py @@ -263,37 +263,29 @@ def test_build_tiered_maps_rejects_invalid_partitions() -> None: @pytest.mark.skipif(not _sm12x_available(), reason="requires an SM120/SM121 GPU") def test_one_grid_block64_avoids_serial_prefill_drift() -> None: - """Keep block-64 routing without grouping K3 and K4 reductions.""" + """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 - reference_tiles = (128, 128, 32, 512) - block64_tiles = (128, 64, 64, 128) + tile_config = (128, 128, 32, 512) tier0_experts, tier1_experts = 6, 2 - def tiers(tile_config: tuple[int, int, int, int]): - return 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), - ) + prepared_tiers = tuple( + _prepared( + experts=experts, + hidden=hidden, + intermediate=intermediate, + bits=bits, + seed=seed, + device=device, + tile_config=tile_config, ) - - reference_tiers = tiers(reference_tiles) - block64_tiers = tiers(block64_tiles) - for reference, candidate in zip(reference_tiers, block64_tiers, strict=True): - assert torch.equal(reference.w13, candidate.w13) - assert torch.equal(reference.w2, candidate.w2) + 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( @@ -302,29 +294,14 @@ def tiers(tile_config: tuple[int, int, int, int]): topk_weights = torch.softmax( torch.randn((m, topk), dtype=torch.float32, device=device), dim=-1 ) - map0 = torch.cat( - ( - torch.arange(tier0_experts, dtype=torch.int32, device=device), - torch.full((tier1_experts,), -1, dtype=torch.int32, device=device), - ) - ) - map1 = torch.cat( - ( - torch.full((tier0_experts,), -1, dtype=torch.int32, device=device), - torch.arange(tier1_experts, dtype=torch.int32, device=device), - ) - ) - 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( - prepared_tiers, - tile_config: tuple[int, int, int, int], block_size_m: int, - ) -> torch.Tensor: + ) -> 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, @@ -342,7 +319,7 @@ def one_grid( buffers = make_mixed_trellis_buffers( launch, device=device, sms=int(props.multi_processor_count) ) - return run_mixed_trellis( + output = run_mixed_trellis( x, prepared_tiers[0], prepared_tiers[1], @@ -354,41 +331,42 @@ def one_grid( launch, buffers, ).clone() - - reference = one_grid(reference_tiers, reference_tiles, 8) - candidate = one_grid(block64_tiers, block64_tiles, 64) - serial_block8 = _serial_tier( - x, reference_tiers[0], topk_weights, topk_ids, map0, block_size_m=8 - ).clone() - serial_block8.add_( - _serial_tier( - x, reference_tiers[1], topk_weights, topk_ids, map1, block_size_m=8 - ) - ) - serial_block64 = _serial_tier( - x, block64_tiers[0], topk_weights, topk_ids, map0, block_size_m=64 - ).clone() - serial_block64.add_( - _serial_tier( - x, block64_tiers[1], topk_weights, topk_ids, map1, block_size_m=64 + 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) - denominator = reference.norm().clamp_min(1.0e-12) - serial8_error = float((serial_block8 - reference).norm() / denominator) - serial64_error = float((serial_block64 - reference).norm() / denominator) - candidate_error = float((candidate - reference).norm() / denominator) + 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_numerics " - f"serial_block8={serial8_error:.9e} " - f"serial_block64={serial64_error:.9e} " - f"one_grid_block64={candidate_error:.9e}" + "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"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 candidate_launch.moe_block_size == 64 + assert candidate_launch.fc2_moe_block_size == 8 + 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 ) - assert serial64_error > 5.0e-5 - assert candidate_error < 2.0e-5 - assert candidate_error < serial64_error * 0.2 - assert abs(candidate_error - serial8_error) < 1.0e-6 @pytest.mark.skipif(not _sm12x_available(), reason="requires an SM120/SM121 GPU") From d218d82f105181c3c0e6ee98bc758812df7429a6 Mon Sep 17 00:00:00 2001 From: derek Date: Sat, 1 Aug 2026 13:05:00 -0400 Subject: [PATCH 3/4] Optimize exact block64 mixed Trellis FC2 scheduling --- .../moe/_shared/kernels/w4a16/kernel.py | 33 ++++- .../_shared/kernels/w4a16/mixed_trellis.py | 133 +++++++++++------- tests/moe/test_w4a16_mixed_trellis.py | 3 + 3 files changed, 120 insertions(+), 49 deletions(-) diff --git a/sparkinfer/moe/_shared/kernels/w4a16/kernel.py b/sparkinfer/moe/_shared/kernels/w4a16/kernel.py index 35c97645..61aa550a 100644 --- a/sparkinfer/moe/_shared/kernels/w4a16/kernel.py +++ b/sparkinfer/moe/_shared/kernels/w4a16/kernel.py @@ -706,6 +706,7 @@ def __init__( fused_topk_sum: bool = False, fused_sum_topk: int = 1, schedule_whole_tiles: bool = False, + schedule_route_block_factor: int = 1, ): if element_dtype not in {"bf16", "fp16"}: raise ValueError(f"unsupported element_dtype {element_dtype!r}") @@ -881,6 +882,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 @@ -1039,6 +1052,7 @@ 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, ) @cute.jit @@ -1294,7 +1308,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 @@ -4591,6 +4605,7 @@ def __init__( 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, @@ -4657,6 +4672,9 @@ def __init__( 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 @@ -4666,6 +4684,18 @@ def __init__( 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 not in ( + 1, + expected_fc2_schedule_factor, + ): + raise ValueError( + "FC2 schedule factor must be one or cover one packed route " + f"block: factor={self.fc2_schedule_route_block_factor}, " + f"expected={expected_fc2_schedule_factor}" + ) self.activation = activation self.activation_is_gated = is_gated self.activation_is_situ = activation == SITU @@ -4784,6 +4814,7 @@ 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, ) self.cta_threads = max(self.fc1.cta_threads, self.fc2.cta_threads) if self.fc1.cta_threads != self.fc2.cta_threads: diff --git a/sparkinfer/moe/_shared/kernels/w4a16/mixed_trellis.py b/sparkinfer/moe/_shared/kernels/w4a16/mixed_trellis.py index 8bdc5753..251cbf33 100644 --- a/sparkinfer/moe/_shared/kernels/w4a16/mixed_trellis.py +++ b/sparkinfer/moe/_shared/kernels/w4a16/mixed_trellis.py @@ -56,6 +56,7 @@ class MixedTrellisCompileResult: fc2_tile_n: int moe_block_size: int fc2_moe_block_size: int + fc2_schedule_route_block_factor: int max_m_blocks: int blocks_per_sm: int sms: int @@ -110,7 +111,7 @@ class MixedTrellisTier(Protocol): class W4A16MixedTrellisKernel: """One cooperative grid over two native Trellis bitrates.""" - ABI_VERSION = 2 + ABI_VERSION = 3 def __init__( self, @@ -153,6 +154,7 @@ def __init__( g.tile_k, g.cta_threads, g.moe_block_size, + g.schedule_route_block_factor, ) for g in gemms ) @@ -160,6 +162,13 @@ def __init__( 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 not in (1, expected_factor): + raise ValueError( + "mixed Trellis FC2 schedule factor must be one or cover one " + f"packed route block: factor={fc2_factor}, expected={expected_factor}" + ) 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: @@ -226,7 +235,11 @@ def _emit_tier_tile( 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.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): @@ -239,57 +252,77 @@ 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, - ) + 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, - ) + 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__( @@ -572,6 +605,7 @@ def make_kernel(num_experts: int, bits: int) -> W4A16FusedMoeKernel: 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=(8 if int(moe_block_size) == 64 else 1), element_dtype="fp16", weight_layout="trellis3_t256", scale_format="e4m3_k32", @@ -720,6 +754,9 @@ def tier_args(experts: int, bits: int): 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 + ), 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 1f461a30..7910dcdb 100644 --- a/tests/moe/test_w4a16_mixed_trellis.py +++ b/tests/moe/test_w4a16_mixed_trellis.py @@ -353,14 +353,17 @@ def one_grid( 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 == 8 assert phase_equal == (True, True, True) assert torch.equal(candidate, reference) assert candidate_launch.local_memory_bytes == 0 From 1c2b052c426ea861c140f0981a7ea78709c9fdde Mon Sep 17 00:00:00 2001 From: derek Date: Sat, 1 Aug 2026 14:48:05 -0400 Subject: [PATCH 4/4] moe: share exact FC2 weights across paired M8 tiles --- .../moe/_shared/kernels/w4a16/kernel.py | 668 +++++++++++++++++- .../_shared/kernels/w4a16/mixed_trellis.py | 110 ++- tests/moe/test_w4a16_mixed_trellis.py | 3 +- 3 files changed, 740 insertions(+), 41 deletions(-) diff --git a/sparkinfer/moe/_shared/kernels/w4a16/kernel.py b/sparkinfer/moe/_shared/kernels/w4a16/kernel.py index 61aa550a..c0cb0e9b 100644 --- a/sparkinfer/moe/_shared/kernels/w4a16/kernel.py +++ b/sparkinfer/moe/_shared/kernels/w4a16/kernel.py @@ -707,6 +707,7 @@ def __init__( 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}") @@ -908,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()) @@ -939,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 @@ -980,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 ) @@ -1053,6 +1072,7 @@ def __cache_key__(self) -> tuple[object, ...]: self.blocks_per_sm, self.schedule_whole_tiles, self.schedule_route_block_factor, + self.paired_m8_routes, ) @cute.jit @@ -1632,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, @@ -1763,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) @@ -1902,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, @@ -1983,12 +2120,201 @@ def _run_tile_m8( tid, output_n_tile, block_valid_rows, + Int32(0), + global_scale_f32, + reduce_slice_count, + reduce_slice_idx, + lock_slot, + 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( @@ -2087,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, @@ -2168,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, @@ -2245,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, @@ -2346,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, @@ -2360,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, @@ -2404,6 +2898,7 @@ def _finish_tile( tid, output_n_tile, block_valid_rows, + metadata_row_base, global_scale_f32, ) else: @@ -2976,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 @@ -3428,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, @@ -3442,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) @@ -3469,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, @@ -3482,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): @@ -3645,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, @@ -3670,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, @@ -3695,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, @@ -3719,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, @@ -3750,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, @@ -3774,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, @@ -4000,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 @@ -4016,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) @@ -4066,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( @@ -4081,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) @@ -4121,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): @@ -4203,6 +4826,7 @@ def _store_tile_m8( c_sh_rd, c_sh_rd_delta, block_valid_rows, + metadata_row_base, store_iters, ) else: @@ -4215,6 +4839,7 @@ def _store_tile_m8( c_sh_rd, c_sh_rd_delta, block_valid_rows, + metadata_row_base, store_iters, ) @@ -4514,6 +5139,7 @@ def _store_tile_large_m( c_sh_rd, c_sh_rd_delta, block_valid_rows, + Int32(0), store_iters, ) else: @@ -4526,6 +5152,7 @@ def _store_tile_large_m( c_sh_rd, c_sh_rd_delta, block_valid_rows, + Int32(0), store_iters, ) @@ -4687,14 +5314,15 @@ def __init__( expected_fc2_schedule_factor = ( self.moe_block_size // self.fc2_moe_block_size ) - if self.fc2_schedule_route_block_factor not in ( - 1, - expected_fc2_schedule_factor, + 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 be one or cover one packed route " - f"block: factor={self.fc2_schedule_route_block_factor}, " - f"expected={expected_fc2_schedule_factor}" + "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 @@ -4815,6 +5443,10 @@ def __init__( 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: @@ -5118,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 @@ -5253,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 251cbf33..8d8c6a51 100644 --- a/sparkinfer/moe/_shared/kernels/w4a16/mixed_trellis.py +++ b/sparkinfer/moe/_shared/kernels/w4a16/mixed_trellis.py @@ -57,6 +57,7 @@ class MixedTrellisCompileResult: 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 @@ -111,7 +112,7 @@ class MixedTrellisTier(Protocol): class W4A16MixedTrellisKernel: """One cooperative grid over two native Trellis bitrates.""" - ABI_VERSION = 3 + ABI_VERSION = 4 def __init__( self, @@ -155,6 +156,7 @@ def __init__( g.cta_threads, g.moe_block_size, g.schedule_route_block_factor, + g.paired_m8_routes, ) for g in gemms ) @@ -164,10 +166,17 @@ def __init__( ) fc2_factor = int(driver.fc2.schedule_route_block_factor) expected_factor = int(driver.moe_block_size // driver.fc2.moe_block_size) - if fc2_factor not in (1, expected_factor): + if fc2_factor < 1 or expected_factor % fc2_factor != 0: raise ValueError( - "mixed Trellis FC2 schedule factor must be one or cover one " - f"packed route block: factor={fc2_factor}, expected={expected_factor}" + "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") @@ -252,15 +261,12 @@ def _emit_tier_tile( gemm = self.tier0.fc1 else: gemm = self.tier0.fc2 - for subtile in cutlass.range_constexpr( - gemm.schedule_route_block_factor - ): + if cutlass.const_expr(gemm.paired_m8_routes): tile_route_block_idx = ( route_block_idx * Int32(gemm.schedule_route_block_factor) - + Int32(subtile) ) - gemm._run_tile( + gemm._run_tile_m8_pair( a_flat, a_alt_flat, t0_b_flat, @@ -280,25 +286,54 @@ def _emit_tier_tile( reduce_tile_count, reduce_slice_count, reduce_slice_idx, - lock_slot - * Int32(gemm.schedule_route_block_factor) - + Int32(subtile), + 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 - for subtile in cutlass.range_constexpr( - gemm.schedule_route_block_factor - ): + if cutlass.const_expr(gemm.paired_m8_routes): tile_route_block_idx = ( route_block_idx * Int32(gemm.schedule_route_block_factor) - + Int32(subtile) ) - gemm._run_tile( + gemm._run_tile_m8_pair( a_flat, a_alt_flat, t1_b_flat, @@ -318,11 +353,43 @@ def _emit_tier_tile( reduce_tile_count, reduce_slice_count, reduce_slice_idx, - lock_slot - * Int32(gemm.schedule_route_block_factor) - + Int32(subtile), + 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__( @@ -605,7 +672,7 @@ def make_kernel(num_experts: int, bits: int) -> W4A16FusedMoeKernel: 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=(8 if int(moe_block_size) == 64 else 1), + 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", @@ -757,6 +824,7 @@ def tier_args(experts: int, bits: int): 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 7910dcdb..32810e62 100644 --- a/tests/moe/test_w4a16_mixed_trellis.py +++ b/tests/moe/test_w4a16_mixed_trellis.py @@ -363,7 +363,8 @@ def one_grid( 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 == 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