dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

[CRITICAL] Enable ixformer native V1/V2 paged attention kernels

Browse Source 
Hardware diagnostics revealed three fatal issues:

1. V1 CRASH: paged_attn.py passes num_kv_heads=4 (int) but ixformer's
   vllm_single_query_cached_kv_attention requires head_mapping as Tensor:
   torch.repeat_interleave(arange(4), 6) = [0,0,0,0,0,0,1,...,3,3,3,3,3,3]
   RuntimeError: Expected Tensor for argument '_4' but found int.
   FIX: Convert int→Tensor in _custom_ops.py paged_attention_v1().

2. V2 NATIVE KERNEL EXISTS but was never called:
   ixformer has vllm_single_query_cached_kv_attention_v2() — a compiled,
   EX-engine-optimized V2 kernel. _custom_ops.py had raise NotImplementedError().
   Our Python V2 (paged_attention_v2_pytorch.py) was a workaround for
   something that already existed in the runtime.
   FIX: Replace NotImplementedError with ixf_F call. V2 signature:
     (output, partition, exp_sums, max_logits, temp_output, query,
      key_cache, value_cache, head_mapping, scale, block_tables,
      context_lens, block_size, max_context_len, alibi_slopes)
   Note 'partition' (int) = max_num_partitions, between output and exp_sums.

3. Triton path: installed at /usr/local/lib/python3.10/ but vllm looks in
   /usr/local/corex/lib64/python3/. Symlink + sys.path fix.

Impact: This replaces ALL Python attention fallbacks with native kernels.
  V1: EX-engine compiled kernel for seq ≤ 8192 (was crashing)
  V2: EX-engine compiled kernel for seq > 8192 (was Python fallback)
  Combined: expect 10-100x speedup on decode path.
This commit is contained in:
Claude
2026-07-31 06:18:32 +00:00
parent 39e32343eb
commit 4867d4f780
2 changed files with 251 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
Dockerfile
View File

@@ -11,6 +11,11 @@ COPY ./paged_attention_v2_triton.py /workspace/paged_attention_v2_triton.py
# Run baseline patches (model registration, xformers fallback, tool parser, etc.)
RUN cd ./qwen3_6_scripts && ./patch_ops.sh
# CRITICAL: Enable ixformer native V1/V2 paged attention kernels.
# Fixes: V1 head_mapping int→Tensor, V2 NotImplementedError → native kernel,
# Triton path mismatch.
RUN python3 /workspace/qwen3_6_scripts/patch_ixformer_native.py
# 1. PagedAttention V2 — fills the NotImplementedError hole
# Enables partitioned attention for long sequences (>8192 tokens)
RUN python3 /workspace/qwen3_6_scripts/patch_paged_attention_v2.py