Is this page helpful?

core.ssm.ops#

Triton kernels for the linear-attention / SSM mixers.

Organized by which recurrence a kernel belongs to:

  • common – shared by every variant (causal convolution, slot-indexed state gather/scatter, autotune determinism).

  • mamba2 – the Mamba2 state-space duality kernels, adapted from vLLM and state-spaces/mamba.

  • gdp – the Gated Delta Product kernels, forked from flash-linear-attention.

The re-exports below are the historical top-level entry points and are kept so existing callers do not have to spell out the subpackage.

This package is internal to Megatron Core. The kernels here are implementation details of the SSM mixers, not a supported kernel library.

Subpackages#

Package Contents#

Data#

__all__

API#

core.ssm.ops.__all__#

[‘mamba_chunk_scan_combined_varlen’, ‘causal_conv1d_varlen_fn’]