core.transformer

Subpackages

• core.transformer.custom_layers
  • core.transformer.custom_layers.batch_invariant_kernels
• core.transformer.moe
  • core.transformer.moe.token_dispatcher
  • core.transformer.moe.moe_layer
  • core.transformer.moe.router
  • core.transformer.moe.router_replay
  • core.transformer.moe.moe_utils
  • core.transformer.moe.fused_a2a
  • core.transformer.moe.token_dispatcher_inference
  • core.transformer.moe.upcycling_utils
  • core.transformer.moe.experts
  • core.transformer.moe.shared_experts

Submodules

• core.transformer.identity_op
• core.transformer.dot_product_attention
• core.transformer.enums
• core.transformer.fsdp_dtensor_checkpoint
• core.transformer.pipeline_parallel_layer_layout
• core.transformer.cuda_graphs
• core.transformer.transformer_block
• core.transformer.utils
• core.transformer.multi_latent_attention
• core.transformer.multi_token_prediction
• core.transformer.transformer_config
• core.transformer.torch_layer_norm
• core.transformer.attention
• core.transformer.transformer_layer
• core.transformer.spec_utils
• core.transformer.torch_norm
• core.transformer.module
• core.transformer.mlp