diff --git a/deepinfer/neuron_ops/neuron_compile.py b/deepinfer/neuron_ops/neuron_compile.py index 1f39433..a671df6 100644 --- a/deepinfer/neuron_ops/neuron_compile.py +++ b/deepinfer/neuron_ops/neuron_compile.py @@ -279,10 +279,27 @@ def _fp8_inputs_as_bytes(gm: torch.fx.GraphModule, example_inputs: Sequence[Any] return inputs, positions +def _print_aliased_inputs(gm: torch.fx.GraphModule, example_inputs: Sequence[Any]) -> None: + """Debug: graph inputs sharing a storage (aot_autograd rejects mutated aliases of different dtypes).""" + from torch.multiprocessing.reductions import StorageWeakRef + + names = [n.target for n in gm.graph.nodes if n.op == "placeholder"] + groups: dict[Any, list[str]] = {} + for name, x in zip(names, example_inputs): + if isinstance(x, torch.Tensor): + key = StorageWeakRef(x.untyped_storage()) + groups.setdefault(key, []).append(f"{name}:{x.dtype}{tuple(x.shape)}@{x.storage_offset()}") + for members in groups.values(): + if len(members) > 1: + print("ALIASED INPUTS:", members, flush=True) + + @register_backend def deepinfer_neuron(gm: torch.fx.GraphModule, example_inputs: Sequence[Any], **kwargs: Any) -> Callable: from torch_neuronx.neuron_dynamo_backend.backend import neuron_backend + if os.environ.get("DEEPINFER_NEURON_ALIAS_DEBUG"): + _print_aliased_inputs(gm, example_inputs) inputs, positions = _fp8_inputs_as_bytes(gm, example_inputs) if legalize_fp8(gm) and os.environ.get("DEEPINFER_FP8_LEGALIZE_DEBUG"): print(gm.graph, flush=True) diff --git a/deepinfer/neuron_ops/nki_attention.py b/deepinfer/neuron_ops/nki_attention.py index 43d1da1..388c6cc 100644 --- a/deepinfer/neuron_ops/nki_attention.py +++ b/deepinfer/neuron_ops/nki_attention.py @@ -330,7 +330,7 @@ def paged_mqa_logits( # table and differ only in their context length. Gather the keys once and score every token in one kernel call # (padding rows past the chunk score garbage on the same keys; their results are never read). page = block_table[0].to(torch.int32).index_select(0, slot) - page = torch.where(pos < lens[0], page, torch.zeros_like(page)).view(1, n) + page = torch.where(pos < lens.max(), page, torch.zeros_like(page)).view(1, n) # the chunk's last token sees most offset = (pos % page_kv).view(1, n) data = gather_rows(kv_cache, page, offset, data_bytes, data_bytes).to(torch.int32) sf = gather_rows(kv_cache, page, offset, SF_BYTES, SF_BYTES, base=page_kv * data_bytes).to(torch.int32) diff --git a/deepinfer/worker/cache_engine.py b/deepinfer/worker/cache_engine.py index ee98852..fe9b771 100644 --- a/deepinfer/worker/cache_engine.py +++ b/deepinfer/worker/cache_engine.py @@ -316,13 +316,20 @@ class _LayerMajorColumns: 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: + def __init__( + self, flat: torch.Tensor, rows: int, col0: int, cols: int, regions: dict[tuple[int, int], torch.Tensor] | None = None + ) -> None: self.flat, self.rows, self.col0, self.cols = flat, rows, col0, cols + self.regions = regions @classmethod - def of(cls, t: torch.Tensor) -> _LayerMajorColumns: + def of(cls, t: torch.Tensor, separate: bool = False) -> _LayerMajorColumns: + """``separate``: give every column region its own allocation instead of a range of ``t``. The state buffer + needs it: a layer's uint8 SWA cache and fp32 compressor state are both region graph inputs, and aot_autograd + compares the storage offsets of differently typed views of one storage in their own element units, so it + sees false overlaps and rejects the mutated aliases ("input mutations on views with different dtypes").""" 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]) + return cls(t.view(-1), t.shape[0], 0, t.shape[1], {} if separate else None) @property def shape(self) -> tuple[int, int]: @@ -336,9 +343,16 @@ class _LayerMajorColumns: 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) + return _LayerMajorColumns(self.flat, self.rows, self.col0 + a, b - a, self.regions) def tensor(self) -> torch.Tensor: + if self.regions is not None: + key = (self.col0, self.cols) + if key not in self.regions: + for c0, n in self.regions: # a separate region must not overlap another one (they would diverge) + assert self.col0 + self.cols <= c0 or c0 + n <= self.col0, (key, (c0, n)) + self.regions[key] = torch.zeros((self.rows, self.cols), dtype=torch.uint8, device=self.flat.device) + return self.regions[key] start = self.col0 * self.rows return self.flat[start : start + self.cols * self.rows].view(self.rows, self.cols) @@ -362,7 +376,7 @@ class KVCacheManager: 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) + state_tensor = _LayerMajorColumns.of(state_tensor, separate=True) 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 effb38c..b0055c1 100644 --- a/deepinfer/worker/neuron_graph_runner.py +++ b/deepinfer/worker/neuron_graph_runner.py @@ -190,11 +190,26 @@ class NeuronGraphRunner(CudaGraphRunner): with hide_cuda(): return self.compiled[batch_size](input_metadata_device) torch.neuron.synchronize() + self._steps = getattr(self, "_steps", 0) + 1 + profiler = None + if self._steps == 8 and (not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0): + import cProfile + + profiler = cProfile.Profile() + profiler.enable() started = time.perf_counter() with hide_cuda(): out = self.compiled[batch_size](input_metadata_device) torch.neuron.synchronize() total = time.perf_counter() - started + if profiler is not None: + import io + import pstats + + profiler.disable() + buf = io.StringIO() + pstats.Stats(profiler, stream=buf).sort_stats("tottime").print_stats(40) + logger.info(f"Neuron step profile (rank 0, step {self._steps}):\n{buf.getvalue()}") timings = regions.take_timings() in_regions = sum(sum(v) for v in timings.values()) if torch.distributed.is_initialized() and torch.distributed.get_rank() != 0: