Add packed paged KV formats with bounded workspace admission

[Gajesh2007]

Summary

Add optional packed full-attention KV storage with K4/V4, K8/V4 and K8/V8 formats, token-local affine groups, FP32 scale/offset metadata and a fixed signed Hadamard basis for Q/K. Native KV remains the default; sliding windows and recurrent state retain their existing representation. K4/V4 uses five effective bits per value including metadata, so its BF16 payload reduction is 3.2× before other memory costs.

Direct attention consumes packed pages without expanding the history. Physical admission separately prices growing storage, page/fixed costs and guaranteed step workspace. Step-owned arenas have real completion fences, including overlapping decode/MTP work; credits are refunded after GPU retirement. Resident reuse and complete checkpoints preserve packed bytes and reject incompatible storage identities before allocation.

An explicit optional prefill mode can reserve additional FP32 reconstruction and output memory for forced fused SDPA on eligible D64/D256 shapes. It falls back to the funded direct path when geometry or budget is unsuitable. Deadline-aware arrivals retire the actual optional step before forecasting remaining work against their original absolute deadline, including real drain time. Native/default direct forecasts are unchanged.

Before

flowchart LR
  R[Request] --> A[Native dtype based admission]
  A --> P[PagedLayerCache writes native K/V]
  P --> T[Native paged attention]
  T --> F[Finalize and release]
  P --> C[Native checkpoint tensors]

Loading

After

flowchart LR
  R[Request and explicit storage config] --> A[AdmissionV2: storage plus guaranteed workspace]
  A --> P[PagedLayerCache: quantize fresh K/V once]
  P --> D[Direct packed attention]
  P --> G{Eligible prefill and extra permit?}
  G -->|yes| S[Step arena plus forced fused SDPA]
  G -->|no| D
  D --> F[GPU fenced finalize then refund]
  S --> F
  N[Deadline arrival with optional step in flight] --> F
  F --> V[Forecast remaining work using original deadline]
  P --> C[Exact packed checkpoint with format identity]

Loading

Validation

On Apple M4 Max, the scoped baseline covered 126 unique native tests across numerical, checkpoint and admission runs. Subsequent overlapping suites include 40 optional-fused integration tests, 23 stride tests and 28 XCTest plus 11 Swift Testing deadline/ownership tests. The last group includes six new arrival cases covering near-full grants, partial prefills, expiration during actual retirement, cancellation and reused-ID ownership, and the unchanged direct path. These overlapping counts are not a new aggregate total.

Raw logs, exact model observations and performance limitations are retained by d-inference #860. Small authored quality probes do not establish broad non-inferiority. The current direct path is slower than native for measured 4K/16K prefill; no native-speed or serving-default claim is made.

Depends on mlx-swift #22, which pins mlx-c #9 and mlx #17 for the explicit fused API and FP32 wide-head tile correction. No merge or rollout is included.