coffee/DeepInfer
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepInfer.git synced 2026-10-28 23:07:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Neuron: split region graphs for prefill chunks only (#26)

Browse Source 
A 128-row chunk's attention ran 122 ms a layer as one graph and 62 ms split at
the stage boundaries (bench3 vs bench4); decode's 1-row graphs are faster
whole. split() now cuts when DEEPINFER_NEURON_REGION_SPLIT (default off now)
or, under DEEPINFER_NEURON_REGION_SPLIT_PREFILL (default on), when the step
has more than one row. Syncing after every region call (the queueing
hypothesis) did not help: 37.7 s TTFT at 512.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
whn09 2026-10-23 21:39:56 +08:00
parent 32e08d78cd
commit d188d1b3d1
3 changed files with 19 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
deepinfer/environ.py
View File

@ -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.

16
deepinfer/neuron_ops/regions.py
View File

@ -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()

2
deepinfer/worker/neuron_graph_runner.py
View File

@ -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)