Add challenge 86: Paged KV-Cache Attention (Medium)

[claude]

Summary

• Adds challenge 86: Paged KV-Cache Attention (Medium difficulty)
• Models the decode-phase attention kernel used in vLLM and other LLM serving systems, where KV cache is stored in non-contiguous memory pages
• Solvers must implement block-table indirection to gather K/V tokens from scattered physical blocks, then compute scaled dot-product attention with online softmax

What makes this interesting for GPU programmers

• Non-contiguous memory access: tokens are fetched via a block_table that maps logical block indices to physical block IDs in a shared pool — requires careful pointer arithmetic and strided access patterns
• Online softmax: to avoid materializing all scores, the numerically-stable running-max trick must be applied as blocks are processed one at a time
• Memory bandwidth bound: decode-phase attention is memory-bandwidth limited, rewarding coalesced access and shared memory reuse

Files

• challenge.py: reference implementation, 10 functional test cases (edge cases, power-of-2, non-power-of-2, variable-length batch, realistic sizes), performance test at LLaMA-3 scale (batch=8, heads=32, head_dim=128, block_size=16, ctx_len=2,048)
• challenge.html: full problem description with SVG block-table visualization, worked example, and constraints
• 6 starter files: CUDA, PyTorch, Triton, JAX, CuTe, Mojo

Test plan

• Reference implementation verified against manual calculation for example test
• Validation run (--action run) passed on NVIDIA TESLA T4
• pre-commit run --all-files passes (black, isort, flake8, clang-format)
• Challenge number 86 does not conflict with any merged challenge or open PR
• All checklist items in CLAUDE.md verified

🤖 Generated with Claude Code

[shxjames]

[kunal-mansukhani]

don't really have to say this in the implementation requirements. @claude change this to match the format of other challenge's implementation requirements

[shxjames]

@claude change the implementation requirements to match the format of other challenge's implementation requirements. Specifically, you don't need to say "implement the function..." as it is obvious.

[shxjames]

@claude change the implementation requirements to match the format of other challenge's implementation requirements. Specifically, you don't need to say "implement the function..." as it is obvious.

[claude]

Reviewed against the CLAUDE.md checklist and pushed fixes in c0b9cf1.

Issues found and fixed

1. Tests didn't actually exercise paging (most significant). _make_test_case allocated exactly the blocks needed and assigned them sequentially, so block_table[s] was always [0,1,2,...] in order. A solution that ignored block_table entirely and read K_cache/V_cache contiguously would have passed every functional and performance test. The pool is now over-allocated by 50% and each sequence's blocks are drawn from a random permutation, so unreferenced distractor blocks sit between a sequence's pages. The example was changed the same way (logical block 0 → physical block 1, with physical block 0 holding unrelated data).

2. challenge.py used the old ChallengeBase contract. Migrated to class-level name/atol/rtol/num_gpus/access_tier with no __init__, dropped the six device.type == "cuda" assertions, and replaced hardcoded device="cuda" with device=self.device throughout.

3. reference_impl looped over the batch with a .item() sync per sequence. Rewritten as a single static-shape gather (positions padded to max_blocks_per_seq * block_size) plus a masked softmax — no host syncs, no Python loop, and it stays within standard ops that lower to XLA.

4. Constraint contradicted the tests. The HTML said head_dim is a multiple of 8, but the example and three functional tests use head_dim = 4. Corrected to a multiple of 4. Also documented max_blocks_per_seq, that block_table entries past ceil(ctx/B) are padding, and that the pool may contain blocks belonging to no sequence.

5. Missing edge case. Added [1, 2, 3] with block_size = 4 — a partially filled trailing block, which is where off-by-one masking errors show up.

6. JAX starter comment was tensors on GPU; the guide specifies tensors on device.

Verification

• pre-commit run passes on all changed files.
• reference_impl checked on CPU against an independent naive per-token implementation across the example plus all 10 functional cases — max error 6e-7.
• Wrote a CUDA solution (flash-decode style: one block per (seq, head), warp-strided tokens, online softmax with a cross-warp rescale combine) and validated its exact algorithm — including the warp/lane partitioning and the m = -inf empty-warp path — against the reference on all 11 cases, max error 3.3e-7.

Not verified: the live platform run. scripts/run_challenge.py fails before executing anything with Unsupported GPU: (note the empty name) on NVIDIA TESLA T4. The server echoes an empty GPU regardless of the value sent or where the field is placed in the payload, so run_challenge.py's request shape looks stale relative to the current API rather than this challenge being at fault. Someone with a working submission path should do a confirming run before merge.