Repository navigation
[feat](kt-kernel): add Intel XPU backend with a SYCL/Level Zero stream bridge - #2251
Draft
siju-samuel wants to merge 2 commits into
Draft
siju-samuel wants to merge 2 commits into
siju-samuel wants to merge 2 commits into
Conversation
…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.
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Adds Intel XPU (
torch.xpu) as a kt-kernel backend, so the CPU expert kernels can be scheduledfrom 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 incpu_backend/vendors/because their torch builds pose astorch.cuda; Ascend supplies its owncallback worker.
torch.xpuhas neither, so there was no path for Intel GPUs: anxputensor fellinto the CUDA branch and handed a
sycl::queue*to a parameter typedcudaStream_t. The existingKTRANSFORMERS_USE_SYCLbackend is a different feature — it offloads GPTQ-INT4 experts to an InteliGPU 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).cpu_backend/xpu_stream_bridge.cppiscompiled against the SYCL and Level Zero headers, through per-source properties, and the target
links
libsycl+libze_loader. No icpx anywhere — unlikeKTRANSFORMERS_USE_SYCL, whichforces icpx on the whole build and drops
-ffast-math.libsyclthat torch itself loads, read from/proc/self/maps, and warns whenthat differs from what it is about to link. Deriving the path instead is unreliable:
sysconfig'sLIBDIRis the base interpreter's, so a venv's ownlib/is never probed, andfind_librarycannot match a soname-only
libsycl.so.9(the usual wheel layout). Either miss puts a second SYCLruntime in the process and has it dereference torch's
sycl::queue*.torch.xpuand every dev file the CMake branch needs, so a hostmissing them keeps building CPU-only rather than failing a build that used to work. An explicit
CPUINFER_USE_XPU=1still errors, naming what is absent. It stands down for an explicitlyrequested 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 behindKT_XPU_BRIDGE:l0(default)ext_oneapi_submit_barrier(), a per-queue worker thread waits on ithost_taskhost_taskhost_taskhost_taskis exactlycudaLaunchHostFunc's semantics and survives XPUGraph capture, so recordingqueues 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 parkederror.
Python —
kt_kernel.accel(current_stream_handle/synchronize/is_capturing) so callerssuch as the SGLang integration need no device branching;
xpubranches inexperts_base(
_wait_device, and a synchronous fallback with a one-time warning when the extension has nobridge, in either direction);
sft/base.pyroutesxputo the synchronous inference path, sincetorch.xpuhas noExternalStream.__init__.pyimports torch before the extension so GPU buildsresolve their runtime library through torch rather than a build-time RPATH, and registers the
bridge's shutdown hook there so direct
CPUInferusers 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_pendingdrain, asynchrony, XPUGraph capture/replay, worker-thread lifecycle by threadname), plus
test_accel_device_helpers.pyandtest_gpu_backend_attr.py, which run on theexisting CPU runner; the latter pins the
gpu_backend#iforder the xpu path depends on.register_xpu_ci,--hw xpuand thexpumarker are wired up. The eight AMX accuracy/bench testsno 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 inNEEDED.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, andfour threads racing
launch_*against ashutdown()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:
KT_XPU_BRIDGE=l0KT_XPU_BRIDGE=host_taskNumerics 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
none of those toolkits. Every new arm is an
#elifappended after the existing chains and noexisting line was modified, so those paths are dead when their macros are set;
g++ -Eper defineset confirms
cudaLaunchHostFuncis still the only active body for CUDA and for the SYCL-iGPUbuild, and that the bridge TU preprocesses to 6 lines in every non-XPU build.
USE_SYCLis tested first in the newgpu_backendchain, because the SYCL-iGPU build definesKTRANSFORMERS_USE_CUDA_HOST_CALLBACKSrather thanKTRANSFORMERS_USE_CUDA. Reordering it wouldsilently change xpu's behaviour, which is what
test_gpu_backend_attr.pyguards.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 buildinterpreter;
setup.pysays so and the README documentsCPUINFER_USE_XPU=1plus--no-deps(the project pins
torch==2.9.1, which would otherwise replace a user's XPU torch).--hw xpuis not exercised in CI yet; the two new CPU tests dorun on the existing runner. Separately,
run_suite.pycannot run for any--hwtoday becausecollect_tests(sanity_check=True)raises on the 21per_commitfiles that have noregister_*_cicall — pre-existing onmain, not touched here.auditwheel repairneeds--exclude libsycl.so.9 --exclude libze_loader.so.1(plus thelibur_*/libmkl_sycl_*set) orit will vendor a second SYCL runtime into the wheel.
KT_XPU_BRIDGE_EVENTShas a floor but no ceiling (a typo failswith 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.mdtutorial following theAscend/SYCL pattern, since this PR documents the backend in
README.mdonly.[fix](kt-kernel): mark KML-only builds as CPU-only) is independent of theXPU 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 norKTRANSFORMERS_CPU_ONLY, so a KML build exposessubmit_with_cuda_streamwith every#ifarm inits 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