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

python/sglang/srt/lora/layers.py

@@ -66,10 +66,13 @@ class BaseLayerWithLoRA(nn.Module):
     def set_lora_info(self, *args):
         pass
 
-    def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int):
+    # Weight slicing derives the shard rank from the wrapped base layer
+    # (base_layer.tp_rank): under DP attention, attention layers are built
+    # on the attn-TP group, so the outer/global TP rank would overshoot.
+    def slice_lora_a_weights(self, A: torch.Tensor):
         pass
 
-    def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int):
+    def slice_lora_b_weights(self, B: torch.Tensor):
         pass
 
 
@@ -234,13 +237,13 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):
 
         return base_output
 
-    def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int):
+    def slice_lora_a_weights(self, A: torch.Tensor):
         # LoRA A weights (rank, vocab_size) are kept unsharded.
         # Each rank does a full embedding lookup; the result is complete
         # on every rank and added to the already all-reduced base output.
         return A
 
-    def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int):
+    def slice_lora_b_weights(self, B: torch.Tensor):
         # LoRA B weights (embedding_dim, rank) are kept unsharded.
         # The base embedding output is all-reduced (full embedding_dim),
         # so LoRA B must also produce full embedding_dim.
@@ -414,12 +417,12 @@ class ParallelLMHeadWithLoRA(BaseLayerWithLoRA):
         """Reset the lm_head pass index after all passes are done."""
         self.lora_backend._lm_head_pass_idx = None
 
-    def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int):
+    def slice_lora_a_weights(self, A: torch.Tensor):
         # LoRA A weights (rank, hidden_size) are kept unsharded.
         # Each rank receives full hidden_states, so A operates on full input.
         return A
 
-    def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int):
+    def slice_lora_b_weights(self, B: torch.Tensor):
         # lm_head is column-parallel: each rank produces vocab_size/tp_size (shard_vocab_size)
         # logits.  LoRA B (vocab_size, rank) must be sliced along the vocab
         # dimension to match the sharded base output.
@@ -491,13 +494,14 @@ class ColumnParallelLinearWithLoRA(BaseLayerWithLoRA):
         output_bias = self.base_layer.bias if self.base_layer.skip_bias_add else None
         return output, output_bias
 
-    def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int):
+    def slice_lora_a_weights(self, A: torch.Tensor):
         return A
 
-    def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int):
+    def slice_lora_b_weights(self, B: torch.Tensor):
+        local_tp_rank = self.base_layer.tp_rank
         shard_size = self.base_layer.output_partition_sizes[0]
-        start_idx = tp_rank * shard_size
+        start_idx = local_tp_rank * shard_size
-        end_idx = (tp_rank + 1) * shard_size
+        end_idx = (local_tp_rank + 1) * shard_size
         B = B[start_idx:end_idx, :]
         return B
 
@@ -587,16 +591,17 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
             )
         return lora_output
 
-    def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int):
+    def slice_lora_a_weights(self, A: torch.Tensor):
         return A
 
-    def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int):
+    def slice_lora_b_weights(self, B: torch.Tensor):
+        local_tp_rank = self.base_layer.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
@@ -613,8 +618,9 @@ class InklingQKVRLinearWithLoRA(MergedColumnParallelLinearWithLoRA):
     are head-partitioned uniformly. LoRA-A stays unsharded (inherited).
     """
 
-    def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int):
+    def slice_lora_b_weights(self, B: torch.Tensor):
         bl = self.base_layer
+        tp_rank = bl.tp_rank
         hd, nkv, nh, dr, tp = (
             bl.inkling_head_dim,
             bl.inkling_num_kv_heads,
@@ -692,19 +698,20 @@ class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
 
         return lora_output
 
-    def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int):
+    def slice_lora_a_weights(self, A: torch.Tensor):
         return A
 
-    def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int) -> torch.Tensor:
+    def slice_lora_b_weights(self, B: torch.Tensor) -> torch.Tensor:
         base_layer = self.base_layer
         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
+        local_tp_rank = base_layer.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
 
@@ -787,13 +794,20 @@ class RowParallelLinearWithLoRA(BaseLayerWithLoRA):
             and not should_skip_mlp_all_reduce()
         )
 
+        # Match the base layer's reduce group: layers built with
+        # use_dp_attention_reduce are sharded over the attn-TP group, so
+        # reducing over the global TP group would mix tokens across DP groups.
+        if self.base_layer.use_dp_attention_reduce:
+            all_reduce = get_parallel().attn_tp_group.all_reduce
+        else:
+            all_reduce = tensor_model_parallel_all_reduce
         lora_active = self.lora_active
         if lora_active and should_reduce:
             lora_a_output = self.lora_backend.run_lora_a_sgemm(
                 input_parallel, self.A_buffer
             )
-            output_ = tensor_model_parallel_all_reduce(output_parallel)
+            output_ = all_reduce(output_parallel)
-            lora_a_output = tensor_model_parallel_all_reduce(lora_a_output)
+            lora_a_output = all_reduce(lora_a_output)
             output_ = self.lora_backend.run_lora_b_sgemm(
                 x=lora_a_output,
                 weights=self.B_buffer,
@@ -805,21 +819,22 @@ class RowParallelLinearWithLoRA(BaseLayerWithLoRA):
             if lora_active:
                 output_parallel = self.apply_lora(output_parallel, input_parallel)
             if should_reduce:
-                output_ = tensor_model_parallel_all_reduce(output_parallel)
+                output_ = all_reduce(output_parallel)
             else:
                 output_ = output_parallel
 
         output_bias = self.base_layer.bias if self.base_layer.skip_bias_add else None
         return output_, output_bias
 
-    def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int):
+    def slice_lora_a_weights(self, A: torch.Tensor):
+        local_tp_rank = self.base_layer.tp_rank
         shard_size = self.base_layer.input_size_per_partition
-        start_idx = tp_rank * shard_size
+        start_idx = local_tp_rank * shard_size
-        end_idx = (tp_rank + 1) * shard_size
+        end_idx = (local_tp_rank + 1) * shard_size
         A = A[:, start_idx:end_idx].contiguous()
         return A
 
-    def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int):
+    def slice_lora_b_weights(self, B: torch.Tensor):
         return B
 
 
@@ -906,10 +921,10 @@ class ReplicatedLinearWithLoRA(BaseLayerWithLoRA):
         output_bias = self.base_layer.bias if self.base_layer.skip_bias_add else None
         return output, output_bias
 
-    def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int):
+    def slice_lora_a_weights(self, A: torch.Tensor):
         return A
 
-    def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int):
+    def slice_lora_b_weights(self, B: torch.Tensor):
         return B
 
 
@@ -1166,10 +1181,10 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
 
         return final_hidden_states
 
-    def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int):
+    def slice_lora_a_weights(self, A: torch.Tensor):
         return A
 
-    def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int):
+    def slice_lora_b_weights(self, B: torch.Tensor):
         return B
 
     def slice_moe_lora_a_weights(

python/sglang/srt/lora/lora_manager.py

@@ -46,6 +46,7 @@ from sglang.srt.lora.utils import (
 )
 from sglang.srt.managers.io_struct import LoRAUpdateOutput
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.utils import get_available_gpu_memory, replace_submodule
 from sglang.srt.utils.hf_transformers_utils import AutoConfig
@@ -82,6 +83,9 @@ class LoRAManager:
         self.device: torch.device = next(self.base_model.parameters()).device
         self.tp_size: int = tp_size
         self.tp_rank: int = tp_rank
+        # Attention projections shard on the attn-TP group; extracted once
+        # here (parallel groups are frozen after init_torch_distributed).
+        self.attn_tp_size: int = get_parallel().attn_tp_size
         self.lora_added_tokens_size: Optional[int] = None
         self.enable_lora_overlap_loading: Optional[bool] = (
             server_args.enable_lora_overlap_loading
@@ -853,6 +857,7 @@ class LoRAManager:
             dtype=self.dtype,
             tp_size=self.tp_size,
             tp_rank=self.tp_rank,
+            attn_tp_size=self.attn_tp_size,
             max_lora_rank=self.max_lora_rank,
             target_modules=self.target_modules,
             base_model=self.base_model,

python/sglang/srt/lora/mem_pool.py

@@ -26,6 +26,7 @@ from sglang.srt.lora.lora import LoRAAdapter
 from sglang.srt.lora.lora_config import LoRAConfig
 from sglang.srt.lora.lora_registry import LoRARef
 from sglang.srt.lora.utils import (
+    ATTN_TP_LORA_MODULE_NAMES,
     EMBEDDING_NAMES,
     REPLICATED_LINEAR_LORA_NAMES,
     ROW_PARALLELISM_LINEAR_LORA_NAMES,
@@ -137,6 +138,7 @@ class LoRAMemoryPool:
         dtype: torch.dtype,
         tp_size: int,
         tp_rank: int,
+        attn_tp_size: int,
         max_lora_rank: int,
         target_modules: Set[str],
         base_model: torch.nn.Module,
@@ -182,10 +184,15 @@ class LoRAMemoryPool:
         # here would yield a 4x-narrower inner dim than the adapter weight
         # (which MoE LoRA modules correctly skip-slice when
         # `moe_tp_size <= 1`), producing a shape-mismatch
-        # assert during weight load. Non-MoE modules still shard by
+        # assert during weight load.
-        # `tp_size` because attention TP is unchanged.
         self.moe_tp_size, self.moe_tp_rank = _get_moe_tp_context()
 
+        # Attention projections shard along the attention TP group, which
+        # under `--enable-dp-attention` is `attn_tp_size = tp_size // dp_size`.
+        # The corresponding LoRA wrappers slice weights by the base layer's
+        # attn_tp-local rank, so the buffer shapes must match that shard.
+        self.attn_tp_size: int = attn_tp_size
+
         # Initialize eviction policy
         self.eviction_policy = get_eviction_policy(eviction_policy)
 
@@ -260,6 +267,20 @@ class LoRAMemoryPool:
         """Whether this buffer belongs to the shared-expert MoE namespace."""
         return module_name.endswith("_shared_moe")
 
+    def _effective_tp_size(self, module_name: str) -> int:
+        """TP width the module's weights are actually sharded along: routed
+        MoE experts shard by `moe_tp_size` (shared experts by the outer
+        `tp_size` at EP=1), attention projections by `attn_tp_size` (smaller
+        than the outer `tp_size` under `--enable-dp-attention`), everything
+        else by the outer `tp_size`."""
+        if self.is_moe_module(module_name) and not self.is_shared_moe_module(
+            module_name
+        ):
+            return self.moe_tp_size
+        if module_name in ATTN_TP_LORA_MODULE_NAMES:
+            return self.attn_tp_size
+        return self.tp_size
+
     @staticmethod
     def _get_num_experts(base_model: torch.nn.Module) -> int:
         cfg = base_model.config
@@ -388,13 +409,7 @@ class LoRAMemoryPool:
             module_name, self.base_hf_config, base_model, layer_idx
         )
         c = get_stacked_multiply(module_name, base_model)
-        # Routed MoE shards along moe_tp_size; shared MoE shards over full TP at EP=1.
+        effective_tp_size = self._effective_tp_size(module_name)
-        effective_tp_size = (
-            self.tp_size
-            if not self.is_moe_module(module_name)
-            or self.is_shared_moe_module(module_name)
-            else self.moe_tp_size
-        )
         if (
             effective_tp_size > 1
             and module_name in ROW_PARALLELISM_LINEAR_LORA_NAMES
@@ -491,13 +506,8 @@ class LoRAMemoryPool:
         _, output_dim = get_hidden_dim(
             module_name, self.base_hf_config, base_model, layer_idx
         )
-        # Same TP-vs-moe-TP sharding rule as get_lora_A_shape above.
+        # Same sharding rule as get_lora_A_shape above.
-        effective_tp_size = (
+        effective_tp_size = self._effective_tp_size(module_name)
-            self.tp_size
-            if not self.is_moe_module(module_name)
-            or self.is_shared_moe_module(module_name)
-            else self.moe_tp_size
-        )
         if (
             effective_tp_size > 1
             and module_name not in ROW_PARALLELISM_LINEAR_LORA_NAMES
@@ -1067,7 +1077,7 @@ class LoRAMemoryPool:
 
                 # Handle standard modules
                 temp_A_buffer[target_module] = module.slice_lora_a_weights(
-                    temp_A_buffer[target_module], self.tp_rank
+                    temp_A_buffer[target_module]
                 )
                 cache_keys = temp_A_cache_keys[target_module]
                 assert cache_keys is not None
@@ -1077,7 +1087,7 @@ class LoRAMemoryPool:
                 )
 
                 temp_B_buffer[target_module] = module.slice_lora_b_weights(
-                    temp_B_buffer[target_module], self.tp_rank
+                    temp_B_buffer[target_module]
                 )
                 cache_keys = temp_B_cache_keys[target_module]
                 assert cache_keys is not None
@@ -1413,7 +1423,7 @@ class LoRAMemoryPool:
                     # Slice B along vocab dimension for this TP rank
                     if self.tp_size > 1:
                         lora_b_weights = lora_lm_head_module.slice_lora_b_weights(
-                            lora_b_weights, self.tp_rank
+                            lora_b_weights
                         )
                         cache_key = append_cache_key_suffix(name, f"tp{self.tp_rank}")
                     else:

python/sglang/srt/lora/utils.py

@@ -340,6 +340,24 @@ REPLICATED_LINEAR_LORA_NAMES = [
     "fc2_latent_proj",
     *DSA_INDEXER_LORA_NAMES,
 ]
+# Attention-projection LoRA modules shard on the attention-TP group, which
+# under `--enable-dp-attention` is `attn_tp_size = tp_size // dp_size` rather
+# than the outer TP size. in_proj / in_proj_qkvz (linear-attention hybrids)
+# belong here too: their layers are built on the attn-TP group (mamba.py and
+# qwen3_5.py take tp_size/tp_rank from attn_tp when dp attention is enabled).
+ATTN_TP_LORA_MODULE_NAMES = frozenset(
+    {
+        "qkv_proj",
+        "qkvr",
+        "q_b_proj",
+        "kv_b_proj",
+        "o_proj",
+        "out_proj",
+        "wo_ud",
+        "in_proj",
+        "in_proj_qkvz",
+    }
+)
 
 # Normalized module names that the LoRA system fully supports
 # (i.e. get_hidden_dim, init_buffers, and init_lora_modules can handle them).

python/sglang/srt/managers/tokenizer_control_mixin.py

@@ -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
+                self.server_args.dp_size == 1 or self.server_args.enable_dp_attention
-            ), "dp_size must be 1 for dynamic lora loading"
+            ), "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
+                self.server_args.dp_size == 1 or self.server_args.enable_dp_attention
-            ), "dp_size must be 1 for dynamic lora loading"
+            ), "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,

test/registered/unit/lora/test_mem_pool_ep_unit.py

@@ -101,10 +101,10 @@ class _FakeRoutedMoeLayer(_FakeFusedMoEWithLoRA, _IdentityMoeSlices):
 
 
 class _FakeDenseLayer:
-    def slice_lora_a_weights(self, weights, _rank):
+    def slice_lora_a_weights(self, weights):
         return weights
 
-    def slice_lora_b_weights(self, weights, _rank):
+    def slice_lora_b_weights(self, weights):
         return weights
 
 
@@ -994,6 +994,7 @@ class TestPoolInitPicksUpEpContext(unittest.TestCase):
                 dtype=torch.bfloat16,
                 tp_size=tp_size,
                 tp_rank=tp_rank,
+                attn_tp_size=tp_size,
                 max_lora_rank=8,
                 target_modules={"qkv_proj"},
                 base_model=base_model,
@@ -1089,6 +1090,9 @@ def _fake_base_model_with_hidden_dim(num_experts: int) -> torch.nn.Module:
                 return cfg.hidden_size, cfg.moe_intermediate_size * 2
             if module_name == "down_proj_moe":
                 return cfg.moe_intermediate_size, cfg.hidden_size
+            if module_name == "in_proj_qkvz":
+                # linear-attention qkvz input projection (column-parallel)
+                return cfg.hidden_size, 4 * cfg.hidden_size
             raise NotImplementedError(module_name)
 
     return _Model()
@@ -1119,6 +1123,9 @@ class TestMoeBufferShardsByMoeTp(unittest.TestCase):
         pool.max_loras_per_batch = 2
         pool.tp_size = tp_size
         pool.tp_rank = 0
+        # Without --enable-dp-attention the attention TP group equals the
+        # outer TP group.
+        pool.attn_tp_size = tp_size
         pool.moe_ep_size = ep_size
         pool.moe_ep_rank = ep_rank
         pool.moe_tp_size = moe_tp_size
@@ -1218,6 +1225,76 @@ class TestMoeBufferShardsByMoeTp(unittest.TestCase):
         self.assertEqual(q_b, (2, 48, 8))
 
 
+class TestAttnModulesShardByAttnTp(unittest.TestCase):
+    """Regression: attention-module LoRA buffers must shard by `attn_tp_size`,
+    not the outer `tp_size`.
+
+    Under `--enable-dp-attention` attention layers are built on the attn_tp
+    group (`attn_tp_size = tp_size // dp_size`), so e.g. MLA `o_proj` holds an
+    attn_tp-local input shard. Sizing the LoRA buffer by the outer `tp_size`
+    would make it narrower than the slice produced by
+    `RowParallelLinearWithLoRA.slice_lora_a_weights` (which slices by the base
+    layer's attn_tp-local rank), failing the shape-match assert at load time.
+    """
+
+    def _pool(self, *, tp_size: int, attn_tp_size: int) -> LoRAMemoryPool:
+        pool = LoRAMemoryPool.__new__(LoRAMemoryPool)
+        pool.max_loras_per_batch = 2
+        pool.tp_size = tp_size
+        pool.tp_rank = 0
+        pool.attn_tp_size = attn_tp_size
+        pool.moe_ep_size = 1
+        pool.moe_ep_rank = 0
+        pool.moe_tp_size = tp_size
+        pool.moe_tp_rank = 0
+        pool.moe_use_local_expert_ids = False
+        pool._num_experts_local = 1
+        pool.experts_shared_outer_loras = False
+        pool.base_hf_config = types.SimpleNamespace(
+            hidden_size=64,
+            num_attention_heads=8,
+            num_key_value_heads=8,
+            head_dim=8,
+            intermediate_size=256,
+            moe_intermediate_size=192,
+        )
+        return pool
+
+    def test_attn_tp_1_keeps_attention_buffers_full_width(self):
+        """tp=4 with attn_tp=1 (dp-attention, dp=4): attention weights are
+        replicated across ranks, so the LoRA buffers must be full-width.
+        """
+        pool = self._pool(tp_size=4, attn_tp_size=1)
+        model = _fake_base_model_with_hidden_dim(num_experts=1)
+        # o_proj is row-parallel: A input_dim = head_dim*num_heads = 64,
+        # undivided under attn_tp=1 (pre-fix: 16).
+        self.assertEqual(pool.get_lora_A_shape("o_proj", model, 8, 0), (2, 8, 64))
+        # qkv_proj is column-parallel: B output_dim = 8 * 24 = 192,
+        # undivided under attn_tp=1 (pre-fix: 48).
+        self.assertEqual(pool.get_lora_B_shape("qkv_proj", model, 8, 0), (2, 192, 8))
+
+    def test_attn_tp_gt1_still_shards_attention_buffers(self):
+        """tp=4 with attn_tp=2 (dp=2): attention weights are sharded 2-way."""
+        pool = self._pool(tp_size=4, attn_tp_size=2)
+        model = _fake_base_model_with_hidden_dim(num_experts=1)
+        self.assertEqual(pool.get_lora_A_shape("o_proj", model, 8, 0), (2, 8, 32))
+        self.assertEqual(pool.get_lora_B_shape("qkv_proj", model, 8, 0), (2, 96, 8))
+
+    def test_linear_attention_in_proj_shards_by_attn_tp(self):
+        """Regression: in_proj_qkvz is built on the attn-TP group under
+        dp-attention (qwen3_5.py passes tp_rank=attn_tp_rank), but it was
+        classified as outer-TP, so with tp=4 / attn_tp=1 its LoRA B buffer
+        came out 4x narrower than the wrapper's attn_tp-local slice and
+        adapter load failed on the shape assert."""
+        pool = self._pool(tp_size=4, attn_tp_size=1)
+        model = _fake_base_model_with_hidden_dim(num_experts=1)
+        # column-parallel: B output_dim = 4*64 = 256, undivided under
+        # attn_tp=1 (pre-fix: divided by the outer tp=4 -> 64).
+        self.assertEqual(
+            pool.get_lora_B_shape("in_proj_qkvz", model, 8, 0), (2, 256, 8)
+        )
+
+
 class TestLoadBufferPassesMoeTpRankToSlice(unittest.TestCase):
     """Regression: `load_lora_weight_to_buffer` must hand `moe_tp_rank` (not
     the outer `tp_rank`) to `slice_moe_lora_{a,b}_weights`.

test/registered/unit/managers/test_lora_update_result_merge.py

@@ -0,0 +1,78 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Unit tests for merging per-rank LoRA update replies from the control fan-out."""
+
+import unittest
+
+from sglang.srt.managers.io_struct import LoRAUpdateOutput
+from sglang.srt.managers.tokenizer_control_mixin import _merge_lora_update_results
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=2, suite="base-a-test-cpu")
+
+
+def _ok(adapters=None) -> LoRAUpdateOutput:
+    return LoRAUpdateOutput(success=True, loaded_adapters=adapters or {})
+
+
+def _err(message, adapters=None) -> LoRAUpdateOutput:
+    return LoRAUpdateOutput(
+        success=False, error_message=message, loaded_adapters=adapters or {}
+    )
+
+
+class TestMergeLoRAUpdateResults(CustomTestCase):
+    def test_all_success_returns_first_rank_result(self):
+        """On success the merge must hand back a rank's own reply: callers
+        mutate result.loaded_adapters in place during LRU eviction, which a
+        synthesized empty result would silently break."""
+        results = [_ok({"a": "path"}), _ok({"a": "path"})]
+        merged = _merge_lora_update_results(results)
+        self.assertIs(merged, results[0])
+        self.assertTrue(merged.success)
+
+    def test_any_rank_failure_wins(self):
+        """Regression guard for the pre-merge behavior of returning
+        results[0]: a failure on a non-zero rank was reported as success,
+        letting the tokenizer-side registry drift from that rank's actual
+        adapter state."""
+        merged = _merge_lora_update_results(
+            [_ok({"a": "path"}), _err("out of memory", {"stale": "path"})]
+        )
+        self.assertFalse(merged.success)
+        self.assertEqual(merged.error_message, "out of memory")
+        self.assertEqual(merged.loaded_adapters, {"stale": "path"})
+
+    def test_duplicate_error_messages_deduplicated(self):
+        """All ranks usually fail identically (e.g. "already loaded"); the
+        joined message must not repeat per rank, but distinct causes must all
+        be kept."""
+        merged = _merge_lora_update_results(
+            [_err("already loaded"), _err("already loaded"), _err("bad rank")]
+        )
+        self.assertFalse(merged.success)
+        self.assertEqual(merged.error_message, "already loaded | bad rank")
+
+    def test_failure_without_message(self):
+        """A rank replying success=False with error_message=None must not
+        crash the join."""
+        merged = _merge_lora_update_results([_err(None), _ok()])
+        self.assertFalse(merged.success)
+        self.assertEqual(merged.error_message, "")
+
+
+if __name__ == "__main__":
+    unittest.main()