forked from ccf-ai-infra/TileOPs-Metax
139 lines
4.0 KiB
Python
139 lines
4.0 KiB
Python
from .attention import (
|
|
FlashAttnBwdPostprocessKernel,
|
|
FlashAttnBwdPreprocessKernel,
|
|
GqaBwdKernel,
|
|
GqaBwdWgmmaPipelinedKernel,
|
|
GqaDecodeKernel,
|
|
GqaDecodePagedKernel,
|
|
GqaFwdKernel,
|
|
GqaFwdWgmmaPipelinedKernel,
|
|
GqaSlidingWindowFwdKernel,
|
|
GqaSlidingWindowFwdWgmmaPipelinedKernel,
|
|
GqaSlidingWindowVarlenFwdKernel,
|
|
GqaSlidingWindowVarlenFwdWgmmaPipelinedKernel,
|
|
MeanPoolingFwdKernel,
|
|
MhaBwdKernel,
|
|
MhaBwdWgmmaPipelinedKernel,
|
|
MhaDecodeKernel,
|
|
MhaDecodePagedKernel,
|
|
MhaFwdKernel,
|
|
MhaFwdWgmmaPipelinedKernel,
|
|
MlaDecodeKernel,
|
|
MlaDecodeWsKernel,
|
|
NSACmpFwdVarlenKernel,
|
|
NSAFwdVarlenKernel,
|
|
NSATopkVarlenKernel,
|
|
SparseMlaKernel,
|
|
)
|
|
from .convolution import Conv1dKernel, Conv2d1x1Kernel, Conv2dKernel, Conv3dKernel
|
|
from .deltanet import DeltaNetBwdKernel, DeltaNetFwdKernel
|
|
from .deltanet_recurrence import DeltaNetDecodeFP32Kernel, DeltaNetDecodeKernel
|
|
from .dropout import DropoutKernel
|
|
from .elementwise import BinaryKernel, FusedGatedKernel, UnaryKernel
|
|
from .engram import EngramDecodeKernel, EngramGateConvBwdKernel, EngramGateConvFwdKernel
|
|
from .fft import FFTC2CKernel
|
|
from .fp8_lighting_indexer import FP8LightingIndexerKernel
|
|
from .fp8_quant import FP8QuantKernel
|
|
from .gated_deltanet import GatedDeltaNetBwdKernel, GatedDeltaNetFwdKernel
|
|
from .gated_deltanet_recurrence import GatedDeltaNetDecodeFP32Kernel, GatedDeltaNetDecodeKernel
|
|
from .gemm import GemmKernel, GemvKernel
|
|
from .gla import GLABwdKernel, GLAFwdKernel
|
|
from .gla_recurrence import GLADecodeFP32Kernel, GLADecodeKernel
|
|
from .grouped_gemm import GroupedGemmKernel
|
|
from .kernel import Kernel
|
|
from .mhc import MHCPostKernel, MHCPreKernel
|
|
from .moe import MoePermuteAlignKernel
|
|
from .norm import (
|
|
BatchNormBwdKernel,
|
|
BatchNormFwdInferKernel,
|
|
BatchNormFwdTrainKernel,
|
|
GroupNormKernel,
|
|
LayerNormKernel,
|
|
RmsNormKernel,
|
|
)
|
|
from .pool import AvgPool1dKernel, AvgPool2dKernel, AvgPool3dKernel
|
|
from .rope import (
|
|
RopeLlama31Kernel,
|
|
RopeLongRopeKernel,
|
|
RopeNeoxKernel,
|
|
RopeNonNeoxKernel,
|
|
RopeYarnKernel,
|
|
)
|
|
from .topk_selector import TopkSelectorKernel
|
|
|
|
__all__ = [
|
|
"BinaryKernel",
|
|
"AvgPool1dKernel",
|
|
"AvgPool2dKernel",
|
|
"AvgPool3dKernel",
|
|
"BatchNormBwdKernel",
|
|
"BatchNormFwdInferKernel",
|
|
"BatchNormFwdTrainKernel",
|
|
"DropoutKernel",
|
|
"Conv1dKernel",
|
|
"Conv2d1x1Kernel",
|
|
"Conv2dKernel",
|
|
"Conv3dKernel",
|
|
"EngramDecodeKernel",
|
|
"EngramGateConvBwdKernel",
|
|
"EngramGateConvFwdKernel",
|
|
"FFTC2CKernel",
|
|
"FusedGatedKernel",
|
|
"FlashAttnBwdPostprocessKernel",
|
|
"FlashAttnBwdPreprocessKernel",
|
|
"FP8LightingIndexerKernel",
|
|
"FP8QuantKernel",
|
|
"GqaSlidingWindowFwdKernel",
|
|
"GqaSlidingWindowFwdWgmmaPipelinedKernel",
|
|
"GqaSlidingWindowVarlenFwdKernel",
|
|
"GqaSlidingWindowVarlenFwdWgmmaPipelinedKernel",
|
|
"DeltaNetBwdKernel",
|
|
"DeltaNetDecodeFP32Kernel",
|
|
"DeltaNetDecodeKernel",
|
|
"DeltaNetFwdKernel",
|
|
"GatedDeltaNetBwdKernel",
|
|
"GatedDeltaNetDecodeFP32Kernel",
|
|
"GatedDeltaNetDecodeKernel",
|
|
"GatedDeltaNetFwdKernel",
|
|
"GLABwdKernel",
|
|
"GLADecodeFP32Kernel",
|
|
"GLADecodeKernel",
|
|
"GLAFwdKernel",
|
|
"GemmKernel",
|
|
"GemvKernel",
|
|
"GqaBwdKernel",
|
|
"GqaBwdWgmmaPipelinedKernel",
|
|
"GqaFwdKernel",
|
|
"GqaFwdWgmmaPipelinedKernel",
|
|
"GroupNormKernel",
|
|
"Kernel",
|
|
"LayerNormKernel",
|
|
"MeanPoolingFwdKernel",
|
|
"MhaBwdKernel",
|
|
"MhaBwdWgmmaPipelinedKernel",
|
|
"MhaFwdKernel",
|
|
"MhaFwdWgmmaPipelinedKernel",
|
|
"NSACmpFwdVarlenKernel",
|
|
"NSAFwdVarlenKernel",
|
|
"NSATopkVarlenKernel",
|
|
"RmsNormKernel",
|
|
"RopeLlama31Kernel",
|
|
"RopeLongRopeKernel",
|
|
"RopeNeoxKernel",
|
|
"RopeNonNeoxKernel",
|
|
"RopeYarnKernel",
|
|
"SparseMlaKernel",
|
|
"TopkSelectorKernel",
|
|
"UnaryKernel",
|
|
"GqaDecodeKernel",
|
|
"GqaDecodePagedKernel",
|
|
"GroupedGemmKernel",
|
|
"MoePermuteAlignKernel",
|
|
"MHCPostKernel",
|
|
"MHCPreKernel",
|
|
"MhaDecodeKernel",
|
|
"MhaDecodePagedKernel",
|
|
"MlaDecodeKernel",
|
|
"MlaDecodeWsKernel",
|
|
]
|