diff --git a/Dockerfile b/Dockerfile index a064b65d..dae3bf39 100644 --- a/Dockerfile +++ b/Dockerfile @@ -3,10 +3,12 @@ FROM git.modelhub.org.cn:9443/enginex-iluvatar/bi100-3.2.3-x86-ubuntu20.04-py3.1 RUN mkdir /workspace WORKDIR /workspace/ -# Copy all scripts and the V2 module +# Copy all scripts, V2 kernels, CCCL-tuned prefill, and muh dispatch COPY ./qwen3_6_scripts /workspace/qwen3_6_scripts COPY ./paged_attention_v2_pytorch.py /workspace/paged_attention_v2_pytorch.py COPY ./paged_attention_v2_triton.py /workspace/paged_attention_v2_triton.py +COPY ./prefix_prefill.py /workspace/prefix_prefill.py +COPY ./muh_dispatch.py /workspace/muh_dispatch.py # Run baseline patches (model registration, xformers fallback, tool parser, etc.) RUN cd ./qwen3_6_scripts && ./patch_ops.sh @@ -23,12 +25,25 @@ RUN python3 /workspace/qwen3_6_scripts/patch_ixformer_native.py # Triton V2 risk: SMEM=32KB zero margin at head_dim=256 BLOCK_N=32. # If Triton V2 crashes, PyTorch V2 (batched bmm, no intermediate tensor # savings but correct) takes over automatically via try/except. +# Deploy Triton V2 kernel into vllm package RUN cp /workspace/paged_attention_v2_triton.py \ /usr/local/corex/lib/python3/dist-packages/vllm/paged_attention_v2_triton.py 2>/dev/null || \ cp /workspace/paged_attention_v2_triton.py \ /usr/local/corex/lib64/python3/dist-packages/vllm/paged_attention_v2_triton.py 2>/dev/null || true RUN python3 /workspace/qwen3_6_scripts/patch_paged_attention_v2.py +# Deploy CCCL-tuned prefix_prefill.py (SM=16: BLOCK=64, NUM_WARPS=4) +RUN cp /workspace/prefix_prefill.py \ + /usr/local/corex/lib/python3/dist-packages/vllm/attention/ops/prefix_prefill.py 2>/dev/null || \ + cp /workspace/prefix_prefill.py \ + /usr/local/corex/lib64/python3/dist-packages/vllm/attention/ops/prefix_prefill.py 2>/dev/null || true + +# Deploy muh_dispatch.py (CCCL-style type dispatch for kernel configs) +RUN cp /workspace/muh_dispatch.py \ + /usr/local/corex/lib/python3/dist-packages/vllm/muh_dispatch.py 2>/dev/null || \ + cp /workspace/muh_dispatch.py \ + /usr/local/corex/lib64/python3/dist-packages/vllm/muh_dispatch.py 2>/dev/null || true + # 2. Triton kernel tuning: BLOCK=64, NUM_WARPS=4 # SMEM: BLOCK_N=64 × head_dim=128 × 2B × 2(K+V) = 32KB ≤ 48KB # Occupancy: 4 warps allows 2 blocks/SM vs 1 at 8 warps diff --git a/muh_dispatch.py b/muh_dispatch.py index 7323a16b..e59c3745 100644 --- a/muh_dispatch.py +++ b/muh_dispatch.py @@ -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: