[diffusion] optimization: support cuda graph for Pi-0.5 prefix encoding (#34256)

docs/cookbook/vla/OpenPI/Pi0.5.mdx

@@ -81,7 +81,7 @@ These registered LeRobot checkpoints dispatch to the native `multimodal_gen` Pi0
 | `parameters.action_dim` | integer, optional | Internal padded action dimension. Defaults to the checkpoint config. |
 | `runtime.return_timing` | boolean, optional | Return stage timing fields. Defaults to `true`. |
 | `runtime.prefix_cache` | boolean or `"auto"`, optional | Enable exact full-prefix lookup for this request when the server has `enable_global_prefix_cache=true`. Defaults to `"auto"`. |
-| `runtime.cuda_graph` | boolean or `"auto"`, optional | Enable the action denoise CUDA graph path for this request when a matching shape bucket is available. Defaults to `"auto"`. |
+| `runtime.cuda_graph` | boolean or `"auto"`, optional | Enable the available prefix and action-denoise CUDA graph paths for this request. Defaults to `"auto"`. |
 | `runtime.output_format` | `"list"` or `"numpy"`, optional | Use `"list"` for JSON compatibility. Use `"numpy"` with msgpack or Python clients to avoid Python-list materialization. Defaults to `"list"`. |
 | `runtime.response_format` | `"envelope"` or `"raw"`, optional | HTTP-only response shape. `"envelope"` returns the generic action envelope. `"raw"` returns the policy payload directly. Defaults to `"envelope"`. |
 
@@ -261,7 +261,7 @@ asyncio.run(main())
 - Request-local `PrefixContext` is always reused across all denoise steps in one request. The prefix K/V is not cloned per step.
 - The optional global prefix cache is a bounded exact-match LRU. It is disabled by default because changing robot frames rarely hit it and enabling it prevents unrelated misses from entering grouped prefix execution. Set `enable_global_prefix_cache=true` for repeated observations, retries, or multiple policy calls over the same camera/state sample; `runtime.prefix_cache` can then disable lookup per request.
 - Partial-prefix reuse is not supported because Pi0.5 combines image and tokenized task/state inputs under full attention. Changing any input can change every deeper-layer prefix K/V tensor. The exact key hashes resized and normalized pixels before SigLIP, plus effective token IDs, token masks, camera masks, model revision, dtype, and parallel layout. Tensor content hashing reuses SRT's CPU/CUDA implementation; hashing the pre-SigLIP input lets an exact hit skip both the vision encoder and prefix transformer.
-- CUDA graph capture targets one action-denoise step by shape bucket, then replays it across the flow-matching loop. Shape buckets include batch size, prefix length, action horizon, action dim, dtype, and parallel layout. With action SP enabled, the graph bucket uses the local action shard length and rank-specific position offset.
+- CUDA graph capture targets single-request prefix encoding and one action-denoise step. Prefix capture uses one bounded input-shape bucket by default; grouped prefixes, prefix TP, CPU offload, and global prefix-cache misses stay eager. The denoise graph is replayed across the flow-matching loop and uses batch size, prefix length, action horizon, action dim, dtype, and parallel layout in its shape signature. With action SP enabled, the denoise bucket uses the local action shard length and rank-specific position offset.
 - Cache-DiT is not used in the default Pi0.5 path. The current robot policy target is numerically lossless inference, while Cache-DiT-style reuse is an image/video DiT approximation that needs separate policy-quality validation before it can be recommended for action control.
 - Do not use CFG parallelism to split the 10 Euler steps. Use it only for independent branches such as multiple candidate actions or future conditional/unconditional branches.
 - Prefix TP uses native SGLang parallel linear layers for the PaliGemma language prefix model when model parallel TP is initialized and the VLA split broadcast group is not active. The action expert does not share that TP layout. The v1 split prefix/action path instead uses the SP group: prefix root computes/broadcasts `PrefixContext`, while action ranks run the SP action path.
@@ -295,6 +295,8 @@ Use this single-GPU config first for 16GB-class robot workstations. It keeps par
   "materialize_dtype": "bf16",
   "enable_global_prefix_cache": false,
   "prefix_cache_max_entries": 0,
+  "enable_prefix_cuda_graph": true,
+  "prefix_cuda_graph_max_entries": 1,
   "enable_action_cuda_graph": true
 }
 ```
@@ -324,6 +326,7 @@ For a moderate fallback, offload cache growth and selected stage-resident module
   "materialize_dtype": "bf16",
   "enable_global_prefix_cache": false,
   "prefix_cache_max_entries": 0,
+  "enable_prefix_cuda_graph": false,
   "enable_action_cuda_graph": false,
   "offload_prefix_image_encoder_after_embed": true,
   "offload_prefix_token_embedding": true,
@@ -340,6 +343,7 @@ If that still does not fit, full prefix layerwise CPU offload keeps every PaliGe
   "materialize_dtype": "bf16",
   "enable_global_prefix_cache": false,
   "prefix_cache_max_entries": 0,
+  "enable_prefix_cuda_graph": false,
   "enable_action_cuda_graph": false,
   "offload_prefix_image_encoder": true,
   "offload_prefix_token_embedding": true,
@@ -353,7 +357,7 @@ Offload validation should be repeated on the target hardware after any dtype or
 
 ### 6.3 Per-Request Controls
 
-For HTTP calls, disable cache or CUDA graph without restarting the server:
+For HTTP calls, disable cache or both CUDA graph paths without restarting the server:
 
 ```python Example
 payload = {
@@ -457,6 +461,7 @@ The following checks were run on H100 GPUs with the native SGLang Pi0.5 path:
 | `lerobot/pi05_base` direct end-to-end | Prefix length `968`, output shape `[1, 50, 32]`, peak allocated memory `12.817 GiB`. |
 | Official OpenPI parity | Against OpenPI PyTorch revision `15a9616`, with the same LeRobot checkpoint revision, observation, and noise: first-step velocity max/mean absolute difference `0.02677` / `0.00344`; production 10-step normalized action `0.00813` / `0.00092`. |
 | Action denoise CUDA graph | Eager 10-step denoise `125.4 ms`; steady graph replay `50.8 ms`; max output difference `0`. |
+| Prefix CUDA graph | On H200 with action graph already enabled, ALOHA batch=1 p50 improved from `48.04 ms` to `42.98 ms` (`1.118x`). Two observations at 5 and 10 steps were bit-exact with graph disabled. One prefix shape bucket added about `48.5 MiB`; batch=4 showed no benefit and stays eager. |
 | Exact full-prefix cache | First prefix pass about `203 ms`; exact cache hit prefix stage about `0.2 ms`. |
 | `lerobot/pi05_libero_base` direct end-to-end | Image keys `image`, `image2`, `empty_camera_0`; state dim `8`; output action dim `7`; output tensor shape `[1, 50, 32]`. |
 | Python grouped execution | ALOHA batch=4 grouped path measured `91.9 ms / 4` on current mixed precision: prefix `18.4 ms`, action denoise `61.2 ms`, preprocess about `2.4 ms` per request. Sequential Python loop batch=4 measured `211.9 ms / 4`. |
@@ -466,7 +471,7 @@ The following checks were run on H100 GPUs with the native SGLang Pi0.5 path:
 | OpenPI/SGLang precision | Official OpenPI JAX inference restores the public GCS checkpoint as bf16 with selected fp32 stability compute and returns float32 actions. The converted OpenPI PyTorch `pi05_aloha` checkpoint keeps `119,720,608` fp32 stability params; SGLang reports the same fp32 set and `3,233,713,264` bf16 runtime params after skipping unused LM heads. |
 | Native attention dtype | Checkpoint source tensors may be fp32, but SGLang finalizes PiGemma and SigLIP compute dtype before native attention backend selection; backend logs showed `Using fa attention backend` for the PiGemma path in the prior run. |
 | 16GB-free Python pressure | With an H100 artificially constrained to `16381 MiB` free before model load, single-GPU bf16 no-offload Python grouped path completed without OOM. Re-run latency after precision or loader changes before using pressure numbers for deployment sizing. |
-| Low-VRAM switches | Disabling prefix cache prevents cache growth across changing robot frames. CUDA graph can stay enabled when the action expert remains resident; disable it only for offload fallback modes. |
+| Low-VRAM switches | Disabling prefix cache prevents cache growth across changing robot frames. Prefix graph residency is bounded by `prefix_cuda_graph_max_entries` (default `1`); set `enable_prefix_cuda_graph=false` or the limit to `0` to save about `48.5 MiB` for the validated ALOHA bucket. Action graph can stay enabled when the action expert remains resident; disable both graph paths for offload fallback modes. |
 | Offload fallback | CPU/offload modes are retained as numerically lossless compatibility fallbacks, but earlier fp32-runtime offload latency numbers are stale after the bf16 dtype correction and should be revalidated before deployment decisions. |
 | Run:ai direct loader | Single-GPU serve streamed `13.5 GiB` safetensors to `cuda:0` in about `1.5 s` and returned `[50, 32]` actions. Distributed direct streaming is now rank-local and should be revalidated on the target split topology; offload ranks with CPU targets still use the safe loader. |
 | OpenPI comparison status | Official OpenPI GCS `pi05_base` is a JAX checkpoint; converted PyTorch eager was validated without `torch.compile`. On 80GB H100, ALOHA OpenPI PyTorch eager was about `125-130 ms` single and about `164 ms / 4` in the direct-model batch path. Current SGLang Python grouped measured `52.4 ms` single and `91.9 ms / 4`; JAX OpenPI was `53.0 ms` single and `59.5 ms / 2` in a short check. |

python/sglang/multimodal_gen/configs/pipeline_configs/pi05.py

@@ -46,6 +46,8 @@ class Pi05PipelineConfig(PipelineConfig):
     image_normalization_std: tuple[float, float, float] = (0.5, 0.5, 0.5)
 
     enable_global_prefix_cache: bool = False
+    enable_prefix_cuda_graph: bool = True
+    prefix_cuda_graph_max_entries: int = 1
     enable_action_cuda_graph: bool = True
     prefix_cache_max_entries: int = 1
     prefix_cache_layout_version: str = "pi05-prefix-v1"

python/sglang/multimodal_gen/runtime/models/vlas/pi05_core.py

@@ -1327,25 +1327,21 @@ class Pi05CoreModel(nn.Module):
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         embs = []
         pad_masks = []
-        att_masks = []
         image_embs = self.paligemma_with_expert.embed_images(images)
         for image_emb, image_mask in zip(image_embs, image_masks, strict=True):
             batch_size, num_image_embs = image_emb.shape[:2]
             embs.append(image_emb)
             pad_masks.append(image_mask[:, None].expand(batch_size, num_image_embs))
-            att_masks += [0] * num_image_embs
 
         lang_emb = self.paligemma_with_expert.embed_language_tokens(tokens)
         # Match OpenPI's Pi0.5 prefix embedding semantics.
         lang_emb = lang_emb * math.sqrt(lang_emb.shape[-1])
         embs.append(lang_emb)
         pad_masks.append(token_masks)
-        att_masks += [0] * lang_emb.shape[1]
 
         embs = torch.cat(embs, dim=1)
         pad_masks = torch.cat(pad_masks, dim=1)
-        att_masks_t = torch.tensor(att_masks, dtype=torch.bool, device=pad_masks.device)
-        att_masks_t = att_masks_t[None, :].expand(pad_masks.shape[0], len(att_masks))
+        att_masks_t = torch.zeros_like(pad_masks)
         return embs, pad_masks, att_masks_t
 
     def embed_suffix(

python/sglang/multimodal_gen/runtime/models/vlas/pi05_policy.py

@@ -38,9 +38,11 @@ from sglang.multimodal_gen.runtime.models.vlas.pi05_core import Pi05CoreModel
 from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import maybe_download_model
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
-from sglang.multimodal_gen.runtime.vla.denoise_cuda_graph import (
+from sglang.multimodal_gen.runtime.vla.cuda_graph import (
     VLADenoiseGraphRunner,
     VLADenoiseGraphSignature,
+    VLAPrefixGraphRunner,
+    VLAPrefixGraphSignature,
 )
 from sglang.multimodal_gen.runtime.vla.observation import (
     VLAObservationBatch,
@@ -161,6 +163,10 @@ class Pi05PolicyModel(nn.Module):
         if self.runtime_role != "all":
             logger.info("Pi05 split runtime role on rank: %s", self.runtime_role)
         self.action_expert = Pi05ActionExpert(config, self.core_model)
+        self.prefix_graph_runner = VLAPrefixGraphRunner(
+            enabled=self._prefix_cuda_graph_enabled(),
+            max_entries=config.prefix_cuda_graph_max_entries,
+        )
         self.graph_runner = VLADenoiseGraphRunner(
             enabled=config.enable_action_cuda_graph
         )
@@ -176,6 +182,25 @@ class Pi05PolicyModel(nn.Module):
             return False
         return get_tp_world_size() > 1
 
+    def _prefix_cuda_graph_enabled(self) -> bool:
+        if (
+            not self.config.enable_prefix_cuda_graph
+            or self.device.type != "cuda"
+            or self.runtime_role not in ("all", "prefix")
+            or self._prefix_tensor_parallel_enabled()
+        ):
+            return False
+        return not any(
+            (
+                self.config.offload_prefix_image_encoder,
+                self.config.offload_prefix_image_encoder_after_embed,
+                self.config.offload_prefix_token_embedding,
+                self.config.offload_prefix_language_layers,
+                self.config.offload_prefix_language_layers_after_prefix,
+                self.config.empty_cache_after_prefix,
+            )
+        )
+
     @staticmethod
     def _to_empty_preserve_buffers(module: nn.Module, *, device: torch.device) -> None:
         buffers = {
@@ -796,6 +821,10 @@ class Pi05PolicyModel(nn.Module):
             return None
         return paligemma.model.language_model
 
+    def _prefix_tensor_parallel_enabled(self) -> bool:
+        language_model = self._prefix_language_model()
+        return language_model is not None and language_model.tensor_parallel
+
     def _prefix_kv_requires_tp_gather(self) -> bool:
         language_model = self._prefix_language_model()
         if language_model is None or not language_model.tensor_parallel:
@@ -822,7 +851,12 @@ class Pi05PolicyModel(nn.Module):
             )
         )
 
-    def encode_prefix(self, observation: VLAObservationBatch) -> PrefixContext:
+    def encode_prefix(
+        self,
+        observation: VLAObservationBatch,
+        *,
+        use_cuda_graph: bool = True,
+    ) -> PrefixContext:
         camera_order = tuple(observation.metadata.get("camera_order", ()))
         images = [
             observation.images[name].to(self.device, dtype=torch.float32)
@@ -844,22 +878,57 @@ class Pi05PolicyModel(nn.Module):
         prefix_full_attention_hint = all(
             bool(observation.image_masks[name].all().item()) for name in camera_order
         ) and bool(token_masks_cpu.all().item())
-        past_key_values, prefix_pad_masks, full_attention = (
-            self.core_model.encode_prefix(
-                images,
-                image_masks,
-                tokens,
-                token_masks,
-                prefix_full_attention_hint=prefix_full_attention_hint,
-                tokens_trimmed=tokens_trimmed,
+
+        image_count = len(images)
+        graph_inputs = tuple([*images, *image_masks, tokens, token_masks])
+
+        def encode(current_inputs: tuple[torch.Tensor, ...]) -> PrefixContext:
+            current_images = list(current_inputs[:image_count])
+            current_image_masks = list(current_inputs[image_count : 2 * image_count])
+            current_tokens = current_inputs[-2]
+            current_token_masks = current_inputs[-1]
+            past_key_values, prefix_pad_masks, full_attention = (
+                self.core_model.encode_prefix(
+                    current_images,
+                    current_image_masks,
+                    current_tokens,
+                    current_token_masks,
+                    prefix_full_attention_hint=prefix_full_attention_hint,
+                    tokens_trimmed=tokens_trimmed,
+                )
             )
+            past_key_values = self._materialize_prefix_kv_for_action(past_key_values)
+            return PrefixContext(
+                past_key_values=past_key_values,
+                prefix_pad_masks=prefix_pad_masks,
+                prefix_len=prefix_pad_masks.shape[1],
+                layout={"full_attention": full_attention},
+            )
+
+        if (
+            not use_cuda_graph
+            or not self.prefix_graph_runner.enabled
+            or observation.batch_size != 1
+        ):
+            return encode(graph_inputs)
+
+        signature = VLAPrefixGraphSignature(
+            batch_size=observation.batch_size,
+            input_shapes=tuple(tuple(tensor.shape) for tensor in graph_inputs),
+            input_dtypes=tuple(
+                str(tensor.dtype).replace("torch.", "") for tensor in graph_inputs
+            ),
+            static_layout=(
+                image_count,
+                prefix_full_attention_hint,
+                tokens_trimmed,
+            ),
+            parallel_layout=self.config.parallel_layout_version,
         )
-        past_key_values = self._materialize_prefix_kv_for_action(past_key_values)
-        return PrefixContext(
-            past_key_values=past_key_values,
-            prefix_pad_masks=prefix_pad_masks,
-            prefix_len=prefix_pad_masks.shape[1],
-            layout={"full_attention": full_attention},
+        return self.prefix_graph_runner.capture_or_run(
+            signature,
+            encode,
+            graph_inputs,
         )
 
     def sample_noise(

python/sglang/multimodal_gen/runtime/pipelines/pi05.py

@@ -59,13 +59,15 @@ class Pi05Pipeline(ComposedPipelineBase):
             or bool(server_args.text_encoder_cpu_offload)
         )
         logger.info(
-            "Pi05 memory config: prefix_cache=%s/%s, action_cuda_graph=%s, "
+            "Pi05 memory config: prefix_cache=%s/%s, cuda_graph=%s/%s/%s, "
             "offload_image=%s, offload_image_after_embed=%s, "
             "offload_tokens=%s, offload_language_layers=%s, "
             "offload_language_after_prefix=%s/%s, "
             "offload_action_after_denoise=%s, empty_cache_after_prefix=%s",
             pipeline_config.enable_global_prefix_cache,
             pipeline_config.prefix_cache_max_entries,
+            pipeline_config.enable_prefix_cuda_graph,
+            pipeline_config.prefix_cuda_graph_max_entries,
             pipeline_config.enable_action_cuda_graph,
             pipeline_config.offload_prefix_image_encoder,
             pipeline_config.offload_prefix_image_encoder_after_embed,

python/sglang/multimodal_gen/runtime/pipelines_core/stages/vla.py

@@ -89,6 +89,10 @@ def _effective_prefix_cache_enabled(
     )
 
 
+def _cuda_graph_enabled(batch: Req) -> bool:
+    return bool(vla_options(batch).get("enable_cuda_graph", True))
+
+
 def _grouped_fingerprint(
     batch: Req,
     server_args: ServerArgs,
@@ -122,6 +126,7 @@ def _grouped_fingerprint(
         batch.action_horizon,
         batch.action_dim,
         batch.num_inference_steps,
+        _cuda_graph_enabled(batch),
     )
 
 
@@ -180,7 +185,10 @@ class VLAPrefixEncodingStage(PipelineStage):
                 vla_state(batch)["observation_batch"] for batch in group_batches
             ]
             grouped_observation = collate_vla_observation_batches(observations)
-            prefix_context = self.policy_model.encode_prefix(grouped_observation)
+            prefix_context = self.policy_model.encode_prefix(
+                grouped_observation,
+                use_cuda_graph=_cuda_graph_enabled(group_batches[0]),
+            )
             prefix_ms = (time.perf_counter() - prefix_start) * 1000
 
             for offset, (index, batch) in enumerate(group):
@@ -299,7 +307,10 @@ class VLAPrefixEncodingStage(PipelineStage):
         prefix_start = time.perf_counter()
 
         # 3. run encoding
-        prefix_context = self.policy_model.encode_prefix(observation)
+        prefix_context = self.policy_model.encode_prefix(
+            observation,
+            use_cuda_graph=_cuda_graph_enabled(batch) and not cache_enabled,
+        )
         if cache_key is not None:
             prefix_context.cache_key_digest = cache_key
 
@@ -371,7 +382,7 @@ class VLAActionDenoisingStage(PipelineStage):
                 prefix_context,
                 noise=observation.noise,
                 num_steps=group_batches[0].num_inference_steps,
-                use_cuda_graph=bool(options.get("enable_cuda_graph", True)),
+                use_cuda_graph=_cuda_graph_enabled(group_batches[0]),
                 generator=None,
             )
             synchronize_vla_action_tensor(actions)
@@ -417,7 +428,7 @@ class VLAActionDenoisingStage(PipelineStage):
                 prefix_context,
                 noise=noise,
                 num_steps=batch.num_inference_steps,
-                use_cuda_graph=bool(options.get("enable_cuda_graph", True)),
+                use_cuda_graph=_cuda_graph_enabled(batch),
                 generator=batch.generator,
             )
             synchronize_vla_action_tensor(actions)

python/sglang/multimodal_gen/runtime/vla/cuda_graph.py

@@ -22,6 +22,22 @@ from sglang.srt.model_executor.runner_utils.pool import (
 logger = init_logger(__name__)
 
 
+@dataclass(frozen=True)
+class VLAPrefixGraphSignature:
+    batch_size: int
+    input_shapes: tuple[tuple[int, ...], ...]
+    input_dtypes: tuple[str, ...]
+    static_layout: tuple[Any, ...]
+    parallel_layout: str
+
+
+@dataclass
+class _CapturedPrefixGraph:
+    graph: torch.cuda.CUDAGraph
+    static_inputs: tuple[torch.Tensor, ...]
+    static_output: PrefixContext
+
+
 @dataclass(frozen=True)
 class VLADenoiseGraphSignature:
     batch_size: int
@@ -76,6 +92,103 @@ def _copy_prefix_context_(dst: PrefixContext, src: PrefixContext) -> None:
     dst.cache_key_digest = src.cache_key_digest
 
 
+class VLAPrefixGraphRunner:
+    """Full CUDA graph runner for VLA prefix encoding shape buckets."""
+
+    def __init__(self, enabled: bool = True, max_entries: int = 1):
+        self.max_entries = max(0, max_entries)
+        self.enabled = enabled and self.max_entries > 0
+        self._captured: dict[VLAPrefixGraphSignature, _CapturedPrefixGraph] = {}
+        self._disabled_signatures: set[VLAPrefixGraphSignature] = set()
+        self._capture_stream: torch.cuda.Stream | None = None
+        self._graph_pool: Any = None
+
+    def _capture(
+        self,
+        signature: VLAPrefixGraphSignature,
+        step_fn: Callable[[tuple[torch.Tensor, ...]], PrefixContext],
+        inputs: tuple[torch.Tensor, ...],
+    ) -> _CapturedPrefixGraph:
+        static_inputs = tuple(tensor.detach().clone() for tensor in inputs)
+        device_module = torch.get_device_module(inputs[0].device)
+        if self._capture_stream is None:
+            self._capture_stream = device_module.Stream(device=inputs[0].device)
+        if self._graph_pool is None:
+            self._graph_pool = get_or_create_global_graph_memory_pool(device_module)
+            set_graph_pool_id(self._graph_pool)
+
+        device_module.synchronize()
+        with device_module.stream(self._capture_stream), torch.inference_mode():
+            step_fn(static_inputs)
+        self._capture_stream.synchronize()
+
+        graph = torch.cuda.CUDAGraph()
+        with (
+            device_module.graph(
+                cuda_graph=graph,
+                pool=self._graph_pool,
+                stream=self._capture_stream,
+            ),
+            torch.inference_mode(),
+        ):
+            static_output = step_fn(static_inputs)
+        self._capture_stream.synchronize()
+        static_output.layout["mutable_graph_output"] = True
+
+        captured = _CapturedPrefixGraph(
+            graph=graph,
+            static_inputs=static_inputs,
+            static_output=static_output,
+        )
+        self._captured[signature] = captured
+        logger.info(
+            "Captured VLA prefix CUDA graph: batch=%d inputs=%s",
+            signature.batch_size,
+            signature.input_shapes,
+        )
+        return captured
+
+    def capture_or_run(
+        self,
+        signature: VLAPrefixGraphSignature,
+        step_fn: Callable[[tuple[torch.Tensor, ...]], PrefixContext],
+        inputs: tuple[torch.Tensor, ...],
+    ) -> PrefixContext:
+        if (
+            not self.enabled
+            or signature in self._disabled_signatures
+            or not inputs
+            or inputs[0].device.type != "cuda"
+        ):
+            return step_fn(inputs)
+
+        captured = self._captured.get(signature)
+        if captured is None and len(self._captured) >= self.max_entries:
+            return step_fn(inputs)
+        try:
+            if captured is None:
+                captured = self._capture(signature, step_fn, inputs)
+            else:
+                for static_input, current_input in zip(
+                    captured.static_inputs, inputs, strict=True
+                ):
+                    static_input.copy_(current_input)
+            captured.graph.replay()
+            torch.get_device_module(inputs[0].device).current_stream(
+                device=inputs[0].device
+            ).synchronize()
+            return captured.static_output
+        except Exception:
+            self._disabled_signatures.add(signature)
+            self._captured.pop(signature, None)
+            logger.warning(
+                "VLA prefix CUDA graph disabled for signature %s",
+                signature,
+                exc_info=True,
+            )
+            return step_fn(inputs)
+
+
 class VLADenoiseGraphRunner:
     """Full CUDA graph runner for one VLA action-denoise step.
 
@@ -97,13 +210,16 @@ class VLADenoiseGraphRunner:
     ) -> None:
         context_id = id(prefix_context.past_key_values)
         context_digest = prefix_context.cache_key_digest
-        if (
+        mutable_graph_output = bool(
+            prefix_context.layout.get("mutable_graph_output", False)
+        )
+        if not mutable_graph_output and (
             context_digest is not None
             and captured.current_context_digest == context_digest
         ):
             captured.current_context_id = context_id
             return
-        if captured.current_context_id == context_id:
+        if not mutable_graph_output and captured.current_context_id == context_id:
             return
         _copy_prefix_context_(captured.static_prefix_context, prefix_context)
         captured.current_context_id = context_id

python/sglang/multimodal_gen/test/unit/test_pi05_runtime_helpers.py

@@ -21,7 +21,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.p
     _preprocess_image,
     _resize_with_pad_image_tensor,
 )
-from sglang.multimodal_gen.runtime.vla.denoise_cuda_graph import (
+from sglang.multimodal_gen.runtime.vla.cuda_graph import (
     VLADenoiseGraphRunner,
     _CapturedDenoiseGraph,
 )
@@ -74,7 +74,7 @@ def test_denoise_graph_skips_prefix_copy_for_same_digest(monkeypatch):
         raise AssertionError("PrefixContext should not be copied on digest hit")
 
     monkeypatch.setattr(
-        "sglang.multimodal_gen.runtime.vla.denoise_cuda_graph._copy_prefix_context_",
+        "sglang.multimodal_gen.runtime.vla.cuda_graph._copy_prefix_context_",
         fail_copy,
     )
 
@@ -83,6 +83,35 @@ def test_denoise_graph_skips_prefix_copy_for_same_digest(monkeypatch):
     assert captured.static_prefix_context.past_key_values[0][0].eq(1.0).all()
 
 
+def test_denoise_graph_copies_mutable_prefix_graph_output():
+    runner = VLADenoiseGraphRunner(enabled=True)
+    static_context = _prefix_context(1.0, None)
+    captured = _CapturedDenoiseGraph(
+        graph=object(),
+        static_prefix_context=static_context,
+        static_x_t=torch.empty(1, 2, 4),
+        static_timestep=torch.empty(1),
+        static_output=torch.empty(1, 2, 4),
+        current_context_id=123,
+    )
+    current_context = _prefix_context(2.0, None)
+    current_context.layout["mutable_graph_output"] = True
+
+    runner._sync_context_if_needed(captured, current_context)
+
+    assert captured.static_prefix_context.past_key_values[0][0].eq(2.0).all()
+
+
+def test_prefix_graph_rejects_tensor_parallel_prefix():
+    model = Pi05PolicyModel.__new__(Pi05PolicyModel)
+    model.config = Pi05PipelineConfig()
+    model.device = torch.device("cuda")
+    model.runtime_role = "all"
+    model._prefix_language_model = lambda: SimpleNamespace(tensor_parallel=True)
+
+    assert not model._prefix_cuda_graph_enabled()
+
+
 def test_runai_direct_gpu_loader_does_not_reject_split_roles(monkeypatch):
     class FakeSafeOpen:
         def __enter__(self):