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

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
+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
@@ -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:
@@ -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)
+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,
+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
+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:
@@ -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:
@@ -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):
@@ -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(
@@ -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
@@ -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
@@ -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
+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
@@ -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))
@@ -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(
@@ -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,
)
@@ -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 (
@@ -0,0 +1,125 @@
"""Does a spec-on vs spec-off divergence need LoRA at all?
Run this when the spec+LoRA harness
(``test/manual/lora/run_spec_lora_matrix.py``) reports a mismatch that
survives its filters. It runs the *same* speculative config with **no
adapters loaded**, comparing spec-off against spec-on greedy outputs and
sampling each side twice so run-to-run instability is visible.
If the base model diverges on the same prompt, the cause is the model and
the speculative path (EAGLE3 topk>1 tree verify in particular), not the LoRA
integration -- which is exactly what it showed for
Qwen3-30B-A3B-Instruct-2507 on prompt #2: DIFFERS, and UNSTABLE-self, with
zero adapters involved.
The harness's own filters cannot answer this: they can tell whether a pair
reproduces within one server config, but not whether the divergence depends
on LoRA being present. That needs this second config.
Usage (from a checkout, with a GPU):
python test/manual/lora/check_spec_baseline_divergence.py
"""
import os
import subprocess
import sys
import time
import requests
PROMPTS = [
"What is the capital of France? Answer in one sentence.",
"List three primary colors.",
"Write a one-sentence story about a brave detective on Mars.",
"Explain what a hash table is in two sentences.",
]
BASE = "http://127.0.0.1:31000"
COMMON = [
"--tp",
"4",
"--moe-runner-backend",
"triton",
"--attention-backend",
"flashinfer",
"--prefill-attention-backend",
"fa4",
"--decode-attention-backend",
"fa4",
"--mem-fraction-static",
"0.8",
]
SPEC = [
"--speculative-algorithm",
"EAGLE3",
"--speculative-draft-model-path",
"lmsys/SGLang-EAGLE3-Qwen3-30B-A3B-Instruct-2507-SpecForge-Nex",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"4",
"--speculative-num-draft-tokens",
"8",
]
def launch(extra):
env = dict(os.environ, SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN="1")
p = subprocess.Popen(
[
sys.executable,
"-m",
"sglang.launch_server",
"--model-path",
"Qwen/Qwen3-30B-A3B-Instruct-2507",
"--port",
"31000",
]
+ COMMON
+ extra,
stdout=open("/scratch/loraspec/logs/attrib_server.log", "a"),
stderr=subprocess.STDOUT,
env=env,
)
for _ in range(120):
time.sleep(10)
try:
if requests.get(BASE + "/health", timeout=3).ok:
return p
except Exception:
pass
raise SystemExit("server did not come up")
def gen():
r = requests.post(
BASE + "/generate",
json={
"text": PROMPTS,
"sampling_params": {"temperature": 0, "max_new_tokens": 32},
},
timeout=900,
)
r.raise_for_status()
return [x["text"] for x in r.json()]
results = {}
for label, extra in [("nospec", []), ("spec", SPEC)]:
p = launch(extra)
try:
results[label] = [gen(), gen()]
finally:
p.terminate()
p.wait(timeout=120)
time.sleep(10)
print("=" * 70)
for i, prompt in enumerate(PROMPTS):
a1, a2 = results["nospec"][0][i], results["nospec"][1][i]
b1, b2 = results["spec"][0][i], results["spec"][1][i]
stable = "stable" if (a1 == a2 and b1 == b2) else "UNSTABLE-self"
verdict = "same" if a1 == b1 else "DIFFERS"
print(f"prompt#{i}: nospec-vs-spec={verdict} ({stable})")
if a1 != b1:
print(f" nospec: {a1!r}")
print(f" spec : {b1!r}")
+876
View File
@@ -0,0 +1,876 @@
"""Manual validation harness: multi-adapter LoRA + EAGLE-family speculative decoding.
The oracle is spec-on vs spec-off *per adapter*, not adapter vs base.
Speculative decoding is lossless with respect to the target it verifies
against, so for greedy sampling:
output(adapter=X, spec=ON) == output(adapter=X, spec=OFF)
Comparing an adapter's output to the base model's only proves LoRA was
applied at all. This harness launches the same weights twice (reference
without spec, then with spec) and runs four checks:
1. parity per adapter, spec-on output == spec-off output (the oracle)
2. distinct each adapter's output differs from base (LoRA really applied)
3. mixed a batch interleaving every adapter matches the solo outputs
(crossed verify segments serve a request the wrong adapter)
4. eager a batch wider than --cuda-graph-max-bs still matches
(exercises the non-cuda-graph target-verify path)
It also reports per-adapter accept length from each response's
meta_info["spec_accept_length"], which is how much speculation the shared
unadapted draft actually buys for that adapter.
Usage:
python test/manual/lora/run_spec_lora_matrix.py --config eagle3-llama31
python test/manual/lora/run_spec_lora_matrix.py --config nextn-qwen35-35b-a3b
python test/manual/lora/run_spec_lora_matrix.py --list
Add --keep-derived to reuse a previously written derived adapter, and
--max-new-tokens / --port to adjust the run.
"""
import argparse
import hashlib
import json
import os
import random
import shutil
import sys
import time
from concurrent import futures
import requests
from sglang.srt.utils import kill_process_tree
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
popen_launch_server,
)
PROMPTS = [
"What is the capital of France? Answer in one sentence.",
"List three primary colors.",
"Write a one-sentence story about a brave detective on Mars.",
"Explain what a hash table is in two sentences.",
]
DERIVED_SUFFIX = "-derived"
# Each config is a runnable server shape. `adapters` are (name, hf_path)
# pairs; `derive_from` names an adapter to synthesize a second, distinct
# adapter from when the base has only one public adapter (needed to make
# mixed-adapter batches meaningful).
CONFIGS = {
# Cheapest: 5 genuinely distinct public adapters, ranks 8-64, no synthesis.
"eagle3-llama31": dict(
model="meta-llama/Llama-3.1-8B-Instruct",
spec_args=[
"--speculative-algorithm=EAGLE3",
"--speculative-draft-model-path=lmsys/sglang-EAGLE3-LLaMA3.1-Instruct-8B",
"--speculative-num-steps=5",
"--speculative-eagle-topk=8",
"--speculative-num-draft-tokens=32",
],
adapters=[
("fact", "algoprog/fact-generation-llama-3.1-8b-instruct-lora"),
("guard", "nvidia/llama-3.1-nemoguard-8b-topic-control"),
("sql", "philschmid/code-llama-3-1-8b-text-to-sql-lora"),
("ocr", "pbevan11/llama-3.1-8b-ocr-correction"),
("zh", "faridlazuarda/valadapt-llama-3.1-8B-it-chinese"),
],
common_args=["--mem-fraction-static=0.7", "--max-lora-rank=64"],
tp=1,
),
# EAGLE topk=1 is the same chain-shaped drafting NEXTN/MTP uses, on 1 GPU
# without needing an MTP checkpoint. cuda-graph-max-bs forces the eager path.
"eagle-topk1-llama2": dict(
model="meta-llama/Llama-2-7b-chat-hf",
spec_args=[
"--speculative-algorithm=EAGLE",
"--speculative-draft-model-path=lmsys/sglang-EAGLE-llama2-chat-7B",
"--speculative-num-steps=3",
"--speculative-eagle-topk=1",
"--speculative-num-draft-tokens=4",
],
adapters=[("norwegian", "RuterNorway/Llama-2-7b-chat-norwegian-LoRa")],
derive_from="norwegian",
common_args=[
"--mem-fraction-static=0.7",
"--max-lora-rank=128",
"--cuda-graph-max-bs=2",
],
tp=1,
),
# Real NEXTN (self-bundled MTP head, no draft path) on a MoE base: the
# densest coverage per GPU-hour — MoE-LoRA cuda-graph buffers, virtual
# experts, and an arch whose MTP head shares the target lm_head module.
"nextn-qwen35-35b-a3b": dict(
model="Qwen/Qwen3.5-35B-A3B",
spec_args=[
"--speculative-algorithm=NEXTN",
"--speculative-num-steps=3",
"--speculative-eagle-topk=1",
"--speculative-num-draft-tokens=4",
],
adapters=[("case", "opherlie/lora-test-case-Qwen3.5-35B-A3B", "dataset")],
derive_from="case",
common_args=[
"--max-lora-rank=64",
"--moe-runner-backend=triton",
"--experts-shared-outer-loras",
"--lora-use-virtual-experts",
"--disable-shared-experts-fusion",
"--mem-fraction-static=0.8",
],
tp=4,
),
# MoE + EAGLE3 with a draft checkpoint matched to the same base variant.
"eagle3-qwen3-30b-a3b": dict(
model="Qwen/Qwen3-30B-A3B-Instruct-2507",
spec_args=[
"--speculative-algorithm=EAGLE3",
"--speculative-draft-model-path=lmsys/SGLang-EAGLE3-Qwen3-30B-A3B-Instruct-2507-SpecForge-Nex",
"--speculative-num-steps=3",
"--speculative-eagle-topk=4",
"--speculative-num-draft-tokens=8",
],
adapters=[
("case", "yushengsu/lora-diff-Qwen3-30B-A3B-Instruct-2507", "dataset")
],
derive_from="case",
common_args=[
"--max-lora-rank=32",
"--moe-runner-backend=triton",
"--experts-shared-outer-loras",
"--attention-backend=flashinfer",
"--prefill-attention-backend=fa4",
"--decode-attention-backend=fa4",
"--mem-fraction-static=0.8",
],
# The draft checkpoint's config derives a 2048 context; the draft
# follows the target's context_length (it reads target KV), so the
# longer-context guard has to be waived for this pairing.
env={"SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN": "1"},
tp=4,
),
# DFLASH: fixed-size block drafting, verified in one TARGET_VERIFY pass.
# Its target is the one base with several genuinely distinct public
# adapters, so no synthesis is needed.
"dflash-llama31": dict(
model="meta-llama/Llama-3.1-8B-Instruct",
spec_args=[
"--speculative-algorithm=DFLASH",
"--speculative-draft-model-path=z-lab/LLaMA3.1-8B-Instruct-DFlash-UltraChat",
"--speculative-num-draft-tokens=4",
],
adapters=[
("fact", "algoprog/fact-generation-llama-3.1-8b-instruct-lora"),
("guard", "nvidia/llama-3.1-nemoguard-8b-topic-control"),
],
common_args=["--mem-fraction-static=0.7", "--max-lora-rank=64"],
# The draft checkpoint derives a 40960 context against the target's
# 131072; the draft follows the target, so waive the guard.
env={"SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN": "1"},
tp=1,
),
# DSPARK with a separate DSpark draft head, on the one base that has two
# distinct public adapters. Only the default static ragged-verify mode is
# supported: the others schedule per-request verify lengths, which the
# uniform-width LoRA segment layout cannot express.
"dspark-inkling": dict(
model="thinkingmachines/Inkling-Small",
spec_args=[
"--speculative-algorithm=DSPARK",
"--speculative-draft-model-path=RadixArk/Inkling-Small-DSpark",
"--speculative-dspark-block-size=5",
],
adapters=[
("gutenberg", "nbeerbower/Inkling-Small-Gutenberg-DPO-LoRA"),
("hemlock", "hemlang/Inkling-Small-Hemlock-SFT-LoRA"),
],
common_args=[
"--max-lora-rank=32",
"--moe-runner-backend=triton",
"--experts-shared-outer-loras",
"--mem-fraction-static=0.8",
],
env={"SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN": "1"},
# A ~500GB checkpoint over network storage does not load inside the
# default launch budget.
launch_timeout=3600,
tp=8,
),
# Flagship MLA + MoE NextN. 8 GPUs.
"nextn-deepseek-v31": dict(
model="deepseek-ai/DeepSeek-V3.1-Base",
spec_args=[
"--speculative-algorithm=NEXTN",
"--speculative-num-steps=3",
"--speculative-eagle-topk=1",
"--speculative-num-draft-tokens=4",
],
adapters=[("case", "yushengsu/lora-diff-DeepSeek-V3.1-Base", "dataset")],
derive_from="case",
common_args=[
"--max-lora-rank=32",
"--moe-runner-backend=triton",
"--experts-shared-outer-loras",
"--attention-backend=flashinfer",
"--prefill-attention-backend=fa4",
"--decode-attention-backend=flashinfer",
"--disable-shared-experts-fusion",
"--mem-fraction-static=0.8",
],
tp=8,
),
}
def log(msg: str) -> None:
print(f"[harness] {msg}", flush=True)
def materialize_adapter(hf_path: str, repo_type: str) -> str:
"""Download an adapter and return its local directory.
Several adapters used here live in *dataset* repos (that is how the
repo's own LoRA logprob-diff tests host them), which the server cannot
resolve as a model id so every adapter is materialized locally and
passed to --lora-paths by path.
"""
from huggingface_hub import snapshot_download
return snapshot_download(repo_id=hf_path, repo_type=repo_type)
def check_adapter_supported(name: str, local_path: str) -> None:
"""Report what each adapter exercises. Nothing here is a rejection: a
draft sharing the target's lm_head gets the unwrapped base layer, so
lm_head and embedding adapters cost accept rate, not correctness."""
config_path = os.path.join(local_path, "adapter_config.json")
try:
with open(config_path) as f:
config = json.load(f)
except Exception as e: # noqa: BLE001 - not every layout ships one
log(f" {name}: could not read adapter_config.json ({e}); skipping precheck")
return
modules = config.get("target_modules")
log(f" {name}: r={config.get('r')} target_modules={modules}")
if isinstance(modules, list):
shared = {"lm_head", "output", "unembed_tokens"} & set(modules)
if shared:
log(f" {name}: targets {sorted(shared)} — shared-lm_head path")
if {"embed_tokens", "vocab_emb", "word_embeddings"} & set(modules):
log(f" {name}: targets embeddings — expect a reduced accept rate")
if config.get("modules_to_save"):
log(f" {name}: WARNING modules_to_save={config['modules_to_save']}")
LM_HEAD_MARKERS = ("lm_head", "unembed_tokens")
def adapter_lm_head_keys(local_path: str):
"""Tensor keys that would land on lm_head (PEFT unembed_tokens keys are
rewritten to lm_head during load)."""
from safetensors import safe_open
weights = os.path.join(local_path, "adapter_model.safetensors")
if not os.path.isfile(weights):
return []
with safe_open(weights, "pt") as f:
return [k for k in f.keys() if any(m in k for m in LM_HEAD_MARKERS)]
def write_adapter_variant(
source_dir: str, out_dir: str, keep: bool, *, negate_b: bool, drop_lm_head: bool
) -> str:
"""Write a modified copy of an adapter.
``drop_lm_head`` removes output-layer tensors, isolating an adapter's
other modules from the shared-lm_head path. The server accepts lm_head
adapters, so this is for narrowing a failure, not a requirement.
``negate_b`` flips the sign of every lora_B tensor, producing a second
adapter of identical shape whose outputs genuinely diverge from the
original without which a "mixed-adapter" batch would compare an adapter
against itself and could not detect crossed verify segments.
"""
if keep and os.path.isdir(out_dir):
log(f"reusing adapter variant at {out_dir}")
return out_dir
from safetensors.torch import load_file, save_file
weights_name = "adapter_model.safetensors"
src_weights = os.path.join(source_dir, weights_name)
if not os.path.isfile(src_weights):
raise SystemExit(f"{source_dir} has no {weights_name}")
if os.path.isdir(out_dir):
shutil.rmtree(out_dir)
os.makedirs(out_dir)
tensors = load_file(src_weights)
dropped = flipped = 0
out = {}
for key, value in tensors.items():
if drop_lm_head and any(m in key for m in LM_HEAD_MARKERS):
dropped += 1
continue
if negate_b and "lora_B" in key:
value = -value
flipped += 1
out[key] = value
save_file(out, os.path.join(out_dir, weights_name))
shutil.copy(
os.path.join(source_dir, "adapter_config.json"),
os.path.join(out_dir, "adapter_config.json"),
)
log(
f"wrote {out_dir} (dropped {dropped} lm_head tensors, "
f"negated {flipped} lora_B tensors)"
)
return out_dir
def launch(config: dict, adapters, base_url: str, with_spec: bool):
other_args = list(config["common_args"])
env = config.get("env")
if config["tp"] > 1:
other_args += [f"--tp={config['tp']}"]
if adapters:
other_args += [
"--enable-lora",
f"--lora-backend={config.get('lora_backend', 'triton')}",
# +1: the base model occupies a pool slot too, and every mixed
# batch below co-batches base requests.
f"--max-loras-per-batch={len(adapters) + 1}",
"--lora-paths",
] + [f"{name}={path}" for name, path in adapters]
else:
# Base arms: --max-lora-rank is a LoRA-only flag and is rejected
# without --enable-lora, so drop it from the common set.
other_args = [a for a in other_args if not a.startswith("--max-lora-rank")]
if with_spec:
other_args += config["spec_args"]
other_args += config.get("extra_args", [])
log(f"launching {'spec' if with_spec else 'reference'} server: {other_args}")
return popen_launch_server(
config["model"],
base_url,
timeout=config.get("launch_timeout", DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH),
other_args=other_args,
env=env,
)
def generate(base_url: str, texts, routes, max_new_tokens: int):
"""One /generate call; returns (texts, per-request accept lengths)."""
response = requests.post(
base_url + "/generate",
json={
"text": texts,
"lora_path": routes,
"sampling_params": {"temperature": 0, "max_new_tokens": max_new_tokens},
},
timeout=1800,
)
response.raise_for_status()
results = response.json()
outputs = [item["text"] for item in results]
accepts = [item["meta_info"].get("spec_accept_length") for item in results]
return outputs, accepts
def min_logprob_gap(base_url: str, route, prompt: str, max_new_tokens: int):
"""Smallest top-2 logprob gap along the greedy path for one request.
A gap of ~0 means the argmax at that position is decided by kernel scan
order, so any change in reduction order (a verify forward computes logits
at a different batch shape than a plain decode) flips the token with no
correctness implication. Used to classify surviving mismatches.
"""
response = requests.post(
base_url + "/generate",
json={
"text": prompt,
"lora_path": route,
"sampling_params": {"temperature": 0, "max_new_tokens": max_new_tokens},
"return_logprob": True,
"top_logprobs_num": 2,
},
timeout=600,
)
response.raise_for_status()
tops = response.json()["meta_info"].get("output_top_logprobs") or []
gaps = [c[0][0] - c[1][0] for c in tops if c and len(c) >= 2]
return min(gaps) if gaps else None
def collect_solo(base_url: str, routes, max_new_tokens: int):
"""Per route: outputs and mean accept length over PROMPTS."""
solo, accept = {}, {}
for route in routes:
outputs, accepts = generate(
base_url, list(PROMPTS), [route] * len(PROMPTS), max_new_tokens
)
solo[route] = outputs
seen = [a for a in accepts if a]
accept[route] = sum(seen) / len(seen) if seen else None
return solo, accept
def report_parity(label: str, reference: dict, actual: dict) -> int:
"""Compare per-route output lists; returns the mismatch count."""
mismatches = 0
for route, expected in reference.items():
for i, (want, got) in enumerate(zip(expected, actual[route])):
if want != got:
mismatches += 1
log(f" MISMATCH [{label}] route={route} prompt#{i}")
log(f" spec-off: {want!r}")
log(f" spec-on : {got!r}")
return mismatches
def build_load(num_prompts: int, input_len: int, seed: int):
"""Fixed synthetic prompts as raw token ids, so every arm sends identical
work and the input length is exact rather than tokenizer-dependent."""
rng = random.Random(seed)
return [
[rng.randint(1000, 10000) for _ in range(input_len)] for _ in range(num_prompts)
]
def one_request(base_url: str, input_ids, route, output_len: int):
started = time.perf_counter()
response = requests.post(
base_url + "/generate",
json={
"input_ids": input_ids,
"lora_path": route,
"sampling_params": {
"temperature": 0,
"max_new_tokens": output_len,
# Fixed-length output: every arm emits exactly output_len
# tokens, so throughput is comparable without length drift.
"ignore_eos": True,
},
},
timeout=3600,
)
response.raise_for_status()
meta = response.json()["meta_info"]
return dict(
latency=time.perf_counter() - started,
output_tokens=meta.get("completion_tokens", output_len),
accept_length=meta.get("spec_accept_length"),
)
def measure(base_url: str, loads, routes, output_len: int, concurrency: int):
"""Drive `concurrency` in-flight requests over the load; return throughput,
accept length, and per-request latency."""
pairs = [(loads[i], routes[i % len(routes)]) for i in range(len(loads))]
started = time.perf_counter()
with futures.ThreadPoolExecutor(max_workers=concurrency) as pool:
results = list(
pool.map(
lambda pair: one_request(base_url, pair[0], pair[1], output_len), pairs
)
)
elapsed = time.perf_counter() - started
accepts = [r["accept_length"] for r in results if r["accept_length"]]
total_out = sum(r["output_tokens"] for r in results)
latencies = sorted(r["latency"] for r in results)
return dict(
output_tps=total_out / elapsed,
accept_length=(sum(accepts) / len(accepts)) if accepts else None,
median_latency=latencies[len(latencies) // 2],
elapsed=elapsed,
requests=len(results),
)
def run_perf(config: dict, args) -> int:
"""Four arms answering two questions: does speculation still pay once LoRA
is on (lora_spec vs lora_nospec), and how much does serving adapters cost
relative to speculating on the plain base model (lora_spec vs base_spec)?
"""
base_url = f"http://127.0.0.1:{args.port}"
adapters = resolve_adapters(config, args)
concurrencies = [int(c) for c in args.concurrency.split(",")]
loads = build_load(args.num_prompts, args.input_len, args.seed)
warmup = build_load(max(concurrencies), args.input_len, args.seed + 1)
# (label, serve LoRA?, speculate?)
arms = [
("base_nospec", False, False),
("base_spec", False, True),
("lora_nospec", True, False),
("lora_spec", True, True),
]
table = {}
for label, with_lora, with_spec in arms:
log("=" * 68)
log(f"arm {label}: lora={with_lora} spec={with_spec}")
process = None
try:
process = launch(
config,
adapters if with_lora else [],
base_url,
with_spec=with_spec,
)
# Multi-adapter arms spread requests round-robin across every
# adapter (the multi-tenant shape); base arms send no adapter.
routes = [name for name, _ in adapters] if with_lora else [None]
measure(base_url, warmup, routes, args.output_len, len(warmup))
for concurrency in concurrencies:
stats = measure(base_url, loads, routes, args.output_len, concurrency)
table[(label, concurrency)] = stats
log(
f" c={concurrency:<4} {stats['output_tps']:8.1f} tok/s "
f"accept={stats['accept_length']} "
f"p50={stats['median_latency']:.2f}s"
)
finally:
if process is not None:
kill_process_tree(process.pid)
report_perf(table, concurrencies, args)
return 0
def report_perf(table, concurrencies, args) -> None:
log("=" * 78)
log(
f"PERF model-shape: {args.num_prompts} reqs x {args.input_len} in / "
f"{args.output_len} out"
)
log("=" * 78)
labels = ["base_nospec", "base_spec", "lora_nospec", "lora_spec"]
log(f"{'concurrency':<12}" + "".join(f"{label:>16}" for label in labels))
for concurrency in concurrencies:
cells = []
for label in labels:
stats = table.get((label, concurrency))
cells.append(f"{stats['output_tps']:.1f}" if stats else "-")
log(
f"{concurrency:<12}" + "".join(f"{cell:>16}" for cell in cells) + " tok/s"
)
log("")
log(f"{'concurrency':<12}{'base accept':>16}{'lora accept':>16}")
for concurrency in concurrencies:
base = table.get(("base_spec", concurrency), {}).get("accept_length")
lora = table.get(("lora_spec", concurrency), {}).get("accept_length")
fmt = lambda v: f"{v:.3f}" if v else "-" # noqa: E731
log(f"{concurrency:<12}{fmt(base):>16}{fmt(lora):>16}")
log("")
log("ratios (>1 is better)")
log(
f"{'concurrency':<12}{'spec on base':>16}{'spec on lora':>16}"
f"{'lora vs base':>16}"
)
for concurrency in concurrencies:
def tps(label):
stats = table.get((label, concurrency))
return stats["output_tps"] if stats else None
base_nospec, base_spec = tps("base_nospec"), tps("base_spec")
lora_nospec, lora_spec = tps("lora_nospec"), tps("lora_spec")
ratio = lambda a, b: f"{a / b:.2f}x" if a and b else "-" # noqa: E731
log(
f"{concurrency:<12}"
f"{ratio(base_spec, base_nospec):>16}"
f"{ratio(lora_spec, lora_nospec):>16}"
f"{ratio(lora_spec, base_spec):>16}"
)
log("")
log(" spec on base = speculation speedup without adapters (the ceiling)")
log(" spec on lora = speculation speedup with adapters served (the ask)")
log(" lora vs base = cost of serving adapters, both speculating")
def resolve_adapters(config: dict, args):
"""Materialize every adapter locally and return (name, local_path) pairs,
deriving a second distinct adapter when the base has only one public one.
Config entries are (name, hf_path) or (name, hf_path, repo_type).
"""
adapters = []
for entry in config["adapters"]:
name, hf_path = entry[0], entry[1]
repo_type = entry[2] if len(entry) > 2 else "model"
local_path = materialize_adapter(hf_path, repo_type)
# A draft sharing the target's lm_head gets the unwrapped base layer,
# so these adapters cost accept rate, not correctness. Keep the
# weights as published so the run exercises that path.
lm_head_keys = adapter_lm_head_keys(local_path)
if lm_head_keys:
log(
f" {name}: carries {len(lm_head_keys)} lm_head tensor(s) "
f"({lm_head_keys[0]}) — exercising the shared-lm_head path"
)
adapters.append((name, local_path))
log(f"prechecking {len(adapters)} adapter(s) against the lm_head rule")
for name, local_path in adapters:
check_adapter_supported(name, local_path)
derive_name = config.get("derive_from")
if derive_name:
source = dict(adapters)[derive_name]
# Key the cache on the source path: several configs name their
# adapter "case", so a name-only key served one model's variant to
# another (mismatched layer counts).
source_key = hashlib.sha1(source.encode()).hexdigest()[:8]
adapters.append(
(
derive_name + DERIVED_SUFFIX,
write_adapter_variant(
source,
os.path.join(
args.derived_dir, f"{derive_name}{DERIVED_SUFFIX}-{source_key}"
),
args.keep_derived,
negate_b=True,
drop_lm_head=False,
),
)
)
return adapters
def collect_shapes(base_url: str, routes, args):
"""Outputs for every (shape, route, prompt) the checks compare.
Three shapes: solo (one route per batch), mixed (every route interleaved
in one batch), wide (a batch larger than the cuda-graph capture, so the
eager verify path runs).
"""
out = {}
solo, accept = collect_solo(base_url, routes, args.max_new_tokens)
for route, texts in solo.items():
for i, text in enumerate(texts):
out[("solo", route, i)] = text
mixed_texts = [p for p in PROMPTS for _ in routes]
mixed_routes = [r for _ in PROMPTS for r in routes]
mixed, _ = generate(base_url, mixed_texts, mixed_routes, args.max_new_tokens)
for text, route, got in zip(mixed_texts, mixed_routes, mixed):
out[("mixed", route, PROMPTS.index(text))] = got
wide = max(args.wide_batch, len(routes) * 2)
wide_routes = [routes[i % len(routes)] for i in range(wide)]
wide_texts = [PROMPTS[i % len(PROMPTS)] for i in range(wide)]
wide_out, _ = generate(base_url, wide_texts, wide_routes, args.max_new_tokens)
for text, route, got in zip(wide_texts, wide_routes, wide_out):
out[("wide", route, PROMPTS.index(text))] = got
return out, accept, solo
def unstable_keys(ref_a: dict, ref_b: dict) -> set:
"""(route, prompt) pairs the reference itself cannot reproduce.
Greedy decoding is not bitwise reproducible here: batch composition
changes reduction order (MoE routing especially), and two runs of the
same server can diverge. Anything unstable *without* speculation is a
property of the model and backend, not of spec+LoRA, so the spec
comparison must exclude it or it reports noise as failure.
"""
unstable = set()
for (shape, route, index), text in ref_a.items():
# run-to-run, identical shape
if ref_b.get((shape, route, index)) != text:
unstable.add((route, index))
# shape-to-shape, same run
if ref_a.get(("solo", route, index)) != text:
unstable.add((route, index))
return unstable
def is_reproducible(base_url: str, route, prompt: str, args, repeats: int = 3) -> bool:
"""Whether one prompt yields the same greedy output across repeats."""
outputs = {
generate(base_url, [prompt], [route], args.max_new_tokens)[0][0]
for _ in range(repeats)
}
return len(outputs) == 1
def compare(ref_a, spec_a, spec_b, noisy, base_url: str, args) -> int:
"""Count genuine spec-on vs spec-off differences.
Skips pairs the reference could not reproduce, pairs where the spec side
disagrees with itself, and differences that land on an argmax tie (where
the token choice is arbitrary in the first place).
"""
log("comparing spec-on against spec-off on reproducible pairs only")
excluded = ties = failures = 0
for key, want in ref_a.items():
shape, route, index = key
# Both sides must be self-consistent before a difference between them
# means anything: the instability set is stochastic, so a single
# sample per side would report noise as a failure.
if (route, index) in noisy or spec_a.get(key) != spec_b.get(key):
excluded += 1
continue
got = spec_a.get(key)
if got == want:
continue
# A pair can look stable in two samples and still be stochastic, so
# re-sample this specific pair before calling it a defect. Verified
# necessary: a prompt that survived the blanket filter turned out to
# diverge spec-on vs spec-off with no adapters loaded at all, and to
# disagree with itself across repeats.
if not is_reproducible(base_url, route, PROMPTS[index], args):
ties += 1
gap = min_logprob_gap(base_url, route, PROMPTS[index], args.max_new_tokens)
log(
f" UNSTABLE [{shape}] route={route} prompt#{index}: the spec "
f"server does not reproduce this prompt across repeats "
f"(min top-2 logprob gap {gap}); not attributable to spec"
)
continue
failures += 1
log(f" MISMATCH [{shape}] route={route} prompt#{index}")
log(f" spec-off: {want!r}")
log(f" spec-on : {got!r}")
log(
f" ({excluded} skipped as baseline-nondeterministic, "
f"{ties} skipped as not reproducible on re-sampling)"
)
return failures
def run(config: dict, args) -> int:
base_url = f"http://127.0.0.1:{args.port}"
adapters = resolve_adapters(config, args)
routes = [None] + [name for name, _ in adapters]
failures = 0
process = None
try:
# Phase 1 — reference: same weights and adapters, speculation off.
# Collected twice to measure how much this model/backend varies on
# its own, which sets the noise floor for the spec comparison.
process = launch(config, adapters, base_url, with_spec=False)
ref_a, _, solo_a = collect_shapes(base_url, routes, args)
ref_b, _, _ = collect_shapes(base_url, routes, args)
noisy = unstable_keys(ref_a, ref_b)
log(
f"baseline noise floor: {len(noisy)} of "
f"{len(routes) * len(PROMPTS)} (route, prompt) pairs are not "
"reproducible without speculation"
)
for route in routes[1:]:
if solo_a[route] == solo_a[None]:
failures += 1
log(
f" FAIL adapter {route} output is identical to base on every prompt"
)
kill_process_tree(process.pid)
process = None
# Phase 2 — speculation on, same three shapes.
process = launch(config, adapters, base_url, with_spec=True)
spec_a, accept, _ = collect_shapes(base_url, routes, args)
spec_b, _, _ = collect_shapes(base_url, routes, args)
failures += compare(ref_a, spec_a, spec_b, noisy, base_url, args)
finally:
if process is not None:
kill_process_tree(process.pid)
log("=" * 68)
log("per-adapter accept length (tokens per verify step; 1.0 = no speedup)")
for route in routes:
value = accept.get(route)
label = "base" if route is None else route
log(f" {label:<24} {value if value is None else round(value, 3)}")
log("=" * 68)
log(f"RESULT: {'PASS' if failures == 0 else f'FAIL ({failures} mismatches)'}")
return 1 if failures else 0
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--config", help=f"one of: {', '.join(CONFIGS)}")
parser.add_argument("--list", action="store_true", help="list configs and exit")
parser.add_argument("--port", type=int, default=30000)
parser.add_argument("--max-new-tokens", type=int, default=32)
parser.add_argument(
"--wide-batch",
type=int,
default=16,
help="batch size for the eager-path check; keep it above --cuda-graph-max-bs",
)
parser.add_argument(
"--lora-backend", help="override the config's LoRA kernel backend"
)
parser.add_argument(
"--extra-arg",
action="append",
default=[],
dest="extra_args",
help="extra server flag, repeatable (e.g. --extra-arg=--enable-lora-overlap-loading)",
)
parser.add_argument("--derived-dir", default="/tmp/sglang-derived-loras")
parser.add_argument("--keep-derived", action="store_true")
parser.add_argument(
"--mode",
choices=["correctness", "perf", "both"],
default="correctness",
help="correctness = the 4 output checks; perf = the 4-arm throughput "
"and accept-length comparison",
)
parser.add_argument(
"--concurrency",
default="1,8,32",
help="comma-separated in-flight request counts to sweep (perf mode)",
)
parser.add_argument("--num-prompts", type=int, default=64)
parser.add_argument("--input-len", type=int, default=512)
parser.add_argument("--output-len", type=int, default=256)
parser.add_argument("--seed", type=int, default=42)
args = parser.parse_args()
if args.list or not args.config:
for name, config in CONFIGS.items():
algo = next(
a.split("=")[1] for a in config["spec_args"] if "algorithm" in a
)
n = len(config["adapters"]) + (1 if config.get("derive_from") else 0)
print(
f"{name:<24} tp={config['tp']} {algo:<7} adapters={n} {config['model']}"
)
return 0
if args.config not in CONFIGS:
raise SystemExit(f"unknown config {args.config!r}; use --list")
started = time.time()
config = dict(CONFIGS[args.config])
if args.lora_backend:
config["lora_backend"] = args.lora_backend
config["extra_args"] = args.extra_args
code = 0
if args.mode in ("correctness", "both"):
code |= run(config, args)
if args.mode in ("perf", "both"):
code |= run_perf(config, args)
log(f"total wall time {time.time() - started:.0f}s")
return code
if __name__ == "__main__":
sys.exit(main())
@@ -0,0 +1,142 @@
"""E2E test for multi-adapter LoRA + EAGLE-family speculative decoding.
Adapters apply to the target model only; one shared draft runs unadapted.
Deliberately asserts serving properties rather than exact output text.
Greedy decoding is not bitwise reproducible across batch shapes or server
restarts here -- reduction order changes flip a token and greedy amplifies
it -- so text equality between configurations is a flaky assertion, not a
correctness oracle. Losslessness is verified out of band by
test/manual/lora/run_spec_lora_matrix.py, which measures that noise floor
first. What this guards is what CI can hold stable:
- the server starts at all with multiple adapters + speculation (it used to
crash loading the target's adapters into the draft model);
- adapters are actually applied during target-verify (adapter output differs
from base output);
- a mixed-adapter batch, and a batch wider than the cuda-graph capture, are
served without error (the eager verify path used to crash on
extend_seq_lens_cpu=None);
- speculation is really running (accept length above 1).
"""
import unittest
import requests
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import (
DEFAULT_DRAFT_MODEL_EAGLE3,
DEFAULT_TARGET_MODEL_EAGLE3,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
register_cuda_ci(est_time=600, stage="base-b", runner_config="1-gpu-large")
PROMPTS = [
"What is the capital of France? Answer in one sentence.",
"List three primary colors.",
"Write a one-sentence story about a brave detective on Mars.",
]
# Ranks 8 and 64: a mixed-rank batch is what the per-request lora_ranks
# indexing has to get right, and uniform ranks would hide a mixup.
ADAPTERS = [
("fact", "algoprog/fact-generation-llama-3.1-8b-instruct-lora"),
("guard", "nvidia/llama-3.1-nemoguard-8b-topic-control"),
]
SAMPLING = {"temperature": 0, "max_new_tokens": 32}
class TestEagle3MultiLoRA(CustomTestCase):
process = None
@classmethod
def setUpClass(cls):
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
DEFAULT_TARGET_MODEL_EAGLE3,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
# Canonical EAGLE3 sglang config, as in
# test/registered/core/test_basic_sanity_eagle3.py: the draft
# checkpoint is fp16, and bf16 + flashinfer cutlass RMSNorm
# hits a dtype mismatch on the draft's input_layernorm.
"--dtype=float16",
"--attention-backend=triton",
"--speculative-algorithm=EAGLE3",
f"--speculative-draft-model-path={DEFAULT_DRAFT_MODEL_EAGLE3}",
"--speculative-num-steps=5",
"--speculative-eagle-topk=8",
"--speculative-num-draft-tokens=32",
"--enable-lora",
"--lora-backend=triton",
"--max-lora-rank=64",
# +1: the base model occupies a memory-pool slot too, and the
# batches below co-batch base requests with adapter requests.
f"--max-loras-per-batch={len(ADAPTERS) + 1}",
"--mem-fraction-static=0.7",
"--lora-paths",
]
+ [f"{name}={path}" for name, path in ADAPTERS],
)
@classmethod
def tearDownClass(cls):
if cls.process is not None:
kill_process_tree(cls.process.pid)
def _generate(self, texts, lora_paths):
response = requests.post(
self.base_url + "/generate",
json={"text": texts, "lora_path": lora_paths, "sampling_params": SAMPLING},
)
self.assertEqual(response.status_code, 200, response.text)
results = response.json()
self.assertEqual(len(results), len(texts))
for item in results:
self.assertTrue(item["text"].strip(), f"empty output in {results}")
return [item["text"] for item in results]
def test_adapters_are_applied_under_speculation(self):
base = self._generate(PROMPTS, [None] * len(PROMPTS))
for name, _ in ADAPTERS:
adapted = self._generate(PROMPTS, [name] * len(PROMPTS))
self.assertNotEqual(
adapted,
base,
f"adapter {name} matched the base model on every prompt; LoRA "
"was likely not applied during target-verify",
)
def test_mixed_adapter_and_wide_batches_are_served(self):
routes = [None] + [name for name, _ in ADAPTERS]
self._generate(
[p for p in PROMPTS for _ in routes],
[r for _ in PROMPTS for r in routes],
)
# Wider than the default cuda-graph capture, so target-verify falls
# back to the eager path.
wide = 24
self._generate(
[PROMPTS[i % len(PROMPTS)] for i in range(wide)],
[routes[i % len(routes)] for i in range(wide)],
)
def test_speculation_is_active(self):
self._generate(PROMPTS, [ADAPTERS[0][0]] * len(PROMPTS))
info = requests.get(self.base_url + "/get_server_info").json()
accept_length = info["internal_states"][0]["avg_spec_accept_length"]
self.assertGreater(
accept_length,
1.0,
f"no drafts accepted with LoRA enabled: {accept_length}",
)
if __name__ == "__main__":
unittest.main(verbosity=3)
@@ -0,0 +1,109 @@
"""Unit tests for LoRA batch-info preparation under TARGET_VERIFY.
TARGET_VERIFY reports is_extend() True but ForwardBatch.init_new leaves
extend_seq_lens / extend_seq_lens_cpu as None (verify is routed through the
decode-style positions branch), while every request carries a uniform
spec_info.draft_token_num-token segment. Regression guards for the paths that
assumed extend fields exist on every is_extend() mode:
- triton eager path: ``max(None)`` TypeError, and 1-token/req segments for a
draft_token_num-token layout (silent adapter mis-segmentation).
- MoE _add_moe_lora_info: ``sum(None)`` TypeError.
- static cuda-graph path: seg_lens are pre-filled at the captured width; a
verify batch of another width must fail loudly, not mis-segment.
"""
import unittest
from types import SimpleNamespace
from unittest.mock import patch
import torch
from sglang.srt.lora.backend.triton_backend import TritonLoRABackend
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=30, stage="base-b", runner_config="1-gpu-small")
def _verify_batch(bs: int, draft_token_num: int) -> SimpleNamespace:
return SimpleNamespace(
forward_mode=ForwardMode.TARGET_VERIFY,
batch_size=bs,
spec_info=SimpleNamespace(draft_token_num=draft_token_num),
extend_seq_lens=None,
extend_seq_lens_cpu=None,
)
class TestLoRASpecVerifyBatchInfo(CustomTestCase):
def _backend(self, max_loras_per_batch: int = 2) -> TritonLoRABackend:
return TritonLoRABackend(
max_loras_per_batch=max_loras_per_batch, device=torch.device("cuda")
)
def _prepare(self, backend, forward_batch, use_cuda_graph: bool):
backend.prepare_lora_batch(
forward_batch,
weight_indices=[0, 1],
lora_ranks=[8, 8],
scalings=[1.0, 1.0],
use_cuda_graph=use_cuda_graph,
)
return backend.batch_info
def test_eager_target_verify_builds_uniform_draft_width_segments(self):
backend = self._backend()
batch_info = self._prepare(
backend, _verify_batch(bs=2, draft_token_num=4), use_cuda_graph=False
)
self.assertEqual(batch_info.max_len, 4)
self.assertEqual(batch_info.seg_lens[:2].tolist(), [4, 4])
self.assertEqual(batch_info.seg_indptr[:3].tolist(), [0, 4, 8])
self.assertEqual(batch_info.num_segments, 2)
def test_graph_path_serves_the_captured_width_and_rejects_others(self):
"""prepare_lora_batch predicts graph use before can_run_graph decides,
so a mismatched width used to silently apply the captured segment
layout to a differently-shaped batch (wrong adapter on wrong rows).
The matching case also pins that bs is rebound per batch, which is
what makes a mixed-adapter batch index the right slots."""
backend = self._backend()
backend.init_cuda_graph_batch_info(max_bs_in_cuda_graph=4, num_tokens_per_req=4)
batch_info = self._prepare(
backend, _verify_batch(bs=2, draft_token_num=4), use_cuda_graph=True
)
self.assertIs(batch_info, backend.cuda_graph_batch_info)
self.assertEqual(batch_info.seg_lens.tolist(), [4, 4, 4, 4])
self.assertEqual(batch_info.bs, 2)
with self.assertRaisesRegex(AssertionError, "width"):
self._prepare(
backend, _verify_batch(bs=2, draft_token_num=8), use_cuda_graph=True
)
def test_moe_lora_info_uses_draft_width_token_counts_for_verify(self):
backend = self._backend()
backend.is_moe_lora = True
captured = {}
def _capture(num_tokens, seg_indptr, lora_ranks, req_to_lora, *args, **kwargs):
captured["num_tokens"] = num_tokens
captured["max_len"] = kwargs["max_len"]
return None, None
with patch(
"sglang.srt.lora.backend.base_backend._compute_moe_lora_info",
side_effect=_capture,
):
self._prepare(
backend, _verify_batch(bs=2, draft_token_num=4), use_cuda_graph=False
)
self.assertEqual(captured["num_tokens"], 8)
self.assertEqual(captured["max_len"], 4)
if __name__ == "__main__":
unittest.main()
@@ -0,0 +1,51 @@
"""The draft runner must not build a LoRA manager.
Adapters apply to the target model only. Every worker is handed the same
published ServerArgs, so `enable_lora` is True for the draft too -- the
decision is the runner's own, keyed on is_draft_worker. Without it the draft
tries to load the target's adapters into the draft model, whose layer count
differs, and startup fails inside LoRAAdapter weight loading.
The LoRA paths downstream then key on `lora_manager is not None` rather than
the config, so a draft runner skips them by construction.
"""
import unittest
from unittest.mock import patch
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
class TestDraftRunnerSkipsLoRA(CustomTestCase):
def _init_lora_called(self, *, is_draft_worker: bool, enable_lora: bool) -> bool:
runner = ModelRunner.__new__(ModelRunner)
runner.is_draft_worker = is_draft_worker
runner.lora_manager = None
with patch.object(ModelRunner, "init_lora_manager") as init_lora:
with patch("sglang.srt.model_executor.model_runner.get_lora") as get_lora:
get_lora.return_value.enable_lora = enable_lora
runner.maybe_init_lora_manager()
return init_lora.called
def test_only_the_target_runner_builds_a_lora_manager(self):
cases = [
(False, True, True), # target + LoRA -> builds one
(True, True, False), # draft + LoRA -> must not
(False, False, False), # LoRA off -> nobody builds one
]
for is_draft_worker, enable_lora, expected in cases:
with self.subTest(draft=is_draft_worker, enable_lora=enable_lora):
self.assertEqual(
self._init_lora_called(
is_draft_worker=is_draft_worker, enable_lora=enable_lora
),
expected,
)
if __name__ == "__main__":
unittest.main()
@@ -69,9 +69,11 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
device="cuda",
gpu_id=0,
is_draft_worker=False,
# A real ModelRunner always has this attribute; the prefill gate
# reads it rather than the process-wide LoRA config.
lora_manager=None,
spec_algorithm=SimpleNamespace(is_eagle=lambda: False),
server_args=SimpleNamespace(
enable_lora=False,
cuda_graph_config=SimpleNamespace(
prefill=SimpleNamespace(bs=[1], backend=Backend.BREAKABLE)
),