diff --git a/README.md b/README.md
index 29d48fa7..2108de27 100644
--- a/README.md
+++ b/README.md
@@ -163,6 +163,7 @@ Models | GPU Machine Type
| Models | GPU Machine Type | Framework | Workload Type | Orchestrator | Link to the recipe |
| ---------------- | ---------------- | --------- | ------------------- | ------------ | ------------------ |
| **DeepSeek R1 671B** | [A4X (NVIDIA GB200)](https://cloud.google.com/compute/docs/accelerator-optimized-machines#a4x-vms) | vLLM (v0.14.0rc1) | Inference | GKE | [Link](./inference/a4x/single-host-serving/vllm/README.md)
+| **Kimi-K3** | [A4X (NVIDIA GB200)](https://cloud.google.com/compute/docs/accelerator-optimized-machines#a4x-vms) | SGLang | Inference | GKE | [Link](./inference/a4x/multi-host-serving/sglang/README.md)
| **Wan2.2 T2V A14B Diffusers** | [A4X (NVIDIA GB200)](https://cloud.google.com/compute/docs/accelerator-optimized-machines#a4x-vms) | SGLang (latest) | Inference | GKE | [Link](./inference/a4x/single-host-serving/sglang/README.md)
| **Wan2.2 I2V A14B Diffusers** | [A4X (NVIDIA GB200)](https://cloud.google.com/compute/docs/accelerator-optimized-machines#a4x-vms) | SGLang (latest) | Inference | GKE | [Link](./inference/a4x/single-host-serving/sglang/README.md)
| **DeepSeek R1 671B** | [A4X (NVIDIA GB200)](https://cloud.google.com/compute/docs/accelerator-optimized-machines#a4x-vms) | TensorRT-LLM (1.3.0rc5) | Inference | GKE | [Link](./inference/a4x/single-host-serving/tensorrt-llm/README.md)
[Link for Using Google Cloud Storage (GCS) as Storage Option]((./inference/a4x/single-host-serving/tensorrt-llm-gcs/README.md))
[Link for Using Lustre as Storage Option]((./inference/a4x/single-host-serving/tensorrt-llm-lustre/README.md))
diff --git a/inference/a4x/multi-host-serving/sglang/values_kimi_k3.yaml b/inference/a4x/multi-host-serving/sglang/values_kimi_k3.yaml
index 1de607e9..ff897b93 100644
--- a/inference/a4x/multi-host-serving/sglang/values_kimi_k3.yaml
+++ b/inference/a4x/multi-host-serving/sglang/values_kimi_k3.yaml
@@ -19,13 +19,25 @@ workload:
value: "16"
- name: PYTORCH_CUDA_ALLOC_CONF
value: "expandable_segments:True"
+ - name: TORCH_NCCL_ENABLE_MONITORING
+ value: "0"
+ - name: TORCH_NCCL_ASYNC_ERROR_HANDLING
+ value: "0"
+ - name: TORCH_DISTRIBUTED_TIMEOUT
+ value: "3600"
+ - name: TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC
+ value: "0"
extraArgs: |
+ --dist-timeout 3600 \
--trust-remote-code \
--reasoning-parser kimi_k3 \
--tool-call-parser kimi_k3 \
--mamba-full-memory-ratio 7.21 \
--dcp-size 16 \
- --mem-fraction-static 0.85
+ --mem-fraction-static 0.85 \
+ --speculative-algorithm DSPARK \
+ --speculative-draft-model-path RadixArk/Kimi-K3-DSpark \
+ --speculative-dspark-block-size 7
volumes:
gcsfuse:
diff --git a/src/helm-charts/a4x/inference-templates/lws-deployment/templates/lws-launcher.yaml b/src/helm-charts/a4x/inference-templates/lws-deployment/templates/lws-launcher.yaml
index 3f9c5ddf..aab813fd 100644
--- a/src/helm-charts/a4x/inference-templates/lws-deployment/templates/lws-launcher.yaml
+++ b/src/helm-charts/a4x/inference-templates/lws-deployment/templates/lws-launcher.yaml
@@ -30,6 +30,24 @@ spec:
automountServiceAccountToken: true
subdomain: {{ .Release.Name }}
restartPolicy: Always
+ {{- if .Values.workload.nodeSelector }}
+ nodeSelector:
+ {{- toYaml .Values.workload.nodeSelector | nindent 10 }}
+ {{- end }}
+ affinity:
+ {{- if .Values.workload.affinity }}
+ {{- toYaml .Values.workload.affinity | nindent 10 }}
+ {{- else }}
+ podAffinity:
+ requiredDuringSchedulingIgnoredDuringExecution:
+ - labelSelector:
+ matchExpressions:
+ - key: app
+ operator: In
+ values:
+ - {{ .Release.Name }}
+ topologyKey: "cloud.google.com/gce-topology-subblock"
+ {{- end }}
resourceClaims:
- name: compute-domain-channel
resourceClaimTemplateName: {{ .Release.Name }}-compute-domain-claim
@@ -159,10 +177,8 @@ spec:
--node-rank ${RANK} \
--dist-init-addr ${LEADER_HOST}:20000 \
--host 0.0.0.0 \
- --port 30100 \
- {{- if .Values.workload.extraArgs }}
- {{ .Values.workload.extraArgs | nindent 14 }}
- {{- end }} 2>&1 | tee /gcs-cache/sglang_server_${RANK}.log
+ --port 30100 {{- if .Values.workload.extraArgs }} \
+ {{ .Values.workload.extraArgs | trim | nindent 14 | trimPrefix "\n" }}{{- end }} 2>&1 | tee /gcs-cache/sglang_server_${RANK}.log
workerTemplate:
metadata:
@@ -184,6 +200,24 @@ spec:
automountServiceAccountToken: true
subdomain: {{ .Release.Name }}
restartPolicy: Always
+ {{- if .Values.workload.nodeSelector }}
+ nodeSelector:
+ {{- toYaml .Values.workload.nodeSelector | nindent 10 }}
+ {{- end }}
+ affinity:
+ {{- if .Values.workload.affinity }}
+ {{- toYaml .Values.workload.affinity | nindent 10 }}
+ {{- else }}
+ podAffinity:
+ requiredDuringSchedulingIgnoredDuringExecution:
+ - labelSelector:
+ matchExpressions:
+ - key: app
+ operator: In
+ values:
+ - {{ .Release.Name }}
+ topologyKey: "cloud.google.com/gce-topology-subblock"
+ {{- end }}
resourceClaims:
- name: compute-domain-channel
resourceClaimTemplateName: {{ .Release.Name }}-compute-domain-claim
@@ -313,7 +347,5 @@ spec:
--node-rank ${RANK} \
--dist-init-addr ${LEADER_HOST}:20000 \
--host 0.0.0.0 \
- --port 30100 \
- {{- if .Values.workload.extraArgs }}
- {{ .Values.workload.extraArgs | nindent 14 }}
- {{- end }} 2>&1 | tee /gcs-cache/sglang_server_${RANK}.log
+ --port 30100 {{- if .Values.workload.extraArgs }} \
+ {{ .Values.workload.extraArgs | trim | nindent 14 | trimPrefix "\n" }}{{- end }} 2>&1 | tee /gcs-cache/sglang_server_${RANK}.log