Skip to content

Add cuDNN variable-length attention path for padded key/value streams #728

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
puririshi98 wants to merge 1 commit into from
Open

Add cuDNN variable-length attention path for padded key/value streams #728

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
2 changes: 1 addition & 1 deletion .github/workflows/test-gpu.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,7 @@ jobs:
--find-links=https://data.pyg.org/whl/torch-2.14.0+cu130.html \
--index-strategy=unsafe-best-match \
--group=test \
.
'.[cudnn]'

- name: GPU sanity check
run: |
Expand Down
2 changes: 2 additions & 0 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,8 @@ pip install git+https://github.com/NVIDIA/structured-data-models.git
> [!NOTE]
> For CUDA workloads, we highly recommend installing [`cudf`](https://docs.rapids.ai/install) as an additional dependency to keep dataframe-style operations on GPU and avoid unnecessary data movement.

The optional `cudnn` extra (`pip install "structured-data-models[cudnn]"`) adds cuDNN variable-length attention for padded inputs, enabled via `sdm.nn.enable_cudnn_varlen()` (Linux only; on other platforms the extra installs nothing and the boolean-mask path is used).

## Model Families

**Tabular Foundation Models:**
Expand Down
2 changes: 2 additions & 0 deletions docs/source/install.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,3 +10,5 @@ pip install git+https://github.com/NVIDIA/structured-data-models.git
```{note}
For CUDA workloads, we highly recommend installing [`cudf`](https://docs.rapids.ai/install) as an additional dependency to keep dataframe-style operations on GPU and avoid unnecessary data movement.
```

The optional `cudnn` extra (`pip install "structured-data-models[cudnn]"`) adds cuDNN variable-length attention for padded inputs, enabled via `sdm.nn.enable_cudnn_varlen()` (Linux only; on other platforms the extra installs nothing and the boolean-mask path is used).
5 changes: 5 additions & 0 deletions pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,11 @@ dependencies = [
"typing_extensions",
]

[project.optional-dependencies]
cudnn = [
"nvidia-cudnn-frontend>=1.23,<1.29; sys_platform == 'linux'",
]

[project.urls]
Download = "https://github.com/NVIDIA/structured-data-models/releases"
Homepage = "https://github.com/NVIDIA/structured-data-models"
Expand Down
3 changes: 3 additions & 0 deletions sdm/nn/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@

"""Neural network modules for structured data models."""

from sdm.nn._cudnn_varlen import cudnn_varlen_stats, enable_cudnn_varlen
from sdm.nn.rope import RotaryEmbedding
from sdm.nn.glu import SwiGLU
from sdm.nn.softplus import SoftplusScale
Expand All @@ -13,6 +14,8 @@


__all__ = [
"cudnn_varlen_stats",
"enable_cudnn_varlen",
"RotaryEmbedding",
"SwiGLU",
"SoftplusScale",
Expand Down
Loading
Loading