dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

sync: align with enginex-vllm-bi100-qwen36 baseline (1902c81f)

Browse Source 
Synced files from EngineX baseline zip (2026-06-30):
- ADD paged_attn.py (root): production paged attention with PyTorch fallback
- ADD launch_service: BI-V100 server startup script with env configuration
- SYNC computility-run.yaml: gpu_memory=0.9, batched_tokens=8192, seq_capture=32768
- SYNC qwen3_6_scripts/paged_attn.py: +311 lines, Triton bypass docs, _forward_decode_pytorch shape docs
- SYNC qwen3_6_scripts/qwen3_5.py: -72 lines, revert optimized MoE prefill to baseline (untested on BI-V100)
- KEEP Dockerfile: repo version has V2/Triton/head256 optimization patches not in baseline

Baseline commit: 1902c81fdd373943f17f5983eb8750758c7f4a69
Source: enginex-vllm-bi100-qwen36-main.zip (dev.modelhub.org.cn)
This commit is contained in:
Claude
2026-07-31 09:43:58 +00:00
parent de7ee4383e
commit c5a0d61851
5 changed files with 535 additions and 169 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
computility-run.yaml
View File

@@ -10,7 +10,7 @@ command:
- --max-model-len
- '100000'
- --gpu-memory-utilization
- '0.95'
- '0.9'
- --trust-remote-code
- -tp
- '4'
@@ -19,10 +19,10 @@ command:
- --disable-log-requests
- --disable-frontend-multiprocessing
- --max-num-batched-tokens
- '16384'
- '8192'
- --enable-chunked-prefill
- --max-seq-len-to-capture
- '65536'
- '32768'
- --enable-auto-tool-choice
- --tool-call-parser
- qwen3_coder

80
launch_service Executable file
View File

@@ -0,0 +1,80 @@
#!/bin/bash
export PYTHONPATH=/usr/local/corex/lib64/python3/dist-packages
export LD_LIBRARY_PATH=/usr/local/corex/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/corex/lib64/python3/dist-packages/bin:/usr/local/openmpi/bin
export JAVA_HOME=/root/apps/jdk1.8.0_411
export JRE_HOME=/root/apps/jdk1.8.0_411/jre
export JMETER_HOME=/root/apps/apache-jmeter-5.6.3
export CLASSPATH=.:/root/apps/jdk1.8.0_411/lib/dt.jar:/root/apps/jdk1.8.0_411/lib/tools.jar:/root/apps/apache-jmeter-5.6.3/lib/ext/ApacheJMeter_core.jar:/root/apps/apache-jmeter-5.6.3/lib/jorphan.jar:/root/apps/apache-jmeter-5.6.3/lib/logkit-2.0.jar:
export PATH=/root/apps/apache-jmeter-5.6.3/bin:/root/apps/jdk1.8.0_411/bin:/usr/local/corex/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/local/corex/lib64/python3/dist-packages/bin:/usr/local/openmpi/bin
/iluvatar/welcome.sh
data
cat /proc/cpuinfo | tail -n 50
ixsmi
unset CUDA_VISIBLE_DEVICES
export
date
DEFAULT_HOST="0.0.0.0"
DEFAULT_PORT="80"
DEFAULT_SERVED_MODEL_NAME="llm"
DEFAULT_MODEL_PATH="/model"
DEFAULT_MAX_MODEL_LEN="10000"
DEFAULT_TENSOR_PARALLEL_SIZE="1"
DEFAULT_MAX_NUM_SEQS="64"
DEFAULT_ENFORCE_EAGER="true"
DEFAULT_DISABLE_LOG_REQUESTS="true"
DEFAULT_PREFIX_CACHING="true"
HOST_VAL=${HOST:-$DEFAULT_HOST}
PORT_VAL=${PORT:-$DEFAULT_PORT}
SERVED_MODEL_NAME_VAL=${SERVED_MODEL_NAME:-$DEFAULT_SERVED_MODEL_NAME}
MODEL_PATH_VAL=${MODEL_PATH:-$DEFAULT_MODEL_PATH}
MAX_MODEL_LEN_VAL=${MAX_MODEL_LEN:-$DEFAULT_MAX_MODEL_LEN}
TENSOR_PARALLEL_SIZE_VAL=${TENSOR_PARALLEL_SIZE:-$DEFAULT_TENSOR_PARALLEL_SIZE}
MAX_NUM_SEQS_VAL=${MAX_NUM_SEQS:-$DEFAULT_MAX_NUM_SEQS}
INCLUDE_ENFORCE_EAGER_FLAG=${ENFORCE_EAGER:-$DEFAULT_ENFORCE_EAGER}
INCLUDE_DISABLE_LOG_REQUESTS_FLAG=${DISABLE_LOG_REQUESTS:-$DEFAULT_DISABLE_LOG_REQUESTS}
INCLUDE_PREFIX_CACHING_FLAG=${PREFIX_CACHING:-$DEFAULT_PREFIX_CACHING}
CMD_ARGS=()
CMD_ARGS+=(--host "$HOST_VAL")
CMD_ARGS+=(--port "$PORT_VAL")
if [[ "$INCLUDE_ENFORCE_EAGER_FLAG" != "false" && "$INCLUDE_ENFORCE_EAGER_FLAG" != "0" ]]; then
CMD_ARGS+=(--enforce-eager)
fi
if [[ "$INCLUDE_DISABLE_LOG_REQUESTS_FLAG" != "false" && "$INCLUDE_DISABLE_LOG_REQUESTS_FLAG" != "0" ]]; then
CMD_ARGS+=(--disable-log-requests)
fi
if [[ "$INCLUDE_PREFIX_CACHING_FLAG" != "false" && "$INCLUDE_PREFIX_CACHING_FLAG" != "0" ]]; then
CMD_ARGS+=(--enable-prefix-caching)
fi
CMD_ARGS+=(--served-model-name "$SERVED_MODEL_NAME_VAL")
CMD_ARGS+=(--model "$MODEL_PATH_VAL")
CMD_ARGS+=(--max-model-len "$MAX_MODEL_LEN_VAL")
CMD_ARGS+=(--tensor-parallel-size "$TENSOR_PARALLEL_SIZE_VAL")
CMD_ARGS+=(--max-num-seqs "$MAX_NUM_SEQS_VAL")
CMD_ARGS+=(--trust-remote-code)
echo "--------------------------------------------------"
echo "Starting VLLM OpenAI API Server..."
echo "Using effective arguments:"
echo " Host (--host): $HOST_VAL"
echo " Port (--port): $PORT_VAL"
echo " Enforce Eager (--enforce-eager):" $([[ "$INCLUDE_ENFORCE_EAGER_FLAG" != "false" && "$INCLUDE_ENFORCE_EAGER_FLAG" != "0" ]] && echo "Enabled" || echo "Disabled (Env: ENFORCE_EAGER=$ENFORCE_EAGER)")
echo " Disable Log Req (--disable-log-requests):" $([[ "$INCLUDE_DISABLE_LOG_REQUESTS_FLAG" != "false" && "$INCLUDE_DISABLE_LOG_REQUESTS_FLAG" != "0" ]] && echo "Enabled" || echo "Disabled (Env: DISABLE_LOG_REQUESTS=$DISABLE_LOG_REQUESTS)")
echo " Served Model Name (--served-model-name): $SERVED_MODEL_NAME_VAL"
echo " Model Path (--model): $MODEL_PATH_VAL"
echo " Max Model Length (--max-model-len): $MAX_MODEL_LEN_VAL"
echo " Tensor Parallel Size (--tensor-parallel-size): $TENSOR_PARALLEL_SIZE_VAL"
echo " Max Num Seqs (--max-num-seqs): $MAX_NUM_SEQS_VAL"
echo "--------------------------------------------------"
echo "Full cmd:"
echo "python3 -m vllm.entrypoints.openai.api_server ${CMD_ARGS[*]}"
echo "--------------------------------------------------"
python3 -m vllm.entrypoints.openai.api_server "${CMD_ARGS[@]}"

243
paged_attn.py Normal file
View File

@@ -0,0 +1,243 @@
from dataclasses import dataclass
from typing import List, Optional, Tuple
import torch
from vllm import _custom_ops as ops
from vllm.attention.ops.prefix_prefill import context_attention_fwd
# 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,
key_cache,
value_cache,
slot_mapping.flatten(),
kv_cache_dtype,
k_scale,
v_scale,
)
@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:
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:
output = torch.empty_like(query)
context_attention_fwd(
query,
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,
k_scale,
v_scale,
alibi_slopes,
sliding_window,
)
return output
@staticmethod
def swap_blocks(
src_kv_cache: torch.Tensor,
dst_kv_cache: torch.Tensor,
src_to_dst: torch.Tensor,
) -> None:
src_key_cache = src_kv_cache[0]
dst_key_cache = dst_kv_cache[0]
ops.swap_blocks(src_key_cache, dst_key_cache, src_to_dst)
src_value_cache = src_kv_cache[1]
dst_value_cache = dst_kv_cache[1]
ops.swap_blocks(src_value_cache, dst_value_cache, src_to_dst)
@staticmethod
def copy_blocks(
kv_caches: List[torch.Tensor],
src_to_dists: torch.Tensor,
) -> None:
key_caches = [kv_cache[0] for kv_cache in kv_caches]
value_caches = [kv_cache[1] for kv_cache in kv_caches]
ops.copy_blocks(key_caches, value_caches, src_to_dists)

303
qwen3_6_scripts/paged_attn.py
View File

@@ -5,6 +5,12 @@ 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
@@ -12,8 +18,17 @@ _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]
@@ -40,8 +55,10 @@ class PagedAttention:
) -> 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)
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
@@ -62,7 +79,7 @@ class PagedAttention:
value,
key_cache,
value_cache,
slot_mapping,
slot_mapping.flatten(),
kv_cache_dtype,
k_scale,
v_scale,
@@ -70,18 +87,34 @@ class PagedAttention:
@staticmethod
def _forward_decode_pytorch(
query, key_cache, value_cache, block_tables, seq_lens, scale
):
"""Pure-PyTorch decode fallback for seq_len > threshold.
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).
Used when ixf_F.paged_attention_v1 cannot handle the sequence length.
Optimized with batched KV gather (no per-block Python loop).
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:
@@ -90,25 +123,33 @@ class PagedAttention:
num_blocks = (seq_len + block_size - 1) // block_size
blk_ids = block_tables[i, :num_blocks]
# Batched gather: one index_select for all 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()
.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()
.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))
attn_w = torch.matmul(q_grouped * scale, k_t.unsqueeze(1))
# [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)
@@ -120,9 +161,10 @@ class PagedAttention:
return output
# BI-V100: Try higher threshold for compiled v1 kernel.
# Compiled kernel is ~100x faster than Python fallback.
_PYTORCH_DECODE_THRESHOLD = 65536
# 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(
@@ -149,17 +191,31 @@ class PagedAttention:
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))
# V2 now works (paged_attention_v2_pytorch), so use the original heuristic
# instead of hardcoding use_v1=True
use_v1 = True
if use_v1:
# Run PagedAttention V1.
ops.paged_attention_v1(
output,
query,
@@ -174,6 +230,7 @@ class PagedAttention:
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),
@@ -212,9 +269,6 @@ class PagedAttention:
)
return output
# Triton prefill: try once, fall back permanently if it fails
_triton_prefill_ok = None
@staticmethod
def forward_prefix(
query: torch.Tensor,
@@ -233,30 +287,9 @@ class PagedAttention:
k_scale: float,
v_scale: float,
) -> torch.Tensor:
# Try Triton kernel if available and not known to fail
if PagedAttention._triton_prefill_ok is not False:
try:
from vllm.triton_utils import HAS_TRITON
if HAS_TRITON:
from vllm.attention.ops.prefix_prefill import context_attention_fwd
output = torch.empty_like(query)
context_attention_fwd(
query, key, value, output, kv_cache_dtype,
key_cache, value_cache, block_tables,
query_start_loc[:-1], seq_lens_tensor, context_lens,
max_query_len, k_scale, v_scale,
alibi_slopes, sliding_window,
)
if PagedAttention._triton_prefill_ok is None:
print("[paged_attn] Triton prefill kernel: SUCCESS", flush=True)
PagedAttention._triton_prefill_ok = True
return output
except Exception as e:
print(f"[paged_attn] Triton prefill failed: {type(e).__name__}: {e}",
flush=True)
print("[paged_attn] Falling back to PyTorch prefill permanently", flush=True)
PagedAttention._triton_prefill_ok = False
# 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,
@@ -276,101 +309,150 @@ class PagedAttention:
seq_lens_tensor: torch.Tensor,
context_lens: torch.Tensor,
) -> torch.Tensor:
"""Pure-PyTorch prefix-attention with pre-gathered KV and Flash-Attention online softmax.
"""Pure-PyTorch prefix-attention with K-tiling (Flash-Attention online softmax).
Optimization over baseline:
- Context KV is gathered ONCE outside the tile loop (one index_select + reshape)
- Tile loop just slices views from the pre-gathered tensor (no per-tile gather)
- Eliminates 194 redundant permute+contiguous calls for 100K context
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.
Memory: O(q_len × tile_sz) per tile — same as baseline.
The full context K/V tensor is ~50MB for 100K tokens, fits in GPU memory.
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]
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
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_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]
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 — PRE-GATHER optimization
# Gather ALL context K/V in ONE shot, then tile via slicing
# ===========================================================
# --------------------------------------------------------------
# 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]}. "
"Capping context to available 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]
# ONE gather for ALL context blocks
ctx_blk_ids = block_tables[i, :num_ctx_blocks]
# 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))
# [num_ctx_blocks, kv_h, d/x, blk_sz, x] → [ctx_tokens, kv_h, d]
ctx_k_all = (key_cache[ctx_blk_ids]
.permute(0, 3, 1, 2, 4)
.contiguous()
.view(-1, num_kv_heads, head_dim))[:ctx_len]
# 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]
ctx_v_all = (value_cache[ctx_blk_ids]
.permute(0, 3, 1, 2)
.contiguous()
.view(-1, num_kv_heads, head_dim))[:ctx_len]
# Pre-transpose for matmul: [kv_h, d, ctx_len] and [kv_h, ctx_len, d]
ctx_k_t = ctx_k_all.permute(1, 2, 0).contiguous().float()
ctx_v_t = ctx_v_all.permute(1, 0, 2).contiguous().float()
# Tile loop: just SLICE from pre-gathered tensors
for tile_start in range(0, ctx_len, tile_sz):
tile_end = min(tile_start + tile_sz, ctx_len)
# Slice (view, no copy)
k_t = ctx_k_t[:, :, tile_start:tile_end].unsqueeze(1)
v_t = ctx_v_t[:, tile_start:tile_end, :].unsqueeze(1)
# 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)
corr = torch.exp(m - m_new)
m.copy_(m_new)
del m_blk, m_new
l.mul_(corr).add_(exp_s.sum(dim=-1))
@@ -378,40 +460,44 @@ class PagedAttention:
torch.matmul(exp_s, v_t))
del exp_s, v_t, corr
del ctx_k_t, ctx_v_t, ctx_k_all, ctx_v_all
# ===========================================================
# Phase 2: Current-chunk tokens (with causal mask)
# ===========================================================
# --------------------------------------------------------------
# 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]
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())
.float()) # [kv_h, 1, d, kc_len]
v_t = (v_blk.permute(1, 0, 2)
.unsqueeze(1)
.float())
.float()) # [kv_h, 1, kc_len, d]
s = torch.matmul(q_seq, k_t)
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)
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)
corr = torch.exp(m - m_new)
m.copy_(m_new)
del m_blk, m_new
l.mul_(corr).add_(exp_s.sum(dim=-1))
@@ -419,6 +505,10 @@ class PagedAttention:
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)
@@ -442,6 +532,7 @@ class PagedAttention:
src_key_cache = src_kv_cache[0]
dst_key_cache = dst_kv_cache[0]
ops.swap_blocks(src_key_cache, dst_key_cache, src_to_dst)
src_value_cache = src_kv_cache[1]
dst_value_cache = dst_kv_cache[1]
ops.swap_blocks(src_value_cache, dst_value_cache, src_to_dst)

72
qwen3_6_scripts/qwen3_5.py
View File

@@ -795,66 +795,18 @@ class Qwen3_5MoeSparseBlock(nn.Module):
# General path (prefill / multi-seq): loop over unique active experts.
# At most T*top_k unique experts, always <= num_experts.
out = torch.zeros_like(hidden_states)
# Optimized prefill MoE: sort tokens by expert for contiguous access.
# Pattern from CCCL segmented sort: partition input by key (expert),
# process each segment on contiguous memory, scatter results back.
#
# Before: for eid in unique_eids: hidden_states[tok_ids] (scattered gather)
# After: sort by expert → each F.linear gets contiguous input
# activation computed in ONE fused op across all token×expert pairs
# index_add_ for scatter-back is one kernel
T_local = hidden_states.shape[0]
two_I = w13.shape[1]
# Flatten token-expert assignments: each token appears top_k times
flat_ids = topk_ids.view(-1) # (T*K,)
flat_weights_f = topk_weights.view(-1) # (T*K,)
flat_tok_idx = torch.arange(
T_local, device=hidden_states.device
).unsqueeze(1).expand(-1, self.top_k).reshape(-1) # (T*K,)
# Sort by expert ID → contiguous memory per expert
sorted_order = flat_ids.argsort(stable=True)
sorted_eids = flat_ids[sorted_order]
sorted_tok_idx = flat_tok_idx[sorted_order]
sorted_weights_f = flat_weights_f[sorted_order]
sorted_tokens = hidden_states[sorted_tok_idx] # (T*K, H) contiguous groups
# Expert boundaries via unique_consecutive
unique_eids_t, counts = torch.unique_consecutive(
sorted_eids, return_counts=True)
cum_counts = counts.cumsum(0)
starts = torch.cat([torch.zeros(1, dtype=torch.long, device=hidden_states.device),
cum_counts[:-1]])
# Grouped up-projection: each expert's tokens are contiguous
TK = T_local * self.top_k
gate_up_all = torch.empty(TK, two_I, dtype=hidden_states.dtype,
device=hidden_states.device)
for idx_e in range(len(unique_eids_t)):
eid_val = unique_eids_t[idx_e].item()
s = starts[idx_e].item()
e = cum_counts[idx_e].item()
gate_up_all[s:e] = F.linear(sorted_tokens[s:e], w13[eid_val])
# Fused activation across ALL token×expert pairs (one kernel)
gate, up = gate_up_all.chunk(2, dim=-1)
act = F.silu(gate) * up
# Grouped down-projection
down_all = torch.empty(TK, hidden_states.shape[-1],
dtype=hidden_states.dtype,
device=hidden_states.device)
for idx_e in range(len(unique_eids_t)):
eid_val = unique_eids_t[idx_e].item()
s = starts[idx_e].item()
e = cum_counts[idx_e].item()
down_all[s:e] = F.linear(act[s:e], w2[eid_val])
# Weighted scatter-add (one kernel)
weighted = down_all * sorted_weights_f.unsqueeze(-1)
out.index_add_(0, sorted_tok_idx, weighted.to(out.dtype))
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()