a2a5dd8f00
bench_triton_prefill.py:
- Split --block into --block (BLOCK_M) and --block-n (BLOCK_N)
- Each (M, N, warps) combo triggers Triton JIT recompilation
- Enables finding asymmetric optima like M=64,N=32 that save SMEM
triton_flash_attention.py:
- Re-add 3 BI-V100 autotune configs (64x32, 32x64, 64x64 with warps=4)
- These were wrongly reverted in 8c1955d -- autotune is zero-risk
run_on_bi100.sh:
- Updated to use asymmetric block search
337 lines
15 KiB
Python
337 lines
15 KiB
Python
#!/usr/bin/env python3
|
||
"""muh/bench_triton_prefill.py — Real Triton kernel benchmark for BI-V100
|
||
|
||
Unlike bench_bi100.py (which called torch.sum and never injected params),
|
||
this script directly invokes the prefix_prefill Triton JIT kernel with
|
||
different (BLOCK, NUM_WARPS) constexpr values. Each combination triggers
|
||
Triton to compile a separate kernel binary — the same mechanism as CCCL's
|
||
`#define TUNE_THREADS_PER_BLOCK N` + recompile.
|
||
|
||
Search space (from CCCL scan analogy):
|
||
BLOCK: [16, 32, 64, 128] — tile size for Q/K/V
|
||
NUM_WARPS: [1, 2, 4, 8] — warps per CTA (threads = warps × 32)
|
||
|
||
Valid combos after SMEM pruning (head_dim=128, fp16):
|
||
SMEM ≈ BLOCK × 128 × 2B × 3(Q+K+V) + BLOCK × BLOCK × 4B(fp32 accum)
|
||
BLOCK=128: SMEM ≈ 128×128×6 + 128×128×4 = 98304 + 65536 = 163840 → OVERFLOW
|
||
BLOCK=64: SMEM ≈ 64×128×6 + 64×64×4 = 49152 + 16384 = 65536 → OVERFLOW at 48KB
|
||
BLOCK=32: SMEM ≈ 32×128×6 + 32×32×4 = 24576 + 4096 = 28672 → FITS
|
||
BLOCK=16: SMEM ≈ 16×128×6 + 16×16×4 = 12288 + 1024 = 13312 → FITS
|
||
|
||
NOTE: Triton's actual SMEM usage differs from this estimate because:
|
||
- Triton uses register tiling, not full SMEM staging
|
||
- Accumulator is in registers, not SMEM
|
||
- K/V cache loads may be streamed, not fully staged
|
||
The real constraint is determined by Triton compiler. Combos that exceed
|
||
SMEM will fail at compile time with a clear error — not silently.
|
||
|
||
Output format (CCCL compatible):
|
||
block_64.warps_4 <speedup_ctx128> <speedup_ctx512> <speedup_ctx2048> <speedup_ctx8192>
|
||
|
||
Usage (ON BI-V100 ONLY — requires GPU):
|
||
python3 muh/bench_triton_prefill.py
|
||
python3 muh/bench_triton_prefill.py --block 32 64 --warps 2 4 8
|
||
python3 muh/bench_triton_prefill.py --output results/prefill_bench.json
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import time
|
||
import json
|
||
import argparse
|
||
from typing import List, Dict, Tuple, Optional
|
||
|
||
# ============================================================
|
||
# This script MUST run on BI-V100. Import torch only at runtime.
|
||
# ============================================================
|
||
|
||
def require_gpu():
|
||
try:
|
||
import torch
|
||
if not torch.cuda.is_available():
|
||
print("ERROR: No CUDA GPU available. This benchmark must run on BI-V100.")
|
||
sys.exit(1)
|
||
dev = torch.cuda.get_device_name(0)
|
||
props = torch.cuda.get_device_properties(0)
|
||
print(f"GPU: {dev}")
|
||
print(f" SMs: {props.multi_processor_count}")
|
||
print(f" SMEM/block: {props.max_shared_memory_per_block} bytes")
|
||
print(f" Total VRAM: {props.total_mem / 1024**3:.1f} GB")
|
||
return torch, props
|
||
except ImportError:
|
||
print("ERROR: torch not available. Run on BI-V100 with vllm environment.")
|
||
sys.exit(1)
|
||
|
||
|
||
def make_test_tensors(torch, batch: int, seq_len: int, ctx_len: int,
|
||
num_heads: int, num_kv_heads: int, head_dim: int,
|
||
block_size: int, dtype):
|
||
"""Create realistic test tensors matching Qwen3.6 dimensions.
|
||
|
||
Qwen3.6-35B-A3B:
|
||
head_dim = 128
|
||
num_heads = 64 (query heads)
|
||
num_kv_heads = 8 (GQA 8:1)
|
||
dtype = bfloat16 or float16
|
||
"""
|
||
device = "cuda"
|
||
|
||
# Query: [total_tokens, num_heads, head_dim]
|
||
total_tokens = batch * seq_len
|
||
q = torch.randn(total_tokens, num_heads, head_dim, device=device, dtype=dtype)
|
||
|
||
# K, V: [total_tokens, num_kv_heads, head_dim]
|
||
k = torch.randn(total_tokens, num_kv_heads, head_dim, device=device, dtype=dtype)
|
||
v = torch.randn(total_tokens, num_kv_heads, head_dim, device=device, dtype=dtype)
|
||
|
||
# Output
|
||
o = torch.empty_like(q)
|
||
|
||
# KV cache: [num_blocks, num_kv_heads, head_dim/x, block_size, x]
|
||
x = 16 // dtype_size(dtype) # vector width
|
||
num_blocks = (batch * (ctx_len + seq_len) + block_size - 1) // block_size + 16
|
||
k_cache = torch.randn(num_blocks, num_kv_heads, head_dim // x, block_size, x,
|
||
device=device, dtype=dtype)
|
||
v_cache = torch.randn(num_blocks, num_kv_heads, head_dim, block_size,
|
||
device=device, dtype=dtype)
|
||
|
||
# Block location table: [batch, max_blocks_per_seq]
|
||
max_blocks = (ctx_len + seq_len + block_size - 1) // block_size
|
||
b_loc = torch.zeros(batch, max_blocks, device=device, dtype=torch.int32)
|
||
for i in range(batch):
|
||
n_blocks = min(max_blocks, num_blocks - i * max_blocks)
|
||
b_loc[i, :n_blocks] = torch.arange(
|
||
i * max_blocks, i * max_blocks + n_blocks, device=device)
|
||
|
||
# Sequence metadata
|
||
b_start_loc = torch.arange(0, batch * seq_len, seq_len, device=device, dtype=torch.int32)
|
||
b_seq_len = torch.full((batch,), seq_len + ctx_len, device=device, dtype=torch.int32)
|
||
b_ctx_len = torch.full((batch,), ctx_len, device=device, dtype=torch.int32)
|
||
|
||
return q, k, v, o, k_cache, v_cache, b_loc, b_start_loc, b_seq_len, b_ctx_len
|
||
|
||
|
||
def dtype_size(dtype):
|
||
import torch
|
||
return torch.tensor([], dtype=dtype).element_size()
|
||
|
||
|
||
def bench_one_config(torch, triton, _fwd_kernel, block_m: int, block_n: int,
|
||
num_warps: int,
|
||
q, k, v, o, k_cache, v_cache, b_loc, b_start_loc,
|
||
b_seq_len, b_ctx_len, max_input_len: int,
|
||
warmup: int = 3, repeats: int = 10) -> Optional[float]:
|
||
"""Benchmark one (BLOCK, NUM_WARPS) configuration.
|
||
|
||
Each call triggers Triton JIT compilation for this specific
|
||
(BLOCK, NUM_WARPS) pair if not already cached. This IS the
|
||
compile-time parameter injection — same as CCCL's #define mechanism.
|
||
|
||
Returns median time in ms, or None if compilation fails (SMEM overflow).
|
||
"""
|
||
Lk = q.shape[-1]
|
||
Lk_padded = triton.next_power_of_2(Lk)
|
||
sm_scale = 1.0 / (Lk ** 0.5)
|
||
batch = b_seq_len.shape[0]
|
||
head = q.shape[1]
|
||
num_queries_per_kv = q.shape[1] // k.shape[1]
|
||
|
||
grid = (batch, head, triton.cdiv(max_input_len, block_m))
|
||
|
||
# Build kernel call args (reused for warmup and timed runs)
|
||
def call_kernel():
|
||
_fwd_kernel[grid](
|
||
q, k, v, k_cache, v_cache, b_loc,
|
||
sm_scale, 1.0, 1.0,
|
||
b_start_loc, b_seq_len, b_ctx_len,
|
||
v_cache.shape[3], k_cache.shape[4],
|
||
o,
|
||
b_loc.stride(0), b_loc.stride(1),
|
||
q.stride(0), q.stride(1), q.stride(2),
|
||
k.stride(0), k.stride(1), k.stride(2),
|
||
v.stride(0), v.stride(1), v.stride(2),
|
||
o.stride(0), o.stride(1), o.stride(2),
|
||
k_cache.stride(0), k_cache.stride(1), k_cache.stride(2),
|
||
k_cache.stride(3), k_cache.stride(4),
|
||
v_cache.stride(0), v_cache.stride(1), v_cache.stride(2),
|
||
v_cache.stride(3),
|
||
num_queries_per_kv=num_queries_per_kv,
|
||
BLOCK_M=block_m,
|
||
BLOCK_DMODEL=Lk,
|
||
BLOCK_DMODEL_PADDED=Lk_padded,
|
||
BLOCK_N=block_n,
|
||
SLIDING_WINDOW=0,
|
||
num_warps=num_warps,
|
||
num_stages=1,
|
||
)
|
||
|
||
# Attempt compilation + warmup
|
||
try:
|
||
for _ in range(warmup):
|
||
call_kernel()
|
||
torch.cuda.synchronize()
|
||
except Exception as e:
|
||
return None # Compilation failed (likely SMEM overflow)
|
||
|
||
# Timed runs
|
||
times = []
|
||
for _ in range(repeats):
|
||
torch.cuda.synchronize()
|
||
start = time.perf_counter()
|
||
call_kernel()
|
||
torch.cuda.synchronize()
|
||
times.append((time.perf_counter() - start) * 1000) # ms
|
||
|
||
times.sort()
|
||
return times[len(times) // 2] # median
|
||
|
||
|
||
def main():
|
||
p = argparse.ArgumentParser(
|
||
description="Triton prefill kernel benchmark with compile-time param injection")
|
||
p.add_argument("--block", type=int, nargs="+", default=[16, 32, 64, 128],
|
||
help="BLOCK_M sizes to test (each triggers Triton recompilation)")
|
||
p.add_argument("--block-n", type=int, nargs="+", default=None,
|
||
help="BLOCK_N sizes (default: same as --block). Use different values for asymmetric search.")
|
||
p.add_argument("--warps", type=int, nargs="+", default=[1, 2, 4, 8],
|
||
help="NUM_WARPS values to test")
|
||
p.add_argument("--ctx-lens", type=int, nargs="+", default=[128, 512, 2048, 8192],
|
||
help="Context lengths to benchmark (problem sizes)")
|
||
p.add_argument("--batch", type=int, default=1, help="Batch size")
|
||
p.add_argument("--seq-len", type=int, default=1, help="New tokens per sequence")
|
||
p.add_argument("--head-dim", type=int, default=128, help="Head dimension (Qwen3.6=128)")
|
||
p.add_argument("--num-heads", type=int, default=64, help="Query heads (Qwen3.6=64)")
|
||
p.add_argument("--num-kv-heads", type=int, default=8, help="KV heads (Qwen3.6=8, GQA)")
|
||
p.add_argument("--block-size", type=int, default=16, help="KV cache block size")
|
||
p.add_argument("--dtype", default="float16", choices=["float16", "bfloat16"])
|
||
p.add_argument("--output", type=str, default=None, help="Save JSON results")
|
||
p.add_argument("--warmup", type=int, default=3)
|
||
p.add_argument("--repeats", type=int, default=10)
|
||
args = p.parse_args()
|
||
|
||
torch, props = require_gpu()
|
||
import triton
|
||
|
||
# Import the actual Triton kernel from prefix_prefill.py
|
||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)) + "/..")
|
||
try:
|
||
from prefix_prefill import _fwd_kernel
|
||
except ImportError:
|
||
# Try vllm path
|
||
from vllm.attention.ops.prefix_prefill import _fwd_kernel
|
||
|
||
dtype = getattr(torch, args.dtype)
|
||
|
||
print(f"\nmuh bench_triton_prefill")
|
||
print(f" Kernel: prefix_prefill._fwd_kernel (Triton JIT)")
|
||
print(f" Mechanism: each (BLOCK, NUM_WARPS) pair → Triton recompilation → different PTX")
|
||
print(f" Model: Qwen3.6 (head_dim={args.head_dim}, heads={args.num_heads}, kv_heads={args.num_kv_heads})")
|
||
block_n_vals = args.block_n if args.block_n else args.block
|
||
print(f" Search: BLOCK_M={args.block} × BLOCK_N={block_n_vals} × WARPS={args.warps} = {len(args.block)*len(block_n_vals)*len(args.warps)} variants")
|
||
print(f" Problem sizes: ctx_len={args.ctx_lens}")
|
||
print()
|
||
|
||
# Collect results
|
||
all_results = []
|
||
baseline_times = {} # ctx_len → time with default config (BLOCK=64, WARPS=4)
|
||
|
||
for ctx_len in args.ctx_lens:
|
||
print(f"--- ctx_len={ctx_len} ---")
|
||
|
||
q, k, v, o, k_cache, v_cache, b_loc, b_start_loc, b_seq_len, b_ctx_len = \
|
||
make_test_tensors(torch, args.batch, args.seq_len, ctx_len,
|
||
args.num_heads, args.num_kv_heads, args.head_dim,
|
||
args.block_size, dtype)
|
||
max_input_len = args.seq_len
|
||
|
||
for block_m in args.block:
|
||
for block_n in block_n_vals:
|
||
for warps in args.warps:
|
||
label = f"bm_{block_m}.bn_{block_n}.w_{warps}"
|
||
|
||
t = bench_one_config(
|
||
torch, triton, _fwd_kernel, block_m, block_n, warps,
|
||
q, k, v, o, k_cache, v_cache, b_loc, b_start_loc,
|
||
b_seq_len, b_ctx_len, max_input_len,
|
||
warmup=args.warmup, repeats=args.repeats,
|
||
)
|
||
|
||
if t is None:
|
||
print(f" {label:35s} COMPILE FAIL (SMEM overflow)")
|
||
all_results.append({
|
||
"block_m": block_m, "block_n": block_n, "warps": warps,
|
||
"ctx_len": ctx_len, "time_ms": None, "status": "compile_fail",
|
||
})
|
||
else:
|
||
if block_m == 64 and block_n == 64 and warps == 4:
|
||
baseline_times[ctx_len] = t
|
||
|
||
speedup = baseline_times.get(ctx_len, t) / t if t > 0 else 0
|
||
marker = " ★" if speedup > 1.05 else " ✗" if speedup < 0.9 else ""
|
||
print(f" {label:35s} {t:8.3f} ms {speedup:6.3f}x{marker}")
|
||
|
||
all_results.append({
|
||
"block_m": block_m, "block_n": block_n, "warps": warps,
|
||
"ctx_len": ctx_len, "time_ms": round(t, 4),
|
||
"speedup": round(speedup, 4), "status": "ok",
|
||
})
|
||
|
||
# Free tensors
|
||
del q, k, v, o, k_cache, v_cache, b_loc, b_start_loc, b_seq_len, b_ctx_len
|
||
torch.cuda.empty_cache()
|
||
|
||
# Summary: best config per ctx_len
|
||
print(f"\n{'='*60}")
|
||
print("BEST per context length:")
|
||
for ctx_len in args.ctx_lens:
|
||
ctx_results = [r for r in all_results if r["ctx_len"] == ctx_len and r["status"] == "ok"]
|
||
if ctx_results:
|
||
best = min(ctx_results, key=lambda r: r["time_ms"])
|
||
print(f" ctx={ctx_len:>5d}: bm={best['block_m']}, bn={best['block_n']}, "
|
||
f"warps={best['warps']}, time={best['time_ms']:.3f}ms, "
|
||
f"speedup={best.get('speedup', 1):.3f}x")
|
||
|
||
# CCCL-format output
|
||
print(f"\nCCCL-format output:")
|
||
for block_m in args.block:
|
||
for block_n in block_n_vals:
|
||
for warps in args.warps:
|
||
label = f"bm_{block_m}.bn_{block_n}.w_{warps}"
|
||
speedups = []
|
||
for ctx_len in args.ctx_lens:
|
||
r = next((r for r in all_results
|
||
if r["block_m"] == block_m and r["block_n"] == block_n
|
||
and r["warps"] == warps
|
||
and r["ctx_len"] == ctx_len and r["status"] == "ok"), None)
|
||
if r and "speedup" in r:
|
||
speedups.append(f"{r['speedup']:.6f}")
|
||
else:
|
||
speedups.append("N/A")
|
||
print(f" {label} {' '.join(speedups)}")
|
||
|
||
if args.output:
|
||
os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)
|
||
with open(args.output, "w") as f:
|
||
json.dump({
|
||
"benchmark": "triton_prefill",
|
||
"kernel": "prefix_prefill._fwd_kernel",
|
||
"mechanism": "Triton JIT constexpr recompilation",
|
||
"gpu": torch.cuda.get_device_name(0),
|
||
"sm_count": props.multi_processor_count,
|
||
"smem_per_block": props.max_shared_memory_per_block,
|
||
"model_config": {
|
||
"head_dim": args.head_dim,
|
||
"num_heads": args.num_heads,
|
||
"num_kv_heads": args.num_kv_heads,
|
||
"dtype": args.dtype,
|
||
},
|
||
"baseline": {"block": 64, "warps": 4},
|
||
"baseline_times": {str(k): v for k, v in baseline_times.items()},
|
||
"results": all_results,
|
||
}, f, indent=2)
|
||
print(f"\nSaved: {args.output}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|