diff --git a/doc/zh/Intel-iGPU-Installation_zh.md b/doc/zh/Intel-iGPU-Installation_zh.md new file mode 100644 index 00000000..8af5d114 --- /dev/null +++ b/doc/zh/Intel-iGPU-Installation_zh.md @@ -0,0 +1,98 @@ +# Intel iGPU XMX 源码安装 + +本页仅介绍源码安装。支持范围、完整启动命令和精度边界见 [使用文档](Intel-iGPU_zh.md)。 +本方案需要 NVIDIA GPU 执行 Attention 等算子，同时需要 Intel GPU 的 Level Zero 驱动和 SYCL 运行时。 + +## 已验证环境 + +Core Ultra 7 356H / Intel `b0a0`、RTX 5060 Laptop 8GB、约30GiB可用物理内存， +Python 3.11、oneAPI DPC++ 2026.1.1、Intel compute runtime 25.18.33578.77、 +Level Zero loader 1.24.0、CUDA Toolkit 13.2、PyTorch 2.9.1、Triton 3.5.1， +Linux 7.0.0-38-generic。这是实测版本记录，不表示任意版本组合都已验证。 + +构建前需要安装： + +- 支持本机 Intel GPU 的内核和 Level Zero 驱动，参照 Intel 的 + [GPU 驱动安装说明](https://dgpu-docs.intel.com/driver/installation.html)。 +- 包含 `icpx`、SYCL runtime、`sycl-ls` 的 Intel oneAPI DPC++ 编译器。 +- 可正常使用 NVIDIA GPU 的 CUDA Toolkit 和 CUDA 版 PyTorch。 +- 源码构建依赖：Git、CMake、Ninja、C/C++工具链、hwloc和NUMA开发库。 + +Ubuntu可先安装系统开发依赖： + +```bash +sudo apt-get update +sudo apt-get install -y build-essential cmake ninja-build git libnuma-dev libhwloc-dev pkg-config +``` + +在独立 Python 3.11 环境中检查两类 GPU： + +```bash +source /opt/intel/oneapi/setvars.sh --force +icpx --version +sycl-ls +nvidia-smi +python -c 'import torch; print(torch.__version__, torch.cuda.is_available())' +ls -l /dev/dri/renderD* +groups +``` + +`sycl-ls` 应列出 Intel Level Zero GPU，当前用户应能读写对应 render node。 +若需加入 `render` 用户组，请按发行版配置并重新登录。`ONEAPI_DEVICE_SELECTOR` 用于选择 SYCL +设备，NVIDIA设备仍由CUDA管理。这里假定当前Python环境已按项目要求安装CUDA版PyTorch。 + +## 通过官方安装入口从源码安装 + +激活 Python 环境后，在源码目录运行根目录的 `install.sh`： + +```bash +git clone --recursive https://github.com/kvcache-ai/ktransformers.git +cd ktransformers +# 切换到包含本后端的提交后运行： +./install.sh --sycl +``` + +`--sycl` 选择 Intel SYCL 构建，安装脚本会完成以下步骤： + +1. 初始化仓库子模块并检查系统依赖；已安装依赖时不重复调用包管理器。 +2. 加载默认位置的 oneAPI 环境，使用 `icpx` 构建本机的 AVX2+SYCL 内核。 +3. 使用已验证的构建设置：关闭 AMX/AVX512、MLA、LTO 和 crash tracing，编译并行度默认为2。 +4. 在构建目录副本中应用已有 llama.cpp 补丁，保留原子模块 checkout。 +5. 构建并安装本仓库的 kt-kernel，再从固定公开 SGLang 源码安装推理服务。 + +Linux 源码构建会将所选 oneAPI 编译器的运行库目录写入扩展的运行时搜索路径，包括 +Intel OpenMP 和 SYCL 库。因此，完成安装后新终端中的普通 `import kt_kernel` 不应依赖 +安装脚本内的临时 `LD_LIBRARY_PATH`。oneAPI 仍需保留在构建时的位置；启动命令中的 +`source .../setvars.sh` 用于初始化完整工具链环境，不会将 oneAPI 打包进扩展。 +安装器还会清除该环境变量并启动新 Python 进程检查扩展能否加载，避免安装时的临时环境掩盖问题。 + +不传 `--sycl` 时保留原来的默认安装流程。 + +如需指定非默认 oneAPI 路径，可设置 `ONEAPI_SETVARS`。如果确实需要调整编译并行度， +可显式设置 `CPUINFER_PARALLEL`，否则保持默认即可。这些是高级覆盖项，不是安装前提。 +已经显式启用其它互斥 GPU 构建后端的终端需先取消相应设置，安装器会明确提示冲突。 + +替换原生扩展前停止模型服务；安装后重启 Python 加载新二进制。`./install.sh kt-kernel --sycl` +可只重建内核；首次安装使用上面的完整入口。 + +## 验证安装 + +```bash +source /opt/intel/oneapi/setvars.sh --force +python -c 'from kt_kernel import kt_kernel_ext as e; print(e.__file__); assert e.moe.SYCLXMXGPTQInt4_MOE.implementation_revision() >= 7' +git submodule status +python -m pip check +``` + +记录实际加载的扩展路径和hash，并运行使用文档中的硬件检查；仅能import不能证明完整功能正常。 +安装后按照[完整启动命令](Intel-iGPU_zh.md#完整启动命令)启动服务。 + +缺少Intel设备时先处理驱动和权限；缺少SYCL共享库时加载构建所用oneAPI； +缺少XMX类或revision过旧时检查当前Python环境和`.so`路径。 +模型需要对称GPTQ INT4、`desc_act=false`和满足要求的group及维度，不能直接传入FP8/BF16权重。 +记录内存和swap状态，保持已验证的chunk及请求上限；全量logprob请求可能比正常生成占用更多显存。 + +如果旧版构建报 `libiomp5.so: cannot open shared object file`，可先在运行推理的同一终端执行 +`source /opt/intel/oneapi/setvars.sh --force`。安装脚本作为子进程运行，不能修改父终端的环境。 +更新到包含运行库路径修复的源码后，执行 `./install.sh kt-kernel --sycl --no-clean` 重建安装； +无需手动复制或软链接运行库。若 oneAPI 已移动或被卸载，则恢复相应工具链或重新构建。 diff --git a/doc/zh/Intel-iGPU_zh.md b/doc/zh/Intel-iGPU_zh.md new file mode 100644 index 00000000..38c9aefe --- /dev/null +++ b/doc/zh/Intel-iGPU_zh.md @@ -0,0 +1,90 @@ +# Intel iGPU XMX 后端：使用方法 + +本后端在 KTransformers 混合推理中将 routed experts 交给 Intel iGPU，Attention、 +Embedding、LM Head 和配置在 NVIDIA 上的专家仍由 CUDA 执行。当前实现不是纯 Intel GPU 模型运行器。 + +已验证硬件为 Core Ultra 7 356H（Panther Lake、Intel `b0a0`）和 RTX 5060 Laptop。 +调优配置针对 Qwen3.5-35B-A3B-GPTQ-Int4：256 experts、top-k8、hidden2048、intermediate512、 +group128，单个 CPUInfer/NUMA pool、单请求。其它平台和模型形状尚未获得相同的性能验证。 + +源码构建、环境版本和驱动检查见 [环境安装](Intel-iGPU-Installation_zh.md)。 + +## 完整启动命令 + +完成源码安装后，激活构建时所用的 Python 环境，从仓库根目录执行以下命令。 +将模型路径和 oneAPI 初始化脚本路径替换为实际位置；`python` 使用当前环境中的解释器。 + +```bash +source /opt/intel/oneapi/setvars.sh --force +MODEL=/path/to/Qwen3.5-35B-A3B-GPTQ-Int4 + +KT_SYCL_XMX_PROFILE=panther-lake \ +SGLANG_DISABLE_CUDNN_CHECK=1 \ +SGLANG_MAMBA_CONV_DTYPE=float16 \ +python -m sglang.launch_server \ + --model "$MODEL" \ + --kt-weight-path "$MODEL" \ + --kt-method SYCL_GPTQ_INT4 \ + --host 127.0.0.1 \ + --port 30000 \ + --kt-cpuinfer 10 \ + --kt-threadpool-count 1 \ + --kt-num-gpu-experts 8 \ + --tensor-parallel-size 1 \ + --dtype float16 \ + --attention-backend triton \ + --chunked-prefill-size 512 \ + --max-prefill-tokens 4096 \ + --max-running-requests 1 \ + --max-total-tokens 8192 \ + --context-length 8192 \ + --mem-fraction-static 0.9 \ + --disable-radix-cache \ + --tool-call-parser qwen3_coder \ + --disable-shared-experts-fusion \ + --watchdog-timeout 1200 \ + --random-seed 20261016 \ + --trust-remote-code +``` + +`KT_SYCL_XMX_PROFILE=panther-lake` 是唯一需要设置的后端配置变量。它自动选择 XMX， +并启用已验证的 INT8 激活量化、decode 整专家融合、prefill M8N32/M16N32 gather/vector +路径和 NUMA 分派。原来的 rows、splits、quant、merge 等细项开关用于调试和对照测试， +日常启动无需逐项配置；graph、CPU input quant 和诊断 profiling 默认关闭。 + +上面另两个 `SGLANG_` 变量属于本机 SGLang 运行环境设置：`SGLANG_DISABLE_CUDNN_CHECK=1` +用于绕过当前 CUDA/cuDNN 组合的版本检查，`SGLANG_MAMBA_CONV_DTYPE=float16` 为本模型 +选择 FP16 Mamba 卷积。这两个设置沿用本机已验证环境，不属于 XMX 的分块调优参数。 +`MODEL` 只是减少路径重复的 shell 变量，无需 export。 + +单 Intel GPU 环境会优先选择 Level Zero 设备，无需指定 `ONEAPI_DEVICE_SELECTOR`； +多 SYCL GPU 时可显式指定设备。 + +推荐配置在算子构造时确定，要求 XMX implementation revision≥7、hidden2048、intermediate512、 +group128；不兼容配置会明确报错。它不会修改进程的环境变量，显式设置的细项参数优先于配置 +默认值。首次使用请在未保留旧调优变量的新终端启动，避免旧实验设置覆盖推荐配置。 + +已验证服务配置为 GPU experts8、CPUInfer10、chunk512、最多一个请求；保持这些设置进行首次验收。 +`--disable-radix-cache` 用于让每次性能请求都实际执行完整 prefill，日常使用可按需求移除。 + +服务就绪后，可在另一个终端检查并发送请求： + +```bash +curl --fail http://127.0.0.1:30000/health +curl --fail http://127.0.0.1:30000/v1/completions \ + -H 'Content-Type: application/json' \ + -d '{"model":"Qwen3.5-35B-A3B-GPTQ-Int4","prompt":"请简要说明混合推理的作用。","max_tokens":64,"temperature":0}' +``` + +## 支持范围和后端选择 + +普通 `--kt-method SYCL_GPTQ_INT4` 仍默认使用 legacy 实现。 +设置 `KT_SYCL_XMX_PROFILE=panther-lake` 会自动选择新实现； +也可单独设置 `KT_SYCL_INT4_BACKEND=xmx` 使用通用执行器并自行调优。 +使用普通 SGLang 命令时，不设置推荐配置并设为 `KT_SYCL_INT4_BACKEND=legacy` 可切回原实现。 +推荐配置与显式 legacy 选择冲突时会报错，不会静默忽略配置。 + +权重必须为对称 GPTQ INT4、`desc_act=false`。通用执行器要求 group size 为正的32倍数， +投影维度满足整除和对齐约束；上面的融合调优配置还要求 hidden2048、intermediate512、group128。 +不支持的模式和配置会明确报错。专家放置 mask 在加载时确定，不支持运行中动态迁移专家， +也不支持将同一专家拆到多个 NUMA part。 diff --git a/doc/zh/SYCL-GPTQ-INT4-Tutorial_zh.md b/doc/zh/SYCL-GPTQ-INT4-Tutorial_zh.md index 8c29ec94..49506264 100644 --- a/doc/zh/SYCL-GPTQ-INT4-Tutorial_zh.md +++ b/doc/zh/SYCL-GPTQ-INT4-Tutorial_zh.md @@ -4,6 +4,9 @@ 这是一个异构推理方案：Intel iGPU 负责未放在 CUDA GPU 上的 GPTQ INT4 experts；Attention、Embedding、LM Head 以及通过 `--kt-num-gpu-experts` 保留的 experts 仍由 CUDA GPU 执行。因此，当前实现不能替代 NVIDIA GPU，也不是纯 Intel GPU 推理后端。 +Panther Lake XMX 的最终使用说明见 [Intel iGPU XMX 后端](Intel-iGPU_zh.md)， +源码构建步骤见 [环境安装](Intel-iGPU-Installation_zh.md)。 + ## 目录 - [支持范围](#支持范围) diff --git a/install.sh b/install.sh index db1343e9..089ccdbb 100755 --- a/install.sh +++ b/install.sh @@ -18,6 +18,7 @@ SUBCOMMANDS: -h, --help Show this help message OPTIONS: + --sycl Build Intel SYCL experts with the tested local configuration --skip-sglang Skip sglang installation (for "all" subcommand) --skip-kt-kernel Skip kt-kernel installation (for "all" subcommand) --editable Install sglang in editable/dev mode (-e) @@ -28,6 +29,9 @@ EXAMPLES: # Full install (recommended) $0 + # Intel iGPU SYCL source install + $0 --sycl + # Install everything in editable mode for development $0 all --editable @@ -83,7 +87,7 @@ read_kt_version() { init_submodules() { log_step "Initializing git submodules" - if [ ! -d "$REPO_ROOT/.git" ]; then + if [ ! -e "$REPO_ROOT/.git" ]; then log_warn "Not a git repository. Skipping submodule init." log_warn "If you need sglang, clone with: git clone --recursive https://github.com/kvcache-ai/ktransformers.git" return 0 @@ -94,6 +98,41 @@ init_submodules() { log_info "Submodules initialized successfully." } +# Configure the optional Intel build once; the ordinary install path is unchanged. +configure_sycl() { + local name + for name in CPUINFER_USE_CUDA CPUINFER_USE_ROCM CPUINFER_USE_MUSA CPUINFER_USE_MACA CPUINFER_USE_ASCEND_NPU; do + case "${!name:-0}" in + 1|ON|on|true|TRUE|yes|YES) + log_error "--sycl conflicts with $name=${!name}; unset it before installing." + exit 1 + ;; + esac + done + local setvars="${ONEAPI_SETVARS:-/opt/intel/oneapi/setvars.sh}" + if [ -r "$setvars" ]; then + # Vendor environment scripts are not guaranteed to support errexit/nounset. + # Validate the compiler after initialization instead of treating their final + # shell status as proof that the toolchain was initialized. + set +eu + source "$setvars" --force + set -eu + elif ! command -v icpx >/dev/null 2>&1; then + log_error "Install oneAPI DPC++ or set ONEAPI_SETVARS to its setvars.sh." + exit 1 + fi + command -v icpx >/dev/null 2>&1 || { log_error "oneAPI initialization did not provide icpx"; exit 1; } + export CPUINFER_USE_SYCL=1 CPUINFER_USE_CUDA=0 + export CPUINFER_CPU_INSTRUCT="${CPUINFER_CPU_INSTRUCT:-AVX2}" + export CPUINFER_ENABLE_AMX="${CPUINFER_ENABLE_AMX:-OFF}" + export CPUINFER_ENABLE_AVX512="${CPUINFER_ENABLE_AVX512:-OFF}" + export CPUINFER_ENABLE_MLA="${CPUINFER_ENABLE_MLA:-OFF}" + export CPUINFER_ENABLE_CPPTRACE="${CPUINFER_ENABLE_CPPTRACE:-OFF}" + export CPUINFER_ENABLE_LTO="${CPUINFER_ENABLE_LTO:-OFF}" + export CPUINFER_LLAMA_BUILD_COPY="${CPUINFER_LLAMA_BUILD_COPY:-ON}" + export CPUINFER_PARALLEL="${CPUINFER_PARALLEL:-2}" +} + # ─── sglang install ─────────────────────────────────────────────────────────── install_sglang() { @@ -110,14 +149,39 @@ install_sglang() { exit 1 fi + local install_target="./python[all]" + if [ "$SYCL_INSTALL" = "1" ]; then + # Build the local kernel first, then resolve SGLang against that exact source + # version. Never point users at a private submodule commit or replace the + # checked-out SGLang source with generated packaging edits. + python3 -c ' +import os +from importlib.metadata import version +from kt_kernel import kt_kernel_ext as e + +if e.moe.SYCLXMXGPTQInt4_MOE.implementation_revision() < 7: + raise RuntimeError("Build kt-kernel first with ./install.sh --sycl") +if version("kt-kernel") != os.environ["SGLANG_KT_VERSION"]: + raise RuntimeError("Installed kt-kernel version does not match this checkout; rebuild with ./install.sh kt-kernel --sycl") +' + mkdir -p "$REPO_ROOT/build" + local prepared + prepared=$(mktemp -d "$REPO_ROOT/build/sglang-sycl.XXXXXX") + python3 "$REPO_ROOT/kt-kernel/scripts/prepare_sglang.py" --output "$prepared/source" + sglang_dir="$prepared/source" + # Routed-expert inference needs the base runtime; diffusion/tracing extras + # remain part of the ordinary all-platform install. + install_target="./python" + fi + cd "$sglang_dir" if [ "$editable" = "1" ]; then log_info "Installing sglang in editable mode..." - pip install -e "./python[all]" + python3 -m pip install -e "$install_target" else log_info "Installing sglang..." - pip install "./python[all]" + python3 -m pip install "$install_target" fi log_info "sglang installed successfully." @@ -139,7 +203,14 @@ install_kt_kernel() { fi cd "$REPO_ROOT/kt-kernel" + if [ "$SYCL_INSTALL" = "1" ]; then + kt_args+=("--manual") + fi bash ./install.sh build "${kt_args[@]}" + if [ "$SYCL_INSTALL" = "1" ]; then + log_info "Checking installed SYCL extension without LD_LIBRARY_PATH" + env -u LD_LIBRARY_PATH -u LD_PRELOAD python3 -c 'from kt_kernel import kt_kernel_ext as e; assert e.moe.SYCLXMXGPTQInt4_MOE.implementation_revision() >= 7, "Installed XMX extension is too old"; print("SYCL runtime library lookup passed")' + fi } # ─── deps install ───────────────────────────────────────────────────────────── @@ -190,6 +261,14 @@ install_all() { # 3. Read version for sglang-kt read_kt_version + if [ "$SYCL_INSTALL" = "1" ]; then + # The source-aligned SGLang metadata depends on this kernel version. + if [ "$skip_kt_kernel" = "0" ]; then install_kt_kernel "${kt_args[@]}"; fi + if [ "$skip_sglang" = "0" ]; then install_sglang "$editable"; fi + log_step "SYCL source installation complete!" + return + fi + # 4. Install sglang if [ "$skip_sglang" = "0" ]; then install_sglang "$editable" @@ -211,6 +290,13 @@ install_all() { # ─── Subcommand dispatcher ──────────────────────────────────────────────────── +SYCL_INSTALL=0 +case "${CPUINFER_USE_SYCL:-0}" in 1|ON|on|true|TRUE|yes|YES) SYCL_INSTALL=1 ;; esac +INSTALL_ARGS=() +for arg in "$@"; do + if [ "$arg" = "--sycl" ]; then SYCL_INSTALL=1; else INSTALL_ARGS+=("$arg"); fi +done +set -- "${INSTALL_ARGS[@]}" SUBCMD="all" if [[ $# -gt 0 ]]; then case "$1" in @@ -232,6 +318,8 @@ if [[ $# -gt 0 ]]; then esac fi +if [ "$SYCL_INSTALL" = "1" ]; then configure_sycl; fi + case "$SUBCMD" in all) install_all "$@" diff --git a/kt-kernel/CMakeLists.txt b/kt-kernel/CMakeLists.txt index a1f340b0..1a19d1ca 100644 --- a/kt-kernel/CMakeLists.txt +++ b/kt-kernel/CMakeLists.txt @@ -504,6 +504,16 @@ add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/../third_party/pybind11 ${CMAKE_CUR # but succeeds `git apply --reverse --check`. # --------------------------------------------------------------------------- set(KT_LLAMA_CPP_DIR ${CMAKE_CURRENT_SOURCE_DIR}/../third_party/llama.cpp) +option(KTRANSFORMERS_LLAMA_BUILD_COPY "Apply vendored llama.cpp patches to a build-local copy" OFF) +if(KTRANSFORMERS_LLAMA_BUILD_COPY) + set(KT_LLAMA_CPP_COPY_ROOT ${CMAKE_CURRENT_BINARY_DIR}/patched_deps) + if(NOT EXISTS "${KT_LLAMA_CPP_COPY_ROOT}/llama.cpp/CMakeLists.txt") + file(COPY "${KT_LLAMA_CPP_DIR}" DESTINATION "${KT_LLAMA_CPP_COPY_ROOT}" PATTERN ".git" EXCLUDE) + endif() + set(KT_LLAMA_CPP_DIR "${KT_LLAMA_CPP_COPY_ROOT}/llama.cpp") + # Headers in the patched copy must precede the original third_party path. + include_directories(BEFORE "${KT_LLAMA_CPP_COPY_ROOT}") +endif() set(KT_LLAMA_CPP_PATCH_DIR ${CMAKE_CURRENT_SOURCE_DIR}/third_party_patches/llama.cpp) file(GLOB KT_LLAMA_CPP_PATCHES "${KT_LLAMA_CPP_PATCH_DIR}/*.patch") list(SORT KT_LLAMA_CPP_PATCHES) @@ -515,16 +525,24 @@ if(KT_LLAMA_CPP_PATCHES) endif() endif() +set(KT_LLAMA_GIT_COMMAND ${GIT_EXECUTABLE}) +if(KTRANSFORMERS_LLAMA_BUILD_COPY) + # Otherwise git discovers the outer repository and can silently skip + # patches whose paths are inside the ignored build directory. + set(KT_LLAMA_GIT_COMMAND ${CMAKE_COMMAND} -E env + "GIT_CEILING_DIRECTORIES=${KT_LLAMA_CPP_COPY_ROOT}" ${GIT_EXECUTABLE}) +endif() + foreach(KT_PATCH IN LISTS KT_LLAMA_CPP_PATCHES) execute_process( - COMMAND ${GIT_EXECUTABLE} apply --check "${KT_PATCH}" + COMMAND ${KT_LLAMA_GIT_COMMAND} apply --check "${KT_PATCH}" WORKING_DIRECTORY ${KT_LLAMA_CPP_DIR} RESULT_VARIABLE KT_PATCH_CHECK OUTPUT_QUIET ERROR_QUIET) if(KT_PATCH_CHECK EQUAL 0) execute_process( - COMMAND ${GIT_EXECUTABLE} apply "${KT_PATCH}" + COMMAND ${KT_LLAMA_GIT_COMMAND} apply "${KT_PATCH}" WORKING_DIRECTORY ${KT_LLAMA_CPP_DIR} RESULT_VARIABLE KT_PATCH_APPLY ERROR_VARIABLE KT_PATCH_APPLY_ERR) @@ -535,7 +553,7 @@ foreach(KT_PATCH IN LISTS KT_LLAMA_CPP_PATCHES) message(STATUS "llama.cpp: applied patch ${KT_PATCH}") else() execute_process( - COMMAND ${GIT_EXECUTABLE} apply --reverse --check "${KT_PATCH}" + COMMAND ${KT_LLAMA_GIT_COMMAND} apply --reverse --check "${KT_PATCH}" WORKING_DIRECTORY ${KT_LLAMA_CPP_DIR} RESULT_VARIABLE KT_PATCH_REVERSE_CHECK OUTPUT_QUIET @@ -543,6 +561,11 @@ foreach(KT_PATCH IN LISTS KT_LLAMA_CPP_PATCHES) if(KT_PATCH_REVERSE_CHECK EQUAL 0) message(STATUS "llama.cpp: patch already applied, skipping ${KT_PATCH}") else() + if(KTRANSFORMERS_LLAMA_BUILD_COPY) + message(FATAL_ERROR + "llama.cpp patch ${KT_PATCH} neither applies nor is already applied to the build copy.\n" + "Remove ${KT_LLAMA_CPP_DIR} and reconfigure to refresh it from the submodule.") + endif() message(FATAL_ERROR "llama.cpp patch ${KT_PATCH} neither applies nor is already applied.\n" "third_party/llama.cpp is probably not at the pinned commit or has local edits.\n" @@ -553,7 +576,7 @@ foreach(KT_PATCH IN LISTS KT_LLAMA_CPP_PATCHES) endif() endforeach() -add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/../third_party/llama.cpp ${CMAKE_CURRENT_BINARY_DIR}/third_party/llama.cpp) +add_subdirectory(${KT_LLAMA_CPP_DIR} ${CMAKE_CURRENT_BINARY_DIR}/third_party/llama.cpp) include_directories(${CMAKE_CURRENT_SOURCE_DIR}/../third_party) if(KTRANSFORMERS_USE_CUDA) @@ -796,6 +819,8 @@ else() message(STATUS "LTO: disabled") endif() +target_link_libraries(${PROJECT_NAME} PRIVATE ${CMAKE_DL_LIBS}) + if(KTRANSFORMERS_CPU_USE_ONEDNN_VNNI) if(KTRANSFORMERS_ONEDNN_SOURCE_DIR) if(NOT EXISTS "${KTRANSFORMERS_ONEDNN_SOURCE_DIR}/CMakeLists.txt") @@ -878,6 +903,10 @@ if(TARGET ${PROJECT_NAME} AND DEFINED ENV{CONDA_PREFIX} AND EXISTS "$ENV{CONDA_P SKIP_BUILD_RPATH OFF ) endif() +if(KTRANSFORMERS_USE_SYCL) + include(cmake/SYCLRuntime.cmake) + kt_add_sycl_runtime_rpath(${PROJECT_NAME}) +endif() if(NOT HOST_IS_X86 AND KTRANSFORMERS_CPU_USE_KML) message(STATUS "KML CPU detected") diff --git a/kt-kernel/bench/bench_sycl_xmx_moe.py b/kt-kernel/bench/bench_sycl_xmx_moe.py new file mode 100644 index 00000000..61a09c22 --- /dev/null +++ b/kt-kernel/bench/bench_sycl_xmx_moe.py @@ -0,0 +1,308 @@ +#!/usr/bin/env python3 +"""Same-input GPTQ INT4 MoE correctness/latency benchmark (requires a SYCL build). + +Run under the kt-kernel Python environment after sourcing oneAPI. No HTTP/token +count estimates: times include CPUInfer submit, complete MoE execution and sync. +""" + +import argparse +import gc +import json +import os +import statistics +import time +from pathlib import Path + +import torch +from kt_kernel import kt_kernel_ext as ext + + +def packed_weights(experts, n, k, gs, generator): + # Generate stored GPTQ nibbles directly, without a resident FP32 weight copy. + q = torch.randint(0, 16, (experts, k // 8, n, 8), generator=generator, dtype=torch.int64) + shifts = torch.arange(8) * 4 + packed = (q << shifts).sum(-1).to(torch.int32).contiguous() + scales = (torch.rand(experts, k // gs, n, generator=generator) * 0.003 + 0.002).contiguous() + return packed, scales + + +def dequant(packed, scales, gs): + k8, n = packed.shape + shifts = torch.arange(8) * 4 + values = ((packed.long().unsqueeze(-1) >> shifts) & 15) - 8 + return values.permute(0, 2, 1).reshape(k8 * 8, n).float() * scales.repeat_interleave(gs, 0) + + +def reference(x, ids, routing, projections, gs, mask=None, physical_map=None, limit=0.0, alpha=0.0): + result = torch.zeros_like(x, dtype=torch.float32) + for e in range(projections[0][0].shape[0]): + logical = int(physical_map[e]) if physical_map is not None else e + if mask is not None and mask[logical]: + continue + where = (ids == e).nonzero() + if not len(where): + continue + tokens, slots = where[:, 0], where[:, 1] + source = x[tokens].float() + gate, up, down = [dequant(p[0][logical], p[1][logical], gs) for p in projections] + # Baseline converts gate/up to BF16 before activation, then down to BF16. + g = (source @ gate).bfloat16().float() + u = (source @ up).bfloat16().float() + if limit > 0: + g = g.clamp(-limit, limit) if alpha > 0 else g.clamp(max=limit) + u = u.clamp(-limit, limit) + mid = ( + (g * torch.sigmoid(g * alpha) * (u + 1) if alpha > 0 else torch.nn.functional.silu(g) * u) + .bfloat16() + .float() + ) + y = (mid @ down).bfloat16().float() * routing[tokens, slots, None] + result.index_add_(0, tokens, y) + return result.bfloat16().float() + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--backends", default="legacy,packed,xmx-fp16,xmx-int8") + p.add_argument("--rows", default="1,16,64,512") + p.add_argument("--experts", type=int, default=256) + p.add_argument("--hidden", type=int, default=2048) + p.add_argument("--intermediate", type=int, default=512) + p.add_argument("--topk", type=int, default=8) + p.add_argument("--group-size", type=int, default=128) + p.add_argument("--model-path", type=Path, help="Use the actual GPTQ model weights instead of random weights") + p.add_argument("--layer", type=int, default=0) + p.add_argument("--threads", type=int, default=10) + p.add_argument("--repeat", type=int, default=10) + p.add_argument("--warmup", type=int, default=3) + p.add_argument("--route-sets", type=int, default=1, help="Rotate expert assignments; 1 measures fixed-route reuse") + p.add_argument("--swiglu-limit", type=float, default=0.0) + p.add_argument("--swiglu-alpha", type=float, default=0.0) + p.add_argument("--check", action="store_true") + p.add_argument( + "--edge-cases", action="store_true", help="duplicate/invalid IDs, mask, permutation and incremental output" + ) + p.add_argument("--output", type=Path) + args = p.parse_args() + backends = args.backends.split(",") + if not set(backends) <= { + "legacy", + "packed", + "avx2", + "xmx-fp16", + "xmx-int8", + "vector", + "dp4a", + "xmx-dp4a", + "dp4a-xmx", + "vector-xmx", + "xmx-vector", + "xmx-split-xmx", + }: + p.error("Unknown backend") + if args.repeat <= 0 or args.warmup < 0 or args.route_sets <= 0 or not 0 < args.topk <= args.experts: + p.error("Invalid repetitions or expert count") + if min(args.hidden, args.intermediate, args.group_size, args.threads) <= 0: + p.error("Dimensions, group size and thread count must be positive") + if args.hidden % args.group_size or args.intermediate % args.group_size: + p.error("Projection dimensions must be divisible by group size") + torch.set_num_threads(4) + generator = torch.Generator().manual_seed(20261014) + rows_list = [int(n) for n in args.rows.split(",")] + if not rows_list or min(rows_list) <= 0: + p.error("Rows must be positive") + if args.model_path: + from safetensors import safe_open + + index = json.loads((args.model_path / "model.safetensors.index.json").read_text())["weight_map"] + actual = [] + for name in ("gate", "up", "down"): + tensors = [[], []] + for expert in range(args.experts): + stem = f"model.language_model.layers.{args.layer}.mlp.experts.{expert}.{name}_proj" + for dest, suffix in zip(tensors, ("qweight", "scales")): + key = stem + "." + suffix + with safe_open(args.model_path / index[key], framework="pt", device="cpu") as f: + t = f.get_tensor(key) + dest.append(t.float().contiguous() if suffix == "scales" else t.clone()) + actual.append(tuple(torch.stack(ts) for ts in tensors)) + proj = actual + else: + proj = [ + packed_weights(args.experts, args.intermediate, args.hidden, args.group_size, generator), + packed_weights(args.experts, args.intermediate, args.hidden, args.group_size, generator), + packed_weights(args.experts, args.hidden, args.intermediate, args.group_size, generator), + ] + cases = [] + for rows in rows_list: + x = (torch.randn(rows, args.hidden, generator=generator) * 0.5).bfloat16().contiguous() + ids = torch.rand(rows, args.experts, generator=generator).argsort(1)[:, : args.topk].contiguous() + routing = torch.softmax(torch.randn(rows, args.topk, generator=generator), 1).contiguous() + if args.edge_cases: + ids[:, 0] = -1 + if args.topk > 2: + ids[:, 2] = ids[:, 1] + cases.append((rows, x, ids, routing)) + mask = torch.zeros(args.experts, dtype=torch.bool) + mapping = torch.arange(args.experts, dtype=torch.int64) + if args.edge_cases: + mask[::4] = True + mapping = torch.randperm(args.experts, generator=generator) + refs = ( + { + r: reference(x, ids, w, proj, args.group_size, mask, mapping, args.swiglu_limit, args.swiglu_alpha) + for r, x, ids, w in cases + } + if args.check + else {} + ) + cpu = ext.CPUInfer(args.threads) + results = [] + for backend in backends: + kernels = { + "vector": ("vector", "vector"), + "dp4a": ("dp4a", "dp4a"), + "xmx-dp4a": ("xmx", "dp4a"), + "dp4a-xmx": ("dp4a", "xmx"), + "xmx-split-xmx": ("xmx-split", "xmx"), + "vector-xmx": ("vector", "xmx"), + "xmx-vector": ("xmx", "vector"), + }.get(backend, ("xmx", "xmx")) + os.environ["KT_SYCL_XMX_GATE_KERNEL"], os.environ["KT_SYCL_XMX_DOWN_KERNEL"] = kernels + os.environ["KT_SYCL_XMX_MODE"] = ( + "fp16" if backend in ("vector", "xmx-fp16", "legacy", "packed", "avx2") else "int8" + ) + cls = {"legacy": "SYCLGPTQInt4_MOE", "packed": "AVXVNNI256GPTQInt4Packed_MOE", "avx2": "AVX2GPTQInt4_MOE"}.get( + backend, "SYCLXMXGPTQInt4_MOE" + ) + config = ext.moe.MOEConfig(args.experts, args.topk, args.hidden, args.intermediate, mask.data_ptr()) + config.pool = cpu.backend_ + config.max_len = max(rows_list) + config.quant_config.bits = 4 + config.quant_config.group_size = args.group_size + config.quant_config.zero_point = False + config.swiglu_limit = args.swiglu_limit + config.swiglu_alpha = args.swiglu_alpha + for name, (qw, sc) in zip(("gate", "up", "down"), proj): + setattr(config, name + "_projs", [[t.data_ptr() for t in qw]]) + setattr(config, name + "_scales", [[t.data_ptr() for t in sc]]) + moe = getattr(ext.moe, cls)(config) + cpu.submit(moe.load_weights_task(mapping.data_ptr())) + cpu.sync() + for rows, x, ids, weights in cases: + out = torch.zeros_like(x) + qlen = torch.tensor([rows], dtype=torch.int32) + routes = [ + torch.where(ids >= 0, (ids + index * args.topk) % args.experts, ids).contiguous() + for index in range(args.route_sets) + ] + + def run(incremental=False, route=0): + cpu.submit( + moe.forward_task( + qlen.data_ptr(), + args.topk, + routes[route].data_ptr(), + weights.data_ptr(), + x.data_ptr(), + out.data_ptr(), + incremental, + ) + ) + cpu.sync() + + for _ in range(args.warmup): + run() + times = [] + for index in range(args.repeat): + start = time.perf_counter_ns() + run(route=index % args.route_sets) + times.append((time.perf_counter_ns() - start) / 1e6) + record = { + "backend": backend, + "rows": rows, + "median_ms": statistics.median(times), + "min_ms": min(times), + "samples_ms": times, + } + record["kernels"] = { + key: os.getenv(key) + for key in ( + "KT_SYCL_XMX_MODE", + "KT_SYCL_XMX_GATE_KERNEL", + "KT_SYCL_XMX_DOWN_KERNEL", + "KT_SYCL_XMX_DECODE_GATE_KERNEL", + "KT_SYCL_XMX_DECODE_DOWN_KERNEL", + "KT_SYCL_XMX_GATE_SPLITS", + "KT_SYCL_XMX_DOWN_SPLITS", + "KT_SYCL_XMX_DECODE_MERGE", + "KT_SYCL_XMX_TRACE", + "KT_SYCL_XMX_FUSED_EXPERT", + "KT_SYCL_XMX_PREFILL_ROWS", + "KT_SYCL_XMX_PREFILL_DOWN_ROWS", + "KT_SYCL_XMX_PREFILL_N", + "KT_SYCL_XMX_PREFILL_LOAD", + "KT_SYCL_XMX_PREFILL_MERGE", + ) + } + if args.check: + run() # Compare the base route, separately from timed requests. + ref = refs[rows] + error = out.float() - ref + record.update( + relative_rmse=(error.square().mean().sqrt() / ref.square().mean().sqrt().clamp_min(1e-12)).item(), + max_abs_error=error.abs().max().item(), + finite=bool(torch.isfinite(out).all()), + ) + quantized = backend not in ("legacy", "avx2", "vector", "xmx-fp16") + if not record["finite"] or record["relative_rmse"] > (0.08 if quantized else 0.025): + raise AssertionError(f"correctness failed: {record}") + if args.edge_cases and backend not in ("legacy", "packed", "avx2"): + previous = out.clone() + run(incremental=True) + expected = (previous.float() * 2).bfloat16() + torch.testing.assert_close(out, expected, rtol=0.02, atol=0.002) + print(json.dumps(record), flush=True) + results.append(record) + del moe + gc.collect() + if args.output: + args.output.parent.mkdir(parents=True, exist_ok=True) + config = vars(args) | { + "output": str(args.output), + "model_path": str(args.model_path) if args.model_path else None, + } + environment = { + key: os.getenv(key) + for key in ( + "KT_SYCL_XMX_MODE", + "KT_SYCL_XMX_ROWS", + "KT_SYCL_XMX_GRAPH", + "KT_SYCL_XMX_HOST_QUANT", + "KT_SYCL_XMX_FUSED_QUANT", + "KT_SYCL_XMX_NUMA_DISPATCH", + "KT_SYCL_XMX_GATE_KERNEL", + "KT_SYCL_XMX_DOWN_KERNEL", + "KT_SYCL_XMX_DECODE_GATE_KERNEL", + "KT_SYCL_XMX_DECODE_DOWN_KERNEL", + "KT_SYCL_XMX_VECTOR_SPLITS", + "KT_SYCL_XMX_GATE_SPLITS", + "KT_SYCL_XMX_DOWN_SPLITS", + "KT_SYCL_XMX_PREFILL_ROWS", + "KT_SYCL_XMX_PREFILL_DOWN_ROWS", + "KT_SYCL_XMX_PREFILL_N", + "KT_SYCL_XMX_PREFILL_LOAD", + "KT_SYCL_XMX_PREFILL_MERGE", + "KT_SYCL_XMX_DECODE_MERGE", + "KT_SYCL_XMX_FUSED_EXPERT", + "KT_SYCL_XMX_TRACE", + "KT_SYCL_XMX_TRACE_EVERY", + "KT_SYCL_PROFILE", + "ONEAPI_DEVICE_SELECTOR", + ) + } + args.output.write_text(json.dumps({"config": config, "environment": environment, "results": results}, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/kt-kernel/cmake/SYCLRuntime.cmake b/kt-kernel/cmake/SYCLRuntime.cmake new file mode 100644 index 00000000..31ebef89 --- /dev/null +++ b/kt-kernel/cmake/SYCLRuntime.cmake @@ -0,0 +1,35 @@ +# The oneAPI driver links its own runtime libraries implicitly. CMake omits +# implicit compiler directories from generated RPATHs, and the conda setup can +# additionally replace the target's install RPATH. Source-built SYCL extensions +# therefore need the selected toolchain's runtime directory recorded explicitly. +function(kt_add_sycl_runtime_rpath target) + if(NOT CMAKE_SYSTEM_NAME STREQUAL "Linux") + return() + endif() + + set(_kt_runtime_dirs) + foreach(_kt_library IN ITEMS libsycl.so libiomp5.so libsvml.so libimf.so libirng.so libintlc.so.5) + execute_process( + COMMAND "${CMAKE_CXX_COMPILER}" "-print-file-name=${_kt_library}" + RESULT_VARIABLE _kt_query_result + OUTPUT_VARIABLE _kt_library_path + OUTPUT_STRIP_TRAILING_WHITESPACE + ERROR_QUIET + ) + if(_kt_query_result EQUAL 0 AND IS_ABSOLUTE "${_kt_library_path}" AND EXISTS "${_kt_library_path}") + get_filename_component(_kt_library_path "${_kt_library_path}" REALPATH) + get_filename_component(_kt_runtime_dir "${_kt_library_path}" DIRECTORY) + list(APPEND _kt_runtime_dirs "${_kt_runtime_dir}") + elseif(_kt_library STREQUAL "libsycl.so" OR _kt_library STREQUAL "libiomp5.so") + message(FATAL_ERROR "SYCL compiler ${CMAKE_CXX_COMPILER} could not locate ${_kt_library}; initialize the matching oneAPI installation before building") + endif() + endforeach() + list(REMOVE_DUPLICATES _kt_runtime_dirs) + # Append after any conda target setup, preserving its existing search path. + # Set both properties because setuptools packages the build output directly, + # while BUILD_WITH_INSTALL_RPATH may already be enabled by the conda setup. + set_property(TARGET ${target} APPEND PROPERTY BUILD_RPATH "${_kt_runtime_dirs}") + set_property(TARGET ${target} APPEND PROPERTY INSTALL_RPATH "${_kt_runtime_dirs}") + set_property(TARGET ${target} PROPERTY SKIP_BUILD_RPATH OFF) + message(STATUS "SYCL runtime RPATH: ${_kt_runtime_dirs}") +endfunction() diff --git a/kt-kernel/ext_bindings.cpp b/kt-kernel/ext_bindings.cpp index 5296f859..772191b1 100644 --- a/kt-kernel/ext_bindings.cpp +++ b/kt-kernel/ext_bindings.cpp @@ -78,6 +78,7 @@ static const bool _is_plain_ = false; #endif #if defined(USE_SYCL) #include "operators/sycl/gptq_int4_sycl-moe.hpp" +#include "operators/sycl/gptq_int4_xmx-moe.hpp" #endif #include // std::vector/std::pair/std::string conversions @@ -482,6 +483,17 @@ void bind_moe_module(py::module_& moe_module, const char* name) { .def("load_weights", &MoeClass::load_weights) .def("forward", &MoeClass::forward_binding); + if constexpr (requires(MoeClass& instance) { + instance.profile_snapshot(); + instance.reset_profile(); + }) { + moe_cls.def("profile_snapshot", &MoeClass::profile_snapshot).def("reset_profile", &MoeClass::reset_profile); + } + + if constexpr (requires { MoeClass::implementation_revision(); }) { + moe_cls.def_static("implementation_revision", &MoeClass::implementation_revision); + } + // Bind write_weight_scale_to_buffer_task for MoE types that support it // Uses SFINAE to detect if MoeClass has write_weight_scale_to_buffer method if constexpr (requires { &MoeClass::write_weight_scale_to_buffer; }) { @@ -532,15 +544,14 @@ void bind_moe_module(py::module_& moe_module, const char* name) { std::uint64_t epoch, std::int64_t layer_id, int expert_count) { if (!transport) throw std::invalid_argument("FP8 layerwise transport is null"); const int tp_size = transport->tp_size(); - transport->run_producer( - epoch, layer_id, expert_count, - [moe, tp_size](int expert_id, const std::vector& w13_weight_ptrs, - const std::vector& w13_scale_ptrs, - const std::vector& w2_weight_ptrs, - const std::vector& w2_scale_ptrs) { - moe->write_weight_scale_to_buffer(tp_size, expert_id, w13_weight_ptrs, w13_scale_ptrs, w2_weight_ptrs, - w2_scale_ptrs); - }); + transport->run_producer(epoch, layer_id, expert_count, + [moe, tp_size](int expert_id, const std::vector& w13_weight_ptrs, + const std::vector& w13_scale_ptrs, + const std::vector& w2_weight_ptrs, + const std::vector& w2_scale_ptrs) { + moe->write_weight_scale_to_buffer(tp_size, expert_id, w13_weight_ptrs, + w13_scale_ptrs, w2_weight_ptrs, w2_scale_ptrs); + }); }, py::arg("transport"), py::arg("epoch"), py::arg("layer_id"), py::arg("expert_count"), py::call_guard()); @@ -587,18 +598,17 @@ PYBIND11_MODULE(kt_kernel_ext, m) { m.attr("FP8_LAYERWISE_CONTROL_BYTES") = kt::layerwise::kFP8LayerwiseControlBytes; m.attr("FP8_LAYERWISE_MAX_TP_SIZE") = kt::layerwise::kFP8LayerwiseMaxTPSize; - m.def("initialize_fp8_layerwise_control", &kt::layerwise::initialize_fp8_layerwise_control, - py::arg("control_ptr"), py::arg("control_size"), py::arg("tp_size")); + m.def("initialize_fp8_layerwise_control", &kt::layerwise::initialize_fp8_layerwise_control, py::arg("control_ptr"), + py::arg("control_size"), py::arg("tp_size")); py::class_>( m, "FP8LayerwiseTransport") .def(py::init&, const std::vector&, const std::vector&, const std::vector&, int, std::uint64_t>(), - py::arg("control_ptr"), py::arg("control_size"), py::arg("rank"), py::arg("tp_size"), - py::arg("cuda_device"), py::arg("local_host_ptrs"), py::arg("local_gpu_ptrs"), - py::arg("all_rank_host_ptrs"), py::arg("expert_nbytes"), py::arg("num_experts"), - py::arg("timeout_ms") = 60000) + py::arg("control_ptr"), py::arg("control_size"), py::arg("rank"), py::arg("tp_size"), py::arg("cuda_device"), + py::arg("local_host_ptrs"), py::arg("local_gpu_ptrs"), py::arg("all_rank_host_ptrs"), + py::arg("expert_nbytes"), py::arg("num_experts"), py::arg("timeout_ms") = 60000) .def("join", &kt::layerwise::FP8LayerwiseTransport::join, py::arg("epoch"), py::arg("layer_id"), py::arg("expert_count"), py::call_guard()) .def( @@ -634,16 +644,14 @@ PYBIND11_MODULE(kt_kernel_ext, m) { #if defined(KTRANSFORMERS_USE_ASCEND_NPU) m.def( - "init_ascend_callback_worker", - []() { kt::ascend::ensure_callback_worker(nullptr); }, + "init_ascend_callback_worker", []() { kt::ascend::ensure_callback_worker(nullptr); }, "Start ACL aclrtProcessReport worker for stream callbacks (Ascend NPU)."); m.def( "subscribe_ascend_stream", [](intptr_t stream_handle) { kt::ascend::ensure_stream_subscribed(reinterpret_cast(stream_handle)); }, - py::arg("stream_handle"), - "Subscribe an aclrtStream with the global callback worker."); + py::arg("stream_handle"), "Subscribe an aclrtStream with the global callback worker."); m.def("shutdown_ascend_callback_worker", &kt::ascend::shutdown_callback_worker, "Stop the ACL callback worker thread."); m.def("is_ascend_callback_worker_running", &kt::ascend::callback_worker_running, @@ -977,8 +985,7 @@ PYBIND11_MODULE(kt_kernel_ext, m) { #if defined(__AVX512BF16__) // SFT MoE with LoRA support (BF16, INT8, INT4, AWQ, K2) bind_moe_sft_module>(moe_module, "AMXBF16_SFT_MOE"); - bind_moe_sft_module>( - moe_module, "MXFP4_SFT_MOE"); + bind_moe_sft_module>(moe_module, "MXFP4_SFT_MOE"); #if defined(__AVX512VBMI__) bind_moe_sft_module>(moe_module, "AMXFP8_SFT_MOE"); #endif @@ -1012,8 +1019,7 @@ PYBIND11_MODULE(kt_kernel_ext, m) { // AVX2 tier of the MXFP4 routed-expert LoRA SFT (same Python name as the AMX build) bind_moe_sft_module>(moe_module, "MXFP4_SFT_MOE"); // AVX2 tier of the RAWINT4 routed-expert LoRA SFT (same Python name as the AMX build) - bind_moe_sft_module>(moe_module, - "AMXInt4_KGroup_SFT_MOE"); + bind_moe_sft_module>(moe_module, "AMXInt4_KGroup_SFT_MOE"); #endif bind_moe_module>(moe_module, "AVX2MXFP8_MOE"); bind_moe_module>(moe_module, @@ -1024,6 +1030,7 @@ PYBIND11_MODULE(kt_kernel_ext, m) { #endif #if defined(USE_SYCL) bind_moe_module>(moe_module, "SYCLGPTQInt4_MOE"); + bind_moe_module(moe_module, "SYCLXMXGPTQInt4_MOE"); #endif #if defined(USE_MOE_KERNEL) diff --git a/kt-kernel/install.sh b/kt-kernel/install.sh index 214348b6..a80269a5 100755 --- a/kt-kernel/install.sh +++ b/kt-kernel/install.sh @@ -89,6 +89,10 @@ EOF install_dependencies() { echo "Checking and installing system dependencies..." + if [ "${CPUINFER_USE_SYCL:-0}" = "1" ] && command -v cmake >/dev/null 2>&1 && command -v pkg-config >/dev/null 2>&1 && pkg-config --exists hwloc numa; then + echo "CMake, hwloc and NUMA development dependencies are already installed." + return 0 + fi # Determine if we need to use sudo SUDO="" diff --git a/kt-kernel/operators/sycl/gptq_int4_sycl-runtime.hpp b/kt-kernel/operators/sycl/gptq_int4_sycl-runtime.hpp index f95e52d0..e1284bf9 100644 --- a/kt-kernel/operators/sycl/gptq_int4_sycl-runtime.hpp +++ b/kt-kernel/operators/sycl/gptq_int4_sycl-runtime.hpp @@ -7,6 +7,7 @@ #include #include +#include #include #include #include @@ -57,6 +58,10 @@ inline sycl::queue& queue() { (void)async_error_state(); static sycl::queue q([] { auto async_handler = [](sycl::exception_list exceptions) { async_error_state().capture(std::move(exceptions)); }; + const char* profile = std::getenv("KT_SYCL_PROFILE"); + const sycl::property_list properties = profile && std::string(profile) == "1" + ? sycl::property_list{sycl::property::queue::enable_profiling{}} + : sycl::property_list{}; try { const auto devices = sycl::device::get_devices(sycl::info::device_type::gpu); @@ -65,10 +70,10 @@ inline sycl::queue& queue() { } for (const auto& device : devices) { if (device.get_platform().get_backend() == sycl::backend::ext_oneapi_level_zero) { - return sycl::queue(device, async_handler); + return sycl::queue(device, async_handler, properties); } } - return sycl::queue(devices.front(), async_handler); + return sycl::queue(devices.front(), async_handler, properties); } catch (const sycl::exception& error) { throw std::runtime_error(std::string("Failed to create the SYCL GPTQ INT4 queue. Check sycl-ls, " "the render-group permission, or ONEAPI_DEVICE_SELECTOR. Original error: ") + diff --git a/kt-kernel/operators/sycl/gptq_int4_vector-kernels.hpp b/kt-kernel/operators/sycl/gptq_int4_vector-kernels.hpp new file mode 100644 index 00000000..d3af28c8 --- /dev/null +++ b/kt-kernel/operators/sycl/gptq_int4_vector-kernels.hpp @@ -0,0 +1,166 @@ +// SPDX-License-Identifier: Apache-2.0 +#ifndef KT_SYCL_GPTQ_INT4_VECTOR_KERNELS_HPP +#define KT_SYCL_GPTQ_INT4_VECTOR_KERNELS_HPP + +#include +#include + +#include "gptq_int4_xmx-kernels.hpp" + +namespace sycl_xmx { + +// GPTQ nibbles use offset binary, not two's-complement S4. Expand four +// nibbles to signed bytes without carries between the packed byte lanes. +ESIMD_INLINE es::simd unpack_s4_bytes(es::simd packed) { + auto v = packed & 0xffffu; + v = (v | (v << 8)) & 0x00ff00ffu; + v = (v | (v << 4)) & 0x0f0f0f0fu; + auto negative = (~v) & 0x08080808u; + v = (v & 0x07070707u) | negative | (negative << 1) | (negative << 2) | (negative << 3) | (negative << 4); + return v.template bit_cast_view(); +} + +ESIMD_INLINE es::simd round_bf16(es::simd value) { + es::simd bits = value.template bit_cast_view(); + bits += 0x7fffu + ((bits >> 16) & 1u); + bits &= 0xffff0000u; + return bits.template bit_cast_view(); +} + +// All candidates read the same N16/K8 layout. Lanes own adjacent output +// columns; work-items split K groups, then sum their FP32 partials in SLM. +// This vector path does not issue DPAS or pad an activation into a matrix. +template +inline sycl::event vector_projection(const Tile* tiles, int tile_count, int tile_rows, const int* row_tokens, + const uint16_t* input, const int8_t* qinput, const float* ascales, + const Weight& first, const Weight* second, uint16_t* activated, float* down, + float limit, float alpha, const sycl::event* dep = nullptr) { + static_assert(Splits == 1 || Splits == 2 || Splits == 4 || Splits == 8); + const int n = first.n, k = first.k, gs = first.gs, groups = k / gs; + const uint32_t* w0 = first.packed.data; + const uint32_t* w1 = second ? second->packed.data : nullptr; + const float* s0 = first.scales.data; + const float* s1 = second ? second->scales.data : nullptr; + const size_t workgroups = static_cast(tile_count) * tile_rows * (n / 16); + return sycl_int4::queue().submit([&](sycl::handler& h) { + if (dep) h.depends_on(*dep); + h.parallel_for(sycl::nd_range<1>(workgroups * Splits, Splits), [=](sycl::nd_item<1> it) SYCL_ESIMD_KERNEL { + if constexpr (Splits > 1) es::slm_init(); + const size_t task = it.get_group_linear_id(); + const int nb = task % (n / 16), split = it.get_local_linear_id(); + const int r = (task / (n / 16)) % tile_rows; + const Tile tile = tiles[task / (static_cast(n / 16) * tile_rows)]; + // Every work-item in this work-group takes the same branch. + if (r >= tile.rows) return; + const int row = GateUp ? row_tokens[tile.begin + r] : tile.begin + r; + const size_t wi = (static_cast(tile.expert) * (n / 16) + nb) * (k / 8) * 16; + const size_t si = (static_cast(tile.expert) * (n / 16) + nb) * groups * 16; + es::simd acc0 = 0.0f, acc1 = 0.0f; + for (int g = split; g < groups; g += Splits) { + es::simd c0 = 0, c1 = 0; + es::simd f0 = 0.0f, f1 = 0.0f; + for (int j = 0; j < gs; j += 8) { + es::simd b0; + b0.copy_from(w0 + wi + (g * gs + j) / 8 * 16); + es::simd b1; + if constexpr (GateUp) b1.copy_from(w1 + wi + (g * gs + j) / 8 * 16); + if constexpr (Q8) { + es::simd a; + a.copy_from(qinput + static_cast(row) * k + g * gs + j); + es::simd av = a.template bit_cast_view(); + es::simd al = static_cast(av[0]), ah = static_cast(av[1]); + c0 = es::dp4a(c0, unpack_s4_bytes(b0), al); + c0 = es::dp4a(c0, unpack_s4_bytes(b0 >> 16), ah); + if constexpr (GateUp) { + c1 = es::dp4a(c1, unpack_s4_bytes(b1), al); + c1 = es::dp4a(c1, unpack_s4_bytes(b1 >> 16), ah); + } + } else { + es::simd bits; + bits.copy_from(input + static_cast(row) * k + g * gs + j); + es::simd u = es::convert(bits) << 16; + es::simd a = u.template bit_cast_view(); +#pragma unroll + for (int x = 0; x < 8; ++x) { + es::simd av = static_cast(a[x]); + auto v0 = es::convert((b0 >> (x * 4)) & 15) - 8.0f; + f0 = sycl::ext::intel::experimental::esimd::fma(av, v0, f0); + if constexpr (GateUp) { + auto v1 = es::convert((b1 >> (x * 4)) & 15) - 8.0f; + f1 = sycl::ext::intel::experimental::esimd::fma(av, v1, f1); + } + } + } + } + es::simd scale; + scale.copy_from(s0 + si + g * 16); + if constexpr (Q8) { + float sa = ascales[row * groups + g]; + acc0 += es::convert(c0) * (scale * sa); + if constexpr (GateUp) { + scale.copy_from(s1 + si + g * 16); + acc1 += es::convert(c1) * (scale * sa); + } + } else { + acc0 += f0 * scale; + if constexpr (GateUp) { + scale.copy_from(s1 + si + g * 16); + acc1 += f1 * scale; + } + } + } + if constexpr (Splits > 1) { + es::slm_block_store(split * 16 * sizeof(float), acc0); + if constexpr (GateUp) es::slm_block_store((Splits + split) * 16 * sizeof(float), acc1); + es::barrier(); + if (split != 0) return; + acc0 = 0.0f; + acc1 = 0.0f; +#pragma unroll + for (int s = 0; s < Splits; ++s) { + acc0 += es::slm_block_load(s * 16 * sizeof(float)); + if constexpr (GateUp) acc1 += es::slm_block_load((Splits + s) * 16 * sizeof(float)); + } + } + if constexpr (GateUp) { + if (limit > 0) { + acc0 = es::min(acc0, limit); + if (alpha > 0) acc0 = es::max(acc0, -limit); + acc1 = es::min(es::max(acc1, -limit), limit); + } + acc0 = acc0 / (1.0f + es::exp(-acc0 * (alpha > 0 ? alpha : 1.0f))) * (alpha > 0 ? acc1 + 1.0f : acc1); + auto rounded = round_bf16(acc0); + es::simd bits = rounded.template bit_cast_view(); + es::simd bf = es::convert(bits >> 16); + bf.copy_to(activated + static_cast(tile.begin + r) * n + nb * 16); + } else { + round_bf16(acc0).copy_to(down + static_cast(tile.begin + r) * n + nb * 16); + } + }); + }); +} + +template +inline sycl::event vector_dispatch(int splits, const Tile* tiles, int tile_count, int tile_rows, const int* row_tokens, + const uint16_t* input, const int8_t* qinput, const float* ascales, + const Weight& first, const Weight* second, uint16_t* activated, float* down, + float limit, float alpha, const sycl::event* dep = nullptr) { +#define KT_VECTOR_DISPATCH(S) \ + vector_projection(tiles, tile_count, tile_rows, row_tokens, input, qinput, ascales, first, second, \ + activated, down, limit, alpha, dep) + switch (splits) { + case 1: + return KT_VECTOR_DISPATCH(1); + case 2: + return KT_VECTOR_DISPATCH(2); + case 4: + return KT_VECTOR_DISPATCH(4); + case 8: + return KT_VECTOR_DISPATCH(8); + default: + throw std::runtime_error("Vector K splits must be 1, 2, 4 or 8"); + } +#undef KT_VECTOR_DISPATCH +} +} // namespace sycl_xmx +#endif diff --git a/kt-kernel/operators/sycl/gptq_int4_xmx-fused-expert.hpp b/kt-kernel/operators/sycl/gptq_int4_xmx-fused-expert.hpp new file mode 100644 index 00000000..83338524 --- /dev/null +++ b/kt-kernel/operators/sycl/gptq_int4_xmx-fused-expert.hpp @@ -0,0 +1,130 @@ +// SPDX-License-Identifier: Apache-2.0 +#ifndef KT_SYCL_GPTQ_INT4_XMX_FUSED_EXPERT_HPP +#define KT_SYCL_GPTQ_INT4_XMX_FUSED_EXPERT_HPP + +#include "gptq_int4_xmx-kernels.hpp" + +namespace sycl_xmx { +// Hardware-specific decode candidate: one work-group owns a complete expert. +// Gate/up -> BF16 SwiGLU -> groupwise Q8 -> down stays inside SLM. Global +// expert merge remains separate; there is no cross-work-group synchronization. +// Only H=2048, I=512 and group=128 are implemented for the measured model. +template +inline sycl::event fused_expert_decode(const Tile* tiles, int tile_count, const int8_t* input, + const float* input_scales, const Weight& gate, const Weight& up, + const Weight& down, float* output, float limit, float alpha, + const sycl::event* dep = nullptr) { + static_assert(KSplits == 1 || KSplits == 2); + constexpr int H = 2048, I = 512, GS = 128, NB = I / 16, Threads = NB * KSplits; + constexpr int PartialBytes = KSplits > 1 ? KSplits * I * 2 * sizeof(float) : 0; + constexpr int Middle = PartialBytes; + constexpr int Quantized = Middle + I * sizeof(uint16_t); + constexpr int Scales = Quantized + I; + constexpr int SLM = Scales + I / GS * sizeof(float); + const auto* gw = gate.packed.data; + const auto* gs = gate.scales.data; + const auto* uw = up.packed.data; + const auto* us = up.scales.data; + const auto* dw = down.packed.data; + const auto* ds = down.scales.data; + return sycl_int4::queue().submit([&](sycl::handler& handler) { + if (dep) handler.depends_on(*dep); + handler.parallel_for( + sycl::nd_range<1>(static_cast(tile_count) * Threads, Threads), + [=](sycl::nd_item<1> item) SYCL_ESIMD_KERNEL { + es::slm_init(); + const Tile tile = tiles[item.get_group_linear_id()]; + const int lane = item.get_local_linear_id(), nb = lane % NB, split = lane / NB; + const size_t wi = (static_cast(tile.expert) * NB + nb) * (H / 8) * 16; + const size_t si = (static_cast(tile.expert) * NB + nb) * (H / GS) * 16; + es::simd gacc = 0, uacc = 0; + for (int group = split; group < H / GS; group += KSplits) { + es::simd gi = 0, ui = 0; + for (int j = 0; j < GS; j += 32) { + es::simd a; + a.copy_from(input + group * GS + j); // single input token for all experts + es::simd b; + b.copy_from(gw + wi + (group * GS + j) / 8 * 16); + b ^= 0x88888888u; + gi = + xmx::dpas<8, 1, int, int, uint32_t, int8_t, xmx::dpas_argument_type::s4, xmx::dpas_argument_type::s8>( + gi, b, a); + b.copy_from(uw + wi + (group * GS + j) / 8 * 16); + b ^= 0x88888888u; + ui = + xmx::dpas<8, 1, int, int, uint32_t, int8_t, xmx::dpas_argument_type::s4, xmx::dpas_argument_type::s8>( + ui, b, a); + } + es::simd scale; + scale.copy_from(gs + si + group * 16); + gacc += es::convert(gi) * (scale * input_scales[group]); + scale.copy_from(us + si + group * 16); + uacc += es::convert(ui) * (scale * input_scales[group]); + } + if constexpr (KSplits > 1) { + es::slm_block_store(lane * 16 * sizeof(float), gacc); + es::slm_block_store((Threads + lane) * 16 * sizeof(float), uacc); + es::barrier(); + if (split == 0) { + gacc = 0; + uacc = 0; +#pragma unroll + for (int part = 0; part < KSplits; ++part) { + gacc += es::slm_block_load((part * NB + nb) * 16 * sizeof(float)); + uacc += es::slm_block_load((Threads + part * NB + nb) * 16 * sizeof(float)); + } + } + } + if (split == 0) { + if (limit > 0) { + gacc = es::min(gacc, limit); + if (alpha > 0) gacc = es::max(gacc, -limit); + uacc = es::min(es::max(uacc, -limit), limit); + } + auto value = gacc / (1.0f + es::exp(-gacc * (alpha > 0 ? alpha : 1.0f))) * (alpha > 0 ? uacc + 1.0f : uacc); + es::simd bits = value.template bit_cast_view(); + bits += 0x7fffu + ((bits >> 16) & 1u); + es::slm_block_store(Middle + nb * 16 * sizeof(uint16_t), es::convert(bits >> 16)); + } + es::barrier(); + if (lane < I / GS) { + const auto bf = es::slm_block_load(Middle + lane * GS * sizeof(uint16_t)); + es::simd bits = es::convert(bf) << 16; + const es::simd values = bits.template bit_cast_view(); + const float scale = es::hmax(es::abs(values)) / 127.0f; + es::slm_scatter(es::simd(Scales + lane * sizeof(float)), es::simd(scale)); + auto quantized = es::rnde(values / (scale > 0 ? scale : 1.0f)); + es::slm_block_store(Quantized + lane * GS, + es::convert(es::min(es::max(quantized, -127.0f), 127.0f))); + } + es::barrier(); + const auto activation_scales = es::slm_block_load(Scales); + for (int out_nb = lane; out_nb < H / 16; out_nb += Threads) { + const size_t wi = (static_cast(tile.expert) * (H / 16) + out_nb) * (I / 8) * 16; + const size_t si = (static_cast(tile.expert) * (H / 16) + out_nb) * (I / GS) * 16; + es::simd accum = 0; + for (int group = 0; group < I / GS; ++group) { + es::simd partial = 0; + for (int j = 0; j < GS; j += 32) { + const auto a = es::slm_block_load(Quantized + group * GS + j); + es::simd b; + b.copy_from(dw + wi + (group * GS + j) / 8 * 16); + b ^= 0x88888888u; + partial = xmx::dpas<8, 1, int, int, uint32_t, int8_t, xmx::dpas_argument_type::s4, + xmx::dpas_argument_type::s8>(partial, b, a); + } + es::simd scale; + scale.copy_from(ds + si + group * 16); + accum += es::convert(partial) * (scale * static_cast(activation_scales[group])); + } + es::simd bits = accum.template bit_cast_view(); + bits += 0x7fffu + ((bits >> 16) & 1u); + bits &= 0xffff0000u; + es::simd rounded = bits.template bit_cast_view(); + rounded.copy_to(output + static_cast(tile.begin) * H + out_nb * 16); + } + }); + }); +} +} // namespace sycl_xmx +#endif diff --git a/kt-kernel/operators/sycl/gptq_int4_xmx-kernels.hpp b/kt-kernel/operators/sycl/gptq_int4_xmx-kernels.hpp new file mode 100644 index 00000000..9def37b3 --- /dev/null +++ b/kt-kernel/operators/sycl/gptq_int4_xmx-kernels.hpp @@ -0,0 +1,369 @@ +// SPDX-License-Identifier: Apache-2.0 +#ifndef KT_SYCL_GPTQ_INT4_XMX_KERNELS_HPP +#define KT_SYCL_GPTQ_INT4_XMX_KERNELS_HPP + +#include +#include +#include + +#include "gptq_int4_sycl-runtime.hpp" + +namespace sycl_xmx { +namespace es = sycl::ext::intel::esimd; +namespace xmx = es::xmx; + +// One tile belongs to one expert. Rows are indices in the routed workspace. +struct Tile { + int expert; + int begin; + int rows; +}; + +template +class USM { + public: + T* data = nullptr; + size_t capacity = 0; + USM() = default; + USM(const USM&) = delete; + USM& operator=(const USM&) = delete; + ~USM() { sycl_int4::usm_free(data); } + void reserve(size_t n) { + if (n <= capacity) return; + T* replacement = sycl_int4::usm_alloc(n, "XMX workspace"); + sycl_int4::usm_free(data); + data = replacement; + capacity = n; + } +}; + +struct Weight { + USM packed; // [expert, N/16, K/8, 16], original GPTQ nibbles + USM scales; // [expert, N/16, K/group_size, 16] + int k = 0, n = 0, gs = 0; + void init(int experts, int n_, int k_, int gs_) { + n = n_; + k = k_; + gs = gs_; + packed.reserve(static_cast(experts) * n * k / 8); + scales.reserve(static_cast(experts) * n * (k / gs)); + } + void from_mat(int expert, const uint32_t* qw, const float* sc) { + if (!qw || !sc) throw std::runtime_error("XMX weight/scale pointer is null"); + uint32_t* dst = packed.data + static_cast(expert) * n * k / 8; + float* sd = scales.data + static_cast(expert) * n * (k / gs); + for (int nb = 0; nb < n / 16; ++nb) { + for (int p = 0; p < k / 8; ++p) + std::memcpy(dst + (static_cast(nb) * (k / 8) + p) * 16, qw + p * n + nb * 16, 64); + for (int g = 0; g < k / gs; ++g) + std::memcpy(sd + (static_cast(nb) * (k / gs) + g) * 16, sc + g * n + nb * 16, 64); + } + } +}; + +ESIMD_INLINE es::simd bf16_load(const uint16_t* p) { + es::simd b; + b.copy_from(p); + es::simd u = es::convert(b) << 16; + return u.template bit_cast_view(); +} + +// DPAS B uses two FP16 K values per dword, then consecutive output columns. +ESIMD_INLINE es::simd dequant_f16(const uint32_t* p, es::simd scale) { + es::simd w; + w.copy_from(p); + es::simd b; +#pragma unroll + for (int pair = 0; pair < 8; ++pair) { + es::simd q = w.template select<16, 1>((pair / 4) * 16); + const int shift = (pair % 4) * 8; + es::simd lo = (es::convert((q >> shift) & 15) - 8.0f) * scale; + es::simd hi = (es::convert((q >> (shift + 4)) & 15) - 8.0f) * scale; + es::simd lh = es::convert(lo), hh = es::convert(hi); + es::simd word = es::convert(lh.template bit_cast_view().read()) | + (es::convert(hh.template bit_cast_view().read()) << 16); + b.template bit_cast_view().template select<16, 1>(pair * 16) = word; + } + return b; +} + +inline sycl::event quantize(const uint16_t* src, int8_t* dst, float* scales, int rows, int k, int gs, + const sycl::event* dep = nullptr) { + const int groups = k / gs; + return sycl_int4::queue().submit([&](sycl::handler& h) { + if (dep) h.depends_on(*dep); + h.parallel_for(sycl::nd_range<1>(static_cast(rows) * groups * 128, 128), [=](sycl::nd_item<1> it) { + int rg = it.get_group(0), lane = it.get_local_id(0); + int row = rg / groups, g = rg % groups; + float maximum = 0.0f; + for (int j = lane; j < gs; j += 128) + maximum = sycl::fmax(maximum, sycl::fabs(sycl_int4::bf16_to_fp32(src[row * k + g * gs + j]))); + maximum = sycl::reduce_over_group(it.get_group(), maximum, sycl::maximum()); + float scale = maximum / 127.0f; + if (lane == 0) scales[rg] = scale; + for (int j = lane; j < gs; j += 128) { + float v = sycl_int4::bf16_to_fp32(src[row * k + g * gs + j]); + dst[row * k + g * gs + j] = + static_cast(sycl::clamp(sycl::rint(scale > 0 ? v / scale : 0), -127.0f, 127.0f)); + } + }); + }); +} + +template +inline sycl::event projection(const Tile* tiles, int tile_count, const int* row_tokens, const uint16_t* input, + const int8_t* qinput, const float* ascales, const Weight& first, const Weight* second, + uint16_t* activated, float* down, float limit, float alpha, + const sycl::event* dep = nullptr, int8_t* qoutput = nullptr, + float* output_scales = nullptr, const int* merge_slots = nullptr, + const float* routing = nullptr, float* merged = nullptr, int topk = 0) { + static_assert(!FusedQuant || (Q8 && GateUp)); + static_assert(KSplits == 1 || KSplits == 2 || KSplits == 4 || KSplits == 8); + static_assert(!FusedQuant || KSplits == 1); + static_assert(!FusedMerge || (Rows == 1 && !GateUp && !FusedQuant && KSplits == 1)); + const int n = first.n, k = first.k, gs = first.gs, groups = k / gs; + const uint32_t* w0 = first.packed.data; + const float* s0 = first.scales.data; + const uint32_t* w1 = second ? second->packed.data : nullptr; + const float* s1 = second ? second->scales.data : nullptr; + auto compute = [=](size_t index, int lane) SYCL_ESIMD_FUNCTION { + if constexpr (FusedQuant) es::slm_init(); + if constexpr (FusedMerge) es::slm_init<8 * 16 * sizeof(float)>(); + if constexpr (KSplits > 1) es::slm_init(); + const int nb = index % (n / 16); + const bool valid = !FusedMerge || index / (n / 16) < static_cast(tile_count); + const Tile tile = valid ? tiles[index / (n / 16)] : Tile{0, 0, 0}; + const size_t wi = (static_cast(tile.expert) * (n / 16) + nb) * (k / 8) * 16; + const size_t si = (static_cast(tile.expert) * (n / 16) + nb) * groups * 16; + es::simd acc0 = 0.0f, acc1 = 0.0f; + for (int g = KSplits > 1 ? lane : 0; g < (valid ? groups : 0); g += KSplits) { + es::simd sc0; + sc0.copy_from(s0 + si + g * 16); + es::simd sc1; + if constexpr (GateUp) sc1.copy_from(s1 + si + g * 16); + if constexpr (Q8) { + es::simd c0 = 0, c1 = 0; + for (int j = 0; j < gs; j += 32) { + es::simd a = 0; +#pragma unroll + for (int r = 0; r < Rows; ++r) { + if (r < tile.rows) { + int row = GateUp ? row_tokens[tile.begin + r] : tile.begin + r; + es::simd v; + v.copy_from(qinput + static_cast(row) * k + g * gs + j); + a.template select<32, 1>(r * 32) = v; + } + } + es::simd b; + b.copy_from(w0 + wi + (g * gs + j) / 8 * 16); + b ^= 0x88888888u; + constexpr int D = Rows > 8 ? 8 : Rows; +#pragma unroll + for (int row_base = 0; row_base < Rows; row_base += D) { + es::simd partial = c0.template select(row_base * 16); + es::simd src = a.template select(row_base * 32); + c0.template select(row_base * 16) = + xmx::dpas<8, D, int, int, uint32_t, int8_t, xmx::dpas_argument_type::s4, xmx::dpas_argument_type::s8>( + partial, b, src); + } + if constexpr (GateUp) { + b.copy_from(w1 + wi + (g * gs + j) / 8 * 16); + b ^= 0x88888888u; +#pragma unroll + for (int row_base = 0; row_base < Rows; row_base += D) { + es::simd partial = c1.template select(row_base * 16); + es::simd src = a.template select(row_base * 32); + c1.template select(row_base * 16) = + xmx::dpas<8, D, int, int, uint32_t, int8_t, xmx::dpas_argument_type::s4, xmx::dpas_argument_type::s8>( + partial, b, src); + } + } + } +#pragma unroll + for (int r = 0; r < Rows; ++r) { + if (r < tile.rows) { + int row = GateUp ? row_tokens[tile.begin + r] : tile.begin + r; + float sa = ascales[row * groups + g]; + acc0.template select<16, 1>(r * 16) += + es::convert(c0.template select<16, 1>(r * 16).read()) * (sc0 * sa); + if constexpr (GateUp) + acc1.template select<16, 1>(r * 16) += + es::convert(c1.template select<16, 1>(r * 16).read()) * (sc1 * sa); + } + } + } else { + for (int j = 0; j < gs; j += 16) { + es::simd a = 0; +#pragma unroll + for (int r = 0; r < Rows; ++r) { + if (r < tile.rows) { + int row = GateUp ? row_tokens[tile.begin + r] : tile.begin + r; + a.template select<16, 1>(r * 16) = + es::convert(bf16_load(input + static_cast(row) * k + g * gs + j)); + } + } + auto b = dequant_f16(w0 + wi + (g * gs + j) / 8 * 16, sc0); + constexpr int D = Rows > 8 ? 8 : Rows; +#pragma unroll + for (int row_base = 0; row_base < Rows; row_base += D) { + es::simd partial = acc0.template select(row_base * 16); + es::simd src = a.template select(row_base * 16); + acc0.template select(row_base * 16) = xmx::dpas<8, D, float>(partial, b, src); + } + if constexpr (GateUp) { + b = dequant_f16(w1 + wi + (g * gs + j) / 8 * 16, sc1); +#pragma unroll + for (int row_base = 0; row_base < Rows; row_base += D) { + es::simd partial = acc1.template select(row_base * 16); + es::simd src = a.template select(row_base * 16); + acc1.template select(row_base * 16) = xmx::dpas<8, D, float>(partial, b, src); + } + } + } + } + // Groups carry distinct scales; integer accumulation never crosses them. + } + if constexpr (KSplits > 1) { + constexpr int bytes = Rows * 16 * sizeof(float); + es::slm_block_store(lane * bytes, acc0); + if constexpr (GateUp) es::slm_block_store((KSplits + lane) * bytes, acc1); + es::barrier(); + if (lane != 0) return; + acc0 = 0.0f; + acc1 = 0.0f; +#pragma unroll + for (int split = 0; split < KSplits; ++split) { + acc0 += es::slm_block_load(split * bytes); + if constexpr (GateUp) acc1 += es::slm_block_load((KSplits + split) * bytes); + } + } +#pragma unroll + for (int r = 0; r < Rows; ++r) { + if (r >= tile.rows) continue; + es::simd value = acc0.template select<16, 1>(r * 16); + if constexpr (GateUp) { + es::simd up = acc1.template select<16, 1>(r * 16); + if (limit > 0) { + value = es::min(value, limit); + if (alpha > 0) value = es::max(value, -limit); + up = es::min(es::max(up, -limit), limit); + } + value = value / (1.0f + es::exp(-value * (alpha > 0 ? alpha : 1.0f))) * (alpha > 0 ? up + 1.0f : up); + } + es::simd bits = value.template bit_cast_view(); + bits += 0x7fffu + ((bits >> 16) & 1u); + es::simd bf = es::convert(bits >> 16); + if constexpr (FusedQuant) { + es::simd u = es::convert(bf) << 16; + es::simd rounded = u.template bit_cast_view(); + acc0.template select<16, 1>(r * 16) = rounded; + float maximum = es::hmax(es::abs(rounded)); + es::slm_scatter(es::simd((r * 8 + lane) * sizeof(float)), es::simd(maximum)); + } else if constexpr (GateUp) { + bf.copy_to(activated + static_cast(tile.begin + r) * n + nb * 16); + } else { + es::simd u = es::convert(bf) << 16; + es::simd rounded = u.template bit_cast_view(); + if constexpr (FusedMerge) + acc0 = rounded; + else + rounded.copy_to(down + static_cast(tile.begin + r) * n + nb * 16); + } + } + if constexpr (FusedMerge) { + // Decode tiles have exactly one row. Eight work-items own a column + // block across routed rows, keeping BF16 rounding before slot weights. + // Inactive padded lanes still reach the work-group barrier. + es::slm_block_store(lane * 16 * sizeof(float), acc0); + es::barrier(); + if (lane == 0) { + es::simd sum = 0.0f; + for (int slot = 0; slot < topk; ++slot) { + const int row = merge_slots[slot]; + if (row >= 0) { + const auto values = es::slm_block_load(row * 16 * sizeof(float)); + sum += values * routing[slot]; + } + } + sum.copy_to(merged + nb * 16); + } + } + if constexpr (FusedQuant) { + // Eight adjacent N=16 work-items jointly own one group of 128 values. + es::barrier(); +#pragma unroll + for (int r = 0; r < Rows; ++r) { + if (r >= tile.rows) continue; + auto maxima = es::slm_block_load(r * 8 * sizeof(float)); + float scale = es::hmax(maxima) / 127.0f; + if (lane == 0) output_scales[(tile.begin + r) * (n / 128) + nb / 8] = scale; + es::simd values = acc0.template select<16, 1>(r * 16); + values = es::rnde(values / (scale > 0 ? scale : 1.0f)); + es::simd quantized = es::convert(es::min(es::max(values, -127.0f), 127.0f)); + quantized.copy_to(qoutput + static_cast(tile.begin + r) * n + nb * 16); + } + } + }; + return sycl_int4::queue().submit([&](sycl::handler& h) { + if (dep) h.depends_on(*dep); + const size_t work = static_cast(tile_count) * (n / 16); + if constexpr (FusedMerge) { + h.parallel_for(sycl::nd_range<1>(static_cast(n / 16) * 8, 8), [=](sycl::nd_item<1> it) SYCL_ESIMD_KERNEL { + const int lane = it.get_local_linear_id(); + compute(static_cast(lane) * (n / 16) + it.get_group_linear_id(), lane); + }); + } else if constexpr (FusedQuant) { + h.parallel_for(sycl::nd_range<1>(work, 8), [=](sycl::nd_item<1> it) SYCL_ESIMD_KERNEL { + compute(it.get_global_linear_id(), it.get_local_linear_id()); + }); + } else if constexpr (KSplits > 1) { + h.parallel_for(sycl::nd_range<1>(work * KSplits, KSplits), [=](sycl::nd_item<1> it) SYCL_ESIMD_KERNEL { + compute(it.get_group_linear_id(), it.get_local_linear_id()); + }); + } else { + h.parallel_for(sycl::range<1>(work), [=](sycl::id<1> id) SYCL_ESIMD_KERNEL { compute(id[0], 0); }); + } + }); +} + +// Preserve the original slot summation order and FP32 precision. A work-item +// handles a contiguous output vector, amortizing route lookup and loads. +template +inline sycl::event merge_vector(const int* slot_rows, const float* weights, const float* down, float* output, int rows, + int topk, int hidden, const sycl::event& dep) { + return sycl_int4::queue().submit([&](sycl::handler& h) { + h.depends_on(dep); + h.parallel_for(sycl::range<1>(static_cast(rows) * (hidden / Width)), [=](sycl::id<1> id) SYCL_ESIMD_KERNEL { + const int token = id[0] / (hidden / Width); + const int col = (id[0] % (hidden / Width)) * Width; + es::simd sum = 0; + for (int slot = 0; slot < topk; ++slot) { + const int row = slot_rows[token * topk + slot]; + if (row >= 0) { + es::simd values; + values.copy_from(down + static_cast(row) * hidden + col); + sum += values * weights[token * topk + slot]; + } + } + sum.copy_to(output + static_cast(token) * hidden + col); + }); + }); +} + +inline sycl::event merge(const int* slot_rows, const float* weights, const float* down, float* output, int rows, + int topk, int hidden, const sycl::event& dep) { + return sycl_int4::queue().submit([&](sycl::handler& h) { + h.depends_on(dep); + h.parallel_for(sycl::range<1>(static_cast(rows) * hidden), [=](sycl::id<1> id) { + int token = id[0] / hidden, col = id[0] % hidden; + float sum = 0; + for (int slot = 0; slot < topk; ++slot) { + int row = slot_rows[token * topk + slot]; + if (row >= 0) sum += down[static_cast(row) * hidden + col] * weights[token * topk + slot]; + } + output[id] = sum; + }); + }); +} +} // namespace sycl_xmx +#endif diff --git a/kt-kernel/operators/sycl/gptq_int4_xmx-moe.hpp b/kt-kernel/operators/sycl/gptq_int4_xmx-moe.hpp new file mode 100644 index 00000000..5ea467fd --- /dev/null +++ b/kt-kernel/operators/sycl/gptq_int4_xmx-moe.hpp @@ -0,0 +1,662 @@ +// SPDX-License-Identifier: Apache-2.0 +#ifndef KT_SYCL_GPTQ_INT4_XMX_MOE_HPP +#define KT_SYCL_GPTQ_INT4_XMX_MOE_HPP + +#include +#include +#include +#include + +#include "../avx2/avx2_bf16_utils.hpp" +#include "../moe-tp.hpp" +#include "gptq_int4_vector-kernels.hpp" +#include "gptq_int4_xmx-fused-expert.hpp" +#include "gptq_int4_xmx-options.hpp" +#include "gptq_int4_xmx-prefill.hpp" +#include "gptq_int4_xmx-profile.hpp" + +// Dedicated single-NUMA executor. No CPU expert-input pools or FP16 weight copy. +class SYCL_XMX_GPTQ_INT4_MOE_TP { + public: + using input_t = ggml_bf16_t; + using output_t = float; + GeneralMOEConfig config; + explicit SYCL_XMX_GPTQ_INT4_MOE_TP(GeneralMOEConfig c, int part = 0) + : config(c), profile_(std::getenv("KT_SYCL_XMX_TRACE"), c.layer_idx) { + if (part != 0 || c.pool->config.subpool_count != 1) + throw std::runtime_error("SYCL XMX currently requires one CPU threadpool/NUMA part"); + if (c.quant_config.bits != 4 || c.quant_config.zero_point) + throw std::runtime_error("SYCL XMX requires symmetric GPTQ INT4 weights"); + int gs = c.quant_config.group_size; + if (gs <= 0 || gs % 32 || c.hidden_size <= 0 || c.intermediate_size <= 0 || c.hidden_size % gs || + c.intermediate_size % gs || c.hidden_size % 16 || c.intermediate_size % 16 || c.expert_num <= 0 || + c.num_experts_per_tok <= 0 || c.max_possible_qlen() <= 0) + throw std::runtime_error("SYCL XMX requires positive dimensions, N aligned to 16 and K/group_size aligned to 32"); + if (options_.panther_lake() && (c.hidden_size != 2048 || c.intermediate_size != 512 || gs != 128)) + throw std::runtime_error("KT_SYCL_XMX_PROFILE=panther-lake requires H=2048 I=512 group=128"); + auto device = sycl_int4::queue().get_device(); + if (!device.has(sycl::aspect::ext_intel_esimd) || !device.has(sycl::aspect::ext_intel_matrix)) + throw std::runtime_error("SYCL XMX requires an Intel GPU with ESIMD and XMX"); + const char* mode = options_.get("KT_SYCL_XMX_MODE"); + std::string value = mode ? mode : "fp16"; + if (value != "fp16" && value != "int8") throw std::runtime_error("KT_SYCL_XMX_MODE must be fp16 or int8"); + q8_ = value == "int8"; + gate_kernel_ = parse_kernel("KT_SYCL_XMX_GATE_KERNEL", Kernel::XMX); + down_kernel_ = parse_kernel("KT_SYCL_XMX_DOWN_KERNEL", Kernel::XMX); + decode_gate_ = parse_kernel("KT_SYCL_XMX_DECODE_GATE_KERNEL", gate_kernel_); + decode_down_ = parse_kernel("KT_SYCL_XMX_DECODE_DOWN_KERNEL", down_kernel_); + vector_splits_ = parse_splits("KT_SYCL_XMX_VECTOR_SPLITS", 4); + gate_splits_ = parse_splits("KT_SYCL_XMX_GATE_SPLITS", vector_splits_); + down_splits_ = parse_splits("KT_SYCL_XMX_DOWN_SPLITS", vector_splits_); + const char* expert = options_.get("KT_SYCL_XMX_FUSED_EXPERT"); + const std::string expert_mode = expert ? expert : "0"; + if (expert_mode != "0" && expert_mode != "1") throw std::runtime_error("KT_SYCL_XMX_FUSED_EXPERT must be 0 or 1"); + fused_expert_ = expert_mode == "1"; + if (fused_expert_ && (!q8_ || c.hidden_size != 2048 || c.intermediate_size != 512 || gs != 128 || + (decode_gate_ != Kernel::XMX && decode_gate_ != Kernel::XMXSplit) || + decode_down_ != Kernel::XMX || (decode_gate_ == Kernel::XMXSplit && gate_splits_ > 2))) + throw std::runtime_error( + "KT_SYCL_XMX_FUSED_EXPERT requires INT8 XMX H=2048 I=512 group=128, unsplit down and gate splits<=2"); + const char* merge = options_.get("KT_SYCL_XMX_DECODE_MERGE"); + const std::string merge_mode = merge ? merge : "separate"; + if (merge_mode != "separate" && merge_mode != "fused") + throw std::runtime_error("KT_SYCL_XMX_DECODE_MERGE must be separate or fused"); + fused_merge_ = merge_mode == "fused"; + if (fused_expert_ && fused_merge_) + throw std::runtime_error("KT_SYCL_XMX_FUSED_EXPERT does not support fused decode merge"); + if (fused_merge_ && (c.num_experts_per_tok > 8 || decode_down_ != Kernel::XMX)) + throw std::runtime_error("KT_SYCL_XMX_DECODE_MERGE=fused requires topk<=8 and unsplit XMX decode down"); + const char* tile_rows = options_.get("KT_SYCL_XMX_ROWS"); + std::string tile = tile_rows ? tile_rows : "8"; + if (tile != "1" && tile != "2" && tile != "4" && tile != "8") + throw std::runtime_error("KT_SYCL_XMX_ROWS must be 1, 2, 4 or 8"); + const char* prefill_rows = options_.get("KT_SYCL_XMX_PREFILL_ROWS"); + const std::string prefill_tile = prefill_rows ? std::string(prefill_rows) : tile; + if (prefill_tile != "1" && prefill_tile != "2" && prefill_tile != "4" && prefill_tile != "8" && + prefill_tile != "16" && prefill_tile != "32") + throw std::runtime_error("KT_SYCL_XMX_PREFILL_ROWS must be 1, 2, 4, 8, 16 or 32"); + prefill_rows_ = std::stoi(prefill_tile); + if (prefill_rows_ > 8 && (!q8_ || gate_kernel_ != Kernel::XMX || down_kernel_ != Kernel::XMX)) + throw std::runtime_error("KT_SYCL_XMX_PREFILL_ROWS > 8 requires INT8 unsplit XMX prefill projections"); + const char* down_rows = options_.get("KT_SYCL_XMX_PREFILL_DOWN_ROWS"); + const std::string down_tile = down_rows ? down_rows : prefill_tile; + if (down_tile != "1" && down_tile != "2" && down_tile != "4" && down_tile != "8" && down_tile != "16" && + down_tile != "32") + throw std::runtime_error("KT_SYCL_XMX_PREFILL_DOWN_ROWS must be 1, 2, 4, 8, 16 or 32"); + prefill_down_rows_ = std::stoi(down_tile); + if (prefill_down_rows_ != prefill_rows_ && (!q8_ || gate_kernel_ != Kernel::XMX || down_kernel_ != Kernel::XMX)) + throw std::runtime_error("Independent prefill down rows require INT8 unsplit XMX projections"); + const char* prefill_n = options_.get("KT_SYCL_XMX_PREFILL_N"); + const std::string prefill_columns = prefill_n ? prefill_n : "16"; + if (prefill_columns != "16" && prefill_columns != "32") + throw std::runtime_error("KT_SYCL_XMX_PREFILL_N must be 16 or 32"); + prefill_n_ = std::stoi(prefill_columns); + if (prefill_n_ == 32 && (!q8_ || gate_kernel_ != Kernel::XMX || down_kernel_ != Kernel::XMX || c.hidden_size % 32 || + c.intermediate_size % 32 || prefill_rows_ > 16 || prefill_down_rows_ > 16)) + throw std::runtime_error("KT_SYCL_XMX_PREFILL_N=32 requires INT8 XMX prefill, N%32=0 and rows<=16"); + const char* load = options_.get("KT_SYCL_XMX_PREFILL_LOAD"); + const std::string load_mode = load ? load : "block"; + if (load_mode != "block" && load_mode != "gather") + throw std::runtime_error("KT_SYCL_XMX_PREFILL_LOAD must be block or gather"); + if (load_mode != "block" && prefill_n_ != 32) + throw std::runtime_error("KT_SYCL_XMX_PREFILL_LOAD=gather requires PREFILL_N=32"); + prefill_load_ = load_mode == "gather" ? sycl_xmx::PrefillLoad::Gather : sycl_xmx::PrefillLoad::Block; + const char* merge_prefill = options_.get("KT_SYCL_XMX_PREFILL_MERGE"); + const std::string prefill_merge = merge_prefill ? merge_prefill : "scalar"; + if (prefill_merge != "scalar" && prefill_merge != "vector") + throw std::runtime_error("KT_SYCL_XMX_PREFILL_MERGE must be scalar or vector"); + vector_prefill_merge_ = prefill_merge == "vector"; + const char* graph = options_.get("KT_SYCL_XMX_GRAPH"); + if (graph && std::string(graph) != "0" && std::string(graph) != "1") + throw std::runtime_error("KT_SYCL_XMX_GRAPH must be 0 or 1"); + use_graph_ = graph && std::string(graph) == "1"; + const char* hq = options_.get("KT_SYCL_XMX_HOST_QUANT"); + if (hq && std::string(hq) != "0" && std::string(hq) != "1") + throw std::runtime_error("KT_SYCL_XMX_HOST_QUANT must be 0 or 1"); + host_quant_ = hq && std::string(hq) == "1"; + const char* fq = options_.get("KT_SYCL_XMX_FUSED_QUANT"); + if (fq && std::string(fq) != "0" && std::string(fq) != "1" && std::string(fq) != "prefill") + throw std::runtime_error("KT_SYCL_XMX_FUSED_QUANT must be 0, 1 or prefill"); + fused_quant_ = fq && std::string(fq) != "0"; + fused_decode_ = fq && std::string(fq) == "1"; + if (fused_quant_ && + (!q8_ || gs != 128 || c.intermediate_size % 128 || gate_kernel_ != Kernel::XMX || !uses_q8(down_kernel_) || + (fused_decode_ && (decode_gate_ != Kernel::XMX || !uses_q8(decode_down_))))) + throw std::runtime_error("SYCL XMX fused quantization requires int8 XMX gate, quantized down and group_size=128"); + if (use_graph_ && profile_.device) + throw std::runtime_error( + "KT_SYCL_XMX_TRACE=device does not profile recorded graph nodes; use TRACE=host or GRAPH=0"); + if (use_graph_ && !device.has(sycl::aspect::ext_oneapi_limited_graph)) + throw std::runtime_error("SYCL XMX command graph is unavailable on this device"); + const int max_rows = c.max_possible_qlen(); + if (max_rows > std::numeric_limits::max() / c.num_experts_per_tok || + max_rows > std::numeric_limits::max() / c.hidden_size) + throw std::runtime_error("SYCL XMX workspace dimensions exceed 32-bit indexing"); + const int max_routed = max_rows * c.num_experts_per_tok; + if (max_routed > std::numeric_limits::max() / c.hidden_size || + max_routed > std::numeric_limits::max() / c.intermediate_size) + throw std::runtime_error("SYCL XMX routed workspace exceeds 32-bit indexing"); + if (use_graph_) graphs_.resize(static_cast(c.num_experts_per_tok + 1) * (c.num_experts_per_tok + 1)); + counts_.resize(c.expert_num); + offsets_.resize(c.expert_num + 1); + cursors_.resize(c.expert_num); + input_.reserve(static_cast(max_rows) * c.hidden_size); + activated_.reserve(static_cast(max_routed) * c.intermediate_size); + down_.reserve(static_cast(max_routed) * c.hidden_size); + output_.reserve(static_cast(max_rows) * c.hidden_size); + row_tokens_.reserve(max_routed); + slot_rows_.reserve(max_routed); + routing_.reserve(max_routed); + tiles_.reserve(max_routed); + down_tiles_.reserve(max_routed); + if (uses_q8(gate_kernel_) || uses_q8(decode_gate_)) { + qi_.reserve(static_cast(max_rows) * c.hidden_size); + si_.reserve(static_cast(max_rows) * c.hidden_size / gs); + } + if (uses_q8(down_kernel_) || uses_q8(decode_down_)) { + qm_.reserve(static_cast(max_routed) * c.intermediate_size); + sm_.reserve(static_cast(max_routed) * c.intermediate_size / gs); + } + if (options_.panther_lake()) printf("SYCL XMX profile=panther-lake\n"); + printf("Created SYCL_XMX_GPTQ_INT4_MOE mode=%s layer=%d\n", value.c_str(), c.layer_idx); + printf("SYCL decode fused_expert=%d merge=%s\n", fused_expert_, merge_mode.c_str()); + printf("SYCL prefill tile gate=M%dN%d down=M%dN%d\n", prefill_rows_, prefill_n_, prefill_down_rows_, prefill_n_); + printf("SYCL prefill load=%s merge=%s\n", load_mode.c_str(), prefill_merge.c_str()); + printf("SYCL projection kernels prefill=%s/%s decode=%s/%s splits=%d/%d\n", kernel_name(gate_kernel_), + kernel_name(down_kernel_), kernel_name(decode_gate_), kernel_name(decode_down_), gate_splits_, down_splits_); + } + + void load_weights(const GeneralMOEConfig& source) { + const int e = config.expert_num, h = config.hidden_size, i = config.intermediate_size; + const int gs = config.quant_config.group_size; + gate_.init(e, i, h, gs); + up_.init(e, i, h, gs); + down_weight_.init(e, h, i, gs); + const auto* map = static_cast(source.physical_to_logical_map); + loaded_.assign(e, false); + for (int physical = 0; physical < e; ++physical) { + uint64_t mapped = map ? map[physical] : static_cast(physical); + if (mapped >= static_cast(e)) throw std::runtime_error("Invalid XMX physical_to_logical_map"); + int logical = static_cast(mapped); + if (source.should_skip_expert(logical)) continue; + auto load = [&](sycl_xmx::Weight& w, void* p, void* s, const auto& ps, const auto& ss) { + const uint32_t* qw = + ps.empty() ? static_cast(p) : static_cast(ps.at(0).at(logical)); + const float* sc = ss.empty() ? static_cast(s) : static_cast(ss.at(0).at(logical)); + if (ps.empty() && qw) qw += static_cast(logical) * w.n * w.k / 8; + if (ss.empty() && sc) sc += static_cast(logical) * w.n * (w.k / w.gs); + w.from_mat(physical, qw, sc); + }; + load(gate_, source.gate_proj, source.gate_scale, source.gate_projs, source.gate_scales); + load(up_, source.up_proj, source.up_scale, source.up_projs, source.up_scales); + load(down_weight_, source.down_proj, source.down_scale, source.down_projs, source.down_scales); + loaded_[physical] = true; + } + } + + void forward(int rows, int topk, const int64_t* ids, const float* weights, const void* input, void* output) { + if (rows < 0 || rows > config.max_possible_qlen() || topk < 0 || topk > config.num_experts_per_tok) + throw std::runtime_error("XMX forward exceeds configured workspace"); + if (rows == 0) return; + if (loaded_.empty()) throw std::runtime_error("XMX weights not loaded"); + const int h = config.hidden_size, i = config.intermediate_size; + using P = sycl_xmx::PipelineProfile; + const auto begin = profile_.enabled ? P::Clock::now() : P::Time{}; + auto checkpoint = begin; + if (profile_.enabled) frame_ = {}; + std::fill(counts_.begin(), counts_.end(), 0); + std::fill(slot_rows_.data, slot_rows_.data + rows * topk, -1); + for (int slot = 0; slot < rows * topk; ++slot) { + int64_t expert = ids[slot]; + if (expert >= 0 && expert < config.expert_num && loaded_[expert]) ++counts_[expert]; + } + offsets_[0] = 0; + for (int e = 0; e < config.expert_num; ++e) { + offsets_[e + 1] = offsets_[e] + counts_[e]; + cursors_[e] = offsets_[e]; + } + for (int slot = 0; slot < rows * topk; ++slot) { + int64_t expert = ids[slot]; + if (expert < 0 || expert >= config.expert_num || !loaded_[expert]) continue; + const int routed = cursors_[expert]++; + row_tokens_.data[routed] = slot / topk; + slot_rows_.data[slot] = routed; + } + int routed = offsets_.back(), nt = 0, tile_rows = rows == 1 ? 1 : prefill_rows_; + for (int e = 0; e < config.expert_num; ++e) { + for (int r = 0; r < counts_[e]; r += tile_rows) + tiles_.data[nt++] = {e, offsets_[e] + r, std::min(tile_rows, counts_[e] - r)}; + } + down_tile_count_ = nt; + if (rows > 1 && prefill_down_rows_ != prefill_rows_) { + down_tile_count_ = 0; + for (int e = 0; e < config.expert_num; ++e) + for (int r = 0; r < counts_[e]; r += prefill_down_rows_) + down_tiles_.data[down_tile_count_++] = {e, offsets_[e] + r, std::min(prefill_down_rows_, counts_[e] - r)}; + } + if (profile_.enabled) { + const auto now = P::Clock::now(); + frame_.host[P::Route] = P::us(checkpoint, now); + frame_.routed = routed; + frame_.tiles = nt; + checkpoint = now; + } + if (nt == 0) { + std::fill_n(static_cast(output), static_cast(rows) * h, 0.0f); + if (profile_.enabled) { + frame_.host[P::Executor] = P::us(begin, P::Clock::now()); + profile_.record(rows, frame_); + } + return; + } + std::memcpy(input_.data, input, static_cast(rows) * h * 2); + std::memcpy(routing_.data, weights, static_cast(rows) * topk * sizeof(float)); + if (profile_.enabled) { + const auto now = P::Clock::now(); + frame_.host[P::InputCopy] = P::us(checkpoint, now); + checkpoint = now; + } + if (uses_q8(decode_gate_) && host_quant_ && rows == 1) { + const int gs = config.quant_config.group_size; + for (int g = 0; g < h / gs; ++g) { + const auto* src = static_cast(input) + g * gs; + float maximum = 0; + for (int j = 0; j < gs; j += 8) { + __m256 v = avx2::load_bf16_to_fp32(src + j); + v = _mm256_and_ps(v, _mm256_castsi256_ps(_mm256_set1_epi32(0x7fffffff))); + alignas(32) float values[8]; + _mm256_store_ps(values, v); + for (float value : values) maximum = std::max(maximum, value); + } + const float scale = maximum / 127.0f; + si_.data[g] = scale; + const __m256 inv = _mm256_set1_ps(scale > 0 ? 1.0f / scale : 0.0f); + for (int j = 0; j < gs; j += 8) { + __m256 v = _mm256_mul_ps(avx2::load_bf16_to_fp32(src + j), inv); + __m256i iv = _mm256_cvtps_epi32(_mm256_round_ps(v, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC)); + alignas(32) int values[8]; + _mm256_store_si256(reinterpret_cast<__m256i*>(values), iv); + for (int r = 0; r < 8; ++r) qi_.data[g * gs + j + r] = static_cast(std::clamp(values[r], -127, 127)); + } + } + } + if (profile_.enabled) frame_.host[P::HostQuant] = P::us(checkpoint, P::Clock::now()); + try { + switch (tile_rows) { + case 1: + run<1>(nt, rows, topk, routed); + break; + case 2: + run<2>(nt, rows, topk, routed); + break; + case 4: + run<4>(nt, rows, topk, routed); + break; + case 16: + run<16>(nt, rows, topk, routed); + break; + case 32: + run<32>(nt, rows, topk, routed); + break; + default: + run<8>(nt, rows, topk, routed); + break; + } + } catch (...) { + // A failed submission may leave earlier stages in flight. Drain before + // the caller reuses the shared workspace or destroys this executor. + try { + sycl_int4::queue().wait_and_throw(); + } catch (...) { + } + throw; + } + if (profile_.enabled) checkpoint = P::Clock::now(); + std::memcpy(output, output_.data, static_cast(rows) * h * sizeof(float)); + if (profile_.enabled) { + const auto now = P::Clock::now(); + frame_.host[P::OutputCopy] = P::us(checkpoint, now); + frame_.host[P::Executor] = P::us(begin, now); + profile_.record(rows, frame_); + } + } + + sycl_xmx::PipelineProfile& pipeline_profile() { return profile_; } + + private: + sycl_xmx::Options options_; + sycl_xmx::PipelineProfile profile_; + sycl_xmx::PipelineProfile::Frame frame_; + void remember(int stage, const sycl::event& event) { + if (profile_.device) { + frame_.events[stage] = event; + frame_.present[stage] = true; + } + } + enum class Kernel { XMX, XMXSplit, Vector, DP4A }; + Kernel parse_kernel(const char* name, Kernel fallback) { + const char* value = options_.get(name); + if (!value || std::string(value) == "inherit") return fallback; + if (std::string(value) == "xmx") return Kernel::XMX; + if (std::string(value) == "xmx-split") return Kernel::XMXSplit; + if (std::string(value) == "vector") return Kernel::Vector; + if (std::string(value) == "dp4a") return Kernel::DP4A; + throw std::runtime_error(std::string(name) + " must be xmx, xmx-split, vector, dp4a or inherit"); + } + static const char* kernel_name(Kernel k) { + return k == Kernel::XMX ? "xmx" : k == Kernel::XMXSplit ? "xmx-split" : k == Kernel::Vector ? "vector" : "dp4a"; + } + int parse_splits(const char* name, int fallback) { + const char* value = options_.get(name); + if (!value) return fallback; + std::string s(value); + if (s == "1" || s == "2" || s == "4" || s == "8") return std::stoi(s); + throw std::runtime_error(std::string(name) + " must be 1, 2, 4 or 8"); + } + bool uses_q8(Kernel k) const { return k == Kernel::DP4A || ((k == Kernel::XMX || k == Kernel::XMXSplit) && q8_); } + bool q8_ = false; + Kernel gate_kernel_ = Kernel::XMX, down_kernel_ = Kernel::XMX; + Kernel decode_gate_ = Kernel::XMX, decode_down_ = Kernel::XMX; + int vector_splits_ = 4, gate_splits_ = 4, down_splits_ = 4; + int prefill_rows_ = 8, prefill_down_rows_ = 8, prefill_n_ = 16, down_tile_count_ = 0; + sycl_xmx::PrefillLoad prefill_load_ = sycl_xmx::PrefillLoad::Block; + bool vector_prefill_merge_ = false; + bool use_graph_ = false; + bool fused_merge_ = false; + bool fused_expert_ = false; + bool host_quant_ = false; + bool fused_quant_ = false; + bool fused_decode_ = false; + using Graph = + sycl::ext::oneapi::experimental::command_graph; + std::vector loaded_; + std::vector counts_, offsets_, cursors_; + sycl_xmx::Weight gate_, up_, down_weight_; + sycl_xmx::USM input_, activated_; + sycl_xmx::USM down_, output_, routing_, si_, sm_; + sycl_xmx::USM qi_, qm_; + sycl_xmx::USM row_tokens_, slot_rows_; + sycl_xmx::USM tiles_, down_tiles_; + // Executable graphs must be destroyed before the USM storage they reference. + std::vector> graphs_; + + template + sycl::event submit_projection(Kernel kernel, const sycl::event* dep, int nt, bool fused = false, bool prefill = false, + const sycl_xmx::Tile* selected_tiles = nullptr) { + const auto* tiles = selected_tiles ? selected_tiles : tiles_.data; + const auto& first = GateUp ? gate_ : down_weight_; + const auto* second = GateUp ? &up_ : nullptr; + const auto* input = GateUp ? input_.data : activated_.data; + const auto* quantized = GateUp ? qi_.data : qm_.data; + const auto* scales = GateUp ? si_.data : sm_.data; + const auto* tokens = GateUp ? row_tokens_.data : nullptr; + auto* activated = GateUp ? activated_.data : nullptr; + auto* down = GateUp ? nullptr : down_.data; + const int splits = GateUp ? gate_splits_ : down_splits_; + if constexpr (Rows <= 16) { + if (prefill && prefill_n_ == 32) { + auto n32 = [&]() { + if constexpr (GateUp) { + if (fused) + return sycl_xmx::prefill_n32(tiles, nt, tokens, quantized, scales, first, second, + nullptr, nullptr, config.swiglu_limit, + config.swiglu_alpha, dep, qm_.data, sm_.data); + } + return sycl_xmx::prefill_n32(tiles, nt, tokens, quantized, scales, first, second, + activated, down, config.swiglu_limit, + config.swiglu_alpha, dep); + }; + if (prefill_load_ != sycl_xmx::PrefillLoad::Block) + return n32.template operator()(); + return n32.template operator()(); + } + } + if constexpr (Rows > 8) { + if constexpr (GateUp) { + if (fused) + return sycl_xmx::projection(tiles, nt, tokens, input, quantized, scales, first, + second, nullptr, nullptr, config.swiglu_limit, + config.swiglu_alpha, dep, qm_.data, sm_.data); + } + return sycl_xmx::projection(tiles, nt, tokens, input, quantized, scales, first, second, + activated, down, config.swiglu_limit, config.swiglu_alpha, dep); + } else { + if (kernel == Kernel::Vector) + return sycl_xmx::vector_dispatch(splits, tiles, nt, Rows, tokens, input, quantized, scales, + first, second, activated, down, config.swiglu_limit, + config.swiglu_alpha, dep); + if (kernel == Kernel::DP4A) + return sycl_xmx::vector_dispatch(splits, tiles, nt, Rows, tokens, input, quantized, scales, first, + second, activated, down, config.swiglu_limit, + config.swiglu_alpha, dep); + if (kernel == Kernel::XMXSplit) { +#define KT_XMX_SPLIT(S) \ + (q8_ ? sycl_xmx::projection(tiles, nt, tokens, input, quantized, scales, first, \ + second, activated, down, config.swiglu_limit, \ + config.swiglu_alpha, dep) \ + : sycl_xmx::projection(tiles, nt, tokens, input, quantized, scales, first, \ + second, activated, down, config.swiglu_limit, \ + config.swiglu_alpha, dep)) + switch (splits) { + case 1: + return KT_XMX_SPLIT(1); + case 2: + return KT_XMX_SPLIT(2); + case 4: + return KT_XMX_SPLIT(4); + default: + return KT_XMX_SPLIT(8); + } +#undef KT_XMX_SPLIT + } + if constexpr (GateUp) { + if (fused) + return sycl_xmx::projection(tiles, nt, tokens, input, quantized, scales, first, + second, nullptr, nullptr, config.swiglu_limit, + config.swiglu_alpha, dep, qm_.data, sm_.data); + } + return q8_ ? sycl_xmx::projection(tiles, nt, tokens, input, quantized, scales, first, second, + activated, down, config.swiglu_limit, config.swiglu_alpha, + dep) + : sycl_xmx::projection(tiles, nt, tokens, input, quantized, scales, first, second, + activated, down, config.swiglu_limit, config.swiglu_alpha, + dep); + } + } + sycl::event submit_prefill_down(Kernel kernel, const sycl::event* dep) { + const auto* tiles = prefill_down_rows_ == prefill_rows_ ? tiles_.data : down_tiles_.data; +#define KT_PREFILL_DOWN(R) submit_projection(kernel, dep, down_tile_count_, false, true, tiles) + switch (prefill_down_rows_) { + case 1: + return KT_PREFILL_DOWN(1); + case 2: + return KT_PREFILL_DOWN(2); + case 4: + return KT_PREFILL_DOWN(4); + case 8: + return KT_PREFILL_DOWN(8); + case 16: + return KT_PREFILL_DOWN(16); + default: + return KT_PREFILL_DOWN(32); + } +#undef KT_PREFILL_DOWN + } + template + sycl::event submit(int nt, int rows, int topk, int routed) { + const Kernel gate = rows == 1 ? decode_gate_ : gate_kernel_; + const Kernel down = rows == 1 ? decode_down_ : down_kernel_; + const bool device_quant = uses_q8(gate) && !(host_quant_ && rows == 1); + const bool use_fused_quant = fused_quant_ && (rows > 1 || fused_decode_); + sycl::event qi; + if (device_quant) + qi = + sycl_xmx::quantize(input_.data, qi_.data, si_.data, rows, config.hidden_size, config.quant_config.group_size); + if (device_quant) remember(sycl_xmx::PipelineProfile::InputQuant, qi); + if constexpr (Rows == 1) { + if (rows == 1 && fused_expert_) { + sycl::event experts; +#define KT_FUSED_EXPERT(S) \ + sycl_xmx::fused_expert_decode(tiles_.data, nt, qi_.data, si_.data, gate_, up_, down_weight_, down_.data, \ + config.swiglu_limit, config.swiglu_alpha, device_quant ? &qi : nullptr) + if (gate == Kernel::XMX || gate_splits_ == 1) + experts = KT_FUSED_EXPERT(1); + else + experts = KT_FUSED_EXPERT(2); +#undef KT_FUSED_EXPERT + remember(sycl_xmx::PipelineProfile::Gate, experts); // includes quantization and down + auto result = sycl_xmx::merge(slot_rows_.data, routing_.data, down_.data, output_.data, rows, topk, + config.hidden_size, experts); + remember(sycl_xmx::PipelineProfile::Merge, result); + return result; + } + } + auto gu = submit_projection(gate, device_quant ? &qi : nullptr, nt, use_fused_quant, rows > 1); + remember(sycl_xmx::PipelineProfile::Gate, gu); + sycl::event middle = gu; + if (uses_q8(down) && !use_fused_quant) { + middle = sycl_xmx::quantize(activated_.data, qm_.data, sm_.data, routed, config.intermediate_size, + config.quant_config.group_size, &gu); + remember(sycl_xmx::PipelineProfile::MiddleQuant, middle); + } + if constexpr (Rows == 1) { + if (rows == 1 && fused_merge_) { + auto result = + q8_ ? sycl_xmx::projection<1, true, false, false, 1, true>( + tiles_.data, nt, nullptr, activated_.data, qm_.data, sm_.data, down_weight_, nullptr, nullptr, + nullptr, 0, 0, &middle, nullptr, nullptr, slot_rows_.data, routing_.data, output_.data, topk) + : sycl_xmx::projection<1, false, false, false, 1, true>( + tiles_.data, nt, nullptr, activated_.data, qm_.data, sm_.data, down_weight_, nullptr, nullptr, + nullptr, 0, 0, &middle, nullptr, nullptr, slot_rows_.data, routing_.data, output_.data, topk); + // Down includes expert merge; no separate Merge event is recorded. + remember(sycl_xmx::PipelineProfile::Down, result); + return result; + } + } + auto d = rows > 1 ? submit_prefill_down(down, &middle) : submit_projection(down, &middle, nt); + remember(sycl_xmx::PipelineProfile::Down, d); + sycl::event result; + if (rows > 1 && vector_prefill_merge_ && config.hidden_size % 256 == 0) + result = sycl_xmx::merge_vector<256>(slot_rows_.data, routing_.data, down_.data, output_.data, rows, topk, + config.hidden_size, d); + else if (rows > 1 && vector_prefill_merge_ && config.hidden_size % 128 == 0) + result = sycl_xmx::merge_vector<128>(slot_rows_.data, routing_.data, down_.data, output_.data, rows, topk, + config.hidden_size, d); + else + result = + sycl_xmx::merge(slot_rows_.data, routing_.data, down_.data, output_.data, rows, topk, config.hidden_size, d); + remember(sycl_xmx::PipelineProfile::Merge, result); + return result; + } + template + void run(int nt, int rows, int topk, int routed) { + if constexpr (Rows == 1) { + if (use_graph_ && rows == 1) { + namespace gx = sycl::ext::oneapi::experimental; + auto& cached = graphs_[topk * (config.num_experts_per_tok + 1) + nt]; + auto& q = sycl_int4::queue(); + if (!cached) { + gx::command_graph graph(q, {gx::property::graph::assume_buffer_outlives_graph{}}); + graph.begin_recording(q); + try { + (void)submit(nt, rows, topk, routed); + } catch (...) { + graph.end_recording(); + throw; + } + graph.end_recording(); + cached = std::make_unique(graph.finalize()); + } + using P = sycl_xmx::PipelineProfile; + const auto begin = profile_.enabled ? P::Clock::now() : P::Time{}; + auto result = q.ext_oneapi_graph(*cached); + const auto submitted = profile_.enabled ? P::Clock::now() : P::Time{}; + sycl_int4::wait_and_throw(result); + if (profile_.enabled) { + frame_.host[P::Submit] = P::us(begin, submitted); + frame_.host[P::Wait] = P::us(submitted, P::Clock::now()); + } + return; + } + } + using P = sycl_xmx::PipelineProfile; + const auto begin = profile_.enabled ? P::Clock::now() : P::Time{}; + auto result = submit(nt, rows, topk, routed); + const auto submitted = profile_.enabled ? P::Clock::now() : P::Time{}; + sycl_int4::wait_and_throw(result); + if (profile_.enabled) { + frame_.host[P::Submit] = P::us(begin, submitted); + frame_.host[P::Wait] = P::us(submitted, P::Clock::now()); + } + } +}; + +template <> +class TP_MOE : public TP_MOE_Common { + public: + using Base = TP_MOE_Common; + explicit TP_MOE(const GeneralMOEConfig& config) : Base(config) { + const sycl_xmx::Options options; + const char* numa = options.get("KT_SYCL_XMX_NUMA_DISPATCH"); + if (numa && std::string(numa) != "0" && std::string(numa) != "1") + throw std::runtime_error("KT_SYCL_XMX_NUMA_DISPATCH must be 0 or 1"); + use_numa_dispatch_ = numa && std::string(numa) == "1"; + } + static int implementation_revision() { return 7; } + void load_weights() override { + this->weights_loaded = false; + this->tps[0]->load_weights(this->config); + this->weights_loaded = true; + } + void forward(int rows, int topk, const int64_t* ids, const float* weights, const void* input, void* output, + bool incremental = false) override { + if (!this->weights_loaded) throw std::runtime_error("XMX weights not loaded"); + using P = sycl_xmx::PipelineProfile; + auto& profile = this->tps[0]->pipeline_profile(); + const auto begin = profile.enabled ? P::Clock::now() : P::Time{}; + if (use_numa_dispatch_) { + // Worker exceptions must be captured here: the generic NUMA distributor + // does not propagate exceptions from its worker threads. + std::exception_ptr failure; + this->config.pool->dispense_backend()->do_numa_job([&](int part) { + try { + this->tps[part]->forward(rows, topk, ids, weights, input, this->local_output_numa[part]); + } catch (...) { + failure = std::current_exception(); + } + }); + if (failure) std::rethrow_exception(failure); + } else { + this->tps[0]->forward(rows, topk, ids, weights, input, this->local_output_numa[0]); + } + const auto converted = profile.enabled ? P::Clock::now() : P::Time{}; + merge_results(rows, output, incremental); + if (profile.enabled && rows > 0) { + const auto end = P::Clock::now(); + profile.record_wrapper(rows, P::us(begin, end), P::us(converted, end)); + } + } + std::map profile_snapshot() { return this->tps[0]->pipeline_profile().snapshot(); } + void reset_profile() { this->tps[0]->pipeline_profile().reset(); } + void forward_binding(intptr_t rows, int topk, intptr_t ids, intptr_t weights, intptr_t input, intptr_t output, + bool incremental) { + if (!rows) throw std::runtime_error("XMX qlen pointer is null"); + forward(*reinterpret_cast(rows), topk, reinterpret_cast(ids), + reinterpret_cast(weights), reinterpret_cast(input), + reinterpret_cast(output), incremental); + } + void merge_results(int rows, void* output) override { merge_results(rows, output, false); } + void merge_results(int rows, void* output, bool incremental) override { + auto* dst = static_cast(output); + const float* src = this->local_output_numa[0]; + for (int j = 0; j < rows * this->config.hidden_size; j += 8) { + __m256 value = _mm256_loadu_ps(src + j); + if (incremental) value = _mm256_add_ps(value, avx2::load_bf16_to_fp32(dst + j)); + avx2::store_fp32_to_bf16(dst + j, value); + } + } + + private: + bool use_numa_dispatch_ = false; +}; +#endif diff --git a/kt-kernel/operators/sycl/gptq_int4_xmx-options.hpp b/kt-kernel/operators/sycl/gptq_int4_xmx-options.hpp new file mode 100644 index 00000000..48e42d40 --- /dev/null +++ b/kt-kernel/operators/sycl/gptq_int4_xmx-options.hpp @@ -0,0 +1,49 @@ +// SPDX-License-Identifier: Apache-2.0 +#ifndef KT_SYCL_GPTQ_INT4_XMX_OPTIONS_HPP +#define KT_SYCL_GPTQ_INT4_XMX_OPTIONS_HPP + +#include +#include +#include +#include +#include + +namespace sycl_xmx { +// Defaults are selected once when constructing an executor. Explicit settings +// remain available for diagnostics; selecting a preset never changes global env. +class Options { + public: + Options() { + const char* value = std::getenv("KT_SYCL_XMX_PROFILE"); + if (value && std::string(value) != "panther-lake") + throw std::runtime_error("KT_SYCL_XMX_PROFILE must be panther-lake or unset"); + panther_lake_ = value != nullptr; + } + bool panther_lake() const { return panther_lake_; } + const char* get(const char* name) const { + if (const char* value = std::getenv(name)) return value; + if (!panther_lake_) return nullptr; + static constexpr std::pair defaults[] = { + {"KT_SYCL_XMX_MODE", "int8"}, + {"KT_SYCL_XMX_DECODE_GATE_KERNEL", "xmx-split"}, + {"KT_SYCL_XMX_GATE_SPLITS", "2"}, + {"KT_SYCL_XMX_DOWN_SPLITS", "1"}, + {"KT_SYCL_XMX_PREFILL_ROWS", "8"}, + {"KT_SYCL_XMX_PREFILL_DOWN_ROWS", "16"}, + {"KT_SYCL_XMX_PREFILL_N", "32"}, + {"KT_SYCL_XMX_PREFILL_LOAD", "gather"}, + {"KT_SYCL_XMX_PREFILL_MERGE", "vector"}, + {"KT_SYCL_XMX_FUSED_EXPERT", "1"}, + {"KT_SYCL_XMX_FUSED_QUANT", "prefill"}, + {"KT_SYCL_XMX_NUMA_DISPATCH", "1"}, + }; + for (const auto& [key, value] : defaults) + if (std::strcmp(key, name) == 0) return value; + return nullptr; + } + + private: + bool panther_lake_ = false; +}; +} // namespace sycl_xmx +#endif diff --git a/kt-kernel/operators/sycl/gptq_int4_xmx-prefill.hpp b/kt-kernel/operators/sycl/gptq_int4_xmx-prefill.hpp new file mode 100644 index 00000000..84520485 --- /dev/null +++ b/kt-kernel/operators/sycl/gptq_int4_xmx-prefill.hpp @@ -0,0 +1,192 @@ +// SPDX-License-Identifier: Apache-2.0 +#ifndef KT_SYCL_GPTQ_INT4_XMX_PREFILL_HPP +#define KT_SYCL_GPTQ_INT4_XMX_PREFILL_HPP +#include "gptq_int4_xmx-kernels.hpp" + +namespace sycl_xmx { +enum class PrefillLoad { Block, Gather }; + +// Routed rows are not contiguous. Gather dwords across the rows in one vector +// instead of issuing a separate 32-byte block load for every DPAS input row. +template +ESIMD_INLINE es::simd prefill_gather(const int8_t* input, es::simd offsets, + es::simd_mask mask, uint32_t column) { + es::simd words = es::gather(reinterpret_cast(input), + es::simd(offsets + column), + mask, es::simd(0)); + return words.template bit_cast_view(); +} + +// Prefill-only N32 tile: two N16 DPAS blocks share activation loads. No +// resident weight conversion; precision/group boundaries match the N16 path. +template +inline sycl::event prefill_n32(const Tile* tiles, int tile_count, const int* row_tokens, const int8_t* input, + const float* input_scales, const Weight& first, const Weight* second, + uint16_t* activated, float* down, float limit, float alpha, + const sycl::event* dep = nullptr, int8_t* qoutput = nullptr, + float* output_scales = nullptr) { + static_assert(!FusedQuant || GateUp); + static_assert(Rows == 1 || Rows == 2 || Rows == 4 || Rows == 8 || Rows == 16); + constexpr int D = Rows > 8 ? 8 : Rows; + const int n = first.n, k = first.k, gs = first.gs, groups = k / gs; + const auto* w0 = first.packed.data; + const auto* s0 = first.scales.data; + const auto* w1 = second ? second->packed.data : nullptr; + const auto* s1 = second ? second->scales.data : nullptr; + return sycl_int4::queue().submit([&](sycl::handler& h) { + if (dep) h.depends_on(*dep); + const size_t work = static_cast(tile_count) * (n / 32); + auto compute = [=](size_t index, int lane) SYCL_ESIMD_FUNCTION { + if constexpr (FusedQuant) es::slm_init(); + const int nb = (index % (n / 32)) * 2; + const Tile tile = tiles[index / (n / 32)]; + es::simd acc[2], up[2]; +#pragma unroll + for (int block = 0; block < 2; ++block) { + acc[block] = 0; + up[block] = 0; + } + es::simd input_offsets = 0; + es::simd_mask input_mask = 0; + if constexpr (Load == PrefillLoad::Gather) { +#pragma unroll + for (int r = 0; r < Rows; ++r) { + if (r < tile.rows) { + const int row = GateUp ? row_tokens[tile.begin + r] : tile.begin + r; + input_offsets.template select<8, 1>(r * 8) = es::simd(row * k, 4); + input_mask.template select<8, 1>(r * 8) = 1; + } + } + } + for (int g = 0; g < groups; ++g) { + es::simd c[2], u[2]; +#pragma unroll + for (int block = 0; block < 2; ++block) { + c[block] = 0; + u[block] = 0; + } + for (int j = 0; j < gs; j += 32) { + es::simd a = 0; + if constexpr (Load == PrefillLoad::Gather) { + a = prefill_gather(input, input_offsets, input_mask, g * gs + j); + } else { +#pragma unroll + for (int r = 0; r < Rows; ++r) { + if (r < tile.rows) { + const int row = GateUp ? row_tokens[tile.begin + r] : tile.begin + r; + es::simd value; + value.copy_from(input + static_cast(row) * k + g * gs + j); + a.template select<32, 1>(r * 32) = value; + } + } + } +#pragma unroll + for (int block = 0; block < 2; ++block) { + const size_t wi = (static_cast(tile.expert) * (n / 16) + nb + block) * (k / 8) * 16; + es::simd b; + b.copy_from(w0 + wi + (g * gs + j) / 8 * 16); + b ^= 0x88888888u; +#pragma unroll + for (int r = 0; r < Rows; r += D) { + es::simd partial = c[block].template select(r * 16); + es::simd av = a.template select(r * 32); + c[block].template select(r * 16) = + xmx::dpas<8, D, int, int, uint32_t, int8_t, xmx::dpas_argument_type::s4, xmx::dpas_argument_type::s8>( + partial, b, av); + } + if constexpr (GateUp) { + b.copy_from(w1 + wi + (g * gs + j) / 8 * 16); + b ^= 0x88888888u; +#pragma unroll + for (int r = 0; r < Rows; r += D) { + es::simd partial = u[block].template select(r * 16); + es::simd av = a.template select(r * 32); + u[block].template select(r * 16) = + xmx::dpas<8, D, int, int, uint32_t, int8_t, xmx::dpas_argument_type::s4, + xmx::dpas_argument_type::s8>(partial, b, av); + } + } + } + } +#pragma unroll + for (int block = 0; block < 2; ++block) { + const size_t si = (static_cast(tile.expert) * (n / 16) + nb + block) * groups * 16; + es::simd sc, su; + sc.copy_from(s0 + si + g * 16); + if constexpr (GateUp) su.copy_from(s1 + si + g * 16); +#pragma unroll + for (int r = 0; r < Rows; ++r) { + if (r < tile.rows) { + const int row = GateUp ? row_tokens[tile.begin + r] : tile.begin + r; + const float sa = input_scales[row * groups + g]; + acc[block].template select<16, 1>(r * 16) += + es::convert(c[block].template select<16, 1>(r * 16).read()) * (sc * sa); + if constexpr (GateUp) + up[block].template select<16, 1>(r * 16) += + es::convert(u[block].template select<16, 1>(r * 16).read()) * (su * sa); + } + } + } + } +#pragma unroll + for (int r = 0; r < Rows; ++r) { + if (r >= tile.rows) continue; + float maximum = 0; +#pragma unroll + for (int block = 0; block < 2; ++block) { + es::simd value = acc[block].template select<16, 1>(r * 16); + if constexpr (GateUp) { + es::simd uv = up[block].template select<16, 1>(r * 16); + if (limit > 0) { + value = es::min(value, limit); + if (alpha > 0) value = es::max(value, -limit); + uv = es::min(es::max(uv, -limit), limit); + } + value = value / (1.0f + es::exp(-value * (alpha > 0 ? alpha : 1.0f))) * (alpha > 0 ? uv + 1 : uv); + } + es::simd bits = value.template bit_cast_view(); + bits += 0x7fffu + ((bits >> 16) & 1u); + bits &= 0xffff0000u; + es::simd rounded = bits.template bit_cast_view(); + if constexpr (FusedQuant) { + acc[block].template select<16, 1>(r * 16) = rounded; + maximum = es::max(maximum, es::hmax(es::abs(rounded))); + } else if constexpr (GateUp) { + es::convert(bits >> 16) + .copy_to(activated + static_cast(tile.begin + r) * n + (nb + block) * 16); + } else { + rounded.copy_to(down + static_cast(tile.begin + r) * n + (nb + block) * 16); + } + } + if constexpr (FusedQuant) + es::slm_scatter(es::simd((r * 4 + lane) * sizeof(float)), es::simd(maximum)); + } + if constexpr (FusedQuant) { + es::barrier(); +#pragma unroll + for (int r = 0; r < Rows; ++r) { + if (r >= tile.rows) continue; + auto maxima = es::slm_block_load(r * 4 * sizeof(float)); + const float scale = es::hmax(maxima) / 127.0f; + if (lane == 0) output_scales[(tile.begin + r) * (n / 128) + nb / 8] = scale; +#pragma unroll + for (int block = 0; block < 2; ++block) { + es::simd value = acc[block].template select<16, 1>(r * 16); + value = es::rnde(value / (scale > 0 ? scale : 1.0f)); + es::convert(es::min(es::max(value, -127.0f), 127.0f)) + .copy_to(qoutput + static_cast(tile.begin + r) * n + (nb + block) * 16); + } + } + } + }; + if constexpr (FusedQuant) { + h.parallel_for(sycl::nd_range<1>(work, 4), [=](sycl::nd_item<1> it) SYCL_ESIMD_KERNEL { + compute(it.get_global_linear_id(), it.get_local_linear_id()); + }); + } else { + h.parallel_for(sycl::range<1>(work), [=](sycl::id<1> id) SYCL_ESIMD_KERNEL { compute(id[0], 0); }); + } + }); +} +} // namespace sycl_xmx +#endif diff --git a/kt-kernel/operators/sycl/gptq_int4_xmx-profile.hpp b/kt-kernel/operators/sycl/gptq_int4_xmx-profile.hpp new file mode 100644 index 00000000..7497cea4 --- /dev/null +++ b/kt-kernel/operators/sycl/gptq_int4_xmx-profile.hpp @@ -0,0 +1,125 @@ +// SPDX-License-Identifier: Apache-2.0 +#ifndef KT_SYCL_GPTQ_INT4_XMX_PROFILE_HPP +#define KT_SYCL_GPTQ_INT4_XMX_PROFILE_HPP + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "gptq_int4_sycl-runtime.hpp" + +namespace sycl_xmx { +// Opt-in diagnostics only. Host and device clocks must never be subtracted +// from one another. Event queries are deferred until the whole chain finishes. +class PipelineProfile { + public: + using Clock = std::chrono::steady_clock; + using Time = Clock::time_point; + static double us(Time start, Time end) { return std::chrono::duration(end - start).count(); } + enum Host { Route, InputCopy, HostQuant, Submit, Wait, OutputCopy, Executor, HostCount }; + enum Device { InputQuant, Gate, MiddleQuant, Down, Merge, DeviceCount }; + struct Frame { + std::array host{}; + std::array events; + std::array present{}; + int routed = 0, tiles = 0; + }; + explicit PipelineProfile(const char* value, int layer = 0) : layer_(layer) { + const std::string mode = value ? value : "0"; + if (mode != "0" && mode != "host" && mode != "device") + throw std::runtime_error("KT_SYCL_XMX_TRACE must be 0, host or device"); + enabled = mode != "0"; + device = mode == "device"; + const char* every = std::getenv("KT_SYCL_XMX_TRACE_EVERY"); + if (every) { + const std::string count(every); + if (count.empty() || count.find_first_not_of("0123456789") != std::string::npos || count.size() > 9) + throw std::runtime_error("KT_SYCL_XMX_TRACE_EVERY must be an integer from 0 to 999999999"); + every_ = std::stoul(count); + } + if (device && !sycl_int4::queue().has_property()) + throw std::runtime_error("KT_SYCL_XMX_TRACE=device requires KT_SYCL_PROFILE=1 before queue construction"); + } + void record(int rows, const Frame& f) { + static constexpr const char* host_names[] = {"route", "input_copy", "host_quant", "submit", + "wait", "output_copy", "executor"}; + static constexpr const char* device_names[] = {"input_quant", "gate", "middle_quant", "down", "merge"}; + std::map sample; + const std::string prefix = rows == 1 ? "decode." : "prefill."; + sample["calls"] = 1; + sample["rows"] = rows; + sample["routed"] = f.routed; + sample["tiles"] = f.tiles; + for (int j = 0; j < HostCount; ++j) sample[std::string("host_") + host_names[j] + "_us"] = f.host[j]; + if (device) { + uint64_t first = std::numeric_limits::max(), last = 0; + double kernels = 0; + for (int j = 0; j < DeviceCount; ++j) { + double elapsed = 0; + if (f.present[j]) { + const auto begin = f.events[j].get_profiling_info(); + const auto end = f.events[j].get_profiling_info(); + elapsed = (end - begin) / 1000.0; + first = std::min(first, begin); + last = std::max(last, end); + sample["device_commands"] += 1; + } + sample[std::string("device_") + device_names[j] + "_us"] = elapsed; + kernels += elapsed; + } + const double span = last ? (last - first) / 1000.0 : 0; + sample["device_kernels_us"] = kernels; + sample["device_span_us"] = span; + // The executor is a serial dependency chain, so there is no overlap. + sample["device_gaps_us"] = std::max(0.0, span - kernels); + } + std::lock_guard lock(mutex_); + for (const auto& [key, value] : sample) totals_[prefix + key] += value; + } + void record_wrapper(int rows, double wrapper, double conversion) { + std::lock_guard lock(mutex_); + const std::string prefix = rows == 1 ? "decode." : "prefill."; + totals_[prefix + "host_wrapper_us"] += wrapper; + totals_[prefix + "host_bf16_us"] += conversion; + if (every_ && static_cast(totals_[prefix + "calls"]) % every_ == 0) { + std::ostringstream report; + report.imbue(std::locale::classic()); + report << "KT_SYCL_PIPELINE {\"layer\":" << layer_ << ",\"phase\":\"" << (rows == 1 ? "decode" : "prefill") + << "\",\"totals\":{"; + bool first = true; + for (const auto& [key, value] : totals_) { + if (key.rfind(prefix, 0) != 0) continue; + if (!first) report << ','; + first = false; + report << '\"' << key.substr(prefix.size()) << "\":" << std::setprecision(12) << value; + } + report << "}}\n"; + // Logging is deliberately opt-in and outside recorded host timings. + std::fputs(report.str().c_str(), stderr); + } + } + std::map snapshot() const { + std::lock_guard lock(mutex_); + return totals_; + } + void reset() { + std::lock_guard lock(mutex_); + totals_.clear(); + } + bool enabled = false, device = false; + + private: + int layer_ = 0; + uint64_t every_ = 0; + mutable std::mutex mutex_; + std::map totals_; +}; +} // namespace sycl_xmx +#endif diff --git a/kt-kernel/python/experts_base.py b/kt-kernel/python/experts_base.py index 647ab294..827a60c1 100644 --- a/kt-kernel/python/experts_base.py +++ b/kt-kernel/python/experts_base.py @@ -225,6 +225,11 @@ class KExpertsCPUBuffer: temp_bs: int = 0 temp_buffer: tuple = tuple() buffer_depth: int = 2 + # Native tasks/captured host callbacks store raw pointers to these CPU-only + # scalars. Unlike DMA buffers, they are not covered by the pinned allocator's + # stream events. Keep one immutable set per shape for the process lifetime; + # replacing temp_buffer must never invalidate a queued task's qlen pointer. + _qlen_cache: Dict[Tuple[int, int], List[torch.Tensor]] = {} @classmethod def get_buffer(cls, hidden_states: torch.Tensor, num_experts_per_tok): @@ -258,10 +263,14 @@ class KExpertsCPUBuffer: torch.zeros((batch_size, hidden_size), device="cpu", pin_memory=pin_memory, dtype=torch.bfloat16) for _ in range(cls.buffer_depth) ] - bsz_tensor_cpu = [ - torch.full((1,), batch_size, device="cpu", dtype=torch.int32, pin_memory=pin_memory) - for _ in range(cls.buffer_depth) - ] + qlen_key = (batch_size, cls.buffer_depth) + if qlen_key not in cls._qlen_cache: + cls._qlen_cache[qlen_key] = [ + torch.full((1,), batch_size, device="cpu", dtype=torch.int32, pin_memory=pin_memory) + for _ in range(cls.buffer_depth) + ] + # These lengths describe allocated shapes, not mutable request metadata. + bsz_tensor_cpu = cls._qlen_cache[qlen_key] output_gpu = [ torch.zeros((batch_size, hidden_size), device=hidden_states.device, dtype=hidden_states.dtype) for _ in range(cls.buffer_depth) @@ -937,4 +946,3 @@ class BaseMoEWrapper(_MoEBase, ABC): KExpertsCPUBuffer.capture_buffers.clear() KExpertsCPUBuffer.temp_bs = 0 KExpertsCPUBuffer.temp_buffer = tuple() - diff --git a/kt-kernel/python/utils/amx.py b/kt-kernel/python/utils/amx.py index f7a02c72..20a60e7a 100644 --- a/kt-kernel/python/utils/amx.py +++ b/kt-kernel/python/utils/amx.py @@ -42,6 +42,7 @@ AVXVNNI256GPTQInt4_MOE = getattr(_moe_mod, "AVXVNNI256GPTQInt4_MOE", None) AVXVNNI256GPTQInt4Packed_MOE = getattr(_moe_mod, "AVXVNNI256GPTQInt4Packed_MOE", None) AVXVNNI256RawInt4_MOE = getattr(_moe_mod, "AVXVNNI256RawInt4_MOE", None) SYCLGPTQInt4_MOE = getattr(_moe_mod, "SYCLGPTQInt4_MOE", None) +SYCLXMXGPTQInt4_MOE = getattr(_moe_mod, "SYCLXMXGPTQInt4_MOE", None) _HAS_AMXINT4_SUPPORT = AMXInt4_MOE is not None _HAS_AMXINT8_SUPPORT = AMXInt8_MOE is not None @@ -114,6 +115,29 @@ def _supports_avxvnni256_rawint4_group_size(group_size: Optional[int]) -> bool: return group_size > 0 and group_size % 32 == 0 and group_size <= _AVXVNNI256_RAW_INT4_MAX_GROUP_SIZE +def _select_sycl_int4_backend(): + """Keep the experimental XMX executor opt-in; never silently fall back.""" + profile = os.getenv("KT_SYCL_XMX_PROFILE") + if profile is not None and profile != "panther-lake": + raise ValueError("KT_SYCL_XMX_PROFILE must be panther-lake or unset") + backend = os.getenv("KT_SYCL_INT4_BACKEND", "xmx" if profile else "legacy").strip().lower() + if profile and backend != "xmx": + raise ValueError("KT_SYCL_XMX_PROFILE=panther-lake requires the xmx backend") + if backend == "legacy": + cls = SYCLGPTQInt4_MOE + elif backend == "xmx": + cls = SYCLXMXGPTQInt4_MOE + else: + raise ValueError("KT_SYCL_INT4_BACKEND must be legacy or xmx") + if cls is None: + raise RuntimeError(f"SYCL {backend} backend not compiled; rebuild kt-kernel with CPUINFER_USE_SYCL=1") + if profile and getattr(cls, "implementation_revision", lambda: 0)() < 7: + raise RuntimeError( + "Panther Lake profile requires XMX revision >= 7; rebuild kt-kernel with ./install.sh --sycl" + ) + return cls + + def _select_gptq_int4_backend(group_size: Optional[int] = None): forced = os.getenv("KT_GPTQ_INT4_BACKEND", "").strip().lower() avxvnni_group_supported = _supports_avxvnni256_gptq_int4_group_size(group_size) @@ -663,9 +687,7 @@ class NativeMoEWrapper(BaseMoEWrapper): "CPUINFER_USE_SYCL=1 using a SYCL compiler such as icpx." ) if method == "NVFP4" and not _HAS_AVX2_MXFP4_SUPPORT: - raise RuntimeError( - "NVFP4 needs the AVX2 FP4 backend (AVX2MXFP4_MOE), which is not compiled in." - ) + raise RuntimeError("NVFP4 needs the AVX2 FP4 backend (AVX2MXFP4_MOE), which is not compiled in.") if method == "MXFP4" and not (_HAS_MXFP4_SUPPORT or _HAS_AVX2_MXFP4_SUPPORT): raise RuntimeError( "MXFP4 backend not available. Required ISA (any one of):\n" @@ -1007,7 +1029,7 @@ class NativeMoEWrapper(BaseMoEWrapper): moe_config.quant_config.group_size = actual_gs moe_config.quant_config.zero_point = False _preflight_sycl_device() - self.moe = SYCLGPTQInt4_MOE(moe_config) + self.moe = _select_sycl_int4_backend()(moe_config) elif self.method == "BF16": # BF16 has no quantization config needed # Prefer AMX backend, fall back to AVX2 @@ -1102,15 +1124,11 @@ class NativeMoEWrapper(BaseMoEWrapper): each rank's local H2D copy of expert ``e``. """ if self.method != "FP8": - raise RuntimeError( - "run_layerwise_fp8_batch is only valid for the block-FP8 NativeMoEWrapper backend" - ) + raise RuntimeError("run_layerwise_fp8_batch is only valid for the block-FP8 NativeMoEWrapper backend") if self.moe is None: raise RuntimeError("MoE instance not initialized; cannot run FP8 layerwise transport.") if not hasattr(self.moe, "run_layerwise_fp8_batch"): - raise NotImplementedError( - "The installed kt-kernel extension does not expose run_layerwise_fp8_batch." - ) + raise NotImplementedError("The installed kt-kernel extension does not expose run_layerwise_fp8_batch.") # The native batch calls the shared NUMA distributor directly. Drain # CPUInfer first so it cannot race a previously queued task against the # non-reentrant distributor state. diff --git a/kt-kernel/scripts/prepare_sglang.py b/kt-kernel/scripts/prepare_sglang.py new file mode 100644 index 00000000..c332bda2 --- /dev/null +++ b/kt-kernel/scripts/prepare_sglang.py @@ -0,0 +1,75 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +"""Prepare pinned SGLang source without modifying its submodule checkout. + +Install the returned directory with ``python -m pip install --no-deps DIR/python`` +inside an already configured KT/SGLang environment. The exported dependency +metadata follows this checkout's kt-kernel version; runtime source is unchanged. +""" + +import argparse +import io +import runpy +import tomllib +import json +from pathlib import Path +import subprocess +import tarfile +import tempfile + +PINNED_SGLANG = "541ddc37cbc92c60dc748db5ff1a2aad0b069a80" +REPO_ROOT = Path(__file__).resolve().parents[2] + + +def prepare(source: Path, output: Path) -> Path: + source, output = source.resolve(), output.absolute() + if output.exists() or output.is_symlink(): + raise ValueError(f"Output already exists; select a new directory: {output}") + revision = subprocess.check_output(["git", "-C", str(source), "rev-parse", "HEAD"], text=True).strip() + if revision != PINNED_SGLANG: + raise ValueError( + f"Expected pinned SGLang {PINNED_SGLANG}, got {revision}; " "initialize the repository submodules first" + ) + output.parent.mkdir(parents=True, exist_ok=True) + # Export tracked files from the exact pin. Untracked/local source edits are + # neither copied nor overwritten. The temporary destination is isolated from + # enclosing git repositories. + archive = subprocess.check_output(["git", "-C", str(source), "archive", "--format=tar", revision]) + with tempfile.TemporaryDirectory(prefix=".kt-sglang-", dir=output.parent) as tmp: + staging = Path(tmp) / "sglang" + staging.mkdir() + with tarfile.open(fileobj=io.BytesIO(archive)) as tar: + tar.extractall(staging, filter="data") + # The pinned submodule can still depend on an older released kernel. + # Keep the exact source pin while aligning only that package constraint + # to the kernel built from this checkout. Do not relax Torch/other pins. + kernel_version = runpy.run_path(str(REPO_ROOT / "version.py"))["__version__"] + project = staging / "python/pyproject.toml" + text = project.read_text() + requirements = tomllib.loads(text)["project"]["dependencies"] + kernel_requirements = [r for r in requirements if r.startswith("kt-kernel==")] + if len(kernel_requirements) != 1 or text.count('"' + kernel_requirements[0] + '"') != 1: + raise ValueError("Expected exactly one pinned kt-kernel dependency") + project.write_text(text.replace('"' + kernel_requirements[0] + '"', '"kt-kernel==' + kernel_version + '"')) + manifest = { + "source_commit": revision, + "kernel_version": kernel_version, + } + (staging / ".kt-sglang-source.json").write_text(json.dumps(manifest, indent=2) + "\n") + staging.rename(output) + return output + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", required=True, type=Path) + args = parser.parse_args() + try: + result = prepare(REPO_ROOT / "third_party/sglang", args.output) + except (ValueError, subprocess.CalledProcessError) as error: + parser.exit(1, f"{error}\n") + print(result) + + +if __name__ == "__main__": + main() diff --git a/kt-kernel/setup.py b/kt-kernel/setup.py index 64718b4e..d0e2197d 100644 --- a/kt-kernel/setup.py +++ b/kt-kernel/setup.py @@ -25,6 +25,7 @@ Environment knobs (export before running pip install .): CPUINFER_ENABLE_AVX512_BF16=OFF ON/OFF -> -DLLAMA_AVX512_BF16 CPUINFER_ENABLE_AVX512_VBMI=OFF ON/OFF -> -DLLAMA_AVX512_VBMI (required for FP8 MoE) CPUINFER_ENABLE_CPPTRACE=ON/OFF ON/OFF -> -DKTRANSFORMERS_ENABLE_CPPTRACE (debug-only) + CPUINFER_LLAMA_BUILD_COPY=ON/OFF Apply llama.cpp patches in the build directory, not the submodule CPUINFER_ENABLE_ONEDNN_VNNI=OFF ON/OFF -> -DKTRANSFORMERS_CPU_USE_ONEDNN_VNNI CPUINFER_ONEDNN_SOURCE_DIR=/path Forward to -DKTRANSFORMERS_ONEDNN_SOURCE_DIR CPUINFER_BLIS_ROOT=/path/to/blis Forward to -DBLIS_ROOT @@ -727,6 +728,7 @@ class CMakeBuild(build_ext): _forward_str_env(cmake_args, "CPUINFER_LTO_JOBS", "CPUINFER_LTO_JOBS") _forward_str_env(cmake_args, "CPUINFER_LTO_MODE", "CPUINFER_LTO_MODE") _forward_bool_env(cmake_args, "CPUINFER_ENABLE_CPPTRACE", "KTRANSFORMERS_ENABLE_CPPTRACE") + _forward_bool_env(cmake_args, "CPUINFER_LLAMA_BUILD_COPY", "KTRANSFORMERS_LLAMA_BUILD_COPY") _forward_bool_env(cmake_args, "CPUINFER_ENABLE_ONEDNN_VNNI", "KTRANSFORMERS_CPU_USE_ONEDNN_VNNI") _forward_str_env(cmake_args, "CPUINFER_ONEDNN_SOURCE_DIR", "KTRANSFORMERS_ONEDNN_SOURCE_DIR") diff --git a/kt-kernel/test/per_commit/test_inference_qlen_lifetime.py b/kt-kernel/test/per_commit/test_inference_qlen_lifetime.py new file mode 100644 index 00000000..8e17c884 --- /dev/null +++ b/kt-kernel/test/per_commit/test_inference_qlen_lifetime.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU-only regressions for batch-size pointers consumed by async native tasks.""" + +import ctypes +import gc +import importlib.util +from pathlib import Path +import weakref +import sys +from types import ModuleType, SimpleNamespace + +import pytest +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +from ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=10, suite="default") + + +@pytest.fixture +def buffers(monkeypatch): + path = Path(__file__).parents[2] / "python/experts_base.py" + spec = importlib.util.spec_from_file_location("kt_kernel._test_experts_base", path) + module = importlib.util.module_from_spec(spec) + # The buffer owner is pure Python; do not require a built native extension. + fake_package = ModuleType("kt_kernel") + fake_package.kt_kernel_ext = SimpleNamespace() + with monkeypatch.context() as imports: + imports.setitem(sys.modules, "kt_kernel", fake_package) + spec.loader.exec_module(module) + cls = module.KExpertsCPUBuffer + # Test object ownership without requiring a CUDA pinned allocator/device. + for name in ("zeros", "full"): + original = getattr(torch, name) + + def cpu_factory(*args, _original=original, **kwargs): + kwargs.pop("pin_memory", None) + return _original(*args, **kwargs) + + monkeypatch.setattr(torch, name, cpu_factory) + monkeypatch.setattr(cls, "capture_bs", []) + monkeypatch.setattr(cls, "capture_buffers", {}) + monkeypatch.setattr(cls, "temp_bs", 0) + monkeypatch.setattr(cls, "temp_buffer", ()) + monkeypatch.setattr(cls, "_qlen_cache", {}, raising=False) + return cls + + +def get_buffer(cls, rows): + return cls.get_buffer(torch.empty(rows, 4, dtype=torch.bfloat16), 2) + + +def test_native_qlen_pointer_survives_temporary_shape_replacement(buffers): + old = get_buffer(buffers, 2) + scalar = weakref.ref(old[5][0]) + raw_pointer = old[5][0].data_ptr() + del old + get_buffer(buffers, 5) + gc.collect() + # A native host callback retains only raw_pointer, not a Python tensor owner. + assert scalar() is not None, "async native qlen pointer lost its owner on shape change" + assert ctypes.c_int32.from_address(raw_pointer).value == 2 + + +def test_qlen_address_is_stable_when_shape_returns(buffers): + old = get_buffer(buffers, 17) + addresses = [tensor.data_ptr() for tensor in old[5]] + get_buffer(buffers, 64) + again = get_buffer(buffers, 17) + assert [tensor.data_ptr() for tensor in again[5]] == addresses + assert all(tensor.item() == 17 for tensor in again[5]) + + +def test_cached_lengths_are_distinct_and_not_mutated_by_other_shapes(buffers): + old = get_buffer(buffers, 512) + get_buffer(buffers, 22) + assert all(tensor.dtype == torch.int32 and tensor.item() == 512 for tensor in old[5]) + assert all(tensor.item() == 22 for tensor in buffers.temp_buffer[5]) + + +def test_repeated_shape_flips_reuse_a_small_immutable_length_cache(buffers): + for _ in range(5): + for rows in (2, 17, 512, 22): + get_buffer(buffers, rows) + assert len(buffers._qlen_cache) == 4 + for (rows, depth), lengths in buffers._qlen_cache.items(): + assert len(lengths) == depth + assert all(length.item() == rows for length in lengths) + + +def test_buffer_depth_is_part_of_length_cache_key(buffers, monkeypatch): + first = get_buffer(buffers, 17)[5] + monkeypatch.setattr(buffers, "buffer_depth", 3) + # Force allocation of a new temporary shape, then return with the new depth. + get_buffer(buffers, 22) + changed = get_buffer(buffers, 17)[5] + assert len(first) == 2 and len(changed) == 3 + assert all(length.item() == 17 for length in first + changed) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/kt-kernel/test/per_commit/test_sycl_install_entry.py b/kt-kernel/test/per_commit/test_sycl_install_entry.py new file mode 100644 index 00000000..87137475 --- /dev/null +++ b/kt-kernel/test/per_commit/test_sycl_install_entry.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Exercise the public installer with isolated dependency/build/pip stand-ins.""" + +import json +import os +from pathlib import Path +import shutil +import subprocess +import sys +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +from ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=10, suite="default") +REPO = Path(__file__).resolve().parents[3] + + +@pytest.fixture +def installer(tmp_path): + root = tmp_path / "checkout with spaces" + root.mkdir() + shutil.copy2(REPO / "install.sh", root / "install.sh") + (root / "version.py").write_text('__version__ = "0.7.2"\n') + (root / "third_party/sglang/python").mkdir(parents=True) + (root / "third_party/sglang/python/pyproject.toml").write_text('[project]\nname="sglang-kt"\n') + (root / "kt-kernel/scripts").mkdir(parents=True) + (root / "kt-kernel/scripts/prepare_sglang.py").write_text( + 'import pathlib,sys\np=pathlib.Path(sys.argv[sys.argv.index("--output")+1]); (p/"python").mkdir(parents=True)\n' + ) + (root / "kt-kernel/install.sh").write_text( + '#!/bin/bash\npython3 -c \'import json,os,sys; print(json.dumps({"kind":"kernel", "args":sys.argv[1:], ' + '"env":{k:v for k,v in os.environ.items() if k.startswith("CPUINFER_")}}))\' "$@" >> "$INSTALL_EVENTS"\n' + ) + tools = tmp_path / "tools" + tools.mkdir() + python = tools / "python3" + python.write_text( + f"#!{sys.executable}\n" + "import json,os,sys\n" + "args=sys.argv[1:]\n" + 'if args[:2]==["-m","pip"]:\n' + ' with open(os.environ["INSTALL_EVENTS"],"a") as f: f.write(json.dumps({"kind":"pip","args":args,"cwd":os.getcwd()})+"\\n")\n' + 'elif args[:1]==["-c"] and "SYCLXMXGPTQInt4_MOE" in args[1]:\n' + ' with open(os.environ["INSTALL_EVENTS"],"a") as f: f.write(json.dumps({"kind":"preflight", "script":args[1], "ld_library_path":os.environ.get("LD_LIBRARY_PATH"), "ld_preload":os.environ.get("LD_PRELOAD")})+"\\n")\n' + 'else: os.execv(os.environ["INSTALL_REAL_PYTHON"],[os.environ["INSTALL_REAL_PYTHON"],*args])\n' + ) + python.chmod(0o755) + (tools / "icpx").write_text("#!/bin/sh\nexit 0\n") + (tools / "icpx").chmod(0o755) + env = {k: v for k, v in os.environ.items() if not k.startswith(("CPUINFER_", "SGLANG_KT_VERSION"))} + env.update( + PATH=str(tools) + os.pathsep + env["PATH"], + INSTALL_EVENTS=str(tmp_path / "events.jsonl"), + INSTALL_REAL_PYTHON=sys.executable, + ONEAPI_SETVARS=str(tmp_path / "absent-setvars"), + ) + return root, env + + +def run(installer, *args, extra=None): + root, env = installer + result = subprocess.run( + ["bash", str(root / "install.sh"), *args], env=env | (extra or {}), text=True, capture_output=True + ) + events = Path(env["INSTALL_EVENTS"]) + return result, [json.loads(line) for line in events.read_text().splitlines()] if events.exists() else [] + + +def test_sycl_all_builds_kernel_before_installing_aligned_source(installer): + result, events = run(installer, "--sycl") + assert result.returncode == 0, result.stdout + result.stderr + assert [e["kind"] for e in events] == ["kernel", "kernel", "preflight", "preflight", "pip"] + build = events[1] + assert build["args"] == ["build", "--manual"] + for key, value in { + "CPUINFER_USE_SYCL": "1", + "CPUINFER_USE_CUDA": "0", + "CPUINFER_CPU_INSTRUCT": "AVX2", + "CPUINFER_ENABLE_AMX": "OFF", + "CPUINFER_LLAMA_BUILD_COPY": "ON", + "CPUINFER_PARALLEL": "2", + }.items(): + assert build["env"][key] == value + assert events[-1]["args"] == ["-m", "pip", "install", "./python"] + assert "/build/sglang-sycl." in events[-1]["cwd"] + + +def test_plain_all_preserves_original_install_order(installer): + result, events = run(installer) + assert result.returncode == 0, result.stdout + result.stderr + assert [e["kind"] for e in events] == ["kernel", "pip", "kernel"] + assert events[1]["args"] == ["-m", "pip", "install", "./python[all]"] + assert events[-1]["args"] == ["build"] + assert "CPUINFER_USE_SYCL" not in events[-1]["env"] + + +def test_sycl_kernel_only_forwards_build_arguments_and_overrides(installer): + result, events = run(installer, "kt-kernel", "--sycl", "--no-clean", extra={"CPUINFER_PARALLEL": "3"}) + assert result.returncode == 0, result.stdout + result.stderr + assert [e["kind"] for e in events] == ["kernel", "preflight"] + assert events[0]["args"] == ["build", "--no-clean", "--manual"] + assert events[0]["env"]["CPUINFER_PARALLEL"] == "3" + + +def test_sycl_conflicting_gpu_backend_fails_before_install(installer): + result, events = run(installer, "--sycl", extra={"CPUINFER_USE_CUDA": "1"}) + assert result.returncode != 0 and "conflicts with CPUINFER_USE_CUDA" in result.stderr + assert not events + + +def test_sycl_sglang_only_uses_source_copy_and_editable_target(installer): + result, events = run(installer, "sglang", "--sycl", "--editable") + assert result.returncode == 0, result.stdout + result.stderr + assert [e["kind"] for e in events] == ["preflight", "pip"] + assert events[-1]["args"] == ["-m", "pip", "install", "-e", "./python"] + + +def test_help_does_not_initialize_toolchain_or_install(installer): + result, events = run(installer, "--sycl", "--help") + assert "--sycl" in result.stdout and not events + + +def test_sycl_post_build_check_does_not_inherit_runtime_library_path(installer): + result, events = run(installer, "kt-kernel", "--sycl", extra={"LD_LIBRARY_PATH": "/unused/oneapi-runtime"}) + assert result.returncode == 0, result.stdout + result.stderr + assert [e["kind"] for e in events] == ["kernel", "preflight"] + assert events[-1]["ld_library_path"] is None + assert events[-1]["ld_preload"] is None + + +@pytest.mark.parametrize("installed_version", ["0.7.2", "0.7.0.post4"]) +def test_sglang_install_checks_kernel_package_version_before_pip(installer, installed_version): + result, events = run(installer, "sglang", "--sycl") + assert result.returncode == 0, result.stdout + result.stderr + assert events[0]["kind"] == "preflight" + extension = SimpleNamespace( + moe=SimpleNamespace(SYCLXMXGPTQInt4_MOE=SimpleNamespace(implementation_revision=lambda: 7)) + ) + package = SimpleNamespace(kt_kernel_ext=extension) + with ( + patch.dict(sys.modules, {"kt_kernel": package}), + patch("importlib.metadata.version", return_value=installed_version), + patch.dict(os.environ, {"SGLANG_KT_VERSION": "0.7.2"}), + ): + if installed_version == "0.7.2": + exec(events[0]["script"], {}) + else: + with pytest.raises(RuntimeError, match="does not match this checkout"): + exec(events[0]["script"], {}) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/kt-kernel/test/per_commit/test_sycl_runtime_rpath.py b/kt-kernel/test/per_commit/test_sycl_runtime_rpath.py new file mode 100644 index 00000000..3ece9196 --- /dev/null +++ b/kt-kernel/test/per_commit/test_sycl_runtime_rpath.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Load build/install artifacts without LD_LIBRARY_PATH after conda RPATH setup.""" + +import json +import os +from pathlib import Path +import shutil +import subprocess +import sys + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +from ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=10, suite="default") +HELPER = Path(__file__).resolve().parents[2] / "cmake/SYCLRuntime.cmake" + + +@pytest.mark.skipif(sys.platform != "linux", reason="ELF runtime search paths") +@pytest.mark.parametrize("install_rpath_during_build", [False, True]) +def test_runtime_directory_survives_conda_rpath_and_install(tmp_path, install_rpath_during_build): + cc = shutil.which("cc") + cmake = shutil.which("cmake") + if not cc or not cmake: + pytest.skip("requires a C compiler and CMake") + runtime = tmp_path / "oneapi runtime with spaces" + runtime.mkdir() + implementation = tmp_path / "runtime.c" + implementation.write_text("int kt_test_runtime(void) { return 42; }\n") + subprocess.run( + [cc, "-shared", "-fPIC", str(implementation), "-Wl,-soname,libiomp5.so", "-o", str(runtime / "libiomp5.so")], + check=True, + ) + shutil.copy2(runtime / "libiomp5.so", runtime / "libsycl.so") + compiler = tmp_path / "compiler" + compiler.write_text( + f"#!{sys.executable}\n" + "import pathlib,sys\n" + f"root=pathlib.Path({str(runtime)!r})\n" + "name=sys.argv[1].split('=',1)[1]; candidate=root/name\n" + "print(candidate if candidate.exists() else name)\n" + ) + compiler.chmod(0o755) + source = tmp_path / "source" + source.mkdir() + (source / "wrapper.c").write_text( + "int kt_test_runtime(void);\nint kt_test_result(void) { return kt_test_runtime(); }\n" + ) + conda = tmp_path / "conda lib" + conda.mkdir() + (source / "CMakeLists.txt").write_text( + "cmake_minimum_required(VERSION 3.16)\nproject(test_runtime C)\n" + f"set(CMAKE_CXX_COMPILER {json.dumps(str(compiler))})\n" + f"set(CMAKE_BUILD_WITH_INSTALL_RPATH {'TRUE' if install_rpath_during_build else 'FALSE'})\n" + "add_library(wrapper SHARED wrapper.c)\n" + f"target_link_libraries(wrapper PRIVATE {json.dumps(str(runtime / 'libiomp5.so'))})\n" + f"set_target_properties(wrapper PROPERTIES BUILD_RPATH {json.dumps(str(conda))} INSTALL_RPATH {json.dumps(str(conda))})\n" + f"include({json.dumps(str(HELPER))})\nkt_add_sycl_runtime_rpath(wrapper)\n" + "install(TARGETS wrapper LIBRARY DESTINATION lib)\n" + ) + env = {k: v for k, v in os.environ.items() if k not in ("LD_LIBRARY_PATH", "LD_PRELOAD", "LIBRARY_PATH")} + build = tmp_path / "build" + subprocess.run([cmake, "-S", str(source), "-B", str(build)], env=env, capture_output=True, text=True, check=True) + subprocess.run([cmake, "--build", str(build)], env=env, capture_output=True, text=True, check=True) + installed = tmp_path / "installed" + subprocess.run( + [cmake, "--install", str(build), "--prefix", str(installed)], + env=env, + capture_output=True, + text=True, + check=True, + ) + for artifact in (build / "libwrapper.so", installed / "lib/libwrapper.so"): + subprocess.run( + [ + sys.executable, + "-c", + "import ctypes,sys; assert ctypes.CDLL(sys.argv[1]).kt_test_result() == 42", + str(artifact), + ], + env=env, + capture_output=True, + text=True, + check=True, + ) diff --git a/kt-kernel/test/per_commit/test_sycl_sglang_source.py b/kt-kernel/test/per_commit/test_sycl_sglang_source.py new file mode 100644 index 00000000..38b2f408 --- /dev/null +++ b/kt-kernel/test/per_commit/test_sycl_sglang_source.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Pinned SGLang exports preserve runtime source and align the kernel dependency.""" + +import importlib.util +import json +from pathlib import Path +import subprocess +import sys + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +from ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=10, suite="default") + +SCRIPT = Path(__file__).resolve().parents[2] / "scripts/prepare_sglang.py" +spec = importlib.util.spec_from_file_location("prepare_sglang_test", SCRIPT) +module = importlib.util.module_from_spec(spec) +spec.loader.exec_module(module) + + +@pytest.fixture +def source(tmp_path, monkeypatch): + root = tmp_path / "source with spaces" + root.mkdir() + subprocess.run(["git", "init", "-q", str(root)], check=True) + (root / "tracked").write_text("original\n") + (root / "python").mkdir() + (root / "python/pyproject.toml").write_text('[project]\ndependencies = ["kt-kernel==0.1", "torch==2.9.1"]\n') + subprocess.run(["git", "-C", str(root), "add", "tracked", "python/pyproject.toml"], check=True) + subprocess.run( + [ + "git", + "-C", + str(root), + "-c", + "user.name=Test", + "-c", + "user.email=test@example.invalid", + "commit", + "-qm", + "fixture", + ], + check=True, + ) + pin = subprocess.check_output(["git", "-C", str(root), "rev-parse", "HEAD"], text=True).strip() + monkeypatch.setattr(module, "PINNED_SGLANG", pin) + return root + + +def test_exports_pin_preserves_runtime_and_dirty_source(source, tmp_path): + (source / "tracked").write_text("local edit\n") + (source / "untracked").write_text("private local file\n") + output = tmp_path / "build parent" / "prepared" + module.prepare(source, output) + assert (output / "tracked").read_text() == "original\n" + assert not (output / "untracked").exists() + assert (source / "tracked").read_text() == "local edit\n" + assert (source / "untracked").read_text() == "private local file\n" + assert not (output / ".git").exists() + version = module.runpy.run_path(str(module.REPO_ROOT / "version.py"))["__version__"] + project = module.tomllib.loads((output / "python/pyproject.toml").read_text()) + assert project["project"]["dependencies"] == [f"kt-kernel=={version}", "torch==2.9.1"] + assert json.loads((output / ".kt-sglang-source.json").read_text()) == { + "source_commit": module.PINNED_SGLANG, + "kernel_version": version, + } + + +def test_rejects_unknown_pin(source, tmp_path, monkeypatch): + monkeypatch.setattr(module, "PINNED_SGLANG", "0" * 40) + output = tmp_path / "prepared" + with pytest.raises(ValueError, match="Expected pinned SGLang"): + module.prepare(source, output) + assert not output.exists() + + +def test_preserves_existing_destination(source, tmp_path): + output = tmp_path / "prepared" + output.mkdir() + (output / "keep").write_text("keep\n") + with pytest.raises(ValueError, match="already exists"): + module.prepare(source, output) + assert (output / "keep").read_text() == "keep\n" + + +@pytest.mark.parametrize("dependencies", [[], ["kt-kernel==0.1", "kt-kernel==0.2"]]) +def test_invalid_kernel_dependency_does_not_publish_partial_tree(source, tmp_path, dependencies): + project = source / "python/pyproject.toml" + project.write_text("[project]\ndependencies = " + json.dumps(dependencies) + "\n") + subprocess.run(["git", "-C", str(source), "add", "python/pyproject.toml"], check=True) + subprocess.run( + [ + "git", + "-C", + str(source), + "-c", + "user.name=Test", + "-c", + "user.email=test@example.invalid", + "commit", + "--amend", + "--no-edit", + "-q", + ], + check=True, + ) + module.PINNED_SGLANG = subprocess.check_output(["git", "-C", str(source), "rev-parse", "HEAD"], text=True).strip() + output = tmp_path / "prepared" + with pytest.raises(ValueError, match="exactly one pinned kt-kernel"): + module.prepare(source, output) + assert not output.exists() + assert (source / "tracked").read_text() == "original\n" + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/kt-kernel/test/per_commit/test_sycl_xmx_dispatch.py b/kt-kernel/test/per_commit/test_sycl_xmx_dispatch.py new file mode 100644 index 00000000..95cb74c8 --- /dev/null +++ b/kt-kernel/test/per_commit/test_sycl_xmx_dispatch.py @@ -0,0 +1,91 @@ +"""Backend selection tests without importing GPU extensions or torch.""" + +import ast +import os +from pathlib import Path +from unittest.mock import patch + +import pytest + +import sys + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +from ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=10, suite="default") + +SOURCE = Path(__file__).parents[2] / "python/utils/amx.py" + + +def select(legacy=object(), xmx=object()): + tree = ast.parse(SOURCE.read_text()) + fn = next(n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == "_select_sycl_int4_backend") + scope = {"os": os, "SYCLGPTQInt4_MOE": legacy, "SYCLXMXGPTQInt4_MOE": xmx} + exec(compile(ast.Module(body=[fn], type_ignores=[]), str(SOURCE), "exec"), scope) + return scope["_select_sycl_int4_backend"] + + +def test_default_remains_legacy(): + legacy = object() + with patch.dict(os.environ, {}, clear=True): + assert select(legacy=legacy)() is legacy + + +def test_xmx_is_explicit_and_normalized(): + xmx = object() + with patch.dict(os.environ, {"KT_SYCL_INT4_BACKEND": " XMX "}): + assert select(xmx=xmx)() is xmx + + +@pytest.mark.parametrize("value", ["auto", "typo", ""]) +def test_invalid_selection_is_not_a_fallback(value): + with patch.dict(os.environ, {"KT_SYCL_INT4_BACKEND": value}): + with pytest.raises(ValueError): + select()() + + +def test_uncompiled_xmx_is_not_a_fallback(): + with patch.dict(os.environ, {"KT_SYCL_INT4_BACKEND": "xmx"}): + with pytest.raises(RuntimeError, match="not compiled"): + select(xmx=None)() + + +@pytest.mark.parametrize("profile", ["", "auto", "typo", "Panther-Lake"]) +def test_invalid_profile_does_not_silently_select_legacy(profile): + with patch.dict(os.environ, {"KT_SYCL_XMX_PROFILE": profile}, clear=True): + with pytest.raises(ValueError, match="KT_SYCL_XMX_PROFILE"): + select()() + + +def test_profile_selects_xmx_without_setting_other_environment(): + class XMX: + @staticmethod + def implementation_revision(): + return 7 + + with patch.dict(os.environ, {"KT_SYCL_XMX_PROFILE": "panther-lake"}, clear=True): + before = dict(os.environ) + assert select(xmx=XMX)() is XMX + assert dict(os.environ) == before + + +def test_profile_rejects_conflicting_backend(): + with patch.dict(os.environ, {"KT_SYCL_XMX_PROFILE": "panther-lake", "KT_SYCL_INT4_BACKEND": "legacy"}, clear=True): + with pytest.raises(ValueError, match="requires the xmx backend"): + select()() + + +@pytest.mark.parametrize("revision", [0, 6]) +def test_profile_rejects_old_extension(revision): + class XMX: + @staticmethod + def implementation_revision(): + return revision + + with patch.dict(os.environ, {"KT_SYCL_XMX_PROFILE": "panther-lake"}, clear=True): + with pytest.raises(RuntimeError, match="revision >= 7"): + select(xmx=XMX)() + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/kt-kernel/test/sycl/check_pipeline_trace.py b/kt-kernel/test/sycl/check_pipeline_trace.py new file mode 100644 index 00000000..9a990b35 --- /dev/null +++ b/kt-kernel/test/sycl/check_pipeline_trace.py @@ -0,0 +1,77 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +"""Fresh-process device trace regression; source oneAPI before running. + +Opt-in standalone check: ONEAPI_DEVICE_SELECTOR=level_zero:gpu python check_pipeline_trace.py +""" + + +def main(): + import os + + os.environ.update( + KT_SYCL_XMX_GATE_KERNEL="xmx", + KT_SYCL_XMX_DOWN_KERNEL="xmx", + KT_SYCL_XMX_HOST_QUANT="0", + KT_SYCL_XMX_NUMA_DISPATCH="1", + KT_SYCL_PROFILE="1", + KT_SYCL_XMX_TRACE="device", + KT_SYCL_XMX_TRACE_EVERY="0", + KT_SYCL_XMX_MODE="int8", + KT_SYCL_XMX_FUSED_QUANT="0", + KT_SYCL_XMX_GRAPH="0", + KT_SYCL_XMX_DECODE_GATE_KERNEL="inherit", + KT_SYCL_XMX_DECODE_DOWN_KERNEL="xmx", + ) + import torch + from kt_kernel import kt_kernel_ext as ext + + cpu = ext.CPUInfer(2) + mask = torch.zeros(4, dtype=torch.bool) + config = ext.moe.MOEConfig(4, 2, 128, 128, mask.data_ptr()) + config.pool = cpu.backend_ + config.max_len = 1 + config.quant_config.bits = 4 + config.quant_config.group_size = 128 + config.quant_config.zero_point = False + qw = torch.full((4, 16, 128), -2004318072, dtype=torch.int32) + sc = torch.ones(4, 1, 128) + for stage in ("gate", "up", "down"): + setattr(config, stage + "_projs", [[t.data_ptr() for t in qw]]) + setattr(config, stage + "_scales", [[t.data_ptr() for t in sc]]) + x = torch.zeros(1, 128, dtype=torch.bfloat16) + ids = torch.tensor([[0, 1]], dtype=torch.int64) + w = torch.full((1, 2), 0.5) + rows = torch.tensor([1], dtype=torch.int32) + out = torch.empty_like(x) + for mode, commands in (("separate", 5), ("fused", 4)): + os.environ["KT_SYCL_XMX_DECODE_MERGE"] = mode + moe = ext.moe.SYCLXMXGPTQInt4_MOE(config) + cpu.submit(moe.load_weights_task()) + cpu.sync() + cpu.submit( + moe.forward_task(rows.data_ptr(), 2, ids.data_ptr(), w.data_ptr(), x.data_ptr(), out.data_ptr(), False) + ) + cpu.sync() + stats = moe.profile_snapshot() + assert stats["decode.calls"] == 1 and stats["decode.device_commands"] == commands, stats + assert stats["decode.device_span_us"] >= stats["decode.device_kernels_us"] + if mode == "separate": + assert stats["decode.device_merge_us"] > 0 + else: + assert stats["decode.device_merge_us"] == 0 + assert not out.count_nonzero() + moe.reset_profile() + assert moe.profile_snapshot() == {} + os.environ["KT_SYCL_XMX_GRAPH"] = "1" + try: + ext.moe.SYCLXMXGPTQInt4_MOE(config) + except RuntimeError as error: + assert "recorded graph nodes" in str(error) + else: + raise AssertionError("Device trace must reject graph node profiling") + print("Device trace commands, totals, reset and graph guard passed") + + +if __name__ == "__main__": + main() diff --git a/kt-kernel/test/sycl/test_prefill_merge.cpp b/kt-kernel/test/sycl/test_prefill_merge.cpp new file mode 100644 index 00000000..f73b2f3e --- /dev/null +++ b/kt-kernel/test/sycl/test_prefill_merge.cpp @@ -0,0 +1,97 @@ +// SPDX-License-Identifier: Apache-2.0 +// Build: icpx -fsycl -std=c++20 -fp-model=precise -O2 test_prefill_merge.cpp -o test_prefill_merge +// Optional event benchmark: KT_SYCL_PROFILE=1 ./test_prefill_merge 50 +#include +#include +#include +#include +#include + +#include "../../operators/sycl/gptq_int4_xmx-kernels.hpp" +using namespace sycl_xmx; +void check(int rows, int topk, int hidden, int samples) { + const int routed = rows * std::max(topk, 1); + USM slots; + USM weights, down, output; + slots.reserve(routed); + weights.reserve(routed); + down.reserve(static_cast(routed) * hidden); + output.reserve(static_cast(rows) * hidden); + std::mt19937 rng(20261015); + for (int j = 0; j < routed; ++j) { + slots.data[j] = j % 7 == 0 ? -1 : (routed - j - 1) / 2; // Invalid and duplicate expert slots. + weights.data[j] = (int(rng() % 255) - 127) * .001f; + } + for (size_t j = 0; j < static_cast(routed) * hidden; ++j) + down.data[j] = (int(rng() % 65535) - 32767) * .0001f; + auto dep = sycl_int4::queue().single_task([] {}); + auto call = [&](int width) { +#define MERGE(W) merge_vector(slots.data, weights.data, down.data, output.data, rows, topk, hidden, dep) + switch (width) { + case 16: + return MERGE(16); + case 32: + return MERGE(32); + case 64: + return MERGE(64); + case 128: + return MERGE(128); + case 256: + return MERGE(256); + case 512: + return MERGE(512); + default: + return merge(slots.data, weights.data, down.data, output.data, rows, topk, hidden, dep); + } +#undef MERGE + }; + call(0).wait_and_throw(); + std::vector reference(output.data, output.data + rows * hidden); + std::vector widths{0, 16, 32, 64, 128, 256, 512}; + widths.erase(std::remove_if(widths.begin(), widths.end(), [=](int width) { return width > hidden; }), widths.end()); + for (int width : widths) { + call(width).wait_and_throw(); + for (int j = 0; j < rows * hidden; ++j) + if (std::bit_cast(reference[j]) != std::bit_cast(output.data[j])) + throw std::runtime_error("Vector merge bitwise mismatch width=" + std::to_string(width)); + } + std::vector times[7]; + for (int s = -5; s < samples; ++s) { + if (samples == 0) break; + std::shuffle(widths.begin(), widths.end(), rng); + for (int width : widths) { + auto event = call(width); + event.wait_and_throw(); + int index = width == 0 ? 0 + : width == 16 ? 1 + : width == 32 ? 2 + : width == 64 ? 3 + : width == 128 ? 4 + : width == 256 ? 5 + : 6; + if (s >= 0) + times[index].push_back((event.get_profiling_info() - + event.get_profiling_info()) / + 1000.); + } + } + for (int j = 0; j < 7 && samples > 0; ++j) { + auto& t = times[j]; + std::sort(t.begin(), t.end()); + std::cout << "rows=" << rows << " topk=" << topk << " width=" << (j == 0 ? 0 : (8 << j)) + << " median_us=" << t[t.size() / 2] << '\n'; + } + std::cout << "Vector merge exact passed rows=" << rows << " topk=" << topk << '\n'; +} +int main(int argc, char** argv) { + try { + const int samples = argc > 1 ? std::stoi(argv[1]) : 0; + if (samples < 0) throw std::runtime_error("samples must be nonnegative"); + check(1, 0, 128, 0); + check(17, 2, 128, 0); + check(512, 8, 2048, samples); + } catch (const std::exception& error) { + std::cerr << error.what() << '\n'; + return 1; + } +} diff --git a/kt-kernel/test/sycl/test_prefill_n32.cpp b/kt-kernel/test/sycl/test_prefill_n32.cpp new file mode 100644 index 00000000..11616d80 --- /dev/null +++ b/kt-kernel/test/sycl/test_prefill_n32.cpp @@ -0,0 +1,115 @@ +// SPDX-License-Identifier: Apache-2.0 +// Build: icpx -fsycl -fp-model=precise -O2 test_prefill_n32.cpp -o test_prefill_n32 +#include +#include +#include + +#include "../../operators/sycl/gptq_int4_xmx-prefill.hpp" +using namespace sycl_xmx; + +template +void check(int gs) { + constexpr int K = 256, N = 256, R = 33, E = 2; + const int groups = K / gs; + std::mt19937 random(258); + Weight w, u; + w.init(E, N, K, gs); + u.init(E, N, K, gs); + std::vector packed(K / 8 * N); + std::vector scales(groups * N); + for (int e = 0; e < E; ++e) { + for (auto* weight : {&w, &u}) { + for (auto& value : packed) value = random(); + for (auto& value : scales) value = .001f + (random() % 64) * .0001f; + weight->from_mat(e, packed.data(), scales.data()); + } + } + USM input, qref, qtest; + USM sa, oref, otest, sref, stest; + USM aref, atest; + USM tokens; + USM old_tiles, new_tiles; + input.reserve(R * E * K); + sa.reserve(R * E * groups); + tokens.reserve(R * E); + oref.reserve(R * E * N); + otest.reserve(R * E * N); + aref.reserve(R * E * N); + atest.reserve(R * E * N); + qref.reserve(R * E * N); + qtest.reserve(R * E * N); + sref.reserve(R * E * (N / 128)); + stest.reserve(R * E * (N / 128)); + for (int r = 0; r < R * E; ++r) { + tokens.data[r] = R * E - r - 1; + for (int k = 0; k < K; ++k) input.data[r * K + k] = int(random() % 255) - 127; + for (int g = 0; g < groups; ++g) sa.data[r * groups + g] = .0001f * (1 + random() % 32); + } + old_tiles.reserve(E * ((R + 7) / 8)); + new_tiles.reserve(E * ((R + Rows - 1) / Rows)); + int a = 0, b = 0; + for (int e = 0; e < E; ++e) { + for (int r = 0; r < R; r += 8) old_tiles.data[a++] = {e, e * R + r, std::min(8, R - r)}; + for (int r = 0; r < R; r += Rows) new_tiles.data[b++] = {e, e * R + r, std::min(Rows, R - r)}; + } + auto first = projection<8, true, Gate, Fused>(old_tiles.data, a, tokens.data, nullptr, input.data, sa.data, w, &u, + aref.data, oref.data, .5f, 1.702f, nullptr, qref.data, sref.data); + sycl_int4::wait_and_throw(first); + auto second = + prefill_n32(new_tiles.data, b, tokens.data, input.data, sa.data, w, &u, atest.data, + otest.data, .5f, 1.702f, nullptr, qtest.data, stest.data); + sycl_int4::wait_and_throw(second); + for (int j = 0; j < R * E * N; ++j) { + bool equal; + if constexpr (Fused) + equal = qref.data[j] == qtest.data[j]; + else if constexpr (Gate) + equal = aref.data[j] == atest.data[j]; + else + equal = oref.data[j] == otest.data[j]; + if (!equal) { + std::cerr << "M=" << Rows << " gate=" << Gate << " fused=" << Fused << " group=" << gs << " index=" << j; + if constexpr (Fused) + std::cerr << " values=" << int(qref.data[j]) << '/' << int(qtest.data[j]); + else if constexpr (Gate) + std::cerr << " values=" << sycl_int4::bf16_to_fp32(aref.data[j]) << '/' + << sycl_int4::bf16_to_fp32(atest.data[j]); + else + std::cerr << " values=" << oref.data[j] << '/' << otest.data[j]; + std::cerr << '\n'; + throw std::runtime_error("N32 projection/activation/quantization mismatch"); + } + } + if constexpr (Fused) + for (int j = 0; j < R * E * N / 128; ++j) + if (sref.data[j] != stest.data[j]) throw std::runtime_error("N32 quantization scale mismatch"); +} + +template +void suite() { + for (int group : {32, 128}) { + check(group); + check(group); + } + check(128); + for (int group : {32, 128}) { + check(group); + check(group); + } + check(128); + std::cout << "N32 M=" << Rows + << " gate/down, distinct experts/groups/scales, reverse row mapping, tail33, block/gather, fused quant " + "exact passed\n"; +} +int main() { + try { + suite<1>(); + suite<2>(); + suite<4>(); + suite<8>(); + suite<16>(); + } catch (const std::exception& e) { + std::cerr << e.what() << '\n'; + return 1; + } +} diff --git a/kt-kernel/test/sycl/test_vector_projection.cpp b/kt-kernel/test/sycl/test_vector_projection.cpp new file mode 100644 index 00000000..463db9cc --- /dev/null +++ b/kt-kernel/test/sycl/test_vector_projection.cpp @@ -0,0 +1,137 @@ +// SPDX-License-Identifier: Apache-2.0 +// Build: icpx -fsycl -O2 test_vector_projection.cpp -o test_vector_projection +#include +#include +#include +#include + +#include "../../operators/sycl/gptq_int4_vector-kernels.hpp" + +using namespace sycl_xmx; + +void check_dp4a() { + USM w; + USM a, output; + w.reserve(16); + a.reserve(16); + output.reserve(16); + std::mt19937 random(20261014); + for (int j = 0; j < 16; ++j) { + w.data[j] = random(); + a.data[j] = random(); + } + auto* wp = w.data; + auto* ap = a.data; + auto* op = output.data; + auto event = sycl_int4::queue().single_task([=]() SYCL_ESIMD_KERNEL { + es::simd packed; + packed.copy_from(wp); + es::simd activations; + activations.copy_from(ap); + es::simd zero = 0; + es::dp4a(zero, unpack_s4_bytes(packed), activations).copy_to(op); + }); + sycl_int4::wait_and_throw(event); + for (int j = 0; j < 16; ++j) { + int expected = 0; + for (int r = 0; r < 4; ++r) { + int weight = static_cast((w.data[j] >> (4 * r)) & 15) - 8; + int activation = static_cast((static_cast(a.data[j]) >> (8 * r)) & 255); + expected += weight * activation; + } + if (output.data[j] != expected) { + std::cerr << "lane=" << j << " w=" << w.data[j] << " a=" << a.data[j] << " expected=" << expected + << " actual=" << output.data[j] << '\n'; + throw std::runtime_error("Signed nibble/DP4A mismatch"); + } + } +} + +template +void check_projection(int group_size, int splits) { + constexpr int K = 256, N = 32, R = 3; + std::mt19937 random(123); + const int groups = K / group_size; + std::vector qw(K / 8 * N), up(K / 8 * N); + std::vector scales(groups * N), up_scales(groups * N); + for (auto& x : qw) x = random(); + for (auto& x : up) x = random(); + for (int g = 0; g < groups; ++g) + for (int n = 0; n < N; ++n) { + scales[g * N + n] = (g + 1) * 0.003f; + up_scales[g * N + n] = (g + 2) * 0.001f; + } + Weight w, uw; + w.init(1, N, K, group_size); + uw.init(1, N, K, group_size); + w.from_mat(0, qw.data(), scales.data()); + uw.from_mat(0, up.data(), up_scales.data()); + USM input, activated; + USM qi; + USM sa, output; + USM tokens; + USM tiles; + input.reserve(R * K); + activated.reserve(R * N); + qi.reserve(R * K); + sa.reserve(R * groups); + output.reserve(R * N); + tiles.reserve(1); + tokens.reserve(R); + tiles.data[0] = {0, 0, R}; + for (int r = 0; r < R; ++r) { + tokens.data[r] = R - r - 1; + for (int k = 0; k < K; ++k) { + float value = (static_cast(random() % 255) - 127) * 0.125f; + input.data[r * K + k] = sycl_int4::fp32_to_bf16(value); + qi.data[r * K + k] = static_cast(random() % 255 - 127); + } + for (int g = 0; g < groups; ++g) sa.data[r * groups + g] = (g + 1) * 0.03f; + } + auto event = vector_dispatch(splits, tiles.data, 1, 8, tokens.data, input.data, qi.data, sa.data, w, &uw, + activated.data, output.data, 0.5f, 1.7f); + sycl_int4::wait_and_throw(event); + for (int r = 0; r < R; ++r) + for (int n = 0; n < N; ++n) { + const int row = GateUp ? tokens.data[r] : r; + double gate = 0, upper = 0; + for (int g = 0; g < groups; ++g) { + double pg = 0, pu = 0; + for (int j = 0; j < group_size; ++j) { + int k = g * group_size + j; + double av = Q8 ? qi.data[row * K + k] : sycl_int4::bf16_to_fp32(input.data[row * K + k]); + pg += av * (static_cast((qw[(k / 8) * N + n] >> (4 * (k % 8))) & 15) - 8); + pu += av * (static_cast((up[(k / 8) * N + n] >> (4 * (k % 8))) & 15) - 8); + } + double activation_scale = Q8 ? sa.data[row * groups + g] : 1.0; + gate += pg * scales[g * N + n] * activation_scale; + upper += pu * up_scales[g * N + n] * activation_scale; + } + if constexpr (GateUp) { + gate = std::clamp(gate, -0.5, 0.5); + upper = std::clamp(upper, -0.5, 0.5); + gate = gate / (1.0 + std::exp(-gate * 1.7)) * (upper + 1.0); + } + float expected = sycl_int4::bf16_to_fp32(sycl_int4::fp32_to_bf16(static_cast(gate))); + float actual = GateUp ? sycl_int4::bf16_to_fp32(activated.data[r * N + n]) : output.data[r * N + n]; + if (std::fabs(expected - actual) > std::max(0.004f * std::fabs(expected), 1e-5f)) + throw std::runtime_error("Vector projection/reference mismatch"); + } +} + +int main() { + try { + check_dp4a(); + for (int group : {32, 128}) + for (int splits : {1, 2, 4, 8}) { + check_projection(group, splits); + check_projection(group, splits); + check_projection(group, splits); + check_projection(group, splits); + std::cout << "vector/DP4A gate/down group=" << group << " splits=" << splits << " passed\n"; + } + } catch (const std::exception& e) { + std::cerr << e.what() << '\n'; + return 1; + } +} diff --git a/kt-kernel/test/sycl/test_xmx_dpas.cpp b/kt-kernel/test/sycl/test_xmx_dpas.cpp new file mode 100644 index 00000000..75e81621 --- /dev/null +++ b/kt-kernel/test/sycl/test_xmx_dpas.cpp @@ -0,0 +1,130 @@ +// SPDX-License-Identifier: Apache-2.0 +// Build: icpx -fsycl -O2 test_xmx_dpas.cpp -o test_xmx_dpas +#include +#include +#include +#include +#include + +#include "../../operators/sycl/gptq_int4_xmx-kernels.hpp" + +namespace es = sycl::ext::intel::esimd; +namespace xm = es::xmx; + +template +void check_signed_dpas() { + auto& q = sycl_int4::queue(); + sycl_xmx::USM a; + sycl_xmx::USM b; + sycl_xmx::USM out; + a.reserve(M * 32); + b.reserve(64); + out.reserve(M * 16); + int8_t weights[32 * 16]; + std::mt19937 gen(20261014 + M); + for (int i = 0; i < M * 32; ++i) a.data[i] = static_cast(gen() % 255) - 127; + std::fill_n(b.data, 64, 0u); + for (int k = 0; k < 32; ++k) + for (int n = 0; n < 16; ++n) { + int v = static_cast(gen() % 16) - 8; + weights[k * 16 + n] = v; + // Convert the GPTQ offset-binary nibble into a two's-complement S4 nibble. + b.data[(k / 8) * 16 + n] |= static_cast((v + 8) ^ 8) << ((k % 8) * 4); + } + auto* ap = a.data; + auto* bp = b.data; + auto* op = out.data; + q.single_task([=]() SYCL_ESIMD_KERNEL { + es::simd av; + av.copy_from(ap); + es::simd bv; + bv.copy_from(bp); + es::simd c = 0; + auto value = + xm::dpas<8, M, int, int, uint32_t, int8_t, xm::dpas_argument_type::s4, xm::dpas_argument_type::s8>(c, bv, av); + value.copy_to(op); + }).wait_and_throw(); + for (int m = 0; m < M; ++m) + for (int n = 0; n < 16; ++n) { + int expected = 0; + for (int k = 0; k < 32; ++k) expected += a.data[m * 32 + k] * weights[k * 16 + n]; + if (out.data[m * 16 + n] != expected) throw std::runtime_error("Signed S8xS4 DPAS mismatch"); + } + std::cout << "signed DPAS M=" << M << " passed\n"; +} + +template +void check_projection() { + constexpr int K = 128, N = 32, GS = 32, R = 3; + std::mt19937 gen(123); + std::vector qw(K / 8 * N, 0); + std::vector scales(K / GS * N); + for (auto& word : qw) word = gen(); + for (int g = 0; g < K / GS; ++g) + for (int n = 0; n < N; ++n) scales[g * N + n] = (g + 1) * 0.007f; + sycl_xmx::Weight weight; + weight.init(1, N, K, GS); + weight.from_mat(0, qw.data(), scales.data()); + sycl_xmx::USM input; + sycl_xmx::USM qi; + sycl_xmx::USM sa, output; + sycl_xmx::USM tile; + input.reserve(R * K); + qi.reserve(R * K); + sa.reserve(R * K / GS); + output.reserve(R * N); + tile.reserve(1); + tile.data[0] = {0, 0, R}; + for (int r = 0; r < R; ++r) + for (int k = 0; k < K; ++k) { + float v = static_cast(gen() % 15) - 7; + input.data[r * K + k] = sycl_int4::fp32_to_bf16(v); + qi.data[r * K + k] = static_cast(v); + } + for (int r = 0; r < R; ++r) + for (int g = 0; g < K / GS; ++g) sa.data[r * K / GS + g] = (g + 1) * 0.2f; + auto event = sycl_xmx::projection<8, Q8, false, false, Splits>(tile.data, 1, nullptr, input.data, qi.data, sa.data, + weight, nullptr, nullptr, output.data, 0, 0); + sycl_int4::wait_and_throw(event); + for (int r = 0; r < R; ++r) + for (int n = 0; n < N; ++n) { + float expected = 0; + for (int g = 0; g < K / GS; ++g) { + int partial = 0; + for (int j = 0; j < GS; ++j) { + int k = g * GS + j; + int w = static_cast((qw[(k / 8) * N + n] >> (4 * (k % 8))) & 15) - 8; + if constexpr (Q8) + partial += qi.data[r * K + k] * w; + else + expected += + sycl_int4::bf16_to_fp32(input.data[r * K + k]) * static_cast(sycl::half(w * scales[g * N + n])); + } + if constexpr (Q8) expected += partial * scales[g * N + n] * sa.data[r * K / GS + g]; + } + expected = sycl_int4::bf16_to_fp32(sycl_int4::fp32_to_bf16(expected)); + if (std::fabs(output.data[r * N + n] - expected) > 0.02f) + throw std::runtime_error("Projection/tail/group-scale mismatch"); + } + std::cout << (Q8 ? "W4A8" : "W4A16") << " projection, partial tile and distinct group scales passed\n"; +} + +int main() { + try { + check_signed_dpas<1>(); + check_signed_dpas<2>(); + check_signed_dpas<4>(); + check_signed_dpas<8>(); + check_projection(); + check_projection(); + check_projection(); + check_projection(); + check_projection(); + check_projection(); + check_projection(); + check_projection(); + } catch (const std::exception& e) { + std::cerr << e.what() << '\n'; + return 1; + } +} diff --git a/kt-kernel/test/test_sycl_xmx_hardware.py b/kt-kernel/test/test_sycl_xmx_hardware.py new file mode 100644 index 00000000..929a06a2 --- /dev/null +++ b/kt-kernel/test/test_sycl_xmx_hardware.py @@ -0,0 +1,760 @@ +"""Opt-in integration checks: source oneAPI and set KT_TEST_SYCL_XMX=1.""" + +import os + +import pytest + +pytestmark = pytest.mark.skipif( + os.getenv("KT_TEST_SYCL_XMX") != "1", reason="requires Intel XMX GPU and SYCL extension" +) + + +@pytest.fixture +def setup(monkeypatch): + import torch + from kt_kernel import kt_kernel_ext as ext + + monkeypatch.setenv("KT_SYCL_XMX_MODE", "int8") + cpu = ext.CPUInfer(2) + mask = torch.zeros(4, dtype=torch.bool) + config = ext.moe.MOEConfig(4, 2, 128, 128, mask.data_ptr()) + config.pool = cpu.backend_ + config.max_len = 17 + config.quant_config.bits = 4 + config.quant_config.group_size = 128 + config.quant_config.zero_point = False + qw = torch.full((4, 16, 128), -2004318072, dtype=torch.int32) # all GPTQ nibbles are 8, hence zero + scales = torch.ones(4, 1, 128) + for name in ("gate", "up", "down"): + setattr(config, name + "_projs", [[t.data_ptr() for t in qw]]) + setattr(config, name + "_scales", [[t.data_ptr() for t in scales]]) + return ext, torch, cpu, config, mask, qw, scales + + +@pytest.mark.parametrize("mode", ["fp16", "int8"]) +@pytest.mark.parametrize("contiguous", [False, True]) +def test_zero_weights_and_incremental_output(setup, monkeypatch, mode, contiguous): + ext, torch, cpu, config, mask, qw, scales = setup + monkeypatch.setenv("KT_SYCL_XMX_MODE", mode) + if contiguous: + for name in ("gate", "up", "down"): + setattr(config, name + "_projs", []) + setattr(config, name + "_scales", []) + setattr(config, name + "_proj", qw.data_ptr()) + setattr(config, name + "_scale", scales.data_ptr()) + moe = ext.moe.SYCLXMXGPTQInt4_MOE(config) + cpu.submit(moe.load_weights_task()) + cpu.sync() + for rows in (1, 17): + qlen = torch.tensor([rows], dtype=torch.int32) + x = torch.zeros(rows, 128, dtype=torch.bfloat16) + ids = torch.tensor([[0, 1]] * rows, dtype=torch.int64) + weights = torch.full((rows, 2), 0.5) + out = torch.ones_like(x) + cpu.submit( + moe.forward_task( + qlen.data_ptr(), 2, ids.data_ptr(), weights.data_ptr(), x.data_ptr(), out.data_ptr(), False + ) + ) + cpu.sync() + assert not out.count_nonzero() + out.fill_(1) + cpu.submit( + moe.forward_task(qlen.data_ptr(), 2, ids.data_ptr(), weights.data_ptr(), x.data_ptr(), out.data_ptr(), True) + ) + cpu.sync() + assert bool((out == 1).all()) + + +def test_invalid_quantization_is_rejected(setup): + ext, _, _, config, *_ = setup + config.quant_config.zero_point = True + with pytest.raises(RuntimeError, match="symmetric GPTQ INT4"): + ext.moe.SYCLXMXGPTQInt4_MOE(config) + + +def test_invalid_mode_is_rejected(setup, monkeypatch): + ext, _, _, config, *_ = setup + monkeypatch.setenv("KT_SYCL_XMX_MODE", "typo") + with pytest.raises(RuntimeError, match="KT_SYCL_XMX_MODE"): + ext.moe.SYCLXMXGPTQInt4_MOE(config) + + +def test_invalid_mapping_is_rejected(setup): + ext, torch, _, config, *_ = setup + mapping = torch.tensor([0, 1, 2, 1 << 40], dtype=torch.int64) + config.physical_to_logical_map = mapping.data_ptr() + moe = ext.moe.SYCLXMXGPTQInt4_MOE(config) + with pytest.raises(RuntimeError, match="physical_to_logical_map"): + moe.load_weights() + + +@pytest.mark.parametrize("numa", ["0", "1"]) +def test_workspace_bounds_are_rejected(setup, monkeypatch, numa): + ext, torch, cpu, config, *_ = setup + monkeypatch.setenv("KT_SYCL_XMX_NUMA_DISPATCH", numa) + moe = ext.moe.SYCLXMXGPTQInt4_MOE(config) + moe.load_weights() + x = torch.zeros(18, 128, dtype=torch.bfloat16) + ids = torch.zeros(18, 2, dtype=torch.int64) + weights = torch.ones(18, 2) + qlen = torch.tensor([18], dtype=torch.int32) + with pytest.raises(RuntimeError, match="workspace"): + moe.forward(qlen.data_ptr(), 2, ids.data_ptr(), weights.data_ptr(), x.data_ptr(), x.data_ptr(), False) + cpu.submit( + moe.forward_task(qlen.data_ptr(), 2, ids.data_ptr(), weights.data_ptr(), x.data_ptr(), x.data_ptr(), False) + ) + with pytest.raises(RuntimeError, match="workspace"): + cpu.sync() + # The queue remains usable after reporting an invalid request. + qlen[0] = 1 + cpu.submit( + moe.forward_task(qlen.data_ptr(), 2, ids.data_ptr(), weights.data_ptr(), x.data_ptr(), x.data_ptr(), False) + ) + cpu.sync() + + +@pytest.mark.parametrize( + "name,value", + [ + ("KT_SYCL_XMX_ROWS", "16"), + ("KT_SYCL_XMX_GRAPH", "yes"), + ("KT_SYCL_XMX_HOST_QUANT", "auto"), + ("KT_SYCL_XMX_FUSED_QUANT", "yes"), + ("KT_SYCL_XMX_NUMA_DISPATCH", "auto"), + ], +) +def test_invalid_tuning_is_rejected(setup, monkeypatch, name, value): + ext, _, _, config, *_ = setup + monkeypatch.setenv(name, value) + with pytest.raises(RuntimeError, match=name): + ext.moe.SYCLXMXGPTQInt4_MOE(config) + + +@pytest.mark.parametrize("fused", ["0", "1", "prefill"]) +@pytest.mark.parametrize("decode_gate", ["xmx", "vector", "dp4a", "xmx-split"]) +def test_graph_reads_changing_inputs_and_routes(setup, monkeypatch, fused, decode_gate): + ext, torch, cpu, config, mask, qw, scales = setup + generator = torch.Generator().manual_seed(20261014) + qw.copy_(torch.randint(-(1 << 31), (1 << 31) - 1, qw.shape, dtype=torch.int32, generator=generator)) + scales.fill_(0.003) + monkeypatch.setenv("KT_SYCL_XMX_FUSED_QUANT", fused) + monkeypatch.setenv("KT_SYCL_XMX_DECODE_GATE_KERNEL", decode_gate) + if fused == "1" and decode_gate != "xmx": + pytest.skip("Full quantization fusion requires XMX gate") + monkeypatch.setenv("KT_SYCL_XMX_GRAPH", "0") + direct = ext.moe.SYCLXMXGPTQInt4_MOE(config) + monkeypatch.setenv("KT_SYCL_XMX_GRAPH", "1") + graph = ext.moe.SYCLXMXGPTQInt4_MOE(config) + for moe in (direct, graph): + cpu.submit(moe.load_weights_task()) + cpu.sync() + qlen = torch.tensor([1], dtype=torch.int32) + for iteration in range(12): + topk = iteration % 2 + 1 + x = torch.randn(1, 128, generator=generator).bfloat16() + ids = torch.randperm(4, generator=generator)[:topk].reshape(1, topk) + if iteration % 3 == 0: + ids[0, 0] = -1 + weights = torch.softmax(torch.randn(1, topk, generator=generator), 1) + outputs = [] + for moe in (direct, graph): + out = torch.empty_like(x) + cpu.submit( + moe.forward_task( + qlen.data_ptr(), topk, ids.data_ptr(), weights.data_ptr(), x.data_ptr(), out.data_ptr(), False + ) + ) + cpu.sync() + outputs.append(out) + torch.testing.assert_close(*outputs, rtol=0, atol=0) + + +@pytest.mark.parametrize("rows", [1, 17]) +@pytest.mark.parametrize("fused_mode", ["1", "prefill"]) +def test_fused_quantization_matches_unfused(setup, monkeypatch, rows, fused_mode): + ext, torch, cpu, config, mask, qw, scales = setup + generator = torch.Generator().manual_seed(42) + qw.copy_(torch.randint(-(1 << 31), (1 << 31) - 1, qw.shape, dtype=torch.int32, generator=generator)) + scales.fill_(0.003) + monkeypatch.setenv("KT_SYCL_XMX_FUSED_QUANT", "0") + direct = ext.moe.SYCLXMXGPTQInt4_MOE(config) + monkeypatch.setenv("KT_SYCL_XMX_FUSED_QUANT", fused_mode) + fused = ext.moe.SYCLXMXGPTQInt4_MOE(config) + qlen = torch.tensor([rows], dtype=torch.int32) + x = torch.randn(rows, 128, generator=generator).bfloat16() + ids = torch.tensor([[0, 1]] * rows, dtype=torch.int64) + weights = torch.full((rows, 2), 0.5) + outputs = [] + for moe in (direct, fused): + cpu.submit(moe.load_weights_task()) + cpu.sync() + out = torch.empty_like(x) + cpu.submit( + moe.forward_task( + qlen.data_ptr(), 2, ids.data_ptr(), weights.data_ptr(), x.data_ptr(), out.data_ptr(), False + ) + ) + cpu.sync() + outputs.append(out) + torch.testing.assert_close(*outputs, rtol=0, atol=0) + + +@pytest.mark.parametrize( + "gate,down", + [ + ("vector", "vector"), + ("dp4a", "dp4a"), + ("xmx", "dp4a"), + ("dp4a", "xmx"), + ("vector", "xmx"), + ("xmx", "vector"), + ("xmx-split", "xmx"), + ("xmx", "xmx-split"), + ], +) +@pytest.mark.parametrize("rows", [1, 7]) +def test_projection_candidates_against_reference(setup, monkeypatch, gate, down, rows): + ext, torch, cpu, config, mask, qw, scales = setup + generator = torch.Generator().manual_seed(73) + qw.copy_(torch.randint(-(1 << 31), (1 << 31) - 1, qw.shape, dtype=torch.int32, generator=generator)) + scales.fill_(0.003) + monkeypatch.setenv("KT_SYCL_XMX_MODE", "int8") + monkeypatch.setenv("KT_SYCL_XMX_FUSED_QUANT", "0") + monkeypatch.setenv("KT_SYCL_XMX_GATE_KERNEL", gate) + monkeypatch.setenv("KT_SYCL_XMX_DOWN_KERNEL", down) + monkeypatch.delenv("KT_SYCL_XMX_DECODE_GATE_KERNEL", raising=False) + monkeypatch.delenv("KT_SYCL_XMX_DECODE_DOWN_KERNEL", raising=False) + moe = ext.moe.SYCLXMXGPTQInt4_MOE(config) + cpu.submit(moe.load_weights_task()) + cpu.sync() + x = torch.randn(rows, 128, generator=generator).bfloat16() + ids = torch.tensor([[0, 1]] * rows, dtype=torch.int64) + weights = torch.full((rows, 2), 0.5) + out = torch.empty_like(x) + qlen = torch.tensor([rows], dtype=torch.int32) + cpu.submit( + moe.forward_task(qlen.data_ptr(), 2, ids.data_ptr(), weights.data_ptr(), x.data_ptr(), out.data_ptr(), False) + ) + cpu.sync() + shift = torch.arange(8) * 4 + deq = (((qw.long()[..., None] >> shift) & 15) - 8).permute(0, 1, 3, 2).reshape(4, 128, 128).float() * 0.003 + expected = torch.zeros(rows, 128) + for e in (0, 1): + projection = x.float() @ deq[e] + middle = (torch.nn.functional.silu(projection) * projection).bfloat16().float() + expected += (middle @ deq[e]).bfloat16().float() * 0.5 + expected = expected.bfloat16().float() + rmse = (out.float() - expected).square().mean().sqrt() / expected.square().mean().sqrt().clamp_min(1e-12) + assert bool(torch.isfinite(out).all()) + assert rmse < 0.04 + + +@pytest.mark.parametrize( + "name,value", + [ + ("KT_SYCL_XMX_GATE_KERNEL", "typo"), + ("KT_SYCL_XMX_DOWN_KERNEL", "auto"), + ("KT_SYCL_XMX_DECODE_GATE_KERNEL", "bad"), + ("KT_SYCL_XMX_VECTOR_SPLITS", "3"), + ("KT_SYCL_XMX_GATE_SPLITS", "0"), + ("KT_SYCL_XMX_DOWN_SPLITS", "16"), + ], +) +def test_invalid_projection_selection_is_rejected(setup, monkeypatch, name, value): + ext, _, _, config, *_ = setup + monkeypatch.setenv(name, value) + with pytest.raises(RuntimeError, match=name): + ext.moe.SYCLXMXGPTQInt4_MOE(config) + + +def test_prefill_fusion_allows_vector_decode_but_rejects_full_fusion(setup, monkeypatch): + ext, _, _, config, *_ = setup + monkeypatch.setenv("KT_SYCL_XMX_DECODE_GATE_KERNEL", "vector") + monkeypatch.setenv("KT_SYCL_XMX_FUSED_QUANT", "prefill") + ext.moe.SYCLXMXGPTQInt4_MOE(config) + monkeypatch.setenv("KT_SYCL_XMX_FUSED_QUANT", "1") + with pytest.raises(RuntimeError, match="fused quantization"): + ext.moe.SYCLXMXGPTQInt4_MOE(config) + + +@pytest.mark.parametrize( + "decode_gate,decode_down", + [ + ("vector", "xmx"), + ("xmx", "vector"), + ("dp4a", "xmx"), + ("xmx", "dp4a"), + ("xmx-split", "xmx"), + ("xmx", "xmx-split"), + ], +) +@pytest.mark.parametrize("host_quant", ["0", "1"]) +def test_decode_overrides_preserve_prefill_and_capture_settings( + setup, monkeypatch, decode_gate, decode_down, host_quant +): + """Decode-only choices must not leak into fused prefill or later executors.""" + ext, torch, cpu, config, mask, qw, scales = setup + generator = torch.Generator().manual_seed(20261014) + qw.copy_(torch.randint(-(1 << 31), (1 << 31) - 1, qw.shape, dtype=torch.int32, generator=generator)) + scales.fill_(0.003) + monkeypatch.setenv("KT_SYCL_XMX_GRAPH", "0") + monkeypatch.setenv("KT_SYCL_XMX_HOST_QUANT", host_quant) + monkeypatch.setenv("KT_SYCL_XMX_GATE_SPLITS", "2") + monkeypatch.setenv("KT_SYCL_XMX_DOWN_SPLITS", "2") + monkeypatch.setenv("KT_SYCL_XMX_GATE_KERNEL", "xmx") + monkeypatch.setenv("KT_SYCL_XMX_DOWN_KERNEL", "xmx") + monkeypatch.setenv("KT_SYCL_XMX_DECODE_GATE_KERNEL", decode_gate) + monkeypatch.setenv("KT_SYCL_XMX_DECODE_DOWN_KERNEL", decode_down) + monkeypatch.setenv("KT_SYCL_XMX_FUSED_QUANT", "prefill") + selected = ext.moe.SYCLXMXGPTQInt4_MOE(config) + cpu.submit(selected.load_weights_task()) + cpu.sync() + + monkeypatch.setenv("KT_SYCL_XMX_DECODE_GATE_KERNEL", "inherit") + monkeypatch.setenv("KT_SYCL_XMX_DECODE_DOWN_KERNEL", "inherit") + for rows in (1, 17): + # The explicit reference matches this stage. It also changes the + # environment after selected captured its own immutable choices. + monkeypatch.setenv("KT_SYCL_XMX_GATE_KERNEL", decode_gate if rows == 1 else "xmx") + monkeypatch.setenv("KT_SYCL_XMX_DOWN_KERNEL", decode_down if rows == 1 else "xmx") + monkeypatch.setenv("KT_SYCL_XMX_FUSED_QUANT", "0" if rows == 1 else "prefill") + explicit = ext.moe.SYCLXMXGPTQInt4_MOE(config) + cpu.submit(explicit.load_weights_task()) + cpu.sync() + x = torch.randn(rows, 128, generator=generator).bfloat16() + ids = torch.randint(0, 4, (rows, 2), generator=generator) + weights = torch.softmax(torch.randn(rows, 2, generator=generator), 1) + qlen = torch.tensor([rows], dtype=torch.int32) + outputs = [] + for moe in (selected, explicit): + out = torch.empty_like(x) + cpu.submit( + moe.forward_task( + qlen.data_ptr(), 2, ids.data_ptr(), weights.data_ptr(), x.data_ptr(), out.data_ptr(), False + ) + ) + cpu.sync() + outputs.append(out) + torch.testing.assert_close(*outputs, rtol=0, atol=0) + + +@pytest.mark.parametrize("mode,quant_fusion", [("fp16", "0"), ("int8", "0"), ("int8", "1"), ("int8", "prefill")]) +@pytest.mark.parametrize("graph", ["0", "1"]) +@pytest.mark.parametrize("topk", [1, 2, 8]) +def test_fused_decode_merge_matches_separate_with_dynamic_routes(setup, monkeypatch, mode, quant_fusion, graph, topk): + ext, torch, cpu, config, mask, qw, scales = setup + config.num_experts_per_tok = topk + generator = torch.Generator().manual_seed(108) + qw.copy_(torch.randint(-(1 << 31), (1 << 31) - 1, qw.shape, dtype=torch.int32, generator=generator)) + scales.fill_(0.003) + monkeypatch.setenv("KT_SYCL_XMX_MODE", mode) + monkeypatch.setenv("KT_SYCL_XMX_GRAPH", graph) + monkeypatch.setenv("KT_SYCL_XMX_DECODE_DOWN_KERNEL", "xmx") + monkeypatch.setenv("KT_SYCL_XMX_FUSED_QUANT", quant_fusion) + monkeypatch.setenv("KT_SYCL_XMX_DECODE_MERGE", "separate") + separate = ext.moe.SYCLXMXGPTQInt4_MOE(config) + monkeypatch.setenv("KT_SYCL_XMX_DECODE_MERGE", "fused") + fused = ext.moe.SYCLXMXGPTQInt4_MOE(config) + for moe in (separate, fused): + cpu.submit(moe.load_weights_task()) + cpu.sync() + for ids_list in ([0, 1], [1, 1], [-1, 1], [3, 0], [-1, -1], [99, 2]): + for rows in (1, 17): + x = torch.randn(rows, 128, generator=generator).bfloat16() + expanded = (ids_list * ((topk + 1) // 2))[:topk] + ids = torch.tensor([expanded] * rows, dtype=torch.int64) + weights = torch.softmax(torch.randn(rows, topk, generator=generator), 1) + qlen = torch.tensor([rows], dtype=torch.int32) + outputs = [] + for moe in (separate, fused): + out = torch.randn(rows, 128, generator=generator).bfloat16() + cpu.submit( + moe.forward_task( + qlen.data_ptr(), topk, ids.data_ptr(), weights.data_ptr(), x.data_ptr(), out.data_ptr(), False + ) + ) + cpu.sync() + outputs.append(out.clone()) + previous = out.clone() + cpu.submit( + moe.forward_task( + qlen.data_ptr(), topk, ids.data_ptr(), weights.data_ptr(), x.data_ptr(), out.data_ptr(), True + ) + ) + cpu.sync() + torch.testing.assert_close(out, (previous.float() * 2).bfloat16(), rtol=0, atol=0) + torch.testing.assert_close(*outputs, rtol=0, atol=0) + + +@pytest.mark.parametrize( + "name,value", + [("KT_SYCL_XMX_TRACE", "yes"), ("KT_SYCL_XMX_DECODE_MERGE", "auto"), ("KT_SYCL_XMX_TRACE_EVERY", "-1")], +) +def test_invalid_pipeline_options_are_rejected(setup, monkeypatch, name, value): + ext, _, _, config, *_ = setup + monkeypatch.setenv(name, value) + with pytest.raises(RuntimeError, match=name): + ext.moe.SYCLXMXGPTQInt4_MOE(config) + + +def test_fused_merge_rejects_unsupported_down_kernel(setup, monkeypatch): + ext, _, _, config, *_ = setup + monkeypatch.setenv("KT_SYCL_XMX_DECODE_MERGE", "fused") + monkeypatch.setenv("KT_SYCL_XMX_DECODE_DOWN_KERNEL", "vector") + with pytest.raises(RuntimeError, match="DECODE_MERGE"): + ext.moe.SYCLXMXGPTQInt4_MOE(config) + + +def test_host_profile_snapshot_reset_and_zero_routes(setup, monkeypatch): + ext, torch, cpu, config, mask, qw, scales = setup + monkeypatch.setenv("KT_SYCL_XMX_TRACE", "host") + monkeypatch.setenv("KT_SYCL_XMX_DECODE_MERGE", "separate") + moe = ext.moe.SYCLXMXGPTQInt4_MOE(config) + cpu.submit(moe.load_weights_task()) + cpu.sync() + for row_count, expert in ((1, 0), (17, -1)): + rows = torch.tensor([row_count], dtype=torch.int32) + ids = torch.full((row_count, 2), expert, dtype=torch.int64) + weights = torch.full((row_count, 2), 0.5) + x = torch.ones(row_count, 128, dtype=torch.bfloat16) + out = torch.empty_like(x) + cpu.submit( + moe.forward_task( + rows.data_ptr(), 2, ids.data_ptr(), weights.data_ptr(), x.data_ptr(), out.data_ptr(), False + ) + ) + cpu.sync() + stats = moe.profile_snapshot() + assert stats["decode.calls"] == stats["prefill.calls"] == 1 + assert stats["decode.routed"] == 2 and stats["prefill.routed"] == 0 + assert stats["decode.host_executor_us"] > 0 + assert stats["decode.host_wrapper_us"] >= stats["decode.host_bf16_us"] + assert not any("device_" in key for key in stats) + moe.reset_profile() + assert moe.profile_snapshot() == {} + + +def test_fused_merge_rejects_topk_above_eight(setup, monkeypatch): + ext, _, _, config, *_ = setup + config.num_experts_per_tok = 9 + monkeypatch.setenv("KT_SYCL_XMX_DECODE_MERGE", "fused") + with pytest.raises(RuntimeError, match="topk<=8"): + ext.moe.SYCLXMXGPTQInt4_MOE(config) + + +def test_device_trace_requires_profiled_queue(setup, monkeypatch): + ext, _, _, config, *_ = setup + monkeypatch.setenv("KT_SYCL_XMX_TRACE", "device") + # The queue property is fixed by this process's first creation. + if os.getenv("KT_SYCL_PROFILE") == "1": + pytest.skip("Queue profiling was explicitly enabled for this test process") + with pytest.raises(RuntimeError, match="requires KT_SYCL_PROFILE=1"): + ext.moe.SYCLXMXGPTQInt4_MOE(config) + + +@pytest.mark.parametrize("splits", ["1", "2"]) +@pytest.mark.parametrize("graph", ["0", "1"]) +@pytest.mark.parametrize("host_quant", ["0", "1"]) +def test_fused_expert_decode_matches_staged_and_preserves_prefill(setup, monkeypatch, splits, graph, host_quant): + from pathlib import Path + import importlib.util + + spec = importlib.util.spec_from_file_location( + "projection_reference", Path(__file__).parents[1] / "bench/bench_sycl_xmx_moe.py" + ) + reference_module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(reference_module) + ext, torch, cpu, _, *_ = setup + monkeypatch.setenv("KT_SYCL_XMX_HOST_QUANT", host_quant) + generator = torch.Generator().manual_seed(47) + mask = torch.tensor([False, False, False, True]) + config = ext.moe.MOEConfig(4, 4, 2048, 512, mask.data_ptr()) + config.pool = cpu.backend_ + config.max_len = 17 + config.quant_config.bits = 4 + config.quant_config.group_size = 128 + config.quant_config.zero_point = False + config.swiglu_limit = 1.5 + config.swiglu_alpha = 1.702 + proj = [ + reference_module.packed_weights(4, 512, 2048, 128, generator), + reference_module.packed_weights(4, 512, 2048, 128, generator), + reference_module.packed_weights(4, 2048, 512, 128, generator), + ] + for stage, (qw, sc) in zip(("gate", "up", "down"), proj): + setattr(config, stage + "_projs", [[t.data_ptr() for t in qw]]) + setattr(config, stage + "_scales", [[t.data_ptr() for t in sc]]) + mapping = torch.tensor([2, 0, 3, 1], dtype=torch.int64) + monkeypatch.setenv("KT_SYCL_XMX_MODE", "int8") + monkeypatch.setenv("KT_SYCL_XMX_DECODE_GATE_KERNEL", "xmx-split") + monkeypatch.setenv("KT_SYCL_XMX_DECODE_DOWN_KERNEL", "xmx") + monkeypatch.setenv("KT_SYCL_XMX_GATE_SPLITS", splits) + monkeypatch.setenv("KT_SYCL_XMX_FUSED_QUANT", "prefill") + monkeypatch.setenv("KT_SYCL_XMX_GRAPH", graph) + monkeypatch.setenv("KT_SYCL_XMX_DECODE_MERGE", "separate") + monkeypatch.setenv("KT_SYCL_XMX_FUSED_EXPERT", "0") + staged = ext.moe.SYCLXMXGPTQInt4_MOE(config) + monkeypatch.setenv("KT_SYCL_XMX_FUSED_EXPERT", "1") + fused = ext.moe.SYCLXMXGPTQInt4_MOE(config) + for moe in (staged, fused): + cpu.submit(moe.load_weights_task(mapping.data_ptr())) + cpu.sync() + for rows in (1, 17): + for route in ([0, 1, 2, -1], [1, 1, 0, 3], [-1, 99, 2, -1], [-1, -1, -1, -1]): + x = torch.randn(rows, 2048, generator=generator).bfloat16() + ids = torch.tensor([route] * rows, dtype=torch.int64) + weights = torch.softmax(torch.randn(rows, 4, generator=generator), 1) + qlen = torch.tensor([rows], dtype=torch.int32) + outputs = [] + for moe in (staged, fused): + out = torch.empty_like(x) + cpu.submit( + moe.forward_task( + qlen.data_ptr(), 4, ids.data_ptr(), weights.data_ptr(), x.data_ptr(), out.data_ptr(), False + ) + ) + cpu.sync() + outputs.append(out.clone()) + previous = out.clone() + cpu.submit( + moe.forward_task( + qlen.data_ptr(), 4, ids.data_ptr(), weights.data_ptr(), x.data_ptr(), out.data_ptr(), True + ) + ) + cpu.sync() + torch.testing.assert_close(out, (previous.float() * 2).bfloat16(), rtol=0, atol=0) + # Same group boundaries/scales and BF16 rounding; allow rare + # quantization ties from ESIMD vs scalar SIMD division. + denominator = outputs[0].float().square().mean().sqrt().clamp_min(1e-12) + rmse = (outputs[0].float() - outputs[1].float()).square().mean().sqrt() / denominator + assert rmse < 0.005 + + +def test_fused_expert_rejects_unsupported_dimensions(setup, monkeypatch): + ext, _, _, config, *_ = setup + monkeypatch.setenv("KT_SYCL_XMX_FUSED_EXPERT", "1") + with pytest.raises(RuntimeError, match="FUSED_EXPERT"): + ext.moe.SYCLXMXGPTQInt4_MOE(config) + + +@pytest.mark.parametrize( + "setting,value", + [("KT_SYCL_XMX_MODE", "fp16"), ("KT_SYCL_XMX_GATE_SPLITS", "4"), ("KT_SYCL_XMX_DECODE_MERGE", "fused")], +) +def test_fused_expert_rejects_incompatible_options(setup, monkeypatch, setting, value): + ext, torch, cpu, _, *_ = setup + mask = torch.zeros(4, dtype=torch.bool) + config = ext.moe.MOEConfig(4, 2, 2048, 512, mask.data_ptr()) + config.pool = cpu.backend_ + config.max_len = 1 + config.quant_config.bits = 4 + config.quant_config.group_size = 128 + config.quant_config.zero_point = False + monkeypatch.setenv("KT_SYCL_XMX_FUSED_EXPERT", "1") + monkeypatch.setenv("KT_SYCL_XMX_GATE_SPLITS", "2") + monkeypatch.setenv("KT_SYCL_XMX_DECODE_GATE_KERNEL", "xmx-split") + monkeypatch.setenv("KT_SYCL_XMX_DECODE_DOWN_KERNEL", "xmx") + monkeypatch.setenv("KT_SYCL_XMX_DECODE_MERGE", "separate") + monkeypatch.setenv(setting, value) + with pytest.raises(RuntimeError, match="FUSED_EXPERT"): + ext.moe.SYCLXMXGPTQInt4_MOE(config) + + +def test_xmx_implementation_revision_is_reported(setup): + ext, _, _, _, *_ = setup + assert ext.moe.SYCLXMXGPTQInt4_MOE.implementation_revision() >= 6 + + +@pytest.mark.parametrize( + "prefill_rows,columns,down_rows,load", + [ + ("16", "16", "16", "block"), + ("32", "16", "32", "block"), + ("8", "32", "8", "block"), + ("16", "32", "16", "block"), + ("8", "32", "16", "block"), + ("8", "32", "16", "gather"), + ("16", "32", "16", "gather"), + ], +) +@pytest.mark.parametrize("fused", ["0", "prefill"]) +@pytest.mark.parametrize("merge", ["scalar", "vector"]) +@pytest.mark.parametrize("incremental", [False, True]) +def test_wide_prefill_matches_original_with_tails( + setup, monkeypatch, prefill_rows, columns, down_rows, load, fused, merge, incremental +): + ext, torch, cpu, config, mask, qw, scales = setup + generator = torch.Generator().manual_seed(150) + qw.copy_(torch.randint(-(1 << 31), (1 << 31) - 1, qw.shape, dtype=torch.int32, generator=generator)) + scales.fill_(0.003) + config.max_len = 65 + monkeypatch.setenv("KT_SYCL_XMX_FUSED_QUANT", fused) + monkeypatch.setenv("KT_SYCL_XMX_GATE_KERNEL", "xmx") + monkeypatch.setenv("KT_SYCL_XMX_DOWN_KERNEL", "xmx") + monkeypatch.setenv("KT_SYCL_XMX_PREFILL_ROWS", "8") + monkeypatch.setenv("KT_SYCL_XMX_PREFILL_N", "16") + monkeypatch.setenv("KT_SYCL_XMX_PREFILL_DOWN_ROWS", "8") + monkeypatch.setenv("KT_SYCL_XMX_PREFILL_LOAD", "block") + monkeypatch.setenv("KT_SYCL_XMX_PREFILL_MERGE", "scalar") + original = ext.moe.SYCLXMXGPTQInt4_MOE(config) + monkeypatch.setenv("KT_SYCL_XMX_PREFILL_LOAD", load) + monkeypatch.setenv("KT_SYCL_XMX_PREFILL_MERGE", merge) + monkeypatch.setenv("KT_SYCL_XMX_PREFILL_ROWS", prefill_rows) + monkeypatch.setenv("KT_SYCL_XMX_PREFILL_N", columns) + monkeypatch.setenv("KT_SYCL_XMX_PREFILL_DOWN_ROWS", down_rows) + wide = ext.moe.SYCLXMXGPTQInt4_MOE(config) + for moe in (original, wide): + cpu.submit(moe.load_weights_task()) + cpu.sync() + for rows in (1, 17, 33, 65): + x = torch.randn(rows, 128, generator=generator).bfloat16() + ids = torch.tensor([[1, 1] if r % 3 == 0 else [r % 4, -1] for r in range(rows)], dtype=torch.int64) + weights = torch.softmax(torch.randn(rows, 2, generator=generator), 1) + qlen = torch.tensor([rows], dtype=torch.int32) + initial = torch.randn(rows, 128, generator=generator).bfloat16() + values = [] + for moe in (original, wide): + out = initial.clone() + cpu.submit( + moe.forward_task( + qlen.data_ptr(), 2, ids.data_ptr(), weights.data_ptr(), x.data_ptr(), out.data_ptr(), incremental + ) + ) + cpu.sync() + values.append(out) + assert torch.isfinite(values[1]).all() + torch.testing.assert_close(values[0], values[1], rtol=0, atol=0) + + +@pytest.mark.parametrize( + "name,value", + [ + ("KT_SYCL_XMX_PREFILL_ROWS", "12"), + ("KT_SYCL_XMX_PREFILL_N", "64"), + ("KT_SYCL_XMX_PREFILL_DOWN_ROWS", "0"), + ("KT_SYCL_XMX_PREFILL_MERGE", "invalid"), + ], +) +def test_invalid_prefill_tiling_is_rejected(setup, monkeypatch, name, value): + ext, _, _, config, *_ = setup + monkeypatch.setenv(name, value) + with pytest.raises(RuntimeError, match=name): + ext.moe.SYCLXMXGPTQInt4_MOE(config) + + +@pytest.mark.parametrize("mode,gate,down", [("fp16", "xmx", "xmx"), ("int8", "vector", "xmx"), ("int8", "xmx", "dp4a")]) +def test_n32_rejects_incompatible_prefill_projections(setup, monkeypatch, mode, gate, down): + ext, _, _, config, *_ = setup + monkeypatch.setenv("KT_SYCL_XMX_MODE", mode) + monkeypatch.setenv("KT_SYCL_XMX_GATE_KERNEL", gate) + monkeypatch.setenv("KT_SYCL_XMX_DOWN_KERNEL", down) + monkeypatch.setenv("KT_SYCL_XMX_PREFILL_N", "32") + with pytest.raises(RuntimeError, match="PREFILL_N"): + ext.moe.SYCLXMXGPTQInt4_MOE(config) + + +@pytest.mark.parametrize("load,columns", [("invalid", "32"), ("gather", "16"), ("packed", "32")]) +def test_prefill_load_rejects_invalid_or_unusable_modes(setup, monkeypatch, load, columns): + ext, _, _, config, *_ = setup + monkeypatch.setenv("KT_SYCL_XMX_PREFILL_LOAD", load) + monkeypatch.setenv("KT_SYCL_XMX_PREFILL_N", columns) + with pytest.raises(RuntimeError, match="PREFILL_LOAD"): + ext.moe.SYCLXMXGPTQInt4_MOE(config) + + +def test_panther_lake_profile_matches_explicit_v6_configuration(monkeypatch): + """Real executor comparison: a preset must preserve the existing arithmetic.""" + import importlib.util + from pathlib import Path + import torch + from kt_kernel import kt_kernel_ext as ext + + for name in list(os.environ): + if name.startswith("KT_SYCL_XMX_"): + monkeypatch.delenv(name) + spec = importlib.util.spec_from_file_location( + "profile_reference", Path(__file__).parents[1] / "bench/bench_sycl_xmx_moe.py" + ) + reference = importlib.util.module_from_spec(spec) + spec.loader.exec_module(reference) + generator = torch.Generator().manual_seed(20261016) + cpu = ext.CPUInfer(2) + mask = torch.zeros(4, dtype=torch.bool) + config = ext.moe.MOEConfig(4, 4, 2048, 512, mask.data_ptr()) + config.pool = cpu.backend_ + config.max_len = 33 + config.quant_config.bits = 4 + config.quant_config.group_size = 128 + config.quant_config.zero_point = False + projections = [ + reference.packed_weights(4, 512, 2048, 128, generator), + reference.packed_weights(4, 512, 2048, 128, generator), + reference.packed_weights(4, 2048, 512, 128, generator), + ] + for name, (qw, scales) in zip(("gate", "up", "down"), projections): + setattr(config, name + "_projs", [[t.data_ptr() for t in qw]]) + setattr(config, name + "_scales", [[t.data_ptr() for t in scales]]) + explicit = { + "KT_SYCL_XMX_MODE": "int8", + "KT_SYCL_XMX_ROWS": "8", + "KT_SYCL_XMX_GATE_KERNEL": "xmx", + "KT_SYCL_XMX_DOWN_KERNEL": "xmx", + "KT_SYCL_XMX_DECODE_GATE_KERNEL": "xmx-split", + "KT_SYCL_XMX_DECODE_DOWN_KERNEL": "xmx", + "KT_SYCL_XMX_GATE_SPLITS": "2", + "KT_SYCL_XMX_DOWN_SPLITS": "1", + "KT_SYCL_XMX_PREFILL_ROWS": "8", + "KT_SYCL_XMX_PREFILL_DOWN_ROWS": "16", + "KT_SYCL_XMX_PREFILL_N": "32", + "KT_SYCL_XMX_PREFILL_LOAD": "gather", + "KT_SYCL_XMX_PREFILL_MERGE": "vector", + "KT_SYCL_XMX_FUSED_EXPERT": "1", + "KT_SYCL_XMX_DECODE_MERGE": "separate", + "KT_SYCL_XMX_FUSED_QUANT": "prefill", + "KT_SYCL_XMX_NUMA_DISPATCH": "1", + "KT_SYCL_XMX_GRAPH": "0", + "KT_SYCL_XMX_HOST_QUANT": "0", + } + with monkeypatch.context() as settings: + for key, value in explicit.items(): + settings.setenv(key, value) + old = ext.moe.SYCLXMXGPTQInt4_MOE(config) + monkeypatch.setenv("KT_SYCL_XMX_PROFILE", "panther-lake") + tuned = ext.moe.SYCLXMXGPTQInt4_MOE(config) + assert not any(key in os.environ for key in explicit) + for moe in (old, tuned): + cpu.submit(moe.load_weights_task()) + cpu.sync() + for rows in (1, 17, 33): + x = torch.randn(rows, 2048, generator=generator).bfloat16() + ids = torch.tensor([[0, 1, 1, -1]] * rows, dtype=torch.int64) + weights = torch.softmax(torch.randn(rows, 4, generator=generator), dim=1) + qlen = torch.tensor([rows], dtype=torch.int32) + outputs = [] + for moe in (old, tuned): + output = torch.empty_like(x) + cpu.submit( + moe.forward_task( + qlen.data_ptr(), 4, ids.data_ptr(), weights.data_ptr(), x.data_ptr(), output.data_ptr(), False + ) + ) + cpu.sync() + outputs.append(output.clone()) + torch.testing.assert_close(*outputs, rtol=0, atol=0) + assert ext.moe.SYCLXMXGPTQInt4_MOE.implementation_revision() >= 7 + + +def test_panther_lake_profile_rejects_unsupported_shape(setup, monkeypatch): + ext, _, _, config, *_ = setup + monkeypatch.setenv("KT_SYCL_XMX_PROFILE", "panther-lake") + with pytest.raises(RuntimeError, match="requires H=2048 I=512 group=128"): + ext.moe.SYCLXMXGPTQInt4_MOE(config) + + +@pytest.mark.parametrize("profile", ["", "auto", "typo"]) +def test_invalid_tuning_profile_is_rejected(setup, monkeypatch, profile): + ext, _, _, config, *_ = setup + monkeypatch.setenv("KT_SYCL_XMX_PROFILE", profile) + with pytest.raises(RuntimeError, match="KT_SYCL_XMX_PROFILE"): + ext.moe.SYCLXMXGPTQInt4_MOE(config)