| 169 | torch.manual_seed(seed) |
| 170 | return _sample_compiled(*args, **kwargs) |
| 171 | |
| 172 | |
| 173 | @nvtx.annotate() |
| 174 | @torch.compile(fullgraph=True) |
| 175 | def sequential_sample_pt( |
| 176 | weights: torch.Tensor, # [V, D] |
| 177 | hidden_states: torch.Tensor, # [n_hidden_states, D] |
| 178 | num_samples: int, |
| 179 | temperature: torch.Tensor, # scalar (0-d) |
| 180 | **kwargs, |
| 181 | ): |
| 182 | device = weights.device |
| 183 | V, D = weights.shape # noqa: N806 |
| 184 | H, D2 = hidden_states.shape # noqa: N806 |
| 185 | if D2 != D: |
| 186 | raise ValueError( |
| 187 | f"hidden_states second dimension ({D2}) must match weights first dimension ({D})" |
| 188 | ) |
| 189 | block_size = 8192 |
| 190 | # compute logits blocks |
| 191 | gumbel_max = torch.full((num_samples, H), float("-inf"), device=device) |
| 192 | gumbel_max_idx = torch.empty(size=(num_samples, H), dtype=torch.long, device=device) |
| 193 | n_blocks = cdiv(V, block_size) |
| 194 | for blk_idx in range(n_blocks): |
| 195 | idx_from = blk_idx * block_size |
| 196 | idx_to = (blk_idx + 1) * block_size |
| 197 | w_blk = weights[idx_from:idx_to, :] # [block_size, D] |
| 198 | logits_blk = hidden_states @ w_blk.T / temperature # [n_hidden_states, block_size] |
| 199 | unif_noise = torch.rand((num_samples, *logits_blk.shape), device=device) |
| 200 | gumbel_noise = -(-unif_noise.log()).log() |
| 201 | new_max, new_max_idx_local = torch.max(logits_blk + gumbel_noise, dim=2) |
| 202 | new_max_idx_global = idx_from + new_max_idx_local |
| 203 | |
| 204 | replace_mask = new_max > gumbel_max |
| 205 | gumbel_max = torch.where(replace_mask, new_max, gumbel_max) |
| 206 | gumbel_max_idx = torch.where(replace_mask, new_max_idx_global, gumbel_max_idx) |