Skip to content

fix: To validate AI decoder RPC slot capacity to avoid gateway output overwrite - #764

Open
kaiqiy-nv wants to merge 1 commit into
NVIDIA:mainfrom
kaiqiy-nv:kaiqiy/fix-ai-decoder-gateway
Open

fix: To validate AI decoder RPC slot capacity to avoid gateway output overwrite#764
kaiqiy-nv wants to merge 1 commit into
NVIDIA:mainfrom
kaiqiy-nv:kaiqiy/fix-ai-decoder-gateway

Conversation

@kaiqiy-nv

Copy link
Copy Markdown
Collaborator

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:

./libs/qec/unittests/test_ai_decoder_quantized_onnx --gtest_filter=AiDecoderQuantizedOnnxSmokeTest.RejectsOversizedGatewayOutput:AiDecoderQuantizedOnnxSmokeTest.WritesGatewayOutputWithinSlot --gtest_color=no

Running main() from /CUDA-QX/build/public/_deps/googletest-src/googletest/src/gtest_main.cc
Note: Google Test filter = AiDecoderQuantizedOnnxSmokeTest.RejectsOversizedGatewayOutput:AiDecoderQuantizedOnnxSmokeTest.WritesGatewayOutputWithinSlot
[==========] Running 2 tests from 1 test suite.
[----------] Global test environment set-up.
[----------] 2 tests from AiDecoderQuantizedOnnxSmokeTest
[ RUN      ] AiDecoderQuantizedOnnxSmokeTest.RejectsOversizedGatewayOutput
[TensorRT] Binding 0: "input" INPUT, dtype=0, elem_bits=32, volume=8, 32 bytes
[TensorRT] Binding 1: "output" OUTPUT, dtype=0, elem_bits=32, volume=9, 36 bytes
[       OK ] AiDecoderQuantizedOnnxSmokeTest.RejectsOversizedGatewayOutput (768 ms)
[ RUN      ] AiDecoderQuantizedOnnxSmokeTest.WritesGatewayOutputWithinSlot
[TensorRT] Binding 0: "input" INPUT, dtype=0, elem_bits=32, volume=8, 32 bytes
[TensorRT] Binding 1: "output" OUTPUT, dtype=0, elem_bits=32, volume=9, 36 bytes
[       OK ] AiDecoderQuantizedOnnxSmokeTest.WritesGatewayOutputWithinSlot (200 ms)
[----------] 2 tests from AiDecoderQuantizedOnnxSmokeTest (969 ms total)

[----------] Global test environment tear-down
[==========] 2 tests from 1 test suite ran. (969 ms total)
[  PASSED  ] 2 tests.

and

ctest --output-on-failure -R test_ai_decoder_quantized_onnx

Test project /CUDA-QX/build/public
    Start 326: test_ai_decoder_quantized_onnx.AiDecoderQuantizedOnnxSmokeTest.Int8QdqRunsWithExpectedNumerics
1/4 Test #326: test_ai_decoder_quantized_onnx.AiDecoderQuantizedOnnxSmokeTest.Int8QdqRunsWithExpectedNumerics ....................   Passed    1.18 sec
    Start 327: test_ai_decoder_quantized_onnx.AiDecoderQuantizedOnnxSmokeTest.Fp8QdqRunsWithExpectedNumerics
2/4 Test #327: test_ai_decoder_quantized_onnx.AiDecoderQuantizedOnnxSmokeTest.Fp8QdqRunsWithExpectedNumerics .....................***Skipped   0.81 sec
    Start 328: test_ai_decoder_quantized_onnx.AiDecoderQuantizedOnnxSmokeTest.RejectsOversizedGatewayOutput
3/4 Test #328: test_ai_decoder_quantized_onnx.AiDecoderQuantizedOnnxSmokeTest.RejectsOversizedGatewayOutput ...   Passed    1.08 sec
    Start 329: test_ai_decoder_quantized_onnx.AiDecoderQuantizedOnnxSmokeTest.WritesGatewayOutputWithinSlot
4/4 Test #329: test_ai_decoder_quantized_onnx.AiDecoderQuantizedOnnxSmokeTest.WritesGatewayOutputWithinSlot ........   Passed    1.10 sec

100% tests passed, 0 tests failed out of 4

Total Test time (real) =   4.19 sec

The following tests did not run:
	327 - test_ai_decoder_quantized_onnx.AiDecoderQuantizedOnnxSmokeTest.Fp8QdqRunsWithExpectedNumerics (Skipped)

@kaiqiy-nv
kaiqiy-nv requested review from bmhowe23 and wsttiger August 4, 2026 09:11
@copy-pr-bot

copy-pr-bot Bot commented Aug 4, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Signed-off-by: Kaiqi Yan <kaiqiy@nvidia.com>
@kaiqiy-nv
kaiqiy-nv force-pushed the kaiqiy/fix-ai-decoder-gateway branch from dca731c to 1bf3014 Compare August 5, 2026 07:32
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant