Shorten verbose CUDA kernel names for legend readability.
(name: str)
| 133 | |
| 134 | |
| 135 | def _shorten_kernel(name: str) -> str: |
| 136 | """Shorten verbose CUDA kernel names for legend readability.""" |
| 137 | if "fused_mm_sample" in name: |
| 138 | return "fmms_kernel" |
| 139 | # FMMS post-kernel reductions (must come before generic triton_red pattern) |
| 140 | if "triton_red_fused_add_gather_max" in name or "triton_per_fused_add_gather_max" in name: |
| 141 | return "reduction" |
| 142 | if "triton_poi_fused_argmax_gather_stack" in name: |
| 143 | return "TP winner select" |
| 144 | # torch.compile fused kernels: split by what they actually compute. |
| 145 | # Names encode the fused ops, e.g. triton_red_fused__softmax_div_ge_... |
| 146 | if "triton_poi_fused" in name or "triton_red_fused" in name or "triton_per_fused" in name: |
| 147 | if "_arange_" in name: |
| 148 | return "arange" |
| 149 | if "_softmax_" in name: |
| 150 | return "softmax + masking" |
| 151 | if "_amax_" in name or "_where_" in name: |
| 152 | return "top-k masking" |
| 153 | if "_div_" in name: |
| 154 | return "temp scaling" |
| 155 | return "torch.compile fused" |
| 156 | if "gemv2T" in name or "gemm" in name: |
| 157 | return "cuBLAS matmul" |
| 158 | if "flashinfer::SamplingFromLogits" in name: |
| 159 | return "FI SamplingFromLogits" |
| 160 | if "flashinfer::TopPSamplingFromProb" in name: |
| 161 | return "FI TopPSamplingFromProb" |
| 162 | if "flashinfer::RadixTopKMask" in name: |
| 163 | return "FI RadixTopKMask" |
| 164 | if "flashinfer::" in name: |
| 165 | start = name.index("flashinfer::") + len("flashinfer::") |
| 166 | end = name.index("<", start) if "<" in name[start:] else len(name) |
| 167 | return f"FI {name[start:end]}" |
| 168 | if "cunn_SoftMax" in name: |
| 169 | return "softmax" |
| 170 | if "distribution_elementwise" in name: |
| 171 | return "rand / multinomial" |
| 172 | if "ArgMaxOps" in name: |
| 173 | return "argmax" |
| 174 | if "reduce_kernel" in name: |
| 175 | return "reduce" |
| 176 | if "vectorized_elementwise" in name or "elementwise_kernel" in name: |
| 177 | return "elementwise" |
| 178 | if "_assert_async" in name: |
| 179 | return "assert_async" |
| 180 | if "direct_copy" in name or "unrolled_elementwise" in name: |
| 181 | return "copy / cast" |
| 182 | if "CatArrayBatchedCopy" in name: |
| 183 | return "all-gather concat" |
| 184 | return name[:40] |
| 185 | |
| 186 | |
| 187 | def _label_bar(ax, rect, text: str, min_height: float = 0.3) -> None: |
nothing calls this directly
no outgoing calls
no test coverage detected