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

Neuron DSpark: bench --nextn step count fix; split-row threshold knob; table attention stays prompt-only

Browse Source 
bench18 (512 in, 128 out, TP=64, nextn=5): acceptance 3.85 tokens a step,
TPOT 465.9 ms. bench21 with the table path for 6-row steps: acceptance 1.27,
TPOT 848 ms, so verification keeps the per-row path. Whole-layer 6-row graphs
crash neuronx-cc on the full model (DEEPINFER_NEURON_SPLIT_MIN_ROWS, default 1
= split as before).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
whn09 2026-10-24 21:14:58 +08:00
parent fdbb81f14f
commit e0510e92ca
3 changed files with 7 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
deepinfer/neuron_ops/nki_attention.py
View File

@ -255,6 +255,8 @@ def flash_mla_decode(
if mask is not None: if mask is not None:
invalid = invalid | mask invalid = invalid | mask
scopes.append((cache, torch.where(invalid, torch.zeros_like(idx), idx), invalid)) scopes.append((cache, torch.where(invalid, torch.zeros_like(idx), idx), invalid))
# Prompt chunks only: DSpark verification steps (6 rows) on the table path lost acceptance (3.85 -> 1.27 tokens
# a step) and ran slower (bench21), so they stay on the per-row path.
if rows == PREFILL_BUCKET and sum(c.shape[0] * c.shape[1] for c, _, _ in scopes) <= _TABLE_MAX_SLOTS: if rows == PREFILL_BUCKET and sum(c.shape[0] * c.shape[1] for c, _, _ in scopes) <= _TABLE_MAX_SLOTS:
out = _table_attention(q.reshape(rows, h_q, d), scopes, softmax_scale, attn_sink) out = _table_attention(q.reshape(rows, h_q, d), scopes, softmax_scale, attn_sink)
elif rows == PREFILL_BUCKET and any(idx.shape[1] >= _SHARED_SCOPE_MIN for _, idx, _ in scopes): elif rows == PREFILL_BUCKET and any(idx.shape[1] >= _SHARED_SCOPE_MIN for _, idx, _ in scopes):

5
deepinfer/neuron_ops/regions.py
View File

@ -119,6 +119,9 @@ def _boundary() -> None:
#: Rows of the step being run (set by NeuronGraphRunner.replay): a prefill chunk has more than one. #: Rows of the step being run (set by NeuronGraphRunner.replay): a prefill chunk has more than one.
_STEP_ROWS = [1] _STEP_ROWS = [1]
#: Steps up to this many rows keep whole-layer graphs (DEEPINFER_NEURON_SPLIT_MIN_ROWS). Whole 6-row layers (DSpark
#: verification) crash neuronx-cc on the full model, so they are split like prompt chunks by default.
_SPLIT_MIN_ROWS = int(os.environ.get("DEEPINFER_NEURON_SPLIT_MIN_ROWS", "1"))
def set_step_rows(rows: int) -> None: def set_step_rows(rows: int) -> None:
@ -131,5 +134,5 @@ def split() -> None:
layer, issue #26), while decode's 1-row graphs are faster whole.""" layer, issue #26), while decode's 1-row graphs are faster whole."""
if not torch.compiler.is_compiling(): if not torch.compiler.is_compiling():
return return
if envs.DEEPINFER_NEURON_REGION_SPLIT or (envs.DEEPINFER_NEURON_REGION_SPLIT_PREFILL and _STEP_ROWS[0] > 1): if envs.DEEPINFER_NEURON_REGION_SPLIT or (envs.DEEPINFER_NEURON_REGION_SPLIT_PREFILL and _STEP_ROWS[0] > _SPLIT_MIN_ROWS):
_boundary() _boundary()

2
examples/neuron/bench_neuron.py
View File

@ -82,7 +82,7 @@ def _run_one(llm, prompt: list[int] | None, params) -> tuple[float, float, int]
unfinished = engine.has_unfinished_requests() unfinished = engine.has_unfinished_requests()
outputs = engine.step() if unfinished else engine.flush_inflight_reqs() outputs = engine.step() if unfinished else engine.flush_inflight_reqs()
now = time.perf_counter() now = time.perf_counter()
steps += bool(outputs) steps += len(outputs) > 0
for out in outputs: for out in outputs:
tokens += 1 tokens += 1
first = first if first is not None else now first = first if first is not None else now