[cherrypick from #36389] [sglang][lora] Support DP attention in LoRA backends

docs/docs/advanced_features/lora.mdx

@@ -1,7 +1,7 @@
 ---
 title: "LoRA Serving"
 metatags:
-    description: "SGLang multi-LoRA serving: S-LoRA and Punica techniques, dynamic adapter loading, GPU pinning, overlap loading, Triton and CSGMV backends."
+    description: "SGLang multi-LoRA serving: S-LoRA and Punica techniques, DP attention, dynamic adapter loading, GPU pinning, overlap loading, Triton and CSGMV backends."
 ---
 SGLang enables the use of [LoRA adapters](https://arxiv.org/abs/2106.09685) with a base model. By incorporating techniques from [S-LoRA](https://arxiv.org/pdf/2311.03285) and [Punica](https://arxiv.org/pdf/2310.18547), SGLang can efficiently support multiple LoRA adapters for different sequences within a single batch of inputs.
 
@@ -412,6 +412,32 @@ server_process, port = launch_server_cmd(
 terminate_process(server_process)
 ```
 
+## LoRA with DP attention
+
+SGLang supports LoRA serving with [data parallel attention (DPA)](/docs/advanced_features/dp_dpa_smg_guide). The model architecture must support DPA, and you must currently select the Triton LoRA backend. Other built-in LoRA backends reject this configuration at startup.
+
+The following NVIDIA CUDA example uses four attention-DP ranks:
+
+```bash
+python3 -m sglang.launch_server \
+    --model-path MODEL_PATH \
+    --tp 4 \
+    --dp-size 4 \
+    --enable-dp-attention \
+    --enable-lora \
+    --lora-backend triton \
+    --lora-paths '{"lora_name":"ADAPTER_NAME","lora_path":"ADAPTER_PATH","pinned":true}'
+```
+
+Keep these constraints in mind:
+
+- `--dp-size` must equal `--tp-size`, so every attention DP rank owns its whole token batch. The server rejects other layouts at startup.
+- All adapters must be GPU-pinned (`pinned: true`), including dynamically loaded adapters. You can pin at most `--max-loras-per-batch - 1` adapters; the remaining slot is reserved for the base model.
+- Decode CUDA graphs are supported. LoRA prefill runs eagerly when DPA is enabled.
+- `--max-loras-per-batch` limits the union of adapters active across all DP ranks in a synchronized forward, not the adapters on one rank alone.
+- Dynamic load and unload requests are sent to every DP rank. An operation succeeds only if every rank reports success. After a partially successful load, retry unloading the same adapter name before loading it again.
+- The built-in LoRA memory pool does not support `--enable-lora-overlap-loading` with DPA. The server rejects that combination rather than allowing DP ranks to assign different GPU slots.
+
 ## LoRA Overlap Loading
 
 

proto/sglang/runtime/v1/sglang.proto

@@ -11,6 +11,7 @@ service SglangService {
   rpc Tokenize(TokenizeRequest) returns (TokenizeResponse);
   rpc Detokenize(DetokenizeRequest) returns (DetokenizeResponse);
   rpc HealthCheck(HealthCheckRequest) returns (HealthCheckResponse);
+  rpc WatchEngineState(WatchEngineStateRequest) returns (stream EngineStateSnapshot);
   rpc GetModelInfo(GetModelInfoRequest) returns (GetModelInfoResponse);
   rpc GetServerInfo(GetServerInfoRequest) returns (GetServerInfoResponse);
   rpc ListModels(ListModelsRequest) returns (ListModelsResponse);
@@ -174,6 +175,23 @@ message HealthCheckResponse {
   bool healthy = 1;
 }
 
+// ---- Engine state ----
+
+message WatchEngineStateRequest {}
+
+// A complete discovery and lifecycle snapshot. The JSON discovery payloads
+// intentionally retain the same wire shape as GetModelInfo/GetServerInfo.
+message EngineStateSnapshot {
+  // Unix time in nanoseconds, captured once for this engine process.
+  uint64 instance_id = 1;
+  // Starts at one and increases for each snapshot from this instance.
+  uint64 revision = 2;
+  bool healthy = 3;
+  bool is_pause = 4;
+  GetModelInfoResponse model_info = 5;
+  GetServerInfoResponse server_info = 6;
+}
+
 // ---- Model info ----
 
 message GetModelInfoRequest {}

python/pyproject.toml

@@ -97,8 +97,8 @@ dependencies = [
   "uvloop",
   "watchfiles",
   "xgrammar==0.2.1",
-  "xxhash",
-  "zstandard",
+  "xxhash",  # /pull_weights delta checksum
+  "zstandard",  # /pull_weights delta codec
 ]
 
 [[tool.uv.index]]
@@ -141,7 +141,7 @@ diffusion-qvg = [
 ]
 
 ray = [
-  "ray[default]>=2.55.1",
+  "ray[default]>=2.56.0",
 ]
 
 tracing = [

python/sglang/kernels/ops/moe/fused_moe_lora_kernel.py

@@ -403,8 +403,11 @@ def _fused_moe_lora_expand(
         IS_PRIMARY=False,
         **expand_config,
     )
+    # the base output may be padded (Marlin), so its slices sit at output width / num_slices
+    slice_stride = output.shape[-1] // num_slices
     for i in range(num_slices):
-        output[:, :, i * N + offset : (i + 1) * N + offset] += b_intermediate_cache1[i]
+        start = i * slice_stride + offset
+        output[:, :, start : start + N] += b_intermediate_cache1[i]
 
 
 @torch.inference_mode()

python/sglang/kernels/ops/moe/virtual_experts.py

@@ -55,7 +55,9 @@ def _fused_virtual_topk_ids_kernel(
     # a real virtual-expert slot belonging to another adapter and trigger OOB
     # loads in downstream LoRA kernels.
     shifted = base + safe_lora * num_experts_for_weight
-    result = tl.where(base < 0, base, shifted)
+    # Reject an out-of-range base ID before it can alias the next adapter.
+    result = tl.where(base >= num_experts_for_weight, -1, shifted)
+    result = tl.where(base < 0, base, result)
     tl.store(virtual_topk_ids_ptr + offs, result, mask=valid)
 
     # Write mask once per row (at first k position)
@@ -517,7 +519,7 @@ def _merged_experts_fused_moe_lora_add_impl(
     output: torch.Tensor,
     hidden_states: torch.Tensor,
     lora_a: torch.Tensor,
-    lora_b: torch.Tensor,
+    lora_b: torch.Tensor | list[torch.Tensor] | tuple[torch.Tensor, ...],
     topk_ids: torch.Tensor,
     topk_weights: torch.Tensor,
     token_lora_mapping: torch.Tensor,
@@ -526,13 +528,30 @@ def _merged_experts_fused_moe_lora_add_impl(
     experts_shared_outer_loras_b: bool,
     routing_cache: dict | None = None,
 ) -> None:
+    """Fused virtual-experts LoRA delta add.
+
+    ``lora_b`` accepts either a single tensor or a sequence of tensors stacked
+    along the output dim. Length-2 is the gate_up case where A has rank ``2*r``
+    (gate's A and up's A concatenated along rank) and each B has rank ``r``.
+    The shrink runs once over the full ``2*r`` rank; the expand runs once per
+    B, each reading its half of the intermediate and writing to its slice of
+    ``output``.
     """
-    1. Prepare virtual expert routing metadata from topk_ids + token_lora_mapping * num_experts.
-    2. Flatten LoRA weights from [max_loras, num_experts, ...] to [max_loras * num_experts, ...].
-    3. Run regular SGLang fused-MoE kernels for LoRA A and LoRA B.
-    4. Mask out tokens with token_lora_mapping == -1 on the add path.
-    """
+    lora_b_list: list[torch.Tensor] = (
+        list(lora_b) if isinstance(lora_b, (list, tuple)) else [lora_b]
+    )
+    n_b = len(lora_b_list)
+    assert n_b in (1, 2), f"lora_b must be length 1 or 2, got {n_b}"
+    b_rank = lora_b_list[0].shape[3]
+    for b in lora_b_list[1:]:
+        assert b.shape == lora_b_list[0].shape, (
+            f"all lora_b tensors must share shape; got {[tuple(t.shape) for t in lora_b_list]}"
+        )
+
     max_loras, _, max_lora_rank, _ = lora_a.shape
+    assert max_lora_rank == n_b * b_rank, (
+        f"lora_a rank {max_lora_rank} != n_b ({n_b}) * lora_b rank {b_rank}"
+    )
     input_top_k = 1 if hidden_states.shape[0] == topk_ids.numel() else topk_ids.shape[1]
 
     def _merge_lora_expert_weight(t: torch.Tensor) -> torch.Tensor:
@@ -616,16 +635,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,
@@ -644,13 +665,24 @@ def _get_routing(
     )
 
     lora_a_virtual = _merge_lora_expert_weight(lora_a)
-    lora_b_virtual = _merge_lora_expert_weight(lora_b)
+    lora_b_virtuals = [_merge_lora_expert_weight(b) for b in lora_b_list]
     num_experts_a = lora_a.shape[1]
-    num_experts_b = lora_b.shape[1]
-
+    num_experts_b = lora_b_list[0].shape[1]
+    half_out = lora_b_list[0].shape[2]
+
+    # 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; under DP attention the mapping must "
+        f"come from the TP-global LoRA routing view (LoRABatchLayout.TP_GLOBAL)"
+    )
+    # Split-K atomic additions must accumulate before rounding to the activation dtype.
     intermediate = torch.zeros(
-        [token_lora_mapping.shape[0], topk_ids.shape[1], max_lora_rank],
-        dtype=hidden_states.dtype,
+        [topk_ids.shape[0], topk_ids.shape[1], max_lora_rank],
+        dtype=torch.float32,
         device=hidden_states.device,
     )
 
@@ -680,7 +712,8 @@ def _get_routing(
         a_stage_config,
     )
 
-    b_stage_config = _get_stage_config(lora_b_virtual, 1)
+    intermediate = intermediate.to(hidden_states.dtype)
+    b_stage_config = _get_stage_config(lora_b_virtuals[0], 1)
     (
         sorted_token_ids,
         expert_ids,
@@ -694,33 +727,53 @@ def _get_routing(
         b_stage_config["BLOCK_SIZE_M"],
     )
 
-    invoke_fused_moe_kernel(
-        intermediate.view(-1, max_lora_rank),
-        lora_b_virtual,
-        None,
-        output,
-        None,
-        None,
-        None,
-        topk_weights,
-        topk_ids,
-        sorted_token_ids,
-        expert_ids,
-        num_tokens_post_padded,
-        mul_routed_weight,
-        1,
-        b_stage_config,
-        tl.bfloat16 if hidden_states.dtype == torch.bfloat16 else tl.float16,
-        False,
-        False,
-        False,
-        False,
-        False,
-        None,
-        fuse_add_to_output=True,
-        add_output_mask=token_lora_mask,
-        router_topk=topk_ids.shape[1],
-    )
+    # n_b expands. For len 1: K=b_rank covers full intermediate, write full output.
+    # For len 2 (gate_up): split intermediate along rank into [gate, up] halves
+    # (each contiguous, K=b_rank=r) and output along last dim into [gate, up]
+    # halves (each of width half_out). Each B in lora_b_virtuals is its own
+    # half's weight tensor, naturally K=b_rank.
+    for b_idx, b_virtual in enumerate(lora_b_virtuals):
+        if n_b == 1:
+            inter_arg = intermediate.view(-1, b_rank)
+            out_arg = output
+        else:
+            inter_arg = (
+                intermediate[..., b_idx * b_rank : (b_idx + 1) * b_rank]
+                .contiguous()
+                .view(-1, b_rank)
+            )
+            out_arg = output[
+                ..., b_idx * half_out : (b_idx + 1) * half_out
+            ].contiguous()
+        invoke_fused_moe_kernel(
+            inter_arg,
+            b_virtual,
+            None,
+            out_arg,
+            None,
+            None,
+            None,
+            topk_weights,
+            topk_ids,
+            sorted_token_ids,
+            expert_ids,
+            num_tokens_post_padded,
+            mul_routed_weight,
+            1,
+            b_stage_config,
+            tl.bfloat16 if hidden_states.dtype == torch.bfloat16 else tl.float16,
+            False,
+            False,
+            False,
+            False,
+            False,
+            None,
+            fuse_add_to_output=True,
+            add_output_mask=token_lora_mask,
+            router_topk=topk_ids.shape[1],
+        )
+        if n_b != 1:
+            output[..., b_idx * half_out : (b_idx + 1) * half_out].copy_(out_arg)
 
 
 def _merged_experts_fused_moe_lora_add_op(
@@ -763,7 +816,7 @@ def merged_experts_fused_moe_lora_add(
     output: torch.Tensor,
     hidden_states: torch.Tensor,
     lora_a: torch.Tensor,
-    lora_b: torch.Tensor,
+    lora_b: torch.Tensor | list[torch.Tensor] | tuple[torch.Tensor, ...],
     topk_ids: torch.Tensor,
     topk_weights: torch.Tensor,
     token_lora_mapping: torch.Tensor,
@@ -772,7 +825,12 @@ def merged_experts_fused_moe_lora_add(
     experts_shared_outer_loras_b: bool,
     routing_cache: dict | None = None,
 ) -> None:
-    """Public API: wraps the registered op with routing_cache support."""
+    """Public API: wraps the registered op with routing_cache support.
+
+    ``lora_b`` accepts a sequence of length 2 for the gate_up case (each B
+    holds one half of the stacked output, rank ``r``, with A's rank ``2*r``);
+    a single tensor is used for the down case.
+    """
     _merged_experts_fused_moe_lora_add_impl(
         output,
         hidden_states,

python/sglang/multimodal_gen/configs/sample/minimax_h3.py

@@ -241,14 +241,15 @@ def _validate(self) -> None:
                 "video/audio denoise loop has no lossless TeaCache contract"
             )
         if self.rollout:
-            raise ValueError(
-                "MiniMax H3 does not support rollout: its coupled video/audio "
-                "scheduler has no SchedulerRLMixin contract"
-            )
+            task = str(self.task or "t2va").lower()
+            if task not in ("t2va",):
+                raise ValueError(
+                    f"MiniMax H3 rollout currently supports task=t2va only, got {task!r}"
+                )
         if self.return_trajectory_latents or self.return_trajectory_decoded:
             raise ValueError(
-                "MiniMax H3 does not support trajectory output for its coupled "
-                "video/audio denoise state"
+                "MiniMax H3 does not support return_trajectory_latents/decoded; "
+                "use rollout=True with rollout_return_dit_trajectory instead"
             )
         seeds = self.seed if isinstance(self.seed, list) else [self.seed]
         for seed in seeds:

python/sglang/multimodal_gen/runtime/entrypoints/post_training/io_struct.py

@@ -79,6 +79,8 @@ class RolloutRequest(BaseModel):
     fps: Optional[int] = None
 
     rollout: bool = True
+    # "uint8": quantise the video engine-side; None: ship unchanged
+    rollout_video_dtype: Optional[str] = None
     rollout_sde_type: str = "sde"
     rollout_noise_level: float = 0.7
     rollout_log_prob_no_const: bool = False