""" 策略：顺序（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()