core.ssm

Subpackages

• core.ssm.gated_delta_net
  • core.ssm.gated_delta_net.gdn
  • core.ssm.gated_delta_net.common
  • core.ssm.gated_delta_net.gdn2
• core.ssm.ops
  • core.ssm.ops.gdp
    • core.ssm.ops.gdp.metadata
    • core.ssm.ops.gdp.chunk_o
    • core.ssm.ops.gdp.fused_recurrent
    • core.ssm.ops.gdp.common
    • core.ssm.ops.gdp.decode_prepare
    • core.ssm.ops.gdp.wy_fast
    • core.ssm.ops.gdp.solve_tril
    • core.ssm.ops.gdp.chunk_h
    • core.ssm.ops.gdp.chunk
    • core.ssm.ops.gdp.scaled_dot_kkt
    • core.ssm.ops.gdp.cumsum
  • core.ssm.ops.mamba2
    • core.ssm.ops.mamba2.ssd_state_passing
    • core.ssm.ops.mamba2.mamba_ssm
    • core.ssm.ops.mamba2.ssd_bmm
    • core.ssm.ops.mamba2.batch_invariant_decode
    • core.ssm.ops.mamba2.ssd_chunk_state
    • core.ssm.ops.mamba2.ssd_chunk_scan
    • core.ssm.ops.mamba2.ssd_combined
  • core.ssm.ops.common
    • core.ssm.ops.common.determinism
    • core.ssm.ops.common.intermediate_extraction
    • core.ssm.ops.common.causal_conv1d_varlen
    • core.ssm.ops.common.causal_conv1d_triton
• core.ssm.context_parallel
  • core.ssm.context_parallel.gdp_common
  • core.ssm.context_parallel.gdp_cutedsl
  • core.ssm.context_parallel.gdp
  • core.ssm.context_parallel.chunkwise

Submodules

• core.ssm.mamba_hybrid_layer_allocation
• core.ssm.packed_seq_helpers
• core.ssm.mamba_layer
• core.ssm.ssm_inference
• core.ssm.mamba_context_parallel
• core.ssm.triton_cache_manager
• core.ssm.mamba_mixer
• core.ssm.utils
• core.ssm.mamba_layer_config
• core.ssm.mlp_layer_config
• core.ssm.gdp_context_parallel
• core.ssm.gated_delta_product
• core.ssm.gdn_layer_config
• core.ssm.mlp_layer
• core.ssm.causal_conv1d
• core.ssm.mamba_block