feat: support TP>1 Domino rollout for DFlash V2 (#37069)

Co-authored-by: Qiaolin-Yu <liin1211@outlook.com>

python/sglang/srt/speculative/dflash_worker_v2.py

@@ -285,7 +285,7 @@ class _SelectorDraftSampler:
 
 
 class _DominoDraftSampler:
-    """Capture-safe TP=1 Domino rollout over a fixed-size draft block."""
+    """Capture-safe Domino rollout over a fixed-size draft block."""
 
     def __init__(
         self,
@@ -299,6 +299,10 @@ class _DominoDraftSampler:
         shift_label,
         max_bs,
         candidate_pool_size,
+        tp_group=None,
+        lm_head_org_vocab_start=0,
+        lm_head_num_org=None,
+        lm_head_num_org_padded=None,
     ):
         self.target_embedding = target_embedding
         self.lm_head_weight = lm_head_weight
@@ -308,6 +312,10 @@ class _DominoDraftSampler:
         self.block_size = int(block_size)
         self.shift_label = bool(shift_label)
         self.candidate_pool_size = int(candidate_pool_size)
+        self.tp_group = tp_group
+        self.lm_head_org_vocab_start = int(lm_head_org_vocab_start)
+        self.lm_head_num_org = lm_head_num_org
+        self.lm_head_num_org_padded = lm_head_num_org_padded
         max_tokens = int(max_bs) * (self.block_size - 1)
         self.out = torch.empty(
             (max_tokens,), dtype=torch.int64, device=lm_head_weight.device
@@ -329,6 +337,11 @@ class _DominoDraftSampler:
             vocab_size=self.vocab_size,
             shift_label=self.shift_label,
             candidate_pool_size=self.candidate_pool_size,
+            tp_group=self.tp_group,
+            lm_head_org_vocab_start=self.lm_head_org_vocab_start,
+            lm_head_num_org=self.lm_head_num_org,
+            lm_head_num_org_padded=self.lm_head_num_org_padded,
+            prefer_tp_candidate_pool=bs > 1,
         )
         self.out[: bs * (self.block_size - 1)].copy_(proposals.reshape(-1))
 
@@ -408,6 +421,7 @@ class DFlashWorkerV2(BaseSpecWorker):
             validate_domino_runtime(
                 device=torch.device(self.device),
                 tp_size=int(get_tp_group().world_size),
+                tp_rank=int(self.ps.tp_rank),
                 target_vocab_size=int(self.model_runner.model_config.vocab_size),
                 draft_vocab_size=int(self.draft_model_runner.model_config.vocab_size),
                 hidden_size=int(self.draft_model.config.hidden_size),
@@ -463,7 +477,8 @@ class DFlashWorkerV2(BaseSpecWorker):
             )
             if self._is_domino:
                 logger.info(
-                    "DFLASH Domino rollout enabled (BF16, TP=1, block-shared candidate pool size=%s).",
+                    "DFLASH Domino rollout enabled (BF16, TP=%s, block-shared candidate pool size=%s).",
+                    int(get_tp_group().world_size),
                     self.domino_candidate_pool_size,
                 )
             logger.info(
@@ -755,16 +770,16 @@ class DFlashWorkerV2(BaseSpecWorker):
             return _eager("quantized lm_head")
         tp_group = get_tp_group()
         if self._is_domino:
-            if tp_group.world_size != 1:
-                return _eager("Domino cuda graph currently requires tp=1")
             prefix_gru = self.draft_model.prefix_gru
             embed_proj = self.draft_model.embed_proj
             if prefix_gru is None or embed_proj is None:
                 return _eager("Domino projector modules are unavailable")
             if self.ps.tp_rank == 0:
                 logger.info(
-                    "DFLASH Domino rollout folded into the draft cuda graph (tp=1)."
+                    "DFLASH Domino rollout folded into the draft cuda graph (tp=%s).",
+                    int(tp_group.world_size),
                 )
+            shard = getattr(lm_head, "shard_indices", None)
             return _DominoDraftSampler(
                 target_embedding=target_model.get_input_embeddings(),
                 lm_head_weight=lm_head.weight,
@@ -775,6 +790,16 @@ class DFlashWorkerV2(BaseSpecWorker):
                 shift_label=self.draft_model.shift_label,
                 max_bs=max(get_exec().graph.cuda_graph_config.decode.bs),
                 candidate_pool_size=self.domino_candidate_pool_size,
+                tp_group=tp_group,
+                lm_head_org_vocab_start=(
+                    int(shard.org_vocab_start_index) if shard is not None else 0
+                ),
+                lm_head_num_org=(
+                    int(shard.num_org_elements) if shard is not None else None
+                ),
+                lm_head_num_org_padded=(
+                    int(shard.num_org_elements_padded) if shard is not None else None
+                ),
             )
         if not hasattr(lm_head, "shard_indices"):
             if tp_group.world_size != 1:
@@ -2287,6 +2312,8 @@ class DFlashWorkerV2(BaseSpecWorker):
             embed_proj = self.draft_model.embed_proj
             if prefix_gru is None or embed_proj is None:
                 raise RuntimeError("DFLASH Domino projector modules are unavailable.")
+            tp_group = get_tp_group()
+            shard = getattr(lm_head, "shard_indices", None)
             draft_next = domino_greedy_rollout(
                 draft_hidden=draft_hidden,
                 bonus_tokens=block_ids[:, 0],
@@ -2297,6 +2324,16 @@ class DFlashWorkerV2(BaseSpecWorker):
                 vocab_size=int(self.model_runner.model_config.vocab_size),
                 shift_label=bool(self.draft_model.shift_label),
                 candidate_pool_size=self.domino_candidate_pool_size,
+                tp_group=tp_group,
+                lm_head_org_vocab_start=(
+                    int(shard.org_vocab_start_index) if shard is not None else 0
+                ),
+                lm_head_num_org=(
+                    int(shard.num_org_elements) if shard is not None else None
+                ),
+                lm_head_num_org_padded=(
+                    int(shard.num_org_elements_padded) if shard is not None else None
+                ),
             )
         elif self._draft_sampler is not None and draft_out.can_run_graph:
             draft_next = self._draft_sampler.out[

python/sglang/srt/speculative/domino_utils.py

@@ -19,10 +19,113 @@ def _domino_gru_cell(
     )
 
 
+# Throughput policy for the logical full-vocabulary tensor. TP>1 uses the
+# compact candidate pool when this gathered representation would be large.
+_DOMINO_TP_FULL_BASE_LOGITS_MAX_BYTES = 96 * 1024 * 1024
+
+
+def _domino_tp_first_ids(
+    local_logits: torch.Tensor,
+    *,
+    org_vocab_start: int,
+    num_org: int,
+    tp_group,
+) -> torch.Tensor:
+    """Select the full-vocab first argmax from contiguous vocab shards."""
+    local_max, local_arg = torch.max(local_logits[:, :num_org], dim=-1)
+    local_ids = local_arg.to(torch.int64) + int(org_vocab_start)
+    tp_size = int(tp_group.world_size)
+    batch_size = int(local_logits.shape[0])
+    gathered_max = torch.empty(
+        (tp_size * batch_size,), dtype=local_max.dtype, device=local_max.device
+    )
+    gathered_ids = torch.empty(
+        (tp_size * batch_size,), dtype=torch.int64, device=local_max.device
+    )
+    tp_group.all_gather_into_tensor(gathered_max, local_max.contiguous())
+    tp_group.all_gather_into_tensor(gathered_ids, local_ids.contiguous())
+    gathered_max = gathered_max.view(tp_size, batch_size)
+    gathered_ids = gathered_ids.view(tp_size, batch_size)
+    best_rank = torch.argmax(gathered_max, dim=0, keepdim=True)
+    return torch.gather(gathered_ids, 0, best_rank).squeeze(0)
+
+
+def _domino_tp_candidate_state(
+    local_feedback_logits: torch.Tensor,
+    *,
+    candidate_pool_size: int,
+    org_vocab_start: int,
+    num_org: int,
+    tp_group,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Build one strict global candidate pool from vocab-sharded base logits."""
+    num_steps, batch_size, local_vocab_size = local_feedback_logits.shape
+    tp_size = int(tp_group.world_size)
+    local_k = min(int(candidate_pool_size), int(local_vocab_size))
+    if int(candidate_pool_size) > tp_size * local_k:
+        raise ValueError("Domino TP shards do not cover the requested candidate pool.")
+
+    pool_scores = local_feedback_logits.amax(dim=0)
+    if num_org < local_vocab_size:
+        pool_scores[:, num_org:].fill_(float("-inf"))
+    local_scores, local_positions = torch.topk(
+        pool_scores, k=local_k, dim=-1, sorted=False
+    )
+    local_ids = local_positions.to(torch.int64) + int(org_vocab_start)
+
+    gathered_scores = torch.empty(
+        (tp_size * batch_size, local_k),
+        dtype=local_scores.dtype,
+        device=local_scores.device,
+    )
+    gathered_ids = torch.empty(
+        (tp_size * batch_size, local_k),
+        dtype=torch.int64,
+        device=local_scores.device,
+    )
+    tp_group.all_gather_into_tensor(gathered_scores, local_scores.contiguous())
+    tp_group.all_gather_into_tensor(gathered_ids, local_ids.contiguous())
+    gathered_scores = (
+        gathered_scores.view(tp_size, batch_size, local_k)
+        .permute(1, 0, 2)
+        .reshape(batch_size, tp_size * local_k)
+    )
+    gathered_ids = (
+        gathered_ids.view(tp_size, batch_size, local_k)
+        .permute(1, 0, 2)
+        .reshape(batch_size, tp_size * local_k)
+    )
+    global_positions = torch.topk(
+        gathered_scores,
+        k=int(candidate_pool_size),
+        dim=-1,
+        sorted=False,
+    ).indices
+    candidate_ids = torch.gather(gathered_ids, 1, global_positions).contiguous()
+
+    owned = (candidate_ids >= int(org_vocab_start)) & (
+        candidate_ids < int(org_vocab_start + num_org)
+    )
+    local_positions = (candidate_ids - int(org_vocab_start)).clamp(
+        0, max(num_org - 1, 0)
+    )
+    candidate_base = torch.gather(
+        local_feedback_logits.transpose(0, 1),
+        2,
+        local_positions[:, None, :].expand(-1, num_steps, -1),
+    )
+    # Each global candidate is owned by exactly one vocab shard, so SUM
+    # reconstructs its base logit without gathering the full vocabulary.
+    candidate_base.masked_fill_(~owned[:, None, :], 0)
+    candidate_base = tp_group.all_reduce(candidate_base.contiguous())
+    return candidate_ids, candidate_base.transpose(0, 1)
+
+
 def validate_domino_runtime(
     *,
     device: torch.device,
     tp_size: int,
+    tp_rank: int,
     target_vocab_size: int,
     draft_vocab_size: int,
     hidden_size: int,
@@ -34,8 +137,14 @@ def validate_domino_runtime(
     """Validate the deliberately narrow correctness-first Domino runtime."""
     if device.type != "cuda":
         raise ValueError(f"DFLASH Domino currently requires CUDA, got {device}.")
-    if int(tp_size) != 1:
+    tp_size = int(tp_size)
-        raise ValueError(f"DFLASH Domino currently requires TP=1, got TP={tp_size}.")
+    if tp_size < 1:
+        raise ValueError(f"DFLASH Domino requires TP>=1, got TP={tp_size}.")
+    tp_rank = int(tp_rank)
+    if not 0 <= tp_rank < tp_size:
+        raise ValueError(
+            f"DFLASH Domino requires 0<=TP rank<TP size, got rank={tp_rank}, size={tp_size}."
+        )
     if int(target_vocab_size) != int(draft_vocab_size):
         raise ValueError(
             "DFLASH Domino requires identical target and draft vocab sizes, "
@@ -49,25 +158,108 @@ def validate_domino_runtime(
             "DFLASH Domino requires target embedding and lm_head weight tensors."
         )
 
-    shard = getattr(lm_head, "shard_indices", None)
+    lm_head_shard = getattr(lm_head, "shard_indices", None)
-    if shard is not None:
+    if lm_head_shard is not None:
-        if int(shard.num_added_elements) != 0:
+        if (
+            int(getattr(lm_head, "num_added_embeddings", 0)) != 0
+            or int(lm_head_shard.num_added_elements) != 0
+        ):
             raise ValueError(
                 "DFLASH Domino does not support added-vocab lm_head shards."
             )
-        if int(shard.org_vocab_start_index) != 0 or int(shard.num_org_elements) != int(
+        if int(getattr(lm_head, "org_vocab_size", target_vocab_size)) != int(
             target_vocab_size
         ):
             raise ValueError(
-                "DFLASH Domino requires the complete target vocabulary on TP=1."
+                "DFLASH Domino lm_head original vocab size does not match the target."
             )
-    elif int(lm_head_weight.shape[0]) != int(target_vocab_size):
+        if int(getattr(lm_head, "tp_size", tp_size)) != tp_size:
-        raise ValueError(
+            raise ValueError(
-            "DFLASH Domino lm_head row count must equal the target vocab size, "
+                "DFLASH Domino lm_head TP size does not match the runtime TP size."
-            f"got rows={int(lm_head_weight.shape[0])}, vocab={target_vocab_size}."
+            )
+        required_shard_fields = (
+            "org_vocab_start_index",
+            "org_vocab_end_index",
+            "num_org_elements",
+            "num_org_elements_padded",
         )
+        missing_shard_fields = [
+            name for name in required_shard_fields if not hasattr(lm_head_shard, name)
+        ]
+        if missing_shard_fields:
+            raise ValueError(
+                "DFLASH Domino lm_head shard metadata is missing: "
+                + ", ".join(missing_shard_fields)
+            )
+        org_vocab_start = int(lm_head_shard.org_vocab_start_index)
+        org_vocab_end = int(lm_head_shard.org_vocab_end_index)
+        num_org = int(lm_head_shard.num_org_elements)
+        num_org_padded = int(lm_head_shard.num_org_elements_padded)
+        if (
+            num_org <= 0
+            or org_vocab_start < 0
+            or org_vocab_end != org_vocab_start + num_org
+            or org_vocab_end > int(target_vocab_size)
+            or num_org_padded < num_org
+        ):
+            raise ValueError("DFLASH Domino lm_head original-vocab shard is invalid.")
+        if int(lm_head_weight.shape[0]) < num_org_padded:
+            raise ValueError(
+                "DFLASH Domino lm_head weight is smaller than its padded vocab shard."
+            )
+        expected_start = tp_rank * num_org_padded
+        expected_end = min(expected_start + num_org_padded, int(target_vocab_size))
+        if (
+            org_vocab_start != expected_start
+            or org_vocab_end != expected_end
+            or num_org != expected_end - expected_start
+        ):
+            raise ValueError(
+                "DFLASH Domino lm_head vocab shard does not match its TP rank."
+            )
+    else:
+        if tp_size != 1:
+            raise ValueError("DFLASH Domino requires lm_head shard metadata for TP>1.")
+        if int(lm_head_weight.shape[0]) != int(target_vocab_size):
+            raise ValueError(
+                "DFLASH Domino lm_head row count must equal the target vocab size, "
+                f"got rows={int(lm_head_weight.shape[0])}, vocab={target_vocab_size}."
+            )
 
-    if int(embedding_weight.shape[0]) < int(target_vocab_size):
+    embedding_shard = getattr(target_embedding, "shard_indices", None)
+    if embedding_shard is not None:
+        if (
+            int(getattr(target_embedding, "num_added_embeddings", 0)) != 0
+            or int(embedding_shard.num_added_elements) != 0
+        ):
+            raise ValueError(
+                "DFLASH Domino does not support added-vocab embedding shards."
+            )
+        if int(getattr(target_embedding, "org_vocab_size", target_vocab_size)) != int(
+            target_vocab_size
+        ):
+            raise ValueError(
+                "DFLASH Domino embedding original vocab size does not match the target."
+            )
+        embedding_tp_size = int(getattr(target_embedding, "tp_size", tp_size))
+        if embedding_tp_size not in (1, tp_size):
+            raise ValueError(
+                "DFLASH Domino embedding TP size does not match the runtime TP size."
+            )
+        required_embedding_rows = (
+            int(target_vocab_size)
+            if embedding_tp_size == 1
+            else int(embedding_shard.num_org_elements_padded)
+        )
+        if int(embedding_weight.shape[0]) < required_embedding_rows:
+            raise ValueError(
+                "DFLASH Domino embedding weight is smaller than its padded vocab shard."
+            )
+    elif int(embedding_weight.shape[0]) < int(target_vocab_size):
+        if tp_size != 1:
+            raise ValueError(
+                "DFLASH Domino requires embedding shard metadata for TP>1."
+            )
         raise ValueError(
             "DFLASH Domino target embedding has fewer rows than the target vocab "
             f"size: rows={int(embedding_weight.shape[0])}, vocab={target_vocab_size}."
@@ -119,6 +311,11 @@ def domino_greedy_rollout(
     vocab_size: int,
     shift_label: bool,
     candidate_pool_size: int,
+    tp_group=None,
+    lm_head_org_vocab_start: int = 0,
+    lm_head_num_org: int | None = None,
+    lm_head_num_org_padded: int | None = None,
+    prefer_tp_candidate_pool: bool | None = None,
 ) -> torch.Tensor:
     """Generate a Domino chain using one block-shared base-logit candidate pool."""
     if draft_hidden.ndim != 3:
@@ -148,24 +345,97 @@ def domino_greedy_rollout(
             "Domino draft hidden states do not contain enough proposal positions."
         )
 
-    weight = lm_head_weight[: int(vocab_size)]
+    tp_size = int(tp_group.world_size) if tp_group is not None else 1
+    if tp_size > 1 and (lm_head_num_org is None or lm_head_num_org_padded is None):
+        raise ValueError(
+            "Domino TP rollout requires local lm_head vocab shard metadata."
+        )
+    local_vocab_size = int(lm_head_num_org_padded) if tp_size > 1 else int(vocab_size)
+    if tp_size > 1:
+        num_org = int(lm_head_num_org)
+        org_vocab_start = int(lm_head_org_vocab_start)
+        if (
+            num_org <= 0
+            or num_org > local_vocab_size
+            or org_vocab_start < 0
+            or org_vocab_start + num_org > int(vocab_size)
+        ):
+            raise ValueError(
+                "Domino TP rollout received an invalid lm_head vocab shard."
+            )
+    if int(lm_head_weight.shape[0]) < local_vocab_size:
+        raise ValueError(
+            "Domino lm_head weight is smaller than its padded vocab shard."
+        )
+    weight = lm_head_weight[:local_vocab_size]
     z_for_logits = z.to(weight.dtype) if z.dtype != weight.dtype else z
     logits_input = (
         z_for_logits.transpose(0, 1)
         .contiguous()
         .view(num_proposals * batch_size, hidden_size)
     )
-    base_logits = F.linear(logits_input, weight).view(num_proposals, batch_size, -1)
+    local_logits = F.linear(logits_input, weight).view(
+        num_proposals, batch_size, local_vocab_size
+    )
+    if prefer_tp_candidate_pool is None:
+        full_base_logits_bytes = (
+            num_proposals * batch_size * int(vocab_size) * local_logits.element_size()
+        )
+        prefer_tp_candidate_pool = (
+            full_base_logits_bytes > _DOMINO_TP_FULL_BASE_LOGITS_MAX_BYTES
+        )
+    use_tp_candidate_pool = (
+        tp_size > 1
+        and prefer_tp_candidate_pool
+        and 0 < candidate_pool_size < int(vocab_size)
+    )
+    first_ids = None
+    candidate_ids = None
+    candidate_base = None
+    if tp_size == 1:
+        base_logits = local_logits[:, :, : int(vocab_size)]
+    elif use_tp_candidate_pool:
+        first_ids = _domino_tp_first_ids(
+            local_logits[0],
+            org_vocab_start=int(lm_head_org_vocab_start),
+            num_org=int(lm_head_num_org),
+            tp_group=tp_group,
+        )
+        if num_proposals > 1:
+            candidate_ids, candidate_base = _domino_tp_candidate_state(
+                local_logits[1:],
+                candidate_pool_size=candidate_pool_size,
+                org_vocab_start=int(lm_head_org_vocab_start),
+                num_org=int(lm_head_num_org),
+                tp_group=tp_group,
+            )
+        base_logits = None
+    else:
+        local_logits_t = local_logits.view(
+            num_proposals * batch_size, local_vocab_size
+        ).T.contiguous()
+        gathered_logits = torch.empty(
+            (tp_size * local_vocab_size, num_proposals * batch_size),
+            dtype=local_logits.dtype,
+            device=local_logits.device,
+        )
+        tp_group.all_gather_into_tensor(gathered_logits, local_logits_t)
+        base_logits = (
+            gathered_logits.T[:, : int(vocab_size)]
+            .contiguous()
+            .view(num_proposals, batch_size, int(vocab_size))
+        )
 
-    first_ids = torch.argmax(base_logits[0], dim=-1).to(torch.long)
+    if first_ids is None:
+        first_ids = torch.argmax(base_logits[0], dim=-1).to(torch.long)
     proposals = [first_ids]
     if num_proposals == 1:
         return first_ids[:, None]
 
-    candidate_ids = None
-    candidate_base = None
     candidate_weight = None
-    if 0 < candidate_pool_size < int(vocab_size):
+    if candidate_ids is not None:
+        candidate_weight = F.embedding(candidate_ids, embed_proj[2].weight)
+    elif 0 < candidate_pool_size < int(vocab_size):
         feedback_logits = base_logits[1:]
         candidate_ids = torch.topk(
             feedback_logits.amax(dim=0),

test/registered/e2e/speculative/test_dflash_domino.py

@@ -14,7 +14,7 @@ from sglang.test.test_utils import (
     popen_launch_server,
 )
 
-register_cuda_ci(est_time=600, stage="base-b", runner_config="1-gpu-small")
+register_cuda_ci(est_time=600, stage="base-b", runner_config="2-gpu-large")
 
 
 class TestDFlashDominoFullVocab(CustomTestCase):
@@ -38,7 +38,7 @@ class TestDFlashDominoFullVocab(CustomTestCase):
                 "--dtype",
                 "bfloat16",
                 "--tp-size",
-                "1",
+                "2",
                 "--attention-backend",
                 "triton",
                 "--speculative-algorithm",
@@ -64,6 +64,7 @@ class TestDFlashDominoFullVocab(CustomTestCase):
         response = requests.get(self.base_url + "/server_info", timeout=10)
         response.raise_for_status()
         state = response.json()["internal_states"][0]
+        self.assertEqual(state["tp_size"], 2)
         self.assertEqual(state["speculative_num_draft_tokens"], 16)
         self.assertFalse(state["disable_overlap_schedule"])
         self.assertEqual(
@@ -71,11 +72,11 @@ class TestDFlashDominoFullVocab(CustomTestCase):
         )
         log = Path(self.server_log.name).read_text()
         self.assertIn(
-            "DFLASH Domino rollout enabled (BF16, TP=1, "
+            "DFLASH Domino rollout enabled (BF16, TP=2, "
             f"block-shared candidate pool size={self.candidate_pool_size}).",
             log,
         )
-        self.assertIn("Domino rollout folded into the draft cuda graph", log)
+        self.assertIn("Domino rollout folded into the draft cuda graph (tp=2)", log)
         self.assertIn(
             "Capture draft verify CUDA graph begin. backend=full, num_tokens_per_req=16,",
             log,

test/registered/unit/spec/test_dflash_domino.py

@@ -6,6 +6,7 @@ from torch import nn
 
 from sglang.srt.models.dflash import DFlashDraftModel
 from sglang.srt.speculative.dflash_utils import parse_dflash_draft_config
+from sglang.srt.speculative.domino_utils import validate_domino_runtime
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
 
@@ -126,5 +127,93 @@ class TestDFlashDominoWeights(CustomTestCase):
             model.load_weights([("prefix_gru.weight_ih_l0", torch.empty(12, 8))])
 
 
+class TestDFlashDominoRuntimeValidation(CustomTestCase):
+    def _modules(self, dtype=torch.bfloat16):
+        embedding = nn.Embedding(31, 8, dtype=dtype)
+        lm_head = nn.Linear(8, 31, bias=False, dtype=dtype)
+        prefix_gru = nn.GRU(8, 4, batch_first=True, bias=False, dtype=dtype)
+        embed_proj = nn.Sequential(
+            nn.Linear(12, 5, bias=False, dtype=dtype),
+            nn.SiLU(),
+            nn.Linear(5, 31, bias=False, dtype=dtype),
+        )
+        return embedding, lm_head, prefix_gru, embed_proj
+
+    def _tp2_modules(self):
+        embedding, lm_head, prefix_gru, embed_proj = self._modules()
+        embedding = nn.Embedding(16, 8, dtype=torch.bfloat16)
+        lm_head = nn.Linear(8, 16, bias=False, dtype=torch.bfloat16)
+        shard = SimpleNamespace(
+            num_added_elements=0,
+            org_vocab_start_index=0,
+            org_vocab_end_index=16,
+            num_org_elements=16,
+            num_org_elements_padded=16,
+        )
+        for module in (embedding, lm_head):
+            module.shard_indices = shard
+            module.org_vocab_size = 31
+            module.tp_size = 2
+            module.num_added_embeddings = 0
+        return embedding, lm_head, prefix_gru, embed_proj
+
+    def _validate(self, **overrides):
+        embedding, lm_head, prefix_gru, embed_proj = overrides.pop(
+            "modules", self._modules()
+        )
+        args = {
+            "device": torch.device("cuda"),
+            "tp_size": 1,
+            "tp_rank": 0,
+            "target_vocab_size": 31,
+            "draft_vocab_size": 31,
+            "hidden_size": 8,
+            "target_embedding": embedding,
+            "lm_head": lm_head,
+            "prefix_gru": prefix_gru,
+            "embed_proj": embed_proj,
+        }
+        args.update(overrides)
+        validate_domino_runtime(**args)
+
+    def test_tp_requires_vocab_shard_metadata(self):
+        with self.assertRaisesRegex(ValueError, "lm_head shard metadata"):
+            self._validate(tp_size=2)
+
+    def test_tp2_vocab_shards_supported(self):
+        self._validate(tp_size=2, modules=self._tp2_modules())
+
+    def test_tp2_incomplete_lm_head_shard_fails(self):
+        modules = self._tp2_modules()
+        modules[1].shard_indices = SimpleNamespace(
+            num_added_elements=0,
+            num_org_elements_padded=16,
+        )
+        with self.assertRaisesRegex(ValueError, "shard metadata is missing"):
+            self._validate(tp_size=2, modules=modules)
+
+    def test_tp_vocab_shard_must_match_rank(self):
+        modules = self._tp2_modules()
+        modules[1].shard_indices.org_vocab_start_index = 1
+        modules[1].shard_indices.org_vocab_end_index = 17
+        with self.assertRaisesRegex(ValueError, "does not match its TP rank"):
+            self._validate(tp_size=2, modules=modules)
+
+    def test_tp1_requires_complete_vocab_shard(self):
+        modules = self._modules()
+        modules[1].shard_indices = SimpleNamespace(
+            num_added_elements=0,
+            org_vocab_start_index=0,
+            org_vocab_end_index=30,
+            num_org_elements=30,
+            num_org_elements_padded=31,
+        )
+        modules[1].org_vocab_size = 31
+        modules[1].tp_size = 1
+        modules[1].num_added_embeddings = 0
+        with self.assertRaisesRegex(ValueError, "does not match its TP rank"):
+            self._validate(modules=modules)
+
+
 if __name__ == "__main__":
     unittest.main()