Parallelise the NOBAG inference TBE over row ranges, not tables (#6256)

[zhaozhul]

Summary:

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

parallel_for_table_threads schedules whole tables, so the makespan can never drop below the single largest table and speedup is hard-capped at sum(rows) / max(rows in one table). Sequence models put nearly all their lookups in one or two features, which makes that cap brutal.

The cap is structural, not a tuning problem

• On blue_reels_vdd model 2113513583 the remote_request_only TBE has 14 features, and two of them (datafm_fb_trait_post_id / datafm_fb_trait_post_owner_rid, the two parallel UIH sequences) carry 19,641 rows each — 99.2% of the work.
• Measured ceiling: 2.02x, at any thread count. The two features are one entry per history event, so they are locked to the same length and the ratio stays ~2.0 regardless of history length or traffic mix.

Fix: parallelise over (table, row-range) chunks

• The enabling property is in the kernel. EmbeddingSpMDMAutovec.cc:393-402, the int4 NOBAG inner loop, is for i in [0, output_size): read indices[i], memcpy one weight row to out, advance. It never reads offsets_or_lengths and holds no cross-row state, so a row range is trivially separable.
• The pooled path right below it does accumulate via offsets_or_lengths[m+1] - offsets_or_lengths[m], and is left on the per-table scheduler.
• A chunk is the existing whole-table call rebased: indices_acc + chunk_first_row, a unit-stride offsets slice starting at the same row, and output_acc + chunk_first_row * elems_D. The whole-table call is just the special case r0=0, r1=rows_of(t), so the sub-call is self-similar to the one that ships today and does not depend on the kernel's internal offset convention.
• Threads write disjoint output slices with no reduction, so results are bitwise identical to serial.

Row-based thread-count heuristic

• calculate_num_threads() used table count as a work proxy (num_tables / 16) — exactly what misclassified this model. 14 tables carrying a 39k-row gather looked "too small to thread" and got 1 thread even with FBGEMM_TBE_MAX_NUM_THREADS=8 set, because 14/16 = 0.
• choose_num_threads_for_rows() derives the count from actual output rows, with the onset knob renamed accordingly (FBGEMM_TBE_MIN_ROWS_PER_THREAD, default 4096).

Default path is unchanged. FBGEMM_TBE_MAX_NUM_THREADS still defaults to a cap of 1, so the default stays serial. Nothing threads unless a deployment opts in.

Incidental: the NOBAG output byte offset was computed in int32_t (offsets_acc[t*B] * elems_D), which overflows past ~8.3M output rows at D=512 int4. Now int64_t.

Results

Synthetic 14-table INT4 NOBAG host-DRAM TBE (D=512, row_alignment=1, 8 GiB) driven by the real per-feature index distribution, 39,617 rows/call. Paired in-process real-vs-balanced comparison, so the ratio is immune to host drift:

threads	before (per-table)	ratio	after (row chunks)	ratio	speedup vs serial
serial	4.337 ms	0.79x	4.500 ms	0.82x	1.00x
2	4.431 ms	1.40x	2.708 ms	0.88x	1.66x
4	2.864 ms	1.49x	1.490 ms	1.03x	3.02x
8	4.505 ms	2.99x	0.916 ms	0.90x	4.91x
14/16	2.754 ms	3.14x	0.790 ms	0.98x	5.70x

• The ratio column is load imbalance (balanced work / real work, same total rows, same table sizes). Per-table scheduling degrades from 0.79x to 2.99x as threads are added — that divergence is the cap. With row chunking it stays ~1.0 at every thread count: the imbalance is gone.
• Scaling efficiency falls from 61% at 8 threads to 36% at 16, consistent with a memory-latency-bound random gather saturating memory-level parallelism — the expected physical limit rather than a scheduling one.

Reviewed By: hanli0612

Differential Revision: D114795659

[meta-codesync]

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