Skip to content

vllm.distributed.aux_output_connector.connector

Scheduler-side control plane for execution auxiliary outputs.

Classes:

AuxOutputSchedulerConnector

Build worker metadata without owning auxiliary output payloads or stores.

Methods:

  • build_connector_meta –

    Build one step's incremental worker metadata.

  • request_finished –

    Queue a request's terminal event and final block hashes.

  • reset –

    Start a new auxiliary output namespace after a prefix-cache reset.

  • take_output –

    Return the accepted R3 rows for one scheduled request.

Source code in vllm/distributed/aux_output_connector/connector.py 
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
class AuxOutputSchedulerConnector:
    """Build worker metadata without owning auxiliary output payloads or stores."""

    def __init__(self) -> None:
        # Number of hashes already sent to the worker for each active request.
        self._sent_hash_counts: dict[str, int] = {}
        # Terminal events are delivered with the next connector metadata.
        self._finished_requests: dict[str, PackedBlockHashes | None] = {}
        self._generation = 0

    def build_connector_meta(
        self,
        scheduler_output: SchedulerOutput,
        requests: dict[str, Request],
    ) -> AuxOutputConnectorMetadata:
        """Build one step's incremental worker metadata."""
        scheduled_requests: dict[str, int] = {}
        block_hashes_by_request: dict[str, PackedBlockHashes] = {}
        for request_id in scheduler_output.num_scheduled_tokens:
            num_sent = self._sent_hash_counts.setdefault(request_id, 0)
            request = requests[request_id]
            packed = self._pack_new_hashes(request.block_hashes, num_sent)
            if packed is not None:
                block_hashes_by_request[request_id] = packed
            self._sent_hash_counts[request_id] = len(request.block_hashes)
            assert request.sampling_params is not None
            scheduled_requests[request_id] = max(
                request.sampling_params.routed_experts_prompt_start,
                0 if request.num_output_tokens == 0 else request.num_tokens - 1,
            )
        # A settled token can complete a hash block after the next async schedule
        # was built. Send a hash-only update if the request was not rescheduled.
        if len(self._sent_hash_counts) != len(scheduled_requests):
            for request_id, num_sent in self._sent_hash_counts.items():
                if request_id in scheduled_requests:
                    continue
                request = requests[request_id]
                packed = self._pack_new_hashes(request.block_hashes, num_sent)
                if packed is None:
                    continue
                block_hashes_by_request[request_id] = packed
                self._sent_hash_counts[request_id] = len(request.block_hashes)
        # Sending transfers ownership of these one-shot events.
        finished_requests = tuple(self._finished_requests)
        block_hashes_by_request.update(
            (request_id, block_hashes)
            for request_id, block_hashes in self._finished_requests.items()
            if block_hashes is not None
        )
        self._finished_requests = {}
        return AuxOutputConnectorMetadata(
            self._generation,
            scheduled_requests,
            block_hashes_by_request,
            finished_requests,
        )

    def take_output(
        self,
        request: Request,
        output: dict[str, AuxRequestOutput] | None,
    ) -> np.ndarray | None:
        """Return the accepted R3 rows for one scheduled request."""
        request_id = request.request_id
        assert output is not None and request_id in output, (
            f"auxiliary output worker output is missing {request_id}"
        )
        request_output = output[request_id]
        token_end = request.num_tokens - 1
        local_end = token_end - request_output.token_start
        if local_end < 0:
            assert not request.is_finished(), (
                "finished auxiliary output output has no accepted token range: "
                f"request={request_id}, token_end={token_end}, "
                f"output_start={request_output.token_start}, "
                "output_end="
                f"{request_output.token_start + len(request_output.rows)}"
            )
            return None
        assert local_end <= len(request_output.rows), (
            "auxiliary output worker output has an invalid token range: "
            f"request={request_id}, token_end={token_end}, "
            f"output_start={request_output.token_start}, "
            f"output_end={request_output.token_start + len(request_output.rows)}"
        )
        return request_output.rows[:local_end]

    def request_finished(self, request: Request) -> None:
        """Queue a request's terminal event and final block hashes."""
        request_id = request.request_id
        num_sent = self._sent_hash_counts.pop(request_id, None)
        if num_sent is None:
            return
        # The next metadata delivers this terminal event to the worker.
        self._finished_requests[request_id] = self._pack_new_hashes(
            request.block_hashes, num_sent
        )

    @staticmethod
    def _pack_new_hashes(
        block_hashes: Sequence[bytes], num_sent: int
    ) -> PackedBlockHashes | None:
        assert num_sent <= len(block_hashes), "KV block-hash history shrank"
        new_hashes = block_hashes[num_sent:]
        if not new_hashes:
            return None
        return PackedBlockHashes(b"".join(new_hashes), len(new_hashes[0]))

    def reset(self) -> None:
        """Start a new auxiliary output namespace after a prefix-cache reset."""
        # The worker drops temporary state on generation changes; resend hashes.
        self._sent_hash_counts.clear()
        self._finished_requests.clear()
        self._generation += 1

build_connector_meta(scheduler_output, requests)

Build one step's incremental worker metadata.

Source code in vllm/distributed/aux_output_connector/connector.py 
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
def build_connector_meta(
    self,
    scheduler_output: SchedulerOutput,
    requests: dict[str, Request],
) -> AuxOutputConnectorMetadata:
    """Build one step's incremental worker metadata."""
    scheduled_requests: dict[str, int] = {}
    block_hashes_by_request: dict[str, PackedBlockHashes] = {}
    for request_id in scheduler_output.num_scheduled_tokens:
        num_sent = self._sent_hash_counts.setdefault(request_id, 0)
        request = requests[request_id]
        packed = self._pack_new_hashes(request.block_hashes, num_sent)
        if packed is not None:
            block_hashes_by_request[request_id] = packed
        self._sent_hash_counts[request_id] = len(request.block_hashes)
        assert request.sampling_params is not None
        scheduled_requests[request_id] = max(
            request.sampling_params.routed_experts_prompt_start,
            0 if request.num_output_tokens == 0 else request.num_tokens - 1,
        )
    # A settled token can complete a hash block after the next async schedule
    # was built. Send a hash-only update if the request was not rescheduled.
    if len(self._sent_hash_counts) != len(scheduled_requests):
        for request_id, num_sent in self._sent_hash_counts.items():
            if request_id in scheduled_requests:
                continue
            request = requests[request_id]
            packed = self._pack_new_hashes(request.block_hashes, num_sent)
            if packed is None:
                continue
            block_hashes_by_request[request_id] = packed
            self._sent_hash_counts[request_id] = len(request.block_hashes)
    # Sending transfers ownership of these one-shot events.
    finished_requests = tuple(self._finished_requests)
    block_hashes_by_request.update(
        (request_id, block_hashes)
        for request_id, block_hashes in self._finished_requests.items()
        if block_hashes is not None
    )
    self._finished_requests = {}
    return AuxOutputConnectorMetadata(
        self._generation,
        scheduled_requests,
        block_hashes_by_request,
        finished_requests,
    )

request_finished(request)

Queue a request's terminal event and final block hashes.

Source code in vllm/distributed/aux_output_connector/connector.py 
131
132
133
134
135
136
137
138
139
140
def request_finished(self, request: Request) -> None:
    """Queue a request's terminal event and final block hashes."""
    request_id = request.request_id
    num_sent = self._sent_hash_counts.pop(request_id, None)
    if num_sent is None:
        return
    # The next metadata delivers this terminal event to the worker.
    self._finished_requests[request_id] = self._pack_new_hashes(
        request.block_hashes, num_sent
    )

reset()

Start a new auxiliary output namespace after a prefix-cache reset.

Source code in vllm/distributed/aux_output_connector/connector.py 
152
153
154
155
156
157
def reset(self) -> None:
    """Start a new auxiliary output namespace after a prefix-cache reset."""
    # The worker drops temporary state on generation changes; resend hashes.
    self._sent_hash_counts.clear()
    self._finished_requests.clear()
    self._generation += 1

take_output(request, output)

Return the accepted R3 rows for one scheduled request.

Source code in vllm/distributed/aux_output_connector/connector.py 
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
def take_output(
    self,
    request: Request,
    output: dict[str, AuxRequestOutput] | None,
) -> np.ndarray | None:
    """Return the accepted R3 rows for one scheduled request."""
    request_id = request.request_id
    assert output is not None and request_id in output, (
        f"auxiliary output worker output is missing {request_id}"
    )
    request_output = output[request_id]
    token_end = request.num_tokens - 1
    local_end = token_end - request_output.token_start
    if local_end < 0:
        assert not request.is_finished(), (
            "finished auxiliary output output has no accepted token range: "
            f"request={request_id}, token_end={token_end}, "
            f"output_start={request_output.token_start}, "
            "output_end="
            f"{request_output.token_start + len(request_output.rows)}"
        )
        return None
    assert local_end <= len(request_output.rows), (
        "auxiliary output worker output has an invalid token range: "
        f"request={request_id}, token_end={token_end}, "
        f"output_start={request_output.token_start}, "
        f"output_end={request_output.token_start + len(request_output.rows)}"
    )
    return request_output.rows[:local_end]

PackedBlockHashes dataclass

Contiguous, self-contained block hashes for scheduler-worker IPC.

Source code in vllm/distributed/aux_output_connector/connector.py 
18
19
20
21
22
23
24
25
26
27
@dataclass
class PackedBlockHashes:
    """Contiguous, self-contained block hashes for scheduler-worker IPC."""

    data: bytes
    item_size: int

    def __iter__(self) -> Iterator[bytes]:
        for start in range(0, len(self.data), self.item_size):
            yield self.data[start : start + self.item_size]