[CRITICAL] Fix V2 cache layout: V1=5D K, V2=4D K with transposed layout
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.
This commit is contained in:
@@ -144,6 +144,26 @@ def patch_custom_ops():
|
|||||||
_PARTITION_SIZE = 512
|
_PARTITION_SIZE = 512
|
||||||
max_num_partitions = (max_seq_len + _PARTITION_SIZE - 1) // _PARTITION_SIZE
|
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(
|
return ixf_F.vllm_single_query_cached_kv_attention_v2(
|
||||||
out,
|
out,
|
||||||
max_num_partitions,
|
max_num_partitions,
|
||||||
@@ -151,8 +171,8 @@ def patch_custom_ops():
|
|||||||
max_logits,
|
max_logits,
|
||||||
tmp_out,
|
tmp_out,
|
||||||
query,
|
query,
|
||||||
key_cache,
|
key_cache_v2,
|
||||||
value_cache,
|
value_cache_v2,
|
||||||
head_mapping,
|
head_mapping,
|
||||||
scale,
|
scale,
|
||||||
block_tables,
|
block_tables,
|
||||||
|
|||||||
Reference in New Issue
Block a user