diff --git a/common/types/pd_planner.go b/common/types/pd_planner.go index 7f0714b1..e924ee03 100644 --- a/common/types/pd_planner.go +++ b/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,6 +398,14 @@ 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, @@ -403,16 +413,15 @@ func PlanPDRecommendation(modelName string, totalParamsB float64, totalExperts, 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), } } diff --git a/common/types/pd_planner_test.go b/common/types/pd_planner_test.go index 81af3f72..bf0f71e4 100644 --- a/common/types/pd_planner_test.go +++ b/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,13 +344,13 @@ 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") } @@ -329,21 +358,21 @@ func TestPlanPDRecommendation_NegativeParams(t *testing.T) { 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) diff --git a/common/types/pd_recommendation.go b/common/types/pd_recommendation.go index 48b01ffd..397453fd 100644 --- a/common/types/pd_recommendation.go +++ b/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"` diff --git a/configs/pd/models.toml b/configs/pd/models.toml index 4d3a81b8..e7a2ad72 100644 --- a/configs/pd/models.toml +++ b/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) ───