One-Forward-Pass Readout Engine for TensorRT-LLM—1,014 tokens in 92.3 ms

Accelerated Computing DGX Spark / GB10 User Forum DGX Spark / GB10 Projects
, , ,
1

I’m sharing a TensorRT-LLM branch I’ve been running across a 3× DGX Spark cluster.

It replaces the usual autoregressive decode loop with a compiled readout engine: the frozen base model runs exactly once per query, its internal state is captured, and a separately trained decoder reads the entire response from that captured state in a single TensorRT execute_async_v3.

Before anybody assumes I am secretly an NVIDIA inference wizard: I have no formal experience in this field. I’m a forward-deployed engineer by trade and started this mostly by tinkering.

Some of the early work was, in the most literal sense, vibe-coded.

The current measurements, however, come from real engines, tests, parity audits, and served requests—not from the vibes. My interpretation of those results may still be wrong, and I would genuinely appreciate corrections from people who work on inference runtimes, compilers, or model architecture professionally.

Everything below was measured on GB10:

  • sm_121

  • compute capability 12.1

  • CUDA 13.0

  • TensorRT 10.14.1.48

  • nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-NVFP4

The model is NemotronHForCausalLM: 52 layers, hybrid Mamba + MLP + attention, hidden size 2,688, vocabulary 131,072.

Repository: https://github.com/lcoleman0422/TensorRT-LLM
Branch: nemoclaw/v9.16-rc14

It is based on upstream commit:

93cb6518b6d6dbd6095748189e626db731f44545

plus three local SM121 fixes.

What it does

A normal decoder performs roughly one model forward step per emitted token.

This path performs one base-model forward pass per request.

The resulting hidden state is captured and projected into a fixed latent representation. A separately trained, non-autoregressive readout decoder then emits an entire answer block from that latent.

There is:

  • no base-model forward pass per answer token;

  • no draft-and-verify loop;

  • no speculative decoder;

  • no second base-model pass on the answer path;

  • no autoregressive fallback hiding behind the curtain.

The compiled engine graph uses ONNX opset 17. Attention is expressed through explicit operators so a fused-attention fast path cannot silently bake in a static sequence shape.

The graph contains:

  • context pooling and projection;

  • a two-stage latent-repair path;

  • functional sinusoidal position terms computed from a position-index input;

  • one explicit-operator transformer mixer layer;

  • chunked full-vocabulary projection;

  • a running top-2 merge;

  • an in-graph first-EOS scan.

The functional positions mean there is no learned position table imposing a fixed readout-length ceiling.

Chunked vocabulary projection

The output vocabulary contains 131,072 tokens.

Instead of materializing the complete vocabulary-logit tensor, the engine evaluates:

8 chunks × 16,384 vocabulary entries

The resident lm_head is bound as an engine input. It is not baked into the engine and is not copied again for every request.

Each chunk updates a running top-2 result, so the full vocabulary-logit tensor is never resident at once.

EOS detection also happens inside the graph using eos_token_id = 11.

Dynamic profile

prompt_len:   [1, 512, 4096]
answer_depth: [1, 8192, 8192]

Autoregressive fallback is hard-disabled.

fallback_to_autoregressive is fixed to False, and the runtime rejects configurations that attempt to enable it.

The path therefore fails closed:

no_autoregressive_fallback_passed: true

That is an important distinction. This is not an ordinary decoder that happens to take a fast path under favorable conditions. If the compiled StateBlock path cannot run, the request does not quietly become an autoregressive request.

Measured results

Latency on a single GB10

Measurement Result
Live real-chat request 1,014 tokens in 92.3 ms
Served 8,192-token envelope 511 ms
Single 8,192-token engine execution approximately 0.6 s
Research-path baseline, concurrent parity audit 1.535 s
Research-path baseline, block-size sweep 4.91 s

The live request completed in 92,335 µs, returned HTTP 200, and stopped on native EOS.

The 8,192-token served envelope completed in 511,356 µs, committed 7,317 tokens, and stopped on native EOS.

The single-operation 8,192-depth engine measurements were:

  • 0.62 seconds during the export-time audit;

  • 0.586 seconds during the parity-time audit.

I am deliberately reporting both research-path baselines rather than selecting the flattering one.

They are two different on-disk measurements of what is nominally the same reference operation:

  • 1.535 seconds in the concurrent parity run;

  • 4.91 seconds in the block-size sweep, with n=2 and σ = 68.6 ms.

I do not currently have an artifact that explains the gap between them. Therefore, I am not claiming one clean universal speedup number.

Token parity against the reference decoder

The compiled engine is tested against the non-compiled reference decoder that trained and defined the readout behavior.

Depth Exact-match fraction Mismatches
256 1.0 0
512 1.0 0
1,024 0.99902 Near-ties only
2,048 1.0 0
4,096 0.99927 Near-ties only
8,192 0.99939 Near-ties only

Every observed mismatch occurred at a top-1/top-2 near-tie.

The maximum observed reference margin at a mismatch was:

0.00384

The audit’s near-tie threshold is:

0.05

These mismatches are consistent with legal floating-point reduction reordering between:

  • the engine’s 64-token-chunked accumulation; and

  • the reference implementation’s full-depth accumulation.

I am therefore classifying them as numerical tie resolution rather than demonstrated semantic divergence.

The distinct-4-gram difference is 0.0 at all six tested depths.

Stop-boundary parity

The EOS boundary matches exactly at:

  • depth 1,024: position 1,013 versus 1,013;

  • depth 8,192: position 6,673 versus 6,673.

At depths 256, 512, 2,048, and 4,096, neither side emits EOS. Both return the no-EOS sentinel:

1 << 20

Those four cases are agreement on the absence of a stop, not matching EOS discoveries.

K-cohort execution

The branch also supports K readout lanes resolved from one packed context forward.

Rows are assigned to lanes using cumulative-sum offsets, with a layout-identity invariant:

total captured rows == Σ per-lane row counts

The current per-lane served readout wall times are:

K Wall time
1 0.17 s
2 0.23 s
4 0.36 s
8 0.71 s

The current maximum supported cohort is:

supported_k_max = 8

A separate packed K-dimensional engine, kcohort_v2, supports:

K:     [1, 8]
depth: [1, 2048]

It took 51.3 seconds to build and executes the entire cohort with one enqueue:

K Packed-engine wall time
1 0.082 s
2 0.121 s
4 0.198 s
8 0.357 s

Minimum token agreement is 0.998, with all observed mismatches satisfying the same near-tie proof.

Important K-cohort caveat

The packed engine is currently classified as:

ENGINE_PROVEN

It is not runner-bound.

The served path currently performs:

  1. one packed cohort context forward;

  2. K separate TensorRT readout enqueues, one per lane.

The packed kcohort_v2 engine has been tested independently, but the production runner does not yet invoke it.

Depths above 2,048 remain on the per-lane v1 engine.

Full self-attention at:

K = 8
D = 8192

would materialize roughly 45 GiB of attention scores on GB10. That is the present memory wall—not a philosophical objection to larger K.

I would especially like to compare notes with anyone experimenting with readout-style inference on Spark who has found a less ridiculous solution to that K=8 attention tensor.

Sampling

A top-8 Gumbel selection path is compiled into an engine using explicit noise and temperature inputs.

At:

temperature = 0

it reduces to exact argmax.

The sampling implementation is classified as:

ENGINE_READY

It is still fail-closed in serving.

Production decoding is currently argmax-only until the runner explicitly binds the sampling-capable engine.

So, no, the production numbers below are not secretly helped by sampling five answers and selecting the nicest one.

Determinism and origin attribution

The base model’s MoE and Mamba implementations contain nondeterministic floating-point reductions.

Demanding byte-identical output across executions would therefore mix two different questions:

  1. Did the captured base-model state change?

  2. Did the readout produce different output from the same captured state?

The runtime records a forward-state digest and applies an origin-aware integrity rule:

  • PASS if the output is identical;

  • PASS if output differs and the captured forward-state digest also differs;

  • FAIL if the same captured forward-state digest produces a different output.

The current serving mode is:

strict_output_digest_match

This does not make the base model deterministic. It attempts to identify whether any divergence originated before or after the readout boundary.

Coherence

Distinct-4-gram measurements for a single-operation readout at each answer depth are:

Depth Distinct-4-gram
256 0.960
512 0.963
1,024 0.973
2,048 0.968
4,096 0.854
8,192 0.754

The degradation at longer depths is real and should not be waved away.

Length performance tracks the length distribution of the decoder’s training data. Extending supervision to responses of approximately 7,373 tokens is what pushed coherent single-operation output toward the full 8,192-position envelope.

This is therefore not “infinite generation for free.”

It is a learned, fixed-envelope readout whose quality remains dependent on:

  • supervision coverage;

  • representation quality;

  • decoder capacity;

  • long-range mixer behavior;

  • the distribution of tasks presented at inference time.

StateBlock captured-state readout

The base model produces a frozen forward state. A separate functional-position decoder reads an answer block from that state without feeding emitted tokens back into the model.

This corresponds to project’s StateBlock and “future-sheet” work: capture the model’s resolved state first, then perform a governed readout or commit from that state instead of repeatedly re-entering the base model.

Active K-path and task-world work

The K-cohort path comes from the project’s branch-state work.

K is not intended to mean “take the top K logits and call it parallel reasoning.” Each lane has explicit identity, captured-layout ownership, sequence-length accounting, and isolation requirements.

In the broader runtime project, that machinery also supports separately evolving task worlds and recurrent/attention branch state. This repository’s packed readout engine is narrower: it proves packed cohort readout, not the entire multi-world runtime in one TensorRT graph.

Compact commit and single-consumer ownership

The readout path produces a compact token-ID answer block and consumes each registered capture exactly once.

That follows the same architectural lineage as the project’s direct-state handoff and CAS-style commit work: state must have an owner, a layout identity, and an unambiguous commit boundary.

The request-scoped capture registry exists because the original process-wide stash violated that discipline and raced under concurrent traffic.

Mamba and hybrid-state lineage

The broader work includes recurrent-state and cooperative Mamba/SSM memory mechanisms.

This readout branch does not claim to replace or reproduce that entire system. It does, however, operate on a hybrid Mamba/attention model and inherits the same basic concern: recurrent, convolutional, attention, and request-layout state must remain associated with the correct request and lane.

Provenance and digest gates

The forward-state digest and fail-closed release gates are implementation relatives of the project’s authenticated-state and origin-tracking work.

They are not proof that the output is correct. They provide evidence about which side of the captured-state boundary produced a divergence.

For people following the numbered internal work, the closest conceptual relatives are:

  • P3: direct state/bus write and avoiding unnecessary serialized handoffs;

  • P7: authenticated state lineage and cryptographic boundary discipline;

  • P32: CAS-style ownership and commit semantics;

  • P38CIP: cooperative Mamba/SSM memory and recurrent-state continuity.

Capture registry

The capture registry may be the most generally useful part of the branch for anyone attempting something similar.

It replaced an early process-wide capture stash that could race under concurrent requests.

Before the model forward pass, the runtime registers an immutable per-request capture plan.

When the hook records the state, the registry verifies:

  • the hook-stamped request identity;

  • per-lane sequence lengths;

  • total captured row count;

  • cohort layout;

  • expected capture mode.

The capture is then consumed exactly once against a layout digest.

Among the 13 recorded capture-registry counters is:

stateblock_capture_cross_request_events

The goal is not merely to notice a bad tensor shape. It is to prove that a valid-looking tensor from request A cannot silently become the readout state for request B.

Canonical paths

Core C++

tensorrt_llm/batch_manager/stateBlock.cpp                       692
tensorrt_llm/batch_manager/stateBlockSemanticDecoder.cpp        818
tensorrt_llm/batch_manager/stateBlockDecoderFeedback.cpp        481
tensorrt_llm/batch_manager/stateBlockFrozenModelAnswerGate.cpp  277
tensorrt_llm/batch_manager/stateBlockCompactCommitKernels.cu     191
include/tensorrt_llm/batch_manager/stateBlockTypes.h             181

Python runtime and serving

tensorrt_llm/runtime/stateblock_functional_readout.py
tensorrt_llm/runtime/stateblock_capture_registry.py
tensorrt_llm/serve/stateblock_functional_session.py
tensorrt_llm/serve/openai_server.py
tensorrt_llm/serve/openai_protocol.py

Export scripts

scripts/stateblock_native_bridge_engine_export.py
scripts/stateblock_native_kcohort_engine_export.py

The first script exports the v1 single-lane engine.

The second exports the packed kcohort_v2 engine.

Build

The wheel is built inside:

nvcr.io/nvidia/pytorch:26.02-py3

for aarch64 and SM121 only.

./scripts/build_wheel.py \
  --trt_root /usr/local/tensorrt \
  --cuda_architectures '121-real' \
  --fast_build \
  --configure_cmake \
  --job_count 16 \
  --extra-cmake-vars \
  'BUILD_DEEP_EP=OFF;BUILD_DEEP_GEMM=OFF;BUILD_FLASH_MLA=OFF'

Two patches are required in that image.

This was the portion of the experiment where the vibes became noticeably less fun.

1. Install ZeroMQ development dependencies

apt-get install -y --no-install-recommends libzmq3-dev pkg-config

2. Replace unavailable static NVRTC targets

The pytorch:26.02-py3 image provides:

libnvrtc.so
libnvrtc-builtins.so

but not the corresponding .a static libraries expected by the build configuration.

I rewrite those targets with:

find cpp/ -name CMakeLists.txt -exec sed -i \
  -e 's/CUDA::nvrtc_builtins_static/CUDA::nvrtc_builtins/g' \
  -e 's/CUDA::nvrtc_static/CUDA::nvrtc/g' {} +

Host-side C++ tests

cmake -S cpp -B cpp/build
cmake --build cpp/build --target <TestName>

Serving

trtllm-serve "$MODEL" \
  --backend pytorch \
  --host 0.0.0.0 \
  --port 8402 \
  --max_seq_len 12288 \
  --trust_remote_code \
  --config k8s/stateblock/serve_config.yaml

The path is gated by:

TRTLLM_STATE_BLOCK_CAPTURE
TRTLLM_STATE_BLOCK_CAPTURE_HIDDEN

It currently requires:

TLLM_WORKER_USE_SINGLE_PROCESS=1

The intended go-live topology is single-process.

A multiprocess-executor embodiment is designed but deferred. Mixed autoregressive and StateBlock traffic in one pod is deliberately out of scope; the intended deployment uses dedicated pods.

Caveats, collected in one place

So nobody has to excavate the repository to find the asterisks:

  • I am a forward-deployed engineer, not an ML-inference researcher, TensorRT specialist, or CUDA compiler engineer.

  • I had no professional background in this field before beginning the project.

  • Yes, portions of the initial implementation were aggressively vibe-coded.

  • The current quantitative claims are backed by recorded runs and artifacts, but my conclusions may still be incomplete or wrong.

  • The compiled decoder is separately trained. This is not the frozen base model magically emitting 8,192 tokens directly from its ordinary LM head in one operation.

  • The system uses a bounded answer envelope, not unconstrained infinite-length generation.

  • Long-output coherence degrades and remains strongly dependent on supervision-length coverage.

  • The packed K-cohort engine is proven independently but is not connected to the served runner.

  • The served K path currently performs one context forward followed by K per-lane TensorRT enqueues.

  • Production decoding is argmax-only.

  • The sampling head is engine-ready but not served.

  • The serving topology is single-process only.

  • Mixed autoregressive and StateBlock traffic in the same pod is unsupported by design.

  • The reported measurements come from the PyTorch/TensorRT path.

  • A separate later live-chat daemon is classified in my own audits as DEMO_REPLAY.

  • None of these performance numbers should be attributed to that daemon.

  • One release audit records release_decision: GO, but its empty blocker list excludes two failing gates.

  • That “GO” should be treated as scoped, not absolute.

In other words: it runs, it has receipts, and it is interesting enough to share.

It is not a declaration that autoregressive decoding is dead, that I have solved language generation, or that three DGX Sparks have granted me tenure.

What I would like feedback on

I would especially appreciate technical criticism or comparisons around:

  • request-scoped hidden-state capture under concurrent traffic;

  • binding a resident lm_head as a TensorRT engine input;

  • chunked vocabulary projection without full-logit materialization;

  • packed K-lane execution;

  • alternatives to the K=8 × D=8,192 self-attention memory wall;

  • practical determinism contracts for hybrid Mamba/MoE models;

  • runner integration for one-enqueue packed cohort execution.

If you are doing readout-style inference on DGX Spark—or have already discovered why this entire approach is secretly cursed—I would genuinely like to compare notes.

2 Likes