Skip to content

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

Merged
yushengsu-thu merged 1 commit into
sgl-project:sglang-milesfrom
nanjiangwill:kimi25
May 25, 2026
Merged

support kimi 2.5/6 lora (logprob diff exist)#25141
yushengsu-thu merged 1 commit into
sgl-project:sglang-milesfrom
nanjiangwill:kimi25

Conversation

@nanjiangwill

@nanjiangwill nanjiangwill commented May 13, 2026

Copy link
Copy Markdown
Contributor

Adds Kimi K2.5 / K2.6 LoRA support to the `sglang-miles` branch:

  • MoE LoRA: `experts_shared_outer_loras` layout. `gate_up_proj_moe` has shared A (`(slot, 1, 2r, hidden)`) + per-expert B; `down_proj_moe` has the inverse. `LoRAMemoryPool` initializes A and B independently per target so a shared Tensor on one side doesn't clobber the per-expert dict on the other.
  • Virtual-experts kernel: `lora_b` as a list. `merged_experts_fused_moe_lora_add` accepts a length-2 `lora_b` for the gated case (one `2r` shrink, two `r` expands); `lora_moe_runners` passes `tuple(lora_b_stacked)` for it.
  • Chunked `load_lora_adapter_from_tensors`. Per-rank serialization: `serialized_tensors: str` -> `serialized_named_tensors: List[Union[str, bytes]]` so each TP rank deserializes its own producer's CUDA-IPC bucket (matches `update_weights_from_tensor`). `Engine` and `BaseTpWorker` plumbed accordingly; `TokenizerControlMixin` relaxes `dp_size == 1` to also allow `enable_dp_attention`.
  • DP-attention safety. Token→lora mapping pads `req_to_lora` with a `-1` sentinel so foreign tokens past `seg_indptr[-1]` don't `searchsorted` past the end / read uninitialised memory.
  • Quant post-process skips LoRA wrappers. `restore_weights_before_load` / `process_weights_after_loading` loops in `ModelRunner` skip `BaseLayerWithLoRA`; the inner base layer is yielded separately by `named_modules` and gets processed there.
  • MLA fused-A cache persistence. `DeepseekV2WeightLoaderMixin` persists `cached_a_proj` across `do_load_weights` calls so chunked updates that split `q_a_proj` / `kv_a_proj_with_mqa` across chunks still produce the fused `fused_qkv_a_proj_with_mqa` once both halves arrive. (Reverted on `sglang-miles` in Revert deepseek_weight_loader.py changes from #25141 #28561 — sender-side pairing in Pair q_a_proj / kv_a_proj_with_mqa in the bridge weight iterator radixark/miles#1361 is the cleaner fix.)
  • Base-layer dim derivation. `BaseLoRABackend` derives MoE intermediate buffer dims from the base FusedMoE rather than quant-specific tensors so it works for any scheme (FP, WNA16, Marlin-packed, …).

`logprob diff exist` in the title flags a residual `train/train_rollout_logprob_abs_diff` drift seen on K2.5 RL runs;


Decomposition for `main`: #28562, #28564, #28566. The MoE post-process LoRA-wrapper skip is blocked on #28565 (upstream prereq). The MLA fused-A cache change is being reverted on `sglang-miles` via #28561.

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

Comment thread python/sglang/srt/lora/triton_ops/virtual_experts.py
Comment thread python/sglang/srt/lora/layers.py Outdated
Comment thread python/sglang/srt/lora/mem_pool.py
Comment thread python/sglang/srt/managers/tp_worker.py
@yushengsu-thu

Copy link
Copy Markdown
Collaborator

these modification should be sent to sglang/main as well or it will becomse useless after we upgrade the sgl

@github-actions github-actions Bot added quant LLM Quantization dependencies Pull requests that update a dependency file blackwell SM100/SM120 diffusion SGLang Diffusion model-gateway jit-kernel labels May 23, 2026
@yushengsu-thu

yushengsu-thu commented May 25, 2026

Copy link
Copy Markdown
Collaborator

the modification are almost lora related only ‎python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py‎ is related to kimi2.5 base.

We can consider if revert the kimi2.5 base related parts. To unblock kimi full-param and lora training support so I merge them first.

@yushengsu-thu
yushengsu-thu merged commit c0cac9e into sgl-project:sglang-miles May 25, 2026
2 checks passed
JessicaJiang-123 pushed a commit to JessicaJiang-123/sglang that referenced this pull request Jun 21, 2026
nanjiangwill added a commit to nanjiangwill/sglang that referenced this pull request Jun 30, 2026
The persistent cached_a_proj and unconditional .detach().clone()
were a workaround for q_a_proj / kv_a_proj_with_mqa landing in
different chunks during chunked weight updates. The proper fix is
on the sender side: pair the two halves before chunking so they
always arrive in the same load_weights call. That fix is in
radixark/miles#1361 (mirrors THUDM/slime#1532).

With sender-side pairing in place, the cached_a_proj only ever
needs to live for one load_weights call, matching the original
upstream behavior.
nanjiangwill added a commit to nanjiangwill/sglang that referenced this pull request Jul 7, 2026
yueming-yuan pushed a commit that referenced this pull request Jul 25, 2026
…h/RL fixes (#25141, #29874, #31251)

Squashes the two follow-up PRs into the MoE-LoRA base commit: both rewrite the
mem_pool sharding helpers this commit introduces and cannot be applied
independently.

Rebased onto v0.5.16:
  - the shard probes keep v0.5.16's shared-MoE-over-full-TP carve-out at EP=1
    (`is_shared_moe_module`) and only fall through to the probed shard for
    non-MoE modules;
  - #29874's `free_lora` is dropped -- v0.5.16 already releases the slot on
    unload through `LoRAMemoryPool.remove_lora()`, which additionally zeroes the
    buffers for graph-captured replay, so `lora_manager.unload` needs no second
    call;
  - #31831's expected_checksums check is likewise already upstream in
    `tp_worker.load_lora_adapter_from_tensors`.

Co-authored-by: Ethan (Yusheng) Su <yushengsu@radixark.ai>
yueming-yuan pushed a commit that referenced this pull request Jul 25, 2026
…h/RL fixes (#25141, #29874, #31251)

Squashes the two follow-up PRs into the MoE-LoRA base commit: both rewrite the
mem_pool sharding helpers this commit introduces and cannot be applied
independently.

Rebased onto v0.5.16:
  - the shard probes keep v0.5.16's shared-MoE-over-full-TP carve-out at EP=1
    (`is_shared_moe_module`) and only fall through to the probed shard for
    non-MoE modules;
  - #29874's `free_lora` is dropped -- v0.5.16 already releases the slot on
    unload through `LoRAMemoryPool.remove_lora()`, which additionally zeroes the
    buffers for graph-captured replay, so `lora_manager.unload` needs no second
    call;
  - #31831's expected_checksums check is likewise already upstream in
    `tp_worker.load_lora_adapter_from_tensors`.

Co-authored-by: Ethan (Yusheng) Su <yushengsu@radixark.ai>
yueming-yuan pushed a commit that referenced this pull request Jul 25, 2026
…h/RL fixes (#25141, #29874, #31251)

Squashes the two follow-up PRs into the MoE-LoRA base commit: both rewrite the
mem_pool sharding helpers this commit introduces and cannot be applied
independently.

Rebased onto v0.5.16:
  - the shard probes keep v0.5.16's shared-MoE-over-full-TP carve-out at EP=1
    (`is_shared_moe_module`) and only fall through to the probed shard for
    non-MoE modules;
  - #29874's `free_lora` is dropped -- v0.5.16 already releases the slot on
    unload through `LoRAMemoryPool.remove_lora()`, which additionally zeroes the
    buffers for graph-captured replay, so `lora_manager.unload` needs no second
    call;
  - #31831's expected_checksums check is likewise already upstream in
    `tp_worker.load_lora_adapter_from_tensors`.

Co-authored-by: Ethan (Yusheng) Su <yushengsu@radixark.ai>
yueming-yuan pushed a commit that referenced this pull request Jul 25, 2026
…h/RL fixes (#25141, #29874, #31251)

Squashes the two follow-up PRs into the MoE-LoRA base commit: both rewrite the
mem_pool sharding helpers this commit introduces and cannot be applied
independently.

Rebased onto v0.5.16:
  - the shard probes keep v0.5.16's shared-MoE-over-full-TP carve-out at EP=1
    (`is_shared_moe_module`) and only fall through to the probed shard for
    non-MoE modules;
  - #29874's `free_lora` is dropped -- v0.5.16 already releases the slot on
    unload through `LoRAMemoryPool.remove_lora()`, which additionally zeroes the
    buffers for graph-captured replay, so `lora_manager.unload` needs no second
    call;
  - #31831's expected_checksums check is likewise already upstream in
    `tp_worker.load_lora_adapter_from_tensors`.

Co-authored-by: Ethan (Yusheng) Su <yushengsu@radixark.ai>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

blackwell SM100/SM120 deepseek dependencies Pull requests that update a dependency file diffusion SGLang Diffusion jit-kernel lora model-gateway quant LLM Quantization

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants