Add selective recall training and reached-prefix coverage metrics

.github/workflows/ci.yml

@@ -71,6 +71,30 @@ jobs:
           actual=$(.venv/bin/python -c 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")')
           test "$actual" = "${{ matrix.python-version }}"
 
+  training-reference:
+    name: Optional DFlash2 CPU training and reference parity
+    runs-on: ubuntu-24.04
+    timeout-minutes: 15
+    steps:
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+      - uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
+        with:
+          python-version: "3.12"
+      - name: Install optional CPU training dependencies
+        run: |
+          python -m pip install torch==2.13.0 --index-url https://download.pytorch.org/whl/cpu
+          python -m pip install -r training/requirements.txt
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+        with:
+          repository: z-lab/dflash
+          ref: 07ebd93db9f472af339b644bb70221ad8428328a
+          path: .dflash-reference
+          persist-credentials: false
+      - name: Run small-model training and numerical parity tests
+        env:
+          DFLASH_REFERENCE: ${{ github.workspace }}/.dflash-reference/dflash/model.py
+        run: python -m unittest discover -s dev/tests/training -t . -v
+
   release-hardware:
     name: Release gate (${{ matrix.runner_label }})
     if: ${{ github.event_name == 'workflow_dispatch' && inputs.release_gate == true }}

.github/workflows/draft-training-mps.yml

@@ -0,0 +1,48 @@
+name: Experimental DFlash2 MPS smoke
+
+on:
+  workflow_dispatch:
+    inputs:
+      runner_label:
+        description: Label of a trusted self-hosted Apple Silicon runner with MPS
+        type: string
+        required: true
+        default: splash-mps-training
+
+permissions:
+  contents: read
+
+jobs:
+  mps-smoke:
+    runs-on: [self-hosted, macOS, ARM64, "${{ inputs.runner_label }}"]
+    timeout-minutes: 30
+    env:
+      DFLASH_REQUIRE_MPS: "1"
+      PYTORCH_ENABLE_MPS_FALLBACK: "0"
+    steps:
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+        with:
+          persist-credentials: false
+      - uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
+        with:
+          python-version: "3.12"
+      - name: Install optional training dependencies
+        run: python -m pip install -r training/requirements.txt
+      - name: Record Apple GPU and library versions
+        run: |
+          python - <<'PY'
+          import json
+          import platform
+          import torch
+          import transformers
+          from draft_training.__main__ import execution_context, require_torch
+          require_torch("mps")
+          print(json.dumps({
+              "execution": execution_context(torch, "mps"),
+              "python": platform.python_version(),
+              "torch": torch.__version__,
+              "transformers": transformers.__version__,
+          }, indent=2))
+          PY
+      - name: Run actual MPS capture, training, resume, export, and evaluation
+        run: python -m unittest discover -s dev/tests/training -t . -v

README.md

@@ -170,6 +170,7 @@ ship precompiled, with no Xcode or local tuning required.
 
 ## More
 
+- [Experimental DFlash2 training](docs/dflash2-training.md): optional source-checkout fine-tuning.
 - [Development](DEVELOPMENT.md): configuration, API, internals, building, and testing.
 - [Issues and feedback](https://github.com/incoai/splash/issues)
 - [Apache-2.0](LICENSE). GGUF kernels include MIT-licensed material from

THIRD_PARTY_NOTICES

@@ -236,3 +236,30 @@ NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
 LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
 OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
 WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+
+DFlash (https://github.com/z-lab/dflash)
+
+The training reference in draft_training/model.py adapts dflash/model.py at
+07ebd93db9f472af339b644bb70221ad8428328a.
+
+MIT License
+
+Copyright (c) 2026 Z Lab
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

dev/Makefile

@@ -88,11 +88,11 @@ test-python: install-development $(TARGET)
 
 check-python: install-development test-python
 	$(PYTHON) -m compileall -q server install/*.py \
-		dev/tools dev/tests dev/benchmarks
+		dev/tools dev/tests dev/benchmarks draft_training
 	$(PYTHON) -m ruff check --config dev/ruff.toml server \
-		install/*.py dev/tools dev/tests dev/benchmarks
+		install/*.py dev/tools dev/tests dev/benchmarks draft_training
 	$(PYTHON) -m ruff format --check --config dev/ruff.toml server \
-		install/*.py dev/tools dev/tests dev/benchmarks
+		install/*.py dev/tools dev/tests dev/benchmarks draft_training
 
 # It parses the server's sources, which need the environment's Python
 # (3.12-3.14); BUILD_ID_PYTHON, the build tools' python3, may be Xcode's 3.9.

dev/tests/training/fixtures.py

@@ -0,0 +1,171 @@
+"""Small independent checkpoints for optional tensor tests; no downloads."""
+
+import json
+from pathlib import Path
+
+
+def tiny_config():
+    return {
+        "architectures": ["DFlash2DraftModel"],
+        "model_type": "qwen3",
+        "hidden_size": 16,
+        "intermediate_size": 32,
+        "num_hidden_layers": 2,
+        "num_attention_heads": 2,
+        "num_key_value_heads": 1,
+        "head_dim": 8,
+        "num_target_layers": 4,
+        "vocab_size": 32,
+        "hidden_act": "silu",
+        "rms_norm_eps": 1e-6,
+        "attention_bias": False,
+        "attention_dropout": 0.0,
+        "max_position_embeddings": 128,
+        "sliding_window": 16,
+        "use_sliding_window": True,
+        "layer_types": ["sliding_attention"] * 2,
+        "is_causal": False,
+        "tie_word_embeddings": False,
+        "rope_parameters": {"rope_type": "default", "rope_theta": 10000.0},
+        "dflash_config": {
+            "block_size": 8,
+            "conv_kernel_size": 2,
+            "conv_group_size": 4,
+            "selector_rank": 4,
+            "selector_top_k": 4,
+            "target_layer_ids": [0, 2],
+            "mask_token_id": 31,
+        },
+    }
+
+
+def tiny_components():
+    import torch
+    from torch import nn
+
+    from draft_training.model import DraftModel
+
+    torch.manual_seed(7)
+    config = tiny_config()
+    model = DraftModel(config).to(dtype=torch.bfloat16)
+    embedding = nn.Embedding(32, 16).to(dtype=torch.bfloat16).requires_grad_(False)
+    head = nn.Linear(16, 32, bias=False).to(dtype=torch.bfloat16).requires_grad_(False)
+    return model, embedding, head
+
+
+def cached_features(directory, tokens=None):
+    import torch
+    from safetensors.torch import save_file
+
+    from draft_training.config import file_digest, write_json
+
+    directory = Path(directory)
+    directory.mkdir(parents=True)
+    torch.manual_seed(13)
+    sample = {
+        "tokens": torch.arange(24) % 30 if tokens is None else tokens,
+        "features": torch.randn(24, 32, dtype=torch.bfloat16),
+        "assistant_mask": torch.arange(24) >= 4,
+    }
+    path = directory / "000000.safetensors"
+    save_file(sample, str(path))
+    manifest = {
+        "binding": {"supervision": "target-generated"},
+        "complete": True,
+        "shards": [{"file": path.name, "sha256": file_digest(path), "tokens": 24}],
+    }
+    write_json(directory / "manifest.json", manifest)
+    return sample, manifest
+
+
+def write_data(path, prompt="hello", answer="response", **extra):
+    Path(path).write_text(
+        json.dumps(
+            {
+                "messages": [
+                    {"role": "user", "content": prompt},
+                    {"role": "assistant", "content": answer},
+                ],
+                **extra,
+            }
+        )
+        + "\n"
+    )
+
+
+def tiny_target(directory):
+    import torch
+    from tokenizers import Tokenizer
+    from tokenizers.models import WordLevel
+    from tokenizers.pre_tokenizers import Whitespace
+    from transformers import (
+        PreTrainedTokenizerFast,
+        Qwen3_5Config,
+        Qwen3_5ForConditionalGeneration,
+    )
+
+    config = Qwen3_5Config(
+        text_config=dict(
+            vocab_size=32,
+            hidden_size=16,
+            intermediate_size=32,
+            num_hidden_layers=4,
+            num_attention_heads=2,
+            num_key_value_heads=1,
+            head_dim=8,
+            linear_key_head_dim=4,
+            linear_value_head_dim=4,
+            linear_num_key_heads=1,
+            linear_num_value_heads=2,
+            linear_conv_kernel_dim=2,
+            full_attention_interval=2,
+            max_position_embeddings=128,
+            rope_parameters={
+                "rope_type": "default",
+                "rope_theta": 10000.0,
+                "partial_rotary_factor": 1.0,
+                "mrope_section": [1, 1, 2],
+            },
+            eos_token_id=None,
+            pad_token_id=0,
+        ),
+        vision_config=dict(
+            depth=1,
+            hidden_size=16,
+            intermediate_size=32,
+            num_heads=2,
+            out_hidden_size=16,
+            num_position_embeddings=16,
+            patch_size=2,
+            temporal_patch_size=2,
+            spatial_merge_size=2,
+        ),
+        image_token_id=28,
+        video_token_id=29,
+    )
+    torch.manual_seed(19)
+    model = Qwen3_5ForConditionalGeneration(config).to(torch.bfloat16).eval()
+    tokenizer = Tokenizer(
+        WordLevel(
+            {
+                "[PAD]": 0,
+                "[UNK]": 1,
+                "user": 2,
+                "assistant": 3,
+                "system": 4,
+                **{f"w{i}": i for i in range(5, 32)},
+            },
+            unk_token="[UNK]",
+        )
+    )
+    tokenizer.pre_tokenizer = Whitespace()
+    fast = PreTrainedTokenizerFast(
+        tokenizer_object=tokenizer, pad_token="[PAD]", unk_token="[UNK]"
+    )
+    fast.chat_template = (
+        "{% for m in messages %}{{ m.role }} {{ m.content }} {% endfor %}"
+        "{% if add_generation_prompt %}assistant {% endif %}"
+    )
+    model.save_pretrained(directory)
+    fast.save_pretrained(directory)
+    return model, fast