Neuron: one KV gather per chunk for dense attention; no layer_idx guards (#26)

- nki_attention._chunk_attention: a prompt chunk run as a decode step gathers a
  dense compressed-KV scope once (every token's list is a prefix of the
  longest) and masks it per row; top-k and SWA scopes stay per row. The
  per-row gather was [128, 8192, 512] for a ratio-2 layer of the full model:
  3.8 GB moved per call and a 3 GB NEFF (model code) that ran HBM out.
- decode_part_timer: the decode timers are no-ops inside region graphs. Their
  f-string labels read layer_idx, so Dynamo guarded on it and retraced every
  layer instead of every layer kind.

Tiny V4.1 on a NeuronCore, bucket 128: short 9/9, long 9/9 (logit cos >= 0.997).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

deepinfer/model_executor/models/v4_modules/attn.py

@@ -43,14 +43,20 @@ from deepinfer.model_executor.models.v4_modules.compress import Compressor
 from deepinfer.model_executor.models.v4_modules.cross_layer import CrossLayerRuntime, SharedLayerTopology
 from deepinfer.model_executor.models.v4_modules.index_attn import MQA
 from deepinfer.model_executor.models.v4_modules.mla_attn import LoraRopeCompressFlashMHAAutoRecomp, runtime_attn_sink
-from deepinfer.model_executor.models.v4_modules.utils import convert_rmsnorm, convert_rope, get_enable_timer_layer, merge_linear_weights
+from deepinfer.model_executor.models.v4_modules.utils import (
+    convert_rmsnorm,
+    convert_rope,
+    decode_part_timer,
+    get_enable_timer_layer,
+    merge_linear_weights,
+)
 from deepinfer.model_executor.models.v4_modules.yoco import is_yoco_consumer_layer
 from deepinfer.neuron_ops import nki_gemm
 from deepinfer.neuron_ops import regions as neuron_regions
 from deepinfer.model_executor.parallel_utils.parallel_state import PG
 from deepinfer.model_executor.parallel_utils.userbuffer import userbuffer
 from deepinfer.model_executor.utils import Movable
-from deepinfer.timer import cuda_graph_timer, cuda_timer
+from deepinfer.timer import cuda_timer
 from deepinfer.worker.cache_engine import KVCache
 
 
@@ -655,11 +661,8 @@ class V4Attention(V4AttentionBase):
                 index_length=dense_index_len,
             )
 
-        second_layer, _, same_layers = get_enable_timer_layer(self.layer_idx)
         fused_kwargs = self._fused_norm_rope_cast_kwargs(input_metadata.positions)
-        with cuda_graph_timer.record(
-            f"decode_core_attn-part-{self.layer_idx}", enabled=self.layer_idx == second_layer, scale=len(same_layers)
-        ):
+        with decode_part_timer("decode_core_attn-part", self):
             context, _ = decode_mla_mqa_sparse(
                 q_b,
                 head_dim_v=self.qk_head_dim,

deepinfer/model_executor/models/v4_modules/index_attn.py

@@ -33,8 +33,8 @@ from deepinfer.model_executor.layers.tilelang_ops.load_paged_kv_cache import loa
 from deepinfer.model_executor.models.v4_modules.compress import KProject
 from deepinfer.model_executor.models.v4_modules.cross_layer import CrossLayerRuntime, SharedLayerTopology
 from deepinfer.model_executor.models.v4_modules.mla_attn import MQAIndexer
-from deepinfer.model_executor.models.v4_modules.utils import get_enable_timer_layer
-from deepinfer.timer import cuda_graph_timer, cuda_timer
+from deepinfer.model_executor.models.v4_modules.utils import decode_part_timer, get_enable_timer_layer
+from deepinfer.timer import cuda_timer
 
 
 def _align_mqa_weights(weights: torch.Tensor) -> torch.Tensor:
@@ -476,10 +476,7 @@ class MQA(nn.Module):
             cross_layer_runtime,
         )
 
-        second_layer, _, same_layers = get_enable_timer_layer(self.layer_idx)
-        with cuda_graph_timer.record(
-            f"decode_index_attn-part-{self.layer_idx}", enabled=self.layer_idx == second_layer, scale=len(same_layers)
-        ):
+        with decode_part_timer("decode_index_attn-part", self):
             p_sum = self._compute_decode_p_sum(
                 (q_fp8or4, q_sf),
                 kv_cache,

deepinfer/model_executor/models/v4_modules/utils.py

@@ -5,6 +5,7 @@ layers; `merge_linear_weights` and `init_dummy_weight` reshape and seed runtime
 The `DeepseekDecoderLayer*` registries feed the timer sampling groups.
 """
 
+import contextlib
 from collections.abc import Callable
 
 import torch
@@ -141,6 +142,19 @@ def create_deepseek_decode_layer(layer_str: str, layer_idx: int) -> None:
     DeepseekDecoderLayerIdx[layer_idx] = layer_str
 
 
+def decode_part_timer(name: str, module: nn.Module) -> contextlib.AbstractContextManager:
+    """``cuda_graph_timer`` for a decode part, sampled on the second layer of its kind. A no-op inside a compiled
+    Neuron region: reading ``module.layer_idx`` there makes Dynamo guard on its value (one recompile per layer
+    instead of one per layer kind)."""
+    if torch.compiler.is_compiling():
+        return contextlib.nullcontext()
+    from deepinfer.timer import cuda_graph_timer
+
+    layer_idx = module.layer_idx
+    second_layer, _, same_layers = get_enable_timer_layer(layer_idx)
+    return cuda_graph_timer.record(f"{name}-{layer_idx}", enabled=layer_idx == second_layer, scale=len(same_layers))
+
+
 def get_enable_timer_layer(layer_idx: int) -> tuple[int, int, list[int]]:
     same_layers = DeepseekDecoderLayerType[DeepseekDecoderLayerIdx[layer_idx]]
     second_layer = same_layers[1] if len(same_layers) > 1 else same_layers[0]

deepinfer/neuron_ops/nki_attention.py

@@ -255,10 +255,63 @@ def flash_mla_decode(
         if mask is not None:
             invalid = invalid | mask
         scopes.append((cache, torch.where(invalid, torch.zeros_like(idx), idx), invalid))
-    out = sparse_decode(q.reshape(rows, h_q, d), scopes, softmax_scale, attn_sink)
+    if rows == PREFILL_BUCKET:
+        out = _chunk_attention(q.reshape(rows, h_q, d), scopes, softmax_scale, attn_sink)
+    else:
+        out = sparse_decode(q.reshape(rows, h_q, d), scopes, softmax_scale, attn_sink)
     return out.view(b, s_q, h_q, d)
 
 
+#: 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
+
+
+def _chunk_attention(
+    q: torch.Tensor,
+    scopes: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]],
+    softmax_scale: float,
+    attn_sink: torch.Tensor | None,
+) -> torch.Tensor:
+    """``sparse_decode`` for a prompt chunk run as a decode step (issue #26): rows are one sequence's tokens.
+
+    A dense compressed-KV scope lists, for every token, the positions before it: the lists are prefixes of the
+    longest one. Gathering every row's list (``[rows, max_model_len / ratio, 512]``) would take GBs and a NEFF of
+    as many DMA descriptors, so such a scope is gathered once and masked per row. Top-k and SWA scopes stay per row.
+    Same numerics as the kernel: fp32 logits and softmax with the sink in the denominator, bf16 P @ V.
+    """
+    rows, heads, d = q.shape
+    qf = q.float()
+    logits, values = [], []
+    for cache, idx, invalid in scopes:
+        if idx.shape[1] >= _SHARED_SCOPE_MIN:
+            longest = (~invalid).sum(dim=1).argmax()
+            kv = gather_paged_kv(cache, idx.index_select(0, longest.view(1))).view(-1, d).float()  # [topk, 512]
+            s = torch.matmul(qf.view(rows * heads, d), kv.transpose(0, 1)).view(rows, heads, -1)
+        else:
+            kv = gather_paged_kv(cache, idx).float()  # [rows, topk, 512]
+            s = torch.bmm(qf, kv.transpose(1, 2))
+        logits.append(torch.where(invalid.unsqueeze(1), float("-inf"), s * softmax_scale))
+        values.append(kv)
+    peak = torch.stack([x.amax(dim=-1) for x in logits], dim=-1).amax(dim=-1)  # [rows, heads]
+    sink = attn_sink.float().view(1, heads) if attn_sink is not None else None
+    if sink is not None:
+        peak = torch.maximum(peak, sink.expand(rows, heads))
+    peak = torch.where(torch.isinf(peak), torch.zeros_like(peak), peak)
+    total = torch.exp(sink - peak) if sink is not None else torch.zeros_like(peak)
+    out = torch.zeros((rows, heads, d), dtype=torch.float32, device=q.device)
+    for s, kv in zip(logits, values, strict=True):
+        p = torch.exp(s - peak.unsqueeze(-1))
+        total = total + p.sum(dim=-1)
+        p = p.to(torch.bfloat16).float()
+        if kv.dim() == 2:
+            out = out + torch.matmul(p.view(rows * heads, -1), kv).view(rows, heads, d)
+        else:
+            out = out + torch.bmm(p, kv)
+    valid_any = torch.stack([(~inv).any(dim=1) for _, _, inv in scopes], dim=-1).any(dim=-1).view(rows, 1, 1)
+    out = out / total.clamp_min(1e-30).unsqueeze(-1)
+    return torch.where(valid_any, out, torch.zeros_like(out)).to(torch.bfloat16)
+
+
 # --------------------------------------------------------------------------- #
 # indexer scores (deep_gemm.fp8_fp4_paged_mqa_logits)
 # --------------------------------------------------------------------------- #