Entry point called by bench.py. Must match reference.cross_entropy_ref signature.
(logits: torch.Tensor, targets: torch.Tensor)
| 149 | |
| 150 | |
| 151 | def kernel_fn(logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: |
| 152 | """Entry point called by bench.py. Must match reference.cross_entropy_ref signature.""" |
| 153 | assert logits.is_cuda and targets.is_cuda |
| 154 | |
| 155 | if logits.dtype != torch.float16: |
| 156 | logits = logits.to(torch.float16) |
| 157 | |
| 158 | mod = _get_module() |
| 159 | loss = mod.cross_entropy_cuda(logits, targets) |
| 160 | |
| 161 | # Loss is always returned as float32 (scalar loss, matches F.cross_entropy) |
| 162 | return loss.to(torch.float32) |
nothing calls this directly
no test coverage detected