[https://nvbugs/6215690][fix] AutoDeploy: FlashInfer 128-byte alignment for Mamba inputs (also addresses nvbugs/6162114) by galagam · Pull Request #14535 · NVIDIA/TensorRT-LLM

galagam · 2026-05-25T16:27:51Z

Description

The decode path requires contiguous aligned to 128 bytes. Add _fi_align() before passing x/B/C to FlashInfer SSM kernel in the decode path.

data.contiguous() is the real fix: with TP >= 8, the hidden-state tensor is a strided view (non-contiguous), which FlashInfer rejects.
This regression was invisible at world_size=4 where the same sharding happened to produce a contiguous layout.

clone() is purely defensive: for the current Mamba2 configs, row sizes
(num_heads * head_dim * dtype_size and n_groups * ssm_state_size * dtype_size)
are always multiples of 128 bytes, so a contiguous slice can never be misaligned.
The clone() path guards against future configs where that assumption breaks.

The extend path does not require contiguous memory and remains as-is.

Test Coverage

The unwaived Ultra test:
-accuracy/test_llm_api_autodeploy.py::TestNemotronUltraV3::test_accuracy[nvfp4-8]

PR Checklist

Please review the following before submitting your PR:

PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

Summary by CodeRabbit

Bug Fixes
- Improved memory alignment handling for tensor operations to enhance performance and stability.
Tests
- Removed test waivers as accuracy tests now pass reliably.

coderabbitai · 2026-05-25T16:32:24Z

📝 Walkthrough

Walkthrough

This PR adds 128-byte memory alignment enforcement for tensors passed to FlashInfer kernels in the Mamba SSM backend. A new _fi_align helper ensures tensors are properly aligned by cloning when necessary, then applies this alignment to input tensors in both extend and decode execution paths before kernel invocation. Test waivers are removed, confirming the alignment fix resolves previously failing accuracy tests.

Changes

FlashInfer Mamba Memory Alignment

Layer / File(s)	Summary
Alignment helper implementation `tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py`	`_fi_align` helper function enforces 128-byte alignment for tensors by returning contiguous tensors when already aligned and cloning contiguous copies when realignment is required.
Alignment applied to extend and decode paths `tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py`	Input tensors (`x_extend`, `B_extend`, `C_extend`) in the extend path and (`x_decode`, `B_decode`, `C_decode`) in the decode path are aligned before calling `_flashinfer_ssm_update`.
Test waivers removed `tests/integration/test_lists/waives.txt`	Waivers for NemotronSuperV3 small functional tests (bf16, fp8) and NemotronUltraV3 nvfp4 accuracy test are removed, indicating the alignment fix resolves previously failing test cases.

Possibly related PRs

NVIDIA/TensorRT-LLM#13711: Modifies the same Mamba FlashInfer backend file, adding extend-path integration for the FlashInfer SSM handler.
NVIDIA/TensorRT-LLM#14514: Modifies the same test waivers file, with potential overlap on test waiver entries.

Suggested reviewers

QiJune

🎯 2 (Simple) | ⏱️ ~12 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title directly describes the main change: adding 128-byte alignment for Mamba inputs in FlashInfer, which aligns with the core implementation in flashinfer_backend_mamba.py and the test unwaiving changes.
Description check	✅ Passed	The PR description clearly explains the issue, solution, test coverage, and confirms the checklist requirements.

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

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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

galagam · 2026-05-25T17:06:24Z

/bot run --stage-list "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-Post-Merge-1,DGX_B200-8_GPUs-AutoDeploy-Post-Merge-1"

tensorrt-cicd · 2026-05-25T17:12:01Z

PR_Github #50229 [ run ] triggered by Bot. Commit: 1c43bea Link to invocation

tensorrt-cicd · 2026-05-25T19:32:52Z

PR_Github #50229 [ run ] completed with state FAILURE. Commit: 1c43bea
/LLM/main/L0_MergeRequest_PR pipeline #39765 (Partly Tested) completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

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

CI Agent Failure Analysis

Link to invocation

galagam · 2026-05-26T06:40:09Z

/bot run --stage-list "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-Post-Merge-1,DGX_B200-8_GPUs-AutoDeploy-Post-Merge-1"

tensorrt-cicd · 2026-05-26T06:45:52Z

PR_Github #50304 [ run ] triggered by Bot. Commit: 9356eeb Link to invocation

tensorrt-cicd · 2026-05-26T09:13:52Z

PR_Github #50304 [ run ] completed with state SUCCESS. Commit: 9356eeb
/LLM/main/L0_MergeRequest_PR pipeline #39834 (Partly Tested) completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

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

CI Agent Failure Analysis

Link to invocation

The decode path requires contiguous aligned to 128 bytes. Add _fi_align() before passing x/B/C to FlashInfer SSM kernel in the decode path. data.contiguous() is the real fix: with TP >= 8, the hidden-state tensor is a strided view (non-contiguous), which FlashInfer rejects. This regression was invisible at world_size=4 where the same sharding happened to produce a contiguous layout. clone() is purely defensive: for the current Mamba2 configs, row sizes (num_heads * head_dim * dtype_size and n_groups * ssm_state_size * dtype_size) are always multiples of 128 bytes, so a contiguous slice can never be misaligned. The clone() path guards against future configs where that assumption breaks. The extend path does not require contiguous memory and remains as-is. Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com>

Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com>

galagam · 2026-05-26T19:35:08Z

/bot run --stage-list "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-Post-Merge-1,DGX_B200-8_GPUs-AutoDeploy-Post-Merge-1"

tensorrt-cicd · 2026-05-26T19:41:14Z

PR_Github #50382 [ run ] triggered by Bot. Commit: 1292c5a Link to invocation

tensorrt-cicd · 2026-05-26T22:09:31Z

PR_Github #50382 [ run ] completed with state SUCCESS. Commit: 1292c5a
/LLM/main/L0_MergeRequest_PR pipeline #39909 (Partly Tested) completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

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

CI Agent Failure Analysis

Link to invocation

Signed-off-by: Gal Hubara-Agam <96368689+galagam@users.noreply.github.com>

galagam · 2026-05-27T06:01:22Z

/bot run --stage-list "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-Post-Merge-1,DGX_B200-8_GPUs-AutoDeploy-Post-Merge-1"

tensorrt-cicd · 2026-05-27T06:07:44Z

PR_Github #50475 [ run ] triggered by Bot. Commit: bf4f68a Link to invocation

tensorrt-cicd · 2026-05-27T08:41:24Z

PR_Github #50475 [ run ] completed with state SUCCESS. Commit: bf4f68a
/LLM/main/L0_MergeRequest_PR pipeline #39985 (Partly Tested) completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

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

CI Agent Failure Analysis

Link to invocation

galagam · 2026-05-27T12:42:30Z

/bot run

tensorrt-cicd · 2026-05-27T12:48:24Z

PR_Github #50537 [ run ] triggered by Bot. Commit: bf4f68a Link to invocation

tensorrt-cicd · 2026-05-27T15:26:50Z

PR_Github #50537 [ run ] completed with state SUCCESS. Commit: bf4f68a
/LLM/main/L0_MergeRequest_PR pipeline #40042 completed with status: 'SUCCESS'

CI Report

Link to invocation

galagam · 2026-05-27T16:02:04Z

/bot run --stage-list "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-Post-Merge-1,DGX_B200-8_GPUs-AutoDeploy-Post-Merge-1"

tensorrt-cicd · 2026-05-27T16:08:31Z

PR_Github #50573 [ run ] triggered by Bot. Commit: bf4f68a Link to invocation

tensorrt-cicd · 2026-05-27T17:11:18Z

PR_Github #50573 [ run ] completed with state FAILURE. Commit: bf4f68a
/LLM/main/L0_MergeRequest_PR pipeline #40074 (Partly Tested) completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

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

CI Agent Failure Analysis

Link to invocation

galagam · 2026-05-27T17:17:31Z

/bot run --stage-list "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-Post-Merge-1"

galagam · 2026-05-27T17:22:37Z

/bot help

github-actions · 2026-05-27T17:22:51Z

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) --high-priority]

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. Supports wildcard * for pattern matching (e.g., "*PerfSanity*" matches all stages containing PerfSanity). Examples: "A10-PyTorch-1, xxx", "PerfSanity". 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. Supports wildcard * for pattern matching. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx", --extra-stage "Post-Merge".

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

--high-priority (OPTIONAL) : Run the pipeline with high priority. This option is restricted to authorized users only and will route the job to a high-priority queue.

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.

tensorrt-cicd · 2026-05-27T17:23:45Z

PR_Github #50586 [ run ] triggered by Bot. Commit: bf4f68a Link to invocation

tensorrt-cicd · 2026-05-27T17:53:50Z

PR_Github #50586 [ run ] completed with state SUCCESS. Commit: bf4f68a
/LLM/main/L0_MergeRequest_PR pipeline #40086 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

galagam · 2026-05-27T18:05:18Z

/bot run reuse-pipeline

tensorrt-cicd · 2026-05-27T18:12:03Z

PR_Github #50594 Bot args parsing error: usage: /bot [-h]
{run,kill,skip,submit,reviewers,reuse-pipeline,reuse-review} ...
/bot: error: unrecognized arguments: reuse-pipeline

Link to invocation

galagam · 2026-05-27T18:16:43Z

/bot reuse-pipeline

tensorrt-cicd · 2026-05-27T18:22:38Z

PR_Github #50595 [ reuse-pipeline ] triggered by Bot. Commit: bf4f68a Link to invocation

tensorrt-cicd · 2026-05-27T18:34:15Z

PR_Github #50595 [ reuse-pipeline ] completed with state SUCCESS. Commit: bf4f68a
Reusing PR_Github #50586 (Partly Tested) for commit bf4f68a

Link to invocation

galagam · 2026-05-27T18:45:49Z

Full bot run succeeded: #14535 (comment)

This failure is from a known regression in a post merge test #14535 (comment)

…nt for Mamba inputs (also addresses nvbugs/6162114) (NVIDIA#14535) Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> Signed-off-by: Gal Hubara-Agam <96368689+galagam@users.noreply.github.com>

galagam requested a review from a team as a code owner May 25, 2026 16:27

galagam requested review from greg-kwasniewski1 and tcherckez-nvidia May 25, 2026 16:27

github-actions Bot assigned galagam May 25, 2026

galagam changed the title ~~[http://nvbugs/6162114][fix] AutoDeploy: FlashInfer 128-byte alignment for Mamba inputs~~ [https://nvbugs/6162114][fix] AutoDeploy: FlashInfer 128-byte alignment for Mamba inputs May 25, 2026

galagam changed the title ~~[https://nvbugs/6162114][fix] AutoDeploy: FlashInfer 128-byte alignment for Mamba inputs~~ [https://nvbugs/6215690][fix] AutoDeploy: FlashInfer 128-byte alignment for Mamba inputs May 25, 2026

galagam changed the title ~~[https://nvbugs/6215690][fix] AutoDeploy: FlashInfer 128-byte alignment for Mamba inputs~~ [nvbugs/6215690][nvbugs/6162114][fix] AutoDeploy: FlashInfer 128-byte alignment for Mamba inputs May 25, 2026

galagam changed the title ~~[nvbugs/6215690][nvbugs/6162114][fix] AutoDeploy: FlashInfer 128-byte alignment for Mamba inputs~~ [https://nvbugs/6215690][https://nvbugs/6162114][fix] AutoDeploy: FlashInfer 128-byte alignment for Mamba inputs May 25, 2026

galagam changed the title ~~[https://nvbugs/6215690][https://nvbugs/6162114][fix] AutoDeploy: FlashInfer 128-byte alignment for Mamba inputs~~ [https://nvbugs/6215690][fix] AutoDeploy: FlashInfer 128-byte alignment for Mamba inputs (also addresses nvbugs/6162114) May 25, 2026

galagam marked this pull request as draft May 26, 2026 02:22

galagam removed the request for review from greg-kwasniewski1 May 26, 2026 02:23

galagam force-pushed the gagam/nvbug-6162114 branch 2 times, most recently from 04df655 to 9356eeb Compare May 26, 2026 06:38

tcherckez-nvidia approved these changes May 26, 2026

View reviewed changes

galagam marked this pull request as ready for review May 26, 2026 15:54

galagam force-pushed the gagam/nvbug-6162114 branch from 9356eeb to e1f06dd Compare May 26, 2026 18:52

galagam added 2 commits May 26, 2026 12:11

also unwaive TestNemotronSuperV3::test_functional_small

1e79f76

Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com>

galagam force-pushed the gagam/nvbug-6162114 branch from aff7b86 to 1e79f76 Compare May 26, 2026 19:14

waive https://nvbugs/6221483

1292c5a

Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com>

Merge branch 'main' into gagam/nvbug-6162114

bf4f68a

Signed-off-by: Gal Hubara-Agam <96368689+galagam@users.noreply.github.com>

galagam merged commit 86c1b0c into NVIDIA:main May 27, 2026
7 checks passed

Uh oh!

Conversation

galagam commented May 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Test Coverage

PR Checklist

GitHub Bot Help

Summary by CodeRabbit

Uh oh!

coderabbitai Bot commented May 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Possibly related PRs

Suggested reviewers

❌ Failed checks (1 warning)

Uh oh!

galagam commented May 25, 2026

Uh oh!

tensorrt-cicd commented May 25, 2026

Uh oh!

tensorrt-cicd commented May 25, 2026

Uh oh!

galagam commented May 26, 2026

Uh oh!

tensorrt-cicd commented May 26, 2026

Uh oh!

tensorrt-cicd commented May 26, 2026

Uh oh!

galagam commented May 26, 2026

Uh oh!

tensorrt-cicd commented May 26, 2026

Uh oh!

tensorrt-cicd commented May 26, 2026

Uh oh!

galagam commented May 27, 2026

Uh oh!

tensorrt-cicd commented May 27, 2026

Uh oh!

tensorrt-cicd commented May 27, 2026

Uh oh!

galagam commented May 27, 2026

Uh oh!

tensorrt-cicd commented May 27, 2026

Uh oh!

tensorrt-cicd commented May 27, 2026

Uh oh!

galagam commented May 27, 2026

Uh oh!

tensorrt-cicd commented May 27, 2026

Uh oh!

tensorrt-cicd commented May 27, 2026

Uh oh!

galagam commented May 27, 2026

Uh oh!

galagam commented May 27, 2026

Uh oh!

github-actions Bot commented May 27, 2026

GitHub Bot Help

kill

skip

reuse-pipeline

Uh oh!

tensorrt-cicd commented May 27, 2026

Uh oh!

tensorrt-cicd commented May 27, 2026

Uh oh!

galagam commented May 27, 2026

Uh oh!

tensorrt-cicd commented May 27, 2026

Uh oh!

galagam commented May 27, 2026

Uh oh!

tensorrt-cicd commented May 27, 2026

Uh oh!

tensorrt-cicd commented May 27, 2026

Uh oh!

galagam commented May 27, 2026

Uh oh!

galagam commented May 25, 2026 •

edited

Loading

coderabbitai Bot commented May 25, 2026 •

edited

Loading