feat: deploy CCCL-tuned prefix_prefill + muh_dispatch + fix SM=16 count
muh_dispatch.py: - Fix missing os/sys imports (was crashing on import) - Fix SM count 50→16 (confirmed via ixsmi, matches hardware.cuh) - Fix C++ struct name lookup to match actual tuning_reduce.cuh names: bi100_plus_float32_o4, bi100_plus_float64_o4, bi100_plus_accum2_o4 (was: bi100_float32_plus_o4 — wrong name, would always fall through to default) Dockerfile: - Add COPY for prefix_prefill.py and muh_dispatch.py - Deploy CCCL-tuned prefix_prefill.py into vllm attention ops (BLOCK=64, NUM_WARPS=4 for BI-V100 SM=16) - Deploy muh_dispatch.py into vllm package for type-dispatched kernel configs - These files were written but never deployed — dead code until now Impact: prefix_prefill.py deployment means the CCCL-derived block sizes actually take effect at runtime. Previously the base image's original prefix_prefill.py (BLOCK=128 for cc>=80, or 64 for cc<80) was used, which is correct for BI-V100 but our version adds explicit SM=16 documentation and the path for future tuning.
This commit is contained in:
@@ -25,6 +25,8 @@ Deploy: cp muh_dispatch.py /usr/local/corex/.../vllm/muh_dispatch.py
|
||||
Then patch paged_attn.py to import and use it.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import torch
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
@@ -38,7 +40,7 @@ class HardwareCapability:
|
||||
warp_size: int = 32
|
||||
max_threads_per_block: int = 1024
|
||||
max_shared_memory_per_block: int = 49152 # 48KB
|
||||
sm_count: int = 50
|
||||
sm_count: int = 16 # CONFIRMED: ixsmi shows 16 SMs per BI-V100 (NOT 50 from spec)
|
||||
memory_bandwidth_gbps: int = 900
|
||||
l2_cache_size_bytes: int = 6 * 1024 * 1024 # 6MB
|
||||
|
||||
@@ -97,12 +99,14 @@ def _read_reduce_config(accum_size: int) -> dict:
|
||||
from gen_patch import extract_bi100_structs
|
||||
structs = extract_bi100_structs(header_path)
|
||||
|
||||
# Select struct by accum_size
|
||||
# Select struct by accum_size — names match tuning_reduce.cuh
|
||||
target_struct = None
|
||||
if accum_size <= 4:
|
||||
target_struct = "bi100_float32_plus_o4"
|
||||
if accum_size <= 2:
|
||||
target_struct = "bi100_plus_accum2_o4"
|
||||
elif accum_size <= 4:
|
||||
target_struct = "bi100_plus_float32_o4"
|
||||
else:
|
||||
target_struct = "bi100_float64_plus_o4"
|
||||
target_struct = "bi100_plus_float64_o4"
|
||||
|
||||
for name, fields in structs:
|
||||
if name == target_struct:
|
||||
|
||||
Reference in New Issue
Block a user