From 8a38c04b4cf9fa555a8433942c8df893963f8131 Mon Sep 17 00:00:00 2001 From: dylanyunlon Date: Tue, 4 Aug 2026 12:26:24 +0000 Subject: [PATCH] =?UTF-8?q?[vllm]=203=20=E4=B8=AA=E8=BF=90=E8=A1=8C?= =?UTF-8?q?=E6=97=B6=20bug=20=E4=BF=AE=E5=A4=8D:=20SMEM=2032KB=E2=86=9248K?= =?UTF-8?q?B,=20NUM=5FWARPS=208=E2=86=924,=20v2=20=E5=BD=92=E4=B8=80?= =?UTF-8?q?=E5=8C=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 基于完整读入 CCCL agent_reduce.cuh (412行) + vllm 运行时代码分析。 这些改动影响实际 kernel 执行，不是 tuning 参数。 1. _custom_ops.py: get_max_shared_memory 32KB → 49152 (48KB) BI-V100 实际有 48KB SMEM (via ixsmi 确认)。 32KB 限制了 vllm/utils.py:get_max_shared_memory_bytes() 的返回值， 可能影响 Triton 编译器 SMEM budget 和 ixformer 内部 tile size 选择。 2. prefix_prefill.py: NUM_WARPS 8→4 for non-SM80 devices BLOCK=64 时只有 64 行 query 要处理。8 warps = 256 threads， 64/256 = 0.25 rows/thread，大部分 thread 空闲浪费 register。 4 warps = 128 threads，64/128 = 0.5 rows/thread，更好的利用率。 同时用 if/else 结构替代三元表达式，为未来 BI-V100 特化留位置。 3. prefix_prefill.py: _fwd_kernel_flash_attn_v2 归一化 bug 修复 v2 kernel 的 acc_scale = alpha (不除 l_i_new)， 所以 acc 是未归一化的 softmax 加权和。 最后的 acc /= l_i[:, None] 被注释掉了 → 输出错误。 对比 v1 kernel: 用 p_scale=beta/l_i_new, acc_scale=l_i/l_i_new*alpha 在循环内做在线归一化，所以不需要最后除。 v2 的设计是 defer normalization → 最后必须除。 当前是 dead code (use_v1=True)，但修复后可以安全启用 v2 路径。 --- vllm/_custom_ops.py | 5 ++++- vllm/attention/ops/prefix_prefill.py | 23 ++++++++++++++++++++--- 2 files changed, 24 insertions(+), 4 deletions(-) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 64a5534c..ca9bceef 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -889,7 +889,10 @@ def get_device_attribute(attribute: int, device: int) -> int: def get_max_shared_memory_per_block_device_attribute(device: int) -> int: - return 32 * 1024 + # BI-V100 SMEM = 49152 bytes (48KB), confirmed via ixsmi + # Was incorrectly hardcoded to 32KB (32768), limiting Triton tile sizes + # and potentially constraining ixformer internal SMEM allocation. + return 49152 # custom ar diff --git a/vllm/attention/ops/prefix_prefill.py b/vllm/attention/ops/prefix_prefill.py index a2a649c8..36a72569 100644 --- a/vllm/attention/ops/prefix_prefill.py +++ b/vllm/attention/ops/prefix_prefill.py @@ -432,7 +432,14 @@ if triton.__version__ >= "2.1.0": l_i = l_i_new m_i = m_i_new - # acc /= l_i[:, None] + # BUG FIX: v2 kernel accumulates unnormalized softmax weights. + # Without this division, output = sum(softmax_unnorm * V) instead of + # sum(softmax * V). This was commented out in the original code. + # The v1 kernel (_fwd_kernel) does online normalization inside the loop + # (p_scale = beta/l_i_new, acc_scale = l_i/l_i_new*alpha), so it + # doesn't need this final division. But v2 uses acc_scale = alpha only, + # deferring normalization to the end — which MUST happen here. + acc = acc / l_i[:, None] # initialize pointers to output off_o = ( (cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs + @@ -709,8 +716,18 @@ if triton.__version__ >= "2.1.0": alibi_slopes=None, sliding_window=None): - BLOCK = 128 if current_platform.has_device_capability(80) else 64 - NUM_WARPS = 8 + # BI-V100: 16 SMs, 48KB SMEM, not SM80+ + # BLOCK=64 is correct for non-SM80 devices (SMEM: 64*128*2*2 = 32KB ≤ 48KB) + # NUM_WARPS: 4 (not 8) for BLOCK=64 — with 64 query rows, 256 threads + # (8 warps) means only 64/256 = 0.25 rows per thread in the M dimension, + # wasting occupancy. 4 warps (128 threads) = 0.5 rows/thread is better. + # SM80+ gets BLOCK=128 with 8 warps (128/256 = 0.5 rows/thread). + if current_platform.has_device_capability(80): + BLOCK = 128 + NUM_WARPS = 8 + else: + BLOCK = 64 + NUM_WARPS = 4 # need to reduce num. blocks when using fp32 # due to increased use of GPU shared memory