d7dcdf3efd
Co-authored-by: liqichao <liqichao@baidu.com> Co-authored-by: chenbong <bhchen@stu.xmu.edu.cn>
64 lines
1.8 KiB
Python
64 lines
1.8 KiB
Python
from .attn import (
|
|
fused_store_cache,
|
|
get_paged_mqa_logits_metadata,
|
|
triton_create_paged_compress_data,
|
|
)
|
|
from .c128_cleanup import clear_unaccepted_c128_draft_states
|
|
from .compress import (
|
|
CompressorDecodePlan,
|
|
CompressorPrefillPlan,
|
|
compress_forward,
|
|
compress_norm_rope_store,
|
|
)
|
|
from .compress_old import fused_norm_rope_inplace
|
|
from .elementwise import (
|
|
fused_k_norm_rope_flashmla,
|
|
fused_q_indexer_rope_first_quant,
|
|
fused_q_indexer_rope_hadamard_fp4_quant,
|
|
fused_q_indexer_rope_hadamard_quant,
|
|
fused_q_norm_rope,
|
|
fused_rope_inplace,
|
|
)
|
|
from .fp8_wo_a import sglang_per_token_group_quant_fp8_dsv4_wo_a
|
|
from .gemm import linear_bf16_fp32
|
|
from .moe import (
|
|
hash_topk,
|
|
mask_topk_ids,
|
|
mega_moe_pre_dispatch,
|
|
silu_and_mul_clamp,
|
|
silu_and_mul_contig_post_quant,
|
|
silu_and_mul_masked_post_quant,
|
|
)
|
|
from .topk import plan_topk_v2, topk_transform_512, topk_transform_512_v2
|
|
from .utils import make_name
|
|
|
|
__all__ = [
|
|
"CompressorDecodePlan",
|
|
"CompressorPrefillPlan",
|
|
"compress_forward",
|
|
"compress_norm_rope_store",
|
|
"clear_unaccepted_c128_draft_states",
|
|
"fused_norm_rope_inplace",
|
|
"fused_store_cache",
|
|
"fused_rope_inplace",
|
|
"fused_q_norm_rope",
|
|
"fused_q_indexer_rope_first_quant",
|
|
"fused_q_indexer_rope_hadamard_fp4_quant",
|
|
"fused_q_indexer_rope_hadamard_quant",
|
|
"fused_k_norm_rope_flashmla",
|
|
"sglang_per_token_group_quant_fp8_dsv4_wo_a",
|
|
"make_name",
|
|
"linear_bf16_fp32",
|
|
"get_paged_mqa_logits_metadata",
|
|
"triton_create_paged_compress_data",
|
|
"topk_transform_512",
|
|
"topk_transform_512_v2",
|
|
"plan_topk_v2",
|
|
"hash_topk",
|
|
"mega_moe_pre_dispatch",
|
|
"mask_topk_ids",
|
|
"silu_and_mul_clamp",
|
|
"silu_and_mul_masked_post_quant",
|
|
"silu_and_mul_contig_post_quant",
|
|
]
|