#!/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())