Skip to content

support kimi 2.5/6 lora (logprob diff exist) #25141

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yushengsu-thu merged 1 commit into from
May 25, 2026
Merged

support kimi 2.5/6 lora (logprob diff exist) #25141

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 6 additions & 5 deletions python/sglang/srt/entrypoints/engine.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1133,15 +1133,16 @@ def load_lora_adapter_from_tensors(
load_format: Optional[str] = None,
):
if load_format == "flattened_bucket":
serialized_tensors = tensors
serialized_named_tensors = list(tensors)
else:
serialized_tensors = MultiprocessingSerializer.serialize(
tensors, output_str=True
)
serialized_named_tensors = [
MultiprocessingSerializer.serialize(tensors, output_str=True)
for _ in range(self.server_args.tp_size)
]
lora_req = LoadLoRAAdapterFromTensorsReqInput(
lora_name=lora_name,
config_dict=config_dict,
serialized_tensors=serialized_tensors,
serialized_named_tensors=serialized_named_tensors,
load_format=load_format,
)
return self.loop.run_until_complete(
Expand Down
11 changes: 7 additions & 4 deletions python/sglang/srt/lora/backend/base_backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ class BaseLoRABackend(LoRABackendLmHeadMixing):
def __init__(self, max_loras_per_batch: int, device: torch.device):
self.max_loras_per_batch = max_loras_per_batch
self.device = device
self.batch_info = None
self.init_lm_head_config()

def run_lora_a_embedding(
Expand Down Expand Up @@ -176,10 +177,12 @@ def init_cuda_graph_moe_buffers(
"""
base = moe_layer.base_layer
top_k = base.top_k
qinfo = moe_layer._quant_info
E, N, _ = qinfo.w13_weight.shape
hidden_dim = qinfo.w2_weight.shape[1]
device = qinfo.w13_weight.device
# Derive dims from the base FusedMoE rather than quant-specific tensors,
# so this works for any scheme (FP, WNA16, Marlin-packed, etc.).
E = base.num_local_experts
hidden_dim = base.hidden_size
N = 2 * base.intermediate_size_per_partition
device = next(base.parameters()).device
dtype = compute_dtype
num_experts = base.num_experts

Expand Down
81 changes: 66 additions & 15 deletions python/sglang/srt/lora/layers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,15 @@ def __init__(
self.weight = self.base_layer.weight
if hasattr(self.base_layer, "bias") and self.base_layer.bias is not None:
self.bias = self.base_layer.bias
if hasattr(self.base_layer, "reduce_results"):
self.reduce_results = self.base_layer.reduce_results
# Alias remaining base-layer parameters onto the wrapper so
# `named_parameters(remove_duplicate=True)` yields them at the outer
# path — weight loaders (e.g. FusedMoE's `w13_weight_packed`) lookup
# names without the `.base_layer.` segment.
for _name, _param in base_layer.named_parameters(recurse=False):
if not hasattr(self, _name):
setattr(self, _name, _param)

def forward(self, x: torch.Tensor):
return self.base_layer.forward(x)
Expand Down Expand Up @@ -207,8 +216,9 @@ def forward(self, input_: torch.Tensor):
):
base_output = self.extra_token_embedding(input_, base_output)

# Apply LoRA if configured
if self.set_lora:
# Apply LoRA if configured. Skip if no batch_info (DP-attention idle
# forward): the base path is correct because no real tokens need LoRA.
if self.set_lora and self.lora_backend.batch_info is not None:
# The backend's run_lora_a_embedding now handles both regular
# and extra tokens efficiently with CUDA graph support
base_output = self.apply_lora(base_output, input_, batch_info)
Expand Down Expand Up @@ -373,8 +383,8 @@ def forward(self, hidden_states: torch.Tensor):
hidden_states, self.weight, bias=getattr(self.base_layer, "bias", None)
)

# Apply LoRA if set
if self.set_lora:
# Apply LoRA if set. Skip in DP-attention idle forward (batch_info unset).
if self.set_lora and self.lora_backend.batch_info is not None:
base_output = self.apply_lora(base_output, hidden_states)

return base_output
Expand Down Expand Up @@ -463,7 +473,7 @@ def forward(self, input_: torch.Tensor):
self.base_layer, input_, bias
)

if self.set_lora:
if self.set_lora and self.lora_backend.batch_info is not None:
output_parallel = self.apply_lora(output_parallel, input_)

if self.base_layer.gather_output:
Expand All @@ -477,9 +487,13 @@ def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int):
return A

def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int):
# See RowParallelLinearWithLoRA.slice_lora_a_weights for why base_layer.tp_rank
# is authoritative: DP-attention makes output_partition_sizes attn_tp-local while
# the caller passes global tp_rank.
local_tp_rank = getattr(self.base_layer, "tp_rank", tp_rank)
shard_size = self.base_layer.output_partition_sizes[0]
start_idx = tp_rank * shard_size
end_idx = (tp_rank + 1) * shard_size
start_idx = local_tp_rank * shard_size
end_idx = (local_tp_rank + 1) * shard_size
B = B[start_idx:end_idx, :]
return B

Expand Down Expand Up @@ -573,12 +587,15 @@ def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int):
return A

def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int):
# base_layer.tp_rank is authoritative under DP-attention: the caller passes
# the global tp_rank but output_partition_sizes is attn_tp-local.
local_tp_rank = getattr(self.base_layer, "tp_rank", tp_rank)
partition_sizes = self.base_layer.output_partition_sizes
output_sizes = self.base_layer.output_sizes
slices = []
offset = 0
for full_size, part_size in zip(output_sizes, partition_sizes):
start_idx = tp_rank * part_size
start_idx = local_tp_rank * part_size
end_idx = start_idx + part_size
slices.append(B[offset + start_idx : offset + end_idx, :])
offset += full_size
Expand Down Expand Up @@ -645,11 +662,14 @@ def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int) -> torch.Tensor:
q_proj_shard_size = base_layer.q_proj_shard_size
kv_proj_shard_size = base_layer.kv_proj_shard_size
num_kv_head_replicas = base_layer.num_kv_head_replicas
# See RowParallelLinearWithLoRA.slice_lora_a_weights for why base_layer.tp_rank
# is authoritative under DP-attention.
local_tp_rank = getattr(base_layer, "tp_rank", tp_rank)

q_start_idx = q_proj_shard_size * tp_rank
q_start_idx = q_proj_shard_size * local_tp_rank
q_end_idx = q_start_idx + q_proj_shard_size

kv_shard_id = tp_rank // num_kv_head_replicas
kv_shard_id = local_tp_rank // num_kv_head_replicas
kv_start_idx = kv_proj_shard_size * kv_shard_id
kv_end_idx = kv_start_idx + kv_proj_shard_size

Expand Down Expand Up @@ -731,7 +751,9 @@ def forward(self, input_: torch.Tensor, skip_all_reduce=False, forward_batch=Non
and not skip_all_reduce
)

if self.set_lora and should_reduce:
# LoRA skipped when batch_info is None (DP-attention idle forward).
have_batch_info = self.lora_backend.batch_info is not None
if self.set_lora and have_batch_info and should_reduce:
lora_a_output = self.lora_backend.run_lora_a_sgemm(
input_parallel, self.A_buffer
)
Expand All @@ -745,7 +767,7 @@ def forward(self, input_: torch.Tensor, skip_all_reduce=False, forward_batch=Non
base_output=output_,
)
else:
if self.set_lora:
if self.set_lora and have_batch_info:
output_parallel = self.apply_lora(output_parallel, input_parallel)
if should_reduce:
output_ = tensor_model_parallel_all_reduce(output_parallel)
Expand All @@ -756,9 +778,15 @@ def forward(self, input_: torch.Tensor, skip_all_reduce=False, forward_batch=Non
return output_, output_bias

def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int):
# Use base_layer.tp_rank (not the argument) so the slicing rank matches
# the partition group the base layer was built on. For MLA o_proj under
# DP-attention, base_layer.tp_rank is attn_tp_rank while the caller
# passes the global tp_rank; input_size_per_partition is already
# attn_tp-sized, so using global tp_rank overshoots to empty.
local_tp_rank = getattr(self.base_layer, "tp_rank", tp_rank)
shard_size = self.base_layer.input_size_per_partition
start_idx = tp_rank * shard_size
end_idx = (tp_rank + 1) * shard_size
start_idx = local_tp_rank * shard_size
end_idx = (local_tp_rank + 1) * shard_size
A = A[:, start_idx:end_idx].contiguous()
return A

Expand Down Expand Up @@ -844,7 +872,7 @@ def apply_lora(self, base_output: torch.Tensor, x: torch.Tensor) -> torch.Tensor
def forward(self, x: torch.Tensor):
bias = self.base_layer.bias if not self.base_layer.skip_bias_add else None
output = self.base_layer.quant_method.apply(self.base_layer, x, bias)
if self.set_lora:
if self.set_lora and self.lora_backend.batch_info is not None:
output = self.apply_lora(output, x)
output_bias = self.base_layer.bias if self.base_layer.skip_bias_add else None
return output, output_bias
Expand Down Expand Up @@ -878,7 +906,19 @@ def __init__(

self.experts_shared_outer_loras: bool = False
self.lora_use_virtual_experts: bool = False
# Forward for the model's own forward-path dispatch — the outer model
# reads several FusedMoE attributes (e.g. `self.experts.moe_runner_config`,
# `self.experts.dispatcher`, `self.experts.num_local_experts`,
# `self.experts.quant_method`) directly on the wrapper. Quant
# post-processing iterators skip LoRA wrappers via
# `isinstance(module, BaseLayerWithLoRA)` so the packed params on the
# inner FusedMoE get processed there, not here.
self.quant_method = base_layer.quant_method
self.moe_runner_config = base_layer.moe_runner_config
self.dispatcher = base_layer.dispatcher
self.num_local_experts = base_layer.num_local_experts
if hasattr(base_layer, "scheme"):
self.scheme = base_layer.scheme

self.tp_size = getattr(base_layer, "moe_tp_size", 1)
self.tp_rank = getattr(base_layer, "moe_tp_rank", 0)
Expand All @@ -903,6 +943,12 @@ def __init__(
and base_layer.quant_method.runner is not None
):
runner_backend = base_layer.quant_method.runner.runner_backend
elif (
hasattr(base_layer, "scheme")
and hasattr(base_layer.scheme, "runner")
and base_layer.scheme.runner is not None
):
runner_backend = base_layer.scheme.runner.runner_backend
else:
runner_backend = MoeRunnerBackend.TRITON

Expand Down Expand Up @@ -1007,6 +1053,8 @@ def forward(self, hidden_states: torch.Tensor, topk_output: TopKOutput, **kwargs
1. After gate_up projection, before activation
2. After down projection, before final reduction
"""
if self.lora_backend.batch_info is None:
return self.base_layer.forward(hidden_states, topk_output, **kwargs)

# Build LoRA info for this batch
lora_info = self._get_lora_info()
Expand Down Expand Up @@ -1034,6 +1082,9 @@ def _forward_with_lora(

# Use pre-computed quant info (doesn't change so not sure why we need to pass it in every time)
quant_info = self._quant_info
quant_info.expert_map = getattr(
base_layer.dispatcher, "local_expert_mapping", None
)

# Run the only lora moe runner (Triton)
combine_input = self._lora_runner.run(
Expand Down
22 changes: 18 additions & 4 deletions python/sglang/srt/lora/lora_moe_runners.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -205,7 +205,14 @@ def _compute_token_lora_mapping(
hidden_states: torch.Tensor,
lora_info: LoRAInfo,
) -> torch.Tensor:
"""Map each token to its LoRA adapter index (-1 for no LoRA)."""
"""Map each token to its LoRA adapter index (-1 for no LoRA).

Under DP-attention, `hidden_states` is the gathered batch (local + foreign
tokens) but `seg_indptr` / `req_to_lora` cover only local requests, so
`searchsorted` on foreign positions would index one past the end. Pad
`req_to_lora` with a -1 sentinel; foreign outputs are discarded by the
DP-attention scatter anyway.
"""
token_positions = torch.arange(
hidden_states.shape[0], device=hidden_states.device, dtype=torch.int32
)
Expand All @@ -214,7 +221,11 @@ def _compute_token_lora_mapping(
token_positions,
right=True,
)
return lora_info.req_to_lora.to(torch.int32)[req_indices]
req_to_lora = lora_info.req_to_lora.to(torch.int32)
req_to_lora_padded = torch.cat(
[req_to_lora, req_to_lora.new_full((1,), -1)], dim=0
)
return req_to_lora_padded[req_indices]


def _compute_lora_alignment(
Expand Down Expand Up @@ -349,7 +360,6 @@ def _add_lora_gate_up_delta(
r = lora_info.max_lora_rank
gate_up_a = lora_info.gate_up_lora_a_weights
gate_up_b = lora_info.gate_up_lora_b_weights

if lora_info.experts_shared_outer_loras and not lora_info.lora_use_virtual_experts:
gate_up_a = gate_up_a.expand(-1, lora_info.num_experts, -1, -1)

Expand All @@ -359,6 +369,8 @@ def _add_lora_gate_up_delta(
if is_gated:
inter_size = gate_up_b.shape[2] // 2
lora_a_stacked = [gate_up_a[:, :, :r, :], gate_up_a[:, :, r : 2 * r, :]]
# B halves are also the tuple form the virtual-experts kernel wants
# (one shrink at K=2*r, two expands at K=r each).
lora_b_stacked = [
gate_up_b[:, :, :inter_size, :],
gate_up_b[:, :, inter_size:, :],
Expand All @@ -372,7 +384,7 @@ def _add_lora_gate_up_delta(
output=intermediate_cache,
hidden_states=hidden_states,
lora_a=gate_up_a,
lora_b=gate_up_b,
lora_b=tuple(lora_b_stacked) if is_gated else gate_up_b,
topk_ids=topk_ids,
topk_weights=topk_weights,
token_lora_mapping=token_lora_mapping,
Expand Down Expand Up @@ -447,6 +459,8 @@ def _add_lora_down_delta(
down_lora_a = lora_info.down_lora_a_weights
down_lora_b = lora_info.down_lora_b_weights
if lora_info.experts_shared_outer_loras and not lora_info.lora_use_virtual_experts:
# fused_moe_lora requires B's expert_dim to match A's; expand the
# shared B view.
down_lora_b = down_lora_b.expand(-1, lora_info.num_experts, -1, -1)

if lora_info.fully_sharded and lora_info.tp_size > 1:
Expand Down
Loading
Loading