Files
AstrAI/csrc/kernels
ViperEkura 4c82d5d84b perf: dispatch non-128-divisible shapes to the 64x64 cta
- a 128x128 CTA that is not exactly tiled (m or n not a multiple of 128) runs its edge tiles on the predicated generic path, and with a single in-flight wave the runtime is the slowest CTA — the edge tiles drag the whole shape down, so prefer_small_cta now also takes n and returns true when 64 divides both dims but 128 does not (measured sweep: 1088^3 big-CTA 76T vs 64x64 93T)
- the 64x64 grid tiles exactly on such shapes and overlaps waves (CTA count jumps 64->100->144 in the 1024-1536 band; the non-divisible points like 1088/1216 sit in deep sawtooth valleys that the divisibility rule lifts to the flat ~100-111T plateau)
- double-non-divisible shapes (e.g. 1000^3) stay on the big CTA: measured 64x64 36.7T vs 128x128 40.0T — both grids carry edge tiles there and the big CTA's efficiency wins
- m <= 64 or n <= 64 also takes the small CTA (a 128-wide CTA wastes more than half its columns on narrow N)

Benchmark: L20, e2e CUDA graph (GPU 7, same GPU as all comparisons): 960^3 60.2->100.1T, 1088^3 77.6->101.6T, 1216^3 69.7->103.5T, 1344^3 84.9->111.1T; 128-divisible shapes unchanged within noise (1024^3 118.4, 1152^3 150.6, 1280^3 106.8, 1536^3 153.7, 2048^3 192.2). 596 pytests pass.
2026-08-26 19:02:50 +08:00
..