Implementing bandwidth metrics for performance tests

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

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.
 #

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);

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=<git> -DPATCH_FILE=<abs path> -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}")

auto_round_extension/ark/auto_round_kernel/patches/sycl_tla_fix_subbyte_reorder_recursion.patch

@@ -0,0 +1,40 @@
+diff --git a/include/cute/algorithm/reorder.hpp b/include/cute/algorithm/reorder.hpp
+index 4c5ae1a..827aa4f 100644
+--- a/include/cute/algorithm/reorder.hpp
++++ b/include/cute/algorithm/reorder.hpp
+@@ -227,7 +227,18 @@ reorder_impl(ReorderDispatchConvertRelayout const&,
+   using NewSrcType = conditional_t<is_subbyte_v<SrcType>, upcast_subbyte_t<SrcType>, DstType>;
+   auto src_c = make_fragment_like<NewSrcType>(src);
+
+-  reorder(src, src_c, slayout, slayout);
++  // 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<NewSrcType>(src(_ri));
++  }
+   reorder(src_c, dst, slayout, dlayout);
+ }
+
+@@ -248,7 +259,14 @@ reorder_impl(ReorderDispatchRelayoutConvert const&,
+   auto dst_c = make_fragment_like<NewDstType>(dst);
+
+   reorder(src, dst_c, slayout, dlayout);
+-  reorder(dst_c, dst, dlayout, dlayout);
++  // 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) {
++    dst(_ri) = static_cast<DstType>(dst_c(_ri));
++  }
+ }
+
+