[Spec Decoding] Add DFlash model and proposer

[aaronzhfeng]

Description

Add DFlash draft model and proposer for block-diffusion speculative decoding on JAX/TPU. DFlash predicts multiple tokens in parallel using discrete diffusion, unlike Eagle3's autoregressive drafting. This follows the same proposer pattern as Eagle3.

This is PR 1 of 3 for DFlash support:

1. This PR: Model, proposer, and unit tests (all new files)
2. Pipeline integration (modifications to existing files)
3. E2E tests and Buildkite CI

New files:

• tpu_inference/models/jax/dflash.py -- DFlash draft model (DFlashForCausalLM)
• tpu_inference/models/jax/qwen3_dflash.py -- Qwen3-specific DFlash variant with attention
• tpu_inference/layers/common/dflash_attention_interface.py -- dflash_concat_attention kernel
• tpu_inference/spec_decode/jax/dflash.py -- DFlashProposer (prepare_inputs, propose, sampling)
• tests/models/jax/test_qwen3_dflash_attention.py -- DFlash attention unit tests
• tests/models/jax/test_qwen3_dflash.py -- target layer ID selection tests
• tests/spec_decode/test_dflash.py -- proposer sampling tests

Tests

• Unit tests for DFlash attention (concat, additive bias, GQA): tests/models/jax/test_qwen3_dflash_attention.py
• Unit tests for target layer ID selection: tests/models/jax/test_qwen3_dflash.py
• Unit tests for proposer sampling: tests/spec_decode/test_dflash.py

Checklist

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have made or will make corresponding changes to any relevant documentation.

[Lumosis]

This is a large PR. Can we break it down into several small PRs to make review easier?

[aaronzhfeng]

Sorry about the large PR. The model, proposer, and attention kernel are tightly coupled (proposer calls model forward, model uses the attention kernel), so splitting them further would leave each PR non-functional on its own. All files here are new additions with no changes to existing code, which should make it easier to review.

Broke the original PR down into 3:

1. This PR (updated): DFlash model, proposer, and unit tests -- all new files, no existing files modified
2. Pipeline integration -- modifications to existing files (speculative_decoding_manager, kv_cache_manager, qwen3, tpu_runner, model_loader)
3. E2E tests + Buildkite CI

PRs 2 and 3 coming shortly.

[kyuyeunk]

Hi @aaronzhfeng! thank you for the contribution. couple of questions

• for those who aren't familiar with DFlash (like myself), can you give a brief overview & maybe a link where we can find out about more info?
• is my understanding correct that this feature is not available in vllm's pytorch model implementation? if so, is there a way for a backend that utilizes vllm's model implementation to leverage this spec decoding?
• can you share a sample command for people to try out this feature while going through the review process?

[aaronzhfeng]

Thanks for taking a look!

DFlash overview: DFlash is a block-diffusion speculative decoding method that predicts multiple tokens in parallel using discrete diffusion, instead of generating them one at a time autoregressively. Given a context, the draft model takes a block of masked/noise positions and denoises them in a single forward pass to produce K candidate tokens simultaneously. This makes drafting O(1) in block size rather than O(K).

Paper: "DFlash: Block Diffusion for Flash Speculative Decoding" (Chen et al., arXiv:2602.06036). The reference GPU implementation is at https://github.com/z-lab/dflash.

PyTorch/vLLM availability: Right now there is no DFlash support in vLLM's PyTorch backend. The DFlash authors have confirmed vLLM integration is still in progress on their end (see z-lab/dflash#6). SGLang has DFlash support via sgl-project/sglang#16818, but this PR would be the first DFlash integration in the vLLM ecosystem. It targets the JAX/TPU backend specifically, since the draft model uses non-causal attention which required a different attention path from the standard causal pipeline. A PyTorch port is feasible but not in scope for this PR.

Sample command: The unit tests in this PR can be run without a full serving setup:

pytest tests/models/jax/test_qwen3_dflash_attention.py
pytest tests/models/jax/test_qwen3_dflash.py
pytest tests/spec_decode/test_dflash.py

End-to-end serving requires the pipeline integration in PR #1869 (already open). Once both are merged, with Qwen3-4B on a TPU v5p-8:

python -m tpu_inference.entrypoint \
  --model Qwen/Qwen3-4B \
  --speculative_config '{"model": "z-lab/Qwen3-4B-DFlash-b16", "num_speculative_tokens": 15, "method": "dflash", "draft_tensor_parallel_size": 1}'

[kyuyeunk]

can you elaborate what kind of feature is missing from existing attention implementation that it requires its own separate code? if it's due to bi-directional attention, we already have an implementation for that.

[kyuyeunk]

in general, i think this function is lacking a lot of comments explaning what each line does.

[kyuyeunk]

@Lumosis can you help take a look at spec decoding part?