From 41bd1a884b5aeb993e49b218c0443467465912a3 Mon Sep 17 00:00:00 2001 From: hgt312 Date: Tue, 31 Mar 2026 12:12:13 +0000 Subject: [PATCH] fix(trace): canonicalize non-aliased output names Always lower non-aliased results as output{idx}, even when tracing assigned an intermediate name. Keep alias naming in NKIPy lowering, add direct NEFF naming coverage, and improve Spike I/O mismatch diagnostics without teaching Spike alias policy. --- nkipy/src/nkipy/core/trace.py | 8 +- spike/src/spike/spike_model.py | 20 +++++ tests/unit/test_alias.py | 145 +++++++++++++++++++++++++++++++++ 3 files changed, 168 insertions(+), 5 deletions(-) diff --git a/nkipy/src/nkipy/core/trace.py b/nkipy/src/nkipy/core/trace.py index 2d44a4b6..b3fc5b03 100644 --- a/nkipy/src/nkipy/core/trace.py +++ b/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] diff --git a/spike/src/spike/spike_model.py b/spike/src/spike/spike_model.py index b9aa42fe..8726322b 100644 --- a/spike/src/spike/spike_model.py +++ b/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, ( diff --git a/tests/unit/test_alias.py b/tests/unit/test_alias.py index ebc5f471..728d08bd 100644 --- a/tests/unit/test_alias.py +++ b/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"])