willychan21/ParallelKernelBench_Problems
ParallelKernelBench (benchmark) Reference problems for ParallelKernelBench: a benchmark for LLM-generated multi-GPU CUDA kernels. This dataset contains 87 reference implementations in reference/ and the input tensor specification in utils/input_output_tensors.py. Files Path Description data/problems.parquet One row per problem (tabular access) reference/*.py Reference solution() implementations utils/input_output_tensors.py Input/output tensor… See the full description on the dataset page: https://huggingface.co/datasets/willychan21/ParallelKernelBench_Problems.
0263
1import torch2import torch.distributed as dist3from typing import Union, Tuple, Optional4 5def solution(6 gate_logits: Union[torch.Tensor, Tuple[torch.Tensor, ...]],7 num_experts: int,8 top_k: int = 2,9 attention_mask: Optional[torch.Tensor] = None,10) -> torch.Tensor:11 if isinstance(gate_logits, (tuple, list)):12 compute_device = gate_logits[0].device13 concatenated_gate_logits = torch.cat(14 [layer_gate.to(compute_device) for layer_gate in gate_logits], dim=015 )16 else:17 compute_device = gate_logits.device18 concatenated_gate_logits = gate_logits19 20 routing_weights = torch.nn.functional.softmax(concatenated_gate_logits, dim=-1)21 _, selected_experts = torch.topk(routing_weights, top_k, dim=-1)22 23 expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)24 25 if attention_mask is None:26 tokens_per_expert = torch.mean(expert_mask.float(), dim=0)27 router_prob_per_expert = torch.mean(routing_weights, dim=0)28 else:29 batch_size, sequence_length = attention_mask.shape30 num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length)31 32 expert_attention_mask = (33 attention_mask[None, :, :, None, None]34 .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts))35 .reshape(-1, top_k, num_experts)36 .to(compute_device)37 )38 tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum(expert_attention_mask, dim=0)39 40 router_per_expert_attention_mask = (41 attention_mask[None, :, :, None]42 .expand((num_hidden_layers, batch_size, sequence_length, num_experts))43 .reshape(-1, num_experts)44 .to(compute_device)45 )46 router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum(router_per_expert_attention_mask, dim=0)47 48 overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0))49 overall_loss = overall_loss * num_experts50 51 if dist.is_available() and dist.is_initialized():52 dist.all_reduce(overall_loss, op=dist.ReduceOp.SUM)53 overall_loss = overall_loss / dist.get_world_size()54 55 return overall_loss56