Evaluate the model's perplexity on the test set using batch processing. It is expected that model is already on the correct device.

Per-group PCA via eigh with 1% Tikhonov diagonal damping. Returns (G, full_dim, retained_dim).

Rotates half the hidden dims of the input.

r""" TODO let's just use the original freqcis computation to not have the view transpose + reshape! This is not optimized! Applies Rotary

Evaluate the model's perplexity on the test set using batch processing. It is expected that model is already on the correct device.

Take the input signals ("activations") for a layer, run the layer forward.

Run PCA on a list of batched data. Returns the eigenvalues and eigenvectors.

Applies Rotary Position Embedding to the query and key tensors. Args: q (`torch.Tensor`): The query tensor. k (`torch.Tensor`): T

Mean per-sequence next-token NLL, exp'd; logits chunked along seq axis for memory.

Capture each layer's attention input (CPU bf16) via forward pre-hooks. Independent mode. If ``capture_final=True``, also captures the OUTPUT of t

Get the dataset from the HuggingFace datasets library. Args: name: The name of the HuggingFace dataset to load. Must be one of "wiki

Load wikitext2 / alpaca / pg19 as HF datasets, normalised to a single 'text' column.

Recursively map tensors to device and dtype.

Recursively map tensors to device and dtype.

Get a DataLoader from a test dataset. This dataloader should be used when comparing WikiText2 perplexities with other papers, e.g. SparseGPT (arx

Comes from AiHarness

This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, num_key_value_heads, seqlen, he

Rotates half the hidden dims of the input.

Sync all GPUs to make sure all operations are finished, needed for correct benchmarking of latency/throughput.

Sync all GPUs to make sure all operations are finished, needed for correct benchmarking of latency/throughput.

Expand a GQLA-packed kv_b_proj weight to MLA layout. Source: ``(num_kv_heads * per_head_kv, kv_lora_rank)`` — one (W_UK|W_UV) block per KV gr

Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token first unpad the input, th

Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape `(batch_size, key_value_l

Slice per-group ``(G, gs*d, gs*d)`` covs from the all-pair ``(H*d, H*d)`` cov. For group ``g``, the in-group head order is ``perm[g*gs:(g+1)*gs]`

Run GC and clear GPU memory.

All-pair K/V covariance ``(H * d, H * d)``, fp64. Drives data-driven head grouping. Same kv_a chain as ``collect_kv_grams``; differs only in that

Stream the kv_a chain over calibration; accumulate per-group K/V covariance (fp64). If ``token_weights`` is provided (one ``(S,)`` fp32 tensor pe

Compress + absorb for arbitrary head groupings (e.g. similarity-driven). For each new GQLA head ``h_new = g*gs + i`` with original head ``ori

Build compressed kv_b_proj + absorbed q_b_proj (nope rows) and o_proj (per-head v cols).

Nuclear-norm-of-cross-cov head similarity, weighted K + V contributions. For each pair ``(h, h')``, ``||Sigma^{h,h'}||_*`` is the negated optimal

Per-token NLL weights from the cached last-layer hidden states. Standard per-layer PCA accumulates ``Sigma_X = X.T @ X`` with uniform weights,

Quick stats for logging Hessian-PCA token weights.

Take the input signals ("activations") for a layer, run the layer forward.

Get the dataset from the HuggingFace datasets library. Args: name: The name of the HuggingFace dataset to load. Must be one of "wiki

Seed-and-grow balanced grouping: pick the best unassigned pair to seed each group, then iteratively add the unassigned head with highest sum-simil

Replace kv_b_proj with the shrunken Linear; copy absorbed q/o weights; flip kv_groups.