Spaces:
Running on Zero
Running on Zero
File size: 7,447 Bytes
5f7afef | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | {
"slug": "attention-head-split-merge",
"title": "手撕多头张量拆分与重组 (Multi-Head Tensor Split and Merge)",
"difficulty": "Hard",
"category": "Attention",
"tags": [
"张量操作",
"多头注意力",
"维度变换"
],
"description_md": "### 题目描述\n在注意力机制中，频繁进行 `[B, S, D] -> [B, H, S, d]` 拆分与逆向合并。\n\n### 函数签名\n```python\nclass HeadReshaper:\n @staticmethod\n def split_heads(x: torch.Tensor, num_heads: int) -> torch.Tensor:\n ...\n\n @staticmethod\n def merge_heads(x: torch.Tensor) -> torch.Tensor:\n ...\n```\n\n### 示例\n```python\n# 输入\nx = [[[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]]]\nnum_heads = 2\n\n# 输出 (Tensor, shape=[1, 2, 2, 2], dtype=float32)\n[[[[1.0, 2.0], [5.0, 6.0]], [[3.0, 4.0], [7.0, 8.0]]]]\n```\n\n### 参数规格与约束\n- **`x`**：`torch.Tensor`，测试规格 `shape=[2, 4, 128], dtype=float32`\n- **`num_heads`**：`int`\n- **返回值**：`torch.Tensor`，输出规格 `shape=[2, 8, 4, 16], dtype=float32`\n- **精度容差**：`atol=0.0001, rtol=0.0001`",
"entry_point": "HeadReshaper",
"starter_code": "import torch\n\nclass HeadReshaper:\n @staticmethod\n def split_heads(x: torch.Tensor, num_heads: int) -> torch.Tensor:\n # 输入 x: [B, S, D]，按 num_heads 拆分并重排为 [B, H, S, d]\n pass\n\n @staticmethod\n def merge_heads(x: torch.Tensor) -> torch.Tensor:\n # 输入 x: [B, H, S, d]，逆向合并重排为 [B, S, H * d]\n pass\n",
"reference_code": "import torch\n\nclass HeadReshaper:\n @staticmethod\n def split_heads(x: torch.Tensor, num_heads: int) -> torch.Tensor:\n B, S, D = x.shape\n d = D // num_heads\n return x.view(B, S, num_heads, d).transpose(1, 2)\n\n @staticmethod\n def merge_heads(x: torch.Tensor) -> torch.Tensor:\n B, H, S, d = x.shape\n return x.transpose(1, 2).contiguous().view(B, S, H * d)\n",
"forbidden_modules": [
"os",
"sys",
"subprocess",
"shutil",
"socket",
"requests"
],
"forbidden_calls": [],
"time_limit": 3.0,
"memory_limit_mb": 512,
"atol": 0.0001,
"rtol": 0.0001,
"test_code": "import torch\ncases = [\n (\"标准多头张量拆合\", 2, 4, 8, 16),\n (\"单 Token 解码步拆合 (B=1, S=1, H=4, d=32)\", 1, 1, 4, 32),\n (\"奇数序列长与奇数批次 (B=3, S=7, H=8, d=16)\", 3, 7, 8, 16),\n]\nfor idx, (name, b, s, h, d) in enumerate(cases, 1):\n x = torch.randn(b, s, h * d)\n t0 = time.time()\n try:\n user_split = user_obj.split_heads(x, h)\n user_merge = user_obj.merge_heads(user_split)\n match, err, max_diff = compare_outputs(user_merge, x, atol=atol, rtol=rtol)\n passed, msg = match, err\n except Exception as e:\n passed, msg, max_diff, user_merge = False, str(e), None, None\n cases_results.append({\n \"case_id\": idx, \"case_name\": name, \"passed\": passed,\n \"time_ms\": round((time.time() - t0) * 1000, 2), \"message\": msg,\n \"user_output_summary\": summarize_obj(user_merge) if user_merge is not None else None,\n \"expected_summary\": summarize_obj(x), \"max_diff\": max_diff,\n })\n",
"is_featured": false,
"editorial_md": "#### 💡 核心算法原理与数学推导\n**手撕多头张量拆分与重组 (Multi-Head Tensor Split and Merge)**（`HeadReshaper`）聚焦于 **Attention** 中的关键计算逻辑：\n\n在注意力机制中，频繁进行 `[B, S, D] -> [B, H, S, d]` 拆分与逆向合并。\n\n在实现中，函数接收参数 `x` (`torch.Tensor`)、`num_heads` (`int`)，目标是计算出类型与形状严格符合 `torch.Tensor` 契约的结果。\n\n#### 📐 输入输出张量变换与步骤拆解\n- **输入基准规格**：`x` 规格为 `torch.Tensor(shape=[2, 4, 128], dtype=float32)`，`num_heads` 规格为 `int` (示例值: `8`)。\n- **期望输出规格**：`torch.Tensor(shape=[2, 8, 4, 16], dtype=float32)`。\n\n**具体推导与计算步骤如下：**\n1. 计算中间量 `B, H, S, d`：通过表达式 `x.shape` 完成张量或数值变换。\n2. 返回最终结果：输出 `x.transpose(1, 2).contiguous().view(B, S, H * d)`，严格对齐函数签名要求的返回类型与维度。\n\n#### 📊 复杂度分析\n- **时间复杂度**：$O(N \\cdot M \\cdot D)$，其中 $N, M$ 为矩阵行/列样本规模、$D$ 为内积特征维度，由矩阵乘法算子主导。\n- **空间复杂度**：$O(N \\cdot M)$，用于存储输出矩阵及中间内积张量。\n\n#### ⚠️ 关键实现细节与易错点\n- **除零保护与数值平滑 (`eps`)**：在执行分母归一化、标准差缩放或对数运算时，必须按公式位置准确加上 `eps`（例如在平方根内 `sqrt(var + eps)` 还是分母外 `+ eps`），防止全零或常量输入触发 `NaN` / `Inf`。\n- **内存连续性与维度重排顺序**：在执行 `transpose` 或 `permute` 交换多头/序列维度后，张量在底层内存中变为非连续布局；若后续调用 `.view(...)` 合并维度，必须先调用 `.contiguous()` 或改用 `.reshape(...)`。\n- **整数整除与返回类型契约**：涉及块索引、分片大小或显存字节换算时，需严格区分 Python 浮点除法 `/` 与整数整除 `//`，确保返回的索引、步数或字节数类型与函数签名一致。\n\n#### 💻 官方标准参考实现\n```python\nimport torch\n\nclass HeadReshaper:\n @staticmethod\n def split_heads(x: torch.Tensor, num_heads: int) -> torch.Tensor:\n B, S, D = x.shape\n d = D // num_heads\n return x.view(B, S, num_heads, d).transpose(1, 2)\n\n @staticmethod\n def merge_heads(x: torch.Tensor) -> torch.Tensor:\n B, H, S, d = x.shape\n return x.transpose(1, 2).contiguous().view(B, S, H * d)\n```",
"difficulty_breakdown": {
"concept": 4,
"implementation": 4,
"math_derivation": 4,
"boundary_complexity": 4
},
"test_cases": [
{
"name": "标准多头张量拆合",
"input_desc": "真实参数契约用例 (x, num_heads)",
"params": [
{
"name": "x",
"type": "torch.Tensor",
"value": "torch.randn(2, 4, 128)"
},
{
"name": "num_heads",
"type": "int",
"value": "8"
}
],
"raw_inputs": {
"x": "torch.randn(2, 4, 128)",
"num_heads": "8"
}
},
{
"name": "单 Token 解码步拆合 (B=1, S=1, H=4, d=32)",
"input_desc": "真实参数契约用例 (x, num_heads)",
"params": [
{
"name": "x",
"type": "torch.Tensor",
"value": "torch.randn(1, 1, 128)"
},
{
"name": "num_heads",
"type": "int",
"value": "4"
}
],
"raw_inputs": {
"x": "torch.randn(1, 1, 128)",
"num_heads": "4"
}
},
{
"name": "奇数序列长与奇数批次 (B=3, S=7, H=8, d=16)",
"input_desc": "真实参数契约用例 (x, num_heads)",
"params": [
{
"name": "x",
"type": "torch.Tensor",
"value": "torch.randn(3, 7, 128)"
},
{
"name": "num_heads",
"type": "int",
"value": "8"
}
],
"raw_inputs": {
"x": "torch.randn(3, 7, 128)",
"num_heads": "8"
}
}
]
} |