diff --git a/.github/workflows/xpu-max1100.yml b/.github/workflows/xpu-max1100.yml
index b9768caf6d26..cbca3a291f6e 100644
--- a/.github/workflows/xpu-max1100.yml
+++ b/.github/workflows/xpu-max1100.yml
@@ -72,7 +72,7 @@ jobs:
         export I_MPI_SHM=off
         pytest --verbose accelerator/*
         pytest --verbose autotuning/*
-        pytest --verbose model_parallelism/*
+        pytest --verbose v1/autotp/*
         pytest --verbose monitor/*
         pytest --verbose utils/*
         pytest --verbose runtime/test_ds_config_model.py
diff --git a/deepspeed/inference/engine.py b/deepspeed/inference/engine.py
index 6e5188ba5928..f247be89a02c 100755
--- a/deepspeed/inference/engine.py
+++ b/deepspeed/inference/engine.py
@@ -25,6 +25,7 @@
 from ..module_inject import LinearAllreduce, LinearLayer, Normalize, ReplaceWithTensorSlicing
 from deepspeed.accelerator import get_accelerator
 from ..module_inject.policy import TransformerPolicy
+from deepspeed.module_inject.tp_shard import AutoTPMeta
 from ..module_inject.auto_tp import AutoTP
 
 from ..module_inject.replace_policy import generic_policies
@@ -218,52 +219,44 @@ def build_alibi_tensor(self):
             if hasattr(self.module.transformer, 'build_alibi_tensor'):
                 # The heads must be sliced with the same tensor-parallel group that partitioned
                 # the attention weights, so bind it rather than letting the helper guess.
-                num_heads = self._get_model_head_count(self.module.transformer)
-                num_kv_heads = self._get_model_kv_head_count(self.module.transformer, num_heads)
-                shard_sizes = get_head_shard_sizes(num_heads, self.mp_group, num_kv_heads)
-                self.module.transformer.build_alibi_tensor = functools.partial(build_bloom_alibi_tensor,
-                                                                               mp_group=self.mp_group,
-                                                                               head_shard_sizes=shard_sizes,
-                                                                               total_num_heads=num_heads)
+                meta = self._autotp_meta(self.module.transformer)
+                shard_sizes = get_head_shard_sizes(meta, self.mp_group)
+                self.module.transformer.build_alibi_tensor = functools.partial(
+                    build_bloom_alibi_tensor,
+                    mp_group=self.mp_group,
+                    head_shard_sizes=shard_sizes,
+                    total_num_heads=meta.num_attention_heads)
             if hasattr(self.module.transformer, 'build_mpt_alibi_tensor'):
-                num_heads = self._get_model_head_count(self.module.transformer)
-                num_kv_heads = self._get_model_kv_head_count(self.module.transformer, num_heads)
+                meta = self._autotp_meta(self.module.transformer)
                 install_head_sharded_helper(self.module.transformer, 'build_mpt_alibi_tensor', build_mpt_alibi_tensor,
-                                            self.mp_group, num_heads, num_kv_heads)
+                                            meta, self.mp_group)
         if hasattr(self.module, 'model'):
             if hasattr(self.module.model, 'get_alibi_mask'):
-                num_heads = self._get_model_head_count(self.module.model)
-                num_kv_heads = self._get_model_kv_head_count(self.module.model, num_heads)
-                install_head_sharded_helper(self.module.model, 'get_alibi_mask', get_alibi_mask, self.mp_group,
-                                            num_heads, num_kv_heads)
+                meta = self._autotp_meta(self.module.model)
+                install_head_sharded_helper(self.module.model, 'get_alibi_mask', get_alibi_mask, meta, self.mp_group)
 
     def build_attn_bias(self):
         if hasattr(self.module, 'transformer'):
             if hasattr(self.module.transformer, '_attn_bias'):
-                num_heads = self._get_model_head_count(self.module.transformer)
-                num_kv_heads = self._get_model_kv_head_count(self.module.transformer, num_heads)
-                install_head_sharded_helper(self.module.transformer, '_attn_bias', build_mpt_atten_bias_tensor,
-                                            self.mp_group, num_heads, num_kv_heads)
-
-    def _get_model_head_count(self, module):
-        for source in (module, getattr(module, "config", None), getattr(self.module, "config", None)):
-            if source is None:
-                continue
-            for name in ("num_heads", "n_heads", "n_head", "num_attention_heads"):
-                value = getattr(source, name, None)
-                if value is not None:
-                    return value
-        raise ValueError(f"Cannot determine the attention head count for {module.__class__.__name__}.")
-
-    def _get_model_kv_head_count(self, module, num_heads):
-        for source in (module, getattr(module, "config", None), getattr(self.module, "config", None)):
-            if source is None:
-                continue
-            for name in ("num_key_value_heads", "num_kv_heads", "n_head_kv", "kv_n_heads"):
-                value = getattr(source, name, None)
-                if value is not None:
-                    return value
-        return num_heads
+                meta = self._autotp_meta(self.module.transformer)
+                install_head_sharded_helper(self.module.transformer, '_attn_bias', build_mpt_atten_bias_tensor, meta,
+                                            self.mp_group)
+
+    def _autotp_meta(self, module):
+        # from_model_config extracts from a single config object, but the head counts may live on
+        # the module, its config, or the top-level model config. Probe each source via the shared
+        # from_model_config and merge per field -- num_attention_heads is mandatory (alibi needs
+        # it), num_kv_heads is None for non-GQA. Subsumes the former
+        # _get_model_head_count / _get_model_kv_head_count pair.
+        metas = [
+            AutoTPMeta.from_model_config(s)
+            for s in (module, getattr(module, "config", None), getattr(self.module, "config", None)) if s is not None
+        ]
+        num_heads = next((m.num_attention_heads for m in metas if m.num_attention_heads is not None), None)
+        if num_heads is None:
+            raise ValueError(f"Cannot determine the attention head count for {module.__class__.__name__}.")
+        num_kv_heads = next((m.num_kv_heads for m in metas if m.num_kv_heads is not None), None)
+        return AutoTPMeta(num_attention_heads=num_heads, num_kv_heads=num_kv_heads)
 
     def _pre_forward_hook(self, module, *inputs, **kwargs):
         if self.use_cuda_events:
diff --git a/deepspeed/module_inject/auto_tp.py b/deepspeed/module_inject/auto_tp.py
index 6c967d5ef9f1..3204551dc3b0 100755
--- a/deepspeed/module_inject/auto_tp.py
+++ b/deepspeed/module_inject/auto_tp.py
@@ -15,7 +15,7 @@
 from .layers import *
 from deepspeed.accelerator import get_accelerator
 from .fusedqkv_utils import require_tp_fused_qkvw
-from deepspeed.module_inject.tp_shard import get_shard_size, get_shard_size_list
+from deepspeed.module_inject.tp_shard import AutoTPMeta, get_shard_size, get_shard_size_list
 from deepspeed.utils import groups
 from deepspeed.utils.logging import log_dist, print_dist
 from deepspeed.module_inject.layers import is_autotp_training_mode
@@ -205,7 +205,9 @@ def __init__(self,
                  linear_layer_setting,
                  orig_layer_impl,
                  keep_module_on_host=False,
-                 partition_config: Optional[AutoTPConfig] = None):
+                 partition_config: Optional[AutoTPConfig] = None,
+                 model_config=None,
+                 tp_grain_size: int = 1):
         self.module = module
         self.all_reduce_linears = all_reduce_linears
         self.prefix = prefix
@@ -213,6 +215,9 @@ def __init__(self,
 
         self.mp_size = None
         self.mp_group = None
+        # Per-model TP metadata threaded through every layer / helper so each AutoTP instance
+        # shards by its own model's kv-head / grain values.
+        self.tp_meta = AutoTPMeta.from_model_config(model_config, tp_grain_size)
         self.linear_layer_setting = linear_layer_setting
         self.orig_layer_impl = orig_layer_impl
         self.linear_policies = None
@@ -378,14 +383,14 @@ def _replace(self, child, name, conv_linear_layer):
         # For Yuan model
         if 'Yuan' in str(self.module):
             if 'v_proj' in name:
-                return Yuan_LinearLayer(child, self.mp_group)
+                return Yuan_LinearLayer(child, self.mp_group, tp_meta=self.tp_meta)
 
             elif 'o_proj' in name:
-                return Yuan_LinearAllreduce(child, self.mp_group)
+                return Yuan_LinearAllreduce(child, self.mp_group, tp_meta=self.tp_meta)
 
         # For MLP including chunk layer.
         if 'gate_up_proj' in name or ('dense_h_to_4h' in name and 'GLM' in str(self.module)):
-            return GateUpPack_LinearLayer(child, self.mp_group)
+            return GateUpPack_LinearLayer(child, self.mp_group, tp_meta=self.tp_meta)
             # For Arctic model, bypass to all_reduce replacement for w2 weights
         arctic_w2_all_reduce_linear = False
         if 'Arctic' in str(self.module) and 'w2' in name:
@@ -398,21 +403,21 @@ def _replace(self, child, name, conv_linear_layer):
 
             setattr(child, "replaced", True)
             if self.conv_linear_layer:
-                return Conv_LinearALlreduce(child, self.mp_group, name=name)
+                return Conv_LinearALlreduce(child, self.mp_group, name=name, tp_meta=self.tp_meta)
             elif name == "lm_head" or name == 'embed_out':
-                return LmHeadLinearAllreduce(child, self.mp_group)
+                return LmHeadLinearAllreduce(child, self.mp_group, tp_meta=self.tp_meta)
 
-            return LinearAllreduce(child, self.mp_group, name=name)
+            return LinearAllreduce(child, self.mp_group, name=name, tp_meta=self.tp_meta)
         else:
 
             setattr(child, "replaced", True)
             if self.conv_linear_layer:
-                conv_LinearLayer(child, self.mp_group)
+                conv_LinearLayer(child, self.mp_group, tp_meta=self.tp_meta)
             elif require_tp_fused_qkvw(name, self.mp_size):
                 #Check and handle fused qkv for TP
-                return fused_LinearLayer(child, self.mp_group, fused_module=self.module)
+                return fused_LinearLayer(child, self.mp_group, fused_module=self.module, tp_meta=self.tp_meta)
 
-            return LinearLayer(child, self.mp_group, name=name)
+            return LinearLayer(child, self.mp_group, name=name, tp_meta=self.tp_meta)
 
     def _replace_with_config(self, child, name):
         """
@@ -450,10 +455,10 @@ def _replace_with_config(self, child, name):
     def _create_row_parallel_layer(self, module, spec: TPLayerSpec, name: str):
         """Create row-parallel layer (AllReduce after forward)."""
         if self.conv_linear_layer:
-            return Conv_LinearALlreduce(module, self.mp_group, name=name)
+            return Conv_LinearALlreduce(module, self.mp_group, name=name, tp_meta=self.tp_meta)
         # Check for lm_head / embed_out
         if name == "lm_head" or name == 'embed_out':
-            return LmHeadLinearAllreduce(module, self.mp_group)
+            return LmHeadLinearAllreduce(module, self.mp_group, tp_meta=self.tp_meta)
 
         if spec.shape is not None:
             return SubParamLinearAllreduce(
@@ -462,17 +467,22 @@ def _create_row_parallel_layer(self, module, spec: TPLayerSpec, name: str):
                 shape=spec.shape,
                 partition_dim=spec.get_partition_dim(),
                 name=name,
+                tp_meta=self.tp_meta,
             )
-        return LinearAllreduce(module, self.mp_group, name=name)
+        return LinearAllreduce(module, self.mp_group, name=name, tp_meta=self.tp_meta)
 
     def _create_column_parallel_layer(self, module, spec: TPLayerSpec, name: str):
         """Create column-parallel layer (AllReduce in backward)."""
         if self.conv_linear_layer:
-            return conv_LinearLayer(module, self.mp_group, name=name, gather_output=spec.gather_output)
+            return conv_LinearLayer(module,
+                                    self.mp_group,
+                                    name=name,
+                                    gather_output=spec.gather_output,
+                                    tp_meta=self.tp_meta)
         # Only use fused-QKV heuristics when no partition_config is provided.
         elif self.partition_config is None and require_tp_fused_qkvw(name, self.mp_size):
             # Check and handle fused qkv for TP
-            return fused_LinearLayer(module, self.mp_group, fused_module=self.module)
+            return fused_LinearLayer(module, self.mp_group, fused_module=self.module, tp_meta=self.tp_meta)
         if spec.shape is not None:
             if spec.gather_output:
                 raise NotImplementedError("AutoTP gather_output does not yet support shaped sub-parameter layers.")
@@ -482,8 +492,9 @@ def _create_column_parallel_layer(self, module, spec: TPLayerSpec, name: str):
                 shape=spec.shape,
                 partition_dim=spec.get_partition_dim(),
                 name=name,
+                tp_meta=self.tp_meta,
             )
-        return LinearLayer(module, self.mp_group, name=name, gather_output=spec.gather_output)
+        return LinearLayer(module, self.mp_group, name=name, gather_output=spec.gather_output, tp_meta=self.tp_meta)
 
     def _configure_gathered_column_tie_fallbacks(self):
         """Configure a replicated fallback for gathered output layers tied to embeddings."""
@@ -561,7 +572,7 @@ def _slice_embedding(self, child, name, conv_linear_layer):
         mp_replace = ReplaceWithTensorSlicing(mp_group=self.mp_group)
 
         original_shape = tuple(child.weight.shape)
-        partition_sizes = get_shard_size_list(original_shape[1], self.mp_size, name)
+        partition_sizes = get_shard_size_list(original_shape[1], self.mp_size, self.tp_meta, name)
         if hasattr(child.weight, 'ds_tensor'):
             data = child.weight.ds_tensor.data.split(partition_sizes, dim=1)
         else:
@@ -650,7 +661,7 @@ def update_mp_params(self, child, name=None):
                 param_val = getattr(child, param)
                 # get_shard_size selects its partitioning strategy from the module name, so the
                 # attributes must be sharded under the same name as the weights they describe.
-                setattr(child, param, get_shard_size(param_val, self.mp_size, name, rank=tp_index))
+                setattr(child, param, get_shard_size(param_val, self.mp_size, self.tp_meta, name, rank=tp_index))
         setattr(child, "replaced", True)
 
     def update_linear_policies(self):
@@ -761,21 +772,6 @@ def _replace_module(self, r_module, prev_name='', prev_class_name=''):
                 self.update_mp_params(child, name)
         return r_module
 
-    @staticmethod
-    def get_model_num_kv_heads(config):
-        num_kv_heads = None
-        # multi_query_group_num is for chatglm2 & chatglm3
-        kv_head_names = [
-            'multi_query_group_num', 'num_kv_heads', 'num_key_value_heads', 'num_attention_heads', 'n_heads',
-            'attention_heads'
-        ]
-        for name in kv_head_names:
-            if hasattr(config, name):
-                num_kv_heads = getattr(config, name)
-                if num_kv_heads is not None:
-                    break
-        return num_kv_heads
-
     def _replace_last_linear_module(self, r_module):
         if hasattr(r_module, "lm_head"):
             name = "lm_head"
diff --git a/deepspeed/module_inject/auto_tp_model_utils.py b/deepspeed/module_inject/auto_tp_model_utils.py
index 8b4a3a6c852e..422bc651fbc3 100644
--- a/deepspeed/module_inject/auto_tp_model_utils.py
+++ b/deepspeed/module_inject/auto_tp_model_utils.py
@@ -8,7 +8,7 @@
 from deepspeed import comm as dist
 import torch
 from typing import Optional
-from deepspeed.module_inject.tp_shard import get_shard_size_list
+from deepspeed.module_inject.tp_shard import AutoTPMeta, get_shard_size_list
 
 
 class _HeadCountProxy:
@@ -27,16 +27,12 @@ def __setattr__(self, name, value):
             setattr(self._module, name, value)
 
 
-def get_head_shard_sizes(num_heads, mp_group=None, num_kv_heads=None):
+def get_head_shard_sizes(meta: AutoTPMeta, mp_group=None):
     tp_world_size = dist.get_world_size(group=mp_group)
-    return get_shard_size_list(
-        num_heads,
-        tp_world_size,
-        num_kv_heads=num_kv_heads,
-    )
+    return get_shard_size_list(meta.num_attention_heads, tp_world_size, meta)
 
 
-def install_head_sharded_helper(module, name, wrapper, mp_group=None, num_heads=None, num_kv_heads=None):
+def install_head_sharded_helper(module, name, wrapper, meta, mp_group=None):
     """Give ``module`` a head-slicing wrapper around one of its own methods.
 
     The wrapper is bound to this instance instead of installed on its class. A class-wide patch
@@ -48,10 +44,15 @@ def install_head_sharded_helper(module, name, wrapper, mp_group=None, num_heads=
     if original_name not in module.__dict__:
         # Wrapping an already wrapped instance would make it delegate to itself.
         setattr(module, original_name, getattr(module, name))
-    shard_sizes = get_head_shard_sizes(num_heads, mp_group, num_kv_heads) if num_heads is not None else None
+    total_num_heads = meta.num_attention_heads
+    shard_sizes = get_head_shard_sizes(meta, mp_group) if total_num_heads is not None else None
     setattr(
         module, name,
-        functools.partial(wrapper, module, mp_group=mp_group, head_shard_sizes=shard_sizes, total_num_heads=num_heads))
+        functools.partial(wrapper,
+                          module,
+                          mp_group=mp_group,
+                          head_shard_sizes=shard_sizes,
+                          total_num_heads=total_num_heads))
 
 
 def _head_shard(num_heads, mp_group=None, head_shard_sizes=None, total_num_heads=None):
@@ -64,7 +65,9 @@ def _head_shard(num_heads, mp_group=None, head_shard_sizes=None, total_num_heads
     tp_world_size = dist.get_world_size(group=mp_group)
     tp_index = dist.get_rank(group=mp_group)
     full_num_heads = total_num_heads if total_num_heads is not None else num_heads
-    shard_sizes = head_shard_sizes or get_shard_size_list(full_num_heads, tp_world_size)
+    # The fallback split is only reached when no per-head shard sizes were recorded, i.e. when
+    # the model is not GQA-aware, so an even split (default meta, no kv-heads) is correct here.
+    shard_sizes = head_shard_sizes or get_shard_size_list(full_num_heads, tp_world_size, AutoTPMeta())
     if len(shard_sizes) != tp_world_size or sum(shard_sizes) != full_num_heads:
         raise ValueError(f"Head shard sizes {shard_sizes} do not partition {full_num_heads} heads across "
                          f"{tp_world_size} tensor-parallel ranks.")
diff --git a/deepspeed/module_inject/fusedqkv_utils.py b/deepspeed/module_inject/fusedqkv_utils.py
index 22aa1186a5b1..1e7d9963d3c5 100644
--- a/deepspeed/module_inject/fusedqkv_utils.py
+++ b/deepspeed/module_inject/fusedqkv_utils.py
@@ -4,7 +4,7 @@
 # DeepSpeed Team
 import torch
 from deepspeed.utils.logging import warning_once
-from deepspeed.module_inject.tp_shard import get_shard_size, get_shard_size_list, get_num_kv_heads, get_n_embd, get_num_attention_heads
+from deepspeed.module_inject.tp_shard import AutoTPMeta, get_shard_size, get_shard_size_list
 
 
 def split_by_qkvlist_and_refuse(qkv_list, split_size, split_dim=0, cat_dim=0):
@@ -51,7 +51,7 @@ def get_fused_qkv_type(module):
     return FUSED_QKV_TYPE_DICT[max(module_name_matches, key=len)]
 
 
-def fused_qkv_subparam_sizes(module, weight_shape):
+def fused_qkv_subparam_sizes(module, weight_shape, meta: AutoTPMeta):
     """Sizes of the sub-parameters a fused qkv weight is cut into, or None.
 
     ``prepare_tp_fused_qkvw`` splits most layouts into q/k/v (or a single block) and
@@ -68,15 +68,15 @@ def fused_qkv_subparam_sizes(module, weight_shape):
     if fused_type == 'bloomtype':
         return (total_size, )
     if fused_type in ('glmtype', 'qwentype'):
-        if get_num_kv_heads() == 2:
-            hidden_dim = get_n_embd()
-            kv_dim = (total_size - hidden_dim) // get_num_kv_heads()
+        if meta.num_kv_heads == 2:
+            hidden_dim = meta.n_embd
+            kv_dim = (total_size - hidden_dim) // meta.num_kv_heads
             return (hidden_dim, kv_dim, kv_dim)
         third = total_size // 3
         return (third, third, third)
     if fused_type == 'phi3type':
-        head_dim = weight_shape[1] // get_num_attention_heads()
-        kv_dim = get_num_kv_heads() * head_dim
+        head_dim = weight_shape[1] // meta.num_attention_heads
+        kv_dim = meta.num_kv_heads * head_dim
         return (total_size - 2 * kv_dim, kv_dim, kv_dim)
     # codegentype interleaves blocks across ranks and bigcodetype replicates the kv block,
     # so neither is a per-sub-parameter split.
@@ -102,79 +102,83 @@ def set_fused_qkv_shard_state(module, shard_widths, tp_index):
     module.attn.split_size = query_width
 
 
-def prepare_tp_fused_qkvw(module, src, mp_size, gpu_index):
+def prepare_tp_fused_qkvw(module, src, mp_size, gpu_index, meta: AutoTPMeta):
 
     if src is None:
         return
 
-    def _codegen_type_transpose(input, mp_size, codegen_mp_num=4):
+    def _codegen_type_transpose(input, mp_size, codegen_mp_num=4, meta=meta):
         # codegen_mp_num defined in https://github.com/huggingface/transformers/blob/main/src/transformers/models/codegen/modeling_codegen.py
-        assert get_num_kv_heads() % (
+        assert meta.num_kv_heads % (
             mp_size * codegen_mp_num) == 0, "codgen autoTP requires num_kv_heads % (mp_size*codegen_mp_num) == 0"
         #input : [3*hidden_dim, hidden_dim](weight) or [3*hidden_dim](bias)
 
         shape = input.shape
-        dst_shape = get_shard_size(shape[0], mp_size, rank=gpu_index)
+        dst_shape = get_shard_size(shape[0], mp_size, meta, rank=gpu_index)
         num_mp_blocks = input.reshape(codegen_mp_num, shape[0] // codegen_mp_num, shape[1])
 
         #num_mp_blocks : [codegen_mp_num, 3*hidden_dim/codegen_mp_num, :]
         src_split = list(torch.split(num_mp_blocks, num_mp_blocks.shape[1] // 3, dim=1))
         src_split = [x.reshape(codegen_mp_num * mp_size, -1, shape[1]) for x in src_split]
 
-        split_fusedqkv = split_by_qkvlist_and_refuse(src_split, get_shard_size(shape[0] // 3, mp_size, rank=gpu_index),
-                                                     0, 1)
+        split_fusedqkv = split_by_qkvlist_and_refuse(src_split,
+                                                     get_shard_size(shape[0] // 3, mp_size, meta, rank=gpu_index), 0,
+                                                     1)
         tp_fuseqkv_weight = torch.cat(split_fusedqkv, dim=0).reshape(shape[0], -1)
 
         return tp_fuseqkv_weight[gpu_index * dst_shape:(gpu_index + 1) * dst_shape]
 
-    def _glm_type_transpose(input, mp_size):
+    def _glm_type_transpose(input, mp_size, meta=meta):
         #input : [3*hidden_dim, hidden_dim](weight) or [3*hidden_dim](bias)
 
         # For chatglm2 & chatglm3(kv_heads=2), need to special handle.
-        if get_num_kv_heads() == 2:
+        if meta.num_kv_heads == 2:
             shape = input.shape
-            hidden_dim = get_n_embd()
-            kv_dim = (shape[0] - hidden_dim) // get_num_kv_heads()
+            hidden_dim = meta.n_embd
+            kv_dim = (shape[0] - hidden_dim) // meta.num_kv_heads
             q = input[:hidden_dim]
             k = input[hidden_dim:hidden_dim + kv_dim]
             v = input[hidden_dim + kv_dim:]
-            q_split = q.split(get_shard_size_list(q.shape[0], mp_size), dim=0)
-            k_split = k.split(get_shard_size_list(k.shape[0], mp_size), dim=0)
-            v_split = v.split(get_shard_size_list(v.shape[0], mp_size), dim=0)
+            q_split = q.split(get_shard_size_list(q.shape[0], mp_size, meta), dim=0)
+            k_split = k.split(get_shard_size_list(k.shape[0], mp_size, meta), dim=0)
+            v_split = v.split(get_shard_size_list(v.shape[0], mp_size, meta), dim=0)
             return torch.cat((q_split[gpu_index], k_split[gpu_index], v_split[gpu_index]), dim=0)
         else:
             shape = input.shape
             src_split = torch.split(input, shape[0] // 3, dim=0)
 
-            split_fusedqkv = split_by_qkvlist_and_refuse(src_split, get_shard_size_list(shape[0] // 3, mp_size))
+            split_fusedqkv = split_by_qkvlist_and_refuse(src_split, get_shard_size_list(shape[0] // 3, mp_size, meta))
             return split_fusedqkv[gpu_index]
 
-    def _bloom_type_transpose(input, mp_size):
+    def _bloom_type_transpose(input, mp_size, meta=meta):
         shape = input.shape
 
-        split_fusedqkv = input.split(get_shard_size_list(shape[0], mp_size), dim=0)
+        split_fusedqkv = input.split(get_shard_size_list(shape[0], mp_size, meta), dim=0)
         return split_fusedqkv[gpu_index]
 
-    def _bigcode_type_transpose(input, mp_size):
-        n_embd = get_n_embd()
+    def _bigcode_type_transpose(input, mp_size, meta=meta):
+        n_embd = meta.n_embd
+        # A missing hidden size would slice as input[:None], handing the whole fused weight to q
+        # and leaving kv empty rather than failing.
+        assert n_embd is not None
         q = input[:n_embd]
         kv = input[n_embd:]
         shape = q.shape
-        split_q = q.split(get_shard_size_list(shape[0], mp_size), dim=0)
+        split_q = q.split(get_shard_size_list(shape[0], mp_size, meta), dim=0)
         return torch.cat((split_q[gpu_index], kv), dim=0)
 
-    def _phi3_type_transpose(input, mp_size):
-        num_kv_heads = get_num_kv_heads()
-        num_heads = get_num_attention_heads()
+    def _phi3_type_transpose(input, mp_size, meta=meta):
+        num_kv_heads = meta.num_kv_heads
+        num_heads = meta.num_attention_heads
         hidden_size = input.shape[1]
         head_dim = hidden_size // num_heads
         q_pos = input.shape[0] - 2 * num_kv_heads * head_dim
         q = input[:q_pos]
         k = input[q_pos:q_pos + num_kv_heads * head_dim]
         v = input[q_pos + num_kv_heads * head_dim:]
-        split_q = q.split(get_shard_size_list(q.shape[0], mp_size), dim=0)
-        split_k = k.split(get_shard_size_list(k.shape[0], mp_size), dim=0)
-        split_v = v.split(get_shard_size_list(v.shape[0], mp_size), dim=0)
+        split_q = q.split(get_shard_size_list(q.shape[0], mp_size, meta), dim=0)
+        split_k = k.split(get_shard_size_list(k.shape[0], mp_size, meta), dim=0)
+        split_v = v.split(get_shard_size_list(v.shape[0], mp_size, meta), dim=0)
         return torch.cat((split_q[gpu_index], split_k[gpu_index], split_v[gpu_index]), dim=0)
 
     def _transpose_fused_qkvw(src, mp_size, fused_qkv_type=None, module=None):
@@ -216,15 +220,15 @@ def shard_value_with_share_qk(
         bias,
         rank,
         world_size,
-        shard_value=True  # True -> shard_value; False -> shard_oproj
-):
+        shard_value,  # True -> shard_value; False -> shard_oproj
+        meta: AutoTPMeta):
     if shard_value:
         total_size = weight.shape[0]
         weight_cat_dim = 0
     else:
         total_size = weight.shape[1]
         weight_cat_dim = 1
-    num_heads = get_num_kv_heads()
+    num_heads = meta.num_kv_heads
     head_dim = total_size // num_heads
     assert (num_heads % world_size == 0)
     if world_size > num_heads // 2:
diff --git a/deepspeed/module_inject/layers.py b/deepspeed/module_inject/layers.py
index 22623b42c73c..3a779c968fcf 100644
--- a/deepspeed/module_inject/layers.py
+++ b/deepspeed/module_inject/layers.py
@@ -11,7 +11,7 @@
 from torch.nn import functional as F
 from torch.nn.parameter import Parameter
 from deepspeed.accelerator import get_accelerator
-from deepspeed.module_inject.tp_shard import get_shard_size_list
+from deepspeed.module_inject.tp_shard import AutoTPMeta, get_shard_size_list
 from deepspeed.utils.logging import log_dist_once
 from deepspeed.runtime.zero.utils import is_zero_param
 from abc import ABC, abstractmethod
@@ -348,6 +348,10 @@ def __init__(self, mp_group: Optional[dist.ProcessGroup], **kwargs: Any):
         if kwargs.get('name') is not None:
             self.name = kwargs.get('name')  # Set the layer name if provided.
 
+        # Per-model TP metadata threaded from AutoTP; defaults for layers built outside it
+        # (e.g. the from_weights back-compat constructor).
+        self.tp_meta: AutoTPMeta = kwargs.get('tp_meta') or AutoTPMeta()
+
     @classmethod
     def set_keep_module_on_host(cls, value: bool):
         """
@@ -473,13 +477,11 @@ def is_training_mode(self):
     def _freeze_partition_sizes(self, total_size):
         """Resolve the tensor parallel split of this layer once, while the layer is built.
 
-        ``get_shard_size_list`` reads the process-wide tp_shard globals (``num_kv_heads``,
-        ``tp_grain_size``), which a later ``init_inference`` call or a second AutoTP model
-        overwrites. The split is part of the checkpoint contract, so it is resolved here and
-        every later consumer -- the forward gather, the parameter gather and the checkpoint
-        metadata -- reads the cached value rather than querying those globals again.
+        The split depends on this model's kv-head/grain metadata (``self.tp_meta``), so it is
+        resolved here and every later consumer -- the forward gather, the parameter gather and
+        the checkpoint metadata -- reads the cached value rather than re-deriving it.
         """
-        self._partition_sizes = tuple(get_shard_size_list(total_size, self.tp_world_size, self.name))
+        self._partition_sizes = tuple(get_shard_size_list(total_size, self.tp_world_size, self.tp_meta, self.name))
         return self._partition_sizes
 
     @torch.no_grad()
@@ -964,7 +966,8 @@ def _freeze_partition_sizes(self, total_size):
         self._subparam_sizes = subparam_sizes
         self._subparam_shard_widths = None
         if subparam_sizes is not None:
-            self._subparam_shard_widths = _subparam_shard_widths(subparam_sizes, self.tp_world_size, shard_name)
+            self._subparam_shard_widths = _subparam_shard_widths(subparam_sizes, self.tp_world_size, self.tp_meta,
+                                                                 shard_name)
 
     def _subparam_shape_spec(self, logical_shape):
         shape_spec = list(logical_shape)
@@ -1084,8 +1087,8 @@ def __init__(self, module, mp_group, skip_partition=False, **kwargs):
         self.fused_module = FusedModuleWrapper(kwargs.get('fused_module'))
         # prepare_tp_fused_qkvw takes its own shard sizes without a layer name, so the widths
         # describing its split must be resolved the same way.
-        self._subparam_layout_spec = (fused_qkv_subparam_sizes(kwargs.get('fused_module'),
-                                                               tuple(module.weight.shape)), None)
+        self._subparam_layout_spec = (fused_qkv_subparam_sizes(kwargs.get('fused_module'), tuple(module.weight.shape),
+                                                               kwargs['tp_meta']), None)
         super().__init__(module, mp_group, skip_partition, **kwargs)
 
     def _freeze_partition_sizes(self, total_size):
@@ -1101,7 +1104,8 @@ def _tp_partition_unsupported_layout(self, params_list):
             if param is None:
                 return
 
-            _partition = prepare_tp_fused_qkvw(self.fused_module.module, param, self.tp_world_size, self.tp_index)
+            _partition = prepare_tp_fused_qkvw(self.fused_module.module, param, self.tp_world_size, self.tp_index,
+                                               self.tp_meta)
 
             _partition = self.move(_partition).detach()
 
@@ -1118,13 +1122,15 @@ def _tp_partition(self, params_list):
             weight = params_list[0]
         elif len(params_list) == 2:
             weight, bias = params_list[0], params_list[1]
-        _partition = weight.data.split(get_shard_size_list(weight.shape[0], self.tp_world_size, self.name),
+        _partition = weight.data.split(get_shard_size_list(weight.shape[0], self.tp_world_size, self.tp_meta,
+                                                           self.name),
                                        dim=1)[self.tp_index]
         _partition = self.move(_partition).detach()
         weight.data = _partition
 
         if bias is not None:
-            _partition = bias.data.split(get_shard_size_list(weight.shape[1], self.tp_world_size, self.name),
+            _partition = bias.data.split(get_shard_size_list(weight.shape[1], self.tp_world_size, self.tp_meta,
+                                                             self.name),
                                          dim=0)[self.tp_index]
             _partition = self.move(_partition).detach()
 
@@ -1141,7 +1147,7 @@ class Yuan_LinearAllreduce(LinearAllreduce):
     @torch.no_grad()
     def _tp_partition(self, params_list):
         weight, bias = shard_value_with_share_qk(params_list[0].data, params_list[1], self.tp_index,
-                                                 self.tp_world_size, False)
+                                                 self.tp_world_size, False, self.tp_meta)
         params_list[0].data = weight
         if bias is not None:
             params_list[1].data = bias
@@ -1175,7 +1181,7 @@ class Yuan_LinearLayer(LinearLayer):
     @torch.no_grad()
     def _tp_partition(self, params_list):
         weight, bias = shard_value_with_share_qk(params_list[0].data, params_list[1], self.tp_index,
-                                                 self.tp_world_size, True)
+                                                 self.tp_world_size, True, self.tp_meta)
         params_list[0].data = self.move(weight).detach()
         if bias is not None:
             params_list[1].data = self.move(bias).detach()
@@ -1220,7 +1226,7 @@ def _tp_partition(self, params_list):
                 return
             param.data = param.data.transpose(-1, -2).contiguous()
 
-            _partition = param.split(get_shard_size_list(param.shape[0], self.tp_world_size, self.name),
+            _partition = param.split(get_shard_size_list(param.shape[0], self.tp_world_size, self.tp_meta, self.name),
                                      dim=1)[self.tp_index]
 
             _partition = self.move(_partition).detach()
@@ -1517,7 +1523,7 @@ def _bias_subparam_shape_spec(output_shape, bias_partition_dim, subparam_sizes):
     return tuple(shape_spec)
 
 
-def _subparam_shard_widths(subparam_sizes, tp_world_size, name=None):
+def _subparam_shard_widths(subparam_sizes, tp_world_size, meta: AutoTPMeta, name=None):
     """Per-rank width of each sub-parameter, as one list per sub-parameter.
 
     Sub-parameters follow the same deterministic split as ordinary layers, so a fused
@@ -1525,7 +1531,7 @@ def _subparam_shard_widths(subparam_sizes, tp_world_size, name=None):
     """
     widths = []
     for size in subparam_sizes:
-        per_rank = get_shard_size_list(size, tp_world_size, name)
+        per_rank = get_shard_size_list(size, tp_world_size, meta, name)
         if min(per_rank) == 0:
             # Those ranks contribute zeros to the row-parallel all-reduce, so the result stays
             # correct and this matches how separate q/k/v projections already behave. Serving
@@ -1664,9 +1670,10 @@ def __init__(self, module, mp_group, shape, partition_dim=0, **kwargs):
          self._bias_partition_dim) = _infer_subparam_logical_shapes(self._orig_weight_shape, self.shape,
                                                                     self.partition_dim, self.name)
         # Resolve the per-rank widths once, for the same reason _freeze_partition_sizes does:
-        # get_shard_size_list reads process-wide globals that a later model can overwrite.
+        # the split depends on this model's tp_meta.
         self._subparam_shard_widths = _subparam_shard_widths(
-            self._subparam_sizes or (self._logical_shape[self.partition_dim], ), self.tp_world_size, self.name)
+            self._subparam_sizes or (self._logical_shape[self.partition_dim], ), self.tp_world_size, self.tp_meta,
+            self.name)
         self._bias_shape_spec = _bias_subparam_shape_spec(self._output_shape, self._bias_partition_dim,
                                                           self._subparam_sizes)
         if self.bias is not None and self.bias.numel() != _shape_prod(self._output_shape):
@@ -1796,9 +1803,10 @@ def __init__(self, module, mp_group, shape, partition_dim=1, **kwargs):
          self._bias_partition_dim) = _infer_subparam_logical_shapes(self._orig_weight_shape, self.shape,
                                                                     self.partition_dim, self.name)
         # Resolve the per-rank widths once, for the same reason _freeze_partition_sizes does:
-        # get_shard_size_list reads process-wide globals that a later model can overwrite.
+        # the split depends on this model's tp_meta.
         self._subparam_shard_widths = _subparam_shard_widths(
-            self._subparam_sizes or (self._logical_shape[self.partition_dim], ), self.tp_world_size, self.name)
+            self._subparam_sizes or (self._logical_shape[self.partition_dim], ), self.tp_world_size, self.tp_meta,
+            self.name)
 
         if self._should_materialize_tp_partition():
             self._tp_partition([self.weight, self.bias])
diff --git a/deepspeed/module_inject/replace_module.py b/deepspeed/module_inject/replace_module.py
index 263369fc0484..35b03587c6bd 100644
--- a/deepspeed/module_inject/replace_module.py
+++ b/deepspeed/module_inject/replace_module.py
@@ -17,7 +17,6 @@
 from .layers import TensorParallelOcShardConv2d, TensorParallelIcShardConv2d
 from deepspeed.module_inject.layers import is_autotp_training_mode
 from deepspeed import comm as dist
-from deepspeed.module_inject.tp_shard import set_num_kv_heads, set_n_embd, set_num_attention_heads, set_tp_grain_size
 
 from .load_checkpoint import load_model_with_checkpoint
 import time
@@ -278,7 +277,19 @@ def replace_wo_policy(module, all_reduce_linears, prefix="", state_dict=None):
         if hasattr(config, 'get_partition_config_object'):
             partition_config = config.get_partition_config_object()
 
-        # 1. Create AutoTP object
+        # 1. Resolve which sub-config describes the layers being replaced (e.g. Mllama vision
+        # vs text); AutoTP reads its TP metadata from it.
+        if hasattr(model_config, "vision_config"):
+            if "MllamaVisionEncoderLayer" in str(module):
+                meta_config = model_config.vision_config
+            elif hasattr(model_config, "text_config"):
+                meta_config = model_config.text_config
+            else:
+                meta_config = model_config
+        else:
+            meta_config = model_config
+
+        # 2. Create AutoTP object; it derives its tp_meta from meta_config.
         _autotp = AutoTP(module,
                          all_reduce_linears,
                          prefix,
@@ -286,48 +297,17 @@ def replace_wo_policy(module, all_reduce_linears, prefix="", state_dict=None):
                          linear_layer_setting,
                          orig_layer_impl,
                          config.keep_module_on_host,
-                         partition_config=partition_config)
+                         partition_config=partition_config,
+                         model_config=meta_config,
+                         tp_grain_size=config.tensor_parallel.tp_grain_size)
 
-        # 2. Set the tensor parallelism config
+        # 3. Set the tensor parallelism config
         _autotp.set_tensor_parallel_config(config.tensor_parallel.tp_size, config.tensor_parallel.tp_group)
 
-        # 3. Try to get num_key_heads from model_config.num_key_value_heads
-        if hasattr(model_config, "vision_config"):
-            if "MllamaVisionEncoderLayer" in str(module):
-                num_kv_heads = _autotp.get_model_num_kv_heads(model_config.vision_config)
-            elif hasattr(model_config, "text_config"):
-                num_kv_heads = _autotp.get_model_num_kv_heads(model_config.text_config)
-            else:
-                num_kv_heads = _autotp.get_model_num_kv_heads(model_config)
-        else:
-            num_kv_heads = _autotp.get_model_num_kv_heads(model_config)
-
-        # 4. When we have num_kv_heads defined, uneven division is possible, otherwise enforce even division
-        set_num_kv_heads(num_kv_heads)
-
-        # 4.1 Get n_embd
-        n_embd = None
-        multi_query_n_embd_names = ['n_embd', 'hidden_size']
-        for name in multi_query_n_embd_names:
-            if hasattr(model_config, name):
-                n_embd = getattr(model_config, name)
-            if n_embd != None:
-                break
-
-        # 4.2 set n_embd
-        set_n_embd(n_embd)
-
-        # 4.3 set attention_heads
-        if hasattr(model_config, 'num_attention_heads'):
-            set_num_attention_heads(getattr(model_config, 'num_attention_heads'))
-
-        # 4.4 set tp_grain_size
-        set_tp_grain_size(config.tensor_parallel.tp_grain_size)
-
-        # 5. Set linear policies
+        # 4. Set linear policies
         _autotp.update_linear_policies()
 
-        # 6. Replace modules
+        # 5. Replace modules
         if "lm_head" in all_reduce_linears or "embed_out" in all_reduce_linears:
             return _autotp._replace_last_linear_module(module)
         return _autotp._replace_module(module)
diff --git a/deepspeed/module_inject/tp_shard.py b/deepspeed/module_inject/tp_shard.py
index f976ff22bc33..092a5bcb4a78 100644
--- a/deepspeed/module_inject/tp_shard.py
+++ b/deepspeed/module_inject/tp_shard.py
@@ -3,69 +3,110 @@
 
 # DeepSpeed Team
 
-from deepspeed import comm as dist
-
-# Defaults for optional TP globals. These can be overridden by setters.
-num_kv_heads = None
-num_attention_heads = None
-n_embd = None
-tp_grain_size = 1
-
-
-def set_num_kv_heads(num):
-    global num_kv_heads
-    num_kv_heads = num
-
-
-def set_num_attention_heads(num):
-    global num_attention_heads
-    num_attention_heads = num
+from dataclasses import dataclass
+from typing import Optional
 
+from deepspeed import comm as dist
 
-def set_n_embd(num):
-    global n_embd
-    n_embd = num
+# Attribute names that carry the key/value head count, in probe order. Modern canonical names
+# first, then legacy aliases kept for older configs / checkpoints, then the query-head names as
+# an implicit "kv == q heads" default for plain (non-GQA) transformers. Shared by AutoTP and the
+# inference engine so the two never diverge on which model families they recognize.
+_KV_HEAD_ATTRS = (
+    'num_key_value_heads',  # llama-class, mistral, qwen2, gemma, phi3, dbrx top-level, ...
+    'num_kv_heads',  # falcon, qwen3_moe
+    'multi_query_group_num',  # chatglm2 / chatglm3 (not in stock transformers)
+    'n_head_kv',  # legacy Falcon custom-code; deprecated after transformers 4.33 (-> num_kv_heads)
+    'kv_n_heads',  # dbrx nested attn_config; top-level config exposes num_key_value_heads (>= transformers 4.40)
+    'num_attention_heads',  # plain transformer: kv == q heads
+    'n_heads',
+    'attention_heads',
+)
+
+
+def _kv_head_count_from(config) -> Optional[int]:
+    """First non-None kv-head count found on ``config`` via :data:`_KV_HEAD_ATTRS`, else ``None``."""
+    for name in _KV_HEAD_ATTRS:
+        value = getattr(config, name, None)
+        if value is not None:
+            return value
+    return None
 
 
-def set_tp_grain_size(num):
-    global tp_grain_size
-    tp_grain_size = num
+# Attribute names that carry the (query) attention head count, in probe order.
+_ATTN_HEAD_ATTRS = ('num_attention_heads', 'num_heads', 'n_heads', 'n_head', 'attention_heads')
 
 
-def get_num_kv_heads():
-    global num_kv_heads
-    if 'num_kv_heads' in globals():
-        return num_kv_heads
+def _attention_head_count_from(config) -> Optional[int]:
+    """First non-None attention head count found on ``config`` via :data:`_ATTN_HEAD_ATTRS`, else ``None``."""
+    for name in _ATTN_HEAD_ATTRS:
+        value = getattr(config, name, None)
+        if value is not None:
+            return value
     return None
 
 
-def get_num_attention_heads():
-    global num_attention_heads
-    return num_attention_heads
-
+@dataclass(frozen=True)
+class AutoTPMeta:
+    """Per-model tensor-parallel metadata AutoTP derives from the model config.
 
-def get_shard_size(total_size, mp_size, name=None, rank=None, mp_group=None, num_kv_heads=None):
+    Each model carries its own kv-head / grain values so its layers, fused-QKV repacking and
+    checkpoint conversion shard consistently, and more than one AutoTP model can live in the
+    same process.
+    """
+    num_kv_heads: Optional[int] = None
+    num_attention_heads: Optional[int] = None
+    n_embd: Optional[int] = None
+    tp_grain_size: int = 1
+
+    @classmethod
+    def from_model_config(cls, model_config, tp_grain_size: int = 1) -> "AutoTPMeta":
+        """The single source of truth for reading kv-head / hidden / attention-head counts.
+
+        ``model_config`` may be ``None`` for callers that build a synthetic AutoTP with no real
+        model (some unit tests); they get an empty meta and the resulting even-grain split.
+        """
+        if model_config is None:
+            return cls(tp_grain_size=tp_grain_size)
+        # Multimodal configs (e.g. vision-language models) keep the head counts only under
+        # text_config while the outer config carries the modality heads, so descend to the
+        # text config when it exists.
+        text_config = getattr(model_config, "text_config", None)
+        if text_config is not None:
+            model_config = text_config
+        num_kv_heads = _kv_head_count_from(model_config)
+        n_embd = None
+        for name in ('n_embd', 'hidden_size'):
+            if hasattr(model_config, name):
+                n_embd = getattr(model_config, name)
+            if n_embd is not None:
+                break
+        num_attention_heads = _attention_head_count_from(model_config)
+        return cls(num_kv_heads=num_kv_heads,
+                   num_attention_heads=num_attention_heads,
+                   n_embd=n_embd,
+                   tp_grain_size=tp_grain_size)
+
+
+def get_shard_size(total_size, mp_size, meta: AutoTPMeta, name=None, rank=None, mp_group=None):
     """Size of one shard of ``total_size`` split across a tensor-parallel group of ``mp_size``.
 
+    ``meta`` carries this model's ``num_kv_heads`` / ``tp_grain_size`` so the split is stable
+    for the lifetime of the model instead of depending on whichever AutoTP model was loaded
+    last.
+
     ``rank`` is the rank *within the tensor-parallel group*, i.e. in ``[0, mp_size)``, matching
     ``dist.get_rank(group=mp_group)`` and the index used by ``get_shard_size_list``. It is not a
     global rank.
-
-    ``num_kv_heads`` overrides the module-level global of the same name when set. Passing it
-    explicitly lets callers split fused sub-parameters (Q/K/V) against their respective head
-    counts without reimplementing the KV-head-aligned partition logic at the call site. When
-    ``None`` (default), the function falls back to the global ``num_kv_heads`` set via
-    ``set_num_kv_heads`` for backward compatibility.
     """
-    if num_kv_heads is None:
-        num_kv_heads = globals()["num_kv_heads"]
+    num_kv_heads = meta.num_kv_heads
+    tp_grain_size = meta.tp_grain_size
     last_linear = ["lm_head", "embed_out"]
     # MoE MLP layer use near even division will get better perf.
     moe_mlp_layer = ["gate_proj", "up_proj", "down_proj", "w1", "w2", "w3"]
     not_moe_mlp_layer = True
     if name != None and any(s in str(name) for s in moe_mlp_layer):
         not_moe_mlp_layer = False
-    # When we have num_kv_heads defined, uneven division is possible, otherwise enforce near even division
     if rank is None:
         if mp_group is not None:
             rank = dist.get_rank(group=mp_group)
@@ -75,6 +116,7 @@ def get_shard_size(total_size, mp_size, name=None, rank=None, mp_group=None, num
                 raise ValueError("get_shard_size requires a group-local rank or process group when mp_size "
                                  f"({mp_size}) differs from the distributed world size ({world_size}).")
             rank = dist.get_rank()
+    # A known kv-head count allows uneven division, otherwise enforce near even division.
     if num_kv_heads is not None and total_size % num_kv_heads == 0 and "mlp" not in str(name) and \
         str(name) not in last_linear and not_moe_mlp_layer:
         my_slices = (num_kv_heads // mp_size) + (1 if rank < (num_kv_heads % mp_size) else 0)
@@ -93,21 +135,14 @@ def get_shard_size(total_size, mp_size, name=None, rank=None, mp_group=None, num
             return total_size // mp_size + (1 if rank < (total_size % mp_size) else 0)
 
 
-def get_n_embd():
-    global n_embd
-    return n_embd
-
-
-def get_shard_size_list(total_size, mp_size, name=None, num_kv_heads=None):
+def get_shard_size_list(total_size, mp_size, meta: AutoTPMeta, name=None):
     shard_sizes = []
-    if num_kv_heads is None:
-        num_kv_heads = globals()["num_kv_heads"]
     for i in range(mp_size):
-        shard_sizes.append(get_shard_size(total_size, mp_size, name, i, num_kv_heads=num_kv_heads))
+        shard_sizes.append(get_shard_size(total_size, mp_size, meta, name, i))
     # Shards must tile the dimension exactly, otherwise the partitioned weights no longer
     # reconstruct the original tensor.
     assert sum(shard_sizes) == total_size, (
         f"AutoTP shard sizes {shard_sizes} for layer '{name}' do not sum to the dimension size "
-        f"{total_size} with tp_size={mp_size}, tp_grain_size={tp_grain_size} and "
-        f"num_kv_heads={num_kv_heads}.")
+        f"{total_size} with tp_size={mp_size}, tp_grain_size={meta.tp_grain_size} and "
+        f"num_kv_heads={meta.num_kv_heads}.")
     return shard_sizes
diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
index 1c590ccf90ce..c1c202e20313 100644
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -742,20 +742,15 @@ def _apply_autotp_partitioning(self, model, tp_config):
             partition_config = tp_config.get_partition_config_object()
 
         model_config = getattr(model, "config", None)
-        # The direct Hugging Face tp_plan path bypasses replace_transformer_layer, which
-        # normally initializes the shard-size globals that AutoTP layers consult. Without
-        # them attention projections are split by grain size and can be cut mid-head, so
-        # the model's later reshape onto head_dim fails.
-        from deepspeed.module_inject.tp_shard import set_num_kv_heads, set_n_embd, set_num_attention_heads
-        from deepspeed.module_inject.tp_shard import set_tp_grain_size
-
-        # 1. Try to get num_key_heads from model_config.num_key_value_heads
-        if hasattr(model_config, "text_config"):
-            num_kv_heads = AutoTP.get_model_num_kv_heads(model_config.text_config)
-        else:
-            num_kv_heads = AutoTP.get_model_num_kv_heads(model_config)
+        # AutoTP derives its per-model sharding metadata from the model config; the warning
+        # below only needs the kv-head count, which that same metadata already carries.
+        # from_model_config descends into text_config itself, so multimodal outer configs
+        # work here too.
+        from deepspeed.module_inject.tp_shard import AutoTPMeta
+
+        num_kv_heads = AutoTPMeta.from_model_config(model_config).num_kv_heads
 
-        # 2. Ranks beyond the KV head count get no attention shard. This still computes the
+        # Ranks beyond the KV head count get no attention shard. This still computes the
         # correct result because the row-parallel all-reduce sums their empty contribution,
         # but attention work concentrates on the first num_kv_heads ranks.
         if num_kv_heads is not None and tp_size > num_kv_heads:
@@ -766,28 +761,6 @@ def _apply_autotp_partitioning(self, model, tp_config):
                 ranks=[0],
                 level=logging.WARNING)
 
-        # 3. When we have num_kv_heads defined, uneven division is possible, otherwise enforce even division
-        set_num_kv_heads(num_kv_heads)
-
-        # 3.1 Get n_embd
-        n_embd = None
-        multi_query_n_embd_names = ['n_embd', 'hidden_size']
-        for name in multi_query_n_embd_names:
-            if hasattr(model_config, name):
-                n_embd = getattr(model_config, name)
-            if n_embd != None:
-                break
-
-        # 3.2 set n_embd
-        set_n_embd(n_embd)
-
-        # 3.3 set attention_heads
-        if hasattr(model_config, 'num_attention_heads'):
-            set_num_attention_heads(getattr(model_config, 'num_attention_heads'))
-
-        # 3.4 set tp_grain_size
-        set_tp_grain_size(tp_config.tensor_parallel.tp_grain_size)
-
         from deepspeed.runtime.tensor_parallel.config import _get_hf_tp_plan
         hf_tp_plan = _get_hf_tp_plan(model)
 
@@ -799,7 +772,9 @@ def _apply_autotp_partitioning(self, model, tp_config):
                             linear_layer_setting=(torch.nn.Linear, torch.nn.Embedding),
                             orig_layer_impl=None,
                             keep_module_on_host=tp_config.keep_module_on_host,
-                            partition_config=partition_config)
+                            partition_config=partition_config,
+                            model_config=model_config,
+                            tp_grain_size=tp_config.tensor_parallel.tp_grain_size)
             autotp.set_tensor_parallel_config(tp_size, tp_config.tensor_parallel.tp_group)
             autotp.update_linear_policies()
             autotp._replace_module(model)
@@ -838,6 +813,8 @@ def _apply_autotp_partitioning(self, model, tp_config):
                     orig_layer_impl=None,
                     keep_module_on_host=tp_config.keep_module_on_host,
                     partition_config=tp_plan_config,
+                    model_config=model_config,
+                    tp_grain_size=tp_config.tensor_parallel.tp_grain_size,
                 )
                 autotp.set_tensor_parallel_config(tp_size, tp_config.tensor_parallel.tp_group)
                 autotp.update_linear_policies()
diff --git a/deepspeed/sequence/layer.py b/deepspeed/sequence/layer.py
index 1ae18c14ced8..2dea696918f5 100644
--- a/deepspeed/sequence/layer.py
+++ b/deepspeed/sequence/layer.py
@@ -12,9 +12,32 @@
 
 import deepspeed.comm as dist
 from deepspeed.accelerator import get_accelerator
-from deepspeed.module_inject.tp_shard import get_shard_size_list, set_num_kv_heads, get_num_kv_heads
+from deepspeed.module_inject.tp_shard import AutoTPMeta, get_shard_size_list
 from deepspeed.utils import groups
 
+# Ulysses sequence parallelism keeps its own kv-head count, memoized on the first uneven
+# all-to-all. The state lives here, independent of AutoTP.
+# TODO: this is process-wide and never reset, so the first model to take the uneven path locks
+# every later Ulysses call into it -- a second model with a different head count would then be
+# split against the first one's value. Moving it onto the attention instance means threading the
+# total head count through _SeqAllToAll and its backward pass, because the gather direction
+# cannot recover it from the tensor shape alone.
+_ulysses_num_kv_heads = None
+
+
+def set_ulysses_num_kv_heads(num):
+    global _ulysses_num_kv_heads
+    _ulysses_num_kv_heads = num
+
+
+def get_ulysses_num_kv_heads():
+    return _ulysses_num_kv_heads
+
+
+def _ulysses_meta():
+    return AutoTPMeta(num_kv_heads=_ulysses_num_kv_heads)
+
+
 try:
     from torchembed._triton import fused_rope_forward as _torchembed_rope_forward
     _torchembed_available = True
@@ -134,7 +157,7 @@ def uneven_heads_all2all(input, scatter_idx, gather_idx, batch_dim_idx, group):
     assert batch_dim_idx in [0, 1], "batch_dim_idx must be either 0 or 1"
 
     if not (scatter_idx < 2):
-        input_splits = get_shard_size_list(inp_shape[scatter_idx], seq_world_size)
+        input_splits = get_shard_size_list(inp_shape[scatter_idx], seq_world_size, _ulysses_meta())
         input = input.transpose(0, scatter_idx).contiguous()
         local_heads = input_splits[groups._get_sequence_parallel_rank()]
         output_splits = [local_heads] * seq_world_size
@@ -168,7 +191,7 @@ def uneven_heads_all2all(input, scatter_idx, gather_idx, batch_dim_idx, group):
         elif batch_dim_idx == 1:  #s,b,h
             input = input.transpose(1, 2).contiguous()  #s,h,b
         seq_len, h, batch_size = input.shape
-        num_local_heads_list = get_shard_size_list(get_num_kv_heads(), seq_world_size)
+        num_local_heads_list = get_shard_size_list(get_ulysses_num_kv_heads(), seq_world_size, _ulysses_meta())
         local_heads = num_local_heads_list[groups._get_sequence_parallel_rank()]
         h_dim = h // local_heads
         local_seq_len = seq_len // seq_world_size
@@ -179,7 +202,7 @@ def uneven_heads_all2all(input, scatter_idx, gather_idx, batch_dim_idx, group):
         coeff = local_seq_len_with_heads // local_heads  #per head: dim size of local_seq_len*hdim
 
         #uneven seq_world_size coeff,  total_heads/local_heads.
-        heads_scale_coeff = get_num_kv_heads() / local_heads
+        heads_scale_coeff = get_ulysses_num_kv_heads() / local_heads
 
         output_splits = [num_local_heads * coeff for num_local_heads in num_local_heads_list]
         output_buff_d1_size = int(heads_scale_coeff * local_seq_len_with_heads)
@@ -196,9 +219,9 @@ def uneven_heads_all2all(input, scatter_idx, gather_idx, batch_dim_idx, group):
         #total_num_large_heads=sum([2,2,2])=7
         #total_num_small_heads=sum([1])=1
 
-        chunk_num_heads_small = get_num_kv_heads() // seq_world_size  # even heads compatible
+        chunk_num_heads_small = get_ulysses_num_kv_heads() // seq_world_size  # even heads compatible
         chunk_num_heads_large = chunk_num_heads_small + 1
-        num_chunk_heads_large = get_num_kv_heads() % seq_world_size
+        num_chunk_heads_large = get_ulysses_num_kv_heads() % seq_world_size
         num_chunk_heads_small = seq_world_size - num_chunk_heads_large
         total_num_large_heads = num_chunk_heads_large * chunk_num_heads_large
         total_num_small_heads = num_chunk_heads_small * chunk_num_heads_small
@@ -243,14 +266,14 @@ def single_all_to_all(input, scatter_idx, gather_idx, batch_dim_idx, group, asyn
     # we only need num_heads once
     num_heads = input.shape[2]
 
-    if get_num_kv_heads() is not None or (num_heads % seq_world_size != 0 and not scatter_idx < 2):
+    if get_ulysses_num_kv_heads() is not None or (num_heads % seq_world_size != 0 and not scatter_idx < 2):
         # Assuming here that the number of heads for q is consistent with kv
         # If not, additional logic is required for cases like GQA
-        if get_num_kv_heads() is None:
+        if get_ulysses_num_kv_heads() is None:
             assert num_heads > seq_world_size, f"Number of heads ({num_heads}) must be larger than sequence parallel size ({seq_world_size})"
             # set heads at first call by num_total_heads.
-            # then use ``get_num_kv_heads() is not None`` to re-entry uneven path.
-            set_num_kv_heads(num_heads)
+            # then use ``get_ulysses_num_kv_heads() is not None`` to re-entry uneven path.
+            set_ulysses_num_kv_heads(num_heads)
         assert async_op == False, "uneven head sp does not support async op"
         return uneven_heads_all2all(input, scatter_idx, gather_idx, batch_dim_idx, group)
 
diff --git a/tests/unit/checkpoint/test_autotp_uc_checkpoint.py b/tests/unit/checkpoint/test_autotp_uc_checkpoint.py
index efc85905cdba..af615fa2cc27 100644
--- a/tests/unit/checkpoint/test_autotp_uc_checkpoint.py
+++ b/tests/unit/checkpoint/test_autotp_uc_checkpoint.py
@@ -1105,7 +1105,7 @@ def test_save_convert_load_uneven_lm_head(self, tmpdir):
 
         tp_group = groups.get_tensor_model_parallel_group()
         tp_rank = groups.get_tensor_model_parallel_rank()
-        shard_sizes = get_shard_size_list(vocab_size, self.world_size, "lm_head")
+        shard_sizes = list(engine.module.lm_head._partition_sizes)
         assert shard_sizes == [51, 50], shard_sizes
         assert engine.module.lm_head.weight.shape[0] == shard_sizes[tp_rank]
 
@@ -1194,7 +1194,7 @@ def test_save_convert_load_uneven_row_parallel(self, tmpdir):
         # Column and row parallelism must shard the same dimension identically.
         assert attn.q_proj.weight.shape[0] == attn.o_proj.weight.shape[1]
         head_dim = hidden_dim // num_heads
-        head_shards = get_shard_size_list(num_heads, self.world_size, attn.q_proj.name)
+        head_shards = get_shard_size_list(num_heads, self.world_size, attn.q_proj.tp_meta, attn.q_proj.name)
         assert head_shards == [2, 2, 1, 1], head_shards
         dim_shards = [h * head_dim for h in head_shards]
         assert attn.q_proj.weight.shape[0] == dim_shards[tp_rank], (attn.q_proj.weight.shape, dim_shards)
@@ -1383,7 +1383,7 @@ def test_strict_correctness(self, tmpdir):
 
         tp_group = groups.get_tensor_model_parallel_group()
         tp_rank = groups.get_tensor_model_parallel_rank()
-        shard_sizes = get_shard_size_list(vocab_size, self.world_size, "lm_head")
+        shard_sizes = list(engine.module.lm_head._partition_sizes)
         assert sum(shard_sizes) == vocab_size, shard_sizes
         assert shard_sizes == [4, 3, 3], shard_sizes
         assert engine.module.lm_head.weight.shape[0] == shard_sizes[tp_rank]
@@ -1506,7 +1506,7 @@ def test_strict_correctness(self, tmpdir):
         assert attn.q_proj.weight.shape[0] == attn.o_proj.weight.shape[1], (attn.q_proj.weight.shape,
                                                                             attn.o_proj.weight.shape)
         head_dim = hidden_dim // num_heads
-        head_shards = get_shard_size_list(num_heads, self.world_size, attn.q_proj.name)
+        head_shards = get_shard_size_list(num_heads, self.world_size, attn.q_proj.tp_meta, attn.q_proj.name)
         assert head_shards == [2, 1, 1], head_shards
         dim_shards = [h * head_dim for h in head_shards]
         assert attn.q_proj.weight.shape[0] == dim_shards[tp_rank], (attn.q_proj.weight.shape, dim_shards)
diff --git a/tests/unit/module_inject/test_fused_repartition.py b/tests/unit/module_inject/test_fused_repartition.py
index 5bf54d1a207b..8819a3400ff2 100644
--- a/tests/unit/module_inject/test_fused_repartition.py
+++ b/tests/unit/module_inject/test_fused_repartition.py
@@ -2,62 +2,42 @@
 # SPDX-License-Identifier: Apache-2.0
 
 # DeepSpeed Team
-"""Fused AutoTP layers must keep the shard widths that were frozen when they were built.
+"""Fused AutoTP layers keep the shard widths frozen when they were built.
 
-``tp_shard`` keeps the grain size and the kv head count in process-wide globals that every
-AutoTP model overwrites while it is being replaced.  A layer that re-derives its split at
-partition time therefore cuts the weight differently once a second model has been loaded,
-which silently disagrees with the gather and with the universal checkpoint metadata.
+Each AutoTP model carries its own :class:`AutoTPMeta` (kv-head count, grain size, ...), so a
+layer's split stays aligned with its gather and its universal checkpoint metadata regardless
+of what other models do.  These tests build layers with an explicit ``tp_meta`` and check the
+split is stable.
 """
 
 import pytest
 import torch
 
-from deepspeed.module_inject import tp_shard
 from deepspeed.module_inject.layers import GateUpPack_LinearLayer, fused_LinearLayer
+from deepspeed.module_inject.tp_shard import AutoTPMeta
 
 
-@pytest.fixture
-def clean_tp_shard_globals():
-    yield
-    tp_shard.set_tp_grain_size(1)
-    tp_shard.set_num_kv_heads(None)
-
-
-def _build_gate_up_layer(out_features, tp_world_size, tp_index):
-    layer = GateUpPack_LinearLayer(torch.nn.Linear(3, out_features, bias=False), mp_group=None, name="dense_h_to_4h")
+def _build_gate_up_layer(out_features, tp_world_size, tp_index, meta):
+    layer = GateUpPack_LinearLayer(torch.nn.Linear(3, out_features, bias=False),
+                                   mp_group=None,
+                                   name="dense_h_to_4h",
+                                   tp_meta=meta)
     layer.tp_world_size = tp_world_size
     layer.tp_index = tp_index
     layer._freeze_partition_sizes(out_features)
     return layer
 
 
-def test_gate_up_partition_ignores_later_grain_size_changes(clean_tp_shard_globals):
-    tp_shard.set_tp_grain_size(1)
-    layer = _build_gate_up_layer(out_features=10, tp_world_size=2, tp_index=0)
-    assert layer._subparam_shard_widths == [[3, 2], [3, 2]]
-
-    full_weight = torch.arange(30, dtype=torch.float32).view(10, 3)
-
-    first = torch.nn.Parameter(full_weight.clone())
-    layer._tp_partition([first, None])
-
-    # A second AutoTP model would install its own grain size before this layer is gathered.
-    tp_shard.set_tp_grain_size(4)
-    second = torch.nn.Parameter(full_weight.clone())
-    layer._tp_partition([second, None])
-
-    assert tuple(second.shape) == (6, 3)
-    assert torch.equal(first.data, second.data)
-
-
-def test_gate_up_partition_covers_the_whole_weight(clean_tp_shard_globals):
-    tp_shard.set_tp_grain_size(1)
+def test_gate_up_partition_covers_the_whole_weight():
+    meta = AutoTPMeta(tp_grain_size=1)
     full_weight = torch.arange(30, dtype=torch.float32).view(10, 3)
 
     shards = []
     for tp_index in range(2):
-        layer = _build_gate_up_layer(out_features=10, tp_world_size=2, tp_index=tp_index)
+        layer = _build_gate_up_layer(out_features=10, tp_world_size=2, tp_index=tp_index, meta=meta)
+        # The gate and up halves (5 each) split over 2 ranks as [3, 2]; the layer freezes these
+        # widths from its own tp_meta so the partition is deterministic.
+        assert layer._subparam_shard_widths == [[3, 2], [3, 2]]
         param = torch.nn.Parameter(full_weight.clone())
         layer._tp_partition([param, None])
         shards.append(param.data)
@@ -81,25 +61,23 @@ def __init__(self, split_size):
         self.attn = _QWenAttention(split_size)
 
 
-def _build_qwen_attn_layer(block, hidden, tp_world_size, tp_index):
+def _build_qwen_attn_layer(block, hidden, tp_world_size, tp_index, meta):
     layer = fused_LinearLayer(torch.nn.Linear(hidden, 3 * hidden, bias=False),
                               mp_group=None,
                               skip_partition=True,
                               name="c_attn",
-                              fused_module=block)
+                              fused_module=block,
+                              tp_meta=meta)
     layer.tp_world_size = tp_world_size
     layer.tp_index = tp_index
     layer._freeze_partition_sizes(3 * hidden)
     return layer
 
 
-def test_qwen_split_size_follows_the_frozen_shard_width(clean_tp_shard_globals):
-    tp_shard.set_tp_grain_size(1)
-    tp_shard.set_num_kv_heads(4)
-    tp_shard.set_n_embd(12)
-
+def test_qwen_split_size_follows_the_frozen_shard_width():
+    meta = AutoTPMeta(num_kv_heads=4, n_embd=12, tp_grain_size=1)
     block = _QWenBlock(split_size=12)
-    layer = _build_qwen_attn_layer(block, hidden=12, tp_world_size=4, tp_index=3)
+    layer = _build_qwen_attn_layer(block, hidden=12, tp_world_size=4, tp_index=3, meta=meta)
 
     weight = torch.nn.Parameter(torch.zeros(36, 12))
     layer._tp_partition([weight, None])
@@ -109,29 +87,26 @@ def test_qwen_split_size_follows_the_frozen_shard_width(clean_tp_shard_globals):
     assert weight.shape[0] == 3 * block.attn.split_size
 
 
-def test_qwen_rejects_a_tensor_parallel_size_that_empties_a_rank(clean_tp_shard_globals):
-    tp_shard.set_tp_grain_size(1)
-    tp_shard.set_num_kv_heads(4)
-    tp_shard.set_n_embd(12)
+def test_qwen_rejects_a_tensor_parallel_size_that_empties_a_rank():
+    meta = AutoTPMeta(num_kv_heads=4, n_embd=12, tp_grain_size=1)
 
     with pytest.raises(RuntimeError, match="empty query/key/value shard"):
-        _build_qwen_attn_layer(_QWenBlock(split_size=12), hidden=12, tp_world_size=16, tp_index=15)
+        _build_qwen_attn_layer(_QWenBlock(split_size=12), hidden=12, tp_world_size=16, tp_index=15, meta=meta)
 
 
 class _CodeGenBlock(torch.nn.Module):
     pass
 
 
-def test_interleaved_fused_layout_refuses_to_gather(clean_tp_shard_globals):
-    tp_shard.set_tp_grain_size(1)
-    tp_shard.set_num_kv_heads(8)
-    tp_shard.set_n_embd(4)
+def test_interleaved_fused_layout_refuses_to_gather():
+    meta = AutoTPMeta(num_kv_heads=8, n_embd=4, tp_grain_size=1)
 
     layer = fused_LinearLayer(torch.nn.Linear(4, 24, bias=False),
                               mp_group=None,
                               skip_partition=True,
                               name="qkv_proj",
-                              fused_module=_CodeGenBlock())
+                              fused_module=_CodeGenBlock(),
+                              tp_meta=meta)
     layer.tp_world_size = 2
     layer.tp_index = 0
     layer._freeze_partition_sizes(24)
diff --git a/tests/unit/module_inject/test_tp_partition_config_path.py b/tests/unit/module_inject/test_tp_partition_config_path.py
index 76fbb6d28e20..5b268f2a6df5 100644
--- a/tests/unit/module_inject/test_tp_partition_config_path.py
+++ b/tests/unit/module_inject/test_tp_partition_config_path.py
@@ -86,6 +86,7 @@ def capture(self, child, full_name):
             linear_layer_setting=None,
             orig_layer_impl=None,
             partition_config=config,
+            model_config=getattr(model, "config", None),
         )
         autotp._replace_module(model)
     finally:
@@ -152,6 +153,7 @@ def _build_gathered_lm_head_autotp(model, mp_size=1):
         linear_layer_setting=None,
         orig_layer_impl=None,
         partition_config=config,
+        model_config=getattr(model, "config", None),
     )
     autotp.set_tensor_parallel_config(mp_size, None)
     autotp.update_linear_policies()
diff --git a/tests/unit/module_inject/test_tp_shard.py b/tests/unit/module_inject/test_tp_shard.py
index dc72ecde8bfb..b37de1425942 100644
--- a/tests/unit/module_inject/test_tp_shard.py
+++ b/tests/unit/module_inject/test_tp_shard.py
@@ -4,28 +4,19 @@
 # DeepSpeed Team
 
 import pytest
+from types import SimpleNamespace
 
 from deepspeed.module_inject import tp_shard
-from deepspeed.module_inject.tp_shard import get_shard_size, get_shard_size_list, set_num_kv_heads, set_tp_grain_size
-
-
-@pytest.fixture(autouse=True)
-def restore_tp_shard_globals():
-    # tp_grain_size and num_kv_heads are process wide, so leaking them would change how
-    # unrelated tests partition their layers.
-    grain_size, kv_heads = tp_shard.tp_grain_size, tp_shard.num_kv_heads
-    yield
-    set_tp_grain_size(grain_size)
-    set_num_kv_heads(kv_heads)
+from deepspeed.module_inject.tp_shard import AutoTPMeta, get_shard_size, get_shard_size_list
 
 
 @pytest.mark.parametrize("total_size,tp_size", [(50257, 2), (50257, 8), (151936, 8), (32000, 4)])
 def test_grain_quantized_shards_tile_the_dimension(total_size, tp_size):
     # A vocabulary that is not a multiple of tp_grain_size used to lose its tail to the grain
     # quantization, so the shards no longer reconstructed the embedding table.
-    set_tp_grain_size(64)
+    meta = AutoTPMeta(tp_grain_size=64)
 
-    shard_sizes = get_shard_size_list(total_size, tp_size, "lm_head")
+    shard_sizes = get_shard_size_list(total_size, tp_size, meta, "lm_head")
 
     assert sum(shard_sizes) == total_size
     # Only the rank that absorbs the sub-grain tail gives up its alignment.
@@ -33,59 +24,61 @@ def test_grain_quantized_shards_tile_the_dimension(total_size, tp_size):
 
 
 def test_uneven_shards_without_grain_quantization():
-    assert get_shard_size_list(101, 2, "lm_head") == [51, 50]
+    assert get_shard_size_list(101, 2, AutoTPMeta(), "lm_head") == [51, 50]
 
 
-def test_kv_head_shards_tile_the_dimension():
-    set_num_kv_heads(6)
+def test_meta_descends_into_multimodal_text_config():
+    # Vision-language outer configs keep the head counts only under text_config; reading the
+    # outer config directly would lose them and fall back to an even-grain split.
+    outer = SimpleNamespace(text_config=SimpleNamespace(num_key_value_heads=4, num_attention_heads=8, hidden_size=64))
 
-    # 6 kv heads over 4 ranks gives 2/2/1/1 heads, so 384 hidden splits as 128/128/64/64.
-    assert get_shard_size_list(384, 4, "layers.0.self_attn.q_proj") == [128, 128, 64, 64]
+    meta = AutoTPMeta.from_model_config(outer)
 
+    assert meta == AutoTPMeta(num_kv_heads=4, num_attention_heads=8, n_embd=64)
 
-def test_process_group_resolves_noncontiguous_group_rank(monkeypatch):
-    set_tp_grain_size(64)
-    tp_group = object()
-    monkeypatch.setattr(tp_shard.dist, "get_rank", lambda group=None: 1 if group is tp_group else 2)
 
-    shard_sizes = get_shard_size_list(50257, 2, "lm_head")
-    assert get_shard_size(50257, 2, "lm_head", mp_group=tp_group) == shard_sizes[1]
+def test_meta_descends_only_when_text_config_is_set():
+    plain = SimpleNamespace(num_key_value_heads=3, num_attention_heads=6, hidden_size=32)
 
+    meta = AutoTPMeta.from_model_config(plain)
 
-def test_shard_size_refuses_to_guess_subgroup_rank(monkeypatch):
-    monkeypatch.setattr(tp_shard.dist, "get_world_size", lambda: 4)
+    assert meta == AutoTPMeta(num_kv_heads=3, num_attention_heads=6, n_embd=32)
 
-    with pytest.raises(ValueError, match="group-local rank or process group"):
-        get_shard_size(12, 2)
 
+# 6 kv heads over 4 ranks gives 2/2/1/1 heads, so an attention projection splits 384 as
+# 128/128/64/64. Every other layer kind keeps the near-even 96/96/96/96 split, either because
+# the dimension has no head structure or because even shards serve it better.
+@pytest.mark.parametrize("name,expected", [
+    ("layers.0.self_attn.q_proj", [128, 128, 64, 64]),
+    ("layers.0.mlp.dense_h_to_4h", [96, 96, 96, 96]),
+    ("lm_head", [96, 96, 96, 96]),
+    ("embed_out", [96, 96, 96, 96]),
+    ("layers.0.experts.w1", [96, 96, 96, 96]),
+])
+def test_kv_head_split_applies_only_to_attention(name, expected):
+    meta = AutoTPMeta(num_kv_heads=6)
 
-def test_explicit_num_kv_heads_is_used():
-    assert get_shard_size_list(
-        384,
-        4,
-        "self_attn.q_proj",
-        num_kv_heads=6,
-    ) == [128, 128, 64, 64]
+    assert get_shard_size_list(384, 4, meta, name) == expected
 
 
-def test_explicit_num_kv_heads_matches_global_value():
-    set_num_kv_heads(6)
-    expected = get_shard_size_list(384, 4, "self_attn.q_proj")
+def test_kv_head_split_needs_a_divisible_dimension():
+    # 385 is not a multiple of the 6 kv heads, so there is no head-aligned split to make.
+    meta = AutoTPMeta(num_kv_heads=6)
 
-    actual = get_shard_size_list(
-        384,
-        4,
-        "self_attn.q_proj",
-        num_kv_heads=6,
-    )
+    assert get_shard_size_list(385, 4, meta, "layers.0.self_attn.q_proj") == [97, 96, 96, 96]
 
-    assert actual == expected
+
+def test_process_group_resolves_noncontiguous_group_rank(monkeypatch):
+    meta = AutoTPMeta(tp_grain_size=64)
+    tp_group = object()
+    monkeypatch.setattr(tp_shard.dist, "get_rank", lambda group=None: 1 if group is tp_group else 2)
+
+    shard_sizes = get_shard_size_list(50257, 2, meta, "lm_head")
+    assert get_shard_size(50257, 2, meta, "lm_head", mp_group=tp_group) == shard_sizes[1]
 
 
-def test_uneven_shards_without_grain_quantization_no_kv_heads_used():
-    assert get_shard_size_list(
-        101,
-        2,
-        "lm_head",
-        num_kv_heads=None,
-    ) == [51, 50]
+def test_shard_size_refuses_to_guess_subgroup_rank(monkeypatch):
+    monkeypatch.setattr(tp_shard.dist, "get_world_size", lambda: 4)
+
+    with pytest.raises(ValueError, match="group-local rank or process group"):
+        get_shard_size(12, 2, AutoTPMeta())
diff --git a/tests/unit/runtime/tensor_parallel/test_autotp_universal_checkpoint.py b/tests/unit/runtime/tensor_parallel/test_autotp_universal_checkpoint.py
index fe996e267f3c..2a1159359241 100644
--- a/tests/unit/runtime/tensor_parallel/test_autotp_universal_checkpoint.py
+++ b/tests/unit/runtime/tensor_parallel/test_autotp_universal_checkpoint.py
@@ -12,6 +12,7 @@
                                             VOCABULARY_PARAMETER_PATTERNS, DS_AUTOTP_UC_META,
                                             UNIVERSAL_CHECKPOINT_VERSION_VALUE)
 from deepspeed.checkpoint.universal_checkpoint import _narrow_sub_params, _resolve_autotp_partition
+from deepspeed.module_inject.tp_shard import AutoTPMeta
 from deepspeed.module_inject.layers import (_build_param_uc_restore_meta, _get_param_uc_conversion_meta,
                                             _subparam_shard_widths, GateUpPack_LinearLayer, LinearAllreduce,
                                             LinearLayer, SubParamLinearAllreduce, SubParamLinearLayer,
@@ -340,7 +341,8 @@ class CodeGenBlock(torch.nn.Module):
                               mp_group=None,
                               skip_partition=True,
                               fused_module=fused_module,
-                              name="qkv_proj")
+                              name="qkv_proj",
+                              tp_meta=AutoTPMeta())
     model = torch.nn.Module()
     model.qkv_proj = layer
 
@@ -403,7 +405,7 @@ def test_sub_param_shard_widths_round_trip_with_zero_width_ranks(tp_world_size):
     # More ranks than kv heads leaves some ranks holding none of a sub-parameter. Those empty
     # shards still have to tile the sub-parameter and survive a restore round trip.
     sub_param_sizes = (8, 2, 2)
-    widths = _subparam_shard_widths(sub_param_sizes, tp_world_size)
+    widths = _subparam_shard_widths(sub_param_sizes, tp_world_size, AutoTPMeta())
 
     assert any(width == 0 for per_rank in widths for width in per_rank)
     for size, per_rank in zip(sub_param_sizes, widths):
@@ -474,27 +476,35 @@ def test_sub_param_layer_materializes_zero_width_final_dimension(layer_cls):
 
 
 def test_lm_head_forward_uses_frozen_partition_sizes():
-    # The weight columns were cut when the layer was built. A second AutoTP model overwrites the
-    # process-wide tp_shard globals, so re-deriving the split in forward would slice the input
-    # differently than the weight and the row-parallel all-reduce would hide the mismatch.
-    from deepspeed.module_inject import tp_shard
+    # The weight columns were cut when the layer was built; forward must slice the input with
+    # the same frozen partition sizes rather than re-derive them.
     from deepspeed.module_inject.layers import LmHeadLinearAllreduce
+    from deepspeed.module_inject.tp_shard import AutoTPMeta
 
-    grain_size, kv_heads = tp_shard.tp_grain_size, tp_shard.num_kv_heads
-    try:
-        tp_shard.set_tp_grain_size(1)
-        tp_shard.set_num_kv_heads(None)
-        layer = LmHeadLinearAllreduce(torch.nn.Linear(101, 8, bias=False), mp_group=None)
-        layer.tp_world_size = 2
-        layer.tp_index = 1
-        frozen = layer._freeze_partition_sizes(101)
-        assert frozen == (51, 50)
-
-        # A later model narrows the grain, which would change a recomputed split.
-        tp_shard.set_tp_grain_size(64)
-        layer.weight.data = torch.zeros(8, frozen[1])
-
-        layer(torch.zeros(1, 1, 101))
-    finally:
-        tp_shard.set_tp_grain_size(grain_size)
-        tp_shard.set_num_kv_heads(kv_heads)
+    layer = LmHeadLinearAllreduce(torch.nn.Linear(101, 8, bias=False),
+                                  mp_group=None,
+                                  tp_meta=AutoTPMeta(tp_grain_size=1))
+    layer.tp_world_size = 2
+    layer.tp_index = 1
+    frozen = layer._freeze_partition_sizes(101)
+    assert frozen == (51, 50)
+    assert layer.tp_meta.tp_grain_size == 1
+    layer.weight.data = torch.zeros(8, frozen[1])
+
+    # A second layer of a different model is then built and initialized in the same process:
+    # a coarser grain (vocabulary sharding) and a different kv-head count (GQA attention
+    # sharding). Neither may leak into the first layer, whose meta and frozen split describe
+    # its own model only.
+    second = LmHeadLinearAllreduce(torch.nn.Linear(101, 8, bias=False),
+                                   mp_group=None,
+                                   tp_meta=AutoTPMeta(tp_grain_size=64, num_kv_heads=1))
+    second.tp_world_size = 2
+    second.tp_index = 1
+    assert second._freeze_partition_sizes(101) == (64, 37)
+
+    assert layer.tp_meta.tp_grain_size == 1
+    assert layer.tp_meta.num_kv_heads is None
+    assert layer._freeze_partition_sizes(101) == frozen
+    assert layer._partition_sizes == frozen
+
+    layer(torch.zeros(1, 1, 101))
diff --git a/tests/unit/sequence_parallelism/test_ulysses.py b/tests/unit/sequence_parallelism/test_ulysses.py
index abb9899892a8..72e8a81c8048 100644
--- a/tests/unit/sequence_parallelism/test_ulysses.py
+++ b/tests/unit/sequence_parallelism/test_ulysses.py
@@ -16,7 +16,7 @@
 from unit.util import skip_on_arch
 from unit.simple_model import *
 from deepspeed.utils import groups
-from deepspeed.module_inject.tp_shard import get_shard_size_list
+from deepspeed.module_inject.tp_shard import AutoTPMeta, get_shard_size_list
 #Use mesh device to create data and sequence parallel group
 
 
@@ -205,7 +205,7 @@ def seq_batch_heads_hash(d0, d1, h, offset_d0=0, offset_d1=0, offset_h=0):
             d0_indices = torch.arange(s2h_tensor.shape[0]).reshape(-1, 1, 1, 1)
             d1_indices = torch.arange(s2h_tensor.shape[1]).reshape(1, -1, 1, 1)
             h_indices = torch.arange(s2h_tensor.shape[2]).reshape(1, 1, -1, 1)
-            shard_list = get_shard_size_list(num_heads, groups._get_sequence_parallel_world_size())
+            shard_list = get_shard_size_list(num_heads, groups._get_sequence_parallel_world_size(), AutoTPMeta())
             head_offset = sum(shard_list[:groups._get_sequence_parallel_rank()])
             s2h_truth = torch.zeros_like(s2h_tensor)
             s2h_truth[:] = seq_batch_heads_hash(d0_indices, d1_indices, h_indices, 0, 0, head_offset)
diff --git a/tests/unit/model_parallelism/test_autotp_custom_patterns.py b/tests/unit/v1/autotp/test_autotp_custom_patterns.py
similarity index 81%
rename from tests/unit/model_parallelism/test_autotp_custom_patterns.py
rename to tests/unit/v1/autotp/test_autotp_custom_patterns.py
index a4104f0ede4d..63c714b42723 100644
--- a/tests/unit/model_parallelism/test_autotp_custom_patterns.py
+++ b/tests/unit/v1/autotp/test_autotp_custom_patterns.py
@@ -19,7 +19,7 @@
 from deepspeed.checkpoint.constants import (DS_AUTOTP_UC_META, PARAMETER_WITH_ROW_PARALLELISM_PATTERNS,
                                             TP_REPLICATED_PARAMETER_PATTERNS)
 from deepspeed.module_inject.autotp_config import AutoTPConfig
-from deepspeed.module_inject.tp_shard import get_shard_size, get_shard_size_list, set_num_kv_heads
+from deepspeed.module_inject.tp_shard import AutoTPMeta, get_shard_size, get_shard_size_list
 from deepspeed.module_inject.auto_tp import AutoTP
 from deepspeed.module_inject.auto_tp_model_utils import (build_bloom_alibi_tensor, build_mpt_alibi_tensor,
                                                          get_alibi_mask, install_head_sharded_helper)
@@ -170,7 +170,8 @@ def apply_autotp_with_partition_config(model, tp_size, partition_config):
                     linear_layer_setting=None,
                     orig_layer_impl=None,
                     keep_module_on_host=False,
-                    partition_config=autotp_config)
+                    partition_config=autotp_config,
+                    model_config=getattr(model, "config", None))
     autotp.set_tensor_parallel_config(tp_size, groups.get_tensor_model_parallel_group())
     autotp.update_linear_policies()
     autotp._replace_module(model)
@@ -502,11 +503,8 @@ def test_update_mp_params_uses_group_local_rank(monkeypatch):
     child.num_heads = 12
 
     monkeypatch.setattr(dist, "get_rank", lambda group=None: 1 if group is tp_group else 0)
-    set_num_kv_heads(3)
-    try:
-        autotp.update_mp_params(child)
-    finally:
-        set_num_kv_heads(None)
+    autotp.tp_meta = AutoTPMeta(num_kv_heads=3)
+    autotp.update_mp_params(child)
 
     # Three KV groups split as [2, 1], so the second TP rank owns four query heads.
     assert child.num_heads == 4
@@ -523,15 +521,13 @@ def test_update_mp_params_shards_attributes_like_their_weights(monkeypatch):
     child.hidden_size = 12
 
     monkeypatch.setattr(dist, "get_rank", lambda group=None: 1 if group is tp_group else 0)
-    set_num_kv_heads(3)
-    try:
-        autotp.update_mp_params(child, "model.layers.0.mlp")
-    finally:
-        set_num_kv_heads(None)
+    meta = AutoTPMeta(num_kv_heads=3)
+    autotp.tp_meta = meta
+    autotp.update_mp_params(child, "model.layers.0.mlp")
 
     # MLP layers are excluded from the KV-head split, so the attribute has to follow the same
     # near-even split that the MLP weights use rather than the [2, 1] KV-group split.
-    assert child.hidden_size == get_shard_size(12, 2, "model.layers.0.mlp", rank=1)
+    assert child.hidden_size == get_shard_size(12, 2, meta, "model.layers.0.mlp", rank=1)
 
 
 def test_update_mp_params_preserves_unsharded_high_dimensional_modules(monkeypatch):
@@ -544,11 +540,8 @@ def test_update_mp_params_preserves_unsharded_high_dimensional_modules(monkeypat
     child.proj = nn.Conv3d(3, child.embed_dim, kernel_size=1)
 
     monkeypatch.setattr(dist, "get_rank", lambda group=None: 1 if group is tp_group else 0)
-    set_num_kv_heads(3)
-    try:
-        autotp.update_mp_params(child, "model.visual.patch_embed")
-    finally:
-        set_num_kv_heads(None)
+    autotp.tp_meta = AutoTPMeta(num_kv_heads=3)
+    autotp.update_mp_params(child, "model.visual.patch_embed")
 
     assert child.embed_dim == 12
 
@@ -569,11 +562,8 @@ def test_replace_module_preserves_metadata_for_unsharded_subtree(monkeypatch):
     visual.attn.proj = nn.Linear(12, 12)
 
     monkeypatch.setattr(dist, "get_rank", lambda group=None: 1 if group is tp_group else 0)
-    set_num_kv_heads(None)
-    try:
-        autotp._replace_module(visual, "model.visual")
-    finally:
-        set_num_kv_heads(None)
+    autotp.tp_meta = AutoTPMeta()
+    autotp._replace_module(visual, "model.visual")
 
     assert visual.attn.num_heads == 16
 
@@ -589,11 +579,8 @@ def test_update_mp_params_follows_actual_tp_parameter_metadata(monkeypatch):
     setattr(child.proj.weight, DS_AUTOTP_UC_META, {})
 
     monkeypatch.setattr(dist, "get_rank", lambda group=None: 1 if group is tp_group else 0)
-    set_num_kv_heads(None)
-    try:
-        autotp.update_mp_params(child, "model.layers.0.mlp")
-    finally:
-        set_num_kv_heads(None)
+    autotp.tp_meta = AutoTPMeta()
+    autotp.update_mp_params(child, "model.layers.0.mlp")
 
     assert child.hidden_size == 6
 
@@ -603,6 +590,7 @@ def test_sliced_embedding_publishes_row_partition_metadata(monkeypatch):
     autotp = object.__new__(AutoTP)
     autotp.mp_group = tp_group
     autotp.mp_size = 2
+    autotp.tp_meta = AutoTPMeta()
     embedding = nn.Embedding(5, 4)
 
     monkeypatch.setattr(dist, "get_rank", lambda group=None: 1 if group is tp_group else 0)
@@ -625,66 +613,70 @@ def test_mpt_alibi_covers_every_head_of_an_uneven_split(self):
         init_tp_engine(tp_size=2)
 
         num_heads = 5
-        set_num_kv_heads(num_heads)
-        try:
 
-            class MptTransformer(nn.Module):
+        class MptTransformer(nn.Module):
 
-                def build_mpt_alibi_tensor(self, heads, sequence_length, alibi_bias_max=8, device=None):
-                    return torch.arange(heads, dtype=torch.float32).view(heads, 1, 1).expand(heads, 1, sequence_length)
+            def build_mpt_alibi_tensor(self, heads, sequence_length, alibi_bias_max=8, device=None):
+                return torch.arange(heads, dtype=torch.float32).view(heads, 1, 1).expand(heads, 1, sequence_length)
 
-            transformer = MptTransformer()
-            install_head_sharded_helper(transformer, 'build_mpt_alibi_tensor', build_mpt_alibi_tensor)
+        transformer = MptTransformer()
+        install_head_sharded_helper(transformer,
+                                    'build_mpt_alibi_tensor',
+                                    build_mpt_alibi_tensor,
+                                    meta=AutoTPMeta(num_attention_heads=num_heads, num_kv_heads=num_heads))
 
-            alibi = transformer.build_mpt_alibi_tensor(num_heads, 3)
+        alibi = transformer.build_mpt_alibi_tensor(num_heads, 3)
 
-            # AutoTP splits 5 heads over 2 ranks as [3, 2]; an even split would give every rank
-            # 2 heads and drop the last one entirely.
-            expected_heads = get_shard_size_list(num_heads, dist.get_world_size())
-            offset = sum(expected_heads[:dist.get_rank()])
-            assert alibi.shape[0] == expected_heads[dist.get_rank()]
-            torch.testing.assert_close(alibi[:, 0, 0].cpu(),
-                                       torch.arange(offset, offset + alibi.shape[0], dtype=torch.float32))
-        finally:
-            set_num_kv_heads(None)
+        # AutoTP splits 5 heads over 2 ranks as [3, 2]; an even split would give every rank
+        # 2 heads and drop the last one entirely.
+        expected_heads = get_shard_size_list(num_heads, dist.get_world_size(), AutoTPMeta(num_kv_heads=num_heads))
+        offset = sum(expected_heads[:dist.get_rank()])
+        assert alibi.shape[0] == expected_heads[dist.get_rank()]
+        torch.testing.assert_close(alibi[:, 0, 0].cpu(),
+                                   torch.arange(offset, offset + alibi.shape[0], dtype=torch.float32))
 
     def test_head_sharded_helper_leaves_the_class_untouched(self):
         skip_on_device()
         init_tp_engine(tp_size=2)
 
         num_heads = 5
-        set_num_kv_heads(num_heads)
-        try:
-
-            class MptTransformer(nn.Module):
 
-                def build_mpt_alibi_tensor(self, heads, sequence_length, alibi_bias_max=8, device=None):
-                    return torch.arange(heads, dtype=torch.float32).view(heads, 1, 1).expand(heads, 1, sequence_length)
-
-            class MptSubclass(MptTransformer):
-                pass
-
-            first = MptTransformer()
-            install_head_sharded_helper(first, 'build_mpt_alibi_tensor', build_mpt_alibi_tensor)
-            expected = first.build_mpt_alibi_tensor(num_heads, 3)
-
-            # Injecting a second model of the same architecture must not make either of them
-            # delegate to the other's wrapper.
-            second = MptTransformer()
-            install_head_sharded_helper(second, 'build_mpt_alibi_tensor', build_mpt_alibi_tensor)
-            torch.testing.assert_close(second.build_mpt_alibi_tensor(num_heads, 3), expected)
-            torch.testing.assert_close(first.build_mpt_alibi_tensor(num_heads, 3), expected)
+        class MptTransformer(nn.Module):
 
-            # A model of the same class that was never injected keeps its own method, and a
-            # subclass of it inherits that method rather than an installed wrapper.
-            plain = MptTransformer()
-            assert plain.build_mpt_alibi_tensor(num_heads, 3).shape[0] == num_heads
+            def build_mpt_alibi_tensor(self, heads, sequence_length, alibi_bias_max=8, device=None):
+                return torch.arange(heads, dtype=torch.float32).view(heads, 1, 1).expand(heads, 1, sequence_length)
 
-            derived = MptSubclass()
-            install_head_sharded_helper(derived, 'build_mpt_alibi_tensor', build_mpt_alibi_tensor)
-            torch.testing.assert_close(derived.build_mpt_alibi_tensor(num_heads, 3), expected)
-        finally:
-            set_num_kv_heads(None)
+        class MptSubclass(MptTransformer):
+            pass
+
+        first = MptTransformer()
+        install_head_sharded_helper(first,
+                                    'build_mpt_alibi_tensor',
+                                    build_mpt_alibi_tensor,
+                                    meta=AutoTPMeta(num_attention_heads=num_heads, num_kv_heads=num_heads))
+        expected = first.build_mpt_alibi_tensor(num_heads, 3)
+
+        # Injecting a second model of the same architecture must not make either of them
+        # delegate to the other's wrapper.
+        second = MptTransformer()
+        install_head_sharded_helper(second,
+                                    'build_mpt_alibi_tensor',
+                                    build_mpt_alibi_tensor,
+                                    meta=AutoTPMeta(num_attention_heads=num_heads, num_kv_heads=num_heads))
+        torch.testing.assert_close(second.build_mpt_alibi_tensor(num_heads, 3), expected)
+        torch.testing.assert_close(first.build_mpt_alibi_tensor(num_heads, 3), expected)
+
+        # A model of the same class that was never injected keeps its own method, and a
+        # subclass of it inherits that method rather than an installed wrapper.
+        plain = MptTransformer()
+        assert plain.build_mpt_alibi_tensor(num_heads, 3).shape[0] == num_heads
+
+        derived = MptSubclass()
+        install_head_sharded_helper(derived,
+                                    'build_mpt_alibi_tensor',
+                                    build_mpt_alibi_tensor,
+                                    meta=AutoTPMeta(num_attention_heads=num_heads, num_kv_heads=num_heads))
+        torch.testing.assert_close(derived.build_mpt_alibi_tensor(num_heads, 3), expected)
 
     def test_head_sharded_helper_freezes_the_models_split(self):
         skip_on_device()
@@ -696,29 +688,23 @@ def build_mpt_alibi_tensor(self, heads, sequence_length, alibi_bias_max=8, devic
                 return torch.arange(heads, dtype=torch.float32).view(heads, 1, 1).expand(heads, 1, sequence_length)
 
         num_heads = 6
-        set_num_kv_heads(3)
-        try:
-            transformer = MptTransformer()
-            install_head_sharded_helper(transformer,
-                                        'build_mpt_alibi_tensor',
-                                        build_mpt_alibi_tensor,
-                                        num_heads=num_heads,
-                                        num_kv_heads=3)
-
-            # Initializing another model can replace this process-wide setting. The first
-            # model's helper must keep the [4, 2] split frozen with its weights.
-            set_num_kv_heads(2)
-            expected_sizes = [4, 2]
-            # AutoTP replaces the model's public head count with this rank's local count.
-            local_num_heads = expected_sizes[dist.get_rank()]
-            alibi = transformer.build_mpt_alibi_tensor(local_num_heads, 3)
-
-            offset = sum(expected_sizes[:dist.get_rank()])
-            assert alibi.shape[0] == expected_sizes[dist.get_rank()]
-            torch.testing.assert_close(alibi[:, 0, 0].cpu(),
-                                       torch.arange(offset, offset + alibi.shape[0], dtype=torch.float32))
-        finally:
-            set_num_kv_heads(None)
+        transformer = MptTransformer()
+        install_head_sharded_helper(transformer,
+                                    'build_mpt_alibi_tensor',
+                                    build_mpt_alibi_tensor,
+                                    meta=AutoTPMeta(num_attention_heads=num_heads, num_kv_heads=3))
+
+        # The helper freezes the [4, 2] split from this model's own num_kv_heads at install
+        # time.
+        expected_sizes = [4, 2]
+        # AutoTP replaces the model's public head count with this rank's local count.
+        local_num_heads = expected_sizes[dist.get_rank()]
+        alibi = transformer.build_mpt_alibi_tensor(local_num_heads, 3)
+
+        offset = sum(expected_sizes[:dist.get_rank()])
+        assert alibi.shape[0] == expected_sizes[dist.get_rank()]
+        torch.testing.assert_close(alibi[:, 0, 0].cpu(),
+                                   torch.arange(offset, offset + alibi.shape[0], dtype=torch.float32))
 
     def test_bloom_alibi_uses_original_total_after_injection(self):
         skip_on_device()
@@ -753,7 +739,10 @@ def get_alibi_mask(self, tensor, sequence_length):
                                     dtype=torch.float32).view(-1, 1, 1).expand(-1, sequence_length, sequence_length)
 
         model = AlibiModel()
-        install_head_sharded_helper(model, 'get_alibi_mask', get_alibi_mask, num_heads=5, num_kv_heads=5)
+        install_head_sharded_helper(model,
+                                    'get_alibi_mask',
+                                    get_alibi_mask,
+                                    meta=AutoTPMeta(num_attention_heads=5, num_kv_heads=5))
 
         shard_sizes = [3, 2]
         model.n_head = shard_sizes[dist.get_rank()]
@@ -780,24 +769,22 @@ def test_gate_up_gather_restores_sub_param_order(self):
 
         hidden_dim = 4
         gate_up_dim = 6
-        set_num_kv_heads(3)
-        try:
-            torch.manual_seed(17)
-            linear = nn.Linear(hidden_dim,
-                               gate_up_dim * 2,
-                               bias=False,
-                               dtype=preferred_dtype(),
-                               device=get_accelerator().current_device_name())
-            full_weight = deepcopy(linear.weight.data)
-            layer = GateUpPack_LinearLayer(deepcopy(linear), groups.get_tensor_model_parallel_group())
-
-            # The gate and the up halves are each cut in two, so a rank-order concatenation of
-            # the shards would interleave them instead of restoring the original weight.
-            gathered = nn.Parameter(layer.weight.data.clone())
-            layer.gather_params([gathered, None])
-            torch.testing.assert_close(gathered.data, full_weight)
-        finally:
-            set_num_kv_heads(None)
+        torch.manual_seed(17)
+        linear = nn.Linear(hidden_dim,
+                           gate_up_dim * 2,
+                           bias=False,
+                           dtype=preferred_dtype(),
+                           device=get_accelerator().current_device_name())
+        full_weight = deepcopy(linear.weight.data)
+        layer = GateUpPack_LinearLayer(deepcopy(linear),
+                                       groups.get_tensor_model_parallel_group(),
+                                       tp_meta=AutoTPMeta(num_kv_heads=3))
+
+        # The gate and the up halves are each cut in two, so a rank-order concatenation of
+        # the shards would interleave them instead of restoring the original weight.
+        gathered = nn.Parameter(layer.weight.data.clone())
+        layer.gather_params([gathered, None])
+        torch.testing.assert_close(gathered.data, full_weight)
 
     def test_gate_up_fused_weight_partition(self):
         skip_on_device()
@@ -951,29 +938,26 @@ def test_gather_uses_the_layers_own_shard_widths(self):
 
         hidden_dim = 8
         head_size = 12
-        set_num_kv_heads(3)
-        try:
-            torch.manual_seed(7)
-            linear = nn.Linear(hidden_dim,
-                               head_size * 3,
-                               bias=True,
-                               dtype=preferred_dtype(),
-                               device=get_accelerator().current_device_name())
-            full_weight = deepcopy(linear.weight.data)
-            full_bias = deepcopy(linear.bias.data)
-
-            layer = SubParamLinearLayer(deepcopy(linear),
-                                        groups.get_tensor_model_parallel_group(),
-                                        shape=((head_size, head_size, head_size), -1),
-                                        partition_dim=0,
-                                        name="self_attn.qkv_proj")
-            assert layer._subparam_shard_widths == [[8, 4], [8, 4], [8, 4]]
+        torch.manual_seed(7)
+        linear = nn.Linear(hidden_dim,
+                           head_size * 3,
+                           bias=True,
+                           dtype=preferred_dtype(),
+                           device=get_accelerator().current_device_name())
+        full_weight = deepcopy(linear.weight.data)
+        full_bias = deepcopy(linear.bias.data)
 
-            layer.gather_params([layer.weight, layer.bias])
-            torch.testing.assert_close(layer.weight.data, full_weight)
-            torch.testing.assert_close(layer.bias.data, full_bias)
-        finally:
-            set_num_kv_heads(None)
+        layer = SubParamLinearLayer(deepcopy(linear),
+                                    groups.get_tensor_model_parallel_group(),
+                                    shape=((head_size, head_size, head_size), -1),
+                                    partition_dim=0,
+                                    name="self_attn.qkv_proj",
+                                    tp_meta=AutoTPMeta(num_kv_heads=3))
+        assert layer._subparam_shard_widths == [[8, 4], [8, 4], [8, 4]]
+
+        layer.gather_params([layer.weight, layer.bias])
+        torch.testing.assert_close(layer.weight.data, full_weight)
+        torch.testing.assert_close(layer.bias.data, full_bias)
 
     def test_gqa_uneven_qkv_fused_forward(self):
         skip_on_device()
@@ -1001,3 +985,7 @@ def test_gqa_uneven_qkv_fused_forward(self):
         gathered_output = gather_subparam_output(tp_output, (q_size, k_size, v_size),
                                                  groups.get_tensor_model_parallel_group())
         assert_close_for_preferred_dtype(gathered_output, full_output)
+
+
+# Multi-model (teacher + student in one process) and multimodal-config regressions live in
+# tests/unit/v1/autotp/test_autotp_multiple_models.py so they run on the GPU workflow too.
diff --git a/tests/unit/v1/autotp/test_autotp_multiple_models.py b/tests/unit/v1/autotp/test_autotp_multiple_models.py
new file mode 100644
index 000000000000..5f7d2d22d550
--- /dev/null
+++ b/tests/unit/v1/autotp/test_autotp_multiple_models.py
@@ -0,0 +1,105 @@
+# Copyright (c) DeepSpeed Team.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+"""Regression tests for multiple AutoTP models living in the same process.
+
+The kv-head / grain sharding state used to be process-wide globals in tp_shard, so a second
+AutoTP model overwrote the first one's metadata (#8231). AutoTPMeta is per-model now; these
+tests hold that line.
+"""
+
+import torch
+from types import SimpleNamespace
+
+from unit.common import DistributedTest
+
+import deepspeed.comm as dist
+from deepspeed.utils import groups
+from deepspeed.module_inject.auto_tp import AutoTP
+from deepspeed.module_inject.autotp_config import AutoTPConfig
+from deepspeed.module_inject.tp_shard import get_shard_size_list
+
+# Only q_proj is partitioned; matches the naming the kv-head split logic looks for.
+PARTITION_CONFIG = {
+    "use_default_specs": False,
+    "layer_specs": [{
+        "patterns": [".*q_proj\\.weight$"],
+        "partition_type": "column",
+    }],
+}
+
+
+def apply_autotp(model, tp_size, partition_config):
+    groups._init_tp_mesh_device(tensor_model_parallel_size=tp_size)
+    autotp = AutoTP(module=model,
+                    all_reduce_linears=[],
+                    prefix="",
+                    state_dict=None,
+                    linear_layer_setting=None,
+                    orig_layer_impl=None,
+                    keep_module_on_host=False,
+                    partition_config=AutoTPConfig.from_dict(partition_config),
+                    model_config=getattr(model, "config", None))
+    autotp.set_tensor_parallel_config(tp_size, groups.get_tensor_model_parallel_group())
+    autotp.update_linear_policies()
+    autotp._replace_module(model)
+    return model
+
+
+class AttentionOnlyModel(torch.nn.Module):
+    """Minimal stand-in for a decoder layer, named so AutoTP treats it as attention."""
+
+    class _Attention(torch.nn.Module):
+
+        def __init__(self, hidden_dim, proj_dim):
+            super().__init__()
+            self.q_proj = torch.nn.Linear(hidden_dim, proj_dim, bias=False)
+
+    def __init__(self, hidden_dim, proj_dim, num_kv_heads):
+        super().__init__()
+        self.self_attn = AttentionOnlyModel._Attention(hidden_dim, proj_dim)
+        # AutoTPMeta probes attributes by name, so any config-like object will do.
+        self.config = SimpleNamespace(num_key_value_heads=num_kv_heads,
+                                      num_attention_heads=num_kv_heads,
+                                      hidden_size=hidden_dim)
+
+
+class TestAutoTPMultipleModels(DistributedTest):
+    world_size = 2
+    reuse_dist_env = False
+
+    def test_a_second_model_does_not_reshard_the_first(self):
+        # 3 kv heads over 2 ranks is uneven ([128, 64] of 192), while 2 kv heads divides evenly
+        # ([96, 96]). The second model's split is what a clobbered kv-head count would give the
+        # first one, so the two are distinguishable.
+        teacher = AttentionOnlyModel(hidden_dim=64, proj_dim=192, num_kv_heads=3)
+        teacher = apply_autotp(teacher, tp_size=2, partition_config=PARTITION_CONFIG)
+        teacher_layer = teacher.self_attn.q_proj
+        teacher_split = list(teacher_layer._partition_sizes)
+        assert teacher_split == [128, 64]
+
+        # A second model with a different kv-head count is built in the same process.
+        student = AttentionOnlyModel(hidden_dim=64, proj_dim=192, num_kv_heads=2)
+        student = apply_autotp(student, tp_size=2, partition_config=PARTITION_CONFIG)
+        assert list(student.self_attn.q_proj._partition_sizes) == [96, 96]
+
+        # The teacher still describes, and re-derives, its own split.
+        assert teacher_layer.tp_meta.num_kv_heads == 3
+        assert list(teacher_layer._partition_sizes) == teacher_split
+        assert get_shard_size_list(192, 2, teacher_layer.tp_meta, teacher_layer.name) == teacher_split
+        assert teacher_layer.weight.shape[0] == teacher_split[dist.get_rank()]
+
+    def test_multimodal_outer_config_keeps_kv_head_split(self):
+        # Multimodal configs keep the head counts only under text_config; passing the outer
+        # config to AutoTP would lose num_kv_heads and fall back to an even-grain split that
+        # cuts through KV heads.
+        model = AttentionOnlyModel(hidden_dim=64, proj_dim=192, num_kv_heads=3)
+        text_config = model.config
+        model.config = SimpleNamespace(text_config=text_config)
+
+        model = apply_autotp(model, tp_size=2, partition_config=PARTITION_CONFIG)
+        layer = model.self_attn.q_proj
+        assert layer.tp_meta.num_kv_heads == 3
+        assert list(layer._partition_sizes) == [128, 64]
+        assert layer.weight.shape[0] == [128, 64][dist.get_rank()]
diff --git a/tests/unit/model_parallelism/test_autotp_training.py b/tests/unit/v1/autotp/test_autotp_training.py
similarity index 99%
rename from tests/unit/model_parallelism/test_autotp_training.py
rename to tests/unit/v1/autotp/test_autotp_training.py
index f694ef18efcf..6cf2e35e7bed 100644
--- a/tests/unit/model_parallelism/test_autotp_training.py
+++ b/tests/unit/v1/autotp/test_autotp_training.py
@@ -17,7 +17,6 @@
 from contextlib import contextmanager
 from torch import nn
 from deepspeed.module_inject.layers import LinearAllreduce, LinearLayer, set_autotp_mode, is_autotp_training_mode
-from deepspeed.module_inject.tp_shard import get_shard_size_list
 from unit.checkpoint.common import compare_lr_scheduler_states, compare_optimizer_states
 import os
 from deepspeed.runtime.utils import is_model_parallel_parameter
@@ -502,7 +501,7 @@ def run_tp_layer_fwd_bwd(tp_size,
         loss.backward()
 
         expected_out = torch_out
-        output_partition_sizes = get_shard_size_list(torch_out.shape[-1], tp_size, linear.name)
+        output_partition_sizes = list(linear._partition_sizes)
         tp_rank = groups.get_tensor_model_parallel_rank()
         if not gather_output:
             shard_offset = sum(output_partition_sizes[:tp_rank])
@@ -689,7 +688,7 @@ def test_uneven_linear_gather_params(self):
         total_params = sum(p.numel() for p in torch_linear.parameters())
         tp_layer = LinearLayer(deepcopy(torch_linear), groups.get_tensor_model_parallel_group())
         tp_rank = groups.get_tensor_model_parallel_rank()
-        output_partition_sizes = get_shard_size_list(output_dim, tp_size, tp_layer.name)
+        output_partition_sizes = list(tp_layer._partition_sizes)
         expected_tp_params = output_partition_sizes[tp_rank] * (hidden_dim + 1)
 
         assert expected_tp_params == sum(p.numel() for p in tp_layer.parameters())
@@ -778,7 +777,7 @@ def test_consolidated_checkpoint(self):
         engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
 
         tp_rank = groups.get_tensor_model_parallel_rank()
-        output_partition_sizes = get_shard_size_list(vocab_size, self.world_size, "lm_head")
+        output_partition_sizes = list(engine.module.lm_head._partition_sizes)
         assert isinstance(engine.module.lm_head, LinearLayer)
         assert engine.module.lm_head.gather_output
         assert engine.module.lm_head.weight.shape == (output_partition_sizes[tp_rank], hidden_dim)
diff --git a/tests/unit/model_parallelism/test_configurable_parallel_mp.py b/tests/unit/v1/autotp/test_configurable_parallel_mp.py
similarity index 100%
rename from tests/unit/model_parallelism/test_configurable_parallel_mp.py
rename to tests/unit/v1/autotp/test_configurable_parallel_mp.py
diff --git a/tests/unit/model_parallelism/test_configurable_parallel_pp.py b/tests/unit/v1/autotp/test_configurable_parallel_pp.py
similarity index 100%
rename from tests/unit/model_parallelism/test_configurable_parallel_pp.py
rename to tests/unit/v1/autotp/test_configurable_parallel_pp.py
diff --git a/tests/unit/model_parallelism/test_tp_plan_e2e.py b/tests/unit/v1/autotp/test_tp_plan_e2e.py
similarity index 100%
rename from tests/unit/model_parallelism/test_tp_plan_e2e.py
rename to tests/unit/v1/autotp/test_tp_plan_e2e.py
diff --git a/tests/unit/model_parallelism/test_tp_plan_real_models.py b/tests/unit/v1/autotp/test_tp_plan_real_models.py
similarity index 98%
rename from tests/unit/model_parallelism/test_tp_plan_real_models.py
rename to tests/unit/v1/autotp/test_tp_plan_real_models.py
index b21a57c5d9c7..aa20f6933e47 100644
--- a/tests/unit/model_parallelism/test_tp_plan_real_models.py
+++ b/tests/unit/v1/autotp/test_tp_plan_real_models.py
@@ -9,7 +9,6 @@
 import deepspeed
 from deepspeed.accelerator import get_accelerator
 from deepspeed.module_inject.layers import LinearLayer
-from deepspeed.module_inject.tp_shard import get_shard_size_list
 from deepspeed.runtime.tensor_parallel.config import _get_hf_tp_plan
 from deepspeed.utils import groups
 from unit.common import DistributedTest
@@ -129,7 +128,7 @@ def test_qwen2_tp_plan_with_uneven_vocab(self):
         engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=ds_config)
 
         tp_rank = groups.get_tensor_model_parallel_rank()
-        output_partition_sizes = get_shard_size_list(config.vocab_size, 2, "lm_head")
+        output_partition_sizes = list(model.lm_head._partition_sizes)
         assert engine.autotp_size() == 2
         assert isinstance(model.lm_head, LinearLayer)
         assert model.lm_head.gather_output
diff --git a/tests/unit/v1/moe/test_autoep_autotp_runtime.py b/tests/unit/v1/moe/test_autoep_autotp_runtime.py
index 6d0e79a95c8a..b09f1b9d5547 100644
--- a/tests/unit/v1/moe/test_autoep_autotp_runtime.py
+++ b/tests/unit/v1/moe/test_autoep_autotp_runtime.py
@@ -112,6 +112,7 @@ def __init__(self):
         linear_layer_setting=None,
         orig_layer_impl=None,
         partition_config=None,
+        model_config=getattr(model, "config", None),
     )
 
     calls = []