Decode LoRA kernels are optimized by srujan-thomdapu · Pull Request #550 · sgl-project/sgl-kernel-xpu

srujan-thomdapu · 2026-10-15T06:48:07Z

Summary

Rewrites the chunked-SGMV LoRA decode kernels (chunked_sgmv_lora_shrink_fwd,
chunked_sgmv_lora_expand_fwd) as gather / scatter grouped GEMMs on sycl-tla's GemmUniversalAdapter
/ GemmUniversal, with custom LoRA collectives (gathered-A mainloop, scatter epilogues, on-device
segment scheduler). The value-level input checks of all LoRA ops now run only in RelWithDebInfo
builds.

Effective bandwidth goes from ~47 GB/s to ~220 GB/s (shrink) and ~270 GB/s (expand). Both ops beat the
Triton csgmv kernels (sglang's deployed XPU defaults) by 2.7–4.8x. In a default build, no LoRA op issues a device-to-host sync per call, except embedding_lora_a_fwd when max_len is omitted.

Performance

Intel Arc B580, triton.testing.do_bench (L2 flushed between runs),
benchmark/bench_chunked_sgmv_lora_{shrink,expand}_fwd.py. Triton uses sglang's deployed XPU csgmv
defaults (chunk 16; shrink BLOCK_N=16 / BLOCK_K=256, expand BLOCK_N=64 / BLOCK_K=16).

Mean effective bandwidth (GB/s; 9 cases, tok = 50 … 8192):

op	dtype	before	after	Triton	after vs Triton (geomean)
shrink	bf16	~48	222	49	4.8x, 9/9 cases
shrink	fp16	~47	214	52	4.2x, 9/9 cases
expand	bf16	~48	270	93	2.7x, 9/9 cases
expand	fp16	~47	265	94	2.7x, 8/9 cases*

* The one loss is the script's first-timed case (tok=50). Timed on its own, it is 1.5x faster than
Triton.

Behaviour changes

Value checks are debug-build only. In other builds, an invalid seg_indptr, index range,
permutation or offset table is undefined behaviour (wrong output or a GPU fault) instead of a
TORCH_CHECK error. Build with RelWithDebInfo to validate inputs.
Pass one segment per adapter. The kernels tile M themselves, so callers should pass coalesced
segments. This requirement is unchanged.
Python API: unchanged except for docstrings and a new optional keyword-only max_len on
embedding_lora_a_fwd.

Testing

Default build: 292 passed, 21 skipped. The skips are the value-check cases, which run only
against a RelWithDebInfo build with SGL_KERNEL_DEBUG_CHECKS=1.
New shape coverage in the chunked tests: K tails, odd K/N, odd ranks, uneven and unaligned slice
starts, and misaligned pointers.
New debug_checks pytest marker (tests/conftest.py): the existing prefill value-check cases
carry it.

Copilot

🟡 Changes recommended

Debug-test collection can fail without an XPU, and the synchronization-free claim needs qualification.

2 open findings

What changed in this PR

Optimizes Intel XPU LoRA decode by replacing separate gather, matrix multiplication, and scatter steps with gather-fused DPAS kernels, while preserving Python signatures.

Changes:

Adds optimized shrink and expand kernels with alignment and tail handling.
Moves value-level validation behind USE_DEBUG_CHECKS.
Extends shape tests and benchmarks, and updates API documentation.

File	Description
tests/test_sgemm_lora_b_fwd.py	Marks debug-only validation cases.
tests/test_sgemm_lora_a_fwd.py	Marks debug-only validation cases.
tests/test_qkv_lora_b_fwd.py	Marks debug-only validation cases.
tests/test_lora_debug_checks.py	Adds shared validation tests.
tests/test_gate_up_lora_b_fwd.py	Marks debug-only validation cases.
tests/test_chunked_sgmv_lora_shrink_fwd.py	Tests tails and misaligned inputs.
tests/test_chunked_sgmv_lora_expand_fwd.py	Tests irregular slices and alignment.
tests/conftest.py	Registers and skips debug-check tests.
src/sycl/SGEMMLoraBFwd.cpp	Uses conditional validation helpers.
src/sycl/SGEMMLoraAFwd.cpp	Uses conditional validation helpers.
src/sycl/QKVLoraBFwd.cpp	Gates value and offset checks.
src/sycl/kernels/lora/device/chunked_sgmv_lora_shrink_fwd_types.hpp	Defines split-K tile settings.
src/sycl/kernels/lora/device/chunked_sgmv_lora_shrink_fwd_runner.hpp	Launches alignment-aware shrink kernels.
src/sycl/kernels/lora/device/chunked_sgmv_lora_shrink_fwd_kernel.hpp	Implements gather-fused shrink.
src/sycl/kernels/lora/device/chunked_sgmv_lora_shrink_fwd_dispatch.hpp	Updates shrink launch declarations.
src/sycl/kernels/lora/device/chunked_sgmv_lora_expand_fwd_types.hpp	Defines wide tile settings.
src/sycl/kernels/lora/device/chunked_sgmv_lora_expand_fwd_runner.hpp	Launches alignment-aware expand kernels.
src/sycl/kernels/lora/device/chunked_sgmv_lora_expand_fwd_kernel.hpp	Implements gather-fused expand.
src/sycl/kernels/lora/device/chunked_sgmv_lora_expand_fwd_dispatch.hpp	Updates expand launch declarations.
src/sycl/kernels/lora/common/lora_dpas_primitives.hpp	Adds shared DPAS and memory helpers.
src/sycl/kernels/lora/common/lora_debug_checks.hpp	Centralizes conditional value checks.
src/sycl/GateUpLoraBFwd.cpp	Uses conditional validation helpers.
src/sycl/EmbeddingLoraAFwd.cpp	Gates validation; retains launch-sizing sync.
src/sycl/ChunkedSgmvLoraShrinkFwd.cpp	Dispatches fused shrink without sorted copies.
src/sycl/ChunkedSgmvLoraExpandFwd.cpp	Dispatches fused expand without sorted copies.
src/sycl/chunked_sgmv_lora_shrink_fwd_kernel.cpp.in	Updates generated shrink launchers.
src/sycl/chunked_sgmv_lora_expand_fwd_kernel.cpp.in	Updates generated expand launchers.
src/sgl_kernel_debug_config.h.in	Defines generated validation flag.
src/CMakeLists.txt	Adds debug-check build option.
src/ChunkedSgmvLoraShrinkFwdXe20.cmake	Registers split-K tile.
src/ChunkedSgmvLoraExpandFwdXe20.cmake	Registers wide tile.
python/sgl_kernel/lora.py	Documents fused kernels and coalesced segments.
benchmark/bench_chunked_sgmv_lora_shrink_fwd.py	Adds cases and adjusts Triton baseline.
benchmark/bench_chunked_sgmv_lora_expand_fwd.py	Adds cases and adjusts Triton baseline.

🧠 Review effort: Balanced

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Copilot

🟡 Changes recommended

The required embedding launch-bound parameter breaks existing API calls and needs a backward-compatible interface.

2 open findings

1 resolved since last review

Unconditional embedding item() causes device-to-host synchronization

🧠 Review effort: Balanced

pralay-das · 2026-10-22T03:49:59Z

+"""Debug-only value checks of the LoRA kernels (USE_DEBUG_CHECKS=ON builds).
+


why this test file is needed ?

I have kept all of the checks which require device-host copies under DEBUG build mode. This test file tests those cases when the repo is built under DEBUG MODE ON. I placed those debug checks related to all lora kernels in a separate file.

I think for checking data(configs) under debug mode flag is okay, but no need to create test file for that, if the debug mode flag is used in workload enablement then still it is fine, otherwise it will be cluttered, and maintaining in CI will be a burden, I think it need to be removed.

Done. The file is removed as per the recommendation.

pralay-das · 2026-10-23T09:35:37Z

 endif()
+# Delivered as a generated header, not a global -D: only sources that include it
+# (via lora_debug_checks.hpp) recompile when USE_DEBUG_CHECKS is toggled.
+set(SGL_KERNEL_DEBUG_CHECKS ${USE_DEBUG_CHECKS})


Instead of creating new debug flags, could you use RelWithDebInfo, which is currently available and you will enable your checks by simply add SKBUILD_CMAKE_BUILD_TYPE=RelWithDebInfo with your build cmd.

Done. Removed USE_DEBUG_CHECKS and the generated config header. The value checks are now compiled in when CMAKE_BUILD_TYPE is RelWithDebInfo.

airMeng

Requesting changes to fix the SGL_WERROR=ON build regression described in the inline comment.

airMeng

Re-reviewed the latest revision. The earlier embedding API compatibility and license-comment/warnings-as-errors findings are fixed; fresh targeted Release and RelWithDebInfo builds both pass with -Werror. Requesting changes for the newly reproduced default embedding launch-bound regression described inline. The PR's existing 16,384-token, one-token-per-segment test fails in both modes on B60/Xe20 with icpx 2026.1 and Torch 2.14.0+xpu; the same input succeeds with an explicit tight bound. Verification used isolated targeted source builds rather than a complete production wheel.

airMeng · 2026-10-29T20:53:02Z

+template <typename T>
+struct ToCuteElementType;
+template <>
+struct ToCuteElementType<at::BFloat16> {
+  using type = cute::bfloat16_t;
+};
+template <>
+struct ToCuteElementType<at::Half> {
+  using type = cute::half_t;
+};


Could we reuse the existing at::native::xpu::ToCutlassElementType<T> from common/group_gemm_types.hpp:79–93 instead of introducing a second dtype mapper here? The two fp16/bf16 specializations are identical: sycl-tla’s cute/numeric/numeric_types.hpp imports cutlass::half_t and cutlass::bfloat16_t into cute, so ToCuteElementType does not provide different types or behavior. Reusing the existing trait would keep the Torch-to-kernel dtype mapping in one place. If importing the full prefill type-bundle header creates undesirable dependencies, the shared trait can be factored into a lightweight common header.

Done. The two specializations were identical (cute::half_t / bfloat16_t are the cutlass types), so ToCutlassElementType now lives in a small shared header, common/lora_element_type.hpp, used by both the prefill and decode type bundles. I didn't include group_gemm_types.hpp directly to avoid pulling the prefill epilogue/mainloop headers into the decode TUs.

conflicts resolved

…20 extension in file names as these files are not specific to only Xe20

…for the checks

pralay-das

check few comments mentioned below, incorporate those in your next PR, otherwise LGTM.

pralay-das · 2026-10-27T09:38:32Z

+      sgk,
+      n_groups);
+  const int64_t local = 16 * sgk;
+  sycl_kernel_submit(row_tiles * n_groups * local, local, queue, kernel);


use launch<kerenl,grf_size> instead of sycl_kernel_submit

pralay-das · 2026-10-30T09:44:55Z

+  template <class FrgTensorC>
+  CUTLASS_DEVICE static void fixup(Params const&, WorkTileInfo const&, FrgTensorC&, uint32_t, uint32_t) {}


pralay-das · 2026-10-30T09:45:35Z

+  CUTLASS_HOST_DEVICE static uint32_t get_work_k_tile_start(WorkTileInfo const&) {
+    return 0u;
+  }


pralay-das · 2026-10-30T09:57:39Z

+          tSrB(i) = ElementB::bitcast(n < p.N && k < p.K ? p.b[static_cast<int64_t>(n) * p.K + k] : uint16_t(0));
+        }
+        load_a(c);
+        cute::gemm(mma, tSrA, tSrB, accum);


try to use prefetch before gemm to get more performance. though you can incorporate in your next PR.

srujan-thomdapu requested a review from pralay-das October 15, 2026 06:48

srujan-thomdapu self-assigned this Oct 15, 2026

srujan-thomdapu added the run-ci label Oct 15, 2026

srujan-thomdapu requested a review from airMeng as a code owner October 15, 2026 06:48

srujan-thomdapu added the perf Performance-related; enables baseline recording. label Oct 15, 2026

srujan-thomdapu requested a review from kareemshaik80 as a code owner October 15, 2026 06:48

Copilot AI balanced review requested due to automatic review settings October 15, 2026 06:48

Copilot AI reviewed Oct 15, 2026

View reviewed changes

Comment thread src/sycl/EmbeddingLoraAFwd.cpp Resolved

Comment thread tests/test_lora_debug_checks.py Resolved

Copilot started reviewing on behalf of srujan-thomdapu October 15, 2026 07:33

View session

srujan-thomdapu removed the perf Performance-related; enables baseline recording. label Oct 15, 2026

airMeng requested a balanced review from Copilot October 15, 2026 15:49

Copilot started reviewing on behalf of airMeng October 15, 2026 15:50

View session

Copilot AI reviewed Oct 15, 2026

View reviewed changes

Comment thread python/sgl_kernel/lora.py Resolved

srujan-thomdapu mentioned this pull request Oct 16, 2026

[Feat] [Intel] lora: Add intel_xpu LoRA backend sgl-project/sglang#44769

Open

srujan-thomdapu force-pushed the lora_kernels/lora_kernels_perf_opt branch from 0b81027 to 12fd7b7 Compare October 21, 2026 06:19

pralay-das reviewed Oct 22, 2026

View reviewed changes

srujan-thomdapu force-pushed the lora_kernels/lora_kernels_perf_opt branch from 12fd7b7 to 454ee5d Compare October 23, 2026 06:33

pralay-das reviewed Oct 23, 2026

View reviewed changes

srujan-thomdapu force-pushed the lora_kernels/lora_kernels_perf_opt branch from 454ee5d to 25bac18 Compare October 26, 2026 04:51

airMeng requested changes Oct 26, 2026

View reviewed changes

Comment thread src/sycl/kernels/lora/common/lora_debug_checks.hpp Resolved

srujan-thomdapu force-pushed the lora_kernels/lora_kernels_perf_opt branch from 25bac18 to 2a195fb Compare October 27, 2026 05:26

airMeng requested changes Oct 27, 2026

View reviewed changes

Comment thread src/sycl/EmbeddingLoraAFwd.cpp Resolved

srujan-thomdapu force-pushed the lora_kernels/lora_kernels_perf_opt branch 2 times, most recently from 85d2f31 to c9efb9a Compare October 29, 2026 08:25

airMeng approved these changes Oct 29, 2026

View reviewed changes

airMeng reviewed Oct 29, 2026

View reviewed changes

srujan-thomdapu added 4 commits October 30, 2026 07:09

lora kernels are optimized

5edb00e

conflicts resolved

Removed the last sync line in embedding lora a fwd kernel. Removed Xe…

fbab4bf

…20 extension in file names as these files are not specific to only Xe20

#undef is added at the end

0cb9bbf

test_lora_debug_checks.py file is removed

b8bab9a

srujan-thomdapu added 5 commits October 30, 2026 07:09

removed the flag USE_DEBUG_CHECKS and instead RelWithDebInfo is used …

4b88a7f

…for the checks

reviewer comments are addressed

4091e77

Code alignment is done as per the recommondation

9cbc16a

longest segment fallback when max_len param is not passed

04b7049

decode kernels reuse toCutlassElementType as prefill kernels do

519cb51

srujan-thomdapu force-pushed the lora_kernels/lora_kernels_perf_opt branch from c9efb9a to 519cb51 Compare October 30, 2026 07:10

pralay-das approved these changes Oct 30, 2026

View reviewed changes

pralay-das merged commit 3a5da0a into main Oct 30, 2026
6 checks passed

pralay-das deleted the lora_kernels/lora_kernels_perf_opt branch October 30, 2026 10:10

		"""Debug-only value checks of the LoRA kernels (USE_DEBUG_CHECKS=ON builds).

		template <class FrgTensorC>
		CUTLASS_DEVICE static void fixup(Params const&, WorkTileInfo const&, FrgTensorC&, uint32_t, uint32_t) {}

Conversation

srujan-thomdapu commented Oct 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Performance

Behaviour changes

Testing

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

airMeng left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

airMeng left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pralay-das left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

srujan-thomdapu commented Oct 15, 2026 •

edited

Loading

pralay-das left a comment •

edited

Loading