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

1from ..utils.config_utils import arch_version 

2 

3__all__ = [] 

4 

5if arch_version == 300: 

6 from .gcu300.cross_entropy_loss import cross_entropy_loss 

7 from .gcu300.flash_mla import flash_mla 

8 from .gcu300.fused_add_rms_norm import fused_add_rms_norm 

9 from .gcu300.gelu_and_mul import gelu_and_mul 

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

11 from .gcu300.silu_and_mul import silu_and_mul 

12 from .gcu300.skip_layernorm import skip_layer_norm 

13 

14 __all__ = [ 

15 "apply_rotary_pos_emb", 

16 "cross_entropy_loss", 

17 "flash_mla", 

18 "fused_add_rms_norm", 

19 "gelu_and_mul", 

20 "silu_and_mul", 

21 "skip_layer_norm", 

22 ] 

23 

24elif arch_version == 400 or arch_version == 410: 

25 from .gcu400.cross_entropy_loss import cross_entropy_loss 

26 from .gcu400.flash_mla import flash_mla 

27 from .gcu400.fused_add_rms_norm import fused_add_rms_norm 

28 from .gcu400.gelu_and_mul import gelu_and_mul 

29 from .gcu400.outer import outer 

30 from .gcu400.silu_and_mul import silu_and_mul 

31 from .gcu400.skip_layernorm import skip_layer_norm 

32 

33 __all__ = [ 

34 "cross_entropy_loss", 

35 "flash_mla", 

36 "fused_add_rms_norm", 

37 "gelu_and_mul", 

38 "outer", 

39 "silu_and_mul", 

40 "skip_layer_norm", 

41 ]