Skip to content

[https://nvbugs/6533916][fix] Make sparse attention example runnable - #17055

Merged
lfr-0531 merged 2 commits into
NVIDIA:mainfrom
lfr-0531:user/fanrongl/fix-sparse-attention-example-6533916
Aug 3, 2026
Merged

[https://nvbugs/6533916][fix] Make sparse attention example runnable#17055
lfr-0531 merged 2 commits into
NVIDIA:mainfrom
lfr-0531:user/fanrongl/fix-sparse-attention-example-6533916

Conversation

@lfr-0531

@lfr-0531 lfr-0531 commented Jul 30, 2026

Copy link
Copy Markdown
Collaborator

Dev Engineer Review

  • Updated the sparse attention example to use the public nvidia/Llama-3.1-8B-Instruct-FP8 model by default.
  • Made --input_file optional and added a deterministic needle-in-a-haystack prompt for zero-input execution.
  • Retained JSONL input support for custom long-context workloads.
  • Ensured the built-in prompt exceeds the default prompt_budget=2048 and exercises sparse attention.
  • Added an Apache-2.0 SPDX header and expanded usage documentation.
  • Scoped changes to one Python example. No public APIs, kernels, runtimes, dependencies, or test files changed.
  • Validation included formatting checks, source compilation, CLI-default checks, prompt-content and token-length checks, and mocked validation of both input paths.
  • Full GPU inference remains unverified because libnvinfer.so.10 is 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-FP8 Hugging Face ID and makes --input_file optional. 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

  • Reproduced the original FileNotFoundError from the example's invalid default input path on the latest main branch used for this fix.
  • pre-commit run --files examples/llm-api/llm_sparse_attention.py
  • git diff --check
  • Compiled the modified source with Python and confirmed the loaded module path points to the current working-tree script.
  • Validated that the built-in prompt contains 128 log entries and exactly one access-code occurrence.
  • Validated with the default model tokenizer that the built-in prompt is 2,842 tokens: greater than prompt_budget=2048 and less than max_seq_len=10240.
  • Ran mocked end-to-end validation of both the built-in prompt path and the optional JSONL input path.

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-compatible or api-breaking. For api-breaking, include BREAKING in 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.

@lfr-0531
lfr-0531 requested a review from a team as a code owner July 30, 2026 07:41
@lfr-0531
lfr-0531 requested review from chang-l and nv-guomingz July 30, 2026 07:41
@coderabbitai

coderabbitai Bot commented Jul 30, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

Walkthrough

The 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.

Changes

Sparse attention example

Layer / File(s) Summary
Example documentation and CLI defaults
examples/llm-api/llm_sparse_attention.py
Usage documentation, license content, built-in prompt definitions, and --model_path/--input_file defaults are updated.
Conditional prompt execution and output handling
examples/llm-api/llm_sparse_attention.py
run_llm selects built-in or JSONL prompts and prints references only when available.

Estimated code review effort: 2 (Simple) | ~10 minutes

Suggested reviewers: bowenfu, chang-l

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
Loading
🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Title check ✅ Passed The title follows the required NVBugs and type format and clearly summarizes the fix that makes the sparse attention example runnable.
Description check ✅ Passed The description explains the issue, solution, test coverage, limitations, and checklist status in sufficient detail.
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

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

@lfr-0531
lfr-0531 force-pushed the user/fanrongl/fix-sparse-attention-example-6533916 branch from 4edb499 to 59f6dbf Compare July 30, 2026 08:12

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

📥 Commits

Reviewing files that changed from the base of the PR and between 4edb499 and 59f6dbf.

📒 Files selected for processing (1)
  • examples/llm-api/llm_sparse_attention.py

Comment thread examples/llm-api/llm_sparse_attention.py Outdated
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
@lfr-0531
lfr-0531 force-pushed the user/fanrongl/fix-sparse-attention-example-6533916 branch from 59f6dbf to 2ef156d Compare July 30, 2026 08:20
@lfr-0531
lfr-0531 marked this pull request as draft July 30, 2026 16:07
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
@lfr-0531
lfr-0531 requested a review from heyuhhh July 31, 2026 04:19
@lfr-0531
lfr-0531 marked this pull request as ready for review July 31, 2026 04:19
@lfr-0531

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

📥 Commits

Reviewing files that changed from the base of the PR and between 2ef156d and 9deddbb.

📒 Files selected for processing (1)
  • examples/llm-api/llm_sparse_attention.py

Comment thread examples/llm-api/llm_sparse_attention.py
@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #62952 [ run ] triggered by Bot. Commit: 9deddbb Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #62952 [ run ] completed with state SUCCESS. Commit: 9deddbb
/LLM/main/L0_MergeRequest_PR pipeline #51068 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@lfr-0531

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #63014 [ run ] triggered by Bot. Commit: 9deddbb Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #63014 [ run ] completed with state SUCCESS. Commit: 9deddbb
/LLM/main/L0_MergeRequest_PR pipeline #51121 completed with status: 'SUCCESS'

CI Report

Link to invocation

@BowenFu BowenFu left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@lfr-0531
lfr-0531 enabled auto-merge (squash) August 3, 2026 08:26
@lfr-0531
lfr-0531 merged commit 7e2fca0 into NVIDIA:main Aug 3, 2026
15 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants