mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-30 15:57:20 +08:00
8e488bbd55
A prompt-only engine tray (PE) and a decode-only engine tray (CE) behind the OpenAI-compatible frontend from serve/openai_server.py: - serve/pd_server.py: --role pe / --role ce. The CE frontend renders the prompt, picks a CE rank (affinity or least-loaded routing) and POSTs /prefill to the PE. The PE prefills with the host prefix cache on, exports the prompt's KV blocks and the full SWA state in the pe_do_save_kvcache hook, streams the first sampled token back in the HTTP response and pushes the KV bytes to the CE rank over TCP. The CE rank admits the sequence into decode via Scheduler.add_prompted_sequence and reclaims PE sequences through Scheduler.free_finished_seqs. - deepinfer/hicache/pd_transfer.py: export/admit of KV pages and state through the hicache page-copy kernels. The CE keeps the pinned receive buffers alive until the admit copy kernels finish (the kernels read them by raw pointer; the caching host allocator otherwise reused the memory while the copy was queued -> cudaErrorLaunchFailure). - serve/pd_run.sh start/stop/logs over ssh; serve/check_pd.py compares outputs against a single-engine server token by token; serve/nccl_xtray_test.py reports the cross-host NCCL transport. Measured on 4 GB300 (PE) + 4 GB300 (CE) against one 4-GPU tray with the AgentX replay, outputs identical: ITL p90 stays within 8.6-15 ms from conc 8 to 256 (single tray 12-98 ms), P90 interactivity 66-117 tok/s/user where the single tray falls to 10; per total chip PD is behind up to conc 64 and ahead at 256 (doc/benchmark_gb300.md). Depends on the host prefix cache (#144), the OpenAI server (#147) and the scheduler bindings (add_prompted_sequence / free_finished_seqs), which are included as the first three commits. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01LFPG29daR1FhchJBTjiEKr
31 lines
912 B
Python
31 lines
912 B
Python
"""Cross-tray NCCL check: run one rank on each of two trays and report the transport NCCL picks + bandwidth.
|
|
|
|
# on tray A (rank 0) and tray B (rank 1), same MASTER_ADDR:
|
|
MASTER_ADDR=<trayA> MASTER_PORT=29999 RANK=0 WORLD_SIZE=2 CUDA_VISIBLE_DEVICES=0 NCCL_DEBUG=INFO python serve/nccl_xtray_test.py
|
|
"""
|
|
|
|
import os
|
|
import time
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
dist.init_process_group("nccl")
|
|
rank = dist.get_rank()
|
|
torch.cuda.set_device(0)
|
|
x = torch.ones(256 << 20, dtype=torch.uint8, device="cuda") # 256 MiB
|
|
for _ in range(3):
|
|
dist.all_reduce(x)
|
|
torch.cuda.synchronize()
|
|
t0 = time.perf_counter()
|
|
n = 10
|
|
for _ in range(n):
|
|
dist.all_reduce(x)
|
|
torch.cuda.synchronize()
|
|
dt = (time.perf_counter() - t0) / n
|
|
print(
|
|
f"rank {rank} on {os.uname().nodename}: all_reduce 256 MiB avg {dt * 1e3:.2f} ms -> {x.numel() / dt / 2**30:.1f} GiB/s algbw",
|
|
flush=True,
|
|
)
|
|
dist.destroy_process_group()
|