From c5a0d6185164e2698a50840cc09ddd7da8219b33 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 31 Jul 2026 09:43:58 +0000 Subject: [PATCH] sync: align with enginex-vllm-bi100-qwen36 baseline (1902c81f) Synced files from EngineX baseline zip (2026-06-30): - ADD paged_attn.py (root): production paged attention with PyTorch fallback - ADD launch_service: BI-V100 server startup script with env configuration - SYNC computility-run.yaml: gpu_memory=0.9, batched_tokens=8192, seq_capture=32768 - SYNC qwen3_6_scripts/paged_attn.py: +311 lines, Triton bypass docs, _forward_decode_pytorch shape docs - SYNC qwen3_6_scripts/qwen3_5.py: -72 lines, revert optimized MoE prefill to baseline (untested on BI-V100) - KEEP Dockerfile: repo version has V2/Triton/head256 optimization patches not in baseline Baseline commit: 1902c81fdd373943f17f5983eb8750758c7f4a69 Source: enginex-vllm-bi100-qwen36-main.zip (dev.modelhub.org.cn) --- computility-run.yaml | 6 +- launch_service | 80 +++++++++ paged_attn.py | 243 +++++++++++++++++++++++++++ qwen3_6_scripts/paged_attn.py | 303 ++++++++++++++++++++++------------ qwen3_6_scripts/qwen3_5.py | 72 ++------ 5 files changed, 535 insertions(+), 169 deletions(-) create mode 100755 launch_service create mode 100644 paged_attn.py diff --git a/computility-run.yaml b/computility-run.yaml index 286cc10c..6447abd0 100644 --- a/computility-run.yaml +++ b/computility-run.yaml @@ -10,7 +10,7 @@ command: - --max-model-len - '100000' - --gpu-memory-utilization - - '0.95' + - '0.9' - --trust-remote-code - -tp - '4' @@ -19,10 +19,10 @@ command: - --disable-log-requests - --disable-frontend-multiprocessing - --max-num-batched-tokens - - '16384' + - '8192' - --enable-chunked-prefill - --max-seq-len-to-capture - - '65536' + - '32768' - --enable-auto-tool-choice - --tool-call-parser - qwen3_coder diff --git a/launch_service b/launch_service new file mode 100755 index 00000000..0086b855 --- /dev/null +++ b/launch_service @@ -0,0 +1,80 @@ +#!/bin/bash + +export PYTHONPATH=/usr/local/corex/lib64/python3/dist-packages +export LD_LIBRARY_PATH=/usr/local/corex/lib64:/usr/local/openmpi/lib +export PATH=/usr/local/corex/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/local/corex/lib64/python3/dist-packages/bin:/usr/local/openmpi/bin +export JAVA_HOME=/root/apps/jdk1.8.0_411 +export JRE_HOME=/root/apps/jdk1.8.0_411/jre +export JMETER_HOME=/root/apps/apache-jmeter-5.6.3 +export CLASSPATH=.:/root/apps/jdk1.8.0_411/lib/dt.jar:/root/apps/jdk1.8.0_411/lib/tools.jar:/root/apps/apache-jmeter-5.6.3/lib/ext/ApacheJMeter_core.jar:/root/apps/apache-jmeter-5.6.3/lib/jorphan.jar:/root/apps/apache-jmeter-5.6.3/lib/logkit-2.0.jar: +export PATH=/root/apps/apache-jmeter-5.6.3/bin:/root/apps/jdk1.8.0_411/bin:/usr/local/corex/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/local/corex/lib64/python3/dist-packages/bin:/usr/local/openmpi/bin +/iluvatar/welcome.sh + +data +cat /proc/cpuinfo | tail -n 50 +ixsmi +unset CUDA_VISIBLE_DEVICES +export +date + +DEFAULT_HOST="0.0.0.0" +DEFAULT_PORT="80" +DEFAULT_SERVED_MODEL_NAME="llm" +DEFAULT_MODEL_PATH="/model" +DEFAULT_MAX_MODEL_LEN="10000" +DEFAULT_TENSOR_PARALLEL_SIZE="1" +DEFAULT_MAX_NUM_SEQS="64" +DEFAULT_ENFORCE_EAGER="true" +DEFAULT_DISABLE_LOG_REQUESTS="true" +DEFAULT_PREFIX_CACHING="true" + +HOST_VAL=${HOST:-$DEFAULT_HOST} +PORT_VAL=${PORT:-$DEFAULT_PORT} +SERVED_MODEL_NAME_VAL=${SERVED_MODEL_NAME:-$DEFAULT_SERVED_MODEL_NAME} +MODEL_PATH_VAL=${MODEL_PATH:-$DEFAULT_MODEL_PATH} +MAX_MODEL_LEN_VAL=${MAX_MODEL_LEN:-$DEFAULT_MAX_MODEL_LEN} +TENSOR_PARALLEL_SIZE_VAL=${TENSOR_PARALLEL_SIZE:-$DEFAULT_TENSOR_PARALLEL_SIZE} +MAX_NUM_SEQS_VAL=${MAX_NUM_SEQS:-$DEFAULT_MAX_NUM_SEQS} +INCLUDE_ENFORCE_EAGER_FLAG=${ENFORCE_EAGER:-$DEFAULT_ENFORCE_EAGER} +INCLUDE_DISABLE_LOG_REQUESTS_FLAG=${DISABLE_LOG_REQUESTS:-$DEFAULT_DISABLE_LOG_REQUESTS} +INCLUDE_PREFIX_CACHING_FLAG=${PREFIX_CACHING:-$DEFAULT_PREFIX_CACHING} + +CMD_ARGS=() +CMD_ARGS+=(--host "$HOST_VAL") +CMD_ARGS+=(--port "$PORT_VAL") + +if [[ "$INCLUDE_ENFORCE_EAGER_FLAG" != "false" && "$INCLUDE_ENFORCE_EAGER_FLAG" != "0" ]]; then + CMD_ARGS+=(--enforce-eager) +fi +if [[ "$INCLUDE_DISABLE_LOG_REQUESTS_FLAG" != "false" && "$INCLUDE_DISABLE_LOG_REQUESTS_FLAG" != "0" ]]; then + CMD_ARGS+=(--disable-log-requests) +fi +if [[ "$INCLUDE_PREFIX_CACHING_FLAG" != "false" && "$INCLUDE_PREFIX_CACHING_FLAG" != "0" ]]; then + CMD_ARGS+=(--enable-prefix-caching) +fi + +CMD_ARGS+=(--served-model-name "$SERVED_MODEL_NAME_VAL") +CMD_ARGS+=(--model "$MODEL_PATH_VAL") +CMD_ARGS+=(--max-model-len "$MAX_MODEL_LEN_VAL") +CMD_ARGS+=(--tensor-parallel-size "$TENSOR_PARALLEL_SIZE_VAL") +CMD_ARGS+=(--max-num-seqs "$MAX_NUM_SEQS_VAL") +CMD_ARGS+=(--trust-remote-code) + +echo "--------------------------------------------------" +echo "Starting VLLM OpenAI API Server..." +echo "Using effective arguments:" +echo " Host (--host): $HOST_VAL" +echo " Port (--port): $PORT_VAL" +echo " Enforce Eager (--enforce-eager):" $([[ "$INCLUDE_ENFORCE_EAGER_FLAG" != "false" && "$INCLUDE_ENFORCE_EAGER_FLAG" != "0" ]] && echo "Enabled" || echo "Disabled (Env: ENFORCE_EAGER=$ENFORCE_EAGER)") +echo " Disable Log Req (--disable-log-requests):" $([[ "$INCLUDE_DISABLE_LOG_REQUESTS_FLAG" != "false" && "$INCLUDE_DISABLE_LOG_REQUESTS_FLAG" != "0" ]] && echo "Enabled" || echo "Disabled (Env: DISABLE_LOG_REQUESTS=$DISABLE_LOG_REQUESTS)") +echo " Served Model Name (--served-model-name): $SERVED_MODEL_NAME_VAL" +echo " Model Path (--model): $MODEL_PATH_VAL" +echo " Max Model Length (--max-model-len): $MAX_MODEL_LEN_VAL" +echo " Tensor Parallel Size (--tensor-parallel-size): $TENSOR_PARALLEL_SIZE_VAL" +echo " Max Num Seqs (--max-num-seqs): $MAX_NUM_SEQS_VAL" +echo "--------------------------------------------------" +echo "Full cmd:" +echo "python3 -m vllm.entrypoints.openai.api_server ${CMD_ARGS[*]}" +echo "--------------------------------------------------" + +python3 -m vllm.entrypoints.openai.api_server "${CMD_ARGS[@]}" diff --git a/paged_attn.py b/paged_attn.py new file mode 100644 index 00000000..988f9032 --- /dev/null +++ b/paged_attn.py @@ -0,0 +1,243 @@ +from dataclasses import dataclass +from typing import List, Optional, Tuple + +import torch + +from vllm import _custom_ops as ops + +from vllm.attention.ops.prefix_prefill import context_attention_fwd + +# Should be the same as PARTITION_SIZE in `paged_attention_v2_launcher`. +_PARTITION_SIZE = 512 + + +@dataclass +class PagedAttentionMetadata: + """Metadata for PagedAttention.""" + # (batch_size,). The length of sequences (entire tokens seen so far) per + # sequence. + seq_lens_tensor: Optional[torch.Tensor] + # Maximum sequence length in the batch. 0 if it is prefill-only batch. + max_decode_seq_len: int + # (batch_size, max_blocks_per_seq). + # Block addresses per sequence. (Seq id -> list of physical block) + # E.g., [0, 1, 2] means tokens are stored in 0th, 1st, and 2nd blocks + # in the kv cache. Each block can contain up to block_size tokens. + # 2nd dimensions are padded up to max_blocks_per_seq if it is cuda-graph + # captured. + block_tables: Optional[torch.Tensor] + + +class PagedAttention: + + @staticmethod + def get_supported_head_sizes() -> List[int]: + return [64, 80, 96, 112, 120, 128, 192, 256] + + @staticmethod + def get_kv_cache_shape( + num_blocks: int, + block_size: int, + num_kv_heads: int, + head_size: int, + ) -> Tuple[int, ...]: + return (2, num_blocks, block_size * num_kv_heads * head_size) + + @staticmethod + def split_kv_cache( + kv_cache: torch.Tensor, + num_kv_heads: int, + head_size: int, + ) -> Tuple[torch.Tensor, torch.Tensor]: + x = 16 // kv_cache.element_size() + num_blocks = kv_cache.shape[1] + + key_cache = kv_cache[0] + key_cache = key_cache.view(num_blocks, num_kv_heads, head_size // x, + -1, x) + value_cache = kv_cache[1] + value_cache = value_cache.view(num_blocks, num_kv_heads, head_size, -1) + return key_cache, value_cache + + @staticmethod + def write_to_paged_cache( + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + slot_mapping: torch.Tensor, + kv_cache_dtype: str, + k_scale: float, + v_scale: float, + ) -> None: + ops.reshape_and_cache( + key, + value, + key_cache, + value_cache, + slot_mapping.flatten(), + kv_cache_dtype, + k_scale, + v_scale, + ) + + @staticmethod + def forward_decode( + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + max_seq_len: int, + kv_cache_dtype: str, + num_kv_heads: int, + scale: float, + alibi_slopes: Optional[torch.Tensor], + k_scale: float, + v_scale: float, + tp_rank: int = 0, + blocksparse_local_blocks: int = 0, + blocksparse_vert_stride: int = 0, + blocksparse_block_size: int = 64, + blocksparse_head_sliding_step: int = 0, + ) -> torch.Tensor: + if blocksparse_vert_stride is not None and blocksparse_vert_stride > 1: + # use blocksparse paged attention + block_size = value_cache.size(-1) + assert (blocksparse_block_size > 0 and + blocksparse_block_size % block_size == 0), \ + (f"{blocksparse_block_size=} needs to be a multiple of" + f"{block_size=} used in block_tables.") + + output = torch.empty_like(query) + block_size = value_cache.shape[3] + num_seqs, num_heads, head_size = query.shape + max_num_partitions = ((max_seq_len + _PARTITION_SIZE - 1) // + _PARTITION_SIZE) + # NOTE(woosuk): We use a simple heuristic to decide whether to use + # PagedAttention V1 or V2. If the number of partitions is 1, we use + # V1 to avoid the overhead of reduction. Also, if the number of + # sequences or heads is large, we use V1 since there is enough work + # to parallelize. + # TODO(woosuk): Tune this heuristic. + # For context len > 8192, use V2 kernel to avoid shared memory shortage. + use_v1 = (max_seq_len <= 8192 + and (max_num_partitions == 1 or num_seqs * num_heads > 512)) + use_v1 = True + if use_v1: + # Run PagedAttention V1. + ops.paged_attention_v1( + output, + query, + key_cache, + value_cache, + num_kv_heads, + scale, + block_tables, + seq_lens, + block_size, + max_seq_len, + alibi_slopes, + ) + else: + # Run PagedAttention V2. + assert _PARTITION_SIZE % block_size == 0 + tmp_output = torch.empty( + size=(num_seqs, num_heads, max_num_partitions, head_size), + dtype=output.dtype, + device=output.device, + ) + exp_sums = torch.empty( + size=(num_seqs, num_heads, max_num_partitions), + dtype=torch.float32, + device=output.device, + ) + max_logits = torch.empty_like(exp_sums) + ops.paged_attention_v2( + output, + exp_sums, + max_logits, + tmp_output, + query, + key_cache, + value_cache, + num_kv_heads, + scale, + block_tables, + seq_lens, + block_size, + max_seq_len, + alibi_slopes, + kv_cache_dtype, + k_scale, + v_scale, + tp_rank, + blocksparse_local_blocks, + blocksparse_vert_stride, + blocksparse_block_size, + blocksparse_head_sliding_step, + ) + return output + + @staticmethod + def forward_prefix( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + kv_cache_dtype: str, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + block_tables: torch.Tensor, + query_start_loc: torch.Tensor, + seq_lens_tensor: torch.Tensor, + context_lens: torch.Tensor, + max_query_len: int, + alibi_slopes: Optional[torch.Tensor], + sliding_window: Optional[int], + k_scale: float, + v_scale: float, + ) -> torch.Tensor: + output = torch.empty_like(query) + context_attention_fwd( + query, + key, + value, + output, + kv_cache_dtype, + key_cache, + value_cache, + block_tables, + # query_start_loc is (batch_size + 1,) + query_start_loc[:-1], + seq_lens_tensor, + context_lens, + max_query_len, + k_scale, + v_scale, + alibi_slopes, + sliding_window, + ) + return output + + @staticmethod + def swap_blocks( + src_kv_cache: torch.Tensor, + dst_kv_cache: torch.Tensor, + src_to_dst: torch.Tensor, + ) -> None: + src_key_cache = src_kv_cache[0] + dst_key_cache = dst_kv_cache[0] + ops.swap_blocks(src_key_cache, dst_key_cache, src_to_dst) + + src_value_cache = src_kv_cache[1] + dst_value_cache = dst_kv_cache[1] + ops.swap_blocks(src_value_cache, dst_value_cache, src_to_dst) + + @staticmethod + def copy_blocks( + kv_caches: List[torch.Tensor], + src_to_dists: torch.Tensor, + ) -> None: + key_caches = [kv_cache[0] for kv_cache in kv_caches] + value_caches = [kv_cache[1] for kv_cache in kv_caches] + ops.copy_blocks(key_caches, value_caches, src_to_dists) diff --git a/qwen3_6_scripts/paged_attn.py b/qwen3_6_scripts/paged_attn.py index 43afef4a..85904895 100644 --- a/qwen3_6_scripts/paged_attn.py +++ b/qwen3_6_scripts/paged_attn.py @@ -5,6 +5,12 @@ import torch import traceback from vllm import _custom_ops as ops +# from vllm.attention.ops.prefix_prefill import context_attention_fwd +# NOTE: context_attention_fwd (Triton kernel from prefix_prefill.py) is NOT +# imported here. On Iluvatar BI-V100 that kernel hangs the GPU card +# permanently. Chunked-prefill / prefix-caching attention is handled by +# _forward_prefix_pytorch below (pure PyTorch, no Triton dependency). + # Should be the same as PARTITION_SIZE in `paged_attention_v2_launcher`. _PARTITION_SIZE = 512 @@ -12,8 +18,17 @@ _PARTITION_SIZE = 512 @dataclass class PagedAttentionMetadata: """Metadata for PagedAttention.""" + # (batch_size,). The length of sequences (entire tokens seen so far) per + # sequence. seq_lens_tensor: Optional[torch.Tensor] + # Maximum sequence length in the batch. 0 if it is prefill-only batch. max_decode_seq_len: int + # (batch_size, max_blocks_per_seq). + # Block addresses per sequence. (Seq id -> list of physical block) + # E.g., [0, 1, 2] means tokens are stored in 0th, 1st, and 2nd blocks + # in the kv cache. Each block can contain up to block_size tokens. + # 2nd dimensions are padded up to max_blocks_per_seq if it is cuda-graph + # captured. block_tables: Optional[torch.Tensor] @@ -40,8 +55,10 @@ class PagedAttention: ) -> Tuple[torch.Tensor, torch.Tensor]: x = 16 // kv_cache.element_size() num_blocks = kv_cache.shape[1] + key_cache = kv_cache[0] - key_cache = key_cache.view(num_blocks, num_kv_heads, head_size // x, -1, x) + key_cache = key_cache.view(num_blocks, num_kv_heads, head_size // x, + -1, x) value_cache = kv_cache[1] value_cache = value_cache.view(num_blocks, num_kv_heads, head_size, -1) return key_cache, value_cache @@ -62,7 +79,7 @@ class PagedAttention: value, key_cache, value_cache, - slot_mapping, + slot_mapping.flatten(), kv_cache_dtype, k_scale, v_scale, @@ -70,18 +87,34 @@ class PagedAttention: @staticmethod def _forward_decode_pytorch( - query, key_cache, value_cache, block_tables, seq_lens, scale - ): - """Pure-PyTorch decode fallback for seq_len > threshold. + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + scale: float, + ) -> torch.Tensor: + """Pure-PyTorch decode attention for long contexts (no hardware kernel). - Used when ixf_F.paged_attention_v1 cannot handle the sequence length. - Optimized with batched KV gather (no per-block Python loop). + paged_attention_v1 hangs on BI-V100 when max_seq_len > ~32K due to + shared memory limits. For decode, q_len=1 per sequence so no Q-tiling + is needed — the attention weight tensor is [H, 1, seq_len] which is + trivially small (~5 MB at 50K). + + Shapes + ------ + query : [num_seqs, num_heads, head_dim] + key_cache : [num_blocks, num_kv_heads, head_dim//x, block_size, x] + value_cache : [num_blocks, num_kv_heads, head_dim, block_size] + block_tables: [num_seqs, max_blocks_per_seq] + seq_lens : [num_seqs] """ num_seqs, num_heads, head_dim = query.shape num_kv_heads = key_cache.shape[1] block_size = value_cache.shape[3] gqa_ratio = num_heads // num_kv_heads orig_dtype = query.dtype + output = torch.empty_like(query) try: @@ -90,25 +123,33 @@ class PagedAttention: num_blocks = (seq_len + block_size - 1) // block_size blk_ids = block_tables[i, :num_blocks] - # Batched gather: one index_select for all blocks + # Gather K: [kv_h, head_dim, seq_len] fp32 — no GQA expansion. + # With kv_h=1 and seq_len=100K this is 98 MB vs 586 MB if expanded. k_t = (key_cache[blk_ids] .permute(0, 3, 1, 2, 4) .contiguous() .view(-1, num_kv_heads, head_dim))[:seq_len] \ - .permute(1, 2, 0).contiguous().float() + .permute(1, 2, 0).contiguous().float() # [kv_h, d, seq_len] + # Gather V: [kv_h, seq_len, head_dim] fp32 v_t = (value_cache[blk_ids] .permute(0, 3, 1, 2) .contiguous() .view(-1, num_kv_heads, head_dim))[:seq_len] \ - .permute(1, 0, 2).contiguous().float() + .permute(1, 0, 2).contiguous().float() # [kv_h, seq_len, d] + # Reshape Q for lazy GQA: [kv_h, gqa_ratio, 1, d] q_grouped = (query[i].float() .view(num_kv_heads, gqa_ratio, head_dim) .unsqueeze(2)) - attn_w = torch.matmul(q_grouped * scale, k_t.unsqueeze(1)) + # [kv_h, gqa_ratio, 1, seq_len] + attn_w = torch.matmul( + q_grouped * scale, # [kv_h, gqa, 1, d] + k_t.unsqueeze(1)) # [kv_h, 1, d, seq_len] attn_w = torch.softmax(attn_w, dim=-1) + + # [kv_h, gqa_ratio, 1, d] → [num_heads, head_dim] out_i = torch.matmul(attn_w, v_t.unsqueeze(1)) output[i] = out_i.view(num_heads, head_dim).to(orig_dtype) @@ -120,9 +161,10 @@ class PagedAttention: return output - # BI-V100: Try higher threshold for compiled v1 kernel. - # Compiled kernel is ~100x faster than Python fallback. - _PYTORCH_DECODE_THRESHOLD = 65536 + # paged_attention_v1 on BI-V100 fails for long contexts. + # Route on actual sequence length (seq_lens.max()), not the max_seq_len + # parameter which is inflated to max_model_len in CUDA graph mode. + _PYTORCH_DECODE_THRESHOLD = 32768 @staticmethod def forward_decode( @@ -149,17 +191,31 @@ class PagedAttention: return PagedAttention._forward_decode_pytorch( query, key_cache, value_cache, block_tables, seq_lens, scale) + if blocksparse_vert_stride is not None and blocksparse_vert_stride > 1: + # use blocksparse paged attention + block_size = value_cache.size(-1) + assert (blocksparse_block_size > 0 and + blocksparse_block_size % block_size == 0), \ + (f"{blocksparse_block_size=} needs to be a multiple of" + f"{block_size=} used in block_tables.") + output = torch.empty_like(query) block_size = value_cache.shape[3] num_seqs, num_heads, head_size = query.shape max_num_partitions = ((max_seq_len + _PARTITION_SIZE - 1) // _PARTITION_SIZE) - + # NOTE(woosuk): We use a simple heuristic to decide whether to use + # PagedAttention V1 or V2. If the number of partitions is 1, we use + # V1 to avoid the overhead of reduction. Also, if the number of + # sequences or heads is large, we use V1 since there is enough work + # to parallelize. + # TODO(woosuk): Tune this heuristic. + # For context len > 8192, use V2 kernel to avoid shared memory shortage. use_v1 = (max_seq_len <= 8192 and (max_num_partitions == 1 or num_seqs * num_heads > 512)) - # V2 now works (paged_attention_v2_pytorch), so use the original heuristic - # instead of hardcoding use_v1=True + use_v1 = True if use_v1: + # Run PagedAttention V1. ops.paged_attention_v1( output, query, @@ -174,6 +230,7 @@ class PagedAttention: alibi_slopes, ) else: + # Run PagedAttention V2. assert _PARTITION_SIZE % block_size == 0 tmp_output = torch.empty( size=(num_seqs, num_heads, max_num_partitions, head_size), @@ -212,9 +269,6 @@ class PagedAttention: ) return output - # Triton prefill: try once, fall back permanently if it fails - _triton_prefill_ok = None - @staticmethod def forward_prefix( query: torch.Tensor, @@ -233,30 +287,9 @@ class PagedAttention: k_scale: float, v_scale: float, ) -> torch.Tensor: - # Try Triton kernel if available and not known to fail - if PagedAttention._triton_prefill_ok is not False: - try: - from vllm.triton_utils import HAS_TRITON - if HAS_TRITON: - from vllm.attention.ops.prefix_prefill import context_attention_fwd - output = torch.empty_like(query) - context_attention_fwd( - query, key, value, output, kv_cache_dtype, - key_cache, value_cache, block_tables, - query_start_loc[:-1], seq_lens_tensor, context_lens, - max_query_len, k_scale, v_scale, - alibi_slopes, sliding_window, - ) - if PagedAttention._triton_prefill_ok is None: - print("[paged_attn] Triton prefill kernel: SUCCESS", flush=True) - PagedAttention._triton_prefill_ok = True - return output - except Exception as e: - print(f"[paged_attn] Triton prefill failed: {type(e).__name__}: {e}", - flush=True) - print("[paged_attn] Falling back to PyTorch prefill permanently", flush=True) - PagedAttention._triton_prefill_ok = False - + # NOTE: The Triton context_attention_fwd kernel hangs on Iluvatar + # BI-V100 hardware (same class of issue as cudnnFlashAttnForward). + # Use a pure-PyTorch fallback that reads the paged KV cache directly. return PagedAttention._forward_prefix_pytorch( query, key, value, key_cache, value_cache, @@ -276,101 +309,150 @@ class PagedAttention: seq_lens_tensor: torch.Tensor, context_lens: torch.Tensor, ) -> torch.Tensor: - """Pure-PyTorch prefix-attention with pre-gathered KV and Flash-Attention online softmax. + """Pure-PyTorch prefix-attention with K-tiling (Flash-Attention online softmax). - Optimization over baseline: - - Context KV is gathered ONCE outside the tile loop (one index_select + reshape) - - Tile loop just slices views from the pre-gathered tensor (no per-tile gather) - - Eliminates 194 redundant permute+contiguous calls for 100K context + Memory complexity: O(q_len), independent of kv_len. + With chunked prefill (q_len ≤ max_num_batched_tokens = 4096) peak + per layer ≈ 96 MB regardless of context length. - Memory: O(q_len × tile_sz) per tile — same as baseline. - The full context K/V tensor is ~50MB for 100K tokens, fits in GPU memory. + Algorithm: Flash Attention online softmax. + Q is reshaped once to [kv_h, gqa, q_len, d] (24 MB) and held for all + K-tiles. For each tile a running (m, l, o) accumulator is updated — + the [q_len × kv_len] attention matrix is NEVER materialised in full. + + Tile budget (kv_h=1, gqa=6, q_len=4096, tile=256 tokens): + q_seq [1, 6, 4096, 256] fp32 24 MB (held all tiles) + o_acc same shape 24 MB (held all tiles) + s same shape 24 MB (per tile, freed before exp_s) + exp_s same shape 24 MB (per tile, brief overlap with s) + Peak ≈ 96 MB (s and exp_s briefly coexist during update). + + Shapes + ------ + query : [total_q_tokens, num_q_heads, head_dim] + key : [total_q_tokens, num_kv_heads, head_dim] + value : [total_q_tokens, num_kv_heads, head_dim] + key_cache : [num_blocks, num_kv_heads, head_dim//x, block_size, x] + value_cache : [num_blocks, num_kv_heads, head_dim, block_size] + block_tables : [batch_size, max_blocks_per_seq] + query_start_loc: [batch_size + 1] + seq_lens_tensor: [batch_size] total length (context + query) + context_lens : [batch_size] tokens already in KV cache """ try: + # Paged-block tiles for context phase. + # tile_sz = _BLOCKS_PER_TILE × block_size (e.g. 16×16 = 256 tokens). + # Score tensor [kv_h, gqa, q_len, tile_sz] fp32 = 24 MB per tile. + # Same tile size reused for the current-chunk phase. _BLOCKS_PER_TILE = 32 - batch_size = seq_lens_tensor.shape[0] - num_q_heads = query.shape[1] + batch_size = seq_lens_tensor.shape[0] + num_q_heads = query.shape[1] num_kv_heads = key_cache.shape[1] - head_dim = query.shape[2] - gqa_ratio = num_q_heads // num_kv_heads - block_size = value_cache.shape[3] - tile_sz = _BLOCKS_PER_TILE * block_size - scale = head_dim ** -0.5 - orig_dtype = query.dtype - output = torch.empty_like(query) - dev = query.device + head_dim = query.shape[2] + gqa_ratio = num_q_heads // num_kv_heads + block_size = value_cache.shape[3] + tile_sz = _BLOCKS_PER_TILE * block_size + scale = head_dim ** -0.5 + orig_dtype = query.dtype + output = torch.empty_like(query) + dev = query.device for i in range(batch_size): ctx_len = int(context_lens[i].item()) q_start = int(query_start_loc[i].item()) - q_end = int(query_start_loc[i + 1].item()) - q_len = q_end - q_start + q_end = int(query_start_loc[i + 1].item()) + q_len = q_end - q_start - q_i = query[q_start:q_end] - k_i = key[q_start:q_end] + q_i = query[q_start:q_end] # [q_len, q_h, d] + k_i = key [q_start:q_end] # [q_len, kv_h, d] v_i = value[q_start:q_end] + # Q reshaped and scaled once; held for all K-tiles. + # [kv_h, gqa, q_len, d] fp32 — 24 MB for q_len=4096, d=256 q_seq = (q_i.permute(1, 0, 2) .float() .view(num_kv_heads, gqa_ratio, q_len, head_dim) .mul_(scale)) + # Flash-Attention online-softmax accumulators. + # m, l : [kv_h, gqa, q_len] fp32 — <0.1 MB + # o : [kv_h, gqa, q_len, d] fp32 — 24 MB m = torch.full((num_kv_heads, gqa_ratio, q_len), float('-inf'), dtype=torch.float32, device=dev) l = torch.zeros_like(m) o = torch.zeros((num_kv_heads, gqa_ratio, q_len, head_dim), dtype=torch.float32, device=dev) - # =========================================================== - # Phase 1: Context tokens — PRE-GATHER optimization - # Gather ALL context K/V in ONE shot, then tile via slicing - # =========================================================== + # -------------------------------------------------------------- + # Phase 1 — context tokens (positions 0 … ctx_len-1). + # + # Every context key has absolute position < ctx_len; every + # query has position ≥ ctx_len. k_pos < q_pos is always True + # → no causal mask needed for pure context tiles. + # -------------------------------------------------------------- if ctx_len > 0: num_ctx_blocks = (ctx_len + block_size - 1) // block_size + # Safety: if block_tables is too narrow this indicates a + # prefix_cache_hit + chunked-prefill bug in model_runner.py + # (Case 1 leaves prefix_cache_hit=True but block_table is + # only computed_block_nums, not the full context blocks). + # patch_model_runner.py fixes the root cause; this guard + # prevents a zero-dim amax() crash if it still slips through. if num_ctx_blocks > block_tables.shape[1]: print( f"[paged_attn WARNING] seq {i}: num_ctx_blocks={num_ctx_blocks} " - f"> block_tables.shape[1]={block_tables.shape[1]}. " - "Capping context to available blocks.", + f"> block_tables.shape[1]={block_tables.shape[1]}, ctx_len={ctx_len}. " + "Block table is undersized (prefix_cache_hit bug). " + "Capping context to available blocks — attention may be incorrect.", file=sys.stderr, flush=True) num_ctx_blocks = block_tables.shape[1] + for tile_blk in range(0, num_ctx_blocks, _BLOCKS_PER_TILE): + blk_end = min(tile_blk + _BLOCKS_PER_TILE, num_ctx_blocks) + blk_ids = block_tables[i, tile_blk:blk_end] - # ONE gather for ALL context blocks - ctx_blk_ids = block_tables[i, :num_ctx_blocks] + # Gather K/V for this tile. + # key_cache [blk_ids]: [n, kv_h, d//x, blk_sz, x] + # value_cache[blk_ids]: [n, kv_h, d, blk_sz] + k_tile = (key_cache[blk_ids] + .permute(0, 3, 1, 2, 4) + .contiguous() + .view(-1, num_kv_heads, head_dim)) + v_tile = (value_cache[blk_ids] + .permute(0, 3, 1, 2) + .contiguous() + .view(-1, num_kv_heads, head_dim)) - # [num_ctx_blocks, kv_h, d/x, blk_sz, x] → [ctx_tokens, kv_h, d] - ctx_k_all = (key_cache[ctx_blk_ids] - .permute(0, 3, 1, 2, 4) - .contiguous() - .view(-1, num_kv_heads, head_dim))[:ctx_len] + # Trim padding in the last block of the tile. + valid = (min(blk_end * block_size, ctx_len) + - tile_blk * block_size) + k_tile = k_tile[:valid] # [valid, kv_h, d] + v_tile = v_tile[:valid] - ctx_v_all = (value_cache[ctx_blk_ids] - .permute(0, 3, 1, 2) - .contiguous() - .view(-1, num_kv_heads, head_dim))[:ctx_len] - - # Pre-transpose for matmul: [kv_h, d, ctx_len] and [kv_h, ctx_len, d] - ctx_k_t = ctx_k_all.permute(1, 2, 0).contiguous().float() - ctx_v_t = ctx_v_all.permute(1, 0, 2).contiguous().float() - - # Tile loop: just SLICE from pre-gathered tensors - for tile_start in range(0, ctx_len, tile_sz): - tile_end = min(tile_start + tile_sz, ctx_len) - - # Slice (view, no copy) - k_t = ctx_k_t[:, :, tile_start:tile_end].unsqueeze(1) - v_t = ctx_v_t[:, tile_start:tile_end, :].unsqueeze(1) + # k_t: [kv_h, 1, d, valid] (broadcast over gqa_ratio) + # v_t: [kv_h, 1, valid, d] + k_t = (k_tile.permute(1, 0, 2) + .unsqueeze(1) + .transpose(-1, -2) + .float()) + v_t = (v_tile.permute(1, 0, 2) + .unsqueeze(1) + .float()) + del k_tile, v_tile + # Scores: [kv_h, gqa, q_len, valid] s = torch.matmul(q_seq, k_t) del k_t + # No causal mask: all context keys precede all queries. + # Online softmax update — Flash-Attention Algorithm 1. + # exp_s = s - new_max (in-place exp after del s) m_blk = s.amax(dim=-1) m_new = torch.maximum(m, m_blk) exp_s = s - m_new.unsqueeze(-1) del s exp_s.exp_() - corr = torch.exp(m - m_new) + corr = torch.exp(m - m_new) m.copy_(m_new) del m_blk, m_new l.mul_(corr).add_(exp_s.sum(dim=-1)) @@ -378,40 +460,44 @@ class PagedAttention: torch.matmul(exp_s, v_t)) del exp_s, v_t, corr - del ctx_k_t, ctx_v_t, ctx_k_all, ctx_v_all - - # =========================================================== - # Phase 2: Current-chunk tokens (with causal mask) - # =========================================================== + # -------------------------------------------------------------- + # Phase 2 — current-chunk tokens (positions ctx_len … ctx_len+q_len-1). + # + # Causal mask: query at relative position j sees key at relative + # position k only when k ≤ j. Tiles of tile_sz tokens each. + # -------------------------------------------------------------- for kc_start in range(0, q_len, tile_sz): kc_end = min(kc_start + tile_sz, q_len) + kc_len = kc_end - kc_start - k_blk = k_i[kc_start:kc_end] + k_blk = k_i[kc_start:kc_end] # [kc_len, kv_h, d] v_blk = v_i[kc_start:kc_end] k_t = (k_blk.permute(1, 0, 2) .unsqueeze(1) .transpose(-1, -2) - .float()) + .float()) # [kv_h, 1, d, kc_len] v_t = (v_blk.permute(1, 0, 2) .unsqueeze(1) - .float()) + .float()) # [kv_h, 1, kc_len, d] - s = torch.matmul(q_seq, k_t) + s = torch.matmul(q_seq, k_t) # [kv_h, gqa, q_len, kc_len] del k_t + # Causal mask: key at (kc_start+k) must not exceed query j. k_rel = torch.arange(kc_start, kc_end, device=dev) q_rel = torch.arange(q_len, device=dev) - mask = k_rel.unsqueeze(0) > q_rel.unsqueeze(1) + mask = k_rel.unsqueeze(0) > q_rel.unsqueeze(1) # [q_len, kc_len] s.masked_fill_(mask.unsqueeze(0).unsqueeze(0), float('-inf')) del mask, k_rel, q_rel + # Online softmax update (identical to context phase). m_blk = s.amax(dim=-1) m_new = torch.maximum(m, m_blk) exp_s = s - m_new.unsqueeze(-1) del s exp_s.exp_() - corr = torch.exp(m - m_new) + corr = torch.exp(m - m_new) m.copy_(m_new) del m_blk, m_new l.mul_(corr).add_(exp_s.sum(dim=-1)) @@ -419,6 +505,10 @@ class PagedAttention: torch.matmul(exp_s, v_t)) del exp_s, v_t, corr + # -------------------------------------------------------------- + # Finalize: normalize running output by normalization factor. + # o: [kv_h, gqa, q_len, d] → [q_len, q_h, d] + # -------------------------------------------------------------- o.div_(l.unsqueeze(-1)) output[q_start:q_end] = ( o.view(num_q_heads, q_len, head_dim) @@ -442,6 +532,7 @@ class PagedAttention: src_key_cache = src_kv_cache[0] dst_key_cache = dst_kv_cache[0] ops.swap_blocks(src_key_cache, dst_key_cache, src_to_dst) + src_value_cache = src_kv_cache[1] dst_value_cache = dst_kv_cache[1] ops.swap_blocks(src_value_cache, dst_value_cache, src_to_dst) diff --git a/qwen3_6_scripts/qwen3_5.py b/qwen3_6_scripts/qwen3_5.py index b63ab5bd..ca427609 100644 --- a/qwen3_6_scripts/qwen3_5.py +++ b/qwen3_6_scripts/qwen3_5.py @@ -795,66 +795,18 @@ class Qwen3_5MoeSparseBlock(nn.Module): # General path (prefill / multi-seq): loop over unique active experts. # At most T*top_k unique experts, always <= num_experts. out = torch.zeros_like(hidden_states) - # Optimized prefill MoE: sort tokens by expert for contiguous access. - # Pattern from CCCL segmented sort: partition input by key (expert), - # process each segment on contiguous memory, scatter results back. - # - # Before: for eid in unique_eids: hidden_states[tok_ids] (scattered gather) - # After: sort by expert → each F.linear gets contiguous input - # activation computed in ONE fused op across all token×expert pairs - # index_add_ for scatter-back is one kernel - - T_local = hidden_states.shape[0] - two_I = w13.shape[1] - - # Flatten token-expert assignments: each token appears top_k times - flat_ids = topk_ids.view(-1) # (T*K,) - flat_weights_f = topk_weights.view(-1) # (T*K,) - flat_tok_idx = torch.arange( - T_local, device=hidden_states.device - ).unsqueeze(1).expand(-1, self.top_k).reshape(-1) # (T*K,) - - # Sort by expert ID → contiguous memory per expert - sorted_order = flat_ids.argsort(stable=True) - sorted_eids = flat_ids[sorted_order] - sorted_tok_idx = flat_tok_idx[sorted_order] - sorted_weights_f = flat_weights_f[sorted_order] - sorted_tokens = hidden_states[sorted_tok_idx] # (T*K, H) contiguous groups - - # Expert boundaries via unique_consecutive - unique_eids_t, counts = torch.unique_consecutive( - sorted_eids, return_counts=True) - cum_counts = counts.cumsum(0) - starts = torch.cat([torch.zeros(1, dtype=torch.long, device=hidden_states.device), - cum_counts[:-1]]) - - # Grouped up-projection: each expert's tokens are contiguous - TK = T_local * self.top_k - gate_up_all = torch.empty(TK, two_I, dtype=hidden_states.dtype, - device=hidden_states.device) - for idx_e in range(len(unique_eids_t)): - eid_val = unique_eids_t[idx_e].item() - s = starts[idx_e].item() - e = cum_counts[idx_e].item() - gate_up_all[s:e] = F.linear(sorted_tokens[s:e], w13[eid_val]) - - # Fused activation across ALL token×expert pairs (one kernel) - gate, up = gate_up_all.chunk(2, dim=-1) - act = F.silu(gate) * up - - # Grouped down-projection - down_all = torch.empty(TK, hidden_states.shape[-1], - dtype=hidden_states.dtype, - device=hidden_states.device) - for idx_e in range(len(unique_eids_t)): - eid_val = unique_eids_t[idx_e].item() - s = starts[idx_e].item() - e = cum_counts[idx_e].item() - down_all[s:e] = F.linear(act[s:e], w2[eid_val]) - - # Weighted scatter-add (one kernel) - weighted = down_all * sorted_weights_f.unsqueeze(-1) - out.index_add_(0, sorted_tok_idx, weighted.to(out.dtype)) + unique_eids = topk_ids.view(-1).unique().tolist() + for eid in unique_eids: + eid = int(eid) + mask = (topk_ids == eid) # (T, top_k) + tok_ids, topk_pos = mask.nonzero(as_tuple=True) + tokens = hidden_states[tok_ids] # (n, H) + gate_up = F.linear(tokens, w13[eid]) # (n, 2*I) + gate, up = gate_up.chunk(2, dim=-1) + act = F.silu(gate) * up # (n, I) + expert_out = F.linear(act, w2[eid]) # (n, H) + weights = topk_weights[tok_ids, topk_pos].unsqueeze(-1) + out.index_add_(0, tok_ids, (expert_out * weights).to(out.dtype)) return out # partial, all-reduce done in forward()