code: def test_attention_fwd_bwd(
      block,
      dtype,
      input_scale=1.0,
      scale=1 / 8.0,
      n_ctx=256,
      batch_size=2,
      n_heads=2,
  ):
      capability = torch.cuda.get_device_capability()
      if capability[0] < 7:
          pytest.skip("Only test tl.dot() on devices with sm >= 70")
  
      # inputs
      qkv_shape = (batch_size, n_heads, n_ctx, 64)
      qkvs = [
          torch.nn.Parameter(
              input_scale * torch.randn(qkv_shape), requires_grad=True
          )
          .to(dtype)
          .cuda()
          for _ in range(3)
      ]
  
      # Triton:
      n_blocks = n_ctx // block
      layout = torch.tril(
          torch.ones([n_heads, n_blocks, n_blocks], dtype=torch.long)
      )
      query, key, value = [x.clone() for x in qkvs]
      query.retain_grad()
      key.retain_grad()
      value.retain_grad()
      attn_out = triton_attention(
          layout, block, query=query, key=key, value=value, scale=scale
      )
      # ad hoc loss
      loss = (attn_out**2).mean()
      loss.backward()
      grads = [query.grad, key.grad, value.grad]
  
      # Torch version:
      torch_q, torch_k, torch_v = [x.clone() for x in qkvs]
      attn_mask = torch.ones([n_ctx, n_ctx], device="cuda", dtype=dtype)
      attn_mask = torch.tril(attn_mask, diagonal=0)
      attn_mask = 1e6 * (
          -1 + (attn_mask.reshape((1, 1, n_ctx, n_ctx)).cuda())
      )
      torch_q.retain_grad()
      torch_k.retain_grad()
      torch_v.retain_grad()
      scores = scale * torch.einsum("bhsd,bhtd->bhst", torch_q, torch_k)
      scores = scores + attn_mask
      probs = torch.softmax(scores, dim=-1)
      torch_attn_out = torch.einsum("bhst,bhtd->bhsd", probs, torch_v)
      # ad hoc loss
      torch_loss = (torch_attn_out**2).mean()
      torch_loss.backward()
      torch_grads = [torch_q.grad, torch_k.grad, torch_v.grad]
  
      # comparison
      # print(f"Triton loss {loss} and torch loss {torch_loss}.  Also checking grads...")
      torch.testing.assert_close(loss, torch_loss, atol=1e-3, rtol=0)
      for g1, g2 in zip(grads, torch_grads):
          torch.testing.assert_close(g1, g2)
  
file_path: triton/python/test/unit/operators/test_blocksparse.py
imports: import pytest
  import torch
  import triton
  import triton.ops
label_desc: Use image node_id 6 for calling a global function with example usage: test_attention_fwd_bwd(block, dtype, input_scale, scale, n_ctx, batch_size, n_heads) without return types