[mkumargarg]

Multi-Device Support - PR #527

Overview

This PR implements multi-architecture wheel support for the sgl-kernel-xpu project, enabling a single wheel to run on multiple Intel GPU architectures (Xe2/BMG and Xe3P/CRI) without requiring separate architecture-specific builds per deployment.

Problem Statement

Previously, each build was architecture-specific due to SYCL_INTEL_TARGET being a compile-time switch that affects common kernel sources. This meant:

• Separate wheels required for each GPU architecture
• Complex deployment and dependency management
• Users had to know their GPU architecture upfront

Solution

Implement a "fat wheel" strategy that:

1. Builds two separate wheel payloads (one per architecture: _xe20 for BMG, _xe35 for CRI)
2. Merges them into a single multi-architecture wheel
3. Uses a runtime architecture probe to dynamically select the correct backend at import time

Key Changes

1. Architecture Detection System

• include/sgl_arch_query.h (NEW): Unified header shared by probe and torch ops

  • device_architecture(): Gets raw device architecture via SYCL API
  • device_capability(): Maps architecture → (major, minor) version tuple

• src/arch/sgl_arch_probe.cpp (NEW): Lightweight probe library

  • Built independently of SYCL_INTEL_TARGET
  • Exposes C entry point sgl_probe_device_capability()
  • Installed at package root for dynamic backend selection

• src/sycl/Device.cpp (REFACTORED):

  • Removed duplicate architecture logic
  • Now delegates to shared sgl_arch_query.h

2. Per-Architecture Package Layout

• python/sgl_kernel/_xe20/__init__.py (NEW): Xe2/BMG backend stub
• python/sgl_kernel/_xe35/__init__.py (NEW): Xe3P/CRI backend stub
• Compiled binaries (*.so) install into respective architecture directories
• Enables disjoint binary directories for safe merging

3. Runtime Backend Selection

• python/sgl_kernel/_arch.py (NEW): Architecture detection and backend loader

  • probe_capability(): Calls C probe to get device (major, minor)
  • select_backend(): Maps capability to backend subpackage name
  • load_backend(): Imports and registers architecture-appropriate backend
  • Supports SGL_KERNEL_ARCH env var override for debugging/CI

• python/sgl_kernel/__init__.py (MODIFIED):

  • Tries direct import first (single-arch builds)
  • Falls back to _arch.load_backend() for multi-arch wheels

4. Wheel Merging Tool

• scripts/build_fat_wheel.py (NEW): Merges per-architecture wheels
  • Validates disjoint architecture directories
  • Recomputes wheel RECORD manifest
  • Produces canonical wheel filename
  • Includes validation that probe library is present
• scripts/build_multiarch_wheel.py (NEW): build and merge multiple wheels in single command

5. Build System Updates

• CMakeLists.txt (MODIFIED):

  • Torch import error detection (missing --no-isolation)
  • SGL_ARCH_PKG_DIR variable controls per-arch install location

• src/BuildOnLinux.cmake (MODIFIED):

  • Build sgl_arch_probe independently
  • Install to package root (not per-arch)
  • Route common_ops and kernel libraries to per-arch directories

Build & Test Instructions

Building Multi-Architecture Wheel

source /PATH/TO/ONEAPI/setvars.sh

# Pass 1 — Build for Xe2 (BMG)
python -m build --wheel --no-isolation \
  -C build-dir=build \
  -C cmake.define.DPCPP_SYCL_TARGET=bmg \
  -o dist_bmg

# Pass 2 — Build for Xe3P (CRI)
python -m build --wheel --no-isolation \
  -C build-dir=build \
  -C cmake.define.DPCPP_SYCL_TARGET=cri \
  -o dist_cri

# Merge into single multi-architecture wheel
python scripts/build_fat_wheel.py dist_fat dist_bmg/*.whl dist_cri/*.whl


Alternatively, "python scripts/build_multiarch_wheel.py" might be used to do #Pass 1, #Pass 2 and merging of multiple wheel files (BMG and CRI) into single wheel file.

# Install and verify
pip install --force-reinstall dist_fat/*.whl
python -c "import sgl_kernel, sgl_kernel._arch as a; print(a.probe_capability(), a.select_backend())"

[Copilot]

🟡 Changes recommended

The unguarded CRI enum breaks supported older-toolchain builds, and relative output paths can prevent the build helper from locating merged wheels.

2 open issues

• Guard CRI enum usage for older SYCL compiler versions · New
• Resolve output directory before merger execution · New

What changed in this PR

This PR enables a single sgl-kernel-xpu wheel to support BMG and CRI GPUs by packaging separate architecture backends and selecting one at import time.

Changes:

• Adds shared architecture detection and runtime backend loading.
• Moves compiled libraries into per-architecture package directories.
• Adds scripts to build and merge architecture-specific wheels.

File	Description
src/​sycl/​Device.cpp	Delegates device queries to the shared helper.
src/​BuildOnLinux.cmake	Builds the probe and installs libraries by architecture.
src/​arch/​sgl_arch_probe.cpp	Exposes a C entry point for architecture detection.
scripts/​build_multiarch_wheel.py	Orchestrates builds, merging, and optional installation.
scripts/​build_fat_wheel.py	Merges wheel payloads and regenerates RECORD.
python/​sgl_kernel/​_xe35/​__init__.py	Adds the CRI backend package stub.
python/​sgl_kernel/​_xe20/​__init__.py	Adds the BMG backend package stub.
python/​sgl_kernel/​_arch.py	Detects the architecture and loads its backend.
python/​sgl_kernel/​__init__.py	Adds fallback loading for per-architecture backends.
include/​sgl_arch_query.h	Centralizes architecture-to-capability mapping.
CMakeLists.txt	Selects backend install paths and checks Torch import failures.

_{Review effort: Balanced}

[airMeng]

Reviewed the latest revision. The older-toolchain CRI enum guard is fixed, and the actual architecture probe builds with the production CMake/GNU-host path on oneAPI 2026.1 and correctly detects BMG. Requesting changes for the relative output-directory issue described inline, confirming the existing Copilot finding. I reproduced it with actual merger subprocesses and synthetic architecture wheel payloads; an absolute output directory succeeds. Full dual-architecture wheel installation and CRI execution were not tested.

[airMeng]

[P2] Resolve the output directory before invoking the merger.

run() executes the merger with cwd=PROJECT_ROOT, whereas this glob resolves a relative args.out_dir against the caller's working directory. Invoking python /path/to/repo/scripts/build_multiarch_wheel.py --skip-build --out-dir result from outside the repository creates repo/result/<wheel>.whl, then searches caller/result and exits with expected one merged wheel ... found []; --install is never reached. I confirmed this with the actual merger subprocess: the relative path reproduces the failure and an absolute path succeeds. Please normalize the output directory to an absolute path once and reuse it for merging, discovery and installation.

[mkumargarg]

Fixed: args.out_dir (and args.build_dir, which had the same problem with -C build-dir=) are now converted to absolute paths once, right after parse_args(). The merger, the wheel lookup and --install all use that path.