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