Neuron: separate state regions, chunk-wide indexer context (issue #26)

- The state buffer's column regions get their own allocations on Neuron:
  a layer's uint8 SWA cache and fp32 compressor state are both region graph
  inputs, and aot_autograd compares offsets of differently typed views of one
  storage in their own element units (false overlap -> 'input mutations on
  views with different dtypes').
- paged_mqa_logits chunk path: gather the keys the chunk's last token sees
  (lens.max()), not the first token's.
- Debug: DEEPINFER_NEURON_ALIAS_DEBUG prints aliased graph inputs; a cProfile
  of decode step 8 under DEEPINFER_NEURON_REGION_TIMING.

Tiny V4.1 on a NeuronCore, bucket 128: short 8/9 (= eager), long 715 tokens
in 6 chunks 9/9 (eager 8/9) in 26 s instead of 978 s of eager prefill.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

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)

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)

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)

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: