test: add scan tuning verification — union SMEM model from agent_scan.cuh
Tests scan SMEM safety (7/7 pass), delay parameter scaling from SM100, and tile size comparison. Key insight: agent_scan.cuh _TempStorage is a UNION — BlockLoad, BlockStore, BlockScan share SMEM. Peak = max(tile, scan_scratch), NOT tile + scan_scratch. 8B structs at 99% SMEM utilization (48640/49152) are valid under union model. Initial test had false SMEM overflow alarm (used sum model).
This commit is contained in:
dylanyunlon
205
muh/tests/test_scan_tuning.py
Normal file
205
muh/tests/test_scan_tuning.py
Normal file
@@ -0,0 +1,205 @@
|
||||
#!/usr/bin/env python3
|
||||
"""test_scan_tuning.py — Verify muh scan tuning against CCCL ground truth
|
||||
|
||||
Scan is SMEM-critical (unlike reduce which is register-limited):
|
||||
BlockLoad stages data from global→SMEM: tile = threads × items × type_size
|
||||
BlockScan uses SMEM scratch: ~threads × sizeof(AccumT) for RAKING
|
||||
Lookback uses tile_state SMEM: negligible (global memory)
|
||||
|
||||
Total SMEM ≈ threads × items × type_size + threads × accum_size
|
||||
Must be ≤ 48KB (49152 bytes) on BI-V100.
|
||||
|
||||
CCCL reference: cub/benchmarks/bench/scan/exclusive/sum.cu
|
||||
ipt_22.tpb_384.ns_1904.dcid_6.l2w_830.trp_1.ld_0 1.148 0.997 1.140 1.463
|
||||
|
||||
This benchmark comment tells us SM100's best config for 4B types is:
|
||||
items=22, threads=384, delay=1904ns, algo=exponential_backon_jitter_window,
|
||||
l2w=830ns, transpose=true, load=DEFAULT
|
||||
→ tile = 384 × 22 × 4 = 33792B (69% of 48KB) ✓
|
||||
|
||||
BI-V100 differences:
|
||||
1. SM count: 16 vs 148 → fewer CTAs, larger tiles beneficial
|
||||
2. L2 cache: 6MB vs 50MB → l2w needs recalibration
|
||||
3. Lookback delay: 16 CTAs contend less → shorter delays possible
|
||||
"""
|
||||
|
||||
import sys, os
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
|
||||
|
||||
SMEM_LIMIT = 49152
|
||||
WARP_SIZE = 32
|
||||
SM_COUNT = 16
|
||||
|
||||
def scale_mem_bound(nom_t, nom_i, ts, max_smem=SMEM_LIMIT):
|
||||
items = max(1, min(nom_i * 4 // ts, nom_i * 2))
|
||||
spt = ts * items
|
||||
if spt > 0:
|
||||
raw = max_smem // spt
|
||||
max_t = ((raw + 31) // 32) * 32
|
||||
else:
|
||||
max_t = nom_t
|
||||
threads = min(nom_t, max_t)
|
||||
if threads < 32:
|
||||
threads = 32
|
||||
return items, threads
|
||||
|
||||
|
||||
# bi100 scan structs from tuning_scan.cuh
|
||||
BI100_SCAN = {
|
||||
# Lookback structs: tile = threads × items × type_size (BlockLoad staging)
|
||||
"bi100_lookback_1B_o4": {"threads": 512, "items": 18, "type_size": 1,
|
||||
"delay_algo": "exponential_backon", "delay_ns": 175, "l2w": 270},
|
||||
"bi100_lookback_2B_o4": {"threads": 512, "items": 13, "type_size": 2,
|
||||
"delay_algo": "exponential_backon", "delay_ns": 175, "l2w": 270},
|
||||
"bi100_lookback_4B_o4": {"threads": 384, "items": 22, "type_size": 4,
|
||||
"delay_algo": "exponential_backon_jitter_window", "delay_ns": 952, "l2w": 415},
|
||||
"bi100_lookback_8B_o4": {"threads": 320, "items": 19, "type_size": 8,
|
||||
"delay_algo": "exponential_backon_jitter_window", "delay_ns": 386, "l2w": 426},
|
||||
"bi100_lookback_1B_o8": {"threads": 384, "items": 14, "type_size": 1,
|
||||
"delay_algo": "exponential_backon_jitter_window", "delay_ns": 760, "l2w": 465},
|
||||
"bi100_lookback_4B_o8": {"threads": 320, "items": 19, "type_size": 4,
|
||||
"delay_algo": "exponential_backon_jitter_window", "delay_ns": 478, "l2w": 330},
|
||||
"bi100_lookback_8B_o8": {"threads": 320, "items": 19, "type_size": 8,
|
||||
"delay_algo": "exponential_backon_jitter_window", "delay_ns": 478, "l2w": 330},
|
||||
}
|
||||
|
||||
# SM100 best from CCCL benchmark annotations
|
||||
SM100_SCAN = {
|
||||
"4B_o4_best": {"threads": 384, "items": 22, "type_size": 4,
|
||||
"speedup": [1.148442, 0.997167, 1.139902, 1.462651],
|
||||
"annotation": "ipt_22.tpb_384.ns_1904.dcid_6.l2w_830.trp_1.ld_0"},
|
||||
"4B_o4_alt1": {"threads": 512, "items": 18, "type_size": 4,
|
||||
"speedup": [1.188818, 1.005682, 1.173041, 1.305288],
|
||||
"annotation": "ipt_18.tpb_512.ns_768.dcid_7.l2w_820.trp_1.ld_0"},
|
||||
"8B_o4_best": {"threads": 416, "items": 14, "type_size": 8,
|
||||
"speedup": [1.107210, 1.000000, 1.100637, 1.307692],
|
||||
"annotation": "ipt_14.tpb_384.ns_228.dcid_7.l2w_775.trp_1.ld_1"},
|
||||
}
|
||||
|
||||
|
||||
def test_scan_smem():
|
||||
"""Test SMEM safety for all bi100 scan structs.
|
||||
|
||||
CRITICAL INSIGHT from agent_scan.cuh:
|
||||
_TempStorage is a UNION of {load, store, scan+prefix}
|
||||
Peak SMEM = max(load_tile, store_tile, scan_scratch + prefix_scratch)
|
||||
NOT load_tile + scan_scratch (that was the bug in previous test version)
|
||||
|
||||
BlockLoad::TempStorage for WARP_TRANSPOSE:
|
||||
= threads × items × type_size (the tile staging buffer)
|
||||
BlockScan::TempStorage for WARP_SCANS:
|
||||
≈ num_warps × type_size (one value per warp) — much smaller than tile
|
||||
TilePrefixCallback::TempStorage:
|
||||
= a few integers for lookback state — negligible
|
||||
|
||||
So: peak SMEM ≈ tile = threads × items × type_size
|
||||
"""
|
||||
print("=== scan SMEM safety (union model from agent_scan.cuh) ===")
|
||||
all_ok = True
|
||||
for name, cfg in BI100_SCAN.items():
|
||||
tile = cfg["threads"] * cfg["items"] * cfg["type_size"]
|
||||
# WARP_SCANS scratch: one value per warp
|
||||
num_warps = cfg["threads"] // WARP_SIZE
|
||||
scan_scratch = num_warps * cfg["type_size"]
|
||||
# prefix callback: ~16 bytes
|
||||
prefix_scratch = 16
|
||||
# Union: peak = max(tile, scan+prefix)
|
||||
peak_smem = max(tile, scan_scratch + prefix_scratch)
|
||||
pct = peak_smem / SMEM_LIMIT * 100
|
||||
ok = peak_smem <= SMEM_LIMIT
|
||||
if not ok:
|
||||
all_ok = False
|
||||
status = "PASS" if ok else "FAIL"
|
||||
print(f" {status}: {name:30s} tile={tile:6d}B (peak, union) "
|
||||
f"scan_alt={scan_scratch+prefix_scratch:5d}B → peak={peak_smem:6d}B ({pct:.0f}%)")
|
||||
return all_ok
|
||||
|
||||
|
||||
def test_scan_delay_scaling():
|
||||
"""Test that BI-V100 delay parameters are scaled correctly from SM100.
|
||||
|
||||
BI-V100 vs SM100:
|
||||
- SM count: 16 vs 148 → 9.25× fewer CTAs → less tile_state contention
|
||||
- L2: 6MB vs 50MB → 8.3× smaller → l2w should be proportionally higher
|
||||
- But L2/SM: 384KB vs 338KB → BI-V100 has MORE L2 per SM
|
||||
|
||||
CCCL scaling heuristic (from existing muh comments):
|
||||
delay_ns *= 0.5 (fewer CTAs → less contention → shorter delay)
|
||||
l2w *= 0.6 (L2/SM is similar, but total L2 is smaller)
|
||||
"""
|
||||
print("\n=== scan delay parameter scaling ===")
|
||||
|
||||
# SM100 reference
|
||||
sm100_delay = 1904 # ns (from best scan benchmark)
|
||||
sm100_l2w = 830 # ns
|
||||
|
||||
# BI-V100 4B_o4 struct
|
||||
bi100 = BI100_SCAN["bi100_lookback_4B_o4"]
|
||||
|
||||
# Expected: delay_ns ≈ 1904 × 0.5 = 952
|
||||
expected_ns = int(sm100_delay * 0.5)
|
||||
actual_ns = bi100["delay_ns"]
|
||||
ns_ok = abs(actual_ns - expected_ns) <= 100 # tolerance
|
||||
|
||||
# Expected: l2w ≈ 830 × 0.5 = 415
|
||||
expected_l2w = int(sm100_l2w * 0.5)
|
||||
actual_l2w = bi100["l2w"]
|
||||
l2w_ok = abs(actual_l2w - expected_l2w) <= 100
|
||||
|
||||
print(f" SM100 best: delay={sm100_delay}ns l2w={sm100_l2w}ns")
|
||||
print(f" BI100 4B: delay={actual_ns}ns (expected ~{expected_ns}) "
|
||||
f"{'PASS' if ns_ok else 'WARN'}")
|
||||
print(f" BI100 4B: l2w={actual_l2w}ns (expected ~{expected_l2w}) "
|
||||
f"{'PASS' if l2w_ok else 'WARN'}")
|
||||
print(f" Note: delay scaling is heuristic — real values need BI-V100 benchmark")
|
||||
|
||||
return True # Warn only, don't fail
|
||||
|
||||
|
||||
def test_scan_vs_sm100_tile_size():
|
||||
"""Compare tile sizes between SM100 and BI-V100."""
|
||||
print("\n=== scan tile size comparison ===")
|
||||
|
||||
for key in ["4B_o4_best", "8B_o4_best"]:
|
||||
sm = SM100_SCAN[key]
|
||||
sm_tile = sm["threads"] * sm["items"] * sm["type_size"]
|
||||
|
||||
# Find matching bi100 struct
|
||||
ts = sm["type_size"]
|
||||
bi_key = f"bi100_lookback_{ts}B_o4"
|
||||
if bi_key in BI100_SCAN:
|
||||
bi = BI100_SCAN[bi_key]
|
||||
bi_tile = bi["threads"] * bi["items"] * bi["type_size"]
|
||||
ratio = bi_tile / sm_tile
|
||||
|
||||
print(f" {key}:")
|
||||
print(f" SM100: t={sm['threads']:3d} i={sm['items']:2d} → tile={sm_tile:6d}B")
|
||||
print(f" BI100: t={bi['threads']:3d} i={bi['items']:2d} → tile={bi_tile:6d}B "
|
||||
f"(ratio={ratio:.2f}×)")
|
||||
print(f" SM100 speedup: {sm['speedup']}")
|
||||
|
||||
# BI-V100 wants same or larger tiles (16 SMs need more work per CTA)
|
||||
if ratio < 0.5:
|
||||
print(f" WARN: BI-V100 tile is less than half of SM100 — may be too small")
|
||||
|
||||
|
||||
def main():
|
||||
print("muh scan tuning verification")
|
||||
print("CCCL ref: cub/benchmarks/bench/scan/exclusive/sum.cu")
|
||||
print("=" * 60)
|
||||
|
||||
results = []
|
||||
results.append(("scan SMEM safety", test_scan_smem()))
|
||||
results.append(("delay scaling", test_scan_delay_scaling()))
|
||||
test_scan_vs_sm100_tile_size()
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
all_pass = all(r for _, r in results)
|
||||
for name, ok in results:
|
||||
print(f" {'PASS' if ok else 'FAIL'}: {name}")
|
||||
print(f"\nOverall: {'ALL PASS' if all_pass else 'SOME FAILURES'}")
|
||||
return 0 if all_pass else 1
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Reference in New Issue
Block a user