From 263b5264ebe74a9439ba01f6faa4d18f745cd3fd Mon Sep 17 00:00:00 2001 From: Ruodi Lu Date: Thu, 23 Jul 2026 04:17:15 +0000 Subject: [PATCH] [None][feat] Support DeepSeek-V4 in layer_wise_benchmarks DeepSeek-V4's multi-head hyper-connection (mHC) decoder layers do not fit the generic single-layer harness in tools/layer_wise_benchmarks/runner.py: - create_kv_cache_manager: the is_mla branch omitted vocab_size, which DeepseekV4CacheManager.__init__ requires (keyword-only, no default). - create_run_pack: V4 layers take the initial residual as hc_state shaped [num_tokens, hc_mult, hidden_size] and their MoE routing needs input_ids; both are now synthesized when the config exposes hc_mult. - run_pack(check=True): guarded the NaN/Inf/zero check with isinstance since mHC layers return an HCState, not a Tensor. Non-V4 models are unaffected (all V4 paths gate on hc_mult / MLA class). Usage for V4: --moe-backend DEEPGEMM and a --layer-indices slice that starts at layer 0 and covers both compress ratios (e.g. 0,1,2,3). Signed-off-by: Ruodi Lu --- .../tools/layer_wise_benchmarks/runner.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/tools/layer_wise_benchmarks/runner.py b/tensorrt_llm/tools/layer_wise_benchmarks/runner.py index f63634f463d4..e3bd34d60a7f 100644 --- a/tensorrt_llm/tools/layer_wise_benchmarks/runner.py +++ b/tensorrt_llm/tools/layer_wise_benchmarks/runner.py @@ -661,6 +661,21 @@ def create_run_pack( ) kwargs = {} + # DeepSeek-V4 (multi-head hyper-connection) decoder layers take the initial residual + # as ``hc_state`` shaped ``[num_tokens, hc_mult, hidden_size]`` (not a 2D hidden-states + # tensor), and their MoE routing requires ``input_ids``. Both are absent from the + # generic single-layer harness, so synthesize them when the model exposes ``hc_mult``. + hc_mult = getattr(pretrained_config, "hc_mult", None) + if hc_mult is not None: + hidden_states = hidden_states.unsqueeze(1).expand(-1, hc_mult, -1).contiguous() + kwargs["input_ids"] = torch.randint( + 0, + pretrained_config.vocab_size, + (batch_size * seq_len_q,), + dtype=torch.int32, + device="cuda", + ) + if is_nemotron_hybrid(pretrained_config) or is_qwen3_hybrid(pretrained_config): mamba_metadata = Mamba2Metadata( attn_metadata.max_num_requests, @@ -678,7 +693,7 @@ def run_pack(*, check=False): hidden_states_out, residual_out = self.model( position_ids, hidden_states, attn_metadata, residual, **kwargs ) - if check: + if check and isinstance(hidden_states_out, torch.Tensor): if hidden_states_out.isnan().any(): raise ValueError("Has nan, please fix weights initialization") if hidden_states_out.isinf().any(): @@ -816,6 +831,7 @@ def create_kv_cache_manager( dtype=kv_cache_dtype, spec_config=None, layer_mask=layer_mask, + vocab_size=config.vocab_size, sparse_attention_config=model_config.sparse_attention_config, pretrained_config=model_config.pretrained_config, )