[https://nvbugs/5669671][fix] Support GuidedDecoder with sharded logits (pick #10698)

[syuoni]

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

Summary by CodeRabbit

Release Notes

New Features

• Guided decoding now supports distributed vocabulary configurations with rank-based tensor partitioning for multi-rank inference deployments.

Performance

• Optimized logits masking kernels to improve execution efficiency in mask application operations.

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

[syuoni]

/bot run --disable-fail-fast

[coderabbitai]

📝 Walkthrough

Walkthrough

The pull request refactors bitmask parameter handling in CUDA kernels by replacing size-based parameters with stride-based parameters, and extends the guided decoder to support vocabulary sharding across multiple ranks via rank-aware bitmask slicing.

Changes

Cohort / File(s)	Summary
CUDA Kernel Refactoring cpp/tensorrt_llm/kernels/logitsBitmask.cu, cpp/tensorrt_llm/kernels/logitsBitmask.h	Replaced bitmaskSize parameter with bitmaskStride in contiguousLogitsBitmaskKernel and related dispatch functions; removed bitmaskSize entirely from logitsBitmaskKernel. Updated all kernel invocations and template instantiations to reflect the new stride-based parameter semantics.
PyTorch Binding Updates cpp/tensorrt_llm/thop/logitsBitmaskOp.cpp	Added constexpr bits-per-mask-element constant, replaced hardcoded bitmaskSize computation with bitmask.stride(0) calls, converted type casts from reinterpret_cast to static_cast, and added validity check for bitmask size when vocabulary is not sharded.
Guided Decoder Rank Support tensorrt_llm/_torch/pyexecutor/guided_decoder.py	Added rank parameter to GuidedDecoder and CapturableGuidedDecoder constructors; extended _apply_bitmask logic to slice bitmask and optional d2t tensors based on rank and vocabulary partitioning for sharded vocabulary scenarios.
Executor Configuration tensorrt_llm/_torch/pyexecutor/py_executor_creator.py	Propagated mapping.rank to guided decoder initialization via kwargs, enabling rank-aware vocabulary sharding in the decoder initialization path.

Sequence Diagram(s)

sequenceDiagram
    actor Executor
    participant GuidedDecoder
    participant Kernel
    
    Executor->>GuidedDecoder: __init__(rank=0, vocab_size_padded, ...)
    GuidedDecoder->>GuidedDecoder: Store rank as public attribute
    Executor->>Executor: Build bitmask and optional d2t tensors
    Executor->>GuidedDecoder: _apply_bitmask(bitmask, d2t, ...)
    
    alt Sharded Vocabulary
        GuidedDecoder->>GuidedDecoder: Compute vocab partition boundaries<br/>based on rank
        GuidedDecoder->>GuidedDecoder: Slice bitmask[bitmask_start:bitmask_end]
        GuidedDecoder->>GuidedDecoder: Slice d2t[d2t_start:d2t_end]
    else Non-Sharded
        GuidedDecoder->>GuidedDecoder: Use full bitmask and d2t
    end
    
    GuidedDecoder->>Kernel: invokeContiguousLogitsBitmask(<br/>logits, sliced_bitmask,<br/>tokenMask, d2t, bitmaskStride, ...)
    Kernel->>Kernel: Apply sliced bitmask to logits<br/>using stride semantics

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

🚥 Pre-merge checks | ✅ 1 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The PR description is largely incomplete, containing only the template structure with empty Description and Test Coverage sections and an unchecked checklist.	Add a clear explanation of the issue being fixed and the solution approach. Include specific test coverage details that validate the sharded logits functionality in GuidedDecoder.
Docstring Coverage	⚠️ Warning	Docstring coverage is 15.38% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly indicates a fix for supporting GuidedDecoder with sharded logits and references a cherry-pick from PR #10698, directly relating to the changeset.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

---

🧹 Recent nitpick comments

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

1-2: Update copyright year to 2025.

The file has been modified in this PR, so the copyright year should be updated from 2024 to 2025 to reflect the latest meaningful modification, per coding guidelines.

Suggested fix

 /*
- * Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2024-2025, NVIDIA CORPORATION.  All rights reserved.
  *

---

📜 Recent 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 bc5094c and a6e3ff4.

📒 Files selected for processing (5)

• cpp/tensorrt_llm/kernels/logitsBitmask.cu
• cpp/tensorrt_llm/kernels/logitsBitmask.h
• cpp/tensorrt_llm/thop/logitsBitmaskOp.cpp
• tensorrt_llm/_torch/pyexecutor/guided_decoder.py
• tensorrt_llm/_torch/pyexecutor/py_executor_creator.py

🧰 Additional context used

📓 Path-based instructions (6)

**/*.{cpp,cc,cxx,h,hpp,hxx,cu,cuh}

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

**/*.{cpp,cc,cxx,h,hpp,hxx,cu,cuh}: Closing braces of namespaces should have a comment saying the namespace it closes (e.g., } // namespace foo)
Prefer const or constexpr variables over #defines whenever possible
A variable that is not modified after its initialization should be declared as const
For naming of constants in C++, use uppercase snakecase with prefix 'k' (e.g., kDIGIT_NUM)
Except for 0, nullptr, true, and false, all other literals should only be used for variable initialization and not in comparisons or expressions
Use Allman indentation style for brace notation in C++ code
Put the semicolon for an empty for or while loop in a new line
The statement forming the body of a switch, while, do..while, or for statement must be a compound statement (use brace-delimited statements)
If and else statements should always be followed by brace-delimited statements, even if empty or a single statement
C++ filenames should use camelCase with first letter lowercase (e.g., thisIsAFilename.cpp)
All types (including class names) in C++ should use PascalCase with uppercase first letter (e.g., FooBarClass)
Local variables, methods, and namespaces in C++ should use camelCase with first letter lowercase (e.g., localFooBar)
Non-magic-number global variables that are non-static and not defined in anonymous namespace should use camelCase prefixed with 'g' (e.g., gDontUseGlobalFoos)
Non-magic-number global variables that are static or defined in an anonymous namespace should use camelCase prefixed with 's' (e.g., sMutableStaticGlobal)
Locally visible static variables should use camelCase with 's' as the first letter (e.g., static std::once_flag sFlag;)
Public, private, and protected class member variables should use camelCase prefixed with 'm' (e.g., mNbFooValues)
Do not use Hungarian notation in C++ except for 'apps hungarian' (e.g., 'nb' to indicate count: mNbLayers)
If a constructor parameter name conflicts with a public me...

Files:

• cpp/tensorrt_llm/thop/logitsBitmaskOp.cpp
• cpp/tensorrt_llm/kernels/logitsBitmask.cu
• cpp/tensorrt_llm/kernels/logitsBitmask.h

**/*.{cpp,cc,cxx,cu}

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

**/*.{cpp,cc,cxx,cu}: Use smart pointers for allocating objects on the heap in C++
Prefer unique_ptr for single resource ownership and shared_ptr for shared resource ownership in C++. Use weak_ptr only in exceptional cases
In C++ function calls where parameters are not obvious, use inline C comments to document the parameter (e.g., doSomeOperation(/* checkForErrors = */ false);)
Use the least forceful cast necessary in C++, or no cast if possible
Casting a pointer to void* in C++ should be implicit (except if removing const)
Casting in C++ should not remove any const or volatile qualification from the type of a pointer or reference
Do not use C-style casts (other than void casts) and functional notation casts (other than explicit constructor calls) in C++
Casting from void* to T* in C++ should be done with static_cast, not reinterpret_cast
Use reinterpret_cast in C++ as a last resort, where const_cast and static_cast won't work
Avoid dynamic_cast in C++
Do not use assignment operator in C++ subexpressions (e.g., x = y = z or if (x = y))
When practical, a C++ switch statement controlled by an enum should have a case for each enum value and not have a default clause
C++ switch statements should be well structured as structured multi-way branches, not as 'glorified gotos'
In C++ switch statements, prohibit fall-through except from one case label to another. Each case clause must be terminated with a break or throw
Do not end a C++ case clause with return; use break or throw instead
If a C++ switch clause is a compound statement, put the break inside the braces
Do not use C library functions in C++ whenever possible. Use C++ alternatives like brace initialization or std::fill_n() instead of memset()

Files:

• cpp/tensorrt_llm/thop/logitsBitmaskOp.cpp
• cpp/tensorrt_llm/kernels/logitsBitmask.cu

**/*.{h,hpp,hxx,cpp,cc,cxx,cu,cuh}

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

All C++ class templates, function templates, class template member functions, and class template static members must be instantiated at least once

Files:

• cpp/tensorrt_llm/thop/logitsBitmaskOp.cpp
• cpp/tensorrt_llm/kernels/logitsBitmask.cu
• cpp/tensorrt_llm/kernels/logitsBitmask.h

**/*.{cpp,cc,cxx,h,hpp,hxx,cu,cuh,py}

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

All TensorRT-LLM source files (.cpp, .h, .cu, .py, and other source files) should contain an NVIDIA copyright header with the year of latest meaningful modification

Files:

• cpp/tensorrt_llm/thop/logitsBitmaskOp.cpp
• tensorrt_llm/_torch/pyexecutor/guided_decoder.py
• cpp/tensorrt_llm/kernels/logitsBitmask.cu
• tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
• cpp/tensorrt_llm/kernels/logitsBitmask.h

**/*.py

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

**/*.py: The code developed for TensorRT-LLM should conform to Python 3.8+
Indent Python code with 4 spaces. Do not use tabs
Always maintain the namespace when importing Python modules, even if only one class or function from a module is used
Python filenames should use snake_case (e.g., some_file.py)
Python classes should use PascalCase (e.g., class SomeClass)
Python functions and methods should use snake_case (e.g., def my_awesome_function():)
Python local variables should use snake_case, with prefix k for variable names that start with a number (e.g., k_99th_percentile)
Python global variables should use upper snake_case with prefix G (e.g., G_MY_GLOBAL)
Python constants should use upper snake_case (e.g., MY_CONSTANT)
Avoid shadowing variables declared in an outer scope in Python
Initialize all externally visible members of a Python class in the constructor
For Python interfaces that may be used outside a file, prefer docstrings over comments
Use comments in Python for code within a function, or interfaces that are local to a file
Use Google-style docstrings for Python classes and functions, which can be parsed by Sphinx
Python attributes and variables can be documented inline with the format """<type>: Description"""
Avoid using reflection in Python when functionality can be easily achieved without reflection
When using try-except blocks in Python, limit the except clause to the smallest set of errors possible
When using try-except blocks in Python to handle multiple possible variable types (duck-typing), keep the body of the try as small as possible and use the else block for the main logic

Files:

• tensorrt_llm/_torch/pyexecutor/guided_decoder.py
• tensorrt_llm/_torch/pyexecutor/py_executor_creator.py

**/*.{h,hpp,hxx}

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

**/*.{h,hpp,hxx}: Follow Doxygen rules for documenting new C++ class interfaces and function prototypes. Use //! for C++-style single-line comments and //!< for class members
Use a preprocessor guard in C++ header files with the format TRTLLM_<FILENAME>_H, where the filename is in uppercase with no underscores, no prefix underscores, and no trailing underscores

Files:

• cpp/tensorrt_llm/kernels/logitsBitmask.h

🧠 Learnings (13)

📚 Learning: 2025-09-02T13:42:44.885Z

Learnt from: pcastonguay
Repo: NVIDIA/TensorRT-LLM PR: 7455
File: tensorrt_llm/_torch/pyexecutor/py_executor.py:1852-1860
Timestamp: 2025-09-02T13:42:44.885Z
Learning: In MPI communication within TensorRT-LLM pipeline parallelism, different communication types (tokens, logits, termination sync) must use disjoint tag namespaces to avoid message routing collisions when using the same source/destination patterns.

Applied to files:

• cpp/tensorrt_llm/thop/logitsBitmaskOp.cpp
• cpp/tensorrt_llm/kernels/logitsBitmask.cu

📚 Learning: 2025-09-19T21:28:13.751Z

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.

Applied to files:

• cpp/tensorrt_llm/thop/logitsBitmaskOp.cpp
• cpp/tensorrt_llm/kernels/logitsBitmask.cu

📚 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:

• cpp/tensorrt_llm/thop/logitsBitmaskOp.cpp
• cpp/tensorrt_llm/kernels/logitsBitmask.cu

📚 Learning: 2025-08-21T02:41:10.565Z

Learnt from: djns99
Repo: NVIDIA/TensorRT-LLM PR: 7104
File: cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_gemm_kernels.h:141-145
Timestamp: 2025-08-21T02:41:10.565Z
Learning: In TensorRT-LLM MOE GEMM kernels (cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_gemm_kernels.h), the stride_act and stride_weight pointers in TmaWarpSpecializedGroupedGemmInput are intentionally declared as void* rather than typed pointers because the actual stride type is determined at runtime based on factors like the swap_ab flag and layout decisions. This runtime type determination makes compile-time type safety impossible, so void* is the correct approach.

Applied to files:

• cpp/tensorrt_llm/thop/logitsBitmaskOp.cpp

📚 Learning: 2025-12-19T06:31:54.973Z

Learnt from: nvyocox
Repo: NVIDIA/TensorRT-LLM PR: 10117
File: tensorrt_llm/_torch/auto_deploy/transform/library/fuse_rope_attention.py:336-339
Timestamp: 2025-12-19T06:31:54.973Z
Learning: In tensorrt_llm/_torch/auto_deploy/transform/library/fuse_rope_attention.py, the cast to torch.float16 for qkv_node before creating the AttentionPlugin is intentional and required because DriveOS LLM expects float16 dtype specifically. This should not be changed to preserve original dtype or made configurable for bfloat16 models in the DriveOS LLM ONNX export path.

Applied to files:

• cpp/tensorrt_llm/thop/logitsBitmaskOp.cpp

📚 Learning: 2025-12-12T03:27:08.565Z

Learnt from: tongyuantongyu
Repo: NVIDIA/TensorRT-LLM PR: 9655
File: tensorrt_llm/_torch/pyexecutor/sampler.py:3031-3031
Timestamp: 2025-12-12T03:27:08.565Z
Learning: In files under tensorrt_llm/_torch/pyexecutor, avoid accessing torch.Tensor objects inside for-loops when iterating over requests. Convert batched tensors to Python lists beforehand using tensor.tolist(), and then iterate over those lists. This improves performance by reducing tensor-bound operations inside hot loops. Apply this pattern to similar code paths that process batches to access simple Python data structures (lists) inside loops.

Applied to files:

• tensorrt_llm/_torch/pyexecutor/guided_decoder.py
• tensorrt_llm/_torch/pyexecutor/py_executor_creator.py

📚 Learning: 2025-09-23T15:01:00.070Z

Learnt from: nv-lschneider
Repo: NVIDIA/TensorRT-LLM PR: 7910
File: cpp/tensorrt_llm/kernels/nccl_device/config.cu:15-17
Timestamp: 2025-09-23T15:01:00.070Z
Learning: In TensorRT-LLM NCCL device kernels, the <sstream> header is not needed as an explicit include in config.cu because it's provided transitively through other headers. Local compilation testing confirms this works without the explicit include.

Applied to files:

• cpp/tensorrt_llm/kernels/logitsBitmask.cu

📚 Learning: 2025-11-14T11:22:03.729Z

Learnt from: nzmora-nvidia
Repo: NVIDIA/TensorRT-LLM PR: 9163
File: tensorrt_llm/_torch/auto_deploy/custom_ops/quant.py:107-113
Timestamp: 2025-11-14T11:22:03.729Z
Learning: In TensorRT-LLM AutoDeploy custom ops, when adding hardware capability checks to select between kernel implementations (e.g., cuBLAS vs. CUDA kernel), use descriptive variable names that identify the specific GPU architectures or families being targeted (e.g., `is_blackwell_geforce_or_ada`) rather than generic names like `enable_cuda_core`. This makes it clear that the code is selecting an implementation path based on hardware capabilities, not enabling/disabling hardware features.

Applied to files:

• cpp/tensorrt_llm/kernels/logitsBitmask.cu

📚 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:

• cpp/tensorrt_llm/kernels/logitsBitmask.cu

📚 Learning: 2025-08-20T07:43:36.447Z

Learnt from: ChristinaZ
Repo: NVIDIA/TensorRT-LLM PR: 7068
File: cpp/tensorrt_llm/kernels/moeTopKFuncs.cuh:169-172
Timestamp: 2025-08-20T07:43:36.447Z
Learning: In TensorRT-LLM MOE kernels, when processing up to 128 experts across 32 threads, each thread handles at most 4 experts (N < 5 constraint), where N represents candidates per thread rather than total system capacity.

Applied to files:

• cpp/tensorrt_llm/kernels/logitsBitmask.cu

📚 Learning: 2025-08-26T06:07:02.166Z

Learnt from: shaharmor98
Repo: NVIDIA/TensorRT-LLM PR: 7231
File: tensorrt_llm/_torch/pyexecutor/_util.py:504-509
Timestamp: 2025-08-26T06:07:02.166Z
Learning: In tensorrt_llm/_torch/pyexecutor/_util.py, when calling model_engine.set_lora_model_config(), pass model_binding_config.mlp_hidden_size directly without multiplying by mapping.tp_size, as the mlp_hidden_size from get_bindings_model_config() is already the per-TP rank value needed for LoRA weight packaging.

Applied to files:

• tensorrt_llm/_torch/pyexecutor/py_executor_creator.py

📚 Learning: 2025-08-19T12:45:11.997Z

Learnt from: amitz-nv
Repo: NVIDIA/TensorRT-LLM PR: 7033
File: tensorrt_llm/_torch/pyexecutor/model_engine.py:0-0
Timestamp: 2025-08-19T12:45:11.997Z
Learning: In tensorrt_llm/_torch/pyexecutor/model_engine.py, DoRA (Delta Orthogonal Rank Adaptation) functionality was removed from the PyTorch flow to eliminate issues with inverted DoRA detection logic. The original is_dora condition was checking if scaling_vec_pointer == 0, which was potentially incorrect.

Applied to files:

• tensorrt_llm/_torch/pyexecutor/py_executor_creator.py

📚 Learning: 2025-08-09T20:57:04.084Z

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

Applied to files:

• cpp/tensorrt_llm/kernels/logitsBitmask.h

🧬 Code graph analysis (1)

tensorrt_llm/_torch/pyexecutor/py_executor_creator.py (4)

tensorrt_llm/_torch/models/modeling_utils.py (1)

• vocab_size_padded (531-532)

tensorrt_llm/mapping.py (2)

• rank (199-200)
• rank (203-210)

tensorrt_llm/_torch/distributed/communicator.py (2)

• rank (89-90)
• rank (551-552)

tensorrt_llm/_torch/device_mesh.py (1)

• rank (36-37)

⏰ 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 (13)

tensorrt_llm/_torch/pyexecutor/py_executor_creator.py (1)

505-523: LGTM!

The rank parameter is correctly propagated to the guided decoder constructors via kwargs. The mapping.rank value is properly initialized earlier in the function (line 305) and will enable rank-aware bitmask slicing for sharded vocabulary scenarios.

cpp/tensorrt_llm/thop/logitsBitmaskOp.cpp (4)

26-29: LGTM!

Good addition of the constant in an anonymous namespace, consistent with the kernel implementation.

---

42-46: LGTM!

The validation is correctly conditional on d2t not being present. When d2t is provided, the bitmask may be a slice (for sharded vocab scenarios), so the size check would be inappropriate.

---

63-63: LGTM!

Switching from reinterpret_cast to static_cast for void* to T* conversions aligns with C++ best practices as per coding guidelines.

---

83-99: LGTM!

Using bitmask.stride(0) instead of a computed size correctly handles cases where the bitmask tensor is a slice (non-contiguous in dimension 1), which is essential for the sharded vocabulary support. The static_cast usage for data_ptr() is also correct.

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

34-35: LGTM!

The parameter rename from bitmaskSize to bitmaskStride accurately reflects the semantic change and is consistent with the implementation in logitsBitmask.cu.

tensorrt_llm/_torch/pyexecutor/guided_decoder.py (3)

144-154: LGTM!

The rank parameter is properly added with a sensible default of 0 and stored as an instance attribute. This enables rank-aware bitmask slicing for distributed scenarios.

---

310-331: LGTM!

The bitmask slicing logic correctly handles vocabulary sharding:

1. tp_size is derived from the ratio of full vocabulary to logits dimension
2. tp_rank = self.rank % tp_size correctly maps the global rank to the TP-local rank
3. The bitmask slice [bitmask_start:bitmask_end] extracts the appropriate portion for this rank
4. When tp_size=1 (no sharding), the full bitmask is used as expected
5. The d2t tensor is also correctly sliced when present

The assertions provide good validation for misconfiguration detection.

---

439-451: LGTM!

CapturableGuidedDecoder correctly accepts the rank parameter and propagates it to the parent class constructor using keyword arguments, maintaining consistency with the base class interface.

cpp/tensorrt_llm/kernels/logitsBitmask.cu (4)

67-68: LGTM!

The bitmaskSize parameter was correctly removed from logitsBitmaskKernel since this non-contiguous version accesses bitmask via per-batch pointers and never used the size parameter.

---

184-201: LGTM!

The contiguousLogitsBitmaskKernel correctly uses bitmaskStride for row-major indexing. The calculation bitmask + batchIdx * bitmaskStride + blockOffset / kBitsPerMaskElement properly computes the row start plus column offset within the bitmask tensor.

---

224-230: LGTM!

The d2t branch correctly uses bitmaskStride for row indexing: bitmask[batchIdx * bitmaskStride + d2tOffset / kBitsPerMaskElement]. This maintains consistency with the non-d2t path.

---

294-302: LGTM!

Excellent documentation clarifying the two scenarios where bitmaskStride differs from the computed bitmask size:

1. EAGLE3-style "pruned" logits with d2t indirection
2. Vocabulary sharding across ranks

This makes the API contract clear for callers.

_{✏️ Tip: You can disable this entire section by setting review_details to false in your review settings.}

---

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

[tensorrt-cicd]

PR_Github #32225 [ run ] triggered by Bot. Commit: a6e3ff4

[syuoni]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #32242 [ run ] triggered by Bot. Commit: a6e3ff4

[syuoni]

/bot kill

[tensorrt-cicd]

PR_Github #32288 [ kill ] triggered by Bot. Commit: a6e3ff4

[tensorrt-cicd]

PR_Github #32288 [ kill ] completed with state SUCCESS. Commit: a6e3ff4
Successfully killed previous jobs for commit a6e3ff4

[syuoni]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #32294 [ run ] triggered by Bot. Commit: a6e3ff4

[tensorrt-cicd]

PR_Github #32294 [ run ] completed with state SUCCESS. Commit: a6e3ff4
/LLM/release-1.2/L0_MergeRequest_PR pipeline #21 completed with status: 'SUCCESS'