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

Open

shoumikhin wants to merge 90 commits into

pytorch:mainfrom

shoumikhin:cuda-graph-data-dependent

Open

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 wants to merge 90 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
  cf265f0
- List the CUDA index shims in the wheel's shared library check
  shoumikhin
  committed
  bf85d27
- Test aoti_torch_get_size
  shoumikhin
  committed
  46cfe6f
- Test more CUDA index shim and CUDA graph cases
  shoumikhin
  committed
  bb354c1
- Run a CUDA graph capture call without a graph when the capture fails
  shoumikhin
  committed
  4233d0b
- Say which dim differs when a CUDA graph replay refuses an input
  shoumikhin
  committed
  303393e
- Log an expected CUDA graph capture failure at Info
  shoumikhin
  committed
  0932472
- Record CUDA graph allocations only after the first launch
  shoumikhin
  committed
  50700ef
- Check the new CUDA graph input shapes in the recapture and replay tests
  shoumikhin
  committed
  6ca35a1
- Delete the antialiased resize torch.compile test
  shoumikhin
  committed
  d473a0d
- Say why the antialiased resize avoids a conditional
  shoumikhin
  committed
  a882031
- Refuse a hand-built data name in AOTI backends that keep data in the .pte
  shoumikhin
  committed
  a773a35
- Warn that CUDA data names differing only in case can share a file
  shoumikhin
  committed
  a76b9de
- Correct comments on the CUDA index shims and set_output_data_ptr
  shoumikhin
  committed
  b5b99de
- Seed the CUDA random shims from std::random_device
  shoumikhin
  committed
  a677f3d
- Run a CUDA method that keeps host state without a graph
  shoumikhin
  committed
  a21e73b
- Check the AOTI data name before compiling
  shoumikhin
  committed
  cdf7ed0
- Check delegation and restore deterministic mode in the CUDA index shim tests
  shoumikhin
  committed
  21decf8
- Fix the CUDA graph capture guard comment
  shoumikhin
  committed
  1e0cae1
- Copy through a boolean mask with where in the CUDA backend
  shoumikhin
  committed
  27fc353
- Keep ops with more than 6 dims out of the XNNPACK delegate
  shoumikhin
  committed
  bbc2239
- Drop tensor checks before the CUDA backend compiles a program
  shoumikhin
  committed
  2066a73
- Test the CUDA masked copy pass and the capture retry end to end
  shoumikhin
  committed
  cffa98d
- Remove a size check only when it tests the index's count
  shoumikhin
  committed
  5b307d8
- Say why a CUDA method keeps its graph off for host constants
  shoumikhin
  committed
  7311573
- Keep destroy() with the other public CUDA backend methods
  shoumikhin
  committed
  f2ce3a4
- Draw the CUDA RNG seed before allocating its state
  shoumikhin
  committed
  51d29e6
- Format the RemoveAssertAsyncPass import with ufmt
  shoumikhin
  committed
  0b171d6
- Load float64 weights in the CUDA backend
  shoumikhin
  committed
  1f6412a
- Decompose aten._reshape_copy in the CUDA backend
  shoumikhin
  committed
  40c11b5
- Add aoti_torch_dtype_float64 to aoti_cuda_shims.lib
  shoumikhin
  committed
  f052fa1
- Retrigger CI
  shoumikhin
  committed
  0ff1f96
- Add a CUDA shim for scatter_reduce with amax, amin, prod and mean
  shoumikhin
  committed
  39a8943
- Retrigger CI
  shoumikhin
  committed
  38468fb
- Decompose varlen_attn in the CUDA backend
  shoumikhin
  committed
  eca5be6
- Test the int64 scatter sum in the CUDA backend instead of expecting it to fail
  shoumikhin
  committed
  b39a9ed
- Use a model PyTorch 2.14 can export in the CUDA tensor check test
  shoumikhin
  committed
  ddaa65d
- Check that the CUDA mask tests try a graph capture
  shoumikhin
  committed
  9832675
- Pin the masked copy pass's integer index and shared picks guards
  shoumikhin
  committed
  e1535cc
- Compare masks without recursion in the CUDA masked copy pass
  shoumikhin
  committed
  38a8306
- Accept float64 in the CUDA shims that still refused it
  shoumikhin
  committed
  8ca54cc
- Count a scatter mean in the tensor's dtype, as ATen does
  shoumikhin
  committed
  91a4f7a
- Retry the CUDA random state setup when its seed copy fails
  shoumikhin
  committed
  4a4c758
- List every exported CUDA shim in the Windows import library
  shoumikhin
  committed
  11724c1
- Keep buffer updates across calls in AOTI backends with PyTorch 2.14
  shoumikhin
  committed
  7158c02
- Test that the CUDA missing-kernel check sees the scatter fallback
  shoumikhin
  committed
  902c4c7
- Drop a redundant branch from the masked copy pass and describe what it accepts
  shoumikhin
  committed
  f1ade88
- Read a float64 scalar in the CUDA float64 shim test
  shoumikhin
  committed
  e38f0f0
- Use a scatter the CUDA backend lacks in the missing-kernel test
  shoumikhin
  committed
  9775fe1
- Return a contiguous output from the CUDA varlen_attn decomposition
  shoumikhin
  committed
  6d4cd43
- Refuse varlen_attn windows on which Flash attention versions differ
  shoumikhin
  committed
  1c46039
- Keep NaN values outside the key sequences out of varlen_attn outputs
  shoumikhin
  committed
  493c7bd
- Correct the comments on the AOTI folding fix and the scatter kernel test
  shoumikhin
  committed
  f64c5f3
- Test the count check guard in the masked copy pass
  shoumikhin
  committed
  58532ea
- Correct two comments in the CUDA index shim tests
  shoumikhin
  committed
  3d3a3c8
- Test a scatter mean count that wraps to zero
  shoumikhin
  committed
  c986794
- Say the CUDA random state setup failed, not its allocation
  shoumikhin
  committed
  0aac975
- Add the float16 scalar read to the CUDA shims
  shoumikhin
  committed
  41ab368
- Move the XNNPACK rank check to its own pull request
  shoumikhin
  committed
  ce8c75c
- Make two varlen_attn decomposition comments exact
  shoumikhin
  committed
  fedd43f
- Say that float accumulation in the CUDA index_put shim can still round away small terms
  shoumikhin
  committed
  0d2a06d
- Say what loading CUDA weights waits for in the stream docs
  shoumikhin
  committed
  8206ba1
- Index the CUDA mask test with a negative int32 column
  shoumikhin
  committed
  b8e7aca
- Say the CUDA scatter_reduce sum loses less to rounding, not nothing
  shoumikhin
  committed
  f213adf
- Drop the unused libgomp dependency from CUDA backend programs
  shoumikhin
  committed
  c25dd63
- Keep keys past the end out of the CUDA Triton SDPA softmax
  shoumikhin
  committed
  51788e2
- Remove the CUDA backend's temporary library when a load fails
  shoumikhin
  committed
  54b3d3c
- Point the CUDA backend's libgomp TODO at the PyTorch fix
  shoumikhin
  committed
  05a407f
- Add a CUDA shim for scatter with a tensor source
  shoumikhin
  committed
  644ed5d
- Let the CUDA scatter shims read an int32 index
  shoumikhin
  committed
  0931251
- Write each element once in the CUDA scatter shim
  shoumikhin
  committed
  9510d4f
- Name the calling CUDA scatter shim in its errors
  shoumikhin
  committed
  ba4e077
- Say which ATen kernel the CUDA scatter bounds check matches
  shoumikhin
  committed
  da8e023
- Run a CUDA graph capture that reads a device value without a graph
  shoumikhin
  committed
  e279513
- Rewrite masked copies inside torch.cond branches too
  shoumikhin
  committed
  e6aba30
- Keep the CUDA compile patches in place while any thread lowers
  shoumikhin
  committed
  a909f5d
- Pass a float read with .item() to Triton kernels as a float
  shoumikhin
  committed
  420f65e
- Say the AOTI folding switch covers the whole program
  shoumikhin
  committed
  31e1685
- Test the CUDA scatter shims on more layouts, dtypes and edge cases
  shoumikhin
  committed
  15acfa6
- Run the varlen_attn and _reshape_copy programs in the CUDA tests
  shoumikhin
  committed
  87730de
- Keep PyTorch's internal test helpers out of the CUDA index shim tests
  shoumikhin
  committed
  9102e04
- Assign the CUDA compile patches explicitly
  shoumikhin
  committed
  46e21f3
- Test that the masked copy pass keeps two random masks apart
  shoumikhin
  committed
  cb9ec88
- Check that a failed CUDA graph capture is tried only once
  shoumikhin
  committed
  e707127
- Lower the _reshape_copy and varlen_attn programs on ROCm again
  shoumikhin
  committed
  2edba93
- Use a stable softmax in the CUDA varlen_attn decomposition
  shoumikhin
  committed
  d7f5d66
- Keep masked copies whose layouts differ
  shoumikhin
  committed
  fc8280e
- Tidy the masked copy pass and its tests
  shoumikhin
  committed
  f4f8aa7
- Check that the CUDA backend runs RemoveAssertAsyncPass
  shoumikhin
  committed
  72ee3cc
- Clarify CUDA backend names and comments
  shoumikhin
  committed
  4c89f6a

Commits

Commits on Nov 2, 2026