Repository navigation

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

Merged

shoumikhin merged 89 commits into

pytorch:mainfrom

shoumikhin:cuda-graph-data-dependent

Nov 2, 2026

Merged

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
shoumikhin merged 89 commits into
pytorch:mainfrom
shoumikhin:cuda-graph-data-dependent

Commits on Nov 2, 2026
- Refuse a CUDA graph replay when an input changes shape
  shoumikhin
  committed
  960ea18
- List the CUDA index shims in the wheel's shared library check
  shoumikhin
  committed
  5f3ce38
- Test aoti_torch_get_size
  shoumikhin
  committed
  c4cb8cf
- Test more CUDA index shim and CUDA graph cases
  shoumikhin
  committed
  301e0e1
- Run a CUDA graph capture call without a graph when the capture fails
  shoumikhin
  committed
  03e2c82
- Say which dim differs when a CUDA graph replay refuses an input
  shoumikhin
  committed
  bf99895
- Log an expected CUDA graph capture failure at Info
  shoumikhin
  committed
  5202db7
- Record CUDA graph allocations only after the first launch
  shoumikhin
  committed
  1d638d8
- Check the new CUDA graph input shapes in the recapture and replay tests
  shoumikhin
  committed
  474e661
- Delete the antialiased resize torch.compile test
  shoumikhin
  committed
  69d8749
- Say why the antialiased resize avoids a conditional
  shoumikhin
  committed
  8eafbbd
- Refuse a hand-built data name in AOTI backends that keep data in the .pte
  shoumikhin
  committed
  ce87740
- Warn that CUDA data names differing only in case can share a file
  shoumikhin
  committed
  84d1468
- Correct comments on the CUDA index shims and set_output_data_ptr
  shoumikhin
  committed
  f86630c
- Seed the CUDA random shims from std::random_device
  shoumikhin
  committed
  4b7dabd
- Run a CUDA method that keeps host state without a graph
  shoumikhin
  committed
  26d0612
- Check the AOTI data name before compiling
  shoumikhin
  committed
  bfca1cd
- Check delegation and restore deterministic mode in the CUDA index shim tests
  shoumikhin
  committed
  19a6331
- Fix the CUDA graph capture guard comment
  shoumikhin
  committed
  f332a88
- Copy through a boolean mask with where in the CUDA backend
  shoumikhin
  committed
  7eaf9ea
- Drop tensor checks before the CUDA backend compiles a program
  shoumikhin
  committed
  ebca400
- Test the CUDA masked copy pass and the capture retry end to end
  shoumikhin
  committed
  9e380c8
- Remove a size check only when it tests the index's count
  shoumikhin
  committed
  3c73496
- Say why a CUDA method keeps its graph off for host constants
  shoumikhin
  committed
  423ebaf
- Keep destroy() with the other public CUDA backend methods
  shoumikhin
  committed
  ccc1259
- Draw the CUDA RNG seed before allocating its state
  shoumikhin
  committed
  984df09
- Format the RemoveAssertAsyncPass import with ufmt
  shoumikhin
  committed
  1bfa4de
- Load float64 weights in the CUDA backend
  shoumikhin
  committed
  76f519a
- Decompose aten._reshape_copy in the CUDA backend
  shoumikhin
  committed
  ebcde7b
- Add aoti_torch_dtype_float64 to aoti_cuda_shims.lib
  shoumikhin
  committed
  2fa5105
- Retrigger CI
  shoumikhin
  committed
  f9d29a1
- Add a CUDA shim for scatter_reduce with amax, amin, prod and mean
  shoumikhin
  committed
  89bebfa
- Retrigger CI
  shoumikhin
  committed
  f23afa8
- Decompose varlen_attn in the CUDA backend
  shoumikhin
  committed
  da1db73
- Test the int64 scatter sum in the CUDA backend instead of expecting it to fail
  shoumikhin
  committed
  8a4ebbf
- Use a model PyTorch 2.14 can export in the CUDA tensor check test
  shoumikhin
  committed
  4120761
- Check that the CUDA mask tests try a graph capture
  shoumikhin
  committed
  1a0704a
- Pin the masked copy pass's integer index and shared picks guards
  shoumikhin
  committed
  7084c6f
- Compare masks without recursion in the CUDA masked copy pass
  shoumikhin
  committed
  0869711
- Accept float64 in the CUDA shims that still refused it
  shoumikhin
  committed
  e29e302
- Count a scatter mean in the tensor's dtype, as ATen does
  shoumikhin
  committed
  fe11987
- Retry the CUDA random state setup when its seed copy fails
  shoumikhin
  committed
  9e707a6
- List every exported CUDA shim in the Windows import library
  shoumikhin
  committed
  4341022
- Keep buffer updates across calls in AOTI backends with PyTorch 2.14
  shoumikhin
  committed
  05e7d7a
- Test that the CUDA missing-kernel check sees the scatter fallback
  shoumikhin
  committed
  b66b64e
- Drop a redundant branch from the masked copy pass and describe what it accepts
  shoumikhin
  committed
  951d55b
- Read a float64 scalar in the CUDA float64 shim test
  shoumikhin
  committed
  911db99
- Use a scatter the CUDA backend lacks in the missing-kernel test
  shoumikhin
  committed
  4a8bc18
- Return a contiguous output from the CUDA varlen_attn decomposition
  shoumikhin
  committed
  900d0ee
- Refuse varlen_attn windows on which Flash attention versions differ
  shoumikhin
  committed
  5c5d262
- Keep NaN values outside the key sequences out of varlen_attn outputs
  shoumikhin
  committed
  e9a232f
- Correct the comments on the AOTI folding fix and the scatter kernel test
  shoumikhin
  committed
  9680613
- Test the count check guard in the masked copy pass
  shoumikhin
  committed
  c3a6f91
- Correct two comments in the CUDA index shim tests
  shoumikhin
  committed
  a4767c3
- Test a scatter mean count that wraps to zero
  shoumikhin
  committed
  89c4d5f
- Say the CUDA random state setup failed, not its allocation
  shoumikhin
  committed
  d6f4a34
- Add the float16 scalar read to the CUDA shims
  shoumikhin
  committed
  1f884d9
- Make two varlen_attn decomposition comments exact
  shoumikhin
  committed
  3cd283f
- Say that float accumulation in the CUDA index_put shim can still round away small terms
  shoumikhin
  committed
  eebec99
- Say what loading CUDA weights waits for in the stream docs
  shoumikhin
  committed
  ba1c036
- Index the CUDA mask test with a negative int32 column
  shoumikhin
  committed
  f342450
- Say the CUDA scatter_reduce sum loses less to rounding, not nothing
  shoumikhin
  committed
  a9fc94a
- Drop the unused libgomp dependency from CUDA backend programs
  shoumikhin
  committed
  53a0623
- Keep keys past the end out of the CUDA Triton SDPA softmax
  shoumikhin
  committed
  5a94bfa
- Remove the CUDA backend's temporary library when a load fails
  shoumikhin
  committed
  4dc4d5a
- Point the CUDA backend's libgomp TODO at the PyTorch fix
  shoumikhin
  committed
  624b83a
- Add a CUDA shim for scatter with a tensor source
  shoumikhin
  committed
  4f53e3d
- Let the CUDA scatter shims read an int32 index
  shoumikhin
  committed
  96e111d
- Write each element once in the CUDA scatter shim
  shoumikhin
  committed
  f021c76
- Name the calling CUDA scatter shim in its errors
  shoumikhin
  committed
  f6108de
- Say which ATen kernel the CUDA scatter bounds check matches
  shoumikhin
  committed
  c1baac9
- Run a CUDA graph capture that reads a device value without a graph
  shoumikhin
  committed
  ab9ed1c
- Rewrite masked copies inside torch.cond branches too
  shoumikhin
  committed
  10f3e2f
- Keep the CUDA compile patches in place while any thread lowers
  shoumikhin
  committed
  2b350d9
- Pass a float read with .item() to Triton kernels as a float
  shoumikhin
  committed
  ba5d868
- Say the AOTI folding switch covers the whole program
  shoumikhin
  committed
  d996e09
- Test the CUDA scatter shims on more layouts, dtypes and edge cases
  shoumikhin
  committed
  2a09294
- Run the varlen_attn and _reshape_copy programs in the CUDA tests
  shoumikhin
  committed
  cb41bb9
- Keep PyTorch's internal test helpers out of the CUDA index shim tests
  shoumikhin
  committed
  b2137cd
- Assign the CUDA compile patches explicitly
  shoumikhin
  committed
  ec7046c
- Test that the masked copy pass keeps two random masks apart
  shoumikhin
  committed
  19a8028
- Check that a failed CUDA graph capture is tried only once
  shoumikhin
  committed
  e922a55
- Lower the _reshape_copy and varlen_attn programs on ROCm again
  shoumikhin
  committed
  7e59c0a
- Use a stable softmax in the CUDA varlen_attn decomposition
  shoumikhin
  committed
  8f9dc8f
- Keep masked copies whose layouts differ
  shoumikhin
  committed
  fa053a7
- Tidy the masked copy pass and its tests
  shoumikhin
  committed
  88f1538
- Check that the CUDA backend runs RemoveAssertAsyncPass
  shoumikhin
  committed
  b508df4
- Clarify CUDA backend names and comments
  shoumikhin
  committed
  fd5948e
- Compare the varlen_attn decomposition with the op on CUDA only
  shoumikhin
  committed
  fab02c4

Commits

Commits on Nov 2, 2026