This document is relevant for: Trn2, Trn3

vLLM Neuron Logit Validation Design#

Overview#

The vLLM Neuron logit validation framework provides logit-level validation and visualization for comparing model outputs across different hardware platforms. This enables detection of numerical drift and hardware-specific differences.

Core Components#

Logit Validation Engine#

Purpose: Compare raw logits between reference and target models with configurable tolerances.

API (vllm_neuron.accuracy.logit_validation.logit_validation):

logit_validation(
    input_ids: List[List[int]],
    generate_fn: Callable[
        [torch.Tensor], Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]
    ],
    expected_logits: torch.Tensor,
    tol_map: dict = None,
    baseline_logits: torch.Tensor = None,
    visualize: bool = False,
    save_logits: bool = True,
    output_dir: str = "validation_output",
) -> bool

Visualization System#

Purpose: Generate interactive HTML reports for logit comparison results.

API (vllm_neuron.accuracy.logit_visualization.visualize_logit_results):

visualize_logit_results(
    results: List[List[Dict]],
    output_dir: str = "logit_validation_plots",
    max_tokens: int = 5,
    save_logits: bool = True,
) -> None

Direct Logits Comparison (vllm_neuron.accuracy.logit_visualization.visualize_logits):

visualize_logits(
    target_logits: torch.Tensor,
    baseline_logits: torch.Tensor,
    output_dir: str = "logit_plots",
    tol_map: dict = None,
) -> None

Tolerance Configuration#

Default Tolerances (vllm_neuron.accuracy.constants.DEFAULT_TOLERANCE_MAP), keyed by top-k slice as (atol, rtol):

DEFAULT_TOLERANCE_MAP = {
    "5": (1e-5, 0.011),
    "50": (1e-5, 0.02),
    "1000": (1e-5, 0.03),
    "all": (1e-5, 0.05),
}

Teacher Forcing Strategy#

Both the reference and target models are fed the same reference tokens at every decode step, rather than each sampling its own continuation. This isolates hardware precision effects from context drift: without teacher forcing, a single divergent token early in generation would cascade into a completely different continuation, making a per-token logit comparison meaningless.

File Outputs#

Validation Artifacts:

  • reference_logits.pt: Expected logits tensor

  • target_logits.pt: Actual model outputs

Visualization Files:

  • logit_analysis_b{batch}.html: Interactive plots

  • logit_analysis_summary.json: Results summary

Expected Visualization Output#

Logit Visualization Example

Subgraph Analysis#

Error Map Analysis (Top)

  • Tracks divergence values across top-k tokens for different k values (k=5, k=50, k=1000, k=all)

  • The divergence shows the relative error between target and reference (top-k logits)

  • If there’s token divergence, it will be marked as vertical orange lines

  • X-axis: token id (0-indexed)

  • Y-axis: Divergence value (for different k values)

Relative Difference Analysis (Middle-Top)

  • Measures Top1 logit - Top2 logit difference for both expected and actual models

  • The diff (between Top1 and Top2 logit) needs to be compared against unit of least precision (ULP)

  • When the diff (Top1 - Top2 < ULP), it could lead to wrong token selection due to precision issues.

  • X-axis: token id (0-indexed)

  • Y-axis: Relative difference (always larger than 0)

Top1 Indices Comparison (Middle-Bottom)

  • Compares which tokens are predicted as most likely (Top1) between models

  • Lines show expected (cyan) vs actual (pink) Top1 token indices

  • Diverging lines indicate different token selection

  • X-axis: token id (0-indexed)

  • Y-axis: Token index (0-N vocabulary range)

Top2 Logit Values Analysis (Bottom)

  • Shows the actual logit values for top-2 logit values

  • The prediction values should match between expected and actual prediction. It helps validate numerical precision across hardware platforms

  • X-axis: token id (0-indexed)

  • Y-axis: Logit value (this can be cross checked with the dtype precision graph)

Usage Examples#

Basic validation (visualization disabled by default):

passed = logit_validation(
    input_ids=input_ids,
    generate_fn=target_fn,
    expected_logits=reference_logits,
    visualize=False,  # set to True to also write an HTML plot
    save_logits=True,
    output_dir="validation_results",
)

Direct Logits Visualization#

visualize_logits(
    target_logits=neuron_logits,
    baseline_logits=cpu_logits,
    output_dir="logit_comparison",
)

Multi-Prompt Logit Validation#

Features: Validate logits across multiple prompts with aggregate thresholds.

API (vllm_neuron.accuracy.logit_validation.multi_prompt_logit_validation):

multi_prompt_logit_validation(
    prompts_input_ids: List[List[List[int]]],
    generate_fn: Callable[
        [torch.Tensor], Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]
    ],
    prompts_expected_logits: List[torch.Tensor],
    prompts_baseline_logits: List[torch.Tensor] = None,  # enables three-way mode
    aggregate_config: dict = None,
    tol_map: dict = None,
) -> MultiPromptValidationResult

MultiPromptValidationResult (from vllm_neuron.accuracy.types) has fields:

  • passed: bool — True if all prompts pass two-way validation

  • per_prompt_results: List[Tuple[bool, Optional[List[List[dict]]]]] — per-prompt (passed, results)

  • aggregate_metrics: Dict[str, Any] — cross-prompt aggregate threshold results

Two-Way vs Three-Way Comparison#

Both logit_validation and multi_prompt_logit_validation run in one of two modes, selected by whether a baseline is supplied (baseline_logits / prompts_baseline_logits).

Two-way (no baseline): compares the target (Neuron) logits directly against whatever reference the caller passes as expected_logits, using the fixed top-k tolerance map. For each top-k slice the relative error must stay under that slice’s threshold (the rtol entry of the (atol, rtol) tuple in DEFAULT_TOLERANCE_MAP).

The reference is the caller’s choice. The on-device logit test flow passes the FP32 goldens (fp32_logits) as expected_logits for its two-way check — i.e. it compares Neuron directly against FP32 ground truth.

expected (caller's reference, e.g. HF FP32)  ──► compared against ──►  target (Neuron)
pass if per-top-k relative error < tol_map[k] rtol

Three-way (baseline supplied): compares three models to separate vllm-neuron bugs from dtype-inherent BF16 noise:

  • FP32 (baseline) — HuggingFace FP32 on CPU (ground truth; drives teacher forcing)

  • BF16 (expected) — HuggingFace BF16 on CPU (the dtype-noise reference)

  • Neuron (target) — vllm-neuron BF16 on hardware (under test)

Errors are measured relative to the FP32 baseline, so the target’s error can be compared against the BF16-vs-FP32 error. If the target’s error looks like BF16 noise, it passes even when it exceeds the static two-way thresholds:

base error = |FP32 − HF BF16|      (dtype-inherent noise)
tgt  error = |FP32 − Neuron BF16|  (what we're testing)
pass if the target error is statistically indistinguishable from the baseline

The overall verdict is an OR of three independent routes — it passes if any of these holds:

  • all prompts pass the static two-way top-k thresholds, or

  • aggregate σ-ratio ≤ 1.0 (agg_sigma_ratio_threshold), or

  • aggregate BC ≥ 0.99 (agg_bc_threshold).

So the σ-ratio / BC checks can rescue a test whose errors exceed the fixed two-way thresholds but are still statistically within BF16 noise. The per-token L-inf / L2 ratios (default multiplier 1.5) are reported as diagnostics and for interpreting where a target deviates, but they do not gate the aggregate verdict. Interpreting a failure (e.g. a large token-0 ratio indicating a prefill issue) is covered in Accuracy debugging design → LogitValPlugin.

Example Scripts#

  • examples/vllm_neuron/accuracy/run_logit_validation_offline.py - Offline serving (LLM class)

  • examples/vllm_neuron/accuracy/run_logit_validation_online.py - Online serving (OpenAI API)

Golden Caching#

Computing HuggingFace reference logits (FP32 baseline + BF16 expected) is expensive and deterministic for a given model + prompts + output length. The golden caching system avoids recomputation across test runs.

Architecture#

Two-tier cache: local disk → S3 (optional).

  • vllm_neuron/utils/golden_cache.py: Generic get_or_compute_goldens() with content-addressed keys (hash of key_config dict).

  • vllm_neuron/accuracy/goldens/reference_logits.py: get_cached_reference_goldens() wraps the generic cache with logit-validation-specific defaults (builds the key_config, threads model_loader/FP8 params into generate_three_way_reference_logits). Exported from vllm_neuron.accuracy.goldens. The model_loader parameter lets callers inject a model-specific override loader.

Each cached golden is stored as a directory containing:

  • golden.pt: Tensor data (FP32 logits, BF16 logits, input IDs)

  • metadata.json: Framework version, key config (human-readable)

Cache location: $VLLM_NEURON_GOLDEN_CACHE_DIR (default: ~/.cache/vllm/neuron/goldens).

Usage#

from vllm_neuron.accuracy.goldens import get_cached_reference_goldens

goldens = get_cached_reference_goldens(
    model_id="meta-llama/Llama-3.2-1B-Instruct",
    model_checkpoint=checkpoint_path,
    model_config=config,
    tokenizer=tokenizer,
    prompts=["The capital of France is", ...],
    output_length=16,
)

# goldens["fp32_logits"]  - FP32 baseline (list, one entry per prompt)
# goldens["dtype_logits"] - BF16 expected, teacher-forced (list, one per prompt)
# goldens["input_ids"]    - Tokenized prompts (list, one per prompt)

# Normalize each prompt's token ids to List[List[int]] before validating.
prompts_input_ids = [
    [list(ids) for ids in prompt] for prompt in goldens["input_ids"]
]

result = multi_prompt_logit_validation(
    prompts_input_ids=prompts_input_ids,
    generate_fn=neuron_generate_fn,
    prompts_expected_logits=goldens["dtype_logits"],
    prompts_baseline_logits=goldens["fp32_logits"],
)
print("PASSED" if result.passed else "FAILED")

Key Config#

The cache key is a hash of fields that affect the golden output:

  • model_id, prompts, output_length, dtype

  • kv_cache_dtype (only when not “auto”)

  • Fields like tensor_parallel_size are excluded (they don’t affect reference logits)

See vllm_neuron/utils/golden_cache.py for the full key-hashing and cache-assembly logic.

This document is relevant for: Trn2, Trn3