From cedd2377725e6991436083703abce1ed4207e669 Mon Sep 17 00:00:00 2001 From: whn09 Date: Tue, 27 Oct 2026 11:50:34 +0800 Subject: [PATCH] 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 --- deepinfer/model_executor/layers/linear.py | 6 ++++++ deepinfer/model_executor/models/decoder_layer.py | 7 ++++++- deepinfer/model_executor/models/v4_modules/attn.py | 5 +++++ deepinfer/neuron_ops/nki_gemm.py | 10 +++++++++- 4 files changed, 26 insertions(+), 2 deletions(-) diff --git a/deepinfer/model_executor/layers/linear.py b/deepinfer/model_executor/layers/linear.py index b362d4d..7a9ae99 100644 --- a/deepinfer/model_executor/layers/linear.py +++ b/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) diff --git a/deepinfer/model_executor/models/decoder_layer.py b/deepinfer/model_executor/models/decoder_layer.py index 3ad5af9..70d97ba 100644 --- a/deepinfer/model_executor/models/decoder_layer.py +++ b/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 diff --git a/deepinfer/model_executor/models/v4_modules/attn.py b/deepinfer/model_executor/models/v4_modules/attn.py index 98a73a4..b5e2171 100644 --- a/deepinfer/model_executor/models/v4_modules/attn.py +++ b/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 diff --git a/deepinfer/neuron_ops/nki_gemm.py b/deepinfer/neuron_ops/nki_gemm.py index 7b197c1..465ea24 100644 --- a/deepinfer/neuron_ops/nki_gemm.py +++ b/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