diff --git a/sparkinfer/moe/_shared/kernels/w4a16/kernel.py b/sparkinfer/moe/_shared/kernels/w4a16/kernel.py index a2bfc11e..5c1f188a 100644 --- a/sparkinfer/moe/_shared/kernels/w4a16/kernel.py +++ b/sparkinfer/moe/_shared/kernels/w4a16/kernel.py @@ -215,6 +215,8 @@ def _fake_m_for_specialization(size_m: int) -> int: (64, 128, 128), (128, 64, 128), ) +_TC_DECODE_ULTRA_FC2_TILE_K = 32 +_TC_DECODE_ULTRA_FC2_TILE_N = 512 def _covering_count(total: int, quantum: int) -> int: @@ -403,6 +405,43 @@ def _candidate_tile_fits( return smem_bytes <= int(max_shared_mem) +def _tc_decode_ultra_fc2_tile_fits( + *, + problem_n: int, + problem_k: int, + cta_m_blocks: int, + tile_n: int, + tile_k: int, + max_shared_mem: int, + scale_format: str, + weight_layout: str, +) -> bool: + """Validate the one TC-decode tile below the generic tile-k floor.""" + if ( + int(tile_n) != _TC_DECODE_ULTRA_FC2_TILE_N + or int(tile_k) != _TC_DECODE_ULTRA_FC2_TILE_K + ): + return False + scale_group_size = _scale_group_size(scale_format) + if ( + int(problem_n) % int(tile_n) != 0 + or int(problem_k) % int(tile_k) != 0 + or int(problem_k) % scale_group_size != 0 + or int(tile_k) % scale_group_size != 0 + ): + return False + return ( + _shared_memory_footprint( + cta_m_blocks=cta_m_blocks, + tile_n=tile_n, + tile_k=tile_k, + scale_format=scale_format, + weight_layout=weight_layout, + ) + <= int(max_shared_mem) + ) + + def _select_tile_config( *, problem_m: int, @@ -6202,20 +6241,18 @@ def compile_w4a16_fused_moe( and (fc2_mn_after // 2) <= fc1_mn_after and fc1_mn_after <= int(sms) ): - ultra_fc2_tile_k = 32 - ultra_smem = _shared_memory_footprint( + if _tc_decode_ultra_fc2_tile_fits( + problem_n=hidden_size, + problem_k=intermediate_size, cta_m_blocks=_covering_count(moe_block_size, 16), - tile_n=512, - tile_k=ultra_fc2_tile_k, + tile_n=_TC_DECODE_ULTRA_FC2_TILE_N, + tile_k=_TC_DECODE_ULTRA_FC2_TILE_K, + max_shared_mem=int(max_shared_mem) - 512, scale_format=scale_format, weight_layout=weight_layout, - ) - if ( - int(intermediate_size) % ultra_fc2_tile_k == 0 - and ultra_smem <= int(max_shared_mem) - 512 ): - fc2_tile_n = 512 - fc2_tile_k = ultra_fc2_tile_k + fc2_tile_n = _TC_DECODE_ULTRA_FC2_TILE_N + fc2_tile_k = _TC_DECODE_ULTRA_FC2_TILE_K fc2_cta_threads = 256 if force_tile_config is not None: # Explicit (fc1_tile_k, fc1_tile_n, fc2_tile_k, fc2_tile_n) pin. The @@ -6238,10 +6275,21 @@ def compile_w4a16_fused_moe( ("fc1", fc1_cols, hidden_size, fc1_tile_n, fc1_tile_k), ("fc2", hidden_size, intermediate_size, fc2_tile_n, fc2_tile_k), ): - if not _candidate_tile_fits( + cta_m_blocks = _covering_count(moe_block_size, 16) + ultra_fc2_fits = name == "fc2" and _tc_decode_ultra_fc2_tile_fits( problem_n=forced_pn, problem_k=forced_pk, - cta_m_blocks=_covering_count(moe_block_size, 16), + cta_m_blocks=cta_m_blocks, + tile_n=forced_tn, + tile_k=forced_tk, + max_shared_mem=int(max_shared_mem) - 512, + scale_format=scale_format, + weight_layout=weight_layout, + ) + generic_tile_fits = _candidate_tile_fits( + problem_n=forced_pn, + problem_k=forced_pk, + cta_m_blocks=cta_m_blocks, tile_n=forced_tn, tile_k=forced_tk, cta_threads=fc1_cta_threads, @@ -6249,7 +6297,8 @@ def compile_w4a16_fused_moe( scale_format=scale_format, weight_layout=weight_layout, allow_logical_tail=allow_native_logical_tail, - ): + ) + if not ultra_fc2_fits and not generic_tile_fits: raise ValueError( f"force_tile_config {name} tile " f"(tile_k={forced_tk}, tile_n={forced_tn}) does not fit " diff --git a/tests/moe/test_w4a16_e2e.py b/tests/moe/test_w4a16_e2e.py index ca5868c5..5e381341 100644 --- a/tests/moe/test_w4a16_e2e.py +++ b/tests/moe/test_w4a16_e2e.py @@ -961,6 +961,59 @@ def _spy_compile(*args, **kwargs): _assert_matches_oracle(actual, expected, activation=activation) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_w4a16_tc_decode_ultra_fc2_tile_round_trips_forced_repin() -> None: + """A GB10-selected ultra-wide FC2 tile must survive custom-op re-pinning.""" + props = torch.cuda.get_device_properties(0) + compile_kwargs = { + "size_m": 1, + "hidden_size": 4096, + "intermediate_size": 1024, + "num_experts": 256, + "top_k": 6, + "activation": "silu", + "apply_router_weight_on_input": False, + "zero_fc2_output": False, + "moe_block_size": 8, + "max_m_blocks": 6, + "element_dtype": "bf16", + # The wave-balance override is specific to the 48-SM GB10 geometry. + "sms": 48, + "max_shared_mem": int(props.shared_memory_per_block_optin), + "weight_layout": "packed", + "scale_format": "e8m0_k32", + "w13_layout": "w13", + "direct_topk_routes": True, + "tc_decode_fused_sum": True, + } + + auto = compile_w4a16_fused_moe(**compile_kwargs) + tile_config = ( + int(auto.fc1_tile_k), + int(auto.fc1_tile_n), + int(auto.fc2_tile_k), + int(auto.fc2_tile_n), + ) + assert tile_config == (64, 256, 32, 512) + + repinned = compile_w4a16_fused_moe( + **compile_kwargs, + force_tile_config=tile_config, + ) + assert ( + int(repinned.fc1_tile_k), + int(repinned.fc1_tile_n), + int(repinned.fc2_tile_k), + int(repinned.fc2_tile_n), + ) == tile_config + + with pytest.raises(ValueError, match="force_tile_config fc2 tile"): + compile_w4a16_fused_moe( + **(compile_kwargs | {"hidden_size": 3840}), + force_tile_config=tile_config, + ) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @pytest.mark.parametrize("m", [1, 2, 3, 6, 8]) def test_w4a16_tc_decode_preplanned_launch_matches_oracle(m: int) -> None: