ITADN

KvCacheManagerV2 segfault on ~256k prompts with Nemotron-3-Ultra when max_seq_len=1M and avg_seq_len unset

#17926Opendsingal0 创建于 10 小时前
KV-Cache ManagementPytorch
D
dsingal0commented
## 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 条评论