Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
29 changes: 29 additions & 0 deletions docs/source/builder/writing-kernels.md
Original file line number Diff line number Diff line change
Expand Up @@ -545,6 +545,35 @@ $ nix run .#ciTests.torch210-cxx11-cpu-x86_64-linux
When running the tests on a non-NixOS systems, make sure that
[the CUDA driver library can be found](https://danieldk.eu/Software/Nix/Nix-CUDA-on-non-NixOS-systems#solutions).

### Tests that need access to internal APIs

If the tests need to access a module, not exposed in the public API
surface, do not use `importlib`. For example:

```py
...

importlib.import_module(f"{flash_attn4.__name__}.compute_block_sparsity")
BlockSparsityKernel = flash_attn4.compute_block_sparsity.BlockSparsityKernel
BlockSparseTensors = flash_attn4.block_sparsity.BlockSparseTensors
```

This pattern should be avoided. Instead, create a separate private
module named `_private_for_testing` and include the modules needed in the tests inside it. Then
expose that private module from the public API surface. So, the
`__init__.py` would look like so:

```py
from . import _private_for_tesing

__all__ = [
"_private_for_testing",
...
]
```

Where `_private_for_testing` would contain the modules needed for tests.

## Kernel docs

We provide a utility to generate a system card for a given kernel, utilizing
Expand Down
9 changes: 7 additions & 2 deletions docs/source/kernel-requirements.md
Original file line number Diff line number Diff line change
Expand Up @@ -298,8 +298,8 @@ fail on the build variants that were not rebuilt.

### Exceptions

A version bump is not needed when every existing build variant is
replaced by the new build:
For the following types of kernels, some API changes do not require
bumping the version:

- Python-only (noarch) kernels, e.g. `torch-cuda`. All build variants get
updated, so this issue does not exist.
Expand All @@ -310,6 +310,11 @@ replaced by the new build:
- Torch stable ABI kernels, as long as the CUDA versions that get built
overlap with the current build variants.

In the following cases no version bump is required for these types of kernels:

- Adding new functions/methods/classes.
- Adding new arguments with default values to a function/method/constructor, as long as the default values do not change previous behavior.

## Native Python module

Kernels will typically contain a native Python module with precompiled
Expand Down
Loading