Baseline: matmul for logits followed by argmax. Returns [n_hidden_states, 1].
(
weights: torch.Tensor, # [V, D]
hidden_states: torch.Tensor, # [n_hidden_states, D]
num_samples: int, # ignored (always returns 1 sample per row)
temperature: torch.Tensor, # ignored (greedy)
tp: "TPInfo" = TP1,
**_kwargs,
)
| 137 | |
| 138 | @nvtx.annotate() |
| 139 | def greedy_sample( |
| 140 | weights: torch.Tensor, # [V, D] |
| 141 | hidden_states: torch.Tensor, # [n_hidden_states, D] |
| 142 | num_samples: int, # ignored (always returns 1 sample per row) |
| 143 | temperature: torch.Tensor, # ignored (greedy) |
| 144 | tp: "TPInfo" = TP1, |
| 145 | **_kwargs, |
| 146 | ) -> torch.Tensor: |
| 147 | """Baseline: matmul for logits followed by argmax. Returns [n_hidden_states, 1].""" |
| 148 | logits = hidden_states @ weights.T # [n_hidden_states, V] |
| 149 | if tp.size > 1: |
| 150 | logits = _allgather_logits(logits) |
| 151 | return logits.argmax(dim=-1, keepdim=True) # [n_hidden_states, 1] |
| 152 | |
| 153 | |
| 154 | greedy_sample_compiled = nvtx.annotate()(torch.compile(greedy_sample, fullgraph=True)) |
nothing calls this directly
no test coverage detected