[NVIDIA #6] TensorRT-LLM as ultra-low-latency AHE eval backend on GB10

[AugustChaoTW]

Summary

Add NVIDIA TensorRT-LLM as a fourth AHE eval backend optimized for maximum throughput and minimum latency on the GB10 (Blackwell SM 12.1). TensorRT-LLM applies kernel fusion, FP8 quantization, and in-flight batching, achieving 2–4× higher throughput than vLLM for AHE eval batch workloads.

GB10 Blackwell Advantages for TRT-LLM

Feature	GB10 Benefit
FP8 Tensor Cores	Native FP8 inference (Blackwell-first)
121.6 GB VRAM	Full 70B models in FP8 without quantization loss
NVLink Bandwidth	High KV cache transfer for paged attention
CUDA 13.0	Latest TRT-LLM kernels optimized for SM 12.1

Installation

# TensorRT-LLM for Blackwell (CUDA 13)
pip install tensorrt-llm --extra-index-url https://pypi.nvidia.com
# Or via NGC container (recommended for production):
docker pull nvcr.io/nvidia/tensorrt-llm:latest

Architecture

ahe/run_eval.py --backend trtllm-local
        ↓
ensure_trtllm_server(model, engine_dir)
        ↓
TRT-LLM triton server (port 8001)
        ↓
triton_client POST /v2/models/ensemble/generate
        ↓
FP8 inference on GB10 with paged KV cache + prefix caching

Changes

1. owl.yaml — trtllm-local backend

backends:
  trtllm-local:
    command: python3
    args:
      - "ahe/trtllm_client.py"
      - "{prompt}"
    env:
      TRTLLM_SERVER_URL: "http://localhost:8001"
      TRTLLM_MODEL: "ensemble"

2. ahe/trtllm_client.py — new thin client

"""TensorRT-LLM Triton Inference Server client for AHE eval."""
import sys, json
import tritonclient.http as httpclient

def generate(prompt: str, server_url: str, model: str = "ensemble",
             max_tokens: int = 4096) -> str:
    client = httpclient.InferenceServerClient(url=server_url.replace("http://",""))
    # Prepare input tensors
    input_ids = tokenize(prompt)
    inputs = [httpclient.InferInput("input_ids", input_ids.shape, "INT32")]
    inputs[0].set_data_from_numpy(input_ids)

    outputs = [httpclient.InferRequestedOutput("output_ids")]
    response = client.infer(model, inputs, outputs=outputs)
    return detokenize(response.as_numpy("output_ids")[0])

if __name__ == "__main__":
    prompt = sys.argv[1] if len(sys.argv) > 1 else ""
    result = generate(prompt,
                      server_url=os.environ.get("TRTLLM_SERVER_URL","http://localhost:8001"),
                      model=os.environ.get("TRTLLM_MODEL","ensemble"))
    print(result)

3. Engine build script scripts/build_trtllm_engine.sh

#!/bin/bash
# Build TRT-LLM engine for Qwen2.5-32B on GB10 (SM 12.1)
MODEL=${1:-"Qwen/Qwen2.5-32B-Instruct"}
OUTPUT_DIR=${2:-"/tmp/trtllm-engines/qwen32b"}

python3 -m tensorrt_llm.commands.build \
  --checkpoint_dir "$MODEL" \
  --output_dir "$OUTPUT_DIR" \
  --gemm_plugin float16 \
  --use_paged_context_fmha enable \
  --paged_kv_cache enable \         # Prefix KV Cache
  --max_batch_size 32 \
  --max_input_len 8192 \
  --max_output_len 4096 \
  --tp_size 1 \
  --workers 1

Benchmark Expectations

Backend	Throughput (tokens/s)	TTFT (1K prefix)
Ollama (qwen3.5-35b)	~80	~2s
vLLM (prefix cached)	~200	~0.3s
TRT-LLM FP8 (GB10)	~600+	~0.1s

Higher throughput means AHE eval suites run 3–8× faster, enabling more frequent harness iteration.

Prerequisites

• TensorRT-LLM installed (pip install tensorrt-llm)
• tritonclient installed (pip install tritonclient[http])
• Engine built for target model (see scripts/build_trtllm_engine.sh)
• Triton Inference Server running with TRT-LLM backend

Acceptance Criteria

• pip install tensorrt-llm succeeds on CUDA 13.0 / GB10
• scripts/build_trtllm_engine.sh builds a valid engine for one model
• python ahe/run_eval.py --suite terminal --backend trtllm-local runs end-to-end
• Throughput benchmark documented in issue comments
• Paged KV cache + prefix caching flags confirmed active
• Graceful error if Triton server not running

Related Issues

• Builds on: [[NVIDIA #2] Add vLLM local backend with RadixAttention prefix caching to AHE eval #50] vLLM backend (same owl.yaml backend pattern)
• Synergy: [[Prefix-KV #1] Stable-first system prompt ordering for prefix cache efficiency #42]–[[Prefix-KV #7] Research: Radix Tree memory index for prefix-aware retrieval #48] Prefix-KV application layer (stable prefix → TRT-LLM cache hits)

Effort: 1 week | Priority: P3

[AugustChaoTW]

GPU-agnostic requirement: This feature must degrade gracefully without an NVIDIA GPU. See #56 for the full compatibility contract, fallback implementation details, and CI requirements that must pass before this issue can be merged.