Skip to content

perf: reduce split W4A16 routing and output overhead #5186

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
zianglih wants to merge 3 commits into
base: main
Choose a base branch
from
+73 −36
Open

perf: reduce split W4A16 routing and output overhead #5186

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
968380d
perf: specialize split W4A16 fused finalize on top-k
zianglih Sep 13, 2026
b90b347
perf: specialize split W4A16 epilogue bounds
zianglih Sep 13, 2026
88ff7db
perf: use complete W4A16 subtiles for scratch stores
zianglih Sep 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion flashinfer/fused_moe/cute_dsl/blackwell/moe_w4a16.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,8 @@ def _get_compiled_kernel(
):
mma_tiler_m, route_tile, mma_tiler_k = mma_tiler_mnk
use_2cta_instrs = mma_tiler_m == 256
cta_tile_m = mma_tiler_m // (2 if use_2cta_instrs else 1)
m_cluster_aligned = m > 0 and m % (cta_tile_m * cluster_shape_mn[0]) == 0
transform_fragment_size = (
128 if activation_type is not None or k == mma_tiler_k else 32
)
Expand All @@ -137,6 +139,7 @@ def _get_compiled_kernel(
situ_beta,
situ_linear_beta,
use_fused_finalize,
top_k,
enable_pdl,
use_clc_scheduler,
mma_tiler_m,
Expand All @@ -145,6 +148,7 @@ def _get_compiled_kernel(
cluster_shape_mn,
raster_along_m,
transform_fragment_size,
m_cluster_aligned,
)
compiled = _kernel_cache.get(cache_key)
if compiled is None:
Expand All @@ -167,6 +171,7 @@ def _get_compiled_kernel(
use_clc_scheduler=use_clc_scheduler,
raster_along_m=raster_along_m,
transform_fragment_size=transform_fragment_size,
m_cluster_aligned=m_cluster_aligned,
)
compiled = cute.compile(
kernel.wrapper,
Expand Down Expand Up @@ -350,7 +355,6 @@ def _run_grouped_gemm(
n,
k,
num_tokens,
top_k,
stream=stream,
)

Expand Down
46 changes: 32 additions & 14 deletions flashinfer/fused_moe/cute_dsl/blackwell/moe_w4a16_kernel.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,7 @@ def __init__(
use_clc_scheduler: bool,
raster_along_m: bool,
transform_fragment_size: int,
m_cluster_aligned: bool,
):
"""Initialize the W4A16 grouped GEMM configuration."""
self.group_count = group_count
Expand All @@ -116,6 +117,7 @@ def __init__(
self.use_clc_scheduler = use_clc_scheduler
self.raster_along_m = raster_along_m
self.transform_fragment_size = transform_fragment_size
self.m_cluster_aligned = m_cluster_aligned
if activation_type is None:
if situ_beta is not None or situ_linear_beta is not None:
raise ValueError("SiTU parameters require an activation")
Expand Down Expand Up @@ -570,7 +572,7 @@ def wrapper(
n: cutlass.Int64,
k: cutlass.Int64,
num_tokens: cutlass.Int64,
top_k: cutlass.Int64,
top_k: cutlass.Constexpr,
max_active_clusters: cutlass.Constexpr,
stream: cuda.CUstream,
):
Expand Down Expand Up @@ -2211,19 +2213,28 @@ def kernel(
hidden_base = (
work_tile.cta_coord_m * self.cta_tile_shape_mnk[0]
)
valid_elements = (
cutlass.Int64(final_output.shape[0]) - hidden_base
)
if valid_elements > 0:
valid_elements = cutlass.Int64(self.cta_tile_shape_mnk[0])
if cutlass.const_expr(not self.m_cluster_aligned):
valid_elements = (
cutlass.Int64(final_output.shape[0]) - hidden_base
)
if (
cutlass.const_expr(self.m_cluster_aligned)
or valid_elements > 0
):
scatter_out = cute.domain_offset(
(hidden_base, reduce_token_idx, 0), final_output
)
copy_elements = cutlass.Int32(
cutlass.min(
cutlass.Int64(self.cta_tile_shape_mnk[0]),
valid_elements,
)
self.cta_tile_shape_mnk[0]
)
if cutlass.const_expr(not self.m_cluster_aligned):
copy_elements = cutlass.Int32(
cutlass.min(
cutlass.Int64(self.cta_tile_shape_mnk[0]),
valid_elements,
)
)
blk_reduce_bf16(
scatter_out,
sFinalize[(reduce_route, None)],
Expand All @@ -2233,7 +2244,9 @@ def kernel(
cute.arch.cp_async_bulk_commit_group()
cute.arch.cp_async_bulk_wait_group(0, read=True)
self.epilog_sync_barrier.arrive_and_wait()
elif tma_distance_to_boundary >= self.cta_tile_shape_mnk[1]:
elif (
tma_distance_to_boundary >= (subtile_idx + 1) * self.epi_tile_n
):
# Convert to C type
acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load()
if cutlass.const_expr(not self.fuse_activation):
Expand Down Expand Up @@ -2279,10 +2292,15 @@ def kernel(
m_thr_slice = m_thr_offset[(None, None, None, subtile_idx)]
for i in cutlass.range(cute.size(tCpC), unroll_full=True):
tCpC[i] = (
m_thr_slice[(i)][0]
+ work_tile.cta_coord_m * self.cta_tile_shape_mnk_c[0]
< cute.size(tensor_c.shape[0])
) and (m_thr_slice[(i)][1] < work_tile.distance_to_boundary)
m_thr_slice[(i)][1] < work_tile.distance_to_boundary
)
if cutlass.const_expr(not self.m_cluster_aligned):
tCpC[i] = (
m_thr_slice[(i)][0]
+ work_tile.cta_coord_m
* self.cta_tile_shape_mnk_c[0]
< cute.size(tensor_c.shape[0])
) and tCpC[i]
# Store C to global memory
cute.copy(
simt_atom,
Expand Down
57 changes: 36 additions & 21 deletions tests/moe/test_cute_dsl_fused_moe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1411,19 +1411,28 @@ def test_weight_scale_mapping(
),
],
)
@pytest.mark.parametrize(
"half_tile_tail", [False, True], ids=["tail1", "tail-half-plus1"]
)
@pytest.mark.parametrize("top_k", [2, 3])
@pytest.mark.parametrize("use_fused_finalize", [False, True])
def test_route_tile_boundary_accuracy(
self,
route_tile: int,
gemm1_tactic: tuple,
gemm2_tactic: tuple,
half_tile_tail: bool,
top_k: int,
use_fused_finalize: bool,
):
from flashinfer.fused_moe.cute_dsl.blackwell.moe_w4a16 import (
launch_w4a16_moe,
)
from flashinfer.fused_moe.cute_dsl.tuner import W4A16_MOE_TACTICS

num_tokens, hidden_size, intermediate_size = route_tile + 1, 256, 512
num_experts, top_k = 8, 2
num_tokens = route_tile + (route_tile // 2 + 1 if half_tile_tail else 1)
hidden_size, intermediate_size = 256, 512
num_experts = 8
tensors = create_moe_tensors(
num_tokens=num_tokens,
hidden_size=hidden_size,
Expand All @@ -1432,7 +1441,7 @@ def test_route_tile_boundary_accuracy(
num_local_experts=num_experts,
top_k=top_k,
)
# Give two experts one full route tile and one boundary tile each.
# Give each selected expert one full route tile and one boundary tile.
tensors["token_selected_experts"][:] = torch.arange(
top_k, device=tensors["token_selected_experts"].device
)
Expand All @@ -1456,7 +1465,7 @@ def test_route_tile_boundary_accuracy(
moe_output=torch.empty(
(num_tokens, hidden_size), dtype=torch.bfloat16, device="cuda"
),
use_fused_finalize=False,
use_fused_finalize=use_fused_finalize,
enable_pdl=False,
activation_type=ActivationType.Swiglu,
tactic=tactic,
Expand Down Expand Up @@ -1600,15 +1609,20 @@ def test_deterministic_finalize_numerical_accuracy(
)

@pytest.mark.parametrize(
"quant_mode, use_per_token_activation",
_MOE_QUANT_MODE_CASES,
"quant_mode,use_per_token_activation,use_fused_finalize,hidden_sizes",
[
pytest.param("w4a4", False, True, (256, 384), id="w4a4-per-tensor"),
pytest.param("w4a4", True, True, (256, 384), id="w4a4-per-token"),
pytest.param("w4a16", False, False, (256, 384, 256), id="w4a16-ordinary"),
pytest.param("w4a16", False, True, (256, 384, 256), id="w4a16-fused"),
],
)
@pytest.mark.parametrize("hidden_size", [256, 384])
def test_finalize_handles_cluster_padding_and_partial_tiles(
self,
quant_mode: str,
use_per_token_activation: bool,
hidden_size: int,
use_fused_finalize: bool,
hidden_sizes: tuple[int, ...],
monkeypatch: pytest.MonkeyPatch,
):
from flashinfer.autotuner import AutoTuner
Expand All @@ -1624,8 +1638,8 @@ def test_finalize_handles_cluster_padding_and_partial_tiles(
((256, 256), (2, 2), False),
)
elif quant_mode == "w4a16":
# W4A16 clusters 128-wide M CTAs in pairs, so hidden=384 leaves a
# padding peer.
# The same cache sees a full cluster, a padding peer, then a full
# cluster again. Keep route tails in both finalize modes.
tail_config = (
((256, 128, 256), (2, 1), True),
((256, 128, 256), (2, 1), True),
Expand All @@ -1639,17 +1653,18 @@ def choose_tail_config(
return runners[0], tail_config

monkeypatch.setattr(AutoTuner, "choose_one", choose_tail_config)
self._run_numerical_accuracy(
activation_type=ActivationType.Relu2,
num_tokens=128,
top_k=2,
hidden_size=hidden_size,
intermediate_size=512,
num_experts=8,
quant_mode=quant_mode,
use_per_token_activation=use_per_token_activation,
use_fused_finalize=True,
)
for hidden_size in hidden_sizes:
self._run_numerical_accuracy(
activation_type=ActivationType.Relu2,
num_tokens=128,
top_k=2,
hidden_size=hidden_size,
intermediate_size=512,
num_experts=8,
quant_mode=quant_mode,
use_per_token_activation=use_per_token_activation,
use_fused_finalize=use_fused_finalize,
)

def _run_numerical_accuracy(
self,
Expand Down
Loading