[kevalmorabia97]

What does this PR do?

Type of change: CI / tests (follow-up to #2832)

Related: #2870 (independent, targets main) runs GPU CI only for PRs targeting main.

• Worker pool: a failing rank is now reported before its teardown runs. In gpu_megatron, teardown starts with dist.barrier(), which blocked against a still-running peer, so the parent never heard of the failure and the 30 s grace period could not start (it waited for the pytest timeout). New test fails without the fix. Also the remaining Fix post-merge unit CI and the gpu_megatron nightly timeout #2832 review nits (fixture-scope note, comments) and the sub-millisecond get_live_pool test moved to tests/unit.
• Wheel cache for the gpu job: mamba_ssm, causal-conv1d and fast-hadamard-transform were compiled from sdist on every run (~10 min of the job, e.g. 4 + 6 min on 10-16). They are now pinned in noxfile.py and built into a wheelhouse that actions/cache keeps per container image (key includes the noxfile.py hash); a wheel is only built when the wheelhouse lacks it. A plain wheelhouse rather than pip's or uv's cache: pip does not cache wheels it builds from a local source tree (the C++20-patched sdist) or reliably from a git URL, its cache lives under HOME (/github/home in container jobs), and the GPU sessions use the container's own pip (venv_backend="none"), not uv. A directory of wheels keyed on the image and noxfile.py is explicit and is saved only once complete.
• CUDA-extension cache never worked for the nightly: across 25 nightlies the gpu job never found its torch_extensions cache (205-350 s of extension compile each night; gpu_megatron hit about half the nights). actions/cache saves only when the job succeeds, and the nightly was failing or timing out every night, so main's cache scope was never populated (PR runs save into their own pull-request/N scope, which the nightly cannot read). The extension cache and the wheelhouse are now restored with actions/cache/restore and saved with actions/cache/save under always() (new .github/actions/save-extensions).
• Container bumps: pytorch 26.05/26.07 -> 26.09. TensorRT-LLM stays at rc20: no later release runs the multimodal quickstart that test_qwen_vl uses (rc21-rc23 import MultimodalConfig, which tensorrt_llm.llmapi does not export yet; rc24+ import fla, which the image lacks). NeMo 26.08 is already the newest public image. vLLM stays at v0.30.0: v0.31.0 breaks the vLLM plugin (test_kimi_gemm_rs_ar_bypasses_enabled_fake_quantizers: gemm_rs_ar has no .apply; test_flashinfer_metadata_builder_patch_stashes_common_metadata: FlashInferMetadataBuilder.cache_config; test_tiny_llama_quantize times out) and needs plugin changes by its owners. The gpu_vllm job on v0.20.0 (legacy FusedMoE registration, vLLM < 0.24) is dropped; v0.30.0 is the only vLLM job, and its timeout goes from 15 to 30 min: on PRs (1 GPU) its tests take ~10 min plus ~5 min to pull the image, so it timed out twice with all tests passing; the 2-GPU nightly also runs the TP2 tests and needs ~22 min.
• Extension cache key: keyed only on the C++/CUDA sources and their build scripts, so editing a Triton kernel no longer forces a rebuild (that rebuild pushed gpu_vllm past its 15 min limit). The key now also includes the job name: example and GPU jobs on the same image shared one key, and an example job that builds no extensions finished first and saved an empty cache. gpu and gpu_megatron then restored it as an exact hit, recompiled on every run (test_cuda_ext took 35-46 s in every gpu run of this PR), and could never save their own (Unable to reserve cache ... another job may be creating this cache). This also affects main. The save step now skips when no .so was built, so a run cancelled before building anything cannot claim the key with an empty cache. The tensorrt:26.05 images used by the ONNX/diffusers examples stay: torch-tensorrt (<2.13) needs libnvinfer.so.10, which newer images drop.
• mamba_ssm / causal-conv1d on pytorch:26.09: their setup.py hardcodes -std=c++17 but the 26.09 torch headers need C++20 (ATen requires C++20), so the wheels are built with the standard patched to C++20 (checked to compile and import against torch 2.13 headers).
• NVFP4 scale conversion without TensorRT-LLM's auto_deploy (removed in rc29, so the next bump would otherwise break): ModelOpt's NVFP4 dequantize/export paths (and gpu_trtllm collection) failed with ModuleNotFoundError. The cutlass-swizzled block-scale conversion is a tensor permutation, now implemented in ModelOpt as the inverse of swizzle_nvfp4_scales (same math as TensorRT-LLM's helper; bit-identical to rc20's on 6 shapes) with a CPU round-trip test.
• conv3d implicit-GEMM JIT build: implicit_gemm_cuda.py passed -std=c++17, which fails against the 26.09 torch headers (165 GPU test failures); torch's load() now picks the standard.
• FSDP2 + layerwise calibration bug: persistent_materialization patched FSDPParamGroup.unshard/reshard on the class, so while one layer was materialized every FSDP module skipped unshard; a full forward then ran on sharded DTensor params (test_layerwise_calibrate_fsdp2 failed on torch 2.14 with 'FSDPParam' object has no attribute '_unsharded_param'). Now patched only on the param groups inside the layer, and only when the layer itself is FSDP-wrapped (otherwise nothing was gathered for its FSDP submodules, so they must still unshard). CPU/gloo regression test fails without the fix; a second test covers a shared fully_shard([a, b]) group, nested contexts, a pre-existing instance override and a non-FSDP layer (fails without the group dedup or the FSDP-wrapped check).
• EAGLE training crash (found by the new container): hf_eagle.py compiled _eagle_forward with mode="max-autotune", which enables inductor CUDA graphs; recording them intermittently fails (cudaStreamCaptureStatusInvalidated / "graph recording observed an input tensor deallocate"). It failed test_llama_eagle3 on 5 of the last 10 nightlies and on the pytorch:26.09 dispatch run. Now max-autotune-no-cudagraphs (same autotuning, no graphs). Training-speed impact not measured.
• NVFP4 FP8-sweep kernel compile time: the kernel was autotuned over three tile configs, so first use compiled all three (~70 s per dtype on the pytorch:26.09 runner), pushing the MSE-sweep export and calibration tests past the 120 s timeout. It now uses the config that was fastest on B300 (BLOCKS_PER_PROGRAM=64, 8 warps; on Ada all three were within ~8% from 256² to 16384x8192). Cold first call drops from 13.6 s to 5.4 s locally, which also helps users' first MSE calibration in each process.
• Stale xfails: removed the five xfail(raises=RuntimeError) markers in test_megatron_preprocess_data.py; they xpassed in all 151 PR runs and every nightly that reached them over the last week (one of them never touches HF).
• Dedupe: removed tests/gpu_megatron/_extensions/test_torch_extensions.py, an exact copy of tests/gpu/_extensions/ (the megatron conftest already prebuilds the extensions).
• Puzzletron e2e trimmed: Puzzletron is not actively maintained, so CI runs only Nemotron-3-Nano (hybrid MoE), gpt-oss and Qwen3-VL; the other six models are manual (~40 s each, ~4 min of the gpu job). Run them with --run-manual.
• New 2-GPU nightly failures (from tests merged Oct 19): gpu_vllm's new TP2 tests (Sync the amax of vLLM's tensor-parallel linears across TP ranks #2842) failed because vLLM needs 160 MiB of /dev/shm and the container had Docker's default 64 MiB; the GPU test container now gets --shm-size=2gb, like the example jobs. test_distill_llm_sft_hf_chat[True] (Megatron-Bridge: GGML PTQ, QAD and fake-quant export for large MoE and VLM models #2873, packed SFT on a tiny MoE at TP2) failed on the sm_120 runner with "No dot product attention backend is available": the tiny model's head_dim of 2 leaves only TE's unfused attention, which the 2-GPU sm_120 runner rejects for packed input (it passes on Ada); the packed case now uses head_dim 8 so the flash and cuDNN backends also qualify. Both only run on the 2-GPU nightly, not in single-GPU PR CI.
• gpu_megatron timeout 60 -> 75 min: PR runs (1 GPU) take ~42 min, but the nightly runs every test on 2 GPUs: 46.5 min on Oct 19, then cancelled at 60 min with 99% done on Oct 20, and at 89% in this branch's dispatch run.
• Not included: destroying each test's NCCL groups in the worker teardown. An earlier commit here did that to stop communicators piling up in the persistent workers, but on the 1-GPU PR runners the branch then hit CUDA illegal memory accesses in random test_megatron.py tests (3 of 7 full-suite runs: test_dsa_kv_cache_quant, test_fp8_real_quantize, test_layerwise_calibrate_vlm_text_forward_loop), never with main's conftest (0 of 2 bisect runs, none on main). Reverted (6bc0119); the 2-GPU full suite takes the same time without it (43:36 vs 44:21), so the 60→75 min job timeout covers it.
• Unit CI timeout on main (Oct 20): test_hf_dflash2.py::test_overfits_a_single_batch took 51-61 s on most unit runners (2-7 s on others) and hit its 60 s timeout, also in main's nightly; both accuracies pass 0.9 by step ~21, so it now runs 30 steps instead of 60 (8 s on one CPU thread, was 16 s).
• test_mamba_search_space timeout 120 -> 150 s: across ~135 PR gpu_megatron jobs on Oct 16-17 it usually took 64-76 s, ~15% took 85-103 s on slow runners, and it reached 112, 119 and 123 s (twice; one timed out here). Nightlies were ~72 s.

Testing

• Pool: the new test (rank 0 fails, rank 1 hangs, teardown barrier) fails without the fix (queue.Empty after the 30 s job timeout) and passes with it; CPU unit tests that use the shared pool pass.
• Cache helper: stub-session check, cold cache runs pip wheel per missing package, warm cache runs none and installs the cached wheels.
• Dispatched on this branch: Regression tests on pytorch:26.09 passed; Unit tests passed. PR CI on 9d9ab38: everything green except gpu (2 FP8-sweep timeouts, fixed above); gpu_megatron passed in 32.5 min (60 min limit); the torch-extension cache was restored.
• FP8 sweep: tests/gpu/torch/quantization/test_nvfp4_fp8_sweep_kernel.py and the accelerate MSE tests pass locally (Ada, torch 2.14); the hf_ptq.py export case was not run locally.
• PR CI on f222d46: all checks green except gpu_vllm (all 190 tests passed, then hit the old 15 min limit) and gpu_megatron (test_mamba_search_space at 123 s, timeout raised above). The gpu job restored the wheelhouse and installed all three wheels without compiling.
• FSDP: tests/unit/torch/quantization/test_utils.py passes on CPU; GPU FSDP2/layerwise/export tests pass in nemo:26.08 (56 passed, 1 skipped).
• Puzzletron/mamba: in nemo:26.08 the six puzzletron models skip as manual and test_mamba_search_space passes with the new marker.
• PR CI on f1311ea (per-job cache keys): all checks green. Each GPU job started cold on its new key and saved a populated extension cache (gpu 7 MB, gpu_megatron 5 MB, gpu_trtllm 5 MB, gpu_vllm 2 MB; before, gpu and gpu_megatron only ever had a 0 MB one). Re-running gpu_trtllm restored its cache: test_cuda_ext 33.2 s -> 3.1 s and the pytest step 202 s -> 79 s.
• Eagle mode: torch._inductor.list_mode_options confirms max-autotune sets triton.cudagraphs and the new mode does not; tests/unit/torch/speculative (339) and tests/gpu/torch/speculative (14) pass.
• 2-GPU nightly fixes (8f58d40): not reproduced locally (the packed test passes on Ada). Dispatched on this branch (2-GPU runners): Example tests all 13 lanes passed, including test_distill_llm_sft_hf_chat[True]; in GPU tests gpu and gpu_trtllm passed and both gpu_vllm TP2 tests passed, but gpu_vllm hit its 20 min limit at 51% of tests (raised to 30 min below) and gpu_megatron hit its 60 min limit at 89% (raised to 75 min below). Re-dispatched GPU tests on f25ff98: gpu, gpu_trtllm and gpu_vllm passed (gpu_vllm 211 passed in 15 min); gpu_megatron 402 passed and failed only test_fp8_real_quantize (the Blackwell flake above).

Looked at but left alone: the 11 unified-export cases (each hits a different fuse-flag branch). Shared worker pools across megatron modules and a preloaded forkserver (importing modelopt initializes CUDA) did not work out.

• Is this change backward compatible?: ✅
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: N/A
• Did you get Claude approval on this PR?: ❌

🤖 Generated with Claude Code

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

• Last reviewed commit: 27235c7
• Latest commit, not reviewed: cb1cac9

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/Model-Optimizer/.coderabbit.yaml
• Review profile: CHILL
• Plan: Enterprise
• Run ID: 003096da-2c11-4858-88da-ee5e5dca6398

📥 Commits

Reviewing files that changed from the base of the PR and between e0dafde and 27235c7.

📒 Files selected for processing: 6

• .github/workflows/gpu_tests.yml
• modelopt/onnx/quantization/autotune/benchmark.py
• noxfile.py
• tests/_test_utils/torch/distributed/utils.py
• tests/gpu/torch/utils/test_distributed_worker_pool.py
• tests/unit/torch/utils/test_get_live_pool.py

🚧 Files skipped from review as they are similar to previous changes: 1

• tests/unit/torch/utils/test_get_live_pool.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 7 remain after this review.

📝 Walkthrough

The pull request updates test container images and cache workflows, adds local NVFP4 scale conversion, changes TensorRT flag handling and Eagle’s compile mode, and reports distributed-worker errors before teardown. Related tests and fixture documentation are updated.

Changes: CI images and dependency caches

Layer / File(s)	Summary
Test container versions .github/workflows/example_tests.yml, .github/workflows/gpu_tests.yml, .github/workflows/regression_tests.yml	Updates PyTorch and TensorRT-LLM container versions used in test jobs.
GPU dependency wheel builds and cache noxfile.py, .github/workflows/gpu_tests.yml	Builds or reuses pinned GPU dependency wheels and installs them in the GPU session. Adds wheelhouse cache restore and save steps.
Extension cache restore and save .github/actions/cache-extensions/action.yml, .github/actions/save-extensions/action.yml, .github/workflows/_example_tests_runner.yml, .github/workflows/gpu_tests.yml	Exposes extension cache hit and key outputs, adds a cache save action, and wires cache saves into example and GPU workflows.

NVFP4 scale conversion

Layer / File(s)	Summary
Local NVFP4 scale conversion modelopt/torch/quantization/qtensor/nvfp4_tensor.py, modelopt/torch/export/unified_export_hf.py, tests/gpu_trtllm/torch/quantization/backends/test_nvfp4_gemm.py, tests/unit/torch/quantization/test_nvfp4_tensor.py	Adds a local CUTLASS scale-layout converter and uses it in export and dequantization. Adds conversion tests and imports the helper from ModelOpt in the GPU test.

Distributed worker pool

Layer / File(s)	Summary
Report worker errors before teardown tests/_test_utils/torch/distributed/utils.py	Sends job errors before teardown and avoids replacing or duplicating the reported job error with a teardown error.
Worker failure and pool tests tests/gpu/torch/utils/test_distributed_worker_pool.py, tests/unit/torch/utils/test_get_live_pool.py, tests/gpu_megatron/conftest.py	Tests failure reporting while teardown blocks and tests live-pool reuse and broken-pool replacement. Updates fixture-scope documentation.

Runtime compatibility and compile mode

Layer / File(s)	Summary
Conditional TensorRT network flags modelopt/onnx/quantization/autotune/benchmark.py	Adds EXPLICIT_BATCH only when TensorRT exposes the flag. The strongly typed flag remains enabled.
Eagle forward compile mode modelopt/torch/speculative/plugins/hf_eagle.py	Changes _eagle_forward to use max-autotune-no-cudagraphs. Other compile targets are unchanged.

GPU test references

Layer / File(s)	Summary
GPU test comments and references tests/gpu/torch/speculative/plugins/test_hf_dflash.py, tests/gpu_megatron/_extensions/test_torch_extensions.py, tests/gpu_megatron/conftest.py	Revises the DFlash layer-placement comment, removes an extension-test forwarding file, and updates the documented extension-test location.

Priority: ➖ Normal

Estimated code review effort: 3 (Moderate) | ~25 minutes

Change: Other

Sequence Diagram(s)

sequenceDiagram
  participant GPUWorkflow
  participant WheelhouseCache
  participant NoxGPU
  GPUWorkflow->>WheelhouseCache: restore wheelhouse using image and noxfile key
  GPUWorkflow->>NoxGPU: run GPU session with MODELOPT_WHEELHOUSE
  NoxGPU->>WheelhouseCache: reuse or build and install selected wheels
  GPUWorkflow->>WheelhouseCache: save wheelhouse after cache miss and completion

Loading

sequenceDiagram
  participant TestWorkflow
  participant CacheExtensions
  participant SaveExtensions
  TestWorkflow->>CacheExtensions: restore extensions
  CacheExtensions-->>TestWorkflow: return cache hit and key
  TestWorkflow->>SaveExtensions: save extensions after cache miss

Loading

Suggested reviewers: chenhanyu, shengliangxu

Merge Risk: ⚪ Minimal · up to 27235

This change updates CI images, caches, and compatibility handling. No actionable merge-blocking risk remains after review.

Pre-merge checks: 5 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 46.43% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 28 functions across 12 files. (1 skipped:…	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 5 passed

Check name	Status	Explanation
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns		No listed security anti-pattern was introduced. The PR changes four Python files under modelopt and no Python files under examples. Their added diff lines contain no `torch.load(..., weights_only=…
Title check		The title clearly summarizes the main changes: CI cache updates, container image bumps, and GPU test changes. It is concise and related to the changeset.
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.

Full details: Security Anti-Patterns

Explanation

No listed security anti-pattern was introduced. The PR changes four Python files under modelopt and no Python files under examples. Their added diff lines contain no torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), hardcoded trust_remote_code=True, eval()/exec(), or # nosec. No pyproject.toml or requirements.txt file changed.

Full details: Docstring Coverage

Explanation

Docstring coverage is 46.43% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 28 functions across 12 files. (1 skipped: 1 unsupported.)

Resolution

Write docstrings for the functions missing them to satisfy the coverage threshold.

✨ Finishing Touches: 💡 1

📝 Generate docstrings: 💡

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Commit to this branch
• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

[kevalmorabia97]

/claude review

[github-actions]

PR Preview Action v1.8.1
Preview removed because the pull request was closed.
2026-10-20 17:24 UTC

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 78.80%. Comparing base (28a38b6) to head (cb1cac9).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2836      +/-   ##
==========================================
+ Coverage   76.90%   78.80%   +1.89%     
==========================================
  Files         658      658              
  Lines       74023    74029       +6     
==========================================
+ Hits        56930    58335    +1405     
+ Misses      17093    15694    -1399     

Flag	Coverage Δ
examples-diffusers	20.05% <15.38%> (+<0.01%)	⬆️
examples-gpt-oss	13.51% <15.38%> (+<0.01%)	⬆️
examples-hf_ptq	23.19% <26.92%> (-0.04%)	⬇️
examples-llm_distill	13.59% <15.38%> (-0.01%)	⬇️
examples-llm_eval	17.24% <15.38%> (+<0.01%)	⬆️
examples-llm_qat	17.48% <15.38%> (-0.01%)	⬇️
examples-llm_sparsity	15.77% <15.38%> (+<0.01%)	⬆️
examples-megatron_bridge	26.41% <15.38%> (-0.17%)	⬇️
examples-specdec_bench	13.29% <15.38%> (+<0.01%)	⬆️
examples-speculative_decoding	17.47% <15.38%> (-0.06%)	⬇️
examples-torch_onnx	21.40% <15.38%> (+<0.01%)	⬆️
examples-torch_trt	15.24% <15.38%> (+<0.01%)	⬆️
examples-vllm_serve	15.09% <15.38%> (+<0.01%)	⬆️
gpu	59.28% <100.00%> (+8.23%)	⬆️
regression	15.10% <15.38%> (-0.02%)	⬇️
unit	60.46% <84.61%> (+0.03%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[claude]

Claude review — CI follow-up PR

Scope: full review (the trigger comment was a bare /claude review). 11 changed files, all reviewed: modelopt/onnx/quantization/autotune/benchmark.py, noxfile.py, the three workflow files, and the five test files. Note that a two-dot git diff origin/main HEAD in the shallow CI checkout surfaces unrelated main-only changes (the external-draft removal, find_final_norm_path); I scoped the review to the 11 files GitHub reports as this PR's, and ignored that noise.

Findings — CRITICAL: 1, IMPORTANT: 4, SUGGESTION: 3

Most impactful

1. The wheel cache probably doesn't work at all (CRITICAL, noxfile.py:131). The workflow caches /root/.cache/wheelhouse and never sets MODELOPT_WHEELHOUSE, so the noxfile falls back to ~/.cache/wheelhouse. In a GitHub Actions container job HOME is remapped to /github/home, not the image's /root — so the restore fills a directory nothing reads, the glob always misses, and all three packages recompile on every run. That is precisely the ~10 min this PR exists to remove. The sibling .github/actions/cache-extensions avoids this by writing the literal path into $GITHUB_ENV rather than relying on ~; doing the same here is a one-line fix. A single echo "$HOME" in the gpu job confirms or refutes it before merge.

2. A partial wheelhouse can be frozen into the cache key permanently (IMPORTANT, gpu_tests.yml:104). always() is justified for a failing test, but a timeout can also land between wheel builds. Saving a one-of-two wheelhouse, then skipping the save on every subsequent run because cache-hit == 'true', leaves the missing wheel rebuilt forever — and the key only changes when noxfile.py does, so it cannot self-heal. Suggested either a completeness marker or run_id-suffixed keys with restore-keys.

3. The container bump collides with the repo's own cppimport note (IMPORTANT, gpu_tests.yml:47). noxfile.py still says tests/gpu/_extensions/test_onnx_extensions.py "fails for newer containers until cppimport#98", and that test is an unconditional assert with no skip marker. 26.05 → 26.09 is squarely "newer". Either the note is stale and should be deleted in this PR, or the test needs a guard. The PR body says the workflows have not yet been run on the new images, so this is unverified either way.

4. Teardown failures are now silently swallowed on the failure path (IMPORTANT, distributed/utils.py:200-210). The early put is the right fix for the reported hang, but the finally block still assembles a [teardown] … traceback that the if not reported guard then discards. That diagnostic is most valuable exactly in the gpu_megatron scenario this PR targets, where the failing rank's own teardown barrier is the thing that blocks.

5. The vLLM pin the noxfile says must stay in sync is now out of sync (IMPORTANT, example_tests.yml:151). Workflows go to v0.31.0; examples/vllm_serve/Dockerfile and its README stay at 0.30.0, so CI validates a version users building the shipped Dockerfile don't get.

Also raised three non-blocking SUGGESTIONs: glob prefix-collision and a bare IndexError in _install_built_from_source, a thin 5 s timing margin in the new pool test, and a 2024 copyright year on a file added in 2026.

Assessment

The pool fix itself is correct and well-targeted — reporting before teardown is the right resolution for a teardown barrier that blocks against a still-running peer, and the new test genuinely fails without it and terminates fast rather than hanging CI. The EXPLICIT_BATCH guard is right and consistent with the other create_network call site. The dedupe of the symlinked megatron extension test is sound given the conftest already prebuilds the extensions.

Risk: medium, and concentrated in the CI plumbing rather than the library. Nothing here can affect user-facing behavior — benchmark.py is the only modelopt/ change and it only widens TensorRT version tolerance. But the headline benefit of the PR (the wheel cache) looks likely to be a no-op as written, and two of the container bumps are unverified against the images. Running the gpu job on this branch would resolve findings 1 and 3 together and is worth doing before merge.

[yeyu-nvidia]

Reviewed the hf_eagle.py hunk you asked about. The change is right and I'd land it — flaky nightlies beat unmeasured throughput, and max-autotune-no-cudagraphs is the correct knob: same Triton autotuning, no graph capture.

Two observations, neither blocking.

1. Two comments in the same file now describe a constraint that no longer applies

hf_eagle.py:762 and :910 both carry:

# Init rotary_emb outside _eagle_forward so inv_freq is not captured by CUDAGraph.
self.eagle_module._maybe_init_rope(device=...)

Those hoists exist because _eagle_forward was capturing CUDA graphs. With graphs off, the stated reason is gone while the code stays — and the next person to read it will either think graphs are still in play, or move the call back in without realising it was load-bearing for a reason that may return.

I would not move the calls; the hoist is harmless and keeps the door open if graphs are re-enabled. But a word in each comment saying the capture is currently disabled and this guards re-enabling would stop the rationale going stale. Your call whether that belongs in this PR or a follow-up.

2. suppress_errors = True means we would not have heard about this from a user

_activate_torch_compile sets torch._dynamo.config.suppress_errors = True four lines above the change. A compile-time failure degrades to eager silently; what made these failures visible is that cudagraph recording blows up at runtime rather than at compile. So the only reason this surfaced at all is the nightly, and only as a hard crash — a quieter regression in the same block would be invisible.

Not something to fix here. Worth knowing that this file has no signal between "works" and "crashes the test".

On the unmeasured speed impact

Honest to state it, and I would not gate on it. If it does get measured later, the comparison worth having is EAGLE3 training step time with max-autotune vs max-autotune-no-cudagraphs on a real target rather than the tiny CI model — graph capture pays off on many-small-kernel workloads, and EAGLE3's draft forward at training shapes is mostly large GEMMs, so my expectation is the delta is small. If it turns out not to be, the better fix is probably making capture robust (static input buffers) rather than reverting, since the recording failure is a lifetime problem, not an autotuning one.

I did not review the rest of the PR — CI infrastructure, the worker-pool fix, and the quantization changes are outside what you asked me for and outside what I can usefully check.

[yeyu-nvidia]

Approving on the basis of the modelopt/torch/speculative/plugins/hf_eagle.py change only — the one @kevalmorabia97 asked me to look at. I have not reviewed the CI infrastructure, the worker-pool fix or the quantization changes in this PR.

max-autotune → max-autotune-no-cudagraphs is the right call and the right knob: identical Triton autotuning, no graph capture. Trading unmeasured throughput for a nightly that stops failing 5 runs in 10 is a good trade, and the comment records why.

Detail in my earlier comment: #2836 (review) — two non-blocking notes, the stale CUDAGraph rationale now left behind at hf_eagle.py:762 and :910, and suppress_errors = True four lines above meaning this block has no signal between working and crashing the nightly.

[yeyu-nvidia]

Approving on the basis of the hf_eagle.py change only, which is what @kevalmorabia97 asked me to look at. max-autotune -> max-autotune-no-cudagraphs is the right knob: same Triton autotuning, no graph capture. Detail in #2836 (review)

[kevalmorabia97]

@yeyu-nvidia Thanks. Addressed note 1 in cb1cac9: both rotary_emb hoist comments in hf_eagle.py now say capture is currently off and the hoist keeps re-enabling it safe. Code unchanged. Leaving suppress_errors as is, per your note.