[feat](sycl): add opt-in Panther Lake XMX expert backend (#2250)

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

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.