diff --git a/docs/cookbook/vla/OpenPI/Pi0.5.mdx b/docs/cookbook/vla/OpenPI/Pi0.5.mdx index 166eb2f2d..0111bde27 100644 --- a/docs/cookbook/vla/OpenPI/Pi0.5.mdx +++ b/docs/cookbook/vla/OpenPI/Pi0.5.mdx @@ -263,7 +263,18 @@ 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 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. +- CUDA graph capture targets single-request prefix encoding and one action-denoise step. The lossless default keeps exact prompt lengths and one resident prefix graph. Prefix and action graph residency are independently bounded by `prefix_cuda_graph_max_entries` (default `1`) and `action_cuda_graph_max_entries` (default `4`). Without prompt buckets, unseen signatures run eagerly after the cache reaches capacity, avoiding capture churn. The denoise signature includes whether prefix attention is full or masked, plus batch size, prefix length, action horizon, action dim, dtype, and parallel layout. Mask-aware action graphs rebuild position IDs from the current request mask inside the captured graph. With action SP enabled, the denoise bucket uses the local action shard length and rank-specific position offset. +- `prompt_token_buckets` optionally right-pads prompt tokens to fixed sizes such as `[32, 64, 128, 200]`, allowing nearby prompt lengths to reuse both graphs. In this mode the bounded graph caches use LRU replacement and reset evicted CUDA graphs; prompts beyond the largest configured bucket stay exact-length and eager. This is opt-in because padding changes GPU reduction shapes: an H200 sweep over lengths around all four boundaries was structurally correct and bounded to four prefix/action graphs, but non-boundary prompts differed from the exact-length path by up to `0.09589` in normalized action space over five denoise steps. Validate closed-loop policy quality before enabling it; empty buckets preserve the numerically lossless path. + +An opt-in four-bucket serving config is: + +```json File +{ + "prompt_token_buckets": [32, 64, 128, 200], + "prefix_cuda_graph_max_entries": 4, + "action_cuda_graph_max_entries": 4 +} +``` - 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. @@ -299,7 +310,8 @@ Use this single-GPU config first for 16GB-class robot workstations. It keeps par "prefix_cache_max_entries": 0, "enable_prefix_cuda_graph": true, "prefix_cuda_graph_max_entries": 1, - "enable_action_cuda_graph": true + "enable_action_cuda_graph": true, + "action_cuda_graph_max_entries": 1 } ``` @@ -473,7 +485,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. 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. | +| Low-VRAM switches | Disabling prefix cache prevents cache growth across changing robot frames. Prefix graph residency defaults to `1`, action graph residency defaults to `4`, and setting either limit to `0` disables that runner. One validated ALOHA prefix graph used about `48.5 MiB`; action graph residency is now bounded separately. Keep `prompt_token_buckets` empty for the lossless path. | | 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. | diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/pi05.py b/python/sglang/multimodal_gen/configs/pipeline_configs/pi05.py index 6a6624b36..4d19069f4 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/pi05.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/pi05.py @@ -47,8 +47,12 @@ class Pi05PipelineConfig(PipelineConfig): enable_global_prefix_cache: bool = False enable_prefix_cuda_graph: bool = True + # Opt-in prompt buckets shared by prefix and action CUDA graphs. Padding + # changes reduction shapes, so exact prompt lengths remain the default. + prompt_token_buckets: list[int] = field(default_factory=list) prefix_cuda_graph_max_entries: int = 1 enable_action_cuda_graph: bool = True + action_cuda_graph_max_entries: int = 4 prefix_cache_max_entries: int = 1 prefix_cache_layout_version: str = "pi05-prefix-v1" offload_prefix_image_encoder: bool = False @@ -86,6 +90,64 @@ class Pi05PipelineConfig(PipelineConfig): } ) + def __post_init__(self) -> None: + self._validate_cuda_graph_config() + + def _validate_cuda_graph_config(self) -> None: + try: + buckets = list(self.prompt_token_buckets) + except TypeError as exc: + raise ValueError("prompt_token_buckets must contain integers") from exc + if not all( + isinstance(bucket, int) and not isinstance(bucket, bool) + for bucket in buckets + ): + raise ValueError("prompt_token_buckets must contain integers") + if any(bucket <= 0 for bucket in buckets): + raise ValueError("prompt_token_buckets must contain positive lengths") + if sorted(set(buckets)) != buckets: + raise ValueError( + "prompt_token_buckets must be strictly increasing and unique" + ) + if buckets and buckets[-1] > self.max_token_len: + raise ValueError( + "prompt_token_buckets cannot exceed max_token_len " + f"({self.max_token_len}), got {buckets[-1]}" + ) + if self.prefix_cuda_graph_max_entries < 0: + raise ValueError("prefix_cuda_graph_max_entries must be non-negative") + if self.action_cuda_graph_max_entries < 0: + raise ValueError("action_cuda_graph_max_entries must be non-negative") + self.prompt_token_buckets = buckets + + def check_pipeline_config(self) -> None: + super().check_pipeline_config() + self._validate_cuda_graph_config() + + def prefix_cuda_graph_available(self) -> bool: + return bool( + self.enable_prefix_cuda_graph + and self.prefix_cuda_graph_max_entries > 0 + and not any( + ( + self.offload_prefix_image_encoder, + self.offload_prefix_image_encoder_after_embed, + self.offload_prefix_token_embedding, + self.offload_prefix_language_layers, + self.offload_prefix_language_layers_after_prefix, + self.offload_prefix_language_layer_count_after_prefix > 0, + self.empty_cache_after_prefix, + ) + ) + ) + + def action_cuda_graph_available(self) -> bool: + return bool( + self.enable_action_cuda_graph + and self.action_cuda_graph_max_entries > 0 + and not self.offload_action_expert_after_denoise + ) + def supports_dynamic_batching(self): return True diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/action/protocol.py b/python/sglang/multimodal_gen/runtime/entrypoints/action/protocol.py index da8ae3cdd..50ca7603e 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/action/protocol.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/action/protocol.py @@ -164,6 +164,8 @@ def action_metadata(server_args: ServerArgs) -> dict[str, Any]: "policy_family", type(pipeline_config).__name__.removesuffix("PipelineConfig").lower(), ) + prefix_graph_enabled = pipeline_config.prefix_cuda_graph_available() + action_graph_enabled = pipeline_config.action_cuda_graph_available() return { "object": "action.metadata", "model": server_args.served_model_name, @@ -184,6 +186,13 @@ def action_metadata(server_args: ServerArgs) -> dict[str, Any]: "runtime": { "materialize_dtype": pipeline_config.materialize_dtype, "enable_autocast": pipeline_config.enable_autocast, + "cuda_graph": { + "prefix_enabled": prefix_graph_enabled, + "prefix_max_entries": pipeline_config.prefix_cuda_graph_max_entries, + "action_enabled": action_graph_enabled, + "action_max_entries": pipeline_config.action_cuda_graph_max_entries, + "prompt_token_buckets": list(pipeline_config.prompt_token_buckets), + }, "parallelism": { "num_gpus": server_args.num_gpus, "tp_size": server_args.tp_size, @@ -201,11 +210,13 @@ def action_metadata(server_args: ServerArgs) -> dict[str, Any]: "prefix_cache": ( "auto" if pipeline_config.enable_global_prefix_cache else False ), - "cuda_graph": "auto" if pipeline_config.enable_action_cuda_graph else False, + "cuda_graph": ( + "auto" if prefix_graph_enabled or action_graph_enabled else False + ), }, "capabilities": { "exact_prefix_cache": True, - "cuda_graph": pipeline_config.enable_action_cuda_graph, + "cuda_graph": prefix_graph_enabled or action_graph_enabled, "realtime_websocket": True, "openpi_websocket": True, "batch_inputs": False, diff --git a/python/sglang/multimodal_gen/runtime/layers/linear.py b/python/sglang/multimodal_gen/runtime/layers/linear.py index bccc9229f..2ed258277 100644 --- a/python/sglang/multimodal_gen/runtime/layers/linear.py +++ b/python/sglang/multimodal_gen/runtime/layers/linear.py @@ -310,6 +310,73 @@ class ReplicatedLinear(LinearBase): return s +class MergedReplicatedLinear(ReplicatedLinear): + """Packed replicated linear layers with shard-aware weight loading. + + This is the non-tensor-parallel counterpart of + :class:`MergedColumnParallelLinear`. It keeps independently stored logical + projections in one physical weight so eager inference launches one GEMM. + """ + + def __init__( + self, + input_size: int, + output_sizes: list[int], + bias: bool = True, + skip_bias_add: bool = False, + params_dtype: torch.dtype | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + self.output_sizes = output_sizes + super().__init__( + input_size=input_size, + output_size=sum(output_sizes), + bias=bias, + skip_bias_add=skip_bias_add, + params_dtype=params_dtype, + quant_config=quant_config, + output_sizes=output_sizes, + prefix=prefix, + ) + + def weight_loader( + self, + param: Parameter, + loaded_weight: torch.Tensor, + loaded_shard_id: int | str | None = None, + ) -> None: + if loaded_shard_id is None: + return super().weight_loader(param, loaded_weight) + + if isinstance(loaded_shard_id, str): + try: + loaded_shard_id = {"q": 0, "k": 1, "v": 2}[loaded_shard_id] + except KeyError as exc: + raise ValueError(f"Invalid merged shard id: {loaded_shard_id}") from exc + if not 0 <= loaded_shard_id < len(self.output_sizes): + raise ValueError(f"Invalid merged shard id: {loaded_shard_id}") + param_data = param.data + output_dim = getattr(param, "output_dim", None) + if output_dim is not None: + shard_offset = sum(self.output_sizes[:loaded_shard_id]) + shard_size = self.output_sizes[loaded_shard_id] + param_data = param_data.narrow(output_dim, shard_offset, shard_size) + elif getattr(param, "is_metadata", False): + shard_size = loaded_weight.shape[0] + param_data = param_data.narrow(0, loaded_shard_id * shard_size, shard_size) + elif getattr(param, "needs_scalar_to_array", False): + param_data, loaded_weight = adjust_scalar_to_fused_array( + param_data, loaded_weight, loaded_shard_id + ) + if tuple(param_data.shape) != tuple(loaded_weight.shape): + raise ValueError( + f"Tried to load merged shard of size {loaded_weight.size()} " + f"to a parameter slice of size {param_data.size()}" + ) + param_data.copy_(loaded_weight) + + class ColumnParallelLinear(LinearBase): """Linear layer with column parallelism. diff --git a/python/sglang/multimodal_gen/runtime/models/vlas/pi05_core.py b/python/sglang/multimodal_gen/runtime/models/vlas/pi05_core.py index 01b57c309..791b3125f 100644 --- a/python/sglang/multimodal_gen/runtime/models/vlas/pi05_core.py +++ b/python/sglang/multimodal_gen/runtime/models/vlas/pi05_core.py @@ -26,6 +26,7 @@ from sglang.multimodal_gen.runtime.layers.activation import get_act_fn from sglang.multimodal_gen.runtime.layers.attention import LocalAttention, USPAttention from sglang.multimodal_gen.runtime.layers.linear import ( MergedColumnParallelLinear, + MergedReplicatedLinear, QKVParallelLinear, RowParallelLinear, ) @@ -169,31 +170,26 @@ class PiGemmaMLP(nn.Module): self.gate_proj = None self.up_proj = None else: - self.gate_proj = nn.Linear( - self.hidden_size, self.intermediate_size, bias=False - ) - self.up_proj = nn.Linear( - self.hidden_size, self.intermediate_size, bias=False + self.gate_up_proj = MergedReplicatedLinear( + input_size=self.hidden_size, + output_sizes=[self.intermediate_size] * 2, + bias=False, ) self.down_proj = nn.Linear( self.intermediate_size, self.hidden_size, bias=False ) - self.gate_up_proj = None + self.gate_proj = None + self.up_proj = None if config.hidden_act != "gelu_pytorch_tanh": raise ValueError(f"Unsupported PiGemma activation: {config.hidden_act}") self.act_fn = GeluAndMul(approximate="tanh") @property def projection_dtype(self) -> torch.dtype: - if self.tensor_parallel: - return self.gate_up_proj.weight.dtype - return self.up_proj.weight.dtype + return self.gate_up_proj.weight.dtype def forward(self, x: torch.Tensor) -> torch.Tensor: - if self.tensor_parallel: - gate_up = linear_forward(self.gate_up_proj, x) - else: - gate_up = torch.cat([self.gate_proj(x), self.up_proj(x)], dim=-1) + gate_up = linear_forward(self.gate_up_proj, x) return linear_forward(self.down_proj, self.act_fn(gate_up)) @@ -291,19 +287,11 @@ class PiGemmaAttention(nn.Module): else: self.num_heads = config.num_attention_heads self.num_key_value_heads = config.num_key_value_heads - self.q_proj = nn.Linear( - config.hidden_size, - self.num_heads * self.head_dim, - bias=config.attention_bias, - ) - self.k_proj = nn.Linear( - config.hidden_size, - self.num_key_value_heads * self.head_dim, - bias=config.attention_bias, - ) - self.v_proj = nn.Linear( - config.hidden_size, - self.num_key_value_heads * self.head_dim, + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_key_value_heads * self.head_dim + self.qkv_proj = MergedReplicatedLinear( + input_size=config.hidden_size, + output_sizes=[self.q_size, self.kv_size, self.kv_size], bias=config.attention_bias, ) self.o_proj = nn.Linear( @@ -311,9 +299,9 @@ class PiGemmaAttention(nn.Module): config.hidden_size, bias=config.attention_bias, ) - self.qkv_proj = None - self.q_size = self.num_heads * self.head_dim - self.kv_size = self.num_key_value_heads * self.head_dim + self.q_proj = None + self.k_proj = None + self.v_proj = None self.num_key_value_groups = self.num_heads // self.num_key_value_heads self.attn = LocalAttention( num_heads=self.num_heads, @@ -350,9 +338,7 @@ class PiGemmaAttention(nn.Module): @property def projection_dtype(self) -> torch.dtype: - if self.tensor_parallel: - return self.qkv_proj.weight.dtype - return self.q_proj.weight.dtype + return self.qkv_proj.weight.dtype def project_qkv( self, @@ -362,16 +348,11 @@ class PiGemmaAttention(nn.Module): query_shape = (*input_shape, self.num_heads, self.head_dim) kv_shape = (*input_shape, self.num_key_value_heads, self.head_dim) - if self.tensor_parallel: - qkv = linear_forward(self.qkv_proj, hidden_states) - query_states, key_states, value_states = qkv.split( - [self.q_size, self.kv_size, self.kv_size], - dim=-1, - ) - else: - query_states = self.q_proj(hidden_states) - key_states = self.k_proj(hidden_states) - value_states = self.v_proj(hidden_states) + qkv = linear_forward(self.qkv_proj, hidden_states) + query_states, key_states, value_states = qkv.split( + [self.q_size, self.kv_size, self.kv_size], + dim=-1, + ) return ( query_states.view(query_shape).transpose(1, 2), key_states.view(kv_shape).transpose(1, 2), @@ -828,20 +809,22 @@ def create_sinusoidal_pos_embedding( dimension: int, min_period: float, max_period: float, + scaling: torch.Tensor | None = None, ) -> Tensor: if dimension % 2 != 0: raise ValueError(f"dimension ({dimension}) must be divisible by 2") if time.ndim != 1: raise ValueError("time must have shape [batch]") - fraction = torch.linspace( - 0.0, - 1.0, - dimension // 2, - dtype=torch.float64, - device=time.device, - ) - period = min_period * (max_period / min_period) ** fraction - scaling = 1.0 / period * 2 * math.pi + if scaling is None: + fraction = torch.linspace( + 0.0, + 1.0, + dimension // 2, + dtype=torch.float64, + device=time.device, + ) + period = min_period * (max_period / min_period) ** fraction + scaling = 1.0 / period * 2 * math.pi sin_input = scaling[None, :] * time[:, None].to(torch.float64) return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1) @@ -1223,6 +1206,11 @@ class Pi05CoreModel(nn.Module): self.action_out_proj = None self.time_mlp_in = None self.time_mlp_out = None + self.register_buffer( + "_time_embedding_scaling", + torch.empty(0, dtype=torch.float64), + persistent=False, + ) def retain_runtime_components( self, @@ -1280,12 +1268,34 @@ class Pi05CoreModel(nn.Module): self, noisy_actions: torch.Tensor, timestep: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: + if ( + self._time_embedding_scaling.numel() + != self.action_in_proj.out_features // 2 + or self._time_embedding_scaling.device != timestep.device + ): + fraction = torch.linspace( + 0.0, + 1.0, + self.action_in_proj.out_features // 2, + dtype=torch.float64, + device=timestep.device, + ) + period = ( + self.config.time_embedding_min_period + * ( + self.config.time_embedding_max_period + / self.config.time_embedding_min_period + ) + ** fraction + ) + self._time_embedding_scaling = 1.0 / period * 2 * math.pi time_emb = create_sinusoidal_pos_embedding( timestep, self.action_in_proj.out_features, min_period=self.config.time_embedding_min_period, max_period=self.config.time_embedding_max_period, + scaling=self._time_embedding_scaling, ) action_emb = self.action_in_proj( noisy_actions.to(dtype=self.action_in_proj.weight.dtype) @@ -1296,21 +1306,54 @@ class Pi05CoreModel(nn.Module): time_emb = self.time_mlp_out(time_emb) adarms_cond = F.silu(time_emb) - batch_size, action_len = action_emb.shape[:2] + return action_emb, adarms_cond + + def prepare_denoise_layout( + self, + prefix_pad_masks: torch.Tensor, + x_t: torch.Tensor, + prefix_full_attention: bool = False, + *, + action_position_offset: int = 0, + ) -> tuple[torch.Tensor | None, torch.Tensor]: + batch_size, action_len = x_t.shape[:2] pad_masks = torch.ones( batch_size, action_len, dtype=torch.bool, - device=noisy_actions.device, + device=x_t.device, ) att_masks_t = torch.zeros( batch_size, action_len, - dtype=action_emb.dtype, - device=noisy_actions.device, + dtype=x_t.dtype, + device=x_t.device, ) att_masks_t[:, 0] = 1 - return action_emb, pad_masks, att_masks_t, adarms_cond + if prefix_full_attention: + attention_mask = None + else: + prefix_len = prefix_pad_masks.shape[1] + prefix_pad_2d_masks = prefix_pad_masks[:, None, :].expand( + batch_size, action_len, prefix_len + ) + suffix_att_2d_masks = make_att_2d_masks(pad_masks, att_masks_t) + full_att_2d_masks = torch.cat( + [prefix_pad_2d_masks, suffix_att_2d_masks], + dim=2, + ) + # A masked prefix guarantees that the concatenated layout is not + # full attention. Avoid a device-to-host ``.item()`` here so this + # path remains CUDA-graph capturable. + attention_mask = self.prepare_attention_masks_4d( + full_att_2d_masks, + full_attention=False, + ) + prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None] + position_ids = ( + prefix_offsets + action_position_offset + torch.cumsum(pad_masks, dim=1) - 1 + ) + return attention_mask, position_ids def _move_prefix_image_encoder_to_device(self, device: torch.device) -> None: paligemma = self.paligemma_with_expert.paligemma @@ -1384,12 +1427,16 @@ class Pi05CoreModel(nn.Module): ) self._offload_prefix_image_encoder_after_embed() prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1 - prefix_full_attention = bool(prefix_full_attention_hint) - if prefix_full_attention: + if prefix_full_attention_hint is True: + prefix_full_attention = True attention_mask = None else: prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks) - prefix_full_attention = bool(prefix_att_2d_masks.all().item()) + prefix_full_attention = ( + bool(prefix_att_2d_masks.all().item()) + if prefix_full_attention_hint is None + else False + ) attention_mask = self.prepare_attention_masks_4d( prefix_att_2d_masks, full_attention=prefix_full_attention, @@ -1420,32 +1467,17 @@ class Pi05CoreModel(nn.Module): prefix_full_attention: bool = False, *, action_position_offset: int = 0, + denoise_layout: tuple[torch.Tensor | None, torch.Tensor] | None = None, ) -> torch.Tensor: - suffix_embs, suffix_pad_masks, suffix_att_masks, adarms_cond = ( - self.embed_suffix(x_t, timestep) - ) - suffix_len = suffix_pad_masks.shape[1] - batch_size = prefix_pad_masks.shape[0] - prefix_len = prefix_pad_masks.shape[1] - if prefix_full_attention: - attention_mask = None - else: - prefix_pad_2d_masks = prefix_pad_masks[:, None, :].expand( - batch_size, suffix_len, prefix_len + suffix_embs, adarms_cond = self.embed_suffix(x_t, timestep) + if denoise_layout is None: + denoise_layout = self.prepare_denoise_layout( + prefix_pad_masks, + x_t, + prefix_full_attention, + action_position_offset=action_position_offset, ) - suffix_att_2d_masks = make_att_2d_masks(suffix_pad_masks, suffix_att_masks) - full_att_2d_masks = torch.cat( - [prefix_pad_2d_masks, suffix_att_2d_masks], - dim=2, - ) - attention_mask = self.prepare_attention_masks_4d(full_att_2d_masks) - prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None] - position_ids = ( - prefix_offsets - + action_position_offset - + torch.cumsum(suffix_pad_masks, dim=1) - - 1 - ) + attention_mask, position_ids = denoise_layout with set_forward_context(current_timestep=0, attn_metadata=None): outputs_embeds, _ = self.paligemma_with_expert.forward( attention_mask=attention_mask, diff --git a/python/sglang/multimodal_gen/runtime/models/vlas/pi05_policy.py b/python/sglang/multimodal_gen/runtime/models/vlas/pi05_policy.py index 7a7d18ee5..94faba358 100644 --- a/python/sglang/multimodal_gen/runtime/models/vlas/pi05_policy.py +++ b/python/sglang/multimodal_gen/runtime/models/vlas/pi05_policy.py @@ -57,6 +57,11 @@ from sglang.multimodal_gen.runtime.vla.prefix_cache import ( VLADensePrefixCache, VLAPrefixCacheManager, ) +from sglang.multimodal_gen.runtime.vla.prompt_bucketing import ( + bucket_prompt_tokens, + effective_token_length, + select_prompt_token_bucket, +) from sglang.multimodal_gen.utils import set_mixed_precision_policy logger = init_logger(__name__) @@ -83,6 +88,7 @@ class Pi05ActionExpert(nn.Module): timestep: torch.Tensor, *, action_position_offset: int = 0, + denoise_layout: tuple[torch.Tensor | None, torch.Tensor] | None = None, ) -> torch.Tensor: return self.core_model.denoise_step( prefix_context.prefix_pad_masks, @@ -91,6 +97,7 @@ class Pi05ActionExpert(nn.Module): timestep, bool(prefix_context.layout.get("full_attention", False)), action_position_offset=action_position_offset, + denoise_layout=denoise_layout, ) @@ -163,12 +170,16 @@ 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) + evict_graphs_on_miss = self._prompt_token_bucketing_enabled() self.prefix_graph_runner = VLAPrefixGraphRunner( enabled=self._prefix_cuda_graph_enabled(), max_entries=config.prefix_cuda_graph_max_entries, + evict_on_miss=evict_graphs_on_miss, ) self.graph_runner = VLADenoiseGraphRunner( - enabled=config.enable_action_cuda_graph + enabled=config.action_cuda_graph_available(), + max_entries=config.action_cuda_graph_max_entries, + evict_on_miss=evict_graphs_on_miss, ) def _should_use_prefix_tensor_parallel(self) -> bool: @@ -184,22 +195,30 @@ class Pi05PolicyModel(nn.Module): def _prefix_cuda_graph_enabled(self) -> bool: if ( - not self.config.enable_prefix_cuda_graph + not self.config.prefix_cuda_graph_available() 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, - ) + return True + + def _prompt_token_bucketing_enabled(self) -> bool: + graph_path_available = ( + self._prefix_cuda_graph_enabled() + or self.config.action_cuda_graph_available() ) + if ( + not self.config.prompt_token_buckets + or not graph_path_available + or self.device.type != "cuda" + or self.runtime_role not in ("all", "prefix") + or self._prefix_tensor_parallel_enabled() + ): + return False + # Action SP currently requires a full-attention prefix, while prompt + # buckets introduce padding masks. + return get_vla_split_group() is None @staticmethod def _to_empty_preserve_buffers(module: nn.Module, *, device: torch.device) -> None: @@ -785,6 +804,8 @@ class Pi05PolicyModel(nn.Module): def build_prefix_cache_key( self, observation: VLAObservationBatch, + *, + bucket_prompt: bool = False, ) -> str: camera_order = tuple(observation.metadata.get("camera_order", ())) image_hashes = { @@ -793,7 +814,7 @@ class Pi05PolicyModel(nn.Module): masks = { name: bool(mask.item()) for name, mask in observation.image_masks.items() } - token_len = int(observation.token_masks.sum(dim=1).max().item()) + token_len = effective_token_length(observation.token_masks) tokens = ( observation.tokens[:, :token_len] if token_len > 0 else observation.tokens ) @@ -803,6 +824,18 @@ class Pi05PolicyModel(nn.Module): else observation.token_masks ) model_revision = os.path.basename(os.path.normpath(self.model_path)) + bucketing_enabled = bucket_prompt and self._prompt_token_bucketing_enabled() + prompt_bucket = ( + select_prompt_token_bucket(token_len, self.config.prompt_token_buckets) + if bucketing_enabled + else None + ) + if prompt_bucket is not None: + prompt_layout = f"bucket-{prompt_bucket}" + elif bucketing_enabled: + prompt_layout = "bucket-miss-exact" + else: + prompt_layout = "exact" return VLAPrefixCacheManager.make_key( model_revision=model_revision, tokenizer_id=f"{self.config.paligemma_variant}:{self.config.max_token_len}", @@ -811,7 +844,9 @@ class Pi05PolicyModel(nn.Module): token_digest=tensor_fingerprint(tokens), token_mask_digest=tensor_fingerprint(token_masks), masks=masks, - positions_version=self.config.prefix_cache_layout_version, + positions_version=( + f"{self.config.prefix_cache_layout_version}:{prompt_layout}" + ), dtype=str(self.dtype).replace("torch.", ""), parallel_layout_version=self.config.parallel_layout_version, cache_namespace="pi05", @@ -858,6 +893,7 @@ class Pi05PolicyModel(nn.Module): observation: VLAObservationBatch, *, use_cuda_graph: bool = True, + bucket_prompt: bool | None = None, ) -> PrefixContext: camera_order = tuple(observation.metadata.get("camera_order", ())) images = [ @@ -867,19 +903,42 @@ class Pi05PolicyModel(nn.Module): image_masks = [ observation.image_masks[name].to(self.device) for name in camera_order ] - token_len = int(observation.token_masks.sum(dim=1).max().item()) - tokens_trimmed = token_len > 0 - if tokens_trimmed and token_len < observation.tokens.shape[1]: - tokens_cpu = observation.tokens[:, :token_len] - token_masks_cpu = observation.token_masks[:, :token_len] + if bucket_prompt is None: + bucket_prompt = use_cuda_graph + use_prompt_bucket = bucket_prompt and self._prompt_token_bucketing_enabled() + if use_prompt_bucket: + tokens_cpu, token_masks_cpu, token_len, prompt_bucket = ( + bucket_prompt_tokens( + observation.tokens, + observation.token_masks, + self.config.prompt_token_buckets, + ) + ) else: - tokens_cpu = observation.tokens - token_masks_cpu = observation.token_masks + token_len = effective_token_length(observation.token_masks) + prompt_bucket = None + if 0 < token_len < observation.tokens.shape[1]: + tokens_cpu = observation.tokens[:, :token_len] + token_masks_cpu = observation.token_masks[:, :token_len] + else: + tokens_cpu = observation.tokens + token_masks_cpu = observation.token_masks + prompt_bucket_miss = use_prompt_bucket and prompt_bucket is None + preserve_token_shape = token_len > 0 or prompt_bucket is not None tokens = tokens_cpu.to(self.device) token_masks = token_masks_cpu.to(self.device) - prefix_full_attention_hint = all( - bool(observation.image_masks[name].all().item()) for name in camera_order - ) and bool(token_masks_cpu.all().item()) + # Use one masked control flow for every logical length in a bucket. + prefix_full_attention_hint = ( + False + if prompt_bucket is not None + else ( + all( + bool(observation.image_masks[name].all().item()) + for name in camera_order + ) + and bool(token_masks_cpu.all().item()) + ) + ) image_count = len(images) graph_inputs = tuple([*images, *image_masks, tokens, token_masks]) @@ -896,7 +955,7 @@ class Pi05PolicyModel(nn.Module): current_tokens, current_token_masks, prefix_full_attention_hint=prefix_full_attention_hint, - tokens_trimmed=tokens_trimmed, + tokens_trimmed=preserve_token_shape, ) ) past_key_values = self._materialize_prefix_kv_for_action(past_key_values) @@ -904,13 +963,18 @@ class Pi05PolicyModel(nn.Module): past_key_values=past_key_values, prefix_pad_masks=prefix_pad_masks, prefix_len=prefix_pad_masks.shape[1], - layout={"full_attention": full_attention}, + layout={ + "full_attention": full_attention, + "prompt_token_bucket": prompt_bucket, + "cuda_graph_eligible": not prompt_bucket_miss, + }, ) if ( not use_cuda_graph or not self.prefix_graph_runner.enabled or observation.batch_size != 1 + or prompt_bucket_miss ): return encode(graph_inputs) @@ -923,7 +987,8 @@ class Pi05PolicyModel(nn.Module): static_layout=( image_count, prefix_full_attention_hint, - tokens_trimmed, + preserve_token_shape, + prompt_bucket, ), parallel_layout=self.config.parallel_layout_version, ) @@ -957,15 +1022,17 @@ class Pi05PolicyModel(nn.Module): use_cuda_graph: bool = True, action_position_offset: int = 0, action_sp_enabled: bool = False, + denoise_layout: tuple[torch.Tensor | None, torch.Tensor] | None = None, ) -> torch.Tensor: - if not bool(prefix_context.layout.get("full_attention", False)): - use_cuda_graph = False - if not use_cuda_graph: + if not use_cuda_graph or not prefix_context.layout.get( + "cuda_graph_eligible", True + ): return self.action_expert( prefix_context, x_t, timestep, action_position_offset=action_position_offset, + denoise_layout=denoise_layout, ) parallel_layout = self.config.parallel_layout_version if action_sp_enabled: @@ -976,6 +1043,9 @@ class Pi05PolicyModel(nn.Module): signature = VLADenoiseGraphSignature( batch_size=x_t.shape[0], prefix_len=prefix_context.prefix_len, + prefix_full_attention=bool( + prefix_context.layout.get("full_attention", False) + ), action_horizon=x_t.shape[1], action_dim=x_t.shape[2], dtype=str(x_t.dtype).replace("torch.", ""), @@ -992,6 +1062,7 @@ class Pi05PolicyModel(nn.Module): current_x_t, current_timestep, action_position_offset=action_position_offset, + denoise_layout=denoise_layout, ) return self.graph_runner.capture_or_run( @@ -1135,6 +1206,21 @@ class Pi05PolicyModel(nn.Module): if action_sp_enabled: x_t, action_position_offset = self._shard_action_sequence(x_t) + full_attention = bool(prefix_context.layout.get("full_attention", False)) + graph_enabled = ( + use_cuda_graph + and self.graph_runner.enabled + and prefix_context.layout.get("cuda_graph_eligible", True) + ) + denoise_layout = None + if not graph_enabled: + denoise_layout = self.core_model.prepare_denoise_layout( + prefix_context.prefix_pad_masks, + x_t, + full_attention, + action_position_offset=action_position_offset, + ) + dt = -1.0 / num_steps timesteps = torch.linspace( 1.0, @@ -1152,6 +1238,7 @@ class Pi05PolicyModel(nn.Module): use_cuda_graph=use_cuda_graph, action_position_offset=action_position_offset, action_sp_enabled=action_sp_enabled, + denoise_layout=denoise_layout, ) x_t.add_(velocity, alpha=dt) if action_sp_enabled: diff --git a/python/sglang/multimodal_gen/runtime/pipelines/pi05.py b/python/sglang/multimodal_gen/runtime/pipelines/pi05.py index 64f9cb077..ef66d6903 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/pi05.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/pi05.py @@ -59,7 +59,8 @@ class Pi05Pipeline(ComposedPipelineBase): or bool(server_args.text_encoder_cpu_offload) ) logger.info( - "Pi05 memory config: prefix_cache=%s/%s, cuda_graph=%s/%s/%s, " + "Pi05 memory config: prefix_cache=%s/%s, " + "cuda_graph=prefix:%s/%s action:%s/%s buckets:%s, " "offload_image=%s, offload_image_after_embed=%s, " "offload_tokens=%s, offload_language_layers=%s, " "offload_language_after_prefix=%s/%s, " @@ -69,6 +70,8 @@ class Pi05Pipeline(ComposedPipelineBase): pipeline_config.enable_prefix_cuda_graph, pipeline_config.prefix_cuda_graph_max_entries, pipeline_config.enable_action_cuda_graph, + pipeline_config.action_cuda_graph_max_entries, + pipeline_config.prompt_token_buckets, pipeline_config.offload_prefix_image_encoder, pipeline_config.offload_prefix_image_encoder_after_embed, pipeline_config.offload_prefix_token_embedding, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/vla.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/vla.py index 6c6aca470..c17a5b5e7 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/vla.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/vla.py @@ -260,7 +260,10 @@ class VLAPrefixEncodingStage(PipelineStage): """try querying the cache for PrefixContext with prefix cache key built from observations and other keys""" cache_enabled = _effective_prefix_cache_enabled(batch, server_args) if cache_enabled: - cache_key = self.policy_model.build_prefix_cache_key(observation) + cache_key = self.policy_model.build_prefix_cache_key( + observation, + bucket_prompt=_cuda_graph_enabled(batch), + ) cached_context = self.prefix_cache.get(cache_key) else: cache_key = None @@ -299,6 +302,7 @@ class VLAPrefixEncodingStage(PipelineStage): "scope": "global", "mode": "exact", "prefix_len": cached_context.prefix_len, + "prompt_token_bucket": cached_context.layout.get("prompt_token_bucket"), } if split is not None: self._send_prefix_result(batch, split, cached_context) @@ -307,9 +311,11 @@ class VLAPrefixEncodingStage(PipelineStage): prefix_start = time.perf_counter() # 3. run encoding + cuda_graph_enabled = _cuda_graph_enabled(batch) prefix_context = self.policy_model.encode_prefix( observation, - use_cuda_graph=_cuda_graph_enabled(batch) and not cache_enabled, + use_cuda_graph=cuda_graph_enabled and not cache_enabled, + bucket_prompt=cuda_graph_enabled, ) if cache_key is not None: prefix_context.cache_key_digest = cache_key @@ -321,6 +327,7 @@ class VLAPrefixEncodingStage(PipelineStage): "scope": "global" if cache_enabled else "request", "mode": "exact" if cache_enabled else "disabled", "prefix_len": prefix_context.prefix_len, + "prompt_token_bucket": prefix_context.layout.get("prompt_token_bucket"), } # 4. update prefix kv cache @@ -421,7 +428,6 @@ class VLAActionDenoisingStage(PipelineStage): ) elif should_run_action: # broadcast PrefixContext from action root rank to action ranks - options = vla_options(batch) noise = observation.noise if observation is not None else None actions = self.policy_model.sample_actions( observation, diff --git a/python/sglang/multimodal_gen/runtime/vla/cuda_graph.py b/python/sglang/multimodal_gen/runtime/vla/cuda_graph.py index 3550b16dc..5f56265f1 100644 --- a/python/sglang/multimodal_gen/runtime/vla/cuda_graph.py +++ b/python/sglang/multimodal_gen/runtime/vla/cuda_graph.py @@ -2,6 +2,7 @@ from __future__ import annotations +from collections import OrderedDict from dataclasses import dataclass from typing import Any, Callable @@ -42,6 +43,7 @@ class _CapturedPrefixGraph: class VLADenoiseGraphSignature: batch_size: int prefix_len: int + prefix_full_attention: bool action_horizon: int action_dim: int dtype: str @@ -59,6 +61,121 @@ class _CapturedDenoiseGraph: current_context_digest: str | None = None +@dataclass(frozen=True) +class VLAGraphCacheInfo: + size: int + max_entries: int + hits: int + misses: int + captures: int + evictions: int + failures: int + evict_on_miss: bool + + +class _BoundedCaptureCache: + """Bounded LRU owning CUDA graphs and their static buffers.""" + + def __init__(self, name: str, max_entries: int, *, evict_on_miss: bool): + self.name = name + self.max_entries = max(0, int(max_entries)) + self.evict_on_miss = evict_on_miss + self.entries: OrderedDict[Any, Any] = OrderedDict() + self.hits = 0 + self.misses = 0 + self.captures = 0 + self.evictions = 0 + self.failures = 0 + + @staticmethod + def _release(entry: Any) -> None: + reset = getattr(entry.graph, "reset", None) + if callable(reset): + reset() + + def get(self, signature: Any) -> Any | None: + entry = self.entries.get(signature) + if entry is None: + self.misses += 1 + return None + self.hits += 1 + self.entries.move_to_end(signature) + return entry + + def can_admit(self, signature: Any) -> bool: + return self.max_entries > 0 and ( + signature in self.entries + or len(self.entries) < self.max_entries + or self.evict_on_miss + ) + + def prepare_admission(self, signature: Any) -> None: + if ( + signature in self.entries + or len(self.entries) < self.max_entries + or not self.evict_on_miss + ): + return + evicted_signature, evicted = self.entries.popitem(last=False) + self._release(evicted) + self.evictions += 1 + logger.info( + "Evicted VLA %s CUDA graph for signature %s (entries=%d/%d)", + self.name, + evicted_signature, + len(self.entries), + self.max_entries, + ) + + def put(self, signature: Any, entry: Any) -> bool: + if self.max_entries == 0 or not self.can_admit(signature): + self._release(entry) + return False + previous = self.entries.pop(signature, None) + if previous is not None: + self._release(previous) + self.entries[signature] = entry + self.captures += 1 + + if len(self.entries) > self.max_entries: + evicted_signature, evicted = self.entries.popitem(last=False) + self._release(evicted) + self.evictions += 1 + logger.info( + "Evicted VLA %s CUDA graph for signature %s (entries=%d/%d)", + self.name, + evicted_signature, + len(self.entries), + self.max_entries, + ) + return True + + def discard(self, signature: Any) -> None: + entry = self.entries.pop(signature, None) + if entry is not None: + self._release(entry) + + def clear(self) -> None: + for entry in self.entries.values(): + self._release(entry) + self.entries.clear() + + def mark_failure(self) -> None: + self.failures += 1 + + def info(self) -> VLAGraphCacheInfo: + return VLAGraphCacheInfo( + size=len(self.entries), + max_entries=self.max_entries, + hits=self.hits, + misses=self.misses, + captures=self.captures, + evictions=self.evictions, + failures=self.failures, + evict_on_miss=self.evict_on_miss, + ) + + def _clone_past_key_values(past_key_values: Any) -> Any: return VLADensePrefixCache( tuple( @@ -95,14 +212,31 @@ def _copy_prefix_context_(dst: PrefixContext, src: PrefixContext) -> None: 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) + def __init__( + self, + enabled: bool = True, + max_entries: int = 1, + *, + evict_on_miss: bool = False, + ): + self._cache = _BoundedCaptureCache( + "prefix", + max_entries, + evict_on_miss=evict_on_miss, + ) + self.max_entries = self._cache.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 cache_info(self) -> VLAGraphCacheInfo: + return self._cache.info() + + def clear(self) -> None: + self._cache.clear() + self._disabled_signatures.clear() + def _capture( self, signature: VLAPrefixGraphSignature, @@ -140,11 +274,11 @@ class VLAPrefixGraphRunner: static_inputs=static_inputs, static_output=static_output, ) - self._captured[signature] = captured logger.info( - "Captured VLA prefix CUDA graph: batch=%d inputs=%s", + "Captured VLA prefix CUDA graph: batch=%d inputs=%s (capacity=%d)", signature.batch_size, signature.input_shapes, + self.max_entries, ) return captured @@ -162,12 +296,14 @@ class VLAPrefixGraphRunner: ): return step_fn(inputs) - captured = self._captured.get(signature) - if captured is None and len(self._captured) >= self.max_entries: + captured = self._cache.get(signature) + if captured is None and not self._cache.can_admit(signature): return step_fn(inputs) try: if captured is None: + self._cache.prepare_admission(signature) captured = self._capture(signature, step_fn, inputs) + self._cache.put(signature, captured) else: for static_input, current_input in zip( captured.static_inputs, inputs, strict=True @@ -180,7 +316,8 @@ class VLAPrefixGraphRunner: return captured.static_output except Exception: self._disabled_signatures.add(signature) - self._captured.pop(signature, None) + self._cache.discard(signature) + self._cache.mark_failure() logger.warning( "VLA prefix CUDA graph disabled for signature %s", signature, @@ -196,12 +333,32 @@ class VLADenoiseGraphRunner: diffusion BCG and does not capture prefix encoding or token decode. """ - def __init__(self, enabled: bool = True): - self.enabled = enabled - self._captured: dict[VLADenoiseGraphSignature, _CapturedDenoiseGraph] = {} + def __init__( + self, + enabled: bool = True, + max_entries: int = 1, + *, + evict_on_miss: bool = False, + ): + self._cache = _BoundedCaptureCache( + "action-denoise", + max_entries, + evict_on_miss=evict_on_miss, + ) + self.max_entries = self._cache.max_entries + self.enabled = enabled and self.max_entries > 0 self._disabled_signatures: set[VLADenoiseGraphSignature] = set() self._capture_stream: torch.cuda.Stream | None = None self._graph_pool: Any = None + self._capacity_warning_emitted = False + + def cache_info(self) -> VLAGraphCacheInfo: + return self._cache.info() + + def clear(self) -> None: + self._cache.clear() + self._disabled_signatures.clear() + self._capacity_warning_emitted = False def _sync_context_if_needed( self, @@ -279,15 +436,15 @@ class VLADenoiseGraphRunner: current_context_id=id(prefix_context.past_key_values), current_context_digest=prefix_context.cache_key_digest, ) - self._captured[signature] = captured logger.info( "Captured VLA denoise CUDA graph: batch=%d prefix=%d action=%dx%d " - "dtype=%s", + "dtype=%s (capacity=%d)", signature.batch_size, signature.prefix_len, signature.action_horizon, signature.action_dim, signature.dtype, + self.max_entries, ) return captured @@ -305,12 +462,23 @@ class VLADenoiseGraphRunner: if x_t.device.type != "cuda": return step_fn(prefix_context, x_t, timestep) - captured = self._captured.get(signature) + captured = self._cache.get(signature) + if captured is None and not self._cache.can_admit(signature): + if not self._capacity_warning_emitted: + logger.info( + "VLA denoise CUDA graph capacity reached (%d); " + "new signatures run eagerly", + self.max_entries, + ) + self._capacity_warning_emitted = True + return step_fn(prefix_context, x_t, timestep) try: if captured is None: + self._cache.prepare_admission(signature) captured = self._capture( signature, step_fn, prefix_context, x_t, timestep ) + self._cache.put(signature, captured) captured.graph.replay() else: self._sync_context_if_needed(captured, prefix_context) @@ -320,7 +488,8 @@ class VLADenoiseGraphRunner: return captured.static_output except Exception: self._disabled_signatures.add(signature) - self._captured.pop(signature, None) + self._cache.discard(signature) + self._cache.mark_failure() logger.warning( "VLA denoise CUDA graph disabled for signature %s", signature, diff --git a/python/sglang/multimodal_gen/runtime/vla/prompt_bucketing.py b/python/sglang/multimodal_gen/runtime/vla/prompt_bucketing.py new file mode 100644 index 000000000..1d7916b1b --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/vla/prompt_bucketing.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections.abc import Sequence + +import torch +import torch.nn.functional as F + + +def effective_token_length(token_masks: torch.Tensor) -> int: + """Return the last visible token position across the batch.""" + + if token_masks.ndim != 2: + raise ValueError( + f"Pi0.5 token masks must be [batch, seq], got {token_masks.shape}" + ) + if token_masks.shape[1] == 0: + return 0 + + positions = torch.arange( + 1, + token_masks.shape[1] + 1, + device=token_masks.device, + dtype=torch.long, + ) + lengths = torch.where(token_masks.to(torch.bool), positions, 0).amax(dim=1) + return int(lengths.max().item()) + + +def select_prompt_token_bucket( + token_length: int, + buckets: Sequence[int], +) -> int | None: + """Select the smallest configured bucket containing ``token_length``.""" + + if token_length < 0: + raise ValueError("token_length must be non-negative") + return next((int(bucket) for bucket in buckets if token_length <= bucket), None) + + +def bucket_prompt_tokens( + tokens: torch.Tensor, + token_masks: torch.Tensor, + buckets: Sequence[int], + *, + pad_token_id: int = 0, +) -> tuple[torch.Tensor, torch.Tensor, int, int | None]: + """Trim or right-pad prompt tensors to a stable CUDA graph bucket.""" + + if tokens.ndim != 2: + raise ValueError(f"Pi0.5 tokens must be [batch, seq], got {tokens.shape}") + if token_masks.shape != tokens.shape: + raise ValueError( + "Pi0.5 tokens and token masks must have identical shapes, got " + f"{tokens.shape} and {token_masks.shape}" + ) + + logical_length = effective_token_length(token_masks) + bucket = select_prompt_token_bucket(logical_length, buckets) + target_length = bucket if bucket is not None else logical_length + + # Preserve the existing empty-prompt fallback when no bucket is selected. + if target_length == 0 and bucket is None: + target_length = tokens.shape[1] + + if tokens.shape[1] >= target_length: + return ( + tokens[:, :target_length], + token_masks[:, :target_length], + logical_length, + bucket, + ) + + padding = target_length - tokens.shape[1] + return ( + F.pad(tokens, (0, padding), value=pad_token_id), + F.pad(token_masks, (0, padding), value=False), + logical_length, + bucket, + ) diff --git a/python/sglang/multimodal_gen/test/unit/test_parallel_linear_weight_loading.py b/python/sglang/multimodal_gen/test/unit/test_parallel_linear_weight_loading.py index 8d9deb6a2..364b7eb1c 100644 --- a/python/sglang/multimodal_gen/test/unit/test_parallel_linear_weight_loading.py +++ b/python/sglang/multimodal_gen/test/unit/test_parallel_linear_weight_loading.py @@ -3,6 +3,7 @@ import torch from sglang.multimodal_gen.runtime.layers.linear import ( MergedColumnParallelLinear, + MergedReplicatedLinear, QKVParallelLinear, ) from sglang.multimodal_gen.runtime.models.parameter import PerTensorScaleParameter @@ -55,3 +56,28 @@ def test_qkv_parallel_full_scale_vector_loads_all_fused_slots(): layer.weight_loader_v2(param, torch.tensor([0.25, 0.5, 0.75])) assert torch.equal(param.data, torch.tensor([0.25, 0.5, 0.75])) + + +def test_merged_replicated_linear_loads_independent_weight_shards(): + layer = MergedReplicatedLinear(3, [2, 1, 1], bias=False) + first = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + second = torch.tensor([[7.0, 8.0, 9.0]]) + third = torch.tensor([[10.0, 11.0, 12.0]]) + + layer.weight_loader(layer.weight, first, 0) + layer.weight_loader(layer.weight, second, 1) + layer.weight_loader(layer.weight, third, "v") + + assert torch.equal(layer.weight, torch.cat((first, second, third))) + + +def test_merged_replicated_linear_loads_independent_scalar_shards(): + layer = MergedReplicatedLinear(3, [2, 1, 1], bias=False) + scales = torch.nn.Parameter(torch.zeros(3), requires_grad=False) + scales.needs_scalar_to_array = True + + layer.weight_loader(scales, torch.tensor(0.25), "q") + layer.weight_loader(scales, torch.tensor(0.5), "k") + layer.weight_loader(scales, torch.tensor(0.75), "v") + + assert torch.equal(scales, torch.tensor([0.25, 0.5, 0.75])) diff --git a/python/sglang/multimodal_gen/test/unit/test_pi05_action_api.py b/python/sglang/multimodal_gen/test/unit/test_pi05_action_api.py index edd5d54b7..72a0acd4f 100644 --- a/python/sglang/multimodal_gen/test/unit/test_pi05_action_api.py +++ b/python/sglang/multimodal_gen/test/unit/test_pi05_action_api.py @@ -160,7 +160,10 @@ def test_action_metadata_reports_policy_shape_and_capabilities(): action_horizon=10, action_dim=32, output_action_dim=7, + prompt_token_buckets=[32, 64, 128, 200], + prefix_cuda_graph_max_entries=4, enable_action_cuda_graph=True, + action_cuda_graph_max_entries=6, ) metadata = action_metadata(_server_args(config)) @@ -176,6 +179,13 @@ def test_action_metadata_reports_policy_shape_and_capabilities(): assert metadata["output"]["padded_action_dim"] == 32 assert metadata["runtime"]["materialize_dtype"] == "bf16" assert metadata["runtime"]["enable_autocast"] is True + assert metadata["runtime"]["cuda_graph"] == { + "prefix_enabled": True, + "prefix_max_entries": 4, + "action_enabled": True, + "action_max_entries": 6, + "prompt_token_buckets": [32, 64, 128, 200], + } assert metadata["runtime"]["parallelism"]["num_gpus"] == 1 assert metadata["runtime"]["parallelism"]["kv_gather_degree"] == 1 assert metadata["runtime"]["parallelism"]["prefix_strategy"] == "tp" @@ -185,6 +195,20 @@ def test_action_metadata_reports_policy_shape_and_capabilities(): assert metadata["capabilities"]["openpi_websocket"] +def test_action_metadata_reports_effective_graph_availability(): + config = Pi05PipelineConfig( + prefix_cuda_graph_max_entries=0, + offload_action_expert_after_denoise=True, + ) + + metadata = action_metadata(_server_args(config)) + + assert metadata["runtime"]["cuda_graph"]["prefix_enabled"] is False + assert metadata["runtime"]["cuda_graph"]["action_enabled"] is False + assert metadata["defaults"]["cuda_graph"] is False + assert metadata["capabilities"]["cuda_graph"] is False + + def test_action_generation_response_uses_actual_output_parameters(): output = { "request_id": "action-response-1", diff --git a/python/sglang/multimodal_gen/test/unit/test_pi05_runtime_helpers.py b/python/sglang/multimodal_gen/test/unit/test_pi05_runtime_helpers.py index ea81e89a7..d4ff2de69 100644 --- a/python/sglang/multimodal_gen/test/unit/test_pi05_runtime_helpers.py +++ b/python/sglang/multimodal_gen/test/unit/test_pi05_runtime_helpers.py @@ -3,6 +3,7 @@ from types import SimpleNamespace import numpy as np +import pytest import torch from torch import nn @@ -11,6 +12,8 @@ from sglang.multimodal_gen.configs.pipeline_configs.pi05 import Pi05PipelineConf from sglang.multimodal_gen.runtime.models.vlas.pi05_core import ( Pi05CoreModel, Pi05SiglipVisionModel, + create_sinusoidal_pos_embedding, + make_att_2d_masks, ) from sglang.multimodal_gen.runtime.models.vlas.pi05_policy import ( Pi05CheckpointManifest, @@ -22,13 +25,22 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.p ) from sglang.multimodal_gen.runtime.vla.cuda_graph import ( VLADenoiseGraphRunner, + VLADenoiseGraphSignature, + VLAPrefixGraphRunner, + _BoundedCaptureCache, _CapturedDenoiseGraph, ) +from sglang.multimodal_gen.runtime.vla.observation import VLAObservationBatch from sglang.multimodal_gen.runtime.vla.parallel import VLASplitGroup from sglang.multimodal_gen.runtime.vla.prefix_cache import ( PrefixContext, VLADensePrefixCache, ) +from sglang.multimodal_gen.runtime.vla.prompt_bucketing import ( + bucket_prompt_tokens, + effective_token_length, + select_prompt_token_bucket, +) from sglang.srt.models.siglip import SiglipVisionModel from sglang.srt.runtime_context import get_context @@ -97,10 +109,330 @@ def test_denoise_graph_copies_mutable_prefix_graph_output(): ) current_context = _prefix_context(2.0, None) current_context.layout["mutable_graph_output"] = True + current_context.prefix_pad_masks[:, -1] = False runner._sync_context_if_needed(captured, current_context) assert captured.static_prefix_context.past_key_values[0][0].eq(2.0).all() + assert torch.equal( + captured.static_prefix_context.prefix_pad_masks, + current_context.prefix_pad_masks, + ) + + +def _denoise_signature(prefix_len: int) -> VLADenoiseGraphSignature: + return VLADenoiseGraphSignature( + batch_size=1, + prefix_len=prefix_len, + prefix_full_attention=False, + action_horizon=2, + action_dim=4, + dtype="float32", + parallel_layout="single", + ) + + +def test_denoise_graph_capacity_falls_back_without_capturing_new_signature(): + runner = VLADenoiseGraphRunner(enabled=True, max_entries=1) + runner._cache.entries[_denoise_signature(32)] = object() + fake_cuda_tensor = SimpleNamespace(device=SimpleNamespace(type="cuda")) + + result = runner.capture_or_run( + _denoise_signature(64), + lambda *_args: "eager", + _prefix_context(1.0, None), + fake_cuda_tensor, + object(), + ) + + assert result == "eager" + assert list(runner._cache.entries) == [_denoise_signature(32)] + + +def test_zero_denoise_graph_capacity_disables_runner(): + runner = VLADenoiseGraphRunner(enabled=True, max_entries=0) + + assert not runner.enabled + + +class _FakeGraph: + def __init__(self): + self.reset_calls = 0 + + def reset(self): + self.reset_calls += 1 + + +def _fake_capture(): + return SimpleNamespace(graph=_FakeGraph()) + + +def test_graph_cache_evicts_lru_and_releases_graph(): + cache = _BoundedCaptureCache("test", max_entries=2, evict_on_miss=True) + first = _fake_capture() + second = _fake_capture() + third = _fake_capture() + cache.put("first", first) + cache.put("second", second) + + assert cache.get("first") is first + cache.put("third", third) + + assert tuple(cache.entries) == ("first", "third") + assert second.graph.reset_calls == 1 + assert cache.info().evictions == 1 + + cache.clear() + assert first.graph.reset_calls == 1 + assert third.graph.reset_calls == 1 + + +def test_graph_cache_releases_lru_before_new_capture(): + cache = _BoundedCaptureCache("test", max_entries=1, evict_on_miss=True) + first = _fake_capture() + cache.put("first", first) + + cache.prepare_admission("second") + + assert not cache.entries + assert first.graph.reset_calls == 1 + assert cache.info().evictions == 1 + + +def test_non_evicting_graph_cache_rejects_new_signature_at_capacity(): + cache = _BoundedCaptureCache("test", max_entries=1, evict_on_miss=False) + first = _fake_capture() + rejected = _fake_capture() + cache.put("first", first) + + assert not cache.put("second", rejected) + assert tuple(cache.entries) == ("first",) + assert first.graph.reset_calls == 0 + assert rejected.graph.reset_calls == 1 + + +def test_graph_cache_info_tracks_hits_misses_and_failures(): + cache = _BoundedCaptureCache("test", max_entries=1, evict_on_miss=False) + cache.put("first", _fake_capture()) + + assert cache.get("first") is not None + assert cache.get("missing") is None + cache.mark_failure() + + info = cache.info() + assert info.hits == 1 + assert info.misses == 1 + assert info.captures == 1 + assert info.failures == 1 + + +def test_zero_graph_capacity_disables_both_runners(): + assert not VLAPrefixGraphRunner(enabled=True, max_entries=0).enabled + assert not VLADenoiseGraphRunner(enabled=True, max_entries=0).enabled + + +def test_prefix_prompt_bucket_preserves_tokens_and_masks(): + tokens = torch.arange(40).view(1, 40) + token_masks = torch.ones_like(tokens, dtype=torch.bool) + + bucketed_tokens, bucketed_masks, logical_length, bucket = bucket_prompt_tokens( + tokens, + token_masks, + (32, 64, 128, 200), + ) + + assert logical_length == 40 + assert bucket == 64 + assert bucketed_tokens.shape == (1, 64) + assert torch.equal(bucketed_tokens[:, :40], tokens) + assert bucketed_tokens[:, 40:].eq(0).all() + assert bucketed_masks[:, :40].all() + assert not bucketed_masks[:, 40:].any() + + +def test_prefix_prompt_bucket_preserves_mask_holes(): + tokens = torch.arange(6).view(1, 6) + token_masks = torch.tensor([[True, False, True, False, False, False]]) + + bucketed_tokens, bucketed_masks, logical_length, bucket = bucket_prompt_tokens( + tokens, + token_masks, + (4, 8), + ) + + assert logical_length == 3 + assert bucket == 4 + assert torch.equal(bucketed_tokens, tokens[:, :4]) + assert torch.equal(bucketed_masks, token_masks[:, :4]) + + +def test_prompt_bucket_selection_and_exact_tail(): + assert select_prompt_token_bucket(33, (32, 64, 128)) == 64 + assert select_prompt_token_bucket(129, (32, 64, 128)) is None + + tokens = torch.arange(140).view(1, 140) + token_masks = torch.arange(140).view(1, 140) < 129 + exact_tokens, exact_masks, logical_length, bucket = bucket_prompt_tokens( + tokens, + token_masks, + (32, 64, 128), + ) + + assert logical_length == 129 + assert bucket is None + assert exact_tokens.shape == (1, 129) + assert exact_masks.all() + + +def test_effective_token_length_uses_last_visible_position(): + masks = torch.tensor( + [ + [True, False, False, True, False], + [True, True, False, False, False], + ] + ) + + assert effective_token_length(masks) == 4 + + +def test_pi05_graph_config_validation(): + config = Pi05PipelineConfig() + config.update_pipeline_config( + { + "prompt_token_buckets": [16, 48, 96], + "action_cuda_graph_max_entries": 3, + } + ) + assert config.prompt_token_buckets == [16, 48, 96] + assert config.action_cuda_graph_max_entries == 3 + + invalid_configs = ( + ({"prompt_token_buckets": [32, 32]}, "strictly increasing"), + ({"prompt_token_buckets": [64, 32]}, "strictly increasing"), + ({"prompt_token_buckets": [0, 32]}, "positive"), + ({"prompt_token_buckets": [32, 256]}, "max_token_len"), + ({"prefix_cuda_graph_max_entries": -1}, "prefix_cuda_graph"), + ({"action_cuda_graph_max_entries": -1}, "action_cuda_graph"), + ) + for overrides, match in invalid_configs: + with pytest.raises(ValueError, match=match): + Pi05PipelineConfig(**overrides) + + +def test_prefix_cache_key_distinguishes_bucket_layouts_and_mask_holes(): + model = Pi05PolicyModel.__new__(Pi05PolicyModel) + model.config = Pi05PipelineConfig(prompt_token_buckets=[32, 64]) + model.dtype = torch.bfloat16 + model.model_path = "lerobot/pi05_base" + model._prompt_token_bucketing_enabled = lambda: True + common = dict( + metadata={"camera_order": ("front",)}, + images={"front": torch.zeros(1, 3, 2, 2)}, + image_masks={"front": torch.tensor(True)}, + token_masks=torch.tensor([[True, False, False, True, False]]), + ) + first = SimpleNamespace(tokens=torch.tensor([[1, 2, 3, 4, 0]]), **common) + second = SimpleNamespace(tokens=torch.tensor([[1, 2, 3, 9, 0]]), **common) + + exact = model.build_prefix_cache_key(first) + bucketed = model.build_prefix_cache_key(first, bucket_prompt=True) + + assert exact != bucketed + assert model.build_prefix_cache_key(first) != model.build_prefix_cache_key(second) + + +def test_bucket_miss_keeps_action_denoise_eager(): + model = Pi05PolicyModel.__new__(Pi05PolicyModel) + nn.Module.__init__(model) + model.action_expert = lambda _context, x_t, _timestep, **_kwargs: x_t + 1 + model.graph_runner = SimpleNamespace( + capture_or_run=lambda *_args, **_kwargs: pytest.fail( + "bucket misses must not capture action graphs" + ) + ) + context = SimpleNamespace( + layout={"cuda_graph_eligible": False}, + prefix_len=900, + ) + x_t = torch.zeros(1, 50, 32) + + output = model.denoise_step( + context, + x_t, + torch.ones(1), + use_cuda_graph=True, + ) + + torch.testing.assert_close(output, torch.ones_like(x_t)) + + +def _observation_with_token_len(token_len: int) -> VLAObservationBatch: + tokens = torch.arange(200).view(1, 200) + token_masks = torch.arange(200).view(1, 200) < token_len + return VLAObservationBatch( + prompt=["prompt"], + images={"camera": torch.zeros(1, 3, 4, 4)}, + image_masks={"camera": torch.ones(1, dtype=torch.bool)}, + state=None, + noise=None, + tokens=tokens, + token_masks=token_masks, + batch_size=1, + metadata={"camera_order": ("camera",)}, + ) + + +class _RecordingPrefixRunner: + enabled = True + + def __init__(self): + self.calls = [] + + def capture_or_run(self, signature, _step_fn, inputs): + self.calls.append((signature, inputs)) + return signature + + +def _recording_policy(config: Pi05PipelineConfig) -> Pi05PolicyModel: + model = Pi05PolicyModel.__new__(Pi05PolicyModel) + nn.Module.__init__(model) + model.config = config + model.device = torch.device("cpu") + model.prefix_graph_runner = _RecordingPrefixRunner() + model._prompt_token_bucketing_enabled = lambda: bool(config.prompt_token_buckets) + return model + + +def test_default_prefix_graph_keeps_exact_prompt_signatures(): + model = _recording_policy(Pi05PipelineConfig()) + + signature_33 = model.encode_prefix(_observation_with_token_len(33)) + signature_64 = model.encode_prefix(_observation_with_token_len(64)) + + assert signature_33 != signature_64 + assert model.prefix_graph_runner.calls[0][1][-2].shape == (1, 33) + assert model.prefix_graph_runner.calls[1][1][-2].shape == (1, 64) + + +def test_prefix_prompt_lengths_share_bucket_graph_signature(): + model = _recording_policy( + Pi05PipelineConfig( + prompt_token_buckets=[32, 64, 128, 200], + ) + ) + + signature_33 = model.encode_prefix(_observation_with_token_len(33)) + signature_64 = model.encode_prefix(_observation_with_token_len(64)) + + assert signature_33 == signature_64 + for (_, inputs), expected_token_len in zip( + model.prefix_graph_runner.calls, + (33, 64), + strict=True, + ): + assert inputs[-2].shape == (1, 64) + assert inputs[-1].shape == (1, 64) + assert inputs[-1].sum().item() == expected_token_len def test_prefix_graph_rejects_tensor_parallel_prefix(): @@ -113,6 +445,18 @@ def test_prefix_graph_rejects_tensor_parallel_prefix(): assert not model._prefix_cuda_graph_enabled() +def test_prefix_graph_rejects_partial_language_offload(): + model = Pi05PolicyModel.__new__(Pi05PolicyModel) + model.config = Pi05PipelineConfig( + offload_prefix_language_layer_count_after_prefix=1 + ) + model.device = torch.device("cuda") + model.runtime_role = "all" + model._prefix_tensor_parallel_enabled = lambda: False + + assert not model._prefix_cuda_graph_enabled() + + def test_runai_direct_gpu_loader_does_not_reject_split_roles(monkeypatch): class FakeSafeOpen: def __enter__(self): @@ -268,6 +612,133 @@ def test_prefix_language_embedding_matches_openpi_scale(): ) +def test_cached_pi05_sinusoidal_scaling_is_bit_exact(): + time = torch.tensor([0.125, 0.75], dtype=torch.float32) + dimension = 32 + min_period = 4e-3 + max_period = 4.0 + fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=torch.float64) + period = min_period * (max_period / min_period) ** fraction + scaling = 1.0 / period * 2 * torch.pi + + expected = create_sinusoidal_pos_embedding( + time, + dimension, + min_period, + max_period, + ) + actual = create_sinusoidal_pos_embedding( + time, + dimension, + min_period, + max_period, + scaling=scaling, + ) + + assert torch.equal(actual, expected) + + +def test_prepare_denoise_layout_matches_per_step_construction(): + model = Pi05CoreModel.__new__(Pi05CoreModel) + nn.Module.__init__(model) + prefix_pad_masks = torch.tensor( + [[True, True, False], [True, False, False]], dtype=torch.bool + ) + x_t = torch.zeros(2, 4, 7) + + attention_mask, position_ids = model.prepare_denoise_layout( + prefix_pad_masks, + x_t, + action_position_offset=3, + ) + + suffix_pad_masks = torch.ones(2, 4, dtype=torch.bool) + suffix_att_masks = torch.zeros(2, 4) + suffix_att_masks[:, 0] = 1 + expected_2d_mask = torch.cat( + [ + prefix_pad_masks[:, None, :].expand(2, 4, 3), + make_att_2d_masks(suffix_pad_masks, suffix_att_masks), + ], + dim=2, + ) + expected_mask = model.prepare_attention_masks_4d(expected_2d_mask) + expected_positions = torch.tensor([[5, 6, 7, 8], [4, 5, 6, 7]]) + + assert torch.equal(attention_mask, expected_mask) + assert torch.equal(position_ids, expected_positions) + full_attention_mask, full_attention_positions = model.prepare_denoise_layout( + prefix_pad_masks, + x_t, + prefix_full_attention=True, + action_position_offset=3, + ) + assert full_attention_mask is None + assert torch.equal(full_attention_positions, expected_positions) + + +def test_sample_actions_only_hoists_denoise_layout_for_eager(): + model = Pi05PolicyModel.__new__(Pi05PolicyModel) + nn.Module.__init__(model) + model.config = SimpleNamespace(action_horizon=2, action_dim=3) + model.device = torch.device("cpu") + model.graph_runner = SimpleNamespace(enabled=True) + model._offload_action_expert_between_requests = lambda: False + model._can_use_action_sequence_parallel = lambda *_args: False + model.denoise_step = lambda _ctx, x_t, _t, **_kwargs: torch.zeros_like(x_t) + layout_calls = [] + model.core_model = SimpleNamespace( + prepare_denoise_layout=lambda *args, **kwargs: layout_calls.append( + (args, kwargs) + ) + or (None, torch.zeros(1, 2, dtype=torch.long)) + ) + observation = SimpleNamespace(batch_size=1) + prefix_context = _prefix_context(1.0, "prompt") + prefix_context.layout["full_attention"] = True + noise = torch.zeros(1, 2, 3) + + model.sample_actions( + observation, + prefix_context, + noise=noise, + num_steps=2, + use_cuda_graph=True, + ) + assert not layout_calls + + prefix_context.layout["full_attention"] = False + model.sample_actions( + observation, + prefix_context, + noise=noise, + num_steps=2, + use_cuda_graph=True, + ) + assert not layout_calls + + prefix_context.layout["cuda_graph_eligible"] = False + model.sample_actions( + observation, + prefix_context, + noise=noise, + num_steps=2, + use_cuda_graph=True, + ) + assert len(layout_calls) == 1 + + layout_calls.clear() + prefix_context.layout["cuda_graph_eligible"] = True + model.sample_actions( + observation, + prefix_context, + noise=noise, + num_steps=2, + use_cuda_graph=False, + ) + assert len(layout_calls) == 1 + + def test_uint8_resize_rounds_before_normalization(): image = torch.tensor([[[0.0, 1.0], [2.0, 3.0]]]) / 255.0