From ab74c463894d5d81929bbbc3db8a87b095b41abb Mon Sep 17 00:00:00 2001 From: whn09 Date: Tue, 27 Oct 2026 14:11:06 +0800 Subject: [PATCH] neuron: dense prefill attention over small KV tables and prefill table bounds (both opt-in / dormant by default) _dense_attention matches gathered_attn_idx to ~1 bf16 ulp (same error vs an fp64 reference). It only applies to tables <= DEEPINFER_NEURON_DENSE_ATTN_MAX_KV rows, which a prefill chunk's whole-pool table exceeds; bounding that table (DEEPINFER_NEURON_PREFILL_TABLE_BOUNDS) went through the compacting gather and doubled TTFT, so it is off. Co-Authored-By: Claude Opus 5.5 --- deepinfer/environ.py | 4 +++ deepinfer/neuron_ops/nki_attention.py | 34 ++++++++++++++++++++++++- deepinfer/worker/neuron_graph_runner.py | 7 ++++- 3 files changed, 43 insertions(+), 2 deletions(-) diff --git a/deepinfer/environ.py b/deepinfer/environ.py index 383334f..6932118 100644 --- a/deepinfer/environ.py +++ b/deepinfer/environ.py @@ -148,6 +148,10 @@ class envs: 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) + # Bound a prefill chunk's attention table by max_num_seqs requests and the context bucket (issue #29). Off: + # the compacting gather it switches to made a 512-token prefill slower (TTFT 1.67 -> 3.61 s) than decoding + # the whole pool in place. + DEEPINFER_NEURON_PREFILL_TABLE_BOUNDS = EnvVar.bool(default=False) # Neuron only: compile the decode step with fullgraph=True (raise on the first graph break). DEEPINFER_NEURON_FULLGRAPH = EnvVar.bool() diff --git a/deepinfer/neuron_ops/nki_attention.py b/deepinfer/neuron_ops/nki_attention.py index 64a6c37..1829743 100644 --- a/deepinfer/neuron_ops/nki_attention.py +++ b/deepinfer/neuron_ops/nki_attention.py @@ -445,13 +445,45 @@ def _table_attention( sink = torch.full((heads, 1), EMPTY, dtype=torch.float32, device=q.device) else: sink = attn_sink.to(torch.float32).reshape(heads, 1) + lonely = invalid.all(dim=1).view(rows, 1, 1) + if rows >= _DENSE_MIN_ROWS and 0 < kv.shape[0] <= _DENSE_MAX_KV: + out = _dense_attention(q, kv, idx, invalid, softmax_scale, sink) + return torch.where(lonely, torch.zeros_like(out), out) out = _kernel("sparse_attn", "gathered_attn_idx")( q.reshape(rows * heads, d).contiguous(), kv.contiguous(), torch.where(invalid, 0, idx).contiguous(), bias, sink ).view(rows, heads, d) - lonely = invalid.all(dim=1).view(rows, 1, 1) return torch.where(lonely, torch.zeros_like(out), out) +#: DEEPINFER_NEURON_DENSE_ATTN_MAX_KV: prefill chunks whose decoded KV table has at most this many rows attend densely +#: (two matmuls over the table, masked to each row's slots) instead of through gathered_attn_idx, whose per-row +#: gathers and small matmuls cost ~9 ms a layer for a 512-token chunk (issue #29). 0 disables it. +_DENSE_MAX_KV = int(os.environ.get("DEEPINFER_NEURON_DENSE_ATTN_MAX_KV", "4096")) +#: DEEPINFER_NEURON_DENSE_ATTN_MIN_ROWS: the dense path's smallest step (default: prefill chunks only). +_DENSE_MIN_ROWS = int(os.environ.get("DEEPINFER_NEURON_DENSE_ATTN_MIN_ROWS", "65")) + + +def _dense_attention( + q: torch.Tensor, kv: torch.Tensor, idx: torch.Tensor, invalid: torch.Tensor, softmax_scale: float, sink: torch.Tensor +) -> torch.Tensor: + """``gathered_attn_idx``'s result computed over the whole table: logits of bf16 operands accumulated in fp32, + scaled; each row's slots ``idx`` (unique per row) unmasked; the max over the slots only; the sink in the + denominator; probabilities rounded to bf16 for the value product, as the kernel does. ``q [rows, heads, d]``, + ``kv [T, d]`` (keys are values) -> ``[rows, heads, d]`` bf16.""" + rows, heads, d = q.shape + t = kv.shape[0] + slot = torch.where(invalid, torch.full_like(idx, t), idx).long() # invalid slots -> a dump column + allowed = torch.zeros((rows, t + 1), dtype=torch.float32, device=q.device).scatter(1, slot, 1.0)[:, :t] > 0 + kvf = kv.float() + logits = torch.matmul(q.float(), kvf.t()) * softmax_scale # [rows, heads, t] + logits = torch.where(allowed.unsqueeze(1), logits, torch.full_like(logits, EMPTY)) + neg_max = -logits.amax(dim=-1, keepdim=True) + probs = torch.exp(logits + neg_max) + den = probs.sum(dim=-1, keepdim=True) + torch.exp(sink.view(1, heads, 1) + neg_max) + out = torch.matmul(probs.to(torch.bfloat16).float(), kvf) / den + return out.to(torch.bfloat16) + + #: A scope this wide is a dense compressed-KV read (every position up to the token's own), not a top-k selection. _SHARED_SCOPE_MIN = 1024 diff --git a/deepinfer/worker/neuron_graph_runner.py b/deepinfer/worker/neuron_graph_runner.py index 24212a6..1551b70 100644 --- a/deepinfer/worker/neuron_graph_runner.py +++ b/deepinfer/worker/neuron_graph_runner.py @@ -123,6 +123,8 @@ class _LazyKVCache: class NeuronGraphRunner(CudaGraphRunner): def __init__(self, max_batch_size: int, input_on_device: bool = False) -> None: super().__init__(max_batch_size, input_on_device) + #: Requests a step can hold (max_num_seqs): the decode rows are max_num_seqs * (1 + nextn). + self.max_reqs = max(1, max_batch_size // (1 + envs.DEEPINFER_USE_NEXTN_DECODE)) # Prompt chunks run as decode-shaped steps of this many tokens (issue #26): one more graph size. self.prefill_bucket = envs.DEEPINFER_NEURON_PREFILL_BUCKET if self.prefill_bucket > self.max_batch_size: @@ -194,7 +196,10 @@ class NeuronGraphRunner(CudaGraphRunner): from deepinfer.neuron_ops import nki_attention if batch_size == nki_attention.PREFILL_BUCKET: - nki_attention.set_table_bounds(None, None) + # A prompt chunk packs at most max_num_seqs requests: bound the table by those and the step's context + # bucket (without a bound every prefill layer decoded the whole KV pool, issue #29). + ctx = getattr(input_metadata_device, "neuron_ctx_limit", None) + nki_attention.set_table_bounds(self.max_reqs if envs.DEEPINFER_NEURON_PREFILL_TABLE_BOUNDS and ctx else None, ctx) else: reqs = -(-batch_size // (1 + envs.DEEPINFER_USE_NEXTN_DECODE)) nki_attention.set_table_bounds(reqs, getattr(input_metadata_device, "neuron_ctx_limit", None))