perf(pflash): build the Qwen3.5 drafter's causal mask in the graph

[Graffioh]

Summary

Stacked on #746. The branch is #746's current head (555ee2a97) plus one commit, 8fb32b938. Once #746 merges, the diff is that commit only.

The Qwen3.5-0.8B PFlash drafter spent most of its time on a causal mask, not on attention. For every full-attention layer and every 1,024-token ubatch, it filled a [align32(kv_len), align32(n)] F16 mask on the host and uploaded it: about 43 GB of uploads at 120K. The profile put that at 51–58% of drafter time at 79K–118K on the R9700.

This PR makes two changes; output is byte-identical.

• The mask is built on the device, inside the graph. build_qwen35_causal_mask() writes the same bytes the host loop wrote:
  1. a zero fill (a repeat of one F16 element);
  2. the diagonal block over keys [kv_start, kv_pad), from ggml_tri of an F32 −inf square cast into place, which also covers the kv padding;
  3. −inf over the padded query rows.
     Unaligned resume and checkpoint starts are handled; there is no CPU fallback. The extra memory is about 8 MiB at 120K.
• The compute buffer is reserved once. The mask grows with kv_len, so gallocr reallocated the compute buffer on almost every ubatch of the first attention layer: 117 times in 119 ubatches at 120K, replayed on the CPU allocator. On HIP each reallocation is a free plus a malloc of up to ~240 MB. The buffer is now reserved before the layer loop, for the full-attention graph of the ubatch with the largest mask.

Both scorers share the change: the strict head path and the legacy running-max path.

Results

R9700 (gfx1201), 27B IQ4_XS target, v14 probe, native head. The A/B compares the same commit with and without the fix, both with PFLASH_DRAFTER_PROFILE=1.

workload	drafter before → after	TTFT before → after	identical to before
dev-qa 32K (30)	2.51 → 1.59 s (−37%)	11.95 → 11.05 s	30/30
code-debug 98K (4)	18.59 → 8.18 s (−56%)	38.31 → 28.08 s (−27%)	4/4
NIAH 120K (5)	28.24 → 11.63 s (−59%)	50.61 → 34.27 s (−32%)	5/5
NIAH 120K, --prefill-skip-park (5)	28.84 → 11.54 s (−60%)	46.44 → 29.49 s (−36%)	5/5

• "Identical" means answers, selected_chunks, compressed_ids and chunk_scores all match, with the scores bitwise equal (max abs diff 0.0).
• Peak VRAM is equal or lower in every run (e.g. 22.13 → 20.81 GB on code-debug).
• Attention, DeltaNet and scoring times are unchanged. The whole gain is the mask and the reallocation.

Tests

• New test/test_qwen35_causal_mask.cpp, in test_server_unit. It compares the device mask with the host mask byte for byte on the CPU backend.
  • Grid: starts {0, 1, 777, 1024, 5120, 18976, 19937} × n {1, 7, 64, 1000, 1024}, with kv up to 20,961.
  • It fails if the diagonal is shifted by one, or if the padded-row write is skipped.
• test_server_unit: 688 passed, 0 failed, 2 skipped (GPUs hidden).
• HIP: FILL and TRI are supported for F32, the F16 repeat goes through binbcast, and the F32→F16 copy into a strided view goes through cpy. The flash-attention mask stays F16 and contiguous, as before.

Risks

• The mask copies run before the attention because of graph node order, not a data dependency. The KV-cache writes in this file already rely on the same ordering. The opt-in GGML_CUDA_GRAPH_OPT=1 reorder pass would break both.
• The buffer reaches its peak size from layer 0, instead of growing to it by the end of the first attention layer. The peak itself is unchanged, plus about 8 MiB.
• The served A/B ran only the strict head path. The legacy path is covered by the unit test only.

🤖 Generated with Claude Code