deformops 3.0.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (133) hide show
  1. deformops-3.0.0/LICENSE +21 -0
  2. deformops-3.0.0/MANIFEST.in +4 -0
  3. deformops-3.0.0/PKG-INFO +625 -0
  4. deformops-3.0.0/README.md +572 -0
  5. deformops-3.0.0/VERSION +1 -0
  6. deformops-3.0.0/pyproject.toml +448 -0
  7. deformops-3.0.0/setup.cfg +4 -0
  8. deformops-3.0.0/setup.py +331 -0
  9. deformops-3.0.0/sources/deformops/__init__.py +44 -0
  10. deformops-3.0.0/sources/deformops/__init__.pyi +6 -0
  11. deformops-3.0.0/sources/deformops/backends/__init__.py +123 -0
  12. deformops-3.0.0/sources/deformops/backends/cuda/__init__.py +34 -0
  13. deformops-3.0.0/sources/deformops/backends/cuda/_msda_bridge.py +492 -0
  14. deformops-3.0.0/sources/deformops/backends/cuda/_schema.py +71 -0
  15. deformops-3.0.0/sources/deformops/backends/cuda/attn.py +34 -0
  16. deformops-3.0.0/sources/deformops/backends/cuda/conv.py +34 -0
  17. deformops-3.0.0/sources/deformops/backends/cuda/experimental/__init__.py +8 -0
  18. deformops-3.0.0/sources/deformops/backends/cuda/experimental/_query_order.py +216 -0
  19. deformops-3.0.0/sources/deformops/backends/cuda/experimental/msda_half2atomic.cpp +2 -0
  20. deformops-3.0.0/sources/deformops/backends/cuda/experimental/msda_half2atomic.cu +25 -0
  21. deformops-3.0.0/sources/deformops/backends/cuda/experimental/msda_half2atomic.py +19 -0
  22. deformops-3.0.0/sources/deformops/backends/cuda/experimental/msda_half2atomic_serialreduce.cpp +2 -0
  23. deformops-3.0.0/sources/deformops/backends/cuda/experimental/msda_half2atomic_serialreduce.cu +25 -0
  24. deformops-3.0.0/sources/deformops/backends/cuda/experimental/msda_half2atomic_serialreduce.py +19 -0
  25. deformops-3.0.0/sources/deformops/backends/cuda/experimental/msda_hilbert.py +18 -0
  26. deformops-3.0.0/sources/deformops/backends/cuda/experimental/msda_morton.py +18 -0
  27. deformops-3.0.0/sources/deformops/backends/cuda/experimental/msda_scalaratomic.cpp +2 -0
  28. deformops-3.0.0/sources/deformops/backends/cuda/experimental/msda_scalaratomic.cu +25 -0
  29. deformops-3.0.0/sources/deformops/backends/cuda/experimental/msda_scalaratomic.py +19 -0
  30. deformops-3.0.0/sources/deformops/backends/cuda/experimental/msda_scalaratomic_serialreduce.cpp +2 -0
  31. deformops-3.0.0/sources/deformops/backends/cuda/experimental/msda_scalaratomic_serialreduce.cu +25 -0
  32. deformops-3.0.0/sources/deformops/backends/cuda/experimental/msda_scalaratomic_serialreduce.py +19 -0
  33. deformops-3.0.0/sources/deformops/backends/cuda/experimental/msda_serialreduce.cpp +2 -0
  34. deformops-3.0.0/sources/deformops/backends/cuda/experimental/msda_serialreduce.cu +25 -0
  35. deformops-3.0.0/sources/deformops/backends/cuda/experimental/msda_serialreduce.py +19 -0
  36. deformops-3.0.0/sources/deformops/backends/cuda/include/interpolate.cuh +220 -0
  37. deformops-3.0.0/sources/deformops/backends/cuda/include/msda_args.cuh +87 -0
  38. deformops-3.0.0/sources/deformops/backends/cuda/include/msda_backward.cuh +339 -0
  39. deformops-3.0.0/sources/deformops/backends/cuda/include/msda_boxed.cuh +56 -0
  40. deformops-3.0.0/sources/deformops/backends/cuda/include/msda_forward.cuh +206 -0
  41. deformops-3.0.0/sources/deformops/backends/cuda/include/msda_launch.cuh +279 -0
  42. deformops-3.0.0/sources/deformops/backends/cuda/include/msda_shared_memory.cuh +55 -0
  43. deformops-3.0.0/sources/deformops/backends/cuda/include/stable_compat.cuh +101 -0
  44. deformops-3.0.0/sources/deformops/backends/cuda/include/utils.cuh +246 -0
  45. deformops-3.0.0/sources/deformops/backends/cuda/msda.py +51 -0
  46. deformops-3.0.0/sources/deformops/backends/cuda/msda_fp32acc.cpp +2 -0
  47. deformops-3.0.0/sources/deformops/backends/cuda/msda_fp32acc.cu +25 -0
  48. deformops-3.0.0/sources/deformops/backends/cuda/msda_fp32acc.py +19 -0
  49. deformops-3.0.0/sources/deformops/backends/eager/__init__.py +19 -0
  50. deformops-3.0.0/sources/deformops/backends/eager/attn.py +38 -0
  51. deformops-3.0.0/sources/deformops/backends/eager/attn_pytorch.py +250 -0
  52. deformops-3.0.0/sources/deformops/backends/eager/conv.py +56 -0
  53. deformops-3.0.0/sources/deformops/backends/eager/conv_pytorch.py +188 -0
  54. deformops-3.0.0/sources/deformops/backends/eager/msda.py +36 -0
  55. deformops-3.0.0/sources/deformops/backends/eager/msda_pytorch.py +150 -0
  56. deformops-3.0.0/sources/deformops/backends/reference/__init__.py +17 -0
  57. deformops-3.0.0/sources/deformops/backends/reference/attn.py +174 -0
  58. deformops-3.0.0/sources/deformops/backends/reference/conv.py +139 -0
  59. deformops-3.0.0/sources/deformops/backends/reference/msda.py +97 -0
  60. deformops-3.0.0/sources/deformops/backends/triton/__init__.py +26 -0
  61. deformops-3.0.0/sources/deformops/backends/triton/_attn_bridge.py +225 -0
  62. deformops-3.0.0/sources/deformops/backends/triton/_kernels/__init__.py +1 -0
  63. deformops-3.0.0/sources/deformops/backends/triton/_kernels/attn_query_tiled.py +670 -0
  64. deformops-3.0.0/sources/deformops/backends/triton/_kernels/conv_deform2d.py +839 -0
  65. deformops-3.0.0/sources/deformops/backends/triton/_kernels/msda_query_tiled.py +534 -0
  66. deformops-3.0.0/sources/deformops/backends/triton/_launch.py +93 -0
  67. deformops-3.0.0/sources/deformops/backends/triton/_msda_bridge.py +614 -0
  68. deformops-3.0.0/sources/deformops/backends/triton/_vmap.py +7 -0
  69. deformops-3.0.0/sources/deformops/backends/triton/attn.py +75 -0
  70. deformops-3.0.0/sources/deformops/backends/triton/attn_query_tiled.py +21 -0
  71. deformops-3.0.0/sources/deformops/backends/triton/attn_query_tiled_fp32acc.py +22 -0
  72. deformops-3.0.0/sources/deformops/backends/triton/conv.py +79 -0
  73. deformops-3.0.0/sources/deformops/backends/triton/conv_deform2d.py +229 -0
  74. deformops-3.0.0/sources/deformops/backends/triton/experimental/__init__.py +6 -0
  75. deformops-3.0.0/sources/deformops/backends/triton/experimental/_kernels/__init__.py +1 -0
  76. deformops-3.0.0/sources/deformops/backends/triton/experimental/_kernels/msda_point_parallel.py +510 -0
  77. deformops-3.0.0/sources/deformops/backends/triton/experimental/_kernels/msda_query_tiled_autotuned.py +601 -0
  78. deformops-3.0.0/sources/deformops/backends/triton/experimental/_kernels/msda_query_tiled_cached.py +592 -0
  79. deformops-3.0.0/sources/deformops/backends/triton/experimental/_kernels/msda_reduce.py +792 -0
  80. deformops-3.0.0/sources/deformops/backends/triton/experimental/msda_atomic_wide.py +27 -0
  81. deformops-3.0.0/sources/deformops/backends/triton/experimental/msda_block_reduce.py +27 -0
  82. deformops-3.0.0/sources/deformops/backends/triton/experimental/msda_block_reduce_wide.py +27 -0
  83. deformops-3.0.0/sources/deformops/backends/triton/experimental/msda_point_parallel.py +25 -0
  84. deformops-3.0.0/sources/deformops/backends/triton/experimental/msda_query_tiled_autotuned.py +29 -0
  85. deformops-3.0.0/sources/deformops/backends/triton/experimental/msda_query_tiled_autotuned_fp32acc.py +28 -0
  86. deformops-3.0.0/sources/deformops/backends/triton/experimental/msda_query_tiled_cached.py +27 -0
  87. deformops-3.0.0/sources/deformops/backends/triton/experimental/msda_query_tiled_cached_fp32acc.py +25 -0
  88. deformops-3.0.0/sources/deformops/backends/triton/experimental/msda_sort_segscan.py +27 -0
  89. deformops-3.0.0/sources/deformops/backends/triton/experimental/msda_sort_segscan_wide.py +27 -0
  90. deformops-3.0.0/sources/deformops/backends/triton/experimental/msda_warp_reduce.py +27 -0
  91. deformops-3.0.0/sources/deformops/backends/triton/experimental/msda_warp_reduce_wide.py +27 -0
  92. deformops-3.0.0/sources/deformops/backends/triton/msda.py +72 -0
  93. deformops-3.0.0/sources/deformops/backends/triton/msda_query_tiled.py +27 -0
  94. deformops-3.0.0/sources/deformops/backends/triton/msda_query_tiled_fp32acc.py +27 -0
  95. deformops-3.0.0/sources/deformops/core/__init__.py +5 -0
  96. deformops-3.0.0/sources/deformops/core/build.py +1296 -0
  97. deformops-3.0.0/sources/deformops/core/capabilities.py +83 -0
  98. deformops-3.0.0/sources/deformops/core/configs.py +186 -0
  99. deformops-3.0.0/sources/deformops/core/errors.py +47 -0
  100. deformops-3.0.0/sources/deformops/core/runtime.py +91 -0
  101. deformops-3.0.0/sources/deformops/core/sanitize.py +715 -0
  102. deformops-3.0.0/sources/deformops/core/tuning.py +133 -0
  103. deformops-3.0.0/sources/deformops/integrations/__init__.py +9 -0
  104. deformops-3.0.0/sources/deformops/integrations/_ops.py +93 -0
  105. deformops-3.0.0/sources/deformops/integrations/_patch.py +103 -0
  106. deformops-3.0.0/sources/deformops/integrations/mmcv.py +261 -0
  107. deformops-3.0.0/sources/deformops/integrations/torchvision.py +78 -0
  108. deformops-3.0.0/sources/deformops/integrations/transformers.py +282 -0
  109. deformops-3.0.0/sources/deformops/models/__init__.py +42 -0
  110. deformops-3.0.0/sources/deformops/models/deformable_detr/__init__.py +43 -0
  111. deformops-3.0.0/sources/deformops/models/deformable_detr/model.py +84 -0
  112. deformops-3.0.0/sources/deformops/models/deformable_detr/model_patches.py +48 -0
  113. deformops-3.0.0/sources/deformops/models/mask2former/__init__.py +54 -0
  114. deformops-3.0.0/sources/deformops/models/mask2former/model.py +103 -0
  115. deformops-3.0.0/sources/deformops/models/mask2former/model_patches.py +260 -0
  116. deformops-3.0.0/sources/deformops/nn/__init__.py +23 -0
  117. deformops-3.0.0/sources/deformops/nn/attn.py +464 -0
  118. deformops-3.0.0/sources/deformops/nn/conv.py +183 -0
  119. deformops-3.0.0/sources/deformops/nn/msda.py +276 -0
  120. deformops-3.0.0/sources/deformops/py.typed +1 -0
  121. deformops-3.0.0/sources/deformops/registry/__init__.py +46 -0
  122. deformops-3.0.0/sources/deformops/registry/_dispatch.py +378 -0
  123. deformops-3.0.0/sources/deformops/registry/_public_ops.py +67 -0
  124. deformops-3.0.0/sources/deformops/registry/_vmap.py +82 -0
  125. deformops-3.0.0/sources/deformops/registry/attn.py +453 -0
  126. deformops-3.0.0/sources/deformops/registry/conv.py +355 -0
  127. deformops-3.0.0/sources/deformops/registry/msda.py +452 -0
  128. deformops-3.0.0/sources/deformops/version.py +114 -0
  129. deformops-3.0.0/sources/deformops.egg-info/PKG-INFO +625 -0
  130. deformops-3.0.0/sources/deformops.egg-info/SOURCES.txt +131 -0
  131. deformops-3.0.0/sources/deformops.egg-info/dependency_links.txt +1 -0
  132. deformops-3.0.0/sources/deformops.egg-info/requires.txt +41 -0
  133. deformops-3.0.0/sources/deformops.egg-info/top_level.txt +1 -0
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2023 Mobile Perception Systems (MPS) lab at Eindhoven University of Technology (TU/e)
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,4 @@
1
+ # The base version is read from VERSION by setup.py at build time (the project
2
+ # version is declared dynamic), so it must be present in the sdist too: uv and
3
+ # python-build build the wheel from the unpacked sdist.
4
+ include VERSION
@@ -0,0 +1,625 @@
1
+ Metadata-Version: 2.4
2
+ Name: deformops
3
+ Version: 3.0.0
4
+ Summary: Deformable sampling operations for PyTorch.
5
+ Author-email: Kurt Stolle <kurt@computer.org>
6
+ License-Expression: MIT
7
+ Project-URL: Homepage, https://github.com/tue-p8n/deformops
8
+ Project-URL: Repository, https://github.com/tue-p8n/deformops
9
+ Project-URL: Issues, https://github.com/tue-p8n/deformops/issues
10
+ Project-URL: Changelog, https://github.com/tue-p8n/deformops/blob/main/CHANGELOG.md
11
+ Project-URL: Paper, https://openreview.net/forum?id=Q4jZ7zKNKx
12
+ Keywords: perception,computer vision,deep learning,detection,segmentation,deformable detr,mask transformer,vision transformer,deformable convolution,deformable sampling,multi-scale deformable attention,deformconv,deformattn,msda,msdeformattn,pytorch,cuda,triton
13
+ Classifier: Development Status :: 5 - Production/Stable
14
+ Classifier: Environment :: GPU :: NVIDIA CUDA
15
+ Classifier: Intended Audience :: Science/Research
16
+ Classifier: Operating System :: POSIX :: Linux
17
+ Classifier: Programming Language :: Python
18
+ Classifier: Programming Language :: Python :: 3.12
19
+ Classifier: Programming Language :: Python :: 3.13
20
+ Classifier: Programming Language :: Python :: Implementation :: CPython
21
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
22
+ Classifier: Typing :: Typed
23
+ Requires-Python: <3.14,>=3.12
24
+ Description-Content-Type: text/markdown
25
+ License-File: LICENSE
26
+ Requires-Dist: torch>=2.11.0
27
+ Provides-Extra: cpu
28
+ Requires-Dist: torch>=2.11.0; extra == "cpu"
29
+ Requires-Dist: torchvision>=0.24.0; extra == "cpu"
30
+ Provides-Extra: cu126
31
+ Requires-Dist: torch>=2.11.0; extra == "cu126"
32
+ Provides-Extra: cu128
33
+ Requires-Dist: torch>=2.11.0; extra == "cu128"
34
+ Provides-Extra: cu129
35
+ Requires-Dist: torch>=2.11.0; extra == "cu129"
36
+ Provides-Extra: cu130
37
+ Requires-Dist: torch>=2.12.0; extra == "cu130"
38
+ Provides-Extra: cu132
39
+ Requires-Dist: torch>=2.12.0; extra == "cu132"
40
+ Provides-Extra: triton
41
+ Requires-Dist: triton>=3.0; (sys_platform == "linux" and platform_machine == "x86_64") and extra == "triton"
42
+ Provides-Extra: jit
43
+ Requires-Dist: ninja>=1.11; extra == "jit"
44
+ Provides-Extra: mask2former
45
+ Requires-Dist: torchmatch>=1.0.4; python_version >= "3.13" and extra == "mask2former"
46
+ Requires-Dist: transformers>=4.40.0; extra == "mask2former"
47
+ Provides-Extra: deformable-detr
48
+ Requires-Dist: timm>=0.9; extra == "deformable-detr"
49
+ Requires-Dist: transformers>=4.40.0; extra == "deformable-detr"
50
+ Provides-Extra: mmcv
51
+ Requires-Dist: mmcv-lite>=1.7; extra == "mmcv"
52
+ Dynamic: license-file
53
+
54
+ # DeformOps
55
+
56
+ `deformops` is a PyTorch library of three deformable sampling operators:
57
+ single-scale deformable sampling,
58
+ classic deformable convolution,
59
+ and multi-scale deformable attention (MSDA).
60
+ Its CUDA and Triton kernels are tuned for the encoder-scale workload
61
+ that dominates Mask2Former-style training:
62
+ 5k to 22k queries per call in the pixel decoder.
63
+ [Efficient Multi-Scale Deformable Attention on GPUs](https://openreview.net/forum?id=Q4jZ7zKNKx)
64
+ (TMLR 2026)
65
+ describes the kernel design and its measurements.
66
+
67
+ - **Zero-buffer kernels.**
68
+ The multi-scale attention backward writes into the value gradient
69
+ with no intermediate sampling buffer.
70
+ At encoder scale, the paper measures forward and backward peak memory
71
+ 71% below the reference for the CUDA kernel
72
+ and 88% below for the native-bfloat16 Triton kernel.
73
+ - **Triton kernels** for all three operators,
74
+ with multi-scale attention measured in the paper on Ampere and Hopper GPUs.
75
+ - **Automatic backend selection** by device, GPU architecture and training precision;
76
+ see [Backend selection](#backend-selection).
77
+ - **Half precision** (bfloat16 and float16) throughout;
78
+ bfloat16 is the recommended training dtype.
79
+ - **Parallel-safe kernel builds.**
80
+ Concurrent compiles under `pytest-xdist`, multiple GPUs or batch-job ranks
81
+ serialize without deadlocking.
82
+
83
+ ## Installation
84
+
85
+ ```bash
86
+ pip install deformops
87
+ ```
88
+
89
+ `deformops` supports Python 3.12 and 3.13
90
+ and needs PyTorch 2.11 or newer (2.12 or newer for CUDA 13.x).
91
+ A bare install needs only `torch`
92
+ and gives the eager backends on any platform.
93
+ The GPU kernels need Linux x86-64 and a CUDA-capable GPU.
94
+
95
+ The published package is pure Python.
96
+ With `backend=None`, a GPU call runs a Triton kernel compiled at runtime,
97
+ so it needs no CUDA toolkit.
98
+ The CUDA backend needs a compiled kernel library;
99
+ see [Kernel loading](#kernel-loading).
100
+
101
+ Installing from the source distribution with `--no-build-isolation`
102
+ compiles the non-experimental CUDA kernels into the package,
103
+ against the installed CUDA build of torch that later loads them.
104
+ This needs setuptools, a CUDA toolkit with torch's CUDA major version,
105
+ and a visible GPU or `TORCH_CUDA_ARCH_LIST`:
106
+
107
+ ```bash
108
+ pip install --no-binary deformops --no-build-isolation deformops
109
+ DEFORMOPS_EXTENSIONS=none pip install --no-binary deformops deformops # without compiled kernels
110
+ ```
111
+
112
+ Without a CUDA toolkit, the install skips the compile without a message.
113
+ It skips the kernels and prints why in four other cases:
114
+ the build is isolated,
115
+ the toolkit's CUDA major version differs from torch's,
116
+ no GPU is visible and `TORCH_CUDA_ARCH_LIST` is unset,
117
+ or the compile fails.
118
+ The CUDA backend repeats that reason when it finds no library.
119
+ An explicit `DEFORMOPS_EXTENSIONS` or `FORCE_CUDA=1` fails the install instead.
120
+ pip and uv reuse a wheel built from the same source distribution,
121
+ so after changing these settings install with `--no-cache-dir` (pip) or `--no-cache` (uv).
122
+ A host compiler newer than PyTorch's range for the CUDA toolkit fails torch's version check;
123
+ `TORCH_DONT_CHECK_COMPILER_ABI=1` skips it.
124
+
125
+ | Install | What it adds |
126
+ | :--------------------------- | :---------------------------------------------------------------------------- |
127
+ | `deformops[cpu]`, `deformops[cu126]`, `deformops[cu128]`, `deformops[cu129]`, `deformops[cu130]`, `deformops[cu132]` | In a uv-synced checkout, routes `torch` to the CPU or matching CUDA wheel index (`cu130` and `cu132` need torch 2.12 or newer). |
128
+ | `deformops[triton]` | Installs the Triton compiler for a CPU torch. |
129
+ | `deformops[jit]` | Installs `ninja`, so torch can compile a kernel locally. |
130
+ | `deformops[mask2former]` | Installs `transformers`, and `torchmatch` on Python 3.13, for the pre-made Mask2Former graph. |
131
+ | `deformops[deformable-detr]` | Installs `transformers` and `timm` for the pre-made Deformable DETR graph. |
132
+ | `deformops[mmcv]` | Installs `mmcv-lite` for the mmcv integration. |
133
+
134
+ The `cpu` and `cuXXX` extras are mutually exclusive.
135
+ Their index mappings live in `[tool.uv.sources]`, not in published metadata,
136
+ so under pip, choose the build through the PyTorch index:
137
+
138
+ ```bash
139
+ pip install torch --index-url https://download.pytorch.org/whl/cu128 # CUDA 12.8, or /whl/cpu
140
+ pip install deformops
141
+ uv pip install deformops --torch-backend=cu128 # the same, with uv
142
+ ```
143
+
144
+ A CUDA torch already depends on Triton;
145
+ the Triton backends detect it at runtime either way.
146
+ Compiling needs a CUDA toolkit besides `ninja`, and pip does not provide one.
147
+ If either is missing, a compile that `DEFORMOPS_COMPILE` requests raises `KernelBuildError` naming what is missing.
148
+ Extras do not change operator speed.
149
+
150
+ ### Kernel loading
151
+
152
+ The CUDA backend takes its compiled kernel library from the first source that has one:
153
+ the install itself;
154
+ an earlier compile on this machine for the same sources and environment;
155
+ a prebuilt index (`DEFORMOPS_DOWNLOAD`, off by default);
156
+ or a local compile (`DEFORMOPS_COMPILE`, off by default).
157
+ An enabled download that fails raises `KernelDownloadError`,
158
+ which names the variable that skips it.
159
+ With both disabled, the CUDA backend is unavailable:
160
+ `backend=None` moves on with a warning,
161
+ and `backend=Backend.CUDA` raises `KernelNotProvidedError`.
162
+
163
+ Ahead-of-time libraries (prebuilt or from a source-distribution install)
164
+ have no float64 kernels,
165
+ so a float64 call on the CUDA backend raises
166
+ unless the library is a just-in-time build (`DEFORMOPS_COMPILE=always`).
167
+
168
+ The library loads on the first call,
169
+ or earlier through `deformops.backends.load`.
170
+ Run it before loading an exported program that calls the CUDA operators
171
+ (`torch.ops.deformops_cuda.*`):
172
+
173
+ ```python
174
+ from deformops.backends import load
175
+ from deformops.registry import Backend
176
+
177
+ load("msda", Backend.CUDA) # loads the CUDA kernel library
178
+ ```
179
+
180
+ ## Multi-scale deformable attention
181
+
182
+ The layer is `deformops.nn.MultiScaleDeformAttn2d`
183
+ and the functional form `deformops.registry.msda`.
184
+ Shapes below use `B` batch, `G` heads, `D` channels per head,
185
+ `Q` queries, `L` levels and `K` points per level.
186
+
187
+ ```python
188
+ import torch
189
+ from deformops.nn import MultiScaleDeformAttn2d
190
+
191
+ layer = MultiScaleDeformAttn2d(
192
+ dim=256,
193
+ num_heads=8,
194
+ num_levels=4,
195
+ num_points=4,
196
+ ).cuda()
197
+
198
+ query = torch.randn(1, 300, 256).cuda()
199
+ reference = torch.rand(1, 300, 4, 2).cuda() # normalized [0, 1]
200
+ value = torch.randn(1, 5440, 256).cuda() # flattened feature maps
201
+ shapes = torch.tensor([[64, 64], [32, 32], [16, 16], [8, 8]]).cuda()
202
+ starts = torch.tensor([0, 4096, 5120, 5376]).cuda() # prefix sum of H * W
203
+
204
+ out = layer(query, reference, value, shapes, starts)
205
+ ```
206
+
207
+ The functional form takes the sampled operands without the projections;
208
+ its five tensor operands are positional-only.
209
+
210
+ By default, every backend uses the attention weights as given.
211
+ `method="softmax"` normalizes raw logits jointly over levels and points
212
+ inside the kernel.
213
+ The layer requests softmax itself, so this default does not affect it.
214
+
215
+ ```python
216
+ from deformops.registry import msda
217
+
218
+ value = torch.randn(1, 5440, 8, 32).cuda() # (B, sum H*W, G, D)
219
+ locations = torch.rand(1, 300, 8, 4, 4, 2).cuda() # (B, Q, G, L, K, 2)
220
+ weights = torch.rand(1, 300, 8, 4 * 4).cuda() # (B, Q, G, L * K), raw logits
221
+
222
+ out = msda(value, shapes, starts, locations, weights, method="softmax")
223
+ ```
224
+
225
+ ## Single-scale deformable sampling
226
+
227
+ The layers are `deformops.nn.DeformAttn2d` and `deformops.nn.DeformSample2d`;
228
+ the functional form is `deformops.registry.attn`.
229
+ The operator samples and aggregates without the spatial weight of deformable convolution.
230
+ `DeformAttn2d` takes a separate query and value,
231
+ predicts offsets and aggregation weights from the query,
232
+ and projects the value and the output.
233
+ `DeformSample2d` derives the query from its input with a depthwise convolution
234
+ and runs `DeformAttn2d` on it.
235
+
236
+ The device kernel is Triton,
237
+ so it needs no CUDA toolkit, per-CUDA build or downloaded library.
238
+ Off CUDA and for float64, `backend=None` runs the eager implementation.
239
+
240
+ ```python
241
+ from deformops.nn import DeformSample2d
242
+
243
+ layer = DeformSample2d(dim=64, kernel_size=3, groups=4).cuda()
244
+ x = torch.randn(1, 64, 32, 32).cuda() # NCHW, as nn.Conv2d takes
245
+ y = layer(x) # -> (1, 64, 32, 32)
246
+ ```
247
+
248
+ ## Deformable convolution
249
+
250
+ The layer is `deformops.nn.DeformConv2d`;
251
+ the functional form is `deformops.registry.conv`.
252
+ The operator computes classic deformable convolution as `torchvision.ops.deform_conv2d` does:
253
+ a `(C_out, C_in / groups, kH, kW)` weight contracted against the samples,
254
+ v2 with a mask and v1 without.
255
+
256
+ ```python
257
+ from deformops.nn import DeformConv2d
258
+
259
+ layer = DeformConv2d(64, 64, kernel_size=3, padding=1).cuda()
260
+ offset = torch.randn(1, 2 * 9, 32, 32).cuda() # (dy, dx) per kernel point
261
+ mask = torch.rand(1, 9, 32, 32).cuda() # omit for v1
262
+ y = layer(x, offset, mask) # -> (1, 64, 32, 32)
263
+ ```
264
+
265
+ One kernel gathers and contracts,
266
+ so the im2col column buffer `(B, H_out * W_out, C_in, kH * kW)` never exists.
267
+ Accumulation is float32 for every operand dtype.
268
+ TensorFloat-32 (TF32) is off inside the contraction,
269
+ since its rounding error would reach the offset gradients.
270
+ `conv` runs the same kernel over channels-last operands with `deformops` argument names.
271
+
272
+ ## torchvision drop-in
273
+
274
+ `deformops.integrations.torchvision` matches `torchvision.ops.deform_conv2d`'s
275
+ argument order and NCHW layout.
276
+
277
+ ```python
278
+ from deformops.integrations.torchvision import deform_conv2d
279
+
280
+ y = deform_conv2d(x, offset, layer.weight, layer.bias, padding=layer.padding, mask=mask)
281
+ ```
282
+
283
+ ## Backend selection
284
+
285
+ Every operator dispatches through one of three backends,
286
+ named by `deformops.registry.Backend`:
287
+
288
+ ```python
289
+ from deformops.registry import Backend
290
+
291
+ assert set(Backend) == {Backend.CUDA, Backend.TRITON, Backend.EAGER}
292
+ ```
293
+
294
+ Pass one to a layer, or leave `backend=None` for the selection policy:
295
+
296
+ ```python
297
+ from deformops.nn import MultiScaleDeformAttn2d
298
+
299
+ layer = MultiScaleDeformAttn2d(
300
+ dim=256,
301
+ num_heads=8,
302
+ num_levels=4,
303
+ backend=Backend.CUDA,
304
+ )
305
+ ```
306
+
307
+ | Backend | What it is |
308
+ | :--------------- | :------------------------------------------------------------------------- |
309
+ | `Backend.CUDA` | The compiled CUDA kernel, for multi-scale attention. |
310
+ | `Backend.TRITON` | The Triton kernels, compiled at runtime. |
311
+ | `Backend.EAGER` | The eager PyTorch backend; device-agnostic and traceable. |
312
+
313
+ Only the CUDA backend has these limits:
314
+ a per-head channel count divisible by 8,
315
+ 1 to 5 levels,
316
+ 4 or 8 points per level,
317
+ and at most 1024 threads per block.
318
+ A head layout uses `G * D / 8` threads in the forward
319
+ and, when training, `G * D` in the backward, or `G * D / 2` when `D >= 64`.
320
+
321
+ With `backend=None`, CPU tensors run eager.
322
+ CUDA tensors run a GPU kernel or raise:
323
+ MSDA tries Triton, then CUDA;
324
+ single-scale sampling and deformable convolution use Triton.
325
+ Eager runs on a CUDA tensor only with `backend=Backend.EAGER` or for float64.
326
+ A call that skips its first choice warns with `BackendFallbackWarning`,
327
+ naming the skipped backends and why;
328
+ `DEFORMOPS_FALLBACK_WARNINGS=0` silences it.
329
+ When no backend can run, `NoBackendError` lists every candidate and why it was rejected.
330
+
331
+ GPUs below SM 9.0 (streaming-multiprocessor version, the CUDA compute capability)
332
+ have no native bfloat16 atomic add,
333
+ the instruction that lets parallel GPU threads add into one gradient entry.
334
+ The kernels emulate it with a compare-and-swap retry loop,
335
+ which biases the value gradient under contention,
336
+ so half-precision training there runs Triton with the float32 accumulator.
337
+ MSDA and single-scale sampling share this rule,
338
+ and their other calls use native half-precision atomics.
339
+ The choice is memoized per call context,
340
+ so it follows a model to a new device or dtype.
341
+ An explicit `backend=Backend.TRITON` gets the float32 accumulator,
342
+ which is correct on every architecture;
343
+ `fp32acc=False` uses native atomics on SM 9.0 and later.
344
+
345
+ For MSDA, `select_backend` implements this policy:
346
+
347
+ ```python
348
+ from deformops.registry.msda import select_backend
349
+ ```
350
+
351
+ Each operator's `resolve` reports the backend a call would run
352
+ and why it skipped earlier candidates:
353
+
354
+ ```python
355
+ import torch
356
+ from deformops.registry import msda
357
+
358
+ value = torch.randn(1, 80, 4, 8)
359
+ shapes = torch.tensor([[8, 8], [4, 4]])
360
+ starts = torch.tensor([0, 64])
361
+ locations = torch.rand(1, 10, 4, 2, 4, 2)
362
+ weights = torch.rand(1, 10, 4, 2 * 4)
363
+
364
+ print(msda.resolve(value, shapes, starts, locations, weights))
365
+ ```
366
+
367
+ `deformops.registry.msda`, `.conv` and `.attn` hold their operator's backends
368
+ and are callable:
369
+
370
+ ```python
371
+ out = msda(value, shapes, starts, locations, weights, method="softmax", backend=Backend.EAGER)
372
+ ```
373
+
374
+ A `backend=None` call runs a public `torch.library` operator
375
+ (`torch.ops.deformops.msda`, `msda_softmax`, `attn`, `attn_softmax` or `conv`).
376
+ Its one kernel for every device applies the selection policy
377
+ and runs the chosen backend's private operator
378
+ under `torch.ops.deformops_triton` or `torch.ops.deformops_cuda`;
379
+ [operator registration and export](https://github.com/tue-p8n/deformops/blob/main/docs/export.md) explains why.
380
+ Calling a public operator directly selects the same way:
381
+
382
+ ```python
383
+ out = torch.ops.deformops.msda_softmax(value, shapes, starts, locations, weights)
384
+ ```
385
+
386
+ A `backend=None` call with a backend-specific keyword such as `fp32acc`
387
+ selects its backend in Python, since the public operators take none,
388
+ and a traced graph records the private operator.
389
+
390
+ The reference implementations are `msda.reference`, `attn.reference` and `conv.reference`
391
+ in the submodules of `deformops.backends.reference`.
392
+ `Backend` has no member for them, so dispatch never selects them;
393
+ the test suite checks every other backend against them.
394
+
395
+ ## Pre-made models
396
+
397
+ Each model package loads the upstream graph and overrides its deformable attention;
398
+ Mask2Former also overrides the `einsum` mask predictor, which obstructs export,
399
+ and the scipy-backed Hungarian matcher.
400
+ Architecture, weights and post-processing are upstream.
401
+
402
+ | Model | Full graph | deformops override | Extra |
403
+ | :-------------- | :---------------------- | :----------------------------------------------------- | :------------------ |
404
+ | Mask2Former | `models.Mask2Former` | MSDA + mask predictor + torchmatch matcher | `[mask2former]` |
405
+ | Deformable DETR | `models.DeformableDETR` | Hugging Face MSDA core | `[deformable-detr]` |
406
+
407
+ ```bash
408
+ pip install "deformops[mask2former]" # transformers + torchmatch
409
+ pip install "deformops[deformable-detr]" # transformers + timm
410
+ ```
411
+
412
+ Mask2Former needs Python 3.13 and torch 2.11 or newer, through `torchmatch`,
413
+ which publishes no Python 3.12 wheel.
414
+
415
+ The patched graphs compile and export on every backend;
416
+ see [Compile and export](#compile-and-export).
417
+
418
+ The overrides are importable without the full graph:
419
+
420
+ ```python
421
+ from deformops.models.mask2former import (
422
+ PatchedMask2FormerHungarianMatcher,
423
+ PatchedMask2FormerLoss,
424
+ PatchedMask2FormerMaskPredictor,
425
+ PatchedMask2FormerPixelDecoderEncoderMultiscaleDeformableAttention,
426
+ )
427
+ from deformops.models.deformable_detr import PatchedMultiScaleDeformableAttention
428
+
429
+ # Or reuse the library integration directly:
430
+ from deformops.integrations.transformers import MultiScaleDeformableAttention
431
+ from deformops.integrations.mmcv import MultiScaleDeformableAttention as MMCVAttention
432
+ from deformops.integrations.transformers import patch as patch_transformers
433
+ ```
434
+
435
+ The integration `patch(...)` helpers replace matching upstream modules in place
436
+ and raise `RuntimeError` when nothing matches,
437
+ usually because the upstream version renamed its modules.
438
+ `strict=False` accepts a model with nothing to replace.
439
+
440
+ A full Mask2Former loads from a Hugging Face checkpoint (network access)
441
+ and takes a batch of RGB images:
442
+
443
+ ```python
444
+ import torch
445
+ from deformops.models.mask2former import Mask2Former
446
+
447
+ model = Mask2Former.from_pretrained("facebook/mask2former-swin-base-coco-instance").eval()
448
+ pixel_values = torch.randn(1, 3, 384, 384) # (B, 3, H, W), normalized RGB
449
+
450
+ with torch.no_grad():
451
+ outputs = model(pixel_values) # class_queries_logits, masks_queries_logits
452
+ ```
453
+
454
+ `deformops.models.mask2former.DEFAULT_CHECKPOINT` names the base COCO instance checkpoint;
455
+ the other Swin sizes (tiny, small, large) and the COCO panoptic and ADE20K semantic checkpoints load the same way.
456
+
457
+ Models built on mmcv 2.x's own `MultiScaleDeformableAttention`,
458
+ such as the mmdet 3.x Deformable DETR and DINO,
459
+ run on deformops through `deformops.integrations.mmcv.patch(model)`,
460
+ with the `deformops[mmcv]` extra.
461
+
462
+ ## Compile and export
463
+
464
+ Every operator is a registered `torch.library` operator with autograd,
465
+ so models using them work with `torch.compile`, `torch.export` and AOTInductor.
466
+ The test suite runs each recipe below on both pre-made models,
467
+ on every backend and with `backend=None`.
468
+
469
+ `torch.compile(model, fullgraph=True)` compiles inference and the backward pass.
470
+ With `backend=None`, Dynamo records the public operator,
471
+ and the compiled graph holds the private operator selected for the traced call.
472
+ In training mode, both upstream models call `isfinite(...).all()` on their hidden states,
473
+ which breaks the graph in upstream code;
474
+ compile them for training without `fullgraph`.
475
+
476
+ `torch.export.export` keeps each public operator (`torch.ops.deformops.*`)
477
+ and each private operator an explicit backend selects
478
+ (`torch.ops.deformops_triton.*` and `torch.ops.deformops_cuda.*`)
479
+ as a single node.
480
+ The eager backend exports as standard operators.
481
+ `ExportedProgram.run_decompositions()` replaces each public operator
482
+ with the private operator selected on the host that runs the decomposition.
483
+ A saved program names its operators,
484
+ so `torch.export.load` needs deformops installed and `import deformops.registry` run first.
485
+ Mask2Former exports with a dynamic batch size.
486
+ Deformable DETR resizes its pixel mask with the batch in the channel position,
487
+ which fixes the batch size of its export on CUDA.
488
+
489
+ AOTInductor runs those decompositions before compiling,
490
+ so a package holds the kernel selected at build time,
491
+ with the Triton kernels compiled in.
492
+ A package built on the Triton or eager backend,
493
+ or with `backend=None` selecting Triton,
494
+ loads without deformops installed.
495
+ A package that calls the CUDA operators needs their kernel library loaded first
496
+ in the process that loads it:
497
+
498
+ ```python
499
+ import torch
500
+ from deformops.backends import load
501
+ from deformops.models import Mask2Former
502
+ from deformops.registry import Backend
503
+
504
+ model = Mask2Former.from_pretrained(backend=Backend.CUDA).cuda().eval()
505
+ image = torch.randn(2, 3, 384, 384, device="cuda")
506
+ batch = torch.export.Dim("batch", min=1, max=16)
507
+
508
+ with torch.no_grad():
509
+ program = torch.export.export(model, (image,), dynamic_shapes={"pixel_values": {0: batch}})
510
+ torch._inductor.aoti_compile_and_package(program, package_path="mask2former.pt2")
511
+
512
+ # In the process that serves the package:
513
+ load("msda", Backend.CUDA)
514
+ served = torch._inductor.aoti_load_package("mask2former.pt2")
515
+ ```
516
+
517
+ AOTInductor needs `openssl` on `PATH` to hash headers.
518
+
519
+ ## Environment variables
520
+
521
+ | Variable | Effect |
522
+ | :----------------------------- | :------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
523
+ | `DEFORMOPS_DOWNLOAD` | `0` (default) skips the download. `1` downloads a prebuilt CUDA kernel library from `DEFORMOPS_OPS_INDEX_URL` when none is installed or compiled; failure raises. |
524
+ | `DEFORMOPS_COMPILE` | `0` (default) never compiles. `1` compiles when no library is available. `always` uses a library compiled on this machine; `rebuild` recompiles it from scratch. |
525
+ | `DEFORMOPS_FALLBACK_WARNINGS` | `0` silences `BackendFallbackWarning`. |
526
+ | `DEFORMOPS_OPS_INDEX_URL` | Base URL of the prebuilt operator index, for a private mirror. |
527
+ | `DEFORMOPS_OPS_CACHE_DIR` | The cache directory for downloaded libraries. Defaults to `~/.cache/deformops/ops`. |
528
+ | `DEFORMOPS_OPS_ALLOW_INSECURE` | Permits a plaintext `http://` index. Off by default, since the manifest holds the library checksums and plaintext defeats that check. |
529
+ | `DEFORMOPS_OPT_LEVEL` | Optimization level, `0` to `3`, for a runtime JIT build. Unset, the build uses the level the extension declares. |
530
+ | `TORCH_EXTENSIONS_DIR` | Where JIT builds are cached. Unset, torch's default build root, one directory per package and environment. |
531
+ | `DEFORMOPS_VALIDATE` | Enables operand checks that read the level table back to the host. Off by default because they synchronize. |
532
+ | `DEFORMOPS_EXTENSIONS` | Build-time only. `all` compiles the non-experimental CUDA kernels, `everything` adds the experimental ones, `none` compiles nothing, and a comma-separated list names kernels. See [Installation](#installation). |
533
+
534
+ ## Examples
535
+
536
+ The [tests](https://github.com/tue-p8n/deformops/tree/main/tests/deformops)
537
+ hold runnable examples, one file per operator,
538
+ with the pre-made models under [tests/deformops/models](https://github.com/tue-p8n/deformops/tree/main/tests/deformops/models).
539
+
540
+ ## Support
541
+
542
+ Report bugs and propose changes through GitHub issues and pull requests;
543
+ for questions about the method, see the [paper](https://openreview.net/forum?id=Q4jZ7zKNKx).
544
+
545
+ ## Contributing
546
+
547
+ [CONTRIBUTING.md](https://github.com/tue-p8n/deformops/blob/main/CONTRIBUTING.md)
548
+ describes the layout, conventions, commands and CI.
549
+
550
+ ## Citation
551
+
552
+ Cite the paper that introduced each operator you use.
553
+
554
+ ### Multi-scale deformable attention
555
+
556
+ This covers `MultiScaleDeformAttn2d`, `registry.msda`,
557
+ and the Mask2Former, Deformable DETR and mmcv integrations.
558
+ Cite Deformable DETR for the operator and Efficient MSDA for the kernels.
559
+
560
+ ```bibtex
561
+ @inproceedings{zhu2021deformable,
562
+ title = {Deformable {DETR}: Deformable Transformers for End-to-End Object Detection},
563
+ author = {Zhu, Xizhou and Su, Weijie and Lu, Lewei and Li, Bin and Wang, Xiaogang and Dai, Jifeng},
564
+ booktitle = {International Conference on Learning Representations (ICLR)},
565
+ year = {2021},
566
+ url = {https://openreview.net/forum?id=gZ9hCDWe6ke}
567
+ }
568
+
569
+ @article{stolle2026efficient,
570
+ title = {Efficient Multi-Scale Deformable Attention on GPUs},
571
+ author = {Stolle, Kurt H. W.},
572
+ journal = {Transactions on Machine Learning Research},
573
+ year = {2026},
574
+ url = {https://openreview.net/forum?id=Q4jZ7zKNKx}
575
+ }
576
+ ```
577
+
578
+ ### Deformable convolution
579
+
580
+ This covers `DeformConv2d` and `integrations.torchvision.deform_conv2d`.
581
+ Cite v1, and v2 as well if you pass a modulation mask.
582
+
583
+ ```bibtex
584
+ @inproceedings{dai2017deformable,
585
+ title = {Deformable Convolutional Networks},
586
+ author = {Dai, Jifeng and Qi, Haozhi and Xiong, Yuwen and Li, Yi and Zhang, Guodong and Hu, Han and Wei, Yichen},
587
+ booktitle = {Proceedings of the IEEE International Conference on Computer Vision (ICCV)},
588
+ year = {2017},
589
+ pages = {764--773},
590
+ doi = {10.1109/ICCV.2017.89}
591
+ }
592
+
593
+ @inproceedings{zhu2019deformable,
594
+ title = {Deformable {ConvNets} v2: More Deformable, Better Results},
595
+ author = {Zhu, Xizhou and Hu, Han and Lin, Stephen and Dai, Jifeng},
596
+ booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
597
+ year = {2019},
598
+ pages = {9300--9308},
599
+ doi = {10.1109/CVPR.2019.00953}
600
+ }
601
+ ```
602
+
603
+ ### Single-scale deformable sampling
604
+
605
+ This covers `DeformAttn2d` and `DeformSample2d`,
606
+ which implement the deformable convolution v3 (DCNv3) operator introduced in InternImage.
607
+
608
+ ```bibtex
609
+ @inproceedings{wang2023internimage,
610
+ title = {{InternImage}: Exploring Large-Scale Vision Foundation Models with Deformable Convolutions},
611
+ author = {Wang, Wenhai and Dai, Jifeng and Chen, Zhe and Huang, Zhenhang and Li, Zhiqi and Zhu, Xizhou and Hu, Xiaowei and Lu, Tong and Lu, Lewei and Li, Hongsheng and Wang, Xiaogang and Qiao, Yu},
612
+ booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
613
+ year = {2023},
614
+ pages = {14408--14419},
615
+ doi = {10.1109/CVPR52729.2023.01385}
616
+ }
617
+ ```
618
+
619
+ ## Resources
620
+
621
+ - [PyTorch Blog on Triton kernel optimization](https://pytorch.org/blog/accelerating-triton/)
622
+
623
+ ## License
624
+
625
+ MIT. See [LICENSE](https://github.com/tue-p8n/deformops/blob/main/LICENSE).