From d81fc6e7ccee865f56e901b35240d48fcdf22c5a Mon Sep 17 00:00:00 2001
From: Nan Jiang <59716405+nanjiangwill@users.noreply.github.com>
Date: Thu, 16 Jul 2026 17:43:48 -0700
Subject: [PATCH 1/2] [LoRA] DP-attention support for LoRA: idle-forward
 guards, attn-TP-local slicing, per-rank tensor serialization

Skip LoRA application on DP-attention idle forwards via a None batch_info
guard in every LoRA-wrapped layer, and clear stale batch info on idle batches.
Slice LoRA A/B weights by the base layer's attn_tp-local rank and size the
memory-pool buffers of attention modules by attn_tp_size.
Relax the dp_size == 1 asserts for dynamic LoRA load/unload when DP attention
is enabled, and ship load_lora_adapter_from_tensors payloads as one serialized
copy per TP rank, mirroring update_weights_from_tensor.
Derive CUDA-graph MoE LoRA buffer dims from the base FusedMoE so all quant
schemes work.

Ported from the sglang-miles branch (original PR #25141).
---
 python/sglang/srt/entrypoints/engine.py       | 15 ++--
 .../sglang/srt/lora/backend/base_backend.py   | 15 ++--
 python/sglang/srt/lora/layers.py              | 73 ++++++++++++++-----
 python/sglang/srt/lora/lora_manager.py        | 10 +++
 python/sglang/srt/lora/mem_pool.py            | 47 +++++++++---
 python/sglang/srt/lora/utils.py               | 10 +++
 python/sglang/srt/managers/io_struct.py       |  5 +-
 .../srt/managers/tokenizer_control_mixin.py   | 20 ++---
 python/sglang/srt/managers/tp_worker.py       | 13 ++--
 .../srt/model_executor/forward_batch_info.py  |  4 +
 .../rl/test_lora_load_from_tensor.py          |  4 +-
 .../unit/lora/test_mem_pool_ep_unit.py        | 59 +++++++++++++++
 .../managers/test_msgpack_ipc_roundtrip.py    |  2 +-
 13 files changed, 218 insertions(+), 59 deletions(-)

diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py
index 7fd810c4708b..8aeeda503499 100644
--- a/python/sglang/srt/entrypoints/engine.py
+++ b/python/sglang/srt/entrypoints/engine.py
@@ -1138,20 +1138,23 @@ def get_weights_by_name(self, name: str, truncate_size: int = 100):
     def load_lora_adapter_from_tensors(
         self,
         lora_name: str,
-        tensors,
+        tensors: Union[Dict[str, torch.Tensor], List[SerializedTensorPayload]],
         config_dict: Dict,
         load_format: Optional[str] = None,
     ):
         if load_format == "flattened_bucket":
-            serialized_tensors = tensors
-        else:
-            serialized_tensors = MultiprocessingSerializer.serialize(
-                tensors, output_str=True
+            serialized_named_tensors = normalize_serialized_named_tensor_payloads(
+                cast(List[SerializedTensorPayload], tensors)
             )
+        else:
+            serialized_named_tensors = [
+                MultiprocessingSerializer.serialize(tensors)
+                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(
diff --git a/python/sglang/srt/lora/backend/base_backend.py b/python/sglang/srt/lora/backend/base_backend.py
index 16879b547435..74562559d4b1 100644
--- a/python/sglang/srt/lora/backend/base_backend.py
+++ b/python/sglang/srt/lora/backend/base_backend.py
@@ -1,4 +1,4 @@
-from typing import Tuple, Union
+from typing import Optional, Tuple, Union
 
 import torch
 import triton
@@ -22,6 +22,10 @@ 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
+        # Set by prepare_lora_batch() before each forward. Stays None on
+        # DP-attention idle forwards, which skip batch preparation, so the
+        # LoRA layers use it to skip LoRA application there.
+        self.batch_info: Optional[LoRABatchInfo] = None
         self.init_lm_head_config()
         self._is_moe_lora = False
 
@@ -188,10 +192,11 @@ 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.).
+        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
 
diff --git a/python/sglang/srt/lora/layers.py b/python/sglang/srt/lora/layers.py
index f0032b80b606..a79ba7abee46 100644
--- a/python/sglang/srt/lora/layers.py
+++ b/python/sglang/srt/lora/layers.py
@@ -45,6 +45,12 @@ 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
+        # Forward reduce_results so model code that inspects it on the module
+        # (e.g. DeepseekV2AttentionMLA's `assert not self.o_proj.reduce_results`
+        # on DP-attention idle forwards) keeps working when the layer is
+        # LoRA-wrapped.
+        if hasattr(self.base_layer, "reduce_results"):
+            self.reduce_results = self.base_layer.reduce_results
 
     def forward(self, x: torch.Tensor):
         return self.base_layer.forward(x)
@@ -211,8 +217,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 when batch_info is None (DP-attention
+        # idle forward): no real tokens need LoRA.
+        if self.set_lora and 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)
@@ -377,8 +384,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 (no batch_info).
+        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
@@ -467,7 +474,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:
@@ -481,9 +488,14 @@ 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):
+        # Slice by the base layer's own TP rank, not the caller's outer
+        # tp_rank: under DP-attention, attention layers are built on the
+        # attn_tp group, so output_partition_sizes is attn_tp-local and the
+        # outer tp_rank would overshoot the weight.
+        local_tp_rank = self.base_layer.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
 
@@ -577,12 +589,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):
+        # See ColumnParallelLinearWithLoRA.slice_lora_b_weights for why the
+        # base layer's own TP rank is authoritative under DP-attention.
+        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
@@ -649,11 +664,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 ColumnParallelLinearWithLoRA.slice_lora_b_weights for why the
+        # base layer's own TP rank is authoritative under DP-attention.
+        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
 
@@ -736,12 +754,22 @@ def forward(self, input_: torch.Tensor, skip_all_reduce=False, forward_batch=Non
             and not should_skip_mlp_all_reduce()
         )
 
-        if self.set_lora and should_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 is 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
             )
-            output_ = tensor_model_parallel_all_reduce(output_parallel)
-            lora_a_output = tensor_model_parallel_all_reduce(lora_a_output)
+            output_ = all_reduce(output_parallel)
+            lora_a_output = all_reduce(lora_a_output)
             output_ = self.lora_backend.run_lora_b_sgemm(
                 x=lora_a_output,
                 weights=self.B_buffer,
@@ -750,10 +778,10 @@ 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)
+                output_ = all_reduce(output_parallel)
             else:
                 output_ = output_parallel
 
@@ -761,9 +789,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):
+        # Slice by the base layer's own TP rank, not the caller's outer
+        # tp_rank: under DP-attention, attention layers (e.g. MLA o_proj) are
+        # built on the attn_tp group, so input_size_per_partition is already
+        # attn_tp-sized and the outer tp_rank would overshoot to an empty
+        # slice.
+        local_tp_rank = self.base_layer.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
 
@@ -849,7 +883,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
@@ -1026,6 +1060,9 @@ 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
         """
+        # DP-attention idle forward: no batch_info, run the base MoE path.
+        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()
diff --git a/python/sglang/srt/lora/lora_manager.py b/python/sglang/srt/lora/lora_manager.py
index a83ca071f956..902a49f0c99c 100644
--- a/python/sglang/srt/lora/lora_manager.py
+++ b/python/sglang/srt/lora/lora_manager.py
@@ -375,6 +375,16 @@ def prepare_lora_batch(self, forward_batch: ForwardBatch):
             lora_ranks[wi] > 0 for wi in weight_indices
         )
 
+    def prepare_idle_lora_batch(self):
+        """Reset per-batch LoRA state for a DP-attention idle forward.
+
+        Idle forwards run the model with zero local tokens and skip
+        prepare_lora_batch(); clearing batch_info makes the LoRA layers fall
+        back to the base path instead of reading the previous batch's stale
+        metadata.
+        """
+        self.lora_backend.batch_info = None
+
     def update_lora_info(self):
         """
         Update all LoRA modules to associate them with the latest memory buffer.
diff --git a/python/sglang/srt/lora/mem_pool.py b/python/sglang/srt/lora/mem_pool.py
index 4b475c6aedd4..188f02b67a42 100644
--- a/python/sglang/srt/lora/mem_pool.py
+++ b/python/sglang/srt/lora/mem_pool.py
@@ -26,6 +26,7 @@
 from sglang.srt.lora.lora_config import LoRAConfig
 from sglang.srt.lora.lora_registry import LoRARef
 from sglang.srt.lora.utils import (
+    ATTENTION_LINEAR_LORA_NAMES,
     EMBEDDING_NAMES,
     REPLICATED_LINEAR_LORA_NAMES,
     ROW_PARALLELISM_LINEAR_LORA_NAMES,
@@ -109,6 +110,17 @@ def _get_moe_tp_context() -> Tuple[int, int]:
         return 1, 0
 
 
+def _get_attn_tp_size(tp_size: int) -> int:
+    """Return `attn_tp_size`, or the outer `tp_size` if the attention TP
+    group is not initialized. Under `--enable-dp-attention` attention weights
+    (e.g. MLA `o_proj`) are sharded along `attn_tp_size = tp_size // dp_size`
+    instead of the outer `tp_size`."""
+    try:
+        return get_parallel().attn_tp_size
+    except Exception:  # pragma: no cover - attention TP group not initialized
+        return tp_size
+
+
 def _moe_runner_keeps_global_expert_ids() -> bool:
     """True if the active MoE runner keeps global `topk_ids` instead of
     remapping to local IDs. Mirrors the predicate in `StandardDispatcher`."""
@@ -181,10 +193,15 @@ def __init__(
         # here would yield a 4x-narrower inner dim than the adapter weight
         # (which `FusedMoEWithLoRA.slice_moe_lora_{a,b}_weights` correctly
         # skip-slices when `moe_tp_size <= 1`), producing a shape-mismatch
-        # assert during weight load. Non-MoE modules still shard by
-        # `tp_size` because attention TP is unchanged.
+        # assert during weight load.
         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 = _get_attn_tp_size(tp_size)
+
         # Initialize eviction policy
         self.eviction_policy = get_eviction_policy(eviction_policy)
 
@@ -254,6 +271,17 @@ def is_moe_module(self, module_name: str) -> bool:
         """Check if module is part of MoE experts."""
         return "moe" in module_name
 
+    def _effective_tp_size(self, module_name: str) -> int:
+        """TP width the module's weights are actually sharded along: MoE
+        experts shard by `moe_tp_size`, 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):
+            return self.moe_tp_size
+        if module_name in ATTENTION_LINEAR_LORA_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
@@ -342,8 +370,9 @@ def _get_standard_shape(
             module_name, self.base_hf_config, base_model, layer_idx
         )
         c = get_stacked_multiply(module_name, base_model)
-        if self.tp_size > 1 and module_name in ROW_PARALLELISM_LINEAR_LORA_NAMES:
-            input_dim = divide(input_dim, self.tp_size)
+        effective_tp_size = self._effective_tp_size(module_name)
+        if effective_tp_size > 1 and module_name in ROW_PARALLELISM_LINEAR_LORA_NAMES:
+            input_dim = divide(input_dim, effective_tp_size)
         return (self.max_loras_per_batch, max_lora_dim * c, input_dim)
 
     def get_lora_A_shape(
@@ -364,10 +393,7 @@ def get_lora_A_shape(
             module_name, self.base_hf_config, base_model, layer_idx
         )
         c = get_stacked_multiply(module_name, base_model)
-        # MoE modules shard along `moe_tp_size`, not the outer `tp_size`.
-        effective_tp_size = (
-            self.moe_tp_size if self.is_moe_module(module_name) else self.tp_size
-        )
+        effective_tp_size = self._effective_tp_size(module_name)
         if (
             effective_tp_size > 1
             and module_name in ROW_PARALLELISM_LINEAR_LORA_NAMES
@@ -458,10 +484,7 @@ def get_lora_B_shape(
         _, output_dim = get_hidden_dim(
             module_name, self.base_hf_config, base_model, layer_idx
         )
-        # MoE modules shard along `moe_tp_size`, not the outer `tp_size`.
-        effective_tp_size = (
-            self.moe_tp_size if self.is_moe_module(module_name) else self.tp_size
-        )
+        effective_tp_size = self._effective_tp_size(module_name)
         if (
             effective_tp_size > 1
             and module_name not in ROW_PARALLELISM_LINEAR_LORA_NAMES
diff --git a/python/sglang/srt/lora/utils.py b/python/sglang/srt/lora/utils.py
index e8301d5b416c..1d7d2f139366 100644
--- a/python/sglang/srt/lora/utils.py
+++ b/python/sglang/srt/lora/utils.py
@@ -323,6 +323,16 @@ def get_target_module_name(full_module_name: str, target_modules: Set[str]) -> s
 
 EMBEDDING_NAMES = ["embed_tokens", "lm_head"]
 ROW_PARALLELISM_LINEAR_LORA_NAMES = ["o_proj", "out_proj", "down_proj", "down_proj_moe"]
+# Attention projections shard along the attention TP group, which under
+# `--enable-dp-attention` is `attn_tp_size = tp_size // dp_size` rather than
+# the outer `tp_size`.
+ATTENTION_LINEAR_LORA_NAMES = [
+    "qkv_proj",
+    "o_proj",
+    "out_proj",
+    "q_b_proj",
+    "kv_b_proj",
+]
 DSA_INDEXER_LORA_NAMES = frozenset(
     {"indexer.wq_b", "indexer.wk", "indexer.weights_proj"}
 )
diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py
index 6f900c7e0850..6ada5dc33aca 100644
--- a/python/sglang/srt/managers/io_struct.py
+++ b/python/sglang/srt/managers/io_struct.py
@@ -2036,7 +2036,10 @@ class LoadLoRAAdapterFromTensorsReqInput(BaseReq, kw_only=True):
     # The PEFT adapter_config.json, already JSON — a tighter type would only add
     # decode strictness with no benefit.
     config_dict: Dict[str, Any]
-    serialized_tensors: str
+    # One serialized copy of the adapter tensors per TP rank; each rank
+    # deserializes only its own copy. Same normalization conventions as
+    # UpdateWeightsFromTensorReqInput.serialized_named_tensors.
+    serialized_named_tensors: Annotated[List[bytes], Base64Bytes()]
     pinned: bool = False
     added_tokens_config: Optional[Dict[str, int]] = None
     lora_id: Optional[str] = None
diff --git a/python/sglang/srt/managers/tokenizer_control_mixin.py b/python/sglang/srt/managers/tokenizer_control_mixin.py
index 86b7b378f896..e3773542c446 100644
--- a/python/sglang/srt/managers/tokenizer_control_mixin.py
+++ b/python/sglang/srt/managers/tokenizer_control_mixin.py
@@ -574,11 +574,9 @@ async def load_lora_adapter(
                     "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,
@@ -653,13 +651,17 @@ async def load_lora_adapter_from_tensors(
                 )
 
             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 from tensors. Lora name=%s",
                 obj.lora_name,
             )
 
+            obj.serialized_named_tensors = normalize_serialized_named_tensor_payloads(
+                obj.serialized_named_tensors
+            )
+
             async with self.lora_update_lock:
                 new_adapter = LoRARef(
                     lora_name=obj.lora_name,
@@ -726,11 +728,9 @@ async def unload_lora_adapter(
                 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,
diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py
index eee3e4fb1e44..0888555d8339 100644
--- a/python/sglang/srt/managers/tp_worker.py
+++ b/python/sglang/srt/managers/tp_worker.py
@@ -209,18 +209,21 @@ def load_lora_adapter_from_tensors(
         self, recv_req: LoadLoRAAdapterFromTensorsReqInput
     ):
         # The LoRA code handles TP sharding internally using slice_lora_a_weights
-        # and slice_lora_b_weights methods (see lora/layers.py:46-49, mem_pool.py:437-440).
+        # and slice_lora_b_weights methods (see lora/layers.py and mem_pool.py).
+        # Each rank deserializes its own serialized copy — same convention as
+        # update_weights_from_tensor above — so the CUDA-IPC ref count on the
+        # producer side drops cleanly after every load.
+        monkey_patch_torch_reductions()
+        serialized = recv_req.serialized_named_tensors[self.ps.tp_rank]
         if recv_req.load_format == "flattened_bucket":
-            flattened_data = MultiprocessingSerializer.deserialize(
-                recv_req.serialized_tensors
-            )
+            flattened_data = MultiprocessingSerializer.deserialize(serialized)
             bucket = FlattenedTensorBucket(
                 flattened_tensor=flattened_data["flattened_tensor"],
                 metadata=flattened_data["metadata"],
             )
             tensors = dict(bucket.reconstruct_tensors())
         else:
-            tensors = MultiprocessingSerializer.deserialize(recv_req.serialized_tensors)
+            tensors = MultiprocessingSerializer.deserialize(serialized)
         result = self.model_runner.load_lora_adapter_from_tensors(
             recv_req.to_ref(),
             tensors,
diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py
index 85204c0d0b60..e57caaefde9a 100644
--- a/python/sglang/srt/model_executor/forward_batch_info.py
+++ b/python/sglang/srt/model_executor/forward_batch_info.py
@@ -808,6 +808,10 @@ def init_new(
 
         if ret.forward_mode.is_idle():
             ret.positions = torch.empty((0,), dtype=torch.int64, device=device)
+            # Idle forwards (DP attention) skip the LoRA batch preparation
+            # below; clear stale batch info so LoRA layers skip application.
+            if model_runner.server_args.enable_lora:
+                model_runner.lora_manager.prepare_idle_lora_batch()
             return ret
 
         # Override the positions with diffusion LLM or spec_info
diff --git a/test/registered/rl/test_lora_load_from_tensor.py b/test/registered/rl/test_lora_load_from_tensor.py
index d2505215a389..c37aab68b00b 100644
--- a/test/registered/rl/test_lora_load_from_tensor.py
+++ b/test/registered/rl/test_lora_load_from_tensor.py
@@ -342,9 +342,11 @@ def test_lora_e2e_load_from_flattened_bucket(self):
         }
         serialized = MultiprocessingSerializer.serialize(bucket_dict, output_str=True)
 
+        # flattened_bucket callers pass one serialized copy per TP rank, same
+        # as Engine.update_weights_from_tensor.
         result = self.engine.load_lora_adapter_from_tensors(
             lora_name="self_cognition_Alice_flattened",
-            tensors=serialized,
+            tensors=[serialized],
             config_dict=self.lora_config_dict,
             load_format="flattened_bucket",
         )
diff --git a/test/registered/unit/lora/test_mem_pool_ep_unit.py b/test/registered/unit/lora/test_mem_pool_ep_unit.py
index 3f9b6dfc4a19..db1c2d977941 100644
--- a/test/registered/unit/lora/test_mem_pool_ep_unit.py
+++ b/test/registered/unit/lora/test_mem_pool_ep_unit.py
@@ -530,6 +530,9 @@ def _pool(
         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
@@ -629,6 +632,62 @@ def test_non_moe_modules_unaffected_by_moe_tp(self):
         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))
+
+
 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`.
diff --git a/test/registered/unit/managers/test_msgpack_ipc_roundtrip.py b/test/registered/unit/managers/test_msgpack_ipc_roundtrip.py
index e3d3e95d067b..0dabbe33023d 100644
--- a/test/registered/unit/managers/test_msgpack_ipc_roundtrip.py
+++ b/test/registered/unit/managers/test_msgpack_ipc_roundtrip.py
@@ -126,7 +126,7 @@ def _checksum_info(tag: str) -> ChecksumInfo:
     "LoadLoRAAdapterFromTensorsReqInput": LoadLoRAAdapterFromTensorsReqInput(
         lora_name="adapter",
         config_dict={"r": 8, "lora_alpha": 16, "target_modules": ["q_proj", "v_proj"]},
-        serialized_tensors="",
+        serialized_named_tensors=[b"tp0-bytes", b"tp1-bytes"],
         added_tokens_config={"<extra>": 32000},
     ),
     "DumperControlReqInput": DumperControlReqInput(method="start", body={"k": "v"}),

From 150464d13481b1304409d360ff11fbbf6dfc87e7 Mon Sep 17 00:00:00 2001
From: "Ethan (Yusheng) Su" <yushengsu@radixark.ai>
Date: Thu, 16 Jul 2026 18:24:09 -0700
Subject: [PATCH 2/2] [LoRA] MoE-LoRA cuda-graph fixes under DP-attention + RL
 adapter-reload fixes

Size MoE-LoRA routing buffers for the DP-gathered token count via
get_gathered_moe_num_tokens (cuda-graph buffers of max_bs*attn_dp_size), demote
to eager prep when the gathered length exceeds them, and stamp the gathered
tail with the single loaded adapter (idle-rank / active-rank stamps).
Free the memory-pool slot on unload for colocate-RL per-step adapter reloads,
fix per-expert A/B init under experts_shared_outer_loras, drop the unsafe
virtual-experts routing trim, and allow DeepSeek V4 in get_dsa_index_n_heads.

Ported from the sglang-miles branch (original PR #29874).
---
 .../sglang/kernels/ops/moe/virtual_experts.py |  33 +++--
 python/sglang/srt/configs/model_config.py     |   5 +-
 .../sglang/srt/lora/backend/base_backend.py   | 128 +++++++++++++++--
 python/sglang/srt/lora/lora_manager.py        |  87 ++++++++++--
 python/sglang/srt/lora/mem_pool.py            |  65 ++++++++-
 .../srt/model_executor/forward_batch_info.py  |  10 +-
 .../unit/lora/test_moe_lora_tail_stamp.py     | 134 ++++++++++++++++++
 7 files changed, 414 insertions(+), 48 deletions(-)
 create mode 100644 test/registered/unit/lora/test_moe_lora_tail_stamp.py

diff --git a/python/sglang/kernels/ops/moe/virtual_experts.py b/python/sglang/kernels/ops/moe/virtual_experts.py
index 23fc8355b6dc..e09435485578 100644
--- a/python/sglang/kernels/ops/moe/virtual_experts.py
+++ b/python/sglang/kernels/ops/moe/virtual_experts.py
@@ -614,16 +614,18 @@ def _get_routing(
             block_size=block_size,
             num_experts=virtual_num_experts,
         )
-        # _align_block_size uses a worst-case padded allocation. Trim the routing buffers
-        # to a tighter upper bound so we keep the real routed work but drop unused padding
-        num_tokens = topk_ids.numel()
-        max_nonempty = min(num_tokens, virtual_num_experts)
-        tight_padded = (
-            triton.cdiv(num_tokens + max_nonempty * (block_size - 1), block_size)
-            * block_size
-        )
-        sorted_token_ids = sorted_token_ids[:tight_padded]
-        expert_ids = expert_ids[: tight_padded // block_size]
+        # NOTE: do NOT trim sorted_token_ids / expert_ids to a tighter upper bound here.
+        # The downstream kernels (_moe_lora_shrink_splitk_kernel, fused_moe_kernel) read
+        # sorted_token_ids[pid_m*BLOCK : +BLOCK] and expert_ids[pid_m] WITHOUT a bounds mask
+        # for every block up to num_tokens_post_padded (a GPU-side count loaded at run time).
+        # num_tokens_post_padded comes from _align_block_size with `virtual_num_experts` buckets
+        # and can exceed a tighter `numel + min(numel,virtual_num_experts)*(block-1)` bound
+        # (most so for shared-outer, where virtual_num_experts = max_loras is small), so trimming
+        # made those unmasked reads land PAST the view. In eager mode the slack still lives inside
+        # the same _align_block_size allocation (garbage, masked out downstream) so it worked; under
+        # CUDA-graph capture/replay the graph mempool packs tensors tightly and that slack may belong
+        # to another pooled tensor / lie past a page -> cudaErrorIllegalInstruction during capture.
+        # Keep the full worst-case-allocated buffers so every unmasked read stays in-allocation.
         expert_ids = fused_sanitize_expert_ids(expert_ids, virtual_num_experts)
         result = (
             sorted_token_ids,
@@ -646,8 +648,17 @@ def _get_routing(
     num_experts_a = lora_a.shape[1]
     num_experts_b = lora_b.shape[1]
 
+    # The kernels index token_lora_mapping / intermediate by token ids up to
+    # topk_ids.shape[0] (the DP-gathered token count under --enable-dp-attention). An
+    # under-sized mapping means unmasked OOB reads/writes that surface as a sticky,
+    # hard-to-attribute CUDA IMA — fail loudly on the host instead.
+    assert token_lora_mapping.shape[0] >= topk_ids.shape[0], (
+        f"token_lora_mapping covers {token_lora_mapping.shape[0]} tokens but the MoE runs on "
+        f"{topk_ids.shape[0]} (DP-gathered?) tokens; mapping was sized before the dp gather "
+        f"length was known (see get_gathered_moe_num_tokens)"
+    )
     intermediate = torch.zeros(
-        [token_lora_mapping.shape[0], topk_ids.shape[1], max_lora_rank],
+        [topk_ids.shape[0], topk_ids.shape[1], max_lora_rank],
         dtype=hidden_states.dtype,
         device=hidden_states.device,
     )
diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py
index 9f3ea8f7c735..a52f69176e15 100644
--- a/python/sglang/srt/configs/model_config.py
+++ b/python/sglang/srt/configs/model_config.py
@@ -207,7 +207,10 @@ def dsa_layer_skips_topk(config: PretrainedConfig, layer_id: int) -> bool:
 
 
 def get_dsa_index_n_heads(config: PretrainedConfig) -> int:
-    assert is_deepseek_dsa(config)
+    # Permit both DSA (V3.2-family) and V4: both carry the indexer (index_n_heads) and this must
+    # match get_dsa_index_head_dim's contract, else LoRA buffer init for indexer.wq_b /
+    # indexer.weights_proj on a V4 model asserts here while indexer.wk (which uses head_dim) succeeds.
+    assert is_deepseek_dsa(config) or is_deepseek_v4(config)
     return config.index_n_heads
 
 
diff --git a/python/sglang/srt/lora/backend/base_backend.py b/python/sglang/srt/lora/backend/base_backend.py
index 74562559d4b1..dc9d3ec3a6e4 100644
--- a/python/sglang/srt/lora/backend/base_backend.py
+++ b/python/sglang/srt/lora/backend/base_backend.py
@@ -7,6 +7,37 @@
 from sglang.srt.lora.backend.lmhead_mixing import LoRABackendLmHeadMixing
 from sglang.srt.lora.utils import LoRABatchInfo, MoELoRABatchInfo
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.runtime_context import get_parallel
+from sglang.srt.utils.common import ceil_align
+
+
+def get_gathered_moe_num_tokens(forward_batch: ForwardBatch, num_tokens: int) -> int:
+    """Token count the MoE-LoRA mapping must cover: gathered under --enable-dp-attention, else per-rank.
+
+    In the eager path prepare_lora_batch runs from ForwardBatch.init_new, BEFORE
+    prepare_mlp_sync_batch assigns forward_batch.global_dp_buffer_len (only the cuda-graph
+    capture path pre-sets it), so that field alone under-sizes the MoE-LoRA token mapping to the
+    per-rank length and the MoE-LoRA kernels index past it (sticky CUDA IMA). When it is unset,
+    derive an upper bound of the gathered length from global_num_tokens_cpu (assigned in
+    init_new before prepare_lora_batch runs), mirroring prepare_mlp_sync_batch's attn-tp/cp
+    alignment; max*n covers both SUM_LEN and MAX_LEN padding modes. Over-allocation is harmless:
+    the kernels index at most the actual gathered length.
+    """
+    if forward_batch.global_dp_buffer_len is not None:
+        return max(forward_batch.global_dp_buffer_len, num_tokens)
+    global_num_tokens = forward_batch.global_num_tokens_cpu
+    if not global_num_tokens:
+        return num_tokens
+    # Local import: a module-level cp_utils import here is circular (see forward_batch_info).
+    from sglang.srt.layers.utils.cp_utils import get_cp_padding_align_size
+
+    attn_tp_size = get_parallel().attn_tp_size
+    cp_align_size = get_cp_padding_align_size()
+    upper = max(
+        ceil_align(ceil_align(t, attn_tp_size), cp_align_size)
+        for t in global_num_tokens
+    ) * len(global_num_tokens)
+    return max(upper, num_tokens)
 
 
 class BaseLoRABackend(LoRABackendLmHeadMixing):
@@ -22,10 +53,16 @@ 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
-        # Set by prepare_lora_batch() before each forward. Stays None on
-        # DP-attention idle forwards, which skip batch preparation, so the
-        # LoRA layers use it to skip LoRA application there.
+        # Set by prepare_lora_batch() before each forward. Stays None until
+        # the first batch is prepared, so the LoRA layers use it to skip LoRA
+        # application before then.
         self.batch_info: Optional[LoRABatchInfo] = None
+        # MoE-LoRA cuda-graph capture buffers; set by init_cuda_graph_moe_buffers().
+        self.moe_cg_buffers: Optional[dict] = None
+        # Tier-1 single-adapter tail stamps (host ints), armed per batch by
+        # LoRAManager.prepare_lora_batch and applied in _add_moe_lora_info.
+        self._idle_rank_active_buffer_id: Optional[int] = None
+        self._single_loaded_buffer_id: Optional[int] = None
         self.init_lm_head_config()
         self._is_moe_lora = False
 
@@ -200,8 +237,18 @@ def init_cuda_graph_moe_buffers(
         dtype = compute_dtype
         num_experts = base.num_experts
 
+        # Under --enable-dp-attention the MoE runs on DP-GATHERED tokens (the global DP buffer of
+        # length up to max_bs * attn_dp_size), not the per-rank batch. The per-token LoRA routing
+        # buffers below are indexed by that gathered token count, so size them for the gathered
+        # maximum; otherwise the MoE-LoRA kernels read token_lora_mapping / sorted_token_ids past a
+        # per-rank-sized buffer -> cudaErrorIllegalInstruction during cuda-graph capture. Expert- and
+        # adapter-indexed buffers (cumsum_buffer, adapter_enabled, lora_ids) are unaffected.
+        max_moe_tokens = max_bs * max(1, get_parallel().attn_dp_size)
+
         block_size_m = 64
-        max_num_tokens_padded = max_bs * top_k + num_experts * (block_size_m - 1)
+        max_num_tokens_padded = max_moe_tokens * top_k + num_experts * (
+            block_size_m - 1
+        )
         max_num_tokens_padded = (
             (max_num_tokens_padded + block_size_m - 1) // block_size_m
         ) * block_size_m
@@ -238,7 +285,7 @@ def init_cuda_graph_moe_buffers(
             # LongTensor.  weight_indices itself must stay int32 because the
             # CUDA moe_lora_align kernel casts it to int32_t*.
             "weight_indices_long": torch.zeros(
-                max_bs, dtype=torch.int64, device=device
+                max_moe_tokens, dtype=torch.int64, device=device
             ),
             "lora_ids": torch.arange(max_loras, dtype=torch.int32, device=device),
             "cumsum_buffer": torch.zeros(
@@ -247,14 +294,14 @@ def init_cuda_graph_moe_buffers(
                 device=device,
             ),
             "token_mask": torch.empty(
-                (max_loras * max_bs * top_k,),
+                (max_loras * max_moe_tokens * top_k,),
                 dtype=torch.int32,
                 device=device,
             ),
             "max_num_tokens_padded": max_num_tokens_padded,
             "max_num_m_blocks": max_num_m_blocks,
             "token_lora_mapping": torch.full(
-                (max_bs,), -1, dtype=torch.int32, device=device
+                (max_moe_tokens,), -1, dtype=torch.int32, device=device
             ),
         }
 
@@ -296,6 +343,19 @@ def _add_moe_lora_info(
             seg_indptr = batch_info.seg_indptr[: num_moe_segments + 1]
             req_to_lora = batch_info.weight_indices[:num_moe_segments]
 
+        # --enable-dp-attention all-gathers tokens into the MoE, so the MoE-LoRA kernels index
+        # token_lora_mapping by the GATHERED token count, not the per-rank num_tokens. Size the
+        # mapping to (an upper bound of) the gathered count so those reads stay in-bounds — see
+        # get_gathered_moe_num_tokens for why global_dp_buffer_len alone is NOT enough in the
+        # eager path (the per-rank segments still fill only [0, num_tokens); the tail stays -1).
+        moe_num_tokens = get_gathered_moe_num_tokens(forward_batch, num_tokens)
+        if batch_info.use_cuda_graph:
+            # Static capture buffers hold max_bs*dp tokens; a REAL replay's gathered length never
+            # exceeds that (the captured graph could not address it), so cap the upper bound at
+            # the buffer size. Batches whose gathered bound exceeds it are demoted to the eager
+            # prep path in LoRAManager.prepare_lora_batch before we get here.
+            moe_num_tokens = min(moe_num_tokens, token_lora_mapping.shape[0])
+
         adapter_enabled, token_lora_mapping = _compute_moe_lora_info(
             num_tokens,
             seg_indptr,
@@ -304,8 +364,34 @@ def _add_moe_lora_info(
             adapter_enabled,
             token_lora_mapping,
             max_len=max_len,
+            mapping_len=moe_num_tokens,
         )
 
+        # Tier-1 (colocate RL, exactly one adapter loaded): the DP-gathered tail
+        # [num_tokens, moe_num_tokens) covers OTHER dp ranks' tokens, which under colocate RL all
+        # use that single adapter. The per-rank fill above only wrote [0, num_tokens), leaving the
+        # tail -1 (adapter-disabled) -> cross-rank gathered tokens would miss the LoRA delta on
+        # this rank's local experts. The stamps below are ARMED BY LoRAManager.prepare_lora_batch
+        # (policy lives there; both are host ints, no GPU sync -> cuda-graph safe, written each
+        # batch in the eager prep path):
+        #   * _single_loaded_buffer_id: armed only when exactly one adapter is loaded AND this
+        #     rank's local requests actively use it -> stamp the tail with that adapter.
+        #   * _idle_rank_active_buffer_id: armed only on a true idle rank (num_tokens == 0) ->
+        #     _compute_moe_lora_info left the whole mapping -1 / adapter_enabled all-0, so stamp
+        #     the whole gathered buffer and enable the adapter.
+        # Multi-adapter batches and base-only local batches arm neither stamp and keep the -1
+        # tail: foreign tokens get base rather than a delta this rank cannot attribute.
+        if moe_num_tokens > num_tokens:
+            if num_tokens > 0:
+                single_bid = self._single_loaded_buffer_id
+                if single_bid is not None:
+                    token_lora_mapping[num_tokens:moe_num_tokens].fill_(single_bid)
+            else:
+                idle_bid = self._idle_rank_active_buffer_id
+                if idle_bid is not None:
+                    token_lora_mapping.fill_(idle_bid)
+                    adapter_enabled[idle_bid] = 1
+
         batch_info.moe_lora_info = MoELoRABatchInfo(
             seg_indptr=seg_indptr,
             req_to_lora=req_to_lora,
@@ -378,16 +464,28 @@ def _compute_moe_lora_info(
     adapter_enabled: torch.Tensor | None,
     token_lora_mapping: torch.Tensor | None,
     max_len: int,
+    mapping_len: int | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
+    # ``num_tokens`` is the PER-RANK fill count (segments cover this rank's tokens). ``mapping_len``
+    # is the length of the token_lora_mapping the MoE-LoRA kernels actually index -- under
+    # --enable-dp-attention the MoE runs on DP-GATHERED tokens (mapping_len = global_dp_buffer_len
+    # >= num_tokens), so the returned mapping must span the gathered count to keep those kernels
+    # in-bounds. The DP-gathered tail [num_tokens, mapping_len) defaults to -1 (adapter-disabled).
+    if mapping_len is None:
+        mapping_len = num_tokens
+    assert mapping_len >= num_tokens
     if token_lora_mapping is not None:
         assert (
-            num_tokens <= token_lora_mapping.shape[0]
-        ), "num_tokens must be less than or equal to the shape of token_lora_mapping"
-        token_lora_mapping = token_lora_mapping[:num_tokens]
+            mapping_len <= token_lora_mapping.shape[0]
+        ), "mapping_len must be less than or equal to the shape of token_lora_mapping"
+        token_lora_mapping = token_lora_mapping[:mapping_len]
     else:
         token_lora_mapping = torch.empty(
-            (num_tokens,), dtype=torch.int32, device=seg_indptr.device
+            (mapping_len,), dtype=torch.int32, device=seg_indptr.device
         )
+    if mapping_len > num_tokens:
+        # clean the gathered tail before the per-rank fill writes [0, num_tokens)
+        token_lora_mapping.fill_(-1)
 
     if adapter_enabled is not None:
         assert (
@@ -445,8 +543,12 @@ def _compute_moe_lora_info(
         torch.searchsorted(seg_indptr.to(torch.int32), token_positions, right=True) - 1
     )
 
-    token_lora_mapping = torch.index_select(
-        weight_indices.to(torch.int32), 0, req_indices, out=token_lora_mapping
+    # Fill only the per-rank prefix [0, num_tokens); the gathered tail keeps the -1 set above.
+    torch.index_select(
+        weight_indices.to(torch.int32),
+        0,
+        req_indices,
+        out=token_lora_mapping[:num_tokens],
     )
 
     return adapter_enabled, token_lora_mapping
diff --git a/python/sglang/srt/lora/lora_manager.py b/python/sglang/srt/lora/lora_manager.py
index 902a49f0c99c..a92058e6c0f3 100644
--- a/python/sglang/srt/lora/lora_manager.py
+++ b/python/sglang/srt/lora/lora_manager.py
@@ -28,7 +28,10 @@
     ParallelLMHead,
     VocabParallelEmbedding,
 )
-from sglang.srt.lora.backend.base_backend import BaseLoRABackend
+from sglang.srt.lora.backend.base_backend import (
+    BaseLoRABackend,
+    get_gathered_moe_num_tokens,
+)
 from sglang.srt.lora.backend.lora_registry import get_backend_from_name
 from sglang.srt.lora.layers import BaseLayerWithLoRA, FusedMoEWithLoRA, get_lora_layer
 from sglang.srt.lora.lora import LoRAAdapter
@@ -286,6 +289,11 @@ def _unload_lora_adapter(self, lora_ref: LoRARef) -> LoRAUpdateOutput:
             del self.loras[lora_ref.lora_id]
             del self.lora_refs[lora_ref.lora_id]
             self.num_pinned_loras -= int(lora_ref.pinned)
+            # Free the memory-pool buffer slot too, so a later load (colocate RL pushes a fresh uid
+            # every step) re-copies the new weights into a cleanly-freed slot instead of leaving a
+            # dangling uid_to_buffer_id entry that makes the reload skip the in-place buffer copy ->
+            # served (cuda-graph) buffer would keep stale weights. Cuda-graph-replay-safe.
+            self.memory_pool.free_lora(lora_ref.lora_id)
         except Exception as e:
             return self.create_lora_update_result(
                 success=False,
@@ -350,6 +358,20 @@ def prepare_lora_batch(self, forward_batch: ForwardBatch):
             and bs <= self.max_bs_in_cuda_graph
             and forward_batch.forward_mode.is_cuda_graph()
         )
+        if use_cuda_graph and self.lora_backend.is_moe_lora:
+            # This flag is a HEURISTIC computed before the runner's real replay decision. Under
+            # --enable-dp-attention a batch that is graph-eligible on THIS rank (small local bs)
+            # can still have a DP-gathered token count exceeding the static MoE capture buffers
+            # (max_bs*dp) when other ranks carry more tokens — such a batch cannot replay the
+            # captured graph and runs eager, so prep the eager (freshly sized) buffers for it
+            # instead of an under-sized static mapping.
+            moe_cg_buffers = self.lora_backend.moe_cg_buffers
+            if (
+                moe_cg_buffers is not None
+                and get_gathered_moe_num_tokens(forward_batch, bs)
+                > moe_cg_buffers["token_lora_mapping"].shape[0]
+            ):
+                use_cuda_graph = False
 
         weight_indices = [0] * len(forward_batch.lora_ids)
         lora_ranks = [0] * self.max_loras_per_batch
@@ -362,6 +384,55 @@ def prepare_lora_batch(self, forward_batch: ForwardBatch):
                 lora = self.loras[uid]
                 lora_ranks[weight_indices[i]] = lora.config.r
                 scalings[weight_indices[i]] = lora.scaling
+
+        local_active = any(lora_ranks[wi] > 0 for wi in weight_indices)
+
+        # MoE-expert LoRA under --enable-dp-attention: the MoE runs on DP-GATHERED tokens, so this
+        # rank's local experts also process OTHER ranks' tokens, whose adapter identity is not known
+        # host-side. Under Tier-1 (exactly one adapter loaded — the colocate RL case) the gathered
+        # tail can be attributed to that single adapter. Two mutually-exclusive stamps, both armed
+        # here (policy) and mechanically applied by the backend in _add_moe_lora_info:
+        #   * idle rank (forward_mode.is_idle(), zero local tokens): inject the adapter into
+        #     lora_ranks/scalings (nothing else carries them into the batch tensors) and record its
+        #     buffer id so the backend stamps the WHOLE gathered mapping and enables the adapter —
+        #     otherwise foreign tokens routed to this rank's experts silently lose the LoRA delta.
+        #   * active rank (local requests DO use the adapter): record the buffer id so the backend
+        #     stamps the gathered tail [num_tokens, moe_num_tokens) with it instead of copying an
+        #     arbitrary local token's slot. Base-only local batches (local_active False) and
+        #     multi-adapter batches arm NOTHING: the tail stays -1 (adapter-disabled), foreign
+        #     tokens get base — never a delta this rank cannot attribute.
+        # Buffer ids are host ints (no GPU sync -> cuda-graph safe). NOTE: gate on
+        # global_num_tokens_cpu too, not just global_dp_buffer_len — the eager path runs
+        # prepare_lora_batch from ForwardBatch.init_new BEFORE prepare_mlp_sync_batch assigns
+        # global_dp_buffer_len (only cuda-graph capture pre-sets it).
+        idle_rank_active_buffer_id = None
+        single_active_buffer_id = None
+        if self.lora_backend.is_moe_lora and (
+            forward_batch.global_dp_buffer_len is not None
+            or len(forward_batch.global_num_tokens_cpu or []) > 1
+        ):
+            loaded = [
+                (uid, bid)
+                for uid, bid in self.memory_pool.uid_to_buffer_id.items()
+                if uid is not None
+                and uid in self.loras
+                and self.loras[uid].config.r > 0
+            ]
+            if len(loaded) == 1:
+                uid, bid = loaded[0]
+                if forward_batch.forward_mode.is_idle():
+                    lora = self.loras[uid]
+                    lora_ranks[bid] = lora.config.r
+                    scalings[bid] = lora.scaling
+                    idle_rank_active_buffer_id = bid
+                elif local_active:
+                    single_active_buffer_id = bid
+
+        # Pass the stamps to the backend (read in _add_moe_lora_info); reset each batch so a
+        # previous batch's stamp never leaks into a later one.
+        self.lora_backend._idle_rank_active_buffer_id = idle_rank_active_buffer_id
+        self.lora_backend._single_loaded_buffer_id = single_active_buffer_id
+
         # Do in-place updates when CUDA graph is enabled and the batch forward mode
         # could use CUDA graph.
         self.lora_backend.prepare_lora_batch(
@@ -371,20 +442,10 @@ def prepare_lora_batch(self, forward_batch: ForwardBatch):
             scalings=scalings,
             use_cuda_graph=use_cuda_graph,
         )
-        self.lora_backend.batch_info.has_active_lora = any(
-            lora_ranks[wi] > 0 for wi in weight_indices
+        self.lora_backend.batch_info.has_active_lora = local_active or (
+            idle_rank_active_buffer_id is not None
         )
 
-    def prepare_idle_lora_batch(self):
-        """Reset per-batch LoRA state for a DP-attention idle forward.
-
-        Idle forwards run the model with zero local tokens and skip
-        prepare_lora_batch(); clearing batch_info makes the LoRA layers fall
-        back to the base path instead of reading the previous batch's stale
-        metadata.
-        """
-        self.lora_backend.batch_info = None
-
     def update_lora_info(self):
         """
         Update all LoRA modules to associate them with the latest memory buffer.
diff --git a/python/sglang/srt/lora/mem_pool.py b/python/sglang/srt/lora/mem_pool.py
index 188f02b67a42..acd92f3a7085 100644
--- a/python/sglang/srt/lora/mem_pool.py
+++ b/python/sglang/srt/lora/mem_pool.py
@@ -772,6 +772,31 @@ def get_available_buffer_slot():
                 self.uid_to_buffer_id[uid] = buffer_id
                 self.buffer_id_to_uid[buffer_id] = uid
 
+    def free_lora(self, uid: Optional[str]) -> None:
+        """Release a resident adapter's buffer slot + bookkeeping (called from unload_lora_adapter).
+
+        Per-step LoRA weight refresh in colocate RL pushes a FRESH uid every step (unload + load).
+        Without freeing the slot here, unload leaves ``uid_to_buffer_id`` pointing at the unloaded
+        adapter, so the next load's in-place buffer copy in ``prepare_lora_batch`` is skipped (the
+        eviction self-heal of a dangling entry is fragile) and the SERVED (cuda-graph) buffer keeps
+        stale weights. Freeing the slot makes the next load re-copy the new weights into the SAME
+        fixed-address buffer slot the captured graph reads -> cuda-graph-replay-safe. Works for both
+        ``max_loras_per_batch == 1`` and ``> 1``: each adapter frees/reloads its own slot, and
+        ``weight_indices`` (attention) / ``token_lora_mapping`` (MoE) are updated in place per step,
+        so the kernels follow whatever slot the reload assigned. Mirrors the eviction path above.
+        """
+        if uid is None:
+            return
+        buffer_id = self.uid_to_buffer_id.pop(uid, None)
+        if buffer_id is None:
+            return
+        self.eviction_policy.remove(uid)
+        if (
+            0 <= buffer_id < len(self.buffer_id_to_uid)
+            and self.buffer_id_to_uid[buffer_id] == uid
+        ):
+            self.buffer_id_to_uid[buffer_id] = EMPTY_SLOT
+
     def load_lora_weight_to_buffer(
         self,
         uid: str,
@@ -871,19 +896,45 @@ def load_lora_weight_tensor(
                 expert_match = re.search(r"experts\.(\d+)\.", name)
 
                 if expert_match:
-                    # Per-expert MoE weight — 2D tensors, one per expert
+                    # Per-expert MoE weight — 2D tensors, one per expert.
+                    # Init A and B INDEPENDENTLY (both buffer and cache_keys). Under
+                    # ``experts_shared_outer_loras`` one side of a projection is a
+                    # shared 3D Tensor (set by the dim()==3 branch below) and the
+                    # other is this per-expert dict: fc1 = shared A + per-expert B,
+                    # fc2 = the opposite. A coupled init keyed off
+                    # ``temp_A_buffer is None`` would either leave the per-expert
+                    # side's cache_keys as None (-> TypeError at
+                    # ``[expert_id] = name``) or clobber the shared side the 3D
+                    # branch already populated. So each side guards its own
+                    # buffer + cache_keys and never touches the other side.
                     target_module = target_module + "_moe"
-                    if temp_A_buffer[target_module] is None:
-                        temp_A_buffer[target_module] = {}
-                        temp_B_buffer[target_module] = {}
-                        temp_A_cache_keys[target_module] = {}
-                        temp_B_cache_keys[target_module] = {}
-
                     expert_id = int(expert_match.group(1))
                     if "lora_A" in name:
+                        assert not isinstance(
+                            temp_A_buffer[target_module], torch.Tensor
+                        ), (
+                            f"{target_module} lora_A already holds a shared-outer 3D "
+                            f"tensor but also got per-expert weight '{name}'; a "
+                            f"projection side must use one layout, not both."
+                        )
+                        if temp_A_buffer[target_module] is None:
+                            temp_A_buffer[target_module] = {}
+                        if temp_A_cache_keys[target_module] is None:
+                            temp_A_cache_keys[target_module] = {}
                         temp_A_buffer[target_module][expert_id] = weights
                         temp_A_cache_keys[target_module][expert_id] = name
                     else:
+                        assert not isinstance(
+                            temp_B_buffer[target_module], torch.Tensor
+                        ), (
+                            f"{target_module} lora_B already holds a shared-outer 3D "
+                            f"tensor but also got per-expert weight '{name}'; a "
+                            f"projection side must use one layout, not both."
+                        )
+                        if temp_B_buffer[target_module] is None:
+                            temp_B_buffer[target_module] = {}
+                        if temp_B_cache_keys[target_module] is None:
+                            temp_B_cache_keys[target_module] = {}
                         temp_B_buffer[target_module][expert_id] = weights
                         temp_B_cache_keys[target_module][expert_id] = name
                 elif "experts" in name and weights.dim() == 3:
diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py
index e57caaefde9a..c412e2b5297c 100644
--- a/python/sglang/srt/model_executor/forward_batch_info.py
+++ b/python/sglang/srt/model_executor/forward_batch_info.py
@@ -808,10 +808,14 @@ def init_new(
 
         if ret.forward_mode.is_idle():
             ret.positions = torch.empty((0,), dtype=torch.int64, device=device)
-            # Idle forwards (DP attention) skip the LoRA batch preparation
-            # below; clear stale batch info so LoRA layers skip application.
+            # Under --enable-dp-attention an IDLE rank still runs its local experts over the
+            # DP-GATHERED tokens of the other ranks, so the (MoE-)LoRA batch info must be
+            # re-prepared for THIS batch: sized to the gathered length and stamped via the
+            # Tier-1 idle-rank path. Skipping it here leaves the backend serving the PREVIOUS
+            # batch's stale token_lora_mapping — undersized (OOB reads/writes in the MoE-LoRA
+            # kernels) and/or pointing foreign tokens at the wrong adapter.
             if model_runner.server_args.enable_lora:
-                model_runner.lora_manager.prepare_idle_lora_batch()
+                model_runner.lora_manager.prepare_lora_batch(ret)
             return ret
 
         # Override the positions with diffusion LLM or spec_info
diff --git a/test/registered/unit/lora/test_moe_lora_tail_stamp.py b/test/registered/unit/lora/test_moe_lora_tail_stamp.py
new file mode 100644
index 000000000000..b10d53389331
--- /dev/null
+++ b/test/registered/unit/lora/test_moe_lora_tail_stamp.py
@@ -0,0 +1,134 @@
+"""Unit tests for the MoE-LoRA DP-gathered tail stamping in _add_moe_lora_info.
+
+Under --enable-dp-attention the MoE runs on DP-GATHERED tokens, so the per-token
+LoRA mapping must cover [0, moe_num_tokens) while the per-rank segments only fill
+[0, num_tokens). These tests pin the Tier-1 (single loaded adapter) semantics of
+the gathered tail [num_tokens, moe_num_tokens):
+
+  * active rank (local requests use the single adapter, ``_single_loaded_buffer_id``
+    armed): tail stamped with that adapter's buffer id;
+  * true idle rank (num_tokens == 0, ``_idle_rank_active_buffer_id`` armed): the
+    whole mapping stamped and the adapter enabled;
+  * base-only local batches, multi-adapter batches, and stale idle stamps on
+    token-bearing batches: tail stays -1 (adapter-disabled) — foreign tokens are
+    never given a delta this rank cannot attribute.
+
+The backend object is stubbed (only the fields _add_moe_lora_info reads are
+populated) so the tests run hermetically without a server or dist groups; the
+gathered length is forced via forward_batch.global_dp_buffer_len.
+
+Usage:
+    python -m pytest test/registered/unit/lora/test_moe_lora_tail_stamp.py -v
+"""
+
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+
+register_cuda_ci(est_time=9, stage="base-b", runner_config="1-gpu-small")
+register_amd_ci(est_time=9, suite="stage-b-test-1-gpu-small-amd")
+
+import unittest
+from types import SimpleNamespace
+
+import torch
+
+from sglang.srt.lora.backend.base_backend import BaseLoRABackend
+from sglang.srt.lora.utils import LoRABatchInfo
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+GATHERED = 16  # forced gathered length (> any local num_tokens used below)
+# buffer slots: 0 = the None/base uid slot (rank 0), 1/2 = adapters (rank 8), 3 = free
+RANKS = [0, 8, 8, 0]
+
+
+def _forward_batch(num_tokens: int):
+    mode = SimpleNamespace(
+        is_extend=lambda: False,
+        is_idle=lambda: num_tokens == 0,
+        is_cuda_graph=lambda: False,
+    )
+    return SimpleNamespace(
+        forward_mode=mode,
+        batch_size=num_tokens,  # decode: 1 token per seq
+        extend_seq_lens_cpu=None,
+        global_dp_buffer_len=GATHERED,
+        global_num_tokens_cpu=[num_tokens, GATHERED - num_tokens],
+    )
+
+
+def _batch_info(weight_indices, num_tokens):
+    bs = len(weight_indices)
+    return LoRABatchInfo(
+        use_cuda_graph=False,
+        bs=bs,
+        num_segments=bs,
+        seg_indptr=torch.arange(bs + 1, dtype=torch.int32, device=DEVICE),
+        weight_indices=torch.tensor(weight_indices, dtype=torch.int32, device=DEVICE),
+        lora_ranks=torch.tensor(RANKS, dtype=torch.int64, device=DEVICE),
+        scalings=torch.ones(len(RANKS), dtype=torch.float, device=DEVICE),
+        max_len=1,
+        seg_lens=torch.ones(bs, dtype=torch.int32, device=DEVICE),
+        permutation=None,
+        expected_tokens=num_tokens,
+    )
+
+
+def _run(weight_indices, num_tokens, idle_bid=None, single_bid=None):
+    stub = SimpleNamespace(
+        is_moe_lora=True,
+        _idle_rank_active_buffer_id=idle_bid,
+        _single_loaded_buffer_id=single_bid,
+    )
+    out = BaseLoRABackend._add_moe_lora_info(
+        stub, _forward_batch(num_tokens), _batch_info(weight_indices, num_tokens)
+    )
+    info = out.moe_lora_info
+    return info.token_lora_mapping.tolist(), info.adapter_enabled.tolist()
+
+
+class TestMoELoRATailStamp(unittest.TestCase):
+    def test_base_only_local_batch_keeps_disabled_tail(self):
+        # All local tokens on the base (None-uid) slot, nothing armed: the
+        # gathered tail must stay -1 and no adapter may be enabled.
+        mapping, enabled = _run([0, 0, 0, 0], num_tokens=4)
+        self.assertTrue(all(x == -1 for x in mapping[4:]), mapping)
+        self.assertEqual(enabled, [0, 0, 0, 0])
+
+    def test_stale_idle_stamp_ignored_on_token_bearing_batch(self):
+        # Defense in depth: an idle stamp must only be consumed when the rank is
+        # truly idle (num_tokens == 0), never on a batch with local tokens.
+        mapping, enabled = _run([0, 0, 0, 0], num_tokens=4, idle_bid=2)
+        self.assertTrue(all(x == -1 for x in mapping[4:]), mapping)
+        self.assertEqual(enabled[2], 0)
+
+    def test_multi_adapter_batch_keeps_disabled_tail(self):
+        # Two adapters used locally: the foreign tokens' adapter identity is
+        # unknowable host-side, so the tail must stay -1 (no mis-stamping with
+        # whatever the last local token happened to use).
+        mapping, _ = _run([1, 1, 2, 2], num_tokens=4)
+        self.assertTrue(all(x == -1 for x in mapping[4:]), mapping)
+
+    def test_idle_rank_stamps_whole_mapping_and_enables_adapter(self):
+        # True idle rank under Tier-1: the whole gathered mapping is stamped with
+        # the single loaded adapter and the adapter is enabled, so foreign tokens
+        # routed to this rank's experts get the LoRA delta.
+        mapping, enabled = _run([], num_tokens=0, idle_bid=2)
+        self.assertTrue(all(x == 2 for x in mapping), mapping)
+        self.assertEqual(enabled[2], 1)
+
+    def test_active_rank_stamps_tail_with_single_adapter(self):
+        # Local tokens actively use the single loaded adapter: the tail is
+        # stamped with its buffer id (not a copy of an arbitrary local slot).
+        mapping, enabled = _run([2, 2], num_tokens=2, single_bid=2)
+        self.assertTrue(all(x == 2 for x in mapping[2:]), mapping)
+        self.assertEqual(enabled[2], 1)
+
+    def test_single_adapter_loaded_but_local_base_keeps_disabled_tail(self):
+        # An adapter is loaded but this rank's local tokens are all base (the
+        # manager arms no stamp in this case): tail stays -1.
+        mapping, enabled = _run([0, 0], num_tokens=2)
+        self.assertTrue(all(x == -1 for x in mapping[2:]), mapping)
+        self.assertEqual(enabled, [0, 0, 0, 0])
+
+
+if __name__ == "__main__":
+    unittest.main()