Functions — github.com/WaveSpeedAI/QuantumAttention

hub / github.com/WaveSpeedAI/QuantumAttention / functions

Functions76 in github.com/WaveSpeedAI/QuantumAttention

⨍ Functions 76 ◇ Types & classes 3

↓ 7 callersFunctionrequire_dense_memory

(x, num_dims=None)

src/quantum_attn/inductor/kernels/mm_common.py:14

↓ 5 callersFunctioncan_use_attention

( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: Optional[torch.Tenso

src/quantum_attn/nn.py:282

↓ 4 callersFunction_dynamically_quantize_fp8

(t: torch.Tensor, *, reduction_dim=-1)

src/quantum_attn/nn.py:14

↓ 3 callersFunction_attention_wrapper

( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: Optional[torch.Tenso

src/quantum_attn/nn.py:310

↓ 3 callersFunction_fp8_attention_wrapper

( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: Optional[torch.Tenso

src/quantum_attn/nn.py:394

↓ 2 callersFunction_attention_forward

( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: Optional[torch.Tenso

src/quantum_attn/ops.py:17

↓ 2 callersFunction_fp8_attention_forward

( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, scale_q: Optional[torch.Tensor]

src/quantum_attn/ops.py:64

↓ 2 callersFunction_test_attn_func

(B, H, S_Q, S_KV, D, dtype, device, is_causal, force_eager_fallback, is_fp8=False)

tests/test_interface.py:32

↓ 2 callersFunction_test_benchmark_attn_func

(D, dtype, device, is_causal, is_fp8=False)

tests/test_interface.py:90

↓ 2 callersFunctionacc_type

(dtype)

src/quantum_attn/inductor/kernels/mm_common.py:36

↓ 2 callersFunctionattn_func

( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: Optional[torch.Tenso

src/quantum_attn/quantum_attn_interface.py:41

↓ 2 callersFunctionflash_attention

(query, key, value, is_causal=False)

tests/test_interface.py:13

↓ 2 callersFunctionfp8_attention

( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: Optional[torch.Tenso

src/quantum_attn/nn.py:433

↓ 2 callersFunctionhas_triton_package

()

src/quantum_attn/utils/checks.py:23

↓ 2 callersFunctionload_tk_attention_module

(dtype, is_fp8=False)

src/quantum_attn/tk/attention.py:639

↓ 1 callersFunction_define_composite_implicit_autograd_op

(namespace, name, signature)

src/quantum_attn/quantum_attn_interface.py:21

↓ 1 callersFunction_pre_check_can_use_tk_tma_attention

(device)

src/quantum_attn/nn.py:209

↓ 1 callersFunction_pre_check_can_use_triton_tma_attention

(device)

src/quantum_attn/nn.py:244

↓ 1 callersFunction_tk_tma_supported_head_dim

(n: Union[int, torch.SymInt])

src/quantum_attn/nn.py:48

↓ 1 callersFunction_triton_tma_sdpa_supported_head_dim

Returns true if the head dim is supported by FlexAttention

src/quantum_attn/nn.py:130

↓ 1 callersFunction_validate_tk_tma_input

( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: Optional[torch.Tenso

src/quantum_attn/nn.py:52

↓ 1 callersFunction_validate_triton_tma_sdpa_input

( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: Optional[torch.Tenso

src/quantum_attn/nn.py:135

↓ 1 callersFunctionattention

( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: Optional[torch.Tenso

src/quantum_attn/nn.py:325

↓ 1 callersFunctionattention_fn

()

tests/test_interface.py:104

↓ 1 callersFunctionattention_heuristic_configs

( head_dim, B, H, N_CTX_Q, N_CTX_K, is_causal=False, layout=None, optimize_blo

src/quantum_attn/inductor/kernels/attention.py:673

↓ 1 callersFunctioncudnn_sdpa

(query, key, value, is_causal=False)

tests/test_interface.py:18

↓ 1 callersFunctionearly_attention_config_prune

(configs, query, key, value)

src/quantum_attn/inductor/kernels/attention.py:765

↓ 1 callersFunctionfetch_requirements

()

setup.py:47

↓ 1 callersFunctionfp8_attention

(query, key, value, is_causal=False)

tests/test_interface.py:27

↓ 1 callersFunctionfp8_attn_func

( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: Optional[torch.Tenso

src/quantum_attn/quantum_attn_interface.py:101

↓ 1 callersFunctionfp8_token_wise_attn_func

( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: Optional[torch.Tenso

src/quantum_attn/quantum_attn_interface.py:179

↓ 1 callersFunctiongenerate_attention_template_choices

( choices, query, key, value, scale_q=None, scale_k=None, attn_mask=None, drop

src/quantum_attn/inductor/kernels/attention.py:799

↓ 1 callersFunctionget_attention_layout

( query, dtype, )

src/quantum_attn/inductor/kernels/attention.py:788

↓ 1 callersFunctionget_device_shared_memory

(device=0)

src/quantum_attn/inductor/kernels/mm_common.py:49

↓ 1 callersFunctionhas_triton_language

(attr)

src/quantum_attn/utils/checks.py:40

↓ 1 callersFunctionis_git_directory

(path=".")

setup.py:21

↓ 1 callersFunctionis_nvidia_cuda

()

src/quantum_attn/utils/checks.py:57

↓ 1 callersFunctionmm_options

(c, sym_m, sym_n, sym_k, layout, b_prologue_cast_type=None, optimize_block_size=True)

src/quantum_attn/inductor/kernels/mm_common.py:60

↓ 1 callersFunctionreduce_block_size_for_cuda

(BLOCK_M, BLOCK_N, m, n, device=None, b=1)

src/quantum_attn/inductor/kernels/mm_common.py:84

↓ 1 callersFunctiontorch_cuda_version

()

src/quantum_attn/utils/checks.py:67

↓ 1 callersFunctiontuned_attention_forward

( query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, *, scale=

src/quantum_attn/inductor/kernels/attention.py:888

↓ 1 callersFunctionvanilla_attention

(query, key, value, is_causal=False)

tests/test_interface.py:23

Function_

( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: Optional[torch.Tenso

src/quantum_attn/ops.py:49

Function_define_quantum_attn_composite_implicit_autograd_op

(name, signature)

src/quantum_attn/quantum_attn_interface.py:33

Functionaten_attention_forward_kernel

( query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, *, scale=

src/quantum_attn/inductor/kernels/attention.py:83

Functionaten_fp8_attention_forward_kernel

( query, key, value, scale_q, scale_k, attn_mask=None, dropout_p=0.0, is_causa

src/quantum_attn/inductor/kernels/attention.py:111

Functionattention_forward

( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: Optional[torch.Tenso

src/quantum_attn/ops.py:33

Functionattn_func_with_fallback

( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: Optional[torch.Tenso

src/quantum_attn/quantum_attn_interface.py:66

Functioncan_use_tk_tma_attention

( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: Optional[torch.Tenso

src/quantum_attn/nn.py:219

Functioncan_use_triton_tma_attention

( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: Optional[torch.Tenso

src/quantum_attn/nn.py:258

Functioncuda_capability_compare

(op, major, minor, *, device=None)

src/quantum_attn/utils/checks.py:61

Functioncudnn_sdpa_fn

()

tests/test_interface.py:118

Functiondecorator

(fn)

src/quantum_attn/quantum_attn_interface.py:22

Functiondynamically_quantize_fp8

(t: torch.Tensor, *, reduction_dim=-1)

src/quantum_attn/nn.py:22

Functionfa_fn

()

tests/test_interface.py:115

Functionfp8_attention_forward

( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, scale_q: Optional[torch.Tensor]

src/quantum_attn/ops.py:99

Functionfp8_attention_forward

( query, key, value, scale_q=None, scale_k=None, attn_mask=None, dropout_p=0.0,

src/quantum_attn/inductor/kernels/attention.py:1041

Functionfp8_attn_func_with_fallback

( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: Optional[torch.Tenso

src/quantum_attn/quantum_attn_interface.py:134

Functionfp8_token_wise_attn_func_with_fallback

( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: Optional[torch.Tenso

src/quantum_attn/quantum_attn_interface.py:209

Functionget_constant_attr

(module, attr)

src/quantum_attn/utils/checks.py:10

Functionget_tk_include_dir

()

src/quantum_attn/tk/utils.py:6

Functionhas_triton_tma_support

()

src/quantum_attn/utils/checks.py:48

Functionis_8bit_type

(dtype)

src/quantum_attn/utils/types.py:5

Functionis_fp8_type

(dtype)

src/quantum_attn/utils/types.py:1

Functionmy_local_scheme

(version)

setup.py:38

Functionpersistent_attention_grid

(b, h, s, d, meta)

src/quantum_attn/inductor/kernels/attention.py:143

Functionsdpa_dispatcher

(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, *, scale=None)

src/quantum_attn/quantum_attn_interface.py:37

Functiontest_attn_func

(B, H, S_Q, S_KV, D, dtype, device, is_causal, force_eager_fallback)

tests/test_interface.py:72

Functiontest_benchmark_attn_func

(D, dtype, device, is_causal, is_fp8=False)

tests/test_interface.py:145

Functiontest_benchmark_fp8_attn_func

(D, dtype, device, is_causal)

tests/test_interface.py:153

Functiontest_fp8_attn_func

(B, H, S_Q, S_KV, D, dtype, device, is_causal, force_eager_fallback)

tests/test_interface.py:85

Functiontk_attention_forward_kernel

( query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, *, scale=

src/quantum_attn/inductor/kernels/attention.py:29

Functiontk_fp8_attention_forward_kernel

( query, key, value, scale_q, scale_k, attn_mask=None, dropout_p=0.0, is_causa

src/quantum_attn/inductor/kernels/attention.py:55

Functiontorch_cuda_version_compare

(op, major, minor)

src/quantum_attn/utils/checks.py:74

Functiontorch_version_compare

(op, v)

src/quantum_attn/utils/checks.py:18

Functiontriton_version_compare

(op, v)

src/quantum_attn/utils/checks.py:32