78a0ebd516
Hardware testing confirmed: V1: K=[blocks, kv_heads, head_dim/x, block_size, x] (5D), V=[blocks, kv_heads, head_dim, block_size] (4D) → OK V2: K=[blocks, kv_heads, block_size, head_dim] (4D), V=[blocks, kv_heads, block_size, head_dim] (4D) → OK V2 with V1's layout → FAIL (Expected key_cache.dim()==4, value_cache.size(3)==head_size) V1 and V2 use DIFFERENT cache memory layouts in ixformer. V2 patch now converts cache on the fly before calling native kernel: K: permute(0,1,3,2,4).reshape → [B,H,bs,d] V: permute(0,1,3,2).contiguous → [B,H,bs,d] This is a view+reshape for K (no copy if contiguous) and a transpose+contiguous for V. The cost is one V copy per decode step, but this enables the native compiled V2 kernel which is 10-100x faster than the Python fallback it replaces.
267 lines
9.5 KiB
Python
267 lines
9.5 KiB
Python
"""
|
|
patch_ixformer_native.py — Enable ixformer's native V1/V2 paged attention kernels
|
|
===================================================================================
|
|
|
|
CRITICAL FIXES discovered from hardware diagnostics:
|
|
|
|
1. V1 head_mapping: paged_attn.py passes num_kv_heads (int=4), but
|
|
ixformer_torch_ops.vllm_single_query_cached_kv_attention requires a Tensor.
|
|
The Tensor is: torch.repeat_interleave(torch.arange(num_kv_heads, device), num_queries_per_kv)
|
|
For Qwen3.6: [0,0,0,0,0,0, 1,1,1,1,1,1, 2,2,2,2,2,2, 3,3,3,3,3,3]
|
|
|
|
2. V2 native kernel EXISTS in ixformer (vllm_single_query_cached_kv_attention_v2)
|
|
but _custom_ops.py has raise NotImplementedError(). The 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, use_sqrt_alibi)
|
|
Note: V2 has an extra 'partition' (int) parameter that V1 doesn't.
|
|
|
|
3. Triton 2.3.1 is installed at /usr/local/lib/python3.10/site-packages/triton
|
|
but vllm (at /usr/local/corex/lib64/python3/dist-packages/vllm) says
|
|
"Triton not installed" — path mismatch.
|
|
|
|
This patch fixes all three by modifying _custom_ops.py and the Triton import.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
|
|
|
|
VLLM_ROOT = "/usr/local/corex/lib64/python3/dist-packages/vllm"
|
|
CUSTOM_OPS_PATH = os.path.join(VLLM_ROOT, "_custom_ops.py")
|
|
|
|
|
|
def patch_custom_ops():
|
|
"""Fix V1 head_mapping and replace V2 NotImplementedError with native ixformer kernel."""
|
|
with open(CUSTOM_OPS_PATH, "r") as f:
|
|
content = f.read()
|
|
|
|
changes = []
|
|
|
|
# --- Fix 1: V1 head_mapping must be a Tensor ---
|
|
# Current code passes head_mapping directly through.
|
|
# paged_attn.py passes num_kv_heads (int).
|
|
# We need to convert int → Tensor inside _custom_ops.py.
|
|
old_v1 = '''def paged_attention_v1(
|
|
output,
|
|
query,
|
|
key_cache,
|
|
value_cache,
|
|
head_mapping,
|
|
scale,
|
|
block_tables,
|
|
context_lens,
|
|
block_size,
|
|
max_context_len,
|
|
alibi_slopes=None,
|
|
kv_cache_dtype=None,
|
|
):
|
|
return ixf_F.vllm_single_query_cached_kv_attention(
|
|
output,
|
|
query,
|
|
key_cache,
|
|
value_cache,
|
|
head_mapping,
|
|
scale,
|
|
block_tables,
|
|
context_lens,
|
|
block_size,
|
|
max_context_len,
|
|
alibi_slopes,
|
|
)'''
|
|
|
|
new_v1 = '''def paged_attention_v1(
|
|
output,
|
|
query,
|
|
key_cache,
|
|
value_cache,
|
|
head_mapping,
|
|
scale,
|
|
block_tables,
|
|
context_lens,
|
|
block_size,
|
|
max_context_len,
|
|
alibi_slopes=None,
|
|
kv_cache_dtype=None,
|
|
):
|
|
# BI-V100 fix: ixformer requires head_mapping as Tensor, not int.
|
|
# paged_attn.py passes num_kv_heads (int). Convert here.
|
|
if isinstance(head_mapping, int):
|
|
num_kv_heads = head_mapping
|
|
num_heads = query.shape[1]
|
|
num_queries_per_kv = num_heads // num_kv_heads
|
|
head_mapping = torch.repeat_interleave(
|
|
torch.arange(num_kv_heads, dtype=torch.int32, device=query.device),
|
|
num_queries_per_kv)
|
|
return ixf_F.vllm_single_query_cached_kv_attention(
|
|
output,
|
|
query,
|
|
key_cache,
|
|
value_cache,
|
|
head_mapping,
|
|
scale,
|
|
block_tables,
|
|
context_lens,
|
|
block_size,
|
|
max_context_len,
|
|
alibi_slopes,
|
|
)'''
|
|
|
|
if old_v1 in content:
|
|
content = content.replace(old_v1, new_v1, 1)
|
|
changes.append("V1: Added int→Tensor conversion for head_mapping")
|
|
else:
|
|
print(" [warn] V1 function body not found as expected — may already be patched")
|
|
|
|
# --- Fix 2: V2 replace NotImplementedError with native ixformer kernel ---
|
|
# ixformer signature:
|
|
# vllm_single_query_cached_kv_attention_v2(
|
|
# 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, use_sqrt_alibi)
|
|
# _PARTITION_SIZE = 512 (from paged_attn.py)
|
|
|
|
old_v2 = ''' blocksparse_block_size: int = 64,
|
|
blocksparse_head_sliding_step: int = 0,
|
|
) -> None:
|
|
raise NotImplementedError()'''
|
|
|
|
new_v2 = ''' blocksparse_block_size: int = 64,
|
|
blocksparse_head_sliding_step: int = 0,
|
|
) -> None:
|
|
# BI-V100: Use ixformer native V2 kernel instead of NotImplementedError.
|
|
# Convert num_kv_heads (int) → head_mapping (Tensor)
|
|
if isinstance(num_kv_heads, int):
|
|
num_heads = query.shape[1]
|
|
num_queries_per_kv = num_heads // num_kv_heads
|
|
head_mapping = torch.repeat_interleave(
|
|
torch.arange(num_kv_heads, dtype=torch.int32, device=query.device),
|
|
num_queries_per_kv)
|
|
else:
|
|
head_mapping = num_kv_heads
|
|
|
|
_PARTITION_SIZE = 512
|
|
max_num_partitions = (max_seq_len + _PARTITION_SIZE - 1) // _PARTITION_SIZE
|
|
|
|
# V2 native kernel expects different cache layout than V1:
|
|
# V1: K=[blocks, kv_heads, head_dim/x, block_size, x] (5D), V=[blocks, kv_heads, head_dim, block_size] (4D)
|
|
# V2: K=[blocks, kv_heads, block_size, head_dim] (4D), V=[blocks, kv_heads, block_size, head_dim] (4D)
|
|
# Convert on the fly. This is a view/permute, not a data copy (for contiguous inputs).
|
|
if key_cache.dim() == 5:
|
|
# K: [B, H, d/x, bs, x] → [B, H, bs, d]
|
|
B, H, dx, bs, xp = key_cache.shape
|
|
key_cache_v2 = key_cache.permute(0, 1, 3, 2, 4).reshape(B, H, bs, dx * xp)
|
|
elif key_cache.dim() == 4 and key_cache.shape[3] != query.shape[2]:
|
|
# K: [B, H, d, bs] → [B, H, bs, d]
|
|
key_cache_v2 = key_cache.permute(0, 1, 3, 2).contiguous()
|
|
else:
|
|
key_cache_v2 = key_cache
|
|
|
|
if value_cache.dim() == 4 and value_cache.shape[3] != query.shape[2]:
|
|
# V: [B, H, d, bs] → [B, H, bs, d]
|
|
value_cache_v2 = value_cache.permute(0, 1, 3, 2).contiguous()
|
|
else:
|
|
value_cache_v2 = value_cache
|
|
|
|
return ixf_F.vllm_single_query_cached_kv_attention_v2(
|
|
out,
|
|
max_num_partitions,
|
|
exp_sum,
|
|
max_logits,
|
|
tmp_out,
|
|
query,
|
|
key_cache_v2,
|
|
value_cache_v2,
|
|
head_mapping,
|
|
scale,
|
|
block_tables,
|
|
seq_lens,
|
|
block_size,
|
|
max_seq_len,
|
|
alibi_slopes,
|
|
)'''
|
|
|
|
if old_v2 in content:
|
|
content = content.replace(old_v2, new_v2, 1)
|
|
changes.append("V2: Replaced NotImplementedError with ixformer native kernel")
|
|
elif "raise NotImplementedError()" in content and "paged_attention_v2" in content:
|
|
print(" [warn] V2 NotImplementedError found but exact match failed")
|
|
else:
|
|
print(" [warn] V2 may already be patched")
|
|
|
|
with open(CUSTOM_OPS_PATH, "w") as f:
|
|
f.write(content)
|
|
|
|
for c in changes:
|
|
print(f" [ok] {c}")
|
|
|
|
|
|
def patch_triton_path():
|
|
"""Fix Triton import path so vllm can find it."""
|
|
# Triton is at /usr/local/lib/python3.10/site-packages/triton
|
|
# vllm checks for triton at import time in importing.py
|
|
triton_path = "/usr/local/lib/python3.10/site-packages"
|
|
importing_path = os.path.join(VLLM_ROOT, "importing.py")
|
|
|
|
if not os.path.exists(importing_path):
|
|
# Try to find it
|
|
for root, dirs, files in os.walk(VLLM_ROOT):
|
|
if "importing.py" in files:
|
|
importing_path = os.path.join(root, "importing.py")
|
|
break
|
|
|
|
if os.path.exists(importing_path):
|
|
with open(importing_path, "r") as f:
|
|
content = f.read()
|
|
|
|
if triton_path not in content and "Triton not installed" in content:
|
|
# Add sys.path insertion before the triton import check
|
|
old_import = "import triton"
|
|
new_import = f"import sys; sys.path.insert(0, '{triton_path}'); import triton"
|
|
if old_import in content:
|
|
content = content.replace(old_import, new_import, 1)
|
|
with open(importing_path, "w") as f:
|
|
f.write(content)
|
|
print(f" [ok] Triton path fix: added {triton_path} to sys.path")
|
|
else:
|
|
print(" [warn] Could not find 'import triton' in importing.py")
|
|
else:
|
|
print(" [skip] Triton path already fixed or not needed")
|
|
else:
|
|
print(f" [warn] importing.py not found at {importing_path}")
|
|
|
|
# Also create a symlink as backup
|
|
triton_src = "/usr/local/lib/python3.10/site-packages/triton"
|
|
triton_dst = "/usr/local/corex/lib64/python3/dist-packages/triton"
|
|
if os.path.exists(triton_src) and not os.path.exists(triton_dst):
|
|
try:
|
|
os.symlink(triton_src, triton_dst)
|
|
print(f" [ok] Symlinked triton → corex dist-packages")
|
|
except Exception as e:
|
|
print(f" [warn] Symlink failed: {e}")
|
|
|
|
|
|
def main():
|
|
print("=== patch_ixformer_native: Enable native V1/V2 kernels ===")
|
|
print()
|
|
|
|
print("--- Fix 1+2: _custom_ops.py V1 head_mapping + V2 native kernel ---")
|
|
patch_custom_ops()
|
|
|
|
print("\n--- Fix 3: Triton path ---")
|
|
patch_triton_path()
|
|
|
|
print("\n=== Summary ===")
|
|
print("V1: head_mapping int→Tensor conversion (fixes RuntimeError)")
|
|
print("V2: ixformer native kernel (replaces NotImplementedError)")
|
|
print(" ixf_F.vllm_single_query_cached_kv_attention_v2()")
|
|
print(" partition = max_num_partitions (PARTITION_SIZE=512)")
|
|
print("Triton: sys.path fix + symlink for vllm import")
|
|
print()
|
|
print("Done.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|