(args: Args)
| 56 | |
| 57 | |
| 58 | def run_proton_profile(args: Args) -> None: |
| 59 | v, d, h = args.vocab_size, HIDDEN_SIZE, args.n_hidden_states |
| 60 | print(f"GPU: {get_gpu_name()}") |
| 61 | print(f"V={v}, D={d}, n_hidden_states={h}, n_samples={args.n_samples}") |
| 62 | |
| 63 | weights = torch.randn((v, d), dtype=torch.bfloat16, device=device) |
| 64 | hidden_states = torch.randn((h, d), dtype=torch.bfloat16, device=device) |
| 65 | temperature = torch.tensor(1.0, device=device) |
| 66 | num_sms = torch.cuda.get_device_properties("cuda").multi_processor_count |
| 67 | max_grid_size_v = triton.cdiv(v, MIN_BLOCK_SIZE_V) |
| 68 | |
| 69 | # Pre-allocate output buffers (reused across iterations) |
| 70 | maxs = torch.empty((args.n_samples, max_grid_size_v, h), dtype=torch.bfloat16, device=device) |
| 71 | maxs_idx = torch.empty_like(maxs, dtype=torch.long) |
| 72 | logits_out = torch.empty((v, h), dtype=torch.float32, device=device) |
| 73 | |
| 74 | def grid(meta): |
| 75 | gv = triton.cdiv(v, meta["BLOCK_SIZE_V"]) |
| 76 | gh = triton.cdiv(h, meta["BLOCK_SIZE_H"]) |
| 77 | return (min(num_sms, gv * gh),) |
| 78 | |
| 79 | kernel_kwargs = dict( |
| 80 | weights_ptr=weights, |
| 81 | hidden_states_ptr=hidden_states, |
| 82 | max_out_ptr=maxs, |
| 83 | max_out_idx_ptr=maxs_idx, |
| 84 | symm_mem_buffer_ptrs=maxs, |
| 85 | vocab_size=v, |
| 86 | hidden_size=d, |
| 87 | n_hidden_states=h, |
| 88 | num_samples=args.n_samples, |
| 89 | temperature_ptr=temperature, |
| 90 | seed=42, |
| 91 | max_grid_size_v=max_grid_size_v, |
| 92 | storage_offset_maxs_idx=0, |
| 93 | tp_rank=0, |
| 94 | tp_world_size=1, |
| 95 | logits_out_ptr=logits_out, |
| 96 | WARP_SPECIALIZE=supports_warp_specialization_cached(), |
| 97 | NUM_SMS=num_sms, |
| 98 | GREEDY_SAMPLING=False, |
| 99 | RETURN_LOGITS=False, |
| 100 | FAN_OUT_TP=False, |
| 101 | ) |
| 102 | |
| 103 | # Start Proton before the first kernel call. When using the TTGIR override |
| 104 | # workflow, the proton dialect must be registered before the overridden TTGIR |
| 105 | # (which contains proton.record ops) is parsed. |
| 106 | mode = Default( |
| 107 | sampling_strategy=SAMPLING_STRATEGY.SELECTIVE, |
| 108 | sampling_options="0", # only profile warp 0 |
| 109 | buffer_type=BUFFER_TYPE.GLOBAL, |
| 110 | ) |
| 111 | proton.start(name="kernel", data="trace", backend="instrumentation", mode=mode) |
| 112 | |
| 113 | print("Profiling...") |
| 114 | fused_mm_sample_triton_kernel[grid](**kernel_kwargs) |
| 115 | torch.cuda.synchronize() |
no test coverage detected