feat(fast): add opt-in fused D256 SDPA

[Gajesh2007]

Summary

• add a default-off force_fused control to scaled dot-product attention
• instantiate normal Steel full attention for FP16/BF16 head_dim=256
• preserve default D256 routing to the faster composed path
• reject unsupported CPU, float32, shape, and transformed (vmap) requests rather than silently falling back
• keep D256 off the unqualified NAX path and validate pipeline threadgroup limits
• extend the D128 V-tile synchronization fence to all BD >= 128

This is a bounded-memory control, not an automatic speed route. M4 Max measurements show composed attention remains faster, so the default is unchanged.

Before

flowchart LR
  A[D256 full-attention request] --> B[automatic dispatch]
  B --> C[composed QK -> softmax -> PV]
  C --> D[score transient grows with qL x kL]
  E[caller needs bounded memory] --> C

Loading

After

flowchart LR
  A[D256 full-attention request] --> B{force_fused?}
  B -- no --> C[unchanged composed path]
  B -- yes --> D[validate GPU, dtype, shape, mask]
  D --> E[FP16/BF16 Steel D256]
  D -- unsupported --> F[explicit error]
  G[vmap of forced primitive] --> F

Loading

Code Flow

flowchart LR
  subgraph Before
    A1[fast.scaled_dot_product_attention] --> B1[use_fallback]
    B1 --> C1[full dims 64/80/128 only]
  end
  subgraph After
    A2[legacy overload or force-aware overload] --> B2[capability validation]
    B2 --> C2[primitive state includes force_fused]
    C2 --> D2[normal Steel D256 no-JIT/JIT symbol]
    C2 --> E2[vmap rejects forced execution]
  end

Loading

Verification

• no-JIT Metal build and metallib generation
• full fast-SDPA tests across FP16/BF16 masks, tails, sinks, and long K
• D256 float32 and CPU rejection controls
• positional C++ stream compatibility
• forced-vmap rejection
• benchmark matrix for default versus generic fused on M4 Max

Performance Note

The generic fused path reduces attention transient memory but is slower on M4 Max. It remains opt-in so runtimes can make the memory/throughput decision explicitly.