[Model] GR00T-N1.7: support request-level batching

recipes/NVIDIA/GR00T-N1.7.md

@@ -49,11 +49,11 @@ vllm serve nvidia/GR00T-N1.7-3B \
 
 Notes:
 
-- Only `max_num_seqs: 1` is supported (configured in the deploy YAML). The
-  policy is markovian (`state_history_length=1`, `reset()` returns `{}`) — the
-  reason for the cap is that `Gr00tPolicy` does its own per-sample
-  (un)batching inside `get_action` and is not integrated with vLLM's
-  continuous batching path.
+- The deploy YAML defaults to `max_num_seqs: 1`, serving one request at a
+  time. The policy is markovian (`state_history_length=1`, `reset()` returns
+  `{}`), and `Gr00tN1d7Pipeline` implements the request-batch contract, so
+  raising `max_num_seqs` in an overlay batches concurrent sessions into one
+  policy call — see "Optional: batch concurrent sessions" below.
 - This pipeline is a thin wrapper around the upstream HF policy
   (`AutoModel.from_pretrained` + `enforce_eager`, `tensor_parallel_size=1`,
   pipeline `load_weights` is a no-op). The standard diffusion accelerators
@@ -91,6 +91,55 @@ The test sends a synthetic two-frame DROID observation and checks:
 - All action values are finite float32
 - Reset response is `"reset successful"`
 
+## Optional: batch concurrent sessions
+
+The bundled deploy config serves one request at a time (`max_num_seqs: 1`):
+with N robots connected, every observation waits for the requests ahead of it,
+so latency grows linearly with N while throughput stays flat.
+`Gr00tN1d7Pipeline` supports request-level batching — the scheduler groups
+in-flight requests into one wave and the pipeline serves the whole wave with a
+single `Gr00tPolicy.get_action()` call. Opt in with an overlay:
+
+```yaml
+# gr00t_batch.yaml
+base_config: vllm_omni/deploy/Gr00tN1d7.yaml
+stages:
+  - stage_id: 0
+    max_num_seqs: 8
+```
+
+```bash
+vllm serve nvidia/GR00T-N1.7-3B --omni --deploy-config gr00t_batch.yaml
+```
+
+Measured on 1xA30 (eager, DP=1, 30 requests/client after 20 warmup requests),
+closed-loop OpenPI WebSocket clients:
+
+| concurrent sessions | `max_num_seqs: 1` (default) | `max_num_seqs: 8` |
+| --- | --- | --- |
+| 1 | 6.56 req/s, 152 ms | 6.76 req/s, 148 ms |
+| 2 | 6.68 req/s, 297 ms | 6.94 req/s, 286 ms |
+| 4 | 6.71 req/s, 589 ms | 10.89 req/s, 366 ms |
+| 8 | 6.77 req/s, 1165 ms | 13.84 req/s, 576 ms |
+
+Notes:
+
+- Batches form opportunistically: requests that arrive while a wave is running
+  are grouped into the next wave, so a single client is served exactly as
+  before (wave size 1, no added wait). `request_batch_max_wait_ms` (default 0)
+  can trade first-request latency for fuller waves.
+- Actions for a request depend only on that request's observation — a
+  neighbor's content never leaks into your actions. Batched execution does
+  change bf16 kernel shapes, which shifts actions on the order of 5e-2 versus
+  a solo run (1e-5 when the same comparison is done in fp32) — well inside the
+  policy's own seed-to-seed sampling spread (~8e-1). With a fixed
+  `GR00T_NOISE_SEED`, per-request noise additionally depends on the wave size,
+  so bit-exact reproducibility holds per wave shape, not across load levels.
+  The single-session golden test is unaffected (wave size 1).
+- Batching grows the per-wave activation footprint only modestly: waves of 7
+  peak at ~6.4 GiB reserved vs ~5.9 GiB for single-request serving on A30
+  (the model dominates; per-observation activations are small).
+
 ## Notes
 
 - **Do not change the model-specific `policy_server_config` values.** `action_horizon`,

tests/diffusion/models/gr00t/test_pipeline.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from types import SimpleNamespace
 
 import numpy as np
@@ -6,18 +9,19 @@
 from vllm_omni.diffusion.models.gr00t import pipeline_gr00t
 from vllm_omni.diffusion.models.gr00t.pipeline_gr00t import Gr00tN1d7Pipeline
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 class FakeGr00tPolicy:
-    instances = []
+    instances: list["FakeGr00tPolicy"] = []
 
     def __init__(self, **kwargs):
         self.kwargs = kwargs
         self.reset_calls = 0
-        self.seen_obs = None
+        self.calls = []
         self.embodiment_tag = SimpleNamespace(value="fake_embodiment")
         self.language_key = "annotation.language.language_instruction"
         self.modality_configs = {
@@ -41,11 +45,14 @@ def __init__(self, **kwargs):
         FakeGr00tPolicy.instances.append(self)
 
     def get_action(self, obs):
-        self.seen_obs = obs
-        return {
-            "arm": np.array([[[1.0, 2.0]]], dtype=np.float64),
-            "gripper": [[[3.0]]],
-        }, {"latency_ms": 1.0}
+        self.calls.append(obs)
+        batch_size = len(next(iter(obs["video"].values())))
+        # Echo each sample's state marker into its actions so tests can check
+        # that batched outputs are sliced back to the right request.
+        marker = np.asarray(obs["state"]["joint"], dtype=np.float64)[:, 0, 0]
+        arm = np.broadcast_to(marker[:, None, None], (batch_size, 2, 2)).copy()
+        gripper = np.broadcast_to(marker[:, None, None], (batch_size, 2, 1)).tolist()
+        return {"arm": arm, "gripper": gripper}, {"latency_ms": 1.0}
 
     def reset(self):
         self.reset_calls += 1
@@ -70,6 +77,26 @@ def _pipeline():
     return Gr00tN1d7Pipeline(od_config=od_config)
 
 
+def _robot_request(
+    request_id: str, *, marker: float, prompt: str = "pick the cube", side: int = 8, reset: bool = False
+):
+    return OmniDiffusionRequest(
+        prompt="pick",
+        request_id=request_id,
+        sampling_params=OmniDiffusionSamplingParams(
+            extra_args={
+                "robot_obs": {
+                    "images": {"cam": np.full((1, 1, side, side, 3), int(marker), dtype=np.uint8)},
+                    "state": {"joint": np.full((1, 1, 2), marker, dtype=np.float32)},
+                    "prompt": prompt,
+                    "session_id": f"session-{request_id}",
+                },
+                "reset": reset,
+            }
+        ),
+    )
+
+
 def test_pipeline_initializes_local_policy():
     pipeline = _pipeline()
 
@@ -79,40 +106,135 @@ def test_pipeline_initializes_local_policy():
     assert policy.kwargs["strict"] is False
     assert pipeline.weights_sources == ()
     assert pipeline.load_weights(iter(())) == set()
+    assert Gr00tN1d7Pipeline.supports_request_batch is True
 
 
 def test_forward_returns_dict_actions_in_output():
     pipeline = _pipeline()
-    req = OmniDiffusionRequest(
+    batch = DiffusionRequestBatch(requests=[_robot_request("req", marker=1.0, reset=True)])
+
+    outputs = pipeline.forward(batch)
+
+    assert len(outputs) == 1
+    output = outputs[0]
+    assert output.error is None
+    actions = output.output["actions"]
+    assert set(actions) == {"arm", "gripper"}
+    assert actions["arm"].dtype == np.float32
+    np.testing.assert_allclose(actions["arm"], np.ones((1, 2, 2), dtype=np.float32))
+    policy = FakeGr00tPolicy.instances[0]
+    seen_obs = policy.calls[0]
+    assert "video" in seen_obs
+    assert seen_obs["language"] == {"annotation.language.language_instruction": [["pick the cube"]]}
+    assert "images" not in seen_obs
+    assert "prompt" not in seen_obs
+    assert "session_id" not in seen_obs
+    assert policy.reset_calls == 1
+
+
+def test_forward_merges_wave_into_single_policy_call():
+    pipeline = _pipeline()
+    batch = DiffusionRequestBatch(
+        requests=[
+            _robot_request("req-a", marker=1.0, prompt="pick the cube"),
+            _robot_request("req-b", marker=2.0, prompt="open the drawer"),
+        ]
+    )
+
+    outputs = pipeline.forward(batch)
+
+    policy = FakeGr00tPolicy.instances[0]
+    assert len(policy.calls) == 1
+    merged_obs = policy.calls[0]
+    assert merged_obs["video"]["cam"].shape == (2, 1, 8, 8, 3)
+    assert merged_obs["state"]["joint"].shape == (2, 1, 2)
+    assert merged_obs["language"] == {
+        "annotation.language.language_instruction": [["pick the cube"], ["open the drawer"]]
+    }
+
+    assert [output.error for output in outputs] == [None, None]
+    np.testing.assert_allclose(outputs[0].output["actions"]["arm"], np.ones((1, 2, 2), dtype=np.float32))
+    np.testing.assert_allclose(outputs[1].output["actions"]["arm"], np.full((1, 2, 2), 2.0, dtype=np.float32))
+    for output in outputs:
+        assert output.output["actions"]["gripper"].shape == (1, 2, 1)
+        assert output.output["actions"]["gripper"].dtype == np.float32
+
+
+def test_forward_wave_isolates_malformed_request():
+    pipeline = _pipeline()
+    bad_request = OmniDiffusionRequest(
         prompt="pick",
-        request_id="req",
+        request_id="req-bad",
+        sampling_params=OmniDiffusionSamplingParams(extra_args={"robot_obs": "not-a-dict"}),
+    )
+    batch = DiffusionRequestBatch(requests=[_robot_request("req-a", marker=3.0), bad_request])
+
+    outputs = pipeline.forward(batch)
+
+    assert outputs[0].error is None
+    np.testing.assert_allclose(outputs[0].output["actions"]["arm"], np.full((1, 2, 2), 3.0, dtype=np.float32))
+    assert "robot_obs must be a dict" in outputs[1].error
+    # The valid request still runs through the single-observation path.
+    assert len(FakeGr00tPolicy.instances[0].calls) == 1
+
+
+def test_forward_wave_serves_dummy_and_real_requests():
+    pipeline = _pipeline()
+    dummy = OmniDiffusionRequest(
+        prompt="dummy run",
+        request_id="dummy_req_id",
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+    )
+    batch = DiffusionRequestBatch(requests=[dummy, _robot_request("req-a", marker=4.0)])
+
+    outputs = pipeline.forward(batch)
+
+    dummy_actions = outputs[0].output["actions"]
+    assert dummy_actions["arm"].shape == (1, 2, 2)
+    assert not dummy_actions["arm"].any()
+    np.testing.assert_allclose(outputs[1].output["actions"]["arm"], np.full((1, 2, 2), 4.0, dtype=np.float32))
+
+
+def test_forward_wave_survives_garbage_modality():
+    pipeline = _pipeline()
+    garbage = OmniDiffusionRequest(
+        prompt="pick",
+        request_id="req-garbage",
         sampling_params=OmniDiffusionSamplingParams(
             extra_args={
                 "robot_obs": {
                     "images": {"cam": np.zeros((1, 1, 8, 8, 3), dtype=np.uint8)},
-                    "state": {"joint": np.zeros((1, 1, 2), dtype=np.float32)},
-                    "prompt": "pick the cube",
-                    "session_id": "session-a",
-                },
-                "reset": True,
+                    "state": 5,
+                    "prompt": "pick",
+                }
             }
         ),
     )
+    batch = DiffusionRequestBatch(requests=[_robot_request("req-a", marker=6.0), garbage])
 
-    output = pipeline.forward(req)
+    outputs = pipeline.forward(batch)
 
-    assert output.error is None
-    actions = output.output["actions"]
-    assert set(actions) == {"arm", "gripper"}
-    assert actions["arm"].dtype == np.float32
-    np.testing.assert_allclose(actions["arm"], np.array([[[1.0, 2.0]]], dtype=np.float32))
-    policy = FakeGr00tPolicy.instances[0]
-    assert "video" in policy.seen_obs
-    assert policy.seen_obs["language"] == {"annotation.language.language_instruction": [["pick the cube"]]}
-    assert "images" not in policy.seen_obs
-    assert "prompt" not in policy.seen_obs
-    assert "session_id" not in policy.seen_obs
-    assert policy.reset_calls == 1
+    assert outputs[0].error is None
+    np.testing.assert_allclose(outputs[0].output["actions"]["arm"], np.full((1, 2, 2), 6.0, dtype=np.float32))
+    assert "GR00T policy inference failed" in outputs[1].error
+
+
+def test_forward_wave_falls_back_on_incompatible_observations():
+    pipeline = _pipeline()
+    batch = DiffusionRequestBatch(
+        requests=[
+            _robot_request("req-a", marker=1.0, side=8),
+            _robot_request("req-b", marker=2.0, side=4),
+        ]
+    )
+
+    outputs = pipeline.forward(batch)
+
+    # Video shapes cannot be concatenated, so each request gets its own call.
+    assert len(FakeGr00tPolicy.instances[0].calls) == 2
+    assert [output.error for output in outputs] == [None, None]
+    np.testing.assert_allclose(outputs[0].output["actions"]["arm"], np.ones((1, 2, 2), dtype=np.float32))
+    np.testing.assert_allclose(outputs[1].output["actions"]["arm"], np.full((1, 2, 2), 2.0, dtype=np.float32))
 
 
 def test_dummy_warmup_returns_shape_correct_zero_actions():
@@ -123,15 +245,16 @@ def test_dummy_warmup_returns_shape_correct_zero_actions():
         sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
     )
 
-    output = pipeline.forward(req)
+    outputs = pipeline.forward(DiffusionRequestBatch(requests=[req]))
 
-    assert output.error is None
-    actions = output.output["actions"]
+    assert len(outputs) == 1
+    assert outputs[0].error is None
+    actions = outputs[0].output["actions"]
     assert set(actions) == {"arm", "gripper"}
     assert actions["arm"].shape == (1, 2, 2)
     assert actions["gripper"].shape == (1, 2, 1)
     assert not actions["arm"].any()
-    assert FakeGr00tPolicy.instances[0].seen_obs is None
+    assert FakeGr00tPolicy.instances[0].calls == []
 
 
 def test_reset_delegates_to_policy():

tests/engine/test_async_omni_engine_stage_init.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import concurrent.futures
 import contextlib
 import importlib
@@ -513,6 +516,39 @@ def _capture_client(model, config, metadata, stage_init_timeout, batch_size, use
     }
 
 
+def test_initialize_diffusion_stage_default_batch_size_keeps_yaml_max_num_seqs(monkeypatch):
+    """The default batch_size (1) must not clobber the deploy YAML's max_num_seqs.
+
+    Online serving never passes diffusion_batch_size, so the stage YAML is its
+    only way to enable request batching.
+    """
+    import vllm_omni.diffusion.stage_diffusion_client as client_mod
+    import vllm_omni.engine.stage_init_utils as init_mod
+
+    od_config = types.SimpleNamespace(max_num_seqs=8)
+    captured: dict[str, object] = {}
+    monkeypatch.setattr(init_mod, "build_diffusion_config", lambda *args: od_config)
+
+    def _capture_client(model, config, metadata, stage_init_timeout, batch_size, use_inline):
+        captured.update(batch_size=batch_size)
+        return object()
+
+    monkeypatch.setattr(client_mod, "create_diffusion_client", _capture_client)
+
+    init_mod.initialize_diffusion_stage(
+        0,
+        "dummy-model",
+        types.SimpleNamespace(),
+        _make_diffusion_metadata(0),
+        stage_init_timeout=12,
+        batch_size=1,
+        use_inline=True,
+    )
+
+    assert od_config.max_num_seqs == 8
+    assert captured["batch_size"] == 8
+
+
 def test_launch_diffusion_stage_replica_applies_batch_size_to_config(monkeypatch):
     import vllm_omni.diffusion.stage_diffusion_client as client_mod
     import vllm_omni.diffusion.stage_diffusion_proc as proc_mod

vllm_omni/deploy/Gr00tN1d7.yaml

@@ -10,6 +10,10 @@ async_chunk: false
 
 stages:
   - stage_id: 0
+    # Kept at 1 so the default deployment serves one request at a time.
+    # Gr00tN1d7Pipeline supports request-level batching, so raising this in an
+    # overlay lets concurrent OpenPI sessions share one policy forward.
+    # Opt-in instructions are in recipes/NVIDIA/GR00T-N1.7.md.
     max_num_seqs: 1
     enforce_eager: true
     model_class_name: Gr00tN1d7Pipeline