Coverage for src/flag_gems/runtime/backend/_enflame/fused/__init__.py: 0%

1from .. import arch_version 

2 

3__all__ = [] 

4 

5if arch_version == 300: 

6 from .gcu300.concat_and_cache_mla import concat_and_cache_mla # noqa: F401 

7 from .gcu300.cross_entropy_loss import cross_entropy_loss 

8 from .gcu300.flash_mla import flash_mla 

9 from .gcu300.fused_add_rms_norm import fused_add_rms_norm 

10 from .gcu300.gelu_and_mul import gelu_and_mul 

11 from .gcu300.rotary_embedding import apply_rotary_pos_emb # noqa: F401 

12 from .gcu300.silu_and_mul import silu_and_mul 

13 from .gcu300.skip_layernorm import skip_layer_norm 

14 

15 __all__ = [ 

16 "apply_rotary_pos_emb", 

17 "silu_and_mul", 

18 "gelu_and_mul", 

19 "cross_entropy_loss", 

20 "flash_mla", 

21 "skip_layer_norm", 

22 "fused_add_rms_norm", 

23 "concat_and_cache_mla", 

24 ] 

25elif arch_version == 400 or arch_version == 410: 

26 from .gcu400.bincount import bincount 

27 from .gcu400.cross_entropy_loss import cross_entropy_loss 

28 from .gcu400.flash_mla import flash_mla 

29 from .gcu400.fused_add_rms_norm import fused_add_rms_norm 

30 from .gcu400.gelu_and_mul import gelu_and_mul 

31 from .gcu400.outer import outer 

32 from .gcu400.rotary_embedding import apply_rotary_pos_emb # noqa: F401 

33 from .gcu400.silu_and_mul import silu_and_mul 

34 from .gcu400.skip_layernorm import skip_layer_norm 

35 from .gcu400.sparse_attention import sparse_attn_triton 

36 

37 __all__ = [ 

38 "apply_rotary_pos_emb", 

39 "gelu_and_mul", 

40 "silu_and_mul", 

41 "cross_entropy_loss", 

42 "flash_mla", 

43 "moe_sum", 

44 "outer", 

45 "fused_add_rms_norm", 

46 "skip_layer_norm", 

47 "rwkv_ka_fusion", 

48 "bincount", 

49 "sparse_attn_triton", 

50 ]