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

17 "gelu_and_mul", 

18 "cross_entropy_loss", 

19 "flash_mla", 

20 "skip_layer_norm", 

21 "fused_add_rms_norm", 

22 ] 

23elif arch_version == 400 or arch_version == 410: 

24 from .gcu400.cross_entropy_loss import cross_entropy_loss 

25 from .gcu400.flash_mla import flash_mla 

26 from .gcu400.fused_add_rms_norm import fused_add_rms_norm 

27 from .gcu400.gelu_and_mul import gelu_and_mul 

28 from .gcu400.outer import outer 

29 from .gcu400.silu_and_mul import silu_and_mul 

30 from .gcu400.skip_layernorm import skip_layer_norm 

31 

32 __all__ = [ 

33 "gelu_and_mul", 

34 "silu_and_mul", 

35 "cross_entropy_loss", 

36 "flash_mla", 

37 "outer", 

38 "fused_add_rms_norm", 

39 "skip_layer_norm", 

40 ]