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

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

Browse Source
This commit is contained in:
Mick
2026-08-11 09:14:53 +08:00
committed by GitHub
parent df986c4d5e
commit aeab1de1de
8 changed files with 264 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/configs/pipeline_configs/pi05.py
+2
View File
@@ -46,6 +46,8 @@ class Pi05PipelineConfig(PipelineConfig):
image_normalization_std: tuple[float, float, float] = (0.5, 0.5, 0.5)
enable_global_prefix_cache: bool = False
enable_prefix_cuda_graph: bool = True
prefix_cuda_graph_max_entries: int = 1
enable_action_cuda_graph: bool = True
prefix_cache_max_entries: int = 1
prefix_cache_layout_version: str = "pi05-prefix-v1"
python/sglang/multimodal_gen/runtime/models/vlas/pi05_core.py
+1 -5
View File
@@ -1327,25 +1327,21 @@ class Pi05CoreModel(nn.Module):
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
embs = []
pad_masks = []
att_masks = []
image_embs = self.paligemma_with_expert.embed_images(images)
for image_emb, image_mask in zip(image_embs, image_masks, strict=True):
batch_size, num_image_embs = image_emb.shape[:2]
embs.append(image_emb)
pad_masks.append(image_mask[:, None].expand(batch_size, num_image_embs))
att_masks += [0] * num_image_embs
lang_emb = self.paligemma_with_expert.embed_language_tokens(tokens)
# Match OpenPI's Pi0.5 prefix embedding semantics.
lang_emb = lang_emb * math.sqrt(lang_emb.shape[-1])
embs.append(lang_emb)
pad_masks.append(token_masks)
att_masks += [0] * lang_emb.shape[1]
embs = torch.cat(embs, dim=1)
pad_masks = torch.cat(pad_masks, dim=1)
att_masks_t = torch.tensor(att_masks, dtype=torch.bool, device=pad_masks.device)
att_masks_t = att_masks_t[None, :].expand(pad_masks.shape[0], len(att_masks))
att_masks_t = torch.zeros_like(pad_masks)
return embs, pad_masks, att_masks_t
def embed_suffix(
python/sglang/multimodal_gen/runtime/models/vlas/pi05_policy.py
+85 -16
View File
@@ -38,9 +38,11 @@ from sglang.multimodal_gen.runtime.models.vlas.pi05_core import Pi05CoreModel
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import maybe_download_model
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.runtime.vla.denoise_cuda_graph import (
from sglang.multimodal_gen.runtime.vla.cuda_graph import (
VLADenoiseGraphRunner,
VLADenoiseGraphSignature,
VLAPrefixGraphRunner,
VLAPrefixGraphSignature,
)
from sglang.multimodal_gen.runtime.vla.observation import (
VLAObservationBatch,
@@ -161,6 +163,10 @@ class Pi05PolicyModel(nn.Module):
if self.runtime_role != "all":
logger.info("Pi05 split runtime role on rank: %s", self.runtime_role)
self.action_expert = Pi05ActionExpert(config, self.core_model)
self.prefix_graph_runner = VLAPrefixGraphRunner(
enabled=self._prefix_cuda_graph_enabled(),
max_entries=config.prefix_cuda_graph_max_entries,
)
self.graph_runner = VLADenoiseGraphRunner(
enabled=config.enable_action_cuda_graph
)
@@ -176,6 +182,25 @@ class Pi05PolicyModel(nn.Module):
return False
return get_tp_world_size() > 1
def _prefix_cuda_graph_enabled(self) -> bool:
if (
not self.config.enable_prefix_cuda_graph
or self.device.type != "cuda"
or self.runtime_role not in ("all", "prefix")
or self._prefix_tensor_parallel_enabled()
):
return False
return not any(
(
self.config.offload_prefix_image_encoder,
self.config.offload_prefix_image_encoder_after_embed,
self.config.offload_prefix_token_embedding,
self.config.offload_prefix_language_layers,
self.config.offload_prefix_language_layers_after_prefix,
self.config.empty_cache_after_prefix,
)
)
@staticmethod
def _to_empty_preserve_buffers(module: nn.Module, *, device: torch.device) -> None:
buffers = {
@@ -796,6 +821,10 @@ class Pi05PolicyModel(nn.Module):
return None
return paligemma.model.language_model
def _prefix_tensor_parallel_enabled(self) -> bool:
language_model = self._prefix_language_model()
return language_model is not None and language_model.tensor_parallel
def _prefix_kv_requires_tp_gather(self) -> bool:
language_model = self._prefix_language_model()
if language_model is None or not language_model.tensor_parallel:
@@ -822,7 +851,12 @@ class Pi05PolicyModel(nn.Module):
)
)
def encode_prefix(self, observation: VLAObservationBatch) -> PrefixContext:
def encode_prefix(
self,
observation: VLAObservationBatch,
*,
use_cuda_graph: bool = True,
) -> PrefixContext:
camera_order = tuple(observation.metadata.get("camera_order", ()))
images = [
observation.images[name].to(self.device, dtype=torch.float32)
@@ -844,22 +878,57 @@ class Pi05PolicyModel(nn.Module):
prefix_full_attention_hint = all(
bool(observation.image_masks[name].all().item()) for name in camera_order
) and bool(token_masks_cpu.all().item())
past_key_values, prefix_pad_masks, full_attention = (
self.core_model.encode_prefix(
images,
image_masks,
tokens,
token_masks,
prefix_full_attention_hint=prefix_full_attention_hint,
tokens_trimmed=tokens_trimmed,
image_count = len(images)
graph_inputs = tuple([*images, *image_masks, tokens, token_masks])
def encode(current_inputs: tuple[torch.Tensor, ...]) -> PrefixContext:
current_images = list(current_inputs[:image_count])
current_image_masks = list(current_inputs[image_count : 2 * image_count])
current_tokens = current_inputs[-2]
current_token_masks = current_inputs[-1]
past_key_values, prefix_pad_masks, full_attention = (
self.core_model.encode_prefix(
current_images,
current_image_masks,
current_tokens,
current_token_masks,
prefix_full_attention_hint=prefix_full_attention_hint,
tokens_trimmed=tokens_trimmed,
)
)
past_key_values = self._materialize_prefix_kv_for_action(past_key_values)
return PrefixContext(
past_key_values=past_key_values,
prefix_pad_masks=prefix_pad_masks,
prefix_len=prefix_pad_masks.shape[1],
layout={"full_attention": full_attention},
)
if (
not use_cuda_graph
or not self.prefix_graph_runner.enabled
or observation.batch_size != 1
):
return encode(graph_inputs)
signature = VLAPrefixGraphSignature(
batch_size=observation.batch_size,
input_shapes=tuple(tuple(tensor.shape) for tensor in graph_inputs),
input_dtypes=tuple(
str(tensor.dtype).replace("torch.", "") for tensor in graph_inputs
),
static_layout=(
image_count,
prefix_full_attention_hint,
tokens_trimmed,
),
parallel_layout=self.config.parallel_layout_version,
)
past_key_values = self._materialize_prefix_kv_for_action(past_key_values)
return PrefixContext(
past_key_values=past_key_values,
prefix_pad_masks=prefix_pad_masks,
prefix_len=prefix_pad_masks.shape[1],
layout={"full_attention": full_attention},
return self.prefix_graph_runner.capture_or_run(
signature,
encode,
graph_inputs,
)
def sample_noise(
python/sglang/multimodal_gen/runtime/pipelines/pi05.py
+3 -1
View File
@@ -59,13 +59,15 @@ class Pi05Pipeline(ComposedPipelineBase):
or bool(server_args.text_encoder_cpu_offload)
)
logger.info(
"Pi05 memory config: prefix_cache=%s/%s, action_cuda_graph=%s, "
"Pi05 memory config: prefix_cache=%s/%s, cuda_graph=%s/%s/%s, "
"offload_image=%s, offload_image_after_embed=%s, "
"offload_tokens=%s, offload_language_layers=%s, "
"offload_language_after_prefix=%s/%s, "
"offload_action_after_denoise=%s, empty_cache_after_prefix=%s",
pipeline_config.enable_global_prefix_cache,
pipeline_config.prefix_cache_max_entries,
pipeline_config.enable_prefix_cuda_graph,
pipeline_config.prefix_cuda_graph_max_entries,
pipeline_config.enable_action_cuda_graph,
pipeline_config.offload_prefix_image_encoder,
pipeline_config.offload_prefix_image_encoder_after_embed,
python/sglang/multimodal_gen/runtime/pipelines_core/stages/vla.py
+15 -4
View File
@@ -89,6 +89,10 @@ def _effective_prefix_cache_enabled(
)
def _cuda_graph_enabled(batch: Req) -> bool:
return bool(vla_options(batch).get("enable_cuda_graph", True))
def _grouped_fingerprint(
batch: Req,
server_args: ServerArgs,
@@ -122,6 +126,7 @@ def _grouped_fingerprint(
batch.action_horizon,
batch.action_dim,
batch.num_inference_steps,
_cuda_graph_enabled(batch),
)
@@ -180,7 +185,10 @@ class VLAPrefixEncodingStage(PipelineStage):
vla_state(batch)["observation_batch"] for batch in group_batches
]
grouped_observation = collate_vla_observation_batches(observations)
prefix_context = self.policy_model.encode_prefix(grouped_observation)
prefix_context = self.policy_model.encode_prefix(
grouped_observation,
use_cuda_graph=_cuda_graph_enabled(group_batches[0]),
)
prefix_ms = (time.perf_counter() - prefix_start) * 1000
for offset, (index, batch) in enumerate(group):
@@ -299,7 +307,10 @@ class VLAPrefixEncodingStage(PipelineStage):
prefix_start = time.perf_counter()
# 3. run encoding
prefix_context = self.policy_model.encode_prefix(observation)
prefix_context = self.policy_model.encode_prefix(
observation,
use_cuda_graph=_cuda_graph_enabled(batch) and not cache_enabled,
)
if cache_key is not None:
prefix_context.cache_key_digest = cache_key
@@ -371,7 +382,7 @@ class VLAActionDenoisingStage(PipelineStage):
prefix_context,
noise=observation.noise,
num_steps=group_batches[0].num_inference_steps,
use_cuda_graph=bool(options.get("enable_cuda_graph", True)),
use_cuda_graph=_cuda_graph_enabled(group_batches[0]),
generator=None,
)
synchronize_vla_action_tensor(actions)
@@ -417,7 +428,7 @@ class VLAActionDenoisingStage(PipelineStage):
prefix_context,
noise=noise,
num_steps=batch.num_inference_steps,
use_cuda_graph=bool(options.get("enable_cuda_graph", True)),
use_cuda_graph=_cuda_graph_enabled(batch),
generator=batch.generator,
)
synchronize_vla_action_tensor(actions)
python/sglang/multimodal_gen/runtime/vla/denoise_cuda_graph.py → python/sglang/multimodal_gen/runtime/vla/cuda_graph.py
+118 -2
View File
@@ -22,6 +22,22 @@ from sglang.srt.model_executor.runner_utils.pool import (
logger = init_logger(__name__)
@dataclass(frozen=True)
class VLAPrefixGraphSignature:
batch_size: int
input_shapes: tuple[tuple[int, ...], ...]
input_dtypes: tuple[str, ...]
static_layout: tuple[Any, ...]
parallel_layout: str
@dataclass
class _CapturedPrefixGraph:
graph: torch.cuda.CUDAGraph
static_inputs: tuple[torch.Tensor, ...]
static_output: PrefixContext
@dataclass(frozen=True)
class VLADenoiseGraphSignature:
batch_size: int
@@ -76,6 +92,103 @@ def _copy_prefix_context_(dst: PrefixContext, src: PrefixContext) -> None:
dst.cache_key_digest = src.cache_key_digest
class VLAPrefixGraphRunner:
"""Full CUDA graph runner for VLA prefix encoding shape buckets."""
def __init__(self, enabled: bool = True, max_entries: int = 1):
self.max_entries = max(0, max_entries)
self.enabled = enabled and self.max_entries > 0
self._captured: dict[VLAPrefixGraphSignature, _CapturedPrefixGraph] = {}
self._disabled_signatures: set[VLAPrefixGraphSignature] = set()
self._capture_stream: torch.cuda.Stream | None = None
self._graph_pool: Any = None
def _capture(
self,
signature: VLAPrefixGraphSignature,
step_fn: Callable[[tuple[torch.Tensor, ...]], PrefixContext],
inputs: tuple[torch.Tensor, ...],
) -> _CapturedPrefixGraph:
static_inputs = tuple(tensor.detach().clone() for tensor in inputs)
device_module = torch.get_device_module(inputs[0].device)
if self._capture_stream is None:
self._capture_stream = device_module.Stream(device=inputs[0].device)
if self._graph_pool is None:
self._graph_pool = get_or_create_global_graph_memory_pool(device_module)
set_graph_pool_id(self._graph_pool)
device_module.synchronize()
with device_module.stream(self._capture_stream), torch.inference_mode():
step_fn(static_inputs)
self._capture_stream.synchronize()
graph = torch.cuda.CUDAGraph()
with (
device_module.graph(
cuda_graph=graph,
pool=self._graph_pool,
stream=self._capture_stream,
),
torch.inference_mode(),
):
static_output = step_fn(static_inputs)
self._capture_stream.synchronize()
static_output.layout["mutable_graph_output"] = True
captured = _CapturedPrefixGraph(
graph=graph,
static_inputs=static_inputs,
static_output=static_output,
)
self._captured[signature] = captured
logger.info(
"Captured VLA prefix CUDA graph: batch=%d inputs=%s",
signature.batch_size,
signature.input_shapes,
)
return captured
def capture_or_run(
self,
signature: VLAPrefixGraphSignature,
step_fn: Callable[[tuple[torch.Tensor, ...]], PrefixContext],
inputs: tuple[torch.Tensor, ...],
) -> PrefixContext:
if (
not self.enabled
or signature in self._disabled_signatures
or not inputs
or inputs[0].device.type != "cuda"
):
return step_fn(inputs)
captured = self._captured.get(signature)
if captured is None and len(self._captured) >= self.max_entries:
return step_fn(inputs)
try:
if captured is None:
captured = self._capture(signature, step_fn, inputs)
else:
for static_input, current_input in zip(
captured.static_inputs, inputs, strict=True
):
static_input.copy_(current_input)
captured.graph.replay()
torch.get_device_module(inputs[0].device).current_stream(
device=inputs[0].device
).synchronize()
return captured.static_output
except Exception:
self._disabled_signatures.add(signature)
self._captured.pop(signature, None)
logger.warning(
"VLA prefix CUDA graph disabled for signature %s",
signature,
exc_info=True,
)
return step_fn(inputs)
class VLADenoiseGraphRunner:
"""Full CUDA graph runner for one VLA action-denoise step.
@@ -97,13 +210,16 @@ class VLADenoiseGraphRunner:
) -> None:
context_id = id(prefix_context.past_key_values)
context_digest = prefix_context.cache_key_digest
if (
mutable_graph_output = bool(
prefix_context.layout.get("mutable_graph_output", False)
)
if not mutable_graph_output and (
context_digest is not None
and captured.current_context_digest == context_digest
):
captured.current_context_id = context_id
return
if captured.current_context_id == context_id:
if not mutable_graph_output and captured.current_context_id == context_id:
return
_copy_prefix_context_(captured.static_prefix_context, prefix_context)
captured.current_context_id = context_id
python/sglang/multimodal_gen/test/unit/test_pi05_runtime_helpers.py
+31 -2
View File
@@ -21,7 +21,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.p
_preprocess_image,
_resize_with_pad_image_tensor,
)
from sglang.multimodal_gen.runtime.vla.denoise_cuda_graph import (
from sglang.multimodal_gen.runtime.vla.cuda_graph import (
VLADenoiseGraphRunner,
_CapturedDenoiseGraph,
)
@@ -74,7 +74,7 @@ def test_denoise_graph_skips_prefix_copy_for_same_digest(monkeypatch):
raise AssertionError("PrefixContext should not be copied on digest hit")
monkeypatch.setattr(
"sglang.multimodal_gen.runtime.vla.denoise_cuda_graph._copy_prefix_context_",
"sglang.multimodal_gen.runtime.vla.cuda_graph._copy_prefix_context_",
fail_copy,
)
@@ -83,6 +83,35 @@ def test_denoise_graph_skips_prefix_copy_for_same_digest(monkeypatch):
assert captured.static_prefix_context.past_key_values[0][0].eq(1.0).all()
def test_denoise_graph_copies_mutable_prefix_graph_output():
runner = VLADenoiseGraphRunner(enabled=True)
static_context = _prefix_context(1.0, None)
captured = _CapturedDenoiseGraph(
graph=object(),
static_prefix_context=static_context,
static_x_t=torch.empty(1, 2, 4),
static_timestep=torch.empty(1),
static_output=torch.empty(1, 2, 4),
current_context_id=123,
)
current_context = _prefix_context(2.0, None)
current_context.layout["mutable_graph_output"] = True
runner._sync_context_if_needed(captured, current_context)
assert captured.static_prefix_context.past_key_values[0][0].eq(2.0).all()
def test_prefix_graph_rejects_tensor_parallel_prefix():
model = Pi05PolicyModel.__new__(Pi05PolicyModel)
model.config = Pi05PipelineConfig()
model.device = torch.device("cuda")
model.runtime_role = "all"
model._prefix_language_model = lambda: SimpleNamespace(tensor_parallel=True)
assert not model._prefix_cuda_graph_enabled()
def test_runai_direct_gpu_loader_does_not_reject_split_roles(monkeypatch):
class FakeSafeOpen:
def __enter__(self):