[NPU] [Diffusion] support distributed inference pipeline for GLM-Image (#31320)
Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>
This commit is contained in:
co-authored by
Xiaoyu Zhang
parent
803b4fb31c
commit
ecbadf0b4b
@@ -153,6 +153,118 @@ sglang serve --model-path ... --disagg-role server \
|
|||||||
--decoder-urls "tcp://10.0.0.5:35000"
|
--decoder-urls "tcp://10.0.0.5:35000"
|
||||||
```
|
```
|
||||||
|
|
||||||
|
### GLM-Image distributed mode
|
||||||
|
|
||||||
|
GLM-Image can batch AR generation in the head and dispatch the resulting prior
|
||||||
|
tokens to distributed denoiser workers. Each worker runs prompt/glyph preparation,
|
||||||
|
DiT, and VAE decoding locally; no latent or embedding tensors are transferred.
|
||||||
|
|
||||||
|
The following 16-device deployment uses devices 0-1 for the external AR server
|
||||||
|
and devices 2-15 for 14 independent batch-1 Cache-DiT denoisers.
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# Run 14 distributed denoisers
|
||||||
|
DISAGG_SERVER="tcp://127.0.0.1:19655"
|
||||||
|
MODEL_PATH="zai-org/GLM-Image"
|
||||||
|
BASE_MASTER_PORT=29005
|
||||||
|
|
||||||
|
export SGLANG_CACHE_DIT_FN=2
|
||||||
|
export SGLANG_CACHE_DIT_BN=1
|
||||||
|
export SGLANG_CACHE_DIT_WARMUP=4
|
||||||
|
export SGLANG_CACHE_DIT_RDT=0.4
|
||||||
|
export SGLANG_CACHE_DIT_MC=4
|
||||||
|
export SGLANG_CACHE_DIT_TAYLORSEER=true
|
||||||
|
export SGLANG_CACHE_DIT_TS_ORDER=2
|
||||||
|
export SGLANG_CACHE_DIT_ENABLED=true
|
||||||
|
|
||||||
|
worker_pids=()
|
||||||
|
|
||||||
|
cleanup() {
|
||||||
|
trap - EXIT
|
||||||
|
if ((${#worker_pids[@]})); then
|
||||||
|
kill "${worker_pids[@]}" 2>/dev/null || true
|
||||||
|
wait "${worker_pids[@]}" 2>/dev/null || true
|
||||||
|
fi
|
||||||
|
}
|
||||||
|
|
||||||
|
trap cleanup EXIT
|
||||||
|
trap 'exit 130' INT
|
||||||
|
trap 'exit 143' TERM
|
||||||
|
|
||||||
|
for i in $(seq 2 15); do
|
||||||
|
scheduler_port=$((19001 + i))
|
||||||
|
master_port=$((BASE_MASTER_PORT + i))
|
||||||
|
|
||||||
|
sglang serve \
|
||||||
|
--model-path "$MODEL_PATH" \
|
||||||
|
--disagg-role denoiser \
|
||||||
|
--disagg-server-addr "$DISAGG_SERVER" \
|
||||||
|
--srt-encoder-url http://127.0.0.1:30020 \
|
||||||
|
--scheduler-port "$scheduler_port" \
|
||||||
|
--master-port "$master_port" \
|
||||||
|
--num-gpus 1 \
|
||||||
|
--base-gpu-id "$i" \
|
||||||
|
--denoiser-sp 1 \
|
||||||
|
--cfg-parallel-size 1 \
|
||||||
|
--batching-max-size 1 \
|
||||||
|
--dit-cpu-offload false \
|
||||||
|
--attention-backend fa &
|
||||||
|
worker_pids+=("$!")
|
||||||
|
done
|
||||||
|
|
||||||
|
# Stop all denoisers if any worker exits or fails during startup.
|
||||||
|
wait -n "${worker_pids[@]}"
|
||||||
|
```
|
||||||
|
|
||||||
|
Run the external AR server:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
sglang serve \
|
||||||
|
--model-path ./zai-org/GLM-Image/vision_language_encoder/ \
|
||||||
|
--tokenizer-path ./zai-org/GLM-Image/processor/ \
|
||||||
|
--enable-multimodal \
|
||||||
|
--cuda-graph-max-bs 28 \
|
||||||
|
--device npu \
|
||||||
|
--attention-backend ascend \
|
||||||
|
--disable-fast-image-processor \
|
||||||
|
--tp-size 2 \
|
||||||
|
--host 0.0.0.0 \
|
||||||
|
--port 30020 \
|
||||||
|
--mem-fraction-static 0.8
|
||||||
|
```
|
||||||
|
|
||||||
|
Run the public head. `--encoder-urls` and `--decoder-urls` are intentionally
|
||||||
|
omitted for this topology.
|
||||||
|
|
||||||
|
```bash
|
||||||
|
sglang serve \
|
||||||
|
--model-path zai-org/GLM-Image \
|
||||||
|
--disagg-role server \
|
||||||
|
--srt-encoder-url http://127.0.0.1:30020 \
|
||||||
|
--srt-encoder-timeout 300 \
|
||||||
|
--denoiser-urls "tcp://127.0.0.1:19003;tcp://127.0.0.1:19004;tcp://127.0.0.1:19005;tcp://127.0.0.1:19006;tcp://127.0.0.1:19007;tcp://127.0.0.1:19008;tcp://127.0.0.1:19009;tcp://127.0.0.1:19010;tcp://127.0.0.1:19011;tcp://127.0.0.1:19012;tcp://127.0.0.1:19013;tcp://127.0.0.1:19014;tcp://127.0.0.1:19015;tcp://127.0.0.1:19016" \
|
||||||
|
--batching-mode dynamic \
|
||||||
|
--batching-max-size 28 \
|
||||||
|
--batching-delay-ms 30 \
|
||||||
|
--enable-batching-metrics \
|
||||||
|
--host 0.0.0.0 \
|
||||||
|
--port 30052 \
|
||||||
|
--scheduler-port 19655 \
|
||||||
|
--output-path ./outputs
|
||||||
|
```
|
||||||
|
|
||||||
|
Workers return decoded pixels to the head, which saves and serves the final
|
||||||
|
files from `--output-path`. The PR benchmark used
|
||||||
|
[longtext-bench.zip](https://github.com/user-attachments/files/29779516/longtext-bench.zip):
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python fetch_images.py \
|
||||||
|
--base-url http://localhost:30052/v1 \
|
||||||
|
--model GLM-Image-distributed-test \
|
||||||
|
--output-dir generated_images \
|
||||||
|
--max-concurrency 56
|
||||||
|
```
|
||||||
|
|
||||||
## Port Convention
|
## Port Convention
|
||||||
|
|
||||||
Result endpoints are derived deterministically from the head node's `--scheduler-port` (default: 5555):
|
Result endpoints are derived deterministically from the head node's `--scheduler-port` (default: 5555):
|
||||||
@@ -196,6 +308,9 @@ Tensor data between roles (encoder→denoiser, denoiser→decoder) is transferre
|
|||||||
|
|
||||||
**mooncake-transfer-engine** is required for disaggregated diffusion. It provides RDMA for direct GPU-to-GPU data movement.
|
**mooncake-transfer-engine** is required for disaggregated diffusion. It provides RDMA for direct GPU-to-GPU data movement.
|
||||||
|
|
||||||
|
The GLM-Image distributed mode is an exception: it relays only prior token IDs
|
||||||
|
and request metadata over ZMQ and does not require Mooncake.
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
pip install mooncake-transfer-engine
|
pip install mooncake-transfer-engine
|
||||||
```
|
```
|
||||||
|
|||||||
@@ -7,9 +7,14 @@ import pickle
|
|||||||
import threading
|
import threading
|
||||||
import time
|
import time
|
||||||
from collections import deque
|
from collections import deque
|
||||||
from dataclasses import dataclass
|
from concurrent.futures import Future, ThreadPoolExecutor
|
||||||
|
from dataclasses import dataclass, field
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
|
import torch
|
||||||
import zmq
|
import zmq
|
||||||
|
from transformers import AutoProcessor
|
||||||
|
from zmq.utils.monitor import recv_monitor_message
|
||||||
|
|
||||||
from sglang.multimodal_gen.runtime.disaggregation.dispatch_policy import (
|
from sglang.multimodal_gen.runtime.disaggregation.dispatch_policy import (
|
||||||
PoolDispatcher,
|
PoolDispatcher,
|
||||||
@@ -20,6 +25,7 @@ from sglang.multimodal_gen.runtime.disaggregation.request_state import (
|
|||||||
)
|
)
|
||||||
from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
|
from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
|
||||||
from sglang.multimodal_gen.runtime.disaggregation.transport.codec import (
|
from sglang.multimodal_gen.runtime.disaggregation.transport.codec import (
|
||||||
|
send_tensors,
|
||||||
unpack_tensors,
|
unpack_tensors,
|
||||||
)
|
)
|
||||||
from sglang.multimodal_gen.runtime.disaggregation.transport.protocol import (
|
from sglang.multimodal_gen.runtime.disaggregation.transport.protocol import (
|
||||||
@@ -31,11 +37,70 @@ from sglang.multimodal_gen.runtime.disaggregation.transport.protocol import (
|
|||||||
encode_transfer_msg,
|
encode_transfer_msg,
|
||||||
is_transfer_message,
|
is_transfer_message,
|
||||||
)
|
)
|
||||||
|
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import (
|
||||||
|
OutputBatch,
|
||||||
|
Req,
|
||||||
|
)
|
||||||
from sglang.multimodal_gen.runtime.utils.common import get_zmq_socket
|
from sglang.multimodal_gen.runtime.utils.common import get_zmq_socket
|
||||||
|
from sglang.multimodal_gen.runtime.utils.perf_logger import (
|
||||||
|
MemorySnapshot,
|
||||||
|
RequestMetrics,
|
||||||
|
)
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.glm_image import (
|
||||||
|
GlmImageAR,
|
||||||
|
)
|
||||||
|
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
def _deserialize_request_metrics(data: dict | None) -> RequestMetrics | None:
|
||||||
|
if data is None:
|
||||||
|
return None
|
||||||
|
|
||||||
|
metrics = RequestMetrics(request_id=data["request_id"])
|
||||||
|
metrics.stages = data.get("stages", {})
|
||||||
|
metrics.steps = data.get("steps", [])
|
||||||
|
metrics.total_duration_ms = data.get("total_duration_ms", 0.0)
|
||||||
|
for name, snapshot in data.get("memory_snapshots", {}).items():
|
||||||
|
metrics.memory_snapshots[name] = MemorySnapshot(
|
||||||
|
allocated_mb=snapshot.get("allocated_mb", 0.0),
|
||||||
|
reserved_mb=snapshot.get("reserved_mb", 0.0),
|
||||||
|
peak_allocated_mb=snapshot.get("peak_allocated_mb", 0.0),
|
||||||
|
peak_reserved_mb=snapshot.get("peak_reserved_mb", 0.0),
|
||||||
|
peak_host_anon_mb=snapshot.get("peak_host_anon_mb", 0.0),
|
||||||
|
)
|
||||||
|
return metrics
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class _GlmDistributedRequest:
|
||||||
|
"""Track one client request as it moves from AR to a denoiser."""
|
||||||
|
|
||||||
|
client_request_id: str
|
||||||
|
req: Req
|
||||||
|
enqueue_time: float
|
||||||
|
worker_idx: int | None = None
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class _GlmDistributedModeState:
|
||||||
|
"""State used only by the GLM external-AR distributed topology."""
|
||||||
|
|
||||||
|
server_args: "ServerArgs"
|
||||||
|
ar_stage: "GlmImageAR"
|
||||||
|
executor: ThreadPoolExecutor
|
||||||
|
denoiser_worker_available: list[bool]
|
||||||
|
pending_ar_requests: deque[_GlmDistributedRequest] = field(default_factory=deque)
|
||||||
|
pending_denoiser_requests: deque[_GlmDistributedRequest] = field(
|
||||||
|
default_factory=deque
|
||||||
|
)
|
||||||
|
denoiser_requests: dict[str, _GlmDistributedRequest] = field(default_factory=dict)
|
||||||
|
active_ar_batch: tuple[Future, list[_GlmDistributedRequest]] | None = None
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
class _EncoderTTAEntry:
|
class _EncoderTTAEntry:
|
||||||
request_id: str
|
request_id: str
|
||||||
@@ -89,9 +154,11 @@ class DiffusionServer:
|
|||||||
dispatch_policy_name: str = "round_robin",
|
dispatch_policy_name: str = "round_robin",
|
||||||
timeout_s: float = 600.0,
|
timeout_s: float = 600.0,
|
||||||
encoder_capacity: int = 4,
|
encoder_capacity: int = 4,
|
||||||
denoiser_capacity: int = 2,
|
denoiser_capacity_per_worker: int = 2,
|
||||||
decoder_capacity: int = 4,
|
decoder_capacity: int = 4,
|
||||||
p2p_mode: bool = True,
|
p2p_mode: bool = True,
|
||||||
|
server_args=None,
|
||||||
|
glm_distributed_mode_enabled: bool = False,
|
||||||
):
|
):
|
||||||
self._frontend_endpoint = frontend_endpoint
|
self._frontend_endpoint = frontend_endpoint
|
||||||
self._encoder_work_endpoints = encoder_work_endpoints
|
self._encoder_work_endpoints = encoder_work_endpoints
|
||||||
@@ -108,9 +175,9 @@ class DiffusionServer:
|
|||||||
|
|
||||||
self._tracker = RequestTracker()
|
self._tracker = RequestTracker()
|
||||||
self._dispatcher = PoolDispatcher(
|
self._dispatcher = PoolDispatcher(
|
||||||
num_encoders=self._num_encoders,
|
num_encoders=max(1, self._num_encoders),
|
||||||
num_denoisers=self._num_denoisers,
|
num_denoisers=self._num_denoisers,
|
||||||
num_decoders=self._num_decoders,
|
num_decoders=max(1, self._num_decoders),
|
||||||
policy_name=dispatch_policy_name,
|
policy_name=dispatch_policy_name,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -124,7 +191,7 @@ class DiffusionServer:
|
|||||||
|
|
||||||
# FreeBufferSlots per instance
|
# FreeBufferSlots per instance
|
||||||
self._encoder_free_slots = [encoder_capacity] * self._num_encoders
|
self._encoder_free_slots = [encoder_capacity] * self._num_encoders
|
||||||
self._denoiser_free_slots = [denoiser_capacity] * self._num_denoisers
|
self._denoiser_free_slots = [denoiser_capacity_per_worker] * self._num_denoisers
|
||||||
self._decoder_free_slots = [decoder_capacity] * self._num_decoders
|
self._decoder_free_slots = [decoder_capacity] * self._num_decoders
|
||||||
|
|
||||||
# TTA queues per role type
|
# TTA queues per role type
|
||||||
@@ -133,6 +200,24 @@ class DiffusionServer:
|
|||||||
self._decoder_tta: deque[_RoleTTAEntry] = deque()
|
self._decoder_tta: deque[_RoleTTAEntry] = deque()
|
||||||
|
|
||||||
self._transfer_mode = p2p_mode
|
self._transfer_mode = p2p_mode
|
||||||
|
self._glm_distributed_state: _GlmDistributedModeState | None = None
|
||||||
|
|
||||||
|
if glm_distributed_mode_enabled:
|
||||||
|
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.glm_image import (
|
||||||
|
GlmImageAR,
|
||||||
|
)
|
||||||
|
|
||||||
|
processor = AutoProcessor.from_pretrained(
|
||||||
|
server_args.model_path, subfolder="processor"
|
||||||
|
)
|
||||||
|
self._glm_distributed_state = _GlmDistributedModeState(
|
||||||
|
server_args=server_args,
|
||||||
|
ar_stage=GlmImageAR(processor=processor, vision_language_encoder=None),
|
||||||
|
executor=ThreadPoolExecutor(
|
||||||
|
max_workers=1, thread_name_prefix="glm-distributed-ar"
|
||||||
|
),
|
||||||
|
denoiser_worker_available=[False] * self._num_denoisers,
|
||||||
|
)
|
||||||
self._transfer_state: dict[str, _TransferRequestState] = {}
|
self._transfer_state: dict[str, _TransferRequestState] = {}
|
||||||
|
|
||||||
# Per-instance registration: instance_idx -> {session_id, pool_ptr, pool_size}
|
# Per-instance registration: instance_idx -> {session_id, pool_ptr, pool_size}
|
||||||
@@ -197,6 +282,10 @@ class DiffusionServer:
|
|||||||
if self._thread is not None:
|
if self._thread is not None:
|
||||||
self._thread.join(timeout=5.0)
|
self._thread.join(timeout=5.0)
|
||||||
self._thread = None
|
self._thread = None
|
||||||
|
if self._glm_distributed_state is not None:
|
||||||
|
self._glm_distributed_state.executor.shutdown(
|
||||||
|
wait=False, cancel_futures=True
|
||||||
|
)
|
||||||
|
|
||||||
def _event_loop(self) -> None:
|
def _event_loop(self) -> None:
|
||||||
frontend, _ = get_zmq_socket(
|
frontend, _ = get_zmq_socket(
|
||||||
@@ -209,9 +298,16 @@ class DiffusionServer:
|
|||||||
encoder_pushes.append(sock)
|
encoder_pushes.append(sock)
|
||||||
|
|
||||||
denoiser_pushes: list[zmq.Socket] = []
|
denoiser_pushes: list[zmq.Socket] = []
|
||||||
|
denoiser_monitors: list[zmq.Socket] = []
|
||||||
for i, ep in enumerate(self._denoiser_work_endpoints):
|
for i, ep in enumerate(self._denoiser_work_endpoints):
|
||||||
sock, _ = get_zmq_socket(self._context, zmq.PUSH, ep, bind=False)
|
sock, _ = get_zmq_socket(self._context, zmq.PUSH, ep, bind=False)
|
||||||
denoiser_pushes.append(sock)
|
denoiser_pushes.append(sock)
|
||||||
|
if self._glm_distributed_state is not None:
|
||||||
|
denoiser_monitors.append(
|
||||||
|
sock.get_monitor_socket(
|
||||||
|
events=zmq.EVENT_CONNECTED | zmq.EVENT_DISCONNECTED
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
decoder_pushes: list[zmq.Socket] = []
|
decoder_pushes: list[zmq.Socket] = []
|
||||||
for i, ep in enumerate(self._decoder_work_endpoints):
|
for i, ep in enumerate(self._decoder_work_endpoints):
|
||||||
@@ -233,6 +329,8 @@ class DiffusionServer:
|
|||||||
poller.register(encoder_result_pull, zmq.POLLIN)
|
poller.register(encoder_result_pull, zmq.POLLIN)
|
||||||
poller.register(denoiser_result_pull, zmq.POLLIN)
|
poller.register(denoiser_result_pull, zmq.POLLIN)
|
||||||
poller.register(decoder_result_pull, zmq.POLLIN)
|
poller.register(decoder_result_pull, zmq.POLLIN)
|
||||||
|
for monitor in denoiser_monitors:
|
||||||
|
poller.register(monitor, zmq.POLLIN)
|
||||||
|
|
||||||
self._encoder_pushes = encoder_pushes
|
self._encoder_pushes = encoder_pushes
|
||||||
self._denoiser_pushes = denoiser_pushes
|
self._denoiser_pushes = denoiser_pushes
|
||||||
@@ -246,6 +344,7 @@ class DiffusionServer:
|
|||||||
+ encoder_pushes
|
+ encoder_pushes
|
||||||
+ denoiser_pushes
|
+ denoiser_pushes
|
||||||
+ decoder_pushes
|
+ decoder_pushes
|
||||||
|
+ denoiser_monitors
|
||||||
)
|
)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
@@ -266,7 +365,16 @@ class DiffusionServer:
|
|||||||
if decoder_result_pull in events:
|
if decoder_result_pull in events:
|
||||||
self._handle_role_result(decoder_result_pull, RoleType.DECODER)
|
self._handle_role_result(decoder_result_pull, RoleType.DECODER)
|
||||||
|
|
||||||
self._drain_all_queues()
|
for worker_idx, monitor in enumerate(denoiser_monitors):
|
||||||
|
if monitor in events:
|
||||||
|
self._handle_glm_denoiser_monitor_event(worker_idx, monitor)
|
||||||
|
|
||||||
|
if self._glm_distributed_state is not None:
|
||||||
|
self._process_glm_ar_batch_result_if_ready()
|
||||||
|
self._dispatch_glm_ar_batch_if_ready()
|
||||||
|
self._dispatch_glm_denoiser_requests_if_ready()
|
||||||
|
else:
|
||||||
|
self._drain_all_queues()
|
||||||
|
|
||||||
except Exception:
|
except Exception:
|
||||||
logger.exception("DiffusionServer event loop error")
|
logger.exception("DiffusionServer event loop error")
|
||||||
@@ -285,7 +393,9 @@ class DiffusionServer:
|
|||||||
self._handle_transfer_result(frames, role)
|
self._handle_transfer_result(frames, role)
|
||||||
return
|
return
|
||||||
|
|
||||||
if role == RoleType.DECODER:
|
if self._glm_distributed_state is not None and role == RoleType.DENOISER:
|
||||||
|
self._handle_glm_denoiser_result_frames(frames)
|
||||||
|
elif role == RoleType.DECODER:
|
||||||
self._handle_decoder_result_frames(frames)
|
self._handle_decoder_result_frames(frames)
|
||||||
else:
|
else:
|
||||||
# Non-transfer frames from encoder/denoiser are error results
|
# Non-transfer frames from encoder/denoiser are error results
|
||||||
@@ -378,6 +488,22 @@ class DiffusionServer:
|
|||||||
self._tracker.transition(request_id, RequestState.ENCODER_WAITING)
|
self._tracker.transition(request_id, RequestState.ENCODER_WAITING)
|
||||||
except ValueError:
|
except ValueError:
|
||||||
pass
|
pass
|
||||||
|
if self._glm_distributed_state is not None:
|
||||||
|
if (
|
||||||
|
not isinstance(req.prompt, str)
|
||||||
|
or getattr(req, "image_path", None) is not None
|
||||||
|
):
|
||||||
|
self._complete_with_error(
|
||||||
|
request_id,
|
||||||
|
"GLM distributed mode supports one text prompt without image input",
|
||||||
|
)
|
||||||
|
return
|
||||||
|
now = time.monotonic()
|
||||||
|
self._glm_distributed_state.pending_ar_requests.append(
|
||||||
|
_GlmDistributedRequest(request_id, req, now)
|
||||||
|
)
|
||||||
|
return
|
||||||
|
|
||||||
self._encoder_tta.append(
|
self._encoder_tta.append(
|
||||||
_EncoderTTAEntry(
|
_EncoderTTAEntry(
|
||||||
request_id=request_id,
|
request_id=request_id,
|
||||||
@@ -390,11 +516,245 @@ class DiffusionServer:
|
|||||||
request_id,
|
request_id,
|
||||||
)
|
)
|
||||||
|
|
||||||
def _handle_decoder_result_frames(self, frames: list) -> None:
|
def _dispatch_glm_ar_batch_if_ready(self) -> None:
|
||||||
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import (
|
"""Dispatch one compatible batch to the external AR server."""
|
||||||
OutputBatch,
|
state = self._glm_distributed_state
|
||||||
|
assert state is not None
|
||||||
|
if state.active_ar_batch is not None or not state.pending_ar_requests:
|
||||||
|
return
|
||||||
|
|
||||||
|
batch_max_size = (
|
||||||
|
max(1, state.server_args.batching_max_size)
|
||||||
|
if state.server_args.batching_mode == "dynamic"
|
||||||
|
else 1
|
||||||
|
)
|
||||||
|
base = state.pending_ar_requests[0]
|
||||||
|
indices = [0]
|
||||||
|
output_slots = max(1, int(base.req.num_outputs_per_prompt or 1))
|
||||||
|
for index in range(1, len(state.pending_ar_requests)):
|
||||||
|
if output_slots >= batch_max_size:
|
||||||
|
break
|
||||||
|
candidate = state.pending_ar_requests[index]
|
||||||
|
if (base.req.height, base.req.width) == (
|
||||||
|
candidate.req.height,
|
||||||
|
candidate.req.width,
|
||||||
|
):
|
||||||
|
candidate_outputs = max(
|
||||||
|
1, int(candidate.req.num_outputs_per_prompt or 1)
|
||||||
|
)
|
||||||
|
if output_slots + candidate_outputs > batch_max_size:
|
||||||
|
continue
|
||||||
|
indices.append(index)
|
||||||
|
output_slots += candidate_outputs
|
||||||
|
|
||||||
|
waited = time.monotonic() - base.enqueue_time
|
||||||
|
batch_delay_s = state.server_args.batching_delay_ms / 1000.0
|
||||||
|
if output_slots < batch_max_size and waited < batch_delay_s:
|
||||||
|
return
|
||||||
|
|
||||||
|
requests = [state.pending_ar_requests[index] for index in indices]
|
||||||
|
for index in reversed(indices):
|
||||||
|
del state.pending_ar_requests[index]
|
||||||
|
|
||||||
|
for request in requests:
|
||||||
|
try:
|
||||||
|
self._tracker.transition(
|
||||||
|
request.client_request_id, RequestState.ENCODER_RUNNING
|
||||||
|
)
|
||||||
|
except ValueError:
|
||||||
|
pass
|
||||||
|
state.active_ar_batch = (
|
||||||
|
state.executor.submit(
|
||||||
|
state.ar_stage.generate_and_assign_prior_tokens,
|
||||||
|
[request.req for request in requests],
|
||||||
|
state.server_args,
|
||||||
|
device=torch.device("cpu"),
|
||||||
|
),
|
||||||
|
requests,
|
||||||
|
)
|
||||||
|
logger.info(
|
||||||
|
"GLM distributed AR dispatched batch size=%d requests, %d outputs",
|
||||||
|
len(requests),
|
||||||
|
output_slots,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
def _process_glm_ar_batch_result_if_ready(self) -> None:
|
||||||
|
"""Move a completed AR batch into the denoiser dispatch queue."""
|
||||||
|
state = self._glm_distributed_state
|
||||||
|
assert state is not None
|
||||||
|
active_batch = state.active_ar_batch
|
||||||
|
if active_batch is None or not active_batch[0].done():
|
||||||
|
return
|
||||||
|
future, requests = active_batch
|
||||||
|
state.active_ar_batch = None
|
||||||
|
try:
|
||||||
|
future.result()
|
||||||
|
except Exception as error:
|
||||||
|
for request in requests:
|
||||||
|
self._complete_with_error(
|
||||||
|
request.client_request_id, f"GLM AR error: {error}"
|
||||||
|
)
|
||||||
|
return
|
||||||
|
|
||||||
|
group_id = f"glm-distributed::{time.monotonic_ns()}"
|
||||||
|
for request_index, request in enumerate(requests):
|
||||||
|
if self._tracker.get(request.client_request_id) is None:
|
||||||
|
continue
|
||||||
|
try:
|
||||||
|
self._tracker.transition(
|
||||||
|
request.client_request_id, RequestState.ENCODER_DONE
|
||||||
|
)
|
||||||
|
self._tracker.transition(
|
||||||
|
request.client_request_id, RequestState.DENOISING_WAITING
|
||||||
|
)
|
||||||
|
except ValueError:
|
||||||
|
pass
|
||||||
|
denoiser_req = request.req
|
||||||
|
denoiser_req.request_id = f"{group_id}::request::{request_index}"
|
||||||
|
state.denoiser_requests[denoiser_req.request_id] = request
|
||||||
|
state.pending_denoiser_requests.append(request)
|
||||||
|
|
||||||
|
def _dispatch_glm_denoiser_requests_if_ready(self) -> None:
|
||||||
|
"""Dispatch AR-complete requests to available denoisers.
|
||||||
|
|
||||||
|
Each GLM denoiser accepts one request at a time. Dispatch continues until
|
||||||
|
either the pending queue is empty or every connected worker is busy.
|
||||||
|
"""
|
||||||
|
from sglang.multimodal_gen.runtime.disaggregation.scheduler_mixin import (
|
||||||
|
extract_transfer_fields,
|
||||||
|
)
|
||||||
|
|
||||||
|
state = self._glm_distributed_state
|
||||||
|
assert state is not None
|
||||||
|
while state.pending_denoiser_requests:
|
||||||
|
available_slots = [
|
||||||
|
slots if state.denoiser_worker_available[index] else 0
|
||||||
|
for index, slots in enumerate(self._denoiser_free_slots)
|
||||||
|
]
|
||||||
|
worker_idx = self._dispatcher.select_denoiser_with_capacity(available_slots)
|
||||||
|
if worker_idx is None:
|
||||||
|
return
|
||||||
|
request = state.pending_denoiser_requests.popleft()
|
||||||
|
self._denoiser_free_slots[worker_idx] -= 1
|
||||||
|
request.worker_idx = worker_idx
|
||||||
|
tensor_fields, scalar_fields = extract_transfer_fields(request.req)
|
||||||
|
scalar_fields["request_id"] = request.req.request_id
|
||||||
|
send_tensors(
|
||||||
|
self._denoiser_pushes[worker_idx], tensor_fields, scalar_fields
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
self._tracker.transition(
|
||||||
|
request.client_request_id,
|
||||||
|
RequestState.DENOISING_RUNNING,
|
||||||
|
denoiser_instance=worker_idx,
|
||||||
|
)
|
||||||
|
except ValueError:
|
||||||
|
pass
|
||||||
|
logger.info(
|
||||||
|
"GLM distributed mode dispatched request with %d output(s) "
|
||||||
|
"to denoiser[%d]",
|
||||||
|
request.req.num_outputs_per_prompt,
|
||||||
|
worker_idx,
|
||||||
|
)
|
||||||
|
|
||||||
|
def _handle_glm_denoiser_monitor_event(
|
||||||
|
self, worker_idx: int, monitor: zmq.Socket
|
||||||
|
) -> None:
|
||||||
|
"""Update dispatch eligibility when a denoiser connects or disconnects."""
|
||||||
|
state = self._glm_distributed_state
|
||||||
|
assert state is not None
|
||||||
|
event = recv_monitor_message(monitor, flags=zmq.NOBLOCK)["event"]
|
||||||
|
if event == zmq.EVENT_DISCONNECTED:
|
||||||
|
state.denoiser_worker_available[worker_idx] = False
|
||||||
|
self._denoiser_free_slots[worker_idx] = 0
|
||||||
|
replayed_requests = [
|
||||||
|
request
|
||||||
|
for request in state.denoiser_requests.values()
|
||||||
|
if request.worker_idx == worker_idx
|
||||||
|
]
|
||||||
|
for request in replayed_requests:
|
||||||
|
request.worker_idx = None
|
||||||
|
try:
|
||||||
|
self._tracker.transition(
|
||||||
|
request.client_request_id, RequestState.DENOISING_WAITING
|
||||||
|
)
|
||||||
|
except ValueError:
|
||||||
|
pass
|
||||||
|
state.pending_denoiser_requests.extendleft(reversed(replayed_requests))
|
||||||
|
logger.warning(
|
||||||
|
"GLM denoiser[%d] disconnected; requeued %d request(s)",
|
||||||
|
worker_idx,
|
||||||
|
len(replayed_requests),
|
||||||
|
)
|
||||||
|
elif event == zmq.EVENT_CONNECTED:
|
||||||
|
registered = worker_idx in self._denoiser_peers
|
||||||
|
state.denoiser_worker_available[worker_idx] = True
|
||||||
|
if not any(
|
||||||
|
request.worker_idx == worker_idx
|
||||||
|
for request in state.denoiser_requests.values()
|
||||||
|
):
|
||||||
|
self._denoiser_free_slots[worker_idx] = 1
|
||||||
|
logger.info(
|
||||||
|
"GLM denoiser[%d] connected (registered=%s)",
|
||||||
|
worker_idx,
|
||||||
|
registered,
|
||||||
|
)
|
||||||
|
|
||||||
|
def _handle_glm_denoiser_result_frames(self, frames: list) -> None:
|
||||||
|
"""Return decoded denoiser output to the originating HTTP request."""
|
||||||
|
state = self._glm_distributed_state
|
||||||
|
assert state is not None
|
||||||
|
tensor_fields, scalar_fields = unpack_tensors(frames, device="cpu")
|
||||||
|
denoiser_request_id = scalar_fields.get("request_id")
|
||||||
|
request = state.denoiser_requests.pop(denoiser_request_id, None)
|
||||||
|
if request is None:
|
||||||
|
logger.warning(
|
||||||
|
"Unknown GLM distributed denoiser result: %s", denoiser_request_id
|
||||||
|
)
|
||||||
|
return
|
||||||
|
if request.worker_idx is not None:
|
||||||
|
self._denoiser_free_slots[request.worker_idx] = 1
|
||||||
|
|
||||||
|
error = scalar_fields.get("error")
|
||||||
|
output = tensor_fields.get("output")
|
||||||
|
total = max(1, int(request.req.num_outputs_per_prompt or 1))
|
||||||
|
output_size = len(output) if output is not None else None
|
||||||
|
if output_size is not None and output_size != total:
|
||||||
|
error = (
|
||||||
|
f"GLM distributed output size mismatch: got {output_size}, "
|
||||||
|
f"expected {total}"
|
||||||
|
)
|
||||||
|
output = None
|
||||||
|
result = OutputBatch(
|
||||||
|
output=output,
|
||||||
|
error=error,
|
||||||
|
metrics=_deserialize_request_metrics(scalar_fields.get("metrics")),
|
||||||
|
metrics_list=[
|
||||||
|
_deserialize_request_metrics(metrics)
|
||||||
|
for metrics in scalar_fields.get("metrics_list", [])
|
||||||
|
]
|
||||||
|
or None,
|
||||||
|
peak_memory_mb=scalar_fields.get("peak_memory_mb", 0.0),
|
||||||
|
usage=scalar_fields.get("usage"),
|
||||||
|
)
|
||||||
|
with self._lock:
|
||||||
|
identity = self._pending.pop(request.client_request_id, None)
|
||||||
|
if identity is not None:
|
||||||
|
self._frontend.send_multipart([identity, b"", pickle.dumps(result)])
|
||||||
|
try:
|
||||||
|
self._tracker.transition(
|
||||||
|
request.client_request_id, RequestState.DENOISING_DONE
|
||||||
|
)
|
||||||
|
self._tracker.transition(
|
||||||
|
request.client_request_id,
|
||||||
|
RequestState.FAILED if error else RequestState.DONE,
|
||||||
|
error=error,
|
||||||
|
)
|
||||||
|
except ValueError:
|
||||||
|
pass
|
||||||
|
self._tracker.remove(request.client_request_id)
|
||||||
|
|
||||||
|
def _handle_decoder_result_frames(self, frames: list) -> None:
|
||||||
request_id = self._extract_request_id(frames)
|
request_id = self._extract_request_id(frames)
|
||||||
if request_id is None:
|
if request_id is None:
|
||||||
logger.warning("DiffusionServer: decoder result missing request_id")
|
logger.warning("DiffusionServer: decoder result missing request_id")
|
||||||
@@ -521,10 +881,6 @@ class DiffusionServer:
|
|||||||
return None
|
return None
|
||||||
|
|
||||||
def _complete_with_error(self, request_id: str, error_msg: str) -> None:
|
def _complete_with_error(self, request_id: str, error_msg: str) -> None:
|
||||||
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import (
|
|
||||||
OutputBatch,
|
|
||||||
)
|
|
||||||
|
|
||||||
logger.error("DiffusionServer: %s — %s", request_id, error_msg)
|
logger.error("DiffusionServer: %s — %s", request_id, error_msg)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
@@ -578,6 +934,25 @@ class DiffusionServer:
|
|||||||
self._decoder_tta = deque(
|
self._decoder_tta = deque(
|
||||||
e for e in self._decoder_tta if e.request_id not in timed_set
|
e for e in self._decoder_tta if e.request_id not in timed_set
|
||||||
)
|
)
|
||||||
|
if self._glm_distributed_state is not None:
|
||||||
|
state = self._glm_distributed_state
|
||||||
|
state.pending_ar_requests = deque(
|
||||||
|
request
|
||||||
|
for request in state.pending_ar_requests
|
||||||
|
if request.client_request_id not in timed_set
|
||||||
|
)
|
||||||
|
state.pending_denoiser_requests = deque(
|
||||||
|
request
|
||||||
|
for request in state.pending_denoiser_requests
|
||||||
|
if request.client_request_id not in timed_set
|
||||||
|
)
|
||||||
|
for denoiser_request_id, request in list(
|
||||||
|
state.denoiser_requests.items()
|
||||||
|
):
|
||||||
|
if request.client_request_id in timed_set:
|
||||||
|
state.denoiser_requests.pop(denoiser_request_id)
|
||||||
|
if request.worker_idx is not None:
|
||||||
|
self._denoiser_free_slots[request.worker_idx] = 1
|
||||||
|
|
||||||
def _free_slot_for_record(self, record) -> None:
|
def _free_slot_for_record(self, record) -> None:
|
||||||
if (
|
if (
|
||||||
@@ -667,6 +1042,8 @@ class DiffusionServer:
|
|||||||
prealloc = msg.get("preallocated_slots", [])
|
prealloc = msg.get("preallocated_slots", [])
|
||||||
info["free_preallocated_slots"] = list(prealloc)
|
info["free_preallocated_slots"] = list(prealloc)
|
||||||
peers[idx] = info
|
peers[idx] = info
|
||||||
|
if role == RoleType.DENOISER and self._glm_distributed_state is not None:
|
||||||
|
self._glm_distributed_state.denoiser_worker_available[idx] = True
|
||||||
|
|
||||||
logger.info(
|
logger.info(
|
||||||
"DiffusionServer transfer: registered %s[%d] work_endpoint=%s "
|
"DiffusionServer transfer: registered %s[%d] work_endpoint=%s "
|
||||||
@@ -1026,10 +1403,6 @@ class DiffusionServer:
|
|||||||
)
|
)
|
||||||
|
|
||||||
def _transfer_return_to_client_from_msg(self, request_id: str, msg: dict) -> None:
|
def _transfer_return_to_client_from_msg(self, request_id: str, msg: dict) -> None:
|
||||||
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import (
|
|
||||||
OutputBatch,
|
|
||||||
)
|
|
||||||
|
|
||||||
with self._lock:
|
with self._lock:
|
||||||
client_identity = self._pending.pop(request_id, None)
|
client_identity = self._pending.pop(request_id, None)
|
||||||
|
|
||||||
@@ -1054,7 +1427,7 @@ class DiffusionServer:
|
|||||||
def get_stats(self) -> dict:
|
def get_stats(self) -> dict:
|
||||||
with self._lock:
|
with self._lock:
|
||||||
pending_count = len(self._pending)
|
pending_count = len(self._pending)
|
||||||
return {
|
stats = {
|
||||||
"role": "diffusion_server",
|
"role": "diffusion_server",
|
||||||
"transfer_mode": self._transfer_mode,
|
"transfer_mode": self._transfer_mode,
|
||||||
"num_encoders": self._num_encoders,
|
"num_encoders": self._num_encoders,
|
||||||
@@ -1074,3 +1447,16 @@ class DiffusionServer:
|
|||||||
"decoder_peers": len(self._decoder_peers),
|
"decoder_peers": len(self._decoder_peers),
|
||||||
"tracker": self._tracker.snapshot(),
|
"tracker": self._tracker.snapshot(),
|
||||||
}
|
}
|
||||||
|
if self._glm_distributed_state is not None:
|
||||||
|
state = self._glm_distributed_state
|
||||||
|
stats.update(
|
||||||
|
{
|
||||||
|
"glm_ar_queue_depth": len(state.pending_ar_requests),
|
||||||
|
"glm_ar_in_flight": state.active_ar_batch is not None,
|
||||||
|
"glm_denoiser_queue_depth": len(state.pending_denoiser_requests),
|
||||||
|
"glm_denoiser_worker_available": list(
|
||||||
|
state.denoiser_worker_available
|
||||||
|
),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
return stats
|
||||||
|
|||||||
@@ -45,10 +45,14 @@ _VALID_TRANSITIONS: dict[RequestState, set[RequestState]] = {
|
|||||||
RequestState.DENOISING_RUNNING,
|
RequestState.DENOISING_RUNNING,
|
||||||
},
|
},
|
||||||
RequestState.DENOISING_WAITING: {RequestState.DENOISING_RUNNING},
|
RequestState.DENOISING_WAITING: {RequestState.DENOISING_RUNNING},
|
||||||
RequestState.DENOISING_RUNNING: {RequestState.DENOISING_DONE},
|
RequestState.DENOISING_RUNNING: {
|
||||||
|
RequestState.DENOISING_WAITING,
|
||||||
|
RequestState.DENOISING_DONE,
|
||||||
|
},
|
||||||
RequestState.DENOISING_DONE: {
|
RequestState.DENOISING_DONE: {
|
||||||
RequestState.DECODER_WAITING,
|
RequestState.DECODER_WAITING,
|
||||||
RequestState.DECODER_RUNNING,
|
RequestState.DECODER_RUNNING,
|
||||||
|
RequestState.DONE,
|
||||||
},
|
},
|
||||||
RequestState.DECODER_WAITING: {RequestState.DECODER_RUNNING},
|
RequestState.DECODER_WAITING: {RequestState.DECODER_RUNNING},
|
||||||
RequestState.DECODER_RUNNING: {RequestState.DONE},
|
RequestState.DECODER_RUNNING: {RequestState.DONE},
|
||||||
|
|||||||
@@ -30,6 +30,7 @@ from sglang.multimodal_gen.runtime.disaggregation.transport.buffer import (
|
|||||||
)
|
)
|
||||||
from sglang.multimodal_gen.runtime.disaggregation.transport.codec import (
|
from sglang.multimodal_gen.runtime.disaggregation.transport.codec import (
|
||||||
send_tensors,
|
send_tensors,
|
||||||
|
unpack_tensors,
|
||||||
)
|
)
|
||||||
from sglang.multimodal_gen.runtime.disaggregation.transport.engine import (
|
from sglang.multimodal_gen.runtime.disaggregation.transport.engine import (
|
||||||
create_transfer_engine,
|
create_transfer_engine,
|
||||||
@@ -49,6 +50,7 @@ from sglang.multimodal_gen.runtime.disaggregation.transport.protocol import (
|
|||||||
encode_transfer_msg,
|
encode_transfer_msg,
|
||||||
is_transfer_message,
|
is_transfer_message,
|
||||||
)
|
)
|
||||||
|
from sglang.multimodal_gen.runtime.entrypoints.utils import expand_request_outputs
|
||||||
from sglang.multimodal_gen.runtime.pipelines_core import Req
|
from sglang.multimodal_gen.runtime.pipelines_core import Req
|
||||||
from sglang.multimodal_gen.runtime.pipelines_core.diffusion_scheduler_utils import (
|
from sglang.multimodal_gen.runtime.pipelines_core.diffusion_scheduler_utils import (
|
||||||
clone_scheduler_runtime,
|
clone_scheduler_runtime,
|
||||||
@@ -65,6 +67,44 @@ if TYPE_CHECKING:
|
|||||||
|
|
||||||
logger = init_logger(__name__)
|
logger = init_logger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
def _advertised_pool_work_endpoint(server_args) -> str:
|
||||||
|
host = server_args.disagg_p2p_hostname or server_args.host or "127.0.0.1"
|
||||||
|
if host == "0.0.0.0":
|
||||||
|
host = server_args.disagg_p2p_hostname or "127.0.0.1"
|
||||||
|
return server_args.pool_work_endpoint.replace("0.0.0.0", host)
|
||||||
|
|
||||||
|
|
||||||
|
def _expand_glm_distributed_outputs(req: Req) -> list[Req]:
|
||||||
|
"""Split external-AR tokens into requests for sequential denoising."""
|
||||||
|
output_count = max(1, int(req.num_outputs_per_prompt or 1))
|
||||||
|
if output_count == 1:
|
||||||
|
return [req]
|
||||||
|
|
||||||
|
prior_token_ids = req.prior_token_id
|
||||||
|
if not isinstance(prior_token_ids, torch.Tensor) or (
|
||||||
|
prior_token_ids.shape[0] != output_count
|
||||||
|
):
|
||||||
|
actual_count = (
|
||||||
|
prior_token_ids.shape[0]
|
||||||
|
if isinstance(prior_token_ids, torch.Tensor)
|
||||||
|
else type(prior_token_ids).__name__
|
||||||
|
)
|
||||||
|
raise RuntimeError(
|
||||||
|
"Cannot split GLM-Image AR output for distributed inference: "
|
||||||
|
f"expected {output_count} token rows, got {actual_count}."
|
||||||
|
)
|
||||||
|
|
||||||
|
usage_by_output = req.extra.get("usage_by_output")
|
||||||
|
output_reqs = expand_request_outputs(req)
|
||||||
|
for output_index, output_req in enumerate(output_reqs):
|
||||||
|
output_req.prior_token_id = prior_token_ids[output_index : output_index + 1]
|
||||||
|
output_req.extra.pop("usage_by_output", None)
|
||||||
|
if usage_by_output is not None and output_index < len(usage_by_output):
|
||||||
|
output_req.usage = usage_by_output[output_index]
|
||||||
|
return output_reqs
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
# Field extraction: split Req into tensors (transfer buffer) and scalars (JSON)
|
# Field extraction: split Req into tensors (transfer buffer) and scalars (JSON)
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
@@ -116,6 +156,9 @@ _SAMPLING_PARAMS_EXCLUDE_FIELDS = frozenset(
|
|||||||
}
|
}
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Receivers reconstruct base SamplingParams, so only base defaults can be omitted.
|
||||||
|
_BASE_SAMPLING_PARAM_FIELDS = {f.name: f for f in dataclasses.fields(SamplingParams)}
|
||||||
|
|
||||||
|
|
||||||
def _is_tensor_like(value) -> bool:
|
def _is_tensor_like(value) -> bool:
|
||||||
if isinstance(value, torch.Tensor):
|
if isinstance(value, torch.Tensor):
|
||||||
@@ -278,7 +321,8 @@ def extract_transfer_fields(req) -> tuple[dict, dict]:
|
|||||||
value = getattr(sp, name, None)
|
value = getattr(sp, name, None)
|
||||||
if value is None:
|
if value is None:
|
||||||
continue
|
continue
|
||||||
if _is_default(value, f):
|
base_field = _BASE_SAMPLING_PARAM_FIELDS.get(name)
|
||||||
|
if base_field is not None and _is_default(value, base_field):
|
||||||
continue
|
continue
|
||||||
try:
|
try:
|
||||||
scalar_fields[name] = _to_json_serializable(value)
|
scalar_fields[name] = _to_json_serializable(value)
|
||||||
@@ -458,6 +502,23 @@ class SchedulerDisaggMixin:
|
|||||||
|
|
||||||
sa = self.server_args
|
sa = self.server_args
|
||||||
|
|
||||||
|
if self._is_glm_distributed_mode():
|
||||||
|
self._preallocated_slots = {}
|
||||||
|
register_msg = TransferRegisterMsg(
|
||||||
|
role=self._disagg_role.value,
|
||||||
|
work_endpoint=_advertised_pool_work_endpoint(sa),
|
||||||
|
)
|
||||||
|
self._pool_result_push.send_multipart(encode_transfer_msg(register_msg))
|
||||||
|
self._compute_ready_queue = queue.Queue(maxsize=4)
|
||||||
|
self._recv_prefetch_thread = threading.Thread(
|
||||||
|
target=self._recv_prefetch_loop,
|
||||||
|
daemon=True,
|
||||||
|
name="recv-prefetch-glm-distributed-denoiser",
|
||||||
|
)
|
||||||
|
self._recv_prefetch_thread.start()
|
||||||
|
logger.info("GLM distributed denoiser registered")
|
||||||
|
return
|
||||||
|
|
||||||
# Pool size: configurable, default 256 MiB
|
# Pool size: configurable, default 256 MiB
|
||||||
pool_size = getattr(sa, "disagg_transfer_pool_size", 256 * 1024 * 1024)
|
pool_size = getattr(sa, "disagg_transfer_pool_size", 256 * 1024 * 1024)
|
||||||
|
|
||||||
@@ -522,7 +583,7 @@ class SchedulerDisaggMixin:
|
|||||||
session_id=self._transfer_manager.session_id,
|
session_id=self._transfer_manager.session_id,
|
||||||
pool_ptr=self._transfer_manager.pool_data_ptr,
|
pool_ptr=self._transfer_manager.pool_data_ptr,
|
||||||
pool_size=self._transfer_manager.pool_size,
|
pool_size=self._transfer_manager.pool_size,
|
||||||
work_endpoint=sa.pool_work_endpoint,
|
work_endpoint=_advertised_pool_work_endpoint(sa),
|
||||||
preallocated_slots=preallocated_slot_info,
|
preallocated_slots=preallocated_slot_info,
|
||||||
)
|
)
|
||||||
self._pool_result_push.send_multipart(encode_transfer_msg(register_msg))
|
self._pool_result_push.send_multipart(encode_transfer_msg(register_msg))
|
||||||
@@ -623,6 +684,10 @@ class SchedulerDisaggMixin:
|
|||||||
raw_frames = self._pool_work_pull.recv_multipart()
|
raw_frames = self._pool_work_pull.recv_multipart()
|
||||||
frames = [bytes(f) for f in raw_frames]
|
frames = [bytes(f) for f in raw_frames]
|
||||||
|
|
||||||
|
if not is_transfer_message(frames):
|
||||||
|
self._compute_ready_queue.put(("relay_compute", frames))
|
||||||
|
continue
|
||||||
|
|
||||||
msg = decode_transfer_msg(frames)
|
msg = decode_transfer_msg(frames)
|
||||||
msg_type = msg.get("msg_type", "")
|
msg_type = msg.get("msg_type", "")
|
||||||
|
|
||||||
@@ -903,6 +968,18 @@ class SchedulerDisaggMixin:
|
|||||||
self._broadcast_to_all_ranks(("skip",))
|
self._broadcast_to_all_ranks(("skip",))
|
||||||
self._handle_transfer_msg(data)
|
self._handle_transfer_msg(data)
|
||||||
|
|
||||||
|
elif msg_type == "relay_compute":
|
||||||
|
local_device = (
|
||||||
|
f"{current_platform.device_type}:{self.worker.local_rank}"
|
||||||
|
)
|
||||||
|
tensors, scalar_fields = unpack_tensors(data, device=local_device)
|
||||||
|
request_id = scalar_fields.get("request_id", "unknown")
|
||||||
|
req = self._build_disagg_req(scalar_fields, tensors)
|
||||||
|
if is_multi_rank:
|
||||||
|
self._broadcast_to_all_ranks(("compute",))
|
||||||
|
self._broadcast_req_to_all_ranks(req)
|
||||||
|
self._execute_glm_distributed_denoiser_request(req, request_id)
|
||||||
|
|
||||||
self._consecutive_error_count = 0
|
self._consecutive_error_count = 0
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
@@ -1298,11 +1375,16 @@ class SchedulerDisaggMixin:
|
|||||||
(:meth:`_disagg_non_rank0_event_loop`).
|
(:meth:`_disagg_non_rank0_event_loop`).
|
||||||
"""
|
"""
|
||||||
if self._disagg_role == RoleType.DENOISER:
|
if self._disagg_role == RoleType.DENOISER:
|
||||||
# Initialize scheduler timesteps (same as rank 0)
|
if not self._is_glm_distributed_mode():
|
||||||
_init_disagg_request_scheduler(self, req)
|
_init_disagg_request_scheduler(self, req)
|
||||||
|
|
||||||
with self._disagg_trace_dispatch(req):
|
with self._disagg_trace_dispatch(req):
|
||||||
self.worker.execute_forward([req], return_req=True)
|
if self._is_glm_distributed_mode():
|
||||||
|
req.save_output = False
|
||||||
|
req.return_file_paths_only = False
|
||||||
|
self.worker.execute_forward([req])
|
||||||
|
else:
|
||||||
|
self.worker.execute_forward([req], return_req=True)
|
||||||
|
|
||||||
elif self._disagg_role == RoleType.DECODER:
|
elif self._disagg_role == RoleType.DECODER:
|
||||||
req.save_output = False
|
req.save_output = False
|
||||||
@@ -1456,6 +1538,64 @@ class SchedulerDisaggMixin:
|
|||||||
duration_s,
|
duration_s,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
def _is_glm_distributed_mode(self: Scheduler) -> bool:
|
||||||
|
"""Return whether this scheduler is a GLM distributed denoiser."""
|
||||||
|
return (
|
||||||
|
self._disagg_role == RoleType.DENOISER
|
||||||
|
and self.worker.pipeline.pipeline_name == "GlmImagePipeline"
|
||||||
|
and self.server_args.srt_encoder_url is not None
|
||||||
|
)
|
||||||
|
|
||||||
|
def _execute_glm_distributed_denoiser_request(
|
||||||
|
self: Scheduler, req: Req, request_id: str
|
||||||
|
) -> None:
|
||||||
|
"""Run local preparation, DiT, and VAE, then return decoded pixels."""
|
||||||
|
req.save_output = False
|
||||||
|
start_time = time.monotonic()
|
||||||
|
with self._disagg_trace_dispatch(req):
|
||||||
|
if (
|
||||||
|
self.server_args.pipeline_config.supports_sequential_multi_output_inference()
|
||||||
|
and max(1, int(req.num_outputs_per_prompt or 1)) > 1
|
||||||
|
):
|
||||||
|
output_reqs = _expand_glm_distributed_outputs(req)
|
||||||
|
output_batches = list(
|
||||||
|
self.worker.execute_forward_sequentially(output_reqs)
|
||||||
|
)
|
||||||
|
output_batch = self.worker._merge_expanded_output_batches(
|
||||||
|
output_batches
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
output_batch = self.worker.execute_forward([req])
|
||||||
|
|
||||||
|
tensor_fields = {}
|
||||||
|
scalar_fields = {"request_id": request_id}
|
||||||
|
if output_batch.output is not None:
|
||||||
|
tensor_fields["output"] = output_batch.output
|
||||||
|
if output_batch.error is not None:
|
||||||
|
scalar_fields["error"] = output_batch.error
|
||||||
|
if output_batch.usage is not None:
|
||||||
|
scalar_fields["usage"] = output_batch.usage
|
||||||
|
if output_batch.metrics is not None:
|
||||||
|
scalar_fields["metrics"] = output_batch.metrics.to_dict()
|
||||||
|
if output_batch.metrics_list is not None:
|
||||||
|
scalar_fields["metrics_list"] = [
|
||||||
|
metrics.to_dict() if metrics is not None else None
|
||||||
|
for metrics in output_batch.metrics_list
|
||||||
|
]
|
||||||
|
scalar_fields["peak_memory_mb"] = output_batch.peak_memory_mb
|
||||||
|
send_tensors(self._pool_result_push, tensor_fields, scalar_fields)
|
||||||
|
|
||||||
|
if self._disagg_metrics:
|
||||||
|
if output_batch.error:
|
||||||
|
self._disagg_metrics.record_request_failed(request_id)
|
||||||
|
else:
|
||||||
|
self._disagg_metrics.record_request_complete(request_id)
|
||||||
|
logger.debug(
|
||||||
|
"GLM distributed denoiser: processed %s in %.2f s",
|
||||||
|
request_id,
|
||||||
|
time.monotonic() - start_time,
|
||||||
|
)
|
||||||
|
|
||||||
def _disagg_decoder_compute(self: Scheduler, req: Req, request_id: str) -> None:
|
def _disagg_decoder_compute(self: Scheduler, req: Req, request_id: str) -> None:
|
||||||
"""Run decoder compute in transfer mode, send result to DS.
|
"""Run decoder compute in transfer mode, send result to DS.
|
||||||
|
|
||||||
|
|||||||
@@ -3,12 +3,9 @@
|
|||||||
import dataclasses
|
import dataclasses
|
||||||
import multiprocessing as mp
|
import multiprocessing as mp
|
||||||
import os
|
import os
|
||||||
import signal
|
|
||||||
import sys
|
import sys
|
||||||
import threading
|
|
||||||
import time
|
import time
|
||||||
|
|
||||||
import psutil
|
|
||||||
import uvicorn
|
import uvicorn
|
||||||
|
|
||||||
from sglang.multimodal_gen.runtime.disaggregation.orchestrator import (
|
from sglang.multimodal_gen.runtime.disaggregation.orchestrator import (
|
||||||
@@ -24,7 +21,10 @@ from sglang.multimodal_gen.runtime.server_args import (
|
|||||||
prepare_server_args,
|
prepare_server_args,
|
||||||
set_global_server_args,
|
set_global_server_args,
|
||||||
)
|
)
|
||||||
from sglang.multimodal_gen.runtime.utils.common import is_port_available
|
from sglang.multimodal_gen.runtime.utils.common import (
|
||||||
|
is_port_available,
|
||||||
|
kill_process_tree,
|
||||||
|
)
|
||||||
from sglang.multimodal_gen.runtime.utils.logging_utils import configure_logger, logger
|
from sglang.multimodal_gen.runtime.utils.logging_utils import configure_logger, logger
|
||||||
from sglang.multimodal_gen.runtime.utils.trace_wrapper import init_diffusion_tracing
|
from sglang.multimodal_gen.runtime.utils.trace_wrapper import init_diffusion_tracing
|
||||||
from sglang.multimodal_gen.utils import kill_itself_when_parent_died
|
from sglang.multimodal_gen.utils import kill_itself_when_parent_died
|
||||||
@@ -53,45 +53,6 @@ def _find_available_port(
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
def kill_process_tree(parent_pid, include_parent: bool = True, skip_pid: int = None):
|
|
||||||
"""Kill the process and all its child processes."""
|
|
||||||
# Remove sigchld handler to avoid spammy logs.
|
|
||||||
if threading.current_thread() is threading.main_thread():
|
|
||||||
signal.signal(signal.SIGCHLD, signal.SIG_DFL)
|
|
||||||
|
|
||||||
if parent_pid is None:
|
|
||||||
parent_pid = os.getpid()
|
|
||||||
include_parent = False
|
|
||||||
|
|
||||||
try:
|
|
||||||
itself = psutil.Process(parent_pid)
|
|
||||||
except psutil.NoSuchProcess:
|
|
||||||
return
|
|
||||||
|
|
||||||
children = itself.children(recursive=True)
|
|
||||||
for child in children:
|
|
||||||
if child.pid == skip_pid:
|
|
||||||
continue
|
|
||||||
try:
|
|
||||||
child.kill()
|
|
||||||
except psutil.NoSuchProcess:
|
|
||||||
pass
|
|
||||||
|
|
||||||
if include_parent:
|
|
||||||
try:
|
|
||||||
if parent_pid == os.getpid():
|
|
||||||
itself.kill()
|
|
||||||
sys.exit(0)
|
|
||||||
|
|
||||||
itself.kill()
|
|
||||||
|
|
||||||
# Sometime processes cannot be killed with SIGKILL (e.g, PID=1 launched by kubernetes),
|
|
||||||
# so we send an additional signal to kill them.
|
|
||||||
itself.send_signal(signal.SIGQUIT)
|
|
||||||
except psutil.NoSuchProcess:
|
|
||||||
pass
|
|
||||||
|
|
||||||
|
|
||||||
def _process_names(processes) -> str:
|
def _process_names(processes) -> str:
|
||||||
return ", ".join(getattr(p, "name", repr(p)) for p in processes)
|
return ", ".join(getattr(p, "name", repr(p)) for p in processes)
|
||||||
|
|
||||||
@@ -585,12 +546,12 @@ def launch_http_server_only(server_args):
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
def parse_url_string(url_str: str) -> list[str]:
|
def parse_url_string(url_str: str | None) -> list[str]:
|
||||||
"""Parse a semicolon-separated URL string into a list.
|
"""Parse a semicolon-separated URL string into a list.
|
||||||
|
|
||||||
Example: "tcp://10.0.0.1:35000;tcp://10.0.0.2:35000" -> ["tcp://...", "tcp://..."]
|
Example: "tcp://10.0.0.1:35000;tcp://10.0.0.2:35000" -> ["tcp://...", "tcp://..."]
|
||||||
"""
|
"""
|
||||||
return [u.strip() for u in url_str.split(";") if u.strip()]
|
return [u.strip() for u in (url_str or "").split(";") if u.strip()]
|
||||||
|
|
||||||
|
|
||||||
def launch_disagg_server(server_args: ServerArgs):
|
def launch_disagg_server(server_args: ServerArgs):
|
||||||
@@ -605,12 +566,23 @@ def launch_disagg_server(server_args: ServerArgs):
|
|||||||
decoder result: scheduler_port + 3
|
decoder result: scheduler_port + 3
|
||||||
"""
|
"""
|
||||||
configure_logger(server_args)
|
configure_logger(server_args)
|
||||||
|
set_global_server_args(server_args)
|
||||||
|
|
||||||
for name, val in [
|
glm_distributed_mode_enabled = (
|
||||||
("--encoder-urls", server_args.encoder_urls),
|
type(server_args.pipeline_config).__name__ == "GlmImagePipelineConfig"
|
||||||
("--denoiser-urls", server_args.denoiser_urls),
|
and server_args.srt_encoder_url is not None
|
||||||
("--decoder-urls", server_args.decoder_urls),
|
and server_args.encoder_urls is None
|
||||||
]:
|
and server_args.decoder_urls is None
|
||||||
|
)
|
||||||
|
required_urls = [("--denoiser-urls", server_args.denoiser_urls)]
|
||||||
|
if not glm_distributed_mode_enabled:
|
||||||
|
required_urls.extend(
|
||||||
|
[
|
||||||
|
("--encoder-urls", server_args.encoder_urls),
|
||||||
|
("--decoder-urls", server_args.decoder_urls),
|
||||||
|
]
|
||||||
|
)
|
||||||
|
for name, val in required_urls:
|
||||||
if val is None:
|
if val is None:
|
||||||
raise ValueError(f"{name} is required for --disagg-role server")
|
raise ValueError(f"{name} is required for --disagg-role server")
|
||||||
|
|
||||||
@@ -644,6 +616,9 @@ def launch_disagg_server(server_args: ServerArgs):
|
|||||||
decoder_result_ep,
|
decoder_result_ep,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
denoiser_options = (
|
||||||
|
{"denoiser_capacity_per_worker": 1} if glm_distributed_mode_enabled else {}
|
||||||
|
)
|
||||||
diffusion_server = DiffusionServer(
|
diffusion_server = DiffusionServer(
|
||||||
frontend_endpoint=frontend_endpoint,
|
frontend_endpoint=frontend_endpoint,
|
||||||
encoder_work_endpoints=encoder_work_endpoints,
|
encoder_work_endpoints=encoder_work_endpoints,
|
||||||
@@ -654,6 +629,9 @@ def launch_disagg_server(server_args: ServerArgs):
|
|||||||
decoder_result_endpoint=decoder_result_ep,
|
decoder_result_endpoint=decoder_result_ep,
|
||||||
dispatch_policy_name=server_args.disagg_dispatch_policy,
|
dispatch_policy_name=server_args.disagg_dispatch_policy,
|
||||||
timeout_s=float(server_args.disagg_timeout),
|
timeout_s=float(server_args.disagg_timeout),
|
||||||
|
server_args=server_args,
|
||||||
|
glm_distributed_mode_enabled=glm_distributed_mode_enabled,
|
||||||
|
**denoiser_options,
|
||||||
)
|
)
|
||||||
diffusion_server.start()
|
diffusion_server.start()
|
||||||
|
|
||||||
@@ -726,6 +704,44 @@ def launch_disagg_role(server_args: ServerArgs):
|
|||||||
"ulysses_degree": role_par["ulysses_degree"],
|
"ulysses_degree": role_par["ulysses_degree"],
|
||||||
"ring_degree": role_par["ring_degree"],
|
"ring_degree": role_par["ring_degree"],
|
||||||
}
|
}
|
||||||
|
role_tp = role_par["tp_size"] or 1
|
||||||
|
role_sp = role_par["sp_degree"] or 1
|
||||||
|
cfg_degree = (
|
||||||
|
server_args.cfg_parallel_degree if server_args.enable_cfg_parallel else 1
|
||||||
|
)
|
||||||
|
cfg_parallel_explicit = server_args.is_arg_explicitly_set(
|
||||||
|
"enable_cfg_parallel"
|
||||||
|
) or server_args.is_arg_explicitly_set("cfg_parallel_degree")
|
||||||
|
required_devices = role_tp * role_sp * cfg_degree * server_args.dp_size
|
||||||
|
if not cfg_parallel_explicit and (
|
||||||
|
required_devices > server_args.num_gpus
|
||||||
|
or server_args.num_gpus % required_devices != 0
|
||||||
|
):
|
||||||
|
logger.warning(
|
||||||
|
"Disabling auto-enabled CFG parallel for %s role because tp=%d, "
|
||||||
|
"sp=%d, cfg=%d, dp=%d is incompatible with %d devices",
|
||||||
|
role_type.value,
|
||||||
|
role_tp,
|
||||||
|
role_sp,
|
||||||
|
cfg_degree,
|
||||||
|
server_args.dp_size,
|
||||||
|
server_args.num_gpus,
|
||||||
|
)
|
||||||
|
role_overrides["enable_cfg_parallel"] = False
|
||||||
|
role_overrides["cfg_parallel_degree"] = 1
|
||||||
|
cfg_degree = 1
|
||||||
|
required_devices = role_tp * role_sp * server_args.dp_size
|
||||||
|
|
||||||
|
if (
|
||||||
|
required_devices > server_args.num_gpus
|
||||||
|
or server_args.num_gpus % required_devices != 0
|
||||||
|
):
|
||||||
|
raise ValueError(
|
||||||
|
f"Invalid parallelism for {role_type.value} role: "
|
||||||
|
f"tp={role_tp}, sp={role_sp}, cfg={cfg_degree}, "
|
||||||
|
f"dp={server_args.dp_size} requires groups of {required_devices} "
|
||||||
|
f"devices, but num_gpus={server_args.num_gpus}"
|
||||||
|
)
|
||||||
|
|
||||||
base_dict = {
|
base_dict = {
|
||||||
f.name: getattr(server_args, f.name)
|
f.name: getattr(server_args, f.name)
|
||||||
|
|||||||
@@ -12,6 +12,22 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.g
|
|||||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||||
|
|
||||||
|
|
||||||
|
class GlmImageDenoiserDecodingStage(GlmImageDecodingStage):
|
||||||
|
"""Run VAE decoding on the denoiser because this topology has no decoder worker."""
|
||||||
|
|
||||||
|
@property
|
||||||
|
def role_affinity(self) -> RoleType:
|
||||||
|
return RoleType.DENOISER
|
||||||
|
|
||||||
|
|
||||||
|
class GlmImageDenoiserPreparationStage(GlmImageBeforeDenoisingStage):
|
||||||
|
"""Run DiT preparation on the denoiser because AR runs in the server head."""
|
||||||
|
|
||||||
|
@property
|
||||||
|
def role_affinity(self) -> RoleType:
|
||||||
|
return RoleType.DENOISER
|
||||||
|
|
||||||
|
|
||||||
class GlmImagePipeline(LoRAPipeline, ComposedPipelineBase):
|
class GlmImagePipeline(LoRAPipeline, ComposedPipelineBase):
|
||||||
pipeline_name = "GlmImagePipeline"
|
pipeline_name = "GlmImagePipeline"
|
||||||
|
|
||||||
@@ -26,6 +42,10 @@ class GlmImagePipeline(LoRAPipeline, ComposedPipelineBase):
|
|||||||
]
|
]
|
||||||
|
|
||||||
def create_pipeline_stages(self, server_args: ServerArgs):
|
def create_pipeline_stages(self, server_args: ServerArgs):
|
||||||
|
is_glm_distributed_mode = (
|
||||||
|
self._disagg_role == RoleType.DENOISER
|
||||||
|
and server_args.srt_encoder_url is not None
|
||||||
|
)
|
||||||
self.add_stage(
|
self.add_stage(
|
||||||
GlmImageAR(
|
GlmImageAR(
|
||||||
processor=self.get_module("processor"),
|
processor=self.get_module("processor"),
|
||||||
@@ -34,8 +54,13 @@ class GlmImagePipeline(LoRAPipeline, ComposedPipelineBase):
|
|||||||
"glm_image_ar",
|
"glm_image_ar",
|
||||||
)
|
)
|
||||||
|
|
||||||
|
before_denoising_stage_cls = (
|
||||||
|
GlmImageDenoiserPreparationStage
|
||||||
|
if is_glm_distributed_mode
|
||||||
|
else GlmImageBeforeDenoisingStage
|
||||||
|
)
|
||||||
self.add_stage(
|
self.add_stage(
|
||||||
GlmImageBeforeDenoisingStage(
|
before_denoising_stage_cls(
|
||||||
vae=self.get_module("vae"),
|
vae=self.get_module("vae"),
|
||||||
text_encoder=self.get_module("text_encoder"),
|
text_encoder=self.get_module("text_encoder"),
|
||||||
tokenizer=self.get_module("tokenizer"),
|
tokenizer=self.get_module("tokenizer"),
|
||||||
@@ -52,14 +77,22 @@ class GlmImagePipeline(LoRAPipeline, ComposedPipelineBase):
|
|||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
self.add_stage_factory(
|
if is_glm_distributed_mode:
|
||||||
RoleType.DECODER,
|
self.add_stage(
|
||||||
lambda: GlmImageDecodingStage(
|
GlmImageDenoiserDecodingStage(
|
||||||
vae=self.get_module("vae"),
|
vae=self.get_module("vae"), pipeline=self
|
||||||
pipeline=self,
|
),
|
||||||
),
|
"decoding_stage",
|
||||||
"decoding_stage",
|
)
|
||||||
)
|
else:
|
||||||
|
self.add_stage_factory(
|
||||||
|
RoleType.DECODER,
|
||||||
|
lambda: GlmImageDecodingStage(
|
||||||
|
vae=self.get_module("vae"),
|
||||||
|
pipeline=self,
|
||||||
|
),
|
||||||
|
"decoding_stage",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
EntryClass = [GlmImagePipeline]
|
EntryClass = [GlmImagePipeline]
|
||||||
|
|||||||
@@ -270,6 +270,12 @@ class ComposedPipelineBase(ABC):
|
|||||||
extra_allowed_modules = set(
|
extra_allowed_modules = set(
|
||||||
role_to_pipeline_modules.get(role, {}).get(self.pipeline_name, set())
|
role_to_pipeline_modules.get(role, {}).get(self.pipeline_name, set())
|
||||||
)
|
)
|
||||||
|
if (
|
||||||
|
role == RoleType.DENOISER
|
||||||
|
and self.pipeline_name == "GlmImagePipeline"
|
||||||
|
and getattr(self.server_args, "srt_encoder_url", None) is not None
|
||||||
|
):
|
||||||
|
extra_allowed_modules.update({"text_encoder", "tokenizer", "vae"})
|
||||||
|
|
||||||
if role == RoleType.DENOISER and task_name == "ti2v":
|
if role == RoleType.DENOISER and task_name == "ti2v":
|
||||||
if self.pipeline_name in {
|
if self.pipeline_name in {
|
||||||
|
|||||||
@@ -108,6 +108,10 @@ class Req:
|
|||||||
pooled_embeds: list[torch.Tensor] = field(default_factory=list)
|
pooled_embeds: list[torch.Tensor] = field(default_factory=list)
|
||||||
neg_pooled_embeds: list[torch.Tensor] = field(default_factory=list)
|
neg_pooled_embeds: list[torch.Tensor] = field(default_factory=list)
|
||||||
|
|
||||||
|
# GLM-Image autoregressive prior tokens
|
||||||
|
prior_token_id: torch.Tensor | None = None
|
||||||
|
prior_token_image_ids: torch.Tensor | list[torch.Tensor] | None = None
|
||||||
|
|
||||||
# Additional text-related parameters
|
# Additional text-related parameters
|
||||||
max_sequence_length: int | None = None
|
max_sequence_length: int | None = None
|
||||||
prompt_template: dict[str, Any] | None = None
|
prompt_template: dict[str, Any] | None = None
|
||||||
|
|||||||
+33
-20
@@ -15,7 +15,6 @@ from sglang.multimodal_gen.configs.sample.glmimage import (
|
|||||||
GLM_IMAGE_RESOLUTION_ALIGNMENT,
|
GLM_IMAGE_RESOLUTION_ALIGNMENT,
|
||||||
align_glm_image_resolution,
|
align_glm_image_resolution,
|
||||||
)
|
)
|
||||||
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
|
|
||||||
from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
|
from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
|
||||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import (
|
from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import (
|
||||||
ComponentUse,
|
ComponentUse,
|
||||||
@@ -27,6 +26,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.base import (
|
|||||||
StageParallelismType,
|
StageParallelismType,
|
||||||
)
|
)
|
||||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding import DecodingStage
|
from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding import DecodingStage
|
||||||
|
from sglang.multimodal_gen.runtime.platforms import current_platform
|
||||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||||
from sglang.multimodal_gen.runtime.utils.precision import (
|
from sglang.multimodal_gen.runtime.utils.precision import (
|
||||||
@@ -414,7 +414,7 @@ class GlmImageAR(PipelineStage):
|
|||||||
Tuple of the D16 prior token IDs, optional source-image token IDs,
|
Tuple of the D16 prior token IDs, optional source-image token IDs,
|
||||||
and optional usage statistics returned by an external AR server.
|
and optional usage statistics returned by an external AR server.
|
||||||
"""
|
"""
|
||||||
device = get_local_torch_device()
|
device = current_platform.get_local_torch_device()
|
||||||
_validate_glm_image_resolution_alignment(width, height)
|
_validate_glm_image_resolution_alignment(width, height)
|
||||||
|
|
||||||
is_text_to_image = image is None or len(image) == 0
|
is_text_to_image = image is None or len(image) == 0
|
||||||
@@ -516,8 +516,9 @@ class GlmImageAR(PipelineStage):
|
|||||||
height: int,
|
height: int,
|
||||||
width: int,
|
width: int,
|
||||||
server_args: ServerArgs,
|
server_args: ServerArgs,
|
||||||
|
device: Optional[torch.device] = None,
|
||||||
) -> tuple[list[torch.Tensor], list[dict[str, int] | None]]:
|
) -> tuple[list[torch.Tensor], list[dict[str, int] | None]]:
|
||||||
device = get_local_torch_device()
|
device = device or current_platform.get_local_torch_device()
|
||||||
_validate_glm_image_resolution_alignment(width, height)
|
_validate_glm_image_resolution_alignment(width, height)
|
||||||
|
|
||||||
input_ids = []
|
input_ids = []
|
||||||
@@ -577,27 +578,15 @@ class GlmImageAR(PipelineStage):
|
|||||||
usages.append(_extract_srt_usage(item.get("meta_info")))
|
usages.append(_extract_srt_usage(item.get("meta_info")))
|
||||||
return prior_token_ids, usages
|
return prior_token_ids, usages
|
||||||
|
|
||||||
def run_grouped_requests(
|
def generate_and_assign_prior_tokens(
|
||||||
self,
|
self,
|
||||||
batches: list[Req],
|
batches: list[Req],
|
||||||
server_args: ServerArgs,
|
server_args: ServerArgs,
|
||||||
|
device: Optional[torch.device] = None,
|
||||||
) -> list[Req]:
|
) -> list[Req]:
|
||||||
can_batch_ar = (
|
"""Generate one AR batch and assign its tokens and usage to each request."""
|
||||||
len(batches) > 1
|
|
||||||
and server_args.srt_encoder_url is not None
|
|
||||||
and all(
|
|
||||||
isinstance(batch.prompt, str) and batch.image_path is None
|
|
||||||
for batch in batches
|
|
||||||
)
|
|
||||||
)
|
|
||||||
if not can_batch_ar:
|
|
||||||
return super().run_grouped_requests(batches, server_args)
|
|
||||||
|
|
||||||
height = batches[0].height
|
height = batches[0].height
|
||||||
width = batches[0].width
|
width = batches[0].width
|
||||||
if any(batch.height != height or batch.width != width for batch in batches[1:]):
|
|
||||||
return super().run_grouped_requests(batches, server_args)
|
|
||||||
|
|
||||||
start_time = time.time()
|
start_time = time.time()
|
||||||
output_counts = [_num_outputs_per_prompt(batch) for batch in batches]
|
output_counts = [_num_outputs_per_prompt(batch) for batch in batches]
|
||||||
prompts = [
|
prompts = [
|
||||||
@@ -616,6 +605,7 @@ class GlmImageAR(PipelineStage):
|
|||||||
height=height,
|
height=height,
|
||||||
width=width,
|
width=width,
|
||||||
server_args=server_args,
|
server_args=server_args,
|
||||||
|
device=device,
|
||||||
)
|
)
|
||||||
duration = time.time() - start_time
|
duration = time.time() - start_time
|
||||||
logger.info(
|
logger.info(
|
||||||
@@ -642,6 +632,29 @@ class GlmImageAR(PipelineStage):
|
|||||||
output_offset += output_count
|
output_offset += output_count
|
||||||
return batches
|
return batches
|
||||||
|
|
||||||
|
def run_grouped_requests(
|
||||||
|
self,
|
||||||
|
batches: list[Req],
|
||||||
|
server_args: ServerArgs,
|
||||||
|
) -> list[Req]:
|
||||||
|
can_batch_ar = (
|
||||||
|
len(batches) > 1
|
||||||
|
and server_args.srt_encoder_url is not None
|
||||||
|
and all(
|
||||||
|
isinstance(batch.prompt, str) and batch.image_path is None
|
||||||
|
for batch in batches
|
||||||
|
)
|
||||||
|
)
|
||||||
|
if not can_batch_ar:
|
||||||
|
return super().run_grouped_requests(batches, server_args)
|
||||||
|
|
||||||
|
height = batches[0].height
|
||||||
|
width = batches[0].width
|
||||||
|
if any(batch.height != height or batch.width != width for batch in batches[1:]):
|
||||||
|
return super().run_grouped_requests(batches, server_args)
|
||||||
|
|
||||||
|
return self.generate_and_assign_prior_tokens(batches, server_args)
|
||||||
|
|
||||||
def iter_sequential_requests(
|
def iter_sequential_requests(
|
||||||
self, batch: Req, server_args: ServerArgs
|
self, batch: Req, server_args: ServerArgs
|
||||||
) -> Iterator[Req]:
|
) -> Iterator[Req]:
|
||||||
@@ -714,7 +727,7 @@ class GlmImageAR(PipelineStage):
|
|||||||
else:
|
else:
|
||||||
ar_condition_images = None
|
ar_condition_images = None
|
||||||
|
|
||||||
device = get_local_torch_device()
|
device = current_platform.get_local_torch_device()
|
||||||
|
|
||||||
if ar_condition_images is not None:
|
if ar_condition_images is not None:
|
||||||
height = height or ar_condition_images[0].height
|
height = height or ar_condition_images[0].height
|
||||||
@@ -1158,7 +1171,7 @@ class GlmImageBeforeDenoisingStage(PipelineStage):
|
|||||||
height = batch.height
|
height = batch.height
|
||||||
width = batch.width
|
width = batch.width
|
||||||
|
|
||||||
device = get_local_torch_device()
|
device = current_platform.get_local_torch_device()
|
||||||
batch_size = _num_outputs_per_prompt(batch)
|
batch_size = _num_outputs_per_prompt(batch)
|
||||||
max_sequence_length = 1024
|
max_sequence_length = 1024
|
||||||
seed = getattr(batch, "seed", None)
|
seed = getattr(batch, "seed", None)
|
||||||
|
|||||||
@@ -39,6 +39,7 @@ if current_platform.is_npu():
|
|||||||
DEFAULT_STANDALONE_EST_TIME_SECONDS,
|
DEFAULT_STANDALONE_EST_TIME_SECONDS,
|
||||||
FILE_SUITES,
|
FILE_SUITES,
|
||||||
PARAMETRIZED_CASE_GROUPS,
|
PARAMETRIZED_CASE_GROUPS,
|
||||||
|
STANDALONE_FILE_EST_TIMES,
|
||||||
STANDALONE_FILES,
|
STANDALONE_FILES,
|
||||||
STARTUP_OVERHEAD_SECONDS,
|
STARTUP_OVERHEAD_SECONDS,
|
||||||
SUITES,
|
SUITES,
|
||||||
|
|||||||
@@ -0,0 +1,315 @@
|
|||||||
|
"""NPU smoke test for the GLM-Image external-AR distributed topology."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import base64
|
||||||
|
import os
|
||||||
|
import signal
|
||||||
|
import subprocess
|
||||||
|
import time
|
||||||
|
import unittest
|
||||||
|
from concurrent.futures import ThreadPoolExecutor
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import requests
|
||||||
|
import torch
|
||||||
|
|
||||||
|
from sglang.multimodal_gen.test.server.ascend.testcase_configs_npu import (
|
||||||
|
GLM_IMAGE_WEIGHTS_PATH,
|
||||||
|
)
|
||||||
|
from sglang.multimodal_gen.test.test_utils import find_free_port, wait_for_server_health
|
||||||
|
from sglang.test.test_utils import CustomTestCase
|
||||||
|
|
||||||
|
HOST = "127.0.0.1"
|
||||||
|
_LOG_DIR = Path(os.environ.get("SGLANG_TEST_LOG_DIR", "/tmp"))
|
||||||
|
_STARTUP_TIMEOUT_S = float(os.environ.get("SGLANG_GLM_AR_STARTUP_TIMEOUT", "600"))
|
||||||
|
# A3 warm AR (26.711s) plus six sequential denoiser outputs (24.1945s each).
|
||||||
|
_EXPECTED_MAKESPAN_S = 172.0
|
||||||
|
_PERFORMANCE_TOLERANCE = 0.25
|
||||||
|
|
||||||
|
|
||||||
|
def _kill_process_tree(proc: subprocess.Popen) -> None:
|
||||||
|
try:
|
||||||
|
os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
|
||||||
|
except (ProcessLookupError, PermissionError):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
def _tail_log(path: Path, lines: int = 80) -> str:
|
||||||
|
if not path.exists():
|
||||||
|
return f"<no log at {path}>"
|
||||||
|
try:
|
||||||
|
return "\n".join(path.read_text(errors="ignore").splitlines()[-lines:])
|
||||||
|
except OSError as error:
|
||||||
|
return f"<log read failed: {error}>"
|
||||||
|
|
||||||
|
|
||||||
|
def _wait_for_log(path: Path, message: str, timeout: float) -> None:
|
||||||
|
deadline = time.monotonic() + timeout
|
||||||
|
while time.monotonic() < deadline:
|
||||||
|
if path.exists():
|
||||||
|
try:
|
||||||
|
if message in path.read_text(errors="ignore"):
|
||||||
|
return
|
||||||
|
except OSError:
|
||||||
|
pass
|
||||||
|
time.sleep(2)
|
||||||
|
raise TimeoutError(f"Missing {message!r} in {path}:\n{_tail_log(path)}")
|
||||||
|
|
||||||
|
|
||||||
|
class _GlmDistributedCluster:
|
||||||
|
def __init__(self) -> None:
|
||||||
|
self.model_path = Path(GLM_IMAGE_WEIGHTS_PATH)
|
||||||
|
self.ar_port = find_free_port(HOST)
|
||||||
|
self.denoiser_port = find_free_port(HOST)
|
||||||
|
self.head_port = find_free_port(HOST)
|
||||||
|
self.head_scheduler_port = find_free_port(HOST)
|
||||||
|
self.master_port = find_free_port(HOST)
|
||||||
|
self.processes: list[subprocess.Popen] = []
|
||||||
|
self.log_paths = {
|
||||||
|
"ar": _LOG_DIR / "glm_image_distributed_ar.log",
|
||||||
|
"denoiser": _LOG_DIR / "glm_image_distributed_denoiser.log",
|
||||||
|
"head": _LOG_DIR / "glm_image_distributed_head.log",
|
||||||
|
}
|
||||||
|
self._log_handles: list = []
|
||||||
|
|
||||||
|
def __enter__(self) -> _GlmDistributedCluster:
|
||||||
|
if not self.model_path.is_dir():
|
||||||
|
raise RuntimeError(
|
||||||
|
f"GLM-Image ModelScope cache is missing: {self.model_path}"
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
self._start_ar()
|
||||||
|
self._start_denoiser()
|
||||||
|
self._start_head()
|
||||||
|
except Exception:
|
||||||
|
self.stop()
|
||||||
|
raise
|
||||||
|
return self
|
||||||
|
|
||||||
|
def __exit__(self, *exc) -> None:
|
||||||
|
self.stop()
|
||||||
|
|
||||||
|
def _start_process(self, command: list[str], log_name: str) -> None:
|
||||||
|
log_handle = open(self.log_paths[log_name], "w")
|
||||||
|
self._log_handles.append(log_handle)
|
||||||
|
self.processes.append(
|
||||||
|
subprocess.Popen(
|
||||||
|
command,
|
||||||
|
stdout=log_handle,
|
||||||
|
stderr=subprocess.STDOUT,
|
||||||
|
preexec_fn=os.setsid,
|
||||||
|
env=os.environ.copy(),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
def _start_ar(self) -> None:
|
||||||
|
self._start_process(
|
||||||
|
[
|
||||||
|
"sglang",
|
||||||
|
"serve",
|
||||||
|
"--model-path",
|
||||||
|
str(self.model_path / "vision_language_encoder"),
|
||||||
|
"--tokenizer-path",
|
||||||
|
str(self.model_path / "processor"),
|
||||||
|
"--enable-multimodal",
|
||||||
|
"--device",
|
||||||
|
"npu",
|
||||||
|
"--attention-backend",
|
||||||
|
"ascend",
|
||||||
|
"--disable-fast-image-processor",
|
||||||
|
"--tp-size",
|
||||||
|
"1",
|
||||||
|
"--cuda-graph-bs",
|
||||||
|
"2",
|
||||||
|
"--base-gpu-id",
|
||||||
|
"0",
|
||||||
|
"--host",
|
||||||
|
HOST,
|
||||||
|
"--port",
|
||||||
|
str(self.ar_port),
|
||||||
|
],
|
||||||
|
"ar",
|
||||||
|
)
|
||||||
|
self._wait_for_health("ar", self.ar_port)
|
||||||
|
|
||||||
|
def _start_denoiser(self) -> None:
|
||||||
|
self._start_process(
|
||||||
|
[
|
||||||
|
"sglang",
|
||||||
|
"serve",
|
||||||
|
"--model-path",
|
||||||
|
str(self.model_path),
|
||||||
|
"--disagg-role",
|
||||||
|
"denoiser",
|
||||||
|
"--disagg-server-addr",
|
||||||
|
f"tcp://{HOST}:{self.head_scheduler_port}",
|
||||||
|
"--srt-encoder-url",
|
||||||
|
f"http://{HOST}:{self.ar_port}",
|
||||||
|
"--scheduler-port",
|
||||||
|
str(self.denoiser_port),
|
||||||
|
"--master-port",
|
||||||
|
str(self.master_port),
|
||||||
|
"--num-gpus",
|
||||||
|
"1",
|
||||||
|
"--base-gpu-id",
|
||||||
|
"1",
|
||||||
|
"--denoiser-sp",
|
||||||
|
"1",
|
||||||
|
"--cfg-parallel-size",
|
||||||
|
"1",
|
||||||
|
"--batching-max-size",
|
||||||
|
"1",
|
||||||
|
"--dit-cpu-offload",
|
||||||
|
"false",
|
||||||
|
"--attention-backend",
|
||||||
|
"fa",
|
||||||
|
],
|
||||||
|
"denoiser",
|
||||||
|
)
|
||||||
|
|
||||||
|
def _start_head(self) -> None:
|
||||||
|
self._start_process(
|
||||||
|
[
|
||||||
|
"sglang",
|
||||||
|
"serve",
|
||||||
|
"--model-path",
|
||||||
|
str(self.model_path),
|
||||||
|
"--disagg-role",
|
||||||
|
"server",
|
||||||
|
"--denoiser-urls",
|
||||||
|
f"tcp://{HOST}:{self.denoiser_port}",
|
||||||
|
"--srt-encoder-url",
|
||||||
|
f"http://{HOST}:{self.ar_port}",
|
||||||
|
"--batching-mode",
|
||||||
|
"dynamic",
|
||||||
|
"--batching-max-size",
|
||||||
|
"2",
|
||||||
|
"--batching-delay-ms",
|
||||||
|
"100",
|
||||||
|
"--scheduler-port",
|
||||||
|
str(self.head_scheduler_port),
|
||||||
|
"--host",
|
||||||
|
HOST,
|
||||||
|
"--port",
|
||||||
|
str(self.head_port),
|
||||||
|
],
|
||||||
|
"head",
|
||||||
|
)
|
||||||
|
_wait_for_log(
|
||||||
|
self.log_paths["denoiser"], "Role DENOISER ready", _STARTUP_TIMEOUT_S
|
||||||
|
)
|
||||||
|
self._wait_for_health("head", self.head_port)
|
||||||
|
|
||||||
|
def _wait_for_health(self, name: str, port: int) -> None:
|
||||||
|
try:
|
||||||
|
wait_for_server_health(
|
||||||
|
f"http://{HOST}:{port}", path="/v1/models", timeout=_STARTUP_TIMEOUT_S
|
||||||
|
)
|
||||||
|
except Exception as error:
|
||||||
|
raise RuntimeError(
|
||||||
|
f"{name} failed to become healthy:\n{_tail_log(self.log_paths[name])}"
|
||||||
|
) from error
|
||||||
|
|
||||||
|
def stop(self) -> None:
|
||||||
|
for process in self.processes:
|
||||||
|
_kill_process_tree(process)
|
||||||
|
for log_handle in self._log_handles:
|
||||||
|
log_handle.close()
|
||||||
|
self.processes.clear()
|
||||||
|
self._log_handles.clear()
|
||||||
|
|
||||||
|
|
||||||
|
class TestGlmImageDistributedNpu(CustomTestCase):
|
||||||
|
@classmethod
|
||||||
|
def setUpClass(cls) -> None:
|
||||||
|
super().setUpClass()
|
||||||
|
if not hasattr(torch, "npu") or torch.npu.device_count() < 2:
|
||||||
|
raise unittest.SkipTest("requires two Ascend NPUs")
|
||||||
|
cls.cluster = _GlmDistributedCluster()
|
||||||
|
cls.cluster.__enter__()
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def tearDownClass(cls) -> None:
|
||||||
|
if hasattr(cls, "cluster"):
|
||||||
|
for name, path in cls.cluster.log_paths.items():
|
||||||
|
print(f"\n=== [glm-image-distributed] {name} log tail ===")
|
||||||
|
print(_tail_log(path))
|
||||||
|
cls.cluster.stop()
|
||||||
|
super().tearDownClass()
|
||||||
|
|
||||||
|
def _generate(self, prompt: str, n: int = 1) -> list[bytes]:
|
||||||
|
response = requests.post(
|
||||||
|
f"http://{HOST}:{self.cluster.head_port}/v1/images/generations",
|
||||||
|
json={
|
||||||
|
"model": str(self.cluster.model_path),
|
||||||
|
"prompt": prompt,
|
||||||
|
"n": n,
|
||||||
|
"size": "1024x1024",
|
||||||
|
"response_format": "b64_json",
|
||||||
|
},
|
||||||
|
timeout=600,
|
||||||
|
)
|
||||||
|
response.raise_for_status()
|
||||||
|
images = response.json()["data"]
|
||||||
|
self.assertEqual(len(images), n)
|
||||||
|
return [base64.b64decode(image["b64_json"]) for image in images]
|
||||||
|
|
||||||
|
def test_external_ar_batching_multi_output_disaggregation_performance(
|
||||||
|
self,
|
||||||
|
) -> None:
|
||||||
|
self._generate("A warmup landscape")
|
||||||
|
requests_to_generate = [
|
||||||
|
("A mountain sunrise", 2),
|
||||||
|
("A city at night", 1),
|
||||||
|
("A forest lake", 2),
|
||||||
|
("A desert sunset", 1),
|
||||||
|
]
|
||||||
|
|
||||||
|
def generate(request: tuple[str, int]) -> tuple[list[bytes], float]:
|
||||||
|
prompt, output_count = request
|
||||||
|
start_time = time.perf_counter()
|
||||||
|
images = self._generate(prompt, n=output_count)
|
||||||
|
return images, time.perf_counter() - start_time
|
||||||
|
|
||||||
|
start_time = time.perf_counter()
|
||||||
|
with ThreadPoolExecutor(max_workers=4) as executor:
|
||||||
|
results = list(executor.map(generate, requests_to_generate))
|
||||||
|
makespan_s = time.perf_counter() - start_time
|
||||||
|
request_latencies_s = [latency_s for _, latency_s in results]
|
||||||
|
images = [image for request_images, _ in results for image in request_images]
|
||||||
|
|
||||||
|
print(
|
||||||
|
"GLM distributed performance: "
|
||||||
|
f"makespan={makespan_s:.2f}s, "
|
||||||
|
f"request_latencies={[round(value, 2) for value in request_latencies_s]}"
|
||||||
|
)
|
||||||
|
self.assertLessEqual(
|
||||||
|
makespan_s,
|
||||||
|
_EXPECTED_MAKESPAN_S * (1 + _PERFORMANCE_TOLERANCE),
|
||||||
|
"GLM external-AR and DiT overlap performance regressed",
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertTrue(
|
||||||
|
all(images),
|
||||||
|
"each requested output must produce a decoded image",
|
||||||
|
)
|
||||||
|
head_log = self.cluster.log_paths["head"].read_text(errors="ignore")
|
||||||
|
self.assertGreaterEqual(
|
||||||
|
head_log.count(
|
||||||
|
"GLM distributed AR dispatched batch size=2 requests, 2 outputs"
|
||||||
|
),
|
||||||
|
1,
|
||||||
|
"the two n=1 requests must share an external-AR batch",
|
||||||
|
)
|
||||||
|
self.assertGreaterEqual(
|
||||||
|
head_log.count(
|
||||||
|
"GLM distributed AR dispatched batch size=1 requests, 2 outputs"
|
||||||
|
),
|
||||||
|
2,
|
||||||
|
"each n=2 request must preserve both outputs through disaggregation",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
@@ -237,14 +237,26 @@ DEFAULT_EST_TIME_SECONDS = 300.0
|
|||||||
STARTUP_OVERHEAD_SECONDS = 120.0
|
STARTUP_OVERHEAD_SECONDS = 120.0
|
||||||
DEFAULT_STANDALONE_EST_TIME_SECONDS = 300.0
|
DEFAULT_STANDALONE_EST_TIME_SECONDS = 300.0
|
||||||
|
|
||||||
|
STANDALONE_FILES = {
|
||||||
|
"2-npu": [
|
||||||
|
"ascend/test_glm_image_distributed.py",
|
||||||
|
],
|
||||||
|
}
|
||||||
|
|
||||||
|
STANDALONE_FILE_EST_TIMES = {
|
||||||
|
"2-npu": {
|
||||||
|
"ascend/test_glm_image_distributed.py": 900.0,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
SUITES = {
|
SUITES = {
|
||||||
"1-npu": [
|
"1-npu": [
|
||||||
"ascend/test_server_1_npu.py",
|
"ascend/test_server_1_npu.py",
|
||||||
# add new 1-npu test files here
|
*STANDALONE_FILES.get("1-npu", []),
|
||||||
],
|
],
|
||||||
"2-npu": [
|
"2-npu": [
|
||||||
"ascend/test_server_2_npu.py",
|
"ascend/test_server_2_npu.py",
|
||||||
# add new 2-npu test files here
|
*STANDALONE_FILES.get("2-npu", []),
|
||||||
],
|
],
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -258,6 +270,5 @@ PARAMETRIZED_CASE_GROUPS = {
|
|||||||
}
|
}
|
||||||
|
|
||||||
FILE_SUITES = {}
|
FILE_SUITES = {}
|
||||||
STANDALONE_FILES = {}
|
|
||||||
COMPONENT_ACCURACY_SUITES = {}
|
COMPONENT_ACCURACY_SUITES = {}
|
||||||
_UPDATE_WEIGHTS_FROM_DISK_TEST_FILE = None
|
_UPDATE_WEIGHTS_FROM_DISK_TEST_FILE = None
|
||||||
|
|||||||
@@ -73,7 +73,7 @@ class TestGlmImageARSrtBackend(unittest.TestCase):
|
|||||||
|
|
||||||
@patch(
|
@patch(
|
||||||
"sglang.multimodal_gen.runtime.pipelines_core.stages."
|
"sglang.multimodal_gen.runtime.pipelines_core.stages."
|
||||||
"model_specific_stages.glm_image.get_local_torch_device",
|
"model_specific_stages.glm_image.current_platform.get_local_torch_device",
|
||||||
return_value=torch.device("cpu"),
|
return_value=torch.device("cpu"),
|
||||||
)
|
)
|
||||||
@patch(
|
@patch(
|
||||||
@@ -102,7 +102,7 @@ class TestGlmImageARSrtBackend(unittest.TestCase):
|
|||||||
|
|
||||||
@patch(
|
@patch(
|
||||||
"sglang.multimodal_gen.runtime.pipelines_core.stages."
|
"sglang.multimodal_gen.runtime.pipelines_core.stages."
|
||||||
"model_specific_stages.glm_image.get_local_torch_device",
|
"model_specific_stages.glm_image.current_platform.get_local_torch_device",
|
||||||
return_value=torch.device("cpu"),
|
return_value=torch.device("cpu"),
|
||||||
)
|
)
|
||||||
@patch(
|
@patch(
|
||||||
@@ -142,7 +142,7 @@ class TestGlmImageARSrtBackend(unittest.TestCase):
|
|||||||
|
|
||||||
@patch(
|
@patch(
|
||||||
"sglang.multimodal_gen.runtime.pipelines_core.stages."
|
"sglang.multimodal_gen.runtime.pipelines_core.stages."
|
||||||
"model_specific_stages.glm_image.get_local_torch_device",
|
"model_specific_stages.glm_image.current_platform.get_local_torch_device",
|
||||||
return_value=torch.device("cpu"),
|
return_value=torch.device("cpu"),
|
||||||
)
|
)
|
||||||
@patch(
|
@patch(
|
||||||
@@ -182,7 +182,7 @@ class TestGlmImageARSrtBackend(unittest.TestCase):
|
|||||||
|
|
||||||
@patch(
|
@patch(
|
||||||
"sglang.multimodal_gen.runtime.pipelines_core.stages."
|
"sglang.multimodal_gen.runtime.pipelines_core.stages."
|
||||||
"model_specific_stages.glm_image.get_local_torch_device",
|
"model_specific_stages.glm_image.current_platform.get_local_torch_device",
|
||||||
return_value=torch.device("cpu"),
|
return_value=torch.device("cpu"),
|
||||||
)
|
)
|
||||||
@patch(
|
@patch(
|
||||||
@@ -207,7 +207,7 @@ class TestGlmImageARSrtBackend(unittest.TestCase):
|
|||||||
|
|
||||||
@patch(
|
@patch(
|
||||||
"sglang.multimodal_gen.runtime.pipelines_core.stages."
|
"sglang.multimodal_gen.runtime.pipelines_core.stages."
|
||||||
"model_specific_stages.glm_image.get_local_torch_device",
|
"model_specific_stages.glm_image.current_platform.get_local_torch_device",
|
||||||
return_value=torch.device("cpu"),
|
return_value=torch.device("cpu"),
|
||||||
)
|
)
|
||||||
def test_forward_aligns_runtime_dimensions_before_ar_generation(self, _mock_device):
|
def test_forward_aligns_runtime_dimensions_before_ar_generation(self, _mock_device):
|
||||||
@@ -248,7 +248,7 @@ class TestGlmImageARSrtBackend(unittest.TestCase):
|
|||||||
|
|
||||||
@patch(
|
@patch(
|
||||||
"sglang.multimodal_gen.runtime.pipelines_core.stages."
|
"sglang.multimodal_gen.runtime.pipelines_core.stages."
|
||||||
"model_specific_stages.glm_image.get_local_torch_device",
|
"model_specific_stages.glm_image.current_platform.get_local_torch_device",
|
||||||
return_value=torch.device("cpu"),
|
return_value=torch.device("cpu"),
|
||||||
)
|
)
|
||||||
@patch(
|
@patch(
|
||||||
@@ -282,7 +282,7 @@ class TestGlmImageARSrtBackend(unittest.TestCase):
|
|||||||
|
|
||||||
@patch(
|
@patch(
|
||||||
"sglang.multimodal_gen.runtime.pipelines_core.stages."
|
"sglang.multimodal_gen.runtime.pipelines_core.stages."
|
||||||
"model_specific_stages.glm_image.get_local_torch_device",
|
"model_specific_stages.glm_image.current_platform.get_local_torch_device",
|
||||||
return_value=torch.device("cpu"),
|
return_value=torch.device("cpu"),
|
||||||
)
|
)
|
||||||
@patch(
|
@patch(
|
||||||
@@ -349,7 +349,7 @@ class TestGlmImageARSrtBackend(unittest.TestCase):
|
|||||||
|
|
||||||
@patch(
|
@patch(
|
||||||
"sglang.multimodal_gen.runtime.pipelines_core.stages."
|
"sglang.multimodal_gen.runtime.pipelines_core.stages."
|
||||||
"model_specific_stages.glm_image.get_local_torch_device",
|
"model_specific_stages.glm_image.current_platform.get_local_torch_device",
|
||||||
return_value=torch.device("cpu"),
|
return_value=torch.device("cpu"),
|
||||||
)
|
)
|
||||||
def test_generate_prior_tokens_rejects_unaligned_internal_dimensions(
|
def test_generate_prior_tokens_rejects_unaligned_internal_dimensions(
|
||||||
@@ -370,7 +370,7 @@ class TestGlmImageARSrtBackend(unittest.TestCase):
|
|||||||
|
|
||||||
@patch(
|
@patch(
|
||||||
"sglang.multimodal_gen.runtime.pipelines_core.stages."
|
"sglang.multimodal_gen.runtime.pipelines_core.stages."
|
||||||
"model_specific_stages.glm_image.get_local_torch_device",
|
"model_specific_stages.glm_image.current_platform.get_local_torch_device",
|
||||||
return_value=torch.device("cpu"),
|
return_value=torch.device("cpu"),
|
||||||
)
|
)
|
||||||
def test_generate_prior_tokens_batch_rejects_unaligned_internal_dimensions(
|
def test_generate_prior_tokens_batch_rejects_unaligned_internal_dimensions(
|
||||||
|
|||||||
@@ -110,7 +110,9 @@ def test_ar_stage_generates_one_prior_per_requested_output():
|
|||||||
)
|
)
|
||||||
|
|
||||||
with patch.object(
|
with patch.object(
|
||||||
glm_stage, "get_local_torch_device", return_value=torch.device("cpu")
|
glm_stage.current_platform,
|
||||||
|
"get_local_torch_device",
|
||||||
|
return_value=torch.device("cpu"),
|
||||||
):
|
):
|
||||||
result = stage.forward(batch, SimpleNamespace())
|
result = stage.forward(batch, SimpleNamespace())
|
||||||
|
|
||||||
@@ -135,7 +137,9 @@ def test_before_denoising_expands_latents_and_conditions_for_requested_outputs()
|
|||||||
)
|
)
|
||||||
|
|
||||||
with patch.object(
|
with patch.object(
|
||||||
glm_stage, "get_local_torch_device", return_value=torch.device("cpu")
|
glm_stage.current_platform,
|
||||||
|
"get_local_torch_device",
|
||||||
|
return_value=torch.device("cpu"),
|
||||||
):
|
):
|
||||||
result = stage.forward(batch, SimpleNamespace())
|
result = stage.forward(batch, SimpleNamespace())
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user