Skip to content

Fix UVM FP16 performance regression on MI300 #145

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
avbokovoy wants to merge 2 commits into
base: abokovoi/upstream
Choose a base branch
from
+9 −19
Open

Fix UVM FP16 performance regression on MI300 #145

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6f74f37
Revert force FP32 cache type for ROCm
avbokovoy Mar 5, 2026
24190b3
Vectorize half stores
avbokovoy Mar 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 3 additions & 16 deletions fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -772,22 +772,9 @@ def __init__( # noqa C901

self.weights_precision = weights_precision

if torch.cuda.is_available() and torch.version.hip:
# NOTE: It was discovered that FP16 cache precision caused a 500x
# slowdown in performance of split_embedding_nobag_backward_codegen_rowwise_adagrad_unweighted_kernel_warp_per_row_1
# kernel on ROCm, so to work around this, we fix cache precision to
# be FP32 always for the ROCm environment case.
#
# See:
# https://fb.workplace.com/groups/fbgemmusers/permalink/9438488366231860/
cache_precision = SparseType.FP32
self.log("Override cache_precision=SparseType.FP32 on ROCm")
else:
# NOTE: The changes from D65865527 are retained here until we can
# test that the the hack also works for non-ROCm environments.
cache_precision = (
weights_precision if cache_precision is None else cache_precision
)
cache_precision = (
weights_precision if cache_precision is None else cache_precision
)

self.output_dtype: int = output_dtype.as_int()
assert (
Expand Down
9 changes: 6 additions & 3 deletions fbgemm_gpu/include/fbgemm_gpu/utils/float.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,9 +51,12 @@ struct Half4 {

__device__ inline void store(at::Half* p) {
#ifdef USE_ROCM
*reinterpret_cast<unsigned int*>(p) = *reinterpret_cast<unsigned int*>(&a);
*reinterpret_cast<unsigned int*>(p + 2) =
*reinterpret_cast<unsigned int*>(&b);
const unsigned int lo = *reinterpret_cast<unsigned int*>(&a);
const unsigned int hi = *reinterpret_cast<unsigned int*>(&b);
const unsigned long long packed =
static_cast<unsigned long long>(lo) |
(static_cast<unsigned long long>(hi) << 32);
*reinterpret_cast<unsigned long long*>(p) = packed;
#else

#ifndef __HALF2_TO_UI
Expand Down