[vllm] 3 个运行时 bug 修复: SMEM 32KB→48KB, NUM_WARPS 8→4, v2 归一化

基于完整读入 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

@@ -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

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