Kimi k3 model streaming

inference/a4x/multi-host-serving/sglang/README.md

@@ -57,6 +57,7 @@ export NAMESPACE="default"
 export GCS_BUCKET="<YOUR_GCS_BUCKET_NAME>"
 export RELEASE_NAME="sglang-kimi-k3"
 export K8S_SERVICE_ACCOUNT="workload-identity-k8s-sa"
+export MODEL_PATH="gs://${GCS_BUCKET}/huggingface_model_cache/models--moonshotai--Kimi-K3/"
 
 gcloud container clusters get-credentials ${CLUSTER_NAME} --region ${REGION}
 ```
@@ -90,6 +91,7 @@ helm install ${RELEASE_NAME} \
   --namespace ${NAMESPACE} \
   -f values_kimi_k3.yaml \
   --set volumes.gcsfuse.bucketName=${GCS_BUCKET} \
+  --set workload.modelPath=${MODEL_PATH} \
   --set workload.serviceAccountName=${K8S_SERVICE_ACCOUNT} \
   ../../../../src/helm-charts/a4x/inference-templates/lws-deployment
 ```

inference/a4x/multi-host-serving/sglang/values_kimi_k3.yaml

@@ -1,11 +1,24 @@
 workload:
   model: moonshotai/Kimi-K3
+  servedModelName: moonshotai/Kimi-K3
   modelPath: /data/model/huggingface_model_cache/models--moonshotai--Kimi-K3
   image: lmsysorg/sglang:kimi-k3
   framework: sglang
   nnodes: 4
   tpSize: 16
   serviceAccountName: workload-identity-k8s-sa
+  extraPipPackages: "hf_transfer fla-core google-cloud-storage"
+  extraEnv:
+  - name: HF_HOME
+    value: "/gcs-cache/hf_cache"
+  - name: TMPDIR
+    value: "/gcs-cache/tmp"
+  - name: RUNAI_STREAMER_MEMORY_LIMIT
+    value: "34359738368"
+  - name: RUNAI_STREAMER_CONCURRENCY
+    value: "16"
+  - name: PYTORCH_CUDA_ALLOC_CONF
+    value: "expandable_segments:True"
   extraArgs: |
     --trust-remote-code \
     --reasoning-parser kimi_k3 \

src/helm-charts/a4x/inference-templates/lws-deployment/templates/lws-launcher.yaml

@@ -112,6 +112,9 @@ spec:
             value: "PIX"
           - name: NCCL_NET_GDR_C2C
             value: "1"
+          {{- if .Values.workload.extraEnv }}
+          {{- toYaml .Values.workload.extraEnv | nindent 10 }}
+          {{- end }}
           volumeMounts:
           - mountPath: /usr/local/gib
             name: gib
@@ -125,6 +128,7 @@ spec:
           args:
           - |
             set -e
+            mkdir -p /gcs-cache/hf_cache /gcs-cache/tmp || true
             nvidia-smi
             . /usr/local/gib/scripts/set_nccl_env.sh
             {{- if .Values.workload.extraPipPackages }}
@@ -147,6 +151,9 @@ spec:
 
             python3 -m sglang.launch_server \
               --model-path {{ .Values.workload.modelPath | default (printf "/data/model/%s" .Values.workload.model) }} \
+              {{- if .Values.workload.servedModelName }}
+              --served-model-name {{ .Values.workload.servedModelName }} \
+              {{- end }}
               --tp-size {{ .Values.workload.tpSize | default 16 }} \
               --nnodes ${SIZE} \
               --node-rank ${RANK} \
@@ -259,6 +266,9 @@ spec:
             value: "PIX"
           - name: NCCL_NET_GDR_C2C
             value: "1"
+          {{- if .Values.workload.extraEnv }}
+          {{- toYaml .Values.workload.extraEnv | nindent 10 }}
+          {{- end }}
           volumeMounts:
           - mountPath: /usr/local/gib
             name: gib
@@ -272,6 +282,7 @@ spec:
           args:
           - |
             set -e
+            mkdir -p /gcs-cache/hf_cache /gcs-cache/tmp || true
             nvidia-smi
             . /usr/local/gib/scripts/set_nccl_env.sh
             {{- if .Values.workload.extraPipPackages }}
@@ -294,6 +305,9 @@ spec:
 
             python3 -m sglang.launch_server \
               --model-path {{ .Values.workload.modelPath | default (printf "/data/model/%s" .Values.workload.model) }} \
+              {{- if .Values.workload.servedModelName }}
+              --served-model-name {{ .Values.workload.servedModelName }} \
+              {{- end }}
               --tp-size {{ .Values.workload.tpSize | default 16 }} \
               --nnodes ${SIZE} \
               --node-rank ${RANK} \