[https://nvbugs/6533916][fix] Make sparse attention example runnable - #17055
Conversation
WalkthroughThe sparse-attention example now supports a deterministic built-in prompt when no input file is provided. It retains JSONL input handling, updates CLI defaults and usage documentation, and conditionally displays reference outputs. ChangesSparse attention example
Estimated code review effort: 2 (Simple) | ~10 minutes Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant parse_arguments
participant run_llm
participant result_printing
parse_arguments->>run_llm: Provide model_path and optional input_file
run_llm->>run_llm: Select DEFAULT_PROMPTS or JSONL prompts
run_llm->>result_printing: Provide generated text and optional references
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
4edb499 to
59f6dbf
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@examples/llm-api/llm_sparse_attention.py`:
- Around line 27-38: Update the runnable command documentation to include the
--input_file option pointing to the packaged JSONL sample, alongside
--model_path, so it explicitly exercises the long-prompt RocketKV
sparse-attention path.
🪄 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: 4c346b6e-5db3-461f-af47-de5705d382ad
📒 Files selected for processing (1)
examples/llm-api/llm_sparse_attention.py
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
59f6dbf to
2ef156d
Compare
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
|
/bot run --disable-fail-fast |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@examples/llm-api/llm_sparse_attention.py`:
- Around line 77-91: Update the _build_default_prompt function signature with a
str return type annotation, preserving its existing prompt construction and
return behavior.
🪄 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: b951bf3e-3b34-48fb-8999-fbcfe370cdc3
📒 Files selected for processing (1)
examples/llm-api/llm_sparse_attention.py
|
PR_Github #62952 [ run ] triggered by Bot. Commit: |
|
PR_Github #62952 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63014 [ run ] triggered by Bot. Commit: |
|
PR_Github #63014 [ run ] completed with state |
BowenFu
left a comment
There was a problem hiding this comment.
Verified the --input_file path is byte-identical to before — same slicing, ordering, and prompt/reference construction — so only the omitted-file default changes, and LLM.generate normalizes the plain strings the new default passes just like the dicts the file branch still passes.
Not blocking, just flagging: the new --model_path default is the FP8 checkpoint, so a zero-arg run now needs an FP8-capable GPU where the old BF16 default did not. Worth a line in the usage block if you touch this again.
Dev Engineer Review
nvidia/Llama-3.1-8B-Instruct-FP8model by default.--input_fileoptional and added a deterministic needle-in-a-haystack prompt for zero-input execution.prompt_budget=2048and exercises sparse attention.libnvinfer.so.10is unavailable.QA Engineer Review
No test changes.
Description
NVBug 6533916 reports that the documented sparse-attention example fails immediately for external users because its model and input defaults point to internal NVIDIA and repository test paths.
This PR replaces the internal model default with the public
nvidia/Llama-3.1-8B-Instruct-FP8Hugging Face ID and makes--input_fileoptional. Without an input file, the example now builds a deterministic needle-in-a-haystack prompt from numbered expedition logs, hides an access code in one entry, and asks the model to retrieve it. The prompt exceeds the default prompt budget and directly exercises sparse attention. Users can still provide a JSONL file for custom long-context inputs. The change is limited to the example and does not modify public APIs, kernels, dependencies, or runtime behavior outside this script.Related issue: https://nvbugspro.nvidia.com/bug/6533916
Test Coverage
FileNotFoundErrorfrom the example's invalid default input path on the latestmainbranch used for this fix.pre-commit run --files examples/llm-api/llm_sparse_attention.pygit diff --checkprompt_budget=2048and less thanmax_seq_len=10240.Full GPU model inference was not run because the current environment does not provide
libnvinfer.so.10; the reported failure and this fix occur before model initialization.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.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.