[Presets] Add PD Disaggregation Support

[Bihan]

Conducted tests with

1. Multi-Node PD, and
2. Non PD preset

Multi-Node PD
Fleet Used

type: fleet
name: pd-disagg

placement: cluster

ssh_config:
  user: bihan
  identity_file: ~/.ssh/id_rsa
  hosts:
    - 195.242.28.146 # cpu
    - 89.169.113.25 # L40S
    - 195.242.31.145 # L40S
    - 89.169.124.170 # L40S
    - 89.169.113.186 # L40S

Backend: Nebius

Preset Used

type: preset
name: pd-preset-internode-l40s4

repo: meta-llama/Llama-3.2-3B-Instruct


fleets:
  - pd-disagg

trials: 5

min_context_length: 32768
max_ttft: 1000

concurrency: 16
input_tokens: 10000
output_tokens: 512

env:
  - HF_TOKEN

prompt: |
  Rules for every trial:
  - PD disaggregation only, and only INTER-NODE: prefill, decode, and
    router are three different instances. Never two roles on one VM.
  - The router always runs on a dedicated CPU instance (no GPU). Never
    colocate it on a decode (or prefill) GPU node.
  - Use all 4 GPUs across prefill+decode. Do not idle a GPU.
  - For router use CPU node without any GPUs. Do not use GPU node for router.
  - Before choosing that ratio, compute a spec-sheet roofline. Mention 
    in log why the ratio is choosen.

PD disaggregation multi-node results

Trial Summary

trial	layout	change	tok/s	TTFT p50	TTFT cap
1	2P+2D	NIXL/TCP, BF16	219	21.2 s	fail
2	2P+2D	more transfer threads	248	24.4 s	fail
3	2P+2D	mooncake_tcp	561	1.57 s	fail
4	2P+2D	+ FP8 weights	678	1.50 s	fail
5	1P+3D	more decode	396	16.4 s	fail
service	trial 4	dstack PD + round-robin	679	1.52 s	fail

trial-4 was choosen for service. TTFT cap was not meet.

Task To Service conversion
The conversion from trials/4/task.dstack.yml to service/2/service.dstack.yml is correct. The task uses node groups and service uses replica groups. Minor Correction need in prototyping skill: We need to drop startup_order: workers-first from the service YAML — that option only orders master vs workers inside a multi-node task, and on a service it is unused.

Non-PD
NON PD works as expected

Preset Used

type: preset
name: non-pd-preset-01

# The agent picks a compatible variant of the base model
base: meta-llama/Llama-3.2-3B-Instruct

# The number of benchmarked trials
trials: 5

fleets:
  - my-fleet

# The requirements the preset must meet (time to first token is in milliseconds)
min_context_length: 32768
max_ttft: 1000

# The number of simultaneous requests every benchmark uses
concurrency: 4

# The request shape every benchmark uses (defaults to 1024 and 1024)
input_tokens: 10000
output_tokens: 512

# The environment variables the agent may pass to runs
env:
  - HF_TOKEN

Trial Summary

trial	layout	change	tok/s	TTFT p50	TTFT cap
1	1× L40S	vLLM BF16 baseline	203	871 ms	pass
2	1× L40S	vLLM FP8 weights + FP8 KV	304	708 ms	pass
3	1× L40S	vLLM INT4 AWQ + FP8 KV	360	931 ms	pass
4	1× L40S	SGLang INT4 AWQ + FP8 KV	362	946 ms	pass
5	1× L40S	SGLang FP8 + EAGLE spec	307	419 ms	pass
service	trial 4	SGLang INT4 AWQ + FP8 KV	362	946 ms	pass

Task To Service conversion
Worked as expected.

[peterschmidt85]

I pushed changes to skills/dstack/SKILL.md, skills/dstack-prototyping/SKILL.md, and the preset system prompt to make the language around PD disaggregation (node groups / replica groups, replica and job targeting) sharper and more explicit.

[peterschmidt85]

No keep_groups should be needed. Instead, ensure that when a service is serialized, it uses the new groups syntax for replica groups — and not the old one. Then neither this call nor PresetStore.save needs a context flag.

[Bihan]

keep_groups was added because _serialize_legacy_replica_groups dumps the new groups layout as the old replicas/count layout, for CLIs released before groups existed . Preset files are written by that same serializer, so they inherited the old layout too, and the context flag was the way to opt them out.

Unit we remove _serialize_legacy_replica_groups , the alternative I propose is to convert in the preset writers instead.

# store.py
def to_groups_syntax(service: dict[str, Any]) -> dict[str, Any]:
    """ If replicas/count (old layout) exists convert 
          replicas to groups and count to replicas (new layout)"""
    if not isinstance(service.get("replicas"), list):
        return service
    result: dict[str, Any] = {}
    for key, value in service.items():
        if key != "replicas":
            result[key] = value
            continue
        result["groups"] = [
            {("replicas" if k == "count" else k): v for k, v in group.items()}
            for group in value
        ]
    return result

# store.py — PresetStore.save
document = preset.model_dump(mode="json")
document["service"] = to_groups_syntax(document["service"])
content = yaml.safe_dump(document, sort_keys=False)

# export.py
destination.write_text(
    yaml.safe_dump(to_groups_syntax(service.model_dump(mode="json")), sort_keys=False)
)

configurations.py then goes back untouched — no SerializationInfo, no context.

What we can't do is remove _serialize_legacy_replica_groups itself because an old CLI runs dstack apply -f with replicas/count; the new server normalises that into groups on parse via _normalize_legacy_replica_groups, then replies with a RunPlan. Without the serializer that reply carries groups, and responses are parsed with validate_extra_ignore, so the old CLI silently drops it. The old CLI returns a service with neither replicas nor groups. The run becomes wrong.

Plan: deprecate replicas/count in the next release, warning when a user's YAML uses it client-side, then drop _serialize_legacy_replica_groups once pre-groups CLIs are unsupported.

[Bihan]

I am sending a separate PR to resolve this issue. The PR will drop _serialize_legacy_replica_groups and use patching instead.

[Bihan]

This PR #4248 resolves above issue and should be merged first.

[peterschmidt85]

Could you make the PD description in the PR shorter but clearer (if possible)? What to include:

1. The fleet configuration and the preset configuration used.
2. The results of the testing — e.g. both non-PD and PD presets work as expected, at least on a single node:
   • Do you plan to test it on a multi-node fleet without an interconnect?
   • No flaws are detected based on the trial/service records as well as the traces.

Also, please update the docs — remove "Doesn't support PD disaggregation (coming soon)" from mkdocs/docs/concepts/presets.md.