[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. |