Add Inkling attention prologue MXFP8 tau path

[yuankuns]

Summary

• Add Inkling attention prologue XPU log_scaling_tau support for verify/decode/extend q paths.
• Add MXFP8 q quantization plus K/V cache store scale-buffer writes using the upstream [pages, Hkv, 32, page_size / 32, 4] layout.
• Add focused XPU reference tests for scalar tau, MXFP8 verify/decode/extend, valid-slot gating, and no-store behavior.

Tests

• SKIP=black-jupyter python3 -m pre_commit run --files include/sgl_kernel_ops.h python/sgl_kernel/inkling_attn_prologue.py src/sycl/InklingAttnPrologue.cpp src/torch_extension_inkling_attn_prologue.cc tests/test_inkling_attn_prologue.py --show-diff-on-failure
• docker exec sglang-syk bash -lc 'source /opt/intel/oneapi/setvars.sh --force >/dev/null 2>&1; cd /workspace/worktrees/sgl-kernel-xpu/port-inkling-kernel-to-sglang; build_dir=$(cat /tmp/sgl_kernel_attn_mxfp8_build_dir); ninja -C "$build_dir" inkling_attn_prologue_ops'
• docker exec sglang-syk bash -lc 'source /opt/intel/oneapi/setvars.sh --force >/dev/null 2>&1; cd /workspace/worktrees/sgl-kernel-xpu/port-inkling-kernel-to-sglang; build_dir=$(cat /tmp/sgl_kernel_attn_mxfp8_build_dir); pkg_root=$(mktemp -d /tmp/sgl_kernel_attn_mxfp8_pkg.XXXXXX); mkdir -p "$pkg_root/sgl_kernel"; touch "$pkg_root/sgl_kernel/__init__.py"; ln -s "$PWD/python/sgl_kernel/inkling_attn_prologue.py" "$pkg_root/sgl_kernel/inkling_attn_prologue.py"; ln -s "$build_dir/src/inkling_attn_prologue_ops.abi3.so" "$pkg_root/sgl_kernel/inkling_attn_prologue_ops.abi3.so"; LD_LIBRARY_PATH="$build_dir/src:${LD_LIBRARY_PATH:-}" PYTHONPATH="$pkg_root:/workspace/python-targets/py312-xpu:$PWD/tests" ZE_AFFINITY_MASK=0 ONEAPI_DEVICE_SELECTOR=level_zero:gpu /root/miniforge3/envs/py312/bin/python -m pytest -q tests/test_inkling_attn_prologue.py'

Dependency

• Dependent PRs: this PR builds on Add Inkling attention prologue XPU ops #324 and is part of the Inkling XPU stack with Add Inkling sconv XPU ops #325.
• Review note: review the final commit b6aa54b for the MXFP8/tau delta.

[Copilot]

Pull request overview

Adds Inkling attention-prologue support to the XPU extension stack, extending the existing verify/decode/extend kernels with per-token log_scaling_tau handling on the Q path and an MXFP8 quantized-Q + MXFP8 KV-store path (including the upstream scale-buffer layout), plus focused XPU reference tests.

Changes:

• Extend Inkling attention prologue kernels (SYCL + Torch registrations + Python wrapper) with log_scaling_tau for Q and MXFP8 quantization / KV scale-buffer writes.
• Add a scoped inkling_attn_prologue_ops extension target and adjust build/link/install wiring to keep schemas out of common_ops.
• Add comprehensive XPU reference tests covering scalar tau, MXFP8 verify/decode/extend, slot-gating, and no-store behavior.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
tests/test_inkling_attn_prologue.py	Adds reference-model tests for verify/decode/extend, including tau and MXFP8 store/scale-layout checks.
src/torch_extension_inkling_attn_prologue.cc	Registers Inkling attention prologue Torch schemas in a scoped extension module.
src/sycl/InklingAttnPrologue.cpp	Implements SYCL kernels for verify/decode/extend, plus MXFP8 Q quantization and MXFP8 KV store with scale-buffer writes.
src/CMakeLists.txt	Excludes the scoped Inkling extension source from common_ops build compilation unit list.
src/BuildOnLinux.cmake	Builds/installs the scoped inkling_attn_prologue_ops module and adjusts linking so common_ops doesn’t transitively load the Inkling SYCL lib.
python/sgl_kernel/inkling_attn_prologue.py	Python wrapper that loads the scoped extension, routes to MXFP8 vs non-MXFP8 ops, and normalizes tau inputs.
python/sgl_kernel/__init__.py	Exposes Inkling attention prologue entry points at package top-level.
include/sgl_kernel_ops.h	Declares the new Inkling attention prologue C++ entry points (including MXFP8 variants).

Comments suppressed due to low confidence (4)

src/torch_extension_inkling_attn_prologue.cc:32

• In the Torch schema, k_cache and v_cache are both annotated as Tensor(a!) (and k_buf/v_buf are also both Tensor(a!)), implying all four may alias. These are separate mutable tensors and should have distinct alias sets to keep PyTorch’s alias analysis correct.

      "inkling_attn_prologue_decode(Tensor qkvr, Tensor(a!) k_cache, Tensor(a!) v_cache, "
      "Tensor cache_indices, Tensor cache_mask, Tensor k_weight, Tensor v_weight, "
      "Tensor? track_mask, Tensor? track_indices, Tensor q_gamma, Tensor k_gamma, float eps, "
      "Tensor loc, Tensor(a!) k_buf, Tensor(a!) v_buf, int q_off, int k_off, int v_off, "

src/torch_extension_inkling_attn_prologue.cc:42

• In the Torch schema, k_cache and v_cache are both annotated as Tensor(a!) (and k_buf/v_buf are also both Tensor(a!)), implying these mutable tensors may alias each other. They are logically distinct buffers and should use distinct alias sets.

      "inkling_attn_prologue_extend(Tensor qkvr, Tensor(a!) k_cache, Tensor(a!) v_cache, "
      "Tensor cache_indices, Tensor cache_mask, Tensor has_initial_state, Tensor cu, Tensor si, "
      "Tensor k_weight, Tensor v_weight, Tensor? track_rows, Tensor? track_mask, Tensor? track_dst, "
      "Tensor q_gamma, Tensor k_gamma, float eps, Tensor loc, Tensor(a!) k_buf, Tensor(a!) v_buf, "

src/torch_extension_inkling_attn_prologue.cc:59

• In the MXFP8 decode schema, k_cache and v_cache are both annotated as Tensor(a!), implying they may alias. They are distinct mutable tensors and should not share an alias set.

      "inkling_attn_prologue_decode_mxfp8(Tensor qkvr, Tensor(a!) k_cache, Tensor(a!) v_cache, "

src/torch_extension_inkling_attn_prologue.cc:69

• In the MXFP8 extend schema, k_cache and v_cache are both annotated as Tensor(a!), implying they may alias. They are distinct mutable tensors and should not share an alias set.

      "inkling_attn_prologue_extend_mxfp8(Tensor qkvr, Tensor(a!) k_cache, Tensor(a!) v_cache, "