icer/qwen36_01
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

align platform build config with working submission

Browse Source
This commit is contained in:
the magician
2026-07-15 00:49:24 +08:00
parent 62f10c3375
commit af375830c7
9 changed files with 1774 additions and 274 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
computility-run.yaml
View File

@@ -8,7 +8,7 @@ command:
- --served-model-name - --served-model-name
- llm - llm
- --max-model-len - --max-model-len
- '100000' - '245000'
- --gpu-memory-utilization - --gpu-memory-utilization
- '0.9' - '0.9'
- --trust-remote-code - --trust-remote-code
@@ -19,16 +19,34 @@ command:
- --disable-log-requests - --disable-log-requests
- --disable-frontend-multiprocessing - --disable-frontend-multiprocessing
- --max-num-batched-tokens - --max-num-batched-tokens
- '8192' - '16384'
- --enable-chunked-prefill - --enable-chunked-prefill
- --max-seq-len-to-capture - --max-seq-len-to-capture
- '32768' - '245000'
- --enable-auto-tool-choice - --enable-auto-tool-choice
- --tool-call-parser - --tool-call-parser
- qwen3_coder - qwen3_coder
- --reasoning-parser - --reasoning-parser
- qwen3 - qwen3
- --enable-prefix-caching - --enable-prefix-caching
- --chat-template
- /workspace/chat_template_multi_system.jinja
env: env:
- name: VLLM_ENGINE_ITERATION_TIMEOUT_S - name: VLLM_ENGINE_ITERATION_TIMEOUT_S
value: 3600 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'

175
qwen3_6_scripts/chat_template_multi_system.jinja Normal file
View File

@@ -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 %}

118
qwen3_6_scripts/paged_attn.py
View File

@@ -289,7 +289,22 @@ class PagedAttention:
) -> torch.Tensor: ) -> torch.Tensor:
# NOTE: The Triton context_attention_fwd kernel hangs on Iluvatar # NOTE: The Triton context_attention_fwd kernel hangs on Iluvatar
# BI-V100 hardware (same class of issue as cudnnFlashAttnForward). # 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( return PagedAttention._forward_prefix_pytorch(
query, key, value, query, key, value,
key_cache, value_cache, key_cache, value_cache,
@@ -297,6 +312,107 @@ class PagedAttention:
seq_lens_tensor, context_lens, 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 @staticmethod
def _forward_prefix_pytorch( def _forward_prefix_pytorch(
query: torch.Tensor, query: torch.Tensor,

122
qwen3_6_scripts/patch_ops.sh
View File

@@ -24,6 +24,45 @@
# --max-num-batched-tokens 8192 --enable-chunked-prefill --enable-prefix-caching \ # --max-num-batched-tokens 8192 --enable-chunked-prefill --enable-prefix-caching \
# --max-seq-len-to-capture 32768 # --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 ------- # --- paged_attn.py: replace forward_prefix with pure-PyTorch fallback -------
# The Triton context_attention_fwd kernel hangs BI-V100 GPUs permanently # The Triton context_attention_fwd kernel hangs BI-V100 GPUs permanently
# (standard Triton 2.3.1 PTX is not supported by the corex runtime either). # (standard Triton 2.3.1 PTX is not supported by the corex runtime either).
@@ -92,3 +131,86 @@ cp ./cli_args.py /usr/local/corex/lib/python3/dist-packages/vllm/entrypoints/ope
cp ./serving_chat.py /usr/local/corex/lib/python3/dist-packages/vllm/entrypoints/openai/serving_chat.py cp ./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 ./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 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/)"

69
qwen3_6_scripts/patch_xformers_sdpa_seq.py
View File

@@ -29,10 +29,6 @@ flash attention kernel（ixformer / cudnnFlashAttnForward）。
关闭选项），原意是防止 profiling OOM。但 _run_sdpa_fallback 已通过 Q-tiling 关闭选项），原意是防止 profiling OOM。但 _run_sdpa_fallback 已通过 Q-tiling
解决了该问题，chunked prefill 反而会把推理路径从 _run_sdpa_fallback 切换到 解决了该问题，chunked prefill 反而会把推理路径从 _run_sdpa_fallback 切换到
_forward_prefix_pytorch，属于不必要的行为变更，因此一并禁用该自动逻辑。 _forward_prefix_pytorch，属于不必要的行为变更，因此一并禁用该自动逻辑。
同时解除 EngineArgs.create_*_config 中强制 enforce_eager=True 和
disable_custom_all_reduce=True 的硬编码，让 Iluvatar 环境可以实际验证
CUDA Graph / custom all-reduce；需要回退时仍可通过命令行显式传
--enforce-eager --disable-custom-all-reduce。
Deploy: Deploy:
python3 modified_scripts/patch_xformers_sdpa_seq.py python3 modified_scripts/patch_xformers_sdpa_seq.py
@@ -95,13 +91,6 @@ _ARG_NEW_BLOCK = """\
# handles long-context memory without chunked prefill\ # handles long-context memory without chunked prefill\
""" """
_ARG_FORCE_EAGER_OLD = " enforce_eager=True,"
_ARG_FORCE_EAGER_NEW = " enforce_eager=self.enforce_eager,"
_ARG_FORCE_ALLREDUCE_OLD = " disable_custom_all_reduce=True,"
_ARG_FORCE_ALLREDUCE_NEW = (
" disable_custom_all_reduce=self.disable_custom_all_reduce,"
)
FALLBACK_METHOD = ''' FALLBACK_METHOD = '''
def _run_sdpa_fallback( def _run_sdpa_fallback(
self, self,
@@ -225,18 +214,34 @@ NEW_XFORMER_BLOCK = """\
query = query.unsqueeze(0) query = query.unsqueeze(0)
key = key.unsqueeze(0) key = key.unsqueeze(0)
value = value.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) 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)\ return out.view_as(original_query)\
""" """
@@ -286,26 +291,6 @@ def patch_arg_utils(path):
else: else:
print(" [warn] target block not found — check arg_utils.py version") print(" [warn] target block not found — check arg_utils.py version")
if _ARG_FORCE_EAGER_NEW in content:
print(" [skip] enforce_eager already respects CLI")
elif _ARG_FORCE_EAGER_OLD in content:
content = content.replace(_ARG_FORCE_EAGER_OLD,
_ARG_FORCE_EAGER_NEW, 1)
print(" [ok] enforce_eager now respects CLI")
changed = True
else:
print(" [warn] enforce_eager assignment not found")
if _ARG_FORCE_ALLREDUCE_NEW in content:
print(" [skip] custom all-reduce already respects CLI")
elif _ARG_FORCE_ALLREDUCE_OLD in content:
content = content.replace(_ARG_FORCE_ALLREDUCE_OLD,
_ARG_FORCE_ALLREDUCE_NEW, 1)
print(" [ok] custom all-reduce now respects CLI")
changed = True
else:
print(" [warn] custom all-reduce assignment not found")
if changed: if changed:
with open(path, "w") as f: with open(path, "w") as f:
f.write(content) f.write(content)

72
qwen3_6_scripts/protocol.py
View File

@@ -1,5 +1,6 @@
# Adapted from # Adapted from
# https://github.com/lm-sys/FastChat/blob/168ccc29d3f7edc50823016105c024fe2282732a/fastchat/protocol/openai_api_protocol.py # https://github.com/lm-sys/FastChat/blob/168ccc29d3f7edc50823016105c024fe2282732a/fastchat/protocol/openai_api_protocol.py
import json
import time import time
from argparse import Namespace from argparse import Namespace
from typing import Any, Dict, List, Literal, Optional, Union from typing import Any, Dict, List, Literal, Optional, Union
@@ -254,6 +255,15 @@ class ChatCompletionRequest(OpenAIBaseModel):
description=("Additional kwargs to pass to the template renderer. " description=("Additional kwargs to pass to the template renderer. "
"Will be accessible by the chat template."), "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( guided_json: Optional[Union[str, dict, BaseModel]] = Field(
default=None, default=None,
description=("If specified, the output will follow the JSON schema."), description=("If specified, the output will follow the JSON schema."),
@@ -411,6 +421,26 @@ class ChatCompletionRequest(OpenAIBaseModel):
"Each message must have at least one of 'content' or " "Each message must have at least one of 'content' or "
"'reasoning_content'.") "'reasoning_content'.")
msg = {**msg, "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) normalized.append(msg)
data = {**data, "messages": normalized} data = {**data, "messages": normalized}
return data return data
@@ -424,6 +454,48 @@ class ChatCompletionRequest(OpenAIBaseModel):
return data 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") @model_validator(mode="before")
@classmethod @classmethod
def check_logprobs(cls, data): def check_logprobs(cls, data):

488
qwen3_6_scripts/qwen3_5.py
View File

@@ -3,9 +3,6 @@
# Text-only (no VL, no MTP). # Text-only (no VL, no MTP).
from collections import OrderedDict from collections import OrderedDict
from contextlib import contextmanager
import os
import time
from typing import Dict, Iterable, List, Optional, Tuple from typing import Dict, Iterable, List, Optional, Tuple
import torch import torch
@@ -44,60 +41,20 @@ from vllm.model_executor.models.interfaces import HasInnerState, SupportsLoRA
logger = init_logger(__name__) logger = init_logger(__name__)
_ENGINEX_PROFILE_ENABLED = os.getenv("ENGINEX_PROFILE_DECODE", "0") == "1"
_ENGINEX_PROFILE_EVERY = int(os.getenv("ENGINEX_PROFILE_EVERY", "32"))
_ENGINEX_PROFILE_SYNC = os.getenv("ENGINEX_PROFILE_SYNC", "1") != "0"
_ENGINEX_MOE_TINY_IMPL = os.getenv("ENGINEX_MOE_TINY_IMPL", "tokenwise")
_ENGINEX_MOE_TINY_MAX = int(os.getenv("ENGINEX_MOE_TINY_MAX", "4"))
_enginex_profile_stats: Dict[str, List[float]] = {}
_enginex_profile_steps = 0
_enginex_profile_mode = "unknown"
# ---------------------------------------------------------------------------
def _enginex_profile_active() -> bool: # Profiling debug prints (locate where startup hangs after weight load).
return _ENGINEX_PROFILE_ENABLED and torch.cuda.is_available() # Gate on PROF_TRACE=1 (default on, since we are diagnosing a startup hang).
# Each print flushes; grep "[PROF]" in the server log.
# ---------------------------------------------------------------------------
def _enginex_profile_sync() -> None: import os as _os_prof
if _ENGINEX_PROFILE_SYNC: import sys as _sys_prof
torch.cuda.synchronize() import time as _time_prof
_PROF_TRACE = _os_prof.environ.get("PROF_TRACE", "1") != "0"
_t0 = _time_prof.time()
@contextmanager def _prof(msg):
def _enginex_profile(label: str): if _PROF_TRACE:
if not _enginex_profile_active(): print(f"[PROF] {_time_prof.time()-_t0:7.2f}s {msg}", file=_sys_prof.stderr, flush=True)
yield
return
label = f"{_enginex_profile_mode}.{label}"
_enginex_profile_sync()
start = time.perf_counter()
try:
yield
finally:
_enginex_profile_sync()
elapsed_ms = (time.perf_counter() - start) * 1000.0
stat = _enginex_profile_stats.setdefault(label, [0.0, 0.0])
stat[0] += elapsed_ms
stat[1] += 1.0
def _enginex_profile_log(mode: str) -> None:
global _enginex_profile_steps
if not _enginex_profile_active():
return
_enginex_profile_steps += 1
if _enginex_profile_steps % max(_ENGINEX_PROFILE_EVERY, 1) != 0:
return
tp_rank = get_tensor_model_parallel_rank()
parts = []
for label, (total_ms, count) in sorted(
_enginex_profile_stats.items(),
key=lambda item: item[1][0],
reverse=True):
avg_ms = total_ms / max(count, 1.0)
parts.append(f"{label}: total={total_ms:.2f}ms avg={avg_ms:.3f}ms n={int(count)}")
logger.info("[ENGINEX_PROFILE_QWEN] rank=%d steps=%d mode=%s %s",
tp_rank, _enginex_profile_steps, mode, " | ".join(parts))
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -172,11 +129,34 @@ def _torch_chunk_gated_delta_rule(
g = g.cumsum(dim=-1) g = g.cumsum(dim=-1)
decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril() 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) attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask_upper, 0)
for i in range(1, chunk_size): # loop1: 三角求逆 (I - attn)^{-1}。默认走 solve_triangular (batched trsm kernel,
row = attn[..., i, :i].clone() # 替代 63 步 Python 前向替换循环, 真实服务隔离 4-5x, 数值等价 max_abs<1e-7)。
sub = attn[..., :i, :i].clone() # 传 -attn + unitriangular=True 等效于求解 (I-attn)@X=I → X=(I-attn)^{-1}。
attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2) # 注意: 不做运行时 attn.abs().max() 阈值检查(会引入 device→host 同步,
attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device) # 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 value = attn @ v_beta
k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1)) k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1))
@@ -190,18 +170,56 @@ def _torch_chunk_gated_delta_rule(
torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device),
diagonal=1) diagonal=1)
for i in range(total_len // chunk_size): # loop2: chunk 间 state 递推。默认走 chunk-parallel(把串行 scan 改写成
q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i] # 线性矩阵递归 S_i=A_i@S_{i-1}+B_i,重活批量化,串行阶段只对小矩阵 S 做 A@S+B)。
attn_i = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask_upper2, 0) # 数值等价于下方串行实现(已验证 max_abs~1e-7, 3x)。CHUNK_PARALLEL=0 可回退。
v_prime = k_cumdecay[:, :, i] @ last_state import os as _os
v_new = v_i - v_prime NC = total_len // chunk_size
attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_state _done = False
core_out[:, :, i] = attn_inter + attn_i @ v_new if _os.environ.get("CHUNK_PARALLEL", "1") != "0":
last_state = ( try:
last_state * g[:, :, i, -1, None, None].exp() g_last = g[:, :, :, -1] # [b,h,NC]
+ (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]) exp_glast = g_last.exp()
.transpose(-1, -2) @ v_new 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: if not output_final_state:
last_state = None last_state = None
@@ -270,6 +288,20 @@ class Qwen3_5RMSNormGated(nn.Module):
def forward(self, hidden_states: torch.Tensor, def forward(self, hidden_states: torch.Tensor,
gate: torch.Tensor) -> 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 input_dtype = hidden_states.dtype
hs = hidden_states.to(torch.float32) hs = hidden_states.to(torch.float32)
variance = hs.pow(2).mean(-1, keepdim=True) variance = hs.pow(2).mean(-1, keepdim=True)
@@ -417,6 +449,7 @@ class GatedDeltaNet(nn.Module):
mixed_qkv_conv = F.silu(mixed_qkv_conv) mixed_qkv_conv = F.silu(mixed_qkv_conv)
# (1, seq_len, local_conv_dim) # (1, seq_len, local_conv_dim)
mixed_qkv_conv = mixed_qkv_conv.squeeze(0).transpose(0, 1).unsqueeze(0) 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( q, k, v = torch.split(
mixed_qkv_conv, mixed_qkv_conv,
@@ -442,6 +475,7 @@ class GatedDeltaNet(nn.Module):
_DNN_CHUNK = 4096 _DNN_CHUNK = 4096
cur_state = temporal_state[si:si + 1].clone() cur_state = temporal_state[si:si + 1].clone()
core_out_parts = [] 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): for sc_start in range(0, seq_len, _DNN_CHUNK):
sc_end = min(sc_start + _DNN_CHUNK, seq_len) sc_end = min(sc_start + _DNN_CHUNK, seq_len)
c_out, cur_state = _torch_chunk_gated_delta_rule( c_out, cur_state = _torch_chunk_gated_delta_rule(
@@ -455,6 +489,7 @@ class GatedDeltaNet(nn.Module):
use_qk_l2norm_in_kernel=True, use_qk_l2norm_in_kernel=True,
) )
core_out_parts.append(c_out) core_out_parts.append(c_out)
_prof(f" L{self.layer_idx} GDN chunk-scan done")
if cur_state is not None: if cur_state is not None:
temporal_state[si].copy_(cur_state[0]) temporal_state[si].copy_(cur_state[0])
# [1, seq_len, num_v_heads, head_v_dim] # [1, seq_len, num_v_heads, head_v_dim]
@@ -671,48 +706,46 @@ class Qwen3_5FullAttention(nn.Module):
total_tokens = hidden_states.shape[0] total_tokens = hidden_states.shape[0]
# q_proj output includes gate (dim doubled) # q_proj output includes gate (dim doubled)
with _enginex_profile("full_attn.qkv_proj"): qg, _ = self.q_proj(hidden_states) # (total, local_num_heads * head_dim * 2)
qg, _ = self.q_proj(hidden_states) # (total, local_num_heads * head_dim * 2) qg = qg.view(total_tokens, self.local_num_heads, self.head_dim * 2)
qg = qg.view(total_tokens, self.local_num_heads, self.head_dim * 2) q = qg[:, :, :self.head_dim].reshape(total_tokens, -1)
q = qg[:, :, :self.head_dim].reshape(total_tokens, -1) gate = qg[:, :, self.head_dim:].reshape(total_tokens, -1)
gate = qg[:, :, self.head_dim:].reshape(total_tokens, -1)
k, _ = self.k_proj(hidden_states) # (total, proj_kv_heads * head_dim) k, _ = self.k_proj(hidden_states) # (total, proj_kv_heads * head_dim)
v, _ = self.v_proj(hidden_states) v, _ = self.v_proj(hidden_states)
# q_norm on local Q heads # q_norm on local Q heads
with _enginex_profile("full_attn.norm_rope"): q = self.q_norm.forward_cuda(
q = self.q_norm.forward_cuda( q.view(total_tokens, self.local_num_heads, self.head_dim)
q.view(total_tokens, self.local_num_heads, self.head_dim) .contiguous()).view(total_tokens, -1)
.contiguous()).view(total_tokens, -1)
# GQA-aware TP: select rank-local KV head BEFORE k_norm and rope so # GQA-aware TP: select rank-local KV head BEFORE k_norm and rope so
# that ixformer kernels always see num_kv_heads=1 (same as 27B path). # that ixformer kernels always see num_kv_heads=1 (same as 27B path).
# Doing k_norm/rope on 2 KV heads (proj_kv_heads=2) triggers ixformer # Doing k_norm/rope on 2 KV heads (proj_kv_heads=2) triggers ixformer
# paths that can produce NaN; restricting to 1 head avoids the issue. # paths that can produce NaN; restricting to 1 head avoids the issue.
if self.q_per_kv_global is not None: if self.q_per_kv_global is not None:
tp_rank = get_tensor_model_parallel_rank() tp_rank = get_tensor_model_parallel_rank()
kv_idx = (tp_rank * self.local_num_heads) // self.q_per_kv_global kv_idx = (tp_rank * self.local_num_heads) // self.q_per_kv_global
k = (k.view(total_tokens, self.proj_kv_heads, self.head_dim) k = (k.view(total_tokens, self.proj_kv_heads, self.head_dim)
[:, kv_idx, :].contiguous()) # (T, head_dim) — 1 head [:, kv_idx, :].contiguous()) # (T, head_dim) — 1 head
v = (v.view(total_tokens, self.proj_kv_heads, self.head_dim) v = (v.view(total_tokens, self.proj_kv_heads, self.head_dim)
[:, kv_idx, :].contiguous()) # (T, head_dim) — 1 head [:, kv_idx, :].contiguous()) # (T, head_dim) — 1 head
# k_norm on the (now always 1) rank-local KV head # k_norm on the (now always 1) rank-local KV head
k = self.k_norm.forward_cuda( k = self.k_norm.forward_cuda(
k.view(total_tokens, self.local_num_kv_heads, self.head_dim) k.view(total_tokens, self.local_num_kv_heads, self.head_dim)
.contiguous()).view(total_tokens, -1) .contiguous()).view(total_tokens, -1)
# rope: q=(T, local_num_heads*head_dim), k=(T, 1*head_dim) — mirrors 27B # rope: q=(T, local_num_heads*head_dim), k=(T, 1*head_dim) — mirrors 27B
q, k = self.rotary_emb(positions, q, k) q, k = self.rotary_emb(positions, q, k)
with _enginex_profile("full_attn.paged_attention"): _prof(f" L{self.layer_idx} ATTN flash start tokens={q.shape[0]}")
attn_out = self.attn(q, k, v, kv_cache, attn_metadata) 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 # Multiply by sigmoid gate before output projection
with _enginex_profile("full_attn.gate_o_proj"): attn_out = attn_out * torch.sigmoid(gate.float()).to(attn_out.dtype)
attn_out = attn_out * torch.sigmoid(gate.float()).to(attn_out.dtype) output, _ = self.o_proj(attn_out)
output, _ = self.o_proj(attn_out)
return output return output
@@ -816,13 +849,12 @@ class Qwen3_5MoeSparseBlock(nn.Module):
Output is partial (pre-all-reduce), same contract as FusedMoE Output is partial (pre-all-reduce), same contract as FusedMoE
with reduce_results=False. with reduce_results=False.
""" """
# Routing: softmax -> topk -> renormalise # Routing: softmax → topk → renormalise
with _enginex_profile("moe.routing_topk"): routing_weights = torch.softmax(router_logits.float(), dim=-1)
routing_weights = torch.softmax(router_logits.float(), dim=-1) topk_weights, topk_ids = torch.topk(
topk_weights, topk_ids = torch.topk( routing_weights, self.top_k, dim=-1) # (T, top_k)
routing_weights, self.top_k, dim=-1) # (T, top_k) topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) topk_weights = topk_weights.to(hidden_states.dtype)
topk_weights = topk_weights.to(hidden_states.dtype)
w13 = self.experts.w13_weight # (E, 2*I, H) w13 = self.experts.w13_weight # (E, 2*I, H)
w2 = self.experts.w2_weight # (E, H, I) w2 = self.experts.w2_weight # (E, H, I)
@@ -834,114 +866,116 @@ class Qwen3_5MoeSparseBlock(nn.Module):
# gate_up: 1 large GEMM (1,H) × (K*2*I,H)^T → (1, K*2*I) # 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) # down: 1 bmm (K,H,I) @ (K,I,1) → (K,H)
# Total: 3 kernel launches vs previous 16 (top_k*2). # Total: 3 kernel launches vs previous 16 (top_k*2).
with _enginex_profile("moe.routed_decode_experts"): eids = topk_ids[0] # (K,)
eids = topk_ids[0] # (K,) ws = topk_weights[0].to(hidden_states.dtype) # (K,)
ws = topk_weights[0].to(hidden_states.dtype) # (K,) w13_sel = w13[eids] # (K, 2*I, H)
w13_sel = w13[eids] # (K, 2*I, H) w2_sel = w2[eids] # (K, H, I)
w2_sel = w2[eids] # (K, H, I)
H = hidden_states.shape[-1] H = hidden_states.shape[-1]
gate_up = F.linear( gate_up = F.linear(
hidden_states, hidden_states,
w13_sel.reshape(-1, H), # (K*2*I, H) - contiguous after indexing w13_sel.reshape(-1, H), # (K*2*I, H) — contiguous after indexing
) # (1, K*2*I) ) # (1, K*2*I)
gate_up = gate_up.view(self.top_k, -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 gate, up = gate_up.chunk(2, dim=-1) # (K, I) each
act = F.silu(gate) * up # (K, I) act = F.silu(gate) * up # (K, I)
# bmm: (K,H,I) @ (K,I,1) -> (K,H,1) -> (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) expert_out = torch.bmm(w2_sel, act.unsqueeze(-1)).squeeze(-1) # (K, H)
out = (expert_out * ws.unsqueeze(-1)).sum(0, keepdim=True).to( out = (expert_out * ws.unsqueeze(-1)).sum(0, keepdim=True).to(
hidden_states.dtype) # (1, H) hidden_states.dtype) # (1, H)
elif T <= _ENGINEX_MOE_TINY_MAX and _ENGINEX_MOE_TINY_IMPL == "tokenwise":
# Fast path: tiny decode batch. Compute each token with the proven
# T==1 large-F.linear path, avoiding the per-expert Python loop.
# This usually beats batched bmm for very small T because the first
# projection becomes T larger GEMMs instead of T*K tiny GEMMs.
with _enginex_profile("moe.routed_tiny_tokenwise_experts"):
H = hidden_states.shape[-1]
pieces = []
for t in range(T):
eids = topk_ids[t]
ws = topk_weights[t].to(hidden_states.dtype)
w13_sel = w13[eids] # (K, 2*I, H)
w2_sel = w2[eids] # (K, H, I)
gate_up = F.linear(
hidden_states[t:t + 1],
w13_sel.reshape(-1, H),
).view(self.top_k, -1) # (K, 2*I)
gate, up = gate_up.chunk(2, dim=-1)
act = F.silu(gate) * up
expert_out = torch.bmm(
w2_sel, act.unsqueeze(-1)).squeeze(-1) # (K, H)
pieces.append((expert_out * ws.unsqueeze(-1)).sum(0))
out = torch.stack(pieces, dim=0).to(hidden_states.dtype)
elif T <= _ENGINEX_MOE_TINY_MAX:
# Alternative tiny decode batch implementation. Kept for A/B
# testing via ENGINEX_MOE_TINY_IMPL=bmm.
with _enginex_profile("moe.routed_tiny_bmm_experts"):
H = hidden_states.shape[-1]
flat_eids = topk_ids.reshape(-1) # (T*K,)
flat_ws = topk_weights.reshape(-1).to(hidden_states.dtype)
w13_sel = w13[flat_eids] # (T*K, 2*I, H)
w2_sel = w2[flat_eids] # (T*K, H, I)
x = (hidden_states[:, None, :]
.expand(T, self.top_k, H)
.reshape(-1, 1, H)) # (T*K, 1, H)
gate_up = torch.bmm(
x, w13_sel.transpose(1, 2)).squeeze(1) # (T*K, 2*I)
gate, up = gate_up.chunk(2, dim=-1)
act = F.silu(gate) * up # (T*K, I)
expert_out = torch.bmm(
w2_sel, act.unsqueeze(-1)).squeeze(-1) # (T*K, H)
out = (expert_out * flat_ws.unsqueeze(-1)).view(
T, self.top_k, H).sum(1).to(hidden_states.dtype)
else: else:
# General path (prefill / multi-seq): loop over unique active experts. # General path (prefill / multi-seq).
# At most T*top_k unique experts, always <= num_experts. # 优化: sort-by-expert + 连续切片 + 原生 kernel(act_bias_mm/silu_and_mul/
with _enginex_profile("moe.routed_prefill_experts"): # ixf linear),消除逐 expert 的 nonzero/index_add 循环调度开销。
out = torch.zeros_like(hidden_states) # 数值等价于下方纯 PyTorch fallback(已验证 byte-for-byte 一致, 端到端1.36x)。
unique_eids = topk_ids.view(-1).unique().tolist() # 任意异常回退纯 PyTorch。可用 MOE_NATIVE=0 关闭。
for eid in unique_eids: import os as _os
eid = int(eid) if _os.environ.get("MOE_NATIVE", "1") != "0":
mask = (topk_ids == eid) # (T, top_k) try:
tok_ids, topk_pos = mask.nonzero(as_tuple=True) return self._native_experts_sorted(
tokens = hidden_states[tok_ids] # (n, H) hidden_states, w13, w2, topk_weights, topk_ids)
gate_up = F.linear(tokens, w13[eid]) # (n, 2*I) except Exception as _e:
gate, up = gate_up.chunk(2, dim=-1) import sys as _sys, traceback as _tb
act = F.silu(gate) * up # (n, I) print(f"[moe_native FALLBACK] {type(_e).__name__}: {_e}",
expert_out = F.linear(act, w2[eid]) # (n, H) file=_sys.stderr, flush=True)
weights = topk_weights[tok_ids, topk_pos].unsqueeze(-1) _tb.print_exc(file=_sys.stderr)
out.index_add_(0, tok_ids, (expert_out * weights).to(out.dtype)) # ---- 纯 PyTorch fallback (原实现) ----
out = torch.zeros_like(hidden_states)
unique_eids = topk_ids.view(-1).unique().tolist()
for eid in unique_eids:
eid = int(eid)
mask = (topk_ids == eid) # (T, top_k)
tok_ids, topk_pos = mask.nonzero(as_tuple=True)
tokens = hidden_states[tok_ids] # (n, H)
gate_up = F.linear(tokens, w13[eid]) # (n, 2*I)
gate, up = gate_up.chunk(2, dim=-1)
act = F.silu(gate) * up # (n, I)
expert_out = F.linear(act, w2[eid]) # (n, H)
weights = topk_weights[tok_ids, topk_pos].unsqueeze(-1)
out.index_add_(0, tok_ids, (expert_out * weights).to(out.dtype))
return out # partial, all-reduce done in forward() 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: def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
with _enginex_profile("moe.gate"): router_logits, _ = self.gate(hidden_states)
router_logits, _ = self.gate(hidden_states) routed_out = self._pure_pytorch_experts(hidden_states, router_logits)
with _enginex_profile("moe.routed_total"): _prof(f" MoE routed-experts done tokens={hidden_states.shape[0]}")
routed_out = self._pure_pytorch_experts(hidden_states, router_logits)
with _enginex_profile("moe.shared_expert"): gate_up, _ = self.shared_expert_gate_up(hidden_states)
gate_up, _ = self.shared_expert_gate_up(hidden_states) shared_out = self.act_fn(gate_up)
shared_out = self.act_fn(gate_up) shared_out, _ = self.shared_expert_down(shared_out)
shared_out, _ = self.shared_expert_down(shared_out) # Scalar sigmoid gate (Qwen2-MoE / Qwen3.5-MoE style)
# Scalar sigmoid gate (Qwen2-MoE / Qwen3.5-MoE style) gate_score, _ = self.shared_expert_gate(hidden_states) # (T, 1)
gate_score, _ = self.shared_expert_gate(hidden_states) # (T, 1) shared_out = shared_out * torch.sigmoid(gate_score)
shared_out = shared_out * torch.sigmoid(gate_score)
with _enginex_profile("moe.combine"): out = routed_out + shared_out
out = routed_out + shared_out
if self.experts.tp_size > 1: if self.experts.tp_size > 1:
with _enginex_profile("moe.tp_all_reduce"): _prof(f" MoE all_reduce start tp_size={self.experts.tp_size}")
out = tensor_model_parallel_all_reduce(out) out = tensor_model_parallel_all_reduce(out)
_prof(f" MoE all_reduce done")
return out return out
@@ -1001,27 +1035,27 @@ class Qwen3_5DecoderLayer(nn.Module):
) -> Tuple[torch.Tensor, torch.Tensor]: ) -> Tuple[torch.Tensor, torch.Tensor]:
if residual is None: if residual is None:
residual = hidden_states residual = hidden_states
with _enginex_profile("layer.input_norm"): hidden_states = self.input_layernorm(hidden_states)
hidden_states = self.input_layernorm(hidden_states)
else: else:
with _enginex_profile("layer.input_norm"): hidden_states, residual = self.input_layernorm(hidden_states, residual)
hidden_states, residual = self.input_layernorm(hidden_states, residual)
if self.layer_type == "linear_attention": if self.layer_type == "linear_attention":
with _enginex_profile("layer.linear_attention"): _prof(f"L{self.layer_idx} GDN-start tokens={hidden_states.shape[0]}")
hidden_states = self.linear_attn( hidden_states = self.linear_attn(
hidden_states, attn_metadata, conv_state, temporal_state) hidden_states, attn_metadata, conv_state, temporal_state)
_prof(f"L{self.layer_idx} GDN-done")
else: else:
with _enginex_profile("layer.full_attention"): _prof(f"L{self.layer_idx} ATTN-start tokens={hidden_states.shape[0]}")
hidden_states = self.self_attn( hidden_states = self.self_attn(
positions, hidden_states, kv_cache, attn_metadata) positions, hidden_states, kv_cache, attn_metadata)
_prof(f"L{self.layer_idx} ATTN-done")
with _enginex_profile("layer.post_attn_norm"): hidden_states, residual = self.post_attention_layernorm(
hidden_states, residual = self.post_attention_layernorm( hidden_states, residual)
hidden_states, residual)
with _enginex_profile("layer.mlp"): _prof(f"L{self.layer_idx} MLP-start")
hidden_states = self.mlp(hidden_states) hidden_states = self.mlp(hidden_states)
_prof(f"L{self.layer_idx} MLP-done")
return hidden_states, residual return hidden_states, residual
@@ -1058,9 +1092,6 @@ class Qwen3_5Model(nn.Module):
conv_states: torch.Tensor, # (num_linear_layers, batch, ...) conv_states: torch.Tensor, # (num_linear_layers, batch, ...)
temporal_states: torch.Tensor, # (num_linear_layers, batch, ...) temporal_states: torch.Tensor, # (num_linear_layers, batch, ...)
) -> torch.Tensor: ) -> torch.Tensor:
global _enginex_profile_mode
mode = "prefill" if attn_metadata.num_prefill_tokens > 0 else "decode"
_enginex_profile_mode = mode
hidden_states = self.embed_tokens(input_ids) hidden_states = self.embed_tokens(input_ids)
residual = None residual = None
@@ -1088,7 +1119,6 @@ class Qwen3_5Model(nn.Module):
attn_idx += 1 attn_idx += 1
hidden_states, _ = self.norm(hidden_states, residual) hidden_states, _ = self.norm(hidden_states, residual)
_enginex_profile_log(mode)
return hidden_states return hidden_states
@@ -1188,6 +1218,7 @@ class Qwen3_5ForCausalLM(nn.Module, HasInnerState, SupportsLoRA):
intermediate_tensors: Optional[IntermediateTensors] = None, intermediate_tensors: Optional[IntermediateTensors] = None,
**kwargs, **kwargs,
) -> torch.Tensor: ) -> 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.mamba_cache is None:
if self.scheduler_config is not None: if self.scheduler_config is not None:
max_batch_size = _get_graph_batch_size( max_batch_size = _get_graph_batch_size(
@@ -1270,6 +1301,7 @@ class Qwen3_5ForCausalLM(nn.Module, HasInnerState, SupportsLoRA):
self._gdn_prefix_cache.popitem(last=False) self._gdn_prefix_cache.popitem(last=False)
# ── End save ──────────────────────────────────────────────────────────── # ── End save ────────────────────────────────────────────────────────────
_prof(f"ForCausalLM.forward END hidden={tuple(hidden_states.shape)}")
return hidden_states return hidden_states
def compute_logits( def compute_logits(

946
qwen3_6_scripts/xformers.py Normal file
View File

@@ -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

34
worklogs/decode_analysis_report_2026-07-14.md
View File

@@ -731,3 +731,37 @@ ENGINEX_MOE_TINY_MAX=4
1. 测试 `max_num_seqs=4`，利用 `T <= 4` tokenwise fast path，看 aggregate Output TPS 是否继续上涨。 1. 测试 `max_num_seqs=4`，利用 `T <= 4` tokenwise fast path，看 aggregate Output TPS 是否继续上涨。
2. 继续优化 shared expert / linear_attention，因为 routed expert 已经明显改善，decode 下一个大头会逐渐转向 `linear_attention` 和 shared expert。 2. 继续优化 shared expert / linear_attention，因为 routed expert 已经明显改善，decode 下一个大头会逐渐转向 `linear_attention` 和 shared expert。
## 2026-07-15：平台提交构建/启动配置修复
### 背景
官方平台仍提示当前提交的 Docker 镜像构建失败。对比一个已能进入 benchmark-agent 阶段的成功案例仓库 `fangfangssj/enginex-vllm-bi100-qwen36` 后，确认我们的根目录结构没有问题，根目录已有 `Dockerfile` 和 `computility-run.yaml`。
成功案例日志中的 `BENCHMARK_DOCKER_IMAGE_NOT_READY` 只是镜像正在下载，随后已经创建 benchmark-agent 容器；因此该案例至少通过了代码仓库识别、镜像构建和镜像拉取阶段。
### 对齐内容
本次将当前仓库的平台提交相关文件对齐成功案例：
- `computility-run.yaml`
- 上下文从 `100000` 调整为 `245000`，满足官方说明中 235K+ 输入窗口需求。
- `max_num_batched_tokens` 从 `8192` 调整为 `16384`。
- 增加 `--chat-template /workspace/chat_template_multi_system.jinja`，兼容多/乱序 system 消息。
- 增加 `PYTHONPATH`、`LD_LIBRARY_PATH`、`PATH` 和五个 native 优化开关环境变量。
- `qwen3_6_scripts/patch_ops.sh`
- 增加 CoreX/OpenMPI 环境变量。
- 增加 `libcusolver.so` 软链，避免 `torch.linalg.solve_triangular` 在容器内找不到库。
- 增加 multi-system chat template 部署。
- 增加 patch 后验证输出。
- 新增/更新成功案例依赖文件：
- `qwen3_6_scripts/chat_template_multi_system.jinja`
- `qwen3_6_scripts/xformers.py`
- `qwen3_6_scripts/paged_attn.py`
- `qwen3_6_scripts/patch_xformers_sdpa_seq.py`
- `qwen3_6_scripts/protocol.py`
- `qwen3_6_scripts/qwen3_5.py`
### 影响
这次优先目标是恢复平台可构建、可启动、可进入 benchmark-agent 验证流程。为降低平台构建风险，直接采用成功案例中已验证的一组补丁文件；此前本地 MoE tiny-batch 实验代码可能被成功案例的 native MoE/GDN 优化实现覆盖，后续如果平台构建通过，再基于这套可提交基线继续做性能优化分支。