Commit Graph
6 Commits
Author SHA1 Message Date
tzhouamandClaude Opus 5 67bee59fda [Edge][Mobile] Ring-buffer KV and score-space attention for the decode step
Reading the device profile of the first export showed a fifth of every
sliding layer went nowhere useful: Output 9.9%, its quantise conversion
3.1%, and the window Slice 7.4%. The layer took a 512-entry window in,
appended the new key, sliced the oldest off and handed the whole recomputed
window back out -- 21 of 28 layers paying to move a cache-sized tensor
across the graph boundary once per token. Full layers paid a similar 9.5%
concatenating the new key onto the cache.

Both are the same mistake, and both go away by joining in score space
instead of key space: score the cache and the new token separately, then
concatenate two vectors of length n and 1 rather than an n x 256 key matrix.
Nothing cache-sized is copied, and each layer returns just the new K/V entry
for the runtime to drop into its ring slot. Softmax over keys does not care
what order the slots are in and each cached key already carries its own
rotary phase, so a ring buffer needs no reordering -- only a mask for slots
not yet written. The sliding cache correspondingly holds window-1 entries,
since a 512-token window counts the current token (transformers' own sliding
cache holds 511 for the same reason).

Measured on a real Galaxy S25, fp16, ms per layer:

  sliding  1.689 -> 1.483   (-12.2%)
  full     1.559 -> 1.567   (unchanged: the concat saved is spent on the
                             extra score matmul)

Two shapes had to be written as matmuls rather than the obvious elementwise
forms -- the new token's score (broadcast multiply + ReduceSum) and its
value contribution (a two-axis broadcast multiply). Both compile at fp16 and
both make QNN's context-binary converter exit 14 under w4a16; as matmuls
they compile and are the cheaper op on HTP anyway.

Also adds a selectable GELU approximation, and the measurement that closes
the question: erf GELU is 11-13% of every layer, but tanh is *slower* on
this NPU (1.535 vs 1.483 ms) despite holding 65.2 dB. The exact form stays,
which is also the better-fidelity choice.

Parity against the reference implementation is 136.8 dB with identical
top-1/top-5, and a test pins ring and roll to the same output on a shared
history so the device numbers cannot silently measure a different model.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01L1ax5wYLA6XbqEL1i8L9gZ
2026-09-11 17:02:48 +08:00
tzhouamandClaude Opus 5 658341531a [Edge][Mobile] Spark-X2.5 decode-step export, and a compile-safe head gate
Exports one Spark decode step with the KV cache as explicit tensors, sized
per layer type: the 21 sliding layers hold a fixed 512 entries at any
context, the 7 full-attention layers grow. That asymmetry is the whole
reason the architecture is interesting on a phone, so the two types get
separately sized caches, masks and rotary angles.

Verified against the HF reference at 1024 context: 138.5 dB logits SNR,
identical top-1 and top-5.

Profiled on a real Galaxy S25 (Snapdragon 8 Elite) through AI Hub, QNN
context binary on the NPU, ms per layer:

  sliding (512 KV, any ctx)  fp16 1.689  w8a16 0.912  w4a16 0.917
  full @ 1k                  fp16 1.559  w8a16 1.129  w4a16 1.127
  full @ 4k                  fp16 2.964  w8a16 1.681  w4a16 1.662

Composed into a step (21 sliding + 7 full), w4a16 gives 27.2 ms at 1k and
30.9 ms at 4k, against 31.6 and 46.5 ms for the same model with every layer
full-attention -- 1.16x and 1.51x. Worth noting the hybrid is 0.94x at 1k in
fp16: sliding layers rotate all 256 head dims where full layers rotate 64,
and below roughly 2k that RoPE cost outweighs the shorter cache.

Also moves the head-wise attention gate off vLLM's CPU linear dispatch. The
gate is [2048, 8], and any linear that narrow is routed to the sgl-kernel
packed GEMM, whose fake-tensor rule assumes the packed weight keeps N in dim
0; VNNI packing does not at that size, so it reports K and torch.compile
fails tracing the gate multiply. A plain parameter with F.linear gains
nothing less -- a 2048x8 matmul has no use for packed AMX tiles -- and it
keeps the model compilable on CPU and GPU alike.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01L1ax5wYLA6XbqEL1i8L9gZ
2026-09-11 14:49:10 +08:00
tzhouamandClaude Opus 5 73cdc7395a [Edge][Mobile] QNN/TFLite decode-step export for phone accelerators
The Code2Wav decoder already had an on-device export (WP4); the talker and its
code predictor did not, so the codec-token path stopped at the vocoder. This
adds the missing half: a one-token decode step with an explicit per-layer KV
cache, exportable to ONNX and from there to a QNN context binary, TFLite or
ONNX Runtime.

Two layout choices carry the performance, both measured on real Snapdragon
devices through Qualcomm AI Hub (Galaxy S24 = 8 Gen 3, S25 = 8 Elite):

* per-layer cache tensors instead of one stacked [num_layers, ...] tensor that
  every layer slices, and
* grouped-query attention by reshaping the query rather than expanding keys and
  values.

The stacked/expanded form spent ~70% of its NPU cycles in cache select,
reshape and fp32->fp16 conversion; the layout here is bandwidth-bound and runs
the 0.6B talker step in 15.8 ms instead of 36.9 ms on an 8 Elite (fp16,
256-token cache). The code predictor step is 2.25 ms, so talker plus its 15
sub-steps cost about 50 ms of an 80 ms frame.

sanitize_onnx() fixes the three things the dynamo exporter emits that
qairt-converter and AI Hub reject: IsNaN/Where guards, Reshape allowzero=1, and
rank-0 axes initializers on Unsqueeze/Squeeze.

The module docstring records the measured placement table: this graph belongs
on the NPU, the vocoder on the GPU (21.7 ms vs 182 ms per frame), which is 3.3x
better than sending everything to the NPU and makes the pipeline real time on
the phone in fp16. Integer quantization is not the lever it looks like: every
automatic int8 build measured was numerically broken.

Full write-up, job ids and profiles: analysis/mobile_htp_simulation.md section 5.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Dzr1Jusv4s6g765gqS9L92
2026-09-10 16:17:58 +08:00
tzhouamandClaude Fable 5.1 580a201631 [Edge][WP4] Windowed Code2Wav decoder export (torch.export / ExecuTorch) + client backend; results summarizer
vllm_omni/edge/decoder_export.py exports a stateless windowed decoder (context
72 frames + chunk, one program per ramp chunk size) via torch.export (.pt2,
verified here on CPU for chunks 2 and 25 with the real 0.6B decoder: exact
match vs eager, 460 MB fp32 per program) or ExecuTorch XNNPACK (.pte; written
against the executorch 1.x API, not executed: package not installed, install
is permission gated; QNN partitioner export-only/unverified). The WP5 client
gains --decoder-backend executorch --pte-manifest for phone/NPU targets.
benchmarks/tts/edge_results/summarize.py renders all committed JSON as tables.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
2026-09-09 15:02:19 +08:00
tzhouamandClaude Fable 5.1 37c0ea9a6b [Edge][WP8] Per-step overhead attribution, edge scheduling helpers, CPU thread sweep
- vllm_omni/edge/step_stats.py: process-local counters enabled by
  VLLM_OMNI_STEP_STATS_DIR (scheduler schedule/update_from_output, engine-core
  step, chunk hop enqueue->put / put->get / get, orchestrator dispatch), dumped
  per process and merged into an attribution table (share of the 80 ms frame).
  Hooks: stage_engine_core_proc (instrument after construction, dump at exit),
  chunk_transfer_adapter (put/get timing, put->get via the shm lock-file mtime),
  orchestrator dispatch loop. Bench: stream_latency_bench.py --step-stats.
- vllm_omni/edge/scheduling.py: playback slack, prefill-throttle decision and
  the no-preemption KV admission rule (exact bytes/token from the talker KV
  geometry); adapt.derive_overrides now sizes max_num_seqs with it and
  auto_profile passes the geometry read from the model config.
- Engine re-reads parallel_stage_init after apply_deploy_profile so the edge
  profile's orchestrator default takes effect (was read before the profile).
- benchmarks/tts/cpu_thread_sweep.py: CPU-only 8/16/32-thread sweep (+inductor)
  with per-step attribution.
- Tests: tests/core/test_edge_scheduling.py, tests/edge/test_step_stats.py,
  tests/benchmarks/test_cpu_thread_sweep.py, profile default test.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
2026-09-09 14:55:55 +08:00
tzhouamandClaude Fable 5.1 55192de984 [Edge][WP4] Add hardware adaptation: probe, derived overrides, auto profile, calibration
* vllm_omni/edge/hardware_probe.py: HardwareProfile from /proc, /sys and torch
  (ISA flags incl. AVX-512/AMX and NEON/dotprod/i8mm/SVE, physical cores,
  big.LITTLE clusters, NUMA, RAM, CUDA discrete vs Tegra unified memory, NPU
  via override), hardware classes jetson / cuda_discrete / arm64_cpu / x86_cpu /
  npu_phone; VLLM_OMNI_HW_PROFILE=<json> overrides the probe (emulation, tests).
* vllm_omni/edge/adapt.py: derive_overrides -> dtype, absolute KV budgets from
  available memory minus weights, max_num_seqs, per-stage disjoint CPU thread
  ranges on the big cluster, eager/graph switches for small or unified memory,
  adaptive-chunk parameters, timeouts from a load cost model.
* vllm_omni/edge/calibrate.py: chunk-schedule calibration from measured chunk
  arrivals with a per-hardware-fingerprint cache under VLLM_CACHE_ROOT.
* vllm_omni/edge/auto_profile.py + deploy_profile="auto": edge/<model>.yaml ->
  edge/hardware/<class>.yaml -> derived overrides -> cached calibration,
  materialized as a temp deploy YAML (platform overlays are superseded).
* benchmarks/tts/hw_emulation.py: emulated classes (ATEN_CPU_CAPABILITY,
  taskset, profile overrides) driving stream_latency_bench with auto profiles.
* CPU-only tests for probe snapshots, derivation rules, calibration and auto.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
2026-09-09 11:01:40 +08:00