[TRTLLM-14704][feat] Support multi-modal part of K3

.gitignore

@@ -124,3 +124,8 @@ tests/integration/defs/stress_test/artifacts/
 .claude/settings.json
 
 .plans/
+
+# Kimi K3 eval harness per-job runtime configs
+# (run_eval_kimi_k3.sbatch --parallel tep and --dflash rewrites)
+examples/kimi_k3/.eval_tep_runtime.*.yaml
+examples/kimi_k3/.eval_dflash_runtime.*.yaml

docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md

@@ -247,12 +247,12 @@ The response should contain a `choices[0].message.content` field completing the
 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 \
+sbatch examples/kimi_k3/run_eval_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-<job-id>.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:
+The job writes progress and results to `kimi-k3-eval-<job-id>.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 |
 | :-- | --: |

examples/kimi_k3/README.md

@@ -21,12 +21,13 @@ other GPU architectures may be added in a future release.
   [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-<major>.<minor>` path in the
-  commands on this page. Adjust `--cuda_architectures` to the target
-  GPUs (`103-real` for GB300).
+  `build_wheel.py` creates the `.venv-3.12` virtual environment at the
+  repository root (named after the container's Python version). If your
+  container ships a different Python, substitute `.venv-<major>.<minor>`
+  for `.venv-3.12` in every command below and export
+  `TRTLLM_VENV=/path/to/repo/.venv-<major>.<minor>` when submitting the
+  Slurm jobs (they default to the repository-root `.venv-3.12`). 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
@@ -70,7 +71,7 @@ other GPU architectures may be added in a future release.
   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
+  currently pins `flashinfer-python==0.6.14`, so a later
   dependency-resolving TensorRT-LLM install can replace this source revision.
 
 ## Run the model
@@ -93,13 +94,13 @@ should report `True` for all four checks.
 For a full GSM8K evaluation, submit:
 
 ```bash
-sbatch examples/kimi_k3/run_gsm8k_kimi_k3.sbatch \
+sbatch examples/kimi_k3/run_eval_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-<job-id>.log`. If no local dataset path is configured,
+`kimi-k3-eval-<job-id>.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
@@ -116,7 +117,7 @@ possible with different checkpoint or dependency revisions.
 To evaluate with suffix-automaton (SA) speculative decoding, add `--sa`:
 
 ```bash
-sbatch examples/kimi_k3/run_gsm8k_kimi_k3.sbatch \
+sbatch examples/kimi_k3/run_eval_kimi_k3.sbatch \
     --model /path/to/kimi-k3-checkpoint \
     --image /path/to/tensorrt-llm-container.sqsh \
     --sa
@@ -184,7 +185,7 @@ Request a longer allocation if the time is not sufficient for your environment:
 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 \
+sbatch --time=04:00:00 examples/kimi_k3/run_eval_kimi_k3.sbatch \
     --model MODEL --image IMAGE
 ```
 
@@ -194,24 +195,21 @@ 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`):
+KV-cache block reuse is supported 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 \
+sbatch examples/kimi_k3/run_eval_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.
+Block reuse stays off by default because suffix-automaton speculative
+decoding requires the default cache manager, which cannot reuse blocks.
 
 ## Current limitations
 

examples/kimi_k3/perf_sweep/acc_sweep.sbatch

@@ -149,7 +149,7 @@ srun --mpi=pmix --kill-on-bad-exit=1 \
         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).
+        # \$HOME escapes pathlib's exist_ok on NFS (see run_eval_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}\"

examples/kimi_k3/run_dspark_acceptance.sbatch

@@ -0,0 +1,189 @@
+#!/bin/bash
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Kimi K3 DSpark acceptance / speedup A/B on 16 NVIDIA Blackwell GPUs.
+#
+# One submission runs BOTH legs of the weights-day A/B back to back in the
+# same allocation (identical nodes, so the TPOT comparison is clean):
+#   (a) spec-off TPOT reference (recorder off),
+#   (b) DSpark-on measurement (AL, per-position AR, histogram, and — with a
+#       confidence head in the drafter — calibration data).
+#
+#   sbatch examples/kimi_k3/run_dspark_acceptance.sbatch \
+#       --model /path/to/kimi-k3-checkpoint \
+#       --drafter /path/to/dspark-drafter \
+#       --image /path/to/tensorrt-llm-container.sqsh \
+#       --outdir /path/to/results
+#
+# Optional: --num-prompts N (default 64), --max-tokens N (default 256),
+# --confidence-threshold T --confidence-policy P (leave unset for
+# calibration runs; see measure_dspark_acceptance.py), --skip-baseline
+# (rerun only the DSpark leg, e.g. for a threshold sweep).
+#
+# Worktree submits: export REPO=<worktree>, EXTRA_MOUNTS="$MAIN:$MAIN:rw",
+# and submit with --export=ALL (mirrors run_eval_kimi_k3.sbatch).
+#
+#SBATCH --job-name=kimi-k3-dspark-accept
+#SBATCH --partition=batch
+#SBATCH --account=${account}
+#SBATCH --nodes=4
+#SBATCH --ntasks-per-node=4
+#SBATCH --gpus-per-node=4
+#SBATCH --time=04:00:00
+#SBATCH --output=kimi-k3-dspark-accept-%j.log
+
+set -euo pipefail
+
+usage() {
+    echo "Usage: sbatch $0 --model PATH --drafter PATH --image PATH [--outdir DIR]" \
+         "[--num-prompts N] [--max-tokens N] [--confidence-threshold T]" \
+         "[--confidence-policy first_below|cumulative] [--skip-baseline]"
+}
+
+MODEL=""
+DRAFTER=""
+CONTAINER_IMAGE=""
+OUTDIR=""
+NUM_PROMPTS=64
+MAX_TOKENS=256
+CONF_THRESHOLD=""
+CONF_POLICY=first_below
+SKIP_BASELINE=0
+# Abort with a usage error when a value-taking option has no value ($1 is the
+# option name, $2 the remaining positional-argument count of the caller);
+# without this, `set -u` aborts on the unbound $2 instead.
+require_value() {
+    [[ $2 -ge 2 ]] || { echo "error: $1 requires a value" >&2; usage >&2; exit 2; }
+}
+
+while [[ $# -gt 0 ]]; do
+    case "$1" in
+        --model)        require_value --model $#; MODEL=$2; shift 2 ;;
+        --model=*)      MODEL=${1#*=}; shift ;;
+        --drafter)      require_value --drafter $#; DRAFTER=$2; shift 2 ;;
+        --drafter=*)    DRAFTER=${1#*=}; shift ;;
+        --image)        require_value --image $#; CONTAINER_IMAGE=$2; shift 2 ;;
+        --image=*)      CONTAINER_IMAGE=${1#*=}; shift ;;
+        --outdir)       require_value --outdir $#; OUTDIR=$2; shift 2 ;;
+        --outdir=*)     OUTDIR=${1#*=}; shift ;;
+        --num-prompts)  require_value --num-prompts $#; NUM_PROMPTS=$2; shift 2 ;;
+        --max-tokens)   require_value --max-tokens $#; MAX_TOKENS=$2; shift 2 ;;
+        --confidence-threshold) require_value --confidence-threshold $#; CONF_THRESHOLD=$2; shift 2 ;;
+        --confidence-policy)    require_value --confidence-policy $#; CONF_POLICY=$2; shift 2 ;;
+        --skip-baseline) SKIP_BASELINE=1; shift ;;
+        -h|--help)      usage; exit 0 ;;
+        *) echo "error: unknown argument: $1" >&2; usage >&2; exit 2 ;;
+    esac
+done
+
+[[ -n "$MODEL" && -n "$DRAFTER" && -n "$CONTAINER_IMAGE" ]] || { usage >&2; exit 2; }
+[[ -e "$MODEL" ]]   || { echo "error: model path does not exist: $MODEL" >&2; exit 2; }
+[[ -e "$DRAFTER" ]] || { echo "error: drafter path does not exist: $DRAFTER" >&2; exit 2; }
+[[ -e "$CONTAINER_IMAGE" ]] || { echo "error: image path does not exist: $CONTAINER_IMAGE" >&2; exit 2; }
+
+REPO=${REPO:-$SLURM_SUBMIT_DIR}
+# Virtual environment created by build_wheel.py; export TRTLLM_VENV to point
+# at .venv-<major>.<minor> when the container ships another Python version.
+VENV=${TRTLLM_VENV:-$REPO/.venv-3.12}
+OUTDIR=${OUTDIR:-$REPO/dspark-accept-$SLURM_JOB_ID}
+mkdir -p "$OUTDIR"
+
+CONF_ARGS=""
+[[ -n "$CONF_THRESHOLD" ]] && \
+    CONF_ARGS="--confidence-threshold $CONF_THRESHOLD --confidence-policy $CONF_POLICY"
+
+# 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,$DRAFTER:$DRAFTER:ro,$OUTDIR:$OUTDIR:rw"
+if [ -n "${EXTRA_MOUNTS:-}" ]; then
+    MOUNTS+=",$EXTRA_MOUNTS"
+fi
+
+# User cache (GSM8K dataset under ~/.cache/huggingface):
+# --container-mount-home mounts $HOME at /root, but when ~/.cache is a
+# symlink onto a shared filesystem its target dead-ends inside the container
+# (FileNotFoundError /root/.cache/huggingface). Mount the
+# resolved cache root at its host path — the /root/.cache symlink chain
+# then works too — and point HF_HOME at it explicitly. flashinfer/triton
+# JIT caches stay node-local via the /tmp overrides below.
+CACHE_HOST=$(readlink -f "$HOME/.cache" 2>/dev/null || echo "$HOME/.cache")
+if [[ -d "$CACHE_HOST" && "$CACHE_HOST" != "$HOME/.cache" ]]; then
+    MOUNTS+=",$CACHE_HOST:$CACHE_HOST:rw"
+    # Also mount at the symlink's literal target so the in-container
+    # /root/.cache -> <target> chain resolves (the target's parent dirs
+    # do not otherwise exist in the container).
+    CACHE_LINK_TARGET=$(readlink "$HOME/.cache")
+    if [[ -n "$CACHE_LINK_TARGET" && "$CACHE_LINK_TARGET" != "$CACHE_HOST" ]]; then
+        MOUNTS+=",$CACHE_HOST:$CACHE_LINK_TARGET:rw"
+    fi
+fi
+HF_CACHE_HOST="$CACHE_HOST/huggingface"
+
+run_leg() {
+    local leg=$1; shift
+    # %q-quote every forwarded argument: they are spliced into the nested
+    # bash -c string, where a bare $* would re-split values containing
+    # whitespace (e.g. a drafter path with spaces).
+    local escaped_args
+    printf -v escaped_args ' %q' "$@"
+    srun --mpi=pmix \
+        --container-image="$CONTAINER_IMAGE" \
+        --container-mount-home \
+        --container-mounts="$MOUNTS" \
+        bash -c "
+            set -x
+            ulimit -n 65536
+            # Node-local JIT caches (shared-NFS races; see run_eval_kimi_k3.sbatch)
+            export TRITON_CACHE_DIR=/tmp/triton-cache-rank\${SLURM_PROCID:-0}
+            mkdir -p \"\$TRITON_CACHE_DIR\"
+            export FLASHINFER_WORKSPACE_BASE=/tmp/flashinfer-rank\${SLURM_PROCID:-0}
+            export HF_MODULES_CACHE=/tmp/hf-modules-rank\${SLURM_PROCID:-0}
+            export HF_HOME='$HF_CACHE_HOST'
+            # Recorder env must be exported HERE (per rank), not inside the
+            # python script: trtllm-llmapi-launch pre-spawns the MPI worker
+            # ranks, so os.environ changes in the driver never reach the
+            # workers where DFlashWorker lives (symptom: no accept-stats
+            # files). Empty for the spec-off leg (recorder off = clean TPOT).
+            ${LEG_ENV_EXPORT:-true}
+            export PATH=\"$VENV/bin:\$PATH\"
+            # Import tensorrt_llm from \$REPO, not from wherever the venv's
+            # in-place install points (worktree submits differ).
+            export PYTHONPATH=\"$REPO\${PYTHONPATH:+:\$PYTHONPATH}\"
+            exec '$REPO/tensorrt_llm/llmapi/trtllm-llmapi-launch' python3 \
+                '$REPO/examples/kimi_k3/measure_dspark_acceptance.py' \
+                --model '$MODEL' --tp-size 16 \
+                --num-prompts $NUM_PROMPTS --max-tokens $MAX_TOKENS \
+                $escaped_args
+        " 2>&1 | tee "$OUTDIR/$leg.log"
+}
+
+if [[ "$SKIP_BASELINE" -eq 0 ]]; then
+    echo "=== leg (a): spec-off TPOT reference ==="
+    LEG_ENV_EXPORT="" \
+        run_leg spec_off --spec-off --output-json "$OUTDIR/results_spec_off.json"
+fi
+
+echo "=== leg (b): DSpark-on measurement ==="
+LEG_ENV_EXPORT="export TLLM_DFLASH_ACCEPT_STATS_DIR='$OUTDIR/accept-stats'" \
+    run_leg dspark --drafter "$DRAFTER" $CONF_ARGS \
+    --stats-dir "$OUTDIR/accept-stats" \
+    --output-json "$OUTDIR/results_dspark.json"
+
+# Combined speedup summary (host-side python3, stdlib only)
+python3 - "$OUTDIR" <<'EOF'
+import json, os, sys
+outdir = sys.argv[1]
+def load(name):
+    p = os.path.join(outdir, name)
+    return json.load(open(p)) if os.path.exists(p) else None
+ref, spec = load("results_spec_off.json"), load("results_dspark.json")
+summary = {"spec_off": ref, "dspark": spec}
+if ref and spec and ref.get("tpot_proxy_ms") and spec.get("tpot_proxy_ms"):
+    summary["e2e_speedup"] = ref["tpot_proxy_ms"] / spec["tpot_proxy_ms"]
+    print(f"E2E speedup (spec-off TPOT / dspark TPOT): {summary['e2e_speedup']:.3f}x")
+json.dump(summary, open(os.path.join(outdir, "ab_summary.json"), "w"), indent=2)
+print(f"A/B summary: {os.path.join(outdir, 'ab_summary.json')}")
+EOF