Skip to content

A.X-K2 on RBLN NPU

Install the RBLN Driver on the host. An RBLN Portal account is required to install the RBLN Compiler (rebel-compiler). See the installation matrix for compatible package versions.

1
2
3
$ pip install \
  --extra-index-url https://pypi.rbln.ai/simple \
  rebel-compiler==0.12.0

The KV offload (LMCache) and prefill/decode-disaggregated commands need the kv_connectors extra, which installs lmcache-rbln and nixl-rbln. Without it those servers fail at startup.

1
2
3
4
5
$ pip install \
  --extra-index-url https://pypi.rbln.ai/simple \
  --extra-index-url https://wheels.vllm.ai/0.30.0/cpu \
  --extra-index-url https://download.pytorch.org/whl/cpu \
  "vllm-rbln[kv_connectors]==0.12.0"
Guide — overview and configuration

Overview

skt/A.X-K2 is a 688B-parameter Mixture-of-Experts, decoder-only model (256 routed experts + 1 shared, 8 activated per token) using Multi-head Latent Attention (MLA) with Sparse Gated Attention (SGA). On RBLN-CR13 it is served through the RBLN vLLM plugin, aggregated or with prefill and decode disaggregated. The served checkpoint is the NVFP4-quantized skt/A.X-K2-NVFP4, at the full 262,144-token context.

Configuration

Verified on RBLN-CR13. Select a configuration in the recipe above; each server uses 8 devices.

  • Data + Expert Parallel: --data-parallel-size sets the CR13 devices that each hold a model replica, and --enable-expert-parallel shards the experts across them.
  • Pipeline Parallel: --pipeline-parallel-size splits the layers across the devices.
  • Prefill/Decode disaggregated: a prefill server and a decode server run on separate nodes behind the llm-d router, whose helm command follows the two servers in the recipe. For how to deploy the llm-d router and the serving stack on Kubernetes, see Deployment on Kubernetes in the Coding Assistant reference design.

Every command runs on the vllm model implementation, selected with --model-impl vllm. The custom architecture requires --trust-remote-code, which every command carries.

See the Configuration Guide for all engine arguments.

Notes and limitations

  • Verified on RBLN-CR13; the exact serve flags may change.
  • This recipe does not include a configuration for the original FP8 checkpoint (skt/A.X-K2). At 262,144 tokens its KV cache takes 21.1 GiB per rank at --block-size 2048, which an RBLN-CR13 device cannot hold beside its weights, so the server fails with SYS_EBUSY: Lack of device memory. It serves with a shorter context: --max-model-len 163840 brings the KV cache to 13.3 GiB per rank.
  • The batch size, sequence length, and device count are fixed at compilation; changing any of these settings triggers recompilation at the next startup.