Autoregressive KV-cache append and causal attention

Your task

Function signature

kv_cache_attention_step(
    query: torch.Tensor,
    new_key: torch.Tensor,
    new_value: torch.Tensor,
    key_cache: torch.Tensor,
    value_cache: torch.Tensor,
    lengths: torch.Tensor,
    append_mask: torch.Tensor,
    options: dict,
) -> dict[str, torch.Tensor]

Worked example

Two queries see different prefixes

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• query[B,H_q,Q,D], new_key[B,H_k,Q,D], new_value[B,H_k,Q,D_v], key_cache[B,H_k,M,D], value_cache[B,H_k,M,D_v] are finite floating with matching dtype/device. B,H_q,H_k,D,D_v,M\ge 1,Q\ge 0; H_q must be a multiple of H_k.
• Query head h uses KV head \operatorname{floor}(h/(H_q/H_k)); H_q=H_k is ordinary MHA. Inputs are already projected and, if required by the model, position-rotated. No projections or RoPE are performed here. lengths[B] integers specify the old valid cache prefix, 0\le \mathrm{length}\le M.
• append_mask[B,Q] is bool and each row must be a prefix of True values followed by False. Its count n must satisfy \mathrm{length}+n\le M. Write new_key/value at positions \mathrm{length}+t only for True append slots t.
• Return the entire updated cache, preserving every unwritten entry exactly (including unused tail storage); never modify inputs. Updated lengths are \mathrm{length}+n. A valid query t may attend only cache slots j\le \mathrm{old\_length}+t and j<\mathrm{new\_length}. Invalid query rows return zero attention/context.
• Compute \mathrm{scores}(q,k)=\mathrm{scale}\cdot \operatorname{dot}(q,k), with positive scale default 1/\sqrt{D}, and a stable softmax on allowed keys. Return attention[B,H_q,Q,M] and context[B,H_q,Q,D_v]=\mathrm{attention}\,\mathrm{updated\_values}.
• Cache-tail values and future tokens cannot influence earlier query context even though the returned cache transports unused tail entries. Q=0 or rows with n=0 preserve their cache/lengths and return empty/zero query outputs. scale is the only key.
• No dropout, eviction, output projection or neural-model invocation is included; capacity overflow and non-prefix append masks are outside this contract.

Why this matters

Starter code

You write the body of this function in the BioTorch workspace. Your code runs in an isolated PyTorch sandbox; Run checks the public examples and Submit grades the full test suite, including gradients.

import torch


def kv_cache_attention_step(
    query: torch.Tensor,  # (B, Hq, Q, D)
    new_key: torch.Tensor,  # (B, Hk, Q, D)
    new_value: torch.Tensor,  # (B, Hk, Q, Dv)
    key_cache: torch.Tensor,  # (B, Hk, M, D)
    value_cache: torch.Tensor,  # (B, Hk, M, Dv)
    lengths: torch.Tensor,  # (B,)
    append_mask: torch.Tensor,  # (B, Q)
    options: dict,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    key_cache: (B, Hk, M, D).
    value_cache: (B, Hk, M, Dv).
    lengths: (B,).
    attention: (B, Hq, Q, M).
    context: (B, Hq, Q, Dv)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement kv_cache_attention_step")

Sources and further reading

• doi.org
• togethercomputer/stripedhyena · model.py
• togethercomputer/stripedhyena · generation.py
• arxiv.org
• salesforce/progen · modeling_progen.py