deepmodeling / deepmodeling/deepmd-kit

[Code scan] SE-Attention ops can read and write past mismatched tensor shapes

Open
#5,895 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 focused memory-safety review at commit `21b73ab303581f467bba0cae101a880c35cc2d82`.

## Problem

The TensorFlow and PyTorch SE-Attention wrappers validate tensor rank but do not validate the dimensions or element counts that the native implementation indexes.

For PyTorch:

- [forward checks only that `em_x` and `two_embed` are rank 2 and `em` is rank 3](https://github.com/deepmodeling/deepmd-kit/blob/21b73ab303581f467bba0cae101a880c35cc2d82/source/op/pt/tabulate_multi_device.cc#L29-L40)
- [`nloc` and `nnei` are taken exclusively from `em`](https://github.com/deepmodeling/deepmd-kit/blob/21b73ab303581f467bba0cae101a880c35cc2d82/source/op/pt/tabulate_multi_device.cc#L48-L58)
- [backward allocates `dy_dtwo` with `zeros_like(two_embed)`](https://github.com/deepmodeling/deepmd-kit/blob/21b73ab303581f467bba0cae101a880c35cc2d82/source/op/pt/tabulate_multi_device.cc#L857-L874)

The native code nevertheless indexes or clears exactly `nloc * nnei * last_layer_size` `two_embed`/ `dy_dtwo` elements:

- [CPU reads `two_embed[ii * nnei * last_layer_size + jj * last_layer_size + kk]`](https://github.com/deepmodeling/deepmd-kit/blob/21b73ab303581f467bba0cae101a880c35cc2d82/source/lib/src/tabulate.cc#L298-L300)
- [CPU backward clears that full computed size](https://github.com/deepmodeling/deepmd-kit/blob/21b73ab303581f467bba0cae101a880c35cc2d82/source/lib/src/tabulate.cc#L249-L254)
- [GPU backward writes the same computed range](https://github.com/deepmodeling/deepmd-kit/blob/21b73ab303581f467bba0cae101a880c35cc2d82/source/lib/src/gpu/tabulate.cu#L475-L480)

TensorFlow has the same mismatch: [it checks only ranks](https://github.com/deepmodeling/deepmd-kit/blob/21b73ab303581f467bba0cae101a880c35cc2d82/source/op/tf/tabulate_multi_device.cc#L375-L390), while [the backward output is allocated using the unvalidated `two_embed` shape](https://github.com/deepmodeling/deepmd-kit/blob/21b73ab303581f467bba0cae101a880c35cc2d82/source/op/tf/tabulate_multi_device.cc#L442-L468).

The same class of problem applies to `em_x`, the final dimension of `em` (assumed to be 4), `dy`, `descriptor`, `table_info`, and the table size.

## Reproduction

The probes were run inside an `srun --partition=main --gres=gpu:5090:1` allocation.

A PyTorch tensor with logical shape `[1, 1]` was passed as `two_embed` for `em.shape == [1, 3, 4]`. Its backing storage contained three values `[0, 10, 20]`. Although only the first value belongs to the tensor view, forward consumed all three:

```text
logical_two_shape (1, 1) logical_numel 1
backing_storage [0.0, 10.0, 20.0]
output [6.6000000000000005, 33.0, 0.0, 0.0]
```

A C++ canary probe then emulated the one-element backward output followed by two adjacent canary values. The canaries started as `12345` and `67890`; after `tabulate_fusion_se_a_grad_cpu` they had both been overwritten:

```text
logical_output=1.2 canary0=1.2 canary1=1.2
```

Thus the missing validation causes both reads beyond a tensor's logical boundary and writes beyond the allocated output shape. On GPU the corresponding write can become an illegal-address error; on CPU it can silently corrupt another allocation or eventually segfault.

## Impact

Normal descriptor code currently constructs matching shapes, so this is primarily reachable through direct/raw custom-op calls or an upstream shape bug. However, custom ops must reject malformed inputs rather than corrupt process memory.

## Suggested fix

Before taking raw pointers or launching the native kernels, validate at least:

- `em.shape == [nloc, nnei, 4]`;
- `em_x.shape == [nloc, nnei]`;
- `two_embed.numel() == nloc * nnei * last_layer_size`;
- forward/backward descriptor and cotangent shapes;
- `table_info.numel() >= 5`;
- the table contains every interval and feature coefficient that `table_info` and `last_layer_size` can address;
- required device placement is consistent.

Apply equivalent checks to the related SE-A, SE-T, and SE-R raw ops.

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

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.