↓ 2 callersFunction_profile_case(label, fn, q, k, v, cumseq_q, cumseq_kv, max_q, max_kv)
tests/test_varlen_recompile.py:127
↓ 2 callersFunction_run_flex_attn(
q: Tensor,
k: Tensor,
v: Tensor,
block_mask: BlockMask,
scale: float,
torch_compile:
src/natten/backends/flex.py:88
↓ 2 callersMethod_test_permute_unpermute_torch(
self, B, H, S, D, tile_shape, dilation, flip_dims, eps, dtype, device="cuda"
)
tests/test_token_permute.py:48
↓ 2 callersFunctionchoose_fmha_backend(
query: Tensor,
key: Tensor,
value: Tensor,
is_causal: bool,
is_varlen: bool,
torch_c
src/natten/backends/__init__.py:113
↓ 2 callersFunctionget_compatible_fmha_backends(
query: Tensor,
key: Tensor,
value: Tensor,
is_causal: bool,
is_varlen: bool,
torch_c
src/natten/backends/__init__.py:175