diff --git a/deepinfer/environ.py b/deepinfer/environ.py index ccd6bb5..6df4fd9 100644 --- a/deepinfer/environ.py +++ b/deepinfer/environ.py @@ -113,7 +113,9 @@ class envs: # graphs (neuron_ops/regions.py, the vllm-neuron#71 layout); "model" compiles the whole step as one graph. DEEPINFER_NEURON_GRAPH_MODE = EnvVar.str(default="regions") # Neuron only: cut the region graphs further at the stage boundaries marked by neuron_ops.regions.split(). - DEEPINFER_NEURON_REGION_SPLIT = EnvVar.bool(default=True) + DEEPINFER_NEURON_REGION_SPLIT = EnvVar.bool(default=False) + # Neuron only: split the region graphs of prefill chunks (steps of more than one row) regardless. + DEEPINFER_NEURON_REGION_SPLIT_PREFILL = EnvVar.bool(default=True) # Neuron only: synchronize around every region graph call and log per-kind times each decode step (profiling). DEEPINFER_NEURON_REGION_TIMING = EnvVar.bool() # Keep the token embedding (and the text model's unused vision tower) in host memory on Neuron. diff --git a/deepinfer/neuron_ops/regions.py b/deepinfer/neuron_ops/regions.py index cd5bc6d..d71e424 100644 --- a/deepinfer/neuron_ops/regions.py +++ b/deepinfer/neuron_ops/regions.py @@ -117,7 +117,19 @@ def _boundary() -> None: _MARK[0], _MARK[1] = now, _MARK[1] + 1 +#: Rows of the step being run (set by NeuronGraphRunner.replay): a prefill chunk has more than one. +_STEP_ROWS = [1] + + +def set_step_rows(rows: int) -> None: + _STEP_ROWS[0] = rows + + def split() -> None: - """Cut the current region graph here (only while compiling; free in eager).""" - if torch.compiler.is_compiling() and envs.DEEPINFER_NEURON_REGION_SPLIT: + """Cut the current region graph here (only while compiling; free in eager). Always for prefill chunks under + DEEPINFER_NEURON_REGION_SPLIT_PREFILL: neuronx-cc runs the split 128-row attention ~2x faster (62 vs 122 ms a + layer, issue #26), while decode's 1-row graphs are faster whole.""" + if not torch.compiler.is_compiling(): + return + if envs.DEEPINFER_NEURON_REGION_SPLIT or (envs.DEEPINFER_NEURON_REGION_SPLIT_PREFILL and _STEP_ROWS[0] > 1): _boundary() diff --git a/deepinfer/worker/neuron_graph_runner.py b/deepinfer/worker/neuron_graph_runner.py index 3ca70a0..702abbd 100644 --- a/deepinfer/worker/neuron_graph_runner.py +++ b/deepinfer/worker/neuron_graph_runner.py @@ -170,6 +170,7 @@ class NeuronGraphRunner(CudaGraphRunner): else: compiled = torch.compile(forward, backend=backend, fullgraph=fullgraph, dynamic=False) # First call compiles (and runs) this batch size's graph. + regions.set_step_rows(batch_size) input_metadata_device.recursive_copy_(input_metadata_source) with hide_cuda(): compiled(input_metadata_device) @@ -187,6 +188,7 @@ class NeuronGraphRunner(CudaGraphRunner): if not self.input_on_device: input_metadata_device.recursive_copy_(cuda_graph_data.input_metadata_source) batch_size = next(bs for bs, data in self.graphs.items() if data is cuda_graph_data) + regions.set_step_rows(batch_size) if not envs.DEEPINFER_NEURON_REGION_TIMING: with hide_cuda(): return self.compiled[batch_size](input_metadata_device)