MCPcopy Create free account

hub / github.com/SHI-Labs/NATTEN / functions

Functions2,059 in github.com/SHI-Labs/NATTEN

⨍ Functions 2,059◇ Types & classes 573
↓ 7 callersMethodenable_mask
Clears the predicate set efficiently
csrc/include/natten/cuda/fmha/iterators/predicated_tile_iterator_residual_last.h:864
↓ 7 callersMethodget_mask
Gets the mask
csrc/include/natten/cuda/fmha/iterators/predicated_tile_iterator_residual_last.h:876
↓ 7 callersMethodset_mask
Sets the predicate mask, overriding value stored in predicate iterator
csrc/include/natten/cuda/fmha/iterators/predicated_tile_iterator_residual_last.h:870
↓ 6 callersMethod_test_all_dtypes_against_cutlass_2x_fna
( self, batch, heads, head_dim, input_shape, kernel_size,
tests/test_hopper_fna.py:69
↓ 6 callersMethod_test_all_dtypes_against_cutlass_2x_fna
( self, batch, heads, head_dim, input_shape, kernel_size,
tests/test_blackwell_fna.py:70
↓ 6 callersMethod_test_all_dtypes_against_reference
( self, batch, heads, head_dim, head_dim_v, input_shape,
tests/test_fna.py:72
↓ 6 callersMethod_test_determinism
( self, batch, heads, head_dim, input_shape, kernel_size,
tests/test_fna.py:565
↓ 6 callersMethod_test_randsweep
(self, backend, num_tests=1000)
tests/test_attn_merge.py:200
↓ 6 callersMethod_test_randsweep_against_cutlass_2x
( self, na_dim, max_tests=1000, configs_to_test=None )
tests/test_hopper_fna.py:453
↓ 6 callersMethod_test_randsweep_against_cutlass_2x
( self, na_dim, max_tests=1000, configs_to_test=5 )
tests/test_blackwell_fna.py:445
↓ 6 callersMethod_test_randsweep_against_reference
( self, na_dim, max_tests=1000, configs_to_test=10, quick=False, )
tests/test_fna.py:459
↓ 6 callersMethod_test_varlen_backend
( self, batch, heads, head_dim, seqlens_Q_list, seqlens_KV_lis
tests/test_fmha_varlen.py:595
↓ 6 callersMethodapply_mask
csrc/include/natten/cuda/fmha_blackwell/collective/fmha_fusion.hpp:67
↓ 6 callersFunctioncan_run_cutlass_blackwell_fmha
( query: Tensor, key: Tensor, value: Tensor, is_causal: bool, is_varlen: bool, raise_e
src/natten/backends/configs/checks.py:39
↓ 6 callersFunctioncan_run_cutlass_blackwell_fna
( query: Tensor, key: Tensor, value: Tensor, raise_error: bool = False )
src/natten/backends/configs/checks.py:143
↓ 6 callersFunctioncan_run_cutlass_fmha
( query: Tensor, key: Tensor, value: Tensor, is_causal: bool, is_varlen: bool, raise_e
src/natten/backends/configs/checks.py:433
↓ 6 callersFunctioncan_run_cutlass_fna
( query: Tensor, key: Tensor, value: Tensor, raise_error: bool = False )
src/natten/backends/configs/checks.py:514
↓ 6 callersFunctioncan_run_cutlass_hopper_fmha
( query: Tensor, key: Tensor, value: Tensor, is_causal: bool, is_varlen: bool, raise_e
src/natten/backends/configs/checks.py:254
↓ 6 callersFunctioncan_run_cutlass_hopper_fna
( query: Tensor, key: Tensor, value: Tensor, raise_error: bool = False )
src/natten/backends/configs/checks.py:344
↓ 6 callersFunctioncheck_args_against_input
( input_tensor: Tensor, kernel_size: DimensionType, stride: DimensionType, dilation: Dimension
src/natten/utils/checks.py:486
↓ 6 callersFunctioncreate_dim_from_int
(na_dim: int, value: int)
src/natten/utils/tuples.py:46
↓ 6 callersFunctionget_all_backward_configs
( input_tensor: Tensor, )
src/natten/backends/configs/cutlass_hopper/__init__.py:196
↓ 6 callersFunctionget_all_forward_configs
( input_tensor: Tensor, )
src/natten/backends/configs/cutlass_hopper/__init__.py:170
↓ 6 callersFunctionincrIteration
Returns the next block to process
csrc/include/natten/cuda/fna/kernel_backward.h:2429
↓ 6 callersFunctionis_fully_block_sparse_int
csrc/include/natten/cuda/fna_hopper/collective/fna_fusion.hpp:437
↓ 6 callersFunctionis_fully_block_sparse_int
csrc/include/natten/cuda/fna_blackwell/collective/fna_fusion.hpp:503
↓ 6 callersMethodmake_qkv_tensors
( self, device, requires_grad: bool, heads_last: bool, flatten: bool = False )
src/natten/profiling_utils/problem.py:140
↓ 6 callersMethodnum_splits_key_device
csrc/include/natten/cuda/fna/kernel_backward.h:714
↓ 6 callersFunctionqkv_fix_dilation
csrc/include/natten/cuda/reference/mask.hpp:220
↓ 6 callersFunctionqkv_fix_dilation
csrc/include/natten/cuda/fna_hopper/collective/fna_fusion.hpp:204
↓ 6 callersFunctionqkv_fix_dilation
csrc/include/natten/cuda/fna_blackwell/collective/fna_fusion.hpp:250
↓ 6 callersFunctionrun_ops
( query, key, value, d_out, window_size, stride, dilat
src/natten/profiling_utils/profiling.py:105
↓ 6 callersFunctionset_memory_usage_preference
Sets memory usage preference for KV parallelism in `"cutlass-fna"` and `"cutlass-fmha"` backends. Args: pref: Choices are `"default"`
src/natten/context.py:54
↓ 6 callersMethodset_predicates
csrc/include/natten/cuda/fna/iterators/predicated_tile_access_iterator.h:338
↓ 6 callersMethodtest
( self, eps: Union[ float, Tuple[float, float], Tuple[float, Tuple[float, float, f
tests/utils.py:199
↓ 6 callersFunctiontuple_leq
csrc/include/natten/cuda/fna_hopper/collective/fna_common.hpp:43
↓ 6 callersFunctionuse_kv_parallelism_in_fused_na
Sets guards for using KV Parallelism in backpropagation in `"cutlass-fna"`/`"cutlass-fmha"` backends. Warning: Disabling KV paralleli
src/natten/context.py:113
↓ 6 callersFunctionwrite_header_file
(content, path, namespaces, extra_includes=None)
scripts/autogen_hopper_fna_bwd.py:547
↓ 6 callersFunctionwrite_header_file
(content, path, namespaces, extra_includes=None)
scripts/autogen_blackwell_fna.py:539
↓ 6 callersFunctionwrite_header_file
(content, path, namespaces, extra_includes=None)
scripts/autogen_blackwell_fna_bwd.py:538
↓ 6 callersFunctionwrite_header_file
(content, path, namespaces, extra_includes=None)
scripts/autogen_hopper_fna.py:589
↓ 5 callersFunction_reset_everything
(random_seed: int = 42, torch_seed: int = 42)
tests/test_token_permute.py:34
↓ 5 callersMethod_test_randsweep
( self, backend, max_tests=1000, reference: str = "sdpa", determinism:
tests/test_fmha.py:1016
↓ 5 callersMethod_test_varlen_randsweep
( self, backend, max_tests=1000, determinism: bool = False )
tests/test_fmha_varlen.py:658
↓ 5 callersMethodadd_pointer_offset
Adds a pointer offset in units of Element
csrc/include/natten/cuda/fmha/iterators/predicated_tile_access_iterator_residual_last.h:816
↓ 5 callersFunctionallow_flex_compile
Sets guards for Flex Attention + `torch.compile`. Allows using our Flex FNA / Flex FMHA backends with `torch.compile`, meaning you can pass `
src/natten/context.py:161
↓ 5 callersFunctionapply_variable_length_scheduler
csrc/include/natten/cuda/fmha_hopper/collective/fmha_varlen.hpp:74
↓ 5 callersMethodcorrect_qkv_shape
csrc/include/natten/cuda/fna_blackwell/collective/fna_fusion.hpp:293
↓ 5 callersFunctionformat_elapsed
(seconds)
scripts/testing/monitor.py:51
↓ 5 callersMethodgK_strideM
csrc/include/natten/cuda/fmha/kernel_backward.h:642
↓ 5 callersMethodgV_strideM
csrc/include/natten/cuda/fmha/kernel_backward.h:645
↓ 5 callersFunctiongemm_zero_acc
csrc/include/natten/cuda/fmha_blackwell/collective/fmha_common.hpp:61
↓ 5 callersFunctiongemm_zero_acc
csrc/include/natten/cuda/fna_blackwell/collective/fna_common.hpp:61
↓ 5 callersFunctiongetSmallestQueryForKey
csrc/include/natten/cuda/fna/kernel_backward.h:2347
↓ 5 callersMethodget_block_coord
csrc/include/natten/cuda/fmha_blackwell/kernel/fmha_tile_scheduler.hpp:78
↓ 5 callersMethodget_dispatcher
(self)
scripts/autogen_hopper_fna_bwd.py:278
↓ 5 callersMethodget_dispatcher
(self)
scripts/autogen_fna.py:314
↓ 5 callersMethodget_dispatcher
(self)
scripts/autogen_blackwell_fna.py:269
↓ 5 callersMethodget_dispatcher
(self)
scripts/autogen_blackwell_fna_bwd.py:273
↓ 5 callersMethodget_dispatcher
(self)
scripts/autogen_hopper_fna.py:313
↓ 5 callersMethodget_trip_count
csrc/include/natten/cuda/fna_blackwell/collective/fna_fusion.hpp:315
↓ 5 callersMethodis_valid
csrc/include/natten/cuda/fmha_blackwell/kernel/fmha_tile_scheduler.hpp:73
↓ 5 callersMethodload_with_pointer_offset
Loads a fragment from memory
csrc/include/natten/cuda/fmha/iterators/predicated_tile_iterator_residual_last.h:882
↓ 5 callersMethodoperand_B_ref
Returns a TensorRef to the B operand
csrc/include/natten/cuda/fmha/gemm/mma_from_smem.h:233
↓ 5 callersMethodoperand_B_ref
Returns a TensorRef to the B operand
csrc/include/natten/cuda/fna/gemm/mma_from_smem.h:234
↓ 5 callersFunctionprologue
csrc/include/natten/cuda/fmha/gemm/mma_from_smem.h:536
↓ 5 callersFunctionprologue
csrc/include/natten/cuda/fna/gemm/mma_from_smem.h:537
↓ 5 callersMethodset_residual_tile
csrc/include/natten/cuda/fmha/iterators/predicated_tile_access_iterator_residual_last.h:810
↓ 5 callersMethodstore_with_pointer_offset
Store a fragment to memory
csrc/include/natten/cuda/fmha/iterators/predicated_tile_iterator_residual_last.h:900
↓ 5 callersFunctionsupports_float16
(device: torch.device)
src/natten/utils/testing.py:126
↓ 5 callersFunctionvarlen_tensor_checks
( query: Tensor, key: Tensor, value: Tensor, seqlens_Q: Optional[Tensor] = None, seqlens_K
src/natten/utils/checks.py:569
↓ 5 callersFunctionwarp_uniform
csrc/include/natten/cuda/fmha_blackwell/collective/fmha_common.hpp:79
↓ 5 callersFunctionwrite_header_file
(content, path, namespaces, extra_includes=None)
scripts/autogen_fna.py:527
↓ 4 callersFunction_prepare_case
(seqlens_list, total_seqlen, heads, head_dim)
tests/test_varlen_recompile.py:171
↓ 4 callersFunction_read_log_file
(path: str)
tests/test_log.py:75
↓ 4 callersMethod_test_backend_against_torch_sdpa
( self, batch, heads, head_dim, seqlen_q, seqlen_kv, i
tests/test_fmha.py:901
↓ 4 callersFunctionalign_up
csrc/include/natten/cuda/fmha/gemm_kernel_utils.h:123
↓ 4 callersMethodapply_mask
csrc/include/natten/cuda/fna_hopper/collective/fna_fusion.hpp:310
↓ 4 callersMethodapply_padded_mask
csrc/include/natten/cuda/fna_hopper/collective/fna_fusion.hpp:382
↓ 4 callersMethodbefore_softmax
csrc/include/natten/cuda/fmha_hopper/collective/fmha_fusion.hpp:152
↓ 4 callersFunctioncan_run_cutlass_tokperm
(tensor: Tensor)
src/natten/token_permute/cutlass_impl.py:51
↓ 4 callersMethodclear_mask
Clears the predicate set efficiently
csrc/include/natten/cuda/fna/iterators/predicated_tile_access_iterator.h:444
↓ 4 callersMethodcompute_predicates_
Computes predicates based on internally tracked per-thread offset.
csrc/include/natten/cuda/fna/iterators/predicated_tile_access_iterator.h:262
↓ 4 callersFunctionconvert_c_layout_to_a_layout
csrc/include/natten/cuda/fna_hopper/collective/fna_common.hpp:272
↓ 4 callersFunctionconvert_c_layout_to_a_layout
csrc/include/natten/cuda/fmha_hopper/collective/fmha_common.hpp:229
↓ 4 callersMethodcorrect_qkv_shape
csrc/include/natten/cuda/fna_hopper/collective/fna_fusion.hpp:247
↓ 4 callersFunctioncutlass_blackwell_fna_generic
( query: Tensor, key: Tensor, value: Tensor, kernel_size: DimensionTypeOrDed, stride: Dime
src/natten/backends/blackwell_fna.py:332
↓ 4 callersFunctioncutlass_fna_generic
( query: Tensor, key: Tensor, value: Tensor, kernel_size: DimensionTypeOrDed, stride: Dime
src/natten/backends/fna.py:223
↓ 4 callersFunctioncutlass_hopper_fna_generic
( query: Tensor, key: Tensor, value: Tensor, kernel_size: DimensionTypeOrDed, stride: Dime
src/natten/backends/hopper_fna.py:330
↓ 4 callersMethodenable_mask
Clears the predicate set efficiently
csrc/include/natten/cuda/fna/iterators/predicated_tile_access_iterator.h:453
↓ 4 callersFunctionflex_fna_generic
( query: Tensor, key: Tensor, value: Tensor, kernel_size: DimensionTypeOrDed, stride: Dime
src/natten/backends/flex.py:557
↓ 4 callersFunctiongemm_reset_zero_acc
csrc/include/natten/cuda/fmha_blackwell/collective/fmha_common.hpp:43
↓ 4 callersFunctiongemm_reset_zero_acc
csrc/include/natten/cuda/fna_blackwell/collective/fna_common.hpp:43
↓ 4 callersFunctionget_all_backward_configs
( input_tensor: Tensor, )
src/natten/backends/configs/cutlass_blackwell/__init__.py:118
↓ 4 callersFunctionget_all_fmha_backward_configs
( input_tensor: Tensor, )
src/natten/backends/configs/cutlass_hopper/__init__.py:239
↓ 4 callersFunctionget_all_fmha_forward_configs
( input_tensor: Tensor, )
src/natten/backends/configs/cutlass_hopper/__init__.py:223
↓ 4 callersMethodget_dispatcher
(self)
scripts/autogen_reference_fna.py:301
↓ 4 callersMethodget_mask
Gets the mask
csrc/include/natten/cuda/fna/iterators/predicated_tile_access_iterator_residual_last.h:960
↓ 4 callersMethodget_mask
Gets the mask
csrc/include/natten/cuda/fna/iterators/predicated_tile_access_iterator.h:471
↓ 4 callersMethodget_operand_shape
( self, heads: int, dim: int, heads_last: bool, flatten: bool = False )
src/natten/profiling_utils/problem.py:83
← previousnext →101–200 of 2,059, ranked by callers