@@ -23,6 +23,7 @@
 import torch.distributed as dist
 from parameterized import parameterized
 from torch.distributed.device_mesh import init_device_mesh
+from torch.distributed.tensor import DTensor
 
 from transformers import (
     AutoModelForCausalLM,
@@ -124,14 +125,16 @@ def _init_cp_mesh(device_type):
 def _tiny_config(**kwargs):
     """A `LlamaConfig` small enough for CPU tests, in float32 with `sdpa` attention."""
     return LlamaConfig(
-        hidden_size=32,
-        intermediate_size=64,
-        num_hidden_layers=2,
-        num_attention_heads=4,
-        num_key_value_heads=2,
-        vocab_size=99,
-        attn_implementation="sdpa",
-        **kwargs,
+        **{
+            "hidden_size": 32,
+            "intermediate_size": 64,
+            "num_hidden_layers": 2,
+            "num_attention_heads": 4,
+            "num_key_value_heads": 2,
+            "vocab_size": 99,
+            "attn_implementation": "sdpa",
+            **kwargs,
+        }
     )
 
 
@@ -450,6 +453,26 @@ def test_refuses_indivisible_heads(self):
         with self.assertRaisesRegex(ValueError, "must divide both the number of attention heads"):
             self._check(config, cp_size=4)
 
+    def test_head_divisibility_covers_tensor_parallelism(self):
+        # tensor parallelism keeps `num_heads / tp_size` heads on this rank, and those are what cp splits
+        def config_with_heads(num_heads, num_kv_heads):
+            config = LlamaConfig(
+                hidden_size=32,
+                num_attention_heads=num_heads,
+                num_key_value_heads=num_kv_heads,
+                attn_implementation="sdpa",
+            )
+            config.distributed_config = DistributedConfig(tp_size=2, cp_size=2)
+            return config
+
+        self._check(config_with_heads(8, 4), cp_size=2)
+        # 1 key/value head left on each rank cannot be split across 2 cp ranks
+        with self.assertRaisesRegex(ValueError, r"`cp_size \* tp_size` must divide"):
+            self._check(config_with_heads(8, 2), cp_size=2)
+        # so few attention heads left on each rank that 2 cp ranks cannot split them either
+        with self.assertRaisesRegex(ValueError, r"`cp_size \* tp_size` must divide"):
+            self._check(config_with_heads(8, 4), cp_size=4)
+
     def test_accepts_supported_config(self):
         self._check(LlamaConfig(attn_implementation="sdpa", num_attention_heads=8, num_key_value_heads=4))
         # the mask the model builds over the whole sequence already has the sliding window in it
@@ -668,6 +691,44 @@ def _compare_with_unsharded(init_dir, batch, distributed_config, keep=None):
         torch.testing.assert_close(grad, expected, rtol=CP_RTOL, atol=CP_ATOL, msg=f"Gradient mismatch for {name}")
 
 
+def _test_matches_unsharded_with_tp(rank, tp_size, cp_size, fsdp_size, packed):
+    """The tp ranks hold the same batch and the same cp slice, and split the heads before cp splits them again."""
+    # 4 attention heads and 4 key/value heads, both divisible by `tp_size * cp_size` = 4
+    config = _tiny_config(num_key_value_heads=4)
+    batch_size, seq_len = 2, 8
+    with _deterministic_init_model_dir(rank, config, torch.float32) as init_dir:
+        input_ids, position_ids = _tiny_batch(config, batch_size * fsdp_size, seq_len, packed)
+        distributed_config = DistributedConfig(tp_size=tp_size, cp_size=cp_size, fsdp_size=fsdp_size)
+        batch = {"input_ids": input_ids, "labels": input_ids.clone(), "use_cache": False}
+        if position_ids is not None:
+            batch["position_ids"] = position_ids
+        _compare_with_unsharded(init_dir, batch, distributed_config)
+        dist.barrier()
+
+
+def _test_tp_and_cp_compose(rank, tp_size, cp_size):
+    """The tp plan shards the projections of an attention layer, the cp plan registers its attention call."""
+    config = _tiny_config(num_key_value_heads=4)
+    with _deterministic_init_model_dir(rank, config, torch.float32) as init_dir:
+        distributed_config = DistributedConfig(tp_size=tp_size, cp_size=cp_size)
+        model = AutoModelForCausalLM.from_pretrained(init_dir, distributed_config=distributed_config)
+        attention = model.model.layers[0].self_attn
+        # tensor parallelism sharded the projections of the very module whose forward context parallelism wrapped
+        assert isinstance(attention.q_proj.weight, DTensor), type(attention.q_proj.weight)
+        assert hasattr(attention.forward, "__wrapped__"), "the attention style was not installed"
+
+        # and one forward goes through both: the logits of this rank's slice match the unsharded run
+        input_ids, _ = _tiny_batch(config, 2, 8)
+        with torch.no_grad():
+            logits = model(input_ids=input_ids).logits
+        reference = AutoModelForCausalLM.from_pretrained(init_dir)
+        with torch.no_grad():
+            expected = reference(input_ids=input_ids).logits
+        cp_rank = model._mesh_manager.get_mesh("cp").get_local_rank()
+        torch.testing.assert_close(logits, expected.chunk(cp_size, dim=1)[cp_rank], rtol=CP_RTOL, atol=CP_ATOL)
+        dist.barrier()
+
+
 def _test_matches_unsharded(rank, cp_size, fsdp_size, dp_replicate_size, packed, seq_len=8):
     config = _tiny_config()
     if cp_size > 2:  # Ulysses splits the heads over the ranks, so four ranks need four key/value heads
@@ -757,15 +818,16 @@ def _test_custom_attention_style(rank, cp_size):
         dist.barrier()
 
 
-def _test_save_pretrained(rank, cp_size):
+def _test_save_pretrained(rank, cp_size, tp_size=1):
     """Weights saved from a context parallel model reload without a distributed config, unchanged."""
-    config = _tiny_config()
+    config = _tiny_config(num_key_value_heads=4 if tp_size > 1 else 2)
     with _deterministic_init_model_dir(rank, config, torch.float32) as init_dir:
         reference = AutoModelForCausalLM.from_pretrained(init_dir)
         expected = {name: value.clone() for name, value in reference.state_dict().items()}
         del reference
 
-        model = AutoModelForCausalLM.from_pretrained(init_dir, distributed_config=DistributedConfig(cp_size=cp_size))
+        distributed_config = DistributedConfig(cp_size=cp_size, tp_size=tp_size)
+        model = AutoModelForCausalLM.from_pretrained(init_dir, distributed_config=distributed_config)
         with _distributed_tmpdir(rank) as save_dir:
             model.save_pretrained(save_dir, is_main_process=(rank == 0))
             dist.barrier()
@@ -799,10 +861,9 @@ def _test_save_load_distributed_checkpoint(rank, cp_size):
 
 @unittest.skipUnless(is_torch_greater_or_equal("2.10"), "Context parallelism requires torch>=2.10")
 class TestContextParallelConfig(unittest.TestCase):
-    def test_refuses_other_parallelisms(self):
-        for other in ({"pp_size": 2}, {"tp_size": 2}, {"ep_size": 2}):
-            with self.subTest(**other), self.assertRaisesRegex(ValueError, "not supported yet with"):
-                DistributedConfig(cp_size=2, **other)
+    def test_refuses_pipeline_parallelism(self):
+        with self.assertRaisesRegex(ValueError, "not supported yet with `pp_size`"):
+            DistributedConfig(cp_size=2, pp_size=2)
 
     def test_refuses_a_model_without_a_cp_plan(self):
         config = _tiny_config()
@@ -905,6 +966,14 @@ def test_sliding_window_and_full_layers(self):
         )
         run_distributed(world_size=2, backend="gloo")(_test_padded)(2, config, False)
 
+    def test_matches_unsharded_with_tensor_parallelism(self):
+        for packed in (False, True):
+            with self.subTest(packed=packed):
+                run_distributed(world_size=4, backend="gloo")(_test_matches_unsharded_with_tp)(2, 2, 1, packed)
+
+    def test_tensor_parallelism_and_context_parallelism_compose(self):
+        run_distributed(world_size=4, backend="gloo")(_test_tp_and_cp_compose)(2, 2)
+
     def test_inputs_embeds(self):
         run_distributed(world_size=2, backend="gloo")(_test_inputs_embeds)(2)
 
@@ -914,5 +983,8 @@ def test_custom_attention_style(self):
     def test_save_pretrained(self):
         run_distributed(world_size=2, backend="gloo")(_test_save_pretrained)(2)
 
+    def test_save_pretrained_with_tensor_parallelism(self):
+        run_distributed(world_size=4, backend="gloo")(_test_save_pretrained)(2, 2)
+
     def test_save_load_distributed_checkpoint(self):
         run_distributed(world_size=2, backend="gloo")(_test_save_load_distributed_checkpoint)(2)