68d500c960
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).
206 lines
8.6 KiB
Python
206 lines
8.6 KiB
Python
#!/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())
|