This document is relevant for: Trn2, Trn3
Tensor Replacement — Design#
Purpose#
Tensor replacement injects reference tensors into a model’s forward pass at runtime. This lets you swap any named intermediate tensor (e.g., MoE router logits) with a reference value, isolating specific components from hardware numerics during accuracy debugging.
How It Works#
Reference tensors (e.g., from an HF model) are captured once and held in memory (or loaded from disk via reference_captures_path). On each forward pass, the model runner builds replacement tensors by indexing into the reference data using the scheduler’s positions and request IDs. The result is a [bucket_size, feature_dim] tensor that matches Neuron’s runtime layout — bucketed, padded, and ordered by the scheduler.
Replacement tensors are injected via global context (set_active_context / get_replacement_tensor). The model runner sets the context before each forward pass and clears it after.
Each DP replica builds its own replacement context from the same reference captures. The cross-DP all-gather in the MoE kernel combines injected values downstream.
Phase 1: Reference Capture
──────────────────────────
Source model (e.g., HF)
───────────────────────
Capture target tensors per prompt
(e.g. router logits via monkey-patching or TensorCaptureModel)
│
▼
reference_captures: List[Dict[module_name, List[step_tensor]]]
- One entry per prompt
- Step 0 = prefill [1, prompt_len, E]
- Steps 1+ = decode [1, seq_so_far, E]
Phase 2: Init (TensorReplacer)
───────────────────────────────
Flatten step-wise captures → position-indexed tensors
- Prefill: keep all positions from step 0
- Decode: take last token from each step (new position only)
- Result: {prompt_idx: {module: Tensor[total_positions, E]}}
Phase 3: Runtime Injection (per forward pass)
──────────────────────────────────────────────
Model runner Model code
──────────── ──────────
build_context() get_replacement_tensor("...router")
→ index flat tensor at → returns tensor or None
scheduler positions if not None:
→ zero-pad to bucket size use as router logits override
set_active_context(ctx)
set_active_context(None)
Phase 4: Validation
───────────────────
logit_validation() with replacement active
- Token divergence retries work automatically (build_context()
indexes the same flat tensor at new scheduler positions)
Data Flow#
Reference capture: Capture reference tensors from the source model (e.g., monkey-patch HF router forward to record pre-top-k logits). Result:
List[Dict[module_name, List[step_tensor]]]— one entry per prompt, step 0 = prefill, steps 1+ = decode.Flatten at init:
TensorReplacer.__init__flattens step-wise captures into a single[total_positions, feature_dim]tensor per (prompt, module). Prefill positions come from step 0; each decode step contributes its last token (HF recomputes the full sequence each step, so only the last row is new).Build context per forward pass:
build_context()takes the scheduler’sreq_ids,positions, andis_prefillflag, then indexes into the flat reference tensor at each position. Zero-pads padded slots. Returns{module_name: Tensor[bucket_size, E]}.Inject via global context:
get_replacement_tensor()retrieves the tensor from the active context. When present, the replacement logits directly override the router’s computed logits.
Prompt Index Resolution#
Each prompt’s token IDs are hashed at init. When a new request arrives, the model runner calls register_request(req_id, prompt_token_ids) which hashes the request’s tokens at each known prompt length and selects the longest verified match:
Init (from prompt_token_ids):
prompt 0: [10, 20, 30] → hash([10,20,30]) → idx 0
prompt 1: [40, 50] → hash([40,50]) → idx 1
New request arrives with tokens [10, 20, 30, 77, 88]:
try [:3] → hash([10,20,30]) → matches prompt 0, best_len=3
try [:2] → hash([10,20]) → no match
→ longest match is prompt 0 ✓
→ store "req-abc" → 0
Decode batch ["req-abc", "req-def"]:
_resolve_prompt_idx("req-abc") → 0
_resolve_prompt_idx("req-def") → 1
Longest-Match-Wins Strategy#
Prefix matching is required because logit_validation’s teacher-forcing loop re-prompts vLLM with extended tokens (original prompt + generated tokens) as a new request. However, naively returning on the first hash match can cause silent wrong-matches when one stored prompt is a prefix of another.
The algorithm tries all stored prompt lengths, verifies the actual token content (not just the hash), and picks the longest match:
Stored prompts:
prompt 0: [1, 2, 3] (length 3)
prompt 1: [1, 2, 3, 4, 5] (length 5)
Request with [1, 2, 3, 4, 5]:
try [:3] → [1,2,3] matches prompt 0, best_len=3
try [:5] → [1,2,3,4,5] matches prompt 1, best_len=5 (wins)
→ maps to prompt 1 ✓
Retry of prompt 0 with [1, 2, 3, 77, 88]:
try [:3] → [1,2,3] matches prompt 0, best_len=3
try [:5] → [1,2,3,77,88] → no match (hash differs from prompt 1)
→ maps to prompt 0 ✓
Request with [1, 2, 3, 4] (shorter than prompt 1):
try [:3] → [1,2,3] matches prompt 0, best_len=3
plen=5 > len([1,2,3,4])=4 → skipped
→ maps to prompt 0 ✓
Known Corner Case#
The only remaining ambiguity occurs when a retry of a shorter prompt generates tokens that exactly reproduce a longer stored prompt:
Stored prompts:
prompt 0: [1, 2, 3]
prompt 1: [1, 2, 3, 4, 5]
Retry of prompt 0 where model generated [4, 5]:
incoming = [1, 2, 3, 4, 5]
→ longest match picks prompt 1 (length 5) instead of prompt 0
This requires the model to coincidentally generate the exact tokens that complete another stored prompt — unlikely in practice.
Core Components#
TensorReplacer — Context Builder#
Holds flattened reference tensors in memory. Called per forward pass to build replacement tensors from the scheduler’s metadata.
from vllm_neuron.accuracy.tensor_replacement import TensorReplacer
replacer = TensorReplacer(reference_captures, prompt_token_ids=prompt_token_ids)
# Warmup: zero tensors with correct shape for torch.compile tracing
ctx = replacer.warmup_context(num_tokens=256, device=device)
# Inference: build from scheduler metadata
ctx = replacer.build_context(req_ids, positions, is_prefill, device)
build_context() internals:
Prefill: single request per forward pass. Iterate positions in order; each position indexes into the flat reference tensor. First duplicate position marks padding start — remaining slots stay zero.
Decode: multiple requests batched. For each request, look up prompt index via
_resolve_prompt_idx(req_id)(populated byregister_requestat prefill time), then index that prompt’s flat reference tensor at the request’s position. Slots beyond real requests stay zero (batch padding).
set_active_context / get_replacement_tensor — Global Context#
The model runner sets the context before each forward pass and clears it after. Model code calls get_replacement_tensor() to retrieve tensors by module name. Returns None when replacement is not configured.
# Model runner — before each forward
set_active_context(replacer.build_context(req_ids, positions, is_prefill, device))
model(**model_kwargs)
set_active_context(None)
# Inside model — during forward
replacement = get_replacement_tensor("model.layers.0.mlp.router")
# Returns None when not configured → replacement is skipped
Model Integration#
Tensor replacement requires model-specific integration to inject replacement tensors at the correct point in the computation. This can be done by adding get_replacement_tensor() calls in the model’s forward methods.
Integration Pattern 1: No 1:1 Mapping (Fused Downstream)#
When the downstream expects a different format than raw logits (e.g., the CTE prefill kernel expects pre-computed expert affinities), the injection point must replicate the post-logits pipeline to convert replacement logits into the expected format:
# Downstream expects expert_affinities [T, E], not raw logits
if replacement_logits is not None:
# Replicate the routing pipeline: TopK → Softmax → Scatter
top_values, top_indices = torch.topk(
replacement_logits.to(torch.float32),
self.num_experts_per_token, dim=-1,
)
top_values = torch.nn.functional.softmax(top_values, dim=-1)
expert_affinities = torch.zeros(
replacement_logits.shape[0], self.total_num_experts,
device=hidden_states.device, dtype=torch.float32,
).scatter_(1, top_indices, top_values)
else:
expert_affinities = NF.router(hidden_states=hidden_states, ...)
Integration Pattern 2: Direct Interception Point#
When the downstream function already accepts raw logits as a parameter, the replacement tensor can be passed directly without replicating any pipeline logic:
# Add replacement tensor as an optional argument
def _run_moe_block_tkg(
self,
hidden_states: torch.Tensor,
rank: torch.Tensor,
router_logits_override: torch.Tensor | None = None,
):
...
# At the call site:
replacement_logits = get_replacement_tensor(f"model.layers.{self.layer_idx}.mlp.router")
output = self._run_moe_block_tkg(
hidden_states, rank,
router_logits_override=replacement_logits,
)
Layer Index Requirement#
Each module that uses get_replacement_tensor() must know its layer index to construct the correct tensor name. This is typically stored at init:
class GptOssExperts(nn.Module):
def __init__(self, config, layer_idx: int = 0):
super().__init__()
self.layer_idx = layer_idx
...
def forward_decode(self, hidden_states, rank):
# layer_idx identifies which layer's replacement to fetch
replacement = get_replacement_tensor(
f"model.layers.{self.layer_idx}.mlp.router"
)
Token Divergence Retry#
During logit validation with teacher forcing, Neuron may pick a different token than the reference. logit_validation() handles this by extending input_ids with the reference token and retrying.
build_context() is called each forward pass with the scheduler’s current positions, and the flat reference tensor covers all positions up to prompt_len + max_tokens, so the extended positions map directly to existing rows in the flat tensor.
Why Prefix Matching Is Needed#
Each teacher-forcing retry calls llm.generate() with extended tokens, which creates a brand new vLLM request with a new req_id. The scheduler calls register_request(new_req_id, extended_tokens) where extended_tokens = original_prompt + generated_tokens. Since the incoming tokens are longer than what was stored at init, exact matching would fail with KeyError. Prefix matching solves this by slicing the incoming tokens to each known prompt length.
Iteration 1:
llm.generate([{"prompt_token_ids": [1, 2, 3]}])
→ req_id="abc-123"
→ register_request("abc-123", [1, 2, 3]) → exact match prompt 0
Model diverges at token 2, teacher forcing extends:
input_ids = [1, 2, 3, 77, 88]
Iteration 2:
llm.generate([{"prompt_token_ids": [1, 2, 3, 77, 88]}])
→ req_id="def-456" (new request, new ID)
→ register_request("def-456", [1, 2, 3, 77, 88])
→ prefix [:3] = [1,2,3] matches prompt 0 ✓
Note: This is distinct from vLLM scheduler preemption/resume, where the req_id stays the same and register_request is never called again (preempted requests resume via scheduled_cached_reqs.resumed_req_ids, not scheduled_new_reqs).
vLLM Configuration#
Offline (in-memory captures):
from vllm_neuron.model.neuron_config import TensorReplacementConfig
llm = LLM(
model="openai/gpt-oss-20b",
additional_config={
"neuron_config": {
"tensor_replacement": TensorReplacementConfig(
tensors=reference_captures,
prompt_token_ids=all_input_ids,
),
}
},
)
Online (captures and token IDs loaded from disk):
# Save captures and token IDs to disk
torch.save(reference_captures, "/tmp/tensors.pt")
torch.save(all_input_ids, "/tmp/prompt_token_ids.pt")
# Server CLI: pass paths in JSON config
server = start_server(f"""
vllm serve {model_checkpoint}
--additional-config '{{
"neuron_config": {{
"tensor_replacement": {{
"tensors_path": "/tmp/tensors.pt",
"prompt_token_ids_path": "/tmp/prompt_token_ids.pt"
}}
}}
}}'
""")
Speculative Decoding Support#
Tensor replacement works with speculative decoding. During spec decode verification, the target model processes 1 + num_spec_tokens positions per request in a single decode forward pass. build_context() handles this via the tokens_per_req parameter.
Slot-to-Request Mapping#
In standard decode, each slot maps 1:1 to a request. With speculative decoding, each request occupies tokens_per_req consecutive slots (1 verified token + N speculative tokens). The mapping is:
req_idx = slot_idx // tokens_per_req
This allows build_context() to determine which request owns each slot and index into the correct prompt’s flat reference tensor:
Standard decode (tokens_per_req=1, 2 requests, 4 slots with padding):
positions = [5, 3, 0, 0]
slot 0 → req_idx 0//1=0 → position 5 → prompt_0_tensor[5]
slot 1 → req_idx 1//1=1 → position 3 → prompt_1_tensor[3]
slot 2 → req_idx 2//1=2 → 2 >= len(req_ids)=2 → padding (zero)
slot 3 → req_idx 3//1=3 → 3 >= len(req_ids)=2 → padding (zero)
Decode with spec decode (tokens_per_req=4, 1 request):
positions = [8, 9, 10, 11] (verifying 3 draft tokens + 1 bonus)
slot 0 → req_idx 0//4=0 → position 8 → prompt_0_tensor[8]
slot 1 → req_idx 1//4=0 → position 9 → prompt_0_tensor[9]
slot 2 → req_idx 2//4=0 → position 10 → prompt_0_tensor[10]
slot 3 → req_idx 3//4=0 → position 11 → prompt_0_tensor[11]
Decode with spec decode (tokens_per_req=4, 2 requests):
positions = [8, 9, 10, 11, 5, 6, 7, 8]
slot 0 → req_idx 0//4=0 → position 8 → prompt_0_tensor[8]
slot 1 → req_idx 1//4=0 → position 9 → prompt_0_tensor[9]
slot 2 → req_idx 2//4=0 → position 10 → prompt_0_tensor[10]
slot 3 → req_idx 3//4=0 → position 11 → prompt_0_tensor[11]
slot 4 → req_idx 4//4=1 → position 5 → prompt_1_tensor[5]
slot 5 → req_idx 5//4=1 → position 6 → prompt_1_tensor[6]
slot 6 → req_idx 6//4=1 → position 7 → prompt_1_tensor[7]
slot 7 → req_idx 7//4=1 → position 8 → prompt_1_tensor[8]
Slots beyond len(req_ids) * tokens_per_req are padding and stay zero.
The model runner passes tokens_per_req from attn_metadata["decode_token_threshold"] (which equals 1 + num_speculative_tokens when spec decode is active). Without spec decode, tokens_per_req defaults to 1 and the mapping reduces to the standard 1:1 slot-to-request behavior.
This document is relevant for: Trn2, Trn3