From d0585b4a178d34b244506ed3a8948d809e45ede3 Mon Sep 17 00:00:00 2001 From: root Date: Tue, 14 Jul 2026 13:11:21 +0000 Subject: [PATCH] fix perforence --- computility-run.yaml | 30 +- paged_attn.py | 799 ++++++++++++++- .../chat_template_multi_system.jinja | 175 ++++ qwen3_6_scripts/paged_attn.py | 118 ++- qwen3_6_scripts/patch_ops.sh | 122 +++ qwen3_6_scripts/patch_xformers_sdpa_seq.py | 38 +- qwen3_6_scripts/protocol.py | 72 ++ qwen3_6_scripts/qwen3_5.py | 163 ++- qwen3_6_scripts/xformers.py | 946 ++++++++++++++++++ 9 files changed, 2414 insertions(+), 49 deletions(-) create mode 100644 qwen3_6_scripts/chat_template_multi_system.jinja create mode 100644 qwen3_6_scripts/xformers.py diff --git a/computility-run.yaml b/computility-run.yaml index 6447abd..d04963d 100644 --- a/computility-run.yaml +++ b/computility-run.yaml @@ -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' + diff --git a/paged_attn.py b/paged_attn.py index 988f903..a78fedd 100644 --- a/paged_attn.py +++ b/paged_attn.py @@ -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 diff --git a/qwen3_6_scripts/chat_template_multi_system.jinja b/qwen3_6_scripts/chat_template_multi_system.jinja new file mode 100644 index 0000000..3592c48 --- /dev/null +++ b/qwen3_6_scripts/chat_template_multi_system.jinja @@ -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" }} + {%- for tool in tools %} + {{- "\n" }} + {{- tool | tojson }} + {%- endfor %} + {{- "\n" }} + {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n\n\n\nvalue_1\n\n\nThis is the value for the second parameter\nthat can span\nmultiple lines\n\n\n\n\n\nReminder:\n- Function calls MUST follow the specified format: an inner block must be nested within 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' }} + {%- 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('') and content.endswith('')) %} + {%- 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 '' in content %} + {%- set reasoning_content = content.split('')[0].rstrip('\n').split('')[-1].lstrip('\n') %} + {%- set content = content.split('')[-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\n' + reasoning_content + '\n\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\n\n' }} + {%- else %} + {{- '\n\n' }} + {%- endif %} + {%- else %} + {{- '\n\n\n' }} + {%- endif %} + {%- if tool_call.arguments is defined %} + {%- for args_name, args_value in tool_call.arguments|items %} + {{- '\n' }} + {%- set args_value = args_value | string if args_value is string else args_value | tojson | safe %} + {{- args_value }} + {{- '\n\n' }} + {%- endfor %} + {%- endif %} + {{- '\n' }} + {%- endfor %} + {%- endif %} + {{- '<|im_end|>\n' }} + {%- elif message.role == "tool" %} + {%- if loop.previtem and loop.previtem.role != "tool" %} + {{- '<|im_start|>user' }} + {%- endif %} + {{- '\n\n' }} + {{- content }} + {{- '\n' }} + {%- 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 %} + {{- '\n\n\n\n' }} + {%- else %} + {{- '\n' }} + {%- endif %} +{%- endif %} \ No newline at end of file diff --git a/qwen3_6_scripts/paged_attn.py b/qwen3_6_scripts/paged_attn.py index 8590489..2d9ae1b 100644 --- a/qwen3_6_scripts/paged_attn.py +++ b/qwen3_6_scripts/paged_attn.py @@ -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, diff --git a/qwen3_6_scripts/patch_ops.sh b/qwen3_6_scripts/patch_ops.sh index 1035b2f..d6ef3ae 100755 --- a/qwen3_6_scripts/patch_ops.sh +++ b/qwen3_6_scripts/patch_ops.sh @@ -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