Implement the KV cache append step used during autoregressive decoding. Given a single new key vector knew and value vector vnew of shape (D,), write them into row pos of the pre-allocated caches in place:
Kcache[pos,:]←knew,Vcache[pos,:]←vnewThis is the per-step scatter that grows the cache by one token during inference. There is no separate output tensor: the kernel mutates the caches in place. The pattern is a 1D map along D, partitioned into tiles so larger embedding dimensions stay safe with a tail mask.
Input: k_new = [1.0, 2.0], v_new = [3.0, 4.0], k_cache = [[0.0, 0.0], [0.0, 0.0]], v_cache = [[0.0, 0.0], [0.0, 0.0]], pos = 0
Output: k_cache = [[1.0, 2.0], [0.0, 0.0]], v_cache = [[3.0, 4.0], [0.0, 0.0]]
Input: k_new = [5.0, 6.0, 7.0], v_new = [8.0, 9.0, 10.0], k_cache = [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], v_cache = [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], pos = 1
Output: k_cache = [[0.0, 0.0, 0.0], [5.0, 6.0, 7.0], [0.0, 0.0, 0.0]], v_cache = [[0.0, 0.0, 0.0], [8.0, 9.0, 10.0], [0.0, 0.0, 0.0]]
Each program handles one tile along D. The row offset into the caches is pos⋅D+offsets where offsets is the per-tile index vector.
Read knew and vnew at the tile offsets and write them at the row-offset positions in Kcache and Vcache. Both pairs of reads and writes share the same tail mask.
The kernel must not touch any row other than pos. The harness checks the whole cache, so unmodified rows must equal their original values.
torch.float32 on cuda.atol=1e-3 on the full caches after the write.k_cache and v_cache in place; no out tensor is passed.Accepts: array
Accepts: array
Accepts: array
Accepts: array
Accepts: number
Submit a correct solution to see your kernel performance.
Implement the KV cache append step used during autoregressive decoding. Given a single new key vector knew and value vector vnew of shape (D,), write them into row pos of the pre-allocated caches in place:
Kcache[pos,:]←knew,Vcache[pos,:]←vnewThis is the per-step scatter that grows the cache by one token during inference. There is no separate output tensor: the kernel mutates the caches in place. The pattern is a 1D map along D, partitioned into tiles so larger embedding dimensions stay safe with a tail mask.
Input: k_new = [1.0, 2.0], v_new = [3.0, 4.0], k_cache = [[0.0, 0.0], [0.0, 0.0]], v_cache = [[0.0, 0.0], [0.0, 0.0]], pos = 0
Output: k_cache = [[1.0, 2.0], [0.0, 0.0]], v_cache = [[3.0, 4.0], [0.0, 0.0]]
Input: k_new = [5.0, 6.0, 7.0], v_new = [8.0, 9.0, 10.0], k_cache = [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], v_cache = [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], pos = 1
Output: k_cache = [[0.0, 0.0, 0.0], [5.0, 6.0, 7.0], [0.0, 0.0, 0.0]], v_cache = [[0.0, 0.0, 0.0], [8.0, 9.0, 10.0], [0.0, 0.0, 0.0]]
Each program handles one tile along D. The row offset into the caches is pos⋅D+offsets where offsets is the per-tile index vector.
Read knew and vnew at the tile offsets and write them at the row-offset positions in Kcache and Vcache. Both pairs of reads and writes share the same tail mask.
The kernel must not touch any row other than pos. The harness checks the whole cache, so unmodified rows must equal their original values.
torch.float32 on cuda.atol=1e-3 on the full caches after the write.k_cache and v_cache in place; no out tensor is passed.Accepts: array
Accepts: array
Accepts: array
Accepts: array
Accepts: number
Submit a correct solution to see your kernel performance.