[Spec Decode][Spark2.5] Support DSpark

.github/workflows/check.yml

@@ -37,7 +37,7 @@ jobs:
     needs: lint
     strategy:
       matrix:
-        python-version: ["3.10", "3.11", "3.12"]
+        python-version: ["3.11", "3.12"]
     steps:
       - uses: actions/checkout@v7
       - uses: actions/setup-python@v7

CHANGELOG.md

@@ -7,12 +7,29 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- **DSpark draft model**: `spark2_5_dspark.py` provides `Spark2_5DSparkForCausalLM`, shared target embedding/LM head, auxiliary hidden-state projection, `vanilla`/`rnn` Markov heads, an optional confidence head, and draft-weight loading from `draft_model.*`, `mtp.*`, or `model.mtp.*` checkpoints. `__init__.py` registers the draft and restores Spark speculative configuration after the DeepSeek fallback, normalizing `markov_rank` and `target_layer_ids` to their DSpark field names.
+- **DSpark attention**: `spark_triton_attention.py` implements Triton JIT fixed-history sliding draft attention over the context window and the full non-causal draft block. Draft layers default to sliding attention, with per-layer sliding/full attention selection through `dspark_layer_types`. It requires dense FP16/BF16 KV, power-of-two head dimensions, pipeline-parallel size 1, and decode-context-parallel size 1.
+- **DSpark numerical test**: `tests/test_spark_triton_attention.py` compares ragged-request draft attention with a PyTorch reference on CUDA; it is skipped when CUDA is unavailable.
+
+### Changed
+
+- **vLLM dependency baseline**: Model and DSpark integration now target release tag [`v0.32.0`](https://github.com/vllm-project/vllm/tree/v0.32.0). `TESTED_VLLM` and registration logs now identify this tag as the tested baseline. Explicitly install `vllm==0.32.0`; package metadata still declares an unpinned `vllm` dependency.
+- **Target model integration**: Expose auxiliary hidden states and shared embeddings for the vLLM V2 GPU DSpark runner, and exclude `draft_model.*` weights when loading the target model.
+- **Tool parser compatibility**: Update chat utility and streaming response protocol imports for vLLM `0.32.0`.
+- **Documentation**: Expand the README with DSpark checkpoint requirements, serving examples, sampling recommendations, parallelism restrictions, and CUDA test instructions; update third-party license documentation for the runtime baseline and DSpark/Triton additions.
+
+### Fixed
+
+- **Transformers configuration validation**: Assign Spark's per-attention-type RoPE parameters after base configuration validation so the generic flat-parameter validator does not reject them.
+
 ## [0.1.0] - 2026-08-28
 
 ### Added
 
 - **Spark2_5 vLLM Plugin**: Out-of-tree plugin that restores Spark2_5 model support to vLLM, registering `Spark2_5ForCausalLM`, the `spark2_5` Transformers config, and the Spark2_5 XML/KV tool-call parser
-- **Vendored Model Implementation**: Spark2_5 model files vendored from vLLM commit [`81efe78`](https://github.com/vllm-project/vllm/tree/81efe7883f30582696b69f9b9ea93c4819a8c608), including hybrid sliding-window/full attention, head-wise attention gating, grouped-query attention, tensor and pipeline parallelism
+- **Vendored Model Implementation**: Spark2_5 model files adapted from vLLM, including hybrid sliding-window/full attention, head-wise attention gating, grouped-query attention, tensor and pipeline parallelism. Historical extraction provenance is retained in `NOTICE` and source headers; the current runtime baseline is tag `v0.32.0`.
 - **Tool-Call Parser**: Spark2_5 XML/KV parser registered under the vLLM parser name `spark25`, with typed argument extraction (integer, number, boolean, JSON)
 - **Project Governance**: LICENSE (Apache-2.0), NOTICE, THIRD_PARTY_LICENSES.md, CODE-OF-CONDUCT.md, SECURITY.md
 - **GitHub Templates**: PR template (`.github/PULL_REQUEST_TEMPLATE.md`), Dependabot config for pip and github-actions

NOTICE

@@ -3,9 +3,11 @@ Copyright 2026 XHToken Authors
 
 This product includes software vendored from the vLLM project
 (https://github.com/vllm-project/vllm), licensed under the Apache License 2.0.
-The vendored files were taken from vLLM commit
-81efe7883f30582696b69f9b9ea93c4819a8c608 and retain their original
-SPDX copyright headers ("Copyright contributors to the vLLM project").
+The original extraction was based on vLLM commit
+81efe7883f30582696b69f9b9ea93c4819a8c608. The current runtime baseline is
+vLLM tag v0.32.0, pinned in package metadata and requiring Python 3.11 or newer.
+Vendored files retain their original SPDX copyright headers
+("Copyright contributors to the vLLM project").
 
 This product depends on the following third-party packages:
 

README.md

@@ -8,13 +8,15 @@ It enables loading and serving Spark2_5 models without modifying the vLLM source
 tree. The plugin registers:
 
 - `Spark2_5ForCausalLM`, the Spark2_5 model implementation;
+- `Spark2_5DSparkForCausalLM`, the Spark DSpark speculative draft model;
 - the `spark2_5` Transformers configuration (`model_type: "spark2_5"`); and
 - the Spark2_5 XML/KV tool-call parser, registered under the vLLM parser name
   `spark25`.
 
-The implementation is vendored from vLLM commit
-`81efe7883f30582696b69f9b9ea93c4819a8c608` and uses vLLM's existing runtime
-layers. It does not contain custom CUDA kernels or require a separate build.
+The implementation targets vLLM tag
+[`v0.32.0`](https://github.com/vllm-project/vllm/tree/v0.32.0) and uses vLLM's
+existing runtime layers. DSpark sliding draft layers use the plugin's Triton
+JIT kernel; no separate extension build is required.
 
 ## Spark2_5 model capabilities
 
@@ -56,9 +58,35 @@ The MLP combines the gate and up projections with
 down projection. `AutoWeightsLoader` and `WeightsMapper` map Hugging Face
 checkpoint weights into this packed projection layout.
 
+### Spark DSpark integration
+
+The DSpark implementation lives in
+[`spark2_5_dspark.py`](src/vllm_spark2_5_plugin/spark2_5_dspark.py). It projects
+selected target hidden states into a non-causal draft block, shares the target
+embedding and LM head, and loads `draft_model.*`, `mtp.*`, or `model.mtp.*`
+checkpoint weights. It supports `vanilla` and `rnn` Markov heads and an optional
+confidence head.
+
+The checkpoint must supply a positive `dspark_markov_rank` and nonempty,
+unique, ascending, zero-based `dspark_target_layer_ids`. The registration hook
+in `__init__.py` restores Spark after vLLM's DeepSeek fallback and normalizes
+`markov_rank` and `target_layer_ids` to these DSpark field names. Draft and
+target hidden sizes and full vocabularies must match.
+
+[`spark_triton_attention.py`](src/vllm_spark2_5_plugin/spark_triton_attention.py)
+implements fixed-history sliding draft attention. For target context length
+`L` and draft block length `B`, every query attends to keys from
+`max(0, L - sliding_window)` through `L + B - 1`, including the full non-causal
+block. It uses vLLM's bundled Triton and requires dense FP16/BF16 KV and
+power-of-two head dimensions. Target attention keeps vLLM's selected backend.
+Draft layers default to sliding attention; `dspark_layer_types` can choose
+sliding or full attention per layer. DSpark requires pipeline-parallel and
+decode-context-parallel sizes of 1; the target-only pipeline-parallel support
+above does not extend to DSpark.
+
 ## Linux prerequisites
 
-Use a supported Linux Python environment with vLLM already installed. GPU
+Use a supported Linux Python 3.11 or newer environment with vLLM `0.32.0` installed. GPU
 deployments must have a compatible NVIDIA driver and the CUDA/PyTorch runtime
 required by that vLLM installation. The plugin itself adds no system packages
 or compiled extensions.
@@ -78,6 +106,7 @@ plugin through the `vllm.general_plugins` entry-point metadata.
 
 ```bash
 # From the vLLM checkout containing this directory.
+uv pip install "vllm==0.32.0"
 uv pip install -e ./Spark-plugin --no-deps
 ```
 
@@ -123,6 +152,42 @@ vllm serve /path/to/spark2_5-model \
 The parser name is `spark25`, not `spark2_5`: `spark2_5` identifies the model
 configuration, while `spark25` is the name registered in `ToolParserManager`.
 
+To enable Spark DSpark speculative decoding, use a Spark2_5 checkpoint that
+includes DSpark draft weights (`draft_model.*` or `mtp.*`) and the corresponding
+DSpark configuration:
+
+```bash
+vllm serve /path/to/spark2_5-dspark-model \
+  --trust-remote-code \
+  --enable-auto-tool-choice \
+  --tool-call-parser spark25 \
+  --chat-template /path/to/chat_template.jinja \
+  --served-model-name Spark2_5 \
+  --speculative-config '{"method":"dspark","num_speculative_tokens":7}'
+```
+
+When DSpark is enabled and the request's `temperature` is not `0.0`, we
+recommend setting `draft_sample_method` to `probabilistic` so that the draft
+sampler retains a probabilistic proposal distribution for rejection sampling.
+The draft applies temperature; target top-k/top-p and penalties can still make
+the two distributions differ. Set it in the server's
+`--speculative-config`:
+
+```text
+--speculative-config '{"method":"dspark","num_speculative_tokens":7,"draft_sample_method":"probabilistic"}'
+```
+
+The draft weights are loaded from the same checkpoint as the target model, so
+no separate draft `model` path is needed. The plugin registers
+`Spark2_5DSparkForCausalLM` and restores the Spark draft architecture if vLLM's
+DSpark configuration selects its DeepSeek fallback. The override ensures that
+the plugin's target model and tool parser are used even when vLLM already
+provides Spark2_5 support.
+
+The value `7` is an example; match `num_speculative_tokens` to the checkpoint's
+`dspark_block_size` (or `block_size`) when configured. Spark DSpark currently
+requires pipeline-parallel and decode-context-parallel sizes of 1.
+
 For a small local test deployment, consider `--enforce-eager` and lower values
 for `--max-model-len`, `--max-num-batched-tokens`, and
 `--gpu-memory-utilization` as appropriate for the available device.
@@ -194,13 +259,24 @@ concurrency of 500. Decode latency had a mean TPOT of 11.58 ms and a P99 of
 
 ## Compatibility and maintenance
 
-The model implementation depends on vLLM's internal runtime-layer APIs. The
-`TESTED_VLLM` value in `src/vllm_spark2_5_plugin/__init__.py` records the vLLM
-revision from which the implementation was extracted. When upgrading vLLM,
-retest the plugin and re-vendor the model files if an internal API has moved.
+The model and DSpark implementations depend on vLLM's internal runtime APIs.
+The current dependency baseline is vLLM tag `v0.32.0`. Install that version in
+the serving environment before installing the plugin with `--no-deps`:
+
+```bash
+uv pip install "vllm==0.32.0"
+```
 
-The plugin logs both the tested revision and the running vLLM version when it
-registers, making version drift visible in the server log.
+The package metadata pins `vllm==0.32.0` and requires Python 3.11 or newer,
+matching that runtime's minimum Python version. CI tests Python 3.11 and 3.12
+against this pinned dependency. When upgrading vLLM, update the dependency pin
+and retest model registration, tool parsing, and DSpark decoding against the
+new runtime APIs.
+
+The `TESTED_VLLM` constant in `src/vllm_spark2_5_plugin/__init__.py` records
+`v0.32.0`, the current runtime baseline, and appears in registration logs
+alongside the running vLLM version. Historical extraction provenance is retained
+in `NOTICE` and source headers.
 
 ## Tests
 
@@ -211,9 +287,19 @@ the plugin in the test environment, run:
 .venv/bin/python -m pytest Spark-plugin/tests/test_smoke.py -q
 ```
 
-The tests verify both the public registration entry point and a complete
+The smoke tests verify both the public registration entry point and a complete
 Spark2_5 XML tool-call parse.
 
+[`test_spark_triton_attention.py`](tests/test_spark_triton_attention.py) checks
+the fixed-history draft kernel against a PyTorch reference for requests with
+different context lengths. This test requires CUDA and is skipped when CUDA
+is unavailable; it is not an end-to-end DSpark decoding test.
+
+```bash
+# From the plugin repository root, in the vLLM 0.32.0 environment.
+.venv/bin/python -m pytest tests/test_spark_triton_attention.py -q
+```
+
 ## Plugin loading and allowlists
 
 General plugins are loaded automatically by vLLM unless `VLLM_PLUGINS` is set
@@ -264,20 +350,22 @@ This is expected when the installed vLLM contains the model. Set
 
 The entry-point loader catches plugin import failures and continues starting
 vLLM. Check the complete server log for the original traceback, then compare
-the running vLLM version with the vendored reference commit recorded in
-`src/vllm_spark2_5_plugin/__init__.py`.
+the running vLLM version with the current dependency baseline, `0.32.0`.
 
 ## Project layout
 
 ```text
 Spark-plugin/
 +-- src/vllm_spark2_5_plugin/
-|   +-- __init__.py             # register() and entry-point integration
-|   +-- spark2_5.py               # vendored Spark2_5 model
-|   +-- spark2_5_config.py        # Transformers configuration
-|   +-- spark2_5_tool_parser.py   # Spark2_5 XML/KV parser
+|   +-- __init__.py              # registration and DSpark config integration
+|   +-- spark2_5.py              # Spark2_5 target model
+|   +-- spark2_5_dspark.py       # DSpark draft, heads, and weight loading
+|   +-- spark_triton_attention.py # fixed-history draft attention kernel
+|   +-- spark2_5_config.py       # Transformers configuration
+|   +-- spark2_5_tool_parser.py  # Spark2_5 XML/KV parser
 +-- tests/
 |   +-- test_smoke.py           # registration and parser smoke tests
+|   +-- test_spark_triton_attention.py # CUDA kernel numerical test
 +-- pyproject.toml
 ```
 

THIRD_PARTY_LICENSES.md

@@ -11,7 +11,7 @@ The project itself is licensed under the Apache License 2.0 (see `LICENSE`).
 
 | Package | Version / Reference | License |
 |---|---|---|
-| [vLLM](https://github.com/vllm-project/vllm) | commit `81efe788` (vendored) | Apache-2.0 |
+| [vLLM](https://github.com/vllm-project/vllm) | == 0.32.0 ([runtime baseline](https://github.com/vllm-project/vllm/tree/v0.32.0)) | Apache-2.0 |
 | [PyTorch](https://github.com/pytorch/pytorch) | (transitive via vLLM) | BSD-3-Clause |
 | [Transformers](https://github.com/huggingface/transformers) | (transitive via vLLM) | Apache-2.0 |
 | [OpenAI Python SDK](https://github.com/openai/openai-python) | >= 2.25.0 | Apache-2.0 |
@@ -25,15 +25,31 @@ The project itself is licensed under the Apache License 2.0 (see `LICENSE`).
 ## Vendored Code
 
 The following source files under `src/vllm_spark2_5_plugin/` were vendored
-from the [vLLM](https://github.com/vllm-project/vllm) project at commit
-[`81efe7883f30582696b69f9b9ea93c4819a8c608`](https://github.com/vllm-project/vllm/tree/81efe7883f30582696b69f9b9ea93c4819a8c608)
-and are licensed under the Apache License 2.0:
+from the [vLLM](https://github.com/vllm-project/vllm) project and are licensed
+under the Apache License 2.0. The current runtime dependency baseline is
+release tag [`v0.32.0`](https://github.com/vllm-project/vllm/tree/v0.32.0).
+Historical extraction details remain in `NOTICE` and source headers; they do
+not select the current runtime dependency:
 
 - `spark2_5.py` -- Spark2_5 model implementation
 - `spark2_5_config.py` -- Transformers configuration for Spark2_5
 - `spark2_5_tool_parser.py` -- Spark2_5 XML/KV tool-call parser
 
-Each vendored file carries the following SPDX headers:
+The additional DSpark files are:
+
+- `src/vllm_spark2_5_plugin/spark2_5_dspark.py` -- draft model and weight
+  mapping, importing vLLM DSpark heads; retains Apache-2.0/vLLM SPDX headers.
+- `src/vllm_spark2_5_plugin/spark_triton_attention.py` -- plugin fixed-history
+  draft attention, using Triton JIT through vLLM utilities. Triton is licensed
+  under MIT and supplied transitively through vLLM; no separate extension build
+  is required.
+- `tests/test_spark_triton_attention.py` -- CUDA numerical checks for the draft
+  attention implementation, using PyTorch and vLLM.
+
+Package metadata pins `vllm==0.32.0` and requires Python 3.11 or newer, matching
+the selected runtime. This version change adds no new third-party packages.
+
+Each of the three vendored model/config/parser files above carries these headers:
 
 ```
 SPDX-License-Identifier: Apache-2.0

pyproject.toml

@@ -8,10 +8,9 @@ version = "0.1.0"
 description = "Spark2_5 support for vLLM, maintained out of tree"
 readme = "README.md"
 license = "Apache-2.0"
-requires-python = ">=3.10"
-# The implementation is vendored from an unreleased vLLM commit. Keep version
-# selection with the operator and record the exact tested commit in the package.
-dependencies = ["vllm"]
+requires-python = ">=3.11"
+# Internal vLLM APIs must match the implementation and TESTED_VLLM baseline.
+dependencies = ["vllm==0.32.0"]
 
 [project.optional-dependencies]
 test = ["pytest>=8.0"]