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