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 "concat_and_cache_mla", 

18 "cross_entropy_loss", 

19 "flash_mla", 

20 "fused_add_rms_norm", 

21 "gelu_and_mul", 

22 "silu_and_mul", 

23 "skip_layer_norm", 

24 ] 

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

26 from .gcu400.cross_entropy_loss import cross_entropy_loss 

27 from .gcu400.flash_mla import flash_mla 

28 from .gcu400.fused_add_rms_norm import fused_add_rms_norm 

29 from .gcu400.gelu_and_mul import gelu_and_mul 

30 from .gcu400.outer import outer 

31 from .gcu400.silu_and_mul import silu_and_mul 

32 from .gcu400.skip_layernorm import skip_layer_norm 

33 

34 __all__ = [ 

35 "cross_entropy_loss", 

36 "flash_mla", 

37 "fused_add_rms_norm", 

38 "gelu_and_mul", 

39 "outer", 

40 "silu_and_mul", 

41 "skip_layer_norm", 

42 ]