diff --git a/auto_round_extension/ark/auto_round_kernel/__init__.py b/auto_round_extension/ark/auto_round_kernel/__init__.py index b66e180d8..4b5a5726b 100644 --- a/auto_round_extension/ark/auto_round_kernel/__init__.py +++ b/auto_round_extension/ark/auto_round_kernel/__init__.py @@ -2331,8 +2331,9 @@ def _native_fp8_prefill_enabled() -> bool: # # `moe_gemm_decode` and `moe_gemm_prefill` accept identical argument shapes # and dtypes -- the only difference is which underlying SYCL kernel is -# launched (a GEMV variant tuned for 1-2 tokens/expert vs. a Grouped GEMM -# variant tuned for many tokens/expert). Model code that runs through both +# launched (a GEMV variant tuned for smaller total-token workloads vs. a +# Grouped GEMM variant tuned for larger total-token workloads). Model code +# that runs through both # regimes (prefill of a prompt, then autoregressive decode) traditionally # has to keep two call sites and branch on phase. `moe(...)` collapses that # into a single API and auto-selects the right kernel from the token @@ -2340,21 +2341,40 @@ def _native_fp8_prefill_enabled() -> bool: # # Callers that already know the phase (e.g., a model's generation loop knows # whether it's in prefill or decode) should pass it via the `phase` argument -# to avoid the small host-device sync that `phase="auto"` needs to inspect -# `num_tokens_per_expert.max()`. +# to bypass the auto-dispatch heuristic entirely. # --------------------------------------------------------------------------- -# Default tokens-per-expert threshold used by `phase="auto"`. The decode -# GEMV kernel is faster when every expert sees only a handful of tokens -# (TopK >= 1 with batch size 1-4); above that the GEMM-tuned prefill kernel -# wins. The crossover is hardware-dependent but `4` is a conservative default -# that matches the regime `moe_gemm_decode`'s docstring describes -# ("typically only 1-2 tokens", up to top-k * small batch). -_MOE_AUTO_DECODE_MAX_TOKENS_PER_EXPERT = 4 +# Default total-token threshold used by `phase="auto"`: dispatch to decode +# when `activations.shape[0] <= threshold`, otherwise prefill. This threshold +# is hardware-dependent and can be overridden via +# `ARK_MOE_AUTO_DECODE_MAX_TOKENS`. +_MOE_AUTO_DECODE_MAX_TOTAL_TOKENS = 256 _MOE_VALID_PHASES = ("auto", "decode", "prefill") +def _moe_auto_decode_max_total_tokens() -> int: + """Return auto decode threshold from env or the module default. + + ``ARK_MOE_AUTO_DECODE_MAX_TOKENS`` is accepted when it is a positive + integer. Unset/empty/invalid values fall back to + ``_MOE_AUTO_DECODE_MAX_TOTAL_TOKENS``. + """ + env = os.environ.get("ARK_MOE_AUTO_DECODE_MAX_TOKENS") + if env is None: + return _MOE_AUTO_DECODE_MAX_TOTAL_TOKENS + env = env.strip() + if not env: + return _MOE_AUTO_DECODE_MAX_TOTAL_TOKENS + try: + value = int(env) + except ValueError: + return _MOE_AUTO_DECODE_MAX_TOTAL_TOKENS + if value <= 0: + return _MOE_AUTO_DECODE_MAX_TOTAL_TOKENS + return value + + def moe( activations: torch.Tensor, weights: torch.Tensor, @@ -2366,7 +2386,7 @@ def moe( group_size: int = 128, asym: bool = False, phase: str = "auto", - decode_threshold: int = _MOE_AUTO_DECODE_MAX_TOKENS_PER_EXPERT, + decode_threshold: Optional[int] = None, ) -> torch.Tensor: """Unified MoE GEMM entry point that dispatches to decode or prefill. @@ -2387,18 +2407,18 @@ def moe( underlying kernel; see :func:`moe_gemm_decode`. phase: dispatch mode. - * ``"auto"`` (default): inspect ``num_tokens_per_expert.max()`` - and pick decode if every expert sees ``<= decode_threshold`` - tokens, otherwise prefill. This incurs one small host-device - sync per call. + * ``"auto"`` (default): dispatch to decode when + ``activations.shape[0] <= decode_threshold`` (total tokens), + otherwise prefill. * ``"decode"``: always dispatch to :func:`moe_gemm_decode`. Use when the model's generation loop already knows it is in the - decode phase; avoids the sync. + decode phase. * ``"prefill"``: always dispatch to :func:`moe_gemm_prefill`. Use when the model knows it is in the prefill phase. - decode_threshold: ``"auto"`` mode dispatches to decode when - ``num_tokens_per_expert.max() <= decode_threshold``. Defaults to - 4 (the regime the decode GEMV kernel is tuned for). + decode_threshold: Total-token threshold for ``"auto"`` mode. If not + provided, uses ``ARK_MOE_AUTO_DECODE_MAX_TOKENS`` when set to a + valid positive integer, otherwise defaults to 256. Explicit + argument values take precedence over the environment variable. Returns: ``[total_tokens, N]`` in the activations dtype. Bit-identical to the @@ -2408,14 +2428,11 @@ def moe( raise ValueError(f"phase must be one of {_MOE_VALID_PHASES}, got {phase!r}") if phase == "auto": - # `.max().item()` triggers a host-device sync; callers in tight - # decode loops should pass `phase="decode"` explicitly to skip this. - # We tolerate a non-int32 / non-contiguous tensor here because the - # downstream kernel wrappers will normalise it anyway. + threshold = _moe_auto_decode_max_total_tokens() if decode_threshold is None else int(decode_threshold) if num_tokens_per_expert.numel() == 0: raise ValueError("num_tokens_per_expert must be non-empty") - max_tpe = int(num_tokens_per_expert.max().item()) - phase = "decode" if max_tpe <= int(decode_threshold) else "prefill" + total_tokens = int(activations.shape[0]) + phase = "decode" if total_tokens <= threshold else "prefill" if phase == "decode": return moe_gemm_decode( diff --git a/auto_round_extension/ark/test/test_moe_unified.py b/auto_round_extension/ark/test/test_moe_unified.py index 4e769456f..98bea8057 100644 --- a/auto_round_extension/ark/test/test_moe_unified.py +++ b/auto_round_extension/ark/test/test_moe_unified.py @@ -25,8 +25,8 @@ This file checks: - * Dispatch correctness: ``phase="auto"`` picks decode when every expert - sees few tokens and prefill otherwise. + * Dispatch correctness: ``phase="auto"`` picks decode when total tokens are + below threshold and prefill otherwise. * Bit-parity: ``moe(phase="auto")`` matches the kernel it dispatched to. * Explicit-phase parity: ``moe(phase="decode")`` matches ``moe_gemm_decode``, ``moe(phase="prefill")`` matches @@ -88,11 +88,12 @@ def _unified_skip_reason() -> str: # --------------------------------------------------------------------------- -# Small shapes (one decode-shaped, one prefill-shaped) -- keep wall-clock low. +# Small shapes (one decode-sized by total tokens, one prefill-sized by total +# tokens) -- keep wall-clock low. # --------------------------------------------------------------------------- -_DECODE_SHAPE = dict(num_experts=4, tokens_per_expert=[1, 2, 0, 2], N=128, K=256) -_PREFILL_SHAPE = dict(num_experts=4, tokens_per_expert=[16, 8, 0, 20], N=128, K=256) +_AUTO_DECODE_SHAPE = dict(num_experts=4, tokens_per_expert=[64, 64, 64, 64], N=128, K=256) # total_tokens=256 +_AUTO_PREFILL_SHAPE = dict(num_experts=4, tokens_per_expert=[80, 80, 80, 80], N=128, K=256) # total_tokens=320 def _make_int4_sym(E, N, K, group_size, dtype, total_tokens): @@ -163,8 +164,8 @@ def _make_fp8(E, N, K, group_size, dtype, total_tokens, fp8_dtype): class TestMoeUnifiedDispatch: """Tests for the auto-dispatch logic itself.""" - def test_auto_picks_decode_for_small_tokens_per_expert(self): - shape = _DECODE_SHAPE + def test_auto_picks_decode_for_small_total_tokens(self): + shape = _AUTO_DECODE_SHAPE total_tokens = sum(shape["tokens_per_expert"]) E, N, K = shape["num_experts"], shape["N"], shape["K"] group_size = 128 @@ -192,12 +193,12 @@ def test_auto_picks_decode_for_small_tokens_per_expert(self): group_size=group_size, asym=False, ) - # max tokens/expert = 2 (<= default threshold 4) -> dispatched to decode + # total tokens = 256 (<= default threshold 256) -> dispatched to decode # -> output must be bit-identical to moe_gemm_decode. torch.testing.assert_close(out_auto, out_decode, rtol=0, atol=0) - def test_auto_picks_prefill_for_large_tokens_per_expert(self): - shape = _PREFILL_SHAPE + def test_auto_picks_prefill_for_large_total_tokens(self): + shape = _AUTO_PREFILL_SHAPE total_tokens = sum(shape["tokens_per_expert"]) E, N, K = shape["num_experts"], shape["N"], shape["K"] group_size = 128 @@ -228,9 +229,66 @@ def test_auto_picks_prefill_for_large_tokens_per_expert(self): torch.testing.assert_close(out_auto, out_prefill, rtol=0, atol=0) def test_decode_threshold_override(self): - # Same prefill-shaped input but bump the threshold above the max - # tokens/expert -> auto must now pick decode. - shape = _PREFILL_SHAPE + decode_shape = _AUTO_DECODE_SHAPE + decode_total_tokens = sum(decode_shape["tokens_per_expert"]) + E, N, K = decode_shape["num_experts"], decode_shape["N"], decode_shape["K"] + group_size = 128 + dtype = torch.float16 + + activations, packed, scales, _ = _make_int4_sym(E, N, K, group_size, dtype, decode_total_tokens) + ntpe = torch.tensor(decode_shape["tokens_per_expert"], dtype=torch.int32, device="xpu") + + out_auto = ark.moe( + activations, + packed, + ntpe, + scales=scales, + weight_bits=4, + group_size=group_size, + asym=False, + phase="auto", + decode_threshold=128, + ) + out_prefill = ark.moe_gemm_prefill( + activations, + packed, + ntpe, + scales=scales, + weight_bits=4, + group_size=group_size, + asym=False, + ) + torch.testing.assert_close(out_auto, out_prefill, rtol=0, atol=0) + + prefill_shape = _AUTO_PREFILL_SHAPE + prefill_total_tokens = sum(prefill_shape["tokens_per_expert"]) + activations, packed, scales, _ = _make_int4_sym(E, N, K, group_size, dtype, prefill_total_tokens) + ntpe = torch.tensor(prefill_shape["tokens_per_expert"], dtype=torch.int32, device="xpu") + + out_auto = ark.moe( + activations, + packed, + ntpe, + scales=scales, + weight_bits=4, + group_size=group_size, + asym=False, + phase="auto", + decode_threshold=prefill_total_tokens, + ) + out_decode = ark.moe_gemm_decode( + activations, + packed, + ntpe, + scales=scales, + weight_bits=4, + group_size=group_size, + asym=False, + ) + torch.testing.assert_close(out_auto, out_decode, rtol=0, atol=0) + + def test_decode_threshold_env_override(self, monkeypatch): + shape = _AUTO_PREFILL_SHAPE total_tokens = sum(shape["tokens_per_expert"]) E, N, K = shape["num_experts"], shape["N"], shape["K"] group_size = 128 @@ -239,7 +297,7 @@ def test_decode_threshold_override(self): activations, packed, scales, _ = _make_int4_sym(E, N, K, group_size, dtype, total_tokens) ntpe = torch.tensor(shape["tokens_per_expert"], dtype=torch.int32, device="xpu") - max_tpe = max(shape["tokens_per_expert"]) + monkeypatch.setenv("ARK_MOE_AUTO_DECODE_MAX_TOKENS", "512") out_auto = ark.moe( activations, packed, @@ -249,7 +307,6 @@ def test_decode_threshold_override(self): group_size=group_size, asym=False, phase="auto", - decode_threshold=max_tpe + 1, ) out_decode = ark.moe_gemm_decode( activations, @@ -262,8 +319,43 @@ def test_decode_threshold_override(self): ) torch.testing.assert_close(out_auto, out_decode, rtol=0, atol=0) + out_auto = ark.moe( + activations, + packed, + ntpe, + scales=scales, + weight_bits=4, + group_size=group_size, + asym=False, + phase="auto", + decode_threshold=128, + ) + out_prefill = ark.moe_gemm_prefill( + activations, + packed, + ntpe, + scales=scales, + weight_bits=4, + group_size=group_size, + asym=False, + ) + torch.testing.assert_close(out_auto, out_prefill, rtol=0, atol=0) + + monkeypatch.setenv("ARK_MOE_AUTO_DECODE_MAX_TOKENS", "invalid") + out_auto = ark.moe( + activations, + packed, + ntpe, + scales=scales, + weight_bits=4, + group_size=group_size, + asym=False, + phase="auto", + ) + torch.testing.assert_close(out_auto, out_prefill, rtol=0, atol=0) + def test_invalid_phase_raises(self): - shape = _DECODE_SHAPE + shape = _AUTO_DECODE_SHAPE total_tokens = sum(shape["tokens_per_expert"]) E, N, K = shape["num_experts"], shape["N"], shape["K"] group_size = 128 @@ -298,8 +390,8 @@ class TestMoeUnifiedBitParity: @pytest.mark.parametrize( "shape_name,shape", [ - ("decode-shape", _DECODE_SHAPE), - ("prefill-shape", _PREFILL_SHAPE), + ("decode-shape", _AUTO_DECODE_SHAPE), + ("prefill-shape", _AUTO_PREFILL_SHAPE), ], ) def test_fp_unquantized(self, dtype, shape_name, shape): @@ -324,8 +416,8 @@ def test_fp_unquantized(self, dtype, shape_name, shape): @pytest.mark.parametrize( "shape_name,shape", [ - ("decode-shape", _DECODE_SHAPE), - ("prefill-shape", _PREFILL_SHAPE), + ("decode-shape", _AUTO_DECODE_SHAPE), + ("prefill-shape", _AUTO_PREFILL_SHAPE), ], ) def test_int4(self, dtype, asym, shape_name, shape): @@ -353,7 +445,7 @@ def test_int4(self, dtype, asym, shape_name, shape): def test_int8(self, dtype, asym): # Single shape -- the quant path is the same on both shapes, so # iterating both would just slow the test suite down. - shape = _PREFILL_SHAPE + shape = _AUTO_PREFILL_SHAPE E, N, K = shape["num_experts"], shape["N"], shape["K"] total_tokens = sum(shape["tokens_per_expert"]) group_size = 128 @@ -375,7 +467,7 @@ def test_int8(self, dtype, asym): @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("asym", [False, True]) def test_int2(self, dtype, asym): - shape = _PREFILL_SHAPE + shape = _AUTO_PREFILL_SHAPE E, N, K = shape["num_experts"], shape["N"], shape["K"] total_tokens = sum(shape["tokens_per_expert"]) group_size = 128 @@ -397,7 +489,7 @@ def test_int2(self, dtype, asym): @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("fp8_dtype", [torch.float8_e4m3fn, torch.float8_e5m2]) def test_fp8(self, dtype, fp8_dtype): - shape = _PREFILL_SHAPE + shape = _AUTO_PREFILL_SHAPE E, N, K = shape["num_experts"], shape["N"], shape["K"] total_tokens = sum(shape["tokens_per_expert"]) group_size = 128