forked from ccf-ai-infra/TileOPs-Metax
218 lines
6.2 KiB
Python
218 lines
6.2 KiB
Python
from .attention import (
|
|
FlashAttnBwdPostprocessKernel,
|
|
FlashAttnBwdPreprocessKernel,
|
|
GQABwdKernel,
|
|
GQABwdWgmmaPipelinedKernel,
|
|
GQADecodeKernel,
|
|
GQADecodePagedKernel,
|
|
GQAFwdFP8Fa3ContractPtxAccBN224WsTmaVKernel,
|
|
GQAFwdKernel,
|
|
GQAFwdWgmmaPipelinedKernel,
|
|
GQAPrefillFwdKernel,
|
|
GQAPrefillPagedWithFP8KVCacheFwdKernel,
|
|
GQAPrefillPagedWithKVCacheFwdKernel,
|
|
GQAPrefillPagedWithKVCacheRopeAppendKernel,
|
|
GQAPrefillPagedWithKVCacheRopeFwdKernel,
|
|
GQAPrefillWithKVCacheFwdKernel,
|
|
GQAPrefillWithKVCacheRopeAppendKernel,
|
|
GQAPrefillWithKVCacheRopeFwdKernel,
|
|
GQASlidingWindowFwdKernel,
|
|
GQASlidingWindowFwdWgmmaPipelinedKernel,
|
|
GQASlidingWindowVarlenFwdKernel,
|
|
GQASlidingWindowVarlenFwdWgmmaPipelinedKernel,
|
|
MeanPoolingFwdKernel,
|
|
MHABwdKernel,
|
|
MHABwdWgmmaPipelinedKernel,
|
|
MHADecodeKernel,
|
|
MHADecodePagedKernel,
|
|
MHAFwdKernel,
|
|
MHAFwdWgmmaPipelinedKernel,
|
|
MLADecodeKernel,
|
|
MLADecodeWsKernel,
|
|
NSACmpFwdVarlenKernel,
|
|
NSAFwdVarlenKernel,
|
|
NSATopkVarlenKernel,
|
|
SparseMlaKernel,
|
|
)
|
|
from .bmm import BmmFp8Kernel, BmmKernel
|
|
from .convolution import (
|
|
Conv1dKernel,
|
|
Conv1dPointwiseKernel,
|
|
Conv2d1x1Kernel,
|
|
Conv2dKernel,
|
|
Conv3dKernel,
|
|
GroupConv1dKernel,
|
|
GroupConv2dKernel,
|
|
GroupConv3dKernel,
|
|
)
|
|
from .deltanet import DeltaNetBwdKernel, DeltaNetFwdKernel
|
|
from .deltanet_recurrence import (
|
|
DeltaNetDecodeFP32Kernel,
|
|
DeltaNetDecodeKernel,
|
|
DeltaNetDecodeRawCudaFlaStyleKernel,
|
|
)
|
|
from .dropout import DropoutKernel
|
|
from .elementwise import BinaryKernel, FusedGatedKernel, UnaryKernel
|
|
from .engram import EngramDecodeKernel, EngramGateConvBwdKernel, EngramGateConvFwdKernel
|
|
from .fft import FFTC2CKernel
|
|
from .fp8_lightning_indexer import FP8LightningIndexerKernel
|
|
from .fp8_quant import FP8QuantKernel
|
|
from .gated_deltanet import (
|
|
GatedDeltaNetBwdKernel,
|
|
GatedDeltaNetFwdKernel,
|
|
GatedDeltaNetPrefillFwdKernel,
|
|
)
|
|
from .gated_deltanet_recurrence import (
|
|
GatedDeltaNetDecodeFP32Kernel,
|
|
GatedDeltaNetDecodeKernel,
|
|
GatedDeltaNetDecodeRawCudaFlaStyleKernel,
|
|
)
|
|
from .gemm import GemmFp8BlockScaledKernel, GemmFp8EpilogueKernel, GemmKernel, GemvKernel
|
|
from .gemm_maca import GemmMACAKernel
|
|
from .gla import GLABwdKernel, GLAFwdKernel
|
|
from .gla_recurrence import GLADecodeFP32Kernel, GLADecodeKernel
|
|
from .grouped_gemm import GroupedGemmKernel
|
|
from .kernel_base import Kernel
|
|
from .mhc import MHCPostKernel, MHCPreKernel
|
|
from .moe import MoePermuteAlignKernel
|
|
from .norm import (
|
|
BatchNormBwdKernel,
|
|
BatchNormFwdInferKernel,
|
|
BatchNormFwdTrainKernel,
|
|
GroupNormKernel,
|
|
LayerNormKernel,
|
|
RMSNormKernel,
|
|
)
|
|
from .pool import (
|
|
AvgPool1dKernel,
|
|
AvgPool1dSpatialKernel,
|
|
AvgPool2dKernel,
|
|
AvgPool2dSpatialKernel,
|
|
AvgPool3dKernel,
|
|
AvgPool3dSpatialKernel,
|
|
MaxPool1dKernel,
|
|
MaxPool1dWithIndicesKernel,
|
|
MaxPool2dKernel,
|
|
MaxPool2dWithIndicesKernel,
|
|
MaxPool3dKernel,
|
|
MaxPool3dWithIndicesKernel,
|
|
)
|
|
from .quant_swiglu_channel_cast_transpose import QuantSwiGLUChannelCastTransposeKernel
|
|
from .rope import (
|
|
RopeLlama31Kernel,
|
|
RopeLongRopeKernel,
|
|
RopeNeoxKernel,
|
|
RopeNeoxPositionIdsKernel,
|
|
RopeNonNeoxKernel,
|
|
RopeYarnKernel,
|
|
)
|
|
from .topk_selector import TopkSelectorKernel
|
|
|
|
__all__ = [
|
|
"AvgPool1dKernel",
|
|
"AvgPool1dSpatialKernel",
|
|
"AvgPool2dKernel",
|
|
"AvgPool2dSpatialKernel",
|
|
"AvgPool3dKernel",
|
|
"AvgPool3dSpatialKernel",
|
|
"BatchNormBwdKernel",
|
|
"BatchNormFwdInferKernel",
|
|
"BatchNormFwdTrainKernel",
|
|
"BinaryKernel",
|
|
"BmmFp8Kernel",
|
|
"BmmKernel",
|
|
"Conv1dKernel",
|
|
"Conv1dPointwiseKernel",
|
|
"Conv2d1x1Kernel",
|
|
"Conv2dKernel",
|
|
"Conv3dKernel",
|
|
"DeltaNetBwdKernel",
|
|
"DeltaNetDecodeFP32Kernel",
|
|
"DeltaNetDecodeKernel",
|
|
"DeltaNetDecodeRawCudaFlaStyleKernel",
|
|
"DeltaNetFwdKernel",
|
|
"DropoutKernel",
|
|
"EngramDecodeKernel",
|
|
"EngramGateConvBwdKernel",
|
|
"EngramGateConvFwdKernel",
|
|
"FFTC2CKernel",
|
|
"FP8LightningIndexerKernel",
|
|
"FP8QuantKernel",
|
|
"FlashAttnBwdPostprocessKernel",
|
|
"FlashAttnBwdPreprocessKernel",
|
|
"FusedGatedKernel",
|
|
"GLABwdKernel",
|
|
"GLADecodeFP32Kernel",
|
|
"GLADecodeKernel",
|
|
"GLAFwdKernel",
|
|
"GQABwdKernel",
|
|
"GQABwdWgmmaPipelinedKernel",
|
|
"GQADecodeKernel",
|
|
"GQADecodePagedKernel",
|
|
"GQAFwdFP8Fa3ContractPtxAccBN224WsTmaVKernel",
|
|
"GQAFwdKernel",
|
|
"GQAFwdWgmmaPipelinedKernel",
|
|
"GQAPrefillFwdKernel",
|
|
"GQAPrefillPagedWithFP8KVCacheFwdKernel",
|
|
"GQAPrefillPagedWithKVCacheFwdKernel",
|
|
"GQAPrefillPagedWithKVCacheRopeAppendKernel",
|
|
"GQAPrefillPagedWithKVCacheRopeFwdKernel",
|
|
"GQAPrefillWithKVCacheFwdKernel",
|
|
"GQAPrefillWithKVCacheRopeAppendKernel",
|
|
"GQAPrefillWithKVCacheRopeFwdKernel",
|
|
"GQASlidingWindowFwdKernel",
|
|
"GQASlidingWindowFwdWgmmaPipelinedKernel",
|
|
"GQASlidingWindowVarlenFwdKernel",
|
|
"GQASlidingWindowVarlenFwdWgmmaPipelinedKernel",
|
|
"GatedDeltaNetBwdKernel",
|
|
"GatedDeltaNetDecodeFP32Kernel",
|
|
"GatedDeltaNetDecodeKernel",
|
|
"GatedDeltaNetDecodeRawCudaFlaStyleKernel",
|
|
"GatedDeltaNetFwdKernel",
|
|
"GatedDeltaNetPrefillFwdKernel",
|
|
"GemmFp8BlockScaledKernel",
|
|
"GemmFp8EpilogueKernel",
|
|
"GemmKernel",
|
|
"GemmMACAKernel",
|
|
"GemvKernel",
|
|
"GroupConv1dKernel",
|
|
"GroupConv2dKernel",
|
|
"GroupConv3dKernel",
|
|
"GroupNormKernel",
|
|
"GroupedGemmKernel",
|
|
"Kernel",
|
|
"LayerNormKernel",
|
|
"MHABwdKernel",
|
|
"MHABwdWgmmaPipelinedKernel",
|
|
"MHADecodeKernel",
|
|
"MHADecodePagedKernel",
|
|
"MHAFwdKernel",
|
|
"MHAFwdWgmmaPipelinedKernel",
|
|
"MHCPostKernel",
|
|
"MHCPreKernel",
|
|
"MLADecodeKernel",
|
|
"MLADecodeWsKernel",
|
|
"MaxPool1dKernel",
|
|
"MaxPool1dWithIndicesKernel",
|
|
"MaxPool2dKernel",
|
|
"MaxPool2dWithIndicesKernel",
|
|
"MaxPool3dKernel",
|
|
"MaxPool3dWithIndicesKernel",
|
|
"MeanPoolingFwdKernel",
|
|
"MoePermuteAlignKernel",
|
|
"NSACmpFwdVarlenKernel",
|
|
"NSAFwdVarlenKernel",
|
|
"NSATopkVarlenKernel",
|
|
"QuantSwiGLUChannelCastTransposeKernel",
|
|
"RMSNormKernel",
|
|
"RopeLlama31Kernel",
|
|
"RopeLongRopeKernel",
|
|
"RopeNeoxKernel",
|
|
"RopeNeoxPositionIdsKernel",
|
|
"RopeNonNeoxKernel",
|
|
"RopeYarnKernel",
|
|
"SparseMlaKernel",
|
|
"TopkSelectorKernel",
|
|
"UnaryKernel",
|
|
]
|