Skip to content

[feat](kt-kernel): add Intel XPU backend with a SYCL/Level Zero stream bridge - #2251

Draft
siju-samuel wants to merge 2 commits into
kvcache-ai:mainfrom
siju-samuel:feat/kt-kernel-xpu-backend
Draft

siju-samuel wants to merge 2 commits into
kvcache-ai:mainfrom
siju-samuel:feat/kt-kernel-xpu-backend

Conversation

@siju-samuel

@siju-samuel siju-samuel commented Oct 21, 2026 •

Copy link
Copy Markdown

What does this PR do?

Adds Intel XPU (torch.xpu) as a kt-kernel backend, so the CPU expert kernels can be scheduled
from a torch XPU stream the way they are from a CUDA stream today.

kt-kernel hands CPU expert work to the GPU stream with cudaLaunchHostFunc
(cpu_backend/cpuinfer.h). ROCm, MUSA and MACA reuse that call through the alias headers in
cpu_backend/vendors/ because their torch builds pose as torch.cuda; Ascend supplies its own
callback worker. torch.xpu has neither, so there was no path for Intel GPUs: an xpu tensor fell
into the CUDA branch and handed a sycl::queue* to a parameter typed cudaStream_t. The existing
KTRANSFORMERS_USE_SYCL backend is a different feature — it offloads GPTQ-INT4 experts to an Intel
iGPU while the model runs on an NVIDIA GPU, and still needs the CUDA toolkit for this bridge.

XPU is added as a peer backend shaped like Ascend, not as an alias header.

What changes

Build — KTRANSFORMERS_USE_XPU / CPUINFER_USE_XPU (CMakeLists.txt, setup.py,
install.sh, CMakePresets.json, README.md).

  • The CPU kernels keep the host compiler and its flags. Only cpu_backend/xpu_stream_bridge.cpp is
    compiled against the SYCL and Level Zero headers, through per-source properties, and the target
    links libsycl + libze_loader. No icpx anywhere — unlike KTRANSFORMERS_USE_SYCL, which
    forces icpx on the whole build and drops -ffast-math.
  • CMake links the libsycl that torch itself loads, read from /proc/self/maps, and warns when
    that differs from what it is about to link. Deriving the path instead is unreliable: sysconfig's
    LIBDIR is the base interpreter's, so a venv's own lib/ is never probed, and find_library
    cannot match a soname-only libsycl.so.9 (the usual wheel layout). Either miss puts a second SYCL
    runtime in the process and has it dereference torch's sycl::queue*.
  • Auto-detection requires torch.xpu and every dev file the CMake branch needs, so a host
    missing them keeps building CPU-only rather than failing a build that used to work. An explicit
    CPUINFER_USE_XPU=1 still errors, naming what is absent. It stands down for an explicitly
    requested backend, for a detected CUDA toolkit, and for multi-variant (distribution) builds, which
    must stay importable where there is no SYCL runtime.

Bridge — cpu_backend/xpu_stream_bridge.cpp (+ vendors/xpu.h), two modes behind
KT_XPU_BRIDGE:

mode submit edge (GPU→CPU) sync edge (CPU→GPU) cost/edge
l0 (default) ext_oneapi_submit_barrier(), a per-queue worker thread waits on it the same barrier, plus one on a host-visible Level Zero event the worker signals after the drain ~2 µs submit / ~10 µs sync
host_task SYCL host_task SYCL host_task ~15 µs

host_task is exactly cudaLaunchHostFunc's semantics and survives XPUGraph capture, so recording
queues select it automatically. Because host tasks bypass the worker FIFO, a queue switching to that
path drains its FIFO first — otherwise a recorded drain could overtake an l0 submit still waiting on
the GPU, release the GPU early, and let the H2D copy read a buffer the CPU experts had not written.
Exceptions from CPU tasks are parked and rethrown by the next call with the GPU event still
signalled, so the queue cannot hang. shutdown() excludes in-flight launches and clears the parked
error.

Python — kt_kernel.accel (current_stream_handle / synchronize / is_capturing) so callers
such as the SGLang integration need no device branching; xpu branches in experts_base
(_wait_device, and a synchronous fallback with a one-time warning when the extension has no
bridge, in either direction); sft/base.py routes xpu to the synchronous inference path, since
torch.xpu has no ExternalStream. __init__.py imports torch before the extension so GPU builds
resolve their runtime library through torch rather than a build-time RPATH, and registers the
bridge's shutdown hook there so direct CPUInfer users are covered too.

Tests — 12 bridge tests (ordering of both edges in both modes against a torch reference,
event-ring wraparound with a 2-slot ring and no host sync, two concurrent streams, the deferred
allow_n_pending drain, asynchrony, XPUGraph capture/replay, worker-thread lifecycle by thread
name), plus test_accel_device_helpers.py and test_gpu_backend_attr.py, which run on the
existing CPU runner; the latter pins the gpu_backend #if order the xpu path depends on.
register_xpu_ci, --hw xpu and the xpu marker are wired up. The eight AMX accuracy/bench tests
no longer create their weights with device="cuda", so they now run on CPU-only and XPU hosts too.

Verification

Host: 2x Xeon 6980P (AMX) + 8x Arc Pro B60, oneAPI 2026.0, torch 2.14 xpu.

  • XPU build: 31 tests pass, also under random order. CPU-only build: 19 pass + 1 expected
    skip
    , extension reports gpu_backend == "cpu" with no stream bindings and no SYCL in NEEDED.

  • A standalone C++ harness drives the bridge directly in both modes: ordering, event-ring reuse,
    shutdown/restart, a throwing drain and a throwing fn, shutdown() from inside a callback, and
    four threads racing launch_* against a shutdown() loop (~74k edges, ~2k clean "shutting down"
    rejections, no hang and no unexpected exception).

  • 60-layer AMX INT8 decode step, qlen=1, top-8, H=7168, over several runs on a shared box:

    path eager XPUGraph replay
    synchronous (no bridge) 105-130 ms n/a
    KT_XPU_BRIDGE=l0 87-95 ms ~300 ms
    KT_XPU_BRIDGE=host_task 170-300 ms ~290 ms

    Numerics 0.029-0.031 rel. diff against a torch reference throughout (AMX INT8 threshold 0.05).
    Graph replay is ~3x slower than eager because host nodes inside SYCL graphs cost milliseconds, so
    graphs are best left off for this path until Level Zero events work inside recorded graphs.

Notes for reviewers

  • CUDA, ROCm, MUSA, MACA and Ascend were reviewed at the preprocessor level only — this box has
    none of those toolkits. Every new arm is an #elif appended after the existing chains and no
    existing line was modified, so those paths are dead when their macros are set; g++ -E per define
    set confirms cudaLaunchHostFunc is still the only active body for CUDA and for the SYCL-iGPU
    build, and that the bridge TU preprocesses to 6 lines in every non-XPU build.
  • USE_SYCL is tested first in the new gpu_backend chain, because the SYCL-iGPU build defines
    KTRANSFORMERS_USE_CUDA_HOST_CALLBACKS rather than KTRANSFORMERS_USE_CUDA. Reordering it would
    silently change xpu's behaviour, which is what test_gpu_backend_attr.py guards.
  • The ordering test is deliberately shaped so a broken sync edge cannot hide: GPU filler work goes
    before the D2H copies, never between submit and the H2D. With filler in between (an earlier
    version) the CPU experts finish during it and the test passes even with the sync barrier removed
    entirely — verified by substituting a no-op sync edge. Worth preserving if the test is
    restructured. The margin is ~3 ms of CPU MoE against a 0.03 ms copy.
  • pip install . cannot auto-detect XPU because pip's build isolation hides torch from the build
    interpreter; setup.py says so and the README documents CPUINFER_USE_XPU=1 plus --no-deps
    (the project pins torch==2.9.1, which would otherwise replace a user's XPU torch).
  • No XPU CI runner exists, so --hw xpu is not exercised in CI yet; the two new CPU tests do
    run on the existing runner. Separately, run_suite.py cannot run for any --hw today because
    collect_tests(sanity_check=True) raises on the 21 per_commit files that have no
    register_*_ci call — pre-existing on main, not touched here.
  • Before an XPU wheel job is added, its auditwheel repair needs
    --exclude libsycl.so.9 --exclude libze_loader.so.1 (plus the libur_* / libmkl_sycl_* set) or
    it will vendor a second SYCL runtime into the wheel.
  • Known limits, left as follow-up: KT_XPU_BRIDGE_EVENTS has a floor but no ceiling (a typo fails
    with a clear Level Zero error); nothing evicts workers, so a caller rotating torch.xpu.Stream()
    would accumulate one thread and one event pool per queue (SGLang uses one stream, and the
    two-stream test covers the real case); a doc/en/kt-kernel/xpu.md tutorial following the
    Ascend/SYCL pattern, since this PR documents the backend in README.md only.
  • The second commit ([fix](kt-kernel): mark KML-only builds as CPU-only) is independent of the
    XPU work and can be dropped or taken separately. It fixes a latent bug the new test exposed: the
    elseif(KTRANSFORMERS_CPU_USE_KML) branch defines neither a GPU backend macro nor
    KTRANSFORMERS_CPU_ONLY, so a KML build exposes submit_with_cuda_stream with every #if arm in
    its body false — the call returns without enqueueing and the CPU MoE result is silently stale.
    I could not test this on ARM Kunpeng hardware; a KML configure now carries
    KTRANSFORMERS_CPU_ONLY=1, but please confirm the runtime behaviour.

Before submitting

…m bridge

kt-kernel hands its CPU expert work to the GPU stream with cudaLaunchHostFunc.
ROCm, MUSA and MACA reuse that call through alias headers because their torch
builds pose as torch.cuda; Ascend supplies its own callback worker. torch.xpu
has neither, so Intel GPUs had no way to schedule the CPU MoE from a stream:
an xpu tensor fell into the CUDA branch and passed a sycl::queue* where a
cudaStream_t was expected.

This adds XPU as a peer backend, shaped like Ascend rather than as an alias
header.

Build: KTRANSFORMERS_USE_XPU / CPUINFER_USE_XPU. The CPU kernels keep the host
compiler and its flags; only cpu_backend/xpu_stream_bridge.cpp sees the SYCL and
Level Zero headers (per-source properties) and the target links libsycl and
libze_loader. No icpx anywhere, unlike KTRANSFORMERS_USE_SYCL, which switches
the whole build over and drops -ffast-math. CMake links the libsycl that torch
itself loaded, read from /proc/self/maps: deriving it from sysconfig's LIBDIR
misses a venv's own lib (LIBDIR is the base interpreter's) and find_library
cannot match a soname-only libsycl.so.9, either of which would put a second SYCL
runtime in the process and have it dereference torch's sycl::queue*.
Auto-detection requires torch.xpu plus every dev file the CMake branch needs, so
a host without them keeps building CPU-only instead of failing; an explicit
CPUINFER_USE_XPU=1 still errors with what is missing.

Bridge: two modes behind KT_XPU_BRIDGE. "l0" (default) submits a barrier for the
GPU->CPU edge and has a per-queue worker thread wait on it, and for the CPU->GPU
edge makes the queue wait on a host-visible Level Zero event that the worker
signals after the drain; about 12 us per edge. "host_task" uses SYCL host tasks,
which are exactly cudaLaunchHostFunc's semantics and survive XPUGraph capture
but cost far more, so recording queues select it automatically. Exceptions from
CPU tasks are parked and rethrown by the next call with the GPU event still
signalled, so the queue cannot hang.

Python: kt_kernel.accel (current_stream_handle / synchronize / is_capturing) so
callers such as the SGLang integration need no device branching; xpu branches in
experts_base (_wait_device, atexit shutdown, and a synchronous fallback with a
one-time warning when the extension has no bridge); sft/base.py sends xpu to the
synchronous inference path, since torch.xpu has no ExternalStream. __init__.py
imports torch before the extension so GPU builds resolve their runtime through
torch rather than a build-time RPATH.

Tests: 12 bridge tests (ordering of both edges in both modes against a torch
reference, event-ring wraparound with no host sync, two streams, the deferred
allow_n_pending drain, asynchrony, XPUGraph replay, worker-thread lifecycle),
plus test_accel_device_helpers.py and test_gpu_backend_attr.py, which run on the
existing CPU runner and pin the gpu_backend #if order the xpu path depends on.
register_xpu_ci / --hw xpu / the xpu marker are wired up, though no XPU runner
exists yet. The eight AMX accuracy and bench tests no longer need CUDA to build
their weights, so they now run on CPU-only and XPU hosts.

Verified on 2x Xeon 6980P + 8x Arc Pro B60 (oneAPI 2026.0, torch 2.14 xpu):
31 tests pass on the XPU build, 19 plus the expected skip on a CPU-only build,
both also in random order. A 60-layer AMX INT8 decode step takes 87-95 ms with
the l0 bridge against 105-130 ms for the synchronous path; the ranges are over
several runs, since the box is shared. XPU graph replay is about 3x slower than
eager because host nodes inside SYCL graphs cost milliseconds, so graphs are
best left off for this path for now. CUDA, ROCm, MUSA, MACA and Ascend were
reviewed at the preprocessor level only; this box has none of those toolkits.
The backend chain in CMakeLists.txt defines KTRANSFORMERS_CPU_ONLY in its final
else() but not in the elseif(KTRANSFORMERS_CPU_USE_KML) branch above it, so a
KML build (ARM Kunpeng, no GPU) defines neither that macro nor any GPU backend
macro. ext_bindings.cpp binds submit_with_cuda_stream / sync_with_cuda_stream
under #ifndef KTRANSFORMERS_CPU_ONLY, so they are exposed, while every #if arm
in their bodies in cpuinfer.h is false: the calls return without enqueueing
anything and the CPU MoE result is silently whatever was left in the output
buffer.

Declaring the branch CPU-only removes the bindings, so callers that probe with
hasattr() take their synchronous submit/sync path instead of losing the work.

Found by the new test_gpu_backend_attr.py, which asserts that a build reporting
gpu_backend == "cpu" does not expose the stream callbacks. Not reproducible on
this x86 host: a KML configure now carries KTRANSFORMERS_CPU_ONLY=1, but the
runtime behaviour needs an ARM Kunpeng machine to confirm.
@siju-samuel siju-samuel changed the title Feat/kt kernel xpu backend [feat](kt-kernel): add Intel XPU backend with a SYCL/Level Zero stream bridge Oct 21, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant