#1426 dflash: experimental standalone DeepSpec DFlash block drafter for Qwen3.8-Flash-Next (greedy, correctness-first)

Descripción

## Summary

Experimental support for the standalone **DeepSpec DFlash block drafter** for [`PixelML/Qwen3.8-Flash-Next-NVFP4-DFlash`](https://huggingface.co/PixelML/Qwen3.8-Flash-Next-NVFP4-DFlash), integrated alongside Strata's existing MTP speculative decoding path.

```text
Qwen3.8-Flash-Next target (any supported quant)
        +
standalone DFlash GGUF
(BF16, 498M params, no embedding / LM head)
        ↓
one parallel draft pass per cycle
        ↓
existing Strata Verifier
```

DFlash is **opt-in**. Target-only decoding and the existing MTP path are unchanged when `--dflash` is not selected.

The implementation is currently **greedy/correctness-first**, with working speculative acceptance and end-to-end tensor parity against an independent reference.

Full semantics and artifact contract are documented in [`docs/DFLASH.md`](docs/DFLASH.md).

---

## Implemented

### DFlash artifact

Adds a standalone GGUF-based DFlash artifact loaded with:

```text
--dflash FILE.gguf
```

The loader:

- validates all **58 expected tensors**
- currently accepts **BF16 only**
- validates tensor shapes against artifact metadata
- accepts both llama.cpp-style GGUF and raw Hugging Face tensor names
- rejects unsupported DFlash variants/extensions:
  - DFlash2 selector tensors
  - DSpark / Markov extras
  - confidence heads
- rejects artifacts containing their own `embed_tokens` or `lm_head`, since both are bound from the target model

`tools/dflash_gguf.py` converts the source checkpoint to GGUF and performs a bit-exact read-back self-check.

The source safetensors were also checked against the generated GGUF and match **bit-for-bit**.

DFlash is mutually exclusive with `--mtp` and currently requires:

```text
--spec 2..8
```

---

## Target feature taps

The target captures the HC-contracted **2560-wide residual** at boundaries:

```text
[4, 16, 24, 36, 44]
```

corresponding to the checkpoint's trained target layers:

```text
[3, 15, 23, 35, 43] + 1
```

The same feature definition is used by both:

- prompt / prefill
- verifier windows

`tools/dflash_taps.py` provides a cross-path validation gate.

Measured on an IQ3_XXS target with 400-token prompts:

```text
cosine similarity: 0.9966 – 0.9999
```

The semantics were cross-checked against the vLLM DFlash / DSpark reference implementation:

- anchor = bonus token at position `p`
- remaining block rows use the mask token
- non-causal block attention
- feature boundaries `[4,16,24,36,44]`
- attention scale `1/16`

---

## Drafter runtime

The standalone drafter implements the published DeepSpec DFlash architecture:

- 5 dense transformer layers
- hidden size: `2560`
- GQA: `24 Q / 2 KV`
- head dimension: `256`
- MLP intermediate size: `7680`
- full-head **NeoX RoPE**
- RoPE theta: `1e7`
- non-causal block attention
- target embedding and LM head reused directly from the target model
- eager BF16 execution
- single CUDA stream
- greedy proposal path

Feature fusion follows:

```text
rmsnorm(fc(concat(target taps)))
```

The resulting context representation is projected into independent per-layer DFlash K/V state.

---

## Owned DFlash K/V state

Each draft layer owns an independent `QsaState`.

A dedicated `QsaStateInitOptions` path enforces the ownership invariant:

- DFlash K/V pools live inside drafter-owned CUDA arenas
- FP16 K/V
- no elastic K/V membership
- no streaming K/V
- five independent draft-layer states

Because of this ownership requirement:

```text
--dflash + --kv-grow
```

is currently rejected explicitly.

The drafter's VRAM footprint is included in expert-cache reservation/autosizing.

---

## Speculative decode integration

DFlash plugs directly into the existing Strata verifier contract.

Each cycle performs one parallel DFlash proposal pass and builds:

```text
[anchor, K candidates]
```

for the existing `Verifier`.

Rejected candidate cells do not need drafter-side rollback. Position-addressed context/KV cells are overwritten when that position is later committed.

Additional integration guards:

- sampled requests (`temperature > 0`) decode without DFlash
- `--dflash + --serve` is currently rejected
- prompt lookup is disabled while DFlash is active
- `--mtp + --dflash` is rejected
- `--kv-grow + --dflash` is rejected

A sampled request reports explicitly:

```text
dflash: sampled request decoded without the drafter
```

---

## Correctness debugging and fixes

Two independent bugs were found while bringing the implementation to non-zero acceptance.

### 1. Layer-1 anchor-row RoPE corruption

The original tensor-level divergence first appeared at:

```text
L1 q_ row 0, heads 0..13, post-RoPE
```

The producing operation was:

```cpp
dflash_rope_neox_apply(q_, ..., pos_)
```

inside `DFlashDrafter::propose()`.

#### Root cause

`h_pos_` was a single pinned staging buffer reused for two different layouts:

- Q-position staging
- K/V-position staging

The asynchronous device copy for a layer's Q positions could still be pending when the CPU rewrote the first `K * n_head_kv` entries for the K/V layout.

A late DMA could therefore observe the rewritten bytes.

For query row 0 this caused the first 14 of 24 Q heads to be roped at positions:

```text
pos+0 .. pos+6
```

instead of all heads using:

```text
pos+0
```

Rows `1..K-1` remained correct, explaining the original **row-0-only, timing-dependent** symptom.

The same staging hazard could also allow the next prefill chunk to overwrite context position/append metadata while the previous chunk was still consuming it.

#### Fix

The staging memory is now separated into disjoint regions for:

- Q layout
- K/V layout

Rewrites inside each region carry identical values, making any overlap benign.

A synchronization point was also added after fusion's per-chunk position / step copies, where the staged values genuinely change between calls.

Relevant files:

```text
include/strata/core/dflash.hpp
src/core/dflash_runtime.cpp
```

---

### 2. DFlash proposals were not reaching the verifier

After restoring tensor parity, acceptance was still exactly zero.

This turned out to be a separate integration bug in `src/program/generate.cpp`.

The verify-window fill effectively selected:

```cpp
use_mtp ? drafts[i - 1] : (oracle_fixture ? ... : 0)
```

DFlash runs with:

```text
use_mtp == false
```

so the actual verify window was:

```text
[anchor, 0, 0, 0, ...]
```

rather than:

```text
[anchor, dflash_draft_0, dflash_draft_1, ...]
```

Acceptance was therefore zero by construction regardless of drafter quality.

DFlash proposals now use the same draft slot consumed by the existing verifier path.

---

## Parity validation

`tools/dflash_stage_parity.py` now performs full end-to-end reference validation across **all five DFlash layers and every block row**.

The harness:

- validates every row at every stage
- rolls the NumPy reference forward through all five layers
- recomputes attention cell-by-cell from the dumped K/V pools
- fails hard on invalid/out-of-range reference cells
- validates every context K/V cell
- validates every context RoPE position

For the original ~120-token fixture (`pos=120`, `K=6`):

| Stage | Row-0 cosine | Max abs row-0 | Result |
|---|---:|---:|---|
| L0 xn / q raw / q normed / q / k / v | `0.999999–1.000000` | `2.4e-07–1.4e-02` | pass |
| L0 attention | `1.000000` | `3.0e-06` | pass |
| L0 attn+o+res / MLP | `0.999997 / 0.999993` | `8.8e-03 / 4.4e-02` | pass |
| L1 xn / q raw / q normed / q | `0.999993–0.999996` | `1.6e-02–2.6e-02` | pass |
| L1 attention | `1.000000` | `4.8e-06` | pass |
| L2 / L3 / L4, all substages + attention | `0.999988–1.000000` | `≤ 1.2e-01` | pass |
| Final norm | `0.999991` | `1.9e-02` | pass |

```text
FIRST DIVERGENCE (any row): none
```

The previous apparent K/V corruption / NaN result was also traced to the debugging tooling rather than the runtime: the pool dump helper widened FP16 data using a BF16 bit shift instead of a real FP16 conversion.

That dump path now uses a proper `__half2float` conversion, and the reference oracle fails explicitly on invalid cells instead of propagating NaNs.

---

## Speculative acceptance

After both correctness fixes, DFlash produces non-zero acceptance.

| Run | Before | After |
|---|---:|---:|
| Original 119-token fixture, 24 new tokens, `K=6` | `0 / 30` | `10 / 78` (`0.128`) |
| Natural 69-token prompt, 128 new tokens | `0 / 161` | `68 / 413` (`0.165`) |

Observed effective output:

```text
24-token run:   1.71 tokens / round
128-token run:  2.13 tokens / round
```

Acceptance increased with the tested generation length:

```text
~3.0% at 24 generated tokens
→
16.5% at 128 generated tokens
```

No performance optimization has been attempted yet; the current goal is correctness and semantic validation.

---

## Greedy semantic gate

Target-only and target+DFlash decoding at `temperature = 0` are **token-for-token identical** for both the 24-token and 128-token validation runs.

There is one important test-control requirement:

```text
--expert-cache 4283
```

was pinned to the same value for both target-only and DFlash runs.

Without that control, the drafter's ~1 GiB VRAM footprint changes expert-cache autosizing. Some experts can then move between CPU and GPU execution, whose numerical rounding differs slightly and can flip near-tie target argmaxes independently of DFlash.

This caveat is documented in `docs/DFLASH.md`.

---

## Tests

Dedicated DFlash tests pass:

```text
dflash_load_test        OK
dflash_rope_parity      OK
dflash_owned_kv_test    OK
```

Full CTest result:

```text
90 / 92
```

The two remaining failures are pre-existing environment limitations:

- `ple_parity`: Q2_0 PLE table file is absent on this machine
- `expert_multi_test`: host CPU does not provide AVX-512

Neither failure is related to the DFlash changes.

---

## Current limitations

This remains a **correctness-first experimental implementation**.

Currently unsupported:

- sampled DFlash acceptance
- `--serve`
- multi-GPU
- DFlash2
- DSpark extensions
- conversation resume
- elastic `--kv-grow`

Additional constraints:

- greedy decoding only
- drafter attention window defaults to `32768` cells
- trained maximum DFlash query block is 7 candidates
- no performance claims are made yet

The performance numbers published for PixelML's implementation are **not reproduced or claimed by this PR**.

---

## Status

The core DFlash path is now functionally working:

```text
artifact loading
    ✓
target feature capture
    ✓
owned K/V state
    ✓
5-layer DFlash forward
    ✓
full tensor parity
    ✓
verifier integration
    ✓
non-zero acceptance
    ✓
greedy token-identity gate
    ✓
```

The original anchor-row divergence and the zero-acceptance blocker were two separate bugs; both are now identified and fixed.

The next step is performance work and broader validation, without changing the currently verified greedy semantics.