Aller au contenu
login
arrow_backRetour aux issues
poojithdevan4D/pooji-vllm #4

Fused paged attention kernel (remove the per-layer K/V gather)

ecoDébutant help wanted performance

descriptionDescription

The single biggest remaining cost. We gather K/V into a contiguous tensor every layer, every step, then call SDPA. **Where:** `pooji_vllm/llm_engine.py` — `_forward` and `_graph_forward`, the two lines building `K` and `V` from `self.buf[l,0][phys.reshape(-1)]`. **Why it matters:** at batch 32 with 2048-token contexts the gather moves more bytes than the 942 MiB weight read it accompanies. A kernel that reads blocks in place removes that entirely — and also removes the need to zero freshly allocated blocks (`LLMEngine._clear`), since it would never read past the sequence length. **Approach:** Triton is the most approachable route. Start with the decode case only (T=1), fixed block size, and validate against `benchmarks/verify.py` before touching prefill. **Done when:** token-identical output, and `benchmarks/micro.py` shows a decode-step improvement at B=16.
codeOuvre sur GitHub