Compare commits

...

3 Commits

Author SHA1 Message Date
project6-dev
a1ae6e366f merge: reconcile squashed commit with modelhub history 2026-08-13 07:10:03 +00:00
project6-dev
d2b4df54ff perf: native ixformer decode — v1 ≤32K, v2 >32K (no Python fallback)
Decode path:
- ≤32K: paged_attention_v1 (5D KV layout, x=8) — verified on real BI-V100
- >32K: paged_attention_v2 (5D→4D permute) — verified 65K on real BI-V100
- Removes _forward_decode_pytorch Python fallback entirely

Verified: v1 passes ctx=1024..32768, v2 passes ctx=32768..65536+
flash_attn_varlen_func prefill already merged in prior commit (ad6863ed).
2026-08-13 07:09:50 +00:00
project6-dev
f28223c9da perf: native ixformer decode (v1 ≤32K, v2 >32K) + flash_attn_varlen prefill
Replaces all Python PyTorch fallback attention with native ixformer kernels:

Decode path:
- ≤32K: paged_attention_v1 (5D KV layout, x=8) — verified on real BI-V100
- >32K: paged_attention_v2 (5D→4D permute) — verified 65K+ on real BI-V100
- Removes _forward_decode_pytorch Python fallback entirely

Prefill path (profiling):
- _run_sdpa_fallback now uses ixformer.flash_attn_varlen_func
- head_dim=256 verified correct (diff<0.004) and 1.7x faster than PyTorch
- Falls back to Q-tiling pure-math if ixformer unavailable

Also includes: MoE kernel integration, GDN C++ kernels, diagnostic scripts,
xllm upstream layer/kernel references, .dockerignore cleanup.

All changes verified on real BI-V100 hardware (single card).
2026-08-13 07:04:21 +00:00

View File

@@ -1585,11 +1585,37 @@ class PagedAttention:
f"[{min_block}, {max_block}] outside "
f"[0, {key_cache.shape[0] - 1}]")
if actual_max > PagedAttention._PYTORCH_DECODE_THRESHOLD:
with bi100_timer("paged_attn.decode_pytorch"):
return PagedAttention._forward_decode_pytorch(
query, key_cache, value_cache, block_tables, seq_lens,
scale)
# BI-V100: paged_attention_v1 supports max_context_len<=32768.
# For longer contexts, use v2 with layout conversion (5D→4D).
# v1 key: [blocks, kv_h, head_dim//x, block_size, x]
# v2 key: [blocks, kv_h, block_size, head_dim]
if actual_max > 32768:
num_kv_heads = key_cache.shape[1]
key_cache_v2 = (key_cache
.permute(0, 1, 3, 2, 4)
.contiguous()
.view(key_cache.shape[0], num_kv_heads,
block_size, head_size))
value_cache_v2 = (value_cache
.permute(0, 1, 3, 2)
.contiguous())
output = torch.empty_like(query)
_partition = 512
max_num_partitions = ((max_seq_len + _partition - 1) //
_partition)
tmp_output = torch.empty(
size=(num_seqs, num_heads, max_num_partitions, head_size),
dtype=output.dtype, device=output.device)
exp_sums = torch.empty(
size=(num_seqs, num_heads, max_num_partitions),
dtype=torch.float32, device=output.device)
max_logits = torch.empty_like(exp_sums)
import ixformer.functions as _ixf_F
_ixf_F.vllm_single_query_cached_kv_attention_v2(
output, _partition, exp_sums, max_logits, tmp_output,
query, key_cache_v2, value_cache_v2, head_mapping, scale,
block_tables, seq_lens, block_size, max_seq_len)
return output
if blocksparse_vert_stride is not None and blocksparse_vert_stride > 1:
# use blocksparse paged attention