From 6c653df08c70703cad8b47287efba6b8df696eef Mon Sep 17 00:00:00 2001 From: whn09 Date: Tue, 27 Oct 2026 00:14:32 +0800 Subject: [PATCH] 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 --- deepinfer/environ.py | 4 +++- .../layers/tilelang_ops/neuron/get_kv_index.py | 16 ++++++++++++++++ deepinfer/model_executor/models/mtp/dflash.py | 12 +++++++++--- 3 files changed, 28 insertions(+), 4 deletions(-) diff --git a/deepinfer/environ.py b/deepinfer/environ.py index 956f971..54b68bf 100644 --- a/deepinfer/environ.py +++ b/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() diff --git a/deepinfer/model_executor/layers/tilelang_ops/neuron/get_kv_index.py b/deepinfer/model_executor/layers/tilelang_ops/neuron/get_kv_index.py index 0f11c2e..fceb688 100644 --- a/deepinfer/model_executor/layers/tilelang_ops/neuron/get_kv_index.py +++ b/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, diff --git a/deepinfer/model_executor/models/mtp/dflash.py b/deepinfer/model_executor/models/mtp/dflash.py index a4b3f79..510279b 100644 --- a/deepinfer/model_executor/models/mtp/dflash.py +++ b/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)