KvCacheManagerV2 segfault on ~256k prompts with Nemotron-3-Ultra when max_seq_len=1M and avg_seq_len unset
KV-Cache ManagementPytorch
## Summary
Serving `nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-NVFP4` with `trtllm-serve`, `TLLM_ALLOW_LONG_MAX_MODEL_LEN=1`, and `--max_seq_len 1048576` **segfaults inside `KvCacheManagerV2`** on a ~256k-token mid-document needle-in-a-haystack request when `kv_cache_config.avg_seq_len` is **not** set.
With the same setup plus `avg_seq_len: 9000`, the same server successfully completes **800k-token** mid-needle retrieval (`ORANGE-MANGO-42`).
This makes the documented long-context Ultra path unreliable: HF / deployment docs say to set `TLLM_ALLOW_LONG_MAX_MODEL_LEN=1` and `--max_seq_len <seq_len>`, but do not require `avg_seq_len`. Without it, TRT-LLM warns and falls back to `max_seq_len / 2` (= 524288), then crashes.
## Environment
- Image: `nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc25.dev2026081200`
- TensorRT-LLM: `1.3.0rc25.dev2026081200`
- Hardware: 4× NVIDIA B200 (TP=4, EP=4)
- Checkpoint: `nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-NVFP4`
- Backend: PyTorch (`MambaHybridCacheManagerV2`)
## Repro (stock `trtllm-serve` only)
### 1. Config that crashes (`extra-llm-api-config.yml`)
```yaml
backend: pytorch
enable_chunked_prefill: true
disable_overlap_scheduler: true
enable_attention_dp: false
max_batch_size: 8
max_num_tokens: 16384
trust_remote_code: true
cuda_graph_config:
enable_padding: true
max_batch_size: 8
kv_cache_config:
dtype: fp8
free_gpu_memory_fraction: 0.8
enable_block_reuse: true
tokens_per_block: 64
# avg_seq_len intentionally omitted
mamba_ssm_cache_dtype: float16
mamba_ssm_stochastic_rounding: true
mamba_ssm_philox_rounds: 5
mamba_state_config:
periodic_snapshot_interval: 256
additional_snapshot_offsets_from_end:
- 0
moe_config:
backend: TRTLLM
use_low_precision_moe_combine: true
```
### 2. Launch
```bash
export TLLM_ALLOW_LONG_MAX_MODEL_LEN=1
MODEL=nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-NVFP4
trtllm-serve serve "$MODEL" \
--host 0.0.0.0 --port 8000 \
--tp_size 4 --ep_size 4 \
--max_batch_size 8 \
--max_num_tokens 16384 \
--max_seq_len 1048576 \
--trust_remote_code \
--reasoning_parser nemotron-v3 \
--tool_parser qwen3_coder \
--config ./extra-llm-api-config.yml
```
### 3. Client
```python
# save as repro_needle.py
import json, time, urllib.request, urllib.error
URL = "http://127.0.0.1:8000/v1/chat/completions"
MODEL = "nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-NVFP4"
NEEDLE = "The secret code is ORANGE-MANGO-42."
FILLER = "alphabet "
def prompt(n_words: int) -> str:
left, right = n_words // 2, n_words - n_words // 2
hay = FILLER * left + NEEDLE + " " + FILLER * right
return (
"Find the secret code in the document. "
"Reply with ONLY the code string, nothing else.\n\nDOCUMENT:\n" + hay
)
def once(n_words: int):
body = {
"model": MODEL,
"messages": [{"role": "user", "content": prompt(n_words)}],
"max_tokens": 64,
"temperature": 0,
"chat_template_kwargs": {"enable_thinking": False},
}
t0 = time.time()
req = urllib.request.Request(
URL, data=json.dumps(body).encode(),
headers={"Content-Type": "application/json"},
)
try:
with urllib.request.urlopen(req, timeout=1800) as r:
d = json.loads(r.read())
c = d["choices"][0]["message"].get("content")
u = d.get("usage") or {}
print(n_words, "OK", u.get("prompt_tokens"), repr(c), "sec", round(time.time()-t0,1))
except urllib.error.HTTPError as e:
print(n_words, "HTTP", e.code, e.read()[:400])
once(200_000) # succeeds
once(256_000) # segfault / engine dead
```
## Observed results
| Prompt size | `avg_seq_len` unset | `avg_seq_len: 9000` |
| --- | --- | --- |
| ~200k mid-needle | HIT `ORANGE-MANGO-42` | HIT |
| ~256k mid-needle | **engine segfault / MPI worker dies** | HIT |
| ~800k mid-needle | n/a (server already dead) | HIT `ORANGE-MANGO-42` (`prompt_tokens=800053`, ~323s) |
### Crash signature (`avg_seq_len` unset)
Startup warning:
```text
'kv_cache_config.avg_seq_len' is not set for a hybrid Mamba model using KV cache manager V2.
Falling back to max_seq_len / 2=524288 for cache-pool sizing.
```
Then on the ~256k request:
```text
!!!!!!! Segfault encountered !!!!!!!
File "<unknown>", line 0, in tensorrt_llm::batch_manager::kv_cache_manager_v2::Block::tokensPerBlock() const
File "<unknown>", line 0, in tensorrt_llm::batch_manager::kv_cache_manager_v2::addOrGetExistingBlock(...)
File "<unknown>", line 0, in tensorrt_llm::batch_manager::kv_cache_manager_v2::KvCache::_commitBlock(...)
File "<unknown>", line 0, in tensorrt_llm::batch_manager::kv_cache_manager_v2::KvCache::commit(...)
...
Fatal engine error recorded: RuntimeError('MPI worker rank 0 ... exited unexpectedly')
EngineDeadError: Engine has died: RuntimeError: MPI worker rank 0 (pid ...) exited unexpectedly
```
Client then sees:
```text
HTTP 400 {"message":"Engine has died: RuntimeError: MPI worker rank 0 (pid ...) exited unexpectedly", ...}
```
## Workaround
Set an explicit workload-average sequence length, e.g.:
```yaml
kv_cache_config:
avg_seq_len: 9000
# ...rest unchanged...
```
With that single change, 256k and 800k mid-needle requests succeed on the same image/hardware.
## Expected
1. Documented 1M Ultra serving should not segfault when `avg_seq_len` is omitted (or should fail fast with a clear validation error instead of corrupting/killing the engine).
2. Ideally `avg_seq_len` fallback for hybrid V2 should not use `max_seq_len/2` when `max_seq_len` has been raised far above the native HF window via `TLLM_ALLOW_LONG_MAX_MODEL_LEN`.
## Notes
- Overlap scheduler was disabled in this repro (`disable_overlap_scheduler: true`) to isolate the KV-cache manager issue.
- MTP was not enabled.
- Tokenizer still logs `Token indices sequence length is longer than the specified maximum sequence length for this model (800053 > 262144)` even when the request succeeds; that may be a separate HF tokenizer `model_max_length` warning, but it did not prevent successful 800k generation once `avg_seq_len` was set.
0 条评论