| import functools |
| import os |
|
|
| import torch |
| from bench_tirx_prenorm import build_kernel, compare |
|
|
| VALIDATE = os.getenv("MAGI_TIRX_PRENORM_VALIDATE", "0") == "1" |
| checked = set() |
| reported = set() |
|
|
|
|
| @functools.lru_cache(maxsize=1) |
| def executable(): |
| return build_kernel() |
|
|
|
|
| def apply(norm, x, dispatcher): |
| def fallback(): |
| return norm(x, modality_dispatcher=dispatcher) |
|
|
| if ( |
| torch.is_grad_enabled() |
| or norm.training |
| or not x.is_cuda |
| or x.dtype != torch.bfloat16 |
| or x.ndim != 2 |
| or x.shape[1] != 5120 |
| or not x.is_contiguous() |
| or norm.eps != 1e-6 |
| or norm.num_modality not in (1, 3) |
| or norm.weight.dtype != torch.bfloat16 |
| or norm.weight.device != x.device |
| or not norm.weight.is_contiguous() |
| or norm.weight.numel() != norm.num_modality * 5120 |
| ): |
| return fallback() |
|
|
| groups = ( |
| [x.shape[0]] |
| if norm.num_modality == 1 |
| else [int(n) for n in dispatcher.group_size_cpu] |
| ) |
| if ( |
| len(groups) != norm.num_modality |
| or any(n < 0 for n in groups) |
| or sum(groups) != x.shape[0] |
| or x.shape[0] == 0 |
| ): |
| return fallback() |
|
|
| weight = norm.weight |
| version = None if torch.is_inference(weight) else weight._version |
| cached = getattr(norm, "_tirx_prenorm_scale_cache", None) |
| if ( |
| cached is None |
| or cached[0] is not weight |
| or cached[1] != version |
| ): |
| |
| scale = (weight.detach() + 1).reshape( |
| norm.num_modality, 5120 |
| ).contiguous() |
| norm._tirx_prenorm_scale_cache = (weight, version, scale) |
| else: |
| scale = cached[2] |
|
|
| output = torch.empty_like(x) |
| ex = executable() |
| start = 0 |
| for i, size in enumerate(groups): |
| if size: |
| ex( |
| x[start:start + size].reshape(-1), |
| scale[i], |
| output[start:start + size].reshape(-1), |
| size, |
| ) |
| start += size |
|
|
| key = (norm.num_modality, x.device.index) |
| if key not in reported: |
| print( |
| f"TIRX_PRENORM_ACTIVE shape={tuple(x.shape)} " |
| f"groups={groups}", |
| flush=True, |
| ) |
| reported.add(key) |
|
|
| if VALIDATE and key not in checked: |
| print( |
| f"TIRX_PRENORM_CHECK shape={tuple(x.shape)} " |
| f"groups={groups}", |
| flush=True, |
| ) |
| expected = fallback() |
| compare(output, expected) |
| checked.add(key) |
| print("TIRX_PRENORM_CHECK: 真实输入误差检查通过", flush=True) |
|
|
| return output |
|
|
|
|
| def transform_source(source): |
| old = ( |
| "self.pre_norm(hidden_states, " |
| "modality_dispatcher=modality_dispatcher)" |
| ) |
| if source.count(old) != 1: |
| raise RuntimeError("Attention pre_norm入口不匹配，未安装补丁") |
| return source.replace( |
| old, |
| "_tirx_prenorm(self.pre_norm, hidden_states, modality_dispatcher)", |
| 1, |
| ) |
|
|
|
|
| @torch.inference_mode() |
| def prewarm(): |
| if os.getenv("MAGI_TIRX_PREWARM", "0") != "1": |
| return |
| print("TIRX_PREWARM_PRENORM: 开始", flush=True) |
| ex = executable() |
| x = torch.zeros(8 * 5120, device="cuda", dtype=torch.bfloat16) |
| w = torch.ones(5120, device="cuda", dtype=torch.bfloat16) |
| y = torch.empty_like(x) |
| for _ in range(3): |
| ex(x, w, y, 8) |
| torch.cuda.synchronize() |
| print("TIRX_PREWARM_PRENORM: 完成", flush=True) |
|
|
