| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from dataclasses import dataclass, replace |
| 5 | |
| 6 | import numpy as np |
| 7 | import torch |
| 8 | import torch.distributed |
| 9 | |
| 10 | from vllm.distributed.communication_op import tensor_model_parallel_all_gather |
| 11 | from vllm.distributed.parallel_state import get_tp_group |
| 12 | from vllm.triton_utils import tl, triton |
| 13 | from vllm.v1.core.sched.output import GrammarOutput |
| 14 | from vllm.v1.outputs import LogprobsTensors |
| 15 | from vllm.v1.worker.gpu.input_batch import InputBatch |
| 16 | from vllm.v1.worker.gpu.sample.output import SamplerOutput |
| 17 | |
| 18 | |
| 19 | @dataclass |
| 20 | class BatchShardMetadata: |
| 21 | """Collective layout for one sampling step. |
| 22 | |
| 23 | Requests (and their logits) are owner-sorted: all requests owned by |
| 24 | rank 0 first, then rank 1, etc. Within an owner, batch order is |
| 25 | preserved (stable sort). Every field is a pure function of the |
| 26 | replicated idx_mapping and cu_num_logits, so all ranks build identical |
| 27 | plans without communication — which relies on request slots |
| 28 | (idx_mapping) being assigned identically on every rank (see |
| 29 | finish_requests in the model runner). |
| 30 | """ |
| 31 | |
| 32 | tp_size: int |
| 33 | # Per-owner-rank logits counts (all-to-all send splits). |
| 34 | num_logits_per_rank: list[int] |
| 35 | num_local_logits: int |
| 36 | num_local_reqs: int |
| 37 | # The number of request entries contributed by each rank to the |
| 38 | # gathered result (shorter shards are padded up to this value). |
| 39 | max_num_reqs_per_rank: int |
| 40 | # [num_reqs] For batch request i, its source index in the rank-major |
| 41 | # gathered results. Indexing the gathered tensors with this restores |
| 42 | # the original batch order. |
| 43 | gathered_src_indices: torch.Tensor |
| 44 | # The number of entries occupied by each request in the gathered |
| 45 | # tensor. Derived from the global cu_num_logits, so each rank agrees |
| 46 | # ont he gather shapes, including empty shards. |
| 47 | max_num_logits_per_req: int |
| 48 | |
| 49 | |
| 50 | @triton.jit( |
| 51 | do_not_specialize=["num_reqs", "local_logits_start", "max_num_reqs_per_rank"] |
| 52 | ) |
| 53 | def _build_shard_plan_kernel( |
| 54 | idx_mapping_ptr, |
| 55 | cu_num_logits_ptr, |
| 56 | query_start_loc_ptr, |
| 57 | seq_lens_ptr, |
| 58 | sorted_logits_indices_ptr, |
| 59 | gathered_src_indices_ptr, |
| 60 | local_idx_mapping_ptr, |
| 61 | local_cu_num_logits_ptr, |
| 62 | local_logits_indices_ptr, |
| 63 | local_expanded_idx_mapping_ptr, |
| 64 | local_expanded_local_pos_ptr, |
| 65 | local_seq_lens_ptr, |
| 66 | num_reqs, |
| 67 | local_logits_start, |
| 68 | max_num_reqs_per_rank, |
| 69 | TP_SIZE: tl.constexpr, |
| 70 | TP_RANK: tl.constexpr, |
| 71 | PADDED_NUM_REQS: tl.constexpr, |
| 72 | PADDED_NUM_LOGITS_PER_REQ: tl.constexpr, |
| 73 | ): |
| 74 | req_idx = tl.program_id(0) |
| 75 | req_state_idx = tl.load(idx_mapping_ptr + req_idx) |
| 76 | owner = req_state_idx % TP_SIZE |
| 77 | num_logits = tl.load(cu_num_logits_ptr + req_idx + 1) - tl.load( |
| 78 | cu_num_logits_ptr + req_idx |
| 79 | ) |
| 80 | |
| 81 | req_block = tl.arange(0, PADDED_NUM_REQS) |
| 82 | req_mask = req_block < num_reqs |
| 83 | owners = tl.load(idx_mapping_ptr + req_block, mask=req_mask, other=0) % TP_SIZE |
| 84 | num_logits_per_req = tl.load( |
| 85 | cu_num_logits_ptr + req_block + 1, mask=req_mask, other=0 |
| 86 | ) - tl.load(cu_num_logits_ptr + req_block, mask=req_mask, other=0) |
| 87 | |
| 88 | # Owner-sorted layout is rank 0's requests in batch order, then rank 1's, |
| 89 | # etc, so a request's starting offset can be derived from counting. |
| 90 | lower = (owners < owner) & req_mask |
| 91 | earlier = (owners == owner) & (req_block < req_idx) & req_mask |
| 92 | |
| 93 | # Compute the logits row offset for the current request in the sorted |
| 94 | # batch. |
| 95 | req_logits_start = tl.sum(tl.where(lower, num_logits_per_req, 0)) + tl.sum( |
| 96 | tl.where(earlier, num_logits_per_req, 0) |
| 97 | ) |
| 98 | |
| 99 | # Copy the current request logit indices into the sorted batch. |
| 100 | query_start = tl.load(query_start_loc_ptr + req_idx + 1) - num_logits |
| 101 | logit_block = tl.arange(0, PADDED_NUM_LOGITS_PER_REQ) |
| 102 | logit_mask = logit_block < num_logits |
| 103 | logits_indices = (query_start + logit_block).to(tl.int64) |
| 104 | tl.store( |
| 105 | sorted_logits_indices_ptr + req_logits_start + logit_block, |
| 106 | logits_indices, |
| 107 | mask=logit_mask, |
| 108 | ) |
| 109 | |
| 110 | # Write the mapping of output tensor request indices => gathered tensor |
| 111 | # request indices. |
| 112 | # Batch request i's result lives in the padded, gathered results at |
| 113 | # rank_of(i) * max_num_reqs_per_rank + offset_within_rank(i). |
| 114 | local_req_idx = tl.sum(earlier.to(tl.int32)) |
| 115 | tl.store( |
| 116 | gathered_src_indices_ptr + req_idx, |
| 117 | (owner * max_num_reqs_per_rank + local_req_idx).to(tl.int64), |
| 118 | ) |
| 119 | |
| 120 | if req_idx == 0: |
| 121 | tl.store(local_cu_num_logits_ptr, 0) |
| 122 | |
| 123 | if owner == TP_RANK: |
| 124 | local_start = req_logits_start - local_logits_start |
| 125 | tl.store(local_idx_mapping_ptr + local_req_idx, req_state_idx) |
| 126 | tl.store(local_cu_num_logits_ptr + local_req_idx + 1, local_start + num_logits) |
| 127 | tl.store( |
| 128 | local_logits_indices_ptr + local_start + logit_block, |
| 129 | logits_indices, |
| 130 | mask=logit_mask, |
| 131 | ) |
| 132 | tl.store( |
| 133 | local_expanded_idx_mapping_ptr + local_start + logit_block, |
| 134 | req_state_idx, |
| 135 | mask=logit_mask, |
| 136 | ) |
| 137 | tl.store( |
| 138 | local_expanded_local_pos_ptr + local_start + logit_block, |
| 139 | logit_block, |
| 140 | mask=logit_mask, |
| 141 | ) |
| 142 | tl.store(local_seq_lens_ptr + local_req_idx, tl.load(seq_lens_ptr + req_idx)) |
| 143 | |
| 144 | |
| 145 | class BatchSharder: |
| 146 | """Shards the sampler inputs across TP ranks along the batch dimension.""" |
| 147 | |
| 148 | def __init__( |
| 149 | self, |
| 150 | max_num_reqs: int, |
| 151 | max_num_logits_per_req: int, |
| 152 | device: torch.device, |
| 153 | ): |
| 154 | tp_group = get_tp_group() |
| 155 | self.tp_rank = tp_group.rank_in_group |
| 156 | self.tp_size = tp_group.world_size |
| 157 | self.device = device |
| 158 | self._padded_num_reqs = triton.next_power_of_2(max_num_reqs) |
| 159 | self._padded_num_logits_per_req = triton.next_power_of_2(max_num_logits_per_req) |
| 160 | self._num_warps = max(1, min(8, self._padded_num_reqs // 128)) |
| 161 | |
| 162 | def shard_sampler_inputs( |
| 163 | self, |
| 164 | input_batch: InputBatch, |
| 165 | grammar_output: GrammarOutput | None, |
| 166 | ) -> tuple[InputBatch, torch.Tensor, GrammarOutput | None, BatchShardMetadata]: |
| 167 | """Owner-sort the batch and build this rank's local sampler inputs. |
| 168 | |
| 169 | Returns the local sub-batch, the owner-sorted logits_indices (gather |
| 170 | the sampling hidden states with these so `compute_logits_local` emits |
| 171 | logits in all-to-all send order), the local grammar output (None if |
| 172 | this rank owns none of the structured-output requests), and the shard |
| 173 | metadata (used for all-gathering the sampler outputs). |
| 174 | """ |
| 175 | tp_rank = self.tp_rank |
| 176 | tp_size = self.tp_size |
| 177 | num_reqs = input_batch.idx_mapping_np.shape[0] |
| 178 | num_logits = int(input_batch.cu_num_logits_np[-1]) |
| 179 | |
| 180 | # Deterministically assign requests to ranks, round-robin over slot |
| 181 | # indices so ownership stays balanced when the slot allocator fills |
| 182 | # low slots first (partial occupancy). |
| 183 | # NOTE: This mirrors the assignment used in _build_shard_plan_kernel. |
| 184 | req_owner_np = input_batch.idx_mapping_np % tp_size |
| 185 | local_req_indices_np = np.flatnonzero(req_owner_np == tp_rank) |
| 186 | local_idx_mapping_np = input_batch.idx_mapping_np[local_req_indices_np] |
| 187 | num_local_reqs = local_req_indices_np.shape[0] |
| 188 | num_reqs_per_rank_np = np.bincount(req_owner_np, minlength=tp_size) |
| 189 | max_num_reqs_per_rank = int(num_reqs_per_rank_np.max()) if num_reqs else 1 |
| 190 | # Derive the number of logits owned by each rank, as well as the |
| 191 | # local rank. |
| 192 | num_logits_per_req_np = np.diff(input_batch.cu_num_logits_np) |
| 193 | num_logits_per_rank_np = np.bincount( |
| 194 | req_owner_np, weights=num_logits_per_req_np, minlength=tp_size |
| 195 | ).astype(np.int64) |
| 196 | num_local_logits = int(num_logits_per_rank_np[tp_rank]) |
| 197 | local_logits_start = int(num_logits_per_rank_np[:tp_rank].sum()) |
| 198 | local_cu_num_logits_np = np.zeros(num_local_reqs + 1, dtype=np.int32) |
| 199 | np.cumsum( |
| 200 | num_logits_per_req_np[local_req_indices_np], out=local_cu_num_logits_np[1:] |
| 201 | ) |
| 202 | max_num_logits_per_req = int(num_logits_per_req_np.max()) if num_reqs else 1 |
| 203 | local_seq_lens_cpu_upper_bound = input_batch.seq_lens_cpu_upper_bound[ |
| 204 | torch.from_numpy(local_req_indices_np) |
| 205 | ] |
| 206 | |
| 207 | # Shard the input batch GPU tensors. |
| 208 | sorted_logits_indices = torch.empty( |
| 209 | num_logits, dtype=torch.int64, device=self.device |
| 210 | ) |
| 211 | gathered_src_indices = torch.empty( |
| 212 | num_reqs, dtype=torch.int64, device=self.device |
| 213 | ) |
| 214 | local_logits_indices = torch.empty( |
| 215 | num_local_logits, dtype=torch.int64, device=self.device |
| 216 | ) |
| 217 | local_idx_mapping = torch.empty( |
| 218 | num_local_reqs, dtype=torch.int32, device=self.device |
| 219 | ) |
| 220 | local_cu_num_logits = torch.empty( |
| 221 | num_local_reqs + 1, dtype=torch.int32, device=self.device |
| 222 | ) |
| 223 | local_expanded_idx_mapping = torch.empty( |
| 224 | num_local_logits, dtype=torch.int32, device=self.device |
| 225 | ) |
| 226 | local_expanded_local_pos = torch.empty( |
| 227 | num_local_logits, dtype=torch.int32, device=self.device |
| 228 | ) |
| 229 | local_seq_lens = torch.empty( |
| 230 | num_local_reqs, dtype=torch.int32, device=self.device |
| 231 | ) |
| 232 | if num_reqs > 0: |
| 233 | _build_shard_plan_kernel[(num_reqs,)]( |
| 234 | input_batch.idx_mapping, |
| 235 | input_batch.cu_num_logits, |
| 236 | input_batch.query_start_loc, |
| 237 | input_batch.seq_lens, |
| 238 | sorted_logits_indices, |
| 239 | gathered_src_indices, |
| 240 | local_idx_mapping, |
| 241 | local_cu_num_logits, |
| 242 | local_logits_indices, |
| 243 | local_expanded_idx_mapping, |
| 244 | local_expanded_local_pos, |
| 245 | local_seq_lens, |
| 246 | num_reqs, |
| 247 | local_logits_start, |
| 248 | max_num_reqs_per_rank, |
| 249 | TP_SIZE=tp_size, |
| 250 | TP_RANK=tp_rank, |
| 251 | PADDED_NUM_REQS=self._padded_num_reqs, |
| 252 | PADDED_NUM_LOGITS_PER_REQ=self._padded_num_logits_per_req, |
| 253 | num_warps=self._num_warps, |
| 254 | ) |
| 255 | |
| 256 | # Compute the local number of draft tokens. |
| 257 | num_draft_tokens_per_req = None |
| 258 | num_draft_tokens = 0 |
| 259 | if input_batch.num_draft_tokens_per_req is not None: |
| 260 | num_draft_tokens_per_req = input_batch.num_draft_tokens_per_req[ |
| 261 | local_req_indices_np |
| 262 | ] |
| 263 | num_draft_tokens = int(num_draft_tokens_per_req.sum()) |
| 264 | |
| 265 | local_req_ids = [input_batch.req_ids[i] for i in local_req_indices_np.tolist()] |
| 266 | local_batch = replace( |
| 267 | input_batch, |
| 268 | req_ids=local_req_ids, |
| 269 | num_reqs=num_local_reqs, |
| 270 | idx_mapping=local_idx_mapping, |
| 271 | idx_mapping_np=local_idx_mapping_np, |
| 272 | expanded_idx_mapping=local_expanded_idx_mapping, |
| 273 | expanded_local_pos=local_expanded_local_pos, |
| 274 | seq_lens=local_seq_lens, |
| 275 | seq_lens_cpu_upper_bound=local_seq_lens_cpu_upper_bound, |
| 276 | logits_indices=local_logits_indices, |
| 277 | cu_num_logits=local_cu_num_logits, |
| 278 | cu_num_logits_np=local_cu_num_logits_np, |
| 279 | num_draft_tokens=num_draft_tokens, |
| 280 | num_draft_tokens_per_req=num_draft_tokens_per_req, |
| 281 | ) |
| 282 | local_grammar_output = None |
| 283 | if grammar_output is not None: |
| 284 | local_grammar_output = _shard_grammar_output( |
| 285 | grammar_output, input_batch, local_req_ids |
| 286 | ) |
| 287 | metadata = BatchShardMetadata( |
| 288 | tp_size=tp_size, |
| 289 | num_logits_per_rank=num_logits_per_rank_np.tolist(), |
| 290 | num_local_logits=num_local_logits, |
| 291 | num_local_reqs=num_local_reqs, |
| 292 | max_num_reqs_per_rank=max_num_reqs_per_rank, |
| 293 | gathered_src_indices=gathered_src_indices, |
| 294 | max_num_logits_per_req=max_num_logits_per_req, |
| 295 | ) |
| 296 | return local_batch, sorted_logits_indices, local_grammar_output, metadata |
| 297 | |
| 298 | |
| 299 | def all_to_all_logits( |
| 300 | logits_shard: torch.Tensor, |
| 301 | metadata: BatchShardMetadata, |
| 302 | ) -> torch.Tensor: |
| 303 | recv = logits_shard.new_empty( |
| 304 | metadata.tp_size * metadata.num_local_logits, logits_shard.shape[-1] |
| 305 | ) |
| 306 | torch.distributed.all_to_all_single( |
| 307 | recv, |
| 308 | logits_shard.contiguous(), |
| 309 | output_split_sizes=[metadata.num_local_logits] * metadata.tp_size, |
| 310 | input_split_sizes=metadata.num_logits_per_rank, |
| 311 | group=get_tp_group().device_group, |
| 312 | ) |
| 313 | shard_width = recv.shape[-1] |
| 314 | return ( |
| 315 | recv.view(metadata.tp_size, metadata.num_local_logits, shard_width) |
| 316 | .permute(1, 0, 2) |
| 317 | .reshape(metadata.num_local_logits, metadata.tp_size * shard_width) |
| 318 | ) |
| 319 | |
| 320 | |
| 321 | def _shard_grammar_output( |
| 322 | grammar_output: GrammarOutput, |
| 323 | input_batch: InputBatch, |
| 324 | local_req_ids: list[str], |
| 325 | ) -> GrammarOutput | None: |
| 326 | owned = set(local_req_ids) |
| 327 | req_id_to_idx = {req_id: i for i, req_id in enumerate(input_batch.req_ids)} |
| 328 | cu_num_logits_np = input_batch.cu_num_logits_np |
| 329 | |
| 330 | local_ids: list[str] = [] |
| 331 | keep_indices: list[int] = [] |
| 332 | cursor = 0 |
| 333 | for req_id in grammar_output.structured_output_request_ids: |
| 334 | req_idx = req_id_to_idx[req_id] |
| 335 | num_req_logits = int(cu_num_logits_np[req_idx + 1] - cu_num_logits_np[req_idx]) |
| 336 | if req_id in owned: |
| 337 | local_ids.append(req_id) |
| 338 | keep_indices.extend(range(cursor, cursor + num_req_logits)) |
| 339 | cursor += num_req_logits |
| 340 | if not local_ids: |
| 341 | return None |
| 342 | # num_acceptable_drafts is ordered with structured_output_request_ids. |
| 343 | num_acceptable = grammar_output.num_acceptable_drafts |
| 344 | if num_acceptable is not None: |
| 345 | owned = set(local_ids) |
| 346 | ids = grammar_output.structured_output_request_ids |
| 347 | num_acceptable = [n for i, n in zip(ids, num_acceptable) if i in owned] |
| 348 | return GrammarOutput( |
| 349 | structured_output_request_ids=local_ids, |
| 350 | grammar_bitmask=grammar_output.grammar_bitmask[keep_indices], |
| 351 | num_acceptable_drafts=num_acceptable, |
| 352 | ) |
| 353 | |
| 354 | |
| 355 | @triton.jit( |
| 356 | do_not_specialize=["max_num_logits_per_req", "num_src_cols", "packed_stride"] |
| 357 | ) |
| 358 | def _pack_sampler_output_kernel( |
| 359 | packed_ptr, |
| 360 | packed_stride, |
| 361 | sampled_token_ids_ptr, |
| 362 | sampled_token_ids_stride, |
| 363 | num_sampled_ptr, |
| 364 | num_rejected_ptr, |
| 365 | num_nans_ptr, |
| 366 | local_cu_num_logits_ptr, |
| 367 | max_num_logits_per_req, |
| 368 | num_src_cols, |
| 369 | BLOCK_SIZE: tl.constexpr, |
| 370 | ): |
| 371 | req_idx = tl.program_id(0) |
| 372 | cols = tl.arange(0, BLOCK_SIZE) |
| 373 | row_ptr = packed_ptr + req_idx * packed_stride |
| 374 | |
| 375 | token_ids = tl.load( |
| 376 | sampled_token_ids_ptr + req_idx * sampled_token_ids_stride + cols, |
| 377 | mask=cols < num_src_cols, |
| 378 | other=0, |
| 379 | ) |
| 380 | tl.store(row_ptr + cols, token_ids, mask=cols < max_num_logits_per_req) |
| 381 | num_sampled = tl.load(num_sampled_ptr + req_idx) |
| 382 | tl.store(row_ptr + max_num_logits_per_req, num_sampled.to(tl.int64)) |
| 383 | num_rejected = tl.load(num_rejected_ptr + req_idx) |
| 384 | tl.store(row_ptr + max_num_logits_per_req + 1, num_rejected.to(tl.int64)) |
| 385 | if num_nans_ptr is not None: |
| 386 | # num_nans is per logits row; reduce it per request. The replicated |
| 387 | # samplers report per-row counts that downstream zips per request, |
| 388 | # so a per-request sum is the well-defined equivalent. |
| 389 | start = tl.load(local_cu_num_logits_ptr + req_idx) |
| 390 | num_req_logits = tl.load(local_cu_num_logits_ptr + req_idx + 1) - start |
| 391 | nans = tl.load(num_nans_ptr + start + cols, mask=cols < num_req_logits, other=0) |
| 392 | tl.store( |
| 393 | row_ptr + max_num_logits_per_req + 2, tl.sum(nans.to(tl.int64), axis=0) |
| 394 | ) |
| 395 | |
| 396 | |
| 397 | @triton.jit(do_not_specialize=["max_num_logits_per_req", "gathered_stride"]) |
| 398 | def _unpack_gathered_output_kernel( |
| 399 | gathered_ptr, |
| 400 | gathered_stride, |
| 401 | gathered_src_indices_ptr, |
| 402 | sampled_token_ids_ptr, |
| 403 | num_sampled_ptr, |
| 404 | num_rejected_ptr, |
| 405 | num_nans_ptr, |
| 406 | max_num_logits_per_req, |
| 407 | BLOCK_SIZE: tl.constexpr, |
| 408 | ): |
| 409 | req_idx = tl.program_id(0) |
| 410 | src = tl.load(gathered_src_indices_ptr + req_idx) |
| 411 | row_ptr = gathered_ptr + src * gathered_stride |
| 412 | cols = tl.arange(0, BLOCK_SIZE) |
| 413 | mask = cols < max_num_logits_per_req |
| 414 | |
| 415 | token_ids = tl.load(row_ptr + cols, mask=mask) |
| 416 | tl.store( |
| 417 | sampled_token_ids_ptr + req_idx * max_num_logits_per_req + cols, |
| 418 | token_ids, |
| 419 | mask=mask, |
| 420 | ) |
| 421 | num_sampled = tl.load(row_ptr + max_num_logits_per_req) |
| 422 | tl.store(num_sampled_ptr + req_idx, num_sampled.to(tl.int32)) |
| 423 | num_rejected = tl.load(row_ptr + max_num_logits_per_req + 1) |
| 424 | tl.store(num_rejected_ptr + req_idx, num_rejected.to(tl.int32)) |
| 425 | if num_nans_ptr is not None: |
| 426 | num_nans = tl.load(row_ptr + max_num_logits_per_req + 2) |
| 427 | tl.store(num_nans_ptr + req_idx, num_nans.to(tl.int32)) |
| 428 | |
| 429 | |
| 430 | @triton.jit( |
| 431 | do_not_specialize=[ |
| 432 | "max_num_logits_per_req", |
| 433 | "num_logprob_cols", |
| 434 | "num_src_cols", |
| 435 | "send_ids_stride", |
| 436 | "send_logprobs_stride", |
| 437 | "src_ids_stride", |
| 438 | "src_logprobs_stride", |
| 439 | ] |
| 440 | ) |
| 441 | def _pack_logprobs_kernel( |
| 442 | send_ids_ptr, |
| 443 | send_ids_stride, |
| 444 | send_logprobs_ptr, |
| 445 | send_logprobs_stride, |
| 446 | src_ids_ptr, |
| 447 | src_ids_stride, |
| 448 | src_logprobs_ptr, |
| 449 | src_logprobs_stride, |
| 450 | selected_token_ranks_ptr, |
| 451 | local_cu_num_logits_ptr, |
| 452 | max_num_logits_per_req, |
| 453 | num_logprob_cols, |
| 454 | num_src_cols, |
| 455 | BLOCK_SIZE: tl.constexpr, |
| 456 | ): |
| 457 | req_idx = tl.program_id(0) |
| 458 | pos = tl.program_id(1) |
| 459 | start = tl.load(local_cu_num_logits_ptr + req_idx) |
| 460 | if pos < tl.load(local_cu_num_logits_ptr + req_idx + 1) - start: |
| 461 | src = start + pos |
| 462 | dst = req_idx * max_num_logits_per_req + pos |
| 463 | cols = tl.arange(0, BLOCK_SIZE) |
| 464 | # A rank whose sub-batch has no draft tokens runs the regular sampler |
| 465 | # while others run the rejection sampler, and their column counts can |
| 466 | # differ (e.g. logprob_token_ids widening); the two masks pad or |
| 467 | # truncate to the agreed width. |
| 468 | src_mask = cols < num_src_cols |
| 469 | dst_mask = cols < num_logprob_cols |
| 470 | ids = tl.load(src_ids_ptr + src * src_ids_stride + cols, mask=src_mask, other=0) |
| 471 | logprobs = tl.load( |
| 472 | src_logprobs_ptr + src * src_logprobs_stride + cols, |
| 473 | mask=src_mask, |
| 474 | other=float("-inf"), |
| 475 | ) |
| 476 | ids_row = send_ids_ptr + dst * send_ids_stride |
| 477 | tl.store(ids_row + cols, ids.to(tl.int64), mask=dst_mask) |
| 478 | rank = tl.load(selected_token_ranks_ptr + src) |
| 479 | tl.store(ids_row + num_logprob_cols, rank.to(tl.int64)) |
| 480 | tl.store( |
| 481 | send_logprobs_ptr + dst * send_logprobs_stride + cols, |
| 482 | logprobs.to(tl.float32), |
| 483 | mask=dst_mask, |
| 484 | ) |
| 485 | |
| 486 | |
| 487 | @triton.jit( |
| 488 | do_not_specialize=[ |
| 489 | "max_num_logits_per_req", |
| 490 | "num_logprob_cols", |
| 491 | "gathered_ids_stride", |
| 492 | "gathered_logprobs_stride", |
| 493 | ] |
| 494 | ) |
| 495 | def _unpack_logprobs_kernel( |
| 496 | gathered_ids_ptr, |
| 497 | gathered_ids_stride, |
| 498 | gathered_logprobs_ptr, |
| 499 | gathered_logprobs_stride, |
| 500 | gathered_src_indices_ptr, |
| 501 | cu_num_logits_ptr, |
| 502 | logprob_token_ids_ptr, |
| 503 | logprobs_ptr, |
| 504 | selected_token_ranks_ptr, |
| 505 | max_num_logits_per_req, |
| 506 | num_logprob_cols, |
| 507 | BLOCK_SIZE: tl.constexpr, |
| 508 | ): |
| 509 | req_idx = tl.program_id(0) |
| 510 | pos = tl.program_id(1) |
| 511 | start = tl.load(cu_num_logits_ptr + req_idx) |
| 512 | if pos < tl.load(cu_num_logits_ptr + req_idx + 1) - start: |
| 513 | src = tl.load(gathered_src_indices_ptr + req_idx) * max_num_logits_per_req + pos |
| 514 | dst = start + pos |
| 515 | cols = tl.arange(0, BLOCK_SIZE) |
| 516 | mask = cols < num_logprob_cols |
| 517 | src_ids_row = gathered_ids_ptr + src * gathered_ids_stride |
| 518 | ids = tl.load(src_ids_row + cols, mask=mask) |
| 519 | tl.store(logprob_token_ids_ptr + dst * num_logprob_cols + cols, ids, mask=mask) |
| 520 | rank = tl.load(src_ids_row + num_logprob_cols) |
| 521 | tl.store(selected_token_ranks_ptr + dst, rank) |
| 522 | logprobs = tl.load( |
| 523 | gathered_logprobs_ptr + src * gathered_logprobs_stride + cols, mask=mask |
| 524 | ) |
| 525 | tl.store(logprobs_ptr + dst * num_logprob_cols + cols, logprobs, mask=mask) |
| 526 | |
| 527 | |
| 528 | def _gather_logprobs_tensors( |
| 529 | local_output: SamplerOutput | None, |
| 530 | metadata: BatchShardMetadata, |
| 531 | global_batch: InputBatch, |
| 532 | local_batch: InputBatch, |
| 533 | logprobs_dims: tuple[int, int], |
| 534 | device: torch.device, |
| 535 | ) -> LogprobsTensors: |
| 536 | num_logprobs, max_token_ids = logprobs_dims |
| 537 | num_logprob_cols = 1 + max(num_logprobs, max_token_ids) |
| 538 | max_num_logits_per_req = metadata.max_num_logits_per_req |
| 539 | num_send_rows = metadata.max_num_reqs_per_rank * max_num_logits_per_req |
| 540 | block_size = triton.next_power_of_2(num_logprob_cols) |
| 541 | |
| 542 | # Padded slots are never read back: gathered_src_indices only addresses |
| 543 | # rows that an owner rank wrote. |
| 544 | send_ids_ranks = torch.empty( |
| 545 | num_send_rows, num_logprob_cols + 1, dtype=torch.int64, device=device |
| 546 | ) |
| 547 | send_logprobs = torch.empty( |
| 548 | num_send_rows, num_logprob_cols, dtype=torch.float32, device=device |
| 549 | ) |
| 550 | if local_output is not None and local_output.logprobs_tensors is not None: |
| 551 | lp = local_output.logprobs_tensors |
| 552 | _pack_logprobs_kernel[(metadata.num_local_reqs, max_num_logits_per_req)]( |
| 553 | send_ids_ranks, |
| 554 | send_ids_ranks.stride(0), |
| 555 | send_logprobs, |
| 556 | send_logprobs.stride(0), |
| 557 | lp.logprob_token_ids, |
| 558 | lp.logprob_token_ids.stride(0), |
| 559 | lp.logprobs, |
| 560 | lp.logprobs.stride(0), |
| 561 | lp.selected_token_ranks, |
| 562 | local_batch.cu_num_logits, |
| 563 | max_num_logits_per_req, |
| 564 | num_logprob_cols, |
| 565 | lp.logprob_token_ids.shape[1], |
| 566 | BLOCK_SIZE=block_size, |
| 567 | ) |
| 568 | |
| 569 | gathered_ids_ranks = tensor_model_parallel_all_gather(send_ids_ranks, dim=0) |
| 570 | gathered_logprobs = tensor_model_parallel_all_gather(send_logprobs, dim=0) |
| 571 | |
| 572 | num_reqs = global_batch.num_reqs |
| 573 | num_logits = int(global_batch.cu_num_logits_np[-1]) |
| 574 | logprob_token_ids = torch.empty( |
| 575 | num_logits, num_logprob_cols, dtype=torch.int64, device=device |
| 576 | ) |
| 577 | logprobs = torch.empty( |
| 578 | num_logits, num_logprob_cols, dtype=torch.float32, device=device |
| 579 | ) |
| 580 | selected_token_ranks = torch.empty(num_logits, dtype=torch.int64, device=device) |
| 581 | _unpack_logprobs_kernel[(num_reqs, max_num_logits_per_req)]( |
| 582 | gathered_ids_ranks, |
| 583 | gathered_ids_ranks.stride(0), |
| 584 | gathered_logprobs, |
| 585 | gathered_logprobs.stride(0), |
| 586 | metadata.gathered_src_indices, |
| 587 | global_batch.cu_num_logits, |
| 588 | logprob_token_ids, |
| 589 | logprobs, |
| 590 | selected_token_ranks, |
| 591 | max_num_logits_per_req, |
| 592 | num_logprob_cols, |
| 593 | BLOCK_SIZE=block_size, |
| 594 | ) |
| 595 | return LogprobsTensors( |
| 596 | logprob_token_ids=logprob_token_ids, |
| 597 | logprobs=logprobs, |
| 598 | selected_token_ranks=selected_token_ranks, |
| 599 | cu_num_generated_tokens=( |
| 600 | global_batch.cu_num_logits_np.tolist() if num_logits != num_reqs else None |
| 601 | ), |
| 602 | ) |
| 603 | |
| 604 | |
| 605 | def gather_sampler_output( |
| 606 | local_output: SamplerOutput | None, |
| 607 | metadata: BatchShardMetadata, |
| 608 | device: torch.device, |
| 609 | global_batch: InputBatch, |
| 610 | local_batch: InputBatch, |
| 611 | gather_num_nans: bool = False, |
| 612 | logprobs_dims: tuple[int, int] | None = None, |
| 613 | ) -> SamplerOutput: |
| 614 | max_num_logits_per_req = metadata.max_num_logits_per_req |
| 615 | num_packed_cols = max_num_logits_per_req + 2 + (1 if gather_num_nans else 0) |
| 616 | block_size = triton.next_power_of_2(max_num_logits_per_req) |
| 617 | |
| 618 | # Pack the sampler output tensors (excluding logprobs) into a single |
| 619 | # tensor for the all-gather. |
| 620 | packed = torch.empty( |
| 621 | metadata.max_num_reqs_per_rank, |
| 622 | num_packed_cols, |
| 623 | dtype=torch.int64, |
| 624 | device=device, |
| 625 | ) |
| 626 | if local_output is not None: |
| 627 | assert not gather_num_nans or local_output.num_nans is not None |
| 628 | num_src_cols = min( |
| 629 | local_output.sampled_token_ids.shape[1], max_num_logits_per_req |
| 630 | ) |
| 631 | _pack_sampler_output_kernel[(metadata.num_local_reqs,)]( |
| 632 | packed, |
| 633 | packed.stride(0), |
| 634 | local_output.sampled_token_ids, |
| 635 | local_output.sampled_token_ids.stride(0), |
| 636 | local_output.num_sampled, |
| 637 | local_output.num_rejected, |
| 638 | local_output.num_nans if gather_num_nans else None, |
| 639 | local_batch.cu_num_logits if gather_num_nans else None, |
| 640 | max_num_logits_per_req, |
| 641 | num_src_cols, |
| 642 | BLOCK_SIZE=block_size, |
| 643 | ) |
| 644 | |
| 645 | # All-gather the sampler output tensors (excluding logprobs). |
| 646 | gathered = tensor_model_parallel_all_gather(packed, dim=0) |
| 647 | # All-gather the logprobs tensors. |
| 648 | logprobs_tensors = None |
| 649 | if logprobs_dims is not None: |
| 650 | logprobs_tensors = _gather_logprobs_tensors( |
| 651 | local_output, metadata, global_batch, local_batch, logprobs_dims, device |
| 652 | ) |
| 653 | |
| 654 | # Unpack the gathered tensor into a sampler output object. |
| 655 | num_reqs = metadata.gathered_src_indices.shape[0] |
| 656 | sampled_token_ids = torch.empty( |
| 657 | num_reqs, max_num_logits_per_req, dtype=torch.int64, device=device |
| 658 | ) |
| 659 | num_sampled = torch.empty(num_reqs, dtype=torch.int32, device=device) |
| 660 | num_rejected = torch.empty(num_reqs, dtype=torch.int32, device=device) |
| 661 | num_nans = ( |
| 662 | torch.empty(num_reqs, dtype=torch.int32, device=device) |
| 663 | if gather_num_nans |
| 664 | else None |
| 665 | ) |
| 666 | _unpack_gathered_output_kernel[(num_reqs,)]( |
| 667 | gathered, |
| 668 | gathered.stride(0), |
| 669 | metadata.gathered_src_indices, |
| 670 | sampled_token_ids, |
| 671 | num_sampled, |
| 672 | num_rejected, |
| 673 | num_nans, |
| 674 | max_num_logits_per_req, |
| 675 | BLOCK_SIZE=block_size, |
| 676 | ) |
| 677 | return SamplerOutput( |
| 678 | sampled_token_ids=sampled_token_ids, |
| 679 | logprobs_tensors=logprobs_tensors, |
| 680 | num_nans=num_nans, |
| 681 | num_sampled=num_sampled, |
| 682 | num_rejected=num_rejected, |
| 683 | ) |