A.X-K2

Serve A.X-K2 with Dynamo and vLLM on B200, aggregated or disaggregated.

View as Markdown

Deploy SK Telecom’s A.X-K2 with NVIDIA Dynamo, NVFP4 weights, FP8 KV cache, KV-aware routing, and EAGLE3 speculative decoding. Both topologies use four B200 GPUs per worker and support text generation, reasoning, and tool calling.

Choose your deployment target

Topology
Checkpoint skt/A.X-K2-NVFP4Hardware 8x B200, two TP4 aggregate workersPrecision NVFP4 weights, FP8 KV cacheSpeculation EAGLE3, 3 tokensFlashInfer autotuning Disabled
Checkpoint skt/A.X-K2-NVFP4Hardware 12x B200, two TP4 prefill workers and one TP4 decode workerPrecision NVFP4 weights, FP8 KV cacheSpeculation EAGLE3, 3 tokens on both rolesFlashInfer autotuning Disabled on prefill and decode

Prerequisites

  • A Kubernetes cluster with the Dynamo platform installed.
  • Four B200 GPUs and 400 GiB of host memory per worker: eight GPUs total for aggregated serving or twelve for disaggregated serving.
  • A ReadWriteMany storage class for model weights.
  • A Hugging Face token with access to skt/A.X-K2-NVFP4 and skt/A.X-K2-EAGLE3.

Disaggregated serving also requires InfiniBand with one rdma/shared_ib device resource available per worker for NIXL/UCX KV transfer.

Deploy

Run commands from the repository root. Set CONTEXT and NAMESPACE to your cluster context and namespace, and set HF_TOKEN to your Hugging Face token. Create the namespace and token secret:

kubectl --context "${CONTEXT}" create namespace "${NAMESPACE}"
kubectl --context "${CONTEXT}" -n "${NAMESPACE}" create secret generic hf-token-secret \
--from-literal=HF_TOKEN="${HF_TOKEN}"

Set storageClassName in model-cache.yaml to your cluster’s ReadWriteMany storage class. Create the PVC and download both pinned checkpoints:

kubectl --context "${CONTEXT}" -n "${NAMESPACE}" apply -f recipes/ax-k2/model-cache/model-cache.yaml
kubectl --context "${CONTEXT}" -n "${NAMESPACE}" apply -f recipes/ax-k2/model-cache/model-download.yaml
kubectl --context "${CONTEXT}" -n "${NAMESPACE}" wait --for=condition=Complete \
job/axk2-model-download --timeout=14400s

Both profiles and the download job mount model-cache. To use an existing populated PVC, set claimName in the job and selected deployment source to that PVC and skip creating a new claim. Regenerate the selected manifest after editing its Kustomize source.

Deploy two aggregate workers:

export DGD=axk2-agg-b200-chat
kubectl --context "${CONTEXT}" -n "${NAMESPACE}" apply \
-f recipes/ax-k2/vllm/agg-b200-chat/deploy-generic.yaml

The manifest is generated from Kustomize sources. To change a setting, edit kustomize/base/deploy.yaml and regenerate:

python3 scripts/kustomize-matrix.py unfold recipes/ax-k2/vllm/agg-b200-chat/.kustomize-matrix.yaml
python3 scripts/kustomize-matrix.py render recipes/ax-k2/vllm/agg-b200-chat/.kustomize-matrix.yaml

Deploy two prefill workers and one decode worker:

export DGD=axk2-disagg-b200-chat
kubectl --context "${CONTEXT}" -n "${NAMESPACE}" apply \
-f recipes/ax-k2/vllm/disagg-b200-chat/deploy-generic.yaml

The manifest is generated from Kustomize sources. To change a setting, edit kustomize/base/deploy.yaml and regenerate:

python3 scripts/kustomize-matrix.py unfold recipes/ax-k2/vllm/disagg-b200-chat/.kustomize-matrix.yaml
python3 scripts/kustomize-matrix.py render recipes/ax-k2/vllm/disagg-b200-chat/.kustomize-matrix.yaml

Smoke Test

Wait for the selected deployment, then forward its frontend port:

kubectl --context "${CONTEXT}" -n "${NAMESPACE}" wait --for=condition=Ready \
"dynamographdeployment/${DGD}" --timeout=7200s
kubectl --context "${CONTEXT}" -n "${NAMESPACE}" port-forward \
"service/${DGD}-frontend" 8000:8000

In another terminal, verify model discovery and a completion:

curl --fail http://localhost:8000/v1/models
curl --fail http://localhost:8000/v1/chat/completions \
-H 'Content-Type: application/json' \
-d '{"model":"skt/A.X-K2-NVFP4","messages":[{"role":"user","content":"What is 2 + 2?"}],"temperature":0,"max_tokens":4096,"stream":false}'

Model discovery should list skt/A.X-K2-NVFP4; the completion should contain a choices array. Reasoning uses the deepseek_v3 parser, and tool calling uses the hermes parser.

Benchmark

The AIPerf workflow replays an 8K-input / 1K-output chat trace with 70% KV reuse against the aggregated profile at concurrency 32. It preserves expanded configurations, raw reports, and frontend metrics. Keep the default real EAGLE3 acceptance configuration for accuracy evaluation.

Before starting the benchmark, follow the workflow’s trace-staging instructions to copy the Git LFS chat trace from your checkout onto the model-cache PVC. The job reads it through TRACE_FILE and verifies its SHA-256 and request counts before sending traffic.

The checked-in benchmark job targets the aggregate frontend. To benchmark the disaggregated profile, update its endpoint and frontend selector to this deployment before applying the job.

Performance Results

Benchmarking uses synthetic EAGLE3 acceptance length 2.12.

WorkloadFrameworkRecipeSKUConcurrencySystem output tok/s/GPUUser output tok/s (mean)TTFT P50 (seconds)
Chat (15% subset)vLLMAggregated (2 workers)B2001676.4955.830.525
Chat (15% subset)vLLMDisaggregated (2P1D)B2001661.97833.879

Compare All Targets

SettingAggregatedDisaggregated
Workers2 aggregate2 prefill + 1 decode
Total GPUs8x B20012x B200
ParallelismTP4, DP1, expert parallelism disabledSame on both roles
PrecisionNVFP4 weights, FP8 KV cacheSame on both roles
AttentionFLASHINFER_MLA_SPARSESame on both roles
FlashInfer autotuningDisabledDisabled on prefill and decode
EAGLE33 speculative tokens3 speculative tokens on both roles
Async schedulingEnabledPrefill disabled, decode enabled
Context length262,144 tokens262,144 tokens
RoutingKV-awareKV-aware
KV transferN/ANIXL/UCX over InfiniBand

Source