[None] [feat] Optimize the algorithm part of RocketKV

[heyuhhh]

Summary by CodeRabbit

• New Features

  • Added CLI arguments: --topk (replaces --topr), --kt_cache_dtype, and --tokens_per_block for enhanced sparse attention configuration.
  • Support for configurable KT cache data type selection (bfloat16 or float8_e5m2).

• Updates

  • Updated RocketSparseAttentionConfig defaults: window_size (32), kernel_size (63), topk (64), topr (128), prompt_budget (2048), page_size (4).
  • KvCacheConfig now includes tokens_per_block parameter for improved KV cache control.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

Description

In this PR, mainly optimized the algorithm part of RocketKV based on the first triton implements (PR8682). The main optimizations are:

• Optimize several key triton kernels, including paged_kt_cache_bmm and update_kt_cache_ctx kernels
• Support fp8 kt cache for RocketKV in trtllm backend
• Others: the default config settings update and several minor bugs fix

After the optimization, using RocketKV can achieve higher throughput and lower latency than before:

The accuracy is also competitive with TRTLLM baseline on Longbench v1 task:

• TRTLLM: 48.7
• RocketKV: 48.15

Test Coverage

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)

• 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

/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...

Provide a user friendly way for developers to interact with a Jenkins server.

Run /bot [-h|--help] to print this help message.

See details below for each supported subcommand.

Details

run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]

Launch build/test pipelines. All previously running jobs will be killed.

--reuse-test (optional)pipeline-id (OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.

--disable-reuse-test (OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.

--disable-fail-fast (OPTIONAL) : Disable fail fast on build/tests/infra failures.

--skip-test (OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.

--stage-list "A10-PyTorch-1, xxx" (OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.

--gpu-type "A30, H100_PCIe" (OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.

--test-backend "pytorch, cpp" (OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.

--only-multi-gpu-test (OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.

--disable-multi-gpu-test (OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.

--add-multi-gpu-test (OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.

--post-merge (OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.

--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" (OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".

--detailed-log (OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.

--debug (OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in the stage-list parameter to access the appropriate container environment. Note: Does NOT update GitHub check status.

For guidance on mapping tests to stage names, see docs/source/reference/ci-overview.md
and the scripts/test_to_stage_mapping.py helper.

kill

kill

Kill all running builds associated with pull request.

skip

skip --comment COMMENT

Skip testing for latest commit on pull request. --comment "Reason for skipping build/test" is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

reuse-pipeline

reuse-pipeline

Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

[coderabbitai]

📝 Walkthrough

Walkthrough

The PR updates RocketKV sparse attention by replacing the top-r parameter with top-k, introducing KT cache data type selection, and refactoring kernels to use block-oriented indexing. Configuration defaults are adjusted across multiple files, and new parameters kt_cache_dtype and tokens_per_block are introduced to CLI arguments, configuration classes, and kernel signatures throughout the stack.

Changes

Cohort / File(s)	Summary
Example scripts – CLI argument updates examples/llm-api/llm_sparse_attention.py, examples/longbench/eval_longbench_v1.py, examples/longbench/eval_longbench_v2.py	Added CLI arguments --topk (default 64), --kt_cache_dtype (default 'float8_e5m2', choices ['bfloat16', 'float8_e5m2']), and --tokens_per_block (default 64). Renamed --topr to --topk in RocketKV configurations, updating defaults from 90 to 64.
Configuration classes tensorrt_llm/llmapi/llm_args.py	RocketSparseAttentionConfig: updated defaults (window_size: 32, kernel_size: 63, topr: 128, topk: 64, prompt_budget: 2048, page_size: 4), added new field kt_cache_dtype (str, default 'float8_e5m2'). KvCacheConfig: added tokens_per_block parameter.
Sparse kernel implementation tensorrt_llm/_torch/attention_backend/sparse/kernel.py	Replaced tensor-dimension positional computations with block-oriented indexing. Added total_rocket_k_ctx_tokens and prompt_budget parameters to multiple triton kernels. Generalized triton_softmax to support 2D and 3D input tensors. Updated grid calculations and DIM_BLOCK_SIZE logic for improved tiling. Removed dim_pos-based indexing in favor of explicit head/block-based computations.
RocketKV implementation tensorrt_llm/_torch/attention_backend/sparse/rocket.py	Added topr_filter() method to RocketTrtllmAttention and RocketVanillaAttention for top-r gating. Introduced total_rocket_k_ctx_tokens and max_rocket_k_ctx_len tracking during context preparation. Added kt_cache_dtype support in RocketKVCacheManager with dtype-aware buffer management. Updated preprocess_for_gen to return only (q, k) instead of (q, k, dim_pos). Enforced page_size as power of two via assertion.
Unit tests – RocketKV tests/unittest/_torch/attention/sparse/test_rocketkv.py	Updated RocketSparseAttentionConfig instantiations to include kt_cache_dtype parameter. Adjusted page_size expectations (3 → 4). Added dtype-aware KT buffer initialization and explicit dtype casting during data transfer between buffers (float8_e5m2 ↔ bfloat16).
Unit tests – Triton BMM tests/unittest/_torch/attention/sparse/test_triton_bmm.py	Removed dim_pos parameter from pytorch_reference_paged_kt_cache_bmm() signature. Updated output tensor shape from [total_num_heads, 1, total_kt_tokens] to [num_kv_heads, num_heads_per_kv, total_kt_tokens]. Reworked indexing and accumulation logic to match new tensor layout. Removed dim_pos-related test scaffolding.

Sequence Diagram(s)

sequenceDiagram
    participant CLI as CLI Arguments
    participant Config as RocketSparseAttentionConfig
    participant KVCache as KvCacheConfig
    participant KVMgr as RocketKVCacheManager
    participant Kernel as Sparse Kernels

    Note over CLI: --topk, --kt_cache_dtype, --tokens_per_block

    CLI->>Config: topk (int, default 64)
    CLI->>Config: kt_cache_dtype (str, default 'float8_e5m2')
    CLI->>KVCache: tokens_per_block (int, default 64)

    Config->>KVMgr: Pass topk, kt_cache_dtype
    KVCache->>KVMgr: Pass tokens_per_block

    KVMgr->>KVMgr: Set KT cache buffer dtype based on kt_cache_dtype
    KVMgr->>Kernel: total_rocket_k_ctx_tokens, prompt_budget
    
    Note over Kernel: Block-oriented indexing (replaces dim_pos)
    Kernel->>Kernel: Grid config: (batch_size, num_kv_heads, num_KT_token_tiles)

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

• tensorrt_llm/_torch/attention_backend/sparse/kernel.py: Dense refactoring with block-oriented indexing replacing dimension-based positional logic; multiple kernel signatures updated with new parameters; grid calculation changes require careful validation of correctness.
• tensorrt_llm/_torch/attention_backend/sparse/rocket.py: Introduction of topr_filter() method, kt_cache_dtype dtype-aware buffer management, and preprocess_for_gen() return signature change; integration of total_rocket_k_ctx_tokens and max_rocket_k_ctx_len tracking.
• tensorrt_llm/llmapi/llm_args.py: Multiple default value changes and new field addition; verify that defaults align with downstream expectations.
• Test updates: Verification that new tensor shapes, dtype conversions, and parameter propagation are correctly validated across test_rocketkv.py and test_triton_bmm.py.

Pre-merge checks and finishing touches

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 43.33% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.
Description check	⚠️ Warning	PR description provides clear explanation of RocketKV algorithm optimizations, performance improvements, and test coverage. However, the Test Coverage section is completely empty, leaving it unclear what specific tests safeguard the changes.	Add specific test names/paths in the Test Coverage section to clarify which tests validate the kernel optimizations, fp8 cache support, and configuration changes.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'Optimize the algorithm part of RocketKV' clearly and specifically summarizes the main optimization work described in the raw summary and PR description.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

Tip

📝 Customizable high-level summaries are now available in beta!

You can now customize how CodeRabbit generates the high-level summary in your pull requests — including its content, structure, tone, and formatting.

• Provide your own instructions using the high_level_summary_instructions setting.
• Format the summary however you like (bullet lists, tables, multi-section layouts, contributor stats, etc.).
• Use high_level_summary_in_walkthrough to move the summary from the description to the walkthrough section.

Example instruction:

"Divide the high-level summary into five sections:

1. 📝 Description — Summarize the main change in 50–60 words, explaining what was done.
2. 📓 References — List relevant issues, discussions, documentation, or related PRs.
3. 📦 Dependencies & Requirements — Mention any new/updated dependencies, environment variable changes, or configuration updates.
4. 📊 Contributor Summary — Include a Markdown table showing contributions:
   | Contributor | Lines Added | Lines Removed | Files Changed |
5. ✔️ Additional Notes — Add any extra reviewer context.
   Keep each section concise (under 200 words) and use bullet or numbered lists for clarity."

Note: This feature is currently in beta for Pro-tier users, and pricing will be announced later.

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

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

[coderabbitai]

Actionable comments posted: 4

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (2)

tensorrt_llm/llmapi/llm_args.py (1)

1440-1516: Add validation for KvCacheConfig.tokens_per_block to prevent invalid KV cache layouts

tokens_per_block is the logical capacity of a KV cache block and is used in block-count computations (and in downstream logic like paged KT cache / attention windows). If a caller sets it to 0 or a negative value, you will eventually hit division-by-zero or out‑of‑range indexing in those paths.

It’s worth enforcing a positive value at the config level, similar to the existing validators on free_gpu_memory_fraction and max_gpu_total_bytes. For example:

 class KvCacheConfig(StrictBaseModel, PybindMirror):
@@
-    tokens_per_block: int = Field(default=32,
-                                  description="The number of tokens per block.")
+    tokens_per_block: int = Field(
+        default=32,
+        description="The number of tokens per block.")
@@
     @field_validator('max_attention_window')
     @classmethod
     def validate_max_attention_window(cls, v: Optional[List[int]]):
@@
         return v
+
+    @field_validator('tokens_per_block')
+    @classmethod
+    def validate_tokens_per_block(cls, v: int) -> int:
+        if v <= 0:
+            raise ValueError(
+                "kv_cache_config.tokens_per_block must be a positive integer")
+        return v

This keeps misconfigurations from propagating into the KV cache manager where they’re harder to diagnose. Based on learnings

tensorrt_llm/_torch/attention_backend/sparse/kernel.py (1)

894-981: Critical: Mismatch in KT token calculation between wrapper and kernel.

The wrapper function calculates grid dimension z using prompt_budget but the kernel recalculates total_kt_tokens per-batch using num_sparse_tokens from sparse offsets. These can differ, causing the grid to be sized incorrectly relative to actual token counts:

• Wrapper (kernel.py:1018-1020): total_kt_tokens = (prompt_budget + kt_page_size - 1) // kt_page_size
• Kernel (kernel.py:923): total_kt_tokens = (num_sparse_tokens + kt_page_size - 1) // kt_page_size

When per-batch num_sparse_tokens < prompt_budget, the kernel mask at line 925 uses the smaller value while grid.z was sized for the larger value. This creates a correctness risk.

Additionally, no test coverage exists for this kernel. The search found tests for other sparse kernels (test_triton_bmm.py, test_triton_topk.py, test_flash_mla.py) but nothing for rocket_update_kt_cache_ctx.

Fix the calculation to use a consistent total_kt_tokens value, and add comprehensive test coverage for edge cases: partial pages, varying batch sizes, and different head dimensions.

🧹 Nitpick comments (7)

tensorrt_llm/llmapi/llm_args.py (1)

222-240: RocketSparseAttentionConfig RocketKV knobs look consistent; small API polish possible

The new defaults (window_size=32, kernel_size=63, topr=128, topk=64, prompt_budget=2048) and the added page_size / kt_cache_dtype fields are coherent with the example scripts and tests, and supports_backend correctly limits this to the PyTorch backend.

Two optional improvements:

• The page_size name is very close to KvCacheConfig.tokens_per_block. A brief clarification in the page_size description that it is the KT index page size (not KV cache block size) would help avoid confusion now that a separate tokens_per_block knob exists.
• For kt_cache_dtype, consider tightening the type to a Literal['bfloat16', 'float8_e5m2'] or an Enum instead of Optional[str] plus choices, so that Pydantic and static tooling enforce valid values at type level.

tests/unittest/_torch/attention/sparse/test_triton_bmm.py (1)

183-216: Reference paged KT BMM matches new layout; minor cleanups are optional

The reference now correctly produces scores with shape [num_kv_heads, num_heads_per_kv, total_kt_tokens] and uses output_offset = batch_idx * max_kt_tokens, which aligns with how output_offsets/mask are computed in the test, so the layout/indexing look sound.

A couple of non‑essential cleanups you could consider:

• k_vec and k_selected depend only on kv_head_idx, but are recomputed inside the q_head_idx loop; they could be moved one level up.
• Given you intentionally use the same tensor for min/max, k_selected = torch.where(dim_pos_vec, k_vec, k_vec) is a no‑op; keeping the dim_pos_vec comment for documentation but dropping the where call would reduce confusion.

These are purely cosmetic in a test‑only reference path.

examples/llm-api/llm_sparse_attention.py (1)

68-76: RocketKV CLI knobs are wired correctly; consider clarifying KT vs KV dtypes and tokens_per_block

The new arguments --topk, --kt_cache_dtype, and --tokens_per_block are correctly plumbed:

• topk / kt_cache_dtype go into RocketSparseAttentionConfig, matching its updated signature.
• tokens_per_block and kv_cache_dtype go into KvCacheConfig, matching the Pydantic model.

Two small clarity improvements you might consider:

• In the help text for --kt_cache_dtype, call out explicitly that it controls the KT cache precision and is independent of --kv_cache_dtype, which still sets the KV cache dtype.
• For --tokens_per_block, a short note that this affects the PyTorch backend’s KV cache manager page/block size (and is ignored by the TensorRT backend) would help users understand when the knob has an effect.

Behavior-wise this looks good as is.

Also applies to: 117-117, 143-149, 193-200

examples/longbench/eval_longbench_v1.py (1)

153-162: LongBench v1 RocketKV knobs are correctly integrated; only minor doc consistency nits

Here, --topk, --kt_cache_dtype, and --tokens_per_block are all correctly threaded:

• topk / kt_cache_dtype are passed into RocketSparseAttentionConfig when --rocket_sparse is enabled.
• tokens_per_block is passed into KvCacheConfig, aligning with the new KV cache block‑size knob.

If you want to polish further (optional):

• Mirror the --kt_cache_dtype help string from llm_sparse_attention.py to keep CLI UX consistent across examples.
• Consider mentioning in the --tokens_per_block help that this primarily affects the PyTorch backend’s KV cache manager so users don’t expect it to change prebuilt TensorRT engines.

Functionally this looks solid.

Also applies to: 168-173, 325-330, 341-347

tensorrt_llm/_torch/attention_backend/sparse/rocket.py (2)

462-465: Verify contiguity check performance impact.

The contiguity checks with conditional .contiguous() calls are defensive but create new tensors if the input is non-contiguous. If this method is called frequently in the hot path, consider whether upstream callers should be required to provide contiguous tensors to avoid allocation overhead.

Based on learnings: TensorRT-LLM's attention backend typically trusts callers to provide correctly formatted tensors without validation. Consider documenting the contiguity requirement instead of runtime checks, or profile to confirm the overhead is negligible.

---

625-625: Restrictive assertion may limit future extensibility.

The hard assertion assert sparse_attention_config.kt_cache_dtype == 'bfloat16' prevents Vanilla RocketKV from supporting FP8 KT cache. If FP8 support for Vanilla is planned or could be useful, consider replacing this with a capability check or warning.

-        assert sparse_attention_config.kt_cache_dtype == 'bfloat16', "Only bfloat16 kt cache is supported for Vanilla RocketKV"
+        if sparse_attention_config.kt_cache_dtype != 'bfloat16':
+            raise NotImplementedError(f"Vanilla RocketKV currently only supports bfloat16 kt_cache, got {sparse_attention_config.kt_cache_dtype}")

This makes the temporary limitation clearer and more maintainable.

tensorrt_llm/_torch/attention_backend/sparse/kernel.py (1)

1195-1195: Static analysis: Replace lambda with def.

Ruff flags the lambda expression as a style issue. While this works, following the style guideline improves readability.

Apply this diff:

-    grid = lambda meta: (num_gen_tokens, num_kv_heads)
+    def grid(meta):
+        return (num_gen_tokens, num_kv_heads)

However, since meta is unused, you could also use:

-    grid = lambda meta: (num_gen_tokens, num_kv_heads)
+    grid = (num_gen_tokens, num_kv_heads)

if the kernel invocation accepts a tuple directly.

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between a39e8c5 and 63934ff.

📒 Files selected for processing (8)

• examples/llm-api/llm_sparse_attention.py (4 hunks)
• examples/longbench/eval_longbench_v1.py (3 hunks)
• examples/longbench/eval_longbench_v2.py (2 hunks)
• tensorrt_llm/_torch/attention_backend/sparse/kernel.py (13 hunks)
• tensorrt_llm/_torch/attention_backend/sparse/rocket.py (10 hunks)
• tensorrt_llm/llmapi/llm_args.py (1 hunks)
• tests/unittest/_torch/attention/sparse/test_rocketkv.py (5 hunks)
• tests/unittest/_torch/attention/sparse/test_triton_bmm.py (2 hunks)

🧰 Additional context used

🧠 Learnings (10)

📓 Common learnings

Learnt from: thorjohnsen
Repo: NVIDIA/TensorRT-LLM PR: 6910
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:0-0
Timestamp: 2025-08-14T21:04:50.248Z
Learning: In KV cache onboarding logic during prefill in cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp, when calculating which blocks fall within the attention window, use getTokensPerBlock() to advance token indices rather than block->getUniqueTokens().size(), because the calculation needs to consider the post-prefill state where blocks will be filled to capacity, not their current token count.

Learnt from: eopXD
Repo: NVIDIA/TensorRT-LLM PR: 6767
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:0-0
Timestamp: 2025-08-15T06:46:54.897Z
Learning: In cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp addToken function, newly allocated blocks are unshared by design. The beam search path in addToken (when sequence.getNumTokens() > windowSize) is currently broken/non-functional with SWA, so the block allocation doesn't follow a shared-then-unshared pattern.

Learnt from: eopXD
Repo: NVIDIA/TensorRT-LLM PR: 6767
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:0-0
Timestamp: 2025-08-15T06:46:53.813Z
Learning: In the TensorRT-LLM KV cache manager, SWA (Sliding Window Attention) combined with beam search is currently in a broken/non-functional state and is planned for future rework. During preparatory refactoring phases, code related to SWA+beam search may intentionally remain in a non-working state until the broader rework is completed.

📚 Learning: 2025-08-14T15:38:01.771Z

Learnt from: MatthiasKohl
Repo: NVIDIA/TensorRT-LLM PR: 6904
File: cpp/tensorrt_llm/pybind/thop/bindings.cpp:55-57
Timestamp: 2025-08-14T15:38:01.771Z
Learning: In TensorRT-LLM Python bindings, tensor parameter collections like mla_tensor_params and spec_decoding_tensor_params are kept as required parameters without defaults to maintain API consistency, even when it might affect backward compatibility.

Applied to files:

• tensorrt_llm/llmapi/llm_args.py

📚 Learning: 2025-08-14T21:04:50.248Z

Learnt from: thorjohnsen
Repo: NVIDIA/TensorRT-LLM PR: 6910
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:0-0
Timestamp: 2025-08-14T21:04:50.248Z
Learning: In KV cache onboarding logic during prefill in cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp, when calculating which blocks fall within the attention window, use getTokensPerBlock() to advance token indices rather than block->getUniqueTokens().size(), because the calculation needs to consider the post-prefill state where blocks will be filled to capacity, not their current token count.

Applied to files:

• tests/unittest/_torch/attention/sparse/test_rocketkv.py
• examples/longbench/eval_longbench_v1.py
• examples/llm-api/llm_sparse_attention.py
• tests/unittest/_torch/attention/sparse/test_triton_bmm.py
• tensorrt_llm/_torch/attention_backend/sparse/rocket.py
• examples/longbench/eval_longbench_v2.py
• tensorrt_llm/_torch/attention_backend/sparse/kernel.py

📚 Learning: 2025-08-26T09:37:10.463Z

Learnt from: jiaganc
Repo: NVIDIA/TensorRT-LLM PR: 7031
File: tensorrt_llm/bench/dataclasses/configuration.py:90-104
Timestamp: 2025-08-26T09:37:10.463Z
Learning: In TensorRT-LLM's bench configuration, the `get_pytorch_perf_config()` method returns `self.pytorch_config` which is a Dict[str, Any] that can contain default values including `cuda_graph_config`, making the fallback `llm_args["cuda_graph_config"]` safe to use.

Applied to files:

• tests/unittest/_torch/attention/sparse/test_rocketkv.py

📚 Learning: 2025-08-14T15:43:23.107Z

Learnt from: MatthiasKohl
Repo: NVIDIA/TensorRT-LLM PR: 6904
File: tensorrt_llm/_torch/attention_backend/trtllm.py:259-262
Timestamp: 2025-08-14T15:43:23.107Z
Learning: In TensorRT-LLM's attention backend, tensor parameters in the plan() method are assigned directly without validation (dtype, device, contiguity checks). This maintains consistency across all tensor inputs and follows the pattern of trusting callers to provide correctly formatted tensors.

Applied to files:

• tests/unittest/_torch/attention/sparse/test_rocketkv.py
• tensorrt_llm/_torch/attention_backend/sparse/kernel.py

📚 Learning: 2025-08-15T06:46:54.897Z

Learnt from: eopXD
Repo: NVIDIA/TensorRT-LLM PR: 6767
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:0-0
Timestamp: 2025-08-15T06:46:54.897Z
Learning: In cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp addToken function, newly allocated blocks are unshared by design. The beam search path in addToken (when sequence.getNumTokens() > windowSize) is currently broken/non-functional with SWA, so the block allocation doesn't follow a shared-then-unshared pattern.

Applied to files:

• examples/llm-api/llm_sparse_attention.py
• tensorrt_llm/_torch/attention_backend/sparse/kernel.py

📚 Learning: 2025-09-23T14:58:05.372Z

Learnt from: nv-lschneider
Repo: NVIDIA/TensorRT-LLM PR: 7910
File: cpp/tensorrt_llm/kernels/nccl_device/config.cu:42-49
Timestamp: 2025-09-23T14:58:05.372Z
Learning: In TensorRT-LLM NCCL device kernels (cpp/tensorrt_llm/kernels/nccl_device/), the token partitioning intentionally uses ceil-like distribution (same token_per_rank for all ranks) to ensure all ranks launch the same number of blocks. This is required for optimal NCCL device API barrier performance, even though it may launch extra blocks for non-existent tokens on later ranks. Runtime bounds checking in the kernel (blockID validation) handles the overshoot cases.

Applied to files:

• tensorrt_llm/_torch/attention_backend/sparse/kernel.py

📚 Learning: 2025-08-15T06:46:53.813Z

Learnt from: eopXD
Repo: NVIDIA/TensorRT-LLM PR: 6767
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:0-0
Timestamp: 2025-08-15T06:46:53.813Z
Learning: In the TensorRT-LLM KV cache manager, SWA (Sliding Window Attention) combined with beam search is currently in a broken/non-functional state and is planned for future rework. During preparatory refactoring phases, code related to SWA+beam search may intentionally remain in a non-working state until the broader rework is completed.

Applied to files:

• tensorrt_llm/_torch/attention_backend/sparse/kernel.py

📚 Learning: 2025-08-22T01:54:35.850Z

Learnt from: djns99
Repo: NVIDIA/TensorRT-LLM PR: 7104
File: cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h:999-1000
Timestamp: 2025-08-22T01:54:35.850Z
Learning: The `internal_cutlass_kernels` directory in TensorRT-LLM is a mirror of an internal NVIDIA repository and maintains its own implementation and API that may diverge from the public `cutlass_kernels` version. API inconsistencies between these two directories are intentional and by design, not bugs to be fixed.

Applied to files:

• tensorrt_llm/_torch/attention_backend/sparse/kernel.py

📚 Learning: 2025-08-21T09:41:49.347Z

Learnt from: eopXD
Repo: NVIDIA/TensorRT-LLM PR: 6768
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:2010-2045
Timestamp: 2025-08-21T09:41:49.347Z
Learning: In cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp, updateSequenceCacheBlockOffsets is specifically for updating bookkeeping when blocks are added during the context phase, not for refreshing offsets after detach operations. During detach operations, GenerationRequest::removeFrontBlock handles the necessary cache block bookkeeping internally.

Applied to files:

• tensorrt_llm/_torch/attention_backend/sparse/kernel.py

🧬 Code graph analysis (6)

tensorrt_llm/llmapi/llm_args.py (1)

tensorrt_llm/_torch/attention_backend/flashinfer.py (1)

• page_size (200-204)

tests/unittest/_torch/attention/sparse/test_rocketkv.py (1)

tensorrt_llm/llmapi/llm_args.py (1)

• RocketSparseAttentionConfig (222-250)

examples/longbench/eval_longbench_v1.py (2)

tensorrt_llm/_torch/attention_backend/trtllm.py (1)

• tokens_per_block (654-658)

tensorrt_llm/functional.py (1)

• topk (7318-7414)

examples/llm-api/llm_sparse_attention.py (2)

tensorrt_llm/_torch/attention_backend/trtllm.py (1)

• tokens_per_block (654-658)

tensorrt_llm/functional.py (1)

• topk (7318-7414)

tensorrt_llm/_torch/attention_backend/sparse/rocket.py (3)

tensorrt_llm/_torch/attention_backend/flashinfer.py (1)

• page_size (200-204)

tensorrt_llm/_torch/attention_backend/trtllm.py (1)

• tokens_per_block (654-658)

tensorrt_llm/_torch/attention_backend/sparse/dsa.py (1)

• get_cache_bytes_per_token (1569-1589)

tensorrt_llm/_torch/attention_backend/sparse/kernel.py (2)

tensorrt_llm/_torch/attention_backend/trtllm.py (1)

• tokens_per_block (654-658)

cpp/tensorrt_llm/kernels/sparseAttentionKernels.h (1)

• sparse_kv_offsets (33-85)

🪛 Ruff (0.14.5)

tensorrt_llm/_torch/attention_backend/sparse/kernel.py

1195-1195: Do not assign a lambda expression, use a def

Rewrite grid as a def

(E731)

---

1195-1195: Unused lambda argument: meta

(ARG005)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)

• GitHub Check: Pre-commit Check

🔇 Additional comments (7)

tests/unittest/_torch/attention/sparse/test_rocketkv.py (2)

374-386: Verify intentional differences in test configurations.

The test uses different configurations for vanilla vs trtllm:

• page_size=2 for both
• topk=128 for vanilla, topk=64 for trtllm
• topr=96 for both

These differences may be intentional for testing different code paths, but please confirm this is the intended test design rather than an oversight.

---

468-473: FP8 initialization loses statistical properties without explicit scaling.

The naive conversion from float16 to float8_e5m2 does not preserve the normal distribution's properties. Quantization noise is large due to FP8 having only 2 mantissa bits and can change the sample mean and especially variance/tails. Per-tensor or per-channel scaling should be used when converting to FP8 to reduce bias and clipping.

The current implementation (lines 468-473) should either:

1. Apply explicit per-tensor scaling before the conversion, or
2. Use an alternative initialization approach that directly supports float8_e5m2, or
3. Document why unscaled conversion is acceptable for this specific test

⛔ Skipped due to learnings

Learnt from: jhaotingc
Repo: NVIDIA/TensorRT-LLM PR: 7856
File: cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp:159-166
Timestamp: 2025-09-19T21:28:13.751Z
Learning: In TensorRT-LLM blockScaleMoe routing (cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu), the DeepSeek routing method performs reinterpret_cast<float*>(routingLogits) at line 89, which could cause issues if routing_logits are BF16. However, Qwen3-FP8 models use RenormalizeNaive routing method and are not affected by this dtype casting issue.

Learnt from: nvchenghaoz
Repo: NVIDIA/TensorRT-LLM PR: 8469
File: tensorrt_llm/_torch/auto_deploy/transform/library/rms_norm.py:180-182
Timestamp: 2025-10-20T17:09:21.560Z
Learning: In tensorrt_llm/_torch/auto_deploy/transform/library/rms_norm.py, the _gated_rmsnorm_replacement function does not need to cast the output of torch.ops.auto_deploy.torch_rmsnorm_gated back to the input dtype, even though the custom op returns fp32. The dtype handling is managed elsewhere or the fp32 output is acceptable for downstream consumers.

tensorrt_llm/_torch/attention_backend/sparse/rocket.py (2)

48-49: Good validation: page_size must be power of 2.

This assertion enforces a critical constraint that the block-oriented indexing relies on. The error will catch misconfigurations early.

---

258-267: Properly handles empty valid batch case.

Setting max_rocket_k_ctx_len=0 and total_rocket_k_ctx_tokens=0 when valid_batch_size=0 prevents undefined behavior from calling .max().item() on an empty tensor slice.

tensorrt_llm/_torch/attention_backend/sparse/kernel.py (3)

508-520: Softmax kernel generalized for 2D and 3D inputs.

The updated logic correctly handles both input shapes:

• 2D: [num_heads, total_k_tokens] with len_per_seq=1
• 3D: [num_heads, len_per_seq, total_k_tokens]

The grid calculation (num_heads, batch_size * len_per_seq) properly accounts for both cases.

---

1062-1152: Review comment is unsupported by evidence.

The dim_pos computation via tl.sum(q_values, axis=0) > 0 is mathematically equivalent to and directly validated by the test reference (test_triton_bmm.py lines 200-201), which implements the identical logic. The test confirms this is the expected behavior, not a deviation.

The claim of a "critical logic change from the previous dim_pos-based approach" lacks supporting evidence. The current implementation IS the dim_pos-based approach and aligns perfectly with test expectations. No alternative previous implementation was found in the codebase.

The shape analysis (Q: [Q_BLOCK_SIZE, DIM_BLOCK_SIZE], KT: [DIM_BLOCK_SIZE, KT_BLOCK_SIZE], output: [Q_BLOCK_SIZE, KT_BLOCK_SIZE]) and output indexing logic are correct as stated.

Likely an incorrect or invalid review comment.

---

1190-1197: Output shape mismatch breaks downstream function compatibility.

The function returns [num_kv_heads, num_heads_per_kv, total_kt_tokens], but triton_rocket_reduce_scores (rocket.py:528) expects [num_kv_heads * num_heads_per_kv, 1, total_kt_tokens]. The dimensions are incompatible and no reshape operation exists between the calls (rocket.py:512-528). A reshape is needed before passing scores to triton_rocket_reduce_scores, or the function should flatten the first two dimensions before returning.

⛔ Skipped due to learnings

Learnt from: sklevtsov-nvidia
Repo: NVIDIA/TensorRT-LLM PR: 3294
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_tma_warp_specialized_input.cu:118-127
Timestamp: 2025-08-09T20:57:04.084Z
Learning: In the CUTLASS MoE finalize fusion implementation (cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_tma_warp_specialized_input.cu), when setting `fused_finalize_epilogue.stride_final_output` with shape `(hidden_size, num_output_tokens, 1)`, the `num_rows_in_final_output` should be set to `num_output_tokens` (not `hidden_size`) because of a swap+transpose operation that maps rows of the output tensor to `hidden_size` and columns to `num_output_tokens`.

Learnt from: ixlmar
Repo: NVIDIA/TensorRT-LLM PR: 7294
File: tensorrt_llm/_torch/pyexecutor/sampler.py:887-891
Timestamp: 2025-08-28T10:25:22.370Z
Learning: In tensorrt_llm/_torch/pyexecutor/sampler.py, the draft_probs and target_probs tensors have shapes [1, steps] not [steps, vocab_size] as might be expected, making the .squeeze(0) operations appropriate for removing the batch dimension of size 1.

[heyuhhh]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #25173 [ run ] triggered by Bot. Commit: 63934ff

[tensorrt-cicd]

PR_Github #25173 [ run ] completed with state SUCCESS. Commit: 63934ff
/LLM/main/L0_MergeRequest_PR pipeline #19033 completed with status: 'FAILURE'

[heyuhhh]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #25277 [ run ] triggered by Bot. Commit: c1e0355

[tensorrt-cicd]

PR_Github #25277 [ run ] completed with state SUCCESS. Commit: c1e0355
/LLM/main/L0_MergeRequest_PR pipeline #19122 completed with status: 'FAILURE'

[heyuhhh]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #25323 [ run ] triggered by Bot. Commit: 6520577

[tensorrt-cicd]

PR_Github #25323 [ run ] completed with state SUCCESS. Commit: 6520577
/LLM/main/L0_MergeRequest_PR pipeline #19156 completed with status: 'FAILURE'

[heyuhhh]

/bot run

[tensorrt-cicd]

PR_Github #25365 [ run ] triggered by Bot. Commit: 6520577

[tensorrt-cicd]

PR_Github #25365 [ run ] completed with state SUCCESS. Commit: 6520577
/LLM/main/L0_MergeRequest_PR pipeline #19186 completed with status: 'FAILURE'

[heyuhhh]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #25559 [ run ] triggered by Bot. Commit: 7f77a2a

[tensorrt-cicd]

PR_Github #25559 [ run ] completed with state SUCCESS. Commit: 7f77a2a
/LLM/main/L0_MergeRequest_PR pipeline #19357 completed with status: 'FAILURE'

[heyuhhh]

/bot run

[tensorrt-cicd]

PR_Github #25637 [ run ] triggered by Bot. Commit: 7f77a2a

[tensorrt-cicd]

PR_Github #25637 [ run ] completed with state FAILURE. Commit: 7f77a2a
LLM/main/L0_MergeRequest_PR #19425 (Blue Ocean) completed with status: ABORTED

[lfr-0531]

/bot run

[tensorrt-cicd]

PR_Github #26088 [ run ] triggered by Bot. Commit: 4f9ff81

[tensorrt-cicd]

PR_Github #26088 [ run ] completed with state SUCCESS. Commit: 4f9ff81
/LLM/main/L0_MergeRequest_PR pipeline #19808 completed with status: 'FAILURE'

[heyuhhh]

/bot run

[tensorrt-cicd]

PR_Github #26145 [ run ] triggered by Bot. Commit: 4f9ff81

[tensorrt-cicd]

PR_Github #26145 [ run ] completed with state SUCCESS. Commit: 4f9ff81
/LLM/main/L0_MergeRequest_PR pipeline #19857 completed with status: 'FAILURE'

[lfr-0531]

/bot run

[tensorrt-cicd]

PR_Github #26186 [ run ] triggered by Bot. Commit: 4f9ff81

[tensorrt-cicd]

PR_Github #26186 [ run ] completed with state SUCCESS. Commit: 4f9ff81
/LLM/main/L0_MergeRequest_PR pipeline #19891 completed with status: 'FAILURE'

[heyuhhh]

/bot run

[tensorrt-cicd]

PR_Github #26215 [ run ] triggered by Bot. Commit: 4d9d6c8

[tensorrt-cicd]

PR_Github #26215 [ run ] completed with state SUCCESS. Commit: 4d9d6c8
/LLM/main/L0_MergeRequest_PR pipeline #19917 completed with status: 'FAILURE'

[heyuhhh]

/bot run

[tensorrt-cicd]

PR_Github #26234 [ run ] triggered by Bot. Commit: 65187c0

[tensorrt-cicd]

PR_Github #26234 [ run ] completed with state SUCCESS. Commit: 65187c0
/LLM/main/L0_MergeRequest_PR pipeline #19933 completed with status: 'FAILURE'

[heyuhhh]

/bot run

[tensorrt-cicd]

PR_Github #26245 [ run ] triggered by Bot. Commit: 65187c0

[tensorrt-cicd]

PR_Github #26245 [ run ] completed with state SUCCESS. Commit: 65187c0
/LLM/main/L0_MergeRequest_PR pipeline #19942 completed with status: 'FAILURE'

[heyuhhh]

/bot run

[tensorrt-cicd]

PR_Github #26263 [ run ] triggered by Bot. Commit: 65187c0

[tensorrt-cicd]

PR_Github #26263 [ run ] completed with state FAILURE. Commit: 65187c0
/LLM/main/L0_MergeRequest_PR pipeline #19958 completed with status: 'FAILURE'

[lfr-0531]

/bot run

[tensorrt-cicd]

PR_Github #26320 [ run ] triggered by Bot. Commit: 65187c0

[tensorrt-cicd]

PR_Github #26320 [ run ] completed with state SUCCESS. Commit: 65187c0
/LLM/main/L0_MergeRequest_PR pipeline #19986 completed with status: 'SUCCESS'