feat(attention): ArcFlash — Arc's own attention kernel, rung 0

[heydryft]

The name

Per TAXONOMY.md's rule that no subsystem is left without an absolute parent:

Arc → ArcInfer → ArcAttention → ArcFlash
                                ├── ArcFlash/Plan    shape → path + the D18 engagement record
                                ├── ArcFlash/Fold    GQA/MQA expansion as a reshape, never a copy
                                ├── ArcFlash/Tile    query-row tiling under a byte budget
                                └── ArcFlash/Cute    the fused CuTeDSL d=512 kernel — NOT BUILT

ArcAttention/Flash was already taken and means the wrong thing — it is the capability envelope for Dao's FA2/FA3, a wrapper around someone else's kernel. Reusing "Flash" for ours would guarantee the two get conflated in exactly the claims that matter. ArcFlash/Cute is named so rungs 3–6 have a home; no code claims it exists.

What rung 0 does

Removes the two allocations that made long prompts fail rather than merely slow, at the two sites that performed the GQA expansion:

• backends/sinks.rs — every DeepSeek-V4 layer, since head_dim 512 is outside the fused sinks kernel's {64,80,96,112,128,192,256}.
• run_attention_noflash — repeat_kv now runs inside the cuBLASLt arm instead of ahead of all three paths.

ArcFlash/Fold expresses the expansion as a reshape. repeat_kv maps x[b,kv,s,d] → y[b, kv*n_rep + rep, s, d] — heads are kv-major — so a contiguous q[B,H,Tq,D] reshapes to [B,Hkv,g*Tq,D] with no data movement, and both GEMM operands then carry batch B*Hkv. Both facts checked against the pinned candle (aeonmindai/candle @ c3bb5bf), not recalled: reshape is a view when contiguous, and broadcast_matmul is not an alternative because candle concretises the broadcast operand (its own // TODO: Avoid concretising the broadcasted matrixes). The trick is already in-tree at t_q == 1 (absorbed_mqa_decode); this generalises it to any t_q and any GQA ratio.

ArcFlash/Tile bounds the score matrix to one query tile (ARC_FLASH_TILE_BYTES, default 32 MiB). Tiling query rows is exact, not an approximation: each output row is an independent softmax, so there is no online rescaling, and sinks — a per-head constant in each row's denominator — are untouched.

Numbers, labelled

ARITHMETIC (not measured). One layer at V4 prefill, B=1 H=64 D=512 T=1400 bf16:

	Old	ArcFlash rung 0
GQA expansion of K and V	175 MiB	0
score matrices	478 MiB	64 MiB
folded query tile + K/V	—	15 MiB
peak	~653 MiB	~78 MiB

≈8.3×, pinned by a test that also refuses to pass if the shape does not actually tile.

MEASURED (CPU only — D14): 386 mistralrs-core lib tests green, including the 17 pre-existing dsv4_attention tests, which now execute through ArcFlash.

NOT MEASURED: every GPU number. No wall-clock, VRAM or tok/s claim is made here.

D18 — engagement asserted, both directions

note() logs ArcAttention: first dispatch on <path> once per path: present in the arm that uses it, absent in the arm that does not. The test instrument is a #[cfg(test)] thread-local rather than the process-wide atomics, because cargo test runs test fns in parallel and a control arm would otherwise read other tests' dispatches — that was a real failure during development, not a hypothetical.

Mutations — four applied, four caught

Mutation	Caught by
GQA heads treated rep-major instead of kv-major	fold_matches_repeat_kv_across_gqa_ratios (diff 1.5)
mask does not follow the query tile	query_tiling_is_exact_and_the_mask_follows_the_tile
sink dropped from the softmax denominator	2 ArcFlash tests and 3 pre-existing V4 tests
vendor flash disabled, everything falls to Tile	the D18 control arm

The third matters most: the existing V4 suite genuinely covers this path, so ArcFlash is not tested only by its own author's tests.

What this does NOT do

• It does not make FA2 stop being the default for GQA models. That default exists because the pinned candle-flash-attn-v3 has two specific defects — NULL tile_count_semaphore on dense causal, and a window clamp that rewrites causal to full attention on varlen causal (Fix causal attention in candle-flash-attn-v3: NULL tile_count_semaphore crash (dense) + silently non-causal varlen huggingface/candle#3606). Both live in aeonmindai/candle, a different repository, so they cannot be fixed from this branch. fa3_preferred() is already the single switch.
• It is not the fused kernel. Rungs 3–5.

For main to run

arc-tools/arcflash/rung0_gpu_validate.sh — gates on compiling the CUDA-cfg'd Rust that a macOS cargo check cannot reach, then measures ~200-word and ~1055-word prompts and checks D18 in both directions. Needs ARC_V4_MODEL and ARC_GQA_MODEL.

Surfaced, not shipped

• dsv4_attention.rs builds k_cat and v_cat with two identical Tensor::cat calls despite V == K being a V4 invariant — a full redundant copy of the union per layer per step.
• chunked_attention narrows a rank-4 mask on dim 2 unconditionally; a mask that is size-1 on the query axis would be mis-selected. ArcFlash's narrow_mask_rows guards this; the older helper does not.

[github-actions]

Code Metrics Report

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 Language              Files        Lines         Code     Comments       Blanks
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 C Header                  5          305          210           52           43
 CSS                       2         1181         1036           34          111
 CUDA                     72        24645        17702         4215         2728
 Dockerfile                1           39           22            8            9
 JavaScript               16         3546         2676          482          388
 Jinja2                    7          694          656            5           33
 JSON                     74         4600         4597            0            3
 Makefile                  1            6            5            0            1
 Metal Shading Lan|       33        12224         9431         1142         1651
 PowerShell                1          300          227           30           43
 Python                  143        14830        12217          797         1816
 Shell                    24         5924         4071         1457          396
 Plain Text                4         3801            0         2479         1322
 TOML                     33         1485         1292           43          150
 YAML                      3           25           23            2            0
─────────────────────────────────────────────────────────────────────────────────
 HTML                      4         2687         2604           43           40
 |- CSS                    2          543          479           37           27
 |- JavaScript             1         1233         1215           12            6
 (Total)                             4463         4298           92           73
─────────────────────────────────────────────────────────────────────────────────
 Jupyter Notebooks         4          122           83           23           16
 |- Markdown               1           60           30           22            8
 |- Python                 1          122          113            1            8
 (Total)                              304          226           46           32
─────────────────────────────────────────────────────────────────────────────────
 Markdown                194        38699            0        29715         8984
 |- BASH                  71         1632         1192          324          116
 |- C                      3           17           17            0            0
 |- CUDA                   2           84           56           16           12
 |- JSON                  18          708          708            0            0
 |- PowerShell             1            1            1            0            0
 |- Python                23         1008          787          113          108
 |- Rust                  65         2048         1713           77          258
 |- TOML                   6          207          164            0           43
 |- YAML                   4           38           33            5            0
 (Total)                            44442         4671        30250         9521
─────────────────────────────────────────────────────────────────────────────────
 Rust                    663       313742       271458        14550        27734
 |- Markdown             478        24544          471        21171         2902
 (Total)                           338286       271929        35721        30636
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 Total                  1284       461100       335289        76855        48956
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

[heydryft]

Measurement methodology for ArcFlash's before/after, and the number that prices it.

From tonight's profiler run on the H200 (PR #113, qtip2b, MTP depth 3, uniform
prompts, admission OFF, same binary and step count in every cell):

B	decode step	model_forward (device)	mla_attn (device)	moe (device)	host self
1	113.4 ms	104.9	58.0	29.1	9.2
8	288.5 ms	253.0	152.3	72.0	39.2
32	714.9 ms	622.3	398.9	157.9	107.8

Marginal cost of one more sequence in a decode step is 19.4 ms, and
mla_attn is 11.0 ms of it — 57%. Attention is the single largest term
both at b=1 (55% of the forward) and in the slope. That is what a fused sinks
path at head_dim 512 is worth attacking, stated as a measured quantity rather
than a moat argument.

Please take the before/after the same way, so the two numbers are comparable:

1. Build with -p mistralrs-cli and assert the running server logs the SHA you
   built (git revision:), before driving anything.
2. ARC_PROFILE=1 ARC_PROFILE_SELFTEST=1 ARC_PROFILE_WARMUP=8 ARC_PROFILE_STEPS=60 ARC_PROFILE_LABEL=<cell> ARC_PROFILE_OUT=<dir>.
3. 🔴 Check the selftest line before quoting any device number. It must read
   ratio=NN.Nx — PASS. Until PR perf(profiler): make the ragged-admission branch measurable — three channels on the front-alignment path #113's device::timer_for fix, this always
   read NOT RUN (no CUDA event timer attached) — a NULL CUstream is CUDA's
   legacy default stream, not "no stream", and the guard rejected it, so
   every device and sync column ever produced on this fork was unmeasured.
   Also require unresolved_device_spans: 0 and violations: [].
4. Read mla_attn.device_ns ÷ decode_steps, not wall. The span is one
   device_span around self.attn.forward in forward_4d, aggregated over all
   43 layers (calls/step == 43 is the engagement check).
5. Run B=1, 8, 32 uniform in one process each and report the slope, not one
   batch: a kernel that helps at b=1 and not in the slope does not move the
   fleet number.
6. Profiler overhead measured at −0.14% aggregate, so no correction is
   needed — but take one no-profiler cell anyway and say so.

Cells, spans and the exact numbers above are in PR #113.

[heydryft]

Your name-collision finding is fixed at the interface — PR #120, master-based
and independent of my other work, so it cherry-picks like the timer fix did.

You were right to void the run, and right that the diagnosis didn't match: your
guard said unreached node, the truth was wrong node. The crate gave you no
safe alternative — node() took a full dotted path and nothing else existed —
so a name match was the natural thing to write. That's my interface's fault, not
your consumer's.

What changed:

• Node::branch — "prompt" / "decode", so the two copies are distinguishable
  by attribute instead of by parsing the path string.
• Profile::resolve("mla_attn") now ERRORS on an ambiguous name and lists
  every candidate with its branch and calls. Profile::resolve_in("decode", "mla_attn") is the call you want.
• Node::verdict() → Unreachable / NotEnteredThisWindow / BelowTimerFloor
  / Measured. You and one other chain hand-rolled that three-way distinction
  the same night; it's one call now.

Recommended guard, replacing the calls == 0 check:

let attn = profile.resolve_in("decode", "mla_attn")?;   // never the prefill copy
assert_eq!(attn.verdict(), Verdict::Measured);          // says WHICH zero if not
assert_eq!(attn.calls / decode_steps, 43);              // engagement: one per layer

I did not make the report prefer the decode node. Preferring one is still
guessing — it would just be wrong less often, and hand the next consumer with a
prefill question the same silent failure in the other direction.

Also in #120: the device-timer fix onto master. 66a61259d sits on a branch
based on unmerged work, so master still has the null-stream bug — if you
rebase onto master rather than carrying your cherry-pick, take #120 with it or
your device columns go dark again.

Good call holding the rebuild until your current run finishes. The two-binaries
confound is exactly what made my first ON/OFF pair unanswerable tonight.