From e2b0265d7a87493fd5559ebb10f667d774a0c0f9 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 8 Jul 2026 06:33:08 +0000 Subject: [PATCH 01/17] feat: add bandwidth metrics to MoE perf scripts --- .../ark/test/test_moe_decode_perf.py | 103 ++++++++++++++++-- .../ark/test/test_moe_prefill_perf.py | 101 ++++++++++++++--- 2 files changed, 176 insertions(+), 28 deletions(-) diff --git a/auto_round_extension/ark/test/test_moe_decode_perf.py b/auto_round_extension/ark/test/test_moe_decode_perf.py index cf9b9a74c..d09ec6ba3 100644 --- a/auto_round_extension/ark/test/test_moe_decode_perf.py +++ b/auto_round_extension/ark/test/test_moe_decode_perf.py @@ -150,6 +150,25 @@ def _xpu_time_ms(fn, warmup: int = WARMUP, iters: int = ITERS) -> float: return timings[len(timings) // 2] +def _estimate_bandwidth_gbps(total_bytes: int, elapsed_ms: float) -> float: + if elapsed_ms <= 0: + return float("nan") + return total_bytes / (elapsed_ms * 1e-3) / 1e9 + + +def _estimate_moe_decode_bytes(activations, weights=None, scales=None, zeros=None, output_numel: int | None = None) -> int: + total_bytes = activations.numel() * activations.element_size() + if weights is not None: + total_bytes += weights.numel() * weights.element_size() + if scales is not None: + total_bytes += scales.numel() * scales.element_size() + if zeros is not None: + total_bytes += zeros.numel() * zeros.element_size() + if output_numel is not None: + total_bytes += output_numel * activations.element_size() + return total_bytes + + def _default_moe_decode(activations, dequant_weights, num_tokens_per_expert): """Default XPU MoE decode baseline: per-expert torch matmul loop. @@ -282,21 +301,29 @@ def _xpu_cleanup_between_tests(): def _print_header(title: str) -> None: print() - print("=" * 110) + print("=" * 130) print(title) - print(f"{'shape':<18}{'N':>7}{'K':>7}{'tokens':>8}" f"{'baseline(ms)':>16}{'ark(ms)':>14}{'speedup':>12}") - print("-" * 110) + print( + f"{'shape':<18}{'N':>7}{'K':>7}{'tokens':>8}" + f"{'baseline(ms)':>16}{'ark(ms)':>14}{'ark BW GB/s':>15}{'speedup':>12}" + ) + print("-" * 130) -def _print_row(label, N, K, total_tokens, base_ms, ark_ms): +def _print_row(label, N, K, total_tokens, base_ms, ark_ms, ark_bw=None): """Print a benchmark row. ``speedup`` is ``baseline / ark`` -- a pure matmul-vs-matmul comparison against the per-expert ``A @ W.T`` baseline running on already-dequantized - weights. + weights. ``ark_bw`` estimates the bandwidth of the timed fused-kernel + computation in GB/s. """ speedup = base_ms / ark_ms if ark_ms > 0 else float("nan") - print(f"{label:<18}{N:>7}{K:>7}{total_tokens:>8}" f"{base_ms:>16.4f}{ark_ms:>14.4f}{speedup:>11.2f}x") + if ark_bw is None: + ark_bw_col = f"{'--':>15}" + else: + ark_bw_col = f"{ark_bw:>14.2f} " + print(f"{label:<18}{N:>7}{K:>7}{total_tokens:>8}" f"{base_ms:>16.4f}{ark_ms:>14.4f}{ark_bw_col}{speedup:>11.2f}x") # --------------------------------------------------------------------------- @@ -326,7 +353,11 @@ def test_perf_fp(self, dtype): base_ms = _xpu_time_ms(lambda: _default_moe_decode(activations, weights, ntpe)) ark_ms = _xpu_time_ms(lambda: ark.moe_gemm_decode(activations, weights, ntpe, weight_bits=16)) - _print_row(label, N, K, total_tokens, base_ms, ark_ms) + ark_bw = _estimate_bandwidth_gbps( + _estimate_moe_decode_bytes(activations, weights=weights, output_numel=total_tokens * N), + ark_ms, + ) + _print_row(label, N, K, total_tokens, base_ms, ark_ms, ark_bw=ark_bw) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("asym", [False, True]) @@ -367,7 +398,17 @@ def test_perf_int4(self, dtype, asym): asym=asym, ) ) - _print_row(label, N, K, total_tokens, base_ms, ark_ms) + ark_bw = _estimate_bandwidth_gbps( + _estimate_moe_decode_bytes( + activations, + weights=packed, + scales=scales, + zeros=zeros, + output_numel=total_tokens * N, + ), + ark_ms, + ) + _print_row(label, N, K, total_tokens, base_ms, ark_ms, ark_bw=ark_bw) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("asym", [False, True]) @@ -408,7 +449,17 @@ def test_perf_int8(self, dtype, asym): asym=asym, ) ) - _print_row(label, N, K, total_tokens, base_ms, ark_ms) + ark_bw = _estimate_bandwidth_gbps( + _estimate_moe_decode_bytes( + activations, + weights=packed, + scales=scales, + zeros=zeros, + output_numel=total_tokens * N, + ), + ark_ms, + ) + _print_row(label, N, K, total_tokens, base_ms, ark_ms, ark_bw=ark_bw) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("asym", [False, True]) @@ -449,7 +500,17 @@ def test_perf_int2(self, dtype, asym): asym=asym, ) ) - _print_row(label, N, K, total_tokens, base_ms, ark_ms) + ark_bw = _estimate_bandwidth_gbps( + _estimate_moe_decode_bytes( + activations, + weights=packed, + scales=scales, + zeros=zeros, + output_numel=total_tokens * N, + ), + ark_ms, + ) + _print_row(label, N, K, total_tokens, base_ms, ark_ms, ark_bw=ark_bw) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("fp8_dtype", [torch.float8_e4m3fn, torch.float8_e5m2]) @@ -481,7 +542,16 @@ def test_perf_fp8(self, dtype, fp8_dtype): asym=False, ) ) - _print_row(label, N, K, total_tokens, base_ms, ark_ms) + ark_bw = _estimate_bandwidth_gbps( + _estimate_moe_decode_bytes( + activations, + weights=packed, + scales=scales, + output_numel=total_tokens * N, + ), + ark_ms, + ) + _print_row(label, N, K, total_tokens, base_ms, ark_ms, ark_bw=ark_bw) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("fp8_dtype", [torch.float8_e4m3fn, torch.float8_e5m2]) @@ -552,7 +622,16 @@ def test_perf_fp8_per_tensor(self, dtype, fp8_dtype): asym=False, ) ) - _print_row(label, N, K, total_tokens, base_ms, ark_ms) + ark_bw = _estimate_bandwidth_gbps( + _estimate_moe_decode_bytes( + activations, + weights=packed, + scales=scales, + output_numel=total_tokens * N, + ), + ark_ms, + ) + _print_row(label, N, K, total_tokens, base_ms, ark_ms, ark_bw=ark_bw) if __name__ == "__main__": diff --git a/auto_round_extension/ark/test/test_moe_prefill_perf.py b/auto_round_extension/ark/test/test_moe_prefill_perf.py index c322cb66b..26ee06f6f 100644 --- a/auto_round_extension/ark/test/test_moe_prefill_perf.py +++ b/auto_round_extension/ark/test/test_moe_prefill_perf.py @@ -197,6 +197,21 @@ def _xpu_time_ms(fn, warmup: int = WARMUP, iters: int = ITERS) -> float: return timings[len(timings) // 2] +def _estimate_bandwidth_gbps(total_bytes: int, elapsed_ms: float) -> float: + if elapsed_ms <= 0: + return float("nan") + return total_bytes / (elapsed_ms * 1e-3) / 1e9 + + +def _estimate_dequant_bytes(packed, scales, zeros, dequant) -> int: + total_bytes = packed.numel() * packed.element_size() + total_bytes += scales.numel() * scales.element_size() + if zeros is not None: + total_bytes += zeros.numel() * zeros.element_size() + total_bytes += dequant.numel() * dequant.element_size() + return total_bytes + + def _build_bmm_pad_layout(activations, num_tokens_per_expert, num_experts): """Pack ``[total_tokens, K]`` activations into a ``[E, M_max, K]`` buffer. @@ -501,6 +516,8 @@ def _print_header(title: str) -> None: * ``ark(ms)``: default ARK path (dequant workspace + grouped GEMM, or fused-dequant if ``ARK_MOE_PREFILL_FUSED_FP8=1`` is set in the env for FP8 rows). + * ``deq BW GB/s``: estimated bandwidth for the timed dequant stage in + quantized rows (``--`` for FP rows with no dequant cost). * ``native(ms)`` / ``native TFLOPS``: FP8 rows only. ARK path with ``ARK_MOE_PREFILL_NATIVE_FP8=1`` — the fused scalar native-FP8 GEMM that skips the ``[E, K, N]`` bf16/fp16 workspace and folds the @@ -515,12 +532,12 @@ def _print_header(title: str) -> None: ``speedup``. Prints ``--`` when no ``dpas(ms)`` was measured. """ print() - width = 200 + width = 215 print("=" * width) print(title) print( f"{'shape':<22}{'E':>4}{'N':>7}{'K':>7}{'tokens':>8}" - f"{'baseline(ms)':>16}{'ark(ms)':>14}{'speedup':>12}{'TFLOPS':>10}" + f"{'baseline(ms)':>16}{'ark(ms)':>14}{'deq BW GB/s':>16}{'speedup':>12}{'TFLOPS':>10}" f"{'native(ms)':>14}{'native TFLOPS':>16}" f"{'dpas(ms)':>14}{'dpas TFLOPS':>16}{'dpas speedup':>14}" ) @@ -537,6 +554,7 @@ def _print_row( deq_ms, ark_ms, tflops, + deq_bw=None, native_ms=None, native_tflops=None, dpas_ms=None, @@ -558,6 +576,10 @@ def _print_row( """ del deq_ms # no longer displayed; kept in signature for caller compatibility speedup = base_ms / ark_ms if ark_ms > 0 else float("nan") + if deq_bw is None: + deq_bw_col = f"{'--':>16}" + else: + deq_bw_col = f"{deq_bw:>15.2f} " if native_ms is None: native_col = f"{'--':>14}" native_tflops_col = f"{'--':>16}" @@ -575,7 +597,7 @@ def _print_row( dpas_speedup_col = f"{dpas_speedup:>13.2f}x" print( f"{label:<22}{E:>4}{N:>7}{K:>7}{total_tokens:>8}" - f"{base_ms:>16.4f}{ark_ms:>14.4f}{speedup:>11.2f}x{tflops:>9.1f}" + f"{base_ms:>16.4f}{ark_ms:>14.4f}{deq_bw_col}{speedup:>11.2f}x{tflops:>9.1f}" f"{native_col}{native_tflops_col}" f"{dpas_col}{dpas_tflops_col}{dpas_speedup_col}" ) @@ -666,6 +688,10 @@ def test_perf_int4(self, dtype, asym): deq_ms = _xpu_time_ms(lambda: _dequant_int4_asym(packed, scales, zeros, group_size).to(dtype)) else: deq_ms = _xpu_time_ms(lambda: _dequant_int4_sym(packed, scales, group_size).to(dtype)) + deq_bw = _estimate_bandwidth_gbps( + _estimate_dequant_bytes(packed, scales, zeros, dequant), + deq_ms, + ) base_ms = _xpu_time_ms(lambda: _default_moe_prefill(act_padded, dequant)) # Default ARK path (dequant + GEMM). INT4-sym is DPAS-accelerated @@ -733,9 +759,20 @@ def test_perf_int4(self, dtype, asym): os.environ.pop("ARK_MOE_PREFILL_DPAS_INT8", None) else: os.environ["ARK_MOE_PREFILL_DPAS_INT8"] = prev_dpas_int8 - + _print_row( - label, E, N, K, total_tokens, base_ms, deq_ms, ark_ms, tflops, dpas_ms=dpas_ms, dpas_tflops=dpas_tflops + label, + E, + N, + K, + total_tokens, + base_ms, + deq_ms, + ark_ms, + tflops, + deq_bw=deq_bw, + dpas_ms=dpas_ms, + dpas_tflops=dpas_tflops, ) activations = ntpe = act_padded = w_float = scales = zeros = packed = dequant = None @@ -773,6 +810,10 @@ def test_perf_int8(self, dtype, asym): deq_ms = _xpu_time_ms(lambda: _dequant_int8_asym(packed, scales, zeros, group_size).to(dtype)) else: deq_ms = _xpu_time_ms(lambda: _dequant_int8_sym(packed, scales, group_size).to(dtype)) + deq_bw = _estimate_bandwidth_gbps( + _estimate_dequant_bytes(packed, scales, zeros, dequant), + deq_ms, + ) base_ms = _xpu_time_ms(lambda: _default_moe_prefill(act_padded, dequant)) # Default ARK path (dequant + GEMM). Force @@ -827,9 +868,20 @@ def test_perf_int8(self, dtype, asym): os.environ.pop("ARK_MOE_PREFILL_DPAS_INT8", None) else: os.environ["ARK_MOE_PREFILL_DPAS_INT8"] = prev_dpas - + _print_row( - label, E, N, K, total_tokens, base_ms, deq_ms, ark_ms, tflops, dpas_ms=dpas_ms, dpas_tflops=dpas_tflops + label, + E, + N, + K, + total_tokens, + base_ms, + deq_ms, + ark_ms, + tflops, + deq_bw=deq_bw, + dpas_ms=dpas_ms, + dpas_tflops=dpas_tflops, ) activations = ntpe = act_padded = w_float = scales = zeros = packed = dequant = None @@ -867,6 +919,10 @@ def test_perf_int2(self, dtype, asym): deq_ms = _xpu_time_ms(lambda: _dequant_int2_asym(packed, scales, zeros, group_size).to(dtype)) else: deq_ms = _xpu_time_ms(lambda: _dequant_int2_sym(packed, scales, group_size).to(dtype)) + deq_bw = _estimate_bandwidth_gbps( + _estimate_dequant_bytes(packed, scales, zeros, dequant), + deq_ms, + ) base_ms = _xpu_time_ms(lambda: _default_moe_prefill(act_padded, dequant)) ark_ms = _xpu_time_ms( lambda: ark.moe_gemm_prefill( @@ -884,7 +940,7 @@ def test_perf_int2(self, dtype, asym): flops = _compute_moe_flops(total_tokens, K, N, E) tflops = flops / (ark_ms * 1e-3) / 1e12 - _print_row(label, E, N, K, total_tokens, base_ms, deq_ms, ark_ms, tflops) + _print_row(label, E, N, K, total_tokens, base_ms, deq_ms, ark_ms, tflops, deq_bw=deq_bw) activations = ntpe = act_padded = w_float = scales = zeros = packed = dequant = None _release_xpu_memory() @@ -912,6 +968,10 @@ def test_perf_fp8(self, dtype, fp8_dtype): dequant = _dequant_fp8(packed, scales, group_size, dtype) deq_ms = _xpu_time_ms(lambda: _dequant_fp8(packed, scales, group_size, dtype)) + deq_bw = _estimate_bandwidth_gbps( + _estimate_dequant_bytes(packed, scales, None, dequant), + deq_ms, + ) base_ms = _xpu_time_ms(lambda: _default_moe_prefill(act_padded, dequant)) # Default ARK path (respects the caller's env, i.e. dequant or @@ -1003,10 +1063,11 @@ def test_perf_fp8(self, dtype, fp8_dtype): deq_ms, ark_ms, tflops, - native_ms, - native_tflops, - dpas_ms, - dpas_tflops, + deq_bw=deq_bw, + native_ms=native_ms, + native_tflops=native_tflops, + dpas_ms=dpas_ms, + dpas_tflops=dpas_tflops, ) activations = ntpe = act_padded = w_float = scales = packed = dequant = None @@ -1056,16 +1117,20 @@ def test_perf_fp8_per_tensor(self, dtype, fp8_dtype): amax = w_float.reshape(E, -1).abs().amax(dim=1).clamp_min(1e-8) scales = (amax / fp8_finfo_max).to(torch.float32) # [E] fp32 packed = (w_float / scales.reshape(E, 1, 1)).to(fp8_dtype) - + # Baseline dequant: cast fp8 -> fp32 -> apply per-tensor scale -> # cast to act dtype, then transpose to [E, N, K] which is what # `_default_moe_prefill` (single torch.bmm) consumes. def _do_dequant(): dequant_KN = packed.to(torch.float32) * scales.reshape(E, 1, 1) return dequant_KN.transpose(1, 2).contiguous().to(dtype) - + dequant_NK = _do_dequant() deq_ms = _xpu_time_ms(_do_dequant) + deq_bw = _estimate_bandwidth_gbps( + _estimate_dequant_bytes(packed, scales, None, dequant_NK), + deq_ms, + ) base_ms = _xpu_time_ms(lambda: _default_moe_prefill(act_padded, dequant_NK)) ark_ms = _xpu_time_ms( @@ -1081,7 +1146,7 @@ def _do_dequant(): flops = _compute_moe_flops(total_tokens, K, N, E) tflops = flops / (ark_ms * 1e-3) / 1e12 - _print_row(label, E, N, K, total_tokens, base_ms, deq_ms, ark_ms, tflops) + _print_row(label, E, N, K, total_tokens, base_ms, deq_ms, ark_ms, tflops, deq_bw=deq_bw) activations = ntpe = act_padded = w_float = scales = packed = dequant_NK = None _release_xpu_memory() @@ -1134,6 +1199,10 @@ def _do_dequant(): dequant_NK = _do_dequant() deq_ms = _xpu_time_ms(_do_dequant) + deq_bw = _estimate_bandwidth_gbps( + _estimate_dequant_bytes(packed, scales, None, dequant_NK), + deq_ms, + ) base_ms = _xpu_time_ms(lambda: _default_moe_prefill(act_padded, dequant_NK)) ark_ms = _xpu_time_ms( @@ -1149,7 +1218,7 @@ def _do_dequant(): flops = _compute_moe_flops(total_tokens, K, N, E) tflops = flops / (ark_ms * 1e-3) / 1e12 - _print_row(label, E, N, K, total_tokens, base_ms, deq_ms, ark_ms, tflops) + _print_row(label, E, N, K, total_tokens, base_ms, deq_ms, ark_ms, tflops, deq_bw=deq_bw) activations = ntpe = act_padded = w_float = scales = packed = dequant_NK = None _release_xpu_memory() From c7a32494ac05f7306f4b851214c30a768c267b5f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 8 Jul 2026 06:49:15 +0000 Subject: [PATCH 02/17] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- auto_round_extension/ark/test/test_moe_decode_perf.py | 4 +++- auto_round_extension/ark/test/test_moe_prefill_perf.py | 8 ++++---- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/auto_round_extension/ark/test/test_moe_decode_perf.py b/auto_round_extension/ark/test/test_moe_decode_perf.py index d09ec6ba3..b9051e445 100644 --- a/auto_round_extension/ark/test/test_moe_decode_perf.py +++ b/auto_round_extension/ark/test/test_moe_decode_perf.py @@ -156,7 +156,9 @@ def _estimate_bandwidth_gbps(total_bytes: int, elapsed_ms: float) -> float: return total_bytes / (elapsed_ms * 1e-3) / 1e9 -def _estimate_moe_decode_bytes(activations, weights=None, scales=None, zeros=None, output_numel: int | None = None) -> int: +def _estimate_moe_decode_bytes( + activations, weights=None, scales=None, zeros=None, output_numel: int | None = None +) -> int: total_bytes = activations.numel() * activations.element_size() if weights is not None: total_bytes += weights.numel() * weights.element_size() diff --git a/auto_round_extension/ark/test/test_moe_prefill_perf.py b/auto_round_extension/ark/test/test_moe_prefill_perf.py index 26ee06f6f..657965b2f 100644 --- a/auto_round_extension/ark/test/test_moe_prefill_perf.py +++ b/auto_round_extension/ark/test/test_moe_prefill_perf.py @@ -759,7 +759,7 @@ def test_perf_int4(self, dtype, asym): os.environ.pop("ARK_MOE_PREFILL_DPAS_INT8", None) else: os.environ["ARK_MOE_PREFILL_DPAS_INT8"] = prev_dpas_int8 - + _print_row( label, E, @@ -868,7 +868,7 @@ def test_perf_int8(self, dtype, asym): os.environ.pop("ARK_MOE_PREFILL_DPAS_INT8", None) else: os.environ["ARK_MOE_PREFILL_DPAS_INT8"] = prev_dpas - + _print_row( label, E, @@ -1117,14 +1117,14 @@ def test_perf_fp8_per_tensor(self, dtype, fp8_dtype): amax = w_float.reshape(E, -1).abs().amax(dim=1).clamp_min(1e-8) scales = (amax / fp8_finfo_max).to(torch.float32) # [E] fp32 packed = (w_float / scales.reshape(E, 1, 1)).to(fp8_dtype) - + # Baseline dequant: cast fp8 -> fp32 -> apply per-tensor scale -> # cast to act dtype, then transpose to [E, N, K] which is what # `_default_moe_prefill` (single torch.bmm) consumes. def _do_dequant(): dequant_KN = packed.to(torch.float32) * scales.reshape(E, 1, 1) return dequant_KN.transpose(1, 2).contiguous().to(dtype) - + dequant_NK = _do_dequant() deq_ms = _xpu_time_ms(_do_dequant) deq_bw = _estimate_bandwidth_gbps( From 4feac83b909407ab97fc323fb7bec4b58a9005b3 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 8 Jul 2026 07:13:13 +0000 Subject: [PATCH 03/17] Apply remaining changes --- .../ark/test/test_moe_decode_perf.py | 78 +++++++++++++++---- 1 file changed, 65 insertions(+), 13 deletions(-) diff --git a/auto_round_extension/ark/test/test_moe_decode_perf.py b/auto_round_extension/ark/test/test_moe_decode_perf.py index b9051e445..c5cc084a4 100644 --- a/auto_round_extension/ark/test/test_moe_decode_perf.py +++ b/auto_round_extension/ark/test/test_moe_decode_perf.py @@ -37,6 +37,8 @@ --all-shapes """ +import math + import auto_round_kernel import pytest import torch @@ -157,17 +159,57 @@ def _estimate_bandwidth_gbps(total_bytes: int, elapsed_ms: float) -> float: def _estimate_moe_decode_bytes( - activations, weights=None, scales=None, zeros=None, output_numel: int | None = None + activations, + weights=None, + scales=None, + zeros=None, + num_tokens_per_expert=None, + output_numel: int | None = None, ) -> int: - total_bytes = activations.numel() * activations.element_size() - if weights is not None: - total_bytes += weights.numel() * weights.element_size() - if scales is not None: - total_bytes += scales.numel() * scales.element_size() - if zeros is not None: - total_bytes += zeros.numel() * zeros.element_size() - if output_numel is not None: - total_bytes += output_numel * activations.element_size() + """Estimate bytes touched by the active-expert decode work. + + The fused decode kernel only processes experts that receive at least one + routed token. Counting the full weight tensor for every call overstates the + amount of data moved and produces unrealistically high bandwidth numbers. + """ + if num_tokens_per_expert is None: + total_bytes = activations.numel() * activations.element_size() + if weights is not None: + total_bytes += weights.numel() * weights.element_size() + if scales is not None: + total_bytes += scales.numel() * scales.element_size() + if zeros is not None: + total_bytes += zeros.numel() * zeros.element_size() + if output_numel is not None: + total_bytes += output_numel * activations.element_size() + return total_bytes + + if hasattr(num_tokens_per_expert, "tolist"): + num_tokens_per_expert = num_tokens_per_expert.tolist() + else: + num_tokens_per_expert = list(num_tokens_per_expert) + + total_bytes = 0 + activation_bytes_per_elem = activations.element_size() + output_bytes_per_elem = activations.element_size() + activation_cols = activations.shape[1] + output_cols = weights.shape[1] if weights is not None else None + if output_numel is not None and output_cols is None: + output_cols = output_numel // max(sum(num_tokens_per_expert), 1) + + for num_tokens in num_tokens_per_expert: + if num_tokens <= 0: + continue + total_bytes += num_tokens * activation_cols * activation_bytes_per_elem + if weights is not None: + total_bytes += math.prod(weights.shape[1:]) * weights.element_size() + if scales is not None: + total_bytes += math.prod(scales.shape[1:]) * scales.element_size() + if zeros is not None: + total_bytes += math.prod(zeros.shape[1:]) * zeros.element_size() + if output_cols is not None: + total_bytes += num_tokens * output_cols * output_bytes_per_elem + return total_bytes @@ -317,8 +359,8 @@ def _print_row(label, N, K, total_tokens, base_ms, ark_ms, ark_bw=None): ``speedup`` is ``baseline / ark`` -- a pure matmul-vs-matmul comparison against the per-expert ``A @ W.T`` baseline running on already-dequantized - weights. ``ark_bw`` estimates the bandwidth of the timed fused-kernel - computation in GB/s. + weights. ``ark_bw`` estimates the effective bandwidth of the active-expert + slices touched by the fused kernel in GB/s. """ speedup = base_ms / ark_ms if ark_ms > 0 else float("nan") if ark_bw is None: @@ -356,7 +398,12 @@ def test_perf_fp(self, dtype): base_ms = _xpu_time_ms(lambda: _default_moe_decode(activations, weights, ntpe)) ark_ms = _xpu_time_ms(lambda: ark.moe_gemm_decode(activations, weights, ntpe, weight_bits=16)) ark_bw = _estimate_bandwidth_gbps( - _estimate_moe_decode_bytes(activations, weights=weights, output_numel=total_tokens * N), + _estimate_moe_decode_bytes( + activations, + weights=weights, + num_tokens_per_expert=ntpe, + output_numel=total_tokens * N, + ), ark_ms, ) _print_row(label, N, K, total_tokens, base_ms, ark_ms, ark_bw=ark_bw) @@ -406,6 +453,7 @@ def test_perf_int4(self, dtype, asym): weights=packed, scales=scales, zeros=zeros, + num_tokens_per_expert=ntpe, output_numel=total_tokens * N, ), ark_ms, @@ -457,6 +505,7 @@ def test_perf_int8(self, dtype, asym): weights=packed, scales=scales, zeros=zeros, + num_tokens_per_expert=ntpe, output_numel=total_tokens * N, ), ark_ms, @@ -508,6 +557,7 @@ def test_perf_int2(self, dtype, asym): weights=packed, scales=scales, zeros=zeros, + num_tokens_per_expert=ntpe, output_numel=total_tokens * N, ), ark_ms, @@ -549,6 +599,7 @@ def test_perf_fp8(self, dtype, fp8_dtype): activations, weights=packed, scales=scales, + num_tokens_per_expert=ntpe, output_numel=total_tokens * N, ), ark_ms, @@ -629,6 +680,7 @@ def test_perf_fp8_per_tensor(self, dtype, fp8_dtype): activations, weights=packed, scales=scales, + num_tokens_per_expert=ntpe, output_numel=total_tokens * N, ), ark_ms, From 23a217bbc2427ef60041660fb6d578a670e3e401 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 8 Jul 2026 08:07:32 +0000 Subject: [PATCH 04/17] feat: surface DPAS dequant BW in MoE prefill perf output --- .../ark/test/test_moe_prefill_perf.py | 34 ++++++++++--------- 1 file changed, 18 insertions(+), 16 deletions(-) diff --git a/auto_round_extension/ark/test/test_moe_prefill_perf.py b/auto_round_extension/ark/test/test_moe_prefill_perf.py index 657965b2f..ed1e033a6 100644 --- a/auto_round_extension/ark/test/test_moe_prefill_perf.py +++ b/auto_round_extension/ark/test/test_moe_prefill_perf.py @@ -516,8 +516,9 @@ def _print_header(title: str) -> None: * ``ark(ms)``: default ARK path (dequant workspace + grouped GEMM, or fused-dequant if ``ARK_MOE_PREFILL_FUSED_FP8=1`` is set in the env for FP8 rows). - * ``deq BW GB/s``: estimated bandwidth for the timed dequant stage in - quantized rows (``--`` for FP rows with no dequant cost). + * ``dpas deq BW GB/s``: estimated bandwidth for the dequant stage that + feeds the DPAS path in quantized rows (``--`` for FP rows with no + dequant cost). * ``native(ms)`` / ``native TFLOPS``: FP8 rows only. ARK path with ``ARK_MOE_PREFILL_NATIVE_FP8=1`` — the fused scalar native-FP8 GEMM that skips the ``[E, K, N]`` bf16/fp16 workspace and folds the @@ -532,12 +533,12 @@ def _print_header(title: str) -> None: ``speedup``. Prints ``--`` when no ``dpas(ms)`` was measured. """ print() - width = 215 + width = 225 print("=" * width) print(title) print( f"{'shape':<22}{'E':>4}{'N':>7}{'K':>7}{'tokens':>8}" - f"{'baseline(ms)':>16}{'ark(ms)':>14}{'deq BW GB/s':>16}{'speedup':>12}{'TFLOPS':>10}" + f"{'baseline(ms)':>16}{'ark(ms)':>14}{'dpas deq BW GB/s':>20}{'speedup':>12}{'TFLOPS':>10}" f"{'native(ms)':>14}{'native TFLOPS':>16}" f"{'dpas(ms)':>14}{'dpas TFLOPS':>16}{'dpas speedup':>14}" ) @@ -554,7 +555,7 @@ def _print_row( deq_ms, ark_ms, tflops, - deq_bw=None, + dpas_deq_bw=None, native_ms=None, native_tflops=None, dpas_ms=None, @@ -565,7 +566,8 @@ def _print_row( ``speedup`` is ``baseline / ark`` -- the fused kernel's speedup over the matmul-only baseline (weights pre-dequantized). ``deq_ms`` is accepted for signature compatibility with existing callers but is - no longer displayed. + no longer displayed. ``dpas_deq_bw`` is printed in the dedicated DPAS + dequant-bandwidth column when available. ``native_ms`` / ``native_tflops`` are printed for FP8 rows where the native fused kernel was benchmarked, and left blank otherwise. @@ -576,10 +578,10 @@ def _print_row( """ del deq_ms # no longer displayed; kept in signature for caller compatibility speedup = base_ms / ark_ms if ark_ms > 0 else float("nan") - if deq_bw is None: - deq_bw_col = f"{'--':>16}" + if dpas_deq_bw is None: + dpas_deq_bw_col = f"{'--':>20}" else: - deq_bw_col = f"{deq_bw:>15.2f} " + dpas_deq_bw_col = f"{dpas_deq_bw:>19.2f} " if native_ms is None: native_col = f"{'--':>14}" native_tflops_col = f"{'--':>16}" @@ -597,7 +599,7 @@ def _print_row( dpas_speedup_col = f"{dpas_speedup:>13.2f}x" print( f"{label:<22}{E:>4}{N:>7}{K:>7}{total_tokens:>8}" - f"{base_ms:>16.4f}{ark_ms:>14.4f}{deq_bw_col}{speedup:>11.2f}x{tflops:>9.1f}" + f"{base_ms:>16.4f}{ark_ms:>14.4f}{dpas_deq_bw_col}{speedup:>11.2f}x{tflops:>9.1f}" f"{native_col}{native_tflops_col}" f"{dpas_col}{dpas_tflops_col}{dpas_speedup_col}" ) @@ -770,7 +772,7 @@ def test_perf_int4(self, dtype, asym): deq_ms, ark_ms, tflops, - deq_bw=deq_bw, + dpas_deq_bw=deq_bw, dpas_ms=dpas_ms, dpas_tflops=dpas_tflops, ) @@ -879,7 +881,7 @@ def test_perf_int8(self, dtype, asym): deq_ms, ark_ms, tflops, - deq_bw=deq_bw, + dpas_deq_bw=deq_bw, dpas_ms=dpas_ms, dpas_tflops=dpas_tflops, ) @@ -940,7 +942,7 @@ def test_perf_int2(self, dtype, asym): flops = _compute_moe_flops(total_tokens, K, N, E) tflops = flops / (ark_ms * 1e-3) / 1e12 - _print_row(label, E, N, K, total_tokens, base_ms, deq_ms, ark_ms, tflops, deq_bw=deq_bw) + _print_row(label, E, N, K, total_tokens, base_ms, deq_ms, ark_ms, tflops, dpas_deq_bw=deq_bw) activations = ntpe = act_padded = w_float = scales = zeros = packed = dequant = None _release_xpu_memory() @@ -1063,7 +1065,7 @@ def test_perf_fp8(self, dtype, fp8_dtype): deq_ms, ark_ms, tflops, - deq_bw=deq_bw, + dpas_deq_bw=deq_bw, native_ms=native_ms, native_tflops=native_tflops, dpas_ms=dpas_ms, @@ -1146,7 +1148,7 @@ def _do_dequant(): flops = _compute_moe_flops(total_tokens, K, N, E) tflops = flops / (ark_ms * 1e-3) / 1e12 - _print_row(label, E, N, K, total_tokens, base_ms, deq_ms, ark_ms, tflops, deq_bw=deq_bw) + _print_row(label, E, N, K, total_tokens, base_ms, deq_ms, ark_ms, tflops, dpas_deq_bw=deq_bw) activations = ntpe = act_padded = w_float = scales = packed = dequant_NK = None _release_xpu_memory() @@ -1218,7 +1220,7 @@ def _do_dequant(): flops = _compute_moe_flops(total_tokens, K, N, E) tflops = flops / (ark_ms * 1e-3) / 1e12 - _print_row(label, E, N, K, total_tokens, base_ms, deq_ms, ark_ms, tflops, deq_bw=deq_bw) + _print_row(label, E, N, K, total_tokens, base_ms, deq_ms, ark_ms, tflops, dpas_deq_bw=deq_bw) activations = ntpe = act_padded = w_float = scales = packed = dequant_NK = None _release_xpu_memory() From db84594d00c2219b8477d23070a0c0a68c506b15 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 8 Jul 2026 08:41:44 +0000 Subject: [PATCH 05/17] refactor: optimize int4/int2 dequant helpers --- auto_round_extension/ark/test/test_moe.py | 62 +++++++++++------------ 1 file changed, 30 insertions(+), 32 deletions(-) diff --git a/auto_round_extension/ark/test/test_moe.py b/auto_round_extension/ark/test/test_moe.py index 8bad20cf9..011ce8762 100644 --- a/auto_round_extension/ark/test/test_moe.py +++ b/auto_round_extension/ark/test/test_moe.py @@ -218,31 +218,44 @@ def _pack_int4_asym(w_float, scales, zeros, group_size): return packed +def _unpack_int4_values(packed, *, signed: bool): + packed = packed.to(torch.uint8) + low = (packed & 0x0F).to(torch.int16) + high = ((packed >> 4) & 0x0F).to(torch.int16) + if signed: + low = torch.where(low >= 8, low - 16, low) + high = torch.where(high >= 8, high - 16, high) + values = torch.stack((low, high), dim=-1) + return values.reshape(*packed.shape[:-1], packed.shape[-1] * 2) + + +def _unpack_int2_values(packed, *, signed: bool): + packed = packed.to(torch.uint8) + q0 = (packed & 0x03).to(torch.int16) + q1 = ((packed >> 2) & 0x03).to(torch.int16) + q2 = ((packed >> 4) & 0x03).to(torch.int16) + q3 = ((packed >> 6) & 0x03).to(torch.int16) + if signed: + q0 = torch.where(q0 >= 2, q0 - 4, q0) + q1 = torch.where(q1 >= 2, q1 - 4, q1) + q2 = torch.where(q2 >= 2, q2 - 4, q2) + q3 = torch.where(q3 >= 2, q3 - 4, q3) + values = torch.stack((q0, q1, q2, q3), dim=-1) + return values.reshape(*packed.shape[:-1], packed.shape[-1] * 4) + + def _dequant_int4_sym(packed, scales, group_size): """Inverse of _pack_int4_sym. Returns [E, N, K] in scales.dtype.""" E, N, K_half = packed.shape K = K_half * 2 - low = (packed & 0x0F).to(torch.int8) - high = ((packed >> 4) & 0x0F).to(torch.int8) - # Sign extend 4-bit -> 8-bit - low = torch.where(low >= 8, low - 16, low) - high = torch.where(high >= 8, high - 16, high) - q = torch.empty(E, N, K, dtype=torch.int8, device=packed.device) - q[..., 0::2] = low - q[..., 1::2] = high - q = q.reshape(E, N, K // group_size, group_size).to(scales.dtype) + q = _unpack_int4_values(packed, signed=True).reshape(E, N, K // group_size, group_size).to(scales.dtype) return (q * scales.unsqueeze(-1)).reshape(E, N, K) def _dequant_int4_asym(packed, scales, zeros, group_size): E, N, K_half = packed.shape K = K_half * 2 - low = (packed & 0x0F).to(torch.int32) - high = ((packed >> 4) & 0x0F).to(torch.int32) - q = torch.empty(E, N, K, dtype=torch.int32, device=packed.device) - q[..., 0::2] = low - q[..., 1::2] = high - q = q.reshape(E, N, K // group_size, group_size).to(scales.dtype) + q = _unpack_int4_values(packed, signed=False).reshape(E, N, K // group_size, group_size).to(scales.dtype) deq = (q - zeros.to(scales.dtype).unsqueeze(-1)) * scales.unsqueeze(-1) return deq.reshape(E, N, K) @@ -346,29 +359,14 @@ def _pack_int2_asym(w_float, scales, zeros, group_size): def _dequant_int2_sym(packed, scales, group_size): E, N, K_q = packed.shape K = K_q * 4 - p = packed.to(torch.int32) - fields = torch.empty(E, N, K, dtype=torch.int32, device=packed.device) - fields[..., 0::4] = p & 0x3 - fields[..., 1::4] = (p >> 2) & 0x3 - fields[..., 2::4] = (p >> 4) & 0x3 - fields[..., 3::4] = (p >> 6) & 0x3 - # Sign-extend 2-bit (>=2 means negative). - fields = torch.where(fields >= 2, fields - 4, fields).to(scales.dtype) - fields = fields.reshape(E, N, K // group_size, group_size) + fields = _unpack_int2_values(packed, signed=True).reshape(E, N, K // group_size, group_size).to(scales.dtype) return (fields * scales.unsqueeze(-1)).reshape(E, N, K) def _dequant_int2_asym(packed, scales, zeros, group_size): E, N, K_q = packed.shape K = K_q * 4 - p = packed.to(torch.int32) - fields = torch.empty(E, N, K, dtype=torch.int32, device=packed.device) - fields[..., 0::4] = p & 0x3 - fields[..., 1::4] = (p >> 2) & 0x3 - fields[..., 2::4] = (p >> 4) & 0x3 - fields[..., 3::4] = (p >> 6) & 0x3 - fields = fields.to(scales.dtype) - fields = fields.reshape(E, N, K // group_size, group_size) + fields = _unpack_int2_values(packed, signed=False).reshape(E, N, K // group_size, group_size).to(scales.dtype) deq = (fields - zeros.to(scales.dtype).unsqueeze(-1)) * scales.unsqueeze(-1) return deq.reshape(E, N, K) From 25c3b611337befed47c7d4bd994bc6d66f61d6b0 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 8 Jul 2026 08:43:06 +0000 Subject: [PATCH 06/17] Apply remaining changes --- auto_round_extension/ark/test/test_moe.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/auto_round_extension/ark/test/test_moe.py b/auto_round_extension/ark/test/test_moe.py index 011ce8762..3e4187460 100644 --- a/auto_round_extension/ark/test/test_moe.py +++ b/auto_round_extension/ark/test/test_moe.py @@ -248,14 +248,16 @@ def _dequant_int4_sym(packed, scales, group_size): """Inverse of _pack_int4_sym. Returns [E, N, K] in scales.dtype.""" E, N, K_half = packed.shape K = K_half * 2 - q = _unpack_int4_values(packed, signed=True).reshape(E, N, K // group_size, group_size).to(scales.dtype) + q = _unpack_int4_values(packed, signed=True) + q = q.reshape(E, N, K // group_size, group_size).to(scales.dtype) return (q * scales.unsqueeze(-1)).reshape(E, N, K) def _dequant_int4_asym(packed, scales, zeros, group_size): E, N, K_half = packed.shape K = K_half * 2 - q = _unpack_int4_values(packed, signed=False).reshape(E, N, K // group_size, group_size).to(scales.dtype) + q = _unpack_int4_values(packed, signed=False) + q = q.reshape(E, N, K // group_size, group_size).to(scales.dtype) deq = (q - zeros.to(scales.dtype).unsqueeze(-1)) * scales.unsqueeze(-1) return deq.reshape(E, N, K) @@ -359,14 +361,16 @@ def _pack_int2_asym(w_float, scales, zeros, group_size): def _dequant_int2_sym(packed, scales, group_size): E, N, K_q = packed.shape K = K_q * 4 - fields = _unpack_int2_values(packed, signed=True).reshape(E, N, K // group_size, group_size).to(scales.dtype) + fields = _unpack_int2_values(packed, signed=True) + fields = fields.reshape(E, N, K // group_size, group_size).to(scales.dtype) return (fields * scales.unsqueeze(-1)).reshape(E, N, K) def _dequant_int2_asym(packed, scales, zeros, group_size): E, N, K_q = packed.shape K = K_q * 4 - fields = _unpack_int2_values(packed, signed=False).reshape(E, N, K // group_size, group_size).to(scales.dtype) + fields = _unpack_int2_values(packed, signed=False) + fields = fields.reshape(E, N, K // group_size, group_size).to(scales.dtype) deq = (fields - zeros.to(scales.dtype).unsqueeze(-1)) * scales.unsqueeze(-1) return deq.reshape(E, N, K) From c03590318907efecf2c9fc2cf95219998ac1b3fb Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 8 Jul 2026 11:48:10 +0000 Subject: [PATCH 07/17] refactor: optimize triton dequant tiling --- .../triton/triton_utils/dequant.py | 61 +++++++------------ .../triton/triton_utils_zp/dequant.py | 60 +++++++----------- 2 files changed, 44 insertions(+), 77 deletions(-) diff --git a/auto_round_extension/triton/triton_utils/dequant.py b/auto_round_extension/triton/triton_utils/dequant.py index e1564daa3..f7ff344a5 100644 --- a/auto_round_extension/triton/triton_utils/dequant.py +++ b/auto_round_extension/triton/triton_utils/dequant.py @@ -51,7 +51,6 @@ def make_dequant_configs(block_sizes, num_warps): DEFAULT_DEQUANT_CONFIGS = make_dequant_configs([128, 256, 512, 1024], [4, 8]) -@triton.autotune(DEFAULT_DEQUANT_CONFIGS, key=["numels"]) @triton.jit def dequant_kernel_248( g_idx_ptr, @@ -64,57 +63,40 @@ def dequant_kernel_248( bits: tl.constexpr, outfeatures: tl.constexpr, num_groups: tl.constexpr, - X_BLOCK: tl.constexpr, + BLOCK_COLS: tl.constexpr, ): - # Block indexing - xoffset = tl.program_id(0) * X_BLOCK - x_index = xoffset + tl.arange(0, X_BLOCK) - xmask = x_index < numels - row_idx = x_index // outfeatures - col_idx = x_index % outfeatures + row_idx = tl.program_id(0) + col_block = tl.program_id(1) elements_per_feature: tl.constexpr = 32 // bits + col_offsets = col_block * BLOCK_COLS + tl.arange(0, BLOCK_COLS) + col_mask = col_offsets < outfeatures - # Load parameters - g_idx = tl.load(g_idx_ptr + (row_idx), None, eviction_policy="evict_last") - qweights = tl.load( - qweight_ptr + (col_idx + (outfeatures * (row_idx // elements_per_feature))), - None, - ) - - wf_weights = (row_idx % elements_per_feature) * bits - - wf_zeros = (col_idx % elements_per_feature) * bits - + g_idx = tl.load(g_idx_ptr + row_idx, None, eviction_policy="evict_last") tmp1 = g_idx + num_groups tmp2 = g_idx < 0 tl.device_assert(g_idx >= 0, "index out of bounds: 0 <= tmp0 < 0") - groups = tl.where(tmp2, tmp1, g_idx) # tmp3 are g_idx + groups = tl.where(tmp2, tmp1, g_idx) - scales = tl.load(scales_ptr + (col_idx + (outfeatures * groups)), None).to(tl.float32) + qweight_offsets = col_offsets + (outfeatures * (row_idx // elements_per_feature)) + qweights = tl.load(qweight_ptr + qweight_offsets, mask=col_mask, other=0) - # Unpack weights - weights = qweights >> wf_weights # bit shift qweight - - weights = weights & maxq + wf_weights = (row_idx % elements_per_feature) * bits + weights = (qweights >> wf_weights) & maxq - # Unpack zeros qzero_ncols: tl.constexpr = outfeatures // elements_per_feature - qzeros = tl.load( - qzeros_ptr + ((qzero_ncols * groups) + (col_idx // elements_per_feature)), - None, - eviction_policy="evict_last", - ) - zeros = qzeros >> wf_zeros - zeros = zeros & maxq + qzero_idx = (qzero_ncols * groups) + col_block + qzeros = tl.load(qzeros_ptr + qzero_idx, None, eviction_policy="evict_last") + wf_zeros = (col_offsets % elements_per_feature) * bits + zeros = (qzeros >> wf_zeros) & maxq + + scales = tl.load(scales_ptr + (col_offsets + (outfeatures * groups)), mask=col_mask, other=0).to(tl.float32) - # # ##Dequantize - # zeros = zeros + 1 weights = weights - zeros weights = weights.to(tl.float32) weights = scales * weights - tl.store(out_ptr + (x_index), weights, mask=xmask) + tl.store(out_ptr + (row_idx * outfeatures + col_offsets), weights, mask=col_mask) def dequant248_core(qweight, scales, qzeros, g_idx, bits, maxq=None, input_dtype=torch.float16): @@ -126,9 +108,9 @@ def dequant248_core(qweight, scales, qzeros, g_idx, bits, maxq=None, input_dtype infeatures = g_idx.shape[0] out = torch.empty((infeatures, outfeatures), device=qweight.device, dtype=input_dtype) - numels = out.numel() maxq = 2**bits - 1 if maxq is None else maxq - grid = lambda meta: (triton.cdiv(numels, meta["X_BLOCK"]),) # noqa: E731 + block_cols = 32 // bits + grid = (infeatures, triton.cdiv(outfeatures, block_cols)) dequant_kernel_248[grid]( g_idx, @@ -136,11 +118,12 @@ def dequant248_core(qweight, scales, qzeros, g_idx, bits, maxq=None, input_dtype qweight, qzeros, out, - numels, + out.numel(), maxq=maxq, bits=bits, outfeatures=outfeatures, num_groups=num_groups, + BLOCK_COLS=block_cols, ) return out diff --git a/auto_round_extension/triton/triton_utils_zp/dequant.py b/auto_round_extension/triton/triton_utils_zp/dequant.py index b4de5fb8f..dcbfbc674 100644 --- a/auto_round_extension/triton/triton_utils_zp/dequant.py +++ b/auto_round_extension/triton/triton_utils_zp/dequant.py @@ -51,7 +51,6 @@ def make_dequant_configs(block_sizes, num_warps): DEFAULT_DEQUANT_CONFIGS = make_dequant_configs([128, 256, 512, 1024], [4, 8]) -@triton.autotune(DEFAULT_DEQUANT_CONFIGS, key=["numels"]) @triton.jit def dequant_kernel_248( g_idx_ptr, @@ -64,57 +63,41 @@ def dequant_kernel_248( bits: tl.constexpr, outfeatures: tl.constexpr, num_groups: tl.constexpr, - X_BLOCK: tl.constexpr, + BLOCK_COLS: tl.constexpr, ): - # Block indexing - xoffset = tl.program_id(0) * X_BLOCK - x_index = xoffset + tl.arange(0, X_BLOCK) - xmask = x_index < numels - row_idx = x_index // outfeatures - col_idx = x_index % outfeatures + row_idx = tl.program_id(0) + col_block = tl.program_id(1) elements_per_feature: tl.constexpr = 32 // bits + col_offsets = col_block * BLOCK_COLS + tl.arange(0, BLOCK_COLS) + col_mask = col_offsets < outfeatures - # Load parameters - g_idx = tl.load(g_idx_ptr + (row_idx), None, eviction_policy="evict_last") - qweights = tl.load( - qweight_ptr + (col_idx + (outfeatures * (row_idx // elements_per_feature))), - None, - ) - - wf_weights = (row_idx % elements_per_feature) * bits - - wf_zeros = (col_idx % elements_per_feature) * bits - + g_idx = tl.load(g_idx_ptr + row_idx, None, eviction_policy="evict_last") tmp1 = g_idx + num_groups tmp2 = g_idx < 0 tl.device_assert(g_idx >= 0, "index out of bounds: 0 <= tmp0 < 0") - groups = tl.where(tmp2, tmp1, g_idx) # tmp3 are g_idx + groups = tl.where(tmp2, tmp1, g_idx) - scales = tl.load(scales_ptr + (col_idx + (outfeatures * groups)), None).to(tl.float32) + qweight_offsets = col_offsets + (outfeatures * (row_idx // elements_per_feature)) + qweights = tl.load(qweight_ptr + qweight_offsets, mask=col_mask, other=0) - # Unpack weights - weights = qweights >> wf_weights # bit shift qweight - - weights = weights & maxq + wf_weights = (row_idx % elements_per_feature) * bits + weights = (qweights >> wf_weights) & maxq - # Unpack zeros qzero_ncols: tl.constexpr = outfeatures // elements_per_feature - qzeros = tl.load( - qzeros_ptr + ((qzero_ncols * groups) + (col_idx // elements_per_feature)), - None, - eviction_policy="evict_last", - ) - zeros = qzeros >> wf_zeros - zeros = zeros & maxq + qzero_idx = (qzero_ncols * groups) + col_block + qzeros = tl.load(qzeros_ptr + qzero_idx, None, eviction_policy="evict_last") + wf_zeros = (col_offsets % elements_per_feature) * bits + zeros = (qzeros >> wf_zeros) & maxq + + scales = tl.load(scales_ptr + (col_offsets + (outfeatures * groups)), mask=col_mask, other=0).to(tl.float32) - ##Dequantize zeros = zeros + 1 weights = weights - zeros weights = weights.to(tl.float32) weights = scales * weights - tl.store(out_ptr + (x_index), weights, mask=xmask) + tl.store(out_ptr + (row_idx * outfeatures + col_offsets), weights, mask=col_mask) def dequant248_core(qweight, scales, qzeros, g_idx, bits, maxq=None, input_dtype=torch.float16): @@ -126,9 +109,9 @@ def dequant248_core(qweight, scales, qzeros, g_idx, bits, maxq=None, input_dtype infeatures = g_idx.shape[0] out = torch.empty((infeatures, outfeatures), device=qweight.device, dtype=input_dtype) - numels = out.numel() maxq = 2**bits - 1 if maxq is None else maxq - grid = lambda meta: (triton.cdiv(numels, meta["X_BLOCK"]),) # noqa: E731 + block_cols = 32 // bits + grid = (infeatures, triton.cdiv(outfeatures, block_cols)) dequant_kernel_248[grid]( g_idx, @@ -136,11 +119,12 @@ def dequant248_core(qweight, scales, qzeros, g_idx, bits, maxq=None, input_dtype qweight, qzeros, out, - numels, + out.numel(), maxq=maxq, bits=bits, outfeatures=outfeatures, num_groups=num_groups, + BLOCK_COLS=block_cols, ) return out From 1393aba56a0db40b6a44ab6a07e45df4af6c3a97 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 8 Jul 2026 11:51:52 +0000 Subject: [PATCH 08/17] fix: correct packed qzero loads in Triton dequant --- auto_round_extension/triton/triton_utils/dequant.py | 4 ++-- auto_round_extension/triton/triton_utils_zp/dequant.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/auto_round_extension/triton/triton_utils/dequant.py b/auto_round_extension/triton/triton_utils/dequant.py index f7ff344a5..c26845074 100644 --- a/auto_round_extension/triton/triton_utils/dequant.py +++ b/auto_round_extension/triton/triton_utils/dequant.py @@ -85,8 +85,8 @@ def dequant_kernel_248( weights = (qweights >> wf_weights) & maxq qzero_ncols: tl.constexpr = outfeatures // elements_per_feature - qzero_idx = (qzero_ncols * groups) + col_block - qzeros = tl.load(qzeros_ptr + qzero_idx, None, eviction_policy="evict_last") + qzero_idx = (qzero_ncols * groups) + (col_offsets // elements_per_feature) + qzeros = tl.load(qzeros_ptr + qzero_idx, mask=col_mask, other=0, eviction_policy="evict_last") wf_zeros = (col_offsets % elements_per_feature) * bits zeros = (qzeros >> wf_zeros) & maxq diff --git a/auto_round_extension/triton/triton_utils_zp/dequant.py b/auto_round_extension/triton/triton_utils_zp/dequant.py index dcbfbc674..c5bc260c7 100644 --- a/auto_round_extension/triton/triton_utils_zp/dequant.py +++ b/auto_round_extension/triton/triton_utils_zp/dequant.py @@ -85,8 +85,8 @@ def dequant_kernel_248( weights = (qweights >> wf_weights) & maxq qzero_ncols: tl.constexpr = outfeatures // elements_per_feature - qzero_idx = (qzero_ncols * groups) + col_block - qzeros = tl.load(qzeros_ptr + qzero_idx, None, eviction_policy="evict_last") + qzero_idx = (qzero_ncols * groups) + (col_offsets // elements_per_feature) + qzeros = tl.load(qzeros_ptr + qzero_idx, mask=col_mask, other=0, eviction_policy="evict_last") wf_zeros = (col_offsets % elements_per_feature) * bits zeros = (qzeros >> wf_zeros) & maxq From 939915f0afd94703eea03d5355d12a407d8d70e5 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 8 Jul 2026 14:21:05 +0000 Subject: [PATCH 09/17] revert: restore original Triton int4/int2 dequant kernels Per feedback, the Triton dequant kernels must not be modified. This reverts the tiling changes to triton_utils/dequant.py and triton_utils_zp/dequant.py back to their original implementation. The XPU-path dequant performance work stays in the ARK moe prefill test helpers (contiguous stack+reshape unpack instead of strided scatter writes) plus the added dequant-bandwidth metric. --- .../triton/triton_utils/dequant.py | 61 ++++++++++++------- .../triton/triton_utils_zp/dequant.py | 60 +++++++++++------- 2 files changed, 77 insertions(+), 44 deletions(-) diff --git a/auto_round_extension/triton/triton_utils/dequant.py b/auto_round_extension/triton/triton_utils/dequant.py index c26845074..e1564daa3 100644 --- a/auto_round_extension/triton/triton_utils/dequant.py +++ b/auto_round_extension/triton/triton_utils/dequant.py @@ -51,6 +51,7 @@ def make_dequant_configs(block_sizes, num_warps): DEFAULT_DEQUANT_CONFIGS = make_dequant_configs([128, 256, 512, 1024], [4, 8]) +@triton.autotune(DEFAULT_DEQUANT_CONFIGS, key=["numels"]) @triton.jit def dequant_kernel_248( g_idx_ptr, @@ -63,40 +64,57 @@ def dequant_kernel_248( bits: tl.constexpr, outfeatures: tl.constexpr, num_groups: tl.constexpr, - BLOCK_COLS: tl.constexpr, + X_BLOCK: tl.constexpr, ): - row_idx = tl.program_id(0) - col_block = tl.program_id(1) + # Block indexing + xoffset = tl.program_id(0) * X_BLOCK + x_index = xoffset + tl.arange(0, X_BLOCK) + xmask = x_index < numels + row_idx = x_index // outfeatures + col_idx = x_index % outfeatures elements_per_feature: tl.constexpr = 32 // bits - col_offsets = col_block * BLOCK_COLS + tl.arange(0, BLOCK_COLS) - col_mask = col_offsets < outfeatures - g_idx = tl.load(g_idx_ptr + row_idx, None, eviction_policy="evict_last") + # Load parameters + g_idx = tl.load(g_idx_ptr + (row_idx), None, eviction_policy="evict_last") + qweights = tl.load( + qweight_ptr + (col_idx + (outfeatures * (row_idx // elements_per_feature))), + None, + ) + + wf_weights = (row_idx % elements_per_feature) * bits + + wf_zeros = (col_idx % elements_per_feature) * bits + tmp1 = g_idx + num_groups tmp2 = g_idx < 0 tl.device_assert(g_idx >= 0, "index out of bounds: 0 <= tmp0 < 0") - groups = tl.where(tmp2, tmp1, g_idx) + groups = tl.where(tmp2, tmp1, g_idx) # tmp3 are g_idx - qweight_offsets = col_offsets + (outfeatures * (row_idx // elements_per_feature)) - qweights = tl.load(qweight_ptr + qweight_offsets, mask=col_mask, other=0) + scales = tl.load(scales_ptr + (col_idx + (outfeatures * groups)), None).to(tl.float32) - wf_weights = (row_idx % elements_per_feature) * bits - weights = (qweights >> wf_weights) & maxq + # Unpack weights + weights = qweights >> wf_weights # bit shift qweight - qzero_ncols: tl.constexpr = outfeatures // elements_per_feature - qzero_idx = (qzero_ncols * groups) + (col_offsets // elements_per_feature) - qzeros = tl.load(qzeros_ptr + qzero_idx, mask=col_mask, other=0, eviction_policy="evict_last") - wf_zeros = (col_offsets % elements_per_feature) * bits - zeros = (qzeros >> wf_zeros) & maxq + weights = weights & maxq - scales = tl.load(scales_ptr + (col_offsets + (outfeatures * groups)), mask=col_mask, other=0).to(tl.float32) + # Unpack zeros + qzero_ncols: tl.constexpr = outfeatures // elements_per_feature + qzeros = tl.load( + qzeros_ptr + ((qzero_ncols * groups) + (col_idx // elements_per_feature)), + None, + eviction_policy="evict_last", + ) + zeros = qzeros >> wf_zeros + zeros = zeros & maxq + # # ##Dequantize + # zeros = zeros + 1 weights = weights - zeros weights = weights.to(tl.float32) weights = scales * weights - tl.store(out_ptr + (row_idx * outfeatures + col_offsets), weights, mask=col_mask) + tl.store(out_ptr + (x_index), weights, mask=xmask) def dequant248_core(qweight, scales, qzeros, g_idx, bits, maxq=None, input_dtype=torch.float16): @@ -108,9 +126,9 @@ def dequant248_core(qweight, scales, qzeros, g_idx, bits, maxq=None, input_dtype infeatures = g_idx.shape[0] out = torch.empty((infeatures, outfeatures), device=qweight.device, dtype=input_dtype) + numels = out.numel() maxq = 2**bits - 1 if maxq is None else maxq - block_cols = 32 // bits - grid = (infeatures, triton.cdiv(outfeatures, block_cols)) + grid = lambda meta: (triton.cdiv(numels, meta["X_BLOCK"]),) # noqa: E731 dequant_kernel_248[grid]( g_idx, @@ -118,12 +136,11 @@ def dequant248_core(qweight, scales, qzeros, g_idx, bits, maxq=None, input_dtype qweight, qzeros, out, - out.numel(), + numels, maxq=maxq, bits=bits, outfeatures=outfeatures, num_groups=num_groups, - BLOCK_COLS=block_cols, ) return out diff --git a/auto_round_extension/triton/triton_utils_zp/dequant.py b/auto_round_extension/triton/triton_utils_zp/dequant.py index c5bc260c7..b4de5fb8f 100644 --- a/auto_round_extension/triton/triton_utils_zp/dequant.py +++ b/auto_round_extension/triton/triton_utils_zp/dequant.py @@ -51,6 +51,7 @@ def make_dequant_configs(block_sizes, num_warps): DEFAULT_DEQUANT_CONFIGS = make_dequant_configs([128, 256, 512, 1024], [4, 8]) +@triton.autotune(DEFAULT_DEQUANT_CONFIGS, key=["numels"]) @triton.jit def dequant_kernel_248( g_idx_ptr, @@ -63,41 +64,57 @@ def dequant_kernel_248( bits: tl.constexpr, outfeatures: tl.constexpr, num_groups: tl.constexpr, - BLOCK_COLS: tl.constexpr, + X_BLOCK: tl.constexpr, ): - row_idx = tl.program_id(0) - col_block = tl.program_id(1) + # Block indexing + xoffset = tl.program_id(0) * X_BLOCK + x_index = xoffset + tl.arange(0, X_BLOCK) + xmask = x_index < numels + row_idx = x_index // outfeatures + col_idx = x_index % outfeatures elements_per_feature: tl.constexpr = 32 // bits - col_offsets = col_block * BLOCK_COLS + tl.arange(0, BLOCK_COLS) - col_mask = col_offsets < outfeatures - g_idx = tl.load(g_idx_ptr + row_idx, None, eviction_policy="evict_last") + # Load parameters + g_idx = tl.load(g_idx_ptr + (row_idx), None, eviction_policy="evict_last") + qweights = tl.load( + qweight_ptr + (col_idx + (outfeatures * (row_idx // elements_per_feature))), + None, + ) + + wf_weights = (row_idx % elements_per_feature) * bits + + wf_zeros = (col_idx % elements_per_feature) * bits + tmp1 = g_idx + num_groups tmp2 = g_idx < 0 tl.device_assert(g_idx >= 0, "index out of bounds: 0 <= tmp0 < 0") - groups = tl.where(tmp2, tmp1, g_idx) + groups = tl.where(tmp2, tmp1, g_idx) # tmp3 are g_idx - qweight_offsets = col_offsets + (outfeatures * (row_idx // elements_per_feature)) - qweights = tl.load(qweight_ptr + qweight_offsets, mask=col_mask, other=0) + scales = tl.load(scales_ptr + (col_idx + (outfeatures * groups)), None).to(tl.float32) - wf_weights = (row_idx % elements_per_feature) * bits - weights = (qweights >> wf_weights) & maxq + # Unpack weights + weights = qweights >> wf_weights # bit shift qweight - qzero_ncols: tl.constexpr = outfeatures // elements_per_feature - qzero_idx = (qzero_ncols * groups) + (col_offsets // elements_per_feature) - qzeros = tl.load(qzeros_ptr + qzero_idx, mask=col_mask, other=0, eviction_policy="evict_last") - wf_zeros = (col_offsets % elements_per_feature) * bits - zeros = (qzeros >> wf_zeros) & maxq + weights = weights & maxq - scales = tl.load(scales_ptr + (col_offsets + (outfeatures * groups)), mask=col_mask, other=0).to(tl.float32) + # Unpack zeros + qzero_ncols: tl.constexpr = outfeatures // elements_per_feature + qzeros = tl.load( + qzeros_ptr + ((qzero_ncols * groups) + (col_idx // elements_per_feature)), + None, + eviction_policy="evict_last", + ) + zeros = qzeros >> wf_zeros + zeros = zeros & maxq + ##Dequantize zeros = zeros + 1 weights = weights - zeros weights = weights.to(tl.float32) weights = scales * weights - tl.store(out_ptr + (row_idx * outfeatures + col_offsets), weights, mask=col_mask) + tl.store(out_ptr + (x_index), weights, mask=xmask) def dequant248_core(qweight, scales, qzeros, g_idx, bits, maxq=None, input_dtype=torch.float16): @@ -109,9 +126,9 @@ def dequant248_core(qweight, scales, qzeros, g_idx, bits, maxq=None, input_dtype infeatures = g_idx.shape[0] out = torch.empty((infeatures, outfeatures), device=qweight.device, dtype=input_dtype) + numels = out.numel() maxq = 2**bits - 1 if maxq is None else maxq - block_cols = 32 // bits - grid = (infeatures, triton.cdiv(outfeatures, block_cols)) + grid = lambda meta: (triton.cdiv(numels, meta["X_BLOCK"]),) # noqa: E731 dequant_kernel_248[grid]( g_idx, @@ -119,12 +136,11 @@ def dequant248_core(qweight, scales, qzeros, g_idx, bits, maxq=None, input_dtype qweight, qzeros, out, - out.numel(), + numels, maxq=maxq, bits=bits, outfeatures=outfeatures, num_groups=num_groups, - BLOCK_COLS=block_cols, ) return out From 4312b8f8f7d288a1fc76322c0f970c839eb8bbe0 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 9 Jul 2026 00:51:56 +0000 Subject: [PATCH 10/17] feat: measure internal moe_gemm_prefill dequant for deq_bw Add a dequant-only entry point (`moe_gemm_prefill_dequant`) that runs only the internal `[E,N,K_packed] -> [E,K,N]` weight-dequant stage of `moe_gemm_prefill` (the `dequant_to_KN` kernel) without the Grouped GEMM, so its throughput can be benchmarked in isolation. - C++: `moe_gemm_prefill_dequant` free fn in sycl_tla_moe_mixed.hpp - C++: wrapper + pybind registration in ark.cpp - Python: `moe_gemm_prefill_dequant` wrapper returning [E,K,N] weights - Test: `deq_bw` for INT4/INT8/INT2/FP8 rows now measures the interface's internal dequant instead of the PyTorch reference dequant; rename column to `deq BW GB/s` and document the new semantics. --- .../ark/auto_round_kernel/__init__.py | 108 ++++++++++++++ .../ark/auto_round_kernel/ark.cpp | 16 ++ .../wrapper/include/sycl_tla_moe_mixed.hpp | 36 +++++ .../ark/test/test_moe_prefill_perf.py | 137 +++++++++++++----- 4 files changed, 264 insertions(+), 33 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/__init__.py b/auto_round_extension/ark/auto_round_kernel/__init__.py index 512fab6a5..9e11338f7 100644 --- a/auto_round_extension/ark/auto_round_kernel/__init__.py +++ b/auto_round_extension/ark/auto_round_kernel/__init__.py @@ -2234,6 +2234,114 @@ def moe_gemm_prefill( return outputs +def moe_gemm_prefill_dequant( + activations: torch.Tensor, + weights: torch.Tensor, + num_tokens_per_expert: torch.Tensor, + *, + scales: Optional[torch.Tensor] = None, + zeros: Optional[torch.Tensor] = None, + weight_bits: int = 4, + group_size: int = 128, + asym: bool = False, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Run only the internal weight-dequant stage of :func:`moe_gemm_prefill`. + + :func:`moe_gemm_prefill` fuses an on-device weight-dequant step + (``[E, N, K_packed]`` quantized weights -> ``[E, K, N]`` ``act_dtype`` + workspace) with the subsequent Grouped GEMM, so the dequant cost cannot + be observed from the fused call alone. This entry point launches *only* + that dequant kernel and returns the materialised ``[E, K, N]`` weights, + letting callers benchmark the interface's internal dequant throughput in + isolation. The dequant numerics are identical to the generic + (non-DPAS) path of :func:`moe_gemm_prefill`. + + Args mirror :func:`moe_gemm_prefill` (quantized paths only). ``activations`` + is used only to derive the target dtype/device and ``K``; its values are + not read. INT2/INT4/INT8 (sym/asym) and FP8 weights are supported. + + Args: + activations: ``[total_tokens, K]`` fp16/bf16 tensor. Only its dtype, + device, and ``K`` are used. + weights: quantized ``[E, N, K_packed]`` weights (same contract as + :func:`moe_gemm_prefill`). + num_tokens_per_expert: ``[E]`` int32. Experts with 0 tokens are + skipped by the kernel (their workspace rows are left untouched). + scales: ``[E, N, K // group_size]`` act-dtype scales (required). + zeros: ``[E, N, K // group_size]`` act-dtype zeros (required when + ``asym=True``). + weight_bits: 2, 4, or 8. FP8 is inferred from the weight dtype. + group_size: group along K (default 128). + asym: unsigned encoding with ``zeros`` when ``True``. + out: optional pre-allocated ``[E, K, N]`` act-dtype destination. When + provided it is written in place and returned (useful for timing + loops that must exclude allocation cost). When ``None`` a fresh + tensor is allocated. + + Returns: + ``[E, K, N]`` dequantized weights in the activations dtype. + """ + activations, weights, scales, zeros, num_tokens_per_expert, weight_dtype, total_tokens, N, K, num_experts = ( + _validate_moe_quant_args( + activations, + weights, + num_tokens_per_expert, + scales=scales, + zeros=zeros, + weight_bits=weight_bits, + group_size=group_size, + asym=asym, + api_name="moe_gemm_prefill_dequant", + ) + ) + + is_unquantized = (weight_bits == 16) and (weights.dtype == activations.dtype) + if is_unquantized: + raise ValueError("moe_gemm_prefill_dequant: only quantized weights have an internal dequant stage") + + lib = get_lib(activations) + stream = get_stream(activations) + + if out is None: + dequant_workspace = torch.empty((num_experts, K, N), device=activations.device, dtype=activations.dtype) + else: + if tuple(out.shape) != (num_experts, K, N): + raise ValueError(f"out shape {tuple(out.shape)} != expected {(num_experts, K, N)}") + if out.dtype != activations.dtype: + raise ValueError("out dtype must match activations dtype") + if not out.is_contiguous(): + raise ValueError("out must be contiguous") + dequant_workspace = out + + if not hasattr(lib, "moe_gemm_prefill_dequant"): + raise RuntimeError( + "moe_gemm_prefill_dequant: the C++ backend was built without the " + "`moe_gemm_prefill_dequant` symbol. Rebuild auto_round_extension with " + "sycl-tla support to enable the internal-dequant benchmark entry point." + ) + + scales_ptr = scales.data_ptr() if scales is not None else 0 + zeros_ptr = zeros.data_ptr() if zeros is not None else 0 + + lib.moe_gemm_prefill_dequant( + stream, + weights.data_ptr(), + scales_ptr, + zeros_ptr, + dequant_workspace.data_ptr(), + cvt_dtype(activations.dtype), + weight_dtype, + N, + K, + group_size, + num_tokens_per_expert.data_ptr(), + num_experts, + bool(asym), + ) + return dequant_workspace + + # --------------------------------------------------------------------------- # `moe_gemm_prefill` dequant-workspace cache. # diff --git a/auto_round_extension/ark/auto_round_kernel/ark.cpp b/auto_round_extension/ark/auto_round_kernel/ark.cpp index 279083b26..71cb39d98 100755 --- a/auto_round_extension/ark/auto_round_kernel/ark.cpp +++ b/auto_round_extension/ark/auto_round_kernel/ark.cpp @@ -432,6 +432,21 @@ static void moe_gemm_prefill_wrapper(torch_ptr stream, torch_ptr activations, to total_tokens, asym); } +// Weight-dequant-only sibling of `moe_gemm_prefill`: runs the internal +// `[E, N, K_packed] -> [E, K, N]` weight-dequant stage into +// `dequant_workspace` without the Grouped GEMM, so callers can benchmark +// the interface's internal dequant throughput in isolation. +static void moe_gemm_prefill_dequant_wrapper(torch_ptr stream, torch_ptr weights, torch_ptr scales, torch_ptr zeros, + torch_ptr dequant_workspace, int act_dtype, int weight_dtype, int N, int K, + int group_size, torch_ptr num_tokens_per_expert, int num_experts, + bool asym) { + ark::moe_gemm_prefill_dequant((sycl::queue*)stream, (void*)weights, scales ? (void*)scales : nullptr, + zeros ? (void*)zeros : nullptr, + dequant_workspace ? (void*)dequant_workspace : nullptr, (BTLA_DTYPE)(act_dtype), + (BTLA_DTYPE)(weight_dtype), N, K, group_size, (int*)num_tokens_per_expert, num_experts, + asym); +} + // Variant A: FP8 per-tensor DPAS grouped GEMM (mirrors vllm-xpu-kernels' // `cutlass_grouped_gemm_xe2_impl` FP8 branch). `scales` is [E] FP32. // Weights are [E, K, N] row-major uint8. STATUS: NEEDS-HARDWARE-VALIDATION. @@ -726,6 +741,7 @@ PYBIND11_MODULE(PY_NAME, m) { m.def("moe_gemm", &ark::moe_gemm_wrapper); m.def("moe_gemm_decode", &ark::moe_gemm_decode_wrapper); m.def("moe_gemm_prefill", &ark::moe_gemm_prefill_wrapper); + m.def("moe_gemm_prefill_dequant", &ark::moe_gemm_prefill_dequant_wrapper); m.def("moe_gemm_prefill_fp8_dpas", &ark::moe_gemm_prefill_fp8_dpas_wrapper); m.def("moe_gemm_prefill_int_dpas", &ark::moe_gemm_prefill_int_dpas_wrapper); m.def("matmul_sycl_tla", &ark::matmul_sycl_tla); diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_mixed.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_mixed.hpp index e6a3ececc..c936e6eb1 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_mixed.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_mixed.hpp @@ -1095,6 +1095,42 @@ inline void moe_gemm_prefill(sycl::queue* q, void* activations, void* weights, v } } +// ---------------------------------------------------------------------------- +// Weight-dequant-only entry point for `moe_gemm_prefill`. +// +// Runs *only* the internal weight-dequantization stage that +// `moe_gemm_prefill` performs on its generic (non-DPAS) path -- i.e. it +// materialises the `[E, K, N]` act-dtype weights into `dequant_workspace` +// via `dequant_to_KN` and returns without launching the subsequent Grouped +// GEMM. This exists so callers can benchmark the interface's internal +// dequant throughput in isolation (the fused `moe_gemm_prefill` call cannot +// expose the dequant time on its own because the GEMM immediately consumes +// the workspace). The dequant numerics are bit-identical to the generic +// path of `moe_gemm_prefill`. +// +// Weight/scale/zero layouts match `moe_gemm_prefill` (decode-style +// `[E, N, K_packed]` weights, `[E, N, K/group_size]` scales/zeros). The +// destination `dequant_workspace` must be an `[E, K, N]` act-dtype buffer. +// ---------------------------------------------------------------------------- +inline void moe_gemm_prefill_dequant(sycl::queue* q, void* weights, void* scales, void* zeros, + void* dequant_workspace, BTLA_DTYPE act_dtype, BTLA_DTYPE weight_dtype, int N, + int K, int group_size, int* num_tokens_per_expert, int num_experts, bool asym) { + if (dequant_workspace == nullptr) { + throw std::invalid_argument("moe_gemm_prefill_dequant: dequant_workspace must be non-null"); + } + if (act_dtype == BTLA_DTYPE::F16) { + moe_mixed_detail::dequant_to_KN(q, weights, scales, zeros, static_cast(dequant_workspace), + weight_dtype, num_experts, N, K, group_size, asym, + num_tokens_per_expert); + } else if (act_dtype == BTLA_DTYPE::BF16) { + using BF = sycl::ext::oneapi::bfloat16; + moe_mixed_detail::dequant_to_KN(q, weights, scales, zeros, static_cast(dequant_workspace), weight_dtype, + num_experts, N, K, group_size, asym, num_tokens_per_expert); + } else { + throw std::invalid_argument("moe_gemm_prefill_dequant: act_dtype must be F16 or BF16"); + } +} + // ---------------------------------------------------------------------------- // MoE prefill Grouped GEMM -- FP8 per-tensor DPAS (Variant A). // diff --git a/auto_round_extension/ark/test/test_moe_prefill_perf.py b/auto_round_extension/ark/test/test_moe_prefill_perf.py index ed1e033a6..0a05bd09b 100644 --- a/auto_round_extension/ark/test/test_moe_prefill_perf.py +++ b/auto_round_extension/ark/test/test_moe_prefill_perf.py @@ -212,6 +212,51 @@ def _estimate_dequant_bytes(packed, scales, zeros, dequant) -> int: return total_bytes +def _internal_dequant_ms_bw( + activations, + packed, + num_tokens_per_expert, + *, + scales, + zeros, + weight_bits, + group_size, + asym, +): + """Time the weight-dequant stage *inside* ``ark.moe_gemm_prefill``. + + ``moe_gemm_prefill`` fuses an on-device weight-dequant step + (``[E, N, K_packed]`` quantized weights -> ``[E, K, N]`` act-dtype + workspace) with the subsequent Grouped GEMM, so the dequant cost cannot + be read off the fused call. This helper launches *only* that dequant + kernel via the dedicated ``ark.moe_gemm_prefill_dequant`` entry point + (which runs the same ``dequant_to_KN`` kernel the generic prefill path + uses, excluding the GEMM) and returns ``(deq_ms, deq_bw)`` where + ``deq_bw`` is the estimated dequant bandwidth in GB/s. + + A destination ``[E, K, N]`` buffer is pre-allocated once and reused across + timing iterations so the measured time excludes allocation overhead. + """ + E, N = packed.shape[0], packed.shape[1] + K = activations.shape[1] + out = torch.empty((E, K, N), device=activations.device, dtype=activations.dtype) + deq_ms = _xpu_time_ms( + lambda: ark.moe_gemm_prefill_dequant( + activations, + packed, + num_tokens_per_expert, + scales=scales, + zeros=zeros, + weight_bits=weight_bits, + group_size=group_size, + asym=asym, + out=out, + ) + ) + deq_bw = _estimate_bandwidth_gbps(_estimate_dequant_bytes(packed, scales, zeros, out), deq_ms) + return deq_ms, deq_bw + + def _build_bmm_pad_layout(activations, num_tokens_per_expert, num_experts): """Pack ``[total_tokens, K]`` activations into a ``[E, M_max, K]`` buffer. @@ -516,9 +561,15 @@ def _print_header(title: str) -> None: * ``ark(ms)``: default ARK path (dequant workspace + grouped GEMM, or fused-dequant if ``ARK_MOE_PREFILL_FUSED_FP8=1`` is set in the env for FP8 rows). - * ``dpas deq BW GB/s``: estimated bandwidth for the dequant stage that - feeds the DPAS path in quantized rows (``--`` for FP rows with no - dequant cost). + * ``deq BW GB/s``: estimated bandwidth of the weight-dequant stage + *inside* ``moe_gemm_prefill`` -- the on-device + ``[E, N, K_packed] -> [E, K, N]`` dequant kernel it runs before the + Grouped GEMM, measured in isolation via ``ark.moe_gemm_prefill_dequant`` + (per-group quantized rows: INT4/INT8/INT2/FP8). This is the internal + dequant throughput being optimized, *not* the PyTorch reference dequant + that feeds the ``baseline`` column. For the per-tensor DPAS schemes + (which fuse dequant into the GEMM and have no separable dequant stage) + the column reflects the reference-dequant estimate instead. * ``native(ms)`` / ``native TFLOPS``: FP8 rows only. ARK path with ``ARK_MOE_PREFILL_NATIVE_FP8=1`` — the fused scalar native-FP8 GEMM that skips the ``[E, K, N]`` bf16/fp16 workspace and folds the @@ -538,7 +589,7 @@ def _print_header(title: str) -> None: print(title) print( f"{'shape':<22}{'E':>4}{'N':>7}{'K':>7}{'tokens':>8}" - f"{'baseline(ms)':>16}{'ark(ms)':>14}{'dpas deq BW GB/s':>20}{'speedup':>12}{'TFLOPS':>10}" + f"{'baseline(ms)':>16}{'ark(ms)':>14}{'deq BW GB/s':>20}{'speedup':>12}{'TFLOPS':>10}" f"{'native(ms)':>14}{'native TFLOPS':>16}" f"{'dpas(ms)':>14}{'dpas TFLOPS':>16}{'dpas speedup':>14}" ) @@ -566,8 +617,9 @@ def _print_row( ``speedup`` is ``baseline / ark`` -- the fused kernel's speedup over the matmul-only baseline (weights pre-dequantized). ``deq_ms`` is accepted for signature compatibility with existing callers but is - no longer displayed. ``dpas_deq_bw`` is printed in the dedicated DPAS - dequant-bandwidth column when available. + no longer displayed. ``dpas_deq_bw`` is printed in the dedicated + ``deq BW GB/s`` column and, for the per-group quantized rows, holds the + bandwidth of the internal ``moe_gemm_prefill`` weight-dequant stage. ``native_ms`` / ``native_tflops`` are printed for FP8 rows where the native fused kernel was benchmarked, and left blank otherwise. @@ -684,15 +736,18 @@ def test_perf_int4(self, dtype, asym): dequant = _dequant_int4_sym(packed, scales, group_size).to(dtype) # ``dequant`` is already [E, N, K] -- matches the baseline contract. - # We still time dequant separately (measured but no longer reported) - # so callers of ``_print_row`` keep their existing signature. - if asym: - deq_ms = _xpu_time_ms(lambda: _dequant_int4_asym(packed, scales, zeros, group_size).to(dtype)) - else: - deq_ms = _xpu_time_ms(lambda: _dequant_int4_sym(packed, scales, group_size).to(dtype)) - deq_bw = _estimate_bandwidth_gbps( - _estimate_dequant_bytes(packed, scales, zeros, dequant), - deq_ms, + # ``deq_ms`` / ``deq_bw`` measure the weight-dequant stage *inside* + # ``moe_gemm_prefill`` (its internal ``[E, N, K_packed] -> [E, K, N]`` + # kernel), not the PyTorch reference dequant used for the baseline. + deq_ms, deq_bw = _internal_dequant_ms_bw( + activations, + packed, + ntpe, + scales=scales, + zeros=zeros, + weight_bits=4, + group_size=group_size, + asym=asym, ) base_ms = _xpu_time_ms(lambda: _default_moe_prefill(act_padded, dequant)) @@ -808,13 +863,17 @@ def test_perf_int8(self, dtype, asym): packed = _pack_int8_sym(w_float, scales, group_size) dequant = _dequant_int8_sym(packed, scales, group_size).to(dtype) - if asym: - deq_ms = _xpu_time_ms(lambda: _dequant_int8_asym(packed, scales, zeros, group_size).to(dtype)) - else: - deq_ms = _xpu_time_ms(lambda: _dequant_int8_sym(packed, scales, group_size).to(dtype)) - deq_bw = _estimate_bandwidth_gbps( - _estimate_dequant_bytes(packed, scales, zeros, dequant), - deq_ms, + # Internal ``moe_gemm_prefill`` weight-dequant stage (not the + # PyTorch reference dequant, which only feeds the bmm baseline). + deq_ms, deq_bw = _internal_dequant_ms_bw( + activations, + packed, + ntpe, + scales=scales, + zeros=zeros, + weight_bits=8, + group_size=group_size, + asym=asym, ) base_ms = _xpu_time_ms(lambda: _default_moe_prefill(act_padded, dequant)) @@ -917,13 +976,17 @@ def test_perf_int2(self, dtype, asym): packed = _pack_int2_sym(w_float, scales, group_size) dequant = _dequant_int2_sym(packed, scales, group_size).to(dtype) - if asym: - deq_ms = _xpu_time_ms(lambda: _dequant_int2_asym(packed, scales, zeros, group_size).to(dtype)) - else: - deq_ms = _xpu_time_ms(lambda: _dequant_int2_sym(packed, scales, group_size).to(dtype)) - deq_bw = _estimate_bandwidth_gbps( - _estimate_dequant_bytes(packed, scales, zeros, dequant), - deq_ms, + # Internal ``moe_gemm_prefill`` weight-dequant stage (not the + # PyTorch reference dequant, which only feeds the bmm baseline). + deq_ms, deq_bw = _internal_dequant_ms_bw( + activations, + packed, + ntpe, + scales=scales, + zeros=zeros, + weight_bits=2, + group_size=group_size, + asym=asym, ) base_ms = _xpu_time_ms(lambda: _default_moe_prefill(act_padded, dequant)) ark_ms = _xpu_time_ms( @@ -969,10 +1032,18 @@ def test_perf_fp8(self, dtype, fp8_dtype): packed = _pack_fp8(w_float, scales, group_size, fp8_dtype) dequant = _dequant_fp8(packed, scales, group_size, dtype) - deq_ms = _xpu_time_ms(lambda: _dequant_fp8(packed, scales, group_size, dtype)) - deq_bw = _estimate_bandwidth_gbps( - _estimate_dequant_bytes(packed, scales, None, dequant), - deq_ms, + # Internal ``moe_gemm_prefill`` FP8 weight-dequant stage (its + # ``[E, N, K] -> [E, K, N]`` dequant kernel), not the PyTorch + # reference dequant which only feeds the bmm baseline. + deq_ms, deq_bw = _internal_dequant_ms_bw( + activations, + packed, + ntpe, + scales=scales, + zeros=None, + weight_bits=8, + group_size=group_size, + asym=False, ) base_ms = _xpu_time_ms(lambda: _default_moe_prefill(act_padded, dequant)) From a126a83f17cdfbe2e6a250f4429ef6d28c924fcb Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 9 Jul 2026 02:23:30 +0000 Subject: [PATCH 11/17] feat: add single-pass INT2-sym fused DPAS MoE prefill path (S2) --- .../wrapper/include/sycl_tla_moe_mixed.hpp | 38 + .../include/sycl_tla_moe_prefill_s2_dpas.hpp | 676 ++++++++++++++++++ 2 files changed, 714 insertions(+) create mode 100644 auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s2_dpas.hpp diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_mixed.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_mixed.hpp index c936e6eb1..a22210cdf 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_mixed.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_mixed.hpp @@ -45,6 +45,7 @@ #include "sycl_tla_moe_prefill_fp8_dpas.hpp" #include "sycl_tla_moe_prefill_int_dpas.hpp" #include "sycl_tla_moe_prefill_s4_dpas.hpp" +#include "sycl_tla_moe_prefill_s2_dpas.hpp" #include "sycl_tla_moe_prefill_fp8_native.hpp" #include "sycl_tla_moe_prefill_fused.hpp" @@ -908,6 +909,43 @@ inline void moe_gemm_prefill(sycl::queue* q, void* activations, void* weights, v return; } + // S2-sym single-pass DPAS grouped GEMM (Variant B: per-K-group scale, + // in-register crumb->act upcast, XMX MMA). Reads packed `[E, N, K/4]` + // uint8_t crumbs directly and folds the upcast into the DPAS mainloop + // via CuTe's `reorder(tBrB, tCrB)`, so the B-side global traffic is a + // quarter of the S2->S8 upcast branch below. Opt-in default via + // `ARK_MOE_PREFILL_DPAS_S2` (default ON); silent fallback to the + // S2->S8 upcast branch (which is itself gated by + // `ARK_MOE_PREFILL_DPAS_INT8`) or to the generic dequant path if the + // shape gate rejects the tile geometry. + // + // STATUS: NEEDS-HARDWARE-VALIDATION. See + // `sycl_tla_moe_prefill_s2_dpas.hpp` for the port's provenance & the + // on-hardware TODOs (chief among them: `NumericArrayConverter + // ` availability in the pinned + // cutlass-sycl). + if (weight_dtype == BTLA_DTYPE::S2_CLIP && !asym && + moe_dpas_s2::moe_prefill_dpas_s2_enabled() && + moe_dpas_s2::moe_prefill_dpas_s2_pergroup_shape_ok(N, K, group_size) && + (act_dtype == BTLA_DTYPE::F16 || act_dtype == BTLA_DTYPE::BF16)) { + if (act_dtype == BTLA_DTYPE::F16) { + using ScalarT = sycl::half; + moe_dpas_s2::moe_prefill_s2_dpas_per_group_dispatch( + q, static_cast(activations), + static_cast(weights), + static_cast(scales), static_cast(outputs), + num_tokens_per_expert, num_experts, N, K, group_size, total_tokens); + } else { + using ScalarT = sycl::ext::oneapi::bfloat16; + moe_dpas_s2::moe_prefill_s2_dpas_per_group_dispatch( + q, static_cast(activations), + static_cast(weights), + static_cast(scales), static_cast(outputs), + num_tokens_per_expert, num_experts, N, K, group_size, total_tokens); + } + return; + } + // INT4-sym / INT2-sym via INT8 DPAS. Rather than dequantise packed // nibbles/crumbs into a bf16/fp16 `[E, K, N]` workspace and then run a // bf16 x bf16 GEMM, we upcast the low-bit-width weights to `int8_t` in diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s2_dpas.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s2_dpas.hpp new file mode 100644 index 000000000..9ad009d5f --- /dev/null +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s2_dpas.hpp @@ -0,0 +1,676 @@ +// SYCL MoE Prefill — S2 (sym) mixed-input DPAS Grouped GEMM (Variant B) +// (Fork of `sycl_tla_moe_prefill_s4_dpas.hpp` per-group mainloop) +// +// STATUS: NEEDS-HARDWARE-VALIDATION -- untested single-pass port. Precedence +// and gating in `sycl_tla_moe_mixed.hpp` fall back to the S2->S8 upcast + +// INT8 DPAS path when this header's env gate is off, so a regression can +// be neutralised at runtime without a rebuild. +// --------------------------------------------------------------------------- +// Design rationale +// ---------------- +// The prior S2-sym prefill path materialised `[E, N, K]` `int8_t` upcast +// bytes into the `dequant_workspace` (`launch_upcast_int2_sym_to_int8`) +// and then handed the buffer to `moe_prefill_int_dpas_per_group_dispatch`. +// The upcast pass writes `E * N * K` bytes and the mainloop then reads +// them back through L2 -- essentially quadrupling the B-side global-memory +// traffic vs. a direct packed-crumb read (a 2-bit weight is one quarter of +// an int8 byte). This header removes the round-trip: the mainloop reads +// packed `[E, N, K/4]` `uint8_t` (four 2-bit fields per byte, sym-signed +// [-2, 1]) and upcasts to the activation dtype in registers via the same +// `cute::reorder(tBrB, tCrB)` machinery the S4 header uses -- the only +// substantive difference is that CuTe/cutlass-sycl's +// `NumericArrayConverter` unpacks 2-bit +// fields four-per-byte from the loaded fragment. The B-side global load +// is a quarter of the bytes (`E * N * K / 4`) so the mainloop is +// bandwidth-bound on a quarter of the INT8-upcast traffic. +// +// Numerical parity +// ---------------- +// The auto-round S2_CLIP encoding packs four 2-bit sym-signed fields per +// byte: field j (0..3) at K index 4*i+j occupies bits [2j+1 : 2j] and is +// sign-extended from [-2, 1]. This is bit-identical to `cutlass::int2b_t` +// (== `integer_subbyte<2, true>`) storage, so reinterpret-casting the +// packed `uint8_t` pointer to `cutlass::int2b_t*` and letting the +// `NumericArrayConverter` sign-extend each field reproduces exactly the +// values that `moe_dequant::decode_int2_quad` computes on the +// decode / dequant paths. Asym S2 is intentionally NOT supported here (see +// the S4 header's asym rationale) and falls through to the generic dequant +// + `moe_gemm` path in `sycl_tla_moe_mixed.hpp`. +// +// Path taxonomy +// ------------- +// This header exposes ONE fused S2 grouped-GEMM entry point: +// +// * Variant B (s2, per-K-group, sym) -- weight scales `[E, N, K/GS]` +// in the activation dtype (half / bfloat16). Weights `[E, N, K/4]` +// packed uint8_t (4 crumbs per byte, sym-signed [-2, 1]). The +// launcher passes `LayoutKindB='C'` so `MoEGEMM_s2<>` XOR-flips to +// `'R'` inside `make_moe_tensor`, matching the physical `[N, K/4]` +// row-major storage. +// +// Variant A (per-tensor) is intentionally not implemented -- auto-round +// never ships a per-tensor S2 scale. +// +// On-hardware open questions (must be resolved on first build): +// 1. Whether the pinned cutlass-sycl exposes a +// `NumericArrayConverter` +// specialisation. Upstream cutlass ships `int2b_t` under +// `cutlass/integer_subbyte.h`; cutlass-sycl may need the same +// pull-in for the 2-bit converter. If the converter is missing, +// `reorder(tBrB, tCrB)` will fail to instantiate at compile time -- +// the failure mode is a "no matching function" error localised to +// the `reorder(tBrB, tCrB)` line in `xe_gemm_s2_pergroup`. In that +// case disable `ARK_MOE_PREFILL_DPAS_S2` (falls back to the S2->S8 +// upcast + INT8 DPAS path) until the converter is available. +// 2. Whether `XE_DPAS_TT<8, float, ElementA, ElementA>` (the atom used +// by the S4 / INT8 headers, keeping A/B as the SAME activation +// dtype after the upcast) is still the right atom here. It should +// be -- once `reorder` has upcast the packed-crumb fragment to +// `ElementA` the atom's operand dtypes match A exactly. +// 3. Whether CuTe's block-2d copy atom for a 2-bit storage type loads +// the packed bytes correctly (tile_k = 32 elements -> 8 bytes per +// SG per k_tile, deduced from `sizeof_bits::value == 2`). +// +// Copyright (C) 2026 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include +#include +#include +#include + +#ifdef ARK_XPU +#include +#endif + +#if defined(ARK_XPU) && defined(ARK_SYCL_TLA) + +// The S4 DPAS header pulls in the full cutlass-sycl / CuTe include set +// (via `sycl_tla_moe_prefill_int_dpas.hpp` -> `..._fp8_dpas.hpp`) and +// defines the policy classes, scalar-type mapping (`cute_scalar`), and +// per-group scale machinery we need. Include it first so we inherit all +// of that; the S2 driver lives in a sibling namespace to avoid ODR +// clashes. +#include "sycl_tla_moe_prefill_s4_dpas.hpp" + +// `cutlass::int2b_t` is the storage-narrow signed 2-bit type upstream +// cutlass / cutlass-sycl uses to trigger the packed-crumb copy / +// `NumericArrayConverter` code paths. Its `sizeof` is defined as 1 byte +// (one storage byte holds four elements); CuTe uses the type's +// `bits_per_element` trait, not `sizeof`, to derive copy/atom strides. +#include "cutlass/integer_subbyte.h" + +namespace ark { +namespace moe_dpas_s2 { + +using namespace cute; + +// Re-export the FP8/S4 policy classes / helpers into the S2 namespace so +// callsites read cleanly. The classes carry no per-dtype state. +using ::ark::moe_dpas_fp8::cute_scalar; +using ::ark::moe_dpas_fp8::cute_scalar_t; +using ::ark::moe_dpas_fp8::dpas_w8a16_policy; +using ::ark::moe_dpas_fp8::dpas_w8a16_policy_m_16; +using ::ark::moe_dpas_fp8::dpas_w8a16_policy_m_32; +using ::ark::moe_dpas_fp8::make_moe_tensor; + +// --------------------------------------------------------------------------- +// Variant B -- per-K-group S2 (sym) mainloop. +// +// Structurally identical to `moe_dpas_s4::xe_gemm_s4_pergroup<>` except +// `ElementB` is required to be `cutlass::int2b_t`. The 2-bit-per-element +// storage triggers CuTe's packed-crumb copy atom and the +// `NumericArrayConverter` specialisation inside +// `reorder(tBrB, tCrB)`, which decodes each byte into four sign-extended +// `ElementA` (bf16/fp16) values in-register. The B fragment size is +// unchanged in element units -- CuTe deduces it from the MMA tile shape +// and the `bits_per_element` trait on `int2b_t`, so a `tile_k` of 32 +// loads 8 bytes per SG per k_tile. Per-group scale reload / fold cadence +// is bit-identical to the S4 per-group path. +// --------------------------------------------------------------------------- + +template +CUTE_DEVICE void xe_gemm_s2_pergroup( + ATensor const& A, // (M,K) -- ElementA (bf16/fp16) + BTensor const& B, // (N,K) -- cutlass::int2b_t (packed crumbs) + const ElementS* Scales, + const ElementBI* Bias, + DTensor& C, // (M,N) -- ElementA + Coord blk_coord, + TiledMMA const& mma) { + using TA = typename ATensor::element_type; + using TB = typename BTensor::element_type; + static_assert(std::is_same_v, + "xe_gemm_s2_pergroup: ElementB must be cutlass::int2b_t (sym only)"); + static constexpr int group_size = GroupSize; + static constexpr int sg_local_range = 16; + auto item = sycl::ext::oneapi::this_work_item::get_nd_item<3>(); + auto wg_m = get<0>(blk_coord); + auto wg_n = get<1>(blk_coord); + int local_id = item.get_local_linear_id(); + + Tensor cA = make_identity_tensor(A.shape()); + Tensor cB = make_identity_tensor(B.shape()); + Tensor cC = make_identity_tensor(C.shape()); + + auto wg_tile = mma.tile_mnk(); + auto wg_coord = make_coord(wg_m, wg_n, 0); + + Tensor gA = local_tile(cA, select<0, 2>(wg_tile), make_coord(wg_m, _)); + Tensor gB = local_tile(cB, select<1, 2>(wg_tile), make_coord(wg_n, _)); + Tensor gC = local_tile(cC, wg_tile, wg_coord, Step<_1, _1, X>{}); + + auto copy_a = get_block_2d_copy_A(mma, A); + auto copy_b = get_block_2d_copy_B(mma, B); + auto copy_c = get_block_2d_copy_D(mma, C); + + auto thr_mma = mma.get_slice(local_id); + auto thr_copy_a = copy_a.get_slice(local_id); + auto thr_copy_b = copy_b.get_slice(local_id); + auto thr_copy_c = copy_c.get_slice(local_id); + + auto tCrA = thr_mma.partition_sg_fragment_A(gA(_, _, 0)); + auto tCrB = thr_mma.partition_sg_fragment_B(gB(_, _, 0)); + + auto tArA = thr_copy_a.partition_sg_fragment_D(gA(_, _, 0)); + auto tBrB = thr_copy_b.partition_sg_fragment_D(gB(_, _, 0)); + + Tensor tAgA = thr_copy_a.partition_S(gA); + Tensor tBgB = thr_copy_b.partition_S(gB); + + auto tCrC = thr_mma.partition_sg_fragment_C(gC); + auto tCrC_out = thr_copy_c.partition_sg_fragment_S(gC); + auto tCgC = thr_copy_c.partition_D(gC); + + auto prefetch_a = make_block_2d_prefetch(copy_a); + auto prefetch_b = make_block_2d_prefetch(copy_b); + + auto thr_prefetch_A = prefetch_a.get_slice(local_id); + auto thr_prefetch_B = prefetch_b.get_slice(local_id); + + auto pAgA = thr_prefetch_A.partition_S(gA); + auto pBgB = thr_prefetch_B.partition_S(gB); + + // Prefetch distance mirrors `xe_gemm_s4_pergroup<>` for now. On-hardware + // perf tuning may want to grow `prefetch_dist` further on the packed S2 + // path since the B stream is a quarter of the INT8-upcast bandwidth. + const int prefetch_dist = 3; + const int prefetch_dist_scale = 3; + constexpr auto barrier_scope = ScopeWorkgroup; + int k_tile_count = ceil_div(shape<1>(A), get<2>(wg_tile)); + int k_tile_prefetch = 0; + + static constexpr auto ATOM_M = get<1>(typename TiledMMA::ThrLayoutVMNK{}.shape()); + static constexpr auto ATOM_N = get<2>(typename TiledMMA::ThrLayoutVMNK{}.shape()); + + static constexpr auto tile_m = get<0>(wg_tile); + static constexpr auto tile_n = get<1>(wg_tile); + static constexpr auto tile_k = get<2>(wg_tile); + + static constexpr auto SG_M = tile_m / ATOM_M; + static constexpr auto SG_N = tile_n / ATOM_N; + + static constexpr int sg_n_strides = SG_N / sg_local_range; + + auto n_tile_start = wg_n * tile_n; + auto sg_local_n_coord = cutlass::get_sub_group_id() % ATOM_N; + auto sg_local_m_coord = cutlass::get_sub_group_id() / ATOM_N; + int sg_local_id = cutlass::get_sub_group_local_id(); + int n_sg_start = sg_local_n_coord * SG_N; + int m_sg_start = sg_local_m_coord * SG_M; + int m_tile_start = wg_m * tile_m; + int group_num = get<1>(A.shape()) / group_size; + + // Group-local accumulator: same fragment shape as `tCrC`, cleared at + // every scale-group boundary and folded into `tCrC` with a per-N-column + // scale before being reset. Mirrors the S4 per-group path exactly. + auto tCrC_group = thr_mma.partition_sg_fragment_C(gC); + + clear(tCrC); + clear(tCrC_group); + + // Per-SG per-N scale cache. Same layout / semantics as the S4 path. + float sg_scale[sg_n_strides]; + + CUTE_UNROLL + for (; k_tile_prefetch < prefetch_dist; k_tile_prefetch++) { + prefetch(prefetch_a, pAgA(_, _, _, k_tile_prefetch)); + prefetch(prefetch_b, pBgB(_, _, _, k_tile_prefetch)); + } + CUTLASS_PRAGMA_UNROLL + for (int pg = 0; pg < prefetch_dist_scale; ++pg) { + if (pg * group_size < shape<1>(A)) { + auto next_scales_tensor = make_tensor( + make_gmem_ptr(reinterpret_cast( + Scales + (n_tile_start + n_sg_start) * group_num + pg)), + make_layout(make_shape(Int{}, Int<1>{}), + make_stride(group_num, Int<1>{}))); + auto prefetch_scales = make_block_2d_prefetch<1>( + make_shape(Int{}, Int<1>{}), next_scales_tensor); + auto thr_prefetch_scales = prefetch_scales.get_slice(sg_local_id); + auto pSgS = thr_prefetch_scales.partition_S( + make_identity_tensor(make_shape(Int{}, Int<1>{}))); + prefetch(prefetch_scales, pSgS(_, 0, 0)); + } + } + + for (int k_tile = 0; k_tile < k_tile_count; k_tile++, k_tile_prefetch++) { + barrier_arrive(barrier_scope); + + copy(copy_a, tAgA(_, _, _, k_tile), tArA); + copy(copy_b, tBgB(_, _, _, k_tile), tBrB); + + // Group-boundary scale reload. Same math as the S4 per-group path: + // `tile_k` is expressed in element units (crumbs), not bytes, so + // `k_tile * tile_k` is the reduction position in *element* space and + // the modulo test against `group_size` matches the scale-tensor + // layout `[E, N, K/group_size]` unchanged. + if (k_tile * tile_k % group_size == 0) { + int group_idx = (k_tile * tile_k) / group_size; + CUTLASS_PRAGMA_UNROLL + for (int sn = 0; sn < sg_n_strides; ++sn) { + int sg_local_n = sn * sg_local_range + sg_local_id; + sg_scale[sn] = static_cast( + Scales[(n_tile_start + n_sg_start + sg_local_n) * group_num + group_idx]); + } + + if ((group_idx + prefetch_dist_scale) * group_size < shape<1>(A)) { + auto next_scales_tensor = make_tensor( + make_gmem_ptr(reinterpret_cast( + Scales + (n_tile_start + n_sg_start) * group_num + + group_idx + prefetch_dist_scale)), + make_layout(make_shape(Int{}, Int<1>{}), + make_stride(group_num, Int<1>{}))); + auto prefetch_scales = make_block_2d_prefetch<1>( + make_shape(Int{}, Int<1>{}), next_scales_tensor); + auto thr_prefetch_scales = prefetch_scales.get_slice(sg_local_id); + auto pSgS = thr_prefetch_scales.partition_S( + make_identity_tensor(make_shape(Int{}, Int<1>{}))); + prefetch(prefetch_scales, pSgS(_, 0, 0)); + } + } + + if (k_tile_prefetch < k_tile_count) { + prefetch(prefetch_a, pAgA(_, _, _, k_tile_prefetch)); + prefetch(prefetch_b, pBgB(_, _, _, k_tile_prefetch)); + } + + // `reorder` performs the in-register `int2b_t -> ElementA` unpack + // + sign-extend + cast via `cutlass::NumericArrayConverter< + // ElementA, cutlass::int2b_t, N>`. Once `tCrB` carries bf16/fp16 + // values it is compatible with the same DPAS atom used by the S4 / + // INT8 per-group paths. See the header preamble open-question (1) -- + // if the pinned cutlass-sycl is missing this converter specialisation + // this line is where the build fails. + reorder(tArA, tCrA); + reorder(tBrB, tCrB); + + // HOT MAINLOOP -- MMA accumulates into `tCrC_group`. Per-N scale is + // applied ONCE at the end of the group in the fold block below. + cute::gemm(mma, tCrA, tCrB, tCrC_group); + + // Group-boundary fold. Fires when either (a) the NEXT k_tile would + // start a new scale group, or (b) we've reached the last k_tile of + // the K reduction (tail-group protection). + const bool is_group_end = (((k_tile + 1) * tile_k) % group_size == 0) || + (k_tile + 1 == k_tile_count); + if (is_group_end) { + CUTLASS_PRAGMA_UNROLL + for (int sn = 0; sn < sg_n_strides; ++sn) { + float s = sg_scale[sn]; + CUTLASS_PRAGMA_UNROLL + for (int sm = 0; sm < SG_M; ++sm) { + const int idx = sn * SG_M + sm; + tCrC(idx) += tCrC_group(idx) * s; + tCrC_group(idx) = 0.0f; + } + } + } + + barrier_wait(barrier_scope); + } + + if (Bias != nullptr) { + CUTLASS_PRAGMA_UNROLL + for (int sn = 0; sn < sg_n_strides; ++sn) { + int sg_local_n = sn * sg_local_range + sg_local_id; + float b_float = Bias[n_tile_start + n_sg_start + sg_local_n]; + CUTLASS_PRAGMA_UNROLL + for (int sm = 0; sm < SG_M; ++sm) { + tCrC(sn * SG_M + sm) += b_float; + } + } + } + + reorder(tCrC, tCrC_out); + copy(copy_c, tCrC_out, tCgC); +} + +// --------------------------------------------------------------------------- +// Persistent scheduler (fork of `moe_dpas_s4::MoEGEMM_s4<>`). +// +// The only substantive difference vs. the S4 scheduler is the per-expert +// `B_offset` computation: `int2b_t` storage is 2 bits per element, so +// `expert_id * gemm_n * gemm_k` elements = `expert_id * gemm_n * gemm_k / +// 4` bytes. Pointer arithmetic on `ElementB*` (where +// `sizeof(cutlass::int2b_t) == 1` byte and each byte holds four elements) +// means adding `B_offset` in *element* units to the raw pointer would +// quadruple-count -- so we divide by 4 here. +// --------------------------------------------------------------------------- + +template +CUTE_DEVICE void MoEGEMM_s2(const ElementA* Activations, + const ElementB* Weights, + const ElementS* Scales, + const ElementBI* Bias, + ElementD* Outputs, TiledMMA const& mma, + const int* rows_per_expert, + const int32_t num_experts, + const int32_t group_size, const int32_t gemm_n, + const int32_t gemm_k, int32_t* atomic_buffer, + const sycl::local_accessor& slm_mem_const) { + constexpr char actual_layout_of_B = LayoutKindB ^ ('R' ^ 'C'); + + auto item = sycl::ext::oneapi::this_work_item::get_nd_item<3>(); + auto wg_tile = mma.tile_mnk(); + auto wg_tile_m = get<0>(wg_tile); + auto wg_tile_n = get<1>(wg_tile); + + int group_id = item.get_group_linear_id(); + int gemm_n_pad = (gemm_n + wg_tile_n - 1) / wg_tile_n * wg_tile_n; + int group_m_id = (group_id * wg_tile_n) / gemm_n_pad; + int group_range = item.get_group_range(1); + int local_id = item.get_local_linear_id(); + + if (group_id == 0 && local_id == 0) { + auto atm = sycl::atomic_ref( + atomic_buffer[0]); + atm.store(0); + } + + int pre_rows = 0; + int pre_tiles = 0; + + int32_t* slm_mem = static_cast( + slm_mem_const.template get_multi_ptr().get()); + + for (int i = 0; i < num_experts; ++i) { + int gemm_m = rows_per_expert[i]; + int cumsum_rows_for_experts = pre_rows + gemm_m; + int cumsum_tiles_for_experts = + (gemm_m + wg_tile_m - 1) / wg_tile_m + pre_tiles; + + if (group_m_id >= cumsum_tiles_for_experts) { + pre_rows = cumsum_rows_for_experts; + pre_tiles = cumsum_tiles_for_experts; + continue; + } + + int expert_id = i; + // 2-bit branch: `gemm_n * gemm_k` is the *element* count for one + // expert; the packed byte count is a quarter of that. `ElementB*` + // pointer arithmetic advances by whole bytes, so we quarter the + // offset here. `gemm_k` is guaranteed a multiple of 4 by the shape + // gate below. + int64_t B_offset = static_cast(expert_id) * + static_cast(gemm_n) * + static_cast(gemm_k) / 4; + + ElementA* ptr_A_curr_batch = + const_cast(Activations) + pre_rows * gemm_k; + ElementB* ptr_B_curr_batch = const_cast(Weights) + B_offset; + ElementD* ptr_D_curr_batch = Outputs + pre_rows * gemm_n; + + // Per-group scale offset in ElementS units. Layout `[E, N, K/GS]`. + ElementS* ptr_Scales_curr_batch = nullptr; + int64_t scale_expert_stride = + static_cast(gemm_n) * (gemm_k / group_size); + ptr_Scales_curr_batch = + const_cast(Scales) + expert_id * scale_expert_stride; + + ElementBI* ptr_Bias_curr_batch = nullptr; + if (Bias != static_cast(nullptr)) { + ptr_Bias_curr_batch = const_cast(Bias) + expert_id * gemm_n; + } + + auto A_tensor = make_moe_tensor(ptr_A_curr_batch, + gemm_m, gemm_k); + // B tensor extents are in *element* units (crumbs) -- CuTe deduces + // the byte-space stride from `sizeof_bits::value == 2`. + auto B_tensor = make_moe_tensor( + ptr_B_curr_batch, gemm_n, gemm_k); + auto D_tensor = make_moe_tensor(ptr_D_curr_batch, + gemm_m, gemm_n); + + while (group_m_id < cumsum_tiles_for_experts) { + int n_coord = (group_id * wg_tile_n) % gemm_n_pad / wg_tile_n; + int m_coord = (group_m_id - pre_tiles); + auto tile_coord = make_coord(m_coord, n_coord, _, 0); + +// Per-K-group dispatch on the runtime group_size, mirrors the S4 header's +// macro. Only per-group is supported for S2 (no per-tensor variant is +// instantiated here). +#define ARK_MOE_DPAS_S2_GROUP_CALLER(GS) \ + xe_gemm_s2_pergroup( \ + A_tensor, B_tensor, ptr_Scales_curr_batch, ptr_Bias_curr_batch, \ + D_tensor, tile_coord, mma); + if (group_size == 32) { + ARK_MOE_DPAS_S2_GROUP_CALLER(32) + } else if (group_size == 64) { + ARK_MOE_DPAS_S2_GROUP_CALLER(64) + } else if (group_size == 128) { + ARK_MOE_DPAS_S2_GROUP_CALLER(128) + } else if (group_size == 256) { + ARK_MOE_DPAS_S2_GROUP_CALLER(256) + } +#undef ARK_MOE_DPAS_S2_GROUP_CALLER + + if (local_id == 0) { + slm_mem[0] = cutlass::atomicAdd(atomic_buffer, 1); + } + item.barrier(sycl::access::fence_space::local_space); + group_id = group_range + slm_mem[0]; + group_m_id = (group_id * wg_tile_n) / gemm_n_pad; + } + pre_rows = cumsum_rows_for_experts; + pre_tiles = cumsum_tiles_for_experts; + } +} + +// --------------------------------------------------------------------------- +// Launcher (fork of `moe_dpas_s4::MoEGEMMLauncher_s4<>`). +// --------------------------------------------------------------------------- + +template +class DpasGemmS2Name; + +template +void MoEGEMMLauncher_s2(sycl::queue& stream, const ElementA* activations, + const ElementB* weights, const ElementS* scales, + const ElementBI* bias, ElementD* outputs, + const int gemm_n, const int gemm_k, + const int* rows_per_expert, const int num_experts, + const int group_size, int32_t* atomic_buffer) { + using ElementA_non_CV = cutlass::platform::remove_cv_t; + // DPAS atom keeps its bf16/fp16 x bf16/fp16 -> fp32 shape; the S2 B + // tensor is upcast to ElementA in `reorder(tBrB, tCrB)` in the mainloop + // before entering the MMA atom. See open question #2 in the header + // preamble. + auto op = XE_DPAS_TT<8, float, ElementA_non_CV, ElementA_non_CV>{}; + + using WGTile = typename policy::WGTile; + using SGLayout = typename policy::SGLayout; + using MMA = typename TiledMMAHelper, Layout, + SGLayout>::TiledMMA; + auto mma = MMA{}; + + int sm_count = + cutlass::KernelHardwareInfo::query_device_multiprocessor_count(0); + auto MaxThreadsPerWorkgroup = size(mma); + + static constexpr int MaxThreadsPerSM = 512; + if (MaxThreadsPerSM % MaxThreadsPerWorkgroup != 0) { + throw std::runtime_error( + "moe_prefill_s2_dpas: MaxThreadsPerSM must be divisible by " + "MaxThreadsPerWorkgroup"); + } + + sycl::range<3> local(1, 1, MaxThreadsPerWorkgroup); + sycl::range<3> global(1, sm_count * MaxThreadsPerSM / MaxThreadsPerWorkgroup, 1); + + namespace syclex = sycl::ext::oneapi::experimental; + namespace intelex = sycl::ext::intel::experimental; + + syclex::properties kernel_props{syclex::sub_group_size<16>, + intelex::grf_size<256>}; + + using GmemTiledCopyA = typename policy::GmemTiledCopyA; + using GmemTiledCopyB = typename policy::GmemTiledCopyB; + using GmemTiledCopyD = typename policy::GmemTiledCopyD; + + auto event = stream.submit([&](sycl::handler& cgh) { + sycl::local_accessor local_mem(sycl::range<1>(1), cgh); + cgh.parallel_for>( + sycl::nd_range<3>{global * local, local}, kernel_props, [=](auto) { + MoEGEMM_s2( + activations, weights, scales, bias, outputs, mma, + rows_per_expert, num_experts, group_size, gemm_n, gemm_k, + atomic_buffer, local_mem); + }); + }); + + EventManager::getInstance().addEvent(event); + event.wait(); +} + +// --------------------------------------------------------------------------- +// Host-side driver: S2 Variant B (per-K-group, sym only). +// +// Weight layout `[E, N, K/4]` row-major packed uint8_t -- four sym-signed +// 2-bit fields per byte, matching the auto-round S2_CLIP encoding. On +// entry the raw uint8_t pointer is reinterpret-cast to +// `cutlass::int2b_t*` so CuTe's copy atom / `NumericArrayConverter` can +// decode the packed crumbs directly. `LayoutKindB='C'` at the launcher +// level so `MoEGEMM_s2<>` XOR-flips to `'R'` inside `make_moe_tensor`, +// matching the physical `[N, K/4]` row-major storage. +// Scales `[E, N, K/group_size]` in act dtype (half / bfloat16). +// --------------------------------------------------------------------------- + +template +void moe_prefill_s2_dpas_per_group_dispatch( + sycl::queue* q, const ScalarT* activations, const uint8_t* weights_NKp, + const ScalarT* scales, ScalarT* outputs, + const int* num_tokens_per_expert, int E, int N, int K, int group_size, + int total_tokens) { + if (E == 0 || N == 0 || K == 0 || total_tokens == 0) return; + if (K % group_size != 0) { + throw std::invalid_argument( + "moe_prefill_s2_dpas(per-group): K must be a multiple of group_size"); + } + if ((K & 0x3) != 0) { + throw std::invalid_argument( + "moe_prefill_s2_dpas(per-group): K must be a multiple of 4 (packed crumbs)"); + } + + compat::set_default_queue(*q); + + // Map the caller-facing SYCL native half/bfloat16 to the CUTLASS type CUTE + // has DPAS-atom specializations for. See `cute_scalar` in the FP8 header. + using ElementA = cute_scalar_t; + const auto* activations_ca = + reinterpret_cast(activations); + const auto* scales_ca = reinterpret_cast(scales); + auto* outputs_ca = reinterpret_cast(outputs); + // `cutlass::int2b_t` is a 2-bit-per-element storage type whose sizeof + // is 1 (one byte holds four elements). CuTe reads the bit-width from + // `sizeof_bits::value == 2` when computing copy strides, so + // reinterpret-casting the packed uint8_t pointer is safe. + const auto* weights_i2 = + reinterpret_cast(weights_NKp); + + int A_avg_M = total_tokens / E; + + int32_t* atomic_buffer = sycl::malloc_device(1, *q); + if (atomic_buffer == nullptr) { + throw std::runtime_error( + "moe_prefill_s2_dpas(per-group): failed to allocate atomic buffer"); + } + +#define ARK_DPAS_S2_PG_LAUNCH_SYM(policy) \ + MoEGEMMLauncher_s2<'R', 'C', policy>( \ + *q, activations_ca, weights_i2, scales_ca, \ + static_cast(nullptr), outputs_ca, N, K, \ + num_tokens_per_expert, E, group_size, atomic_buffer); + + if (A_avg_M <= 8) { + ARK_DPAS_S2_PG_LAUNCH_SYM(dpas_w8a16_policy_m_16); + } else if (A_avg_M <= 32) { + ARK_DPAS_S2_PG_LAUNCH_SYM(dpas_w8a16_policy_m_32); + } else { + ARK_DPAS_S2_PG_LAUNCH_SYM(dpas_w8a16_policy); + } +#undef ARK_DPAS_S2_PG_LAUNCH_SYM + + sycl::free(atomic_buffer, *q); +} + +// --------------------------------------------------------------------------- +// Env-flag helper -- `ARK_MOE_PREFILL_DPAS_S2` (default ON, semantics +// identical to `moe_prefill_dpas_s4_enabled`). Decoupled from +// `ARK_MOE_PREFILL_DPAS_S4` / `ARK_MOE_PREFILL_DPAS_INT8` so this new +// single-pass path can be disabled in isolation if it regresses -- +// switching S2 off falls back to the S2->S8 upcast + INT8 DPAS path which +// is itself gated by `ARK_MOE_PREFILL_DPAS_INT8`. +// +// Truthy values (case-insensitive): "1", "true", "on", "yes". Explicit +// "0" / "false" / "off" / "no" disable. Re-read on every call so +// benchmarks / tests can toggle the path in-process. +// --------------------------------------------------------------------------- +inline bool moe_prefill_dpas_s2_enabled() { + const char* env = std::getenv("ARK_MOE_PREFILL_DPAS_S2"); + if (env == nullptr) return true; // default ON + std::string s(env); + for (auto& c : s) c = static_cast(std::tolower(static_cast(c))); + if (s == "0" || s == "false" || s == "off" || s == "no") return false; + return true; +} + +// --------------------------------------------------------------------------- +// Shape preconditions for the per-K-group S2 dispatcher. Mirrors the S4 +// per-group predicate but requires K to be a multiple of 4 so packed +// crumbs never straddle a byte boundary -- guaranteed for auto-round +// S2_CLIP callers by construction, checked here for defensive parity with +// the launcher's throw. +// --------------------------------------------------------------------------- +inline bool moe_prefill_dpas_s2_pergroup_shape_ok(int N, int K, + int group_size) { + if (N <= 0 || K <= 0 || group_size <= 0) return false; + if (N % 64 != 0) return false; + if (K % 32 != 0) return false; + if ((K & 0x3) != 0) return false; + if (K % group_size != 0) return false; + if ((group_size & 0x3) != 0) return false; + if (group_size != 32 && group_size != 64 && group_size != 128 && + group_size != 256) { + return false; + } + return true; +} + +} // namespace moe_dpas_s2 +} // namespace ark + +#endif // ARK_XPU && ARK_SYCL_TLA From 18c28975161927060e98f27004a6122be1e07a52 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 9 Jul 2026 02:27:09 +0000 Subject: [PATCH 12/17] feat: add env-gated 2-token weight-reuse int4/int2 MoE decode kernels --- .../wrapper/include/sycl_tla_moe_decode.hpp | 384 ++++++++++++++++++ 1 file changed, 384 insertions(+) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp index 1b8a4feb0..a59c33308 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp @@ -35,7 +35,9 @@ #include #include +#include #include +#include #include "bestla/bestla.h" #include "sycl_tla_moe_dequant.hpp" @@ -86,9 +88,46 @@ class MoEDecodeKernelInt8; template class MoEDecodeKernelInt2; +// Multi-token (2 tokens per sub-group) variants for the packed low-bit +// paths. See `launch_int4_multitoken` / `launch_int2_multitoken` and the +// `ARK_MOE_DECODE_MULTI_TOKEN` env gate for the weight-reuse rationale. +template +class MoEDecodeKernelInt4MT; + +template +class MoEDecodeKernelInt2MT; + template class MoEDecodeKernelFP8; +// ---------------------------------------------------------------------------- +// Env gate: `ARK_MOE_DECODE_MULTI_TOKEN` (default OFF). +// +// When enabled AND `total_tokens >= 2`, the packed INT4 / INT2 decode +// paths use a 2-tokens-per-sub-group kernel: each sub-group lane owns one +// output N-column for a PAIR of tokens and streams the (BW-dominant) +// packed weight row once per iteration for both tokens. When the two +// tokens map to the SAME expert the second token's weight loads hit L1 +// (same address) -- roughly halving real weight traffic vs. the +// one-token-per-sub-group kernel; when they map to different experts the +// result is still correct (the second load is a real fetch of that +// expert's row). Decode (dequant) is recomputed per token but is cheap. +// +// OFF by default: the win depends on decode batch shape / routing, so it +// is opt-in and A/B-measurable in-process (re-read on every call). +// Truthy (case-insensitive): "1", "true", "on", "yes". +// ---------------------------------------------------------------------------- +inline bool moe_decode_multi_token_enabled() { + const char* env = std::getenv("ARK_MOE_DECODE_MULTI_TOKEN"); + if (env == nullptr) return false; + std::string s(env); + for (auto& c : s) c = static_cast(std::tolower(static_cast(c))); + if (s == "1" || s == "true" || s == "on" || s == "yes") return true; + return false; +} + +template +class MoEDecodeKernelFP8; // ---------------------------------------------------------------------------- // FP8 weight dequantization primitives + host-side env-var reader live in // `sycl_tla_moe_dequant.hpp` so the prefill (mixed-input Grouped GEMM) and @@ -539,6 +578,284 @@ void launch_int2(sycl::queue* q, const ScalarT* activations, const uint8_t* weig }); } +// ---------------------------------------------------------------------------- +// INT4 (S4_CLIP) GEMV, 2 tokens per sub-group (weight-reuse variant). +// +// Numerically identical to `launch_int4` -- same `decode_int4_pair` bit +// ops, same fp32 accumulation order per token. The only difference is the +// work decomposition: dim0 of the launch grid is `ceil(total_tokens / 2)` +// and each sub-group lane processes a PAIR of tokens (t0 = 2*p, t1 = +// 2*p+1) for the same output N-column. The packed weight row for each +// token is streamed inside one loop iteration; when both tokens route to +// the same expert the two loads share an address so the second is an L1 +// hit, halving real weight traffic. When the pair straddles an expert +// boundary the second load is a real fetch of that expert's row, keeping +// the result correct. The odd tail token (when total_tokens is odd) is +// handled by clamping t1 to t0 and suppressing its store. +// ---------------------------------------------------------------------------- +template +void launch_int4_multitoken(sycl::queue* q, const ScalarT* activations, const uint8_t* weights, + const ScalarT* scales, const ScalarT* zeros, ScalarT* outputs, + const int* expert_id_per_token, int total_tokens, int N, int K, int group_size) { + if (N % N_TILE != 0) { + throw std::invalid_argument("moe_gemm_decode(int4): N must be a multiple of 16"); + } + if (K % group_size != 0 || (group_size & 1) != 0) { + throw std::invalid_argument("moe_gemm_decode(int4): K must be a multiple of group_size and group_size must be even"); + } + if (Asym && zeros == nullptr) { + throw std::invalid_argument("moe_gemm_decode(int4): zeros pointer required when asym=true"); + } + if (total_tokens == 0) return; + + const int n_tiles = N / N_TILE; + const int num_groups_k = K / group_size; + const int k_packed = K / 2; + const int token_pairs = (total_tokens + 1) / 2; + + sycl::range<2> global{static_cast(token_pairs), static_cast(n_tiles * SG_SIZE)}; + sycl::range<2> local{1, static_cast(SG_SIZE)}; + + q->parallel_for>( + sycl::nd_range<2>(global, local), + [=](sycl::nd_item<2> it) [[intel::reqd_sub_group_size(SG_SIZE)]] { + const int pair = static_cast(it.get_global_id(0)); + const int n_tile = static_cast(it.get_group(1)); + const int lane = static_cast(it.get_local_id(1)); + const int n_global = n_tile * N_TILE + lane; + + const int t0 = 2 * pair; + const int t1_raw = t0 + 1; + const bool has_t1 = t1_raw < total_tokens; + const int t1 = has_t1 ? t1_raw : t0; // clamp tail; store suppressed + + const int e0 = expert_id_per_token[t0]; + const int e1 = expert_id_per_token[t1]; + + const ScalarT* act0 = activations + static_cast(t0) * K; + const ScalarT* act1 = activations + static_cast(t1) * K; + const uint8_t* w0 = weights + (static_cast(e0) * N + static_cast(n_global)) * k_packed; + const uint8_t* w1 = weights + (static_cast(e1) * N + static_cast(n_global)) * k_packed; + const ScalarT* s0 = scales + (static_cast(e0) * N + static_cast(n_global)) * num_groups_k; + const ScalarT* s1 = scales + (static_cast(e1) * N + static_cast(n_global)) * num_groups_k; + const ScalarT* z0 = Asym ? zeros + (static_cast(e0) * N + static_cast(n_global)) * num_groups_k + : nullptr; + const ScalarT* z1 = Asym ? zeros + (static_cast(e1) * N + static_cast(n_global)) * num_groups_k + : nullptr; + + float acc0 = 0.0f; + float acc1 = 0.0f; + for (int g = 0; g < num_groups_k; ++g) { + const float scale0 = static_cast(s0[g]); + const float scale1 = static_cast(s1[g]); + float zero0 = 0.0f, zero1 = 0.0f; + if constexpr (Asym) { + zero0 = static_cast(z0[g]); + zero1 = static_cast(z1[g]); + } + const int k_base = g * group_size; + constexpr int CHUNK = 16; + using ActVec = sycl::vec; + using PackVec = sycl::vec; + static_assert(sizeof(ScalarT) == sizeof(uint16_t), "ScalarT must be a 16-bit floating type"); + const int chunk_end = (group_size / CHUNK) * CHUNK; + int kk = 0; + for (; kk < chunk_end; kk += CHUNK) { + const ActVec av0 = *reinterpret_cast(act0 + k_base + kk); + const ActVec av1 = *reinterpret_cast(act1 + k_base + kk); + const PackVec pv0 = *reinterpret_cast(w0 + (k_base + kk) / 2); + const PackVec pv1 = *reinterpret_cast(w1 + (k_base + kk) / 2); +#pragma unroll + for (int b = 0; b < CHUNK / 2; ++b) { + int q0a, q1a, q0b, q1b; + decode_int4_pair(pv0[b], q0a, q1a); + decode_int4_pair(pv1[b], q0b, q1b); + float w0a, w1a, w0b, w1b; + if constexpr (Asym) { + w0a = (static_cast(q0a) - zero0) * scale0; + w1a = (static_cast(q1a) - zero0) * scale0; + w0b = (static_cast(q0b) - zero1) * scale1; + w1b = (static_cast(q1b) - zero1) * scale1; + } else { + w0a = static_cast(q0a) * scale0; + w1a = static_cast(q1a) * scale0; + w0b = static_cast(q0b) * scale1; + w1b = static_cast(q1b) * scale1; + } + const float a0_0 = static_cast(sycl::bit_cast(static_cast(av0[2 * b]))); + const float a0_1 = static_cast(sycl::bit_cast(static_cast(av0[2 * b + 1]))); + const float a1_0 = static_cast(sycl::bit_cast(static_cast(av1[2 * b]))); + const float a1_1 = static_cast(sycl::bit_cast(static_cast(av1[2 * b + 1]))); + acc0 += a0_0 * w0a; + acc0 += a0_1 * w1a; + acc1 += a1_0 * w0b; + acc1 += a1_1 * w1b; + } + } + for (; kk < group_size; kk += 2) { + const uint8_t p0 = w0[(k_base + kk) / 2]; + const uint8_t p1 = w1[(k_base + kk) / 2]; + int q0a, q1a, q0b, q1b; + decode_int4_pair(p0, q0a, q1a); + decode_int4_pair(p1, q0b, q1b); + float w0a, w1a, w0b, w1b; + if constexpr (Asym) { + w0a = (static_cast(q0a) - zero0) * scale0; + w1a = (static_cast(q1a) - zero0) * scale0; + w0b = (static_cast(q0b) - zero1) * scale1; + w1b = (static_cast(q1b) - zero1) * scale1; + } else { + w0a = static_cast(q0a) * scale0; + w1a = static_cast(q1a) * scale0; + w0b = static_cast(q0b) * scale1; + w1b = static_cast(q1b) * scale1; + } + acc0 += static_cast(act0[k_base + kk]) * w0a; + acc0 += static_cast(act0[k_base + kk + 1]) * w1a; + acc1 += static_cast(act1[k_base + kk]) * w0b; + acc1 += static_cast(act1[k_base + kk + 1]) * w1b; + } + } + + outputs[static_cast(t0) * N + n_global] = static_cast(acc0); + if (has_t1) { + outputs[static_cast(t1) * N + n_global] = static_cast(acc1); + } + }); +} + +// ---------------------------------------------------------------------------- +// INT2 (S2_CLIP) GEMV, 2 tokens per sub-group (weight-reuse variant). +// Numerically identical to `launch_int2`; work decomposition mirrors +// `launch_int4_multitoken` (see its comment for the weight-reuse rationale). +// ---------------------------------------------------------------------------- +template +void launch_int2_multitoken(sycl::queue* q, const ScalarT* activations, const uint8_t* weights, + const ScalarT* scales, const ScalarT* zeros, ScalarT* outputs, + const int* expert_id_per_token, int total_tokens, int N, int K, int group_size) { + if (N % N_TILE != 0) { + throw std::invalid_argument("moe_gemm_decode(int2): N must be a multiple of 16"); + } + if ((K & 0x3) != 0) { + throw std::invalid_argument("moe_gemm_decode(int2): K must be a multiple of 4"); + } + if (K % group_size != 0 || (group_size & 0x3) != 0) { + throw std::invalid_argument( + "moe_gemm_decode(int2): K must be a multiple of group_size and group_size must be a multiple of 4"); + } + if (Asym && zeros == nullptr) { + throw std::invalid_argument("moe_gemm_decode(int2): zeros pointer required when asym=true"); + } + if (total_tokens == 0) return; + + const int n_tiles = N / N_TILE; + const int num_groups_k = K / group_size; + const int k_packed = K / 4; + const int token_pairs = (total_tokens + 1) / 2; + + sycl::range<2> global{static_cast(token_pairs), static_cast(n_tiles * SG_SIZE)}; + sycl::range<2> local{1, static_cast(SG_SIZE)}; + + q->parallel_for>( + sycl::nd_range<2>(global, local), + [=](sycl::nd_item<2> it) [[intel::reqd_sub_group_size(SG_SIZE)]] { + const int pair = static_cast(it.get_global_id(0)); + const int n_tile = static_cast(it.get_group(1)); + const int lane = static_cast(it.get_local_id(1)); + const int n_global = n_tile * N_TILE + lane; + + const int t0 = 2 * pair; + const int t1_raw = t0 + 1; + const bool has_t1 = t1_raw < total_tokens; + const int t1 = has_t1 ? t1_raw : t0; + + const int e0 = expert_id_per_token[t0]; + const int e1 = expert_id_per_token[t1]; + + const ScalarT* act0 = activations + static_cast(t0) * K; + const ScalarT* act1 = activations + static_cast(t1) * K; + const uint8_t* w0 = weights + (static_cast(e0) * N + static_cast(n_global)) * k_packed; + const uint8_t* w1 = weights + (static_cast(e1) * N + static_cast(n_global)) * k_packed; + const ScalarT* s0 = scales + (static_cast(e0) * N + static_cast(n_global)) * num_groups_k; + const ScalarT* s1 = scales + (static_cast(e1) * N + static_cast(n_global)) * num_groups_k; + const ScalarT* z0 = Asym ? zeros + (static_cast(e0) * N + static_cast(n_global)) * num_groups_k + : nullptr; + const ScalarT* z1 = Asym ? zeros + (static_cast(e1) * N + static_cast(n_global)) * num_groups_k + : nullptr; + + float acc0 = 0.0f; + float acc1 = 0.0f; + for (int g = 0; g < num_groups_k; ++g) { + const float scale0 = static_cast(s0[g]); + const float scale1 = static_cast(s1[g]); + float zero0 = 0.0f, zero1 = 0.0f; + if constexpr (Asym) { + zero0 = static_cast(z0[g]); + zero1 = static_cast(z1[g]); + } + const int k_base = g * group_size; + constexpr int CHUNK = 16; + using ActVec = sycl::vec; + using PackVec = sycl::vec; + static_assert(sizeof(ScalarT) == sizeof(uint16_t), "ScalarT must be a 16-bit floating type"); + const int chunk_end = (group_size / CHUNK) * CHUNK; + int kk = 0; + for (; kk < chunk_end; kk += CHUNK) { + const ActVec av0 = *reinterpret_cast(act0 + k_base + kk); + const ActVec av1 = *reinterpret_cast(act1 + k_base + kk); + const PackVec pv0 = *reinterpret_cast(w0 + (k_base + kk) / 4); + const PackVec pv1 = *reinterpret_cast(w1 + (k_base + kk) / 4); +#pragma unroll + for (int b = 0; b < CHUNK / 4; ++b) { + int qa[4], qb[4]; + decode_int2_quad(pv0[b], qa); + decode_int2_quad(pv1[b], qb); +#pragma unroll + for (int j = 0; j < 4; ++j) { + float wa, wb; + if constexpr (Asym) { + wa = (static_cast(qa[j]) - zero0) * scale0; + wb = (static_cast(qb[j]) - zero1) * scale1; + } else { + wa = static_cast(qa[j]) * scale0; + wb = static_cast(qb[j]) * scale1; + } + const float a0 = static_cast(sycl::bit_cast(static_cast(av0[4 * b + j]))); + const float a1 = static_cast(sycl::bit_cast(static_cast(av1[4 * b + j]))); + acc0 += a0 * wa; + acc1 += a1 * wb; + } + } + } + for (; kk < group_size; kk += 4) { + const uint8_t p0 = w0[(k_base + kk) / 4]; + const uint8_t p1 = w1[(k_base + kk) / 4]; + int qa[4], qb[4]; + decode_int2_quad(p0, qa); + decode_int2_quad(p1, qb); +#pragma unroll + for (int j = 0; j < 4; ++j) { + float wa, wb; + if constexpr (Asym) { + wa = (static_cast(qa[j]) - zero0) * scale0; + wb = (static_cast(qb[j]) - zero1) * scale1; + } else { + wa = static_cast(qa[j]) * scale0; + wb = static_cast(qb[j]) * scale1; + } + acc0 += static_cast(act0[k_base + kk + j]) * wa; + acc1 += static_cast(act1[k_base + kk + j]) * wb; + } + } + } + + outputs[static_cast(t0) * N + n_global] = static_cast(acc0); + if (has_t1) { + outputs[static_cast(t1) * N + n_global] = static_cast(acc1); + } + }); +} + // ---------------------------------------------------------------------------- // FP8 (E4M3 / E5M2) GEMV with group-wise scale (no zero-point). // @@ -663,6 +980,40 @@ inline void moe_gemm_decode(sycl::queue* q, void* activations, void* weights, vo } if (weight_dtype == BTLA_DTYPE::S4_CLIP) { + // Opt-in 2-tokens-per-sub-group weight-reuse path. Numerically + // identical to the one-token kernel; halves real weight traffic for + // same-expert token pairs. Gated by `ARK_MOE_DECODE_MULTI_TOKEN`. + if (moe_decode_detail::moe_decode_multi_token_enabled() && total_tokens >= 2) { + if (act_dtype == BTLA_DTYPE::F16) { + if (asym) { + moe_decode_detail::launch_int4_multitoken( + q, static_cast(activations), static_cast(weights), + static_cast(scales), static_cast(zeros), + static_cast(outputs), expert_id_per_token_buf, total_tokens, N, K, group_size); + } else { + moe_decode_detail::launch_int4_multitoken( + q, static_cast(activations), static_cast(weights), + static_cast(scales), static_cast(zeros), + static_cast(outputs), expert_id_per_token_buf, total_tokens, N, K, group_size); + } + return; + } else if (act_dtype == BTLA_DTYPE::BF16) { + using BF = sycl::ext::oneapi::bfloat16; + if (asym) { + moe_decode_detail::launch_int4_multitoken( + q, static_cast(activations), static_cast(weights), + static_cast(scales), static_cast(zeros), static_cast(outputs), + expert_id_per_token_buf, total_tokens, N, K, group_size); + } else { + moe_decode_detail::launch_int4_multitoken( + q, static_cast(activations), static_cast(weights), + static_cast(scales), static_cast(zeros), static_cast(outputs), + expert_id_per_token_buf, total_tokens, N, K, group_size); + } + return; + } + // Unsupported act_dtype falls through to the one-token path below. + } if (act_dtype == BTLA_DTYPE::F16) { if (asym) { moe_decode_detail::launch_int4( @@ -727,6 +1078,39 @@ inline void moe_gemm_decode(sycl::queue* q, void* activations, void* weights, vo } if (weight_dtype == BTLA_DTYPE::S2_CLIP) { + // Opt-in 2-tokens-per-sub-group weight-reuse path (see the S4_CLIP + // block above). Gated by `ARK_MOE_DECODE_MULTI_TOKEN`. + if (moe_decode_detail::moe_decode_multi_token_enabled() && total_tokens >= 2) { + if (act_dtype == BTLA_DTYPE::F16) { + if (asym) { + moe_decode_detail::launch_int2_multitoken( + q, static_cast(activations), static_cast(weights), + static_cast(scales), static_cast(zeros), + static_cast(outputs), expert_id_per_token_buf, total_tokens, N, K, group_size); + } else { + moe_decode_detail::launch_int2_multitoken( + q, static_cast(activations), static_cast(weights), + static_cast(scales), static_cast(zeros), + static_cast(outputs), expert_id_per_token_buf, total_tokens, N, K, group_size); + } + return; + } else if (act_dtype == BTLA_DTYPE::BF16) { + using BF = sycl::ext::oneapi::bfloat16; + if (asym) { + moe_decode_detail::launch_int2_multitoken( + q, static_cast(activations), static_cast(weights), + static_cast(scales), static_cast(zeros), static_cast(outputs), + expert_id_per_token_buf, total_tokens, N, K, group_size); + } else { + moe_decode_detail::launch_int2_multitoken( + q, static_cast(activations), static_cast(weights), + static_cast(scales), static_cast(zeros), static_cast(outputs), + expert_id_per_token_buf, total_tokens, N, K, group_size); + } + return; + } + // Unsupported act_dtype falls through to the one-token path below. + } if (act_dtype == BTLA_DTYPE::F16) { if (asym) { moe_decode_detail::launch_int2( From 32be7a7d05df5930c4e7b8b491eed4731e8f4c28 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 9 Jul 2026 02:29:31 +0000 Subject: [PATCH 13/17] docs: document S2 single-pass prefill + multi-token decode env gates (EN/CN) --- .../ark/test/README_MOE_PREFILL_PERF.md | 49 +++++++++++++++++-- .../ark/test/README_MOE_PREFILL_PERF_CN.md | 43 +++++++++++++++- 2 files changed, 87 insertions(+), 5 deletions(-) diff --git a/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md b/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md index ff97deaa1..e783ef0de 100644 --- a/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md +++ b/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md @@ -298,7 +298,48 @@ which forces `ARK_MOE_PREFILL_DPAS_S4=1` + exclusively exercised, at the same production shapes as `test_accuracy_int4`, with tolerance `rtol=atol=1e-1`. -## FP8 per-expert (per-tensor) perf tests +## INT2-sym Prefill Paths (opt-in env flags) + +The INT2 sym prefill path mirrors the INT4-sym design one bit-width down: +weights are packed `[E, N, K/4]` `uint8_t` (four 2-bit sym-signed fields +per byte, `[-2, 1]`). A single-pass **DPAS S2** mainloop folds the +crumb→`act_dtype` upcast into the DPAS pipeline so the standalone dequant +pass is eliminated. asym S2 always falls through to the dequant path. + +| Precedence | Env flag | Kernel | +| ----------- | --------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| 1 (highest) | `ARK_MOE_PREFILL_DPAS_S2` unset or truthy (**default ON**) | **S2-sym single-pass DPAS mixed-input mainloop.** Reads packed `[E, N, K/4]` `uint8_t` crumbs directly and folds the S2→`act_dtype` upcast into the DPAS mainloop via CuTe's `reorder(tBrB, tCrB)` (which relies on `NumericArrayConverter`). B-side global traffic is a quarter of the S8 path. Per-K-group scale is applied through the same deferred group-boundary fold as INT4/INT8. Implemented in `sycl_tla_moe_prefill_s2_dpas.hpp`. **Status: NEEDS-HARDWARE-VALIDATION** (untested port). | +| 2 (fallback)| `ARK_MOE_PREFILL_DPAS_S2=0` and `ARK_MOE_PREFILL_DPAS_INT8` truthy (**default ON**) | **S2→S8 upcast + shared INT8 DPAS mainloop.** Two-pass: `launch_upcast_int2_sym_to_int8` writes an `[E, N, K]` `int8_t` view of the dequant workspace, then the standard INT8 per-group DPAS mainloop consumes it. Robust but pays the ~E·N·K byte round-trip vs. path 1. Implemented in `sycl_tla_moe_mixed.hpp` + `sycl_tla_moe_prefill_int_dpas.hpp`. | +| 3 (default) | `ARK_MOE_PREFILL_DPAS_S2=0` and `ARK_MOE_PREFILL_DPAS_INT8=0` | v1 dequant kernel (`sycl_tla_moe_mixed.hpp::launch_dequant_int2`) followed by the stock bf16/fp16 grouped GEMM. Handles both sym and asym. | + +**S2 DPAS path shape preconditions** — the `moe_gemm_prefill` dispatcher +silently falls back to precedence 2 (then 3) whenever any of these fail: + +- `N % 64 == 0` (BN) +- `K % 32 == 0` (BK) +- `K % 4 == 0` (four crumbs per packed byte never straddle a byte boundary) +- `K % group_size == 0` +- `group_size % 4 == 0` (crumb quad never straddles a group boundary) +- `group_size ∈ {32, 64, 128, 256}` +- `asym == false` (asym S2 is out of scope for the DPAS path) + +## Decode weight-reuse (multi-token, opt-in) + +The packed INT4/INT2 **decode** (GEMV) kernels already fuse dequant inline +(no separate dequant pass), so their throughput is dominated by packed +weight bandwidth. When more than one token is being decoded, +`ARK_MOE_DECODE_MULTI_TOKEN=1` selects a 2-tokens-per-sub-group variant: +each sub-group lane owns one output N-column for a *pair* of tokens and +streams the packed weight row once per iteration for both tokens. When +both tokens route to the same expert the second token's weight load hits +L1 (same address), roughly halving real weight traffic; when the pair +straddles an expert boundary the second load is a real fetch of that +expert's row, keeping the result numerically identical to the one-token +kernel. Default **OFF** (the win depends on decode batch shape / routing); +covered by `sycl_tla_moe_decode.hpp::launch_int4_multitoken` / +`launch_int2_multitoken`. + + `test_perf_fp8_per_tensor` benchmarks the Variant A DPAS path against the single-`torch.bmm` baseline for the **one-FP32-scalar-per-expert** @@ -363,5 +404,7 @@ the same production shapes, with the standard INT8 tolerance (`rtol=atol=1e-1`). **Status: NEEDS-HARDWARE-VALIDATION** (untested port; sym-only for -Phase 1 — per-group and asym INT4 / INT2 DPAS are follow-up phases -that will reuse the same mainloop skeleton with an added unpack step). +Phase 1. Per-group INT4 (S4) and INT2 (S2) DPAS now exist as single-pass +paths — see the "INT4-sym / INT2-sym Prefill Paths" sections above. asym +INT4 / INT2 DPAS remain follow-up phases that will reuse the same +mainloop skeleton with an added activation-sum precompute). diff --git a/auto_round_extension/ark/test/README_MOE_PREFILL_PERF_CN.md b/auto_round_extension/ark/test/README_MOE_PREFILL_PERF_CN.md index 732cfec70..c8a48f2d4 100644 --- a/auto_round_extension/ark/test/README_MOE_PREFILL_PERF_CN.md +++ b/auto_round_extension/ark/test/README_MOE_PREFILL_PERF_CN.md @@ -228,6 +228,43 @@ S4-sym 有两条独立的 DPAS 路径;asym S4 始终回退到 dequant 路径。 `ARK_MOE_PREFILL_DPAS_INT8=0`,专门验证单遍 mainloop 路径,形状矩阵与 `test_accuracy_int4` 一致,容差 `rtol=atol=1e-1`。 +## INT2-sym Prefill 路径(opt-in env 开关) + +INT2 sym 路径与 INT4-sym 设计同构,只是位宽再降一档:权重按 +`[E, N, K/4]` `uint8_t` 打包(每字节四个 2-bit 对称有符号字段, +`[-2, 1]`)。单遍 **DPAS S2** mainloop 把 crumb→`act_dtype` 上转折叠进 +DPAS 流水线,从而消除独立的 dequant 遍。asym S2 始终回退到 dequant 路径。 + +| 优先级 | Env 开关 | Kernel | +| ------------ | --------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| 1 (最高) | `ARK_MOE_PREFILL_DPAS_S2` 未设置或为真值(**默认开启**) | **S2-sym 单遍 DPAS 混合输入 mainloop**。直接读取 packed `[E, N, K/4]` `uint8_t` crumb,通过 CuTe `reorder(tBrB, tCrB)`(依赖 `NumericArrayConverter`)在寄存器中把 S2 上转到 `act_dtype`。B 侧 global 带宽是 S8 路径的四分之一。Per-K-group scale 使用与 INT4/INT8 相同的组边界延迟折叠。实现于 `sycl_tla_moe_prefill_s2_dpas.hpp`。**状态:NEEDS-HARDWARE-VALIDATION**(未经测试的移植)。 | +| 2 (回退) | `ARK_MOE_PREFILL_DPAS_S2=0` 且 `ARK_MOE_PREFILL_DPAS_INT8` 为真值(**默认开启**) | **S2→S8 上转 + 共享 INT8 DPAS mainloop**。两遍:`launch_upcast_int2_sym_to_int8` 把权重写成 `[E, N, K]` `int8_t`(复用 dequant workspace),再由标准 INT8 per-group DPAS mainloop 消费。相较路径 1 需要付出 ~E·N·K 字节的 workspace 往返。实现于 `sycl_tla_moe_mixed.hpp` + `sycl_tla_moe_prefill_int_dpas.hpp`。 | +| 3 (默认回退) | `ARK_MOE_PREFILL_DPAS_S2=0` 且 `ARK_MOE_PREFILL_DPAS_INT8=0` | v1 dequant kernel(`sycl_tla_moe_mixed.hpp::launch_dequant_int2`)后接标准 bf16/fp16 grouped GEMM。同时支持 sym 与 asym。 | + +**S2 DPAS 路径形状前置条件** — 任何条件不满足时,`moe_gemm_prefill` +分发器会静默回退到优先级 2(再回退到 3): + +- `N % 64 == 0` (BN) +- `K % 32 == 0` (BK) +- `K % 4 == 0`(每字节四个 crumb 不跨越字节边界) +- `K % group_size == 0` +- `group_size % 4 == 0`(crumb 四元组不会跨越组边界) +- `group_size ∈ {32, 64, 128, 256}` +- `asym == false`(asym S2 不在 DPAS 路径的支持范围内) + +## Decode 权重复用(multi-token,opt-in) + +打包的 INT4/INT2 **decode**(GEMV)kernel 已经在寄存器中内联融合了 +dequant(没有独立的 dequant 遍),因此其吞吐由打包权重带宽主导。当同时 +decode 多个 token 时,`ARK_MOE_DECODE_MULTI_TOKEN=1` 会选择 +2-tokens-per-sub-group 变体:每个 sub-group lane 为一*对* token 负责同一个 +输出 N 列,并在每次迭代中把打包权重行只流一次供两个 token 共用。当两个 +token 路由到同一专家时,第二个 token 的权重加载命中 L1(同一地址),把真实 +权重流量大致减半;当这一对跨越专家边界时,第二次加载是对该专家权重行的真实 +读取,结果与单 token kernel 数值完全一致。默认**关闭**(收益取决于 decode +的 batch 形状 / 路由);由 `sycl_tla_moe_decode.hpp::launch_int4_multitoken` / +`launch_int2_multitoken` 覆盖。 + ## FP8 per-expert (per-tensor) 性能测试 `test_perf_fp8_per_tensor` 提供 Variant A DPAS 路径的性能表格,对应 @@ -286,5 +323,7 @@ pytest -v -s test_moe_prefill_perf.py::TestMoEGemmPrefillPerf::test_perf_int8_pe (`rtol=atol=1e-1`)。 **状态:NEEDS-HARDWARE-VALIDATION**(未在硬件上验证过的移植;Phase 1 -仅支持 sym。per-group / asym 的 INT4 / INT2 DPAS 是后续阶段,将复用 -同一份 mainloop 骨架,只在其中追加 unpack 步骤)。 +仅支持 sym。per-group 的 INT4 (S4) 与 INT2 (S2) DPAS 现已作为单遍路径存在 +——见上文 "INT4-sym / INT2-sym Prefill 路径" 小节。asym 的 INT4 / INT2 +DPAS 仍为后续阶段,将复用同一份 mainloop 骨架,只在其中追加 activation-sum +预计算)。 From 9441fdfe0c431600fd90b6defddf03f541cc4b3f Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 9 Jul 2026 02:30:37 +0000 Subject: [PATCH 14/17] test: add S2 single-pass DPAS prefill accuracy parity test --- .../ark/test/README_MOE_PREFILL_PERF.md | 6 ++ .../ark/test/README_MOE_PREFILL_PERF_CN.md | 6 ++ .../ark/test/test_moe_prefill_accuracy.py | 61 +++++++++++++++++++ 3 files changed, 73 insertions(+) diff --git a/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md b/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md index e783ef0de..a5beb4ba3 100644 --- a/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md +++ b/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md @@ -323,6 +323,12 @@ silently falls back to precedence 2 (then 3) whenever any of these fail: - `group_size ∈ {32, 64, 128, 256}` - `asym == false` (asym S2 is out of scope for the DPAS path) +Accuracy parity is covered by +`test_moe_prefill_accuracy.py::test_accuracy_int2_dpas_per_group`, which +forces `ARK_MOE_PREFILL_DPAS_S2=1` + `ARK_MOE_PREFILL_DPAS_INT8=0` so the +single-pass mainloop is exclusively exercised, at the same production +shapes as `test_accuracy_int2`, with tolerance `rtol=atol=1.5e-1`. + ## Decode weight-reuse (multi-token, opt-in) The packed INT4/INT2 **decode** (GEMV) kernels already fuse dequant inline diff --git a/auto_round_extension/ark/test/README_MOE_PREFILL_PERF_CN.md b/auto_round_extension/ark/test/README_MOE_PREFILL_PERF_CN.md index c8a48f2d4..21489c5e2 100644 --- a/auto_round_extension/ark/test/README_MOE_PREFILL_PERF_CN.md +++ b/auto_round_extension/ark/test/README_MOE_PREFILL_PERF_CN.md @@ -252,6 +252,12 @@ DPAS 流水线,从而消除独立的 dequant 遍。asym S2 始终回退到 dequa - `group_size ∈ {32, 64, 128, 256}` - `asym == false`(asym S2 不在 DPAS 路径的支持范围内) +精度对齐由 +`test_moe_prefill_accuracy.py::test_accuracy_int2_dpas_per_group` +覆盖,该用例强制 `ARK_MOE_PREFILL_DPAS_S2=1` + +`ARK_MOE_PREFILL_DPAS_INT8=0`,专门验证单遍 mainloop 路径,形状矩阵与 +`test_accuracy_int2` 一致,容差 `rtol=atol=1.5e-1`。 + ## Decode 权重复用(multi-token,opt-in) 打包的 INT4/INT2 **decode**(GEMV)kernel 已经在寄存器中内联融合了 diff --git a/auto_round_extension/ark/test/test_moe_prefill_accuracy.py b/auto_round_extension/ark/test/test_moe_prefill_accuracy.py index c8c1ee168..3a3dbf6cd 100644 --- a/auto_round_extension/ark/test/test_moe_prefill_accuracy.py +++ b/auto_round_extension/ark/test/test_moe_prefill_accuracy.py @@ -687,6 +687,67 @@ def test_accuracy_int4_dpas_per_group(self, dtype): else: os.environ["ARK_MOE_PREFILL_DPAS_INT8"] = prev_int8 + @pytest.mark.skipif(bool(_QUANT_PREFILL_SKIP), reason=_QUANT_PREFILL_SKIP or "ok") + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + def test_accuracy_int2_dpas_per_group(self, dtype): + """Parity of the S2-sym single-pass DPAS mixed-input mainloop. + + Sibling of :meth:`test_accuracy_int4_dpas_per_group` -- uses the + standard ``moe_gemm_prefill`` call path with + ``ARK_MOE_PREFILL_DPAS_S2=1`` (default ON) and + ``ARK_MOE_PREFILL_DPAS_INT8=0`` so the two-pass S2->S8 upcast + fallback is disabled and we exercise the single-pass mainloop + exclusively. The C++ dispatcher should pick the S2 DPAS branch + for shapes that satisfy ``N%64==0 && K%32==0 && K%4==0 && + K%group_size==0 && group_size%4==0 && group_size in + {32,64,128,256}`` and silently fall back to the dequant path + otherwise, so this test is checking parity, not that the DPAS + branch is exercised. + + STATUS: NEEDS-HARDWARE-VALIDATION. + """ + group_size = 128 + prev_s2 = os.environ.get("ARK_MOE_PREFILL_DPAS_S2") + prev_int8 = os.environ.get("ARK_MOE_PREFILL_DPAS_INT8") + os.environ["ARK_MOE_PREFILL_DPAS_S2"] = "1" + os.environ["ARK_MOE_PREFILL_DPAS_INT8"] = "0" + try: + for label, E, tpe, N, K in PREFILL_SHAPES: + if K % group_size != 0 or K % 4 != 0: + continue + total_tokens = sum(tpe) + activations = torch.randn(total_tokens, K, dtype=dtype, device="xpu") + w_float = (torch.randn(E, N, K, dtype=torch.float32, device="xpu") * 0.1).to(dtype) + scales = torch.empty(E, N, K // group_size, dtype=dtype, device="xpu") + packed = _pack_int2_sym(w_float, scales, group_size) + dequant = _dequant_int2_sym(packed, scales, group_size).to(dtype) + ntpe = torch.tensor(tpe, dtype=torch.int32, device="xpu") + + out = ark.moe_gemm_prefill( + activations, + packed, + ntpe, + scales=scales, + weight_bits=2, + group_size=group_size, + asym=False, + ) + + ref = _reference_moe_prefill(activations, dequant, ntpe) + assert out.shape == (total_tokens, N), f"{label}: bad shape {out.shape}" + torch.testing.assert_close( + out, ref, msg=lambda m, lbl=label: f"[{lbl}] {m}", **_tol_for_dtype(_TOL_INT2, dtype) + ) + finally: + if prev_s2 is None: + os.environ.pop("ARK_MOE_PREFILL_DPAS_S2", None) + else: + os.environ["ARK_MOE_PREFILL_DPAS_S2"] = prev_s2 + if prev_int8 is None: + os.environ.pop("ARK_MOE_PREFILL_DPAS_INT8", None) + else: + os.environ["ARK_MOE_PREFILL_DPAS_INT8"] = prev_int8 + if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) From f98767dc92e7cb909883b1d35e2cb80d58832134 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 9 Jul 2026 03:00:59 +0000 Subject: [PATCH 15/17] fix: patch sycl-tla to break sub-byte cute::reorder SYCL recursion (ARK XPU) --- .../ark/auto_round_kernel/CMakeLists.txt | 16 ++++++++ .../patches/apply_patch.cmake | 40 +++++++++++++++++++ ...cl_tla_fix_subbyte_reorder_recursion.patch | 36 +++++++++++++++++ 3 files changed, 92 insertions(+) create mode 100644 auto_round_extension/ark/auto_round_kernel/patches/apply_patch.cmake create mode 100644 auto_round_extension/ark/auto_round_kernel/patches/sycl_tla_fix_subbyte_reorder_recursion.patch diff --git a/auto_round_extension/ark/auto_round_kernel/CMakeLists.txt b/auto_round_extension/ark/auto_round_kernel/CMakeLists.txt index 5d40df7a4..7b0ab42d4 100755 --- a/auto_round_extension/ark/auto_round_kernel/CMakeLists.txt +++ b/auto_round_extension/ark/auto_round_kernel/CMakeLists.txt @@ -84,10 +84,26 @@ if(ARK_XPU AND ARK_SYCL_TLA) set(SYCL_TLA_GIT_REPOSITORY "https://github.com/luoyu-intel/sycl-tla.git" CACHE STRING "sycl-tla git repository") set(SYCL_TLA_GIT_TAG "260630" CACHE STRING "sycl-tla git tag/commit") + # Patch the pinned sycl-tla tree to break an infinite template recursion in + # cute::reorder for sub-byte weight types (e.g. int4/uint4 MoE dequant). Without + # this, the identity-relayout "convert" step re-selects ReorderDispatchConvertRelayout + # for a sub-byte source and the SYCL device compiler rejects it with + # "SYCL kernel cannot call a recursive function". Keeping the fix as an in-repo patch + # lets us stay on the reproducible pinned tag until it is fixed upstream. + find_package(Git REQUIRED) + set(SYCL_TLA_PATCH + "${CMAKE_CURRENT_LIST_DIR}/patches/sycl_tla_fix_subbyte_reorder_recursion.patch") + FetchContent_Declare( sycl_tla GIT_REPOSITORY ${SYCL_TLA_GIT_REPOSITORY} GIT_TAG ${SYCL_TLA_GIT_TAG} + # Apply the patch only when it has not already been applied, so re-configuring an + # existing build tree does not fail with "patch does not apply". + PATCH_COMMAND ${CMAKE_COMMAND} + -DGIT_EXECUTABLE=${GIT_EXECUTABLE} + -DPATCH_FILE=${SYCL_TLA_PATCH} + -P "${CMAKE_CURRENT_LIST_DIR}/patches/apply_patch.cmake" ) # Map target to device name for -Xs flag diff --git a/auto_round_extension/ark/auto_round_kernel/patches/apply_patch.cmake b/auto_round_extension/ark/auto_round_kernel/patches/apply_patch.cmake new file mode 100644 index 000000000..5cbb6b5e2 --- /dev/null +++ b/auto_round_extension/ark/auto_round_kernel/patches/apply_patch.cmake @@ -0,0 +1,40 @@ +# Idempotently apply a git patch to a fetched dependency source tree. +# +# Usage: +# cmake -DGIT_EXECUTABLE= -DPATCH_FILE= -P apply_patch.cmake +# +# The script is safe to run multiple times: if the patch is already applied +# (detected via `git apply --reverse --check`), it does nothing instead of +# failing with "patch does not apply". + +if(NOT GIT_EXECUTABLE) + set(GIT_EXECUTABLE "git") +endif() + +if(NOT PATCH_FILE) + message(FATAL_ERROR "apply_patch.cmake: PATCH_FILE must be provided") +endif() + +# If the patch reverse-applies cleanly, it is already present -> nothing to do. +execute_process( + COMMAND ${GIT_EXECUTABLE} apply --reverse --check --ignore-whitespace "${PATCH_FILE}" + RESULT_VARIABLE _already_applied + OUTPUT_QUIET + ERROR_QUIET +) + +if(_already_applied EQUAL 0) + message(STATUS "apply_patch.cmake: patch already applied, skipping: ${PATCH_FILE}") + return() +endif() + +execute_process( + COMMAND ${GIT_EXECUTABLE} apply --ignore-whitespace "${PATCH_FILE}" + RESULT_VARIABLE _apply_result +) + +if(NOT _apply_result EQUAL 0) + message(FATAL_ERROR "apply_patch.cmake: failed to apply patch: ${PATCH_FILE}") +endif() + +message(STATUS "apply_patch.cmake: applied patch: ${PATCH_FILE}") diff --git a/auto_round_extension/ark/auto_round_kernel/patches/sycl_tla_fix_subbyte_reorder_recursion.patch b/auto_round_extension/ark/auto_round_kernel/patches/sycl_tla_fix_subbyte_reorder_recursion.patch new file mode 100644 index 000000000..860bd175f --- /dev/null +++ b/auto_round_extension/ark/auto_round_kernel/patches/sycl_tla_fix_subbyte_reorder_recursion.patch @@ -0,0 +1,36 @@ +diff --git a/include/cute/algorithm/reorder.hpp b/include/cute/algorithm/reorder.hpp +index 4c5ae1a..22d3ed1 100644 +--- a/include/cute/algorithm/reorder.hpp ++++ b/include/cute/algorithm/reorder.hpp +@@ -227,7 +227,15 @@ reorder_impl(ReorderDispatchConvertRelayout const&, + using NewSrcType = conditional_t, upcast_subbyte_t, DstType>; + auto src_c = make_fragment_like(src); + +- reorder(src, src_c, slayout, slayout); ++ // Pure element-wise type conversion (identical source/destination TV-layouts). ++ // Performing it directly instead of recursing through reorder()/choose_xe_reorder_impl() ++ // avoids an infinite template recursion: for a sub-byte SrcType the identity relayout ++ // would re-select ReorderDispatchConvertRelayout and call reorder() with the same ++ // sub-byte source again, which the SYCL device compiler rejects as a recursive kernel. ++ CUTE_UNROLL ++ for (int _ri = 0; _ri < int(size(src)); ++_ri) { ++ src_c(_ri) = static_cast(src(_ri)); ++ } + reorder(src_c, dst, slayout, dlayout); + } + +@@ -248,7 +256,13 @@ reorder_impl(ReorderDispatchRelayoutConvert const&, + auto dst_c = make_fragment_like(dst); + + reorder(src, dst_c, slayout, dlayout); +- reorder(dst_c, dst, dlayout, dlayout); ++ // Pure element-wise type conversion (identical source/destination TV-layouts); see the ++ // note in reorder_impl(ReorderDispatchConvertRelayout, ...). Doing it directly avoids the ++ // sub-byte DstType path re-selecting ReorderDispatchRelayoutConvert and recursing. ++ CUTE_UNROLL ++ for (int _ri = 0; _ri < int(size(dst)); ++_ri) { ++ dst(_ri) = static_cast(dst_c(_ri)); ++ } + } + + From ea8f933e01fb2bf90a299456aecc73ad6d42109f Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 9 Jul 2026 03:04:33 +0000 Subject: [PATCH 16/17] docs: clarify semantic equivalence in sycl-tla reorder patch comments --- ...cl_tla_fix_subbyte_reorder_recursion.patch | 24 +++++++++++-------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/patches/sycl_tla_fix_subbyte_reorder_recursion.patch b/auto_round_extension/ark/auto_round_kernel/patches/sycl_tla_fix_subbyte_reorder_recursion.patch index 860bd175f..3f4953f06 100644 --- a/auto_round_extension/ark/auto_round_kernel/patches/sycl_tla_fix_subbyte_reorder_recursion.patch +++ b/auto_round_extension/ark/auto_round_kernel/patches/sycl_tla_fix_subbyte_reorder_recursion.patch @@ -1,17 +1,20 @@ diff --git a/include/cute/algorithm/reorder.hpp b/include/cute/algorithm/reorder.hpp -index 4c5ae1a..22d3ed1 100644 +index 4c5ae1a..827aa4f 100644 --- a/include/cute/algorithm/reorder.hpp +++ b/include/cute/algorithm/reorder.hpp -@@ -227,7 +227,15 @@ reorder_impl(ReorderDispatchConvertRelayout const&, +@@ -227,7 +227,18 @@ reorder_impl(ReorderDispatchConvertRelayout const&, using NewSrcType = conditional_t, upcast_subbyte_t, DstType>; auto src_c = make_fragment_like(src); - reorder(src, src_c, slayout, slayout); -+ // Pure element-wise type conversion (identical source/destination TV-layouts). -+ // Performing it directly instead of recursing through reorder()/choose_xe_reorder_impl() -+ // avoids an infinite template recursion: for a sub-byte SrcType the identity relayout -+ // would re-select ReorderDispatchConvertRelayout and call reorder() with the same -+ // sub-byte source again, which the SYCL device compiler rejects as a recursive kernel. ++ // The convert step passes identical source/destination TV-layouts (slayout, slayout), ++ // so the reorder degenerates to a pure element-wise type conversion with no layout ++ // transformation; a direct static_cast over the fragment is therefore semantically ++ // equivalent. Doing it directly (instead of recursing through reorder() / ++ // choose_xe_reorder_impl()) avoids an infinite template recursion: for a sub-byte ++ // SrcType the identity relayout would re-select ReorderDispatchConvertRelayout and call ++ // reorder() with the same sub-byte source again, which the SYCL device compiler rejects ++ // as a recursive kernel. + CUTE_UNROLL + for (int _ri = 0; _ri < int(size(src)); ++_ri) { + src_c(_ri) = static_cast(src(_ri)); @@ -19,13 +22,14 @@ index 4c5ae1a..22d3ed1 100644 reorder(src_c, dst, slayout, dlayout); } -@@ -248,7 +256,13 @@ reorder_impl(ReorderDispatchRelayoutConvert const&, +@@ -248,7 +259,14 @@ reorder_impl(ReorderDispatchRelayoutConvert const&, auto dst_c = make_fragment_like(dst); reorder(src, dst_c, slayout, dlayout); - reorder(dst_c, dst, dlayout, dlayout); -+ // Pure element-wise type conversion (identical source/destination TV-layouts); see the -+ // note in reorder_impl(ReorderDispatchConvertRelayout, ...). Doing it directly avoids the ++ // The convert step passes identical TV-layouts (dlayout, dlayout), so it is a pure ++ // element-wise type conversion; see the note in ++ // reorder_impl(ReorderDispatchConvertRelayout, ...). Doing it directly avoids the + // sub-byte DstType path re-selecting ReorderDispatchRelayoutConvert and recursing. + CUTE_UNROLL + for (int _ri = 0; _ri < int(size(dst)); ++_ri) { From 9d5c8c43469b7d6b0b6a286f41226bbe75e93cb3 Mon Sep 17 00:00:00 2001 From: "Dong, Bo1" Date: Fri, 10 Jul 2026 11:25:25 +0800 Subject: [PATCH 17/17] add the code Signed-off-by: Dong, Bo1 --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 054f98350..94be8e867 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -# 1.5.1