Commit Graph

2 Commits

Author SHA1 Message Date
Claude
a2a5dd8f00 feat: asymmetric BLOCK_M/BLOCK_N search + re-add BI-V100 autotune configs
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
2026-08-03 11:18:18 +00:00
dylanyunlon
327497dc17 feat(muh): add bench_triton_prefill.py — real compile-time param injection benchmark
Unlike bench_bi100.py which called torch.sum() without injecting params:
- Directly invokes prefix_prefill._fwd_kernel Triton JIT kernel
- Each (BLOCK, NUM_WARPS) constexpr pair triggers Triton recompilation
  into a different kernel binary — same mechanism as CCCL #define TUNE_*
- Combos that exceed SMEM fail at compile time (caught, reported as COMPILE FAIL)
- Measures actual kernel execution time per compiled variant
- Outputs speedup vs baseline (BLOCK=64, WARPS=4) in CCCL format

Search space: BLOCK=[16,32,64,128] × WARPS=[1,2,4,8] = 16 variants
Problem sizes: ctx_len=[128,512,2048,8192] (Qwen3.6 typical workloads)
Test tensors match Qwen3.6: head_dim=128, num_heads=64, num_kv_heads=8 (GQA)

Requires GPU — will error immediately if no CUDA device available.
Reports GPU properties (SM count, SMEM, VRAM) to confirm BI-V100 hardware.
2026-08-03 10:46:19 +00:00