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 <noreply@anthropic.com>

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

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
 

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