diff --git a/deepinfer/environ.py b/deepinfer/environ.py index 5450c14..b09f076 100644 --- a/deepinfer/environ.py +++ b/deepinfer/environ.py @@ -114,6 +114,8 @@ class envs: DEEPINFER_NEURON_GRAPH_MODE = EnvVar.str(default="regions") # Neuron only: cut the region graphs further at the stage boundaries marked by neuron_ops.regions.split(). DEEPINFER_NEURON_REGION_SPLIT = EnvVar.bool(default=True) + # Keep the token embedding (and the text model's unused vision tower) in host memory on Neuron. + DEEPINFER_NEURON_HOST_EMBED = EnvVar.bool(default=True) # Neuron only: compile the decode step with fullgraph=True (raise on the first graph break). DEEPINFER_NEURON_FULLGRAPH = EnvVar.bool() diff --git a/deepinfer/model_executor/layers/tilelang_ops/fused_transform.py b/deepinfer/model_executor/layers/tilelang_ops/fused_transform.py index dfb0513..6981cf8 100644 --- a/deepinfer/model_executor/layers/tilelang_ops/fused_transform.py +++ b/deepinfer/model_executor/layers/tilelang_ops/fused_transform.py @@ -275,7 +275,8 @@ def fused_norm_rope_hadamard_cast_store_kvcache( block_size, cache = split_kv_cache(cache_out.kv_cache, cache_layout) cache_quant, cache_sf = cache.quantized_data, cache.sf_data assert cache_quant.stride(1) % cache_layout.stride_alignment == 0 - assert cache_quant.stride(0) % 32 == 0 + # Page alignment is a GPU/Ascend kernel requirement; Neuron's layer-major caches pack pages tightly. + assert PLATFORM.is_neuron or cache_quant.stride(0) % 32 == 0 out_bf16 = _restore_single_head_axis(cache_out.out, x.shape, original_ndim) bf16_stage = cache_out.out_stage if out_bf16 is not None: diff --git a/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_transform.py b/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_transform.py index 70a5185..2e2ec03 100644 --- a/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_transform.py +++ b/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_transform.py @@ -373,11 +373,13 @@ def fused_norm_rope_hadamard_cast_store_kvcache_tl(num_heads: int, head_dim: int # Neuron graphs: join the scale byte pairs into the int16 storage arithmetically (low byte first), as # neuron/norm.py does. Writing them through a uint8 view that the same graph reads back as int16 # crashes neuronx-cc (NCC_IINL902, issue #25). - e = scale.to(torch.int32) + # Pair in fp32 and narrow once (int32 interleaved slices -> int16 crash codegen, NCC_IBCG902). + f = scale.to(torch.float32) if num_scales % 2: - e = torch.nn.functional.pad(e, (0, 1)) - v = e[:, 0::2] | (e[:, 1::2] << 8) - osf[:, : v.shape[-1]].copy_(torch.where(v >= 32768, v - 65536, v).to(torch.int16)) + f = torch.nn.functional.pad(f, (0, 1)) + lo, hi = f[:, 0::2], f[:, 1::2] + v = lo + hi * 256.0 - torch.where(hi >= 128.0, 65536.0, 0.0) + osf[:, : v.shape[-1]].copy_(v.to(torch.int16)) else: osf.view(torch.uint8)[:, :num_scales].copy_(scale) diff --git a/deepinfer/model_executor/layers/tilelang_ops/neuron/norm.py b/deepinfer/model_executor/layers/tilelang_ops/neuron/norm.py index 2bd5771..6cf42e0 100644 --- a/deepinfer/model_executor/layers/tilelang_ops/neuron/norm.py +++ b/deepinfer/model_executor/layers/tilelang_ops/neuron/norm.py @@ -161,11 +161,13 @@ def norm_tl( if use_packed_ue8m0 and sf.dtype == torch.int16: # int16 storage (compiled regions): join exponent byte pairs arithmetically, low byte first. - e = sf_values.to(torch.int32) - if e.shape[-1] % 2: - e = torch.nn.functional.pad(e, (0, 1)) - v = e[..., 0::2] | (e[..., 1::2] << 8) - sf.copy_(torch.where(v >= 32768, v - 65536, v).to(torch.int16).view(sf.shape)) + # Pair the bytes in fp32 (exact below 2^17) and narrow once: int32 interleaved slices narrowed to int16 + # crash neuronx-cc codegen (NCC_IBCG902, see nki_norm.pack_ue8m0_int16). + f = sf_values.to(torch.float32) + if f.shape[-1] % 2: + f = torch.nn.functional.pad(f, (0, 1)) + lo, hi = f[..., 0::2], f[..., 1::2] + sf.copy_((lo + hi * 256.0 - torch.where(hi >= 128.0, 65536.0, 0.0)).to(torch.int16).view(sf.shape)) elif use_packed_ue8m0 and sf.ndim == 4: # col-major packed view: (B, M, num_groups // pack, pack) bytes sf.copy_(sf_values.view(*lead, sf.shape[-2], sf.shape[-1])) diff --git a/deepinfer/model_executor/layers/vocab_parallel_embedding.py b/deepinfer/model_executor/layers/vocab_parallel_embedding.py index 6b6169a..f2a57a3 100644 --- a/deepinfer/model_executor/layers/vocab_parallel_embedding.py +++ b/deepinfer/model_executor/layers/vocab_parallel_embedding.py @@ -20,6 +20,9 @@ class VocabParallelEmbedding(torch.nn.Module): self.weight = reconstruct_param(weight) def forward(self, input_: torch.Tensor) -> torch.Tensor: + if self.weight.device != input_.device: + # Host-resident table (Neuron, see DeepseekForCausalLM.load_model): move the looked-up rows only. + return gather(self.weight, input_.cpu()).to(input_.device) output = gather(self.weight, input_) return output diff --git a/deepinfer/model_executor/model_loader.py b/deepinfer/model_executor/model_loader.py index d55adcd..4af10db 100644 --- a/deepinfer/model_executor/model_loader.py +++ b/deepinfer/model_executor/model_loader.py @@ -76,6 +76,8 @@ def get_model( # The model device is "cuda", or the mapped Neuron target (neuron:0, or the host in Neuron CPU mode). device_type = PLATFORM.device.type if PLATFORM.is_neuron else "cuda" for n, p in model.named_parameters(): + if PLATFORM.is_neuron and p.device.type == "cpu" and n.startswith("model.embed_tokens."): + continue # host-resident on Neuron (DeepseekForCausalLM._neuron_host_resident) assert p.device.type == device_type, f"{n} not on {device_type}" if hasattr(p, "sf"): assert p.sf.device.type == device_type, f"{n}.sf not on {device_type}" diff --git a/deepinfer/model_executor/models/deepseek.py b/deepinfer/model_executor/models/deepseek.py index de66a64..1946b18 100644 --- a/deepinfer/model_executor/models/deepseek.py +++ b/deepinfer/model_executor/models/deepseek.py @@ -148,7 +148,7 @@ class DeepseekModel(nn.Module): ) def allocate_output(self, num_tokens: int) -> CausalLMForwardOutput: - hidden_states = self.embed_tokens.weight.new_zeros((num_tokens, self.hidden_size)) + hidden_states = torch.zeros((num_tokens, self.hidden_size), dtype=self.embed_tokens.weight.dtype, device="cuda") return CausalLMForwardOutput(hidden_states=hidden_states) def forward( @@ -475,6 +475,8 @@ class DeepseekForCausalLM(nn.Module): for nm, m in gpt.named_modules(): for pn, p in m.named_parameters(recurse=False): device = engram_device if getattr(p, "is_engram", False) else "cuda" + if PLATFORM.is_neuron and cls._neuron_host_resident(nm): + device = "cpu" packed_fp4 = ("experts.linear1" in nm or "experts.linear2" in nm) and model_cfg.gpt.mxfp4 alias_source = hf_coverage.aliases.get(f"{nm}.{pn}") if alias_source is not None: @@ -526,6 +528,15 @@ class DeepseekForCausalLM(nn.Module): return model + @classmethod + def _neuron_host_resident(cls, module_name: str) -> bool: + """Parameters kept in host memory on Neuron, where a NeuronCore has 24 GiB of HBM: the token embedding (1.2 GiB; + a lookup moves only its rows, see VocabParallelEmbedding) and the vision tower the text model drops after + loading (1.3 GiB that would otherwise stay in the device caching allocator).""" + if not envs.DEEPINFER_NEURON_HOST_EMBED: + return False + return module_name.endswith("emb.tok_emb") or module_name.startswith("vision_pipeline") + @classmethod def from_model_cfg(cls, gpt: nn.Module, model_cfg: AttrDict) -> nn.Module: return cls(gpt, model_cfg.gpt) diff --git a/deepinfer/model_executor/models/deepseek_vl.py b/deepinfer/model_executor/models/deepseek_vl.py index ef98df8..ea2bf79 100644 --- a/deepinfer/model_executor/models/deepseek_vl.py +++ b/deepinfer/model_executor/models/deepseek_vl.py @@ -38,6 +38,11 @@ class DeepseekVLForCausalLM(DeepseekForCausalLM): vision_source=model.vision_pipeline.aligner, ) + @classmethod + def _neuron_host_resident(cls, module_name: str) -> bool: + # The vision tower runs here: only the token embedding stays on the host. + return super()._neuron_host_resident(module_name) and not module_name.startswith("vision_pipeline") + @staticmethod def _split_cached_prefix( segment_input: "SegmentInput", diff --git a/deepinfer/model_executor/models/v4_modules/engram.py b/deepinfer/model_executor/models/v4_modules/engram.py index 0f30312..314de18 100644 --- a/deepinfer/model_executor/models/v4_modules/engram.py +++ b/deepinfer/model_executor/models/v4_modules/engram.py @@ -27,6 +27,9 @@ from deepinfer.model_executor.parallel_utils.userbuffer import userbuffer def _pack_ue8m0_to_int(x: torch.Tensor) -> torch.Tensor: """Pack the ue8m0 scales represented in fp32 into int32 (CUDA) / int16 (ASC), 4/2 per group.""" assert x.dtype == torch.float32 and x.size(-1) % 4 == 0 + if PLATFORM.is_neuron and x.device.type != "cpu": + # One-time load step on the host: eager device ops on this large tensor compile a big NEFF per shape. + return _pack_ue8m0_to_int(x.cpu()).to(x.device) x_int32 = x.view(torch.int32) assert (x_int32 >= 0).all() and (x_int32 & 0x7FFFFF == 0).all() return torch.floor_divide(x_int32, 1 << 23).to(torch.uint8).view(PLATFORM.fp8_sf_type) diff --git a/deepinfer/model_executor/models/v4_modules/mlp.py b/deepinfer/model_executor/models/v4_modules/mlp.py index c8ed50a..af26b66 100644 --- a/deepinfer/model_executor/models/v4_modules/mlp.py +++ b/deepinfer/model_executor/models/v4_modules/mlp.py @@ -190,6 +190,12 @@ class DeepseekMoE(nn.Module): and PG.is_within_nvl_domain(PG.expert_parallel_group()) ) + if PLATFORM.is_neuron: + # A registered buffer (see neuron_rank_tensor): Dynamo lifts buffers as graph inputs, while a plain tensor + # attribute was constant-folded into the graph (one NEFF per rank again). + self.register_buffer( + "_neuron_rank_t", torch.tensor([PG.expert_parallel_group().rank()], dtype=torch.int32), persistent=False + ) self.ep_buffer: EPBuffer | None = None self.sym_buffer: deep_gemm.SymmBuffer | None = None self.ep_group = PG.expert_parallel_group() @@ -326,6 +332,11 @@ class DeepseekMoE(nn.Module): use_fp8_dispatch=self.use_fp8_act, ) + def neuron_rank_tensor(self) -> torch.Tensor: + """This rank's EP index as a device tensor: compiled graphs read it as an input, so every rank traces the + same graph and neuronx-cc compiles it once (a Python int would be baked in, one NEFF per rank).""" + return self._neuron_rank_t + def _eplb_layer(self) -> eplb.ExpertLoadBalancerLayerState: return eplb.current().layer(self.layer_idx) @@ -379,21 +390,23 @@ class DeepseekMoE(nn.Module): else: gate_force_random = input_metadata.gate_force_random - eplb_layer = self._eplb_layer() + # Neuron without redundant experts: the logical->physical map is the identity, so skip it. Reading it indexes + # a Python list by layer_idx, which made every layer's MoE region graph its own (Dynamo guard on layer_idx). + eplb_layer = None if PLATFORM.is_neuron and self.num_extra_experts == 0 else self._eplb_layer() gate_kwargs = dict( num_topk=self.num_topk, use_shared_as_routed=False, num_shared_experts=0, routed_scaling_factor=self.routed_scaling_factor, - ep_rank=self.ep_rank, + ep_rank=self.neuron_rank_tensor() if PLATFORM.is_neuron else self.ep_rank, mask=input_metadata.valid_token_mask, bias=self.e_score_correction_bias, image_bias=self.e_score_correction_bias_vl if images_seq_mask is not None else None, image_token_mask=images_seq_mask, force_random=gate_force_random, - to_physical_map=eplb_layer.logical_to_physical_map, - logical_count=eplb_layer.logical_count, + to_physical_map=None if eplb_layer is None else eplb_layer.logical_to_physical_map, + logical_count=None if eplb_layer is None else eplb_layer.logical_count, scoring_func=self.scoring_func, out=out, ) diff --git a/deepinfer/neuron_ops/fp8.py b/deepinfer/neuron_ops/fp8.py index 4fd1bd1..8ebb669 100644 --- a/deepinfer/neuron_ops/fp8.py +++ b/deepinfer/neuron_ops/fp8.py @@ -15,9 +15,11 @@ identically whenever the exponent field is below 15, i.e. for |x| < 256. So on N from __future__ import annotations +import functools import os import shlex +import numpy as np import torch FP8_MAX = float(os.environ.get("DEEPINFER_NEURON_FP8_MAX", "224")) @@ -65,13 +67,37 @@ _CHUNK_BYTES = 1 << 26 def _halve_inplace(data: torch.Tensor) -> None: - flat = data.view(-1) - for start in range(0, flat.numel(), _CHUNK_BYTES): # bounded int16 temporaries (Engram tables are ~200 GB) + """In place, on the host in fixed-size chunks (bounded int16 temporaries: the Engram tables are ~200 GB). + Device weights make one host round trip: eager device ops recompile per shape / view offset, which made this + step take over half an hour for the full model.""" + if data.device.type != "cpu": + host = data.cpu() + _halve_inplace(host) + data.copy_(host.to(data.device)) + return + flat = data.view(-1).view(torch.uint8) + table = _halve_table().numpy() + for start in range(0, flat.numel(), _CHUNK_BYTES): chunk = flat[start : start + _CHUNK_BYTES] - chunk.copy_(halve_e4m3fn_bytes(chunk)) + # A clone per chunk: .numpy() on the parameter itself would pin its storage as non-resizable, and the + # runtime adaptation later frees merged weights with untyped_storage().resize_(0). + arr = chunk.clone().numpy() + np.take(table, arr, out=arr) # one 256-entry lookup per byte + chunk.copy_(torch.from_numpy(arr)) + + +@functools.cache +def _halve_table() -> torch.Tensor: + # Explicitly on the host: model loading runs under a default-device context that maps to the NeuronCore. + return halve_e4m3fn_bytes(torch.arange(256, dtype=torch.int32, device="cpu").to(torch.uint8)) def _double_scale(sf: torch.Tensor) -> None: + if sf.device.type != "cpu": # small: on the host, no per-shape eager compiles on the device + host = sf.cpu() + _double_scale(host) + sf.copy_(host.to(sf.device)) + return if sf.dtype in (torch.float32, torch.bfloat16, torch.float16): sf.mul_(2) return diff --git a/deepinfer/neuron_ops/nki_gemm.py b/deepinfer/neuron_ops/nki_gemm.py index 5d29f17..6dc933d 100644 --- a/deepinfer/neuron_ops/nki_gemm.py +++ b/deepinfer/neuron_ops/nki_gemm.py @@ -47,8 +47,11 @@ def prepare_model(model: torch.nn.Module) -> int: 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 - scale = _kernel_scale(p.data, sf.data if isinstance(sf, torch.nn.Parameter) else sf) + # On the host: the scale tensors are small, and eager device ops would compile once per weight shape. + sf_host = (sf.data if isinstance(sf, torch.nn.Parameter) else sf).cpu() + scale = _kernel_scale(p.data, sf_host) if scale is not None: + scale = scale.to(p.device) p._nki_scale = scale # The weight bytes as their own uint8 tensor (same storage): a dtype bitcast of a graph input feeding the # kernel crashes neuronx-cc (InferNonlocalTensors), a uint8 input does not. diff --git a/deepinfer/neuron_ops/nki_norm.py b/deepinfer/neuron_ops/nki_norm.py index c15d2fd..c9a5146 100644 --- a/deepinfer/neuron_ops/nki_norm.py +++ b/deepinfer/neuron_ops/nki_norm.py @@ -77,11 +77,17 @@ def ue8m0_exponent(v: torch.Tensor) -> torch.Tensor: def pack_ue8m0_int16(e: torch.Tensor) -> torch.Tensor: - """int32 exponent bytes ``[..., g]`` -> int16 words ``[..., ceil(g / 2)]``, low byte first (arithmetic only).""" + """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)) - v = e[..., 0::2] + e[..., 1::2] * 256 - return torch.where(v >= 32768, v - 65536, v).to(torch.int16) + 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( diff --git a/deepinfer/neuron_ops/regions.py b/deepinfer/neuron_ops/regions.py index 8a1b684..d5795b0 100644 --- a/deepinfer/neuron_ops/regions.py +++ b/deepinfer/neuron_ops/regions.py @@ -7,6 +7,7 @@ fullgraph. Layers of the same kind trace to the same graph, so neuronx-cc's NEFF from __future__ import annotations +import functools from collections.abc import Callable from typing import Any @@ -33,25 +34,33 @@ def raise_dynamo_limits() -> None: setattr(torch._dynamo.config, name, 1 << 30) +_ENTRIES: dict[tuple[type, str], Callable[..., Any]] = {} + + def entry(module: torch.nn.Module, name: str) -> Callable[..., Any]: - """``module.`` as its own fullgraph region, compiled once per module (cached on it).""" - cache = module.__dict__.setdefault("_neuron_region_entries", {}) - fn = cache.get(name) + """``type(module).`` as a region, compiled once per (class, method) and called with the module. + + One compiled callable for all layers of a class: Dynamo then reuses the same graph (and the same loaded NEFF) + for every layer whose guards match (same layer kind and shapes). A per-module compile gave every layer its own + compiled artifact and NEFF load -- the same code loaded 43 times filled ~6 GB of HBM with model code. + """ + key = (type(module), name) + fn = _ENTRIES.get(key) if fn is None: from deepinfer.neuron_ops.transfer_to_neuron import hide_cuda raise_dynamo_limits() # fullgraph=False: split() marks sub-region boundaries inside the entry (smaller graphs for neuronx-cc). compiled = torch.compile( - getattr(module, name), backend=envs.DEEPINFER_NEURON_COMPILE_BACKEND, fullgraph=False, dynamic=False + getattr(type(module), name), backend=envs.DEEPINFER_NEURON_COMPILE_BACKEND, fullgraph=False, dynamic=False ) - def fn(*args: Any, **kwargs: Any) -> Any: + def fn(mod: torch.nn.Module, *args: Any, **kwargs: Any) -> Any: with hide_cuda(): - return compiled(*args, **kwargs) + return compiled(mod, *args, **kwargs) - cache[name] = fn - return fn + _ENTRIES[key] = fn + return functools.partial(fn, module) @torch._dynamo.disable diff --git a/deepinfer/neuron_ops/static_moe.py b/deepinfer/neuron_ops/static_moe.py index 0fd96c4..8745ab0 100644 --- a/deepinfer/neuron_ops/static_moe.py +++ b/deepinfer/neuron_ops/static_moe.py @@ -148,7 +148,10 @@ def _nki_moe_forward( num_local = moe.num_local_physical_experts num_tokens = x.shape[0] cap = max(moe.max_num_decoding_tokens, num_tokens) - idx, w = topk_idx.to(torch.int32), topk_weights.to(torch.float32) + x_local = x + # Local expert ids relative to this rank, from a rank *tensor*: the graph stays identical across ranks. + idx = topk_idx.to(torch.int32) + w = topk_weights.to(torch.float32) if size > 1: pad = cap - num_tokens x = torch.nn.functional.pad(x, (0, 0, 0, pad)) @@ -156,8 +159,9 @@ def _nki_moe_forward( w = torch.nn.functional.pad(w, (0, 0, 0, pad)) x, idx, w = (_all_gather(t, group, size) for t in (x, idx, w)) acc = torch.zeros((x.shape[0], moe.hidden_size), dtype=torch.float32, device=x.device) + local_idx = idx - moe.neuron_rank_tensor() * num_local for e in range(num_local): - w_e = torch.where(idx == rank * num_local + e, w, torch.zeros_like(w)).sum(-1) + w_e = torch.where(local_idx == e, w, torch.zeros_like(w)).sum(-1) gate_up = nki_gemm.linear_qdq(x, moe.gate_up_proj_weight, group=e) act = swiglu_forward_and_per_token_cast( gate_up, fmt="bf16", num_per_channels=None, topk_weights=w_e, clamp_value=moe.swiglu_clamp_value @@ -165,8 +169,7 @@ def _nki_moe_forward( acc = acc + nki_gemm.linear_qdq(act, moe.down_proj_weight, group=e).float() out = _reduce_scatter(acc, group, size, rank)[:num_tokens] if moe.has_shared: - shared = nki_gemm.linear_qdq(x[:num_tokens] if size == 1 else x.narrow(0, rank * cap, num_tokens), - moe.shared_gate_up_proj.weight) + shared = nki_gemm.linear_qdq(x_local, moe.shared_gate_up_proj.weight) shared = swiglu_forward_and_per_token_cast( shared, fmt="bf16", num_per_channels=None, clamp_value=moe.swiglu_clamp_value ) diff --git a/deepinfer/neuron_ops/vendor/tile_kernels/moe/moe_topk_gate_forward_kernel.py b/deepinfer/neuron_ops/vendor/tile_kernels/moe/moe_topk_gate_forward_kernel.py index 93c1dc9..446f32f 100644 --- a/deepinfer/neuron_ops/vendor/tile_kernels/moe/moe_topk_gate_forward_kernel.py +++ b/deepinfer/neuron_ops/vendor/tile_kernels/moe/moe_topk_gate_forward_kernel.py @@ -22,7 +22,7 @@ from tile_kernels.moe.sqrt_softplus import sqrt_softplus _LARGE_PRIME = 23333 -def _hash_uniform(a: torch.Tensor, b: torch.Tensor, seed: int) -> torch.Tensor: +def _hash_uniform(a: torch.Tensor, b: torch.Tensor, seed) -> torch.Tensor: """Deterministic uniform [0, 1) from int32 coordinates (an integer hash; no generator, traceable).""" h = a * 73856093 + b * 19349663 + seed * 83492791 h = h ^ (h >> 13) @@ -31,14 +31,19 @@ def _hash_uniform(a: torch.Tensor, b: torch.Tensor, seed: int) -> torch.Tensor: return (h & 0xFFFFFF).to(torch.float32) * (1.0 / (1 << 24)) +def _seed(ep_rank): + # A tensor rank keeps compiled graphs identical across ranks (one NEFF for all instead of one per rank). + return ep_rank.to(torch.int32) if isinstance(ep_rank, torch.Tensor) else int(ep_rank) + + def _force_random(num_tokens, num_physical_topk, num_physical_experts, max_physical_experts, ep_rank, device): # NEURON: random routing for the forced tokens from an integer hash instead of a seeded CPU generator, so it # stays on the device and inside the decode graph. token = torch.arange(num_tokens, dtype=torch.int32, device=device)[:, None] cols = torch.arange(max_physical_experts, dtype=torch.int32, device=device) - keys = _hash_uniform(token, cols[None, :], int(ep_rank)) + keys = _hash_uniform(token, cols[None, :], _seed(ep_rank)) weights = _hash_uniform(token, torch.arange(num_physical_topk, dtype=torch.int32, device=device)[None, :] + max_physical_experts, - int(ep_rank)) + _seed(ep_rank)) keys = torch.where(cols[None, :] < num_physical_experts, keys, torch.full_like(keys, 2.0)) idx = torch.sort(keys, dim=1, stable=True).indices[:, :num_physical_topk] return idx.to(torch.int64), weights diff --git a/deepinfer/worker/base.py b/deepinfer/worker/base.py index 158e26a..31fc4e4 100644 --- a/deepinfer/worker/base.py +++ b/deepinfer/worker/base.py @@ -18,7 +18,7 @@ import numpy as np import torch import torch.distributed as dist -from deepinfer import envs +from deepinfer import PLATFORM, envs from deepinfer._C.engine import clone_and_zero_states_and_blocks, set_engram_lookup_table from deepinfer.config import CacheConfig, ModelConfig, ParallelConfig, SchedulerConfig from deepinfer.engine.engine_type import EngineType @@ -163,6 +163,8 @@ class WorkerBase: clone_blocks: np.ndarray[np.int32], zero_states: np.ndarray[np.int32], ) -> tuple[np.ndarray[np.int32], np.ndarray[np.int32]]: + if PLATFORM.is_neuron: + return _neuron_clone_and_zero(self.cache_engine.gpu_cache, clone_states, clone_blocks, zero_states) return clone_and_zero_states_and_blocks( self.cache_engine.gpu_states, self.cache_engine.gpu_blocks, clone_states, clone_blocks, zero_states ) @@ -355,3 +357,38 @@ def _bind_numa(local_rank: int) -> None: bind_numa(which_numa(local_rank)) except Exception: logger.warning("NUMA 绑定失败,内存落位退化为 first-touch,可能导致 H2D 性能下降", exc_info=True) + + +def _neuron_clone_and_zero( + gpu_cache: object, clone_states: np.ndarray, clone_blocks: np.ndarray, zero_states: np.ndarray +) -> tuple[np.ndarray, np.ndarray]: + """``clone_and_zero_states_and_blocks`` for Neuron's layer-major cache layout (cache_engine._LayerMajorColumns): + a block or state slot is spread over every layer's view, so copy / zero row ``i`` of each view instead of one + contiguous block. Same contract as the C++ version: returns the cloned source ids (states, blocks).""" + block_views, state_views, seen = [], [], set() + for cache in gpu_cache.cache_list: + for name in ("ica_cache", "ca_cache"): + t = getattr(cache, name) + if t is not None and id(t) not in seen: + seen.add(id(t)) + block_views.append(t) + for name in ("swa_cache", "ica_state", "ca_state"): + t = getattr(cache, name) + if t is not None and id(t) not in seen: + seen.add(id(t)) + state_views.append(t) + released_states, released_blocks = [], [] + for src, dst in clone_states.reshape(-1, 2) if clone_states.size else []: + if dst: + for t in state_views: + t[int(dst)].copy_(t[int(src)]) + released_states.append(int(src)) + for src, dst in clone_blocks.reshape(-1, 2) if clone_blocks.size else []: + if dst: + for t in block_views: + t[int(dst)].copy_(t[int(src)]) + released_blocks.append(int(src)) + for i in zero_states.reshape(-1) if zero_states.size else []: + for t in state_views: + t[int(i)].zero_() + return np.asarray(released_states, dtype=np.int32), np.asarray(released_blocks, dtype=np.int32) diff --git a/deepinfer/worker/cache_engine.py b/deepinfer/worker/cache_engine.py index 83ee80f..ee98852 100644 --- a/deepinfer/worker/cache_engine.py +++ b/deepinfer/worker/cache_engine.py @@ -306,6 +306,50 @@ class StateBlock: self.state_size_to_disk = self.swa_store_total_bytes_inference_pad +class _LayerMajorColumns: + """Neuron: a ``[rows, cols]`` byte buffer laid out column-region-major (vllm-neuron's per-layer layout). + + The engine slices each layer's cache as a column range ``[:, a:b]`` of a block-major (or state-slot-major) buffer. + Here the same bytes hold the range ``[a, b)`` as one contiguous ``[rows, b - a]`` region at ``flat[a*rows:b*rows]``: + every layer view is then contiguous and has no per-layer offset inside it, so region graphs take the views as + inputs and same-kind layers share one compiled NEFF (strided column views were made contiguous -- copied -- by + the executor on every call, and slicing them in-graph baked each layer's offsets into its own NEFF). + """ + + def __init__(self, flat: torch.Tensor, rows: int, col0: int, cols: int) -> None: + self.flat, self.rows, self.col0, self.cols = flat, rows, col0, cols + + @classmethod + def of(cls, t: torch.Tensor) -> _LayerMajorColumns: + assert t.dim() == 2 and t.dtype == torch.uint8 and t.is_contiguous() + return cls(t.view(-1), t.shape[0], 0, t.shape[1]) + + @property + def shape(self) -> tuple[int, int]: + return (self.rows, self.cols) + + def stride(self, dim: int | None = None) -> int | tuple[int, int]: + return (self.cols, 1) if dim is None else (self.cols, 1)[dim] + + def __getitem__(self, idx: tuple[slice, slice]) -> _LayerMajorColumns: + rows, cols = idx + assert rows == slice(None), idx + a, b, step = cols.indices(self.cols) + assert step == 1 + return _LayerMajorColumns(self.flat, self.rows, self.col0 + a, b - a) + + def tensor(self) -> torch.Tensor: + start = self.col0 * self.rows + return self.flat[start : start + self.cols * self.rows].view(self.rows, self.cols) + + def view(self, *shape: int) -> torch.Tensor: + return self.tensor().view(*shape) + + def as_strided(self, size: tuple[int, ...], stride: tuple[int, ...]) -> torch.Tensor: + # The block-major strides asked for describe this same logical layout; the region is contiguous. + return self.tensor().view(*size) + + class KVCacheManager: """管理 KV cache 的分配和布局。 @@ -315,6 +359,10 @@ class KVCacheManager: def __init__(self, cfg: AttrDict, kv_tensor: torch.Tensor | None, state_tensor: torch.Tensor | None) -> None: self.cfg = cfg + if PLATFORM.is_neuron and isinstance(kv_tensor, torch.Tensor): + kv_tensor = _LayerMajorColumns.of(kv_tensor) + if PLATFORM.is_neuron and isinstance(state_tensor, torch.Tensor): + state_tensor = _LayerMajorColumns.of(state_tensor) self.kv_tensor = kv_tensor model_metadata = ModelMetadata(cfg) diff --git a/deepinfer/worker/neuron_graph_runner.py b/deepinfer/worker/neuron_graph_runner.py index 93cc267..8b68edc 100644 --- a/deepinfer/worker/neuron_graph_runner.py +++ b/deepinfer/worker/neuron_graph_runner.py @@ -41,7 +41,12 @@ class _GraphKVCaches: """ def __init__(self, kv_caches: GPUCacheInterface) -> None: + self.cache_list = kv_caches.cache_list self.roots: list[torch.Tensor] = [] + if envs.DEEPINFER_NEURON_GRAPH_MODE == "regions": + # The cache engine lays every layer view out contiguously on Neuron (_LayerMajorColumns); the region + # graphs take them as inputs as they are. + return self.recipes: list[dict[str, tuple | None]] = [] for cache in kv_caches.cache_list: recipe = {} @@ -82,15 +87,16 @@ class _GraphKVCaches: def __getitem__(self, layer_idx: int) -> KVCache | _LazyKVCache: if envs.DEEPINFER_NEURON_GRAPH_MODE == "regions": - # Views built on attribute access, i.e. inside the layer's region graph (roots are its inputs). - return _LazyKVCache(self, layer_idx) + # The layer's own contiguous, non-overlapping views as region graph inputs (like vllm-neuron): slicing + # them from the roots inside the graph baked each layer's byte offsets in, one NEFF per layer. + return self.cache_list[layer_idx] tensors = {} for name, recipe in self.recipes[layer_idx].items(): tensors[name] = None if recipe is None else self._build(self.roots[recipe[0]], recipe) return KVCache(**tensors) def is_valid_layeridx(self, layer_idx: int) -> bool: - return layer_idx < len(self.recipes) + return layer_idx < len(self.cache_list) class _LazyKVCache: diff --git a/examples/neuron/offline_inference_neuron.py b/examples/neuron/offline_inference_neuron.py index 16e35df..ef375fb 100644 --- a/examples/neuron/offline_inference_neuron.py +++ b/examples/neuron/offline_inference_neuron.py @@ -40,6 +40,8 @@ parser.add_argument("--max_tokens", default=16, type=int) parser.add_argument("--num_kv_blocks", default=400, type=int, help="fixed KV pool per rank; the scheduler admits a request only if max_model_len tokens fit, per running sequence") parser.add_argument("--max_model_len", default=16384, type=int, help="the candidate-block filter needs >= candidate_topk_blocks * candidate_block_size (2048 * 8)") parser.add_argument("--max_num_batched_tokens", default=1024, type=int) +parser.add_argument("--max_num_seqs", default=1, type=int, + help="decode batch; every captured batch size adds one set of region NEFFs per layer (HBM: model code)") args = parser.parse_args() dp_size = args.ep @@ -77,7 +79,7 @@ def main(dp_size: int, dp_rank: int) -> None: data_parallel_size=dp_size, max_num_batched_tokens=args.max_num_batched_tokens, max_model_len=args.max_model_len, - max_num_seqs=len(PROMPTS), + max_num_seqs=args.max_num_seqs, gpu_memory_utilization=0.85, lookahead_tokens=0, ) diff --git a/tests/neuron/compile/test_nki_hc.py b/tests/neuron/compile/test_nki_hc.py index 66f19de..658ef4b 100644 --- a/tests/neuron/compile/test_nki_hc.py +++ b/tests/neuron/compile/test_nki_hc.py @@ -345,3 +345,11 @@ def test_shifted_post_pre_chain(cfg: str, n: int, mode: str) -> None: close(got[4], ref[4]) npc = get_fp8_cast_params()["num_per_channels"] torch.testing.assert_close(_dequant(got[5], got[6], npc), _dequant(ref[5], ref[6], npc), rtol=0.13, atol=1e-2) + + +def test_pack_ue8m0_int16_exhaustive() -> None: + """Every exponent byte pair packs to the same int16 word as the byte view (host only).""" + e = torch.stack(torch.meshgrid(torch.arange(256), torch.arange(256), indexing="ij"), -1).reshape(-1, 2).to(torch.int32) + words = nki_norm.pack_ue8m0_int16(e) + assert words.dtype == torch.int16 and words.shape == (65536, 1) + assert torch.equal(words.contiguous().view(torch.uint8).view(-1, 2).to(torch.int32), e)