diff --git a/iron/operators/strided_copy/design.py b/iron/operators/strided_copy/design.py index 4e71bc5f5..b43ee1810 100644 --- a/iron/operators/strided_copy/design.py +++ b/iron/operators/strided_copy/design.py @@ -86,12 +86,9 @@ def strided_copy( np.dtype[dtype], ] - # input_offset_parameter (and output_offset_parameter) is the name of an - # aiex.scratchpad_parameter used to patch the DMA BD base address at runtime. The - # statically-computed offset is used as the base; the parameter's value is - # additively combined onto it inside the BD address registers via UPDATE_REG. - # The host writes the byte offset into the ctrl scratchpad before each - # dispatch via ParameterScratchpad. + # Patches the DMA BD base address at runtime: the static offset is the base, + # and UPDATE_REG adds the parameter onto it. The host writes an element count, + # which the firmware scales by elemBytes. in_offset_param = ( ScratchpadParameter(input_offset_parameter, np.int32) if input_offset_parameter is not None diff --git a/iron/operators/strided_copy/test.py b/iron/operators/strided_copy/test.py index e36328105..5187eebc7 100644 --- a/iron/operators/strided_copy/test.py +++ b/iron/operators/strided_copy/test.py @@ -2,11 +2,14 @@ # SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import numpy as np import pytest +import torch from iron.operators.strided_copy.op import StridedCopy from iron.operators.strided_copy.reference import generate_golden_reference -from iron.common.test_utils import run_test +from iron.common.sequence import OperatorSequence +from iron.common.test_utils import run_test, verify_buffer # Llama's KV-cache write, shrunk: the cache is (n_kv_groups, seq, head_dim) and one # token's keys land in slot t of every group. SEQ is 128 rather than the real 2048 to @@ -95,6 +98,53 @@ def test_strided_copy(kwargs, aie_context): assert not errors, f"Test failed with errors: {errors}" +@pytest.mark.supported_devices("npu2") +def test_strided_copy_cache_offset_parameter(aie_context): + """dispatch="fused" is the only mode with a ctrl scratchpad. + + The whole cache is checked rather than the target slot: a mis-scaled addend + (elements vs bytes) lands the write in a different slot. + """ + base_kwargs = _kv_slot(SEQ, 0) + op = StridedCopy( + **base_kwargs, output_offset_parameter="cache_offset", context=aie_context + ) + seq = OperatorSequence( + "strided_copy_cache_offset_cov", + [(op, "in", "out")], + input_args=["in"], + output_args=["out"], + dispatch="fused", + context=aie_context, + ) + seq.compile() + run = seq.get_callable() + assert run.params is not None, "cache_offset did not produce a ctrl scratchpad" + + out_buf = run.get_buffer("out") + expected = torch.zeros(base_kwargs["output_buffer_size"], dtype=torch.bfloat16) + out_buf.torch_view()[: expected.numel()] = expected + out_buf.to("npu") + + for i, slot in enumerate((0, 5, SEQ - 1)): + golden = generate_golden_reference( + **base_kwargs, output_offset_addend=slot * HEAD_DIM, seed=i + 1 + ) + expected += golden["output"] + + in_buf = run.get_buffer("in") + in_buf.torch_view()[: golden["input"].numel()] = golden["input"] + in_buf.to("npu") + + run.params.write("cache_offset", np.int32(slot * HEAD_DIM)) + run.params.sync() + run() + + out = out_buf.torch_view()[: expected.numel()].clone() + errors = verify_buffer(out, "out", expected, rel_tol=0.0, abs_tol=0.0) + assert not errors, f"slot {slot}: {errors}" + + def test_transfer_size_not_dividing_per_channel_share_is_rejected(aie_context): """A BD shorter than the ObjectFifo object hangs the device, so it must not compile.