neuron: fp8-activation linears take the NKI kernel (Engram projection); drop unused fp8 hidden copies before region splits; log weights without NKI scales

An fp8 tuple input fell back to the deep_gemm shim, which dequantized the whole fp8 weight in the graph on
every call (Engram fused_kv_projection: 9 ms of an Engram layer's prefill graph, GpSimd POW for the scales).
Now decoded to bf16 and run through the NKI kernel. TTFT 1.71 -> 1.66 s (bucket 512), tokens unchanged.

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

deepinfer/model_executor/layers/linear.py

@@ -154,6 +154,12 @@ class ReplicatedLinear(torch.nn.Module):
                 # The NKI kernel quantizes the bf16 input itself (same 1x32 fp8 round trip as per_token_cast).
                 return nki_gemm.linear_qdq(x, self.weight)
             x = tile_kernels.quant.per_token_cast(x, **get_fp8_cast_params())
+        elif out_acc is None and PLATFORM.is_neuron and nki_gemm.has_kernel(self.weight) and x[0].shape[-1] % 128 == 0:
+            # An fp8 activation (e.g. Engram's rows fetched from the host): decode it exactly to bf16 and take the NKI
+            # kernel, whose 1x32 re-quantization is idempotent on it. Through the deep_gemm shim the whole fp8
+            # weight was dequantized in the graph on every call (9 ms of a prefill layer graph, issue #29).
+            k = x[0].shape[-1]
+            return nki_gemm.linear_qdq(nki_gemm._activation_bf16(x[0], x[1], k).reshape(-1, k), self.weight)
 
         # FP8 GEMM.
         weight = (self.weight, self.weight.sf)

deepinfer/model_executor/models/decoder_layer.py

@@ -3,7 +3,7 @@ from collections.abc import Callable
 import torch
 from torch import nn
 
-from deepinfer import AttrDict, envs
+from deepinfer import PLATFORM, AttrDict, envs
 from deepinfer.logger import init_logger
 from deepinfer.model_executor.input_metadata import InputMetadata
 from deepinfer.model_executor.input_metadata.yoco import YocoMetadataMixin
@@ -29,6 +29,7 @@ from deepinfer.model_executor.models.v4_modules.yoco import (
     yoco_shared_hidden_required,
 )
 from deepinfer.model_executor.utils import Movable
+from deepinfer.neuron_ops import nki_gemm
 from deepinfer.neuron_ops import regions as neuron_regions
 from deepinfer.timer import cuda_graph_timer, cuda_timer
 from deepinfer.worker.cache_engine import KVCache
@@ -170,6 +171,10 @@ class DeepseekDecoderLayer(HCLayerMixin):
         hidden_states, hidden_states_bf16, residual = self._attn_half(
             hidden_states, kv_cache, input_metadata, residual, engram_input, cross_layer_runtime, dflash_input
         )
+        if PLATFORM.is_neuron and nki_gemm.has_kernel(getattr(self.mlp, "gate_up_proj_weight", None)):
+            # The static MoE's NKI path reads only the bf16 hidden: drop the ffn norm's fp8 copy before the split,
+            # where it would be a graph output (and its encode would run) for nothing.
+            hidden_states = hidden_states_bf16
         neuron_regions.split(coarse=True)
         return self.mlp(hidden_states, hidden_states_bf16, input_metadata), residual
 

deepinfer/model_executor/models/v4_modules/attn.py

@@ -847,6 +847,11 @@ class V4Attention(V4AttentionBase):
         cross_layer_runtime: CrossLayerRuntime,  # values held across layers (shared_hidden / pre_rope_kv / topk)
     ) -> torch.Tensor:
         hidden_states, hidden_states_bf16 = hidden_states.move(), hidden_states_bf16.move()
+        if PLATFORM.is_neuron and nki_gemm.has_kernel(getattr(self.q_a_kv_swa_proj, "weight", None)):
+            # The NKI projections read only the bf16 hidden. Kept live, the norm's fp8 copy crosses the region splits
+            # below as a graph output, so its encode (integer-free but ~10 elementwise passes, much of it on GpSimd)
+            # runs on every layer for nothing (issue #29).
+            hidden_states = hidden_states_bf16
         # shared_hidden is a bf16 tensor held across layers for the long term (not a Movable); it is written
         # into the runtime by the first YOCO consumer layer
         shared_hidden_bf16 = cross_layer_runtime.shared_hidden

deepinfer/neuron_ops/nki_gemm.py

@@ -43,7 +43,8 @@ def prepare_model(model: torch.nn.Module) -> int:
     if not ENABLED:
         return 0
     count = 0
-    for p in model.parameters():
+    skipped = []
+    for name, p in model.named_parameters():
         sf = getattr(p, "sf", None)
         if sf is None or p.dtype not in (torch.float8_e4m3fn, torch.int8) or p.dim() not in (2, 3):
             continue
@@ -57,6 +58,13 @@ def prepare_model(model: torch.nn.Module) -> int:
             # kernel crashes neuronx-cc (InferNonlocalTensors), a uint8 input does not.
             p._nki_weight = p.data.view(torch.uint8)
             count += 1
+        else:
+            skipped.append(f"{name}{tuple(p.shape)}")
+    if skipped:
+        # These stay on the deep_gemm shims: each call dequantizes the whole weight inside the graph.
+        from deepinfer.logger import init_logger
+
+        init_logger(__name__).warning(f"Neuron: {len(skipped)} block-quant weights without the NKI kernels: {skipped[:12]}")
     return count