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

Split #32584 into 2/2: [LoRA] Shard attention LoRA by attn-TP and allow dynamic LoRA with dp attention (#32708)

Browse Source
This commit is contained in:
Ethan (Yusheng) Su
2026-07-31 15:37:13 -07:00
committed by GitHub
parent 55b6769b0e
commit 3c5f115741
7 changed files with 287 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/tokenizer_control_mixin.py
+34 -11
View File
@@ -123,6 +123,27 @@ _COMMUNICATOR_SPECS = [
]
def _merge_lora_update_results(results: List[LoRAUpdateOutput]) -> LoRAUpdateOutput:
"""Merge the per-rank replies of a LoRA load/unload fan-out into one result.
The operation succeeded only if every rank succeeded. Reporting a partial
failure as success would let the tokenizer-side LoRA registry drift from
the ranks that failed, so failures win: their deduplicated error messages
are joined, and loaded_adapters reflects the first failed rank.
"""
failed = [r for r in results if not r.success]
if not failed:
return results[0]
error_messages = list(
dict.fromkeys(r.error_message for r in failed if r.error_message)
)
return LoRAUpdateOutput(
success=False,
error_message=" | ".join(error_messages),
loaded_adapters=failed[0].loaded_adapters,
)
class TokenizerControlMixin:
"""Mixin for TokenizerManager's control-plane operations (weights, cache, lora,
profile, internal state, etc.) -- everything that talks to the scheduler via
@@ -557,7 +578,9 @@ class TokenizerControlMixin:
# Initiate the actual unloading operation at the backend processes only after all
# ongoing requests using this LoRA adapter are finished.
await self.lora_registry.wait_for_unload(lora_id)
result = (await self.update_lora_adapter_communicator(obj))[0]
result = _merge_lora_update_results(
await self.update_lora_adapter_communicator(obj)
)
return result
@@ -574,11 +597,9 @@ class TokenizerControlMixin:
"LoRA is not enabled. Please set `--enable-lora` to enable LoRA."
)
# TODO (lifuhuang): Remove this after we verify that dynamic lora loading works
# with dp_size > 1.
assert (
self.server_args.dp_size == 1
), "dp_size must be 1 for dynamic lora loading"
self.server_args.dp_size == 1 or self.server_args.enable_dp_attention
), "dp_size must be 1 or dp attention must be enabled for dynamic lora loading"
logger.info(
"Start load Lora adapter. Lora name=%s, path=%s",
obj.lora_name,
@@ -595,7 +616,9 @@ class TokenizerControlMixin:
# Trigger the actual loading operation at the backend processes.
obj.lora_id = new_adapter.lora_id
result = (await self.update_lora_adapter_communicator(obj))[0]
result = _merge_lora_update_results(
await self.update_lora_adapter_communicator(obj)
)
# Register the LoRA adapter only after loading is successful.
if result.success:
@@ -671,7 +694,9 @@ class TokenizerControlMixin:
pinned=obj.pinned,
)
obj.lora_id = new_adapter.lora_id
result = (await self.update_lora_adapter_communicator(obj))[0]
result = _merge_lora_update_results(
await self.update_lora_adapter_communicator(obj)
)
if result.success:
await self.lora_registry.register(new_adapter)
@@ -730,11 +755,9 @@ class TokenizerControlMixin:
obj.lora_name is not None
), "lora_name must be provided to unload LoRA adapter"
# TODO (lifuhuang): Remove this after we verify that dynamic lora loading works
# with dp_size > 1.
assert (
self.server_args.dp_size == 1
), "dp_size must be 1 for dynamic lora loading"
self.server_args.dp_size == 1 or self.server_args.enable_dp_attention
), "dp_size must be 1 or dp attention must be enabled for dynamic lora loading"
logger.info(
"Start unload Lora adapter. Lora name=%s",
obj.lora_name,