Add spec decoding and multi-nodepool support to kimi-k3

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) <br> <br> [Link for Using Google Cloud Storage (GCS) as Storage Option]((./inference/a4x/single-host-serving/tensorrt-llm-gcs/README.md)) <br> <br> [Link for Using Lustre as Storage Option]((./inference/a4x/single-host-serving/tensorrt-llm-lustre/README.md))

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:

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