MCPcopy Create free account

hub / github.com/OpenNLPLab/lightning-attention / functions

Functions53 in github.com/OpenNLPLab/lightning-attention

⨍ Functions 53◇ Types & classes 7↳ Endpoints 3
↓ 11 callersMethodbackward
( ctx, dy )
lightning_attn/ops/triton/srmsnorm.py:125
↓ 5 callersFunction_build_slope_tensor
(n_attention_heads: int)
lightning_attn/utils/utils.py:6
↓ 2 callersFunctionflash_wrapper
(q, k, v, causal=True)
benchmarks/benchmark_lightning2.py:35
↓ 2 callersFunctionget_memory
(device)
lightning_attn/utils/utils.py:34
↓ 2 callersFunctionget_slopes_power_of_2
(n)
lightning_attn/utils/utils.py:8
↓ 2 callersFunctionis_support
(dim)
lightning_attn/ops/lightning_attn_interface.py:8
↓ 2 callersFunctionlightning_attn_func
(q, k, v, s=None, variant="chunk_loop")
lightning_attn/ops/lightning_attn_interface.py:16
↓ 2 callersFunctionlinear_attn
(q, k, v, s=None)
lightning_attn/ops/torch/linear_attention.py:27
↓ 2 callersFunctionxformer_wrapper
(q, k, v, causal=True)
benchmarks/benchmark_lightning2.py:43
↓ 1 callersMethod_norm
(self, x)
lightning_attn/ops/torch/srmsnorm.py:14
↓ 1 callersFunctionget_full_mask
(n, slopes)
lightning_attn/ops/torch/linear_attention.py:15
↓ 1 callersFunctionget_mask
(n, slope=1)
lightning_attn/ops/torch/linear_attention.py:4
↓ 1 callersFunctionget_slopes
(n)
lightning_attn/utils/utils.py:7
↓ 1 callersFunctionnext_power_of_2
(n)
lightning_attn/ops/lightning_attn_interface.py:12
Method__init__
( self, embed_dim: int, num_heads: int, bias: bool = False, norm_type:
examples/module/tnl_attn.py:13
Method__init__
(self, dim: int, eps: float = 1e-6)
lightning_attn/ops/torch/srmsnorm.py:10
Method__init__
(self, dim: int, eps: float = 1e-6)
lightning_attn/ops/triton/srmsnorm.py:179
Function_bwd_diag_kernel
( Q, K, V, S, DO, DQ, DK, DV, b: tl.constexpr, h: tl.constexpr, n:
lightning_attn/ops/triton/lightning_attn2_parallel.py:382
Function_bwd_dkv_parallel
( Q, DO, S, DKV, b: tl.constexpr, h: tl.constexpr, n: tl.constexpr, d: tl.cons
lightning_attn/ops/triton/lightning_attn2_parallel.py:566
Function_bwd_dkv_reduce
( Q, DO, S, DKV, b: tl.constexpr, h: tl.constexpr, n: tl.constexpr, d: tl.cons
lightning_attn/ops/triton/lightning_attn2_parallel.py:651
Function_bwd_inter_kernel
( Q, K, V, DO, DQ, DK, DV, b: tl.constexpr, h: tl.constexpr, n: tl.con
lightning_attn/ops/triton/lightning_attn2_no_decay.py:173
Function_bwd_inter_kernel
( Q, K, V, S, DO, DQ, DK, DV, b: tl.constexpr, h: tl.constexpr, n:
lightning_attn/ops/triton/lightning_attn2.py:187
Function_bwd_intra_kernel
( Q, K, V, DO, DQ, DK, DV, b: tl.constexpr, h: tl.constexpr, n: tl.con
lightning_attn/ops/triton/lightning_attn2_no_decay.py:81
Function_bwd_intra_kernel
( Q, K, V, S, DO, DQ, DK, DV, b: tl.constexpr, h: tl.constexpr, n:
lightning_attn/ops/triton/lightning_attn2.py:89
Function_bwd_none_diag_kernel
( Q, K, V, S, DO, DQ, DK, DV, KV, DKV, b: tl.constexpr, h: tl.
lightning_attn/ops/triton/lightning_attn2_parallel.py:705
Function_fwd_diag_kernel
( Q, K, V, Out, S, b: tl.constexpr, h: tl.constexpr, n: tl.constexpr, d: t
lightning_attn/ops/triton/lightning_attn2_parallel.py:10
Function_fwd_kernel
( Q, K, V, Out, b: tl.constexpr, h: tl.constexpr, n: tl.constexpr, d: tl.const
lightning_attn/ops/triton/lightning_attn2_no_decay.py:9
Function_fwd_kernel
( Q, K, V, Out, S, # log lambda b: tl.constexpr, h: tl.constexpr, n: tl.const
lightning_attn/ops/triton/lightning_attn2.py:9
Function_fwd_kv_parallel
( K, V, S, KV, b: tl.constexpr, h: tl.constexpr, n: tl.constexpr, d: tl.conste
lightning_attn/ops/triton/lightning_attn2_parallel.py:113
Function_fwd_kv_reduce
( K, V, S, KV, b: tl.constexpr, h: tl.constexpr, n: tl.constexpr, d: tl.conste
lightning_attn/ops/triton/lightning_attn2_parallel.py:199
Function_fwd_none_diag_kernel
( Q, K, V, Out, S, KV, b: tl.constexpr, h: tl.constexpr, n: tl.constexpr,
lightning_attn/ops/triton/lightning_attn2_parallel.py:253
Methodbackward
(ctx, do)
lightning_attn/ops/triton/lightning_attn2_parallel.py:949
Methodbackward
(ctx, do)
lightning_attn/ops/triton/lightning_attn2_no_decay.py:404
Methodbackward
(ctx, do)
lightning_attn/ops/triton/lightning_attn2.py:437
Functionbench_memory
(b, h, n, d, dtype, device, mode, provider)
benchmarks/benchmark_lightning2.py:160
Functionbench_memory
(b, n, d, dtype, device, mode, provider)
benchmarks/benchmark_srmsnorm.py:107
Functionbench_speed
(b, h, n, d, dtype, device, mode, provider)
benchmarks/benchmark_lightning2.py:94
Functionbench_speed
(b, n, d, dtype, device, mode, provider)
benchmarks/benchmark_srmsnorm.py:52
Methodforward
( self, x, **kwargs, )
examples/module/tnl_attn.py:40
Methodforward
(self, x)
lightning_attn/ops/torch/srmsnorm.py:17
Methodforward
(ctx, q, k, v, s)
lightning_attn/ops/triton/lightning_attn2_parallel.py:837
Methodforward
(ctx, q, k, v)
lightning_attn/ops/triton/lightning_attn2_no_decay.py:369
Methodforward
(ctx, q, k, v, s)
lightning_attn/ops/triton/lightning_attn2.py:400
Methodforward
(ctx, x, eps)
lightning_attn/ops/triton/srmsnorm.py:78
Methodforward
(self, x)
lightning_attn/ops/triton/srmsnorm.py:184
Functionget_params
()
tests/ops/test_lightning2.py:8
Functionget_params
()
tests/ops/test_srmsnorm.py:8
Functionget_params
()
tests/ops/test_lightning2_no_decay.py:7
Functionsrms_norm_bwd_dx_fused
( DX, DY, X, V, stride, N, # META-parameters BLOCK_SIZE_N: tl.constexpr, )
lightning_attn/ops/triton/srmsnorm.py:39
Functionsrms_norm_fw
(X, Y, V, stride, N, eps, BLOCK_SIZE_N: tl.constexpr)
lightning_attn/ops/triton/srmsnorm.py:13
Functiontest_lightning2
(b, h, n, d, e, dtype)
tests/ops/test_lightning2.py:36
Functiontest_lightning2
(b, h, n, d, e, dtype)
tests/ops/test_lightning2_no_decay.py:34
Functiontest_srmsnorm
(b, n, d, dtype)
tests/ops/test_srmsnorm.py:20