[LoRA] MoE-LoRA CUDA-graph fixes under DP attention + RL adapter-reload fixes

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,
     )

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
 
 

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(

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
@@ -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,6 +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 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
 
@@ -188,15 +229,26 @@ 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
 
+        # 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
@@ -233,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(
@@ -242,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
             ),
         }
 
@@ -291,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,
@@ -299,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,
@@ -373,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 (
@@ -440,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