class TransformersDenseTPSP:
"""Install two local fusion patterns, retaining the original forward path."""
def __init__(self, model: nn.Module, vllm_config: "VllmConfig") -> None:
# TODO: Add TPSP support for quantized models.
if vllm_config.quant_config is not None or vllm_config.lora_config is not None:
raise ValueError("TPSP does not support quantization or LoRA")
self.model = model
get_decoder = getattr(model, "get_decoder", None)
decoder = get_decoder() if callable(get_decoder) else model
groups: list[TPSPOpsGroup] = []
rewrites = []
try:
if any(
getattr(getattr(model, "config", None), name, False)
for name in ("output_attentions", "output_hidden_states")
):
raise ValueError("TPSP does not support intermediate HF outputs")
layers_name = _module_name(decoder, _DECODER_LAYER_NAMES)
layers = getattr(decoder, layers_name)
if not isinstance(layers, nn.ModuleList) or not layers:
raise ValueError("TPSP requires a nonempty decoder ModuleList")
final_name = _module_name(decoder, _FINAL_NORM_NAMES)
_validate_decoder(decoder, layers_name, final_name)
input_names = [_module_name(layer, _INPUT_NORM_NAMES) for layer in layers]
norms = [
layer.get_submodule(name) for layer, name in zip(layers, input_names)
]
norms.append(decoder.get_submodule(final_name))
self.input_norm = norms[0]
for index, layer in enumerate(layers):
post_name = _module_name(layer, _POST_ATTENTION_NORM_NAMES)
parents = (
_projection_parent(layer, "o_proj"),
_projection_parent(layer, "down_proj"),
)
rewrites.append(
_rewrite_layer(layer, (input_names[index], post_name), parents)
)
for name, parent, norm in zip(
("o_proj", "down_proj"),
parents,
(layer.get_submodule(post_name), norms[index + 1]),
strict=True,
):
groups.append(
TPSPOpsGroup(
name, layer.get_submodule(f"{parent}.{name}"), norm
)
)
modules = [
self.input_norm,
*(m for g in groups for m in (g.projection, g.norm)),
]
if len({id(module) for module in modules}) != len(modules):
raise ValueError("TPSP does not support shared projections or norms")
norm_types = {
id(norm): get_tpsp_norm_type(norm)
for norm in (self.input_norm, *(group.norm for group in groups))
}
hidden_size = self.input_norm.weight.numel()
for group in groups:
first = groups[0 if group.name == "o_proj" else 1]
projection = group.projection
if (
projection.weight.ndim != 2
or projection.weight.shape != first.projection.weight.shape
or getattr(
projection,
"output_size",
getattr(projection, "out_features", None),
)
!= hidden_size
or group.norm.weight.numel() != hidden_size
or norm_types[id(group.norm)] != norm_types[id(first.norm)]
):
raise ValueError(
f"Incompatible TPSP {group.name} shapes or normalization"
)
except (
AttributeError,
IndexError,
OSError,
SyntaxError,
TypeError,
ValueError,
) as exc:
raise UnsupportedTPSPModelError(
f"unsupported {type(decoder).__name__} fusion region: {exc}"
) from exc
self.groups = groups[:2]
self._projections = [group.projection for group in groups]
self.max_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
self.microchunk_tokens = vllm_config.model_config.tpsp_microchunk_tokens
self.threshold_tokens = vllm_config.model_config.tpsp_threshold_tokens
self.context: TPSPContext | None = None
self.profile_attempted = False
_bind_norm(self.input_norm, groups[0].projection)
for index, group in enumerate(groups):
_bind_projection(group, norm_types[id(group.norm)])
next_projection = (
groups[index + 1].projection if index + 1 < len(groups) else None
)
_bind_norm(group.norm, next_projection)
for layer, rewritten in zip(layers, rewrites, strict=True):
_bind(layer, MethodType(rewritten, layer))
from vllm.distributed.parallel_state import get_tp_group
from vllm.platforms import current_platform
self.backend: TPSPBackend | None = None
backend_cls = current_platform.get_tpsp_backend_cls()
if backend_cls is None:
logger.warning("Transformers TPSP using regular forward: no backend")
return
tp_group = get_tp_group()
device = torch.device(vllm_config.device_config.device)
if device.type != "cpu" and device.index is None:
device = torch.device(device.type, torch.accelerator.current_device_index())
backend = backend_cls(
tp_group.device_group.group_name,
device,
)
try:
if backend.open(
dtype=vllm_config.model_config.dtype,
tp_size=tp_group.world_size,
hidden_size=hidden_size,
max_batched_tokens=self.max_batched_tokens,
group_name=tp_group.device_group.group_name,
device=device,
):
self.backend = backend
finally:
if self.backend is None:
backend.close()
@torch.inference_mode()
def profile(self) -> None:
if self.profile_attempted:
return
self.profile_attempted = True
backend = self.backend
if backend is None:
logger.info("TPSP profile skipped: backend unavailable, success=False")
return
start_time = time.perf_counter()
logger.info(
"TPSP profile started: microchunk_tokens=%s, threshold_tokens=%s, "
"max_batched_tokens=%s",
self.microchunk_tokens if self.microchunk_tokens is not None else "default",
self.threshold_tokens if self.threshold_tokens is not None else "default",
self.max_batched_tokens,
)
status = "failed"
try:
self.model.eval()
# Weights are now loaded: prepare every layer's backend-specific layout
# outside forward; profiling below only uses two representative groups.
backend.prepare_weights(self._projections)
self.context = backend.profile(
self.groups,
self.max_batched_tokens,
microchunk_tokens=self.microchunk_tokens,
threshold_tokens=self.threshold_tokens,
)
status = (
"enabled"
if self.context is not None and self.context.profile.enabled
else "fallback"
)
finally:
try:
if self.context is None:
backend.close()
self.backend = None
finally:
logger.info(
"TPSP profile finished: status=%s, success=%s, elapsed=%.3f s",
status,
status == "enabled",
time.perf_counter() - start_time,
)
@contextmanager
def forward_context(self, num_tokens: int) -> Iterator[None]:
context = self.context
state = None
if context is not None and context.profile.is_active(num_tokens):
state = _ForwardState(context, self.input_norm)
token = _FORWARD_STATE.set(state)
try:
yield
if state is not None and state.expected is not None:
raise RuntimeError("TPSP forward did not consume every fused stage")
finally:
_FORWARD_STATE.reset(token)