How should I practice Transformer KV Cache Operations?

Start by reading the full statement, solve the visible examples, then compare the pattern against related Nvidia phone screen problems in FastPrep.

Is Transformer KV Cache Operations associated with Nvidia?

FastPrep tags this as a Nvidia phone screen coding practice problem when the problem catalog has that company and stage context.

Problem · Design

Transformer KV Cache Operations

Learn this problem

● Medium

NvidiaFULLTIMEPHONE SCREEN

Problem statement

Implement a Transformer key-value cache configured with layerCount layers, headCount attention heads, and headDimension values per head. Each sequence ID owns an independent ordered token cache at each layer. A token stores one key tensor and one value tensor, each flattened in row-major order into exactly headCount * headDimension integers.

Process these operations in order:

APPEND|sequenceId|layer|keyValues|valueValues: append one token to the sequence and layer. The two value fields are comma-separated flattened tensors. Return the token's zero-based index in that layer.
READ|sequenceId|layer|startInclusive|endExclusive: return the requested ordered token slice. Serialize each token as keyValues/valueValues and separate tokens with semicolons. Return an empty string for an empty slice or a missing sequence.
RESET|sequenceId: remove every cached token for the sequence across all layers and return the total number removed. A missing sequence returns 0. A later append for that sequence starts again at index 0.

Return one result string for every operation, in operation order.

Function

processTransformerKvCache(layerCount: int, headCount: int, headDimension: int, operations: String[]) → String[]

Examples

Example 1

layerCount = 2headCount = 1headDimension = 2operations = ["APPEND|chat|0|1,2|10,20","APPEND|chat|0|3,4|30,40","APPEND|chat|1|5,6|50,60","READ|chat|0|0|2","RESET|chat","READ|chat|0|0|0"]return = ["0", "1", "0", "1,2/10,20;3,4/30,40", "3", ""]

Token indices are local to a sequence and layer. Reset removes two tokens from layer 0 and one from layer 1, after which the read is empty.

Example 2

layerCount = 2headCount = 1headDimension = 1operations = ["APPEND|a|0|7|70","APPEND|b|0|8|80","READ|a|0|0|1","RESET|b","APPEND|b|0|9|90","READ|b|0|0|1"]return = ["0", "0", "7/70", "1", "0", "9/90"]

Sequences a and b are independent. Resetting b does not affect a, and the next token for b receives index 0.

Example 3

layerCount = 1headCount = 2headDimension = 2operations = ["READ|missing|0|0|0","RESET|missing","APPEND|s|0|1,-2,3,-4|5,6,7,8","READ|s|0|0|1"]return = ["", "0", "0", "1,-2,3,-4/5,6,7,8"]

Missing-sequence operations are empty or zero. The final read preserves the row-major flattened key and value tensors exactly.

Constraints

1 <= layerCount <= 128
1 <= headCount <= 128
1 <= headDimension <= 256
1 <= operations.length <= 10000
There are at most 1000 active sequence IDs and at most 100000 cached tokens at any time.
Every layer is in [0, layerCount - 1]. Each appended key and value contains exactly headCount * headDimension signed integers.
For an existing sequence and layer, every read satisfies 0 <= startInclusive <= endExclusive <= tokenCount.