APPENDIX A

Compression and indexing

Chapter 3 followed a sliding window into attention; chapter 4 allocated blocks without explaining their compressed rows. Here we follow the other read path: completed compressed KV, a smaller index used to select rows, and the temporary state needed to produce the next row.

python3 labs/26_compression.py

C2 rows (before BF16/norm/RoPE): [('6', '8'), ('8', '3'), ('3', '5')]
Chunks [0,1), [1,4), [4,7): valid/reserved rows [(0, 1), (2, 2), (1, 1)]
Pending token positions: [6]
Six-position decode: ring 7 preserves every completed pair
Ring 6 corrupts first pair: ('42/5', '8')
Checks: 64 chunk partitions agree exactly; undersized ring fails
Default C2 state payload: 7 * 2 * 512 * 4 = 28672 bytes

Seven inputs produce three completed pairs and one pending input. Splitting the inputs into chunks of lengths 1, 3 and 3 does not change those pairs. A six-entry ring, however, lets the last new write overwrite the first pair’s old input. Seven entries preserve it. Compression reduces the completed cache; it still needs live state between steps.

Scope and order. We first explain the gated mean and its boundaries, then follow V4Attention through the compressor and indexer to the final attention read. A second experiment separates scoring, candidate restriction and address conversion. Finally, we account for bytes and distinguish the objects shared across layers. The case is a sparse C2 layer that owns its main KV and index-K, without YOCO compaction or CP sharding; later sections identify where those other paths differ.

A.1 A pair is complete only at its boundary

There are two kinds of compression in this discussion. MLA compresses a feature vector, as chapter 3 explained. C2 combines two token positions into one row. It does not halve the number of channels in that row. C1 keeps one row per position. The released Compressor accepts only ratios 1 and 2 and rejects the optional compress_overlap and use_ape settings [87].

For latent component $d$, the C2 operation over positions $2j,2j+1$ is

The score $s$ is a learned gate; $v$ is the projected latent. Each channel has its own weights over the two positions. This is not attention over the whole context, nor an unweighted average. In v4_modules/mla_attn.py, SoftmaxAPEGatedColumnParallelLora constructs an A projection for values and a Gate projection for scores. Its name retains APE, although the runtime rejects that option. The single-position C1 softmax is one, so C1 needs no Gate [87].

The first lab pair has values $(0,10)$ and $(8,2)$. It substitutes positive integer weights $(1,3)$ and $(3,1)$ for the exponentials. The result is $(6,8)$: the second token dominates channel zero, while the first dominates channel one. Fractions let us check this mathematical operation exactly. They do not reproduce floating-point exponentials or BF16 rounding.

The lab implements that mean with rational weights:

18def compress(rows):
19    return tuple(sum(F(w[d]) * v[d] for w, v in rows)
20                 / sum(w[d] for w, _ in rows) for d in range(2))

The source reference reduces over positions, then casts the result:

199result = (latent_slice * score_slice.softmax(dim=0)).sum(dim=0).to(torch.bfloat16)
200kv_compressed[compressed_start + blk_id - qstart // compress_ratio] = result

The production wrapper makes the C1/C2 distinction explicit:

89if self.kv_compress_ratio == 1:
90    assert projected_kv.dtype == torch.bfloat16 and state_cache is None
91    if input_metadata.mirror_cp_reverse_index is not None:
92        projected_kv = gather(projected_kv.contiguous(), input_metadata.mirror_cp_reverse_index)
93    return projected_kv
94
95assert projected_kv.dtype == torch.float32 and state_cache is not None
96score, latent = projected_kv.chunk(2, dim=-1)

C1 receives BF16 and no compression state. Its optional gather restores CP’s input ordering. C2 receives FP32 score/latent pairs and requires state. The following branch calls the prefill or decode helper with the same ratio. These helpers return BF16 rows before the compressor’s normalisation and RoPE.

A.1.1 Chunk boundaries and placeholder rows

Let $s$ be the first new token position and $e$ the exclusive end of a chunk. For ratio $r$, the number of newly completed rows is

The prefill reference iterates exactly those compressed indices. If a pair began before $s$, it reads that input from the state ring; otherwise it reads from the current projection. It subsequently retains recent score/latent pairs in the ring for later chunks [88].

For $[0,1)$, no pair is complete. For $[1,4)$, pairs $(0,1)$ and $(2,3)$ finish: position zero must survive the first chunk. For $[4,7)$, only $(4,5)$ finishes. Position six remains pending. Lab 26 tries every choice of cuts between the seven inputs and compares the completed rows and pending suffix with processing the whole sequence.

Physical output shape is not the valid-row count. The prefill metadata reserves $\text{max}(1,n_c)$ rows per request, even when $n_c=0$. Its placeholder has position zero and store slot $-1$; the reference compression output is zero-filled there. Decode instead reserves one output row per query position, with zero output and no valid store slot for a position that does not complete a group. The C++ builder implements these conventions in prepare_tensor_input_v4.cpp:46–91. A placeholder is not an extra cached token [89].

A.2 Temporary state must survive speculative writes

A C2 cache row cannot replace the score and latent of an incomplete pair. The live ca_state stores both, in FP32, with shape

The axis of length two is score versus latent, not K versus V. ca_cache holds completed compressed KV in paged blocks. Their lifetimes and byte formats differ.

The decode reference writes every new position before computing completed rows. Its assertion states the required capacity:

223compress_token = compress_ratio
224state_cache_size = state_cache.shape[1]
225assert state_cache.shape[1] >= compress_ratio + max_seqlen_q - 1
226
227# store: each token writes to cache
228for idx in range(total_q):
229    sidx = state_block_idx[idx].item()
230    cpos = (context_lens[idx].item() - 1) % state_cache_size
231    state_cache[sidx, cpos, 0, ...] = score[idx, ...]
232    state_cache[sidx, cpos, 1, ...] = latent[idx, ...]
233
234# compute: check each token's position
235kv_compressed = torch.zeros(total_q, dim, dtype=torch.bfloat16, device=score.device)
236for idx in range(total_q):
237    abs_pos = context_lens[idx].item() - 1
238    if abs_pos % compress_ratio == compress_ratio - 1:

For a group of $r$ positions and up to $q_{\text{max}}$ new positions per state slot, retaining $r-1$ old inputs alongside all new inputs requires $S\ge r+q_{\text{max}}-1$. This is the same write-before-read hazard encountered with SWA, applied to a different ring.

In the lab, position zero is old; positions 1 through 6 are written in one step. With $S=6$, position six takes position zero’s slot before pair $(0,1)$ is read. With $S=7$, no required input is overwritten. The incorrect first component becomes $42∕5$, rather than 6. No future position belongs in that pair, even though it has already been written to memory.

The cache engine derives its ring from the backtracking allowance, and separately decides whether an indexer needs its own state:

123def _needs_index_state(attn_module: AttrDict) -> bool:
124    if attn_module.cfg.get("index_topk", 0) <= 0:
125        return False
126    index_cfg = attn_module.cfg.get("attn_index")
127    return index_cfg is None or not index_cfg.get("use_kproj", False)
128
129
130@lru_cache(8)
131def get_inference_compress_block_size(compress_ratio: int) -> int:
132    """Compression ring plus rollback history; environment changes require clearing the cache."""
133    return compress_ratio + max(0, envs.DEEPINFER_MAX_BACKTRACKING_TOKENS - 1)

With the default nextn=5, the backtracking allowance is at least six. For C2 the default ring is $2+(6-1)=7$, not the SWA ring’s 134. At dimension 512 its score/latent payload is $7\times 2\times 512\times 4=28,672$ bytes per owning layer per state slot. That excludes other state regions and any layout padding. Rejecting a draft changes which positions are valid; it does not require clearing every byte, as chapter 4 established.

A.3 One compressed result feeds two caches

The main attention module connects ownership to computation:

515assert wait_projected_kv is not None
516assert self.PackedWkv is not None
517return self.PackedWkv.forward_store_kv_cache(
518    wait_projected_kv(),
519    kv_cache.ca_cache,
520    kv_cache.ca_state,
521    input_metadata,
522    return_pre_rope_kv=self.has_local_index,
523)

The result forks after normalisation and before rotary position encoding. The main branch applies RoPE, casts to the main cache format and writes ca_cache. A local indexer also needs the normalised pre-RoPE latent:

158    cast_args=CastArgs.from_cache_format(cache_format),
159    out_args=KVCacheOutArgs(
160        kv_cache,
161        input_metadata.slot_mapping_store_ca[self.ratio_key],
162        out=pre_rope_kv,
163        out_stage="norm",
164        layout=cache_format,
165    ),
166)
167return pre_rope_kv[:, 0] if pre_rope_kv is not None else None

out_stage="norm" identifies the retained stage. The surrounding call supplies the compressed positions and RoPE parameters. For compressed row $j$, the C++ metadata uses position $\mathit{jr}$, the beginning of its group, while a query keeps its original token position. Equal row numbers therefore do not imply equal positional coordinates [89]. Figure A.1 shows where the two caches separate.

The index branch projects that latent to a smaller key. Its store method requires that no independent index compression state be present:

140if not self.store_index_k:
141    return
142assert index_k_input is not None
143assert state_cache is None
144index_k = self.index_k_module.forward_pre_norm(index_k_input)
145norm = self.index_k_module.norm
146fused_norm_rope_hadamard_cast_store_kvcache(
147    index_k,
148    norm_args=NormArgs(norm.weight, norm.variance_epsilon),

The model configuration selects KProject with attn_index.use_kproj=True. The cache engine therefore leaves ica_state=None on this V4.1 path. The field remains in the generic interface for a separately compressed index path. It is not a second accumulator that every sparse layer allocates. ica_cache still exists: absence of live index state does not mean absence of completed index keys [90].

This is a model-driven requirement. A block allocator alone cannot decide that the index keys derive from the main compressed latent, that C1 needs no pending pair, or that a consumer layer reuses another layer’s cache. Those decisions come from the model configuration and its runtime adapters.

A.4 Select rows, then translate their addresses

An indexer makes a cheaper scoring pass before main attention. The runtime describes its score for key $k_j$ as

These are learned index queries and head weights, not the main attention softmax probabilities. ReLU means $\text{max}(0,x)$. DeepGEMM computes the logits; DeepSelect supplies the imported topk. Their kernel internals remain outside this walkthrough [90].

python3 labs/27_indexing.py

Scores: [2, 6, 3, 0, 300] | valid prefix: 4
Top-2 over valid rows: [1, 2]
Candidate block scores: [6, 3] | retained block: 0
Top-2 within retained block: [0, 1]
One valid row, three output slots: [0, -1, -1]
Prefill at gathered offset 10: [11, 12]
Decode logical rows: [1, 2, 4, -1] -> physical slots: [29, 30, 8, -1]

The lab uses two query heads, $(1,0)$ and $(0,1)$, with weights 1 and 2. Key $(0,3)$ scores 6, while $(1,1)$ scores 3. The fifth key would score 300, but it is outside the valid prefix. The chosen rows are 1 and 2, not 4. The calculations omit projection, RoPE and quantisation and use a stated index-based tie rule; they do not reproduce the device top-k implementation.

The teaching score uses ordinary integer arithmetic:

16def score(key):
17    return sum(w * max(0, sum(a * b for a, b in zip(q, key)))
18               for w, q in zip(WEIGHTS, Q))

A.4.1 Validity comes before ranking

A query at original position $p$ can see $\lfloor (p+1)∕r\rfloor$ completed compressed rows. The metadata supplies that count per query:

591for (int r : args.sparse_compress_ratios) {
592    py::str name = std::to_string(r);
593    auto klen = get_or_create_tensor(klen_dict, has_out, name, {total_q}, torch::kInt32);
594    auto p_klen = klen.mutable_data_ptr<int32_t>();
595    i_q = 0;
596    for (auto p_seq : seq_metadata_list) {
597        for (int32_t j = p_seq->seqlen_k - p_seq->seqlen_q; j < p_seq->seqlen_k; ++j, ++i_q) {
598            p_klen[i_q] = (j + 1) / r;
599        }
600    }
601    TORCH_CHECK_EQ(i_q, total_q);

For C2, query zero sees no completed row; query one sees row zero; query two still sees only row zero. The most recent unpaired token remains available through the layer’s SWA path. A large score in an unwritten or future slot must not enter the selection.

Decode passes an explicit valid range to top-k:

494topk_index = topk(
495    p_sum,
496    self.topk_tokens,
497    sorted=False,
498    end=topk_lens,  # We must pass `end` to the topk kernel since OOB p_sum is not cleaned.
499    indices_type=torch.int32,
500    sorted_index=True,
501    idx_oob_fill_value=-1,
502    return_value=False,
503)[1]
504if candidate_blocks is not None:
505    restore_candidate_indices_inplace(topk_index, candidate_blocks, candidate_block_size)

end excludes invalid score storage, whose contents need not have been cleared. idx_oob_fill_value=-1 fills unused output lanes. The request for sorted_index=True concerns index ordering, not decreasing score order. The lab returns a set of chosen rows in index order and checks its total score against all pairs of valid rows. It also checks empty and undersized ranges.

A.4.2 Candidate blocks are another restriction

Some layers first reduce index scores by candidate block, using a maximum, and select blocks for later layers to search. This grouping is distinct from the allocator’s 128-token page. In the lab, two index keys form a candidate block: block scores are 6 and 3. Retaining only the first block changes token top-2 from $[1,2]$ to $[0,1]$. Thus top-k within candidates need not equal top-k over all valid keys, even with exact arithmetic.

The runtime’s candidate consumer accepts block sizes 8 or 16 in compressed-row units and imposes further width constraints; the lab’s block size two and single retained block are deliberately smaller teaching values. A source layer fills block_candidate_idxs; a consumer computes logits only within that set, then restores compact candidate offsets to logical row indices. Comparing this model with unrestricted attention would require a model-quality check, not merely equality between two implementations of restricted top-k [90, 91].

Prefill limits the size of its intermediate score matrix with compute_chunk_size. It derives the row width from either the full compressed context or the candidate set, applies the memory budget, rounds to multiples of eight and caps the chunk at 8,192 queries. That is a temporary logit-storage budget, not the number of keys finally selected.

A.4.3 The same logical row has two physical representations

For prefill, _load_prefill_kv gathers keys into a concatenated sequence buffer. Top-k adds the request’s ks offset, or restores candidate indices with that offset. With base 10, logical rows 1 and 2 become 11 and 12.

Decode keeps paged storage. With $B$ original-token positions per allocator block, a ratio-$r$ cache has $B_c=B∕r$ compressed rows in that block. Logical row $j$ maps to

The released indexer performs this conversion after candidate offsets have been restored:

506# topk outputs the logical index of the compressed token within the seq, which is converted in place into the kvcache
507# physical slot through block_tables (-1 is kept)
508convert_kv_index(
509    topk_index.unsqueeze(1),
510    compressed_block_size,
511    topk_lens.unsqueeze(1),
512    input_metadata.block_tables,
513)
514return topk_index

The lab’s smaller page has $B=8,r=2,B_c=4$, with block table $[7,2]$. Rows 1, 2 and 4 map to slots 29, 30 and 8. Its last lane represents padding and remains $-1$. The CUDA conversion checks the lane’s position against the valid count; callers must not put an arbitrary invalid index in a lane declared valid. The lab’s sentinel guard is a teaching convenience, not a transcription of that kernel [92].

Finally, main attention reads the SWA and compressed branches together:

640if self.use_swa:
641    attn_kv_args = dict(
642        key_cache=kv_cache.swa_cache,
643        indices_in_kvcache=kv_swa_index,
644        key_cache_extra=kv_cache.ca_cache,
645        indices_in_kvcache_extra=kv_ca_index,
646        index_length_extra=dense_index_len,
647    )

The following decode_mla_mqa_sparse call goes to FlashMLA. It receives ca_cache, not the smaller ica_cache, as its extra key/value storage. This completes the path from compression state to the actual attention read [93].

A.5 Count rows, payloads and owners separately

C2 halves the completed row count relative to C1. Low precision separately changes each row’s byte count. The index format specifies 128 FP4 elements and one one-byte UE8M0 scale per 32 elements:

129FP4_INDEX_KV_FORMAT = KVCacheFormat(
130    quantized_dtype="e2m1",
131    quantized_elements=128,
132    elements_per_sf=32,
133    bf16_elements=0,
134    do_cast=False,
135    sf_format="ue8m0",
136    note="V4/V41, index kv.",
137)

The same file describes main formats. Here E2M1 and E4M3 name floating-point encodings; UE8M0 is the unsigned exponent scale encoding. The relevant accounting is payload bits plus the stored scales:

The main FP4 format uses one E4M3 scale per 16 values; its scale format differs from the index format. Main/SWA caches interleave data and scales within each row, whereas the index cache uses block-wide data and scale sections. The 68-byte figure is a logical per-row byte budget, not a claim that each physical index row is laid out as 64 immediately adjacent data bytes followed by four scales. Alignment and shared ownership affect total allocation [24, 90].

A 128-position C2 block has 64 completed-row slots. An owning main cache uses $64\times 528=33,792$ payload bytes with FP8, or 18,432 with FP4; an owning index cache uses $64\times 68=4,352$. The cache engine rounds regions to its alignment requirements and does not allocate another copy for a sharing consumer. The live 28,672-byte compression-state example is a separate per-request allocation. These equations explain chapter 4’s rows without claiming a measured reduction in model error or end-to-end latency.

A.6 Sharing has several meanings

Chapter 9 separated YOCO’s shared hidden state from segmented sharing. Here are the latter’s distinct objects:

For a four-layer teaching example, main-KV sources might be $[0,0,0,0]$, index-K sources $[0,0,2,2]$, and index sources $[0,1,2,3]$. Each finer segment lies inside a coarser one. An index-K source cannot project from a main KV belonging to a different segment. The topology checks this refinement:

181@staticmethod
182def _validate_map_refines(fine_name: str, fine: list[int], coarse_name: str, coarse: list[int]) -> None:
183    # Under the segmented assumption, coarse[fine[l]] == coarse[l] is equivalent to every segment of fine
184    # falling inside a single segment of coarse.
185    for layer_id, fine_src in enumerate(fine):
186        assert coarse[fine_src] == coarse[layer_id], (
187            f"shared {fine_name} map must refine {coarse_name} map: layer {layer_id} -> {fine_name} source {fine_src}, "
188            f"but {coarse_name} sources differ: {coarse[layer_id]} != {coarse[fine_src]}"
189        )

This is only an example of valid map nesting, not a replacement model configuration. Other checks require compatible layer settings, preceding source layers and a fresh KV source at the YOCO boundary. The topology is static; CrossLayerRuntime carries pre-RoPE KV, loaded index keys, indices and candidate data from source to consumer during a forward call. Under YOCO compaction, full-layout main/index KV may already have been written, so the compact query path skips those stores. Chapter 9 covered that boundary [91, 93].

A.7 Experiments to try next

The labs check rational compression, state lifetime and index/address arithmetic; they do not establish accelerator equivalence or acceptable generation quality. The released compression tests compare device outputs with a tensor reference and check state tensors. We inspected those assertions without running the tests.

1. Change the first pair’s weights so each channel has equal weights. Which outputs become ordinary averages? Keep the chunk checks unchanged.
2. Explain why a zero-valid-row prefill chunk still reserves an output row. Which metadata prevents that placeholder from becoming a real cache entry?
3. For C2 with at most four new positions, derive a sufficient compression ring size and construct the write that would break a smaller ring.
4. Give an example where restricting candidates preserves token top-k, and one where it changes it. What would a model-quality evaluation need in addition?
5. Trace one selected compressed row through ica_cache, top-k, the block table and ca_cache. Which steps change values, and which only change addresses or the set of rows read?

The reading path distinguishes pending state, completed KV and selection keys. Use those distinctions to locate the next reader of a row and check its validity.