| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Shared Qwen4Exp n-gram embedding storage with device and pinned-host backends. |
| 4 | |
| 5 | Both the NVIDIA and AMD Qwen4Exp implementations use these classes so the (large) |
| 6 | n-gram embedding table can be kept in pinned host memory and looked up through |
| 7 | Unified Virtual Addressing on any CUDA-alike platform. |
| 8 | """ |
| 9 | |
| 10 | import threading |
| 11 | from abc import ABC, abstractmethod |
| 12 | from typing import ClassVar |
| 13 | |
| 14 | import torch |
| 15 | import torch.nn.functional as F |
| 16 | from torch import nn |
| 17 | |
| 18 | from vllm.compilation.breakable_cudagraph import eager_break_during_capture |
| 19 | from vllm.distributed import get_dp_group, get_etp_group, get_tp_group |
| 20 | from vllm.forward_context import DPMetadata, get_forward_context |
| 21 | from vllm.logger import init_logger |
| 22 | from vllm.model_executor.layers.quantization.base_config import ( |
| 23 | QuantizationConfig, |
| 24 | QuantizeMethodBase, |
| 25 | ) |
| 26 | from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501 |
| 27 | CompressedTensorsConfig, |
| 28 | ) |
| 29 | from vllm.model_executor.layers.quantization.fp8 import Fp8Config |
| 30 | from vllm.model_executor.layers.quantization.modelopt import ( |
| 31 | ModelOptMixedPrecisionConfig, |
| 32 | ModelOptQuantConfigBase, |
| 33 | ) |
| 34 | from vllm.model_executor.layers.quantization.utils.fp8_utils import ( |
| 35 | create_fp8_scale_parameter, |
| 36 | ) |
| 37 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 38 | is_layer_skipped, |
| 39 | ) |
| 40 | from vllm.model_executor.parameter import ( |
| 41 | ModelWeightParameter, |
| 42 | PerTensorScaleParameter, |
| 43 | ) |
| 44 | from vllm.model_executor.utils import set_weight_attrs |
| 45 | from vllm.triton_utils import tl, triton |
| 46 | from vllm.utils.platform_utils import is_uva_available |
| 47 | from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor |
| 48 | |
| 49 | from .ple import PLEVocabParallelEmbedding |
| 50 | |
| 51 | logger = init_logger(__name__) |
| 52 | |
| 53 | |
| 54 | class Qwen4ExpPLEEmbedding(PLEVocabParallelEmbedding, ABC): |
| 55 | """ETP-sharded PLE table shared by device and pinned-host backends.""" |
| 56 | |
| 57 | supports_prefetch: ClassVar[bool] = False |
| 58 | |
| 59 | def __init__( |
| 60 | self, |
| 61 | num_embeddings: int, |
| 62 | embedding_dim: int, |
| 63 | *, |
| 64 | params_dtype: torch.dtype, |
| 65 | padding_size: int, |
| 66 | prefix: str, |
| 67 | embedding_method: "Qwen4ExpPLEEmbeddingMethod", |
| 68 | num_ngram_heads: int = 1, |
| 69 | max_total_tokens: int = 0, |
| 70 | data_parallel_rank: int = 0, |
| 71 | ) -> None: |
| 72 | del num_ngram_heads, max_total_tokens |
| 73 | super().__init__( |
| 74 | num_embeddings, |
| 75 | embedding_dim, |
| 76 | params_dtype=params_dtype, |
| 77 | padding_size=padding_size, |
| 78 | prefix=prefix, |
| 79 | quant_method=embedding_method, |
| 80 | parallel_group=get_etp_group(), |
| 81 | ) |
| 82 | self.embedding_method = embedding_method |
| 83 | self.data_parallel_rank = data_parallel_rank |
| 84 | tp_size = get_tp_group().world_size |
| 85 | if self.tp_size % tp_size: |
| 86 | raise ValueError( |
| 87 | "ETP size must be divisible by TP size, but got " |
| 88 | f"ETP={self.tp_size} and TP={tp_size}" |
| 89 | ) |
| 90 | self.etp_data_parallel_size = self.tp_size // tp_size |
| 91 | |
| 92 | @abstractmethod |
| 93 | def allocate_embedding_weight( |
| 94 | self, |
| 95 | num_embeddings: int, |
| 96 | embedding_dim: int, |
| 97 | dtype: torch.dtype, |
| 98 | ) -> torch.Tensor: |
| 99 | """Allocate storage for the complete embedding weight.""" |
| 100 | raise NotImplementedError |
| 101 | |
| 102 | def dequantize( |
| 103 | self, |
| 104 | embeddings: torch.Tensor, |
| 105 | output_dtype: torch.dtype, |
| 106 | ) -> torch.Tensor: |
| 107 | """Delegate storage-format conversion to the embedding method.""" |
| 108 | return self.embedding_method.dequantize(self, embeddings, output_dtype) |
| 109 | |
| 110 | def _get_dp_gather_slot(self, local_num_tokens: int) -> tuple[int, int]: |
| 111 | """Return the per-DP slot size and this rank's slot offset.""" |
| 112 | if self.etp_data_parallel_size == 1: |
| 113 | return local_num_tokens, 0 |
| 114 | dp_metadata: DPMetadata | None = get_forward_context().dp_metadata |
| 115 | if dp_metadata is None: |
| 116 | raise RuntimeError("ETP spanning DP requires DP token metadata") |
| 117 | group_start = (self.data_parallel_rank // self.etp_data_parallel_size) * ( |
| 118 | self.etp_data_parallel_size |
| 119 | ) |
| 120 | group_end = group_start + self.etp_data_parallel_size |
| 121 | token_counts = dp_metadata.num_tokens_across_dp_cpu.tolist() |
| 122 | group_counts = token_counts[group_start:group_end] |
| 123 | slot_size = max(group_counts) |
| 124 | dp_rank = get_dp_group().rank_in_group |
| 125 | return slot_size, dp_rank * slot_size |
| 126 | |
| 127 | def _gather_dp_ids( |
| 128 | self, |
| 129 | ngram_ids: torch.Tensor, |
| 130 | slot_size: int, |
| 131 | ) -> torch.Tensor: |
| 132 | """Gather DP-local IDs that share one ETP-sharded PLE table.""" |
| 133 | if self.etp_data_parallel_size == 1: |
| 134 | return ngram_ids |
| 135 | if ngram_ids.shape[0] < slot_size: |
| 136 | padding = ngram_ids.new_zeros( |
| 137 | slot_size - ngram_ids.shape[0], ngram_ids.shape[1] |
| 138 | ) |
| 139 | ngram_ids = torch.cat((ngram_ids, padding), dim=0) |
| 140 | return get_dp_group().all_gather(ngram_ids, dim=0) |
| 141 | |
| 142 | def _select_embeddings( |
| 143 | self, |
| 144 | embeddings: torch.Tensor, |
| 145 | local_num_tokens: int, |
| 146 | slot_offset: int, |
| 147 | ) -> torch.Tensor: |
| 148 | """Select this DP rank's rows from the ETP-reduced embeddings.""" |
| 149 | if self.etp_data_parallel_size == 1: |
| 150 | return embeddings |
| 151 | return embeddings.narrow(0, slot_offset, local_num_tokens) |
| 152 | |
| 153 | @abstractmethod |
| 154 | def start_prefetch( |
| 155 | self, |
| 156 | hidden_states: torch.Tensor, |
| 157 | ngram_ids: torch.Tensor, |
| 158 | ) -> None: |
| 159 | """Start an asynchronous lookup when supported.""" |
| 160 | raise NotImplementedError |
| 161 | |
| 162 | |
| 163 | class Qwen4ExpPLEEmbeddingMethod(QuantizeMethodBase): |
| 164 | """Quantization interface shared by resident and pinned PLE tables.""" |
| 165 | |
| 166 | # PLE post-load processing only validates scales in their current storage. |
| 167 | requires_device_loading: bool = False |
| 168 | |
| 169 | @staticmethod |
| 170 | def from_quant_config( |
| 171 | quant_config: QuantizationConfig | None, |
| 172 | prefix: str, |
| 173 | embedding_dtype: str | None = None, |
| 174 | ) -> "Qwen4ExpPLEEmbeddingMethod": |
| 175 | """Select the concrete PLE embedding format for a layer.""" |
| 176 | if embedding_dtype == "float8_e4m3fn": |
| 177 | return Qwen4ExpPLEFp8EmbeddingMethod() |
| 178 | if quant_config is None: |
| 179 | return Qwen4ExpPLEUnquantizedEmbeddingMethod() |
| 180 | if isinstance(quant_config, ModelOptMixedPrecisionConfig): |
| 181 | if quant_config._resolve_quant_algo(prefix) == "FP8": |
| 182 | return Qwen4ExpPLEFp8EmbeddingMethod() |
| 183 | return Qwen4ExpPLEUnquantizedEmbeddingMethod() |
| 184 | if isinstance( |
| 185 | quant_config, ModelOptQuantConfigBase |
| 186 | ) and quant_config.is_layer_excluded(prefix): |
| 187 | return Qwen4ExpPLEUnquantizedEmbeddingMethod() |
| 188 | if ( |
| 189 | isinstance(quant_config, CompressedTensorsConfig) |
| 190 | and quant_config.get_scheme_dict(None, layer_name=prefix) is None |
| 191 | ): |
| 192 | return Qwen4ExpPLEUnquantizedEmbeddingMethod() |
| 193 | if not isinstance(quant_config, Fp8Config): |
| 194 | raise NotImplementedError( |
| 195 | "Qwen4Exp PLE embedding does not support quantization config " |
| 196 | f"{type(quant_config).__name__}" |
| 197 | ) |
| 198 | |
| 199 | ignored_layers = quant_config.ignored_layers |
| 200 | if is_layer_skipped( |
| 201 | prefix, |
| 202 | ignored_layers, |
| 203 | quant_config.packed_modules_mapping, |
| 204 | match_mode=quant_config.ignored_layers_match_mode, |
| 205 | ): |
| 206 | return Qwen4ExpPLEUnquantizedEmbeddingMethod() |
| 207 | # PLE checkpoint shards form one runtime embedding parameter. |
| 208 | shard_prefix = f"{prefix}.shard_" |
| 209 | if any(name.startswith(shard_prefix) for name in ignored_layers): |
| 210 | return Qwen4ExpPLEUnquantizedEmbeddingMethod() |
| 211 | if not quant_config.is_checkpoint_fp8_serialized: |
| 212 | raise NotImplementedError( |
| 213 | "Qwen4Exp PLE embedding only supports serialized FP8 checkpoints" |
| 214 | ) |
| 215 | return Qwen4ExpPLEFp8EmbeddingMethod() |
| 216 | |
| 217 | def apply( |
| 218 | self, |
| 219 | layer: nn.Module, |
| 220 | x: torch.Tensor, |
| 221 | bias: torch.Tensor | None = None, |
| 222 | ) -> torch.Tensor: |
| 223 | raise NotImplementedError("PLE weights only support embedding lookup") |
| 224 | |
| 225 | def embedding(self, layer: nn.Module, input_: torch.Tensor) -> torch.Tensor: |
| 226 | return F.embedding(input_, layer.weight) |
| 227 | |
| 228 | @abstractmethod |
| 229 | def dequantize( |
| 230 | self, |
| 231 | layer: nn.Module, |
| 232 | embeddings: torch.Tensor, |
| 233 | output_dtype: torch.dtype, |
| 234 | ) -> torch.Tensor: |
| 235 | """Convert looked-up PLE rows to the activation dtype.""" |
| 236 | raise NotImplementedError |
| 237 | |
| 238 | |
| 239 | class Qwen4ExpPLEUnquantizedEmbeddingMethod(Qwen4ExpPLEEmbeddingMethod): |
| 240 | """Unquantized PLE embedding storage and lookup semantics.""" |
| 241 | |
| 242 | def create_weights( |
| 243 | self, |
| 244 | layer: Qwen4ExpPLEEmbedding, |
| 245 | input_size_per_partition: int, |
| 246 | output_partition_sizes: list[int], |
| 247 | input_size: int, |
| 248 | output_size: int, |
| 249 | params_dtype: torch.dtype, |
| 250 | **extra_weight_attrs, |
| 251 | ) -> None: |
| 252 | del input_size, output_size |
| 253 | weight = nn.Parameter( |
| 254 | layer.allocate_embedding_weight( |
| 255 | sum(output_partition_sizes), |
| 256 | input_size_per_partition, |
| 257 | params_dtype, |
| 258 | ), |
| 259 | requires_grad=False, |
| 260 | ) |
| 261 | set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0}) |
| 262 | set_weight_attrs(weight, extra_weight_attrs) |
| 263 | layer.register_parameter("weight", weight) |
| 264 | |
| 265 | def dequantize( |
| 266 | self, |
| 267 | layer: nn.Module, |
| 268 | embeddings: torch.Tensor, |
| 269 | output_dtype: torch.dtype, |
| 270 | ) -> torch.Tensor: |
| 271 | del layer, output_dtype |
| 272 | return embeddings |
| 273 | |
| 274 | |
| 275 | class Qwen4ExpPLEFp8EmbeddingMethod(Qwen4ExpPLEEmbeddingMethod): |
| 276 | """FP8 PLE embedding with one global checkpoint scale.""" |
| 277 | |
| 278 | def create_weights( |
| 279 | self, |
| 280 | layer: Qwen4ExpPLEEmbedding, |
| 281 | input_size_per_partition: int, |
| 282 | output_partition_sizes: list[int], |
| 283 | input_size: int, |
| 284 | output_size: int, |
| 285 | params_dtype: torch.dtype, |
| 286 | **extra_weight_attrs, |
| 287 | ) -> None: |
| 288 | del input_size, output_size, params_dtype |
| 289 | weight_loader = extra_weight_attrs.get("weight_loader") |
| 290 | weight = ModelWeightParameter( |
| 291 | data=layer.allocate_embedding_weight( |
| 292 | sum(output_partition_sizes), |
| 293 | input_size_per_partition, |
| 294 | torch.float8_e4m3fn, |
| 295 | ), |
| 296 | input_dim=1, |
| 297 | output_dim=0, |
| 298 | weight_loader=weight_loader, |
| 299 | ) |
| 300 | layer.register_parameter("weight", weight) |
| 301 | |
| 302 | weight_scale = create_fp8_scale_parameter( |
| 303 | PerTensorScaleParameter, |
| 304 | output_partition_sizes, |
| 305 | input_size_per_partition, |
| 306 | None, |
| 307 | weight_loader, |
| 308 | scale_dtype=torch.float32, |
| 309 | ) |
| 310 | layer.register_parameter("weight_scale", weight_scale) |
| 311 | |
| 312 | def process_weights_after_loading(self, layer: nn.Module) -> None: |
| 313 | """Reject FP8 PLE checkpoints without a global scale.""" |
| 314 | sentinel = torch.finfo(torch.float32).min |
| 315 | if torch.any(layer.weight_scale == sentinel): |
| 316 | raise ValueError("FP8 PLE checkpoint is missing its global scale") |
| 317 | |
| 318 | def dequantize( |
| 319 | self, |
| 320 | layer: nn.Module, |
| 321 | embeddings: torch.Tensor, |
| 322 | output_dtype: torch.dtype, |
| 323 | ) -> torch.Tensor: |
| 324 | weight_scale = getattr(layer, "weight_scale", None) |
| 325 | if weight_scale is None: |
| 326 | raise RuntimeError("FP8 PLE embedding is missing its global scale") |
| 327 | if weight_scale.device != embeddings.device: |
| 328 | raise RuntimeError("FP8 PLE embedding scale must be on the output device") |
| 329 | return embeddings.to(output_dtype) * weight_scale.to(output_dtype) |
| 330 | |
| 331 | |
| 332 | class Qwen4ExpPLEDeviceEmbedding(Qwen4ExpPLEEmbedding): |
| 333 | """PLE table allocated on the active model device.""" |
| 334 | |
| 335 | def allocate_embedding_weight( |
| 336 | self, |
| 337 | num_embeddings: int, |
| 338 | embedding_dim: int, |
| 339 | dtype: torch.dtype, |
| 340 | ) -> torch.Tensor: |
| 341 | """Allocate the complete PLE weight on the active device.""" |
| 342 | return torch.empty(num_embeddings, embedding_dim, dtype=dtype) |
| 343 | |
| 344 | def start_prefetch( |
| 345 | self, |
| 346 | hidden_states: torch.Tensor, |
| 347 | ngram_ids: torch.Tensor, |
| 348 | ) -> None: |
| 349 | """Resident embedding prefetch is a no-op.""" |
| 350 | return None |
| 351 | |
| 352 | def forward(self, ngram_ids: torch.Tensor) -> torch.Tensor: |
| 353 | """Gather ETP inputs, look up embeddings, and select local rows.""" |
| 354 | slot_size, slot_offset = self._get_dp_gather_slot(ngram_ids.shape[0]) |
| 355 | gathered_ids = self._gather_dp_ids(ngram_ids, slot_size) |
| 356 | embeddings = super().forward(gathered_ids) |
| 357 | return self._select_embeddings( |
| 358 | embeddings, |
| 359 | ngram_ids.shape[0], |
| 360 | slot_offset, |
| 361 | ) |
| 362 | |
| 363 | |
| 364 | @triton.jit |
| 365 | def _lookup_ple_embedding_from_pinned_kernel( |
| 366 | weight_ptr, |
| 367 | ids_ptr, |
| 368 | output_ptr, |
| 369 | embedding_dim, |
| 370 | tp_vocab_start, |
| 371 | tp_vocab_end, |
| 372 | BLOCK_D: tl.constexpr, |
| 373 | ): |
| 374 | """Look up TP-owned PLE rows through a CUDA view of pinned host memory.""" |
| 375 | row_id = tl.program_id(0) |
| 376 | global_idx = tl.load(ids_ptr + row_id) |
| 377 | in_range = (global_idx >= tp_vocab_start) & (global_idx < tp_vocab_end) |
| 378 | local_idx = tl.where(in_range, global_idx - tp_vocab_start, 0) |
| 379 | offsets = tl.arange(0, BLOCK_D) |
| 380 | store_mask = offsets < embedding_dim |
| 381 | load_mask = store_mask & in_range |
| 382 | values = tl.load( |
| 383 | weight_ptr + local_idx * embedding_dim + offsets, |
| 384 | mask=load_mask, |
| 385 | other=0.0, |
| 386 | ) |
| 387 | tl.store( |
| 388 | output_ptr + row_id * embedding_dim + offsets, |
| 389 | values, |
| 390 | mask=store_mask, |
| 391 | ) |
| 392 | |
| 393 | |
| 394 | class Qwen4ExpPLEPinnedHostEmbedding(Qwen4ExpPLEEmbedding): |
| 395 | """PLE table loaded into pinned CPU memory and looked up through UVA.""" |
| 396 | |
| 397 | supports_prefetch: ClassVar[bool] = True |
| 398 | |
| 399 | def __init__( |
| 400 | self, |
| 401 | num_embeddings: int, |
| 402 | embedding_dim: int, |
| 403 | *, |
| 404 | params_dtype: torch.dtype, |
| 405 | padding_size: int, |
| 406 | prefix: str, |
| 407 | embedding_method: Qwen4ExpPLEEmbeddingMethod, |
| 408 | num_ngram_heads: int = 1, |
| 409 | max_total_tokens: int = 0, |
| 410 | data_parallel_rank: int = 0, |
| 411 | ) -> None: |
| 412 | if not is_uva_available(): |
| 413 | raise RuntimeError("Engram CPU offload requires UVA support") |
| 414 | super().__init__( |
| 415 | num_embeddings, |
| 416 | embedding_dim, |
| 417 | params_dtype=params_dtype, |
| 418 | padding_size=padding_size, |
| 419 | prefix=prefix, |
| 420 | embedding_method=embedding_method, |
| 421 | num_ngram_heads=num_ngram_heads, |
| 422 | max_total_tokens=max_total_tokens, |
| 423 | data_parallel_rank=data_parallel_rank, |
| 424 | ) |
| 425 | self._uva_weight = get_accelerator_view_from_cpu_tensor(self.weight) |
| 426 | self._block_d = triton.next_power_of_2(self.embedding_dim) |
| 427 | self._prefetch_stream: torch.cuda.Stream | None = None |
| 428 | self._prefetch_buffer: torch.Tensor | None = None |
| 429 | self._prefetch_alloc_lock = threading.Lock() |
| 430 | self._prefetch_rows = max_total_tokens * self.etp_data_parallel_size |
| 431 | self._num_ngram_heads = num_ngram_heads |
| 432 | self._output_dim = num_ngram_heads * self.embedding_dim |
| 433 | |
| 434 | def allocate_embedding_weight( |
| 435 | self, |
| 436 | num_embeddings: int, |
| 437 | embedding_dim: int, |
| 438 | dtype: torch.dtype, |
| 439 | ) -> torch.Tensor: |
| 440 | """Allocate the complete PLE weight directly in pinned CPU memory.""" |
| 441 | return torch.empty( |
| 442 | num_embeddings, |
| 443 | embedding_dim, |
| 444 | dtype=dtype, |
| 445 | device="cpu", |
| 446 | pin_memory=True, |
| 447 | ) |
| 448 | |
| 449 | def _lookup( |
| 450 | self, |
| 451 | input_ids: torch.Tensor, |
| 452 | output: torch.Tensor | None = None, |
| 453 | ) -> torch.Tensor: |
| 454 | """Look up local ETP rows while preserving the weight storage dtype.""" |
| 455 | expected_shape = (*input_ids.shape, self.embedding_dim) |
| 456 | if output is None: |
| 457 | output = torch.empty( |
| 458 | expected_shape, |
| 459 | dtype=self.weight.dtype, |
| 460 | device=input_ids.device, |
| 461 | ) |
| 462 | elif ( |
| 463 | tuple(output.shape) != expected_shape |
| 464 | or output.dtype != self.weight.dtype |
| 465 | or output.device != input_ids.device |
| 466 | ): |
| 467 | raise ValueError( |
| 468 | "PLE prefetch output must match the input shape, weight dtype, " |
| 469 | "and input device" |
| 470 | ) |
| 471 | |
| 472 | flat_ids = input_ids.reshape(-1).long() |
| 473 | if flat_ids.numel(): |
| 474 | _lookup_ple_embedding_from_pinned_kernel[(flat_ids.numel(),)]( |
| 475 | self._uva_weight, |
| 476 | flat_ids, |
| 477 | output, |
| 478 | self.embedding_dim, |
| 479 | self.shard_indices.org_vocab_start_index, |
| 480 | self.shard_indices.org_vocab_end_index, |
| 481 | BLOCK_D=self._block_d, |
| 482 | ) |
| 483 | return output |
| 484 | |
| 485 | def sync_lookup(self, ngram_ids: torch.Tensor) -> torch.Tensor: |
| 486 | """Synchronous UVA lookup for platforms without prefetch wiring.""" |
| 487 | slot_size, slot_offset = self._get_dp_gather_slot(ngram_ids.shape[0]) |
| 488 | gathered_ids = self._gather_dp_ids(ngram_ids, slot_size) |
| 489 | embeddings = self._lookup(gathered_ids) |
| 490 | embeddings = self._reduce_etp_embeddings(embeddings) |
| 491 | return self._select_embeddings( |
| 492 | embeddings, |
| 493 | ngram_ids.shape[0], |
| 494 | slot_offset, |
| 495 | ) |
| 496 | |
| 497 | def _reduce_etp_embeddings(self, embeddings: torch.Tensor) -> torch.Tensor: |
| 498 | """Combine pinned lookup results owned by different ETP ranks.""" |
| 499 | if self.tp_size == 1: |
| 500 | return embeddings |
| 501 | assert self.parallel_group is not None |
| 502 | if embeddings.dtype in (torch.float8_e4m3fn, torch.float8_e5m2): |
| 503 | # Each vocabulary row has one owner, so reduce the raw FP8 bytes. |
| 504 | reduced = self.parallel_group.all_reduce(embeddings.view(torch.int8)) |
| 505 | return reduced.view(embeddings.dtype) |
| 506 | return self.parallel_group.all_reduce(embeddings) |
| 507 | |
| 508 | @eager_break_during_capture |
| 509 | def start_prefetch( |
| 510 | self, |
| 511 | hidden_states: torch.Tensor, |
| 512 | ngram_ids: torch.Tensor, |
| 513 | ) -> None: |
| 514 | """Gather ETP IDs and launch their UVA lookup on the side stream.""" |
| 515 | buffer = self._prefetch_buffer |
| 516 | if buffer is None: |
| 517 | # First use allocates. The eager profile run always precedes |
| 518 | # cudagraph capture, so allocation never happens mid-capture; |
| 519 | # the lock keeps concurrent first callers from tearing the |
| 520 | # stream/buffer pair. |
| 521 | with self._prefetch_alloc_lock: |
| 522 | buffer = self._prefetch_buffer |
| 523 | if buffer is None: |
| 524 | if torch.cuda.is_current_stream_capturing(): |
| 525 | raise RuntimeError( |
| 526 | "pinned PLE prefetch buffer must be allocated " |
| 527 | "eagerly, before cudagraph capture" |
| 528 | ) |
| 529 | self._prefetch_stream = torch.cuda.Stream( |
| 530 | device=self._uva_weight.device |
| 531 | ) |
| 532 | buffer = torch.empty( |
| 533 | self._prefetch_rows, |
| 534 | self._num_ngram_heads, |
| 535 | self.embedding_dim, |
| 536 | dtype=self.weight.dtype, |
| 537 | device=self._uva_weight.device, |
| 538 | ) |
| 539 | self._prefetch_buffer = buffer |
| 540 | prefetch_stream = self._prefetch_stream |
| 541 | if prefetch_stream is None: |
| 542 | raise RuntimeError("pinned PLE prefetch stream was not allocated") |
| 543 | slot_size, _ = self._get_dp_gather_slot(ngram_ids.shape[0]) |
| 544 | gathered_ids = self._gather_dp_ids(ngram_ids, slot_size) |
| 545 | if gathered_ids.shape[0] > buffer.shape[0]: |
| 546 | raise ValueError( |
| 547 | f"pinned PLE prefetch buffer holds {buffer.shape[0]} rows, " |
| 548 | f"but the batch needs {gathered_ids.shape[0]}" |
| 549 | ) |
| 550 | active_output = buffer[: gathered_ids.shape[0]] |
| 551 | prefetch_stream.wait_stream(torch.cuda.current_stream()) |
| 552 | gathered_ids.record_stream(prefetch_stream) |
| 553 | with torch.cuda.stream(prefetch_stream): |
| 554 | self._lookup(gathered_ids, output=active_output) |
| 555 | |
| 556 | @eager_break_during_capture |
| 557 | def _finalize_prefetch( |
| 558 | self, |
| 559 | prefetch_output: torch.Tensor, |
| 560 | output: torch.Tensor, |
| 561 | ) -> None: |
| 562 | """Join the side stream, reduce ETP shards, and select local rows.""" |
| 563 | prefetch_stream = self._prefetch_stream |
| 564 | if prefetch_stream is None: |
| 565 | raise RuntimeError("pinned PLE finalize requires a prior start_prefetch") |
| 566 | torch.cuda.current_stream().wait_stream(prefetch_stream) |
| 567 | slot_size, slot_offset = self._get_dp_gather_slot(output.shape[0]) |
| 568 | active_output = prefetch_output[: slot_size * self.etp_data_parallel_size] |
| 569 | embeddings = self._reduce_etp_embeddings(active_output) |
| 570 | embeddings = self._select_embeddings( |
| 571 | embeddings, |
| 572 | output.shape[0], |
| 573 | slot_offset, |
| 574 | ) |
| 575 | output.copy_(embeddings.flatten(-2)) |
| 576 | |
| 577 | def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| 578 | """Finish the pinned lookup into graph-owned output storage.""" |
| 579 | buffer = self._prefetch_buffer |
| 580 | if buffer is None: |
| 581 | raise RuntimeError("pinned PLE lookup requires a prior start_prefetch") |
| 582 | output = buffer.new_empty((hidden_states.shape[0], self._output_dim)) |
| 583 | self._finalize_prefetch(buffer, output) |
| 584 | return output |