Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 17 additions & 8 deletions common/types/pd_planner.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
Expand Down Expand Up @@ -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
Expand All @@ -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),
}
}
59 changes: 44 additions & 15 deletions common/types/pd_planner_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand Down
4 changes: 2 additions & 2 deletions common/types/pd_recommendation.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"`
}

Expand All @@ -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"`
Expand Down
60 changes: 42 additions & 18 deletions configs/pd/models.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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
Expand All @@ -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"
Expand All @@ -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
Expand Down Expand Up @@ -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) ───

Expand Down
Loading