Coverage for src/flag_gems/fused/__init__.py: 100%

1from flag_gems.fused.apply_repetition_penalties import apply_repetition_penalties 

2from flag_gems.fused.bincount import bincount 

3from flag_gems.fused.concat_and_cache_mla import concat_and_cache_mla 

4from flag_gems.fused.cross_entropy_loss import cross_entropy_loss 

5from flag_gems.fused.cutlass_scaled_mm import cutlass_scaled_mm 

6from flag_gems.fused.DSA.bin_topk import bucket_sort_topk 

7from flag_gems.fused.FLA import ( 

8 chunk_gated_delta_rule_fwd, 

9 fused_recurrent_gated_delta_rule_fwd, 

10) 

11from flag_gems.fused.flash_mla import flash_mla 

12from flag_gems.fused.flashmla_sparse import flash_mla_sparse_fwd 

13from flag_gems.fused.fused_add_rms_norm import fused_add_rms_norm 

14from flag_gems.fused.fused_moe import ( 

15 dispatch_fused_moe_kernel, 

16 fused_experts_impl, 

17 inplace_fused_experts, 

18 invoke_fused_moe_triton_kernel, 

19 outplace_fused_experts, 

20) 

21from flag_gems.fused.geglu import dgeglu, geglu 

22from flag_gems.fused.gelu_and_mul import gelu_and_mul 

23from flag_gems.fused.grouped_topk import grouped_topk 

24from flag_gems.fused.instance_norm import instance_norm 

25from flag_gems.fused.mhc import ( 

26 mhc_bwd, 

27 mhc_bwd_ref, 

28 mhc_post, 

29 mhc_pre, 

30 sinkhorn_forward, 

31) 

32from flag_gems.fused.moe_align_block_size import ( 

33 moe_align_block_size, 

34 moe_align_block_size_triton, 

35) 

36from flag_gems.fused.moe_sum import moe_sum 

37from flag_gems.fused.outer import outer 

38from flag_gems.fused.reglu import dreglu, reglu 

39from flag_gems.fused.reshape_and_cache import reshape_and_cache 

40from flag_gems.fused.reshape_and_cache_flash import reshape_and_cache_flash 

41from flag_gems.fused.rotary_embedding import apply_rotary_pos_emb 

42from flag_gems.fused.rwkv_ka_fusion import rwkv_ka_fusion 

43from flag_gems.fused.rwkv_mm_sparsity import rwkv_mm_sparsity 

44from flag_gems.fused.silu_and_mul import silu_and_mul, silu_and_mul_out 

45from flag_gems.fused.skip_layernorm import skip_layer_norm 

46from flag_gems.fused.sparse_attention import sparse_attn_triton 

47from flag_gems.fused.swiglu import dswiglu, swiglu 

48from flag_gems.fused.topk_softmax import topk_softmax 

49from flag_gems.fused.weight_norm import weight_norm 

50 

51__all__ = [ 

52 "apply_repetition_penalties", 

53 "apply_rotary_pos_emb", 

54 "bincount", 

55 "bucket_sort_topk", 

56 "chunk_gated_delta_rule_fwd", 

57 "concat_and_cache_mla", 

58 "cutlass_scaled_mm", 

59 "cross_entropy_loss", 

60 "dispatch_fused_moe_kernel", 

61 "dgeglu", 

62 "dreglu", 

63 "dswiglu", 

64 "flash_mla", 

65 "flash_mla_sparse_fwd", 

66 "fused_add_rms_norm", 

67 "fused_experts_impl", 

68 "fused_recurrent_gated_delta_rule_fwd", 

69 "geglu", 

70 "gelu_and_mul", 

71 "grouped_topk", 

72 "inplace_fused_experts", 

73 "instance_norm", 

74 "invoke_fused_moe_triton_kernel", 

75 "mhc_bwd", 

76 "mhc_bwd_ref", 

77 "mhc_post", 

78 "mhc_pre", 

79 "moe_sum", 

80 "moe_align_block_size", 

81 "moe_align_block_size_triton", 

82 "outer", 

83 "outplace_fused_experts", 

84 "reglu", 

85 "reshape_and_cache", 

86 "reshape_and_cache_flash", 

87 "rwkv_ka_fusion", 

88 "rwkv_mm_sparsity", 

89 "silu_and_mul", 

90 "silu_and_mul_out", 

91 "sinkhorn_forward", 

92 "skip_layer_norm", 

93 "swiglu", 

94 "topk_softmax", 

95 "weight_norm", 

96 "sparse_attn_triton", 

97]