Repository navigation
[Feat] [Intel] lora: Add intel_xpu LoRA backend - #44769
srujan-thomdapu wants to merge 5 commits into
Conversation
This comment has been minimized.
This comment has been minimized.
| ) | ||
|
|
||
| # Prefill (contiguous): XPU B-expand. | ||
| return xpu_qkv_lora_b_fwd( |
There was a problem hiding this comment.
These fused B kernels, and the decode chunked expand, read slice p at column p * max_rank and ignore lora_ranks, but LoRAMemoryPool packs a rank-r adapter's stacked A at stride r. Any adapter below --max-lora-rank gets wrong k, v and up deltas; triton masks by rank. Rank-aware kernels, or a max_rank-spaced pool layout?
There was a problem hiding this comment.
Good catch. Fixed with a max_rank-spaced pool layout rather than rank-aware kernels. The XPU kernels deliberately contract over the full max_rank, so the backend now declares uses_max_rank_spaced_lora_layout and the pool places each stacked A slice at p * max_lora_rank, matching what MoE A already does. This also fixed a related issue: the XPU kernels need zeroed rank tails, but the pool didn't zero them on LRU slot reuse, so the slot is now cleared before load. Triton keeps the packed layout. Added a regression test in test_mem_pool_ep_unit.py.
| num_slices=1, | ||
| base_output=base_output, | ||
| ) | ||
| return xpu_sgemm_lora_b_fwd( |
There was a problem hiding this comment.
These kernels also need A rows and B columns past the adapter's rank zeroed, but slot eviction reloads only the first r, so a rank-8 adapter in a slot a rank-16 one used keeps the old tail and part of its delta leaks in. The pool could zero the B tail on load for this backend, like the SGLANG_EXPERIMENTAL_LORA_OPTI path does.
There was a problem hiding this comment.
Same change addressed in the previous comment covers this too.
| return xpu_qkv_lora_b_fwd( | ||
| input_x=lora_a_output, | ||
| qkv_lora_b=qkv_lora_b, | ||
| output_offset=output_offset, |
There was a problem hiding this comment.
n_slices reaches the shrink and the decode expand but not this call, and qkv_lora_b_fwd is hard-wired to three slices, so DeepSeek MLA fused_qkv_a_proj_with_mqa (n_slices=2) fails prefill with "output_offset must have 4 elements". Could prefill use chunked_sgmv_lora_expand_fwd with permutation=None?
There was a problem hiding this comment.
Yes. Prefill now routes to chunked_sgmv_lora_expand_fwd with permutation=None (its contiguous fast path) whenever n_slices != 3, so MLA fused_qkv_a_proj_with_mqa gets the 2-slice expand with its own output_offset. The standard 3-slice q/k/v prefill keeps the CUTLASS qkv_lora_b_fwd. Both kernels derive the segment count from seg_indptr, so the prefill CUDA graph is unaffected.
| input_ids=input_ids, | ||
| weights=weights, | ||
| vocab_size=vocab_size, | ||
| max_len=bi.max_len, |
There was a problem hiding this comment.
embedding_lora_a_fwd in sglang-kernel-xpu 0.4.0, the wheel pinned in pyproject_xpu.toml, has no max_len, so any adapter on embed_tokens raises TypeError here. That and the per-call .item() checks that break XPU graph capture are only fixed in unmerged sgl-kernel-xpu#550, so the pin bump has to land with this PR.
There was a problem hiding this comment.
Agreed. This PR depends on sgl-kernel-xpu#550. Once it merges, I will bump the pyproject_xpu.toml pin in this PR.
| permutation=torch.zeros(max_tokens, dtype=torch.int32), | ||
| ) | ||
|
|
||
| def init_prefill_cuda_graph_batch_info(self, max_num_tokens: int): |
There was a problem hiding this comment.
LoRAManager.init_prefill_cuda_graph_batch_info passes max_num_requests= as a keyword, so --cuda-graph-backend-prefill breakable or full with this backend dies with TypeError at capture. Accept it, and probably size the slots from it like triton does instead of the fixed 32.
There was a problem hiding this comment.
Fixed. init_prefill_cuda_graph_batch_info now takes max_num_requests and sizes the slots from it like triton (max_num_tokens when None); the fixed 32-slot constant is gone, so the graph covers the captured request capacity instead of falling back to eager prefill past 32 requests.
| else: | ||
| max_len = ( | ||
| # Calculate max_len from the CPU copy to avoid D2H transfer. | ||
| max(forward_batch.extend_seq_lens_cpu) |
There was a problem hiding this comment.
TARGET_VERIFY counts as is_extend() but leaves extend_seq_lens_cpu None, so eager verify with EAGLE plus LoRA (bs over the graph cap, or graphs off) crashes on max(None). Triton uses get_batch_token_counts and generate_sequence_lengths; this is an older copy of its prepare_lora_batch, so subclassing it fixes that.
There was a problem hiding this comment.
Done. IntelXpuLoRABackend now subclasses TritonLoRABackend.
| ) | ||
|
|
||
| DENSE_BACKENDS = [*BACKENDS, "torch_native"] | ||
| DENSE_BACKENDS = [*BACKENDS, "torch_native", "intel_xpu"] |
There was a problem hiding this comment.
DENSE_BACKENDS only feeds the accuracy and batching tests, which start with rank-8 adapters, so max_lora_rank equals the rank and nothing is evicted. The embedding, eviction and dynamic-load tests run rank 8 under max 16 and would hit the stride and max_len failures; can intel_xpu get variants there like torch_native?
There was a problem hiding this comment.
Added intel_xpu variants for dynamic load, eviction (lru/fifo) and embedding. The eviction and embedding cases inferred max rank = adapter rank, and the embedding test targets q/v rather than embed_tokens, so on their own they wouldn't have hit either failure. To cover them directly I added test_lora_below_max_rank_matches_hf (rank 8 under max 16 vs HF, all dense backends; catches the stride bug) and test_eviction_mixed_rank_intel_xpu (rank-8 reload into a slot a rank-16 adapter used; catches stale tails).
c60cc7c to
c5bb3e3
Compare
This comment has been minimized.
This comment has been minimized.
dbf5196 to
eee0377
Compare
[Feat] [Intel] lora: Add intel_xpu LoRA backend
Motivation
On Intel XPU, multi-LoRA serving currently has to use the generic
triton,csgmv, ortorch_nativebackends. This PR adds anintel_xpuLoRA backend that runs the native Intel XPU LoRA kernels from sgl-kernel-xpu on both the prefill and decode paths.Usage:
python -m sglang.launch_server \ --model-path <base-model> \ --device xpu \ --enable-lora \ --lora-paths <name>=<adapter-path> \ --lora-backend intel_xpuDependency: sgl-kernel-xpu
The backend imports these ops from
sgl_kernel(sgl-kernel-xpu):embedding_lora_a_fwdsgemm_lora_a_fwdsgemm_lora_b_fwdqkv_lora_b_fwdgate_up_lora_b_fwdchunked_sgmv_lora_shrink_fwdchunked_sgmv_lora_expand_fwdThese kernels especially decode kernels are optimized. To get the best performance, build sgl-kernel-xpu from the PR below:
This PR should merge after the above PR got merged in sgl-kernel-xpu.
Modifications
New backend:
python/sglang/srt/lora/backend/intel_xpu_backend.py(IntelXpuLoRABackend)permutation is None): per-shape XPU kernels (sgemm_lora_a_fwd,sgemm_lora_b_fwd,qkv_lora_b_fwd,gate_up_lora_b_fwd).compute_sgemm_routing), and the chunked-SGMV shrink/expand pair are gather fused DPAS kernels. The shrink output layout(num_tokens, num_slices * max_rank)is exactly the expand's input, so decode stays on XPU from A to B. The expand fuses per-adapter scaling and thebase_outputresidual.supports_prefill_cuda_graph = True, 32 fixed prefill segment slots). Batches with more requests fall back to eager prefill.LoRABatchInfolayout.Registration and CLI
lora/backend/lora_registry.py: registerintel_xpu.arg_groups/choices.py: addLORA_BACKEND_CHOICES(includingintel_xpu).arg_groups/fields/lora.pynow reads--lora-backendchoices from it instead of a hardcoded list, andserver_args.pyre-exports it with the other choice lists.Validation
arg_groups/validation_hook.py:validate_device_lora_backendrejects--lora-backend intel_xpuwithout--device xpuat argument-check time. Without it, the error only appears later as an import failure insideLoRAManagerinit. This mirrors the existing--sampling-backend intel_xpucheck.Tests
test/registered/unit/server_args/test_server_args.py:TestIntelXpuLoRABackend(CPU CI) checks that the backend is registered and selectable from the CLI, rejected off XPU, and accepted on XPU.test/manual/lora/test_lora_xpu_basics.py: addintel_xputoDENSE_BACKENDS, so it runs the HF+PEFT logprob comparison and the multi-LoRA batching / drainer routing tests next totriton,csgmv, andtorch_native.Accuracy Tests
Models and adapters are the ones defined for
test/manual/lora/test_lora_backend.py. Each case runs 4 engines (SRT+LoRA, SRT base, HF+LoRA, HF base) with--lora-backend intel_xpuand compares top input/output logprobs against HF plus ROUGE-L of the decoded text. fp16,max_new_tokens=32, 2 prompts per case (values are prompt 0 / prompt 1); limits: logprob diff< 1e-1, ROUGE-L1.0.Speed Tests and Profiling
Llama-3.3-70B-Instruct
Setup
meta-llama/Llama-3.3-70B-Instruct,--tp-size 8on 8 Intel XPU cards,--dtype bfloat16,--attention-backend triton,--disable-radix-cache, graph capture on.adamkarvonen/llama3.3_70b_shuffled_e3_v2,jeongseokoh/llama3.3_70b_sft,jeongseokoh/llama3.3_70b_speed60_bos,vineetdaniels/llama33_70b_icd_cpt),--max-loras-per-batch 4, adapters drawn uniformly per request.baseis the same model served with LoRA disabled.--random-range-ratio 1.0,ignore_eoson, so every request emits exactlyoutput_lentokens), request rate inf, fixed concurrency. Every run completed all requests with exactlynum_requests * output_lengenerated tokens.prefill_4k_4lora: 512 in / 8 out, 128 requests, concurrency 8,
--chunked-prefill-size 4096,--max-running-requests 8csgmvtritonintel_xpu*intel_xpuvscsgmv: total throughput +13.1%, mean TTFT -18.0%.decode_128_4lora: 16 in / 64 out, 3000 requests, concurrency 128,
--max-running-requests 128csgmvtritonintel_xpu*intel_xpuvscsgmv: total throughput +17.5%, mean TPOT -24.0%.mixed_1k_4lora: 1024 in / 128 out, 400 requests, concurrency 20,
--chunked-prefill-size 4096,--max-running-requests 20csgmvtritonintel_xpuintel_xpuvscsgmv: total throughput +22.0%, mean TPOT -21.2%, mean TTFT -11.1%.Cost of attaching adapters. Each row is the LoRA run relative to the same base model served with no adapters, so a negative throughput delta is the expected price of LoRA.
csgmvtritonintel_xpucsgmvtritonintel_xpucsgmvtritonintel_xpuIn
prefill_4k_4loraeach request emits only 8 tokens, so mean TPOT is dominated by decode steps queued behind other requests' prefill chunks (the base run shows the same effect: P99 ITL 1888 ms); total throughput and TTFT are the meaningful metrics there.In
decode_128_4lora, mean TTFT is higher withintel_xputhan withcsgmv(2299 vs 1579 ms, +46 %). Each request has a 16-token prompt, so TTFT there is mostly time spent queueing for a slot in the full decode batch rather than prefill compute (base TTFT also grows with concurrency), and it is 23 % ofintel_xpu's end-to-end latency. Mean E2E latency is still 15 % lower thancsgmv's (10140 vs 11895 ms); total throughput, TPOT, and E2E latency are the meaningful metrics for decode scenarios.The base (no-LoRA) path is unchanged by this PR: it only adds the
intel_xpuLoRA backend, its--lora-backendchoice, and the--device xpucheck.Llama-3.1-8B-Instruct
Setup
meta-llama/Llama-3.1-8B-Instruct,--tp-size 8on 8 Intel XPU cards,--dtype bfloat16,--attention-backend triton,--disable-radix-cache, graph capture on.adamkarvonen/llama3_1_8b_on_llama3_1_8b_counterfactual_24k_e1_kl0,dsfsi/llama_3.1_8b_instruct-lora-r64-eng-{amh,hau,lug,yor,zul},r-three/lora_baseline_10_lr3e-4_step100_rank64_race_high,r-three/lora_baseline_lr3e-4_step400_rank64_qa_wikidata),--max-loras-per-batch 8, adapters drawn uniformly per request.baseis the same model served with LoRA disabled.ignore_eos, request rate inf, fixed concurrency). Prefill scenarios: long input, 8 output tokens, chunked-prefill size grown with the input. Decode scenarios: 16 in / 96 out, decode batch = concurrency,--chunked-prefill-size 8192 --max-prefill-tokens 16384. Every run completed all requests with exactlynum_requests * output_lengenerated tokens.Total throughput (tok/s)
csgmvtritonintel_xpuintel_xpuvscsgmvintel_xpuvstritonCost of attaching adapters. Each row is the LoRA run relative to the same base model served with no adapters.
csgmvtritonintel_xpucsgmvtritonintel_xpucsgmvtritonintel_xpucsgmvtritonintel_xpucsgmvtritonintel_xpucsgmvtritonintel_xpucsgmvtritonintel_xpucsgmvtritonintel_xpucsgmvtritonintel_xpucsgmvtritonintel_xpucsgmvtritonintel_xpucsgmvtritonintel_xpuThe prefill rows emit only 8 tokens per request, so mean TPOT there is dominated by decode steps queued behind prefill chunks; total throughput and TTFT are the meaningful metrics.
In the decode rows, mean TTFT is higher with
intel_xputhan withcsgmv(+25 % to +36 %). Each request has a 16-token prompt, so TTFT there is mostly time spent queueing for a slot in the full decode batch rather than prefill compute (base TTFT also grows with concurrency, 133 ms at 64 to 1968 ms at 2048), and it is 7 % to 14 % ofintel_xpu's end-to-end latency. Mean E2E latency is still 14 % to 17 % lower thancsgmv's at every concurrency; total throughput, TPOT, and E2E latency are the meaningful metrics for decode scenarios.Llama-3.1-8B-Instruct: full metrics per scenario
prefill_4k_8lora: 512 in / 8 out, 512 requests, concurrency 16,
--max-running-requests 16,--chunked-prefill-size 4096csgmvtritonintel_xpuprefill_8k_8lora: 1024 in / 8 out, 512 requests, concurrency 16,
--max-running-requests 16,--chunked-prefill-size 8192csgmvtritonintel_xpuprefill_12k_8lora: 1536 in / 8 out, 512 requests, concurrency 16,
--max-running-requests 16,--chunked-prefill-size 12288csgmvtritonintel_xpuprefill_16k_8lora: 2048 in / 8 out, 512 requests, concurrency 16,
--max-running-requests 16,--chunked-prefill-size 16384csgmvtritonintel_xpuprefill_24k_8lora: 3072 in / 8 out, 512 requests, concurrency 24,
--max-running-requests 24,--chunked-prefill-size 24576csgmvtritonintel_xpuprefill_65k_8lora: 4096 in / 8 out, 384 requests, concurrency 32,
--max-running-requests 32,--chunked-prefill-size 65536csgmvtritonintel_xpudecode_64_8lora: 16 in / 96 out, 6000 requests, concurrency 64,
--max-running-requests 64csgmvtritonintel_xpudecode_128_8lora: 16 in / 96 out, 6000 requests, concurrency 128,
--max-running-requests 128csgmvtritonintel_xpudecode_256_8lora: 16 in / 96 out, 6000 requests, concurrency 256,
--max-running-requests 256csgmvtritonintel_xpudecode_512_8lora: 16 in / 96 out, 4000 requests, concurrency 512,
--max-running-requests 512csgmvtritonintel_xpudecode_1024_8lora: 16 in / 96 out, 6000 requests, concurrency 1024,
--max-running-requests 1024csgmvtritonintel_xpudecode_2048_8lora: 16 in / 96 out, 8000 requests, concurrency 2048,
--max-running-requests 2048csgmvtritonintel_xpuCI States
Latest PR Test (Base): ❌ Run #40483093338
Latest PR Test (Extra): ❌ Run #40483093556
Latest PR Test (AMD ROCm 10): ❌ Run #40483093157