ONNX-less direct builder cannot resolve decoder tensors for multimodal cosmos3_edge, and renaming them yields a numerically wrong engine

[filipemartinsubrobotics]

Component: experimental/builder (ONNX-less / direct engine builder)
Version: TensorRT Edge-LLM v0.10.1 (e8b2952), built from source
Checkpoint: W4A16_AWQ quantise of nvidia/Cosmos3-Edge (model_type: cosmos3_edge)
Impact: the direct path cannot produce an LLM engine for this checkpoint. The ONNX
path (tensorrt-edgellm-export → llm_build) works and produces a running engine, so
this blocks only the ONNX-less frontend.

Reproduced independently on two machines: a Jetson Orin Nano 8 GB (SM87, JetPack 7.2,
CUDA 13.2, TensorRT 10.16.2.10) and a DGX B300 (SM103, CUDA 13.0, TensorRT 10.13.2).

1. Summary

registry.py binds a single LLM component definition to both the multimodal
cosmos3_edge model type and the text-only cosmos3_edge_text. That definition
addresses decoder tensors as model.layers.*, which is correct only for the text-only
checkpoint. In a multimodal cosmos3_edge checkpoint the decoder lives under
model.language_model.*, so weight lookup fails on the first layer.

2. Reproduction

python -m experimental.builder.cli \
  --model-dir  <Cosmos3-Edge-INT4-AWQ checkpoint> \
  --engine-dir out/ \
  --plugin-path build/libNvInfer_edgellm_plugin.so \
  --components llm \
  --max-batch-size 1 --max-input-len 4096 --max-kv-cache-capacity 4096 --verbose

Observed

INFO  experimental.builder.core.builder: Detected CUDA device 0 with compute capability 8.7
DEBUG builder.ops.transformer: building layer 1/28
Traceback (most recent call last):
  File "experimental/builder/cli.py", line 347, in main
  File "experimental/builder/cli.py", line 247, in _build
  File "experimental/builder/cli.py", line 199, in _build_one
  File "experimental/builder/core/builder.py", line 368, in build_engine
    build_model(net, bundle, cfg, weights, args)
  File "experimental/builder/ops/transformer.py", line 352, in forward
    hidden_states, present = layer(hidden_states, ...)
  ...
KeyError: "checkpoint tensor not found: 'model.layers.0.input_layernorm.weight'"

Fails at layer 1 of 28. No engine is written. Identical failure on SM103.

Evidence

The checkpoint declares model_type: cosmos3_edge and carries no model.layers.*
tensors. Top-level prefixes in model.safetensors, by tensor count:

prefix	tensors
model.language_model.*	562
model.visual.*	437
model.projector.*	6
lm_head.*	1 (+ scale)

experimental/builder/models/registry.py:499-512 maps both model types to the same
component:

ModelFamily(
    "cosmos3",
    {
        model_type:
        _set(Component.LLM, Component.VISUAL, Component.UND_PREFILL,
             Component.GEN, Component.VAE_ENCODER)
        for model_type in ("cosmos3_omni", "cosmos3_edge", "cosmos3")
    } | {"cosmos3_edge_text": _set(Component.LLM)},
    {
        Component.LLM:
        _component("cosmos3.modeling_cosmos3_reasoner_text",
                   "Cosmos3ReasonerForCausalLM"),
        ...

Cosmos3ReasonerForCausalLM is therefore used for both the nested and the flat layout,
while addressing tensors as though the layout were always flat.

Expected

Either the LLM component resolves the decoder root per model type —
model.language_model for cosmos3_edge / cosmos3_omni, model for
cosmos3_edge_text — or the two model types get separate component definitions.

ModelComponent.__init__(self, ctx, prefix) already threads a prefix, so the plumbing
for the first option appears to exist; what is missing is a model-type-dependent value
for it.

3. Second data point: renaming the tensors is not a sufficient fix

On the SM103 machine we did not stop at the KeyError. We renamed the checkpoint
tensors into the flat layout the component expects — model.language_model.layers.N.*
→ layers.N.*, plus self_attn.q_proj→to_q (and k/v/o) and mlp.fc1→up_proj,
mlp.fc2→down_proj. That gets past the lookup and builds an engine.

The engine is numerically wrong. The same checkpoint through the ONNX path
(llm_build + llm_inference) on the same GPU answers correctly; the direct-builder
engine emits degenerate text.

prompt (identical formatted string in both, 18 tokens):
  <|im_start|>system<|im_end|><|im_start|>user\nHello<|im_end|><|im_start|>assistant\n<think></think>

ONNX path      -> 'Hello'   (argmax 22177)
direct builder -> '.'       (argmax 1046)

So a prefix-only fix may unblock the build while still producing bad output, turning a
loud KeyError into silent garbage. We would suggest a numeric check as part of the fix.

What we ruled out

Everything the direct builder consumes is identical to the ONNX that works. We
compared the builder's own assembly (weight_packing.int4.repack_modelopt_awq) against
the ONNX initializers:

tensor (layer 0)	result
q_proj / o_proj INT4 bytes	2097152 / 2097152 identical
k_proj / v_proj INT4 bytes	1048576 / 1048576 identical
mlp.up_proj / mlp.down_proj INT4 bytes	9437184 / 9437184 identical
all weight_scale	exactly equal after the builder's .T
all six pre_quant_scale	equal
input_layernorm / post_attention_layernorm	equal
embedding table vs embedding.safetensors	byte-identical

Plugin attributes match the ONNX node-for-node, including the asymmetric ones:

q_proj    gemm_n=2048 gemm_k=2048 group_size=128   w(4096,512)  scales(16,2048)
k/v_proj  gemm_n=1024 gemm_k=2048 group_size=128   w(2048,512)  scales(16,1024)
up_proj   gemm_n=9216 gemm_k=2048 group_size=128   w(18432,512) scales(16,9216)
down_proj gemm_n=2048 gemm_k=9216 group_size=128   w(18432,512) scales(72,2048)

Graph shape is the same in both: x * pre_quant_scale -> Int4GroupwiseGemmPluginV2(w, scales).
The two runtime config.json files have no differing shared values (same
rope_theta, rope_scaling.mrope_section, head counts, head_dim, vocab_size). The
embedding binding resolves correctly (__embedding__ <- embed_tokens.weight, fp16,
[131072, 2048]).

Also not the cause:

• Not one kernel — fails identically with --int4-gemm-plugin-version 1 and 2.
• Not the runtime weight binding — fails with weights externalized and baked
  (--externalize-weights embedding).
• Not the driver — the builder's engine produces the same wrong text under
  llm_inference as under experimental.server.LLM.
• Not the transformer graph — an FP16 engine through the same direct builder on the
  same GPU is correct (405 tok/s, coherent output), so RMSNorm, RoPE and the attention
  path are fine. That FP16 run used the original checkpoint, whose decoder tensors are
  already flat, so it needed no renaming: it exonerates the graph, not the rename.
  The rename is instead evidenced by the byte-identity table above, where every renamed
  tensor matches the ONNX initializer the working engine uses.

Where it first diverges

With EDGELLM_DUMP_LOGITS_KVCACHE_LAYERS=28 on both engines, same prompt, sliced to the
dumped context_lengths (18):

layer 0   K rel=2.31e+00   V rel=1.98e+01     <- already wrong
layer 1   K rel=9.40e-01   V rel=8.87e-01
layer 2   K rel=1.43e+00   V rel=2.39e+00
...
logits    rel=8.98e-01     argmax 22177 vs 1046

(rel = max|diff| normalised by max|ONNX|.) Layer-0 V is the informative one: it
carries no RoPE and sits before the MLP, so it depends only on the embedding, the input
layernorm and v_proj — all three verified identical above. Its magnitude is ~20x too
large: ONNX V range [-0.490, 0.499], direct-builder V range [-9.883, 7.922].

What we did not establish

We could not localise it further. The shipped debugger dumps KV and logits only, not
intermediate activations, so we cannot see whether the RMSNorm output or the plugin
result is the first wrong value — that appears to need a small runtime patch. We also
only have SM87 and SM103, and the SM87 machine cannot build this path at all because of
the bug in section 1, so the numeric half is single-architecture. Given FP16 is correct
on the same part and the inputs are byte-identical, we would not expect it to be
architecture-specific, but we have not shown that.

4. Two documentation issues found alongside

1. pip install ".[builder]" refers to an extra that does not exist in v0.10.1's
   pyproject.toml. Only tools and server are defined. Following the documented
   line fails.

2. Falling back to .[tools] pins torch==2.13.0 from PyPI, which on aarch64/Tegra
   is not a CUDA-enabled build. On Jetson the working recipe is a
   --system-site-packages venv (which picks up the system tensorrt, matching the
   version the C++ was built against), the package installed with --no-deps, and the
   JetPack torch supplied on PYTHONPATH.

   For reference, the direct builder's actual third-party surface is small — numpy and
   tensorrt under experimental/builder, plus torch and safetensors for
   tensorrt_edgellm — so the heavy tools pin set is not required to run it.

5. Environments

Jetson Orin Nano 8 GB, L4T R39.2 (JetPack 7.2), SM 8.7
  Edge-LLM v0.10.1, cmake -DCMAKE_BUILD_TYPE=Release -DEMBEDDED_TARGET=jetson-orin
           -DCUDA_CTK_VERSION=13.2 -DENABLE_CUTE_DSL=ALL
           -DCMAKE_CUDA_ARCHITECTURES=87 -DTRT_PACKAGE_DIR=/usr
  TensorRT 10.16.2.10-1+cuda13.2 (system python bindings, same version)
  torch 2.12.0+cu132, python 3.12.3

DGX B300, SM 10.3
  Edge-LLM v0.10.1 with #205 and #207 applied,
           -DENABLE_CUTE_DSL="fmha;int4_fp16_gemm" -DCMAKE_CUDA_ARCHITECTURES=103
  CUDA 13.0, TensorRT 10.13.2

-DENABLE_CUTE_DSL was set as the direct-builder docs require, and the same plugin
library builds and loads correctly for the ONNX path on both machines.

---

Found during the NVIDIA / OpenHackathons / Oracle Open Models Codefest 2026, while
bringing up an INT4-AWQ Cosmos3-Edge reasoner for an offline-first search-and-rescue
robotics entry (Team UBR Stack). The Jetson Orin Nano is the production target; the
B300 is a bench machine used for evaluation.