Question: Support different numbers of NPU in Prefill and Decode instance with PD-Disaggregation scenario?

[zuiaikugua]

Thanks for your excellent work!

Currently, I'm studying the PD-disaggregation scenario. I found there is single_node_pd_instance.json in the cluster_configure directory. In this configuration, the prefill instance and decode instance have the same number of NPUs. And the PD-Disaggregation example works!

I tried to simulate the PD-disaggregation with different NPUs in Prefill and Decode instance. For example, there are 4 NPUs in Prefill instance and 8 NPUs in Decode instance.
In my understanding, LLMServingSim adopts the AstraSim simulator as backend. The AstraSim uses the basic topology (Ring, FullConnect, Switch) to compose a multi-dimension topology. In the PD-disaggregation scenario (Prefill has 4 NPUs, Decode has 4 NPUs), LLMServingSim constructs a two-dimension topology as the following.

topology: [FullyConnected, FullyConnected]
npus_count: [4, 2]
bandwidth: [112.0, 112.0]
latency: [0.0, 0.0]

In a two-dimension topology, the first dimension must have the same number of nodes in AstraSim network.

When the Prefill and Decode instance have different numbers of NPUs, how can we configure the AstraSim network? Is it possible to build the network to support this scenario in AstraSim?

[JaehongCS20]

Thanks for your interest and for raising this point.

As you noted, we have been aware of this issue since working on this project. As a current workaround, we are building a larger network and only using the routes that are needed.

If I remember correctly, config_builder.py constructs the configuration by taking the NPU count of the smallest instance as the unit and then scaling the others by multiples of that unit to match the overall topology size. This allows the setup to work with Astra-Sim, but as you pointed out, the collective communication behavior may not be fully accurate in this case.

A proper solution would likely require modifying Astra-Sim itself, and we have not gone that far yet.

If you notice any issues with the PD configuration, please let us know.

[JaehongCS20]

While working on the DP+EP support, I also found that Astra-sim can communicate across specific dimensions using the involved_dim variable. I’ll look into addressing this issue together as part of the DP+EP work.

[zuiaikugua]

In my understanding, the AstraSim also supports the collective communication within same communication group. So we may can classify the different instances NPUs by group id, which should be added in the ET node, to support the PD-disaggregation with non-equal NPUs.

[JaehongCS20]

We were able to get EP+DP working by modifying the code so that collectives can be restricted to a specific dimension using involved_dim.

However, when we started looking into heterogeneous-sized instances, it seems much harder to support in Astra-Sim because the current topology assumptions are symmetric. Because of that, our current idea is to run Astra-Sim with a separate thread per instance, rather than trying to represent all heterogeneous instances inside a single symmetric topology.

This should be reasonable for our use case since, at the moment, there is no communication across instances anyway. In the P/D setup, we simulate the transfer time in the prefill instance for simplicity.

[zuiaikugua]

Do you mean initializing multiple Astra-Sim for multi-instance? If yes, I suspect there might be time issue. In the discrete event simulation, we must make sure there is only one timeline, multiple AstraSim would violate this rule or need to synchronize the timeline among simulators.

[horser1]

Thanks for the previous discussions on this issue!

I see from the earlier conversation that you mentioned potential solutions like using involved_dim to restrict collectives to specific dimensions, and possibly running separate Astra-Sim instances per instance for heterogeneous-sized setups.

I'm wondering if there has been any progress or resolution on this issue since those discussions? Has this been implemented in any recent updates or is there a planned timeline for addressing the heterogeneous PD-disaggregation scenario?

Thanks for your thoughts!

[JaehongCS20]

Hi, thanks for bringing this up.

I’m working on it. I’ll try to make it public around mid-May.

As @zuiaikugua mentioned, there may be sync-related issues here.
I’ll think through how we can address this with minimal overhead

[horser1]

Thanks for your working. There are still a few points I would like to ask for clarification

My understanding is that the communication between prefill and decode instances should not necessarily be modeled as normal collective communication. In PD disaggregation, the main communication path between a prefill instance and a decode instance is more like:

prefill worker computes KV cache
    -> KV cache transfer
    -> decode worker receives KV cache
    -> decode continues generation

In other words, P-D communication seems closer to request-level or layer-wise KV cache migration, rather than TP/EP-style collectives such as all-reduce, all-gather, or all-to-all. Collective communication should mainly happen inside each instance, for example:

Prefill instance internal communication:
  TP=4 topology, modeled by ASTRA-Sim

Decode instance internal communication:
  TP=8 topology, modeled by ASTRA-Sim

Prefill -> Decode communication:
  explicit KV cache transfer event/link

From this perspective, for a configuration such as:

Prefill instance: 4 NPUs
Decode instance: 8 NPUs
Total: 12 NPUs

it seems more natural to build separate intra-instance topologies for the prefill and decode instances, and then model the P->D KV movement as a separate KV transfer event. The memory model or system simulator could then compute the transfer cost based on the corresponding inter-instance bandwidth, latency, and topology.

This also seems consistent with the high-level description in the paper: for PD disaggregation, the Execution Planner can instantiate multiple prefill and decode MSGs; the prefill MSG scheduler inserts layer-wise KV cache transfer operations; and these KV movements are evaluated by the memory model and System Simulator according to network topology and bandwidth.

However, from the current example configuration, it looks like a single global regular ASTRA-Sim topology may need to be constructed first, for example:

"topology": ["FullyConnected", "FullyConnected"],
"npus_count": [4, 2]

This naturally represents two 4-NPU groups, e.g., one 4-NPU prefill instance and one 4-NPU decode instance. But if the prefill instance has 4 NPUs and the decode instance has 8 NPUs, it seems difficult to express this with the same regular 2D topology, because the group size along the first dimension must be the same:

Wanted:
  group0 = 4 NPUs  // prefill
  group1 = 8 NPUs  // decode

But a regular 2D topology seems to represent:
  group0 = 4 NPUs
  group1 = 4 NPUs
  ...

So I would like to ask: does the current implementation of LLMServingSim 2.0 place both intra-instance TP collectives and inter-instance KV transfers into a single global ASTRA-Sim topology / System Simulator backend? If so, does supporting asymmetric PD instance sizes such as P=4 NPUs and D=8 NPUs require additional configuration support or code changes?

Alternatively, is the recommended modeling approach something like the following?

1. Build a 4-NPU internal topology for the prefill MSG;
2. Build an 8-NPU internal topology for the decode MSG;
3. Model P->D KV cache transfer as a separate inter-instance link/event;
4. If the TP degrees differ, additionally model KV resharding / layout conversion cost.

I am not sure whether my understanding of the current implementation is correct, and I would really appreciate any clarification on the intended modeling approach. Thanks!

[JaehongCS20]

Yes, as you pointed out, the original design assumes that all instances are mapped into a single ASTRA-Sim topology.

Although we initially had a more symmetric setting in mind, asymmetric configurations can also be represented to some extent. For example, if we have TP2 and TP4 instances, a current workaround is to use the smallest TP group, i.e., TP2, as the unit size and configure npus_count as something like [2, 3]. This made the simulation run in our experiments, although we have not fully validated this setup yet.

One additional detail is that, even if P and D are configured with the same number of NPUs, for example TP4, the P side is internally instantiated in ASTRA-Sim with twice the given number of NPUs. Therefore, the effective NPU allocation between P and D becomes different inside ASTRA-Sim.

This is because P and D are independent and need to run different static graphs in ASTRA-Sim. For instance, in the TP4 case, ASTRA-Sim is configured with 8 NPUs, which are divided into two 4-NPU groups. The KV transfer latency is then accounted for by modeling one group sending KV cache to the other group layer by layer. After this transfer is completed, the request is moved to the scheduler of the decode instance.

So the KV transfer path itself is not the main issue. The main challenge for heterogeneous configurations is whether collectives work correctly under a topology modeled as [unit_npu_size * unit_count]. This was the motivation for considering running ASTRA-Sim instances in parallel.

[zuiaikugua]

Yes, as you pointed out, the original design assumes that all instances are mapped into a single ASTRA-Sim topology.

Although we initially had a more symmetric setting in mind, asymmetric configurations can also be represented to some extent. For example, if we have TP2 and TP4 instances, a current workaround is to use the smallest TP group, i.e., TP2, as the unit size and configure npus_count as something like [2, 3]. This made the simulation run in our experiments, although we have not fully validated this setup yet.

One additional detail is that, even if P and D are configured with the same number of NPUs, for example TP4, the P side is internally instantiated in ASTRA-Sim with twice the given number of NPUs. Therefore, the effective NPU allocation between P and D becomes different inside ASTRA-Sim.

This is because P and D are independent and need to run different static graphs in ASTRA-Sim. For instance, in the TP4 case, ASTRA-Sim is configured with 8 NPUs, which are divided into two 4-NPU groups. The KV transfer latency is then accounted for by modeling one group sending KV cache to the other group layer by layer. After this transfer is completed, the request is moved to the scheduler of the decode instance.

So the KV transfer path itself is not the main issue. The main challenge for heterogeneous configurations is whether collectives work correctly under a topology modeled as [unit_npu_size * unit_count]. This was the motivation for considering running ASTRA-Sim instances in parallel.

Regarding the TP2 + TP4 example that you mentioned, it's partial "asymmetric". In my understanding, using the topology [2, 3] and involved_dims might not correctly simulate collective communication of TP4, as the involed_dims is 2 for the TP4.
But I suggest that you could try to add a communication group, like group1: [0, 1], group2: [2, 3, 4, 5], to limit the communication scope of TP4.