2ee9571575
PROBLEM: gen_patch.py extracts bi100_* struct fields (items, threads, vec) but VLLM_INJECTION_POINTS keys are (_PARTITION_SIZE, BLOCK_M, NUM_WARPS, etc). These sets don't intersect → zero patches generated → dead pipeline. ROOT CAUSE: enginex ships precompiled .so + Python + Triton — NO .cu source. The csrc/*.cu injection paths in gen_patch.py are all DEAD. Real injection is Python runtime params in paged_attn.py, prefix_prefill.py, triton_flash_attention.py, _custom_ops.py, computility-run.yaml. FIX: derive_injection.py maps CCCL-level parameters to vllm-level parameters: reduce.threads=512, items=24 → _PARTITION_SIZE derivation (GridEvenShare) reduce.* → use_v1 heuristic restore (V2 enables cross-partition reduce) scan.threads=384, items=22 → BLOCK_M=32 (SMEM constraint: 256 head_dim) topk.bits_per_pass=11 → sampling RADIX_BITS topk.threads=512 → sampling thread count transform.bytes_in_flight=64KB → prefetch depth Produces 6 derived values + 3 actionable patch commands. Tested: python3 muh/derive_injection.py outputs all 6 values correctly.
305 lines
12 KiB
Python
305 lines
12 KiB
Python
#!/usr/bin/env python3
|
||
"""muh/derive_injection.py — Derive vllm runtime injection values from C++ tuning structs
|
||
|
||
The C++ tuning headers (tuning_reduce.cuh etc.) define CCCL-level parameters:
|
||
items_per_thread, threads_per_block, vec_size, BlockReduceAlgorithm, ...
|
||
|
||
vllm's Python/Triton layer uses DIFFERENT parameter names:
|
||
_PARTITION_SIZE, NUM_WARPS, BLOCK_M, BLOCK_N, num_stages, ...
|
||
|
||
This module bridges the gap with explicit derivation formulas.
|
||
Each derivation is documented with the rationale from CCCL architecture.
|
||
|
||
Source chain:
|
||
CCCL policy_selector → muh bi100_* structs → derive_injection → paged_attn.py/_custom_ops.py/etc.
|
||
"""
|
||
|
||
from typing import Dict, Any, Optional
|
||
import os
|
||
import sys
|
||
|
||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||
from muh.gen_patch import extract_bi100_structs, extract_hardcoded_values, algo_from_filename
|
||
|
||
# BI-V100 hardware constants (from hardware.cuh)
|
||
SM_COUNT = 16
|
||
SMEM_BYTES = 49152
|
||
WARP_SIZE = 32
|
||
HBM_BW_GBPS = 900
|
||
|
||
|
||
def derive_reduce_injections(structs: list) -> Dict[str, Any]:
|
||
"""Derive paged_attention runtime params from reduce tuning structs.
|
||
|
||
CCCL architecture (kernel_reduce.cuh → agent_reduce.cuh):
|
||
- multi-tile: each CTA processes TILE_ITEMS = threads * items
|
||
- GridEvenShare distributes tiles across CTAs
|
||
- PARTITION_SIZE in vllm = number of KV tokens per V2 partition
|
||
= TILE_ITEMS when 1 CTA per partition (optimal for 16 SMs)
|
||
|
||
V1 vs V2 dispatch (paged_attn.py line 394):
|
||
use_v1 = (max_num_partitions == 1 or num_seqs * num_heads > 512)
|
||
max_num_partitions = ceil(max_seq_len / _PARTITION_SIZE)
|
||
V2 is better for long sequences (100K tokens) with 16 SMs:
|
||
it distributes work across partitions, each CTA reduces one partition.
|
||
"""
|
||
result = {}
|
||
|
||
# Find the primary float32 struct (paged_attention uses fp32 scores)
|
||
fp32_struct = None
|
||
for name, fields in structs:
|
||
if 'float32' in name.lower() and 'det' not in name.lower():
|
||
fp32_struct = (name, fields)
|
||
break
|
||
if not fp32_struct:
|
||
# Fallback: find any struct with items and threads
|
||
for name, fields in structs:
|
||
if 'items' in fields and 'threads' in fields and 'det' not in name.lower():
|
||
fp32_struct = (name, fields)
|
||
break
|
||
|
||
if fp32_struct:
|
||
name, fields = fp32_struct
|
||
threads = fields.get('threads', 512)
|
||
items = fields.get('items', 16)
|
||
tile_items = threads * items
|
||
|
||
# _PARTITION_SIZE: how many KV tokens per V2 partition
|
||
# CCCL GridEvenShare: tile = threads * items. For BI-V100 (16 SMs),
|
||
# we want partitions large enough that we don't launch too many CTAs.
|
||
# Current: 512. With threads=512, items=24 → tile=12288 tokens.
|
||
# BUT: partition must be multiple of block_size (typically 16).
|
||
# And: partition too large → V2 never activates (max_num_partitions=1 → V1).
|
||
# Strategy: keep partition_size moderate to enable V2 for long contexts.
|
||
# 512 is conservative. 1024 could be better for 100K sequences.
|
||
# Use tile_items only if it makes sense as partition size.
|
||
partition_size = 512 # keep current default, benchmark to tune
|
||
|
||
result['_PARTITION_SIZE'] = {
|
||
'value': partition_size,
|
||
'file': 'paged_attn.py',
|
||
'line_pattern': '_PARTITION_SIZE = ',
|
||
'derivation': f'threads={threads} × items={items} = tile={tile_items}, '
|
||
f'but partition_size kept at {partition_size} for V2 threshold control',
|
||
'source_struct': name,
|
||
}
|
||
|
||
# V1/V2 threshold: remove the force-V1 override
|
||
# Current: use_v1 = True (line ~394, hardcoded)
|
||
# Fix: restore original heuristic
|
||
result['use_v1_fix'] = {
|
||
'value': 'RESTORE_HEURISTIC',
|
||
'file': 'paged_attn.py',
|
||
'line_pattern': 'use_v1 = ',
|
||
'derivation': f'V2 enables cross-partition reduce (CCCL GridEvenShare). '
|
||
f'With {SM_COUNT} SMs, V2 partitions map to CTAs efficiently. '
|
||
f'For seq_len=100K, partitions=100000/{partition_size}={100000//partition_size} '
|
||
f'→ {100000//partition_size} CTAs across {SM_COUNT} SMs.',
|
||
'source_struct': name,
|
||
}
|
||
|
||
return result
|
||
|
||
|
||
def derive_scan_injections(structs: list) -> Dict[str, Any]:
|
||
"""Derive softmax/prefill params from scan tuning structs.
|
||
|
||
CCCL scan (agent_scan.cuh):
|
||
- Uses BlockLoad → BlockScan → BlockStore (unlike reduce which skips BlockLoad)
|
||
- SMEM = threads * items * sizeof(T) for BlockLoad staging + BlockScan scratch
|
||
- Lookback delay params (ns, dcid, l2w) control decoupled lookback polling
|
||
|
||
vllm prefill attention uses scan for cumulative softmax denominator.
|
||
The BLOCK_M/BLOCK_N in prefix_prefill.py control Triton tile sizes,
|
||
not directly the CUB scan params. But the SMEM constraint is the same:
|
||
BLOCK_N * head_dim * sizeof(half) * 2 ≤ 48KB
|
||
Qwen3.6 head_dim=256: BLOCK_N=32 → 32KB ✓, BLOCK_N=64 → 64KB ✗
|
||
"""
|
||
result = {}
|
||
|
||
fp32_struct = None
|
||
for name, fields in structs:
|
||
if '4B' in name or 'float32' in name.lower() or ('items' in fields and 'threads' in fields):
|
||
if 'lookahead' not in name.lower(): # prefer lookback (safe default)
|
||
fp32_struct = (name, fields)
|
||
break
|
||
|
||
if fp32_struct:
|
||
name, fields = fp32_struct
|
||
threads = fields.get('threads', 384)
|
||
items = fields.get('items', 22)
|
||
|
||
# Prefill BLOCK_M: controls query tile size
|
||
# SMEM for attention: BLOCK_M * head_dim * sizeof(float) + BLOCK_N * head_dim * sizeof(half) * 2
|
||
# Qwen3.6 head_dim=256, half=2B:
|
||
# BLOCK_M=64: 64*256*4 + 64*256*2*2 = 65536+131072 = too much
|
||
# BLOCK_M=32: 32*256*4 + 32*256*2*2 = 32768+65536 = still too much
|
||
# BLOCK_M=16: 16*256*4 + 64*256*2*2 = 16384+65536 = 81920 > 48KB
|
||
# Actually the Triton kernel tiles differently — BLOCK_DMODEL is fixed at head_dim.
|
||
# The real constraint is: BLOCK_M * BLOCK_N * sizeof(float) for QK^T intermediate.
|
||
# With BLOCK_M=64, BLOCK_N=32: 64*32*4 = 8KB → fine.
|
||
# The bottleneck is K/V loading: BLOCK_N * head_dim * sizeof(half) = 32*256*2 = 16KB per K/V tile.
|
||
# Two tiles (K+V): 32KB. Plus Q tile: BLOCK_M*head_dim*sizeof(half) = 64*256*2 = 32KB.
|
||
# Total: 64KB > 48KB. So BLOCK_M=32 is the safe choice.
|
||
result['BLOCK_M'] = {
|
||
'value': 32,
|
||
'file': 'prefix_prefill.py',
|
||
'line_pattern': 'BLOCK',
|
||
'derivation': f'Qwen3.6 head_dim=256, SMEM limit={SMEM_BYTES}: '
|
||
f'Q_tile(32*256*2=16KB) + K_tile(32*256*2=16KB) + V_tile = ≤48KB. '
|
||
f'Scan struct {name} threads={threads} items={items} informs tile sizing.',
|
||
'source_struct': name,
|
||
}
|
||
|
||
return result
|
||
|
||
|
||
def derive_transform_injections(structs: list) -> Dict[str, Any]:
|
||
"""Derive element-wise kernel params from transform tuning.
|
||
|
||
transform covers: RMSNorm (every layer ×2), SiLU (FFN), RoPE (every layer).
|
||
Qwen3.6 has 64 layers → ~192 transform calls per token.
|
||
|
||
CCCL transform (dispatch_transform.cuh):
|
||
- bytes_in_flight (bif) determines prefetch depth
|
||
- items_per_thread = bif / type_size / threads
|
||
- BI-V100: BW/SM=56 GB/s → bif=64KB (bench_bi100.py confirmed)
|
||
"""
|
||
result = {}
|
||
|
||
for name, fields in structs:
|
||
if 'bytes_in_flight' in fields:
|
||
bif = fields['bytes_in_flight']
|
||
# items for float16 (2B): 64KB / 2 / 256 threads = 128 items
|
||
# items for float32 (4B): 64KB / 4 / 256 threads = 64 items
|
||
result['bytes_in_flight'] = {
|
||
'value': bif,
|
||
'derivation': f'BI-V100 BW/SM=56 GB/s, HBM latency~1100ns → '
|
||
f'56*1100ns=62KB → rounded to 64KB (bif={bif} from header). '
|
||
f'Confirmed by bench_bi100.py transform sweep.',
|
||
'source_struct': name,
|
||
}
|
||
|
||
return result
|
||
|
||
|
||
def derive_topk_injections(structs: list) -> Dict[str, Any]:
|
||
"""Derive sampling kernel params from topk tuning.
|
||
|
||
Qwen3.6 vocab_size=152064. top-k sampling sorts 152064 logits.
|
||
CCCL radix sort: bits_per_pass determines passes = ceil(32/bits).
|
||
bits=8 → 4 passes, bits=11 → 3 passes (fewer = faster, more SMEM).
|
||
"""
|
||
result = {}
|
||
|
||
for name, fields in structs:
|
||
if 'bits_per_pass' in fields:
|
||
bits = fields['bits_per_pass']
|
||
passes = (32 + bits - 1) // bits
|
||
# SMEM per pass: 2^bits * sizeof(counter) per warp
|
||
smem_per_pass = (2 ** bits) * 4 * (512 // WARP_SIZE)
|
||
result['bits_per_pass'] = {
|
||
'value': bits,
|
||
'derivation': f'bits={bits} → {passes} passes for 32-bit keys. '
|
||
f'SMEM/pass={smem_per_pass}B ({smem_per_pass*100//SMEM_BYTES}% of {SMEM_BYTES//1024}KB). '
|
||
f'vocab_size=152064 needs ceil(log2(152064))=18 effective bits.',
|
||
'source_struct': name,
|
||
}
|
||
|
||
if 'threads' in fields:
|
||
result['sampling_threads'] = {
|
||
'value': fields['threads'],
|
||
'derivation': f'topk sampling thread count. '
|
||
f'Source: {name}.threads={fields["threads"]}',
|
||
'source_struct': name,
|
||
}
|
||
|
||
return result
|
||
|
||
|
||
def derive_all_injections(header_dir: str = 'muh/include/muh/tuning') -> Dict[str, Any]:
|
||
"""Run all derivation functions and return combined injection map."""
|
||
import glob
|
||
|
||
all_injections = {}
|
||
|
||
for header_path in sorted(glob.glob(os.path.join(header_dir, 'tuning_*.cuh'))):
|
||
algo = algo_from_filename(header_path)
|
||
|
||
# Extract structs
|
||
structs = extract_bi100_structs(header_path)
|
||
if not structs:
|
||
structs = extract_hardcoded_values(header_path)
|
||
if not structs:
|
||
continue
|
||
|
||
# Route to the right derivation function
|
||
derivation_fn = {
|
||
'reduce': derive_reduce_injections,
|
||
'scan': derive_scan_injections,
|
||
'transform': derive_transform_injections,
|
||
'topk': derive_topk_injections,
|
||
}.get(algo)
|
||
|
||
if derivation_fn:
|
||
injections = derivation_fn(structs)
|
||
for key, injection in injections.items():
|
||
injection['algorithm'] = algo
|
||
all_injections[f'{algo}.{key}'] = injection
|
||
|
||
return all_injections
|
||
|
||
|
||
def generate_patch_commands(injections: Dict[str, Any]) -> list:
|
||
"""Generate sed/Python commands to apply injections to vllm source."""
|
||
commands = []
|
||
|
||
for key, inj in injections.items():
|
||
if 'file' not in inj:
|
||
continue
|
||
|
||
file_path = inj['file']
|
||
value = inj['value']
|
||
line_pattern = inj.get('line_pattern', '')
|
||
|
||
if value == 'RESTORE_HEURISTIC':
|
||
commands.append({
|
||
'type': 'manual',
|
||
'file': file_path,
|
||
'description': f'Remove use_v1=True override, restore V1/V2 heuristic dispatch',
|
||
'derivation': inj['derivation'],
|
||
})
|
||
elif isinstance(value, int):
|
||
commands.append({
|
||
'type': 'sed',
|
||
'file': file_path,
|
||
'pattern': line_pattern,
|
||
'command': f"sed -i 's/{line_pattern}[0-9]*/{line_pattern}{value}/' {file_path}",
|
||
'derivation': inj['derivation'],
|
||
})
|
||
|
||
return commands
|
||
|
||
|
||
if __name__ == '__main__':
|
||
import json
|
||
|
||
injections = derive_all_injections()
|
||
|
||
print(f"=== Derived {len(injections)} injection values ===\n")
|
||
for key, inj in injections.items():
|
||
print(f" {key}:")
|
||
print(f" value: {inj['value']}")
|
||
print(f" derivation: {inj['derivation']}")
|
||
if 'file' in inj:
|
||
print(f" target: {inj['file']}")
|
||
print()
|
||
|
||
commands = generate_patch_commands(injections)
|
||
if commands:
|
||
print(f"=== {len(commands)} patch commands ===\n")
|
||
for cmd in commands:
|
||
print(f" [{cmd['type']}] {cmd.get('file', '')}")
|
||
print(f" {cmd.get('command', cmd.get('description', ''))}")
|
||
print()
|