[Don't Review] Align BF16 prefill and decode arithmetic in deterministic mode

docs/source/features/attention.md

@@ -34,6 +34,50 @@ The TRT-LLM backend, `TrtllmAttention`, serves as the default backend and suppor
 1. **Fused QKV Input**: It can accept a single QKV tensor as input, which is more efficient compared to using separate Q, K, and V tensors.
 2. **FP8 Output**: It supports outputting the attention result in FP8 format, fusing quantization into the attention computation process.
 
+## Deterministic prefill and generation arithmetic
+
+For BF16 causal self-attention on SM100/SM103, set
+`FORCE_ATTENTION_KERNEL_DETERMINISTIC=1` before constructing the engine to use
+the `deterministic` FMHA library. It uses the same FP32 reduction for each query
+in context and generation, independent of the number of queries or unrelated
+requests in the batch. It reuses the ordinary paged-KV preprocessing, RoPE and
+cache updates. The library participates in the default FMHA selection order;
+an explicit `TLLM_FMHA_LIBS` list must include `deterministic` to select it.
+
+This implementation supports BF16 Q/K/V and KV cache, head dimensions 64, 128
+and 256, 32- or 64-token KV pages, at most 32 query heads per KV head,
+a full causal attention window, and beam width 1.
+
+`FORCE_DETERMINISTIC=1` also selects this attention library. In the PyTorch
+backend it additionally uses a fixed K reduction for unquantized BF16 linear
+layers, including local tensor-parallel shards, and for `Qwen3NextGate` and
+`DeepseekV3Gate` router projections. The DeepSeek V3 gate, also used by
+Nemotron-H, retains FP32 router output. BF16 CUTLASS MoE on
+SM100/SM103 uses a common fallback GEMM tactic across token counts and disables
+fused finalize. This keeps the expert GEMM and activation arithmetic identical
+when an autotuning cache contains different tactics for prefill and decode.
+
+With `allreduce_strategy="NCCL"`, BF16, FP16 and FP32 tensor-parallel reductions
+gather rank contributions and add them in rank order using FP32 accumulators.
+Unquantized RMS normalization and residual RMS normalization use a fixed
+per-row reduction as well. The gather needs a temporary buffer with one input
+copy per rank and increases communication volume. Quantized fused epilogues
+and other collective strategies retain their existing implementations.
+
+GDN and Mamba2 use shared recurrent prefill/decode arithmetic with FP32 SSM
+state. Set `kv_cache_config.mamba_ssm_cache_dtype="float32"` to use those paths.
+Mamba2 also uses the same causal convolution implementation in both phases.
+Lower precision SSM state and speculative decoding retain their existing
+implementations. The attention-only flag does not change GEMM or recurrent
+arithmetic, so it cannot by itself make an entire model's logits invariant.
+
+The fixed reductions trade throughput for reproducibility, particularly for
+long prefills. Benchmark the intended sequence lengths before enabling them
+in a serving configuration. These paths do not establish whole-model
+equivalence for other quantized expert GEMMs, collective strategies outside the
+supported NCCL paths or recurrent architectures beyond GDN and Mamba2. Compare the same
+token sequence when checking prefill and generation logprobs.
+
 ## Implement a New Attention Backend
 
 You can implement a new attention backend to integrate other attention libraries.

examples/llm-api/README.md

@@ -2,6 +2,61 @@
 
 Please refer to the [official documentation](https://nvidia.github.io/TensorRT-LLM/llm-api/) including [customization](https://nvidia.github.io/TensorRT-LLM/examples/customization.html) for detailed information and usage guidelines regarding the LLM API.
 
+## Measure prefill/decode logprob consistency
+
+```bash
+python3 llm_prefill_decode_consistency.py \
+    --model /path/to/checkpoint --tp-size 1 \
+    --prompt-lengths 127 128 129 511 512 513 --max-tokens 64 \
+    --output /tmp/prefill-decode.json
+```
+
+The probe generates once and scores the exact prompt and response tokens with
+full prefill twice. It normalizes both paths' logits with the same FP32 CPU
+`log_softmax`; the API-reported logprobs are retained as a separate normalization
+check. It saves token IDs, per-token scores, error percentiles, and tail fractions.
+The first response token comes from the initial prefill and is reported separately
+from the subsequent decode tokens. Repeated prefill scoring measures prefill
+repeatability; this probe does not measure batch invariance or decode determinism.
+
+The baseline uses one request at a time, eager execution, NCCL allreduce, FP32 SSM
+state, no prefix reuse, no chunked prefill, and no speculative decoding. Extra LLM
+options can be supplied through `--config options.json`; options that break the
+controlled comparison are rejected. Cache capacity is based on available memory;
+hybrid KV/SSM models can need substantially more cache memory than a token-only
+estimate suggests. Check runtime logs for reductions of the requested sequence
+limit when overriding the cache budget. Prompt text is encoded without a chat template
+and repeated/truncated to the specified lengths. For custom inputs, use `--prompts`
+with a JSON list of objects containing `id` and either `prompt` or `prompt_token_ids`.
+
+This is a developer diagnostic, with no assumed numerical tolerance. Once a model
+and platform have a measured baseline, `--max-abs-dlogprob VALUE` makes exceeding
+that bound on any decode token fail the run. A report whose status is `running`
+is incomplete and must not be treated as a passing result. GPU results depend on
+the checkpoint, precision, backends, and hardware recorded in the report.
+
+The CPU scoring tests and GPU regressions live in
+`tests/unittest/llmapi/test_prefill_decode_consistency*.py` and
+`test_prefill_decode_cache_consistency.py`. The GPU tests build a fixed, two-layer
+BF16 Llama checkpoint locally. They check scoring alignment and repeated prefill,
+then compare fixed continuations across whole/chunked prefill, cold/warm prefix
+caches, unrelated intervening requests, and eager/CUDA graph execution. Cache
+hits and actual context calls are checked so these paths cannot silently go
+untested. Default-path tests use a tolerance specific to that small checkpoint.
+On SM100/SM103, deterministic-mode tests select the TRTLLM attention backend and
+require exact score equality across chunking and CUDA graphs. The executor
+disables prefix reuse under `FORCE_DETERMINISTIC=1`; actual prefix hits are
+covered separately by the default-mode tests.
+
+For the supported BF16, TP=1 runtime paths, enable `FORCE_DETERMINISTIC=1` before
+starting the engine and select `attn_backend="TRTLLM"`. GDN and Mamba2 also require
+`kv_cache_config.mamba_ssm_cache_dtype="float32"`. The mode fixes GEMM reductions,
+causal attention, recurrent state updates and BF16 CUTLASS expert tactics across
+prefill and decode. See the [deterministic attention documentation](../../docs/source/features/attention.md#deterministic-prefill-and-generation-arithmetic)
+for attention shape limits. These computations can substantially slow long
+prefills; benchmark the intended workload. This mode does not establish a
+whole-model equality guarantee for quantization, MLA or multi-GPU reductions.
+
 
 ## Run the advanced usage example script: