neuron: the grouped wv projection takes the bf16 context through the NKI qdq kernel instead of an in-graph fp8 encode (TTFT 1009 -> 942 ms, step 437 -> 414 ms, tokens identical)

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

deepinfer/model_executor/layers/attention_backend.py

@@ -88,6 +88,7 @@ def decode_mla_mqa_sparse(
     indices_in_kvcache_extra: torch.Tensor | None = None,
     index_length_extra: torch.Tensor | None = None,
     fused_norm_rope_cast_kwargs: dict | None = None,
+    neuron_dense: dict | None = None,
 ) -> tuple[torch.Tensor | tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
     """
     Args:
@@ -117,6 +118,7 @@ def decode_mla_mqa_sparse(
             extra_indices_in_kvcache=indices_in_kvcache_extra,
             extra_topk_length=index_length_extra,
             **fused_norm_rope_cast_kwargs,
+            **({"neuron_dense": neuron_dense} if neuron_dense is not None else {}),
         )
         return (output_fp8, output_sf), lse
 
@@ -139,6 +141,7 @@ def decode_mla_mqa_sparse(
         topk_length=index_length,
         extra_topk_length=index_length_extra,
         enable_batch_invariant=batch_invariant,
+        **({"neuron_dense": neuron_dense} if neuron_dense is not None else {}),
     )
     output, lse = output.squeeze(1), lse.squeeze(-1)
     if padded_heads is not None:

deepinfer/model_executor/model_loader.py

@@ -85,6 +85,8 @@ def get_model(
         from deepinfer.neuron_ops import nki_gemm
 
         logger.info(f"Neuron: {nki_gemm.prepare_model(model)} weights take the NKI block-quant GEMMs")
+        if os.environ.get("DEEPINFER_NEURON_BF16_NORMS", "1") == "1":
+            logger.info(f"Neuron: {nki_gemm.bf16_norms(model)} norms feed only NKI kernels and skip their fp8 copy")
         from deepinfer.neuron_ops import static_moe
 
         logger.info(f"Neuron: {static_moe.prepare_shared_slices(model)} MoE layers compute a slice of the shared expert")

deepinfer/model_executor/models/decoder_layer.py

@@ -176,6 +176,14 @@ class DeepseekDecoderLayer(HCLayerMixin):
             # where it would be a graph output (and its encode would run) for nothing.
             hidden_states = hidden_states_bf16
         neuron_regions.split(coarse=True)
+        capacity = envs.DEEPINFER_NEURON_MOE_CAPACITY
+        if capacity and hidden_states_bf16.shape[0] == envs.DEEPINFER_NEURON_PREFILL_BUCKET:
+            # A prefill chunk: the experts over their routed rows only; the overflow check is a graph break, and an
+            # overflowed chunk reruns the dense path.
+            out, overflow = self.mlp.forward_capacity(hidden_states, hidden_states_bf16, input_metadata)
+            if bool(overflow):
+                out = self.mlp(hidden_states, hidden_states_bf16, input_metadata)
+            return out, residual
         return self.mlp(hidden_states, hidden_states_bf16, input_metadata), residual
 
     def _attn_half(

deepinfer/model_executor/models/v4_modules/attn.py

@@ -200,6 +200,23 @@ class V4AttentionBase(nn.Module):
     def _wv_and_o_proj(self, context: torch.Tensor | tuple[torch.Tensor, torch.Tensor], positions: torch.Tensor) -> torch.Tensor:
         if not self.use_fused_norm_rope_attn_rope_cast:
             self.rotary_emb(positions, context[..., -self.qk_rope_head_dim :], conjugate=True)
+        if (
+            PLATFORM.is_neuron
+            and not isinstance(context, tuple)
+            and not self.use_fp8_einsum_output
+            and self.einsum_out_clamp_value is None
+            and nki_gemm.has_kernel(self.wv_weight)
+        ):
+            # The NKI kernel quantizes the bf16 context itself (1x32, as per_token_cast): no fp8 encode in the graph.
+            ctx = context.view(context.shape[0], self.n_wv_group, self.wv_group_dim)
+            wv = torch.stack([nki_gemm.linear_qdq(ctx[:, h], self.wv_weight, group=h) for h in range(self.n_wv_group)], dim=1)
+            out = self.out_proj(wv.flatten(1, 2))
+            tp_group = getattr(self, "tp_group", None)
+            if tp_group is not None:
+                import torch.distributed._functional_collectives as funcol
+
+                out = funcol.wait_tensor(funcol.all_reduce(out.float(), "sum", tp_group)).to(out.dtype)
+            return out
         if isinstance(context, tuple):
             context_fp8, context_sf = context
             gran = 32
@@ -745,6 +762,9 @@ class V4Attention(V4AttentionBase):
             )
 
         fused_kwargs = self._fused_norm_rope_cast_kwargs(input_metadata.positions)
+        dense = self._neuron_dense_chunk(q_b, kv_ca_index, kv_cache, input_metadata)
+        if dense is not None:
+            attn_kv_args["neuron_dense"] = dense
         with decode_part_timer("decode_core_attn-part", self):
             context, _ = decode_mla_mqa_sparse(
                 q_b,
@@ -758,6 +778,37 @@ class V4Attention(V4AttentionBase):
             )
         return context
 
+    def _neuron_dense_chunk(
+        self, q_b: torch.Tensor, kv_ca_index: torch.Tensor | None, kv_cache: KVCache, input_metadata: InputMetadataV4
+    ) -> dict | None:
+        """Neuron: the dense-attention inputs for a single request's prompt chunk (nki_attention.dense_chunk), or None
+        to keep the gathered kernel. A top-k scope qualifies only when the context bucket's compressed keys fit the
+        top-k width (then the indexer selects all of them)."""
+        if not (PLATFORM.is_neuron and self.use_swa and self.has_packed_kv and kv_ca_index is not None):
+            return None
+        from deepinfer.neuron_ops import nki_attention
+
+        if not (nki_attention.PREFILL_DENSE and nki_attention.single_request_chunk() and q_b.shape[0] == nki_attention.PREFILL_BUCKET):
+            return None
+        ctx = getattr(input_metadata, "neuron_ctx_limit", None) or self.max_model_len
+        block_tables = input_metadata.block_tables
+        if block_tables is None or kv_cache.ca_cache is None:
+            return None
+        ratio = int(self.ratio_key)
+        if self.is_sparse:
+            # The indexer ranks the first klen compressed keys; when the bucket's keys fit the width, it keeps all.
+            klen = (getattr(input_metadata, "klen", None) or {}).get(self.ratio_key)
+            if klen is None or ctx // ratio > kv_ca_index.shape[-1]:
+                return None
+            offset_len = torch.stack([torch.zeros_like(klen), klen], dim=1)
+        else:
+            offset_len = (input_metadata.kv_ca_index_offset_len or {}).get(self.ratio_key)
+            if offset_len is None:
+                return None
+        block_tokens = kv_cache.ca_cache.shape[1] * ratio
+        blocks = min(block_tables.shape[1], -(-ctx // block_tokens))
+        return {"block_table": block_tables[0], "blocks": blocks, "offset_len": offset_len}
+
     def _store_yoco_full_kv(
         self,
         cross_layer_runtime: CrossLayerRuntime,

deepinfer/neuron_ops/nki_attention.py

@@ -272,9 +272,21 @@ def flash_mla_decode(
     extra_indices: torch.Tensor | None,
     extra_topk_length: torch.Tensor | None,
     softmax_scale: float,
+    dense: dict | None = None,
 ) -> torch.Tensor:
     """The output of ``flash_mla_with_kvcache`` (decode, ``q [b, s_q, h_q, 512]``, ``indices [b, s_q, topk]``, -1 =
-    invalid, ``topk_length [b]`` masks the tail) -> ``[b, s_q, h_q, 512]`` bf16."""
+    invalid, ``topk_length [b]`` masks the tail) -> ``[b, s_q, h_q, 512]`` bf16. ``dense``: see dense_chunk."""
+    if dense is not None and extra_k_cache is not None:
+        if _DENSE_CHECK:
+            dense = {**dense, "extra_indices": extra_indices, "extra_topk_length": extra_topk_length}
+        out = dense_chunk(q, k_cache, indices, topk_length, extra_k_cache, dense, softmax_scale, attn_sink)
+        if not _DENSE_CHECK:
+            return out
+        ref = flash_mla_decode(
+            q, k_cache, indices, topk_length, attn_sink, extra_k_cache, extra_indices, extra_topk_length, softmax_scale
+        )
+        _report_dense_check(out, ref, dense.get("_own"))
+        return ref
     b, s_q, h_q, d = q.shape
     rows = b * s_q
     scopes, bounds = [], []
@@ -336,6 +348,18 @@ def set_table_bounds(reqs: int | None, ctx: int | None) -> None:
     _TABLE_BOUNDS["reqs"], _TABLE_BOUNDS["ctx"] = reqs, ctx
 
 
+#: Whether the step being replayed is a prompt chunk of at most one request (graph runner, max_num_seqs == 1).
+_SINGLE = [False]
+
+
+def set_single_request_chunk(single: bool) -> None:
+    _SINGLE[0] = single
+
+
+def single_request_chunk() -> bool:
+    return _SINGLE[0]
+
+
 def _table_blocks(blocks: int, page: int) -> int:
     reqs, ctx = _TABLE_BOUNDS["reqs"], _TABLE_BOUNDS["ctx"]
     if not reqs or not ctx:
@@ -470,10 +494,16 @@ def _dense_attention(
     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
+    rows = q.shape[0]
     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
+    return _dense_masked(q, kv, allowed, softmax_scale, sink)
+
+
+def _dense_masked(q: torch.Tensor, kv: torch.Tensor, allowed: torch.Tensor, softmax_scale: float, sink: torch.Tensor) -> torch.Tensor:
+    """``_dense_attention`` given the ``[rows, T]`` mask of each row's slots."""
+    heads = q.shape[1]
     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))
@@ -484,6 +514,133 @@ def _dense_attention(
     return out.to(torch.bfloat16)
 
 
+#: DEEPINFER_NEURON_PREFILL_DENSE: a single request's prompt chunk attends densely over its own SWA ring and
+#: compressed blocks, the mask built from positions (see dense_chunk), when every visible compressed key fits the
+#: top-k width; otherwise (and for decode) the gathered kernel.
+PREFILL_DENSE = os.environ.get("DEEPINFER_NEURON_PREFILL_DENSE", "1") == "1"
+
+
+#: DEEPINFER_NEURON_PREFILL_DENSE_CHECK: also run the gathered kernel, log the largest difference, use the kernel's.
+_DENSE_CHECK = os.environ.get("DEEPINFER_NEURON_PREFILL_DENSE_CHECK") == "1"
+
+
+@torch._dynamo.disable
+def _report_dense_check(out: torch.Tensor, ref: torch.Tensor, own: torch.Tensor | None = None) -> None:
+    o, r = out.float().cpu().flatten(0, 1), ref.float().cpu().flatten(0, 1)
+    d = (o - r).abs()
+    if own is not None:
+        keep = own.cpu().reshape(-1)
+        per_row = d.amax(dim=(-1, -2))
+        worst = torch.argsort(per_row, descending=True)[:4]
+        print(
+            f"dense prefill check (own rows {int(keep.sum())}): max |diff| {per_row[keep].max().item():.4g}, "
+            f"rows differing {(per_row[keep] > 0).sum().item()}; padding rows max {per_row[~keep].max().item() if (~keep).any() else 0:.4g}; "
+            f"worst rows {[(int(i), bool(keep[i]), round(per_row[i].item(), 4)) for i in worst]}",
+            flush=True,
+        )
+        return
+    print(
+        f"dense prefill check: max |diff| {d.max().item():.4g}, mean {d.mean().item():.3g}, ref max {r.abs().max().item():.4g}, "
+        f"rows differing {(d.amax(dim=(-1, -2)) > 0).sum().item()}/{d.shape[0] * d.shape[1]}",
+        flush=True,
+    )
+
+
+def dense_chunk(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    indices: torch.Tensor,
+    topk_length: torch.Tensor | None,
+    extra_k_cache: torch.Tensor,
+    hint: dict,
+    softmax_scale: float,
+    attn_sink: torch.Tensor | None,
+) -> torch.Tensor:
+    """``flash_mla_decode`` for one request's prompt chunk, as two dense matmuls over its own keys.
+
+    The gathered kernel attends each row to its ``window`` SWA slots and its ``topk`` compressed slots with a gather
+    and small matmuls per row (~9 ms a layer for a 512-row chunk). Here the tables are the request's SWA ring
+    (``k_cache[state]``) and its first ``hint["blocks"]`` compressed blocks in logical order (``hint["block_table"]``),
+    and each row's slots are given by arithmetic instead of a scatter: SWA slot ``s`` when ``(s - first) mod ring <
+    count`` (``get_kv_swa_index_decode``: ``count`` consecutive ring slots from the row's first index), compressed key
+    ``t`` when ``offset <= t < offset + len`` (``hint["offset_len"]``). For a top-k scope the caller guarantees the
+    visible keys fit the width, so the selection is all of them. ``q [b, s_q, h, 512]`` -> ``[b, s_q, h, 512]`` bf16.
+    """
+    b, s_q, heads, d = q.shape
+    rows = b * s_q
+    device = q.device
+    ring = k_cache.shape[1]
+    idx = indices.reshape(rows, -1).to(torch.int32)
+    invalid = idx < 0
+    mask = _length_mask(topk_length, b, s_q, idx.shape[1], device)
+    if mask is not None:
+        invalid = invalid | mask
+    count = (~invalid).to(torch.int32).sum(dim=1, keepdim=True)
+    first = idx[:, :1]
+    state = (idx.amax() // ring).reshape(1)  # one request: one ring
+    swa = _gather_blocks(k_cache, state)  # [ring, 512]
+    slots = torch.arange(ring, dtype=torch.int32, device=device).view(1, ring)
+    delta = slots - (first - first // ring * ring)
+    delta = torch.where(delta < 0, delta + ring, delta)
+    own = first // ring == state  # padding rows of the chunk index a dummy slot of another ring: no keys
+    if _DENSE_CHECK:
+        hint["_own"] = own
+    allowed_swa = (delta < count) & own
+    blocks = hint["block_table"][: hint["blocks"]].to(torch.int32)
+    blocks = blocks.clamp(0, extra_k_cache.shape[0] - 1)  # unallocated entries: masked below by the key count
+    ca = _gather_blocks(extra_k_cache, blocks)  # [blocks * page, 512]
+    keys = torch.arange(ca.shape[0], dtype=torch.int32, device=device).view(1, -1)
+    offset_len = hint["offset_len"].to(torch.int32)
+    lo, n = offset_len[:rows, :1], offset_len[:rows, 1:2]
+    allowed_ca = (keys >= lo) & (keys < lo + n) & own
+    kv = torch.nan_to_num(torch.cat([swa, ca], dim=0), nan=0.0, posinf=0.0, neginf=0.0)  # unwritten slots: never NaN
+    allowed = torch.cat([allowed_swa, allowed_ca], dim=1)
+    if attn_sink is None:
+        sink = torch.full((heads, 1), EMPTY, dtype=torch.float32, device=device)
+    else:
+        sink = attn_sink.to(torch.float32).reshape(heads, 1)
+    if _DENSE_CHECK and hint.get("extra_indices") is not None:
+        # The masks the indices give (scatter), for comparison.
+        ref_swa = torch.zeros((rows, ring + 1), dtype=torch.float32, device=device).scatter(
+            1, torch.where(invalid | (idx // ring != state), ring, idx - state * ring).long(), 1.0)[:, :ring] > 0
+        eidx = hint["extra_indices"].reshape(rows, -1).to(torch.int32)
+        einv = eidx < 0
+        emask = _length_mask(hint.get("extra_topk_length"), b, s_q, eidx.shape[1], device)
+        if emask is not None:
+            einv = einv | emask
+        page = extra_k_cache.shape[1]
+        inv = torch.full((extra_k_cache.shape[0] + 1,), -1, dtype=torch.int32, device=device).scatter(
+            0, blocks.long(), torch.arange(blocks.shape[0], dtype=torch.int32, device=device))
+        safe = torch.where(einv, torch.zeros_like(eidx), eidx)
+        logical = inv.index_select(0, (safe // page).reshape(-1).long()).reshape(eidx.shape) * page + safe % page
+        bad = einv | (logical < 0)
+        ref_ca = torch.zeros((rows, ca.shape[0] + 1), dtype=torch.float32, device=device).scatter(
+            1, torch.where(bad, ca.shape[0], logical).long(), 1.0)[:, : ca.shape[0]] > 0
+        _report_masks(allowed_swa, ref_swa & own, allowed_ca, ref_ca & own, first, count, lo, n)
+    out = _dense_masked(q.reshape(rows, heads, d), kv, allowed, softmax_scale, sink)
+    out = torch.where(allowed.any(dim=1).view(rows, 1, 1), out, torch.zeros_like(out))
+    return out.reshape(b, s_q, heads, d)
+
+
+@torch._dynamo.disable
+def _report_masks(a_swa, r_swa, a_ca, r_ca, first, count, lo, n) -> None:  # noqa: ANN001
+    a_swa, r_swa, a_ca, r_ca = (x.cpu() for x in (a_swa, r_swa, a_ca, r_ca))
+    first, count, lo, n = (x.cpu().reshape(-1) for x in (first, count, lo, n))
+    bs = (a_swa != r_swa).sum(1)
+    bc = (a_ca != r_ca).sum(1)
+    rows = [int(i) for i in torch.nonzero(bs + bc).reshape(-1)[:4]]
+    print(
+        f"dense mask check: swa rows wrong {(bs > 0).sum().item()}/{bs.numel()}, ca rows wrong {(bc > 0).sum().item()}; "
+        + "; ".join(
+            f"r{r}: first {first[r].item()} count {count[r].item()} ours_swa {a_swa[r].sum().item()} ref_swa {r_swa[r].sum().item()} "
+            f"lo {lo[r].item()} n {n[r].item()} ours_ca {a_ca[r].sum().item()} ref_ca {r_ca[r].sum().item()} "
+            f"ref_ca_max {int(torch.nonzero(r_ca[r]).max()) if r_ca[r].any() else -1}"
+            for r in rows
+        ),
+        flush=True,
+    )
+
+
 #: 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/neuron_ops/nki_gemm.py

@@ -146,3 +146,25 @@ def linear_qdq_gated(x: torch.Tensor, w: torch.Tensor, group: int, flag: torch.T
 
 def has_kernel(w: torch.Tensor | None) -> bool:
     return ENABLED and w is not None and getattr(w, "_nki_scale", None) is not None and w.device.type == "neuron"
+
+
+def bf16_norms(model: torch.nn.Module) -> int:
+    """Stop the attention / ffn norms casting to fp8 where every consumer is an NKI kernel (prepare_model first).
+
+    The kernels quantize the bf16 hidden themselves, so the norm's fp8 copy and its scales were dead -- but they
+    crossed the graph boundary after the mHC pre block as outputs, and their encode (floor / select / power over the
+    whole hidden, mostly on Vector and GpSimd) was about half of that block (3.6 ms of a 512-row prefill chunk,
+    twice a layer). Returns the count of norms switched to bf16.
+    """
+    count = 0
+    for layer in model.modules():
+        attn, mlp = getattr(layer, "self_attn", None), getattr(layer, "mlp", None)
+        pairs = (
+            (getattr(layer, "attn_layernorm", None), has_kernel(getattr(getattr(attn, "q_a_kv_swa_proj", None), "weight", None))),
+            (getattr(layer, "ffn_layernorm", None), has_kernel(getattr(mlp, "gate_up_proj_weight", None)) and not getattr(mlp, "use_mega_moe", False)),
+        )
+        for norm, nki_only in pairs:
+            if nki_only and getattr(norm, "use_fp8", False):
+                norm.use_fp8 = False
+                count += 1
+    return count

deepinfer/neuron_ops/static_moe.py

@@ -339,7 +339,16 @@ def _nki_moe_forward(
         acc = acc + _shared_slice(moe, x).float()
     regions.split()
     hier = getattr(moe, "_neuron_hier_ar", None)
-    out = _all_reduce(acc, group, hier) if replicated else _reduce_scatter(acc, group, size, rank)[:num_tokens]
+    flag_reduced = None
+    if capacity and replicated and size > 1:
+        # The ranks' overflow flags ride the output's all-reduce as one extra row (a separate 64-rank all-reduce
+        # cost ~3.4 ms a layer): every rank then sees the sum, so all of them take the same path.
+        flag = (overflow if overflow is not None else torch.zeros((), dtype=torch.bool, device=acc.device)).to(torch.float32)
+        acc = torch.cat([acc, flag.reshape(1, 1).expand(1, acc.shape[1])], dim=0)
+        out = _all_reduce(acc, group, hier)
+        out, flag_reduced = out[:-1], out[-1, 0]
+    else:
+        out = _all_reduce(acc, group, hier) if replicated else _reduce_scatter(acc, group, size, rank)[:num_tokens]
     regions.split()
     if moe.has_shared and not sliced:
         shared = nki_gemm.linear_qdq(x_local, moe.shared_gate_up_proj.weight)
@@ -347,6 +356,8 @@ def _nki_moe_forward(
             shared, fmt="bf16", num_per_channels=None, clamp_value=moe.swiglu_clamp_value
         )
         out = out + nki_gemm.linear_qdq(shared, moe.shared_down_proj.weight).float()
+    if capacity and flag_reduced is not None:
+        return out.to(torch.bfloat16), flag_reduced > 0
     if capacity:
         flag = (overflow if overflow is not None else torch.zeros((), dtype=torch.bool, device=out.device)).to(torch.float32)
         if size > 1:

deepinfer/neuron_ops/vendor/flash_mla/flash_mla_interface.py

@@ -233,6 +233,7 @@ def _decode_core(
     d_v: int,
     softmax_scale: float,
     rms_norm_scale_factor: Optional[torch.Tensor],  # [b, s_q, h_q]
+    neuron_dense: Optional[dict] = None,  # DeepInfer Neuron: a prompt chunk's dense attention (nki_attention.dense_chunk)
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     b, s_q, h_q, d_qk = q.shape
     rows = b * s_q
@@ -243,7 +244,7 @@ def _decode_core(
         assert (extra_k_cache is None) == (extra_indices_in_kvcache is None)
         out = _nki.flash_mla_decode(
             q, k_cache, indices_in_kvcache, topk_length, attn_sink,
-            extra_k_cache, extra_indices_in_kvcache, extra_topk_length, softmax_scale,
+            extra_k_cache, extra_indices_in_kvcache, extra_topk_length, softmax_scale, neuron_dense,
         )
         return out, torch.full((b, s_q, h_q), float("nan"), dtype=torch.float32, device=q.device)
 
@@ -312,6 +313,7 @@ def flash_mla_with_kvcache(
     topk_length: Optional[torch.Tensor] = None,
     extra_topk_length: Optional[torch.Tensor] = None,
     enable_batch_invariant: bool = False,
+    neuron_dense: Optional[dict] = None,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     """Sparse MLA decoding with a paged quantized KV cache.
 
@@ -414,6 +416,7 @@ def flash_mla_with_kvcache(
         head_dim_v,
         softmax_scale,
         None,
+        neuron_dense=neuron_dense,
     )
     return out, lse.transpose(1, 2).contiguous()
 

deepinfer/neuron_ops/vendor/flash_mla/fused_norm_rope_attn_rope_cast.py

@@ -144,6 +144,7 @@ def decode(
     extra_k_cache: Optional[torch.Tensor] = None,
     extra_indices_in_kvcache: Optional[torch.Tensor] = None,
     extra_topk_length: Optional[torch.Tensor] = None,
+    neuron_dense: Optional[dict] = None,
 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """Returns (out_fp8, out_sf, lse [s_q, h_q]).
 
@@ -170,6 +171,7 @@ def decode(
         d_v,
         sm_scale,
         None if scale_factor is None else scale_factor.view(s_q, 1, h_q),
+        neuron_dense=neuron_dense,
     )
     out = _rope(out.view(s_q, h_q, d_v), True, cos_sin_cache, token_positions, rope_dim)
     out_fp8, out_sf = _cast_out(

deepinfer/worker/neuron_graph_runner.py

@@ -200,7 +200,9 @@ class NeuronGraphRunner(CudaGraphRunner):
             # 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)
+            nki_attention.set_single_request_chunk(self.max_reqs == 1)
         else:
+            nki_attention.set_single_request_chunk(False)
             reqs = -(-batch_size // (1 + envs.DEEPINFER_USE_NEXTN_DECODE))
             nki_attention.set_table_bounds(reqs, getattr(input_metadata_device, "neuron_ctx_limit", None))
         if not envs.DEEPINFER_NEURON_REGION_TIMING: