code/models/common/modules/lazy_buffer.py

# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0

"""
LazyBuffer: Lazy device buffer allocation for mutable state tensors.

This module has NO torch dependency - it accepts any tensor-like object that
ttnn.from_torch() can handle (duck typing with string type hints).

Design principles:
- Mirrors LazyWeight's allocation contract (source + from_torch() parameters)
- Designed for buffers that are mutated in-place after allocation, NOT immutable model weights
- No disk caching: mutable buffers would corrupt state across instances
- No fingerprinting: without caching, there is no cache to invalidate
- Explicit parameters over hidden closures (IDE-friendly)
- Duck typing for source tensors (no torch import)

See also: LazyWeight in models/common/modules/lazy_weight.py
"""

from dataclasses import dataclass, field, replace
from typing import Optional

import ttnn


@dataclass
class LazyBuffer:
    """
    Lazy-allocated device buffer for mutable state tensors.

    Mirrors LazyWeight's allocation contract (source + from_torch() parameters) but is
    designed for buffers that are mutated in-place after allocation, NOT immutable model weights.

    Key differences from LazyWeight:
    - No disk caching: The device data is overwritten in-place via output_tensor= during
      decode (e.g., penalty masks, token counts). Caching a mutable buffer would cause
      state corruption if loaded by another instance.
    - No fingerprinting: Without caching, there is no cache to invalidate.
    - _value caching is safe: The ttnn.Tensor *handle* returned by get_device_buffer()
      never changes — only the on-device data changes via output_tensor= writes.
      So "allocate once, return same handle" is correct for mutable buffers.

    The only thing that makes these different from weights is post-allocation mutability
    and no disk caching. If a buffer becomes read-only in a future refactor, it can be
    promoted to a LazyWeight with caching enabled.

    See also: LazyWeight in models/common/modules/lazy_weight.py

    Example usage:
        # Fully specified at construction
        buf = LazyBuffer(
            source=torch.zeros(32, 128256, dtype=torch.int32),
            device=mesh_device,
            dtype=ttnn.int32,
            layout=ttnn.TILE_LAYOUT,
            memory_config=ttnn.DRAM_MEMORY_CONFIG,
        )
        tt_tensor = buf.get_device_buffer()  # allocates on first call

        # Refresh device data without reallocation
        buf.update(torch.ones(32, 128256, dtype=torch.int32))

        # Partial construction — resolve later
        buf = LazyBuffer(source=torch.zeros(32, 1))
        buf = resolve_lazy_buffer(buf, device=mesh_device, dtype=ttnn.int32)
    """

    # Source: initial host tensor values (e.g., torch.zeros, torch.ones).
    # Duck-typed — string annotation avoids torch import at module level.
    source: "torch.Tensor"

    # from_torch() parameters — same fields as LazyWeight (minus cache_dir_weight_name, pad_value).
    # Unlike LazyWeight, mesh_mapper stores a pre-built mapper (e.g., ShardTensor2dMesh)
    # rather than a MeshMapperConfig, because LazyBuffer has no caching/fingerprinting.
    dtype: Optional[ttnn.DataType] = ttnn.int32
    layout: Optional[ttnn.Layout] = ttnn.TILE_LAYOUT
    device: Optional[ttnn.MeshDevice] = None
    mesh_mapper: object = None  # Pre-built mapper (ShardTensor2dMesh, etc.) or None for replicate
    memory_config: Optional[ttnn.MemoryConfig] = None

    # Cached device tensor handle (allocated once, device data mutated in-place)
    _value: Optional[ttnn.Tensor] = field(default=None, repr=False)

    def _get_mesh_mapper(self):
        """Get mesh mapper for from_torch(). Shared by get_device_buffer() and update()."""
        if self.mesh_mapper is not None:
            return self.mesh_mapper
        return ttnn.replicate_tensor_to_mesh_mapper(self.device)

    def _from_torch_args(self, *, device):
        """
        Build the full from_torch() kwargs. Used by both get_device_buffer() and update()
        to ensure the same dtype/layout/mesh_mapper/memory_config are used consistently.
        Only ``device`` differs: real device for allocation, None for host-side update.
        """
        return dict(
            dtype=self.dtype,
            layout=self.layout,
            device=device,
            mesh_mapper=self._get_mesh_mapper(),
            memory_config=self.memory_config,
        )

    def get_device_buffer(self) -> ttnn.Tensor:
        """Allocate on first call, return cached handle thereafter."""
        if self._value is not None:
            return self._value

        if self.device is None:
            raise ValueError("device must be set before materializing buffer")
        if self.layout is None:
            raise ValueError("layout must be set before materializing buffer")

        self._value = ttnn.from_torch(
            self.source,
            **self._from_torch_args(device=self.device),
        )
        return self._value

    def update(self, new_source: "torch.Tensor") -> None:
        """
        Overwrite the device buffer contents with a new source tensor, without reallocating.

        If the buffer has not yet been materialized (get_device_buffer not called), this
        simply replaces self.source for future materialization.

        If the buffer IS already materialized, this performs an in-place device update
        using the SAME from_torch() args as the original allocation (dtype, layout,
        mesh_mapper, memory_config) but with device=None to create a host tensor::

            host_tt = ttnn.from_torch(new_source, **same_args, device=None)
            ttnn.copy_host_to_device_tensor(host_tt, self._value)

        The ttnn.Tensor handle (self._value) is preserved — no DRAM reallocation.

        This encapsulates the pattern seen in:
        - TTPenalties._copy_host_to_device (tt_penalties.py:157-159)
        - SeedManager.get_new_values (generator.py:382-383)
        """
        self.source = new_source
        if self._value is not None:
            host_tt = ttnn.from_torch(
                new_source,
                **self._from_torch_args(device=None),
            )
            ttnn.copy_host_to_device_tensor(host_tt, self._value)

    def release(self) -> None:
        """Release the materialized device buffer and allow later reload."""
        if self._value is None:
            return
        value = self._value
        ttnn.deallocate(value)
        self._value = None

    def is_resolved(self) -> bool:
        """Check if all required fields for materialization are set."""
        return self.device is not None and self.dtype is not None and self.layout is not None


def resolve_lazy_buffer(buf: LazyBuffer, **kwargs) -> LazyBuffer:
    """Resolve None fields of ``buf`` with the given kwargs; do not override non-None fields."""
    to_set = {k: v for k, v in kwargs.items() if getattr(buf, k, None) is None}
    return replace(buf, **to_set)