Diff Coverage

        set_inputs_layout = suffix == "input"
        set_outputs_layout = not set_inputs_layout
        register_cell = cell_dict.get(prefix)
        if register_cell is None:
            raise ValueError(f"Cannot find target cell {prefix!r} in sharding_plan")

        _set_layouts(register_cell, value, set_inputs_layout, set_outputs_layout)
        _register_cell_hook(register_cell, set_inputs_layout, set_outputs_layout)

        if len(self_layout.mesh_shape) != 1 or not expanded_actions:
            return False
        first_action = expanded_actions[0]
        if first_action[0] != "int" or first_action[-1] != 0:
            return False
        if not all(
                _is_full_slice_action(action, global_shape)
                for action in expanded_actions[1:]
        ):
            return False
        if not isinstance(placement, Shard) or isinstance(placement, StridedShard):
            return False
        return placement.is_shard(0)

    @staticmethod
    def _infer_shard_dim0_int(self_layout, expanded_actions, global_shape, kind):

def _make_bias_free_forward(owner: nn.Module, original: Callable[..., Any]) -> Callable[..., Any]:
    """Create a forward wrapper that temporarily hides one module bias."""

    @functools.wraps(original)
    def bias_free_forward(*args: Any, **kwargs: Any) -> Any:
        """Run the owner's forward with its bias temporarily hidden."""
        bias = owner.bias
        try:
            owner._parameters["bias"] = None  # pylint: disable=protected-access
            return original(*args, **kwargs)
        finally:
            owner._parameters["bias"] = bias  # pylint: disable=protected-access

    return bias_free_forward


def _install_bias_suppression(module, spec):
    """D-22: make each defer-listed Linear run bias-free inside the region.

    for param_path in spec._deferred_bias_params:  # pylint: disable=protected-access
        owner_path = param_path.rpartition(".")[0]
        owner = module.get_submodule(owner_path) if owner_path else module
        original = owner.forward
        owner.forward = _make_bias_free_forward(owner, original)


def _add_bias_to_primary_output(output, bias, module_name):
    """Add a deferred bias to the primary Tensor while preserving output structure."""

def _is_named_layer(item: Any) -> bool:
    """Return whether an item is a valid ``(fqn, module)`` pair."""
    if not isinstance(item, tuple) or len(item) != 2:
        return False
    name, layer = item
    return isinstance(name, str) and isinstance(layer, nn.Module)


def _normalize_declared_layers(declared: Any) -> list[tuple[str, nn.Module]]:
    """Normalize the model-owned decoder-layer contract."""

    for index, item in enumerate(declared):
        if isinstance(item, nn.Module):
            layers.append((str(index), item))
            continue
        if _is_named_layer(item):
            layers.append(item)
            continue
        raise TypeError(
            "get_compile_layers() entries must be modules or (fqn, module) pairs, "

            keep_kwargs, peer_kwargs = _prepare_head_tail_sdpa_kwargs(
                call_kwargs, q, k, cp_mesh
            )

            def _call_original_sdpa(query, key, value, attention_kwargs):
                return original_sdpa(query, key, value, **attention_kwargs)

            return head_tail_load_balance_attention(
                _call_original_sdpa,
                q,

                cp_mesh,
                ulysses_degree,
            )

            def _call_original_sdpa(call_query, call_key, call_value, call_kwargs):
                return original_sdpa(
                    call_query, call_key, call_value, **call_kwargs
                )

            return hybrid_cp_attention(

                **attention_kwargs: Any) -> Any:
            """Route one intercepted FlexAttention call through Hybrid CP."""
            fired["hit"] = True

            def _call_original_flex(call_query, call_key, call_value, call_kwargs):
                return original_flex_attention(
                    call_query, call_key, call_value, **call_kwargs
                )

            return hybrid_cp_attention(