Repository: kaito-project/kaito

Issue:
vLLM sometimes doesn't work when `PYTORCH_CUDA_ALLOC_CONF` is set to `expandable_segments:True`
**Is your feature request related to a problem? Please describe.**

vLLM deployment crashed with the following error message:

![Image](https://github.com/user-attachments/assets/69329358-9efd-45c3-b459-4719f377a92f)

Found this from vLLM repository: https://github.com/vllm-project/vllm/issues/9046.

It works fine for tensor parallelism = 1 but anything >1 would crash vLLM.

**Describe the solution you'd like**

Remove the environment variable from vLLM deployment but not sure what effect it would have.

**Describe alternatives you've considered**

**Additional context**

Hints from issue comments:
@zhuangqh: The main idea of adding this env was to reduce memory Fragments.

```yaml
apiVersion: kaito.sh/v1beta1
kind: Workspace
metadata:
  name: workspace-phi-4-mini
resource:
  instanceType: "Standard_NC12s_v3"
  labelSelector:
    matchLabels:
      apps: phi-4
inference:
  preset:
    name: phi-4-mini-instruct 
  config: "inference-params"
---
apiVersion: v1
kind: ConfigMap
metadata:
  name: inference-params
data:
  inference_config.yaml: |
    # Maximum number of steps to find the max available seq len fitting in the GPU memory.
    max_probe_steps: 6

    vllm:
      cpu-offload-gb: 0
      gpu-memory-utilization: 0.95
      swap-space: 4
      max-model-len: 2048
```

I can successfully run this workspace with 2 tensor parallelism. In fact, we have a e2e test to cover this case.
Can you confirm how to reproduce this issue? do you run it on multi nodes?@chewong: I am using one pre-provisioned `Standard_NC48ads_A100_v4` (2 x A100) instance in my setup. The rest of the setup is the same as yours. Let me give it a try again.@chewong: Tested with 1 `Standard_NC12s_v3` + TP = 2 + `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` and it works with no problem. I am starting to suspect that only certain GPUs are affected.@chewong: ![Image](https://github.com/user-attachments/assets/154ec801-03e8-4c9e-b37b-439b2f543805)

I got the same error when running TP=2 with two H100s on a node, which confirms my previous observation that this issue only affects certain GPU models. Current affected models are A100s and H100s.

Edit: again, it works after removing the environment variable/setting the value to False:

![Image](https://github.com/user-attachments/assets/0e1e55c4-4509-42d1-af55-394a25924ee9)

![Image](https://github.com/user-attachments/assets/e35cd1da-2bd8-4af6-9466-6f0ca37ad460)



Solution patch:
diff --git a/pkg/workspace/manifests/manifests.go b/pkg/workspace/manifests/manifests.go
index f4ef3c9..ba2863a 100644
--- a/pkg/workspace/manifests/manifests.go
+++ b/pkg/workspace/manifests/manifests.go
@@ -100,11 +100,6 @@ func GenerateStatefulSetManifest(workspaceObj *kaitov1beta1.Workspace, imageName
 	labelselector := &v1.LabelSelector{
 		MatchLabels: selector,
 	}
-	// Add PYTORCH_CUDA_ALLOC_CONF environment variable
-	envVars = append(envVars, corev1.EnvVar{
-		Name:  "PYTORCH_CUDA_ALLOC_CONF",
-		Value: "expandable_segments:True",
-	})
 
 	ss := &appsv1.StatefulSet{
 		ObjectMeta: v1.ObjectMeta{
@@ -248,11 +243,6 @@ func GenerateDeploymentManifest(workspaceObj *kaitov1beta1.Workspace, revisionNu
 	labelselector := &v1.LabelSelector{
 		MatchLabels: selector,
 	}
-	// Add PYTORCH_CUDA_ALLOC_CONF environment variable
-	envVars = append(envVars, corev1.EnvVar{
-		Name:  "PYTORCH_CUDA_ALLOC_CONF",
-		Value: "expandable_segments:True",
-	})
 
 	pullerContainers, pullerEnvVars, pullerVolumes := GeneratePullerContainers(workspaceObj, volumeMount)
 	envVars = append(envVars, pullerEnvVars...)



Test patch:
diff --git a/pkg/workspace/inference/preset-inferences_test.go b/pkg/workspace/inference/preset-inferences_test.go
index f7dd955..1d7cc80 100644
--- a/pkg/workspace/inference/preset-inferences_test.go
+++ b/pkg/workspace/inference/preset-inferences_test.go
@@ -52,10 +52,6 @@ func TestCreatePresetInference(t *testing.T) {
 			// So expected cmd consists of shell command and inference file
 			expectedCmd: "/bin/sh -c python3 /workspace/vllm/inference_api.py --tensor-parallel-size=2 --served-model-name=mymodel --gpu-memory-utilization=0.90 --kaito-config-file=/mnt/config/inference_config.yaml",
 			hasAdapters: false,
-			expectedEnvVars: []corev1.EnvVar{{
-				Name:  "PYTORCH_CUDA_ALLOC_CONF",
-				Value: "expandable_segments:True",
-			}},
 		},
 
 		"test-model-no-parallel/vllm": {
@@ -72,10 +68,6 @@ func TestCreatePresetInference(t *testing.T) {
 			// So expected cmd consists of shell command and inference file
 			expectedCmd: "/bin/sh -c python3 /workspace/vllm/inference_api.py --kaito-config-file=/mnt/config/inference_config.yaml --gpu-memory-utilization=0.90",
 			hasAdapters: false,
-			expectedEnvVars: []corev1.EnvVar{{
-				Name:  "PYTORCH_CUDA_ALLOC_CONF",
-				Value: "expandable_segments:True",
-			}},
 		},
 
 		"test-model-no-lora-support/vllm": {
@@ -92,10 +84,6 @@ func TestCreatePresetInference(t *testing.T) {
 			// So expected cmd consists of shell command and inference file
 			expectedCmd: "/bin/sh -c python3 /workspace/vllm/inference_api.py --kaito-config-file=/mnt/config/inference_config.yaml --gpu-memory-utilization=0.90",
 			hasAdapters: false,
-			expectedEnvVars: []corev1.EnvVar{{
-				Name:  "PYTORCH_CUDA_ALLOC_CONF",
-				Value: "expandable_segments:True",
-			}},
 		},
 
 		"test-model-with-adapters/vllm": {
@@ -112,9 +100,6 @@ func TestCreatePresetInference(t *testing.T) {
 			hasAdapters:    true,
 			expectedVolume: "adapter-volume",
 			expectedEnvVars: []corev1.EnvVar{{
-				Name:  "PYTORCH_CUDA_ALLOC_CONF",
-				Value: "expandable_segments:True",
-			}, {
 				Name:  "Adapter-1",
 				Value: "0.5",
 			}},
@@ -134,10 +119,6 @@ func TestCreatePresetInference(t *testing.T) {
 			// So expected cmd consists of shell command and inference file
 			expectedCmd: "/bin/sh -c accelerate launch /workspace/tfs/inference_api.py",
 			hasAdapters: false,
-			expectedEnvVars: []corev1.EnvVar{{
-				Name:  "PYTORCH_CUDA_ALLOC_CONF",
-				Value: "expandable_segments:True",
-			}},
 		},
 
 		"test-model-with-adapters": {
@@ -154,9 +135,6 @@ func TestCreatePresetInference(t *testing.T) {
 			hasAdapters:    true,
 			expectedVolume: "adapter-volume",
 			expectedEnvVars: []corev1.EnvVar{{
-				Name:  "PYTORCH_CUDA_ALLOC_CONF",
-				Value: "expandable_segments:True",
-			}, {
 				Name:  "Adapter-1",
 				Value: "0.5",
 			}},
@@ -183,9 +161,6 @@ func TestCreatePresetInference(t *testing.T) {
 						},
 					},
 				},
-			}, {
-				Name:  "PYTORCH_CUDA_ALLOC_CONF",
-				Value: "expandable_segments:True",
 			}},
 		},
 
@@ -210,9 +185,6 @@ func TestCreatePresetInference(t *testing.T) {
 						},
 					},
 				},
-			}, {
-				Name:  "PYTORCH_CUDA_ALLOC_CONF",
-				Value: "expandable_segments:True",
 			}},
 		},
 	}