code: def test_flash_attn_bwd_varlen_overflow(d, causal, dtype):
      """We previously had a bug where not masking elements beyond seqlen_k caused NaN in dQ,
      in the case where seqlen % 128 != 0 or varlen.
      """
      device = "cuda"
      # set seed
      torch.random.manual_seed(0)
      nheads = 5
      q_cuseqlen = torch.tensor(
          [0, 76, 110, 256], device=device, dtype=torch.int32
      )
      k_cuseqlen = torch.tensor(
          [0, 1, 2, 3], device=device, dtype=torch.int32
      )
      Mq = 256
      Mk = 3
  
      q = torch.randn([Mq, nheads, d], dtype=dtype, device=device) * 3
      k, v = [
          torch.randn([Mk, nheads, d], dtype=dtype, device=device) * 3
          for _ in range(2)
      ]
      q.requires_grad_(True)
      k.requires_grad_(True)
      v.requires_grad_(True)
  
      out = flash_attn_varlen_func(
          q, k, v, q_cuseqlen, k_cuseqlen, Mq, Mk, causal=causal
      )
      g = torch.randn_like(out)
      out.backward(g)
  
      assert not q.grad.isnan().any()
      assert not k.grad.isnan().any()
      assert not v.grad.isnan().any()
  
imports: import math
  import pytest
  import torch
  import torch.nn.functional
  from einops import rearrange, repeat
  from flash_attn import flash_attn_func, flash_attn_kvpacked_func, flash_attn_qkvpacked_func, flash_attn_varlen_func, flash_attn_varlen_kvpacked_func, flash_attn_varlen_qkvpacked_func, flash_attn_with_kvcache
  from flash_attn.bert_padding import pad_input, unpad_input
  from flash_attn.flash_attn_interface import _get_block_size_n
  from flash_attn.layers.rotary import apply_rotary_emb
label_desc: Use image node_id 24 for calling a global function with example usage: test_flash_attn_bwd_varlen_overflow(d, causal, dtype) without return types