MCPcopy Create free account

hub / github.com/deepseek-ai/FlashMLA / functions

Functions363 in github.com/deepseek-ai/FlashMLA

⨍ Functions 363◇ Types & classes 181
Methodget_grid_shape
csrc/sm100/prefill/dense/kernel/fmha_causal_tile_scheduler.hpp:78
Methodget_grid_shape
csrc/sm100/prefill/dense/kernel/fmha_causal_tile_scheduler.hpp:168
Methodget_grid_shape
csrc/sm100/prefill/dense/kernel/sm100_fmha_bwd_mla_kernel_tma_warpspecialized.hpp:1821
Methodget_grid_shape
csrc/sm100/prefill/dense/kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp:232
Methodget_grid_shape
csrc/sm100/prefill/dense/kernel/fmha_kernel_bwd_convert.hpp:92
Methodget_grid_shape
csrc/sm100/prefill/dense/kernel/sm100_fmha_bwd_kernel_tma_warpspecialized.hpp:1827
Methodget_grid_shape
csrc/sm100/prefill/dense/kernel/fmha_kernel_bwd_sum_OdO.hpp:93
Methodget_grid_shape
csrc/sm100/prefill/dense/kernel/fmha_tile_scheduler.hpp:64
Methodget_grid_shape
csrc/sm100/prefill/dense/kernel/fmha_tile_scheduler.hpp:131
Methodget_masked_trip_count
csrc/sm100/prefill/dense/collective/fmha_fusion.hpp:88
Methodget_masked_trip_count
csrc/sm100/prefill/dense/collective/fmha_fusion.hpp:140
Methodget_masked_trip_count
csrc/sm100/prefill/dense/collective/fmha_fusion.hpp:218
Functionget_max
csrc/sm100/prefill/sparse/common_subroutine.h:171
Methodget_meta
csrc/api/sparse_decode.h:92
Methodget_meta
csrc/api/sparse_decode.h:126
Methodget_meta
csrc/api/sparse_decode.h:168
Functionget_optional_tensor_ptr
csrc/kerutils/include/kerutils/supplemental/torch_tensors.h:41
Functionget_peer_addr
csrc/sm90/helpers.h:144
Functionget_peer_addr
csrc/sm90/decode/sparse_fp8/components/helpers.h:104
Methodget_problem_shape_O
csrc/sm100/prefill/dense/collective/sm100_fmha_fwd_epilogue_tma_warpspecialized.hpp:96
Functionget_s_from_p
csrc/sm100/prefill/sparse/common_subroutine.h:187
Functionget_sm_id
csrc/sm90/helpers.h:136
Functionget_static_enum_name
csrc/api/common.h:103
Functionget_tensor_ptr
csrc/kerutils/include/kerutils/supplemental/torch_tensors.h:30
Methodget_trip_count
csrc/sm100/prefill/dense/collective/fmha_fusion.hpp:198
Methodget_unmasked_trip_count
csrc/sm100/prefill/dense/collective/fmha_fusion.hpp:100
Methodget_unmasked_trip_count
csrc/sm100/prefill/dense/collective/fmha_fusion.hpp:152
Methodget_unmasked_trip_count
csrc/sm100/prefill/dense/collective/fmha_fusion.hpp:234
Methodget_workspace_size
Gets the workspace size
csrc/sm100/prefill/dense/device/fmha.hpp:97
Methodget_workspace_size
Gets the workspace size
csrc/sm100/prefill/dense/device/fmha_device_bwd.hpp:220
Methodget_workspace_size
csrc/sm100/prefill/dense/kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp:223
Methodget_workspace_size
csrc/sm100/prefill/dense/kernel/fmha_kernel_bwd_convert.hpp:77
Methodget_workspace_size
csrc/sm100/prefill/dense/kernel/fmha_kernel_bwd_sum_OdO.hpp:76
Methodinit_barriers
Helper function to initialize barriers
csrc/sm100/prefill/dense/common/pipeline_mla.hpp:72
Methodinitialize
Initializes GEMM state from arguments.
csrc/sm100/prefill/dense/device/fmha.hpp:153
Methodinitialize_split
Initializes state from arguments.
csrc/sm100/prefill/dense/device/fmha_device_bwd.hpp:237
Methodinitialize_workspace
csrc/sm100/prefill/dense/kernel/sm100_fmha_bwd_mla_kernel_tma_warpspecialized.hpp:367
Methodinitialize_workspace
csrc/sm100/prefill/dense/kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp:224
Methodinitialize_workspace
csrc/sm100/prefill/dense/kernel/fmha_kernel_bwd_convert.hpp:78
Methodinitialize_workspace
csrc/sm100/prefill/dense/kernel/sm100_fmha_bwd_kernel_tma_warpspecialized.hpp:360
Methodinitialize_workspace
csrc/sm100/prefill/dense/kernel/fmha_kernel_bwd_sum_OdO.hpp:77
Functionint4_max
csrc/sm100/helpers.h:13
Functionint4_min
csrc/sm100/helpers.h:18
Methodis_initialized
csrc/sm100/prefill/dense/device/fmha.hpp:72
Functionis_no_cooldown
()
tests/lib.py:404
Methodis_valid
csrc/sm100/prefill/dense/kernel/fmha_causal_tile_scheduler.hpp:82
Methodis_valid
csrc/sm100/prefill/dense/kernel/fmha_causal_tile_scheduler.hpp:173
Methodis_valid
csrc/sm100/prefill/dense/kernel/fmha_tile_scheduler.hpp:136
Functionlaunch_tma_copy
csrc/sm90/helpers.h:154
Functionlaunch_tma_copy
csrc/sm90/decode/sparse_fp8/components/helpers.h:60
Methodload
csrc/sm100/prefill/dense/collective/sm100_fmha_fwd_mainloop_tma_warpspecialized.hpp:242
Methodload
csrc/sm100/prefill/dense/collective/sm100_fmha_mla_fwd_mainloop_tma_warpspecialized.hpp:260
Methodload
csrc/sm100/prefill/dense/collective/sm100_fmha_mla_load_tma_warpspecialized.hpp:157
Methodload
csrc/sm100/prefill/dense/kernel/sm100_fmha_bwd_mla_kernel_tma_warpspecialized.hpp:439
Functionload_128b_from_gmem
csrc/sm90/decode/sparse_fp8/components/dequant.h:54
Functionload_64b_from_gmem
csrc/sm90/decode/sparse_fp8/components/dequant.h:93
Functionload_indices_and_generate_mask
Load K/V indices from global memory, and generate validity mask Each thread loads 8 indices Should be called by lanes 0 ~ (BLOCK_TOPK/8) */
csrc/sm100/prefill/sparse/common_subroutine.h:13
Functionmake_gather_tensor
csrc/sm100/prefill/dense/common/gather_tensor.hpp:155
Methodmake_layout
csrc/sm100/prefill/dense/common/gather_tensor.hpp:130
Functionmake_stride_helper
csrc/kerutils/include/kerutils/host/host.h:147
Functionmake_tensor_map
A wrapper for make_tensor_map
csrc/kerutils/include/kerutils/host/host.h:82
Methodmaximum_active_blocks
Computes the maximum number of active blocks per multiprocessor
csrc/sm100/prefill/dense/device/fmha.hpp:111
Methodmma
csrc/sm100/prefill/dense/collective/sm100_fmha_fwd_mainloop_tma_warpspecialized.hpp:258
Methodmma
csrc/sm100/prefill/dense/collective/sm100_fmha_mla_fwd_mainloop_tma_warpspecialized.hpp:276
Methodmma
csrc/sm100/prefill/dense/kernel/sm100_fmha_bwd_mla_kernel_tma_warpspecialized.hpp:667
Functionnon_contiguousify
(tensor: torch.Tensor)
tests/kernelkit/generate.py:22
Functionnv_bfloat162 fp8x2_to_bf16x2_with_scale
Convert 2x fp8_e4m3 to 2x bf16 with scaling
csrc/sm100/helpers.h:24
Methodoffset_by
csrc/utils.h:75
Methodoperator *
csrc/sm100/prefill/dense/common/pow_2.hpp:52
Functionoperator%
csrc/sm100/prefill/dense/common/pow_2.hpp:72
Methodoperator()
Launches the kernel after first constructing Params internal state from supplied arguments.
csrc/sm100/prefill/dense/device/fmha.hpp:259
Methodoperator()
csrc/sm100/prefill/dense/common/gather_tensor.hpp:56
Methodoperator()
csrc/sm100/prefill/dense/common/gather_tensor.hpp:78
Methodoperator()
csrc/sm100/prefill/dense/kernel/sm100_fmha_bwd_mla_kernel_tma_warpspecialized.hpp:1483
Methodoperator()
csrc/sm100/prefill/dense/kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp:254
Methodoperator()
csrc/sm100/prefill/dense/kernel/fmha_kernel_bwd_convert.hpp:141
Methodoperator()
csrc/sm100/prefill/dense/kernel/sm100_fmha_bwd_kernel_tma_warpspecialized.hpp:1489
Methodoperator()
csrc/sm100/prefill/dense/kernel/fmha_kernel_bwd_sum_OdO.hpp:107
Methodoperator*
csrc/sm100/prefill/dense/common/gather_tensor.hpp:101
Functionoperator/
csrc/sm100/prefill/dense/common/pow_2.hpp:67
Functionoperator<
csrc/sm100/prefill/dense/common/pow_2.hpp:77
Functionoptional_cast_to_bf16_and_cast_back
(tensor: torch.Tensor)
tests/kernelkit/precision.py:23
Methodpair<uint32_t, bool> get
csrc/utils.h:67
Methodprefetch_tma_descriptors
csrc/sm100/prefill/dense/collective/sm100_fmha_fwd_mainloop_tma_warpspecialized.hpp:236
Methodprefetch_tma_descriptors
csrc/sm100/prefill/dense/collective/sm100_fmha_fwd_epilogue_tma_warpspecialized.hpp:145
Methodprefetch_tma_descriptors
csrc/sm100/prefill/dense/collective/sm100_fmha_load_tma_warpspecialized.hpp:141
Methodprefetch_tma_descriptors
csrc/sm100/prefill/dense/collective/sm100_fmha_mla_fwd_mainloop_tma_warpspecialized.hpp:254
Methodprefetch_tma_descriptors
csrc/sm100/prefill/dense/collective/sm100_fmha_mla_load_tma_warpspecialized.hpp:149
Functionprint
csrc/sm100/prefill/dense/collective/fmha_fusion.hpp:391
Methodprint
csrc/sm100/prefill/dense/common/gather_tensor.hpp:61
Methodprint
csrc/sm100/prefill/dense/common/gather_tensor.hpp:83
Functionproducer_acquire
csrc/sm100/prefill/dense/common/pipeline_mla.hpp:171
Functionproducer_acquire_bytes
csrc/sm100/prefill/dense/common/pipeline_mla.hpp:176
Functionprofiler_range_start_marker_kernel
()
tests/kernelkit/bench.py:17
Methodquantize
csrc/sm100/prefill/dense/kernel/sm100_fmha_bwd_mla_kernel_tma_warpspecialized.hpp:421
Methodquantize
csrc/sm100/prefill/dense/kernel/sm100_fmha_bwd_kernel_tma_warpspecialized.hpp:414
Functionquantize_k_cache
Quantize the k-cache For more detail about the layout of K/V, please refer to comments in flash_mla_interface.py
tests/quant.py:20
Methodreduce
csrc/sm100/prefill/dense/kernel/sm100_fmha_bwd_mla_kernel_tma_warpspecialized.hpp:1386
Methodreduce
csrc/sm100/prefill/dense/kernel/sm100_fmha_bwd_kernel_tma_warpspecialized.hpp:1392
Functionref_sparse_attn_decode
A reference implementation of sparse decoding attention in PyTorch
tests/ref.py:55
← previousnext →201–300 of 363, ranked by callers