[https://nvbugs/6729495][fix] Restore Nemotron-H weight-update coverage

jenkins/L0_Test.groovy

@@ -5114,21 +5114,7 @@ def runLLMTestlistOnPlatformImpl(pipeline, platform, testList, config=VANILLA_CO
             trtllm_utils.llmExecStepWithRetry(pipeline, script: "cd ${llmSrc} && pip3 install -r requirements-grpc-smg.txt")
             if (stageName.contains("-Ray-")) {
                 trtllm_utils.llmExecStepWithRetry(pipeline, script: "pip3 install ray[default]==2.55.1")
-                // TODO(dlfw-26.08): reinstate causal-conv1d and mamba-ssm once
-                // upstream publishes wheels built against this base image's torch.
-                // They used to be installed here, from
-                //   causal-conv1d v1.6.2  causal_conv1d-1.6.1+cu13torch26.04cxx11abiTRUE
-                //   mamba v2.3.0          mamba_ssm-2.3.0+cu13torch26.01cxx11abiTRUE
-                // but the newest builds upstream offers target torch 26.07 and 26.04,
-                // so on DLFW 26.08 the extension loads with an undefined c10 symbol,
-                // materialize_cow_storage(StorageImpl&). A broken install is worse
-                // than none: transformers gates its causal_conv1d import on a
-                // package-metadata probe, which a broken install still passes, so
-                // modeling_qwen3_5_moe raises at import and every test collected from
-                // a module that imports it dies as a collection error. Absent, the
-                // gate says no and the model falls back to its Python path -- slower,
-                // and it OOMs on Nemotron-H, which is why
-                // test_llm_update_weights_nemotron_h is waived under nvbugs/6729495.
+                trtllm_utils.llmExecStepWithRetry(pipeline, script: "bash ${llmSrc}/jenkins/scripts/install_mamba.sh")
             }
             if (!skipInstallWheel) {
                 trtllm_utils.llmExecStepWithRetry(pipeline, script: "cd ${llmPath} && pip3 install --force-reinstall --no-deps TensorRT-LLM/tensorrt_llm-*.whl")

jenkins/scripts/install_mamba.sh

@@ -0,0 +1,40 @@
+#!/bin/bash
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+set -euo pipefail
+
+build_dir=$(mktemp -d)
+trap 'rm -rf "$build_dir"' EXIT
+
+curl --fail --location --retry 3 \
+    https://github.com/Dao-AILab/causal-conv1d/archive/refs/tags/v1.7.0.tar.gz \
+    --output "$build_dir/causal-conv1d.tar.gz"
+# mamba-ssm 2.3.0 avoids the newer Mamba3/quack dependency chain.
+curl --fail --location --retry 3 \
+    https://github.com/state-spaces/mamba/archive/refs/tags/v2.3.0.tar.gz \
+    --output "$build_dir/mamba.tar.gz"
+tar -xzf "$build_dir/causal-conv1d.tar.gz" -C "$build_dir"
+tar -xzf "$build_dir/mamba.tar.gz" -C "$build_dir"
+
+# PyTorch 2.14 headers require C++20 in both C++ and CUDA translation units.
+sed -i 's/-std=c++17/-std=c++20/g' "$build_dir/mamba-2.3.0/setup.py"
+
+# Compile against the installed torch, including its headers and C++ ABI.
+export CAUSAL_CONV1D_FORCE_BUILD=TRUE
+export MAMBA_FORCE_BUILD=TRUE
+export MAX_JOBS="${MAX_JOBS:-8}"
+python3 -m pip install --no-build-isolation --no-deps --no-cache-dir \
+    "$build_dir/causal-conv1d-1.7.0" "$build_dir/mamba-2.3.0"
+
+python3 - <<'PY'
+import torch
+import causal_conv1d
+import causal_conv1d_cuda
+import mamba_ssm
+import selective_scan_cuda
+from transformers import Qwen3_5MoeForConditionalGeneration
+
+print(f"torch={torch.__version__}, CUDA={torch.version.cuda}")
+print(f"causal-conv1d={causal_conv1d.__version__}, mamba-ssm={mamba_ssm.__version__}")
+PY

jenkins/scripts/slurm_install.sh

@@ -74,23 +74,12 @@ slurm_install_setup() {
         nvidia-smi && nvidia-smi -q && nvidia-smi topo -m
         if [[ $pytestCommand == *--run-ray* ]]; then
             retry_command --timeout 2700 pip3 install --retries 10 "ray[default]==2.55.1"
-            # TODO(dlfw-26.08): reinstate causal-conv1d and mamba-ssm once upstream
-            # publishes wheels built against this base image's torch. They used to be
-            # installed here, from
-            #   causal-conv1d v1.6.2  causal_conv1d-1.6.1+cu13torch26.04cxx11abiTRUE
-            #   mamba v2.3.0          mamba_ssm-2.3.0+cu13torch26.01cxx11abiTRUE
-            # but the newest builds upstream offers target torch 26.07 and 26.04, so
-            # on DLFW 26.08 the extension loads with an undefined c10 symbol,
-            # materialize_cow_storage(StorageImpl&). A broken install is worse than
-            # none: transformers gates its causal_conv1d import on a package-metadata
-            # probe, which a broken install still passes, so modeling_qwen3_5_moe
-            # raises at import and every test collected from a module that imports it
-            # dies as a collection error. Absent, the gate says no and the model falls
-            # back to its Python path -- slower, and it OOMs on Nemotron-H, which is
-            # why test_llm_update_weights_nemotron_h is waived under nvbugs/6729495.
         fi
         retry_command --timeout 2700 bash -c "cd $llmSrcNode && pip3 install --retries 10 -r requirements-dev.txt"
         retry_command --timeout 2700 bash -c "cd $llmSrcNode && pip3 install --retries 10 -r requirements-grpc-smg.txt"
+        if [[ $pytestCommand == *--run-ray* ]]; then
+            retry_command --timeout 2700 bash "$llmSrcNode/jenkins/scripts/install_mamba.sh"
+        fi
         retry_command --timeout 2700 bash -c "cd $resourcePathNode && pip3 install --retries 10 --force-reinstall --no-deps TensorRT-LLM/tensorrt_llm-*.whl"
         gpuUuids=$(nvidia-smi -q | grep "GPU UUID" | awk '{print $4}' | tr '\n' ',' || true)
         hostNodeName="${HOST_NODE_NAME:-$(hostname -f || hostname)}"

tests/integration/test_lists/waives.txt

@@ -294,7 +294,6 @@ unittest/_torch/multi_gpu/test_minimax_m3_piecewise.py::test_minimax_m3_piecewis
 unittest/_torch/multimodal/test_mm_encoder_standalone.py::test_epd_disagg_mm_hash_kv_cache_reuse[prompts0] SKIP (https://nvbugs/6812085)
 unittest/_torch/multimodal/test_mm_encoder_standalone.py::test_epd_disagg_mm_hash_kv_cache_reuse[prompts1] SKIP (https://nvbugs/6842946)
 unittest/_torch/multimodal/test_mm_encoder_standalone.py::test_pd_disagg_multimodal_no_reuse_when_disabled SKIP (https://nvbugs/6813129)
-unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py::test_llm_update_weights_nemotron_h SKIP (https://nvbugs/6729495)
 unittest/_torch/sampler/test_beam_search_speculative_d2h.py::test_speculative_d2h_predictor_hit_is_sync_free SKIP (https://nvbugs/6859731)
 unittest/_torch/speculative/test_eagle3.py SKIP (https://nvbugs/6849051)
 unittest/_torch/thop/parallel/test_fp4_linear.py::test_fp4_gemm_bias_per_backend[mnk2-cublaslt] SKIP (https://nvbugs/6581067)

tests/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py

@@ -1135,21 +1135,21 @@ def test_llm_update_weights_nemotron_h():
 
     Requires mamba-ssm and causal-conv1d to be importable: without them HF
     falls back to the naive Python selective_scan path, which OOMs on
-    Nemotron-H and produces unmatched logits. The Ray CI stage used to install
-    both next to ray, and will again once upstream ships wheels for this base
-    image's torch -- see the TODO in jenkins/scripts/slurm_install.sh."""
+    Nemotron-H and produces unmatched logits. Ray CI builds both against the
+    installed torch with jenkins/scripts/install_mamba.sh."""
     try:
         import causal_conv1d  # noqa: F401
         import mamba_ssm  # noqa: F401
     except ImportError as e:
         # Fail loudly here rather than let the naive fallback OOM further in,
         # which is a much harder failure to read.
         pytest.fail(
-            f"{e.name} is not installed, so the mamba fast path is unavailable. "
-            "The Ray CI stage stopped installing mamba-ssm and causal-conv1d "
-            "because upstream has no wheel for this base image's torch; "
-            "see the TODO in jenkins/scripts/slurm_install.sh."
+            f"The mamba fast path is unavailable: {e}. "
+            "Build and install the extensions against the current torch with "
+            "bash jenkins/scripts/install_mamba.sh."
         )
+    from transformers.models.nemotron_h import modeling_nemotron_h
+
     model_dir = str(llm_models_root() / "NVIDIA-Nemotron-3-Nano-30B-A3B-BF16")
     num_hidden_layers = 7
     # NemotronHConfig derives num_hidden_layers from ``layers_block_type``
@@ -1160,6 +1160,9 @@ def test_llm_update_weights_nemotron_h():
     hf_model = RefHFModelWithIPCHandles(
         model_dir, num_hidden_layers=num_hidden_layers, layers_block_type=layers_block_type
     )
+    assert modeling_nemotron_h.is_fast_path_available, (
+        "HF Nemotron-H did not enable the Mamba fast path despite the installed extensions."
+    )
     tokenizer = AutoTokenizer.from_pretrained(model_dir)
     # Nemotron-H's Mamba state dominates the cache budget; 0.25 of free memory
     # leaves enough room for HF (resident on cuda:0 + replicas on cuda:1..3)