Revert "bwd pass (pytorch#164504)"

This reverts commit f36f372.

Reverted pytorch#164504 on behalf of https://github.com/jeffdaily due to CI had been clean for both cuda and rocm before merge, broke post merge? ([comment](pytorch#164504 (comment)))

Author: pytorchmergebot

test/test_varlen_attention.py

@@ -5,12 +5,11 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from torch.nn.attention.varlen import varlen_attn
+from torch.nn.attention import varlen_attn
 from torch.testing._internal.common_cuda import PLATFORM_SUPPORTS_FLASH_ATTENTION
 from torch.testing._internal.common_device_type import instantiate_device_type_tests
 from torch.testing._internal.common_nn import NNTestCase
 from torch.testing._internal.common_utils import parametrize, run_tests
-from torch.utils._python_dispatch import TorchDispatchMode
 
 
 VarlenShape = namedtuple(
@@ -24,18 +23,6 @@
 }
 
 
-class OpLoggingMode(TorchDispatchMode):
-    """Logging mode that captures all dispatched operations"""
-
-    def __init__(self):
-        self.called_ops = []
-
-    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
-        op_name = str(func)
-        self.called_ops.append(op_name)
-        return func(*args, **(kwargs or {}))
-
-
 class AttentionBlock(nn.Module):
     def __init__(
         self, embed_dim: int, num_heads: int, device: torch.device, dtype: torch.dtype
@@ -52,9 +39,12 @@ def __init__(
             embed_dim, embed_dim, bias=False, device=device, dtype=dtype
         )
 
-    def get_varlen_qkv(
+    def forward_varlen(
         self,
         x_packed: torch.Tensor,
+        cu_seq: torch.Tensor,
+        max_len: int,
+        is_causal: bool = False,
     ):
         qkv = self.qkv_proj(x_packed)
         q, k, v = qkv.chunk(3, dim=-1)
@@ -63,51 +53,24 @@ def get_varlen_qkv(
         k = k.view(-1, self.num_heads, self.head_dim)
         v = v.view(-1, self.num_heads, self.head_dim)
 
-        return q, k, v
-
-    def forward_varlen(
-        self,
-        x_packed: torch.Tensor,
-        cu_seq: torch.Tensor,
-        max_len: int,
-        is_causal: bool = False,
-    ):
-        q, k, v = self.get_varlen_qkv(x_packed)
-
-        attn_out = varlen_attn(q, k, v, cu_seq, cu_seq, max_len, max_len, is_causal)
+        attn_out = varlen_attn(
+            q, k, v, cu_seq, cu_seq, max_len, max_len, is_causal=is_causal
+        )
         attn_out = attn_out.view(-1, self.embed_dim)
 
         return self.out_proj(attn_out)
 
-    def forward_sdpa(
-        self,
-        x_padded: torch.Tensor,
-        seq_lengths: torch.Tensor,
-        dtype: torch.dtype,
-        is_causal: bool = False,
-    ):
+    def forward_sdpa(self, x_padded: torch.Tensor, is_causal: bool = False):
         batch_size, seq_len, _ = x_padded.shape
 
         qkv = self.qkv_proj(x_padded)
         q, k, v = qkv.chunk(3, dim=-1)
 
-        mask = (
-            torch.arange(seq_len, device=x_padded.device)[None, :]
-            < seq_lengths[:, None]
-        )
-
-        attn_mask = mask[:, None, None, :].expand(
-            batch_size, self.num_heads, seq_len, seq_len
-        )
-
         q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
         k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
         v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
 
-        attn_out = F.scaled_dot_product_attention(
-            q, k, v, attn_mask=attn_mask, is_causal=is_causal
-        )
-
+        attn_out = F.scaled_dot_product_attention(q, k, v, is_causal=is_causal)
         attn_out = (
             attn_out.transpose(1, 2)
             .contiguous()
@@ -128,9 +91,7 @@ def create_variable_length_batch(
     seq_lengths = torch.tensor(seq_lengths, device=device)
     total_tokens = seq_lengths.sum().item()
 
-    x_packed = torch.randn(
-        total_tokens, shape.embed_dim, device=device, dtype=dtype, requires_grad=True
-    )
+    x_packed = torch.randn(total_tokens, shape.embed_dim, device=device, dtype=dtype)
 
     cu_seq = torch.zeros(shape.batch_size + 1, device=device, dtype=torch.int32)
     cu_seq[1:] = seq_lengths.cumsum(0)
@@ -145,7 +106,6 @@ def create_variable_length_batch(
         end_idx = start_idx + seq_len
         x_padded[i, :seq_len] = x_packed[start_idx:end_idx]
         start_idx = end_idx
-    x_padded = x_padded.clone().detach().requires_grad_()
 
     return {
         "seq_lengths": seq_lengths,
@@ -173,11 +133,7 @@ def test_basic_functionality(self, device, dtype):
 
         total_tokens = shape.batch_size * shape.max_seq_len
         x_packed = torch.randn(
-            total_tokens,
-            shape.embed_dim,
-            device=device,
-            dtype=dtype,
-            requires_grad=True,
+            total_tokens, shape.embed_dim, device=device, dtype=dtype
         )
         cu_seq = torch.tensor(
             [0, shape.max_seq_len, total_tokens], device=device, dtype=torch.int32
@@ -191,128 +147,6 @@ def test_basic_functionality(self, device, dtype):
         self.assertEqual(output.device, torch.device(device))
         self.assertEqual(output.dtype, dtype)
 
-        varlen_grad_out = torch.ones_like(output)
-
-        varlen_grad = torch.autograd.grad(
-            outputs=output,
-            inputs=x_packed,
-            grad_outputs=varlen_grad_out,
-            retain_graph=True,
-            create_graph=False,
-            allow_unused=False,
-        )[0]
-
-        self.assertIsNotNone(varlen_grad)
-        self.assertEqual(varlen_grad.shape, x_packed.shape)
-        self.assertEqual(varlen_grad.dtype, x_packed.dtype)
-
-    @unittest.skipIf(
-        not PLATFORM_SUPPORTS_FLASH_ATTENTION, "Flash Attention not supported"
-    )
-    @parametrize("dtype", [torch.bfloat16, torch.float16])
-    def test_custom_op_compliance(self, device, dtype):
-        torch.manual_seed(42)
-
-        shape = VarlenShape(batch_size=2, max_seq_len=512, embed_dim=1024, num_heads=16)
-
-        attention_block = AttentionBlock(
-            shape.embed_dim, shape.num_heads, device, dtype
-        )
-
-        total_tokens = shape.batch_size * shape.max_seq_len
-        x_packed = torch.randn(
-            total_tokens,
-            shape.embed_dim,
-            device=device,
-            dtype=dtype,
-        )
-        cu_seq = torch.tensor(
-            [0, shape.max_seq_len, total_tokens], device=device, dtype=torch.int32
-        )
-
-        q, k, v = attention_block.get_varlen_qkv(x_packed)
-
-        torch.library.opcheck(
-            torch.ops.torch_attn._varlen_attn,
-            (q, k, v, cu_seq, cu_seq, shape.max_seq_len, shape.max_seq_len, False),
-        )
-
-        out, lse, rng_state = torch.ops.torch_attn._varlen_attn(
-            q, k, v, cu_seq, cu_seq, shape.max_seq_len, shape.max_seq_len, False
-        )
-        grad_out = torch.randn_like(out)
-
-        # we don't support double backward
-        # skipping test_autograd_registration, test_aot_dispatch_dynamic, test_aot_dispatch_static
-        torch.library.opcheck(
-            torch.ops.torch_attn._varlen_attn_backward,
-            (
-                grad_out,
-                q,
-                k,
-                v,
-                out,
-                lse,
-                cu_seq,
-                cu_seq,
-                shape.max_seq_len,
-                shape.max_seq_len,
-                False,
-                rng_state,
-            ),
-            test_utils=["test_schema", "test_faketensor"],
-        )
-
-    @unittest.skipIf(
-        not PLATFORM_SUPPORTS_FLASH_ATTENTION, "Flash Attention not supported"
-    )
-    @parametrize("dtype", [torch.bfloat16, torch.float16])
-    def test_custom_op_registration(self, device, dtype):
-        torch.manual_seed(42)
-
-        shape = VarlenShape(batch_size=2, max_seq_len=512, embed_dim=1024, num_heads=16)
-
-        attention_block = AttentionBlock(
-            shape.embed_dim, shape.num_heads, device, dtype
-        )
-
-        total_tokens = shape.batch_size * shape.max_seq_len
-        x_packed = torch.randn(
-            total_tokens,
-            shape.embed_dim,
-            device=device,
-            dtype=dtype,
-            requires_grad=True,
-        )
-        cu_seq = torch.tensor(
-            [0, shape.max_seq_len, total_tokens], device=device, dtype=torch.int32
-        )
-
-        compiled_forward = torch.compile(
-            attention_block.forward_varlen, backend="eager", fullgraph=True
-        )
-        with OpLoggingMode() as mode:
-            output = compiled_forward(
-                x_packed, cu_seq, shape.max_seq_len, is_causal=False
-            )
-
-            varlen_grad_out = torch.ones_like(output)
-            _ = torch.autograd.grad(
-                outputs=output,
-                inputs=x_packed,
-                grad_outputs=varlen_grad_out,
-                retain_graph=True,
-                create_graph=False,
-                allow_unused=False,
-            )[0]
-
-        called_ops = mode.called_ops
-
-        custom_ops_called = any(
-            "torch_attn._varlen_attn" in op for op in called_ops
-        ) and any("torch_attn._varlen_attn_backward" in op for op in called_ops)
-        assert custom_ops_called
-
     @unittest.skipIf(
         not PLATFORM_SUPPORTS_FLASH_ATTENTION, "Flash Attention not supported"
     )
@@ -338,10 +172,7 @@ def test_varlen_vs_sdpa(self, device, dtype, is_causal):
             is_causal=is_causal,
         )
         sdpa_output = attention_block.forward_sdpa(
-            variable_length_batch_data["x_padded"],
-            variable_length_batch_data["seq_lengths"],
-            dtype=dtype,
-            is_causal=is_causal,
+            variable_length_batch_data["x_padded"], is_causal=is_causal
         )
 
         tolerances = default_tolerances[dtype]
@@ -355,44 +186,6 @@ def test_varlen_vs_sdpa(self, device, dtype, is_causal):
             torch.testing.assert_close(varlen_seq, sdpa_seq, **tolerances)
             start_idx = end_idx
 
-        varlen_grad_out = torch.ones_like(varlen_output)
-
-        sdpa_grad_out = torch.zeros_like(sdpa_output)
-
-        start_idx = 0
-        for i, seq_len in enumerate(variable_length_batch_data["seq_lengths"]):
-            end_idx = start_idx + seq_len
-            sdpa_grad_out[i, :seq_len] = varlen_grad_out[start_idx:end_idx]
-            start_idx = end_idx
-
-        varlen_grad = torch.autograd.grad(
-            outputs=varlen_output,
-            inputs=variable_length_batch_data["x_packed"],
-            grad_outputs=varlen_grad_out,
-            retain_graph=True,
-            create_graph=False,
-            allow_unused=False,
-        )[0]
-
-        sdpa_grad = torch.autograd.grad(
-            outputs=sdpa_output,
-            inputs=variable_length_batch_data["x_padded"],
-            grad_outputs=sdpa_grad_out,
-            retain_graph=True,
-            create_graph=False,
-            allow_unused=False,
-        )[0]
-
-        start_idx = 0
-        for i, seq_len in enumerate(variable_length_batch_data["seq_lengths"]):
-            end_idx = start_idx + seq_len
-
-            varlen_grad_seq = varlen_grad[start_idx:end_idx]
-            sdpa_grad_seq = sdpa_grad[i, :seq_len]
-
-            torch.testing.assert_close(varlen_grad_seq, sdpa_grad_seq, **tolerances)
-            start_idx = end_idx
-
 
 device_types = ("cuda",)
 

torch/nn/attention/__init__.py

@@ -14,11 +14,14 @@
     SDPAParams,
 )
 
+from .varlen import varlen_attn
+
 
 __all__: list[str] = [
     "SDPBackend",
     "sdpa_kernel",
     "WARN_FOR_UNFUSED_KERNELS",
+    "varlen_attn",
 ]
 
 # Note: [SDPA warnings]