Fix alias output naming and improve Spike I/O diagnostics

nkipy/src/nkipy/core/trace.py

@@ -235,7 +235,8 @@ def _mark_hlo_outputs(self, code: HLOModule, ret, param_tensor_refs):
                 copy_tensor = ctx.build_op("copy", [bt], bt.shape, bt.dtype)
                 ret[i] = NKIPyTensorRef(copy_tensor, name="")
 
-        # Assign output names and build AliasInfo list
+        # Assign canonical names and build AliasInfo metadata. Alias outputs
+        # keep the parameter name; every other output uses its tuple position.
         for idx, r in enumerate(ret):
             if not isinstance(r, NKIPyTensorRef):
                 raise RuntimeError(f"Unexpected return value type: {type(r)}")
@@ -251,10 +252,7 @@ def _mark_hlo_outputs(self, code: HLOModule, ret, param_tensor_refs):
                     )
                 )
                 r.backend_tensor.name = param_name
-
-            # N.B.: the name "output{idx}" is specific
-            # it avoids variable folding in HLO lowering in Neuron Compiler
-            if not r.backend_tensor.name:
+            else:
                 r.backend_tensor.name = f"output{idx}"
 
         result_tensors = [r.backend_tensor for r in ret]

spike/src/spike/spike_model.py

@@ -196,7 +196,27 @@ def _check_dtype_compatibility(
             )
 
     def _validate_io(self, inputs, outputs):
+        """Validate that caller-supplied I/O dicts match the compiled NEFF.
+
+        Checks tensor names, shapes, dtypes, and core placement.  Raises
+        ``ValueError`` with the expected NEFF names on any name mismatch,
+        so callers get actionable diagnostics instead of a bare ``KeyError``.
+        """
         model_core_id = self.model_ref.core_id
+
+        unknown_inputs = set(inputs) - set(self.input_tensors_info)
+        if unknown_inputs:
+            raise ValueError(
+                f"Unknown input(s) {unknown_inputs} for model '{self.name}'. "
+                f"Expected inputs: {list(self.input_tensors_info.keys())}"
+            )
+        unknown_outputs = set(outputs) - set(self.output_tensors_info)
+        if unknown_outputs:
+            raise ValueError(
+                f"Unknown output(s) {unknown_outputs} for model '{self.name}'. "
+                f"Expected outputs: {list(self.output_tensors_info.keys())}"
+            )
+
         for k, v in inputs.items():
             tensor_core_id = v.tensor_ref.core_id
             assert tensor_core_id == model_core_id, (

tests/unit/test_alias.py

@@ -31,6 +31,13 @@ def nkipy_kernel_multi_alias(a_input, b_input, c_input):
     return a_input, c_input
 
 
+def nkipy_kernel_named_intermediate(a_input, b_input):
+    """Kernel that returns an op result with an existing intermediate name."""
+    out = np.add(a_input, b_input)
+    out.backend_tensor.name = "intermediate0"
+    return out
+
+
 def nkipy_kernel_no_return(a_input, b_input):
     """Kernel that mutates a_input but does not return anything."""
     a_input[0, :] = b_input[1, :]
@@ -204,5 +211,143 @@ def test_mixed_return_alias(trace_mode):
         trace_and_compile(nkipy_kernel_mixed_return, trace_mode, A.copy(), B)
 
 
+def test_non_alias_outputs_are_renamed_to_output_names():
+    """Non-aliased outputs must be renamed even if tracing assigned a temp name."""
+
+    a = ((np.random.rand(128, 512) - 0.5) * 2).astype(np.float16)
+    b = ((np.random.rand(128, 512) - 0.5) * 2).astype(np.float16)
+
+    traced = NKIPyKernel.trace(nkipy_kernel_named_intermediate, backend="hlo")
+    hlo = traced.specialize(a, b)
+
+    assert len(hlo.outputs) == 1
+    assert hlo.outputs[0].name == "output0"
+
+
+# ------------------------------------------------------------------ #
+# Output naming contract tests
+#
+# These verify the exact I/O names that end up in compiled NEFFs,
+# which callers must match when invoking kernel(inputs={...}, outputs={...}).
+# The patterns below cover callers that bind directly to compiled NEFF names.
+# ------------------------------------------------------------------ #
+
+
+def test_alias_output_naming_simple():
+    """Direct alias: mutated param returned by identity keeps param name.
+
+    Pattern: update_kv_cache(kv_cache: mutable) -> kv_cache
+    Expected NEFF:
+        input  = "a_input.must_alias_input"
+        output = "a_input"
+    """
+    a = ((np.random.rand(128, 512) - 0.5) * 2).astype(np.float16)
+    b = ((np.random.rand(128, 512) - 0.5) * 2).astype(np.float16)
+
+    traced = NKIPyKernel.trace(nkipy_kernel_single_alias, backend="hlo")
+    hlo = traced.specialize(a, b)
+
+    # 1 output: the aliased param
+    assert len(hlo.outputs) == 1
+    assert hlo.outputs[0].name == "a_input"
+
+    # Input param renamed with alias suffix
+    param_names = [p.name for p in hlo.parameters]
+    assert "a_input.must_alias_input" in param_names
+    assert "b_input" in param_names
+
+
+def test_alias_auto_appended_output_naming():
+    """Broken-identity alias: mutated param NOT returned → auto-appended.
+
+    Pattern: prefill_post_moe_fn(output: mutable, ...) where output is mutated
+    but the function returns a *different* computed value. The tracer
+    auto-appends the original mutated param as an extra output.
+
+    Expected NEFF:
+        inputs  = "a_input.must_alias_input", "b_input"
+        outputs = "output0" (the sum), "a_input" (auto-appended alias)
+    """
+    a = ((np.random.rand(128, 512) - 0.5) * 2).astype(np.float16)
+    b = ((np.random.rand(128, 512) - 0.5) * 2).astype(np.float16)
+
+    traced = NKIPyKernel.trace(nkipy_kernel_mixed_return, backend="hlo")
+    hlo = traced.specialize(a, b)
+
+    # 2 outputs: user-returned sum (output0) + auto-appended alias (a_input)
+    assert len(hlo.outputs) == 2
+    assert hlo.outputs[0].name == "output0"
+    assert hlo.outputs[1].name == "a_input"
+
+    # Alias metadata
+    assert len(hlo.aliases) == 1
+    assert hlo.aliases[0].param_name == "a_input"
+    assert hlo.aliases[0].output_index == 1
+    assert hlo.aliases[0].is_user_returned is False
+
+    # Input param renamed
+    param_names = [p.name for p in hlo.parameters]
+    assert "a_input.must_alias_input" in param_names
+
+
+def nkipy_kernel_alias_with_multiple_outputs(a_input, b_input, c_input):
+    """Kernel that aliases a_input and c_input, returns them plus a computed value.
+
+    Pattern: fused pre_moe graph returning (kv_cache, hidden, topk, ...)
+    where kv_cache is aliased but hidden/topk/... are not.
+    """
+    a_input[0:1, :] = b_input[0:1, :]
+    c_input[2:3, :] = b_input[2:3, :]
+    computed = np.add(a_input, b_input)
+    return a_input, computed, c_input
+
+
+def test_alias_mixed_with_non_alias_outputs():
+    """Multiple outputs where some are aliased and some are not.
+
+    Expected NEFF:
+        inputs  = "a_input.must_alias_input", "b_input", "c_input.must_alias_input"
+        outputs = "a_input" (alias), "output1" (computed), "c_input" (alias)
+    """
+    a = ((np.random.rand(128, 512) - 0.5) * 2).astype(np.float16)
+    b = ((np.random.rand(128, 512) - 0.5) * 2).astype(np.float16)
+    c = ((np.random.rand(128, 512) - 0.5) * 2).astype(np.float16)
+
+    traced = NKIPyKernel.trace(nkipy_kernel_alias_with_multiple_outputs, backend="hlo")
+    hlo = traced.specialize(a, b, c)
+
+    assert len(hlo.outputs) == 3
+    # Aliased outputs keep param names; non-aliased get output{idx}
+    assert hlo.outputs[0].name == "a_input"
+    assert hlo.outputs[1].name == "output1"
+    assert hlo.outputs[2].name == "c_input"
+
+    # 2 aliases
+    assert len(hlo.aliases) == 2
+    alias_names = {a.param_name for a in hlo.aliases}
+    assert alias_names == {"a_input", "c_input"}
+
+
+def test_no_return_alias_output_naming():
+    """Mutation-only kernel: auto-appended alias is the sole output.
+
+    Expected NEFF:
+        input  = "a_input.must_alias_input", "b_input"
+        output = "a_input"
+    """
+    a = ((np.random.rand(128, 512) - 0.5) * 2).astype(np.float16)
+    b = ((np.random.rand(128, 512) - 0.5) * 2).astype(np.float16)
+
+    traced = NKIPyKernel.trace(nkipy_kernel_no_return, backend="hlo")
+    hlo = traced.specialize(a, b)
+
+    assert len(hlo.outputs) == 1
+    assert hlo.outputs[0].name == "a_input"
+
+    param_names = [p.name for p in hlo.parameters]
+    assert "a_input.must_alias_input" in param_names
+    assert "b_input" in param_names
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])