[ROCm] Fix packed-bag pooling truncation in nbit inference forward - #166
Open
aryaman-gupta wants to merge 91 commits into
Open
[ROCm] Fix packed-bag pooling truncation in nbit inference forward#166aryaman-gupta wants to merge 91 commits into
aryaman-gupta wants to merge 91 commits into
Conversation
…#6144) Summary: Pull Request resolved: pytorch#6144 Restructures `fbgemm_get_warning_flags()` so the nvcc and hipcc warning lists are derived from the same source as the host CC list, and prints all three resolved lists in the `BLOCK_PRINT` summaries. **Produces, does not apply.** The new lists are exported and printed but not wired into `target_compile_options` or `HIPCC_OPTIONS`. Activation is a separate change, so a bad activation can be reverted without losing this refactor. `CC_FLAGS` behaviour is unchanged. Structure now: - `_cc_common` -- portable flags (`-Wall -Wextra -Werror`) - `_cc_clang_only` -- new, intentionally empty, and guarded so it is appended to the host list only when the host compiler is clang - `_cc_suppressions_common` / `_cc_suppressions_clang` / `_cc_suppressions_gcc`, appended last so suppressions win - `_cc`, `_nvcc` (each entry wrapped `-Xcompiler=<flag>`), `_hipcc` `ARG_EXTRA_CC_FLAGS` deliberately stays first, preserving the pre-refactor quirk that per-target extras remain overridable by the shared suppressions. **`_hipcc` is deliberately not `${_cc}`.** hipcc is always clang, but the ROCm CI matrix builds with gcc as the host on half its legs. A host-conditional derivation would silently drop every clang-only flag from HIP device compilation on those legs -- the one surface where these flags reach device code. `_hipcc` is therefore assembled from clang-shaped inputs unconditionally. `_hipcc` also includes `_cc_suppressions_common`. Omitting it would hand hipcc `-Werror` while withholding `-Wno-deprecated-declarations`, `-Wno-strict-aliasing`, `-Wno-sign-compare`, `-Wno-vla` and the `-Wno-error=*` entries the host path relies on; applying such a list would fail immediately. Those are portable `-Wno-*` that hipcc, being clang, accepts. The clang suppressions are split by the clang version that made them necessary. The CXX path applies them gated on `CMAKE_CXX_COMPILER_VERSION` exactly as before; the hipcc path takes all of them, because those gates describe the host compiler and say nothing about hipcc's clang, which is a separate and generally newer toolchain. Before this list is ever applied, confirm the hipcc in the supported ROCm versions accepts every entry -- an unknown `-Wno-*` becomes an error under `-Werror`. Reviewed By: spcyppt Differential Revision: D115745559 fbshipit-source-id: 99e597612cbaa451f7eda211a3c197cf101affac
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/3056 Adds `-Waddress`, `-Wenum-compare`, `-Wmisleading-indentation` and `-Wparentheses` to `_cc_common` in `fbgemm_get_warning_flags()`. All four are implied by `-Wall`, which this list already sets, so no new diagnostic is expected. Listing them explicitly keeps them enabled if `-Wall` is ever narrowed. A comment says so, so a future reader does not delete them as redundant. **Placement matters.** These went in `_cc_common` (both compilers) rather than `_cc_clang_only` because all four are genuine gcc flags as well. The CI matrix builds with gcc too, and a clang-only flag placed in `_cc_common` makes gcc emit "unrecognized command line option", which `-Werror` turns into a build failure. The comment now states this rule at the list. Reviewed By: spcyppt Differential Revision: D115748663 fbshipit-source-id: 284c678d0eb5d464d124215529e6b53c1af3e150
Summary: Pull Request resolved: pytorch#6074 X-link: https://github.com/facebookresearch/FBGEMM/pull/2975 The four tensor copies in tensor_copy_chunk are independent, so dispatch each under its own sibling FBGEMM_DISPATCH_* instead of nesting weights -> indices -> identities -> runtime_meta. Two wins, no behavior change: 1. Readability: the flat structure drops the value_t / index_t / id_t / rm_t aliases that existed only to dodge scalar_t name-shadowing between the nested lambdas -- each copy now just uses scalar_t and reads top-to-bottom instead of 3 levels deep. 2. Fewer template instantiations: nesting stamps each inner copy once per outer type (multiplicative, ~4 x 2); siblings stamp each once per its own type set (additive, 4 + 2) -> smaller binary, faster compile. Reviewed By: chouxi Differential Revision: D113533821 fbshipit-source-id: c7fd9834636163f7f9fb3f70719607b120cc1e57
…orch#6167) Summary: Pull Request resolved: pytorch#6167 X-link: https://github.com/facebookresearch/FBGEMM/pull/2978 Replace the per-iteration raw std::thread in the non-blocking stream() dispatch with a persistent named size-1 folly::CPUThreadPoolExecutor + SemiFuture. Kills per-iter thread create/join churn, names the thread in traces, and removes the raw-thread std::terminate risk (folly captures exceptions into the future) Reviewed By: chouxi Differential Revision: D113669440 fbshipit-source-id: 9c8d837d06bf16260836eecedca0b16899bb5f5d
…orch#6076) Summary: Pull Request resolved: pytorch#6076 X-link: https://github.com/facebookresearch/FBGEMM/pull/2979 The consumer/ship path drained a hand-rolled folly::UMPMCQueue (weights_to_stream_queue_) with a pool of raw std::thread consumers that ran a 10ms-backoff polling loop try_dequeue(), sleep(10ms) on empty, retry and exited via a stop_ latch. This diff replaces both the queue and the raw thread pool with a persistent folly::CPUThreadPoolExecutor: each ship item is posted with add(...), and the dtor join()s the executor so all pending + in-flight ship tasks drain before teardown. Reviewed By: chouxi Differential Revision: D113669447 fbshipit-source-id: fa7bab0e9eb4fbb26c2ae2f7326cccd42eef33f6
…rch#6079) Summary: Pull Request resolved: pytorch#6079 X-link: https://github.com/facebookresearch/FBGEMM/pull/2982 Introduce the copy-thread pools for BOTH lanes in one diff so the interface is two-thread-group-ready from the start (per review): the UVM cache-hit copy pool (copy_executor_, sized res_num_copy_threads) and the dedicated HBM-miss drain pools (hbm_copy_executor_/hbm_dispatch_executor_, sized res_num_hbm_copy_threads), selected via a use_hbm route param on stream(). Folds the former D113835623 (HBM drain isolation) into this diff so adding the second group later is not an interface change. Reviewed By: chouxi Differential Revision: D113752426 fbshipit-source-id: 11170273ba9338f1820c90f320bb23fdccc47762
Summary: Pull Request resolved: pytorch#6168 X-link: https://github.com/facebookresearch/FBGEMM/pull/3059 Ship all shard co_setEmbeddings RPCs concurrently via collectAllRange instead of serially in a for loop; each shard isolates its own RPC failure so one dead shard does not cancel its siblings. ---- Note this is the final behavioral changes for streamer, the new architecture is described here P2447227530 Reviewed By: chouxi Differential Revision: D113917793 fbshipit-source-id: 4e89133fb6c07917ef43a4d1c0134db3d465a466
…ch#6153) Summary: Pull Request resolved: pytorch#6153 Adds `-Wmove`, `-Wpessimizing-move`, `-Wunused-label` and `-Wunused-local-typedefs` to `fbgemm_get_warning_flags()`. All four are `-Wall`-implied, so no new diagnostic is expected. **This is the first group that is not uniformly portable, and the split matters.** Placement was determined by compiling a probe with each compiler rather than by reading documentation. **The table below is about whether each compiler *accepts the flag*, not about diagnostics in FBGEMM's code.** "rejects" means the compiler refuses the command-line option itself; no source change is implied, and the `bucket` column is the mitigation. | flag | does g++ 11.5 accept it? | does clang 22 accept it? | bucket (= the mitigation) | | --- | --- | --- | --- | | `-Wpessimizing-move` | yes | yes | `_cc_common` | | `-Wunused-label` | yes | yes | `_cc_common` | | `-Wunused-local-typedefs` | yes | yes | `_cc_common` | | `-Wmove` | **rejects** (unknown option) | yes | `_cc_clang_only` -- never appended on a gcc host | | `-Wunused-local-typedef` | **rejects** (unknown option) | yes | **not used at all** -- the portable plural is used instead | Note g++ rejects an unknown `-W` as a **hard error**, not a warning, so `-Werror` is not even required to break the build. **On the typedef flag spelling.** clang's singular `-Wunused-local-typedef` is refused by g++. The plural `-Wunused-local-typedefs` is accepted by *both* compilers and produces the same diagnostic -- verified by compiling an unused-typedef probe under each and confirming a real warning rather than an "unrecognized option" message. Using the plural in `_cc_common` therefore buys coverage on the gcc leg, where the singular would have forced the flag into `_cc_clang_only` and covered only half the matrix. `_hipcc` continues to include `_cc_clang_only` unconditionally, which is correct and unchanged: hipcc is always clang regardless of the host compiler. Reviewed By: cthi Differential Revision: D115766842 fbshipit-source-id: 5c9a694a7f11652a9144262cedfe3267da06bbc1
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/3063 Pull Request resolved: pytorch#6150 Adds `-Wunused-value` to `_cc_common` in `fbgemm_get_warning_flags()`. Both compilers accept it and both diagnose; it is `-Wall`-implied (clang enables it by default, gcc via `-Wall`), so it goes in `_cc_common` and is inert on the host surface. Verified the suppression lists still come after it -- `-Wno-strict-aliasing` lands at index 13/14 against the flag's index 10. Its real value is on HIP **device** code, which reaches it via `_hipcc`, and only once that list is applied to `HIPCC_OPTIONS`. Until then this flag is host-only and should produce nothing. Reviewed By: cthi Differential Revision: D115768198 fbshipit-source-id: e8ff6844a788d1479eaf1bd3de3562cae6cdd842
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/3064 This PR adds a new **Community Ports** section to the FBGEMM README and introduces [fbgemm-ascend](https://gitcode.com/Ascend/fbgemm-ascend), a community-maintained Ascend NPU implementation of FBGEMM_GPU operators. The new section provides: - The fbgemm-ascend repository and documentation - Prebuilt wheel links on GitCode Releases and PyPI - pip and local wheel installation commands - Source build instructions fbgemm-ascend is developed and maintained independently by the Ascend team. This PR follows the discussion in pytorch#6149. ## Testing Documentation-only change. Verified the Markdown formatting and links. Pull Request resolved: pytorch#6164 Reviewed By: spcyppt Differential Revision: D116404791 Pulled By: q10 fbshipit-source-id: 63775e08005936ac5c97f613d8cffa29fb08ba5f
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/3058 Moves the `fbgemm_get_warning_flags()` call in `gpu_cpp_library()` above the "Build the Library" section, so the HIPCC list exists before `hip_add_library()` runs. **Reordering only -- no behaviour change, and nothing consumes the NVCC/HIPCC lists yet.** Why it has to move: `hip_add_library()` takes `HIPCC_OPTIONS` as a *creation-time* argument (a legacy FindHIP concept), not a target property that can be set afterwards. Today the flags are computed ~60 lines after the library is created, so a later change cannot pass them without this reorder. **Only the flag computation moves; the definitions stay.** The same block also contains target_compile_definitions(${lib_name} PUBLIC ASMJIT_STATIC PUBLIC FBGEMM_STATIC) for MSVC static builds, and `target_compile_definitions` requires the target to exist. Moving the block wholesale would make `gpu_cpp_library()` fail on MSVC with a "no TARGET" error, and `lib_name` is not even assigned until inside the section being moved past. So the block is split, the remaining section is retitled "Compilation Definitions", and both halves carry a comment pointing at the other so the split is not mistaken for an accident. Incidentally collapses a duplication: the MSVC and non-MSVC branches each called `fbgemm_get_warning_flags()` with byte-identical arguments and differed only in which output they assigned to `lib_cc_flags`. Now one call, then a two-line `if(MSVC)` selection. The function is pure -- it only writes its output variables -- so this is equivalent. Differential Revision: D115795635 fbshipit-source-id: 2e2d5a3b6bdd5358fd9662d63789c7fc9775dca4
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/3067 Pull Request resolved: pytorch#6173 Applies `_nvcc_warning_flags` to CUDA sources in `gpu_cpp_library()`. **This is the first change that actually turns a warning surface on:** `.cu` compilation previously received only `TORCH_CUDA_OPTIONS` -- no `-W` flags at all. The effective change is four lines: ```cmake if(NOT MSVC) target_compile_options(${lib_name} PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:${_nvcc_warning_flags}>) endif() ``` This covers the `.cu` **host pass** only. nvcc cannot warn device code -- that is architectural, and why HIP is the device-code gate. **Uses the list form rather than a per-flag `foreach`.** A `;`-list inside a generator expression was measured not to leak to other languages, so `foreach` would be a no-op refactor and would make this line inconsistent with the CXX line directly above it. **Guarded against MSVC.** `_nvcc_warning_flags` is derived from the gcc/clang list unconditionally, so on a Windows CUDA build nvcc would forward `-Wextra`, `-Wno-strict-aliasing` and friends to `cl.exe`, which does not accept them. The host CXX path already handles this by selecting `_msvc_flags` into `lib_cc_flags`; there is no MSVC-shaped equivalent for the nvcc list today, so the CUDA application is skipped there rather than fabricating one. **Includes one source fix.** Turning this surface on produced exactly one hard error across 316 targets: fp8fp8bf16_fast_gemv.cu:232: error: 'b' may be used uninitialized [-Werror=maybe-uninitialized] `-Wmaybe-uninitialized` is a gcc-only warning with no clang equivalent, and CI's nvcc host compiler is gcc. The diagnostic is a false positive -- `b` is assigned only on the `is_batched` path and read only under the same condition, so the code is correct; gcc cannot correlate the two checks. The fix is an initializer rather than a suppression: `unsigned int b = 0, m, n, k;` costs nothing and documents why. The only other diagnostic in the log is 378 x `-Wunused-parameter`, which is pre-existing `-Wextra` output demoted by `-Wno-error=unused-parameter` and unrelated. Differential Revision: D115805271 fbshipit-source-id: c1b5c92c1e5e70dfae0976b82fddb45c82033b29
…try (pytorch#6172) Summary: Pull Request resolved: pytorch#6172 X-link: https://github.com/facebookresearch/FBGEMM/pull/3062 Renames the UVM cache-hit lane's copy/dispatch executors and their thread-count knob to `uvm_hit_*`, for symmetry with the `hbm_*` lane added in D113752426. C++ only -- 5 files, no Python. The torchbind arg rename `res_num_copy_threads` -> `res_num_uvm_hit_copy_threads` is technically BC-breaking (`Argument::isBackwardCompatibleWith` compares names before anything else), but no caller is affected: every construction site in fbsource is positional with 7 of 11 args, so params 8-11 always take their `torch::arg` defaults, and nothing anywhere names the knob. Arity is unchanged, and a name change is only observable to keyword callers. Worth noting the asymmetry with the BC freeze a few lines below in the same registration block, which keeps `join_stream_tensor_copy_thread` exposed under its old name for the frozen fbgemm clones in published models. That freeze applies because *method* names bind by name unconditionally; *arg* names bind only for keyword invocation, and there are none. Doing this now rather than later is deliberate. The descendant diffs turn this knob into a `fused_params` string key in deployed model configs, and an unconsumed `fused_params` key is silently dropped on the SSD path -- so a rename after that point would quietly revert a tuned knob to its default instead of failing loudly. Reviewed By: chouxi Differential Revision: D113848362 fbshipit-source-id: 7f9ef05efab009781c70005a6e4d62063260b5c6
Summary: The docs CI job runs Sphinx linkcheck over the documentation and fails the build on any broken link. `rocm.docs.amd.com` rate-limits requests from the shared-IP CI runners, so its two links (BuildInstructions, Releases) intermittently come back `429 Too Many Requests` and fail the docs job on PRs that touch no documentation at all. The checker's own retry already backs off (`-rate limited- ... sleeping...`) and still gets throttled. Add the domain to `linkcheck_ignore`, alongside the existing gcc bugzilla entry that exists for the same reason. While reformatting the list, also escape the unescaped `.`s in that pre-existing gcc pattern -- these entries are regexes, so a bare `.` silently matches any character. The links remain in the docs; they are just no longer a CI signal, since their availability tracks AMD's rate limiter rather than the health of these pages. ___ Differential Revision: D116532863 fbshipit-source-id: 893d4a11bbd3d01884863ccca77b6d59e8bf9383
Summary: Pull Request resolved: https://github.com/facebookresearch/FBGEMM/pull/3073 Reviewed By: zmpyzmpy Differential Revision: D116417022 fbshipit-source-id: de9605a67baa71b3f3bfa6d66b3ba6a1e5cf765c
Summary: Pull Request resolved: pytorch#6183 Reviewed By: zmpyzmpy Differential Revision: D116416678 fbshipit-source-id: a720050ef96d0158d56ddf3b2f1bd593f62cd716
Summary: Pull Request resolved: https://github.com/facebookresearch/FBGEMM/pull/3075 Reviewed By: zmpyzmpy Differential Revision: D116416842 fbshipit-source-id: 725a495dc38d1901c5c2f157c425c5950dd61e56
Summary: Pull Request resolved: pytorch#6182 X-link: https://github.com/facebookresearch/FBGEMM/pull/3072 Reviewed By: zmpyzmpy Differential Revision: D116416498 fbshipit-source-id: 9c9295078d1d4ee078526d30079e3fca88e7a679
Summary: Pull Request resolved: pytorch#6179 X-link: https://github.com/facebookresearch/FBGEMM/pull/3049 ### Request generation - Add `generate_requests_for_grouped_tables` to `fbgemm_gpu.tbe.utils` and export it from the package. - Model `T` and `Es` as physical-table metadata while `Ls` and `feature_table_map` describe logical lookup features. - Allow `len(Ls) > T` so multiple fixed-`B` lookups can share a physical table. - Generate each feature's indices within `Es[feature_table_map[feature]]` and construct offsets for its configured bag size. - Support uniform and Zipf indices plus optional per-sample weights. - Keep file-backed requests, VBE, reuse, pruning, and dtype conversion out of the new focused helper. - Validate mapping lengths, table-index bounds, physical-table coverage, pooling factors, and table sizes. ### Triton TBE benchmark - Add `--feature-table-map`, where entry `i` maps logical feature `i` to a zero-based physical table index. - Default to the existing one-feature-per-table mapping when the option is omitted. - Use `generate_requests_for_grouped_tables` for every non-VBE request. - Pass the feature mapping to both `SplitTableBatchedEmbeddingBagsCodegen` and `TritonTableBatchedEmbeddingBags`. - Calculate feature dimensions, accessed bytes, bag-size hints, and backward gradient shape from the logical feature mapping. - Keep VBE on `generate_requests` and require its mapping to remain one-to-one. - Reject nonzero `--reuse` for grouped-table request generation rather than silently ignoring it. Reviewed By: JChunX Differential Revision: D115828870 fbshipit-source-id: c6f29bae5e98fee34a8e7dceb7758f2f437f734d
aryaman-gupta
changed the title
avbokovoy
approved these changes
Aug 19, 2026
avbokovoy
left a comment
avbokovoy
left a comment
There was a problem hiding this comment.
Strange thing, but I recall addressing this issue a while ago. Anyway, nice catch and great changes
| self._execute_nan_zero_fill(weights_ty, D, output_dtype, weighted) | ||
|
|
||
| @unittest.skipIf(*gpu_unavailable) | ||
| def test_nbit_forward_packed_bags_uneven_pooling(self) -> None: |
There was a problem hiding this comment.
Should it be guarded with @skipIfNotRocm?
Author
There was a problem hiding this comment.
I previously thought the test could be useful on CUDA also, but you're right that it doesn't make sense to have a test with packed_bags in the name running on CUDA. Done in 7e12fd6
Summary: Pull Request resolved: pytorch#6178 X-link: https://github.com/facebookresearch/FBGEMM/pull/3070 ## Request generation - Extend `generate_requests_for_grouped_tables` to accept either a scalar batch size or a list of per-feature batch sizes. - Preserve the existing fixed-batch behavior for scalar input, including leaving `Bs_per_feature_per_rank` unset. - For list input, validate that every batch size is positive and that the list length matches the logical feature count. - Generate feature-major lengths and offsets using each feature's batch size. - Populate `Bs_per_feature_per_rank` as `[[B0], [B1], ...]`, modeling the common one-rank VLE/VBE workload. - Continue generating each feature's indices within the embedding count of its mapped physical table, with uniform or Zipf distributions and optional per-sample weights. ## Triton TBE benchmark - Allow `--batch-size` to be either one value or a comma-separated list when `--vbe` is selected. A scalar value is replicated across all logical features. - Use the grouped-table request generator for both fixed-batch and VLE/VBE requests. - Model the common VLE/VBE path as one feature per table and one rank. The benchmark infers the identity feature-to-table mapping, so VLE/VBE commands do not need `--feature-table-map` or `--vbe-num-ranks`. - Remove randomized `--sigma-B` and multi-rank `--vbe-num-ranks` request generation from this benchmark path. - Reject multiple `--batch-size` values unless `--vbe` is enabled. - Calculate accessed weight bytes as the sum of each feature's `batch_size * embedding_dim * bag_size` contribution. - Report per-feature batch sizes while preserving the existing scalar `B` field for non-VBE benchmark statistics. Reviewed By: JChunX Differential Revision: D115936874 fbshipit-source-id: 7c82b54624577baeb31c4ff69ef449aa6cb8de74
…empty buffer (pytorch#6169) Summary: Pull Request resolved: pytorch#6169 X-link: https://github.com/facebookresearch/FBGEMM/pull/3060 On aarch64 `TransposedRequantizeTest` reports 821 of 841 cases as FAILING. None of them is a real numeric defect -- the test is comparing against a buffer that was never written. The implementation under test, `trRequantizeOpt`, is written entirely in AVX2 intrinsics (`src/spmmUtilsAvx2.cc`: `__m256`, `_mm256_*`) and has no portable or arm counterpart. To keep the arm build compiling, the `TESTCODE` macro is split on `#ifdef __aarch64__` so the arm variant computes only the reference and never calls `trRequantizeOpt`. But the test body still ends in an unconditional ASSERT_EQ(output_ref, output_test) << "reference doesn't match with test"; so on arm it compares a correctly-computed reference against a zero-initialized `output_test`. That fails for every parameterization whose reference output is not all zeros -- hence 821, with the ~20 "passes" being the cases that happen to requantize to all zeros. A representative failure: TransposedRequantizeTest.cc:300: Failure Expected equality of these values: output_ref Which is: { 169, 179, 74, 136 } output_test Which is: { 0, 0, 0, 0 } This adds a `GTEST_SKIP()` at the top of the test body on aarch64, so the suite reports what is actually true: the kernel is unported on this architecture, not broken. The `#ifdef` split in `TESTCODE` stays -- the body after `GTEST_SKIP()` is still compiled, so `trRequantizeOpt` must remain unreferenced on arm -- and now carries a comment saying why it exists. This is deliberately a truth-in-reporting fix, not a port. Actually running this suite on arm requires a portable `trRequantizeOpt`, which is a substantial piece of work (a full requantization kernel with per-tensor and per-out-channel granularity, symmetric/asymmetric act and weight, optional bias and relu) and is left for a follow-up. ___ Differential Revision: D116385150 fbshipit-source-id: 9ef1c73880049bda6cf73dfb0a3abe93ca83d2ed
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/3076 Pull Request resolved: pytorch#6185 # TLDR Adds the TBE **inference** benchmark tooling: a runner, a shape sweep, a trace-driven replay, and an analyzer, plus the benchmark-side options they need. The sweep covers 48 distinct inference shapes (96 runs, each shape at two Zipf alphas). # Changes - `tbe_inference_benchmark.py` - `--print-kernel-summary` on `nbit_device` and `nbit_device_with_spec`, matching the training benchmark. `kineto_trace_profiler` takes a `print_summary` flag and prints the top-10 kernel table - kernel-name filter changed from `embedding_codegen_forward` to `codegen_forward`. The nobag kernels are named `..._split_embedding_nobag_codegen_forward_...`, so the old substring missed them entirely and the `kernel_time > 0` assert fired. Both spellings match the pruned-lookup kernels, so nothing new is double-counted - that assert now names the kernel it looked for instead of failing bare - `warmup_runs >= runs_of_iters` is raised to `warmup_runs + 1` with a warning, instead of leaving zero timed runs - dropped the duplicate `--warmup-runs` / `--warmup-ms` declarations on `nbit_device`; the documented ones further down are the live pair - `split_table_batched_embeddings_benchmark.py` - `--print-kernel-summary` help said "Whether the table is weighted or not" - `run_tbe_inference_benchmark.sh` - single-shape runner: arch/precision/shape flags, ncu, trace export, per-run log. Unknown flags are rejected rather than silently ending argv parsing (`* ) break` used to discard everything after the first unrecognised token, `-c/--arch` included), and every value-taking flag is checked before `shift 2`, which otherwise spins the loop forever on a missing value - `run_tbe_inference_sweep.sh` - two shape families, selected with `--grid all|large-table|many-tables`: | family | shape | grid | shapes | runs | | -- | | large-table | T=2, D=256, E=1e8, int8; long bags | B(3) x L(4) | 12 | 24 | | many-tables | B=2048; int4 and int8, small/medium tables | (dtype,D)(6) x L(3) x E(2) | 36 | 72 | - each shape runs at alpha 1 (uniform) and 1.15 (skewed), which is what doubles runs over shapes - unknown flags are still forwarded to the runner, as documented; the runner now rejects the ones it does not know, so a typo surfaces instead of silently truncating the argument list - `run_tbe_inference_from_trace.sh` - replays shapes captured from a trace rather than a grid - `analyze_bench_inference.py` + `BUCK` target - parses the inference traces/logs into a comparison table, same shape as the training analyzers - `external/run_tbe_inference_amd.sh` - standalone OSS/AMD runner; `run_gis_bench_amd.sh` moved alongside it under `external/` # Usage ``` bash ${HOME}/fbsource/fbcode/ai_codesign/nonprod/supadchaya/scripts/fbgemm/run_tbe_inference_sweep.sh -c mi350 bash ${HOME}/fbsource/fbcode/ai_codesign/nonprod/supadchaya/scripts/fbgemm/run_tbe_inference_sweep.sh --grid large-table -c mi350 python3 ${HOME}/fbsource/fbcode/ai_codesign/nonprod/supadchaya/scripts/fbgemm/analyze_bench_inference.py --dir ./baseline --dir ./new --names baseline,new ``` Reviewed By: q10 Differential Revision: D101591048 fbshipit-source-id: ccd321aacf6cfca13a4db77c5cb89e8a061d129e
…ch#6155) Summary: Pull Request resolved: pytorch#6155 X-link: https://github.com/facebookresearch/FBGEMM/pull/3051 Subplan 02c chunk 02c.2 (T169200065), §4.3. Passes the warning set to `hip_add_library()` with `-Werror` stripped, so ROCm CI reports diagnostics on HIP **device** code without failing the build. OSS HIP compilation has never had `-Wall`/`-Wextra` at all -- it received only `-Wno-*`. This is the surface Richard's ask is actually about: hipcc is clang, so unlike nvcc it warns device code. ```cmake set(_hipcc_recon ${_hipcc_warning_flags}) list(REMOVE_ITEM _hipcc_recon -Werror) hip_add_library(${lib_name} ${args_TYPE} ... HIPCC_OPTIONS ${_hipcc_recon} ${HIP_HCC_FLAGS} ${lib_hipcc_system_includes} ${args_HIPCC_FLAGS}) ``` **Warn-only is the one deliberate exception to straight-to-`-Werror`** in this project, and it is narrow: the surface has never been compiled with warnings, and there is no local ROCm build to size the work with. 02c.4 deletes the two `_hipcc_recon` lines and passes `_hipcc_warning_flags` directly. The comment says so at the point of use, so the temporary form is not mistaken for the final one. **Ordering is load-bearing.** The warning flags go *before* `HIP_HCC_FLAGS`, because the tail of `HIP_HCC_FLAGS` is the HIP-specific `-Wno-*` block appended by `RocmSetup.cmake` -- `-Wno-format`, `-Wno-cuda-compat`, `-Wno-unused-result` and friends. If `-Wall` came after them it would re-enable exactly what HIP deliberately turned off. There is a comment against reordering. Note the call already carried `${lib_hipcc_system_includes}` from subplan 01, which §4.3's snippet predates; it is preserved in position. **Verified the §2 caveat before landing.** `Hip.cmake:96` pulls `${CMAKE_CXX_FLAGS}` into HIP compilation, and the plan warns that a warning flag leaking in there would arrive at HIP twice in different positions. Audited `setup.py`, which is what populates it: `cxx_flags` contains include paths, `-stdlib=libstdc++`, `-fopenmp=libgomp`, `-DROCM_VERSION=...` and similar -- **no `-W` flag anywhere**. No double arrival. **Addressed two review findings before export (V2).** 1. **The BLOCK_PRINT was lying about what HIP receives.** It still labelled both GPU lists "produced not applied" and printed `_hipcc_warning_flags`. As of this chunk the HIP branch applies `_hipcc_recon` (the same list minus `-Werror`), and as of 02c.1 the NVCC list is applied too. A reader of the configure-time output would have seen `-Werror` and concluded ROCm was error-gated when it is warn-only. That matters more here than anywhere else: the BLOCK_PRINT is the only instrument for reading what reached each surface from a CI log, so the census this chunk exists to produce would have been interpreted against the wrong flag list. Now prints three labelled entries -- NVCC as applied, the HIPCC full set, and the HIPCC set AS APPLIED -- so the withheld flag is visible rather than implied. 2. **`REMOVE_ITEM -Werror` only stripped the bare token.** The old comment claimed it "stays correct if the warning list is ever reshaped"; it does not. A targeted `-Werror=<name>` would survive and the surface would silently stop being warn-only. Replaced with `list(FILTER ... EXCLUDE REGEX "^-Werror")`, which deliberately does not match `-Wno-error=<name>` -- those are inert once `-Werror` is gone and removing them would be a behaviour change. Also corrected a third stale comment ("nothing consumes the NVCC/HIPCC ones yet"), which 02c.1 and this chunk both falsified. **The filter change is behaviour-neutral today**, which matters because it must not alter what the census measures: ``` list shape old REMOVE_ITEM new FILTER today n=6, no -Werror* n=6, no -Werror* (identical) hypothetical + -Werror=shadow n=7, SURVIVES <-bug n=6, removed -Wno-error=* entries preserved preserved ``` --- **V3 -- also strips `-Werror` from the nvcc list, to fix the OSS CUDA CI break.** 02c.1 (D115805271, landed) began forwarding the host warning set to nvcc as `-Xcompiler=<flag>`, `-Werror` included. That breaks the OSS CUDA build: ``` /tmp/tmpxft_00020cd7_00000000-6_jagged_unique_indices.compute_100f.cudafe1.stub.c:22:1107: error: integer constant is so large that it is unsigned [-Werror] ``` The failing translation unit is **nvcc's own generated stub**, not a file in this repo. nvcc writes template arguments into the stub as bare decimal literals, and `jagged_unique_indices.cu:624` instantiates ```cpp compute_hash_size_kernel<index_t, std::numeric_limits<index_t>::min()> ``` which for `index_t = int64_t` puts `9223372036854775808` in the stub. That does not fit a signed 64-bit type, so GCC warns, and `-Xcompiler=-Werror` promotes it. **This warning cannot be demoted.** Every other noisy warning here is handled with a targeted `-Wno-error=<name>`; this one has no `-W<name>` at all. Verified directly: ``` $ gcc -c -Werror -Wno-error=overflow wtest.c wtest.c:1:15: error: integer constant is so large that it is unsigned [-Werror] ``` Note the bare `[-Werror]` -- no option name in the brackets, and the demotion has no effect. So the only lever is `-Werror` itself, and nvcc gets the same reconnaissance treatment the HIPCC list already gets in this diff: ```cmake set(_nvcc_recon ${_nvcc_warning_flags}) list(FILTER _nvcc_recon EXCLUDE REGEX "^-Xcompiler=-Werror") ``` The anchor is `^-Xcompiler=-Werror` rather than `^-Werror`, because the nvcc entries are wrapped. Same rationale as the HIPCC filter otherwise: it catches a future `-Werror=<name>` and deliberately does not match `-Xcompiler=-Wno-error=<name>`, which is inert once `-Werror` is gone. The BLOCK_PRINT now shows the NVCC full set and the NVCC as-applied set, matching the treatment the HIPCC lists already got in V2. **Not fixed here, deliberately.** The NTTP at `jagged_unique_indices.cu:577` (`template <typename index_t, auto min_value>`) carries no information not already derivable from `index_t` -- it is used once, at line 600, as `index_t t_max = min_value`. Dropping it would shorten the mangled name and remove this particular literal. But it would not make `-Werror` safe on nvcc in general: the constants live in generated stubs across the whole CUDA surface, and the second diagnostic in the same log is on a `PackedTensorAccessor32<double, ...>` typedef that this repo does not parameterize. Sizing that is what the reconnaissance is for. Reviewed By: cthi Differential Revision: D115811780 fbshipit-source-id: 8a18c08b222cff0bcba538ac0d7b03dbed3ad9a6
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/3080 Pull Request resolved: pytorch#6174 Adds `-Wuninitialized`, `-Wimplicit-fallthrough`, `-Winfinite-recursion` and `-Wself-assign`. **These do not all go in one bucket.** Measured with a compiler probe: | flag | g++ 11.5 | clang 22 | bucket | | --- | --- | --- | --- | | `-Wuninitialized` | OK | OK | `_cc_common` | | `-Wimplicit-fallthrough` | OK | OK | `_cc_common` | | `-Winfinite-recursion` | **hard error** | OK | `_cc_clang_only` | | `-Wself-assign` | **hard error** | OK | `_cc_clang_only` | "init and control flow" is a semantic grouping, not a compiler claim -- putting all four in `_cc_common` would have broken the gcc leg. **`-Winfinite-recursion` needed more checking than the others.** There is no `-Wno-infinite-recursion` anywhere in the OSS CMake, and no per-file properties on `PackMatrix.cc`. Rather than assume that was an oversight, I measured: `-Winfinite-recursion` is `-Wall`-implied in clang (0 diagnostics at baseline, 1 with `-Wall`, 1 explicit), and g++ does not have the warning even under `-Wall`. So the clang leg has had it active all along and is green, which is exactly why no suppression exists. Adding it explicitly is a no-op. Differential Revision: D115842468 fbshipit-source-id: be3aecb0ab575e3e17d38e56be3ae4eb22bb0255
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/3071 The GenAI wheels build matrix is generated by pytorch/test-infra@main, which recently began emitting a `rocm7.14` variant. FBGEMM's Nova build scripts do not know that version: `.github/scripts/nova_dir.bash` handles rocm 7.0, 6.4, 6.3 and 6.2, and unknown versions take a fallback path that completes without placing a wheel in `dist/`, so the job fails on `pip install ... dist/` ("Neither 'setup.py' nor 'pyproject.toml' found") -- deterministically, on every pull request, with the `upload` job cascading. Filter the variant out of the matrix until real rocm 7.14 support is added, exactly as this filter step already drops CUDA 13.2/13.4 ("not supported by FBGEMM yet ... Drop until each is supported"). Scoped to the GenAI workflow: the non-GenAI wheels workflow's rocm 7.14 job passes. The alternative -- and the eventual right fix -- is adding a rocm7.14 case to `nova_dir.bash`, which is a decision about which gfx targets 7.14 wheels should build for; leaving that judgement to the fbgemm maintainers rather than guessing here. Differential Revision: D116566622 fbshipit-source-id: dc9cd7799b0ccd7afe91b111f27f160be2ca1a81
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/3083 Pull Request resolved: pytorch#5712 ### TLDR; This diff adds the `hip_mixed_d_warp` kernel to improve TBE backward performance on AMD. It complements `hip_warp` with mixed-D (tables with non-uniform embedding dims) and VBE support. ### Warp kernel selection flag `FBGEMM_TBE_ROCM_HIP_BACKWARD_KERNEL=1` ``` if FBGEMM_TBE_ROCM_HIP_BACKWARD_KERNEL=1 if (vbe || mixed_D) → launch hip_mixed_d_warp (introduced in this diff) else !mixed_D && D in [64, 128, 160, 192, 256, 320] && weights_on_HBM → launch hip_warp (introduced previously - no change for this diff) else → regular split_warp kernel ``` ### Warp kernel optimizations - New `hip_mixed_d_warp` kernel for mixed-D and VBE backward on ROCm - Batched run processing: each outer iteration handles `num_unroll = kThreadGroupSize` runs, with one lane loading each run's metadata and `SUBWARP_SHFL_SYNC` broadcasting it to the group, replacing per-run scalar loads - Momentum value preloading (`split_precomputation_preload`) to eliminate separate global-memory reads - Per-table state (`weights_offset`, `weights_placement`, `D_offsets`, `hash_size`, unique-indices offset) and optimizer-state pointers/values are resolved once per batch and passed to an overloaded `table_update_kernel`, instead of being re-derived per run - AMD `__builtin_amdgcn_readlane` for efficient broadcasting - Small-D template override (`max_D <= 128`) for the warp kernel - Backward block-size fix: `num_warp_per_row_groups = (kBackwardMaxThreads / 2) / 32` with grid cap `blockSize.x * blockSize.y`, keeping 4 AMD wavefronts per block (matching baseline `blockDim.y = 4`) and removing an earlier backward regression - Block dims are materialized only after the shared-memory sizing step, which can shrink `num_warp_per_row_groups`; the `max_D <= 128` override updates the group width so the launch always matches the kernel actually selected - The `hip_mixed_d_warp` optimizations apply when `num_unique_runs > 0 && total_L <= 2 * num_unique_runs`, evaluated on device against the current iteration; otherwise the kernel processes runs one at a time as before. This replaces an earlier host-side heuristic that cached the *previous* backward's unique-run count in a pinned host buffer, removing the pinned allocation, the per-call async D2H copy, and a `thread_local` cache keyed on a raw data pointer - Removed a blocking device-to-host sync (`sorted_linear_indices_num_runs[0].item<int32_t>()`) from the `hip_warp` grid computation, using the host-side `total_unique_indices` instead This diff should not impact any performance on NVIDIA GPUs. ### Support Variants - pooled rowwise adagrad split TBE - Not supported on dense, nobag and global weight decay variants. ### Performance MI350 (gfx950), median, Backward CTA+Warp kernel (us), mixed-D configs (`bench_list` + VBE). The wins are gated behind `FBGEMM_TBE_ROCM_HIP_BACKWARD_KERNEL`, which is off by default. Aggregate over all 42 swept configs: | kernel | no flag | with flag | | --- | --- | --- | | Backward CTA | 67810 -> 67702 (+0.2%) | 60033 -> 60230 (-0.3%) | | Backward Warp | 92191 -> 92216 (-0.0%) | 89612 -> 75259 (**+16.0%**) | | Backward CTA+Warp | 160001 -> 159918 (+0.1%) | 149645 -> 135489 (**+9.5%**) | | Total CUDA | 225069 -> 224938 (+0.1%) | 214802 -> 200502 (**+6.7%**) | Flag off is neutral; the entire gain is in the warp kernel, as expected since the CTA kernel is untouched. **No flag** (baseline vs this diff; + = this diff faster) | config | baseline | this diff | delta | | --- | --- | --- | --- | | bench_103 T=2 Ds=8 | 65 | 64 | +1.5% | | bench_103 T=2 Ds=8 | 64 | 63 | +1.6% | | bench_50 T=2 Ds=128 | 126 | 125 | +0.8% | | bench_61 T=6 Ds=128 | 131 | 133 | -1.5% | | bench_0 T=10 Ds=128 | 318 | 318 | +0.0% | | bench_5 T=10 Ds=20/128 | 332 | 331 | +0.3% | | bench_13 T=12 Ds=24/128 | 278 | 279 | -0.4% | | bench_2 T=14 Ds=12/24/128 | 306 | 306 | +0.0% | | bench_1 T=18 Ds=20/128 | 239 | 238 | +0.4% | | bench_3 T=20 Ds=20/24/128 | 351 | 351 | +0.0% | | bench_27 T=22 Ds=20/24/128 | 341 | 341 | +0.0% | | VBE a=1 | 135 | 132 | +2.2% | | VBE a=1.15 | 139 | 138 | +0.7% | No-flag numbers are expected to show neither gains nor regressions. The optimization is gated behind `FBGEMM_TBE_ROCM_HIP_BACKWARD_KERNEL`; with the flag off the dispatch falls back to the regular `split_warp` kernel, which this diff does not modify. Any movement in this table is run-to-run noise. **With flag** (baseline vs this diff; + = this diff faster) | config | baseline | this diff | delta | | --- | --- | --- | --- | | bench_103 T=2 Ds=8 | 63 | 63 | +0.0% | | bench_103 T=2 Ds=8 | 62 | 63 | -1.6% | | bench_50 T=2 Ds=128 | 98 | 95 | +3.1% | | bench_61 T=6 Ds=128 | 59 | 63 | -6.8% | | bench_0 T=10 Ds=128 | 167 | 170 | -1.8% | | bench_5 T=10 Ds=20/128 | 331 | 302 | +8.8% | | bench_13 T=12 Ds=24/128 | 279 | 205 | +26.5% | | bench_2 T=14 Ds=12/24/128 | 306 | 220 | +28.1% | | bench_1 T=18 Ds=20/128 | 238 | 179 | +24.8% | | bench_3 T=20 Ds=20/24/128 | 351 | 249 | +29.1% | | bench_27 T=22 Ds=20/24/128 | 341 | 239 | +29.9% | | VBE a=1 | 133 | 107 | +19.5% | | VBE a=1.15 | 138 | 117 | +15.2% | **Flag vs no flag** (flag speedup %, + = flag faster; delta in percentage points) | config | baseline | this diff | delta (pp) | | --- | --- | --- | --- | | bench_103 T=2 Ds=8 | +3.1% | +1.6% | -1.5 | | bench_103 T=2 Ds=8 | +3.1% | +0.0% | -3.1 | | bench_50 T=2 Ds=128 | +22.2% | +24.0% | +1.8 | | bench_61 T=6 Ds=128 | +55.0% | +52.6% | -2.3 | | bench_0 T=10 Ds=128 | +47.5% | +46.5% | -0.9 | | bench_5 T=10 Ds=20/128 | +0.3% | +8.8% | +8.5 | | bench_13 T=12 Ds=24/128 | -0.4% | +26.5% | +26.9 | | bench_2 T=14 Ds=12/24/128 | +0.0% | +28.1% | +28.1 | | bench_1 T=18 Ds=20/128 | +0.4% | +24.8% | +24.4 | | bench_3 T=20 Ds=20/24/128 | +0.0% | +29.1% | +29.1 | | bench_27 T=22 Ds=20/24/128 | +0.0% | +29.9% | +29.9 | | VBE a=1 | +1.5% | +18.9% | +17.5 | | VBE a=1.15 | +0.7% | +15.2% | +14.5 | Pull Request resolved: pytorch#5074 Test Plan: On an MI350 (gfx950) host: buck2 run @//mode/opt-amd-gpu //deeplearning/fbgemm/fbgemm_gpu/test/tbe:backward_adagrad buck2 run @//mode/opt-amd-gpu //deeplearning/fbgemm/fbgemm_gpu/test/tbe:backward_optimizers buck2 run @//mode/opt-amd-gpu //deeplearning/fbgemm/fbgemm_gpu/test/tbe:forward bash ${HOME}/fbsource/fbcode/ai_codesign/nonprod/supadchaya/scripts/fbgemm/run_fbgemm_tests.sh --test tbe Benchmark: full TBE backward sweep (`bench_list` + `tbe_sweep` + VBE) on MI350, baseline and this diff run back-to-back in one session, flag off and flag on, 0 failed configs. Numbers above are medians. The mean is unreliable for `tbe_sweep T=10 D=1024 B=131072`, which is bimodal (a few iterations run ~2.5x slow, giving an 84994 us mean against a 32531 us median); the median is stable for that config and both flag states are flat there. Full per-config perf (mean + median): https://docs.google.com/spreadsheets/d/1d9kvO782OY-HLc27iB-zf9gklTAH6bPFO2FVa44DzXk/edit Reviewed By: q10 Differential Revision: D102946325 Pulled By: spcyppt fbshipit-source-id: b934e9fa69fc3deb71eea88a6ddf982966050c15
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/3065 Pull Request resolved: pytorch#6175 Adds `-Wnull-conversion`, `-Wvexing-parse` and `-Wstring-concatenation`. Two of the three are clang-only; the third is portable. Measured: | flag | g++ 11.5 | clang 22 | bucket | | --- | --- | --- | --- | | `-Wnull-conversion` | hard error | OK | `_cc_clang_only` | | `-Wstring-concatenation` | hard error | OK | `_cc_clang_only` | | `-Wvexing-parse` | **OK** | OK | `_cc_common` | `-Wvexing-parse` goes in the portable bucket. Putting it in `_cc_clang_only` would work but would silently cover only half the CI matrix -- the same coverage loss avoided earlier with `-Wunused-local-typedefs`. When a flag is portable, the portable bucket is strictly better. This is the third group whose expected bucketing did not survive measurement, which is the argument for keeping the compiler-probe step per group rather than trusting the grouping. Reviewed By: cthi Differential Revision: D115842465 fbshipit-source-id: 48944cf1b777e62ccadc3353428e961b46282276
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/3111 Pull Request resolved: pytorch#6225 Adds `-Wunused-exception-parameter`. This warning finds a `catch` block that does not use its parameter. GCC does not know the flag. Clang knows it. The flag goes in the clang-only list. Reviewed By: spcyppt Differential Revision: D117101712 fbshipit-source-id: 0378ba5a79f4a70aac8a7ab0b29a949ac12ef92f
Summary: Pull Request resolved: pytorch#6190 X-link: https://github.com/facebookresearch/FBGEMM/pull/3079 `fbgemm::masked_select_jagged_1d` was registered for CPU, MTIA and Meta only, so any pipeline reaching it on CUDA died with a dispatch error. It is reached from the ranking preproc path via `datafm_feature_value_filter` -> `fused_mask_conversion_sparse`, which makes it a hard blocker for running that preproc on the trainer GPU. Two kernels: - a per-row length kernel that differences the mask prefix at row boundaries; - a coalesced one-thread-per-element compaction writing to `mask_prefix[i]`. Because rows tile the input contiguously, that prefix *is* the destination index, so no per-row correction is needed and input order is preserved - which the CPU kernel also guarantees and callers rely on. Reviewed By: q10 Differential Revision: D115327729 fbshipit-source-id: 0e9fefb902cdec0fcd8257901d2deed7c7862534
Summary: Pull Request resolved: pytorch#6236 X-link: https://github.com/facebookresearch/FBGEMM/pull/3108 # TLDR The "Not enough bits to accommodate B" check gave no numbers, so hitting it told you nothing about how far over the limit you were or why the limit was what it was. Add the operands to the message. `b_t_map` packs the feature index `t` and batch index `b` into one 32-bit word. `T` is fixed at TBE init, so its bit width is known up front and the remainder goes to `B` (see D69387123). When `max_B` exceeds that remainder, this is the check that fires. # Changes - `codegen/training/pt2/embedding_split_host_pt2_autograd_template.cpp` - `TORCH_SYM_CHECK` for `max_B_ <= info_B_mask` now reports: - `max_B_` — the offending value - `info_B_mask` — the actual maximum allowed - `info_B_num_bits` — the bits left for `B` - `T` — the feature count the bit split was derived from - a `[VBE]` / `[const B]` prefix via `{{ "VBE" if vbe else "const B" }}` Example rendered message: ``` [VBE] Not enough bits to accommodate B: max_B = 40000000 exceeds max allowed B = 33554431 (info_B_num_bits = 25, derived from T = 101 features) ``` # Notes - `TORCH_SYM_CHECK(cond, ...)` is variadic and forwards to `TORCH_CHECK` (`c10/core/SymBool.h:101`), so no `c10::str` wrapping is needed. - `max_B_` and `T` are both `c10::SymInt`, which has an `operator<<` (`c10/core/SymInt.h:503`). Both are already streamed in the Kineto annotation a few lines above, so this adds no new requirement. - `T` is labelled "features", not tables: it is `len(feature_table_map)` (`split_table_batched_embeddings_ops_training.py:951`), which is `>=` the table count when tables are shared across features. Reviewed By: q10 Differential Revision: D117470122 fbshipit-source-id: f80d456f740dc4d8bdc337f75c15c4b9e962d70c
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/3123 Pull Request resolved: pytorch#6224 `tbe:ssd_split_tbe_training` fails with: ``` hypothesis.errors.DeadlineExceeded: Test took 59423.60ms, which exceeds the deadline of 200.00ms. ``` on `test_ssd_cache_flush`. A single example legitimately takes ~59s -- it drives the full SSD cache prefetch workflow with excessive flushing -- so the 200ms default deadline is meaningless here. `test_ssd_cache_flush` is the only test in `SSDSplitTBETrainingTest` with a `given` but no `settings` at all, so it silently inherited Hypothesis's default deadline. Every sibling cache test already carries: ``` settings( verbosity=Verbosity.verbose, max_examples=MAX_PIPELINE_EXAMPLES, deadline=None ) ``` (`test_ssd_cache_implicit_prefetch`, `test_ssd_cache_explicit_prefetch`, `test_ssd_cache_pipeline_before_fwd`, `test_ssd_cache_pipeline_between_fwd_bwd`). This adds the same decorator but with an explicit `deadline=timedelta(minutes=1)` rather than the siblings' `deadline=None`, so the test still has an upper bound and a genuine hang is caught rather than running forever. **Caveat on the 1 minute bound.** Hypothesis applies `deadline` *per example*, not to the whole test. The slowest example measured so far took 59,423ms, which leaves roughly 577ms -- about 1% -- of headroom under a 60,000ms deadline. The verification run below passed, but this margin is thin enough that the test is likely to flake on a loaded or slower host. If that happens, the options are to raise the bound (2-3 minutes would still catch a real hang) or fall back to `deadline=None` to match the sibling tests. No production or kernel code changes -- this only corrects the test's Hypothesis configuration. The assertions themselves are untouched. Reviewed By: q10 Differential Revision: D117003216 fbshipit-source-id: dbd749ad56baeea2b6a44b408f56a3dce93d4b5d
Summary: Pull Request resolved: pytorch#6234 NOTE: No linked task. Please associate a task with this diff. X-link: https://github.com/facebookresearch/FBGEMM/pull/3121 Adds `-Wheader-hygiene`. This warning finds a `using namespace` directive in a header file. The host build does not see this warning. The host build uses `-isystem` for the third-party headers, and `-isystem` stops the warning completely. The device build does see it, because HIP compiles the same headers without `-isystem`. The change therefore adds `-Wno-error=header-hygiene` to the HIP list. The warning stays visible, but it does not stop the build. Remove the `-Wno-error=` line after the third-party headers become system headers. Reviewed By: spcyppt Differential Revision: D117101708 fbshipit-source-id: 2dff2846c8c5e8b0ecac8c657b0c3549f554f22d
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/3125 Pull Request resolved: pytorch#6239 NOTE: No linked task. Please associate a task with this diff. Adds `-Wshadow` with `-Wno-error=shadow`. The warning becomes visible, but it does not stop the build. GCC 11.5 and Clang 22 both accept the flag and the `-Wno-error=` form. A compiler test shows this. Both go in the portable list. The sub-options `-Wshadow-all`, `-Wshadow-uncaptured-local` and `-Wshadow-field-in-constructor` are clang-only. This change does not use them. The file already contained `-Wno-error=shadow` in a list for Clang versions above 17. That line did nothing, because `-Wshadow` was absent. A `-Wno-error=` line has no effect if its warning is not on. This change removes the old line and adds the flag and the `-Wno-error=` form together. Reviewed By: cthi Differential Revision: D117101706 fbshipit-source-id: 1e412c4f14dff85c9b5c532b3f78685fae58ac74
…nd/codegen (D1) (pytorch#6241) Summary: Pull Request resolved: pytorch#6241 Backend (C++/CUDA codegen) half of the example_counter pruning-state work, rebased onto the landed D113869217 (TBE backward-arg TensorList packing). Adds the optional per-sample `example_counter` state to the existing `rowwise_adagrad_with_counter` optimizer at the C++/codegen layer (no new optimizer, no new WeightDecayMode): - optimizers.py: example_counter (is_optional) + use_example_counter + example_counter_halflife on the op; gated example_counter update in the GPU precomputation (segment_length passed as float, no static_cast); V1 schema string frozen. - kernel templates (cta / warp / device / standalone): segment_length threading (float) + guarded example_counter pointer setup (nullptr when off). The CTA long-run path passes run_length unconditionally — the weight/state update runs exactly once per run (sole CTA, or last CTA via grad_accum_counter), so there is no double-count and split runs are not missed. - host templates (gpu + cpu): synthesize inert example_counter defaults on the frozen V1 path (mirrors the learning_rate_tensor V1 shim); both shims key on the same use_example_counter sentinel. - CPU backward template: emit torch-schema bool casing (False/True) in the generated m.def; TORCH_WARN_ONCE when use_example_counter is set on CPU (example_counter is maintained on GPU only). - MTIA (fbgemm_mtia_ops.cpp): the 4 hand-written rowwise_adagrad_with_counter *_pt2_mtia_wrapper registrations mirror this op's flat PT2 schema, so they are extended to accept (and ignore) the 7 new example_counter args, matching the generated 59-arg schema. example_counter is GPU-only pruning state; MTIA does not maintain it. Without this, the op fails registration (59 vs 52) in every binary linking torch_mtia. Mirrors how D113869217 synced this same file for the packing. The Python-generating templates (lookup_args / invoker) and all runtime Python live in the stacked frontend diff D114133366. Existing optimizers are byte-unchanged; the counter path is byte-identical when use_example_counter is unset. Reviewed By: gchalump, q10, spcyppt, trirpi Differential Revision: D112412933 fbshipit-source-id: 8a63b4c97e72f4cb62f66d35f6aa7af929b4af82
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/3126 NOTE: No linked task. Please associate a task with this diff. Pull Request resolved: pytorch#6218 Adds `-Wzero-as-null-pointer-constant` with the matching `-Wno-error=` form. The warning becomes visible, but it does not stop the build. GCC 11.5 and Clang 22 both accept the flag and the `-Wno-error=` form. A compiler test shows this. Both go in the portable list. The warning finds a literal `0` that the code uses as a null pointer. Use `nullptr` instead. Reviewed By: spcyppt Differential Revision: D117101711 fbshipit-source-id: 3d55879e1c8072a5ac8264e7e5293d9a10f5d9ad
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/3119 Pull Request resolved: pytorch#6232 NOTE: No linked task. Please associate a task with this diff. Move `bench_utils.py` from the benchmark-only source directory into the canonical `fbgemm_gpu/bench` Python package. This prevents the source checkout from shadowing the wheel-installed module and ensures `setuptools.find_packages()` includes `fbgemm_gpu.bench.bench_utils`. Update the ten benchmark entry points that used the old local import to use the canonical package import. Keep the existing root `//deeplearning/fbgemm/fbgemm_gpu:bench_utils` target as the owning Buck provider so downstream labels remain unchanged, and remove the now-unused child BUCK boundary. Reviewed By: gchalump, spcyppt Differential Revision: D117744258 fbshipit-source-id: 4d9ca38e039c923f90e1ebc422a28a83a7a08eae
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/3127 Pull Request resolved: pytorch#6223 NOTE: No linked task. Please associate a task with this diff. Update the fbcode and xplat FBGEMM OSS Composable Kernel pins from `7fe50dc3da2069d6645d9deb8c017a876472a977` to `fcc9372c009c8e0a23fece77b582da83b04a654f`, exactly matching the MSLK pin introduced by parent D117441615. CK 1.2 places `hip_check_error` in namespace `ck`, while the existing internal CK dependency exposes it globally. Import `ck` locally in `flush_icache_ck` so the same source builds with both header layouts. Keep the existing FBGEMM BUCK dependencies unchanged because migrating every CK extension target would require unrelated CK API changes. Reviewed By: spcyppt Differential Revision: D117467079 fbshipit-source-id: f4cc426f8ac48c045c51fdd950843ba6b3c7c2d4
…er — backend/codegen (D1)" (pytorch#6242) Summary: Pull Request resolved: pytorch#6242 X-link: https://github.com/facebookresearch/FBGEMM/pull/3128 Original commit changeset: 8a63b4c97e72 Original Phabricator Diff: D112412933 Reviewed By: ljyuva83 Differential Revision: D117968388 fbshipit-source-id: 15480808c475c337ced7e9e6a15c366995b7c1b7
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/3129 NOTE: No linked task. Please associate a task with this diff. Pull Request resolved: pytorch#6230 Adds `-Wshorten-64-to-32` with the matching `-Wno-error=` form. The warning becomes visible, but it does not stop the build. Clang knows the flag. GCC does not know it. A compiler test shows this. The flag and the `-Wno-error=` form both go in the clang-only list. Keep the `-Wno-error=` form in the clang-only list. GCC stops with an error if it gets `-Wno-error=shorten-64-to-32`, because GCC does not know the warning. GCC gives this message: error: '-Wno-error=shorten-64-to-32': no option '-Wshorten-64-to-32' The warning finds a 64-bit value that the code puts into a 32-bit variable. Reviewed By: cthi Differential Revision: D117101707 fbshipit-source-id: 4cad32ca9719bdb9e916e1361c0a8f93f8349009
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/3120 Pull Request resolved: pytorch#6233 The RES (raw embedding streaming) HBM-lane tests in res_enabled_tables_test.py reach torch.ops.fbgemm.masked_index_select, which is registered in the SSD split-embeddings-cache extension (src/ssd_split_embeddings_cache/). That directory is compiled only by the internal Buck build; the OSS CMake wheel excludes it entirely, so the op is absent in OSS. Every HBM-lane test therefore dies on an _OpNamespace lookup ("'fbgemm' object has no attribute 'masked_index_select'"), plus cascading _res_* attribute errors, turning the OSS 'main' GPU CI red (30 failed / 13 passed). Guard the single chokepoint all HBM-lane tests construct through (_build_mixed_tbe, reached directly and via _drain_tbe/_device_tbe) with a SkipTest when masked_index_select is not registered in the running wheel. On any build that has the op (internal, and OSS once the extension is shipped) the guard is a no-op and all tests run unchanged. The non-HBM allowlist tests use _build_tbe, never touch the op, and continue to run everywhere. This is the fast unblock for OSS trunk/nightly. The proper long-term fix is to add src/ssd_split_embeddings_cache/ to the OSS CMake build so the op is present. Reviewed By: FriedCosey Differential Revision: D117750928 fbshipit-source-id: a133fb867d01c7e1d24d44aebc4f6fb9ac0664cc
…torch#6243) Summary: Pull Request resolved: pytorch#6243 X-link: https://github.com/facebookresearch/FBGEMM/pull/3130 `permute_multi_embedding_function_gpu` requires `permutes`, `in_shapes` and `out_shapes` to already sit on `pooled_embs[0]`'s device: ``` TENSORS_ON_SAME_DEVICE(permutes, pooled_embs[0]); TENSORS_ON_SAME_DEVICE(permutes, in_shapes); TENSORS_ON_SAME_DEVICE(permutes, out_shapes); ``` A statically exported (Sigmoid / nativert) graph cannot always satisfy that. Those three are graph **buffers**: their device is resolved once, when the weights are materialized, against whatever card the loading thread was on. `pooled_embs` is an **activation** produced by an all-to-one whose target device is index-less `cuda`, so it follows whichever runtime is executing the request. On a 4-card host the buffers sit on cuda:0 while `pooled_embs` arrives on cuda:0..3, and roughly three quarters of requests abort: ``` arg0 pooled_embs: GenericList [c10::Half[377, 161628]cuda:1, ] arg1 permutes: Tensor int[395, 6]cuda:0 arg2 in_shapes: Tensor int[1]cuda:0 arg3 out_shapes: Tensor int[11]cuda:0 permutes must be on the same device as pooled_embs[0]! ``` Relocate rather than fail. The three tensors are small constant index tensors -- 9,480 + 4 + 44 = 9,528 B here, about 0.008% of the 116 MiB of pooled embeddings they describe, and ~1/12000th of what the preceding all-to-one already moves across cards. The same function already relocates its own `in_ptr`/`out_ptr` a few lines below, and `kt_regroup_arguments_gpu` does it for these very tensors, so this makes the op internally consistent rather than introducing a new idiom. Relaxing a check that used to be fatal can hide a real misplacement, so the relocation path emits `TORCH_WARN_ONCE` naming both devices. A caller that was already correct sees no warning and no copy -- the lambda returns the tensor unchanged when the devices match. `permute_multi_embedding` and `regroup_keyed_tensor` both route through this function, so both are covered. Not changed: the CPU implementation, which has no such check, and `mtia/tools/experimental/Tritor/kernel_gen/docs/fbgemm/permute_multi_embedding.md`, which documents the same-device requirement as a *wrapper* contract for MTIA kernel generation and is unaffected by the CUDA implementation relaxing it. Reviewed By: q10 Differential Revision: D117957553 fbshipit-source-id: 168189b38981d2ea4e996f4688e5f3eab71a9c89
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/3081 ## Why The `linux.rocm.gpu.gfx942.1` runner label is no longer served: the ROCm GPU test jobs queue and are cancelled without running a single step, so there is currently no ROCm GPU test coverage. The build jobs still pass, which is why this went unnoticed. ## What - Switch the three ROCm GPU jobs to `linux.rocm.gpu.ecosystem.mi350.1`. - Restore gfx950 as a build target. Without it the wheel carries no MI350 code object and the first kernel launch segfaults before any test can run, so the runner change alone is not enough. - Disable the tests that fail on the new runner. This is a firefighting change to restore the signal. It deliberately does not try to fix the underlying test failures — it disables them so the pipeline can go green, and each is tracked separately with a fix in progress. The changes in this PR are derived in part from pytorch#6119 Pull Request resolved: pytorch#6186 Reviewed By: cthi Differential Revision: D116693017 Pulled By: q10 fbshipit-source-id: a23370fecfbca16ed9f68a28c758b678d0e583a0
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/3133 Pull Request resolved: pytorch#6246 NOTE: No linked task. Please associate a task with this diff. Add a reusable `torch_tensor_to_same_device(...)` function beside the tensor device-validation helpers. It returns the source unchanged when devices match; otherwise it warns and moves the source to the target device using caller-selected blocking semantics. Refactor the metadata relocation introduced by D117957553 to use the shared helper. Using an inline function avoids macro identifier capture and makes the value-returning contract explicit. Reviewed By: georgiaphillips Differential Revision: D118157970 fbshipit-source-id: 1b3630f1bfd82d4bc56383e136b0c6707a0cc907
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/3134 NOTE: No linked task. Please associate a task with this diff. Make CPU warning behavior portable when optional compiler features are disabled. Guard OpenMP directives in the mirrored `Utils.cc` sources and generated quantized CPU template, while removing invalid CUDA-style unroll pragmas from the jagged CPU helper. OpenMP-enabled behavior and the serial fallback remain unchanged. Reviewed By: gchalump Differential Revision: D117024953 fbshipit-source-id: f1ae953427cb128461ff43f0af4eb5abb6be4d1b
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/3131 Pull Request resolved: pytorch#6240 NOTE: No linked task. Please associate a task with this diff. Adds `-Wunused-variable`, `-Wunused-const-variable`, and `-Wunused-but-set-variable` to the portable OSS warning list. The first and third are already implied by `-Wall`/`-Wextra`; explicitly naming them also makes the existing `-Wno-error=unused-but-set-variable` escapes effective. `-Wunused-const-variable` adds GCC coverage. Marks configuration- and architecture-dependent header constants and mask tables `[[maybe_unused]]` when some GCC translation units intentionally omit them. This includes `avx2_ps_or_epi32_combined_mask`, which is unused in ARM CPU builds but used by x86 vectorized paths. The existing asmjit target exclusion remains because asmjit is third-party code. Reviewed By: cthi Differential Revision: D117103728 fbshipit-source-id: d9586ac4d535fe3ac7a795eb1fbb859be3e96f7c
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/3137 Pull Request resolved: pytorch#6250 NOTE: No linked task. Please associate a task with this diff. Resolve GPU warnings for parameters that are read only in particular CUDA or ROCm configurations. Annotate `KernelLauncher::checkThreadCountNotExceeded()`'s grid argument and the ROCm-only sparse group-index launch-selection values without changing signatures or launch behavior. Reviewed By: spcyppt Differential Revision: D117024949 fbshipit-source-id: a613be08a792b5c6618e45353a2156d266455a9c
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/3136 Pull Request resolved: pytorch#6249 NOTE: No linked task. Please associate a task with this diff. Adds four warnings: -Wdeprecated-implementations -Wsemicolon-before-method-body -Wimport-preprocessor-directive-pedantic -Wincompatible-function-pointer-types GCC does not know these four flags. Clang knows them. The flags go in the clang-only list. These four warnings do nothing for C++, CUDA or HIP code. They apply to Objective-C methods and to C function pointers. The change adds them so that the two flag lists stay the same. A later comparison of the two lists then gives an empty result. The change adds no `-Wno-error=` line. These warnings cannot occur, so no `-Wno-error=` line is necessary. ___ Differential Revision: D117103727 fbshipit-source-id: aa0fc41ca89fd58295ff2adbf2d53b96cb741ff4
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/3138 Pull Request resolved: pytorch#6251 NOTE: No linked task. Please associate a task with this diff. Make generated Torch schemas independent of C++ declaration rendering. The CPU backward template now consumes `split_function_schemas` directly instead of rewriting `split_function_args`, preventing declaration-only attributes such as `[[maybe_unused]]` from entering `m.def(...)` strings. Reviewed By: spcyppt Differential Revision: D117038101 fbshipit-source-id: c9476ee5cf93b5e3e07fcab0774bc53301ed3e91
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/3139 Pull Request resolved: pytorch#6252 `__verify_pytorch_gpu_integration` in .github/scripts/utils_pytorch.bash computed `torch_cuda_available` but never used it, then called `torch.cuda.get_device_capability()` and `get_device_name()` unconditionally. On the CPU lane (torch built without CUDA) these raise "Torch not compiled with CUDA enabled", failing the whole GPU-integration check and the fbgemm_gpu_ci_cpu job. Guard both device-property queries behind `torch_cuda_available == True`; print an explicit "N/A (CUDA not available)" placeholder on CPU so the report block still renders. Surfaced during the torch-2.14 CI triage (T277878947); independent of any torch version. ___ Differential Revision: D118316051 fbshipit-source-id: 20948c5c1c6036aceba889fcab8e33a3a299fc40
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/3140 Pull Request resolved: pytorch#6253 NOTE: No linked task. Please associate a task with this diff. Adds four warnings: -Wambiguous-reversed-operator -Wbitwise-instead-of-logical -Wunreachable-code-fallthrough -Wunused-local-typedef GCC does not know these four flags. Clang knows them. The flags go in the clang-only list. These four warnings can occur in this code. `-Wbitwise-instead-of-logical` finds `&` where `&&` is correct. `-Wunused-local-typedef` finds a typedef that the code does not use. The change adds no `-Wno-error=` line, so these warnings stop the build. If one of them occurs too often, add a `-Wno-error=` line to the clang-only list. Do not add it to the portable list, because GCC stops with an error if it gets `-Wno-error=` for a warning that it does not know. ___ Differential Revision: D117103729 fbshipit-source-id: 0858c073f8db54d7ed8bfd28723f8e46c89a5f11
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/3141 Pull Request resolved: pytorch#6254 Release prep for FBGEMM OSS v1.9.0 (PyTorch 2.14). Adds the "1.9": "2.14" entry to _fbgemm_torch_compat_table so fbgemm_gpu.__init__ can derive the version <-> torch compatibility at runtime and emit the correct warning for mismatched torch builds. Docs (Releases.rst) compat table + release-notes rows are handled separately in the Phase 3 docs update once the stable matrix is locked. ___ Differential Revision: D118333185 fbshipit-source-id: 82c22f26e6f035a3aba17c17fd131f496aac4cfb
In PackedMode the accumulate and store stages map lanes to bags at uint granularity while the load stage uses uint4 granularity, so the two need different entries of Ls[]. The kernel translated between them by shuffling Ls[] in place, inside the L_start loop. That corrupts the load stage, which keeps reading Ls[] on later passes, and corrupts the shuffle source lanes themselves, so from the second pass onwards the longer bag of a packed pair is truncated to its partner's pooling length. Compute the accumulate-stage lengths once into a separate Ls_acc[], leaving Ls[] intact for the load stage. The mapping is loop-invariant, so this also removes OutputRowsPerThread shuffles from every pass of the loop. Also make max_Ls wave-uniform under PackedMode. The L_start loop is wave-collective, so a per-lane bound lets the short-bag lanes exit while the remaining lanes still shuffle against them. Per-row validity checks against Ls[]/Ls_acc[] still bound each lane's own work, and a divergent wave already executes the union of all lanes' iterations, so this adds none. Fixes the INT4 D=160 subtests of test_nbit_forward_nan_zero_fill. Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
max_Ls is wave-uniform under PackedMode, so min(InputRowsInFlight, max_Ls - L_start) already evaluates identically on every lane and the shuffle that followed it is the identity. nbit_forward_test.py is unchanged at 12 passed / 8 subtests passed, and the divergent-Ls probes stay clean. Mean of 3 runs on MI350X, INT4 pooled forward: D=160 packed, ragged 65.40 -> 64.97 us D=240 packed, ragged 64.97 -> 64.50 us D=160 packed, uniform 45.73 -> 45.83 us D=1024 not packed 139.80 -> 139.40 us The non-packed config cannot be affected by this change yet moves by a similar amount, so treat ~0.3% as the noise floor and the change as perf-neutral. Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
…gths test_nbit_forward_nan_zero_fill only covers bag lengths that happen to agree between packed neighbours, so it does not catch a kernel that applies one bag's L to its partner. Add a case with alternating short/long lengths, which is the shape that exposes it. Every row is 1.0 and nothing is pruned, so each bag must sum to its own L. Verified red/green: passes on the fixed kernel, and fails without the wave-uniform max_Ls (max abs diff 60.0, 272/1280 elements). Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
Replace the hand-rolled butterfly with warp_reduce_max() from utils/find_qparams.cuh, which is the same shfl_xor reduction over kWarpSize and is already in scope via embedding_forward_template_helpers.cuh. Guard test_nbit_forward_packed_bags_uneven_pooling with skipIfNotRocm: bag packing is a ROCm-only path, so elsewhere the test would pass without exercising what its name describes. Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
aryaman-gupta
force-pushed
the
aryaman/fix-packed-bag-pooling
branch
from
September 2, 2026 11:43
7e12fd6 to
c50ea0b
Compare
These were skipped on ROCm in da743f0 while the CI migration was in flight. This change fixes the underlying failure, so the skip comes out with it. Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
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
14 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.
Problem
On ROCm,
TBE_ROCM_INFERENCE_PACKED_BAGSpacks several bags into a single wave. When those bags have different pooling lengths, the longer one is silently truncated to its partner's length and the embedding output is wrong.It only surfaces when a bag's length exceeds
InputRowsInFlight, and it disappears entirely when packed bags happen to share a length — which is why it went unnoticed.Cause
The accumulate/store stages map lanes to bags differently from the load stage (uint vs uint4 granularity). The kernel translated
Ls[]between the two by shuffling it in place inside theL_startloop. That corrupts the load stage, which keeps readingLs[]on later passes, and also corrupts the shuffle source lanes themselves — so from the second pass onwards every lane broadcasts an already-permuted value.Fix
Ls_acc[]before the loop, leavingLs[]intact for the load stage. Being loop-invariant, this also removes shuffles from every pass.max_Lswave-uniform under PackedMode. The loop is wave-collective (syncwarp,shfl_sync), so a per-lane bound let short-bag lanes exit while the remaining lanes still shuffled against them.Both changes are required — either one alone still fails.
Testing (MI350X / gfx950, ROCm 7.1)
nbit_forward_test.py: 2 failed → 12 passed, 8 subtests passedCUDA and the
nobagpath are unaffected — PackedMode is ROCm-and-pooled-only.