[https://nvbugs/6419139][test] Guard against CUDA context creation at import#15985
Open
thorjohnsen wants to merge 1 commit into
Open
[https://nvbugs/6419139][test] Guard against CUDA context creation at import#15985thorjohnsen wants to merge 1 commit into
thorjohnsen wants to merge 1 commit into
Conversation
Contributor
📝 WalkthroughWalkthroughDeepGEMM PDL initialization is refactored from an eager, import-time call into a lazy, one-time initializer guarded by a module-level flag. The initializer is now invoked at the start of FP8 quantization, MoE grouped GEMM, and sparse attention metadata initialization. New tests verify no CUDA side effects occur at import time. ChangesLazy DeepGEMM PDL Initialization
Estimated code review effort: 2 (Simple) | ~15 minutes Sequence Diagram(s)sequenceDiagram
participant Caller as FP8Quantize / MoEGemm / DSAMetadata
participant Init as _init_deep_gemm_pdl
participant DeepGEMM as deep_gemm
Caller->>Init: _init_deep_gemm_pdl()
alt not yet initialized and CUDA available
Init->>DeepGEMM: set_pdl(get_env_enable_pdl())
Init->>Init: set _deep_gemm_pdl_initialized = True
else already initialized
Init-->>Caller: no-op
end
Caller->>DeepGEMM: proceed with GEMM/quantize/metadata setup
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
tests/unittest/others/test_import_side_effects.py (1)
78-93: 🎯 Functional Correctness | 🔵 TrivialCoverage gap: no positive test that PDL init actually fires on first DeepGEMM use.
The two tests here only verify the flag stays
Falseat import time; there's no test exercising one of the three call sites (_fp8_quantize_1x128_ue8m0,deepgemm_fp8_group_blockwise_gemm,DSAtrtllmAttentionMetadata.__init__) and asserting_deep_gemm_pdl_initializedbecomesTrueafterward. As per path instructions, coverage here is currently insufficient for verifying the "deferred, not removed" behavior.Do you want me to draft a CUDA-gated test that calls
_init_deep_gemm_pdl()(or one of the entry points) directly and asserts the flag flips toTrue?🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/others/test_import_side_effects.py` around lines 78 - 93, Add a positive CUDA-gated test that exercises one of the real DeepGEMM PDL entry points and verifies deferred initialization actually happens on first use. In tests/unittest/others/test_import_side_effects.py, extend the existing import-side-effect coverage by invoking a call site such as _fp8_quantize_1x128_ue8m0, deepgemm_fp8_group_blockwise_gemm, or DSAtrtllmAttentionMetadata.__init__ after import, then assert _deep_gemm_pdl_initialized flips from False to True. Keep the current import-time tests, but add this runtime assertion so the behavior is verified as “lazy initialization,” not just “never initializes.”Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unittest/others/test_import_side_effects.py`:
- Around line 78-79: The test functions are missing explicit return type
annotations, so update the affected test definitions to include a None return
type. Apply this to the named test helper(s) in test_import_side_effects.py,
including test_deep_gemm_pdl_init_is_lazy and the other test referenced in the
review, keeping the function behavior unchanged while conforming to the function
annotation guideline.
- Around line 43-62: The import-side-effect test script assumes
CUDA_VISIBLE_DEVICES is always a numeric index, but it may be a UUID-style
selector and fail before the NVML check runs. Update _NO_CONTEXT_SCRIPT to
resolve the NVML device handle from the CUDA_VISIBLE_DEVICES token in a way that
works for both numeric indices and non-numeric identifiers, so the subprocess
can still reach the tensorrt_llm import assertion.
---
Nitpick comments:
In `@tests/unittest/others/test_import_side_effects.py`:
- Around line 78-93: Add a positive CUDA-gated test that exercises one of the
real DeepGEMM PDL entry points and verifies deferred initialization actually
happens on first use. In tests/unittest/others/test_import_side_effects.py,
extend the existing import-side-effect coverage by invoking a call site such as
_fp8_quantize_1x128_ue8m0, deepgemm_fp8_group_blockwise_gemm, or
DSAtrtllmAttentionMetadata.__init__ after import, then assert
_deep_gemm_pdl_initialized flips from False to True. Keep the current
import-time tests, but add this runtime assertion so the behavior is verified as
“lazy initialization,” not just “never initializes.”
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: bc7dbc13-7210-4c89-bd92-c16b925648e2
📒 Files selected for processing (4)
tensorrt_llm/_torch/attention_backend/sparse/dsa.pytensorrt_llm/_torch/custom_ops/torch_custom_ops.pytensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.pytests/unittest/others/test_import_side_effects.py
Comment on lines
+43
to
+62
| _NO_CONTEXT_SCRIPT = r""" | ||
| import os | ||
| import sys | ||
|
|
||
| import pynvml | ||
|
|
||
| pynvml.nvmlInit() | ||
| # NVML indices are physical and ignore CUDA_VISIBLE_DEVICES; the test runner | ||
| # pins CUDA_VISIBLE_DEVICES to a single physical index, so use that one. | ||
| physical = int(os.environ["CUDA_VISIBLE_DEVICES"]) | ||
| handle = pynvml.nvmlDeviceGetHandleByIndex(physical) | ||
|
|
||
| import tensorrt_llm # noqa: F401 (the import under test) | ||
|
|
||
| procs = pynvml.nvmlDeviceGetComputeRunningProcesses(handle) | ||
| mine = [p for p in procs if p.pid == os.getpid()] | ||
| if mine: | ||
| used = (mine[0].usedGpuMemory or 0) >> 20 | ||
| sys.exit(f"import tensorrt_llm created a CUDA context using {used} MiB") | ||
| """ |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check how CI configures CUDA_VISIBLE_DEVICES for GPU test jobs.
rg -n "CUDA_VISIBLE_DEVICES" --type=yaml
rg -n "CUDA_VISIBLE_DEVICES" -g '*.groovy' -g '*.sh'Repository: NVIDIA/TensorRT-LLM
Length of output: 212
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- test_import_side_effects.py (outline) ---'
ast-grep outline tests/unittest/others/test_import_side_effects.py --view expanded || true
echo
echo '--- test_import_side_effects.py (lines 1-220) ---'
sed -n '1,220p' tests/unittest/others/test_import_side_effects.py
echo
echo '--- search for CUDA_VISIBLE_DEVICES in repo ---'
rg -n "CUDA_VISIBLE_DEVICES" .
echo
echo '--- likely CI / workflow files mentioning GPU or nvidia ---'
find . -path '*/.github/workflows/*' -o -path '*/.gitlab-ci*' -o -path '*/Jenkinsfile*' -o -path '*/.azure/*' -o -path '*/.circleci/*' | sed 's#^\./##' | sort
echo
echo '--- search workflows for GPU env setup ---'
rg -n "CUDA_VISIBLE_DEVICES|pynvml|nvml|gpu" .github/workflows .gitlab-ci* Jenkinsfile* .azure .circleci 2>/dev/null || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 19103
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- target file with line numbers ---'
cat -n tests/unittest/others/test_import_side_effects.py | sed -n '1,220p'
echo
echo '--- repo files under .github/workflows ---'
if [ -d .github/workflows ]; then
find .github/workflows -maxdepth 1 -type f | sed 's#^\./##' | sort
else
echo '(none)'
fi
echo
echo '--- focused search in workflows / ci-like files ---'
rg -n "CUDA_VISIBLE_DEVICES|pynvml|nvml|nvidia-smi|GPU" .github/workflows .gitlab-ci* .azure .circleci scripts tests 2>/dev/null || true
echo
echo '--- test_import_side_effects.py references ---'
rg -n "test_import_side_effects|NO_CONTEXT_SCRIPT|_run_in_subprocess|pynvml" tests -g '*.py'Repository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- usages of pynvml ---'
rg -n "\bpynvml\b|import pynvml" . --glob '!**/.git/**'
echo
echo '--- dependency manifests mentioning pynvml/nvml ---'
rg -n "pynvml|nvidia-ml-py|nvml" pyproject.toml requirements*.txt setup.py setup.cfg tox.ini .github/workflows tests scripts 2>/dev/null || true
echo
echo '--- files that look like test dependency docs ---'
find . -maxdepth 2 \( -name 'requirements*.txt' -o -name 'pyproject.toml' -o -name 'setup.py' -o -name 'environment*.yml' -o -name 'dev*.txt' \) | sed 's#^\./##' | sortRepository: NVIDIA/TensorRT-LLM
Length of output: 32601
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- jenkins/L0_Test.groovy around CUDA_VISIBLE_DEVICES handling ---'
sed -n '820,855p' jenkins/L0_Test.groovy
echo
echo '--- l0-test workflow references ---'
rg -n "L0_Test|CUDA_VISIBLE_DEVICES|gpu-type|gpu" .github/workflows/l0-test.yml jenkins/L0_Test.groovy 2>/dev/null || true
echo
echo '--- any docs that mention UUID-style CUDA_VISIBLE_DEVICES ---'
rg -n "CUDA_VISIBLE_DEVICES.*UUID|UUID.*CUDA_VISIBLE_DEVICES|GPU-" tests docs examples jenkins .github 2>/dev/null || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 19683
CUDA_VISIBLE_DEVICES needs a non-numeric fallback. int(os.environ["CUDA_VISIBLE_DEVICES"]) will fail for UUID-style device selectors, so this subprocess can abort before it checks the import side effect. Resolve the NVML handle from the token instead of assuming a bare index.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unittest/others/test_import_side_effects.py` around lines 43 - 62, The
import-side-effect test script assumes CUDA_VISIBLE_DEVICES is always a numeric
index, but it may be a UUID-style selector and fail before the NVML check runs.
Update _NO_CONTEXT_SCRIPT to resolve the NVML device handle from the
CUDA_VISIBLE_DEVICES token in a way that works for both numeric indices and
non-numeric identifiers, so the subprocess can still reach the tensorrt_llm
import assertion.
Comment on lines
+78
to
+79
| def test_deep_gemm_pdl_init_is_lazy(): | ||
| """DeepGEMM PDL setup must not run at import (nvbug 6419139). |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Missing return type annotations on test functions.
As per coding guidelines, "Always annotate functions with return types (use None if no return)."
🧹 Proposed fix
-def test_deep_gemm_pdl_init_is_lazy():
+def test_deep_gemm_pdl_init_is_lazy() -> None:-def test_import_creates_no_cuda_context():
+def test_import_creates_no_cuda_context() -> None:Also applies to: 91-91
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unittest/others/test_import_side_effects.py` around lines 78 - 79, The
test functions are missing explicit return type annotations, so update the
affected test definitions to include a None return type. Apply this to the named
test helper(s) in test_import_side_effects.py, including
test_deep_gemm_pdl_init_is_lazy and the other test referenced in the review,
keeping the function behavior unchanged while conforming to the function
annotation guideline.
Source: Coding guidelines
… import Between 1.3.0rc19 and rc20, an import-time deep_gemm.set_pdl() call (added in NVIDIA#15402) instantiated DeepGEMM's DeviceRuntime (cuBLASLt handle + 32 MiB workspace tensor), creating a CUDA context (~0.5-1.2 GiB with loaded modules, arch-dependent) in every process importing tensorrt_llm: the trtllm-bench parent (which never launches a kernel) and every MPI worker on its default device before torch.cuda.set_device. Those contexts are resident when the KV cache pool is sized from free GPU memory, shrinking the pool (RTX 6000D: 48.33 -> 47.39 GiB, max concurrent requests 513 -> 503) and causing ~10% throughput regressions on memory-tight GPUs (this bug and nvbugs 6405760, 6418453, 6419078) via a long low-batch scheduling tail. The root cause was fixed on main by NVIDIA#15632, which defers PDL configuration to PyTorchModelEngine init. Add a regression test asserting that importing tensorrt_llm neither configures DeepGEMM PDL nor creates a CUDA context, so the import path stays free of device side effects. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
thorjohnsen
force-pushed
the
fix/nvbug-6419139-deep-gemm-pdl-import-context
branch
from
July 6, 2026 17:28
deac457 to
a9b53e6
Compare
thorjohnsen
changed the title
liji-nv
approved these changes
Jul 16, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
@coderabbitai summary
Description
Adds a regression test guarding the root cause of nvbug 6419139 (~10% throughput regression on RTX 6000D for
llama_v3.3_nemotron_super_49b_fp8, 1.3.0rc19 → rc20; same root cause as nvbugs 6405760, 6418453, 6419078).Between rc19 and rc20, #15402 added an import-time
deep_gemm.set_pdl()call.set_pdl()instantiates DeepGEMM'sDeviceRuntime(cuBLASLt handle + 32 MiB workspace tensor), creating a CUDA context (~0.5–1.2 GiB with loaded modules, arch-dependent) in every process importing tensorrt_llm — including thetrtllm-benchparent (which never launches a kernel) and every MPI worker on its default device beforetorch.cuda.set_device(). Those contexts are resident when the KV cache pool is sized from free GPU memory, shrinking the pool (RTX 6000D: 48.33 → 47.39 GiB, max concurrent requests 513 → 503). With the QA benchmark's 512 uniform requests, 9 spill into a ~750-iteration low-batch tail = +10% inference time; GPUs with more headroom lose the same memory but show no regression, which made this look GPU-specific.The root cause was independently fixed on main by #15632 (
_configure_deep_gemm_pdl()deferred toPyTorchModelEngineinit) — this PR originally carried an equivalent fix and has been rebased to keep only the regression test, so the import path stays free of device side effects going forward.Verification of the mechanism (rc19/rc20 release wheels, 2×H100, exact QA bench command): available KV cache memory 44.14 GiB (rc19) vs 43.59 GiB (rc20, deterministic across 9 runs); disabling only the import-time call restores 44.14 GiB exactly;
import tensorrt_llmcreates a 550 MiB context on rc20 and none on rc19.Test Coverage
tests/unittest/others/test_import_side_effects.pytest_deep_gemm_pdl_configuration_is_lazy— assertsimport tensorrt_llmdoes not configure DeepGEMM PDL (no GPU required).test_import_creates_no_cuda_context— assertsimport tensorrt_llmcreates no CUDA context (pynvml, GPU-gated).PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
🤖 Generated with Claude Code