clef / code /models /tt_transformers /tt /generator_vllm.py
Add files using upload-large-folder tool
b025706 verified 25 days ago
56.1 kB
# SPDX-FileCopyrightText: © 2024 Tenstorrent USA, Inc.

# SPDX-License-Identifier: Apache-2.0

from types import SimpleNamespace
from typing import List, Mapping, Union

import torch
from loguru import logger
from PIL.Image import Image
from tqdm import tqdm
from vllm.model_executor.models.gemma3_mm import (
    Gemma3DummyInputsBuilder,
    Gemma3MultiModalProcessor,
    Gemma3ProcessingInfo,
)
from vllm.model_executor.models.interfaces import SupportsMultiModal
from vllm.model_executor.models.mistral3 import (
    Mistral3DummyInputsBuilder,
    Mistral3MultiModalProcessor,
    Mistral3ProcessingInfo,
)
from vllm.multimodal import MULTIMODAL_REGISTRY
from vllm.multimodal.processing import BaseDummyInputsBuilder

try:
    # vLLM >= 0.24.0 exposes MultiModalDataDict from vllm.inputs; older
    # versions export it from vllm.multimodal.inputs.
    from vllm.inputs import MultiModalDataDict
except ImportError:
    from vllm.multimodal.inputs import MultiModalDataDict

import ttnn
from models.common.llama_models import create_vision_mask
from models.common.utility_functions import is_wormhole_b0, nearest_32
from models.tt_transformers.tt.generator import Generator, create_submeshes
from models.tt_transformers.tt.model import Transformer
from models.tt_transformers.tt.model_config import DecodersPrecision, ModelArgs, TensorGroup


def allocate_vllm_kv_cache_per_layer(per_layer_specs, dp_model: List[Transformer], tt_cache_path):
    """Allocate KV cache tensors with optional cross-layer DRAM sharing.
    Args:
        per_layer_specs: list of ``(kv_cache_shape, dtype, tensor_idx)``
            triples, one per layer in model layer-index order. Layers with
            the same ``tensor_idx`` share one underlying TT tensor — this
            is upstream's HMA tensor-sharing layout (e.g. for Gemma3 5:1,
            one full-attention layer and several sliding-window layers
            collapse to a single DRAM buffer; per-group block tables keep
            their slot accesses disjoint at runtime). Layers with unique
            ``tensor_idx`` get their own buffer.
        dp_model: list of replicated TT model handles, one per data-parallel
            submesh.
        tt_cache_path: path used for on-disk weight cache file naming.
    Returns:
        ``list[submesh][layer_idx][k_or_v]`` of TT tensors. Multiple
        ``layer_idx`` entries may refer to the same underlying tensor
        objects when they share a ``tensor_idx``.
    """
    submesh_devices = [model.mesh_device for model in dp_model]
    kv_cache = []
    for mesh_idx, submesh in enumerate(submesh_devices):
        # tensor_idx -> [k, v] ttnn handles; reused across all layers that
        # share a buffer.
        unique_buffers: dict[int, list] = {}
        kv_tt = []
        for layer_num, (kv_cache_shape, dtype, tensor_idx) in enumerate(
            tqdm(per_layer_specs, desc=f"Allocating TT kv caches for each layer (submesh {mesh_idx+1})")
        ):
            existing = unique_buffers.get(tensor_idx)
            if existing is not None:
                kv_tt.append(existing)
                continue
            cache_kv = torch.zeros(kv_cache_shape, dtype=dtype)
            # Get the dtype for the kv cache based on the configured optimizations in the model
            if dp_model[mesh_idx].args.optimizations is not None:
                kv_cache_dtype = dp_model[mesh_idx].args.optimizations.get_tensor_dtype(
                    decoder_id=layer_num, tensor=TensorGroup.KV_CACHE
                )
            else:
                logger.info("No dtype specified for the model KV cache - defaulting to ttnn.bfloat8_b.")
                kv_cache_dtype = None
            # Set default to bfloat8_b when no optimizations are configured
            kv_cache_dtype = ttnn.bfloat8_b if kv_cache_dtype is None else kv_cache_dtype
            kv_tt_i = [
                ttnn.as_tensor(
                    cache_kv,
                    device=submesh,
                    # TODO: this could be ShardTensorToMesh, removing the need for vLLM to know about TP for num_kv_heads.
                    # Could affect other calculations which use TTCacheEngine.num_kv_heads, though.
                    mesh_mapper=ttnn.ReplicateTensorToMesh(submesh),
                    layout=ttnn.TILE_LAYOUT,
                    memory_config=ttnn.DRAM_MEMORY_CONFIG,
                    dtype=kv_cache_dtype,
                    # Separate cache files for K and V to avoid collision.
                    # ``tensor_idx`` distinguishes shared buffers that have the
                    # same shape but back different layer subsets.
                    cache_file_name=tt_cache_path / f"empty_{kv}cache_paged_attention{kv_cache_shape}_t{tensor_idx}",
                )
                for kv in ["k", "v"]
            ]

            unique_buffers[tensor_idx] = kv_tt_i
            kv_tt.append(kv_tt_i)
        kv_cache.append(kv_tt)
    return kv_cache


def allocate_vllm_kv_cache(kv_cache_shape, dtype, num_layers, dp_model: List[Transformer], tt_cache_path):
    """Uniform-shape KV cache allocator for non-hybrid models.
    Hybrid attention models should use :func:`allocate_vllm_kv_cache_per_layer`,
    which takes a per-layer ``(shape, dtype, tensor_idx)`` list so layers
    can share DRAM buffers per upstream's HMA tensor-sharing model.
    """
    return allocate_vllm_kv_cache_per_layer(
        [(kv_cache_shape, dtype, i) for i in range(num_layers)],
        dp_model=dp_model,
        tt_cache_path=tt_cache_path,
    )


class HybridAttentionForCausalLM(Generator):
    """vLLM wrapper base for hybrid attention models.
    Models with mixed sliding-window + full-attention layers (Gemma3,
    Gemma4, GPT-OSS, ...) inherit from this class instead of plain
    :class:`Generator` so they can opt in to upstream's hybrid kv cache
    manager. The shared ``get_kv_cache_spec`` classmethod here builds
    the per-layer KV cache spec from ``hf_config.text_config.layer_types``
    — the standard HF convention used by all of these models — emitting
    ``SlidingWindowSpec`` for sliding layers and ``FullAttentionSpec``
    for full-attention layers.
    Subclasses are responsible for the model-specific pieces:
    * ``initialize_vllm_model``: load the underlying TT model.
    * ``prefill_forward`` / ``decode_forward``: consume the
      ``page_tables_per_layer`` list (one tensor per decoder layer, layer-
      aligned with the model's ``self.layers``) and pass each entry to its
      corresponding attention layer. The plugin pre-expands
      ``block_tables_per_group`` into this per-layer view at submission
      time so bridges don't have to re-derive vLLM's group construction
      order — see ``TTModelRunner._block_tables_per_layer``.
    * ``allocate_kv_cache_per_layer``: typically just delegates to
      :func:`allocate_vllm_kv_cache_per_layer` with the model handles.
    Until a subclass overrides them, ``prefill_forward`` and
    ``decode_forward`` raise :class:`NotImplementedError` to make the
    contract explicit. Legacy (non-hybrid) models never see the
    ``page_tables_per_layer`` kwarg — vLLM's plugin opts in via the
    presence of ``get_kv_cache_spec`` on the model class.
    """

    # Keep this in sync with get_kv_cache_spec below and with the TT vLLM
    # worker's token-budget calculation. While SlidingWindowSpec is disabled,
    # vLLM produces one full-attention KV group and the legacy single page_table
    # path is sufficient. When SlidingWindowSpec is restored, flip this back so
    # warmup also exercises the per-layer persistent page-table path.
    _HYBRID_KV_CACHE_GROUPS_ENABLED = False

    @classmethod
    def get_kv_cache_spec(cls, vllm_config):
        """Build per-layer KVCacheSpec from HF config ``layer_types``.
        Returns a dict keyed by ``model.layers.<idx>.self_attn`` (the
        upstream attention-layer naming convention vLLM's KVCacheGroup
        machinery understands) so :func:`_parse_layer_index` on the
        runner side can map each spec back to its model layer index.
        """
        from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE

        # SlidingWindowSpec import intentionally dropped; restore alongside the
        # branch below when re-enabling kv cache groups.
        from vllm.v1.kv_cache_interface import FullAttentionSpec

        model_config = vllm_config.model_config
        cache_config = vllm_config.cache_config

        hf_config = model_config.hf_config
        text_config = getattr(hf_config, "text_config", hf_config)
        layer_types = getattr(text_config, "layer_types", None)
        if layer_types is None:
            raise ValueError(
                f"{cls.__name__}.get_kv_cache_spec requires "
                "hf_config.text_config.layer_types (one of 'full_attention' / "
                "'sliding_attention' per layer); none found on this model"
            )
        num_kv_heads = model_config.get_num_kv_heads(vllm_config.parallel_config)
        head_size = model_config.get_head_size()
        dtype = (
            model_config.dtype
            if cache_config.cache_dtype == "auto"
            else STR_DTYPE_TO_TORCH_DTYPE[cache_config.cache_dtype]
        )
        block_size = cache_config.block_size

        common = dict(
            block_size=block_size,
            num_kv_heads=num_kv_heads,
            head_size=head_size,
            dtype=dtype,
        )

        # SlidingWindowSpec is temporarily disabled: TT-side decode passes the
        # absolute position to paged_update_cache / paged_sdpa_decode, but vLLM
        # zero-pads the sliding group's page_table past sliding_window/block_size
        # entries, so positions beyond the sliding window collapse onto physical
        # block 0 and silently corrupt the cache. Emit FullAttentionSpec for every
        # layer so vLLM allocates a max_model_len cache per layer; the SDPA op's
        # own sliding_window_size kwarg still trims attention correctly on the
        # read side.
        spec_per_layer = {}
        for i, lt in enumerate(layer_types):
            name = f"model.layers.{i}.self_attn"
            if lt not in ("sliding_attention", "full_attention"):
                raise ValueError(
                    f"Unsupported layer_type {lt!r} at layer {i} on "
                    f"{cls.__name__}; expected 'full_attention' or "
                    "'sliding_attention'"
                )
            spec_per_layer[name] = FullAttentionSpec(**common)
        return spec_per_layer

    def prefill_forward(self, *args, **kwargs):
        raise NotImplementedError(
            f"{type(self).__name__} must override prefill_forward to consume "
            "`page_tables_per_layer` and pass each entry to the matching "
            "attention layer."
        )

    def decode_forward(self, *args, **kwargs):
        raise NotImplementedError(
            f"{type(self).__name__} must override decode_forward to consume "
            "`page_tables_per_layer` and pass each entry to the matching "
            "attention layer."
        )

    def allocate_kv_cache_per_layer(self, per_layer_specs):
        return allocate_vllm_kv_cache_per_layer(per_layer_specs, dp_model=self.model, tt_cache_path=self.cache_path)

    @staticmethod
    def _reload_per_layer_page_tables(kwargs) -> bool:
        """Whether the explicit contract requests per-layer page-table upload."""
        return bool(kwargs["reload_inputs"] or kwargs["reload_page_table"])

    def _ensure_page_tables_per_layer(self, page_tables_per_layer, page_table):
        """When invoked outside the vLLM hybrid plugin (e.g. by warmup
        which only knows about the legacy single ``page_table``), optionally
        broadcast the single page table to a per-layer list.
        Broadcasting is only correct while hybrid KV cache groups are enabled:
        trace capture then needs to exercise the per-layer code path inside
        ``Transformer.forward`` so replay reads the persistent per-layer device
        tensors updated before each call. While hybrid groups are temporarily
        disabled, all layers use one full-attention KV group, so we intentionally
        keep warmup/runtime on the legacy single-page-table path.
        """
        if page_tables_per_layer is not None or page_table is None or not self._HYBRID_KV_CACHE_GROUPS_ENABLED:
            return page_tables_per_layer
        # Broadcast the same torch tensor across every layer in every
        # submesh — content is identical, persistent allocation gives each
        # layer its own device tensor at a stable address.
        num_layers = len(self.model[0].layers)
        return [page_table] * num_layers

    def _chunk_page_tables_per_dp(self, page_tables_per_layer):
        """Split a global per-layer list along DP into one per-layer list
        per submesh.
        The plugin pads each per-layer table to the global
        ``(max_num_seqs * data_parallel, max_num_blocks_per_req)`` shape
        (see ``TTModelRunner._block_tables_per_layer``); warmup likewise
        builds a global-batch tensor. ``Generator.decode_forward`` already
        does ``torch.chunk(page_table, self.data_parallel, 0)`` for the
        legacy single-page-table path before the per-submesh
        ``prepare_inputs_decode``; the hybrid bridge has to do the
        equivalent so each submesh's ``_page_tables_to_ttnn`` receives a
        per-DP slice whose batch dim matches the submesh's K/V tensors.
        Without this, ``paged_update_cache`` asserts a batch-size mismatch
        on multi-DP runs.
        """
        if page_tables_per_layer is None:
            return None
        dp = self.data_parallel
        if dp <= 1:
            return [page_tables_per_layer]
        per_submesh = [list() for _ in range(dp)]
        for pt in page_tables_per_layer:
            if pt is None or isinstance(pt, ttnn.Tensor):
                # Already-resolved or absent entries pass through unchanged
                # to every submesh — chunking only applies to torch tensors
                # carrying global batch.
                for s in per_submesh:
                    s.append(pt)
                continue
            chunks = torch.chunk(pt, dp, dim=0)
            for s, c in zip(per_submesh, chunks):
                s.append(c)
        return per_submesh

    def _route_per_layer_page_tables(self, per_submesh_page_tables):
        """Stash each submesh's per-layer page-table list on its model
        handle for the duration of a forward call.
        ``Generator``'s prefill/decode paths invoke
        ``model[i].ttnn_prefill_forward`` / ``ttnn_decode_forward`` from
        many sites (warmup, trace capture, traced replay, etc.) without
        forwarding an arbitrary kwarg. Threading the per-layer list through
        every site would be a wide change for a feature only this hybrid
        bridge consumes, so we use a localised attribute injection: each
        model reads ``getattr(self, "_active_page_tables_per_layer", None)``
        when its own kwarg is None. ``per_submesh_page_tables[i]`` is the
        per-layer slice that submesh ``i`` should see; ``None`` clears the
        stash entirely (legacy fallback).
        """

        class _Stash:
            def __init__(self, models, per_submesh):
                self._models = models
                self._per_submesh = per_submesh

            def __enter__(self):
                if self._per_submesh is None:
                    return
                for m, value in zip(self._models, self._per_submesh):
                    m._active_page_tables_per_layer = value

            def __exit__(self, *_):
                if self._per_submesh is None:
                    return
                for m in self._models:
                    if hasattr(m, "_active_page_tables_per_layer"):
                        del m._active_page_tables_per_layer

        return _Stash(self.model, per_submesh_page_tables)


def initialize_vllm_text_transformer(
    hf_config,
    tt_data_parallel,
    mesh_device,
    max_batch_size,
    max_seq_len,
    n_layers=None,
    dtype=ttnn.bfloat8_b,
    optimizations=DecodersPrecision.performance,
):
    submesh_devices = create_submeshes(mesh_device, tt_data_parallel)
    # Load model args, weights
    model_args = []
    for submesh in submesh_devices:
        model_args_i = ModelArgs(
            submesh,
            instruct=(
                "Instruct" in hf_config._name_or_path or "DeepSeek-R1-Distill-Llama-70B" in hf_config._name_or_path
            ),
            max_batch_size=max_batch_size // tt_data_parallel,
            optimizations=lambda model_args: optimizations(model_args.n_layers, model_args.model_name),
            max_seq_len=max_seq_len,
        )

        assert model_args_i.model_name.replace("-", "") in hf_config._name_or_path.replace(
            "-", ""
        ), f"The model specified in vLLM ({hf_config._name_or_path}) does not match the model name ({model_args_i.model_name}) with model weights ({model_args_i.CKPT_DIR})."
        if n_layers is not None:
            model_args_i.n_layers = n_layers

        model_args.append(model_args_i)

    state_dict = model_args[0].load_state_dict()

    tt_model = []
    for i, submesh in enumerate(submesh_devices):
        tt_model_i = Transformer(
            args=model_args[i],
            mesh_device=submesh,
            dtype=dtype,
            state_dict=state_dict,
            weight_cache_path=model_args[i].weight_cache_path(dtype),
            use_paged_kv_cache=True,
        )
        tt_model.append(tt_model_i)

    return tt_model, model_args


class DummyInputsBuilder(BaseDummyInputsBuilder):
    """
    We don't need to implement a dummy input builder since we don't do profiling in vLLM.
    Create callable class just for processor registration.
    """

    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
        raise NotImplementedError

    def get_dummy_mm_data(
        self,
        seq_len: int,
        mm_counts: Mapping[str, int],
    ) -> MultiModalDataDict:
        raise NotImplementedError


class CustomNamespace(SimpleNamespace):
    def __contains__(self, key):
        return key in self.__dict__


@MULTIMODAL_REGISTRY.register_processor(
    Mistral3MultiModalProcessor,
    info=Mistral3ProcessingInfo,
    dummy_inputs=Mistral3DummyInputsBuilder,
)
class Mistral3ForConditionalGeneration(Generator, SupportsMultiModal):
    decode_input_update_contract = 1

    model_capabilities = {
        "supports_prefix_caching": False,
        "supports_sample_on_device": True,
        "max_device_top_k": 32,
    }

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.MISTRAL_IMAGE_TOKEN_ID = 151655
        self.max_gen_len = self.model_args[0].max_seq_len - 1

    @classmethod
    def initialize_vllm_model(
        cls,
        hf_config,
        mesh_device,
        max_batch_size,
        max_seq_len=131072,
        tt_data_parallel=1,
        optimizations: str = None,
    ):
        assert optimizations is None, "Custom optimizations are not supported for this model"
        from models.tt_transformers.demo.simple_vision_demo import create_multimodal_model

        max_seq_len = 1024 * 128

        submesh_devices = create_submeshes(mesh_device, tt_data_parallel)

        model_args = []
        model = []
        state_dict = None

        for submesh in submesh_devices:
            model_args_i, model_i, state_dict = create_multimodal_model(
                mesh_device=submesh,
                max_batch_size=max_batch_size // tt_data_parallel,
                max_seq_len=max_seq_len,
                use_paged_kv_cache=True,
                checkpoint=state_dict,
            )
            model_args.append(model_args_i)
            model.append(model_i)

        return cls(model, model_args, mesh_device)

    @property
    def cache_path(self):
        return self.model_args[0].model_cache_path

    def prefill_forward(self, *args, **kwargs):
        self.tokenizer = self.model_args[0].tokenizer
        pad_token_id = self.tokenizer.pad_token_id

        tokens = kwargs["tokens"]
        prompt_lens = kwargs["prompt_lens"]
        inputs = CustomNamespace()
        inputs.input_ids = tokens
        data = kwargs.get("images", None)
        for i in range(tokens.shape[0]):
            tokens[i][prompt_lens[i] :] = pad_token_id
        pixel_values, image_sizes = None, None

        if data and hasattr(data[0], "pixel_values"):
            pixel_values = [im.pixel_values for im in data if hasattr(im, "pixel_values")]
            image_sizes = [im.image_sizes for im in data if hasattr(im, "image_sizes")]

        page_table = kwargs.get("page_table", None)
        kv_cache = kwargs.get("kv_cache", None)

        return super().prefill_forward_text(
            tokens=inputs.input_ids,
            page_table=page_table,
            kv_cache=kv_cache,
            prompt_lens=prompt_lens,
            pixel_values=pixel_values if pixel_values else None,
            image_sizes=image_sizes if image_sizes else None,
        )

    def decode_forward(self, *args, **kwargs):
        return super().decode_forward(*args, **kwargs)

    def allocate_kv_cache(self, *args, **kwargs):
        return allocate_vllm_kv_cache(*args, **kwargs, dp_model=self.model, tt_cache_path=self.cache_path)


# Mllama is currently not supported in vLLM V1.
# TODO: Remove or re-enable when Mllama is supported in vLLM V1.
# @MULTIMODAL_REGISTRY.register_processor(
#     MllamaMultiModalProcessor, info=TT_MllamaProcessingInfo, dummy_inputs=DummyInputsBuilder
# )
class MllamaForConditionalGeneration(Generator, SupportsMultiModal):
    decode_input_update_contract = 1

    # Class-level capabilities
    # Note: Mllama doesn't support prefix caching (it's V0 only)
    # decode_forward calls decode_forward_llama_vision and discards anything
    # but logits, so sampling_params never reach a sampler — explicitly
    # declare on-device sampling unsupported.
    model_capabilities = {
        "supports_prefix_caching": False,
        "supports_async_decode": False,
        "supports_sample_on_device": False,
    }

    @classmethod
    def get_max_tokens_all_users(
        cls,
        model_name: str = "",
        num_devices: int = 1,
        tt_data_parallel: int = 1,
        **kwargs,
    ) -> int:
        """Returns config-specific all-user KV-cache token capacity."""
        devices_per_dp_cache = num_devices // tt_data_parallel
        is_wormhole = is_wormhole_b0()

        # Llama90B on WH T3K
        if "Llama-3.2-90B" in model_name and devices_per_dp_cache == 8 and is_wormhole:
            return 65_536
        return super().get_max_tokens_all_users(
            model_name=model_name,
            num_devices=num_devices,
            tt_data_parallel=tt_data_parallel,
            **kwargs,
        )

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.MLLAMA_IMAGE_TOKEN_ID = 128256
        self.max_gen_len = self.model_args[0].max_seq_len - 1  # TODO: double check what this should be

    @classmethod
    def initialize_vllm_model(
        cls, hf_config, mesh_device, max_batch_size, max_seq_len, tt_data_parallel=1, optimizations: str = None
    ):
        assert optimizations is None, "Custom optimizations are not supported for this model"
        from models.tt_transformers.demo.simple_vision_demo import create_multimodal_model

        submesh_devices = create_submeshes(mesh_device, tt_data_parallel)

        model_args = []
        model = []
        state_dict = None

        for submesh in submesh_devices:
            model_args_i, model_i, state_dict = create_multimodal_model(
                mesh_device=submesh,
                max_batch_size=max_batch_size // tt_data_parallel,
                max_seq_len=max_seq_len,
                use_paged_kv_cache=True,
                checkpoint=state_dict,
            )
            model_args.append(model_args_i)
            model.append(model_i)

        return cls(model, model_args, mesh_device)

    @property
    def cache_path(self):
        return self.model_args[0].model_cache_path

    @property
    def max_cross_attn_tokens(self):
        return self.model_args[0].vision_max_num_chunks * nearest_32(self.model_args[0].vision_chunk_ntok)

    def prefill_forward(
        self,
        tokens: torch.Tensor,
        images: Union[List[Image], List[List[Image]]],
        page_table: torch.Tensor,
        kv_cache,
        prompt_lens,
        cross_page_table: torch.Tensor,
    ):
        """
        Replaces prefill_forward from Generator with a version that supports mask creation.
        """
        batch = tokens.shape[0]

        vision_images = []
        vision_masks = []
        total_lens = []
        for user_id in range(batch):
            image = images[user_id]
            if isinstance(image, list):
                assert len(image) == 1, "Only one image is supported for each user in the batch"
                image = image[0]
            vision_images.append([image] if image else None)
            prompt_tokens = [int(tokens[user_id, i]) for i in range(prompt_lens[user_id])]
            vision_masks.append(create_vision_mask(prompt_tokens, self.MLLAMA_IMAGE_TOKEN_ID) if image else None)
            total_lens.append(prompt_lens[user_id] + self.max_gen_len)

        return super().prefill_forward(
            vision_images,
            vision_masks,
            tokens,
            None,
            total_lens,
            prompt_lens,
            page_table=page_table,
            kv_cache=kv_cache,
            cross_page_table=cross_page_table,
        )

    def decode_forward(self, *args, **kwargs):
        reload_inputs = kwargs.pop("reload_inputs")
        reload_page_table = kwargs.pop("reload_page_table")
        reload_sampling_params = kwargs.pop("reload_sampling_params")
        reset_sampling_state = kwargs.pop("reset_sampling_state")
        # Mllama has no persistent model state keyed by vLLM's condensed
        # decode slots. Contract v1 nevertheless supplies slot_remap for every
        # decode so adapters with recurrent state can consume it.
        kwargs.pop("slot_remap", None)
        if not reload_inputs or reload_page_table or reload_sampling_params or reset_sampling_state:
            raise ValueError("Mllama requires a full host-input reload and has no " "device sampling state")
        logits = super().decode_forward_llama_vision(*args, **kwargs)
        if isinstance(logits, tuple):
            return logits[0]
        else:
            return logits

    def allocate_kv_cache(self, *args, **kwargs):
        return allocate_vllm_kv_cache(*args, **kwargs, dp_model=self.model, tt_cache_path=self.cache_path)


class LlamaForCausalLM(Generator):
    decode_input_update_contract = 1

    # Class-level capabilities
    model_capabilities = {
        "supports_prefix_caching": True,
        "supports_async_decode": True,
        "supports_sample_on_device": True,
        "max_device_top_k": 32,
        # prefill_forward_single_user_text already routes a nonzero start_pos to
        # the chunked SDPA, and the generator floors the offset to what that op
        # needs, so a prompt split across engine steps needs no new prefill code.
        "supports_chunked_prefill": True,
    }

    @classmethod
    def get_max_tokens_all_users(
        cls,
        model_name: str = "",
        num_devices: int = 1,
        tt_data_parallel: int = 1,
        **kwargs,
    ) -> int:
        """Returns config-specific all-user KV-cache token capacity."""
        devices_per_dp_cache = num_devices // tt_data_parallel
        is_wormhole = is_wormhole_b0()

        # Llama8B on N150
        if "Llama-3.1-8B" in model_name and devices_per_dp_cache == 1 and is_wormhole:
            return 32_768
        return super().get_max_tokens_all_users(
            model_name=model_name,
            num_devices=num_devices,
            tt_data_parallel=tt_data_parallel,
            **kwargs,
        )

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    @classmethod
    def initialize_vllm_model(
        cls,
        hf_config,
        mesh_device,
        max_batch_size,
        max_seq_len,
        n_layers=None,
        tt_data_parallel=1,
        optimizations: str = "performance",
    ):
        hf_model_name = hf_config._name_or_path
        if (
            ("3.1-8B" in hf_model_name or "3.2-11B" in hf_model_name)
            and mesh_device.get_num_devices() == 1
            and is_wormhole_b0()
        ):
            MAX_PROMPT_LEN = 32768
            if max_seq_len > MAX_PROMPT_LEN:
                raise ValueError(
                    f"TT-LLama8B and TT-Llama11B do not support max_model_len greater than {MAX_PROMPT_LEN} on N150 "
                    f"(received {max_seq_len}). Set --max_model_len to {MAX_PROMPT_LEN} or lower in vLLM."
                )

        tt_model, model_args = initialize_vllm_text_transformer(
            hf_config,
            tt_data_parallel,
            mesh_device,
            max_batch_size,
            max_seq_len=max_seq_len,
            n_layers=n_layers,
            dtype=ttnn.bfloat8_b,
            optimizations=(
                DecodersPrecision.from_string(optimizations)
                if optimizations is not None
                else DecodersPrecision.performance
            ),
        )
        return cls(tt_model, model_args, mesh_device)

    @property
    def cache_path(self):
        return self.model_args[0].model_cache_path

    def prefill_forward(self, *args, **kwargs):
        return super().prefill_forward_text(*args, **kwargs)

    def decode_forward(self, *args, **kwargs):
        return super().decode_forward(*args, **kwargs)

    def allocate_kv_cache(self, *args, **kwargs):
        return allocate_vllm_kv_cache(*args, **kwargs, dp_model=self.model, tt_cache_path=self.cache_path)


class QwenForCausalLM(Generator):
    decode_input_update_contract = 1

    # Class-level capabilities
    model_capabilities = {
        "supports_prefix_caching": True,
        "supports_async_decode": True,
        "supports_sample_on_device": True,
        "max_device_top_k": 32,
        # prefill_forward_single_user_text already routes a nonzero start_pos to
        # the chunked SDPA, and the generator floors the offset to what that op
        # needs, so a prompt split across engine steps needs no new prefill code.
        "supports_chunked_prefill": True,
    }

    @classmethod
    def get_max_tokens_all_users(
        cls,
        model_name: str = "",
        num_devices: int = 1,
        tt_data_parallel: int = 1,
        **kwargs,
    ) -> int:
        """Returns config-specific all-user KV-cache token capacity."""
        devices_per_dp_cache = num_devices // tt_data_parallel
        is_wormhole = is_wormhole_b0()

        # Qwen3-8B on N150 (same constraint as Llama8B-N150)
        if "Qwen3-8B" in model_name and devices_per_dp_cache == 1 and is_wormhole:
            return 32_768
        # DeepSeek-R1-Distill-Qwen-14B / Qwen2.5-14B on N300
        if (
            ("DeepSeek-R1-Distill-Qwen-14B" in model_name or "Qwen2.5-14B" in model_name)
            and devices_per_dp_cache == 2
            and is_wormhole
        ):
            return 65_536
        return super().get_max_tokens_all_users(
            model_name=model_name,
            num_devices=num_devices,
            tt_data_parallel=tt_data_parallel,
            **kwargs,
        )

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    @classmethod
    def initialize_vllm_model(
        cls,
        hf_config,
        mesh_device,
        max_batch_size,
        max_seq_len,
        n_layers=None,
        tt_data_parallel=1,
        optimizations: str = "performance",
    ):
        tt_model, model_args = initialize_vllm_text_transformer(
            hf_config,
            tt_data_parallel,
            mesh_device,
            max_batch_size,
            max_seq_len=max_seq_len,
            n_layers=n_layers,
            dtype=ttnn.bfloat8_b,
            optimizations=(
                DecodersPrecision.from_string(optimizations)
                if optimizations is not None
                else DecodersPrecision.performance
            ),
        )
        return cls(tt_model, model_args, mesh_device)

    @property
    def cache_path(self):
        return self.model_args[0].model_cache_path

    def prefill_forward(self, *args, **kwargs):
        return super().prefill_forward_text(*args, **kwargs)

    def decode_forward(self, *args, **kwargs):
        return super().decode_forward(*args, **kwargs)

    def allocate_kv_cache(self, *args, **kwargs):
        return allocate_vllm_kv_cache(*args, **kwargs, dp_model=self.model, tt_cache_path=self.cache_path)


class CohereForCausalLM(Generator):
    """Command-R / Command-A family (HF model_type "cohere").
    Minimal text-only dense wrapper, same shape as QwenForCausalLM. All
    Cohere-specific behaviour (parallel decoder block, LayerNorm instead of
    RMSNorm, LM-head logit_scale 0.0625, tied embeddings) is handled by the
    model_type=="cohere" family dispatch in model.py / model_config.py (P5a).
    Canada Quant Labs (org-internal) - bounty tt-metal#49307 track.
    """

    # Class-level capabilities
    model_capabilities = {
        "supports_prefix_caching": True,
        "supports_async_decode": True,
        "supports_sample_on_device": True,
    }

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    @classmethod
    def initialize_vllm_model(
        cls,
        hf_config,
        mesh_device,
        max_batch_size,
        max_seq_len,
        n_layers=None,
        tt_data_parallel=1,
        optimizations: str = "performance",
    ):
        tt_model, model_args = initialize_vllm_text_transformer(
            hf_config,
            tt_data_parallel,
            mesh_device,
            max_batch_size,
            max_seq_len=max_seq_len,
            n_layers=n_layers,
            dtype=ttnn.bfloat8_b,
            optimizations=DecodersPrecision.from_string(optimizations)
            if optimizations is not None
            else DecodersPrecision.performance,
        )
        return cls(tt_model, model_args, mesh_device)

    @property
    def cache_path(self):
        return self.model_args[0].model_cache_path

    def prefill_forward(self, *args, **kwargs):
        return super().prefill_forward_text(*args, **kwargs)

    def decode_forward(self, *args, **kwargs):
        return super().decode_forward(*args, **kwargs)

    def allocate_kv_cache(self, *args, **kwargs):
        return allocate_vllm_kv_cache(*args, **kwargs, dp_model=self.model, tt_cache_path=self.cache_path)


class MistralForCausalLM(Generator):
    decode_input_update_contract = 1

    # Class-level capabilities
    model_capabilities = {
        "supports_prefix_caching": True,
        "supports_async_decode": True,
        "supports_sample_on_device": True,
        "max_device_top_k": 32,
        # prefill_forward_single_user_text already routes a nonzero start_pos to
        # the chunked SDPA, and the generator floors the offset to what that op
        # needs, so a prompt split across engine steps needs no new prefill code.
        "supports_chunked_prefill": True,
    }

    @classmethod
    def get_max_tokens_all_users(
        cls,
        model_name: str = "",
        num_devices: int = 1,
        tt_data_parallel: int = 1,
        **kwargs,
    ) -> int:
        """Returns config-specific all-user KV-cache token capacity."""
        devices_per_dp_cache = num_devices // tt_data_parallel
        is_wormhole = is_wormhole_b0()

        # Mistral-7B on N150
        if "Mistral-7B" in model_name and devices_per_dp_cache == 1 and is_wormhole:
            return 65_536
        return super().get_max_tokens_all_users(
            model_name=model_name,
            num_devices=num_devices,
            tt_data_parallel=tt_data_parallel,
            **kwargs,
        )

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    @classmethod
    def initialize_vllm_model(
        cls,
        hf_config,
        mesh_device,
        max_batch_size,
        max_seq_len,
        n_layers=None,
        tt_data_parallel=1,
        optimizations: str = "performance",
    ):
        tt_model, model_args = initialize_vllm_text_transformer(
            hf_config,
            tt_data_parallel,
            mesh_device,
            max_batch_size,
            max_seq_len=max_seq_len,
            n_layers=n_layers,
            dtype=ttnn.bfloat8_b,
            optimizations=(
                DecodersPrecision.from_string(optimizations)
                if optimizations is not None
                else DecodersPrecision.performance
            ),
        )
        return cls(tt_model, model_args, mesh_device)

    @property
    def cache_path(self):
        return self.model_args[0].model_cache_path

    def prefill_forward(self, *args, **kwargs):
        return super().prefill_forward_text(*args, **kwargs)

    def decode_forward(self, *args, **kwargs):
        return super().decode_forward(*args, **kwargs)

    def allocate_kv_cache(self, *args, **kwargs):
        return allocate_vllm_kv_cache(*args, **kwargs, dp_model=self.model, tt_cache_path=self.cache_path)


@MULTIMODAL_REGISTRY.register_processor(
    Gemma3MultiModalProcessor,
    info=Gemma3ProcessingInfo,
    dummy_inputs=Gemma3DummyInputsBuilder,
)
class Gemma3ForConditionalGeneration(HybridAttentionForCausalLM, SupportsMultiModal):
    """Gemma3 multimodal — hybrid attention (sliding-window + full).
    Gemma3's text decoder alternates ``sliding_attention`` and
    ``full_attention`` per ``hf_config.text_config.layer_types`` (a 5:1
    ratio in the 27B variant), so the bridge inherits from
    :class:`HybridAttentionForCausalLM` to opt into vLLM's hybrid kv cache
    manager. Sliding-window layers index a smaller paged pool than
    full-attention layers — the per-layer KV cache shape difference is
    where the asymmetric-hybrid memory savings live, and was the original
    motivation for kv-cache-groups (it's what unblocks the 62-layer × 107
    MB-per-layer DRAM OOM on T3K seen in run 25437459815).
    Mirrors the ``GptOssForCausalLM`` plumbing: ``prefill_forward`` /
    ``decode_forward`` stash ``page_tables_per_layer`` on each
    ``self.model[i]`` for the duration of a single
    ``super().{prefill_forward_text,decode_forward}`` call. The underlying
    ``Transformer.ttnn_*_forward`` (which ``TtGemmaModel`` inherits)
    picks up the stash via ``_active_page_tables_per_layer`` and routes
    each layer's attention to its own page table.
    """

    decode_input_update_contract = 1

    # Class-level capabilities
    model_capabilities = {
        "supports_prefix_caching": False,
        "supports_async_decode": True,
        "supports_sample_on_device": True,
        "max_device_top_k": 32,
    }

    @classmethod
    def get_max_tokens_all_users(
        cls,
        model_name: str = "",
        num_devices: int = 1,
        tt_data_parallel: int = 1,
        **kwargs,
    ) -> int:
        """Returns config-specific all-user KV-cache token capacity."""
        devices_per_dp_cache = num_devices // tt_data_parallel
        is_wormhole = is_wormhole_b0()

        # gemma-3-4b on wormhole configurations with up to 2 devices per DP shard
        if "gemma-3-4b" in model_name.lower() and devices_per_dp_cache in (1, 2) and is_wormhole:
            return 65_536
        return super().get_max_tokens_all_users(
            model_name=model_name,
            num_devices=num_devices,
            tt_data_parallel=tt_data_parallel,
            **kwargs,
        )

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    @classmethod
    def initialize_vllm_model(
        cls,
        hf_config,
        mesh_device,
        max_batch_size,
        max_seq_len=131072,
        n_layers=None,
        tt_data_parallel=1,
        optimizations: str = "performance",
    ):
        from models.demos.multimodal.gemma3.demo.vision_demo import create_multimodal_model

        optimizations = (
            DecodersPrecision.from_string(optimizations) if optimizations is not None else DecodersPrecision.performance
        )

        submesh_devices = create_submeshes(mesh_device, tt_data_parallel)

        model_args = []
        model = []
        state_dict = None

        for submesh in submesh_devices:
            model_args_i, model_i, state_dict = create_multimodal_model(
                mesh_device=submesh,
                max_batch_size=max_batch_size // tt_data_parallel,
                max_seq_len=max_seq_len,
                use_paged_kv_cache=True,
                checkpoint=state_dict,
                optimizations=lambda model_args: optimizations(model_args.n_layers, model_args.model_name),
            )
            model_args.append(model_args_i)
            model.append(model_i)

        return cls(model, model_args, mesh_device)

    @property
    def cache_path(self):
        return self.model_args[0].model_cache_path

    def prefill_forward(self, *args, page_tables_per_layer=None, **kwargs):
        # While hybrid KV cache groups are disabled (one full-attention group
        # for every layer), the per-layer page-table routing inside this
        # bridge is buggy for users_row_sharded models: it shards page tables
        # naively by mesh row, which doesn't match the gpt-oss
        # slot // max_local_batch_size → row mapping and produces null
        # content on the rows whose page-table chunks point at the wrong
        # KV blocks. Until a row-aware per-layer routing lands, skip the
        # hybrid path entirely and let the legacy single page_table flow
        # through Generator.prefill_forward_text reach the model untouched.
        if not self._HYBRID_KV_CACHE_GROUPS_ENABLED:
            return super().prefill_forward_text(*args, **kwargs)
        page_tables_per_layer = self._ensure_page_tables_per_layer(page_tables_per_layer, kwargs.get("page_table"))
        per_submesh = self._chunk_page_tables_per_dp(page_tables_per_layer)
        # Push the per-layer block IDs into the persistent device buffers
        # *before* entering ``Generator.prefill_forward_text`` — that path
        # may execute a captured trace, which reads block IDs from the
        # persistent addresses and forbids in-trace writes. Allocation
        # itself happens lazily in ``Transformer._page_tables_to_ttnn``
        # the first time the inner forward runs (warmup compile).
        if per_submesh is not None:
            for m, pt_for_submesh in zip(self.model, per_submesh):
                m.update_persistent_per_layer_page_tables(pt_for_submesh)
        with self._route_per_layer_page_tables(per_submesh):
            return super().prefill_forward_text(**kwargs)

    def decode_forward(self, *args, page_tables_per_layer=None, **kwargs):
        # See prefill_forward note above. Skip the hybrid path while
        # _HYBRID_KV_CACHE_GROUPS_ENABLED is False.
        if not self._HYBRID_KV_CACHE_GROUPS_ENABLED:
            return super(HybridAttentionForCausalLM, self).decode_forward(*args, **kwargs)
        page_tables_per_layer = self._ensure_page_tables_per_layer(page_tables_per_layer, kwargs.get("page_table"))
        per_submesh = self._chunk_page_tables_per_dp(page_tables_per_layer)
        if per_submesh is not None and self._reload_per_layer_page_tables(kwargs):
            for m, pt_for_submesh in zip(self.model, per_submesh):
                m.update_persistent_per_layer_page_tables(pt_for_submesh)
        with self._route_per_layer_page_tables(per_submesh):
            # Skip ``HybridAttentionForCausalLM.decode_forward``, which is a
            # NotImplementedError placeholder; route to ``Generator``'s
            # actual decode implementation.
            return super(HybridAttentionForCausalLM, self).decode_forward(*args, **kwargs)

    def allocate_kv_cache(self, *args, **kwargs):
        return allocate_vllm_kv_cache(*args, **kwargs, dp_model=self.model, tt_cache_path=self.cache_path)


class Exaone4_5_ForConditionalGeneration(HybridAttentionForCausalLM):
    """EXAONE-4.5 (text decoder) for vLLM integration.
    Named after the HF architecture string. Note that the vLLM TT plugin
    does NOT pick this class up automatically: the plugin prepends ``TT``
    to the checkpoint architecture and its built-in registry has no EXAONE
    entry, so serving requires registering
    ``TTExaone4_5_ForConditionalGeneration`` -> this class in the vLLM
    ``ModelRegistry`` of the engine-core process. tt-inference-server does
    this via an ``EXTRA_MODELS_DIR`` plugin bundle (see
    tenstorrent/tt-inference-server#5023); direct plugin users must supply
    an equivalent registration and set ``--max-model-len 131072``: the
    checkpoint advertises 262144 tokens, but prefill above 131072 needs
    sliding-window chunked prefill, which is not supported yet.
    The checkpoint is multimodal, but ModelArgs runs its text decoder only
    (``force_text_only``): hybrid LLLG sliding/full attention (via
    :class:`HybridAttentionForCausalLM`), post-norm decoder, NoPE
    full-attention layers. Vision inputs are not yet wired through vLLM —
    this serves text-only requests.
    """

    # Class-level capabilities
    model_capabilities = {
        "supports_prefix_caching": False,  # Sliding window => no prefix caching
        "supports_async_decode": True,
        "supports_sample_on_device": True,
    }

    # One prefill chunk on P150x8 (MAX_PREFILL_CHUNK_SIZES["EXAONE-4.5-33B"]).
    MAX_SUPPORTED_SEQ_LEN = 128 * 1024

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    @classmethod
    def initialize_vllm_model(
        cls,
        hf_config,
        mesh_device,
        max_batch_size,
        max_seq_len,
        n_layers=None,
        tt_data_parallel=1,
        optimizations: str = "performance",
    ):
        # Prompts longer than one prefill chunk take the chunked path, where
        # the sliding-window attention layers raise NotImplementedError, so
        # reject the checkpoint's 262144-token default up front.
        if max_seq_len > cls.MAX_SUPPORTED_SEQ_LEN:
            raise ValueError(
                f"EXAONE-4.5 currently supports max_seq_len <= {cls.MAX_SUPPORTED_SEQ_LEN}; "
                f"got {max_seq_len}. Sliding-window chunked prefill is not supported; "
                "pass --max-model-len 131072 to vLLM."
            )
        tt_model, model_args = initialize_vllm_text_transformer(
            hf_config,
            tt_data_parallel,
            mesh_device,
            max_batch_size,
            max_seq_len=max_seq_len,
            n_layers=n_layers,
            dtype=ttnn.bfloat8_b,
            optimizations=DecodersPrecision.from_string(optimizations)
            if optimizations is not None
            else DecodersPrecision.performance,
        )
        return cls(tt_model, model_args, mesh_device)

    @property
    def cache_path(self):
        return self.model_args[0].model_cache_path

    def prefill_forward(self, *args, page_tables_per_layer=None, **kwargs):
        # While hybrid KV cache groups are disabled upstream (every layer is
        # one full-attention group), skip the per-layer page-table routing
        # and let the legacy single page_table flow through
        # ``Generator.prefill_forward_text`` — same rationale as
        # ``GptOssForCausalLM.prefill_forward``.
        if not self._HYBRID_KV_CACHE_GROUPS_ENABLED:
            return super().prefill_forward_text(*args, **kwargs)
        page_tables_per_layer = self._ensure_page_tables_per_layer(page_tables_per_layer, kwargs.get("page_table"))
        per_submesh = self._chunk_page_tables_per_dp(page_tables_per_layer)
        if per_submesh is not None:
            for m, pt_for_submesh in zip(self.model, per_submesh):
                m.update_persistent_per_layer_page_tables(pt_for_submesh)
        with self._route_per_layer_page_tables(per_submesh):
            return super().prefill_forward_text(*args, **kwargs)

    def decode_forward(self, *args, page_tables_per_layer=None, **kwargs):
        # See prefill_forward note above.
        if not self._HYBRID_KV_CACHE_GROUPS_ENABLED:
            return super(HybridAttentionForCausalLM, self).decode_forward(*args, **kwargs)
        page_tables_per_layer = self._ensure_page_tables_per_layer(page_tables_per_layer, kwargs.get("page_table"))
        per_submesh = self._chunk_page_tables_per_dp(page_tables_per_layer)
        if per_submesh is not None:
            for m, pt_for_submesh in zip(self.model, per_submesh):
                m.update_persistent_per_layer_page_tables(pt_for_submesh)
        with self._route_per_layer_page_tables(per_submesh):
            # ``HybridAttentionForCausalLM.decode_forward`` is a placeholder;
            # route to ``Generator``'s actual decode implementation.
            return super(HybridAttentionForCausalLM, self).decode_forward(*args, **kwargs)

    def allocate_kv_cache(self, *args, **kwargs):
        return allocate_vllm_kv_cache(*args, **kwargs, dp_model=self.model, tt_cache_path=self.cache_path)


class GptOssForCausalLM(HybridAttentionForCausalLM):
    """GPT-OSS model for vLLM integration.
    GPT-OSS is a hybrid attention model — its layers alternate between
    full attention and sliding-window attention per ``hf_config.layer_types``.
    Inheriting from :class:`HybridAttentionForCausalLM` opts into vLLM's
    hybrid kv cache manager so sliding-window layers can index a smaller
    paged pool than full-attention layers, recovering the asymmetric-hybrid
    memory waste described in vLLM's hybrid kv cache manager design.
    The bridge accepts ``page_tables_per_layer`` from the plugin (one tensor
    per decoder layer, layer-aligned with the underlying TT model's
    ``self.layers``) and stashes it on each TT model handle as
    ``_active_page_tables_per_layer`` so the model's ``ttnn_prefill_forward``
    / ``ttnn_decode_forward`` pick it up without us having to thread the
    kwarg through every call site in :class:`Generator`. The attribute is
    cleared on the way out so a subsequent legacy single-page-table call
    isn't accidentally affected.
    """

    decode_input_update_contract = 1

    # Class-level capabilities
    model_capabilities = {
        "supports_prefix_caching": False,  # Sliding window => no prefix caching
        "supports_async_decode": True,
        "supports_sample_on_device": True,
        "max_device_top_k": 32,
    }

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def prefill_forward(self, *args, page_tables_per_layer=None, **kwargs):
        # While hybrid KV cache groups are disabled (one full-attention group
        # for every layer), the per-layer page-table routing inside this
        # bridge is buggy for users_row_sharded models: it shards page tables
        # naively by mesh row, which doesn't match the gpt-oss
        # slot // max_local_batch_size → row mapping and produces null
        # content on the rows whose page-table chunks point at the wrong
        # KV blocks. Until a row-aware per-layer routing lands, skip the
        # hybrid path entirely and let the legacy single page_table flow
        # through Generator.prefill_forward_text reach the model untouched.
        if not self._HYBRID_KV_CACHE_GROUPS_ENABLED:
            return super().prefill_forward_text(*args, **kwargs)
        page_tables_per_layer = self._ensure_page_tables_per_layer(page_tables_per_layer, kwargs.get("page_table"))
        per_submesh = self._chunk_page_tables_per_dp(page_tables_per_layer)
        # See ``Gemma3ForConditionalGeneration.prefill_forward`` for why
        # the persistent-buffer update has to happen *before* the inner
        # decode/prefill path that may run captured traces.
        if per_submesh is not None:
            for m, pt_for_submesh in zip(self.model, per_submesh):
                m.update_persistent_per_layer_page_tables(pt_for_submesh)
        with self._route_per_layer_page_tables(per_submesh):
            return super().prefill_forward_text(*args, **kwargs)

    def decode_forward(self, *args, page_tables_per_layer=None, **kwargs):
        # See prefill_forward note above. Skip the hybrid path while
        # _HYBRID_KV_CACHE_GROUPS_ENABLED is False.
        if not self._HYBRID_KV_CACHE_GROUPS_ENABLED:
            return super(HybridAttentionForCausalLM, self).decode_forward(*args, **kwargs)
        page_tables_per_layer = self._ensure_page_tables_per_layer(page_tables_per_layer, kwargs.get("page_table"))
        per_submesh = self._chunk_page_tables_per_dp(page_tables_per_layer)
        if per_submesh is not None and self._reload_per_layer_page_tables(kwargs):
            for m, pt_for_submesh in zip(self.model, per_submesh):
                m.update_persistent_per_layer_page_tables(pt_for_submesh)
        with self._route_per_layer_page_tables(per_submesh):
            # Skip ``HybridAttentionForCausalLM.decode_forward``, which is a
            # NotImplementedError placeholder; route to ``Generator``'s
            # actual decode implementation.
            return super(HybridAttentionForCausalLM, self).decode_forward(*args, **kwargs)

    @classmethod
    def initialize_vllm_model(
        cls,
        hf_config,
        mesh_device,
        max_batch_size,
        max_seq_len,
        n_layers=None,
        tt_data_parallel=1,
        optimizations: str = None,
    ):
        assert optimizations is None, "Custom optimizations are not supported for this model"
        from models.demos.gpt_oss.tt.common import create_tt_model

        model_args = []
        model = []
        state_dict = None
        # GPT-OSS throughput profile uses user-row sharding on
        # multi-row meshes with large max batch sizes (e.g., 128 on 4x8).
        # This must be selected at model init time to ensure correct sharding
        # and input preparation.
        users_row_sharded = bool(mesh_device.shape[0] > 1 and max_batch_size > 32)
        if users_row_sharded:
            # For users_row_sharded, we internally manage DP=4 in attention so we don't need to create submeshes
            tt_data_parallel = 1
        submesh_devices = create_submeshes(mesh_device, tt_data_parallel)
        for submesh in submesh_devices:
            # Use the existing create_tt_model function
            model_args_i, model_i, _, state_dict = create_tt_model(
                mesh_device=submesh,
                max_batch_size=max_batch_size // tt_data_parallel,
                max_seq_len=max_seq_len,
                paged_attention_config=None,
                dtype=ttnn.bfloat8_b,
                state_dict=state_dict,
                num_layers=n_layers,
                mesh_config=None,
                create_kv_cache=False,
                users_row_sharded=users_row_sharded,
                use_throughput_experts=submesh.shape[0] > 1 and (max_batch_size > 1),
            )

            model_args.append(model_args_i)
            model.append(model_i)

        return cls(model, model_args, mesh_device)

    @property
    def cache_path(self):
        return self.model_args[0].weight_cache_path(ttnn.bfloat8_b)

    # prefill_forward / decode_forward are defined above with the
    # per-layer page-table stash; allocate_kv_cache_per_layer is inherited
    # from HybridAttentionForCausalLM.
    def allocate_kv_cache(self, *args, **kwargs):
        return allocate_vllm_kv_cache(*args, **kwargs, dp_model=self.model, tt_cache_path=self.cache_path)