add log

README.md

@@ -1,118 +1,46 @@
-# 天数智芯 天垓100 文本生成引擎（基于 vLLM 优化）
-
-本项目是为**天数智芯-天垓100**加速卡深度优化的高性能文本生成推理引擎，基于开源 **vLLM** 框架进行架构级适配与增强，率先实现对 **Qwen3 系列**等最新大模型的高效支持。通过引入 **Prefix Caching**、PagedAttention 等先进优化技术，显著提升吞吐与响应速度，同时提供标准 **OpenAI 兼容 API 接口**，便于无缝集成现有应用生态。
-
-## 支持模型
-
-- **Qwen3**
-- **Llama3**
-- **DeepSeek-R1-Distill**
-- 其他兼容 vLLM 的 HuggingFace 模型（持续扩展中）
-
-> 模型下载地址：[https://modelscope.cn/models/Qwen](https://modelscope.cn/models/Qwen)  
-
----
-
-## Quick Start
-
-### 1. 模型下载
-
-从 ModelScope 下载所需模型（以 Qwen2.5-7B-Instruct 为例）：
-
-```bash
-modelscope download --model qwen/Qwen2.5-7B-Instruct README.md --local_dir /mnt/models/Qwen2.5-7B-Instruct
-```
-
-> ⚠️ 请确保模型路径在后续 Docker 启动时正确挂载。
-
----
-
-### 2. 拉取并构建 Docker 镜像
-
-我们提供已预装天垓100驱动与vLLM优化版本的Docker镜像：
+# 天数智芯 天垓100 文本生成引擎（基于 vLLM 优化适配Qwen3.6-35B-A3B）
 
 ```
 # 本地构建
-docker build -t enginex-iluvatar-vllm:bi100 -f Dockerfile .
+docker build -t enginex-iluvatar-vllm:bi100-qwen3.6 -f Dockerfile .
 ```
 
----
 
-### 3. 启动服务容器
+启动容器镜像
 
-```bash
-docker run -it --rm -p 8000:80 \
-  --name vllm-iluvatar \
-  -v /mnt/models/Qwen2.5-7B-Instruct:/model:ro \
-  --privileged \
-  -e TENSOR_PARALLEL_SIZE=1 \
-  -e PREFIX_CACHING=true \
-  -e MAX_MODEL_LEN=10000 \
-  enginex-iluvatar-vllm:bi100
+下载Qwen3.6-35B-A3B模型，并且需要将模型的config.json文件中architectures字段改成
+```json
+"architectures": [
+        "Qwen3_5MoeForCausalLM"
+    ]
 ```
 
-> ✅ 参数说明：
-> - `PREFIX_CACHING=true`: 启用 Prefix Caching 优化，显著提升多请求共享前缀的推理效率
-> - `MAX_MODEL_LEN=10000`: 支持长上下文推理
-> - `--privileged`: 确保天垓100设备可见
-
----
-
-## 4. 测试服务（使用 OpenAI 兼容接口）
-
-服务启动后，可通过标准 OpenAI SDK 或 `curl` 进行测试。
-
-### 示例：文本生成请求
-
 ```bash
-curl http://localhost:8000/v1/chat/completions \
+docker run -dit --network=host --ipc=host \
+  -v /usr/src:/usr/src -v /lib/modules:/lib/modules -v /dev:/dev --privileged \
+  -v /mnt/disk1/models/Qwen3.6-35B-A3B:/model:ro --entrypoint=python3 \
+  -e CUDA_VISIBLE_DEVICES=4,5,6,7 -e VLLM_ENGINE_ITERATION_TIMEOUT_S=3600 \
+  enginex-iluvatar-vllm:bi100-qwen3.6 \
+  -m vllm.entrypoints.openai.api_server \
+  --model /model --port 1111 --served-model-name llm \
+  --max-model-len 100000 --trust-remote-code -tp 4 --gpu-memory-utilization 0.95 \
+  --max-num-seqs 1 --disable-log-requests --disable-frontend-multiprocessing \
+  --max-num-batched-tokens 4096 --enable-chunked-prefill \
+  --max-seq-len-to-capture 32768 --enable-auto-tool-choice \
+  --tool-call-parser qwen3_coder --reasoning-parser qwen3
+```
+
+请求
+```bash
+curl http://localhost:1111/v1/chat/completions \
   -H "Content-Type: application/json" \
   -d '{
-    "model": "qwen3-8b",
+    "model": "llm",
     "messages": [
       {"role": "system", "content": "You are a helpful assistant."},
-      {"role": "user", "content": "请用中文介绍一下上海的特点。"}
+      {"role": "user", "content": "Can you tell me the story of Snow White?"}
     ],
-    "temperature": 0.7,
-    "max_tokens": 512
+    "max_tokens": 200,
+    "temperature": 0.7
   }'
-```
-
-### 使用 OpenAI Python SDK（需安装 `openai>=1.0`）
-
-```python
-from openai import OpenAI
-
-client = OpenAI(base_url="http://localhost:8000/v1", api_key="none")
-
-response = client.chat.completions.create(
-    model="qwen3-8b",
-    messages=[
-        {"role": "system", "content": "You are a helpful assistant."},
-        {"role": "user", "content": "请简要介绍杭州的特色文化。"}
-    ],
-    max_tokens=512,
-    temperature=0.7
-)
-
-print(response.choices[0].message.content)
-```
-
----
-
-## 测试结果对比（A100 vs 天垓100）
-
-### 测试数据集
-
-[chat_dataset_v0.json](chat_dataset_v0.json)
-
-### 测试结果
-
-在相同模型和输入条件下，测试平均输出速度（单位：字每秒），结果如下：
-
-| 模型 | 天垓100 输出速度 | Nvidia A100 输出速度 |
-|--------|--------------------------|-------------------------------|
-| Qwen2.5-7B-Instruct | 36.8 | 112.4 |
-| Qwen2.5-1.5B-Instruct-AWQ | 72.4 | 100.8 |
-| Qwen/Qwen1.5-32B-Chat | 12.4 | 55.7 |
-
+```

README_qwen3_6.md

@@ -1,42 +0,0 @@
-# 天数智芯 天垓100 文本生成引擎（基于 vLLM 优化适配Qwen3.6-27B）
-
-```
-# 本地构建
-docker build -t enginex-iluvatar-vllm:bi100-qwen3.6 -f Dockerfile .
-```
-
-
-启动容器镜像
-
-下载Qwen3.6-27B模型，并且需要将模型的config.json文件中architectures字段改成
-```json
-"architectures": [
-        "Qwen3_5ForCausalLM"
-    ]
-```
-
-```bash
-docker run -dit --network=host --ipc=host \
-  -v /usr/src:/usr/src -v /lib/modules:/lib/modules -v /dev:/dev --privileged \
-  --name vllm-iluvatar \
-  -v /mnt/models/Qwen3.6-27B:/model:ro --entrypoint=python3 \
-  enginex-iluvatar-vllm:bi100 \
-  -m vllm.entrypoints.openai.api_server \
-  --model /model --port 1111 --served-model-name llm \
-  --max-model-len 10000 --enforce-eager --trust-remote-code -tp 4 --gpu-memory-utilization 0.95
-```
-
-请求
-```bash
-curl http://localhost:1111/v1/chat/completions \
-  -H "Content-Type: application/json" \
-  -d '{
-    "model": "llm",
-    "messages": [
-      {"role": "system", "content": "You are a helpful assistant."},
-      {"role": "user", "content": "Can you tell me the story of Snow White?"}
-    ],
-    "max_tokens": 200,
-    "temperature": 0.7
-  }'
-```

computility-run.yaml

@@ -0,0 +1,53 @@
+concurrency: 1
+command:
+    - python3
+    - -m
+    - vllm.entrypoints.openai.api_server
+    - --model
+    - /model
+    - --served-model-name
+    - llm
+    - --max-model-len
+    - '245000'
+    - --gpu-memory-utilization
+    - '0.9'
+    - --trust-remote-code
+    - -tp
+    - '4'
+    - --max-num-seqs
+    - '1'
+    - --disable-log-requests
+    - --disable-frontend-multiprocessing
+    - --max-num-batched-tokens
+    - '16384'
+    - --enable-chunked-prefill
+    - --max-seq-len-to-capture
+    - '245000'
+    - --enable-auto-tool-choice
+    - --tool-call-parser
+    - qwen3_coder
+    - --reasoning-parser
+    - qwen3
+    - --enable-prefix-caching
+    - --chat-template
+    - /workspace/chat_template_multi_system.jinja
+env:
+    - name: VLLM_ENGINE_ITERATION_TIMEOUT_S
+      value: 3600
+    - 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
+    - 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'
+

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

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<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>" }}
+    {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n</IMPORTANT>' }}
+    {%- if messages[0].role == 'system' %}
+        {%- set content = render_content(messages[0].content, false, true)|trim %}
+        {%- if content %}
+            {{- '\n\n' + content }}
+        {%- endif %}
+    {%- endif %}
+    {{- '<|im_end|>\n' }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {%- set content = render_content(messages[0].content, false, true)|trim %}
+        {{- '<|im_start|>system\n' + content + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" %}
+        {%- set content = render_content(message.content, false)|trim %}
+        {%- if not(content.startswith('<tool_response>') and content.endswith('</tool_response>')) %}
+            {%- set ns.multi_step_tool = false %}
+            {%- set ns.last_query_index = index %}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if ns.multi_step_tool %}
+    {{- raise_exception('No user query found in messages.') }}
+{%- endif %}
+{%- for message in messages %}
+    {%- set content = render_content(message.content, true)|trim %}
+    {%- if message.role == "system" %}
+        {%- if not loop.first %}
+            {{- raise_exception('System message must be at the beginning.') }}
+        {%- endif %}
+    {%- elif message.role == "user" %}
+        {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- set reasoning_content = reasoning_content|trim %}
+        {%- if (preserve_thinking is defined and preserve_thinking is true) or (loop.index0 > ns.last_query_index) %}
+            {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content + '\n</think>\n\n' + content }}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' + content }}
+        {%- endif %}
+        {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if tool_call.function is defined %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {%- if loop.first %}
+                    {%- if content|trim %}
+                        {{- '\n\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                    {%- else %}
+                        {{- '<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                    {%- endif %}
+                {%- else %}
+                    {{- '\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                {%- endif %}
+                {%- if tool_call.arguments is defined %}
+                    {%- for args_name, args_value in tool_call.arguments|items %}
+                        {{- '<parameter=' + args_name + '>\n' }}
+                        {%- set args_value = args_value | string if args_value is string else args_value | tojson | safe %}
+                        {{- args_value }}
+                        {{- '\n</parameter>\n' }}
+                    {%- endfor %}
+                {%- endif %}
+                {{- '</function>\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.previtem and loop.previtem.role != "tool" %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- if not loop.last and loop.nextitem.role != "tool" %}
+            {{- '<|im_end|>\n' }}
+        {%- elif loop.last %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- else %}
+        {{- raise_exception('Unexpected message role.') }}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+    {%- if enable_thinking is defined and enable_thinking is false %}
+        {{- '<think>\n\n</think>\n\n' }}
+    {%- else %}
+        {{- '<think>\n' }}
+    {%- endif %}
+{%- endif %}

qwen3_6_scripts/chat_utils.py

@@ -113,6 +113,11 @@ class ConversationMessage(TypedDict, total=False):
     tool_calls: Optional[Iterable[ChatCompletionMessageToolCallParam]]
     """The tool calls generated by the model, such as function calls."""
 
+    reasoning_content: Optional[str]
+    """Reasoning / thinking content for assistant messages.
+    Passed directly to the chat template (Qwen3 reads message.reasoning_content
+    natively) instead of being manually wrapped in <think>...</think>."""
+
 
 ModalityStr = Literal["image", "audio", "video"]
 _T = TypeVar("_T")
@@ -480,15 +485,13 @@ def _parse_chat_message_content(
             if "tool_calls" in parsed_msg:
                 result_msg["tool_calls"] = list(parsed_msg["tool_calls"])
 
-            # Prepend reasoning_content as <think>...</think> so the model
-            # sees its own chain-of-thought in multi-turn conversations.
-            reasoning = message.get("reasoning_content")  # type: ignore[arg-type]
+            # Pass reasoning content as a dedicated field so the chat template
+            # can render it natively (Qwen3: message.reasoning_content branch).
+            # Accept both "reasoning" (new vllm) and "reasoning_content" (ours).
+            reasoning = (message.get("reasoning")  # type: ignore[arg-type]
+                         or message.get("reasoning_content"))  # type: ignore[arg-type]
             if reasoning and isinstance(reasoning, str):
-                existing = result_msg.get("content") or ""
-                result_msg["content"] = (
-                    f"<think>{reasoning}</think>\n\n{existing}"
-                    if existing else f"<think>{reasoning}</think>"
-                )
+                result_msg["reasoning_content"] = reasoning
 
         elif role == "tool":
             parsed_msg = _ToolParser(message)

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,
@@ -393,6 +509,20 @@ class PagedAttention:
                 # --------------------------------------------------------------
                 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]

qwen3_6_scripts/patch_model_runner.py

@@ -0,0 +1,78 @@
+"""
+Fix: prefix_cache_hit stays True for chunked-prefill chunk 2+ even when past cache.
+
+Root cause:
+  model_runner.py _compute_for_prefix_cache_hit has three cases:
+    Case 1: prefix_cache_len <= context_len  → "already past cache, do normal"
+    Case 2: context_len < prefix_cache_len < seq_len  → partial hit, correct
+    Case 3: seq_len <= prefix_cache_len  → full hit, reduce to 1 token
+
+  Case 1 does nothing (leaves prefix_cache_hit = True). Then in utils.py:
+    if inter_data.prefix_cache_hit:
+        block_table = computed_block_nums   ← ONLY the original prefix blocks!
+
+  But context_len > prefix_cache_len means chunk 1 tokens (between prefix_cache_len
+  and context_len) are ALSO in KV cache and need to be in block_table.
+  block_table = computed_block_nums misses all chunk-1 blocks.
+
+  In _forward_prefix_pytorch:
+    num_ctx_blocks = ceil(context_len / block_size)  # e.g. 268
+    block_tables.shape[1] = len(computed_block_nums)  # e.g. 12  <-- too small!
+    At tile_blk >= 12: blk_ids is empty → k_t shape [..., 0] → amax crash.
+
+Fix:
+  Set prefix_cache_hit = False for Case 1, so utils.py falls through to:
+    elif chunked_prefill_enabled:
+        block_table = block_tables[seq_id]   ← full block table (prefix + chunk1)
+"""
+
+import re
+import sys
+
+CANDIDATE_PATHS = [
+    "/usr/local/corex/lib64/python3/dist-packages/vllm/worker/model_runner.py",
+    "/usr/local/corex/lib/python3/dist-packages/vllm/worker/model_runner.py",
+]
+
+OLD_BLOCK = """\
+        if prefix_cache_len <= context_len:
+            # We already passed the cache hit region,
+            # so do normal computation.
+            pass"""
+
+NEW_BLOCK = """\
+        if prefix_cache_len <= context_len:
+            # We already passed the cache hit region,
+            # so do normal computation.
+            # Must clear prefix_cache_hit so _add_seq_group uses the full
+            # block_tables (prefix + previous-chunk blocks) instead of only
+            # computed_block_nums (prefix only).  Without this, block_tables
+            # passed to _forward_prefix_pytorch is too narrow for context_len,
+            # causing an empty blk_ids slice and a zero-dim amax() crash.
+            inter_data.prefix_cache_hit = False"""
+
+import os
+
+patched = False
+for path in CANDIDATE_PATHS:
+    if not os.path.exists(path):
+        continue
+    with open(path, "r") as f:
+        src = f.read()
+    if OLD_BLOCK not in src:
+        if NEW_BLOCK in src:
+            print(f"[patch_model_runner] already patched: {path}")
+            patched = True
+            break
+        print(f"[patch_model_runner] WARNING: expected block not found in {path}, skipping")
+        continue
+    patched_src = src.replace(OLD_BLOCK, NEW_BLOCK, 1)
+    with open(path, "w") as f:
+        f.write(patched_src)
+    print(f"[patch_model_runner] patched Case-1 prefix_cache_hit fix in: {path}")
+    patched = True
+    break
+
+if not patched:
+    print("[patch_model_runner] ERROR: could not find model_runner.py at any known path", file=sys.stderr)
+    sys.exit(1)

qwen3_6_scripts/patch_ops.sh

@@ -8,8 +8,6 @@
 #     are already correct for standard Triton 2.3.1 — do NOT overwrite them.
 #   - DO NOT install BI-V150 corex Triton 2.1.0 (pkgs/triton): that causes
 #     GPU hang on BI-V100 because the Triton CUDA PTX kernels are incompatible.
-#
-# Important Note: Qwen3.6-27B must apply TP=4,PP=2 combination in order to deploy using 8 GPUs
 
 # Recommended server start command for TP=4 support 100K, need chunked prefill
 # CUDA_VISIBLE_DEVICES="4,5,6,7" VLLM_ENGINE_ITERATION_TIMEOUT_S=3600 python3 -m vllm.entrypoints.openai.api_server \
@@ -17,6 +15,53 @@
 #     --max-model-len 100000 --enforce-eager --trust-remote-code -tp 4 --gpu-memory-utilization 0.95 \
 #     --max-num-seqs 1 --disable-log-requests --disable-frontend-multiprocessing \
 #     --max-num-batched-tokens 4096 --enable-chunked-prefill
+#
+# With prefix caching (GDN align-mode, requires chunked prefill):
+# CUDA_VISIBLE_DEVICES="4,5,6,7" VLLM_ENGINE_ITERATION_TIMEOUT_S=3600 python3 -m vllm.entrypoints.openai.api_server \
+#     --model /workspace/models/Qwen3.6-35B-A3B --port 1111 --served-model-name llm \
+#     --max-model-len 150000 --trust-remote-code -tp 4 --gpu-memory-utilization 0.90 \
+#     --max-num-seqs 1 --disable-log-requests --disable-frontend-multiprocessing \
+#     --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
@@ -26,6 +71,15 @@
 # when context length is high
 cp ./paged_attn.py /usr/local/corex/lib/python3/dist-packages/vllm/attention/ops/paged_attn.py
 
+# --- model_runner.py: fix prefix_cache_hit stays True in chunked-prefill chunk 2+ ---
+# Bug: _compute_for_prefix_cache_hit Case 1 (prefix_cache_len <= context_len)
+# leaves prefix_cache_hit=True. Then _add_seq_group uses block_table=computed_block_nums
+# (only the original prefix blocks), ignoring chunk-1 KV cache blocks.
+# _forward_prefix_pytorch then gets an undersized block_tables and crashes with
+# "amax(): Expected reduction dim -1 to have non-zero size" on the 2nd tile.
+# Fix: set prefix_cache_hit=False for Case 1 so the full block_tables is used.
+python3 ./patch_model_runner.py
+
 # --- transformers: Qwen3_5 tokenizer / model files --------------------------
 pip install transformers==4.55.3 -i https://pypi.tuna.tsinghua.edu.cn/simple
 cp -r ./qwen3_5 /usr/local/lib/python3.10/site-packages/transformers/models/
@@ -42,8 +96,15 @@ python3 ./patch_vllm_qwen3_5.py
 # returns _cached_all_token_ids[-0:] == [0:] (the ENTIRE prompt+output list).
 # Each prefill chunk step adds prompt_len to previous_num_tokens, so a 10K
 # prompt processed in 3 chunks inflates completion_tokens by ~30K.
+# Also adds num_cached_tokens field to RequestMetrics for prefix-cache stats.
 cp ./sequence.py /usr/local/corex/lib/python3/dist-packages/vllm/sequence.py
 
+# --- scheduler.py: record num_cached_tokens in RequestMetrics ----------------
+# Sets seq_group.metrics.num_cached_tokens = prefix_cache_len on first prefill
+# when --enable-prefix-caching is active, so serving_chat.py can report it in
+# usage.prompt_tokens_details.cached_tokens (OpenAI-compatible API response).
+cp ./scheduler.py /usr/local/corex/lib/python3/dist-packages/vllm/core/scheduler.py
+
 # --- xformers: bypass cudnnFlashAttnForward (head_dim=256 > 128 limit) ------
 # Injects _run_sdpa_fallback (pure matmul+softmax) into xformers.py.
 # Required because head_dim=256 > 128 and ixformer flash attention either
@@ -70,3 +131,86 @@ cp ./cli_args.py /usr/local/corex/lib/python3/dist-packages/vllm/entrypoints/ope
 cp ./serving_chat.py /usr/local/corex/lib/python3/dist-packages/vllm/entrypoints/openai/serving_chat.py
 cp ./api_server.py /usr/local/corex/lib/python3/dist-packages/vllm/entrypoints/openai/api_server.py
 cp ./chat_utils.py /usr/local/corex/lib/python3/dist-packages/vllm/entrypoints/chat_utils.py
+
+# --- chat_template: multi_system 兼容 (合并多/乱序 system 到首位) -------------
+# 数据集 34 条请求含多个或乱序 system 消息, 原模型 chat_template.jinja 对非首位
+# system raise "System message must be at the beginning." → 15 条 4xx, 威胁成功率≥99%。
+# 模型目录 (/root/public-storage 或容器内 /model) 是只读挂载, 不能改原模板。
+# 改后模板随 qwen3_6_scripts 一起 COPY 进来 (本目录 ./chat_template_multi_system.jinja),
+# 这里 cp 到 /workspace/ (Dockerfile 已 mkdir /workspace), 启动时用
+# --chat-template /workspace/chat_template_multi_system.jinja 指向它 (见 computility-run.yaml)。
+mkdir -p /workspace
+CHAT_TMPL=/workspace/chat_template_multi_system.jinja
+if [ -f "./chat_template_multi_system.jinja" ]; then
+    cp "./chat_template_multi_system.jinja" "$CHAT_TMPL"
+    echo "[patch_ops] deployed multi_system chat template → $CHAT_TMPL"
+else
+    echo "[patch_ops] WARNING: ./chat_template_multi_system.jinja not found in $(pwd)" >&2
+fi
+
+# --- verification: five optimization flags present in deployed package -------
+VLLM_PKG=/usr/local/corex/lib64/python3/dist-packages/vllm
+echo ""
+echo "=== patch_ops verification ==="
+_fail=0
+check() {  # check  <grep_pattern>  <file>  <label>
+    if grep -q "$1" "$2" 2>/dev/null; then
+        echo "  [ok]   $3"
+    else
+        echo "  [MISS] $3  ($1 in $2)"; _fail=1
+    fi
+}
+check "MOE_NATIVE"        "$VLLM_PKG/model_executor/models/qwen3_5.py" "MOE_NATIVE (qwen3_5.py)"
+check "_native_experts_sorted" "$VLLM_PKG/model_executor/models/qwen3_5.py" "_native_experts_sorted"
+check "CHUNK_PARALLEL"   "$VLLM_PKG/model_executor/models/qwen3_5.py" "CHUNK_PARALLEL"
+check "LOOP1_NATIVE"     "$VLLM_PKG/model_executor/models/qwen3_5.py" "LOOP1_NATIVE"
+check "RMSNORM_NATIVE"   "$VLLM_PKG/model_executor/models/qwen3_5.py" "RMSNORM_NATIVE"
+check "PREFIX_FLASH"     "$VLLM_PKG/attention/ops/paged_attn.py"      "PREFIX_FLASH (paged_attn.py)"
+check "_forward_prefix_flash" "$VLLM_PKG/attention/ops/paged_attn.py" "_forward_prefix_flash"
+
+# libcusolver link present (else LOOP1_NATIVE silently falls back at runtime)
+if [ -e /opt/sw_home/local/cuda/lib64/libcusolver.so ]; then
+    echo "  [ok]   libcusolver.so link (LOOP1_NATIVE runtime)"
+else
+    echo "  [MISS] libcusolver.so link — LOOP1_NATIVE will fall back to serial loop"; _fail=1
+fi
+
+# multi_system chat template deployed (else multi-system requests 4xx at runtime)
+if [ -f /workspace/chat_template_multi_system.jinja ]; then
+    echo "  [ok]   /workspace/chat_template_multi_system.jinja (multi_system fix)"
+else
+    echo "  [MISS] /workspace/chat_template_multi_system.jinja — multi-system requests will 4xx"; _fail=1
+fi
+
+# solve_triangular actually callable (catches a broken cusolver link early)
+if python3 -c "import torch,os;os.environ.setdefault('CUDA_VISIBLE_DEVICES','0');a=torch.eye(4,device='cuda');torch.linalg.solve_triangular(-a,a,upper=False,unitriangular=True)" 2>/dev/null; then
+    echo "  [ok]   torch.linalg.solve_triangular callable"
+else
+    echo "  [MISS] torch.linalg.solve_triangular not callable — check libcusolver"; _fail=1
+fi
+
+if [ $_fail -eq 0 ]; then
+    echo "  => all five optimizations deployed and enabled"
+else
+    echo "  => WARNING: some checks failed — review above" >&2
+fi
+
+# --- server start command lives in computility-run.yaml ----------------------
+# patch_ops.sh only configures the environment + deploys patched files. It does
+# NOT start the server. The server launch command (python3 -m vllm... with all
+# CLI args) and the env block (PYTHONPATH / LD_LIBRARY_PATH / PATH / five
+# *_NATIVE flags / VLLM_ENGINE_ITERATION_TIMEOUT_S) are in ../computility-run.yaml.
+#
+# After running this script, start the service via the computility runner, which
+# must run from a NON-qwen3_6_scripts directory (this dir's xformers.py shadows
+# the real xformers package).
+#
+# Sanity check after startup: the log must NOT contain any "* FALLBACK" line
+# during cold prefill (prefix_flash FALLBACK on cache-hit requests is expected —
+# the pad-ratio router intentionally falls back to pure PyTorch there).
+#   grep "FALLBACK" <server_log>   # empty during cold prefill = all five active
+# 80K cold prefill TTFT ≈ 57s with all five on (cached=0).
+echo ""
+echo "=== patch complete. Start the server with the computility runner using: ==="
+echo "    ../computility-run.yaml"
+echo "  (env + command are defined there; run from outside qwen3_6_scripts/)"

qwen3_6_scripts/patch_xformers_sdpa_seq.py

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

qwen3_6_scripts/protocol.py

@@ -1,5 +1,6 @@
 # Adapted from
 # https://github.com/lm-sys/FastChat/blob/168ccc29d3f7edc50823016105c024fe2282732a/fastchat/protocol/openai_api_protocol.py
+import json
 import time
 from argparse import Namespace
 from typing import Any, Dict, List, Literal, Optional, Union
@@ -99,11 +100,16 @@ class ModelList(OpenAIBaseModel):
     data: List[ModelCard] = Field(default_factory=list)
 
 
+class PromptTokensDetails(OpenAIBaseModel):
+    cached_tokens: int = 0
+
+
 class UsageInfo(OpenAIBaseModel):
     prompt_tokens: int = 0
     total_tokens: int = 0
     completion_tokens: Optional[int] = 0
     reasoning_tokens: Optional[int] = None
+    prompt_tokens_details: Optional[PromptTokensDetails] = None
 
 
 class RequestResponseMetadata(BaseModel):
@@ -249,6 +255,15 @@ class ChatCompletionRequest(OpenAIBaseModel):
         description=("Additional kwargs to pass to the template renderer. "
                      "Will be accessible by the chat template."),
     )
+    thinking: Optional[Any] = Field(
+        default=None,
+        description=(
+            "Per-request thinking switch (OpenAI-style top-level). Accepts: "
+            "{'type':'disabled'} / {'type':'enabled'} / bool true/false. "
+            "Translated to chat_template_kwargs.enable_thinking (bool) so the "
+            "Qwen3.6 template's enable_thinking branch controls reasoning. "
+            "Null/absent = server default (thinking on via --reasoning-parser)."),
+    )
     guided_json: Optional[Union[str, dict, BaseModel]] = Field(
         default=None,
         description=("If specified, the output will follow the JSON schema."),
@@ -315,12 +330,20 @@ class ChatCompletionRequest(OpenAIBaseModel):
             prompt_logprobs = self.top_logprobs
 
         guided_json_object = None
-        if (self.response_format is not None
-                and self.response_format.type == "json_object"):
-            guided_json_object = True
+        guided_json_from_schema = None
+        if self.response_format is not None:
+            if self.response_format.type == "json_object":
+                guided_json_object = True
+            elif (self.response_format.type == "json_schema"
+                  and self.response_format.json_schema is not None
+                  and self.response_format.json_schema.json_schema is not None):
+                guided_json_from_schema = \
+                    self.response_format.json_schema.json_schema
 
         guided_decoding = GuidedDecodingParams.from_optional(
-            json=self._get_guided_json_from_tool() or self.guided_json,
+            json=(self._get_guided_json_from_tool()
+                  or self.guided_json
+                  or guided_json_from_schema),
             regex=self.guided_regex,
             choice=self.guided_choice,
             grammar=self.guided_grammar,
@@ -373,6 +396,55 @@ class ChatCompletionRequest(OpenAIBaseModel):
 
         return None
 
+    @model_validator(mode="before")
+    @classmethod
+    def normalize_messages(cls, data):
+        """Normalize incoming messages before pydantic union validation.
+
+        Real-world clients (e.g. from other providers) send assistant tool_call
+        messages with content=null, which fails the strict Union type check.
+        Replace null content with "" so validation passes.
+        reasoning_content is intentionally kept — chat_utils.py wraps it as
+        <think>...</think> for multi-turn reasoning history.
+        """
+        messages = data.get("messages")
+        if not isinstance(messages, list):
+            return data
+        normalized = []
+        for msg in messages:
+            if not isinstance(msg, dict):
+                normalized.append(msg)
+                continue
+            if msg.get("content") is None:
+                if msg.get("reasoning_content") is None:
+                    raise ValueError(
+                        "Each message must have at least one of 'content' or "
+                        "'reasoning_content'.")
+                msg = {**msg, "content": ""}
+            # tool_calls arguments: dict -> JSON string.
+            # Many clients/datasets send function.arguments as a dict (e.g.
+            # {"cmd":"ls"}), but upstream ChatCompletionMessageParam strictly
+            # requires a JSON string. Convert here so validation passes (matches
+            # OpenAI's lenient acceptance of both forms). Without this, dataset
+            # requests with tool messages 4xx with
+            # "function.arguments: Input should be a valid string".
+            tool_calls = msg.get("tool_calls")
+            if isinstance(tool_calls, list):
+                new_tc = []
+                for tc in tool_calls:
+                    if isinstance(tc, dict) and isinstance(tc.get("function"), dict):
+                        fn = tc["function"]
+                        args = fn.get("arguments")
+                        if isinstance(args, dict):
+                            fn = {**fn, "arguments": json.dumps(
+                                args, ensure_ascii=False)}
+                            tc = {**tc, "function": fn}
+                    new_tc.append(tc)
+                msg = {**msg, "tool_calls": new_tc}
+            normalized.append(msg)
+        data = {**data, "messages": normalized}
+        return data
+
     @model_validator(mode="before")
     @classmethod
     def validate_stream_options(cls, data):
@@ -382,6 +454,48 @@ class ChatCompletionRequest(OpenAIBaseModel):
 
         return data
 
+    @model_validator(mode="before")
+    @classmethod
+    def validate_thinking(cls, data):
+        """Translate top-level `thinking` → chat_template_kwargs.enable_thinking.
+
+        Accepts {'type':'disabled'|'enabled'} or bool. The Qwen3.6 template's
+        `enable_thinking is false` branch (chat_template.jinja:149) gates the
+        reasoning prefix, so a bool is what actually controls it. Merged into
+        chat_template_kwargs (request-level value wins over any pre-existing
+        enable_thinking). Without this, top-level `thinking` 4xx with
+        extra_forbidden (OpenAIBaseModel has extra="forbid").
+        """
+        thinking = data.get("thinking")
+        if thinking is None:
+            return data
+        # Normalize to a bool.
+        enable: Optional[bool]
+        if isinstance(thinking, bool):
+            enable = thinking
+        elif isinstance(thinking, dict):
+            t = str(thinking.get("type", "")).lower()
+            if t == "disabled":
+                enable = False
+            elif t == "enabled" or t == "auto":
+                enable = True
+            elif t == "":
+                enable = None
+            else:
+                raise ValueError(
+                    f"thinking.type must be 'enabled'/'disabled'/'auto', "
+                    f"got {thinking.get('type')!r}")
+        else:
+            raise ValueError(
+                f"thinking must be bool or {{'type':'disabled'|'enabled'}}, "
+                f"got {type(thinking).__name__}")
+        if enable is not None:
+            ctk = dict(data.get("chat_template_kwargs") or {})
+            ctk["enable_thinking"] = enable
+            data["chat_template_kwargs"] = ctk
+        data.pop("thinking", None)
+        return data
+
     @model_validator(mode="before")
     @classmethod
     def check_logprobs(cls, data):
@@ -609,12 +723,18 @@ class CompletionRequest(OpenAIBaseModel):
         echo_without_generation = self.echo and self.max_tokens == 0
 
         guided_json_object = None
-        if (self.response_format is not None
-                and self.response_format.type == "json_object"):
-            guided_json_object = True
+        guided_json_from_schema = None
+        if self.response_format is not None:
+            if self.response_format.type == "json_object":
+                guided_json_object = True
+            elif (self.response_format.type == "json_schema"
+                  and self.response_format.json_schema is not None
+                  and self.response_format.json_schema.json_schema is not None):
+                guided_json_from_schema = \
+                    self.response_format.json_schema.json_schema
 
         guided_decoding = GuidedDecodingParams.from_optional(
-            json=self.guided_json,
+            json=self.guided_json or guided_json_from_schema,
             regex=self.guided_regex,
             choice=self.guided_choice,
             grammar=self.guided_grammar,

qwen3_6_scripts/qwen3_5.py

@@ -2,7 +2,8 @@
 # Pure-PyTorch DeltaNet (no fla / causal_conv1d dependency).
 # Text-only (no VL, no MTP).
 
-from typing import Iterable, List, Optional, Tuple
+from collections import OrderedDict
+from typing import Dict, Iterable, List, Optional, Tuple
 
 import torch
 import torch.nn.functional as F
@@ -41,6 +42,21 @@ from vllm.model_executor.models.interfaces import HasInnerState, SupportsLoRA
 logger = init_logger(__name__)
 
 
+# ---------------------------------------------------------------------------
+# Profiling debug prints (locate where startup hangs after weight load).
+# Gate on PROF_TRACE=1 (default on, since we are diagnosing a startup hang).
+# Each print flushes; grep "[PROF]" in the server log.
+# ---------------------------------------------------------------------------
+import os as _os_prof
+import sys as _sys_prof
+import time as _time_prof
+_PROF_TRACE = _os_prof.environ.get("PROF_TRACE", "1") != "0"
+_t0 = _time_prof.time()
+def _prof(msg):
+    if _PROF_TRACE:
+        print(f"[PROF] {_time_prof.time()-_t0:7.2f}s {msg}", file=_sys_prof.stderr, flush=True)
+
+
 # ---------------------------------------------------------------------------
 # Pure-PyTorch DeltaNet kernels (fallbacks from transformers 5.2.0)
 # ---------------------------------------------------------------------------
@@ -113,11 +129,34 @@ def _torch_chunk_gated_delta_rule(
     g = g.cumsum(dim=-1)
     decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril()
     attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask_upper, 0)
-    for i in range(1, chunk_size):
-        row = attn[..., i, :i].clone()
-        sub = attn[..., :i, :i].clone()
-        attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2)
-    attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device)
+    # loop1: 三角求逆 (I - attn)^{-1}。默认走 solve_triangular (batched trsm kernel,
+    # 替代 63 步 Python 前向替换循环, 真实服务隔离 4-5x, 数值等价 max_abs<1e-7)。
+    # 传 -attn + unitriangular=True 等效于求解 (I-attn)@X=I → X=(I-attn)^{-1}。
+    # 注意: 不做运行时 attn.abs().max() 阈值检查(会引入 device→host 同步,
+    # 600次调用下额外耗时~1-2s)。真实模型 attn 值域远在安全范围内。
+    # LOOP1_NATIVE=0 可回退串行循环。
+    import os as _os
+    _loop1_done = False
+    if _os.environ.get("LOOP1_NATIVE", "1") != "0":
+        try:
+            _eye = torch.eye(chunk_size, dtype=attn.dtype, device=attn.device)
+            _prof(f"    chunk-gdn loop1 solve_triangular start shape={tuple(attn.shape)}")
+            attn = torch.linalg.solve_triangular(
+                -attn, _eye, upper=False, unitriangular=True)
+            _prof(f"    chunk-gdn loop1 solve_triangular done")
+            _loop1_done = True
+        except Exception as _e:
+            import sys as _sys, traceback as _tb
+            print(f"[loop1_native FALLBACK] {type(_e).__name__}: {_e}",
+                  file=_sys.stderr, flush=True)
+            _tb.print_exc(file=_sys.stderr)
+
+    if not _loop1_done:
+        for i in range(1, chunk_size):
+            row = attn[..., i, :i].clone()
+            sub = attn[..., :i, :i].clone()
+            attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2)
+        attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device)
     value = attn @ v_beta
     k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1))
 
@@ -131,18 +170,56 @@ def _torch_chunk_gated_delta_rule(
         torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device),
         diagonal=1)
 
-    for i in range(total_len // chunk_size):
-        q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i]
-        attn_i = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask_upper2, 0)
-        v_prime = k_cumdecay[:, :, i] @ last_state
-        v_new = v_i - v_prime
-        attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_state
-        core_out[:, :, i] = attn_inter + attn_i @ v_new
-        last_state = (
-            last_state * g[:, :, i, -1, None, None].exp()
-            + (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None])
-            .transpose(-1, -2) @ v_new
-        )
+    # loop2: chunk 间 state 递推。默认走 chunk-parallel(把串行 scan 改写成
+    # 线性矩阵递归 S_i=A_i@S_{i-1}+B_i,重活批量化,串行阶段只对小矩阵 S 做 A@S+B)。
+    # 数值等价于下方串行实现(已验证 max_abs~1e-7, 3x)。CHUNK_PARALLEL=0 可回退。
+    import os as _os
+    NC = total_len // chunk_size
+    _done = False
+    if _os.environ.get("CHUNK_PARALLEL", "1") != "0":
+        try:
+            g_last = g[:, :, :, -1]                                  # [b,h,NC]
+            exp_glast = g_last.exp()
+            attn_i_all = (query @ key.transpose(-1, -2) * decay_mask
+                          ).masked_fill(mask_upper2, 0)              # [b,h,NC,cs,cs]
+            P_i = k_cumdecay                                         # [b,h,NC,cs,k]
+            Ktil = key * (g_last[..., None, None] - g[..., None]).exp()   # [b,h,NC,cs,k]
+            Qg = query * g.exp()[..., None]                         # [b,h,NC,cs,k]
+            eye = torch.eye(k_dim, device=query.device, dtype=query.dtype)
+            A = exp_glast[..., None, None] * eye - Ktil.transpose(-1, -2) @ P_i  # [b,h,NC,k,k]
+            B = Ktil.transpose(-1, -2) @ value                      # [b,h,NC,k,v]
+            M = Qg - attn_i_all @ P_i                               # [b,h,NC,cs,k]
+            N = attn_i_all @ value                                  # [b,h,NC,cs,v]
+            _prof(f"    chunk-gdn loop2 parallel batched done NC={NC}")
+            S = last_state
+            S_starts = torch.empty(batch, num_heads, NC, k_dim, v_dim,
+                                   device=query.device, dtype=query.dtype)
+            for i in range(NC):
+                S_starts[:, :, i] = S
+                S = A[:, :, i] @ S + B[:, :, i]
+            _prof(f"    chunk-gdn loop2 serial scan done ({NC} iters)")
+            last_state = S
+            core_out = M @ S_starts + N                             # [b,h,NC,cs,v]
+            _done = True
+        except Exception as _e:
+            import sys as _sys, traceback as _tb
+            print(f"[chunk_parallel FALLBACK] {type(_e).__name__}: {_e}",
+                  file=_sys.stderr, flush=True)
+            _tb.print_exc(file=_sys.stderr)
+
+    if not _done:
+        for i in range(NC):
+            q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i]
+            attn_i = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask_upper2, 0)
+            v_prime = k_cumdecay[:, :, i] @ last_state
+            v_new = v_i - v_prime
+            attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_state
+            core_out[:, :, i] = attn_inter + attn_i @ v_new
+            last_state = (
+                last_state * g[:, :, i, -1, None, None].exp()
+                + (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None])
+                .transpose(-1, -2) @ v_new
+            )
 
     if not output_final_state:
         last_state = None
@@ -211,6 +288,20 @@ class Qwen3_5RMSNormGated(nn.Module):
 
     def forward(self, hidden_states: torch.Tensor,
                 gate: torch.Tensor) -> torch.Tensor:
+        # 默认用原生 ixformer rms_norm(替代纯 PyTorch pow/mean/rsqrt),
+        # 再乘 silu(gate)。已验证与下方纯 PyTorch 等价(max_abs~1e-2 fp16)。
+        # RMSNORM_NATIVE=0 或异常回退纯 PyTorch。
+        import os as _os3
+        if _os3.environ.get("RMSNORM_NATIVE", "1") != "0":
+            try:
+                from vllm import _custom_ops as _ops
+                out = torch.empty_like(hidden_states)
+                _ops.rms_norm(out, hidden_states.contiguous(),
+                              self.weight, self.variance_epsilon)
+                return (out.to(torch.float32)
+                        * F.silu(gate.to(torch.float32))).to(hidden_states.dtype)
+            except Exception:
+                pass
         input_dtype = hidden_states.dtype
         hs = hidden_states.to(torch.float32)
         variance = hs.pow(2).mean(-1, keepdim=True)
@@ -358,6 +449,7 @@ class GatedDeltaNet(nn.Module):
                 mixed_qkv_conv = F.silu(mixed_qkv_conv)
                 # (1, seq_len, local_conv_dim)
                 mixed_qkv_conv = mixed_qkv_conv.squeeze(0).transpose(0, 1).unsqueeze(0)
+                _prof(f"  L{self.layer_idx} GDN conv1d done seq_len={seq_len}")
 
                 q, k, v = torch.split(
                     mixed_qkv_conv,
@@ -383,6 +475,7 @@ class GatedDeltaNet(nn.Module):
                 _DNN_CHUNK = 4096
                 cur_state = temporal_state[si:si + 1].clone()
                 core_out_parts = []
+                _prof(f"  L{self.layer_idx} GDN chunk-scan start (DNN_CHUNK={_DNN_CHUNK}, nchunks={(seq_len + _DNN_CHUNK - 1)//_DNN_CHUNK})")
                 for sc_start in range(0, seq_len, _DNN_CHUNK):
                     sc_end = min(sc_start + _DNN_CHUNK, seq_len)
                     c_out, cur_state = _torch_chunk_gated_delta_rule(
@@ -396,6 +489,7 @@ class GatedDeltaNet(nn.Module):
                         use_qk_l2norm_in_kernel=True,
                     )
                     core_out_parts.append(c_out)
+                _prof(f"  L{self.layer_idx} GDN chunk-scan done")
                 if cur_state is not None:
                     temporal_state[si].copy_(cur_state[0])
                 # [1, seq_len, num_v_heads, head_v_dim]
@@ -420,9 +514,6 @@ class GatedDeltaNet(nn.Module):
 
         else:
             # Decode: one token per sequence
-            with open("/tmp/vllm_decode_debug.log", "a") as _f:
-                _f.write(f"[deltanet decode] layer={self.layer_idx} num_seqs={hidden_states.shape[0]}\n")
-                _f.flush()
             num_seqs = hidden_states.shape[0]
             weight_2d = self.conv1d_weight.squeeze(1)
 
@@ -452,17 +543,47 @@ class GatedDeltaNet(nn.Module):
             q = q.repeat_interleave(self.head_expand_ratio, dim=2)
             k = k.repeat_interleave(self.head_expand_ratio, dim=2)
 
-            core_out, last_state = _torch_recurrent_gated_delta_rule(
-                q, k, v, g, beta,
-                initial_state=temporal_state,
-                output_final_state=True,
-                use_qk_l2norm_in_kernel=True,
+            # Inlined decode recurrent step (seq_len=1).
+            # Replaces _torch_recurrent_gated_delta_rule to avoid 5 transpose+
+            # contiguous+float32 copies, core_out allocation, and Python loop.
+            # Uses bmm/baddbmm_ to eliminate 3 large (B,H,k,v) intermediate tensors.
+            # temporal_state: (B, H_v, k_dim, v_dim) float32 — updated in-place.
+            orig_dtype = q.dtype
+            _scale = self.head_k_dim ** -0.5
+
+            q_t = _l2norm(q.squeeze(1)).float() * _scale   # (B, H_v, k_dim)
+            k_t = _l2norm(k.squeeze(1)).float()             # (B, H_v, k_dim)
+            v_t = v.squeeze(1).float()                      # (B, H_v, v_dim)
+            g_t = g.squeeze(1).float().exp_()               # (B, H_v)
+            bt  = beta.squeeze(1).float()                   # (B, H_v)
+
+            # Decay state in-place: (B, H_v, k_dim, v_dim) *= scalar per head
+            temporal_state.mul_(g_t[:, :, None, None])
+
+            # Reshape to batched-matmul layout: (B*H_v, k_dim, v_dim)
+            ts_flat = temporal_state.view(-1, self.head_k_dim, self.head_v_dim)
+            BH = ts_flat.shape[0]
+
+            # kv_mem = k_t @ temporal_state  shape: (B*H_v, 1, k_dim) @ (B*H_v, k_dim, v_dim)
+            kv_mem = torch.bmm(
+                k_t.view(BH, 1, self.head_k_dim), ts_flat
+            ).view(num_seqs, local_num_v, self.head_v_dim)  # (B, H_v, v_dim)
+
+            delta = (v_t - kv_mem) * bt[:, :, None]         # (B, H_v, v_dim)
+
+            # State update: temporal_state += outer(k_t, delta)  fused, no intermediate
+            ts_flat.baddbmm_(
+                k_t.view(BH, self.head_k_dim, 1),
+                delta.view(BH, 1, self.head_v_dim),
             )
-            if last_state is not None:
-                temporal_state.copy_(last_state)
+
+            # Output: core_out = q_t @ updated temporal_state
+            core_out = torch.bmm(
+                q_t.view(BH, 1, self.head_k_dim), ts_flat
+            ).view(num_seqs, local_num_v, self.head_v_dim).to(orig_dtype)
+            # core_out: (B, H_v, v_dim) = (num_seqs, local_num_v, head_v_dim) already
 
             z = z_all.reshape(num_seqs, local_num_v, self.head_v_dim)
-            core_out = core_out.reshape(num_seqs, local_num_v, self.head_v_dim)
             normed = self.norm(
                 core_out.reshape(-1, self.head_v_dim),
                 z.reshape(-1, self.head_v_dim))
@@ -618,7 +739,9 @@ class Qwen3_5FullAttention(nn.Module):
         # rope: q=(T, local_num_heads*head_dim), k=(T, 1*head_dim) — mirrors 27B
         q, k = self.rotary_emb(positions, q, k)
 
+        _prof(f"  L{self.layer_idx} ATTN flash start tokens={q.shape[0]}")
         attn_out = self.attn(q, k, v, kv_cache, attn_metadata)
+        _prof(f"  L{self.layer_idx} ATTN flash done")
 
         # Multiply by sigmoid gate before output projection
         attn_out = attn_out * torch.sigmoid(gate.float()).to(attn_out.dtype)
@@ -740,33 +863,46 @@ class Qwen3_5MoeSparseBlock(nn.Module):
         if T == 1:
             # Fast path: single token (decode).
             # Batched GEMM: replace top_k separate F.linear calls with 2 fused ops.
-            # gate_up: 1 large GEMM  (1,H) × (H, K*2*I) → (1, K*2*I)
-            # down:    1 einsum      (K,I) × (K,H,I)    → (K,H)
+            # gate_up: 1 large GEMM  (1,H) × (K*2*I,H)^T → (1, K*2*I)
+            # down:    1 bmm         (K,H,I) @ (K,I,1)    → (K,H)
             # Total: 3 kernel launches vs previous 16 (top_k*2).
             eids    = topk_ids[0]                              # (K,)
             ws      = topk_weights[0].to(hidden_states.dtype)  # (K,)
             w13_sel = w13[eids]                                # (K, 2*I, H)
             w2_sel  = w2[eids]                                 # (K, H, I)
 
-            H   = hidden_states.shape[-1]
-            K2I = w13_sel.shape[1]  # K * (2*I) after reshape
+            H = hidden_states.shape[-1]
 
             gate_up = F.linear(
                 hidden_states,
-                w13_sel.reshape(-1, H).contiguous(),           # (K*2*I, H)
+                w13_sel.reshape(-1, H),                        # (K*2*I, H) — contiguous after indexing
             )                                                  # (1, K*2*I)
             gate_up = gate_up.view(self.top_k, -1)             # (K, 2*I)
             gate, up = gate_up.chunk(2, dim=-1)                # (K, I) each
             act = F.silu(gate) * up                            # (K, I)
 
-            # einsum: result[k,h] = sum_i act[k,i] * w2_sel[k,h,i]
-            expert_out = torch.einsum('ki,khi->kh', act, w2_sel)  # (K, H)
+            # bmm: (K,H,I) @ (K,I,1) → (K,H,1) → (K,H)
+            expert_out = torch.bmm(w2_sel, act.unsqueeze(-1)).squeeze(-1)  # (K, H)
 
             out = (expert_out * ws.unsqueeze(-1)).sum(0, keepdim=True).to(
                 hidden_states.dtype)                           # (1, H)
         else:
-            # General path (prefill / multi-seq): loop over unique active experts.
-            # At most T*top_k unique experts, always <= num_experts.
+            # General path (prefill / multi-seq).
+            # 优化: sort-by-expert + 连续切片 + 原生 kernel(act_bias_mm/silu_and_mul/
+            # ixf linear),消除逐 expert 的 nonzero/index_add 循环调度开销。
+            # 数值等价于下方纯 PyTorch fallback(已验证 byte-for-byte 一致, 端到端1.36x)。
+            # 任意异常回退纯 PyTorch。可用 MOE_NATIVE=0 关闭。
+            import os as _os
+            if _os.environ.get("MOE_NATIVE", "1") != "0":
+                try:
+                    return self._native_experts_sorted(
+                        hidden_states, w13, w2, topk_weights, topk_ids)
+                except Exception as _e:
+                    import sys as _sys, traceback as _tb
+                    print(f"[moe_native FALLBACK] {type(_e).__name__}: {_e}",
+                          file=_sys.stderr, flush=True)
+                    _tb.print_exc(file=_sys.stderr)
+            # ---- 纯 PyTorch fallback (原实现) ----
             out = torch.zeros_like(hidden_states)
             unique_eids = topk_ids.view(-1).unique().tolist()
             for eid in unique_eids:
@@ -783,9 +919,50 @@ class Qwen3_5MoeSparseBlock(nn.Module):
 
         return out  # partial, all-reduce done in forward()
 
+    def _native_experts_sorted(self, hidden_states, w13, w2, topk_weights, topk_ids):
+        """sort-by-expert + 原生 ixformer kernel 的 MoE expert 计算。
+        w13=(E,2I,H), w2=(E,H,I)。返回 partial(pre-all-reduce)。"""
+        import ixformer.functions as _ixf
+        dev = hidden_states.device
+        T = hidden_states.shape[0]
+        top_k = topk_ids.shape[1]
+        E = w13.shape[0]
+
+        flat_e = topk_ids.reshape(-1)                                    # (T*K,)
+        flat_w = topk_weights.reshape(-1)                                # (T*K,)
+        flat_tok = torch.arange(T, device=dev).repeat_interleave(top_k)  # (T*K,)
+
+        order = torch.argsort(flat_e)
+        se = flat_e[order]
+        sw = flat_w[order]
+        stok = flat_tok[order]
+        gathered = hidden_states.index_select(0, stok).contiguous()      # (T*K, H) 已按 expert 连续
+
+        counts = torch.bincount(se, minlength=E)
+        offs = torch.cat([torch.zeros(1, device=dev, dtype=torch.long),
+                          counts.cumsum(0)]).tolist()
+
+        # 就地把每个 expert 段的输出写回 gathered(复用显存,避免再分配 T*K×H,
+        # 否则 profiling 阶段峰值显存↑ → KV cache 容量不足)。
+        active = (counts > 0).nonzero().flatten().tolist()
+        for eid in active:
+            eid = int(eid)
+            s, e = int(offs[eid]), int(offs[eid + 1])
+            seg = gathered[s:e].contiguous()                             # (n, H)
+            gu = _ixf.act_bias_mm(seg, w13[eid], None, scale=1,
+                                  act_type="none", trans_format="TN")    # (n, 2I)
+            act = _ixf.silu_and_mul(gu.contiguous())                     # (n, I)
+            gathered[s:e] = _ixf.linear(act.contiguous(), w2[eid], None) # (n, H) 写回
+
+        gathered = gathered * sw.unsqueeze(-1).to(gathered.dtype)
+        out = torch.zeros_like(hidden_states)
+        out.index_add_(0, stok, gathered.to(out.dtype))
+        return out
+
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         router_logits, _ = self.gate(hidden_states)
         routed_out = self._pure_pytorch_experts(hidden_states, router_logits)
+        _prof(f"  MoE routed-experts done tokens={hidden_states.shape[0]}")
 
         gate_up, _ = self.shared_expert_gate_up(hidden_states)
         shared_out = self.act_fn(gate_up)
@@ -796,7 +973,9 @@ class Qwen3_5MoeSparseBlock(nn.Module):
 
         out = routed_out + shared_out
         if self.experts.tp_size > 1:
+            _prof(f"  MoE all_reduce start tp_size={self.experts.tp_size}")
             out = tensor_model_parallel_all_reduce(out)
+            _prof(f"  MoE all_reduce done")
         return out
 
 
@@ -861,16 +1040,22 @@ class Qwen3_5DecoderLayer(nn.Module):
             hidden_states, residual = self.input_layernorm(hidden_states, residual)
 
         if self.layer_type == "linear_attention":
+            _prof(f"L{self.layer_idx} GDN-start tokens={hidden_states.shape[0]}")
             hidden_states = self.linear_attn(
                 hidden_states, attn_metadata, conv_state, temporal_state)
+            _prof(f"L{self.layer_idx} GDN-done")
         else:
+            _prof(f"L{self.layer_idx} ATTN-start tokens={hidden_states.shape[0]}")
             hidden_states = self.self_attn(
                 positions, hidden_states, kv_cache, attn_metadata)
+            _prof(f"L{self.layer_idx} ATTN-done")
 
         hidden_states, residual = self.post_attention_layernorm(
             hidden_states, residual)
 
+        _prof(f"L{self.layer_idx} MLP-start")
         hidden_states = self.mlp(hidden_states)
+        _prof(f"L{self.layer_idx} MLP-done")
 
         return hidden_states, residual
 
@@ -1007,6 +1192,15 @@ class Qwen3_5ForCausalLM(nn.Module, HasInnerState, SupportsLoRA):
         # Lazy initialised in first forward call
         self.mamba_cache: Optional[MambaCacheManager] = None
 
+        # GDN prefix state cache (align mode): stores (conv_states, temporal_states) snapshots
+        # at KV-block boundaries so that prefix-cache-hit requests can restore correct GDN state.
+        # Key: tuple of physical block IDs covering the cached prefix
+        # Value: (conv_states_cpu, temporal_states_cpu) each of shape (num_gdn_layers, ...)
+        self._gdn_prefix_cache: OrderedDict = OrderedDict()
+        self._gdn_prefix_cache_max: int = 16   # ~16 × 16 MB ≈ 256 MB CPU RAM
+        self._block_size: int = (cache_config.block_size
+                                  if cache_config is not None else 16)
+
     def _get_mamba_cache_shape(self):
         tp_size = get_tensor_model_parallel_world_size()
         # Each sequence's state is stored in float32
@@ -1024,6 +1218,7 @@ class Qwen3_5ForCausalLM(nn.Module, HasInnerState, SupportsLoRA):
         intermediate_tensors: Optional[IntermediateTensors] = None,
         **kwargs,
     ) -> torch.Tensor:
+        _prof(f"ForCausalLM.forward START input_ids={tuple(input_ids.shape)} prefill={getattr(attn_metadata,'num_prefill_tokens',0)}")
         if self.mamba_cache is None:
             if self.scheduler_config is not None:
                 max_batch_size = _get_graph_batch_size(
@@ -1043,9 +1238,70 @@ class Qwen3_5ForCausalLM(nn.Module, HasInnerState, SupportsLoRA):
         # temporal_states: (num_linear_layers, batch, local_num_v, k_dim, v_dim)
         conv_states, temporal_states = mamba_tensors
 
+        # ── GDN prefix-cache align mode: inject saved state on prefix hit ─────
+        # Conditions: prefill pass, batch=1, context_len > 0 (prefix cached or
+        # previous chunk already processed), block_tables available.
+        # We always attempt a lookup: for subsequent chunked-prefill chunks the
+        # key matches our own saved state (same data already in slot → no-op).
+        # For a true cross-request prefix hit the key matches a previous request.
+        _is_single_seq_prefill = (
+            attn_metadata is not None
+            and attn_metadata.num_prefill_tokens > 0
+            and conv_states.shape[1] == 1               # batch == 1
+            and getattr(attn_metadata, 'context_lens_tensor', None) is not None
+            and getattr(attn_metadata, 'block_tables', None) is not None
+            and attn_metadata.block_tables.numel() > 0
+        )
+        if _is_single_seq_prefill:
+            context_len = int(attn_metadata.context_lens_tensor[0].item())
+            if context_len > 0:
+                num_prefix_blocks = context_len // self._block_size
+                if (num_prefix_blocks > 0
+                        and attn_metadata.block_tables.shape[1] >= num_prefix_blocks):
+                    lookup_key = tuple(
+                        attn_metadata.block_tables[0, :num_prefix_blocks]
+                        .cpu().tolist())
+                    if lookup_key in self._gdn_prefix_cache:
+                        saved_conv, saved_temporal = self._gdn_prefix_cache[lookup_key]
+                        conv_states[:, 0].copy_(
+                            saved_conv.to(conv_states.device), non_blocking=True)
+                        temporal_states[:, 0].copy_(
+                            saved_temporal.to(temporal_states.device), non_blocking=True)
+                        self._gdn_prefix_cache.move_to_end(lookup_key)
+                        logger.debug("GDN prefix cache hit: prefix_len=%d blocks=%d",
+                                     context_len, num_prefix_blocks)
+        # ── End inject ──────────────────────────────────────────────────────────
+
         hidden_states = self.model(
             input_ids, positions, kv_caches, attn_metadata,
             conv_states, temporal_states)
+
+        # ── GDN prefix-cache align mode: save state after this prefill chunk ───
+        # Save state keyed by ALL complete KV blocks processed so far.
+        # Next requests reusing this prefix will restore from here.
+        if _is_single_seq_prefill:
+            context_len = int(attn_metadata.context_lens_tensor[0].item())
+            query_len = attn_metadata.num_prefill_tokens
+            total_processed = context_len + query_len
+            num_complete_blocks = total_processed // self._block_size
+            if (num_complete_blocks > 0
+                    and attn_metadata.block_tables.shape[1] >= num_complete_blocks):
+                save_key = tuple(
+                    attn_metadata.block_tables[0, :num_complete_blocks]
+                    .cpu().tolist())
+                # Move to end (LRU: most recent = last) and update value
+                if save_key in self._gdn_prefix_cache:
+                    self._gdn_prefix_cache.move_to_end(save_key)
+                self._gdn_prefix_cache[save_key] = (
+                    conv_states[:, 0].cpu().clone(),
+                    temporal_states[:, 0].cpu().clone(),
+                )
+                # Evict oldest entries beyond max
+                while len(self._gdn_prefix_cache) > self._gdn_prefix_cache_max:
+                    self._gdn_prefix_cache.popitem(last=False)
+        # ── End save ────────────────────────────────────────────────────────────
+
+        _prof(f"ForCausalLM.forward END hidden={tuple(hidden_states.shape)}")
         return hidden_states
 
     def compute_logits(
@@ -1226,6 +1482,35 @@ class Qwen3_5MoeForCausalLM(Qwen3_5ForCausalLM):
                 weight_loader(param, loaded_weight)
                 continue
 
+            # --- Individual expert weights (FT checkpoint: experts.{i}.{proj}.weight) ---
+            # Standard transformers fine-tuning saves each expert separately instead of
+            # the pre-merged (num_experts, ...) tensors in the original checkpoint.
+            if ".mlp.experts." in name:
+                parts = name.split(".mlp.experts.", 1)
+                expert_rest = parts[1]          # e.g. "0.gate_proj.weight"
+                dot_pos = expert_rest.find(".")
+                if dot_pos > 0 and expert_rest[:dot_pos].isdigit():
+                    eid = int(expert_rest[:dot_pos])
+                    proj_raw = expert_rest[dot_pos + 1:]
+                    proj = proj_raw[:-7] if proj_raw.endswith(".weight") else proj_raw
+                    prefix = parts[0]           # e.g. "model.layers.0"
+                    if proj == "gate_proj":
+                        w13_name = f"{prefix}.mlp.experts.w13_weight"
+                        if w13_name in params_dict:
+                            param = params_dict[w13_name]
+                            param.weight_loader(param, loaded_weight, "w1_weight", "w1", eid)
+                    elif proj == "up_proj":
+                        w13_name = f"{prefix}.mlp.experts.w13_weight"
+                        if w13_name in params_dict:
+                            param = params_dict[w13_name]
+                            param.weight_loader(param, loaded_weight, "w3_weight", "w3", eid)
+                    elif proj == "down_proj":
+                        w2_name = f"{prefix}.mlp.experts.w2_weight"
+                        if w2_name in params_dict:
+                            param = params_dict[w2_name]
+                            param.weight_loader(param, loaded_weight, "w2_weight", "w2", eid)
+                    continue
+
             # --- Stacked / standard weights ---
             for param_name, weight_name, shard_id in stacked_params_mapping:
                 if weight_name not in name:

qwen3_6_scripts/sequence.py

@@ -119,6 +119,7 @@ class RequestMetrics:
     scheduler_time: Optional[float] = None
     model_forward_time: Optional[float] = None
     model_execute_time: Optional[float] = None
+    num_cached_tokens: Optional[int] = None
 
 
 class SequenceDataDelta(

qwen3_6_scripts/serving_chat.py

@@ -25,7 +25,7 @@ from vllm.entrypoints.openai.protocol import (
     ChatCompletionResponseChoice, ChatCompletionResponseStreamChoice,
     ChatCompletionStreamResponse, ChatMessage, DeltaFunctionCall, DeltaMessage,
     DeltaToolCall, ErrorResponse, FunctionCall, RequestResponseMetadata,
-    ToolCall, UsageInfo)
+    PromptTokensDetails, ToolCall, UsageInfo)
 from vllm.entrypoints.openai.serving_engine import (BaseModelPath,
                                                     LoRAModulePath,
                                                     OpenAIServing,
@@ -179,6 +179,16 @@ class OpenAIServingChat(OpenAIServing):
             logger.exception("Error in loading multi-modal data")
             return self.create_error_response(str(e))
 
+        # n > max_num_seqs deadlock guard: scheduler uses break (not continue)
+        # when can_schedule(num_new_seqs=n) fails, so an n that exceeds
+        # max_num_seqs permanently blocks the entire waiting queue with no error.
+        _sched_cfg = await self.engine_client.get_scheduler_config()
+        _max_seqs = _sched_cfg.max_num_seqs
+        if request.n is not None and request.n > _max_seqs:
+            return self.create_error_response(
+                f"n={request.n} exceeds max_num_seqs={_max_seqs}. "
+                f"Use n<={_max_seqs} or omit n.")
+
         # validation for OpenAI tools
         # tool_choice = "required" is not supported
         if request.tool_choice == "required":
@@ -284,12 +294,12 @@ class OpenAIServingChat(OpenAIServing):
         if request.stream:
             return self.chat_completion_stream_generator(
                 request, result_generator, request_id, conversation, tokenizer,
-                request_metadata)
+                request_metadata, raw_request=raw_request)
 
         try:
             return await self.chat_completion_full_generator(
                 request, result_generator, request_id, conversation, tokenizer,
-                request_metadata)
+                request_metadata, raw_request=raw_request)
         except ValueError as e:
             # TODO: Use a vllm-specific Validation Error
             return self.create_error_response(str(e))
@@ -307,6 +317,7 @@ class OpenAIServingChat(OpenAIServing):
         conversation: List[ConversationMessage],
         tokenizer: AnyTokenizer,
         request_metadata: RequestResponseMetadata,
+        raw_request: Optional[Request] = None,
     ) -> AsyncGenerator[str, None]:
         model_name = self.base_model_paths[0].name
         created_time = int(time.time())
@@ -318,6 +329,7 @@ class OpenAIServingChat(OpenAIServing):
         previous_num_tokens = [0] * num_choices
         finish_reason_sent = [False] * num_choices
         num_prompt_tokens = 0
+        num_cached_tokens: Optional[int] = None
 
         if isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam):
             tool_choice_function_name = request.tool_choice.function.name
@@ -379,12 +391,37 @@ class OpenAIServingChat(OpenAIServing):
                 yield "data: [DONE]\n\n"
                 return
 
+        # Background task: poll is_disconnected() every 300 ms and abort the
+        # engine request as soon as the client goes away.  This catches the
+        # case where the HTTP layer (Starlette/uvicorn) does not actively read
+        # the receive channel during streaming, so is_disconnected() in
+        # iterate_with_cancellation never fires during fast decode.
+        _disconnect_watcher: Optional[asyncio.Task] = None
+        if raw_request is not None:
+            async def _watch_disconnect() -> None:
+                try:
+                    while True:
+                        if await raw_request.is_disconnected():
+                            logger.info(
+                                "Client disconnected (decode watcher), "
+                                "aborting request %s", request_id)
+                            await self.engine_client.abort(request_id)
+                            return
+                        await asyncio.sleep(0.3)
+                except asyncio.CancelledError:
+                    pass
+            _disconnect_watcher = asyncio.ensure_future(_watch_disconnect())
+
         try:
             async for res in result_generator:
                 if res.prompt_token_ids is not None:
                     num_prompt_tokens = len(res.prompt_token_ids)
                     if res.encoder_prompt_token_ids is not None:
                         num_prompt_tokens += len(res.encoder_prompt_token_ids)
+                if (num_cached_tokens is None
+                        and res.metrics is not None
+                        and res.metrics.num_cached_tokens is not None):
+                    num_cached_tokens = res.metrics.num_cached_tokens
 
                 # We need to do it here, because if there are exceptions in
                 # the result_generator, it needs to be sent as the FIRST
@@ -560,9 +597,14 @@ class OpenAIServingChat(OpenAIServing):
                     # if the message delta is None (e.g. because it was a
                     # "control token" for tool calls or the parser otherwise
                     # wasn't ready to send a token, then
-                    #   get the next token without streaming a chunk
+                    #   get the next token without streaming a chunk.
+                    # However, if this is the finish token we must NOT skip —
+                    # the finish block updates reasoning_token_counts, sets
+                    # finish_reason_sent, and flushes the final usage chunk.
                     if delta_message is None:
-                        continue
+                        if output.finish_reason is None:
+                            continue
+                        delta_message = DeltaMessage()
 
                     if output.finish_reason is None:
                         # Send token-by-token response for each request.n
@@ -686,6 +728,9 @@ class OpenAIServingChat(OpenAIServing):
                     completion_tokens=completion_tokens,
                     total_tokens=num_prompt_tokens + completion_tokens,
                     reasoning_tokens=total_reasoning,
+                    prompt_tokens_details=(
+                        PromptTokensDetails(cached_tokens=num_cached_tokens)
+                        if num_cached_tokens is not None else None),
                 )
 
                 final_usage_chunk = ChatCompletionStreamResponse(
@@ -708,11 +753,27 @@ class OpenAIServingChat(OpenAIServing):
                 total_tokens=num_prompt_tokens + num_completion_tokens,
                 reasoning_tokens=total_reasoning)
 
+        except asyncio.CancelledError:
+            # Client disconnected via CancelledError path; abort engine request.
+            await self.engine_client.abort(request_id)
+            return
         except ValueError as e:
             # TODO: Use a vllm-specific Validation Error
             logger.error("error in chat completion stream generator: %s", e)
             data = self.create_streaming_error_response(str(e))
             yield f"data: {data}\n\n"
+        finally:
+            # Stop the disconnect watcher (it may already be done if it fired).
+            if _disconnect_watcher is not None and not _disconnect_watcher.done():
+                _disconnect_watcher.cancel()
+                try:
+                    await _disconnect_watcher
+                except asyncio.CancelledError:
+                    pass
+            # Covers GeneratorExit when Starlette calls aclose() on disconnect
+            # during decode (tokens arrive fast so CancelledError path is not
+            # always triggered). abort() is a no-op for already-finished requests.
+            await self.engine_client.abort(request_id)
         # Send the final done message after all response.n are finished
         yield "data: [DONE]\n\n"
 
@@ -724,17 +785,47 @@ class OpenAIServingChat(OpenAIServing):
         conversation: List[ConversationMessage],
         tokenizer: AnyTokenizer,
         request_metadata: RequestResponseMetadata,
+        raw_request: Optional[Request] = None,
     ) -> Union[ErrorResponse, ChatCompletionResponse]:
 
         model_name = self.base_model_paths[0].name
         created_time = int(time.time())
         final_res: Optional[RequestOutput] = None
 
+        # Background watcher: same logic as the streaming path — polls
+        # is_disconnected() every 300 ms so that a client disconnect during
+        # non-streaming decode is caught even when uvicorn isn't actively
+        # reading the receive channel.
+        _disconnect_watcher: Optional[asyncio.Task] = None
+        if raw_request is not None:
+            async def _watch_disconnect() -> None:
+                try:
+                    while True:
+                        if await raw_request.is_disconnected():
+                            logger.info(
+                                "Client disconnected (non-stream watcher), "
+                                "aborting request %s", request_id)
+                            await self.engine_client.abort(request_id)
+                            return
+                        await asyncio.sleep(0.3)
+                except asyncio.CancelledError:
+                    pass
+            _disconnect_watcher = asyncio.ensure_future(_watch_disconnect())
+
         try:
             async for res in result_generator:
                 final_res = res
         except asyncio.CancelledError:
+            await self.engine_client.abort(request_id)
             return self.create_error_response("Client disconnected")
+        finally:
+            if _disconnect_watcher is not None and not _disconnect_watcher.done():
+                _disconnect_watcher.cancel()
+                try:
+                    await _disconnect_watcher
+                except asyncio.CancelledError:
+                    pass
+            await self.engine_client.abort(request_id)
 
         assert final_res is not None
 
@@ -876,11 +967,16 @@ class OpenAIServingChat(OpenAIServing):
             total_reasoning_tokens = sum(
                 rp.count_reasoning_tokens(list(output.token_ids))
                 for output in final_res.outputs)
+        num_cached_tokens = (final_res.metrics.num_cached_tokens
+                             if final_res.metrics is not None else None)
         usage = UsageInfo(
             prompt_tokens=num_prompt_tokens,
             completion_tokens=num_generated_tokens,
             total_tokens=num_prompt_tokens + num_generated_tokens,
             reasoning_tokens=total_reasoning_tokens,
+            prompt_tokens_details=(
+                PromptTokensDetails(cached_tokens=num_cached_tokens)
+                if num_cached_tokens is not None else None),
         )
 
         request_metadata.final_usage_info = usage

qwen3_6_scripts/xformers.py

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