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