diff --git a/deepinfer/model_executor/model_loader.py b/deepinfer/model_executor/model_loader.py index 4ebbc5d..dae736a 100644 --- a/deepinfer/model_executor/model_loader.py +++ b/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") diff --git a/deepinfer/model_executor/models/decoder_layer.py b/deepinfer/model_executor/models/decoder_layer.py index 70d97ba..786ea25 100644 --- a/deepinfer/model_executor/models/decoder_layer.py +++ b/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( diff --git a/deepinfer/neuron_ops/nki_gemm.py b/deepinfer/neuron_ops/nki_gemm.py index 465ea24..fd214c4 100644 --- a/deepinfer/neuron_ops/nki_gemm.py +++ b/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 diff --git a/deepinfer/neuron_ops/static_moe.py b/deepinfer/neuron_ops/static_moe.py index 025e0f9..3d163ad 100644 --- a/deepinfer/neuron_ops/static_moe.py +++ b/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: