Class SyntheticInputs

src/fused_mm_sampling/testing.py:15–24

In the tensor-parallel case V is sharded among ranks

from .tp_info import TP1, TPInfo


@dataclass
class SyntheticInputs:
    """In the tensor-parallel case V is sharded among ranks"""

    weights: torch.Tensor  # [V, D], bfloat16
    hidden_states: torch.Tensor  # [n_hidden_states, D], bfloat16
    logits: (
        torch.Tensor
    )  # [n_hidden_states, V], float32 (the exact logits before bf16 quantization)
    vocab_size: int
    hidden_size: int