[Bug]: EAGLE3 draft on a hybrid Mamba target crashes KV cache estimation ("LlamaConfig is not a supported hybrid Mamba config")

System Info

• CPU architecture: aarch64
• GPU: NVIDIA GB200 (2 nodes, 8 ranks, TP8/PP1)
• TensorRT-LLM: 1.3.0rc24 (commit 1cef02e901); code path unchanged on main c7b8ec26fa
• Model: Qwen/Qwen3.5-397B-A17B with an EAGLE3 draft head, served via trtllm-serve
• CUDA / driver / OS: 13.3 / 580.173.02 / Ubuntu
• Container: NGC PyTorch 26.06

Who can help?

No response

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

Serve a hybrid-linear target (Qwen3.5 / Qwen3-Next / Nemotron-hybrid) with a one-model EAGLE3 draft head and attention DP disabled:

# extra_llm_api_config.yaml
enable_attention_dp: false
speculative_config:
  decoding_type: Eagle
  max_draft_len: 3
  speculative_model_dir: AQ-MedAI/Qwen3.5-397B-A17B-eagle3
  eagle3_one_model: true
kv_cache_config:
  dtype: auto

trtllm-serve Qwen/Qwen3.5-397B-A17B --tp_size 8 --extra_llm_api_options extra_llm_api_config.yaml

Startup fails during KV cache sizing:

[executor][RANK 1] Failed to initialize executor on rank 1:
LlamaConfig is not a supported hybrid Mamba config

Call path (line numbers from main c7b8ec26fa):

• _util.py:775 KvCacheCreator._get_kv_size_per_token(), EAGLE3/MTP branch
  -> _per_manager_cache_cost(self._kv_cache_manager_cls, effective_draft_config, ...)
• the target class here is a Mamba hybrid manager, the config is the draft's LlamaConfig
• mamba_cache_manager.py: MambaHybridCacheManager*.get_cache_size_per_token()
  -> _estimate_mamba_hybrid_cache_cost() -> _get_local_mamba_cache_layout()
  -> config_utils.py:437 extract_mamba_kv_cache_params() raises ValueError

Reached only when _should_create_separate_draft_kv_cache() is True, i.e. a one-model spec mode with a separate draft KV cache, no attention DP, on the last PP rank. With enable_attention_dp: true the branch is skipped, which is why attention DP hides the crash.

Expected behavior

KV cache capacity estimation succeeds and costs the draft layers with the KV cache manager class that is actually used for them at runtime, the same way _create_one_model_draft_kv_cache_manager() resolves it.

actual behavior

Estimation costs the draft layers with the target's manager class while passing the draft's model config, so a hybrid Mamba target plus a plain transformer EAGLE3 draft head raises ValueError and the worker dies before any draft KV cache is built.

additional notes

Only the estimate is wrong; the runtime allocation path is already correct:

• _create_one_model_draft_kv_cache_manager() (_util.py:1454 on main) resolves the class from effective_draft_config via get_kv_cache_manager_cls() + _validate_or_fallback_kv_cache_manager_v2().
• The neighboring external-drafter branch in _get_kv_size_per_token() does the
  same and even carries the comment "may differ from target (e.g. hybrid target
  • plain transformer draft)".

Only the EAGLE3/MTP branch still uses self._kv_cache_manager_cls. Resolving the
class there the same way makes the estimate and the allocation agree; that fix
was verified on 1.3.0rc24 by patching site-packages, after which the server
starts and serves.

Suggested change (main, _util.py:775):

elif self._mapping.is_last_pp_rank():
    # EAGLE3/MTP: draft layers only on last PP rank
    draft_kv_cache_manager_cls = get_kv_cache_manager_cls(
        effective_draft_config,
        draft_kv_cache_config,
        is_disagg=self._is_disagg)
    draft_kv_cache_manager_cls = self._validate_or_fallback_kv_cache_manager_v2(
        draft_kv_cache_manager_cls, effective_draft_config,
        draft_kv_cache_config)
    total += self._per_manager_cache_cost(
        draft_kv_cache_manager_cls,
        effective_draft_config,
        draft_kv_cache_config,
        num_layers=self._get_num_draft_layers(),
        is_draft=True)

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.

Start in _util.py:775 at KvCacheCreator._get_kv_size_per_token() and compare the EAGLE3/MTP branch with the neighboring external-drafter branch. Follow _create_one_model_draft_kv_cache_manager() and mamba_cache_manager.py to confirm the manager must match the effective draft config; done means KV sizing succeeds and the supplied trtllm-serve reproduction starts without the ValueError.