(
weights: torch.Tensor, # [V, D]
hidden_states: torch.Tensor, # [H, D]
num_samples: int,
temperature: torch.Tensor, # scalar (0-d)
**kwargs,
)
| 53 | |
| 54 | |
| 55 | def fused_mm_sample_helion( |
| 56 | weights: torch.Tensor, # [V, D] |
| 57 | hidden_states: torch.Tensor, # [H, D] |
| 58 | num_samples: int, |
| 59 | temperature: torch.Tensor, # scalar (0-d) |
| 60 | **kwargs, |
| 61 | ) -> torch.Tensor: |
| 62 | temperature = temperature.reshape(1) # Helion kernel needs 1D tensor for indexing |
| 63 | V = weights.size(0) # noqa: N806 |
| 64 | H = hidden_states.size(0) # noqa: N806 |
| 65 | n_tiles = helion.cdiv(V, BLOCK_SIZE_V) |
| 66 | hs_t = hidden_states.T.contiguous() # [D, H] |
| 67 | results = [] |
| 68 | for i in range(num_samples): |
| 69 | tile_maxs = torch.full((n_tiles, H), float("-inf"), device=weights.device) |
| 70 | tile_max_idxs = torch.empty((n_tiles, H), dtype=torch.long, device=weights.device) |
| 71 | seed = torch.randint(0, 2**31, (1,)).item() + i |
| 72 | fused_sample_helion_kernel(weights, hs_t, tile_maxs, tile_max_idxs, temperature, seed) |
| 73 | # Stage 2: reduce across tiles |
| 74 | best_tiles = tile_maxs.argmax(dim=0) # [H] |
| 75 | sample_idx = tile_max_idxs.gather(dim=0, index=best_tiles.unsqueeze(0)).squeeze(0) |
| 76 | results.append(sample_idx) |
| 77 | return torch.stack(results, dim=1) # [H, num_samples] |
| 78 | |
| 79 | |
| 80 | @helion.kernel( |
nothing calls this directly
no test coverage detected