[Ruishenl]

Summary:

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

Context

pack_segments converts concatenated variable-length sequences into a dense padded tensor. The legacy GPU path launches an exclusive prefix sum and then packs a flat grid of scalar output elements. For small and medium sequence counts, the extra launch, repeated index arithmetic, and scalar memory accesses are significant overheads.

This diff

• Makes the fused prefix-and-pack kernel the default CUDA and ROCm path for num_seq <= 7168, the largest sampled boundary that improved every tested payload on H100 and MI350.
• Maps grid.y to one sequence, computes its input offset with a block reduction, and performs copy, truncation, and padding in the same kernel.
• Selects TPB64/128/256/512 from device coverage and vector work per thread, with a two-warp minimum, TPB512 maximum, and eight target loads per thread.
• Selects the widest safe aligned load/store type from 128, 64, 32, or scalar bits.
• Retains the legacy prefix-sum path above the measured boundary and changes its grid-stride index to int64_t, so outputs larger than INT_MAX elements remain correct.
• Extends the benchmark CLI with implementation, sequence-count, shape, and cache-flush controls, and adds boundary, vector-width, and large-grid correctness coverage.

The output contract is unchanged. The fused path preserves the assertion that the first input dimension equals the sum of segment lengths.

Performance

On H100, the target model improved QPS by 6.4% and latency by 6.1% versus legacy. TPB1024 regressed 7 of 15 long-shape cases, so the selector remains capped at TPB512.

On MI350X, fixed-TPB sweeps found TPB512 best for the target (num_seq, max_length, cell_size) shapes (35, 1500, 256) and (35, 2000, 256). The sequence-count sweep found 7168 to be the largest sampled all-shape fused win; tiny payloads crossed over near 7680, while large rows remained substantially faster with the fused path beyond that point.

Model 1134083489, snapshot 0, subnet merge_request_only, batch size 2048, was measured for 1,000 iterations in three interleaved legacy/selector pairs on MI350X:

• Legacy median: 1.57 ms, 1,307,644 QPS, 23.92 TFLOP/s
• Final selector median: 1.47 ms, 1,397,127 QPS, 25.55 TFLOP/s
• Improvement: -6.37% latency and +6.84% QPS

An earlier four-pair run measured +13.5% QPS. Candidate throughput was effectively unchanged; the newer legacy baseline was faster, so the three-pair result above is the conservative headline.

Interactive benchmark comparison (white background): https://pxl.cl/d59Rg

Original before trace:
{F1995546236}

Original fused-kernel trace:
{F1995545397}

Reviewed By: adelesun, q10

Differential Revision: D118393046

[meta-codesync]

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