ef6abf3dc7
Adds ALL files needed for Dockerfile build:
- qwen3_6_scripts/ (baseline patches + our optimizations)
- vllm/ (full vllm package)
- paged_attention_v2_pytorch.py (V2 with single-bmm optimization)
- Dockerfile + computility-run.yaml
Our optimizations vs baseline:
1. paged_attn.py: pre-gathered context KV (eliminates 194 gather calls),
Triton try/fallback, V2 heuristic, threshold 32K→64K
2. paged_attention_v2_pytorch.py: fills NotImplementedError,
single-bmm Phase 1 (195 launches → 3)
3. patch_enable_triton.py: HAS_TRITON=True with safety fallback
4. patch_triton_tuning.py: BLOCK=64, NUM_WARPS=4 for BI-V100
5. computility-run.yaml: gpu-memory-utilization 0.9→0.95,
max-num-batched-tokens 8192→16384
This repo can now be submitted to dev.modelhub.org.cn as-is.
47 lines
1.5 KiB
Python
47 lines
1.5 KiB
Python
import functools
|
|
from typing import Dict
|
|
|
|
|
|
@functools.lru_cache
|
|
def _get_op_configs(op_type: str, batch: int, hidden_size: int):
|
|
# TODO: add optimal configurations
|
|
return None
|
|
|
|
|
|
def _check_divisibility(hidden_size: int):
|
|
# The bgmv_expand kernel requires that the hidden_size be divisible by
|
|
# the number below.
|
|
divisibility = [2, 4, 8, 16, 32, 64]
|
|
divisibility.sort(reverse=True)
|
|
for div in divisibility:
|
|
if hidden_size % div == 0:
|
|
return div
|
|
# hidden_size is an odd number
|
|
return 1
|
|
|
|
|
|
def _get_default_config(op_type: str, batch: int, hidden_size: int):
|
|
if op_type == "expand":
|
|
return {
|
|
"BLOCK_N": 256,
|
|
"SPLIT_N": _check_divisibility(hidden_size),
|
|
"num_warps": 8
|
|
}
|
|
else:
|
|
return {"BLOCK_K": 256, "SPLIT_K": 64, "num_warps": 8}
|
|
|
|
|
|
def get_lora_op_configs(op_type: str, batch: int,
|
|
hidden_size: int) -> Dict[str, int]:
|
|
"""Inspired by `fused_moe_kernel`
|
|
The return value will be a dictionary mapping an irregular grid of batch
|
|
sizes and hidden_size to configurations of the bgmv-related kernel.
|
|
NOTE: It currently only supports the default configuration. We plan to
|
|
generate optimal configurations for different hardware in the future using
|
|
scripts similar to `benchmark_moe.py`.
|
|
"""
|
|
config = _get_op_configs(op_type, batch, hidden_size)
|
|
if not config:
|
|
config = _get_default_config(op_type, batch, hidden_size)
|
|
return config
|