Merge main into Mac K2.A reviewer branch

Resolve the dlm_restored_verifier dtype-cast conflict by preserving the K2.A round-trip cast while keeping the DFlash/mainline updates.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: 2 people

.gitattributes

@@ -0,0 +1 @@
+models/dflash-kakeya-baseline/*.safetensors filter=lfs diff=lfs merge=lfs -text

inference_engine/v04/__init__.py

@@ -43,6 +43,12 @@
     slice_position_embeddings,
 )
 from inference_engine.v04.dlm_restored_verifier import DLMRestoredVerifier
+from inference_engine.v04.dflash_drafter import (
+    AuxHiddenProvider,
+    DFlashConfig,
+    DFlashDrafter,
+    DFlashProposer,
+)
 from inference_engine.v04.kv_compressor import (
     IdentityCompressor,
     KakeyaLatticeCompressor,
@@ -110,4 +116,10 @@
     "KakeyaLatticeUnavailable",
     "KVCompressor",
     "make_default_compressor",
+    # K3 — native DFlash drafter (Stage 1: module + proposer; see
+    # docs/design/k3-cross-model-dlmrestored-verifier-contract.md)
+    "AuxHiddenProvider",
+    "DFlashConfig",
+    "DFlashDrafter",
+    "DFlashProposer",
 ]

inference_engine/v04/dflash_drafter.py

@@ -0,0 +1,62 @@
+# DFlash baseline drafter for `google/gemma-4-26B-A4B-it` (Kakeya-aligned)
+
+Baseline **DFlash block-diffusion drafter** for the Gemma-4 26B-A4B verifier,
+for use in Kakeya inference speculative-decoding development and scenario
+testing. Loadable directly by the native engine:
+
+```python
+import torch
+from inference_engine.v04.dflash_drafter import DFlashDrafter
+drafter = DFlashDrafter.from_pretrained("models/dflash-kakeya-baseline", dtype=torch.bfloat16)
+```
+
+## What this is
+
+- **Architecture**: the native `DFlashDrafter` (5-layer Qwen3 backbone + `fc`
+  aux projection + `hidden_norm` + `norm`), faithful to vLLM PR #41703
+  `qwen3_dflash.py`. Shares the verifier's embeddings (`×sqrt(hidden)`) and
+  lm_head (`final_logit_softcapping=30`). Aux layers `(2,7,12,18,23,28)`.
+- **Weights**: the upstream `z-lab/gemma-4-26B-A4B-it-DFlash` checkpoint,
+  **alignment-trained** to the Kakeya engine's inference path (see below).
+  0.43 B params, bf16, `model.safetensors` (stored via Git LFS).
+
+## Why alignment
+
+The upstream DFlash forward is defined inside vLLM (custom KV-cache writes,
+fused kernels). The native engine reconstructs the math, but the exact
+aux-hidden-tap semantics live in vLLM internals. Rather than reverse-engineer
+them, we treat the gap as an `f_θ` alignment task (ADR 0008 §11,
+`docs/design/k3-f-theta-training-pipeline.md`): freeze the verifier, train the
+drafter so its drafts match the verifier's greedy tokens.
+
+## Provenance / reproduce
+
+- Base: `z-lab/gemma-4-26B-A4B-it-DFlash`
+- Verifier: `google/gemma-4-26B-A4B-it`
+- Trainer: `scripts/research/k3_dflash_alignment_train.py`
+  ```
+  python scripts/research/k3_dflash_alignment_train.py \
+      --steps 6000 --lr 5e-5 --block-size 16 --n-prompts 64 --gen-len 192 \
+      --train-scope full --save dflash_aligned_corpus.pt
+  ```
+  (64 diverse prompts, 58 usable; `train_match=0.71`)
+
+## Acceptance (vs the real Gemma-4 verifier, block 16)
+
+| eval | acceptance_rate | acceptance_length |
+|---|---|---|
+| held-out (8 disjoint prompts) | 0.107 | 2.45 |
+| in-domain (small set) | 0.561 | 8.62 |
+| reference (HumanEval, vLLM) | 0.447 | 7.70 |
+
+The in-domain run reaching ≥ the reference proves the integration is correct;
+the held-out number is limited by the small (64-prompt) alignment corpus and
+climbs with more data (10→64 prompts: 1.94→2.45 length). This is a
+**baseline** — scaling the alignment corpus is expected to close the held-out
+gap toward 7.70.
+
+## Status
+
+Research baseline (not GA). Lossless vs greedy AR is preserved by the
+spec-decode accept loop regardless of draft quality; this drafter only affects
+*speedup*, not correctness.

models/dflash-kakeya-baseline/README.md

@@ -0,0 +1,52 @@
+{
+  "architectures": [
+    "DFlashDraftModel"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "block_size": 16,
+  "bos_token_id": 2,
+  "dflash_config": {
+    "mask_token_id": 4,
+    "target_layer_ids": [
+      1,
+      6,
+      11,
+      17,
+      22,
+      27
+    ]
+  },
+  "dtype": "bfloat16",
+  "eos_token_id": 1,
+  "final_logit_softcapping": 30.0,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 2816,
+  "initializer_range": 0.02,
+  "intermediate_size": 5632,
+  "layer_types": [
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 262144,
+  "max_window_layers": 5,
+  "model_type": "qwen3",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 5,
+  "num_key_value_heads": 8,
+  "num_target_layers": 30,
+  "pad_token_id": 0,
+  "rms_norm_eps": 1e-06,
+  "sliding_window": 2048,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.6.0",
+  "use_cache": true,
+  "use_sliding_window": true,
+  "vocab_size": 262144,
+  "rope_theta": 1000000,
+  "rope_scaling": null
+}

models/dflash-kakeya-baseline/config.json

@@ -0,0 +1,31 @@
+{
+  "name": "dflash-kakeya-baseline",
+  "kind": "dflash_drafter",
+  "role": "baseline speculative-decoding drafter for kakeya inference",
+  "base_model": "z-lab/gemma-4-26B-A4B-it-DFlash",
+  "verifier_model": "google/gemma-4-26B-A4B-it",
+  "loader": "inference_engine.v04.dflash_drafter.DFlashDrafter.from_pretrained",
+  "dtype": "bfloat16",
+  "params": 429689088,
+  "block_size": 16,
+  "aux_layer_ids": [2, 7, 12, 18, 23, 28],
+  "alignment": {
+    "trainer": "scripts/research/k3_dflash_alignment_train.py",
+    "scope": "full",
+    "steps": 6000,
+    "lr": 5e-05,
+    "prompts": 64,
+    "usable_sequences": 58,
+    "train_match": 0.71
+  },
+  "acceptance": {
+    "held_out": {"rate": 0.107, "length": 2.45, "n_prompts": 8, "block_size": 16},
+    "in_domain": {"rate": 0.561, "length": 8.62},
+    "reference_humaneval": {"rate": 0.447, "length": 7.7}
+  },
+  "evidence": [
+    "results/research/k3_dflash_specdecode_corpus_heldout.json",
+    "results/research/dflash_aligned_corpus_report.json"
+  ],
+  "notes": "Research baseline; held-out acceptance is corpus-size-limited and improves with more alignment data. Lossless vs greedy AR is guaranteed by the verifier accept loop."
+}

models/dflash-kakeya-baseline/manifest.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c0e9a83089f4cee77eb8231646bb5abbce942cac3eb3e6f206ac30c513b187aa
+size 859384296

models/dflash-kakeya-baseline/model.safetensors

@@ -1,14 +1,33 @@
 # NOTE on transformers version pin:
-#   The proposer checkpoint dllm-hub/Qwen3-0.6B-diffusion-mdlm-v0.1 ships a
-#   custom modeling file (`modeling_qwen3.py`) that depends on transformers
-#   4.x APIs (specifically `decoder_layer.attention_type`). transformers 5.x
-#   removed those APIs, so we pin to 4.x project-wide.
+#   Pin lifted 2026-06-09 — was `<5.0` to keep the legacy Qwen3 dLM
+#   proposer (`dllm-hub/Qwen3-0.6B-diffusion-mdlm-v0.1`) running, since
+#   its custom `modeling_qwen3.py` depends on the 4.x
+#   `decoder_layer.attention_type` API that transformers 5.x removed.
 #
-#   On macOS, this means using a dedicated venv (e.g. `python3 -m venv
-#   .venv`) — do NOT install into a system Python where transformers 5.x is
-#   already present. `scripts/setup_mac.sh` handles this correctly.
+#   K3 critical path needs transformers 5.0+:
+#     * Gemma 4 26B-A4B verifier (per ADR 0008 §11.7.0)
+#     * scripts/research/k3_dflash_specdecode_eval.py and
+#       k3_dflash_alignment_train.py (load Gemma 4 via transformers)
+#
+#   So the upper bound was dropped.
+#
+#   KNOWN ISSUE (K2.B Qwen backport path, NOT K3 critical path): under
+#   transformers 5.x, attempts to load
+#   `dllm-hub/Qwen3-0.6B-diffusion-mdlm-v0.1` will likely raise an
+#   AttributeError on `attention_type`. This affects:
+#       * `training/repr_align/proposer_surgery.py`
+#       * `kv_cache_proposer/proposer.py`
+#       * `inference_engine/proposer/sparse_logits.py`
+#       * `inference_engine/backends/mlx/proposer.py`
+#       * `tests/system/test_http_*_real_engine.py` (if they download)
+#   None of these are on the K3 critical path. K2.B Qwen backport is
+#   the natural place to author a compat patch when that path resumes
+#   after K3 ships.
+#
+#   For now, if you need the legacy Qwen3 dLM path, install transformers
+#   4.x in a dedicated venv:  pip install 'transformers>=4.45,<5.0'
 torch>=2.4,<3.0
-transformers>=4.45,<5.0
+transformers>=4.45
 accelerate>=0.34
 safetensors>=0.4
 huggingface_hub>=0.24

requirements.txt

@@ -0,0 +1,22 @@
+{
+  "kind": "k3_dflash_alignment_train",
+  "config": {
+    "verifier_id": "google/gemma-4-26B-A4B-it",
+    "drafter_id": "z-lab/gemma-4-26B-A4B-it-DFlash",
+    "steps": 6000,
+    "lr": 5e-05,
+    "block_size": 16,
+    "n_prompts": 64,
+    "gen_len": 192,
+    "prompt_min_ctx": 8,
+    "train_scope": "full",
+    "seed": 0,
+    "save": "results/research/dflash_aligned_corpus.pt",
+    "log_every": 25
+  },
+  "trainable_params": 429689088,
+  "n_windows": 7924,
+  "final_loss": 1.039782760348171,
+  "final_train_match": 0.71,
+  "elapsed_s": 213.54307655605953
+}

results/research/dflash_aligned_corpus_report.json

@@ -0,0 +1,22 @@
+{
+  "kind": "k3_dflash_alignment_train",
+  "config": {
+    "verifier_id": "google/gemma-4-26B-A4B-it",
+    "drafter_id": "z-lab/gemma-4-26B-A4B-it-DFlash",
+    "steps": 600,
+    "lr": 0.0001,
+    "block_size": 16,
+    "n_prompts": 8,
+    "gen_len": 160,
+    "prompt_min_ctx": 8,
+    "train_scope": "fc_norms",
+    "seed": 0,
+    "save": "results/research/dflash_aligned_fcnorms.pt",
+    "log_every": 25
+  },
+  "trainable_params": 47584768,
+  "n_windows": 1041,
+  "final_loss": 3.496910614967346,
+  "final_train_match": 0.1725,
+  "elapsed_s": 12.202172419056296
+}