project_6

dylanyunlong/project_6

Fork 0

Commit Graph

Author	SHA1	Message	Date
dylanyunlon	15ef28e863	[OPT] Vectorize paged_attention_v2 — eliminate block-gather for-loop Before: 3 nested Python for-loops for seq_idx: (1 iteration at max_num_seqs=1) for block_idx: (6250 iterations at seq_len=100K, block_size=16) key_cache[physical_block] + permute + reshape per block for part_idx: (195 iterations at seq_len=100K, PARTITION=512) torch.einsum per partition After: 1 seq loop (trivial) + batched gather + bmm partition loop for seq_idx: (1 iteration — same) key_cache[blk_ids] (ONE index_select for all 6250 blocks) .permute().reshape() (ONE reshape for entire sequence) for part_idx: (195 iterations, each uses torch.bmm) torch.bmm (batched over all heads simultaneously) Key changes: - Block gather: block-by-block Python loop → single key_cache[blk_ids] Eliminates 6250 Python iterations for 100K sequence - GQA: repeat_interleave (allocates) → expand (view, zero-copy) - Partition attn: torch.einsum → torch.bmm (more efficient for batched) - Phase 2 reduction: unchanged (already vectorized) The block_idx loop was the real killer: 6250 Python-level tensor operations (index + permute + reshape + slice) per decode step. Now it's one operation.	2026-07-30 15:44:46 +00:00
Claude	9cb7f9d037	[OPT] PagedAttention V2 implementation — fill the NotImplementedError hole The single biggest performance bottleneck in the baseline: paged_attention_v2 = raise NotImplementedError() paged_attn.py: use_v1 = True (hardcoded to avoid calling V2) V1 limitation: processes entire KV sequence in one kernel launch. For seq_len=100K, this is a single massive attention computation. V2: splits into PARTITION_SIZE=512 chunks, runs them in parallel, then reduces with log-sum-exp. 195 parallel partitions vs 1. Implementation (paged_attention_v2_pytorch.py): Phase 1: Per-partition attention - For each (seq, head, partition): compute QK^T, softmax, weighted V sum - Store partial: tmp_output, exp_sums, max_logits (per partition) Phase 2: Cross-partition reduction (log-sum-exp) - global_max = max(max_logits across partitions) - rescale = exp(partition_max - global_max) × partition_exp_sum - output = Σ (rescale / total_sum) × partition_output This is the same algorithm as vllm's paged_attention_v2_kernel.cu: - The reduction pattern is identical to CCCL's block_reduce_warp_reductions (combine partial statistics from independent segments) - The online softmax tiling is the same as Flash Attention's partitioning Integration: - patch_paged_attention_v2.py patches _custom_ops.py and paged_attn.py - Removes use_v1=True hardcode → V2 used for seq_len > 8192 - Dockerfile adds the patch step This is a PyTorch implementation (no CUDA compilation needed). Next step: if /usr/local/corex/ has ixcc or nvcc-compatible compiler, replace with compiled CUDA kernel for further speedup.	2026-07-30 15:40:14 +00:00

Author

SHA1

Message

Date

dylanyunlon

15ef28e863

[OPT] Vectorize paged_attention_v2 — eliminate block-gather for-loop

Before: 3 nested Python for-loops
  for seq_idx:           (1 iteration at max_num_seqs=1)
    for block_idx:       (6250 iterations at seq_len=100K, block_size=16)
      key_cache[physical_block] + permute + reshape per block
    for part_idx:        (195 iterations at seq_len=100K, PARTITION=512)
      torch.einsum per partition

After: 1 seq loop (trivial) + batched gather + bmm partition loop
  for seq_idx:           (1 iteration — same)
    key_cache[blk_ids]   (ONE index_select for all 6250 blocks)
    .permute().reshape() (ONE reshape for entire sequence)
    for part_idx:        (195 iterations, each uses torch.bmm)
      torch.bmm          (batched over all heads simultaneously)

Key changes:
  - Block gather: block-by-block Python loop → single key_cache[blk_ids]
    Eliminates 6250 Python iterations for 100K sequence
  - GQA: repeat_interleave (allocates) → expand (view, zero-copy)
  - Partition attn: torch.einsum → torch.bmm (more efficient for batched)
  - Phase 2 reduction: unchanged (already vectorized)

The block_idx loop was the real killer: 6250 Python-level tensor operations
(index + permute + reshape + slice) per decode step. Now it's one operation.

2026-07-30 15:44:46 +00:00

Claude

9cb7f9d037

[OPT] PagedAttention V2 implementation — fill the NotImplementedError hole

The single biggest performance bottleneck in the baseline:
paged_attention_v2 = raise NotImplementedError()
paged_attn.py: use_v1 = True (hardcoded to avoid calling V2)

V1 limitation: processes entire KV sequence in one kernel launch.
For seq_len=100K, this is a single massive attention computation.
V2: splits into PARTITION_SIZE=512 chunks, runs them in parallel,
then reduces with log-sum-exp. 195 parallel partitions vs 1.

Implementation (paged_attention_v2_pytorch.py):
  Phase 1: Per-partition attention
    - For each (seq, head, partition): compute QK^T, softmax, weighted V sum
    - Store partial: tmp_output, exp_sums, max_logits (per partition)
  Phase 2: Cross-partition reduction (log-sum-exp)
    - global_max = max(max_logits across partitions)
    - rescale = exp(partition_max - global_max) × partition_exp_sum
    - output = Σ (rescale / total_sum) × partition_output

This is the same algorithm as vllm's paged_attention_v2_kernel.cu:
  - The reduction pattern is identical to CCCL's block_reduce_warp_reductions
    (combine partial statistics from independent segments)
  - The online softmax tiling is the same as Flash Attention's partitioning

Integration:
  - patch_paged_attention_v2.py patches _custom_ops.py and paged_attn.py
  - Removes use_v1=True hardcode → V2 used for seq_len > 8192
  - Dockerfile adds the patch step

This is a PyTorch implementation (no CUDA compilation needed).
Next step: if /usr/local/corex/ has ixcc or nvcc-compatible compiler,
replace with compiled CUDA kernel for further speedup.

2026-07-30 15:40:14 +00:00

2 Commits