Skip to content

Commit 87e0749

Browse files
xingguo01xingguo01
authored
Arm backend: add TinyStories-42M Ethos-U KV-cache demo (pytorch#22433)
Add a ctx64 KV-cache export flow for TinyStories-42M using linear w8a16 quantization and Wikitext calibration. Document FVP validation and deployment within the Alif E8 125 MiB model-storage budget. AI-assisted-by: Codex cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani Signed-off-by: Xingguo Li <xingguo.li@arm.com>
1 parent 291a6cc commit 87e0749

2 files changed

Lines changed: 245 additions & 0 deletions

File tree

‎examples/arm/tinystories42m_example_ethos_u/README.md‎

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,136 @@
1+
# TinyStories-42M KV-cache generation on Ethos-U85
2+
3+
This example exports the `stories42M.pt` checkpoint as a static-shape,
4+
token-at-a-time KV-cache model with linear `w8a16` PT2E quantization. It reuses
5+
the validated SmolLM2 Ethos-U semihosting runner and host generation helper.
6+
7+
## Set up
8+
9+
From the ExecuTorch repository root:
10+
11+
```bash
12+
source ../optimum-executorch/.venv_et/bin/activate
13+
source examples/arm/arm-scratch/setup_path.sh
14+
15+
export CC=/usr/bin/gcc-12
16+
export CXX=/usr/bin/g++-12
17+
export CUDAHOSTCXX="${CXX}"
18+
```
19+
20+
Place these files in one directory:
21+
22+
```text
23+
data/tinystories42m/stories42M.pt
24+
data/tinystories42m/params.json
25+
data/tinystories42m/tokenizer.model
26+
```
27+
28+
The checkpoint is available from
29+
`https://huggingface.co/karpathy/tinyllamas/resolve/main/stories42M.pt`.
30+
Set `TINYSTORIES42M_ARTIFACT_DIR` when the files live elsewhere.
31+
32+
Use these parameters for the 42M checkpoint:
33+
34+
```json
35+
{
36+
"dim": 512,
37+
"multiple_of": 32,
38+
"n_heads": 8,
39+
"n_kv_heads": 8,
40+
"n_layers": 8,
41+
"norm_eps": 1e-05,
42+
"vocab_size": 32000
43+
}
44+
```
45+
46+
## Export ctx64 plain KV cache
47+
48+
```bash
49+
bash examples/arm/tinystories42m_example_ethos_u/export_tinystories42m_ethosu.sh
50+
```
51+
52+
The default export uses:
53+
54+
```text
55+
model.use_kv_cache=True
56+
model.quantize_kv_cache=False
57+
model.enable_dynamic_shape=False
58+
export.max_seq_length=64
59+
export.max_context_length=64
60+
quantization.pt2e_quantize=ethosu_16a8w
61+
quantization.quantize_scope=linear
62+
quantization.calibration_tasks=[wikitext]
63+
quantization.calibration_limit=1
64+
quantization.calibration_seq_length=64
65+
debug.generate_full_logits=False
66+
```
67+
68+
Although the KV model consumes one token per invocation, `max_seq_length=64`
69+
describes the exported sequence capacity. PT2E uses one 64-token Wikitext
70+
sample for the quick validation calibration.
71+
72+
## KV-cache configuration
73+
74+
Use FP32 KV-cache buffers for Alif E8 deployment. These buffers retain the key
75+
and value tensors computed for earlier tokens and are updated for each new
76+
token. Board measurements found decoding with FP32 KV-cache buffers
77+
approximately 6x faster than with statically quantized int8 KV-cache buffers.
78+
At ctx64, storing the KV-cache buffers as int8 reduces their runtime size from
79+
approximately 2 MiB to 0.5 MiB, but does not reduce the PTE stored in OSPI. For
80+
implementation details, see the
81+
[SmolLM2 static KV-cache section](../smollm2_example_ethos_u/README.md#7-experimental-static-quantized-kv-cache).
82+
83+
## Alif E8 deployment
84+
85+
The exported PTE is approximately 103 MiB, so it fits within the current
86+
125 MiB model-storage budget for the Alif E8 Ethos-U85 configuration. Deploy
87+
the PTE with the Alif board-specific ExecuTorch firmware runner; the embedded
88+
Corstone-320 semihosting runner below is only for FVP validation and is not the
89+
Alif firmware image.
90+
91+
The board integration must provide the token and `input_pos` int32 inputs,
92+
retain the KV-cache state between decoding steps, and perform token selection
93+
and decoding in the application. Firmware, runtime arenas, KV-cache buffers,
94+
and tokenizer storage are additional to the PTE size and must be checked
95+
against the board's runtime memory configuration.
96+
97+
## Build the FVP runner
98+
99+
```bash
100+
ARTIFACT_DIR="${TINYSTORIES42M_ARTIFACT_DIR:-data/tinystories42m}"
101+
102+
bash examples/arm/smollm2_example_ethos_u/build_executor_runner_semihosting.sh \
103+
--pte="${ARTIFACT_DIR}/tinystories42m_ethosu_u85_256_ctx64_kv_w8a16.pte" \
104+
--output=tinystories42m_ethosu_ctx64/cmake-out \
105+
--target=ethos-u85-256 \
106+
--system_config=Ethos_U85_SYS_DRAM_Mid \
107+
--memory_mode=Sram_Only \
108+
--method_pool_size=0x01000000 \
109+
--scratch_pool_size=0x00800000 \
110+
--input_file_pool_size=0x00100000
111+
```
112+
113+
## Greedy smoke test
114+
115+
```bash
116+
ARTIFACT_DIR="${TINYSTORIES42M_ARTIFACT_DIR:-data/tinystories42m}"
117+
118+
python examples/arm/smollm2_example_ethos_u/generate_sampled.py \
119+
--fvp=examples/arm/arm-scratch/FVP-corstone320/models/Linux64_GCC-9.3/FVP_Corstone_SSE-320 \
120+
--runner=tinystories42m_ethosu_ctx64/cmake-out/arm_executor_runner \
121+
--embedded-pte \
122+
--tokenizer="${ARTIFACT_DIR}/tokenizer.model" \
123+
--prompt="Once upon a time in a small village," \
124+
--window=64 \
125+
--max-context-length=64 \
126+
--use-kv-cache \
127+
--max-new-tokens=2 \
128+
--temperature=0 \
129+
--top-p=0.9 \
130+
--repetition-penalty=1.1 \
131+
--timeout=1200
132+
```
133+
134+
This smoke test validates export, runner integration, stateful token/position
135+
inputs, and decoding. Add prompt-set generation and perplexity only after this
136+
path produces a coherent continuation.

‎examples/arm/tinystories42m_example_ethos_u/export_tinystories42m_ethosu.sh‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
#!/usr/bin/env bash
2+
# Copyright 2026 Arm Limited and/or its affiliates.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
7+
set -euo pipefail
8+
9+
script_dir=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)
10+
repo_root=$(cd "${script_dir}/../../.." && pwd)
11+
12+
artifact_dir="${TINYSTORIES42M_ARTIFACT_DIR:-${repo_root}/data/tinystories42m}"
13+
checkpoint="${artifact_dir}/stories42M.pt"
14+
params="${artifact_dir}/params.json"
15+
tokenizer="${artifact_dir}/tokenizer.model"
16+
output_dir="${artifact_dir}"
17+
output_name="tinystories42m_ethosu_u85_256_ctx64_kv_w8a16.pte"
18+
calibration_limit=1
19+
calibration_seq_length=64
20+
max_seq_length=64
21+
max_context_length=64
22+
target="ethos-u85-256"
23+
system_config="Ethos_U85_SYS_DRAM_Mid"
24+
memory_mode="Sram_Only"
25+
26+
usage() {
27+
cat <<EOF
28+
Usage: $(basename "$0") [options]
29+
30+
Options:
31+
--checkpoint=PATH TinyStories-42M checkpoint.
32+
--params=PATH Model params JSON.
33+
--tokenizer=PATH SentencePiece tokenizer.model.
34+
--output_dir=DIR Export directory. Default: ${output_dir}
35+
--output_name=NAME PTE filename. Default: ${output_name}
36+
--calibration_limit=N Wikitext calibration sample count. Default: ${calibration_limit}
37+
--calibration_seq_length=N PT2E calibration length. Default: ${calibration_seq_length}
38+
--max_seq_length=N Export sequence capacity. Default: ${max_seq_length}
39+
--max_context_length=N KV-cache capacity. Default: ${max_context_length}
40+
--target=NAME Ethos-U target. Default: ${target}
41+
--system_config=NAME Vela system config. Default: ${system_config}
42+
--memory_mode=NAME Vela memory mode. Default: ${memory_mode}
43+
EOF
44+
}
45+
46+
for arg in "$@"; do
47+
case "${arg}" in
48+
-h|--help) usage; exit 0 ;;
49+
--checkpoint=*) checkpoint="${arg#*=}" ;;
50+
--params=*) params="${arg#*=}" ;;
51+
--tokenizer=*) tokenizer="${arg#*=}" ;;
52+
--output_dir=*) output_dir="${arg#*=}" ;;
53+
--output_name=*) output_name="${arg#*=}" ;;
54+
--calibration_limit=*) calibration_limit="${arg#*=}" ;;
55+
--calibration_seq_length=*) calibration_seq_length="${arg#*=}" ;;
56+
--max_seq_length=*) max_seq_length="${arg#*=}" ;;
57+
--max_context_length=*) max_context_length="${arg#*=}" ;;
58+
--target=*) target="${arg#*=}" ;;
59+
--system_config=*) system_config="${arg#*=}" ;;
60+
--memory_mode=*) memory_mode="${arg#*=}" ;;
61+
*)
62+
echo "Unknown option: ${arg}" >&2
63+
usage
64+
exit 1
65+
;;
66+
esac
67+
done
68+
69+
for path in "${checkpoint}" "${params}" "${tokenizer}"; do
70+
if [[ ! -f "${path}" ]]; then
71+
echo "Missing required model file: ${path}" >&2
72+
exit 1
73+
fi
74+
done
75+
76+
if [[ "${output_name}" != *.pte ]]; then
77+
output_name="${output_name}.pte"
78+
fi
79+
if [[ "${output_name}" = /* ]]; then
80+
output_path="${output_name}"
81+
else
82+
output_path="${output_dir}/${output_name}"
83+
fi
84+
mkdir -p "${output_dir}" "$(dirname -- "${output_path}")"
85+
cd "${repo_root}"
86+
87+
python -m extension.llm.export.export_llm \
88+
base.model_class=stories110m \
89+
base.checkpoint="${checkpoint}" \
90+
base.params="${params}" \
91+
base.tokenizer_path="${tokenizer}" \
92+
export.output_dir="${output_dir}" \
93+
export.output_name="${output_path}" \
94+
export.max_seq_length="${max_seq_length}" \
95+
export.max_context_length="${max_context_length}" \
96+
quantization.pt2e_quantize=ethosu_16a8w \
97+
quantization.quantize_scope=linear \
98+
quantization.calibration_tasks="[wikitext]" \
99+
quantization.calibration_limit="${calibration_limit}" \
100+
quantization.calibration_seq_length="${calibration_seq_length}" \
101+
backend.ethosu.enabled=True \
102+
backend.ethosu.target="${target}" \
103+
backend.ethosu.system_config="${system_config}" \
104+
backend.ethosu.memory_mode="${memory_mode}" \
105+
model.use_kv_cache=True \
106+
model.quantize_kv_cache=False \
107+
model.enable_dynamic_shape=False \
108+
debug.verbose=True \
109+
debug.generate_full_logits=False

0 commit comments

Comments
 (0)