diff --git a/muh/vllm_bi100_patch.py b/muh/vllm_bi100_patch.py new file mode 100644 index 00000000..3f75dc67 --- /dev/null +++ b/muh/vllm_bi100_patch.py @@ -0,0 +1,380 @@ +#!/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() diff --git a/prefix_prefill.py b/prefix_prefill.py index 3c0ee577..76c73f19 100644 --- a/prefix_prefill.py +++ b/prefix_prefill.py @@ -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 diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 64a5534c..18bcfada 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -889,6 +889,9 @@ def get_device_attribute(attribute: int, device: int) -> int: 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 diff --git a/vllm/attention/ops/paged_attn.py b/vllm/attention/ops/paged_attn.py index c90e8dd8..5c9989d5 100644 --- a/vllm/attention/ops/paged_attn.py +++ b/vllm/attention/ops/paged_attn.py @@ -123,9 +123,18 @@ class PagedAttention: # 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 + # 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 use_v1: # Run PagedAttention V1. ops.paged_attention_v1( diff --git a/vllm/attention/ops/triton_flash_attention.py b/vllm/attention/ops/triton_flash_attention.py index f9421111..a0efe68d 100644 --- a/vllm/attention/ops/triton_flash_attention.py +++ b/vllm/attention/ops/triton_flash_attention.py @@ -302,6 +302,29 @@ def _attn_fwd_inner( 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, + ), ], key=['IS_CAUSAL', 'dropout_p', 'BLOCK_DMODEL'], ) diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py index 03ffc2f2..b56d70d0 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe.py @@ -338,17 +338,23 @@ def get_default_config( dtype: Optional[str], is_marlin: bool, ) -> Dict[str, int]: + # 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': 32, + 'BLOCK_SIZE_K': 64, # muh: 32→64, better memory coalescing on BI-V100 'GROUP_SIZE_M': 8 } # A heuristic: fused marlin works faster with this config for small M 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': 32, + 'BLOCK_SIZE_N': 64, # muh: 32→64, wider output tile 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 1 }