fix(tf/pt): validate GPU tabulation sizes

[njzjz-bot]

Summary

• validate 1 <= last_layer_size <= 1024 before every TensorFlow CUDA/ROCm tabulation wrapper call
• move the four existing grad-gradient checks ahead of their kernel launches
• cover all SeA, SeAtten, SeT, and SeR forward/gradient/grad-gradient GPU paths, including zero width

Why existing tests missed this

There were no direct TensorFlow tabulation-op GPU tests. Existing PyTorch tabulation coverage uses normal small embedding widths and does not execute these TensorFlow kernels, while real compressed models are also normally well below the device limit. The four existing grad-gradient checks looked like validation but ran only after the GPU wrapper had already launched, and forward/first-gradient paths had no checks at all.

Validation

• cmake --build source/build --target deepmd_op -j2
• constructed all 12 direct-op regression graphs against the fresh build-tree TensorFlow extension
• pytest source/tests/tf/test_tabulate_gpu_size_validation.py source/tests/tf/test_tabulate.py -q (2 passed, 2 GPU-only tests skipped locally because no usable GPU is exposed)
• ruff format .
• ruff check .
• clang-format dry run with --Werror
• git diff --check

Fixes #5654

Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh

Summary by CodeRabbit

• Bug Fixes

  • GPU tabulation operations now reject invalid last_layer_size values outside the supported range (must be between 1 and 1024).
  • Validation is standardized and performed before GPU execution for affected forward/gradient paths.

• Tests

  • Added GPU-focused unit tests to confirm invalid last_layer_size (zero and oversized) is rejected across supported tabulation operation variants, including gradient paths.

[coderabbitai]

📝 Walkthrough

Walkthrough

TensorFlow and PyTorch GPU tabulation wrappers now reject last_layer_size values outside 1–1024 before launching CUDA/ROCm kernels. GPU-gated tests cover oversized and zero values across forward, gradient, and grad-gradient paths.

Changes

GPU tabulation validation

Layer / File(s)	Summary
TensorFlow GPU validation and coverage source/op/tf/tabulate_multi_device.cc, source/tests/tf/test_tabulate_gpu_size_validation.py	Adds shared prelaunch validation to TensorFlow GPU wrappers and tests standardized errors for invalid sizes.
PyTorch GPU validation and coverage source/op/pt/tabulate_multi_device.cc, source/tests/pt/test_tabulate_gpu_size_validation.py	Adds shared prelaunch validation to PyTorch GPU wrappers and tests forward and autograd paths for invalid sizes.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Suggested labels: CUDA, ROCM

Suggested reviewers: njzjz

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Out of Scope Changes check	⚠️ Warning	The PyTorch tabulation validation and tests are not part of linked issue #5654, which only requests TensorFlow GPU validation.	Move the PyTorch changes into a separate PR or link a PyTorch issue, and keep this PR focused on the TensorFlow GPU validation.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	The TensorFlow changes validate last_layer_size before GPU launches across forward, grad, and grad-grad paths and add matching tests.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the main change: adding GPU tabulation size validation across TensorFlow and PyTorch.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[codecov]

Codecov Report

❌ Patch coverage is 0% with 28 lines in your changes missing coverage. Please review.
✅ Project coverage is 78.94%. Comparing base (cc908a8) to head (8183850).
⚠️ Report is 10 commits behind head on master.

Files with missing lines	Patch %	Lines
source/op/tf/tabulate_multi_device.cc	0.00%	15 Missing ⚠️
source/op/pt/tabulate_multi_device.cc	0.00%	13 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5846      +/-   ##
==========================================
- Coverage   79.03%   78.94%   -0.09%     
==========================================
  Files        1055     1069      +14     
  Lines      122233   124094    +1861     
  Branches     4401     4523     +122     
==========================================
+ Hits        96607    97969    +1362     
- Misses      24061    24510     +449     
- Partials     1565     1615      +50     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[wanghan-iapcm]

The TF-side fix itself is correct: the validation is GPU-scoped (called inside if (device == "GPU"), so CPU runs with large last_layer_size are unaffected even on a CUDA build), the 1..1024 bound matches every kernel's blockDim.x, and moving it before the launch (and extending it to the forward/first-grad ops) is exactly right — it resolves a defect CodeRabbit flagged on #5600 that was never applied. The reject tests will run in the merge-queue CUDA job.

Blocking on one thing: the identical bug is still live in the PyTorch op, source/op/pt/tabulate_multi_device.cc. There the grad-grad TORCH_CHECK(last_layer_size <= 1024, ...) runs after the kernel launch (

deepmd-kit/source/op/pt/tabulate_multi_device.cc

Lines 174 to 182 in e5fdff0

	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,
	dz_dy_dtwo, nloc, nnei, last_layer_size, is_sorted);
	#else
	throw std::runtime_error(
	"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,

, and likewise at L328/L472/L596), and the forward / first-gradient PT ops (e.g. the tabulate_fusion_se_a_gpu launch at

deepmd-kit/source/op/pt/tabulate_multi_device.cc

Lines 58 to 71 in e5fdff0

	const int64_t nnei = em_tensor.size(1);
	// compute
	if (device == "GPU") {
	#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);
	#else
	throw std::runtime_error(
	"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
	} else if (device == "CPU") {
	deepmd::tabulate_fusion_se_a_cpu(descriptor, table, table_info, em_x, em,
	two_embed, nloc, nnei, last_layer_size);

) have no check at all — the same silent SIGFPE / bad-launch risk #5654 was filed for. Please either apply the same before-launch validation to the PT op in this PR, or open a tracking issue so the vulnerability class is actually closed rather than half-fixed. Details inline.

[coderabbitai]

Actionable comments posted: 2

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@source/tests/pt/test_tabulate_gpu_size_validation.py`:
- Around line 66-72: Extend test_rejects_zero_width to also validate a negative
width, such as -1, through _assert_rejected, covering the implementation’s
complete last_layer_size <= 0 boundary without changing the existing zero-width
assertion.
- Around line 74-85: The test test_gradient_paths_reject_oversized_width
currently asserts only the forward failure, so it does not exercise gradient
paths. Restructure the tests to use valid forward inputs and explicitly cover
the required tabulate_fusion_se_* _backward and _grad_grad wrappers for PyTorch,
including SeT, SeTTebd, and SeR, while asserting oversized-width rejection in
each relevant gradient operation.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro Plus

Run ID: df27c21c-ccdd-4184-986e-55c3cb1212c7

📥 Commits

Reviewing files that changed from the base of the PR and between 6ce049d and 03d6501.

📒 Files selected for processing (2)

• source/op/pt/tabulate_multi_device.cc
• source/tests/pt/test_tabulate_gpu_size_validation.py

[wanghan-iapcm]

Mirroring the validation rather than deferring it to a tracking issue is the better of the two options I offered, and the PT side is more thorough than what I asked for. Verified at 81838508b:

• ValidateGpuLastLayerSize rejects <= 0 as well as > 1024, and the comment explains why the bound exists (block dimension / dynamic shared memory sizing) rather than just restating it.
• Twelve call sites, which matches your count exactly. I checked a forward and three grad-grad wrappers (lines 72, 186, 332, 601): every one sits inside if (device == "GPU") and immediately precedes the launch, so CPU runs with a large last_layer_size keep working. That scoping was the main thing I wanted to confirm.
• All four post-launch TORCH_CHECK(last_layer_size <= 1024) calls are gone. Removing them was right: sitting after the launch they could never prevent the fault they described, and their presence is a good part of why this survived review on #5600.

On why no test caught this: real configurations use last-layer widths in the tens, so nothing in the suite approaches 1024, and the defect lives entirely inside the device == "GPU" branch that CPU CI never enters. The new test crosses exactly that intersection -- 1025 and 0 across all five forward wrappers plus the descriptor-derived autograd path -- and would fail before the fix, since the forwards had no check and reach the launch instead of raising.

One thing to flag, and it is a correction to my own earlier comment rather than a request. I said the reject tests would run in the merge-queue CUDA job. On this PR Test Python on CUDA is skipping, so neither the new PT tests nor the TF ones from the first commit have actually executed -- Pass testing on CUDA goes green in 8s because it is the aggregate gate over skipped jobs. The gating itself is unavoidable, since a GPU launch guard cannot be tested without a GPU. If the merge queue does not run these either, it would be worth applying the Test CUDA label once before merging so the guards are observed working at least a single time.

Approving; this closes #5654 on both backends.

[njzjz]

I'll ask the agent to fix the following errors:

=========================== short test summary info ============================
SUBFAILED(op='se_a', last_layer_size=-1) source/tests/pt/test_tabulate_gpu_size_validation.py::TestTabulateGpuSizeValidation::test_rejects_nonpositive_width - AssertionError: "last_layer_size must be between 1 and 1024" does not match "Trying to create tensor with negative dimension -1: [1, 4, -1]"
SUBFAILED(op='se_atten', last_layer_size=-1) source/tests/pt/test_tabulate_gpu_size_validation.py::TestTabulateGpuSizeValidation::test_rejects_nonpositive_width - AssertionError: "last_layer_size must be between 1 and 1024" does not match "Trying to create tensor with negative dimension -1: [1, 4, -1]"
SUBFAILED(op='se_t', last_layer_size=-1) source/tests/pt/test_tabulate_gpu_size_validation.py::TestTabulateGpuSizeValidation::test_rejects_nonpositive_width - AssertionError: "last_layer_size must be between 1 and 1024" does not match "Trying to create tensor with negative dimension -1: [1, -1]"
SUBFAILED(op='se_t_tebd', last_layer_size=-1) source/tests/pt/test_tabulate_gpu_size_validation.py::TestTabulateGpuSizeValidation::test_rejects_nonpositive_width - AssertionError: "last_layer_size must be between 1 and 1024" does not match "Trying to create tensor with negative dimension -1: [1, 1, 1, -1]"
SUBFAILED(op='se_r', last_layer_size=-1) source/tests/pt/test_tabulate_gpu_size_validation.py::TestTabulateGpuSizeValidation::test_rejects_nonpositive_width - AssertionError: "last_layer_size must be between 1 and 1024" does not match "Trying to create tensor with negative dimension -1: [1, 1, -1]"
SUBFAILED(op='se_a_forward', last_layer_size=1025) source/tests/tf/test_tabulate_gpu_size_validation.py::TestTabulateGpuSizeValidation::test_rejects_oversized_width_for_all_gpu_paths - AssertionError: InvalidArgumentError not raised
SUBFAILED(op='se_a_grad_grad', last_layer_size=1025) source/tests/tf/test_tabulate_gpu_size_validation.py::TestTabulateGpuSizeValidation::test_rejects_oversized_width_for_all_gpu_paths - AssertionError: InvalidArgumentError not raised
SUBFAILED(op='se_atten_forward', last_layer_size=1025) source/tests/tf/test_tabulate_gpu_size_validation.py::TestTabulateGpuSizeValidation::test_rejects_oversized_width_for_all_gpu_paths - AssertionError: InvalidArgumentError not raised
SUBFAILED(op='se_atten_grad_grad', last_layer_size=1025) source/tests/tf/test_tabulate_gpu_size_validation.py::TestTabulateGpuSizeValidation::test_rejects_oversized_width_for_all_gpu_paths - AssertionError: InvalidArgumentError not raised
SUBFAILED(op='se_t_forward', last_layer_size=1025) source/tests/tf/test_tabulate_gpu_size_validation.py::TestTabulateGpuSizeValidation::test_rejects_oversized_width_for_all_gpu_paths - AssertionError: InvalidArgumentError not raised
SUBFAILED(op='se_t_grad_grad', last_layer_size=1025) source/tests/tf/test_tabulate_gpu_size_validation.py::TestTabulateGpuSizeValidation::test_rejects_oversized_width_for_all_gpu_paths - AssertionError: InvalidArgumentError not raised
SUBFAILED(op='se_r_forward', last_layer_size=1025) source/tests/tf/test_tabulate_gpu_size_validation.py::TestTabulateGpuSizeValidation::test_rejects_oversized_width_for_all_gpu_paths - AssertionError: InvalidArgumentError not raised
SUBFAILED(op='se_r_grad', last_layer_size=1025) source/tests/tf/test_tabulate_gpu_size_validation.py::TestTabulateGpuSizeValidation::test_rejects_oversized_width_for_all_gpu_paths - AssertionError: InvalidArgumentError not raised
SUBFAILED(op='se_r_grad_grad', last_layer_size=1025) source/tests/tf/test_tabulate_gpu_size_validation.py::TestTabulateGpuSizeValidation::test_rejects_oversized_width_for_all_gpu_paths - AssertionError: InvalidArgumentError not raised
SUBFAILED(op='se_a_forward', last_layer_size=0) source/tests/tf/test_tabulate_gpu_size_validation.py::TestTabulateGpuSizeValidation::test_rejects_zero_width - AssertionError: InvalidArgumentError not raised