From d3b4f84d3a7220833908d99b494dd38cccdfe3d6 Mon Sep 17 00:00:00 2001
From: Kazuaki Ishizaki <ishizaki@jp.ibm.com>
Date: Tue, 7 Jul 2026 07:49:28 +0000
Subject: [PATCH 1/5] initial commit

Signed-off-by: Kazuaki Ishizaki <ishizaki@jp.ibm.com>
---
 .../auto_generate_module_config.py            | 249 +++++++++++++++++-
 1 file changed, 248 insertions(+), 1 deletion(-)

diff --git a/utils/module_discovery/auto_generate_module_config.py b/utils/module_discovery/auto_generate_module_config.py
index e32182ce..4e5e7a4c 100644
--- a/utils/module_discovery/auto_generate_module_config.py
+++ b/utils/module_discovery/auto_generate_module_config.py
@@ -184,6 +184,150 @@ def _process_pytree_structure(value: Any, name: str) -> Dict[str, Any] | None:
     return None
 
 
+def _resolve_layer_idx(module: Any) -> int | None:
+    """Find a module's decoder layer index for indexing into the KV cache.
+
+    In transformers >=5 the ``layer_idx`` lives on the attention submodule, not
+    on the DecoderLayer itself, so we check the layer first (older layouts /
+    other archs) and fall back to ``self_attn.layer_idx``. Returns ``None`` when
+    neither exists (e.g. a norm/MLP module that never touches the KV cache).
+    """
+    layer_idx = getattr(module, "layer_idx", None)
+    if layer_idx is not None:
+        return layer_idx
+    self_attn = getattr(module, "self_attn", None)
+    if self_attn is not None:
+        return getattr(self_attn, "layer_idx", None)
+    return None
+
+
+def _extract_cache_info(
+    past_key_values: Any, name: str, layer_idx: int, config: Any = None
+) -> Dict[str, Any] | None:
+    """Snapshot one layer's populated K/V from a ``Cache`` for a decode step.
+
+    A DecoderLayer receives ``past_key_values`` as a live
+    :class:`~transformers.cache_utils.Cache`, not raw tensors. During prefill
+    the layer's slot is empty (``keys is None``), so there is nothing to record
+    and this returns ``None`` — the layer then runs its
+    ``if past_key_values is not None: past_key_values.update(...)`` branch on
+    freshly computed K/V only, which is the correct prefill behaviour. During
+    decode the slot already holds ``past_len`` tokens; we snapshot that layer's
+    ``keys``/``values`` so the module test can rebuild an equivalent Cache and
+    exercise the same "attend over past + new token" path. Without this, the
+    decode invocation would replay with ``past_key_values=None`` and silently
+    degrade to a 1-token self-attention (the ``update`` branch never runs).
+
+    Transformers >=5 stores per-layer K/V under ``.layers[i].keys/.values``
+    (the older flat ``key_cache``/``value_cache`` lists are gone).
+
+    Only :class:`~transformers.cache_utils.StaticCache` is recorded. A
+    fixed-size StaticCache has a fully specified per-layer K/V shape that the
+    test side can reconstruct deterministically; a growable ``DynamicCache``
+    (the default when no cache is passed in) has no such fixed shape, so we warn
+    and skip it rather than emit a cache the test cannot faithfully rebuild.
+    Drive the generator with an explicit StaticCache to capture decode state.
+
+    Args:
+        past_key_values: The live cache passed to the DecoderLayer.
+        name: The kwarg name (``"past_key_values"``).
+        layer_idx: The layer whose K/V slot to snapshot.
+        config: The model config the cache was built from. In transformers >=5
+            a StaticCache no longer exposes ``.config``, but its ``__init__``
+            requires one, so the test side needs ``config_path`` +
+            ``config_kwargs`` to reconstruct it. Pass the DecoderLayer's config
+            (e.g. ``module.self_attn.config``).
+
+    Returns:
+        A cache spec dict, or ``None`` when the cache is not a StaticCache, this
+        layer's slot is empty (prefill), or it exposes no usable per-layer K/V.
+    """
+    cache_cls = type(past_key_values)
+    if cache_cls.__name__ != "StaticCache":
+        logger.warning(
+            "past_key_values is a %s, not StaticCache; skipping cache capture. "
+            "Drive the generator with an explicit StaticCache to record the "
+            "decode KV state (see generate_gpt_oss_20b_config.py).",
+            cache_cls.__name__,
+        )
+        return None
+
+    layers = getattr(past_key_values, "layers", None)
+    if layers is None or layer_idx >= len(layers):
+        return None
+
+    layer_cache = layers[layer_idx]
+    keys = getattr(layer_cache, "keys", None)
+    values = getattr(layer_cache, "values", None)
+
+    # Empty slot -> prefill call; nothing populated to record.
+    if not isinstance(keys, torch.Tensor) or not isinstance(values, torch.Tensor):
+        return None
+
+    # StaticCache allocates the full max_cache_len up front, so keys/values are
+    # [B, num_kv_heads, max_cache_len, head_dim] with only the first past_len
+    # positions populated. Record just that populated slice so the shape means
+    # "the real past" and the test side can prime a cache by a single update()
+    # of past_len tokens (not the whole fixed allocation of mostly-zeros).
+    #
+    # Use the PER-LAYER length, not past_key_values.get_seq_length(): at a
+    # decode step the whole-cache length reflects layers already updated this
+    # pass, so for layer i>0 (whose slot hasn't been updated yet at pre-hook
+    # time) it reads one token too long. layer_cache.get_seq_length() reports
+    # just this layer's populated past, which is the same across layers.
+    try:
+        past_len = int(layer_cache.get_seq_length())
+    except Exception:
+        try:
+            past_len = int(past_key_values.get_seq_length())
+        except Exception:
+            past_len = keys.shape[-2]
+    if past_len <= 0:
+        return None
+    keys = keys[:, :, :past_len, :]
+    values = values[:, :, :past_len, :]
+
+    cache_info: Dict[str, Any] = {
+        "name": name,
+        "type": "cache",
+        "cache_path": f"{cache_cls.__module__}.{cache_cls.__name__}",
+        "layer_idx": layer_idx,
+        # StaticCache.__init__ needs max_cache_len (the fixed allocation), which
+        # is not derivable from the (sliced) K/V shape. Record it so the test
+        # side rebuilds a cache of the same allocation before priming it.
+        "max_cache_len": getattr(past_key_values, "max_cache_len", None),
+        # keys/values carry real past tokens; the test rebuilds a cache of the
+        # same seq length via update(). "key"/"value" are not special tensor
+        # names, so they default to random init (see _is_special_tensor).
+        "key": _extract_tensor_info(keys, f"{name}_key"),
+        "value": _extract_tensor_info(values, f"{name}_value"),
+    }
+
+    # Snapshot the config so the test side can construct the concrete Cache with
+    # matching dimensions (num_kv_heads, head_dim, ...). transformers >=5 no
+    # longer exposes StaticCache.config, so we take the config passed in from
+    # the DecoderLayer; fall back to the cache's own attribute for older builds.
+    if config is None:
+        config = getattr(past_key_values, "config", None)
+    if config is not None:
+        config_cls = type(config)
+        config_kwargs = {}
+        for attr in [
+            "hidden_size",
+            "num_attention_heads",
+            "num_key_value_heads",
+            "head_dim",
+            "num_hidden_layers",
+            "max_position_embeddings",
+        ]:
+            if hasattr(config, attr):
+                config_kwargs[attr] = getattr(config, attr)
+        cache_info["config_path"] = f"{config_cls.__module__}.{config_cls.__name__}"
+        cache_info["config_kwargs"] = config_kwargs
+
+    return cache_info
+
+
 class ModuleInfoCapture:
     """Captures module information during forward pass using hooks."""
 
@@ -376,7 +520,23 @@ def hook(module, args, kwargs):
             # Analyze keyword arguments using pytree
             for key, value in kwargs.items():
                 if key in ("past_key_values", "past_key_value"):
-                    continue  # Skip - not needed for module-level tests
+                    # A live Cache object can't go through the tensor-spec
+                    # pytree path. For a decode step we snapshot this layer's
+                    # populated K/V so the module test can rebuild an equivalent
+                    # cache and drive the real "attend over past + new token"
+                    # path; for prefill the slot is empty and this records
+                    # nothing (equivalent to past_key_values=None).
+                    layer_idx = _resolve_layer_idx(module)
+                    if layer_idx is not None and value is not None:
+                        layer_config = getattr(
+                            getattr(module, "self_attn", None), "config", None
+                        ) or getattr(module, "config", None)
+                        cache_info = _extract_cache_info(
+                            value, "past_key_values", layer_idx, config=layer_config
+                        )
+                        if cache_info is not None:
+                            invocation_inputs.append(cache_info)
+                    continue
                 input_info = _process_pytree_structure(value, key)
                 if input_info:
                     invocation_inputs.append(input_info)
@@ -480,6 +640,15 @@ def _extract_pattern(input_info: Dict[str, Any]) -> Dict[str, Any]:
             - Single tensor: {"name": "arg_0", "shape": [...], "dtype": ..., ...}
             - Container: {"name": "arg_0", "type": "list/tuple/dict/pytree", "items": [...]}
             """
+            # A KV cache: distinct pattern so prefill (no cache) and decode
+            # (cache present) never collapse into one invocation signature.
+            if input_info.get("type") == "cache":
+                return {
+                    "type": "cache",
+                    "cache_path": input_info.get("cache_path"),
+                    "key_shape": input_info.get("key", {}).get("shape"),
+                    "value_shape": input_info.get("value", {}).get("shape"),
+                }
             # Check if this is a container with items
             if "type" in input_info and "items" in input_info:
                 # Container (list, tuple, dict, pytree)
@@ -655,10 +824,80 @@ def _convert_captured_input_to_sample_input(inp_spec: Dict[str, Any]) -> Dict[st
 
         return {"tensor_list": tensor_list}
 
+    elif inp_type == "cache":
+        # A KV cache: emit cache_path + per-layer key/value tensor specs so the
+        # test side can rebuild a concrete Cache and prime it via update(),
+        # reproducing the decode path (attend over past + new token).
+        cache_spec: Dict[str, Any] = {
+            "cache_path": inp_spec["cache_path"],
+            "layer_idx": inp_spec["layer_idx"],
+            "key": _tensor_info_to_spec(inp_spec["key"], f"{inp_name}_key"),
+            "value": _tensor_info_to_spec(inp_spec["value"], f"{inp_name}_value"),
+        }
+        if inp_spec.get("max_cache_len") is not None:
+            cache_spec["max_cache_len"] = inp_spec["max_cache_len"]
+        if "config_path" in inp_spec:
+            cache_spec["config_path"] = inp_spec["config_path"]
+            cache_spec["config_kwargs"] = inp_spec.get("config_kwargs", {})
+        return {"cache": cache_spec}
+
     else:
         return {"value": None}
 
 
+def _validate_cache_mask_consistency(
+    invocation_inputs: List[Dict[str, Any]], module_name: str
+) -> None:
+    """Warn if a cached (decode) invocation lacks a mask that can cover the past.
+
+    When an invocation carries a KV cache, the test side rebuilds a Cache primed
+    with ``past_len`` tokens and drives a decode forward. That forward also needs
+    an ``attention_mask`` whose key/value axis is at least ``past_len`` (the
+    cache's populated length) — otherwise the mask and the cache disagree about
+    how many past tokens exist and the replayed decode attends over the wrong
+    span. This is a generation-time sanity check (logged, not fatal) so a
+    malformed invocation is visible rather than silently emitted.
+
+    K/V key shape is ``[B, num_kv_heads, head_dim, past_len]`` (past_len last).
+    A 4-D ``attention_mask`` is ``[B, 1, q_len, kv_len]`` (kv_len last). We only
+    require ``past_len <= kv_len`` since a fixed-length cache (e.g. StaticCache)
+    reports its allocation, not its populated length, in the mask.
+    """
+    cache_spec = None
+    mask_spec = None
+    for inp in invocation_inputs:
+        if inp.get("type") == "cache":
+            cache_spec = inp
+        elif inp.get("name") == "attention_mask":
+            mask_spec = inp
+
+    if cache_spec is None:
+        return  # prefill invocation — nothing to check
+
+    if mask_spec is None:
+        logger.warning(
+            "%s: decode invocation has a KV cache but no attention_mask; "
+            "the replayed decode cannot mask the cached past correctly.",
+            module_name,
+        )
+        return
+
+    key_shape = cache_spec.get("key", {}).get("shape")
+    mask_shape = mask_spec.get("shape")
+    if not key_shape or not mask_shape:
+        return
+    past_len = key_shape[-1]
+    kv_len = mask_shape[-1]
+    if kv_len < past_len:
+        logger.warning(
+            "%s: attention_mask kv_len=%d < cached past_len=%d; mask cannot "
+            "cover the cached past for the decode step.",
+            module_name,
+            kv_len,
+            past_len,
+        )
+
+
 def _build_module_entry_dict(module_info: Dict[str, Any]) -> Dict[str, Any]:
     """
     Build a module entry dictionary for YAML generation.
@@ -708,6 +947,14 @@ def _build_module_entry_dict(module_info: Dict[str, Any]) -> Dict[str, Any]:
             else:
                 forward_kwargs[inp_name] = converted
 
+        # A decode invocation carries a KV cache; verify it also carries an
+        # attention_mask whose key/value length can cover the cached past, so
+        # the test side rebuilds a self-consistent (mask, cache) pair rather
+        # than a decode step that silently attends over the wrong span.
+        _validate_cache_mask_consistency(
+            invocation_inputs, module_info.get("name", "<unknown>")
+        )
+
         forward_inputs_list.append(
             {
                 "args": forward_args if forward_args else [],

From fe9c90225af884407b39ed75c58cf5fcd0e7f011 Mon Sep 17 00:00:00 2001
From: Kazuaki Ishizaki <ishizaki@jp.ibm.com>
Date: Thu, 6 Aug 2026 04:53:56 +0000
Subject: [PATCH 2/5] add URL to class definition

Signed-off-by: Kazuaki Ishizaki <ishizaki@jp.ibm.com>
---
 .../auto_generate_module_config.py            | 101 +++++++++++++++++-
 1 file changed, 98 insertions(+), 3 deletions(-)

diff --git a/utils/module_discovery/auto_generate_module_config.py b/utils/module_discovery/auto_generate_module_config.py
index 4e5e7a4c..a3c25e3e 100644
--- a/utils/module_discovery/auto_generate_module_config.py
+++ b/utils/module_discovery/auto_generate_module_config.py
@@ -14,10 +14,12 @@
 
 import argparse
 import hashlib
+import inspect
 import json
 import logging
+import os
 from pathlib import Path
-from typing import Any, Dict, List, Set, Tuple
+from typing import Any, Dict, List, Optional, Set, Tuple
 
 import torch
 import yaml
@@ -201,6 +203,83 @@ def _resolve_layer_idx(module: Any) -> int | None:
     return None
 
 
+def _class_source_location(cls: type) -> Tuple[Optional[str], Optional[int]]:
+    """Return (source file, first line of the class definition) for ``cls``.
+
+    Resolved from the live class object while the hook still holds the module
+    instance. ``module_path`` alone is not enough: a model loaded with
+    ``trust_remote_code`` lives in a dynamically created module that cannot be
+    re-imported by name later. Returns ``(None, None)`` when no source is
+    retrievable (C extension, class synthesized at runtime).
+    """
+    try:
+        source_file = inspect.getsourcefile(cls)
+        _, lineno = inspect.getsourcelines(cls)
+    except (OSError, TypeError):
+        return None, None
+    return source_file, lineno
+
+
+def _shorten_source_path(path: str) -> str:
+    """Trim an absolute source path down to something environment-independent.
+
+    A ``site-packages``/``dist-packages`` install becomes ``<pkg>/...`` so the
+    generated YAML does not hard-code the generating machine's venv layout.
+    Paths outside a site install are returned unchanged.
+    """
+    parts = Path(path).parts
+    for marker in ("site-packages", "dist-packages"):
+        if marker in parts:
+            return str(Path(*parts[parts.index(marker) + 1 :]))
+    return path
+
+
+def _get_transformers_ref() -> str:
+    """Git ref used in generated transformers source URLs.
+
+    Mirrors ``utils/model_ops/utils/torchop_yaml.py``: ``TRANSFORMERS_VERSION``
+    overrides, otherwise the installed version becomes a ``vX.Y.Z`` release tag.
+    A dev/editable install ("5.0.0.dev0") has no such tag, so it falls back to
+    ``main`` rather than emitting a dead link.
+    """
+    version = os.getenv("TRANSFORMERS_VERSION")
+    if version:
+        return version
+    try:
+        import transformers
+    except ImportError:
+        return "main"
+    return (
+        "main" if "dev" in transformers.__version__ else f"v{transformers.__version__}"
+    )
+
+
+_TRANSFORMERS_BLOB_URL = "https://github.com/huggingface/transformers/blob"
+
+
+def _source_reference(
+    source_file: Optional[str], lineno: Optional[int]
+) -> Optional[str]:
+    """Render a captured source location as a human-followable reference.
+
+    A file inside an installed ``transformers`` package becomes a GitHub blob
+    URL pinned to the installed version, matching the scheme
+    ``torchop_yaml._convert_transformers_path_to_url`` uses. Any other package
+    (torch, vLLM, a trust_remote_code module) degrades to a venv-relative
+    ``path:line``, since there is no single upstream repo to point at.
+    """
+    if not source_file:
+        return None
+    rel = _shorten_source_path(source_file)
+    # rel differing from the input means the file came from a site install, so
+    # a leading "transformers/" component is the installed transformers package
+    # (and maps onto src/transformers/... in the upstream repo layout).
+    if rel != source_file and rel.startswith("transformers/"):
+        anchor = f"#L{lineno}" if lineno else ""
+        return f"{_TRANSFORMERS_BLOB_URL}/{_get_transformers_ref()}/src/{rel}{anchor}"
+    return f"{rel}:{lineno}" if lineno else rel
+
+
 def _extract_cache_info(
     past_key_values: Any, name: str, layer_idx: int, config: Any = None
 ) -> Dict[str, Any] | None:
@@ -497,10 +576,14 @@ def hook(module, args, kwargs):
             if unique_module_name not in self.module_data:
                 self.seen_module_configs.add(config_signature)
 
+                cls = module.__class__
+                source_file, source_lineno = _class_source_location(cls)
                 self.module_data[unique_module_name] = {
                     "name": unique_module_name,
                     "module_type": module_type,
-                    "module_path": f"{module.__class__.__module__}.{module.__class__.__name__}",
+                    "module_path": f"{cls.__module__}.{cls.__name__}",
+                    "source_file": source_file,
+                    "source_lineno": source_lineno,
                     "example_instance": module_name,
                     "constructor_args": constructor_info["constructor_args"],
                     "constructor_kwargs": constructor_info["constructor_kwargs"],
@@ -964,11 +1047,23 @@ def _build_module_entry_dict(module_info: Dict[str, Any]) -> Dict[str, Any]:
 
     forward_inputs = forward_inputs_list
 
+    # Record where the class is defined so a reader of the generated YAML can
+    # jump straight to the source. Appended to the free-text description rather
+    # than emitted as its own key, so the entry stays within the shape the OOT
+    # framework's include schema accepts. Absent for captures that carry no
+    # source location (the vLLM generator builds module_info dicts by hand).
+    description = f"Module: {module_info['module_path']}"
+    location = _source_reference(
+        module_info.get("source_file"), module_info.get("source_lineno")
+    )
+    if location:
+        description = f"{description} (defined at {location})"
+
     # Build module entry
     entry = {
         "name": module_info["name"],
         "module_path": module_info["module_path"],
-        "description": f"Module: {module_info['module_path']}",
+        "description": description,
         "constructor_inputs": {
             "args": constructor_args if constructor_args else [],
             "kwargs": constructor_kwargs if constructor_kwargs else {},

From 064b8238b48199a5c24c265d0b6889793ff0ccd4 Mon Sep 17 00:00:00 2001
From: Kazuaki Ishizaki <ishizaki@jp.ibm.com>
Date: Fri, 7 Aug 2026 03:17:51 +0000
Subject: [PATCH 3/5] skip DecoderLayer

Signed-off-by: Kazuaki Ishizaki <ishizaki@jp.ibm.com>
---
 tests/test_modules_custom.py | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/tests/test_modules_custom.py b/tests/test_modules_custom.py
index 74e77f35..e664449c 100644
--- a/tests/test_modules_custom.py
+++ b/tests/test_modules_custom.py
@@ -451,6 +451,13 @@ def test_eager_vs_compile(self, device, dtype, module_info, training):
         3. Runs module in compile mode on Spyre
         4. Compares outputs between eager CPU, eager Spyre, and compile Spyre
         """
+
+        cls_name = module_info.module_cls.__name__
+        if cls_name.endswith("DecoderLayer"):
+            self.skipTest(
+                f"{cls_name}: skip the whole decoder layer due to time-consuming module"
+            )
+
         module_inputs = module_info.module_inputs_func(
             module_info, device=device, dtype=dtype, requires_grad=False, training=False
         )
@@ -581,6 +588,13 @@ def test_with_cpu(self, device, dtype, module_info, training):
         """
         run_compile = os.getenv("TEST_COMPILE_WITH_CPU", "1") == "1"
         run_eager = os.getenv("TEST_EAGER_WITH_CPU", "0") == "1"
+
+        cls_name = module_info.module_cls.__name__
+        if cls_name.endswith("DecoderLayer"):
+            self.skipTest(
+                f"{cls_name}: skip the whole decoder layer due to time-consuming module"
+            )
+
         module_inputs = module_info.module_inputs_func(
             module_info,
             device=device,

From ef357530c8bef5c9b206b6b320856bd1469951ae Mon Sep 17 00:00:00 2001
From: Kazuaki Ishizaki <ishizaki@jp.ibm.com>
Date: Fri, 7 Aug 2026 03:18:44 +0000
Subject: [PATCH 4/5] add Attention and DecoderLayer modules

Signed-off-by: Kazuaki Ishizaki <ishizaki@jp.ibm.com>
---
 .../Ministral-3-14B-Instruct-2512.yaml        | 192 ++++++++----
 .../Mistral-Small-3.2-24B-Instruct-2506.yaml  | 276 ++++++++++--------
 .../granite_3_3_8b_instruct_spyre.yaml        | 144 +++++++--
 .../module_tests/granite_4_1_8b_spyre.yaml    | 214 ++++++++------
 4 files changed, 524 insertions(+), 302 deletions(-)

diff --git a/tests/configs/module_tests/Ministral-3-14B-Instruct-2512.yaml b/tests/configs/module_tests/Ministral-3-14B-Instruct-2512.yaml
index 40c9601b..7fef1034 100644
--- a/tests/configs/module_tests/Ministral-3-14B-Instruct-2512.yaml
+++ b/tests/configs/module_tests/Ministral-3-14B-Instruct-2512.yaml
@@ -1,10 +1,7 @@
 # Auto-generated unified test configuration for Ministral_3_14B_Instruct_2512
 # Generated by auto_generate_module_config.py
 # Format compatible with PyTorch's test_modules.py (using edits.modules.include)
-# NOTE: hidden_size is temporarily set to 4096 (diagnostic) instead of the
-# model's real 5120 to test whether a non-power-of-2 reduction dimension is
-# causing sdsc_fused_mean kernel timeouts on Spyre. Revert to 5120 per
-# https://huggingface.co/mistralai/Ministral-3-14B-Instruct-2512/blob/main/config.json#L30
+
 test_suite_config:
   files:
     - path: ${TORCH_ROOT}/test/test_modules.py
@@ -21,17 +18,16 @@ test_suite_config:
               include: &id001
                 - name: Ministral3Model_d7945082
                   module_path: transformers.models.ministral3.modeling_ministral3.Ministral3Model
-                  description: 'Module: transformers.models.ministral3.modeling_ministral3.Ministral3Model'
+                  description: 'Module: transformers.models.ministral3.modeling_ministral3.Ministral3Model (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/ministral3/modeling_ministral3.py#L338)'
                   constructor_inputs:
                     args:
                       - config_path: transformers.models.ministral3.configuration_ministral3.Ministral3Config
                         config_kwargs:
-                          hidden_size: 4096
+                          hidden_size: 5120
                           num_attention_heads: 32
                           num_key_value_heads: 8
                           intermediate_size: 16384
                           max_position_embeddings: 262144
-                          num_hidden_layers: 1  # default is 40; capped to avoid eager-on-Spyre timeout in test_eager_vs_compile
                           _attn_implementation: sdpa
                     kwargs: {}
                   forward_inputs:
@@ -46,15 +42,15 @@ test_suite_config:
                             device: spyre
                             init: randint
                             init_args:
-                              high: 10000
+                              high: 1
                         inputs_embeds:
                           tensor:
-                            shape: [1, 128, 4096]
+                            shape: [1, 128, 5120]
                             stride: null
                             storage_offset: 0
                             dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                     - args: []
                       kwargs:
                         attention_mask:
@@ -66,23 +62,23 @@ test_suite_config:
                             device: spyre
                             init: randint
                             init_args:
-                              high: 10000
+                              high: 1
                         inputs_embeds:
                           tensor:
-                            shape: [1, 1, 4096]
+                            shape: [1, 1, 5120]
                             stride: null
                             storage_offset: 0
                             dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                 - name: Ministral3RotaryEmbedding_483e0b9c
                   module_path: transformers.models.ministral3.modeling_ministral3.Ministral3RotaryEmbedding
-                  description: 'Module: transformers.models.ministral3.modeling_ministral3.Ministral3RotaryEmbedding'
+                  description: 'Module: transformers.models.ministral3.modeling_ministral3.Ministral3RotaryEmbedding (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/ministral3/modeling_ministral3.py#L273)'
                   constructor_inputs:
                     args:
                       - config_path: transformers.models.ministral3.configuration_ministral3.Ministral3Config
                         config_kwargs:
-                          hidden_size: 4096
+                          hidden_size: 5120
                           num_attention_heads: 32
                           num_key_value_heads: 8
                           intermediate_size: 16384
@@ -92,12 +88,12 @@ test_suite_config:
                   forward_inputs:
                     - args:
                         - tensor:
-                            shape: [1, 128, 4096]
+                            shape: [1, 128, 5120]
                             stride: null
                             storage_offset: 0
                             dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs:
                         position_ids:
                           tensor:
@@ -108,15 +104,15 @@ test_suite_config:
                             device: spyre
                             init: randint
                             init_args:
-                              high: 10000
+                              high: 1
                     - args:
                         - tensor:
-                            shape: [1, 1, 4096]
+                            shape: [1, 1, 5120]
                             stride: null
                             storage_offset: 0
                             dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs:
                         position_ids:
                           tensor:
@@ -127,15 +123,15 @@ test_suite_config:
                             device: spyre
                             init: randint
                             init_args:
-                              high: 10000
+                              high: 1
                 - name: Ministral3DecoderLayer_layer0
                   module_path: transformers.models.ministral3.modeling_ministral3.Ministral3DecoderLayer
-                  description: 'Module: transformers.models.ministral3.modeling_ministral3.Ministral3DecoderLayer'
+                  description: 'Module: transformers.models.ministral3.modeling_ministral3.Ministral3DecoderLayer (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/ministral3/modeling_ministral3.py#L213)'
                   constructor_inputs:
                     args:
                       - config_path: transformers.models.ministral3.configuration_ministral3.Ministral3Config
                         config_kwargs:
-                          hidden_size: 4096
+                          hidden_size: 5120
                           num_attention_heads: 32
                           num_key_value_heads: 8
                           intermediate_size: 16384
@@ -146,12 +142,12 @@ test_suite_config:
                   forward_inputs:
                     - args:
                         - tensor:
-                            shape: [1, 128, 4096]
+                            shape: [1, 128, 5120]
                             stride: null
                             storage_offset: 0
                             dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs:
                         position_ids:
                           tensor:
@@ -162,7 +158,7 @@ test_suite_config:
                             device: spyre
                             init: randint
                             init_args:
-                              high: 10000
+                              high: 1
                         position_embeddings:
                           tensor_list:
                             - shape: [1, 128, 128]
@@ -172,7 +168,7 @@ test_suite_config:
                               device: spyre
                               init: randint
                               init_args:
-                                high: 10000
+                                high: 1
                             - shape: [1, 128, 128]
                               stride: null
                               storage_offset: 0
@@ -180,16 +176,26 @@ test_suite_config:
                               device: spyre
                               init: randint
                               init_args:
-                                high: 10000
+                                high: 1
                     - args:
                         - tensor:
-                            shape: [1, 1, 4096]
+                            shape: [1, 1, 5120]
                             stride: null
                             storage_offset: 0
                             dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs:
+                        attention_mask:
+                          tensor:
+                            shape: [1, 1, 1, 2048]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bool
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
                         position_ids:
                           tensor:
                             shape: [1, 1]
@@ -199,7 +205,34 @@ test_suite_config:
                             device: spyre
                             init: randint
                             init_args:
-                              high: 10000
+                              high: 1
+                        past_key_values:
+                          cache:
+                            cache_path: transformers.cache_utils.StaticCache
+                            layer_idx: 0
+                            key:
+                              shape: [1, 8, 128, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: xavier
+                            value:
+                              shape: [1, 8, 128, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: xavier
+                            max_cache_len: 2048
+                            config_path: transformers.models.ministral3.configuration_ministral3.Ministral3Config
+                            config_kwargs:
+                              hidden_size: 5120
+                              num_attention_heads: 32
+                              num_key_value_heads: 8
+                              head_dim: 128
+                              num_hidden_layers: 40
+                              max_position_embeddings: 262144
                         position_embeddings:
                           tensor_list:
                             - shape: [1, 1, 128]
@@ -209,7 +242,7 @@ test_suite_config:
                               device: spyre
                               init: randint
                               init_args:
-                                high: 10000
+                                high: 1
                             - shape: [1, 1, 128]
                               stride: null
                               storage_offset: 0
@@ -217,41 +250,41 @@ test_suite_config:
                               device: spyre
                               init: randint
                               init_args:
-                                high: 10000
-                - name: Ministral3RMSNorm_4096
+                                high: 1
+                - name: Ministral3RMSNorm_5120
                   module_path: transformers.models.ministral3.modeling_ministral3.Ministral3RMSNorm
-                  description: 'Module: transformers.models.ministral3.modeling_ministral3.Ministral3RMSNorm'
+                  description: 'Module: transformers.models.ministral3.modeling_ministral3.Ministral3RMSNorm (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/ministral3/modeling_ministral3.py#L192)'
                   constructor_inputs:
                     args:
-                      - value: 4096
+                      - value: 5120
                     kwargs: {}
                   forward_inputs:
                     - args:
                         - tensor:
-                            shape: [1, 128, 4096]
+                            shape: [1, 128, 5120]
                             stride: null
                             storage_offset: 0
                             dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs: {}
                     - args:
                         - tensor:
-                            shape: [1, 1, 4096]
+                            shape: [1, 1, 5120]
                             stride: null
                             storage_offset: 0
                             dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs: {}
                 - name: Ministral3Attention_layer0
                   module_path: transformers.models.ministral3.modeling_ministral3.Ministral3Attention
-                  description: 'Module: transformers.models.ministral3.modeling_ministral3.Ministral3Attention'
+                  description: 'Module: transformers.models.ministral3.modeling_ministral3.Ministral3Attention (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/ministral3/modeling_ministral3.py#L110)'
                   constructor_inputs:
                     args:
                       - config_path: transformers.models.ministral3.configuration_ministral3.Ministral3Config
                         config_kwargs:
-                          hidden_size: 4096
+                          hidden_size: 5120
                           num_attention_heads: 32
                           num_key_value_heads: 8
                           intermediate_size: 16384
@@ -264,12 +297,12 @@ test_suite_config:
                       kwargs:
                         hidden_states:
                           tensor:
-                            shape: [1, 128, 4096]
+                            shape: [1, 128, 5120]
                             stride: null
                             storage_offset: 0
                             dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                         position_ids:
                           tensor:
                             shape: [1, 128]
@@ -279,7 +312,7 @@ test_suite_config:
                             device: spyre
                             init: randint
                             init_args:
-                              high: 10000
+                              high: 1
                         position_embeddings:
                           tensor_list:
                             - shape: [1, 128, 128]
@@ -289,7 +322,7 @@ test_suite_config:
                               device: spyre
                               init: randint
                               init_args:
-                                high: 10000
+                                high: 1
                             - shape: [1, 128, 128]
                               stride: null
                               storage_offset: 0
@@ -297,17 +330,27 @@ test_suite_config:
                               device: spyre
                               init: randint
                               init_args:
-                                high: 10000
+                                high: 1
                     - args: []
                       kwargs:
                         hidden_states:
                           tensor:
-                            shape: [1, 1, 4096]
+                            shape: [1, 1, 5120]
                             stride: null
                             storage_offset: 0
                             dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
+                        attention_mask:
+                          tensor:
+                            shape: [1, 1, 1, 2048]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bool
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
                         position_ids:
                           tensor:
                             shape: [1, 1]
@@ -317,7 +360,34 @@ test_suite_config:
                             device: spyre
                             init: randint
                             init_args:
-                              high: 10000
+                              high: 1
+                        past_key_values:
+                          cache:
+                            cache_path: transformers.cache_utils.StaticCache
+                            layer_idx: 0
+                            key:
+                              shape: [1, 8, 128, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: xavier
+                            value:
+                              shape: [1, 8, 128, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: xavier
+                            max_cache_len: 2048
+                            config_path: transformers.models.ministral3.configuration_ministral3.Ministral3Config
+                            config_kwargs:
+                              hidden_size: 5120
+                              num_attention_heads: 32
+                              num_key_value_heads: 8
+                              head_dim: 128
+                              num_hidden_layers: 40
+                              max_position_embeddings: 262144
                         position_embeddings:
                           tensor_list:
                             - shape: [1, 1, 128]
@@ -327,7 +397,7 @@ test_suite_config:
                               device: spyre
                               init: randint
                               init_args:
-                                high: 10000
+                                high: 1
                             - shape: [1, 1, 128]
                               stride: null
                               storage_offset: 0
@@ -335,15 +405,15 @@ test_suite_config:
                               device: spyre
                               init: randint
                               init_args:
-                                high: 10000
+                                high: 1
                 - name: Ministral3MLP_eb806b7d
                   module_path: transformers.models.ministral3.modeling_ministral3.Ministral3MLP
-                  description: 'Module: transformers.models.ministral3.modeling_ministral3.Ministral3MLP'
+                  description: 'Module: transformers.models.ministral3.modeling_ministral3.Ministral3MLP (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/ministral3/modeling_ministral3.py#L176)'
                   constructor_inputs:
                     args:
                       - config_path: transformers.models.ministral3.configuration_ministral3.Ministral3Config
                         config_kwargs:
-                          hidden_size: 4096
+                          hidden_size: 5120
                           num_attention_heads: 32
                           num_key_value_heads: 8
                           intermediate_size: 16384
@@ -353,25 +423,25 @@ test_suite_config:
                   forward_inputs:
                     - args:
                         - tensor:
-                            shape: [1, 128, 4096]
+                            shape: [1, 128, 5120]
                             stride: null
                             storage_offset: 0
                             dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs: {}
                     - args:
                         - tensor:
-                            shape: [1, 1, 4096]
+                            shape: [1, 1, 5120]
                             stride: null
                             storage_offset: 0
                             dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs: {}
                 - name: SiLUActivation_d2f532e9
                   module_path: transformers.activations.SiLUActivation
-                  description: 'Module: transformers.activations.SiLUActivation'
+                  description: 'Module: transformers.activations.SiLUActivation (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/activations.py#L92)'
                   constructor_inputs:
                     args: []
                     kwargs: {}
@@ -383,7 +453,7 @@ test_suite_config:
                             storage_offset: 0
                             dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs: {}
                     - args:
                         - tensor:
@@ -392,7 +462,7 @@ test_suite_config:
                             storage_offset: 0
                             dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs: {}
     - path: ${TORCH_DEVICE_ROOT}/tests/test_modules_custom.py
       unlisted_test_mode: skip
diff --git a/tests/configs/module_tests/Mistral-Small-3.2-24B-Instruct-2506.yaml b/tests/configs/module_tests/Mistral-Small-3.2-24B-Instruct-2506.yaml
index 5b5ba33b..14f5cbc6 100644
--- a/tests/configs/module_tests/Mistral-Small-3.2-24B-Instruct-2506.yaml
+++ b/tests/configs/module_tests/Mistral-Small-3.2-24B-Instruct-2506.yaml
@@ -1,11 +1,6 @@
 # Auto-generated unified test configuration for Mistral_Small_3_2_24B_Instruct_2506
 # Generated by auto_generate_module_config.py
 # Format compatible with PyTorch's test_modules.py (using edits.modules.include)
-#
-# NOTE: hidden_size is temporarily set to 4096 (diagnostic) instead of the
-# model's real 5120 to test whether a non-power-of-2 reduction dimension is
-# causing sdsc_fused_mean kernel timeouts on Spyre. Revert to 5120 per
-# https://huggingface.co/mistralai/Mistral-Small-3.2-24B-Instruct-2506/blob/main/config.json#L14
 
 test_suite_config:
   files:
@@ -23,17 +18,16 @@ test_suite_config:
               include: &id001
                 - name: MistralModel_b95f1163
                   module_path: transformers.models.mistral.modeling_mistral.MistralModel
-                  description: 'Module: transformers.models.mistral.modeling_mistral.MistralModel'
+                  description: 'Module: transformers.models.mistral.modeling_mistral.MistralModel (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/mistral/modeling_mistral.py#L327)'
                   constructor_inputs:
                     args:
                       - config_path: transformers.models.mistral.configuration_mistral.MistralConfig
                         config_kwargs:
-                          hidden_size: 4096
+                          hidden_size: 5120
                           num_attention_heads: 32
                           num_key_value_heads: 8
                           intermediate_size: 32768
                           max_position_embeddings: 131072
-                          num_hidden_layers: 1  # default is 40; capped to avoid eager-on-Spyre timeout in test_eager_vs_compile
                           _attn_implementation: sdpa
                     kwargs: {}
                   forward_inputs:
@@ -48,15 +42,15 @@ test_suite_config:
                             device: spyre
                             init: randint
                             init_args:
-                              high: 10000
+                              high: 1
                         inputs_embeds:
                           tensor:
-                            shape: [1, 128, 4096]
+                            shape: [1, 128, 5120]
                             stride: null
                             storage_offset: 0
-                            dtype: torch.float32
+                            dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                     - args: []
                       kwargs:
                         attention_mask:
@@ -68,23 +62,23 @@ test_suite_config:
                             device: spyre
                             init: randint
                             init_args:
-                              high: 10000
+                              high: 1
                         inputs_embeds:
                           tensor:
-                            shape: [1, 1, 4096]
+                            shape: [1, 1, 5120]
                             stride: null
                             storage_offset: 0
-                            dtype: torch.float32
+                            dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                 - name: MistralRotaryEmbedding_0af204ee
                   module_path: transformers.models.mistral.modeling_mistral.MistralRotaryEmbedding
-                  description: 'Module: transformers.models.mistral.modeling_mistral.MistralRotaryEmbedding'
+                  description: 'Module: transformers.models.mistral.modeling_mistral.MistralRotaryEmbedding (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/mistral/modeling_mistral.py#L262)'
                   constructor_inputs:
                     args:
                       - config_path: transformers.models.mistral.configuration_mistral.MistralConfig
                         config_kwargs:
-                          hidden_size: 4096
+                          hidden_size: 5120
                           num_attention_heads: 32
                           num_key_value_heads: 8
                           intermediate_size: 32768
@@ -94,44 +88,50 @@ test_suite_config:
                   forward_inputs:
                     - args:
                         - tensor:
-                            shape: [1, 128, 4096]
+                            shape: [1, 128, 5120]
                             stride: null
                             storage_offset: 0
-                            dtype: torch.float32
+                            dtype: torch.bfloat16
                             device: spyre
-                            init: randn
-                        - tensor:
+                            init: xavier
+                      kwargs:
+                        position_ids:
+                          tensor:
                             shape: [1, 128]
                             stride: null
                             storage_offset: 0
                             dtype: torch.int64
                             device: spyre
-                            init: randn
-                      kwargs: {}
+                            init: randint
+                            init_args:
+                              high: 1
                     - args:
                         - tensor:
-                            shape: [1, 1, 4096]
+                            shape: [1, 1, 5120]
                             stride: null
                             storage_offset: 0
-                            dtype: torch.float32
+                            dtype: torch.bfloat16
                             device: spyre
-                            init: randn
-                        - tensor:
+                            init: xavier
+                      kwargs:
+                        position_ids:
+                          tensor:
                             shape: [1, 1]
                             stride: null
                             storage_offset: 0
                             dtype: torch.int64
                             device: spyre
-                            init: randn
-                      kwargs: {}
+                            init: randint
+                            init_args:
+                              high: 1
                 - name: MistralDecoderLayer_layer0
                   module_path: transformers.models.mistral.modeling_mistral.MistralDecoderLayer
-                  description: 'Module: transformers.models.mistral.modeling_mistral.MistralDecoderLayer'
+                  description: 'Module: transformers.models.mistral.modeling_mistral.MistralDecoderLayer (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/mistral/modeling_mistral.py#L202)'
                   constructor_inputs:
                     args:
                       - config_path: transformers.models.mistral.configuration_mistral.MistralConfig
                         config_kwargs:
-                          hidden_size: 4096
+                          hidden_size: 5120
                           num_attention_heads: 32
                           num_key_value_heads: 8
                           intermediate_size: 32768
@@ -142,12 +142,12 @@ test_suite_config:
                   forward_inputs:
                     - args:
                         - tensor:
-                            shape: [1, 128, 4096]
+                            shape: [1, 128, 5120]
                             stride: null
                             storage_offset: 0
-                            dtype: torch.float32
+                            dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs:
                         position_ids:
                           tensor:
@@ -158,116 +158,133 @@ test_suite_config:
                             device: spyre
                             init: randint
                             init_args:
-                              high: 10000
-                        cache_position:
-                          tensor:
-                            shape: [128]
-                            stride: null
-                            storage_offset: 0
-                            dtype: torch.int64
-                            device: spyre
-                            init: randint
-                            init_args:
-                              high: 10000
+                              high: 1
                         position_embeddings:
                           tensor_list:
                             - shape: [1, 128, 128]
                               stride: null
                               storage_offset: 0
-                              dtype: torch.float32
+                              dtype: torch.bfloat16
                               device: spyre
                               init: randint
                               init_args:
-                                high: 10000
+                                high: 1
                             - shape: [1, 128, 128]
                               stride: null
                               storage_offset: 0
-                              dtype: torch.float32
+                              dtype: torch.bfloat16
                               device: spyre
                               init: randint
                               init_args:
-                                high: 10000
+                                high: 1
                     - args:
                         - tensor:
-                            shape: [1, 1, 4096]
+                            shape: [1, 1, 5120]
                             stride: null
                             storage_offset: 0
-                            dtype: torch.float32
+                            dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs:
-                        position_ids:
+                        attention_mask:
                           tensor:
-                            shape: [1, 1]
+                            shape: [1, 1, 1, 2048]
                             stride: null
                             storage_offset: 0
-                            dtype: torch.int64
+                            dtype: torch.bool
                             device: spyre
                             init: randint
                             init_args:
-                              high: 10000
-                        cache_position:
+                              high: 1
+                        position_ids:
                           tensor:
-                            shape: [1]
+                            shape: [1, 1]
                             stride: null
                             storage_offset: 0
                             dtype: torch.int64
                             device: spyre
                             init: randint
                             init_args:
-                              high: 10000
+                              high: 1
+                        past_key_values:
+                          cache:
+                            cache_path: transformers.cache_utils.StaticCache
+                            layer_idx: 0
+                            key:
+                              shape: [1, 8, 128, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: xavier
+                            value:
+                              shape: [1, 8, 128, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: xavier
+                            max_cache_len: 2048
+                            config_path: transformers.models.mistral.configuration_mistral.MistralConfig
+                            config_kwargs:
+                              hidden_size: 5120
+                              num_attention_heads: 32
+                              num_key_value_heads: 8
+                              head_dim: 128
+                              num_hidden_layers: 40
+                              max_position_embeddings: 131072
                         position_embeddings:
                           tensor_list:
                             - shape: [1, 1, 128]
                               stride: null
                               storage_offset: 0
-                              dtype: torch.float32
+                              dtype: torch.bfloat16
                               device: spyre
                               init: randint
                               init_args:
-                                high: 10000
+                                high: 1
                             - shape: [1, 1, 128]
                               stride: null
                               storage_offset: 0
-                              dtype: torch.float32
+                              dtype: torch.bfloat16
                               device: spyre
                               init: randint
                               init_args:
-                                high: 10000
-                - name: MistralRMSNorm_4096
+                                high: 1
+                - name: MistralRMSNorm_5120
                   module_path: transformers.models.mistral.modeling_mistral.MistralRMSNorm
-                  description: 'Module: transformers.models.mistral.modeling_mistral.MistralRMSNorm'
+                  description: 'Module: transformers.models.mistral.modeling_mistral.MistralRMSNorm (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/mistral/modeling_mistral.py#L181)'
                   constructor_inputs:
                     args:
-                      - value: 4096
+                      - value: 5120
                     kwargs: {}
                   forward_inputs:
                     - args:
                         - tensor:
-                            shape: [1, 128, 4096]
+                            shape: [1, 128, 5120]
                             stride: null
                             storage_offset: 0
-                            dtype: torch.float32
+                            dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs: {}
                     - args:
                         - tensor:
-                            shape: [1, 1, 4096]
+                            shape: [1, 1, 5120]
                             stride: null
                             storage_offset: 0
-                            dtype: torch.float32
+                            dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs: {}
                 - name: MistralAttention_layer0
                   module_path: transformers.models.mistral.modeling_mistral.MistralAttention
-                  description: 'Module: transformers.models.mistral.modeling_mistral.MistralAttention'
+                  description: 'Module: transformers.models.mistral.modeling_mistral.MistralAttention (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/mistral/modeling_mistral.py#L121)'
                   constructor_inputs:
                     args:
                       - config_path: transformers.models.mistral.configuration_mistral.MistralConfig
                         config_kwargs:
-                          hidden_size: 4096
+                          hidden_size: 5120
                           num_attention_heads: 32
                           num_key_value_heads: 8
                           intermediate_size: 32768
@@ -280,12 +297,12 @@ test_suite_config:
                       kwargs:
                         hidden_states:
                           tensor:
-                            shape: [1, 128, 4096]
+                            shape: [1, 128, 5120]
                             stride: null
                             storage_offset: 0
-                            dtype: torch.float32
+                            dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                         position_ids:
                           tensor:
                             shape: [1, 128]
@@ -295,119 +312,136 @@ test_suite_config:
                             device: spyre
                             init: randint
                             init_args:
-                              high: 10000
-                        cache_position:
-                          tensor:
-                            shape: [128]
-                            stride: null
-                            storage_offset: 0
-                            dtype: torch.int64
-                            device: spyre
-                            init: randint
-                            init_args:
-                              high: 10000
+                              high: 1
                         position_embeddings:
                           tensor_list:
                             - shape: [1, 128, 128]
                               stride: null
                               storage_offset: 0
-                              dtype: torch.float32
+                              dtype: torch.bfloat16
                               device: spyre
                               init: randint
                               init_args:
-                                high: 10000
+                                high: 1
                             - shape: [1, 128, 128]
                               stride: null
                               storage_offset: 0
-                              dtype: torch.float32
+                              dtype: torch.bfloat16
                               device: spyre
                               init: randint
                               init_args:
-                                high: 10000
+                                high: 1
                     - args: []
                       kwargs:
                         hidden_states:
                           tensor:
-                            shape: [1, 1, 4096]
+                            shape: [1, 1, 5120]
                             stride: null
                             storage_offset: 0
-                            dtype: torch.float32
+                            dtype: torch.bfloat16
                             device: spyre
-                            init: randn
-                        position_ids:
+                            init: xavier
+                        attention_mask:
                           tensor:
-                            shape: [1, 1]
+                            shape: [1, 1, 1, 2048]
                             stride: null
                             storage_offset: 0
-                            dtype: torch.int64
+                            dtype: torch.bool
                             device: spyre
                             init: randint
                             init_args:
-                              high: 10000
-                        cache_position:
+                              high: 1
+                        position_ids:
                           tensor:
-                            shape: [1]
+                            shape: [1, 1]
                             stride: null
                             storage_offset: 0
                             dtype: torch.int64
                             device: spyre
                             init: randint
                             init_args:
-                              high: 10000
+                              high: 1
+                        past_key_values:
+                          cache:
+                            cache_path: transformers.cache_utils.StaticCache
+                            layer_idx: 0
+                            key:
+                              shape: [1, 8, 128, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: xavier
+                            value:
+                              shape: [1, 8, 128, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: xavier
+                            max_cache_len: 2048
+                            config_path: transformers.models.mistral.configuration_mistral.MistralConfig
+                            config_kwargs:
+                              hidden_size: 5120
+                              num_attention_heads: 32
+                              num_key_value_heads: 8
+                              head_dim: 128
+                              num_hidden_layers: 40
+                              max_position_embeddings: 131072
                         position_embeddings:
                           tensor_list:
                             - shape: [1, 1, 128]
                               stride: null
                               storage_offset: 0
-                              dtype: torch.float32
+                              dtype: torch.bfloat16
                               device: spyre
                               init: randint
                               init_args:
-                                high: 10000
+                                high: 1
                             - shape: [1, 1, 128]
                               stride: null
                               storage_offset: 0
-                              dtype: torch.float32
+                              dtype: torch.bfloat16
                               device: spyre
                               init: randint
                               init_args:
-                                high: 10000
+                                high: 1
                 - name: MistralMLP_f82d9546
                   module_path: transformers.models.mistral.modeling_mistral.MistralMLP
-                  description: 'Module: transformers.models.mistral.modeling_mistral.MistralMLP'
+                  description: 'Module: transformers.models.mistral.modeling_mistral.MistralMLP (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/mistral/modeling_mistral.py#L35)'
                   constructor_inputs:
                     args:
                       - config_path: transformers.models.mistral.configuration_mistral.MistralConfig
                         config_kwargs:
-                          hidden_size: 4096
+                          hidden_size: 5120
                           num_attention_heads: 32
                           num_key_value_heads: 8
                           intermediate_size: 32768
                           max_position_embeddings: 131072
-                          _attn_implementation: null
+                          _attn_implementation: sdpa
                     kwargs: {}
                   forward_inputs:
                     - args:
                         - tensor:
-                            shape: [1, 128, 4096]
+                            shape: [1, 128, 5120]
                             stride: null
                             storage_offset: 0
-                            dtype: torch.float32
+                            dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs: {}
                     - args:
                         - tensor:
-                            shape: [1, 1, 4096]
+                            shape: [1, 1, 5120]
                             stride: null
                             storage_offset: 0
-                            dtype: torch.float32
+                            dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs: {}
                 - name: SiLUActivation_d2f532e9
                   module_path: transformers.activations.SiLUActivation
-                  description: 'Module: transformers.activations.SiLUActivation'
+                  description: 'Module: transformers.activations.SiLUActivation (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/activations.py#L92)'
                   constructor_inputs:
                     args: []
                     kwargs: {}
@@ -417,18 +451,18 @@ test_suite_config:
                             shape: [1, 128, 32768]
                             stride: null
                             storage_offset: 0
-                            dtype: torch.float32
+                            dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs: {}
                     - args:
                         - tensor:
                             shape: [1, 1, 32768]
                             stride: null
                             storage_offset: 0
-                            dtype: torch.float32
+                            dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs: {}
     - path: ${TORCH_DEVICE_ROOT}/tests/test_modules_custom.py
       unlisted_test_mode: skip
diff --git a/tests/configs/module_tests/granite_3_3_8b_instruct_spyre.yaml b/tests/configs/module_tests/granite_3_3_8b_instruct_spyre.yaml
index ea0d26a7..46456cc9 100644
--- a/tests/configs/module_tests/granite_3_3_8b_instruct_spyre.yaml
+++ b/tests/configs/module_tests/granite_3_3_8b_instruct_spyre.yaml
@@ -18,7 +18,7 @@ test_suite_config:
               include: &id001
                 - name: GraniteRotaryEmbedding_64892356
                   module_path: transformers.models.granite.modeling_granite.GraniteRotaryEmbedding
-                  description: 'Module: transformers.models.granite.modeling_granite.GraniteRotaryEmbedding'
+                  description: 'Module: transformers.models.granite.modeling_granite.GraniteRotaryEmbedding (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/granite/modeling_granite.py#L302)'
                   constructor_inputs:
                     args:
                       - config_path: transformers.models.granite.configuration_granite.GraniteConfig
@@ -38,7 +38,7 @@ test_suite_config:
                             storage_offset: 0
                             dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs:
                         position_ids:
                           tensor:
@@ -49,7 +49,7 @@ test_suite_config:
                             device: spyre
                             init: randint
                             init_args:
-                              high: 10000
+                              high: 1
                     - args:
                         - tensor:
                             shape: [1, 1, 4096]
@@ -57,7 +57,7 @@ test_suite_config:
                             storage_offset: 0
                             dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs:
                         position_ids:
                           tensor:
@@ -68,10 +68,10 @@ test_suite_config:
                             device: spyre
                             init: randint
                             init_args:
-                              high: 10000
+                              high: 1
                 - name: GraniteDecoderLayer_layer0
                   module_path: transformers.models.granite.modeling_granite.GraniteDecoderLayer
-                  description: 'Module: transformers.models.granite.modeling_granite.GraniteDecoderLayer'
+                  description: 'Module: transformers.models.granite.modeling_granite.GraniteDecoderLayer (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/granite/modeling_granite.py#L219)'
                   constructor_inputs:
                     args:
                       - config_path: transformers.models.granite.configuration_granite.GraniteConfig
@@ -92,7 +92,7 @@ test_suite_config:
                             storage_offset: 0
                             dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs:
                         position_ids:
                           tensor:
@@ -103,7 +103,7 @@ test_suite_config:
                             device: spyre
                             init: randint
                             init_args:
-                              high: 10000
+                              high: 1
                         position_embeddings:
                           tensor_list:
                             - shape: [1, 128, 128]
@@ -113,7 +113,7 @@ test_suite_config:
                               device: spyre
                               init: randint
                               init_args:
-                                high: 10000
+                                high: 1
                             - shape: [1, 128, 128]
                               stride: null
                               storage_offset: 0
@@ -121,7 +121,7 @@ test_suite_config:
                               device: spyre
                               init: randint
                               init_args:
-                                high: 10000
+                                high: 1
                     - args:
                         - tensor:
                             shape: [1, 1, 4096]
@@ -129,8 +129,18 @@ test_suite_config:
                             storage_offset: 0
                             dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs:
+                        attention_mask:
+                          tensor:
+                            shape: [1, 1, 1, 2048]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bool
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
                         position_ids:
                           tensor:
                             shape: [1, 1]
@@ -140,7 +150,33 @@ test_suite_config:
                             device: spyre
                             init: randint
                             init_args:
-                              high: 10000
+                              high: 1
+                        past_key_values:
+                          cache:
+                            cache_path: transformers.cache_utils.StaticCache
+                            layer_idx: 0
+                            key:
+                              shape: [1, 8, 128, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: xavier
+                            value:
+                              shape: [1, 8, 128, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: xavier
+                            max_cache_len: 2048
+                            config_path: transformers.models.granite.configuration_granite.GraniteConfig
+                            config_kwargs:
+                              hidden_size: 4096
+                              num_attention_heads: 32
+                              num_key_value_heads: 8
+                              num_hidden_layers: 40
+                              max_position_embeddings: 131072
                         position_embeddings:
                           tensor_list:
                             - shape: [1, 1, 128]
@@ -150,7 +186,7 @@ test_suite_config:
                               device: spyre
                               init: randint
                               init_args:
-                                high: 10000
+                                high: 1
                             - shape: [1, 1, 128]
                               stride: null
                               storage_offset: 0
@@ -158,10 +194,10 @@ test_suite_config:
                               device: spyre
                               init: randint
                               init_args:
-                                high: 10000
+                                high: 1
                 - name: GraniteRMSNorm_4096
                   module_path: transformers.models.granite.modeling_granite.GraniteRMSNorm
-                  description: 'Module: transformers.models.granite.modeling_granite.GraniteRMSNorm'
+                  description: 'Module: transformers.models.granite.modeling_granite.GraniteRMSNorm (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/granite/modeling_granite.py#L182)'
                   constructor_inputs:
                     args:
                       - value: 4096
@@ -174,7 +210,7 @@ test_suite_config:
                             storage_offset: 0
                             dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs: {}
                     - args:
                         - tensor:
@@ -183,11 +219,11 @@ test_suite_config:
                             storage_offset: 0
                             dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs: {}
                 - name: GraniteAttention_layer0
                   module_path: transformers.models.granite.modeling_granite.GraniteAttention
-                  description: 'Module: transformers.models.granite.modeling_granite.GraniteAttention'
+                  description: 'Module: transformers.models.granite.modeling_granite.GraniteAttention (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/granite/modeling_granite.py#L114)'
                   constructor_inputs:
                     args:
                       - config_path: transformers.models.granite.configuration_granite.GraniteConfig
@@ -210,7 +246,7 @@ test_suite_config:
                             storage_offset: 0
                             dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                         position_ids:
                           tensor:
                             shape: [1, 128]
@@ -220,7 +256,7 @@ test_suite_config:
                             device: spyre
                             init: randint
                             init_args:
-                              high: 10000
+                              high: 1
                         position_embeddings:
                           tensor_list:
                             - shape: [1, 128, 128]
@@ -230,7 +266,7 @@ test_suite_config:
                               device: spyre
                               init: randint
                               init_args:
-                                high: 10000
+                                high: 1
                             - shape: [1, 128, 128]
                               stride: null
                               storage_offset: 0
@@ -238,7 +274,7 @@ test_suite_config:
                               device: spyre
                               init: randint
                               init_args:
-                                high: 10000
+                                high: 1
                     - args: []
                       kwargs:
                         hidden_states:
@@ -248,7 +284,17 @@ test_suite_config:
                             storage_offset: 0
                             dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
+                        attention_mask:
+                          tensor:
+                            shape: [1, 1, 1, 2048]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bool
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
                         position_ids:
                           tensor:
                             shape: [1, 1]
@@ -258,7 +304,33 @@ test_suite_config:
                             device: spyre
                             init: randint
                             init_args:
-                              high: 10000
+                              high: 1
+                        past_key_values:
+                          cache:
+                            cache_path: transformers.cache_utils.StaticCache
+                            layer_idx: 0
+                            key:
+                              shape: [1, 8, 128, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: xavier
+                            value:
+                              shape: [1, 8, 128, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: xavier
+                            max_cache_len: 2048
+                            config_path: transformers.models.granite.configuration_granite.GraniteConfig
+                            config_kwargs:
+                              hidden_size: 4096
+                              num_attention_heads: 32
+                              num_key_value_heads: 8
+                              num_hidden_layers: 40
+                              max_position_embeddings: 131072
                         position_embeddings:
                           tensor_list:
                             - shape: [1, 1, 128]
@@ -268,7 +340,7 @@ test_suite_config:
                               device: spyre
                               init: randint
                               init_args:
-                                high: 10000
+                                high: 1
                             - shape: [1, 1, 128]
                               stride: null
                               storage_offset: 0
@@ -276,10 +348,10 @@ test_suite_config:
                               device: spyre
                               init: randint
                               init_args:
-                                high: 10000
+                                high: 1
                 - name: GraniteMLP_77d2613c
                   module_path: transformers.models.granite.modeling_granite.GraniteMLP
-                  description: 'Module: transformers.models.granite.modeling_granite.GraniteMLP'
+                  description: 'Module: transformers.models.granite.modeling_granite.GraniteMLP (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/granite/modeling_granite.py#L203)'
                   constructor_inputs:
                     args:
                       - config_path: transformers.models.granite.configuration_granite.GraniteConfig
@@ -299,7 +371,7 @@ test_suite_config:
                             storage_offset: 0
                             dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs: {}
                     - args:
                         - tensor:
@@ -308,11 +380,11 @@ test_suite_config:
                             storage_offset: 0
                             dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs: {}
                 - name: SiLUActivation_d2f532e9
                   module_path: transformers.activations.SiLUActivation
-                  description: 'Module: transformers.activations.SiLUActivation'
+                  description: 'Module: transformers.activations.SiLUActivation (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/activations.py#L92)'
                   constructor_inputs:
                     args: []
                     kwargs: {}
@@ -324,7 +396,7 @@ test_suite_config:
                             storage_offset: 0
                             dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs: {}
                     - args:
                         - tensor:
@@ -333,7 +405,7 @@ test_suite_config:
                             storage_offset: 0
                             dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs: {}
     - path: ${TORCH_DEVICE_ROOT}/tests/test_modules_custom.py
       unlisted_test_mode: skip
@@ -352,6 +424,14 @@ test_suite_config:
               include: *id001
   global:
     supported_dtypes:
+      - name: float16
+        precision:
+          atol: 0.005
+          rtol: 0.005
+      - name: float32
+        precision:
+          atol: 0.001
+          rtol: 0.001
       - name: bfloat16
         precision:
           atol: 0.005
diff --git a/tests/configs/module_tests/granite_4_1_8b_spyre.yaml b/tests/configs/module_tests/granite_4_1_8b_spyre.yaml
index 79e1aa56..802d5764 100644
--- a/tests/configs/module_tests/granite_4_1_8b_spyre.yaml
+++ b/tests/configs/module_tests/granite_4_1_8b_spyre.yaml
@@ -18,7 +18,7 @@ test_suite_config:
               include: &id001
                 - name: GraniteRotaryEmbedding_64892356
                   module_path: transformers.models.granite.modeling_granite.GraniteRotaryEmbedding
-                  description: 'Module: transformers.models.granite.modeling_granite.GraniteRotaryEmbedding'
+                  description: 'Module: transformers.models.granite.modeling_granite.GraniteRotaryEmbedding (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/granite/modeling_granite.py#L302)'
                   constructor_inputs:
                     args:
                       - config_path: transformers.models.granite.configuration_granite.GraniteConfig
@@ -36,36 +36,42 @@ test_suite_config:
                             shape: [1, 128, 4096]
                             stride: null
                             storage_offset: 0
-                            dtype: torch.float32
+                            dtype: torch.bfloat16
                             device: spyre
-                            init: randn
-                        - tensor:
+                            init: xavier
+                      kwargs:
+                        position_ids:
+                          tensor:
                             shape: [1, 128]
                             stride: null
                             storage_offset: 0
                             dtype: torch.int64
                             device: spyre
-                            init: randn
-                      kwargs: {}
+                            init: randint
+                            init_args:
+                              high: 1
                     - args:
                         - tensor:
                             shape: [1, 1, 4096]
                             stride: null
                             storage_offset: 0
-                            dtype: torch.float32
+                            dtype: torch.bfloat16
                             device: spyre
-                            init: randn
-                        - tensor:
+                            init: xavier
+                      kwargs:
+                        position_ids:
+                          tensor:
                             shape: [1, 1]
                             stride: null
                             storage_offset: 0
                             dtype: torch.int64
                             device: spyre
-                            init: randn
-                      kwargs: {}
+                            init: randint
+                            init_args:
+                              high: 1
                 - name: GraniteDecoderLayer_layer0
                   module_path: transformers.models.granite.modeling_granite.GraniteDecoderLayer
-                  description: 'Module: transformers.models.granite.modeling_granite.GraniteDecoderLayer'
+                  description: 'Module: transformers.models.granite.modeling_granite.GraniteDecoderLayer (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/granite/modeling_granite.py#L219)'
                   constructor_inputs:
                     args:
                       - config_path: transformers.models.granite.configuration_granite.GraniteConfig
@@ -84,9 +90,9 @@ test_suite_config:
                             shape: [1, 128, 4096]
                             stride: null
                             storage_offset: 0
-                            dtype: torch.float32
+                            dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs:
                         position_ids:
                           tensor:
@@ -97,85 +103,101 @@ test_suite_config:
                             device: spyre
                             init: randint
                             init_args:
-                              high: 10000
-                        cache_position:
-                          tensor:
-                            shape: [128]
-                            stride: null
-                            storage_offset: 0
-                            dtype: torch.int64
-                            device: spyre
-                            init: randint
-                            init_args:
-                              high: 10000
+                              high: 1
                         position_embeddings:
                           tensor_list:
                             - shape: [1, 128, 128]
                               stride: null
                               storage_offset: 0
-                              dtype: torch.float32
+                              dtype: torch.bfloat16
                               device: spyre
                               init: randint
                               init_args:
-                                high: 10000
+                                high: 1
                             - shape: [1, 128, 128]
                               stride: null
                               storage_offset: 0
-                              dtype: torch.float32
+                              dtype: torch.bfloat16
                               device: spyre
                               init: randint
                               init_args:
-                                high: 10000
+                                high: 1
                     - args:
                         - tensor:
                             shape: [1, 1, 4096]
                             stride: null
                             storage_offset: 0
-                            dtype: torch.float32
+                            dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs:
-                        position_ids:
+                        attention_mask:
                           tensor:
-                            shape: [1, 1]
+                            shape: [1, 1, 1, 2048]
                             stride: null
                             storage_offset: 0
-                            dtype: torch.int64
+                            dtype: torch.bool
                             device: spyre
                             init: randint
                             init_args:
-                              high: 10000
-                        cache_position:
+                              high: 1
+                        position_ids:
                           tensor:
-                            shape: [1]
+                            shape: [1, 1]
                             stride: null
                             storage_offset: 0
                             dtype: torch.int64
                             device: spyre
                             init: randint
                             init_args:
-                              high: 10000
+                              high: 1
+                        past_key_values:
+                          cache:
+                            cache_path: transformers.cache_utils.StaticCache
+                            layer_idx: 0
+                            key:
+                              shape: [1, 8, 128, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: xavier
+                            value:
+                              shape: [1, 8, 128, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: xavier
+                            max_cache_len: 2048
+                            config_path: transformers.models.granite.configuration_granite.GraniteConfig
+                            config_kwargs:
+                              hidden_size: 4096
+                              num_attention_heads: 32
+                              num_key_value_heads: 8
+                              num_hidden_layers: 40
+                              max_position_embeddings: 131072
                         position_embeddings:
                           tensor_list:
                             - shape: [1, 1, 128]
                               stride: null
                               storage_offset: 0
-                              dtype: torch.float32
+                              dtype: torch.bfloat16
                               device: spyre
                               init: randint
                               init_args:
-                                high: 10000
+                                high: 1
                             - shape: [1, 1, 128]
                               stride: null
                               storage_offset: 0
-                              dtype: torch.float32
+                              dtype: torch.bfloat16
                               device: spyre
                               init: randint
                               init_args:
-                                high: 10000
+                                high: 1
                 - name: GraniteRMSNorm_4096
                   module_path: transformers.models.granite.modeling_granite.GraniteRMSNorm
-                  description: 'Module: transformers.models.granite.modeling_granite.GraniteRMSNorm'
+                  description: 'Module: transformers.models.granite.modeling_granite.GraniteRMSNorm (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/granite/modeling_granite.py#L182)'
                   constructor_inputs:
                     args:
                       - value: 4096
@@ -186,22 +208,22 @@ test_suite_config:
                             shape: [1, 128, 4096]
                             stride: null
                             storage_offset: 0
-                            dtype: torch.float32
+                            dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs: {}
                     - args:
                         - tensor:
                             shape: [1, 1, 4096]
                             stride: null
                             storage_offset: 0
-                            dtype: torch.float32
+                            dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs: {}
                 - name: GraniteAttention_layer0
                   module_path: transformers.models.granite.modeling_granite.GraniteAttention
-                  description: 'Module: transformers.models.granite.modeling_granite.GraniteAttention'
+                  description: 'Module: transformers.models.granite.modeling_granite.GraniteAttention (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/granite/modeling_granite.py#L114)'
                   constructor_inputs:
                     args:
                       - config_path: transformers.models.granite.configuration_granite.GraniteConfig
@@ -222,9 +244,9 @@ test_suite_config:
                             shape: [1, 128, 4096]
                             stride: null
                             storage_offset: 0
-                            dtype: torch.float32
+                            dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                         position_ids:
                           tensor:
                             shape: [1, 128]
@@ -234,35 +256,25 @@ test_suite_config:
                             device: spyre
                             init: randint
                             init_args:
-                              high: 10000
-                        cache_position:
-                          tensor:
-                            shape: [128]
-                            stride: null
-                            storage_offset: 0
-                            dtype: torch.int64
-                            device: spyre
-                            init: randint
-                            init_args:
-                              high: 10000
+                              high: 1
                         position_embeddings:
                           tensor_list:
                             - shape: [1, 128, 128]
                               stride: null
                               storage_offset: 0
-                              dtype: torch.float32
+                              dtype: torch.bfloat16
                               device: spyre
                               init: randint
                               init_args:
-                                high: 10000
+                                high: 1
                             - shape: [1, 128, 128]
                               stride: null
                               storage_offset: 0
-                              dtype: torch.float32
+                              dtype: torch.bfloat16
                               device: spyre
                               init: randint
                               init_args:
-                                high: 10000
+                                high: 1
                     - args: []
                       kwargs:
                         hidden_states:
@@ -270,50 +282,76 @@ test_suite_config:
                             shape: [1, 1, 4096]
                             stride: null
                             storage_offset: 0
-                            dtype: torch.float32
+                            dtype: torch.bfloat16
                             device: spyre
-                            init: randn
-                        position_ids:
+                            init: xavier
+                        attention_mask:
                           tensor:
-                            shape: [1, 1]
+                            shape: [1, 1, 1, 2048]
                             stride: null
                             storage_offset: 0
-                            dtype: torch.int64
+                            dtype: torch.bool
                             device: spyre
                             init: randint
                             init_args:
-                              high: 10000
-                        cache_position:
+                              high: 1
+                        position_ids:
                           tensor:
-                            shape: [1]
+                            shape: [1, 1]
                             stride: null
                             storage_offset: 0
                             dtype: torch.int64
                             device: spyre
                             init: randint
                             init_args:
-                              high: 10000
+                              high: 1
+                        past_key_values:
+                          cache:
+                            cache_path: transformers.cache_utils.StaticCache
+                            layer_idx: 0
+                            key:
+                              shape: [1, 8, 128, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: xavier
+                            value:
+                              shape: [1, 8, 128, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: xavier
+                            max_cache_len: 2048
+                            config_path: transformers.models.granite.configuration_granite.GraniteConfig
+                            config_kwargs:
+                              hidden_size: 4096
+                              num_attention_heads: 32
+                              num_key_value_heads: 8
+                              num_hidden_layers: 40
+                              max_position_embeddings: 131072
                         position_embeddings:
                           tensor_list:
                             - shape: [1, 1, 128]
                               stride: null
                               storage_offset: 0
-                              dtype: torch.float32
+                              dtype: torch.bfloat16
                               device: spyre
                               init: randint
                               init_args:
-                                high: 10000
+                                high: 1
                             - shape: [1, 1, 128]
                               stride: null
                               storage_offset: 0
-                              dtype: torch.float32
+                              dtype: torch.bfloat16
                               device: spyre
                               init: randint
                               init_args:
-                                high: 10000
+                                high: 1
                 - name: GraniteMLP_77d2613c
                   module_path: transformers.models.granite.modeling_granite.GraniteMLP
-                  description: 'Module: transformers.models.granite.modeling_granite.GraniteMLP'
+                  description: 'Module: transformers.models.granite.modeling_granite.GraniteMLP (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/granite/modeling_granite.py#L203)'
                   constructor_inputs:
                     args:
                       - config_path: transformers.models.granite.configuration_granite.GraniteConfig
@@ -331,22 +369,22 @@ test_suite_config:
                             shape: [1, 128, 4096]
                             stride: null
                             storage_offset: 0
-                            dtype: torch.float32
+                            dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs: {}
                     - args:
                         - tensor:
                             shape: [1, 1, 4096]
                             stride: null
                             storage_offset: 0
-                            dtype: torch.float32
+                            dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs: {}
                 - name: SiLUActivation_d2f532e9
                   module_path: transformers.activations.SiLUActivation
-                  description: 'Module: transformers.activations.SiLUActivation'
+                  description: 'Module: transformers.activations.SiLUActivation (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/activations.py#L92)'
                   constructor_inputs:
                     args: []
                     kwargs: {}
@@ -356,18 +394,18 @@ test_suite_config:
                             shape: [1, 128, 12800]
                             stride: null
                             storage_offset: 0
-                            dtype: torch.float32
+                            dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs: {}
                     - args:
                         - tensor:
                             shape: [1, 1, 12800]
                             stride: null
                             storage_offset: 0
-                            dtype: torch.float32
+                            dtype: torch.bfloat16
                             device: spyre
-                            init: randn
+                            init: xavier
                       kwargs: {}
     - path: ${TORCH_DEVICE_ROOT}/tests/test_modules_custom.py
       unlisted_test_mode: skip

From d7496c5345cb5c8f661e56f0986dcdcda22f07e8 Mon Sep 17 00:00:00 2001
From: Kazuaki Ishizaki <ishizaki@jp.ibm.com>
Date: Fri, 7 Aug 2026 03:22:33 +0000
Subject: [PATCH 5/5] add more models

Signed-off-by: Kazuaki Ishizaki <ishizaki@jp.ibm.com>
---
 .github/workflows/test_pull_request.yaml      |   3 +
 .../Meta-Llama-3.1-8B-Instruct_spyre.yaml     | 442 ++++++++++
 .../Qwen2.5-7B-Instruct_spyre.yaml            | 438 ++++++++++
 .../module_tests/gpt_oss_20b_spyre.yaml       | 761 ++++++++++++++++++
 4 files changed, 1644 insertions(+)
 create mode 100644 tests/configs/module_tests/Meta-Llama-3.1-8B-Instruct_spyre.yaml
 create mode 100644 tests/configs/module_tests/Qwen2.5-7B-Instruct_spyre.yaml
 create mode 100644 tests/configs/module_tests/gpt_oss_20b_spyre.yaml

diff --git a/.github/workflows/test_pull_request.yaml b/.github/workflows/test_pull_request.yaml
index 0e40e4b7..45357805 100644
--- a/.github/workflows/test_pull_request.yaml
+++ b/.github/workflows/test_pull_request.yaml
@@ -553,8 +553,11 @@ jobs:
         config:
           - granite_4_1_8b_spyre.yaml
           - granite_3_3_8b_instruct_spyre.yaml
+	  - Meta-Llama-3.1-8B-Instruct_spyre.yaml
           - Ministral-3-14B-Instruct-2512.yaml
           - Mistral-Small-3.2-24B-Instruct-2506.yaml
+	  - Qwen2.5-7B-Instruct_spyre.yaml
+	  - gpt_oss_20b_spyre.yaml
 
     steps:
       - name: Checkout composite actions
diff --git a/tests/configs/module_tests/Meta-Llama-3.1-8B-Instruct_spyre.yaml b/tests/configs/module_tests/Meta-Llama-3.1-8B-Instruct_spyre.yaml
new file mode 100644
index 00000000..44bd88e1
--- /dev/null
+++ b/tests/configs/module_tests/Meta-Llama-3.1-8B-Instruct_spyre.yaml
@@ -0,0 +1,442 @@
+# Auto-generated unified test configuration for Meta_Llama_3_1_8B_Instruct
+# Generated by auto_generate_module_config.py
+# Format compatible with PyTorch's test_modules.py (using edits.modules.include)
+
+test_suite_config:
+  files:
+    - path: ${TORCH_ROOT}/test/test_modules.py
+      unlisted_test_mode: skip
+      tests:
+        - names:
+            - '*TestModule*::test_forward'
+          mode: xfail
+          tags:
+            - model__Meta_Llama_3_1_8B_Instruct
+          no_grad: true
+          edits:
+            modules:
+              include: &id001
+                - name: LlamaRotaryEmbedding_b501b76c
+                  module_path: transformers.models.llama.modeling_llama.LlamaRotaryEmbedding
+                  description: 'Module: transformers.models.llama.modeling_llama.LlamaRotaryEmbedding (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/llama/modeling_llama.py#L73)'
+                  constructor_inputs:
+                    args:
+                      - config_path: transformers.models.llama.configuration_llama.LlamaConfig
+                        config_kwargs:
+                          hidden_size: 4096
+                          num_attention_heads: 32
+                          num_key_value_heads: 8
+                          intermediate_size: 14336
+                          max_position_embeddings: 131072
+                          _attn_implementation: sdpa
+                    kwargs: {}
+                  forward_inputs:
+                    - args:
+                        - tensor:
+                            shape: [1, 128, 4096]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                      kwargs:
+                        position_ids:
+                          tensor:
+                            shape: [1, 128]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.int64
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                    - args:
+                        - tensor:
+                            shape: [1, 1, 4096]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                      kwargs:
+                        position_ids:
+                          tensor:
+                            shape: [1, 1]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.int64
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                - name: LlamaDecoderLayer_layer0
+                  module_path: transformers.models.llama.modeling_llama.LlamaDecoderLayer
+                  description: 'Module: transformers.models.llama.modeling_llama.LlamaDecoderLayer (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/llama/modeling_llama.py#L292)'
+                  constructor_inputs:
+                    args:
+                      - config_path: transformers.models.llama.configuration_llama.LlamaConfig
+                        config_kwargs:
+                          hidden_size: 4096
+                          num_attention_heads: 32
+                          num_key_value_heads: 8
+                          intermediate_size: 14336
+                          max_position_embeddings: 131072
+                          _attn_implementation: sdpa
+                    kwargs:
+                      layer_idx: 0
+                  forward_inputs:
+                    - args:
+                        - tensor:
+                            shape: [1, 128, 4096]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                      kwargs:
+                        position_embeddings:
+                          tensor_list:
+                            - shape: [1, 128, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: randint
+                              init_args:
+                                high: 1
+                            - shape: [1, 128, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: randint
+                              init_args:
+                                high: 1
+                        position_ids:
+                          tensor:
+                            shape: [1, 128]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.int64
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                    - args:
+                        - tensor:
+                            shape: [1, 1, 4096]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                      kwargs:
+                        attention_mask:
+                          tensor:
+                            shape: [1, 1, 1, 2048]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bool
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                        position_embeddings:
+                          tensor_list:
+                            - shape: [1, 1, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: randint
+                              init_args:
+                                high: 1
+                            - shape: [1, 1, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: randint
+                              init_args:
+                                high: 1
+                        position_ids:
+                          tensor:
+                            shape: [1, 1]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.int64
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                        past_key_values:
+                          cache:
+                            cache_path: transformers.cache_utils.StaticCache
+                            layer_idx: 0
+                            key:
+                              shape: [1, 8, 128, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: xavier
+                            value:
+                              shape: [1, 8, 128, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: xavier
+                            max_cache_len: 2048
+                            config_path: transformers.models.llama.configuration_llama.LlamaConfig
+                            config_kwargs:
+                              hidden_size: 4096
+                              num_attention_heads: 32
+                              num_key_value_heads: 8
+                              head_dim: 128
+                              num_hidden_layers: 32
+                              max_position_embeddings: 131072
+                - name: LlamaRMSNorm_4096
+                  module_path: transformers.models.llama.modeling_llama.LlamaRMSNorm
+                  description: 'Module: transformers.models.llama.modeling_llama.LlamaRMSNorm (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/llama/modeling_llama.py#L52)'
+                  constructor_inputs:
+                    args:
+                      - value: 4096
+                    kwargs: {}
+                  forward_inputs:
+                    - args:
+                        - tensor:
+                            shape: [1, 128, 4096]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                      kwargs: {}
+                    - args:
+                        - tensor:
+                            shape: [1, 1, 4096]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                      kwargs: {}
+                - name: LlamaAttention_layer0
+                  module_path: transformers.models.llama.modeling_llama.LlamaAttention
+                  description: 'Module: transformers.models.llama.modeling_llama.LlamaAttention (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/llama/modeling_llama.py#L224)'
+                  constructor_inputs:
+                    args:
+                      - config_path: transformers.models.llama.configuration_llama.LlamaConfig
+                        config_kwargs:
+                          hidden_size: 4096
+                          num_attention_heads: 32
+                          num_key_value_heads: 8
+                          intermediate_size: 14336
+                          max_position_embeddings: 131072
+                          _attn_implementation: sdpa
+                    kwargs:
+                      layer_idx: 0
+                  forward_inputs:
+                    - args: []
+                      kwargs:
+                        hidden_states:
+                          tensor:
+                            shape: [1, 128, 4096]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                        position_ids:
+                          tensor:
+                            shape: [1, 128]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.int64
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                        position_embeddings:
+                          tensor_list:
+                            - shape: [1, 128, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: randint
+                              init_args:
+                                high: 1
+                            - shape: [1, 128, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: randint
+                              init_args:
+                                high: 1
+                    - args: []
+                      kwargs:
+                        hidden_states:
+                          tensor:
+                            shape: [1, 1, 4096]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                        attention_mask:
+                          tensor:
+                            shape: [1, 1, 1, 2048]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bool
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                        position_ids:
+                          tensor:
+                            shape: [1, 1]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.int64
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                        past_key_values:
+                          cache:
+                            cache_path: transformers.cache_utils.StaticCache
+                            layer_idx: 0
+                            key:
+                              shape: [1, 8, 128, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: xavier
+                            value:
+                              shape: [1, 8, 128, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: xavier
+                            max_cache_len: 2048
+                            config_path: transformers.models.llama.configuration_llama.LlamaConfig
+                            config_kwargs:
+                              hidden_size: 4096
+                              num_attention_heads: 32
+                              num_key_value_heads: 8
+                              head_dim: 128
+                              num_hidden_layers: 32
+                              max_position_embeddings: 131072
+                        position_embeddings:
+                          tensor_list:
+                            - shape: [1, 1, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: randint
+                              init_args:
+                                high: 1
+                            - shape: [1, 1, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: randint
+                              init_args:
+                                high: 1
+                - name: LlamaMLP_231f37b9
+                  module_path: transformers.models.llama.modeling_llama.LlamaMLP
+                  description: 'Module: transformers.models.llama.modeling_llama.LlamaMLP (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/llama/modeling_llama.py#L171)'
+                  constructor_inputs:
+                    args:
+                      - config_path: transformers.models.llama.configuration_llama.LlamaConfig
+                        config_kwargs:
+                          hidden_size: 4096
+                          num_attention_heads: 32
+                          num_key_value_heads: 8
+                          intermediate_size: 14336
+                          max_position_embeddings: 131072
+                          _attn_implementation: sdpa
+                    kwargs: {}
+                  forward_inputs:
+                    - args:
+                        - tensor:
+                            shape: [1, 128, 4096]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                      kwargs: {}
+                    - args:
+                        - tensor:
+                            shape: [1, 1, 4096]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                      kwargs: {}
+                - name: SiLUActivation_d2f532e9
+                  module_path: transformers.activations.SiLUActivation
+                  description: 'Module: transformers.activations.SiLUActivation (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/activations.py#L92)'
+                  constructor_inputs:
+                    args: []
+                    kwargs: {}
+                  forward_inputs:
+                    - args:
+                        - tensor:
+                            shape: [1, 128, 14336]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                      kwargs: {}
+                    - args:
+                        - tensor:
+                            shape: [1, 1, 14336]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                      kwargs: {}
+    - path: ${TORCH_DEVICE_ROOT}/tests/test_modules_custom.py
+      unlisted_test_mode: skip
+      tests:
+        - names:
+            - '*TestModuleCustom*::test_with_cpu'
+            - '*TestModuleCustom*::test_eager_vs_compile'
+            - '*TestModuleCustom*::test_layout_stride'
+          mode: xfail
+          tags:
+            - model__Meta_Llama_3_1_8B_Instruct
+            - custom_tests
+          no_grad: true
+          edits:
+            modules:
+              include: *id001
+  global:
+    supported_dtypes:
+      - name: float16
+        precision:
+          atol: 0.005
+          rtol: 0.005
+      - name: float32
+        precision:
+          atol: 0.001
+          rtol: 0.001
+      - name: bfloat16
+        precision:
+          atol: 0.005
+          rtol: 0.005
+    input_config:
+      seed: 123
diff --git a/tests/configs/module_tests/Qwen2.5-7B-Instruct_spyre.yaml b/tests/configs/module_tests/Qwen2.5-7B-Instruct_spyre.yaml
new file mode 100644
index 00000000..9cbd5414
--- /dev/null
+++ b/tests/configs/module_tests/Qwen2.5-7B-Instruct_spyre.yaml
@@ -0,0 +1,438 @@
+# Auto-generated unified test configuration for Qwen2_5_7B_Instruct
+# Generated by auto_generate_module_config.py
+# Format compatible with PyTorch's test_modules.py (using edits.modules.include)
+
+test_suite_config:
+  files:
+    - path: ${TORCH_ROOT}/test/test_modules.py
+      unlisted_test_mode: skip
+      tests:
+        - names:
+            - '*TestModule*::test_forward'
+          mode: xfail
+          tags:
+            - model__Qwen2_5_7B_Instruct
+          no_grad: true
+          edits:
+            modules:
+              include: &id001
+                - name: Qwen2RotaryEmbedding_99b30225
+                  module_path: transformers.models.qwen2.modeling_qwen2.Qwen2RotaryEmbedding
+                  description: 'Module: transformers.models.qwen2.modeling_qwen2.Qwen2RotaryEmbedding (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/qwen2/modeling_qwen2.py#L51)'
+                  constructor_inputs:
+                    args:
+                      - config_path: transformers.models.qwen2.configuration_qwen2.Qwen2Config
+                        config_kwargs:
+                          hidden_size: 3584
+                          num_attention_heads: 28
+                          num_key_value_heads: 4
+                          intermediate_size: 18944
+                          max_position_embeddings: 32768
+                          _attn_implementation: sdpa
+                    kwargs: {}
+                  forward_inputs:
+                    - args:
+                        - tensor:
+                            shape: [1, 128, 3584]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                        - tensor:
+                            shape: [1, 128]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.int64
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                      kwargs: {}
+                    - args:
+                        - tensor:
+                            shape: [1, 1, 3584]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                        - tensor:
+                            shape: [1, 1]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.int64
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                      kwargs: {}
+                - name: Qwen2DecoderLayer_layer0
+                  module_path: transformers.models.qwen2.modeling_qwen2.Qwen2DecoderLayer
+                  description: 'Module: transformers.models.qwen2.modeling_qwen2.Qwen2DecoderLayer (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/qwen2/modeling_qwen2.py#L269)'
+                  constructor_inputs:
+                    args:
+                      - config_path: transformers.models.qwen2.configuration_qwen2.Qwen2Config
+                        config_kwargs:
+                          hidden_size: 3584
+                          num_attention_heads: 28
+                          num_key_value_heads: 4
+                          intermediate_size: 18944
+                          max_position_embeddings: 32768
+                          _attn_implementation: sdpa
+                    kwargs:
+                      layer_idx: 0
+                  forward_inputs:
+                    - args:
+                        - tensor:
+                            shape: [1, 128, 3584]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                      kwargs:
+                        position_embeddings:
+                          tensor_list:
+                            - shape: [1, 128, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: randint
+                              init_args:
+                                high: 1
+                            - shape: [1, 128, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: randint
+                              init_args:
+                                high: 1
+                        position_ids:
+                          tensor:
+                            shape: [1, 128]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.int64
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                    - args:
+                        - tensor:
+                            shape: [1, 1, 3584]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                      kwargs:
+                        attention_mask:
+                          tensor:
+                            shape: [1, 1, 1, 2048]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bool
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                        position_embeddings:
+                          tensor_list:
+                            - shape: [1, 1, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: randint
+                              init_args:
+                                high: 1
+                            - shape: [1, 1, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: randint
+                              init_args:
+                                high: 1
+                        position_ids:
+                          tensor:
+                            shape: [1, 1]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.int64
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                        past_key_values:
+                          cache:
+                            cache_path: transformers.cache_utils.StaticCache
+                            layer_idx: 0
+                            key:
+                              shape: [1, 4, 128, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: xavier
+                            value:
+                              shape: [1, 4, 128, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: xavier
+                            max_cache_len: 2048
+                            config_path: transformers.models.qwen2.configuration_qwen2.Qwen2Config
+                            config_kwargs:
+                              hidden_size: 3584
+                              num_attention_heads: 28
+                              num_key_value_heads: 4
+                              num_hidden_layers: 28
+                              max_position_embeddings: 32768
+                - name: Qwen2RMSNorm_3584
+                  module_path: transformers.models.qwen2.modeling_qwen2.Qwen2RMSNorm
+                  description: 'Module: transformers.models.qwen2.modeling_qwen2.Qwen2RMSNorm (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/qwen2/modeling_qwen2.py#L248)'
+                  constructor_inputs:
+                    args:
+                      - value: 3584
+                    kwargs: {}
+                  forward_inputs:
+                    - args:
+                        - tensor:
+                            shape: [1, 128, 3584]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                      kwargs: {}
+                    - args:
+                        - tensor:
+                            shape: [1, 1, 3584]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                      kwargs: {}
+                - name: Qwen2Attention_layer0
+                  module_path: transformers.models.qwen2.modeling_qwen2.Qwen2Attention
+                  description: 'Module: transformers.models.qwen2.modeling_qwen2.Qwen2Attention (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/qwen2/modeling_qwen2.py#L186)'
+                  constructor_inputs:
+                    args:
+                      - config_path: transformers.models.qwen2.configuration_qwen2.Qwen2Config
+                        config_kwargs:
+                          hidden_size: 3584
+                          num_attention_heads: 28
+                          num_key_value_heads: 4
+                          intermediate_size: 18944
+                          max_position_embeddings: 32768
+                          _attn_implementation: sdpa
+                    kwargs:
+                      layer_idx: 0
+                  forward_inputs:
+                    - args: []
+                      kwargs:
+                        hidden_states:
+                          tensor:
+                            shape: [1, 128, 3584]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                        position_ids:
+                          tensor:
+                            shape: [1, 128]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.int64
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                        position_embeddings:
+                          tensor_list:
+                            - shape: [1, 128, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: randint
+                              init_args:
+                                high: 1
+                            - shape: [1, 128, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: randint
+                              init_args:
+                                high: 1
+                    - args: []
+                      kwargs:
+                        hidden_states:
+                          tensor:
+                            shape: [1, 1, 3584]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                        attention_mask:
+                          tensor:
+                            shape: [1, 1, 1, 2048]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bool
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                        position_ids:
+                          tensor:
+                            shape: [1, 1]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.int64
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                        past_key_values:
+                          cache:
+                            cache_path: transformers.cache_utils.StaticCache
+                            layer_idx: 0
+                            key:
+                              shape: [1, 4, 128, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: xavier
+                            value:
+                              shape: [1, 4, 128, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: xavier
+                            max_cache_len: 2048
+                            config_path: transformers.models.qwen2.configuration_qwen2.Qwen2Config
+                            config_kwargs:
+                              hidden_size: 3584
+                              num_attention_heads: 28
+                              num_key_value_heads: 4
+                              num_hidden_layers: 28
+                              max_position_embeddings: 32768
+                        position_embeddings:
+                          tensor_list:
+                            - shape: [1, 1, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: randint
+                              init_args:
+                                high: 1
+                            - shape: [1, 1, 128]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: randint
+                              init_args:
+                                high: 1
+                - name: Qwen2MLP_ce0fbd40
+                  module_path: transformers.models.qwen2.modeling_qwen2.Qwen2MLP
+                  description: 'Module: transformers.models.qwen2.modeling_qwen2.Qwen2MLP (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/qwen2/modeling_qwen2.py#L35)'
+                  constructor_inputs:
+                    args:
+                      - config_path: transformers.models.qwen2.configuration_qwen2.Qwen2Config
+                        config_kwargs:
+                          hidden_size: 3584
+                          num_attention_heads: 28
+                          num_key_value_heads: 4
+                          intermediate_size: 18944
+                          max_position_embeddings: 32768
+                          _attn_implementation: sdpa
+                    kwargs: {}
+                  forward_inputs:
+                    - args:
+                        - tensor:
+                            shape: [1, 128, 3584]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                      kwargs: {}
+                    - args:
+                        - tensor:
+                            shape: [1, 1, 3584]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                      kwargs: {}
+                - name: SiLUActivation_d2f532e9
+                  module_path: transformers.activations.SiLUActivation
+                  description: 'Module: transformers.activations.SiLUActivation (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/activations.py#L92)'
+                  constructor_inputs:
+                    args: []
+                    kwargs: {}
+                  forward_inputs:
+                    - args:
+                        - tensor:
+                            shape: [1, 128, 18944]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                      kwargs: {}
+                    - args:
+                        - tensor:
+                            shape: [1, 1, 18944]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                      kwargs: {}
+    - path: ${TORCH_DEVICE_ROOT}/tests/test_modules_custom.py
+      unlisted_test_mode: skip
+      tests:
+        - names:
+            - '*TestModuleCustom*::test_with_cpu'
+            - '*TestModuleCustom*::test_eager_vs_compile'
+            - '*TestModuleCustom*::test_layout_stride'
+          mode: xfail
+          tags:
+            - model__Qwen2_5_7B_Instruct
+            - custom_tests
+          no_grad: true
+          edits:
+            modules:
+              include: *id001
+  global:
+    supported_dtypes:
+      - name: float16
+        precision:
+          atol: 0.005
+          rtol: 0.005
+      - name: float32
+        precision:
+          atol: 0.001
+          rtol: 0.001
+      - name: bfloat16
+        precision:
+          atol: 0.005
+          rtol: 0.005
+    input_config:
+      seed: 123
diff --git a/tests/configs/module_tests/gpt_oss_20b_spyre.yaml b/tests/configs/module_tests/gpt_oss_20b_spyre.yaml
new file mode 100644
index 00000000..08d488ad
--- /dev/null
+++ b/tests/configs/module_tests/gpt_oss_20b_spyre.yaml
@@ -0,0 +1,761 @@
+# Auto-generated unified test configuration for gpt_oss_20b
+# Generated by auto_generate_module_config.py
+# Format compatible with PyTorch's test_modules.py (using edits.modules.include)
+
+test_suite_config:
+  files:
+    - path: ${TORCH_ROOT}/test/test_modules.py
+      unlisted_test_mode: skip
+      tests:
+        - names:
+            - '*TestModule*::test_forward'
+          mode: xfail
+          tags:
+            - model__gpt_oss_20b
+          no_grad: true
+          edits:
+            modules:
+              include: &id001
+                - name: GptOssRotaryEmbedding_39c59122
+                  module_path: transformers.models.gpt_oss.modeling_gpt_oss.GptOssRotaryEmbedding
+                  description: 'Module: transformers.models.gpt_oss.modeling_gpt_oss.GptOssRotaryEmbedding (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/gpt_oss/modeling_gpt_oss.py#L154)'
+                  constructor_inputs:
+                    args:
+                      - config_path: transformers.models.gpt_oss.configuration_gpt_oss.GptOssConfig
+                        config_kwargs:
+                          hidden_size: 2880
+                          num_attention_heads: 64
+                          num_key_value_heads: 8
+                          intermediate_size: 2880
+                          max_position_embeddings: 131072
+                          _attn_implementation: eager
+                    kwargs: {}
+                  forward_inputs:
+                    - args:
+                        - tensor:
+                            shape: [1, 128, 2880]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                        - tensor:
+                            shape: [1, 128]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.int64
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                      kwargs: {}
+                    - args:
+                        - tensor:
+                            shape: [1, 1, 2880]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                        - tensor:
+                            shape: [1, 1]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.int64
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                      kwargs: {}
+                - name: GptOssDecoderLayer_layer0
+                  module_path: transformers.models.gpt_oss.modeling_gpt_oss.GptOssDecoderLayer
+                  description: 'Module: transformers.models.gpt_oss.modeling_gpt_oss.GptOssDecoderLayer (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/gpt_oss/modeling_gpt_oss.py#L354)'
+                  constructor_inputs:
+                    args:
+                      - config_path: transformers.models.gpt_oss.configuration_gpt_oss.GptOssConfig
+                        config_kwargs:
+                          hidden_size: 2880
+                          num_attention_heads: 64
+                          num_key_value_heads: 8
+                          intermediate_size: 2880
+                          max_position_embeddings: 131072
+                          _attn_implementation: eager
+                    kwargs:
+                      layer_idx: 0
+                  forward_inputs:
+                    - args:
+                        - tensor:
+                            shape: [1, 128, 2880]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                      kwargs:
+                        attention_mask:
+                          tensor:
+                            shape: [1, 1, 128, 128]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                        position_embeddings:
+                          tensor_list:
+                            - shape: [1, 128, 32]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: randint
+                              init_args:
+                                high: 1
+                            - shape: [1, 128, 32]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: randint
+                              init_args:
+                                high: 1
+                        position_ids:
+                          tensor:
+                            shape: [1, 128]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.int64
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                    - args:
+                        - tensor:
+                            shape: [1, 128, 2880]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                      kwargs:
+                        attention_mask:
+                          tensor:
+                            shape: [1, 1, 128, 2048]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                        position_embeddings:
+                          tensor_list:
+                            - shape: [1, 128, 32]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: randint
+                              init_args:
+                                high: 1
+                            - shape: [1, 128, 32]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: randint
+                              init_args:
+                                high: 1
+                        position_ids:
+                          tensor:
+                            shape: [1, 128]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.int64
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                    - args:
+                        - tensor:
+                            shape: [1, 1, 2880]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                      kwargs:
+                        attention_mask:
+                          tensor:
+                            shape: [1, 1, 1, 128]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                        position_embeddings:
+                          tensor_list:
+                            - shape: [1, 1, 32]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: randint
+                              init_args:
+                                high: 1
+                            - shape: [1, 1, 32]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: randint
+                              init_args:
+                                high: 1
+                        position_ids:
+                          tensor:
+                            shape: [1, 1]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.int64
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                        past_key_values:
+                          cache:
+                            cache_path: transformers.cache_utils.StaticCache
+                            layer_idx: 0
+                            key:
+                              shape: [1, 8, 128, 64]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: xavier
+                            value:
+                              shape: [1, 8, 128, 64]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: xavier
+                            max_cache_len: 2048
+                            config_path: transformers.models.gpt_oss.configuration_gpt_oss.GptOssConfig
+                            config_kwargs:
+                              hidden_size: 2880
+                              num_attention_heads: 64
+                              num_key_value_heads: 8
+                              head_dim: 64
+                              num_hidden_layers: 24
+                              max_position_embeddings: 131072
+                    - args:
+                        - tensor:
+                            shape: [1, 1, 2880]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                      kwargs:
+                        attention_mask:
+                          tensor:
+                            shape: [1, 1, 1, 2048]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                        position_embeddings:
+                          tensor_list:
+                            - shape: [1, 1, 32]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: randint
+                              init_args:
+                                high: 1
+                            - shape: [1, 1, 32]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: randint
+                              init_args:
+                                high: 1
+                        position_ids:
+                          tensor:
+                            shape: [1, 1]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.int64
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                        past_key_values:
+                          cache:
+                            cache_path: transformers.cache_utils.StaticCache
+                            layer_idx: 1
+                            key:
+                              shape: [1, 8, 128, 64]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: xavier
+                            value:
+                              shape: [1, 8, 128, 64]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: xavier
+                            max_cache_len: 2048
+                            config_path: transformers.models.gpt_oss.configuration_gpt_oss.GptOssConfig
+                            config_kwargs:
+                              hidden_size: 2880
+                              num_attention_heads: 64
+                              num_key_value_heads: 8
+                              head_dim: 64
+                              num_hidden_layers: 24
+                              max_position_embeddings: 131072
+                - name: GptOssRMSNorm_2880
+                  module_path: transformers.models.gpt_oss.modeling_gpt_oss.GptOssRMSNorm
+                  description: 'Module: transformers.models.gpt_oss.modeling_gpt_oss.GptOssRMSNorm (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/gpt_oss/modeling_gpt_oss.py#L52)'
+                  constructor_inputs:
+                    args:
+                      - value: 2880
+                    kwargs: {}
+                  forward_inputs:
+                    - args:
+                        - tensor:
+                            shape: [1, 128, 2880]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                      kwargs: {}
+                    - args:
+                        - tensor:
+                            shape: [1, 1, 2880]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                      kwargs: {}
+                - name: GptOssAttention_layer0
+                  module_path: transformers.models.gpt_oss.modeling_gpt_oss.GptOssAttention
+                  description: 'Module: transformers.models.gpt_oss.modeling_gpt_oss.GptOssAttention (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/gpt_oss/modeling_gpt_oss.py#L282)'
+                  constructor_inputs:
+                    args:
+                      - config_path: transformers.models.gpt_oss.configuration_gpt_oss.GptOssConfig
+                        config_kwargs:
+                          hidden_size: 2880
+                          num_attention_heads: 64
+                          num_key_value_heads: 8
+                          intermediate_size: 2880
+                          max_position_embeddings: 131072
+                          _attn_implementation: eager
+                    kwargs:
+                      layer_idx: 0
+                  forward_inputs:
+                    - args: []
+                      kwargs:
+                        hidden_states:
+                          tensor:
+                            shape: [1, 128, 2880]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                        attention_mask:
+                          tensor:
+                            shape: [1, 1, 128, 128]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                        position_ids:
+                          tensor:
+                            shape: [1, 128]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.int64
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                        position_embeddings:
+                          tensor_list:
+                            - shape: [1, 128, 32]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: randint
+                              init_args:
+                                high: 1
+                            - shape: [1, 128, 32]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: randint
+                              init_args:
+                                high: 1
+                    - args: []
+                      kwargs:
+                        hidden_states:
+                          tensor:
+                            shape: [1, 128, 2880]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                        attention_mask:
+                          tensor:
+                            shape: [1, 1, 128, 2048]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                        position_ids:
+                          tensor:
+                            shape: [1, 128]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.int64
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                        position_embeddings:
+                          tensor_list:
+                            - shape: [1, 128, 32]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: randint
+                              init_args:
+                                high: 1
+                            - shape: [1, 128, 32]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: randint
+                              init_args:
+                                high: 1
+                    - args: []
+                      kwargs:
+                        hidden_states:
+                          tensor:
+                            shape: [1, 1, 2880]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                        attention_mask:
+                          tensor:
+                            shape: [1, 1, 1, 128]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                        position_ids:
+                          tensor:
+                            shape: [1, 1]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.int64
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                        past_key_values:
+                          cache:
+                            cache_path: transformers.cache_utils.StaticCache
+                            layer_idx: 0
+                            key:
+                              shape: [1, 8, 128, 64]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: xavier
+                            value:
+                              shape: [1, 8, 128, 64]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: xavier
+                            max_cache_len: 2048
+                            config_path: transformers.models.gpt_oss.configuration_gpt_oss.GptOssConfig
+                            config_kwargs:
+                              hidden_size: 2880
+                              num_attention_heads: 64
+                              num_key_value_heads: 8
+                              head_dim: 64
+                              num_hidden_layers: 24
+                              max_position_embeddings: 131072
+                        position_embeddings:
+                          tensor_list:
+                            - shape: [1, 1, 32]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: randint
+                              init_args:
+                                high: 1
+                            - shape: [1, 1, 32]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: randint
+                              init_args:
+                                high: 1
+                    - args: []
+                      kwargs:
+                        hidden_states:
+                          tensor:
+                            shape: [1, 1, 2880]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                        attention_mask:
+                          tensor:
+                            shape: [1, 1, 1, 2048]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                        position_ids:
+                          tensor:
+                            shape: [1, 1]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.int64
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                        past_key_values:
+                          cache:
+                            cache_path: transformers.cache_utils.StaticCache
+                            layer_idx: 1
+                            key:
+                              shape: [1, 8, 128, 64]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: xavier
+                            value:
+                              shape: [1, 8, 128, 64]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: xavier
+                            max_cache_len: 2048
+                            config_path: transformers.models.gpt_oss.configuration_gpt_oss.GptOssConfig
+                            config_kwargs:
+                              hidden_size: 2880
+                              num_attention_heads: 64
+                              num_key_value_heads: 8
+                              head_dim: 64
+                              num_hidden_layers: 24
+                              max_position_embeddings: 131072
+                        position_embeddings:
+                          tensor_list:
+                            - shape: [1, 1, 32]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: randint
+                              init_args:
+                                high: 1
+                            - shape: [1, 1, 32]
+                              stride: null
+                              storage_offset: 0
+                              dtype: torch.bfloat16
+                              device: spyre
+                              init: randint
+                              init_args:
+                                high: 1
+                - name: GptOssMLP_ac9bcb56
+                  module_path: transformers.models.gpt_oss.modeling_gpt_oss.GptOssMLP
+                  description: 'Module: transformers.models.gpt_oss.modeling_gpt_oss.GptOssMLP (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/gpt_oss/modeling_gpt_oss.py#L138)'
+                  constructor_inputs:
+                    args: []
+                    kwargs: {}
+                  forward_inputs:
+                    - args:
+                        - tensor:
+                            shape: [1, 128, 2880]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                      kwargs: {}
+                    - args:
+                        - tensor:
+                            shape: [1, 1, 2880]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                      kwargs: {}
+                - name: GptOssTopKRouter_dab87a5a
+                  module_path: transformers.models.gpt_oss.modeling_gpt_oss.GptOssTopKRouter
+                  description: 'Module: transformers.models.gpt_oss.modeling_gpt_oss.GptOssTopKRouter (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/gpt_oss/modeling_gpt_oss.py#L122)'
+                  constructor_inputs:
+                    args: []
+                    kwargs: {}
+                  forward_inputs:
+                    - args:
+                        - tensor:
+                            shape: [128, 2880]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                      kwargs: {}
+                    - args:
+                        - tensor:
+                            shape: [1, 2880]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                      kwargs: {}
+                - name: GptOssExperts_03652cd6
+                  module_path: transformers.models.gpt_oss.modeling_gpt_oss.GptOssExperts
+                  description: 'Module: transformers.models.gpt_oss.modeling_gpt_oss.GptOssExperts (defined at https://github.com/huggingface/transformers/blob/v5.14.1/src/transformers/models/gpt_oss/modeling_gpt_oss.py#L73)'
+                  constructor_inputs:
+                    args:
+                      - config_path: transformers.models.gpt_oss.configuration_gpt_oss.GptOssConfig
+                        config_kwargs:
+                          hidden_size: 2880
+                          num_attention_heads: 64
+                          num_key_value_heads: 8
+                          intermediate_size: 2880
+                          max_position_embeddings: 131072
+                          _attn_implementation: eager
+                    kwargs: {}
+                  forward_inputs:
+                    - args:
+                        - tensor:
+                            shape: [128, 2880]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                        - tensor:
+                            shape: [128, 4]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.int64
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 4
+                        - tensor:
+                            shape: [128, 4]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                      kwargs: {}
+                    - args:
+                        - tensor:
+                            shape: [1, 2880]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                        - tensor:
+                            shape: [1, 4]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.int64
+                            device: spyre
+                            init: randint
+                            init_args:
+                              high: 1
+                        - tensor:
+                            shape: [1, 4]
+                            stride: null
+                            storage_offset: 0
+                            dtype: torch.bfloat16
+                            device: spyre
+                            init: xavier
+                      kwargs: {}
+    - path: ${TORCH_DEVICE_ROOT}/tests/test_modules_custom.py
+      unlisted_test_mode: skip
+      tests:
+        - names:
+            - '*TestModuleCustom*::test_with_cpu'
+            - '*TestModuleCustom*::test_eager_vs_compile'
+            - '*TestModuleCustom*::test_layout_stride'
+          mode: xfail
+          tags:
+            - model__gpt_oss_20b
+            - custom_tests
+          no_grad: true
+          edits:
+            modules:
+              include: *id001
+  global:
+    supported_dtypes:
+      - name: float16
+        precision:
+          atol: 0.005
+          rtol: 0.005
+      - name: float32
+        precision:
+          atol: 0.001
+          rtol: 0.001
+      - name: bfloat16
+        precision:
+          atol: 0.005
+          rtol: 0.005
+    input_config:
+      seed: 123