[RFC] DFlash2 for Qwen3.8 on consumer Blackwell: integration boundary before we propose a PR

Summary

We have a working DFlash2 (V2) speculative decoder for Qwen3.8-27B NVFP4 on a single RTX 5090
(SM120) and measured +46.25% aggregate decode throughput on a short-context greedy workload.
Before proposing any PR we would like maintainer direction on where each piece belongs, because
several of our controls are deliberately experimental and should not become public API by
accident.

This is a design question, not a patch submission. We are explicitly not proposing to upstream
our branch as-is: it combines five separable concerns and is based on a grafted older commit.

Measured result

Single RTX 5090, SM120, 32 GiB. Greedy. Median of five trials, 15 concurrent requests × 512
completion tokens.

Configuration at batch 15	Median aggregate decode tok/s	Trial range
Target only	770.46	768.52–771.80
DFlash2 profile	1126.77	1093.36–1199.87

Reproducer shape:

GPU: RTX 5090, SM120, 32 GiB
Target: RadixArk/Qwen3.8-27B-NVFP4
Draft:  z-lab/Qwen3.8-27B-DFlash2
Backend: PyTorch;  max_seq_len 4096;  max_batch_size 15;  max_num_tokens 4096
KV dtype: fp8;  DFlash max_draft_len: 7;  sampling: greedy
Peak memory: 29.21 GiB of 32 GiB

What this does NOT establish

Stated plainly so the number is not over-read: no long-context, non-greedy, multi-GPU,
cancellation-stress or mixed-shape soak qualification has been done. The result is one model, one
GPU, one workload shape, greedy only.

Peak throughput also currently depends on an external inference-only FlashAttention 2.8.4 build
for SM120, which is not part of our branch. A PyTorch SDPA compatibility path exists and matched
the DFlash geometry within max abs error 0.0078125, but is slower (368.06 vs 414.48 tok/s at batch
8). Any upstream discussion should treat native SM120 attention as an explicit dependency.

Relationship to existing issues

• #14843 — AutoDeploy: DFlash speculative decoding support
• #16005 — Paged / KV-manager-resident storage for the DFlash draft context, replacing the dense
  buffer. Directly relevant: our private-cache attention is the same design area, and we would
  rather align with whatever storage model you intend than introduce a second one.
• #16277 — test_dflash_qwen3_5_4b flaky-marginal mean-accepted

Questions

1. Integration boundary. Should DFlash2 be a new decoding_type alongside DFlash, or a
capability flag on the existing one? Our port touches _torch/speculative/dflash.py as well as
adding dflash2.py.

2. Private-cache attention. Should this be expressed through an existing TensorRT-LLM
attention abstraction, or remain a DFlash-owned module? See #16005 — if paged/KV-manager-resident
storage is the intended direction, we would build to that instead.

3. Hybrid recurrent state. The Qwen target is hybrid (gated delta-net). We select cached GDN
replay where viable and avoid a full speculative recurrent-state snapshot, plus context-slot
mapping, accepted-token state handling, rollback and KV rewind. Is there an intended lifecycle for
recurrent state under speculation that we should conform to?

4. Draft checkpoint schema. We extract draft tensors from a published checkpoint and share the
target embedding and lm-head. Is there a preferred schema for auxiliary-model checkpoints?

5. Configuration ownership. Our working controls are environment variables and should not
survive as public API. Where should each live?

• shape-dependent selector policy (ours: on at batch 15, off at batch 8 — worth 3.65% at 15,
  and negative at 8)
• greedy-only CUDA-graph capture (avoids a graph-capture path that fails an SM120 capability
  check)
• single active speculative generation graph, with recapture at an idle admission boundary
  (~0.5–0.6 s, and liable to thrash on oscillating shapes — needs hysteresis)
• draft-weight quantization at load (NVFP4 draft linears; we would prefer a prequantized draft
  checkpoint and typed config upstream)

6. Scope of a first PR. Our intended sequence is: generic checkpoint prefix filtering first
(independent, no DFlash dependency), then the core decoder with focused tests, and only then the
policy/graph/quantization concerns as separate proposals. Does that ordering suit you?

A generalization we think is worth stating

Two of these are not DFlash-specific and may matter to other speculative methods:

• Address-safe graph-cache governance. Graph captures embed workspace addresses, so eviction
  and recapture need to be explicit runtime concepts rather than per-feature workarounds.
• Auxiliary-model resource governance. Quantizing the draft model helped mainly by changing
  the feasible concurrency envelope — it was not faster at every isolated shape. Evaluating an
  auxiliary transformation by kernel time alone would have rejected the change that produced the
  headline gain.

We are happy to split this into separate RFCs if you prefer; we filed one issue to avoid
fragmenting the design discussion before you have set a direction.

Start by reading _torch/speculative/dflash.py and the related issues #14843, #16005, and #16277. The first step is to get maintainer decisions on the integration boundary, cache and recurrent-state lifecycle, checkpoint schema, and configuration ownership; done means an agreed split and ordering for focused follow-up proposals.