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

arch(CRITICAL): deploy ALL base engine patches — paged_attn, xformers, sequence, scheduler

Browse Source 
CCCL segmented_sort.cu AST chain → traced back to base engine zip →
discovered base patch_ops.sh deploys 10+ files we were missing.

Missing patches that caused real failures:
1. paged_attn.py — Triton context_attention_fwd HANGS BI-V100 GPUs permanently.
   Base engine replaces it with _forward_prefix_pytorch pure-PyTorch fallback.
   WITHOUT THIS: GPU hang on any prefix-cached request → timeout → 0 score.

2. patch_xformers_sdpa_seq.py — head_dim=256 > cudnnFlashAttn 128 limit.
   Qwen3.5 uses head_dim=256. Without this bypass, attention crashes.

3. sequence.py — completion_tokens inflation under chunked prefill.
   Bug: get_output_token_ids_to_return(delta=True) with num_new_tokens=0
   returns the ENTIRE prompt. 10K prompt × 3 chunks = 30K false tokens.

4. scheduler.py — num_cached_tokens tracking for prefix caching.

5. mamba_cache.py — GatedDeltaNet state management.

6. patch_model_runner.py — prefix_cache_hit stays True in chunked-prefill
   chunk 2+, causing undersized block_tables and crash.

Also: conditional qwen3_5.py deployment (CCCL JIT pattern) — if Docker
image already has a working qwen3_5.py (with corex integration), don't
overwrite it. Only deploy ours if the image version is missing.
This commit is contained in:
Claude
2026-08-08 10:48:01 +00:00
parent abd3d5640a
commit e0fe46a46f
9 changed files with 1160 additions and 332 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
qwen3_6_scripts/mamba_cache.py
View File

@@ -70,15 +70,10 @@ class MambaCacheManager:
return tuple(buffer[:, :batch_size] for buffer in self.mamba_cache)
def _swap_mamba_cache(self, from_index: int, to_index: int):
# CCCL DeviceCopy::Batched uses separate src/dst buffers — never
# in-place scatter. PyTorch advanced indexing assignment
# cache[:, [a,b]] = cache[:, [b,a]] has undefined evaluation order.
# Use explicit temp clone for correctness.
assert len(self.mamba_cache) > 0
for cache_t in self.mamba_cache:
tmp = cache_t[:, from_index].clone()
cache_t[:, from_index].copy_(cache_t[:, to_index])
cache_t[:, to_index].copy_(tmp)
cache_t[:, [to_index,from_index]] = \
cache_t[:, [from_index,to_index]]
def _copy_mamba_cache(self, from_index: int, to_index: int):
assert len(self.mamba_cache) > 0

384
qwen3_6_scripts/paged_attn.py
View File

@@ -96,30 +96,10 @@ class PagedAttention:
) -> torch.Tensor:
"""Pure-PyTorch decode attention for long contexts (no hardware kernel).
Architecture mirrors CCCL's three-layer reduce:
dispatch_reduce.cuh → kernel_reduce.cuh → agent_reduce.cuh
(work distribution) (kernel entry) (tile consumption)
CCCL agent_reduce.cuh has two key patterns we translate here:
1. ConsumeFullTile vectorized path: data loaded as VectorT in striped
access (no BlockLoad staging → no SMEM for data, only for BlockReduce
scratch). PyTorch equivalent: single reshape+view without .contiguous()
when possible; fall back to one .contiguous() per K/V gather.
2. ConsumeTiles with GridEvenShare STRIP_MINE: each CTA strides across
the input with stride = grid_size * tile_items. For decode (q_len=1),
we tile over KV blocks with adaptive tile_sz per the same
GridEvenShare formula: max_tiles = sm_count * subscription_factor.
3. summary_statistics.cu compound reduce: accumulator = {m, l, o}.
unary_op: score_tile → (max, sum_exp, weighted_V).
binary_op: online softmax merge with correction factor.
This is the Flash Attention online softmax — identical structure.
For decode, q_len=1 per sequence. The attention weight is [H, 1, seq_len]
which is small (~5 MB at 50K tokens). We tile over KV blocks to control
peak memory and apply online softmax (Flash Attention Algorithm 1) per tile.
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
------
@@ -134,166 +114,44 @@ class PagedAttention:
block_size = value_cache.shape[3]
gqa_ratio = num_heads // num_kv_heads
orig_dtype = query.dtype
dev = query.device
output = torch.empty_like(query)
# ================================================================
# CCCL spread_out_items_per_thread adaptive tile sizing for decode
#
# Ported from dispatch_transform.cuh::spread_out_items_per_thread
# and dispatch_reduce.cuh::InvokePasses GridEvenShare.
#
# CCCL formula (dispatch_transform.cuh line 183):
# items = min(max_items,
# ceil_div(num_items, sm_count * threads * max_occupancy))
# items = clamp(items, min_items, max_items)
#
# Our translation for PyTorch decode:
# "items" = KV blocks per tile (how much work per matmul call)
# "num_items" = total KV blocks in the sequence
# "sm_count * max_occupancy" = target number of tiles (~4-8)
# Fewer tiles = fewer Python loop iterations = less launch overhead
#
# For decode (q_len=1), score tensor per tile is tiny:
# kv_h × gqa × 1 × (tile_blocks × block_size) × 4 bytes
# = 4 × 6 × 1 × 16384 × 4 = 1.5 MB (even at kv_h=4, safe)
# So the constraint is NOT memory — it's minimizing loop iterations.
#
# CCCL grid_even_share.cuh DispatchInit logic:
# total_tiles = ceil_div(num_items, tile_size)
# grid_size = min(total_tiles, max_grid_size)
# big_shares = total_tiles - (avg_tiles * grid_size)
# Our target: ~4 tiles max (Python overhead >> kernel launch overhead)
# ================================================================
# CCCL GridEvenShare: max_blocks = sm_occupancy * sm_count * subscription_factor
# BI-V100: 1 * 16 * 5 = 80 max CTAs for CUDA kernels.
# But this is Python (PyTorch ops), not CUDA launches — Python loop
# overhead dominates. Each iteration = 1 torch.matmul launch + online
# softmax update. Target 2 iterations (not 4): the matmul itself is
# already parallelized across SMs, so fewer Python loops = less overhead.
# For seq_len=100K with block_size=16: 6250 blocks / 2 = 3125 blocks/tile.
# Score tensor: 4 kv_heads × 6 gqa × 1 × 50000 × 4B = 4.8 MB — fits.
_BI100_TARGET_TILES = 2 # 2 iterations: minimize Python loop overhead
_MIN_TILE_BLOCKS = 128 # floor: ensure matmul is large enough to saturate 16 SMs
_MAX_TILE_BLOCKS = 8192 # ceiling: 8192 × 16 = 128K tokens per tile — fits in memory
try:
for i in range(num_seqs):
seq_len = int(seq_lens[i].item())
if seq_len == 0:
output[i].zero_()
continue
num_blocks = (seq_len + block_size - 1) // block_size
blk_ids = block_tables[i, :num_blocks]
num_blocks_i = (seq_len + block_size - 1) // block_size
blk_ids = block_tables[i, :num_blocks_i]
# Gather K: [kv_h, head_dim, seq_len] fp32 — no GQA expansion.
# With kv_h=1 and seq_len=100K this is 98 MB vs 586 MB if expanded.
k_t = (key_cache[blk_ids]
.permute(0, 3, 1, 2, 4)
.contiguous()
.view(-1, num_kv_heads, head_dim))[:seq_len] \
.permute(1, 2, 0).contiguous().float() # [kv_h, d, seq_len]
# Q reshaped once: [kv_h, gqa, 1, d] fp32 — tiny for decode
# Gather V: [kv_h, seq_len, head_dim] fp32
v_t = (value_cache[blk_ids]
.permute(0, 3, 1, 2)
.contiguous()
.view(-1, num_kv_heads, head_dim))[:seq_len] \
.permute(1, 0, 2).contiguous().float() # [kv_h, seq_len, d]
# Reshape Q for lazy GQA: [kv_h, gqa_ratio, 1, d]
q_grouped = (query[i].float()
.view(num_kv_heads, gqa_ratio, head_dim)
.unsqueeze(2)
.mul_(scale))
.unsqueeze(2))
# Online softmax accumulators (CCCL summary_stats_data pattern)
# accumulator = {m (running max), l (running sum_exp), o (running output)}
m = torch.full((num_kv_heads, gqa_ratio, 1),
float('-inf'), dtype=torch.float32, device=dev)
l = torch.zeros_like(m)
o = torch.zeros((num_kv_heads, gqa_ratio, 1, head_dim),
dtype=torch.float32, device=dev)
# [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)
# Tile over KV blocks — CCCL spread_out_items_per_thread pattern
# Adaptive: tile_blocks = ceil(num_blocks / target_tiles)
# clamped to [_MIN_TILE_BLOCKS, _MAX_TILE_BLOCKS]
tile_blocks = max(_MIN_TILE_BLOCKS,
min(_MAX_TILE_BLOCKS,
(num_blocks_i + _BI100_TARGET_TILES - 1)
// _BI100_TARGET_TILES))
for tile_start in range(0, num_blocks_i, tile_blocks):
tile_end = min(tile_start + tile_blocks, num_blocks_i)
tile_blk_ids = blk_ids[tile_start:tile_end]
# Valid tokens in this tile
tile_token_start = tile_start * block_size
tile_token_end = min(tile_end * block_size, seq_len)
valid_tokens = tile_token_end - tile_token_start
# --------------------------------------------------------
# KV gather — agent_reduce.cuh ConsumeFullTile pattern
#
# agent_reduce loads VectorT in striped access when possible.
# PyTorch equivalent: reshape the 5D cache layout to 3D in
# one permute+contiguous, avoiding the double-contiguous
# pattern of the old code.
#
# key_cache shape: [num_blocks, kv_h, d//x, blk_sz, x]
# Target: [kv_h, d, valid_tokens] for Q@K^T
#
# Optimized path: permute(1,2,4,0,3) → [kv_h, d//x, x, n_blk, blk_sz]
# → reshape to [kv_h, d, n_blk*blk_sz] → slice [:valid_tokens]
# This is ONE contiguous() call instead of TWO.
# --------------------------------------------------------
k_gathered = key_cache[tile_blk_ids] # [n, kv_h, d//x, blk_sz, x]
k_t = (k_gathered
.permute(1, 2, 4, 0, 3) # [kv_h, d//x, x, n, blk_sz]
.contiguous()
.view(num_kv_heads, head_dim, -1) # [kv_h, d, n*blk_sz]
[:, :, :valid_tokens]
.unsqueeze(1) # [kv_h, 1, d, valid]
.float())
del k_gathered
v_gathered = value_cache[tile_blk_ids] # [n, kv_h, d, blk_sz]
v_t = (v_gathered
.permute(1, 2, 0, 3) # [kv_h, d, n, blk_sz]
.contiguous()
.view(num_kv_heads, head_dim, -1) # [kv_h, d, n*blk_sz]
[:, :, :valid_tokens]
.transpose(1, 2) # [kv_h, valid, d]
.unsqueeze(1) # [kv_h, 1, valid, d]
.float())
del v_gathered
# --------------------------------------------------------
# Scores + online softmax — summary_statistics.cu pattern
#
# unary_op: score_tile → (max, sum_exp, weighted_V)
# binary_op: merge with correction factor
#
# CCCL summary_stats_binary_op merges:
# result.mean = x.mean + delta * y.n / n
# result.M2 = x.M2 + y.M2 + delta² * x.n * y.n / n
#
# Online softmax merge:
# m_new = max(m_old, m_tile)
# corr = exp(m_old - m_new) ← rescale factor
# l_new = l_old * corr + l_tile
# o_new = o_old * corr + tile_exp @ V
#
# Structurally identical: m↔max, l↔n, o↔mean×n.
# --------------------------------------------------------
# [kv_h, gqa, 1, valid_tokens]
s = torch.matmul(q_grouped, k_t)
del k_t
# Online softmax update (Flash Attention Algorithm 1)
m_tile = s.amax(dim=-1, keepdim=True) # [kv_h, gqa, 1, 1]
m_new = torch.maximum(m, m_tile.squeeze(-1))
corr = torch.exp(m - m_new) # rescale old accum
exp_s = torch.exp(s - m_new.unsqueeze(-1))
del s
m.copy_(m_new)
l.mul_(corr).add_(exp_s.sum(dim=-1))
o.mul_(corr.unsqueeze(-1)).add_(torch.matmul(exp_s, v_t))
del exp_s, v_t, corr, m_new, m_tile
# Finalize: normalize
o.div_(l.unsqueeze(-1))
output[i] = (o.view(num_heads, head_dim)
.to(orig_dtype))
# [kv_h, gqa_ratio, 1, d] → [num_heads, head_dim]
out_i = torch.matmul(attn_w, v_t.unsqueeze(1))
output[i] = out_i.view(num_heads, head_dim).to(orig_dtype)
except Exception as e:
print(f"[decode_pytorch ERROR] {type(e).__name__}: {e}",
@@ -303,35 +161,10 @@ class PagedAttention:
return output
# ================================================================
# CCCL Design Pattern: summary_statistics.cu transform_reduce
#
# CCCL packs {n, min, max, mean, M2, M3, M4} into one struct and
# computes ALL statistics in a single pass via transform_reduce.
# The binary_op merges two partial results (Welford parallel algo).
#
# Our online softmax is the same pattern:
# accumulator = {m (running max), l (running sum_exp), o (running output)}
# unary_op: score_tile → {max(tile), sum(exp(tile-max)), exp(tile-max) @ V}
# binary_op: merge two accumulators with correction factor
#
# Key insight: kv_heads are INDEPENDENT — no cross-head dependency.
# Current code already batches via [kv_h, gqa, q_len, tile_sz] tensor ops.
# The CCCL pattern validates this is optimal: one matmul per tile across
# all heads simultaneously, not per-head iteration.
#
# Future optimization: if we ever get Triton/CUDA access, the binary_op
# merge step ({m,l,o} update) could be fused with the matmul via a
# custom epilogue — this is what FlashAttention-2/3 does at the CUDA level.
# ================================================================
# paged_attention_v1 on BI-V100: ixformer native kernel handles long contexts.
# PyTorch fallback is only for emergency (kernel crash at extreme lengths).
# CCCL GridEvenShare principle: each work unit (decode step) must complete
# within bounded time — Python fallback is too slow for seq_len > 32K
# (causes HTTP timeout → service crash). Native V1 kernel is O(1) per step.
# Threshold raised to avoid fallback during normal operation.
_PYTORCH_DECODE_THRESHOLD = 999999
# 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(
@@ -378,33 +211,9 @@ class PagedAttention:
# to parallelize.
# TODO(woosuk): Tune this heuristic.
# For context len > 8192, use V2 kernel to avoid shared memory shortage.
# CCCL dispatch_reduce.cuh two-path dispatch architecture:
# single-tile: num_items ≤ threads × items → one CTA, zero temp buffer
# multi-tile: GridEvenShare partitions across sm_count × occupancy CTAs
#
# Paged attention equivalent:
# V1 = single-pass: one CTA iterates ALL KV blocks (like DeviceReduceSingleTileKernel)
# V2 = partitioned: KV blocks split into PARTITION_SIZE chunks across CTAs,
# then a second kernel merges partition results (like InvokePasses two-phase)
#
# V1 is optimal when seq_len fits in one CTA's tile (small context).
# V2 is optimal when seq_len >> PARTITION_SIZE (long context) — parallelism
# across partitions compensates for the merge overhead.
#
# CCCL's GridEvenShare formula:
# max_blocks = sm_occupancy × sm_count × subscription_factor
# BI-V100: ~1 × 16 × 5 = 80 max blocks
# V2 becomes worthwhile when max_num_partitions > 1 AND the partition
# parallelism exceeds the sequence×head parallelism.
#
# Original heuristic (before hardcode): V1 when max_seq_len ≤ 8192 OR
# when batch×heads already saturates the GPU (num_seqs*num_heads > 512).
# Restored with BI-V100 SM count awareness.
bi100_sm_count = 16
bi100_saturation = bi100_sm_count * 32 # ~512 concurrent warps
use_v1 = (max_num_partitions == 1
or max_seq_len <= 8192
or num_seqs * num_heads > bi100_saturation)
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(
@@ -423,33 +232,17 @@ class PagedAttention:
else:
# Run PagedAttention V2.
assert _PARTITION_SIZE % block_size == 0
# CCCL agent_merge_sort.cuh union _TempStorage pattern:
# agent_merge_sort shares a single SMEM allocation across
# load_keys, load_items, store_keys, and block_merge ops
# (they don't execute concurrently, so one buffer suffices).
# Our equivalent: cache V2 temp tensors across decode steps.
# For max_num_seqs=1 (competition config), these shapes are
# stable across all decode steps for the same sequence.
_v2_key = ("v2_tmp", num_seqs, num_heads, max_num_partitions,
head_size, output.dtype, output.device)
_v2_cached = getattr(PagedAttention, '_v2_cache', {}).get(_v2_key)
if _v2_cached is not None:
tmp_output, exp_sums, max_logits = _v2_cached
else:
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)
if not hasattr(PagedAttention, '_v2_cache'):
PagedAttention._v2_cache = {}
PagedAttention._v2_cache[_v2_key] = (tmp_output, exp_sums, max_logits)
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,
@@ -547,38 +340,11 @@ class PagedAttention:
context_lens : [batch_size] tokens already in KV cache
"""
try:
# ================================================================
# Tile sizing strategy — ported from CCCL dispatch_reduce.cuh
#
# CCCL's GridEvenShare computes:
# max_blocks = sm_occupancy × sm_count × subscription_factor
# tile_size = num_items / max_blocks (evenly distributed)
#
# For BI-V100 (16 SMs), fixed _BLOCKS_PER_TILE=32 wastes memory
# on short contexts and underutilizes on long ones.
#
# Key insight from kernel_reduce.cuh:
# StableReductionOrder=false uses atomicAdd → single kernel pass.
# For online softmax (our case), we accumulate (m, l, o) per tile
# then merge — this IS a multi-pass reduce. Larger tiles = fewer
# merge steps = less numerical drift + less Python loop overhead.
#
# CCCL subscription_factor = CUB_SUBSCRIPTION_FACTOR(0) = 5
# Effective: 16 SM × 1 CTA/SM × 5 = 80 concurrent tiles max.
# But Python loop overhead dominates, so we want FEWER, LARGER tiles.
#
# Strategy: target ~4-8 tiles per context phase.
# Fewer tiles → fewer matmul calls → less launch overhead.
# SMEM constraint: score tensor [kv_h, gqa, q_len, tile_sz] fp32
# must not cause OOM. With q_len=4096, kv_h=1, gqa=6:
# tile_sz=1024 → 1×6×4096×1024×4 = 96 MB (too much)
# tile_sz=512 → 48 MB (borderline)
# tile_sz=256 → 24 MB (safe)
# For decode (q_len=1): tile_sz=4096 → only 96 KB (always safe)
# ================================================================
_SMEM_BUDGET_BYTES = 256 * 1024 * 1024 # 256 MB score tensor budget
# CCCL GridEvenShare: fewer tiles = fewer iterations = less overhead
# BI-V100 has 32 GB HBM per card; 256 MB temporary is safe.
# 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]
@@ -586,6 +352,7 @@ class PagedAttention:
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)
@@ -601,36 +368,6 @@ class PagedAttention:
k_i = key [q_start:q_end] # [q_len, kv_h, d]
v_i = value[q_start:q_end]
# CCCL spread_out_items_per_thread adaptive tile sizing.
#
# Two constraints compete:
# 1. Memory: score tensor [kv_h, gqa, q_len, tile_sz] × 4 ≤ budget
# 2. Iteration count: want ~4-8 tiles to minimize Python overhead
#
# CCCL dispatch_transform.cuh::spread_out_items_per_thread:
# items = ceil_div(num_items, sm_count * threads * occupancy)
# items = clamp(items, min_items, max_items)
#
# Our translation: tile_sz = max context tokens / target_tiles,
# then clamp by memory budget.
score_row_bytes = num_kv_heads * gqa_ratio * q_len * 4
if score_row_bytes > 0:
mem_max_tokens = _SMEM_BUDGET_BYTES // score_row_bytes
mem_max_tokens = (mem_max_tokens // block_size) * block_size
else:
mem_max_tokens = block_size * 256
total_kv_tokens = ctx_len + q_len
# spread_out: target 4 tiles for context, 4 for current chunk
spread_tile = max(block_size,
(total_kv_tokens + 3) // 4)
# Round to block_size
spread_tile = (spread_tile // block_size) * block_size
spread_tile = max(spread_tile, block_size)
# Clamp by memory budget
tile_sz = min(spread_tile, mem_max_tokens)
tile_sz = max(tile_sz, block_size) # floor
# 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)
@@ -654,11 +391,14 @@ class PagedAttention:
# query has position ≥ ctx_len. k_pos < q_pos is always True
# → no causal mask needed for pure context tiles.
# --------------------------------------------------------------
# Convert token-based tile_sz to block count for iteration
blocks_per_tile = tile_sz // block_size
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} "
@@ -667,8 +407,8 @@ class PagedAttention:
"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)
for tile_blk in range(0, num_ctx_blocks, _BLOCKS_PER_TILE):
blk_end = min(tile_blk + _BLOCKS_PER_TILE, num_ctx_blocks)
blk_ids = block_tables[i, tile_blk:blk_end]
# Gather K/V for this tile.

78
qwen3_6_scripts/patch_model_runner.py Normal file
View File

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

57
qwen3_6_scripts/patch_ops.sh
View File

@@ -68,8 +68,22 @@ fi
# but the actual module file may be missing (causes ModuleNotFoundError
# on startup: "No module named 'vllm.model_executor.models.qwen3_5'").
# Deploy our qwen3_5.py so the module can be imported.
cp ./qwen3_5.py "$VLLM/model_executor/models/qwen3_5.py" 2>/dev/null && \
echo "[patch_ops] qwen3_5.py deployed (model module)" || true
# CCCL JIT pattern: check if image already has a working qwen3_5.py
# (Sub168's image had one with corex_gdn/corex_moe integration).
# Only deploy ours if the image's version is missing or broken.
_NATIVE_QW="$VLLM/model_executor/models/qwen3_5.py"
if [ -f "$_NATIVE_QW" ]; then
_SZ=$(wc -c < "$_NATIVE_QW" 2>/dev/null || echo 0)
if [ "$_SZ" -gt 1000 ]; then
echo "[patch_ops] qwen3_5.py EXISTS in image ($_SZ bytes) — NOT overwriting (corex native)"
else
cp ./qwen3_5.py "$_NATIVE_QW" 2>/dev/null && \
echo "[patch_ops] qwen3_5.py deployed (image version too small: $_SZ bytes)" || true
fi
else
cp ./qwen3_5.py "$VLLM/model_executor/models/qwen3_5.py" 2>/dev/null && \
echo "[patch_ops] qwen3_5.py deployed (not found in image)" || true
fi
# 2b. Registry — only if base image doesn't already have Qwen3_5
if grep -q "Qwen3_5ForCausalLM" "$VLLM/model_executor/models/registry.py" 2>/dev/null; then
@@ -79,10 +93,37 @@ else
echo "[patch_ops] registry.py deployed" || true
fi
# 2c. paged_attn.py — CRITICAL: Triton context_attention_fwd hangs BI-V100.
# Base engine comment: "The Triton context_attention_fwd kernel hangs BI-V100
# GPUs permanently. Our paged_attn.py bypasses it via _forward_prefix_pytorch."
cp ./paged_attn.py "$VLLM/attention/ops/paged_attn.py" 2>/dev/null && \
echo "[patch_ops] paged_attn.py deployed (Triton hang bypass)" || true
# 2d. patch_model_runner.py — fix prefix_cache_hit in chunked-prefill chunk 2+
python3 ./patch_model_runner.py 2>&1 || echo "[patch_ops] WARNING: model_runner patch failed (non-fatal)"
# 2e. mamba_cache.py — required for GatedDeltaNet state management
cp ./mamba_cache.py "$VLLM/model_executor/models/mamba_cache.py" 2>/dev/null && \
echo "[patch_ops] mamba_cache.py deployed" || true
# 2f. sequence.py — fix completion_tokens inflation under chunked prefill
cp ./sequence.py "$VLLM/sequence.py" 2>/dev/null && \
echo "[patch_ops] sequence.py deployed (token count fix)" || true
# 2g. scheduler.py — record num_cached_tokens in RequestMetrics
cp ./scheduler.py "$VLLM/core/scheduler.py" 2>/dev/null && \
echo "[patch_ops] scheduler.py deployed (cache metrics)" || true
# 2h. xformers — bypass cudnnFlashAttn (head_dim=256 > 128 limit)
python3 ./patch_xformers_sdpa_seq.py 2>&1 || echo "[patch_ops] WARNING: xformers seq patch failed"
python3 ./patch_xformers_sdpa_batch.py 2>&1 || echo "[patch_ops] WARNING: xformers batch patch failed"
echo "[patch_ops] xformers patches applied"
# 3. Tool parser
mkdir -p "$VLLM/entrypoints/openai/tool_parsers" 2>/dev/null || true
cp ./qwen3coder_tool_parser.py "$VLLM/entrypoints/openai/tool_parsers/" 2>/dev/null || true
cp ./tool_parsers_init.py "$VLLM/entrypoints/openai/tool_parsers/__init__.py" 2>/dev/null || true
python3 ./patch_vllm_tool_parser.py 2>&1 || echo "[patch_ops] WARNING: tool parser registry patch failed"
echo "[patch_ops] tool parser deployed"
# 4. Reasoning parser
@@ -108,7 +149,17 @@ for P in /usr/local/corex/lib/python3/dist-packages/vllm \
done
if [ -n "$VLLM2" ]; then
echo "[patch_ops] Second vllm at: $VLLM2"
cp ./qwen3_5.py "$VLLM2/model_executor/models/qwen3_5.py" 2>/dev/null || true
_NATIVE_QW2="$VLLM2/model_executor/models/qwen3_5.py"
if [ -f "$_NATIVE_QW2" ]; then
_SZ2=$(wc -c < "$_NATIVE_QW2" 2>/dev/null || echo 0)
if [ "$_SZ2" -gt 1000 ]; then
echo "[patch_ops] VLLM2 qwen3_5.py EXISTS ($_SZ2 bytes) — NOT overwriting"
else
cp ./qwen3_5.py "$_NATIVE_QW2" 2>/dev/null || true
fi
else
cp ./qwen3_5.py "$_NATIVE_QW2" 2>/dev/null || true
fi
if ! grep -q "Qwen3_5ForCausalLM" "$VLLM2/model_executor/models/registry.py" 2>/dev/null; then
cp ./registry.py "$VLLM2/model_executor/models/registry.py" 2>/dev/null || true
fi

79
qwen3_6_scripts/patch_vllm_tool_parser.py Normal file
View File

@@ -0,0 +1,79 @@
"""
Patches vLLM 0.6.3 to register Qwen3CoderToolParser under the name "qwen3_coder".
Deploy steps on the remote machine (already called by patch_ops.sh):
1. cp qwen3coder_tool_parser.py \
/usr/local/corex/lib/python3/dist-packages/vllm/entrypoints/openai/tool_parsers/
2. python3 patch_vllm_tool_parser.py
Usage after patching:
--tool-call-parser qwen3_coder --enable-auto-tool-choice
"""
import os
VLLM_ROOT = "/usr/local/corex/lib/python3/dist-packages/vllm"
TOOL_PARSERS_DIR = f"{VLLM_ROOT}/entrypoints/openai/tool_parsers"
INIT_FILE = f"{TOOL_PARSERS_DIR}/__init__.py"
def patch_file(path, replacements):
with open(path, "r") as f:
content = f.read()
patched = False
for old, new in replacements:
if new in content:
print(f" [skip] already patched: {repr(new[:70])}")
continue
if old not in content:
print(f" [warn] anchor not found: {repr(old[:70])}")
continue
content = content.replace(old, new, 1)
patched = True
print(f" [ok] patched: {repr(old[:50])} -> {repr(new[:50])}")
if patched:
with open(path, "w") as f:
f.write(content)
def main():
if not os.path.isdir(TOOL_PARSERS_DIR):
raise FileNotFoundError(
f"Tool parsers directory not found: {TOOL_PARSERS_DIR}\n"
"Verify the vLLM installation path.")
print(f"=== Patching {INIT_FILE} ===")
patch_file(INIT_FILE, [
(
"from .mistral_tool_parser import MistralToolParser",
"from .mistral_tool_parser import MistralToolParser\n"
"from .qwen3coder_tool_parser import Qwen3CoderToolParser",
),
(
'"MistralToolParser", "Internlm2ToolParser", "Llama3JsonToolParser"\n]',
'"MistralToolParser", "Internlm2ToolParser", "Llama3JsonToolParser",\n'
' "Qwen3CoderToolParser"\n]',
),
])
print("\n=== Verification ===")
try:
import importlib.util
spec = importlib.util.spec_from_file_location(
"qwen3coder_tool_parser",
f"{TOOL_PARSERS_DIR}/qwen3coder_tool_parser.py",
)
mod = importlib.util.module_from_spec(spec)
print(f" Module spec loaded: {spec.name}")
print(" (full import requires torch/vllm runtime — skipping exec)")
except Exception as e:
print(f" [warn] spec check failed: {e}")
print("\nDone. Start vLLM server with:")
print(" --tool-call-parser qwen3_coder --enable-auto-tool-choice")
if __name__ == "__main__":
main()

192
qwen3_6_scripts/patch_xformers_sdpa_batch.py Normal file
View File

@@ -0,0 +1,192 @@
"""
策略：批量（block-diagonal）fallback — 纯 PyTorch 数学实现
=============================================================
构建块对角 causal mask，对整批序列一次 matmul + softmax，
完全绕开所有硬件 flash attention kernel。
背景：
ixformer flshattF: head_dim > 128 报错拒绝
cudnnFlashAttnForward: 接受 head_dim=256，但数值结果错误（输出全"!"）
两者大概率是同一硬件单元，ixformer 提前拦截了硬件不支持的配置。
纯 matmul 路径完全绕开硬件 flash attention，数值正确。
优点：
数值正确。
并发请求 prefill attention 在 GPU 上真正并行（一次大 matmul）。
缺点：
峰值显存 = total_tokens² × H × dtype_size
total_tokens 受 --max-num-batched-tokens 控制，max-model-len 控制不住。
内存参考（fp16，H_local=6，--max-num-batched-tokens=T）：
T=2048 → 峰值 ~50 MB
T=4096 → 峰值 ~200 MB
T=8192 → 峰值 ~800 MB
T=16384 → 峰值 ~3.2 GB
Deploy:
python3 modified_scripts/patch_xformers_sdpa_batch.py
"""
XFORMERS_PATH = (
"/usr/local/corex/lib64/python3/dist-packages/"
"vllm/attention/backends/xformers.py"
)
FALLBACK_METHOD = '''
def _run_sdpa_fallback(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
attn_metadata: "XFormersMetadata",
) -> torch.Tensor:
"""批量纯数学 attention fallback。
构建块对角 causal mask（等价于 ixformer BlockDiagonalCausalMask），
对整批序列一次 matmul + softmax，GPU 并行处理所有序列。
块对角 mask 结构（seq1 len=3，seq2 len=2）：
s1,0 s1,1 s1,2 s2,0 s2,1
s1,0 [ 0 -inf -inf -inf -inf ]
s1,1 [ 0 0 -inf -inf -inf ]
s1,2 [ 0 0 0 -inf -inf ]
s2,0 [-inf -inf -inf 0 -inf ]
s2,1 [-inf -inf -inf 0 0 ]
softmax 在 float32 下计算防止 float16 溢出，结果转回原始 dtype。
Args:
query : [1, total_prefill_tokens, num_heads, head_dim]
key : [1, total_prefill_tokens, num_kv_heads, head_dim]
value : [1, total_prefill_tokens, num_kv_heads, head_dim]
Returns:
[1, total_prefill_tokens, num_heads, head_dim]
"""
assert attn_metadata.seq_lens is not None
orig_dtype = query.dtype
total_tokens = query.shape[1]
# ── 构建块对角 causal mask [T, T] ────────────────────────────────
# 全部初始化为 -inf，再对每条序列的对角块填入下三角 0
mask = torch.full(
(total_tokens, total_tokens),
float("-inf"),
dtype=torch.float32,
device=query.device,
)
start = 0
for seq_len in attn_metadata.seq_lens:
end = start + seq_len
mask[start:end, start:end] = torch.tril(
torch.zeros(seq_len, seq_len,
dtype=torch.float32, device=query.device)
)
start = end
# ── [1, H, T, D]，.contiguous() ──────────────────────────────────
q_all = query.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0)
k_all = key.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0)
v_all = value.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0)
# ── GQA：展开 KV heads ────────────────────────────────────────────
if k_all.shape[1] != q_all.shape[1]:
n = q_all.shape[1] // k_all.shape[1]
k_all = k_all.repeat_interleave(n, dim=1).contiguous()
v_all = v_all.repeat_interleave(n, dim=1).contiguous()
# ── 纯数学 attention（float32 防溢出）────────────────────────────
# [1, H, T, T]
attn_w = torch.matmul(q_all.float(), k_all.float().transpose(-2, -1))
attn_w = attn_w * self.scale
attn_w = attn_w + mask # 加法广播：mask [T,T] → [1, H, T, T]
attn_w = torch.softmax(attn_w, dim=-1)
out = torch.matmul(attn_w, v_all.float()).to(orig_dtype)
# [1, H, T, D] → [1, T, H, D]
return out.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0)
'''
OLD_XFORMER_BLOCK = """\
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)
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)\
"""
NEW_XFORMER_BLOCK = """\
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 > 128:
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)\
"""
INJECT_ANCHOR = " def _run_memory_efficient_xformers_forward("
def patch_file(path):
with open(path, "r") as f:
content = f.read()
changed = False
if "_run_sdpa_fallback" in content:
print(" [skip] _run_sdpa_fallback already present")
elif INJECT_ANCHOR not in content:
print(" [warn] inject anchor not found")
else:
content = content.replace(INJECT_ANCHOR, FALLBACK_METHOD + INJECT_ANCHOR, 1)
print(" [ok] injected _run_sdpa_fallback (batch, pure-math)")
changed = True
if NEW_XFORMER_BLOCK in content:
print(" [skip] dispatch block already patched")
elif OLD_XFORMER_BLOCK in content:
content = content.replace(OLD_XFORMER_BLOCK, NEW_XFORMER_BLOCK, 1)
print(" [ok] patched dispatch block")
changed = True
else:
print(" [warn] dispatch block anchor not found")
if changed:
with open(path, "w") as f:
f.write(content)
print(f" Written: {path}")
def main():
print("=== patch_xformers_sdpa_batch (batch, pure-math) ===")
print(f"Target: {XFORMERS_PATH}")
patch_file(XFORMERS_PATH)
print("\nDone.")
if __name__ == "__main__":
main()

191
qwen3_6_scripts/patch_xformers_sdpa_batch_kernel.py Normal file
View File

@@ -0,0 +1,191 @@
"""
策略：批量（block-diagonal）— F.scaled_dot_product_attention，可走硬件 kernel
=============================================================================
构建块对角 causal mask，对整批序列一次 F.scaled_dot_product_attention。
与 patch_xformers_sdpa_batch.py（纯 matmul）的区别：
SDPA 会根据 PyTorch/驱动能力分发到最优 kernel（Flash Attention /
mem-efficient attention / math fallback），而不是固定走 cublas matmul。
历史说明：
该方案最早因输出全"!"而被弃用，后续排查确认"!"由 mamba_cache.py bug
引起，与 attention 实现无关。当前恢复此方案用于性能对比测试。
已知硬件限制（BI-V100）：
cudnnFlashAttnForward 不支持 is_causal=True（报错）。
本实现使用 is_causal=False + 显式块对角 additive mask 规避此限制。
若 SDPA 仍分发到有问题的 kernel，回退到 patch_xformers_sdpa_batch.py。
优点（vs 纯 matmul）：
SDPA 可分发到 Flash Attention kernel → O(L) 显存、更快的 CUDA kernel。
缺点：
依赖硬件 kernel 行为，若 kernel 有 bug 则数值错误（需与 matmul 版对比验证）。
Deploy:
python3 modified_scripts/patch_xformers_sdpa_batch_kernel.py
"""
XFORMERS_PATH = (
"/usr/local/corex/lib64/python3/dist-packages/"
"vllm/attention/backends/xformers.py"
)
FALLBACK_METHOD = '''
def _run_sdpa_fallback(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
attn_metadata: "XFormersMetadata",
) -> torch.Tensor:
"""批量 F.scaled_dot_product_attention fallback（可走硬件 kernel）。
构建块对角 causal mask，对整批序列一次 SDPA 调用。
SDPA 可分发到 Flash Attention / mem-efficient attention kernel。
is_causal=False + 显式 additive mask，规避 cudnnFlashAttnForward
不支持 is_causal=True 的限制。
块对角 mask（seq1 len=3，seq2 len=2）：
s1,0 s1,1 s1,2 s2,0 s2,1
s1,0 [ 0 -inf -inf -inf -inf ]
s1,1 [ 0 0 -inf -inf -inf ]
s1,2 [ 0 0 0 -inf -inf ]
s2,0 [-inf -inf -inf 0 -inf ]
s2,1 [-inf -inf -inf 0 0 ]
Args:
query : [1, total_prefill_tokens, num_heads, head_dim]
key : [1, total_prefill_tokens, num_kv_heads, head_dim]
value : [1, total_prefill_tokens, num_kv_heads, head_dim]
Returns:
[1, total_prefill_tokens, num_heads, head_dim]
"""
import torch.nn.functional as F
assert attn_metadata.seq_lens is not None
orig_dtype = query.dtype
total_tokens = query.shape[1]
# ── 块对角 causal mask [T, T] ─────────────────────────────────────
mask = torch.full(
(total_tokens, total_tokens),
float("-inf"),
dtype=orig_dtype,
device=query.device,
)
start = 0
for seq_len in attn_metadata.seq_lens:
end = start + seq_len
mask[start:end, start:end] = torch.tril(
torch.zeros(seq_len, seq_len, dtype=orig_dtype, device=query.device)
)
start = end
# ── [1, H, T, D] ──────────────────────────────────────────────────
q_all = query.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0)
k_all = key.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0)
v_all = value.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0)
# ── GQA：展开 KV heads ────────────────────────────────────────────
if k_all.shape[1] != q_all.shape[1]:
n = q_all.shape[1] // k_all.shape[1]
k_all = k_all.repeat_interleave(n, dim=1).contiguous()
v_all = v_all.repeat_interleave(n, dim=1).contiguous()
# ── F.scaled_dot_product_attention（可走硬件 kernel）─────────────
# is_causal=False：避免 cudnnFlashAttnForward "not support causal mode"
# attn_mask 传 additive float mask（非 bool），SDPA 选择 math/kernel 路径
out = F.scaled_dot_product_attention(
q_all, k_all, v_all,
attn_mask=mask,
dropout_p=0.0,
is_causal=False,
scale=self.scale,
)
# [1, H, T, D] → [1, T, H, D]
return out.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0)
'''
OLD_XFORMER_BLOCK = """\
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)
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)\
"""
NEW_XFORMER_BLOCK = """\
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 > 128:
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)\
"""
INJECT_ANCHOR = " def _run_memory_efficient_xformers_forward("
def patch_file(path):
with open(path, "r") as f:
content = f.read()
changed = False
if "_run_sdpa_fallback" in content:
print(" [skip] _run_sdpa_fallback already present")
elif INJECT_ANCHOR not in content:
print(" [warn] inject anchor not found")
else:
content = content.replace(INJECT_ANCHOR, FALLBACK_METHOD + INJECT_ANCHOR, 1)
print(" [ok] injected _run_sdpa_fallback (batch, F.sdpa kernel)")
changed = True
if NEW_XFORMER_BLOCK in content:
print(" [skip] dispatch block already patched")
elif OLD_XFORMER_BLOCK in content:
content = content.replace(OLD_XFORMER_BLOCK, NEW_XFORMER_BLOCK, 1)
print(" [ok] patched dispatch block")
changed = True
else:
print(" [warn] dispatch block anchor not found")
if changed:
with open(path, "w") as f:
f.write(content)
print(f" Written: {path}")
def main():
print("=== patch_xformers_sdpa_batch_kernel (batch, F.sdpa + kernel dispatch) ===")
print(f"Target: {XFORMERS_PATH}")
patch_file(XFORMERS_PATH)
print("\nDone.")
if __name__ == "__main__":
main()

321
qwen3_6_scripts/patch_xformers_sdpa_seq.py Normal file
View File

@@ -0,0 +1,321 @@
"""
策略：顺序（per-sequence）fallback — 纯 PyTorch 数学实现
==========================================================
逐条序列用 matmul + softmax 手写 attention，完全绕开所有硬件
flash attention kernel（ixformer / cudnnFlashAttnForward）。
背景：
Iluvatar cudnnFlashAttnForward 存在两个已知问题：
1. 不支持 is_causal=True（报错）
2. 使用 attn_mask 路径时数值结果不正确（静默错误，输出全为"!"）
与华为昇腾 910B4 上 llama.cpp --flash-attn off 修复同类问题的原理相同。
纯数学路径（matmul + softmax）在任何 PyTorch 后端上结果都正确。
优点：
数值正确，不依赖任何硬件特定 attention kernel。
峰值显存 = max(seq_len)² × H × dtype_size，由 --max-model-len 控制。
缺点：
并发请求的 prefill attention 串行执行。
O(L²) 显存（无 flash attention 的 O(L) 优化）。
内存参考（fp16，H_local=6）：
max-model-len=4096 → 峰值 ~200 MB
max-model-len=8192 → 峰值 ~800 MB
max-model-len=16384 → 峰值 ~3.2 GB
额外 patch（arg_utils.py）：
vllm 0.6.3 在 max_model_len > 32K 时会自动开启 chunked prefill（无命令行
关闭选项），原意是防止 profiling OOM。但 _run_sdpa_fallback 已通过 Q-tiling
解决了该问题，chunked prefill 反而会把推理路径从 _run_sdpa_fallback 切换到
_forward_prefix_pytorch，属于不必要的行为变更，因此一并禁用该自动逻辑。
Deploy:
python3 modified_scripts/patch_xformers_sdpa_seq.py
"""
XFORMERS_PATH = (
"/usr/local/corex/lib64/python3/dist-packages/"
"vllm/attention/backends/xformers.py"
)
ARG_UTILS_PATH = (
"/usr/local/corex/lib64/python3/dist-packages/"
"vllm/engine/arg_utils.py"
)
LOGITS_PROC_PATH = (
"/usr/local/corex/lib64/python3/dist-packages/"
"vllm/model_executor/layers/logits_processor.py"
)
# _apply_logits_processors crashes when seq_groups is None (intermediate
# chunked-prefill chunks on the driver rank). Add an early-return guard.
_LP_OLD_BLOCK = """\
def _apply_logits_processors(
logits: torch.Tensor,
sampling_metadata: SamplingMetadata,
) -> torch.Tensor:
found_logits_processors = False\
"""
_LP_NEW_BLOCK = """\
def _apply_logits_processors(
logits: torch.Tensor,
sampling_metadata: SamplingMetadata,
) -> torch.Tensor:
if sampling_metadata.seq_groups is None: # intermediate chunked-prefill chunk
return logits
found_logits_processors = False\
"""
# vllm 0.6.3 自动开启 chunked prefill 的原始块
_ARG_OLD_BLOCK = """\
if (is_gpu and not use_sliding_window and not use_spec_decode
and not self.enable_lora
and not self.enable_prompt_adapter):
self.enable_chunked_prefill = True
logger.warning(
"Chunked prefill is enabled by default for models with "
"max_model_len > 32K. Currently, chunked prefill might "
"not work with some features or models. If you "
"encounter any issues, please disable chunked prefill "
"by setting --enable-chunked-prefill=False.")\
"""
_ARG_NEW_BLOCK = """\
if (is_gpu and not use_sliding_window and not use_spec_decode
and not self.enable_lora
and not self.enable_prompt_adapter):
pass # skip auto-enable: Q-tiling in _run_sdpa_fallback
# handles long-context memory without chunked prefill\
"""
FALLBACK_METHOD = '''
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]
'''
OLD_XFORMER_BLOCK = """\
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)
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)\
"""
NEW_XFORMER_BLOCK = """\
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 > 128:
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)\
"""
INJECT_ANCHOR = " def _run_memory_efficient_xformers_forward("
def patch_file(path):
with open(path, "r") as f:
content = f.read()
changed = False
if "_run_sdpa_fallback" in content:
print(" [skip] _run_sdpa_fallback already present")
elif INJECT_ANCHOR not in content:
print(" [warn] inject anchor not found")
else:
content = content.replace(INJECT_ANCHOR, FALLBACK_METHOD + INJECT_ANCHOR, 1)
print(" [ok] injected _run_sdpa_fallback (sequential, pure-math)")
changed = True
if NEW_XFORMER_BLOCK in content:
print(" [skip] dispatch block already patched")
elif OLD_XFORMER_BLOCK in content:
content = content.replace(OLD_XFORMER_BLOCK, NEW_XFORMER_BLOCK, 1)
print(" [ok] patched dispatch block")
changed = True
else:
print(" [warn] dispatch block anchor not found")
if changed:
with open(path, "w") as f:
f.write(content)
print(f" Written: {path}")
def patch_arg_utils(path):
with open(path, "r") as f:
content = f.read()
changed = False
if "skip auto-enable: Q-tiling" in content:
print(" [skip] chunked-prefill auto-enable already disabled")
elif _ARG_OLD_BLOCK in content:
content = content.replace(_ARG_OLD_BLOCK, _ARG_NEW_BLOCK, 1)
print(" [ok] disabled chunked-prefill auto-enable for 32K+")
changed = True
else:
print(" [warn] target block not found — check arg_utils.py version")
if changed:
with open(path, "w") as f:
f.write(content)
print(f" Written: {path}")
def patch_logits_processor(path):
with open(path, "r") as f:
content = f.read()
changed = False
if "intermediate chunked-prefill chunk" in content:
print(" [skip] seq_groups=None guard already present")
elif _LP_OLD_BLOCK in content:
content = content.replace(_LP_OLD_BLOCK, _LP_NEW_BLOCK, 1)
print(" [ok] added seq_groups=None guard in _apply_logits_processors")
changed = True
else:
print(" [warn] target block not found — check logits_processor.py version")
if changed:
with open(path, "w") as f:
f.write(content)
print(f" Written: {path}")
def main():
print("=== patch_xformers_sdpa_seq (sequential, pure-math) ===")
print(f"Target: {XFORMERS_PATH}")
patch_file(XFORMERS_PATH)
print("\n=== patch_arg_utils (disable chunked-prefill auto-enable) ===")
print(f"Target: {ARG_UTILS_PATH}")
patch_arg_utils(ARG_UTILS_PATH)
print("\n=== patch_logits_processor (seq_groups=None guard for chunked prefill) ===")
print(f"Target: {LOGITS_PROC_PATH}")
patch_logits_processor(LOGITS_PROC_PATH)
print("\nDone.")
if __name__ == "__main__":
main()

181
qwen3_6_scripts/patch_xformers_sdpa_seq_kernel.py Normal file
View File

@@ -0,0 +1,181 @@
"""
策略：顺序（per-sequence）— F.scaled_dot_product_attention，可走硬件 kernel
=============================================================================
逐条序列调用 F.scaled_dot_product_attention，is_causal=False + 显式因果 mask。
与 patch_xformers_sdpa_seq.py（纯 matmul）的区别：
SDPA 可分发到 Flash Attention / mem-efficient attention kernel，
而纯 matmul 固定走 cublas。
硬件限制（BI-V100）：
cudnnFlashAttnForward 不支持 is_causal=True（直接报错）。
必须使用 is_causal=False + 显式 additive causal mask。
每条序列单独构造上三角 -inf mask，peak 显存 = max(seq_len)² × dtype，
比 batch 版的 total_tokens² 小得多。
与 batch_kernel 的对比：
seq_kernel: 显存小，peak = max_single_seq²；并发 prefill 串行排队
batch_kernel: 显存大，peak = total_tokens²；并发 prefill 一次并行处理，
通过 --max-num-batched-tokens 控制 total_tokens 上限
Deploy:
python3 modified_scripts/patch_xformers_sdpa_seq_kernel.py
"""
XFORMERS_PATH = (
"/usr/local/corex/lib64/python3/dist-packages/"
"vllm/attention/backends/xformers.py"
)
FALLBACK_METHOD = '''
def _run_sdpa_fallback(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
attn_metadata: "XFormersMetadata",
) -> torch.Tensor:
"""顺序 F.scaled_dot_product_attention fallback（可走硬件 kernel）。
逐条序列调用 SDPA，is_causal=False + 显式上三角 additive mask。
cudnnFlashAttnForward 不支持 is_causal=True，必须用显式 mask。
逐序列构造 mask，peak 显存 = max(seq_len)² × dtype（远小于 batch 版）。
Args:
query : [1, total_prefill_tokens, num_heads, head_dim]
key : [1, total_prefill_tokens, num_kv_heads, head_dim]
value : [1, total_prefill_tokens, num_kv_heads, head_dim]
Returns:
[1, total_prefill_tokens, num_heads, head_dim]
"""
import torch.nn.functional as F
assert attn_metadata.seq_lens is not None
orig_dtype = query.dtype
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)
start = 0
for seq_len in attn_metadata.seq_lens:
end = start + seq_len
# [1, H, L, D]
q_s = q_flat[start:end].permute(1, 0, 2).contiguous().unsqueeze(0)
k_s = k_flat[start:end].permute(1, 0, 2).contiguous().unsqueeze(0)
v_s = v_flat[start:end].permute(1, 0, 2).contiguous().unsqueeze(0)
# GQA：展开 KV heads
if k_s.shape[1] != q_s.shape[1]:
n = q_s.shape[1] // k_s.shape[1]
k_s = k_s.repeat_interleave(n, dim=1).contiguous()
v_s = v_s.repeat_interleave(n, dim=1).contiguous()
# 逐序列因果 mask [L, L]，上三角 -inf
causal_mask = torch.tril(
torch.zeros(seq_len, seq_len, dtype=orig_dtype, device=q_s.device)
)
causal_mask = causal_mask.masked_fill(
torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool,
device=q_s.device), diagonal=1),
float("-inf"),
)
# is_causal=False + 显式 mask，规避 cudnnFlashAttnForward 不支持 is_causal=True
out_s = F.scaled_dot_product_attention(
q_s, k_s, v_s,
attn_mask=causal_mask,
dropout_p=0.0,
is_causal=False,
scale=self.scale,
)
# [1, H, L, D] → [L, H, D]
output[start:end] = out_s.squeeze(0).permute(1, 0, 2).to(orig_dtype)
start = end
return output.unsqueeze(0) # [1, T, H, D]
'''
OLD_XFORMER_BLOCK = """\
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)
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)\
"""
NEW_XFORMER_BLOCK = """\
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 > 128:
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)\
"""
INJECT_ANCHOR = " def _run_memory_efficient_xformers_forward("
def patch_file(path):
with open(path, "r") as f:
content = f.read()
changed = False
if "_run_sdpa_fallback" in content:
print(" [skip] _run_sdpa_fallback already present")
elif INJECT_ANCHOR not in content:
print(" [warn] inject anchor not found")
else:
content = content.replace(INJECT_ANCHOR, FALLBACK_METHOD + INJECT_ANCHOR, 1)
print(" [ok] injected _run_sdpa_fallback (seq, F.sdpa kernel)")
changed = True
if NEW_XFORMER_BLOCK in content:
print(" [skip] dispatch block already patched")
elif OLD_XFORMER_BLOCK in content:
content = content.replace(OLD_XFORMER_BLOCK, NEW_XFORMER_BLOCK, 1)
print(" [ok] patched dispatch block")
changed = True
else:
print(" [warn] dispatch block anchor not found")
if changed:
with open(path, "w") as f:
f.write(content)
print(f" Written: {path}")
def main():
print("=== patch_xformers_sdpa_seq_kernel (seq, F.sdpa + kernel dispatch) ===")
print(f"Target: {XFORMERS_PATH}")
patch_file(XFORMERS_PATH)
print("\nDone.")
if __name__ == "__main__":
main()