Skip to content

llm: cap decode-program layers so deep small-dim models don't hit the ANE op ceiling - #254

Merged
sbryngelson merged 1 commit into
mainfrom
fix/deep-model-decode-chunking
Aug 31, 2026
Merged

llm: cap decode-program layers so deep small-dim models don't hit the ANE op ceiling#254
sbryngelson merged 1 commit into
mainfrom
fix/deep-model-decode-chunking

Conversation

@sbryngelson

Copy link
Copy Markdown
Owner

Problem

af.load_llm("HuggingFaceTB/SmolLM2-360M-Instruct") decoded with RuntimeError: 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_chunks splits 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 = 24 and 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

  • Real SmolLM2-360M-Instruct now decodes coherently: "a powerful tool that can help you analyze and visualize complex data..." (was rc=-1).
  • Synthetic 32-layer model generates.
  • Segmented decode produces the same greedy tokens as a forced single-program run.

Tests (tests/test_decode_chunking.py):

  • off-device: _layer_chunks caps at _chunk_max_layers and covers every layer; shallow models stay single-chunk
  • on-device (requires_ane): 32-layer model decodes without rc-error; segmented vs single-program greedy match

test_llm.py + new tests: 19 passed, no regressions.

… 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.
@sbryngelson
sbryngelson merged commit 2346c02 into main Aug 31, 2026
19 checks passed
@sbryngelson
sbryngelson deleted the fix/deep-model-decode-chunking branch August 31, 2026 00:34
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