backward performance optimization for MI350 by liligwu · Pull Request #4925 · pytorch/FBGEMM

liligwu · 2025-09-24T00:52:30Z

bwd performance optimization for ROCm.
Fix numerical issues

netlify · 2025-09-24T00:52:38Z

✅ Deploy Preview for pytorch-fbgemm-docs ready!

Name	Link
🔨 Latest commit	`c6b0a88`
🔍 Latest deploy log	https://app.netlify.com/projects/pytorch-fbgemm-docs/deploys/6913709656572c000849087c
😎 Deploy Preview	https://deploy-preview-4925--pytorch-fbgemm-docs.netlify.app
📱 Preview on mobile	Toggle QR Code... Use your smartphone camera to open QR code link.

To edit notification comments on pull requests, go to your Netlify project configuration.

facebook-github-bot · 2025-09-24T02:17:01Z

@haoyuz has imported this pull request. If you are a Meta employee, you can view this in D83116315.

q10 · 2025-10-02T00:18:18Z

@liligwu we're seeing

OSError: libtbb.so.12: cannot open shared object file: No such file or directory

We already install tbb here, so it might just be an issue of updating the build scripts to put libtbb in the LD_LIBRARY_PATH

meta-codesync · 2025-10-13T20:54:01Z

@q10 has imported this pull request. If you are a Meta employee, you can view this in D83116315.

liligwu · 2025-10-13T20:55:59Z

@liligwu we're seeing
OSError: libtbb.so.12: cannot open shared object file: No such file or directory
We already install tbb here, so it might just be an issue of updating the build scripts to put libtbb in the LD_LIBRARY_PATH

Hi @q10 , sorry I missed your message.
I actually have this commit 4d2bfdd that links tbb explicitly, it works in your container. Do you have any suggestions for fixing this issue in CI, please?

BTW, we discovered a numerical issue in 986cceb and reverted it in 85417b4. It unblocks merging the bwd optimization first.

Thank you.

q10 · 2025-10-14T03:05:50Z

@liligwu we're seeing
OSError: libtbb.so.12: cannot open shared object file: No such file or directory
We already install tbb here, so it might just be an issue of updating the build scripts to put libtbb in the LD_LIBRARY_PATH
Hi @q10 , sorry I missed your message. I actually have this commit 4d2bfdd that links tbb explicitly, it works in your container. Do you have any suggestions for fixing this issue in CI, please?

BTW, we discovered a numerical issue in 986cceb and reverted it in 85417b4. It unblocks merging the bwd optimization first.

Thank you.

I think this commit only addresses the build step, where we need to link to tbb. However, for runtime, you might need to do a find in $CONDA_PREFIX from inside the container, and manually update LD_LIBRARY_PATH, or create a symlink (something like

FBGEMM/.github/scripts/utils_build.bash

Line 383 in 0d49628

    
           (print_exec ln -s "${conda_prefix}/lib/librhash.so" "${conda_prefix}/lib/librhash.so.0") || return 1

).

liligwu · 2025-10-15T22:47:28Z

@liligwu we're seeing
OSError: libtbb.so.12: cannot open shared object file: No such file or directory
We already install tbb here, so it might just be an issue of updating the build scripts to put libtbb in the LD_LIBRARY_PATH
Hi @q10 , sorry I missed your message. I actually have this commit 4d2bfdd that links tbb explicitly, it works in your container. Do you have any suggestions for fixing this issue in CI, please?
BTW, we discovered a numerical issue in 986cceb and reverted it in 85417b4. It unblocks merging the bwd optimization first.
Thank you.
I think this commit only addresses the build step, where we need to link to tbb. However, for runtime, you might need to do a find in $CONDA_PREFIX from inside the container, and manually update LD_LIBRARY_PATH, or create a symlink (something like

FBGEMM/.github/scripts/utils_build.bash

Line 383 in 0d49628

(print_exec ln -s "${conda_prefix}/lib/librhash.so" "${conda_prefix}/lib/librhash.so.0") || return 1

).

Hi @q10 , I can see a few line below of your example that links tbb

FBGEMM/.github/scripts/utils_build.bash

Line 388 in 0d49628

    
           (print_exec ln -s "${conda_prefix}/lib/libtbb.so.12" "${conda_prefix}/lib/libtbb.so") || return 1

One more thing is that installing tbb may not be sufficient.
for example, in CentOS we dnf install -y tbb-devel tbb

ionuthristodorescu

I will also send some diffs on Slack.

ionuthristodorescu · 2025-10-22T03:24:51Z


                    // Compute shared memory size for cta_per_row
                    constexpr auto kCacheAccBytes = sizeof(at::acc_type<cache_t, true>);
-                    int32_t num_cta_per_row_groups = kMaxThreads / kWarpSize;


Is this line and the one below common between CUDA and ROCm ? If yes, we should the {% if rocm %} guards around them.

This is common.

{% if rocm %} guards are already applied. Do we need additional changes on here?

Let me take a look and review in the final version so that we're on the same page.

ionuthristodorescu · 2025-10-22T03:26:10Z

                    // Compute shared memory size for cta_per_row
                    constexpr auto kCacheAccBytes = sizeof(at::acc_type<cache_t, true>);
-                    int32_t num_cta_per_row_groups = kMaxThreads / kWarpSize;
+                    int32_t total_L = indices.numel();


See above comment - total_L seems used only in the ROCm path, so move it under ifdef USE_ROCM ?

Moved int32_t total_L = indices.numel(); doen under {% if is_rocm %}

I couldn't find the fixes, not sure if it's another PR I missed or it hasn't been pushed. Would you mind sharing the latest version with all the fixes?

The fix regarding this is already on line 1065 on this PR. This display is outdated.

ionuthristodorescu · 2025-10-22T03:28:09Z


                const bool use_deterministic_algorithms = at::globalContext().deterministicAlgorithms();
-                const int max_segment_length_per_cta = use_deterministic_algorithms ? INT_MAX : 1024;
+                const int max_segment_length_per_cta = use_deterministic_algorithms ? INT_MAX : 4096;


This seems to affect the CUDA regular path as well - please use a {% if rocm %} guard to select between 1024 and 4096.

Replaced {% if is_rocm %} onto here

ionuthristodorescu · 2025-10-22T03:29:44Z

                    constexpr auto kCacheAccBytes = sizeof(at::acc_type<cache_t, true>);
-                    int32_t num_cta_per_row_groups = kMaxThreads / kWarpSize;
+                    int32_t total_L = indices.numel();
+                    #ifdef USE_ROCM


USE_ROCM is a ROCm-specific, could we guard it with a {% if rocm %} so it does not bleed into CUDA codegen ?

Replaced {% if is_rocm %} onto here

ionuthristodorescu · 2025-10-22T03:30:10Z

                    FBGEMM_LAUNCH_KERNEL(
                        backward_cta_per_row_kernel,
                        cta_per_row_grid_size,
+                        // (64, 2)


Do we need this comment ?

ionuthristodorescu · 2025-10-22T03:33:16Z

    TORCH_CHECK_EQ(grad_outputs.size(), 1);

-    constexpr int32_t max_segment_length_per_warp = 32;
+    constexpr int32_t max_segment_length_per_warp = 16384;


This path seems common with regular CUDA, could we guard it with a {% if rocm %} guard to select between 32 and 16384?

The max_segment_length_per_warp passed by host will be modified on embedding_split_host_pt2_autograd_template.cpp later. Reverting max_segment_length_per_warp on here back to 32.

ionuthristodorescu · 2025-10-22T03:33:37Z

    const auto permute_output_dim_0_1 =
        ctx->saved_data["permute_output_dim_0_1"].toBool();

-    constexpr int32_t max_segment_length_per_warp = 32;


This path seems common with regular CUDA, could we guard it with a {% if rocm %} guard to select between 32 and 16384?

The max_segment_length_per_warp passed by host will be modified on embedding_split_host_pt2_autograd_template.cpp later. Reverting max_segment_length_per_warp on here back to 32.

ionuthristodorescu · 2025-10-22T03:35:11Z

    TORCH_CHECK(aux_tensor[IDX_LXU_CACHE_LOCATIONS].has_value(), "lxu_cache_locations should have value.");
    const auto lxu_cache_locations = aux_tensor[IDX_LXU_CACHE_LOCATIONS].value();
    const auto is_experimental = aux_bool[IDX_IS_EXPERIMENTAL_TBE];
+    const auto mixed_D = aux_bool[IDX_MIXED_D];


This path seems common with regular CUDA, could we guard it with a {% if rocm %} guard to select between 32 and 16384?

Applied {% if is_rocm %} onto here

I don't think rocm needs to be guarded here. mixed_D has been passed to all. It was not used in the forward function, so it's just being saved for backward.

Please use static_cast<bool>(aux_bool[IDX_MIXED_D]);

then you can replace this line to ctx->saved_data["mixed_D"] = mixed_D

Applied the suggested code onto latest PR

ionuthristodorescu · 2025-10-22T03:37:19Z

+    // Workaround. Should not be upstreamed in any way.
+    // Redistribute all cta_per_row work to warp_per_row.
+    int32_t total_L = indices.numel();
+    {%- if (not nobag) and 


Could we add a {% if rocm %} guard around this code, total_L is used only in the USE_ROCM path and USE_ROCM is ROCm specific (so if we do not guard it with {% if rocm %} it will be codegen'ed for regular CUDA paths as well).

By adding a jinja, do you mean only applying {% if rocm %} around the total_L, i.e. @ionuthristodorescu

{% if rocm %} int32_t total_L = indices.numel(); {%- endif %}?

Or changing from using USE_ROCM to {% if rocm %} entirely?

The source code will generate but it shouldn't be compiled. Are we trying to double check here or does it cause any issues, because CUDA path should not compile this. Besides, I think rocm is not passed in this file as a global variable, so the condition will always be false, and total_L will never show up in the source code.

I agree with you. Either USE_ROCM or jinja {% if is_rocm %} is fine for us. Please let us know which one we stick with.

On this file, we should stick with USE_ROCM. I think your original change is correct. rocm is not defined in this file, so jinja will see it as false, and int32_t total_L = indices.numel(); may not show up in the generated source code. Let me review in the final version.

Updated and pushed to this PR.

ionuthristodorescu · 2025-10-22T03:41:20Z

            )
            {%- endif %}

-            for (auto j = 0; j < kWarpSize && l_start + j < L; ++j) {


Here we should use a {% if rocm %} guard to select between rolled / unrolled versions of the loop on regular, non-ROCm paths.

Applied {% if is_rocm %} onto here

spcyppt · 2025-10-26T04:35:44Z

    TORCH_CHECK(aux_tensor[IDX_LXU_CACHE_LOCATIONS].has_value(), "lxu_cache_locations should have value.");
    const auto lxu_cache_locations = aux_tensor[IDX_LXU_CACHE_LOCATIONS].value();
    const auto is_experimental = aux_bool[IDX_IS_EXPERIMENTAL_TBE];
+    const auto mixed_D = aux_bool[IDX_MIXED_D];


I don't think rocm needs to be guarded here. mixed_D has been passed to all. It was not used in the forward function, so it's just being saved for backward.

Please use static_cast<bool>(aux_bool[IDX_MIXED_D]);

then you can replace this line to ctx->saved_data["mixed_D"] = mixed_D

spcyppt · 2025-10-26T04:43:40Z

+    // Workaround. Should not be upstreamed in any way.
+    // Redistribute all cta_per_row work to warp_per_row.
+    int32_t total_L = indices.numel();
+    {%- if (not nobag) and 


The source code will generate but it shouldn't be compiled. Are we trying to double check here or does it cause any issues, because CUDA path should not compile this. Besides, I think rocm is not passed in this file as a global variable, so the condition will always be false, and total_L will never show up in the source code.

spcyppt · 2025-10-26T04:47:34Z

                self.pooling_mode != PoolingMode.NONE
            ), "Mixed dimension tables only supported for pooling tables."
-
+        self.mixed_D = mixed_D


On this change, I assume mixed_D needs to be accessible as the module parameter? Would it cause any issues if it's not self.mixed_D? Asking this to check if we need to split the PRs for backend (C++ source code and backend codegen) and frontend changes (split_table_batched_embeddings_ops_training.py).

Optimization on warp_per_row and cta_per_row kernels will not be activated if self.mixed_D is not present.

mixed_D is set to be true by default through the pipeline. In this case, we should change mixed_D to be False and let the frontend pass the correct mixed_D.

Done in d2596c7

spcyppt · 2025-10-26T05:28:42Z


                    // Compute shared memory size for cta_per_row
                    constexpr auto kCacheAccBytes = sizeof(at::acc_type<cache_t, true>);
-                    int32_t num_cta_per_row_groups = kMaxThreads / kWarpSize;


This is common.

spcyppt · 2025-10-26T05:28:48Z

                    // Compute shared memory size for cta_per_row
                    constexpr auto kCacheAccBytes = sizeof(at::acc_type<cache_t, true>);
-                    int32_t num_cta_per_row_groups = kMaxThreads / kWarpSize;
+                    int32_t total_L = indices.numel();


I couldn't find the fixes, not sure if it's another PR I missed or it hasn't been pushed. Would you mind sharing the latest version with all the fixes?

spcyppt · 2025-11-03T21:17:10Z

            "_nobag" if nobag else "",
        )
        for nobag in [
-            True,


no nobag for gen_embedding_backward_split_{}{}_device_kernel_hip.hip?

Not now. The nobag support is planned in near future

spcyppt · 2025-11-03T21:27:44Z

+            }
+            {%- endif %}
+            for (; j < kWarpSize && l_start + j < L; ++j) {
+            {%- else %} // if is_rocm


nit: change to {#-/* if is_rocm*/#}, otherwise it would be misleading in the generated code

Done in 585300d

spcyppt · 2025-11-03T21:32:16Z

+                    }
+                }
+            }
+            {%- endif %}


nit: add {#-/* if not ssd and not dense and not use_vec_blocking and not vbe */#}

Done in 585300d

spcyppt · 2025-11-03T21:36:09Z

+                }
+            }
+            {%- endif %}
+            for (; j < kWarpSize && l_start + j < L; ++j) {


line 337 and 339, does it really make a difference?

Otherwise, line 337-340 could just be {%- endif %}{#-/* if is_rocm */#}, since line 339 (i.e., for (auto j = 0; j < kWarpSize && l_start + j < L; ++j) {` will be common for both rocm and cuda.

Done in bf143c7. Now it's unified, so CUDA uses it as a main loop while ROCm uses this loop to handle the tailing iterations

spcyppt · 2025-11-04T20:50:05Z

 {%- endif %}

-{%- if is_rocm and not is_index_select and optimizer == "rowwise_adagrad" and not dense and not is_gwd_kernel and not vbe and not ssd %}
+{%- if is_optimized_hip_kernel_supported_mode %}


No nobag support?

Not now. The nobag support is planned in near future

spcyppt · 2025-11-04T20:52:24Z

+                        }
+                    {%- else %}
+                        int32_t num_cta_per_row_groups = kMaxThreads / kWarpSize;
+                        int32_t work_group_size = kMaxThreads;


make work_group_size const for cuda

const int32_t work_group_size = kMaxThreads;

Done in e076556

spcyppt · 2025-11-04T21:01:59Z

                self.pooling_mode != PoolingMode.NONE
            ), "Mixed dimension tables only supported for pooling tables."
-
+        self.mixed_D = mixed_D


mixed_D is set to be true by default through the pipeline. In this case, we should change mixed_D to be False and let the frontend pass the correct mixed_D.

q10 · 2025-11-04T23:24:08Z

@liligwu we're seeing
OSError: libtbb.so.12: cannot open shared object file: No such file or directory
We already install tbb here, so it might just be an issue of updating the build scripts to put libtbb in the LD_LIBRARY_PATH
Hi @q10 , sorry I missed your message. I actually have this commit 4d2bfdd that links tbb explicitly, it works in your container. Do you have any suggestions for fixing this issue in CI, please?
BTW, we discovered a numerical issue in 986cceb and reverted it in 85417b4. It unblocks merging the bwd optimization first.
Thank you.
I think this commit only addresses the build step, where we need to link to tbb. However, for runtime, you might need to do a find in $CONDA_PREFIX from inside the container, and manually update LD_LIBRARY_PATH, or create a symlink (something like

FBGEMM/.github/scripts/utils_build.bash

Line 383 in 0d49628

(print_exec ln -s "${conda_prefix}/lib/librhash.so" "${conda_prefix}/lib/librhash.so.0") || return 1

).
Hi @q10 , I can see a few line below of your example that links tbb

FBGEMM/.github/scripts/utils_build.bash

Line 388 in 0d49628

(print_exec ln -s "${conda_prefix}/lib/libtbb.so.12" "${conda_prefix}/lib/libtbb.so") || return 1

One more thing is that installing tbb may not be sufficient. for example, in CentOS we dnf install -y tbb-devel tbb

In regards to the TBB issue, I think this might be fixed with PR 5087. Could you rebase on latest main after this PR 5087 lands?

warp per row wg change

…nerator.cpp

avbokovoy · 2025-11-06T10:04:06Z

@liligwu we're seeing
OSError: libtbb.so.12: cannot open shared object file: No such file or directory
We already install tbb here, so it might just be an issue of updating the build scripts to put libtbb in the LD_LIBRARY_PATH
Hi @q10 , sorry I missed your message. I actually have this commit 4d2bfdd that links tbb explicitly, it works in your container. Do you have any suggestions for fixing this issue in CI, please?
BTW, we discovered a numerical issue in 986cceb and reverted it in 85417b4. It unblocks merging the bwd optimization first.
Thank you.
I think this commit only addresses the build step, where we need to link to tbb. However, for runtime, you might need to do a find in $CONDA_PREFIX from inside the container, and manually update LD_LIBRARY_PATH, or create a symlink (something like

FBGEMM/.github/scripts/utils_build.bash

Line 383 in 0d49628

(print_exec ln -s "${conda_prefix}/lib/librhash.so" "${conda_prefix}/lib/librhash.so.0") || return 1

).
Hi @q10 , I can see a few line below of your example that links tbb

FBGEMM/.github/scripts/utils_build.bash

Line 388 in 0d49628

(print_exec ln -s "${conda_prefix}/lib/libtbb.so.12" "${conda_prefix}/lib/libtbb.so") || return 1

One more thing is that installing tbb may not be sufficient. for example, in CentOS we dnf install -y tbb-devel tbb
In regards to the TBB issue, I think this might be fixed with PR 5087. Could you rebase on latest main after this PR 5087 lands?

Rebase is done

spcyppt · 2025-11-10T04:39:05Z

@liligwu we're seeing
OSError: libtbb.so.12: cannot open shared object file: No such file or directory
We already install tbb here, so it might just be an issue of updating the build scripts to put libtbb in the LD_LIBRARY_PATH
Hi @q10 , sorry I missed your message. I actually have this commit 4d2bfdd that links tbb explicitly, it works in your container. Do you have any suggestions for fixing this issue in CI, please?
BTW, we discovered a numerical issue in 986cceb and reverted it in 85417b4. It unblocks merging the bwd optimization first.
Thank you.
I think this commit only addresses the build step, where we need to link to tbb. However, for runtime, you might need to do a find in $CONDA_PREFIX from inside the container, and manually update LD_LIBRARY_PATH, or create a symlink (something like

FBGEMM/.github/scripts/utils_build.bash

Line 383 in 0d49628

(print_exec ln -s "${conda_prefix}/lib/librhash.so" "${conda_prefix}/lib/librhash.so.0") || return 1

).
Hi @q10 , I can see a few line below of your example that links tbb

FBGEMM/.github/scripts/utils_build.bash

Line 388 in 0d49628

(print_exec ln -s "${conda_prefix}/lib/libtbb.so.12" "${conda_prefix}/lib/libtbb.so") || return 1

One more thing is that installing tbb may not be sufficient. for example, in CentOS we dnf install -y tbb-devel tbb
In regards to the TBB issue, I think this might be fixed with PR 5087. Could you rebase on latest main after this PR 5087 lands?
Rebase is done

@liligwu @avbokovoy @q10 We still see the tbb errors on CPU CI? https://github.com/pytorch/FBGEMM/actions/runs/19134142334/job/54730556036

q10 · 2025-11-10T11:02:31Z

@liligwu we're seeing
OSError: libtbb.so.12: cannot open shared object file: No such file or directory
We already install tbb here, so it might just be an issue of updating the build scripts to put libtbb in the LD_LIBRARY_PATH
Hi @q10 , sorry I missed your message. I actually have this commit 4d2bfdd that links tbb explicitly, it works in your container. Do you have any suggestions for fixing this issue in CI, please?
BTW, we discovered a numerical issue in 986cceb and reverted it in 85417b4. It unblocks merging the bwd optimization first.
Thank you.
I think this commit only addresses the build step, where we need to link to tbb. However, for runtime, you might need to do a find in $CONDA_PREFIX from inside the container, and manually update LD_LIBRARY_PATH, or create a symlink (something like

FBGEMM/.github/scripts/utils_build.bash

Line 383 in 0d49628

(print_exec ln -s "${conda_prefix}/lib/librhash.so" "${conda_prefix}/lib/librhash.so.0") || return 1

).
Hi @q10 , I can see a few line below of your example that links tbb

FBGEMM/.github/scripts/utils_build.bash

Line 388 in 0d49628

(print_exec ln -s "${conda_prefix}/lib/libtbb.so.12" "${conda_prefix}/lib/libtbb.so") || return 1

One more thing is that installing tbb may not be sufficient. for example, in CentOS we dnf install -y tbb-devel tbb
In regards to the TBB issue, I think this might be fixed with PR 5087. Could you rebase on latest main after this PR 5087 lands?
Rebase is done
@liligwu @avbokovoy @q10 We still see the tbb errors on CPU CI? https://github.com/pytorch/FBGEMM/actions/runs/19134142334/job/54730556036

We had missed one case, and this should be addressed with PR 5106. This PR should be rebased on latest main once PR 5106 lands

spcyppt · 2025-11-10T05:08:46Z

    const auto lxu_cache_locations = aux_tensor[IDX_LXU_CACHE_LOCATIONS].value();
    const auto is_experimental = aux_bool[IDX_IS_EXPERIMENTAL_TBE];
+    const auto mixed_D = static_cast<bool>(aux_bool[IDX_MIXED_D]);
    {%- endif %}


actually, this is not needed in forward since you use mixed_D in backward. I already fixed this in the internal diff.

spcyppt · 2025-11-10T05:10:00Z

    {%- if not nobag %}
    ctx->saved_data["max_D"] = max_D;
-    ctx->saved_data["mixed_D"] = static_cast<bool>(aux_bool[IDX_MIXED_D]);
+    ctx->saved_data["mixed_D"] = mixed_D;


so this should be ctx->saved_data["mixed_D"] = static_cast<bool>(aux_bool[IDX_MIXED_D]);

I already fixed this in the internal diff.

spcyppt · 2025-11-10T05:12:24Z

    constexpr int32_t BT_block_size = 64;
-    constexpr int32_t max_segment_length_per_warp = 64;
+    int32_t max_segment_length_per_warp = 64;
+    int32_t total_L = indices.numel();


this should be moved under line 1072, since is only used in line 1076 which only exists under {%- if (not nobag) and (optimizer == "rowwise_adagrad") and (not vbe) and (not is_gwd) and (not ssd) and (not is_index_select) and (not dense) %}

I already fixed it internal diff.

spcyppt · 2025-11-11T00:51:06Z

 #include <cassert>
-
+#ifdef USE_ROCM
+#include "fbgemm_gpu/rocm/split_embeddings_common.h"


@liligwu @avbokovoy
This causes problem with internal downstream builds which the code already has const auto BLOCK_SIZE = 64; defined. Adding split_embeddings_common.h which also #define BLOCK_SIZE, caused compilation failure.

Resolved in c6b0a88

Summary: bwd performance optimization for ROCm. Fix numerical issues Reviewed By: leitian Differential Revision: D83116315 Pulled By: spcyppt

Summary: Pull Request resolved: pytorch#5121 bwd performance optimization for ROCm. Fix numerical issues Pull Request resolved: pytorch#4925 Reviewed By: leitian Differential Revision: D83116315 Pulled By: spcyppt

q10 · 2025-11-21T23:52:48Z

Hmm, it looks like we are still seeing tbb linking issues - https://github.com/pytorch/FBGEMM/actions/runs/19376974534/job/56094185131?pr=4925

Im guessing that we need both the tbb installation and the explicit tbb linking. @liligwu sorry for the back and forth, can I ask you to add back the tbb linking in the cmake, and rebase to latest commit on main?

meta-codesync · 2025-12-04T23:05:53Z

@q10 merged this pull request in 7cbaff6.

meta-cla Bot added the cla signed label Sep 24, 2025

facebook-github-bot added the module: rocm label Sep 24, 2025

liligwu changed the title ~~forward performance tuning for MI350~~ backward performance optimization for MI350 Oct 20, 2025

ionuthristodorescu reviewed Oct 22, 2025

View reviewed changes

liligwu force-pushed the mi350_dev branch from f8ee546 to 5d73b9c Compare October 27, 2025 19:08

spcyppt reviewed Oct 28, 2025

View reviewed changes

Bernard-Liu mentioned this pull request Oct 31, 2025

Several kernel optimization from aiter team #5074

Open

spcyppt reviewed Nov 4, 2025

View reviewed changes

avbokovoy and others added 14 commits November 6, 2025 09:15

Add gfx950 build support + fp16 fix + index type fix

523a317

Change int64_t to index_t as template parameters in load_raw_per_warp

aee3078

Implement llvm fp16 buffer load for gfx950

5a1ac2e

Fix c-style half to float cast

7856903

Patch 256 half stores

e1e246a

cta_per_row workgroup optim

6a99fe0

Added mi350 guards

349a7b5

Fix index overflow in row load

1178cd1

cta_per_row workgroup reduce by 4 optim

606ad34

Fix mixed_D frontend to backend connection

a22ddeb

changed max_segment_length_per_cta to 4096

6775452

added rocm guards and removed comment

90e6ba7

clean debug statements in Hip.cmake

a9073ac

Merge pull request #121

9a16e12

warp per row wg change

avbokovoy and others added 14 commits November 6, 2025 09:15

Guard f16 llvm intrinsics with ROCm >=7.0

b61bd19

Fix buffer offset for emb_dim == 160

c38ff6f

Remove sanity check

e201e8b

address the potential lint issues and revert the change in indices_ge…

aaaf80c

…nerator.cpp

addresss code style issue

b8aea67

removed guard rocm on mixed_D and refactored mixed_D var assignment

b9a7759

Remove general load/store methods

a4b4431

Move weight type check to compile-time

5d4f2cd

Switch to 256B stores for float type

d3b7d7a

removed jinj is_rocm on total_L as USE_ROCM is already applied

878d00f

Change mixed_D default value to false

d2596c7

Make const work_group_size for CUDA

e076556

Add jinja comments to grad_indice_weights kernel

585300d

Remove redundand comment

e0db2f1

avbokovoy force-pushed the mi350_dev branch from 990df49 to e0db2f1 Compare November 6, 2025 10:02

Unify cuda and rocm loops

bf143c7

spcyppt reviewed Nov 11, 2025

View reviewed changes

Added BLOCK_SIZE_ROCM

c6b0a88

revert the link to tbb

122a583

meta-codesync Bot closed this in 7cbaff6 Dec 4, 2025

facebook-github-bot added the Merged label Dec 4, 2025

q10 added feature:tbe contributor:amd labels Jan 7, 2026

Uh oh!

Conversation

liligwu commented Sep 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

netlify Bot commented Sep 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

✅ Deploy Preview for pytorch-fbgemm-docs ready!

Uh oh!

facebook-github-bot commented Sep 24, 2025

Uh oh!

q10 commented Oct 2, 2025

Uh oh!

meta-codesync Bot commented Oct 13, 2025

Uh oh!

liligwu commented Oct 13, 2025

Uh oh!

q10 commented Oct 14, 2025

Uh oh!

liligwu commented Oct 15, 2025

Uh oh!

ionuthristodorescu left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kudomcho Oct 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

liligwu commented Sep 24, 2025 •

edited

Loading

netlify Bot commented Sep 24, 2025 •

edited

Loading

kudomcho Oct 22, 2025 •

edited

Loading