fix: correct MinInferenceVRAMGB and TotalVRAMGB calculation in PD planner (#1236)

Co-authored-by: cemeng <cemengzhang@gmail.com>

Author: csg-pr-bot

common/types/pd_planner.go

@@ -346,13 +346,15 @@ const DefaultGPUsPerNode = 8
 //   - precision: inference precision (fp8, bf16, etc.)
 //   - hiddenSize: hidden size from config.json (used to estimate NonMoEParamsB)
 //   - numHiddenLayers: number of hidden layers from config.json
+//   - minInferenceVRAMGB: minimum VRAM per GPU to load and run inference, typically
+//     from metadata.mini_gpu_memory_gb. When <= 0, falls back to DefaultGPUUnitGB (80).
 //
 // The function:
 //  1. Estimates NonMoEParamsB from hidden_size and num_hidden_layers when available,
 //     otherwise falls back to the active/total expert ratio heuristic.
 //  2. Plans prefill and decode configurations using PlanPD.
 //  3. Returns a PDRecommendation suitable for storage as JSONB.
-func PlanPDRecommendation(modelName string, totalParamsB float64, totalExperts, activeExperts int, precision string, hiddenSize, numHiddenLayers int) (*PDRecommendation, error) {
+func PlanPDRecommendation(modelName string, totalParamsB float64, totalExperts, activeExperts int, precision string, hiddenSize, numHiddenLayers int, minInferenceVRAMGB float64) (*PDRecommendation, error) {
 	if totalParamsB <= 0 {
 		return nil, fmt.Errorf("total params must be positive, got %f", totalParamsB)
 	}
@@ -396,23 +398,30 @@ func PlanPDRecommendation(modelName string, totalParamsB float64, totalExperts,
 		return nil, fmt.Errorf("failed to plan PD for model %s: %w", modelName, err)
 	}
 
+	// When minInferenceVRAMGB is 0 (metadata.mini_gpu_memory_gb not populated),
+	// keep it as 0 so the missing value is visible rather than masked by a
+	// default. PD planning checks still pass; only the VRAM fields are
+	// inaccurate until the value is resolved from TOML or metadata.
+	if minInferenceVRAMGB < 0 {
+		minInferenceVRAMGB = 0
+	}
+
 	return &PDRecommendation{
 		ModelName:          modelName,
 		TotalParamsB:       spec.TotalParamsB,
 		NonMoEParamsB:      spec.NonMoEParamsB,
 		TotalExperts:       spec.TotalExperts,
 		ActiveExperts:      activeExperts,
 		Precision:          spec.Precision,
-		MinInferenceVRAMGB: DefaultGPUUnitGB,
-		Prefill:            planResultToRoleConfig(prefill),
-		Decode:             planResultToRoleConfig(decode),
+		MinInferenceVRAMGB: minInferenceVRAMGB,
+		Prefill:            planResultToRoleConfig(prefill, minInferenceVRAMGB),
+		Decode:             planResultToRoleConfig(decode, minInferenceVRAMGB),
 	}, nil
 }
 
 // planResultToRoleConfig converts a PDPlanResult into a PDRoleConfig.
-// TotalVRAMGB is computed as MinInferenceVRAMGB * TotalGPUs by the caller
-// (PlanPDRecommendation sets MinInferenceVRAMGB = DefaultGPUUnitGB).
-func planResultToRoleConfig(result PDPlanResult) PDRoleConfig {
+// TotalVRAMGB is computed as minInferenceVRAMGB * Pods.
+func planResultToRoleConfig(result PDPlanResult, minInferenceVRAMGB float64) PDRoleConfig {
 	pods := result.LWSSize
 	if pods < 1 {
 		pods = 1
@@ -423,6 +432,6 @@ func planResultToRoleConfig(result PDPlanResult) PDRoleConfig {
 		DP:          result.DP,
 		TotalGPUs:   result.TotalGPUs,
 		Pods:        pods,
-		TotalVRAMGB: float64(result.TotalGPUs) * DefaultGPUUnitGB,
+		TotalVRAMGB: minInferenceVRAMGB * float64(pods),
 	}
 }

common/types/pd_planner_test.go

@@ -256,27 +256,56 @@ func TestPlanPD_FullMatrix(t *testing.T) {
 }
 
 func TestPlanPDRecommendation_DeepSeekV3(t *testing.T) {
-	rec, err := PlanPDRecommendation("deepseek-v3", 671, 256, 8, "fp8", 7168, 61)
+	// minInferenceVRAMGB=772.0 matches configs/pd/models.toml for DeepSeek-V3.
+	rec, err := PlanPDRecommendation("deepseek-v3", 671, 256, 8, "fp8", 7168, 61, 772.0)
 	require.NoError(t, err)
 	require.NotNil(t, rec)
 	require.Equal(t, 671.0, rec.TotalParamsB)
 	require.Equal(t, 256, rec.TotalExperts)
 	require.Equal(t, 8, rec.ActiveExperts)
 	require.Equal(t, "fp8", rec.Precision)
-	require.Equal(t, 80.0, rec.MinInferenceVRAMGB)
+	require.Equal(t, 772.0, rec.MinInferenceVRAMGB)
 	require.Greater(t, rec.Prefill.TotalGPUs, 0)
 	require.Greater(t, rec.Decode.TotalGPUs, 0)
-	require.Greater(t, rec.Prefill.TotalVRAMGB, 0.0)
+	// TotalVRAMGB = MinInferenceVRAMGB * Pods
+	require.Equal(t, 772.0*float64(rec.Prefill.Pods), rec.Prefill.TotalVRAMGB)
+	require.Equal(t, 772.0*float64(rec.Decode.Pods), rec.Decode.TotalVRAMGB)
 	require.Greater(t, rec.Prefill.Pods, 0)
 	require.Greater(t, rec.Decode.Pods, 0)
 
-	t.Logf("DeepSeek-V3 80G: Prefill TP=%d EP=%d GPUs=%d Pods=%d | Decode TP=%d EP=%d GPUs=%d Pods=%d",
-		rec.Prefill.TP, rec.Prefill.EP, rec.Prefill.TotalGPUs, rec.Prefill.Pods,
-		rec.Decode.TP, rec.Decode.EP, rec.Decode.TotalGPUs, rec.Decode.Pods)
+	t.Logf("DeepSeek-V3 80G: Prefill TP=%d EP=%d GPUs=%d Pods=%d VRAM=%.0f | Decode TP=%d EP=%d GPUs=%d Pods=%d VRAM=%.0f",
+		rec.Prefill.TP, rec.Prefill.EP, rec.Prefill.TotalGPUs, rec.Prefill.Pods, rec.Prefill.TotalVRAMGB,
+		rec.Decode.TP, rec.Decode.EP, rec.Decode.TotalGPUs, rec.Decode.Pods, rec.Decode.TotalVRAMGB)
+}
+
+func TestPlanPDRecommendation_MinInferenceVRAMZero(t *testing.T) {
+	// When minInferenceVRAMGB is 0 (metadata.mini_gpu_memory_gb not populated),
+	// keep it as 0 so the missing value is visible. PD planning still succeeds;
+	// only the VRAM fields are 0 until the value is resolved from TOML/metadata.
+	rec, err := PlanPDRecommendation("deepseek-v3", 671, 256, 8, "fp8", 7168, 61, 0)
+	require.NoError(t, err)
+	require.NotNil(t, rec)
+	require.Equal(t, 0.0, rec.MinInferenceVRAMGB)
+	// TotalVRAMGB = 0 * Pods = 0
+	require.Equal(t, 0.0, rec.Prefill.TotalVRAMGB)
+	require.Equal(t, 0.0, rec.Decode.TotalVRAMGB)
+	// PD planning checks should still pass
+	require.Greater(t, rec.Prefill.TotalGPUs, 0)
+	require.Greater(t, rec.Decode.TotalGPUs, 0)
+}
+
+func TestPlanPDRecommendation_MinInferenceVRAMNegative(t *testing.T) {
+	// Negative values should be clamped to 0.
+	rec, err := PlanPDRecommendation("deepseek-v3", 671, 256, 8, "fp8", 7168, 61, -10)
+	require.NoError(t, err)
+	require.NotNil(t, rec)
+	require.Equal(t, 0.0, rec.MinInferenceVRAMGB)
+	require.Equal(t, 0.0, rec.Prefill.TotalVRAMGB)
+	require.Equal(t, 0.0, rec.Decode.TotalVRAMGB)
 }
 
 func TestPlanPDRecommendation_GLM52(t *testing.T) {
-	rec, err := PlanPDRecommendation("glm-5.2", 1000, 256, 8, "fp8", 8192, 80)
+	rec, err := PlanPDRecommendation("glm-5.2", 1000, 256, 8, "fp8", 8192, 80, 0)
 	require.NoError(t, err)
 	require.NotNil(t, rec)
 	require.Equal(t, 1000.0, rec.TotalParamsB)
@@ -285,7 +314,7 @@ func TestPlanPDRecommendation_GLM52(t *testing.T) {
 }
 
 func TestPlanPDRecommendation_Qwen3_30B_A3B(t *testing.T) {
-	rec, err := PlanPDRecommendation("qwen3-30b-a3b", 30, 128, 8, "bf16", 2048, 48)
+	rec, err := PlanPDRecommendation("qwen3-30b-a3b", 30, 128, 8, "bf16", 2048, 48, 0)
 	require.NoError(t, err)
 	require.NotNil(t, rec)
 	require.Equal(t, 30.0, rec.TotalParamsB)
@@ -294,7 +323,7 @@ func TestPlanPDRecommendation_Qwen3_30B_A3B(t *testing.T) {
 
 func TestPlanPDRecommendation_UnknownModelMoE(t *testing.T) {
 	// Unknown MoE model with expert info from config.json
-	rec, err := PlanPDRecommendation("custom/moe-model", 200, 64, 8, "bf16", 4096, 32)
+	rec, err := PlanPDRecommendation("custom/moe-model", 200, 64, 8, "bf16", 4096, 32, 0)
 	require.NoError(t, err)
 	require.NotNil(t, rec)
 	require.Equal(t, 200.0, rec.TotalParamsB)
@@ -306,7 +335,7 @@ func TestPlanPDRecommendation_UnknownModelMoE(t *testing.T) {
 }
 
 func TestPlanPDRecommendation_DenseModel(t *testing.T) {
-	rec, err := PlanPDRecommendation("custom-dense-70b", 70, 0, 0, "bf16", 0, 0)
+	rec, err := PlanPDRecommendation("custom-dense-70b", 70, 0, 0, "bf16", 0, 0, 0)
 	require.NoError(t, err)
 	require.NotNil(t, rec)
 	require.Equal(t, 70.0, rec.TotalParamsB)
@@ -315,35 +344,35 @@ func TestPlanPDRecommendation_DenseModel(t *testing.T) {
 }
 
 func TestPlanPDRecommendation_ZeroParams(t *testing.T) {
-	_, err := PlanPDRecommendation("test", 0, 256, 8, "fp8", 0, 0)
+	_, err := PlanPDRecommendation("test", 0, 256, 8, "fp8", 0, 0, 0)
 	require.Error(t, err)
 	require.Contains(t, err.Error(), "total params must be positive")
 }
 
 func TestPlanPDRecommendation_NegativeParams(t *testing.T) {
-	_, err := PlanPDRecommendation("test", -10, 256, 8, "fp8", 0, 0)
+	_, err := PlanPDRecommendation("test", -10, 256, 8, "fp8", 0, 0, 0)
 	require.Error(t, err)
 	require.Contains(t, err.Error(), "total params must be positive")
 }
 
 func TestPlanPDRecommendation_PrecisionOverride(t *testing.T) {
 	// Known model with explicit precision override
 	// Use a smaller model that fits in 80GB with bf16
-	rec, err := PlanPDRecommendation("qwen3-30b-a3b", 30, 128, 8, "bf16", 2048, 48)
+	rec, err := PlanPDRecommendation("qwen3-30b-a3b", 30, 128, 8, "bf16", 2048, 48, 0)
 	require.NoError(t, err)
 	require.Equal(t, "bf16", rec.Precision)
 }
 
 func TestPlanPDRecommendation_ExpertOverride(t *testing.T) {
 	// Override expert count from config.json
-	rec, err := PlanPDRecommendation("deepseek-v3", 671, 512, 16, "fp8", 7168, 61)
+	rec, err := PlanPDRecommendation("deepseek-v3", 671, 512, 16, "fp8", 7168, 61, 0)
 	require.NoError(t, err)
 	require.Equal(t, 512, rec.TotalExperts)
 	require.Equal(t, 16, rec.ActiveExperts)
 }
 
 func TestPlanPDRecommendation_DPField(t *testing.T) {
-	rec, err := PlanPDRecommendation("deepseek-v3", 671, 256, 8, "fp8", 7168, 61)
+	rec, err := PlanPDRecommendation("deepseek-v3", 671, 256, 8, "fp8", 7168, 61, 0)
 	require.NoError(t, err)
 	// DP should default to 1 for all recommendations
 	require.Equal(t, 1, rec.Prefill.DP)

common/types/pd_recommendation.go

@@ -75,7 +75,7 @@ type PDRoleConfig struct {
 	// PDConfig.PrefillReplicas/DecodeReplicas (default 1), which HPA scales up/down.
 	Pods int `json:"pods"`
 	// TotalVRAMGB is the total VRAM required for this role, computed as
-	// MinInferenceVRAMGB * TotalGPUs. Used for VRAM validation and hardware splitting ratio.
+	// MinInferenceVRAMGB * Pods. Used for VRAM validation and hardware splitting ratio.
 	TotalVRAMGB float64 `json:"total_vram_gb"`
 }
 
@@ -97,7 +97,7 @@ type PDRecommendation struct {
 	// Precision is the inference precision.
 	Precision string `json:"precision"`
 	// MinInferenceVRAMGB is the minimum VRAM per GPU required to load and run inference.
-	// TotalVRAMGB for each role is computed as MinInferenceVRAMGB * TotalGPUs.
+	// TotalVRAMGB for each role is computed as MinInferenceVRAMGB * Pods.
 	MinInferenceVRAMGB float64 `json:"min_inference_vram_gb"`
 	// Prefill is the recommended prefill configuration.
 	Prefill PDRoleConfig `json:"prefill"`

configs/pd/models.toml

@@ -20,7 +20,7 @@
 #   pods        - number of pods per LWS group (LWS Size)
 #
 # total_vram_gb is NOT stored in TOML — it is computed in code as:
-#   total_vram_gb = min_inference_vram_gb * total_gpus
+#   total_vram_gb = min_inference_vram_gb * pods
 #
 # ─── TP/EP/DP Convention ───
 # models.toml stores TP and EP with the SAME value (n), DP=1.
@@ -36,13 +36,13 @@
 # ─── DeepSeek ───
 
 [[models]]
-model_name = "deepseek-ai/DeepSeek-V3"
-total_params_b = 671.0
-non_moe_params_b = 37.0
+model_name = "deepseek-ai/DeepSeek-V4-Flash-DSpark"
+total_params_b = 284.0
+non_moe_params_b = 13.0
 total_experts = 256
-active_experts = 8
+active_experts = 6
 precision = "fp8"
-min_inference_vram_gb = 772.0
+min_inference_vram_gb = 327.0
 
 [models.prefill]
 tp = 8
@@ -52,11 +52,11 @@ total_gpus = 8
 pods = 1
 
 [models.decode]
-tp = 16
-ep = 16
+tp = 8
+ep = 8
 dp = 1
-total_gpus = 16
-pods = 2
+total_gpus = 8
+pods = 1
 
 [[models]]
 model_name = "deepseek-ai/DeepSeek-V4-Flash"
@@ -81,6 +81,30 @@ dp = 1
 total_gpus = 8
 pods = 1
 
+[[models]]
+model_name = "deepseek-ai/DeepSeek-V4-Pro-DSpark"
+total_params_b = 1600.0
+non_moe_params_b = 49.0
+total_experts = 384
+active_experts = 6
+precision = "fp8"
+min_inference_vram_gb = 1840.0
+
+[models.prefill]
+tp = 16
+ep = 16
+dp = 1
+total_gpus = 16
+pods = 2
+
+[models.decode]
+tp = 32
+ep = 32
+dp = 1
+total_gpus = 32
+pods = 4
+
+
 [[models]]
 model_name = "deepseek-ai/DeepSeek-V4-Pro"
 total_params_b = 1600.0
@@ -233,18 +257,18 @@ precision = "bf16"
 min_inference_vram_gb = 33.0
 
 [models.prefill]
-tp = 2
-ep = 2
+tp = 1
+ep = 1
 dp = 1
-total_gpus = 2
-pods = 2
+total_gpus = 1
+pods = 1
 
 [models.decode]
-tp = 2
-ep = 2
+tp = 1
+ep = 1
 dp = 1
-total_gpus = 2
-pods = 2
+total_gpus = 1
+pods = 1
 
 # ─── Kimi (Moonshot) ───