strided_copy: cover the runtime offset parameter, fix its comment - #168
Open
atassis wants to merge 1 commit into
Open
strided_copy: cover the runtime offset parameter, fix its comment#168atassis wants to merge 1 commit into
atassis wants to merge 1 commit into
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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(aParameterScratchpad), with the staticoutput_offsetas the base. #158 (now merged) names this as a known gap in its own body: every arm it added bakes the offset in atcompile 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 asingle-entry
OperatorSequence(dispatch="fused")(the only dispatch modewith a ctrl scratchpad) and drives
cache_offsetthroughParameterScratchpadacross three token positions on one compiled programand 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 --checkandreuse lintpass onboth 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'sgenerate_golden_referencefor the same three slots usedin 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/onself-hosted NPU2 hardware, which is where
supported_devices("npu2")letsthis arm actually execute; I have not seen that run.