feat(models): add native Nemotron-H implementation

docs/scaling.md

@@ -112,7 +112,7 @@ For DeepEP, set `type = "deepep"` and tune `num_sms` plus optional `token_chunk_
 
 CP shards a single sequence across multiple GPUs along the token dimension — for long-context sequences. Prefer `ulysses`: it gets the most throughput and is the only style that works for hybrid linear-attention/Mamba models (Qwen3.5, NemotronH) and for VLMs. Each model class declares its supported styles in `cp_support`; an unsupported `cp_style`, or a model with no CP support at all (DeepSeek V4, GPT-OSS), is rejected at setup.
 
-`ulysses` head-shards Q/K/V, so the CP degree must divide `num_attention_heads`. GQA models with fewer KV heads than the CP degree (e.g. NemotronH: 32 query heads, 2 KV heads) are supported via KV-head replication; the CP degree must then be a multiple of `num_key_value_heads`. Hybrid Mamba layers head-shard independently (`cp_mamba`), which requires the CP degree to divide `mamba_num_heads` and `n_groups`.
+`ulysses` head-shards Q/K/V, so the CP degree must divide `num_attention_heads`. GQA models with fewer KV heads than the CP degree (e.g. NemotronH: 32 query heads, 2 KV heads) are supported via KV-head replication; the CP degree must then be a multiple of `num_key_value_heads`. Nemotron-H redistributes its Mamba activations between sequence- and head-parallel layouts in the model-owned Mamba path; the CP degree must divide `mamba_num_heads` and `n_groups`.
 
 ```toml
 [trainer.model]

pyproject.toml

@@ -62,6 +62,7 @@ gpu = [
     "dion ; sys_platform == 'linux'",
     "tilelang>=0.1.8 ; sys_platform == 'linux'",
     "flash-linear-attention ; sys_platform == 'linux'",
+    "mamba-ssm>=2.3.0 ; sys_platform == 'linux'",
     "nvidia-ml-py>=12.575.51 ; sys_platform == 'linux'",
     "torchao>=0.18.0 ; sys_platform == 'linux'",
     "modelexpress==0.3.0 ; sys_platform == 'linux'",
@@ -139,6 +140,8 @@ environments = [
 # Override torch's pinned cuDNN to fix the Conv3d performance regression.
 # See: https://github.com/pytorch/pytorch/issues/166122
 override-dependencies = [
+    # Keep vLLM's TileLang version; Nemotron uses Mamba-2's Triton kernels.
+    "tilelang==0.1.12",
     # Shared by vLLM, TileLang, and the pinned FlashAttention 4 runtime.
     "apache-tvm-ffi==0.1.11",
     "nvidia-cudnn-cu13>=9.15 ; sys_platform == 'linux'",
@@ -285,12 +288,13 @@ prime-kernels = [
     { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/prime_kernels-0.1.0+cu130torch2.13.0-cp312-cp312-linux_aarch64.whl", marker = "platform_machine == 'aarch64'" },
 ]
 flash-linear-attention = { git = "https://github.com/fla-org/flash-linear-attention" }
+mamba-ssm = { git = "https://github.com/state-spaces/mamba", rev = "f577286" }
 flash_attn_3 = { index = "pytorch-cu130-test" }
 
 [tool.uv.extra-build-dependencies]
 flash-attn = [{ requirement = "torch", match-runtime = true }]
 flash-attn-3 = [{ requirement = "torch", match-runtime = true }]
-mamba-ssm = [{ requirement = "torch", match-runtime = true }]
+mamba-ssm = ["torch>=2.13.0"]
 
 [tool.uv.extra-build-variables]
 flash-attn = { FLASH_ATTENTION_SKIP_CUDA_BUILD = "TRUE" }

skills/install/SKILL.md

@@ -59,10 +59,10 @@ See the `kernels` skill.
 ### NemotronH (Mamba SSD kernels)
 
 ```bash
-CUDA_HOME=/usr/local/cuda uv pip install mamba-ssm
+uv sync --extra gpu
 ```
 
-Requires `nvcc`. Without `mamba-ssm`, NemotronH falls back to HF's pure-PyTorch SSD path, which computes softplus in bf16 and yields ~0.4 KL divergence vs vLLM. Do **not** install `causal-conv1d` unless your GPU arch matches the prebuilt kernels — the code falls back to `nn.Conv1d` when it's absent.
+Native NemotronH uses Mamba2 Triton SSD kernels and FLA's packed causal convolution. The pinned `mamba-ssm` source installs without `nvcc`; no separate `causal-conv1d` installation is needed. `uv sync --all-extras` includes these dependencies.
 
 ### Trainer DeepEP backend
 

src/prime_rl/_compat.py

@@ -47,8 +47,7 @@ def _patched_lazy_load_kernel(kernel_name, mapping=_hub_kernels._KERNEL_MODULE_M
         try:
             return _original_lazy_load_kernel(kernel_name, mapping)
         except OfflineModeIsEnabled:
-            # Return None so NemotronH skips hub kernels; prime-rl's
-            # _patch_mamba2_use_triton_ssd() uses mamba_ssm directly.
+            # Return None so callers can use their non-hub implementation.
             mapping[kernel_name] = None
             return None
 

src/prime_rl/trainer/model.py

@@ -149,7 +149,7 @@ def _patch_qwen3_5_linear_attn_varlen():
 
     HF's forward hardcodes seq_idx=None for causal_conv1d and omits cu_seqlens
     for chunk_gated_delta_rule, so packed RL training sees ~0.23 Mismatch KL vs
-    vLLM (target <0.01). Mirrors the NemotronH mamba fix.
+    vLLM (target <0.01).
     """
     import torch.nn.functional as F
     from transformers.models.qwen3_5.modeling_qwen3_5 import (
@@ -630,15 +630,6 @@ def get_model(
 
     logger.debug(f"Loaded model config ({model_config.to_dict()})")
 
-    # NemotronH: transformers' Mamba2 mixer __init__ calls lazy_load_kernel("mamba-ssm" /
-    # "causal-conv1d") whenever config.use_mamba_kernels is set. That hub-kernel path is gated only
-    # by whether the `kernels` package is importable (NOT by USE_HUB_KERNELS) and resolves from the
-    # HF Hub, which hard-crashes under HF_HUB_OFFLINE=1. prime-rl swaps in its own mamba_ssm Triton
-    # SSD kernels via _patch_mamba2_use_triton_ssd, so the hub kernels are redundant; disable them
-    # to keep model init offline-safe.
-    if getattr(model_config, "model_type", "") == "nemotron_h":
-        model_config.use_mamba_kernels = False
-
     if config.debug.num_layers is not None:
         # VLM configs nest num_hidden_layers under text_config
         target_config = getattr(model_config, "text_config", model_config)

src/prime_rl/trainer/models/conversion_ops.py

@@ -27,10 +27,9 @@
 * **Every op is present-guarded** — it no-ops when its inputs are absent — so
   the same per-layer op list can be emitted for every layer and simply skips
   dense / non-matching layers, exactly as the imperative loops did.
-* **"Backward" means "what the imperative ``*_to_hf`` did"**, not a strict
-  mathematical inverse: a few conversions are intentionally lossy (e.g.
-  NemotronH shifts a router bias one way and does not undo it). Such ops carry
-  an explicit backward (see :class:`MapValue`).
+* **"Backward" means "what the imperative ``*_to_hf`` did"**, not necessarily
+  a strict mathematical inverse. Value transforms can provide an explicit
+  backward through :class:`MapValue`.
 """
 
 from __future__ import annotations

src/prime_rl/trainer/models/layers/ulysses_attn.py

@@ -15,9 +15,8 @@
     [S/cp, H,    D]  ◀── all-to-all ──   [S, H/cp, D]
                          (seq ↔ heads)
 
-Key benefit: the attention kernel itself does not need to be CP-aware. The
-all-to-all is purely on Q/K/V tensors, so this works out of the box with
-softmax flash-attn, linear attention, mamba, etc., without rewriting kernels.
+The softmax-attention kernel itself does not need to be CP-aware. Hybrid models
+use the same sequence/head redistribution in their owned recurrent modules.
 
 GQA models with fewer KV heads than cp_size (e.g. NemotronH: 32 query heads, 2
 KV heads) are handled by replicating each KV head cp_size / num_key_value_heads