[TRTLLM-15465][feat] Support SA speculative decoding under disaggregated serving for Kimi K3

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`).

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).

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 <model> --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

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

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

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)

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]

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]