metal: accelerate M5 Max indexed prefill

[rinaldofesta]

M5 Max indexed-prefill attention dropped from 88.619 to 78.608 ms/layer across 21 indexed layers. This PR also replaces the old streaming top-512 path with an exact fused two-run kernel for partial M5 prefill widths.

Changes

• Stage 16 raw or compressed K/V rows per threadgroup in the existing dual-head indexed-attention kernel. The row order and online softmax arithmetic stay unchanged.
• Enable the rb16 specialization only on M5 Max, for the existing 64-head, top-512 prefill shape with at least 32 tokens.
• Fuse the two sorted runs used by top-512 when the compressed width is 1025 through 2048 on M5 devices.
• Keep rollback switches for both paths:
  • DS4_METAL_DISABLE_INDEXED_ATTN_DUAL_RB16=1
  • DS4_METAL_DISABLE_TOPK_FUSED512=1

Performance

Attention stage profile at 6144 tokens, 21 indexed layers:

	Baseline	Candidate	Delta
Mean	88.619 ms	78.608 ms	-11.30%
Median	88.249 ms	78.931 ms	-10.56%
Layer wins		21/21

8192-token prefill, counterbalanced pairs:

Order	Baseline	Candidate	Delta
Baseline, candidate	513.34 tok/s	542.67 tok/s	+5.71%
Candidate, baseline	551.00 tok/s	552.95 tok/s	+0.35%
Two-run mean	532.17 tok/s	547.81 tok/s	+2.94%

Top-512 A/B at width 2048 x 6144 rows: 4.470 to 4.209 ms median, -5.8%.

The 6144-token full-prefill runs had large thermal variation, from 503.9 to 722.6 tok/s. The candidate won both no-dump orders, but I am not using that spread for a rate claim.

Exactness and tests

make -j8 ds4_test ds4-bench
./ds4_test --metal-kernels

• Fused top-512: 0/270336 mismatches, widths 1537 and 2048, unique and heavy-tie inputs, four repeats.
• Dual rb16 attention: 0/2162688 float mismatches against the row-1 kernel, two repeats.
• Full 6144-token logits: byte-identical with and without both optimizations.
• SHA-256: ef2c428e514839cf348ab445465770afe7d260197b15593cee302ea978fd77f9.
• git diff --check: clean.

Scope

Aggregate make test status remains red on this M5 Max: five fixture assertions reproduce byte-identically at clean HEAD and with both optimized paths disabled. They are the existing short_code_completion step-0 token/logprob cases and long_story_4096 top-5 overlap case.

SSD streaming, distributed inference, and CUDA were outside this Metal-only validation run.