Laya on AXERA NPU

Ready-to-run deployment package for the Laya typed-decision model family on AX650 / NPU3.

• Runtime: packaged axllm run
• Target: AX650 / AX8850, aarch64, NPU3
• Graph shape: batch 1, sequence length 256, up to 4 options
• Decision primitives: choice, score, and noul
• Included checkpoints: English, multilingual, and typed decisions
• Included assets: AXModel files, Hugging Face tokenizers, runtime configs, sample requests, sample outputs, bin/axllm, and isolated AX650/PyTorch Python runners

Laya is a bidirectional decision model. It evaluates user-defined questions over text or structured JSON state without generating text.

Supported Platform

• AX650 / AX8850
• NPU3

Checkpoint Selection

Directory	Backbone	Recommended use
english/	ModernBERT-large	English routing, guardrails, moderation, and support triage
multilingual/	mmBERT-base	Chinese and other multilingual inputs
typed-decisions/	ModernBERT-large	Invoice, security, customer-service, and agent-trace workflows

The packaged runtime does not automatically route between checkpoints. Select the directory that matches the input language or workflow.

Model Inputs and Outputs

User-Level Input

Users submit a JSON object with two fields:

• state: the text or structured record to evaluate, such as a support ticket, security incident, email, or agent trace.
• questions: named decision definitions. Each question contains a decision type, an instruction, and the candidate criteria when applicable.

The runtime supports three decision types:

Type	Meaning	Returned value
choice	Select one label from 2 to 4 candidates	Selected label and probability distribution
score	Evaluate an ordered scale containing 2 to 4 levels	Expected score, probability distribution, and level legend
noul	Estimate whether a statement is true	Probability of true

For every question, axllm serializes the request into this token sequence:

[CLS] <question type and instruction> [SEP]
[MASK] <option 0> [MASK] <option 1> ... [SEP]
<state> [SEP]

Each [MASK] position represents one candidate answer. A request containing four questions runs four forwards while keeping the same AXModel resident.

AXModel Tensor Interface

All three packaged AXModels use the same fixed tensor interface.

Tensor	Direction	Dtype	Shape	Meaning
input_ids	Input	S32	[1, 256]	Token IDs for the instruction, candidates, and state
attention_mask	Input	S32	[1, 256]	1 for valid tokens and 0 for padding
marker_pos	Input	S32	[1, 4]	Token position of each candidate's [MASK] marker
marker_mask	Input	S32	[1, 4]	Marks which of the four candidate slots are active
qtype	Input	S32	[1]	0 = choice, 1 = score, 2 = noul
logits	Output	FP32	[1, 4]	Candidate decision logits; only active candidate slots are used
act_logits	Output	FP32	[1, 2]	Auxiliary act-versus-escalate logits

The CPU postprocessor applies the packaged temperature values and returns:

• choice: the highest-probability label plus all active option probabilities.
• score: the expected ordinal value, computed as the probability-weighted level index.
• noul: the probability that the statement is true.
• confidence: distribution concentration for choice/score, or the larger side of the binary probability for noul.
• action.act_probability: the auxiliary head's tendency to act automatically. Validate this field on application data before using it as an automation gate.
• npu_latency_ms: NPU forward time for that question.

Performance

Measured on AX650 / NPU3 with /opt/bin/ax_run_model, two warmup runs and ten measured runs. Each AXModel invocation evaluates one question. Latency excludes model loading and CPU tokenization.

Checkpoint	Fixed input	Average NPU latency
english	B1, S256, K4	69.991 ms
multilingual	B1, S256, K4	27.722 ms
typed-decisions	B1, S256, K4	69.990 ms

Representative four-question axllm requests measured 280.4 ms for English, 110.9 ms for multilingual, and 280.3 ms for typed decisions. These totals are the sum of NPU forward time reported by the runtime.

Why the Multilingual AXModel Is Larger but Faster

AXModel file size and inference latency measure different costs.

Checkpoint	Vocabulary	Token-embedding parameters	Hidden size	Encoder layers	FFN intermediate size
english	50,368	51.6M	1,024	28	2,624
multilingual	256,000	196.6M	768	22	1,152
typed-decisions	50,368	51.6M	1,024	28	2,624

The multilingual checkpoint needs a much larger vocabulary to cover many languages. Its 256,000-by-768 token-embedding table is stored as 16-bit data in the compiled graph and accounts for a large part of the AXModel file and CMM footprint. During one inference, the embedding operator gathers only the rows referenced by the 256 input token IDs; it does not compute over all 256,000 vocabulary entries.

Most NPU time is spent in the encoder layers after the embedding lookup. The multilingual backbone has fewer layers, a smaller hidden dimension, and a much smaller feed-forward dimension. A rough projection-plus-feed-forward compute proxy is about 3.1 times larger for the English backbone, while the measured latency ratio is 2.52 times (69.991 / 27.722). Kernel scheduling and fixed operator overhead account for the difference between the rough compute ratio and measured time.

The larger multilingual file therefore reflects stored vocabulary capacity, while its lower latency reflects a lighter encoder computation.

Runtime Footprint

Checkpoint	AXModel file	Runtime CMM
english	481.13 MiB	481.63 MiB
multilingual	508.61 MiB	508.98 MiB
typed-decisions	481.13 MiB	481.63 MiB

Only one checkpoint needs to be loaded for a single axllm process. The complete package is approximately 1.48 GiB before repository metadata.

Package Layout

.
├── README.md
├── LICENSE
├── NOTICE
├── SHA256SUMS
├── bin/
│   └── axllm
├── demo/
│   ├── app.py
│   ├── backend.py
│   ├── scenarios.py
│   ├── requirements.txt
│   └── README.md
├── python/
│   ├── README.md
│   ├── ax650/
│   │   ├── infer.py
│   │   └── requirements.txt
│   └── pytorch/
│       ├── infer.py
│       └── requirements.txt
├── english/
│   ├── config.json
│   ├── model.axmodel
│   ├── sample_request.json
│   ├── sample_output.json
│   └── tokenizer/
├── multilingual/
│   ├── config.json
│   ├── model.axmodel
│   ├── sample_request.json
│   ├── sample_output.json
│   └── tokenizer/
├── typed-decisions/
│   ├── config.json
│   ├── model.axmodel
│   ├── sample_request.json
│   ├── sample_output.json
│   └── tokenizer/
└── LICENSES/
    └── AXLLM-BSD-3-Clause.txt

Every checkpoint directory is self-contained. Keep its config, AXModel, and tokenizer together.

Download the Package

mkdir -p AXERA-TECH/Laya
cd AXERA-TECH/Laya
hf download AXERA-TECH/Laya --local-dir .

Verify the downloaded runtime artifacts:

sha256sum -c SHA256SUMS

Run with Python

The Python files are isolated by backend under python/ax650/ and python/pytorch/. Both accept the same state and questions JSON schema as axllm run.

AX650 Python Inference

python/ax650/infer.py loads the packaged AXModel directly through PyAXEngine. On the board, create an environment and install the tokenizer dependencies plus an axengine wheel from the PyAXEngine releases page:

python3 -m venv .venv-ax650
source .venv-ax650/bin/activate
python -m pip install -r python/ax650/requirements.txt
python -m pip install /path/to/axengine-*.whl

Run the Chinese sample:

python python/ax650/infer.py multilingual \
  --input multilingual/sample_request.json

Use english, multilingual, or typed-decisions as the model directory. Omit --input to keep the AXModel and PyAXEngine session resident while reading one complete JSON request per line:

{
  tr -d '\n' < multilingual/sample_request.json; echo
  tr -d '\n' < multilingual/sample_request.json; echo
  echo /exit
} | python python/ax650/infer.py multilingual

The Python runner reports python_latency_ms for each decision and total_python_latency_ms for the request. These are wall-clock measurements around InferenceSession.run and include Python/CFFI call overhead.

Validated on an AX650 board with PyAXEngine and the packaged sample requests:

Checkpoint	First four-decision request	Resident second request	Steady single-decision range	Result check
english	293.65 ms	—	71.04–71.62 ms	Labels and probabilities match axllm
multilingual	123.05 ms	114.88 ms	28.66–28.82 ms	Labels and probabilities match axllm
typed-decisions	292.66 ms	—	71.01–71.07 ms	Labels and probabilities match axllm

The first decision after model creation includes one-time initialization overhead. Resident mode removes that effect from subsequent requests.

Original PyTorch Reference

python/pytorch/infer.py uses the original FP checkpoint on CPU, CUDA, or MPS. The selected upstream checkpoint is downloaded to the Hugging Face cache on first use; this AX650 repository does not duplicate the original FP weights.

Create an environment and install the pinned Laya package:

python3 -m venv .venv-pytorch
source .venv-pytorch/bin/activate
python -m pip install -r python/pytorch/requirements.txt

Run the multilingual sample on an automatically selected CUDA, MPS, or CPU device:

python python/pytorch/infer.py \
  --variant multilingual \
  --input multilingual/sample_request.json \
  --device auto

Use english, multilingual, or typed-decisions for --variant. To keep one PyTorch checkpoint resident, omit --input and submit one complete JSON request per line:

{
  tr -d '\n' < multilingual/sample_request.json; echo
  tr -d '\n' < multilingual/sample_request.json; echo
  echo /exit
} | python python/pytorch/infer.py --variant multilingual --device auto

The reference runner sets the original checkpoint to sequence length 256 and option budget 128 to match the packaged AXModel input layout. Its answer schema matches the AXERA runtime, with additional usage and python_runtime metadata.

Python and AXERA inference execute the same typed-decision model logic, tokenizer construction, temperature scaling, and CPU postprocessing. They use different compute paths:

Runner	Model data	Compute device	Multi-question execution	Timing field
python/ax650/infer.py	Compiled and quantized AXModel	AX650 / NPU3 through PyAXEngine	One NPU forward per question	python_latency_ms wall time
bin/axllm run	Compiled and quantized AXModel	AX650 / NPU3 through the native C++ runtime	One NPU forward per question	npu_latency_ms kernel time
python/pytorch/infer.py	Original PyTorch FP checkpoint	CPU, CUDA, or MPS	All questions are batched into one PyTorch forward	python_runtime.latency_ms wall time

PyAXEngine and axllm use the same AXModel, tensor construction, temperature scaling, and answer postprocessing, so their labels and probabilities should match within floating-point display precision. Their main differences are the language/runtime integration and timing scope. The original PyTorch runner uses FP weights, batches questions differently, and is not expected to produce bit-identical probabilities after AXModel quantization.

Gradio Demo on the Board

The independent demo/ directory adapts the layout and scenario patterns from the upstream Laya Gradio Space and connects them to the packaged AXModels through python/ax650/infer.py. It runs directly on the AX650 board and includes five tabs:

• Chinese customer-support triage
• Typed security-incident decisions
• LLM prompt guardrails
• RAG passage filtering
• A free-form state + questions playground

Demo Preview

This example runs the multilingual checkpoint on a Chinese enterprise-customer complaint about repeated invoice charges. One request produces four typed decisions: route the case to billing, assign urgency 1.84 / 2, detect a 99.6% refund intent, and estimate a 95.5% churn signal. The demo then turns these signals into a readable workflow: raise the case priority, start a refund review, and notify the retention team.

The screenshot reports 123.37 ms for the four accumulated NPU calls and 140.37 ms for the complete request. These values cover the whole four-question workflow; each decision card also shows its individual inference time and confidence.

After installing the AX650 Python dependencies and PyAXEngine described above, install the pinned Gradio version:

source .venv-ax650/bin/activate
python -m pip install -r demo/requirements.txt

Start the board service and preload all three checkpoints so tab switching does not pay model-load time:

LAYA_PRELOAD=english,multilingual,typed-decisions \
GRADIO_SERVER_NAME=0.0.0.0 \
GRADIO_SERVER_PORT=7860 \
python demo/app.py

Open http://<board-ip>:7860. If the board is reachable only through an SSH server, forward the port from that server:

ssh -L 7860:<board-ip>:7860 <ssh-server>

Then open http://127.0.0.1:7860. NPU access is serialized across browser requests. Set LAYA_PRELOAD=multilingual to reduce startup time and CMM usage; the other checkpoints will load on their first request.

Run on the Board

The packaged binary uses the AX650 on-chip backend and the system AXERA runtime libraries.

chmod +x ./bin/axllm

Detailed Example: English Customer-Support Triage

This example evaluates a customer-support ticket. The customer reports a duplicate invoice charge, requests a refund today, and threatens to cancel the service.

Input

The state field contains the business record. The four entries under questions ask the model to select a handling department, assign an urgency score, detect refund intent, and detect cancellation risk.

{
  "state": {
    "body": "Invoice 4411 was charged twice. Please refund the duplicate today or we will cancel our plan."
  },
  "questions": {
    "department": {
      "type": "choice",
      "instructions": "Which team should handle this request?",
      "criteria": {
        "billing": "payments, invoices, charges, refunds",
        "technical": "bugs, outages, API or account access failures",
        "sales": "pricing, demos, contracts or purchases",
        "other": "general requests or resolved issues"
      }
    },
    "urgency": {
      "type": "score",
      "instructions": "How urgent is this request?",
      "criteria": [
        "routine, no deadline",
        "needs attention soon",
        "service blocked or deadline today"
      ]
    },
    "refund": {
      "type": "noul",
      "instructions": "Does the customer explicitly request a refund?"
    },
    "churn": {
      "type": "noul",
      "instructions": "Does the customer threaten to cancel or leave?"
    }
  }
}

Run the packaged English checkpoint:

./bin/axllm run ./english --input ./english/sample_request.json

AX650 Output

The following is the complete output recorded on the AX650 board:

{
  "model": "laya-english",
  "answers": {
    "department": {
      "type": "choice",
      "choice": "billing",
      "probabilities": {
        "billing": 0.9689645773705189,
        "technical": 0.008899344656499457,
        "sales": 0.010587676002932096,
        "other": 0.011548401970049494
      },
      "confidence": 0.8757531196892554,
      "action": {
        "act_probability": 1.0
      },
      "npu_latency_ms": 70.223844
    },
    "urgency": {
      "type": "score",
      "probabilities": {
        "0": 0.03266581800541651,
        "1": 0.5424619734608962,
        "2": 0.4248722085336873
      },
      "legend": {
        "0": "routine, no deadline",
        "1": "needs attention soon",
        "2": "service blocked or deadline today"
      },
      "score": 1.3922063905282709,
      "confidence": 0.2652274294222525,
      "action": {
        "act_probability": 1.0
      },
      "npu_latency_ms": 70.111539
    },
    "refund": {
      "type": "noul",
      "noul": 0.8548107451739729,
      "confidence": 0.8548107451739729,
      "action": {
        "act_probability": 1.0
      },
      "npu_latency_ms": 70.021244
    },
    "churn": {
      "type": "noul",
      "noul": 0.844980400947017,
      "confidence": 0.844980400947017,
      "action": {
        "act_probability": 1.0
      },
      "npu_latency_ms": 70.019742
    }
  },
  "total_npu_latency_ms": 280.37636899999995
}

How to Read the Output

• department is a choice decision. The selected label is billing with probability 0.9690. The remaining probabilities show how strongly the model rejected the technical, sales, and other queues.
• department.confidence is 0.8758. For choice and score questions, confidence measures how concentrated the probability distribution is.
• urgency is a score decision with levels 0, 1, and 2. The result 1.3922 is the probability-weighted expected level, between "needs attention soon" and "service blocked or deadline today."
• refund.noul is 0.8548, meaning an 85.48% estimated probability that the customer explicitly requested a refund.
• churn.noul is 0.8450, meaning an 84.50% estimated probability that the customer threatened to cancel or leave.
• action.act_probability comes from the auxiliary act-versus-escalate head. Validate this value on application data before using it to authorize an automatic action.
• npu_latency_ms is the NPU time for one question. total_npu_latency_ms is the sum for all four questions and excludes model loading and CPU tokenization.

A support system can use this result to route the ticket to billing, raise its priority, start a refund-review workflow, and notify a retention team. Production automation thresholds should be calibrated with representative application data.

Other Checkpoints

The other packaged checkpoints use the same request schema and output fields.

Checkpoint	Sample scenario	Run command	Representative result
multilingual	The same duplicate-charge request and questions written in Chinese	./bin/axllm run ./multilingual --input ./multilingual/sample_request.json	billing; urgency 1.9359/2; refund 0.9925; churn 0.9400
typed-decisions	An active production API key exposed in a public issue	./bin/axllm run ./typed-decisions --input ./typed-decisions/sample_request.json	contain; severity 1.7371/2; exposure 0.4807; immediate containment 0.6667

The complete validated input and output for each checkpoint are stored beside the model as sample_request.json and sample_output.json.

Resident JSON Lines Mode

Omit --input to keep one checkpoint resident. Submit one complete JSON request per line and enter /exit to stop.

{
  jq -c . ./multilingual/sample_request.json
  jq -c . ./multilingual/sample_request.json
  echo /exit
} | ./bin/axllm run ./multilingual

Request Format

A request contains one state value and a named questions object.

{
  "state": {
    "body": "Invoice 4411 was charged twice. Please refund the duplicate today."
  },
  "questions": {
    "department": {
      "type": "choice",
      "instructions": "Which team should handle this request?",
      "criteria": {
        "billing": "payments, invoices, charges, refunds",
        "technical": "bugs, outages, API or account failures",
        "sales": "pricing, demos, contracts or purchases",
        "other": "general requests"
      }
    },
    "urgency": {
      "type": "score",
      "instructions": "How urgent is this request?",
      "criteria": [
        "routine",
        "needs attention soon",
        "service blocked or deadline today"
      ]
    },
    "refund": {
      "type": "noul",
      "instructions": "Does the customer explicitly request a refund?"
    }
  }
}

Decision Primitives

Primitive	Input criteria	Output
choice	Object or list containing 2 to 4 options	Selected label and probability for each option
score	Ordered list containing 2 to 4 levels	Expected score, level probabilities, and legend
noul	Optional descriptions for false and true	Probability that the statement is true

Validated Examples

The package includes the exact requests and complete JSON outputs used for board validation.

Checkpoint	Scenario	Representative result
english	Duplicate charge, refund request, cancellation threat	billing 0.9690; refund 0.8548; churn 0.8450
multilingual	Chinese duplicate charge and cancellation intent	billing 1.0000; refund 0.9925; churn 0.9400
typed-decisions	Active production API key exposed publicly	contain 0.4064; severity 1.7371; immediate containment 0.6667

Packaged Graph Constraints

The AX650 graphs in this release have fixed input shapes:

• batch size: 1
• sequence length: 256 tokens
• maximum options: 4
• input dtype: signed 32-bit integer
• outputs: four decision logits and two action logits

The runtime performs one NPU forward per question. The upstream checkpoints support longer contexts and batched questions, but those layouts are outside this release. axllm serve is not exposed for these models; use axllm run.

Probabilities can shift after quantization or when the deployment domain differs from the checkpoint calibration data. Validate thresholds on representative application data before automating high-impact actions.

Conversion References

If you need the original model files or want to rebuild the deployment artifacts, start with:

• Original Hugging Face models: Laya, Laya Multilingual, and Laya Typed Decisions
• Original source project: NandhaKishorM/laya
• AXERA runtime project: AXERA-TECH/ax-llm

Licenses

The Laya source and model family are distributed under Apache License 2.0. The packaged axllm binary is based on BSD-3-Clause licensed code; its license text is included in LICENSES/AXLLM-BSD-3-Clause.txt.

Discussion

• GitHub Issues
• QQ group: 139953715