"""RMSNorm (+ residual add, + fused fp8 cast) and the clamped SwiGLU through the vllm-neuron#71 NKI kernels (issue #13). The pure-torch shims (``tilelang_ops/neuron/norm.py``, ``tile_kernels.quant.swiglu_forward_kernel``) emulate the Ascend kernels bit for bit: a 1-ulp compensated rstd built from error-free fma tricks, a 4x64-lane square-sum order, ue8m0 exponents read out of the fp32 bits. Inside compiled region graphs those formulations crash neuronx-cc; on the device the norm instead runs ``nki/rmsnorm.rms_norm`` (fp32 interior, rsqrt + one Newton step) and the SwiGLU ``nki/swiglu.swiglu_fused``. Results agree with the shims to bf16 rounding, not bit for bit. The fp8 cast after the norm stays in torch, but without dtype bitcasts: the ue8m0 exponent ``ceil(log2(v))`` comes from ``log2`` corrected against exact ``exp2`` powers of two, and the int16 scale words are joined arithmetically. ``DEEPINFER_NEURON_NKI_NORM=0`` / ``DEEPINFER_NEURON_NKI_SWIGLU=0`` restore the shims. """ from __future__ import annotations import os from typing import Any import torch NORM_ENABLED = os.environ.get("DEEPINFER_NEURON_NKI_NORM", "1") == "1" SWIGLU_ENABLED = os.environ.get("DEEPINFER_NEURON_NKI_SWIGLU", "1") == "1" #: ``swiglu`` without a clamp: a limit no finite bf16 activation reaches. _NO_CLAMP = 3.0e38 _KERNELS: dict = {} def _kernel(module: str, name: str): # noqa: ANN202 """``wrap_nki(nki..)``, built once (a dict, not functools.cache, which dynamo warns about).""" k = _KERNELS.get((module, name)) if k is None: from torch_neuronx.nki_hop import wrap_nki from deepinfer.neuron_ops.nki import rmsnorm, swiglu k = _KERNELS[(module, name)] = wrap_nki(getattr({"rmsnorm": rmsnorm, "swiglu": swiglu}[module], name)) return k def _const(value: float, like: torch.Tensor) -> torch.Tensor: return torch.full((1, 1), value, dtype=torch.float32, device=like.device) # --------------------------------------------------------------------------------------------------------------- # RMSNorm def norm_supported(x: torch.Tensor, extra_sf_out: torch.Tensor | None) -> bool: return NORM_ENABLED and x.device.type == "neuron" and x.dtype in (torch.bfloat16, torch.float32) and extra_sf_out is None def rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float, out_scale: float = 1.0) -> torch.Tensor: """``(x.float() * rsqrt(mean(x.float()^2) + eps) * (weight.float() * out_scale)).to(x.dtype)``, any leading shape.""" dim = x.shape[-1] w = weight.float() if out_scale != 1.0: w = w * out_scale out = _kernel("rmsnorm", "rms_norm")(x.reshape(-1, dim), w.reshape(1, dim), _const(eps, x)) return out.view(x.shape) def ue8m0_exponent(v: torch.Tensor) -> torch.Tensor: """Biased exponent of ``2 ** ceil(log2(v))`` for positive normal fp32 ``v``, as int32 -- the shim's ``((bits(v) - 1) >> 23) + 1`` without reading the bits (a dtype bitcast in a region graph crashes neuronx-cc). ``exp2`` of an integer is exact on the device; ``log2`` is only trusted to within one.""" e = torch.floor(torch.log2(v)) e = torch.where(torch.exp2(e) > v, e - 1, e) e = torch.where(torch.exp2(e + 1) <= v, e + 1, e) e = torch.where(torch.exp2(e) == v, e, e + 1) return e.to(torch.int32) + 127 def pack_ue8m0_int16(e: torch.Tensor) -> torch.Tensor: """int32 exponent bytes ``[..., g]`` -> int16 words ``[..., ceil(g / 2)]``, low byte first (arithmetic only). The word is assembled in fp32 (exact: |v| < 2**17) and converted fp32 -> int16 once. An int32 -> int16 narrowing of the strided (``0::2`` / ``1::2``) int32 result -- with or without a select -- makes neuronx-cc emit a same-shape reinterpret between 4- and 2-byte dtypes and die in codegen (NCC_IBCG902 "Cannot reinterpret to dtype with different size", issue #13) when the exponents come from the NKI RMSNorm output.""" if e.shape[-1] % 2: e = torch.nn.functional.pad(e, (0, 1)) f = e.to(torch.float32) lo, hi = f[..., 0::2], f[..., 1::2] return (lo + hi * 256.0 - torch.where(hi >= 128.0, 65536.0, 0.0)).to(torch.int16) def fp8_cast( result: torch.Tensor, num_per_channels: int, round_sf: bool, use_packed_ue8m0: bool, sf_dtype: torch.dtype ) -> tuple[torch.Tensor, torch.Tensor]: """The norm's per-``num_per_channels`` e4m3 cast of ``result`` ``[..., n]`` (the shim's numerics). Returns the e4m3 tensor and the scales: int16-packed ue8m0 words ``[..., n / npc / 2]`` with ``use_packed_ue8m0``, fp32 ``[..., n / npc]`` otherwise.""" from deepinfer.model_executor.layers.tilelang_ops.common import FP8_AMAX_MARGIN from deepinfer.neuron_ops.fp8 import FP8_MAX q_in = result.float() lead, n = q_in.shape[:-1], q_in.shape[-1] groups = q_in.view(*lead, n // num_per_channels, num_per_channels) clamped = torch.clamp_min(groups.abs().amax(dim=-1), FP8_AMAX_MARGIN) if round_sf: exponent = ue8m0_exponent(clamped * (1.0 / FP8_MAX)) # fp32 product, as the shim inverse = torch.exp2((127 - exponent).to(torch.float32)) if use_packed_ue8m0: assert sf_dtype == torch.int16, f"{sf_dtype=}" sf = pack_ue8m0_int16(exponent) else: sf = torch.exp2((exponent - 127).to(torch.float32)) else: inverse = torch.full_like(clamped, FP8_MAX) / clamped sf = clamped / FP8_MAX q = torch.clamp(groups * inverse.unsqueeze(-1), -FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn) return q.view(q_in.shape), sf def norm( x: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor | None, eps: float, out: torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None, fmt: str, num_per_channels: int, round_sf: bool, use_packed_ue8m0: bool, return_bf16: bool, out_scale: float, ) -> Any: """``tilelang_ops.norm.norm`` on the device (same arguments and result packing; no ``extra_sf_out``). The scale layout is Neuron's: row-major ``[..., packs]`` int16 for packed ue8m0 whether or not the column-major layout was asked for (``_alloc_fp8_sf``'s Neuron branch), fp32 ``[..., groups]`` otherwise.""" do_fp8_cast = fmt != "" sf_out = None if isinstance(out, tuple): out, sf_out = out if residual is not None: summed = x + residual # in x.dtype, as the shim residual_out = summed else: summed, residual_out = x, None result = rms_norm(summed, weight, eps, out_scale) if do_fp8_cast: from deepinfer import PLATFORM q, sf = fp8_cast(result, num_per_channels, round_sf, use_packed_ue8m0, PLATFORM.fp8_sf_type) if out is not None: out.copy_(q.view(out.shape)) q = out if sf_out is not None: sf_out.copy_(sf.view(sf_out.shape)) sf = sf_out main: Any = (q, sf) else: if out is not None: out.copy_(result.view(out.shape)) result = out main = result if return_bf16: main = (main, result.to(torch.bfloat16)) if residual is not None: main = (main, residual_out) return main # --------------------------------------------------------------------------------------------------------------- # SwiGLU def swiglu_supported(x: torch.Tensor, fmt: str, clamped_count: torch.Tensor | None, psum: torch.Tensor | None) -> bool: """bf16 in, bf16 (or fp8 via the shim's cast) out; the clamp counters and the expert prefix-sum layout stay on the shim. fp32 output keeps the shim's fp32 activation.""" return ( SWIGLU_ENABLED and x.device.type == "neuron" and x.dtype == torch.bfloat16 and fmt in ("bf16", "e4m3") and clamped_count is None and psum is None and x.dim() == 2 and x.shape[1] % 2 == 0 ) def swiglu( x: torch.Tensor, topk_weights: torch.Tensor | None, routed_scaling_factor: float | None, clamp_value: float | None, ) -> torch.Tensor: """``silu(min(g, c)) * clamp(u, -c, c) [* w * routed_scaling_factor]`` for ``x = [g | u]``, bf16 ``[rows, h]``.""" if topk_weights is None: w = _const(1.0, x) else: w = topk_weights.to(torch.float32).reshape(-1, 1) if routed_scaling_factor is not None: w = w * routed_scaling_factor limit = _const(_NO_CLAMP if clamp_value is None else float(clamp_value), x) return _kernel("swiglu", "swiglu_fused")(x.contiguous(), w, limit)