perf(speculative): speed up DFlash/DSpark drafter training 3.12x

[h-guo18]

Summary

Three independent optimisations to the DFlash-family training step, found by attributing every GPU kernel to the CPU scope that launched it. 1.690 → 0.541 s/step (3.12×). Each is separately motivated and separately measured, and they are three commits so they can be reviewed and reverted independently.

The three changes

1. FlexAttention over a block-sparse mask (2.38×). The draft handed SDPA a materialized [B, 1, Q, KV] additive mask. That disqualifies every fused backend — FlashAttention rejects arbitrary masks and caps head_dim at 256 — so PyTorch fell through to the cutlass memory-efficient backend, whose only kernel is sm80. Every DFlash-family drafter was running an Ampere attention kernel on Blackwell at ~2.5% of peak; that one backward kernel was 59% of the step. The predicate is ~33% dense and expresses exactly as a BlockMask.

Two portability constraints are baked into dflash_flex_attention.py, both found the hard way on torch 2.11 / sm103:

• head_dim > 256 must pin 32×32 kernel tiles — 64×32 and 64×64 give misaligned address, 128×32 gives unspecified launch failure.
• The BlockMask block size must divide the kernel's BLOCK_M/BLOCK_N. head_dim > 256 pins 32×32 so 64 is safe, but head_dim <= 256 lets Inductor autotune up to BLOCK_M=128, where a 64-wide mask block raises "Q and KV block size must be divisible by BLOCK_M and BLOCK_N". Gemma-4 dodged this only because head_dim 512 pins the tiles; every head_dim <= 256 drafter would have hit it.

enable_gqa=True is deliberately not used — forward matches, but backward took 568 ms, 10× the pre-repeated path. K/V are pre-repeated instead.

2. Teacher rows on demand, one device sync instead of five (1.13×). The offline/streaming path built full-sequence [B, seq, vocab] base logits and gathered a few thousand rows out of them; DFlashBaseModelOutput now carries the post-final-norm base hidden and takes defer_lm_head, so the loss projects only the rows it reads (~17 GB/step of traffic at this shape). The Markov bias now folds in place rather than allocating a third [N, vocab] tensor — safe because neither lm_head nor markov_w2 saves its output for backward, and the uncorrected logits are dead whenever dflash_report_acc is off. And five per-step .item() calls became one torch.stack([...]).tolist(): each is a full device synchronize, and five of them chop up the window in which DDP's all-reduce hides behind backward. That sync change alone was worth 39 ms of NCCL exposure.

3. Tensor.split instead of slicing (1.16×). _tvd_per_token chunked its row axis with final_logits[i : i + chunk_size]. Slicing and splitting agree on every value, so no numerical test can tell them apart — only the backward graph can, and there the difference is asymptotic: a slice per chunk creates one SliceBackward0 each, and every one zero-fills a full [N, vocab] tensor, giving O(n_chunks × N × vocab). split creates a single SplitBackward0 whose backward is one cat: O(N × vocab). At this shape that 16× amplification was 93.2 ms/step, flat across the two earlier rounds and invisible because it was never broken out of the elementwise bucket.

Correctness

• 113 tests pass (5 new FlexAttention GPU, 11 new chunking, 97 pre-existing).
• Change 3 is bitwise identical at the production shape: per-token TVD, grad(hidden) and grad(markov_w2.weight) all compare exactly equal.
• End to end over 40 logged steps, max |after − before| on the loss is 0.018, against a 0.025 spread among three earlier runs — inside the pipeline's own run-to-run noise.

The new tests close two gaps. TestTvdPerTokenChunking parametrises chunk_size over values that divide N and values that leave a ragged tail: every existing DSpark test runs 96 rows against the default chunk_size of 1024, so all of them fit in a single chunk and no chunk-boundary bug could fail CI. test_chunks_via_split_not_slice walks the autograd graph and asserts SplitBackward is present and SliceBackward is not — pinning an optimisation that has no observable numerical signature.

Scope

Change 1 benefits DFlash, DSpark and Domino — the branch sits inside the parent's _build_draft_attention_mask, not at its three call sites, so all three plugins get it and the existing monkeypatch-spy tests keep observing one mask build per forward. The sm80 fallback affects any drafter passing a dense mask; the magnitude is largest at head_dim > 256.

Changes 2 and 3 are DSpark-only (the code is in hf_dspark.py). Nothing in them is Gemma-4-specific; change 3's magnitude scales with rows × vocab, so a 32k-vocab drafter sees proportionally less.

All three are off by default or behaviour-preserving: dflash_use_flex_attention defaults to False, and changes 2 and 3 alter no config surface.

Operational prerequisite

TRITON_CACHE_DIR must be node-local once FlexAttention is enabled. It compiles Triton templates on the first backward, and pointing that cache at a shared filesystem makes concurrent ranks race — 16 ranks on a lustre-backed cache died with OSError: [Errno 14] Bad address reading a .json mid-write.

Limits and follow-up

All numbers come from one model on one cluster at one shape. Not checked here: whether hf_dflash.py::_compute_loss has the same [N, vocab] slicing pattern that change 3 fixes.

---

Based on #2186; the diff above is only the delta on top of it.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

---

_{Comment @coderabbitai help to get the list of available commands.}