Skip to content

strided_copy: cover the runtime offset parameter, fix its comment - #168

Open
atassis wants to merge 1 commit into
amd:develfrom
atassis:nr-strided-copy-cov
Open

strided_copy: cover the runtime offset parameter, fix its comment#168
atassis wants to merge 1 commit into
amd:develfrom
atassis:nr-strided-copy-cov

Conversation

@atassis

@atassis atassis commented Aug 29, 2026

Copy link
Copy Markdown
Contributor

Problem

llama_npu.py:319 is strided_copy's only production caller, and it drives the
KV-cache write's offset at runtime through output_offset_parameter (a
ParameterScratchpad), with the static output_offset as the base. #158 (now merged) names this as a known gap in its own body: every arm it added bakes the offset in at
compile time instead, so the operator's only production path has no test
coverage. Separately, design.py's comment on this parameter says the host
writes a byte offset into the scratchpad; it writes an element count, and the
firmware multiplies by elemBytes.

Fix

Adds test_strided_copy_cache_offset_parameter, gated
@pytest.mark.supported_devices("npu2"). It wraps StridedCopy in a
single-entry OperatorSequence(dispatch="fused") (the only dispatch mode
with a ctrl scratchpad) and drives cache_offset through
ParameterScratchpad across three token positions on one compiled program
and one persistent run handle, the same handle reuse the decode path
depends on. Each iteration checks the full cache buffer, not just the target
slot: a mis-scaled addend lands the write in the wrong slot, which a
target-slot-only check would miss. Also corrects the design.py comment.

Test/Evidence

I did not run this on hardware. black --check and reuse lint pass on
both changed files. This op has no kernel, so the only device-free check
available is the reference math itself: a standalone script that calls
reference.py's generate_golden_reference for the same three slots used
in the test confirms their write regions are disjoint (1536 of 65536
elements, no overlap), each token's data round-trips exactly through the
accumulated buffer at its own indices, and doubling the addend (the shape a
byte-vs-element bug would take) produces a different result rather than one
silently masked by the gate.
CI's krackan-small job runs pytest -m "not extensive" iron/operators/ on
self-hosted NPU2 hardware, which is where supported_devices("npu2") lets
this arm actually execute; I have not seen that run.

llama_npu.py:319 drives the KV-cache write through
output_offset_parameter="cache_offset" at runtime; every existing arm bakes
the offset in at compile time instead, so the operator's only production
path had no coverage. amd#158 disclosed this as a known gap.

Adds a fused-dispatch OperatorSequence test that drives cache_offset through
ParameterScratchpad across three token positions on one persistent run
handle, checking the full cache buffer each time so a mis-scaled addend
cannot land in the wrong slot undetected.

Also fixes design.py's comment: the host writes an element count into the
scratchpad, not a byte offset. Left out of amd#158 to keep it scoped.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant