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

1from flag_gems.fused.apply_repetition_penalties import apply_repetition_penalties 

2from flag_gems.fused.concat_and_cache_mla import concat_and_cache_mla 

3from flag_gems.fused.cross_entropy_loss import cross_entropy_loss 

4from flag_gems.fused.cutlass_scaled_mm import cutlass_scaled_mm 

5from flag_gems.fused.FLA import ( 

6 chunk_gated_delta_rule_fwd, 

7 fused_recurrent_gated_delta_rule_fwd, 

8) 

9from flag_gems.fused.flash_mla import flash_mla 

10from flag_gems.fused.fused_add_rms_norm import fused_add_rms_norm 

11from flag_gems.fused.fused_moe import fused_experts_impl, invoke_fused_moe_triton_kernel 

12from flag_gems.fused.geglu import dgeglu, geglu 

13from flag_gems.fused.gelu_and_mul import gelu_and_mul 

14from flag_gems.fused.grouped_topk import grouped_topk 

15from flag_gems.fused.instance_norm import instance_norm 

16from flag_gems.fused.moe_align_block_size import ( 

17 moe_align_block_size, 

18 moe_align_block_size_triton, 

19) 

20from flag_gems.fused.moe_sum import moe_sum 

21from flag_gems.fused.outer import outer 

22from flag_gems.fused.reglu import dreglu, reglu 

23from flag_gems.fused.reshape_and_cache import reshape_and_cache 

24from flag_gems.fused.reshape_and_cache_flash import reshape_and_cache_flash 

25from flag_gems.fused.rotary_embedding import apply_rotary_pos_emb 

26from flag_gems.fused.rwkv_ka_fusion import rwkv_ka_fusion 

27from flag_gems.fused.rwkv_mm_sparsity import rwkv_mm_sparsity 

28from flag_gems.fused.silu_and_mul import silu_and_mul, silu_and_mul_out 

29from flag_gems.fused.skip_layernorm import skip_layer_norm 

30from flag_gems.fused.swiglu import dswiglu, swiglu 

31from flag_gems.fused.topk_softmax import topk_softmax 

32from flag_gems.fused.weight_norm import weight_norm 

33 

34__all__ = [ 

35 "apply_repetition_penalties", 

36 "apply_rotary_pos_emb", 

37 "chunk_gated_delta_rule_fwd", 

38 "concat_and_cache_mla", 

39 "cutlass_scaled_mm", 

40 "cross_entropy_loss", 

41 "dgeglu", 

42 "dreglu", 

43 "dswiglu", 

44 "flash_mla", 

45 "fused_add_rms_norm", 

46 "fused_experts_impl", 

47 "fused_recurrent_gated_delta_rule_fwd", 

48 "geglu", 

49 "gelu_and_mul", 

50 "grouped_topk", 

51 "instance_norm", 

52 "invoke_fused_moe_triton_kernel", 

53 "moe_sum", 

54 "moe_align_block_size", 

55 "moe_align_block_size_triton", 

56 "outer", 

57 "reglu", 

58 "reshape_and_cache", 

59 "reshape_and_cache_flash", 

60 "rwkv_ka_fusion", 

61 "rwkv_mm_sparsity", 

62 "silu_and_mul", 

63 "silu_and_mul_out", 

64 "skip_layer_norm", 

65 "swiglu", 

66 "topk_softmax", 

67 "weight_norm", 

68]