Perf(dflash): compile create_block_mask to avoid dense mask materialization (#731)

## Purpose

DFlash/DSpark rebuild the attention mask every training step.
Uncompiled, `create_block_mask` evaluates the mask over the full Q × KV
grid. For`--speculator-type dspark` on Qwen3-8B (`--max-anchors 3072
--block-size 8
--total-seq-len 4096`):
```
    Q  = max_anchors × block_size  = 3072 × 8       = 24576
    KV = total_seq_len + Q         = 4096 + 24576   = 28672
```
So a 24576 × 28672 dense grid (~7 GB) is built and thrown away every
step.

`torch.compile` fuses the mask evaluation with the block-sparsity
reduction, so the BlockMask is produced without ever materializing that
grid: peak mask memory drops from ~7 GB to ~0, and the mask is
bit-identical.

**Not applicable to EAGLE3/PEAGLE** — their autoregressive drafting
builds only a sequence-sized mask, so there's nothing to fix
(benchmarked: no gain).



## Measured (DSpark, Qwen3-8B, 5-layer draft, seq 4096, 2×A100)

| metric      | before (dense)   | after (compiled) |
| ----------- | ---------------- | ---------------- |
| mask build  | 29.8 ms / 7.05 GB | 2.9 ms / ~0 GB   |
| forward     | 357.9 ms         | 317.3 ms (−11%)  |
| step        | 1087.5 ms        | 1047.9 ms (−3.6%) |

## Tests

just `torch.compile` a function `create_block_mask`. No behavior change.

## Reference

https://pytorch.org/blog/flexattention/

(the blog is stale, now create_block_mask is compilable)

## Checklist

I have filled in:

- [x] The purpose of the PR, such as "Fix some issue (link existing
issues this PR will resolve)".
- [x] The test plan/results, such as providing test command and pasting
the results.
- [ ] (Optional) The necessary documentation update.
- [x] I (a human) have written or reviewed the code in this pr to the
best of my ability.

Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>

Author: WindChimeRan

src/speculators/models/dflash/core.py

@@ -25,6 +25,10 @@
 
 logger = logging.getLogger(__name__)
 
+# Compile so the mask builds block-sparse instead of materializing DFlash's huge
+# dense [Q, KV] grid every step. (No benefit for EAGLE3's small autoregressive mask.)
+_compiled_create_block_mask = torch.compile(create_block_mask)
+
 
 @SpeculatorModel.register("dflash")
 class DFlashDraftModel(DraftVocabMixin, SpeculatorModel):
@@ -58,7 +62,7 @@ def __init__(
             )
         self._attn_impl = config.transformer_layer_config._attn_implementation  # noqa: SLF001
         self._create_mask_fn = (
-            create_block_mask
+            _compiled_create_block_mask
             if self._attn_impl == "simple_flex_attention"
             else create_float_mask
             if self._attn_impl == "eager"