Implement pre-sorting, caching and contigous warp processing in group_index_select#144
Open
avbokovoy wants to merge 126 commits into
Open
Implement pre-sorting, caching and contigous warp processing in group_index_select#144avbokovoy wants to merge 126 commits into
avbokovoy wants to merge 126 commits into
Conversation
aryaman-gupta
requested changes
Mar 5, 2026
aryaman-gupta
left a comment
aryaman-gupta
left a comment
There was a problem hiding this comment.
The PR introduces crucial optimizations for the group_index_select_or_add_2d_kernel. The majority of the code is clean and the separation of ROCm and CUDA codepaths has been done well.
Most of these changes were already reviewed in #139 . I have left a few comments that I think should be looked at before merging. Some of these are design choices, and the PR could proceed with merging even if the code is not modified,
avbokovoy
force-pushed
the
abokovoi/group-index-sort-and-cache-opt
branch
from
March 11, 2026 12:21
3754ab4 to
1248c83
Compare
aryaman-gupta
approved these changes
Mar 11, 2026
aryaman-gupta
left a comment
aryaman-gupta
left a comment
There was a problem hiding this comment.
The points raised in #144 (review) have been addressed and the PR is ready for merging
avbokovoy
force-pushed
the
abokovoi/group-index-sort-and-cache-opt
branch
from
April 1, 2026 13:12
ed3fc4b to
f5f9e70
Compare
Summary: Pull Request resolved: pytorch#5799 X-link: https://github.com/facebookresearch/FBGEMM/pull/2727 D104827588 (length kernel) and f2e9cba2e279 (delinearize kernel) both introduced an unstated "indices must lie in [0, per_feature_hash)" assumption that the op never enforced before. ZCH callers that pass raw IDs to `jagged_unique_indices` (dedup runs before the hash-to-bucket remap) tripped this on the last feature in prod, producing `sum(output_lengths) < unique_indices.numel()` and crashing downstream `dist.all_to_all_single` with "Split sizes doesn't match total dim 0 size". Two fixes in `jagged_unique_indices.cu`: 1. `unique_indices_length_kernel`: when the group's `hash_end == hash_size_cumsum.size(0) - 1` (last feature), set `hi_pos = linear_unique_indices.size(0)` instead of doing the upper-bound binary search. The OOB tail of the last feature sorts past `hash_size_cumsum[T]` and the binary search would otherwise stop short of it. 2. Replace `delinearize_unique_from_sorted_kernel` (gather: `unique_indices[i] = v - hash_size_cumsum[t]`) with the original `delinearize_unique_index_kernel` (scatter: `unique_indices[reverse_index[i]] = indices[i]`). The gather form misattributes OOB values whose linearized key exceeds `hash_size_cumsum[T]` to a phantom feature `t = T` and emits `v - hash_size_cumsum[T]` instead of the original index value. The scatter form is index-bound-agnostic by construction. Note `reverse_index` is `int64_t` (from the cub pipeline) rather than templated `index_t` (which was the case under `at::_unique`). Contract enforcement: added a `CUDA_KERNEL_ASSERT` in `linearize_index_flat_kernel` that fires when an intermediate feature (`t < T - 1`) with `per_feature_hash > 0` has `idx >= per_feature_hash`. Intermediate-feature OOB causes silent per-feature count drift (counts leak from `t` to `t+1` while the total is preserved) — this has been broken since the op was written, but no caller hit it, so the assert surfaces violations rather than silently corrupting downstream embedding lookups. The assert exempts (a) the last feature (legitimately supported) and (b) merged/masked features with `per_feature_hash == 0` (the `hash_size_offsets` indirection pattern used by `multi_keys` and `zch_huge_hash_size`). The pipeline contract doc comment above `jagged_unique_indices_cuda` is updated to enumerate the three cases. Reviewed By: q10 Differential Revision: D106305472 fbshipit-source-id: 1ccfb44e5a08ddf257e1484473323f7a1accd6ba
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/2780 Pull Request resolved: pytorch#5859 - Remove CUDA 12.9.1 support from FBGEMM, following PyTorch Reviewed By: cthi Differential Revision: D107971273 fbshipit-source-id: 04f754d52137bb7c884dfe8e19e517540eaf121a
Summary: Pull Request resolved: pytorch#5858 X-link: https://github.com/facebookresearch/FBGEMM/pull/2773 This PR implements direct HBM->LDS stores in tbe inference kernel. There are 2 major changes: 1. Rows data isn't loaded in-place, instead we store pointers to global memory and store the actual data w.r.t. the predicate into LDS. In case predicate is false, we pre-allocate small chunk of static device memory of 16B once, fill it with zeros, and fallback to this chunk 2. HBM->LDS 16B loads are implemented for ROCm >= 7.0 and MI350. We can expand the support range to MI30* through 4B loads, however it doesn't bring any performance benefits because we'll have to introduce an overhead of addresses transposition and 4x more load operations. You can find out the reference implementation here: pytorch@fe52557. Due to pre-7.2 ROCm features, we are forced to used assembly inline to get 16B loads to work, so manual synchronization was added. In case of ROCm >= 7.2, we use proper intrinsics to handle memory synchronization. This change brings ~10% performance boost on average for weighted and unweighted cases. We may try to push it further by doing async loads for indices weights. Pull Request resolved: pytorch#5348 Reviewed By: spcyppt Differential Revision: D91496421 Pulled By: q10 fbshipit-source-id: acbbc2287b87ba2ecab1524c1300bb431667cc34
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/2781 generate_opcheck_tests records a "xsuccess" status when a previously-failing opcheck variant now passes -- i.e. the entry is stale and should be removed. Removing these de-clutters the failures dicts and drops the problematic-test count with zero behavioral risk (the tests already pass). 12 stale xsuccess entries removed: - jagged: fbgemm::jagged_index_select (KeyedJaggedIndexSelectTest aot_dispatch_dynamic + faketensor for test_keyed_jagged_index_select_dim1) -> {} - sparse: fbgemm::generic_histogram_binning_calibration_by_feature (4), fbgemm::group_index_select_dim0 (2), fbgemm::segment_sum_csr (2) -> {} - quantize: fbgemm::FloatToFP8RowwiseQuantized (2) -> {} Emptied keys are left as {} per the existing convention for healthy ops. Legitimate xfail entries (data-dependent / dynamic-output-shape variants) are untouched. Part of the fbgemm_dev test-health remediation for T191384137. Reviewed By: henrylhtsang Differential Revision: D107927030 fbshipit-source-id: ca72cd34d26db7a854e54a3a7c1f47009e9c9b28
Summary: Pull Request resolved: pytorch#5856 Fixes a ROCm/AMD GPU deadlock in gemv_quantized_fp8_fp8 (experimental/gen_ai/src/quantize/fast_gemv/include/fast_gemv.cuh), found by auditing fbgemm_gpu for the same barrier-divergence pattern fixed in D107554507. In the single-warp branch (blockDim.x <= WARP_SIZE), the `return` was nested inside `if (tid == 0)`, so only thread 0 exited while lanes 1..N-1 fell through to the block-wide __syncthreads(). On ROCm (WARP_SIZE=64) the blockDim.x in {32,64} instantiations take this branch, so the waiting lanes deadlock (latent UB on NVIDIA). This diff moves the `return` out of the `if (tid == 0)` block so all threads in a single-warp block return together and none reach __syncthreads() -- matching the sibling gemv_bf16 / gemv_quantized_bf16_fp8 / gemv_quantized_int4 kernels in the same file. The multi-warp path is unchanged. Reviewed By: henrylhtsang Differential Revision: D107947000 fbshipit-source-id: b3aa48de725eab6068ee5532f10106d76c062f60
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/2743 Pull Request resolved: pytorch#5817 Add out-of-band dirty bit tracking to FixedBlockPool using a striped concurrent hash set. The dirty bit is tracked separately from MetaHeader to maintain backward compatibility with serialization. Blocks are marked dirty when their embedding data is modified, and the bit is cleared on allocate/deallocate. This is foundational for the DRAM+SSD composite backend to know which DRAM blocks need to be flushed to SSD. Reviewed By: EddyLXJ Differential Revision: D107291829 fbshipit-source-id: bc06ee629ffc09556021ef339b9b86c45a21c4cc
…ytorch#5861) Summary: Pull Request resolved: pytorch#5861 X-link: https://github.com/facebookresearch/FBGEMM/pull/2782 block_bucketize_sparse_features_meta unconditionally returned a Tensor for the 5th output (unbucketize_permute), but the real CPU/CUDA kernels return it only when `sequence` is True (CPU sparse_ops_cpu.cpp:1175/1262/1280; CUDA sparse_block_bucketize_features.cu:831-833/956-957; schema sparse_ops_cpu.cpp:3822 -> (Tensor, Tensor, Tensor?, Tensor?, Tensor?)). This Tensor-vs-None mismatch is the root cause behind the block_bucketize test_faketensor__* opcheck xfails. Gate output 5 on `sequence` (the param is already in scope) so the meta matches real kernel semantics. outputs 3 (weights) and 4 (pos) were already correct. The corresponding block_bucketize test_faketensor__* xfails in sparse/failures_dict.json are intentionally left in place: once GPU CI confirms they now pass they will surface as (tolerated) xsuccess and can be removed in a follow-up cleanup. Removing them pre-emptively would hard-fail if any variant still mismatches. Part of the fbgemm_dev test-health remediation for T191384137. Reviewed By: henrylhtsang Differential Revision: D107927031 fbshipit-source-id: 02a28b1bbdd593332e1eced9f8dce27c80adb3c2
Summary: Pull Request resolved: pytorch#5862 The py3_14t-rocm7_2 manywheel build/upload legs were failing in OSS CI (3 FAILING signals on the fbgemm_dev surface, T191384137). Free-threaded Python 3.14 is not yet supported for the ROCm wheel build, mirroring the already-excluded 3.13t+rocm combo. Extend the existing Nova matrix exclusion from gpu_arch_type:rocm;python_version:3.13t to gpu_arch_type:rocm;python_version:3.13t,3.14t in the three wheel workflows that build ROCm (build_wheels_linux_x86, build_wheels_genai_linux_x86, build_wheels_linux_aarch64). This drops the failing py3.14t ROCm legs until upstream free-threaded ROCm support lands. Reviewed By: henrylhtsang Differential Revision: D107927032 fbshipit-source-id: bf19d134ba27d773eb222afa142109c9521f293d
Summary: Pull Request resolved: pytorch#5863 X-link: https://github.com/facebookresearch/FBGEMM/pull/2784 test_scatter_add_padded_tokens was failing as a TIMEOUT (not a kernel bug). Two causes, both confirmed by running on GPU (B200 via RE): 1. torch.compile recompilation explosion. The Hypothesis sweep ran the compiled=True path by calling torch.compile(scatter_add_padded_tokens) per example. dynamo recompiled on every dtype/shape/expert-count combo (e.g. "tensor 'token_counts' dtype mismatch. expected Long, actual Int"; then "in_tokens ... Int vs BFloat16"), blew past config.recompile_limit (8), after which every call recompiled (~1-2 min each) -> timeout. 2. The eager sweep itself was too slow at max_examples=100: each example re-autotunes the triton kernel for a new (EP, E, T_BUCKET, D) key and runs an expensive Python reference over dim=5120. Fixes: - Split the compiled path into a dedicated test_scatter_add_padded_tokens_compiled that exercises torch.compile on a small fixed set of configs (one compile per dtype), compiled once via an lru_cache'd torch.compile(..., dynamic=True) singleton. - The parametrized sweep (test_scatter_add_padded_tokens) now runs eager only, with max_examples reduced to 20. - Make token_counts consistently int32 (bincount returns int64) for parity with the balanced path and the kernel. - Extract the shared body into _run_scatter_add_padded_tokens. Completes the MoE bucket for T191384137 (the shuffling off-by-one was fixed in D107927038). Reviewed By: henrylhtsang Differential Revision: D107936151 fbshipit-source-id: ec5f31a4bdc69b14e41f6eb4992f82db3486e34c
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/2774 Add an `enable_dirty_tracking` option to `SynchronizedShardedMap` and forward it to the pool constructor, and thread the same parameter through `InferenceFixedBlockPool` so the shared template instantiates for both pool types. Defaults to false, so existing behavior is unchanged. This lets callers (e.g. the FeatureEvict SSD writeback path) construct pools that track dirty blocks. Reviewed By: EddyLXJ Differential Revision: D107967620 fbshipit-source-id: 68bdbd2e391192b212fd4b177ca172356282f916
…rch#5857) Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/2779 Pull Request resolved: pytorch#5857 Fixes a ROCm/AMD GPU deadlock in get_pruning_lengths-style barrier divergence, found by auditing fbgemm_gpu for the same pattern fixed in D107554507. In quantize_float_to_mx4_kernel (src/quantize_ops/quantize_mx.cuh), threads whose linear_tid >= total_elems hit an early `return` before the block-wide __syncthreads() calls. The kernel is launched with multiple groups per block (blockDim.y = num_groups_per_block = MAX_THREADS / mx_group_size) and gridDim_x = div_round_up(total_num_groups, num_groups_per_block), so when total_num_groups is not a multiple of num_groups_per_block the last block contains whole trailing groups whose threads all return early while the active groups wait at __syncthreads() -> deadlock on ROCm (latent UB on NVIDIA). This diff removes the early return and instead masks per-thread work with an `active = linear_tid < total_elems` flag: the input load, the smem combine writes, and the output write are guarded, while all threads fall through to every __syncthreads(). Inactive lanes contribute the minimum biased exponent (0) so the max reductions are unaffected. Behavior is identical for active groups. The sibling dequantize_mx4_to_float_kernel has no barrier after its early return and is left unchanged. ## Where the hanging barrier actually is (not apparent from the diff) Unlike the gen_ai amax/quantize twin (D107946896), the barriers here are in the kernel body, but which ones hang and under what condition is still subtle: - Unconditional block-wide barriers: __syncthreads() at the even/odd 4-bit combine steps (post-fix lines ~205 and ~212). These are ALWAYS reached, so they are the primary hang points. - Conditional block-wide barriers: __syncthreads() at lines ~141 and ~159 are only compiled in the has_multiple_warps_in_group == true template instantiation (mx_group_size > WARP_SIZE); for mx_group_size <= WARP_SIZE only the unconditional pair applies. - (The fbgemm_gpu::syncwarp() at line ~148 is warp-scoped, not a block barrier, and is not the hang point.) Why a subset of the block skips them (the deadlock condition): - The block is 2D: blockDim = (mx_group_size, num_groups_per_block) with num_groups_per_block = MAX_THREADS / mx_group_size, so a block holds MULTIPLE groups along threadIdx.y. The __syncthreads() calls are block-wide (all mx_group_size * num_groups_per_block threads) even though they are only logically needed within a single group's smem region. - Because total_elems = total_num_groups * mx_group_size, the early-return condition linear_tid >= total_elems is WHOLE-GROUP: an entire threadIdx.y row is either fully active or fully inactive (never partial). - gridDim_x = div_round_up(total_num_groups, num_groups_per_block). When total_num_groups is not a multiple of num_groups_per_block, the LAST block contains whole trailing (inactive) groups whose threads all returned, while the active groups in the same block reach the block-wide __syncthreads() and wait forever -> ROCm hang. Corollary / dependency: with num_groups_per_block == 1 (one group per block) the bug would not trigger, since the early return would be block-uniform. It is the multi-group-per-block launch that turns the per-group early return into a partial-block barrier divergence. Reviewed By: henrylhtsang Differential Revision: D107946777 fbshipit-source-id: 29f6cca1441a49828e1b753c9f214feb0eb0d4aa
…rch#5864) Summary: Pull Request resolved: pytorch#5864 First foundational diff of the BC-safe TBE type migration (config types only). - tbe/config/embedding_config.py becomes the canonical home for the config types; adds get_new_embedding_location() free-fn parity with the EmbeddingLocation.from_device_and_clf() classmethod. - Pins __module__ of the four TorchScript/pickle identity-serialized config types (EmbeddingLocation, PoolingMode, ComputeDevice, RecordCacheMetrics) back to fbgemm_gpu.split_table_batched_embeddings_ops_common so already-published TorchScript graphs and RankAGI training-TBE pickles still resolve/repackage. torch.package never rewrites a class __module__, so the pin survives interning. - split_table_batched_embeddings_ops_common.py drops its duplicate config-type definitions and re-exports them from the leaf, keeping the legacy import path working. Cache/SSD types remain defined here and migrate in follow-up diffs. - BUCK: split_table_batched_embeddings_ops_common deps += tbe:tbe_config. This is the safe, acyclic version of the config portion of D103477971 (which caused SEV S669532). See plan fbgemm_tbe_type_migration_bc_safe sections A.2/A.4. Reviewed By: henrylhtsang Differential Revision: D107684319 fbshipit-source-id: 42a887f29612c8bcf432838aea64c0a6bab8b7ce
Summary: Pull Request resolved: pytorch#5865 X-link: https://github.com/facebookresearch/FBGEMM/pull/2786 Second foundational diff of the BC-safe TBE type migration (cache types). - tbe/cache/cache_config.py is the canonical home for CacheAlgorithm, CacheState, MultiPassPrefetchConfig, UVMCacheStatsIndex, DEFAULT_ASSOC. Pins CacheAlgorithm.__module__ back to split_table_batched_embeddings_ops_common so already-published scripted TorchScript graphs / torchrec CacheParams pickles still resolve and repackage (torch.package never rewrites a class __module__). - tbe/cache/__init__.py guards the heavy .split_embeddings_cache_ops import with try/except so the new lightweight :tbe_cache_config target loads cleanly. - split_table_batched_embeddings_ops_common.py drops its duplicate cache-type definitions and re-exports them from the leaf; construct_cache_state becomes a thin forwarder to CacheState.construct(). Legacy import path preserved. - BUCK: add lightweight :tbe_cache_config (cache __init__ + cache_config.py); :tbe_cache_ops now omits cache_config.py and deps :tbe_cache_config; ops_common deps += :tbe_cache_config. Acyclic (leaf-or-shim -> tbe_cache_config -> tbe_config; ops_common -> tbe_cache_config). Safe, acyclic, identity-preserving version of the cache portion of D103477971 (which caused SEV S669532). See plan fbgemm_tbe_type_migration_bc_safe A.2/A.4. Reviewed By: henrylhtsang Differential Revision: D107684313 fbshipit-source-id: 8d0e1e0a62544beca8ed634a5e0f5bc51802417b
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/2791 Pull Request resolved: pytorch#5871 The `group-index-select-2d-bench` subcommand declared the CLI options `--input-dims`, `--input-strides`, and `--index-dim` but never implemented the body that uses them, so passing the flags was a silent no-op: every invocation ran the default homogeneous synthetic shape regardless of the ragged flags. This diff implements the ragged path so the bench actually measures the per-group production shapes (matching the echen40 reference scripts), and unifies the forward + backward benchmarking step for both families: - Add `_gen_synthetic_group_inputs` and `_gen_ragged_group_inputs` helpers, each returning `(input_group, indices_group)`. - Branch on `--input-dims` (empty/whitespace normalizes to the synthetic path; `--index-dim` is required for ragged; `--input-strides` is optional and defaults to contiguous strides per group). - Unify forward/backward on `torch.ops.fbgemm.group_index_select_dim0` with `torch.cat(output_group, dim=1)`; drop the single-tensor PyTorch `torch.index_select` reference (cannot represent heterogeneous columns). - Align the D96328337 run harness: pass `--input-precision`/`--sort-indices` for ragged configs and record `sort_indices=false` (echen40 parity). The measured kernel (group_index_select_or_add_2d_kernel) is unaffected by the concat axis, so the existing synthetic A/B results remain valid. Reviewed By: spcyppt Differential Revision: D108184412 fbshipit-source-id: fef51b3d51cebcc0d6df89de1a3474dbc4f08f8a
…ytorch#5869) Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/2789 Pull Request resolved: pytorch#5869 block_bucketize_test had 14+ FAILING test_faketensor__ / test_aot_dispatch_dynamic__ opcheck variants (populate_bucketized_permute*, block_bucketize_sparse_features_inference). Two root causes: 1. block_bucketize_test.py was MISSING `import fbgemm_gpu.sparse_ops` -- every sibling sparse test file has it, and fbgemm_gpu/__init__ only imports it in OSS. Without it the Python meta/FakeTensor impls in sparse_ops.py were never registered in the test process, so every faketensor/aot_dispatch opcheck variant errored ("could not find the abstract impl ... import fbgemm_gpu.sparse_ops"). The xfail'd variants "passed" via that error; the non-xfail'd ones (populate, inference) failed. 2. populate_bucketized_permute had NO meta at all, and the inference / 2d_weights metas returned unbucketize_permute (output 5) unconditionally, but the real CPU/CUDA kernels only return it when `sequence` is set (sparse_block_bucketize_features.cu:956). Fixes: - Add `import fbgemm_gpu.sparse_ops` to block_bucketize_test.py so the metas load. - Add populate_bucketized_permute_meta (returns empty_like(bucket_mapping), matching native_empty_like in sparse_ops_cpu.cpp:1117) and register it via impl_abstract. - Gate output 5 on `sequence` in block_bucketize_sparse_features_inference_meta and block_bucketize_sparse_features_2d_weights_meta (same fix as D107927031 for the base op). - With the metas now correctly loaded, 22 previously-xfail'd block_bucketize_sparse_features faketensor/aot_dispatch entries became unexpected successes -> removed them from sparse/failures_dict.json (kept the 2 BlockBucketize2DWeightsTest *_vs_original entries that still genuinely fail). Part of the fbgemm_dev test-health remediation for T191384137. Reviewed By: henrylhtsang Differential Revision: D108108673 fbshipit-source-id: 00ccaa84d537bf14b52e652bf4f1e68e7cfd33c7
…h#5870) Summary: Pull Request resolved: pytorch#5870 Same root cause as D108108673: block_bucketize_2d_weights_test.py was missing `import fbgemm_gpu.sparse_ops`, so the Python meta/FakeTensor impls were never registered for that test process and every faketensor/aot_dispatch opcheck variant errored ("could not find the abstract impl"). The xfail'd variants "passed" via the error; the genuine ones failed. Adding the import (which all sibling sparse test files already have) loads the metas -- including the block_bucketize_sparse_features_2d_weights_meta sequence-gating fix from D108108673. With the metas loaded, all 16 previously-xfail'd BlockBucketize2DWeightsTest faketensor/aot_dispatch entries became unexpected successes, so they are removed from sparse/failures_dict.json: - fbgemm::block_bucketize_sparse_features_2d_weights: 16 -> 0 - fbgemm::block_bucketize_sparse_features: 2 -> 0 (the *_2d_weights_vs_original entries, which opcheck the base op from the 2DWeights test, also now pass) Part of the fbgemm_dev test-health remediation for T191384137. ___ Differential Revision: D108114892 fbshipit-source-id: f64e19559ff505850014422b0ae3de409819cad0
…5875) Summary: Pull Request resolved: pytorch#5875 X-link: https://github.com/facebookresearch/FBGEMM/pull/2794 Add ROCm-specific `subwarp_reduce_add` template function using `__shfl_xor` intrinsics for subwarp reductions, and redefine `GROUP_REDUCE_ALL_SUM` to use it on ROCm. This replaces the NVIDIA `warpReduceAllSum` path which requires `shfl_sync_mask` not available on AMD GPUs. Affects both warp and CTA backward kernels via the shared optimizer device kernel template. | T | Ds / D | B | Baseline (No flag) | This diff (No flag) | Δ % (No flag) | Baseline (With flag) | Diff (With flag) | Δ % (With flag) | 2 | 8 | — | 65 | 65 | 0.00% | 65 | 65 | 0.00% | 10 | 128 | — | 394 | 376 | 5.00% | 394 | 376 | 5.00% | 14 | 12/24/128 | — | 373 | 360 | 3.00% | 374 | 357 | 5.00% | 18 | 20/128 | — | 309 | 298 | 4.00% | 308 | 298 | 3.00% | 22 | 20/24/128 | — | 424 | 405 | 4.00% | 425 | 405 | 5.00% | 1 | 8 | 2048 | 102 | 100 | 2.00% | 103 | 101 | 2.00% | 1 | 8 | 4096 | 180 | 177 | 2.00% | 181 | 177 | 2.00% | 1 | 8 | 131072 | 1485 | 1492 | 0.00% | 1483 | 1487 | 0.00% | 1 | 128 | 2048 | 115 | 112 | 3.00% | 114 | 114 | 0.00% | 1 | 128 | 4096 | 203 | 200 | 1.00% | 222 | 217 | 2.00% | 1 | 128 | 131072 | 1800 | 1790 | 1.00% | 1900 | 1738 | 9.00% Reviewed By: q10 Differential Revision: D107779306 fbshipit-source-id: c0f4dffa411c05005141297731f19633281139d6
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/2795 Third foundational diff of the BC-safe TBE type migration: make tbe.ssd.ssd_config the canonical home for the SSD/KVZCH types and turn ops_common into a pure re-export shim. - tbe/ssd/ssd_config.py is canonical for BackendType, EvictionPolicy, EnrichmentType, EnrichmentResponseFormat, EnrichmentPolicy, KVZCHParams, KVZCHTBEConfig (init-time only; no __module__ pinning required). - tbe/ssd/__init__.py guards the heavy .common/.inference/.inference_serving/.training imports (re-raising anything but the specific absent submodule) so the new lightweight :tbe_ssd_config target loads cleanly under torch.package. - split_table_batched_embeddings_ops_common.py is now a pure re-export shim: config + cache types eager, SSD/KVZCH lazy via __getattr__ to avoid a tbe.ssd -> .training -> ops_training import cycle; plus the construct_cache_state() forwarder. Legacy import path preserved. - BUCK: add lightweight :tbe_ssd_config; :ssd_split_table_batched_embeddings_ops drops ssd_config.py and deps :tbe_ssd_config. ops_common deps += :tbe_ssd_config and is marked autodeps-skip so the tbe leaf config deps (only reachable lazily) stay pinned into every publisher's torch.package closure (SEV-prevention hardening, S669532/S672649). Acyclic: ops_common -> {tbe_config, tbe_cache_config, tbe_ssd_config}. Completes the acyclic, identity-preserving redo of D103477971 (SEV S669532) for the foundational fbgemm-internal layer. Reviewed By: henrylhtsang Differential Revision: D107684315 fbshipit-source-id: 3fdb8135d6be630b18a9ee91150c1ae4a76fe104
Summary: Pull Request resolved: pytorch#5881 X-link: https://github.com/facebookresearch/FBGEMM/pull/2800 Tier-1 follow-up to D104903707 / D104937969 — apply the canonical ROCm grid-overflow cap to a `split_embeddings_cache/` launch site identified by the audit at: `/home/bensonma415/.llms/plans/permute_metric_layout_etc_rocm_grid_overflow_audit.plan.md` `lru_cache_insert_kernel` is launched with grid `div_round_up(N, kMaxThreads / kWarpSize)` and block `dim3(kWarpSize=32, kMaxThreads/kWarpSize=32) = 1024` threads, where `N = cache_set_sorted_unique_indices.numel()`. Total threads ~= 32 * N. For `N >= 2^27` (~134 M unique indices, reachable in large TBE workloads), total threads exceed HIP's 2^32 hard limit, tripping `KernelLauncher::checkThreadCountNotExceeded`. The kernel already grid-strides over `n` (line 48), so capping the host-side grid is correctness-preserving. NVIDIA codegen sees the host-side `#ifdef USE_ROCM` block as a plain alias and remains a no-op delta in PTX/SASS. The fix targets only the `lock_cache_line=False` branch. The `lock_cache_line=True` branch uses `div_round_up(SM_cnt, ALL_TO_PREFETCH_SM_RATIO)` which is already SM-bounded; the cap on its branch is a no-op. Reviewed By: spcyppt Differential Revision: D105282095 fbshipit-source-id: dde6db74cbff6113bb69aebe81d36268b2d599bf
Summary: Pull Request resolved: pytorch#5879 X-link: https://github.com/facebookresearch/FBGEMM/pull/2798 dense_to_jagged_forward derives the output's leading dim (total_L) from offsets.back().max() with no validation. A corrupted/garbage offset produced a huge negative value (~-1.4e18) that flowed straight into at::empty and aborted with an opaque "Trying to create tensor with negative dimension" error, seen in a real ROCm training crash through fbgemm_gpu::dense_to_jagged. Add a TORCH_CHECK_VALUE guard in the CPU and CUDA forward paths so the op fails fast with the offending value and a clear message instead of an allocator abort. The guard runs only in the concrete kernels; symbolic/torch.export tracing still uses the Python abstract impl. Reviewed By: haoyuz Differential Revision: D108330675 fbshipit-source-id: 7a06c405f88106e6883d58e872222632b0d04347
Summary: Pull Request resolved: pytorch#5886 X-link: https://github.com/facebookresearch/FBGEMM/pull/2807 Register `DispatchKey::Autograd` and `DispatchKey::Meta` for `dense_embedding_codegen_lookup_function` in the codegen template. Previously, the `{% if not dense %}` guard in `embedding_backward_split_host_template.cpp` skipped these registrations for the dense variant, causing `test_autograd_registration` and `test_faketensor` opcheck tests to fail. The non-dense split ops already register the same function at both Autograd and Meta dispatch keys (the autograd.Function works even without autograd enabled, and redispatches to lower-level ops that have their own Meta kernels). The dense variant uses the same pattern. This also reverts the `failures_dict_fast.json` xfail entries and `additional_decorators` workaround from `BackwardDeterminismTest` since they are no longer needed. Reviewed By: q10 Differential Revision: D107814730 fbshipit-source-id: e67be3e28681e2ac8ac843e9b71fea1e4a081f53
… typing (pytorch#5834) Summary: Pull Request resolved: pytorch#5834 X-link: meta-pytorch/torchrec#4320 X-link: https://github.com/facebookresearch/FBGEMM/pull/2757 The SymInt magic-method typing change added strict annotations to the previously-untyped `torch._check`, `torch._check_is_size`, and `torch._constrain_as_size` helpers (`cond: bool | SymBool`, `i`/`symbol: IntLikeType`). Call sites that pass a `Tensor.item()` result into these helpers now fail type checking, because `Tensor.item()` is typed `bool | float | int`, which is not assignable to `IntLikeType` or `bool | SymBool`. The values are genuinely int/bool-like here (tensor sizes and boolean comparison results); the error comes purely from `item()`'s necessarily-broad return type. This wraps the offending arguments in `typing.cast(int, ...)` / `typing.cast(bool, ...)`. `typing.cast` is erased at runtime, so behavior is unchanged. We deliberately do not use `int(...)` / `bool(...)`: under `torch.compile` / `torch.export` these `.item()` values are unbacked `SymInt` / `SymBool`, and forcing `int()` / `bool()` would trigger `GuardOnDataDependentSymNode` and alter runtime behavior. Reviewed By: bobrenjc93 Differential Revision: D107556825 fbshipit-source-id: b3be6de347c996fde3e602e63145331e6be70cdc
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/2832 Pull Request resolved: pytorch#5913 Add observability for the two queue depths in the RES streaming pipeline. (1) `weights_to_stream_queue_` (MPSC stream queue in fbgemm) emits `stream_mpsc_depth` after every dequeue. (2) `tensor_to_dedup_queues_[i]` (per-shard dedup queues in training_ps) emit `dedup_depth.shard_{i}` per shard plus a `dedup_depth.total` rollup, bumped at the top of `co_stream_tensors`. Backpressure today is invisible until a trainer crashes; these signals let oncall see saturation trends and per-shard hot-spots in advance. Both surfaces use OBC via `TrainingPsOdsLogger::bumpKeyGauge` (emits P50+P99 — the right stats for depth percentiles), on the `raw_embedding_streaming` category. The handler already owned `ods_logger_`; the fbgemm-side `RawEmbeddingStreamer` reuses the `TrainingPsOdsLogger` member added by the parent diff D107811590 (constructed only when streaming is enabled). No fb303: OBC reaches ODS via the host-level agent without per-process export config. (An earlier revision used an fb303 `ResGauges.h` helper for the fbgemm gauge; removed per review feedback on the sibling silent-failure diff.) Stacked on D107811590 (silent-failure ODS counters), which introduces the shared OBC logger member on `RawEmbeddingStreamer`. Both feed the master observability initiative T269497764. Design note: per-shard cardinality is N OBC keys x 3 stats; `num_dedup_threads_` is configurable but typically small (4-8); `.total` is for at-a-glance, per-shard for hot-spot debug. Is it caused by our diff? No. Two cases, both exonerating: 1. If the ffmpeg error is in an FBGEMM OSS-CI job (all the FAILED ones): our fbgemm code is entirely behind #ifdef FBGEMM_FBCODE, which the OSS build never defines — so OSS compiles a byte-identical translation unit to pre-diff. An ffmpeg build/link issue there is the OSS environment (the established FBGEMM OSS-CI flake), not our code. 2. If it's in a citadel/ci_workflows WARNING (videorecs_ranking, silvertorch, etc.): those are tagged WARNING_PREEXISTING_ISSUE — CI itself confirms they failed on the base revision before your change. Pre-existing = not us. Plus the clinchers: our change is OBC counter logging (bumpKey/bumpKeyGauge) in training_ps + the RES streamer — zero source/build/runtime path to ffmpeg (a third-party video codec lib). And the MVAI e2e ran green on this exact binary. So: the FFMPEG signal is not caused by your diff — it's either OSS-CI infra or a pre-existing base-rev warning. Reviewed By: FriedCosey Differential Revision: D108312495 fbshipit-source-id: 81381e7ca1609c98ca13cb530d60924981335c3a
) Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/2878 Pull Request resolved: pytorch#5966 The pytorch/test-infra build matrix (generate_binary_build_matrix.yml@main) now emits CUDA 13.2 (cu132) coordinates. FBGEMM does not support 13.2 yet, so every cu132 job in build_wheels_linux_x86 fails in the test-infra validate step with "ERROR: Unknown distribution ${DISTRIBUTION}". Filter cu132 out in the filter-matrix step by adding a `gpu_arch_version:13.2` filter group, until 13.2 is supported. Mirrors the intent of D104335566. Reviewed By: q10 Differential Revision: D110247364 fbshipit-source-id: 31e636ee6a36408f4102c225c4a0aa043d2df582
…smatch (pytorch#5963) Summary: Pull Request resolved: pytorch#5963 X-link: https://github.com/facebookresearch/FBGEMM/pull/2875 Add `__attribute__((visibility("hidden")))` to the `KernelLauncher` struct in `kernel_launcher.cuh`. This prevents the dynamic linker from exporting `launch_kernel` weak symbols, forcing each `.so` to use its own copy. **Root cause:** On CUDA 12.8 with sm_90+, the dynamic linker resolves `index_select_ops.so`'s `launch_kernel` to the main binary's copy (from `embedding_ops_training_gpu` via `embedding_ops`). This causes `__cudaPushCallConfiguration` (in the main binary) and `__cudaPopCallConfiguration` (in `index_select_ops.so`) to use different statically-linked CUDA runtime copies with separate TLS stacks. Pop finds an empty stack, returns non-zero, and `find_long_segments` kernel is silently skipped — causing zero gradients for runs >= 32 in `batch_index_select_dim0` backward. **Why this fix works:** `visibility("hidden")` prevents `launch_kernel` from being exported, so the dynamic linker cannot resolve it cross-module. Each `.so` uses its own copy — push and pop stay in the same TLS. This is structurally guaranteed regardless of the dependency tree or compiler inlining behavior. **Why not the other fixes:** - D103575890 (pre-build args at call site): Targeted fix for `find_long_segments` only — other kernels with the same pattern would need similar fixes. - D103575878 (`launch_kernel_with_transform`): Works empirically but depends on compiler inlining `launch_kernel<..., PackedTensorAccessor32...>` — not structurally guaranteed. Reviewed By: q10 Differential Revision: D104175840 fbshipit-source-id: 9d3aa18dcaf794b6659ab6d4af913ac18cf8aa5f
…pytorch#5968) Summary: Pull Request resolved: pytorch#5968 X-link: https://github.com/facebookresearch/FBGEMM/pull/2880 Fixes 100+ `*-type-checking` failures blamed on D109465368 (the backout of D107684317). Reviewed By: gchalump Differential Revision: D109954014 fbshipit-source-id: 9ffe31cde174edd05052040d22c78ef4b3c4035a
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/2873 Pull Request resolved: pytorch#5961 Register aggregate ops for DRAM_SSD composite backend metrics: ssd_num_lookups, ssd_num_hits, ssd_num_writes, dram_num_ids, and ssd_hit_rate. Reviewed By: EddyLXJ Differential Revision: D110075583 fbshipit-source-id: dfe2f789592ef2173bed5ebc4e59080c91fef55f
Summary: Pull Request resolved: pytorch#5965 ## Context This stack of diffs improves FBGEMM embedding table storage by integrating KVZCH SSD caching into `EmbeddingRocksDB`, enabling a hybrid **DRAM + SSD** approach. It adds support for metadata-only operations, cache lifecycle management, and better stats/performance—aiming to make large embedding table storage/retrieval more efficient. ## This Diff **[KVZCH][SSD] Return zero optimizer state for DRAM_SSD embedding cache** ### Summary This diff ensures **zero optimizer state is returned specifically for `DRAM_SSD` when running in _embedding cache mode_**. ### Key Changes - Adds `embedding_cache_mode` to `SSDTableBatchedEmbeddingBags` (in `kv_backend_test.py`) to toggle embedding cache behavior. - Introduces `_generate_zero_optimizer_states` (in `training.py`) to construct zero-filled optimizer-state tensors based on `optimizer.state_size_table`. ### Why this is correct / when it applies - **Only applies to embedding cache mode for `DRAM_SSD`.** - In embedding cache mode, the **embedding values are not backpropagated**, so their **optimizer state is always 0**. Reviewed By: EddyLXJ Differential Revision: D110243528 fbshipit-source-id: 8ed689b3b1facdf87942d7c3d3cd617505187499
…ytorch#5954) Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/2884 Fixes a use-after-free window between the eviction loop and inplace update / inference reads in the DramKV embedding cache. ## The bug - Eviction **freed blocks first** (mempool lock), then **erased keys from the shard map** (`wlock`) — leaving a window where a freed block was still reachable via the map. - Inplace update **dropped its shard `rlock` before writing** to weight blocks, so the evict loop could free a block mid-write. ## The fix (two coupled changes) 1. **Erase keys before freeing blocks**: pick victims (mempool lock) → erase keys (`wlock`, excludes readers) → free blocks once unreachable. 2. **Hold the shard `rlock` across the inplace hit writes** (lookup + write). While a reader holds the `rlock`, it blocks the evict loop's `wlock`, so the key cannot be erased and the block cannot be freed underneath an in-flight write. Pull Request resolved: pytorch#5954 Reviewed By: henrylhtsang Differential Revision: D110231575 Pulled By: q10 fbshipit-source-id: c440dd88fbc2f20e1a2e41f2524c5552127c125a
Summary: Pull Request resolved: pytorch#5973 X-link: https://github.com/facebookresearch/FBGEMM/pull/2886 HIP enforces a hard limit of 2^32 total threads per launch (unlike CUDA, which silently wraps; see ROCm/hip#2253). `jagged_softmax_forward_cuda` launches `jagged_softmax_kernel` with grid `dim3(D, min(B, kMaxBlockYDim=65535), 1)` and block 128. Total launch threads = `D * min(B, 65535) * 128`; for the canonical embedding shape `D >= 513, B = 65535` that is 4,295,032,320 > 2^32, which trips `KernelLauncher::checkThreadCountNotExceeded` on ROCm. The kernel already grid-strides over both `b` (line 36) and `d` (line 45), so capping the host-side x-dim grid is correctness-preserving. y is already clamped to `kMaxBlockYDim = 65535`. This change applies the same one-line `#ifdef USE_ROCM` cap pattern as parent fixes D104903707 / D104937969 in `sparse_ops/`. Audit: /home/bensonma415/.llms/plans/jagged_and_more_rocm_grid_overflow_audit.plan.md Reviewed By: henrylhtsang Differential Revision: D105096879 fbshipit-source-id: a673834245cdc1fc6bcbc9098ffb4b77c3f1c79d
Summary: Pull Request resolved: pytorch#5974 X-link: https://github.com/facebookresearch/FBGEMM/pull/2887 HIP enforces a hard limit of 2^32 total threads per launch (unlike CUDA, which silently wraps; see ROCm/hip#2253). `jagged_softmax_backward_cuda` launches `jagged_softmax_backward_kernel` with grid `dim3(D, min(B, kMaxBlockYDim=65535), 1)` and block 128. Total launch threads = `D * min(B, 65535) * 128`; for `D >= 513, B = 65535` that is 4,295,032,320 > 2^32 and trips `KernelLauncher::checkThreadCountNotExceeded` on ROCm. The kernel already grid-strides over both `b` (line 37) and `d` (line 46), so capping the host-side x-dim grid is correctness-preserving. y is already clamped to `kMaxBlockYDim = 65535`. Same one-line `#ifdef USE_ROCM` cap pattern as parent fixes D104903707 / D104937969 in `sparse_ops/`. Audit: /home/bensonma415/.llms/plans/jagged_and_more_rocm_grid_overflow_audit.plan.md Reviewed By: henrylhtsang Differential Revision: D105097144 fbshipit-source-id: a2cfd3d9721b638081a94f8def79ccd01615a1a3
…torch#5976) Summary: Pull Request resolved: pytorch#5976 X-link: https://github.com/facebookresearch/FBGEMM/pull/2888 PyTorch dropped CUDA 12.8 (cu128) starting with torch 2.12 -- 2.12 and 2.13 publish no cu128 wheels. Verified on download.pytorch.org (test channel): torch-2.13 has cu126=108 and cu130=108 wheels but cu128=0 (torch-2.12 cu128=0; cu128's last version was 2.11 with 42 wheels). FBGEMM's release workflows still offered "12.8.1" as a CUDA option, a leftover from the CUDA-12.8 era. Building/publishing a cu128 FBGEMM wheel for the v1.8.0 release (torch 2.13) is impossible -- there is no torch cu128 to build against -- so the option can only ever fail. Remove it so the release matrix reflects PyTorch and nobody can trigger a dead cu128 build. Changes (options + genai build/test matrices): - fbgemm_gpu_release_cuda.yml: [ "12.6.3", "12.8.1", "13.0.2" ] -> [ "12.6.3", "13.0.2" ] - fbgemm_gpu_release_genai.yml: same drop in the cuda-version-publish option and both build/test cuda-version matrices. Default stays 13.0.2. Net release CUDA set: cu126 (legacy) + cu130 (stable) -- matching torch 2.13's PyPI matrix. PyTorch source (authoritative): - "Introducing CUDA 13.2 and Deprecating CUDA 12.8 (Release 2.12)", Andrey Talman (PyTorch Dev Infra), 2026-04-01: https://fb.workplace.com/groups/671272153461300/permalink/1976674139587755/ "CUDA 12.8 is being deprecated and removed from CI/CD pipelines and binary build matrices ... CUDA 13.0 remains the stable variant published to PyPI ... CUDA 12.6 remains available as the legacy build." - RFC: pytorch/pytorch#178665 (CUDA support matrix for Release 2.12) - dev-discuss: https://dev-discuss.pytorch.org/t/introducing-cuda-13-2-and-deprecating-cuda-12-8-release-2-12/3337 - Policy: pytorch/pytorch RELEASE.md#pytorch-cuda-support-matrix (legacy / stable / experimental tiers) PyTorch's torch 2.13 CUDA tiers: 12.6 legacy, 13.0 stable, 13.2 experimental. FBGEMM v1.8.0 ships 12.6 + 13.0; 13.2 (experimental) is deferred until the FBGEMM cu132 build is nightly-green (tracked separately; cu132 was filtered in D110247364). ___ Differential Revision: D110492570 fbshipit-source-id: 03d2bccee832f22ad0e7a0d3b028afdfa72a7e51
Summary: Pull Request resolved: pytorch#5849 X-link: https://github.com/facebookresearch/FBGEMM/pull/2767 ## What Parallelizes the per-table loop in the CPU TBE forward kernel (IntNBitTableBatchedEmbeddingBagsCodegen::forward) across tables. The per-table loop is embarrassingly parallel — each table reads its own weight slice and writes a disjoint slice of the output — so fanning tables out across threads gives near-linear speedup on table-heavy inference models. Two env knobs control it: - FBGEMM_TBE_MAX_NUM_THREADS (default 1): the thread-count cap. =1 keeps the unchanged sequential behavior; =N>1 distributes tables over up to N OpenMP threads with dynamic scheduling (good load balancing when table sizes are skewed). - FBGEMM_TBE_MIN_TABLES_PER_THREAD (default 16): work-granularity for the per-call thread count (see the work-size guard below). ## Work-size guard (skip threading for small few-table lookups) A uniform thread count *regresses* small few-table TBE lookups: splitting a ~7–13 table loop across threads adds OpenMP fork/join + cross-call contention with no parallelism benefit. (Measured on dpa: its three small remote_ro_event lookups slowed down under threading, while its single 358-table remote_ro lookup sped up.) choose_table_threads() therefore both GATES and GRADES the per-call thread count on the work size via a single formula: N = clamp(work_units / G, 1, cap) (cap=FBGEMM_TBE_MAX_NUM_THREADS, G=FBGEMM_TBE_MIN_TABLES_PER_THREAD) A call with fewer than 2*G tables runs serial; larger calls scale one thread per G tables, up to the cap. With the default G=16 the threading onset is 32 tables, matching the gate validated in ICE A/B testing. Set FBGEMM_TBE_MIN_TABLES_PER_THREAD=1 to restore unconditional threading. The pure decision (choose_table_threads_impl) is split into fbgemm_gpu/utils/embedding_cpu_threading.h so it can be unit tested without touching the process environment. ## Default behavior is unchanged When FBGEMM_TBE_MAX_NUM_THREADS<=1 (the default) choose_table_threads short-circuits to 1 before FBGEMM_TBE_MIN_TABLES_PER_THREAD is even read, and the helper takes an early return that runs the loop body sequentially with no try/catch wrapper and no thread-local-state guard. So the default path is functionally identical to the pre-change code: same iteration order, the DEVICE-placement TORCH_CHECK in its original per-table position, same error semantics, and the same generated machine code for the body. The only always-on changes are mechanical and behavior-preserving (loop-local `weights` pointer, int64 loop index). ## Design - Raw OpenMP (#pragma omp parallel + omp for, schedule(dynamic)) rather than at::parallel_for, so TBE gets its own thread count independent of the global intra-op pool / OMP_NUM_THREADS (predictors run with OMP_NUM_THREADS=1). - The cap and granularity are read once from the env vars and cached (thread-safe static init); the per-call thread count from choose_table_threads() is clamped to the number of tables. ## Correctness - Removed the function-scoped `weights_acc` pointer, which every iteration overwrote — a data race once the loop is parallel. Replaced with a loop-local pointer (identical pointer value). Every other variable in the loop body is already loop-local, and each table writes a disjoint output slice (output_acc + D_start), so results are bitwise-identical to the sequential path. - The per-table DEVICE-placement TORCH_CHECK stays in its original position. In the threaded path it — like any other throw from the loop body (kernel errors, at::arange checks) — is captured (first one wins) and rethrown after the join, so no exception escapes the OpenMP region. - Worker threads restore the caller's at::ThreadLocalState (dispatch keys, grad/inference mode, autocast, ...), so ATen calls inside the loop (e.g. at::arange in the nobag path) run with the correct thread-local context. ## Verification - Builds clean (mode/opt); confirmed OpenMP is actually enabled for this target by inspecting the compiled object — gen_*_codegen_cpu.cpp.pic.o references __kmpc_fork_call / omp_get_num_threads (the pragma is NOT a no-op). - New unit tests in this diff: - test/tbe:embedding_cpu_threading_test — C++ tests of the pure choose_table_threads_impl decision: cap=1 (no env var) is ALWAYS serial regardless of work/granularity; onset at 2*G=32; grading scales to the cap; G=1 threads every call. - test/tbe:nbit_forward_threading — runs the real CPU int-nbit forward in separate processes under FBGEMM_TBE_MAX_NUM_THREADS in {unset, 1, 2, 4} (and FBGEMM_TBE_MIN_TABLES_PER_THREAD) and asserts the outputs are BITWISE identical (torch.equal) — i.e. table-threading does not change the result. - nbit_forward CPU unit tests pass with FBGEMM_TBE_MAX_NUM_THREADS=4, including test_nbit_forward_cpu_with_table_sharing (non-monotonic weights_offsets). ## Benchmark (INT4, B=512, D=128, E=100K, L=20, FP16/SUM, iters=100, 3-run avg) | Tables | Threads | Avg us | BW (GB/s) | Speedup | Efficiency | | 8 | 1 | 1,436 | 4.38 | --- | --- | | 8 | 2 | 1,042 | 6.03 | 1.38x | 69% | | 8 | 4 | 844 | 7.46 | 1.70x | 43% | | 8 | 8 | 773 | 8.14 | 1.86x | 23% | | 32 | 1 | 4,797 | 5.25 | --- | --- | | 32 | 2 | 2,830 | 8.90 | 1.69x | 85% | | 32 | 4 | 2,003 | 12.60 | 2.40x | 60% | | 32 | 8 | 1,633 | 15.49 | 2.94x | 37% | | 64 | 1 | 10,132 | 4.97 | --- | --- | | 64 | 2 | 6,767 | 7.44 | 1.50x | 75% | | 64 | 4 | 4,864 | 10.35 | 2.08x | 52% | | 64 | 8 | 4,033 | 12.50 | 2.51x | 31% | 2 threads is the efficiency sweet spot (69-85%); efficiency falls off at higher counts due to fixed fork/join overhead per call. This matches the production recommendation FBGEMM_TBE_MAX_NUM_THREADS=2 (+7.3% QPS measured in ICE). Microbenchmark kernel speedups overpredict end-to-end QPS (Amdahl: CPU TBE is a small fraction of total inference latency). Reviewed By: helloguo, q10 Differential Revision: D102867249 fbshipit-source-id: d98151ff7b9bd6af34636dee227b920775efa58e
Summary: Unblock the OSS CI docs/lint/ci jobs, which all failed in conda env setup. 1) setup_miniconda: drop --update-deps from the base conda update. The docs (build-docs), lint (run-lint), and ci_cpu/ci_cuda jobs -- everything that calls setup_miniconda -- failed because the base conda destroyed itself during env setup. Every subsequent `conda` call died with: ModuleNotFoundError: No module named 'conda' surfacing as `conda clean` failing (lint) and as broken conda post-link scripts when installing graphviz/doxygen (docs: gdk-pixbuf/gtk3 LinkError). Root cause: conda update -n base -c conda-forge --override-channels --update-deps -y conda conda-forge periodically publishes a base `conda` whose --update-deps solve bumps base Python to a version conda itself is not built against, orphaning conda's own site-packages. This is independent of the job's matrix Python version -- it happens in the base env before the matrix Python is used. (An earlier attempt to pin the docs/lint matrices to 3.13 was a no-op and was reverted.) Dropping --update-deps leaves the base interpreter intact while still updating conda. The gtk3/gdk-pixbuf docs failure was a downstream symptom of the same broken base conda (its post-link scripts invoke base conda), so this one change fixes both. 2) Drop the dead JSS hyperlink from the Xor128() doc comment in QuantUtilsAvx2.h. With base conda healthy, the docs build completes and only the Sphinx linkcheck fails, on the Xor128() reference which linked JSS's deprecated URL scheme (https://www.jstatsoft.org/v08/i14/paper) that the runner cannot reach. Remove the hyperlink and keep a plain-text citation for Marsaglia's "Xorshift RNGs" (JSS 8(14), 2003; doi:10.18637/jss.v008.i14). No live link means nothing for linkcheck to probe, so no linkcheck_ignore entry is needed. Note: the release wheel builds are unaffected -- they use test-infra's conda setup (setup-miniconda: false), not this function. Reviewed By: q10 Differential Revision: D110251102 fbshipit-source-id: edd2782887bbd5225bba3c687ec39241fb28a4f0
…ch#5978) Summary: Pull Request resolved: pytorch#5978 Addresses AMD benchmarking feedback on the group_index_select_2d benchmark: 1. The reported per-kernel numbers are derived from the exported Kineto trace. group_index_select_2d_bench wrapped the entire run (warm-ups + measured) in a bare profile(), so every iteration -- warm-ups included -- landed in the trace and inflated the sample count and standard deviation. Switch to a schedule-based profiling pass (schedule(wait=0, warmup=W, active=A) with prof.step()) that records ONLY the measured iterations, mirroring the existing index_select_dim0 pattern in the same file. 2. Make warm-up / measured iteration counts configurable via --num-warmups and --iters, and raise the default warm-ups from 10 to 50 for more reliable small-config measurements. Reviewed By: cthi Differential Revision: D109389728 fbshipit-source-id: 52cd7c39ff14310b8c6975aa0a674a0705b86e25
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/2892 HIP enforces a hard limit of 2^32 total threads per launch (unlike CUDA, which silently wraps; see ROCm/hip#2253). `jagged_acc_weights_and_counts_cu` launches `accumulate_weights_and_counts_kernel` with grid `ceil(total_elements / kMaxThreads=1024)` and block 1024. Total launch threads ~= `total_elements`; for `total_elements > 2^32` the launch trips `KernelLauncher::checkThreadCountNotExceeded` on ROCm. The kernel already grid-strides over `i` at lines 300-301 (`for (int i = bid * kMaxThreads + tid; i < total_elements; i += kMaxThreads * gridDim.x)`), so capping the host-side grid unconditionally on ROCm is correctness-preserving. Same one-line `#ifdef USE_ROCM` cap pattern as parent fixes D104903707 / D104937969 in `sparse_ops/`. Audit: /home/bensonma415/.llms/plans/jagged_and_more_rocm_grid_overflow_audit.plan.md Reviewed By: cthi Differential Revision: D105097581 fbshipit-source-id: 3eef8c8372163cfb8c9f320889971ada6e8c661e
…_index_select kernel
…h AutogradContext
…p16 if batch size is small (#158)
avbokovoy
force-pushed
the
abokovoi/group-index-sort-and-cache-opt
branch
from
July 6, 2026 16:23
5844b46 to
475c44d
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
16 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Follow-up of #139
The differences are:
#ifdef USE_ROCMusage in favor ofif constexpr (OPT_BOOL).