trtllm-codebase-exploration

TensorRT-LLM Codebase Exploration Guide

Why This Matters

TRT-LLM is a large codebase (~500K lines) with many reusable abstractions. The most common source of wasted effort is reimplementing something that already exists. On the short-seq MHA branch, ~250 lines were written across 4 iterations before discovering that a 10-line dispatch to an existing method (forward_context_default) was the right solution.

Rule of thumb: Spend 30 minutes reading existing code before writing 1 line of new code.

MANDATORY: Ignore the TensorRT backend, focus on the PyTorch backend

Step-by-Step Exploration Workflow

Step 1: Map the Class You're Modifying

Before adding code to a class, understand its full structure:

# List all methods (not just forward*)
grep -n "def " tensorrt_llm/_torch/attention/attention.py | head -50
 
# List all attributes set in __init__
grep

Step 2: Trace Existing Forward Methods

Read EVERY forward method in the class. Understand what each one does, what inputs it expects, and what backends it uses.

Ask yourself:

• Does any existing forward method already compute what I need?
• Can I dispatch to an existing method by setting up the right state?
• What would I need to change (attributes, guards, assertions) to reuse it?

Step 3: Search for Existing Backends and Utilities

Step 4: Check What the Fused Kernels Handle

Many operations you might implement manually are already handled by fused C++ kernels:

Common surprise: When rope_fusion=True (apply_rotary_emb=False), the fused attention kernel handles RoPE internally via latent_cache. Writing custom RoPE code in Python is unnecessary and will double-apply RoPE.

Step 5: Check Assertions and Invariants

Existing assertions may need updating when you add a new code path. Don't work around them — change them if your new path makes them invalid:

Example: DSA models had assert self.mha is None. When adding short-seq MHA (which creates self.mha for DSA models), the assertion was changed to assert self.mqa is not None — the actual invariant being tested.

Step 6: Understand Weight Layouts

Weight layouts often differ between HuggingFace checkpoints and TRT-LLM's loaded format:

Critical for tests: Always initialize test weights in the loaded layout, not the HF checkpoint layout.

Step 7: Trace Method Limitations

After identifying a method to reuse, understand what it does NOT handle:

Ask yourself:

• What scenarios does this method handle? (fresh prefill? cached KV? chunked context?)
• What scenarios does it NOT handle?
• Is there a higher-level dispatcher that routes to this method for the correct subset of cases?
• If I call this method directly, which scenarios will I silently mishandle?

Example: forward_context_default() handles fresh prefill but does NOT attend over cached KV tokens. forward_context() is the dispatcher that routes to forward_context_default, forward_context_with_cached_kv, or forward_context_with_chunked_prefill based on context state and SM version. Calling forward_context_default directly during chunked context silently drops cached tokens.

Key Discovery Patterns

Pattern: "Can I Reuse an Existing Forward Method?"

1. Read the target forward method (e.g., forward_context_default)
2. Compare it to what your new code path needs to do
3. If >70% overlap, dispatch to the existing method instead of writing a new one
4. Adjust attributes/state in __init__ to make the dispatch work

Pattern: "Is This Already Handled by a Fused Kernel?"

1. Check if the operation is in the attention backend's scope
2. Check the apply_rotary_emb / rope_fusion flag
3. Check latent_cache handling
4. If the fused kernel handles it, DON'T reimplement in Python

Pattern: "Am I Calling the Right Abstraction Level?"

1. Identify the method you plan to call
2. Search for methods that CALL this method — there may be a dispatcher above it
3. Check if the dispatcher handles edge cases your direct call would miss
4. Prefer calling the dispatcher over the specific handler

Pattern: "Does a Utility Already Exist?"

1. Search tensorrt_llm/_torch/utils.py for compiled helpers
2. Search tensorrt_llm/_torch/modules/ for module-level utilities
3. Search test fixtures in tests/unittest/_torch/ for test setup patterns

Common Exploration Mistakes

File Reference for Exploration