feat(plugin): DeepSeek-V4.1-Flash on the vLLM plugin backend

.github/benchmark/oot_models_accuracy.json

@@ -465,5 +465,35 @@
     "accuracy_baseline": 0.9447,
     "accuracy_baseline_model": "zai-org/GLM-5.2-FP8",
     "_baseline_note": "20-shot GSM8K is lossless for MTP; threshold follows GLM-5.2-FP8 until MXFP4 MTP-specific CI baseline is calibrated."
+  },
+  {
+    "model_name": "DeepSeek-V4.1-Flash TP4",
+    "model_path": "deepseek-ai/DeepSeek-V4.1-Flash",
+    "extra_args": "--tensor-parallel-size 4 --tokenizer-mode deepseek_v4 --kv-cache-dtype auto --gpu-memory-utilization 0.9 --max-model-len 8192 --max-num-batched-tokens 8192 --max-num-seqs 64",
+    "env_vars": "VLLM_USE_BREAKABLE_CUDAGRAPH=1\nATOM_V41_EXPERIMENTAL_CUDAGRAPH=1",
+    "runner": "atom-plugin-acc-validation-runner-vllm",
+    "test_level": "nightly",
+    "priority": "P0",
+    "lm_eval_num_fewshot": 20,
+    "accuracy_test_threshold": 0.91,
+    "accuracy_baseline": 0.9303,
+    "accuracy_baseline_model": "deepseek-ai/DeepSeek-V4.1-Flash",
+    "_baseline_note": "DeepSeek-V4.1 on the vLLM plugin. Three of these arguments are forced, not chosen: the runtime takes --kv-cache-dtype auto/bfloat16/nvfp4 and refuses the fp8 this harness supplies by default, prefix caching is refused because a block-table hit restores the compressed pages but not the per-request CSA2 state read alongside them (the harness already passes --no-enable-prefix-caching), and the two env vars are what keep cudagraph_mode at the FULL_AND_PIECEWISE the harness asks for instead of falling back to eager. max-model-len is 8192 because the checkpoint's own 1M would price the KV pool out of the box this was measured on. Baseline 0.9303 flexible-extract (strict identical) is one full 1319-question 20-shot run of this exact recipe, driven by .github/scripts/atom_oot_test.sh itself rather than a stand-in. One run, so read it as a point and not as a noise estimate: the 0.91 bar leaves 2pp, and a single red should be rerun before it is believed. This is the non-speculative cell; the DSpark one below shares every other argument, so the pair isolates speculation. Coverage boundary, since max-model-len reads like one and is not: 8192 is sized for the KV pool (the checkpoint's own max_position_embeddings is 1048576), not for what the indexer sees. Measured, the 20-shot prompts are 3411 tokens at p50 and 3736 at most, so the token-level selection is exercised -- index_topk is 512, picked out of ~3.4k, past a sliding_window of 128. The block-level candidate stage is not: candidate_topk_blocks (2048) x candidate_block_size (8) is 16384 tokens, which 8192 cannot reach, so that stage never discards anything here. Raising max-model-len would NOT close that -- gsm8k's prompts are 3.4k whatever the limit is. Covering it needs a long-context task, not a bigger flag."
+  },
+  {
+    "model_name": "DeepSeek-V4.1-Flash DSpark TP4",
+    "model_path": "deepseek-ai/DeepSeek-V4.1-Flash",
+    "extra_args": "--tensor-parallel-size 4 --tokenizer-mode deepseek_v4 --kv-cache-dtype auto --gpu-memory-utilization 0.9 --max-model-len 8192 --max-num-batched-tokens 8192 --max-num-seqs 64 --speculative-config '{\"method\":\"dspark\",\"num_speculative_tokens\":5,\"kv_cache_dtype\":\"fp8_ds_mla\"}'",
+    "env_vars": "VLLM_USE_BREAKABLE_CUDAGRAPH=1\nATOM_V41_EXPERIMENTAL_CUDAGRAPH=1",
+    "runner": "atom-plugin-acc-validation-runner-vllm",
+    "test_level": "nightly",
+    "priority": "P0",
+    "lm_eval_num_fewshot": 20,
+    "accuracy_test_threshold": 0.91,
+    "accuracy_baseline": 0.9253,
+    "accuracy_baseline_model": "deepseek-ai/DeepSeek-V4.1-Flash",
+    "mtp_accept_threshold": 0.35,
+    "mtp_per_pos_threshold": "0.55,0.40,0.25,0.18,0.10",
+    "_baseline_note": "DeepSeek-V4.1 on the vLLM plugin. Three of these arguments are forced, not chosen: the runtime takes --kv-cache-dtype auto/bfloat16/nvfp4 and refuses the fp8 this harness supplies by default, prefix caching is refused because a block-table hit restores the compressed pages but not the per-request CSA2 state read alongside them (the harness already passes --no-enable-prefix-caching), and the two env vars are what keep cudagraph_mode at the FULL_AND_PIECEWISE the harness asks for instead of falling back to eager. max-model-len is 8192 because the checkpoint's own 1M would price the KV pool out of the box this was measured on. DSpark is fixed at 5 draft tokens by the model's config gate, and the draft is the target checkpoint's own mtp.* stages, so there is no separate draft model to point at. The draft needs its own kv_cache_dtype in the speculative config: it writes fp8_ds_mla onto the CacheConfig it would otherwise share with the target, which the target's gate then rejects. Baseline 0.9253 is the mean of two full 1319-question 20-shot runs of this recipe on fresh servers, 0.9249 and 0.9257 flexible-extract (strict within 0.0008 of flexible both times). Two runs, 0.08pp apart -- read that as a lower bound on the noise, not the noise. The non-speculative cell above measured 0.9303 the same way; the 0.5pp between them is well inside the 0.0072 standard error either run reports. The acceptance floors are the gate that matters here: verification is lossless w.r.t. the target, so a draft head that reads the wrong hidden states, or an output projection no checkpoint ever filled, leaves every answer correct and shows up only in these numbers. Measured 0.4729 and 0.4705 overall across the two runs, per position 0.7785/0.6243/0.4299/0.3214/0.2102 and 0.7898/0.6411/0.4248/0.3063/0.1906, against 0.008 overall when the target's LM head was not shared into the draft. The floors sit roughly 70% of measured, clear of both run-to-run movement and anything a broken draft can reach. Coverage boundary, since max-model-len reads like one and is not: 8192 is sized for the KV pool (the checkpoint's own max_position_embeddings is 1048576), not for what the indexer sees. Measured, the 20-shot prompts are 3411 tokens at p50 and 3736 at most, so the token-level selection is exercised -- index_topk is 512, picked out of ~3.4k, past a sliding_window of 128. The block-level candidate stage is not: candidate_topk_blocks (2048) x candidate_block_size (8) is 16384 tokens, which 8192 cannot reach, so that stage never discards anything here. Raising max-model-len would NOT close that -- gsm8k's prompts are 3.4k whatever the limit is. Covering it needs a long-context task, not a bigger flag."
   }
 ]

.github/scripts/check_mtp_acceptance.sh

@@ -0,0 +1,102 @@
+#!/usr/bin/env bash
+# The speculative-decoding acceptance gate, for whoever runs a spec cell.
+#
+# gsm8k accuracy CANNOT guard a speculative cell. Verification is lossless with
+# respect to the target model, so a draft head that has been broken outright --
+# reading the wrong hidden states, or an output projection no checkpoint ever
+# filled -- leaves every answer correct and shows up only as acceptance falling
+# to the floor. This is the only gate that catches that.
+#
+# It lives here because it was written twice, and the copies had already
+# diverged in the way that matters: the PR-side one gated the overall rate and
+# not the per-position rates, so PR CI could not see a per-position collapse --
+# a draft that still accepts its first token and nothing after it keeps a
+# respectable overall number.
+#
+# `atom_oot_test.sh` records both into the result JSON, scraped from the live
+# /metrics endpoint during the gsm8k run.
+#
+# Usage:
+#   check_mtp_acceptance.sh RESULT_JSON OVERALL_THRESHOLD [PER_POS_THRESHOLDS]
+#
+# PER_POS_THRESHOLDS is a comma-separated list, one per draft position; absent
+# or empty means the per-position rates are not gated.
+#
+# Reads the JSON with jq when it is there and falls back to python3: the
+# nightly host runner has no python, the PR path runs inside a container that
+# does, and the gate has to be the same number either way.
+set -uo pipefail
+
+result_file=${1:-}
+overall_threshold=${2:-}
+per_pos_thresholds=${3:-}
+
+if [[ -z "${result_file}" || -z "${overall_threshold}" ]]; then
+    echo "usage: $(basename "$0") RESULT_JSON OVERALL_THRESHOLD [PER_POS_THRESHOLDS]" >&2
+    exit 2
+fi
+if [[ ! -f "${result_file}" ]]; then
+    echo "ERROR: no results JSON for the MTP acceptance check: ${result_file}" >&2
+    exit 2
+fi
+echo "RESULT_FILE: ${result_file}"
+
+# read_metric overall | read_metric position <index>
+#
+# MTP_GATE_READER pins the reader (jq|python); the default picks whichever is
+# installed. It exists so the two can be compared against each other -- a gate
+# whose answer depends on which one the runner happens to have is not a gate.
+read_metric() {
+    local what=$1 idx=${2:-0}
+    local reader=${MTP_GATE_READER:-auto}
+    if [[ "${reader}" != "python" ]] && command -v jq >/dev/null 2>&1; then
+        case "${what}" in
+        overall)
+            jq -r '.atom_ci_metadata.mtp_acceptance_overall // empty' "${result_file}"
+            ;;
+        position)
+            jq -r ".atom_ci_metadata.mtp_per_pos_acceptance[${idx}] // empty" "${result_file}"
+            ;;
+        esac
+    else
+        python3 "${BASH_SOURCE%/*}/read_mtp_metric.py" "${result_file}" "${what}" "${idx}"
+    fi
+}
+
+overall=$(read_metric overall)
+if [[ -z "${overall}" ]]; then
+    echo "ERROR: mtp_acceptance_overall missing -- spec-decode /metrics were not captured during the run. Treating as a failure." >&2
+    exit 1
+fi
+
+fail=0
+echo "MTP overall acceptance: ${overall} (threshold ${overall_threshold})"
+if awk -v v="${overall}" -v t="${overall_threshold}" 'BEGIN {exit !(v < t)}'; then
+    echo "FAIL: overall acceptance ${overall} < threshold ${overall_threshold}"
+    fail=1
+fi
+
+if [[ -n "${per_pos_thresholds}" ]]; then
+    IFS=',' read -ra thresholds <<<"${per_pos_thresholds}"
+    idx=0
+    for threshold in "${thresholds[@]}"; do
+        position=$(read_metric position "${idx}")
+        if [[ -z "${position}" ]]; then
+            echo "FAIL: position ${idx} acceptance missing (expected >= ${threshold})"
+            fail=1
+        else
+            echo "MTP position ${idx} acceptance: ${position} (threshold ${threshold})"
+            if awk -v v="${position}" -v t="${threshold}" 'BEGIN {exit !(v < t)}'; then
+                echo "FAIL: position ${idx} acceptance ${position} < threshold ${threshold}"
+                fail=1
+            fi
+        fi
+        idx=$((idx + 1))
+    done
+fi
+
+if [[ "${fail}" -eq 1 ]]; then
+    echo "MTP acceptance gate FAILED." >&2
+    exit 1
+fi
+echo "MTP acceptance gate PASSED."

.github/scripts/plugin_ci/run_vllm_ci.sh

@@ -48,6 +48,10 @@ CLIENT_COMMAND="${MATRIX_CLIENT_COMMAND:-}"
 ENV_VARS="${MATRIX_ENV_VARS:-}"
 LM_EVAL_NUM_FEWSHOT="${MATRIX_LM_EVAL_NUM_FEWSHOT:-3}"
 ACCURACY_TEST_THRESHOLD="${MATRIX_ACCURACY_TEST_THRESHOLD:-0.0}"
+# Empty unless the cell speculates: an acceptance floor only means
+# something where there is a draft.
+MTP_ACCEPT_THRESHOLD="${MATRIX_MTP_ACCEPT_THRESHOLD:-}"
+MTP_PER_POS_THRESHOLD="${MATRIX_MTP_PER_POS_THRESHOLD:-}"
 
 if [[ -z "${MODEL_NAME}" || -z "${MODEL_PATH}" ]]; then
   echo "ERROR: MATRIX_MODEL_NAME and MATRIX_MODEL_PATH are required" >&2
@@ -262,3 +266,12 @@ if [[ "${result}" -eq 1 ]]; then
 fi
 
 echo "Accuracy test passed for ${MODEL_NAME}"
+
+# The acceptance gate, shared with the nightly workflow. It was written twice
+# and the copies diverged: this one checked the overall rate only, so a
+# per-position collapse -- a draft still landing its first token and nothing
+# after it -- passed here while the nightly caught it.
+if [[ -n "${MTP_ACCEPT_THRESHOLD}" ]]; then
+  "${REPO_ROOT}/.github/scripts/check_mtp_acceptance.sh" \
+    "${result_file}" "${MTP_ACCEPT_THRESHOLD}" "${MTP_PER_POS_THRESHOLD}"
+fi

.github/scripts/read_mtp_metric.py

@@ -0,0 +1,38 @@
+# SPDX-License-Identifier: MIT
+"""Read one MTP acceptance number out of an accuracy result JSON.
+
+The jq fallback for `check_mtp_acceptance.sh`. The nightly host runner has jq
+and no python; the PR path runs in a container with python and no jq, and the
+gate has to read the same number either way.
+
+Prints the value, or an empty line when it is absent -- which the caller
+treats as a failure rather than a pass, since a missing acceptance means the
+spec-decode metrics were never captured.
+"""
+
+import json
+import sys
+
+
+def main() -> int:
+    if len(sys.argv) < 3:
+        print("usage: read_mtp_metric.py RESULT_JSON overall|position [INDEX]")
+        return 2
+    path, what = sys.argv[1], sys.argv[2]
+    index = int(sys.argv[3]) if len(sys.argv) > 3 else 0
+
+    with open(path, encoding="utf-8") as handle:
+        meta = json.load(handle).get("atom_ci_metadata", {})
+
+    if what == "overall":
+        value = meta.get("mtp_acceptance_overall")
+    else:
+        positions = meta.get("mtp_per_pos_acceptance") or []
+        value = positions[index] if index < len(positions) else None
+
+    print("" if value is None else value)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

.github/workflows/atom-vllm-accuracy-validation.yaml

@@ -1433,55 +1433,16 @@ jobs:
         run: |
           # MTP acceptance is recorded into the result JSON by atom_oot_test.sh
           # (scraped from the live vLLM /metrics endpoint during the gsm8k run).
-          # gsm8k accuracy alone CANNOT guard MTP: speculative decoding is loss-
-          # less w.r.t. the target model, so a broken draft head leaves accuracy
-          # unchanged and only craters acceptance/throughput. This step is the
-          # only gate that catches an MTP regression. Uses jq+awk (no python on
-          # the host runner), mirroring the gsm8k accuracy check above.
+          # The gate itself is shared with the PR path -- it was written twice
+          # and the copies diverged, the PR one having dropped the
+          # per-position check.
           result_file=$(ls -1t oot_accuracy_results/*.json 2>/dev/null | head -n 1)
-          if [ -z "$result_file" ] || [ ! -f "$result_file" ]; then
+          if [ -z "$result_file" ]; then
             echo "ERROR: No results JSON file found for MTP acceptance check."
             exit 2
           fi
-          echo "RESULT_FILE: $result_file"
-
-          overall=$(jq -r '.atom_ci_metadata.mtp_acceptance_overall // empty' "$result_file")
-          if [ -z "$overall" ]; then
-            echo "ERROR: mtp_acceptance_overall missing — spec-decode /metrics were not captured during the run. Treating as a failure."
-            exit 1
-          fi
-
-          fail=0
-          echo "MTP overall acceptance: $overall (threshold ${MTP_ACCEPT_THRESHOLD})"
-          if awk -v v="$overall" -v t="${MTP_ACCEPT_THRESHOLD}" 'BEGIN {exit !(v < t)}'; then
-            echo "FAIL: overall acceptance $overall < threshold ${MTP_ACCEPT_THRESHOLD}"
-            fail=1
-          fi
-
-          if [ -n "${MTP_PER_POS_THRESHOLD}" ]; then
-            IFS=',' read -ra pos_thresholds <<< "${MTP_PER_POS_THRESHOLD}"
-            idx=0
-            for th in "${pos_thresholds[@]}"; do
-              p=$(jq -r ".atom_ci_metadata.mtp_per_pos_acceptance[$idx] // empty" "$result_file")
-              if [ -z "$p" ]; then
-                echo "FAIL: position $idx acceptance missing (expected >= $th)"
-                fail=1
-              else
-                echo "MTP position $idx acceptance: $p (threshold $th)"
-                if awk -v v="$p" -v t="$th" 'BEGIN {exit !(v < t)}'; then
-                  echo "FAIL: position $idx acceptance $p < threshold $th"
-                  fail=1
-                fi
-              fi
-              idx=$((idx + 1))
-            done
-          fi
-
-          if [ "$fail" -eq 1 ]; then
-            echo "MTP acceptance gate FAILED."
-            exit 1
-          fi
-          echo "MTP acceptance gate PASSED."
+          .github/scripts/check_mtp_acceptance.sh \
+            "$result_file" "${MTP_ACCEPT_THRESHOLD}" "${MTP_PER_POS_THRESHOLD}"
 
       - name: Collect summary
         if: ${{ steps.validation-window.outputs.should_run == 'true' && success() }}

.github/workflows/atom-vllm-test.yaml

@@ -163,6 +163,36 @@ jobs:
             lm_eval_num_fewshot: 5
             accuracy_test_threshold: 0.94
             runner: spur-runner-mi355x-8gpu
+          - display_name: "DeepSeek-V4.1-Flash DSpark TP4"
+            model_name: "DeepSeek-V4.1-Flash-DSpark"
+            model_path: "deepseek-ai/DeepSeek-V4.1-Flash"
+            # `--kv-cache-dtype auto`: the V4.1 runtime takes auto/bfloat16/
+            # nvfp4 and refuses fp8, which is what the harness would otherwise
+            # supply. `--no-enable-prefix-caching` comes from the harness and
+            # is also required here: a block-table hit restores the compressed
+            # pages but not the per-request CSA2 state read alongside them.
+            # DSpark is fixed at 5 draft tokens by the model's own config gate,
+            # and the draft is the target checkpoint's `mtp.*` stages, so there
+            # is no separate draft_model_path.
+            extra_args: "--tensor-parallel-size 4 --tokenizer-mode deepseek_v4 --kv-cache-dtype auto --gpu-memory-utilization 0.9 --max-model-len 8192 --max-num-batched-tokens 8192 --max-num-seqs 64 --speculative-config '{\"method\":\"dspark\",\"num_speculative_tokens\":5,\"kv_cache_dtype\":\"fp8_ds_mla\"}'"
+            env_vars: |
+              VLLM_USE_BREAKABLE_CUDAGRAPH=1
+              ATOM_V41_EXPERIMENTAL_CUDAGRAPH=1
+            lm_eval_num_fewshot: 20
+            accuracy_test_threshold: 0.91
+            # The gate that actually guards the draft. Verification is lossless,
+            # so accuracy stays put when the draft head breaks; measured 0.4729
+            # overall on this recipe, against 0.008 when the target's LM head
+            # was not shared into the draft and 0.02 when the aux hidden states
+            # were the wrong quantity. 0.35 sits well clear of both.
+            mtp_accept_threshold: 0.35
+            # Per position, which the overall rate alone cannot see: a draft
+            # that still lands its first token and nothing after it keeps a
+            # respectable overall number. Measured 0.7785/0.6243/0.4299/
+            # 0.3214/0.2102 and 0.7898/0.6411/0.4248/0.3063/0.1906 over two
+            # runs; these floors sit near 70% of that.
+            mtp_per_pos_threshold: "0.55,0.40,0.25,0.18,0.10"
+            runner: spur-runner-mi355x-8gpu
           - display_name: "Qwen3.5-397B-A17B-MXFP4 TP4"
             model_name: "Qwen3.5-397B-A17B-MXFP4"
             model_path: "amd/Qwen3.5-397B-A17B-MXFP4"
@@ -211,6 +241,8 @@ jobs:
           MATRIX_ENV_VARS: ${{ matrix.env_vars }}
           MATRIX_LM_EVAL_NUM_FEWSHOT: ${{ matrix.lm_eval_num_fewshot }}
           MATRIX_ACCURACY_TEST_THRESHOLD: ${{ matrix.accuracy_test_threshold }}
+          MATRIX_MTP_ACCEPT_THRESHOLD: ${{ matrix.mtp_accept_threshold }}
+          MATRIX_MTP_PER_POS_THRESHOLD: ${{ matrix.mtp_per_pos_threshold }}
           PR_BASE_SHA: ${{ github.event.pull_request.base.sha }}
           PR_HEAD_SHA: ${{ github.event.pull_request.head.sha }}
           SKIP_DOCKER_LOGIN: ${{ github.event.pull_request.head.repo.fork && '1' || '0' }}

atom/model_ops/attentions/deepseek_v41/backend.py

@@ -217,6 +217,10 @@ def __init__(self, model_runner):
             self.block_table_cols,
             self.device,
             pages=self._page_bound(model_runner.config.gpu_memory_utilization),
+            # A candidate-consuming layer's plane is as wide as the candidate
+            # list, which the checkpoint fixes and `max_model_len` does not
+            # bound. Short contexts make it the widest plane in the model.
+            candidate_blocks=int(getattr(self.config, "candidate_topk_blocks", 0) or 0),
         )
         logger.info(
             "V4.1 score workspace: %s logits, %.2f GiB",