Instructions to use modilify/Modilify-Mk2 with libraries, inference providers, notebooks, and local apps. Follow these links to get started.

Libraries

How to use modilify/Modilify-Mk2 with Transformers:

# Use a pipeline as a high-level helper
from transformers import pipeline

pipe = pipeline("text-generation", model="modilify/Modilify-Mk2", trust_remote_code=True)
messages = [
    {"role": "user", "content": "Who are you?"},
]
pipe(messages)

# pip install -U transformers accelerate
# Load model directly
from transformers import AutoModel
model = AutoModel.from_pretrained("modilify/Modilify-Mk2", trust_remote_code=True, device_map="auto")

Notebooks
Google Colab
Kaggle
Local Apps Settings

vLLM

How to use modilify/Modilify-Mk2 with vLLM:

Install from pip and serve model

# Install vLLM from pip:
pip install vllm
# Start the vLLM server:
vllm serve "modilify/Modilify-Mk2"
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:8000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "modilify/Modilify-Mk2",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Use Docker

docker model run hf.co/modilify/Modilify-Mk2

SGLang

How to use modilify/Modilify-Mk2 with SGLang:

Install from pip and serve model

# Install SGLang from pip:
pip install sglang
# Start the SGLang server:
python3 -m sglang.launch_server \
    --model-path "modilify/Modilify-Mk2" \
    --host 0.0.0.0 \
    --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "modilify/Modilify-Mk2",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Use Docker images

docker run --gpus all \
    --shm-size 32g \
    -p 30000:30000 \
    -v ~/.cache/huggingface:/root/.cache/huggingface \
    --env "HF_TOKEN=<secret>" \
    --ipc=host \
    lmsysorg/sglang:latest \
    python3 -m sglang.launch_server \
        --model-path "modilify/Modilify-Mk2" \
        --host 0.0.0.0 \
        --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "modilify/Modilify-Mk2",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Docker Model Runner
How to use modilify/Modilify-Mk2 with Docker Model Runner:
```
docker model run hf.co/modilify/Modilify-Mk2
```

Modilify-Mk2 / vocab_ops.py

ydy9038074

Publish Modilify Mk2 step 4400

4bbff32 verified 19 days ago

Raw History Blame Contribute Delete

25.1 kB

"""Exact memory-bounded vocabulary projection, sampling, and backward."""

from __future__ import annotations

import math
from collections.abc import Sequence

import torch
from torch.nn import functional as F


def _stable_max(values: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """Return ``max`` indices that stay in-range on MPS NaN/-inf rows."""

    best, index = values.max(dim=-1)
    width = values.shape[-1]
    if width <= 0:
        return best, index
    return best, index.clamp(0, width - 1)


class _ChunkedVocabStatistics(torch.autograd.Function):
    """One projection pass for CE and exact Gumbel-max categorical sampling."""

    @staticmethod
    def forward(
        ctx,
        hidden: torch.Tensor,
        weight: torch.Tensor,
        labels: torch.LongTensor,
        softcap: float,
        temperature: float,
        chunk_size: int,
        top_k: int | None = 40,
        min_p: float | None = 0.05,
    ) -> tuple[
        torch.Tensor,
        torch.Tensor,
        torch.LongTensor,
        torch.Tensor,
        torch.Tensor,
        torch.LongTensor,
        torch.Tensor,
    ]:
        if hidden.ndim != 2 or weight.ndim != 2:
            raise ValueError("Chunked vocabulary tensors must be matrices.")
        if hidden.shape[1] != weight.shape[1]:
            raise ValueError("Hidden and vocabulary projection dimensions differ.")
        if temperature <= 0 or chunk_size <= 0:
            raise ValueError("Temperature and vocabulary chunk size must be positive.")
        labels = labels.to(device=hidden.device, dtype=torch.long)
        if labels.shape != (hidden.shape[0],):
            raise ValueError("Labels must contain one value per hidden row.")

        rows = hidden.shape[0]
        row_indices = torch.arange(rows, device=hidden.device)
        valid = labels.ge(0)
        raw_log_z = torch.full(
            (rows,), -torch.inf, device=hidden.device, dtype=torch.float32
        )
        sample_log_z = torch.full_like(raw_log_z, -torch.inf)
        gold_score = torch.zeros_like(raw_log_z)
        best_gumbel = torch.full_like(raw_log_z, -torch.inf)
        selected_score = torch.zeros_like(raw_log_z)
        selected = torch.zeros(rows, device=hidden.device, dtype=torch.long)
        greedy_score = torch.full_like(raw_log_z, -torch.inf)
        greedy = torch.zeros(rows, device=hidden.device, dtype=torch.long)
        moment_max = torch.full_like(raw_log_z, -torch.inf)
        moment_sum = torch.zeros_like(raw_log_z)
        moment_weighted = torch.zeros_like(raw_log_z)
        vocab_size = weight.shape[0]

        use_constrained = top_k is not None and top_k > 0
        chunk_cand_scores = []
        chunk_cand_tokens = []

        with torch.no_grad():
            for start in range(0, vocab_size, int(chunk_size)):
                stop = min(start + int(chunk_size), vocab_size)
                raw = F.linear(hidden, weight[start:stop])
                scores = (
                    torch.tanh(raw.float() / float(softcap))
                    * float(softcap)
                )
                raw_log_z = torch.logaddexp(
                    raw_log_z,
                    torch.logsumexp(scores, dim=-1),
                )
                sample_scores = scores / float(temperature)
                sample_log_z = torch.logaddexp(
                    sample_log_z,
                    torch.logsumexp(sample_scores, dim=-1),
                )

                in_chunk = valid & labels.ge(start) & labels.lt(stop)
                local_gold = (labels - start).clamp(0, stop - start - 1)
                gold_score = torch.where(
                    in_chunk,
                    scores[row_indices, local_gold],
                    gold_score,
                )

                if use_constrained:
                    k = min(int(top_k), stop - start)
                    c_score, c_idx = torch.topk(sample_scores, k=k, dim=-1)
                    chunk_cand_scores.append(c_score)
                    chunk_cand_tokens.append(c_idx + start)
                else:
                    uniform = torch.rand(
                        sample_scores.shape,
                        device=sample_scores.device,
                        dtype=torch.float32,
                    ).clamp_(
                        min=torch.finfo(torch.float32).tiny,
                        max=1.0 - torch.finfo(torch.float32).eps,
                    )
                    gumbel_scores = sample_scores - torch.log(-torch.log(uniform))
                    chunk_best, chunk_index = _stable_max(gumbel_scores)
                    replace_best = chunk_best.gt(best_gumbel)
                    candidate_score = sample_scores.gather(
                        1, chunk_index[:, None]
                    ).squeeze(-1)
                    best_gumbel = torch.maximum(best_gumbel, chunk_best)
                    selected = torch.where(
                        replace_best,
                        chunk_index + start,
                        selected,
                    )
                    selected_score = torch.where(
                        replace_best,
                        candidate_score,
                        selected_score,
                    )

                chunk_max, chunk_argmax = _stable_max(sample_scores)
                replace_greedy = chunk_max.gt(greedy_score)
                greedy_score = torch.maximum(greedy_score, chunk_max)
                greedy = torch.where(
                    replace_greedy,
                    chunk_argmax + start,
                    greedy,
                )
                shifted = torch.exp(sample_scores - chunk_max[:, None])
                chunk_sum = shifted.sum(dim=-1)
                chunk_weighted = (shifted * sample_scores).sum(dim=-1)
                merged_max = torch.maximum(moment_max, chunk_max)
                previous_scale = torch.exp(moment_max - merged_max)
                chunk_scale = torch.exp(chunk_max - merged_max)
                moment_sum = (
                    moment_sum * previous_scale + chunk_sum * chunk_scale
                )
                moment_weighted = (
                    moment_weighted * previous_scale
                    + chunk_weighted * chunk_scale
                )
                moment_max = merged_max

        if use_constrained:
            all_scores = torch.cat(chunk_cand_scores, dim=-1)
            all_tokens = torch.cat(chunk_cand_tokens, dim=-1)
            global_k = min(int(top_k), all_scores.shape[-1])
            cand_scores, global_idx = torch.topk(all_scores, k=global_k, dim=-1)
            cand_tokens = all_tokens.gather(dim=-1, index=global_idx)
            if min_p is not None and min_p > 0:
                thresh = greedy_score + math.log(float(min_p))
                valid_mask = cand_scores >= thresh[:, None]
                eligible = cand_scores.masked_fill(~valid_mask, -float("inf"))
            else:
                eligible = cand_scores
            uniform = torch.rand(
                eligible.shape,
                device=eligible.device,
                dtype=torch.float32,
            ).clamp_(
                min=torch.finfo(torch.float32).tiny,
                max=1.0 - torch.finfo(torch.float32).eps,
            )
            gumbel = eligible - torch.log(-torch.log(uniform))
            _, chosen_idx = _stable_max(gumbel)
            selected = cand_tokens.gather(dim=-1, index=chosen_idx[:, None]).squeeze(-1)
            selected_score = cand_scores.gather(dim=-1, index=chosen_idx[:, None]).squeeze(-1)

        raw_nll = torch.where(
            valid,
            raw_log_z - gold_score,
            torch.zeros_like(raw_log_z),
        )
        temperature_nll = torch.where(
            valid,
            sample_log_z - gold_score / float(temperature),
            torch.zeros_like(sample_log_z),
        )
        confidence = torch.exp(selected_score - sample_log_z).clamp_(0.0, 1.0)
        greedy_confidence = torch.exp(greedy_score - sample_log_z).clamp_(0.0, 1.0)
        entropy = sample_log_z - moment_weighted / moment_sum.clamp_min(
            torch.finfo(torch.float32).tiny
        )

        ctx.save_for_backward(
            hidden,
            weight,
            labels,
            selected,
            raw_log_z,
            sample_log_z,
            confidence,
            entropy,
        )
        ctx.softcap = float(softcap)
        ctx.temperature = float(temperature)
        ctx.chunk_size = int(chunk_size)
        ctx.mark_non_differentiable(selected, greedy, greedy_confidence)
        return (
            raw_nll,
            temperature_nll,
            selected,
            confidence,
            entropy,
            greedy,
            greedy_confidence,
            sample_log_z - selected_score,
        )

    @staticmethod
    def backward(
        ctx,
        grad_raw_nll: torch.Tensor | None,
        grad_temperature_nll: torch.Tensor | None,
        grad_selected: torch.Tensor | None,
        grad_confidence: torch.Tensor | None,
        grad_entropy: torch.Tensor | None,
        grad_greedy: torch.Tensor | None,
        grad_greedy_confidence: torch.Tensor | None,
        grad_sampled_nll: torch.Tensor | None,
    ):
        del grad_selected, grad_greedy, grad_greedy_confidence
        (
            hidden,
            weight,
            labels,
            selected,
            raw_log_z,
            sample_log_z,
            confidence,
            entropy,
        ) = ctx.saved_tensors
        grad_raw_nll = (
            torch.zeros_like(raw_log_z)
            if grad_raw_nll is None
            else grad_raw_nll.float()
        )
        grad_temperature_nll = (
            torch.zeros_like(sample_log_z)
            if grad_temperature_nll is None
            else grad_temperature_nll.float()
        )
        grad_confidence = (
            torch.zeros_like(confidence)
            if grad_confidence is None
            else grad_confidence.float()
        )
        valid = labels.ge(0)
        row_indices = torch.arange(hidden.shape[0], device=hidden.device)
        grad_hidden = torch.zeros_like(hidden, dtype=torch.float32)
        confidence_scale = (
            grad_confidence * confidence / ctx.temperature
        )
        if grad_sampled_nll is not None:
            confidence_scale = confidence_scale - grad_sampled_nll.float() / ctx.temperature
        hidden_fp32 = hidden.float()
        for start in range(0, weight.shape[0], ctx.chunk_size):
            stop = min(start + ctx.chunk_size, weight.shape[0])
            # The B16 training shape is 4,096 rows by a 32,768-wide vocabulary
            # chunk.  MPS BF16 GEMM can return NaNs for this backward-only
            # projection even when every saved activation is finite.  Recompute
            # scores and accumulate the exact VJP in FP32; cast the finished
            # hidden gradient back to the model dtype only once below.
            weight_fp32 = weight[start:stop].float()
            scores = torch.tanh(
                F.linear(hidden_fp32, weight_fp32) / ctx.softcap
            ) * ctx.softcap
            # Non-deterministic MPS GEMM can make a recomputed score slightly
            # exceed the saved log-partition.  Clamp the exponent so softmax
            # probabilities stay in [0, 1] instead of overflowing to Inf.
            raw_probability = torch.exp(
                (scores - raw_log_z[:, None]).clamp(max=0.0)
            )
            sample_logit = (
                scores / ctx.temperature - sample_log_z[:, None]
            ).clamp(max=0.0)
            sample_probability = torch.exp(sample_logit)
            score_gradient = (
                raw_probability * (grad_raw_nll * valid)[:, None]
                + sample_probability
                * (grad_temperature_nll * valid)[:, None]
                / ctx.temperature
                - sample_probability * confidence_scale[:, None]
            )
            if grad_entropy is not None:
                grad_entropy_f = grad_entropy.float()
                entropy_grad = -(
                    grad_entropy_f / ctx.temperature
                )[:, None] * sample_probability * (
                    sample_logit + entropy[:, None]
                )
                score_gradient = score_gradient + entropy_grad

            gold_in_chunk = valid & labels.ge(start) & labels.lt(stop)
            gold_local = (labels - start).clamp(0, stop - start - 1)
            score_gradient[row_indices, gold_local] -= (
                grad_raw_nll * gold_in_chunk
                + grad_temperature_nll
                * gold_in_chunk
                / ctx.temperature
            )
            selected_in_chunk = selected.ge(start) & selected.lt(stop)
            selected_local = (selected - start).clamp(0, stop - start - 1)
            score_gradient[row_indices, selected_local] += (
                confidence_scale * selected_in_chunk
            )
            score_gradient *= 1.0 - (scores / ctx.softcap).square()
            score_gradient = torch.nan_to_num(
                score_gradient, nan=0.0, posinf=0.0, neginf=0.0
            )
            chunk_gradient = F.linear(
                score_gradient,
                weight_fp32.transpose(0, 1),
            )
            grad_hidden.add_(chunk_gradient)

        grad_hidden = torch.nan_to_num(
            grad_hidden, nan=0.0, posinf=0.0, neginf=0.0
        )
        return (
            grad_hidden.to(hidden.dtype),
            None,
            None,
            None,
            None,
            None,
            None,
            None,
        )


def chunked_vocab_training_statistics(
    hidden: torch.Tensor,
    weight: torch.Tensor,
    labels: torch.LongTensor,
    *,
    softcap: float,
    temperature: float,
    chunk_size: int,
    top_k: int | None = 40,
    min_p: float | None = 0.05,
    return_sampled_nll: bool = False,
) -> tuple[
    torch.Tensor,
    torch.Tensor,
    torch.LongTensor,
    torch.Tensor,
    torch.Tensor,
    torch.LongTensor,
    torch.Tensor,
]:
    """Return exact losses plus sampled and greedy proposal statistics."""

    shape = labels.shape
    outputs = _ChunkedVocabStatistics.apply(
        hidden.reshape(-1, hidden.shape[-1]),
        weight,
        labels.reshape(-1),
        float(softcap),
        float(temperature),
        int(chunk_size),
        top_k,
        min_p,
    )
    (
        raw_nll,
        temperature_nll,
        proposal,
        confidence,
        entropy,
        greedy,
        greedy_confidence,
        sampled_nll,
    ) = outputs
    result = (
        raw_nll.view(shape),
        temperature_nll.view(shape),
        proposal.view(shape),
        confidence.view(shape),
        entropy.view(shape),
        greedy.view(shape),
        greedy_confidence.view(shape),
    )
    return result + (sampled_nll.view(shape),) if return_sampled_nll else result


@torch.no_grad()
def chunked_vocab_statistics(
    hidden: torch.Tensor,
    weight: torch.Tensor,
    *,
    softcap: float,
    temperature: float,
    chunk_size: int,
    repetition_token_mask: torch.BoolTensor | None = None,
    repetition_penalty: float = 1.0,
    sampling_generators: Sequence[torch.Generator] | None = None,
    top_k: int | None = 40,
    min_p: float | None = 0.05,
) -> tuple[
    torch.LongTensor,
    torch.Tensor,
    torch.Tensor,
    torch.LongTensor,
    torch.Tensor,
]:
    """Return exact inference statistics, optionally penalizing seen tokens."""

    if not math.isfinite(repetition_penalty) or repetition_penalty <= 0:
        raise ValueError("`repetition_penalty` must be a finite positive number.")

    # Preserve the original custom-autograd path exactly when disabled. In
    # particular, this keeps the same chunk-local RNG draws and default output.
    if repetition_penalty != 1.0 or sampling_generators is not None:
        if hidden.ndim < 2 or weight.ndim != 2:
            raise ValueError("Chunked vocabulary tensors must have matrix features.")
        if hidden.shape[-1] != weight.shape[1]:
            raise ValueError("Hidden and vocabulary projection dimensions differ.")
        if temperature <= 0 or chunk_size <= 0:
            raise ValueError("Temperature and vocabulary chunk size must be positive.")
        expected_mask_shape = (hidden.shape[0], weight.shape[0])
        if repetition_penalty != 1.0:
            if repetition_token_mask is None or repetition_token_mask.shape != expected_mask_shape:
                raise ValueError(
                    "`repetition_token_mask` must have shape [batch, vocabulary]."
                )
            if repetition_token_mask.dtype != torch.bool:
                raise ValueError("`repetition_token_mask` must be a boolean tensor.")
        if sampling_generators is not None and len(sampling_generators) != hidden.shape[0]:
            raise ValueError(
                "`sampling_generators` must contain one generator per batch row."
            )

        output_shape = hidden.shape[:-1]
        flat_hidden = hidden.reshape(-1, hidden.shape[-1])
        rows_per_batch = math.prod(hidden.shape[1:-1])
        batch_indices = torch.arange(
            flat_hidden.shape[0], device=hidden.device
        ).div(rows_per_batch, rounding_mode="floor")
        sample_log_z = torch.full(
            (flat_hidden.shape[0],),
            -torch.inf,
            device=hidden.device,
            dtype=torch.float32,
        )
        best_gumbel = torch.full_like(sample_log_z, -torch.inf)
        selected_score = torch.zeros_like(sample_log_z)
        selected = torch.zeros(
            flat_hidden.shape[0], device=hidden.device, dtype=torch.long
        )
        greedy_score = torch.full_like(sample_log_z, -torch.inf)
        greedy = torch.zeros_like(selected)
        moment_max = torch.full_like(sample_log_z, -torch.inf)
        moment_sum = torch.zeros_like(sample_log_z)
        moment_weighted = torch.zeros_like(sample_log_z)

        use_constrained = top_k is not None and top_k > 0
        chunk_cand_scores = []
        chunk_cand_tokens = []

        for start in range(0, weight.shape[0], int(chunk_size)):
            stop = min(start + int(chunk_size), weight.shape[0])
            raw = F.linear(flat_hidden, weight[start:stop])
            scores = torch.tanh(raw.float() / float(softcap)) * float(softcap)
            if repetition_penalty != 1.0:
                assert repetition_token_mask is not None
                seen = repetition_token_mask[:, start:stop].index_select(
                    0, batch_indices
                )
                penalized = torch.where(
                    scores < 0,
                    scores * float(repetition_penalty),
                    scores / float(repetition_penalty),
                )
                scores = torch.where(seen, penalized, scores)
            sample_scores = scores / float(temperature)
            sample_log_z = torch.logaddexp(
                sample_log_z, torch.logsumexp(sample_scores, dim=-1)
            )

            if use_constrained:
                k = min(int(top_k), stop - start)
                c_score, c_idx = torch.topk(sample_scores, k=k, dim=-1)
                chunk_cand_scores.append(c_score)
                chunk_cand_tokens.append(c_idx + start)
            else:
                if sampling_generators is None:
                    uniform = torch.rand(
                        sample_scores.shape,
                        device=sample_scores.device,
                        dtype=torch.float32,
                    )
                else:
                    per_request_shape = (
                        rows_per_batch,
                        sample_scores.shape[-1],
                    )
                    uniform = torch.cat(
                        [
                            torch.rand(
                                per_request_shape,
                                device=sample_scores.device,
                                dtype=torch.float32,
                                generator=generator,
                            )
                            for generator in sampling_generators
                        ],
                        dim=0,
                    )
                uniform = uniform.clamp_(
                    min=torch.finfo(torch.float32).tiny,
                    max=1.0 - torch.finfo(torch.float32).eps,
                )
                gumbel_scores = sample_scores - torch.log(-torch.log(uniform))
                chunk_best, chunk_index = _stable_max(gumbel_scores)
                replace_best = chunk_best.gt(best_gumbel)
                candidate_score = sample_scores.gather(
                    1, chunk_index[:, None]
                ).squeeze(-1)
                best_gumbel = torch.maximum(best_gumbel, chunk_best)
                selected = torch.where(replace_best, chunk_index + start, selected)
                selected_score = torch.where(
                    replace_best, candidate_score, selected_score
                )

            chunk_max, chunk_argmax = _stable_max(sample_scores)
            replace_greedy = chunk_max.gt(greedy_score)
            greedy_score = torch.maximum(greedy_score, chunk_max)
            greedy = torch.where(replace_greedy, chunk_argmax + start, greedy)
            shifted = torch.exp(sample_scores - chunk_max[:, None])
            chunk_sum = shifted.sum(dim=-1)
            chunk_weighted = (shifted * sample_scores).sum(dim=-1)
            merged_max = torch.maximum(moment_max, chunk_max)
            previous_scale = torch.exp(moment_max - merged_max)
            chunk_scale = torch.exp(chunk_max - merged_max)
            moment_sum = moment_sum * previous_scale + chunk_sum * chunk_scale
            moment_weighted = (
                moment_weighted * previous_scale + chunk_weighted * chunk_scale
            )
            moment_max = merged_max

        if use_constrained:
            all_scores = torch.cat(chunk_cand_scores, dim=-1)
            all_tokens = torch.cat(chunk_cand_tokens, dim=-1)
            global_k = min(int(top_k), all_scores.shape[-1])
            cand_scores, global_idx = torch.topk(all_scores, k=global_k, dim=-1)
            cand_tokens = all_tokens.gather(dim=-1, index=global_idx)
            if min_p is not None and min_p > 0:
                thresh = greedy_score + math.log(float(min_p))
                valid_mask = cand_scores >= thresh[:, None]
                eligible = cand_scores.masked_fill(~valid_mask, -float("inf"))
            else:
                eligible = cand_scores
            if sampling_generators is None:
                uniform = torch.rand(
                    eligible.shape,
                    device=eligible.device,
                    dtype=torch.float32,
                )
            else:
                per_request_shape = (
                    rows_per_batch,
                    eligible.shape[-1],
                )
                uniform = torch.cat(
                    [
                        torch.rand(
                            per_request_shape,
                            device=eligible.device,
                            dtype=torch.float32,
                            generator=generator,
                        )
                        for generator in sampling_generators
                    ],
                    dim=0,
                )
            uniform = uniform.clamp_(
                min=torch.finfo(torch.float32).tiny,
                max=1.0 - torch.finfo(torch.float32).eps,
            )
            gumbel = eligible - torch.log(-torch.log(uniform))
            _, chosen_idx = _stable_max(gumbel)
            selected = cand_tokens.gather(dim=-1, index=chosen_idx[:, None]).squeeze(-1)
            selected_score = cand_scores.gather(dim=-1, index=chosen_idx[:, None]).squeeze(-1)

        confidence = torch.exp(selected_score - sample_log_z).clamp_(0.0, 1.0)
        greedy_confidence = torch.exp(greedy_score - sample_log_z).clamp_(0.0, 1.0)
        entropy = sample_log_z - moment_weighted / moment_sum.clamp_min(
            torch.finfo(torch.float32).tiny
        )
        return (
            selected.view(output_shape),
            confidence.view(output_shape),
            entropy.view(output_shape),
            greedy.view(output_shape),
            greedy_confidence.view(output_shape),
        )

    labels = torch.full(
        hidden.shape[:-1],
        -100,
        device=hidden.device,
        dtype=torch.long,
    )
    (
        _,
        _,
        proposal,
        confidence,
        entropy,
        greedy,
        greedy_confidence,
    ) = chunked_vocab_training_statistics(
        hidden,
        weight,
        labels,
        softcap=softcap,
        temperature=temperature,
        chunk_size=chunk_size,
        top_k=top_k,
        min_p=min_p,
    )
    return proposal, confidence, entropy, greedy, greedy_confidence


__all__ = [
    "chunked_vocab_statistics",
    "chunked_vocab_training_statistics",
]

	"""Exact memory-bounded vocabulary projection, sampling, and backward."""

	from __future__ import annotations

	import math
	from collections.abc import Sequence

	import torch
	from torch.nn import functional as F


	def _stable_max(values: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
	"""Return ``max`` indices that stay in-range on MPS NaN/-inf rows."""

	best, index = values.max(dim=-1)
	width = values.shape[-1]
	if width <= 0:
	return best, index
	return best, index.clamp(0, width - 1)


	class _ChunkedVocabStatistics(torch.autograd.Function):
	"""One projection pass for CE and exact Gumbel-max categorical sampling."""

	@staticmethod
	def forward(
	ctx,
	hidden: torch.Tensor,
	weight: torch.Tensor,
	labels: torch.LongTensor,
	softcap: float,
	temperature: float,
	chunk_size: int,
	top_k: int \| None = 40,
	min_p: float \| None = 0.05,
	) -> tuple[
	torch.Tensor,
	torch.Tensor,
	torch.LongTensor,
	torch.Tensor,
	torch.Tensor,
	torch.LongTensor,
	torch.Tensor,
	]:
	if hidden.ndim != 2 or weight.ndim != 2:
	raise ValueError("Chunked vocabulary tensors must be matrices.")
	if hidden.shape[1] != weight.shape[1]:
	raise ValueError("Hidden and vocabulary projection dimensions differ.")
	if temperature <= 0 or chunk_size <= 0:
	raise ValueError("Temperature and vocabulary chunk size must be positive.")
	labels = labels.to(device=hidden.device, dtype=torch.long)
	if labels.shape != (hidden.shape[0],):
	raise ValueError("Labels must contain one value per hidden row.")

	rows = hidden.shape[0]
	row_indices = torch.arange(rows, device=hidden.device)
	valid = labels.ge(0)
	raw_log_z = torch.full(
	(rows,), -torch.inf, device=hidden.device, dtype=torch.float32
	)
	sample_log_z = torch.full_like(raw_log_z, -torch.inf)
	gold_score = torch.zeros_like(raw_log_z)
	best_gumbel = torch.full_like(raw_log_z, -torch.inf)
	selected_score = torch.zeros_like(raw_log_z)
	selected = torch.zeros(rows, device=hidden.device, dtype=torch.long)
	greedy_score = torch.full_like(raw_log_z, -torch.inf)
	greedy = torch.zeros(rows, device=hidden.device, dtype=torch.long)
	moment_max = torch.full_like(raw_log_z, -torch.inf)
	moment_sum = torch.zeros_like(raw_log_z)
	moment_weighted = torch.zeros_like(raw_log_z)
	vocab_size = weight.shape[0]

	use_constrained = top_k is not None and top_k > 0
	chunk_cand_scores = []
	chunk_cand_tokens = []

	with torch.no_grad():
	for start in range(0, vocab_size, int(chunk_size)):
	stop = min(start + int(chunk_size), vocab_size)
	raw = F.linear(hidden, weight[start:stop])
	scores = (
	torch.tanh(raw.float() / float(softcap))
	* float(softcap)
	)
	raw_log_z = torch.logaddexp(
	raw_log_z,
	torch.logsumexp(scores, dim=-1),
	)
	sample_scores = scores / float(temperature)
	sample_log_z = torch.logaddexp(
	sample_log_z,
	torch.logsumexp(sample_scores, dim=-1),
	)

	in_chunk = valid & labels.ge(start) & labels.lt(stop)
	local_gold = (labels - start).clamp(0, stop - start - 1)
	gold_score = torch.where(
	in_chunk,
	scores[row_indices, local_gold],
	gold_score,
	)

	if use_constrained:
	k = min(int(top_k), stop - start)
	c_score, c_idx = torch.topk(sample_scores, k=k, dim=-1)
	chunk_cand_scores.append(c_score)
	chunk_cand_tokens.append(c_idx + start)
	else:
	uniform = torch.rand(
	sample_scores.shape,
	device=sample_scores.device,
	dtype=torch.float32,
	).clamp_(
	min=torch.finfo(torch.float32).tiny,
	max=1.0 - torch.finfo(torch.float32).eps,
	)
	gumbel_scores = sample_scores - torch.log(-torch.log(uniform))
	chunk_best, chunk_index = _stable_max(gumbel_scores)
	replace_best = chunk_best.gt(best_gumbel)
	candidate_score = sample_scores.gather(
	1, chunk_index[:, None]
	).squeeze(-1)
	best_gumbel = torch.maximum(best_gumbel, chunk_best)
	selected = torch.where(
	replace_best,
	chunk_index + start,
	selected,
	)
	selected_score = torch.where(
	replace_best,
	candidate_score,
	selected_score,
	)

	chunk_max, chunk_argmax = _stable_max(sample_scores)
	replace_greedy = chunk_max.gt(greedy_score)
	greedy_score = torch.maximum(greedy_score, chunk_max)
	greedy = torch.where(
	replace_greedy,
	chunk_argmax + start,
	greedy,
	)
	shifted = torch.exp(sample_scores - chunk_max[:, None])
	chunk_sum = shifted.sum(dim=-1)
	chunk_weighted = (shifted * sample_scores).sum(dim=-1)
	merged_max = torch.maximum(moment_max, chunk_max)
	previous_scale = torch.exp(moment_max - merged_max)
	chunk_scale = torch.exp(chunk_max - merged_max)
	moment_sum = (
	moment_sum * previous_scale + chunk_sum * chunk_scale
	)
	moment_weighted = (
	moment_weighted * previous_scale
	+ chunk_weighted * chunk_scale
	)
	moment_max = merged_max

	if use_constrained:
	all_scores = torch.cat(chunk_cand_scores, dim=-1)
	all_tokens = torch.cat(chunk_cand_tokens, dim=-1)
	global_k = min(int(top_k), all_scores.shape[-1])
	cand_scores, global_idx = torch.topk(all_scores, k=global_k, dim=-1)
	cand_tokens = all_tokens.gather(dim=-1, index=global_idx)
	if min_p is not None and min_p > 0:
	thresh = greedy_score + math.log(float(min_p))
	valid_mask = cand_scores >= thresh[:, None]
	eligible = cand_scores.masked_fill(~valid_mask, -float("inf"))
	else:
	eligible = cand_scores
	uniform = torch.rand(
	eligible.shape,
	device=eligible.device,
	dtype=torch.float32,
	).clamp_(
	min=torch.finfo(torch.float32).tiny,
	max=1.0 - torch.finfo(torch.float32).eps,
	)
	gumbel = eligible - torch.log(-torch.log(uniform))
	_, chosen_idx = _stable_max(gumbel)
	selected = cand_tokens.gather(dim=-1, index=chosen_idx[:, None]).squeeze(-1)
	selected_score = cand_scores.gather(dim=-1, index=chosen_idx[:, None]).squeeze(-1)

	raw_nll = torch.where(
	valid,
	raw_log_z - gold_score,
	torch.zeros_like(raw_log_z),
	)
	temperature_nll = torch.where(
	valid,
	sample_log_z - gold_score / float(temperature),
	torch.zeros_like(sample_log_z),
	)
	confidence = torch.exp(selected_score - sample_log_z).clamp_(0.0, 1.0)
	greedy_confidence = torch.exp(greedy_score - sample_log_z).clamp_(0.0, 1.0)
	entropy = sample_log_z - moment_weighted / moment_sum.clamp_min(
	torch.finfo(torch.float32).tiny
	)

	ctx.save_for_backward(
	hidden,
	weight,
	labels,
	selected,
	raw_log_z,
	sample_log_z,
	confidence,
	entropy,
	)
	ctx.softcap = float(softcap)
	ctx.temperature = float(temperature)
	ctx.chunk_size = int(chunk_size)
	ctx.mark_non_differentiable(selected, greedy, greedy_confidence)
	return (
	raw_nll,
	temperature_nll,
	selected,
	confidence,
	entropy,
	greedy,
	greedy_confidence,
	sample_log_z - selected_score,
	)

	@staticmethod
	def backward(
	ctx,
	grad_raw_nll: torch.Tensor \| None,
	grad_temperature_nll: torch.Tensor \| None,
	grad_selected: torch.Tensor \| None,
	grad_confidence: torch.Tensor \| None,
	grad_entropy: torch.Tensor \| None,
	grad_greedy: torch.Tensor \| None,
	grad_greedy_confidence: torch.Tensor \| None,
	grad_sampled_nll: torch.Tensor \| None,
	):
	del grad_selected, grad_greedy, grad_greedy_confidence
	(
	hidden,
	weight,
	labels,
	selected,
	raw_log_z,
	sample_log_z,
	confidence,
	entropy,
	) = ctx.saved_tensors
	grad_raw_nll = (
	torch.zeros_like(raw_log_z)
	if grad_raw_nll is None
	else grad_raw_nll.float()
	)
	grad_temperature_nll = (
	torch.zeros_like(sample_log_z)
	if grad_temperature_nll is None
	else grad_temperature_nll.float()
	)
	grad_confidence = (
	torch.zeros_like(confidence)
	if grad_confidence is None
	else grad_confidence.float()
	)
	valid = labels.ge(0)
	row_indices = torch.arange(hidden.shape[0], device=hidden.device)
	grad_hidden = torch.zeros_like(hidden, dtype=torch.float32)
	confidence_scale = (
	grad_confidence * confidence / ctx.temperature
	)
	if grad_sampled_nll is not None:
	confidence_scale = confidence_scale - grad_sampled_nll.float() / ctx.temperature
	hidden_fp32 = hidden.float()
	for start in range(0, weight.shape[0], ctx.chunk_size):
	stop = min(start + ctx.chunk_size, weight.shape[0])
	# The B16 training shape is 4,096 rows by a 32,768-wide vocabulary
	# chunk. MPS BF16 GEMM can return NaNs for this backward-only
	# projection even when every saved activation is finite. Recompute
	# scores and accumulate the exact VJP in FP32; cast the finished
	# hidden gradient back to the model dtype only once below.
	weight_fp32 = weight[start:stop].float()
	scores = torch.tanh(
	F.linear(hidden_fp32, weight_fp32) / ctx.softcap
	) * ctx.softcap
	# Non-deterministic MPS GEMM can make a recomputed score slightly
	# exceed the saved log-partition. Clamp the exponent so softmax
	# probabilities stay in [0, 1] instead of overflowing to Inf.
	raw_probability = torch.exp(
	(scores - raw_log_z[:, None]).clamp(max=0.0)
	)
	sample_logit = (
	scores / ctx.temperature - sample_log_z[:, None]
	).clamp(max=0.0)
	sample_probability = torch.exp(sample_logit)
	score_gradient = (
	raw_probability * (grad_raw_nll * valid)[:, None]
	+ sample_probability
	* (grad_temperature_nll * valid)[:, None]
	/ ctx.temperature
	- sample_probability * confidence_scale[:, None]
	)
	if grad_entropy is not None:
	grad_entropy_f = grad_entropy.float()
	entropy_grad = -(
	grad_entropy_f / ctx.temperature
	)[:, None] * sample_probability * (
	sample_logit + entropy[:, None]
	)
	score_gradient = score_gradient + entropy_grad

	gold_in_chunk = valid & labels.ge(start) & labels.lt(stop)
	gold_local = (labels - start).clamp(0, stop - start - 1)
	score_gradient[row_indices, gold_local] -= (
	grad_raw_nll * gold_in_chunk
	+ grad_temperature_nll
	* gold_in_chunk
	/ ctx.temperature
	)
	selected_in_chunk = selected.ge(start) & selected.lt(stop)
	selected_local = (selected - start).clamp(0, stop - start - 1)
	score_gradient[row_indices, selected_local] += (
	confidence_scale * selected_in_chunk
	)
	score_gradient *= 1.0 - (scores / ctx.softcap).square()
	score_gradient = torch.nan_to_num(
	score_gradient, nan=0.0, posinf=0.0, neginf=0.0
	)
	chunk_gradient = F.linear(
	score_gradient,
	weight_fp32.transpose(0, 1),
	)
	grad_hidden.add_(chunk_gradient)

	grad_hidden = torch.nan_to_num(
	grad_hidden, nan=0.0, posinf=0.0, neginf=0.0
	)
	return (
	grad_hidden.to(hidden.dtype),
	None,
	None,
	None,
	None,
	None,
	None,
	None,
	)


	def chunked_vocab_training_statistics(
	hidden: torch.Tensor,
	weight: torch.Tensor,
	labels: torch.LongTensor,
	*,
	softcap: float,
	temperature: float,
	chunk_size: int,
	top_k: int \| None = 40,
	min_p: float \| None = 0.05,
	return_sampled_nll: bool = False,
	) -> tuple[
	torch.Tensor,
	torch.Tensor,
	torch.LongTensor,
	torch.Tensor,
	torch.Tensor,
	torch.LongTensor,
	torch.Tensor,
	]:
	"""Return exact losses plus sampled and greedy proposal statistics."""

	shape = labels.shape
	outputs = _ChunkedVocabStatistics.apply(
	hidden.reshape(-1, hidden.shape[-1]),
	weight,
	labels.reshape(-1),
	float(softcap),
	float(temperature),
	int(chunk_size),
	top_k,
	min_p,
	)
	(
	raw_nll,
	temperature_nll,
	proposal,
	confidence,
	entropy,
	greedy,
	greedy_confidence,
	sampled_nll,
	) = outputs
	result = (
	raw_nll.view(shape),
	temperature_nll.view(shape),
	proposal.view(shape),
	confidence.view(shape),
	entropy.view(shape),
	greedy.view(shape),
	greedy_confidence.view(shape),
	)
	return result + (sampled_nll.view(shape),) if return_sampled_nll else result


	@torch.no_grad()
	def chunked_vocab_statistics(
	hidden: torch.Tensor,
	weight: torch.Tensor,
	*,
	softcap: float,
	temperature: float,
	chunk_size: int,
	repetition_token_mask: torch.BoolTensor \| None = None,
	repetition_penalty: float = 1.0,
	sampling_generators: Sequence[torch.Generator] \| None = None,
	top_k: int \| None = 40,
	min_p: float \| None = 0.05,
	) -> tuple[
	torch.LongTensor,
	torch.Tensor,
	torch.Tensor,
	torch.LongTensor,
	torch.Tensor,
	]:
	"""Return exact inference statistics, optionally penalizing seen tokens."""

	if not math.isfinite(repetition_penalty) or repetition_penalty <= 0:
	raise ValueError("`repetition_penalty` must be a finite positive number.")

	# Preserve the original custom-autograd path exactly when disabled. In
	# particular, this keeps the same chunk-local RNG draws and default output.
	if repetition_penalty != 1.0 or sampling_generators is not None:
	if hidden.ndim < 2 or weight.ndim != 2:
	raise ValueError("Chunked vocabulary tensors must have matrix features.")
	if hidden.shape[-1] != weight.shape[1]:
	raise ValueError("Hidden and vocabulary projection dimensions differ.")
	if temperature <= 0 or chunk_size <= 0:
	raise ValueError("Temperature and vocabulary chunk size must be positive.")
	expected_mask_shape = (hidden.shape[0], weight.shape[0])
	if repetition_penalty != 1.0:
	if repetition_token_mask is None or repetition_token_mask.shape != expected_mask_shape:
	raise ValueError(
	"`repetition_token_mask` must have shape [batch, vocabulary]."
	)
	if repetition_token_mask.dtype != torch.bool:
	raise ValueError("`repetition_token_mask` must be a boolean tensor.")
	if sampling_generators is not None and len(sampling_generators) != hidden.shape[0]:
	raise ValueError(
	"`sampling_generators` must contain one generator per batch row."
	)

	output_shape = hidden.shape[:-1]
	flat_hidden = hidden.reshape(-1, hidden.shape[-1])
	rows_per_batch = math.prod(hidden.shape[1:-1])
	batch_indices = torch.arange(
	flat_hidden.shape[0], device=hidden.device
	).div(rows_per_batch, rounding_mode="floor")
	sample_log_z = torch.full(
	(flat_hidden.shape[0],),
	-torch.inf,
	device=hidden.device,
	dtype=torch.float32,
	)
	best_gumbel = torch.full_like(sample_log_z, -torch.inf)
	selected_score = torch.zeros_like(sample_log_z)
	selected = torch.zeros(
	flat_hidden.shape[0], device=hidden.device, dtype=torch.long
	)
	greedy_score = torch.full_like(sample_log_z, -torch.inf)
	greedy = torch.zeros_like(selected)
	moment_max = torch.full_like(sample_log_z, -torch.inf)
	moment_sum = torch.zeros_like(sample_log_z)
	moment_weighted = torch.zeros_like(sample_log_z)

	use_constrained = top_k is not None and top_k > 0
	chunk_cand_scores = []
	chunk_cand_tokens = []

	for start in range(0, weight.shape[0], int(chunk_size)):
	stop = min(start + int(chunk_size), weight.shape[0])
	raw = F.linear(flat_hidden, weight[start:stop])
	scores = torch.tanh(raw.float() / float(softcap)) * float(softcap)
	if repetition_penalty != 1.0:
	assert repetition_token_mask is not None
	seen = repetition_token_mask[:, start:stop].index_select(
	0, batch_indices
	)
	penalized = torch.where(
	scores < 0,
	scores * float(repetition_penalty),
	scores / float(repetition_penalty),
	)
	scores = torch.where(seen, penalized, scores)
	sample_scores = scores / float(temperature)
	sample_log_z = torch.logaddexp(
	sample_log_z, torch.logsumexp(sample_scores, dim=-1)
	)

	if use_constrained:
	k = min(int(top_k), stop - start)
	c_score, c_idx = torch.topk(sample_scores, k=k, dim=-1)
	chunk_cand_scores.append(c_score)
	chunk_cand_tokens.append(c_idx + start)
	else:
	if sampling_generators is None:
	uniform = torch.rand(
	sample_scores.shape,
	device=sample_scores.device,
	dtype=torch.float32,
	)
	else:
	per_request_shape = (
	rows_per_batch,
	sample_scores.shape[-1],
	)
	uniform = torch.cat(
	[
	torch.rand(
	per_request_shape,
	device=sample_scores.device,
	dtype=torch.float32,
	generator=generator,
	)
	for generator in sampling_generators
	],
	dim=0,
	)
	uniform = uniform.clamp_(
	min=torch.finfo(torch.float32).tiny,
	max=1.0 - torch.finfo(torch.float32).eps,
	)
	gumbel_scores = sample_scores - torch.log(-torch.log(uniform))
	chunk_best, chunk_index = _stable_max(gumbel_scores)
	replace_best = chunk_best.gt(best_gumbel)
	candidate_score = sample_scores.gather(
	1, chunk_index[:, None]
	).squeeze(-1)
	best_gumbel = torch.maximum(best_gumbel, chunk_best)
	selected = torch.where(replace_best, chunk_index + start, selected)
	selected_score = torch.where(
	replace_best, candidate_score, selected_score
	)

	chunk_max, chunk_argmax = _stable_max(sample_scores)
	replace_greedy = chunk_max.gt(greedy_score)
	greedy_score = torch.maximum(greedy_score, chunk_max)
	greedy = torch.where(replace_greedy, chunk_argmax + start, greedy)
	shifted = torch.exp(sample_scores - chunk_max[:, None])
	chunk_sum = shifted.sum(dim=-1)
	chunk_weighted = (shifted * sample_scores).sum(dim=-1)
	merged_max = torch.maximum(moment_max, chunk_max)
	previous_scale = torch.exp(moment_max - merged_max)
	chunk_scale = torch.exp(chunk_max - merged_max)
	moment_sum = moment_sum * previous_scale + chunk_sum * chunk_scale
	moment_weighted = (
	moment_weighted * previous_scale + chunk_weighted * chunk_scale
	)
	moment_max = merged_max

	if use_constrained:
	all_scores = torch.cat(chunk_cand_scores, dim=-1)
	all_tokens = torch.cat(chunk_cand_tokens, dim=-1)
	global_k = min(int(top_k), all_scores.shape[-1])
	cand_scores, global_idx = torch.topk(all_scores, k=global_k, dim=-1)
	cand_tokens = all_tokens.gather(dim=-1, index=global_idx)
	if min_p is not None and min_p > 0:
	thresh = greedy_score + math.log(float(min_p))
	valid_mask = cand_scores >= thresh[:, None]
	eligible = cand_scores.masked_fill(~valid_mask, -float("inf"))
	else:
	eligible = cand_scores
	if sampling_generators is None:
	uniform = torch.rand(
	eligible.shape,
	device=eligible.device,
	dtype=torch.float32,
	)
	else:
	per_request_shape = (
	rows_per_batch,
	eligible.shape[-1],
	)
	uniform = torch.cat(
	[
	torch.rand(
	per_request_shape,
	device=eligible.device,
	dtype=torch.float32,
	generator=generator,
	)
	for generator in sampling_generators
	],
	dim=0,
	)
	uniform = uniform.clamp_(
	min=torch.finfo(torch.float32).tiny,
	max=1.0 - torch.finfo(torch.float32).eps,
	)
	gumbel = eligible - torch.log(-torch.log(uniform))
	_, chosen_idx = _stable_max(gumbel)
	selected = cand_tokens.gather(dim=-1, index=chosen_idx[:, None]).squeeze(-1)
	selected_score = cand_scores.gather(dim=-1, index=chosen_idx[:, None]).squeeze(-1)

	confidence = torch.exp(selected_score - sample_log_z).clamp_(0.0, 1.0)
	greedy_confidence = torch.exp(greedy_score - sample_log_z).clamp_(0.0, 1.0)
	entropy = sample_log_z - moment_weighted / moment_sum.clamp_min(
	torch.finfo(torch.float32).tiny
	)
	return (
	selected.view(output_shape),
	confidence.view(output_shape),
	entropy.view(output_shape),
	greedy.view(output_shape),
	greedy_confidence.view(output_shape),
	)

	labels = torch.full(
	hidden.shape[:-1],
	-100,
	device=hidden.device,
	dtype=torch.long,
	)
	(
	_,
	_,
	proposal,
	confidence,
	entropy,
	greedy,
	greedy_confidence,
	) = chunked_vocab_training_statistics(
	hidden,
	weight,
	labels,
	softcap=softcap,
	temperature=temperature,
	chunk_size=chunk_size,
	top_k=top_k,
	min_p=min_p,
	)
	return proposal, confidence, entropy, greedy, greedy_confidence


	__all__ = [
	"chunked_vocab_statistics",
	"chunked_vocab_training_statistics",
	]