Return a reasonable default shape for a given op_type when parsing fails. Based on the 'large' size from bench.py KERNEL_CONFIGS.
(op_type: str)
| 221 | |
| 222 | |
| 223 | def get_default_shape(op_type: str) -> Dict[str, int]: |
| 224 | """ |
| 225 | Return a reasonable default shape for a given op_type when parsing fails. |
| 226 | Based on the 'large' size from bench.py KERNEL_CONFIGS. |
| 227 | """ |
| 228 | defaults: Dict[str, Dict[str, int]] = { |
| 229 | "matmul": {"M": 2048, "N": 2048, "K": 2048}, |
| 230 | "flash_attention": {"batch": 2, "heads": 32, "seq_len": 1024, "head_dim": 64}, |
| 231 | "layernorm": {"batch": 4096, "dim": 2048}, |
| 232 | "softmax": {"rows": 4096, "cols": 4096}, |
| 233 | "cross_entropy": {"batch": 4096, "vocab": 32000}, |
| 234 | "fused_mlp": {"batch": 2048, "dim": 2048, "hidden": 5504}, |
| 235 | "rmsnorm": {"M": 4096, "N": 4096}, |
| 236 | "reduce": {"M": 4096, "N": 4096}, |
| 237 | "rotary_embedding": {"batch": 2, "heads": 32, "seq_len": 1024, "head_dim": 128}, |
| 238 | } |
| 239 | return defaults.get(op_type, {"M": 2048, "N": 2048}) |
| 240 | |
| 241 | |
| 242 | # --------------------------------------------------------------------------- |