From 7c027f8eaac142d906969729525577c897ff49b9 Mon Sep 17 00:00:00 2001
From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com>
Date: Thu, 13 Aug 2026 17:12:23 +0800
Subject: [PATCH 1/3] feat(rollout): add SGLang native group sampling

---
 relax/engine/rollout/sglang_rollout.py        | 326 +++++++++++----
 relax/utils/arguments.py                      |   9 +
 .../multimodal/run-qwen3-vl-4B-8xgpu.sh       |   7 +
 .../test_sglang_native_group_sampling.py      | 370 ++++++++++++++++++
 4 files changed, 635 insertions(+), 77 deletions(-)
 create mode 100644 tests/engine/rollout/test_sglang_native_group_sampling.py

diff --git a/relax/engine/rollout/sglang_rollout.py b/relax/engine/rollout/sglang_rollout.py
index 62854c038..35a21fa8f 100644
--- a/relax/engine/rollout/sglang_rollout.py
+++ b/relax/engine/rollout/sglang_rollout.py
@@ -52,6 +52,69 @@
 logger = get_logger(__name__)
 
 
+def _native_group_sampling_outputs(output: Any, expected_count: int) -> list[dict[str, Any]]:
+    """Validate and normalize the non-streaming SGLang ``n`` response."""
+    if not isinstance(output, list):
+        raise TypeError(f"Native group sampling must return a list, got {type(output)!r}")
+    if len(output) != expected_count:
+        raise ValueError(
+            "Native group sampling returned an unexpected number of outputs: "
+            f"expected {expected_count}, got {len(output)}"
+        )
+    if not all(isinstance(item, dict) for item in output):
+        raise TypeError("Native group sampling responses must be dictionaries")
+    return output
+
+
+def _native_group_sampling_params(sampling_params: dict[str, Any], group_size: int) -> dict[str, Any]:
+    """Create request-local sampling parameters for one prompt group."""
+    if group_size < 2:
+        raise ValueError(f"Native group sampling requires at least two samples, got {group_size}")
+    params = sampling_params.copy()
+    params["n"] = group_size
+    return params
+
+
+def _native_group_sampling_media_payload(encoded_mm: dict[str, list]) -> dict[str, list[list]]:
+    """Wrap one prompt's media as a batch-of-one for SGLang ``n > 1``.
+
+    SGLang promotes a single input to batch mode for parallel sampling. Without
+    this outer batch dimension, a prompt containing multiple images is
+    interpreted as a batch with one prompt per image and fails validation.
+    """
+    return {key: [value] for key, value in encoded_mm.items()}
+
+
+def _native_group_sampling_eligible(args: Namespace, group: list[Sample], evaluation: bool) -> bool:
+    """Limit native ``n`` to standard first-turn group-RM rollouts."""
+    if not getattr(args, "sglang_native_group_sampling", False):
+        return False
+    if evaluation or len(group) < 2 or not getattr(args, "group_rm", False):
+        return False
+    if getattr(args, "partial_rollout", False) or getattr(args, "use_slime_router", False):
+        return False
+    if getattr(args, "use_opd", False) or getattr(args, "use_rollout_routing_replay", False):
+        return False
+    if getattr(args, "lora_rank", 0) > 0 and getattr(args, "lora_adapter_mode", False):
+        return False
+    if getattr(args, "sglang_enable_deterministic_inference", False):
+        return False
+    if getattr(args, "custom_generate_function_path", None) is not None:
+        return False
+    if any(getattr(sample, "generate_function_path", None) is not None for sample in group):
+        return False
+    first = group[0]
+    if any(sample.prompt != first.prompt for sample in group[1:]):
+        return False
+    if any(sample.multimodal_inputs is not first.multimodal_inputs for sample in group[1:]):
+        return False
+    if any(sample.tokens or sample.rollout_tokens or sample.response for sample in group):
+        return False
+    if any(sample.loss_mask is not None for sample in group):
+        return False
+    return all(sample.status in (Sample.Status.PENDING, Sample.Status.ABORTED) for sample in group)
+
+
 # Misuse guard for the per-request permit contract. Set while the session-level
 # lock (GenerateState.semaphore) is held so that a legacy custom function that
 # wrongly calls inference_permit()/post_generate() fails loudly instead of
@@ -294,6 +357,79 @@ async def _encode_multimodal_inputs(multimodal_inputs: dict) -> tuple[dict[str,
     return encoded, monotonic() - t_start
 
 
+def _apply_sglang_output_tokens(
+    args: Namespace,
+    state: "GenerateState",
+    sample: Sample,
+    output: dict[str, Any],
+    evaluation: bool,
+) -> None:
+    """Apply the token-level part of a standard SGLang response."""
+    if "output_token_logprobs" in output["meta_info"]:
+        new_response_tokens = [item[1] for item in output["meta_info"]["output_token_logprobs"]]
+        new_response_log_probs = [item[0] for item in output["meta_info"]["output_token_logprobs"]]
+    else:
+        new_response_tokens = list(output["output_ids"])
+        new_response_log_probs = []
+
+    if state.opd_manager and not evaluation:
+        new_response_tokens, new_response_log_probs = state.opd_manager.parse_rollout_logprobs(
+            output["meta_info"], new_response_tokens, new_response_log_probs
+        )
+
+    while hasattr(state.tokenizer, "image_token_id") and state.tokenizer.image_token_id in new_response_tokens:
+        index = new_response_tokens.index(state.tokenizer.image_token_id)
+        new_response_tokens[index] = state.tokenizer.pad_token_id
+        logger.warning(
+            "Image token found in output tokens, replaced with pad_token_id. Consider updating the model's stop "
+            "condition to stop at image_token_id if you want to avoid this."
+        )
+
+    while hasattr(state.tokenizer, "audio_token_id") and state.tokenizer.audio_token_id in new_response_tokens:
+        index = new_response_tokens.index(state.tokenizer.audio_token_id)
+        new_response_tokens[index] = state.tokenizer.pad_token_id
+        logger.warning(
+            "Audio token found in output tokens, replaced with pad_token_id. Consider updating the model's stop "
+            "condition to stop at audio_token_id if you want to avoid this."
+        )
+
+    while hasattr(state.tokenizer, "video_token_id") and state.tokenizer.video_token_id in new_response_tokens:
+        index = new_response_tokens.index(state.tokenizer.video_token_id)
+        new_response_tokens[index] = state.tokenizer.pad_token_id
+        logger.warning(
+            "Video token found in output tokens, replaced with pad_token_id. Consider updating the model's stop "
+            "condition to stop at video_token_id if you want to avoid this."
+        )
+
+    if state.processor is not None:
+        from relax.utils.data.processing_utils import sanitize_kimi_k25_response_tokens
+
+        sanitized = sanitize_kimi_k25_response_tokens(state.processor, new_response_tokens)
+        if sanitized is not new_response_tokens:
+            replaced = sum(1 for a, b in zip(new_response_tokens, sanitized, strict=True) if a != b)
+            if replaced:
+                logger.warning(
+                    f"K2.x: replaced {replaced} stray <|media_pad|> token(s) in rollout response with pad_token_id."
+                )
+            new_response_tokens = sanitized
+
+    sample.tokens = sample.tokens + new_response_tokens
+    sample.rollout_tokens = sample.rollout_tokens + new_response_tokens
+    sample.response_length += len(new_response_tokens)
+    sample.response += output["text"]
+
+    if sample.loss_mask is not None:
+        assert args.partial_rollout and args.mask_offpolicy_in_partial_rollout
+        sample.loss_mask += [1] * len(new_response_tokens)
+
+    if sample.rollout_log_probs is None:
+        sample.rollout_log_probs = []
+    sample.rollout_log_probs += new_response_log_probs
+
+    if state.opd_manager and not evaluation:
+        state.opd_manager.after_rollout(sample, output)
+
+
 async def generate(
     args: Namespace, sample: Sample, sampling_params: dict[str, Any], evaluation: bool = False
 ) -> Sample:
@@ -400,73 +536,7 @@ async def generate(
 
         sample = await postprocess_sample_with_radix_tree(args, sample, output)
     else:
-        if "output_token_logprobs" in output["meta_info"]:
-            new_response_tokens = [item[1] for item in output["meta_info"]["output_token_logprobs"]]
-            new_response_log_probs = [item[0] for item in output["meta_info"]["output_token_logprobs"]]
-        else:
-            new_response_tokens = output["output_ids"]
-            new_response_log_probs = []
-
-        if state.opd_manager and not evaluation:
-            new_response_tokens, new_response_log_probs = state.opd_manager.parse_rollout_logprobs(
-                output["meta_info"], new_response_tokens, new_response_log_probs
-            )
-
-        while hasattr(state.tokenizer, "image_token_id") and state.tokenizer.image_token_id in new_response_tokens:
-            index = new_response_tokens.index(state.tokenizer.image_token_id)
-            new_response_tokens[index] = state.tokenizer.pad_token_id
-            logger.warning(
-                "Image token found in output tokens, replaced with pad_token_id. Consider updating the model's stop condition to stop at image_token_id if you want to avoid this."
-            )
-
-        while hasattr(state.tokenizer, "audio_token_id") and state.tokenizer.audio_token_id in new_response_tokens:
-            index = new_response_tokens.index(state.tokenizer.audio_token_id)
-            new_response_tokens[index] = state.tokenizer.pad_token_id
-            logger.warning(
-                "Audio token found in output tokens, replaced with pad_token_id. Consider updating the model's stop condition to stop at audio_token_id if you want to avoid this."
-            )
-
-        while hasattr(state.tokenizer, "video_token_id") and state.tokenizer.video_token_id in new_response_tokens:
-            index = new_response_tokens.index(state.tokenizer.video_token_id)
-            new_response_tokens[index] = state.tokenizer.pad_token_id
-            logger.warning(
-                "Video token found in output tokens, replaced with pad_token_id. Consider updating the model's stop condition to stop at video_token_id if you want to avoid this."
-            )
-
-        # K2.x tokenizers don't expose image_token_id but reserve <|media_pad|>
-        # for vision input slots. A hallucinated <|media_pad|> in the response
-        # inflates num_placeholders past sum(feature_lengths) in the bridge,
-        # forcing dynamic expansion → broadcast → 233 GiB OOM. Replace in-place
-        # so positional accounting matches sglang's per-token logprobs.
-        if state.processor is not None:
-            from relax.utils.data.processing_utils import sanitize_kimi_k25_response_tokens
-
-            sanitized = sanitize_kimi_k25_response_tokens(state.processor, new_response_tokens)
-            if sanitized is not new_response_tokens:
-                replaced = sum(1 for a, b in zip(new_response_tokens, sanitized, strict=True) if a != b)
-                if replaced:
-                    logger.warning(
-                        f"K2.x: replaced {replaced} stray <|media_pad|> token(s) in rollout response with pad_token_id."
-                    )
-                new_response_tokens = sanitized
-
-        # Update sample with tokens directly - avoiding re-tokenization
-        sample.tokens = sample.tokens + new_response_tokens
-        sample.rollout_tokens = sample.rollout_tokens + new_response_tokens
-        sample.response_length += len(new_response_tokens)
-        sample.response += output["text"]
-
-        # When partial rollout and masking off policy is enabled, update the loss mask
-        if sample.loss_mask is not None:
-            assert args.partial_rollout and args.mask_offpolicy_in_partial_rollout
-            sample.loss_mask += [1] * len(new_response_tokens)
-
-        if sample.rollout_log_probs is None:
-            sample.rollout_log_probs = []
-        sample.rollout_log_probs += new_response_log_probs
-
-        if state.opd_manager and not evaluation:
-            state.opd_manager.after_rollout(sample, output)
+        _apply_sglang_output_tokens(args, state, sample, output, evaluation)
 
     if "routed_experts" in output["meta_info"]:
         sample.rollout_routed_experts = np.frombuffer(
@@ -491,6 +561,105 @@ async def generate(
     return sample
 
 
+async def _generate_native_group(
+    args: Namespace,
+    state: "GenerateState",
+    group: list[Sample],
+    sampling_params: dict[str, Any],
+) -> list[Sample]:
+    """Generate one first-turn group with SGLang's native parallel sampling."""
+    first = group[0]
+    tokenizer_prompt_ids = state.tokenizer.encode(first.prompt, add_special_tokens=False)
+    processor_prompt_ids = tokenizer_prompt_ids
+    image_processor_time = 0.0
+    has_media = first.multimodal_inputs is not None and any(
+        first.multimodal_inputs.get(key) for key in ("images", "videos", "audio")
+    )
+    if state.processor and has_media:
+        processor_prompt_ids, train_inputs, image_processor_time = await _run_image_processor(
+            state, args, first.prompt, first.multimodal_inputs
+        )
+        for sample in group:
+            sample.multimodal_train_inputs = train_inputs
+
+    params = _native_group_sampling_params(sampling_params, len(group))
+    if params["max_new_tokens"] == 0:
+        for sample in group:
+            sample.status = Sample.Status.TRUNCATED
+        return group
+
+    payload: dict[str, Any] = {
+        "sampling_params": params,
+        "return_logprob": True,
+        "input_ids": tokenizer_prompt_ids,
+    }
+
+    mm_encode_time = 0.0
+    if has_media:
+        encoded_mm = getattr(first, "_pre_encoded_mm", None)
+        if encoded_mm is None:
+            encoded_mm, mm_encode_time = await _encode_multimodal_inputs(first.multimodal_inputs)
+        else:
+            mm_encode_time = getattr(first, "_pre_encoded_mm_elapsed", 0.0)
+        payload.update(_native_group_sampling_media_payload(encoded_mm))
+
+    for sample in group:
+        sample.tokens = list(processor_prompt_ids)
+        sample.rollout_tokens = list(tokenizer_prompt_ids)
+
+    headers = None
+    if args.sglang_router_policy == "consistent_hashing" and first.session_id:
+        headers = {"X-SMG-Routing-Key": first.session_id}
+    elif getattr(args, "slime_router_sticky", False) and first.group_index is not None:
+        headers = {"X-SMG-Routing-Key": str(first.group_index)}
+
+    request_start = monotonic()
+    try:
+        async with state.semaphore:
+            token = _holding_session_lock.set(True)
+            try:
+                if state.aborted:
+                    for sample in group:
+                        sample.status = Sample.Status.ABORTED
+                    return group
+                with state.dp_rank_context() as _:
+                    output = await post(
+                        f"http://{args.sglang_router_ip}:{args.sglang_router_port}/generate",
+                        payload,
+                        headers=headers,
+                    )
+            finally:
+                _holding_session_lock.reset(token)
+    except GenerationAborted:
+        for sample in group:
+            sample.status = Sample.Status.ABORTED
+        return group
+    request_time = monotonic() - request_start
+
+    outputs = _native_group_sampling_outputs(output, len(group))
+    for sample, sample_output in zip(group, outputs, strict=True):
+        post_start = monotonic()
+        _apply_sglang_output_tokens(args, state, sample, sample_output, evaluation=False)
+        if "routed_experts" in sample_output["meta_info"]:
+            sample.rollout_routed_experts = np.frombuffer(
+                pybase64.b64decode(sample_output["meta_info"]["routed_experts"].encode("ascii")),
+                dtype=np.int32,
+            ).reshape(len(sample.tokens) - 1, args.num_layers, args.moe_router_topk)
+        sample.update_from_meta_info(args, sample_output["meta_info"])
+        sample.metadata["_timing"] = {
+            "generate": request_time,
+            "post_generate": monotonic() - post_start,
+            **({"image_processor": image_processor_time} if image_processor_time else {}),
+            **({"mm_encode": mm_encode_time} if mm_encode_time else {}),
+        }
+
+    for sample in group:
+        for attr in ("_pre_encoded_mm", "_pre_encoded_mm_elapsed"):
+            if hasattr(sample, attr):
+                delattr(sample, attr)
+    return group
+
+
 async def _dispatch_generate(
     state: "GenerateState",
     args: Namespace,
@@ -657,17 +826,20 @@ async def generate_and_rm_group(
             sample._pre_encoded_mm = encoded_mm
             sample._pre_encoded_mm_elapsed = t_enc
 
-    tasks = []
-    for idx, sample in enumerate(group):
-        current_sampling_params = sampling_params.copy()
-        if getattr(args, "sglang_enable_deterministic_inference", False):
-            seed = state.group_sampling_seeds[idx]
-            current_sampling_params["sampling_seed"] = seed
-        tasks.append(
-            asyncio.create_task(generate_and_rm(args, sample, current_sampling_params, evaluation=evaluation))
-        )
+    if _native_group_sampling_eligible(args, group, evaluation):
+        group = await _generate_native_group(args, state, group, sampling_params)
+    else:
+        tasks = []
+        for idx, sample in enumerate(group):
+            current_sampling_params = sampling_params.copy()
+            if getattr(args, "sglang_enable_deterministic_inference", False):
+                seed = state.group_sampling_seeds[idx]
+                current_sampling_params["sampling_seed"] = seed
+            tasks.append(
+                asyncio.create_task(generate_and_rm(args, sample, current_sampling_params, evaluation=evaluation))
+            )
 
-    group = await asyncio.gather(*tasks)
+        group = await asyncio.gather(*tasks)
 
     # eval should still compute group reward even if abort was triggered by a concurrent rollout
     if (not state.aborted or evaluation) and args.group_rm:
diff --git a/relax/utils/arguments.py b/relax/utils/arguments.py
index 22430e883..c4283b707 100644
--- a/relax/utils/arguments.py
+++ b/relax/utils/arguments.py
@@ -709,6 +709,15 @@ def add_rollout_arguments(parser):
                 default=1.0,
                 help="the temperature for the inference engine during rollout.",
             )
+            parser.add_argument(
+                "--sglang-native-group-sampling",
+                action="store_true",
+                default=False,
+                help=(
+                    "Use one SGLang /generate request with sampling_params.n equal to the prompt group size "
+                    "for standard first-turn group-RM rollouts."
+                ),
+            )
             parser.add_argument(
                 "--rollout-top-p", type=float, default=1.0, help="the top-p for the inference engine during rollout."
             )
diff --git a/scripts/training/multimodal/run-qwen3-vl-4B-8xgpu.sh b/scripts/training/multimodal/run-qwen3-vl-4B-8xgpu.sh
index 27d97353a..8b138e89c 100644
--- a/scripts/training/multimodal/run-qwen3-vl-4B-8xgpu.sh
+++ b/scripts/training/multimodal/run-qwen3-vl-4B-8xgpu.sh
@@ -55,6 +55,13 @@ ROLLOUT_ARGS=(
    --rollout-temperature 1.0
 )
 
+if [ "${SGLANG_NATIVE_GROUP_SAMPLING:-0}" = "1" ]; then
+   ROLLOUT_ARGS+=(--group-rm --sglang-native-group-sampling)
+   ROLLOUT_ARGS+=(--sglang-router-policy "${SGLANG_ROUTER_POLICY:-round_robin}")
+elif [ -n "${SGLANG_ROUTER_POLICY:-}" ]; then
+   ROLLOUT_ARGS+=(--sglang-router-policy "${SGLANG_ROUTER_POLICY}")
+fi
+
 PERF_ARGS=(
    --tensor-model-parallel-size 2
    --sequence-parallel
diff --git a/tests/engine/rollout/test_sglang_native_group_sampling.py b/tests/engine/rollout/test_sglang_native_group_sampling.py
new file mode 100644
index 000000000..ab313e47f
--- /dev/null
+++ b/tests/engine/rollout/test_sglang_native_group_sampling.py
@@ -0,0 +1,370 @@
+# Copyright (c) 2026 Relax Authors. All Rights Reserved.
+
+import argparse
+import asyncio
+import contextlib
+from argparse import Namespace
+from pathlib import Path
+
+import pytest
+
+from relax.utils.types import Sample
+
+
+INFERENCE_IMPORT_ERROR = ""
+try:
+    from sglang.srt.managers.io_struct import GenerateReqInput
+
+    from relax.engine.rollout import sglang_rollout
+    from relax.engine.rollout.sglang_rollout import (
+        _generate_native_group,
+        _holding_session_lock,
+        _native_group_sampling_eligible,
+        _native_group_sampling_media_payload,
+        _native_group_sampling_outputs,
+        _native_group_sampling_params,
+    )
+    from relax.utils.arguments import get_slime_extra_args_provider
+
+    HAS_INFERENCE_DEPS = True
+except (ImportError, OSError, RuntimeError) as exc:
+    HAS_INFERENCE_DEPS = False
+    INFERENCE_IMPORT_ERROR = f"{type(exc).__name__}: {exc}"
+
+pytestmark = pytest.mark.skipif(
+    not HAS_INFERENCE_DEPS,
+    reason=f"Ray/SGLang inference stack unavailable: {INFERENCE_IMPORT_ERROR}",
+)
+
+
+class _Tokenizer:
+    def encode(self, prompt, add_special_tokens=False):
+        assert add_special_tokens is False
+        return [11, 12]
+
+
+class _NativeState:
+    def __init__(self, *, processor=None, aborted: bool = False) -> None:
+        self.aborted = aborted
+        self.dp_entries = 0
+        self.opd_manager = None
+        self.processor = processor
+        self.semaphore = asyncio.Semaphore(1)
+        self.tokenizer = _Tokenizer()
+
+    @contextlib.contextmanager
+    def dp_rank_context(self):
+        self.dp_entries += 1
+        yield 0
+
+
+def _native_args(**overrides) -> Namespace:
+    values = {
+        "sglang_native_group_sampling": True,
+        "group_rm": True,
+        "partial_rollout": False,
+        "mask_offpolicy_in_partial_rollout": False,
+        "use_slime_router": False,
+        "use_opd": False,
+        "use_rollout_routing_replay": False,
+        "lora_rank": 0,
+        "lora_adapter_mode": False,
+        "sglang_enable_deterministic_inference": False,
+        "custom_generate_function_path": None,
+        "sglang_router_ip": "router.test",
+        "sglang_router_port": 30000,
+        "sglang_router_policy": None,
+        "slime_router_sticky": False,
+        "sglang_speculative_algorithm": None,
+    }
+    values.update(overrides)
+    return Namespace(**values)
+
+
+def _response(token_id: int, text: str, *, finish_reason: str = "stop") -> dict:
+    return {
+        "text": text,
+        "meta_info": {
+            "output_token_logprobs": [(-0.1, token_id, text)],
+            "finish_reason": {"type": finish_reason},
+            "cached_tokens": 2,
+            "prompt_tokens": 2,
+        },
+    }
+
+
+def test_native_group_sampling_params_sets_n_without_mutating_input() -> None:
+    sampling_params = {"temperature": 1.0, "max_new_tokens": 32}
+
+    result = _native_group_sampling_params(sampling_params, 8)
+
+    assert result == {"temperature": 1.0, "max_new_tokens": 32, "n": 8}
+    assert sampling_params == {"temperature": 1.0, "max_new_tokens": 32}
+
+
+def test_native_group_sampling_params_requires_multiple_samples() -> None:
+    with pytest.raises(ValueError, match="at least two"):
+        _native_group_sampling_params({}, 1)
+
+
+def test_native_group_sampling_media_payload_adds_batch_dimension() -> None:
+    encoded_mm = {
+        "image_data": ["image-a", "image-b"],
+        "audio_data": ["audio-a"],
+    }
+
+    assert _native_group_sampling_media_payload(encoded_mm) == {
+        "image_data": [["image-a", "image-b"]],
+        "audio_data": [["audio-a"]],
+    }
+    assert encoded_mm == {
+        "image_data": ["image-a", "image-b"],
+        "audio_data": ["audio-a"],
+    }
+
+
+def test_native_group_sampling_media_payload_matches_sglang_batch_contract() -> None:
+    payload = _native_group_sampling_media_payload({"image_data": ["image-a", "image-b"]})
+    request = GenerateReqInput(
+        input_ids=[11, 12],
+        image_data=payload["image_data"],
+        sampling_params={"n": 2, "max_new_tokens": 8},
+    )
+
+    request.normalize_batch_and_arguments()
+
+    assert request.input_ids == [[11, 12], [11, 12]]
+    assert request.image_data == [
+        ["image-a", "image-b"],
+        ["image-a", "image-b"],
+    ]
+    assert request.modalities == ["multi-images", "multi-images"]
+
+
+def test_native_group_sampling_outputs_requires_native_list_shape() -> None:
+    outputs = [{"text": "a"}, {"text": "b"}]
+
+    assert _native_group_sampling_outputs(outputs, 2) is outputs
+
+    with pytest.raises(TypeError, match="must return a list"):
+        _native_group_sampling_outputs(outputs[0], 2)
+    with pytest.raises(ValueError, match="unexpected number"):
+        _native_group_sampling_outputs(outputs, 8)
+
+
+def test_native_group_sampling_outputs_rejects_non_dict_items() -> None:
+    with pytest.raises(TypeError, match="responses must be dictionaries"):
+        _native_group_sampling_outputs([{"text": "a"}, "not-a-response"], 2)
+
+
+def test_native_group_sampling_cli_is_opt_in() -> None:
+    parser = argparse.ArgumentParser(add_help=False, allow_abbrev=False)
+    get_slime_extra_args_provider()(parser)
+
+    defaults, _ = parser.parse_known_args([])
+    enabled, _ = parser.parse_known_args(["--sglang-native-group-sampling"])
+
+    assert defaults.sglang_native_group_sampling is False
+    assert enabled.sglang_native_group_sampling is True
+
+
+def test_qwen3_vl_launcher_couples_native_sampling_with_group_rm_and_balanced_routing() -> None:
+    launcher = Path(__file__).parents[3] / "scripts/training/multimodal/run-qwen3-vl-4B-8xgpu.sh"
+    source = launcher.read_text()
+    guard = 'if [ "${SGLANG_NATIVE_GROUP_SAMPLING:-0}" = "1" ]; then'
+    before_native, separator, after_native = source.partition(guard)
+    native_branch, else_separator, explicit_override_branch = after_native.partition("elif")
+
+    assert separator == guard
+    assert else_separator == "elif"
+    assert "--group-rm" not in before_native
+    assert "--group-rm --sglang-native-group-sampling" in native_branch
+    assert "${SGLANG_ROUTER_POLICY:-round_robin}" in native_branch
+    assert "${SGLANG_ROUTER_POLICY}" in explicit_override_branch
+
+
+@pytest.mark.parametrize(
+    "overrides",
+    [
+        {"sglang_native_group_sampling": False},
+        {"group_rm": False},
+        {"partial_rollout": True},
+        {"use_slime_router": True},
+        {"use_opd": True},
+        {"use_rollout_routing_replay": True},
+        {"lora_rank": 1, "lora_adapter_mode": True},
+        {"sglang_enable_deterministic_inference": True},
+        {"custom_generate_function_path": "custom.generate"},
+    ],
+)
+def test_native_group_sampling_eligibility_rejects_incompatible_modes(overrides) -> None:
+    group = [Sample(prompt="same"), Sample(prompt="same")]
+    assert not _native_group_sampling_eligible(_native_args(**overrides), group, evaluation=False)
+
+
+def test_native_group_sampling_eligibility_rejects_evaluation_or_single_sample() -> None:
+    group = [Sample(prompt="same"), Sample(prompt="same")]
+    assert not _native_group_sampling_eligible(_native_args(), group, evaluation=True)
+    assert not _native_group_sampling_eligible(_native_args(), group[:1], evaluation=False)
+
+
+def test_native_group_sampling_eligibility_requires_homogeneous_fresh_group() -> None:
+    shared_media = {"images": [object(), object()], "videos": [], "audio": []}
+    group = [
+        Sample(prompt="same", multimodal_inputs=shared_media),
+        Sample(prompt="same", multimodal_inputs=shared_media),
+    ]
+    assert _native_group_sampling_eligible(_native_args(), group, evaluation=False)
+
+    group[1].prompt = "different"
+    assert not _native_group_sampling_eligible(_native_args(), group, evaluation=False)
+    group[1].prompt = "same"
+
+    group[1].multimodal_inputs = {"images": [object()], "videos": [], "audio": []}
+    assert not _native_group_sampling_eligible(_native_args(), group, evaluation=False)
+    group[1].multimodal_inputs = shared_media
+
+    group[1].tokens = [1]
+    assert not _native_group_sampling_eligible(_native_args(), group, evaluation=False)
+    group[1].tokens = []
+
+    group[1].loss_mask = []
+    assert not _native_group_sampling_eligible(_native_args(), group, evaluation=False)
+    group[1].loss_mask = None
+
+    group[1].generate_function_path = "custom.generate"
+    assert not _native_group_sampling_eligible(_native_args(), group, evaluation=False)
+    group[1].generate_function_path = None
+
+    group[1].status = Sample.Status.FAILED
+    assert not _native_group_sampling_eligible(_native_args(), group, evaluation=False)
+
+
+async def test_generate_native_group_uses_one_n_request_and_maps_outputs(monkeypatch) -> None:
+    state = _NativeState()
+    args = _native_args()
+    group = [Sample(prompt="same") for _ in range(8)]
+    sampling_params = {"temperature": 1.0, "max_new_tokens": 32}
+    requests = []
+
+    async def fake_post(url, payload, headers=None):
+        assert state.semaphore.locked()
+        assert _holding_session_lock.get() is True
+        requests.append((url, payload, headers))
+        return [_response(100 + index, f"r{index}") for index in range(8)]
+
+    monkeypatch.setattr(sglang_rollout, "post", fake_post)
+
+    result = await _generate_native_group(args, state, group, sampling_params)
+
+    assert result is group
+    assert len(requests) == 1
+    url, payload, headers = requests[0]
+    assert url == "http://router.test:30000/generate"
+    assert payload["sampling_params"] == {"temperature": 1.0, "max_new_tokens": 32, "n": 8}
+    assert payload["return_logprob"] is True
+    assert payload["input_ids"] == [11, 12]
+    assert headers is None
+    assert sampling_params == {"temperature": 1.0, "max_new_tokens": 32}
+    assert state.dp_entries == 1
+    assert not state.semaphore.locked()
+    assert _holding_session_lock.get() is False
+
+    for index, sample in enumerate(group):
+        assert sample.tokens == [11, 12, 100 + index]
+        assert sample.rollout_tokens == [11, 12, 100 + index]
+        assert sample.response == f"r{index}"
+        assert sample.response_length == 1
+        assert sample.rollout_log_probs == [-0.1]
+        assert sample.status == Sample.Status.COMPLETED
+        assert set(sample.metadata["_timing"]) == {"generate", "post_generate"}
+
+
+async def test_generate_native_group_encodes_shared_multimodal_input_once(monkeypatch) -> None:
+    state = _NativeState(processor=object())
+    args = _native_args()
+    shared_media = {"images": [object(), object()], "videos": [], "audio": []}
+    group = [Sample(prompt="same", multimodal_inputs=shared_media) for _ in range(2)]
+    train_inputs = {"pixel_values": object()}
+    calls = {"processor": 0, "encode": 0, "post": 0}
+
+    async def fake_processor(state_arg, args_arg, prompt, multimodal_inputs):
+        assert state_arg is state
+        assert args_arg is args
+        assert prompt == "same"
+        assert multimodal_inputs is shared_media
+        calls["processor"] += 1
+        return [21, 22, 23], train_inputs, 0.25
+
+    async def fake_encode(multimodal_inputs):
+        assert multimodal_inputs is shared_media
+        calls["encode"] += 1
+        return {"image_data": ["encoded-image-a", "encoded-image-b"]}, 0.5
+
+    async def fake_post(url, payload, headers=None):
+        calls["post"] += 1
+        assert payload["image_data"] == [["encoded-image-a", "encoded-image-b"]]
+        return [_response(101, "a"), _response(102, "b")]
+
+    monkeypatch.setattr(sglang_rollout, "_run_image_processor", fake_processor)
+    monkeypatch.setattr(sglang_rollout, "_encode_multimodal_inputs", fake_encode)
+    monkeypatch.setattr(sglang_rollout, "post", fake_post)
+
+    await _generate_native_group(args, state, group, {"max_new_tokens": 32})
+
+    assert calls == {"processor": 1, "encode": 1, "post": 1}
+    for sample in group:
+        assert sample.tokens[:3] == [21, 22, 23]
+        assert sample.rollout_tokens[:2] == [11, 12]
+        assert sample.multimodal_train_inputs is train_inputs
+        assert sample.metadata["_timing"]["image_processor"] == 0.25
+        assert sample.metadata["_timing"]["mm_encode"] == 0.5
+
+
+async def test_generate_native_group_uses_consistent_hash_routing_key(monkeypatch) -> None:
+    state = _NativeState()
+    args = _native_args(sglang_router_policy="consistent_hashing")
+    group = [Sample(prompt="same", session_id="group-session") for _ in range(2)]
+    observed_headers = []
+
+    async def fake_post(url, payload, headers=None):
+        observed_headers.append(headers)
+        return [_response(101, "a"), _response(102, "b")]
+
+    monkeypatch.setattr(sglang_rollout, "post", fake_post)
+
+    await _generate_native_group(args, state, group, {"max_new_tokens": 32})
+
+    assert observed_headers == [{"X-SMG-Routing-Key": "group-session"}]
+
+
+async def test_generate_native_group_marks_entire_group_aborted_before_post(monkeypatch) -> None:
+    state = _NativeState(aborted=True)
+    group = [Sample(prompt="same") for _ in range(2)]
+
+    async def unexpected_post(url, payload, headers=None):
+        raise AssertionError("aborted group must not issue a request")
+
+    monkeypatch.setattr(sglang_rollout, "post", unexpected_post)
+
+    result = await _generate_native_group(_native_args(), state, group, {"max_new_tokens": 32})
+
+    assert all(sample.status == Sample.Status.ABORTED for sample in result)
+    assert not state.semaphore.locked()
+    assert _holding_session_lock.get() is False
+
+
+async def test_generate_native_group_rejects_wrong_response_cardinality(monkeypatch) -> None:
+    state = _NativeState()
+    group = [Sample(prompt="same") for _ in range(2)]
+
+    async def fake_post(url, payload, headers=None):
+        return [_response(101, "only-one")]
+
+    monkeypatch.setattr(sglang_rollout, "post", fake_post)
+
+    with pytest.raises(ValueError, match="expected 2, got 1"):
+        await _generate_native_group(_native_args(), state, group, {"max_new_tokens": 32})
+
+    assert not state.semaphore.locked()
+    assert _holding_session_lock.get() is False

From d6a28b455ce3a4cfb961c3c19ddef117cf7a2a56 Mon Sep 17 00:00:00 2001
From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com>
Date: Fri, 14 Aug 2026 21:32:16 +0800
Subject: [PATCH 2/3] fix(rollout): harden native group sampling
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

# 🐛 Bug Fix

## Prevent pre-encoded media leak into sample dicts

- Wrap the group generate branch in try/finally so the temporary
  pre-encoded media attributes are always dropped before samples reach
  Sample.to_dict(), covering the per-sample fallback, abort and error
  paths alike
- Previously the base64 media blobs could leak into the data buffer
  via __dict__ on early returns

## Keep aborted groups eligible for the native path

- Revert tokens, rollout_tokens and multimodal_train_inputs to their
  pre-call state when the request is aborted before any output, so
  retried groups no longer silently degrade to per-sample requests

## Exclude speculative decoding from native eligibility

- The draft/verify path is not validated with parallel sampling, so
  keep the per-sample fanout when a speculative algorithm is set

## Add visibility for the native/fallback decision

- Log once when native group sampling activates and warn once when the
  flag is enabled but a training group is ineligible
- Warn at startup when flag combinations make every group ineligible

---

# ♻️ Refactor

## Remove dead code and document semantics

- Drop the unreachable routed_experts branch: the native payload never
  requests it and routing replay is excluded by eligibility
- Document group-shared timing duplication and session-permit
  concurrency semantics

---

# ✅ Tests

## Cover abort, cleanup and eligibility edge cases

- Mid-flight abort maps n outputs to aborted samples with partial
  tokens retained
- Abort before dispatch restores a fresh, still-eligible group
- Pre-encoded media is stripped on success and on fallback failure
- Speculative decoding joins the incompatible-modes matrix
- Ineligible-group warning fires only once
---
 relax/engine/rollout/sglang_rollout.py        |  95 ++++++++---
 relax/utils/arguments.py                      |  29 ++++
 .../test_sglang_native_group_sampling.py      | 155 ++++++++++++++++++
 3 files changed, 253 insertions(+), 26 deletions(-)

diff --git a/relax/engine/rollout/sglang_rollout.py b/relax/engine/rollout/sglang_rollout.py
index 35a21fa8f..112b4cc0f 100644
--- a/relax/engine/rollout/sglang_rollout.py
+++ b/relax/engine/rollout/sglang_rollout.py
@@ -99,6 +99,11 @@ def _native_group_sampling_eligible(args: Namespace, group: list[Sample], evalua
         return False
     if getattr(args, "sglang_enable_deterministic_inference", False):
         return False
+    # Speculative decoding: SGLang serves n>1 by internally duplicating the
+    # request, and the draft/verify path is not validated with parallel
+    # sampling here, so keep the well-tested per-sample fanout.
+    if getattr(args, "sglang_speculative_algorithm", None):
+        return False
     if getattr(args, "custom_generate_function_path", None) is not None:
         return False
     if any(getattr(sample, "generate_function_path", None) is not None for sample in group):
@@ -115,6 +120,17 @@ def _native_group_sampling_eligible(args: Namespace, group: list[Sample], evalua
     return all(sample.status in (Sample.Status.PENDING, Sample.Status.ABORTED) for sample in group)
 
 
+_NATIVE_GROUP_SAMPLING_LOGGED: set[str] = set()
+
+
+def _log_native_group_sampling_once(key: str, message: str, warn: bool = False) -> None:
+    """Log the native-path activation/fallback decision once per process."""
+    if key in _NATIVE_GROUP_SAMPLING_LOGGED:
+        return
+    _NATIVE_GROUP_SAMPLING_LOGGED.add(key)
+    (logger.warning if warn else logger.info)(message)
+
+
 # Misuse guard for the per-request permit contract. Set while the session-level
 # lock (GenerateState.semaphore) is held so that a legacy custom function that
 # wrongly calls inference_permit()/post_generate() fails loudly instead of
@@ -607,6 +623,17 @@ async def _generate_native_group(
         sample.tokens = list(processor_prompt_ids)
         sample.rollout_tokens = list(tokenizer_prompt_ids)
 
+    def _revert_to_fresh_group() -> None:
+        # Nothing was generated: restore the pre-call state (empty tokens, no
+        # train inputs) so the aborted group still passes
+        # _native_group_sampling_eligible when it is retried, instead of
+        # silently degrading to the per-sample fallback.
+        for sample in group:
+            sample.tokens = []
+            sample.rollout_tokens = []
+            sample.multimodal_train_inputs = None
+            sample.status = Sample.Status.ABORTED
+
     headers = None
     if args.sglang_router_policy == "consistent_hashing" and first.session_id:
         headers = {"X-SMG-Routing-Key": first.session_id}
@@ -615,12 +642,15 @@ async def _generate_native_group(
 
     request_start = monotonic()
     try:
+        # NOTE: one session permit covers the whole group here, while the
+        # per-sample path holds one permit per in-flight sample. With the same
+        # --sglang-server-concurrency the native path therefore admits up to
+        # len(group) times more concurrent decode streams.
         async with state.semaphore:
             token = _holding_session_lock.set(True)
             try:
                 if state.aborted:
-                    for sample in group:
-                        sample.status = Sample.Status.ABORTED
+                    _revert_to_fresh_group()
                     return group
                 with state.dp_rank_context() as _:
                     output = await post(
@@ -631,8 +661,7 @@ async def _generate_native_group(
             finally:
                 _holding_session_lock.reset(token)
     except GenerationAborted:
-        for sample in group:
-            sample.status = Sample.Status.ABORTED
+        _revert_to_fresh_group()
         return group
     request_time = monotonic() - request_start
 
@@ -640,23 +669,17 @@ async def _generate_native_group(
     for sample, sample_output in zip(group, outputs, strict=True):
         post_start = monotonic()
         _apply_sglang_output_tokens(args, state, sample, sample_output, evaluation=False)
-        if "routed_experts" in sample_output["meta_info"]:
-            sample.rollout_routed_experts = np.frombuffer(
-                pybase64.b64decode(sample_output["meta_info"]["routed_experts"].encode("ascii")),
-                dtype=np.int32,
-            ).reshape(len(sample.tokens) - 1, args.num_layers, args.moe_router_topk)
         sample.update_from_meta_info(args, sample_output["meta_info"])
+        # NOTE: "generate", "image_processor" and "mm_encode" are shared by the
+        # whole group but recorded on every sample, so perf_detail sums count
+        # each shared phase len(group) times; compare against the per-sample
+        # baseline accordingly.
         sample.metadata["_timing"] = {
             "generate": request_time,
             "post_generate": monotonic() - post_start,
             **({"image_processor": image_processor_time} if image_processor_time else {}),
             **({"mm_encode": mm_encode_time} if mm_encode_time else {}),
         }
-
-    for sample in group:
-        for attr in ("_pre_encoded_mm", "_pre_encoded_mm_elapsed"):
-            if hasattr(sample, attr):
-                delattr(sample, attr)
     return group
 
 
@@ -826,20 +849,40 @@ async def generate_and_rm_group(
             sample._pre_encoded_mm = encoded_mm
             sample._pre_encoded_mm_elapsed = t_enc
 
-    if _native_group_sampling_eligible(args, group, evaluation):
-        group = await _generate_native_group(args, state, group, sampling_params)
-    else:
-        tasks = []
-        for idx, sample in enumerate(group):
-            current_sampling_params = sampling_params.copy()
-            if getattr(args, "sglang_enable_deterministic_inference", False):
-                seed = state.group_sampling_seeds[idx]
-                current_sampling_params["sampling_seed"] = seed
-            tasks.append(
-                asyncio.create_task(generate_and_rm(args, sample, current_sampling_params, evaluation=evaluation))
+    try:
+        if _native_group_sampling_eligible(args, group, evaluation):
+            _log_native_group_sampling_once(
+                "active",
+                f"Native group sampling active: one n={len(group)} request per prompt group.",
             )
+            group = await _generate_native_group(args, state, group, sampling_params)
+        else:
+            if getattr(args, "sglang_native_group_sampling", False) and not evaluation:
+                _log_native_group_sampling_once(
+                    "fallback",
+                    "Native group sampling is enabled but this group is ineligible; using the "
+                    "per-sample fallback (see _native_group_sampling_eligible for the gating).",
+                    warn=True,
+                )
+            tasks = []
+            for idx, sample in enumerate(group):
+                current_sampling_params = sampling_params.copy()
+                if getattr(args, "sglang_enable_deterministic_inference", False):
+                    seed = state.group_sampling_seeds[idx]
+                    current_sampling_params["sampling_seed"] = seed
+                tasks.append(
+                    asyncio.create_task(generate_and_rm(args, sample, current_sampling_params, evaluation=evaluation))
+                )
 
-        group = await asyncio.gather(*tasks)
+            group = await asyncio.gather(*tasks)
+    finally:
+        # The pre-encoded blobs are only a hint for generate(); drop them so the
+        # base64 payloads can never leak into Sample.to_dict() (data buffer /
+        # transfer queue) via __dict__, on success, abort and error paths alike.
+        for sample in group:
+            for attr in ("_pre_encoded_mm", "_pre_encoded_mm_elapsed"):
+                if hasattr(sample, attr):
+                    delattr(sample, attr)
 
     # eval should still compute group reward even if abort was triggered by a concurrent rollout
     if (not state.aborted or evaluation) and args.group_rm:
diff --git a/relax/utils/arguments.py b/relax/utils/arguments.py
index c4283b707..e5b9f363f 100644
--- a/relax/utils/arguments.py
+++ b/relax/utils/arguments.py
@@ -3352,6 +3352,35 @@ def slime_validate_args(args):
     if args.use_rollout_routing_replay:
         args.use_routing_replay = True
 
+    if getattr(args, "sglang_native_group_sampling", False):
+        # Keep in sync with _native_group_sampling_eligible() in
+        # relax/engine/rollout/sglang_rollout.py: any of these makes every group
+        # ineligible, so the flag would silently do nothing.
+        native_group_blockers = []
+        if not getattr(args, "group_rm", False):
+            native_group_blockers.append("--group-rm is not set")
+        if getattr(args, "partial_rollout", False):
+            native_group_blockers.append("--partial-rollout is set")
+        if getattr(args, "use_slime_router", False):
+            native_group_blockers.append("--use-slime-router is set")
+        if getattr(args, "use_opd", False):
+            native_group_blockers.append("OPD is enabled")
+        if getattr(args, "use_rollout_routing_replay", False):
+            native_group_blockers.append("--use-rollout-routing-replay is set")
+        if getattr(args, "lora_rank", 0) > 0 and getattr(args, "lora_adapter_mode", False):
+            native_group_blockers.append("LoRA adapter mode is enabled")
+        if getattr(args, "sglang_enable_deterministic_inference", False):
+            native_group_blockers.append("--sglang-enable-deterministic-inference is set")
+        if getattr(args, "sglang_speculative_algorithm", None):
+            native_group_blockers.append("--sglang-speculative-algorithm is set")
+        if getattr(args, "custom_generate_function_path", None) is not None:
+            native_group_blockers.append("--custom-generate-function-path is set")
+        if native_group_blockers:
+            logger.warning(
+                "--sglang-native-group-sampling is enabled but every rollout group will fall back "
+                f"to per-sample requests because: {'; '.join(native_group_blockers)}."
+            )
+
     if args.custom_config_path:
         with open(args.custom_config_path) as f:
             data = yaml.safe_load(f) or {}
diff --git a/tests/engine/rollout/test_sglang_native_group_sampling.py b/tests/engine/rollout/test_sglang_native_group_sampling.py
index ab313e47f..776938b82 100644
--- a/tests/engine/rollout/test_sglang_native_group_sampling.py
+++ b/tests/engine/rollout/test_sglang_native_group_sampling.py
@@ -16,6 +16,7 @@
     from sglang.srt.managers.io_struct import GenerateReqInput
 
     from relax.engine.rollout import sglang_rollout
+    from relax.engine.rollout.request_permit import GenerationAborted
     from relax.engine.rollout.sglang_rollout import (
         _generate_native_group,
         _holding_session_lock,
@@ -194,6 +195,7 @@ def test_qwen3_vl_launcher_couples_native_sampling_with_group_rm_and_balanced_ro
         {"use_rollout_routing_replay": True},
         {"lora_rank": 1, "lora_adapter_mode": True},
         {"sglang_enable_deterministic_inference": True},
+        {"sglang_speculative_algorithm": "EAGLE"},
         {"custom_generate_function_path": "custom.generate"},
     ],
 )
@@ -368,3 +370,156 @@ async def fake_post(url, payload, headers=None):
 
     assert not state.semaphore.locked()
     assert _holding_session_lock.get() is False
+
+
+async def test_generate_native_group_abort_before_post_keeps_group_native_eligible(monkeypatch) -> None:
+    state = _NativeState(aborted=True)
+    args = _native_args()
+    group = [Sample(prompt="same") for _ in range(2)]
+
+    async def unexpected_post(url, payload, headers=None):
+        raise AssertionError("aborted group must not issue a request")
+
+    monkeypatch.setattr(sglang_rollout, "post", unexpected_post)
+
+    result = await _generate_native_group(args, state, group, {"max_new_tokens": 32})
+
+    for sample in result:
+        assert sample.status == Sample.Status.ABORTED
+        assert sample.tokens == []
+        assert sample.rollout_tokens == []
+        assert sample.multimodal_train_inputs is None
+    assert _native_group_sampling_eligible(args, result, evaluation=False)
+
+
+async def test_generate_native_group_reverts_group_when_request_is_aborted(monkeypatch) -> None:
+    state = _NativeState()
+    args = _native_args()
+    group = [Sample(prompt="same") for _ in range(2)]
+
+    async def aborted_post(url, payload, headers=None):
+        raise GenerationAborted("router draining")
+
+    monkeypatch.setattr(sglang_rollout, "post", aborted_post)
+
+    result = await _generate_native_group(args, state, group, {"max_new_tokens": 32})
+
+    assert not state.semaphore.locked()
+    assert _holding_session_lock.get() is False
+    for sample in result:
+        assert sample.status == Sample.Status.ABORTED
+        assert sample.tokens == []
+        assert sample.rollout_tokens == []
+    assert _native_group_sampling_eligible(args, result, evaluation=False)
+
+
+async def test_generate_native_group_maps_mid_flight_abort_outputs(monkeypatch) -> None:
+    state = _NativeState()
+    args = _native_args()
+    group = [Sample(prompt="same") for _ in range(2)]
+
+    async def fake_post(url, payload, headers=None):
+        return [
+            _response(101, "partial-a", finish_reason="abort"),
+            _response(102, "partial-b", finish_reason="abort"),
+        ]
+
+    monkeypatch.setattr(sglang_rollout, "post", fake_post)
+
+    result = await _generate_native_group(args, state, group, {"max_new_tokens": 32})
+
+    for index, sample in enumerate(result):
+        assert sample.status == Sample.Status.ABORTED
+        assert sample.tokens == [11, 12, 101 + index]
+        assert sample.response == ("partial-a", "partial-b")[index]
+    # Partially generated samples keep their tokens and must retry per-sample.
+    assert not _native_group_sampling_eligible(args, result, evaluation=False)
+
+
+async def test_generate_and_rm_group_strips_pre_encoded_media_on_success(monkeypatch) -> None:
+    state = _NativeState(processor=object())
+    args = _native_args()
+    shared_media = {"images": [object()], "videos": [], "audio": []}
+    group = [Sample(prompt="same", multimodal_inputs=shared_media) for _ in range(2)]
+    calls = {"encode": 0}
+
+    async def fake_processor(state_arg, args_arg, prompt, multimodal_inputs):
+        return [21, 22], {"pixel_values": object()}, 0.25
+
+    async def fake_encode(multimodal_inputs):
+        calls["encode"] += 1
+        return {"image_data": ["encoded"]}, 0.5
+
+    async def fake_post(url, payload, headers=None):
+        assert payload["image_data"] == [["encoded"]]
+        return [_response(101, "a"), _response(102, "b")]
+
+    async def fake_rm(args_arg, group_arg):
+        return [1.0] * len(group_arg)
+
+    monkeypatch.setattr(sglang_rollout, "GenerateState", lambda _args: state)
+    monkeypatch.setattr(sglang_rollout, "_run_image_processor", fake_processor)
+    monkeypatch.setattr(sglang_rollout, "_encode_multimodal_inputs", fake_encode)
+    monkeypatch.setattr(sglang_rollout, "post", fake_post)
+    monkeypatch.setattr(sglang_rollout, "batched_async_rm", fake_rm)
+
+    result = await sglang_rollout.generate_and_rm_group(args, group, {"max_new_tokens": 32})
+
+    assert calls["encode"] == 1
+    for sample in result:
+        assert sample.reward == 1.0
+        assert not hasattr(sample, "_pre_encoded_mm")
+        assert not hasattr(sample, "_pre_encoded_mm_elapsed")
+        assert "_pre_encoded_mm" not in sample.to_dict()
+
+
+async def test_generate_and_rm_group_strips_pre_encoded_media_when_fallback_fails(monkeypatch) -> None:
+    state = _NativeState()
+    args = _native_args(sglang_native_group_sampling=False)
+    shared_media = {"images": [object()], "videos": [], "audio": []}
+    group = [Sample(prompt="same", multimodal_inputs=shared_media)]
+
+    async def fake_encode(multimodal_inputs):
+        return {"image_data": ["encoded"]}, 0.5
+
+    async def failing_generate_and_rm(args_arg, sample, sampling_params, evaluation=False):
+        raise RuntimeError("worker exploded")
+
+    monkeypatch.setattr(sglang_rollout, "GenerateState", lambda _args: state)
+    monkeypatch.setattr(sglang_rollout, "_encode_multimodal_inputs", fake_encode)
+    monkeypatch.setattr(sglang_rollout, "generate_and_rm", failing_generate_and_rm)
+
+    with pytest.raises(RuntimeError, match="worker exploded"):
+        await sglang_rollout.generate_and_rm_group(args, group, {"max_new_tokens": 32})
+
+    for sample in group:
+        assert not hasattr(sample, "_pre_encoded_mm")
+        assert not hasattr(sample, "_pre_encoded_mm_elapsed")
+
+
+async def test_generate_and_rm_group_warns_once_for_ineligible_groups(monkeypatch) -> None:
+    state = _NativeState()
+    args = _native_args(partial_rollout=True)
+    calls = {"generate": 0}
+    warnings = []
+
+    async def fake_generate_and_rm(args_arg, sample, sampling_params, evaluation=False):
+        calls["generate"] += 1
+        return sample
+
+    async def fake_rm(args_arg, group_arg):
+        return [0.0] * len(group_arg)
+
+    monkeypatch.setattr(sglang_rollout, "GenerateState", lambda _args: state)
+    monkeypatch.setattr(sglang_rollout, "generate_and_rm", fake_generate_and_rm)
+    monkeypatch.setattr(sglang_rollout, "batched_async_rm", fake_rm)
+    monkeypatch.setattr(sglang_rollout, "_NATIVE_GROUP_SAMPLING_LOGGED", set())
+    monkeypatch.setattr(sglang_rollout.logger, "warning", lambda message, *a, **k: warnings.append(message))
+
+    for _ in range(2):
+        group = [Sample(prompt="same"), Sample(prompt="same")]
+        await sglang_rollout.generate_and_rm_group(args, group, {"max_new_tokens": 32})
+
+    assert calls["generate"] == 4
+    assert len(warnings) == 1
+    assert "ineligible" in warnings[0]

From bdadf0004f37578c58a16bca01e2ac1ff93790c3 Mon Sep 17 00:00:00 2001
From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com>
Date: Mon, 17 Aug 2026 01:29:31 +0800
Subject: [PATCH 3/3] fix(rollout): skip empty media pre-encoding
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

# 🐛 Bug Fix

## Guard shared multimodal encoding

- Skip group-level pre-encoding when media payloads are empty
- Preserve text-only rollout behavior for empty multimodal mappings

---

# ✅ Tests

## Cover empty media groups

- Verify empty multimodal payloads bypass the encoder
- Confirm reward assignment continues through the fallback path
---
 relax/engine/rollout/sglang_rollout.py        |  3 ++-
 .../test_sglang_native_group_sampling.py      | 25 +++++++++++++++++++
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/relax/engine/rollout/sglang_rollout.py b/relax/engine/rollout/sglang_rollout.py
index 112b4cc0f..fa57f8d3b 100644
--- a/relax/engine/rollout/sglang_rollout.py
+++ b/relax/engine/rollout/sglang_rollout.py
@@ -843,7 +843,8 @@ async def generate_and_rm_group(
     # data_source), encode once and attach the result to every sample so that
     # generate() picks up the pre-encoded data instead of re-encoding per sample.
     first_mm = getattr(group[0], "multimodal_inputs", None)
-    if first_mm is not None and all(getattr(s, "multimodal_inputs", None) is first_mm for s in group[1:]):
+    has_media = first_mm is not None and any(first_mm.get(key) for key in ("images", "videos", "audio"))
+    if has_media and all(getattr(s, "multimodal_inputs", None) is first_mm for s in group[1:]):
         encoded_mm, t_enc = await _encode_multimodal_inputs(first_mm)
         for sample in group:
             sample._pre_encoded_mm = encoded_mm
diff --git a/tests/engine/rollout/test_sglang_native_group_sampling.py b/tests/engine/rollout/test_sglang_native_group_sampling.py
index 776938b82..309e2d8f6 100644
--- a/tests/engine/rollout/test_sglang_native_group_sampling.py
+++ b/tests/engine/rollout/test_sglang_native_group_sampling.py
@@ -497,6 +497,31 @@ async def failing_generate_and_rm(args_arg, sample, sampling_params, evaluation=
         assert not hasattr(sample, "_pre_encoded_mm_elapsed")
 
 
+async def test_generate_and_rm_group_skips_pre_encoding_for_empty_media(monkeypatch) -> None:
+    state = _NativeState()
+    args = _native_args(sglang_native_group_sampling=False)
+    group = [Sample(prompt="same", multimodal_inputs={})]
+
+    async def unexpected_encode(multimodal_inputs):
+        raise AssertionError("empty multimodal inputs must not be encoded")
+
+    async def fake_generate_and_rm(args_arg, sample, sampling_params, evaluation=False):
+        return sample
+
+    async def fake_rm(args_arg, group_arg):
+        return [0.0] * len(group_arg)
+
+    monkeypatch.setattr(sglang_rollout, "GenerateState", lambda _args: state)
+    monkeypatch.setattr(sglang_rollout, "_encode_multimodal_inputs", unexpected_encode)
+    monkeypatch.setattr(sglang_rollout, "generate_and_rm", fake_generate_and_rm)
+    monkeypatch.setattr(sglang_rollout, "batched_async_rm", fake_rm)
+
+    result = await sglang_rollout.generate_and_rm_group(args, group, {"max_new_tokens": 32})
+
+    assert result[0] is group[0]
+    assert result[0].reward == 0.0
+
+
 async def test_generate_and_rm_group_warns_once_for_ineligible_groups(monkeypatch) -> None:
     state = _NativeState()
     args = _native_args(partial_rollout=True)