minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] optimization: optimize Pi0.5 inference and bounded graph serving (#34599)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-08-29 14:47:32 +08:00
committed by GitHub
parent f8f501f2e8
commit 0e1146d04f
13 changed files with 1188 additions and 137 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/cookbook/vla/OpenPI/Pi0.5.mdx
+15 -3
View File
@@ -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. |
python/sglang/multimodal_gen/configs/pipeline_configs/pi05.py
+62
View File
@@ -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
python/sglang/multimodal_gen/runtime/entrypoints/action/protocol.py
+13 -2
View File
@@ -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,
python/sglang/multimodal_gen/runtime/layers/linear.py
+67
View File
@@ -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.
python/sglang/multimodal_gen/runtime/models/vlas/pi05_core.py
+116 -84
View File
@@ -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,
python/sglang/multimodal_gen/runtime/models/vlas/pi05_policy.py
+116 -29
View File
@@ -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:
python/sglang/multimodal_gen/runtime/pipelines/pi05.py
+4 -1
View File
@@ -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,
python/sglang/multimodal_gen/runtime/pipelines_core/stages/vla.py
+9 -3
View File
@@ -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,
python/sglang/multimodal_gen/runtime/vla/cuda_graph.py
+184 -15
View File
@@ -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,
python/sglang/multimodal_gen/runtime/vla/prompt_bucketing.py
+81
View File
@@ -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,
)
python/sglang/multimodal_gen/test/unit/test_parallel_linear_weight_loading.py
+26
View File
@@ -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]))
python/sglang/multimodal_gen/test/unit/test_pi05_action_api.py
+24
View File
@@ -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",
python/sglang/multimodal_gen/test/unit/test_pi05_runtime_helpers.py
+471
View File
@@ -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