[Diffusion][sglang-miles] refactor: rollout controller as a scheduler-step plugin instead of mixins

benchmark/kernels/fused_moe_triton/benchmark_sglang_fused_moe_triton.py

@@ -38,6 +38,7 @@ def fused_moe_triton_api(
 ):
     topk_op = TopK(
         top_k=topk,
+        layer_id=0,
         renormalize=False,
         use_grouped_topk=False,
         output_format=TopKOutputFormat.TRITON_KERNEL,

docs/docs/advanced_features/sglang_for_rl.mdx

@@ -608,6 +608,7 @@ A request can outlive a version change: the RL flow retracts it with `pause_gene
 ```
 
 - Ranges are half-open `[start, end)` over output-token indices, prompt excluded. The usual case is a single span.
+- `--enable-prefill-weight-versions` adds `prefill_weight_versions`, the same span list over the prompt tokens, so KV reused from the prefix cache under older weights is visible.
 
 ## Easy To Postpone Generation
 

python/pyproject.toml

@@ -74,7 +74,7 @@ dependencies = [
   "regex",
   "requests",
   "scipy",
-  "sentencepiece",
+  "sentencepiece==0.2.1",  # 0.2.2 rejects null pieces in InternVL tokenizers.
   "setproctitle",
   "sgl-deep-ep==0.1.2",
   "sgl-deep-gemm==0.2.0",
@@ -142,7 +142,7 @@ diffusion-qvg = [
 ]
 
 ray = [
-  "ray[default]>=2.55.1",
+  "ray[default]>=2.56.0",
 ]
 
 tracing = [

python/pyproject_cpu.toml

@@ -54,7 +54,7 @@ dependencies = [
   "regex",
   "requests",
   "scipy",
-  "sentencepiece",
+  "sentencepiece==0.2.1",  # 0.2.2 rejects null pieces in InternVL tokenizers.
   "setproctitle",
   "smg-grpc-servicer>=0.9.0",
   "soundfile==0.13.1",

python/pyproject_npu.toml

@@ -55,7 +55,7 @@ dependencies = [
   "regex",
   "requests",
   "scipy",
-  "sentencepiece",
+  "sentencepiece==0.2.1",  # 0.2.2 rejects null pieces in InternVL tokenizers.
   "setproctitle",
   "smg-grpc-servicer>=0.9.0",
   "soundfile==0.13.1",

python/pyproject_other.toml

@@ -56,7 +56,7 @@ runtime_base = [
   "regex",
   "requests",
   "scipy",
-  "sentencepiece",
+  "sentencepiece==0.2.1",  # 0.2.2 rejects null pieces in InternVL tokenizers.
   "setproctitle",
   "smg-grpc-servicer>=0.9.0",
   "soundfile==0.13.1",

python/pyproject_xpu.toml

@@ -54,7 +54,7 @@ dependencies = [
   "regex",
   "requests",
   "scipy",
-  "sentencepiece",
+  "sentencepiece==0.2.1",  # 0.2.2 rejects null pieces in InternVL tokenizers.
   "setproctitle",
   "sglang-kernel-xpu @ https://github.com/sgl-project/whl/releases/download/v0.3.0+xpu/sglang_kernel_xpu-0.3.0-cp310-abi3-manylinux_2_38_x86_64.whl",
   "smg-grpc-servicer>=0.9.0",

python/sglang/kernels/ops/layernorm/mhc.py

@@ -1178,14 +1178,17 @@ def mhc_pre(
     else:
         if num_tokens <= 2048:
             assert n_splits == 1
-            if hc_hidden_size == 16384:
-                hidden_block = 256
-            elif hc_hidden_size == 28672:
-                hidden_block = 128
+            split_size = hc_hidden_size // n_splits_pre
+            for hidden_block in (256, 128):
+                if (
+                    hc_hidden_size % hidden_block == 0
+                    and split_size % hidden_block == 0
+                ):
+                    break
             else:
                 raise NotImplementedError(
-                    f"mhc_pre splitk kernel only supports hc_hidden_size in {{16384, 28672}}, "
-                    f"got {hc_hidden_size}"
+                    f"mhc_pre splitk kernel needs a hidden_block in {{256, 128}} that divides "
+                    f"hc_hidden_size // split_k = {split_size}, got hc_hidden_size={hc_hidden_size}"
                 )
             kernel_0, _ = mhc_pre_gemm_sqrsum_splitk_kernel(
                 hc_mult3,

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; mapping was sized before the dp gather "
+        f"length was known (see get_gathered_moe_num_tokens)"
+    )
+    # 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

@@ -262,14 +262,15 @@ def _validate(self) -> None:
                 self.spectrum_params = SpectrumParams()
             apply_h3_spectrum_param_defaults(self.spectrum_params)
         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/diffusion_generator.py

@@ -214,7 +214,8 @@ def _check_remote_scheduler(self):
             raise ConnectionError(
                 f"Could not connect to remote scheduler at "
                 f"{self.server_args.scheduler_endpoint} with `local mode` as False. "
-                "Please ensure the server is running."
+                "Please ensure the server is running, binds a reachable "
+                "--scheduler-host, and uses the same --scheduler-port as this client."
             )
         logger.info(
             f"Successfully connected to remote scheduler at "