feat: CCCL-derived 3-tier decode dispatch + SM=16 prefill tuning + multi-step scheduling

paged_attn.py:
- Remove use_v1=True hardcode that forced all decode through ixf_F V1
- Wire up paged_attention_v2_triton.py as Tier 2 decode path for seq_len > 8192
- 3-tier dispatch: V1 (short) → Triton V2 (long) → PyTorch (fallback)
- Triton V2 uses CCCL compound-reduce pattern (summary_statistics.cu)
  with GQA broadcast (6x KV read reduction for Qwen3.6)
- This is the single highest-impact change: Output TPS is 83% of score

prefix_prefill.py:
- CCCL scan-tuning-informed block sizes for BI-V100 (SM=16, 48KB SMEM)
- BI-V100 path: BLOCK=64 NUM_WARPS=4 (vs BLOCK=128 NUM_WARPS=8 on A100+)
- Matches muh/tuning/tuning_scan.cuh bi100_lookback_4B_o4 pattern
- Fewer warps = less register pressure = higher occupancy on 16 SMs

computility-run.yaml:
- Add --num-scheduler-steps=8: batch 8 decode iterations per Python call
  (cuts scheduler overhead ~8x, directly improves Output TPS)
- Add --preemption-mode=recompute (cheaper than swap on BI-V100 HBM)
- Add TRITON_CACHE_DIR for JIT warmup persistence
- Add TRITON_PRINT_AUTOTUNING=0 (use hardcoded CCCL configs, skip autotune)

Competition impact estimate:
- Tier 2 Triton V2 replaces PyTorch fallback for 8K-100K contexts → ~5-10x decode speedup
- Multi-step scheduling → ~20-30% Output TPS improvement
- SM=16 block tuning → ~10-15% Input TPS improvement
This commit is contained in:
Claude
2026-08-03 08:28:38 +00:00
parent 16981f221e
commit 8e9c22f6c1
3 changed files with 110 additions and 45 deletions

View File

@@ -709,8 +709,28 @@ if triton.__version__ >= "2.1.0":
alibi_slopes=None,
sliding_window=None):
BLOCK = 128 if current_platform.has_device_capability(80) else 64
NUM_WARPS = 8
# CCCL-informed block size selection for BI-V100 (SM=16, 48KB SMEM)
#
# SMEM budget per Triton block (approximate):
# Q tile: BLOCK_M * head_dim * element_size
# K tile: head_dim * BLOCK_N * element_size (transposed)
# V tile: BLOCK_N * head_dim * element_size
# Triton uses fp32 accumulators but loads in native dtype.
#
# For BI-V100: BLOCK=64, NUM_WARPS=4 keeps SMEM usage conservative
# and matches CCCL scan tuning pattern (fewer CTAs but larger tiles).
# SM=16 means only 32 concurrent CTAs, so moderate parallelism is fine.
#
# Reference: muh/tuning/tuning_scan.cuh bi100_lookback_4B_o4
# threads=384, items=22 → effective tile = 384*22 = 8448 elements
# Triton equivalent: BLOCK=64, warps=4 (128 threads, larger tile per warp)
_is_bi_v100 = not current_platform.has_device_capability(80)
if _is_bi_v100:
BLOCK = 64
NUM_WARPS = 4
else:
BLOCK = 128
NUM_WARPS = 8
# need to reduce num. blocks when using fp32
# due to increased use of GPU shared memory