diff --git a/agents/ascend-kernel-developer.md b/agents/ascend-kernel-developer.md index f527fbea..cc76c14a 100644 --- a/agents/ascend-kernel-developer.md +++ b/agents/ascend-kernel-developer.md @@ -50,6 +50,39 @@ Phase 6: 全量用例验证 Phase 7: Trace 记录 (trace-recorder) ``` +## Hook 机制说明 + +本项目的 `.claude/settings.json` 已配置 **toolUse hook**，用于拦截 agent 对 skill 相关脚本的 Bash 调用。 + +### 被拦截的脚本 + +| 类别 | 脚本 | 说明 | +|------|------|------| +| 退化检测 | `validate_tilelang_impl.py` | TileLang AST 退化检测 | +| 退化检测 | `validate_ascendc_impl.py` | AscendC AST 退化检测 | +| 评测脚本 | `evaluate_tilelang.sh` | TileLang 功能验证 | +| 评测脚本 | `evaluate_ascendc.sh` | AscendC 功能验证 | +| 构建脚本 | `utils/build_ascendc.py` | AscendC kernel 编译 | +| 验证脚本 | `utils/verification_ascendc.py` | AscendC 正确性验证 | +| 验证脚本 | `utils/verification_tilelang.py` | TileLang 正确性验证 | +| 性能测试 | `performance.py` | 性能对比测试 | +| 批处理 | `batch_run_performance.sh` | 批量性能测试 | + +### Hook 行为 + +1. **拦截**: 当 agent 通过 Bash tool 调用上述脚本时，hook 自动拦截 +2. **替换执行**: 由 `.claude/hooks/skill_script_hook.py` 接管执行 +3. **等待完成**: hook 等待脚本实际执行完毕（同步阻塞） +4. **返回结果**: 将 exit code、stdout、stderr 以 JSON 格式返回给 agent +5. **Agent 继续**: agent 收到结果后才继续下一步 + +### 配置位置 + +- Hook 脚本: `.claude/hooks/skill_script_hook.py` +- Hook 配置: `.claude/settings.json` + +> **注意**: 非拦截命令（如普通 `ls`、`cp`、`python` 调用其他脚本）会透传执行，不受影响。 + ### 退化检测脚本 | 阶段 | 脚本路径 | 说明 | diff --git a/benchmarks/NPUKernelBench/level1/11_GroupNorm.py b/benchmarks/NPUKernelBench/level1/11_GroupNorm.py index 01295f11..3d44bb79 100644 --- a/benchmarks/NPUKernelBench/level1/11_GroupNorm.py +++ b/benchmarks/NPUKernelBench/level1/11_GroupNorm.py @@ -2,6 +2,7 @@ import torch.nn as nn import json import os +import random class Model(nn.Module): """ @@ -34,13 +35,13 @@ def get_input_groups(): input_groups = [] for idx, case in enumerate(cases): inputs = case["inputs"] - + dtype_map = { "float32": torch.float32, "float16": torch.float16, "bfloat16": torch.bfloat16, } - + x_info = inputs[0] dtype = dtype_map[x_info["dtype"]] if idx % 2 == 0: @@ -53,7 +54,7 @@ def get_input_groups(): num_groups = None weight = None bias = None - + for inp in inputs[1:]: if inp["name"] == "num_groups": num_groups = inp["value"] @@ -61,7 +62,7 @@ def get_input_groups(): weight = torch.randn(inp["shape"], dtype=dtype) elif inp["name"] == "bias": bias = torch.randn(inp["shape"], dtype=dtype) - + input_groups.append([x, num_groups, weight, bias]) return input_groups diff --git a/benchmarks/NPUKernelBench/level1/25_NLLLoss.py b/benchmarks/NPUKernelBench/level1/25_NLLLoss.py index 841bbba2..538e2aa0 100644 --- a/benchmarks/NPUKernelBench/level1/25_NLLLoss.py +++ b/benchmarks/NPUKernelBench/level1/25_NLLLoss.py @@ -33,13 +33,13 @@ def get_input_groups(): json_path = os.path.join(os.path.dirname(__file__), "25_NLLLoss.json") with open(json_path, "r") as f: cases = [json.loads(line) for line in f if line.strip()] - + input_groups = [] for idx, case in enumerate(cases): inputs = case["inputs"] input_info = inputs[0] target_info = inputs[1] - + dtype_map = { "float32": torch.float32, "float16": torch.float16, @@ -56,11 +56,11 @@ def get_input_groups(): target_range = target_info.get("range", [0, input_info["shape"][1] - 1]) target = torch.randint(target_range[0], target_range[1] + 1, tuple(target_info["shape"]), dtype=torch.int64) - + weight = None ignore_index = -100 reduction = "mean" - + for inp in inputs[2:]: if inp["name"] == "weight": weight = torch.randn(inp["shape"], dtype=dtype) @@ -68,7 +68,7 @@ def get_input_groups(): ignore_index = inp["value"] elif inp["name"] == "reduction": reduction = inp["value"] - + input_groups.append([input_tensor, target, weight, ignore_index, reduction]) return input_groups diff --git a/benchmarks/NPUKernelBench/level2/2_GroupNormSwish.py b/benchmarks/NPUKernelBench/level2/2_GroupNormSwish.py index feaf118e..35149dff 100644 --- a/benchmarks/NPUKernelBench/level2/2_GroupNormSwish.py +++ b/benchmarks/NPUKernelBench/level2/2_GroupNormSwish.py @@ -3,6 +3,7 @@ import torch_npu import json import os +import random class Model(nn.Module): """ diff --git a/hooks/skill_script_hook.py b/hooks/skill_script_hook.py new file mode 100755 index 00000000..7c2bf526 --- /dev/null +++ b/hooks/skill_script_hook.py @@ -0,0 +1,204 @@ +#!/usr/bin/env python3 +""" +Skill 脚本执行 Hook (PreToolUse 协议)。 + +对 Bash 工具调用进行拦截： +- 命中 INTERCEPTED_PATTERNS: 直接由本 hook 执行脚本，并以 PreToolUse 协议返回 + permissionDecision=deny + additionalContext，把执行结果（exit_code/stdout/stderr） + 注入回 agent 的上下文，阻止 harness 二次执行。 +- 不命中: 输出 permissionDecision=allow，让 harness 正常执行（不在 hook 中重复执行）。 + +输入：从 stdin 读取 harness PreToolUse JSON： + {"tool_name": "Bash", "tool_input": {"command": "..."}, ...} + +输出：单行 PreToolUse JSON 到 stdout。 +""" + +import json +import os +import re +import subprocess +import sys +import time + + +# 每个条目: (interpreter_regex, script_basename_regex) +# 只匹配“某个解释器 + 脚本路径”这种真正的脚本调用， +# 不匹配命令字符串里仅出现脚本名的情况（如 cat / grep / echo JSON 等）。 +_PY_INTERP = r"(?:python|python3|python3\.\d+)" +_SH_INTERP = r"(?:bash|sh)" + +INTERCEPTED_PATTERNS = [ + (_PY_INTERP, r"validate_tilelang_impl\.py"), + (_PY_INTERP, r"validate_ascendc_impl\.py"), + (_SH_INTERP, r"evaluate_tilelang\.sh"), + (_SH_INTERP, r"evaluate_ascendc\.sh"), + (_PY_INTERP, r"performance\.py"), + (_SH_INTERP, r"batch_run_performance\.sh"), + (_PY_INTERP, r"build_ascendc\.py"), + (_PY_INTERP, r"verification_ascendc\.py"), + (_PY_INTERP, r"verification_tilelang\.py"), +] + +# 命令前缀允许的部分（环境变量赋值、cd ... &&、source ... &&） +# 这里只是用于推断"真实的脚本调用"，不需要覆盖所有 shell 语法。 +_LEADING_PREFIX = ( + r"^\s*" + r"(?:export\s+)?" # 可选的 export 关键字 + r"(?:[A-Za-z_][A-Za-z0-9_]*=\S+\s+)*" # ENV=VAL ... + r"(?:cd\s+\S+\s*&&\s*)?" # cd && + r"(?:export\s+)?" # 可选的 export 关键字 + r"(?:[A-Za-z_][A-Za-z0-9_]*=\S+\s+)*" # 再次允许 ENV=VAL + r"(?:&&\s+)*" # 支持 && 链式连接 +) + +# 项目根目录：优先从环境变量获取，fallback 到 __file__ 计算 +# 避免 hook 被以相对路径调用时 __file__ 解析错误 +PROJECT_ROOT = os.environ.get("PROJECT_ROOT") +if not PROJECT_ROOT: + PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + + +def should_intercept(command: str) -> bool: + # 仅当命令的"实际执行段"是 形式时才拦截。 + # 兼容：环境变量前缀、可选的 `cd &&` 前缀。 + for interp, script in INTERCEPTED_PATTERNS: + # 路径里允许出现 / 和非空白字符；脚本名用 \b 边界 + pattern = ( + _LEADING_PREFIX + + interp + + r"\s+" + + r"(?:\S*/)?" + + script + + r"(?:\s|$)" + ) + if re.search(pattern, command): + return True + return False + + +def emit_pretooluse(permission_decision: str, *, reason: str = "", additional_context: str = ""): + output = { + "hookSpecificOutput": { + "hookEventName": "PreToolUse", + "permissionDecision": permission_decision, + } + } + if reason: + output["hookSpecificOutput"]["permissionDecisionReason"] = reason + if additional_context: + output["hookSpecificOutput"]["additionalContext"] = additional_context + print(json.dumps(output, ensure_ascii=False), flush=True) + + +def truncate(text: str, limit: int = 8000) -> str: + if not text: + return "" + if len(text) <= limit: + return text + head = text[: limit // 2] + tail = text[-limit // 2 :] + return f"{head}\n\n... [truncated {len(text) - limit} chars] ...\n\n{tail}" + + +def execute_intercepted(command: str) -> None: + start_time = time.time() + cwd = os.getcwd() + + # 解析路径：若是相对路径且在 cwd 不存在，则尝试 PROJECT_ROOT 下解析 + # 这里我们直接通过 shell 执行原命令，但先确保 cwd 是 PROJECT_ROOT， + # 这样像 "python skills/..." 之类的相对路径能正确解析。 + if os.path.isdir(PROJECT_ROOT): + cwd = PROJECT_ROOT + + try: + proc = subprocess.run( + command, + shell=True, + cwd=cwd, + capture_output=True, + text=True, + timeout=1800, + ) + exit_code = proc.returncode + stdout = proc.stdout or "" + stderr = proc.stderr or "" + except subprocess.TimeoutExpired as e: + exit_code = 124 + stdout = (e.stdout.decode("utf-8", errors="replace") if e.stdout else "") + stderr = (e.stderr.decode("utf-8", errors="replace") if e.stderr else "") + stderr += "\n[HOOK] 命令执行超时（30 分钟）" + except Exception as e: + exit_code = 1 + stdout = "" + stderr = f"[HOOK] 执行异常: {type(e).__name__}: {e}" + + duration_ms = int((time.time() - start_time) * 1000) + + additional_context = ( + f"[skill_script_hook intercepted execution]\n" + f"command: {command}\n" + f"cwd: {cwd}\n" + f"exit_code: {exit_code}\n" + f"duration_ms: {duration_ms}\n" + f"--- stdout ---\n{truncate(stdout)}\n" + f"--- stderr ---\n{truncate(stderr)}\n" + ) + + status = "成功" if exit_code == 0 else "失败" + reason = ( + f"[Hook 拦截提示] 该命令命中 skill_script_hook 拦截规则，已由 hook 代为执行（{status}，exit_code={exit_code}）。" + f"执行结果见下方 additionalContext，原命令已被替换为 no-op，不会重复执行。" + ) + + # 使用 allow + updatedInput 将原命令替换为 no-op，避免 harness 显示 "Error:" 前缀 + output = { + "hookSpecificOutput": { + "hookEventName": "PreToolUse", + "permissionDecision": "allow", + "permissionDecisionReason": reason, + "additionalContext": additional_context, + "updatedInput": { + "command": f"echo '[hook-noop] original command was intercepted and executed by skill_script_hook'" + }, + } + } + print(json.dumps(output, ensure_ascii=False), flush=True) + + +def read_stdin_command(): + try: + raw = sys.stdin.read() + except Exception: + return None, None + if not raw or not raw.strip(): + return None, None + try: + data = json.loads(raw) + except json.JSONDecodeError: + return None, None + if not isinstance(data, dict): + return None, None + tool_name = data.get("tool_name") + tool_input = data.get("tool_input") or {} + command = tool_input.get("command") if isinstance(tool_input, dict) else None + return tool_name, command + + +def main(): + tool_name, command = read_stdin_command() + + # 没有命令或不是 Bash —— allow，由 harness 处理 + if not command or tool_name != "Bash": + emit_pretooluse("allow") + return + + if should_intercept(command): + execute_intercepted(command) + else: + # 不拦截 —— 让 harness 正常处理（不在此重复执行） + emit_pretooluse("allow") + + +if __name__ == "__main__": + main() diff --git a/settings.json b/settings.json new file mode 100644 index 00000000..1d5d3c44 --- /dev/null +++ b/settings.json @@ -0,0 +1,16 @@ +{ + "hooks": { + "PreToolUse": [ + { + "matcher": "Bash", + "hooks": [ + { + "type": "command", + "command": "PROJECT_ROOT=/home/ascendC python3 /home/ascendC/.claude/hooks/skill_script_hook.py", + "timeout": 1830 + } + ] + } + ] + } +} diff --git a/skills/triton/kernel-verifier/scripts/validate_triton_impl.py b/skills/triton/kernel-verifier/scripts/validate_triton_impl.py index d24180ad..27123fc8 100644 --- a/skills/triton/kernel-verifier/scripts/validate_triton_impl.py +++ b/skills/triton/kernel-verifier/scripts/validate_triton_impl.py @@ -137,6 +137,43 @@ def _resolve_call_name(node): return None +def _resolve_getattr_call(node): + """从 getattr(obj, "attr_name") 或 builtins.getattr(...) 调用中还原 (qualifier, attr)。 + + 仅处理第二个参数为字符串字面量的情况。 + 返回 (qualifier, attr) 或 None。 + """ + if not isinstance(node, ast.Call): + return None + func = node.func + is_getattr = False + if isinstance(func, ast.Name) and func.id == "getattr": + is_getattr = True + elif isinstance(func, ast.Attribute) and func.attr == "getattr": + if isinstance(func.value, ast.Name) and func.value.id == "builtins": + is_getattr = True + if not is_getattr: + return None + args = node.args + if len(args) < 2: + return None + obj_node = args[0] + attr_node = args[1] + if not isinstance(attr_node, ast.Constant) or not isinstance(attr_node.value, str): + return None + attr_name = attr_node.value + if isinstance(obj_node, ast.Name): + return (obj_node.id, attr_name) + if isinstance(obj_node, ast.Attribute): + if isinstance(obj_node.value, ast.Name): + return (f"{obj_node.value.id}.{obj_node.attr}", attr_name) + if isinstance(obj_node.value, ast.Attribute): + inner = obj_node.value + if isinstance(inner.value, ast.Name): + return (f"{inner.value.id}.{inner.attr}.{obj_node.attr}", attr_name) + return None + + def _get_subscript_value_name(node): """从 kernel[grid](...) 的 Subscript 节点提取 kernel 名称。""" if isinstance(node, ast.Subscript): @@ -241,6 +278,43 @@ def _count_kernel_launches_in_forward(forward_node): return count +def _check_single_call_violation(line, qual, attr): + """对单个 (qualifier, attr) 调用判断是否违规，返回 violation dict 或 None。""" + if qual == "torch": + if attr not in ALLOWED_TORCH_FUNCS: + return { + "line": line, + "call": f"torch.{attr}", + "reason": f"torch.{attr} 是计算操作，必须在 Triton kernel 中实现", + } + return None + if qual in ("F", "functional", "torch.nn.functional", "nn.functional"): + return { + "line": line, + "call": f"{qual}.{attr}", + "reason": f"{qual}.{attr} 是 PyTorch 计算操作，必须在 Triton kernel 中实现", + } + if qual == "triton" and attr in ALLOWED_TRITON_ATTRS: + return None + if attr in FORBIDDEN_TENSOR_METHODS: + if qual not in ("torch", "F", "triton", "functional", "torch.nn.functional", "nn.functional"): + return { + "line": line, + "call": f"{qual}.{attr}()" if qual else f"{attr}()", + "reason": f"{attr} 是计算操作，必须在 Triton kernel 中实现", + } + return None + if qual == "self": + if attr not in ("forward",): + return { + "line": line, + "call": f"self.{attr}(...)", + "reason": f"self.{attr}() 疑似 nn.Module 前向调用，核心计算必须在 Triton kernel 中实现", + } + return None + return None + + def check_forbidden_torch_ops(forward_node): """检查 forward 中是否使用了禁止的 torch 计算操作或 Python 控制流。 @@ -306,51 +380,31 @@ def check_forbidden_torch_ops(forward_node): qual, attr = resolved - # --- torch.xxx(...) --- - if qual == "torch": - if attr not in ALLOWED_TORCH_FUNCS: - violations.append({ - "line": node.lineno, - "call": f"torch.{attr}", - "reason": f"torch.{attr} 是计算操作，必须在 Triton kernel 中实现", - }) - continue - - # --- F.xxx(...) / functional.xxx(...) --- - if qual in ("F", "functional", "torch.nn.functional", "nn.functional"): - violations.append({ - "line": node.lineno, - "call": f"{qual}.{attr}", - "reason": f"{qual}.{attr} 是 PyTorch 计算操作，必须在 Triton kernel 中实现", - }) - continue - - # --- triton.cdiv 等 —— 允许 --- - if qual == "triton" and attr in ALLOWED_TRITON_ATTRS: - continue - - # --- tensor 方法计算操作 --- - if attr in FORBIDDEN_TENSOR_METHODS: - # 排除已知安全的 qual（torch/F/triton 已在上面处理） - if qual not in ("torch", "F", "triton", "functional", "torch.nn.functional", "nn.functional"): + # --- getattr(obj, "attr") 反射绕过检测 --- + if attr == "getattr" and qual in (None, "builtins"): + getattr_resolved = _resolve_getattr_call(node) + if getattr_resolved is not None: + gqual, gattr = getattr_resolved + getattr_violation = _check_single_call_violation( + node.lineno, gqual, gattr + ) + if getattr_violation is not None: + getattr_violation["call"] = ( + f"getattr({gqual}, '{gattr}')" + ) + violations.append(getattr_violation) + else: violations.append({ "line": node.lineno, - "call": f"{qual}.{attr}()" if qual else f"{attr}()", - "reason": f"{attr} 是计算操作，必须在 Triton kernel 中实现", + "call": "getattr(...)", + "reason": "getattr() 动态属性访问疑似绕过检测，forward() 中禁止使用 getattr 调用 PyTorch 接口", }) continue - # --- self.layer_name(x) —— 禁止 nn.Module 调用 --- - if qual == "self": - # 允许 self.forward() 递归，以及属性访问不在这里（ast.Attribute 不是 Call） - # self.xxx(...) 形式视为 nn.Module 前向调用 - if attr not in ("forward",): - violations.append({ - "line": node.lineno, - "call": f"self.{attr}(...)", - "reason": f"self.{attr}() 疑似 nn.Module 前向调用，核心计算必须在 Triton kernel 中实现", - }) - continue + violation = _check_single_call_violation(node.lineno, qual, attr) + if violation is not None: + violations.append(violation) + continue # --- 规则 B: 如果 forward() 中 kernel 启动次数 > 1，视为 Type3 退化 --- if kernel_launch_count > 1: diff --git a/skills/triton/latency-optimizer/.DS_Store b/skills/triton/latency-optimizer/.DS_Store new file mode 100644 index 00000000..40660e26 Binary files /dev/null and b/skills/triton/latency-optimizer/.DS_Store differ diff --git a/skills/triton/latency-optimizer/SKILL.md b/skills/triton/latency-optimizer/SKILL.md index 8a35630d..52b76dac 100644 --- a/skills/triton/latency-optimizer/SKILL.md +++ b/skills/triton/latency-optimizer/SKILL.md @@ -160,7 +160,7 @@ for n in range(N): sum_val += val # 标量加法 # 特征 3：标量控制流 -if x > 0: # 标量条件，导致 warp divergence +if x > 0: # 标量条件，导致 SIMD 分支分化 result = tl.exp(x) else: result = tl.cos(x) diff --git a/skills/triton/latency-optimizer/references/autotune.md b/skills/triton/latency-optimizer/references/autotune.md index 25b3d7f6..8b61cd37 100644 --- a/skills/triton/latency-optimizer/references/autotune.md +++ b/skills/triton/latency-optimizer/references/autotune.md @@ -4,7 +4,7 @@ Triton autotune 用于自动选择最优的 kernel 配置参数，主要包括影响分核（split）和切块（tiling）大小的参数，主要使用方式如下： -**三种种使用方式：** +**三种使用方式：** | 方式 | 说明 | 适用场景 | |------|------|---------| @@ -114,8 +114,8 @@ triton.Config( | `kwargs` | ✅ | ✅ | 完全支持 | | `num_warps` | ✅ | ❌ | NPU 架构差异，不支持 | | `num_stages` | ✅ | ❌ | NPU 架构差异，不支持 | -| `multibuffer` | ❌ | ✅ | NPU 特有，多缓冲优化 | -| `unit_flag` | ❌ | ✅ | NPU 特有，独立计算单元 | +| `multibuffer` | ❌ | ✅ | NPU 特有，多缓冲优化（通过 kwargs 传入） | +| `unit_flag` | ❌ | ✅ | NPU 特有，独立计算单元（通过 kwargs 传入） | #### 使用示例 @@ -227,13 +227,19 @@ pid = tl.program_id(0) offs_m = pid * BLOCK_M + tl.arange(0, BLOCK_M)# 可以知道BLOCK_M 是 split 参数 mask_m = offs_m < n_rows -# 二维切分 -pid_m = tl.program_id(0) -pid_n = tl.program_id(1) -offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)[:, None]# 可以知道BLOCK_M 是 split 参数 -offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)[None, :]# 可以知道BLOCK_N 是 split 参数 -mask_m = offs_m < n_rows -mask_n = offs_n < n_cols +# 二维切分（Ascend NPU 必须使用一维 Grid，通过一维 pid 映射到二维坐标） +pid = tl.program_id(0) +NUM_BLOCKS_M = tl.cdiv(M, BLOCK_M) +NUM_BLOCKS_N = tl.cdiv(N, BLOCK_N) +NUM_BLOCKS = NUM_BLOCKS_M * NUM_BLOCKS_N +for block_idx in range(pid, NUM_BLOCKS, CORE_NUM): # 交错循环，每个核处理多个块 + pid_m = block_idx // NUM_BLOCKS_N # 行块索引 + pid_n = block_idx % NUM_BLOCKS_N # 列块索引 + # BLOCK_M 和 BLOCK_N 都是 split 参数 + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)[:, None] + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)[None, :] + mask_m = offs_m < M + mask_n = offs_n < N ``` #### step1.2.识别分块参数 @@ -365,7 +371,7 @@ hints 参数说明: * 通过 `hints` 来显示指定轴关系时，`split_params`、`tiling_params`、`low_dim_axes`、`reduction_axes` 必须传入，即使某些参数为空 * 合法的轴名称是 `x/y/z/w/v/t`，仅仅用做关系映射 * `split_params` 和 `tiling_params` 为自动生成 tiling 算法必须的输入，`low_dim_axes` 和 `reduction_axes` 为 tiling 算法的可选输入，用于优化 tiling 效果，留空时 tiling 也能够自动生成，但可能会影响生成的候选 tiling 数量和质量 - * 当用户传入的 configs 不为空时，`auto_gen_config` 默认值为 `False`，如果希望此时也希望自动生成 tiling 配置并与用户传入的 configs 合并，需要显式在 `hints` 中传如入 `"auto_gen_config": True` + * 当用户传入的 configs 不为空时，`auto_gen_config` 默认值为 `False`，如果希望此时也希望自动生成 tiling 配置并与用户传入的 configs 合并，需要显式在 `hints` 中传入 `"auto_gen_config": True` 使用示例： ```python @@ -779,4 +785,4 @@ export TRITON_PRINT_AUTOTUNING=1 **限制：** 进阶用法仅支持 Vector 类算子，不支持 Cube 类算子。 -**优先级：** 自定义 autotune > 半自动 autotune (hints) > 自定义 autotune +**优先级：** 自动 autotune > 半自动 autotune (hints) > 自定义 autotune diff --git a/skills/triton/latency-optimizer/references/avoid_scalar_lowering.md b/skills/triton/latency-optimizer/references/avoid_scalar_lowering.md index 01ccbf53..01c0309c 100644 --- a/skills/triton/latency-optimizer/references/avoid_scalar_lowering.md +++ b/skills/triton/latency-optimizer/references/avoid_scalar_lowering.md @@ -44,12 +44,12 @@ --- -3. 扩展乘法（vmulext） +4. 扩展乘法（vmulext） 始终降级。该操作在 IR 层面只支持 i32，而 i32 触发降级，所以实际上没有向量硬件支持。 --- -4. 累积操作（cumsum / cumprod） +5. 累积操作（cumsum / cumprod） ``` ┌───────────────────────────────────┬────────────────────┬──────────────────────┬──────────────┐ │ 数据类型 │ 累积维度是最后维度 │ 累积维度不是最后维度 │ 多个累积维度 │ @@ -65,7 +65,7 @@ - 多个累积维度不会触发降级（但可能有其他限制）。 --- -5. 归约操作（reduce） +6. 归约操作（reduce） 归约的降级条件取决于芯片架构： diff --git a/skills/triton/latency-optimizer/references/checklist.md b/skills/triton/latency-optimizer/references/checklist.md index 4f8b1e98..c027cbda 100644 --- a/skills/triton/latency-optimizer/references/checklist.md +++ b/skills/triton/latency-optimizer/references/checklist.md @@ -15,7 +15,7 @@ - [ ] 对于除法操作，在不影响精度的情况下，必须使用 fp32 或 int32 数据类型进行计算 ### 4. 模运算规范 -- [ ] 禁止直接使用 `a % b` 操作，必须使用 `a - (a / b) * b` 操作替代 +- [ ] 禁止直接使用 `a % b` 操作，必须使用 `a - (a // b) * b` 操作替代 ### 5. Grid 并行度规范 - [ ] grid 并行数量禁止超过物理核数： diff --git a/skills/triton/latency-optimizer/references/discrete_memory_access.md b/skills/triton/latency-optimizer/references/discrete_memory_access.md index ea309e5c..486d09f8 100644 --- a/skills/triton/latency-optimizer/references/discrete_memory_access.md +++ b/skills/triton/latency-optimizer/references/discrete_memory_access.md @@ -2,7 +2,7 @@ ## 概述 -在 Triton NPU kernel 中，当线程通过非连续或不可预测的索引向量访问全局内存时，会导致访存效率低下，显著降低带宽利用率。先将整块数据读取到share memory，再取非连续或不可预测的索引可以显著提升计算效率。 +在 Triton NPU kernel 中，当线程通过非连续或不可预测的索引向量访问全局内存时，会导致访存效率低下，显著降低带宽利用率。先将整块数据读取到UB（Unified Buffer），再取非连续或不可预测的索引可以显著提升计算效率。 ## 触发条件 @@ -160,14 +160,14 @@ val = tl.load(x_ptr + offset + idx * stride_x, mask=mask) # 直接从global中 offset = tl.load(offset_ptr) # offset是一个完全无法预测的随机标量 idx = tl.load(idx_ptr + rn * stride_idx) # idx是一个完全无法预测的随机值向量 rm = tl.arange(0, M) # rm包含了所有的值，M为x张量的总长度 -x_shared = tl.load(x_ptr + offset_ptr + rm * stride_x) # 将x对应偏移的所有数据从global搬至share -val = tl.gather(x_shared.to(tl.float16), idx, 0).to(tl.int32) # 再从share中select目标值，注意数据类型的切换 +x_shared = tl.load(x_ptr + offset + rm * stride_x) # 将x对应偏移的所有数据从global搬至UB +val = tl.gather(x_shared.to(tl.float16), idx, 0).to(tl.int32) # 再从UB中select目标值，注意数据类型的切换 ``` ### 关键点 1. **识别无法预测的随机值**：溯源`tl.load`的输入索引计算过程，找到是否有无法预测的随机值，例如被`tl.load`读进来的值 -2. **自动优化**：将 `tl.load`的输入指针中的随机值剔除，改为读取一大块内存（注意不能超过share memory限制），然后使用`tl.gather`输入随机值，得到最终需要取的值 +2. **自动优化**：将 `tl.load`的输入指针中的随机值剔除，改为读取一大块内存（注意不能超过UB限制），然后使用`tl.gather`输入随机值，得到最终需要取的值 3. **注意gather的数据类型**：`tl.gather`不支持`int类型`，最好将输入强转成`tl.float16`再执行`tl.gather`,最后再转回原有的数据类型。如果提示精度报错，可以尝试强转成`tl.float32`。 ### 模式 2：循环内通过随机索引访问小查找表 diff --git a/skills/triton/latency-optimizer/references/load-order.md b/skills/triton/latency-optimizer/references/load-order.md index 9048c7e8..1d064ec8 100644 --- a/skills/triton/latency-optimizer/references/load-order.md +++ b/skills/triton/latency-optimizer/references/load-order.md @@ -70,8 +70,6 @@ def AB_load_kernel( tl.store(p_O, b_O) # store B - idx_B = tl.load(p_B_index) - p_B = B + idx_B tl.store(p_B, b_B) ``` diff --git a/skills/triton/latency-optimizer/references/pass-merge.md b/skills/triton/latency-optimizer/references/pass-merge.md index 1d6864a2..af31fa82 100644 --- a/skills/triton/latency-optimizer/references/pass-merge.md +++ b/skills/triton/latency-optimizer/references/pass-merge.md @@ -140,7 +140,6 @@ for col_offset in range(0, n_cols, BLOCK_SIZE): | 1 | ~1 ms | 无循环开销 | | 10 | ~10 ms | 线性增长 | | 100 | ~100 ms | 线性增长 | -| 512 | ~700 ms | **20x 慢于无循环版本** | **原因分析**: 1. **循环展开有限**: 编译器不会激进展开所有循环 diff --git a/skills/triton/latency-optimizer/references/scalar_to_vector.md b/skills/triton/latency-optimizer/references/scalar_to_vector.md index d182d3a9..e8d3c4d5 100644 --- a/skills/triton/latency-optimizer/references/scalar_to_vector.md +++ b/skills/triton/latency-optimizer/references/scalar_to_vector.md @@ -116,9 +116,8 @@ def count_2d(topk_ids_ptr, expert_num_tokens_ptr, num_experts: tl.constexpr, **原始代码（scalar 控制流）** ```python -# 莫格 x = tl.load(x_ptr + offsets, mask=mask) -if x > 0: # 标量条件，导致 warp divergence +if x > 0: # 标量条件，导致 SIMD 分支分化 result = tl.exp(x) else: result = tl.cos(x) @@ -130,17 +129,16 @@ else: x = tl.load(x_ptr + offsets, mask=mask) cond_mask = x > 0 # vector 比较，返回布尔向量 exp_result = tl.exp(x) -log_result = tl.cos(x) -result = cond_mask * exp_result + ~cond_mask * log_result # vector 加法，无分支 +cos_result = tl.cos(x) +result = cond_mask * exp_result + ~cond_mask * cos_result # vector 加法，无分支 ``` **场景二：两个分支的计算逻辑定义域不一样，至少有一个会出现nan或inf等无效输出** **原始代码（scalar 控制流）** ```python -# 莫格 x = tl.load(x_ptr + offsets, mask=mask) -if x > 0: # 标量条件，导致 warp divergence +if x > 0: # 标量条件，导致 SIMD 分支分化 result = tl.exp(x) else: result = tl.log(x) @@ -268,7 +266,7 @@ d = a - (a // b) * b # 公式转换 - **标量广播优化**：10-20% 加速，通过消除标量指令开销 - **标量规约优化**：5-128 倍加速（取决于 vector 并行度，FP16 理论加速比 128 倍） -- **控制流优化**：消除 warp divergence，提升 SIMD 利用率至 90%+ +- **控制流优化**：消除 SIMD 分支分化，提升 Vector 利用率至 90%+ - **整体 kernel 优化**：在 LayerNorm、Softmax 等带宽受限算子中，端到端性能提升 2-3 倍 **实测数据参考**：在 LayerNorm 算子中，将 mean/variance 计算从标量累加改为 vector 分块规约，UB 利用率从 35% 提升至 78%，kernel 执行时间减少 62%。 diff --git a/skills/triton/latency-optimizer/references/vector_core_partition.md b/skills/triton/latency-optimizer/references/vector_core_partition.md index bbbe641f..41c3eb7c 100644 --- a/skills/triton/latency-optimizer/references/vector_core_partition.md +++ b/skills/triton/latency-optimizer/references/vector_core_partition.md @@ -366,7 +366,6 @@ def softmax_kernel( | 策略 | 适用场景 | Grid 大小 | |------|---------|---------| | **一维分核** | 单维度处理（如逐行） | min(N / BLOCK, num_cores) | -| **二维分核** | 矩阵运算（如 matmul） | (M / BM, N / BN) | | **多行并行** | 行级 reduce（如 softmax） | min(M / ROWS_PER_BLOCK, num_cores) | ### 选择依据