diff --git a/examples/kimi_k3/README.md b/examples/kimi_k3/README.md index fc8d75879bd8..14d6023cd713 100644 --- a/examples/kimi_k3/README.md +++ b/examples/kimi_k3/README.md @@ -221,4 +221,4 @@ decoding requires the default cache manager, which cannot reuse blocks. and the TEP16/TEP8 latency recipes are unaffected. Tracked as TRTLLM-14904. - FP8 KV cache (`kv_cache_config.dtype: fp8`) is not yet supported. -- Speculative decoding: suffix-automaton speculation is supported for aggregated serving (`speculative_config: {decoding_type: SA}` in the extra LLM API options). For evaluation, use `eval_extra_llm_options_sa.yaml` (the `--sa` flag of the GSM8K job): that configuration runs with the overlap scheduler off, `max_batch_size` 8, and a matching CUDA-graph `max_batch_size`. Combining speculation with disaggregated serving is not yet supported. +- Speculative decoding: suffix-automaton speculation is supported for aggregated serving (`speculative_config: {decoding_type: SA}` in the extra LLM API options). For evaluation, use `eval_extra_llm_options_sa.yaml` (the `--sa` flag of the GSM8K job): that configuration runs with the overlap scheduler off, `max_batch_size` 8, and a matching CUDA-graph `max_batch_size`. Suffix-automaton speculation also works under disaggregated serving; enable it on the generation server with `examples/kimi_k3/disagg/gen_config.yaml` (SA runs eager with `max_batch_size` ≤ 8; see `examples/kimi_k3/disagg/README.md`). diff --git a/examples/kimi_k3/disagg/README.md b/examples/kimi_k3/disagg/README.md index f07a6d0d870c..be7b4214aa2d 100644 --- a/examples/kimi_k3/disagg/README.md +++ b/examples/kimi_k3/disagg/README.md @@ -11,7 +11,8 @@ for constraints. | File | Purpose | |---|---| | `ctx_config.yaml` | Context-server extra LLM-API options (DEP16, overlap scheduler off, no spec decode) | -| `gen_config_no_sa.yaml` | Generation-server options, no speculative decoding (CUDA graphs ON by default: GSM8K 96.89, 765/2138 tok/s @c64/c256 vs aggregated 643/1972; null `cuda_graph_config` for token-parity debugging). A suffix-automaton (SA) speculative-decoding variant lands together with K3 SA support. | +| `gen_config.yaml` | Generation-server options WITH suffix-automaton (SA) speculative decoding (DEP16, eager) | +| `gen_config_no_sa.yaml` | Generation-server options WITHOUT spec decode — use this first (CUDA graphs ON by default: GSM8K 96.89, 765/2138 tok/s @c64/c256 vs aggregated 643/1972; null `cuda_graph_config` for token-parity debugging) | | `disagg_proxy_config.yaml` | `trtllm-serve disaggregated` proxy config (1 ctx + 1 gen) | | `benchmark_kimi_k3_dep16.yaml` | Config for the SLURM benchmark harness (`examples/disaggregated/slurm/benchmark/submit.py`) | @@ -34,8 +35,8 @@ for constraints. C++ transceiver, which throws at construction for K3's `MixedMambaHybridCacheManager`. - `disable_overlap_scheduler: true` on the ctx server (disagg - requirement) and on the gen server (keeps the smoke runs maximally - comparable across configurations). + requirement) and on the gen server (SA runs eager; also keeps the + SA-off smoke maximally comparable). - `enable_block_reuse: false`, `tokens_per_block: 64`, no chunked prefill, beam width 1 (model requirements). - `max_tokens_in_buffer: 8448` covers the target max ISL of 8192; raise @@ -53,12 +54,14 @@ for constraints. ## Launch sequence (manual, single ctx + single gen) -Each K3 worker spans 16 GPUs (4 NVL72 nodes at 4 GPUs/node). Leave -`UCX_TLS` unset in every worker shell so UCX selects transports itself, -RDMA/verbs included; in particular, unset any container-default -`UCX_TLS=tcp`, which breaks V2 NIXL VRAM registration. Override the -transport list only on clusters whose verbs transports cannot -initialize (see caveats 1 and 4 below for the symptom and the override). +Each K3 worker spans 16 GPUs (4 NVL72 nodes at 4 GPUs/node). Environment +prerequisites for every worker shell (see caveats below for why): + +```bash +export UCX_TLS=tcp,self,sm,cuda_copy,cuda_ipc # on clusters where verbs cannot + # initialize; a container-default + # UCX_TLS=tcp breaks V2 NIXL +``` 1. Start the context server (16-rank MPI world across its 4 nodes): @@ -68,7 +71,7 @@ initialize (see caveats 1 and 4 below for the symptom and the override). --config examples/kimi_k3/disagg/ctx_config.yaml ``` -2. Start the generation server: +2. Start the generation server (SA off first): ```bash trtllm-llmapi-launch trtllm-serve $MODEL_PATH \ @@ -83,7 +86,9 @@ initialize (see caveats 1 and 4 below for the symptom and the override). trtllm-serve disaggregated -c examples/kimi_k3/disagg/disagg_proxy_config.yaml ``` -4. Send OpenAI-compatible requests to the proxy (port 8000). +4. Send OpenAI-compatible requests to the proxy (port 8000). Once the + SA-off path is parity-validated, restart the gen server with + `gen_config.yaml` to enable SA. ## SLURM benchmark harness @@ -101,11 +106,9 @@ python3 examples/disaggregated/slurm/benchmark/submit.py \ - **Gen-only baseline**: set `benchmark.mode: gen_only_no_context` (submit.py exports `TRTLLM_DISAGG_BENCHMARK_GEN_ONLY=1` to the workers) to measure the decode-side ceiling without KV transfer. -- The harness's `start_worker.sh` clears any container-provided - `UCX_TLS`, so workers run with UCX's own transport selection by - default. Clusters that need a transport override (caveat 4) carry it - via `TRTLLM_WORKER_UCX_TLS` in `worker_env_var`, which - `start_worker.sh` re-exports as `UCX_TLS` after the clear. +- The harness's `start_worker.sh` clears `UCX_TLS`; the config carries + the transport pin via `TRTLLM_WORKER_UCX_TLS`, which `start_worker.sh` + re-exports as `UCX_TLS` after the clear. - pyxis/enroot resets image-defined variables (notably `PATH`) at container start, so the config injects the in-place TRT-LLM venv via `TRTLLM_PATH_PREPEND` / `TRTLLM_PYTHONPATH_PREPEND`, applied inside @@ -123,22 +126,22 @@ python3 examples/disaggregated/slurm/benchmark/submit.py \ setup MPI collectives. Not an MPI/pmix or V2 code bug; with `UCX_TLS=tcp,self,sm,cuda_copy,cuda_ipc` V2 NIXL passes multi-node with no code change. -2. **No speculative decoding yet.** These configs run the gen server - without spec decode. Suffix-automaton (SA) speculative decoding for - K3 disagg is validated on the feature branch and lands in a separate - change together with K3 SA support (an SA `gen_config.yaml` variant - ships with it). +2. **SA ships eager here.** SA speculative decoding in disagg is + validated for accuracy (GSM8K parity with aggregated serving) with + CUDA graphs disabled, as configured in `gen_config.yaml`. SA with + CUDA graphs is functional (the MLA latent-cache append under CUDA + graphs handles spec-dec verification), but the disagg SA + graphs + perf points have not been re-measured yet, so `gen_config.yaml` + keeps graphs off. Start with `gen_config_no_sa.yaml` for the first + bring-up on a new cluster, then switch to `gen_config.yaml`. 3. **Matched-DP only.** Keep ctx and gen at identical DEP16 with attention-DP on both sides; heterogeneous parallelism with attention-DP off is rejected (see constraints above). -4. **Cluster environment** (NVL72 nodes): the default is to leave - `UCX_TLS` unset and let UCX pick transports (RDMA/verbs included). - Two exceptions: never run V2 NIXL with a container-default - `UCX_TLS=tcp` (breaks V2 NIXL VRAM registration; unset it), and on - clusters whose verbs transports cannot initialize, pin - `UCX_TLS=tcp,self,sm,cuda_copy,cuda_ipc` (an effective `UCX_TLS=all` - hangs setup there, see caveat 1). The pin excludes RDMA/verbs, so do - not carry it to clusters where verbs works. +4. **Cluster environment** (NVL72 nodes): on clusters where verbs + transports cannot initialize, pin + `UCX_TLS=tcp,self,sm,cuda_copy,cuda_ipc` (`UCX_TLS=all` hangs setup, + see caveat 1) and never run V2 NIXL with a container-default + `UCX_TLS=tcp` (breaks V2 NIXL VRAM registration) — unset/override it. No bounce env override is needed: the byte gate (`TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES`, default 2 MiB) is always cleared by K3 payloads (constraints section above). @@ -153,16 +156,18 @@ python3 examples/disaggregated/slurm/benchmark/submit.py \ the per-token MLA latent; ≥1024 MB for 8k ISL). An undersized region does not error — every transfer silently falls back to a much slower host-staged TCP path. -7. **Prefill capacity and TTFT under burst.** Without chunked prefill, +7. **SA caps gen-side batch size.** SA requires `max_batch_size` ≤ 8 on + the generation server, which bounds per-instance concurrency at + `8 × dp_size` (128 with DEP16). Plan instance counts accordingly. +8. **Prefill capacity and TTFT under burst.** Without chunked prefill, context-server throughput is limited and queued prefills grow TTFT roughly linearly under closed-loop bursts. Rate-match the ctx:gen instance ratio to the expected traffic instead of oversubscribing a single context server. -8. **Startup time.** Weight loading takes tens of minutes per 16-GPU +9. **Startup time.** Weight loading takes tens of minutes per 16-GPU instance before the first token; set health-check, idle-reaper, and job time limits accordingly. The disaggregated proxy does not serve `/v1/models` (404) — point readiness probes at a different endpoint. -9. **`max_num_tokens` coupling.** The generation side must cover - `max_batch_size × (1 + max_draft_len)` scheduled tokens (with spec - decode off, `max_draft_len` is 0); the context side needs - `max_tokens_in_buffer` ≥ max ISL (see constraints above). +10. **`max_num_tokens` coupling.** The generation side must cover + `max_batch_size × (1 + max_draft_len)`; the context side needs + `max_tokens_in_buffer` ≥ max ISL (see constraints above). diff --git a/examples/kimi_k3/disagg/gen_config.yaml b/examples/kimi_k3/disagg/gen_config.yaml new file mode 100644 index 000000000000..10d28d6d27e8 --- /dev/null +++ b/examples/kimi_k3/disagg/gen_config.yaml @@ -0,0 +1,48 @@ +# Kimi K3 disaggregated serving - GENERATION (decode) server extra LLM-API +# options WITH suffix-automaton (SA) speculative decoding +# (`trtllm-serve --config gen_config.yaml`). +# +# DEP16 deployment (attention data-parallel + MoE EP dispatch/combine), +# mirroring examples/kimi_k3/eval_extra_llm_options_sa.yaml: SA runs +# eager (no CUDA graphs, no overlap scheduler) and needs +# max_batch_size <= 8 (SpeculativeState buffers). For the SA-off variant +# (validate KV/KDA transfer before enabling SA), use gen_config_no_sa.yaml. +tensor_parallel_size: 16 +moe_expert_parallel_size: 16 +enable_attention_dp: true +pipeline_parallel_size: 1 +disable_overlap_scheduler: true +enable_chunked_prefill: false +cuda_graph_config: null +max_batch_size: 8 +# Decode-shaped: per scheduled request 1 + max_draft_len tokens; 256 +# leaves generous headroom for warmup/dummy batches. +max_num_tokens: 256 +# Target ISL 8192 + OSL 1024 (+ margin). +max_seq_len: 9236 +kv_cache_config: + enable_block_reuse: false # mandatory for K3 (KDA state) + tokens_per_block: 64 # mandatory for K3 MLA kernels + # Proven by the aggregated DEP16 SA evals on GB300; bounds decode + # concurrency x seq-len. Tune upward once the deployment is stable. + free_gpu_memory_fraction: 0.25 +cache_transceiver_config: + backend: NIXL + # Mandatory for K3: `auto` resolves to the C++ transceiver, which + # throws for K3's MixedMambaHybridCacheManager. + transceiver_runtime: PYTHON + # Must be >= the target max ISL and match the ctx side. + max_tokens_in_buffer: 8448 + # Fabric-VMM bounce buffer (recv side); must match the ctx server. + # Without it, cross-node V2 pool-to-pool transfers fall back to + # ~0.4 GB/s host-staged tcp; 512 MiB measured at ~455 GB/s/GPU. + # Bounce engages automatically above TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES + # (default 2 MiB) — always true for K3 payloads. + # Sized to fit one full request payload (fixed 433 MiB KDA state + + # ~27 KB/token MLA latent; 649 MiB at 8k ISL) -- see ctx_config.yaml; + # at 512 every 8k transfer falls back to the per-fragment tcp path. + kv_cache_bounce_size_mb: 1024 +speculative_config: + decoding_type: SA + max_draft_len: 2 + max_matching_ngram_size: -1 # longest match via suffix automaton diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_sa.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_sa.yaml new file mode 100644 index 000000000000..397defec094d --- /dev/null +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_sa.yaml @@ -0,0 +1,21 @@ +hostname: localhost +model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 +free_gpu_memory_fraction: 0.1 +backend: pytorch +disable_overlap_scheduler: true +context_servers: + num_instances: 1 + tensor_parallel_size: 1 + pipeline_parallel_size: 1 + cache_transceiver_config: + backend: DEFAULT +generation_servers: + num_instances: 1 + tensor_parallel_size: 1 + pipeline_parallel_size: 1 + cache_transceiver_config: + backend: DEFAULT + speculative_config: + decoding_type: SA + max_draft_len: 4 + max_matching_ngram_size: 4 diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_sa_python.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_sa_python.yaml new file mode 100644 index 000000000000..82984ec49b49 --- /dev/null +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_sa_python.yaml @@ -0,0 +1,29 @@ +# Spec-split SA regression on the V2 PYTHON transceiver path: +# identical to disagg_config_sa.yaml (ctx: no speculative_config, +# gen: SA) but with backend NIXL + transceiver_runtime PYTHON on both +# servers -- the runtime combination Kimi K3 disagg is required to use. +# The existing test_disaggregated_sa covers the spec split only on the +# C++ DEFAULT backend. +hostname: localhost +model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 +free_gpu_memory_fraction: 0.1 +backend: pytorch +disable_overlap_scheduler: true +context_servers: + num_instances: 1 + tensor_parallel_size: 1 + pipeline_parallel_size: 1 + cache_transceiver_config: + backend: NIXL + transceiver_runtime: PYTHON +generation_servers: + num_instances: 1 + tensor_parallel_size: 1 + pipeline_parallel_size: 1 + cache_transceiver_config: + backend: NIXL + transceiver_runtime: PYTHON + speculative_config: + decoding_type: SA + max_draft_len: 4 + max_matching_ngram_size: 4 diff --git a/tests/integration/defs/disaggregated/test_disaggregated.py b/tests/integration/defs/disaggregated/test_disaggregated.py index dd8b73467305..f35b49961440 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated.py +++ b/tests/integration/defs/disaggregated/test_disaggregated.py @@ -347,6 +347,10 @@ def get_test_config(test_desc, example_dir, test_root): f"{test_configs_root}/disagg_config_conditional.yaml", "ngram": f"{test_configs_root}/disagg_config_ngram.yaml", + "sa": + f"{test_configs_root}/disagg_config_sa.yaml", + "sa_python": + f"{test_configs_root}/disagg_config_sa_python.yaml", "ctxpp2_genpp2": f"{test_configs_root}/disagg_config_ctxpp2_genpp2.yaml", "ctxtp2_genpp2": @@ -1826,6 +1830,37 @@ def test_disaggregated_ngram(disaggregated_test_root, llm_venv, cwd=llm_venv.get_working_directory()) +@pytest.mark.parametrize("llama_model_root", ['TinyLlama-1.1B-Chat-v1.0'], + indirect=True) +def test_disaggregated_sa(disaggregated_test_root, llm_venv, + disaggregated_example_root, llama_model_root): + setup_model_symlink(llm_venv, llama_model_root, + "TinyLlama/TinyLlama-1.1B-Chat-v1.0") + run_disaggregated_test(disaggregated_example_root, + "sa", + env=llm_venv._new_env, + model_path=llama_model_root, + cwd=llm_venv.get_working_directory()) + + +@pytest.mark.parametrize("llama_model_root", ['TinyLlama-1.1B-Chat-v1.0'], + indirect=True) +def test_disaggregated_sa_python(disaggregated_test_root, llm_venv, + disaggregated_example_root, llama_model_root): + """Spec-split SA (ctx no-spec, gen SA) on the V2 PYTHON transceiver path. + + NIXL + transceiver_runtime PYTHON. The existing test_disaggregated_sa + covers this split only on the C++ DEFAULT backend. + """ + setup_model_symlink(llm_venv, llama_model_root, + "TinyLlama/TinyLlama-1.1B-Chat-v1.0") + run_disaggregated_test(disaggregated_example_root, + "sa_python", + env=llm_venv._new_env, + model_path=llama_model_root, + cwd=llm_venv.get_working_directory()) + + @pytest.mark.skip_less_device(4) @pytest.mark.parametrize("llama_model_root", ['TinyLlama-1.1B-Chat-v1.0'], indirect=True) diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index a4391476fd75..4864fccbd18a 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -885,6 +885,8 @@ disaggregated/test_disaggregated.py::test_disaggregated_mamba_bs1_concurrency2 disaggregated/test_disaggregated.py::test_disaggregated_mixed[TinyLlama-1.1B-Chat-v1.0] disaggregated/test_disaggregated.py::test_disaggregated_multi_gpu[TinyLlama-1.1B-Chat-v1.0] disaggregated/test_disaggregated.py::test_disaggregated_ngram[TinyLlama-1.1B-Chat-v1.0] +disaggregated/test_disaggregated.py::test_disaggregated_sa[TinyLlama-1.1B-Chat-v1.0] +disaggregated/test_disaggregated.py::test_disaggregated_sa_python[TinyLlama-1.1B-Chat-v1.0] disaggregated/test_disaggregated.py::test_disaggregated_overlap[TinyLlama-1.1B-Chat-v1.0] disaggregated/test_disaggregated.py::test_disaggregated_overlap_gen_first[ctx_pp1-TinyLlama-1.1B-Chat-v1.0] disaggregated/test_disaggregated.py::test_disaggregated_overlap_gen_first[ctx_pp4-TinyLlama-1.1B-Chat-v1.0] diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 057d0ae50a3a..f051f8f1989f 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -66,6 +66,7 @@ l0_a10: - disaggregated/test_disaggregated.py::test_disaggregated_conditional[TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated.py::test_disaggregated_benchmark_gen_only_insufficient_kv[TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated.py::test_disaggregated_ngram[TinyLlama-1.1B-Chat-v1.0] + - disaggregated/test_disaggregated.py::test_disaggregated_sa[TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated.py::test_disaggregated_chat_completion_tool_calls[TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_workers.py::test_workers_conditional_disaggregation[TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_workers.py::test_workers_kv_cache_events[TinyLlama-1.1B-Chat-v1.0]