Skip to content

[LoRA] Sentinel-pad token->lora mapping for DP-attention foreign tokens - #28566

Closed
nanjiangwill wants to merge 1 commit into
sgl-project:mainfrom
nanjiangwill:upstream-dp-attn-lora-mapping-sentinel
Closed

[LoRA] Sentinel-pad token->lora mapping for DP-attention foreign tokens#28566
nanjiangwill wants to merge 1 commit into
sgl-project:mainfrom
nanjiangwill:upstream-dp-attn-lora-mapping-sentinel

Conversation

@nanjiangwill

@nanjiangwill nanjiangwill commented Jun 17, 2026

Copy link
Copy Markdown
Contributor

Summary

Under DP-attention, the gathered forward batch contains both local and foreign tokens, but `seg_indptr` / `weight_indices` cover only the local requests. `_compute_moe_lora_info` in `backend/base_backend.py` currently does one of two wrong things at the foreign positions:

  • CUDA-kernel path: `token_lora_mapping` is `torch.empty(...)`, so foreign-position slots stay uninitialized when the kernel only writes to positions within segments. The MoE LoRA hook then dispatches on garbage.
  • Python fallback path: `torch.searchsorted(seg_indptr, ..., right=True) - 1` returns `num_segments` for foreign positions, which makes `torch.index_select(weight_indices, 0, req_indices, ...)` index one past the end of `weight_indices`.

This PR fixes both paths with the same -1-sentinel convention:

  1. Pre-fill `token_lora_mapping` with `-1` before either path runs, so un-segmented slots stay `-1` after the kernel returns.
  2. In the Python fallback, `torch.cat` a `-1` entry onto `weight_indices` and use that padded tensor for the `index_select`, so the same `-1` lands at the foreign positions.

Foreign-token LoRA outputs are discarded by the DP-attention scatter anyway, so mapping them to `-1` ("no LoRA") is the correct local behaviour — it just stops the kernel from acting on uninitialised data / out-of-bounds reads.

Context

Carved out of #25141. The sglang-miles version of this fix lives in a helper called `_compute_token_lora_mapping` in `lora_moe_runners.py`; on `main` the equivalent work is done in `_compute_moe_lora_info` here, so this PR ports the same idea to the on-main code path.

Test plan

  • Existing non-DP-attention LoRA tests — behaviour unchanged (the pre-fill is a no-op when the kernel/fallback overwrite every slot, which is the normal case).
  • DP-attention + MoE LoRA: confirm the kernel and Python paths both place `-1` at foreign positions.
  • Empty / no-segment edge cases — covered by the simplified `num_tokens == 0 or not has_segments` return-early (the prior explicit `fill_(-1)` for the no-segments case is subsumed by the earlier pre-fill).

CI States

Latest PR Test (Base): ❌ Run #27722492929
Latest PR Test (Extra): ❌ Run #27851913609

Under DP-attention, the gathered forward batch contains both local and
foreign tokens, but ``seg_indptr`` / ``weight_indices`` cover only local
requests. ``_compute_moe_lora_info`` was either reading garbage memory
(CUDA-kernel path leaves un-segmented positions uninitialized in the
``torch.empty(...)`` buffer) or going out of bounds (Python fallback's
``index_select`` indexes ``weight_indices[num_segments]``).

Fix both paths with the same -1-sentinel convention:

* Pre-fill ``token_lora_mapping`` with -1 before either path runs, so
  un-segmented slots stay -1 after the kernel returns.
* In the Python fallback, ``cat`` a -1 entry onto ``weight_indices`` and
  pass that to ``index_select`` so the same -1 lands at the foreign
  positions.

Foreign-token outputs are discarded by the DP-attention scatter, so
mapping them to -1 ("no LoRA") is the correct local behaviour — it
just stops the kernel from acting on uninitialised data.

Carved out of sgl-project#25141 (the sglang-miles version of this fix lived in
``_compute_token_lora_mapping``, which doesn't exist on main; the
equivalent here is ``_compute_moe_lora_info``).
@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!

@nanjiangwill

Copy link
Copy Markdown
Contributor Author

not needed anymore

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants