perf(aura): unify cache to activation dtype + wire 2-pass FA decode (closes -58%→-9% gap)

End-to-end FFAI side of the AURA dtype unification (metaltile sigs
0e4cb1a + 3fdadb3, PR 0xClandestine/metaltile#212). Replaces the
intermediate `.dequantMirror` default (originally bf16'd as a stopgap
because the single-pass `aura_flash_sdpa` kernel starved the GPU with
one simdgroup per query) with the right architecture: a single source
of truth in the activation dtype, and the token-parallel FA-2 kernel
pair.

## Cache schema — single source of truth (AURAQuantizedKVCache)
- kNorms / vNorms allocated in `dtype` (was f32-only).
- kCodebook / vCodebook allocated in `dtype` (was f32-only).
- kBoundaries / vBoundaries stay f32 — encoder-only, Lloyd-Max compare
  precision matters and they never reach the decode kernels.
- encodePerHead view stride now keys off `dtype.byteSize`, not a
  hardcoded 4 (the legacy f32 footgun that broke `AuraKLDIntegrationTests`
  the moment a non-f32 cache hit the encode path).

## Loaders (LlamaText / Qwen3Text)
- New `AURACodebook.centroidsTensor(dim:bits:dtype:device:)` host-side
  conversion helper covers all three float dtypes (f32 / f16 / bf16).
- `AURACodebook.boundariesTensor(...)` mirrors the helper for the
  encoder-only boundaries buffer.
- Both Qwen3 and Llama AURA cache builders use the helpers — no more
  copy-pasted f32 `Tensor.empty + copyIn` block per loader.

## Ops surface
- `Ops.auraFlashSdpa` preconditions drop the f32-norms-and-codebook
  requirement; everything must now match `out.dtype` (the activation
  dtype). Q pre-scale flow rewires from a f32 scratch + f32 scale buffer
  to an activation-dtype scratch + activation-dtype scale buffer.
- `AuraFlashScratchCache` keys both scratches on (count, dtype) — was
  keyed on `count` alone with f32 hardcoded. Adds a `partials(...)`
  scratch cache for the 2-pass partials triple.
- `Ops.auraEncode` + `Ops.auraDequantRotated` preconditions drop the
  f32-norms-and-codebook requirement; the dequant-mirror path flows
  through T now too.
- New `Ops.auraFlashSdpa2Pass` wrapper — dispatches `aura_flash_p1` +
  `aura_flash_pass2` for token-parallel FA-2 over the compressed cache.
  Caller-owned partials (mirrors `Ops.sdpaDecode2Pass`).
- New `Ops.supportsAuraFlashSdpa2Pass` predicate.

## Qwen3Layer.forward
- Prefer `Ops.auraFlashSdpa2Pass` when supported, fall back to
  `Ops.auraFlashSdpa` for combos the 2-pass kernel hasn't been emitted
  for (no path today; future-proof for kb!=4 / vb!=2,4 / d!=128).
- Block size 64 — matches the dense `sdpaDecode2Pass` per-block work
  size and saturates the M5 Max class around liveLength ≈ 4K.

## Default — back to `.compressed`
`LoadOptions.auraDecodePath` defaults to `.compressed`. Matches
@ekryski's stance from the PR review — true compressed attention is
FFAI's quantized-attention story and should be the default-path users
load into. The dtype unification + 2-pass FA-2 closes the perf gap that
made the original `.dequantMirror` flip necessary.

## Quality (M5 Max, Qwen3-0.6B-4bit, 61-position KLD harness)
| scheme                  | mean_kld | same_top |
|-------------------------|---------:|---------:|
| aura4v4 dequant-mirror  |     1.42 |      43% |
| aura4v4 2-pass flash    | **1.40** | **48%**  |
| aura4v2 2-pass flash    |     1.69 |      44% |
| aura8v4 (TQ+ recipe)    |    0.018 |      93% |

2-pass compressed flash matches (slightly beats) dequant-mirror on
aura4v4. KLD harness regression gate green for all schemes.

## Perf (M5 Max, Qwen3-0.6B-4bit decode tps, 5-run median)
| KV   | dequant-mirror | compressed (2-pass) | gap    | gap pre-unification |
|------|----------------|---------------------|--------|---------------------|
| 64   | 80.88          | 71.62               | -11.4% |              -15.7% |
| 256  | 77.14          | 67.71               | -12.2% |          **-43.7%** |
| 1024 | 46.87          | 42.73               |  -8.8% |          **-57.8%** |

Long-KV gap collapsed from -57.8% → -8.8%. Single-digit perf delta vs
dequant-mirror with 1.88× cache memory savings preserved (aura4v4 @
maxSeq=4096: 4352 KiB packed+norms vs 8192 KiB mirror).

## Why the C++ canonical pattern is safe
The fp16-stored norms / f32-at-use pattern this PR adopts mirrors the
production C++ `llama.cpp` TQ+ fork — commit b696c5da1 in that fork
shipped fp16 centroid LUTs + float-norm broadcast with measured zero
PPL impact ("Constant half LUT + float norm broadcast remains the
fastest approach on Apple Silicon", ggml-metal.metal:776). Internal
kernel arithmetic stays in f32 via cast-at-load; only the storage
narrows.

## Pass-2 dispatch shape note
`aura_flash_pass2`'s kernel header says `tg = (32, 1, 1) per q_idx`,
which means `q_idx = tgid_x`. Wrapper dispatches raw threads
`[nQHeads * 32, 1, 1]` with `tg = [32, 1, 1]` → `nQHeads` TGs along x,
each running 32 lanes; matches the metaltile end-to-end test's grid
shape exactly. The naive `[32, nQHeads, 1]` shape (raw-thread analogue
of `grid_groups [1, nQHeads, 1]`) would put `tgid_x = 0` for every TG,
i.e. every Q head's reduce reads q_idx=0's partials — produced
garbage same_top=0.0 / mean_kld=12+ output before the fix. Worth a
comment in the wrapper (added).

## Bench infra retained from the original perf pass
`AuraFlashScratchCache` (process-wide static, NSLock-guarded) memoizes
the Q scratch + scale buffer per (shape, dtype, scale) tuple. The
`AuraDecodeBenchIntegrationTests` side-by-side bench grid + memory
footprint asserter (KV=64 / 256 / 1024 + maxSeq=4096) are also kept
as regression catchers.

Author: TheTom

Sources/FFAI/KVCache/AURACodebook.swift

@@ -246,6 +246,54 @@ public enum AURACodebook {
         return base.map { $0 * scale }
     }
 
+    /// Allocate a codebook tensor in the requested activation dtype.
+    /// AURA cache stores codebook in the same dtype as the model
+    /// activations so both encode + decode kernels (which take
+    /// `Tensor<T>` for the codebook) read directly with no per-call
+    /// cast. The Lloyd-Max values themselves are computed in Float;
+    /// narrow dtypes (`bf16`/`f16`) round at the CPU-side host conversion.
+    public static func centroidsTensor(
+        dim: Int, bits: Int, dtype: DType, device: Device = .shared
+    ) -> Tensor {
+        let values = centroids(dim: dim, bits: bits)
+        return writeFloatsToTensor(values, shape: [values.count], dtype: dtype, device: device)
+    }
+
+    /// Allocate a boundaries tensor. Boundaries stay f32 — they're
+    /// encoder-only, used only by `aura_encode` for the branchless
+    /// Lloyd-Max comparison, where precision matters.
+    public static func boundariesTensor(
+        dim: Int, bits: Int, device: Device = .shared
+    ) -> Tensor {
+        let values = boundaries(dim: dim, bits: bits)
+        let t = Tensor.empty(shape: [values.count], dtype: .f32, device: device)
+        t.copyIn(from: values)
+        return t
+    }
+
+    /// CPU-side host conversion from `[Float]` into a tensor of the
+    /// requested float dtype. Used by `centroidsTensor` and any caller
+    /// that needs Lloyd-Max-precise values landed into narrow storage.
+    private static func writeFloatsToTensor(
+        _ values: [Float], shape: [Int],
+        dtype: DType, device: Device
+    ) -> Tensor {
+        let t = Tensor.empty(shape: shape, dtype: dtype, device: device)
+        switch dtype {
+        case .f32:
+            t.copyIn(from: values)
+        case .f16:
+            t.copyIn(from: values.map { Float16($0) })
+        case .bf16:
+            t.copyIn(from: values.map { UInt16(truncatingIfNeeded: $0.bitPattern >> 16) })
+        default:
+            fatalError(
+                "AURACodebook.centroidsTensor: unsupported dtype \(dtype); "
+                    + "AURA cache supports f32 / f16 / bf16")
+        }
+        return t
+    }
+
     /// Bytes-per-token after AURA packing at this bit width and dim.
     /// `ceil(dim * bits / 32) * 4` for the packed u32 array, plus 4
     /// bytes for the f32 per-token norm. Excludes any per-vector DC

Sources/FFAI/KVCache/AURAQuantizedKVCache.swift

@@ -107,16 +107,21 @@ public final class AURAQuantizedKVCache: KVCacheProtocol, @unchecked Sendable {
     /// Π^T in the activation dtype, used to un-rotate the SDPA output
     /// before `oProj`. Aliases `rotationT` when `dtype == .f32`.
     public let rotationDtypeT: Tensor
-    public let kCodebook: Tensor  // [2^keyBits]   f32
-    public let kBoundaries: Tensor  // [2^keyBits-1] f32
-    public let vCodebook: Tensor  // [2^valueBits] f32
-    public let vBoundaries: Tensor  // [2^valueBits-1] f32
+    /// Codebook in the cache dtype. Encode + decode kernels read
+    /// directly with no per-call cast — the dtype unification landed
+    /// when the single-pass `aura_flash_sdpa` kernel was migrated to
+    /// `Tensor<T>` (matches the production C++ TQ+ fork pattern: fp16-
+    /// stored norms / codebook, f32-at-use via cast-at-load).
+    public let kCodebook: Tensor  // [2^keyBits]   dtype
+    public let kBoundaries: Tensor  // [2^keyBits-1] f32 — encoder-only Lloyd-Max thresholds
+    public let vCodebook: Tensor  // [2^valueBits] dtype
+    public let vBoundaries: Tensor  // [2^valueBits-1] f32 — encoder-only
 
     // Per-cache compressed storage.
     public let kPacked: Tensor  // [nKVHeads, maxSeq, kPackedWidth] u32
     public let vPacked: Tensor  // [nKVHeads, maxSeq, vPackedWidth] u32
-    public let kNorms: Tensor  // [nKVHeads, maxSeq] f32
-    public let vNorms: Tensor  // [nKVHeads, maxSeq] f32
+    public let kNorms: Tensor  // [nKVHeads, maxSeq] dtype — encode writes T, decode reads T
+    public let vNorms: Tensor  // [nKVHeads, maxSeq] dtype
 
     // Shared working buffers — bulk-dequant target; reused across layers.
     public let sharedWorkingK: Tensor  // [nKVHeads, maxSeq, headDim] dtype
@@ -192,11 +197,17 @@ public final class AURAQuantizedKVCache: KVCacheProtocol, @unchecked Sendable {
             "AURAQuantizedKVCache: rotationDtype/rotationDtypeT dtype must match cache dtype \(dtype)"
         )
         precondition(
-            kCodebook.dtype == .f32 && kBoundaries.dtype == .f32,
-            "AURAQuantizedKVCache: K codebook/boundaries must be f32")
+            kCodebook.dtype == dtype,
+            "AURAQuantizedKVCache: K codebook dtype must match cache dtype \(dtype)")
         precondition(
-            vCodebook.dtype == .f32 && vBoundaries.dtype == .f32,
-            "AURAQuantizedKVCache: V codebook/boundaries must be f32")
+            kBoundaries.dtype == .f32,
+            "AURAQuantizedKVCache: K boundaries must be f32 (encoder-only)")
+        precondition(
+            vCodebook.dtype == dtype,
+            "AURAQuantizedKVCache: V codebook dtype must match cache dtype \(dtype)")
+        precondition(
+            vBoundaries.dtype == .f32,
+            "AURAQuantizedKVCache: V boundaries must be f32 (encoder-only)")
         precondition(
             sharedWorkingK.shape == [nKVHeads, maxSeq, headDim],
             "AURAQuantizedKVCache: sharedWorkingK shape mismatch")
@@ -232,9 +243,9 @@ public final class AURAQuantizedKVCache: KVCacheProtocol, @unchecked Sendable {
         self.vPacked = Tensor.empty(
             shape: [nKVHeads, maxSeq, vPackedWidth], dtype: .u32, device: device)
         self.kNorms = Tensor.empty(
-            shape: [nKVHeads, maxSeq], dtype: .f32, device: device)
+            shape: [nKVHeads, maxSeq], dtype: dtype, device: device)
         self.vNorms = Tensor.empty(
-            shape: [nKVHeads, maxSeq], dtype: .f32, device: device)
+            shape: [nKVHeads, maxSeq], dtype: dtype, device: device)
 
         // Codec is purely additive in atomic_or terms, so packed slots
         // MUST start zeroed. Norms slots get overwritten per encode but
@@ -377,7 +388,10 @@ public final class AURAQuantizedKVCache: KVCacheProtocol, @unchecked Sendable {
         let inputBytesPerHead = headDim * dtype.byteSize
         let packedBytesPerSlot = packedWidth * 4  // u32
         let packedBytesPerHead = maxSeq * packedBytesPerSlot
-        let normBytesPerHead = maxSeq * 4  // f32
+        // Norms are stored in the cache dtype post-unification — stride
+        // tracks the activation dtype's byte size, not the legacy 4 (f32).
+        let normByteSize = dtype.byteSize
+        let normBytesPerHead = maxSeq * normByteSize
 
         for h in 0 ..< nKVHeads {
             let inputView = Tensor(
@@ -390,8 +404,8 @@ public final class AURAQuantizedKVCache: KVCacheProtocol, @unchecked Sendable {
                 shape: [1, packedWidth], dtype: .u32)
             let normsView = Tensor(
                 buffer: norms.buffer,
-                offset: norms.offset + h * normBytesPerHead + pos * 4,
-                shape: [1], dtype: .f32)
+                offset: norms.offset + h * normBytesPerHead + pos * normByteSize,
+                shape: [1], dtype: dtype)
             Ops.auraEncode(
                 input: inputView, rotation: rotation,
                 boundaries: boundaries, codebook: codebook,

Sources/FFAI/Loader/LoadOptions.swift

@@ -57,25 +57,31 @@ public enum DispatchMode: Sendable {
 /// Only relevant when `LoadOptions.kvCache == .auraQuantized(...)` —
 /// raw / affine caches ignore this setting.
 public enum AURADecodePath: Sendable, Equatable {
-    /// **Default.** Compressed-domain attention via the
-    /// `aura_flash_p1` + `aura_flash_pass2` kernel pair. Q is rotated,
+    /// **Default.** Compressed-domain attention via the 2-pass FA-2
+    /// kernel pair (`aura_flash_p1` + `aura_flash_pass2`) when emitted
+    /// for the (keyBits, valueBits, headDim, dtype) combo, with the
+    /// single-pass `aura_flash_sdpa` as fallback for cells the 2-pass
+    /// kernel hasn't been emitted for. Q is rotated + pre-scaled,
     /// scored directly against the packed K codes (no full-precision
-    /// dequant), then combined with the packed V codes — the kernel
-    /// dequantises per-tile on chip, never materialising a maxSeq-sized
-    /// f16 mirror buffer. Realises AURA's full memory savings (~4× at
-    /// `aura4v2`).
+    /// dequant), and the V codes are dequanted per-tile on chip. The
+    /// `[nKVHeads, maxSeq, headDim]` mirror buffer never materialises,
+    /// realising AURA's memory savings (~1.88× at aura4v4, ~3.7× at
+    /// aura4v2 on Qwen3 d=128).
+    ///
+    /// AURA-dtype unification (metaltile + FFAI joint change) put the
+    /// per-token norms and per-scheme codebook into the activation
+    /// dtype, so encode + both decode kernel paths consume the cache
+    /// buffers directly — no per-call f32 cast on the decode hot path
+    /// and no parallel f32 mirror storage.
     case compressed
 
-    /// Stage 1a behaviour. `prepareForAttention(on:)` dequantises the
-    /// full compressed K/V cache into per-layer shared working buffers
-    /// (`sharedWorkingK` / `sharedWorkingV`, sized
-    /// `[nKVHeads, maxSeq, headDim]`), and the standard
-    /// `Ops.sdpaDecode` reads those. Preserves AURA's quality but
-    /// **gives back the memory savings** — the mirror is the same size
-    /// as a raw fp16 cache. Kept as an opt-in path for A/B benching
-    /// (`compressed` vs `dequantMirror` speed at production shapes)
-    /// and for callers with the memory headroom who want
-    /// matrix-engine SDPA.
+    /// Dequant-mirror path. `prepareForAttention(on:)` materialises
+    /// the full compressed K/V cache into per-layer shared working
+    /// buffers (`sharedWorkingK` / `sharedWorkingV`, sized
+    /// `[nKVHeads, maxSeq, headDim]`) and `Ops.sdpaDecode` reads those.
+    /// Same quality as `.compressed`, **gives back the memory
+    /// savings** — the mirror is the same size as a raw fp16 cache.
+    /// Useful as an A/B baseline against the compressed path.
     case dequantMirror
 }
 
@@ -113,13 +119,13 @@ public struct LoadOptions: Sendable {
     /// entire advertised window, or a smaller value to bound memory.
     public var maxContextLength: Int?
 
-    /// Selects the AURA decode path. Defaults to `.compressed` (Stage
-    /// 1b: attend on packed K/V codes directly via the `aura_flash_*`
-    /// kernel pair — full ~4× memory savings). Set to `.dequantMirror`
-    /// for the Stage 1a path that maintains a full-precision
-    /// `[nKVHeads, maxSeq, headDim]` mirror buffer and runs the
-    /// standard `Ops.sdpaDecode` against it — useful for A/B speed
-    /// benching. Has no effect when `kvCache != .auraQuantized(...)`.
+    /// Selects the AURA decode path. Defaults to `.compressed` — the
+    /// 2-pass FA-2 kernel pair gives token-parallel attention over the
+    /// packed K/V codes directly, with no f16/f32 mirror materialised.
+    /// Set to `.dequantMirror` for an A/B baseline that dequants the
+    /// cache into a per-layer working buffer and runs the standard
+    /// `Ops.sdpaDecode` against it. Has no effect when
+    /// `kvCache != .auraQuantized(...)`.
     public var auraDecodePath: AURADecodePath
 
     public init(

Sources/FFAI/Models/Text/LlamaText.swift

@@ -567,21 +567,18 @@ public final class LlamaModel: LanguageModel {
             // Codebooks are shared across layers; rotations are per-layer
             // (deterministic SRHT seeded by layer index). See Qwen3's
             // matching case for the longer explanation.
-            let kCodebookData = AURACodebook.centroids(dim: headDim, bits: scheme.keyBits)
-            let kBoundariesData = AURACodebook.boundaries(dim: headDim, bits: scheme.keyBits)
-            let vCodebookData = AURACodebook.centroids(dim: headDim, bits: scheme.valueBits)
-            let vBoundariesData = AURACodebook.boundaries(dim: headDim, bits: scheme.valueBits)
-
-            let kCodebook = Tensor.empty(shape: [kCodebookData.count], dtype: .f32, device: device)
-            kCodebook.copyIn(from: kCodebookData)
-            let kBoundaries = Tensor.empty(
-                shape: [kBoundariesData.count], dtype: .f32, device: device)
-            kBoundaries.copyIn(from: kBoundariesData)
-            let vCodebook = Tensor.empty(shape: [vCodebookData.count], dtype: .f32, device: device)
-            vCodebook.copyIn(from: vCodebookData)
-            let vBoundaries = Tensor.empty(
-                shape: [vBoundariesData.count], dtype: .f32, device: device)
-            vBoundaries.copyIn(from: vBoundariesData)
+            // Codebook in cache dtype (matches encode/decode kernel
+            // signatures — no per-call cast). Boundaries stay f32:
+            // encoder-only and precision-sensitive at the Lloyd-Max
+            // comparison.
+            let kCodebook = AURACodebook.centroidsTensor(
+                dim: headDim, bits: scheme.keyBits, dtype: dtype, device: device)
+            let kBoundaries = AURACodebook.boundariesTensor(
+                dim: headDim, bits: scheme.keyBits, device: device)
+            let vCodebook = AURACodebook.centroidsTensor(
+                dim: headDim, bits: scheme.valueBits, dtype: dtype, device: device)
+            let vBoundaries = AURACodebook.boundariesTensor(
+                dim: headDim, bits: scheme.valueBits, device: device)
 
             let sharedK = Tensor.empty(
                 shape: [nKVHeads, cap, headDim],

Sources/FFAI/Models/Text/Qwen3Text.swift

@@ -316,8 +316,13 @@ public final class Qwen3Layer: Module {
 
         // Decode path selection:
         //   • `.compressed` (default) — score Q directly against the
-        //     packed K codes via `aura_flash_sdpa`. V is dequanted per
-        //     tile on chip; no maxSeq-sized mirror is materialised.
+        //     packed K codes. V is dequanted per tile on chip; no
+        //     maxSeq-sized mirror is materialised. Two kernel variants:
+        //       – `auraFlashSdpa2Pass` when supported (token-parallel
+        //         FA-2; one TG per (q_head, block) — saturates the GPU
+        //         at long context).
+        //       – `auraFlashSdpa` (single-pass, one TG per q_head)
+        //         fallback for combos the 2-pass kernel hasn't emitted.
         //     `kvStride = maxSeq` (NOT `length`) — the per-head row
         //     stride is the allocated cache stride, not the live row
         //     count. Passing `length` aliases reads across head
@@ -328,6 +333,44 @@ public final class Qwen3Layer: Module {
         // Non-AURA caches always take the dequantMirror branch.
         let attnOut: Tensor
         if let auraCache = cache as? AURAQuantizedKVCache,
+            auraCache.decodePath == .compressed,
+            Ops.supportsAuraFlashSdpa2Pass(
+                keyBits: auraCache.scheme.keyBits,
+                valueBits: auraCache.scheme.valueBits,
+                headDim: headDim, dtype: h.dtype)
+        {
+            // FA-2 block tile. 64 is the canonical choice — matches
+            // the dense `sdpaDecode2Pass` per-block work size and gives
+            // ~16 q-heads × ceil(maxSeq/64) blocks of token-parallelism
+            // (saturates the M5 Max class around liveLength ≈ 4K).
+            let blockSize = 64
+            let maxBlocks = (auraCache.maxSeq + blockSize - 1) / blockSize
+            let partials = AuraFlashScratchCache.partials(
+                nQHeads: nHeads, maxBlocks: maxBlocks,
+                headDim: headDim, dtype: h.dtype)
+            let outTensor = Tensor.empty(
+                shape: [nHeads, headDim], dtype: h.dtype, device: device)
+            Ops.auraFlashSdpa2Pass(
+                q: qForSdpa,
+                kPacked: auraCache.kPacked, kNorms: auraCache.kNorms,
+                kCodebook: auraCache.kCodebook,
+                vPacked: auraCache.vPacked, vNorms: auraCache.vNorms,
+                vCodebook: auraCache.vCodebook,
+                into: outTensor,
+                nQHeads: nHeads, nKVHeads: nKVHeads, headDim: headDim,
+                kPackedWidth: auraCache.kPackedWidth,
+                vPackedWidth: auraCache.vPackedWidth,
+                liveLength: auraCache.length, kvStride: auraCache.maxSeq,
+                keyBits: auraCache.scheme.keyBits,
+                valueBits: auraCache.scheme.valueBits,
+                scale: scale,
+                blockSize: blockSize,
+                partialO: partials.partialO,
+                partialM: partials.partialM,
+                partialL: partials.partialL,
+                on: cmd)
+            attnOut = outTensor
+        } else if let auraCache = cache as? AURAQuantizedKVCache,
             auraCache.decodePath == .compressed,
             Ops.supportsAuraFlashSdpa(
                 keyBits: auraCache.scheme.keyBits,
@@ -338,9 +381,9 @@ public final class Qwen3Layer: Module {
                 shape: [nHeads, headDim], dtype: h.dtype, device: device)
             // Qwen3 dense has no attention sinks — the kernel gates the
             // sinks load on `hasSinks=false`, but the wrapper still
-            // requires an f32 buffer to satisfy the dtype precondition.
+            // requires a same-dtype buffer to satisfy the precondition.
             let sinksScratch = Tensor.empty(
-                shape: [nHeads], dtype: .f32, device: device)
+                shape: [nHeads], dtype: h.dtype, device: device)
             Ops.auraFlashSdpa(
                 q: qForSdpa, sinks: sinksScratch,
                 kPacked: auraCache.kPacked, kNorms: auraCache.kNorms,
@@ -645,21 +688,17 @@ public final class Qwen3Model: LanguageModel {
             // are per-layer: each Π_l is an SRHT matrix seeded by the
             // layer index, matching the AURA paper's "fresh rotation per
             // tensor" recipe for de-correlating activation statistics.
-            let kCodebookData = AURACodebook.centroids(dim: headDim, bits: scheme.keyBits)
-            let kBoundariesData = AURACodebook.boundaries(dim: headDim, bits: scheme.keyBits)
-            let vCodebookData = AURACodebook.centroids(dim: headDim, bits: scheme.valueBits)
-            let vBoundariesData = AURACodebook.boundaries(dim: headDim, bits: scheme.valueBits)
-
-            let kCodebook = Tensor.empty(shape: [kCodebookData.count], dtype: .f32, device: device)
-            kCodebook.copyIn(from: kCodebookData)
-            let kBoundaries = Tensor.empty(
-                shape: [kBoundariesData.count], dtype: .f32, device: device)
-            kBoundaries.copyIn(from: kBoundariesData)
-            let vCodebook = Tensor.empty(shape: [vCodebookData.count], dtype: .f32, device: device)
-            vCodebook.copyIn(from: vCodebookData)
-            let vBoundaries = Tensor.empty(
-                shape: [vBoundariesData.count], dtype: .f32, device: device)
-            vBoundaries.copyIn(from: vBoundariesData)
+            // Codebook in cache dtype (matches encode/decode kernel
+            // signatures — no per-call cast). Boundaries stay f32:
+            // encoder-only and precision-sensitive at Lloyd-Max compare.
+            let kCodebook = AURACodebook.centroidsTensor(
+                dim: headDim, bits: scheme.keyBits, dtype: dtype, device: device)
+            let kBoundaries = AURACodebook.boundariesTensor(
+                dim: headDim, bits: scheme.keyBits, device: device)
+            let vCodebook = AURACodebook.centroidsTensor(
+                dim: headDim, bits: scheme.valueBits, dtype: dtype, device: device)
+            let vBoundaries = AURACodebook.boundariesTensor(
+                dim: headDim, bits: scheme.valueBits, device: device)
 
             // Shared working buffers — same pattern as affineQuantized:
             // bulk-dequant target shared across all layers.