fix: To validate AI decoder RPC slot capacity to avoid gateway output overwrite - #764
Open
kaiqiy-nv wants to merge 1 commit into
Open
fix: To validate AI decoder RPC slot capacity to avoid gateway output overwrite#764kaiqiy-nv wants to merge 1 commit into
kaiqiy-nv wants to merge 1 commit into
Conversation
Signed-off-by: Kaiqi Yan <kaiqiy@nvidia.com>
kaiqiy-nv
force-pushed
the
kaiqiy/fix-ai-decoder-gateway
branch
from
August 5, 2026 07:32
dca731c to
1bf3014
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR tries to fix [B] 6557191.
It also adds 2 test cases to expose the bug and verify the fix.
Thanks for looking at this PR.
Note: The new tests are not included in the CI. This is to keep the existing setting. If we want to enable it or other tests that are currently not in the CI, we can do it in a separate PR.
In this suggested fix:
We update the experimental AI decoder gateway to track the total RPC slot size, validate both input and output frames before CUDA graph capture, and pass payload capacity into the gateway kernels for defensive device-side bounds.
In this case, gateway output writes beyond the mapped RPC slot can be avoided when a valid model produces a larger output payload than the request payload.
Added/Modified/Deleted test cases:
RejectsOversizedGatewayOutput:We build a valid TensorRT engine whose output has one more float than its input, use the default input-sized RPC slot assumption, and check that graph capture rejects the oversized response so the bug can be exposed.
WritesGatewayOutputWithinSlot:We use the same asymmetric TensorRT engine with an explicit larger RPC slot, run the gateway graph, compare the response payload, and check an adjacent-slot canary so valid larger-output usage remains supported without corruption.
Tests passed on our local machine as:
and