feat(muh): apply CCCL-derived BI-V100 tuning to 5 vllm Python files

Applied via muh/vllm_bi100_patch.py --conservative:

1. paged_attn.py: removed use_v1=True hardcode, restored V1/V2 heuristic
   with BI-V100 threshold (16384 vs default 8192). SM=16 favors V1 longer.

2. fused_moe.py: BLOCK_SIZE_K 32→64 (better memory coalescing with 900GB/s
   BW), BLOCK_SIZE_N 32→64 for decode path. Qwen3.6 MoE: E≈128, topk=8.

3. _custom_ops.py: SMEM kept at 32KB (conservative mode, pending hardware
   confirmation). Added diagnostic comment.

4. prefix_prefill.py: enhanced BI-V100 block config comment with SMEM
   budget breakdown (BLOCK=64,N=64 → 48KB tight, N=32 → 32KB safe).

5. triton_flash_attention.py: added 2 BI-V100 autotune configs
   (64x32 and 32x64) for SM=16 occupancy characteristics.

CCCL basis: cub/benchmarks/bench/ %RANGE% parameter spaces (reduce 1044
combos, scan 5.4M, topk 1698, transform 25920) → SMEM pruning → policy
selector logic from tuning_*.cuh.

Also includes muh/vllm_bi100_patch.py (713 lines) for reproducible
one-shot patching with --dry-run, --conservative, and --revert modes.
This commit is contained in:
dylanyunlon
2026-08-03 10:27:10 +00:00
parent 094c710efa
commit dc9ac0a757
6 changed files with 431 additions and 6 deletions

View File

@@ -726,8 +726,12 @@ if triton.__version__ >= "2.1.0":
# Triton equivalent: BLOCK=64, warps=4 (128 threads, larger tile per warp)
_is_bi_v100 = not current_platform.has_device_capability(80)
if _is_bi_v100:
# muh: CCCL-informed block selection for BI-V100 (SM=16, SMEM≤48KB)
# SMEM = BLOCK_M*Hd*elem + BLOCK_N*Hd*elem*2(K+V)
# head_dim=128, fp16(2B): BLOCK=64,N=64 → 48KB (100% SMEM, risky)
# Conservative: BLOCK=64,N=32 → 32KB (65% SMEM, safe for 32KB limit)
BLOCK = 64
NUM_WARPS = 4
NUM_WARPS = 4 # 4 warps × 32 = 128 threads; BW-limited at 56 GB/s/SM
else:
BLOCK = 128
NUM_WARPS = 8