Parallelize reorder_batched_ad_indices within a segment

[zhaozhul]

Summary:
X-link: https://github.com/facebookresearch/FBGEMM/pull/3132

reorder_batched_ad_indices_kernel_vec assigns exactly one warp to each (b, t) segment and sizes the grid as ceil(B*T / NUM_WARPS). That is fine when there are thousands of segments, but inference-side KJT projection calls it with very few segments whose lengths are extremely unbalanced.

Found while profiling a CCHSTU ranking predictor on H100 (kineto trace, merge|1). The kernel showed up as grid = 1..6 thread blocks running for ~810 us per call — 1 SM of 132 busy, an effective few GB/s on a 3.35 TB/s card. It cost 286 us/request, 4.5% of all GPU kernel time, and blocked the calling thread (GPURebatchUtils::concatPredictionInputIValues > StaticMemoryBatching::batch > StaticMemoryBatching::projectKJT) for a matching 291 us/request.

Two stacked problems:

• With T = 14 tables and small B, B*T is a few tens, so ceil(B*T / 32) yields 1-6 blocks.
• Segment lengths are heavily skewed. On the model profiled, 2 of the 14 features hold 99.2% of the indices, so even within those blocks roughly 2 warps do all the copying while the rest retire immediately. Runtime is just the time for a single warp to stream ~1 MB.

Fix: parallelize within a segment as well as across segments. The kernel now derives its intra-segment lane index from blockIdx.y * blockDim.x + threadIdx.x and strides by gridDim.y * blockDim.x, and the launch fans each segment out over gridDim.y blocks until the device is filled. The two vectorized tail guards are rebased onto that lane index so exactly one thread still writes each remainder element. blockDim.y is also shrunk to the actual segment count instead of launching idle warps.

Segment lengths live on device, so the fan-out is a fixed heuristic (kBlocksPerSM * #SMs, capped at kMaxSegmentFanout) rather than a data-dependent value — no D2H sync is introduced. The kernel grid-strides over its segment, so any fan-out is correct, and over-provisioned blocks fall straight through. When there are already enough segments to fill the device, gridDim.y resolves to 1 and the launch is identical to before.

Also adds reorder-batched-ad-indices-skew-bench to sparse_ops_benchmark.py. The existing reorder-batched-ad-indices-bench generates uniform lengths and therefore cannot expose this at all — every segment does identical work. The new benchmark reproduces the production shape (T=14, 2 tables holding 99.2% of the indices) and checks the result against the CPU reference before timing.

Differential Revision: D114798440

[meta-codesync]

@zhaozhul has exported this pull request. If you are a Meta employee, you can view the originating Diff in D114798440.