Draft: Support context sparse attention for mqa/gqa - #10788
Conversation
Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>
Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>
Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>
📝 WalkthroughWalkthroughThis pull request refactors sparse attention handling in TensorRT-LLM by transitioning from sparse MLA terminology to generic sparse attention naming conventions. New configuration fields ( Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~55 minutes 🚥 Pre-merge checks | ✅ 1 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (1 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
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}
There was a problem hiding this comment.
Actionable comments posted: 13
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (8)
tensorrt_llm/_torch/attention_backend/trtllm.py (1)
1617-1634: Unbound local when sparse attention config is absent or SkipSoftmax.Line 1617 initializes
sparse_*variables but omitssparse_attn_ctx_indices, yet Line 1694 always passes it intoplan(). Whenself.sparse_attention_configisNoneorSkipSoftmaxAttentionConfig, this will raiseUnboundLocalError.🐛 Initialize sparse_attn_ctx_indices
- sparse_kv_indices, sparse_kv_offsets, sparse_attn_indices, sparse_attn_offsets = None, None, None, None + sparse_kv_indices, sparse_kv_offsets, sparse_attn_indices, sparse_attn_offsets, sparse_attn_ctx_indices = ( + None, None, None, None, None + )cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h (1)
519-536: Resolve the TODO about sparse-attention tile sizing.The unresolved TODO at line 592 questions whether
maxNumHeadsQPerKvInCta = 64is correct for sparse attention in theKeepsMmaAbForGenerationKernelpath. Since sparse attention reaching this kernel path is constrained tonumHeadsQPerKv == 128(validated at line 534), the 64-head tile size should be confirmed against kernel metadata or benchmarks. Either document the reasoning for this value or update the kernel parameters if incorrect.cpp/tensorrt_llm/thop/attentionOp.cpp (1)
1-3: Update the NVIDIA copyright year to 2026.This file has 2026 changes, so the header should reflect the latest meaningful modification year.
As per coding guidelines, please update the header year.📝 Suggested update
- * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & + * SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION &cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/kernelParams.h (1)
1-3: Update the NVIDIA copyright year to 2026.The header should reflect the latest meaningful modification year for this file.
As per coding guidelines, please update the header year.📝 Suggested update
- * Copyright (c) 2020-2024, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2020-2026, NVIDIA CORPORATION. All rights reserved.cpp/tensorrt_llm/kernels/xqaDispatcher.cpp (1)
1-15: Update copyright year to 2026.The file was modified in this PR, but the header still ends at 2024. Please bump it to include 2026. As per coding guidelines, update the year of latest meaningful modification.
✅ Proposed update
- * Copyright (c) 2020-2024, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2020-2026, NVIDIA CORPORATION. All rights reserved.cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaReduction.cu (1)
1-15: Update copyright year to 2026.The header still ends at 2023 though the file is modified in this PR. Please update the year to include 2026. As per coding guidelines, update the year of latest meaningful modification.
✅ Proposed update
- * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2020-2026, NVIDIA CORPORATION. All rights reserved.cpp/tensorrt_llm/thop/attentionOp.h (1)
1-15: Update copyright year to 2026.This header was modified, but the copyright year is still 2025. Please update to 2026. As per coding guidelines, update the year of latest meaningful modification.
✅ Proposed update
- * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved.cpp/tensorrt_llm/common/attentionOp.cpp (1)
1-16: Update copyright year to 2026.This file was modified in this PR, but the header ends at 2025. Please update it to include 2026. As per coding guidelines, update the year of latest meaningful modification.
✅ Proposed update
- * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & + * SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION &
🤖 Fix all issues with AI agents
In `@cpp/tensorrt_llm/common/attentionOp.cpp`:
- Around line 298-299: The code unconditionally sets
AttentionInputLayout::Q_PAGED_KV when useTllmGenSparseAttention() is true, but
useTllmGenSparseAttention() does not check mPagedKVCache and can cause invalid
reads when KV is contiguous; change the guard to require the paged invariant
(e.g., use useTllmGenSparseAttentionPaged() or explicitly check mPagedKVCache)
before assigning Q_PAGED_KV so the paged layout is only used when mPagedKVCache
is true; update the branch that currently calls useTllmGenSparseAttention() (the
one setting AttentionInputLayout::Q_PAGED_KV) to
useTllmGenSparseAttentionPaged() or add && mPagedKVCache to the condition.
In `@cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaRunnerParams.h`:
- Around line 288-291: Update the comments for the fields mSparseAttention and
mSparseTopK to consistently use "sparse attention" instead of "sparse MLA":
replace the comment "Whether to use sparse MLA." above mSparseAttention with
"Whether to use sparse attention." and replace "The top k value for sparse MLA."
above mSparseTopK with "The top-k value for sparse attention." to match the
renamed fields.
In `@cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/kernelParams.h`:
- Around line 151-154: Update the comment for the int32_t member mNumSparseTopk
in kernelParams.h to reflect that this value applies to general sparse attention
(not just MLA); locate the declaration of mNumSparseTopk and replace "The top k
value for sparse MLA." with a concise comment like "The top-k value used for
sparse attention." so the name and intent match the new general usage.
- Around line 718-725: Replace the TODO by making the tileSizePerCtaQ
calculation explicit and tested: keep the default tileSizePerCtaQ =
kernelMeta.mTileSizeQ * numInstsQ, and when kernelMeta.mGroupsHeadsQ is true and
!isSpecDecodingGenerationKernel(options.mKernelType) set tileSizePerCtaQ =
kernelMeta.mTileSizeQ (one instance per CTA) — update params.mNumTokensPerCtaQ
accordingly, add a short comment referencing numInstsQ, tileSizePerCtaQ,
kernelMeta.mGroupsHeadsQ, isSpecDecodingGenerationKernel and
options.mKernelType, and add a tiny unit/integration test or open a follow-up
issue to cover this branch to validate correctness and performance.
- Around line 859-869: The code doesn't initialize params.mSparsityMode when
neither options.mSparseAttention nor options.mUseBlockSparseAttention is true
and uses the magic literal 4 for top-k alignment; update the block around
TLLM_CHECK_WITH_INFO and the subsequent assignments so that the alignment
literal is replaced with a named constant (e.g., kSparseTopKAlignment) and set
params.mSparsityMode to a clear default (e.g., 0 for dense) when neither
options.mSparseAttention nor options.mUseBlockSparseAttention is enabled;
reference the existing symbols options.mSparseAttention,
options.mUseBlockSparseAttention, params.mSparsityMode, params.mNumSparseTopk
and introduce the constant name for the alignment check.
In `@cpp/tensorrt_llm/pybind/thop/bindings.cpp`:
- Around line 67-68: The binding signature for attention() changed: a new
positional parameter sparse_attn_ctx_indices was inserted and
sparse_attn_indices_block_size no longer has a default (making it required);
find every positional call site (e.g., internal callers and external users) and
update them to match the new parameter order or convert to keyword arguments to
avoid misalignment—check calls in
tensorrt_llm/_torch/attention_backend/trtllm.py (lines noted in the review) as a
model, update other modules/scripts/tests that invoke attention() positionally,
and run relevant unit/integration tests to ensure no missing required argument
errors at runtime.
In `@cpp/tensorrt_llm/thop/attentionOp.cpp`:
- Around line 407-421: When op.useTllmGenSparseAttention() is true, add an
explicit guard that validates sparse_attn_ctx_indices.has_value() before
dereferencing it; if not present, return/error with a clear message (or throw)
so run() fails fast instead of UB. Update the block around
useTllmGenSparseAttention()/isMLAEnabled() to check sparse_attn_ctx_indices
first and only call sparse_attn_ctx_indices.value().size(-1) after the guard,
and ensure mRuntimeSparseAttentionParams.sparse_topk is set accordingly.
In `@tensorrt_llm/_torch/attention_backend/sparse/kernel.py`:
- Around line 1843-1854: The code computes block_id = tok // BLOCK_SIZE and only
checks the upper bound (block_id < max_num_blocks_per_req) before computing
bt_ptr and calling tl.load, which allows negative block_id when tok == -1;
update the validity check for accessing block_table to require non-negative
tokens as well (e.g., make valid_block require tok >= 0 in addition to block_id
< max_num_blocks_per_req) so bt_ptr (computed from block_table_ptr, req,
bt_stride0, bt_stride1) is only used when the block index is non-negative and
within max_num_blocks_per_req before calling tl.load.
- Around line 1884-1886: The parameter annotation stride_factor currently uses a
bare default None with type int; change it to use typing.Optional[int] to be
explicit about optionality (e.g., stride_factor: Optional[int] = None),
add/import Optional from typing if not already imported, and update any related
docstring/type comments if present; locate the parameter named stride_factor in
the function signature in kernel.py to apply this change.
- Around line 1937-1941: The 2D branch sets ti_stride0 = 0 which implies
broadcasting token_indices_c across KV heads; add a runtime assertion to ensure
this is only used when num_kv_heads == 1 to prevent silent misuse. Locate the
conditional handling token_indices_c (the if token_indices_c.ndim == 2 block)
and insert an assert or raise that checks num_kv_heads == 1 (or equivalent
variable) before setting ti_stride0 = 0, referencing token_indices_c,
ti_stride0, and num_kv_heads in the message.
In `@tensorrt_llm/_torch/attention_backend/trtllm.py`:
- Around line 1950-1961: Change the stub sparse_attn_ctx_predict to return
Optional[torch.Tensor] instead of Tuple[...] and make it raise
NotImplementedError like the other base predict methods; keep the unused
parameters by adding a "# noqa: ARG002" to the function signature (or the unused
parameter declarations) so Ruff is satisfied, and ensure the docstring/comments
indicate derived classes must implement this (referencing the
TrtllmAttentionMetadata and the sparse_attn_ctx_predict symbol so reviewers can
find it).
In `@tests/unittest/_torch/attention/sparse/test_sparse_attention.py`:
- Around line 303-304: Several unpacked variables in the sparse attention tests
are unused and flagged by Ruff (RUF059): rename num_tokens, device, and
request_ids to _num_tokens, _device, and _request_ids (or prefix with an
underscore) wherever they are unpacked from sparse_attn_ctx_indices or
request-related tuples (e.g., the unpack using sparse_attn_ctx_indices.shape and
any request_ids unpacks around the test helper usages) so the linter recognizes
them as intentionally unused; update all occurrences mentioned (also the other
unpack sites flagged) consistently.
- Around line 129-136: The overrides sparse_kv_predict, sparse_attn_predict, and
sparse_attn_ctx_predict currently declare unused parameters q, k, metadata, and
**kwargs; silence ARG002 by renaming them to _q, _k, _metadata and **_kwargs
(for sparse_attn_ctx_predict drop **_kwargs if not present) so linters treat
them as intentionally unused while preserving signatures and behavior.
♻️ Duplicate comments (1)
cpp/tensorrt_llm/nanobind/thop/bindings.cpp (1)
67-68: Same positional-arg compatibility risk as the pybind binding.See the note on the pybind variant: inserting
sparse_attn_ctx_indicesmid-signature and requiringsparse_attn_indices_block_sizecan break positional callers.
| xqaParams.use_sparse_attention_gen_paged = useTllmGenSparseAttentionPaged(); | ||
| // Skip softmax threshold. |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Find definitions of the sparse attention helpers
rg -n "useTllmGenSparseAttention" cpp/tensorrt_llm/common/attentionOp.cpp -C 3Repository: NVIDIA/TensorRT-LLM
Length of output: 1979
🏁 Script executed:
#!/bin/bash
# Find where these helpers are defined
rg -n "useTllmGenSparseAttention(Paged)?\s*\(" -C 3Repository: NVIDIA/TensorRT-LLM
Length of output: 5102
🏁 Script executed:
#!/bin/bash
# Get the file size and examine the relevant lines
wc -l cpp/tensorrt_llm/common/attentionOp.cppRepository: NVIDIA/TensorRT-LLM
Length of output: 107
🏁 Script executed:
#!/bin/bash
# Get the implementation of useSparseMLA and related helpers
rg -n "useSparseMLA|useSparseAttention" cpp/tensorrt_llm/common/attentionOp.h -C 2Repository: NVIDIA/TensorRT-LLM
Length of output: 706
🏁 Script executed:
#!/bin/bash
# Get broader context around line 2754 to understand the control flow
sed -n '2740,2770p' cpp/tensorrt_llm/common/attentionOp.cppRepository: NVIDIA/TensorRT-LLM
Length of output: 1347
Line 2754-2756 needs a guard: useTllmGenSparseAttention() can be true when mPagedKVCache == false.
The code unconditionally sets AttentionInputLayout::Q_PAGED_KV when useTllmGenSparseAttention() returns true. However, this helper does not check mPagedKVCache—it only checks sparse attention and TllmGen flags. If contiguous KV is in use, the paged layout will cause invalid memory reads.
Line 298 is already safe because useTllmGenSparseAttentionPaged() enforces the paged invariant through useSparseAttention(), which includes && mPagedKVCache. The issue is at lines 2754-2756:
🔧 Required fix
else if (useTllmGenSparseAttention())
+ && mPagedKVCache)
{
fmhaParams.attentionInputLayout = AttentionInputLayout::Q_PAGED_KV;
}🤖 Prompt for AI Agents
In `@cpp/tensorrt_llm/common/attentionOp.cpp` around lines 298 - 299, The code
unconditionally sets AttentionInputLayout::Q_PAGED_KV when
useTllmGenSparseAttention() is true, but useTllmGenSparseAttention() does not
check mPagedKVCache and can cause invalid reads when KV is contiguous; change
the guard to require the paged invariant (e.g., use
useTllmGenSparseAttentionPaged() or explicitly check mPagedKVCache) before
assigning Q_PAGED_KV so the paged layout is only used when mPagedKVCache is
true; update the branch that currently calls useTllmGenSparseAttention() (the
one setting AttentionInputLayout::Q_PAGED_KV) to
useTllmGenSparseAttentionPaged() or add && mPagedKVCache to the condition.
| // Whether to use sparse MLA. | ||
| bool mSparseMla; | ||
| bool mSparseAttention; | ||
| // The top k value for sparse MLA. | ||
| int mSparseMlaTopK; | ||
| int mSparseTopK; |
There was a problem hiding this comment.
Update sparse-attention wording in comments.
Line 288 and Line 290 still say “sparse MLA,” but the fields were renamed to sparse attention. This can mislead future readers.
✏️ Suggested comment update
- // Whether to use sparse MLA.
+ // Whether to use sparse attention.
bool mSparseAttention;
- // The top k value for sparse MLA.
+ // The top-k value for sparse attention.
int mSparseTopK;📝 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.
| // Whether to use sparse MLA. | |
| bool mSparseMla; | |
| bool mSparseAttention; | |
| // The top k value for sparse MLA. | |
| int mSparseMlaTopK; | |
| int mSparseTopK; | |
| // Whether to use sparse attention. | |
| bool mSparseAttention; | |
| // The top-k value for sparse attention. | |
| int mSparseTopK; |
🤖 Prompt for AI Agents
In `@cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaRunnerParams.h` around
lines 288 - 291, Update the comments for the fields mSparseAttention and
mSparseTopK to consistently use "sparse attention" instead of "sparse MLA":
replace the comment "Whether to use sparse MLA." above mSparseAttention with
"Whether to use sparse attention." and replace "The top k value for sparse MLA."
above mSparseTopK with "The top-k value for sparse attention." to match the
renamed fields.
| // The top k value for sparse MLA. | ||
| int32_t mNumSparseTopk; | ||
| // The number of tokensQ per CTA (used for groupsHeadsTokensQ generation kernel). | ||
| int32_t mNumTokensPerCtaQ; |
There was a problem hiding this comment.
Update the mNumSparseTopk comment to match new naming.
The field now applies to general sparse attention, not just MLA.
📝 Comment fix
- // The top k value for sparse MLA.
+ // The top-k value for sparse attention.🤖 Prompt for AI Agents
In `@cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/kernelParams.h` around lines
151 - 154, Update the comment for the int32_t member mNumSparseTopk in
kernelParams.h to reflect that this value applies to general sparse attention
(not just MLA); locate the declaration of mNumSparseTopk and replace "The top k
value for sparse MLA." with a concise comment like "The top-k value used for
sparse attention." so the name and intent match the new general usage.
| int32_t numInstsQ{kernelMeta.mStepQ / kernelMeta.mTileSizeQ}; | ||
| int32_t tileSizePerCtaQ{kernelMeta.mTileSizeQ * numInstsQ}; | ||
| // TODO(yuhangh): check the correctness | ||
| if (kernelMeta.mGroupsHeadsQ && !isSpecDecodingGenerationKernel(options.mKernelType)) | ||
| { | ||
| tileSizePerCtaQ = 1; | ||
| } | ||
| params.mNumTokensPerCtaQ = tileSizePerCtaQ; |
There was a problem hiding this comment.
Resolve the TODO on tileSizePerCtaQ correctness.
This path is performance/accuracy critical; please validate and remove the TODO or track it with a follow-up issue and coverage.
If you want, I can help draft a focused test or a follow-up issue plan.
🤖 Prompt for AI Agents
In `@cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/kernelParams.h` around lines
718 - 725, Replace the TODO by making the tileSizePerCtaQ calculation explicit
and tested: keep the default tileSizePerCtaQ = kernelMeta.mTileSizeQ *
numInstsQ, and when kernelMeta.mGroupsHeadsQ is true and
!isSpecDecodingGenerationKernel(options.mKernelType) set tileSizePerCtaQ =
kernelMeta.mTileSizeQ (one instance per CTA) — update params.mNumTokensPerCtaQ
accordingly, add a short comment referencing numInstsQ, tileSizePerCtaQ,
kernelMeta.mGroupsHeadsQ, isSpecDecodingGenerationKernel and
options.mKernelType, and add a tiny unit/integration test or open a follow-up
issue to cover this branch to validate correctness and performance.
| TLLM_CHECK_WITH_INFO( | ||
| !options.mSparseMla || (options.mSparseMlaTopK % 4) == 0, "SparseMlaTopK must be a multiple of 4"); | ||
| params.mSparseMlaTopK = options.mSparseMlaTopK; | ||
| params.mUseBlockSparseAttention = options.mUseBlockSparseAttention; | ||
| !options.mSparseAttention || (options.mSparseTopK % 4) == 0, "SparseTopK must be a multiple of 4"); | ||
| params.mNumSparseTopk = options.mSparseTopK; | ||
| if (options.mSparseAttention) | ||
| { | ||
| params.mSparsityMode = 0; | ||
| } | ||
| else if (options.mUseBlockSparseAttention) | ||
| { | ||
| params.mSparsityMode = 1; | ||
| } |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# First, find and read the KernelParams structure definition
fd -e h -e hpp "kernelParams" cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha | head -5Repository: NVIDIA/TensorRT-LLM
Length of output: 124
🏁 Script executed:
#!/bin/bash
# Read the KernelParams struct to see mSparsityMode declaration
cat -n cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/kernelParams.h | sed -n '1,100p'Repository: NVIDIA/TensorRT-LLM
Length of output: 4200
🏁 Script executed:
#!/bin/bash
# Find mSparsityMode field declaration in kernelParams.h
rg -n "mSparsityMode" cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/kernelParams.h -A 2 -B 2Repository: NVIDIA/TensorRT-LLM
Length of output: 576
🏁 Script executed:
#!/bin/bash
# Find all uses of mSparsityMode in the codebase to understand how it's used
rg -n "mSparsityMode" cpp/tensorrt_llm/kernels -A 1 -B 1Repository: NVIDIA/TensorRT-LLM
Length of output: 849
🏁 Script executed:
#!/bin/bash
# Check if there are existing constant definitions for sparsity modes
rg -n "SPARSITY_MODE|SPARSE_ATTENTION" cpp/tensorrt_llm/kernels -A 1 -B 1Repository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
#!/bin/bash
# Search for mSparsityMode in all kernel files (.cu files)
rg -n "mSparsityMode" cpp/tensorrt_llm/kernels --type cuRepository: NVIDIA/TensorRT-LLM
Length of output: 89
🏁 Script executed:
#!/bin/bash
# Look for how sparsity modes are handled in the actual kernel code
rg -n "sparsity|Sparsity" cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha --type cu -A 2 -B 2 | head -100Repository: NVIDIA/TensorRT-LLM
Length of output: 89
🏁 Script executed:
#!/bin/bash
# Check the full context around line 859-869 in kernelParams.h
cat -n cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/kernelParams.h | sed -n '850,880p'Repository: NVIDIA/TensorRT-LLM
Length of output: 1550
🏁 Script executed:
#!/bin/bash
# Check how mSparsityMode is actually used in the kernels - look for conditional checks
rg -n "if.*mSparsityMode|mSparsityMode.*==" cpp/tensorrt_llm/kernels -A 3 -B 1Repository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
#!/bin/bash
# Check the struct definition more broadly to see initialization patterns
cat -n cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/kernelParams.h | sed -n '40,200p'Repository: NVIDIA/TensorRT-LLM
Length of output: 7876
🏁 Script executed:
#!/bin/bash
# Search for kernel implementations that read mSparsityMode
find cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha -type f -name "*.cu" -o -name "*.h" | head -20Repository: NVIDIA/TensorRT-LLM
Length of output: 704
🏁 Script executed:
#!/bin/bash
# Look for how KernelParams is used in kernels - search for includes/uses
rg -n "KernelParams" cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha -A 3 -B 1Repository: NVIDIA/TensorRT-LLM
Length of output: 32461
🏁 Script executed:
#!/bin/bash
# Check if mSparsityMode is accessed in kernel launch or template specialization
rg -n "mSparsityMode" cpp/tensorrt_llm/kernels --max-count 50Repository: NVIDIA/TensorRT-LLM
Length of output: 363
🏁 Script executed:
#!/bin/bash
# Check if there's a constructor or initialization pattern for KernelParams
cat -n cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/kernelParams.h | sed -n '750,875p'Repository: NVIDIA/TensorRT-LLM
Length of output: 6974
🏁 Script executed:
#!/bin/bash
# Check how mSparsityMode is actually used in kernel code - search for any comparison or conditional use
rg -n "params\.mSparsityMode|kernelParams.*mSparsityMode" cpp/tensorrt_llm --max-count 20 -B 2 -A 2Repository: NVIDIA/TensorRT-LLM
Length of output: 953
🏁 Script executed:
#!/bin/bash
# Look for any documentation or comments about sparsity mode values
rg -n "sparsity|sparse.*mode" cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha -i -B 1 -A 1 | head -60Repository: NVIDIA/TensorRT-LLM
Length of output: 960
🏁 Script executed:
#!/bin/bash
# Check for the full struct initialization to see if there's a constructor or initializer list
cat -n cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/kernelParams.h | sed -n '650,675p'Repository: NVIDIA/TensorRT-LLM
Length of output: 1301
Initialize mSparsityMode for dense attention paths and replace magic literals with named constants.
KernelParams params at line 658 leaves mSparsityMode uninitialized when neither mSparseAttention nor mUseBlockSparseAttention is enabled. This is undefined behavior if kernel code reads the field unguarded. Additionally, the alignment literal 4 at line 860 should be extracted as a named constant per coding guidelines.
Proposed fix
+ constexpr int32_t kSPARSE_TOPK_ALIGNMENT = 4;
+ constexpr int32_t kSPARSITY_MODE_NONE = -1;
+ constexpr int32_t kSPARSITY_MODE_SPARSE_ATTENTION = 0;
+ constexpr int32_t kSPARSITY_MODE_BLOCK_SPARSE_ATTENTION = 1;
+
+ params.mSparsityMode = kSPARSITY_MODE_NONE;
TLLM_CHECK_WITH_INFO(
- !options.mSparseAttention || (options.mSparseTopK % 4) == 0, "SparseTopK must be a multiple of 4");
+ !options.mSparseAttention || (options.mSparseTopK % kSPARSE_TOPK_ALIGNMENT) == 0, "SparseTopK must be a multiple of 4");
params.mNumSparseTopk = options.mSparseTopK;
if (options.mSparseAttention)
{
- params.mSparsityMode = 0;
+ params.mSparsityMode = kSPARSITY_MODE_SPARSE_ATTENTION;
}
else if (options.mUseBlockSparseAttention)
{
- params.mSparsityMode = 1;
+ params.mSparsityMode = kSPARSITY_MODE_BLOCK_SPARSE_ATTENTION;
}🤖 Prompt for AI Agents
In `@cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/kernelParams.h` around lines
859 - 869, The code doesn't initialize params.mSparsityMode when neither
options.mSparseAttention nor options.mUseBlockSparseAttention is true and uses
the magic literal 4 for top-k alignment; update the block around
TLLM_CHECK_WITH_INFO and the subsequent assignments so that the alignment
literal is replaced with a named constant (e.g., kSparseTopKAlignment) and set
params.mSparsityMode to a clear default (e.g., 0 for dense) when neither
options.mSparseAttention nor options.mUseBlockSparseAttention is enabled;
reference the existing symbols options.mSparseAttention,
options.mUseBlockSparseAttention, params.mSparsityMode, params.mNumSparseTopk
and introduce the constant name for the alignment check.
| stride_factor: | ||
| int = None, # for strided memory layout (with layer interleaving), defaults to BLOCK_SIZE | ||
| layer_id: int = 0, # for layer interleaving layout | ||
| layer_id: int = 0, # for layer interleaving layout |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, check if the file exists and read the relevant lines
head -1890 tensorrt_llm/_torch/attention_backend/sparse/kernel.py | tail -20Repository: NVIDIA/TensorRT-LLM
Length of output: 821
🏁 Script executed:
# Check for typing imports in the file
rg "^import typing" tensorrt_llm/_torch/attention_backend/sparse/kernel.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
# Get the full context around lines 1884-1886
sed -n '1870,1900p' tensorrt_llm/_torch/attention_backend/sparse/kernel.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 1307
🏁 Script executed:
# Search for similar implicit Optional patterns (parameter with : type = None)
rg ":\s*(int|str|float|bool|dict|list|tuple)\s*=\s*None" tensorrt_llm/_torch/attention_backend/sparse/kernel.py | head -20Repository: NVIDIA/TensorRT-LLM
Length of output: 129
🏁 Script executed:
# Check all imports at the top of the file
head -30 tensorrt_llm/_torch/attention_backend/sparse/kernel.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 1067
🏁 Script executed:
# Find all instances of implicit Optional patterns (: type = None)
rg ":\s*(int|str|float|bool|dict|list|tuple|Tensor)\s*=\s*None" tensorrt_llm/_torch/attention_backend/sparse/kernel.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 129
🏁 Script executed:
# Check if typing module is used anywhere in the file
rg "typing\." tensorrt_llm/_torch/attention_backend/sparse/kernel.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
# Count total implicit Optional patterns in the file
rg ":\s*(int|str|float|bool|dict|list|tuple|Tensor)\s*=\s*None" tensorrt_llm/_torch/attention_backend/sparse/kernel.py | wc -lRepository: NVIDIA/TensorRT-LLM
Length of output: 64
🏁 Script executed:
# Get detailed list with context
rg -B2 ":\s*(int|str|float|bool|dict|list|tuple|Tensor)\s*=\s*None" tensorrt_llm/_torch/attention_backend/sparse/kernel.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 265
Fix Optional typing for stride_factor.
stride_factor: int = None is an implicit Optional; use typing.Optional[int] to align with PEP 484 (and maintain Python 3.8 compatibility).
Suggested fix
import math
+import typing
import torch
@@
- stride_factor: int = None, # for strided memory layout (with layer interleaving), defaults to BLOCK_SIZE
+ stride_factor: typing.Optional[int] = None, # for strided memory layout (with layer interleaving), defaults to BLOCK_SIZE🧰 Tools
🪛 Ruff (0.14.13)
1885-1885: PEP 484 prohibits implicit Optional
Convert to T | None
(RUF013)
🤖 Prompt for AI Agents
In `@tensorrt_llm/_torch/attention_backend/sparse/kernel.py` around lines 1884 -
1886, The parameter annotation stride_factor currently uses a bare default None
with type int; change it to use typing.Optional[int] to be explicit about
optionality (e.g., stride_factor: Optional[int] = None), add/import Optional
from typing if not already imported, and update any related docstring/type
comments if present; locate the parameter named stride_factor in the function
signature in kernel.py to apply this change.
| if token_indices_c.ndim == 2: # [num_tokens, NUM_TOPK_TOKENS] num_kv_heads=1 by default | ||
| ti_stride1, ti_stride2 = token_indices_c.stride() | ||
| ti_stride0 = 0 | ||
| else: # [num_kv_heads, num_tokens, NUM_TOPK_TOKENS] | ||
| ti_stride0, ti_stride1, ti_stride2 = token_indices_c.stride() |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, let's check the file exists and examine the specific lines
wc -l tensorrt_llm/_torch/attention_backend/sparse/kernel.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 122
🏁 Script executed:
# Read the code around lines 1937-1941 with extended context
sed -n '1920,1960p' tensorrt_llm/_torch/attention_backend/sparse/kernel.py | cat -nRepository: NVIDIA/TensorRT-LLM
Length of output: 1890
🏁 Script executed:
# Search for the function containing these lines to understand the broader context
sed -n '1900,1950p' tensorrt_llm/_torch/attention_backend/sparse/kernel.py | grep -E "^def |class " -B 50 | head -60Repository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
# Find the function definition that contains line 1937
awk 'NR<=1937' tensorrt_llm/_torch/attention_backend/sparse/kernel.py | grep "^def " | tail -1Repository: NVIDIA/TensorRT-LLM
Length of output: 108
🏁 Script executed:
# Get the full context: look at a wider range to find the function
sed -n '1850,2000p' tensorrt_llm/_torch/attention_backend/sparse/kernel.py | cat -nRepository: NVIDIA/TensorRT-LLM
Length of output: 5823
🏁 Script executed:
# Search for calls to triton_convert_req_index_to_global_index to understand how it's invoked
rg "triton_convert_req_index_to_global_index" -B 5 -A 5Repository: NVIDIA/TensorRT-LLM
Length of output: 5300
🏁 Script executed:
# Check if there's any validation on token_indices shape or num_kv_heads constraint
rg "token_indices.*ndim|token_indices.*shape" -B 3 -A 3Repository: NVIDIA/TensorRT-LLM
Length of output: 4821
🏁 Script executed:
# Look at the kernel implementation to understand how ti_stride0 is used
sed -n '1900,1936p' tensorrt_llm/_torch/attention_backend/sparse/kernel.py | cat -nRepository: NVIDIA/TensorRT-LLM
Length of output: 2020
🏁 Script executed:
# Check for any existing assertions or guards related to token_indices dimensionality
rg "assert.*token_indices|token_indices.*assert" -B 2 -A 2Repository: NVIDIA/TensorRT-LLM
Length of output: 3221
🏁 Script executed:
# Check how topk_indices is created in dsa.py to see if it can be 2D or 3D
grep -B 20 "triton_convert_req_index_to_global_index" tensorrt_llm/_torch/attention_backend/sparse/dsa.py | grep -A 20 "topk_indices"Repository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
# Look at sparse_attn_ctx_indices creation in test
grep -B 20 "triton_convert_req_index_to_global_index" tests/unittest/_torch/attention/sparse/test_sparse_attention.py | grep -A 20 "sparse_attn_ctx_indices"Repository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
# Get more context around the dsa.py call to understand topk_indices shape
sed -n '1,300p' tensorrt_llm/_torch/attention_backend/sparse/dsa.py | tail -200Repository: NVIDIA/TensorRT-LLM
Length of output: 8401
🏁 Script executed:
# Search for where topk_indices is created in dsa.py
rg "topk_indices.*=" tensorrt_llm/_torch/attention_backend/sparse/dsa.py -B 5 -A 5Repository: NVIDIA/TensorRT-LLM
Length of output: 7934
🏁 Script executed:
# Check the test to see what shapes are being used
sed -n '1,100p' tests/unittest/_torch/attention/sparse/test_sparse_attention.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 2791
🏁 Script executed:
# Look for all calls to triton_convert_req_index_to_global_index to check if num_kv_heads is ever passed
rg "triton_convert_req_index_to_global_index" -A 15 -B 2Repository: NVIDIA/TensorRT-LLM
Length of output: 7921
🏁 Script executed:
# Check what sparse_attn_ctx_indices looks like in the test
grep -B 30 "sparse_attn_ctx_indices" tests/unittest/_torch/attention/sparse/test_sparse_attention.py | head -50Repository: NVIDIA/TensorRT-LLM
Length of output: 1748
🏁 Script executed:
# Look for where sparse_attn_ctx_indices is created in the test
grep -B 50 "sparse_attn_ctx_indices," tests/unittest/_torch/attention/sparse/test_sparse_attention.py | grep -A 50 "sparse_attn_ctx_indices"Repository: NVIDIA/TensorRT-LLM
Length of output: 4477
🏁 Script executed:
# Check if sparse_attn_ctx_indices is 2D or 3D in the test
grep -B 10 "sparse_attn_ctx_indices = " tests/unittest/_torch/attention/sparse/test_sparse_attention.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 1166
🏁 Script executed:
# Search more broadly for sparse_attn_ctx_indices creation
rg "sparse_attn_ctx_indices\s*=" tests/unittest/_torch/attention/sparse/ -B 5 -A 5Repository: NVIDIA/TensorRT-LLM
Length of output: 3974
🏁 Script executed:
# Find generate_sparse_attn_ctx_indices to see what shape it creates
rg "def generate_sparse_attn_ctx_indices" -A 20Repository: NVIDIA/TensorRT-LLM
Length of output: 2272
🏁 Script executed:
# Check if there are any calls passing 2D token_indices by looking at topk_indices shape
rg "topk_indices\.shape" -B 5 -A 5Repository: NVIDIA/TensorRT-LLM
Length of output: 1975
🏁 Script executed:
# Also check the dsa.py call more carefully - does it ever pass num_kv_heads?
sed -n '1,200p' tensorrt_llm/_torch/attention_backend/sparse/dsa.py | grep -A 20 "triton_convert_req_index_to_global_index"Repository: NVIDIA/TensorRT-LLM
Length of output: 1300
Add assertion to validate 2D token_indices constraints.
When token_indices is 2D with ti_stride0 = 0, the indices are broadcast across all KV heads. This is only safe for num_kv_heads == 1. Add an assertion to prevent silent misuse:
Suggested guard
if token_indices_c.ndim == 2: # [num_tokens, NUM_TOPK_TOKENS] num_kv_heads=1 by default
+ assert num_kv_heads == 1, "2D token_indices is only supported for num_kv_heads == 1"
ti_stride1, ti_stride2 = token_indices_c.stride()
ti_stride0 = 0🤖 Prompt for AI Agents
In `@tensorrt_llm/_torch/attention_backend/sparse/kernel.py` around lines 1937 -
1941, The 2D branch sets ti_stride0 = 0 which implies broadcasting
token_indices_c across KV heads; add a runtime assertion to ensure this is only
used when num_kv_heads == 1 to prevent silent misuse. Locate the conditional
handling token_indices_c (the if token_indices_c.ndim == 2 block) and insert an
assert or raise that checks num_kv_heads == 1 (or equivalent variable) before
setting ti_stride0 = 0, referencing token_indices_c, ti_stride0, and
num_kv_heads in the message.
| def sparse_attn_ctx_predict( | ||
| self, | ||
| q: torch.Tensor, | ||
| k: Optional[torch.Tensor], | ||
| metadata: TrtllmAttentionMetadata, | ||
| **kwargs, | ||
| ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: | ||
| """ | ||
| Predict sparse attn ctx indices. It's implemented in the derived class. | ||
| """ | ||
| # TODO(yuhangh): complete this part logic | ||
| return None |
There was a problem hiding this comment.
Align sparse_attn_ctx_predict return type and silence Ruff ARG002.
The method currently returns None while annotated as a Tuple[...], but callers expect a single tensor (Line 1632). Ruff also flags unused arguments (ARG002). Consider returning Optional[torch.Tensor] and raising NotImplementedError like other base predict methods, with a # noqa: ARG002 to keep lint clean.
🧹 Suggested stub alignment
- def sparse_attn_ctx_predict(
+ def sparse_attn_ctx_predict( # noqa: ARG002
self,
q: torch.Tensor,
k: Optional[torch.Tensor],
metadata: TrtllmAttentionMetadata,
**kwargs,
- ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
+ ) -> Optional[torch.Tensor]:
"""
Predict sparse attn ctx indices. It's implemented in the derived class.
"""
- # TODO(yuhangh): complete this part logic
- return None
+ raise NotImplementedError📝 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.
| def sparse_attn_ctx_predict( | |
| self, | |
| q: torch.Tensor, | |
| k: Optional[torch.Tensor], | |
| metadata: TrtllmAttentionMetadata, | |
| **kwargs, | |
| ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: | |
| """ | |
| Predict sparse attn ctx indices. It's implemented in the derived class. | |
| """ | |
| # TODO(yuhangh): complete this part logic | |
| return None | |
| def sparse_attn_ctx_predict( # noqa: ARG002 | |
| self, | |
| q: torch.Tensor, | |
| k: Optional[torch.Tensor], | |
| metadata: TrtllmAttentionMetadata, | |
| **kwargs, | |
| ) -> Optional[torch.Tensor]: | |
| """ | |
| Predict sparse attn ctx indices. It's implemented in the derived class. | |
| """ | |
| raise NotImplementedError |
🧰 Tools
🪛 Ruff (0.14.13)
1952-1952: Unused method argument: q
(ARG002)
1953-1953: Unused method argument: k
(ARG002)
1954-1954: Unused method argument: metadata
(ARG002)
1955-1955: Unused method argument: kwargs
(ARG002)
🤖 Prompt for AI Agents
In `@tensorrt_llm/_torch/attention_backend/trtllm.py` around lines 1950 - 1961,
Change the stub sparse_attn_ctx_predict to return Optional[torch.Tensor] instead
of Tuple[...] and make it raise NotImplementedError like the other base predict
methods; keep the unused parameters by adding a "# noqa: ARG002" to the function
signature (or the unused parameter declarations) so Ruff is satisfied, and
ensure the docstring/comments indicate derived classes must implement this
(referencing the TrtllmAttentionMetadata and the sparse_attn_ctx_predict symbol
so reviewers can find it).
| def sparse_kv_predict(self, q, k, metadata, **kwargs): | ||
| return self._sparse_kv_indices, self._sparse_kv_offsets | ||
|
|
||
| def sparse_attn_predict(self, q, k, metadata, **kwargs): | ||
| return self._sparse_attn_indices, self._sparse_attn_offsets | ||
|
|
||
| def sparse_attn_ctx_predict(self, q, k, metadata, **kwargs): | ||
| return self._sparse_attn_ctx_indices |
There was a problem hiding this comment.
Silence unused-argument lints in TestSparseAttention overrides.
q, k, metadata, and kwargs aren’t used in these overrides. Prefix with _ to avoid ARG002 violations and keep lint clean.
✅ Suggested cleanup
- def sparse_kv_predict(self, q, k, metadata, **kwargs):
+ def sparse_kv_predict(self, _q, _k, _metadata, **_kwargs):
return self._sparse_kv_indices, self._sparse_kv_offsets
- def sparse_attn_predict(self, q, k, metadata, **kwargs):
+ def sparse_attn_predict(self, _q, _k, _metadata, **_kwargs):
return self._sparse_attn_indices, self._sparse_attn_offsets
- def sparse_attn_ctx_predict(self, q, k, metadata, **kwargs):
+ def sparse_attn_ctx_predict(self, _q, _k, _metadata, **_kwargs):
return self._sparse_attn_ctx_indices🧰 Tools
🪛 Ruff (0.14.13)
129-129: Unused method argument: q
(ARG002)
129-129: Unused method argument: k
(ARG002)
129-129: Unused method argument: metadata
(ARG002)
129-129: Unused method argument: kwargs
(ARG002)
132-132: Unused method argument: q
(ARG002)
132-132: Unused method argument: k
(ARG002)
132-132: Unused method argument: metadata
(ARG002)
132-132: Unused method argument: kwargs
(ARG002)
135-135: Unused method argument: q
(ARG002)
135-135: Unused method argument: k
(ARG002)
135-135: Unused method argument: metadata
(ARG002)
135-135: Unused method argument: kwargs
(ARG002)
🤖 Prompt for AI Agents
In `@tests/unittest/_torch/attention/sparse/test_sparse_attention.py` around lines
129 - 136, The overrides sparse_kv_predict, sparse_attn_predict, and
sparse_attn_ctx_predict currently declare unused parameters q, k, metadata, and
**kwargs; silence ARG002 by renaming them to _q, _k, _metadata and **_kwargs
(for sparse_attn_ctx_predict drop **_kwargs if not present) so linters treat
them as intentionally unused while preserving signatures and behavior.
| num_kv_heads, num_tokens, num_sparse_tokens = sparse_attn_ctx_indices.shape | ||
| device = sparse_attn_ctx_indices.device |
There was a problem hiding this comment.
Clean up unused unpacked variables flagged by Ruff.
num_tokens, device, and request_ids are unused; prefix with _ to satisfy RUF059.
✅ Suggested cleanup
- num_kv_heads, num_tokens, num_sparse_tokens = sparse_attn_ctx_indices.shape
+ num_kv_heads, _num_tokens, num_sparse_tokens = sparse_attn_ctx_indices.shape- (
- device,
+ (
+ _device,
q,
k,
v,
sparse_kv_indices,
sparse_kv_offsets,
kv_cache_manager,
request_ids,
metadata,
attention,
) = _setup_context_test(s)- device, q, k, v, _, _, kv_cache_manager, request_ids, metadata, _ = _setup_context_test(s)
+ device, q, k, v, _, _, kv_cache_manager, _request_ids, metadata, _ = _setup_context_test(s)Also applies to: 855-866, 971-972
🧰 Tools
🪛 Ruff (0.14.13)
303-303: Unpacked variable num_tokens is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
🤖 Prompt for AI Agents
In `@tests/unittest/_torch/attention/sparse/test_sparse_attention.py` around lines
303 - 304, Several unpacked variables in the sparse attention tests are unused
and flagged by Ruff (RUF059): rename num_tokens, device, and request_ids to
_num_tokens, _device, and _request_ids (or prefix with an underscore) wherever
they are unpacked from sparse_attn_ctx_indices or request-related tuples (e.g.,
the unpack using sparse_attn_ctx_indices.shape and any request_ids unpacks
around the test helper usages) so the linter recognizes them as intentionally
unused; update all occurrences mentioned (also the other unpack sites flagged)
consistently.
|
Close this PR because a new PR which is the same as this PR has merged |
1 participant
Summary by CodeRabbit
New Features
Improvements
✏️ Tip: You can customize this high-level summary in your review settings.
Description
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 thestage-listparameter 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.mdand the
scripts/test_to_stage_mapping.pyhelper.kill
killKill all running builds associated with pull request.
skip
skip --comment COMMENTSkip 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-pipelineReuse 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.