project_6/docs/paged_attention_kernel_architecture.md

# Paged Attention Kernel Architecture for BI-V100

## Derived from CCCL Algorithm Patterns

This document designs a complete paged attention kernel from first principles,
using CCCL's algorithm implementations as the algorithmic foundation.
Every module maps to a proven CCCL pattern.

---

## 1. Problem Definition

Paged attention computes, for each query token in a decode step:

    output[h, d] = softmax(Q[h] · K[t]^T / √d) · V[t]

where K and V are stored in a **paged block table** (non-contiguous physical memory).

**Qwen3.6 parameters:**
- head_dim (d) = 256
- num_heads (H) = 24
- num_kv_heads (kv_h) = 4, GQA ratio = 6
- seq_len (T) = up to 100,000
- block_size = 16 tokens per physical block
- SMEM per block = 48KB

**The challenge:** K/V are scattered across physical blocks.
A naive implementation does 6,250 random memory accesses for 100K tokens.

---

## 2. Algorithm Decomposition (Three Levels from CCCL)

### Level 1: Warp Reduce (from `warp_reduce_shfl.cuh`)

**CCCL pattern:** `shfl.sync.down` butterfly reduction in log2(32) = 5 steps.
Each step: `output = reduction_op(input, ShuffleDown(input, 1 << step))`.

**In attention:** Within one warp (32 threads), each thread holds QK^T scores
for a subset of KV tokens. Warp reduce computes:
- `max_score = warp_reduce(scores, max_op)` — for softmax numerical stability
- `exp_sum = warp_reduce(exp(scores - max_score), plus_op)` — softmax denominator
- `weighted_v = warp_reduce(exp(scores - max_score) * V[t], plus_op)` — numerator

This is a **compound reduction** — the same pattern as CCCL's `summary_statistics.cu`
where (count, mean, M2) are reduced together with a custom binary op.

**Our compound type:**
```
struct attention_partial {
    float max_score;      // running max of QK^T
    float exp_sum;        // sum of exp(score - max_score)
    float weighted_v[D];  // sum of exp(score - max_score) * V
};
```

**Binary op** (from `summary_statistics.cu`):
```
attention_partial combine(attention_partial a, attention_partial b) {
    float new_max = max(a.max_score, b.max_score);
    float scale_a = exp(a.max_score - new_max);
    float scale_b = exp(b.max_score - new_max);
    return {
        new_max,
        scale_a * a.exp_sum + scale_b * b.exp_sum,
        scale_a * a.weighted_v + scale_b * b.weighted_v  // element-wise
    };
}
```

This is exactly the online softmax from Flash Attention.
It's also exactly CCCL's binary reduction op pattern.

### Level 2: Block Reduce (from `block_reduce_warp_reductions.cuh`)

**CCCL pattern:** Each warp produces a `warp_aggregate`. Lane 0 of each warp
writes it to `SMEM warp_aggregates[warp_id]`. Then thread 0 serially reduces
across warps:
```
for (warp_idx = 1; warp_idx < warps; ++warp_idx)
    aggregate = reduction_op(aggregate, warp_aggregates[warp_idx]);
```

**In attention:** One thread block processes one partition of the KV sequence
(e.g., PARTITION_SIZE = 512 tokens). Multiple warps within the block each handle
a chunk of these 512 tokens.

- Warp 0: tokens 0..63 (BLOCK_N=64 at a time, or 32 for head_dim=256)
- Warp 1: tokens 64..127
- ...
- Warp W-1: tokens (W-1)*64..511

Each warp produces an `attention_partial`. Block reduce merges them:
```
__shared__ attention_partial warp_partials[NUM_WARPS];
warp_partials[warp_id] = my_warp_result;
__syncthreads();
if (threadIdx.x == 0) {
    attention_partial block_result = warp_partials[0];
    for (int w = 1; w < NUM_WARPS; w++)
        block_result = combine(block_result, warp_partials[w]);
    // Write block_result to global: tmp_output, exp_sums, max_logits
}
```

**SMEM layout for attention_partial at head_dim=256:**
- max_score: 4 bytes
- exp_sum: 4 bytes  
- weighted_v[256]: 256 × 4 = 1024 bytes
- Total per warp: 1032 bytes
- For 4 warps: 4128 bytes (fits easily in 48KB)

### Level 3: Cross-Partition Coordination (from `agent_scan.cuh` + decoupled lookback)

**CCCL pattern:** `TilePrefixCallbackOp` implements decoupled lookback.
Each tile block:
1. Computes its local aggregate
2. Publishes local aggregate to global `tile_state` (PARTIAL status)
3. Warp 0 looks back through predecessor tiles:
   - If predecessor has INCLUSIVE status → directly use its prefix
   - If predecessor has PARTIAL status → accumulate and keep looking back
4. Once prefix is resolved, update own status to INCLUSIVE

**In attention (V2):** Each partition block has its `attention_partial`.
The cross-partition reduction is simpler than scan because attention
partitions are **commutative** — we don't need prefix sums, just a
global reduce.

But the coordination pattern is the same:
1. Each partition block writes its (max_logit, exp_sum, partial_output) to
   global memory: `tmp_output[seq, head, partition, :]`
2. A separate reduction kernel (or the last partition block) reads all
   partitions and does the final combine.

**Simplification over CCCL's lookback:** Since attention partitions are
independent (no prefix dependency), we don't need the lookback polling loop.
Each partition can run fully independently. The reduction is a simple
parallel reduce over `num_partitions` compound values.

For 100K tokens / 512 partition_size = ~200 partitions.
200 `attention_partial` values × (4 + 4 + 256×4) = 200 × 1032 = ~200KB.
One block can reduce all 200 in registers + SMEM.

---

## 3. Paged K/V Gather (from `block_load.cuh` + `cache_modified_input_iterator.cuh`)

**CCCL pattern:** `BlockLoadWarpTranspose` loads contiguous global memory
into a striped register layout that enables coalesced access. Each thread
loads `ITEMS_PER_THREAD` elements, and the warp transposes them so each
thread gets its tile of the data.

**In paged attention:** K/V are not contiguous — they're indexed through
`block_tables[seq, logical_block] → physical_block`.
- Key cache: `[num_blocks, kv_heads, head_dim/x, block_size, x]`
  where x = 16/sizeof(dtype) is the packing factor
- Value cache: `[num_blocks, kv_heads, head_dim, block_size]`

The gather pattern (from `prefix_prefill.py`, which works on BI-V100):
```
# For BLOCK_N tokens starting at position start_n:
token_ids = start_n + tl.arange(0, BLOCK_N)
logical_blocks = token_ids // block_size
within_block = token_ids % block_size
physical_blocks = tl.load(block_tables + seq * stride + logical_blocks * stride)

# K gather: compute 2D offset array [HEAD_DIM, BLOCK_N]
off_k = (physical_blocks[None, :] * stride_kc_b +
         kv_head * stride_kc_h +
         (offs_d[:, None] // x) * stride_kc_dx +
         within_block[None, :] * stride_kc_bs +
         (offs_d[:, None] % x) * stride_kc_x)
k = tl.load(key_cache + off_k, mask=valid_mask)
```

This is an **indirect gather** — the physical block ID comes from a table lookup.
CCCL's `CacheModifiedInputIterator` handles the cache hint part, but the
indirect indexing is our addition.

**Memory access pattern:**
- block_tables lookup: 1 global read per BLOCK_N tokens (amortized)
- K gather: BLOCK_N × HEAD_DIM / x global reads (scattered by physical block)
- V gather: BLOCK_N × HEAD_DIM global reads (similar scatter)

For BLOCK_N=32, HEAD_DIM=256, x=8: 32 × 32 = 1024 reads for K per iteration.
At 16 bytes per read (128-bit): 16KB per K load.
V is similar. Total per iteration: ~32KB — fits in L2 (6MB on BI-V100).

---

## 4. GQA (Grouped Query Attention) Handling

**The insight:** 6 query heads share 1 KV head. Loading KV once and
computing 6 sets of QK^T scores is 6x more compute-efficient than
loading KV 6 times.

**CCCL analogy:** This is like `BlockReduce` where we have 6 different
reduction operations on the same input data. CCCL doesn't have this exact
pattern, but the principle is: share data loads, parallelize computation.

**Implementation:**
- Each thread block handles one `(seq, kv_head, partition)` triple
- Within the block, 6 query heads are processed simultaneously
- Q vectors: 6 × HEAD_DIM = 6 × 256 = 1536 values in registers (per thread
  this is 1536/32 = 48 registers — feasible)
- K/V: loaded once for the kv_head, broadcast across all 6 query heads
- Scores: 6 × BLOCK_N values per iteration
- Weighted V: 6 × HEAD_DIM per thread's accumulator

This reduces K/V cache reads by 6x (the GQA ratio).

Grid: `(num_seqs, num_kv_heads, num_partitions)` = `(1, 4, 200)` = 800 blocks
instead of `(1, 24, 200)` = 4800 blocks.

Each block does 6x more compute but reads KV only once.

---

## 5. SMEM Budget

For one block processing BLOCK_N=32 KV tokens across 6 query heads:

| Item | Size | Notes |
|------|------|-------|
| K tile [HEAD_DIM, BLOCK_N] | 32×256×2 = 16KB | fp16, loaded from paged cache |
| V tile [BLOCK_N, HEAD_DIM] | 32×256×2 = 16KB | fp16, loaded from paged cache |
| Warp partials [4 warps × attention_partial] | 4×(4+4+256×4) = 4.1KB | For block-level reduce |
| Q vectors [6 × HEAD_DIM] | 6×256×4 = 6KB | In registers ideally, SMEM if spills |
| **Total** | **42.1KB** | **≤ 48KB ✓** |

Tight but feasible. If Q stays in registers (likely with 4 warps × 32 threads
= 128 threads, each handling 6×256/128 = 12 Q values), total SMEM is 36.1KB.

---

## 6. Kernel Launch Configuration

**Phase 1: Partitioned Attention**
- Grid: `(num_seqs, num_kv_heads, num_partitions)`
- Block: `(NUM_WARPS × 32)` = 128 threads (4 warps)
- Each block processes:
  - PARTITION_SIZE = 512 KV tokens
  - 6 query heads (GQA broadcast)
  - Produces 6 × (max_logit, exp_sum, partial_output[256])

**Phase 2: Cross-Partition Reduction**
- Grid: `(num_seqs, num_kv_heads)` 
- Block: 128 threads
- Each block reduces ~200 partitions × 6 query heads
- Uses `combine()` op (same as CCCL `BlockReduce` but with `attention_partial`)

**Phase 1 iterations per block:**
- PARTITION_SIZE / BLOCK_N = 512 / 32 = 16 iterations
- Each iteration: load K[32, 256] + V[32, 256], compute 6×32 scores, update 6 accumulators

---

## 7. Implementation Mapping

| Module | CCCL Source | Our Implementation |
|--------|------------|-------------------|
| Warp-level QK^T + softmax | `warp_reduce_shfl.cuh` | Triton: `tl.sum()` within warp-sized groups |
| Block-level partition reduce | `block_reduce_warp_reductions.cuh` | Triton: shared memory + `tl.reduce()` |
| Cross-partition combine | `agent_scan.cuh` (simplified, no lookback) | Separate reduction kernel |
| Paged K/V gather | `block_load.cuh` + indirect indexing | `prefix_prefill.py` pattern adapted |
| Online softmax | `summary_statistics.cu` binary op | `combine(attention_partial, attention_partial)` |
| GQA broadcast | (no exact CCCL analog) | Multiple Q per KV load |

---

## 8. Why This Design Beats Python V2

Current Python V2 (3 bmm launches + Python overhead):
- gather all KV → permute → contiguous → bmm → reshape → softmax → bmm → reduce
- **Python-CUDA boundary crossed 10+ times per decode step**
- **Full KV tensor materialized in GPU memory** (200MB-2.4GB depending on GQA)

This kernel (2 GPU launches, zero Python-CUDA crossings during compute):
- Phase 1: single kernel, K/V loaded tile-by-tile from paged cache (never materialized)
- Phase 2: single kernel, reduces 200 partitions in SMEM
- **KV cache stays in paged format** — no gather/permute/contiguous overhead
- **GQA broadcast within kernel** — KV loaded once for 6 heads

Expected improvement over Python V2: **10-100x** (eliminating Python overhead
and memory allocation dominates at decode batch_size=1).

Expected improvement over no V2 (V1 only for seq ≤ 8192): **enables long-context
decode** which V1 cannot do due to SMEM overflow at 48KB.

---

## 9. Implementation Priority

1. **Triton implementation** — if Triton works on BI-V100 with BLOCK=32, head_dim=256:
   Use the `prefix_prefill.py` paged gather pattern, add the compound reduction.
   This is the fastest path to a working kernel.

2. **Compiled CUDA kernel** — if `/usr/local/corex/` has a compiler (ixcc):
   Write the kernel in CUDA using the CCCL patterns directly.
   `warp_reduce_shfl` → `__shfl_down_sync` PTX
   `block_reduce` → SMEM warp_aggregates pattern
   Compile with `torch.utils.cpp_extension.load()` at Docker build time.

3. **Python V2** (current) — fallback if neither Triton nor CUDA works:
   Already written, tested, has GQA broadcast optimization.
   This is the floor, not the ceiling.