This repository was archived by the owner on Apr 20, 2026. It is now read-only.
trtllm nsys#231
Draft
karen-sy wants to merge 1 commit into
Draft
Conversation
Contributor
📝 WalkthroughWalkthroughThe PR adds TRTLLM-specific profiling support by introducing a Changes
Sequence Diagram(s)sequenceDiagram
participant Script as Profiling Script
participant Backend as Backend Decision
participant TRTLLM as TRTLLM Path
participant Router as Router Service
participant SGLang as SGLang Path
participant Eval as LM-Eval
Script->>Backend: Check PROFILING_BACKEND env var
alt PROFILING_BACKEND == trtllm
Backend->>TRTLLM: Route to TRTLLM profiling
TRTLLM->>Router: aiperf profile /v1/chat/completions
Router->>Router: Generate synthetic completions
TRTLLM->>TRTLLM: Export raw profile data to nsys_profile
else Default (sglang)
Backend->>SGLang: Route to SGLang profiling
SGLang->>SGLang: python3 -m sglang.bench_serving
SGLang->>Eval: Invoke lm-eval for coverage
Eval->>Eval: Evaluate benchmark results
end
Estimated code review effort🎯 2 (Simple) | ⏱️ ~12 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 inconclusive)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
⚔️ Resolve merge conflicts
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 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: 1
🧹 Nitpick comments (3)
src/srtctl/core/schema.py (1)
645-649: TRTLLM-specific env vars set unconditionally for nsys profiling.
TLLM_PROFILE_START_STOPandTLLM_LLMAPI_ENABLE_NVTXare TRTLLM-specific environment variables, but they're set for any backend whenis_nsysis true. While these variables are harmless for non-TRTLLM backends (they'll be ignored), it would be cleaner to only set them when the backend is actually TRTLLM.However, this method doesn't have access to the backend type. Consider whether this is acceptable as-is (extra env vars are benign) or if the backend type should be passed to this method for conditional logic.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/srtctl/core/schema.py` around lines 645 - 649, The code unconditionally sets TRTLLM-specific env vars (TLLM_PROFILE_START_STOP, TLLM_LLMAPI_ENABLE_NVTX) whenever is_nsys is true; update the logic to only set these when the backend is TRTLLM by either adding a backend_type parameter to the method or using an existing backend attribute on self. Concretely, change the method signature to accept backend_type (or read self.backend_type if available), then wrap the env assignments inside a conditional like: if is_nsys and backend_type == "trtllm" and phase_config and phase_config.start_step is not None and phase_config.stop_step is not None, set TLLM_PROFILE_START_STOP and TLLM_LLMAPI_ENABLE_NVTX; otherwise skip them. Ensure references to is_nsys, phase_config, TLLM_PROFILE_START_STOP, and TLLM_LLMAPI_ENABLE_NVTX are updated accordingly.src/srtctl/benchmarks/scripts/profiling/profile.sh (2)
138-159: Consider making--request-countconfigurable or documenting the hardcoded value.The TRTLLM path uses
--request-count 50while the SGLang path uses--num-prompts 128. If this difference is intentional (e.g., TRTLLM profiling needs fewer requests), a comment explaining the rationale would be helpful. Otherwise, consider using a common environment variable likePROFILE_NUM_REQUESTSfor consistency.Also, the
--tokenizer /model/path assumes the model is mounted at/model/which is consistent with the backend implementations shown in the context snippets.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/srtctl/benchmarks/scripts/profiling/profile.sh` around lines 138 - 159, In the TRTLLM branch of profile.sh (the if checking PROFILING_BACKEND == "trtllm") the aiperf profile call hardcodes --request-count 50; make this configurable by replacing the literal with an environment variable (e.g., use PROFILE_NUM_REQUESTS with a sensible default) and add a short comment explaining the chosen default or why TRTLLM differs from SGLang's --num-prompts; also document or assert the assumed tokenizer mount path (/model/) near the aiperf invocation to match other backend implementations.
176-184:lm-evalonly runs for SGLang backend - verify this is intentional.The
lm-evalevaluation step runs only in the SGLang branch. If this is intentional (e.g., TRTLLM profiling focuses solely on nsys captures), consider adding a comment explaining why. Otherwise, thelm-evalblock could be moved outside the if/else to run for both backends since it uses the OpenAI-compatible/v1/completionsendpoint.Also, the inline
pip installwith output suppression (> /dev/null 2>&1) could hide installation failures. Consider checking the exit code or at least preserving stderr.♻️ Suggested improvement for pip install
- pip install lm-eval tenacity > /dev/null 2>&1 + pip install lm-eval tenacity > /dev/null || echo "Warning: pip install failed"🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/srtctl/benchmarks/scripts/profiling/profile.sh` around lines 176 - 184, The lm-eval block (the python -m lm_eval invocation using base_url=http://${head_node}:${head_port}/v1/completions and variables like ${model_name} and ${PROFILE_CONCURRENCY}) currently executes only in the SGLang branch; decide whether it should run for both backends—if yes, move that entire lm-eval block out of the if/else so both branches run it, otherwise add a concise comment above the block explaining why TRTLLM is intentionally excluded. Also stop silently swallowing pip install failures: run pip install lm-eval tenacity without redirecting both stdout/stderr or check its exit code (e.g., test the install return code) and surface errors to the console before proceeding to python -m lm_eval.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/srtctl/benchmarks/scripts/profiling/profile.sh`:
- Around line 161-188: The script currently overwrites the primary benchmark's
exit status with lm-eval's by assigning exit_code=$? only at the end; modify the
SGLang branch to capture the exit status immediately after running python3 -m
sglang.bench_serving (e.g., save $? into a variable like bench_status), then run
lm-eval, and when computing exit_code prefer the primary benchmark's non-zero
status (set exit_code to bench_status if bench_status is non-zero, otherwise use
the lm-eval status or 0). Update references to the sglang.bench_serving command
and the exit_code variable accordingly so failures in the primary benchmark are
not masked.
---
Nitpick comments:
In `@src/srtctl/benchmarks/scripts/profiling/profile.sh`:
- Around line 138-159: In the TRTLLM branch of profile.sh (the if checking
PROFILING_BACKEND == "trtllm") the aiperf profile call hardcodes --request-count
50; make this configurable by replacing the literal with an environment variable
(e.g., use PROFILE_NUM_REQUESTS with a sensible default) and add a short comment
explaining the chosen default or why TRTLLM differs from SGLang's --num-prompts;
also document or assert the assumed tokenizer mount path (/model/) near the
aiperf invocation to match other backend implementations.
- Around line 176-184: The lm-eval block (the python -m lm_eval invocation using
base_url=http://${head_node}:${head_port}/v1/completions and variables like
${model_name} and ${PROFILE_CONCURRENCY}) currently executes only in the SGLang
branch; decide whether it should run for both backends—if yes, move that entire
lm-eval block out of the if/else so both branches run it, otherwise add a
concise comment above the block explaining why TRTLLM is intentionally excluded.
Also stop silently swallowing pip install failures: run pip install lm-eval
tenacity without redirecting both stdout/stderr or check its exit code (e.g.,
test the install return code) and surface errors to the console before
proceeding to python -m lm_eval.
In `@src/srtctl/core/schema.py`:
- Around line 645-649: The code unconditionally sets TRTLLM-specific env vars
(TLLM_PROFILE_START_STOP, TLLM_LLMAPI_ENABLE_NVTX) whenever is_nsys is true;
update the logic to only set these when the backend is TRTLLM by either adding a
backend_type parameter to the method or using an existing backend attribute on
self. Concretely, change the method signature to accept backend_type (or read
self.backend_type if available), then wrap the env assignments inside a
conditional like: if is_nsys and backend_type == "trtllm" and phase_config and
phase_config.start_step is not None and phase_config.stop_step is not None, set
TLLM_PROFILE_START_STOP and TLLM_LLMAPI_ENABLE_NVTX; otherwise skip them. Ensure
references to is_nsys, phase_config, TLLM_PROFILE_START_STOP, and
TLLM_LLMAPI_ENABLE_NVTX are updated accordingly.
🪄 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: defaults
Review profile: CHILL
Plan: Pro
Run ID: c09e9d5f-d92b-497d-bd88-dc87269d242f
📒 Files selected for processing (3)
src/srtctl/benchmarks/scripts/profiling/profile.shsrc/srtctl/cli/mixins/benchmark_stage.pysrc/srtctl/core/schema.py
Comment on lines
+161
to
188
| else | ||
| # SGLang: use sglang.bench_serving | ||
| python3 -m sglang.bench_serving \ | ||
| --backend sglang \ | ||
| --model "${model_name}" \ | ||
| --host "${head_node}" --port "${head_port}" \ | ||
| --dataset-name random \ | ||
| --max-concurrency "${PROFILE_CONCURRENCY}" \ | ||
| --num-prompts 128 \ | ||
| --random-input-len "${PROFILE_ISL}" \ | ||
| --random-output-len "${PROFILE_OSL}" \ | ||
| --random-range-ratio 1 \ | ||
| --warmup-request 0 | ||
|
|
||
|
|
||
| # lm-eval for additional coverage (uses OpenAI /v1/completions — works for both backends) | ||
| echo "" | ||
| echo "Running lm-eval..." | ||
| pip install lm-eval tenacity > /dev/null 2>&1 | ||
| python -m lm_eval \ | ||
| --model local-completions \ | ||
| --tasks gsm8k \ | ||
| --model_args "base_url=http://${head_node}:${head_port}/v1/completions,model=${model_name},tokenized_requests=False,tokenizer_backend=None,num_concurrent=${PROFILE_CONCURRENCY},timeout=6000,max_retries=1" \ | ||
| --limit 10 | ||
| fi | ||
| fi | ||
|
|
||
| exit_code=$? |
Contributor
There was a problem hiding this comment.
Exit code only captures the last command, potentially masking earlier failures.
In the SGLang branch, exit_code=$? (line 188) captures the exit code of lm-eval, not sglang.bench_serving. If the primary benchmark (bench_serving) fails but lm-eval succeeds, the failure is masked.
Consider capturing and checking the exit code of the primary benchmark command:
🐛 Proposed fix to capture primary benchmark exit code
else
# SGLang: use sglang.bench_serving
python3 -m sglang.bench_serving \
--backend sglang \
--model "${model_name}" \
--host "${head_node}" --port "${head_port}" \
--dataset-name random \
--max-concurrency "${PROFILE_CONCURRENCY}" \
--num-prompts 128 \
--random-input-len "${PROFILE_ISL}" \
--random-output-len "${PROFILE_OSL}" \
--random-range-ratio 1 \
--warmup-request 0
+ bench_exit_code=$?
# lm-eval for additional coverage (uses OpenAI /v1/completions — works for both backends)
echo ""
echo "Running lm-eval..."
pip install lm-eval tenacity > /dev/null 2>&1
python -m lm_eval \
--model local-completions \
--tasks gsm8k \
--model_args "base_url=http://${head_node}:${head_port}/v1/completions,model=${model_name},tokenized_requests=False,tokenizer_backend=None,num_concurrent=${PROFILE_CONCURRENCY},timeout=6000,max_retries=1" \
--limit 10
+
+ # Use bench_serving exit code as primary result
+ exit_code=${bench_exit_code}
fi
fi
-exit_code=$?📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| else | |
| # SGLang: use sglang.bench_serving | |
| python3 -m sglang.bench_serving \ | |
| --backend sglang \ | |
| --model "${model_name}" \ | |
| --host "${head_node}" --port "${head_port}" \ | |
| --dataset-name random \ | |
| --max-concurrency "${PROFILE_CONCURRENCY}" \ | |
| --num-prompts 128 \ | |
| --random-input-len "${PROFILE_ISL}" \ | |
| --random-output-len "${PROFILE_OSL}" \ | |
| --random-range-ratio 1 \ | |
| --warmup-request 0 | |
| # lm-eval for additional coverage (uses OpenAI /v1/completions — works for both backends) | |
| echo "" | |
| echo "Running lm-eval..." | |
| pip install lm-eval tenacity > /dev/null 2>&1 | |
| python -m lm_eval \ | |
| --model local-completions \ | |
| --tasks gsm8k \ | |
| --model_args "base_url=http://${head_node}:${head_port}/v1/completions,model=${model_name},tokenized_requests=False,tokenizer_backend=None,num_concurrent=${PROFILE_CONCURRENCY},timeout=6000,max_retries=1" \ | |
| --limit 10 | |
| fi | |
| fi | |
| exit_code=$? | |
| else | |
| # SGLang: use sglang.bench_serving | |
| python3 -m sglang.bench_serving \ | |
| --backend sglang \ | |
| --model "${model_name}" \ | |
| --host "${head_node}" --port "${head_port}" \ | |
| --dataset-name random \ | |
| --max-concurrency "${PROFILE_CONCURRENCY}" \ | |
| --num-prompts 128 \ | |
| --random-input-len "${PROFILE_ISL}" \ | |
| --random-output-len "${PROFILE_OSL}" \ | |
| --random-range-ratio 1 \ | |
| --warmup-request 0 | |
| bench_exit_code=$? | |
| # lm-eval for additional coverage (uses OpenAI /v1/completions — works for both backends) | |
| echo "" | |
| echo "Running lm-eval..." | |
| pip install lm-eval tenacity > /dev/null 2>&1 | |
| python -m lm_eval \ | |
| --model local-completions \ | |
| --tasks gsm8k \ | |
| --model_args "base_url=http://${head_node}:${head_port}/v1/completions,model=${model_name},tokenized_requests=False,tokenizer_backend=None,num_concurrent=${PROFILE_CONCURRENCY},timeout=6000,max_retries=1" \ | |
| --limit 10 | |
| # Use bench_serving exit code as primary result | |
| exit_code=${bench_exit_code} | |
| fi | |
| fi |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/srtctl/benchmarks/scripts/profiling/profile.sh` around lines 161 - 188,
The script currently overwrites the primary benchmark's exit status with
lm-eval's by assigning exit_code=$? only at the end; modify the SGLang branch to
capture the exit status immediately after running python3 -m
sglang.bench_serving (e.g., save $? into a variable like bench_status), then run
lm-eval, and when computing exit_code prefer the primary benchmark's non-zero
status (set exit_code to bench_status if bench_status is non-zero, otherwise use
the lm-eval status or 0). Update references to the sglang.bench_serving command
and the exit_code variable accordingly so failures in the primary benchmark are
not masked.
karen-sy
marked this pull request as draft
March 30, 2026 17:00
Contributor
Author
|
I should rebase changes onto Ishan main (which also will make it easier to add the other trtllm nsys profiling variants I have for trace- and time-based profiling) |
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 subscribe to this conversation on GitHub.
Already have an account?
Sign in.
1 participant
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.
Changes to enable range-based nsys profiling for TRTLLM deployments
src/srtctl/core/schema.py:• get_nsys_prefix now always uses cuda,nvtx,ucx (ucx added)
• New
continue_benchmark_after_profilngparameter (default: False for backward compat) to not kill benchmark when nsys range finishes. When True, adds--sample=none, --kill none, --wait all• Added
TLLM_PROFILE_START_STOPandTLLM_LLMAPI_ENABLE_NVTXenv vars for TRTLLM• Added validation: torch profiling not supported for trtllm
src/srtctl/benchmarks/scripts/profiling/profile.sh:• TRTLLM: uses aiperf for profiling traffic generation
• SGLang: uses sglang.bench_serving (unchanged)
src/srtctl/cli/mixins/benchmark_stage.py:• Added
PROFILING_BACKENDenv varExample YAML usage for trtllm:
Summary by CodeRabbit
New Features
continue_benchmark_after_profilingconfiguration option to control benchmark lifecycle after profilingImprovements