[None][feat] Add unit tests and revision in block_level kernel for invalid input

[ChristinaZ]

Summary by CodeRabbit

• Bug Fixes

  • Fixed handling of invalid top-k indices in routing operations to prevent memory corruption.
  • Added safeguards to properly validate and record invalid top-k entries instead of corrupting memory.

• Tests

  • Extended test coverage for invalid top-k input scenarios across routing kernel tests.
  • Added new test cases to exercise invalid top-k input handling paths.

Description

For the MoE TRTLLM backend:

1. Add one revision for the block-level kernel for invalid input
2. Add unit tests in the C++ routing unit tests
3. Add unit tests in test_moe.py to test the invalid input e2e

Test Coverage

pytest -vv -s tests/unittest/_torch/thop/parallel/test_moe.py::test_moe_mxe2m1_weights -k "test_invalid_topk_input and use_topk_as_input"

cd cpp/build
make -j$(nproc) google-tests
./tests/unit_tests/kernels/routingKernelsTest

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

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

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

This PR adds validation and defensive handling for invalid top-k expert indices (-1) across the MoE routing infrastructure. It introduces a new hasInvalidTopKInput flag to test parameters, guards memory operations against invalid indices in kernel and host code, and includes test cases to verify correct behavior when top-k results contain invalid entries.

Changes

Cohort / File(s)	Summary
Kernel Guard Against Invalid Expert Indices cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingRenormalize.cu	Added bounds check when filling top-k expert indices; records -1 in mapping for invalid entries instead of writing to shared memory to prevent memory corruption.
Test Infrastructure Base cpp/tests/unit_tests/kernels/routing/routingTest.h	Added hasInvalidTopKInput boolean field to RoutingKernelTestParam; extended constructor signature with validation that rejects hasInvalidTopKInput=true when useTopKAsInput=false; updated toString() output.
Host Test Code cpp/tests/unit_tests/kernels/routing/routingTest.cpp	Replaced auto-typed variables with explicit int32_t and bool types; added guard preventing writes to tokenToIdxInExpertHostPtr when index < 0.
Renormalization Routing Tests cpp/tests/unit_tests/kernels/routing/routingRenormalizeTest.cpp	Updated test parametrization with hasInvalidTopKInput flag; when enabled, forces first half of top-k indices to -1 before packing; added test cases exercising invalid top-k scenarios (BlockLevelParallelization, ClusterLevel, and Top4 variants).
DeepSeek Routing Tests cpp/tests/unit_tests/kernels/routing/routingDeepSeekTest.cpp	Added hasInvalidTopKInput parameter to all test parametrizations with inline comment annotations.
Python MoE Tests tests/unittest/_torch/thop/parallel/test_moe.py	Added test_invalid_topk_input flag; when enabled, appends -1 column to topKIndices and topKLogits, increments top_k by 1 before kernel invocation to test handling of invalid top-k entries.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

• Kernel guard logic: Verify that the bounds check on expert index correctly prevents invalid memory writes and properly records invalid entries in the mapping structure.
• Parameter validation: Ensure the new hasInvalidTopKInput validation rule in routingTest.h constructor (rejecting when useTopKAsInput=false) is correct and handles all code paths.
• Test case generation: Confirm that the logic forcing first half of top-k indices to -1 correctly simulates the invalid input scenario across all test variants.
• Type safety changes in routingTest.cpp: Verify that explicit types and the guard condition index >= 0 don't inadvertently skip valid processing.
• Python test integration: Check that the appended -1 column and incremented top_k value correctly exercise the invalid top-k handling without breaking existing test expectations.

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 7.32% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Title Check	✅ Passed	The PR title "[None][feat] Add unit tests and revision in block_level kernel for invalid input" directly corresponds to the main changes in the changeset. The raw_summary shows modifications to the block-level routing kernel (RoutingRenormalize.cu) to handle invalid top-k inputs, along with new unit tests added across multiple C++ test files and a Python test file. The title is specific, clear, and follows the required format [ticket][type] Summary. A teammate scanning the repository history would understand that this PR adds tests and fixes for invalid input handling in the block-level MoE kernel.
Description Check	✅ Passed	The PR description is mostly complete and follows the required template structure. It includes all required sections: a clear description of the changes (MoE backend revision for invalid input handling and unit tests), a specific Test Coverage section with concrete test commands for both Python and C++ tests, and a completed PR Checklist. The description provides sufficient context about what is being changed and how to verify the changes, avoiding vague or generic language.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

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

---

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

📜 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 6b9b73e and e192735.

📒 Files selected for processing (6)

• cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingRenormalize.cu (1 hunks)
• cpp/tests/unit_tests/kernels/routing/routingDeepSeekTest.cpp (15 hunks)
• cpp/tests/unit_tests/kernels/routing/routingRenormalizeTest.cpp (13 hunks)
• cpp/tests/unit_tests/kernels/routing/routingTest.cpp (1 hunks)
• cpp/tests/unit_tests/kernels/routing/routingTest.h (5 hunks)
• tests/unittest/_torch/thop/parallel/test_moe.py (3 hunks)

🧰 Additional context used

📓 Path-based instructions (8)

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

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

**/*.{h,hpp,hh,hxx,cpp,cxx,cc,cu,cuh}: Namespace closing braces must include a trailing comment with the namespace name (e.g., '} // namespace foo').
Prefer const or constexpr variables over #define for constants.
Declare variables that are not modified after initialization as const.
Avoid magic literals in code; except for 0, nullptr, true, false. Use named constants for comparisons and logic.
Use Allman brace style for formatting.
Place the semicolon of an empty for/while loop on a new line.
Bodies of switch/while/do-while/for must be compound statements (brace-delimited), and if/else must always be followed by brace-delimited statements.
Type names (e.g., classes) must be CamelCase starting with an uppercase letter (e.g., FooBar).
Local variables, methods, and namespaces use lowerCamelCase (e.g., localFooBar).
Non-magic-number global variables that are non-static and not in an anonymous namespace must be lowerCamelCase prefixed with 'g' (e.g., gDontUseGlobalFoos).
Non-magic-number globals that are static or in an anonymous namespace use lowerCamelCase prefixed with 's' (e.g., sMutableStaticGlobal).
Locally visible static variables use lowerCamelCase with 's' prefix (e.g., static std::once_flag sFlag).
Private/protected member variables use 'm' prefix with CamelCase (e.g., mNbFooValues). Public members may omit, but 'm' is encouraged for clarity.
Constants (enums, global constants, static constants, and function-scope magic/literal constants) use uppercase SNAKE_CASE with 'k' prefix (e.g., kDIGIT_NUM).
Function-scope constants that are not magic numbers or literals are named like non-constant variables (e.g., bool const pass = a && b).
If macros are necessary, name them in UPPER_SNAKE_CASE (e.g., FOO_VERSION) and prefer constants over #define.
Use LLVM clang-format; wrap lines at a maximum of 120 columns; use '// clang-format off/on' sparingly with justification.
Use smart pointers for heap allocations; prefer unique_ptr for sole ownership, shared_ptr for shared...

Files:

• cpp/tests/unit_tests/kernels/routing/routingTest.h
• cpp/tests/unit_tests/kernels/routing/routingTest.cpp
• cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingRenormalize.cu
• cpp/tests/unit_tests/kernels/routing/routingDeepSeekTest.cpp
• cpp/tests/unit_tests/kernels/routing/routingRenormalizeTest.cpp

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

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

C++ filenames should be lowerCamelCase (first letter lowercase) and must be case-insensitive unique within a compilation target.

Files:

• cpp/tests/unit_tests/kernels/routing/routingTest.h
• cpp/tests/unit_tests/kernels/routing/routingTest.cpp
• cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingRenormalize.cu
• cpp/tests/unit_tests/kernels/routing/routingDeepSeekTest.cpp
• cpp/tests/unit_tests/kernels/routing/routingRenormalizeTest.cpp

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

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

Use only spaces, no tabs; indent with 4 spaces.

Files:

• cpp/tests/unit_tests/kernels/routing/routingTest.h
• cpp/tests/unit_tests/kernels/routing/routingTest.cpp
• tests/unittest/_torch/thop/parallel/test_moe.py
• cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingRenormalize.cu
• cpp/tests/unit_tests/kernels/routing/routingDeepSeekTest.cpp
• cpp/tests/unit_tests/kernels/routing/routingRenormalizeTest.cpp

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

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

Document new class interfaces and function prototypes with Doxygen; use //! for single-line and //!< for members.

Files:

• cpp/tests/unit_tests/kernels/routing/routingTest.h

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

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

**/*.{h,hpp,hh,hxx,cpp,cxx,cc}: Prefer anonymous namespaces over 'static' for internal linkage of functions.
All templates (class/function/member/static) must be instantiated at least once; non-POD classes should have private data members.

Files:

• cpp/tests/unit_tests/kernels/routing/routingTest.h
• cpp/tests/unit_tests/kernels/routing/routingTest.cpp
• cpp/tests/unit_tests/kernels/routing/routingDeepSeekTest.cpp
• cpp/tests/unit_tests/kernels/routing/routingRenormalizeTest.cpp

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

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

Use include guards named 'TRTLLM_<FILE_NAME_IN_CAPS_WITH_UNDERSCORES>_H' (no leading or trailing underscore; directory names excluded).

Files:

• cpp/tests/unit_tests/kernels/routing/routingTest.h

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

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

Prepend the NVIDIA Apache-2.0 copyright header with current year to the top of all source files (e.g., .cpp, .h, .cu, .py).

Files:

• cpp/tests/unit_tests/kernels/routing/routingTest.h
• cpp/tests/unit_tests/kernels/routing/routingTest.cpp
• tests/unittest/_torch/thop/parallel/test_moe.py
• cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingRenormalize.cu
• cpp/tests/unit_tests/kernels/routing/routingDeepSeekTest.cpp
• cpp/tests/unit_tests/kernels/routing/routingRenormalizeTest.cpp

**/*.py

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

**/*.py: Python code must target Python 3.8+.
Indent Python code with 4 spaces; do not use tabs.
Maintain module namespace when importing; prefer 'from package.subpackage import foo' then 'foo.SomeClass()' instead of importing the class directly.
Python filenames should be snake_case (e.g., some_file.py).
Python classes use PascalCase names.
Functions and methods use snake_case names.
Local variables use snake_case; prefix 'k' for variables that start with a number (e.g., k_99th_percentile).
Global variables use upper SNAKE_CASE prefixed with 'G' (e.g., G_MY_GLOBAL).
Constants use upper SNAKE_CASE (e.g., MY_CONSTANT).
Avoid shadowing variables from an outer scope.
Initialize all externally visible members of a class in the constructor.
Prefer docstrings for interfaces that may be used outside a file; comments for in-function or file-local interfaces.
Use Google-style docstrings for classes and functions (Sphinx-parsable).
Document attributes and variables inline so they render under the class/function docstring.
Avoid reflection when a simpler, explicit approach suffices (e.g., avoid dict(**locals()) patterns).
In try/except, catch the most specific exceptions possible.
For duck-typing try/except, keep the try body minimal and use else for the main logic.

Files:

• tests/unittest/_torch/thop/parallel/test_moe.py

🧠 Learnings (1)

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

Learnt from: jhaotingc
PR: NVIDIA/TensorRT-LLM#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/kernels/trtllmGenKernels/blockScaleMoe/RoutingRenormalize.cu

🧬 Code graph analysis (1)

cpp/tests/unit_tests/kernels/routing/routingRenormalizeTest.cpp (1)

cpp/tests/unit_tests/kernels/routing/routingDeepSeekTest.cpp (8)

• param (47-152)
• param (47-47)
• param (154-165)
• param (154-154)
• param (167-176)
• param (167-167)
• param (203-209)
• param (203-204)

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

cpp/tests/unit_tests/kernels/routing/routingTest.cpp (1)

162-169: LGTM: Defensive handling for invalid top-k indices.

The explicit type declarations improve code clarity and align with coding guidelines. The guard at lines 165-168 correctly prevents writing to tokenToIdxInExpertHostPtr when the index is invalid (-1), which is essential for the invalid top-k input handling introduced in this PR.

cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingRenormalize.cu (1)

133-143: LGTM: Proper handling of invalid top-k indices.

The conditional guard correctly prevents writing to smemKIdx when expertIdx is -1 (invalid), and appropriately records the invalid entry in mPtrExpandedIdxToPermutedIdx. This aligns with the test scenarios introduced in the PR and prevents potential memory corruption from invalid offset calculations.

cpp/tests/unit_tests/kernels/routing/routingDeepSeekTest.cpp (1)

220-220: LGTM: Consistent parameter propagation.

All test cases have been updated to include the hasInvalidTopKInput parameter. The tests that enable this flag (lines 264, 319, 374) correctly pair it with useTopKAsInput=true, which aligns with the validation added in routingTest.h that requires useTopKAsInput when hasInvalidTopKInput is enabled.

Also applies to: 231-231, 242-242, 253-253, 264-264, 275-275, 286-286, 297-297, 308-308, 319-319, 330-330, 341-341, 352-352, 363-363, 374-374

tests/unittest/_torch/thop/parallel/test_moe.py (2)

2042-2052: LGTM: Correct augmentation of top-k inputs with invalid entries.

The code properly augments topKIndices and topKLogits with an extra column of -1 values to simulate invalid top-k entries. This aligns with the C++ test implementation and enables end-to-end testing of invalid top-k input handling.

---

2198-2199: LGTM: Correct adjustment of top_k for extended input.

The increment of top_k by 1 correctly accounts for the extra invalid column added to the top-k inputs when test_invalid_topk_input is true. This ensures the kernel processes the full extended top-k set including the invalid entry.

cpp/tests/unit_tests/kernels/routing/routingTest.h (2)

248-248: LGTM: Proper field addition and initialization.

The hasInvalidTopKInput field is correctly added to the struct with proper default initialization, constructor parameter, and member initialization. This follows C++ best practices and maintains consistency with other test parameters.

Also applies to: 274-275, 286-286

---

325-328: LGTM: Appropriate validation and debug support.

The validation at lines 325-328 correctly enforces that hasInvalidTopKInput can only be enabled when useTopKAsInput is true, which makes logical sense since you can only inject invalid indices when providing top-k inputs directly. The toString update provides helpful debugging information.

Also applies to: 350-353

cpp/tests/unit_tests/kernels/routing/routingRenormalizeTest.cpp (2)

130-135: LGTM: Test data generation for invalid top-k scenarios.

The code correctly injects invalid expert indices (-1) for the first topK/2 + 1 entries when hasInvalidTopKInput is enabled. This creates test coverage for handling invalid top-k inputs. The +1 ensures at least one invalid entry even for small top-k values (e.g., topK=2 would have 2 invalid entries).

---

207-386: LGTM: Comprehensive test coverage for invalid top-k inputs.

The parameter updates maintain backward compatibility by explicitly setting hasInvalidTopKInput=false for existing tests. The new test cases (BlockLevelParallelizationWithInvalidTopKInput, ClusterLevelParallelizationWithInvalidTopKInput, etc.) provide good coverage across different parallelization levels and top-k values. All new tests correctly pair useTopKAsInput=true with hasInvalidTopKInput=true, satisfying the validation constraint.

[ChristinaZ]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #22762 [ run ] triggered by Bot. Commit: e192735

[tensorrt-cicd]

PR_Github #22762 [ run ] completed with state SUCCESS. Commit: e192735
/LLM/main/L0_MergeRequest_PR pipeline #17164 completed with status: 'FAILURE'

[ChristinaZ]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #22864 [ run ] triggered by Bot. Commit: 209554b

[tensorrt-cicd]

PR_Github #22864 [ run ] completed with state SUCCESS. Commit: 209554b
/LLM/main/L0_MergeRequest_PR pipeline #17242 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.