Skip to content

serving.md

Latest commit

 

History

History
147 lines (112 loc) · 8.02 KB

serving.md

File metadata and controls

147 lines (112 loc) · 8.02 KB
Raw

Serving

Modes

Colocated

Scheduler + workers run in colocated mode, with each worker performing both prefill and decode.

uv run python -m sangam.entrypoints.launch \
    --model GSAI-ML/LLaDA-8B-Instruct \
    --mode colocated \
    --gpus 0

Hybrid (Conditional Disaggregation)

Dedicated prefill workers and colocated workers which initially only process decodes. When prefill workers are overloaded, --enable-hybrid-prefill-overflow lets prefills fall through to colocated workers for local prefill+decode.

uv run python -m sangam.entrypoints.launch \
    --model GSAI-ML/LLaDA-8B-Instruct \
    --mode hybrid \
    --prefill-gpus 0 \
    --hybrid-colocated-gpus 1 \

Overflow is off by default, so hybrid mode runs in a pure disaggregated manner unless you pass --enable-hybrid-prefill-overflow.

Invocation Notes

A list of all CLI arguments with helptexts can be found using the command:

uv run python -m sangam.entrypoints.launch --help

Key CLI Arguments

Core

Argument Default Description
--mode colocated Serving mode: colocated or hybrid
--model GSAI-ML/LLaDA-8B-Instruct HuggingFace model name
--scheduler-port 50051 gRPC port for the scheduler
--gpus 0,1 Comma-separated GPU IDs for colocated workers
--prefill-gpus 0,1 Comma-separated GPU IDs for prefill workers (hybrid)
--hybrid-colocated-gpus 1,2,3 Comma-separated GPU IDs for colocated workers (hybrid)
--base-worker-port 20100 Starting port for worker gRPC servers
--master-addr localhost torch.distributed master address
--master-port 29500 torch.distributed master port
--max-batch-size 128 Max decode steps to batch per forward pass
--block-length 32 Tokens per generation block for every request handled by this server
--max-gen-len None Fixed prompt+mask length; when set, must be a positive multiple of --block-length
--mask-id None Token ID for masked positions; falls back to the HF config mask_token_id
--colocated-sticky-worker false Pin each request to its first-assigned worker

Scheduling & topology

Argument Default Description
--max-tokens-per-iteration 4096 Colocated mode token budget per scheduler iteration (decode + optional prefill admission)
--max-prefill-tokens-per-batch 4096 Max total token count across all requests in a single prefill forward pass
--prefill-scheduler-policy least_outstanding_prefill_tokens Prefill assignment policy: round_robin, least_outstanding_prefill_tokens, or colocated-only least_outstanding_requests, least_request_length_sum, balanced_length_clustering
--prefill-queue-policy arrival_order Per-worker prefill queue ordering: arrival_order or fewest_remaining_blocks
--decode-scheduler-policy max_free_memory Decode assignment policy: round_robin, max_free_memory, balanced_length_clustering, or topology-aware topology_guarded_memory in hybrid mode
--decode-grouping-slack-ratio 0.10 Slack ratio for decode balanced_length_clustering; workers within min_projected_sum * (1 + slack) remain eligible before clustering preference is applied
--kv-fast-pairs 0-1,2-3,4-5,6-7 Undirected fast GPU link pairs for topology-aware decode routing in hybrid mode, for example 0-1,2-3
--kv-topology-alpha 0.0 Prefer a fast-link decode worker only when its free KV pages are at least alpha * mem_best.free_pages
--enable-hybrid-prefill-overflow / --no-enable-hybrid-prefill-overflow on in hybrid mode Hybrid-only: allow overloaded prefill workers to overflow to colocated workers for local prefill+decode. On by default in hybrid mode; pass --no-enable-hybrid-prefill-overflow to queue such requests as pending instead
--prefill-overload-threshold 16384 Hybrid-only: outstanding prefill tokens per worker above which requests overflow to colocated workers

KV cache

Argument Default Description
--kv-page-size 16 Tokens per KV cache page
--kv-max-pages auto (per model) Max KV cache pages per decode worker. Auto-selected from the model's HF architecture when omitted (Dream 49152, LLaDA 5632); pass an explicit value to override

CUDA graphs

Argument Default Description
--enable-cuda-graphs / --no-enable-cuda-graphs on Capture and replay CUDA graphs for decode-only batches. On by default; pass --no-enable-cuda-graphs to disable
--cuda-graph-batch-sizes 1,2,4,8,16,24,32,40,48,56,64 Comma-separated decode batch sizes to capture, capped at --max-batch-size

Metrics

Argument Default Description
--metrics-output-dir benchmark_output/<timestamp> Directory for CSV and plot output
--disable-metrics false Disable metrics collection
--enable-individual-batch-metrics false Enable per-batch raw CSV export (worker_batch_metrics.csv)

Additional operation-level metrics flags (--enable-operation-metrics, --op-metrics-layer-id, --export-partial-metrics) are covered in the metrics doc.

Generation Parameters

Parameter Description
gen_length Total tokens to generate; must be divisible by the server --block-length
temperature Sampling temperature (0.0 = greedy via Gumbel max)
unmasking_strategy random, conf_threshold, conf_quota, or conf_dynamic
confidence_threshold Required when unmasking_strategy=conf_threshold
fixed_unmask_quota Required when unmasking_strategy=conf_quota
dynamic_unmask_factor Required when unmasking_strategy=conf_dynamic
request_seed Optional stable sampling seed (a top-level request field, not part of sampling_parameters); if omitted, the scheduler derives a deterministic fallback from request contents and sampling parameters

Request Submission

Use scripts/submit_request.py as the minimal client. A request consists of:

  • prompt_token_ids: fully tokenized prompt payload
  • gen_length: number of completion tokens to generate; it must be divisible by the server --block-length
  • request_seed: optional top-level request field (not nested under sampling_parameters) carrying a per-request sampling seed; identical explicit seeds reproduce the same nonzero-temperature sampling across schedulers
  • sampling_parameters: nested sampling config carrying temperature, unmasking_strategy, and any strategy-specific optional fields

If request_seed is omitted, the scheduler hashes the request payload, sampling parameters, and server block length to derive a deterministic fallback seed before the first worker enqueue. That means duplicate implicit requests intentionally share randomness within a given server configuration.

Polling returns the full current sequence plus a string status. The scheduler currently uses these statuses:

  • PENDING
  • PREFILLING
  • WAITING_DECODE
  • DECODING
  • WAITING_NEXT_BLOCK
  • COMPLETED
  • ERROR
  • NOT_FOUND for an unknown request id

Treat COMPLETED and ERROR as terminal states.

Troubleshooting

If launch succeeds but requests do not complete, check these first:

  • Python/CUDA environment: this project requires Python 3.13, CUDA 12.8, and compatible GPU drivers
  • Dependency resolution: uv sync must be able to use the configured FlashInfer package source from pyproject.toml
  • Ports: ensure --scheduler-port and --base-worker-port are free and do not overlap with other local runs
  • Worker topology: hybrid mode needs both --prefill-gpus and --hybrid-colocated-gpus; colocated mode needs --gpus
  • Scheduler readiness: if workers have not registered yet, requests can remain pending
  • CUDA/NCCL issues: a worker process can start but fail to make progress if inter-GPU communication is misconfigured
  • Terminal status: inspect ERROR responses from Poll; NOT_FOUND means the request id is unknown to the scheduler