NVIDIA / NVIDIA/TensorRT-Edge-LLM

[Regression][Jetson AGX Thor] Layer-0 V tensor corruption in v0.9.1; v0.4.0 works correctly on the same JetPack 7.1 system

Open
#151 6 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
563
Forks
135
Avg merge
14h 13m
Merged PRs (30d)
1

Description

The same Jetson AGX Thor system running JetPack 7.1 produces
correct results with TensorRT Edge-LLM v0.4.0.

TensorRT Edge-LLM v0.9.1 reproduces the layer-0 V-cache corruption
on the same device and software stack.

This suggests an Edge-LLM regression introduced after v0.4.0,
rather than a general JetPack 7.1 compatibility issue.

I am seeing a reproducible numerical correctness issue on Jetson AGX Thor with TensorRT Edge-LLM v0.9.1.

The issue was initially observed as invalid LLM output, but layer-by-layer validation isolates the first divergence to:

decoder layer 0
V projection → AttentionPlugin → V KV cache

The V projection itself is correct. However, in the full graph, the V cache is severely corrupted.

Most importantly, if the V projection tensor is marked as an additional ONNX graph output, the issue completely disappears.

This suggests a possible TensorRT/Myelin tensor lifetime, buffer reuse, or materialization issue around the V tensor.

Environment
Device: Jetson AGX Thor
Jetson Linux: R38.4
JetPack: 7.1
CUDA: 13.0.48
TensorRT: 10.13.3.9
TensorRT Edge-LLM: v0.9.1
Commit: 7f061f21f0a581ba234a1e233c9315b89d8e47d6
Target: jetson-thor
Architecture artifact: aarch64/sm_110
ENABLE_CUTE_DSL=ALL

The repository and submodules are complete.

The problem is reproducible in FP16, so NVFP4 quantization is not the root cause.

The original Hugging Face/PyTorch model runs correctly on the same Thor device.

First numerical divergence

Layer 0 V cache:

cosine = 0.015065519
max_abs = 13.1431694
mean_abs = 0.255735815
shape = [8, 22, 128]
dtype = FP16 vs FP16
NaN/Inf = 0 / 0

Layer 0 K cache from the same execution is essentially correct:

cosine ≈ 1.0
V projection isolation

Testing layer 0 V projection independently:

TensorRT vs PyTorch

cosine = 0.999999881
max_abs = 0.00390625
mean_abs = 0.000058087

So the V projection weights/GEMM themselves appear correct.

Critical observation
Normal full graph

Without exposing the V intermediate tensor:

V cache vs PyTorch V projection

cosine = 0.015077175
Mark V projection as graph output

If ONNX linear_2 (V projection) is marked as an additional graph output:

TensorRT V projection vs PyTorch:

cosine = 0.999998927

And:

V cache vs TensorRT V projection:

cosine = 0.999999881
max_abs = 0

The corruption completely disappears.

Control experiments:

mark Q output → FAIL
mark K output → FAIL
mark V output → PASS

No model weights, inputs, or operator parameters are changed.

CuTe DSL A/B

I also rebuilt Edge-LLM with:

ENABLE_CUTE_DSL=OFF

and confirmed the CuTe DSL FMHA symbol is absent.

The issue still reproduces:

layer 0 V cosine ≈ 0.01507

Therefore this does not appear to be caused by the CuTe DSL FMHA path.

compute-sanitizer

Running:

compute-sanitizer --tool racecheck

on the non-CuTe-DSL build reports:

4 errors
48 warnings

inside the Myelin-generated kernel_mha path.

I cannot confirm that these reports are directly causal, but they appear consistent with the observed behavior.

Expected behavior

The V cache should numerically agree with the PyTorch reference regardless of whether the intermediate V tensor is marked as a graph output.

Actual behavior
V not exposed:
cosine ≈ 0.015

V exposed as graph output:
cosine ≈ 1.0

This makes the issue look like a TensorRT/Myelin optimization, memory lifetime, or buffer reuse problem before AttentionPlugin consumes the V tensor.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Reproduce the issue on Jetson AGX Thor with the stated JetPack, TensorRT, and Edge-LLM versions, comparing the normal graph with one exposing ONNX linear_2. Inspect the AttentionPlugin input path and the Myelin-generated kernel_mha behavior, using compute-sanitizer as described. Done means the layer-0 V cache matches the reference without exposing the intermediate tensor.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, python
Domain
embedded-iot, machine-learning, performance
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.