Corey James- 6lines added to the decoder definition to expose the hidden state
- 128tokens per input across every extraction build
Going one level below the serving stack
Inference Fabric needed to read a language model's internal representations for millions of sequences, and serving stacks are built to produce tokens. NVIDIA TensorRT-LLM carries decoding machinery that a single read has no use for, so I went one level down to raw TensorRT and built the extraction path around the engine itself.
Marking the hidden state in the model graph
The first step was in the model graph. I added lines to the decoder definition in my TensorRT-LLM source build that mark the hidden state after the final normalisation and before the output head as an extra engine output. Any decoder engine built from that tree exposes the representation the downstream work needs.
Lit: the added lines. The generation path is unchanged.
Extraction engines that only read
The engines followed in a sequence of builds on a dense 7B-class model with FP8 weights and an FP8 GEMM plugin. The first kept a key-value cache and generated a single token. The final ones have the cache disabled and generate nothing, because a single read needs neither and pays for both in memory. They run one forward pass over a whole batch, and dropping the cache is what frees the memory for batches that size.
The build recipe is part of the consumer onboarding procedure and each engine comes with a contract describing its limits, so every extraction engine is a reproducible artefact, built per card for the Blackwell generation the cards belong to.
A runtime that keeps the GPU fed
The runtime that drives them is a C++ layer at the TensorRT boundary, called from Rust. For extraction it sets dynamic input shapes per batch, checks that the marked output has the shape it expects, runs one enqueue of the whole engine per pass, and hands the hidden states to the fabric's transport, which carries tensors without interpreting them.
The same layer keeps the GPU fed: memory comes from a pinned allocator that pre-faults every page and is never allocated on the hot path after warm-up, the compute stream is never synchronised, and host and device coordinate through a stream-ordered doorbell, a path used for the extraction backend and kept off for the vendor's executor, which reads its requests on the CPU. I checked both rules with NVIDIA Nsight Systems timelines.
Feeding the vector store
Extraction jobs come through the same token-aware broker as generation; see A message broker that understands tokens. The extracted vectors flow into the fabric's GPU-resident vector store, where hand-written CUDA kernels search them and score cluster quality.
I build the runtime discipline that keeps the result saturated and safe. TensorRT and its kernels are NVIDIA's; the graph change, the engine builds, the runtime around them and the extraction pipeline are mine.
I drop to NVIDIA's lower-level engine when the standard serving path is the wrong tool, and change the vendor graph only where it buys a capability.