Module predictor_ref 

2026-09-25: Host reference for the predictor kernels, with the kernels’ tensor layouts; tests/predictor_parity.rs compares the GPU results against it.

Owner: storage, high-speed swap. Invariants: none beyond the types.

Layouts:

• q: [num_q_heads, head_dim]; k_block: [block_size, num_kv_heads, head_dim];
• p: [head_dim, r]; q_proj: [num_q_heads, r];
• one block’s low-rank keys: [num_kv_heads, block_size, r], and a sequence of them [num_blocks, num_kv_heads, block_size, r];
• scores: [num_blocks], the max over query heads and tokens of the dot product.

Functions

ground_truth_block_weights

2026-09-25: Attention weight of each block: for every query head, the softmax over all tokens summed over the block’s tokens; the max over query heads.

predictor_score_ref

project_kv_block_ref

2026-09-25: Project every token of one block; output [num_kv_heads, block_size, r].

project_q_ref