OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-06 00:20:18 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity
f20bf77874
pytorch/test
History
Angel Li 476b149a00 bwd pass (#164504) 
**Summary**
This implements the backward pass for the Varlen API and registers `_varlen_attn()` as a custom op.

**Benchmarking**

To benchmark, we compare runtime and TFLOPs against the current SDPA approach with padding.

Settings:

- 1 H100 machine
- `batch_size=8`, `max_seq_len=2048`, `embed_dim=1024`, `num_heads=16`
- dtype `torch.bfloat16`
- `is_causal=False`
- for variable length, we set sequences to be random multiples of 64 up to `max_seq_len`
- 100 runs

|        | Variable Length API | SDPA     |
|--------|--------------------|----------|
| Runtime | 0.8189142608642578 ms       | 3.263883056640625 ms  |
| TFLOPs | 268.652       | 158.731  |

We can see that runtime for Varlen is >3x faster

**Testing**

Run `python test/test_varlen_attention.py` for unit tests where we verify basic functionality and confirm numerical match between varlen gradients vs SDPA.

For custom op testing, `test_custom_op_registration` uses logging mode to verify that `_varlen_attn()` was called and tests with `torch.compile`. `test_custom_op_compliances` uses `torch.library.opcheck()` to verify.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/164504
Approved by: https://github.com/drisspg
2025-10-30 03:46:37 +00:00
..
ao/sparsity Enable all PIE rules on ruff (#165814) 2025-10-18 07:36:18 +00:00
autograd
backends/xeon
benchmark_utils
compiled_autograd_skips
cpp [2/N][Fix] Fix typo in test folder (#166374) 2025-10-29 03:02:07 +00:00
cpp_api_parity
cpp_extensions [user-streams] Support streams as contexts (#164507) 2025-10-29 04:46:08 +00:00
custom_backend
custom_operator
distributed Enable local tensor mode for DTensor attention and convolution tests (#166406) 2025-10-30 02:48:02 +00:00
distributions Enable PLW0127 in ruff (#165851) 2025-10-21 03:30:57 +00:00
dynamo Fix comparing inductor actual strides vs bw graph for activations should not throw DDE. (#166277) 2025-10-30 00:34:05 +00:00
dynamo_expected_failures [dynamo] Fix ListIterator tracking mutations to original list (#166350) 2025-10-29 21:54:37 +00:00
dynamo_skips Add functions to setup PrivateUse1 as a python backend device. (#157859) 2025-10-01 21:32:59 +00:00
error_messages
expect [fx] Optimize torch.fx.Node.replace_all_uses_with (#165889) 2025-10-25 03:44:41 +00:00
export Revert "[pytree] add treespec_{leaf,tuple,dict} functions for args_spec modification (#160843)" 2025-10-29 22:46:48 +00:00
forward_backward_compatibility Revert "[BC-Breaking] Remove long-deprecated casting functions from native_functions.yaml (#164641)" 2025-10-09 15:42:51 +00:00
functorch [2/N][Fix] Fix typo in test folder (#166374) 2025-10-29 03:02:07 +00:00
fx [2/N][Fix] Fix typo in test folder (#166374) 2025-10-29 03:02:07 +00:00
higher_order_ops Update gm.print_readable to include Annotation (#165397) 2025-10-28 13:54:38 +00:00
inductor [xpu][test] Reuse native_mm and mix_order_reduction for Intel GPU. (#166384) 2025-10-30 03:38:35 +00:00
inductor_expected_failures [dynamo] Graph break on on user-defined class in compiled region (#161670) 2025-09-10 04:39:20 +00:00
inductor_skips
jit [2/N][Fix] Fix typo in test folder (#166374) 2025-10-29 03:02:07 +00:00
jit_hooks
lazy [2/N][Fix] Fix typo in test folder (#166374) 2025-10-29 03:02:07 +00:00
mobile [2/N][Fix] Fix typo in test folder (#166374) 2025-10-29 03:02:07 +00:00
nn [Conv1d] Check overflow before we compute padding size. (#162363) 2025-10-29 03:27:20 +00:00
onnx [ONNX] Create fake implementations for onnx ops; fix boolean mask in attention (#165780) 2025-10-29 04:51:49 +00:00
optim Enable PLW0127 in ruff (#165851) 2025-10-21 03:30:57 +00:00
package [2/N][Fix] Fix typo in test folder (#166374) 2025-10-29 03:02:07 +00:00
profiler [2/N][Fix] Fix typo in test folder (#166374) 2025-10-29 03:02:07 +00:00
quantization [2/N] Fix unused loop variables (#166500) 2025-10-29 08:30:35 +00:00
scripts
strobelight/examples
test_img
torch_np [2/N] Fix unused loop variables (#166500) 2025-10-29 08:30:35 +00:00
typing [2/N][Fix] Fix typo in test folder (#166374) 2025-10-29 03:02:07 +00:00
xpu [Intel GPU] Xpu matmul implementation for complex dtype (#160867) 2025-10-25 17:13:13 +00:00
_test_bazel.py
allowlist_for_publicAPI.json Revert "[BC-Breaking] Remove long-deprecated casting functions from native_functions.yaml (#164641)" 2025-10-09 15:42:51 +00:00
bench_mps_ops.py [BE] Remove macos-13 guard from bench_mps_ops (#159732) 2025-08-03 20:53:58 +00:00
conftest.py [2/N][Fix] Fix typo in test folder (#166374) 2025-10-29 03:02:07 +00:00
create_dummy_torchscript_model.py
HowToWriteTestsUsingFileCheck.md
linear.py
load_torchscript_model.py
minioptest_failures_dict.json
mkl_verbose.py
mkldnn_verbose.py
pytest_shard_custom.py
run_doctests.sh
run_test.py [2/N][Fix] Fix typo in test folder (#166374) 2025-10-29 03:02:07 +00:00
simulate_nccl_errors.py
slow_tests.json Update slow tests (#165894) 2025-10-27 11:42:14 +00:00
test_accelerator.py Add UT for torch.accelerator memory-related API (#155200) 2025-08-08 17:41:22 +00:00
test_ao_sparsity.py [testing] Add test owner labels for some ao sparse tests (#163203) 2025-09-18 16:08:13 +00:00
test_appending_byte_serializer.py
test_autocast.py [AMP][Refactor] Autocast dtype handling to simplify device-specific c… (#165221) 2025-10-28 06:21:29 +00:00
test_autograd_fallback.py Fix TestAutogradFallback flaky tests under Dynamo: migrate to lib._destroy() (#159443) 2025-07-30 19:30:55 +00:00
test_autograd.py [2/N][Fix] Fix typo in test folder (#166374) 2025-10-29 03:02:07 +00:00
test_autoload.py
test_binary_ufuncs.py Enable ruff rule E721 (#165162) 2025-10-13 01:48:55 +00:00
test_bundled_images.py
test_bundled_inputs.py Pyrefly suppressions 7/n (#164913) 2025-10-08 07:27:17 +00:00
test_ci_sanity_check_fail.py
test_comparison_utils.py
test_compile_benchmark_util.py
test_complex.py Pyrefly suppressions 7/n (#164913) 2025-10-08 07:27:17 +00:00
test_content_store.py [ROCm] Enabling several UTs (#161715) 2025-09-09 15:49:21 +00:00
test_cpp_api_parity.py Revert "Stop parsing command line arguments every time common_utils is imported. (#156703)" 2025-08-04 20:37:39 +00:00
test_cpp_extensions_aot.py Replace setup.py install with pip install (#156711) 2025-09-29 15:15:10 +00:00
test_cpp_extensions_jit.py Move from/to to torch::stable::detail (#164956) 2025-10-21 02:59:46 +00:00
test_cpp_extensions_mtia_backend.py
test_cpp_extensions_stream_and_event.py
test_cuda_expandable_segments.py
test_cuda_multigpu.py [ROCm] Enabling several UTs (#161715) 2025-09-09 15:49:21 +00:00
test_cuda_nvml_based_avail.py More ruff SIM fixes (#164695) 2025-10-09 03:24:50 +00:00
test_cuda_primary_ctx.py [ROCm][CI] skip TestCudaPrimaryCtx.test_set_device_0 (#163693) 2025-09-23 22:15:10 +00:00
test_cuda_sanitizer.py
test_cuda_trace.py
test_cuda.py [ROCm][CUDA] add unit test utility busy_wait_for_flag (#166218) 2025-10-29 22:40:23 +00:00
test_custom_ops.py [Bugfix][vLLM] Explicitly do not support instead of crashing for named tuples in infer schema (#165191) 2025-10-14 14:18:42 +00:00
test_dataloader.py Enable PLW0127 in ruff (#165851) 2025-10-21 03:30:57 +00:00
test_datapipe.py Enable all PIE rules on ruff (#165814) 2025-10-18 07:36:18 +00:00
test_decomp.py address DDE in matmul decomp (#166541) 2025-10-30 03:19:29 +00:00
test_determination.py
test_dispatch.py
test_dlpack.py [Fix] Completely remove stride normalization on DLPack Tensor (#164161) 2025-10-14 17:17:11 +00:00
test_dynamic_shapes.py Support python slicing with tensor inputs. (#165074) 2025-10-29 01:18:45 +00:00
test_expanded_weights.py Revert "Stop parsing command line arguments every time common_utils is imported. (#156703)" 2025-08-04 20:37:39 +00:00
test_extension_utils.py
test_fake_tensor.py [FakeTensor] Supplement the relevant logic for converting conv1d to conv2d in meta_conv (#160408) 2025-09-26 15:45:02 +00:00
test_file_check.py
test_flop_counter.py [ROCm] Enable MI355 CI on PRs, and run full set of UTs on PRs (#160215) 2025-10-09 18:03:12 +00:00
test_foreach.py Remove old ROCm version check in tests (#164245) 2025-10-06 22:42:01 +00:00
test_function_schema.py
test_functional_autograd_benchmark.py
test_functional_optim.py
test_functionalization_of_rng_ops.py
test_functionalization.py
test_futures.py
test_fx_experimental.py [8/N] Apply ruff UP035 rule (#165214) 2025-10-15 03:18:57 +00:00
test_fx_passes.py
test_fx_reinplace_pass.py
test_fx.py [2/N][Fix] Fix typo in test folder (#166374) 2025-10-29 03:02:07 +00:00
test_hop_infra.py
test_hub.py
test_import_stats.py
test_indexing.py [2/N][Fix] Fix typo in test folder (#166374) 2025-10-29 03:02:07 +00:00
test_itt.py Enable PLW0127 in ruff (#165851) 2025-10-21 03:30:57 +00:00
test_jit_autocast.py Enable all SIM rules except disabled ones (#164645) 2025-10-17 07:27:11 +00:00
test_jit_disabled.py
test_jit_fuser_legacy.py Stop parsing command line arguments every time common_utils is imported. (#156703) 2025-10-02 22:22:04 +00:00
test_jit_fuser_te.py Enable all PIE rules on ruff (#165814) 2025-10-18 07:36:18 +00:00
test_jit_fuser.py Stop parsing command line arguments every time common_utils is imported. (#156703) 2025-10-02 22:22:04 +00:00
test_jit_legacy.py Stop parsing command line arguments every time common_utils is imported. (#156703) 2025-10-02 22:22:04 +00:00
test_jit_llga_fuser.py
test_jit_profiling.py
test_jit_simple.py
test_jit_string.py More ruff SIM fixes (#164695) 2025-10-09 03:24:50 +00:00
test_jit.py Enable PLW0127 in ruff (#165851) 2025-10-21 03:30:57 +00:00
test_jiterator.py
test_kernel_launch_checks.py [testing] Add test owner labels for some cuda? tests (#163296) 2025-09-26 18:26:56 +00:00
test_legacy_vmap.py
test_license.py
test_linalg.py [2/N][Fix] Fix typo in test folder (#166374) 2025-10-29 03:02:07 +00:00
test_logging.py
test_masked.py
test_maskedtensor.py
test_matmul_cuda.py Revert "[ROCm] Enable group gemm through CK (#166334)" 2025-10-29 16:45:02 +00:00
test_meta.py [Inductor] support mixed dtype in the native_layer_norm_backward meta function (#159830) 2025-09-17 20:29:12 +00:00
test_metal.py
test_mkl_verbose.py
test_mkldnn_fusion.py
test_mkldnn_verbose.py
test_mkldnn.py [2/N][Fix] Fix typo in test folder (#166374) 2025-10-29 03:02:07 +00:00
test_mobile_optimizer.py
test_model_exports_to_core_aten.py Replace export_for_training with export (#162396) 2025-09-10 14:19:34 +00:00
test_module_tracker.py
test_modules.py [2/N][Fix] Fix typo in test folder (#166374) 2025-10-29 03:02:07 +00:00
test_monitor.py
test_mps.py [2/N][Fix] Fix typo in test folder (#166374) 2025-10-29 03:02:07 +00:00
test_multiprocessing_spawn.py Change forkserver test to only run below 3.13.8 (#165667) 2025-10-16 19:34:10 +00:00
test_multiprocessing.py Enable PLW0127 in ruff (#165851) 2025-10-21 03:30:57 +00:00
test_namedtensor.py Fix tensor creation with empty names crash (#163957) 2025-09-26 17:41:00 +00:00
test_namedtuple_return_api.py
test_native_functions.py
test_native_mha.py
test_nestedtensor.py [CI] Move the periodic debug tests to newer runner (#165158) 2025-10-21 21:28:12 +00:00
test_nn.py [2/N][Fix] Fix typo in test folder (#166374) 2025-10-29 03:02:07 +00:00
test_nnapi.py
test_numa_binding.py Replace NUMA inheritance approach (#166026) 2025-10-29 03:58:44 +00:00
test_numba_integration.py [testing] Add test owner labels for some cuda? tests (#163296) 2025-09-26 18:26:56 +00:00
test_numpy_interop.py Enable all SIM rules except disabled ones (#164645) 2025-10-17 07:27:11 +00:00
test_opaque_obj_v2.py [opaque_obj_v2] PyObject custom op schema type (#165004) 2025-10-14 20:21:04 +00:00
test_opaque_obj.py [opaque_obj] Add make_fx tracing support (#163278) 2025-10-08 09:09:16 +00:00
test_openmp.py
test_ops_fwd_gradients.py
test_ops_gradients.py
test_ops_jit.py
test_ops.py [2/N][Fix] Fix typo in test folder (#166374) 2025-10-29 03:02:07 +00:00
test_optim.py Fix LBFGS wolfe max iteration (#161488) 2025-09-16 12:07:50 +00:00
test_out_dtype_op.py
test_overrides.py Revert "[Reland] Use std::string_view in torchgen (#158625)" 2025-09-13 07:52:50 +00:00
test_package.py
test_per_overload_api.py
test_prims.py [Bugfix] Match eager stride semantics for cloned tensors with preserve_format in compile (#163017) 2025-09-19 19:41:33 +00:00
test_privateuseone_python_backend.py Add functions to setup PrivateUse1 as a python backend device. (#157859) 2025-10-01 21:32:59 +00:00
test_proxy_tensor.py Fix refine_ranges corner case (#164075) (#164846) 2025-10-08 18:42:37 +00:00
test_pruning_op.py Enable all SIM rules except disabled ones (#164645) 2025-10-17 07:27:11 +00:00
test_public_bindings.py [CUDA] Add experimental green context support for SM carveout (#159104) 2025-10-22 21:38:52 +00:00
test_python_dispatch.py [2/N][Fix] Fix typo in test folder (#166374) 2025-10-29 03:02:07 +00:00
test_pytree.py Revert "[pytree] add treespec_{leaf,tuple,dict} functions for args_spec modification (#160843)" 2025-10-29 22:46:48 +00:00
test_quantization.py [2/N][Fix] Fix typo in test folder (#166374) 2025-10-29 03:02:07 +00:00
test_reductions.py Enable all PIE rules on ruff (#165814) 2025-10-18 07:36:18 +00:00
test_rename_privateuse1_to_existing_device.py
test_scaled_matmul_cuda.py [2/N][Fix] Fix typo in test folder (#166374) 2025-10-29 03:02:07 +00:00
test_scatter_gather_ops.py fix deterministic scatter_add path for multi-d tensors (#162866) 2025-09-15 06:50:00 +00:00
test_schema_check.py [inductor] slow test some Windows UTs. (#160267) 2025-08-10 18:35:42 +00:00
test_segment_reductions.py Enable all SIM rules except disabled ones (#164645) 2025-10-17 07:27:11 +00:00
test_serialization.py Quick fix of torch.save memory leak (#165204) 2025-10-27 07:50:58 +00:00
test_set_default_mobile_cpu_allocator.py
test_shape_ops.py [2/N][Fix] Fix typo in test folder (#166374) 2025-10-29 03:02:07 +00:00
test_show_pickle.py
test_sort_and_select.py Add dtype checks in meta dispatch for various ordering ops (#159556) 2025-08-14 17:06:27 +00:00
test_sparse_csr.py Enable PLW0127 in ruff (#165851) 2025-10-21 03:30:57 +00:00
test_sparse_semi_structured.py [CI] Move the periodic debug tests to newer runner (#165158) 2025-10-21 21:28:12 +00:00
test_sparse.py [2/N][Fix] Fix typo in test folder (#166374) 2025-10-29 03:02:07 +00:00
test_spectral_ops.py
test_stateless.py [2/N][Fix] Fix typo in test folder (#166374) 2025-10-29 03:02:07 +00:00
test_static_runtime.py Enable all PIE rules on ruff (#165814) 2025-10-18 07:36:18 +00:00
test_subclass.py
test_sympy_utils.py [1/N] Apply UP035 rule in tests (#163947) 2025-09-29 01:42:01 +00:00
test_tensor_creation_ops.py test for #165446 (#165853) 2025-10-23 00:08:18 +00:00
test_tensorboard.py Add back Windows and macOS to tensorboard tests (#166389) 2025-10-29 04:34:57 +00:00
test_tensorexpr_pybind.py
test_tensorexpr.py [2/N][Fix] Fix typo in test folder (#166374) 2025-10-29 03:02:07 +00:00
test_testing.py [1/N] Apply UP035 rule in tests (#163947) 2025-09-29 01:42:01 +00:00
test_throughput_benchmark.py
test_torch.py Fix incomplete torch.cdist tests (#166507) 2025-10-29 17:11:07 +00:00
test_torchfuzz_repros.py Enable all SIM rules except disabled ones (#164645) 2025-10-17 07:27:11 +00:00
test_transformers.py [2/N][Fix] Fix typo in test folder (#166374) 2025-10-29 03:02:07 +00:00
test_type_hints.py Pyrefly suppressions 7/n (#164913) 2025-10-08 07:27:17 +00:00
test_type_info.py Enable PLW0127 in ruff (#165851) 2025-10-21 03:30:57 +00:00
test_type_promotion.py Enable PLW0127 in ruff (#165851) 2025-10-21 03:30:57 +00:00
test_typing.py
test_unary_ufuncs.py disable jiterator for complex tan and tanh (#165250) 2025-10-29 04:59:01 +00:00
test_utils_config_module.py
test_utils_filelock.py
test_utils.py Enable PLW0127 in ruff (#165851) 2025-10-21 03:30:57 +00:00
test_varlen_attention.py bwd pass (#164504) 2025-10-30 03:46:37 +00:00
test_view_ops.py Enable PLW0127 in ruff (#165851) 2025-10-21 03:30:57 +00:00
test_vulkan.py
test_weak.py
test_xnnpack_integration.py Enable all PIE rules on ruff (#165814) 2025-10-18 07:36:18 +00:00
test_xpu.py Introduce a new API torch.xpu.set_per_process_memory_fraction (#165510) 2025-10-29 03:24:52 +00:00