Diff Coverage

})


def _direct_params(module, placement):
    return {
        name: {"tp": placement}
        for name, _ in module.named_parameters(recurse=False)
    }

    }


def _linear(module, *, param, in_src, in_dst, out_src, out_dst):
    return ModuleShardingSpec(
        params=_direct_params(module, param),
        in_src={"input": {"tp": in_src}},
        in_dst={"input": {"tp": in_dst}},
        out_src={"output": {"tp": out_src}},

def _is_dsa_attention(module) -> bool:
    """Match the DSA/MLA attention contract without using a model name."""
    if getattr(module, "attention_type", None) in _ATTENTION_TYPES:
        return True
    if any(
        name.startswith("param_sink_")
        for name, _ in module.named_parameters(recurse=False)
    ):
        return True
    child_names = {name for name, _ in module.named_children()}
    return bool(_DISTINCTIVE_LEAVES.intersection(child_names))


def _matching_attention_roots(named_modules) -> frozenset[str]:
    return frozenset(

    is intentionally not a communication boundary in this template, so declare
    a parameter-only replicated spec.  Returns None when the module has no sink
    parameters.
    """
    sink_params = {
        name: {"tp": Replicate()}
        for name, _ in module.named_parameters(recurse=False)
        if name.startswith("param_sink_")
    }
    if not sink_params:
        return None
    return ModuleShardingSpec(params=sink_params, is_boundary=False)


def _lm_head_spec(module) -> ModuleShardingSpec:
    """Vocab-parallel LM-head projection consumed after the SP sequence gather.

    The integration gathers the language-model SP output before its
    multi-token prediction vocabulary heads, so the vocab-parallel projection
    consumes a replicated sequence (not Shard(1)).
    """
    return ModuleShardingSpec(
        params=_direct_params(module, Shard(0)),
        in_src={"input": {"tp": Replicate()}},
        in_dst={"input": {"tp": Replicate()}},
        out_src={"output": {"tp": Shard(-1)}},

    entering the language-model SP boundary.  Keep the vocab-parallel embedding
    output replicated here; the parent language-model boundary performs the
    sequence scatter after multimodal fusion.
    """
    return ModuleShardingSpec(
        params=_direct_params(module, Shard(0)),
        in_src={"hidden_states": {"tp": Replicate()}},
        in_dst={"hidden_states": {"tp": Replicate()}},
        out_src={"output": {"tp": Partial()}},

def _head_sharded_projection_spec(fqn, module, leaf) -> ModuleShardingSpec:
    """Head-sharded DSA projection leaf (linear_qb / index_linear_qb / merge)."""
    spec = _linear(
        module, param=Shard(0), in_src=Shard(1), in_dst=Replicate(),
        out_src=Shard(-1), out_dst=Shard(-1))
    # DSA keeps num_heads/num_index_heads on the parent attention module while
    # the head-sharded weight lives in this leaf boundary.  Tag one canonical

    # DSA keeps num_heads/num_index_heads on the parent attention module while
    # the head-sharded weight lives in this leaf boundary.  Tag one canonical
    # Q projection so D-17 adjusts the owner exactly once after TP parameter
    # sharding.
    if leaf == "linear_qb":
        spec._head_count_owner = fqn.rsplit(".", 1)[0]  # pylint: disable=protected-access
    return spec


def _output_projection_spec(fqn, module, named_modules) -> ModuleShardingSpec:
    """Output projection whose reduce-scatter axis follows the runtime layout.

    implementation instead of the FQN: regular decoder layers may mix DSA and
    MLA attention, so an ``mtp_block`` name alone cannot determine the runtime
    layout.
    """
    parent_fqn = fqn.rsplit(".", 1)[0]
    attention_type = getattr(named_modules.get(parent_fqn), "attention_type", None)
    if attention_type == "mla":
        sequence_dim = 0
    elif attention_type in {"gqa", "dsa"}:
        sequence_dim = 1
    else:
        # Retain the legacy fallback for architecture-compatible test doubles
        # or external modules that do not expose attention_type.
        sequence_dim = 0 if ".mtp_block." in fqn else 1
    return _linear(
        module, param=Shard(1), in_src=Shard(-1), in_dst=Shard(-1),
        out_src=Partial(), out_dst=Shard(sequence_dim))

def _layernorm_spec(module) -> ModuleShardingSpec:
    """Sequence-parallel q/k layernorm boundary."""
    return ModuleShardingSpec(
        params=_direct_params(module, Replicate()),
        in_src={"hidden_states": {"tp": Shard(1)}},
        in_dst={"hidden_states": {"tp": Shard(1)}},
        out_src={"output": {"tp": Shard(1)}},

def _rotary_spec() -> ModuleShardingSpec:
    """Replicated rotary-embedding input contract."""
    inputs = {name: {"tp": Replicate()} for name in ("t", "cos", "sin")}
    return ModuleShardingSpec(params={}, in_src=inputs, in_dst=inputs, out_src={}, out_dst={})


def _sparse_indexer_spec() -> ModuleShardingSpec:
    """Contract for the sparse-lightning indexer KLL-loss leaf."""
    src = {
        "index_query": {"tp": Replicate()},
        "index_key": {"tp": Replicate()},
        "merge_weight": {"tp": Replicate()},
        "query": {"tp": Shard(1)},

        "key_rope": {"tp": Replicate()},
        "actual_seq_qlen": {"tp": Replicate()},
        "actual_seq_klen": {"tp": Replicate()},
    }
    return ModuleShardingSpec(
        params={}, in_src=src,
        in_dst={name: {"tp": Replicate()} for name in src},
        out_src={}, out_dst={})

def _attention_leaf_spec(fqn, module, leaf, named_modules) -> Optional[ModuleShardingSpec]:
    """Build the boundary spec for one attention-leaf FQN (None if unhandled)."""
    if leaf in {"linear_qb", "index_linear_qb", "linear_merge_weight"}:
        return _head_sharded_projection_spec(fqn, module, leaf)
    if leaf in {"linear_qkv", "index_linear_k"}:
        return _linear(
            module, param=Replicate(), in_src=Shard(1), in_dst=Shard(1),
            out_src=Shard(1), out_dst=Shard(1))
    if leaf == "linear_kvb":
        return _linear(
            module, param=Shard(0), in_src=Replicate(), in_dst=Replicate(),
            out_src=Shard(-1), out_dst=Shard(-1))
    if leaf == "linear_proj":
        return _output_projection_spec(fqn, module, named_modules)
    if leaf in {"q_layernorm", "k_layernorm", "index_k_layernorm"}:
        return _layernorm_spec(module)
    if leaf in {"rotary_emb", "gather_rotary_emb"}:
        return _rotary_spec()
    if leaf == "sparse_lightning_indexer_kllloss":
        return _sparse_indexer_spec()
    return None


def build_dsa_specs(model) -> Dict[str, ModuleShardingSpec]:
    """Materialize DSA leaf-boundary specs for a structurally matched model.

    preserve sequence parallelism, some shard index/query heads, and
    ``linear_kvb.weight`` is consumed directly.  The template is selected by
    this module contract instead of ``config.architectures``/``model_type``.
    """
    specs = {}
    named_modules = dict(model.named_modules())
    attention_roots = _matching_attention_roots(named_modules)
    if not attention_roots:
        return specs

    for fqn, module in named_modules.items():
        if fqn in attention_roots:
            sink_spec = _attention_sink_spec(module)
            if sink_spec is not None:
                specs[fqn] = sink_spec

        # The integration gathers the language-model SP output before its
        # multi-token prediction vocabulary heads, so the vocab-parallel
        # projection consumes a replicated sequence (not Shard(1)).
        if fqn == "lm_head" or fqn.endswith(".lm_head"):
            specs[fqn] = _lm_head_spec(module)
            continue
        if fqn.endswith(".language_model.embed_tokens"):
            specs[fqn] = _embed_tokens_spec(module)
            continue
        if not re.search(_ATTN, fqn) or not any(
            fqn.startswith(f"{root}.") for root in attention_roots
        ):
            continue
        leaf = fqn.rsplit(".", 1)[-1]
        leaf_spec = _attention_leaf_spec(fqn, module, leaf, named_modules)
        if leaf_spec is not None:
            specs[fqn] = leaf_spec
    return specs

})


def _direct_params(module):
    return {
        name: {"tp": Replicate()}
        for name, _ in module.named_parameters(recurse=False)
    }

    """Return whether *model* contains an independently owned MHC subtree."""
    for fqn, module in model.named_modules():
        if not _is_mhc_module_fqn(fqn):
            continue
        if any(module.named_parameters(recurse=False)):
            return True
    return False


def build_mhc_specs(model) -> Dict[str, ModuleShardingSpec]:

    Its coefficients and small projections are therefore replicated across TP
    ranks.  Each physical owner is registered separately so an outer module
    never claims parameters from a nested module subtree.
    """
    specs = {}
    for fqn, module in model.named_modules():
        if not _is_mhc_module_fqn(fqn):
            continue
        params = _direct_params(module)
        if params:
            specs[fqn] = ModuleShardingSpec(
                params=params,
                is_boundary=False,
            )
    return specs

            continue
        children = dict(module.named_children())
        if not _MTP_CHILDREN.issubset(children):
            continue
        if any(children["prev_proj"].named_parameters(recurse=False)):
            matches[fqn] = module
    return matches


def matches_mtp_template(model: Any) -> bool:

    Returns:
        Fully declared sharding specs keyed by ``prev_proj`` module FQN.
    """
    specs = {}
    for layer_fqn, layer in _matching_mtp_layers(model).items():
        prev_proj = layer.prev_proj
        params = {
            name: {"tp": Replicate()}
            for name, _ in prev_proj.named_parameters(recurse=False)
        }
        specs[f"{layer_fqn}.prev_proj"] = ModuleShardingSpec(
            params=params,
            in_src={"input": {"tp": Shard(1)}},
            in_dst={"input": {"tp": Shard(1)}},
            out_src={"output": {"tp": Shard(1)}},

            in_dst={"input": {"tp": Shard(1)}},
            out_src={"output": {"tp": Shard(1)}},
            out_dst={"output": {"tp": Shard(1)}},
        )
    return specs

                    # template records that parent explicitly; update it after
                    # sharding the tagged leaf.  update_module_head_counts is
                    # idempotent, so a shared/canonical owner remains safe if
                    # a manually assembled plan tags it more than once.
                    owner_module = _resolve_module(model, head_count_owner)
                    tp_size = mesh["tp"].size() if "tp" in plan.mesh_dim_names else 1
                    update_module_head_counts(
                        owner_module, tp_size, head_count_owner)
                else:
                    maybe_update_head_counts(
                        module,

        """
        for name, matcher, builder in _STRUCTURAL_TEMPLATE_PROVIDERS:
            if not matcher(model):
                continue
            specs = builder(model)
            logger.debug(
                "Matched structural sharding template %s (%d specs)",
                name,
                len(specs),
            )
            for fqn, spec in specs.items():
                plan.modules[fqn] = spec

    def _finalize_boundary_specs(
        self,
        plan: ShardingPlan,

    for fqn, module in model.named_modules():
        shared_expert = getattr(module, "shared_expert", None)
        if shared_expert is None:
            continue
        if not all(hasattr(module, name) for name in ("experts", "token_dispatcher")):
            continue
        if not all(hasattr(shared_expert, name) for name in _SHARED_EXPERT_LINEARS):
            continue
        matches[f"{fqn}.shared_expert" if fqn else "shared_expert"] = shared_expert
    return matches


def matches_shared_expert_template(model: Any) -> bool:

    TP ranks participate in the EP dispatch domain, so these projections must
    not additionally shard their hidden dimension.  Each rank independently
    processes its existing sequence shard and returns the same layout.
    """
    specs = {}
    for shared_fqn, shared_expert in _matching_shared_experts(model).items():
        for linear_name in _SHARED_EXPERT_LINEARS:
            linear = getattr(shared_expert, linear_name)
            params = {
                name: {"tp": Replicate()}
                for name, _ in linear.named_parameters(recurse=False)
            }
            specs[f"{shared_fqn}.{linear_name}"] = ModuleShardingSpec(
                params=params,
                in_src={"hidden_states": {"tp": Shard(1)}},
                in_dst={"hidden_states": {"tp": Shard(1)}},
                out_src={"output": {"tp": Shard(1)}},

                in_dst={"hidden_states": {"tp": Shard(1)}},
                out_src={"output": {"tp": Shard(1)}},
                out_dst={"output": {"tp": Shard(1)}},
            )
    return specs

        label_smoothing,
        _is_floating_torch,
    )

    input_local = input_tensor
    if input_dtensor is not None:
        input_local = _get_local_tensor(input_dtensor)

    local_vocab_size = input_local.shape[-1]

    if input_local.ndim > 2:
        input_local = input_local.reshape(-1, local_vocab_size)
        target = target.reshape(-1)
    # else:
    #     raise ValueError(
    #         "input must be a DTensor when using loss_parallel. "
    #         f"Got type: {type(input_tensor)}"