3722503dee
Complete rewrite of qwen3_6_scripts/paged_attn.py with 4 optimizations:
1. _forward_prefix_pytorch: Pre-gather ALL context K/V outside tile loop
Before: each of 195 tiles does key_cache[blk_ids].permute().contiguous()
After: ONE key_cache[all_ctx_blk_ids].permute().contiguous() upfront,
tile loop just does ctx_k_t[:, :, start:end] (view, no copy)
Eliminates 194 redundant gather+permute+contiguous calls per prefill.
2. forward_decode: V2 enabled via original heuristic
Before: use_v1 = True (hardcoded, V2 was NotImplementedError)
After: V2 works (paged_attention_v2_pytorch), use vllm's heuristic:
seq_len > 8192 → V2 (partitioned, better parallelism)
seq_len <= 8192 → V1 (single-block, less overhead)
3. forward_prefix: Triton try/fallback
First call attempts Triton context_attention_fwd (if HAS_TRITON).
If it hangs/errors, permanently falls back to PyTorch.
If it works: 10-50x prefill improvement.
4. _PYTORCH_DECODE_THRESHOLD: 32768 → 65536
Keeps more decode requests on the fast compiled v1 kernel.
All changes are safe: Triton has try/except, V2 fallback exists,
threshold can be lowered back if v1 crashes at 64K.