deepmodeling / deepmodeling/deepmd-kit

[Code scan] Validate TensorFlow tabulation GPU sizes before launching kernels

Open
#5,654 0 comments 0 reactions 0 assignees View on GitHub
bug
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.