llm: cap decode-program layers so deep small-dim models don't hit the ANE op ceiling - #254
Merged
Merged
Conversation
… ANE op ceiling The resident-state decode program chunks layers only by baked weight bytes (_chunk_bytes). A deep but narrow model (e.g. SmolLM2-360M: 32 layers, dim 960) stays under that byte cap, so all 32 layers compiled into one program -- which then failed execute with rc=-1. The ceiling here is op-count, not bytes: it trips at ~31 layers regardless of dim (measured L30 OK / L32 fails at dim 128 and 960 alike), while TinyLlama's 22 layers and Qwen2.5's 24 are fine. Add _chunk_max_layers=24 and cap layers-per-chunk by it as well as by bytes, so deep models segment. Real SmolLM2-360M now decodes coherently (was rc=-1); segmented decode still matches the single-program greedy tokens. Tests: off-device chunk-cap checks + on-device deep-model decode and a segmented-vs-single greedy-match.
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
af.load_llm("HuggingFaceTB/SmolLM2-360M-Instruct")decoded withRuntimeError: execute failed rc=-1. Prefill worked; the failure was the resident-state decode step program (llm.py:601), not the lm_head.Root cause:
_layer_chunkssplits the decode program into chunks by baked weight bytes only (_chunk_bytes ~1.6GB). SmolLM2 is deep but narrow (32 layers, dim 960), so all 32 layers stay under the byte cap and compile into one program -- which exceeds an op-count ANE ceiling and fails execute.I isolated the ceiling with synthetic random models: it trips at ~31 layers regardless of dim (L30 OK / L32 fails at both dim 128 and dim 960), which is why TinyLlama (22 layers) and Qwen2.5-0.5B (24) work while SmolLM2 (32) did not.
Fix
Add
_chunk_max_layers = 24and cap layers-per-chunk by it as well as by bytes. Deep models now segment across decode programs (the machinery already exists; it just wasn't triggered by byte size). 24 sits safely below the measured ~31 ceiling with margin for heavier mixers.Verification
SmolLM2-360M-Instructnow decodes coherently: "a powerful tool that can help you analyze and visualize complex data..." (wasrc=-1).Tests (
tests/test_decode_chunking.py):_layer_chunkscaps at_chunk_max_layersand covers every layer; shallow models stay single-chunkrequires_ane): 32-layer model decodes without rc-error; segmented vs single-program greedy matchtest_llm.py+ new tests: 19 passed, no regressions.