Skip to content

fix(tf/pt): validate GPU tabulation sizes #5846

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
njzjz-bot wants to merge 5 commits into
base: master
Choose a base branch
from
+444 −36
Open

fix(tf/pt): validate GPU tabulation sizes #5846

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
6ce049d
fix(tf): validate GPU tabulation sizes
Jul 16, 2026
17b44cb
Merge remote-tracking branch 'origin/master' into fix/tf-tabulate-gpu…
Jul 27, 2026
03d6501
fix(pt): validate GPU tabulation sizes before launch
Jul 27, 2026
8183850
test(pt): cover tabulation gradient size guards
Jul 30, 2026
b88dd44
fix(tabulate): validate GPU widths before allocation
njzjz-bot Jul 31, 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
49 changes: 37 additions & 12 deletions source/op/pt/tabulate_multi_device.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,26 @@ void GetTensorDevice(const torch::Tensor& t, std::string& str) {
}
}

static void ValidateGpuLastLayerSize(const int64_t last_layer_size) {
// The GPU tabulation kernels use this dimension either as the block size or
// to size dynamic shared memory, so reject invalid models before any launch:
// an out-of-range value overflows the block dimension (or divides by zero)
// inside the launch configuration.
TORCH_CHECK(last_layer_size > 0 && last_layer_size <= 1024,
"last_layer_size must be between 1 and 1024 for GPU tabulation, "
"but got ",
last_layer_size);
}

static void ValidateGpuLastLayerSizeBeforeAllocation(
const torch::Tensor& table_tensor, const int64_t last_layer_size) {
// Forward wrappers allocate descriptor tensors using last_layer_size, so a
// negative GPU width must be rejected before torch::empty sees the shape.
if (table_tensor.device().is_cuda()) {
ValidateGpuLastLayerSize(last_layer_size);
}
}

template <typename FPTYPE>
void TabulateFusionSeAForward(const torch::Tensor& table_tensor,
const torch::Tensor& table_info_tensor,
Expand Down Expand Up @@ -58,6 +78,7 @@ void TabulateFusionSeAForward(const torch::Tensor& table_tensor,
const int64_t nnei = em_tensor.size(1);
// compute
if (device == "GPU") {
ValidateGpuLastLayerSize(last_layer_size);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
deepmd::tabulate_fusion_se_a_gpu(descriptor, table, table_info, em_x, em,
two_embed, nloc, nnei, last_layer_size);
Expand Down Expand Up @@ -110,6 +131,7 @@ void TabulateFusionSeAGradForward(const torch::Tensor& table_tensor,
const int64_t last_layer_size = descriptor_tensor.size(2);
// compute
if (device == "GPU") {
ValidateGpuLastLayerSize(last_layer_size);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
deepmd::tabulate_fusion_se_a_grad_gpu(dy_dem_x, dy_dem, dy_dtwo, table,
table_info, em_x, em, two_embed, dy,
Expand Down Expand Up @@ -170,6 +192,7 @@ void TabulateFusionSeAGradGradForward(const torch::Tensor& table_tensor,
const int64_t last_layer_size = descriptor_tensor.size(2);
// compute
if (device == "GPU") {
ValidateGpuLastLayerSize(last_layer_size);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
deepmd::tabulate_fusion_se_a_grad_grad_gpu(
dz_dy, table, table_info, em_x, em, two_embed, dz_dy_dem_x, dz_dy_dem,
Expand All @@ -179,9 +202,6 @@ void TabulateFusionSeAGradGradForward(const torch::Tensor& table_tensor,
"The input tensor is on the GPU, but the GPU support for the "
"customized OP library is not enabled.");
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
TORCH_CHECK(last_layer_size <= 1024,
"In the process of model compression, the size of the "
"last layer of embedding net must be less than 1024!");
} else if (device == "CPU") {
deepmd::tabulate_fusion_se_a_grad_grad_cpu(
dz_dy, table, table_info, em_x, em, two_embed, dz_dy_dem_x, dz_dy_dem,
Expand Down Expand Up @@ -221,6 +241,7 @@ void TabulateFusionSeTForward(const torch::Tensor& table_tensor,
const int64_t nnei_j = em_tensor.size(2);
// compute
if (device == "GPU") {
ValidateGpuLastLayerSize(last_layer_size);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
deepmd::tabulate_fusion_se_t_gpu(descriptor, table, table_info, em_x, em,
nloc, nnei_i, nnei_j, last_layer_size);
Expand Down Expand Up @@ -266,6 +287,7 @@ void TabulateFusionSeTGradForward(const torch::Tensor& table_tensor,
const int64_t last_layer_size = descriptor_tensor.size(1);
// compute
if (device == "GPU") {
ValidateGpuLastLayerSize(last_layer_size);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
deepmd::tabulate_fusion_se_t_grad_gpu(dy_dem_x, dy_dem, table, table_info,
em_x, em, dy, nloc, nnei_i, nnei_j,
Expand Down Expand Up @@ -316,6 +338,7 @@ void TabulateFusionSeTGradGradForward(const torch::Tensor& table_tensor,
const int64_t last_layer_size = descriptor_tensor.size(1);
// compute
if (device == "GPU") {
ValidateGpuLastLayerSize(last_layer_size);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
deepmd::tabulate_fusion_se_t_grad_grad_gpu(dz_dy, table, table_info, em_x,
em, dz_dy_dem_x, dz_dy_dem, nloc,
Expand All @@ -325,9 +348,6 @@ void TabulateFusionSeTGradGradForward(const torch::Tensor& table_tensor,
"The input tensor is on the GPU, but the GPU support for the "
"customized OP library is not enabled.");
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
TORCH_CHECK(last_layer_size <= 1024,
"In the process of model compression, the size of the "
"last layer of embedding net must be less than 1024!");
} else if (device == "CPU") {
deepmd::tabulate_fusion_se_t_grad_grad_cpu(dz_dy, table, table_info, em_x,
em, dz_dy_dem_x, dz_dy_dem, nloc,
Expand Down Expand Up @@ -368,6 +388,7 @@ void TabulateFusionSeTTebdForward(const torch::Tensor& table_tensor,
const int64_t nnei_j = em_tensor.size(2);
// compute
if (device == "GPU") {
ValidateGpuLastLayerSize(last_layer_size);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
deepmd::tabulate_fusion_se_t_tebd_gpu(descriptor, table, table_info, em_x,
em, nloc, nnei_i, nnei_j,
Expand Down Expand Up @@ -414,6 +435,7 @@ void TabulateFusionSeTTebdGradForward(const torch::Tensor& table_tensor,

// compute
if (device == "GPU") {
ValidateGpuLastLayerSize(last_layer_size);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
deepmd::tabulate_fusion_se_t_tebd_grad_gpu(dy_dem_x, table, table_info,
em_x, em, dy, nloc, nnei_i,
Expand Down Expand Up @@ -460,6 +482,7 @@ void TabulateFusionSeTTebdGradGradForward(
const int64_t last_layer_size = descriptor_tensor.size(3);
// compute
if (device == "GPU") {
ValidateGpuLastLayerSize(last_layer_size);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
deepmd::tabulate_fusion_se_t_tebd_grad_grad_gpu(
dz_dy, table, table_info, em_x, em, dz_dy_dem_x, nloc, nnei_i, nnei_j,
Expand All @@ -469,9 +492,6 @@ void TabulateFusionSeTTebdGradGradForward(
"The input tensor is on the GPU, but the GPU support for the "
"customized OP library is not enabled.");
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
TORCH_CHECK(last_layer_size <= 1024,
"In the process of model compression, the size of the "
"last layer of embedding net must be less than 1024!");
} else if (device == "CPU") {
deepmd::tabulate_fusion_se_t_tebd_grad_grad_cpu(
dz_dy, table, table_info, em_x, em, dz_dy_dem_x, nloc, nnei_i, nnei_j,
Expand Down Expand Up @@ -505,6 +525,7 @@ void TabulateFusionSeRForward(const torch::Tensor& table_tensor,
const int64_t nnei = em_tensor.size(1);
// compute
if (device == "GPU") {
ValidateGpuLastLayerSize(last_layer_size);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
deepmd::tabulate_fusion_se_r_gpu(descriptor, table, table_info, em, nloc,
nnei, last_layer_size);
Expand Down Expand Up @@ -545,6 +566,7 @@ void TabulateFusionSeRGradForward(const torch::Tensor& table_tensor,
const int64_t last_layer_size = descriptor_tensor.size(2);
// compute
if (device == "GPU") {
ValidateGpuLastLayerSize(last_layer_size);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
deepmd::tabulate_fusion_se_r_grad_gpu(dy_dem, table, table_info, em, dy,
nloc, nnei, last_layer_size);
Expand Down Expand Up @@ -585,6 +607,7 @@ void TabulateFusionSeRGradGradForward(const torch::Tensor& table_tensor,
const int64_t last_layer_size = descriptor_tensor.size(2);
// compute
if (device == "GPU") {
ValidateGpuLastLayerSize(last_layer_size);
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
deepmd::tabulate_fusion_se_r_grad_grad_gpu(
dz_dy, table, table_info, em, dz_dy_dem, nloc, nnei, last_layer_size);
Expand All @@ -593,9 +616,6 @@ void TabulateFusionSeRGradGradForward(const torch::Tensor& table_tensor,
"The input tensor is on the GPU, but the GPU support for the "
"customized OP library is not enabled.");
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
TORCH_CHECK(last_layer_size <= 1024,
"In the process of model compression, the size of the "
"last layer of embedding net must be less than 1024!");
} else if (device == "CPU") {
deepmd::tabulate_fusion_se_r_grad_grad_cpu(
dz_dy, table, table_info, em, dz_dy_dem, nloc, nnei, last_layer_size);
Expand Down Expand Up @@ -776,6 +796,7 @@ class TabulateFusionSeAOp
const torch::Tensor& em_x_tensor,
const torch::Tensor& em_tensor,
int64_t last_layer_size) {
ValidateGpuLastLayerSizeBeforeAllocation(table_tensor, last_layer_size);
// allocate output tensors
auto options = torch::TensorOptions()
.dtype(table_tensor.dtype())
Expand Down Expand Up @@ -962,6 +983,7 @@ class TabulateFusionSeAttenOp
const torch::Tensor& two_embed_tensor,
int64_t last_layer_size,
bool is_sorted) {
ValidateGpuLastLayerSizeBeforeAllocation(table_tensor, last_layer_size);
// allocate output tensors
auto options = torch::TensorOptions()
.dtype(table_tensor.dtype())
Expand Down Expand Up @@ -1130,6 +1152,7 @@ class TabulateFusionSeTOp
const torch::Tensor& em_x_tensor,
const torch::Tensor& em_tensor,
int64_t last_layer_size) {
ValidateGpuLastLayerSizeBeforeAllocation(table_tensor, last_layer_size);
// allocate output tensors
auto options = torch::TensorOptions()
.dtype(table_tensor.dtype())
Expand Down Expand Up @@ -1283,6 +1306,7 @@ class TabulateFusionSeROp
const torch::Tensor& table_info_tensor,
const torch::Tensor& em_tensor,
int64_t last_layer_size) {
ValidateGpuLastLayerSizeBeforeAllocation(table_tensor, last_layer_size);
// allocate output tensors
auto options = torch::TensorOptions()
.dtype(table_tensor.dtype())
Expand Down Expand Up @@ -1441,6 +1465,7 @@ class TabulateFusionSeTTebdOp
const torch::Tensor& em_x_tensor,
const torch::Tensor& em_tensor,
int64_t last_layer_size) {
ValidateGpuLastLayerSizeBeforeAllocation(table_tensor, last_layer_size);
// allocate output tensors
auto options = torch::TensorOptions()
.dtype(table_tensor.dtype())
Expand Down
Loading