diff --git a/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md b/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md new file mode 100644 index 000000000000..dfa4daac7549 --- /dev/null +++ b/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md @@ -0,0 +1,361 @@ +# Deployment Guide for Kimi K3 on TensorRT LLM - Blackwell + +## Introduction + +This deployment guide provides step-by-step instructions for running the Kimi K3 model using TensorRT LLM on NVIDIA Blackwell GPUs. The deployment configurations and results in this guide were validated on NVIDIA GB300 NVL GPUs. It covers the complete setup required; from preparing the model weights and building the software environment to configuring TensorRT LLM parameters, launching the multi-node server, and validating inference output and accuracy. + +Kimi K3 is a large hybrid Mixture-of-Experts (MoE) model. Its 93 decoder layers interleave two attention families: most layers use Kimi Delta Attention (KDA), a linear-attention mixer that keeps a constant-size recurrent state per sequence, while every fourth layer uses full Multi-head Latent Attention (MLA). The MoE block routes each token to 16 of 896 experts; the routed expert weights are stored in MXFP4 (group size 32) while the attention, shared-expert, dense-MLP, and LM-head weights stay in BF16. In TensorRT LLM the model is served through the `KimiK3ForConditionalGeneration` / `KimiLinearForCausalLM` architectures (the text decoder path). + +This guide uses Slurm and the `trtllm-llmapi-launch` multi-node launcher. The configuration walkthrough focuses on the **DEP16** high-throughput deployment: attention runs data-parallel across 16 ranks and the 896 experts shard across the expert-parallel group (`enable_attention_dp: true`, `moe_expert_parallel_size: 16`). The repository also provides the **TEP16** low-latency deployment and an 8-GPU deployment, **TEP8**. + +## Prerequisites + +* GPU: NVIDIA Blackwell GPUs. DEP16 and TEP16 use 16 GPUs; the TEP8 recipe uses 8 GPUs. These deployment recipes were validated on GB300 NVL GPUs. The repository's Slurm examples assume 4 GPUs per node. Other GPU architectures are not currently supported. +* Multi-node launcher: Slurm with the pyxis/enroot container plugin (or an equivalent MPI launcher) to start one rank per GPU across the nodes. +* High-speed inter-node interconnect (e.g., NVLink/InfiniBand) for the expert-parallel traffic. +* Shared filesystem visible to all nodes for the repository, the model weights, and the configuration file. +* OS: Linux +* Drivers: CUDA Driver 580 or later +* Container runtime with NVIDIA GPU support on each node + +## Models + +* A complete Hugging Face-format Kimi K3 checkpoint and tokenizer, e.g. [moonshotai/Kimi-K3](https://huggingface.co/moonshotai/Kimi-K3) downloaded from the Hugging Face Hub (`model_type: kimi_k3`, architecture `KimiK3ForConditionalGeneration`). The checkpoint ships the routed-expert weights pre-quantized to MXFP4; no additional quantization step is needed. + +The checkpoint and the configuration file must live on a shared filesystem visible to all nodes. The examples below use `/path/to/kimi-k3-checkpoint` — replace it with your local path. + +## Feature Support Notes + +* **Blackwell only.** NVIDIA Blackwell GPUs are supported. The configurations and results in this guide were validated on NVIDIA GB300 NVL GPUs. Support for other GPU architectures may be added in a future release. +* **High-throughput and low-latency deployments are provided.** DEP16 (`enable_attention_dp: true`, `moe_expert_parallel_size: 16`) is the high-throughput deployment. TEP16 (`enable_attention_dp: false`, `moe_expert_parallel_size: 16`) is the low-latency deployment. An 8-GPU deployment, TEP8 (`enable_attention_dp: false`, `moe_expert_parallel_size: 8`), is also provided. Select the deployment and concurrency appropriate for your workload. +* **CUDA graphs and the overlap scheduler are enabled.** The performance-sweep recipes set `disable_overlap_scheduler: false` and enable CUDA graphs. DEP16 additionally sets `cuda_graph_config.enable_padding: true`. +* **Chunked prefill is supported and enabled** (`enable_chunked_prefill: true`), so prompts longer than `max_num_tokens` are scheduled across multiple steps. +* **`kv_cache_config.tokens_per_block` must be `64`** — required by the MLA (576, 512) trtllm-gen generation kernel. +* **Speculative decoding and disaggregated serving are not yet available** for Kimi K3; support is under development. See the "Current limitations" section of `examples/kimi_k3/README.md`. + +## Deployment Steps + +### Build TensorRT LLM from Source + +Kimi K3 support currently requires TensorRT LLM built from source and installed in place. Inside the TensorRT LLM container, from the repository root: + +```bash +python3 scripts/build_wheel.py --cuda_architectures 103-real --skip_building_wheel --yes +.venv-3.12/bin/python -m pip install -e . +``` + +`build_wheel.py` creates the virtual environment at the repository root, named after the container's Python version: `.venv-3.12` for the current containers (Python 3.12). If your container ships a different Python, substitute the matching `.venv-.` path in the commands on this page. Adjust `--cuda_architectures` to the target GPUs (`103-real` for GB300). The multi-node jobs below run TensorRT LLM from this in-place environment, so build and install with the repository at the same path the jobs use. + +Kimi K3 additionally depends on `fla` and `einops`, installed into the same in-place environment (these dependencies might be removed in future releases, replaced by other kernels): + +```bash +.venv-3.12/bin/python -m pip install fla-core einops +``` + +To use the optimized CuTeDSL MLA kernel, install the FlashInfer revision used by the Kimi K3 example into the same in-place environment: + +```bash +.venv-3.12/bin/python -u -m pip install --force-reinstall --no-deps \ + --no-build-isolation \ + "flashinfer-python[cu13] @ git+https://github.com/PerkzZheng/flashinfer-k3.git@b6cc594918baf76c40c3a6236fd53f0f8fb9d2dc" +``` + +The `packaging>=24.2` requirement of this source build is already satisfied by `requirements.txt`. The TensorRT LLM environment already provides FlashInfer's runtime dependencies. The `--no-deps` option prevents `pip` from replacing the pinned PyTorch, Triton, CUDA, and CuTeDSL packages. Install FlashInfer after TensorRT LLM because a later dependency-resolving TensorRT LLM installation can replace this source revision with the currently pinned `flashinfer-python==0.6.16`. + +For general build-from-source instructions see [https://nvidia.github.io/TensorRT-LLM/latest/installation/build-from-source.html](https://nvidia.github.io/TensorRT-LLM/latest/installation/build-from-source.html). + +### Recommended Performance Settings + +**Treat these as a starting point and tune the parameters for your workload.** + +Create a YAML configuration file on the shared filesystem. The settings below are the tested DEP16 configuration, matching `examples/kimi_k3/eval_extra_llm_options.yaml`: + +```shell +EXTRA_LLM_API_FILE=/path/to/kimi-k3-config.yml + +cat << EOF > ${EXTRA_LLM_API_FILE} +# Kimi K3 DEP16 deployment: attention runs data-parallel and the 896 +# experts shard across the expert-parallel group. +tensor_parallel_size: 16 +enable_attention_dp: true +moe_expert_parallel_size: 16 +max_batch_size: 32 +max_num_tokens: 8192 +max_seq_len: 8192 +trust_remote_code: true +disable_overlap_scheduler: false +enable_chunked_prefill: true +cuda_graph_config: + enable_padding: true + max_batch_size: 32 +moe_config: + max_num_tokens: 33024 + use_low_precision_moe_combine: true +kv_cache_config: + enable_block_reuse: false + free_gpu_memory_fraction: 0.25 + tokens_per_block: 64 +EOF +``` + +Notes: + +* The configuration uses `free_gpu_memory_fraction: 0.25` to leave runtime headroom. Each KDA layer also keeps a per-request recurrent state outside the paged KV pool, so reduce this value if the deployment runs out of memory. +* This YAML is specifically the DEP16 high-throughput configuration. For the TEP16 low-latency and TEP8 8-GPU configurations, use `examples/kimi_k3/perf_sweep/perf_sweep.sbatch`; changing only `enable_attention_dp` does not reproduce those recipes. + +### Launch the TensorRT LLM Server + +Kimi K3 is launched through the `trtllm-llmapi-launch` wrapper, which sets up the multi-rank (MPI/Slurm) environment that the parallel server requires. The wrapper is run once per rank by Slurm (`srun`), with one task (rank) per GPU. The example below launches the server across 4 nodes (`-N 4`), 4 GPUs per node (`--ntasks-per-node 4`, 16 ranks total), using the YAML file to drive parallelism, batching, and the KV-cache constraints: + +```bash +MODEL=/path/to/kimi-k3-checkpoint +REPO=/path/to/TensorRT-LLM # in-place build from the previous step + +srun -N 4 \ + --ntasks 16 --ntasks-per-node 4 \ + --mpi=pmix --gpus-per-node=4 \ + --container-image=/path/to/tensorrt-llm-container.sqsh \ + --container-mount-home \ + --container-mounts=${REPO}:${REPO},${MODEL}:${MODEL},${EXTRA_LLM_API_FILE}:${EXTRA_LLM_API_FILE}:ro \ + bash -c " + ulimit -n 65536 + + # Node-local Triton cache: ~/.triton on shared NFS races across + # ranks during autotune JIT (stale file handles). + export TRITON_CACHE_DIR=/tmp/triton-cache-rank\${SLURM_PROCID:-0} + mkdir -p \"\$TRITON_CACHE_DIR\" + + # Node-local flashinfer cache: the default (\$HOME) is shared NFS and + # races across ranks during cubin download (stale file handles). + export FLASHINFER_WORKSPACE_BASE=/tmp/flashinfer-rank\${SLURM_PROCID:-0} + + # Run from the in-place installation created by build_wheel.py. + export PATH=\"${REPO}/.venv-3.12/bin:\$PATH\" + + exec ${REPO}/tensorrt_llm/llmapi/trtllm-llmapi-launch \ + trtllm-serve ${MODEL} \ + --host 0.0.0.0 --port 8000 \ + --config ${EXTRA_LLM_API_FILE} + " +``` + +> [!NOTE] +> Adjust `-N`, `--ntasks`, `--ntasks-per-node`, and `--gpus-per-node` to match your cluster's GPUs-per-node; the total number of tasks (ranks) must equal `tensor_parallel_size` (`16`). Add the partition / account / time flags (`-p`, `-A`, `-t`) required by your Slurm setup, and ensure the repository, checkpoint, and YAML paths resolve identically on all nodes (shared filesystem). + +TensorRT LLM will load weights and select the best kernels during startup. The server is successfully launched when the following log is shown: + +```log +INFO: Started server process [xxxxx] +INFO: Waiting for application startup. +INFO: Application startup complete. +INFO: Uvicorn running on http://localhost:8000 (Press CTRL+C to quit) +``` + +### Quick Start Without a Server (LLM API) + +For a first functional check without standing up a server, the repository ships a ready-made multi-node quick-start job that loads the model once, runs four prompts, and reports whether each response contains the expected text (a successful run reports `True` for all four checks): + +```bash +sbatch examples/kimi_k3/quick_start_kimi_k3.sbatch \ + --model /path/to/kimi-k3-checkpoint \ + --image /path/to/tensorrt-llm-container.sqsh +``` + +See [`examples/kimi_k3/README.md`](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/kimi_k3) for details, including how to override the Slurm partition, account, and time limit. + +### Key YAML Options + +These options are set within the YAML file passed to `trtllm-serve` via the `--config` argument. Only the Kimi-K3-relevant knobs are described here; see the [`TorchLlmArgs` class](https://nvidia.github.io/TensorRT-LLM/llm-api/reference.html#tensorrt_llm.llmapi.TorchLlmArgs) for the full reference. + +#### `tensor_parallel_size` + +* **Description:** The width of the parallel group. Set it to `16` for DEP16 and TEP16, or `8` for TEP8. + +#### `enable_attention_dp` + +* **Description:** Set it to `true` for the DEP16 high-throughput deployment. The TEP16 low-latency and TEP8 8-GPU deployments set it to `false`. + +#### `moe_expert_parallel_size` + +* **Description:** The expert-parallel width for the MoE layers. Set it to `16` for DEP16 and TEP16, or `8` for TEP8. + +#### `cuda_graph_config` + +* **Description:** The standard DEP16, TEP16, and TEP8 performance recipes enable CUDA graphs. DEP16 also sets `enable_padding: true`, which pads generation batches to a captured size. Set the CUDA-graph `max_batch_size` to the same value as the top-level `max_batch_size`. + +#### `disable_overlap_scheduler` + +* **Description:** Set it to `false` to enable the overlap scheduler, as used by the standard DEP16, TEP16, and TEP8 performance recipes. + +#### `moe_config` + +* **Options:** + + * `max_num_tokens`: The DEP16 performance recipe uses `33024` to reduce serialized all-to-all chunk rounds during prefill. + * `use_low_precision_moe_combine`: Set to `true` in the DEP16 performance recipe to use a low-precision expert-parallel combine payload. + +#### `enable_chunked_prefill` + +* **Description:** The recommended DEP16 configuration sets this option to `true`. Prompts longer than `max_num_tokens` are then scheduled across multiple prefill steps; this is verified on 16-GPU DEP16 at GSM8K parity with the unchunked baseline. + +#### `kv_cache_config` + +* **Options:** + * `enable_block_reuse`: Off by default; set to `true` to enable prefix-cache reuse across requests. + * `mamba_state_config.periodic_snapshot_interval`: With block reuse on, the KDA recurrent state is snapshotted every this many tokens so prefix hits can restore it (default `0` = snapshots off; hybrid models only expose reusable prefixes at snapshot boundaries, so set e.g. `256` for block reuse to engage; see `examples/kimi_k3/eval_extra_llm_options_reuse.yaml`). + * `tokens_per_block`: Must be `64`, required by the MLA (576, 512) trtllm-gen generation kernel. + * `free_gpu_memory_fraction`: Fraction of free GPU memory reserved for the paged KV cache after model load. The configuration above uses `0.25` to leave runtime headroom. Lower it if you hit out-of-memory errors. + +#### `trust_remote_code` + +* **Description:** Required to load the Kimi K3 configuration and tokenizer code shipped with the checkpoint. + +## Testing API Endpoint + +The server (the OpenAI-compatible REST endpoint) runs on the rank-0 node, listening on port `8000`. Send requests to that node's hostname or IP; `localhost` only works from the rank-0 node itself. + +### Health Check + +```shell +curl -s -o /dev/null -w "Status: %{http_code}\n" "http://localhost:8000/health" +``` + +When the `Status: 200` code is returned, the server is ready for queries. Note that the very first query may take longer due to initialization and compilation. + +### Basic Test + +> **Note:** The `/v1/chat/completions` endpoint requires the Kimi K3 chat template and serving parsers, which are being added in [TRTLLM-14814](https://github.com/NVIDIA/TensorRT-LLM/pull/17327). Until that change lands, use the `/v1/completions` endpoint with a plain `prompt` string instead. + +After the TensorRT LLM server is set up and shows `Application startup complete`, you can send requests to the server: + +```shell +curl http://localhost:8000/v1/chat/completions -H "Content-Type: application/json" -d '{ + "model": "/path/to/kimi-k3-checkpoint", + "messages": [ + { + "role": "user", + "content": "Where is New York?" + } + ], + "max_tokens": 64, + "temperature": 0 +}' -w "\n" +``` + +The response should contain a `choices[0].message.content` field completing the request, plus a `usage` section with the token counts. + +## Running Evaluations to Verify Accuracy (Optional) + +The repository ships a ready-made multi-node GSM8K evaluation job built on `trtllm-eval` with the tested DEP16 settings: + +```bash +sbatch examples/kimi_k3/run_gsm8k_kimi_k3.sbatch \ + --model /path/to/kimi-k3-checkpoint \ + --image /path/to/tensorrt-llm-container.sqsh +``` + +The job writes progress and results to `kimi-k3-gsm8k-.log` in the submission directory. If no local dataset path is configured, `trtllm-eval` downloads GSM8K from the Hugging Face Hub. The completed log contains a results table with the normalized GSM8K exact-match scores. With the tested checkpoint and the settings in this example, users should expect approximately: + +| Filter | Exact match | +| :-- | --: | +| Flexible extract | 96.51 | +| Strict match | 96.44 | + +The expected average accuracy is approximately 96.47. Small differences (roughly ±0.5 points) are possible with different checkpoint or dependency revisions. + +## Benchmarking Performance + +### Run the End-to-End Performance Sweep + +Use `examples/kimi_k3/perf_sweep/perf_sweep.sbatch` to reproduce the performance measurements. The script generates the tuned serving configuration for the selected deployment mode, launches the server, waits for it to become healthy, runs a warmup, and measures each requested concurrency: + +```bash +sbatch examples/kimi_k3/perf_sweep/perf_sweep.sbatch \ + --mode dep16 \ + --model /path/to/kimi-k3-checkpoint \ + --image /path/to/tensorrt-llm-container.sqsh \ + --isl 8192 \ + --osl 1024 \ + --concurrencies "64 128 256" +``` + +Set `--mode` to `dep16`, `tep16`, or `tep8`. The script derives `max_seq_len` from the requested input and output lengths as `ISL + OSL + 128`, and writes the server log, client log, generated YAML, and per-concurrency JSON results to the job's output directory. + +### Benchmark an Existing Server + +Launch the server with the deployment YAML that you want to measure, as described in [Launch the TensorRT LLM Server](#launch-the-tensorrt-llm-server). From a machine that can reach the rank-0 server, choose the input length, output length, and concurrency for your workload: + +```bash +MODEL=/path/to/kimi-k3-checkpoint +HOST=localhost +PORT=8000 +ISL=8192 +OSL=1024 +CONCURRENCIES="64 128 256" +NUM_ROUNDS=5 +RESULT_DIR=/path/to/kimi-k3-results + +mkdir -p "${RESULT_DIR}" + +for concurrency in ${CONCURRENCIES}; do + python3 -m tensorrt_llm.serve.scripts.benchmark_serving \ + --model "${MODEL}" \ + --host "${HOST}" \ + --port "${PORT}" \ + --dataset-name random \ + --random-ids \ + --tokenize-on-client \ + --random-input-len "${ISL}" \ + --random-output-len "${OSL}" \ + --num-prompts "$((concurrency * NUM_ROUNDS))" \ + --max-concurrency "${concurrency}" \ + --ignore-eos \ + --trust-remote-code \ + --save-result \ + --result-dir "${RESULT_DIR}" \ + --result-filename "concurrency_${concurrency}.json" +done +``` + +Before running the benchmark, adjust `ISL` and `OSL` for the server configuration being measured. Their sum must not exceed the server's `max_seq_len`. To reproduce the 8K input / 1K output results below, use the performance-sweep script, which sets `max_seq_len: 9344` (`ISL + OSL + 128`). + +### Sample Measured Results + +The following results were measured on GB300 NVL GPUs with the 8K/1K sweep. All requests completed. `TPS/user` is `1000 / median TPOT`; it represents the median decode rate seen by one active user. Your results may vary with checkpoint, software, and cluster I/O revisions. + +For the exact server settings and environment variables used to collect the results, refer to `examples/kimi_k3/perf_sweep/perf_sweep.sbatch`. + +| Mode | GPUs | Concurrency | Output tok/s | Output tok/s/GPU | Median TTFT (ms) | Median TPOT (ms) | TPS/user | +| :-- | --: | --: | --: | --: | --: | --: | --: | +| TEP16 | 16 | 4 | 201.2 | 12.6 | 1769.2 | 18.12 | 55.18 | +| DEP16 | 16 | 64 | 1669.1 | 104.3 | 4512.3 | 34.06 | 29.36 | +| DEP16 | 16 | 256 | 3823.9 | 239.0 | 4586.7 | 62.24 | 16.07 | + +The TEP16 point illustrates the low-latency deployment at small concurrency. The DEP16 points illustrate aggregate-throughput scaling at higher concurrencies. + +### Key Metrics + +#### Time to First Token (TTFT) + * The typical time elapsed from when a request is sent until the first output token is generated. + +#### Time Per Output Token (TPOT) and Inter-Token Latency (ITL) + * TPOT is the typical time required to generate each token *after* the first one. + * ITL is the typical time delay between the completion of one token and the completion of the next. + * Both TPOT and ITL ignore TTFT. + +#### End-to-End (E2E) Latency + * The typical total time from when a request is submitted until the final token of the response is received. + +#### Total Token Throughput + * The combined rate at which the system processes both input (prompt) tokens and output (generated) tokens. + +## Troubleshooting Tips + +* **Multi-node startup hangs or ranks can't find each other:** Verify that the repository, the checkpoint, and the YAML resolve identically on all nodes (shared filesystem), that the total Slurm task count equals `tensor_parallel_size` (`16`), and that the inter-node interconnect is healthy. +* **Stale file handle / cache errors during startup:** Point `TRITON_CACHE_DIR` and `FLASHINFER_WORKSPACE_BASE` at node-local storage (e.g., `/tmp`) as shown in the launch command; the shared-NFS defaults race across ranks during autotune JIT and cubin download. +* **CUDA out-of-memory errors:** Reduce `kv_cache_config.free_gpu_memory_fraction`, `max_batch_size`, or `max_seq_len`. Remember that the KDA layers keep per-request FP32 recurrent state outside the paged KV pool. +* **Job reaches its Slurm time limit before the model loads:** Weight loading for a checkpoint of this size is dominated by filesystem throughput; request a longer allocation (`--time=...`), particularly with cold caches or a busy filesystem. +* **Model fails to load:** Make sure `trust_remote_code: true` is set and that the checkpoint path is a complete Hugging Face-format Kimi K3 checkpoint with its tokenizer files. +* **`fla` / `einops` import errors:** Install the extra dependencies into the in-place environment: `.venv-3.12/bin/python -m pip install fla-core einops`. +* **Accuracy looks off:** Confirm `tokens_per_block: 64` (a hard constraint for Kimi K3), and that the remaining settings match the tested configuration above. +* **GPU utilization:** For performance issues, check GPU utilization with `nvidia-smi` on the compute nodes while the server is running. diff --git a/docs/source/deployment-guide/index.rst b/docs/source/deployment-guide/index.rst index eaa43fde0d97..006f63fe87a2 100644 --- a/docs/source/deployment-guide/index.rst +++ b/docs/source/deployment-guide/index.rst @@ -36,5 +36,6 @@ The deployment guides below provide more detailed instructions for serving speci deployment-guide-for-qwen3-on-trtllm.md deployment-guide-for-qwen3.5-on-trtllm.md deployment-guide-for-kimi-k2-thinking-on-trtllm.md + deployment-guide-for-kimi-k3-on-trtllm.md deployment-guide-for-glm-5-on-trtllm.md deployment-guide-for-minimax-m3-on-trtllm.md diff --git a/examples/kimi_k3/README.md b/examples/kimi_k3/README.md new file mode 100644 index 000000000000..800631bbd5d5 --- /dev/null +++ b/examples/kimi_k3/README.md @@ -0,0 +1,196 @@ +# Kimi K3 + +This example runs Kimi K3 with TensorRT-LLM. It includes an LLM API quick +start and configuration for GSM8K evaluation. + +## Hardware support + +Only NVIDIA Blackwell GPUs are currently supported and tested. Support for +other GPU architectures may be added in a future release. + +## Prerequisites + +- TensorRT-LLM built from this repository and installed. Inside + the TensorRT-LLM container, from the repository root: + + ```bash + python3 scripts/build_wheel.py --cuda_architectures 103-real --skip_building_wheel --yes + .venv-3.12/bin/python -m pip install --no-deps -e . + ``` + Using editable mode is recommended for development and testing; see + [build from source](../../docs/source/installation/build-from-source.md) + for details. + + `build_wheel.py` creates the virtual environment at the repository + root, named after the container's Python version: `.venv-3.12` for the + current containers (Python 3.12). If your container ships a different + Python, substitute the matching `.venv-.` path in the + commands on this page. Adjust `--cuda_architectures` to the target + GPUs (`103-real` for GB300). +- A complete Hugging Face-format Kimi K3 checkpoint and tokenizer, e.g. + [moonshotai/Kimi-K3](https://huggingface.co/moonshotai/Kimi-K3) downloaded + from the Hugging Face Hub (the example scripts take a local filesystem + path). +- A Slurm cluster with 16 NVIDIA Blackwell GPUs and a TensorRT-LLM container + image. +- The `fla-core` and `einops` packages, installed into the same in-place + environment. Note: these dependencies might be removed in a future + release, replaced by other kernels. + + ```bash + .venv-3.12/bin/python -m pip install fla-core einops + ``` +- To use the optimized CuTeDSL MLA kernel, install the following FlashInfer + revision into the same in-place environment after installing TensorRT-LLM: + + ```bash + .venv-3.12/bin/python -u -m pip install --force-reinstall --no-deps \ + --no-build-isolation \ + "flashinfer-python[cu13] @ git+https://github.com/PerkzZheng/flashinfer-k3.git@b6cc594918baf76c40c3a6236fd53f0f8fb9d2dc" + ``` + + The `packaging>=24.2` requirement of this source build is already + satisfied by `requirements.txt`. The TensorRT-LLM environment already + provides FlashInfer's runtime + dependencies; `--no-deps` prevents pip from replacing its pinned PyTorch, + Triton, CUDA, and CuTeDSL packages. Install FlashInfer last: TensorRT-LLM + currently pins `flashinfer-python==0.6.16`, so a later + dependency-resolving TensorRT-LLM install can replace this source revision. + +## Run the model + +Kimi K3 requires a multi-node launch. From the repository root, submit the +quick-start Slurm job with the checkpoint and container paths: + +```bash +sbatch examples/kimi_k3/quick_start_kimi_k3.sbatch \ + --model /path/to/kimi-k3-checkpoint \ + --image /path/to/tensorrt-llm-container.sqsh +``` + +Slurm prints the submitted job ID immediately. After the job starts, its output +is written to `kimi-k3-quick-start-.log` in the submission directory. +The model is loaded once, then the log shows four prompts, their generated +text, and whether each response contains the expected text. A successful run +should report `True` for all four checks. + +For a full GSM8K evaluation, submit: + +```bash +sbatch examples/kimi_k3/run_gsm8k_kimi_k3.sbatch \ + --model /path/to/kimi-k3-checkpoint \ + --image /path/to/tensorrt-llm-container.sqsh +``` + +This job writes progress and results to +`kimi-k3-gsm8k-.log`. If no local dataset path is configured, +`trtllm-eval` downloads GSM8K from the Hugging Face Hub. The completed log +contains a results table with the normalized GSM8K exact-match scores. With +the tested checkpoint and the settings in this example, users should expect +approximately: + +| Filter | Exact match | +| :-- | --: | +| Flexible extract | 96.51 | +| Strict match | 96.44 | + +The expected average accuracy is approximately 96.47. Small differences are +possible with different checkpoint or dependency revisions. + +For serving performance, use the standard sweep under +`examples/kimi_k3/perf_sweep/` (this supersedes the older +`run_serving_benchmark_kimi_k3.sbatch` single-recipe benchmark). It submits +the 17-point 8K/1K sweep across the three tuned serving recipes — `tep16` +(latency, c 1–16), `tep8` (interactive, 8 GPUs, c 1–16), and `dep16` +(throughput, c 16–1024): + +```bash +examples/kimi_k3/perf_sweep/submit_perf_sweep.sh \ + --model /path/to/kimi-k3-checkpoint \ + --image /path/to/tensorrt-llm-container.sqsh +``` + +All jobs of a comparison batch are submitted together on purpose: a weight +load overlapping another job's measurement window depresses DEP16 c>=128 +points by ~30%. Use `--jobs "tep16 tep8 dep16-lo dep16-hi"` to select a +subset and `--dry-run` to inspect the sbatch commands. + +Guard the same three recipes with the GSM8K accuracy sweep whenever the +serving configs or kernels change (expect ~96.5 +/- 0.5 per recipe): + +```bash +examples/kimi_k3/perf_sweep/submit_acc_sweep.sh \ + --model /path/to/kimi-k3-checkpoint \ + --image /path/to/tensorrt-llm-container.sqsh +``` + +Results land in per-job `kimi-k3-sweep--.log` files and result +JSONs in the submission directory (run from a fresh results folder). + +Scheduler options must precede the script path; model and image arguments +follow it. The scripts default to the `batch` partition and use an `${account}` +placeholder. Change the corresponding `#SBATCH` settings +to match your cluster, or override them when submitting the job. For example: + +```bash +sbatch --partition=PARTITION --account=ACCOUNT --time=04:00:00 \ + SCRIPT --model MODEL --image IMAGE +``` + +## Troubleshooting + +### The job reaches its time limit + +The default time limits are intentionally aggressive to make the jobs easier +to schedule: 40 minutes for the quick start and two hours for GSM8K. +Depending on how fast your cluster's filesystem loads the weights, a job may +be terminated before producing its final results, particularly with cold +runtime caches or a busy filesystem. + +Request a longer allocation if the time is not sufficient for your environment: + +```bash +sbatch --time=02:00:00 examples/kimi_k3/quick_start_kimi_k3.sbatch \ + --model MODEL --image IMAGE + +sbatch --time=04:00:00 examples/kimi_k3/run_gsm8k_kimi_k3.sbatch \ + --model MODEL --image IMAGE +``` + +## Chunked prefill and KV-cache block reuse + +Chunked prefill is supported and enabled by default in this example +(`enable_chunked_prefill: true` in the quick start and in +`eval_extra_llm_options.yaml`). + +KV-cache block reuse is supported. The LLM API enables it by default +(`KvCacheConfig.enable_block_reuse` defaults to `true`), but the example +configurations here explicitly disable it and treat reuse as an opt-in: +set `kv_cache_config.enable_block_reuse: true`, or use the example +flags — `--enable-block-reuse` for the quick start and `--reuse` for the +GSM8K job (which selects `eval_extra_llm_options_reuse.yaml`): + +```bash +sbatch examples/kimi_k3/quick_start_kimi_k3.sbatch \ + --model MODEL --image IMAGE --enable-block-reuse + +sbatch examples/kimi_k3/run_gsm8k_kimi_k3.sbatch \ + --model MODEL --image IMAGE --reuse +``` + +Unless one of the flags above is passed, these examples run with block +reuse disabled; the tested evaluation and serving configurations use the +default cache manager. + +## Current limitations + +- Pipeline parallelism is not supported. +- **Known performance limitation at DEP16 saturation** (attention-DP + + EP16 throughput recipe): the 8K/1K serving sweep loses several percent + of output throughput at concurrency ≥ 128 (up to ~15% at concurrency + 1024) relative to earlier development measurements; concurrency ≤ 64 + and the TEP16/TEP8 latency recipes are unaffected. Tracked as + TRTLLM-14904. +- FP8 KV cache (`kv_cache_config.dtype: fp8`) is not yet supported. +- Speculative decoding is not yet supported. +- Disaggregated serving is not yet supported. diff --git a/examples/kimi_k3/eval_extra_llm_options.yaml b/examples/kimi_k3/eval_extra_llm_options.yaml new file mode 100644 index 000000000000..813fe8e0dbe8 --- /dev/null +++ b/examples/kimi_k3/eval_extra_llm_options.yaml @@ -0,0 +1,19 @@ +tensor_parallel_size: 16 +enable_attention_dp: true +moe_expert_parallel_size: 16 +max_batch_size: 32 +max_num_tokens: 8192 +max_seq_len: 8192 +trust_remote_code: true +disable_overlap_scheduler: false +enable_chunked_prefill: true +cuda_graph_config: + enable_padding: true + max_batch_size: 32 +moe_config: + max_num_tokens: 33024 + use_low_precision_moe_combine: true +kv_cache_config: + enable_block_reuse: false + free_gpu_memory_fraction: 0.25 + tokens_per_block: 64 diff --git a/examples/kimi_k3/eval_extra_llm_options_reuse.yaml b/examples/kimi_k3/eval_extra_llm_options_reuse.yaml new file mode 100644 index 000000000000..5ab1f698c071 --- /dev/null +++ b/examples/kimi_k3/eval_extra_llm_options_reuse.yaml @@ -0,0 +1,30 @@ +# Identical to eval_extra_llm_options.yaml except for the reuse-specific +# keys: kv_cache_config.enable_block_reuse is flipped to true and +# mamba_state_config.periodic_snapshot_interval is set. Keep every other +# key in sync with the base file when editing either. +tensor_parallel_size: 16 +enable_attention_dp: true +moe_expert_parallel_size: 16 +max_batch_size: 32 +max_num_tokens: 8192 +max_seq_len: 8192 +trust_remote_code: true +disable_overlap_scheduler: false +enable_chunked_prefill: true +cuda_graph_config: + enable_padding: true + max_batch_size: 32 +moe_config: + max_num_tokens: 33024 + use_low_precision_moe_combine: true +kv_cache_config: + enable_block_reuse: true + free_gpu_memory_fraction: 0.25 + tokens_per_block: 64 + # Restore the pre-rework snapshot cadence: upstream deprecated + # mamba_state_cache_interval (default 256) into + # mamba_state_config.periodic_snapshot_interval (default 0 = snapshots + # off), and hybrid models only expose reusable prefixes at snapshot + # boundaries — without this, block reuse silently never engages. + mamba_state_config: + periodic_snapshot_interval: 256 diff --git a/examples/kimi_k3/perf_sweep/acc_sweep.sbatch b/examples/kimi_k3/perf_sweep/acc_sweep.sbatch new file mode 100644 index 000000000000..01d655adfd83 --- /dev/null +++ b/examples/kimi_k3/perf_sweep/acc_sweep.sbatch @@ -0,0 +1,177 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Kimi K3 ACCURACY sweep (GSM8K) for the serving recipes measured by +# perf_sweep.sbatch: same per-mode yaml, same env, same tp/batch/token CLI +# flags (keep the mode blocks below in sync with perf_sweep.sbatch by hand). +# One in-process trtllm-eval run per job; expect ~96.5 +/- 0.5 on the full +# 1319-question set. +# +# Usage: +# sbatch [--nodes=N ...] acc_sweep.sbatch \ +# --mode dep16|tep16|tep8 --model PATH --image PATH [--repo PATH] [--tag NAME] +# +# The #SBATCH partition and account below are placeholders — edit them or +# override on the sbatch command line (sbatch --partition=... --account=...). +# +#SBATCH --job-name=kimi-k3-acc +#SBATCH --partition=batch +#SBATCH --account=${account} +#SBATCH --nodes=4 +#SBATCH --ntasks-per-node=4 +#SBATCH --gpus-per-node=4 +#SBATCH --time=02:00:00 +#SBATCH --output=kimi-k3-acc-%j.log + +set -euo pipefail + +MODE=""; MODEL=""; CONTAINER_IMAGE=""; TAG=""; REPO_ARG=""; ISL=8192; OSL=1024 +while [[ $# -gt 0 ]]; do + case "$1" in + --mode) MODE=$2; shift 2 ;; + --model) MODEL=$2; shift 2 ;; + --image) CONTAINER_IMAGE=$2; shift 2 ;; + --tag) TAG=$2; shift 2 ;; + --repo) REPO_ARG=$2; shift 2 ;; + *) echo "error: unknown argument: $1" >&2; exit 2 ;; + esac +done +[[ -n "$MODE" && -n "$MODEL" && -n "$CONTAINER_IMAGE" ]] || { + echo "error: --mode, --model and --image are required" >&2; exit 2; } +[[ -e "$MODEL" && -e "$CONTAINER_IMAGE" ]] || { echo "error: bad model/image path" >&2; exit 2; } + +# Same batch shape and tuned env as perf_sweep.sbatch (see that script +# for the per-knob rationale). +case "$MODE" in + dep16) + TP_SIZE=16; MAX_BATCH_SIZE=64 + MODE_ENV="KIMI_K3_FP8_WEIGHT_READ=1" + ;; + tep16) + TP_SIZE=16; MAX_BATCH_SIZE=16 + MODE_ENV="KIMI_K3_FP8_WEIGHT_READ=1 KIMI_K3_FP8_WEIGHT_READ_KDA=0 KIMI_K3_FP8_WEIGHT_READ_GATE_UP=1" + # BF16 KDA read + FP8 gate_up at tp16 measured ~40% better + # tps/user at c1 than the base FP8 recipe, at equal-or-better + # GSM8K accuracy. Do NOT add the separated-routing env here: at + # tp16 it costs accuracy on its own and, combined with the FP8 + # gate_up read, collapses GSM8K entirely; that pair is validated + # for tep8 only. + ;; + tep8) + TP_SIZE=8; MAX_BATCH_SIZE=16 + MODE_ENV="KIMI_K3_FP8_WEIGHT_READ=1 KIMI_K3_FP8_WEIGHT_READ_GATE_UP=1 TLLM_TRTLLMGEN_FORCE_SEPARATED_ROUTING=1" + ;; + *) echo "error: unknown --mode: $MODE (dep16|tep16|tep8)" >&2; exit 2 ;; +esac + +MAX_SEQ_LEN=$((ISL + OSL + 128)) +MAX_NUM_TOKENS=$((ISL + MAX_BATCH_SIZE)) + +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +REPO=${REPO_ARG:-$(cd "$SCRIPT_DIR/../../.." && pwd)} +WORK_DIR=$SLURM_SUBMIT_DIR/kimi-k3-acc-$MODE${TAG:+-$TAG}-$SLURM_JOB_ID +mkdir -p "$WORK_DIR" + +# Same serving yaml as perf_sweep.sbatch, verbatim. +CONFIG_FILE=$WORK_DIR/serve_config.yaml +case "$MODE" in + dep16) + cat > "$CONFIG_FILE" <<'EOF' +enable_attention_dp: true +moe_expert_parallel_size: 16 +disable_overlap_scheduler: false +enable_chunked_prefill: true +cuda_graph_config: + max_batch_size: 64 + enable_padding: true +moe_config: + max_num_tokens: 33024 + use_low_precision_moe_combine: true +kv_cache_config: + enable_block_reuse: false + free_gpu_memory_fraction: 0.6 +stream_interval: 10 +EOF + ;; + tep16) + cat > "$CONFIG_FILE" <<'EOF' +enable_attention_dp: false +moe_expert_parallel_size: 16 +disable_overlap_scheduler: false +enable_chunked_prefill: true +cuda_graph_config: + max_batch_size: 16 +kv_cache_config: + enable_block_reuse: false + free_gpu_memory_fraction: 0.25 +stream_interval: 10 +EOF + ;; + tep8) + cat > "$CONFIG_FILE" <<'EOF' +enable_attention_dp: false +moe_expert_parallel_size: 8 +disable_overlap_scheduler: false +enable_chunked_prefill: true +cuda_graph_config: + max_batch_size: 16 +kv_cache_config: + enable_block_reuse: false + free_gpu_memory_fraction: 0.25 +stream_interval: 10 +EOF + ;; +esac + +CACHE_DIR=${CACHE_DIR:-$SLURM_SUBMIT_DIR/.kimi_k3_cache} +mkdir -p "$CACHE_DIR" +ENV_CMD="true" +[[ -n "$MODE_ENV" ]] && ENV_CMD="export $MODE_ENV" +MOUNTS="$REPO:$REPO:rw,$MODEL:$MODEL:ro,$WORK_DIR:$WORK_DIR:rw,$CACHE_DIR:$CACHE_DIR:rw" + +srun --mpi=pmix --kill-on-bad-exit=1 \ + --container-image="$CONTAINER_IMAGE" \ + --container-mount-home \ + --container-mounts="$MOUNTS" \ + bash -c " + set -x + ulimit -n 65536 + export PYTHONNOUSERSITE=1 + $ENV_CMD + export TRITON_CACHE_DIR=/tmp/triton-cache-rank\${SLURM_PROCID:-0} + mkdir -p \"\$TRITON_CACHE_DIR\" + export FLASHINFER_WORKSPACE_BASE=$CACHE_DIR/fi_workspace + export FLASHINFER_CUBIN_DIR=$CACHE_DIR/fi_cubins + # HF dataset/hub cache: \$HOME/.cache is dangling on remote ranks in + # some cluster images, so keep HF under CACHE_DIR (matches + # perf_sweep.sbatch); the gsm8k dataset is served from this cache. + export HF_HOME=\${HF_HOME:-$CACHE_DIR/hf_home} + mkdir -p \"\$FLASHINFER_WORKSPACE_BASE\" \"\$FLASHINFER_CUBIN_DIR\" \"\$HF_HOME\" + # Node-local HF remote-code modules cache: racing mkdir over shared + # \$HOME escapes pathlib's exist_ok on NFS (see run_gsm8k_kimi_k3). + export HF_MODULES_CACHE=/tmp/hf-modules-rank\${SLURM_PROCID:-0} + # Running partial score every N responses (early failure signal). + export TLLM_EVAL_PARTIAL_SCORES_EVERY=\"\${TLLM_EVAL_PARTIAL_SCORES_EVERY:-100}\" + export PATH=\"$REPO/.venv-3.12/bin:\$PATH\" + export PYTHONPATH=\"$REPO\${PYTHONPATH:+:\$PYTHONPATH}\" + # Offline compute nodes cannot reach the Hub: point lm_eval at the + # pre-downloaded snapshot when present (falls back to the hub id). + GSM8K_DATASET_PATH=\${GSM8K_DATASET_PATH:-\$(ls -d \"\$HF_HOME\"/hub/datasets--openai--gsm8k/snapshots/*/ 2>/dev/null | head -1)} + DATASET_ARGS=\"\" + [ -n \"\$GSM8K_DATASET_PATH\" ] && DATASET_ARGS=\"--dataset_path \$GSM8K_DATASET_PATH\" + exec '$REPO/tensorrt_llm/llmapi/trtllm-llmapi-launch' python3 \ + '$REPO/.venv-3.12/bin/trtllm-eval' \ + --model \"$MODEL\" \ + --backend pytorch \ + --tp_size $TP_SIZE \ + --max_batch_size $MAX_BATCH_SIZE \ + --max_num_tokens $MAX_NUM_TOKENS \ + --max_seq_len $MAX_SEQ_LEN \ + --trust_remote_code \ + --config '$CONFIG_FILE' \ + gsm8k \$DATASET_ARGS + " 2>&1 | tee "$WORK_DIR/eval.log" + +grep -E "average accuracy|exact_match" "$WORK_DIR/eval.log" | tail -4 || true +echo "Accuracy run complete; log in $WORK_DIR/eval.log" diff --git a/examples/kimi_k3/perf_sweep/perf_sweep.sbatch b/examples/kimi_k3/perf_sweep/perf_sweep.sbatch new file mode 100644 index 000000000000..db8dfa81f43b --- /dev/null +++ b/examples/kimi_k3/perf_sweep/perf_sweep.sbatch @@ -0,0 +1,264 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Kimi K3 serving perf sweep (8K-in/1K-out): one trtllm-serve launch +# per job, all requested concurrencies measured against it. +# +# Modes bake the tuned serving recipes (measured on GB300 NVL; treat them +# as starting points for your own cluster): +# tep8 TP8 attention + EP8 MoE, 8 GPU / 2 nodes (interactive, c 1-16) +# tep16 TP16 attention + EP16 MoE, 16 GPU / 4 nodes (low latency, c 1-4) +# dep16 attention-DP + EP16 MoE, 16 GPU / 4 nodes (throughput, c 16-1024) +# +# Usage: +# sbatch [--nodes=N ...] perf_sweep.sbatch \ +# --mode dep16|tep16|tep8 --model PATH --image PATH \ +# [--repo PATH] [--concurrencies "16 32 64"] [--isl N] [--osl N] [--tag NAME] +# +# Submit ALL jobs of a comparison batch at the same time: a concurrent +# multi-hundred-GB weight load depresses in-flight DEP16 c>=128 +# measurements by ~30%. +# +# The #SBATCH partition and account below are placeholders — edit them or +# override on the sbatch command line (sbatch --partition=... --account=...). +# +#SBATCH --job-name=kimi-k3-sweep +#SBATCH --partition=batch +#SBATCH --account=${account} +#SBATCH --nodes=4 +#SBATCH --ntasks-per-node=4 +#SBATCH --gpus-per-node=4 +#SBATCH --time=02:00:00 +#SBATCH --output=kimi-k3-sweep-%j.log + +set -euo pipefail + +MODE=""; MODEL=""; CONTAINER_IMAGE=""; TAG=""; ISL=8192; OSL=1024 +CONCURRENCIES=""; REPO_ARG=""; PORT=8000 +while [[ $# -gt 0 ]]; do + case "$1" in + --mode) MODE=$2; shift 2 ;; + --model) MODEL=$2; shift 2 ;; + --image) CONTAINER_IMAGE=$2; shift 2 ;; + --tag) TAG=$2; shift 2 ;; + --isl) ISL=$2; shift 2 ;; + --osl) OSL=$2; shift 2 ;; + --concurrencies) CONCURRENCIES=$2; shift 2 ;; + --repo) REPO_ARG=$2; shift 2 ;; + *) echo "error: unknown argument: $1" >&2; exit 2 ;; + esac +done +[[ -n "$MODE" && -n "$MODEL" && -n "$CONTAINER_IMAGE" ]] || { + echo "error: --mode, --model and --image are required" >&2; exit 2; } +[[ -e "$MODEL" && -e "$CONTAINER_IMAGE" ]] || { echo "error: bad model/image path" >&2; exit 2; } + +# Per-mode batch shape and the tuned env. MAX_BATCH_SIZE is PER-RANK for +# dep16 (attention DP), GLOBAL for the TEP modes. +# KIMI_K3_FP8_WEIGHT_READ=1 opt-in FP8 (lossy) read of the replicated +# BF16 weights — default OFF in the loader, so every mode sets it +# explicitly: all recipes below were tuned and measured with it on. +# KIMI_K3_FP8_WEIGHT_READ_KDA=0 keep the KDA q/k/v/g/o projections in +# BF16 within the enabled master (the most precision-sensitive slice). +# KIMI_K3_FP8_WEIGHT_READ_GATE_UP=1 FP8 shared-expert gate_up read at TEP +# (with the master on it defaults on only under attention DP): +8-9% +# with stream_interval >= 10. +# TLLM_TRTLLMGEN_FORCE_SEPARATED_ROUTING=1 host-side routing; the in-cubin +# top-k tier register-spills at decode batch 5-64 (33 us/layer). +case "$MODE" in + dep16) + TP_SIZE=16; MAX_BATCH_SIZE=64 + DEFAULT_CONCURRENCIES="16 32 64 128 256 512 1024" + MODE_ENV="KIMI_K3_FP8_WEIGHT_READ=1" + ;; + tep16) + TP_SIZE=16; MAX_BATCH_SIZE=16 + DEFAULT_CONCURRENCIES="1 2 4 8 16" + MODE_ENV="KIMI_K3_FP8_WEIGHT_READ=1 KIMI_K3_FP8_WEIGHT_READ_KDA=0 KIMI_K3_FP8_WEIGHT_READ_GATE_UP=1" + # BF16 KDA read + FP8 gate_up at tp16 measured ~40% better + # tps/user at c1 than the base FP8 recipe, at equal-or-better + # GSM8K accuracy. Do NOT add the separated-routing env here: at + # tp16 it costs accuracy on its own and, combined with the FP8 + # gate_up read, collapses GSM8K entirely; that pair is validated + # for tep8 only. + ;; + tep8) + TP_SIZE=8; MAX_BATCH_SIZE=16 + DEFAULT_CONCURRENCIES="1 2 4 8 16" + MODE_ENV="KIMI_K3_FP8_WEIGHT_READ=1 KIMI_K3_FP8_WEIGHT_READ_GATE_UP=1 TLLM_TRTLLMGEN_FORCE_SEPARATED_ROUTING=1" + ;; + *) echo "error: unknown --mode: $MODE (dep16|tep16|tep8)" >&2; exit 2 ;; +esac +CONCURRENCIES=${CONCURRENCIES:-$DEFAULT_CONCURRENCIES} + +MAX_SEQ_LEN=$((ISL + OSL + 128)) +MAX_NUM_TOKENS=$((ISL + MAX_BATCH_SIZE)) + +# Default repo = the checkout containing this script (expects .venv-3.12 +# built in place, custom FlashInfer installed; see examples/kimi_k3/README.md). +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +REPO=${REPO_ARG:-$(cd "$SCRIPT_DIR/../../.." && pwd)} +WORK_DIR=$SLURM_SUBMIT_DIR/kimi-k3-sweep-$MODE${TAG:+-$TAG}-$SLURM_JOB_ID +mkdir -p "$WORK_DIR" + +# Tuned serving config per mode. stream_interval 10 removes the ~18 ms/token +# host streaming floor (TPOT metric unaffected: (latency-TTFT)/(len-1)). +CONFIG_FILE=$WORK_DIR/serve_config.yaml +case "$MODE" in + dep16) + # moe max_num_tokens 33024: fewer serialized A2A chunk rounds per + # prefill wave (TTFT -4-6%). use_low_precision_moe_combine: FP8 A2A + # combine payload, +3.6-4.6% at c64-128 (GSM8K 96.74). enable_padding: + # ADP graphs are all-or-nothing without it. + cat > "$CONFIG_FILE" <<'EOF' +enable_attention_dp: true +moe_expert_parallel_size: 16 +disable_overlap_scheduler: false +enable_chunked_prefill: true +cuda_graph_config: + max_batch_size: 64 + enable_padding: true +moe_config: + max_num_tokens: 33024 + use_low_precision_moe_combine: true +kv_cache_config: + enable_block_reuse: false + free_gpu_memory_fraction: 0.6 +stream_interval: 10 +EOF + ;; + tep16) + cat > "$CONFIG_FILE" <<'EOF' +enable_attention_dp: false +moe_expert_parallel_size: 16 +disable_overlap_scheduler: false +enable_chunked_prefill: true +cuda_graph_config: + max_batch_size: 16 +kv_cache_config: + enable_block_reuse: false + free_gpu_memory_fraction: 0.25 +stream_interval: 10 +EOF + ;; + tep8) + cat > "$CONFIG_FILE" <<'EOF' +enable_attention_dp: false +moe_expert_parallel_size: 8 +disable_overlap_scheduler: false +enable_chunked_prefill: true +cuda_graph_config: + max_batch_size: 16 +kv_cache_config: + enable_block_reuse: false + free_gpu_memory_fraction: 0.25 +stream_interval: 10 +EOF + ;; +esac + +# Shared warm caches: without a pre-populated FLASHINFER_CUBIN_DIR, ranks on +# egress-less compute nodes can hang fetching cubins during FMHA warmup. +CACHE_DIR=${CACHE_DIR:-$SLURM_SUBMIT_DIR/.kimi_k3_cache} +mkdir -p "$CACHE_DIR" +ENV_CMD="true" +[[ -n "$MODE_ENV" ]] && ENV_CMD="export $MODE_ENV" + +MOUNTS="$REPO:$REPO:rw,$MODEL:$MODEL:ro,$WORK_DIR:$WORK_DIR:rw,$CACHE_DIR:$CACHE_DIR:rw" +HEAD_NODE=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | awk 'NR == 1') + +SERVER_PID="" +cleanup() { [[ -n "$SERVER_PID" ]] && kill "$SERVER_PID" 2>/dev/null || true; } +trap cleanup EXIT + +srun --mpi=pmix --kill-on-bad-exit=1 \ + --container-image="$CONTAINER_IMAGE" \ + --container-mount-home \ + --container-mounts="$MOUNTS" \ + --output="$WORK_DIR/serve.log" \ + bash -c " + set -x + ulimit -n 65536 + export PYTHONNOUSERSITE=1 + $ENV_CMD + export TRITON_CACHE_DIR=/tmp/triton-cache-rank\${SLURM_PROCID:-0} + mkdir -p \"\$TRITON_CACHE_DIR\" + export FLASHINFER_WORKSPACE_BASE=$CACHE_DIR/fi_workspace + export FLASHINFER_CUBIN_DIR=$CACHE_DIR/fi_cubins + export HF_HOME=\${HF_HOME:-$CACHE_DIR/hf_home} + mkdir -p \"\$FLASHINFER_WORKSPACE_BASE\" \"\$FLASHINFER_CUBIN_DIR\" \"\$HF_HOME\" + export PATH=\"$REPO/.venv-3.12/bin:\$PATH\" + export PYTHONPATH=\"$REPO\${PYTHONPATH:+:\$PYTHONPATH}\" + # Non-leader MPI workers inherit their TRT-LLM log level from the + # environment rather than the trtllm-serve CLI. + export TLLM_LOG_LEVEL=\${TLLM_LOG_LEVEL:-info} + exec '$REPO/tensorrt_llm/llmapi/trtllm-llmapi-launch' trtllm-serve \ + \"$MODEL\" \ + --backend pytorch --host 0.0.0.0 --port $PORT \ + --tp_size $TP_SIZE --max_batch_size $MAX_BATCH_SIZE \ + --max_num_tokens $MAX_NUM_TOKENS --max_seq_len $MAX_SEQ_LEN \ + --trust_remote_code --config '$CONFIG_FILE' + " & +SERVER_PID=$! + +# Self-cancel if not healthy so the allocation never idles. +deadline=$((SECONDS + 1500)) +until curl -sf -o /dev/null "http://127.0.0.1:$PORT/health"; do + if ! kill -0 "$SERVER_PID" 2>/dev/null || (( SECONDS >= deadline )); then + echo "error: server not healthy within 25 min (see $WORK_DIR/serve.log)" >&2 + scancel "$SLURM_JOB_ID"; sleep 120; exit 1 + fi + sleep 15 +done +echo "Server healthy; sweeping c = $CONCURRENCIES (see $WORK_DIR/client.log)." + +# One warmup pass then ONE measured run per point: num_prompts = c x 5 +# (c < 256) or c x 3, at max-concurrency = c. Metric of record: +# tps/user = 1000 / median_tpot_ms from each concurrency_.json. +srun --overlap --kill-on-bad-exit=1 --nodes=1 --ntasks=1 -w "$HEAD_NODE" \ + --container-image="$CONTAINER_IMAGE" \ + --container-mount-home \ + --container-mounts="$MOUNTS" \ + --output="$WORK_DIR/client.log" \ + bash -c " + set -x + export PYTHONNOUSERSITE=1 + # Same cache redirects as the server block: importing tensorrt_llm + # pulls in flashinfer, which creates \$HOME-based workspace/cache + # dirs, and \$HOME resolves to an unwritable /root in this container. + export TRITON_CACHE_DIR=/tmp/triton-cache-client + export FLASHINFER_WORKSPACE_BASE=$CACHE_DIR/fi_workspace + export FLASHINFER_CUBIN_DIR=$CACHE_DIR/fi_cubins + export HF_HOME=\${HF_HOME:-$CACHE_DIR/hf_home} + export HF_MODULES_CACHE=/tmp/hf-modules-client + mkdir -p \"\$TRITON_CACHE_DIR\" \"\$FLASHINFER_WORKSPACE_BASE\" \\ + \"\$FLASHINFER_CUBIN_DIR\" \"\$HF_HOME\" \"\$HF_MODULES_CACHE\" + export PATH=\"$REPO/.venv-3.12/bin:\$PATH\" + export PYTHONPATH=\"$REPO\${PYTHONPATH:+:\$PYTHONPATH}\" + # Same cache redirects as the server step: importing tensorrt_llm on + # the client pulls in flashinfer, whose workspace defaults under + # \$HOME/.cache — a dangling path on some cluster images. + export FLASHINFER_WORKSPACE_BASE=$CACHE_DIR/fi_workspace + export FLASHINFER_CUBIN_DIR=$CACHE_DIR/fi_cubins + export HF_HOME=\${HF_HOME:-$CACHE_DIR/hf_home} + export TRITON_CACHE_DIR=/tmp/triton-cache-client + mkdir -p \"\$TRITON_CACHE_DIR\" + run_client() { + python3 -m tensorrt_llm.serve.scripts.benchmark_serving \ + --model \"$MODEL\" --host 127.0.0.1 --port $PORT \ + --dataset-name random --random-ids --tokenize-on-client \ + --random-input-len $ISL --random-output-len $OSL \ + --ignore-eos --trust-remote-code \"\$@\" + } + for c in $CONCURRENCIES; do + reps=5; [ \$c -ge 256 ] && reps=3 + run_client --num-prompts \$c --max-concurrency \$c + run_client --num-prompts \$((c * reps)) --max-concurrency \$c \ + --save-result --result-dir '$WORK_DIR' \ + --result-filename \"concurrency_\${c}.json\" + done + " + +echo "Sweep complete; results in $WORK_DIR" +kill "$SERVER_PID" 2>/dev/null || true +wait "$SERVER_PID" 2>/dev/null || true diff --git a/examples/kimi_k3/perf_sweep/submit_acc_sweep.sh b/examples/kimi_k3/perf_sweep/submit_acc_sweep.sh new file mode 100755 index 000000000000..00ac07d8dedc --- /dev/null +++ b/examples/kimi_k3/perf_sweep/submit_acc_sweep.sh @@ -0,0 +1,63 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Submit the Kimi K3 GSM8K ACCURACY runs for the three perf-sweep serving +# recipes (see acc_sweep.sbatch; expect ~96.5 +/- 0.5 each): +# tep16 4 nodes 2:00 tep8 2 nodes 2:00 +# dep16 4 nodes 2:00 +# +# Partition/account come from the acc_sweep.sbatch placeholders — edit +# them there or export SBATCH_PARTITION / SBATCH_ACCOUNT before submitting. +# On Slurm clusters with block topology, export USE_SLURM_SEGMENTS=1 to +# also request one NVLink segment per node group (sbatch --segment). +# +# Usage: submit_acc_sweep.sh --model PATH --image PATH [--repo PATH] \ +# [--jobs "tep16 tep8 dep16"] [--dry-run] +set -euo pipefail + +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +SBATCH_SCRIPT=$SCRIPT_DIR/acc_sweep.sbatch + +MODEL=""; IMAGE=""; REPO=""; JOBS="tep16 tep8 dep16"; DRY_RUN=0 +while [[ $# -gt 0 ]]; do + case "$1" in + --model) MODEL=$2; shift 2 ;; + --image) IMAGE=$2; shift 2 ;; + --repo) REPO=$2; shift 2 ;; + --jobs) JOBS=$2; shift 2 ;; + --dry-run) DRY_RUN=1; shift ;; + *) echo "error: unknown argument: $1" >&2; exit 2 ;; + esac +done +[[ -n "$MODEL" && -n "$IMAGE" ]] || { + echo "Usage: $0 --model PATH --image PATH [--repo PATH] [--jobs ...] [--dry-run]" >&2; exit 2; } + +REPO_ARGS=() +[[ -n "$REPO" ]] && REPO_ARGS=(--repo "$REPO") + +submit() { + local mode=$1 nodes=$2 segment=$3 + local opts=(--job-name="kimi-k3-acc-$mode" --output="kimi-k3-acc-$mode-%j.log" + --nodes="$nodes" --time="02:00:00") + [[ -n "$segment" && "${USE_SLURM_SEGMENTS:-0}" == 1 ]] && opts+=(--segment="$segment") + local args=(--mode "$mode" --model "$MODEL" --image "$IMAGE" "${REPO_ARGS[@]}") + if (( DRY_RUN )); then + echo "DRY RUN: sbatch ${opts[*]} $SBATCH_SCRIPT ${args[*]}" + else + # Checked assignment so a failed sbatch stops the sweep (set -e); + # echo "$(sbatch ...)" would discard the failure. + local result + result=$(sbatch "${opts[@]}" "$SBATCH_SCRIPT" "${args[@]}") + echo "$mode: $result" + fi +} + +for job in $JOBS; do + case "$job" in + tep16) submit tep16 4 "4" ;; + tep8) submit tep8 2 "2" ;; + dep16) submit dep16 4 "4" ;; + *) echo "error: unknown job: $job (tep16|tep8|dep16)" >&2; exit 2 ;; + esac +done diff --git a/examples/kimi_k3/perf_sweep/submit_perf_sweep.sh b/examples/kimi_k3/perf_sweep/submit_perf_sweep.sh new file mode 100644 index 000000000000..9b6bd232644b --- /dev/null +++ b/examples/kimi_k3/perf_sweep/submit_perf_sweep.sh @@ -0,0 +1,71 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Submit the standard Kimi K3 17-point serving sweep (tuned recipes): +# tep16 c 1 2 4 8 16 4 nodes 1:30 +# tep8 c 1 2 4 8 16 2 nodes 1:30 +# dep16-lo c 16 32 64 128 256 4 nodes 2:30 +# dep16-hi c 512 1024 4 nodes 2:30 +# +# All four jobs are submitted TOGETHER on purpose: weight loads overlapping +# another job's measurement window depress DEP16 c>=128 points by ~30% +# (loads overlapping loads are harmless). Run from a fresh results folder. +# +# Partition/account come from the perf_sweep.sbatch placeholders — edit +# them there or export SBATCH_PARTITION / SBATCH_ACCOUNT before submitting. +# On Slurm clusters with block topology, export USE_SLURM_SEGMENTS=1 to +# also request one NVLink segment per node group (sbatch --segment). +# +# Usage: submit_perf_sweep.sh --model PATH --image PATH [--repo PATH] \ +# [--jobs "tep16 tep8 dep16-lo dep16-hi"] [--dry-run] +set -euo pipefail + +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +SBATCH_SCRIPT=$SCRIPT_DIR/perf_sweep.sbatch + +MODEL=""; IMAGE=""; REPO=""; JOBS="tep16 tep8 dep16-lo dep16-hi"; DRY_RUN=0 +while [[ $# -gt 0 ]]; do + case "$1" in + --model) MODEL=$2; shift 2 ;; + --image) IMAGE=$2; shift 2 ;; + --repo) REPO=$2; shift 2 ;; + --jobs) JOBS=$2; shift 2 ;; + --dry-run) DRY_RUN=1; shift ;; + *) echo "error: unknown argument: $1" >&2; exit 2 ;; + esac +done +[[ -n "$MODEL" && -n "$IMAGE" ]] || { + echo "Usage: $0 --model PATH --image PATH [--repo PATH] [--jobs ...] [--dry-run]" >&2; exit 2; } + +REPO_ARGS=() +[[ -n "$REPO" ]] && REPO_ARGS=(--repo "$REPO") + +submit() { + local name=$1 mode=$2 tag=$3 nodes=$4 walltime=$5 segment=$6 concurrencies=$7 + local opts=(--job-name="kimi-k3-sweep-$name" --output="kimi-k3-sweep-$name-%j.log" + --nodes="$nodes" --time="$walltime") + [[ -n "$segment" && "${USE_SLURM_SEGMENTS:-0}" == 1 ]] && opts+=(--segment="$segment") + local args=(--mode "$mode" --model "$MODEL" --image "$IMAGE" + --concurrencies "$concurrencies" "${REPO_ARGS[@]}") + [[ -n "$tag" ]] && args+=(--tag "$tag") + if (( DRY_RUN )); then + echo "DRY RUN: sbatch ${opts[*]} $SBATCH_SCRIPT ${args[*]}" + else + # Checked assignment so a failed sbatch stops the sweep (set -e); + # echo "$(sbatch ...)" would discard the failure. + local result + result=$(sbatch "${opts[@]}" "$SBATCH_SCRIPT" "${args[@]}") + echo "$name: $result" + fi +} + +for job in $JOBS; do + case "$job" in + tep16) submit tep16 tep16 "" 4 "01:30:00" "4" "1 2 4 8 16" ;; + tep8) submit tep8 tep8 "" 2 "01:30:00" "2" "1 2 4 8 16" ;; + dep16-lo) submit dep16-lo dep16 "lo" 4 "02:30:00" "4" "16 32 64 128 256" ;; + dep16-hi) submit dep16-hi dep16 "hi" 4 "02:30:00" "4" "512 1024" ;; + *) echo "error: unknown job: $job" >&2; exit 2 ;; + esac +done diff --git a/examples/kimi_k3/quick_start_kimi_k3.py b/examples/kimi_k3/quick_start_kimi_k3.py new file mode 100644 index 000000000000..c1db17bc9272 --- /dev/null +++ b/examples/kimi_k3/quick_start_kimi_k3.py @@ -0,0 +1,89 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Run Kimi K3 with the TensorRT-LLM LLM API. + +Example: + trtllm-llmapi-launch python3 examples/kimi_k3/quick_start_kimi_k3.py \ + --model /path/to/kimi-k3-checkpoint --tp-size 16 +""" + +import argparse + +from tensorrt_llm import LLM, SamplingParams +from tensorrt_llm.llmapi import CudaGraphConfig, KvCacheConfig, MambaStateConfig + +SAMPLES = [ + ("The capital of France is", "Paris"), + ("1 + 1 = 2, 2 + 2 = 4, 4 + 4 =", "8"), + ("Water is made of hydrogen and", "oxygen"), + ( + "Question: Natalia sold clips to 48 of her friends in April, and " + "then she sold half as many clips in May. How many clips did " + "Natalia sell altogether in April and May?\nAnswer:", + "#### 72", + ), +] + + +def parse_arguments() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--model", + required=True, + help="Path to a local Kimi K3 checkpoint", + ) + parser.add_argument( + "--tp-size", + type=int, + default=16, + help="Number of GPUs used for expert parallelism (default: 16).", + ) + parser.add_argument( + "--enable-block-reuse", + action="store_true", + help="Enable KV-cache block reuse (unified-pool hybrid cache " + "manager with KDA recurrent-state snapshots; prefix-cache hits " + "skip recomputing shared prompt prefixes).", + ) + return parser.parse_args() + + +def main() -> None: + args = parse_arguments() + llm = LLM( + model=args.model, + tensor_parallel_size=args.tp_size, + enable_attention_dp=True, + moe_expert_parallel_size=args.tp_size, + trust_remote_code=True, + max_batch_size=8, + max_seq_len=4096, + max_num_tokens=4096, + enable_chunked_prefill=True, + cuda_graph_config=CudaGraphConfig(max_batch_size=8), + kv_cache_config=KvCacheConfig( + enable_block_reuse=args.enable_block_reuse, + free_gpu_memory_fraction=0.7, + # Hybrid models only expose reusable prefixes at recurrent-state + # snapshot boundaries, and periodic snapshots default to off; + # without this, block reuse would silently never engage. + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=256) + if args.enable_block_reuse + else MambaStateConfig(), + ), + ) + + sampling_params = SamplingParams(max_tokens=64, temperature=0.0) + prompts = [prompt for prompt, _ in SAMPLES] + try: + for output, (_, expected) in zip(llm.generate(prompts, sampling_params), SAMPLES): + generated_text = output.outputs[0].text + print(f"Prompt: {output.prompt!r}") + print(f"Generated text: {generated_text!r}") + print(f"Contains expected text {expected!r}: {expected in generated_text}\n") + finally: + llm.shutdown() + + +if __name__ == "__main__": + main() diff --git a/examples/kimi_k3/quick_start_kimi_k3.sbatch b/examples/kimi_k3/quick_start_kimi_k3.sbatch new file mode 100644 index 000000000000..336753c5305d --- /dev/null +++ b/examples/kimi_k3/quick_start_kimi_k3.sbatch @@ -0,0 +1,116 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Kimi K3 quick start on 16 NVIDIA Blackwell GPUs. +# +# Submit the job with the checkpoint and container paths: +# +# sbatch examples/kimi_k3/quick_start_kimi_k3.sbatch \ +# --model /path/to/kimi-k3-checkpoint \ +# --image /path/to/tensorrt-llm-container.sqsh +# +# The #SBATCH partition and account below are placeholders — edit them or +# override on the sbatch command line (sbatch --partition=... --account=...). +# REPO defaults to the submit directory (override by exporting REPO) and +# must contain the built TensorRT-LLM checkout. +# +#SBATCH --job-name=kimi-k3-quick-start +#SBATCH --partition=batch +#SBATCH --account=${account} +#SBATCH --nodes=4 +#SBATCH --ntasks-per-node=4 +#SBATCH --gpus-per-node=4 +#SBATCH --time=00:40:00 +#SBATCH --output=kimi-k3-quick-start-%j.log + +set -euo pipefail + +usage() { + echo "Usage: sbatch $0 --model PATH --image PATH [--enable-block-reuse]" +} + +MODEL="" +CONTAINER_IMAGE="" +FEATURE_ARGS="" +while [[ $# -gt 0 ]]; do + case "$1" in + --enable-block-reuse) + # Forwarded verbatim to quick_start_kimi_k3.py (see its --help). + FEATURE_ARGS="$FEATURE_ARGS $1" + shift + ;; + --model) + [[ $# -ge 2 ]] || { echo "error: --model requires a value" >&2; usage >&2; exit 2; } + MODEL=$2 + shift 2 + ;; + --model=*) + MODEL=${1#*=} + shift + ;; + --image) + [[ $# -ge 2 ]] || { echo "error: --image requires a value" >&2; usage >&2; exit 2; } + CONTAINER_IMAGE=$2 + shift 2 + ;; + --image=*) + CONTAINER_IMAGE=${1#*=} + shift + ;; + -h|--help) + usage + exit 0 + ;; + *) + echo "error: unknown argument: $1" >&2 + usage >&2 + exit 2 + ;; + esac +done + +[[ -n "$MODEL" ]] || { echo "error: --model is required" >&2; usage >&2; exit 2; } +[[ -n "$CONTAINER_IMAGE" ]] || { echo "error: --image is required" >&2; usage >&2; exit 2; } +[[ -e "$MODEL" ]] || { echo "error: model path does not exist: $MODEL" >&2; exit 2; } +[[ -e "$CONTAINER_IMAGE" ]] || { echo "error: image path does not exist: $CONTAINER_IMAGE" >&2; exit 2; } + +REPO=${REPO:-$SLURM_SUBMIT_DIR} + +MOUNTS="$REPO:$REPO:rw,$MODEL:$MODEL:ro" +# Optional extra container mounts (srun --container-mounts syntax) +if [ -n "${EXTRA_MOUNTS:-}" ]; then + MOUNTS+=",$EXTRA_MOUNTS" +fi + +srun --mpi=pmix \ + --container-image="$CONTAINER_IMAGE" \ + --container-mount-home \ + --container-mounts="$MOUNTS" \ + bash -c " + set -x + ulimit -n 65536 + + # Node-local Triton cache: ~/.triton on shared NFS races across + # ranks during autotune JIT (stale file handles). + export TRITON_CACHE_DIR=/tmp/triton-cache-rank\${SLURM_PROCID:-0} + mkdir -p \"\$TRITON_CACHE_DIR\" + + # Node-local flashinfer cache: the default (\$HOME) is shared NFS and + # races across ranks during cubin download (stale file handles). + export FLASHINFER_WORKSPACE_BASE=/tmp/flashinfer-rank\${SLURM_PROCID:-0} + + # Run python3 from the in-place installation created in the README + # Prerequisites section (build_wheel.py creates .venv-3.12). + export PATH=\"$REPO/.venv-3.12/bin:\$PATH\" + + # Import tensorrt_llm from $REPO, not from wherever the venv's + # in-place install points (for git-worktree submits those differ; + # without this the job silently tests the main checkout's code). + export PYTHONPATH=\"$REPO\${PYTHONPATH:+:\$PYTHONPATH}\" + + exec '$REPO/tensorrt_llm/llmapi/trtllm-llmapi-launch' python3 \ + '$REPO/examples/kimi_k3/quick_start_kimi_k3.py' \ + --model \"$MODEL\" \ + --tp-size 16$FEATURE_ARGS + " diff --git a/examples/kimi_k3/run_gsm8k_kimi_k3.sbatch b/examples/kimi_k3/run_gsm8k_kimi_k3.sbatch new file mode 100644 index 000000000000..0526f1c7c512 --- /dev/null +++ b/examples/kimi_k3/run_gsm8k_kimi_k3.sbatch @@ -0,0 +1,158 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Kimi K3 GSM8K evaluation on 16 NVIDIA Blackwell GPUs. +# +# Submit the job with the checkpoint and container paths: +# +# sbatch examples/kimi_k3/run_gsm8k_kimi_k3.sbatch \ +# --model /path/to/kimi-k3-checkpoint \ +# --image /path/to/tensorrt-llm-container.sqsh +# +# Pass --reuse to additionally enable KV-cache block reuse +# (eval_extra_llm_options_reuse.yaml); chunked prefill is enabled by +# default. +# +# REPO defaults to the submit directory (override by exporting REPO), and +# must contain the built TensorRT-LLM checkout. The #SBATCH partition and +# account below are placeholders — edit them or override on the sbatch +# command line (sbatch --partition=... --account=...), along with any +# other scheduler-specific options your cluster requires. +# +#SBATCH --job-name=kimi-k3-gsm8k +#SBATCH --partition=batch +#SBATCH --account=${account} +#SBATCH --nodes=4 +#SBATCH --ntasks-per-node=4 +#SBATCH --gpus-per-node=4 +#SBATCH --time=02:00:00 +#SBATCH --output=kimi-k3-gsm8k-%j.log + +set -euo pipefail + +usage() { + echo "Usage: sbatch $0 --model PATH --image PATH [--reuse]" +} + +MODEL="" +CONTAINER_IMAGE="" +MODE=default +while [[ $# -gt 0 ]]; do + case "$1" in + --reuse) + MODE=reuse + shift + ;; + --model) + [[ $# -ge 2 ]] || { echo "error: --model requires a value" >&2; usage >&2; exit 2; } + MODEL=$2 + shift 2 + ;; + --model=*) + MODEL=${1#*=} + shift + ;; + --image) + [[ $# -ge 2 ]] || { echo "error: --image requires a value" >&2; usage >&2; exit 2; } + CONTAINER_IMAGE=$2 + shift 2 + ;; + --image=*) + CONTAINER_IMAGE=${1#*=} + shift + ;; + -h|--help) + usage + exit 0 + ;; + *) + echo "error: unknown argument: $1" >&2 + usage >&2 + exit 2 + ;; + esac +done + +[[ -n "$MODEL" ]] || { echo "error: --model is required" >&2; usage >&2; exit 2; } +[[ -n "$CONTAINER_IMAGE" ]] || { echo "error: --image is required" >&2; usage >&2; exit 2; } +[[ -e "$MODEL" ]] || { echo "error: model path does not exist: $MODEL" >&2; exit 2; } +[[ -e "$CONTAINER_IMAGE" ]] || { echo "error: image path does not exist: $CONTAINER_IMAGE" >&2; exit 2; } + +REPO=${REPO:-$SLURM_SUBMIT_DIR} + +# Per-mode LLM options. All engine options live in the selected yaml +# (single source of truth); trtllm-eval lets explicit CLI flags override +# --config, so none are duplicated here. +# - reuse: KV-cache block reuse (enable_block_reuse: true in the yaml). +case "$MODE" in + reuse) + EVAL_CONFIG=$REPO/examples/kimi_k3/eval_extra_llm_options_reuse.yaml + ;; + *) + EVAL_CONFIG=$REPO/examples/kimi_k3/eval_extra_llm_options.yaml + ;; +esac + +# Mount every input at its host path so the paths are valid inside too. +# EXTRA_MOUNTS (comma-separated src:dst:flags) lets git-worktree submits +# also mount the built main checkout that the worktree's artifact +# symlinks resolve into (submit with --export=ALL). +MOUNTS="$REPO:$REPO:rw,$MODEL:$MODEL:ro" +# Optional extra container mounts (srun --container-mounts syntax). Needed +# e.g. when $REPO is a git worktree whose build artifacts (.venv-3.12, +# bindings *.so) are symlinks into the main checkout. +if [ -n "${EXTRA_MOUNTS:-}" ]; then + MOUNTS+=",$EXTRA_MOUNTS" +fi + +srun --mpi=pmix \ + --container-image="$CONTAINER_IMAGE" \ + --container-mount-home \ + --container-mounts="$MOUNTS" \ + bash -c " + set -x + ulimit -n 65536 + + # Node-local Triton cache: ~/.triton on shared NFS races across + # ranks during autotune JIT (stale file handles). + export TRITON_CACHE_DIR=/tmp/triton-cache-rank\${SLURM_PROCID:-0} + mkdir -p \"\$TRITON_CACHE_DIR\" + + # Node-local flashinfer cache: the default (\$HOME) is shared NFS and + # races across ranks during cubin download (stale file handles). + # These caches are cold on every run (~minutes of JIT per rank); + # the GSM8K dataset cache is \$HOME/.cache/huggingface, provided by + # --container-mount-home above. + export FLASHINFER_WORKSPACE_BASE=/tmp/flashinfer-rank\${SLURM_PROCID:-0} + + # Node-local HF remote-code modules cache: all ranks share \$HOME + # (mounted at /root), and racing mkdir of + # ~/.cache/huggingface/modules at startup escapes pathlib's + # exist_ok on NFS (stale attribute cache fails the is_dir() + # recheck after EEXIST) -> 'FileExistsError: /root/.cache' spam + # from every rank. Only the trust_remote_code modules move + # node-local; the dataset/hub cache stays shared for reuse. + export HF_MODULES_CACHE=/tmp/hf-modules-rank\${SLURM_PROCID:-0} + + # Run trtllm-eval from the in-place installation created in the + # README Prerequisites section (build_wheel.py creates .venv-3.12). + export PATH=\"$REPO/.venv-3.12/bin:\$PATH\" + + # Log a running partial score every N completed responses (0 = off). + export TLLM_EVAL_PARTIAL_SCORES_EVERY=\"\${TLLM_EVAL_PARTIAL_SCORES_EVERY:-100}\" + # Cap in-flight requests: may cost some throughput vs submit-all, but yields steady partial scores (early failure signal) instead of one burst at the end (0 = off). + export TLLM_EVAL_MAX_IN_FLIGHT=\"\${TLLM_EVAL_MAX_IN_FLIGHT:-0}\" + + # Import tensorrt_llm from $REPO, not from wherever the venv's + # in-place install points (for git-worktree submits those differ; + # without this the job silently tests the main checkout's code). + export PYTHONPATH=\"$REPO\${PYTHONPATH:+:\$PYTHONPATH}\" + + exec '$REPO/tensorrt_llm/llmapi/trtllm-llmapi-launch' python3 \ + '$REPO/.venv-3.12/bin/trtllm-eval' \ + --model \"$MODEL\" \ + --backend pytorch \ + --config '$EVAL_CONFIG' \ + gsm8k + " diff --git a/examples/kimi_k3/run_serving_benchmark_kimi_k3.sbatch b/examples/kimi_k3/run_serving_benchmark_kimi_k3.sbatch new file mode 100644 index 000000000000..33f55312798e --- /dev/null +++ b/examples/kimi_k3/run_serving_benchmark_kimi_k3.sbatch @@ -0,0 +1,233 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Kimi K3 serving benchmark on 16 NVIDIA Blackwell GPUs. +# +# Brings up trtllm-serve with attention DP + 16-way expert parallelism, then +# drives the in-repo benchmark client (tensorrt_llm/serve/scripts/ +# benchmark_serving.py) with synthetic random 8K-input/1K-output requests. +# Per concurrency C: one warmup pass (C requests, discarded), then one +# measured pass (10*C requests) saved as JSON. +# +# Submit from the repository root of a built TensorRT-LLM checkout (see the +# README Prerequisites section): +# +# sbatch examples/kimi_k3/run_serving_benchmark_kimi_k3.sbatch \ +# --model /path/to/kimi-k3-checkpoint \ +# --image /path/to/tensorrt-llm-container.sqsh +# +# Optional arguments: --isl N (default 8192), --osl N (default 1024), +# --concurrencies "64 256". REPO defaults to the submit directory (override +# by exporting REPO) and must contain the built TensorRT-LLM checkout. +# Results and server/client logs go to kimi-k3-serving-benchmark-/ +# in the submission directory. The #SBATCH partition and account below are +# placeholders — edit them or override on the sbatch command line +# (sbatch --partition=... --account=...). +# +# Note: server initialization and weight loading take 10-15 minutes before +# timed benchmark traffic starts; if your cluster reaps idle-GPU jobs, grant +# this job an exemption via your site's mechanism. +# +#SBATCH --job-name=kimi-k3-serving-benchmark +#SBATCH --partition=batch +#SBATCH --account=${account} +#SBATCH --nodes=4 +#SBATCH --ntasks-per-node=4 +#SBATCH --gpus-per-node=4 +#SBATCH --time=03:00:00 +#SBATCH --output=kimi-k3-serving-benchmark-%j.log + +set -euo pipefail + +usage() { + echo "Usage: sbatch $0 --model PATH --image PATH [--isl N] [--osl N] [--concurrencies \"64 256\"]" +} + +MODEL="" +CONTAINER_IMAGE="" +ISL=8192 +OSL=1024 +CONCURRENCIES="64 256" +PORT=8000 +while [[ $# -gt 0 ]]; do + case "$1" in + --model|--image|--isl|--osl|--concurrencies) + [[ $# -ge 2 ]] || { echo "error: $1 requires a value" >&2; usage >&2; exit 2; } + VALUE=$2 + OPT=$1 + shift 2 + ;; + --model=*|--image=*|--isl=*|--osl=*|--concurrencies=*) + VALUE=${1#*=} + OPT=${1%%=*} + shift + ;; + -h|--help) + usage + exit 0 + ;; + *) + echo "error: unknown argument: $1" >&2 + usage >&2 + exit 2 + ;; + esac + case "$OPT" in + --model) MODEL=$VALUE ;; + --image) CONTAINER_IMAGE=$VALUE ;; + --isl) ISL=$VALUE ;; + --osl) OSL=$VALUE ;; + --concurrencies) CONCURRENCIES=$VALUE ;; + esac +done + +[[ -n "$MODEL" ]] || { echo "error: --model is required" >&2; usage >&2; exit 2; } +[[ -n "$CONTAINER_IMAGE" ]] || { echo "error: --image is required" >&2; usage >&2; exit 2; } +[[ -e "$MODEL" ]] || { echo "error: model path does not exist: $MODEL" >&2; exit 2; } +[[ -e "$CONTAINER_IMAGE" ]] || { echo "error: image path does not exist: $CONTAINER_IMAGE" >&2; exit 2; } +[[ "$ISL" =~ ^[0-9]+$ && "$OSL" =~ ^[0-9]+$ ]] || { echo "error: --isl/--osl must be integers" >&2; exit 2; } +for concurrency in $CONCURRENCIES; do + [[ "$concurrency" =~ ^[0-9]+$ ]] || { echo "error: bad concurrency: $concurrency" >&2; exit 2; } +done + +# 9344 at the default 8192/1024, matching the verified runs exactly. +MAX_SEQ_LEN=$((ISL + OSL + 128)) + +REPO=${REPO:-$SLURM_SUBMIT_DIR} +WORK_DIR=$SLURM_SUBMIT_DIR/kimi-k3-serving-benchmark-$SLURM_JOB_ID +mkdir -p "$WORK_DIR" + +# Serving config; free_gpu_memory_fraction 0.25 leaves headroom for the +# KDA/SSM state cache on top of the model weights. +CONFIG_FILE=$WORK_DIR/serve_config.yaml +cat > "$CONFIG_FILE" <<'EOF' +enable_attention_dp: true +moe_expert_parallel_size: 16 +disable_overlap_scheduler: false +cuda_graph_config: + max_batch_size: 32 +enable_chunked_prefill: true +kv_cache_config: + enable_block_reuse: false + free_gpu_memory_fraction: 0.25 +EOF + +# Shared warm caches, persistent across runs: without a pre-populated +# FLASHINFER_CUBIN_DIR, ranks on egress-less compute nodes can hang fetching +# cubins during the TRTLLM-Gen FMHA warmup. +CACHE_DIR=${CACHE_DIR:-$SLURM_SUBMIT_DIR/.kimi_k3_cache} +mkdir -p "$CACHE_DIR" + +MOUNTS="$REPO:$REPO:rw,$MODEL:$MODEL:ro,$WORK_DIR:$WORK_DIR:rw,$CACHE_DIR:$CACHE_DIR:rw" +export PYTHONPATH="$REPO${PYTHONPATH:+:$PYTHONPATH}" +HEAD_NODE=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | awk 'NR == 1') + +SERVER_PID="" +cleanup() { [[ -n "$SERVER_PID" ]] && kill "$SERVER_PID" 2>/dev/null || true; } +trap cleanup EXIT + +# Server: all 16 ranks; rank 0 (this node) binds the HTTP endpoint. +# --kill-on-bad-exit=1: any failed rank kills the whole step immediately. +srun --mpi=pmix --kill-on-bad-exit=1 \ + --container-image="$CONTAINER_IMAGE" \ + --container-mount-home \ + --container-mounts="$MOUNTS" \ + --output="$WORK_DIR/serve.log" \ + bash -c " + set -x + ulimit -n 65536 + + # Node-local Triton cache: ~/.triton on shared NFS races across + # ranks during autotune JIT (stale file handles). + export TRITON_CACHE_DIR=/tmp/triton-cache-rank\${SLURM_PROCID:-0} + mkdir -p \"\$TRITON_CACHE_DIR\" + # Shared warm caches (see the CACHE_DIR comment above). + export FLASHINFER_WORKSPACE_BASE=$CACHE_DIR/fi_workspace + export FLASHINFER_CUBIN_DIR=$CACHE_DIR/fi_cubins + export HF_HOME=\${HF_HOME:-$CACHE_DIR/hf_home} + mkdir -p \"\$FLASHINFER_WORKSPACE_BASE\" \"\$FLASHINFER_CUBIN_DIR\" \"\$HF_HOME\" + + # Exercise CommunicationFactory AUTO selection rather than inheriting + # a caller-side forced communication strategy. + unset TRTLLM_FORCE_COMM_METHOD + + # In-place venv of the built checkout (build_wheel.py creates .venv-3.12). + export PATH=\"$REPO/.venv-3.12/bin:\$PATH\" + + # The leader receives --log_level through trtllm-serve. Spawned MPI + # workers read TLLM_LOG_LEVEL instead, so export it for all ranks. + export TLLM_LOG_LEVEL=\${TLLM_LOG_LEVEL:-info} + + exec '$REPO/tensorrt_llm/llmapi/trtllm-llmapi-launch' trtllm-serve \ + \"$MODEL\" \ + --backend pytorch \ + --host 0.0.0.0 \ + --port $PORT \ + --tp_size 16 \ + --max_batch_size 32 \ + --max_num_tokens 8192 \ + --max_seq_len $MAX_SEQ_LEN \ + --trust_remote_code \ + --config '$CONFIG_FILE' + " & +SERVER_PID=$! + +# Wait up to 25 min for health; on timeout or server death, cancel the job +# so the allocation never sits on idle GPUs. +deadline=$((SECONDS + 1500)) +until curl -sf -o /dev/null "http://127.0.0.1:$PORT/health"; do + if ! kill -0 "$SERVER_PID" 2>/dev/null || (( SECONDS >= deadline )); then + echo "error: server not healthy within 25 min (see $WORK_DIR/serve.log); cancelling job" >&2 + scancel "$SLURM_JOB_ID" + sleep 120 + exit 1 + fi + sleep 15 +done +echo "Server is healthy; starting benchmark sweep (see $WORK_DIR/client.log)." + +# Client on the head node, alongside the server. The random dataset is +# generated client-side (--random-ids --tokenize-on-client) so the server +# sees exactly ISL input tokens; --ignore-eos fixes output length at OSL. +srun --overlap --kill-on-bad-exit=1 --nodes=1 --ntasks=1 -w "$HEAD_NODE" \ + --container-image="$CONTAINER_IMAGE" \ + --container-mount-home \ + --container-mounts="$MOUNTS" \ + --output="$WORK_DIR/client.log" \ + bash -c " + set -x + export PATH=\"$REPO/.venv-3.12/bin:\$PATH\" + + run_client() { + python3 -m tensorrt_llm.serve.scripts.benchmark_serving \ + --model \"$MODEL\" \ + --host 127.0.0.1 \ + --port $PORT \ + --dataset-name random \ + --random-ids \ + --tokenize-on-client \ + --random-input-len $ISL \ + --random-output-len $OSL \ + --ignore-eos \ + --trust-remote-code \ + \"\$@\" + } + + for concurrency in $CONCURRENCIES; do + echo \"=== Warmup: concurrency \$concurrency ===\" + run_client --num-prompts \$concurrency --max-concurrency \$concurrency + + echo \"=== Measured: concurrency \$concurrency ===\" + run_client \ + --num-prompts \$((concurrency * 10)) \ + --max-concurrency \$concurrency \ + --save-result \ + --result-dir '$WORK_DIR' \ + --result-filename \"concurrency_\${concurrency}.json\" + done + " + +echo "Benchmark complete; results in $WORK_DIR" +kill "$SERVER_PID" 2>/dev/null || true +wait "$SERVER_PID" 2>/dev/null || true diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py b/tests/unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py index f4d2144bc151..58448da45733 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py @@ -75,6 +75,9 @@ def test_on_update_kv_lens_rebuilds_stale_map(): md = object.__new__(DSAtrtllmAttentionMetadata) md.kv_cache_manager = None md._num_generations = 0 + # __init__ (bypassed by object.__new__) defaults this to False; + # on_update_kv_lens() reads it since #16925. + md.in_mtp_draft_loop = False # Stub collaborators unrelated to the map rebuild (test_dsa_indexer.py style). md.kv_lens_cuda = torch.tensor([100, 200, 300], dtype=torch.int32, device=device) md._compute_kv_lens_row_reorder = Mock()