neuron: drafter forward as one region, compiled dflash metadata fill; verification + drafts regions on by default

vllm-sonnet 512 (Trn2-Hy4): decode step 660 -> 591 ms, TPOT 221.7 -> 206.5 ms (p50), TTFT 4.02 -> 3.93 s,
tokens/step identical (3.02/2.82/2.65). Tiny golden 9/9.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

deepinfer/environ.py

@@ -145,7 +145,9 @@ class envs:
     # Neuron only: run the sampler's process_sampling as one compiled graph instead of ~1900 eager ops a step.
     DEEPINFER_NEURON_COMPILE_SAMPLER = EnvVar.bool(default=True)
     # The decode verification (greedy batch, no structured output) as one compiled region (issue #29).
-    DEEPINFER_NEURON_COMPILE_VERIFY = EnvVar.bool(default=False)
+    DEEPINFER_NEURON_COMPILE_VERIFY = EnvVar.bool(default=True)
+    # The DSpark draft forward (embedding, draft layers, hc head) as one region instead of one per layer (issue #29).
+    DEEPINFER_NEURON_DRAFTS_REGION = EnvVar.bool(default=True)
     # Neuron only: compile the decode step with fullgraph=True (raise on the first graph break).
     DEEPINFER_NEURON_FULLGRAPH = EnvVar.bool()
 

deepinfer/model_executor/layers/tilelang_ops/neuron/get_kv_index.py

@@ -47,6 +47,9 @@ def load_kv_index_tl(index_length: int, block_size: int, has_block_tables: bool)
     return _load_kv_index
 
 
+_COMPILED: dict = {}
+
+
 def get_fused_dflash_get_input_tl(
     swa_block_size: int,
     padded_kv_len: int,
@@ -58,6 +61,19 @@ def get_fused_dflash_get_input_tl(
 
     def _fill(slot_mapping_swa, positions, ftp, last_valid_offset, newx_out, positions_out, input_ids_out,
               valid_out, force_out, kv_index_out, kv_len, window_size):  # noqa: ANN001
+        if slot_mapping_swa.device.type == "neuron" and not torch.compiler.is_compiling():
+            # One compiled graph instead of ~40 eager NEFF launches a decode step (~13 ms of host dispatch, issue #29).
+            key = (swa_block_size, padded_kv_len, has_valid_token_mask, has_gate_force_random)
+            if key not in _COMPILED:
+                from deepinfer import envs
+                from deepinfer.neuron_ops import neuron_compile  # noqa: F401  (registers the "deepinfer_neuron" backend)
+
+                _COMPILED[key] = torch.compile(_fill, backend=envs.DEEPINFER_NEURON_COMPILE_BACKEND, dynamic=False)
+            from deepinfer.neuron_ops.transfer_to_neuron import hide_cuda
+
+            with hide_cuda():
+                return _COMPILED[key](slot_mapping_swa, positions, ftp, last_valid_offset, newx_out, positions_out,
+                                      input_ids_out, valid_out, force_out, kv_index_out, kv_len, window_size)
         block = newx_out.shape[1]
         newx, pos, ids, valid, force, kv_index = get_fused_dflash_get_input_ref(
             slot_mapping_swa, positions, ftp, last_valid_offset, swa_block_size, block, window_size,

deepinfer/model_executor/models/mtp/dflash.py

@@ -257,9 +257,10 @@ class DFlash(nn.Module):
         """Run DFlash draft layers on the fixed-stride metadata built by build_dflash_input."""
         draft_hiddens = self._embed_draft_init_tokens(golden_tokens)
         residual = None
-        use_regions = neuron_regions.enabled(input_metadata)
+        # Neuron: each draft layer is one region graph, shared by the draft layers (same kind and shapes) -- unless
+        # this whole method is one (DEEPINFER_NEURON_DRAFTS_REGION, run_drafts): then the layers are traced inline.
+        use_regions = neuron_regions.enabled(input_metadata) and not torch.compiler.is_compiling()
         for layer in self.layers:
-            # Neuron: each draft layer is one region graph, shared by the draft layers (same kind and shapes).
             forward = neuron_regions.entry(layer, "forward_drafts") if use_regions else layer.forward_drafts
             draft_hiddens, residual = forward(
                 draft_hiddens=draft_hiddens,
@@ -446,7 +447,12 @@ class DFlash(nn.Module):
             self.dflash_block_size,
             self.window_size,
         )
-        hidden_states_before_norm = self.forward_drafts(golden_tokens, dflash_input_metadata, nextn_kv_caches)
+        forward_drafts = self.forward_drafts
+        if envs.DEEPINFER_NEURON_DRAFTS_REGION and neuron_regions.enabled(dflash_input_metadata):
+            # The embedding, the draft layers and the hc head as one region (eagerly ~25 launches around the layers'
+            # graphs, each with its host dispatch, issue #29).
+            forward_drafts = neuron_regions.entry(self, "forward_drafts")
+        hidden_states_before_norm = forward_drafts(golden_tokens, dflash_input_metadata, nextn_kv_caches)
         if driver_input is None:
             return None
         return self.sample_drafts(hidden_states_before_norm, driver_input)