diff --git a/examples/disaggregated/slurm/benchmark/run_benchmark.sh b/examples/disaggregated/slurm/benchmark/run_benchmark.sh index 8fa89815704c..154bf7840d7c 100644 --- a/examples/disaggregated/slurm/benchmark/run_benchmark.sh +++ b/examples/disaggregated/slurm/benchmark/run_benchmark.sh @@ -5,6 +5,22 @@ set -e set -u trap 'echo "Error occurred at line $LINENO"; exit 1' ERR +# Container runtimes (pyxis/enroot) reset image-defined variables like PATH +# at container start, so values passed via srun --export are lost for them. +# Allow the launcher config to prepend entries from inside the container. +# The client_cmds srun passes these via --export with single-quoted values +# (submit.py convert_envs_to_str), and srun keeps the quotes literal — strip them. +TRTLLM_PATH_PREPEND="${TRTLLM_PATH_PREPEND:-}" +TRTLLM_PATH_PREPEND="${TRTLLM_PATH_PREPEND#\'}"; TRTLLM_PATH_PREPEND="${TRTLLM_PATH_PREPEND%\'}" +TRTLLM_PYTHONPATH_PREPEND="${TRTLLM_PYTHONPATH_PREPEND:-}" +TRTLLM_PYTHONPATH_PREPEND="${TRTLLM_PYTHONPATH_PREPEND#\'}"; TRTLLM_PYTHONPATH_PREPEND="${TRTLLM_PYTHONPATH_PREPEND%\'}" +if [ -n "${TRTLLM_PATH_PREPEND:-}" ]; then + export PATH="${TRTLLM_PATH_PREPEND}:${PATH}" +fi +if [ -n "${TRTLLM_PYTHONPATH_PREPEND:-}" ]; then + export PYTHONPATH="${TRTLLM_PYTHONPATH_PREPEND}${PYTHONPATH:+:${PYTHONPATH}}" +fi + # Add parameter validation if [ "$#" -lt 10 ]; then echo "Error: Missing required arguments, got $# arguments, args: $@" diff --git a/examples/disaggregated/slurm/benchmark/start_server.sh b/examples/disaggregated/slurm/benchmark/start_server.sh index ff8e5aa90277..690748a7d816 100644 --- a/examples/disaggregated/slurm/benchmark/start_server.sh +++ b/examples/disaggregated/slurm/benchmark/start_server.sh @@ -5,4 +5,26 @@ set -x config_file=$1 -trtllm-serve disaggregated -c ${config_file} -t 7200 -r 7200 +# Container runtimes (pyxis/enroot) reset image-defined variables like PATH +# at container start, so values passed via srun --export are lost for them. +# Allow the launcher config to prepend entries from inside the container. +# srun --export keeps any quotes in the exported values literal; strip them. +TRTLLM_PATH_PREPEND="${TRTLLM_PATH_PREPEND:-}" +TRTLLM_PATH_PREPEND="${TRTLLM_PATH_PREPEND#\'}"; TRTLLM_PATH_PREPEND="${TRTLLM_PATH_PREPEND%\'}" +TRTLLM_PYTHONPATH_PREPEND="${TRTLLM_PYTHONPATH_PREPEND:-}" +TRTLLM_PYTHONPATH_PREPEND="${TRTLLM_PYTHONPATH_PREPEND#\'}"; TRTLLM_PYTHONPATH_PREPEND="${TRTLLM_PYTHONPATH_PREPEND%\'}" +if [ -n "${TRTLLM_PATH_PREPEND:-}" ]; then + export PATH="${TRTLLM_PATH_PREPEND}:${PATH}" +fi +if [ -n "${TRTLLM_PYTHONPATH_PREPEND:-}" ]; then + export PYTHONPATH="${TRTLLM_PYTHONPATH_PREPEND}${PYTHONPATH:+:${PYTHONPATH}}" +fi + +# In-place (.pth-style) TRT-LLM installs may lack the trtllm-serve console +# script; fall back to the module entry point in that case. +trtllm_serve_cmd="trtllm-serve" +if ! command -v trtllm-serve >/dev/null 2>&1; then + trtllm_serve_cmd="python3 -m tensorrt_llm.commands.serve" +fi + +${trtllm_serve_cmd} disaggregated -c ${config_file} -t 7200 -r 7200 diff --git a/examples/disaggregated/slurm/benchmark/start_worker.sh b/examples/disaggregated/slurm/benchmark/start_worker.sh index 0a5b5897b773..3f695bd48a32 100644 --- a/examples/disaggregated/slurm/benchmark/start_worker.sh +++ b/examples/disaggregated/slurm/benchmark/start_worker.sh @@ -38,8 +38,30 @@ else export CUDA_VISIBLE_DEVICES=${cuda_devices} fi -# Clear UCX_TLS for specific clusters -unset UCX_TLS +# Container runtimes (pyxis/enroot) reset image-defined variables like PATH +# at container start, so values passed via srun --export are lost for them. +# Allow the launcher config to prepend entries from inside the container. +# srun --export keeps any quotes in the exported values literal; strip them. +TRTLLM_PATH_PREPEND="${TRTLLM_PATH_PREPEND:-}" +TRTLLM_PATH_PREPEND="${TRTLLM_PATH_PREPEND#\'}"; TRTLLM_PATH_PREPEND="${TRTLLM_PATH_PREPEND%\'}" +TRTLLM_PYTHONPATH_PREPEND="${TRTLLM_PYTHONPATH_PREPEND:-}" +TRTLLM_PYTHONPATH_PREPEND="${TRTLLM_PYTHONPATH_PREPEND#\'}"; TRTLLM_PYTHONPATH_PREPEND="${TRTLLM_PYTHONPATH_PREPEND%\'}" +if [ -n "${TRTLLM_PATH_PREPEND:-}" ]; then + export PATH="${TRTLLM_PATH_PREPEND}:${PATH}" +fi +if [ -n "${TRTLLM_PYTHONPATH_PREPEND:-}" ]; then + export PYTHONPATH="${TRTLLM_PYTHONPATH_PREPEND}${PYTHONPATH:+:${PYTHONPATH}}" +fi + +# Clear UCX_TLS for specific clusters. Some clusters instead need an +# explicit transport list (e.g. NVL72 nodes whose verbs transports cannot +# initialize): set TRTLLM_WORKER_UCX_TLS in worker_env_var to re-pin +# UCX_TLS here, after the container-provided value is cleared. +if [ -n "${TRTLLM_WORKER_UCX_TLS:-}" ]; then + export UCX_TLS="${TRTLLM_WORKER_UCX_TLS}" +else + unset UCX_TLS +fi echo "SLURM_PROCID: ${SLURM_PROCID}, hostname: $(hostname), instance_id: ${instance_id}" echo "CUDA_VISIBLE_DEVICES: ${CUDA_VISIBLE_DEVICES}" @@ -63,7 +85,14 @@ else nsys_prefix="nsys profile -o ${nsys_file} -f true -t cuda,nvtx,python-gil -c cudaProfilerApi --cuda-graph-trace node --capture-range-end=stop --gpu-metrics-devices=none" fi +# In-place (.pth-style) TRT-LLM installs may lack the trtllm-serve console +# script; fall back to the module entry point in that case. +trtllm_serve_cmd="trtllm-serve" +if ! command -v trtllm-serve >/dev/null 2>&1; then + trtllm_serve_cmd="python3 -m tensorrt_llm.commands.serve" +fi + ${nsys_prefix} trtllm-llmapi-launch ${numa_bind_cmd} \ - trtllm-serve ${model_path} \ + ${trtllm_serve_cmd} ${model_path} \ --host $(hostname) --port ${port} \ --config ${config_file} diff --git a/examples/disaggregated/slurm/cache_transceiver_test/configs/kda_payload_kimi_k3.yaml b/examples/disaggregated/slurm/cache_transceiver_test/configs/kda_payload_kimi_k3.yaml new file mode 100644 index 000000000000..674c275890a8 --- /dev/null +++ b/examples/disaggregated/slurm/cache_transceiver_test/configs/kda_payload_kimi_k3.yaml @@ -0,0 +1,83 @@ +# Kimi K3 KDA-payload cache-transceiver micro-benchmark. +# Run from the repository root: +# python3 examples/disaggregated/slurm/cache_transceiver_test/submit.py \ +# -c examples/disaggregated/slurm/cache_transceiver_test/configs/kda_payload_kimi_k3.yaml +# +# Synthetic KV geometry sized so ONE request's per-rank transfer equals the +# exact per-request KDA state payload of Kimi K3: 454,459,392 bytes +# (69 layers x (conv [3*96*128, 4] bf16 + delta [96,128,128] fp32)). +# Per-rank bytes/token = 69 layers * kvFactor(2) * (24/4 kv heads) * 128 * 2B +# = 211,968 B -> 2144 tokens == 454,459,392 B exactly. +# request_lengths 512 / 1024 also match MLA-latent payloads at ~4k / ~8k ISL +# (bf16) for the KDA-vs-MLA comparison. +# +# Fill in the <...> placeholders for your cluster before submitting. + +slurm: + partition: "" + account: "" + job_time: "00:30:00" + job_name: "ctt_kda_payload" + extra_args: "--gpus-per-node=4" + +hardware: + gpus_per_node: 4 + +environment: + container_image: "" + # Mount the TRT-LLM checkout into the container. + container_mount: "" # Format: path1:path1,path2:path2 + # TRT-LLM comes from the checkout's venv (PATH below) + checkout code + # (PYTHONPATH below); skip pip install. + trtllm_repo: "" + # Output directory (results.json, logs/, csv/); point at any writable + # scratch location. + work_dir: "" + trtllm_wheel_path: "" + build_wheel: false + cuda_architectures: "" + +test_matrix: + combinations: + # NIXL/PYTHON dropped for now: V2 _exchange_rank_info mpi_allgather can + # hang under this harness (srun --mpi=pmix). + - {backend: "NIXL", runtime: "CPP"} + - {backend: "UCX", runtime: "CPP"} + cache_manager_versions: ["V1"] + # 512 tok = 108.5 MB/rank (~ MLA bf16 @4k ISL), 1024 = 217.1 MB (~ MLA @8k), + # 2144 = 454.46 MB = exact K3 per-request KDA payload. + request_lengths: [512, 1024, 2144] + num_requests_per_length: 8 + warmup_requests: 2 + +kv_cache: + num_layers: 69 + num_kv_heads: 24 # /tp4 -> 6 heads per rank + head_dim: 128 + tokens_per_block: 32 + dtype: "HALF" + max_tokens_in_buffer: 4096 + +parallel: + ctx_tp: 4 + ctx_pp: 1 + gen_tp: 4 + gen_pp: 1 + +ucx_env_sweep: + # Single sweep. PATH/PYTHONPATH are exported inside the container prelude so + # python3 resolves to the checkout's venv (system-site-packages=true, + # so container torch/mpi4py are visible) and the checkout on PYTHONPATH + # wins. Keep your cluster's SLURM bin dir on PATH so host-side srun works. + # The UCX_TLS pin (no verbs) is for clusters where verbs transports cannot + # initialize on the compute nodes; drop the pin where verbs works. + - name: "venv_no_verbs" + env: + UCX_TLS: "tcp,self,sm,cuda_copy,cuda_ipc" + PATH: ":/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:" + PYTHONPATH: "" + +run: + timeout_per_cell_s: 120 + max_sweep_s: 900 + capture_proto_info: true diff --git a/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py b/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py index c4ebd8c6460c..3bfb64a8dcbc 100644 --- a/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py +++ b/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py @@ -287,6 +287,35 @@ class _TransferError(Exception): pass +def _wait_ctx_complete(xcvr, rid, runtime): + """Block until this ctx request's send finishes (or errors). + + * C++ transceiver: check_context_transfer_status(None) is a true block-all; + a single call suffices (see the comment at the call site). + * PYTHON (V2) transceiver: even under block_all, each TxSession wait is + bounded by kv_transfer_sender_future_timeout_ms (default 1000 ms). A + slower peer handshake makes the call log "TxSession ... timed out" and + return with the request still DISAGG_CONTEXT_TRANS_IN_PROGRESS -- in the + real executor that is benign (the session stays open and is re-polled + every iteration), but returning here would free and refill the KV blocks + mid-flight, so the receiver reads the NEXT request's pattern (verify FAIL + on every request except the last). So poll until + this rid lands in the completed/failed lists. Collectively safe: every + ctx rank loops on the same rid and the per-call consensus makes all + ranks observe completion on the same iteration. The per-cell + signal.alarm and the hang detector bound the loop. + """ + if runtime != "PYTHON": + xcvr.check_context_transfer_status(None) + return + while True: + completed, failed = xcvr.check_context_transfer_status(None) + if rid in failed: + raise _TransferError(f"ctx transfer failed for rid={rid}") + if rid in completed: + return + + def _wait_gen_complete(xcvr, req, runtime): """Block until this gen request's receive finishes (or errors). @@ -372,7 +401,9 @@ def run_one_request( # transfer is still in progress. NIXL/UCX cold-start connection setup can # exceed that, so the harness would free the request mid-transfer, leaving # the gen side hung and the ctx sender thread asserting on a freed session. - xcvr.check_context_transfer_status(None) + # The PYTHON runtime additionally needs a poll loop on top of block_all; + # see _wait_ctx_complete. + _wait_ctx_complete(xcvr, rid, runtime) state = req.state tensorrt_llm.logger.info(f"[ctx r{rank}] rid={rid}: transfer DONE (send), state={state}") free_sequence(kvm, req, kv_handle, use_v2) @@ -648,6 +679,13 @@ def arm_watchdog(combination_idx, reqlen_idx, what): backend=backend, transceiver_runtime=(None if runtime == "CPP" else "PYTHON"), max_tokens_in_buffer=cfg["kv_cache"]["max_tokens_in_buffer"], + # PYTHON (V2) only; 0 keeps bounce off. With bounce on, the KV data + # rides a fabric-VMM staging buffer (CU_MEM_HANDLE_TYPE_FABRIC), which + # is what lets UCX pick cuda_ipc across NVL72 nodes -- direct + # pool-to-pool transfers from non-fabric allocations fall back to + # much slower host-staged tcp, so enable bounce for cross-node + # transfers inside an NVLink domain. + kv_cache_bounce_size_mb=int(cfg["kv_cache"].get("bounce_size_mb", 0)), ) # Build the cache manager + transceiver ONCE per case (the manager is diff --git a/examples/kimi_k3/README.md b/examples/kimi_k3/README.md index 00349c6d6d43..e58c618b5628 100644 --- a/examples/kimi_k3/README.md +++ b/examples/kimi_k3/README.md @@ -193,4 +193,3 @@ default cache manager. 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). Combining speculation with disaggregated serving is not yet supported. -- Disaggregated serving is not yet supported. diff --git a/examples/kimi_k3/disagg/README.md b/examples/kimi_k3/disagg/README.md new file mode 100644 index 000000000000..8597bd8fe643 --- /dev/null +++ b/examples/kimi_k3/disagg/README.md @@ -0,0 +1,165 @@ +# Kimi K3 disaggregated serving (ctx/gen split) + +Configuration pair + deployment wiring for running Kimi K3 with separate +context (prefill) and generation (decode) servers. Status: **validated +end-to-end on hardware** (GB300 NVL72, 1 ctx + 1 gen, DEP16 both sides, +GSM8K accuracy parity with aggregated serving) — see the caveats section +for constraints. + +## Files + +| 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. | +| `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`) | + +## K3 constraints baked into the configs + +- **EP-only parallelism on BOTH sides**: `ep_size == tp_size`, no PP, no + TP on linears. Deployed as DEP-N (`enable_attention_dp: true`). +- **Matched ctx/gen parallelism (DEP16 = DEP16)** for now. Heterogeneous + ctx/gen TP with attention-DP *off* would silently corrupt memory for + K3's replicated KDA state and is rejected at peer registration — do + not deviate. Hetero DEP with attention-DP on both sides is believed + correct but unvalidated. +- **Ctx sizing = DEP16**: DEP16 is the smallest verified fit + (~193 GiB weights/rank on GB300). A DEP8 ctx is estimated at + ~273 GiB/rank for weights alone (extrapolating the 1.5 TB checkpoint: + replicated share ~113 GiB + experts/8), leaving no activation headroom + on GB300 (288 GiB) and not fitting GB200 (186 GiB). Treat DEP8-ctx as + ruled out on GB200 and an open (likely negative) question on GB300. +- **`transceiver_runtime: PYTHON` is mandatory**: `auto` resolves to the + 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). +- `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 + it together with `max_num_tokens`/`max_seq_len` for longer ISL. +- **`kv_cache_bounce_size_mb: 1024` on both sides**: the V2 transceiver's default + pool-to-pool path cannot use inter-node cuda_ipc on MNNVL (the KV pool + is a plain, non-fabric allocation) and falls back to ~0.4 GB/s + host-staged tcp; the fabric-VMM bounce buffer restores cuda_ipc/MNNVL + eligibility (measured ~455 GB/s/GPU). Bounce engages automatically for payloads above `TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES` (default 2 MiB) — always true for K3's ~433 MiB per-request state. + The region must fit ONE request's full payload: fixed 433 MiB KDA + state + ~27 KB/token MLA latent (649 MiB at 8k ISL). A + 512 MiB value makes every 8k transfer fall back to the per-fragment + tcp path (`[kv-bounce] in-place: transfer 649MiB exceeds the 512MiB + bounce region`). + +## Launch sequence (manual, single ctx + single gen) + +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): + + ```bash + trtllm-llmapi-launch trtllm-serve $MODEL_PATH \ + --host --port 8001 \ + --config examples/kimi_k3/disagg/ctx_config.yaml + ``` + +2. Start the generation server: + + ```bash + trtllm-llmapi-launch trtllm-serve $MODEL_PATH \ + --host --port 8002 \ + --config examples/kimi_k3/disagg/gen_config_no_sa.yaml + ``` + +3. Edit `disagg_proxy_config.yaml` (worker URLs = the head nodes above), + then start the proxy: + + ```bash + trtllm-serve disaggregated -c examples/kimi_k3/disagg/disagg_proxy_config.yaml + ``` + +4. Send OpenAI-compatible requests to the proxy (port 8000). + +## SLURM benchmark harness + +`benchmark_kimi_k3_dep16.yaml` drives the full orchestration (worker +config generation, node allocation, proxy, benchmark client): + +```bash +python3 examples/disaggregated/slurm/benchmark/submit.py \ + -c examples/kimi_k3/disagg/benchmark_kimi_k3_dep16.yaml --dry-run # inspect +python3 examples/disaggregated/slurm/benchmark/submit.py \ + -c examples/kimi_k3/disagg/benchmark_kimi_k3_dep16.yaml # submit +``` + +- Set `benchmark.dataset_file` before an e2e submission. +- **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 `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 + the container by `start_worker.sh` / `start_server.sh` / + `run_benchmark.sh`. + +## Current caveats (read before running) + +1. **V2 transceiver "MPI hang" — root-caused, environmental (RESOLVED + with the env pins).** A reported multi-node hang in + `KvCacheTransceiverV2._exchange_rank_info` → `mpi_allgather` was a + downstream symptom of `UCX_TLS=all` on nodes where `ud_verbs` cannot + initialize: the broken transport wedges native NIXL/UCX agent init + asymmetrically per rank, and the healthy ranks park forever in the + 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). +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): 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). +5. **Transfer payload**: each request moves a fixed 433.4 MiB (~454.5 MB) + KDA state blob ctx → gen in addition to the MLA latent KV + (~27 KB/token). + Within an NVL72 domain this is ~0.9 ms/request (measured; not a + bottleneck), but off-fabric paths would pay 11–23 ms — keep ctx and + gen inside one NVL72 domain. +6. **Bounce-buffer sizing cliff (silent).** Size `kv_cache_bounce_size_mb` + to the largest single request's full KV payload (fixed KDA state plus + 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, + 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 + 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). diff --git a/examples/kimi_k3/disagg/benchmark_kimi_k3_dep16.yaml b/examples/kimi_k3/disagg/benchmark_kimi_k3_dep16.yaml new file mode 100644 index 000000000000..bb0db632d279 --- /dev/null +++ b/examples/kimi_k3/disagg/benchmark_kimi_k3_dep16.yaml @@ -0,0 +1,136 @@ +# Kimi K3 disaggregated benchmark-harness config for +# examples/disaggregated/slurm/benchmark/submit.py, modeled on +# examples/wide_ep/slurm_scripts/kimi-k2-thinking.yaml with K3's +# constraints baked in (EP-only DEP16 on both sides, no chunked prefill, +# no KV block reuse, tokens_per_block 64, NIXL + Python V2 transceiver). +# Fill in the <...> placeholders for your cluster before submitting. +# +# Usage (from the repository root): +# python3 examples/disaggregated/slurm/benchmark/submit.py \ +# -c examples/kimi_k3/disagg/benchmark_kimi_k3_dep16.yaml [--dry-run] +# +# Gen-only baseline: set benchmark.mode to "gen_only_no_context" +# (submit.py then exports TRTLLM_DISAGG_BENCHMARK_GEN_ONLY=1). + +slurm: + script_file: "disaggr_torch.slurm" + partition: "" + account: "" + job_time: "02:00:00" + job_name: "k3-disagg-dep16" + extra_args: "--gpus-per-node=4" + numa_bind: true # GB200/GB300 NVL72 + +benchmark: + mode: "e2e" # e2e | gen_only | gen_only_no_context (gen-only baseline) + use_nv_sa_benchmark: false + multi_round: 4 + benchmark_ratio: 0.8 + streaming: true + concurrency_list: "16" + input_length: 8192 + output_length: 1024 + # Point at a benchmark dataset (e.g. one generated by + # benchmarks/cpp/prepare_dataset.py) before submitting an e2e run. + dataset_file: "" + # Benchmark client env: same in-container venv injection as the workers + # (see the environment section notes). + env_var: + TRTLLM_PATH_PREPEND: "" + TRTLLM_PYTHONPATH_PREPEND: "" + +hardware: + gpus_per_node: 4 + num_ctx_servers: 1 + num_gen_servers: 1 + +environment: + # Mount the TRT-LLM checkout and the K3 checkpoint into the container. + container_mount: "" # Format: path1:path1,path2:path2 + container_image: "" + model_path: "" + # TRT-LLM comes from the checkout's in-place install (venv on PATH, + # repo root on PYTHONPATH via worker/server env below); skip the + # harness's pip install. + trtllm_repo: "" + build_wheel: false + trtllm_wheel_path: "" + work_dir: "" + # Cluster notes: + # - On clusters where verbs UCX transports cannot initialize on the + # compute nodes, UCX_TLS must be pinned — UCX_TLS=all can wedge + # native NIXL init asymmetrically, leaving the surviving ranks hung + # in the V2 setup MPI collectives. start_worker.sh clears + # UCX_TLS (a container-provided UCX_TLS=tcp also breaks V2 NIXL VRAM + # registration), so the pin is carried via TRTLLM_WORKER_UCX_TLS and + # re-exported by start_worker.sh after the clear. + # - pairs with the workers' + # kv_cache_bounce_size_mb (fabric-VMM bounce; measured + # ~455 GB/s/GPU). The default gate of 96 blocks would silently skip + # bounce for requests under 6144 tokens at tokens_per_block=64, + # falling back to ~0.4 GB/s host-staged tcp. + # - pyxis/enroot reset image-defined PATH at container start, so the + # in-place venv is injected via TRTLLM_PATH_PREPEND / + # TRTLLM_PYTHONPATH_PREPEND, which the harness scripts apply inside + # the container (setting PATH= here would be silently dropped). + worker_env_var: "TLLM_LOG_LEVEL=INFO TRTLLM_SERVER_DISABLE_GC=1 TRTLLM_WORKER_DISABLE_GC=1 ENROOT_ALLOW_DEV=yes TRTLLM_WORKER_UCX_TLS=tcp,self,sm,cuda_copy,cuda_ipc TRTLLM_PATH_PREPEND= TRTLLM_PYTHONPATH_PREPEND=" + server_env_var: "TRTLLM_SERVER_DISABLE_GC=1 TRTLLM_PATH_PREPEND= TRTLLM_PYTHONPATH_PREPEND=" + +worker_config: + # GEN: DEP16, no speculative decoding, CUDA graphs on (matches + # examples/kimi_k3/disagg/gen_config_no_sa.yaml). + gen: + 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: + max_batch_size: 32 + max_batch_size: 32 + max_num_tokens: 256 + max_seq_len: 9236 + print_iter_log: true + kv_cache_config: + enable_block_reuse: false # mandatory for K3 (KDA state) + tokens_per_block: 64 # mandatory for K3 MLA kernels + free_gpu_memory_fraction: 0.25 + cache_transceiver_config: + backend: NIXL + transceiver_runtime: PYTHON # mandatory: C++ path throws for K3 + max_tokens_in_buffer: 8448 + # Fabric-VMM bounce: needed for cross-node V2 transfers on clusters + # where inter-node cuda_ipc requires fabric allocations (else + # ~0.4 GB/s host-staged tcp); measured ~455 GB/s/GPU with + # (worker_env_var above). + # 1024 fits one full request payload (433 MiB KDA state + ~27 KB/token + # MLA latent = 649 MiB at 8k ISL); at 512 every 8k transfer falls back + # to the per-fragment tcp path. + kv_cache_bounce_size_mb: 1024 + trust_remote_code: true + # CTX: matched DEP16 (see ../disagg/ctx_config.yaml for the sizing + # rationale; hetero ctx/gen parallelism is not supported yet). + ctx: + tensor_parallel_size: 16 + moe_expert_parallel_size: 16 + enable_attention_dp: true + pipeline_parallel_size: 1 + disable_overlap_scheduler: true # mandatory for disagg ctx servers + enable_chunked_prefill: false + cuda_graph_config: null + max_batch_size: 1 + max_num_tokens: 8448 + max_seq_len: 8212 + print_iter_log: true + kv_cache_config: + enable_block_reuse: false + tokens_per_block: 64 + free_gpu_memory_fraction: 0.25 + cache_transceiver_config: + backend: NIXL + transceiver_runtime: PYTHON + max_tokens_in_buffer: 8448 + # Must match the gen side; see the gen block's comment. + kv_cache_bounce_size_mb: 1024 + trust_remote_code: true diff --git a/examples/kimi_k3/disagg/ctx_config.yaml b/examples/kimi_k3/disagg/ctx_config.yaml new file mode 100644 index 000000000000..4be9ff24473b --- /dev/null +++ b/examples/kimi_k3/disagg/ctx_config.yaml @@ -0,0 +1,67 @@ +# Kimi K3 disaggregated serving - CONTEXT (prefill) server extra LLM-API +# options (`trtllm-serve --config ctx_config.yaml`). +# +# Sizing: DEP16 (attention data-parallel, experts sharded across the +# 16-rank EP group). K3 is EP-only (ep_size == tp_size, no PP, no TP on +# linears), so the ctx server must also be a DEP-N deployment. DEP16 is +# the smallest deployment verified to fit (~193 GiB weights/rank on +# GB300); a DEP8 ctx would need ~273 GiB/rank for weights alone +# (extrapolated from the 1.5 TB checkpoint), which does not leave +# activation headroom on GB300 (288 GiB) and cannot fit GB200 (186 GiB). +# Ctx/gen parallelism must match: heterogeneous ctx/gen TP is rejected +# for K3's replicated KDA recurrent state. +# +# Hard K3 constraints baked in: no chunked prefill, no KV block reuse +# (KDA state), tokens_per_block=64 (MLA latent layout), and +# disable_overlap_scheduler=true (mandatory for any disagg ctx server). +# K3 ships a custom (trust-remote-code) tokenizer; without this the +# server starts tokenizer-less and 400s every string prompt with +# "tokenizer is required to tokenize string prompt". +trust_remote_code: true +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 +# Prefill-shaped limits: one 8k-ISL request per attention-DP rank. +# max_num_tokens must cover the full ISL (no chunked prefill for K3). +max_batch_size: 1 +max_num_tokens: 8448 +max_seq_len: 8212 +kv_cache_config: + # Mandatory for K3 (per-request KDA recurrent state). + enable_block_reuse: false + # Mandatory for K3 MLA (576, 512) trtllm-gen kernels. + tokens_per_block: 64 + # Conservative starting point proven by the aggregated DEP16 evals; the + # ctx server only holds in-flight prefills (MLA latent KV is small: + # ~27 KB/token/request), so keep this low to leave prefill activation + # headroom. Tune upward only if KV allocation fails. + free_gpu_memory_fraction: 0.25 +cache_transceiver_config: + backend: NIXL + # K3 requires the Python V2 transceiver: `auto` resolves to the C++ + # runtime, which throws at construction for K3's + # MixedMambaHybridCacheManager (LinearAttentionMetadata not found). + transceiver_runtime: PYTHON + # Must be >= the target max ISL (8192 here). + max_tokens_in_buffer: 8448 + # Fabric-VMM bounce buffer, needed for cross-node transfers on NVLink + # domains: V2 sends pool-to-pool by default, but the KV pool is a + # plain (non-fabric) allocation, so inter-node cuda_ipc/MNNVL is + # impossible and UCX falls back to ~0.4 GB/s host-staged tcp. 512 MiB + # measured at ~455 GB/s/GPU. Engages automatically for payloads above + # TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES (default 2 MiB) — always true for + # K3's ~433 MiB per-request recurrent state; no env override needed. + # + # Sizing: the region must fit ONE request's full payload = fixed KDA + # state (433 MiB) + MLA latent for the ISL (~27 KB/token, ~216 MiB at + # 8k). At 512 every 8k request is 649 MiB and silently falls back to + # the per-fragment tcp path ("[kv-bounce] in-place: transfer 649MiB + # exceeds the 512MiB bounce region"). 1024 + # covers the 8448-token max ISL with headroom. + kv_cache_bounce_size_mb: 1024 +# No speculative_config on the ctx server: SA drafting runs on the +# generation side only (ctx produces the prompt KV + first token). diff --git a/examples/kimi_k3/disagg/disagg_proxy_config.yaml b/examples/kimi_k3/disagg/disagg_proxy_config.yaml new file mode 100644 index 000000000000..b8e83983e694 --- /dev/null +++ b/examples/kimi_k3/disagg/disagg_proxy_config.yaml @@ -0,0 +1,23 @@ +# Kimi K3 disaggregated serving - proxy configuration for +# `trtllm-serve disaggregated -c disagg_proxy_config.yaml`. +# +# Modeled on examples/disaggregated/disagg_config.yaml. The proxy routes +# each request to a context server for prefill, then to a generation +# server for decode; the KV + KDA state moves ctx -> gen through the +# cache transceiver configured in the worker configs (NIXL, Python V2 +# runtime - see ctx_config.yaml / gen_config_no_sa.yaml). +# +# Replace hostnames/ports with the actual worker endpoints (each K3 +# worker spans 16 GPUs = 4 NVL72 nodes; the URL is rank 0's node). +hostname: localhost +port: 8000 +model: +backend: "pytorch" +context_servers: + num_instances: 1 + urls: + - "localhost:8001" +generation_servers: + num_instances: 1 + urls: + - "localhost:8002" diff --git a/examples/kimi_k3/disagg/gen_config_no_sa.yaml b/examples/kimi_k3/disagg/gen_config_no_sa.yaml new file mode 100644 index 000000000000..ccb1f7812a2f --- /dev/null +++ b/examples/kimi_k3/disagg/gen_config_no_sa.yaml @@ -0,0 +1,50 @@ +# Kimi K3 disaggregated serving - GENERATION (decode) server extra LLM-API +# options WITHOUT speculative decoding +# (`trtllm-serve --config gen_config_no_sa.yaml`). +# +# Default: CUDA graphs ON. Validated end-to-end: GSM8K 96.89 (in the +# aggregated band) and 765/2138 out tok/s at c64/c256 — ~3x the eager +# TPOT and above aggregated serving throughput despite the overlap +# scheduler being off (a disagg requirement on the ctx side). +# For strict token-parity debugging (attribution-clean SA-on/SA-off or +# disagg-vs-agg comparisons), null out the cuda graph config to match +# the SA config's eager regime, and place ctx and gen in SEPARATE NVL72 +# domains (co-located placement selects different comm algorithms and +# diverges token-level on a fixed prompt subset; accuracy-neutral). +# K3 ships a custom (trust-remote-code) tokenizer; without this the +# server starts tokenizer-less and 400s every string prompt with +# "tokenizer is required to tokenize string prompt". +trust_remote_code: true +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: + max_batch_size: 32 +# No SA SpeculativeState buffer cap here; 32 matches the aggregated +# non-SA DEP16 eval config. +max_batch_size: 32 +max_num_tokens: 256 +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 + 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 + 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. + # at tokens_per_block=64). See ctx_config.yaml. + # 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 diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/config.py b/tensorrt_llm/_torch/disaggregation/native/bounce/config.py index cfdcd4e45557..ae6691ce6da6 100644 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/config.py +++ b/tensorrt_llm/_torch/disaggregation/native/bounce/config.py @@ -23,26 +23,25 @@ _MIB = 1024 * 1024 -# Test/advanced override for the block-count gate below (users only tune the bounce size). Read on -# the generation side, so set it there; unset uses the default. -_MIN_BLOCKS_ENV = "TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS" +# Test/advanced overrides for the size gates below (users only tune the bounce size). Read on the +# generation side, so set them there; unset uses the defaults. +_MIN_BYTES_ENV = "TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES" # byte gate for recurrent-state payloads +_MIN_BLOCKS_ENV = "TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS" # block-count gate for plain-KV payloads -def _env_min_blocks(default: int) -> int: - """Read the gate from the env, defensively: unset or malformed falls back to the default (never +def _env_int_gate(name: str, default: int) -> int: + """Read a gate from the env, defensively: unset or malformed falls back to the default (never crashing), and the value is clamped to at least 1.""" - raw = os.environ.get(_MIN_BLOCKS_ENV) + raw = os.environ.get(name) if raw is None: return default try: value = int(raw) except ValueError: - logger.warning(f"{_MIN_BLOCKS_ENV}={raw!r} is not an integer; using default {default}") + logger.warning(f"{name}={raw!r} is not an integer; using default {default}") return default if value < 1: - logger.warning( - f"{_MIN_BLOCKS_ENV}={value} < 1; clamping to 1 (bounce always clears the gate)" - ) + logger.warning(f"{name}={value} < 1; clamping to 1 (bounce always clears the gate)") return 1 return value @@ -103,20 +102,45 @@ def fit_within_free( return capacity_bytes +# Byte gate for transfers that carry recurrent (mamba/KDA) state. The cost this gate guards scales +# with BYTES, not blocks: the block-count gate (96 blocks, calibrated for 128-token blocks) +# silently skipped a 433 MiB Kimi-K3 transfer (67 blocks of 32 tokens plus the non-paged KDA state) +# and dropped it onto the ~0.4 GB/s host-staged fallback, a ~1000x cliff. Break-even is small: +# bounce adds one gather plus one scatter copy (device-local, ~hundreds of GB/s) and ~0.1 ms of +# fixed launch/reservation overhead, while the in-place path can be as slow as ~0.4 GB/s +# inter-node — 2 MiB in-place at that rate is ~5 ms vs well under 1 ms bounced. Below 2 MiB the +# fixed overhead dominates and arena slots are better kept for large transfers. Heuristic, tunable +# via TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES. +DEFAULT_MIN_BYTES = 2 * _MIB + + @dataclass class Config: sizing: Sizing = field(default_factory=FixedSizing) # how much memory to reserve (pluggable) chunk_mb: int = 32 # physical chunk size; a large chunk keeps the write to a single descriptor - # skip bounce below this many blocks (roughly 12k tokens at 128 per block); heuristic, tunable + # Which gate applies depends on the payload (see VmmBounceTransport.reserve): transfers that + # carry recurrent state use min_bytes; plain-KV transfers keep the original min_blocks gate so + # pre-existing bounce deployments see no behavior change. + # byte gate for recurrent-state payloads (see DEFAULT_MIN_BYTES for the rationale) + min_bytes: int = DEFAULT_MIN_BYTES + # block-count gate for plain-KV payloads (roughly 12k tokens at 128 per block); heuristic, + # tunable via TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS min_blocks: int = 96 -def config_from_size(size_mb: int, min_blocks: Optional[int] = None) -> Optional[Config]: +def config_from_size( + size_mb: int, min_blocks: Optional[int] = None, min_bytes: Optional[int] = None +) -> Optional[Config]: """Build a bounce config from a per-region size in MiB, or None to leave bounce off (size <= 0). - Size is both the capacity and the on/off switch. min_blocks is the gate below which a transfer - stays on the per-block path; when unset it comes from the env, else the default.""" + Size is both the capacity and the on/off switch. min_bytes (recurrent-state payloads) and + min_blocks (plain-KV payloads) are the gates below which a transfer stays on the per-block + path; when unset they come from the env, else the defaults.""" if size_mb is None or size_mb <= 0: return None if min_blocks is None: - min_blocks = _env_min_blocks(Config.min_blocks) - return Config(sizing=FixedSizing(capacity_mb=size_mb), min_blocks=min_blocks) + min_blocks = _env_int_gate(_MIN_BLOCKS_ENV, Config.min_blocks) + if min_bytes is None: + min_bytes = _env_int_gate(_MIN_BYTES_ENV, Config.min_bytes) + return Config( + sizing=FixedSizing(capacity_mb=size_mb), min_bytes=min_bytes, min_blocks=min_blocks + ) diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/core.py b/tensorrt_llm/_torch/disaggregation/native/bounce/core.py index 1878fa90fd3a..0eb621d5d189 100644 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/core.py +++ b/tensorrt_llm/_torch/disaggregation/native/bounce/core.py @@ -192,8 +192,17 @@ def release_send(self, slot_id) -> None: """Release a send region after its write completes.""" @abstractmethod - def reserve(self, recv_req, num_writers: int = 1, *, timeout: Optional[float] = None) -> bool: - """Reserve a region and record its address for the senders. False falls back to per-fragment.""" + def reserve( + self, + recv_req, + num_writers: int = 1, + *, + timeout: Optional[float] = None, + extra_bytes: int = 0, + ) -> bool: + """Reserve a region and record its address for the senders. False falls back to per-fragment. + ``extra_bytes`` sizes the sender's non-paged payload (mamba/KDA recurrent state) that rides + the same coalesced write as the KV blocks.""" @abstractmethod def writer_base(self, rid_slice, writer_index: int) -> Optional[int]: diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py b/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py index a9597b7cbb2a..5d88485ec8b1 100644 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py +++ b/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py @@ -37,7 +37,7 @@ from tensorrt_llm._utils import CUASSERT from .buffer import SlotAllocator -from .config import SizingContext, fit_within_free +from .config import DEFAULT_MIN_BYTES, SizingContext, fit_within_free from .core import BounceTransport, Disposition, Settlement, TransferContext from .gather_scatter import Plan, gather_contiguous, scatter_contiguous @@ -83,6 +83,7 @@ def from_config( capacity_bytes=capacity_bytes, phys_chunk_size=chunk, block_bytes_per_group=block_bytes_per_group, + min_bytes=cfg.min_bytes, min_blocks=cfg.min_blocks, ) @@ -94,6 +95,7 @@ def __init__( capacity_bytes: int, phys_chunk_size: int, block_bytes_per_group: List[int], + min_bytes: int = DEFAULT_MIN_BYTES, min_blocks: int = 96, quarantine_grace_s: float = _QUARANTINE_GRACE_S, name: str = "kv_bounce", @@ -102,8 +104,11 @@ def __init__( self._device_id = device_id # The byte size of one cache block, listed for each attention layer group. self._block_bytes_per_group = list(block_bytes_per_group) - # Below this many blocks, skip bounce: coalescing only pays off for long context (the default - # is roughly twelve thousand tokens; a heuristic, and tunable). + # Size gates below which bounce is skipped: coalescing only pays off once the transfer is + # large enough to beat the gather+scatter overhead (see config.DEFAULT_MIN_BYTES for the + # rationale). min_bytes applies to payloads carrying recurrent (mamba/KDA) state; + # min_blocks applies to plain-KV payloads (see the gate in reserve()). + self._min_bytes = min_bytes self._min_blocks = min_blocks # how long an orphaned region is held out of reuse; must outlast the worst in-flight write self._quarantine_grace_s = quarantine_grace_s @@ -182,9 +187,10 @@ def _reserve_and_gather(self, write_meta, *, timeout): total = int(write_meta.sizes.sum()) res = self._send_alloc.reserve(total, timeout=timeout) if res is None: - logger.debug( + logger.warning_once( f"[kv-bounce] in-place: no send region space for {total // _MIB}MiB within {timeout}s " - f"(sender backpressure); falling back" + f"(sender backpressure); falling back", + key="kv-bounce-send-backpressure", ) return None slot_id, src_addr = res @@ -209,29 +215,79 @@ def release_send(self, slot_id) -> None: self._send_alloc.release(slot_id) @staticmethod - def _skip_bounce(reason: str, *, warn_key: Optional[str] = None) -> bool: + def _skip_bounce(reason: str, *, warn_key: str) -> bool: """Log why a transfer falls back to the per-fragment path and return False, so the guards - above stay one line each.""" - msg = f"[kv-bounce] in-place: {reason}" - logger.warning_once(msg, key=warn_key) if warn_key else logger.debug(msg) + above stay one line each. Every reason logs at warning once per key: silently skipping + bounce can be a ~1000x bandwidth cliff (host-staged tcp vs cuda_ipc), so the first skip per + distinct reason must be visible at the default log level.""" + logger.warning_once(f"[kv-bounce] in-place: {reason}", key=warn_key) return False def reserve( - self, recv_req, num_writers: int = 1, *, timeout: Optional[float] = _RESERVE_TIMEOUT_S + self, + recv_req, + num_writers: int = 1, + *, + timeout: Optional[float] = _RESERVE_TIMEOUT_S, + extra_bytes: int = 0, ) -> bool: """Reserve a region and create its state, recording the address for the senders. Returns False to fall back to the per-fragment path. A fan-in splits the region evenly, so the total - must divide across the writers.""" - nblocks = sum(int(a.size) for a in recv_req.block_ids_per_layer_groups) - if nblocks < self._min_blocks: - return self._skip_bounce(f"{nblocks} blocks < min {self._min_blocks} (too small)") + must divide across the writers. ``extra_bytes`` is the non-paged payload the sender appends + to the same coalesced write (mamba/KDA recurrent state, sized by the receiver via + ``MambaPolicy.payload_bytes``); the region must cover it or the write would overrun into the + neighboring slot.""" total = 0 for g, block_ids in enumerate(recv_req.block_ids_per_layer_groups): - if g >= len(self._block_bytes_per_group): - return self._skip_bounce(f"layer group {g} has no known slot size (e.g. mamba)") + if int(block_ids.size) == 0: + # Nothing to transfer for this group. Hybrid models (Kimi K3: KDA + MLA) always + # carry a trailing empty entry for the mamba layer group — its state is not paged + # and rides extra_bytes instead — so an empty group must not disable bounce. + continue + known = g < len(self._block_bytes_per_group) and self._block_bytes_per_group[g] + if not known: + return self._skip_bounce( + f"layer group {g} has blocks but no known slot size", + warn_key="kv-bounce-unknown-slot-size", + ) total += int(block_ids.size) * self._block_bytes_per_group[g] + if extra_bytes > 0 and num_writers > 1: + # Each fan-in writer appends its own recurrent-state fragments, whose sizes may differ + # per writer (PP stages hold different mamba layers), breaking the equal region split. + return self._skip_bounce( + f"fan-in across {num_writers} senders with {extra_bytes}B of recurrent state; the " + f"equal split cannot account for per-writer state fragments", + warn_key="kv-bounce-mamba-fanin", + ) + total += int(extra_bytes) if total <= 0: - return self._skip_bounce(f"computed transfer size {total} <= 0") + return self._skip_bounce( + f"computed transfer size {total} <= 0", warn_key="kv-bounce-nonpositive-size" + ) + # Which size gate applies depends on the payload. Payloads carrying recurrent (mamba/KDA) + # state gate on BYTES: the cost the gate guards (falling back to the slow per-fragment + # path) scales with bytes, and a block count is meaningless for the non-paged state (the + # Kimi-K3 regression: a 433 MiB transfer of 67 small blocks plus KDA state failed a + # 96-block gate and fell onto the ~0.4 GB/s host-staged path). Plain-KV payloads keep the + # original block-count gate so pre-existing bounce deployments (opted in via + # kv_cache_bounce_size_mb) see no change in which transfers use the arena. + # TODO(TRTLLM-15194): investigate whether the byte-only gate is + # safe (or better) for plain-KV payloads too, so this special case can be removed and + # both payload kinds share one gate. + nblocks = sum(int(a.size) for a in recv_req.block_ids_per_layer_groups) + if extra_bytes > 0: + if total < self._min_bytes: + return self._skip_bounce( + f"{total}B ({nblocks} blocks + recurrent state) < min {self._min_bytes}B " + f"(too small; tune TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES)", + warn_key="kv-bounce-below-min-bytes", + ) + elif nblocks < self._min_blocks: + return self._skip_bounce( + f"{nblocks} blocks < min {self._min_blocks} (too small; tune " + f"TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS)", + warn_key="kv-bounce-below-min-blocks", + ) if num_writers > 1 and total % num_writers != 0: return self._skip_bounce( f"fan-in {total}B across {num_writers} senders is not an even split " @@ -262,7 +318,8 @@ def reserve( res = self._recv_alloc.reserve(total, timeout=timeout) if res is None: return self._skip_bounce( - f"no recv region space for {total // _MIB}MiB within {timeout}s (backpressure)" + f"no recv region space for {total // _MIB}MiB within {timeout}s (backpressure)", + warn_key="kv-bounce-recv-backpressure", ) slot_id, addr = res recv_req.bounce_dst_base = addr @@ -436,7 +493,12 @@ def release_send(self, slot_id) -> None: pass def reserve( - self, recv_req, num_writers: int = 1, *, timeout: Optional[float] = _RESERVE_TIMEOUT_S + self, + recv_req, + num_writers: int = 1, + *, + timeout: Optional[float] = _RESERVE_TIMEOUT_S, + extra_bytes: int = 0, ) -> bool: return False @@ -526,8 +588,12 @@ def decode_result_tail(message): def block_bytes_per_group(page_table) -> list: - """Byte size of one cache block for each leading attention layer group, stopping at the first - non-attention group.""" + """Byte size of one cache block for each layer group, aligned with the layer-group indices a + recv request uses. Non-attention groups (mamba/KDA recurrent state) hold ``None``: they carry + no paged blocks (their KVSlice entry is always empty — see ``_create_kv_slice``) and their + payload is sized separately via ``MambaPolicy.payload_bytes``. Keeping them as placeholders + instead of truncating means a trailing (or hypothetically interleaved) mamba group can never + shift an attention group off the end of this list and poison the bounce gate.""" from tensorrt_llm._torch.disaggregation.resource.page import AttentionLayerGroup from tensorrt_llm._torch.disaggregation.resource.utils import get_physical_pool @@ -535,6 +601,7 @@ def block_bytes_per_group(page_table) -> list: out: list = [] for lg_idx, lg in enumerate(page_table.layer_groups): if not isinstance(lg, AttentionLayerGroup): - break + out.append(None) + continue out.append(int(get_physical_pool(page_table, lg_idx, 0).slot_bytes)) return out diff --git a/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py b/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py index a96fc3e8f496..f7952c29ad1a 100644 --- a/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py @@ -335,6 +335,106 @@ class MambaPolicy: Dispatch mappers and build frags for Mamba state transfer. """ + @staticmethod + def _find_mamba_layer_group( + page_table: Optional[KVCachePageTable], + ) -> Optional[MambaLayerGroup]: + if page_table is None: + return None + return next( + (lg for lg in page_table.layer_groups if isinstance(lg, MambaLayerGroup)), + None, + ) + + @staticmethod + def validate_peer_compatible( + self_ri: RankInfo, + peer_ri: RankInfo, + self_page_table: Optional[KVCachePageTable], + peer_page_table: Optional[KVCachePageTable], + ) -> None: + """Validate recurrent-state (Mamba/KDA) layout compatibility with a peer. + + Analogue of the C++ ``rnnCacheFormatter inquireSupport`` gate for the + V2 path: reject at peer-registration time instead of corrupting + memory at transfer time. Raises ``ValueError`` naming the mismatched + field. + + The core invariant checked is *global* (TP-aggregated) state size: + for a TP-sharded state, ``per_rank_bytes * mamba_tp`` is + TP-invariant, so it must match between peers even when their TP + sizes differ. Models with *replicated* recurrent state (e.g. Kimi K3 + KDA, whose per-rank state is pre-scaled to full size when + attention-DP is off) violate this invariant under heterogeneous TP — + exactly the configuration where the TP-mismatch mappers would + compute shard offsets past the end of the slot, silently corrupting + the replicated state — and are therefore rejected here. + """ + self_mlg = MambaPolicy._find_mamba_layer_group(self_page_table) + peer_mlg = MambaPolicy._find_mamba_layer_group(peer_page_table) + if self_mlg is None or peer_mlg is None: + # Under pipeline parallelism each rank publishes only its own + # stage's layers, so a hybrid model can pair a rank holding + # recurrent layers with a peer stage holding none. The transfer + # path (build_mamba_frags) intersects the two layer sets and + # moves nothing when either side has no recurrent layers, so + # there is nothing to validate for this pair. + return + + # Layer sets are NOT required to match: with pipeline parallelism the + # two sides may partition layers differently, and the transfer covers + # exactly the intersection (empty intersection moves no state). The + # invariants below are per-layer-slot quantities, uniform across a + # model's recurrent layers, so they apply regardless of which layers + # overlap. + if ( + self_mlg.ssm_bytes_per_head is not None + and peer_mlg.ssm_bytes_per_head is not None + and self_mlg.ssm_bytes_per_head != peer_mlg.ssm_bytes_per_head + ): + # TP-invariant: head_dim * d_state * element_size. A mismatch + # means different state shape or SSM cache dtype. + raise ValueError( + "MambaPolicy.validate_peer_compatible: ssm_bytes_per_head differs " + f"(local={self_mlg.ssm_bytes_per_head}, peer={peer_mlg.ssm_bytes_per_head}); " + "check head_dim / d_state / mamba_ssm_cache_dtype" + ) + + self_tp, _ = MambaPolicy._mamba_tp(self_ri) + peer_tp, _ = MambaPolicy._mamba_tp(peer_ri) + + def _check_global(field: str, self_bytes: int, peer_bytes: int) -> None: + if self_bytes * self_tp != peer_bytes * peer_tp: + raise ValueError( + f"MambaPolicy.validate_peer_compatible: global (TP-aggregated) {field} " + f"differs: local {self_bytes} bytes/rank x mamba_tp={self_tp} vs " + f"peer {peer_bytes} bytes/rank x mamba_tp={peer_tp}. Per-rank state " + "sizes are inconsistent with a TP-sharded layout across the two " + "sides; either the state shape/dtype differs, or the model keeps a " + "replicated (non-TP-sharded) recurrent state (e.g. Kimi K3 KDA), " + "which supports heterogeneous ctx/gen TP only with attention-DP " + "enabled on both sides." + ) + + _check_global( + "ssm slot_bytes", self_mlg.ssm_states.slot_bytes, peer_mlg.ssm_states.slot_bytes + ) + _check_global( + "conv slot_bytes", self_mlg.conv_states.slot_bytes, peer_mlg.conv_states.slot_bytes + ) + + if self_mlg.conv_section_bytes is not None and peer_mlg.conv_section_bytes is not None: + if len(self_mlg.conv_section_bytes) != len(peer_mlg.conv_section_bytes): + raise ValueError( + "MambaPolicy.validate_peer_compatible: conv section count differs " + f"(local={len(self_mlg.conv_section_bytes)}, " + f"peer={len(peer_mlg.conv_section_bytes)})" + ) + for i, (s, p) in enumerate( + zip(self_mlg.conv_section_bytes, peer_mlg.conv_section_bytes) + ): + _check_global(f"conv_section_bytes[{i}]", s, p) + @staticmethod def _mamba_tp(ri: RankInfo) -> Tuple[int, int]: """Return (mamba_effective_tp_size, mamba_effective_tp_rank). @@ -492,6 +592,35 @@ def build_mamba_frags( return src_frags, dst_frags, kv_sizes + @staticmethod + def payload_bytes( + sender_page_table: KVCachePageTable, + receiver_page_table: KVCachePageTable, + dst_slot: Optional[int], + sender_ri: RankInfo, + receiver_ri: RankInfo, + ) -> int: + """Total recurrent-state bytes one sender appends to the KV write for a request. + + Mirrors the sender's ``collect_frags`` call in ``_build_kv_write_meta`` with the same + argument roles (sender as self/src), so the receiver can size a bounce region for the + exact bytes the coalesced write will carry. Slot indices only shift pointers, never + fragment sizes, so a dummy ``src_slot`` stands in for the sender's state index (which + the receiver does not know at reserve time). Returns 0 when either side has no mamba + layer group or the receiver has no state slot for the request. + """ + if dst_slot is None: + return 0 + _, _, sizes = MambaPolicy.collect_frags( + self_page_table=sender_page_table, + peer_page_table=receiver_page_table, + src_slot=0, + dst_slot=dst_slot, + self_ri=sender_ri, + peer_ri=receiver_ri, + ) + return int(sum(sizes)) + @staticmethod def collect_frags( self_page_table: KVCachePageTable, diff --git a/tensorrt_llm/_torch/disaggregation/native/peer.py b/tensorrt_llm/_torch/disaggregation/native/peer.py index 6b8f40232f1d..485432787d0d 100644 --- a/tensorrt_llm/_torch/disaggregation/native/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/peer.py @@ -23,6 +23,7 @@ from tensorrt_llm._torch.disaggregation.base.region import RegionMapperBase from tensorrt_llm._torch.disaggregation.native.auxiliary import AuxTransferLayout from tensorrt_llm._torch.disaggregation.native.mixers.attention.peer import AttentionPolicy +from tensorrt_llm._torch.disaggregation.native.mixers.ssm.peer import MambaPolicy from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 from tensorrt_llm._torch.disaggregation.resource.page import ( @@ -152,6 +153,16 @@ def _check_peer_compatible(self, peer_ri: RankInfo) -> bool: if not self._attention_policy.check_peer_compatible(peer_ri): return False + # Recurrent-state (Mamba/KDA) layout gate. Raises ValueError with a + # field-level diagnostic instead of returning False, so the precise + # mismatch reaches the caller of register(). + MambaPolicy.validate_peer_compatible( + self._ri, + peer_ri, + self._self_ext_cache.page_table if self._self_ext_cache is not None else None, + peer_ri.page_table, + ) + self_layers = sum(self._ri.layer_num_per_pp) peer_layers = sum(peer_ri.layer_num_per_pp) if self_layers != peer_layers: diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 4f7fc7edf19e..673d088e8443 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -1659,7 +1659,24 @@ def dispatch_task(self, task: KVRecvTask): allow_bounce = task.expected_transfers == 1 or ( sender_dp_rank is not None and self._fanin_bounce_safe(topo_overlap, peer_infos) ) - bounced = allow_bounce and self._bounce.reserve(receiver_req, task.expected_transfers) + # Recurrent (mamba/KDA) state rides the SAME coalesced write as the KV blocks (the sender + # appends its MambaPolicy fragments in _build_kv_write_meta), so the bounce region must be + # sized for those bytes too or the write would overrun into the neighboring slot. + extra_bytes = 0 + if receiver_req.mamba_state_index is not None: + if peer_infos.page_table is None: + allow_bounce = False # cannot size the sender's recurrent-state payload + else: + extra_bytes = MambaPolicy.payload_bytes( + sender_page_table=peer_infos.page_table, + receiver_page_table=self._registrar.self_extractor.page_table, + dst_slot=receiver_req.mamba_state_index, + sender_ri=peer_infos, + receiver_ri=self._registrar.self_rank_info, + ) + bounced = allow_bounce and self._bounce.reserve( + receiver_req, task.expected_transfers, extra_bytes=extra_bytes + ) session = self._get_session(task._unique_rid) if session is None: raise RuntimeError( @@ -1716,6 +1733,19 @@ def _get_sender_info(self, params: DisaggregatedParams) -> RankInfo: finally: messenger.stop() + # Recurrent-state (Mamba/KDA) layout gate on the receiver side. + # The sender-side check (PeerRegistrar.register) runs in the + # sender's listener thread, where exceptions are only logged, so + # reject here — before REGISTER_RANK_INFO is even sent — to fail + # the first gen request loudly instead of hanging on a transfer + # the sender will never serve. + MambaPolicy.validate_peer_compatible( + self._registrar.self_rank_info, + sender_info, + self._registrar.self_extractor.page_table, + sender_info.page_table, + ) + for endpoint in sender_info.sender_endpoints: dealer = self._get_or_connect_dealer(endpoint) rank_info = self._registrar.self_rank_info diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index 6968159a90d7..0d7e3429424b 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -94,7 +94,18 @@ def __init__( self._device_id = torch.cuda.current_device() logger.info(f"device_id: {self._device_id} in KvCacheTransceiverV2") + # Setup interleaves MPI collectives (broadcast/allgather below) with + # per-rank native NIXL/UCX initialization inside TransferWorker. A rank + # that blocks or dies in the native phase leaves its peers stuck in the + # next collective, so log each phase per rank to make the blocking + # rank/phase identifiable from the logs (see the 'setup:' lines). + rank = self._dist.rank + logger.info(f"KvCacheTransceiverV2 setup: rank={rank} broadcast instance name (collective)") self._instance_name = self._broadcast_instance_name() + logger.info( + f"KvCacheTransceiverV2 setup: rank={rank} creating TransferWorker " + "(native NIXL agent init + KV memory registration)" + ) self._transfer_worker = TransferWorker( TransferWorkerConfig( kv_cache_manager=kv_cache_manager, @@ -106,14 +117,22 @@ def __init__( max_concurrent_sessions=max(1, int(kv_cache_manager.max_batch_size)) * 20000, tx_timeout_s=self._sender_future_timeout_ms / 1000.0, rx_timeout_s=self.kv_transfer_timeout_ms / 1000.0, - # Size 0 turns bounce off; the block-count gate is internal (tuned via env). + # Size 0 turns bounce off; the per-transfer size gates are internal (tuned via + # env: TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS for plain-KV payloads, + # TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES for recurrent-state payloads). bounce=bounce_config_from_size(cache_transceiver_config.kv_cache_bounce_size_mb), ) ) + logger.info( + f"KvCacheTransceiverV2 setup: rank={rank} TransferWorker ready; " + "broadcast context endpoint (collective)" + ) self._dp_rank = mapping.tp_rank if mapping.enable_attention_dp else 0 self._context_info_endpoint = self._broadcast_context_endpoint() self._init_sync_policy() + logger.info(f"KvCacheTransceiverV2 setup: rank={rank} exchange rank info (collective)") self._exchange_rank_info() + logger.info(f"KvCacheTransceiverV2 setup: rank={rank} complete") self._send_sessions: Dict[int, TxSessionBase] = {} self._recv_sessions: Dict[int, RxSessionBase] = {} @@ -317,10 +336,21 @@ def _slice_num_bytes(self, slice: KVSlice) -> int: return 0 total = 0 for lg_id, block_ids in enumerate(slice.block_ids_per_layer_groups): - if block_ids is None or block_ids.size == 0: - continue lg = pt.layer_groups[lg_id] if isinstance(lg, MambaLayerGroup): + # Fixed-size recurrent state (mamba/KDA): one slot per layer in + # each of the conv and ssm pools, independent of token count. + # For hybrid models (e.g. Kimi K3) this blob can dominate + # short-prompt transfers, so it must be counted. The caller's + # tp_size scaling then yields total bytes moved across ranks + # (exact for sharded state; for replicated state every rank + # pair moves a full copy, so it matches bytes on the wire). + if slice.mamba_state_index is not None: + total += len(lg.mamba_layer_offsets) * ( + lg.conv_states.slot_bytes + lg.ssm_states.slot_bytes + ) + continue + if block_ids is None or block_ids.size == 0: continue n = int((block_ids >= 0).sum()) if n == 0: diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index dc7b753a89ef..4c04d7a9e507 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -163,14 +163,11 @@ def get_kv_cache_manager_cls( # Mixed manager (separate KV / recurrent-state pools) stays the # default. SA speculative decoding is validated on the Mixed # manager's SpeculativeState scratch path only; reuse + SA is - # unvalidated. - if is_kimi_linear(config) and not use_v2: - if is_disagg: - # Fail fast instead of bypassing the disagg transceiver - # validation below with an unvalidated route. - raise NotImplementedError( - "Disaggregated serving is not supported for Kimi K3 yet " - "(TRTLLM-14815).") + # unvalidated. Disaggregated serving (TRTLLM-14815) routes through + # the shared hybrid transceiver validation below: the Python NIXL + # transceiver selects the Mixed manager, whose KDA recurrent/conv + # states transfer through the bounce buffer. + if is_kimi_linear(config) and not use_v2 and not is_disagg: if kv_cache_config.enable_block_reuse: logger.info( "Using CppMambaHybridCacheManager for Kimi K3 hybrid " @@ -591,6 +588,10 @@ def _get_model_kv_cache_manager_cls( cache_transceiver_config=self._cache_transceiver_config) cls = self._fallback_if_unsupported_kv_cache_manager_v2( cls, model_config, kv_cache_config) + if is_hybrid_linear(model_config.pretrained_config): + logger.info_once( + f"Selected hybrid KV cache manager: {cls.__name__}", + key=f"hybrid_kv_cache_manager_{cls.__name__}") # Compatibility managers do not support MTP block reuse. Warn at the # routing site so users see the concrete manager selected for the # incompatible combination. diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_overlap_transceiver_runtime_python_bounce.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_overlap_transceiver_runtime_python_bounce.yaml index c6573753cbea..b4b3cd1234b9 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_overlap_transceiver_runtime_python_bounce.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_overlap_transceiver_runtime_python_bounce.yaml @@ -3,7 +3,7 @@ # Same as disagg_config_overlap_transceiver_runtime_python.yaml, plus the bounce switch on both the # context (sender) and generation (receiver) cache_transceiver_config: # kv_cache_bounce_size_mb: >0 turns bounce on and sizes the per-region fabric-VMM arena. -# The block-count gate below which a transfer keeps the per-block path is lowered to 1 via the +# Plain-KV transfers gate on block count; the gate is lowered to 1 via the # TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS env (set by the test) so the ordinary short test prompts still # take the coalesced-bounce WRITE path (the production default of 96 would need a ~2k-token prompt). # GB200/GB300 only, since the bounce arena is fabric (MNNVL) VMM memory. diff --git a/tests/integration/defs/kimi_k3_disagg_parity.py b/tests/integration/defs/kimi_k3_disagg_parity.py new file mode 100644 index 000000000000..f327899870a2 --- /dev/null +++ b/tests/integration/defs/kimi_k3_disagg_parity.py @@ -0,0 +1,792 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 two-endpoint parity harness (aggregated vs disaggregated). + +Compares a REFERENCE deployment against a CANDIDATE deployment through +their OpenAI-compatible ``/v1/completions`` endpoints and produces a +parity report: + +* per-prompt first-token agreement (in disagg the first token is produced + by the ctx server and handed off — a mismatch points at the + KV/KDA-state transfer or the first-token handoff), +* longest-common-prefix (LCP) length of the greedy token streams, +* logprob drift along the shared prefix + near-tie classification at the + first divergence (reuses ``_compare_logits_parity`` from + ``kimi_k3_sa_harness.py`` — same tolerances, same aggregate + "non-ties must not dominate" rule), +* optional GSM8K accuracy diff (runs ``lm_eval --model local-completions`` + against both endpoints, the same flow as + ``examples/disaggregated/slurm/benchmark/submit.py``). + +The harness only needs HTTP access to the two endpoints — no GPUs, no +tensorrt_llm runtime — so it runs from a login node against servers +launched elsewhere. + +DRY-RUN RECIPE (two aggregated servers; debugs the harness before disagg +transfer works — expect near-perfect parity): + + # server A (reference) and server B (candidate): identical aggregated + # deployments of the same checkpoint on two ports, e.g. + trtllm-serve $KIMI_K3_CKPT --backend pytorch --port 8000 \ + --config examples/kimi_k3/eval_extra_llm_options.yaml & + trtllm-serve $KIMI_K3_CKPT --backend pytorch --port 8001 \ + --config examples/kimi_k3/eval_extra_llm_options.yaml & + + python tests/integration/defs/kimi_k3_disagg_parity.py \ + --reference http://localhost:8000 --candidate http://localhost:8001 \ + --extra-prompts 16 --report-json parity_dryrun.json + +REAL RECIPE (aggregated vs disagg proxy): + + # reference: aggregated DEP16 deployment (as in + # examples/kimi_k3/run_gsm8k_kimi_k3.sbatch, but served) + trtllm-serve $KIMI_K3_CKPT --backend pytorch --port 8000 ... + + # candidate: ctx + gen workers behind the disagg proxy + trtllm-serve disaggregated -c disagg_config.yaml # port 9000 + + python tests/integration/defs/kimi_k3_disagg_parity.py \ + --reference http://ref-host:8000 --candidate http://proxy-host:9000 \ + --extra-prompts 48 --gsm8k --gsm8k-limit 200 \ + --report-json parity_disagg.json + +Notes: +* Requests are sent sequentially (batch of 1) to both endpoints so batch + composition cannot perturb the comparison. +* Per-token logprobs are requested via the completions ``logprobs`` field + (PyTorch backend supports top-k dicts). Endpoints that reject + ``detokenize=false`` or ``logprobs`` degrade gracefully: the harness + falls back to decoded-token-string comparison, then to text-only, and + reports which capability level each endpoint provided. One-engine + spec-dec samplers (SA) do not emit per-token logprobs — run the parity + probe with spec dec off first, or accept token-level-only + parity for SA-on runs. +* GSM8K mode shells out to ``lm_eval`` (must be installed, e.g. + ``pip install lm-eval[api]``) and diffs the exact_match metrics between + the two endpoints (default tolerance 0.02). + +Self-test (no servers, canned responses — exercises the comparison and +classification logic): + + python tests/integration/defs/kimi_k3_disagg_parity.py --self-test + +Exit code 0 = PASS, 1 = FAIL (any parity failure, endpoint error, or +GSM8K delta above tolerance). +""" + +import argparse +import glob +import json +import os +import subprocess +import sys +import tempfile +import urllib.error +import urllib.request +from types import SimpleNamespace + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from kimi_k3_sa_harness import ( # noqa: E402 + PROMPTS_AND_CHECKS, + _compare_logits_parity, + _parity_prompts, +) + +# Request-payload variants, tried in order per endpoint until one is +# accepted (capability is cached per endpoint afterwards): +# ids+logprobs : token_ids (detokenize=false) + per-token logprobs +# str+logprobs : decoded token strings + per-token logprobs +# text-only : plain text completion (weakest comparison) +_VARIANTS = ("ids+logprobs", "str+logprobs", "text-only") + + +def _http_json(url, payload=None, timeout=300): + data = None + headers = {} + if payload is not None: + data = json.dumps(payload).encode() + headers["Content-Type"] = "application/json" + req = urllib.request.Request(url, data=data, headers=headers) + with urllib.request.urlopen(req, timeout=timeout) as resp: + return json.loads(resp.read()) + + +def _served_model(base_url, timeout): + """Model id from /v1/models, or None if the endpoint does not expose it. + + The `trtllm-serve disaggregated` proxy only implements the completions + routes (no /v1/models -> HTTP 404), so a missing model listing must not + be fatal: the caller falls back to the other endpoint's model name or + to --model. (Dry runs against two aggregated servers do not hit this: + both expose /v1/models.) + """ + try: + models = _http_json(f"{base_url}/v1/models", timeout=timeout) + return models["data"][0]["id"] + except urllib.error.HTTPError as e: + print(f"[parity] NOTE: {base_url}/v1/models unavailable (HTTP {e.code})") + return None + # URLError/OSError: connection-level failures. KeyError/IndexError/ + # TypeError: malformed payload shapes ([], {"data": None}, + # {"data": [None]}, ...). ValueError covers json.JSONDecodeError from + # non-JSON response bodies. + except (urllib.error.URLError, OSError, KeyError, IndexError, TypeError, ValueError) as e: + print(f"[parity] NOTE: {base_url}/v1/models unavailable ({e})") + return None + + +class Endpoint: + """One OpenAI-compatible completions endpoint with capability fallback.""" + + def __init__(self, base_url, model, max_tokens, logprobs_k, timeout): + self.base_url = base_url.rstrip("/") + self.model = model + self.max_tokens = max_tokens + self.logprobs_k = logprobs_k + self.timeout = timeout + self.variant = None # discovered on first successful request + + def _payload(self, prompt, variant): + payload = { + "model": self.model, + "prompt": prompt, + "max_tokens": self.max_tokens, + "temperature": 0.0, + "stream": False, + } + if variant != "text-only" and self.logprobs_k > 0: + payload["logprobs"] = self.logprobs_k + if variant == "ids+logprobs": + payload["detokenize"] = False + return payload + + def complete(self, prompt): + """Return a normalized completion record for one prompt.""" + variants = [self.variant] if self.variant else list(_VARIANTS) + if self.logprobs_k <= 0: + variants = ["text-only"] + last_error = None + for variant in variants: + try: + rsp = _http_json( + f"{self.base_url}/v1/completions", + self._payload(prompt, variant), + timeout=self.timeout, + ) + except urllib.error.HTTPError as e: + # 4xx = capability rejection -> try the next variant; + # anything else is a real endpoint failure. + if 400 <= e.code < 500 and variant != variants[-1]: + last_error = e + continue + raise + if self.variant is None: + self.variant = variant + if variant != _VARIANTS[0]: + print( + f"[parity] NOTE: {self.base_url} degraded to " + f"'{variant}' requests ({last_error})" + ) + return _normalize_choice(rsp["choices"][0]) + raise last_error + + +def _normalize_choice(choice): + """Normalize a CompletionResponseChoice dict for comparison. + + Returns a namespace with: + text str ('' when detokenize=false) + ids per-position token identities: ints (token_ids), + else decoded token strings, else None + token_logprobs chosen-token logprob per position, or None + top_logprobs per-position {token_str: logprob} dict list, or None + """ + lp = choice.get("logprobs") or {} + tokens = lp.get("tokens") or None + token_logprobs = lp.get("token_logprobs") or None + top_logprobs = lp.get("top_logprobs") or None + ids = choice.get("token_ids") or tokens + return SimpleNamespace( + text=choice.get("text", ""), + ids=ids, + tokens=tokens, + token_logprobs=token_logprobs, + top_logprobs=top_logprobs, + ) + + +def _position_logprob_dict(chosen_key, chosen_lp, top_map): + """Build a per-position {token: logprob} dict for _compare_logits_parity. + + ``top_map`` (decoded-token-string keyed, from CompletionLogProbs) + includes the chosen token itself; when re-keying the chosen entry + under ``chosen_key`` (a token id), drop exactly one entry carrying + the chosen logprob so the top-2 near-tie gap is not computed against + a duplicate of the winner. Alternate entries get synthetic keys — + only the chosen key is ever looked up. + """ + d = {} + if top_map: + remaining = list(top_map.items()) + for i, (_, v) in enumerate(remaining): + if v == chosen_lp: + del remaining[i] + break + for j, (_, v) in enumerate(remaining): + d[("alt", j)] = v + d[chosen_key] = chosen_lp + return d + + +def _to_parity_namespace(comp): + """Adapt a normalized completion to _compare_logits_parity's shape. + + Returns a namespace with token_ids + per-position logprob dicts, or + None when the endpoint returned no aligned per-token logprobs. + """ + if comp.ids is None or comp.token_logprobs is None or len(comp.token_logprobs) != len(comp.ids): + return None + logprobs = [] + for i, (tid, lp) in enumerate(zip(comp.ids, comp.token_logprobs)): + top = comp.top_logprobs[i] if comp.top_logprobs and i < len(comp.top_logprobs) else None + logprobs.append(_position_logprob_dict(tid, lp, top)) + return SimpleNamespace(token_ids=list(comp.ids), logprobs=logprobs) + + +def _lcp(a, b): + n = 0 + while n < min(len(a), len(b)) and a[n] == b[n]: + n += 1 + return n + + +def compare_pair(ref, cand, prompt, failures, lp_tol, tie_tol): + """Compare one prompt's completions; returns the per-prompt record. + + Classification (reference logprobs available): + identical / benign (near-tie flip) / non_tie / drift — exactly the + kimi_k3_sa_harness semantics. Without logprobs, only token-level + facts are recorded (outcome 'divergence_unclassified'). + """ + ref_ids = list(ref.ids) if ref.ids is not None else None + cand_ids = list(cand.ids) if cand.ids is not None else None + record = {"prompt": prompt} + if ref_ids is None or cand_ids is None: + # Text-only endpoints: character-level LCP is the best available. + match = ref.text == cand.text + lcp = _lcp(ref.text, cand.text) + record.update( + comparison="text", + text_match=match, + char_lcp=lcp, + outcome="identical" if match else "divergence_unclassified", + ) + if not match: + print( + f"[parity] {prompt!r}: text mismatch at char {lcp} " + "(token/logprob detail unavailable)" + ) + return record + + lcp = _lcp(ref_ids, cand_ids) + full = lcp == len(ref_ids) == len(cand_ids) + record.update( + comparison="tokens", + first_token_match=lcp > 0, + lcp=lcp, + ref_len=len(ref_ids), + cand_len=len(cand_ids), + full_match=full, + ) + + ref_ns = _to_parity_namespace(ref) + cand_ns = _to_parity_namespace(cand) + if ref_ns is not None and cand_ns is not None: + outcome = _compare_logits_parity( + ref_ns, cand_ns, prompt, failures, tol=lp_tol, tie_tol=tie_tol + ) + # Mean |delta| of chosen-token logprobs over the shared prefix — + # a drift trend indicator below the hard tolerance. + deltas = [abs(ref.token_logprobs[i] - cand.token_logprobs[i]) for i in range(lcp)] + if deltas: + record["mean_abs_lp_delta"] = sum(deltas) / len(deltas) + record["max_abs_lp_delta"] = max(deltas) + elif full: + outcome = "identical" + else: + outcome = "divergence_unclassified" + print( + f"[parity] {prompt!r}: token divergence at position {lcp} " + "(logprobs unavailable — cannot classify near-tie vs real)" + ) + record["outcome"] = outcome + + if lcp == 0: + # In disagg the first token comes from the ctx server via the + # handoff; a confident first-token flip is the KV/KDA-transfer + # bug signature. Only a certified near-tie is excusable. + if outcome == "benign": + print( + f"[parity] WARNING: first-token near-tie flip for " + f"{prompt!r} (benign, but watch the aggregate)" + ) + else: + failures.append( + f"first-token mismatch for {prompt!r} " + f"(outcome={outcome}; disagg first-token handoff or " + "KV/state transfer suspect)" + ) + return record + + +def _aggregate(outcomes, failures): + """Apply the kimi_k3_sa_harness systemic check. + + Scattered non-tie flips are rounding; a real transfer/state bug makes + them dominate. + """ + divergences = [o for o in outcomes if o in ("benign", "non_tie")] + non_ties = outcomes.count("non_tie") + unclassified = outcomes.count("divergence_unclassified") + print( + f"[parity] summary: {len(outcomes)} prompts, " + f"{outcomes.count('identical')} identical, " + f"{len(divergences)} classified divergences ({non_ties} non-tie), " + f"{unclassified} unclassified, {outcomes.count('drift')} drift" + ) + if non_ties >= 2 and non_ties > 0.25 * max(len(divergences), 1): + failures.append( + f"non-tie divergences dominate: {non_ties}/{len(divergences)} " + "divergences exceeded the tie bound (systemic — suspect " + "transfer/state bug)" + ) + + +def run_gsm8k(base_url, model, out_dir, limit, concurrency, timeout, trust_remote_code=False): + """Run lm_eval GSM8K against one endpoint; returns {metric: value}. + + Mirrors the accuracy flow of examples/disaggregated/slurm/benchmark/ + submit.py (local-completions against /v1/completions). ``model`` must + be tokenizer-resolvable for lm_eval (a local checkpoint dir or HF repo + id) — the bare served-model name from /v1/models usually is not; pass + --tokenizer in that case. + """ + os.makedirs(out_dir, exist_ok=True) + model_args = ( + f"model={model},base_url={base_url}/v1/completions," + f"num_concurrent={concurrency},max_retries=3," + f"tokenized_requests=false,timeout={timeout}," + "max_gen_toks=256,max_length=4096" + ) + if trust_remote_code: + # e.g. Kimi K3's tiktoken-based tokenizer ships as checkpoint code + model_args += ",trust_remote_code=true" + cmd = [ + "lm_eval", + "--model", + "local-completions", + "--tasks", + "gsm8k", + "--model_args", + model_args, + "--log_samples", + "--output_path", + out_dir, + ] + if limit: + cmd += ["--limit", str(limit)] + print(f"[parity] running: {' '.join(cmd)}") + subprocess.run(cmd, check=True) + results = sorted( + glob.glob(os.path.join(out_dir, "**", "results*.json"), recursive=True), + key=os.path.getmtime, + ) + if not results: + raise RuntimeError(f"lm_eval produced no results json in {out_dir}") + with open(results[-1]) as f: + gsm8k = json.load(f)["results"]["gsm8k"] + return {k: v for k, v in gsm8k.items() if k.startswith("exact_match") and "stderr" not in k} + + +def _diff_gsm8k(ref_scores, cand_scores, tol, failures): + report = {} + for metric, ref_val in sorted(ref_scores.items()): + cand_val = cand_scores.get(metric) + delta = None if cand_val is None else cand_val - ref_val + report[metric] = {"reference": ref_val, "candidate": cand_val, "delta": delta} + print( + f"[parity] gsm8k {metric}: reference={ref_val:.4f} " + f"candidate={cand_val:.4f} delta={delta:+.4f}" + if cand_val is not None + else f"[parity] gsm8k {metric}: candidate missing" + ) + if delta is None or abs(delta) > tol: + failures.append( + f"gsm8k {metric} delta {delta} exceeds tolerance {tol} " + f"(reference {ref_val} vs candidate {cand_val})" + ) + return report + + +def _parse_args(argv): + parser = argparse.ArgumentParser( + description=__doc__.split("\n\n")[0], formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument( + "--reference", + help="base URL of the reference deployment (e.g. aggregated: http://host:8000)", + ) + parser.add_argument( + "--candidate", + help="base URL of the candidate deployment (e.g. disagg proxy: http://host:9000)", + ) + parser.add_argument( + "--model", default=None, help="served model name (default: query /v1/models)" + ) + parser.add_argument( + "--tokenizer", + default=None, + help="tokenizer path/repo for lm_eval GSM8K (a local checkpoint " + "dir; default: the --model value, which must then be " + "tokenizer-resolvable)", + ) + parser.add_argument("--max-tokens", type=int, default=64) + parser.add_argument( + "--extra-prompts", + type=int, + default=16, + help="extra deterministic parity prompts on top of " + "the builtin set (kimi_k3_sa_harness pool)", + ) + parser.add_argument( + "--prompts-file", + default=None, + help="newline-separated prompt file replacing the builtin prompt set", + ) + parser.add_argument( + "--logprobs", + type=int, + default=5, + help="top-k logprobs to request per token (0 disables logprob-level parity)", + ) + parser.add_argument( + "--lp-tol", + type=float, + default=float(os.environ.get("KIMI_K3_SPEC_LP_TOL", "1.0")), + help="shared-prefix logprob drift tolerance", + ) + parser.add_argument( + "--tie-tol", + type=float, + default=float(os.environ.get("KIMI_K3_SPEC_TIE_TOL", "0.3")), + help="near-tie top-2 logprob gap bound at divergence", + ) + parser.add_argument( + "--tie-strict", action="store_true", help="fail (not warn) on any single non-tie divergence" + ) + parser.add_argument("--timeout", type=int, default=600, help="per-request timeout in seconds") + parser.add_argument( + "--gsm8k", + action="store_true", + help="also run lm_eval GSM8K against both endpoints and diff scores", + ) + parser.add_argument( + "--gsm8k-only", action="store_true", help="skip the token/logprob probe; GSM8K diff only" + ) + parser.add_argument( + "--gsm8k-limit", type=int, default=None, help="lm_eval --limit (default: full GSM8K)" + ) + parser.add_argument("--gsm8k-concurrency", type=int, default=8) + parser.add_argument( + "--gsm8k-trust-remote-code", + action="store_true", + help="pass trust_remote_code=true to lm_eval (checkpoints whose " + "tokenizer ships as remote code, e.g. Kimi K3)", + ) + parser.add_argument( + "--gsm8k-tol", type=float, default=0.02, help="max tolerated |accuracy delta|" + ) + parser.add_argument( + "--work-dir", default=None, help="directory for lm_eval outputs (default: temp dir)" + ) + parser.add_argument( + "--report-json", default=None, help="write the full parity report to this path" + ) + parser.add_argument( + "--self-test", + action="store_true", + help="run the built-in comparison-logic self-test (no servers needed)", + ) + args = parser.parse_args(argv) + if not args.self_test and not (args.reference and args.candidate): + parser.error("--reference and --candidate are required (unless --self-test)") + return args + + +def main(argv=None) -> int: + args = _parse_args(argv) + if args.self_test: + return _self_test() + if args.tie_strict: + os.environ["KIMI_K3_SPEC_TIE_STRICT"] = "1" + + failures = [] + report = {"reference": args.reference, "candidate": args.candidate} + + if args.model: + model = args.model + else: + model = _served_model(args.reference, args.timeout) + cand_model = _served_model(args.candidate, args.timeout) + # Either side may lack /v1/models (e.g. the disagg proxy); fall + # back to the side that reports one. + model = model if model is not None else cand_model + if model is None: + print("[parity] ERROR: neither endpoint exposes /v1/models; pass --model explicitly") + return 1 + if cand_model is not None and cand_model != model: + print( + f"[parity] NOTE: endpoints serve different model names " + f"({model!r} vs {cand_model!r}); sending {model!r} to both, " + "pass --model to override" + ) + report["model"] = model + + if not args.gsm8k_only: + if args.prompts_file: + with open(args.prompts_file) as f: + prompts = [line.rstrip("\n") for line in f if line.strip()] + else: + prompts = [p for p, _ in PROMPTS_AND_CHECKS] + [ + p for p, _ in _parity_prompts(args.extra_prompts) + ] + print( + f"[parity] probing {len(prompts)} prompts x " + f"{args.max_tokens} max_tokens, logprobs={args.logprobs}" + ) + + endpoints = {} + for role, url in (("reference", args.reference), ("candidate", args.candidate)): + endpoints[role] = Endpoint(url, model, args.max_tokens, args.logprobs, args.timeout) + + records, outcomes = [], [] + for prompt in prompts: + try: + ref = endpoints["reference"].complete(prompt) + cand = endpoints["candidate"].complete(prompt) + except (urllib.error.URLError, OSError, KeyError) as e: + failures.append(f"endpoint error for {prompt!r}: {e}") + records.append({"prompt": prompt, "error": str(e)}) + continue + record = compare_pair(ref, cand, prompt, failures, args.lp_tol, args.tie_tol) + records.append(record) + outcomes.append(record["outcome"]) + _aggregate(outcomes, failures) + report["prompts"] = records + report["capability"] = {role: ep.variant for role, ep in endpoints.items()} + + if args.gsm8k or args.gsm8k_only: + work_dir = args.work_dir or tempfile.mkdtemp(prefix="k3-parity-") + scores = {} + for role, url in (("reference", args.reference), ("candidate", args.candidate)): + scores[role] = run_gsm8k( + url, + args.tokenizer or model, + os.path.join(work_dir, f"gsm8k_{role}"), + args.gsm8k_limit, + args.gsm8k_concurrency, + args.timeout, + args.gsm8k_trust_remote_code, + ) + report["gsm8k"] = _diff_gsm8k( + scores["reference"], scores["candidate"], args.gsm8k_tol, failures + ) + + if args.report_json: + with open(args.report_json, "w") as f: + json.dump(report, f, indent=2, default=str) + print(f"[parity] report written to {args.report_json}") + + if failures: + print("[parity] FAIL") + for f in failures: + print(f" - {f}") + return 1 + print("[parity] PASS") + return 0 + + +# --------------------------------------------------------------------------- +# Self-test: canned completions through the real comparison pipeline. +# --------------------------------------------------------------------------- + + +def _fake(ids, lps, top=None, text=""): + return SimpleNamespace(text=text, ids=ids, tokens=None, token_logprobs=lps, top_logprobs=top) + + +def _self_test() -> int: + lp_tol, tie_tol = 1.0, 0.3 + checks = [] + + def check(name, cond): + checks.append((name, cond)) + print(f"[self-test] {'PASS' if cond else 'FAIL'}: {name}") + + top_confident = [{"a": -0.05, "b": -3.5}] * 4 + top_tie = [{"a": -0.60, "b": -0.72}] * 4 + + # 1. identical streams + logprobs -> identical, no failures. + f = [] + r = compare_pair( + _fake([1, 2, 3], [-0.1, -0.2, -0.3], top_confident), + _fake([1, 2, 3], [-0.1, -0.2, -0.3], top_confident), + "identical", + f, + lp_tol, + tie_tol, + ) + check( + "identical outcome", + r["outcome"] == "identical" and not f and r["full_match"] and r["lcp"] == 3, + ) + + # 2. near-tie divergence mid-stream -> benign, no failures. + f = [] + r = compare_pair( + _fake([1, 2, 3], [-0.1, -0.2, -0.6], top_tie), + _fake([1, 2, 4], [-0.1, -0.2, -0.7], top_tie), + "near-tie", + f, + lp_tol, + tie_tol, + ) + check( + "near-tie -> benign", + r["outcome"] == "benign" and not f and r["lcp"] == 2 and r["first_token_match"], + ) + + # 3. logprob drift on the shared prefix -> drift + failure. + f = [] + r = compare_pair( + _fake([1, 2, 3], [-0.1, -0.2, -0.3], top_confident), + _fake([1, 2, 3], [-0.1, -5.2, -0.3], top_confident), + "drift", + f, + lp_tol, + tie_tol, + ) + check("drift fails", r["outcome"] == "drift" and len(f) == 1) + + # 4. confident (non-tie) divergence mid-stream -> warning only. + f = [] + r = compare_pair( + _fake([1, 2, 3], [-0.1, -0.2, -0.05], top_confident), + _fake([1, 2, 4], [-0.1, -0.2, -3.5], top_confident), + "non-tie", + f, + lp_tol, + tie_tol, + ) + check("single non-tie warns", r["outcome"] == "non_tie" and not f) + + # 5. confident first-token mismatch -> hard failure. + f = [] + r = compare_pair( + _fake([1, 2], [-0.05, -0.2], top_confident), + _fake([9, 2], [-3.5, -0.2], top_confident), + "first-token", + f, + lp_tol, + tie_tol, + ) + check("first-token mismatch fails", r["lcp"] == 0 and len(f) >= 1) + + # 6. no logprobs available -> unclassified divergence, first-token + # mismatch still fails (cannot be excused without logprobs). + f = [] + r = compare_pair(_fake([5, 6], None), _fake([7, 6], None), "no-logprobs", f, lp_tol, tie_tol) + check("logprob-less mismatch fails", r["outcome"] == "divergence_unclassified" and len(f) == 1) + + # 7. text-only endpoints -> character comparison. + f = [] + r = compare_pair( + _fake(None, None, text="hello world"), + _fake(None, None, text="hello there"), + "text-only", + f, + lp_tol, + tie_tol, + ) + check( + "text-only unclassified", + r["comparison"] == "text" and r["outcome"] == "divergence_unclassified", + ) + + # 8. aggregate rule: dominating non-ties -> systemic failure. + f = [] + _aggregate(["non_tie", "non_tie", "benign"], f) + check("dominating non-ties fail", len(f) == 1) + f = [] + _aggregate(["benign"] * 8 + ["non_tie", "identical"], f) + check("scattered non-tie passes", not f) + + # 9. chosen-token duplicate is dropped from the near-tie gap. + d = _position_logprob_dict(42, -0.05, {"a": -0.05, "b": -3.5}) + top = sorted(d.values(), reverse=True) + check( + "top-2 gap excludes chosen duplicate", len(d) == 2 and abs((top[0] - top[1]) - 3.45) < 1e-9 + ) + + # 10. gsm8k diff: within-tolerance passes, above-tolerance fails. + f = [] + _diff_gsm8k({"exact_match,strict-match": 0.90}, {"exact_match,strict-match": 0.89}, 0.02, f) + ok_within = not f + f = [] + _diff_gsm8k({"exact_match,strict-match": 0.90}, {"exact_match,strict-match": 0.80}, 0.02, f) + check("gsm8k tolerance", ok_within and len(f) == 1) + + # 11. _served_model returns None for malformed /v1/models payloads, + # invalid JSON, and connection failures; extracts the id otherwise. + global _http_json + orig_http_json = _http_json + bad_responses = [ + [], + {"data": None}, + {"data": [None]}, + {"data": [{}]}, + json.JSONDecodeError("Expecting value", "not-json", 0), + urllib.error.URLError("connection refused"), + ] + try: + results = [] + for rsp in bad_responses: + + def _canned(url, payload=None, timeout=0, _rsp=rsp): + if isinstance(_rsp, Exception): + raise _rsp + return _rsp + + _http_json = _canned + results.append(_served_model("http://stub", timeout=1)) + + def _good(url, payload=None, timeout=0): + return {"data": [{"id": "the-model"}]} + + _http_json = _good + served = _served_model("http://stub", timeout=1) + finally: + _http_json = orig_http_json + check("_served_model malformed payloads -> None", all(r is None for r in results)) + check("_served_model well-formed payload", served == "the-model") + + failed = [name for name, cond in checks if not cond] + if failed: + print(f"[self-test] FAIL ({len(failed)}/{len(checks)}): {failed}") + return 1 + print(f"[self-test] PASS ({len(checks)} checks)") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/unittest/disaggregated/region/test_aux.py b/tests/unittest/disaggregated/region/test_aux.py index 54d07252cbd6..7c2528fbf81c 100644 --- a/tests/unittest/disaggregated/region/test_aux.py +++ b/tests/unittest/disaggregated/region/test_aux.py @@ -3,7 +3,12 @@ import numpy as np import pytest -from tensorrt_llm._torch.disaggregation.native.auxiliary import AuxBuffer, AuxBufferMeta, AuxSlot +from tensorrt_llm._torch.disaggregation.native.auxiliary import ( + AuxBuffer, + AuxBufferMeta, + AuxSlot, + build_aux_transfer_layout, +) pytestmark = pytest.mark.cpu_only @@ -128,3 +133,55 @@ def test_fill_slot_unallocated_raises(): with pytest.raises(ValueError, match="not currently allocated"): buf.fill_slot(0, mock_request) + + +def test_aux_buffer_zero_max_draft_len_round_trip(): + """AuxBuffer round-trips with max_draft_len=0. + + A ctx server without speculative_config builds an AuxBuffer with + max_draft_len=0 (ctx/gen spec-config split). fill/get must round-trip + with an empty draft-token list. + """ + buf = AuxBuffer(max_slot_num=4, beam_width=1, max_draft_len=0, device="cpu") + slot = buf.alloc_slot() + + mock_request = MagicMock() + mock_request.get_last_tokens.return_value = [42] + mock_request.py_draft_tokens = [] + mock_request.prompt_len = 16 + mock_request.cached_tokens = 0 + mock_request.py_disaggregated_params = None + + buf.fill_slot(slot.id, mock_request) + first_tokens, draft_tokens = buf.get_slot_tokens(slot.id) + assert first_tokens == [42] + assert draft_tokens == [] + + # The draft buffer's per-slot item size is 0. + assert buf.meta.item_sizes[1] == 0 + + # Overfilling draft tokens must be rejected, not silently truncated. + mock_request.py_draft_tokens = [1] + with pytest.raises(ValueError, match="exceeds `max_draft_len`"): + buf.fill_slot(slot.id, mock_request) + + +def test_aux_transfer_layout_ctx_no_spec_gen_sa(): + """Ctx/gen spec split: ctx max_draft_len=0, gen max_draft_len=2. + + A ctx server without speculative_config publishes a zero-size draft-token + buffer while an SA/NGram gen server sizes it by max_draft_len. The + empty-source entry must be dropped from the transfer layout (nothing to + send), while the other buffers keep each side's own slot stride. + """ + ctx = AuxBuffer(max_slot_num=4, beam_width=1, max_draft_len=0, device="cpu") + gen = AuxBuffer(max_slot_num=4, beam_width=1, max_draft_len=2, device="cpu") + layout = build_aux_transfer_layout(ctx.meta, gen.meta) + + # Buffer order: [first_tokens, draft_tokens, token_counts, prompt_token_counts]; + # the draft_tokens entry (index 1) is dropped. + expected_keep = [0, 2, 3] + np.testing.assert_array_equal(layout.src_item_sizes, ctx.meta.item_sizes[expected_keep]) + np.testing.assert_array_equal(layout.dst_item_sizes, gen.meta.item_sizes[expected_keep]) + np.testing.assert_array_equal(layout.src_base_ptrs, ctx.meta.ptrs[expected_keep]) + np.testing.assert_array_equal(layout.dst_base_ptrs, gen.meta.ptrs[expected_keep]) diff --git a/tests/unittest/disaggregated/test_bounce.py b/tests/unittest/disaggregated/test_bounce.py index be94d873ed61..84a7eda4fe1a 100644 --- a/tests/unittest/disaggregated/test_bounce.py +++ b/tests/unittest/disaggregated/test_bounce.py @@ -36,11 +36,24 @@ try: from tensorrt_llm._torch.disaggregation.native.bounce import buffer as bbuf from tensorrt_llm._torch.disaggregation.native.bounce import impl as btr + from tensorrt_llm._torch.disaggregation.native.mixers.ssm.peer import MambaPolicy _HAVE_TRANSPORT = True except ImportError: # pragma: no cover - CPU-only env without CUDA bindings _HAVE_TRANSPORT = False +# page.py is pure numpy dataclasses, always importable on CPU. +from tensorrt_llm._torch.disaggregation.resource.page import ( + BUFFER_ENTRY_DTYPE, + AttentionLayerGroup, + KVCachePageTable, + LocalLayer, + MambaLayerGroup, + PhysicalPool, + PhysicalPoolGroup, + PoolView, +) + _MIB = 1024 * 1024 @@ -90,7 +103,10 @@ def test_fit_within_free_none_when_too_small(self): def test_config_defaults(self): cfg = bcfg.Config() assert isinstance(cfg.sizing, bcfg.FixedSizing) - assert cfg.chunk_mb == 32 and cfg.min_blocks == 96 + assert cfg.chunk_mb == 32 + # Byte gate for recurrent-state payloads; block gate for plain-KV payloads. + assert cfg.min_bytes == bcfg.DEFAULT_MIN_BYTES == 2 * _MIB + assert cfg.min_blocks == 96 # --------------------------------------------------------------------------- # @@ -106,24 +122,41 @@ def test_positive_enables_with_capacity(self): assert isinstance(cfg, bcfg.Config) assert cfg.sizing.capacity_mb == 2048 + def test_min_bytes_defaults_and_overrides(self, monkeypatch): + monkeypatch.delenv("TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES", raising=False) + assert bcfg.config_from_size(2048).min_bytes == 2 * _MIB # keeps the Config default + assert bcfg.config_from_size(2048, min_bytes=1).min_bytes == 1 # explicit arg overrides + assert bcfg.config_from_size(2048, min_bytes=64 * _MIB).min_bytes == 64 * _MIB + # The gate is tuned for tests via the env override (no user-facing config field). + monkeypatch.setenv("TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES", "1") + assert bcfg.config_from_size(2048).min_bytes == 1 # env override + assert bcfg.config_from_size(2048, min_bytes=250).min_bytes == 250 # arg beats env + def test_min_blocks_defaults_and_overrides(self, monkeypatch): + # The block-count gate (plain-KV payloads) keeps its original default of 96 and honors + # the explicit arg and the env override. monkeypatch.delenv("TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS", raising=False) assert bcfg.config_from_size(2048).min_blocks == 96 # keeps the Config default - assert bcfg.config_from_size(2048, 1).min_blocks == 1 # explicit arg still overrides - assert bcfg.config_from_size(2048, 250).min_blocks == 250 - # The gate is lowered for tests via the env override (no user-facing config field). - monkeypatch.setenv("TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS", "1") - assert bcfg.config_from_size(2048).min_blocks == 1 # env override + assert bcfg.config_from_size(2048, 250).min_blocks == 250 # explicit arg still overrides + monkeypatch.setenv("TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS", "48") + assert bcfg.config_from_size(2048).min_blocks == 48 # env override assert bcfg.config_from_size(2048, 250).min_blocks == 250 # explicit arg beats env - def test_min_blocks_env_is_parsed_defensively(self, monkeypatch): + @pytest.mark.parametrize( + "env,attr,default", + [ + ("TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES", "min_bytes", 2 * _MIB), + ("TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS", "min_blocks", 96), + ], + ) + def test_gate_env_is_parsed_defensively(self, monkeypatch, env, attr, default): # A bad value must not crash setup (falls back to the default); a non-positive value clamps to 1. for bad in ("", "auto", "1.5"): - monkeypatch.setenv("TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS", bad) - assert bcfg.config_from_size(2048).min_blocks == 96 + monkeypatch.setenv(env, bad) + assert getattr(bcfg.config_from_size(2048), attr) == default for nonpos in ("0", "-5"): - monkeypatch.setenv("TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS", nonpos) - assert bcfg.config_from_size(2048).min_blocks == 1 + monkeypatch.setenv(env, nonpos) + assert getattr(bcfg.config_from_size(2048), attr) == 1 # --------------------------------------------------------------------------- # @@ -276,6 +309,7 @@ def __init__(self, capacity_bytes, phys_chunk_size, name="kv_bounce"): self.next_id = 0 self.released = [] self.quarantined = [] + self.reserved_sizes = [] @property def capacity(self): @@ -286,6 +320,7 @@ def reserve(self, size, timeout=None): return None sid = self.next_id self.next_id += 1 + self.reserved_sizes.append(size) return sid, self.base def release(self, slot_id): @@ -301,7 +336,9 @@ def reg_descs(self): return [] -def _make_transport(monkeypatch, block_bytes_per_group, capacity=1 << 30, min_blocks=1): +def _make_transport( + monkeypatch, block_bytes_per_group, capacity=1 << 30, min_bytes=1, min_blocks=1 +): monkeypatch.setattr(btr, "SlotAllocator", _FakeAlloc) monkeypatch.setattr(btr.VmmBounceTransport, "_new_stream", lambda self: 0) monkeypatch.setattr( @@ -316,6 +353,7 @@ def _make_transport(monkeypatch, block_bytes_per_group, capacity=1 << 30, min_bl capacity_bytes=capacity, phys_chunk_size=32 * _MIB, block_bytes_per_group=block_bytes_per_group, + min_bytes=min_bytes, min_blocks=min_blocks, ) @@ -369,9 +407,40 @@ def test_reserve_single_writer_ok(self, monkeypatch): req = _recv_req([1]) # total = 3, num_writers=1 -> no even-split requirement assert t.reserve(req, num_writers=1) is True - def test_reserve_too_small_falls_back(self, monkeypatch): - t = _make_transport(monkeypatch, block_bytes_per_group=[100], min_blocks=96) + def test_reserve_below_min_bytes_falls_back(self, monkeypatch): + # Recurrent-state payloads gate on BYTES: 4 blocks x 100B + 100B state = 500B < 1000B + # falls back, regardless of how many blocks that is. + t = _make_transport(monkeypatch, block_bytes_per_group=[100], min_bytes=1000) + req = _recv_req([4]) + assert t.reserve(req, num_writers=1, extra_bytes=100) is False + assert req.bounce_dst_base is None + + def test_reserve_at_min_bytes_bounces(self, monkeypatch): + # exactly at the threshold (9 x 100B + 100B state = 1000B >= 1000B) the transfer bounces + t = _make_transport(monkeypatch, block_bytes_per_group=[100], min_bytes=1000) + req = _recv_req([9]) + assert t.reserve(req, num_writers=1, extra_bytes=100) is True + assert req.bounce_dst_base == 0x100000 + + def test_reserve_recurrent_payload_ignores_block_gate(self, monkeypatch): + # The K3 regression shape: FEW but LARGE blocks plus recurrent state must clear the byte + # gate even though a block-count gate of 96 fails (67 blocks x 6.5 MiB = 433 MiB). + t = _make_transport( + monkeypatch, + block_bytes_per_group=[int(6.5 * _MIB)], + min_bytes=2 * _MIB, + min_blocks=96, + ) + assert t.reserve(_recv_req([67]), num_writers=1, extra_bytes=1024) is True + + def test_reserve_plain_kv_uses_block_gate(self, monkeypatch): + # Plain-KV payloads (no recurrent state) keep the original block-count gate, and the byte + # gate does not apply to them (96 x 100B = 9600B passes despite min_bytes far above it). + t = _make_transport( + monkeypatch, block_bytes_per_group=[100], min_blocks=96, min_bytes=1 << 20 + ) assert t.reserve(_recv_req([4]), num_writers=1) is False # 4 < 96 blocks + assert t.reserve(_recv_req([96]), num_writers=1) is True def test_reserve_unknown_slot_size_falls_back(self, monkeypatch): t = _make_transport(monkeypatch, block_bytes_per_group=[100]) # only 1 group known @@ -548,6 +617,143 @@ def test_orphan_reservation_quarantines_and_is_idempotent(self, monkeypatch): assert t._recv_alloc.quarantined == [0] +# --------------------------------------------------------------------------- # +# Hybrid layouts (Kimi K3: KDA mamba + MLA attention) — page-table gate, +# trailing-empty-group handling, and recurrent-state (extra_bytes) sizing +# --------------------------------------------------------------------------- # + +# Kimi K3 geometry: 24 MLA layers in one attention layer group; 69 KDA layers +# whose per-layer recurrent state is a bf16 short-conv slot + fp32 delta slot, +# replicated per rank. +_K3_MLA_BLOCK_BYTES = 32 * 576 * 2 * 24 # tpb x (kv_lora_rank+rope) x bf16 x 24 layers = 884,736 +_K3_KDA_LAYERS = 69 +_K3_CONV_SLOT_BYTES = 294_912 # [3*H*hd, W] bf16 per layer +_K3_SSM_SLOT_BYTES = 6_291_456 # [H, hd, hd] fp32 per layer (95.5% of the state) +_K3_KDA_PAYLOAD_BYTES = ( + 454_459_392 # 69 x (conv + delta) per request per rank, from the geometry above +) + + +def _k3_page_table() -> KVCachePageTable: + """A K3-shaped page table exactly as the builders produce it. + + Attention layer group(s) first and the mamba layer group appended LAST + (kv_extractor.py, both ``build_page_table`` and ``_build_page_table_v2``, + append it after every attention group). + """ + attn = AttentionLayerGroup( + pool_group_idx=0, + kv_head_num_per_rank=1, + sliding_window_size=None, + local_layers=[LocalLayer(i, i) for i in range(24)], + pool_views=[PoolView(pool_idx=0, buffer_entries=np.array([], dtype=BUFFER_ENTRY_DTYPE))], + ) + mamba = MambaLayerGroup( + pool_group_idx=1, # dangling by construction: mamba pools live on the LG itself + mamba_layer_offsets={gl: i for i, gl in enumerate(range(24, 24 + _K3_KDA_LAYERS))}, + conv_states=PhysicalPool(0x200000, _K3_CONV_SLOT_BYTES, 8), + ssm_states=PhysicalPool(0x300000, _K3_SSM_SLOT_BYTES, 8), + conv_section_bytes=[_K3_CONV_SLOT_BYTES // 3] * 3, + ssm_bytes_per_head=_K3_SSM_SLOT_BYTES // 4, + ) + return KVCachePageTable( + tokens_per_block=32, + layer_groups=[attn, mamba], + pool_groups=[PhysicalPoolGroup(pools=[PhysicalPool(0x100000, _K3_MLA_BLOCK_BYTES, 128)])], + ) + + +def _k3_rank_info(): + # MambaPolicy only reads tp_size / tp_rank / attention.enable_attention_dp. + return SimpleNamespace(tp_size=1, tp_rank=0, attention=None) + + +@pytest.mark.skipif(not _HAVE_TRANSPORT, reason="bounce.transport import needs CUDA bindings") +class TestHybridK3Bounce: + def test_block_bytes_per_group_keeps_mamba_placeholder(self): + # The mamba group must stay in the list as a placeholder (None), aligned with the + # layer-group indices a recv request uses, instead of truncating the list. + assert btr.block_bytes_per_group(_k3_page_table()) == [_K3_MLA_BLOCK_BYTES, None] + + def test_reserve_engages_on_k3_mixed_layout(self, monkeypatch): + # Regression pin: a K3 recv request always carries a trailing EMPTY entry for + # the mamba layer group (transceiver._create_kv_slice), which used to trip the + # unknown-slot-size guard and silently push every K3 request onto the per-fragment + # (~0.4 GB/s host-staged) path. It must engage bounce, sized for MLA KV + KDA state. + t = _make_transport( + monkeypatch, + block_bytes_per_group=btr.block_bytes_per_group(_k3_page_table()), + min_bytes=2 * _MIB, + ) + req = _recv_req([67, 0]) # 67 MLA blocks (~2144 tokens) + the empty mamba entry + assert t.reserve(req, num_writers=1, extra_bytes=_K3_KDA_PAYLOAD_BYTES) is True + assert req.bounce_dst_base == 0x100000 + # the region covers the MLA KV blocks plus the appended KDA state + assert t._recv_alloc.reserved_sizes == [67 * _K3_MLA_BLOCK_BYTES + _K3_KDA_PAYLOAD_BYTES] + + def test_reserve_nonempty_group_with_unknown_size_still_falls_back(self, monkeypatch): + # Safety direction of the gate is preserved: a group that HAS blocks but no known slot + # size (None placeholder) still disables bounce for the request. + t = _make_transport(monkeypatch, block_bytes_per_group=[_K3_MLA_BLOCK_BYTES, None]) + assert t.reserve(_recv_req([2, 1]), num_writers=1) is False + + def test_reserve_extra_bytes_counts_toward_min_bytes(self, monkeypatch): + t = _make_transport(monkeypatch, block_bytes_per_group=[100], min_bytes=1000) + assert t.reserve(_recv_req([4]), num_writers=1, extra_bytes=599) is False # 999 < 1000 + assert t.reserve(_recv_req([4]), num_writers=1, extra_bytes=600) is True # exactly 1000 + + def test_reserve_extra_bytes_counts_toward_capacity(self, monkeypatch): + # extra_bytes must be part of the reservation, or the sender's coalesced write + # (KV + recurrent state) would overrun the region into the neighboring slot. + t = _make_transport(monkeypatch, block_bytes_per_group=[100], capacity=500) + assert t.reserve(_recv_req([2]), num_writers=1, extra_bytes=400) is False # 600 > 500 + + def test_reserve_extra_bytes_fanin_falls_back(self, monkeypatch): + # Fan-in splits the region equally by writer; per-writer recurrent-state fragments can + # differ (PP stages hold different mamba layers), so state + fan-in falls back. + t = _make_transport(monkeypatch, block_bytes_per_group=[100]) + assert t.reserve(_recv_req([2, 0]), num_writers=2, extra_bytes=64) is False + assert t.reserve(_recv_req([2, 0]), num_writers=2) is True # no state -> fan-in fine + + def test_mamba_payload_bytes_matches_k3_geometry(self): + # Matched-TP replicated KDA state: payload_bytes must reproduce the per-request + # per-rank number derived from K3's KDA geometry (constants above) exactly. + pt = _k3_page_table() + ri = _k3_rank_info() + got = MambaPolicy.payload_bytes( + sender_page_table=pt, receiver_page_table=pt, dst_slot=3, sender_ri=ri, receiver_ri=ri + ) + assert got == _K3_KDA_PAYLOAD_BYTES + + def test_mamba_payload_bytes_matches_collect_frags(self): + # payload_bytes mirrors the sender's collect_frags call: same sizes, slot-independent. + pt = _k3_page_table() + ri = _k3_rank_info() + _, _, sizes = MambaPolicy.collect_frags( + self_page_table=pt, peer_page_table=pt, src_slot=5, dst_slot=3, self_ri=ri, peer_ri=ri + ) + got = MambaPolicy.payload_bytes( + sender_page_table=pt, receiver_page_table=pt, dst_slot=3, sender_ri=ri, receiver_ri=ri + ) + assert got == sum(sizes) > 0 + + def test_mamba_payload_bytes_zero_without_slot_or_group(self): + pt = _k3_page_table() + ri = _k3_rank_info() + assert MambaPolicy.payload_bytes(pt, pt, dst_slot=None, sender_ri=ri, receiver_ri=ri) == 0 + pure_attn = KVCachePageTable( + tokens_per_block=32, + layer_groups=[pt.layer_groups[0]], + pool_groups=pt.pool_groups, + ) + assert ( + MambaPolicy.payload_bytes( + pure_attn, pure_attn, dst_slot=3, sender_ri=ri, receiver_ri=ri + ) + == 0 + ) + + # --------------------------------------------------------------------------- # # SlotAllocator — first-fit reuses out-of-order freed holes (real Buffer mocked) # --------------------------------------------------------------------------- # diff --git a/tests/unittest/disaggregated/test_kda_mamba_transfer.py b/tests/unittest/disaggregated/test_kda_mamba_transfer.py new file mode 100644 index 000000000000..fcef4587c05b --- /dev/null +++ b/tests/unittest/disaggregated/test_kda_mamba_transfer.py @@ -0,0 +1,589 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Kimi K3 KDA-shaped recurrent-state transfer through the Python V2 transceiver. + +Synthetic loopback coverage of the KDA recurrent-state transfer path: + +* KDA slot layout, mapped onto the Mamba cache-manager parametrization the + way ``_util.py`` does for ``kimi_linear``: + - short-conv slot ``[3*H*hd, W]`` **bf16** (qwen3_next ``[Q|K|V]`` + 3-section layout, all sections equal width), + - delta-rule slot ``[H, hd, hd]`` **fp32** (``state_size == head_dim``), + - EP pre-scale: ``num_heads``/``n_groups`` multiplied by ``tp_size`` when + attention-DP is off, so the per-rank state is full-size (replicated). + +* ``test_kda_layer_group_descriptors`` checks that the V2 page table + describes BOTH slots with exact byte sizes and that the matched-TP + ``MambaPolicy`` descriptors tile each layer's slot bytes exactly. + +* ``test_kda_transfer`` performs a real single-node NIXL loopback transfer + and bitwise-compares both dtypes on the gen side. + +* ``test_kda_hetero_tp_rejected`` asserts the peer-registration + guard: with replicated (pre-scaled) state, heterogeneous + ctx/gen TP (attention-DP off) is rejected instead of producing fragment + pointers outside the slot; supported layouts still validate. +""" + +import uuid +from typing import Dict, List + +import pytest +import torch +from test_mamba_transfer import _create_transceivers, _run_concurrent + +import tensorrt_llm +import tensorrt_llm.bindings +import tensorrt_llm.tensorrt_llm_transfer_agent_binding # noqa: F401 +from tensorrt_llm import DisaggregatedParams, Mapping, SamplingParams +from tensorrt_llm._torch.disaggregation.native.mixers.ssm.peer import MambaPolicy +from tensorrt_llm._torch.disaggregation.native.peer import PeerRegistrar +from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo +from tensorrt_llm._torch.disaggregation.resource.kv_extractor import ( + KVRegionExtractorV1, + build_page_table_from_manager, +) +from tensorrt_llm._torch.disaggregation.resource.page import MambaLayerGroup +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestType +from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import MixedMambaHybridCacheManager +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp +from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig, KvCacheConfig + +# --------------------------------------------------------------------------- +# KDA parameters (small stand-ins for Kimi K3's per-layer KDA state) +# --------------------------------------------------------------------------- +NUM_KDA_LAYERS = 4 +KDA_NUM_HEADS = 4 # H +KDA_HEAD_DIM = 32 # hd; KDA state_size == head_dim +KDA_W = 3 # short_conv_kernel_size; manager d_conv = W + 1 +CONV_DTYPE = torch.bfloat16 +SSM_DTYPE = torch.float32 +MAX_BATCH_SIZE = 4 +REQUEST_LENGTHS = [16, 32] + +# Internal: layer 0 is a dummy attention layer required by page table infra; +# layers 1..NUM_KDA_LAYERS are KDA (under test). +_NUM_TOTAL_LAYERS = NUM_KDA_LAYERS + 1 +_KDA_MASK = [False] + [True] * NUM_KDA_LAYERS +_ATTN_MASK = [True] + [False] * NUM_KDA_LAYERS + +# Full per-layer slot byte sizes (replicated on every rank for K3). +CONV_SLOT_BYTES = 3 * KDA_NUM_HEADS * KDA_HEAD_DIM * KDA_W * CONV_DTYPE.itemsize +SSM_SLOT_BYTES = KDA_NUM_HEADS * KDA_HEAD_DIM * KDA_HEAD_DIM * SSM_DTYPE.itemsize + + +def _create_kda_managers( + tp: int, enable_attention_dp: bool = False, max_batch_size: int = MAX_BATCH_SIZE +): + """Create MixedMambaHybridCacheManagers with K3-style KDA slots. + + Mirrors the ``is_kimi_linear`` route in ``_util.py:1855-1915``: KDA is + mapped onto (state_size=hd, conv_kernel=W+1, num_heads=H, n_groups=H, + head_dim=hd, model_type='qwen3_next'), and ``num_heads``/``n_groups`` + are pre-scaled by tp_size when attention-DP is off so the per-rank + state stays full-size (EP-only parallelism, replicated KDA state). + """ + state_tp = tp if not enable_attention_dp else 1 + managers = [] + for rank in range(tp): + mapping = Mapping( + world_size=tp, + rank=rank, + tp_size=tp, + pp_size=1, + enable_attention_dp=enable_attention_dp, + ) + mgr = MixedMambaHybridCacheManager( + mamba_d_state=KDA_HEAD_DIM, + mamba_d_conv=KDA_W + 1, + mamba_num_heads=KDA_NUM_HEADS * state_tp, + mamba_n_groups=KDA_NUM_HEADS * state_tp, + mamba_head_dim=KDA_HEAD_DIM, + mamba_num_layers=NUM_KDA_LAYERS, + mamba_layer_mask=_KDA_MASK, + mamba_cache_dtype=CONV_DTYPE, + mamba_ssm_cache_dtype=SSM_DTYPE, + # dummy attention layer (page table scaffolding) + kv_cache_config=KvCacheConfig( + max_tokens=256 * max_batch_size, + enable_block_reuse=False, + event_buffer_max_size=0, + ), + kv_cache_type=CacheTypeCpp.SELF, + num_layers=1, + layer_mask=_ATTN_MASK, + num_kv_heads=4, + head_dim=64, + tokens_per_block=8, + max_seq_len=256, + max_batch_size=max_batch_size, + mapping=mapping, + dtype=DataType.FLOAT, + model_type="qwen3_next", + ) + managers.append(mgr) + return managers + + +def _get_mamba_layer_group(page_table) -> MambaLayerGroup: + mlgs = [lg for lg in page_table.layer_groups if isinstance(lg, MambaLayerGroup)] + assert len(mlgs) == 1, f"expected exactly one MambaLayerGroup, got {len(mlgs)}" + return mlgs[0] + + +def _layer_slot_start(pool, local_layer_idx: int, slot: int) -> int: + return pool.base_address + (local_layer_idx * pool.num_slots + slot) * pool.slot_bytes + + +def _assert_frags_tile_slots(frags, mlg: MambaLayerGroup, slot: int): + """Assert (ptr, size) frags exactly tile every layer's slot bytes.""" + ptrs, sizes = frags + covered: Dict[int, List[tuple]] = {} + for ptr, size in zip(ptrs, sizes): + pool = None + for cand in (mlg.conv_states, mlg.ssm_states): + lo = cand.base_address + hi = lo + len(mlg.mamba_layer_offsets) * cand.num_slots * cand.slot_bytes + if lo <= ptr < hi: + pool = cand + break + assert pool is not None, f"frag ptr {ptr} outside both mamba pools" + rel = ptr - pool.base_address + layer = rel // (pool.num_slots * pool.slot_bytes) + in_layer = rel - layer * pool.num_slots * pool.slot_bytes + frag_slot = in_layer // pool.slot_bytes + off = in_layer - frag_slot * pool.slot_bytes + assert frag_slot == slot, f"frag targets slot {frag_slot}, expected {slot}" + assert off + size <= pool.slot_bytes, ( + f"frag [{off}, {off + size}) exceeds slot_bytes {pool.slot_bytes}" + ) + covered.setdefault((id(pool), int(layer)), []).append((off, off + size)) + + # Exact tiling per (pool, layer): sorted intervals must be contiguous + # from 0 to slot_bytes with no overlap. + n_layers = len(mlg.mamba_layer_offsets) + for pool in (mlg.conv_states, mlg.ssm_states): + for layer in range(n_layers): + intervals = sorted(covered.get((id(pool), layer), [])) + assert intervals, f"layer {layer} of pool not covered" + pos = 0 + for lo, hi in intervals: + assert lo == pos, f"gap/overlap at byte {pos} (next frag at {lo})" + pos = hi + assert pos == pool.slot_bytes, ( + f"layer {layer}: covered {pos} of {pool.slot_bytes} slot bytes" + ) + + +# --------------------------------------------------------------------------- +# Descriptor-level tests (no NIXL transfer) +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("enable_attention_dp", [True, False], ids=["adp_on", "adp_off_tp1"]) +def test_kda_layer_group_descriptors(enable_attention_dp): + """V2 page table must describe BOTH KDA slots with exact byte extents.""" + mgr = _create_kda_managers(1, enable_attention_dp=enable_attention_dp)[0] + try: + pt = build_page_table_from_manager(mgr) + mlg = _get_mamba_layer_group(pt) + + conv = mgr._impl.mamba_cache.conv + ssm = mgr._impl.mamba_cache.temporal + + # Shapes/dtypes of the backing tensors are KDA-shaped. + assert conv.dtype == CONV_DTYPE + assert ssm.dtype == SSM_DTYPE + assert tuple(conv.shape[2:]) == (3 * KDA_NUM_HEADS * KDA_HEAD_DIM, KDA_W) + assert tuple(ssm.shape[2:]) == (KDA_NUM_HEADS, KDA_HEAD_DIM, KDA_HEAD_DIM) + + # Both pools present, byte-exact (bf16 conv, fp32 delta). + assert mlg.conv_states is not None and mlg.ssm_states is not None + assert mlg.conv_states.base_address == conv.data_ptr() + assert mlg.ssm_states.base_address == ssm.data_ptr() + assert mlg.conv_states.slot_bytes == CONV_SLOT_BYTES + assert mlg.ssm_states.slot_bytes == SSM_SLOT_BYTES + assert mlg.conv_states.num_slots == conv.shape[1] + assert mlg.ssm_states.num_slots == ssm.shape[1] + + # qwen3_next 3-sectioning: equal sections summing to the conv slot. + assert mlg.conv_section_bytes == [CONV_SLOT_BYTES // 3] * 3 + assert sum(mlg.conv_section_bytes) == mlg.conv_states.slot_bytes + assert mlg.ssm_bytes_per_head == KDA_HEAD_DIM * KDA_HEAD_DIM * SSM_DTYPE.itemsize + assert mlg.ssm_states.slot_bytes // mlg.ssm_bytes_per_head == KDA_NUM_HEADS + + # Layer offsets cover exactly the KDA layers. + assert sorted(mlg.mamba_layer_offsets.keys()) == [i for i, m in enumerate(_KDA_MASK) if m] + + # Matched-parallelism frags: 2 pools x layers, full-slot, exact tiling. + ri = RankInfo.from_kv_cache_manager("kda_test", mgr, device_id=0) + src_slot, dst_slot = 0, 1 + src_frags, dst_frags, sizes = MambaPolicy.build_mamba_frags( + mlg, mlg, src_slot, dst_slot, ri, ri + ) + assert len(src_frags) == len(dst_frags) == len(sizes) == 2 * NUM_KDA_LAYERS + _assert_frags_tile_slots((src_frags, sizes), mlg, src_slot) + _assert_frags_tile_slots((dst_frags, sizes), mlg, dst_slot) + + # Full-slot copies at the expected per-layer addresses. + for glid, lid in sorted(mlg.mamba_layer_offsets.items()): + assert _layer_slot_start(mlg.conv_states, lid, src_slot) in src_frags + assert _layer_slot_start(mlg.ssm_states, lid, src_slot) in src_frags + finally: + mgr.shutdown() + + +def test_kda_hetero_tp_rejected(): + """Replicated KDA state + heterogeneous TP (ADP off) must be rejected. + + Hetero ctx/gen TP would silently corrupt replicated state: the TP-mismatch + mappers assume sharded state and would compute shard offsets past the + end of K3's replicated (pre-scaled) slots. The peer-registration gate + (``MambaPolicy.validate_peer_compatible``) must reject this loudly. + """ + ctx_mgrs = _create_kda_managers(2, enable_attention_dp=False) + gen_mgrs = _create_kda_managers(4, enable_attention_dp=False) + ctx_mgr, gen_mgr = ctx_mgrs[0], gen_mgrs[1] + try: + ctx_pt = build_page_table_from_manager(ctx_mgr) + gen_pt = build_page_table_from_manager(gen_mgr) + ctx_ri = RankInfo.from_kv_cache_manager("kda_ctx", ctx_mgr, device_id=0) + gen_ri = RankInfo.from_kv_cache_manager("kda_gen", gen_mgr, device_id=0) + + with pytest.raises(ValueError, match="TP-aggregated"): + MambaPolicy.validate_peer_compatible(ctx_ri, gen_ri, ctx_pt, gen_pt) + + # And via the registrar entry point used at runtime. + registrar = PeerRegistrar(ctx_ri, KVRegionExtractorV1(ctx_pt)) + with pytest.raises(ValueError, match="TP-aggregated"): + registrar.register("kda_gen", 1, gen_ri) + finally: + for mgr in ctx_mgrs + gen_mgrs: + mgr.shutdown() + + +@pytest.mark.parametrize( + "ctx_cfg,gen_cfg", + [ + ((2, False), (2, False)), # matched TP, ADP off (EP pre-scaled) + ((2, True), (4, True)), # heterogeneous DEP with ADP on both sides + ], + ids=["matched_tp2_adp_off", "hetero_dep_adp_on"], +) +def test_kda_peer_validation_accepts_supported_shapes(ctx_cfg, gen_cfg): + """Matched-TP and ADP-on-both-sides layouts must pass peer validation.""" + ctx_tp, ctx_adp = ctx_cfg + gen_tp, gen_adp = gen_cfg + ctx_mgrs = _create_kda_managers(ctx_tp, enable_attention_dp=ctx_adp) + gen_mgrs = _create_kda_managers(gen_tp, enable_attention_dp=gen_adp) + ctx_mgr, gen_mgr = ctx_mgrs[0], gen_mgrs[-1] + try: + ctx_pt = build_page_table_from_manager(ctx_mgr) + gen_pt = build_page_table_from_manager(gen_mgr) + ctx_ri = RankInfo.from_kv_cache_manager("kda_ctx", ctx_mgr, device_id=0) + gen_ri = RankInfo.from_kv_cache_manager("kda_gen", gen_mgr, device_id=0) + MambaPolicy.validate_peer_compatible(ctx_ri, gen_ri, ctx_pt, gen_pt) + finally: + for mgr in ctx_mgrs + gen_mgrs: + mgr.shutdown() + + +def _synthetic_kda_page_table(ssm_slot_bytes: int, conv_slot_bytes: int, layer_ids=None): + """MambaLayerGroup-only page table with fake addresses (no CUDA needed).""" + from tensorrt_llm._torch.disaggregation.resource.page import KVCachePageTable, PhysicalPool + + if layer_ids is None: + layer_ids = range(1, NUM_KDA_LAYERS + 1) + mlg = MambaLayerGroup( + pool_group_idx=0, + mamba_layer_offsets={glid: i for i, glid in enumerate(layer_ids)}, + conv_states=PhysicalPool(base_address=0x1000, slot_bytes=conv_slot_bytes, num_slots=8), + ssm_states=PhysicalPool(base_address=0x2000000, slot_bytes=ssm_slot_bytes, num_slots=8), + conv_section_bytes=[conv_slot_bytes // 3] * 3, + ssm_bytes_per_head=KDA_HEAD_DIM * KDA_HEAD_DIM * SSM_DTYPE.itemsize, + ) + return KVCachePageTable(tokens_per_block=8, layer_groups=[mlg], pool_groups=[]) + + +def _synthetic_rank_info(tp: int, adp: bool): + from tensorrt_llm._torch.disaggregation.native.mixers.attention.spec import AttentionInfo + + return RankInfo( + instance_name="syn", + instance_rank=0, + tp_size=tp, + tp_rank=0, + pp_size=1, + pp_rank=0, + layer_num_per_pp=[_NUM_TOTAL_LAYERS], + sender_endpoints=[], + self_endpoint="", + transfer_engine_info=b"", + attention=AttentionInfo( + kv_heads_per_rank=4, + tokens_per_block=8, + dims_per_head=64, + element_bytes=2, + enable_attention_dp=adp, + is_mla=False, + ), + ) + + +@pytest.mark.parametrize( + "ctx,gen,ok", + [ + # (tp, adp, slot_scale_denominator): replicated K3 state = full slots. + ((2, False, 1), (4, False, 1), False), # hetero TP, ADP off: would corrupt state + ((4, False, 1), (2, False, 1), False), # ...both directions + ((2, True, 1), (2, False, 1), False), # mixed ADP, replicated + ((2, False, 1), (2, False, 1), True), # matched TP + ((2, True, 1), (4, True, 1), True), # hetero DEP, ADP on both + ((2, False, 2), (4, False, 4), True), # sharded state, hetero TP + ], + ids=[ + "reject_hetero_tp_adp_off", + "reject_hetero_tp_adp_off_rev", + "reject_mixed_adp_replicated", + "accept_matched_tp", + "accept_hetero_dep_adp_on", + "accept_sharded_hetero_tp", + ], +) +def test_kda_peer_validation_synthetic_cpu(ctx, gen, ok): + """CPU-only reject/accept matrix for the gap-F1 guard (no CUDA manager).""" + full_ssm = KDA_NUM_HEADS * KDA_HEAD_DIM * KDA_HEAD_DIM * SSM_DTYPE.itemsize + full_conv = 3 * KDA_NUM_HEADS * KDA_HEAD_DIM * KDA_W * CONV_DTYPE.itemsize + + def build(cfg): + tp, adp, denom = cfg + return ( + _synthetic_rank_info(tp, adp), + _synthetic_kda_page_table(full_ssm // denom, full_conv // denom), + ) + + ctx_ri, ctx_pt = build(ctx) + gen_ri, gen_pt = build(gen) + if ok: + MambaPolicy.validate_peer_compatible(ctx_ri, gen_ri, ctx_pt, gen_pt) + else: + with pytest.raises(ValueError, match="TP-aggregated"): + MambaPolicy.validate_peer_compatible(ctx_ri, gen_ri, ctx_pt, gen_pt) + + +def test_kda_peer_validation_allows_pipeline_parallel_layer_split(): + """Peer validation must not require identical layer sets. + + Under pipeline parallelism each rank publishes only its own stage's + layers, and the transfer path takes the intersection of the two layer + sets. Peers with partially overlapping, disjoint, or one-sided + recurrent-layer sets are all legitimate stage pairings and must pass + validation as long as the per-slot size invariants hold. + """ + full_ssm = KDA_NUM_HEADS * KDA_HEAD_DIM * KDA_HEAD_DIM * SSM_DTYPE.itemsize + full_conv = 3 * KDA_NUM_HEADS * KDA_HEAD_DIM * KDA_W * CONV_DTYPE.itemsize + ctx_ri = _synthetic_rank_info(2, False) + gen_ri = _synthetic_rank_info(2, False) + + # Partially overlapping stages (ctx PP split differs from gen's). + ctx_pt = _synthetic_kda_page_table(full_ssm, full_conv, layer_ids=[1, 2, 3]) + gen_pt = _synthetic_kda_page_table(full_ssm, full_conv, layer_ids=[3, 4, 5]) + MambaPolicy.validate_peer_compatible(ctx_ri, gen_ri, ctx_pt, gen_pt) + + # Disjoint stages: this pair simply moves no recurrent state. + gen_pt = _synthetic_kda_page_table(full_ssm, full_conv, layer_ids=[7, 8]) + MambaPolicy.validate_peer_compatible(ctx_ri, gen_ri, ctx_pt, gen_pt) + + # One side holds a recurrent-layer-free stage of the hybrid model + # (no MambaLayerGroup at all): nothing to validate. + from tensorrt_llm._torch.disaggregation.resource.page import KVCachePageTable + + empty_pt = KVCachePageTable(tokens_per_block=8, layer_groups=[], pool_groups=[]) + MambaPolicy.validate_peer_compatible(ctx_ri, gen_ri, ctx_pt, empty_pt) + MambaPolicy.validate_peer_compatible(ctx_ri, gen_ri, empty_pt, gen_pt) + + # Size invariants still apply on the overlap: mismatched global state + # sizes are rejected even when the layer sets differ. + bad_gen_pt = _synthetic_kda_page_table(full_ssm // 2, full_conv // 2, layer_ids=[3, 4, 5]) + with pytest.raises(ValueError, match="TP-aggregated"): + MambaPolicy.validate_peer_compatible(ctx_ri, gen_ri, ctx_pt, bad_gen_pt) + + +# --------------------------------------------------------------------------- +# Loopback transfer over a real NIXL agent (single node) +# --------------------------------------------------------------------------- +def _generate_ground_truth(num_requests: int, seed: int = 20260722): + """Full replicated KDA states per (request, kda_layer), two dtypes.""" + gen = torch.Generator(device="cpu").manual_seed(seed) + results = [] + for _ in range(num_requests): + layers = {} + for i, is_kda in enumerate(_KDA_MASK): + if not is_kda: + continue + layers[i] = { + "conv": torch.rand(3 * KDA_NUM_HEADS * KDA_HEAD_DIM, KDA_W, generator=gen).to( + CONV_DTYPE + ), + "ssm": torch.rand( + KDA_NUM_HEADS, + KDA_HEAD_DIM, + KDA_HEAD_DIM, + generator=gen, + dtype=SSM_DTYPE, + ), + } + results.append(layers) + return results + + +def run_kda_transfer_test(ctx_tp: int, gen_tp: int, enable_attention_dp: bool = False): + """Loopback: matched ctx/gen TP, replicated full-size KDA state per rank.""" + ctx_mgrs = _create_kda_managers(ctx_tp, enable_attention_dp=enable_attention_dp) + gen_mgrs = _create_kda_managers(gen_tp, enable_attention_dp=enable_attention_dp) + ctx_tcs, gen_tcs = [], [] + try: + for mgr in ctx_mgrs + gen_mgrs: + mgr._impl.mamba_cache.conv.zero_() + mgr._impl.mamba_cache.temporal.zero_() + + config = CacheTransceiverConfig( + backend="NIXL", + transceiver_runtime="PYTHON", + max_tokens_in_buffer=512, + ) + ctx_tcs = _create_transceivers(ctx_tp, ctx_mgrs, config) + gen_tcs = _create_transceivers(gen_tp, gen_mgrs, config) + ctx_endpoint = ctx_tcs[0]._context_info_endpoint + + sampling_params = SamplingParams() + ctx_rids, gen_rids, ctx_reqs, gen_reqs = [], [], [], [] + for req_idx, req_len in enumerate(REQUEST_LENGTHS): + unique_rid = uuid.uuid4().int & 0x7FFFFFFFFFFFFFFF + ctx_rid, gen_rid = req_idx * 2, req_idx * 2 + 1 + ctx_rids.append(ctx_rid) + gen_rids.append(gen_rid) + sc = tensorrt_llm.bindings.SamplingConfig(sampling_params._get_sampling_config()) + ctx_req = LlmRequest( + request_id=ctx_rid, + max_new_tokens=1, + input_tokens=list(range(req_len)), + sampling_config=sc, + is_streaming=False, + llm_request_type=LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY, + ) + ctx_req.py_disaggregated_params = DisaggregatedParams(disagg_request_id=unique_rid) + gen_req = LlmRequest( + request_id=gen_rid, + max_new_tokens=1, + input_tokens=list(range(req_len)), + sampling_config=sc, + is_streaming=False, + llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, + ) + gen_req.py_disaggregated_params = DisaggregatedParams( + ctx_request_id=ctx_rid, + ctx_dp_rank=0, + ctx_info_endpoint=ctx_endpoint, + disagg_request_id=unique_rid, + ) + ctx_reqs.append(ctx_req) + gen_reqs.append(gen_req) + + ctx_batch = ScheduledRequests() + ctx_batch.reset_context_requests(ctx_reqs) + for mgr in ctx_mgrs: + mgr.prepare_resources(ctx_batch) + gen_batch = ScheduledRequests() + gen_batch.reset_context_requests(gen_reqs) + for mgr in gen_mgrs: + mgr.prepare_resources(gen_batch) + for req in ctx_reqs + gen_reqs: + req.context_current_position = req.prompt_len + req.add_new_token(req.prompt_len, 0) + for mgr in ctx_mgrs: + mgr.update_resources(ctx_batch) + for mgr in gen_mgrs: + mgr.update_resources(gen_batch) + + # Ground truth: identical full state on every ctx rank (replicated). + ground_truth = _generate_ground_truth(len(REQUEST_LENGTHS)) + for mgr in ctx_mgrs: + for req_idx, rid in enumerate(ctx_rids): + slot = mgr.mamba_cache_index[rid] + for layer_idx in mgr._impl.mamba_layer_offsets: + full = ground_truth[req_idx][layer_idx] + mgr.get_conv_states(layer_idx)[slot] = full["conv"] + mgr.get_ssm_states(layer_idx)[slot] = full["ssm"] + + for rank in range(gen_tp): + for req in gen_reqs: + gen_tcs[rank].request_and_receive_async(req) + for rank in range(ctx_tp): + for req in ctx_reqs: + ctx_tcs[rank].respond_and_send_async(req) + _run_concurrent( + ctx_tcs, lambda tc: tc.check_context_transfer_status(None, mark_complete=True) + ) + _run_concurrent(gen_tcs, lambda tc: tc.check_gen_transfer_status(None)) + + # Transfer-size metric must include the fixed-size KDA state — the actual + # transferred bytes must cover the computed payload size: per rank, + # num_layers * (conv + ssm) slot bytes on top of any KV bytes. + kda_bytes_per_rank = NUM_KDA_LAYERS * (CONV_SLOT_BYTES + SSM_SLOT_BYTES) + rank_factor = 1 if enable_attention_dp else gen_tp + for req in gen_reqs: + assert req.py_kv_cache_xfer_bytes >= kda_bytes_per_rank * rank_factor, ( + f"kv_cache_xfer_bytes={req.py_kv_cache_xfer_bytes} misses the KDA " + f"state payload ({kda_bytes_per_rank} bytes/rank x {rank_factor})" + ) + + # Bitwise comparison on every gen rank (state is replicated). + for gen_rank, mgr in enumerate(gen_mgrs): + for req_idx, rid in enumerate(gen_rids): + slot = mgr.mamba_cache_index[rid] + for layer_idx in mgr._impl.mamba_layer_offsets: + full = ground_truth[req_idx][layer_idx] + for name, getter in ( + ("conv", mgr.get_conv_states), + ("ssm", mgr.get_ssm_states), + ): + torch.testing.assert_close( + getter(layer_idx)[slot].cpu(), + full[name], + rtol=0, + atol=0, + msg=lambda m, n=name, r=gen_rank, ri=req_idx, li=layer_idx: ( + f"{n} mismatch: gen_rank={r} req={ri} layer={li} " + f"ctx_tp={ctx_tp} gen_tp={gen_tp}: {m}" + ), + ) + + finally: + for tc in ctx_tcs + gen_tcs: + tc.shutdown() + for mgr in ctx_mgrs + gen_mgrs: + mgr.shutdown() + + +@pytest.mark.timeout(180) +@pytest.mark.parametrize( + "ctx_tp,gen_tp", + [(1, 1), (2, 2)], + ids=["tp1_tp1", "tp2_tp2_ep_prescaled"], +) +def test_kda_transfer(ctx_tp, gen_tp): + """KDA two-dtype state transfer, matched parallelism, real NIXL loopback.""" + run_kda_transfer_test(ctx_tp, gen_tp) diff --git a/tests/unittest/disaggregated/test_peer.py b/tests/unittest/disaggregated/test_peer.py index 96a3a95186aa..604d4ab8d93e 100644 --- a/tests/unittest/disaggregated/test_peer.py +++ b/tests/unittest/disaggregated/test_peer.py @@ -47,8 +47,16 @@ pytestmark = pytest.mark.cpu_only -def make_page_table(pool_ptrs=None, block_bytes=None, global_layer_ids=None): - """Create a KVCachePageTable for testing.""" +def make_page_table(pool_ptrs=None, block_bytes=None, global_layer_ids=None, mamba_tp=2): + """Create a KVCachePageTable for testing. + + ``mamba_tp`` is the TP degree the synthetic recurrent-state group is + sharded for: per-rank sizes are the ``mamba_tp`` shard of a fixed global + state, so page tables built for different TP degrees stay + peer-compatible under MambaPolicy's global (TP-aggregated) size check. + Pass the rank's effective mamba TP (``make_rankinfo`` defaults to + ``tp_size=2``). + """ if pool_ptrs is None: pool_ptrs = [1234] if block_bytes is None: @@ -87,9 +95,9 @@ def make_page_table(pool_ptrs=None, block_bytes=None, global_layer_ids=None): mamba_lg = MambaLayerGroup( pool_group_idx=1, mamba_layer_offsets={100: 0, 101: 1}, - conv_states=PhysicalPool(base_address=0xA000, slot_bytes=2048, num_slots=128), - ssm_states=PhysicalPool(base_address=0xB000, slot_bytes=4096, num_slots=128), - conv_section_bytes=[512, 256, 256], + conv_states=PhysicalPool(base_address=0xA000, slot_bytes=4096 // mamba_tp, num_slots=128), + ssm_states=PhysicalPool(base_address=0xB000, slot_bytes=8192 // mamba_tp, num_slots=128), + conv_section_bytes=[1024 // mamba_tp, 512 // mamba_tp, 512 // mamba_tp], ssm_bytes_per_head=64, ) pool_groups = [PhysicalPoolGroup(pools=physical_pools)] @@ -558,7 +566,7 @@ def test_peer_registrar_get_kv_map_head_mismatch(): tokens_per_block=16, dims_per_head=8, layer_num_per_pp=[2], - page_table=make_page_table(block_bytes=[2048]), + page_table=make_page_table(block_bytes=[2048], mamba_tp=1), ) reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) mapper = reg.get_kv_map(peer_ri, (0, 0), (0, 0)) @@ -1026,7 +1034,7 @@ def test_peer_registrar_allows_byte_aligned_subbyte_head_mismatch(): element_bytes=0.5, kv_heads_per_rank=4, tp_size=1, - page_table=make_page_table(block_bytes=[2048]), + page_table=make_page_table(block_bytes=[2048], mamba_tp=1), ) reg = _make_peer_registrar(self_ri) reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) @@ -1056,7 +1064,7 @@ def test_peer_registrar_rejects_misaligned_subbyte_head_mismatch(): tokens_per_block=1, dims_per_head=1, tp_size=1, - page_table=make_page_table(block_bytes=[2048]), + page_table=make_page_table(block_bytes=[2048], mamba_tp=1), ) reg = _make_peer_registrar(self_ri) @@ -1066,7 +1074,7 @@ def test_peer_registrar_rejects_misaligned_subbyte_head_mismatch(): def test_peer_registrar_dispatches_nhd_mapper(): self_pt = make_page_table() - peer_pt = make_page_table(block_bytes=[2048]) + peer_pt = make_page_table(block_bytes=[2048], mamba_tp=1) self_pt.layer_groups[0].pool_views[0].mapper_kind = MapperKind.NHD peer_pt.layer_groups[0].pool_views[0].mapper_kind = MapperKind.NHD self_ri = make_rankinfo( @@ -1090,7 +1098,7 @@ def test_peer_registrar_dispatches_nhd_mapper(): def test_peer_registrar_warns_for_nhd_head_mismatch(monkeypatch): self_pt = make_page_table() - peer_pt = make_page_table(block_bytes=[2048]) + peer_pt = make_page_table(block_bytes=[2048], mamba_tp=1) self_pt.layer_groups[0].pool_views[0].mapper_kind = MapperKind.NHD peer_pt.layer_groups[0].pool_views[0].mapper_kind = MapperKind.NHD self_ri = make_rankinfo(kv_heads_per_rank=2, tp_size=2, page_table=self_pt)