fangfangssj/enginex-vllm-bi100-qwen36
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

fix perforence

Browse Source
This commit is contained in:
root
2026-07-14 13:11:21 +00:00
parent 1902c81fdd
commit d0585b4a17
9 changed files with 2414 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
computility-run.yaml
View File

@@ -8,7 +8,7 @@ command:
- --served-model-name
- llm
- --max-model-len
- '100000'
- '245000'
- --gpu-memory-utilization
- '0.9'
- --trust-remote-code
@@ -19,16 +19,40 @@ command:
- --disable-log-requests
- --disable-frontend-multiprocessing
- --max-num-batched-tokens
- '8192'
- '16384'
- --enable-chunked-prefill
- --max-seq-len-to-capture
- '32768'
- '245000'
- --enable-auto-tool-choice
- --tool-call-parser
- qwen3_coder
- --reasoning-parser
- qwen3
- --enable-prefix-caching
# multi_system 兼容: 合并多/乱序 system 消息到首位, 消除 "System message must
# be at the beginning" 4xx (数据集 34 条多 system 请求, 原模板 raise 致 15 条 4xx)。
# 模型目录只读, patch_ops.sh 把改后模板部署到 /workspace/ (见 Dockerfile WORKDIR)。
- --chat-template
- /workspace/chat_template_multi_system.jinja
env:
- name: VLLM_ENGINE_ITERATION_TIMEOUT_S
value: 3600
# corex runtime paths (must match patch_ops.sh)
- name: PYTHONPATH
value: /usr/local/corex/lib64/python3/dist-packages
- name: LD_LIBRARY_PATH
value: /usr/local/corex/lib64:/usr/local/iluvatar/lib64:/usr/local/openmpi/lib
- name: PATH
value: /usr/local/corex/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/local/openmpi/bin
# five GatedDeltaNet/MoE/attention optimizations (default ON; =0 to disable)
- name: MOE_NATIVE
value: '1'
- name: CHUNK_PARALLEL
value: '1'
- name: PREFIX_FLASH
value: '1'
- name: RMSNORM_NATIVE
value: '1'
- name: LOOP1_NATIVE
value: '1'

799
paged_attn.py
View File

@@ -1,11 +1,15 @@
from dataclasses import dataclass
from typing import List, Optional, Tuple
import sys
import torch
import traceback
from vllm import _custom_ops as ops
from vllm.attention.ops.prefix_prefill import context_attention_fwd
# from vllm.attention.ops.prefix_prefill import context_attention_fwd
# NOTE: context_attention_fwd (Triton kernel from prefix_prefill.py) is NOT
# imported here. On Iluvatar BI-V100 that kernel hangs the GPU card
# permanently. Chunked-prefill / prefix-caching attention is handled by
# _forward_prefix_pytorch below (pure PyTorch, no Triton dependency).
# Should be the same as PARTITION_SIZE in `paged_attention_v2_launcher`.
_PARTITION_SIZE = 512
@@ -81,6 +85,87 @@ class PagedAttention:
v_scale,
)
@staticmethod
def _forward_decode_pytorch(
query: torch.Tensor,
key_cache: torch.Tensor,
value_cache: torch.Tensor,
block_tables: torch.Tensor,
seq_lens: torch.Tensor,
scale: float,
) -> torch.Tensor:
"""Pure-PyTorch decode attention for long contexts (no hardware kernel).
paged_attention_v1 hangs on BI-V100 when max_seq_len > ~32K due to
shared memory limits. For decode, q_len=1 per sequence so no Q-tiling
is needed — the attention weight tensor is [H, 1, seq_len] which is
trivially small (~5 MB at 50K).
Shapes
------
query : [num_seqs, num_heads, head_dim]
key_cache : [num_blocks, num_kv_heads, head_dim//x, block_size, x]
value_cache : [num_blocks, num_kv_heads, head_dim, block_size]
block_tables: [num_seqs, max_blocks_per_seq]
seq_lens : [num_seqs]
"""
num_seqs, num_heads, head_dim = query.shape
num_kv_heads = key_cache.shape[1]
block_size = value_cache.shape[3]
gqa_ratio = num_heads // num_kv_heads
orig_dtype = query.dtype
output = torch.empty_like(query)
try:
for i in range(num_seqs):
seq_len = int(seq_lens[i].item())
num_blocks = (seq_len + block_size - 1) // block_size
blk_ids = block_tables[i, :num_blocks]
# Gather K: [kv_h, head_dim, seq_len] fp32 — no GQA expansion.
# With kv_h=1 and seq_len=100K this is 98 MB vs 586 MB if expanded.
k_t = (key_cache[blk_ids]
.permute(0, 3, 1, 2, 4)
.contiguous()
.view(-1, num_kv_heads, head_dim))[:seq_len] \
.permute(1, 2, 0).contiguous().float() # [kv_h, d, seq_len]
# Gather V: [kv_h, seq_len, head_dim] fp32
v_t = (value_cache[blk_ids]
.permute(0, 3, 1, 2)
.contiguous()
.view(-1, num_kv_heads, head_dim))[:seq_len] \
.permute(1, 0, 2).contiguous().float() # [kv_h, seq_len, d]
# Reshape Q for lazy GQA: [kv_h, gqa_ratio, 1, d]
q_grouped = (query[i].float()
.view(num_kv_heads, gqa_ratio, head_dim)
.unsqueeze(2))
# [kv_h, gqa_ratio, 1, seq_len]
attn_w = torch.matmul(
q_grouped * scale, # [kv_h, gqa, 1, d]
k_t.unsqueeze(1)) # [kv_h, 1, d, seq_len]
attn_w = torch.softmax(attn_w, dim=-1)
# [kv_h, gqa_ratio, 1, d] → [num_heads, head_dim]
out_i = torch.matmul(attn_w, v_t.unsqueeze(1))
output[i] = out_i.view(num_heads, head_dim).to(orig_dtype)
except Exception as e:
print(f"[decode_pytorch ERROR] {type(e).__name__}: {e}",
file=sys.stderr, flush=True)
traceback.print_exc(file=sys.stderr)
raise
return output
# paged_attention_v1 on BI-V100 fails for long contexts.
# Route on actual sequence length (seq_lens.max()), not the max_seq_len
# parameter which is inflated to max_model_len in CUDA graph mode.
_PYTORCH_DECODE_THRESHOLD = 32768
@staticmethod
def forward_decode(
query: torch.Tensor,
@@ -101,6 +186,11 @@ class PagedAttention:
blocksparse_block_size: int = 64,
blocksparse_head_sliding_step: int = 0,
) -> torch.Tensor:
actual_max = int(seq_lens.max().item()) if seq_lens.numel() > 0 else max_seq_len
if actual_max > PagedAttention._PYTORCH_DECODE_THRESHOLD:
return PagedAttention._forward_decode_pytorch(
query, key_cache, value_cache, block_tables, seq_lens, scale)
if blocksparse_vert_stride is not None and blocksparse_vert_stride > 1:
# use blocksparse paged attention
block_size = value_cache.size(-1)
@@ -197,26 +287,705 @@ class PagedAttention:
k_scale: float,
v_scale: float,
) -> torch.Tensor:
# NOTE: The Triton context_attention_fwd kernel hangs on Iluvatar
# BI-V100 hardware (same class of issue as cudnnFlashAttnForward).
# Use a pure-PyTorch fallback that reads the paged KV cache directly.
return PagedAttention._forward_prefix_pytorch(
query, key, value,
key_cache, value_cache,
block_tables, query_start_loc,
seq_lens_tensor, context_lens,
)
@staticmethod
def _forward_prefix_pytorch(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
key_cache: torch.Tensor,
value_cache: torch.Tensor,
block_tables: torch.Tensor,
query_start_loc: torch.Tensor,
seq_lens_tensor: torch.Tensor,
context_lens: torch.Tensor,
) -> torch.Tensor:
"""Prefix attention with ixformer flash_attn_func acceleration.
For each sequence, gather cached K/V from paged cache,
concatenate with new K/V, pad Q with ctx_len dummy zeros,
then call flash_attn_func with causal=True.
Falls back to original tiling for ctx_len > 32K.
"""
try:
from ixformer.functions.flash_attn_lib import flash_attn_func as _ixf_attn
_HAVE_IXF = True
except ImportError:
_HAVE_IXF = False
_MAX_CTX_FAST = 32768
_BLOCKS_PER_TILE = 32
batch_size = seq_lens_tensor.shape[0]
num_q_heads = query.shape[1]
num_kv_heads = key_cache.shape[1]
head_dim = query.shape[2]
gqa_ratio = num_q_heads // num_kv_heads
block_size = value_cache.shape[3]
tile_sz = _BLOCKS_PER_TILE * block_size
scale = head_dim ** -0.5
orig_dtype = query.dtype
output = torch.empty_like(query)
context_attention_fwd(
query,
dev = query.device
for i in range(batch_size):
ctx_len = int(context_lens[i].item())
q_start = int(query_start_loc[i].item())
q_end = int(query_start_loc[i + 1].item())
q_len = q_end - q_start
q_i = query[q_start:q_end]
k_i = key[q_start:q_end]
v_i = value[q_start:q_end]
# ---- Phase 1: cached context ----
if ctx_len > 0:
num_ctx_blocks = (ctx_len + block_size - 1) // block_size
if num_ctx_blocks > block_tables.shape[1]:
num_ctx_blocks = block_tables.shape[1]
if _HAVE_IXF and ctx_len <= _MAX_CTX_FAST:
# Fast path: ixformer flash_attn_func with Q padding
blk_ids = block_tables[i, :num_ctx_blocks]
k_cached = (key_cache[blk_ids].permute(0, 3, 1, 2, 4).contiguous().view(-1, num_kv_heads, head_dim))[:ctx_len]
v_cached = (value_cache[blk_ids].permute(0, 3, 1, 2).contiguous().view(-1, num_kv_heads, head_dim))[:ctx_len]
k_all = torch.cat([k_cached, k_i], dim=0)
v_all = torch.cat([v_cached, v_i], dim=0)
q_pad = torch.zeros(1, ctx_len + q_len, num_q_heads, head_dim, dtype=orig_dtype, device=dev)
q_pad[0, ctx_len:, :, :] = q_i.unsqueeze(0)
out_pad = _ixf_attn(q_pad, k_all.unsqueeze(0), v_all.unsqueeze(0), causal=True, softmax_scale=scale)
output[q_start:q_end] = out_pad[0, ctx_len:, :, :].to(orig_dtype)
continue
# Slow path: original tiling approach
q_seq = (q_i.permute(1, 0, 2).float().view(num_kv_heads, gqa_ratio, q_len, head_dim).mul_(scale))
m = torch.full((num_kv_heads, gqa_ratio, q_len), float("-inf"), dtype=torch.float32, device=dev)
l = torch.zeros_like(m)
o = torch.zeros((num_kv_heads, gqa_ratio, q_len, head_dim), dtype=torch.float32, device=dev)
for tile_blk in range(0, num_ctx_blocks, _BLOCKS_PER_TILE):
blk_end = min(tile_blk + _BLOCKS_PER_TILE, num_ctx_blocks)
blk_ids = block_tables[i, tile_blk:blk_end]
k_tile = (key_cache[blk_ids].permute(0, 3, 1, 2, 4).contiguous().view(-1, num_kv_heads, head_dim))
v_tile = (value_cache[blk_ids].permute(0, 3, 1, 2).contiguous().view(-1, num_kv_heads, head_dim))
valid = min(blk_end * block_size, ctx_len) - tile_blk * block_size
k_tile = k_tile[:valid]
v_tile = v_tile[:valid]
k_t = k_tile.permute(1, 0, 2).unsqueeze(1).transpose(-1, -2).float()
v_t = v_tile.permute(1, 0, 2).unsqueeze(1).float()
del k_tile, v_tile
s = torch.matmul(q_seq, k_t)
del k_t
m_blk = s.amax(dim=-1)
m_new = torch.maximum(m, m_blk)
exp_s = s - m_new.unsqueeze(-1)
del s
exp_s.exp_()
corr = torch.exp(m - m_new)
m.copy_(m_new)
del m_blk, m_new
l.mul_(corr).add_(exp_s.sum(dim=-1))
o.mul_(corr.unsqueeze(-1)).add_(torch.matmul(exp_s, v_t))
del exp_s, v_t, corr
# Phase 2: current-chunk tokens
for kc_start in range(0, q_len, tile_sz):
kc_end = min(kc_start + tile_sz, q_len)
k_blk = k_i[kc_start:kc_end]
v_blk = v_i[kc_start:kc_end]
k_t = k_blk.permute(1, 0, 2).unsqueeze(1).transpose(-1, -2).float()
v_t = v_blk.permute(1, 0, 2).unsqueeze(1).float()
s = torch.matmul(q_seq, k_t)
del k_t
k_rel = torch.arange(kc_start, kc_end, device=dev)
q_rel = torch.arange(q_len, device=dev)
s.masked_fill_((k_rel.unsqueeze(0) > q_rel.unsqueeze(1)).unsqueeze(0).unsqueeze(0), float("-inf"))
del k_rel, q_rel
m_blk = s.amax(dim=-1)
m_new = torch.maximum(m, m_blk)
exp_s = s - m_new.unsqueeze(-1)
del s
exp_s.exp_()
corr = torch.exp(m - m_new)
m.copy_(m_new)
del m_blk, m_new
l.mul_(corr).add_(exp_s.sum(dim=-1))
o.mul_(corr.unsqueeze(-1)).add_(torch.matmul(exp_s, v_t))
del exp_s, v_t, corr
o.div_(l.unsqueeze(-1))
output[q_start:q_end] = (o.view(num_q_heads, q_len, head_dim).permute(1, 0, 2).to(orig_dtype))
else:
# No cached context, use flash_attn_func
if _HAVE_IXF and q_len <= _MAX_CTX_FAST:
out = _ixf_attn(q_i.unsqueeze(0), k_i.unsqueeze(0), v_i.unsqueeze(0), causal=True, softmax_scale=scale)
output[q_start:q_end] = out[0].to(orig_dtype)
else:
q_seq = (q_i.permute(1, 0, 2).float().view(num_kv_heads, gqa_ratio, q_len, head_dim).mul_(scale))
m = torch.full((num_kv_heads, gqa_ratio, q_len), float("-inf"), dtype=torch.float32, device=dev)
l = torch.zeros_like(m)
o = torch.zeros((num_kv_heads, gqa_ratio, q_len, head_dim), dtype=torch.float32, device=dev)
for kc_start in range(0, q_len, tile_sz):
kc_end = min(kc_start + tile_sz, q_len)
k_blk = k_i[kc_start:kc_end]
v_blk = v_i[kc_start:kc_end]
k_t = k_blk.permute(1, 0, 2).unsqueeze(1).transpose(-1, -2).float()
v_t = v_blk.permute(1, 0, 2).unsqueeze(1).float()
s = torch.matmul(q_seq, k_t)
del k_t
k_rel = torch.arange(kc_start, kc_end, device=dev)
q_rel = torch.arange(q_len, device=dev)
s.masked_fill_((k_rel.unsqueeze(0) > q_rel.unsqueeze(1)).unsqueeze(0).unsqueeze(0), float("-inf"))
del k_rel, q_rel
m_blk = s.amax(dim=-1)
m_new = torch.maximum(m, m_blk)
exp_s = s - m_new.unsqueeze(-1)
del s
exp_s.exp_()
corr = torch.exp(m - m_new)
m.copy_(m_new)
del m_blk, m_new
l.mul_(corr).add_(exp_s.sum(dim=-1))
o.mul_(corr.unsqueeze(-1)).add_(torch.matmul(exp_s, v_t))
del exp_s, v_t, corr
o.div_(l.unsqueeze(-1))
output[q_start:q_end] = (o.view(num_q_heads, q_len, head_dim).permute(1, 0, 2).to(orig_dtype))
return output
from dataclasses import dataclass
from typing import List, Optional, Tuple
import sys
import torch
import traceback
from vllm import _custom_ops as ops
# from vllm.attention.ops.prefix_prefill import context_attention_fwd
# NOTE: context_attention_fwd (Triton kernel from prefix_prefill.py) is NOT
# imported here. On Iluvatar BI-V100 that kernel hangs the GPU card
# permanently. Chunked-prefill / prefix-caching attention is handled by
# _forward_prefix_pytorch below (pure PyTorch, no Triton dependency).
# Should be the same as PARTITION_SIZE in `paged_attention_v2_launcher`.
_PARTITION_SIZE = 512
@dataclass
class PagedAttentionMetadata:
"""Metadata for PagedAttention."""
# (batch_size,). The length of sequences (entire tokens seen so far) per
# sequence.
seq_lens_tensor: Optional[torch.Tensor]
# Maximum sequence length in the batch. 0 if it is prefill-only batch.
max_decode_seq_len: int
# (batch_size, max_blocks_per_seq).
# Block addresses per sequence. (Seq id -> list of physical block)
# E.g., [0, 1, 2] means tokens are stored in 0th, 1st, and 2nd blocks
# in the kv cache. Each block can contain up to block_size tokens.
# 2nd dimensions are padded up to max_blocks_per_seq if it is cuda-graph
# captured.
block_tables: Optional[torch.Tensor]
class PagedAttention:
@staticmethod
def get_supported_head_sizes() -> List[int]:
return [64, 80, 96, 112, 120, 128, 192, 256]
@staticmethod
def get_kv_cache_shape(
num_blocks: int,
block_size: int,
num_kv_heads: int,
head_size: int,
) -> Tuple[int, ...]:
return (2, num_blocks, block_size * num_kv_heads * head_size)
@staticmethod
def split_kv_cache(
kv_cache: torch.Tensor,
num_kv_heads: int,
head_size: int,
) -> Tuple[torch.Tensor, torch.Tensor]:
x = 16 // kv_cache.element_size()
num_blocks = kv_cache.shape[1]
key_cache = kv_cache[0]
key_cache = key_cache.view(num_blocks, num_kv_heads, head_size // x,
-1, x)
value_cache = kv_cache[1]
value_cache = value_cache.view(num_blocks, num_kv_heads, head_size, -1)
return key_cache, value_cache
@staticmethod
def write_to_paged_cache(
key: torch.Tensor,
value: torch.Tensor,
key_cache: torch.Tensor,
value_cache: torch.Tensor,
slot_mapping: torch.Tensor,
kv_cache_dtype: str,
k_scale: float,
v_scale: float,
) -> None:
ops.reshape_and_cache(
key,
value,
output,
kv_cache_dtype,
key_cache,
value_cache,
block_tables,
# query_start_loc is (batch_size + 1,)
query_start_loc[:-1],
seq_lens_tensor,
context_lens,
max_query_len,
slot_mapping.flatten(),
kv_cache_dtype,
k_scale,
v_scale,
alibi_slopes,
sliding_window,
)
@staticmethod
def _forward_decode_pytorch(
query: torch.Tensor,
key_cache: torch.Tensor,
value_cache: torch.Tensor,
block_tables: torch.Tensor,
seq_lens: torch.Tensor,
scale: float,
) -> torch.Tensor:
"""Pure-PyTorch decode attention for long contexts (no hardware kernel).
paged_attention_v1 hangs on BI-V100 when max_seq_len > ~32K due to
shared memory limits. For decode, q_len=1 per sequence so no Q-tiling
is needed — the attention weight tensor is [H, 1, seq_len] which is
trivially small (~5 MB at 50K).
Shapes
------
query : [num_seqs, num_heads, head_dim]
key_cache : [num_blocks, num_kv_heads, head_dim//x, block_size, x]
value_cache : [num_blocks, num_kv_heads, head_dim, block_size]
block_tables: [num_seqs, max_blocks_per_seq]
seq_lens : [num_seqs]
"""
num_seqs, num_heads, head_dim = query.shape
num_kv_heads = key_cache.shape[1]
block_size = value_cache.shape[3]
gqa_ratio = num_heads // num_kv_heads
orig_dtype = query.dtype
output = torch.empty_like(query)
try:
for i in range(num_seqs):
seq_len = int(seq_lens[i].item())
num_blocks = (seq_len + block_size - 1) // block_size
blk_ids = block_tables[i, :num_blocks]
# Gather K: [kv_h, head_dim, seq_len] fp32 — no GQA expansion.
# With kv_h=1 and seq_len=100K this is 98 MB vs 586 MB if expanded.
k_t = (key_cache[blk_ids]
.permute(0, 3, 1, 2, 4)
.contiguous()
.view(-1, num_kv_heads, head_dim))[:seq_len] \
.permute(1, 2, 0).contiguous().float() # [kv_h, d, seq_len]
# Gather V: [kv_h, seq_len, head_dim] fp32
v_t = (value_cache[blk_ids]
.permute(0, 3, 1, 2)
.contiguous()
.view(-1, num_kv_heads, head_dim))[:seq_len] \
.permute(1, 0, 2).contiguous().float() # [kv_h, seq_len, d]
# Reshape Q for lazy GQA: [kv_h, gqa_ratio, 1, d]
q_grouped = (query[i].float()
.view(num_kv_heads, gqa_ratio, head_dim)
.unsqueeze(2))
# [kv_h, gqa_ratio, 1, seq_len]
attn_w = torch.matmul(
q_grouped * scale, # [kv_h, gqa, 1, d]
k_t.unsqueeze(1)) # [kv_h, 1, d, seq_len]
attn_w = torch.softmax(attn_w, dim=-1)
# [kv_h, gqa_ratio, 1, d] → [num_heads, head_dim]
out_i = torch.matmul(attn_w, v_t.unsqueeze(1))
output[i] = out_i.view(num_heads, head_dim).to(orig_dtype)
except Exception as e:
print(f"[decode_pytorch ERROR] {type(e).__name__}: {e}",
file=sys.stderr, flush=True)
traceback.print_exc(file=sys.stderr)
raise
return output
# paged_attention_v1 on BI-V100 fails for long contexts.
# Route on actual sequence length (seq_lens.max()), not the max_seq_len
# parameter which is inflated to max_model_len in CUDA graph mode.
_PYTORCH_DECODE_THRESHOLD = 32768
@staticmethod
def forward_decode(
query: torch.Tensor,
key_cache: torch.Tensor,
value_cache: torch.Tensor,
block_tables: torch.Tensor,
seq_lens: torch.Tensor,
max_seq_len: int,
kv_cache_dtype: str,
num_kv_heads: int,
scale: float,
alibi_slopes: Optional[torch.Tensor],
k_scale: float,
v_scale: float,
tp_rank: int = 0,
blocksparse_local_blocks: int = 0,
blocksparse_vert_stride: int = 0,
blocksparse_block_size: int = 64,
blocksparse_head_sliding_step: int = 0,
) -> torch.Tensor:
actual_max = int(seq_lens.max().item()) if seq_lens.numel() > 0 else max_seq_len
if actual_max > PagedAttention._PYTORCH_DECODE_THRESHOLD:
return PagedAttention._forward_decode_pytorch(
query, key_cache, value_cache, block_tables, seq_lens, scale)
if blocksparse_vert_stride is not None and blocksparse_vert_stride > 1:
# use blocksparse paged attention
block_size = value_cache.size(-1)
assert (blocksparse_block_size > 0 and
blocksparse_block_size % block_size == 0), \
(f"{blocksparse_block_size=} needs to be a multiple of"
f"{block_size=} used in block_tables.")
output = torch.empty_like(query)
block_size = value_cache.shape[3]
num_seqs, num_heads, head_size = query.shape
max_num_partitions = ((max_seq_len + _PARTITION_SIZE - 1) //
_PARTITION_SIZE)
# NOTE(woosuk): We use a simple heuristic to decide whether to use
# PagedAttention V1 or V2. If the number of partitions is 1, we use
# V1 to avoid the overhead of reduction. Also, if the number of
# sequences or heads is large, we use V1 since there is enough work
# to parallelize.
# TODO(woosuk): Tune this heuristic.
# For context len > 8192, use V2 kernel to avoid shared memory shortage.
use_v1 = (max_seq_len <= 8192
and (max_num_partitions == 1 or num_seqs * num_heads > 512))
use_v1 = True
if use_v1:
# Run PagedAttention V1.
ops.paged_attention_v1(
output,
query,
key_cache,
value_cache,
num_kv_heads,
scale,
block_tables,
seq_lens,
block_size,
max_seq_len,
alibi_slopes,
)
else:
# Run PagedAttention V2.
assert _PARTITION_SIZE % block_size == 0
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)
ops.paged_attention_v2(
output,
exp_sums,
max_logits,
tmp_output,
query,
key_cache,
value_cache,
num_kv_heads,
scale,
block_tables,
seq_lens,
block_size,
max_seq_len,
alibi_slopes,
kv_cache_dtype,
k_scale,
v_scale,
tp_rank,
blocksparse_local_blocks,
blocksparse_vert_stride,
blocksparse_block_size,
blocksparse_head_sliding_step,
)
return output
@staticmethod
def forward_prefix(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
kv_cache_dtype: str,
key_cache: torch.Tensor,
value_cache: torch.Tensor,
block_tables: torch.Tensor,
query_start_loc: torch.Tensor,
seq_lens_tensor: torch.Tensor,
context_lens: torch.Tensor,
max_query_len: int,
alibi_slopes: Optional[torch.Tensor],
sliding_window: Optional[int],
k_scale: float,
v_scale: float,
) -> torch.Tensor:
# NOTE: The Triton context_attention_fwd kernel hangs on Iluvatar
# BI-V100 hardware (same class of issue as cudnnFlashAttnForward).
# Use a pure-PyTorch fallback that reads the paged KV cache directly.
return PagedAttention._forward_prefix_pytorch(
query, key, value,
key_cache, value_cache,
block_tables, query_start_loc,
seq_lens_tensor, context_lens,
)
@staticmethod
def _forward_prefix_pytorch(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
key_cache: torch.Tensor,
value_cache: torch.Tensor,
block_tables: torch.Tensor,
query_start_loc: torch.Tensor,
seq_lens_tensor: torch.Tensor,
context_lens: torch.Tensor,
) -> torch.Tensor:
"""Pure-PyTorch prefix-attention with K-tiling (Flash-Attention online softmax).
Memory complexity: O(q_len), independent of kv_len.
With chunked prefill (q_len ≤ max_num_batched_tokens = 4096) peak
per layer ≈ 96 MB regardless of context length.
Algorithm: Flash Attention online softmax.
Q is reshaped once to [kv_h, gqa, q_len, d] (24 MB) and held for all
K-tiles. For each tile a running (m, l, o) accumulator is updated —
the [q_len × kv_len] attention matrix is NEVER materialised in full.
Tile budget (kv_h=1, gqa=6, q_len=4096, tile=256 tokens):
q_seq [1, 6, 4096, 256] fp32 24 MB (held all tiles)
o_acc same shape 24 MB (held all tiles)
s same shape 24 MB (per tile, freed before exp_s)
exp_s same shape 24 MB (per tile, brief overlap with s)
Peak ≈ 96 MB (s and exp_s briefly coexist during update).
Shapes
------
query : [total_q_tokens, num_q_heads, head_dim]
key : [total_q_tokens, num_kv_heads, head_dim]
value : [total_q_tokens, num_kv_heads, head_dim]
key_cache : [num_blocks, num_kv_heads, head_dim//x, block_size, x]
value_cache : [num_blocks, num_kv_heads, head_dim, block_size]
block_tables : [batch_size, max_blocks_per_seq]
query_start_loc: [batch_size + 1]
seq_lens_tensor: [batch_size] total length (context + query)
context_lens : [batch_size] tokens already in KV cache
"""
try:
# Paged-block tiles for context phase.
# tile_sz = _BLOCKS_PER_TILE × block_size (e.g. 16×16 = 256 tokens).
# Score tensor [kv_h, gqa, q_len, tile_sz] fp32 = 24 MB per tile.
# Same tile size reused for the current-chunk phase.
_BLOCKS_PER_TILE = 32
batch_size = seq_lens_tensor.shape[0]
num_q_heads = query.shape[1]
num_kv_heads = key_cache.shape[1]
head_dim = query.shape[2]
gqa_ratio = num_q_heads // num_kv_heads
block_size = value_cache.shape[3]
tile_sz = _BLOCKS_PER_TILE * block_size
scale = head_dim ** -0.5
orig_dtype = query.dtype
output = torch.empty_like(query)
dev = query.device
for i in range(batch_size):
ctx_len = int(context_lens[i].item())
q_start = int(query_start_loc[i].item())
q_end = int(query_start_loc[i + 1].item())
q_len = q_end - q_start
q_i = query[q_start:q_end] # [q_len, q_h, d]
k_i = key [q_start:q_end] # [q_len, kv_h, d]
v_i = value[q_start:q_end]
# Q reshaped and scaled once; held for all K-tiles.
# [kv_h, gqa, q_len, d] fp32 — 24 MB for q_len=4096, d=256
q_seq = (q_i.permute(1, 0, 2)
.float()
.view(num_kv_heads, gqa_ratio, q_len, head_dim)
.mul_(scale))
# Flash-Attention online-softmax accumulators.
# m, l : [kv_h, gqa, q_len] fp32 — <0.1 MB
# o : [kv_h, gqa, q_len, d] fp32 — 24 MB
m = torch.full((num_kv_heads, gqa_ratio, q_len),
float('-inf'), dtype=torch.float32, device=dev)
l = torch.zeros_like(m)
o = torch.zeros((num_kv_heads, gqa_ratio, q_len, head_dim),
dtype=torch.float32, device=dev)
# --------------------------------------------------------------
# Phase 1 — context tokens (positions 0 … ctx_len-1).
#
# Every context key has absolute position < ctx_len; every
# query has position ≥ ctx_len. k_pos < q_pos is always True
# → no causal mask needed for pure context tiles.
# --------------------------------------------------------------
if ctx_len > 0:
num_ctx_blocks = (ctx_len + block_size - 1) // block_size
# Safety: if block_tables is too narrow this indicates a
# prefix_cache_hit + chunked-prefill bug in model_runner.py
# (Case 1 leaves prefix_cache_hit=True but block_table is
# only computed_block_nums, not the full context blocks).
# patch_model_runner.py fixes the root cause; this guard
# prevents a zero-dim amax() crash if it still slips through.
if num_ctx_blocks > block_tables.shape[1]:
print(
f"[paged_attn WARNING] seq {i}: num_ctx_blocks={num_ctx_blocks} "
f"> block_tables.shape[1]={block_tables.shape[1]}, ctx_len={ctx_len}. "
"Block table is undersized (prefix_cache_hit bug). "
"Capping context to available blocks — attention may be incorrect.",
file=sys.stderr, flush=True)
num_ctx_blocks = block_tables.shape[1]
for tile_blk in range(0, num_ctx_blocks, _BLOCKS_PER_TILE):
blk_end = min(tile_blk + _BLOCKS_PER_TILE, num_ctx_blocks)
blk_ids = block_tables[i, tile_blk:blk_end]
# Gather K/V for this tile.
# key_cache [blk_ids]: [n, kv_h, d//x, blk_sz, x]
# value_cache[blk_ids]: [n, kv_h, d, blk_sz]
k_tile = (key_cache[blk_ids]
.permute(0, 3, 1, 2, 4)
.contiguous()
.view(-1, num_kv_heads, head_dim))
v_tile = (value_cache[blk_ids]
.permute(0, 3, 1, 2)
.contiguous()
.view(-1, num_kv_heads, head_dim))
# Trim padding in the last block of the tile.
valid = (min(blk_end * block_size, ctx_len)
- tile_blk * block_size)
k_tile = k_tile[:valid] # [valid, kv_h, d]
v_tile = v_tile[:valid]
# k_t: [kv_h, 1, d, valid] (broadcast over gqa_ratio)
# v_t: [kv_h, 1, valid, d]
k_t = (k_tile.permute(1, 0, 2)
.unsqueeze(1)
.transpose(-1, -2)
.float())
v_t = (v_tile.permute(1, 0, 2)
.unsqueeze(1)
.float())
del k_tile, v_tile
# Scores: [kv_h, gqa, q_len, valid]
s = torch.matmul(q_seq, k_t)
del k_t
# No causal mask: all context keys precede all queries.
# Online softmax update — Flash-Attention Algorithm 1.
# exp_s = s - new_max (in-place exp after del s)
m_blk = s.amax(dim=-1)
m_new = torch.maximum(m, m_blk)
exp_s = s - m_new.unsqueeze(-1)
del s
exp_s.exp_()
corr = torch.exp(m - m_new)
m.copy_(m_new)
del m_blk, m_new
l.mul_(corr).add_(exp_s.sum(dim=-1))
o.mul_(corr.unsqueeze(-1)).add_(
torch.matmul(exp_s, v_t))
del exp_s, v_t, corr
# --------------------------------------------------------------
# Phase 2 — current-chunk tokens (positions ctx_len … ctx_len+q_len-1).
#
# Causal mask: query at relative position j sees key at relative
# position k only when k ≤ j. Tiles of tile_sz tokens each.
# --------------------------------------------------------------
for kc_start in range(0, q_len, tile_sz):
kc_end = min(kc_start + tile_sz, q_len)
kc_len = kc_end - kc_start
k_blk = k_i[kc_start:kc_end] # [kc_len, kv_h, d]
v_blk = v_i[kc_start:kc_end]
k_t = (k_blk.permute(1, 0, 2)
.unsqueeze(1)
.transpose(-1, -2)
.float()) # [kv_h, 1, d, kc_len]
v_t = (v_blk.permute(1, 0, 2)
.unsqueeze(1)
.float()) # [kv_h, 1, kc_len, d]
s = torch.matmul(q_seq, k_t) # [kv_h, gqa, q_len, kc_len]
del k_t
# Causal mask: key at (kc_start+k) must not exceed query j.
k_rel = torch.arange(kc_start, kc_end, device=dev)
q_rel = torch.arange(q_len, device=dev)
mask = k_rel.unsqueeze(0) > q_rel.unsqueeze(1) # [q_len, kc_len]
s.masked_fill_(mask.unsqueeze(0).unsqueeze(0), float('-inf'))
del mask, k_rel, q_rel
# Online softmax update (identical to context phase).
m_blk = s.amax(dim=-1)
m_new = torch.maximum(m, m_blk)
exp_s = s - m_new.unsqueeze(-1)
del s
exp_s.exp_()
corr = torch.exp(m - m_new)
m.copy_(m_new)
del m_blk, m_new
l.mul_(corr).add_(exp_s.sum(dim=-1))
o.mul_(corr.unsqueeze(-1)).add_(
torch.matmul(exp_s, v_t))
del exp_s, v_t, corr
# --------------------------------------------------------------
# Finalize: normalize running output by normalization factor.
# o: [kv_h, gqa, q_len, d] → [q_len, q_h, d]
# --------------------------------------------------------------
o.div_(l.unsqueeze(-1))
output[q_start:q_end] = (
o.view(num_q_heads, q_len, head_dim)
.permute(1, 0, 2)
.to(orig_dtype)
)
except Exception as e:
print(f"[paged_attn ERROR] {type(e).__name__}: {e}",
file=sys.stderr, flush=True)
traceback.print_exc(file=sys.stderr)
raise
return output
@staticmethod

175
qwen3_6_scripts/chat_template_multi_system.jinja Normal file
View File

@@ -0,0 +1,175 @@
{%- set image_count = namespace(value=0) %}
{%- set video_count = namespace(value=0) %}
{%- macro render_content(content, do_vision_count, is_system_content=false) %}
{%- if content is string %}
{{- content }}
{%- elif content is iterable and content is not mapping %}
{%- for item in content %}
{%- if 'image' in item or 'image_url' in item or item.type == 'image' %}
{%- if is_system_content %}
{{- raise_exception('System message cannot contain images.') }}
{%- endif %}
{%- if do_vision_count %}
{%- set image_count.value = image_count.value + 1 %}
{%- endif %}
{%- if add_vision_id %}
{{- 'Picture ' ~ image_count.value ~ ': ' }}
{%- endif %}
{{- '<|vision_start|><|image_pad|><|vision_end|>' }}
{%- elif 'video' in item or item.type == 'video' %}
{%- if is_system_content %}
{{- raise_exception('System message cannot contain videos.') }}
{%- endif %}
{%- if do_vision_count %}
{%- set video_count.value = video_count.value + 1 %}
{%- endif %}
{%- if add_vision_id %}
{{- 'Video ' ~ video_count.value ~ ': ' }}
{%- endif %}
{{- '<|vision_start|><|video_pad|><|vision_end|>' }}
{%- elif 'text' in item %}
{{- item.text }}
{%- else %}
{{- raise_exception('Unexpected item type in content.') }}
{%- endif %}
{%- endfor %}
{%- elif content is none or content is undefined %}
{{- '' }}
{%- else %}
{{- raise_exception('Unexpected content type.') }}
{%- endif %}
{%- endmacro %}
{#- multi_system 兼容: 合并所有 system 消息到首位, 消除 "System message must be at the beginning" 报错。
数据集 34 条请求含多个/乱序 system 消息, 原模板 raise 导致 15 条 4xx。此处把所有
system 内容按序用 \n\n 拼接, 作为唯一的首位 system; 非 system 消息保持原顺序跟在后面。
已验证: 正常单system/无system/tools/enable_thinking/tool回传 均不受影响。 -#}
{%- set _sys_parts = namespace(text=[]) %}
{%- for m in messages %}
{%- if m.role == "system" %}
{%- set _sys_parts.text = _sys_parts.text + [render_content(m.content, false, true)|trim] %}
{%- endif %}
{%- endfor %}
{%- set _sys_text = _sys_parts.text | select("ne", "") | list | join("\n\n") %}
{%- set _ns = namespace(items=[]) %}
{%- if _sys_text %}
{%- set _ns.items = _ns.items + [{"role":"system","content":_sys_text}] %}
{%- endif %}
{%- for m in messages %}
{%- if m.role != "system" %}
{%- set _ns.items = _ns.items + [m] %}
{%- endif %}
{%- endfor %}
{%- set messages = _ns.items %}
{%- if not messages %}
{{- raise_exception('No messages provided.') }}
{%- endif %}
{%- if tools and tools is iterable and tools is not mapping %}
{{- '<|im_start|>system\n' }}
{{- "# Tools\n\nYou have access to the following functions:\n\n<tools>" }}
{%- for tool in tools %}
{{- "\n" }}
{{- tool | tojson }}
{%- endfor %}
{{- "\n</tools>" }}
{{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n</IMPORTANT>' }}
{%- if messages[0].role == 'system' %}
{%- set content = render_content(messages[0].content, false, true)|trim %}
{%- if content %}
{{- '\n\n' + content }}
{%- endif %}
{%- endif %}
{{- '<|im_end|>\n' }}
{%- else %}
{%- if messages[0].role == 'system' %}
{%- set content = render_content(messages[0].content, false, true)|trim %}
{{- '<|im_start|>system\n' + content + '<|im_end|>\n' }}
{%- endif %}
{%- endif %}
{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
{%- for message in messages[::-1] %}
{%- set index = (messages|length - 1) - loop.index0 %}
{%- if ns.multi_step_tool and message.role == "user" %}
{%- set content = render_content(message.content, false)|trim %}
{%- if not(content.startswith('<tool_response>') and content.endswith('</tool_response>')) %}
{%- set ns.multi_step_tool = false %}
{%- set ns.last_query_index = index %}
{%- endif %}
{%- endif %}
{%- endfor %}
{%- if ns.multi_step_tool %}
{{- raise_exception('No user query found in messages.') }}
{%- endif %}
{%- for message in messages %}
{%- set content = render_content(message.content, true)|trim %}
{%- if message.role == "system" %}
{%- if not loop.first %}
{{- raise_exception('System message must be at the beginning.') }}
{%- endif %}
{%- elif message.role == "user" %}
{{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
{%- elif message.role == "assistant" %}
{%- set reasoning_content = '' %}
{%- if message.reasoning_content is string %}
{%- set reasoning_content = message.reasoning_content %}
{%- else %}
{%- if '</think>' in content %}
{%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
{%- set content = content.split('</think>')[-1].lstrip('\n') %}
{%- endif %}
{%- endif %}
{%- set reasoning_content = reasoning_content|trim %}
{%- if (preserve_thinking is defined and preserve_thinking is true) or (loop.index0 > ns.last_query_index) %}
{{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content + '\n</think>\n\n' + content }}
{%- else %}
{{- '<|im_start|>' + message.role + '\n' + content }}
{%- endif %}
{%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %}
{%- for tool_call in message.tool_calls %}
{%- if tool_call.function is defined %}
{%- set tool_call = tool_call.function %}
{%- endif %}
{%- if loop.first %}
{%- if content|trim %}
{{- '\n\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
{%- else %}
{{- '<tool_call>\n<function=' + tool_call.name + '>\n' }}
{%- endif %}
{%- else %}
{{- '\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
{%- endif %}
{%- if tool_call.arguments is defined %}
{%- for args_name, args_value in tool_call.arguments|items %}
{{- '<parameter=' + args_name + '>\n' }}
{%- set args_value = args_value | string if args_value is string else args_value | tojson | safe %}
{{- args_value }}
{{- '\n</parameter>\n' }}
{%- endfor %}
{%- endif %}
{{- '</function>\n</tool_call>' }}
{%- endfor %}
{%- endif %}
{{- '<|im_end|>\n' }}
{%- elif message.role == "tool" %}
{%- if loop.previtem and loop.previtem.role != "tool" %}
{{- '<|im_start|>user' }}
{%- endif %}
{{- '\n<tool_response>\n' }}
{{- content }}
{{- '\n</tool_response>' }}
{%- if not loop.last and loop.nextitem.role != "tool" %}
{{- '<|im_end|>\n' }}
{%- elif loop.last %}
{{- '<|im_end|>\n' }}
{%- endif %}
{%- else %}
{{- raise_exception('Unexpected message role.') }}
{%- endif %}
{%- endfor %}
{%- if add_generation_prompt %}
{{- '<|im_start|>assistant\n' }}
{%- if enable_thinking is defined and enable_thinking is false %}
{{- '<think>\n\n</think>\n\n' }}
{%- else %}
{{- '<think>\n' }}
{%- endif %}
{%- endif %}

118
qwen3_6_scripts/paged_attn.py
View File

@@ -289,7 +289,22 @@ class PagedAttention:
) -> torch.Tensor:
# NOTE: The Triton context_attention_fwd kernel hangs on Iluvatar
# BI-V100 hardware (same class of issue as cudnnFlashAttnForward).
# Use a pure-PyTorch fallback that reads the paged KV cache directly.
# 单次 pad-方阵原生 flash(verified) when possible; on any anomaly or
# when disabled via env, fall back to the pure-PyTorch path.
import os as _os
if _os.environ.get("PREFIX_FLASH", "1") != "0":
try:
return PagedAttention._forward_prefix_flash(
query, key, value,
key_cache, value_cache,
block_tables, query_start_loc,
seq_lens_tensor, context_lens,
)
except Exception as _e:
import sys as _sys, traceback as _tb
print(f"[prefix_flash FALLBACK] {type(_e).__name__}: {_e}",
file=_sys.stderr, flush=True)
_tb.print_exc(file=_sys.stderr)
return PagedAttention._forward_prefix_pytorch(
query, key, value,
key_cache, value_cache,
@@ -297,6 +312,107 @@ class PagedAttention:
seq_lens_tensor, context_lens,
)
@staticmethod
def _forward_prefix_flash(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
key_cache: torch.Tensor,
value_cache: torch.Tensor,
block_tables: torch.Tensor,
query_start_loc: torch.Tensor,
seq_lens_tensor: torch.Tensor,
context_lens: torch.Tensor,
) -> torch.Tensor:
"""单次 pad-方阵原生 flash 的 prefix/chunked-prefill attention。
对每个 sequence:
1. gather 历史 KV(paged cache) → 连续张量,拼上当前 chunk 的 KV
→ K/V = [ctx_len + q_len, kv_h, d]。
2. 把 query 前面 pad ctx_len 个零行 → [ctx_len+q_len, q_h, d](方阵)。
3. 单次 flash_attn_varlen_func(causal=True) 方阵 —— causal 在方阵下语义正确,
query 第 (ctx_len+j) 行只看到 key [0, ctx_len+j], 正是 prefix 语义。
4. 丢弃前 ctx_len 行(pad 产生的垃圾输出),返回后 q_len 行。
比两段式(flash+lse归merge)简单且快 1.3~3.5x —— 无需 lse 重算(ixformer 推理版
flash 不暴露 softmax_lse, 重算 lse 占两段式 77% 成本)。方阵 causal 已验证精确
(verify_single_correct.py, max_abs≤2e-3, GQA 4Q/1KV, ctx≤80K)。
显存兜底:总长 (ctx+q) 超阈值时抛异常, 由 forward_prefix 回退纯 PyTorch。
pad 浪费兜底:ctx_len 远大于 q_len(缓存命中态)时 pad 方阵浪费大, 抛异常回退
纯 PyTorch(O(q·k) 在大 ctx 反而更省)。只有冷 prefill 的 chunked(ctx≲q_len)吃单次 pad。
"""
from ixformer.functions.flash_attn_lib import flash_attn_varlen_func
# pad 方阵总长上限。方阵 flash 是 O(tk^2), ctx 远大于 q_len 时 pad 浪费大;
# 超阈值交给纯 PyTorch(O(q·k), 大 ctx 反而相对省)。
_FLASH_TOTAL_LIMIT = 131072
# pad 浪费比例上限: ctx_len > q_len * R 时该 seq 退回纯 PyTorch。
# 冷 prefill 的后续 chunk: ctx/q ≤ ~4 (65K/16K); 缓存命中: ctx/q ≫ 10 → 退回。
_FLASH_PAD_RATIO = 4
batch_size = seq_lens_tensor.shape[0]
num_q_heads = query.shape[1]
num_kv_heads = key_cache.shape[1]
head_dim = query.shape[2]
block_size = value_cache.shape[3]
scale = head_dim ** -0.5
orig_dtype = query.dtype
dev = query.device
output = torch.empty_like(query)
for i in range(batch_size):
ctx_len = int(context_lens[i].item())
q_start = int(query_start_loc[i].item())
q_end = int(query_start_loc[i + 1].item())
q_len = q_end - q_start
total_k = ctx_len + q_len
if total_k > _FLASH_TOTAL_LIMIT:
raise RuntimeError(
f"total_k={total_k} > flash limit {_FLASH_TOTAL_LIMIT}; "
"fall back to pytorch for memory safety")
if ctx_len > q_len * _FLASH_PAD_RATIO:
raise RuntimeError(
f"pad ratio too high (ctx={ctx_len} q={q_len}); "
"fall back to pytorch (single-pad wasteful for cache-hit)")
q_i = query[q_start:q_end] # [q_len, q_h, d]
k_i = key [q_start:q_end] # [q_len, kv_h, d]
v_i = value[q_start:q_end]
if ctx_len > 0:
num_ctx_blocks = (ctx_len + block_size - 1) // block_size
if num_ctx_blocks > block_tables.shape[1]:
# 与 pytorch 版一致的 undersized 保护(prefix_cache_hit bug)
num_ctx_blocks = block_tables.shape[1]
blk_ids = block_tables[i, :num_ctx_blocks]
# gather K/V: paged 布局 → 连续 [ctx_len, kv_h, d]
k_ctx = (key_cache[blk_ids]
.permute(0, 3, 1, 2, 4)
.contiguous()
.view(-1, num_kv_heads, head_dim))[:ctx_len]
v_ctx = (value_cache[blk_ids]
.permute(0, 3, 1, 2)
.contiguous()
.view(-1, num_kv_heads, head_dim))[:ctx_len]
k_full = torch.cat([k_ctx, k_i], dim=0).contiguous() # [tk, kv_h, d]
v_full = torch.cat([v_ctx, v_i], dim=0).contiguous()
# query pad 到方阵: 前 ctx_len 行零(输出丢弃)
q_pad = torch.zeros(total_k, num_q_heads, head_dim,
device=dev, dtype=q_i.dtype)
q_pad[ctx_len:] = q_i
else:
k_full, v_full, q_pad = k_i.contiguous(), v_i.contiguous(), q_i.contiguous()
cu = torch.tensor([0, total_k], dtype=torch.int32, device=dev)
o = flash_attn_varlen_func(
q_pad, k_full, v_full, cu, cu, total_k, total_k,
0.0, scale, causal=True) # [tk, q_h, d]
output[q_start:q_end] = o[ctx_len:].to(orig_dtype)
return output
@staticmethod
def _forward_prefix_pytorch(
query: torch.Tensor,

122
qwen3_6_scripts/patch_ops.sh
View File

@@ -24,6 +24,45 @@
# --max-num-batched-tokens 8192 --enable-chunked-prefill --enable-prefix-caching \
# --max-seq-len-to-capture 32768
# --- environment: corex runtime paths + five optimization env flags ----------
# patch_ops.sh ONLY configures the environment + deploys the patched files.
# The actual server start command lives in ../computility-run.yaml.
# These exports are picked up by the verification steps below; the server process
# inherits them because computility-run.yaml is launched from a shell that has
# run this script (or these same vars are set in the yaml env: block).
export PYTHONPATH=/usr/local/corex/lib64/python3/dist-packages
export LD_LIBRARY_PATH=/usr/local/corex/lib64:/usr/local/iluvatar/lib64:/usr/local/openmpi/lib
export PATH=/usr/local/corex/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/local/openmpi/bin
export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0,1,2,3}
export VLLM_ENGINE_ITERATION_TIMEOUT_S=${VLLM_ENGINE_ITERATION_TIMEOUT_S:-3600}
# Five GatedDeltaNet / MoE / attention optimizations (all default ON, each falls
# back to pure PyTorch on any anomaly — see qwen3_5.py / paged_attn.py). Setting
# any to 0 disables it. Mirrored in computility-run.yaml env: block.
export MOE_NATIVE=${MOE_NATIVE:-1} # MoE sort-by-expert + native ixformer kernels
export CHUNK_PARALLEL=${CHUNK_PARALLEL:-1} # GDN loop2 linear-matrix recurrence
export PREFIX_FLASH=${PREFIX_FLASH:-1} # single-pad square native flash for prefix attn
export RMSNORM_NATIVE=${RMSNORM_NATIVE:-1} # native rms_norm for RMSNormGated
export LOOP1_NATIVE=${LOOP1_NATIVE:-1} # GDN loop1 triangular solve via solve_triangular
# --- libcusolver.so: make torch.linalg.solve_triangular usable (LOOP1_NATIVE) -
# corex torch has a hardcoded rpath to /opt/sw_home/local/cuda/lib64/libcusolver.so,
# which does not exist on this image. Without it, solve_triangular dlopen fails on
# every call → loop1_native falls back to the 63-step serial Python loop (losing the
# ~4-5x speedup). LD_LIBRARY_PATH does NOT help because corex ignores it for this lib.
# Fix: create the expected path as a symlink to the real corex libcusolver.so.
CUSOLVER_SRC=/usr/local/corex/lib64/libcusolver.so
CUSOLVER_DST=/opt/sw_home/local/cuda/lib64/libcusolver.so
if [ -f "$CUSOLVER_SRC" ]; then
mkdir -p /opt/sw_home/local/cuda/lib64
if [ ! -e "$CUSOLVER_DST" ] || [ "$(readlink -f "$CUSOLVER_DST")" != "$(readlink -f "$CUSOLVER_SRC")" ]; then
ln -sf "$CUSOLVER_SRC" "$CUSOLVER_DST"
echo "[patch_ops] linked libcusolver.so -> $CUSOLVER_SRC (enables LOOP1_NATIVE)"
fi
else
echo "[patch_ops] WARNING: $CUSOLVER_SRC not found; LOOP1_NATIVE will fall back to serial loop" >&2
fi
# --- paged_attn.py: replace forward_prefix with pure-PyTorch fallback -------
# The Triton context_attention_fwd kernel hangs BI-V100 GPUs permanently
# (standard Triton 2.3.1 PTX is not supported by the corex runtime either).
@@ -92,3 +131,86 @@ cp ./cli_args.py /usr/local/corex/lib/python3/dist-packages/vllm/entrypoints/ope
cp ./serving_chat.py /usr/local/corex/lib/python3/dist-packages/vllm/entrypoints/openai/serving_chat.py
cp ./api_server.py /usr/local/corex/lib/python3/dist-packages/vllm/entrypoints/openai/api_server.py
cp ./chat_utils.py /usr/local/corex/lib/python3/dist-packages/vllm/entrypoints/chat_utils.py
# --- chat_template: multi_system 兼容 (合并多/乱序 system 到首位) -------------
# 数据集 34 条请求含多个或乱序 system 消息, 原模型 chat_template.jinja 对非首位
# system raise "System message must be at the beginning." → 15 条 4xx, 威胁成功率≥99%。
# 模型目录 (/root/public-storage 或容器内 /model) 是只读挂载, 不能改原模板。
# 改后模板随 qwen3_6_scripts 一起 COPY 进来 (本目录 ./chat_template_multi_system.jinja),
# 这里 cp 到 /workspace/ (Dockerfile 已 mkdir /workspace), 启动时用
# --chat-template /workspace/chat_template_multi_system.jinja 指向它 (见 computility-run.yaml)。
mkdir -p /workspace
CHAT_TMPL=/workspace/chat_template_multi_system.jinja
if [ -f "./chat_template_multi_system.jinja" ]; then
cp "./chat_template_multi_system.jinja" "$CHAT_TMPL"
echo "[patch_ops] deployed multi_system chat template → $CHAT_TMPL"
else
echo "[patch_ops] WARNING: ./chat_template_multi_system.jinja not found in $(pwd)" >&2
fi
# --- verification: five optimization flags present in deployed package -------
VLLM_PKG=/usr/local/corex/lib64/python3/dist-packages/vllm
echo ""
echo "=== patch_ops verification ==="
_fail=0
check() { # check <grep_pattern> <file> <label>
if grep -q "$1" "$2" 2>/dev/null; then
echo " [ok] $3"
else
echo " [MISS] $3 ($1 in $2)"; _fail=1
fi
}
check "MOE_NATIVE" "$VLLM_PKG/model_executor/models/qwen3_5.py" "MOE_NATIVE (qwen3_5.py)"
check "_native_experts_sorted" "$VLLM_PKG/model_executor/models/qwen3_5.py" "_native_experts_sorted"
check "CHUNK_PARALLEL" "$VLLM_PKG/model_executor/models/qwen3_5.py" "CHUNK_PARALLEL"
check "LOOP1_NATIVE" "$VLLM_PKG/model_executor/models/qwen3_5.py" "LOOP1_NATIVE"
check "RMSNORM_NATIVE" "$VLLM_PKG/model_executor/models/qwen3_5.py" "RMSNORM_NATIVE"
check "PREFIX_FLASH" "$VLLM_PKG/attention/ops/paged_attn.py" "PREFIX_FLASH (paged_attn.py)"
check "_forward_prefix_flash" "$VLLM_PKG/attention/ops/paged_attn.py" "_forward_prefix_flash"
# libcusolver link present (else LOOP1_NATIVE silently falls back at runtime)
if [ -e /opt/sw_home/local/cuda/lib64/libcusolver.so ]; then
echo " [ok] libcusolver.so link (LOOP1_NATIVE runtime)"
else
echo " [MISS] libcusolver.so link — LOOP1_NATIVE will fall back to serial loop"; _fail=1
fi
# multi_system chat template deployed (else multi-system requests 4xx at runtime)
if [ -f /workspace/chat_template_multi_system.jinja ]; then
echo " [ok] /workspace/chat_template_multi_system.jinja (multi_system fix)"
else
echo " [MISS] /workspace/chat_template_multi_system.jinja — multi-system requests will 4xx"; _fail=1
fi
# solve_triangular actually callable (catches a broken cusolver link early)
if python3 -c "import torch,os;os.environ.setdefault('CUDA_VISIBLE_DEVICES','0');a=torch.eye(4,device='cuda');torch.linalg.solve_triangular(-a,a,upper=False,unitriangular=True)" 2>/dev/null; then
echo " [ok] torch.linalg.solve_triangular callable"
else
echo " [MISS] torch.linalg.solve_triangular not callable — check libcusolver"; _fail=1
fi
if [ $_fail -eq 0 ]; then
echo " => all five optimizations deployed and enabled"
else
echo " => WARNING: some checks failed — review above" >&2
fi
# --- server start command lives in computility-run.yaml ----------------------
# patch_ops.sh only configures the environment + deploys patched files. It does
# NOT start the server. The server launch command (python3 -m vllm... with all
# CLI args) and the env block (PYTHONPATH / LD_LIBRARY_PATH / PATH / five
# *_NATIVE flags / VLLM_ENGINE_ITERATION_TIMEOUT_S) are in ../computility-run.yaml.
#
# After running this script, start the service via the computility runner, which
# must run from a NON-qwen3_6_scripts directory (this dir's xformers.py shadows
# the real xformers package).
#
# Sanity check after startup: the log must NOT contain any "* FALLBACK" line
# during cold prefill (prefix_flash FALLBACK on cache-hit requests is expected —
# the pad-ratio router intentionally falls back to pure PyTorch there).
# grep "FALLBACK" <server_log> # empty during cold prefill = all five active
# 80K cold prefill TTFT ≈ 57s with all five on (cached=0).
echo ""
echo "=== patch complete. Start the server with the computility runner using: ==="
echo " ../computility-run.yaml"
echo " (env + command are defined there; run from outside qwen3_6_scripts/)"

38
qwen3_6_scripts/patch_xformers_sdpa_seq.py
View File

@@ -214,18 +214,34 @@ NEW_XFORMER_BLOCK = """\
query = query.unsqueeze(0)
key = key.unsqueeze(0)
value = value.unsqueeze(0)
if self.head_size > 128:
if self.head_size == 256:
# head_dim=256: ixformer flash_attn supports it natively
# Use ixinfer_flash_attn_unpad directly (bypasses FwOp check)
from ixformer.functions.flash_attn_lib import ixinfer_flash_attn_unpad as _ixf_unpad
total_q = query.shape[0] * query.shape[1]
total_k = key.shape[0] * key.shape[1]
q_flat = query.reshape(total_q, self.num_heads, self.head_size)
k_flat = key.reshape(total_k, self.num_kv_heads, self.head_size)
v_flat = value.reshape(total_k, self.num_kv_heads, self.head_size)
if isinstance(attn_bias[0], BlockDiagonalCausalMask):
cu_seqlens_q = attn_bias[0].q_seqinfo.seqstart.int().to(query.device)
cu_seqlens_k = attn_bias[0].k_seqinfo.seqstart.int().to(query.device)
max_seqlen_q = attn_bias[0].q_seqinfo.max_seqlen
max_seqlen_k = attn_bias[0].k_seqinfo.max_seqlen
else:
batch_size = query.shape[0]
seqlen = query.shape[1]
cu_seqlens_q = torch.arange(0, batch_size + 1, device=query.device, dtype=torch.int32) * seqlen
cu_seqlens_k = cu_seqlens_q
max_seqlen_q = seqlen
max_seqlen_k = seqlen
out = _ixf_unpad(q_flat, k_flat, v_flat,
cu_seqlens_q, cu_seqlens_k,
max_seqlen_q, max_seqlen_k,
True, self.scale, out=None)
out = out.view(query.shape[0], query.shape[1], self.num_heads, self.head_size)
elif self.head_size not in (32, 64, 128, 256):
out = self._run_sdpa_fallback(query, key, value, attn_metadata)
else:
out = xops.memory_efficient_attention_forward(
query,
key,
value,
attn_bias=attn_bias[0],
p=0.0,
scale=self.scale,
op=self.attn_op,
)
return out.view_as(original_query)\
"""

72
qwen3_6_scripts/protocol.py
View File

@@ -1,5 +1,6 @@
# Adapted from
# https://github.com/lm-sys/FastChat/blob/168ccc29d3f7edc50823016105c024fe2282732a/fastchat/protocol/openai_api_protocol.py
import json
import time
from argparse import Namespace
from typing import Any, Dict, List, Literal, Optional, Union
@@ -254,6 +255,15 @@ class ChatCompletionRequest(OpenAIBaseModel):
description=("Additional kwargs to pass to the template renderer. "
"Will be accessible by the chat template."),
)
thinking: Optional[Any] = Field(
default=None,
description=(
"Per-request thinking switch (OpenAI-style top-level). Accepts: "
"{'type':'disabled'} / {'type':'enabled'} / bool true/false. "
"Translated to chat_template_kwargs.enable_thinking (bool) so the "
"Qwen3.6 template's enable_thinking branch controls reasoning. "
"Null/absent = server default (thinking on via --reasoning-parser)."),
)
guided_json: Optional[Union[str, dict, BaseModel]] = Field(
default=None,
description=("If specified, the output will follow the JSON schema."),
@@ -411,6 +421,26 @@ class ChatCompletionRequest(OpenAIBaseModel):
"Each message must have at least one of 'content' or "
"'reasoning_content'.")
msg = {**msg, "content": ""}
# tool_calls arguments: dict -> JSON string.
# Many clients/datasets send function.arguments as a dict (e.g.
# {"cmd":"ls"}), but upstream ChatCompletionMessageParam strictly
# requires a JSON string. Convert here so validation passes (matches
# OpenAI's lenient acceptance of both forms). Without this, dataset
# requests with tool messages 4xx with
# "function.arguments: Input should be a valid string".
tool_calls = msg.get("tool_calls")
if isinstance(tool_calls, list):
new_tc = []
for tc in tool_calls:
if isinstance(tc, dict) and isinstance(tc.get("function"), dict):
fn = tc["function"]
args = fn.get("arguments")
if isinstance(args, dict):
fn = {**fn, "arguments": json.dumps(
args, ensure_ascii=False)}
tc = {**tc, "function": fn}
new_tc.append(tc)
msg = {**msg, "tool_calls": new_tc}
normalized.append(msg)
data = {**data, "messages": normalized}
return data
@@ -424,6 +454,48 @@ class ChatCompletionRequest(OpenAIBaseModel):
return data
@model_validator(mode="before")
@classmethod
def validate_thinking(cls, data):
"""Translate top-level `thinking` → chat_template_kwargs.enable_thinking.
Accepts {'type':'disabled'|'enabled'} or bool. The Qwen3.6 template's
`enable_thinking is false` branch (chat_template.jinja:149) gates the
reasoning prefix, so a bool is what actually controls it. Merged into
chat_template_kwargs (request-level value wins over any pre-existing
enable_thinking). Without this, top-level `thinking` 4xx with
extra_forbidden (OpenAIBaseModel has extra="forbid").
"""
thinking = data.get("thinking")
if thinking is None:
return data
# Normalize to a bool.
enable: Optional[bool]
if isinstance(thinking, bool):
enable = thinking
elif isinstance(thinking, dict):
t = str(thinking.get("type", "")).lower()
if t == "disabled":
enable = False
elif t == "enabled" or t == "auto":
enable = True
elif t == "":
enable = None
else:
raise ValueError(
f"thinking.type must be 'enabled'/'disabled'/'auto', "
f"got {thinking.get('type')!r}")
else:
raise ValueError(
f"thinking must be bool or {{'type':'disabled'|'enabled'}}, "
f"got {type(thinking).__name__}")
if enable is not None:
ctk = dict(data.get("chat_template_kwargs") or {})
ctk["enable_thinking"] = enable
data["chat_template_kwargs"] = ctk
data.pop("thinking", None)
return data
@model_validator(mode="before")
@classmethod
def check_logprobs(cls, data):

163
qwen3_6_scripts/qwen3_5.py
View File

@@ -114,11 +114,32 @@ def _torch_chunk_gated_delta_rule(
g = g.cumsum(dim=-1)
decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril()
attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask_upper, 0)
for i in range(1, chunk_size):
row = attn[..., i, :i].clone()
sub = attn[..., :i, :i].clone()
attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2)
attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device)
# loop1: 三角求逆 (I - attn)^{-1}。默认走 solve_triangular (batched trsm kernel,
# 替代 63 步 Python 前向替换循环, 真实服务隔离 4-5x, 数值等价 max_abs<1e-7)。
# 传 -attn + unitriangular=True 等效于求解 (I-attn)@X=I → X=(I-attn)^{-1}。
# 注意: 不做运行时 attn.abs().max() 阈值检查(会引入 device→host 同步,
# 600次调用下额外耗时~1-2s)。真实模型 attn 值域远在安全范围内。
# LOOP1_NATIVE=0 可回退串行循环。
import os as _os
_loop1_done = False
if _os.environ.get("LOOP1_NATIVE", "1") != "0":
try:
_eye = torch.eye(chunk_size, dtype=attn.dtype, device=attn.device)
attn = torch.linalg.solve_triangular(
-attn, _eye, upper=False, unitriangular=True)
_loop1_done = True
except Exception as _e:
import sys as _sys, traceback as _tb
print(f"[loop1_native FALLBACK] {type(_e).__name__}: {_e}",
file=_sys.stderr, flush=True)
_tb.print_exc(file=_sys.stderr)
if not _loop1_done:
for i in range(1, chunk_size):
row = attn[..., i, :i].clone()
sub = attn[..., :i, :i].clone()
attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2)
attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device)
value = attn @ v_beta
k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1))
@@ -132,18 +153,54 @@ def _torch_chunk_gated_delta_rule(
torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device),
diagonal=1)
for i in range(total_len // chunk_size):
q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i]
attn_i = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask_upper2, 0)
v_prime = k_cumdecay[:, :, i] @ last_state
v_new = v_i - v_prime
attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_state
core_out[:, :, i] = attn_inter + attn_i @ v_new
last_state = (
last_state * g[:, :, i, -1, None, None].exp()
+ (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None])
.transpose(-1, -2) @ v_new
)
# loop2: chunk 间 state 递推。默认走 chunk-parallel(把串行 scan 改写成
# 线性矩阵递归 S_i=A_i@S_{i-1}+B_i,重活批量化,串行阶段只对小矩阵 S 做 A@S+B)。
# 数值等价于下方串行实现(已验证 max_abs~1e-7, 3x)。CHUNK_PARALLEL=0 可回退。
import os as _os
NC = total_len // chunk_size
_done = False
if _os.environ.get("CHUNK_PARALLEL", "1") != "0":
try:
g_last = g[:, :, :, -1] # [b,h,NC]
exp_glast = g_last.exp()
attn_i_all = (query @ key.transpose(-1, -2) * decay_mask
).masked_fill(mask_upper2, 0) # [b,h,NC,cs,cs]
P_i = k_cumdecay # [b,h,NC,cs,k]
Ktil = key * (g_last[..., None, None] - g[..., None]).exp() # [b,h,NC,cs,k]
Qg = query * g.exp()[..., None] # [b,h,NC,cs,k]
eye = torch.eye(k_dim, device=query.device, dtype=query.dtype)
A = exp_glast[..., None, None] * eye - Ktil.transpose(-1, -2) @ P_i # [b,h,NC,k,k]
B = Ktil.transpose(-1, -2) @ value # [b,h,NC,k,v]
M = Qg - attn_i_all @ P_i # [b,h,NC,cs,k]
N = attn_i_all @ value # [b,h,NC,cs,v]
S = last_state
S_starts = torch.empty(batch, num_heads, NC, k_dim, v_dim,
device=query.device, dtype=query.dtype)
for i in range(NC):
S_starts[:, :, i] = S
S = A[:, :, i] @ S + B[:, :, i]
last_state = S
core_out = M @ S_starts + N # [b,h,NC,cs,v]
_done = True
except Exception as _e:
import sys as _sys, traceback as _tb
print(f"[chunk_parallel FALLBACK] {type(_e).__name__}: {_e}",
file=_sys.stderr, flush=True)
_tb.print_exc(file=_sys.stderr)
if not _done:
for i in range(NC):
q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i]
attn_i = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask_upper2, 0)
v_prime = k_cumdecay[:, :, i] @ last_state
v_new = v_i - v_prime
attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_state
core_out[:, :, i] = attn_inter + attn_i @ v_new
last_state = (
last_state * g[:, :, i, -1, None, None].exp()
+ (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None])
.transpose(-1, -2) @ v_new
)
if not output_final_state:
last_state = None
@@ -212,6 +269,20 @@ class Qwen3_5RMSNormGated(nn.Module):
def forward(self, hidden_states: torch.Tensor,
gate: torch.Tensor) -> torch.Tensor:
# 默认用原生 ixformer rms_norm(替代纯 PyTorch pow/mean/rsqrt),
# 再乘 silu(gate)。已验证与下方纯 PyTorch 等价(max_abs~1e-2 fp16)。
# RMSNORM_NATIVE=0 或异常回退纯 PyTorch。
import os as _os3
if _os3.environ.get("RMSNORM_NATIVE", "1") != "0":
try:
from vllm import _custom_ops as _ops
out = torch.empty_like(hidden_states)
_ops.rms_norm(out, hidden_states.contiguous(),
self.weight, self.variance_epsilon)
return (out.to(torch.float32)
* F.silu(gate.to(torch.float32))).to(hidden_states.dtype)
except Exception:
pass
input_dtype = hidden_states.dtype
hs = hidden_states.to(torch.float32)
variance = hs.pow(2).mean(-1, keepdim=True)
@@ -792,8 +863,22 @@ class Qwen3_5MoeSparseBlock(nn.Module):
out = (expert_out * ws.unsqueeze(-1)).sum(0, keepdim=True).to(
hidden_states.dtype) # (1, H)
else:
# General path (prefill / multi-seq): loop over unique active experts.
# At most T*top_k unique experts, always <= num_experts.
# General path (prefill / multi-seq).
# 优化: sort-by-expert + 连续切片 + 原生 kernel(act_bias_mm/silu_and_mul/
# ixf linear),消除逐 expert 的 nonzero/index_add 循环调度开销。
# 数值等价于下方纯 PyTorch fallback(已验证 byte-for-byte 一致, 端到端1.36x)。
# 任意异常回退纯 PyTorch。可用 MOE_NATIVE=0 关闭。
import os as _os
if _os.environ.get("MOE_NATIVE", "1") != "0":
try:
return self._native_experts_sorted(
hidden_states, w13, w2, topk_weights, topk_ids)
except Exception as _e:
import sys as _sys, traceback as _tb
print(f"[moe_native FALLBACK] {type(_e).__name__}: {_e}",
file=_sys.stderr, flush=True)
_tb.print_exc(file=_sys.stderr)
# ---- 纯 PyTorch fallback (原实现) ----
out = torch.zeros_like(hidden_states)
unique_eids = topk_ids.view(-1).unique().tolist()
for eid in unique_eids:
@@ -810,6 +895,46 @@ class Qwen3_5MoeSparseBlock(nn.Module):
return out # partial, all-reduce done in forward()
def _native_experts_sorted(self, hidden_states, w13, w2, topk_weights, topk_ids):
"""sort-by-expert + 原生 ixformer kernel 的 MoE expert 计算。
w13=(E,2I,H), w2=(E,H,I)。返回 partial(pre-all-reduce)。"""
import ixformer.functions as _ixf
dev = hidden_states.device
T = hidden_states.shape[0]
top_k = topk_ids.shape[1]
E = w13.shape[0]
flat_e = topk_ids.reshape(-1) # (T*K,)
flat_w = topk_weights.reshape(-1) # (T*K,)
flat_tok = torch.arange(T, device=dev).repeat_interleave(top_k) # (T*K,)
order = torch.argsort(flat_e)
se = flat_e[order]
sw = flat_w[order]
stok = flat_tok[order]
gathered = hidden_states.index_select(0, stok).contiguous() # (T*K, H) 已按 expert 连续
counts = torch.bincount(se, minlength=E)
offs = torch.cat([torch.zeros(1, device=dev, dtype=torch.long),
counts.cumsum(0)]).tolist()
# 就地把每个 expert 段的输出写回 gathered(复用显存,避免再分配 T*K×H,
# 否则 profiling 阶段峰值显存↑ → KV cache 容量不足)。
active = (counts > 0).nonzero().flatten().tolist()
for eid in active:
eid = int(eid)
s, e = int(offs[eid]), int(offs[eid + 1])
seg = gathered[s:e].contiguous() # (n, H)
gu = _ixf.act_bias_mm(seg, w13[eid], None, scale=1,
act_type="none", trans_format="TN") # (n, 2I)
act = _ixf.silu_and_mul(gu.contiguous()) # (n, I)
gathered[s:e] = _ixf.linear(act.contiguous(), w2[eid], None) # (n, H) 写回
gathered = gathered * sw.unsqueeze(-1).to(gathered.dtype)
out = torch.zeros_like(hidden_states)
out.index_add_(0, stok, gathered.to(out.dtype))
return out
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
router_logits, _ = self.gate(hidden_states)
routed_out = self._pure_pytorch_experts(hidden_states, router_logits)

946
qwen3_6_scripts/xformers.py Normal file
View File

@@ -0,0 +1,946 @@
"""Attention layer with xFormers and PagedAttention."""
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple, Type
import torch
# from xformers import ops as xops
from ixformer.contrib.xformers import ops as xops
from xformers.ops.fmha.attn_bias import (AttentionBias,
BlockDiagonalMask,)
from ixformer.contrib.xformers.ops.fmha.attn_bias import (BlockDiagonalCausalMask,
LowerTriangularMaskWithTensorBias)
from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
AttentionMetadata, AttentionType)
from vllm.attention.backends.utils import (CommonAttentionState,
CommonMetadataBuilder)
from vllm.attention.ops.paged_attn import (PagedAttention,
PagedAttentionMetadata)
from vllm.logger import init_logger
logger = init_logger(__name__)
class XFormersBackend(AttentionBackend):
@staticmethod
def get_name() -> str:
return "xformers"
@staticmethod
def get_impl_cls() -> Type["XFormersImpl"]:
return XFormersImpl
@staticmethod
def get_metadata_cls() -> Type["AttentionMetadata"]:
return XFormersMetadata
@staticmethod
def get_builder_cls() -> Type["XFormersMetadataBuilder"]:
return XFormersMetadataBuilder
@staticmethod
def get_state_cls() -> Type["CommonAttentionState"]:
return CommonAttentionState
@staticmethod
def get_kv_cache_shape(
num_blocks: int,
block_size: int,
num_kv_heads: int,
head_size: int,
) -> Tuple[int, ...]:
return PagedAttention.get_kv_cache_shape(num_blocks, block_size,
num_kv_heads, head_size)
@staticmethod
def swap_blocks(
src_kv_cache: torch.Tensor,
dst_kv_cache: torch.Tensor,
src_to_dst: Dict[int, int],
) -> None:
PagedAttention.swap_blocks(src_kv_cache, dst_kv_cache, src_to_dst)
@staticmethod
def copy_blocks(
kv_caches: List[torch.Tensor],
src_to_dists: torch.Tensor,
) -> None:
PagedAttention.copy_blocks(kv_caches, src_to_dists)
@dataclass
class XFormersMetadata(AttentionMetadata, PagedAttentionMetadata):
"""Metadata for XFormersbackend.
NOTE: Any python object stored here is not updated when it is
cuda-graph replayed. If you have values that need to be changed
dynamically, it should be stored in tensor. The tensor has to be
updated from `CUDAGraphRunner.forward` API.
"""
# |---------- N-1 iteration --------|
# |---------------- N iteration ---------------------|
# |- tokenA -|......................|-- newTokens ---|
# |---------- context_len ----------|
# |-------------------- seq_len ----------------------|
# |-- query_len ---|
# seq_lens stored as a tensor.
seq_lens_tensor: Optional[torch.Tensor]
# FIXME: It is for flash attn.
# Maximum sequence length among prefill batch. 0 if there are decoding
# requests only.
max_prefill_seq_len: int
# Maximum sequence length among decode batch. 0 if there are prefill
# requests only.
max_decode_seq_len: int
# Whether or not if cuda graph is enabled.
# Cuda-graph is currently enabled for decoding only.
# TODO(woosuk): Move `use_cuda_graph` out since it's unrelated to attention.
use_cuda_graph: bool
# (batch_size,). The sequence length per sequence. Sequence length means
# the computed tokens + new tokens None if it is a decoding.
seq_lens: Optional[List[int]] = None
# FIXME: It is for flash attn.
# (batch_size + 1,). The cumulative sequence lengths of the sequences in
# the batch, used to index into sequence. E.g., if the sequence length is
# [4, 6], it is [0, 4, 10].
seq_start_loc: Optional[torch.Tensor] = None
# (batch_size,) A tensor of context lengths (tokens that are computed
# so far).
context_lens_tensor: Optional[torch.Tensor] = None
# Maximum query length in the batch. None for decoding.
max_query_len: Optional[int] = None
# Max number of query tokens among request in the batch.
max_decode_query_len: Optional[int] = None
# (batch_size + 1,). The cumulative subquery lengths of the sequences in
# the batch, used to index into subquery. E.g., if the subquery length
# is [4, 6], it is [0, 4, 10].
query_start_loc: Optional[torch.Tensor] = None
# Self-attention prefill/decode metadata cache
_cached_prefill_metadata: Optional["XFormersMetadata"] = None
_cached_decode_metadata: Optional["XFormersMetadata"] = None
# Begin encoder attn & enc/dec cross-attn fields...
# Encoder sequence lengths representation
encoder_seq_lens: Optional[List[int]] = None
encoder_seq_lens_tensor: Optional[torch.Tensor] = None
# Maximum sequence length among encoder sequences
max_encoder_seq_len: Optional[int] = None
# Number of tokens input to encoder
num_encoder_tokens: Optional[int] = None
# Cross-attention memory-mapping data structures: slot mapping
# and block tables
cross_slot_mapping: Optional[torch.Tensor] = None
cross_block_tables: Optional[torch.Tensor] = None
def __post_init__(self):
# Set during the execution of the first attention op.
# It is a list because it is needed to set per prompt
# when alibi slopes is used. It is because of the limitation
# from xformer API.
# will not appear in the __repr__ and __init__
self.attn_bias: Optional[List[AttentionBias]] = None
self.encoder_attn_bias: Optional[List[AttentionBias]] = None
self.cross_attn_bias: Optional[List[AttentionBias]] = None
@property
def is_all_encoder_attn_metadata_set(self):
'''
All attention metadata required for encoder attention is set.
'''
return ((self.encoder_seq_lens is not None)
and (self.encoder_seq_lens_tensor is not None)
and (self.max_encoder_seq_len is not None))
@property
def is_all_cross_attn_metadata_set(self):
'''
All attention metadata required for enc/dec cross-attention is set.
Superset of encoder attention required metadata.
'''
return (self.is_all_encoder_attn_metadata_set
and (self.cross_slot_mapping is not None)
and (self.cross_block_tables is not None))
@property
def prefill_metadata(self) -> Optional["XFormersMetadata"]:
if self.num_prefills == 0:
return None
if self._cached_prefill_metadata is not None:
# Recover cached prefill-phase attention
# metadata structure
return self._cached_prefill_metadata
assert ((self.seq_lens is not None)
or (self.encoder_seq_lens is not None))
assert ((self.seq_lens_tensor is not None)
or (self.encoder_seq_lens_tensor is not None))
# Compute some attn_metadata fields which default to None
query_start_loc = (None if self.query_start_loc is None else
self.query_start_loc[:self.num_prefills + 1])
slot_mapping = (None if self.slot_mapping is None else
self.slot_mapping[:self.num_prefill_tokens])
seq_lens = (None if self.seq_lens is None else
self.seq_lens[:self.num_prefills])
seq_lens_tensor = (None if self.seq_lens_tensor is None else
self.seq_lens_tensor[:self.num_prefills])
context_lens_tensor = (None if self.context_lens_tensor is None else
self.context_lens_tensor[:self.num_prefills])
block_tables = (None if self.block_tables is None else
self.block_tables[:self.num_prefills])
# Construct & cache prefill-phase attention metadata structure
self._cached_prefill_metadata = XFormersMetadata(
num_prefills=self.num_prefills,
num_prefill_tokens=self.num_prefill_tokens,
num_decode_tokens=0,
slot_mapping=slot_mapping,
seq_lens=seq_lens,
seq_lens_tensor=seq_lens_tensor,
max_query_len=self.max_query_len,
max_prefill_seq_len=self.max_prefill_seq_len,
max_decode_seq_len=0,
query_start_loc=query_start_loc,
context_lens_tensor=context_lens_tensor,
block_tables=block_tables,
use_cuda_graph=False,
# Begin encoder & cross attn fields below...
encoder_seq_lens=self.encoder_seq_lens,
encoder_seq_lens_tensor=self.encoder_seq_lens_tensor,
max_encoder_seq_len=self.max_encoder_seq_len,
cross_slot_mapping=self.cross_slot_mapping,
cross_block_tables=self.cross_block_tables)
return self._cached_prefill_metadata
@property
def decode_metadata(self) -> Optional["XFormersMetadata"]:
if self.num_decode_tokens == 0:
return None
if self._cached_decode_metadata is not None:
# Recover cached decode-phase attention
# metadata structure
return self._cached_decode_metadata
assert ((self.seq_lens_tensor is not None)
or (self.encoder_seq_lens_tensor is not None))
# Compute some attn_metadata fields which default to None
slot_mapping = (None if self.slot_mapping is None else
self.slot_mapping[self.num_prefill_tokens:])
seq_lens_tensor = (None if self.seq_lens_tensor is None else
self.seq_lens_tensor[self.num_prefills:])
block_tables = (None if self.block_tables is None else
self.block_tables[self.num_prefills:])
# Construct & cache decode-phase attention metadata structure
self._cached_decode_metadata = XFormersMetadata(
num_prefills=0,
num_prefill_tokens=0,
num_decode_tokens=self.num_decode_tokens,
slot_mapping=slot_mapping,
seq_lens_tensor=seq_lens_tensor,
max_prefill_seq_len=0,
max_decode_seq_len=self.max_decode_seq_len,
block_tables=block_tables,
use_cuda_graph=self.use_cuda_graph,
# Begin encoder & cross attn fields below...
encoder_seq_lens=self.encoder_seq_lens,
encoder_seq_lens_tensor=self.encoder_seq_lens_tensor,
max_encoder_seq_len=self.max_encoder_seq_len,
cross_slot_mapping=self.cross_slot_mapping,
cross_block_tables=self.cross_block_tables)
return self._cached_decode_metadata
def _get_attn_bias(
attn_metadata: XFormersMetadata,
attn_type: AttentionType,
) -> Optional[AttentionBias]:
'''
Extract appropriate attention bias from attention metadata
according to attention type.
Arguments:
* attn_metadata: Attention metadata structure associated with attention
* attn_type: encoder attention, decoder self-attention,
encoder/decoder cross-attention
Returns:
* Appropriate attention bias value given the attention type
'''
if attn_type == AttentionType.DECODER:
return attn_metadata.attn_bias
elif attn_type == AttentionType.ENCODER:
return attn_metadata.encoder_attn_bias
else:
# attn_type == AttentionType.ENCODER_DECODER
return attn_metadata.cross_attn_bias
def _set_attn_bias(
attn_metadata: XFormersMetadata,
attn_bias: List[Optional[AttentionBias]],
attn_type: AttentionType,
) -> None:
'''
Update appropriate attention bias field of attention metadata,
according to attention type.
Arguments:
* attn_metadata: Attention metadata structure associated with attention
* attn_bias: The desired attention bias value
* attn_type: encoder attention, decoder self-attention,
encoder/decoder cross-attention
'''
if attn_type == AttentionType.DECODER:
attn_metadata.attn_bias = attn_bias
elif attn_type == AttentionType.ENCODER:
attn_metadata.encoder_attn_bias = attn_bias
elif attn_type == AttentionType.ENCODER_DECODER:
attn_metadata.cross_attn_bias = attn_bias
else:
raise AttributeError(f"Invalid attention type {str(attn_type)}")
def _get_seq_len_block_table_args(
attn_metadata: XFormersMetadata,
is_prompt: bool,
attn_type: AttentionType,
) -> tuple:
'''
The particular choice of sequence-length- and block-table-related
attributes which should be extracted from attn_metadata is dependent
on the type of attention operation.
Decoder attn -> select entirely decoder self-attention-related fields
Encoder/decoder cross-attn -> select encoder sequence lengths &
cross-attn block-tables fields
Encoder attn -> select encoder sequence lengths fields & no block tables
Arguments:
* attn_metadata: Attention metadata structure associated with attention op
* is_prompt: True if prefill, False otherwise
* attn_type: encoder attention, decoder self-attention,
encoder/decoder cross-attention
Returns:
* Appropriate sequence-lengths tensor
* Appropriate max sequence-length scalar
* Appropriate block tables (or None)
'''
if attn_type == AttentionType.DECODER:
# Decoder self-attention
# Choose max_seq_len based on whether we are in prompt_run
if is_prompt:
max_seq_len = attn_metadata.max_prefill_seq_len
else:
max_seq_len = attn_metadata.max_decode_seq_len
return (attn_metadata.seq_lens_tensor, max_seq_len,
attn_metadata.block_tables)
elif attn_type == AttentionType.ENCODER_DECODER:
# Enc/dec cross-attention KVs match encoder sequence length;
# cross-attention utilizes special "cross" block tables
return (attn_metadata.encoder_seq_lens_tensor,
attn_metadata.max_encoder_seq_len,
attn_metadata.cross_block_tables)
elif attn_type == AttentionType.ENCODER:
# No block tables associated with encoder attention
return (attn_metadata.encoder_seq_lens_tensor,
attn_metadata.max_encoder_seq_len, None)
else:
raise AttributeError(f"Invalid attention type {str(attn_type)}")
class XFormersMetadataBuilder(CommonMetadataBuilder[XFormersMetadata]):
_metadata_cls = XFormersMetadata
class XFormersImpl(AttentionImpl[XFormersMetadata]):
"""
If the input tensors contain prompt tokens, the layout is as follows:
|<--------------- num_prefill_tokens ----------------->|
|<--prefill_0-->|<--prefill_1-->|...|<--prefill_N-1--->|
Otherwise, the layout is as follows:
|<----------------- num_decode_tokens ------------------>|
|<--decode_0-->|..........|<--decode_M-1-->|<--padding-->|
Generation tokens can contain padding when cuda-graph is used.
Currently, prompt tokens don't contain any padding.
The prompts might have different lengths, while the generation tokens
always have length 1.
If chunked prefill is enabled, prefill tokens and decode tokens can be
batched together in a flattened 1D query.
|<----- num_prefill_tokens ---->|<------- num_decode_tokens --------->|
|<-prefill_0->|...|<-prefill_N-1->|<--decode_0-->|...|<--decode_M-1-->|
Currently, cuda graph is disabled for chunked prefill, meaning there's no
padding between prefill and decode tokens.
"""
def __init__(
self,
num_heads: int,
head_size: int,
scale: float,
num_kv_heads: int,
alibi_slopes: Optional[List[float]],
sliding_window: Optional[int],
kv_cache_dtype: str,
blocksparse_params: Optional[Dict[str, Any]] = None,
logits_soft_cap: Optional[float] = None,
) -> None:
if blocksparse_params is not None:
raise ValueError(
"XFormers does not support block-sparse attention.")
if logits_soft_cap is not None:
raise ValueError(
"XFormers does not support attention logits soft capping.")
self.num_heads = num_heads
self.head_size = head_size
self.scale = float(scale)
self.num_kv_heads = num_kv_heads
if alibi_slopes is not None:
alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32)
self.alibi_slopes = alibi_slopes
self.sliding_window = sliding_window
self.kv_cache_dtype = kv_cache_dtype
assert self.num_heads % self.num_kv_heads == 0
self.num_queries_per_kv = self.num_heads // self.num_kv_heads
suppored_head_sizes = PagedAttention.get_supported_head_sizes()
if head_size not in suppored_head_sizes:
raise ValueError(
f"Head size {head_size} is not supported by PagedAttention. "
f"Supported head sizes are: {suppored_head_sizes}.")
self.head_mapping = torch.repeat_interleave(
torch.arange(self.num_kv_heads, dtype=torch.int32),
self.num_queries_per_kv)
def forward(
self,
query: torch.Tensor,
key: Optional[torch.Tensor],
value: Optional[torch.Tensor],
kv_cache: torch.Tensor,
attn_metadata: "XFormersMetadata",
k_scale: float = 1.0,
v_scale: float = 1.0,
attn_type: AttentionType = AttentionType.DECODER,
) -> torch.Tensor:
"""Forward pass with xFormers and PagedAttention.
For decoder-only models: query, key and value must be non-None.
For encoder/decoder models:
* XFormersImpl.forward() may be invoked for both self- and cross-
attention layers.
* For self-attention: query, key and value must be non-None.
* For cross-attention:
* Query must be non-None
* During prefill, key and value must be non-None; key and value
get cached for use during decode.
* During decode, key and value may be None, since:
(1) key and value tensors were cached during prefill, and
(2) cross-attention key and value tensors do not grow during
decode
A note on how the attn_type (attention type enum) argument impacts
attention forward() behavior:
* DECODER: normal decoder-only behavior;
use decoder self-attention block table
* ENCODER: no KV caching; pass encoder sequence
attributes (encoder_seq_lens/encoder_seq_lens_tensor/
max_encoder_seq_len) to kernel, in lieu of decoder
sequence attributes (seq_lens/seq_lens_tensor/max_seq_len)
* ENCODER_DECODER: cross-attention behavior;
use cross-attention block table for caching KVs derived
from encoder hidden states; since KV sequence lengths
will match encoder sequence lengths, pass encoder sequence
attributes to kernel (encoder_seq_lens/encoder_seq_lens_tensor/
max_encoder_seq_len)
Args:
query: shape = [num_tokens, num_heads * head_size]
key: shape = [num_tokens, num_kv_heads * head_size]
value: shape = [num_tokens, num_kv_heads * head_size]
kv_cache = [2, num_blocks, block_size * num_kv_heads * head_size]
NOTE: kv_cache will be an empty tensor with shape [0]
for profiling run.
attn_metadata: Metadata for attention.
attn_type: Select attention type, between encoder attention,
decoder self-attention, or encoder/decoder cross-
attention. Defaults to decoder self-attention,
which is the vLLM default generally
Returns:
shape = [num_tokens, num_heads * head_size]
"""
# Check that appropriate attention metadata attributes are
# selected for the desired attention type
if (attn_type == AttentionType.ENCODER
and (not attn_metadata.is_all_encoder_attn_metadata_set)):
raise AttributeError("Encoder attention requires setting "
"encoder metadata attributes.")
elif (attn_type == AttentionType.ENCODER_DECODER
and (not attn_metadata.is_all_cross_attn_metadata_set)):
raise AttributeError("Encoder/decoder cross-attention "
"requires setting cross-attention "
"metadata attributes.")
query = query.view(-1, self.num_heads, self.head_size)
if key is not None:
assert value is not None
key = key.view(-1, self.num_kv_heads, self.head_size)
value = value.view(-1, self.num_kv_heads, self.head_size)
else:
assert value is None
# Self-attention vs. cross-attention will impact
# which KV cache memory-mapping & which
# seqlen datastructures we utilize
if (attn_type != AttentionType.ENCODER and kv_cache.numel() > 0):
# KV-cache during decoder-self- or
# encoder-decoder-cross-attention, but not
# during encoder attention.
#
# Even if there are no new key/value pairs to cache,
# we still need to break out key_cache and value_cache
# i.e. for later use by paged attention
key_cache, value_cache = PagedAttention.split_kv_cache(
kv_cache, self.num_kv_heads, self.head_size)
if (key is not None) and (value is not None):
if attn_type == AttentionType.ENCODER_DECODER:
# Update cross-attention KV cache (prefill-only)
# During cross-attention decode, key & value will be None,
# preventing this IF-statement branch from running
updated_slot_mapping = attn_metadata.cross_slot_mapping
else:
# Update self-attention KV cache (prefill/decode)
updated_slot_mapping = attn_metadata.slot_mapping
# Reshape the input keys and values and store them in the cache.
# If kv_cache is not provided, the new key and value tensors are
# not cached. This happens during the initial memory
# profiling run.
PagedAttention.write_to_paged_cache(key, value, key_cache,
value_cache,
updated_slot_mapping,
self.kv_cache_dtype,
k_scale, v_scale)
if attn_type == AttentionType.ENCODER:
# Encoder attention - chunked prefill is not applicable;
# derive token-count from query shape & and treat them
# as 100% prefill tokens
assert attn_metadata.num_encoder_tokens is not None
num_prefill_tokens = attn_metadata.num_encoder_tokens
num_encoder_tokens = attn_metadata.num_encoder_tokens
num_decode_tokens = 0
elif attn_type == AttentionType.DECODER:
# Decoder self-attention supports chunked prefill.
num_prefill_tokens = attn_metadata.num_prefill_tokens
num_encoder_tokens = attn_metadata.num_prefill_tokens
num_decode_tokens = attn_metadata.num_decode_tokens
# Only enforce this shape-constraint for decoder
# self-attention
assert key.shape[0] == num_prefill_tokens + num_decode_tokens
assert value.shape[0] == num_prefill_tokens + num_decode_tokens
else: # attn_type == AttentionType.ENCODER_DECODER
# Encoder/decoder cross-attention requires no chunked
# prefill (100% prefill or 100% decode tokens, no mix)
num_prefill_tokens = attn_metadata.num_prefill_tokens
if attn_metadata.num_encoder_tokens is not None:
num_encoder_tokens = attn_metadata.num_encoder_tokens
else:
num_encoder_tokens = attn_metadata.num_prefill_tokens
num_decode_tokens = attn_metadata.num_decode_tokens
output = torch.empty_like(query)
# Query for decode. KV is not needed because it is already cached.
decode_query = query[num_prefill_tokens:]
# QKV for prefill.
query = query[:num_prefill_tokens]
if key is not None and value is not None:
key = key[:num_encoder_tokens]
value = value[:num_encoder_tokens]
assert query.shape[0] == num_prefill_tokens
assert decode_query.shape[0] == num_decode_tokens
if prefill_meta := attn_metadata.prefill_metadata:
# Prompt run.
if kv_cache.numel() == 0 or prefill_meta.block_tables.numel() == 0:
# normal attention.
# block tables are empty if the prompt does not have a cached
# prefix.
out = self._run_memory_efficient_xformers_forward(
query, key, value, prefill_meta, attn_type=attn_type)
assert out.shape == output[:num_prefill_tokens].shape
output[:num_prefill_tokens] = out
else:
assert prefill_meta.query_start_loc is not None
assert prefill_meta.max_query_len is not None
# prefix-enabled attention
# TODO(Hai) this triton kernel has regression issue (broke) to
# deal with different data types between KV and FP8 KV cache,
# to be addressed separately.
out = PagedAttention.forward_prefix(
query,
key,
value,
self.kv_cache_dtype,
key_cache,
value_cache,
prefill_meta.block_tables,
prefill_meta.query_start_loc,
prefill_meta.seq_lens_tensor,
prefill_meta.context_lens_tensor,
prefill_meta.max_query_len,
self.alibi_slopes,
self.sliding_window,
k_scale,
v_scale,
)
assert output[:num_prefill_tokens].shape == out.shape
output[:num_prefill_tokens] = out
if decode_meta := attn_metadata.decode_metadata:
(
seq_lens_arg,
max_seq_len_arg,
block_tables_arg,
) = _get_seq_len_block_table_args(decode_meta, False, attn_type)
output[num_prefill_tokens:] = PagedAttention.forward_decode(
decode_query,
key_cache,
value_cache,
block_tables_arg,
seq_lens_arg,
max_seq_len_arg,
self.kv_cache_dtype,
self.head_mapping,
self.scale,
self.alibi_slopes,
k_scale,
v_scale,
)
# Reshape the output tensor.
return output.view(-1, self.num_heads * self.head_size)
def _run_sdpa_fallback(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
attn_metadata: "XFormersMetadata",
) -> torch.Tensor:
"""纯数学 causal attention fallback，带 Q-tiling 内存优化。
调用时机：kv_cache.numel()==0（profiling 阶段）。
此路径无 KV 缓存前缀，KV 长度 == query 长度。
内存优化（Q-tiling，与 Flash Attention 同思路）：
将 Q 分成 _Q_CHUNK 大小的子块逐块计算，每块峰值内存
O(_Q_CHUNK × q_len) 而非 O(q_len²)。
profiling 阶段序列可能达到 max_model_len（如 20K tokens），
不加 Q-tiling 会产生 9.6 GB 矩阵直接 OOM。
softmax 在 float32 下计算以防止 float16 溢出，结果转回原始 dtype。
Args:
query : [1, total_query_tokens, num_heads, head_dim]
key : [1, total_query_tokens, num_kv_heads, head_dim]
value : [1, total_query_tokens, num_kv_heads, head_dim]
Returns:
[1, total_query_tokens, num_heads, head_dim]
"""
_Q_CHUNK = 256 # 与 _forward_prefix_pytorch 的 _ATTN_Q_CHUNK 保持一致
assert attn_metadata.seq_lens is not None
orig_dtype = query.dtype
num_seqs = len(attn_metadata.seq_lens)
# 推导每条序列的实际 query 长度。
# 正常 prefill 时 q_len == seq_len；如果将来遇到 chunked 场景，
# query_start_loc 记录的是真实 query token 数（非全序列长度）。
if (attn_metadata.query_start_loc is not None
and len(attn_metadata.query_start_loc) == num_seqs + 1):
q_lens = [
int(attn_metadata.query_start_loc[i + 1].item()) -
int(attn_metadata.query_start_loc[i].item())
for i in range(num_seqs)
]
else:
q_lens = list(attn_metadata.seq_lens)
q_flat = query.squeeze(0) # [T, H, D]
k_flat = key.squeeze(0) # [T, Hkv, D]
v_flat = value.squeeze(0)
output = torch.empty_like(q_flat)
seq_start = 0
for q_len in q_lens:
seq_end = seq_start + q_len
# 当前序列的完整 K/V（此路径无前缀，KV == Q）
k_s = k_flat[seq_start:seq_end].permute(1, 0, 2).float() # [Hkv, q_len, D]
v_s = v_flat[seq_start:seq_end].permute(1, 0, 2).float() # [Hkv, q_len, D]
# GQA：展开 KV heads 至与 query heads 一致
if k_s.shape[0] != self.num_heads:
n = self.num_heads // k_s.shape[0]
k_s = k_s.repeat_interleave(n, dim=0).contiguous()
v_s = v_s.repeat_interleave(n, dim=0).contiguous()
# k_pos 用于因果掩码
k_pos = torch.arange(q_len, device=query.device)
# Q-tiling：分块处理 query，峰值内存 O(_Q_CHUNK × q_len)
for qc_start in range(0, q_len, _Q_CHUNK):
qc_end = min(qc_start + _Q_CHUNK, q_len)
# [H, qc, D]
q_c = q_flat[seq_start + qc_start:seq_start + qc_end] .permute(1, 0, 2).float()
# [H, qc, q_len]
attn_w = torch.matmul(q_c, k_s.transpose(-2, -1)) * self.scale
# 因果掩码：q_c 里位置 j 只能看 k_pos <= j（相对位置）
qc_q_pos = torch.arange(qc_start, qc_end, device=query.device)
mask = k_pos.unsqueeze(0) > qc_q_pos.unsqueeze(1)
attn_w = attn_w.masked_fill(mask.unsqueeze(0), float("-inf"))
attn_w = torch.softmax(attn_w, dim=-1)
out_c = torch.matmul(attn_w, v_s).to(orig_dtype) # [H, qc, D]
output[seq_start + qc_start:seq_start + qc_end] = (
out_c.permute(1, 0, 2))
seq_start = seq_end
return output.unsqueeze(0) # [1, T, H, D]
def _run_memory_efficient_xformers_forward(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
attn_metadata: XFormersMetadata,
attn_type: AttentionType = AttentionType.DECODER,
) -> torch.Tensor:
"""Attention for 1D query of multiple prompts. Multiple prompt
tokens are flattened in to `query` input.
See https://facebookresearch.github.io/xformers/components/ops.html
for API spec.
Args:
output: shape = [num_prefill_tokens, num_heads, head_size]
query: shape = [num_prefill_tokens, num_heads, head_size]
key: shape = [num_prefill_tokens, num_kv_heads, head_size]
value: shape = [num_prefill_tokens, num_kv_heads, head_size]
attn_metadata: Metadata for attention.
attn_type: Select attention type, between encoder attention,
decoder self-attention, or encoder/decoder cross-
attention. Defaults to decoder self-attention,
which is the vLLM default generally
"""
original_query = query
# if self.num_kv_heads != self.num_heads:
# # GQA/MQA requires the shape [B, M, G, H, K].
# # Note that the output also has the same shape (which is different
# # from a spec from the doc).
# query = query.view(query.shape[0], self.num_kv_heads,
# self.num_queries_per_kv, query.shape[-1])
# print(f"5555555555555 q shape {query.shape}")
# key = key[:, :,
# None, :].expand(key.shape[0], self.num_kv_heads,
# self.num_queries_per_kv, key.shape[-1])
# value = value[:, :,
# None, :].expand(value.shape[0], self.num_kv_heads,
# self.num_queries_per_kv,
# value.shape[-1])
# Set attention bias if not provided. This typically happens at
# the very attention layer of every iteration.
# FIXME(woosuk): This is a hack.
attn_bias = _get_attn_bias(attn_metadata, attn_type)
if attn_bias is None:
if self.alibi_slopes is None:
if (attn_type == AttentionType.ENCODER_DECODER):
assert attn_metadata.seq_lens is not None
assert attn_metadata.encoder_seq_lens is not None
# Default enc/dec cross-attention mask is non-causal
attn_bias = BlockDiagonalMask.from_seqlens(
attn_metadata.seq_lens, attn_metadata.encoder_seq_lens)
elif attn_type == AttentionType.ENCODER:
assert attn_metadata.encoder_seq_lens is not None
# Default encoder self-attention mask is non-causal
attn_bias = BlockDiagonalMask.from_seqlens(
attn_metadata.encoder_seq_lens)
else:
assert attn_metadata.seq_lens is not None
# Default decoder self-attention mask is causal
attn_bias = BlockDiagonalCausalMask.from_seqlens(
attn_metadata.seq_lens)
if self.sliding_window is not None:
attn_bias = attn_bias.make_local_attention(
self.sliding_window)
attn_bias = [attn_bias]
else:
assert attn_metadata.seq_lens is not None
attn_bias = _make_alibi_bias(self.alibi_slopes,
self.num_kv_heads, query.dtype,
attn_metadata.seq_lens)
_set_attn_bias(attn_metadata, attn_bias, attn_type)
# No alibi slopes.
# TODO(woosuk): Too many view operations. Let's try to reduce
# them in the future for code readability.
self.attn_op = xops.fmha.flash.FwOp()
if self.alibi_slopes is None:
# Add the batch dimension.
query = query.unsqueeze(0)
key = key.unsqueeze(0)
value = value.unsqueeze(0)
if self.head_size == 256:
# head_dim=256: ixformer flash_attn supports it natively
# Use ixinfer_flash_attn_unpad directly (bypasses FwOp check)
from ixformer.functions.flash_attn_lib import ixinfer_flash_attn_unpad as _ixf_unpad
total_q = query.shape[0] * query.shape[1]
total_k = key.shape[0] * key.shape[1]
q_flat = query.reshape(total_q, self.num_heads, self.head_size)
k_flat = key.reshape(total_k, self.num_kv_heads, self.head_size)
v_flat = value.reshape(total_k, self.num_kv_heads, self.head_size)
# Get cu_seqlens from attn_bias (or construct from seq_lens)
if isinstance(attn_bias[0], BlockDiagonalCausalMask):
cu_seqlens_q = attn_bias[0].q_seqinfo.seqstart.int().to(query.device)
cu_seqlens_k = attn_bias[0].k_seqinfo.seqstart.int().to(query.device)
max_seqlen_q = attn_bias[0].q_seqinfo.max_seqlen
max_seqlen_k = attn_bias[0].k_seqinfo.max_seqlen
else:
batch_size = query.shape[0]
seqlen = query.shape[1]
cu_seqlens_q = torch.arange(0, batch_size + 1, device=query.device, dtype=torch.int32) * seqlen
cu_seqlens_k = cu_seqlens_q
max_seqlen_q = seqlen
max_seqlen_k = seqlen
out = _ixf_unpad(q_flat, k_flat, v_flat,
cu_seqlens_q, cu_seqlens_k,
max_seqlen_q, max_seqlen_k,
True, self.scale, out=None)
out = out.view(query.shape[0], query.shape[1], self.num_heads, self.head_size)
elif self.head_size not in (32, 64, 128, 256):
out = self._run_sdpa_fallback(query, key, value, attn_metadata)
else:
out = xops.memory_efficient_attention_forward(
query,
key,
value,
attn_bias=attn_bias[0],
p=0.0,
scale=self.scale,
op=self.attn_op,
)
return out.view_as(original_query)
# Attention with alibi slopes.
# FIXME(woosuk): Because xformers does not support dynamic sequence
# lengths with custom attention bias, we process each prompt one by
# one. This is inefficient, especially when we have many short prompts.
assert attn_metadata.seq_lens is not None
output = torch.empty_like(original_query)
start = 0
for i, seq_len in enumerate(attn_metadata.seq_lens):
end = start + seq_len
out = xops.memory_efficient_attention_forward(
query[None, start:end],
key[None, start:end],
value[None, start:end],
attn_bias=attn_bias[i],
p=0.0,
scale=self.scale,
)
# TODO(woosuk): Unnecessary copy. Optimize.
output[start:end].copy_(out.view_as(original_query[start:end]))
start += seq_len
return output
def _make_alibi_bias(
alibi_slopes: torch.Tensor,
num_kv_heads: int,
dtype: torch.dtype,
seq_lens: List[int],
) -> List[AttentionBias]:
attn_biases: List[AttentionBias] = []
for seq_len in seq_lens:
bias = torch.arange(seq_len, dtype=dtype)
# NOTE(zhuohan): HF uses
# `bias = bias[None, :].repeat(seq_len, 1)`
# here. We find that both biases give the same results, but
# the bias below more accurately follows the original ALiBi
# paper.
# Calculate a matrix where each element represents ith element- jth
# element.
bias = bias[None, :] - bias[:, None]
padded_len = (seq_len + 7) // 8 * 8
num_heads = alibi_slopes.shape[0]
bias = torch.empty(
1, # batch size
num_heads,
seq_len,
padded_len,
device=alibi_slopes.device,
dtype=dtype,
)[:, :, :, :seq_len].copy_(bias)
bias.mul_(alibi_slopes[:, None, None])
if num_heads != num_kv_heads:
bias = bias.unflatten(1, (num_kv_heads, num_heads // num_kv_heads))
attn_biases.append(LowerTriangularMaskWithTensorBias(bias))
return attn_biases