core.transformer

Subpackages

• core.transformer.custom_layers
  • core.transformer.custom_layers.batch_invariant_kernels
• core.transformer.moe
  • core.transformer.moe.ops
    • core.transformer.moe.ops.paged_stash
  • core.transformer.moe.router_trace
  • core.transformer.moe.inference_routing_mask_kernel
  • core.transformer.moe.shared_experts
  • core.transformer.moe.upcycling_utils
  • core.transformer.moe.fused_a2a
  • core.transformer.moe.moe_logging
  • core.transformer.moe.token_dispatcher
  • core.transformer.moe.moe_layer
  • core.transformer.moe.router_replay
  • core.transformer.moe.batch_invariant
  • core.transformer.moe.paged_stash
  • core.transformer.moe.router
  • core.transformer.moe.moe_utils
  • core.transformer.moe.token_dispatcher_inference
  • core.transformer.moe.experts

Submodules

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