deepmodeling / deepmodeling/deepmd-kit
[Code scan] Validate TensorFlow tabulation GPU sizes before launching kernels
- Dominant language
- Python
- Stars
- 2k
- Forks
- 649
- Avg merge
- 6d 18h
- Merged PRs (30d)
- 15
Description
This issue comes from a Codex global scan of `deepmodeling/deepmd-kit` at commit `73de44b1f94471b2e3bdb6b11f57b34d7bc791bb`.
## Problem
Several TensorFlow tabulation GPU paths validate `last_layer_size <= 1024` only after launching the CUDA/ROCm kernel.
Examples in the grad-grad paths:
- https://github.com/deepmodeling/deepmd-kit/blob/73de44b1f94471b2e3bdb6b11f57b34d7bc791bb/source/op/tf/tabulate_multi_device.cc#L338-L347
- https://github.com/deepmodeling/deepmd-kit/blob/73de44b1f94471b2e3bdb6b11f57b34d7bc791bb/source/op/tf/tabulate_multi_device.cc#L555-L564
- https://github.com/deepmodeling/deepmd-kit/blob/73de44b1f94471b2e3bdb6b11f57b34d7bc791bb/source/op/tf/tabulate_multi_device.cc#L746-L755
- https://github.com/deepmodeling/deepmd-kit/blob/73de44b1f94471b2e3bdb6b11f57b34d7bc791bb/source/op/tf/tabulate_multi_device.cc#L911-L919
The GPU wrappers launch kernels with `last_layer_size` as the block size or shared-memory dimension, for example:
- https://github.com/deepmodeling/deepmd-kit/blob/73de44b1f94471b2e3bdb6b11f57b34d7bc791bb/source/lib/src/gpu/tabulate.cu#L1023-L1048
- https://github.com/deepmodeling/deepmd-kit/blob/73de44b1f94471b2e3bdb6b11f57b34d7bc791bb/source/lib/src/gpu/tabulate.cu#L1085-L1109
So oversized inputs can hit a CUDA runtime failure before TensorFlow returns the intended `InvalidArgument`.
## Impact
Invalid compressed-model tabulation shapes can leave users with low-level GPU launch errors instead of a deterministic TensorFlow validation error. Forward tabulation paths also launch GPU kernels without the same prelaunch bound check.
## Suggested fix
Validate `last_layer_size > 0 && last_layer_size <= 1024` before every GPU tabulation wrapper call that uses it as a kernel launch dimension or shared-memory multiplier.
Add GPU tests that pass `last_layer_size=1025` and assert a clean TensorFlow `InvalidArgument` without launching the kernel.
Contributor guide
Assessment
This issue has not been assessed yet.