(bsz_v, bsz_d, bsz_h)
| 359 | @torch.compile(fullgraph=True) |
| 360 | def _local_reduce( |
| 361 | maxs: torch.Tensor, # [num_samples, n_tiles, H] |
| 362 | maxs_idx: torch.Tensor, # [num_samples, n_tiles, H] |
| 363 | vocab_start_index: int, |
| 364 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 365 | """Reduce across V-tiles (dim=1) on this rank and adjust to global vocab indices.""" |
| 366 | idxs = maxs.max(dim=1).indices # [num_samples, H] |
| 367 | samples = maxs_idx.gather(1, idxs.unsqueeze(1)).squeeze(1) # [num_samples, H] |
| 368 | max_values = maxs.gather(1, idxs.unsqueeze(1)).squeeze(1) # [num_samples, H] |
| 369 | samples += vocab_start_index |
| 370 | return samples.T.contiguous(), max_values.T.contiguous() # [H, num_samples] |
| 371 | |
| 372 | |
| 373 | def clip(low, high, x): |
| 374 | return min(max(x, low), high) |
| 375 | |
| 376 | |
| 377 | def is_config_valid(bsz_v, bsz_d, bsz_h): |
| 378 | # Derive limit from hardware constraints: |
| 379 | # - H100/A100 shared memory: 232448 (from Triton logs) |
| 380 | max_bytes = 232448 |
| 381 | |
| 382 | # Memory usage in kernel: |
| 383 | # - logits_blk: bsz_v * bsz_h * 4 bytes (float32, persists) |
| 384 | # - w_blk: bsz_v * bsz_d * 2 bytes (bfloat16, during matmul) |
| 385 | # - hidden_states_blk: bsz_h * bsz_d * 2 bytes (bfloat16, during matmul) |
| 386 | # - noise: bsz_v * bsz_h * 4 bytes (float32) |
| 387 | # - gumbel_noise: bsz_v * bsz_h * 4 bytes (float32) |
| 388 | |
| 389 | # Peak memory during sampling phase: |
| 390 | # logits_blk + gumbel_noise = bsz_v * bsz_h * (4 + 4) bytes |
nothing calls this directly
no outgoing calls
no test coverage detected