tune heuristics for inner outer scheduler after cached inputs are moved to smem

[liqiangxl]

Issue:
After all inner persistent buffers are moved to shared memory, the register buffer size is 4 and 8 bytes per element for RMS norm backward and layer norm backward. Current heuristics uses the largest possible batch size and leads to performance regressions. For exammpe, in RMS norm at hidden size 20K, will use 20 bathes with 128 threads. The batch size is too large and caused register spills, should consider increase threads per block to reduce batch size.

Fix:
(1) Max threads per block is increased from 256 to 512 due to reduced register pressures.
(2) Generate candidate heuristics using different threads per block and sort these heuristics based on register usage and occupancy.
Generate candidate heuristic:

vectorization_factor = largest possible
For each threads per block {128, 256, 512}
       generate_a_heuristics(threads per block, vectorization_factor)

Sort candidates based on:

        (1) prefer candidate won't cause register spills, extra_regs = avilable - required (based on buffer size)
        if (extra_regs_a > 0 && extra_regs_b < 0) {
          return true;
        } else if (extra_regs_a < 0 && extra_regs_b > 0) {
          return false;
        }
        // (2) prefer candidate with occupancy >= 16 warps, higher occupancy doesn't mean higher peformance but if smaller than
        // 16 warps per sm, increase occupany leads to better performance.
        if (a.warps_per_sm < 16 || b.warps_per_sm < 16) {
          return a.warps_per_sm > b.warps_per_sm;
        }
        // tie breaker, smaller threads_per_block reduces reduction overhead
        return a.threads_per_block < b.threads_per_block;

Results:
Hopper:
(1) layer norm and rms nomr bwd, fp16

(2) layer norm and rms nomr bwd, fp32

Amper: see dashboard

[liqiangxl]

!build

[liqiangxl]

!build --pybench

[liqiangxl]

!build --pybench