Small B/H/S optimizations

[icavan]

Description

Optimize kernel performance for small batch size (B), number of heads (H), and sequence length (S) configurations.

Context

Current kernels are primarily tuned for large-scale configurations. Many real-world inference and fine-tuning workloads operate with small B, H, or S values, where kernel launch overhead and underutilization of SMs become significant. Targeted optimizations for these cases can unlock better performance in production serving scenarios.

Tasks

• Profile kernel performance on small B/H/S configurations
• Identify bottlenecks (launch overhead, SM underutilization, etc.)
• Implement specialized kernel variants or tiling strategies for small configurations
• Add benchmarks covering small B/H/S cases
• Validate correctness

[chli09]

Hi, I'm interested in working on this issue. Since #49 is already exploring the smaller-chunk-size direction, is this still open for parallel work on the other bottlenecks (e.g. launch overhead)? Thanks!

[chli09]

Status update for #11

Hi maintainers,

I ran a (B, H, T) sweep on H100 (fused) vs B200 (modular), comparing against FLA Triton baseline:

Full sweep table

! marks where cuLA loses (speedup < 1.0×). · is no data.

   B  H grid │ T=64       │ T=128      │ T=256      │ T=512      │ T=1024     │ T=2048     │ T=4096     │ T=8192     │ T=16384    │
             │ H100  B200 │ H100  B200 │ H100  B200 │ H100  B200 │ H100  B200 │ H100  B200 │ H100  B200 │ H100  B200 │ H100  B200 │
─────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
   1  4    4 │ 3.57 2.25  │ 3.51 2.24  │ 3.52 2.22  │ 3.46 2.18  │ 3.30 2.17  │ 1.74 2.14  │ 0.91!2.13  │ 0.46!2.14  │ 0.25!1.64  │
   1  8    8 │ 3.58 2.25  │ 3.60 2.26  │ 3.63 2.28  │ 3.55 2.28  │ 3.29 2.18  │ 1.75 2.21  │ 0.90!2.21  │ 0.46!2.85  │ 0.34!1.57  │
   1 16   16 │ 3.60 2.33  │ 3.54 2.12  │ 3.56 2.24  │ 3.53 2.28  │ 3.16 2.13  │ 1.68 2.17  │ 0.87!2.22  │ 0.47!1.76  │ 0.52!1.63  │
   2  4    8 │ 3.54 2.26  │ 3.57 2.22  │ 3.48 2.28  │ 3.51 2.33  │ 3.39 2.21  │ 1.74 2.18  │ 0.91!2.18  │ 0.45!2.14  │ 0.35!1.55  │
   2  8   16 │ 3.51 2.29  │ 3.57 2.25  │ 3.47 2.24  │ 3.44 2.22  │ 3.16 2.21  │ 1.70 2.24  │ 0.87!2.20  │ 0.52!1.80  │ 0.51!1.71  │
   2 16   32 │ 3.55 2.18  │ 3.50 2.25  │ 3.53 2.21  │ 3.43 2.27  │ 3.03 2.20  │ 1.58 2.22  │ 0.90!2.12  │ 0.77!1.72  │ 0.87!1.75  │
   4  4   16 │ 3.53 2.24  │ 3.55 2.22  │ 3.52 2.26  │ 3.45 2.26  │ 3.18 2.17  │ 1.70 2.22  │ 0.88!2.17  │ 0.52!1.80  │ 0.51!1.70  │
   4  8   32 │ 3.50 2.25  │ 3.48 2.26  │ 3.48 2.23  │ 3.46 2.27  │ 3.02 2.23  │ 1.57 2.21  │ 0.89!2.15  │ 0.85!1.81  │ 0.87!1.88  │
   8  4   32 │ 3.55 2.27  │ 3.52 2.27  │ 3.50 2.23  │ 3.43 2.25  │ 2.98 2.17  │ 1.59 2.19  │ 0.91!2.11  │ 0.86!1.84  │ 0.85!1.87  │
─── grid ≥ 64 (≥ H100 SM count = 132) ───
   4 16   64 │ 3.54 2.28  │ 3.47 2.25  │ 3.48 2.25  │ 3.46 2.27  │ 2.61 2.22  │ 1.40 2.05  │ 1.58 1.94  │ 1.29 2.00  │ 1.55 2.05  │
   8  8   64 │ 3.39 2.23  │ 3.54 2.24  │ 3.53 2.24  │ 3.44 2.22  │ 2.66 2.22  │ 1.48 2.21  │ 1.57 2.05  │ 1.46 2.13  │ 1.55 2.16  │
   8 16  128 │ 3.53 2.24  │ 3.53 2.26  │ 3.47 2.25  │ 3.43 2.28  │ 2.01 2.15  │ 2.04 1.95  │ 2.51 2.00  │ 1.97 2.04  │ 2.50 2.05  │
   4 32  128 │ 3.48 2.31  │   ·    ·   │ 3.49 2.77  │   ·    ·   │ 2.56 2.23  │   ·    ·   │ 2.06 1.97  │   ·    ·   │ 2.14 2.04  │
   8 32  256 │ 3.61 2.24  │   ·    ·   │ 3.57 2.26  │   ·    ·   │ 2.10 1.90  │   ·    ·   │ 2.26 2.00  │   ·    ·   │ 2.27 2.05  │
  16 32  512 │ 3.64 2.23  │   ·    ·   │ 1.91 2.24  │   ·    ·   │ 2.09 1.97  │   ·    ·   │ 2.26 2.06  │   ·    ·   │ 2.30 2.06  │
  32 32 1024 │ 3.43 2.26  │   ·    ·   │ 1.61 1.88  │   ·    ·   │ 2.10 1.99  │   ·    ·   │ 2.28 2.05  │   ·    ·   │ 2.31 2.04  │
   4 64  256 │ 3.49 2.23  │   ·    ·   │ 3.56 2.21  │   ·    ·   │ 1.82 1.62  │   ·    ·   │ 1.90 1.67  │   ·    ·   │ 1.94 1.71  │
   8 64  512 │ 3.68 2.19  │   ·    ·   │ 1.91 2.14  │   ·    ·   │ 1.77 1.66  │   ·    ·   │ 1.92 1.70  │   ·    ·   │ 1.95 1.73  │
  16 64 1024 │ 3.41 2.25  │   ·    ·   │ 1.36 1.59  │   ·    ·   │ 1.76 1.65  │   ·    ·   │ 1.93 1.70  │   ·    ·   │ 1.97 1.70  │
  32 64 2048 │ 1.85 1.99  │   ·    ·   │ 1.37 1.62  │   ·    ·   │ 1.76 1.68  │   ·    ·   │ 1.94 1.70  │   ·    ·   │   ·    ·   │

Hopper fused performs poorly at small grid + long T — at B=1, H=4, T=8192 it drops to 0.46× vs FLA, and as bad as 0.25× at T=16384. cuLA wins comfortably (1.4–3.6×) everywhere else.

Interesting observation: B200 modular doesn't show this dip — it stays at 1.55–2.85× across the same problematic cells. Likely because the modular path's intra and fwd_o get chunk-parallel grids naturally, and only chunk_delta_h is grid-limited (smaller share of total time).

Root cause on Hopper

SM90 fused kernel uses grid = B·H. At B=1, H=4 you get 4 CTAs vs 132 SMs —most of the chip sits idle. The fused single-CTA-per-sequence design (probably from FlashInfer's FlatMainloopTmaWarpSpecialized*?) inherently can't fill SMs in this regime, regardless of T.

ChunkWiseParallel Stub？

csrc/kda/sm90/kernel/tile_scheduler.hpp already has:

// compatible interface, for work without ChunkWiseParallel,
// chunk_len equals to seq_len
CUTE_DEVICE int32_t
chunk_len() const { return seq_len; }

The ChunkWiseParallel mention in the comment suggests the original authors anticipated this direction.

Direction I am exploring (work in progress)

I started implementing it: split each (seq, head) along T into N segments, dispatch to parallel CTAs, merge via affine combine (M, h_ext). (Similar to Blackwell design, or "sequence splitting" technique? )

Early measurements: ~1.5–1.8× speedup over the N=1 baseline at the worst-case point (B=1, H=4, T=8192) — narrows the gap from 0.46× to ~0.7–0.8× vs FLA, though still slightly behind FLA in this case.

Question

Is the Hopper small-grid + long-T regime the right focus for #11? Or are there other small-B/H/T configurations you'd rather I prioritize (e.g., very-small-T where launch overhead dominates, varlen patterns, B200 cases)?

[Hyaloid]

Hi, could you please review this pr 86 when you have a chance? This PR adds support for SM90 and may help resolve the issue. I'd appreciate any feedback or suggestions.