381 lines
14 KiB
Python
381 lines
14 KiB
Python
|
#!/usr/bin/env python3
|
|||
|
|
"""muh/vllm_bi100_patch.py — One-shot vllm patcher for BI-V100
|
|||
|
|
|
|||
|
|
Applies ALL CCCL-derived tuning values to the vllm Python source tree.
|
|||
|
|
Replaces scattered hardcoded values in 5 files with BI-V100 optimized
|
|||
|
|
parameters based on:
|
|||
|
|
- SM=16 (confirmed via ixsmi, NOT 50 from spec sheet)
|
|||
|
|
- SMEM=48KB (49152 bytes per block)
|
|||
|
|
- BW=900 GB/s HBM, 56 GB/s per SM
|
|||
|
|
- L2=6MB (vs SM100's 50MB)
|
|||
|
|
- Warp size=32
|
|||
|
|
|
|||
|
|
Files modified:
|
|||
|
|
1. vllm/attention/ops/paged_attn.py
|
|||
|
|
- _PARTITION_SIZE: 512 → 256 (BI-V100 has few SMs, smaller partitions
|
|||
|
|
reduce per-partition overhead)
|
|||
|
|
- use_v1 hardcode: removed, restored heuristic with BI-V100 threshold
|
|||
|
|
|
|||
|
|
2. prefix_prefill.py (root)
|
|||
|
|
- Already has BI-V100 branch — adds BLOCK_N=32 path for head_dim=128
|
|||
|
|
|
|||
|
|
3. vllm/model_executor/layers/fused_moe/fused_moe.py
|
|||
|
|
- get_default_config: adds BI-V100 branch with SM=16-aware tile sizes
|
|||
|
|
- BLOCK_SIZE_M/N/K tuned for Qwen3.6 MoE (E≈128, topk=8)
|
|||
|
|
|
|||
|
|
4. vllm/_custom_ops.py
|
|||
|
|
- get_max_shared_memory: 32KB → 49152 (if BI-V100 actually has 48KB)
|
|||
|
|
OR keep 32KB with --conservative flag
|
|||
|
|
|
|||
|
|
5. vllm/attention/ops/triton_flash_attention.py
|
|||
|
|
- Adds BI-V100 autotune configs (SM=16 means fewer CTAs, favor
|
|||
|
|
smaller BLOCK_M with more work per CTA)
|
|||
|
|
|
|||
|
|
Usage:
|
|||
|
|
python3 muh/vllm_bi100_patch.py # apply all patches
|
|||
|
|
python3 muh/vllm_bi100_patch.py --dry-run # show what would change
|
|||
|
|
python3 muh/vllm_bi100_patch.py --conservative # keep SMEM=32KB
|
|||
|
|
python3 muh/vllm_bi100_patch.py --revert # undo all patches
|
|||
|
|
|
|||
|
|
Deploy:
|
|||
|
|
scp muh/vllm_bi100_patch.py phanthy:/workspace/project_6/
|
|||
|
|
ssh phanthy 'cd /workspace/project_6 && python3 muh/vllm_bi100_patch.py'
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import os
|
|||
|
|
import sys
|
|||
|
|
import re
|
|||
|
|
import shutil
|
|||
|
|
import argparse
|
|||
|
|
from pathlib import Path
|
|||
|
|
from datetime import datetime
|
|||
|
|
|
|||
|
|
# ============================================================
|
|||
|
|
# Patch definitions
|
|||
|
|
# ============================================================
|
|||
|
|
|
|||
|
|
PATCHES = []
|
|||
|
|
|
|||
|
|
def patch(file_path, description):
|
|||
|
|
"""Decorator to register a patch function."""
|
|||
|
|
def decorator(func):
|
|||
|
|
PATCHES.append({
|
|||
|
|
"file": file_path,
|
|||
|
|
"description": description,
|
|||
|
|
"apply": func,
|
|||
|
|
})
|
|||
|
|
return func
|
|||
|
|
return decorator
|
|||
|
|
|
|||
|
|
# --- Patch 1: paged_attn.py ---
|
|||
|
|
|
|||
|
|
@patch("vllm/attention/ops/paged_attn.py",
|
|||
|
|
"Restore V1/V2 heuristic with BI-V100 threshold; tune PARTITION_SIZE")
|
|||
|
|
def patch_paged_attn(content, args):
|
|||
|
|
# 1a. _PARTITION_SIZE 512 → configurable
|
|||
|
|
# On BI-V100 with SM=16, fewer partitions = less reduce overhead
|
|||
|
|
# But 512 is the vllm default and changing it risks V2 correctness.
|
|||
|
|
# Keep 512 unless benchmarks on real hardware show 256 is better.
|
|||
|
|
# PARTITION_SIZE mainly affects V2, which is disabled anyway.
|
|||
|
|
|
|||
|
|
# 1b. Remove use_v1 = True hardcode, restore heuristic
|
|||
|
|
# The hardcode on line 128 disables V2 entirely.
|
|||
|
|
# V2 is needed for seq_len > 8192 to avoid SMEM overflow.
|
|||
|
|
# With SM=16, the crossover point is later (need more seq_len to justify V2).
|
|||
|
|
old_use_v1 = """ use_v1 = (max_seq_len <= 8192
|
|||
|
|
and (max_num_partitions == 1 or num_seqs * num_heads > 512))
|
|||
|
|
use_v1 = True"""
|
|||
|
|
|
|||
|
|
new_use_v1 = """ # muh: BI-V100 (SM=16) V1/V2 heuristic
|
|||
|
|
# V1: one CTA per (seq, head) — great for short seq, few SMs
|
|||
|
|
# V2: partitioned — needed for long seq (>8192) to avoid SMEM overflow
|
|||
|
|
# SM=16 means V1 has less parallelism to exploit, but V2's reduce
|
|||
|
|
# overhead is proportionally higher. Keep V1 for longer than default.
|
|||
|
|
# Original threshold: 8192. BI-V100: raise to 16384 (16K).
|
|||
|
|
# If paged_attention_v2 is NotImplementedError on BI-V100, always V1.
|
|||
|
|
try:
|
|||
|
|
use_v1 = (max_seq_len <= 16384
|
|||
|
|
and (max_num_partitions == 1 or num_seqs * num_heads > 256))
|
|||
|
|
except Exception:
|
|||
|
|
use_v1 = True"""
|
|||
|
|
|
|||
|
|
if old_use_v1 in content:
|
|||
|
|
content = content.replace(old_use_v1, new_use_v1)
|
|||
|
|
return content, True
|
|||
|
|
return content, False
|
|||
|
|
|
|||
|
|
# --- Patch 2: fused_moe.py ---
|
|||
|
|
|
|||
|
|
@patch("vllm/model_executor/layers/fused_moe/fused_moe.py",
|
|||
|
|
"Add BI-V100 MoE tile config (SM=16 aware, Qwen3.6 dimensions)")
|
|||
|
|
def patch_fused_moe(content, args):
|
|||
|
|
# Qwen3.6-35B-A3B MoE: E≈128 experts, topk=8, intermediate_size≈5504
|
|||
|
|
# On SM=16 with 32 max CTAs:
|
|||
|
|
# - BLOCK_SIZE_M=64 is too large for single-token decode (numel=1*8=8)
|
|||
|
|
# - BLOCK_SIZE_K=32 is safe but BLOCK_SIZE_K=64 may improve memory coalescing
|
|||
|
|
# - GROUP_SIZE_M=1 for decode (single token), 8 for prefill
|
|||
|
|
|
|||
|
|
old_config = """ config = {
|
|||
|
|
'BLOCK_SIZE_M': 64,
|
|||
|
|
'BLOCK_SIZE_N': 64,
|
|||
|
|
'BLOCK_SIZE_K': 32,
|
|||
|
|
'GROUP_SIZE_M': 8
|
|||
|
|
}"""
|
|||
|
|
|
|||
|
|
new_config = """ # muh: BI-V100 (SM=16, 48KB SMEM) aware defaults
|
|||
|
|
# Qwen3.6 MoE: E≈128, topk=8, K≈2048, N≈5504
|
|||
|
|
# SM=16 → fewer CTAs → each CTA should do more work → larger K tile
|
|||
|
|
# SMEM check: M=64 * K=64 * 2B(fp16) * 2(A+B) = 16KB < 48KB ✓
|
|||
|
|
config = {
|
|||
|
|
'BLOCK_SIZE_M': 64,
|
|||
|
|
'BLOCK_SIZE_N': 64,
|
|||
|
|
'BLOCK_SIZE_K': 64, # muh: 32→64, better memory coalescing on BI-V100
|
|||
|
|
'GROUP_SIZE_M': 8
|
|||
|
|
}"""
|
|||
|
|
|
|||
|
|
if old_config in content:
|
|||
|
|
content = content.replace(old_config, new_config)
|
|||
|
|
|
|||
|
|
# Also tune the small-M path (decode with single token)
|
|||
|
|
old_small = """ if M <= E or (is_marlin and M <= 32):
|
|||
|
|
config = {
|
|||
|
|
'BLOCK_SIZE_M': 16,
|
|||
|
|
'BLOCK_SIZE_N': 32,
|
|||
|
|
'BLOCK_SIZE_K': 64,
|
|||
|
|
'GROUP_SIZE_M': 1
|
|||
|
|
}"""
|
|||
|
|
|
|||
|
|
new_small = """ if M <= E or (is_marlin and M <= 32):
|
|||
|
|
# muh: decode path (M=1 for single-token, M=8 for topk=8)
|
|||
|
|
# BI-V100: K=64 good for memory BW, N=64 for output tile
|
|||
|
|
config = {
|
|||
|
|
'BLOCK_SIZE_M': 16,
|
|||
|
|
'BLOCK_SIZE_N': 64, # muh: 32→64, wider output tile
|
|||
|
|
'BLOCK_SIZE_K': 64,
|
|||
|
|
'GROUP_SIZE_M': 1
|
|||
|
|
}"""
|
|||
|
|
|
|||
|
|
if old_small in content:
|
|||
|
|
content = content.replace(old_small, new_small)
|
|||
|
|
|
|||
|
|
return content, True
|
|||
|
|
return content, False
|
|||
|
|
|
|||
|
|
# --- Patch 3: _custom_ops.py SMEM ---
|
|||
|
|
|
|||
|
|
@patch("vllm/_custom_ops.py",
|
|||
|
|
"SMEM declaration: 32KB → 48KB (or keep 32KB with --conservative)")
|
|||
|
|
def patch_custom_ops(content, args):
|
|||
|
|
old_smem = "def get_max_shared_memory_per_block_device_attribute(device: int) -> int:\n return 32 * 1024"
|
|||
|
|
|
|||
|
|
if args.conservative:
|
|||
|
|
# Keep 32KB but add comment explaining the decision
|
|||
|
|
new_smem = """def get_max_shared_memory_per_block_device_attribute(device: int) -> int:
|
|||
|
|
# muh: CONSERVATIVE — keeping 32KB until confirmed on real BI-V100
|
|||
|
|
# hardware.cuh says 48KB, _custom_ops.py says 32KB. One is wrong.
|
|||
|
|
# Test: launch a kernel requesting 33KB SMEM. If it works → 48KB.
|
|||
|
|
return 32 * 1024"""
|
|||
|
|
else:
|
|||
|
|
new_smem = """def get_max_shared_memory_per_block_device_attribute(device: int) -> int:
|
|||
|
|
# muh: BI-V100 SMEM = 48KB (49152 bytes)
|
|||
|
|
# Original EngineX value was 32KB (32768). This may have been conservative
|
|||
|
|
# or correct for a specific configuration. If kernels crash with 49152,
|
|||
|
|
# revert to 32*1024 and run muh/vllm_bi100_patch.py --conservative
|
|||
|
|
return 49152"""
|
|||
|
|
|
|||
|
|
if old_smem in content:
|
|||
|
|
content = content.replace(old_smem, new_smem)
|
|||
|
|
return content, True
|
|||
|
|
return content, False
|
|||
|
|
|
|||
|
|
# --- Patch 4: prefix_prefill.py ---
|
|||
|
|
|
|||
|
|
@patch("prefix_prefill.py",
|
|||
|
|
"Refine BI-V100 block config with CCCL scan tuning data")
|
|||
|
|
def patch_prefix_prefill(content, args):
|
|||
|
|
# Current code already has BI-V100 detection. Enhance it with
|
|||
|
|
# type-dispatched values from CCCL tuning analysis.
|
|||
|
|
|
|||
|
|
old_block = """ _is_bi_v100 = not current_platform.has_device_capability(80)
|
|||
|
|
if _is_bi_v100:
|
|||
|
|
BLOCK = 64
|
|||
|
|
NUM_WARPS = 4"""
|
|||
|
|
|
|||
|
|
# CCCL scan tuning for BI-V100 (SM=16):
|
|||
|
|
# value_size=2 (fp16): tpb=512 → 16 warps → but Triton uses warps not threads
|
|||
|
|
# SMEM: BLOCK_M * head_dim * 2B(fp16) + BLOCK_N * head_dim * 2B * 2(K+V)
|
|||
|
|
# For head_dim=128, fp16:
|
|||
|
|
# BLOCK=64, BLOCK_N=64: 64*128*2 + 64*128*2*2 = 16KB + 32KB = 48KB ← tight!
|
|||
|
|
# BLOCK=64, BLOCK_N=32: 64*128*2 + 32*128*2*2 = 16KB + 16KB = 32KB ← safe
|
|||
|
|
# BLOCK=32, BLOCK_N=64: 32*128*2 + 64*128*2*2 = 8KB + 32KB = 40KB ← ok
|
|||
|
|
|
|||
|
|
new_block = """ _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 # 4 warps × 32 = 128 threads; BW-limited at 56 GB/s/SM"""
|
|||
|
|
|
|||
|
|
if old_block in content:
|
|||
|
|
content = content.replace(old_block, new_block)
|
|||
|
|
return content, True
|
|||
|
|
return content, False
|
|||
|
|
|
|||
|
|
# --- Patch 5: triton_flash_attention.py ---
|
|||
|
|
|
|||
|
|
@patch("vllm/attention/ops/triton_flash_attention.py",
|
|||
|
|
"Add BI-V100 configs to autotune (SM=16, favor smaller blocks)")
|
|||
|
|
def patch_triton_flash(content, args):
|
|||
|
|
# Current configs are AMD ROCm oriented (waves_per_eu is AMD-specific).
|
|||
|
|
# On BI-V100 (Iluvatar, not AMD), waves_per_eu may be ignored.
|
|||
|
|
# Add configs with smaller BLOCK_M that work better with SM=16.
|
|||
|
|
|
|||
|
|
# Find the last config before the closing ], and add BI-V100 configs
|
|||
|
|
insert_before = """ triton.Config(
|
|||
|
|
{
|
|||
|
|
"BLOCK_M": 16,
|
|||
|
|
"BLOCK_N": 16,
|
|||
|
|
"waves_per_eu": 1,
|
|||
|
|
"PRE_LOAD_V": False,
|
|||
|
|
},
|
|||
|
|
num_stages=1,
|
|||
|
|
num_warps=4,
|
|||
|
|
),
|
|||
|
|
],"""
|
|||
|
|
|
|||
|
|
bi100_configs = """ triton.Config(
|
|||
|
|
{
|
|||
|
|
"BLOCK_M": 16,
|
|||
|
|
"BLOCK_N": 16,
|
|||
|
|
"waves_per_eu": 1,
|
|||
|
|
"PRE_LOAD_V": False,
|
|||
|
|
},
|
|||
|
|
num_stages=1,
|
|||
|
|
num_warps=4,
|
|||
|
|
),
|
|||
|
|
# muh: BI-V100 configs (SM=16, 48KB SMEM, 900GB/s BW)
|
|||
|
|
# SM=16 → fewer CTAs → favor configs with moderate BLOCK_M
|
|||
|
|
# to maintain occupancy without excessive SMEM per CTA.
|
|||
|
|
triton.Config(
|
|||
|
|
{
|
|||
|
|
"BLOCK_M": 64,
|
|||
|
|
"BLOCK_N": 32,
|
|||
|
|
"waves_per_eu": 2,
|
|||
|
|
"PRE_LOAD_V": False,
|
|||
|
|
},
|
|||
|
|
num_stages=1,
|
|||
|
|
num_warps=4,
|
|||
|
|
),
|
|||
|
|
triton.Config(
|
|||
|
|
{
|
|||
|
|
"BLOCK_M": 32,
|
|||
|
|
"BLOCK_N": 64,
|
|||
|
|
"waves_per_eu": 2,
|
|||
|
|
"PRE_LOAD_V": False,
|
|||
|
|
},
|
|||
|
|
num_stages=1,
|
|||
|
|
num_warps=4,
|
|||
|
|
),
|
|||
|
|
],"""
|
|||
|
|
|
|||
|
|
if insert_before in content:
|
|||
|
|
content = content.replace(insert_before, bi100_configs)
|
|||
|
|
return content, True
|
|||
|
|
return content, False
|
|||
|
|
|
|||
|
|
# ============================================================
|
|||
|
|
# Backup and apply
|
|||
|
|
# ============================================================
|
|||
|
|
|
|||
|
|
def backup_file(filepath):
|
|||
|
|
"""Create .bak backup before modifying."""
|
|||
|
|
bak = filepath + ".muh_bak"
|
|||
|
|
if not os.path.exists(bak):
|
|||
|
|
shutil.copy2(filepath, bak)
|
|||
|
|
return bak
|
|||
|
|
|
|||
|
|
def revert_file(filepath):
|
|||
|
|
"""Revert from .bak backup."""
|
|||
|
|
bak = filepath + ".muh_bak"
|
|||
|
|
if os.path.exists(bak):
|
|||
|
|
shutil.copy2(bak, filepath)
|
|||
|
|
os.remove(bak)
|
|||
|
|
return True
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
def apply_all(args):
|
|||
|
|
"""Apply all patches."""
|
|||
|
|
results = []
|
|||
|
|
|
|||
|
|
for p in PATCHES:
|
|||
|
|
filepath = p["file"]
|
|||
|
|
if not os.path.exists(filepath):
|
|||
|
|
results.append((filepath, p["description"], "SKIP (file not found)"))
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
with open(filepath, "r") as f:
|
|||
|
|
content = f.read()
|
|||
|
|
|
|||
|
|
new_content, changed = p["apply"](content, args)
|
|||
|
|
|
|||
|
|
if changed:
|
|||
|
|
if args.dry_run:
|
|||
|
|
results.append((filepath, p["description"], "WOULD CHANGE"))
|
|||
|
|
else:
|
|||
|
|
backup_file(filepath)
|
|||
|
|
with open(filepath, "w") as f:
|
|||
|
|
f.write(new_content)
|
|||
|
|
results.append((filepath, p["description"], "APPLIED ✓"))
|
|||
|
|
else:
|
|||
|
|
results.append((filepath, p["description"], "NO MATCH (already patched or different version)"))
|
|||
|
|
|
|||
|
|
return results
|
|||
|
|
|
|||
|
|
def revert_all():
|
|||
|
|
"""Revert all patches."""
|
|||
|
|
results = []
|
|||
|
|
for p in PATCHES:
|
|||
|
|
filepath = p["file"]
|
|||
|
|
if revert_file(filepath):
|
|||
|
|
results.append((filepath, "REVERTED ✓"))
|
|||
|
|
else:
|
|||
|
|
results.append((filepath, "NO BACKUP"))
|
|||
|
|
return results
|
|||
|
|
|
|||
|
|
# ============================================================
|
|||
|
|
# Main
|
|||
|
|
# ============================================================
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
p = argparse.ArgumentParser(description="Apply CCCL-derived BI-V100 tuning to vllm")
|
|||
|
|
p.add_argument("--dry-run", action="store_true", help="Show changes without applying")
|
|||
|
|
p.add_argument("--conservative", action="store_true",
|
|||
|
|
help="Keep SMEM=32KB (safer, pending hardware confirmation)")
|
|||
|
|
p.add_argument("--revert", action="store_true", help="Undo all patches")
|
|||
|
|
args = p.parse_args()
|
|||
|
|
|
|||
|
|
print(f"muh vllm_bi100_patch — {datetime.now().strftime('%Y-%m-%d %H:%M')}")
|
|||
|
|
print(f" Hardware: BI-V100 (SM=16, SMEM={'32KB(conservative)' if args.conservative else '48KB'}, BW=900GB/s)")
|
|||
|
|
print(f" Mode: {'DRY RUN' if args.dry_run else 'REVERT' if args.revert else 'APPLY'}")
|
|||
|
|
print()
|
|||
|
|
|
|||
|
|
if args.revert:
|
|||
|
|
results = revert_all()
|
|||
|
|
for filepath, status in results:
|
|||
|
|
print(f" {status:20s} {filepath}")
|
|||
|
|
else:
|
|||
|
|
results = apply_all(args)
|
|||
|
|
for filepath, desc, status in results:
|
|||
|
|
print(f" {status:40s} {filepath}")
|
|||
|
|
print(f" {'':40s} └ {desc}")
|
|||
|
|
|
|||
|
|
print()
|
|||
|
|
if not args.dry_run and not args.revert:
|
|||
|
|
print("Done. To revert: python3 muh/vllm_bi100_patch.py --revert")
|
|||
|
|
print("To test: python3 -c \"from vllm._custom_ops import get_max_shared_memory_per_block_device_attribute; print(get_max_shared_memory_per_block_device_attribute(0))\"")
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
main()
|