From 6ff59c763c821617d5c62e65d80423f7112562a3 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 2 Nov 2026 12:51:19 -0800 Subject: [PATCH 1/2] CUDA backend: CUDA graph fixes for mask indexing, .item() and changed input shapes, masked copies as where, tensor checks dropped before compiling, float64, _reshape_copy, scatter_reduce and varlen_attn, follow-ups (#24326) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## What is wrong today With `enable_cuda_graph_for_method` set, the CUDA backend runs a method a few times normally, captures a CUDA graph on the fourth call, and replays that graph after it. Some calls go wrong: - **A call that cannot be captured fails.** A mask index, such as `x[mask]` or `x[mask] = v`, reads a count back to the host, and a capture cannot record that. The capture call returned an error. An output that is first strided at the capture call failed the same way. A `.item()` on a device tensor, such as a `torch.cond` predicate computed on the GPU, ended the process. - **A replay can return a wrong result with no error.** Replay checked only each input's byte count. An input that changes from `(2, 3, 4, 4)` to `(2, 4, 3, 4)` keeps its byte count, so it replayed the captured shape and returned the wrong output. ## What this change does - Replay compares each input's shape with the captured one and refuses a different shape with `InvalidArgument`. The error names the dim and both sizes. A captured graph only works for the input shapes of the capture call, and the code comment on the graph phases now says so. - When a capture fails, the backend runs the same call again without a graph and returns its result. Graphs stay off for that method, as before. The retry is safe because a capture records kernels without running them, and it happens only if the new graph never started running. The `.item()` shims now refuse to read inside a capture, as the mask index does, so a capture that reads a device value is retried too instead of ending the process. - A method whose program keeps a constant in host memory, such as a `torch.cond` flag updated on every call, now runs without a graph. Host code runs only during the capture, so a graph would keep the branch of the capture call, and a retry after a failed capture would take the other one. The log names the constant found in host memory, or says the constants could not be listed. - An expected capture failure, such as the mask index refusing to run inside a capture, is logged at Info instead of Error, since the call then succeeds. A real failure of the call without a graph is still logged at Error. AOTInductor's generated code still prints its own line for the refused call. - The allocations a captured graph makes are recorded only after its first launch succeeds. Before, a failed first launch left addresses that were never allocated, and a later cleanup passed them to `cudaFree`. A failed capture now cleans up with `CudaGraphState::release()` instead of a copy of it. - These CUDA graph changes do not change the exported program, so older programs get them too. A program that uses the new shims described below (scatter, scatter_reduce, float64 or float16) needs a runtime that has them. Smaller commits: - The CUDA random shims seed from `std::random_device` instead of the clock, so two processes started in the same second no longer draw the same numbers. If `std::random_device` throws, or the copy of the seed to the GPU fails, the shim logs an error and frees what it allocated, and the next call tries again. - The wheel's shared library check now lists the mask index and `index_put` shims. - With PyTorch 2.14, a program that writes a one-element buffer, such as a counter, kept returning its first update, because Inductor folded the update into a constant. The AOTI backends (CUDA and Metal) now turn off Inductor's joint-graph constant folding, for the whole program, when a program writes a buffer. On a small Gemma 4 decoder that writes its KV cache, folding changes two nodes per graph; with it off, export and decode times stayed within run-to-run noise and the generated tokens were identical. - Unit tests for `aoti_torch_get_size`. - Comments on the index shims and on `set_output_data_ptr` that described accumulation, the sort count limit and the output index wrongly are corrected. - The mask index tests move to a new `test_index_shims.py`, since they need the Python runtime. They add a 2-D mask, and `index_put` with repeated indices in deterministic mode (half and bfloat16 accumulation, last write wins, an int32 index). - `test_output_shape_changes_at_capture` ends with a call that only a live graph rejects, so it fails if graphs are off. The KV cache recapture test checks that dropping the graph also clears the captured input shapes. ## Follow-ups to #24314 - Delete `test_align_corners_compiles_with_dynamic_sizes`. It could not fail: it passed with no resize rewrite at all and with the min and max in the align_corners ratio replaced by a plain `if`, because `torch.compile` compiles again for an output size of 1. With torch 2.14 and `CI` set it also failed for an unrelated reason, and it is the one failing test of the CUDA unit test job on main. `test_bicubic_antialias_resize` already runs an exported program at an output height of 1, where a plain `if` breaks. - The comment on that min and max now gives the real reason: an exported program keeps the branch it was traced with, so a plain `if` divides by zero at an output size of 1. - A backend that keeps its data in the `.pte`, such as Metal, now refuses a `data_name` compile spec built by hand too, before compiling anything. Before, only `generate_data_name_compile_spec` refused it, and a hand-built spec was ignored with no name check. - The CUDA docs warn that Windows and macOS usually treat data names that differ only in case as the same file. ## Copies through a boolean mask no longer read the mask count back to the host `out[mask] = src[mask]` ran as a mask index and a mask `index_put`. Each of them read the mask's count back to the host and waited for the GPU, on every call. Transformers' SmolVLM vision embeddings do this for their position ids. When the values are the same mask's picks from a tensor shaped like the destination, every element goes back to where it came from. So a new pass in the CUDA backend writes the copy as `torch.where(mask, src, out)`, which gives the same result, bit for bit. It does so inside `torch.cond` branches and other subgraphs too. It rewrites only when `where` would lay out its result like the destination, since `where` takes its layout from its inputs, the mask first, while `index_put` keeps the destination's layout. It also drops the index and the checks on its size when nothing else reads them and each check tests only the index's count and holds for any mask. A user check on the count, such as `torch._check(x[mask].shape[0] >= 1)`, keeps the index. The pass compares the two masks with a work list, so masks with long histories do not hit Python's recursion limit. On an A100, five calls of the SmolVLM position id code went from 10 count reads and 10 stream waits to none. ## Tensor checks no longer stop the CUDA export The argument checks in `torch.distributions` export as `aten._assert_async.msg`. The CUDA backend compiled each one into a device assert inside a Triton kernel. While compiling, AOTInductor runs every kernel once to pick its launch settings, on made-up buffers: random floats, and zeros for integers and bools. So a check that reads a bool from an earlier kernel saw false and failed, even though it holds for every real input. The export then stopped with "device-side assert triggered". A Gaussian policy head, `MultivariateNormal(mean, scale_tril=...)` with 14 actions, hits this. A new pass in the CUDA backend now removes these checks before compiling, so the CUDA delegate no longer runs them. The rest of ExecuTorch already removes the same checks before a program is saved. A failed device assert could not come back as an error anyway: it breaks the CUDA context for the rest of the process. Checks on sizes (`_assert_scalar`) stay, since they run on the host. ## float64 weights load in the CUDA backend A program with a float64 weight lowered and saved, then failed to load with `Malformed CUDA FQN weight metadata`. This happened with a float64 buffer, and with float64 math that constant propagation folds into one constant, such as a sinusoidal time embedding. The export wrote the weight correctly, but the runtime had no float64 dtype: the weight cache refused it, and the runtime's tensor type and dtype shims had none either. - The runtime's tensor type now has `Double`, with its 8 byte size. The same type wraps a method's inputs and outputs, so a float64 input or output works too. - The shims export `aoti_torch_dtype_float64`, which compiled code calls to name the dtype of a float64 constant. - The weight cache accepts `Double`, and uses 8 bytes when it checks that a weight fits its storage. - The shims that still refused float64 accept it: a float64 `.item()` (`aoti_torch_item_float64`, which did not exist, so the first execute ended the process), the single-element read those entry points share, `index_put` with accumulate, and the sort shim. A float16 `.item()` had the same problem with `aoti_torch_item_float16` and works now too. - A float read with `.item()` reached Triton kernels truncated to an integer, in float32 too, so a scale of 1.5 computed as 1 with no error. AOTInductor declares every symbolic argument of a kernel's launcher `int64_t`. A compile-time patch in the CUDA backend now declares it with the float type the kernel takes, the change that the open pytorch/pytorch#201882 makes in PyTorch. ## `aten._reshape_copy` lowers to the CUDA backend Torch-TensorRT turns `aten.view` into `aten._reshape_copy`. A view that TensorRT leaves to PyTorch reaches the CUDA backend in that form, and the export failed with `missing fallback kernels: aten::_reshape_copy`, since Inductor has no lowering for it and torch registers no decomposition. The CUDA backend now decomposes it the way ATen computes it: a contiguous copy, viewed to the new shape. ## `varlen_attn` lowers to the CUDA backend `torch.nn.attention.varlen.varlen_attn` runs attention over several sequences packed into one tensor. Lowering it failed with `missing fallback kernels: torch_attn::_varlen_attn`, since it calls Flash or cuDNN attention, which this backend does not ship. The CUDA backend now decomposes it into ops Inductor lowers: attention over all the packed tokens, with a mask that keeps each query in its own sequence. The output is contiguous, like the op's output for a query in the usual (tokens, heads, dim) order, so reshaping it before the output projection works. The decomposition matches the op for causal masks, sliding windows, grouped-query attention, and query and key sequences of different lengths. For a query sequence longer than `max_k + 1`, Flash attention 2 (PyTorch's default) and Flash attention 3 give different results for some windows, because version 2 changes the window based on `max_k`. So a window other than `(left, 0)` lowers only when `max_q <= max_k + 1` is known at export time, as in self-attention, and raises a clear error otherwise. `seqused_k` and `block_table` (paged KV) raise a clear error too. It computes the weights with a softmax, which subtracts each row's largest score first, so the weights still add up to one when scores are large. Tokens before the first sequence or past the last one, such as the padding of a fixed-size buffer, do not change any sequence's output, even when they hold NaN. Unlike in Flash attention, a NaN or inf in one sequence's values does reach the other sequences. It holds a float32 score per head for every pair of packed tokens and runs much slower than Flash attention, so it suits packs of a few thousand tokens, such as a few camera images. ## `scatter_reduce` with amax, amin, prod and mean lowers to the CUDA backend Lowering `x.scatter_reduce(dim, index, src, "amax")` failed with `missing fallback kernels: aoti_torch_cuda_scatter_reduce_two_out`, and so did "amin", "prod" and "mean". Inductor writes a scatter itself only when it can add atomically, and leaves the other reductions to an ATen kernel that this backend did not ship. A sum into int64 or bool, a sum into bfloat16 on GPUs older than sm_90, and deterministic mode use the same kernel. - A new shim next to the index shims implements it the way ATen does: `include_self`, elements the index does not name, the count and integer rounding of the mean, NaN in amax and amin, and bool for sum only. Half and bfloat16 sums add in float, as ATen does in deterministic mode. A product rounds to the tensor's dtype at every step and a mean rounds its sum before dividing, as ATen does, so half overflows and integers wrap where ATen's do. A mean also counts its terms in the tensor's dtype, as ATen does, so 128 int8 ones average to 1. - It sorts the writes by the element they land on and reduces each element's values in index order, so the result is deterministic. It also works inside a CUDA graph. The sort is now one helper that `index_put` uses too. - It is slower than ATen's atomic kernel: 1.5 to 3.2 times on an A100 at 512 output rows, 4096 source rows and 1024 columns, and up to 7 times when many values land on few elements. A decomposition into ops Inductor lowers natively was measured first. It has to compare every output element with every source element, and it was 8 to 89 times slower than ATen's kernel. - `aoti_cuda_shims.lib` and the wheel's shared library check list the new shim. The import library now also lists the shims the DLL exported but it left out, such as `aoti_torch_dtype_float16` and the `aoti_torch_item_` getters, so a float16 model lowered for Windows links. ## CUDA programs load without torch on machines with no libgomp On Linux, Inductor adds OpenMP (`-fopenmp`, `-lgomp`) to every link. Unless the compiler passes `--as-needed` to the linker, as Ubuntu's GCC does by default, the compiled library lists `libgomp.so.1` as a dependency even when it calls no OpenMP function. PyTorch ships a libgomp, so this went unnoticed while torch was imported. A process without torch failed to load the program with `libgomp.so.1: cannot open shared object file` on any machine with no system libgomp, such as a minimal container. The CUDA backend now links with `-Wl,--as-needed`, so the linker keeps only the libraries the program uses. A program that does call OpenMP keeps libgomp. PyTorch already does the same for MPS builds on macOS, and pytorch/pytorch#202902 adds it to Inductor for Linux. The torch-free tests in `test_pybindings_cuda_dlpack.py` now run without a system libgomp, and the ones that check which libraries got loaded also check for libgomp. On an A100, a small MLP and CNN exported with this change load without torch and match eager where the loader cannot find libgomp. This patch, and the one that lets Triton templates run on GPUs with few SMs, now stay in place while any thread lowers. Before, when two threads lowered at once, the first one to finish removed the patch from the other, and the rest of that compile ran without it. ## A failed CUDA load no longer leaves its library in the temporary directory To load a method, the CUDA backend writes the program's compiled library to a temporary file, such as `/tmp/executorch_cuda__0.so`. Only `destroy()` removed that file, and `destroy()` never runs when the load fails, so every load that failed after writing the file left it behind. A scope guard now removes the file unless the load succeeds. `test_missing_weights_blob.py` now checks that a load that fails for a missing weights blob leaves nothing in the temporary directory. Without the fix, that check fails. The XNNPACK high-rank permute fix moved to #24380. ## The Triton SDPA leaves keys past the end out of the softmax With `triton_kernel_mode` on (the default), the CUDA backend runs `scaled_dot_product_attention` as its Triton kernel, `triton::sdpa`. For power-of-two head sizes, its main kernel reads keys in tiles of 32, 64 or 128. When the number of keys is not a multiple of the tile, the last tile reaches past the last key. The kernel loaded those positions as zeros but never removed them from the softmax, so without an attention mask or a causal mask each one got a score of 0 and took part of the attention weight. Calls that pass `kv_len` with neither a mask nor `is_causal`, such as the DFlash draft model's global attention layers, had the same problem. On an A100, with the query at 0 and the values at 1, where the right output is 1, one key gave 0.031, 33 keys gave 0.516 and 80 keys gave 0.625. With random inputs and one key, the output was off by 3.4 when the largest value was 3.67. All of these now match eager PyTorch. The kernel now gives every position past the last key, or past `kv_len`, a score of minus infinity, as the other attention kernels in `sdpa.py` already do. The first version of the kernel did this, and a later rewrite dropped it. Checking every tile cost up to 8% where no tile reaches past the end, so only calls that need the check pay for it. A mask or a causal bound already drops those keys, so calls with one skip it. The launcher also passes the largest tile size that divides the number of keys, and tiles that divide it skip the check at compile time. So calls with a mask or a causal bound, and tiles that divide a key count known at export time, compile to the same code as before. Without a mask or a causal bound, a `kv_len` or a symbolic key count checks every tile. `test_triton_sdpa.py` drops its note that called this a kernel limitation, runs its basic test with a single key, and runs every tile without a mask over 80 and 96 keys and over a `kv_len` of 80. These fail without the fix. ## `scatter` with a tensor source lowers to the CUDA backend Lowering `x.scatter(dim, index, src)` failed with a `missing fallback kernels` error that names `aoti_torch_cuda_scatter_src_out` when `src` is bfloat16, int64 or bool. Inductor writes a scatter from a tensor itself only in a dtype it can add atomically, so it leaves int64, bool and, on GPUs older than sm_90, bfloat16 to an ATen kernel, and every dtype in deterministic mode. A mixture-of-experts layer that writes its expert outputs back with `torch.zeros_like(h).scatter(1, index, src)` in bfloat16 hit this on an A100. A new shim next to the `scatter_reduce` one implements it with the same checks, the same sorting of the writes by the element they land on, and the same kernel. The kernel gets a new mode in which only the last write to each element stores its value. So when values land on the same element, the last one in index order wins, as ATen does in deterministic mode; ATen's default CUDA kernel keeps any one of them. As in that default kernel, an index out of range, negative ones included, triggers a device-side assert. In deterministic mode, ATen's scatter counts a negative index from the end instead. The shim takes float, double, half, bfloat16, the integer dtypes and bool. Like ATen, both scatter shims now take an int32 index as well as an int64 one. Mostly because of the sort, the shim is about 10 to 15 times slower than ATen's default kernel for a mixture-of-experts style bfloat16 scatter of 4 to 34 million values on an A100. The Windows import library `aoti_cuda_shims.lib` and the wheel's shared library check now list the new shim. On an A100 with PyTorch 2.14 and 2.15, scatters in float32, float16, bfloat16, int64 and bool lower, run and match eager PyTorch exactly. So does a small mixture-of-experts layer built the same way. `test_scatter_with_tensor_source` fails without the fix, and its int32 index case fails when the shims take only int64. ## Review order Start with the CUDA graph changes: `cuda_backend.cpp` and `cuda_delegate_handle.h` in `backends/cuda/runtime`, the mask and `.item()` shims in `runtime/shims/index.cu` and `memory.cpp`, and their tests in `test_output_stride_rearrange.py`. Then the lowering side: the passes in `backends/cuda/passes`, the decompositions and compile-time patches in `backends/cuda/cuda_backend.py`, and the folding switch in `backends/aoti/aoti_backend.py`. Then the scatter and scatter_reduce shims in `index.cu` with `test_index_shims.py`. The remaining commits are small follow-ups that stand on their own. ## Testing On an A100, with the CUDA backend built from this branch: - `test_index_shims.py`, `test_output_stride_rearrange.py`, `test_cuda_export.py`, `test_antialiased_resize.py`, `test_merge_ptes.py` and `test_cuda_weight_metadata.py` pass. - With the backend built from main, the new CUDA graph tests fail. - The replay error names the dim, for example "input 0 has size 4 in dim 1, captured with 3". - `test_cuda_kv_cache` passes, and fails when `release()` keeps the captured input shapes. - With the min and max replaced by a plain `if`, `test_bicubic_antialias_resize` fails. - With the new data name refusal removed, or moved back behind the compile, the data name test fails. - A `torch.cond` on a host-memory flag matches eager on every call; with a graph it returned the captured branch. - A random program run from two processes started together drew the same numbers in 2 of 3 tries before the seed change, and different ones in 3 of 3 after. - Mask index programs exported from this branch match eager PyTorch on every call, both on a runtime built from main with CUDA graphs off and on this branch's runtime with CUDA graphs on. - The shared shim unit tests pass in a standalone googletest build. - `test_replace_masked_copy_with_where.py` passes, and fails when the pass does nothing, is not registered, or skips its accumulate, same-mask, boolean mask, shared picks or count check. It writes into a destination filled with 7, so it fails if `where` loses the values the mask does not pick. A contiguous destination with a channels_last source, or with a channels_last mask, keeps its `index_put`; before, the rewrite failed to lower when a later `view(-1)` needed the destination's layout. It also fails when the pass walks only the top-level graph, since a masked copy in a `torch.cond` branch then keeps its index, and when it takes two draws of a random mask for the same mask. Two masks computed with over a thousand ops each raised RecursionError before. A check that reads the index's count but tests another size keeps the index; before, the pass removed it. - A masked copy lowered through the CUDA partitioner matches eager for masks that pick some, none and all positions. The test records the graph the pass returns, and fails when the pass is not registered or leaves the index in place. - A module that adds one to a device buffer on every call and reads a mask's count returns the call number on every call, so the capture call that runs again without a graph updates the buffer once. With the retry made to run the call twice, the capture call returns one too many. With PyTorch 2.14 it returned one on every call even without a graph, because Inductor folded the written buffer into a fill; it passes now with 2.14 and with 2.15 and 2.16 nightlies. - On the tiny Gemma 4 decoder from the Gemma 4 31B example's tests, which writes a KV cache, with PyTorch 2.14 on an A100: with the folding off, export took 193 s against 190.5 s with it on, a decode call took 945 µs against 1019 µs (medians of three exports and three runs on a busy machine), and the 41 generated tokens were the same. - The mask tests check that the runtime log reports one failed capture, so they fail when CUDA graphs are off or a failed capture is tried again. `test_host_read_runs_without_a_graph` also runs a module that scales by a `.item()` read; without the capture check in the item shims, the process aborts. - With the debug check on the phase before the retry turned into a release check, the mask tests pass. - `test_distribution_argument_check` passes with PyTorch 2.14 and fails with "device-side assert triggered" without the new pass. It also checks that the backend lists the pass, since without the pass the export fails only when autotuning happens to trip the assert. It writes the `MultivariateNormal` check as `torch._assert_async`, since PyTorch 2.14 cannot export the distribution itself. It also checks that the pass drops the check and the math that only fed it. A Gaussian policy head with hidden sizes 16, 64 and 256 exports, runs, and matches eager PyTorch. - A float64 buffer and a folded float64 constant load and match eager PyTorch. Before, both failed to load. `test_float64_weight_loads` fails without the fix. - `test_reshape_copy` in `test_cuda_export.py` fails without the fix. The one in `test_index_shims.py` runs `_reshape_copy` on a transposed, a sliced and a contiguous input and matches eager PyTorch. - `test_varlen_attn` in `test_cuda_export.py` fails without the fix. It lowers a program that flattens the output, compares the decomposition with the op on CUDA only (on ROCm, varlen_attn runs AOTriton, which lines an explicit window up with the start of each sequence, while the decomposition follows Flash attention on CUDA, which lines it up with the end), including a window in self-attention and NaN values outside the key sequences, and checks that a window the Flash attention versions disagree on raises. It also fails when the output is left non-contiguous, when the window check is removed or refuses every such window, and when the values outside the key sequences are not zeroed. Both `test_varlen_attn` tests include queries and keys filled with 1024, where every key gets the same weight; computing the weights as `exp(scores - logsumexp)` fails them, since at that size float32 rounds the log-sum-exp to a whole number and the weights no longer add up to one. The `test_varlen_attn` in `test_index_shims.py` lowers and runs programs that call `varlen_attn` with no window, a causal mask and the windows (8, 0) and (4, 2), flatten each output, and take `cu_seq_q` and `cu_seq_k` only at run time. They match eager PyTorch. - `test_missing_weights_blob.py`, `test_cuda_weight_reuse.py` and `test_device_placement.py` pass. - `test_scatter_reduce` now runs the sum too, which for int64 goes to the new shim. The test that a missing scatter kernel fails at export now tries to lower a scatter of a scalar in deterministic mode. The backend still lacks that kernel, so the test still checks the scatter fallback now that the scatter_reduce and scatter shims exist. It fails when the check stops watching that fallback. `test_scatter_mean_counts_in_its_dtype` fails with the old count and when a count that wraps to zero is not turned into one. `test_float64_and_float16_values` (a sort and `.item()` in each dtype) and the float64 masked accumulate fail with the old shims, and all pass now. The `.item()` scale is 1.5, so that test also fails without the patch that passes a float to Triton kernels as a float. - `test_scatter_reduce` and `test_sort_shim.py` pass. `test_scatter_reduce` fails without the fix, and also when the shim ignores `include_self`, rounds the integer mean toward zero, or lets a NaN lose in amax and amin. `test_scatter_strides` scatters along the middle dim of a 3-D tensor, given as a negative dim, from a strided source and from transposed operands, and fails when the shims walk only the last two dims. `test_scatter_reduce_edge_cases` covers a bool sum, int32 and half products that wrap and overflow, and an empty index. `test_scatter_reduce_sum_adds_in_float` sums 1000 half and bfloat16 values in deterministic mode, and fails when the sum rounds to the dtype at every step. `test_scatter_reduce_replays_in_a_graph` captures a `scatter_reduce` program in a CUDA graph and matches eager on every replay. - The weight cache, runtime tensor and shim unit tests pass in a standalone googletest build. Without the fix, the weight cache test of a float64 entry fails. - `test_small_gpu_override_survives_a_thread_that_leaves_first` overlaps two threads inside the small-GPU patch, and fails with the old save and restore. - With the seed copy made to fail once, the random shim logs it, fails that call, and seeds on the next one; before, it went on with an unseeded state. Not tested: a failed first graph launch (it cannot be triggered from a test), a `std::random_device` that throws, Windows, macOS case-insensitive file systems, the buffer folding change on Metal, and the ROCm build. No test covers the seed change, since a test cannot make two processes start in the same second. ### Checked on real robot policies On the release/1.6 stack with this change, on A100 and H100: - SmolVLA on the CUDA backend: one device-to-host copy per action chunk (only the action download), matching eager within the usual tolerance. - A Gaussian policy whose `torch.distributions` check used to stop the CUDA export: exports and runs, 0.76 ms per chunk. - A float64 weight that used to fail at load with "Malformed CUDA FQN weight metadata": loads and runs, max difference 0. - Four processes started together: four different random draws, against two out of four before. (cherry picked from commit 59541e11bb4e4013281e4311dfe5d1b31dc0ecf7) --- .ci/scripts/wheel/test_shared_libraries.py | 2 + backends/aoti/aoti_backend.py | 17 + backends/aoti/common_shims_slim.cpp | 4 + backends/aoti/common_shims_slim.h | 1 + backends/aoti/slim/c10/core/ScalarType.h | 14 +- .../slim/c10/core/test/test_scalar_type.cpp | 2 + .../slim/core/test/test_slimtensor_basic.cpp | 6 +- .../aoti/slim/factory/test/test_empty.cpp | 8 +- .../aoti/slim/factory/test/test_from_blob.cpp | 4 +- .../aoti/tests/test_common_shims_slim.cpp | 36 ++ backends/cuda/BUCK | 2 + backends/cuda/cuda_backend.py | 294 +++++++++- backends/cuda/passes/remove_assert_async.py | 50 ++ .../passes/replace_masked_copy_with_where.py | 220 ++++++++ .../test_replace_masked_copy_with_where.py | 245 +++++++++ backends/cuda/runtime/aoti_cuda_shims.lib | Bin 47886 -> 59238 bytes backends/cuda/runtime/cuda_backend.cpp | 155 +++--- backends/cuda/runtime/cuda_delegate_handle.h | 4 +- backends/cuda/runtime/cuda_weight_cache.cpp | 2 + backends/cuda/runtime/shims/index.cu | 503 +++++++++++++++++- backends/cuda/runtime/shims/index.h | 63 +++ backends/cuda/runtime/shims/memory.cpp | 26 + backends/cuda/runtime/shims/memory.h | 8 + backends/cuda/runtime/shims/rand.cu | 44 +- backends/cuda/runtime/shims/sort.cu | 14 + .../cuda/runtime/test/test_cuda_kv_cache.cpp | 2 + .../runtime/test/test_cuda_weight_cache.cpp | 12 +- backends/cuda/tests/test_cuda_export.py | 232 ++++++-- backends/cuda/tests/test_index_shims.py | 369 +++++++++++++ .../cuda/tests/test_missing_weights_blob.py | 43 +- .../tests/test_output_stride_rearrange.py | 109 +++- backends/cuda/tests/test_sort_shim.py | 2 + backends/cuda/tests/test_triton_sdpa.py | 36 +- backends/cuda/triton/kernels/sdpa.py | 28 + extension/cuda/runtime_api.h | 9 + extension/pybindings/README.md | 9 +- .../test/test_pybindings_cuda_dlpack.py | 30 +- 37 files changed, 2377 insertions(+), 228 deletions(-) create mode 100644 backends/cuda/passes/remove_assert_async.py create mode 100644 backends/cuda/passes/replace_masked_copy_with_where.py create mode 100644 backends/cuda/passes/tests/test_replace_masked_copy_with_where.py create mode 100644 backends/cuda/tests/test_index_shims.py diff --git a/.ci/scripts/wheel/test_shared_libraries.py b/.ci/scripts/wheel/test_shared_libraries.py index 85afe0b3718..0e86acbacc8 100644 --- a/.ci/scripts/wheel/test_shared_libraries.py +++ b/.ci/scripts/wheel/test_shared_libraries.py @@ -106,6 +106,8 @@ "aoti_torch_cuda_guard_set_index", "aoti_torch_cuda_index_Tensor", "aoti_torch_index_put_out", + "aoti_torch_cuda_scatter_reduce_two_out", + "aoti_torch_cuda_scatter_src_out", ) _QUANTIZED_KERNEL_SYMBOLS = ( diff --git a/backends/aoti/aoti_backend.py b/backends/aoti/aoti_backend.py index 87a15498de7..a069db4c1fc 100644 --- a/backends/aoti/aoti_backend.py +++ b/backends/aoti/aoti_backend.py @@ -53,6 +53,22 @@ def _checked_data_name(data_name: str) -> str: return data_name +def _written_buffer_options(program: ExportedProgram) -> Dict[str, Any]: + """Inductor options for a program that writes a buffer. + + They turn off Inductor's joint-graph constant folding for the whole + program, not only for the buffers it writes. Before torch 2.15, that + folding can turn the update of a one-element buffer, such as a counter, + into its value at export, so every call writes that same value + (https://github.com/pytorch/pytorch/pull/185077). + + TODO: remove once the pinned torch is 2.15 or newer. + """ + if program.graph_signature.buffers_to_mutate: + return {"joint_graph_constant_folding": False} + return {} + + @experimental( "This API and all of aoti-driven backend related functionality are experimental." ) @@ -292,6 +308,7 @@ def preprocess( device_name = cls.get_device_name() decomposition_table = cls.get_decomposition_table() options = cls.get_aoti_compile_options(compile_specs) + options.update(_written_buffer_options(edge_program)) # Move the edge_program to the target device. Routed through a hook so # backends can keep large non-parameter tensors (e.g. KV-cache buffers) diff --git a/backends/aoti/common_shims_slim.cpp b/backends/aoti/common_shims_slim.cpp index 7a5a53ee744..13d4a6d8f61 100644 --- a/backends/aoti/common_shims_slim.cpp +++ b/backends/aoti/common_shims_slim.cpp @@ -145,6 +145,10 @@ int32_t aoti_torch_dtype_float32() { return 6; // ScalarType::Float } +int32_t aoti_torch_dtype_float64() { + return 7; // ScalarType::Double +} + int32_t aoti_torch_dtype_bfloat16() { return 15; // ScalarType::BFloat16 } diff --git a/backends/aoti/common_shims_slim.h b/backends/aoti/common_shims_slim.h index af275b5514b..cf98a534eb6 100644 --- a/backends/aoti/common_shims_slim.h +++ b/backends/aoti/common_shims_slim.h @@ -83,6 +83,7 @@ aoti_torch_get_device_index(Tensor* tensor, int32_t* ret_device_index); // ============================================================ AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_float32(); +AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_float64(); AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_bfloat16(); AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_float16(); AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_int64(); diff --git a/backends/aoti/slim/c10/core/ScalarType.h b/backends/aoti/slim/c10/core/ScalarType.h index 9a99aecf992..0181ff9e8e8 100644 --- a/backends/aoti/slim/c10/core/ScalarType.h +++ b/backends/aoti/slim/c10/core/ScalarType.h @@ -13,12 +13,14 @@ #include #include +#include #include namespace executorch::backends::aoti::slim::c10 { -// Import BFloat16 from ExecuTorch's portable_type +// Import BFloat16 and Half from ExecuTorch's portable_type using BFloat16 = ::executorch::runtime::etensor::BFloat16; +using Half = ::executorch::runtime::etensor::Half; /// Enum representing the scalar type (dtype) of tensor elements. /// Note: Enum values must match PyTorch's c10::ScalarType for compatibility. @@ -30,7 +32,7 @@ enum class ScalarType : int8_t { Long = 4, // int64_t Half = 5, // float16 Float = 6, // float - // Double = 7, // double - not currently needed + Double = 7, // double // ComplexHalf = 8, // ComplexFloat = 9, // ComplexDouble = 10, @@ -50,6 +52,7 @@ constexpr ScalarType kInt = ScalarType::Int; constexpr ScalarType kLong = ScalarType::Long; constexpr ScalarType kHalf = ScalarType::Half; constexpr ScalarType kFloat = ScalarType::Float; +constexpr ScalarType kDouble = ScalarType::Double; constexpr ScalarType kBool = ScalarType::Bool; constexpr ScalarType kBFloat16 = ScalarType::BFloat16; @@ -72,6 +75,8 @@ inline size_t elementSize(ScalarType t) { return 2; // sizeof(__half) = 2 bytes case ScalarType::Float: return sizeof(float); + case ScalarType::Double: + return sizeof(double); case ScalarType::Bool: return sizeof(bool); case ScalarType::BFloat16: @@ -100,6 +105,8 @@ inline const char* toString(ScalarType t) { return "Half"; case ScalarType::Float: return "Float"; + case ScalarType::Double: + return "Double"; case ScalarType::Bool: return "Bool"; case ScalarType::BFloat16: @@ -116,7 +123,7 @@ inline const char* toString(ScalarType t) { /// @return true if the scalar type is floating point, false otherwise. inline bool isFloatingType(ScalarType t) { return t == ScalarType::Half || t == ScalarType::Float || - t == ScalarType::BFloat16; + t == ScalarType::Double || t == ScalarType::BFloat16; } /// Checks if the scalar type is an integral type (including bool optionally). @@ -157,6 +164,7 @@ inline bool isValidScalarType(ScalarType t) { case ScalarType::Long: case ScalarType::Half: case ScalarType::Float: + case ScalarType::Double: case ScalarType::Bool: case ScalarType::BFloat16: return true; diff --git a/backends/aoti/slim/c10/core/test/test_scalar_type.cpp b/backends/aoti/slim/c10/core/test/test_scalar_type.cpp index 4c06f7ef101..bc09676f2d7 100644 --- a/backends/aoti/slim/c10/core/test/test_scalar_type.cpp +++ b/backends/aoti/slim/c10/core/test/test_scalar_type.cpp @@ -38,6 +38,7 @@ const std::vector kAllScalarTypes = { {ScalarType::Long, 4, 8, "Long", false, true, true, false}, {ScalarType::Half, 5, 2, "Half", true, false, false, false}, {ScalarType::Float, 6, 4, "Float", true, false, false, false}, + {ScalarType::Double, 7, 8, "Double", true, false, false, false}, {ScalarType::Bool, 11, 1, "Bool", false, false, true, true}, {ScalarType::BFloat16, 15, 2, "BFloat16", true, false, false, false}, }; @@ -224,6 +225,7 @@ TEST_F(IsValidScalarTypeTest, AllSupportedTypesAreValid) { EXPECT_TRUE(isValidScalarType(ScalarType::Long)); EXPECT_TRUE(isValidScalarType(ScalarType::Half)); EXPECT_TRUE(isValidScalarType(ScalarType::Float)); + EXPECT_TRUE(isValidScalarType(ScalarType::Double)); EXPECT_TRUE(isValidScalarType(ScalarType::Bool)); EXPECT_TRUE(isValidScalarType(ScalarType::BFloat16)); } diff --git a/backends/aoti/slim/core/test/test_slimtensor_basic.cpp b/backends/aoti/slim/core/test/test_slimtensor_basic.cpp index 29ebe5875b4..b8e44acc22b 100644 --- a/backends/aoti/slim/core/test/test_slimtensor_basic.cpp +++ b/backends/aoti/slim/core/test/test_slimtensor_basic.cpp @@ -508,10 +508,10 @@ TEST(SlimTensorValidationTest, InvalidDtypeUndefined) { ""); } -TEST(SlimTensorValidationTest, InvalidDtypeDouble) { +TEST(SlimTensorValidationTest, InvalidDtypeComplexFloat) { std::vector sizes = {2, 3}; std::vector strides = {3, 1}; - size_t nbytes = 6 * sizeof(double); + size_t nbytes = 6 * 2 * sizeof(float); Storage storage = make_cpu_storage(nbytes); EXPECT_DEATH( @@ -519,7 +519,7 @@ TEST(SlimTensorValidationTest, InvalidDtypeDouble) { std::move(storage), makeArrayRef(sizes), makeArrayRef(strides), - static_cast(7)), // Double = 7 + static_cast(9)), // ComplexFloat = 9 ""); } diff --git a/backends/aoti/slim/factory/test/test_empty.cpp b/backends/aoti/slim/factory/test/test_empty.cpp index f4d0e31da9c..9cc3d6ef178 100644 --- a/backends/aoti/slim/factory/test/test_empty.cpp +++ b/backends/aoti/slim/factory/test/test_empty.cpp @@ -249,7 +249,7 @@ TEST(EmptyStridedTest, InvalidDtypeUndefined) { ""); } -TEST(EmptyStridedTest, InvalidDtypeDouble) { +TEST(EmptyStridedTest, InvalidDtypeComplexFloat) { std::vector sizes = {2, 3}; std::vector strides = {3, 1}; @@ -257,7 +257,7 @@ TEST(EmptyStridedTest, InvalidDtypeDouble) { empty_strided( makeArrayRef(sizes), makeArrayRef(strides), - static_cast(7)), // Double = 7 + static_cast(9)), // ComplexFloat = 9 ""); } @@ -280,9 +280,9 @@ TEST(EmptyTest, InvalidDtypeUndefined) { EXPECT_DEATH(empty({2, 3}, c10::ScalarType::Undefined), ""); } -TEST(EmptyTest, InvalidDtypeDouble) { +TEST(EmptyTest, InvalidDtypeComplexFloat) { EXPECT_DEATH( - empty({2, 3}, static_cast(7)), // Double = 7 + empty({2, 3}, static_cast(9)), // ComplexFloat = 9 ""); } diff --git a/backends/aoti/slim/factory/test/test_from_blob.cpp b/backends/aoti/slim/factory/test/test_from_blob.cpp index d066d7ba79c..40a00c0c682 100644 --- a/backends/aoti/slim/factory/test/test_from_blob.cpp +++ b/backends/aoti/slim/factory/test/test_from_blob.cpp @@ -327,12 +327,12 @@ TEST(FromBlobTest, InvalidDtypeUndefined) { from_blob(external_data, {2, 3}, c10::ScalarType::Undefined), ""); } -TEST(FromBlobTest, InvalidDtypeDouble) { +TEST(FromBlobTest, InvalidDtypeComplexFloat) { constexpr size_t kNumFloats = 6; float external_data[kNumFloats]; EXPECT_DEATH( - from_blob(external_data, {2, 3}, static_cast(7)), ""); + from_blob(external_data, {2, 3}, static_cast(9)), ""); } TEST(FromBlobTest, InvalidDeviceType) { diff --git a/backends/aoti/tests/test_common_shims_slim.cpp b/backends/aoti/tests/test_common_shims_slim.cpp index 2dd8475a4d8..952a79283f2 100644 --- a/backends/aoti/tests/test_common_shims_slim.cpp +++ b/backends/aoti/tests/test_common_shims_slim.cpp @@ -143,6 +143,28 @@ void runGetSizesTest(slim_c10::DeviceType device_type) { delete tensor; } +void runGetSizeTest(slim_c10::DeviceType device_type) { + std::vector sizes = {2, 3}; + std::vector strides = calculateContiguousStrides(sizes); + slim_c10::Device device(device_type, 0); + + Tensor* tensor = new Tensor(slim::empty_strided( + slim::makeArrayRef(sizes), + slim::makeArrayRef(strides), + slim_c10::ScalarType::Float, + device)); + + int64_t size = -1; + EXPECT_EQ(aoti_torch_get_size(tensor, 1, &size), Error::Ok); + EXPECT_EQ(size, 3); + EXPECT_EQ(aoti_torch_get_size(tensor, -2, &size), Error::Ok); + EXPECT_EQ(size, 2); + EXPECT_EQ(aoti_torch_get_size(tensor, 2, &size), Error::InvalidArgument); + EXPECT_EQ(aoti_torch_get_size(tensor, -3, &size), Error::InvalidArgument); + + delete tensor; +} + void runGetStridesTest(slim_c10::DeviceType device_type) { std::vector sizes = {2, 3, 4}; std::vector strides = calculateContiguousStrides(sizes); @@ -432,6 +454,10 @@ TEST_F(CommonShimsSlimTest, GetSizes_CPU) { runGetSizesTest(slim_c10::DeviceType::CPU); } +TEST_F(CommonShimsSlimTest, GetSize_CPU) { + runGetSizeTest(slim_c10::DeviceType::CPU); +} + TEST_F(CommonShimsSlimTest, GetStrides_CPU) { runGetStridesTest(slim_c10::DeviceType::CPU); } @@ -483,6 +509,13 @@ TEST_F(CommonShimsSlimTest, GetSizes_CUDA) { runGetSizesTest(slim_c10::DeviceType::CUDA); } +TEST_F(CommonShimsSlimTest, GetSize_CUDA) { + if (!isCudaAvailable()) { + GTEST_SKIP() << "CUDA not available"; + } + runGetSizeTest(slim_c10::DeviceType::CUDA); +} + TEST_F(CommonShimsSlimTest, GetStrides_CUDA) { if (!isCudaAvailable()) { GTEST_SKIP() << "CUDA not available"; @@ -555,6 +588,7 @@ TEST_F(CommonShimsSlimTest, NullTensorArgument) { EXPECT_EQ( aoti_torch_get_data_ptr(nullptr, &data_ptr), Error::InvalidArgument); EXPECT_EQ(aoti_torch_get_sizes(nullptr, &sizes), Error::InvalidArgument); + EXPECT_EQ(aoti_torch_get_size(nullptr, 0, &numel), Error::InvalidArgument); EXPECT_EQ(aoti_torch_get_strides(nullptr, &strides), Error::InvalidArgument); EXPECT_EQ(aoti_torch_get_dtype(nullptr, &dtype), Error::InvalidArgument); EXPECT_EQ(aoti_torch_get_dim(nullptr, &dim), Error::InvalidArgument); @@ -566,6 +600,7 @@ TEST_F(CommonShimsSlimTest, NullReturnPointer) { EXPECT_EQ(aoti_torch_get_data_ptr(tensor, nullptr), Error::InvalidArgument); EXPECT_EQ(aoti_torch_get_sizes(tensor, nullptr), Error::InvalidArgument); + EXPECT_EQ(aoti_torch_get_size(tensor, 0, nullptr), Error::InvalidArgument); EXPECT_EQ(aoti_torch_get_strides(tensor, nullptr), Error::InvalidArgument); EXPECT_EQ(aoti_torch_get_dtype(tensor, nullptr), Error::InvalidArgument); EXPECT_EQ(aoti_torch_get_dim(tensor, nullptr), Error::InvalidArgument); @@ -663,6 +698,7 @@ TEST_F(CommonShimsSlimTest, ConsistentPointerReturn) { TEST_F(CommonShimsSlimTest, DTypeConstants) { // Verify dtype constants match expected PyTorch ScalarType values EXPECT_EQ(aoti_torch_dtype_float32(), 6); // ScalarType::Float + EXPECT_EQ(aoti_torch_dtype_float64(), 7); // ScalarType::Double EXPECT_EQ(aoti_torch_dtype_bfloat16(), 15); // ScalarType::BFloat16 EXPECT_EQ(aoti_torch_dtype_int64(), 4); // ScalarType::Long EXPECT_EQ(aoti_torch_dtype_int32(), 3); // ScalarType::Int diff --git a/backends/cuda/BUCK b/backends/cuda/BUCK index 7f88126cd5f..7e8d4c88675 100644 --- a/backends/cuda/BUCK +++ b/backends/cuda/BUCK @@ -79,8 +79,10 @@ fbcode_target( "passes/__init__.py", "passes/move_cond_predicate_to_cpu.py", "passes/lower_offgraph_kv.py", + "passes/remove_assert_async.py", "passes/replace_fused_sdpa.py", "passes/replace_int64_floordiv.py", + "passes/replace_masked_copy_with_where.py", ], visibility = [ "//executorch/backends/cuda/...", diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 243437faa8a..3f368cdd37e 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -12,6 +12,7 @@ import math import os import shutil +import sys import threading import typing import weakref @@ -19,6 +20,9 @@ from typing import Any, Dict, final, List, Optional import torch + +# Registers torch_attn._varlen_attn, which the decomposition table names. +import torch.nn.attention.varlen # noqa: F401 from executorch.backends.aoti.aoti_backend import AotiBackend from executorch.backends.cuda.autotune.cuda_graph_timing import cuda_graph_timing from executorch.backends.cuda.autotune.inputs import autotune_input_scenarios @@ -39,12 +43,16 @@ from executorch.backends.cuda.passes.move_cond_predicate_to_cpu import ( MoveCondPredicateToCpuPass, ) +from executorch.backends.cuda.passes.remove_assert_async import RemoveAssertAsyncPass from executorch.backends.cuda.passes.replace_fused_sdpa import ( ReplaceFusedSdpaWithSdpaPass, ) from executorch.backends.cuda.passes.replace_int64_floordiv import ( ReplaceInt64FloorDivWithFloatPass, ) +from executorch.backends.cuda.passes.replace_masked_copy_with_where import ( + ReplaceMaskedCopyWithWherePass, +) from executorch.backends.cuda.target_smem import target_smem_context from executorch.backends.cuda.triton.kernels import triton_mutation_cache from executorch.backends.cuda.triton.replacement_pass import ( @@ -57,6 +65,7 @@ from torch._decomp import decomposition_table from torch._inductor.decomposition import conv1d_to_conv2d from torch._prims_common import suggest_memory_format +from torch.fx.experimental.symbolic_shapes import statically_known_true from torch.nn.attention import SDPBackend # --------------------------------------------------------------------------- @@ -124,11 +133,70 @@ def _has_rmask(_self) -> bool: IndexingOptions.has_rmask = orig_has_rmask -_SMALL_GPU_GUARD = threading.local() +class _Depth(threading.local): + def __init__(self) -> None: + self.value = 0 -@contextlib.contextmanager -def _allow_triton_templates_on_small_gpus(): +class _CompilePatch: + """A patch of Inductor that changes behavior only for the threads inside + ``scope()``. + + ``install`` applies the patch and returns the function that removes it. + The patched code calls the ``active`` function it is given to decide, on + each call, whether the calling thread asked for the change. Threads can + lower at the same time, so the first scope to enter installs the patch + and the last one to leave removes it: a thread that leaves first does not + take the patch away from one still compiling. + """ + + def __init__(self, install: typing.Callable[..., typing.Callable[[], None]]): + self._install = install + self._remove: typing.Callable[[], None] = lambda: None + self._lock = threading.Lock() + self._scopes = 0 + self._depth = _Depth() + + def _active(self) -> bool: + return self._depth.value > 0 + + @contextlib.contextmanager + def scope(self) -> typing.Iterator[None]: + with self._lock: + if self._scopes == 0: + self._remove = self._install(self._active) + self._scopes += 1 + self._depth.value += 1 + try: + yield + finally: + self._depth.value -= 1 + with self._lock: + self._scopes -= 1 + if self._scopes == 0: + self._remove() + + +def _install_small_gpu_override(active): + from torch._inductor import utils as inductor_utils + + original = inductor_utils.is_big_gpu + + def is_big_gpu(index_or_device=0) -> bool: + return active() or original(index_or_device) + + inductor_utils.is_big_gpu = is_big_gpu + + def remove() -> None: + inductor_utils.is_big_gpu = original + + return remove + + +_SMALL_GPU_OVERRIDE = _CompilePatch(_install_small_gpu_override) + + +def _allow_triton_templates_on_small_gpus() -> typing.ContextManager[None]: """Let Inductor use Triton templates on GPUs below its 68-SM floor. This backend autotunes with Triton only, so otherwise a matmul or convolution has no Triton choices left. @@ -139,23 +207,97 @@ def _allow_triton_templates_on_small_gpus(): TODO: remove once Inductor skips the floor when ATen is not an allowed backend (https://github.com/pytorch/pytorch/issues/141690). """ - from torch._inductor import utils as inductor_utils + return _SMALL_GPU_OVERRIDE.scope() - orig_is_big_gpu = inductor_utils.is_big_gpu - def _is_big_gpu(index_or_device=0) -> bool: - return getattr(_SMALL_GPU_GUARD, "active", False) or orig_is_big_gpu( - index_or_device +def _install_as_needed_links(active): + from torch._inductor import cpp_builder + + original = cpp_builder.get_cpp_torch_device_options + + def device_options(device_type, aot_mode=False, *args, **kwargs): + definitions, include_dirs, cflags, ldflags, *rest = original( + device_type, aot_mode, *args, **kwargs ) + if ( + aot_mode + and active() + and sys.platform.startswith("linux") + and not torch._inductor.config.aot_inductor.link_libtorch + and torch._inductor.config.aot_inductor.cross_target_platform is None + ): + ldflags = [*ldflags, "Wl,--as-needed"] + return (definitions, include_dirs, cflags, ldflags, *rest) + + cpp_builder.get_cpp_torch_device_options = device_options + + def remove() -> None: + cpp_builder.get_cpp_torch_device_options = original + + return remove + + +_AS_NEEDED_LINKS = _CompilePatch(_install_as_needed_links) - prev_active = getattr(_SMALL_GPU_GUARD, "active", False) - inductor_utils.is_big_gpu = _is_big_gpu - _SMALL_GPU_GUARD.active = True - try: - yield - finally: - _SMALL_GPU_GUARD.active = prev_active - inductor_utils.is_big_gpu = orig_is_big_gpu + +def _link_libraries_as_needed() -> typing.ContextManager[None]: + """Record only the shared libraries the compiled program calls into. + + Inductor puts OpenMP on every Linux link. Unless the compiler asks the + linker to link as needed, as Ubuntu's GCC does by default, the linker + records libgomp even when nothing in the program calls OpenMP. Torch ships + a libgomp, which hides this whenever torch is loaded, but a process without + torch then fails to load the program on a machine with no system libgomp. + ``--as-needed`` drops only libraries nothing references, so a program that + does call OpenMP keeps it. Links with libtorch are left alone: a library + there can be needed without being referenced, because its static + initializers register something. + + Only threads inside this context get the flag. + + TODO: remove once the pinned PyTorch has the same flag in Inductor + (https://github.com/pytorch/pytorch/pull/202902). + """ + return _AS_NEEDED_LINKS.scope() + + +def _install_float_kernel_args(active): + from torch._inductor.codegen.cpp_wrapper_gpu import ( + DeferredTritonCallWrapper, + TRITON_SIGNATURE_TO_CPP, + ) + + original = DeferredTritonCallWrapper._get_cpp_param_type + + def get_cpp_param_type(self, name, arg_type, signature=None): + param = original(self, name, arg_type, signature) + # PyTorch 2.14 does not pass the signature on every path. + kind = (signature or (self.triton_meta or {}).get("signature", {})).get(name) + if active() and param == f"int64_t {name}" and kind in ("fp32", "fp64"): + return f"{TRITON_SIGNATURE_TO_CPP[kind]} {name}" + return param + + DeferredTritonCallWrapper._get_cpp_param_type = get_cpp_param_type + + def remove() -> None: + DeferredTritonCallWrapper._get_cpp_param_type = original + + return remove + + +_FLOAT_KERNEL_ARGS = _CompilePatch(_install_float_kernel_args) + + +def _pass_float_scalars_to_triton_as_floats() -> typing.ContextManager[None]: + """Declare a float that a Triton kernel takes, such as a value read with + ``.item()``, as a float in the C++ function that launches the kernel. + AOTInductor declares every such argument int64_t, so 1.5 reaches the + kernel as 1. + + TODO: remove once the pinned PyTorch takes the type from the kernel's + signature (https://github.com/pytorch/pytorch/pull/201882). + """ + return _FLOAT_KERNEL_ARGS.scope() def _full_zeros_preserving_strides(x: torch.Tensor, device) -> torch.Tensor: @@ -284,8 +426,6 @@ def _cpu_clone_preserve_strides(x: torch.Tensor) -> torch.Tensor: # `triton_heuristics.py` (clones for benchmark — must stay on GPU for # Triton). Discriminate by caller frame so we only force CPU clones for # the buffer-lifting path. - import sys - caller = sys._getframe(1).f_code.co_name if caller == "_unlift_graph": # KV-cache buffers are emptied (storage resize_(0)) by the low-memory @@ -554,6 +694,110 @@ def _conv_transpose_as_conv( return conv(x, weight, bias, [1] * n, [0] * n, dilation, groups) +def _reshape_copy(x: torch.Tensor, size: List[int]) -> torch.Tensor: + """ + Torch-TensorRT turns a view into _reshape_copy, and a view that TensorRT + leaves to PyTorch reaches this backend in that form. Inductor has no + lowering for it and this backend no kernel. ATen computes it as a + contiguous copy viewed to the new shape. + """ + return x.clone(memory_format=torch.contiguous_format).view(size) + + +def _varlen_attn( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + cu_seq_q: torch.Tensor, + cu_seq_k: Optional[torch.Tensor], + max_q: int, + max_k: int, + is_causal: bool = False, + scale: Optional[float] = None, + window_size: Optional[List[int]] = None, + enable_gqa: bool = False, + seqused_k: Optional[torch.Tensor] = None, + block_table: Optional[torch.Tensor] = None, + num_splits: Optional[int] = None, + backend: Optional[int] = None, +) -> typing.Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + varlen_attn runs Flash or cuDNN attention, and this backend has a kernel + for neither. Attend over all the packed tokens instead, with a mask that + keeps each query in its own sequence. This holds a float32 score per head + for every pair of packed tokens, 1 GiB for 16 heads and 4096 tokens. As in + Flash attention, a causal mask or a window lines up the ends of each query + and key sequence, and a query with no key to attend to gets zeros and an + infinite log-sum-exp. Unlike in Flash attention, a NaN or inf in one + sequence's values also reaches the other sequences. + """ + if cu_seq_k is None or seqused_k is not None or block_table is not None: + raise NotImplementedError( + "The CUDA backend lowers varlen_attn only with cu_seq_k and without " + "seqused_k or block_table." + ) + left, right = window_size if window_size is not None else (-1, -1) + if is_causal: + right = 0 + # Flash attention 2, PyTorch's default, drops a window side of max_k or + # more and gives a left side alone a right side of max_k, while later + # versions keep the window. That only matters for a query sequence + # longer than max_k + 1. + if (right > 0 or (right < 0 and left >= 0)) and not statically_known_true( + max_q <= max_k + 1 + ): + raise NotImplementedError( + "The CUDA backend lowers a varlen_attn window other than (left, 0) " + "only when max_q <= max_k + 1 is known, since Flash attention " + "versions differ on it for longer query sequences. Got " + f"window_size={window_size}, max_q={max_q}, max_k={max_k}." + ) + + def sequence_ends(cu: torch.Tensor, total: int): + # Each token's index in cu of its sequence's end, and its (negative) + # distance from that end. Tokens before cu[0] get the index 0, and + # tokens past cu[-1] get cu.numel() and, through the clamp, their + # distance from cu[-1]. So each group forms one more sequence. + token = torch.arange(total, dtype=cu.dtype, device=cu.device) + end = torch.searchsorted(cu, token, right=True) + return end, token - cu[end.clamp(max=cu.numel() - 1)] + + end_q, from_end_q = sequence_ends(cu_seq_q, query.size(0)) + end_k, from_end_k = sequence_ends(cu_seq_k, key.size(0)) + allowed = end_q[:, None] == end_k[None, :] + offset = from_end_k[None, :] - from_end_q[:, None] + if left >= 0: + allowed = allowed & (offset >= -left) + if right >= 0: + allowed = allowed & (offset <= right) + # Masked keys get zero weight, but NaN * 0 is NaN. Zero the values outside + # the key sequences so that a NaN or inf there stays out of them. + outside = (end_k == 0) | (end_k == cu_seq_k.numel()) + value = value.masked_fill(outside[:, None, None], 0) + + if enable_gqa: + key, value = ( + t.repeat_interleave(query.size(1) // t.size(1), dim=1) for t in (key, value) + ) + if scale is None: + scale = query.size(-1) ** -0.5 + scores = torch.einsum("qhd,khd->hqk", query.float(), key.float()) * scale + scores = scores.masked_fill(~allowed, float("-inf")) + lse = scores.logsumexp(-1) + empty = lse == float("-inf") + # Not exp(scores - lse): at large scores, float32 rounds lse as coarsely + # as the scores, and the weights no longer add up to one. + probs = scores.softmax(-1).masked_fill(empty[..., None], 0) + lse = lse.masked_fill(empty, float("inf")) + out = torch.einsum("hqk,khd->qhd", probs, value.float()).to(query.dtype) + rng_state = torch.zeros(2, dtype=torch.uint64, device=query.device) + # einsum gives a permuted view. The op lays out its output in the + # query's dimension order, so a query in the usual (T, H, D) order, even + # a slice of a packed tensor, gives a contiguous output. Later views of + # it rely on that. + return out.contiguous(), lse, rng_state + + def _linear_kernel(x: torch.Tensor) -> torch.Tensor: return torch.clamp(1 - x.abs(), min=0) @@ -632,8 +876,8 @@ def _antialiased_resize( each output's taps, one dim at a time. Inductor has no lowering for these resizes, so it would call ATen's kernel, which this backend does not ship. """ - # The backend's tensors have no float64, and ATen's CUDA kernel has no - # integer version. + # The rewrite computes in float32, so float64 would lose precision, and + # ATen's CUDA kernel has no integer version. if input.dtype not in (torch.float32, torch.float16, torch.bfloat16): return NotImplemented # Half types are computed in float and rounded once, so they can differ @@ -970,6 +1214,8 @@ def get_supported_fallback_kernels(cls) -> Dict[str, Any]: "aoti_torch_cuda_randint_low_out": None, "aoti_torch_cuda_index_Tensor": None, "aoti_torch_index_put_out": None, + "aoti_torch_cuda_scatter_reduce_two_out": None, + "aoti_torch_cuda_scatter_src_out": None, } @classmethod @@ -980,6 +1226,8 @@ def get_decomposition_table(cls) -> Dict[Any, Any]: # torch._check_tensor_all exports to this op; Inductor has no lowering # for it and the shim no fallback kernel. torch.ops.aten._is_all_true.default: torch.ops.aten.all.default, + torch.ops.aten._reshape_copy.default: _reshape_copy, + torch.ops.torch_attn._varlen_attn.default: _varlen_attn, } for resize, support, kernel in ( (torch.ops.aten._upsample_bilinear2d_aa, 1, _linear_kernel), @@ -1017,9 +1265,11 @@ def get_custom_passes(cls, compile_specs: List[CompileSpec]) -> List[typing.Any] ) triton_kernel_mode = mode passes = [ + RemoveAssertAsyncPass(), MoveCondPredicateToCpuPass(), ReplaceInt64FloorDivWithFloatPass(), ReplaceFusedSdpaWithSdpaPass(), + ReplaceMaskedCopyWithWherePass(), ] cell_layout = any( spec.key == OFFGRAPH_KV_COMPILE_SPEC @@ -1189,6 +1439,10 @@ def _combined(): stack.enter_context(torch.nn.attention.sdpa_kernel([SDPBackend.MATH])) stack.enter_context(triton_mutation_cache.enabled()) stack.enter_context(target_smem_context()) + # Keep an unused libgomp out of the program, so it loads in a + # process without torch. See `_link_libraries_as_needed`. + stack.enter_context(_link_libraries_as_needed()) + stack.enter_context(_pass_float_scalars_to_triton_as_floats()) # On ROCm, is_big_gpu gates on architecture, not SM count. if torch.version.hip is None: stack.enter_context(_allow_triton_templates_on_small_gpus()) diff --git a/backends/cuda/passes/remove_assert_async.py b/backends/cuda/passes/remove_assert_async.py new file mode 100644 index 00000000000..08f9964d263 --- /dev/null +++ b/backends/cuda/passes/remove_assert_async.py @@ -0,0 +1,50 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import torch +from torch.fx import GraphModule +from torch.fx.passes.infra.pass_base import PassBase, PassResult + + +class RemoveAssertAsyncPass(PassBase): + """ + Removes tensor checks (``aten._assert_async.msg``) from the program the + CUDA backend compiles, along with the math that only fed them. + + ``torch._assert_async(cond, msg)`` exports as this op, and the argument + checks in ``torch.distributions`` use it. Inductor compiles each one into + a device assert inside a Triton kernel. AOTInductor then runs every kernel + once at compile time to pick its launch settings, on made-up buffers: + random floats, and zeros for integers and bools. A check that reads a bool + computed by an earlier kernel sees false there and fails, even when it + holds for every real input. A failed device assert also breaks the CUDA + context, so the export cannot go on. + + ExecuTorch already removes these checks from the rest of a program before + it is saved (see ``RemoveGraphAssertsPass``), so the delegate now does the + same. Unlike that pass, this one keeps the size checks, such as + ``aten._assert_scalar``: they run on the host, so they cannot break the + CUDA context. + """ + + def call(self, graph_module: GraphModule) -> PassResult: + modified = False + for module in graph_module.modules(): + if not isinstance(module, GraphModule): + continue + asserts = [ + node + for node in module.graph.nodes + if node.op == "call_function" + and node.target == torch.ops.aten._assert_async.msg + ] + for node in asserts: + module.graph.erase_node(node) + if asserts: + module.graph.eliminate_dead_code() + module.recompile() + modified = True + return PassResult(graph_module, modified) diff --git a/backends/cuda/passes/replace_masked_copy_with_where.py b/backends/cuda/passes/replace_masked_copy_with_where.py new file mode 100644 index 00000000000..e00b1ae4fd7 --- /dev/null +++ b/backends/cuda/passes/replace_masked_copy_with_where.py @@ -0,0 +1,220 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +""" +Graph Transformation Pass for Masked Copies. + +Rewrites a copy through a boolean mask into a select: + + out[mask] = src[mask] -> out = torch.where(mask, src, out) + +A mask index has a result size that depends on the data, so ``index.Tensor`` +and ``index_put`` with a mask run in the backend's index shims, and each call +reads the mask's count back to the host, which waits for the GPU. When the +values written are the same mask's picks from a tensor shaped like the +destination, every element goes back where it came from, so ``where`` gives +the same result, element for element, with no count and no host read. +""" + +from typing import Any, List, Optional + +import sympy +import torch +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dialects.edge._ops import EdgeOpOverload +from torch.fx import GraphModule, Node +from torch.fx.experimental.symbolic_shapes import statically_known_true +from torch.fx.passes.infra.pass_base import PassBase, PassResult +from torch.utils._sympy.value_ranges import bound_sympy, ValueRanges + +_INDEX = exir_ops.edge.aten.index.Tensor +_INDEX_PUT = exir_ops.edge.aten.index_put.default +_ASSERT_SCALAR = torch.ops.aten._assert_scalar.default + + +def _same_value(a: Any, b: Any) -> bool: + """ + Whether a and b are the same node, or the same deterministic op applied to + the same values, as with ``x[m.view(-1)] = y[m.view(-1)]``. + """ + # A work list, not recursion: two long mask computations would pass the + # recursion limit, and each node pair is compared once, since shared + # inputs would otherwise be compared again on every path that reaches them. + pending = [(a, b)] + compared = set() + while pending: + a, b = pending.pop() + if isinstance(a, (list, tuple)) and isinstance(b, (list, tuple)): + if len(a) != len(b): + return False + pending.extend(zip(a, b)) + continue + if isinstance(a, Node) != isinstance(b, Node): + return False + if not isinstance(a, Node): + if a != b: + return False + continue + if a is b or (a, b) in compared: + continue + compared.add((a, b)) + if a.op != "call_function" or a.target != b.target: + return False + if ( + isinstance(a.target, (EdgeOpOverload, torch._ops.OpOverload)) + and torch.Tag.nondeterministic_seeded in a.target.tags + ): + return False + if a.kwargs.keys() != b.kwargs.keys(): + return False + pending.append((a.args, b.args)) + pending.extend((a.kwargs[k], b.kwargs[k]) for k in a.kwargs) + return True + + +class ReplaceMaskedCopyWithWherePass(PassBase): + """ + Replace ``index_put(out, indices, index(src, indices))`` with + ``where(mask, src, out)`` when ``indices`` is one boolean mask (after any + number of ``None``), ``src`` and ``out`` have the same shape and dtype, + ``where`` lays out its result like ``out``, and the put does not + accumulate. + + The two masks may be one node or two equal computations. The ``index`` is + then removed too if nothing else reads it but the range checks export adds + on its size, since those hold for any mask. + """ + + def call(self, graph_module: GraphModule) -> PassResult: + modified = False + for module in graph_module.modules(): + if not isinstance(module, GraphModule): + continue + replaced = False + for node in list(module.graph.nodes): + mask = self._mask_of_copy(node) + if mask is None: + continue + self._replace(module, node, mask) + replaced = True + if replaced: + module.graph.lint() + module.recompile() + modified = True + return PassResult(graph_module, modified) + + @staticmethod + def _mask_of_copy(node: Node) -> Optional[Node]: + """The mask of an index_put that copies through it, else None.""" + if node.op != "call_function" or node.target != _INDEX_PUT: + return None + out, indices, values = node.args[:3] + accumulate = ( + node.args[3] if len(node.args) > 3 else node.kwargs.get("accumulate", False) + ) + if accumulate or not isinstance(values, Node) or values.target != _INDEX: + return None + src, src_indices = values.args + if not _same_value(src_indices, indices): + return None + masks = [entry for entry in indices if entry is not None] + if len(masks) != 1: + return None + mask_val = masks[0].meta.get("val") + out_val = out.meta.get("val") + src_val = src.meta.get("val") + if not all(isinstance(v, torch.Tensor) for v in (mask_val, out_val, src_val)): + return None + if mask_val.dtype != torch.bool or mask_val.dim() == 0: + return None + if src_val.dtype != out_val.dtype or src_val.dim() != out_val.dim(): + return None + if not all( + statically_known_true(a == b) for a, b in zip(src_val.shape, out_val.shape) + ): + return None + # index_put keeps out's layout, but where takes its layout from its + # inputs, the mask first. + condition = mask_val + for _ in range(out_val.dim() - indices.index(masks[0]) - mask_val.dim()): + condition = condition.unsqueeze(-1) + layout = torch.where(condition, src_val, out_val).stride() + if not all( + statically_known_true(a == b) for a, b in zip(layout, out_val.stride()) + ): + return None + return masks[0] + + def _replace(self, graph_module: GraphModule, node: Node, mask: Node) -> None: + graph = graph_module.graph + out, indices, values = node.args[:3] + src = values.args[0] + position = indices.index(mask) + # The mask covers the dims after the leading Nones; broadcasting lines + # it up from the right, so only the dims after it need a size-1 dim. + trailing = out.meta["val"].dim() - position - mask.meta["val"].dim() + with graph.inserting_before(node): + condition = mask + for _ in range(trailing): + condition_val = condition.meta["val"].unsqueeze(-1) + condition = graph.call_function( + exir_ops.edge.aten.unsqueeze_copy.default, args=(condition, -1) + ) + condition.meta["val"] = condition_val + where = graph.call_function( + exir_ops.edge.aten.where.self, args=(condition, src, out) + ) + where.meta = node.meta.copy() + node.replace_all_uses_with(where) + graph.erase_node(node) + dead = self._only_range_checked(values, mask, position) + if dead is not None: + for user in reversed(dead): + graph.erase_node(user) + + @staticmethod + def _only_range_checked( + index: Node, mask: Node, position: int + ) -> Optional[List[Node]]: + """ + The index and everything that reads it, in graph order, if all that + reads it is symbolic math on its sizes and checks that hold for every + count, its size at ``position``, in [0, mask.numel()], else None. + """ + count = index.meta["val"].shape[position] + if not isinstance(count, torch.SymInt): + return None + numel = mask.meta["val"].numel() + if isinstance(numel, torch.SymInt): + numel = numel.node.shape_env.bound_sympy(numel.node.expr).upper + within = ValueRanges(0, numel) + found = [index] + pending = list(index.users) + while pending: + user = pending.pop() + if user in found: + continue + if user.target == _ASSERT_SCALAR: + condition = user.args[0] + expr = ( + condition.meta.get("val").node.expr + if isinstance(condition, Node) + and isinstance(condition.meta.get("val"), torch.SymBool) + else None + ) + if expr is None or expr.free_symbols != {count.node.expr}: + return None + bounds = bound_sympy(expr, {count.node.expr: within}) + if not (bounds.is_singleton() and bounds.lower == sympy.true): + return None + # Computing a symbolic value has no effect of its own, and its + # readers are visited in turn. + elif not isinstance(user.meta.get("val"), (torch.SymInt, torch.SymBool)): + return None + found.append(user) + pending.extend(user.users) + order = {n: i for i, n in enumerate(index.graph.nodes)} + return sorted(found, key=order.__getitem__) diff --git a/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py b/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py new file mode 100644 index 00000000000..cefc8f11b80 --- /dev/null +++ b/backends/cuda/passes/tests/test_replace_masked_copy_with_where.py @@ -0,0 +1,245 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import operator +import unittest + +import torch +from executorch.backends.cuda.cuda_backend import CudaBackend +from executorch.backends.cuda.passes.replace_masked_copy_with_where import ( + ReplaceMaskedCopyWithWherePass, +) +from executorch.exir import to_edge +from executorch.exir.dialects._ops import ops as exir_ops +from torch.export import Dim, export + + +def _targets(graph_module): + return [n.target for n in graph_module.graph.nodes if n.op == "call_function"] + + +class TestReplaceMaskedCopyWithWherePass(unittest.TestCase): + def test_registered_in_cuda_backend(self): + self.assertIn( + "ReplaceMaskedCopyWithWherePass", + {type(p).__name__ for p in CudaBackend.get_custom_passes([])}, + ) + + def test_masked_copy_becomes_where(self): + """Also removes the index and the checks on its size.""" + + class Copy(torch.nn.Module): + def forward(self, x, mask): + # Not zeros, so a lost destination value shows. + out = torch.full_like(x, 7.0) + # A leading None and a trailing dim, and the mask computed twice, + # as Transformers' SmolVLM writes its position ids. + out[:, mask.view(-1)] = x[:, mask.view(-1)] + return out + + x = torch.randn(2, 5, 3) + masks = ( + torch.tensor([True, False, True, True, False]), + torch.zeros(5, dtype=torch.bool), + torch.ones(5, dtype=torch.bool), + ) + program = to_edge(export(Copy(), (x, masks[0]))).exported_program() + ReplaceMaskedCopyWithWherePass()(program.graph_module) + + targets = _targets(program.graph_module) + self.assertIn(exir_ops.edge.aten.where.self, targets) + for removed in ( + exir_ops.edge.aten.index.Tensor, + exir_ops.edge.aten.index_put.default, + torch.ops.aten._assert_scalar.default, + ): + self.assertNotIn(removed, targets) + for mask in masks: + torch.testing.assert_close( + program.module()(x, mask), Copy()(x, mask), rtol=0, atol=0 + ) + + def test_long_equal_masks(self): + """Two long computations of one mask, with inputs shared many times.""" + + class Copy(torch.nn.Module): + def forward(self, x, mask): + def compute(m): + for _ in range(1200): + m = m.logical_not() + for _ in range(30): + m = m & m + return m + + out = torch.zeros_like(x) + out[compute(mask)] = x[compute(mask)] + return out + + x = torch.randn(4, 3) + mask = torch.tensor([True, False, True, False]) + program = to_edge(export(Copy(), (x, mask))).exported_program() + ReplaceMaskedCopyWithWherePass()(program.graph_module) + self.assertIn(exir_ops.edge.aten.where.self, _targets(program.graph_module)) + + def test_masked_copy_in_a_cond_branch(self): + class Copy(torch.nn.Module): + def forward(self, x, mask): + def copy(x, mask): + out = torch.zeros_like(x) + out[mask] = x[mask] + return out + + return torch.cond( + x.sum() > 0, copy, lambda x, mask: x.clone(), (x, mask) + ) + + x = torch.randn(4, 3) + mask = torch.tensor([True, False, True, False]) + program = to_edge(export(Copy(), (x, mask))).exported_program() + ReplaceMaskedCopyWithWherePass()(program.graph_module) + + cond = next( + n + for n in program.graph_module.graph.nodes + if n.target is torch.ops.higher_order.cond + ) + branch = getattr(program.graph_module, cond.args[1].target) + self.assertIn(exir_ops.edge.aten.where.self, _targets(branch)) + self.assertNotIn(exir_ops.edge.aten.index_put.default, _targets(branch)) + for picked in (x.abs(), -x.abs()): + torch.testing.assert_close( + program.module()(picked, mask), Copy()(picked, mask), rtol=0, atol=0 + ) + + def test_other_writes_are_kept(self): + class Kept(torch.nn.Module): + def forward(self, x, mask, idx): + checked = torch.zeros_like(x) + # First, since export attaches the check to the first x[mask]. + picked = x[mask] + # A check on the count that does not hold for every mask. + torch._check(picked.shape[0] >= 1) + checked[mask] = picked + scaled = torch.zeros_like(x) + # The values are not the mask's picks of a tensor. + scaled[mask] = x[mask] * 2 + summed = torch.ones_like(x) + summed.index_put_((mask,), x[mask], accumulate=True) + other = torch.zeros_like(x) + other[mask] = x[mask.logical_not()] + # An integer index is not a mask, so where cannot take it. + rows = torch.zeros_like(x) + rows[idx] = x[idx] + # The picks are also read elsewhere, so the index stays. + shared = torch.zeros_like(x) + read = x[mask] + shared[mask] = read + return scaled, checked, summed, other, rows, shared, read.sum() + + x = torch.randn(4, 3) + mask = torch.tensor([True, False, True, False]) + idx = torch.tensor([2, 0]) + program = to_edge(export(Kept(), (x, mask, idx))).exported_program() + ReplaceMaskedCopyWithWherePass()(program.graph_module) + + targets = _targets(program.graph_module) + self.assertEqual(targets.count(exir_ops.edge.aten.index_put.default), 4) + self.assertEqual(targets.count(exir_ops.edge.aten.where.self), 2) + # The checked count and the shared picks keep their index. + self.assertEqual(targets.count(exir_ops.edge.aten.index.Tensor), 6) + for got, expected in zip( + program.module()(x, mask, idx), Kept()(x, mask, idx), strict=True + ): + torch.testing.assert_close(got, expected, rtol=0, atol=0) + + def test_different_layouts_are_kept(self): + """index_put keeps the destination's layout, but where may take the + mask's or the source's, and a later view(-1) then fails.""" + + class Source(torch.nn.Module): + def forward(self, x, mask): + src = x.contiguous(memory_format=torch.channels_last) + out = torch.zeros_like(x, memory_format=torch.contiguous_format) + out[mask] = src[mask] + return out + + class Mask(torch.nn.Module): + def forward(self, x, y): + mask = y.contiguous(memory_format=torch.channels_last) > 0 + out = torch.zeros_like(x) + out[mask] = x[mask] + return out + + x = torch.randn(4, 3, 5, 6) + for model, inputs in ( + (Source(), (x, torch.tensor([True, False, True, True]))), + (Mask(), (x, torch.randn(4, 3, 5, 6))), + ): + program = to_edge(export(model, inputs)).exported_program() + ReplaceMaskedCopyWithWherePass()(program.graph_module) + targets = _targets(program.graph_module) + self.assertIn(exir_ops.edge.aten.index_put.default, targets) + self.assertNotIn(exir_ops.edge.aten.where.self, targets) + + def test_random_masks_are_kept(self): + """Two draws of a random mask are different masks.""" + + class Copy(torch.nn.Module): + def forward(self, x): + out = torch.zeros_like(x) + out[torch.rand(4) > 0.5] = x[torch.rand(4) > 0.5] + return out + + program = to_edge(export(Copy(), (torch.randn(4, 3),))).exported_program() + ReplaceMaskedCopyWithWherePass()(program.graph_module) + self.assertIn( + exir_ops.edge.aten.index_put.default, _targets(program.graph_module) + ) + + def test_check_on_another_size_keeps_the_index(self): + """A check that reads the count but tests another size does not hold + for every mask.""" + + class Copy(torch.nn.Module): + def forward(self, x, mask, y): + out = torch.zeros_like(x) + out[mask] = x[mask] + return out, y * 2 + + x = torch.randn(4, 3) + mask = torch.tensor([True, False, True, False]) + y = torch.randn(3) + program = to_edge( + export(Copy(), (x, mask, y), dynamic_shapes=(None, None, {0: Dim("n")})) + ).exported_program() + graph = program.graph_module.graph + index = next( + n for n in graph.nodes if n.target == exir_ops.edge.aten.index.Tensor + ) + y_node = next(n for n in graph.nodes if n.name == "y") + # count * 0 + y.size(0) <= 4 reads the count, but whether it holds + # depends only on y's size. + with graph.inserting_before(index.next): + count = graph.call_function(torch.ops.aten.sym_size.int, (index, 0)) + size = graph.call_function(torch.ops.aten.sym_size.int, (y_node, 0)) + zero = graph.call_function(operator.mul, (count, 0)) + total = graph.call_function(operator.add, (zero, size)) + check = graph.call_function(operator.le, (total, 4)) + graph.call_function(torch.ops.aten._assert_scalar.default, (check, "")) + count.meta["val"] = index.meta["val"].shape[0] + size.meta["val"] = y_node.meta["val"].shape[0] + zero.meta["val"] = count.meta["val"] * 0 + total.meta["val"] = zero.meta["val"] + size.meta["val"] + check.meta["val"] = total.meta["val"] <= 4 + ReplaceMaskedCopyWithWherePass()(program.graph_module) + + targets = _targets(program.graph_module) + self.assertIn(exir_ops.edge.aten.where.self, targets) + self.assertIn(exir_ops.edge.aten.index.Tensor, targets) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/cuda/runtime/aoti_cuda_shims.lib b/backends/cuda/runtime/aoti_cuda_shims.lib index aa97264e523c2155b7d27a28f75db138096369e8..ad54a4fdb24e1e8c3395f0245d36a2d9881b5ef3 100644 GIT binary patch delta 3028 zcmZ{l4@}g_6~||mf4eL!{}FcirN; zo~VaWl&8c$rzbh1@rpuuSdB?r+Zto+#XxHRT{POei|3_!@shSlNm~2nHw$b9H~H*L z-n{wF`^}q~cjRBXr}u-7I@=w&b~Qu$RfIFfjxc7aWh^aQGBb%WLFN<5Telb!%x#s- zcQ7WfW=L#bGA3}`kmU9=CRlh*lGn?aprA`q*xt-oA+fwvvO-}@P!ugGe#w}?JtA2} zlLhbZmaN^ym|*=zNksu;g6cF$Z7gGgO=FVH4;T|P3`*LY_Aw^#7D;?GQLuv+73})0 zr2Qykg1s9go%0zJ9QcQ%i~IzKk4uhHPy%R>^yV`rIAxZcp?O3*JNFBw(MQgL3ri%I zLKzcW{zB5<&zNA~pk!z@V}ffWD!4I6a+AJP@Jk9taEk&G{BEt}&OF8hcgg4OSP^5t zzbCoh!l$xmigHhSwk z-p1zYO`htOjg_8l)s1qG*-f6gC$IKys866>yvP&AtHpllHvP-P&r^zVf1Pu zJQ}b>2dr4Zi(i1~CZE94CE+C)_&x60Y>~j(V_v{U)S)5rIG2dGLwkFJ22(>z`zs7G7jgB?bf56VZL2mqr zCro=ez+GM>zwqVi$oT)0pu3=jGHU)cWR*CK*vN8@|d=Lfc<T5OL66Cs9UQKE{nmxs{Zr{}7@zvYV4ip{Y^`b7;XBUp>v|`L_gWPk z@X?ARJ0syyy#a@}=5R=FlU6>Po|gpNyNzsNdldZEOP*aFk&x#zV%hdA4r6<*a{jJ; zWDSo%%MLq-d!1JK9YqHs@zGA&v5N<-GP85KXdPPXSM7Sdy?ZH#J%_CL=WI&RTsXFO zHV$;W%VEh8E6)9h*1Fvt3k60!_U@xk)E|?{gsz?#Tz_B%hx=fac_=**1M3fjVL?{` zhn`-F=@8A1KN*8@hpimeoT7NdwfgdO4F0j3B6NP@x7<7%gJ*ik691_c{{!UHdp;Jc zPAD9{>XSa)2JRE_IQJwe-Me5#V>)>iT{OX1L@bV-&gU?ENk*o#Gsuce!M~hc%whlM z3J&@x#rpn8INKM21LtX(;w%1*7`Yk=>n?>u+XW-;|BN;wevm{i&_<|Ga%mPeUR=iE z*^qxDKOCm*{oI7^%an>c*Q_!XKfN9an@vW{=%-ZdzLAByT32ziEAS#}4`7x5rKbb& zeUNzJkLWF;R*3TyWXrW<>6Ja4FT~NSgaA;@pgxg1)i65p9_u2U2Y89h}-+1cQm#a%G7ElWEuaZ9Gm zYk~K+zYNtP)(Yl$*W=3w%2pIDjpFs%4$-cDMi`Oy%K(^v~vnFx2AS zSL}3gZ-StSsy0Cl`yLV4>eDHsEL(^NFtsrPf80EdQH|y3s?}%m6N|YVj!{8oTG^Z? zFs@Fg@00l0fxQ}Ot@#3PzOGXja;1EqhPT=@3TjIfI8m?9T%m}@_dDj()PdC;W}{9q z1>70dtl?)J78>qM;V@ftT5_1Luhohk&7T^N#0c~tAK$~Rf4phKE^QV))6G$*w5zI! zA#seaGxDHLrM=vBw9`u89(V9EhweUCe#ANE)-Z&pZcZhSi)W-vccx>jW zV7&EiC|wcuda_graph_state.phase == CudaGraphPhase::Warmup && handle->cuda_graph_state.warmup_remaining == 0); - bool graph_ran = false; - Error error = - run_once(handle, args, n_inputs, n_outputs, is_capture_step, graph_ran); - if (error != Error::Ok && is_capture_step && !graph_ran) { + bool graph_launched = false; + Error error = run_once( + handle, args, n_inputs, n_outputs, is_capture_step, graph_launched); + if (error != Error::Ok && is_capture_step && !graph_launched) { // A capture records kernels without running them, so until the new graph // runs, the inputs and any state the program updates are untouched. A // program that keeps state in host memory, where host code does run @@ -855,12 +872,71 @@ class ET_EXPERIMENTAL CudaBackend final n_inputs, n_outputs, /*is_capture_step=*/false, - graph_ran); + graph_launched); } ET_CHECK_OK_OR_RETURN_ERROR(error); return commit_offgraph_kv_step(handle, kv_step_width); } + void destroy(DelegateHandle* handle_) const override { + if (handle_ == nullptr) { + return; + } + cuda::CudaDelegateHandle* handle = (cuda::CudaDelegateHandle*)handle_; + + mutable_state_forget_handle(handle); + if (handle->kv_cache != nullptr) { + handle->kv_cache->forget_handle(handle); + } + + // NOTE: AOTInductorModelContainerDelete does not work correctly with + // multiple .so files. Deleting one container frees shared resources, + // which causes segmentation faults when attempting to delete other + // containers. As a workaround, we skip explicit container deletion and + // defer cleanup to the OS. The corresponding shared library must remain + // loaded as well: the leaked container still owns objects whose code and + // process-wide state live in that library, so dlclose/FreeLibrary can + // invalidate them and crash during multi-method teardown. + // TODO(gasoonjia): Find a proper solution for safe container deletion. + // AOTInductorModelContainerDelete(handle->container_handle); + + // Remove the temporary shared library file + if (!handle->so_path.empty()) { + std::error_code remove_error; + std::filesystem::remove(handle->so_path, remove_error); + ET_CHECK_OR_LOG_ERROR( + !remove_error, + "Failed to remove temporary shared library %s: %s", + handle->so_path.c_str(), + remove_error.message().c_str()); + } + + delete handle; + + // The allocator lets the device pool keep freed memory so that repeated + // delegate execution does not pay to map it again. Nothing is running on + // this backend once the last handle is gone, so hand that memory back + // rather than hold it for the life of the process. + if (live_handles_.fetch_sub(1, std::memory_order_acq_rel) == 1) { + // Only frees the driver has already observed can be released, so without + // this the trim gives back nothing rather than less. A device wait rather + // than a stream wait because the frees went to whichever stream was + // current when the storage was released, which is not necessarily one + // this handle recorded, and teardown is not bound to the thread that ran + // the method either. + const cudaError_t sync_err = cudaDeviceSynchronize(); + if (sync_err != cudaSuccess) { + ET_LOG( + Error, + "cudaDeviceSynchronize before releasing the pool failed: %s.", + cudaGetErrorString(sync_err)); + (void)cudaGetLastError(); + } + // This also releases unused blocks cached by other allocator users. + CudaAllocator::release_cached_memory(-1); + } + } + private: // Names a constant of the program in host memory, or says that the // constants could not be listed, since unknown counts as host memory. @@ -886,15 +962,15 @@ class ET_EXPERIMENTAL CudaBackend final } // Runs the program once. On the capture step it captures the CUDA graph - // instead, then launches it once, and sets graph_ran once that launch is - // queued. + // instead, then launches it once, and sets graph_launched once that launch + // is queued. Error run_once( cuda::CudaDelegateHandle* handle, Span args, size_t n_inputs, size_t n_outputs, bool is_capture_step, - bool& graph_ran) const { + bool& graph_launched) const { // Wrrap them as SlimTensors for AOTI execution without any host<->device // copies. std::vector slim_inputs(n_inputs); @@ -1210,7 +1286,7 @@ class ET_EXPERIMENTAL CudaBackend final Internal, "cudaGraphLaunch (first replay) failed: %s", cudaGetErrorString(gerr)); - graph_ran = true; + graph_launched = true; handle->cuda_graph_state.note_graph_allocations(); // Copy capture-step outputs into the planned GPU ET buffers (D2D). @@ -1288,67 +1364,6 @@ class ET_EXPERIMENTAL CudaBackend final return Error::Ok; } - public: - void destroy(DelegateHandle* handle_) const override { - if (handle_ == nullptr) { - return; - } - cuda::CudaDelegateHandle* handle = (cuda::CudaDelegateHandle*)handle_; - - mutable_state_forget_handle(handle); - if (handle->kv_cache != nullptr) { - handle->kv_cache->forget_handle(handle); - } - - // NOTE: AOTInductorModelContainerDelete does not work correctly with - // multiple .so files. Deleting one container frees shared resources, - // which causes segmentation faults when attempting to delete other - // containers. As a workaround, we skip explicit container deletion and - // defer cleanup to the OS. The corresponding shared library must remain - // loaded as well: the leaked container still owns objects whose code and - // process-wide state live in that library, so dlclose/FreeLibrary can - // invalidate them and crash during multi-method teardown. - // TODO(gasoonjia): Find a proper solution for safe container deletion. - // AOTInductorModelContainerDelete(handle->container_handle); - - // Remove the temporary shared library file - if (!handle->so_path.empty()) { - std::error_code remove_error; - std::filesystem::remove(handle->so_path, remove_error); - ET_CHECK_OR_LOG_ERROR( - !remove_error, - "Failed to remove temporary shared library %s: %s", - handle->so_path.c_str(), - remove_error.message().c_str()); - } - - delete handle; - - // The allocator lets the device pool keep freed memory so that repeated - // delegate execution does not pay to map it again. Nothing is running on - // this backend once the last handle is gone, so hand that memory back - // rather than hold it for the life of the process. - if (live_handles_.fetch_sub(1, std::memory_order_acq_rel) == 1) { - // Only frees the driver has already observed can be released, so without - // this the trim gives back nothing rather than less. A device wait rather - // than a stream wait because the frees went to whichever stream was - // current when the storage was released, which is not necessarily one - // this handle recorded, and teardown is not bound to the thread that ran - // the method either. - const cudaError_t sync_err = cudaDeviceSynchronize(); - if (sync_err != cudaSuccess) { - ET_LOG( - Error, - "cudaDeviceSynchronize before releasing the pool failed: %s.", - cudaGetErrorString(sync_err)); - (void)cudaGetLastError(); - } - // This also releases unused blocks cached by other allocator users. - CudaAllocator::release_cached_memory(-1); - } - } - - private: mutable std::mutex cuda_graph_method_mutex_; std::string cuda_graph_method_; diff --git a/backends/cuda/runtime/cuda_delegate_handle.h b/backends/cuda/runtime/cuda_delegate_handle.h index ead44960a46..3bf7cdae5f6 100644 --- a/backends/cuda/runtime/cuda_delegate_handle.h +++ b/backends/cuda/runtime/cuda_delegate_handle.h @@ -104,8 +104,8 @@ struct CudaWeightStorage { // // The transition flow is: // Disabled ──(if CUDA graph is enabled for this method, the device -// has memory pools, and every constant is in device -// memory)──▶ Warmup +// has memory pools, and every constant is known to be in +// device memory)──▶ Warmup // Warmup ──(after `warmup_remaining` execute() calls)──▶ Replay // Warmup ──(if the capture fails, or an output is not laid out // like its ETensor)──▶ Disabled diff --git a/backends/cuda/runtime/cuda_weight_cache.cpp b/backends/cuda/runtime/cuda_weight_cache.cpp index 104b4461318..ae32c4282bd 100644 --- a/backends/cuda/runtime/cuda_weight_cache.cpp +++ b/backends/cuda/runtime/cuda_weight_cache.cpp @@ -124,6 +124,7 @@ bool is_supported_dtype(int32_t dtype) { case 4: // Long case 5: // Half case 6: // Float + case 7: // Double case 11: // Bool case 15: // BFloat16 return true; @@ -358,6 +359,7 @@ Error CudaWeightCache::validate_view(const Entry& entry) { item_size = 4; break; case aoti::slim::c10::ScalarType::Long: + case aoti::slim::c10::ScalarType::Double: item_size = 8; break; default: diff --git a/backends/cuda/runtime/shims/index.cu b/backends/cuda/runtime/shims/index.cu index 4b696cb03b9..7701bb8528b 100644 --- a/backends/cuda/runtime/shims/index.cu +++ b/backends/cuda/runtime/shims/index.cu @@ -17,8 +17,11 @@ #include #include #include +#include #include #include +#include +#include #include #include @@ -762,6 +765,49 @@ Error gather( return Error::Ok; } +// Stably sorts count (key, value) pairs by the low end_bit bits of the key. +Error sort_pairs( + const int64_t* keys, + int64_t* sorted_keys, + const int64_t* values, + int64_t* sorted_values, + int count, + int end_bit, + Scratch& scratch, + cudaStream_t stream, + const char* op) { + size_t temp_bytes = 0; + ET_CUDA_CHECK_OR_RETURN_ERROR(cub::DeviceRadixSort::SortPairs( + nullptr, + temp_bytes, + keys, + sorted_keys, + values, + sorted_values, + count, + 0, + end_bit, + stream)); + void* temp = scratch.allocate(temp_bytes); + ET_CHECK_OR_RETURN_ERROR( + temp != nullptr, + MemoryAllocationFailed, + "%s: allocating the sort scratch failed", + op); + ET_CUDA_CHECK_OR_RETURN_ERROR(cub::DeviceRadixSort::SortPairs( + temp, + temp_bytes, + keys, + sorted_keys, + values, + sorted_values, + count, + 0, + end_bit, + stream)); + return Error::Ok; +} + template Error put( const IndexPlan& plan, @@ -801,35 +847,16 @@ Error put( index_offsets_kernel<<>>( plan, points, out_offsets, values_offsets); ET_CUDA_KERNEL_LAUNCH_CHECK_OR_RETURN_ERROR(); - const auto count = static_cast(points); - size_t temp_bytes = 0; - ET_CUDA_CHECK_OR_RETURN_ERROR(cub::DeviceRadixSort::SortPairs( - nullptr, - temp_bytes, + ET_CHECK_OK_OR_RETURN_ERROR(sort_pairs( out_offsets, sorted_out_offsets, values_offsets, sorted_values_offsets, - count, - 0, + static_cast(points), static_cast(sizeof(int64_t) * 8), - stream)); - void* temp = scratch.allocate(temp_bytes); - ET_CHECK_OR_RETURN_ERROR( - temp != nullptr, - MemoryAllocationFailed, - "index_put: allocating the sort scratch failed"); - ET_CUDA_CHECK_OR_RETURN_ERROR(cub::DeviceRadixSort::SortPairs( - temp, - temp_bytes, - out_offsets, - sorted_out_offsets, - values_offsets, - sorted_values_offsets, - count, - 0, - static_cast(sizeof(int64_t) * 8), - stream)); + scratch, + stream, + "index_put")); put_sorted_kernel <<>>( plan, @@ -843,6 +870,372 @@ Error put( return Error::Ok; } +// Assign is a plain scatter: each write replaces the one before it. +enum class Reduction { Sum, Prod, Mean, Max, Min, Assign }; + +// A scatter visits every element of index. Dim `dim` of out is addressed by +// the index value, every other dim of out and src by the element's position. +struct ScatterPlan { + int ndim; + int dim; + int64_t sizes[kMaxDims]; // of index + int64_t index_strides[kMaxDims]; + int64_t out_strides[kMaxDims]; + int64_t src_strides[kMaxDims]; + int64_t dim_size; // of out + int64_t max_out_offset; + bool index_is_int32; +}; + +__global__ void scatter_offsets_kernel( + const ScatterPlan plan, + int64_t points, + const void* __restrict__ index, + int64_t* __restrict__ out_offsets, + int64_t* __restrict__ src_offsets) { + for (int64_t point = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + point < points; + point += static_cast(gridDim.x) * blockDim.x) { + int64_t linear = point; + int64_t index_offset = 0; + int64_t out_offset = 0; + int64_t src_offset = 0; + for (int d = plan.ndim - 1; d >= 0; --d) { + const int64_t coordinate = linear % plan.sizes[d]; + linear /= plan.sizes[d]; + index_offset += coordinate * plan.index_strides[d]; + src_offset += coordinate * plan.src_strides[d]; + if (d != plan.dim) { + out_offset += coordinate * plan.out_strides[d]; + } + } + const int64_t position = plan.index_is_int32 + ? static_cast(index)[index_offset] + : static_cast(index)[index_offset]; + // As in ATen's default CUDA kernel, a negative index is out of range, not + // counted from the end. + CUDA_KERNEL_ASSERT_MSG( + position >= 0 && position < plan.dim_size, "index out of bounds"); + out_offsets[point] = out_offset + position * plan.out_strides[plan.dim]; + src_offsets[point] = src_offset; + } +} + +// As ATen's safe_max and safe_min, a NaN on either side wins. +template +__device__ __forceinline__ AccT +combine(Reduction reduction, AccT result, AccT value) { + bool value_is_nan = false; + if constexpr (std::is_floating_point_v) { + value_is_nan = value != value; + } + switch (reduction) { + case Reduction::Prod: + return result * value; + case Reduction::Max: + return value_is_nan || result < value ? value : result; + case Reduction::Min: + return value_is_nan || value < result ? value : result; + default: + return result + value; + } +} + +// After a stable sort of the points by the element they write, the points of +// one element are a run in index order, and each element is written once, so +// the result does not depend on scheduling. An assignment keeps the last point +// of its run, as ATen's deterministic scatter does. Any other reduction walks +// its run in one thread. A sum adds in AccT, as ATen does in deterministic +// mode, so half and bfloat16 lose less to rounding. A product rounds to T at +// every step and a mean rounds its sum to T before dividing, as ATen does, so +// half overflows to infinity and integers wrap where ATen's do. +template +__global__ void scatter_reduce_sorted_kernel( + int64_t points, + const int64_t* __restrict__ out_offsets, + const int64_t* __restrict__ src_offsets, + T* __restrict__ out, + const T* __restrict__ src, + Reduction reduction, + bool include_self, + T identity) { + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + i < points; + i += static_cast(gridDim.x) * blockDim.x) { + const int64_t offset = out_offsets[i]; + if (reduction == Reduction::Assign) { + if (i + 1 == points || out_offsets[i + 1] != offset) { + out[offset] = src[src_offsets[i]]; + } + continue; + } + if (i > 0 && out_offsets[i - 1] == offset) { + continue; + } + auto result = static_cast(include_self ? out[offset] : identity); + int64_t count = include_self ? 1 : 0; + for (int64_t j = i; j < points && out_offsets[j] == offset; ++j, ++count) { + result = + combine(reduction, result, static_cast(src[src_offsets[j]])); + if (reduction == Reduction::Prod) { + result = static_cast(static_cast(result)); + } + } + if (reduction == Reduction::Mean) { + const auto sum = static_cast(static_cast(result)); + // ATen counts in T too, so a large count rounds or wraps, and a count + // that wraps to zero becomes one. + auto divisor = static_cast(static_cast(count)); + if (divisor == AccT(0)) { + divisor = AccT(1); + } + if constexpr (std::is_floating_point_v) { + result = sum / divisor; + } else { + // Integers round down, as ATen's division with rounding_mode="floor". + const AccT quotient = sum / divisor; + result = quotient - + (quotient * divisor != sum && (sum < 0) != (divisor < 0)); + } + } + out[offset] = static_cast(result); + } +} + +// What ATen writes to each element the index names before it reduces into +// it, when include_self is false. +template +T reduction_identity(Reduction reduction) { + switch (reduction) { + case Reduction::Prod: + return static_cast(AccT(1)); + case Reduction::Max: + if constexpr (std::is_floating_point_v) { + return static_cast(-std::numeric_limits::infinity()); + } else { + return std::numeric_limits::lowest(); + } + case Reduction::Min: + if constexpr (std::is_floating_point_v) { + return static_cast(std::numeric_limits::infinity()); + } else { + return std::numeric_limits::max(); + } + default: + return static_cast(AccT(0)); + } +} + +template +Error scatter_reduce( + const char* name, + const ScatterPlan& plan, + int64_t points, + const void* index, + void* out, + const void* src, + Reduction reduction, + bool include_self, + Scratch& scratch, + cudaStream_t stream) { + // The sort below takes the count as an int. + ET_CHECK_OR_RETURN_ERROR( + points <= INT_MAX, + NotSupported, + "%s: more than %d index elements are not supported", + name, + INT_MAX); + auto* offsets = + static_cast(scratch.allocate(4 * points * sizeof(int64_t))); + ET_CHECK_OR_RETURN_ERROR( + offsets != nullptr, + MemoryAllocationFailed, + "%s: allocating the offsets failed", + name); + int64_t* out_offsets = offsets; + int64_t* src_offsets = offsets + points; + int64_t* sorted_out_offsets = offsets + 2 * points; + int64_t* sorted_src_offsets = offsets + 3 * points; + scatter_offsets_kernel<<>>( + plan, points, index, out_offsets, src_offsets); + ET_CUDA_KERNEL_LAUNCH_CHECK_OR_RETURN_ERROR(); + // Sorting only the bits an offset into out can use is several times faster. + int end_bit = 1; + while (end_bit < 63 && (plan.max_out_offset >> end_bit) != 0) { + ++end_bit; + } + ET_CHECK_OK_OR_RETURN_ERROR(sort_pairs( + out_offsets, + sorted_out_offsets, + src_offsets, + sorted_src_offsets, + static_cast(points), + end_bit, + scratch, + stream, + name)); + scatter_reduce_sorted_kernel + <<>>( + points, + sorted_out_offsets, + sorted_src_offsets, + static_cast(out), + static_cast(src), + reduction, + include_self, + reduction_identity(reduction)); + ET_CUDA_KERNEL_LAUNCH_CHECK_OR_RETURN_ERROR(); + return Error::Ok; +} + +// Runs the scatter shim called `name`, with the checks ATen applies to +// aten::scatter and aten::scatter_reduce. +Error scatter( + const char* name, + Tensor* out, + Tensor* self, + int64_t dim, + Tensor* index, + Tensor* src, + Reduction reduction, + bool include_self) { + ET_CHECK_OR_RETURN_ERROR( + out != nullptr && index != nullptr && src != nullptr, + InvalidArgument, + "%s: null argument", + name); + ET_CHECK_OR_RETURN_ERROR( + out == self, + NotSupported, + "%s: out must be self, as AOTInductor calls it", + name); + ET_CHECK_OR_RETURN_ERROR( + out->is_cuda() && index->is_cuda() && src->is_cuda(), + NotSupported, + "%s: tensors must be on the CUDA device", + name); + ET_CHECK_OR_RETURN_ERROR( + src->dtype() == out->dtype(), + InvalidArgument, + "%s: src and self must have the same dtype", + name); + ET_CHECK_OR_RETURN_ERROR( + out->dtype() != ScalarType::Bool || reduction == Reduction::Sum || + reduction == Reduction::Assign, + NotSupported, + "%s: bool supports only sum, as in ATen", + name); + + // As in ATen, a 0-dim tensor counts as one of size 1. + const auto ndim = std::max(out->dim(), 1); + ET_CHECK_OR_RETURN_ERROR( + ndim <= kMaxDims, + NotSupported, + "%s: tensors with more than %d dims are not supported", + name, + kMaxDims); + if (dim < 0) { + dim += ndim; + } + ET_CHECK_OR_RETURN_ERROR( + dim >= 0 && dim < ndim, InvalidArgument, "%s: dim out of range", name); + // ATen checks the index only when it has elements. + if (index->numel() == 0) { + return Error::Ok; + } + ET_CHECK_OR_RETURN_ERROR( + index->dtype() == ScalarType::Long || index->dtype() == ScalarType::Int, + InvalidArgument, + "%s: index must be int64 or int32", + name); + ET_CHECK_OR_RETURN_ERROR( + std::max(index->dim(), 1) == ndim && + std::max(src->dim(), 1) == ndim, + InvalidArgument, + "%s: index and src must have as many dims as self", + name); + auto size = [](const Tensor* tensor, int64_t d) -> int64_t { + return tensor->dim() == 0 ? 1 : tensor->size(d); + }; + auto stride = [](const Tensor* tensor, int64_t d) -> int64_t { + return tensor->dim() == 0 ? 0 : tensor->stride(d); + }; + ScatterPlan plan{}; + plan.ndim = static_cast(ndim); + plan.dim = static_cast(dim); + plan.dim_size = size(out, dim); + plan.index_is_int32 = index->dtype() == ScalarType::Int; + int64_t points = 1; + for (int64_t d = 0; d < ndim; ++d) { + plan.max_out_offset += + std::max(size(out, d) - 1, 0) * stride(out, d); + const int64_t index_size = size(index, d); + ET_CHECK_OR_RETURN_ERROR( + index_size <= size(src, d) && (d == dim || index_size <= size(out, d)), + InvalidArgument, + "%s: index is larger than src, or than self, at dim %" PRId64, + name, + d); + plan.sizes[d] = index_size; + plan.index_strides[d] = stride(index, d); + plan.out_strides[d] = stride(out, d); + plan.src_strides[d] = stride(src, d); + points *= index_size; + } + + const auto device = static_cast(out->device_index()); + auto stream_result = getCurrentCUDAStream(out->device_index()); + ET_CHECK_OR_RETURN_ERROR( + stream_result.ok(), Internal, "%s: failed to get CUDA stream", name); + cudaStream_t stream = stream_result.get(); + Scratch scratch(device, stream); + auto run = [&](auto value, auto accumulator) { + using T = decltype(value); + using AccT = decltype(accumulator); + return scatter_reduce( + name, + plan, + points, + index->data_ptr(), + out->data_ptr(), + src->data_ptr(), + reduction, + include_self, + scratch, + stream); + }; + switch (out->dtype()) { + case ScalarType::Float: + return run(float{}, float{}); + case ScalarType::Double: + return run(double{}, double{}); + case ScalarType::Half: + return run(__half{}, float{}); + case ScalarType::BFloat16: + return run(__nv_bfloat16{}, float{}); + case ScalarType::Long: + return run(int64_t{}, int64_t{}); + case ScalarType::Int: + return run(int32_t{}, int64_t{}); + case ScalarType::Short: + return run(int16_t{}, int64_t{}); + case ScalarType::Char: + return run(int8_t{}, int64_t{}); + case ScalarType::Byte: + return run(uint8_t{}, int64_t{}); + case ScalarType::Bool: + return run(bool{}, bool{}); + default: + ET_LOG( + Error, + "%s: unsupported dtype %d", + name, + static_cast(out->dtype())); + return Error::NotSupported; + } +} + } // namespace #ifdef __cplusplus @@ -1024,6 +1417,9 @@ AOTITorchError aoti_torch_index_put_out( case ScalarType::Float: return put( plan, numel, may_repeat, dst, src, scratch, stream); + case ScalarType::Double: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); case ScalarType::Half: return put<__half, float, true>( plan, numel, may_repeat, dst, src, scratch, stream); @@ -1057,6 +1453,65 @@ AOTITorchError aoti_torch_index_put_out( } } +AOTITorchError aoti_torch_cuda_scatter_src_out( + Tensor* out, + Tensor* self, + int64_t dim, + Tensor* index, + Tensor* src) { + return scatter( + "aoti_torch_cuda_scatter_src_out", + out, + self, + dim, + index, + src, + Reduction::Assign, + /*include_self=*/false); +} + +AOTITorchError aoti_torch_cuda_scatter_reduce_two_out( + Tensor* out, + Tensor* self, + int64_t dim, + Tensor* index, + Tensor* src, + const char* reduce, + int32_t include_self) { + ET_CHECK_OR_RETURN_ERROR( + reduce != nullptr, + InvalidArgument, + "aoti_torch_cuda_scatter_reduce_two_out: null argument"); + const std::string_view name(reduce); + Reduction reduction; + if (name == "sum") { + reduction = Reduction::Sum; + } else if (name == "prod") { + reduction = Reduction::Prod; + } else if (name == "mean") { + reduction = Reduction::Mean; + } else if (name == "amax") { + reduction = Reduction::Max; + } else if (name == "amin") { + reduction = Reduction::Min; + } else { + ET_LOG( + Error, + "aoti_torch_cuda_scatter_reduce_two_out: unknown reduction %s", + reduce); + return Error::InvalidArgument; + } + return scatter( + "aoti_torch_cuda_scatter_reduce_two_out", + out, + self, + dim, + index, + src, + reduction, + include_self != 0); +} + #ifdef __cplusplus } #endif diff --git a/backends/cuda/runtime/shims/index.h b/backends/cuda/runtime/shims/index.h index 3cdfc64e3f9..7a328f57478 100644 --- a/backends/cuda/runtime/shims/index.h +++ b/backends/cuda/runtime/shims/index.h @@ -76,6 +76,69 @@ AOTI_SHIM_EXPORT AOTITorchError aoti_torch_index_put_out( Tensor* values, bool accumulate); +/** + * Writes src into out at index along dim, with the semantics of + * aten::scatter.src. + * + * Inductor calls this fallback, with out and self the same tensor, for a + * tensor source of a dtype its generated code cannot add atomically (int64, + * bool, and bfloat16 on GPUs older than sm_90), and in deterministic mode. + * When values land on the same element, the last one in index order wins, as + * in ATen's deterministic mode; ATen's default CUDA kernel keeps any one of + * them. As in that default kernel, an index out of range, negative ones + * included, triggers a device-side assert. + * + * @param out Tensor to update (CUDA device, up to 16 dims) + * @param self Must be out; AOTInductor passes the tensor it updates as both + * @param dim Dimension of out that index addresses + * @param index int64 or int32 tensor with as many dims as out. In every dim it + * is no larger than src, and in every dim but dim no larger than out + * @param src Values to write, same dtype as out. Float, double, half, + * bfloat16, the integer dtypes and bool are supported. + * @return AOTITorchError + */ +AOTI_SHIM_EXPORT AOTITorchError aoti_torch_cuda_scatter_src_out( + Tensor* out, + Tensor* self, + int64_t dim, + Tensor* index, + Tensor* src); + +/** + * Reduces src into out at index along dim, with the semantics of + * aten::scatter_reduce.two. + * + * Inductor calls this fallback, with out and self the same tensor, for every + * reduction but "sum", for a sum into a dtype its generated code cannot add + * atomically, and in deterministic mode. Values that land on the same element + * are reduced in index order, so the result is deterministic. Half and + * bfloat16 sums add in float, as in ATen's deterministic mode. A product + * rounds to the dtype of out at every step and a mean rounds its sum before + * dividing, as in ATen. As in ATen's default CUDA kernel, an index out of + * range, negative ones included, triggers a device-side assert. + * + * @param out Tensor to update (CUDA device, up to 16 dims) + * @param self Must be out; AOTInductor passes the tensor it updates as both + * @param dim Dimension of out that index addresses + * @param index int64 or int32 tensor with as many dims as out. In every dim it + * is no larger than src, and in every dim but dim no larger than out + * @param src Values to reduce, same dtype as out. Float, double, half, + * bfloat16 and the integer dtypes are supported, and bool for "sum" only, as + * in ATen. + * @param reduce "sum", "prod", "mean", "amax" or "amin" + * @param include_self If zero, the elements index names are reduced from src + * alone; the others keep their values either way + * @return AOTITorchError + */ +AOTI_SHIM_EXPORT AOTITorchError aoti_torch_cuda_scatter_reduce_two_out( + Tensor* out, + Tensor* self, + int64_t dim, + Tensor* index, + Tensor* src, + const char* reduce, + int32_t include_self); + #ifdef __cplusplus } #endif diff --git a/backends/cuda/runtime/shims/memory.cpp b/backends/cuda/runtime/shims/memory.cpp index 86be88e4833..a9a39a7c7b1 100644 --- a/backends/cuda/runtime/shims/memory.cpp +++ b/backends/cuda/runtime/shims/memory.cpp @@ -67,6 +67,25 @@ AOTITorchError narrow_to(From value, To* ret_value) { template AOTITorchError read_one_element(const SlimTensor* tensor, T* ret_value) { +#if defined(CUDA_AVAILABLE) + // A graph capture cannot record the copy to the host below, and the copy + // aborts the process when it fails, so refuse first. Info, not Error: the + // backend then retries the call without a graph. + if (tensor->is_cuda()) { + auto stream = getCurrentCUDAStream(tensor->device_index()); + ET_CHECK_OK_OR_RETURN_ERROR(stream.error()); + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + ET_CUDA_CHECK_OR_RETURN_ERROR( + cudaStreamIsCapturing(stream.get(), &capture)); + if (capture != cudaStreamCaptureStatusNone) { + ET_LOG( + Info, + "reading a single element: a device value cannot be read inside a " + "CUDA graph capture"); + return Error::NotSupported; + } + } +#endif switch (tensor->dtype()) { case ScalarType::Byte: return narrow_to(tensor->item(), ret_value); @@ -80,6 +99,11 @@ AOTITorchError read_one_element(const SlimTensor* tensor, T* ret_value) { return narrow_to(tensor->item(), ret_value); case ScalarType::Float: return narrow_to(tensor->item(), ret_value); + case ScalarType::Double: + return narrow_to(tensor->item(), ret_value); + case ScalarType::Half: + return narrow_to( + static_cast(tensor->item()), ret_value); case ScalarType::BFloat16: return narrow_to( static_cast(tensor->item()), ret_value); @@ -470,6 +494,8 @@ ET_CUDA_DEFINE_ITEM_SHIM(int16, int16_t) ET_CUDA_DEFINE_ITEM_SHIM(int32, int32_t) ET_CUDA_DEFINE_ITEM_SHIM(int64, int64_t) ET_CUDA_DEFINE_ITEM_SHIM(float32, float) +ET_CUDA_DEFINE_ITEM_SHIM(float64, double) +ET_CUDA_DEFINE_ITEM_SHIM(float16, c10::Half) ET_CUDA_DEFINE_ITEM_SHIM(bfloat16, c10::BFloat16) ET_CUDA_DEFINE_ITEM_SHIM(bool, bool) diff --git a/backends/cuda/runtime/shims/memory.h b/backends/cuda/runtime/shims/memory.h index b1b369a58d9..ef51ff5eef0 100644 --- a/backends/cuda/runtime/shims/memory.h +++ b/backends/cuda/runtime/shims/memory.h @@ -256,6 +256,14 @@ aoti_torch_item_int64(SlimTensor* tensor, int64_t* ret_value); AOTI_SHIM_EXPORT AOTITorchError aoti_torch_item_float32(SlimTensor* tensor, float* ret_value); +/// See aoti_torch_item_uint8. +AOTI_SHIM_EXPORT AOTITorchError +aoti_torch_item_float64(SlimTensor* tensor, double* ret_value); + +/// See aoti_torch_item_uint8. +AOTI_SHIM_EXPORT AOTITorchError +aoti_torch_item_float16(SlimTensor* tensor, c10::Half* ret_value); + /// See aoti_torch_item_uint8. AOTI_SHIM_EXPORT AOTITorchError aoti_torch_item_bfloat16(SlimTensor* tensor, c10::BFloat16* ret_value); diff --git a/backends/cuda/runtime/shims/rand.cu b/backends/cuda/runtime/shims/rand.cu index 9916fe8fe0c..85a64e1ffd3 100644 --- a/backends/cuda/runtime/shims/rand.cu +++ b/backends/cuda/runtime/shims/rand.cu @@ -19,9 +19,10 @@ #include #include -#include #include #include +#include +#include #include namespace executorch::backends::cuda { @@ -59,12 +60,20 @@ static std::once_flag g_rng_init_flag; // Initialize RNG state on the given stream. // Must be called during warmup (before graph capture). Subsequent calls -// from any thread are no-ops thanks to std::call_once. A failed allocation -// throws out of the call_once, which leaves it not done, so the next call -// tries again rather than every later launch using a null state. +// from any thread are no-ops thanks to std::call_once. A failure throws out +// of the call_once, which leaves it not done, so the next call tries again +// rather than every later launch using a null state. bool ensure_rng_init(cudaStream_t stream) { try { std::call_once(g_rng_init_flag, [&]() { + RngState h; + // From std::random_device: a seed from the clock repeats for every + // process started in the same second. Drawn before allocating, so a + // throw here leaves nothing allocated. + std::random_device entropy; + h.seed = (static_cast(entropy()) << 32) | entropy(); + h.counter = 0; + h.base_scratch = 0; // The current device: the stream callers pass from getCurrentCUDAStream // may have been created on it whatever index they looked up, and the // cudaMalloc fallback on a device without memory pools follows this @@ -74,19 +83,22 @@ bool ensure_rng_init(cudaStream_t stream) { if (!rng.ok()) { throw std::bad_alloc(); } - d_rng = static_cast(rng.get()); - RngState h; - h.seed = static_cast(time(nullptr)); - h.counter = 0; - h.base_scratch = 0; - cudaMemcpyAsync( - d_rng, &h, sizeof(RngState), cudaMemcpyHostToDevice, stream); // Synchronize to ensure the copy completes before we return // (the host-side RngState `h` is on the stack). - cudaStreamSynchronize(stream); + cudaError_t err = cudaMemcpyAsync( + rng.get(), &h, sizeof(RngState), cudaMemcpyHostToDevice, stream); + if (err == cudaSuccess) { + err = cudaStreamSynchronize(stream); + } + if (err != cudaSuccess) { + (void)CudaAllocator::deallocate_stream_ordered(rng.get(), -1, stream); + throw std::runtime_error(cudaGetErrorString(err)); + } + d_rng = static_cast(rng.get()); }); - } catch (const std::bad_alloc&) { - ET_LOG(Error, "rand: allocating the RNG state failed"); + } catch (const std::exception& e) { + // std::random_device throws an implementation-defined std::exception. + ET_LOG(Error, "rand: initializing the RNG state failed: %s", e.what()); return false; } return true; @@ -233,7 +245,7 @@ AOTITorchError aoti_torch_cuda_rand( ET_CHECK_OR_RETURN_ERROR( ensure_rng_init(stream), MemoryAllocationFailed, - "aoti_torch_cuda_rand: failed to allocate the RNG state"); + "aoti_torch_cuda_rand: failed to set up the RNG state"); constexpr int kThreads = 256; int blocks = static_cast((numel + kThreads - 1) / kThreads); @@ -296,7 +308,7 @@ AOTITorchError aoti_torch_cuda_randint_low_out( ET_CHECK_OR_RETURN_ERROR( ensure_rng_init(stream), MemoryAllocationFailed, - "aoti_torch_cuda_randint_low_out: failed to allocate the RNG state"); + "aoti_torch_cuda_randint_low_out: failed to set up the RNG state"); int64_t range = high - low; int64_t* out_data = static_cast(out->data_ptr()); diff --git a/backends/cuda/runtime/shims/sort.cu b/backends/cuda/runtime/shims/sort.cu index e3bb259bbc7..30d1cf39e08 100644 --- a/backends/cuda/runtime/shims/sort.cu +++ b/backends/cuda/runtime/shims/sort.cu @@ -211,6 +211,9 @@ AOTITorchError aoti_torch_cuda_sort_stable( case c10_slim::ScalarType::Float: elem_size = sizeof(float); break; + case c10_slim::ScalarType::Double: + elem_size = sizeof(double); + break; case c10_slim::ScalarType::BFloat16: elem_size = sizeof(__nv_bfloat16); break; @@ -416,6 +419,17 @@ AOTITorchError aoti_torch_cuda_sort_stable( stream); break; } + case c10_slim::ScalarType::Double: { + sort_slice_impl( + static_cast(values_base) + offset, + idx_ptr, + sort_size, + desc, + is_stable, + device_idx, + stream); + break; + } case c10_slim::ScalarType::BFloat16: { sort_slice_impl( static_cast<__nv_bfloat16*>(values_base) + offset, diff --git a/backends/cuda/runtime/test/test_cuda_kv_cache.cpp b/backends/cuda/runtime/test/test_cuda_kv_cache.cpp index 948a53c0c1b..92ee8406509 100644 --- a/backends/cuda/runtime/test/test_cuda_kv_cache.cpp +++ b/backends/cuda/runtime/test/test_cuda_kv_cache.cpp @@ -932,6 +932,7 @@ TEST_F(CudaKVCacheTest, GrowthDropsACapturedGraphOnlyWhenStorageMoves) { ASSERT_EQ(cudaMalloc(&static_input, 16), cudaSuccess); graph.static_input_ptrs = {static_input}; graph.static_input_nbytes = {16}; + graph.static_input_sizes = {{4}}; ASSERT_EQ(kv.commit_step(3), Error::Ok); // Fits the current storage: the graph keeps replaying. @@ -950,6 +951,7 @@ TEST_F(CudaKVCacheTest, GrowthDropsACapturedGraphOnlyWhenStorageMoves) { EXPECT_EQ(graph.warmup_remaining, 1); EXPECT_TRUE(graph.static_input_ptrs.empty()); EXPECT_TRUE(graph.static_input_nbytes.empty()); + EXPECT_TRUE(graph.static_input_sizes.empty()); EXPECT_EQ(graph.graph_exec, nullptr); // Once recaptured, later steps that fit replay it untouched. diff --git a/backends/cuda/runtime/test/test_cuda_weight_cache.cpp b/backends/cuda/runtime/test/test_cuda_weight_cache.cpp index daa0a41bf0f..a7b86656fc7 100644 --- a/backends/cuda/runtime/test/test_cuda_weight_cache.cpp +++ b/backends/cuda/runtime/test/test_cuda_weight_cache.cpp @@ -275,9 +275,19 @@ TEST(CudaWeightCacheTest, RejectsTruncationAndTrailingData) { Error::InvalidProgram); } +TEST(CudaWeightCacheTest, ParsesDoubleWeight) { + const std::vector bytes = serialized_metadata(7); // Double + cuda::CudaWeightCache::Metadata metadata; + ASSERT_EQ( + cuda::CudaWeightCache::parse(bytes.data(), bytes.size(), metadata), + Error::Ok); + ASSERT_EQ(metadata.entries.size(), 1u); + EXPECT_EQ(metadata.entries[0].dtype, 7); +} + TEST(CudaWeightCacheTest, RejectsUnsupportedDtype) { const std::vector bytes = - serialized_metadata(7); // Double is unsupported. + serialized_metadata(9); // ComplexFloat is unsupported. cuda::CudaWeightCache::Metadata metadata; EXPECT_EQ( cuda::CudaWeightCache::parse(bytes.data(), bytes.size(), metadata), diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index e23127dbb06..90f45591fa0 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -17,10 +17,12 @@ from executorch.backends.aoti.aoti_backend import AotiBackend from executorch.backends.cuda.cuda_backend import CudaBackend from executorch.backends.cuda.cuda_partitioner import CudaPartitioner +from executorch.backends.cuda.passes.remove_assert_async import RemoveAssertAsyncPass from executorch.examples.models.toy_model import SdpaModule from executorch.exir import EdgeCompileConfig, schema, to_edge_transform_and_lower from executorch.exir.backend.compile_spec_schema import CompileSpec from torch.export import export +from torch.nn.attention.varlen import varlen_attn class TestCudaBackendCompileOptions(unittest.TestCase): @@ -211,6 +213,31 @@ def test_small_gpu_override_answers_only_its_own_thread(self): seen["inside"] = inductor_utils.is_big_gpu(0) self.assertEqual({"inside": True, "other": False}, seen) + def test_small_gpu_override_survives_a_thread_that_leaves_first(self): + from executorch.backends.cuda.cuda_backend import ( + _allow_triton_templates_on_small_gpus, + ) + from torch._inductor import utils as inductor_utils + + entered, leave = threading.Event(), threading.Event() + + def overlap(): + with _allow_triton_templates_on_small_gpus(): + entered.set() + leave.wait(timeout=10) + + with mock.patch.object( + inductor_utils, "is_big_gpu", return_value=False + ) as is_big_gpu: + other = threading.Thread(target=overlap) + other.start() + entered.wait(timeout=10) + with _allow_triton_templates_on_small_gpus(): + leave.set() + other.join(timeout=10) + self.assertTrue(inductor_utils.is_big_gpu(0)) + self.assertIs(inductor_utils.is_big_gpu, is_big_gpu) + def test_cuda_include_ptx_compile_spec(self): with mock.patch.object( CudaBackend, "_setup_cuda_environment_for_fatbin", return_value=True @@ -515,6 +542,100 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: CheckModule().eval(), (torch.randn(4, 8),), strict=False ) + def test_distribution_argument_check(self): + """A check that holds for every input must not fail while AOTI autotunes.""" + + class GaussianHead(torch.nn.Module): + def __init__(self): + super().__init__() + self.mean = torch.nn.Linear(16, 14) + self.log_std = torch.nn.Linear(16, 14) + + def forward(self, x: torch.Tensor, eps: torch.Tensor) -> torch.Tensor: + std = torch.clamp(torch.exp(self.log_std(x)), 1e-5, 10.0) + scale_tril = torch.diag_embed(std.sqrt()) + # The check MultivariateNormal makes on scale_tril, as newer + # PyTorch exports it. With 14 actions, Inductor computes each + # half in a kernel of its own, so the assert reads bools that + # autotuning makes up. + lower = (scale_tril.tril() == scale_tril).flatten(-2).min(-1)[0] + positive = (scale_tril.diagonal(dim1=-2, dim2=-1) > 0).min(-1)[0] + torch._assert_async((lower & positive).all(), "invalid scale_tril") + return self.mean(x) + (scale_tril @ eps[..., None])[..., 0] + + inputs = (torch.randn(1, 16), torch.randn(1, 14)) + self._export_to_cuda_with_lower(GaussianHead().eval(), inputs) + + # The export above fails without the pass only when autotuning trips + # the assert, so check that the backend runs it. + self.assertIn( + RemoveAssertAsyncPass, + {type(p) for p in CudaBackend.get_custom_passes([])}, + ) + # The pass drops the check and the math that only fed it. + graph_module = export(GaussianHead().eval(), inputs).graph_module + RemoveAssertAsyncPass()(graph_module) + targets = {n.target for n in graph_module.graph.nodes} + self.assertNotIn(torch.ops.aten._assert_async.msg, targets) + self.assertNotIn(torch.ops.aten.tril.default, targets) + self.assertIn(torch.ops.aten.linear.default, targets) + + def test_reshape_copy(self): + """Torch-TensorRT hands back the views it leaves to PyTorch as _reshape_copy.""" + + class ReshapeCopyModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.ops.aten._reshape_copy(x.transpose(0, 1), [-1, 4]) * 2 + + self._export_to_cuda_with_lower(ReshapeCopyModule(), (torch.randn(8, 4, 4),)) + + def test_varlen_attn(self): + """varlen_attn calls Flash attention, which the CUDA shims lack.""" + + class VarlenAttention(torch.nn.Module): + def forward(self, q, k, v, cu_seq): + # Flattening views the output, which needs the op's layout. + return varlen_attn(q, k, v, cu_seq, cu_seq, 16, 16).flatten(1) + + q, k, v = ( + torch.randn(32, 4, 64, dtype=torch.bfloat16, device="cuda") + for _ in range(3) + ) + cu_seq = torch.tensor([0, 16, 32], dtype=torch.int32, device="cuda") + self._export_to_cuda_with_lower(VarlenAttention(), (q, k, v, cu_seq)) + + # Causal masks over query sequences shorter and longer than their key + # sequences (some queries then see no key, and NaN values outside the + # key sequences must stay out), grouped-query attention with an empty + # sequence, and a window in self-attention. + op = torch.ops.torch_attn._varlen_attn.default + short = torch.tensor([0, 4, 16], dtype=torch.int32, device="cuda") + empty = torch.tensor([0, 16, 16, 32], dtype=torch.int32, device="cuda") + shifted = torch.tensor([2, 6, 16], dtype=torch.int32, device="cuda") + nan_outside = v.clone() + nan_outside[:2] = nan_outside[16:] = float("nan") + big = torch.full_like(q, 1024) + cases = [ + (q[:16], k, v, short, cu_seq, 12, 16, True, None, [3, -1]), + (q, k, nan_outside, cu_seq, shifted, 16, 12, True, None, [-1, 0]), + (q, k[:, :2], v[:, :2], empty, empty, 16, 16, False, 0.5, None, True), + (q, k, v, cu_seq, cu_seq, 16, 16, False, None, [4, 2]), + # Equal scores so large that float32 rounds their log-sum-exp. + (big, big, v, cu_seq, cu_seq, 16, 16), + ] + table = CudaBackend.get_decomposition_table() + # On ROCm, varlen_attn runs AOTriton, which lines an explicit window up + # with the start of each sequence. The decomposition follows Flash + # attention on CUDA, which lines it up with the end, so compare on CUDA. + if torch.version.hip is None: + for args in cases: + for actual, expected in zip(table[op](*args)[:2], op(*args)[:2]): + torch.testing.assert_close(actual, expected, atol=1e-2, rtol=1e-2) + # Flash attention versions differ on this window, since a query sequence + # can be longer than max_k + 1. + with self.assertRaisesRegex(NotImplementedError, "window"): + table[op](q, k[:16], v[:16], cu_seq, short, 16, 12, False, None, [3, -1]) + def test_conv2d(self): """Test CUDA export for 2D convolution.""" @@ -825,63 +946,92 @@ def test_triton_kernel_mode_off(self): ) def test_missing_scatter_kernel_fails_at_export(self): - """Inductor leaves int64 scatter_add to an ATen kernel the CUDA shim lacks.""" + """The scatter fallback writes its kernel call itself, and the check for + missing kernels still sees it. In deterministic mode Inductor leaves a + scatter of a scalar to that fallback, and the CUDA shims lack its + kernel.""" - class Int64ScatterAdd(torch.nn.Module): + class ValueScatter(torch.nn.Module): def forward(self, counts: torch.Tensor, idx: torch.Tensor): - return counts.scatter_add(0, idx, torch.ones_like(idx)) + return counts.scatter(0, idx, 1.0) - inputs = (torch.zeros(8, dtype=torch.long), torch.randint(0, 8, (32,))) - with self.assertRaises(RuntimeError) as context: - self._export_to_cuda_with_lower(Int64ScatterAdd(), inputs) - self.assertIn("aoti_torch_cuda_scatter_reduce_two_out", str(context.exception)) + inputs = (torch.zeros(8), torch.randint(0, 8, (32,))) + deterministic = torch.are_deterministic_algorithms_enabled() + warn_only = torch.is_deterministic_algorithms_warn_only_enabled() + torch.use_deterministic_algorithms(True) + try: + with self.assertRaises(RuntimeError) as context: + self._export_to_cuda_with_lower(ValueScatter(), inputs) + finally: + torch.use_deterministic_algorithms(deterministic, warn_only=warn_only) + self.assertIn("aoti_torch_cuda_scatter_value_out", str(context.exception)) @unittest.skipIf( torch.version.hip is not None, "loads through the Python runtime, which the ROCm build does not include", ) - def test_mask_index_and_index_put(self): - """A mask index makes the result size depend on the data, and Inductor - has no atomic add for int64, so these run in the index shims.""" + def test_fused_sdpa(self): + """Decomposing attention on CUDA, as Torch-TensorRT does, gives fused + variants that the backend has no kernel for.""" from executorch.runtime import Runtime - class MaskIndex(torch.nn.Module): - def forward(self, x, mask, column, idx, counts): - out = x.clone() - out[mask] = x[mask] * 2 - # The mask and column index dims 0 and 2, which are not adjacent. - picked = x[mask, :, column].sum(0) - # A transposed input gives a result that is not row-major. - rows = x.transpose(1, 2)[mask] * 3 - ones = torch.ones_like(idx) - counts = counts.index_put((idx,), ones, accumulate=True) - return out, picked, rows, counts - - model = MaskIndex() - x = torch.randn(4, 5, 6) - # Negative and int32: the shims wrap the index and read it as int32. - column = torch.tensor([-4], dtype=torch.int32) - idx = torch.tensor([0, 3, 3, 1, 3]) - # Nonzero counts check that repeated writes add to the old values. - counts = torch.arange(8) * 5 - masks = ( - torch.tensor([True, False, True, True]), - torch.zeros(4, dtype=torch.bool), - torch.ones(4, dtype=torch.bool), - ) - edge = self._export_to_cuda_with_lower( - model, (x, masks[0], column, idx, counts) + class Attention(torch.nn.Module): + def forward(self, q, k, v, mask): + attention = torch.nn.functional.scaled_dot_product_attention + # Flash attention with fewer key heads and a custom scale. + grouped = attention( + q, + k[:, :2, :16], + v[:, :2, :16], + is_causal=True, + scale=0.2, + enable_gqa=True, + ) + # Efficient attention, with the mask turned into a bias. + masked = attention(q, k[:, :, :16], v[:, :, :16], attn_mask=mask) + # Flash attention of one query over many keys, which runs in + # the Triton decode kernel. + decode = attention(q[:, :, :1], k, v) + return grouped, masked, decode + + model = Attention() + inputs = ( + torch.randn(1, 4, 16, 64, dtype=torch.bfloat16, device="cuda"), + torch.randn(1, 4, 256, 64, dtype=torch.bfloat16, device="cuda"), + torch.randn(1, 4, 256, 64, dtype=torch.bfloat16, device="cuda"), + torch.ones(16, 16, dtype=torch.bool, device="cuda").tril(), ) + program = export(model, inputs, strict=True).run_decompositions() + fused = [ + node + for node in program.graph.nodes + if "aten._scaled_dot_product_" in str(node.target) + ] + self.assertEqual(len(fused), 3) + with self.assertLogs( + "executorch.backends.cuda.triton.replacement_pass", "INFO" + ) as logs: + edge = to_edge_transform_and_lower( + program, + partitioner=[ + CudaPartitioner( + [CudaBackend.generate_method_name_compile_spec("forward")] + ) + ], + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ) + # The masked call keeps a float bias, which the Triton kernels do not take. + output = "\n".join(logs.output) + self.assertIn("Replaced 2 nodes with Triton kernels", output) + self.assertIn("Using split-K decode SDPA", output) with tempfile.TemporaryDirectory() as directory: path = os.path.join(directory, "model.pte") with open(path, "wb") as f: edge.to_executorch().write_to_file(f) method = Runtime.get().load_program(path).load_method("forward") - for mask in masks: - inputs = (x, mask, column, idx, counts) - outputs = method.execute(list(inputs)) - for output, expected in zip(outputs, model(*inputs), strict=True): - torch.testing.assert_close(output.cpu(), expected) + outputs = method.execute([t.cpu() for t in inputs]) + for output, reference in zip(outputs, model(*inputs), strict=True): + torch.testing.assert_close(output, reference.cpu(), atol=2e-2, rtol=2e-2) @unittest.skipIf( torch.version.hip is not None, diff --git a/backends/cuda/tests/test_index_shims.py b/backends/cuda/tests/test_index_shims.py new file mode 100644 index 00000000000..466bbd0fa80 --- /dev/null +++ b/backends/cuda/tests/test_index_shims.py @@ -0,0 +1,369 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Ops that Inductor does not lower by itself in the CUDA backend. + +A mask index makes a result size depend on the data, so Inductor leaves it, +and an index_put with a mask, to the backend's own index shims. Inductor also +leaves them every scatter and scatter_reduce it cannot write with an atomic +add. varlen_attn and _reshape_copy have no Inductor lowering, so the backend +decomposes them. These tests lower real programs, run them through the Python +runtime, and compare with eager PyTorch on CUDA. +""" + +import contextlib +import os +import tempfile +import unittest +from typing import Iterator, Sequence, Tuple +from unittest.mock import patch + +import torch +from executorch.backends.cuda.cuda_backend import CudaBackend +from executorch.backends.cuda.cuda_partitioner import CudaPartitioner +from executorch.backends.cuda.passes.replace_masked_copy_with_where import ( + ReplaceMaskedCopyWithWherePass, +) +from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower +from executorch.exir.dialects._ops import ops as exir_ops +from torch.export import export +from torch.nn.attention.varlen import varlen_attn + + +@contextlib.contextmanager +def _deterministic() -> Iterator[None]: + """Deterministic mode, which sends more scatters and index_puts to the + shims.""" + enabled = torch.are_deterministic_algorithms_enabled() + warn_only = torch.is_deterministic_algorithms_warn_only_enabled() + torch.use_deterministic_algorithms(True) + try: + yield + finally: + torch.use_deterministic_algorithms(enabled, warn_only=warn_only) + + +@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA") +@unittest.skipIf( + torch.version.hip is not None, + "loads through the Python runtime, which the ROCm build does not include", +) +class TestIndexShims(unittest.TestCase): + def _run( + self, + model: torch.nn.Module, + cases: Sequence[Tuple[torch.Tensor, ...]], + **tolerance: float, + ) -> None: + """Lower model for cases[0], then compare every case with eager CUDA.""" + spec = CudaBackend.generate_method_name_compile_spec("forward") + lowered = to_edge_transform_and_lower( + export(model, cases[0]), + partitioner=[CudaPartitioner([spec])], + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ) + # Everything runs in the delegate, the index ops included. + self.assertEqual( + { + n.target.__name__ + for n in lowered.exported_program().graph.nodes + if n.op == "call_function" + }, + {"executorch_call_delegate", "getitem"}, + ) + from executorch.runtime import Runtime + + with tempfile.TemporaryDirectory() as directory: + path = os.path.join(directory, "model.pte") + with open(path, "wb") as f: + lowered.to_executorch().write_to_file(f) + method = Runtime.get().load_program(path).load_method("forward") + for inputs in cases: + outputs = method.execute(list(inputs)) + expected = model(*(t.cuda() for t in inputs)) + for output, reference in zip(outputs, expected, strict=True): + torch.testing.assert_close( + output.cpu(), reference.cpu(), equal_nan=True, **tolerance + ) + + def test_mask_index_and_index_put(self): + """index_put with accumulate goes to the shim too: Inductor has no atomic + add for int64.""" + + class MaskIndex(torch.nn.Module): + def forward(self, x, mask, column, idx, counts): + out = x.clone() + out[mask] = x[mask] * 2 + # The mask and column index dims 0 and 2, which are not adjacent. + picked = x[mask, :, column].sum(0) + # A transposed input gives a result that is not row-major. + rows = x.transpose(1, 2)[mask] * 3 + ones = torch.ones_like(idx) + counts = counts.index_put((idx,), ones, accumulate=True) + wide = x.double() + wide = wide.index_put((mask,), wide[mask], accumulate=True) + return out, picked, rows, counts, wide + + x = torch.randn(4, 5, 6) + # Negative and int32: the shims wrap the index and read it as int32. + column = torch.tensor([-4], dtype=torch.int32) + idx = torch.tensor([0, 3, 3, 1, 3]) + # Nonzero counts check that repeated writes add to the old values. + counts = torch.arange(8) * 5 + masks = ( + torch.tensor([True, False, True, True]), + torch.zeros(4, dtype=torch.bool), + torch.ones(4, dtype=torch.bool), + ) + self._run(MaskIndex(), [(x, mask, column, idx, counts) for mask in masks]) + + def test_two_dim_mask(self): + class TwoDimMask(torch.nn.Module): + def forward(self, x, mask): + return (x[mask] * 2,) + + torch.manual_seed(0) + x = torch.randn(4, 5, 6) + masks = (torch.rand(4, 5) > 0.5, torch.zeros(4, 5, dtype=torch.bool)) + self._run(TwoDimMask(), [(x, mask) for mask in masks]) + + def test_masked_copy_lowers_to_where(self): + """Outputs alone cannot show the rewrite: a masked copy that keeps its + index still gives the right values through the shims.""" + + class Copy(torch.nn.Module): + def forward(self, x, mask): + out = torch.zeros_like(x) + out[:, mask.view(-1)] = x[:, mask.view(-1)] + return (out,) + + rewritten = [] + call = ReplaceMaskedCopyWithWherePass.call + + def record(self, graph_module): + result = call(self, graph_module) + rewritten.extend( + n.target + for n in result.graph_module.graph.nodes + if n.op == "call_function" + ) + return result + + x = torch.randn(2, 5, 3) + masks = ( + torch.tensor([True, False, True, True, False]), + torch.zeros(5, dtype=torch.bool), + torch.ones(5, dtype=torch.bool), + ) + with patch.object(ReplaceMaskedCopyWithWherePass, "call", record): + self._run(Copy(), [(x, mask) for mask in masks]) + self.assertIn(exir_ops.edge.aten.where.self, rewritten) + self.assertNotIn(exir_ops.edge.aten.index.Tensor, rewritten) + + def test_index_put_with_repeated_indices(self): + """Deterministic mode sends half, bfloat16 and int32-index writes to the shim.""" + + class Put(torch.nn.Module): + def forward(self, x, idx, values): + summed = x.index_put((idx,), values, accumulate=True) + # Without accumulate, the last of the repeated writes wins. + written = x.index_put((idx.int(),), values) + return summed, written + + idx = torch.tensor([0, 3, 3, 1, 3]) + with _deterministic(): + for dtype in (torch.float16, torch.bfloat16): + with self.subTest(dtype=dtype): + x = torch.arange(8, dtype=dtype) + values = torch.arange(1, 6, dtype=dtype) + self._run(Put(), [(x, idx, values)]) + + def test_scatter_reduce(self): + """Inductor adds atomically but has no atomic max, min or product, so + those reductions, and the mean, go to the shim. So does a sum it cannot + add atomically: int64, or bfloat16 below sm_90.""" + + class ScatterReduce(torch.nn.Module): + def forward(self, x, index, src): + return tuple( + x.scatter_reduce(1, index, src, reduce, include_self=include) + for reduce in ("sum", "amax", "amin", "prod", "mean") + for include in (True, False) + ) + + # Smaller than src, with repeats, and column 4 is never written. + index = torch.tensor([[0, 2, 2, 5], [1, 1, 1, 0], [5, 3, 0, 5]]) + torch.manual_seed(0) + for dtype in (torch.float32, torch.float16, torch.bfloat16, torch.int64): + with self.subTest(dtype=dtype): + # Small integers keep every sum and product exact in any order. + x = torch.randint(-4, 5, (5, 6)).to(dtype) + src = torch.randint(-4, 5, (4, 5)).to(dtype) + if dtype.is_floating_point: + # Lands on x[1, 1]; every reduction keeps the NaN. + src[1, 2] = float("nan") + self._run(ScatterReduce(), [(x, index, src)]) + + def test_scatter_with_tensor_source(self): + """Inductor writes a scatter from a tensor itself only in a dtype it can + add atomically, so int64, bool and, below sm_90, bfloat16 go to the + shim. Of repeated writes the last in index order wins, as in eager's + deterministic mode, which also sends every dtype to the shim. As in + ATen, the index can be int32.""" + + class Scatter(torch.nn.Module): + def forward(self, x, index, src): + return (x.scatter(1, index, src),) + + # Smaller than src, with repeats, and column 4 is never written. + index = torch.tensor([[0, 2, 2, 5], [1, 1, 1, 0], [5, 3, 0, 5]]) + x = torch.full((5, 6), -1) + # Repeated writes differ from the write before them, also as bool. + src = torch.arange(20).reshape(4, 5) + with _deterministic(): + for dtype, index_dtype in ( + (torch.bfloat16, torch.int64), + (torch.int64, torch.int32), + (torch.bool, torch.int64), + (torch.float32, torch.int64), + (torch.float16, torch.int64), + ): + with self.subTest(dtype=dtype, index_dtype=index_dtype): + values = src % 2 if dtype == torch.bool else src + inputs = (x.to(dtype), index.to(index_dtype), values.to(dtype)) + self._run(Scatter(), [inputs]) + + def test_scatter_strides(self): + """Scatters along the middle dim of a 3-D tensor, given as a negative + dim, from a strided source and from transposed operands, so the shims + walk strides that differ between their tensors.""" + + class Strided(torch.nn.Module): + def forward(self, x, index, src): + picked = src[:, ::2] + return ( + x.scatter(-2, index, picked), + x.scatter_reduce(-2, index, picked, "amax"), + x.transpose(0, 2).scatter( + 1, index.transpose(0, 2), picked.transpose(0, 2) + ), + ) + + torch.manual_seed(0) + x = torch.randn(4, 6, 5) + index = torch.randint(0, 6, (3, 4, 5)) + src = torch.randn(3, 8, 5) + with _deterministic(): + self._run(Strided(), [(x, index, src)]) + + def test_scatter_mean_counts_in_its_dtype(self): + """ATen counts a mean's terms in the tensor's own dtype, so 128 int8 + terms count as -128.""" + + class Mean(torch.nn.Module): + def forward(self, x, index, src): + return tuple( + x.scatter_reduce(0, index, src, "mean", include_self=include) + for include in (True, False) + ) + + x = torch.zeros(2, dtype=torch.int8) + # With self, the second element counts 256, which wraps to 0. + index = torch.tensor([0] * 128 + [1] * 255) + self._run(Mean(), [(x, index, torch.ones(383, dtype=torch.int8))]) + + def test_scatter_reduce_edge_cases(self): + """A bool sum, products that overflow int32 and half, and an empty + index.""" + + class EdgeCases(torch.nn.Module): + def forward(self, flags, ints, index, empty): + halves = ints.half() + return ( + flags.scatter_reduce(0, index, flags, "sum"), + # As in ATen, int32 wraps and half overflows to infinity. + ints.scatter_reduce(0, index, ints, "prod"), + halves.scatter_reduce(0, index, halves, "prod"), + halves.scatter_reduce(0, empty, halves[:0], "mean"), + ) + + flags = torch.tensor([True, False, False, True, False]) + ints = torch.tensor([300, -7, 46341, 300, 5], dtype=torch.int32) + index = torch.tensor([0, 2, 2, 0, 1]) + empty = torch.zeros(0, dtype=torch.long) + self._run(EdgeCases(), [(flags, ints, index, empty)]) + + def test_scatter_reduce_sum_adds_in_float(self): + """Deterministic mode sends half and bfloat16 sums to the shim, which + adds in float, as ATen does there, so many small terms add up.""" + + class Sum(torch.nn.Module): + def forward(self, index, src): + return tuple( + torch.zeros(3, dtype=dtype, device=src.device).scatter_reduce( + 0, index, src.to(dtype), "sum" + ) + for dtype in (torch.float16, torch.bfloat16) + ) + + with _deterministic(): + self._run(Sum(), [(torch.arange(1000) % 3, torch.full((1000,), 0.1))]) + + def test_float64_and_float16_values(self): + class Values(torch.nn.Module): + def forward(self, x, scale): + # Rows this long go to the sort shim, not to Triton. + values, indices = torch.sort(x, stable=True) + return values * scale.item(), indices + + torch.manual_seed(0) + for dtype in (torch.float64, torch.float16): + with self.subTest(dtype=dtype): + x = torch.randn(4, 1000, dtype=dtype) + # A fraction, which a float passed to Triton as an integer + # loses. + scale = torch.tensor(1.5, dtype=dtype) + self._run(Values(), [(x, scale)]) + + def test_reshape_copy(self): + """Torch-TensorRT hands back the views it leaves to PyTorch as + _reshape_copy.""" + + class ReshapeCopy(torch.nn.Module): + def forward(self, x): + reshape = torch.ops.aten._reshape_copy + return ( + reshape(x.transpose(0, 1), [-1, 4]), + reshape(x[:, 1:3], [8, -1]), + reshape(x, [4, 32]), + ) + + self._run(ReshapeCopy(), [(torch.randn(8, 4, 4),)]) + + def test_varlen_attn(self): + """varlen_attn calls Flash attention, which the CUDA shims lack.""" + + class Attention(torch.nn.Module): + def forward(self, q, k, v, cu_seq): + # Flattening views each output, which needs the op's layout. + return tuple( + varlen_attn( + q, k, v, cu_seq, cu_seq, 16, 16, window_size=window + ).flatten(1) + for window in ((-1, -1), (-1, 0), (8, 0), (4, 2)) + ) + + torch.manual_seed(0) + q, k, v = (torch.randn(28, 4, 64, dtype=torch.bfloat16) for _ in range(3)) + # Equal scores so large that float32 rounds their log-sum-exp. + big = torch.full_like(q, 1024) + # The sequence boundaries come only at run time. + cases = [ + (q, k, v, torch.tensor([0, 12, 28], dtype=torch.int32)), + (q, k, v, torch.tensor([0, 16, 28], dtype=torch.int32)), + (big, big, v, torch.tensor([0, 12, 28], dtype=torch.int32)), + ] + self._run(Attention(), cases, atol=1e-2, rtol=1e-2) diff --git a/backends/cuda/tests/test_missing_weights_blob.py b/backends/cuda/tests/test_missing_weights_blob.py index 6f1308b20c3..0a524bd7353 100644 --- a/backends/cuda/tests/test_missing_weights_blob.py +++ b/backends/cuda/tests/test_missing_weights_blob.py @@ -23,6 +23,7 @@ import tempfile import unittest from typing import Optional +from unittest.mock import patch import torch from executorch.backends.cuda.cuda_backend import CudaBackend @@ -167,10 +168,16 @@ def test_load_reports_not_found_when_blob_is_absent(self) -> None: # exception carries only the method name and the error number, so the # cause is asserted through the rewrite check above rather than here. program = Runtime.get().load_program(path) - with self.assertRaisesRegex( - RuntimeError, r"Failed to load method forward, error: 0x:?20" + # The failed load leaves no handle to destroy, so it has to remove the + # library it wrote itself. An empty temporary directory shows that. + with tempfile.TemporaryDirectory() as tmp, patch.dict( + os.environ, {"TMPDIR": tmp} ): - program.load_method("forward") + with self.assertRaisesRegex( + RuntimeError, r"Failed to load method forward, error: 0x:?20" + ): + program.load_method("forward") + self.assertEqual(os.listdir(tmp), []) def test_load_fails_with_another_programs_blob(self) -> None: from executorch.runtime import Runtime @@ -222,6 +229,36 @@ def test_programs_with_their_own_data_names_share_a_folder(self) -> None: module(value).cpu(), ) + def test_float64_weight_loads(self) -> None: + """The weight cache must accept a float64 weight the export wrote.""" + + class Float64Scale(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer( + "scale", torch.linspace(1.0, 2.0, 64, dtype=torch.float64) + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return (x.to(torch.float64) * self.scale).to(x.dtype) + + from executorch.runtime import Runtime + + with tempfile.TemporaryDirectory() as outdir: + module = Float64Scale() + path = self._lower(outdir, module=module) + (blob,) = [f for f in os.listdir(outdir) if f.endswith(".ptd")] + method = ( + Runtime.get() + .load_program(path, data_path=os.path.join(outdir, blob)) + .load_method("forward") + ) + # The same draw _lower made for its example input. + torch.manual_seed(0) + value = torch.randn(4, 64, device="cuda") + out = method.execute([value.cpu()]) + torch.testing.assert_close(out[0].cpu(), module(value).cpu()) + def test_load_succeeds_when_the_model_has_no_constants(self) -> None: """A model with nothing to bind still loads without its weights blob. diff --git a/backends/cuda/tests/test_output_stride_rearrange.py b/backends/cuda/tests/test_output_stride_rearrange.py index 2cbde6cc6f8..a53156cd3f0 100644 --- a/backends/cuda/tests/test_output_stride_rearrange.py +++ b/backends/cuda/tests/test_output_stride_rearrange.py @@ -18,6 +18,7 @@ python -m pytest backends/cuda/tests/test_output_stride_rearrange.py -v """ +import contextlib import glob import os import subprocess @@ -292,6 +293,29 @@ def forward(self, x, mask): return x.index_put((mask,), torch.ones(4), accumulate=True) +class _ScaledByCount(nn.Module): + def forward(self, x, mask): + return x * mask.sum().item() + + +class _ScatterMax(nn.Module): + def forward(self, x, index, src): + return x.scatter_reduce(1, index, src, "amax") + + +class _MaskedCounter(nn.Module): + """Counts its calls in a device buffer, and reads a mask's count. The + buffer has one element, the case torch 2.14 can fold into a constant.""" + + def __init__(self): + super().__init__() + self.register_buffer("calls", torch.zeros(1), persistent=False) + + def forward(self, x, mask): + self.calls.add_(1) + return self.calls.clone(), x[mask] + + class _HostFlag(nn.Module): """A torch.cond on a flag the program keeps in host memory and updates on every call.""" @@ -315,6 +339,29 @@ def forward(self, x): return (x * 2).contiguous(memory_format=self.memory_format) +class _RuntimeLog: + """Collects the runtime's log lines. They go to sys.stderr while the + caller holds the GIL, and straight to file descriptor 2 otherwise.""" + + def __enter__(self): + sys.stderr.flush() + self._file = tempfile.TemporaryFile(mode="w+") + self._saved = os.dup(2) + os.dup2(self._file.fileno(), 2) + self._redirect = contextlib.redirect_stderr(self._file) + self._redirect.__enter__() + return self + + def __exit__(self, *exc): + self._redirect.__exit__(*exc) + self._file.flush() + os.dup2(self._saved, 2) + os.close(self._saved) + self._file.seek(0) + self.text = self._file.read() + self._file.close() + + class TestStridedOutputWithCudaGraph(unittest.TestCase): """A CUDA graph replay copies raw bytes, so a strided output must not reach it.""" @@ -413,7 +460,7 @@ def test_input_shape_changes_after_capture(self): self._assert_matches_eager(module, x, out) else: # InvalidArgument, from the replay's input shape check. - self.assertIn("error 0x12", str(out)) + self.assertIn("execute() failed with error 0x12", str(out)) def test_host_state_runs_without_a_graph(self): # A graph would keep the branch of the capture call, and a failed @@ -422,7 +469,6 @@ def test_host_state_runs_without_a_graph(self): torch.manual_seed(0) module = _HostFlag().cuda().eval() eager = _HostFlag().cuda().eval() - eager.load_state_dict(module.state_dict()) inputs = [torch.randn(1, rows, 64) for rows in (1, 1, 1, 50, 50, 50)] rows = torch.export.Dim("rows", min=1, max=64) outputs = self._run_with_cuda_graph( @@ -434,22 +480,63 @@ def test_host_state_runs_without_a_graph(self): expected = eager(x.cuda())[0].cpu() torch.testing.assert_close(out, expected) - def test_mask_runs_without_a_graph(self): - # A mask index reads its count back to the host, which a capture - # cannot record, so the capture call has to run without a graph. + def test_host_read_runs_without_a_graph(self): + # A mask index reads its count back to the host, and .item() a value, + # which a capture cannot record, so the capture call has to run + # without a graph. torch.manual_seed(0) inputs = [(torch.randn(8, 4), torch.rand(8) > 0.5) for _ in range(6)] - for module in (_MaskedRows().eval(), _MaskedAdd().eval()): - outputs = self._run_with_cuda_graph(module, inputs[0], inputs) + for module in ( + _MaskedRows().eval(), + _MaskedAdd().eval(), + _ScaledByCount().eval(), + ): + with _RuntimeLog() as log: + outputs = self._run_with_cuda_graph(module, inputs[0], inputs) + # Eager calls give the same values, so only the log shows the + # capture was tried, and only once. + self.assertEqual(log.text.count("capture failed for 'forward'"), 1) for i, (args, out) in enumerate(zip(inputs, outputs)): with self.subTest(module=type(module).__name__, call=i): self.assertIsInstance(out, torch.Tensor) torch.testing.assert_close(out, module(*args)) + def test_scatter_reduce_replays_in_a_graph(self): + # The shim sorts its writes in scratch memory, which the capture has to + # record as well. + torch.manual_seed(0) + index = torch.tensor([[0, 2, 2, 5], [1, 1, 1, 0], [5, 3, 0, 5]]) + inputs = [(torch.randn(5, 6), index, torch.randn(4, 5)) for _ in range(6)] + module = _ScatterMax().eval() + with _RuntimeLog() as log: + outputs = self._run_with_cuda_graph(module, inputs[0], inputs) + self.assertIn("captured and instantiated for 'forward'", log.text) + for i, (args, out) in enumerate(zip(inputs, outputs)): + with self.subTest(call=i): + self.assertIsInstance(out, torch.Tensor) + torch.testing.assert_close(out, module(*args)) + + def test_mask_retry_updates_state_once(self): + # The capture call runs again without a graph, so its update to the + # count must not land twice. + torch.manual_seed(0) + inputs = [(torch.randn(8, 4), torch.rand(8) > 0.5) for _ in range(6)] + with _RuntimeLog() as log: + outputs = self._run_with_cuda_graph( + _MaskedCounter().eval(), inputs[0], inputs + ) + self.assertEqual(log.text.count("capture failed for 'forward'"), 1) + for i, out in enumerate(outputs): + with self.subTest(call=i): + self.assertIsInstance(out, torch.Tensor) + torch.testing.assert_close(out, torch.tensor([i + 1.0])) + def test_output_shape_changes_at_capture(self): torch.manual_seed(0) # The capture call has a larger batch than the warmup calls before it. - inputs = [torch.randn(batch, 3, 4, 4) for batch in (1, 1, 1, 2, 2, 2)] + # The last call goes back to the warmup batch, which only a live graph + # rejects, because a replay needs the captured input sizes. + inputs = [torch.randn(batch, 3, 4, 4) for batch in (1, 1, 1, 2, 2, 2, 1)] batch = torch.export.Dim("batch", min=1, max=8) for memory_format in (torch.contiguous_format, torch.channels_last): module = _Double(memory_format).eval() @@ -458,7 +545,11 @@ def test_output_shape_changes_at_capture(self): ) for i, (x, out) in enumerate(zip(inputs, outputs)): with self.subTest(memory_format=memory_format, call=i): - self._assert_matches_eager(module, x, out) + if i == 6: + # InvalidArgument, from the replay's input size check. + self.assertIn("execute() failed with error 0x12", str(out)) + else: + self._assert_matches_eager(module, x, out) if __name__ == "__main__": diff --git a/backends/cuda/tests/test_sort_shim.py b/backends/cuda/tests/test_sort_shim.py index d81d9e9e187..4241405b9a0 100644 --- a/backends/cuda/tests/test_sort_shim.py +++ b/backends/cuda/tests/test_sort_shim.py @@ -43,6 +43,8 @@ "aoti_torch_cuda_randint_low_out", "aoti_torch_cuda_index_Tensor", "aoti_torch_index_put_out", + "aoti_torch_cuda_scatter_reduce_two_out", + "aoti_torch_cuda_scatter_src_out", } ) diff --git a/backends/cuda/tests/test_triton_sdpa.py b/backends/cuda/tests/test_triton_sdpa.py index 8b878af8a4e..54aac766409 100644 --- a/backends/cuda/tests/test_triton_sdpa.py +++ b/backends/cuda/tests/test_triton_sdpa.py @@ -83,10 +83,8 @@ def _max_abs_error(out, ref): HEAD_DIMS_NON_POW2 = [80, 96] # Sequence length pairs (seqlen_q, seqlen_kv) adapted from FlashAttention -# Note: Lk must be >= BLOCK_N (32) to avoid unmasked zero-padding in K -# that dilutes softmax. This is a pre-existing kernel limitation for very -# short KV sequences. SEQLEN_PAIRS_BASIC = [ + (1, 1), (1, 64), (1, 128), (1, 512), @@ -357,6 +355,38 @@ def test_every_tile_packed_or_not_matches(self): with self.subTest(Lq=Lq, tile=tile): self.assertTrue(torch.equal(unpacked, packed)) + def test_every_tile_drops_keys_past_the_end(self): + """Each standard-kernel tile without a mask: keys at or past the key + count, or past kv_len, stay out of the softmax.""" + import triton + + sdpa_module = importlib.import_module( + "executorch.backends.cuda.triton.kernels.sdpa" + ) + kernel = sdpa_module._sdpa_fwd_kernel + B, H, Lq, D = 1, 4, 16, 64 + configs = sdpa_module._sdpa_fwd_prune( + kernel.configs, None, HEAD_DIM=D, NUM_GROUPS=1 + ) + # No tile size divides 80 keys, only the 32-key tiles divide 96, and + # kv_len is only known on the device. + for Lk, kv_len in ((80, None), (96, None), (128, 80)): + torch.manual_seed(Lk) + q = torch.randn(B, H, Lq, D, dtype=torch.bfloat16, device="cuda") + k = torch.randn(B, H, Lk, D, dtype=torch.bfloat16, device="cuda") + v = torch.randn(B, H, Lk, D, dtype=torch.bfloat16, device="cuda") + valid = Lk if kv_len is None else kv_len + ref = _reference_sdpa(q, k[:, :, :valid], v[:, :, :valid]) + kwargs = {} + if kv_len is not None: + kwargs["kv_len"] = torch.tensor([kv_len], device="cuda") + for config in configs: + single = triton.autotune(configs=[config], key=kernel.keys)(kernel.fn) + with self.subTest(Lk=Lk, kv_len=kv_len, config=str(config)): + with mock.patch.object(sdpa_module, "_sdpa_fwd_kernel", single): + out = self.sdpa(q, k, v, **kwargs) + self.assertLess(_max_abs_error(out, ref), MAX_ABS_TOL) + def test_mha_never_packs(self): sdpa_module = importlib.import_module( "executorch.backends.cuda.triton.kernels.sdpa" diff --git a/backends/cuda/triton/kernels/sdpa.py b/backends/cuda/triton/kernels/sdpa.py index 3dc2b45695e..cbc16fc85f3 100644 --- a/backends/cuda/triton/kernels/sdpa.py +++ b/backends/cuda/triton/kernels/sdpa.py @@ -514,6 +514,7 @@ def _sdpa_fwd_kernel_body( # noqa: C901 NUM_GROUPS: tl.constexpr, PACK_GQA: tl.constexpr, MASK_IS_CAUSAL: tl.constexpr = False, + KV_LEN_MULTIPLE: tl.constexpr = 1, ): """ Shared kernel body for SDPA forward pass. @@ -526,6 +527,9 @@ def _sdpa_fwd_kernel_body( # noqa: C901 When False, the grid iterates over H_q heads and each program handles one Q head with simple h_kv = h_q // NUM_GROUPS remapping. + + KV_LEN_MULTIPLE: kv_len is a multiple of this. When BLOCK_N divides it, + no tile reaches past kv_len, so the kernel skips the key bound check. """ pid_m = tl.program_id(axis=0) pid_bh = tl.program_id(axis=1) @@ -632,6 +636,17 @@ def _sdpa_fwd_kernel_body( # noqa: C901 k = tl.load(k_ptrs, mask=k_mask, other=0.0).to(tl.bfloat16) qk = (tl.dot(q, tl.trans(k)).to(tl.float32) * sm_scale).to(tl.float32) + # Keys past kv_len load as zeros, which would score 0 rather than -inf. + # A mask or a causal bound already drops them. + if ( + not (HAS_MASK or IS_CAUSAL or MASK_IS_CAUSAL) + and KV_LEN_MULTIPLE % BLOCK_N != 0 + ): + qk = tl.where( + offs_n[None, :] < kv_len, + qk, + tl.full(qk.shape, -float("inf"), dtype=tl.float32), + ) if HAS_MASK: qk = tl.where( @@ -753,6 +768,7 @@ def _sdpa_prefill_prune(configs, nargs, **kwargs): for config in _SDPA_PREFILL_CONFIGS for pack in (False, True) ] +_SDPA_BLOCK_N_SIZES = {config.kwargs["BLOCK_N"] for config in _SDPA_FWD_CONFIGS} def _sdpa_fwd_prune(configs, nargs, **kwargs): @@ -820,6 +836,7 @@ def _sdpa_fwd_kernel( BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, MASK_IS_CAUSAL: tl.constexpr = False, + KV_LEN_MULTIPLE: tl.constexpr = 1, ): # Packed, the grid walks the KV heads; unpacked, the query heads. if PACK_GQA: @@ -866,6 +883,7 @@ def _sdpa_fwd_kernel( NUM_GROUPS=NUM_GROUPS, PACK_GQA=PACK_GQA, MASK_IS_CAUSAL=MASK_IS_CAUSAL, + KV_LEN_MULTIPLE=KV_LEN_MULTIPLE, ) @@ -1115,6 +1133,15 @@ def grid(meta): return (triton.cdiv(L_q * num_groups, meta["BLOCK_M"]), B * H_kv) return (triton.cdiv(L_q, meta["BLOCK_M"]), B * H_q) + # The largest tile size that divides L_kv. A kv_len or a symbolic length is + # only known at run time, so those always check bounds. A mask or a causal + # bound makes the check unneeded, so those calls keep 1 and share kernels. + kv_len_multiple = 1 + if isinstance(L_kv, int) and not (HAS_KV_LEN or HAS_MASK or is_causal): + kv_len_multiple = max( + (n for n in _SDPA_BLOCK_N_SIZES if L_kv % n == 0), default=1 + ) + # Single autotuned kernel: the config set spans BLOCK_M in {16..128}, each # with and without pack GQA, and `_sdpa_prefill_prune` keeps only # non-spilling tiles for this HEAD_DIM, so the autotuner picks a @@ -1157,6 +1184,7 @@ def grid(meta): HEAD_DIM=D, NUM_GROUPS=num_groups, MASK_IS_CAUSAL=mask_is_causal, + KV_LEN_MULTIPLE=kv_len_multiple, ) diff --git a/extension/cuda/runtime_api.h b/extension/cuda/runtime_api.h index 91a21e53689..ffb53e97db5 100644 --- a/extension/cuda/runtime_api.h +++ b/extension/cuda/runtime_api.h @@ -21,6 +21,7 @@ using cudaGraphExec_t = hipGraphExec_t; using cudaMemcpyKind = hipMemcpyKind; using cudaMemoryType = hipMemoryType; using cudaStreamCaptureMode = hipStreamCaptureMode; +using cudaStreamCaptureStatus = hipStreamCaptureStatus; using cudaStream_t = hipStream_t; inline constexpr cudaError_t cudaSuccess = hipSuccess; @@ -34,6 +35,8 @@ inline constexpr cudaMemoryType cudaMemoryTypeDevice = hipMemoryTypeDevice; inline constexpr cudaMemoryType cudaMemoryTypeManaged = hipMemoryTypeManaged; inline constexpr cudaStreamCaptureMode cudaStreamCaptureModeRelaxed = hipStreamCaptureModeRelaxed; +inline constexpr cudaStreamCaptureStatus cudaStreamCaptureStatusNone = + hipStreamCaptureStatusNone; inline constexpr unsigned long long cudaGraphInstantiateFlagAutoFreeOnLaunch = hipGraphInstantiateFlagAutoFreeOnLaunch; @@ -211,6 +214,12 @@ inline cudaError_t cudaStreamEndCapture( return hipStreamEndCapture(stream, graph); } +inline cudaError_t cudaStreamIsCapturing( + cudaStream_t stream, + cudaStreamCaptureStatus* status) { + return hipStreamIsCapturing(stream, status); +} + inline cudaError_t cudaStreamSynchronize(cudaStream_t stream) { return hipStreamSynchronize(stream); } diff --git a/extension/pybindings/README.md b/extension/pybindings/README.md index c665c7657c2..352c7a6f93a 100644 --- a/extension/pybindings/README.md +++ b/extension/pybindings/README.md @@ -155,8 +155,8 @@ only host inputs and outputs waits depends on where it runs: In both cases, the CUDA delegate copies an output whose memory layout has to change through the legacy default stream. Loading a method's CUDA weights, which a first call may do, copies each one with a plain `cudaMemcpy`, so it waits for -the legacy default stream and for blocking streams; weights an older export -keeps in one blob wait for the whole device. +the legacy default stream and for blocking streams. An older export keeps its +weights in one blob, and copying that blob waits for the whole device. The waits below, for CUDA inputs and outputs, are for the whole device (`cudaDeviceSynchronize`), not for one stream, so they also wait for unrelated @@ -201,10 +201,7 @@ The CUDA build workflow runs `test_pybindings_cuda_dlpack.py` in both output modes on a GPU. It checks device allocations, DLPack stream ordering, channels last views, cloning, and borrowed output lifetimes using a CUDA delegate. Set `EXECUTORCH_PYBIND_TEST_REQUIRE_CUDA=1` when running this test to treat missing -GPU hardware or a missing CUDA backend as a failure. The CUDA delegate's kernels -need the OpenMP runtime (`libgomp.so.1` on Linux). PyTorch ships one; without PyTorch, -install it from the system, or the two tests that run without PyTorch fail -under this setting. +GPU hardware or a missing CUDA backend as a failure. # Link Backends diff --git a/extension/pybindings/test/test_pybindings_cuda_dlpack.py b/extension/pybindings/test/test_pybindings_cuda_dlpack.py index 0792e6c924f..b7422ada832 100644 --- a/extension/pybindings/test/test_pybindings_cuda_dlpack.py +++ b/extension/pybindings/test/test_pybindings_cuda_dlpack.py @@ -562,7 +562,7 @@ def exit_elsewhere(): self._close_after_teardown(stream) def test_torch_free_caller_shares_one_stream_across_programs(self): - self._require_libgomp_without_torch() + self._require_linux() programs = [ self._export_to_directory( f"torch_free_{index}", module, (torch.ones(4, 8),) @@ -615,9 +615,9 @@ def load(path, data): np.testing.assert_allclose(outputs[0], np.maximum(value, 0), atol=1e-6) np.testing.assert_allclose(outputs[1], np.tanh(value), atol=1e-6) assert "torch" not in sys.modules -assert re.search( - r"lib(?:torch|aten|c10)", Path("/proc/self/maps").read_text().lower() -) is None +# libgomp as well, since these programs never call OpenMP. +maps = Path("/proc/self/maps").read_text().lower() +assert re.search(r"/lib(?:torch|aten|c10|gomp)[^/]*$", maps, re.M) is None """ arguments = [item or "" for program in programs for item in program] result = subprocess.run( @@ -1264,27 +1264,16 @@ def chain(): torch.from_dlpack(final), torch.full(self.shape, expected) ) - def _require_libgomp_without_torch(self): + def _require_linux(self): # The torch-free chain checks loaded libraries through /proc. if not sys.platform.startswith("linux"): self.skipTest("needs Linux") - # The CUDA delegate's compiled kernels link libgomp, which a host may - # provide only inside the torch package. - probe = subprocess.run( - [sys.executable, "-c", "import ctypes; ctypes.CDLL('libgomp.so.1')"], - capture_output=True, - ) - if probe.returncode != 0: - message = "libgomp.so.1 is not loadable without torch" - if os.environ.get("EXECUTORCH_PYBIND_TEST_REQUIRE_CUDA") == "1": - raise AssertionError(message) - self.skipTest(message) def test_device_chain_runs_with_torch_blocked(self): # The device-resident chain path, in a process where `import torch` # fails: upload host data, run, feed the output straight into a second # run on the GPU, then read the result back. - self._require_libgomp_without_torch() + self._require_linux() script = """ import importlib.abc import re @@ -1318,9 +1307,9 @@ def find_spec(self, fullname, path=None, target=None): np.asarray(second.to_host()), np.full(shape, 18.0, dtype=np.float32) ) assert "torch" not in sys.modules -assert re.search( - r"lib(?:torch|aten|c10)", Path("/proc/self/maps").read_text().lower() -) is None +# libgomp as well, since these programs never call OpenMP. +maps = Path("/proc/self/maps").read_text().lower() +assert re.search(r"/lib(?:torch|aten|c10|gomp)[^/]*$", maps, re.M) is None print("Device chain ran without torch") """ environment = os.environ.copy() @@ -1349,7 +1338,6 @@ def find_spec(self, fullname, path=None, target=None): def test_device_input_works_beside_a_torch_without_cuda(self): # A torch built without CUDA cannot select a CUDA device. Inputs from # other producers must not depend on it. - self._require_libgomp_without_torch() script = """ import sys import types From d8823260cf74845ce05fa3fbee474675a8ed4888 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 3 Nov 2026 09:42:13 -0800 Subject: [PATCH 2/2] Keep release-only dtype_element_size export and drop duplicate test after pick --- backends/cuda/runtime/aoti_cuda_shims.lib | Bin 59238 -> 60078 bytes backends/cuda/tests/test_cuda_export.py | 68 ---------------------- 2 files changed, 68 deletions(-) diff --git a/backends/cuda/runtime/aoti_cuda_shims.lib b/backends/cuda/runtime/aoti_cuda_shims.lib index ad54a4fdb24e1e8c3395f0245d36a2d9881b5ef3..9801a1d05fe47dae35b33ce08fce088fa71db8dd 100644 GIT binary patch delta 1460 zcmaLVe@xX?6aesZv7S%i;dqZD;U{r605}2W^#jbWQl2aTke*tpgggx(EgV;#sRod&JC*R|0Mhe8 zip?89dcmMXP60?!ca`W7J3#aZKy176>Hz@hwOnP*8vxRJo3eqErPpsM@izgagbPYy zBY?E2SlOHhAf+ZNX)gju+x(Rb`)>dlEb9v;XFq_Hw^{jsHS-ffq>^&B`EOhjfV-NK%z8UBH7^ARv*%Le><;5B+oQ*a zPKytDTUHqjnYyk==(x_}Pf7CvHqwKpIDrG%mj71p3O$2S3#@2&82)#su&h7|%Bd1` zJKHAc?tA>LGdDtD_jaAu8&D;Jn^ijI|a8qD=?~6Cs#aYceVN9 zKODVJeD2V8&Ti%G;~gfP{L+^u+F}Jh=$sy0{%KF@>{u?a=ZuvSig;b@IZxVidMOY0 z>qo=wvf%G?jFfba2b|D7-St?H8H2hXCu1tLB9LY0X}o^rRkeKJmf?7e&WcT%T{vL^EuLn&D1^=Ca~p- z{;)T8)?Vd@lX6WX_i%nDF@TZgKnfXqNnqf( nmG)obiZwS(bnn-N0x#dx=#K>(gjhkfH@MdfTuhHmB)k0!8HfV> delta 1335 zcmXZadr*{B6aesZVWMq^Z;{;$Lc>>M>EQSAW`$9)6(=UfDD6dqF;@0)z;Ufe2HA3S zP|$_p;0CP^Ad)r-krT0lkt>d?=36&PG&UJR|EQ*!nyJar>0Iu9GrxcCJ?EZt&;5FS zh@QL{-MHA5laW1p>MJt?fE!@OdVu6~Wlkc1H18K>egHsPP^F|~0!YSeC2bTy%JeHa z-2hVV0cD8?K+3OE98MR2!wImmSSd^ekcy@$tNsL#p1Z89;biHh8l|KfKq`GhDa!|t zUQbrmTLGlE|5G-E0HlguWs_?=z$W%yky6Qt(iSc%RSzq!MgVEsT4hHHfK>apQpa)9 zo@S+yo070X@hk<9T9Xtn=Si*u-vCJM94vJ{rF6vrNWSw*cQ=4^WT(>iB7oG-hf4k@ zl|jC#G{ntF0d7b-S)v3N0!U{#hCKu4zK1Cnn*gLM94Lk6DkFCQr0c!PO`et%UZ;#F z0Z0=e3^U7(~t55HZ$2jdFTGNs#(C;zmUMa8UwpeMV zw@~1&R)d<~;qqp0BHiD=T;QAg_#y{6W^%ukuD0m{mmcKO9lR1~vr_YC8G=MdlfY-Y zV(7I)ynnh=r(<2R)}eTs!KnZSgv z47y*+)mwTV#fH8I$a8dwz<-WJM)~1U8`kv4Q2EyxoY4rf9=FiG zVTVBNM5J@~u#GmK;=6=TMkZNy%8G5%Oq72nQ&4BHm0$RrUC{ey`Eh@~VCU1$>r)jl zc)^Nm<724jvM%W1OZ*#bzLw5GS3H99M)=InLk6XGaNqMsrc?BFp5cZc)6|dg$4^!Y z-r#;SMs+IO$o)EK>llaLy{Xf^Fi)_0Orz3U#R4ydb-FRmE5)}o%9~gquzNg}N}k|g UeL|zZCR6w|e_cdHch*Jy4<%gEIRF3v diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 90f45591fa0..7d1a0e91fea 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -408,74 +408,6 @@ def _export_to_cuda_with_lower( return edge_program_manager - @unittest.skipIf( - torch.version.hip is not None, - "loads through the Python runtime, which the ROCm build does not include", - ) - def test_fused_sdpa(self): - """Decomposing attention on CUDA, as Torch-TensorRT does, gives fused - variants that the backend has no kernel for.""" - from executorch.runtime import Runtime - - class Attention(torch.nn.Module): - def forward(self, q, k, v, mask): - attention = torch.nn.functional.scaled_dot_product_attention - # Flash attention with fewer key heads and a custom scale. - grouped = attention( - q, - k[:, :2, :16], - v[:, :2, :16], - is_causal=True, - scale=0.2, - enable_gqa=True, - ) - # Efficient attention, with the mask turned into a bias. - masked = attention(q, k[:, :, :16], v[:, :, :16], attn_mask=mask) - # Flash attention of one query over many keys, which runs in - # the Triton decode kernel. - decode = attention(q[:, :, :1], k, v) - return grouped, masked, decode - - model = Attention() - inputs = ( - torch.randn(1, 4, 16, 64, dtype=torch.bfloat16, device="cuda"), - torch.randn(1, 4, 256, 64, dtype=torch.bfloat16, device="cuda"), - torch.randn(1, 4, 256, 64, dtype=torch.bfloat16, device="cuda"), - torch.ones(16, 16, dtype=torch.bool, device="cuda").tril(), - ) - program = export(model, inputs, strict=True).run_decompositions() - fused = [ - node - for node in program.graph.nodes - if "aten._scaled_dot_product_" in str(node.target) - ] - self.assertEqual(len(fused), 3) - with self.assertLogs( - "executorch.backends.cuda.triton.replacement_pass", "INFO" - ) as logs: - edge = to_edge_transform_and_lower( - program, - partitioner=[ - CudaPartitioner( - [CudaBackend.generate_method_name_compile_spec("forward")] - ) - ], - compile_config=EdgeCompileConfig(_check_ir_validity=False), - ) - # The masked call keeps a float bias, which the Triton kernels do not take. - output = "\n".join(logs.output) - self.assertIn("Replaced 2 nodes with Triton kernels", output) - self.assertIn("Using split-K decode SDPA", output) - with tempfile.TemporaryDirectory() as directory: - path = os.path.join(directory, "model.pte") - with open(path, "wb") as f: - edge.to_executorch().write_to_file(f) - method = Runtime.get().load_program(path).load_method("forward") - outputs = method.execute([t.cpu() for t in inputs]) - for output, reference in zip(outputs, model(*inputs), strict=True): - torch.testing.assert_close(output, reference.cpu(), atol=2e-2, rtol=2e-2) - - def test_simple_add(self): """Test CUDA export for simple element-wise addition."""