From a754d77777c1c8f8513b64aae0d7bcb38642749f Mon Sep 17 00:00:00 2001 From: "tong.qiu2@amd.com" Date: Tue, 27 Oct 2026 01:45:53 +0000 Subject: [PATCH 1/3] Revert "fix init moe (#6)" This reverts commit e9a444789ffb476b14360154c55a3a501ffd7cc7. --- aiter/ops/flydsl/fused_moe_gfx942.py | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/aiter/ops/flydsl/fused_moe_gfx942.py b/aiter/ops/flydsl/fused_moe_gfx942.py index 07b99463f42..933aa074ed0 100644 --- a/aiter/ops/flydsl/fused_moe_gfx942.py +++ b/aiter/ops/flydsl/fused_moe_gfx942.py @@ -931,16 +931,11 @@ def _run_prefill( if output_padding_bytes is not None else 0 ) - # Opus may coalesce duplicate expert IDs in zero-weight HIP padding rows. - # Their omitted route slots are not written by invert_sorted_ids. Map those - # slots to a dedicated zero row instead of dereferencing uninitialized IDs. - zero_route_row = sorted_expert_ids.shape[0] * config.BLOCK_M gemm2_out = torch.empty( - [zero_route_row + 1, output_row_size], + [sorted_expert_ids.shape[0] * config.BLOCK_M, output_row_size], dtype=hidden_states.dtype, device=hidden_states.device, ) - gemm2_out[zero_route_row].zero_() down_tile_n = { "default": 128, "1x4_64x256": 256, @@ -1000,9 +995,8 @@ def _run_prefill( else: _launch(down_kernel, *down_args) - loc_ids = torch.full( + loc_ids = torch.empty( [problem.batch, problem.topk], - zero_route_row, dtype=torch.int32, device=hidden_states.device, ) From 905d4728d068de03dc010458272fc3df45d4a28e Mon Sep 17 00:00:00 2001 From: Luo Cheng Date: Sun, 25 Oct 2026 09:41:41 +0000 Subject: [PATCH 2/3] use Opus aux sort for the sorted_sum route lookup to avoid call cost to `invert_sorted_ids` (cherry picked from commit 4641dd780d15e288a64d66b5bd8f5fb0c07dda32) --- aiter/ops/flydsl/fused_moe_gfx942.py | 27 +- .../moe_gemm_2stage_gfx942/__init__.py | 2 - .../moe_gemm_2stage_gfx942/moe_reduce.py | 79 +-- op_tests/test_flydsl_moe_gfx942.py | 632 ++++++++++++++++++ 4 files changed, 649 insertions(+), 91 deletions(-) create mode 100644 op_tests/test_flydsl_moe_gfx942.py diff --git a/aiter/ops/flydsl/fused_moe_gfx942.py b/aiter/ops/flydsl/fused_moe_gfx942.py index 933aa074ed0..ede5b61d51e 100644 --- a/aiter/ops/flydsl/fused_moe_gfx942.py +++ b/aiter/ops/flydsl/fused_moe_gfx942.py @@ -11,13 +11,12 @@ import aiter from aiter import ActivationType, QuantType -from aiter.fused_moe import moe_sorting +from aiter.fused_moe import AUX_SORT_OPUS, moe_sorting from aiter.fused_moe_registry import FusedMoeRequest from aiter.jit.utils.chip_info import get_gfx from aiter.ops.flydsl.kernels.moe_gemm_2stage_gfx942 import ( flydsl_absmax, flydsl_quant_per_tensor, - invert_sorted_ids, sorted_sum, ) from aiter.ops.flydsl.kernels.moe_gemm_2stage_gfx942.common import ( @@ -843,7 +842,16 @@ def _run_prefill( situ_linear_beta: float, ): gateup_block_m = config.GATEUP_BLOCK_M or config.BLOCK_M - sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, cur_out = moe_sorting( + # Opus also returns each route's sorted row; merged repeated expert IDs get -1. + ( + sorted_ids, + sorted_weights, + sorted_expert_ids, + num_valid_ids, + cur_out, + _, + loc_ids, + ) = moe_sorting( topk_ids, topk_weight, problem.experts, @@ -853,6 +861,7 @@ def _run_prefill( None, None, 0, + output_aux=AUX_SORT_OPUS, ) weight_dtype_str = "bf16" if w1.dtype == torch.bfloat16 else "fp8" act_quant_type_str = problem.quant_type @@ -995,18 +1004,6 @@ def _run_prefill( else: _launch(down_kernel, *down_args) - loc_ids = torch.empty( - [problem.batch, problem.topk], - dtype=torch.int32, - device=hidden_states.device, - ) - invert_sorted_ids(problem.topk)( - sorted_ids, - loc_ids, - num_valid_ids, - sorted_ids.shape[0], - problem.batch, - ) sorted_sum(problem.topk, problem.model_dim, output_padding_bytes)( loc_ids, gemm2_out, cur_out, problem.batch ) diff --git a/aiter/ops/flydsl/kernels/moe_gemm_2stage_gfx942/__init__.py b/aiter/ops/flydsl/kernels/moe_gemm_2stage_gfx942/__init__.py index 35d0aa776e9..9045005841d 100644 --- a/aiter/ops/flydsl/kernels/moe_gemm_2stage_gfx942/__init__.py +++ b/aiter/ops/flydsl/kernels/moe_gemm_2stage_gfx942/__init__.py @@ -12,7 +12,6 @@ from .gemm2 import compile_moe_gemm2 from .moe_reduce import ( compile_moe_reduction, - invert_sorted_ids, sorted_sum, ) from .quant import flydsl_absmax, flydsl_quant_per_tensor @@ -24,6 +23,5 @@ "compile_moe_reduction", "flydsl_absmax", "flydsl_quant_per_tensor", - "invert_sorted_ids", "sorted_sum", ] diff --git a/aiter/ops/flydsl/kernels/moe_gemm_2stage_gfx942/moe_reduce.py b/aiter/ops/flydsl/kernels/moe_gemm_2stage_gfx942/moe_reduce.py index 3cd1985333a..4a5683bba5f 100644 --- a/aiter/ops/flydsl/kernels/moe_gemm_2stage_gfx942/moe_reduce.py +++ b/aiter/ops/flydsl/kernels/moe_gemm_2stage_gfx942/moe_reduce.py @@ -35,6 +35,8 @@ def sorted_sum_kernel(loc_ids: fx.Pointer, A: fx.Pointer, B: fx.Pointer): # preload all TOPK locations loc_ids += batch * TOPK token_locs = [loc_ids[topk] for topk in fx.range_constexpr(TOPK)] + # Sorting gives -1 to routes merged as repeated expert IDs of a token. + valid = [loc >= 0 for loc in token_locs] copy_bits = 128 copy_atom = fx.make_copy_atom(fx.UniversalCopy(copy_bits), A.dtype) @@ -49,7 +51,7 @@ def sorted_sum_kernel(loc_ids: fx.Pointer, A: fx.Pointer, B: fx.Pointer): ) token_ptrs = [ - (A + fx.Int64(token_locs[topk]) * source_row_stride) + (A + fx.Int64(valid[topk].select(token_locs[topk], 0)) * source_row_stride) for topk in fx.range_constexpr(TOPK) ] @@ -63,9 +65,9 @@ def load_atoms(off): return [load_atom(topk, off) for topk in fx.range_constexpr(TOPK)] def reduce_store(dst, frag): - vec_sum = frag[0].load().to(fx.Float32) + vec_sum = valid[0].select(frag[0].load().to(fx.Float32), 0.0) for m in fx.range_constexpr(1, TOPK): - vec_sum += frag[m].load().to(fx.Float32) + vec_sum += valid[m].select(frag[m].load().to(fx.Float32), 0.0) vec_sum = vec_sum.to(dst.dtype) out_frag = fx.make_fragment_like(dst) out_frag.store(vec_sum) @@ -121,74 +123,3 @@ def compile_moe_reduction(*, topk, model_dim, row_padding_bytes=None): """Build and cache the route reduction callable for one static configuration.""" return sorted_sum(topk, model_dim, row_padding_bytes) - -@functools.cache -def _invert_sorted_ids_cached(device_cache_key, TOPK): - del device_cache_key - num_threads = 64 - - @flyc.kernel(known_block_size=[num_threads, 1, 1]) - def invert_sorted_ids_kernel( - sorted_ids: fx.Pointer, - invert: fx.Pointer, - p_num_valid: fx.Pointer, - num_ids: fx.Uint32, - batch_size: fx.Uint32, - ): - batch = fx.block_idx.x - tid = fx.thread_idx.x - slot = batch * num_threads + tid - # Scan only the down-written region [0, num_valid): the tail of sorted_ids - # (>= num_valid) is uninitialized, and its garbage would racily map real - # tokens onto unwritten gemm2_out rows. Read the bound on-device (no host sync). - num_valid = p_num_valid[0].to(fx.Uint32) - if slot < num_valid: - sid = sorted_ids[slot].to(fx.Uint32) - tok_id = sid & 0xFFFFFF - top_id = sid >> 24 - idx = tok_id * TOPK + top_id - if top_id < TOPK and tok_id < batch_size: - invert[idx] = fx.Uint32(slot) - - @flyc.jit - def launch( - sorted_ids: fx.Pointer, - invert: fx.Pointer, - p_num_valid: fx.Pointer, - num_ids: fx.Uint32, - batch_size: fx.Uint32, - stream: fx.Stream, - ): - grid_size = fxh.div_up(num_ids, num_threads) - invert_sorted_ids_kernel( - sorted_ids, invert, p_num_valid, num_ids, batch_size - ).launch(grid=(grid_size, 1, 1), block=(num_threads, 1, 1), stream=stream) - - def callable( - sorted_ids: torch.Tensor, - invert: torch.Tensor, - num_valid: torch.Tensor, - num_ids: int, - batch_size: int, - ): - stream = torch.cuda.current_stream() - _run_compiled( - launch, - _ptr(sorted_ids), - _ptr(invert), - _ptr(num_valid), - fx.Uint32(num_ids), - fx.Uint32(batch_size), - stream, - ) - - callable._flydsl_executable = launch - return callable - - -def invert_sorted_ids(TOPK): - return _invert_sorted_ids_cached(get_device_cache_key(), TOPK) - - -invert_sorted_ids.cache_clear = _invert_sorted_ids_cached.cache_clear -invert_sorted_ids.cache_info = _invert_sorted_ids_cached.cache_info diff --git a/op_tests/test_flydsl_moe_gfx942.py b/op_tests/test_flydsl_moe_gfx942.py new file mode 100644 index 00000000000..28b67ab3756 --- /dev/null +++ b/op_tests/test_flydsl_moe_gfx942.py @@ -0,0 +1,632 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Regression tests for the gfx942 FlyDSL whole-MoE backend. + +Its tuned rows use cu_num=80 (MI308X), so test_moe_2stage.py skips them on +runners with another CU count. This test runs every shipped config through the +registered backend on any gfx942 device, checks the compact Down task table +for both MI308X and MI300X CU counts, and checks the padded sorted reduction. +""" + +import argparse +import functools +import itertools +from unittest.mock import patch + +import pandas as pd +import torch + +import aiter +from aiter import ActivationType, QuantType, dtypes +from aiter.fused_moe import ( + AUX_SORT_OPUS, + fused_topk, + moe_sorting, + torch_moe_stage1, + torch_moe_stage2, +) +from aiter.fused_moe_registry import ( + FusedMoeRequest, + make_fused_moe_impl_kernel_name, + resolve_fused_moe_impl, +) +from aiter.jit.core import AITER_CONFIGS +from aiter.jit.utils.chip_info import get_cu_num, get_gfx +from aiter.ops.flydsl.fused_moe_gfx942 import Config, get_tune_space +from aiter.ops.flydsl.kernels.moe_gemm_2stage.gfx9_fp8 import ( + gemm2_8x1_compact as compact, +) +from aiter.ops.flydsl.kernels.moe_gemm_2stage.gfx9_fp8 import sorted_sum +from aiter.ops.flydsl.kernels.moe_gemm_2stage.gfx9_fp8.common import _ptr +from aiter.ops.flydsl.kernels.tensor_shim import _run_compiled +from aiter.ops.shuffle import shuffle_weight +from aiter.test_common import benchmark, checkAllclose, run_perftest + +torch.set_default_device("cuda") + +IMPL = "flydsl_gfx942" +QUANTS = { + "no": QuantType.No, + "per_tensor": QuantType.per_Tensor, + "ptpc": QuantType.per_Token, +} +ACTIVATIONS = {"silu": ActivationType.Silu, "swiglu": ActivationType.Swiglu} +# MI300X then MI308X, the order in which a cached launcher meets a smaller device. +CU_COUNTS = (304, 80) +GUARD, CANARY = 64, -123 +# FP8 rounding moves an element by at most 0.28 token RMS from the torch reference. +ELEMENT_ERR_LIMIT = 1.0 +# FP8 rounding reaches 0.076 per token; one misrouted route gives about 0.67. +TOKEN_ERR_LIMIT = 0.2 +LOGITS_DIFF_LIMIT = 1e-3 +# Prefill paths of one case share Gate/Up and quantization and agree to BF16 rounding. +PATH_DIFF_LIMIT = 0.1 + + +def guarded_task_buffers(sorted_expert_ids, experts, cu_count): + """allocate_task_buffers with canaries before and after every buffer.""" + full, tail = compact.task_capacities(sorted_expert_ids.numel(), experts, cu_count) + storages, buffers = [], [] + for shape in ((full, 2), (tail, 2), (2,)): + storage = torch.full( + (torch.Size(shape).numel() + 2 * GUARD,), + CANARY, + dtype=torch.int32, + device=sorted_expert_ids.device, + ) + storages.append(storage) + buffers.append(storage[GUARD:-GUARD].view(shape)) + return storages, buffers + + +def check_guards(storages): + for storage in storages: + assert bool((storage[:GUARD] == CANARY).all()), "write before a buffer" + assert bool((storage[-GUARD:] == CANARY).all()), "write after a buffer" + + +def check_tasks(blocks, experts, cu_count, full_tasks, tail_tasks, counts): + """Check the split policy and that each block is covered once by its expert.""" + sizes = [0] * experts + for expert in blocks: + sizes[expert] += 1 + full = sum(size // 4 for size in sizes) + tail = sum(size % 4 for size in sizes) + remainder = full % cu_count + split = 0 + if remainder and remainder * compact._SPLIT_DEN < cu_count * compact._SPLIT_NUM: + split = remainder + # The kernel caps the split at the free tail capacity. + split = min(split, (tail_tasks.shape[0] - tail) // 4) + expected = [(full - split) * 256, (tail + 4 * split) * 64] + assert counts.tolist() == expected, f"counts {counts.tolist()} != {expected}" + coverage = [0] * len(blocks) + for tasks, rows, count in ( + (full_tasks, 256, full - split), + (tail_tasks, 64, tail + 4 * split), + ): + assert count <= tasks.shape[0], "task count exceeds its buffer" + for start, expert in tasks[:count].tolist(): + assert start % 64 == 0 and 0 <= start and start + rows <= 64 * len(blocks) + for block in range(start // 64, (start + rows) // 64): + assert blocks[block] == expert, "task routed to another expert" + coverage[block] += 1 + assert coverage == [1] * len(blocks), "blocks not covered exactly once" + return full - split, tail + 4 * split + + +def task_scenarios(cu_count): + """Map a scenario to (experts, [(expert, M64 blocks)]) in metadata order.""" + # Largest remainder of the final full-task CU round that is still split. + split = (cu_count * compact._SPLIT_NUM - 1) // compact._SPLIT_DEN + return { + "empty": (4, []), + "tails_only": (8, [(1, 1), (3, 2), (6, 3)]), + "single_expert": (1, [(0, 8)]), + "full_round": (2, [(1, 4 * cu_count)]), + "split_limit": (2, [(0, 4 * (cu_count + split)), (1, 3)]), + "no_split": (2, [(0, 4 * (cu_count + split + 1)), (1, 3)]), + # Every expert has one full task and three tails, filling the tail capacity. + "tail_capacity": (split, [(expert, 7) for expert in range(split)]), + # 60 full tasks split into tails with 304 CUs but not with 80. + "cu_dependent": (1, [(0, 240)]), + "experts_2048": (2048, [(expert, 1 + expert % 7) for expert in range(2048)]), + } + + +# One launcher per expert count serves every CU count, like the cached builder. +build_task_table = functools.cache(compact.build_task_table) + + +@benchmark() +def test_task_table(scenario, cu_count, buffer_cu_count): + experts, runs = task_scenarios(cu_count)[scenario] + blocks = [expert for expert, size in runs for _ in range(size)] + # Padding names a real expert, so reading past the valid blocks breaks a run. + metadata = torch.tensor(blocks + [0] * 4, dtype=torch.int32) + valid_rows = torch.tensor([64 * len(blocks)], dtype=torch.int32) + storages, (full_tasks, tail_tasks, counts) = guarded_task_buffers( + metadata, experts, buffer_cu_count + ) + launcher = build_task_table(experts) + pointers = [_ptr(t) for t in (metadata, valid_rows, full_tasks, tail_tasks, counts)] + + def build(): + _run_compiled( + launcher, + *pointers, + cu_count, + tail_tasks.shape[0], + torch.cuda.current_stream(), + ) + return counts + + _, us = run_perftest(build) + full, tail = check_tasks(blocks, experts, cu_count, full_tasks, tail_tasks, counts) + check_guards(storages) + if scenario == "tail_capacity" and buffer_cu_count == cu_count: + assert tail == tail_tasks.shape[0], "tail capacity is not tight" + nbytes = metadata.nbytes + valid_rows.nbytes + 8 * (full + tail) + counts.nbytes + return { + "blocks": len(blocks), + "full tasks": full, + "full capacity": full_tasks.shape[0], + "tail tasks": tail, + "tail capacity": tail_tasks.shape[0], + "us": us, + "TB/s": nbytes / us / 1e6, + } + + +@benchmark() +def test_sorted_sum(token, model_dim, topk, padding): + dtype = dtypes.bf16 + experts = max(4, topk) + ids = (torch.arange(token * topk).reshape(token, topk) % experts).to(torch.int32) + # Sorting keeps one route of a repeated expert ID and gives the others -1. + ids[-1] = 0 + weights = torch.ones((token, topk), dtype=torch.float32) + storage = torch.full((token * model_dim + 2 * GUARD,), CANARY, dtype=dtype) + _, _, sorted_experts, _, output, _, loc_ids = moe_sorting( + ids, + weights, + experts, + model_dim, + dtype, + 64, + output=storage[GUARD:-GUARD].view(token, model_dim), + output_aux=AUX_SORT_OPUS, + ) + valid = loc_ids.view(token, topk) >= 0 + assert int((~valid).sum()) == topk - 1, "repeated expert IDs kept extra routes" + # Specialized Down rows carry padding that holds canaries the sum must skip. + down_output = torch.full( + (sorted_experts.numel() * 64, model_dim + padding // 2), CANARY, dtype=dtype + ) + routes = torch.randn((token, topk, model_dim)).to(dtype) + down_output[loc_ids.view(token, topk)[valid].long(), :model_dim] = routes[valid] + reference = torch.zeros((token, model_dim)) + for route in range(topk): + reference = reference + routes[:, route].float() * valid[:, route, None] + reduce = sorted_sum(topk, model_dim, padding) + + def run(): + reduce(loc_ids, down_output, output, token) + return output + + actual, us = run_perftest(run) + err = checkAllclose( + reference.to(dtype).float(), + actual.float(), + rtol=0, + atol=0, + msg=f"sorted_sum padding={padding}", + ) + assert err == 0, "padded sorted reduction is not exact" + check_guards([storage]) + flops = token * model_dim * (topk - 1) + nbytes = token * model_dim * (topk + 1) * output.element_size() + loc_ids.nbytes + return { + "us": us, + "TFLOPS": flops / us / 1e6, + "TB/s": nbytes / us / 1e6, + "err": err, + } + + +def test_config_strings(): + for encoded in get_tune_space(): + assert Config.from_string(encoded).to_string() == encoded + # Rows tuned before the Down suffix existed keep loading as the default Down. + for legacy in ("16_16_16_False", "64_128_128_True", "64_256_128_True"): + assert Config.from_string(legacy) == Config.from_string(f"{legacy}:default:0") + # Tiles without the direct-decoding flag keep their meaning. + assert Config.from_string("16_16_16_False_False:default:0") == Config.from_string( + "16_16_16_False:default:0" + ) + for invalid in ( + "64_256_128_True:default", + "64_256_128_True:default:0:0", + "64_256_128_Maybe", + "64_256_True", + "64_256_128_True:default:128", + "64_256_128_True:8x2:128", + "16_16_16_False:8x1_compact:128", + "16_16_16_False_Maybe:default:0", + "16_16_16_False_True_False:default:0", + "16_16_16_True_True:default:0", + "16_16_16_False_True:1x4_64x256:128", + ): + try: + Config.from_string(invalid) + except ValueError: + continue + raise AssertionError(f"accepted invalid config {invalid!r}") + + +def route(hidden, experts, topk, routing): + token = hidden.shape[0] + if routing in ("random", "padding"): + weights, ids = fused_topk( + hidden, torch.randn((token, experts), dtype=hidden.dtype), topk, True + ) + if routing == "padding": + # HIP padding rows: zero weights and expert 0 repeated, which sorting merges. + weights[-65:] = 0 + ids[-65:, :-1] = 0 + return weights, ids + # "hot": every token picks the same highest experts and the others stay empty. + ids = torch.arange(experts - topk, experts, dtype=torch.int32) + weights = torch.rand((token, topk), dtype=torch.float32) + return ( + weights / weights.sum(dim=-1, keepdim=True), + ids.expand(token, topk).contiguous(), + ) + + +@functools.lru_cache(maxsize=1) +def make_case( + token, + model_dim, + inter_dim, + experts, + topk, + quant, + activation, + routing, + quantize_activations, +): + """Inputs and torch reference as test_moe_2stage.py builds them.""" + torch.manual_seed(0) + q_type, act_type = QUANTS[quant], ACTIVATIONS[activation] + q_dtype = None if q_type == QuantType.No else dtypes.fp8 + torch_quant = aiter.get_torch_quant(q_type) + + def quantize_activation(x): + # Decode and batch-1 kernels read BF16 activations even with FP8 weights. + if quantize_activations: + return torch_quant(x, quant_dtype=q_dtype) + return x, torch.ones(1, dtype=torch.float32) + + hidden = torch.randn((token, model_dim), dtype=dtypes.bf16) + w1 = torch.randn((experts, 2 * inter_dim, model_dim), dtype=dtypes.bf16) + w2 = torch.randn((experts, model_dim, inter_dim), dtype=dtypes.bf16) + topk_weights, topk_ids = route(hidden, experts, topk, routing) + if q_type == QuantType.per_Tensor: + w1_q, w1_scale = aiter.pertoken_quant(w1.view(experts, -1), quant_dtype=q_dtype) + w2_q, w2_scale = aiter.pertoken_quant(w2.view(experts, -1), quant_dtype=q_dtype) + w1_q, w2_q = w1_q.view(w1.shape), w2_q.view(w2.shape) + else: + w1_q, w1_scale = torch_quant(w1, quant_dtype=q_dtype) + w2_q, w2_scale = torch_quant(w2, quant_dtype=q_dtype) + a1, a1_scale = quantize_activation(hidden) + stage1 = torch_moe_stage1( + a1, + w1_q, + w2_q, + topk_weights, + topk_ids, + dtype=dtypes.bf16, + activation=act_type, + quant_type=q_type, + a1_scale=a1_scale, + w1_scale=w1_scale, + ) + a2, a2_scale = quantize_activation(stage1) + reference = torch_moe_stage2( + a2.view(token, topk, -1), + w1_q, + w2_q, + topk_weights, + topk_ids, + dtype=dtypes.bf16, + quant_type=q_type, + w2_scale=w2_scale, + a2_scale=a2_scale, + ) + request = FusedMoeRequest( + hidden_states=hidden, + w1=shuffle_weight(w1_q, layout=(16, 16)), + w2=shuffle_weight(w2_q, layout=(16, 16)), + topk_weight=topk_weights, + topk_ids=topk_ids, + activation=act_type, + quant_type=q_type, + w1_scale=w1_scale, + w2_scale=w2_scale, + ) + # Holds the first FlyDSL output of the case for the cross-path check. + return request, reference, {} + + +class GuardedTaskBuffers: + """Replaces allocate_task_buffers and keeps the last allocation for checks.""" + + def __call__(self, sorted_expert_ids, experts, cu_count): + self.metadata, self.cu_count = sorted_expert_ids, cu_count + self.storages, self.buffers = guarded_task_buffers( + sorted_expert_ids, experts, cu_count + ) + return self.buffers + + +def token_rms(x): + return x.float().pow(2).mean(dim=-1, keepdim=True).sqrt().clamp_min(1e-6) + + +def token_error(expected, actual): + """Largest per-token relative L2 error.""" + x, y = expected.double(), actual.double() + return float(((x - y).norm(dim=-1) / x.norm(dim=-1).clamp_min(1e-6)).max()) + + +def check_output(reference, output, msg): + """Assert every accuracy check; returns the checkAllclose mismatch ratio.""" + assert bool(torch.isfinite(output).all()), f"{msg}: non-finite output" + # Errors in units of each token's reference RMS, so one tolerance fits all shapes. + rms = token_rms(reference) + err = checkAllclose( + reference.float() / rms, + output.float() / rms, + rtol=0, + atol=ELEMENT_ERR_LIMIT, + msg=msg, + ) + x, y = reference.double(), output.double() + logits_diff = float(1 - 2 * (x * y).sum() / (x * x + y * y).sum()) + token_err = token_error(reference, output) + assert err == 0, f"{msg}: {err:.3%} of elements beyond {ELEMENT_ERR_LIMIT} RMS" + assert logits_diff < LOGITS_DIFF_LIMIT, f"{msg}: logits_diff={logits_diff}" + assert token_err < TOKEN_ERR_LIMIT, f"{msg}: worst token error {token_err}" + return err + + +def run_compact(impl, request, experts, cu_count): + """Run with guarded task buffers and check the task table built for cu_count.""" + buffers = GuardedTaskBuffers() + with patch.object(compact, "allocate_task_buffers", buffers): + output = impl(request) + assert buffers.cu_count == cu_count + sizes = torch.bincount(request.topk_ids.flatten().long(), minlength=experts) + blocks = buffers.metadata[: int(((sizes + 63) // 64).sum())].tolist() + full, tail = check_tasks(blocks, experts, cu_count, *buffers.buffers) + check_guards(buffers.storages) + return output, f"{full}/{tail}" + + +def rms_error(expected, actual, rms): + """Largest elementwise difference in units of the token's reference RMS.""" + return float(((actual.float() - expected.float()).abs() / rms).max()) + + +def check_moe( + token, model_dim, inter_dim, experts, topk, quant, activation, routing, config +): + """Run one config through the registered backend and assert its accuracy.""" + parsed = Config.from_string(config) + request, reference, first_output = make_case( + token, + model_dim, + inter_dim, + experts, + topk, + quant, + activation, + routing, + parsed.use_prefill, + ) + impl = resolve_fused_moe_impl(make_fused_moe_impl_kernel_name(IMPL, config)) + rms = token_rms(reference) + tasks = "-" + if parsed.down_path == "8x1_compact": + device_cu = get_cu_num() + output, tasks = run_compact(impl, request, experts, device_cu) + # The cached launcher must stay in bounds and correct for the other CU count. + for cu_count in CU_COUNTS: + if cu_count != device_cu: + with patch( + "aiter.ops.flydsl.fused_moe_gfx942.get_cu_num", + return_value=cu_count, + ): + other, _ = run_compact(impl, request, experts, cu_count) + diff = rms_error(output, other, rms) + assert diff < PATH_DIFF_LIMIT, f"{config} at {cu_count} CUs: {diff=}" + else: + output = impl(request) + err = check_output(reference, output, config) + # Decode kernels accumulate routes with atomics, so only prefill paths are compared. + if parsed.use_prefill: + first = first_output.setdefault("output", output.clone()) + path_diff = rms_error(first, output, rms) + assert ( + path_diff < PATH_DIFF_LIMIT + ), f"{config}: {path_diff=} from the first path" + return impl, request, tasks, err + + +@benchmark() +def test_moe(token, model_dim, inter_dim, experts, topk, quant, activation, config): + impl, request, tasks, err = check_moe( + token, model_dim, inter_dim, experts, topk, quant, activation, "random", config + ) + _, us = run_perftest(impl, request, num_iters=10, num_warmup=2) + flops = 6 * token * topk * model_dim * inter_dim + # Input and output activations plus both weight tensors. + nbytes = 2 * request.hidden_states.nbytes + request.w1.nbytes + request.w2.nbytes + return { + "full/tail tasks": tasks, + "us": us, + "TFLOPS": flops / us / 1e6, + "TB/s": nbytes / us / 1e6, + "err": err, + } + + +def shipped_rows(): + """Every shipped FlyDSL row as test_moe arguments, grouped by shape.""" + df = pd.read_csv(AITER_CONFIGS.AITER_CONFIG_FMOE_FILE) + df = df[df["kernelName1"].astype(str).str.startswith(f"impl__{IMPL}__")] + if "_tag" in df: + df = df[df["_tag"].fillna("").astype(str).str.strip() == ""] + quants = {"No": "no", "per_Tensor": "per_tensor", "per_Token": "ptpc"} + rows = set() + for _, row in df.iterrows(): + shape = ( + int(row["model_dim"]), + int(row["inter_dim"]), + int(row["expert"]), + int(row["topk"]), + quants[str(row["q_type"]).split(".")[-1]], + str(row["act_type"]).split(".")[-1].lower(), + ) + rows.add((shape, int(row["token"]), str(row["kernelName1"]).split("__", 2)[2])) + return [(token, *shape, config) for shape, token, config in sorted(rows)] + + +def cross_config_cases(rows): + """All shipped prefill configs of a shape on the same inputs, so outputs compare.""" + configs = {} + for _, *shape, config in rows: + if Config.from_string(config).use_prefill: + configs.setdefault(tuple(shape), set()).add(config) + return [ + (16384, *shape, "random", config) + for shape, shape_configs in sorted(configs.items()) + if len(shape_configs) > 1 + for config in sorted(shape_configs) + ] + + +# Functional checks only, not timed: +# (token, model_dim, inter_dim, experts, topk, quant, activation, routing, config) +EDGE_CASES = [ + # The 2048-expert limit of the compact task table, with all or 8 experts active. + *( + ( + 4096, + 512, + 256, + 2048, + 8, + "ptpc", + "silu", + routing, + f"64_256_128_True:{path}:128", + ) + for routing in ("random", "hot") + for path in ("8x1_compact", "1x4_64x256") + ), + # One expert with 64 full tasks: kept on 80 CUs, split into tails on 304. + *( + (16384, 1024, 384, 1, 1, "per_tensor", "silu", "hot", config) + for config in ( + "64_256_128_True:8x1_compact:128", + "64_256_128_True:1x4_64x256:128", + "64_256_128_True:default:0", + ) + ), + # Every row padding of the specialized Down paths, reduced with the same stride. + *( + ( + 16384, + 2048, + 512, + 257, + 9, + "ptpc", + "silu", + "random", + f"64_256_128_True:{path}:{padding}", + ) + for path in ("8x1_compact", "1x4_64x256") + for padding in (0, 32, 64) + ), + # Padding rows whose merged routes reach the reduction as -1. + *( + (2048, 2560, 320, 513, 11, quant, "silu", "padding", config) + for quant, config in ( + ("no", "64_128_128_True:default:0"), + ("ptpc", "64_128_256_True:1x4_64x256:128"), + ) + ), +] + + +def summarize(name, rows): + aiter.logger.info( + "%s summary (markdown):\n%s", name, pd.DataFrame(rows).to_markdown(index=False) + ) + + +def main(): + if get_gfx() != "gfx942": + aiter.logger.warning( + "gfx942 FlyDSL whole-MoE backend unsupported on %s; skipping", get_gfx() + ) + return + parser = argparse.ArgumentParser( + formatter_class=argparse.RawTextHelpFormatter, + description=__doc__, + ) + parser.add_argument( + "-b", + "--batch", + type=int, + nargs="*", + default=None, + help="Only time the shipped rows with these token counts (default: every row).", + ) + args = parser.parse_args() + + test_config_strings() + summarize( + "Compact task table", + [ + test_task_table(scenario, cu_count, cu_count) + for cu_count in CU_COUNTS + for scenario in task_scenarios(cu_count) + ] + # Buffers sized for 80 CUs under a 304-CU split: the device guard caps it. + + [test_task_table("cu_dependent", *CU_COUNTS)], + ) + summarize( + "Padded sorted reduction", + [ + test_sorted_sum(token, model_dim, topk, padding) + for token, (model_dim, topk), padding in itertools.product( + (65, 2048), + ((2048, 9), (2560, 11), (4096, 8), (6144, 4)), + (0, 32, 64, 128), + ) + ], + ) + rows = shipped_rows() + for case in cross_config_cases(rows) + EDGE_CASES: + check_moe(*case) + timed = [row for row in rows if not args.batch or row[0] in args.batch] + summarize("Whole MoE", [test_moe(*row) for row in timed]) + + +if __name__ == "__main__": + main() From 0f85400312fb99872cb03e516292de0b3e3596ad Mon Sep 17 00:00:00 2001 From: apinge Date: Tue, 27 Oct 2026 01:53:46 +0000 Subject: [PATCH 3/3] Backport Opus aux sorting dependencies and gfx942 regressions The cherry-picked gfx942 prefill change needs AUX_SORT_OPUS and initialized reverse maps on this branch. Request Opus explicitly while retaining the caller's sorting policy and the legacy boolean aux behavior. Backport the multi-phase aux outputs from fe4814c64b1f7bfe62742659eaf641b13d4f348d and reverse-map initialization/validation from dcb7c7fac07cdd9d0b4578a414af3bea0c8f0dc2. Remove inversion from this branch's AOT precompile entry point. Adapt the cherry-picked tests to existing kernel paths, four-field config suffixes, and compile-time CU specialization. Add deterministic reverse-map, all-skipped-route, BF16/FP8 zero-weight padding, and graph replay regressions. Retain FlyDSL 0.3.2 and the existing kernel compilation wrappers. Validation on MI308X/gfx942: 23 targeted pytest cases; 32 sorted reductions; 4 random/padding whole-MoE checks; task-table and compact-MoE smoke checks; BF16/FP8 AOT precompile; Black, Ruff, and git diff --check. --- aiter/fused_moe.py | 13 +- aiter/ops/flydsl/fused_moe_gfx942.py | 17 +- .../moe_gemm_2stage_gfx942/moe_reduce.py | 1 - csrc/include/moe_sorting_opus.h | 64 +++++++ csrc/py_itfs_cu/moe_sorting_opus_kernels.cu | 8 + op_tests/test_flydsl_moe_gfx942.py | 87 ++++----- op_tests/test_flydsl_moe_opus_aux.py | 181 ++++++++++++++++++ 7 files changed, 301 insertions(+), 70 deletions(-) create mode 100644 op_tests/test_flydsl_moe_opus_aux.py diff --git a/aiter/fused_moe.py b/aiter/fused_moe.py index a859150ea0c..7e23a6d7052 100644 --- a/aiter/fused_moe.py +++ b/aiter/fused_moe.py @@ -79,6 +79,8 @@ def is_flydsl_available(): # tuned-CSV rows routed to the port (exactly the codegen'd shapes). # opus / ck -> never use adaptive (legacy; ck still needs AITER_USE_CK_MOE_SORTING) _MOE_SORT_BACKEND = os.environ.get("AITER_MOE_SORT_BACKEND", "auto").lower() +# Request Opus auxiliary maps without entering the shape-specific adaptive sort. +AUX_SORT_OPUS = "opus" _ACT_TYPE_DISABLED_KEY = "__ignore__" _SWIGLU_MXFP4_BF16_BOUND = int(os.environ.get("GPTOSS_SWIGLU_MXFP4_BF16_BOUND", "256")) _MOE_A8W4_BYPASS_QUANT = os.environ.get("AITER_MOE_A8W4_BYPASS_QUANT", "0") == "1" @@ -197,7 +199,11 @@ def _moe_sorting_impl( device = topk_ids.device M, topk = topk_ids.shape - if output_aux and _MOE_SORT_BACKEND not in ("opus", "ck"): + if ( + output_aux + and output_aux != AUX_SORT_OPUS + and _MOE_SORT_BACKEND not in ("opus", "ck") + ): # adaptive (fused) sort emits the a4w4 extras (m_indices + reverse_sorted) # plus the atomic zero-init; opus single-pass aux is the env-gated fallback. return _adaptive_moe_sort( @@ -239,7 +245,10 @@ def _moe_sorting_impl( aux_reverse_sorted = None if output_aux: use_opus = True - dispatch_policy = 1 + # Opus aux callers keep their policy, including auto/multi-phase prefill. + # Boolean aux callers retain the legacy single-kernel path. + if output_aux != AUX_SORT_OPUS: + dispatch_policy = 1 aux_m_indices = torch.empty( max_num_tokens_padded, dtype=dtypes.i32, device=device ) diff --git a/aiter/ops/flydsl/fused_moe_gfx942.py b/aiter/ops/flydsl/fused_moe_gfx942.py index ede5b61d51e..7615f73fbba 100644 --- a/aiter/ops/flydsl/fused_moe_gfx942.py +++ b/aiter/ops/flydsl/fused_moe_gfx942.py @@ -209,8 +209,7 @@ def unsupported_reason(self, problem: "_Problem") -> str | None: else: if problem.inter_dim not in (192, 256, 320, 384, 512, 640): return ( - "8x1 paths require inter_dim in " - "{192, 256, 320, 384, 512, 640}" + "8x1 paths require inter_dim in " "{192, 256, 320, 384, 512, 640}" ) expected_block_m = 256 if self.down_path == "8x1" else 64 if self.BLOCK_M != expected_block_m or gateup_block_m != 64: @@ -687,9 +686,7 @@ def compile_kernel(*, stage, alg, block_m, block_n, **kwargs): ) metadata_capacity = ( - batch * topk + experts * config.BLOCK_M - topk - + config.BLOCK_M - - 1 + batch * topk + experts * config.BLOCK_M - topk + config.BLOCK_M - 1 ) // config.BLOCK_M full_capacity, tail_capacity = task_capacities( metadata_capacity, @@ -707,16 +704,6 @@ def compile_kernel(*, stage, alg, block_m, block_n, **kwargs): ) compile_launcher(down, *down_args, 0) - invert = invert_sorted_ids(topk) - compile_launcher( - invert._flydsl_executable, - _ptr(int32), - _ptr(int32), - _ptr(int32), - fx.Uint32(batch * topk), - fx.Uint32(batch), - 0, - ) reduction = sorted_sum(topk, model_dim, config.down_output_padding_bytes) compile_launcher( reduction._flydsl_executable, diff --git a/aiter/ops/flydsl/kernels/moe_gemm_2stage_gfx942/moe_reduce.py b/aiter/ops/flydsl/kernels/moe_gemm_2stage_gfx942/moe_reduce.py index 4a5683bba5f..a3497d39248 100644 --- a/aiter/ops/flydsl/kernels/moe_gemm_2stage_gfx942/moe_reduce.py +++ b/aiter/ops/flydsl/kernels/moe_gemm_2stage_gfx942/moe_reduce.py @@ -122,4 +122,3 @@ def sorted_sum(TOPK, N, row_padding_bytes=None): def compile_moe_reduction(*, topk, model_dim, row_padding_bytes=None): """Build and cache the route reduction callable for one static configuration.""" return sorted_sum(topk, model_dim, row_padding_bytes) - diff --git a/csrc/include/moe_sorting_opus.h b/csrc/include/moe_sorting_opus.h index d031294191b..edf8c917ea9 100644 --- a/csrc/include/moe_sorting_opus.h +++ b/csrc/include/moe_sorting_opus.h @@ -1124,6 +1124,14 @@ struct MoeSortingKernel return; } + if(kargs.p_reverse_sorted) + { + opus::index_t* p_reverse_sorted = static_cast(kargs.p_reverse_sorted); + const opus::index_t routes = kargs.tokens * kargs.topk_mdiv.divisor; + for(opus::index_t i = threadIdx.x; i < routes; i += kBlockSize) + p_reverse_sorted[i] = -1; + } + extern __shared__ char smem[]; return moe_align_block_size_kernel_ex( @@ -1362,6 +1370,22 @@ OPUS_D void moe_buf_set_zero_kernel_2d(void* buf, } } +template +OPUS_D void moe_sorting_reverse_sorted_set_invalid(void* p_reverse_sorted, + opus::index_t tokens, + opus::index_t topk) +{ + if(p_reverse_sorted == nullptr) + return; + opus::index_t* p_buf = reinterpret_cast(p_reverse_sorted); + const opus::index_t routes = tokens * topk; + for(opus::index_t i = blockIdx.x * kBlockSize + threadIdx.x; i < routes; + i += gridDim.x * kBlockSize) + { + p_buf[i] = -1; + } +} + } // namespace impl // TODO: tokens could be from @@ -1583,6 +1607,7 @@ struct MoeSortingMultiPhaseKernel_P0_v1 const void* p_topk_ids; // [tokens, topk] const void* p_local_tokens; // [1], if not nullptr, use this as actual tokens void* p_expert_mesh; // [expert, tokens] + void* p_reverse_sorted; opus::index_t tokens; // if p_local_tokens is not nullptr, this indicate the max possible // tokens used for ws/LDS calculation opus::index_t num_experts; @@ -1608,6 +1633,7 @@ struct MoeSortingMultiPhaseKernel_P0_v1 k.p_topk_ids = h.p_topk_ids; k.p_local_tokens = h.p_local_tokens; k.p_expert_mesh = h.p_ws; + k.p_reverse_sorted = h.p_reverse_sorted; k.tokens = h.tokens; k.num_experts = h.num_experts; k.mesh_stride = impl::moe_sorting_mp_mesh_stride(h.tokens); @@ -1624,6 +1650,9 @@ struct MoeSortingMultiPhaseKernel_P0_v1 OPUS_D void operator()(Kargs kargs) const { + impl::moe_sorting_reverse_sorted_set_invalid( + kargs.p_reverse_sorted, kargs.tokens, kargs.topk_mdiv.divisor); + using topk_id_t = opus::vector_t; const topk_id_t* p_topk_ids = reinterpret_cast(kargs.p_topk_ids); @@ -1711,6 +1740,7 @@ struct MoeSortingMultiPhaseKernel_P0_v2 const void* p_local_expert_mask; // [expert] void* p_expert_cumsum; // [expert] + void* p_reverse_sorted; opus::index_t num_experts; }; @@ -1739,6 +1769,7 @@ struct MoeSortingMultiPhaseKernel_P0_v2 k.mesh_stride = impl::moe_sorting_mp_mesh_stride(h.tokens); k.topk_mdiv = opus::mdiv{static_cast(h.topk)}; k.p_local_expert_mask = h.p_local_expert_mask; + k.p_reverse_sorted = h.p_reverse_sorted; k.num_experts = h.num_experts; return k; } @@ -1755,6 +1786,9 @@ struct MoeSortingMultiPhaseKernel_P0_v2 OPUS_D void operator()(Kargs kargs) const { + impl::moe_sorting_reverse_sorted_set_invalid( + kargs.p_reverse_sorted, kargs.tokens, kargs.topk_mdiv.divisor); + constexpr opus::index_t index_pack = Problem::SubTokenTile; // always packed __shared__ char smem[GetSmemSize()]; using topk_id_t = opus::vector_t; @@ -2505,6 +2539,8 @@ struct MoeSortingMultiPhaseKernel_P3 void* p_sorted_weights; void* p_expert_mesh; // [token, expert] void* p_expert_cumsum; + void* p_m_indices; // optional, see MoeSortingMultiPhaseKernel_P23::Kargs + void* p_reverse_sorted; // optional, see MoeSortingMultiPhaseKernel_P23::Kargs opus::index_t tokens; opus::index_t num_experts; @@ -2524,6 +2560,8 @@ struct MoeSortingMultiPhaseKernel_P3 k.p_expert_cumsum = reinterpret_cast( reinterpret_cast(h.p_ws) + impl::moe_sorting_mp_mesh_smem_size(h.tokens, h.num_experts, h.topk)); + k.p_m_indices = h.p_m_indices; + k.p_reverse_sorted = h.p_reverse_sorted; k.tokens = h.tokens; k.num_experts = h.num_experts; k.topk_mdiv = opus::mdiv{static_cast(h.topk)}; @@ -2553,6 +2591,8 @@ struct MoeSortingMultiPhaseKernel_P3 IndexType* p_expert_cumsum = reinterpret_cast(kargs.p_expert_cumsum); const WeightType* p_weights = static_cast(kargs.p_weights); WeightType* p_sorted_weights = reinterpret_cast(kargs.p_sorted_weights); + IndexType* p_m_indices = reinterpret_cast(kargs.p_m_indices); + IndexType* p_reverse_sorted = reinterpret_cast(kargs.p_reverse_sorted); opus::index_t tokens = [&]() { if constexpr(Problem::LocalToken) @@ -2630,6 +2670,11 @@ struct MoeSortingMultiPhaseKernel_P3 #endif p_sorted_weights[e_start + position] = p_weights[i_token * kargs.topk_mdiv.divisor + i_topk]; + if(p_m_indices) + p_m_indices[e_start + position] = i_token & 0x00FFFFFF; + if(p_reverse_sorted) + p_reverse_sorted[i_token * kargs.topk_mdiv.divisor + i_topk] = + e_start + position; } } @@ -2641,6 +2686,8 @@ struct MoeSortingMultiPhaseKernel_P3 p_sorted_token_ids[i] = tokens; #endif p_sorted_weights[i] = static_cast(0.0); + if(p_m_indices) + p_m_indices[i] = tokens; // pad = tokens -> OOB row for gemm1 } } }; @@ -2689,6 +2736,11 @@ struct MoeSortingMultiPhaseKernel_P23 void* p_sorted_weights; void* p_moe_buf; void* p_local_topk_ids; + // Optional a4w4 extras, same contract as the single-kernel + // MoeSortingKernel: without them the MXFP4 callers are forced onto + // dispatch_policy=1 (that single kernel) purely to get these two arrays. + void* p_m_indices; // [total_padded] sorted slot -> token id (pad = tokens) + void* p_reverse_sorted; // [token*topk] (token, slot) -> sorted slot opus::index_t tokens; opus::index_t num_experts; @@ -2724,6 +2776,9 @@ struct MoeSortingMultiPhaseKernel_P23 k.p_moe_buf = h.p_moe_buf; k.p_local_topk_ids = h.p_local_topk_ids; + k.p_m_indices = h.p_m_indices; + k.p_reverse_sorted = h.p_reverse_sorted; + k.tokens = h.tokens; k.num_experts = h.num_experts; k.mesh_stride = impl::moe_sorting_mp_mesh_stride(h.tokens); @@ -2982,6 +3037,8 @@ struct MoeSortingMultiPhaseKernel_P23 IndexType* p_expert_cumsum_smem = s + 4 + 2 * kBlockSize / opus::get_warp_size(); const WeightType* p_weights = static_cast(kargs.p_weights); WeightType* p_sorted_weights = reinterpret_cast(kargs.p_sorted_weights); + IndexType* p_m_indices = reinterpret_cast(kargs.p_m_indices); + IndexType* p_reverse_sorted = reinterpret_cast(kargs.p_reverse_sorted); int eid = blockIdx.x; int wave_id = threadIdx.x / opus::get_warp_size(); @@ -3148,6 +3205,11 @@ struct MoeSortingMultiPhaseKernel_P23 #endif p_sorted_weights[e_start + position] = p_weights[i_token * kargs.topk_mdiv.divisor + i_topk[j]]; + if(p_m_indices) + p_m_indices[e_start + position] = i_token & 0x00FFFFFF; + if(p_reverse_sorted) + p_reverse_sorted[i_token * kargs.topk_mdiv.divisor + + i_topk[j]] = e_start + position; } position += i_show[j]; }); @@ -3225,6 +3287,8 @@ struct MoeSortingMultiPhaseKernel_P23 p_sorted_token_ids[i] = tokens; #endif p_sorted_weights[i] = static_cast(0.0); + if(p_m_indices) + p_m_indices[i] = tokens; // pad = tokens -> OOB row for gemm1 } } } diff --git a/csrc/py_itfs_cu/moe_sorting_opus_kernels.cu b/csrc/py_itfs_cu/moe_sorting_opus_kernels.cu index 0e91f180070..c7335b2b90e 100644 --- a/csrc/py_itfs_cu/moe_sorting_opus_kernels.cu +++ b/csrc/py_itfs_cu/moe_sorting_opus_kernels.cu @@ -51,6 +51,14 @@ void moe_sorting_opus_fwd(aiter_tensor_t& topk_ids, HipDeviceGuard device_guard(topk_ids.device_id); const hipStream_t stream = aiter::getCurrentHIPStream(); + if(reverse_sorted.has_value()) + { + auto& reverse = reverse_sorted.value(); + AITER_CHECK(reverse.dtype() == AITER_DTYPE_i32 && reverse.is_contiguous() && + reverse.numel() >= static_cast(num_tokens) * topk, + "reverse_sorted must be a contiguous int32 buffer for tokens * topk routes"); + } + void* ws_ptr = workspace.has_value() ? workspace.value().data_ptr() : nullptr; moe_sorting_opus( diff --git a/op_tests/test_flydsl_moe_gfx942.py b/op_tests/test_flydsl_moe_gfx942.py index 28b67ab3756..96ac906a5fb 100644 --- a/op_tests/test_flydsl_moe_gfx942.py +++ b/op_tests/test_flydsl_moe_gfx942.py @@ -34,11 +34,13 @@ from aiter.jit.core import AITER_CONFIGS from aiter.jit.utils.chip_info import get_cu_num, get_gfx from aiter.ops.flydsl.fused_moe_gfx942 import Config, get_tune_space -from aiter.ops.flydsl.kernels.moe_gemm_2stage.gfx9_fp8 import ( +from aiter.ops.flydsl.kernels.moe_gemm_2stage_gfx942 import ( gemm2_8x1_compact as compact, ) -from aiter.ops.flydsl.kernels.moe_gemm_2stage.gfx9_fp8 import sorted_sum -from aiter.ops.flydsl.kernels.moe_gemm_2stage.gfx9_fp8.common import _ptr +from aiter.ops.flydsl.kernels.moe_gemm_2stage_gfx942 import sorted_sum +from aiter.ops.flydsl.kernels.moe_gemm_2stage_gfx942.common import ( + torch_tensor_to_pointer as _ptr, +) from aiter.ops.flydsl.kernels.tensor_shim import _run_compiled from aiter.ops.shuffle import shuffle_weight from aiter.test_common import benchmark, checkAllclose, run_perftest @@ -52,7 +54,7 @@ "ptpc": QuantType.per_Token, } ACTIVATIONS = {"silu": ActivationType.Silu, "swiglu": ActivationType.Swiglu} -# MI300X then MI308X, the order in which a cached launcher meets a smaller device. +# Exercise task tables compiled separately for MI300X and MI308X CU counts. CU_COUNTS = (304, 80) GUARD, CANARY = 64, -123 # FP8 rounding moves an element by at most 0.28 token RMS from the torch reference. @@ -66,7 +68,9 @@ def guarded_task_buffers(sorted_expert_ids, experts, cu_count): """allocate_task_buffers with canaries before and after every buffer.""" - full, tail = compact.task_capacities(sorted_expert_ids.numel(), experts, cu_count) + full, tail = compact.task_capacities( + sorted_expert_ids.numel(), experts, cu_count=cu_count + ) storages, buffers = [], [] for shape in ((full, 2), (tail, 2), (2,)): storage = torch.full( @@ -95,10 +99,8 @@ def check_tasks(blocks, experts, cu_count, full_tasks, tail_tasks, counts): tail = sum(size % 4 for size in sizes) remainder = full % cu_count split = 0 - if remainder and remainder * compact._SPLIT_DEN < cu_count * compact._SPLIT_NUM: + if remainder and remainder * 5 < cu_count * 3: split = remainder - # The kernel caps the split at the free tail capacity. - split = min(split, (tail_tasks.shape[0] - tail) // 4) expected = [(full - split) * 256, (tail + 4 * split) * 64] assert counts.tolist() == expected, f"counts {counts.tolist()} != {expected}" coverage = [0] * len(blocks) @@ -119,7 +121,7 @@ def check_tasks(blocks, experts, cu_count, full_tasks, tail_tasks, counts): def task_scenarios(cu_count): """Map a scenario to (experts, [(expert, M64 blocks)]) in metadata order.""" # Largest remainder of the final full-task CU round that is still split. - split = (cu_count * compact._SPLIT_NUM - 1) // compact._SPLIT_DEN + split = (cu_count * 3 - 1) // 5 return { "empty": (4, []), "tails_only": (8, [(1, 1), (3, 2), (6, 3)]), @@ -135,8 +137,8 @@ def task_scenarios(cu_count): } -# One launcher per expert count serves every CU count, like the cached builder. -build_task_table = functools.cache(compact.build_task_table) +# This branch specializes task-table launchers by CU count. +build_task_table = compact.build_task_table @benchmark() @@ -149,15 +151,13 @@ def test_task_table(scenario, cu_count, buffer_cu_count): storages, (full_tasks, tail_tasks, counts) = guarded_task_buffers( metadata, experts, buffer_cu_count ) - launcher = build_task_table(experts) + launcher = build_task_table(experts, cu_count=cu_count) pointers = [_ptr(t) for t in (metadata, valid_rows, full_tasks, tail_tasks, counts)] def build(): _run_compiled( launcher, *pointers, - cu_count, - tail_tasks.shape[0], torch.cuda.current_stream(), ) return counts @@ -165,8 +165,6 @@ def build(): _, us = run_perftest(build) full, tail = check_tasks(blocks, experts, cu_count, full_tasks, tail_tasks, counts) check_guards(storages) - if scenario == "tail_capacity" and buffer_cu_count == cu_count: - assert tail == tail_tasks.shape[0], "tail capacity is not tight" nbytes = metadata.nbytes + valid_rows.nbytes + 8 * (full + tail) + counts.nbytes return { "blocks": len(blocks), @@ -195,9 +193,9 @@ def test_sorted_sum(token, model_dim, topk, padding): model_dim, dtype, 64, - output=storage[GUARD:-GUARD].view(token, model_dim), output_aux=AUX_SORT_OPUS, ) + output = storage[GUARD:-GUARD].view(token, model_dim) valid = loc_ids.view(token, topk) >= 0 assert int((~valid).sum()) == topk - 1, "repeated expert IDs kept extra routes" # Specialized Down rows carry padding that holds canaries the sum must skip. @@ -238,25 +236,18 @@ def run(): def test_config_strings(): for encoded in get_tune_space(): assert Config.from_string(encoded).to_string() == encoded - # Rows tuned before the Down suffix existed keep loading as the default Down. for legacy in ("16_16_16_False", "64_128_128_True", "64_256_128_True"): - assert Config.from_string(legacy) == Config.from_string(f"{legacy}:default:0") - # Tiles without the direct-decoding flag keep their meaning. - assert Config.from_string("16_16_16_False_False:default:0") == Config.from_string( - "16_16_16_False:default:0" - ) + parsed = Config.from_string(legacy) + assert parsed.down_path == "default" + assert parsed.GATEUP_BLOCK_M is None + assert parsed.down_output_padding_bytes is None for invalid in ( "64_256_128_True:default", - "64_256_128_True:default:0:0", "64_256_128_Maybe", "64_256_True", - "64_256_128_True:default:128", - "64_256_128_True:8x2:128", - "16_16_16_False:8x1_compact:128", - "16_16_16_False_Maybe:default:0", - "16_16_16_False_True_False:default:0", - "16_16_16_True_True:default:0", - "16_16_16_False_True:1x4_64x256:128", + "64_256_128_True:8x2:64:128", + "16_16_16_False_Maybe", + "16_16_16_False_True_False", ): try: Config.from_string(invalid) @@ -363,7 +354,11 @@ def quantize_activation(x): class GuardedTaskBuffers: """Replaces allocate_task_buffers and keeps the last allocation for checks.""" - def __call__(self, sorted_expert_ids, experts, cu_count): + def __call__( + self, sorted_expert_ids, experts, *, cu_count=None, min_tail_utilization=0.6 + ): + if cu_count is None: + cu_count = compact.device_cu_count(sorted_expert_ids.device) self.metadata, self.cu_count = sorted_expert_ids, cu_count self.storages, self.buffers = guarded_task_buffers( sorted_expert_ids, experts, cu_count @@ -442,16 +437,6 @@ def check_moe( if parsed.down_path == "8x1_compact": device_cu = get_cu_num() output, tasks = run_compact(impl, request, experts, device_cu) - # The cached launcher must stay in bounds and correct for the other CU count. - for cu_count in CU_COUNTS: - if cu_count != device_cu: - with patch( - "aiter.ops.flydsl.fused_moe_gfx942.get_cu_num", - return_value=cu_count, - ): - other, _ = run_compact(impl, request, experts, cu_count) - diff = rms_error(output, other, rms) - assert diff < PATH_DIFF_LIMIT, f"{config} at {cu_count} CUs: {diff=}" else: output = impl(request) err = check_output(reference, output, config) @@ -532,7 +517,7 @@ def cross_config_cases(rows): "ptpc", "silu", routing, - f"64_256_128_True:{path}:128", + f"64_256_128_True:{path}:64:128", ) for routing in ("random", "hot") for path in ("8x1_compact", "1x4_64x256") @@ -541,9 +526,9 @@ def cross_config_cases(rows): *( (16384, 1024, 384, 1, 1, "per_tensor", "silu", "hot", config) for config in ( - "64_256_128_True:8x1_compact:128", - "64_256_128_True:1x4_64x256:128", - "64_256_128_True:default:0", + "64_256_128_True:8x1_compact:64:128", + "64_256_128_True:1x4_64x256:64:128", + "64_256_128_True:default:64:none", ) ), # Every row padding of the specialized Down paths, reduced with the same stride. @@ -557,7 +542,7 @@ def cross_config_cases(rows): "ptpc", "silu", "random", - f"64_256_128_True:{path}:{padding}", + f"64_256_128_True:{path}:64:{padding}", ) for path in ("8x1_compact", "1x4_64x256") for padding in (0, 32, 64) @@ -566,8 +551,8 @@ def cross_config_cases(rows): *( (2048, 2560, 320, 513, 11, quant, "silu", "padding", config) for quant, config in ( - ("no", "64_128_128_True:default:0"), - ("ptpc", "64_128_256_True:1x4_64x256:128"), + ("no", "64_128_128_True:default:64:none"), + ("ptpc", "64_128_256_True:1x4_64x256:64:128"), ) ), ] @@ -606,9 +591,7 @@ def main(): test_task_table(scenario, cu_count, cu_count) for cu_count in CU_COUNTS for scenario in task_scenarios(cu_count) - ] - # Buffers sized for 80 CUs under a 304-CU split: the device guard caps it. - + [test_task_table("cu_dependent", *CU_COUNTS)], + ], ) summarize( "Padded sorted reduction", diff --git a/op_tests/test_flydsl_moe_opus_aux.py b/op_tests/test_flydsl_moe_opus_aux.py new file mode 100644 index 00000000000..f0603b2f432 --- /dev/null +++ b/op_tests/test_flydsl_moe_opus_aux.py @@ -0,0 +1,181 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Opus reverse-map and zero-weight padding regressions on gfx942.""" + +from contextlib import contextmanager + +import pytest +import torch + +from aiter import ActivationType, QuantType +from aiter.fused_moe import AUX_SORT_OPUS, moe_sorting +from aiter.ops.flydsl.fused_moe_gfx942 import run_flydsl_moe_gfx942 +from aiter.ops.flydsl.kernels.moe_gemm_2stage_gfx942 import sorted_sum +from aiter.ops.shuffle import shuffle_weight + + +@pytest.fixture(autouse=True) +def gfx942_device(): + if not torch.cuda.is_available(): + pytest.skip("requires gfx942") + arch = torch.cuda.get_device_properties(0).gcnArchName + if not arch.startswith("gfx942"): + pytest.skip("requires gfx942") + + +@contextmanager +def poison_uninitialized_memory(): + enabled = torch.are_deterministic_algorithms_enabled() + warn_only = torch.is_deterministic_algorithms_warn_only_enabled() + fill = torch.utils.deterministic.fill_uninitialized_memory + torch.use_deterministic_algorithms(True) + torch.utils.deterministic.fill_uninitialized_memory = True + try: + yield + finally: + torch.use_deterministic_algorithms(enabled, warn_only=warn_only) + torch.utils.deterministic.fill_uninitialized_memory = fill + + +@pytest.mark.parametrize("policy", (0, 1, 2)) +@pytest.mark.parametrize("batch,experts", ((65, 16), (2048, 513), (16346, 513))) +def test_reverse_routes(batch, experts, policy): + topk = 11 + ids = torch.arange(topk, dtype=torch.int32, device="cuda").expand(batch, -1).clone() + ids[-1, :-1] = 0 + weights = torch.ones((batch, topk), dtype=torch.float32, device="cuda") + weights[-1] = 0 + with poison_uninitialized_memory(): + sorted_ids, _, _, valid_rows, _, m_indices, reverse = moe_sorting( + ids, + weights, + experts, + 256, + torch.bfloat16, + 64, + dispatch_policy=policy, + output_aux=AUX_SORT_OPUS, + ) + reverse = reverse.view(batch, topk) + valid = reverse >= 0 + assert int((~valid).sum()) == 9 + assert bool((reverse[~valid] == -1).all()) + assert bool((reverse[valid] < valid_rows[0]).all()) + # Check every emitted inverse against the actual sorted token/topk pair. + packed = sorted_ids[reverse[valid].long()] + slots = torch.arange(batch * topk, device="cuda").view(batch, topk)[valid] + torch.testing.assert_close(packed & 0xFFFFFF, (slots // topk).int()) + torch.testing.assert_close(packed >> 24, (slots % topk).int()) + torch.testing.assert_close(m_indices[reverse[valid].long()], (slots // topk).int()) + + +@pytest.mark.parametrize("padding", (0, 32, 64, 128)) +def test_all_routes_skipped(padding): + batch, topk, dim = 65, 11, 2560 + locs = torch.full((batch, topk), -1, dtype=torch.int32, device="cuda") + # The dummy address is in bounds but contains NaN. Select must yield zero. + rows = torch.full( + (1, dim + padding // 2), float("nan"), dtype=torch.bfloat16, device="cuda" + ) + guard = 64 + storage = torch.full( + (batch * dim + 2 * guard,), -123, dtype=torch.bfloat16, device="cuda" + ) + output = storage[guard:-guard].view(batch, dim) + sorted_sum(topk, dim, padding)(locs, rows, output, batch) + assert bool((output == 0).all()) + assert bool((storage[:guard] == -123).all()) + assert bool((storage[-guard:] == -123).all()) + + +def make_padding_case(batch, masked, quant, routing): + experts, topk, dim, inter = 513, 11, 2560, 320 + fp8 = quant == "fp8" + dtype = torch.float8_e4m3fnuz if fp8 else torch.bfloat16 + w1 = shuffle_weight( + torch.full( + (experts, 2 * inter, dim), + 1 if fp8 else 1 / dim, + dtype=torch.bfloat16, + device="cuda", + ).to(dtype) + ) + w2 = shuffle_weight( + torch.full( + (experts, dim, inter), + 1 if fp8 else 1 / inter, + dtype=torch.bfloat16, + device="cuda", + ).to(dtype) + ) + hidden = torch.ones((batch, dim), dtype=torch.bfloat16, device="cuda") + ids = torch.arange(topk, dtype=torch.int32, device="cuda").expand(batch, -1).clone() + ids[:, -1] = 512 + weights = torch.full((batch, topk), 1 / topk, device="cuda") + weights[-masked:] = 0 + if routing == "mixed": + ids[-masked:, :-1] = 0 + elif routing == "allzero": + ids[-masked:] = 0 + scale1 = torch.full((experts, 2 * inter), 1 / dim, device="cuda") if fp8 else None + scale2 = torch.full((experts, dim), 1 / inter, device="cuda") if fp8 else None + config = "64_128_256_True:1x4_64x256:64:128" if fp8 else "64_128_128_True" + + def run(): + return run_flydsl_moe_gfx942( + hidden, + w1, + w2, + weights, + ids, + ActivationType.Silu, + QuantType.per_Token if fp8 else QuantType.No, + scale1, + scale2, + None, + None, + 0, + config, + ) + + return run, weights, ids + + +def check_padding_output(output, masked): + assert bool(torch.isfinite(output).all()) + if output.shape[0] > masked: + assert float((output[:-masked].float() - 0.7310585786).abs().max()) < 0.005 + assert bool((output[-masked:] == 0).all()) + + +@pytest.mark.parametrize("quant", ("bf16", "fp8")) +@pytest.mark.parametrize( + "batch,masked,routing", + ( + (2048, 1, "mixed"), + (16346, 65, "mixed"), + (2048, 2048, "allzero"), + (2048, 1, "distinct"), + ), +) +def test_zero_weight_padding(batch, masked, quant, routing): + with poison_uninitialized_memory(): + run, _, _ = make_padding_case(batch, masked, quant, routing) + check_padding_output(run(), masked) + + +@pytest.mark.parametrize("quant", ("bf16", "fp8")) +def test_padding_graph_replay(quant): + run, weights, ids = make_padding_case(2048, 65, quant, "mixed") + run() # Compile before capture. + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = run() + for masked in (65, 2048, 1): + ids[:, :-1] = torch.arange(10, dtype=torch.int32, device="cuda") + ids[-masked:, :-1] = 0 + weights.fill_(1 / 11) + weights[-masked:] = 0 + graph.replay() + check_padding_output(output, masked)