Add DFlash spec decode support for NemotronH (Nemotron Ultra 550B)

5 patches to enable DFlash speculative decoding with NemotronH models:

1. nemotron_h.py: Add DFLASH_AUX_HIDDEN_STATE_LAYERS env var hook in
   NemotronHModel.__init__ to auto-set aux hidden state layers during
   extraction/inference. Modify NemotronHForCausalLM.forward() to return
   (hidden_states, aux_hidden_states) tuple for DFlash spec decode mode,
   and save to disk for offline extraction mode (DFLASH_EXTRACT_PATH).

2. qwen2.py: Add relu2 activation support in Qwen2MLP. NemotronH uses
   relu2 (relu squared) instead of silu for MLP activation. Without this,
   DFlash draft model loading fails with 'Unsupported activation: relu2'.

3. flex_attention.py: Fix .view() -> .reshape() for non-contiguous KV
   cache tensors. Required for FLEX_ATTENTION backend with paged KV cache.

4. input_processor.py: Fix tensor boolean ambiguity. NemotronH prompts
   can produce multi-element tensors that fail 'if prompt_ids:' check.

5. gpu_model_runner.py: Add debug logging for aux hidden state tuple
   unpacking in _dummy_run. Helps diagnose DFlash spec decode issues.

Tested with Nemotron Ultra 550B (NemotronH hybrid Mamba2+MoE+attention):
- Offline hidden state extraction: 5000 samples, 656GB
- DFlash draft training: ~4B params, val/pos_1_acc=56.2%
- Inference: 1.66x speedup (233.85 vs 141.18 tok/s) with FLASH_ATTN FA2
- Beats built-in MTP (1.36x / 194.25 tok/s) by 22%
- Lossless: 2/5 samples bit-identical (expected bf16 noise)

Note: B200 (SM100) defaults to FA4 which requires cutlass.cute (Python 3.10
MLIR). Force flash_attn_version=2 via attention_config for DFlash spec decode.

Based on vLLM v0.20.1.

Author: xrsrke

vllm/model_executor/models/nemotron_h.py

@@ -1,6 +1,12 @@
+import os
+
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
 # Adapted from https://github.com/vllm-project/vllm/blob/94d8ec8d2bcb4ec55e33022b313c7e978edf05e1/vllm/model_executor/models/bamba.py
 # Copyright 2024 HuggingFace Inc. team. All rights reserved.
 # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
@@ -64,9 +70,11 @@
     maybe_remap_kv_scale_name,
 )
 from vllm.model_executor.models.interfaces import (
+    EagleModelMixin,
     HasInnerState,
     IsHybrid,
     MixtureOfExperts,
+    SupportsEagle3,
     SupportsLoRA,
     SupportsMambaPrefixCaching,
     SupportsPP,
@@ -539,7 +547,7 @@ def forward(
 
 
 @support_torch_compile
-class NemotronHModel(nn.Module):
+class NemotronHModel(nn.Module, EagleModelMixin):
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
 
@@ -578,6 +586,14 @@ def get_layer(prefix: str):
         self.start_layer, self.end_layer, self.layers = make_layers(
             len(config.hybrid_override_pattern), get_layer, prefix=f"{prefix}.layers"
         )
+
+        # DFlash offline extraction: auto-set aux_hidden_state_layers from env var
+        # This bypasses the need for apply_model() which has serialization issues
+        _dflash_layers = os.environ.get("DFLASH_AUX_HIDDEN_STATE_LAYERS", "")
+        if _dflash_layers:
+            _layers = tuple(int(x) for x in _dflash_layers.split(","))
+            self._set_aux_hidden_state_layers(_layers)
+            logger.info("DFlash: auto-set aux_hidden_state_layers=%s from env", _layers)
         self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
             ["hidden_states", "residual"], config.hidden_size
         )
@@ -605,18 +621,27 @@ def forward(
             hidden_states = intermediate_tensors["hidden_states"]
             residual = intermediate_tensors["residual"]
 
-        for layer in islice(self.layers, self.start_layer, self.end_layer):
+        # EAGLE3 auxiliary hidden state collection (for DFlash spec decoding)
+        aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)
+        for idx, layer in enumerate(
+            islice(self.layers, self.start_layer, self.end_layer)
+        ):
             hidden_states, residual = layer(
                 positions=positions,
                 hidden_states=hidden_states,
                 residual=residual,
             )
+            self._maybe_add_hidden_state(
+                aux_hidden_states, idx + 1, hidden_states, residual
+            )
 
         if not get_pp_group().is_last_rank:
             return IntermediateTensors(
                 {"hidden_states": hidden_states, "residual": residual}
             )
         hidden_states, _ = self.norm_f(hidden_states, residual)
+        if len(aux_hidden_states) > 0:
+            return hidden_states, aux_hidden_states
         return hidden_states
 
     def is_spec_layer(self, config: NemotronHConfig, weight_name: str) -> bool:
@@ -771,6 +796,7 @@ class NemotronHForCausalLM(
     SupportsQuant,
     MixtureOfExperts,
     SupportsMambaPrefixCaching,
+    SupportsEagle3,
 ):
     # Relevant only if self.has_moe is True
     is_non_gated_moe: bool = True
@@ -919,11 +945,42 @@ def forward(
         intermediate_tensors: IntermediateTensors | None = None,
         inputs_embeds: torch.Tensor | None = None,
         **kwargs,
-    ):
-        hidden_states = self.model(
+    ) -> torch.Tensor | IntermediateTensors:
+        model_output = self.model(
             input_ids, positions, intermediate_tensors, inputs_embeds
         )
 
+        # Handle EAGLE3 aux_hidden_states (for offline extraction and DFlash spec decode)
+        if isinstance(model_output, tuple):
+            hidden_states, aux_hidden_states = model_output
+
+            # Extraction mode: save to disk (rank 0 only)
+            _extract_path = os.environ.get("DFLASH_EXTRACT_PATH", "")
+            if _extract_path:
+                import torch.distributed as dist
+                if not dist.is_initialized() or dist.get_rank() == 0:
+                    import safetensors.torch
+                    _num_tokens = hidden_states.shape[0]
+                    stacked = torch.stack(
+                        [t[:_num_tokens].detach().cpu() for t in aux_hidden_states], dim=1
+                    )
+                    _batch_id = id(input_ids) if input_ids is not None else id(hidden_states)
+                    _filename = os.path.join(_extract_path, f"batch_{_batch_id}.safetensors")
+                    if input_ids is not None:
+                        _token_ids = input_ids[:_num_tokens].detach().cpu().reshape(-1)
+                    else:
+                        _token_ids = torch.tensor([], dtype=torch.long)
+                    safetensors.torch.save_file({
+                        "hidden_states": stacked, "token_ids": _token_ids,
+                    }, _filename)
+                # Extraction mode: return just hidden_states
+                return hidden_states
+
+            # DFlash spec decode mode: return the tuple so model runner can extract aux_hidden_states
+            return model_output
+        else:
+            hidden_states = model_output
+
         return hidden_states
 
     def compute_logits(

vllm/model_executor/models/qwen2.py

@@ -104,11 +104,21 @@ def __init__(
             quant_config=quant_config,
             prefix=f"{prefix}.down_proj",
         )
-        if hidden_act != "silu":
+        if hidden_act == "silu":
+            self.act_fn = SiluAndMul()
+        elif hidden_act == "relu2":
+            # Relu2AndMul: relu²(gate) * up, same split as SiluAndMul
+            import torch.nn.functional as F
+            class Relu2AndMul:
+                def __call__(self, x):
+                    d = x.shape[-1] // 2
+                    gate, up = x[..., :d], x[..., d:]
+                    return F.relu(gate).pow(2) * up
+            self.act_fn = Relu2AndMul()
+        else:
             raise ValueError(
-                f"Unsupported activation: {hidden_act}. Only silu is supported for now."
+                f"Unsupported activation: {hidden_act}. Only silu and relu2 are supported."
             )
-        self.act_fn = SiluAndMul()
 
     def forward(self, x):
         gate_up, _ = self.gate_up_proj(x)

vllm/v1/attention/backends/flex_attention.py

@@ -1098,8 +1098,8 @@ def forward(
             key_cache, value_cache = kv_cache.unbind(0)
 
             # View out the block_size dim
-            key_cache = key_cache.view(-1, self.num_kv_heads, self.head_size)
-            value_cache = value_cache.view(-1, self.num_kv_heads, self.head_size)
+            key_cache = key_cache.reshape(-1, self.num_kv_heads, self.head_size)
+            value_cache = value_cache.reshape(-1, self.num_kv_heads, self.head_size)
             query, key_tensor, value_tensor = map(
                 lambda x: self.view_as_4d(x).permute(0, 2, 1, 3),
                 (query, key_cache, value_cache),

vllm/v1/engine/input_processor.py

@@ -455,7 +455,7 @@ def _validate_model_input(
                             f"by setting --limit-mm-per-prompt at startup."
                         )
 
-        if prompt_ids and tokenizer is not None:
+        if prompt_ids is not None and len(prompt_ids) > 0 and tokenizer is not None:
             max_input_id = max(prompt_ids, default=0)
 
             # NOTE: tokenizer.max_token_id is the tokenizer’s vocab size while

vllm/v1/worker/gpu_model_runner.py

@@ -5543,7 +5543,25 @@ def _dummy_run(
                 )
 
             if self.use_aux_hidden_state_outputs:
-                hidden_states, _ = outputs
+                # Debug: check what outputs actually is
+                if not isinstance(outputs, tuple):
+                    print(f"DEBUG _dummy_run: use_aux=True but outputs is {type(outputs)}, not tuple", flush=True)
+                    hidden_states = outputs
+                elif len(outputs) != 2:
+                    print(f"DEBUG _dummy_run: outputs has {len(outputs)} elements, expected 2. Types: {[type(o) for o in outputs]}", flush=True)
+                    hidden_states = outputs[0] if len(outputs) > 0 else outputs
+                else:
+                    hs, aux = outputs
+                    print(f"DEBUG _dummy_run: outputs is 2-tuple, hs type={type(hs)}, aux type={type(aux)}", flush=True)
+                    if not isinstance(hs, torch.Tensor):
+                        print(f"DEBUG _dummy_run: hs is NOT a tensor! Trying outputs[0]... type={type(outputs[0])}", flush=True)
+                        # Maybe the model returns (aux_list, hidden_states) instead
+                        if isinstance(aux, torch.Tensor):
+                            hidden_states = aux
+                        else:
+                            hidden_states = hs[0] if isinstance(hs, list) and len(hs) > 0 else hs
+                    else:
+                        hidden_states = hs
             else:
                 hidden_states = outputs