minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Spec][LoRA] Support multi-adapter LoRA with EAGLE/NEXTN/DFLASH/DSPARK speculative decoding (#34337)

Browse Source
This commit is contained in:
Yanbin Jiang
2026-08-21 14:21:53 -07:00
committed by GitHub
parent 590b11a5ef
commit 7d893255c3
23 changed files with 1487 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/lora/backend/base_backend.py
+6 -11
View File
@@ -5,7 +5,11 @@ import triton
import triton.language as tl
from sglang.srt.lora.backend.lmhead_mixing import LoRABackendLmHeadMixing
from sglang.srt.lora.utils import LoRABatchInfo, MoELoRABatchInfo
from sglang.srt.lora.utils import (
LoRABatchInfo,
MoELoRABatchInfo,
get_batch_token_counts,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
@@ -296,16 +300,7 @@ class BaseLoRABackend(LoRABackendLmHeadMixing):
adapter_enabled = None
token_lora_mapping = None
num_tokens = (
sum(forward_batch.extend_seq_lens_cpu)
if forward_batch.forward_mode.is_extend()
else forward_batch.batch_size
)
max_len = (
max(forward_batch.extend_seq_lens_cpu)
if forward_batch.forward_mode.is_extend()
else 1
)
num_tokens, max_len = get_batch_token_counts(forward_batch)
if (
batch_info.req_seg_indptr is not None
python/sglang/srt/lora/backend/chunked_backend.py
+2 -5
View File
@@ -12,6 +12,7 @@ from sglang.srt.lora.backend.base_backend import BaseLoRABackend
from sglang.srt.lora.utils import (
LoRABatchInfo,
generate_sequence_lengths,
get_batch_token_counts,
get_lm_head_pruned_lens,
merge_and_chunk_segments,
)
@@ -184,11 +185,7 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
Returns:
The determined chunk size
"""
num_tokens = (
forward_batch.extend_num_tokens
if forward_batch.forward_mode.is_extend()
else forward_batch.batch_size
)
num_tokens, _ = get_batch_token_counts(forward_batch)
return self._determine_chunk_size_for_tokens(num_tokens)
def _determine_chunk_size_for_tokens(self, num_tokens: int) -> int:
python/sglang/srt/lora/backend/triton_backend.py
+14 -11
View File
@@ -11,6 +11,8 @@ from sglang.kernels.ops.gemm.sgemm_lora_b import sgemm_lora_b_fwd
from sglang.srt.lora.backend.base_backend import BaseLoRABackend
from sglang.srt.lora.utils import (
LoRABatchInfo,
generate_sequence_lengths,
get_batch_token_counts,
get_lm_head_pruned_lens,
merge_and_chunk_segments,
)
@@ -289,6 +291,15 @@ class TritonLoRABackend(BaseLoRABackend):
self.cuda_graph_batch_info is not None
), "CUDA Graph batch info is not initialized."
batch_info = self.cuda_graph_batch_info
if forward_batch.forward_mode.is_target_verify():
# seg_lens were pre-filled at the captured per-request width
# (stored as max_len); another width would silently
# mis-segment adapters onto the wrong token rows.
assert forward_batch.spec_info.draft_token_num == batch_info.max_len, (
"target-verify width "
f"{forward_batch.spec_info.draft_token_num} does not match "
f"the captured LoRA cuda-graph width {batch_info.max_len}"
)
batch_info.bs = forward_batch.batch_size
batch_info.num_segments = forward_batch.batch_size
elif use_prefill_cuda_graph:
@@ -303,17 +314,9 @@ class TritonLoRABackend(BaseLoRABackend):
batch_info.seg_lens[bs:].zero_()
torch.cumsum(batch_info.seg_lens, dim=0, out=batch_info.seg_indptr[1:])
else:
max_len = (
# Calculate max_len from the CPU copy to avoid D2H transfer.
max(forward_batch.extend_seq_lens_cpu)
if forward_batch.forward_mode.is_extend()
else 1
)
seg_lens = (
forward_batch.extend_seq_lens
if forward_batch.forward_mode.is_extend()
else torch.ones(bs, dtype=torch.int32, device=self.device)
)
# max_len comes from the CPU-side counts to avoid a D2H transfer.
_, max_len = get_batch_token_counts(forward_batch)
seg_lens = generate_sequence_lengths(forward_batch, device=self.device)
seg_indptr = torch.zeros((bs + 1,), dtype=torch.int32, device=self.device)
seg_indptr[1:] = torch.cumsum(seg_lens, dim=0)
python/sglang/srt/lora/layers.py
+5
View File
@@ -31,6 +31,11 @@ from sglang.srt.runtime_context import get_parallel
_SGLANG_EXPERIMENTAL_LORA_OPTI = envs.SGLANG_EXPERIMENTAL_LORA_OPTI.get()
def unwrap_lora_layer(module: nn.Module) -> nn.Module:
"""Return the plain module behind a LoRA wrapper, or the module itself."""
return module.base_layer if isinstance(module, BaseLayerWithLoRA) else module
class BaseLayerWithLoRA(nn.Module):
def __init__(
self,
python/sglang/srt/lora/lora_manager.py
+23 -2
View File
@@ -43,6 +43,7 @@ from sglang.srt.lora.utils import (
auto_detect_lora_target_modules,
get_normalized_target_modules,
get_target_module_name,
warn_if_adapter_targets_embeddings,
)
from sglang.srt.managers.io_struct import LoRAUpdateOutput
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
@@ -50,6 +51,7 @@ from sglang.srt.runtime_context import (
get_exec,
get_lora,
get_parallel,
get_spec,
)
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import get_available_gpu_memory, replace_submodule
@@ -105,6 +107,7 @@ class LoRAManager:
self.lora_strict_loading: bool = getattr(
server_args, "lora_strict_loading", False
)
self.speculative_algorithm: Optional[str] = get_spec().speculative_algorithm
# LoRA backend for running sgemm kernels
logger.info(f"Using {lora_backend} as backend of LoRA kernels.")
@@ -782,6 +785,12 @@ class LoRAManager:
)
lora_adapter.initialize_weights()
warn_if_adapter_targets_embeddings(
lora_name=lora_ref.lora_name,
embedding_layer_names=lora_adapter.embedding_layers.keys(),
speculative_algorithm=self.speculative_algorithm,
)
self.loras[lora_ref.lora_id] = lora_adapter
def load_lora_weights_from_tensors(
@@ -799,6 +808,12 @@ class LoRAManager:
base_model=self.base_model,
)
lora_adapter.initialize_weights_from_tensors(tensors)
warn_if_adapter_targets_embeddings(
lora_name=lora_ref.lora_name,
embedding_layer_names=lora_adapter.embedding_layers.keys(),
speculative_algorithm=self.speculative_algorithm,
)
self.loras[lora_ref.lora_id] = lora_adapter
def load_lora_adapter_from_tensors(
@@ -1035,12 +1050,18 @@ def init_lora_cuda_graph_moe_buffers(
from sglang.srt.lora.layers import FusedMoEWithLoRA
max_bs = get_exec().graph.cuda_graph_config.decode.max_bs
# With spec on, the decode graph captures TARGET_VERIFY batches of
# num_draft_tokens per request, and the buffers below are per-token, so
# they must be sized in tokens rather than requests.
max_tokens = max_bs * (get_spec().speculative_num_draft_tokens or 1)
max_loras = get_lora().max_loras_per_batch
for module in model.modules():
if isinstance(module, FusedMoEWithLoRA):
lora_manager.init_cuda_graph_moe_buffers(max_bs, max_loras, dtype, module)
lora_manager.init_cuda_graph_moe_buffers(
max_tokens, max_loras, dtype, module
)
logger.info(
f"Pre-allocated shared MoE LoRA CUDA graph buffers "
f"(max_bs={max_bs}, max_loras={max_loras})"
f"(max_tokens={max_tokens}, max_loras={max_loras})"
)
break
python/sglang/srt/lora/utils.py
+37
View File
@@ -1,3 +1,4 @@
import logging
from dataclasses import dataclass
from enum import Enum
from typing import Iterable, List, Optional, Set, Tuple, Union
@@ -7,6 +8,29 @@ import torch
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.utils.hf_transformers_utils import AutoConfig
logger = logging.getLogger(__name__)
def warn_if_adapter_targets_embeddings(
lora_name: str,
embedding_layer_names: Iterable[str],
speculative_algorithm: Optional[str],
) -> None:
"""Warn once when an adapter carries embedding weights under EAGLE spec."""
if speculative_algorithm not in ("EAGLE", "EAGLE3"):
return
modules = sorted(embedding_layer_names)
if not modules:
return
logger.warning(
"LoRA adapter '%s' targets embedding modules (%s) while EAGLE-family "
"speculative decoding is enabled. The shared draft consumes their "
"base weights, so those deltas do not influence drafting and may "
"reduce the accept rate. Outputs are unaffected.",
lora_name,
", ".join(modules),
)
@dataclass
class MoELoRABatchInfo:
@@ -465,6 +489,19 @@ def get_lm_head_lora_b_shard_size(output_dim: int, shard_indices=None) -> int:
return output_dim
def get_batch_token_counts(forward_batch: ForwardBatch) -> Tuple[int, int]:
"""(total tokens, max tokens per request) for LoRA segment math."""
mode = forward_batch.forward_mode
if mode.is_decode():
return forward_batch.batch_size, 1
if mode.is_target_verify():
num_tokens_per_req = forward_batch.spec_info.draft_token_num
return forward_batch.batch_size * num_tokens_per_req, num_tokens_per_req
if mode.is_extend():
return forward_batch.extend_num_tokens, max(forward_batch.extend_seq_lens_cpu)
raise ValueError(f"Unsupported forward mode: {mode}")
def generate_sequence_lengths(
forward_batch: ForwardBatch, device: Optional[torch.device] = None
) -> torch.Tensor:
python/sglang/srt/model_executor/forward_batch_info.py
+3 -3
View File
@@ -845,7 +845,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
if ret.forward_mode.is_idle():
ret.positions = torch.empty((0,), dtype=torch.int64, device=device)
if model_runner.server_args.enable_lora:
if model_runner.lora_manager is not None:
model_runner.lora_manager.reset_lora_batch()
return ret
@@ -919,8 +919,8 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
else:
ret._compute_mrope_positions(model_runner, batch)
# Init lora information
if model_runner.server_args.enable_lora:
# Init lora information (None on a draft runner: it is unadapted)
if model_runner.lora_manager is not None:
# In the non-LoRA overlap loading case, we fetch LoRA adapters into the memory pool
# as a batch, right before running the batch
if not model_runner.server_args.enable_lora_overlap_loading:
python/sglang/srt/model_executor/model_runner.py
+5 -1
View File
@@ -319,6 +319,9 @@ class ModelRunner:
# earlier publish.
if not is_draft_worker:
set_global_server_args_for_scheduler(server_args)
# Set by maybe_init_lora_manager; stays None when LoRA is off and on
# draft runners, which serve adapters' target model unadapted.
self.lora_manager: Optional[LoRAManager] = None
self.draft_attention_backend = resolve_draft_attention_backend(
draft_attention_backend=draft_attention_backend,
server_args=server_args,
@@ -751,7 +754,8 @@ class ModelRunner:
self.apply_torch_tp()
def maybe_init_lora_manager(self):
if get_lora().enable_lora:
# Adapters apply to the target model only; the draft runs unadapted.
if get_lora().enable_lora and not self.is_draft_worker:
self.init_lora_manager()
def maybe_enable_batch_invariant_mode(self):
python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py
+1 -1
View File
@@ -271,7 +271,7 @@ def capture_prefill_graph(
return result(eager_runner)
if (
model_runner.server_args.enable_lora
model_runner.lora_manager is not None
and not model_runner.lora_manager.supports_prefill_cuda_graph
):
logger.warning(
python/sglang/srt/model_executor/runner/base_runner.py
+1 -1
View File
@@ -574,7 +574,7 @@ class BaseRunner(ABC):
)
# Optional LoRA metadata.
if mr.server_args.enable_lora:
if mr.lora_manager is not None:
lora_ids = [None] * batch_size
else:
lora_ids = None
python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py
+3 -3
View File
@@ -348,7 +348,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
self._captured_ragged_layouts: dict[int, object] = {}
if self.ragged_verify_mode and (
self.enable_two_batch_overlap
or model_runner.server_args.enable_lora
or model_runner.lora_manager is not None
or self.disable_padding
):
raise ValueError(
@@ -381,7 +381,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
if self.enable_torch_compile:
set_torch_compile_config()
if self.model_runner.server_args.enable_lora:
if self.model_runner.lora_manager is not None:
# Phase 2 of LoRA CUDA graph init: dense LoRA batch metadata.
# Phase 1 (MoE buffers) was handled earlier in ModelRunner via
# lora_manager.init_cuda_graph_moe_buffers().
@@ -927,7 +927,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
spec_info,
)
if self.model_runner.server_args.enable_lora:
if self.model_runner.lora_manager is not None:
# It is safe to capture CUDA graph using empty LoRA id, as the LoRA kernels will always be launched whenever
# `--enable-lora` is set to True (and return immediately if the LoRA id is empty for perf optimization).
lora_ids = [None] * bs
python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py
+1 -1
View File
@@ -256,7 +256,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
# --- model flags ----------------------------------------------
self.quant_config = getattr(model_runner.model, "quant_config", None)
self.is_multimodal = model_runner.model_config.is_multimodal
self.enable_lora = model_runner.server_args.enable_lora
self.enable_lora = model_runner.lora_manager is not None
# Classification/reward forwards branch on return_pooled_hidden_states;
# capture must use the same flag value as replay for those models.
self.capture_return_pooled_hidden_states = not model_runner.is_generation
python/sglang/srt/server_args.py
+64 -4
View File
@@ -319,6 +319,10 @@ RETRACTION_POLICY_CHOICES = ["length", "priority"]
RL_ON_POLICY_TARGET_CHOICES = ["fsdp"]
# Speculative algorithms whose verify forward presents a uniform per-request
# token width, which is what the LoRA segment layout assumes.
_LORA_SPEC_ALGORITHMS = ("EAGLE", "EAGLE3", "DFLASH", "DSPARK")
LORA_BACKEND_CHOICES = ["triton", "csgmv", "ascend", "torch_native"]
ENCODER_TRANSFER_BACKEND_CHOICES = [
@@ -9451,10 +9455,7 @@ class ServerArgs:
)
# Validate compatibility with speculative decoding
if self.speculative_algorithm not in ["NGRAM", None]:
raise ValueError(
"Currently LoRA is only compatible with NGRAM speculative decoding."
)
self._check_lora_speculative_compatibility()
# Parse lora_paths
if isinstance(self.lora_paths, list):
@@ -9559,6 +9560,65 @@ class ServerArgs:
self.lora_drain_wait_threshold >= 0.0
), "--lora-drain-wait-threshold must be non-negative."
def _check_lora_speculative_compatibility(self):
"""Validate LoRA + speculative decoding combinations.
Adapters apply to the target only; a shared draft runs unadapted.
Matches resolved algorithm names (NEXTN has collapsed to EAGLE).
"""
if self.speculative_algorithm in ["NGRAM", None]:
return
if self.speculative_algorithm not in _LORA_SPEC_ALGORITHMS:
promoted = (
" (NEXTN/EAGLE with a Gemma4 assistant draft is automatically "
"promoted to FROZEN_KV_MTP, which does not support LoRA)"
if self.speculative_algorithm == "FROZEN_KV_MTP"
else ""
)
raise ValueError(
"LoRA is only compatible with NGRAM, EAGLE, NEXTN, EAGLE3, "
"DFLASH, or DSPARK speculative decoding, not "
f"{self.speculative_algorithm}{promoted}."
)
ragged_mode = envs.SGLANG_RAGGED_VERIFY_MODE.get()
# Each entry: (is unsupported, why). Reasons are appended to a shared
# prefix so the message names the combination, not just the flag.
unsupported = [
(
self.speculative_algorithm == "DSPARK" and ragged_mode != "static",
f"does not support SGLANG_RAGGED_VERIFY_MODE={ragged_mode!r}: "
"the per-request verify lengths it schedules break the "
"uniform-width LoRA segment layout",
),
(
self.speculative_adaptive,
"does not support --speculative-adaptive: the draft is built "
"from a static ServerArgs snapshot, and the runtime-state "
"swap does not rebuild LoRA cuda-graph metadata",
),
(
"experimental_sgl_trtllm"
in (self.moe_runner_backend, self.speculative_moe_runner_backend),
"does not support the experimental_sgl_trtllm MoE runner: its "
"TopK reads the LoRA config per forward, which the draft "
"resolves against the target's after its own publish ended",
),
(
envs.SGLANG_ENABLE_OVERLAP_PLAN_STREAM.get(),
"does not support SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1: LoRA "
"batch preparation would run on the plan stream, unordered "
"against in-flight forwards",
),
]
for is_unsupported, reason in unsupported:
if is_unsupported:
raise ValueError(
f"LoRA with EAGLE/NEXTN/EAGLE3 speculative decoding {reason}."
)
def validate_buckets_rule(self, arg_name: str, buckets_rule: List[str]):
if not buckets_rule:
return
python/sglang/srt/speculative/dflash_worker_v2.py
+4 -3
View File
@@ -22,6 +22,7 @@ from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.environ import envs
from sglang.srt.layers.logits_processor import should_apply_lm_head_quant_method
from sglang.srt.layers.logprob_processor import compute_spec_logprobs
from sglang.srt.lora.layers import unwrap_lora_layer
from sglang.srt.managers.schedule_batch import ScheduleBatch
from sglang.srt.managers.scheduler import GenerationBatchResult
from sglang.srt.managers.tp_worker import TpModelWorker
@@ -480,7 +481,7 @@ class DFlashWorkerV2(BaseSpecWorker):
if self.block_size <= 1:
return _eager("block_size<=1")
target_model = self._target_worker.model_runner.model
lm_head = getattr(target_model, "lm_head", None)
lm_head = unwrap_lora_layer(getattr(target_model, "lm_head", None))
if lm_head is None:
return _eager("no target lm_head")
@@ -1749,8 +1750,8 @@ class DFlashWorkerV2(BaseSpecWorker):
# --- 1) Draft a fixed block with the draft model.
target_model = self.target_worker.model_runner.model
embed_module = target_model.get_input_embeddings()
lm_head = getattr(target_model, "lm_head", None)
embed_module = unwrap_lora_layer(target_model.get_input_embeddings())
lm_head = unwrap_lora_layer(getattr(target_model, "lm_head", None))
if lm_head is None or not (
hasattr(lm_head, "weight")
or callable(getattr(getattr(lm_head, "quant_method", None), "apply", None))
python/sglang/srt/speculative/dspark_components/dspark_draft.py
+2 -1
View File
@@ -11,6 +11,7 @@ from sglang.kernels.ops.speculative.dspark.dspark_draft_model import (
SampleStepTokens,
)
from sglang.srt.environ import envs
from sglang.srt.lora.layers import unwrap_lora_layer
from sglang.srt.managers.schedule_batch import ScheduleBatch
from sglang.srt.model_executor.forward_batch_info import (
CaptureHiddenMode,
@@ -201,7 +202,7 @@ class DraftBlockProposer:
target_model,
sampling_info,
) -> DraftProposal:
embed_module = target_model.get_input_embeddings()
embed_module = unwrap_lora_layer(target_model.get_input_embeddings())
draft_sampler = self._draft_sampler
all_greedy = sampling_info is None or sampling_info.is_all_greedy
fwd = self._run_forward(
python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py
+5 -2
View File
@@ -12,6 +12,7 @@ from sglang.srt.configs.hybrid_arch import mambaish_config
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.environ import envs
from sglang.srt.layers.logprob_processor import compute_spec_logprobs
from sglang.srt.lora.layers import unwrap_lora_layer
from sglang.srt.managers.schedule_batch import ScheduleBatch
from sglang.srt.managers.scheduler import GenerationBatchResult
from sglang.srt.managers.tp_worker import TpModelWorker
@@ -190,13 +191,15 @@ class DSparkWorkerV2(BaseSpecWorker):
)
else:
target_model = self.target_worker.model_runner.model
lm_head = getattr(target_model, "lm_head", None)
lm_head = unwrap_lora_layer(getattr(target_model, "lm_head", None))
if lm_head is None or not hasattr(lm_head, "weight"):
raise RuntimeError(
"DSpark requires the target model to expose `lm_head` with `weight`."
)
self.draft_model.attach_shared_modules(
embed_tokens=self._resolve_target_embed_tokens(target_model),
embed_tokens=unwrap_lora_layer(
self._resolve_target_embed_tokens(target_model)
),
lm_head=lm_head,
)
python/sglang/srt/speculative/eagle_worker_v2.py
+5 -1
View File
@@ -278,8 +278,12 @@ class EagleDraftWorker(EagleDraftWorkerBase):
self.hot_token_id = None
def init_lm_head(self):
from sglang.srt.lora.layers import unwrap_lora_layer
embed, head = self.target_worker.model_runner.model.get_embed_and_head()
target_lm_head = getattr(self.target_worker.model_runner.model, "lm_head", None)
target_lm_head = unwrap_lora_layer(
getattr(self.target_worker.model_runner.model, "lm_head", None)
)
def maybe_share_target_lm_head():
if (