Repository navigation
Clean up SSM ops directory structure - #6444
Conversation
santhnm2
changed the title
`megatron/core/ssm/ops/` was a flat directory holding the Mamba2 SSD kernels alongside the handful of kernels that are not specific to any recurrence. That was tenable while Mamba2 was the only variant with forked kernels. It stops being tenable as soon as a second variant lands its own fork, because there is then no way to tell from the layout which kernels a given mixer owns. Split into two packages: * `ops/common/` -- the short causal convolution (`causal_conv1d_triton`, `causal_conv1d_varlen`), the slot-indexed state gather/scatter that prefix caching uses (`intermediate_extraction`), and the autotune-config filtering that keeps the kernels deterministic (`determinism`). Every linear-attention mixer needs these; `determinism` was already shared between the SSD kernels and the varlen convolution. * `ops/mamba2/` -- the state-space duality kernels (`ssd_*`, `mamba_ssm`, `batch_invariant_decode`). A later variant with its own forked kernels adds a sibling package next to `mamba2` rather than growing the flat directory. `intermediate_extraction` reads as Mamba-specific from its docstring, but the kernels are a plain conditional gather/scatter over flat state rows with no Mamba math in them, so it belongs with the shared code. Pure move. The top-level `ops` re-exports are unchanged, so callers that go through `megatron.core.ssm.ops` need no edit. Unit tests move to mirror the new layout. Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
b540efb to
43f1592
Compare
|
/ok to test 43f1592 |
guihong-nv
left a comment
guihong-nv
left a comment
There was a problem hiding this comment.
Nice, focused move overall. I left three comments on the bits that seem worth clearing up before it lands.
| from megatron.core.ssm.ops.mamba2.batch_invariant_decode import MambaBatchInvariantDecode | ||
| from megatron.core.ssm.ops.common.causal_conv1d_triton import causal_conv1d_update | ||
| from megatron.core.ssm.ops.common.intermediate_extraction import ( | ||
| scatter_intermediate_conv, | ||
| scatter_intermediate_ssm, | ||
| ) | ||
| from megatron.core.ssm.ops.mamba_ssm import selective_state_update | ||
| from megatron.core.ssm.ops.mamba2.mamba_ssm import selective_state_update |
There was a problem hiding this comment.
Could you run isort on this block? common sorts before mamba2, so the two common imports should come first. The lint job runs isort --check here, so this should fail as written.
| @@ -0,0 +1,10 @@ | |||
| # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. | |||
|
|
|||
| """Triton kernels shared by every linear-attention / SSM mixer. | |||
There was a problem hiding this comment.
Since this package is now described as shared, could you also update common/intermediate_extraction.py's module docstring? It still says “for Mamba intermediate-state extraction,” which makes the file sound Mamba-specific.
There was a problem hiding this comment.
Updated to clarify that the intermediate_extraction.py file is not specific to Mamba
| * `mamba2` -- the Mamba2 state-space duality kernels, adapted from vLLM and | ||
| state-spaces/mamba. | ||
|
|
||
| The re-exports below are the historical top-level entry points and are kept so |
There was a problem hiding this comment.
Are direct submodule imports intentionally considered internal? These re-exports preserve the top-level API, but imports like ops.mamba_ssm.selective_state_update and ops.intermediate_extraction.scatter_intermediate_* disappear without a shim or deprecation path. If those paths are public, could we keep thin re-export modules for a release? Otherwise, it would help to call out in the PR that they’re internal.
There was a problem hiding this comment.
I can include a backwards-compatible shim like this:
# Backwards-compatible aliases for the pre-split module paths, e.g.
# `megatron.core.ssm.ops.ssd_bmm` -> `megatron.core.ssm.ops.mamba2.ssd_bmm`.
# Registering them in `sys.modules` keeps `from megatron.core.ssm.ops.<old> import x`
# working, since importing any submodule imports this package first.
# Remove in megatron-core 0.20.
_MOVED_MODULES = {
"batch_invariant_decode": "mamba2.batch_invariant_decode",
"causal_conv1d_triton": "common.causal_conv1d_triton",
"causal_conv1d_varlen": "common.causal_conv1d_varlen",
"determinism": "common.determinism",
"intermediate_extraction": "common.intermediate_extraction",
"mamba_ssm": "mamba2.mamba_ssm",
"ssd_bmm": "mamba2.ssd_bmm",
"ssd_chunk_scan": "mamba2.ssd_chunk_scan",
"ssd_chunk_state": "mamba2.ssd_chunk_state",
"ssd_combined": "mamba2.ssd_combined",
"ssd_state_passing": "mamba2.ssd_state_passing",
}
class _MovedModule(types.ModuleType):
"""Alias for a module that moved into a subpackage, warning on attribute access."""
def __init__(self, old_name: str, new_name: str):
super().__init__(old_name, f"Deprecated alias for {new_name}.")
self._new_name = new_name
def __getattr__(self, name: str):
warnings.warn(
f"{self.__name__} has moved to {self._new_name}. The old path will be removed
"
"in megatron-core 0.20.",
DeprecationWarning,
stacklevel=2,
)
return getattr(importlib.import_module(self._new_name), name)
for _old, _new in _MOVED_MODULES.items():
_old_path = f"{__name__}.{_old}"
if _old_path not in sys.modules:
sys.modules[_old_path] = _MovedModule(_old_path, f"{__name__}.{_new}")
I suspect there aren't actually any external consumers of these ops though, so I have a slight preference towards breaking the backwards compatibility. Either way is fine though, let me know which you prefer.
There was a problem hiding this comment.
Inference code at this point is considered internal. We disabled the backward compatibility check for Inference.
There was a problem hiding this comment.
Added a comment confirming this
NeMo-RL functional testscluster: Exactly what was tested
Compared against a baseline of the same suite on megatron-core
No failures attributed to this PR. The three red rows (L1 topp/topk and both L2 nano-3.5 tests) match known pre-existing token_mult_prob_error threshold breaks on megatron-core main. Megatron-Bridge was pinned to Posted by the nemo-rl-testing-agent. Re-runs edit this comment in place. |
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
svcnvidia-nemo-ci
added
Final Review
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
|
/ok to test 5fec93f |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/31731634391 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/31761228988 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/31763897419 |
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com> Co-authored-by: Shanmugam Ramasamy <shanmugamr@nvidia.com>
This branch was previously deployed
6 participants
What does this PR do?
Splits
megatron/core/ssm/ops/into two sub-packages to prepare for future linear attention variants:ops/common/-- the short causal convolution (causal_conv1d_triton,causal_conv1d_varlen), the slot-indexed state gather/scatter that prefix caching uses (intermediate_extraction), and the autotune-config filtering that keeps the kernels deterministic (determinism). Every linear-attention mixer needs these;determinismwas already shared between the SSD kernels and the varlen convolution.ops/mamba2/-- the state-space duality kernels (ssd_*,mamba_ssm,batch_invariant_decode).intermediate_extractionreads as Mamba-specific from its docstring, but the kernels are a plain conditional gather/scatter over flat state rows with no Mamba math in them, so it belongs with the shared code.Pure move. The top-level
opsre-exports are unchanged, so callers that go throughmegatron.core.ssm.opsneed no edit. Unit tests move to mirror the new layout.Issue tracking
For PRs from open-source community contributors:
Linked issue:
Contribution process
Pre-checks
Code review
Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!
All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.
Step 1: Mark PR as "Ready for Review"
.github/CODEOWNERS.Final Review might get declined if these requirements are not fulfilled.
Step 2: Final Review
For PRs that change
megatron/core, once all expert reviewers have approved, theFinal Reviewlabel is applied automatically and final reviewers are assigned.For PRs outside
megatron/core, this step is skipped.Step 3: Approved
Once all required reviewers have approved, the
Approvedlabel is applied automatically.Merge
Any member of mcore-engineers will be able to merge your PR.