docs(dflash-npu): correct status — DFlash training blocked on NPU by flex_attention (vllm-project#531)

Walk back the premature 'verified end-to-end' claim: training never completes a
step on NPU. Document the real blocker found by checking the repo issues:

- DFlash (and EAGLE3/PEAGLE) force simple_flex_attention; PyTorch flex_attention
  rejects npu devices (issue vllm-project#531, open). Only MTP uses eager attention.
- @torch.compile on the dflash forward fails triton-ascend codegen with
  block_size=16 (NPU flavor of vllm-project#544 / PR vllm-project#547); TORCHDYNAMO_DISABLE=1 gets past
  it but then hits vllm-project#531.
- add §9 with the algorithm/attention support table and options (train MTP, patch
  DFlash to eager/dense-mask attention, or track upstream).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: Sawyer117

docs/deployment/ascend-npu-dflash-training.md

@@ -5,12 +5,19 @@ End-to-end guide for training a **DFlash** draft model for **Qwen3-8B** on
 target and extracts hidden states on the fly). Hyperparameters are aligned to the
 SpecForge `docs/ascend_npu/run_qwen3_8b_dflash_npu.sh` reference run.
 
-> **Status (June 2026).** ✅ **Verified end-to-end on NPU** (vllm-ascend 0.20.2rc1,
-> CANN 9.0.0): prepare → online vLLM hidden-state extraction → FSDP DFlash training
-> all run and produce loss + checkpoints. The `extract_hidden_states` path works on
-> the vllm-ascend NPU runner. Hyperparameters align to the SpecForge ascend_npu
-> reference (gaps in §2.1 use speculators defaults). Install: see
-> [`ascend-npu-conda.md`](./ascend-npu-conda.md).
+> **Status (June 2026).** ⚠️ **Partially working — DFlash training is BLOCKED on NPU
+> by an upstream bug.** What's verified: install, data prep, and the online vLLM
+> `extract_hidden_states` path (vLLM serves Qwen3-8B on the NPU runner and writes
+> hidden states; the fc dimension / layer alignment is correct). What's blocked: the
+> **DFlash forward** — it uses `flex_attention` (`simple_flex_attention`), which
+> PyTorch rejects on NPU (`FlexAttention is only supported on CUDA, CPU or HPU
+> devices`, issue
+> [#531](https://github.com/vllm-project/speculators/issues/531), open), and its
+> `@torch.compile` fails to codegen on triton-ascend (NPU flavor of
+> [#544](https://github.com/vllm-project/speculators/issues/544) / PR #547). The
+> training loop never completes a step. **EAGLE3 and PEAGLE share the same flex
+> blocker; only MTP uses eager attention and could train on NPU today.** See §9.
+> Install: [`ascend-npu-conda.md`](./ascend-npu-conda.md).
 
 ## 1. Prerequisites
 
@@ -216,7 +223,10 @@ appended; that field name is a vLLM naming artifact, the `method` is still
 > an absolute `$HS_DIR` in both terminals removes all cwd ambiguity; a mismatch
 > silently stalls online training.
 
-## 6. Step 3 — train the DFlash draft (FSDP) ✅ verified
+## 6. Step 3 — train the DFlash draft (FSDP) ⚠️ blocked on NPU (see §9)
+
+The command below is correct (hyperparameters + layer alignment), but on NPU it
+currently crashes in the **first forward** before any step completes — see §9.
 
 ```bash
 # terminal 2, 6 training cards (2-7); vLLM holds 0,1. TASK_QUEUE_ENABLE=2 ok here
@@ -278,3 +288,39 @@ online path is confirmed or ruled out on NPU.
 `./output/qwen3-8b-dflash-npu/checkpoints/<epoch>/` holds a self-contained
 speculator (`config.json`, `model.safetensors`, optimizer/scheduler state). Each is
 deployable in vLLM. (Qwen3 draft vLLM-deploy compatibility — see §2 note ².)
+
+## 9. ⚠️ Known blocker: DFlash forward on NPU
+
+The DFlash draft forward fails on NPU for two compounding reasons (neither is a
+config error):
+
+1. **`flex_attention` is not supported on NPU.** DFlash forces
+   `simple_flex_attention` (`models/dflash/core.py`), and PyTorch's `flex_attention`
+   raises `FlexAttention is only supported on CUDA, CPU or HPU devices. Found input
+   tensors on npu device` — tracked in
+   [#531](https://github.com/vllm-project/speculators/issues/531) (open). **EAGLE3
+   and PEAGLE force the same `simple_flex_attention`, so they're blocked too.**
+2. **`@torch.compile` on the forward fails to codegen on triton-ascend.** With
+   `block_size=16` the generated `copy_full_slice` kernel dies with
+   `NoTritonConfigsError: ... Cannot broadcast [8,16] vs [8,1]` (NPU flavor of
+   [#544](https://github.com/vllm-project/speculators/issues/544) / PR #547).
+   `export TORCHDYNAMO_DISABLE=1` makes `@torch.compile` a no-op (eager forward) and
+   gets past this — but then you hit blocker #1.
+
+**Algorithm support on NPU (by attention impl):**
+
+| Algorithm | Attention | NPU today |
+|---|---|---|
+| DFlash / EAGLE3 / PEAGLE | `simple_flex_attention` | ❌ blocked (#531) |
+| MTP | `eager` (`models/mtp/core.py`) | ✅ viable (no flex) |
+
+**Options:**
+- **Train MTP instead** — uses eager attention, no flex. Different algorithm (needs
+  a native MTP head + `--from-pretrained`), so not a like-for-like DFlash↔SpecForge
+  comparison, but it actually trains on NPU now.
+- **Patch DFlash to an eager / dense-mask attention path** — what #531 needs: replace
+  the flex `BlockMask` (`create_block_mask`) with a materialized dense additive mask
+  and route through `eager_attention_forward`, plus `TORCHDYNAMO_DISABLE=1`.
+  Non-trivial (mask formats differ) but feasible on an editable install; a real
+  upstream contribution.
+- **Track #531 / PR #547** for an upstream fix.