336f3349ca
Keeps ALL infrastructure from the last 70 commits: - paged_attn.py: ixformer native v1/v2 decode dispatch (Output TPS impact) - protocol.py: extra='allow' (fixes ~180 rejected replay requests) - qwen3_5.py: .float() router_logits, chunk_recurrent, index_combine - 14 prebuilt .so (including corex_gdn_chunk_recurrent) - patch_xformers: flash_attn_varlen_func + profiling guard (>32K→Q-tiling) yaml: max-num-seqs=2, TOPK=1, gpu-mem=0.90, max-model-len=131072 No LD_PRELOAD, no expandable_segments, no blocks cap hacks.
475 lines
17 KiB
Python
475 lines
17 KiB
Python
"""
|
||
策略：顺序（per-sequence）fallback — 纯 PyTorch 数学实现
|
||
==========================================================
|
||
逐条序列用 matmul + softmax 手写 attention，完全绕开所有硬件
|
||
flash attention kernel（ixformer / cudnnFlashAttnForward）。
|
||
|
||
背景：
|
||
Iluvatar cudnnFlashAttnForward 存在两个已知问题：
|
||
1. 不支持 is_causal=True（报错）
|
||
2. 使用 attn_mask 路径时数值结果不正确（静默错误，输出全为"!"）
|
||
与华为昇腾 910B4 上 llama.cpp --flash-attn off 修复同类问题的原理相同。
|
||
纯数学路径（matmul + softmax）在任何 PyTorch 后端上结果都正确。
|
||
|
||
优点：
|
||
数值正确，不依赖任何硬件特定 attention kernel。
|
||
峰值显存 = max(seq_len)² × H × dtype_size，由 --max-model-len 控制。
|
||
|
||
缺点：
|
||
并发请求的 prefill attention 串行执行。
|
||
O(L²) 显存（无 flash attention 的 O(L) 优化）。
|
||
|
||
内存参考（fp16，H_local=6）：
|
||
max-model-len=4096 → 峰值 ~200 MB
|
||
max-model-len=8192 → 峰值 ~800 MB
|
||
max-model-len=16384 → 峰值 ~3.2 GB
|
||
|
||
额外 patch（arg_utils.py）：
|
||
vllm 0.6.3 在 max_model_len > 32K 时会自动开启 chunked prefill（无命令行
|
||
关闭选项），原意是防止 profiling OOM。但 _run_sdpa_fallback 已通过 Q-tiling
|
||
解决了该问题，chunked prefill 反而会把推理路径从 _run_sdpa_fallback 切换到
|
||
_forward_prefix_pytorch，属于不必要的行为变更，因此一并禁用该自动逻辑。
|
||
|
||
Deploy:
|
||
python3 modified_scripts/patch_xformers_sdpa_seq.py
|
||
"""
|
||
|
||
from patch_utils import package_root, replace_one_of, replace_once
|
||
|
||
VLLM_ROOT = package_root("vllm")
|
||
XFORMERS_PATH = VLLM_ROOT / "attention" / "backends" / "xformers.py"
|
||
ARG_UTILS_PATH = VLLM_ROOT / "engine" / "arg_utils.py"
|
||
LOGITS_PROC_PATH = (
|
||
VLLM_ROOT / "model_executor" / "layers" / "logits_processor.py")
|
||
OUTLINES_DECODING_PATH = (
|
||
VLLM_ROOT / "model_executor" / "guided_decoding" /
|
||
"outlines_decoding.py")
|
||
|
||
# _apply_logits_processors crashes when seq_groups is None (intermediate
|
||
# chunked-prefill chunks on the driver rank). Add an early-return guard.
|
||
_LP_OLD_BLOCK = """\
|
||
def _apply_logits_processors(
|
||
logits: torch.Tensor,
|
||
sampling_metadata: SamplingMetadata,
|
||
) -> torch.Tensor:
|
||
found_logits_processors = False\
|
||
"""
|
||
|
||
_LP_NEW_BLOCK = """\
|
||
def _apply_logits_processors(
|
||
logits: torch.Tensor,
|
||
sampling_metadata: SamplingMetadata,
|
||
) -> torch.Tensor:
|
||
if sampling_metadata.seq_groups is None: # intermediate chunked-prefill chunk
|
||
return logits
|
||
found_logits_processors = False\
|
||
"""
|
||
|
||
# Outlines' UNESCAPED_STRING accepts raw JSON control characters, including
|
||
# newlines and tabs. The generated text can therefore satisfy the CFG while
|
||
# still failing json.loads(). Use the RFC 8259 string character constraints.
|
||
_JSON_STRING_OLD_BLOCK = """\
|
||
| UNESCAPED_STRING
|
||
| SIGNED_NUMBER -> number
|
||
| "true" -> true
|
||
| "false" -> false
|
||
| "null" -> null
|
||
|
||
array : "[" [value ("," value)*] "]"
|
||
object : "{" [pair ("," pair)*] "}"
|
||
pair : UNESCAPED_STRING ":" value
|
||
|
||
%import common.UNESCAPED_STRING
|
||
%import common.SIGNED_NUMBER
|
||
%import common.WS
|
||
|
||
%ignore WS\
|
||
"""
|
||
|
||
_JSON_STRING_V1_BLOCK = r'''| JSON_STRING
|
||
| SIGNED_NUMBER -> number
|
||
| "true" -> true
|
||
| "false" -> false
|
||
| "null" -> null
|
||
|
||
array : "[" [value ("," value)*] "]"
|
||
object : "{" [pair ("," pair)*] "}"
|
||
pair : JSON_STRING ":" value
|
||
|
||
JSON_STRING: /"(\\["\\\/bfnrt]|\\u[0-9a-fA-F]{4}|[^"\\\x00-\x1f])*"/
|
||
%import common.SIGNED_NUMBER
|
||
%import common.WS
|
||
|
||
%ignore WS'''
|
||
|
||
_JSON_STRING_NEW_BLOCK = r'''| JSON_STRING
|
||
| SIGNED_NUMBER -> number
|
||
| "true" -> true
|
||
| "false" -> false
|
||
| "null" -> null
|
||
|
||
array : "[" _ws [value (_ws "," _ws value)*] _ws "]"
|
||
object : "{" _ws [pair (_ws "," _ws pair)*] _ws "}"
|
||
pair : JSON_STRING _ws ":" _ws value
|
||
_ws : JSON_WS?
|
||
|
||
JSON_STRING: /"(\\["\\\/bfnrt]|\\u[0-9a-fA-F]{4}|[^"\\\x00-\x1f])*"/
|
||
JSON_WS: /[ \t\r\n]{1,4}/
|
||
%import common.SIGNED_NUMBER'''
|
||
|
||
# vllm 0.6.3 自动开启 chunked prefill 的原始块
|
||
_ARG_OLD_BLOCK = """\
|
||
if (is_gpu and not use_sliding_window and not use_spec_decode
|
||
and not self.enable_lora
|
||
and not self.enable_prompt_adapter):
|
||
self.enable_chunked_prefill = True
|
||
logger.warning(
|
||
"Chunked prefill is enabled by default for models with "
|
||
"max_model_len > 32K. Currently, chunked prefill might "
|
||
"not work with some features or models. If you "
|
||
"encounter any issues, please disable chunked prefill "
|
||
"by setting --enable-chunked-prefill=False.")\
|
||
"""
|
||
|
||
_ARG_NEW_BLOCK = """\
|
||
if (is_gpu and not use_sliding_window and not use_spec_decode
|
||
and not self.enable_lora
|
||
and not self.enable_prompt_adapter):
|
||
pass # skip auto-enable: Q-tiling in _run_sdpa_fallback
|
||
# handles long-context memory without chunked prefill\
|
||
"""
|
||
|
||
_MM_PREFIX_OLD_BLOCK = """\
|
||
if model_config.is_multimodal_model:
|
||
if self.enable_prefix_caching:
|
||
logger.warning(
|
||
"--enable-prefix-caching is currently not "
|
||
"supported for multimodal models and has been disabled.")
|
||
self.enable_prefix_caching = False\
|
||
"""
|
||
|
||
_MM_PREFIX_NEW_BLOCK = """\
|
||
if model_config.is_multimodal_model:
|
||
architectures = getattr(model_config.hf_config,
|
||
"architectures", []) or []
|
||
qwen36_native_vision = "Qwen3_5MoeForCausalLM" in architectures
|
||
if self.enable_prefix_caching and qwen36_native_vision:
|
||
logger.info(
|
||
"Keeping prefix caching enabled for the Qwen3.6 native "
|
||
"vision path.")
|
||
elif self.enable_prefix_caching:
|
||
logger.warning(
|
||
"--enable-prefix-caching is currently not "
|
||
"supported for multimodal models and has been disabled.")
|
||
self.enable_prefix_caching = False\
|
||
"""
|
||
|
||
FALLBACK_METHOD = '''
|
||
# --- flash_attn_varlen_func backend (loaded once) ---
|
||
# Import path: ixformer.contrib.vllm_flash_attn (canonical, matches
|
||
# ex_engine/python/corex_fa2.py Tier 1 and ixformer_sdk).
|
||
# Signature ref: ixformer_sdk/contrib/vllm_flash_attn/flash_attn_interface.py
|
||
_flash_varlen_func = None
|
||
_flash_varlen_checked = False
|
||
|
||
@classmethod
|
||
def _get_flash_varlen(cls):
|
||
if not cls._flash_varlen_checked:
|
||
cls._flash_varlen_checked = True
|
||
try:
|
||
from ixformer.contrib.vllm_flash_attn import (
|
||
flash_attn_varlen_func as _fn,
|
||
)
|
||
cls._flash_varlen_func = _fn
|
||
except ImportError:
|
||
pass
|
||
return cls._flash_varlen_func
|
||
|
||
def _run_sdpa_fallback(
|
||
self,
|
||
query: torch.Tensor,
|
||
key: torch.Tensor,
|
||
value: torch.Tensor,
|
||
attn_metadata: "XFormersMetadata",
|
||
) -> torch.Tensor:
|
||
"""Prefill attention fallback for head_dim > 128.
|
||
|
||
Dispatch priority (ref: ex_engine/python/corex_fa2.py):
|
||
1. ixformer flash_attn_varlen_func — fused kernel, O(L) memory
|
||
2. Pure-math Q-tiling fallback — safe for profiling / any HW
|
||
|
||
Profiling guard: when kv_cache is empty (profiling stage), vllm feeds
|
||
a dummy sequence up to max_model_len (131K). flash_attn temp buffers
|
||
at that length can exceed GPU memory. We use Q-tiling for profiling
|
||
(safe, correct, O(chunk × L) memory) and flash_attn for real
|
||
inference (fast, O(L) memory, verified on BI-V100 head_dim=256).
|
||
|
||
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]
|
||
"""
|
||
assert attn_metadata.seq_lens is not None
|
||
orig_dtype = query.dtype
|
||
num_seqs = len(attn_metadata.seq_lens)
|
||
max_seqlen = max(attn_metadata.seq_lens)
|
||
|
||
# Detect profiling: attn_metadata.num_prefill_tokens == total tokens
|
||
# AND no actual KV cache allocated yet (first forward pass).
|
||
# Also guard against very long dummy sequences (profiling uses
|
||
# max_model_len which can be 131K) where flash_attn would OOM.
|
||
_FLASH_SAFE_SEQLEN = 32768 # flash_attn temp buffers safe below this
|
||
is_profiling = (max_seqlen > _FLASH_SAFE_SEQLEN
|
||
and not hasattr(attn_metadata, '_has_real_kv_cache'))
|
||
|
||
# --- Path 1: flash_attn_varlen_func (real inference) ---
|
||
fn = self._get_flash_varlen()
|
||
if fn is not None and not is_profiling:
|
||
try:
|
||
q_flat = query.squeeze(0) # [T, H, D]
|
||
k_flat = key.squeeze(0) # [T, Hkv, D]
|
||
v_flat = value.squeeze(0)
|
||
|
||
cu_seqlens = torch.zeros(
|
||
num_seqs + 1, dtype=torch.int32, device=query.device)
|
||
for i, sl in enumerate(attn_metadata.seq_lens):
|
||
cu_seqlens[i + 1] = cu_seqlens[i] + sl
|
||
|
||
out = fn(
|
||
q=q_flat.to(torch.float16),
|
||
k=k_flat.to(torch.float16),
|
||
v=v_flat.to(torch.float16),
|
||
cu_seqlens_q=cu_seqlens,
|
||
cu_seqlens_k=cu_seqlens,
|
||
max_seqlen_q=max_seqlen,
|
||
max_seqlen_k=max_seqlen,
|
||
softmax_scale=self.scale,
|
||
causal=True,
|
||
)
|
||
return out.to(orig_dtype).unsqueeze(0)
|
||
except Exception:
|
||
pass # fall through to Q-tiling
|
||
|
||
# --- Path 2: Q-tiling (profiling or flash_attn unavailable) ---
|
||
_Q_CHUNK = 256
|
||
|
||
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)
|
||
k_flat = key.squeeze(0)
|
||
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_s = k_flat[seq_start:seq_end].permute(1, 0, 2).float()
|
||
v_s = v_flat[seq_start:seq_end].permute(1, 0, 2).float()
|
||
|
||
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 = torch.arange(q_len, device=query.device)
|
||
|
||
for qc_start in range(0, q_len, _Q_CHUNK):
|
||
qc_end = min(qc_start + _Q_CHUNK, q_len)
|
||
|
||
q_c = q_flat[seq_start + qc_start:seq_start + qc_end] \
|
||
.permute(1, 0, 2).float()
|
||
|
||
attn_w = torch.matmul(q_c, k_s.transpose(-2, -1)) * self.scale
|
||
|
||
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)
|
||
|
||
output[seq_start + qc_start:seq_start + qc_end] = (
|
||
out_c.permute(1, 0, 2))
|
||
|
||
seq_start = seq_end
|
||
|
||
return output.unsqueeze(0)
|
||
|
||
'''
|
||
|
||
OLD_XFORMER_BLOCK = """\
|
||
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)
|
||
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)\
|
||
"""
|
||
|
||
NEW_XFORMER_BLOCK = """\
|
||
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 > 128:
|
||
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)\
|
||
"""
|
||
|
||
INJECT_ANCHOR = " def _run_memory_efficient_xformers_forward("
|
||
|
||
_PREFIX_CALL_OLD_BLOCK = """\
|
||
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,
|
||
)\
|
||
"""
|
||
|
||
_PREFIX_CALL_NEW_BLOCK = """\
|
||
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,
|
||
is_causal_decoder=(attn_type == AttentionType.DECODER),
|
||
)\
|
||
"""
|
||
|
||
|
||
def patch_file(path):
|
||
replace_once(
|
||
path,
|
||
INJECT_ANCHOR,
|
||
FALLBACK_METHOD + INJECT_ANCHOR,
|
||
required=True,
|
||
already_contains="def _run_sdpa_fallback(")
|
||
replace_once(
|
||
path,
|
||
OLD_XFORMER_BLOCK,
|
||
NEW_XFORMER_BLOCK,
|
||
required=True,
|
||
already_contains="out = self._run_sdpa_fallback(query, key, value, attn_metadata)")
|
||
replace_once(
|
||
path,
|
||
_PREFIX_CALL_OLD_BLOCK,
|
||
_PREFIX_CALL_NEW_BLOCK,
|
||
required=True,
|
||
already_contains=(
|
||
"is_causal_decoder=(attn_type == AttentionType.DECODER)"))
|
||
|
||
|
||
def patch_arg_utils(path):
|
||
replace_once(
|
||
path,
|
||
_ARG_OLD_BLOCK,
|
||
_ARG_NEW_BLOCK,
|
||
required=True,
|
||
already_contains="skip auto-enable: Q-tiling")
|
||
replace_once(
|
||
path,
|
||
_MM_PREFIX_OLD_BLOCK,
|
||
_MM_PREFIX_NEW_BLOCK,
|
||
required=True,
|
||
already_contains="Keeping prefix caching enabled for the Qwen3.6")
|
||
|
||
|
||
def patch_logits_processor(path):
|
||
replace_once(
|
||
path,
|
||
_LP_OLD_BLOCK,
|
||
_LP_NEW_BLOCK,
|
||
required=True,
|
||
already_contains="intermediate chunked-prefill chunk")
|
||
|
||
|
||
def patch_outlines_json_grammar(path):
|
||
replace_one_of(
|
||
path,
|
||
[
|
||
(_JSON_STRING_V1_BLOCK, _JSON_STRING_NEW_BLOCK),
|
||
(_JSON_STRING_OLD_BLOCK, _JSON_STRING_NEW_BLOCK),
|
||
],
|
||
required=True,
|
||
already_contains="JSON_WS:")
|
||
|
||
|
||
def main():
|
||
print("=== patch_xformers_sdpa_seq (sequential, pure-math) ===")
|
||
print(f"Target: {XFORMERS_PATH}")
|
||
patch_file(XFORMERS_PATH)
|
||
|
||
print("\n=== patch_arg_utils (disable chunked-prefill auto-enable) ===")
|
||
print(f"Target: {ARG_UTILS_PATH}")
|
||
patch_arg_utils(ARG_UTILS_PATH)
|
||
|
||
print("\n=== patch_logits_processor (seq_groups=None guard for chunked prefill) ===")
|
||
print(f"Target: {LOGITS_PROC_PATH}")
|
||
patch_logits_processor(LOGITS_PROC_PATH)
|
||
|
||
print("\n=== patch_outlines_json_grammar (reject raw control chars) ===")
|
||
print(f"Target: {OUTLINES_DECODING_PATH}")
|
||
patch_outlines_json_grammar(OUTLINES_DECODING_PATH)
|
||
|
||
print("\nDone.")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|