• score + mask 没有报错，为什么 decode 时 attention 结果整批都不对？
• 都是“把 KV head 扩展到 query head 数”，expand() 和 repeat() 为什么显存差了好几倍？
• cache[slots].copy_(values) 跑完了，为什么 KV cache 里什么都没写进去？

广播和索引代码通常只有一行，却可能隐藏三种完全不同的成本：只改变张量元数据、读取并复制数据、生成远大于输入的计算结果。它们还会悄悄改变语义——维度对错了不一定报错，写错了位置也不一定报错。在 LLM 推理里，attention mask、GQA 的 head 扩展、beam 重排和 KV cache 更新都会遇到这些问题。

布局和共享存储的基础见 Tensor Memory。本篇关注：表达式的 shape、读写语义和实际成本是否符合预期。

抽象：每个表达式回答三个问题

读一行广播或索引代码，可以把它拆成三个彼此独立的问题：

• 形状从哪来：输出的每一维由哪个输入的哪一维决定？广播只做数字对齐，语义要靠代码自己保证。
• 存储是谁的：结果是原张量的 view（共享存储、可能零 stride），还是新分配的副本？
• 数据往哪流：buffer[idx] 出现在等号右边是读出（gather 成新张量），出现在左边是写回（scatter 到原存储）。

三个问题分别对应三类错误：形状推错导致静默的语义错误或输出规模爆炸；存储判断错导致多余复制或非法原地写；读写方向判断错导致“写了个寂寞”。下面逐一展开。

广播：对齐的是维度，不是语义

广播从尾部对齐维度；每一对维度的大小必须相等，或者其中一侧为 1，缺失的前导维度按 1 处理。它不会判断某一维到底表示 batch、token 还是 head。规则见官方 Broadcasting semantics。

import torch

x = torch.zeros(4, 1)
y = torch.zeros(4)
z = x + y
assert z.shape == (4, 4)  # y 被视为 (1, 4)，不是 (4, 1)

这里两个输入共八个元素，输出却有十六个。广播本身无需先复制输入，但消费广播结果的算子仍要完成相应规模的计算；普通 eager 加法也会分配输出。编译融合可能省去中间结果的存储，不能据此忽略计算规模。

图中第三部分是推理里最典型的陷阱：score 为 [batch, head, query, key]，每个请求的 key mask 为 [batch, key]。直接相加时，mask 的 batch 维对到了 query 维；decode 时 query 长度恰好为 1，于是不报错，输出却被广播成 [B, H, B, K]。应显式写成 mask[:, None, None, :] 变成 [batch, 1, 1, key]，不要依赖“恰好相等”或“恰好为 1”的维度大小。

expand 与 repeat：逻辑复用和物理复制

x = torch.arange(3).reshape(1, 3)
a = x.expand(4, 3)
b = x.repeat(4, 1)
assert a.stride() == (0, 1)
assert b.stride() == (3, 1)

两者逻辑形状相同，区别全在 stride：a 的四行都读取同一行内存。对它执行需要逐元素独立写入的原地操作是不安全的，可能被 PyTorch 拒绝；应先 clone()。零 stride 和写入限制见 Tensor.expand，复制语义见 Tensor.repeat。

在 GQA 中，把 KV head 扩展到 query head 数量时，expand() 可以省去某一步复制，但后续 reshape() 是否物化、kernel 是否支持零 stride，仍需检查。不能仅凭代码出现 expand() 就认定整个路径零拷贝。支持 GQA 的 attention kernel 可以直接处理 head 分组，避免显式扩展。

索引读取与索引赋值需要分开理解

基础切片读取通常是 view，高级索引读取通常是 copy；但高级索引位于赋值左侧时，会写回原张量，不是先生成一个副本再修改副本。见 Tensor Views。

cache = torch.zeros(8, 2)
slots = torch.tensor([1, 4])
values = torch.ones(2, 2)

selected = cache[slots]   # 高级索引读取：副本
selected.fill_(7)        # cache 不会跟着变
cache[slots] = values    # 高级索引赋值：写回 cache

原因在于 Python 的求值方式：cache[slots] = values 会被翻译成一次 __setitem__，PyTorch 据此调用 index_put_ 直接修改原存储；而 cache[slots].copy_(values) 先执行 __getitem__ 得到 gather 出的临时副本，再对副本 copy_，cache 不变。读写方向由 buffer[idx] 在表达式中的位置决定，而不是由后面跟着什么方法决定。这正对应 KV cache 中“gather 得到的 K/V”和“向 slot 写入新 K/V”的区别。

如果写入下标重复，不能把并发写入想象成 Python 循环里稳定的“最后一次覆盖”。index_put_ 在 accumulate=False 时对重复下标的行为未定义；累加也需另行考虑浮点归约顺序。见 Tensor.index_put_。

动态索引还可能带来同步

整数下标张量的元素数量已由元数据给出；布尔 mask 的选中数量则取决于数据内容。因此不能把所有高级索引都理解成同样的固定成本操作。

CUDA 上的 torch.nonzero() 会引入 host-device 同步，布尔 mask 索引需要先知道选中数量，也会有类似效果。decode 热路径中的 mask.nonzero()，可能既生成动态长度结果，又让 CPU 等待 GPU，还会妨碍 CUDA Graph 捕获。见 torch.nonzero 和 CUDA 异步执行。

在 SGLang 中：KV cache 的一读一写

SGLang 的 KV cache 正是图中那个“按 slot 编号的大 buffer”，读写两侧分别落在不同模块（基于当前 python/sglang/srt/）：

• 写：索引赋值，必要时换成专用 kernel。 mem_cache/memory_pool.py 中 MHATokenToKVPool.set_kv_buffer() 从 loc_info 解出本步新 token 的 slot（即 forward_batch.out_cache_loc），交给 _set_kv_buffer_impl()。CUDA/HIP 上行宽满足条件时调用 store_cache kernel；否则回落到最朴素的 k_cache[indices] = k、v_cache[indices] = v——赋值左侧的高级索引，直接写原存储。CUDA Graph 捕获模式下还会把 V 的写入放到另一条 stream 上，与 K 的写入并行，再用 wait_stream 接回。同文件里还有一处细节：fp8 KV 以 torch.uint8 存储，注释说明原因是 index_put 不支持 float8_e5m2。
• 读：先切片得到下标，再 gather。 layers/attention/torch_native_backend.py 逐请求执行 per_req_tokens = req_to_token[req_pool_idx, start_kv:end_kv]——这是普通切片，拿到的是 slot 编号的 view；随后 k_cache[per_req_tokens] 是高级索引读取，复制出该请求的 K/V，再用 movedim 调整维度顺序（view，不复制）。这一读一写的复制成本，是这个后端只适合做对照与兜底的原因之一（见 Kernel 与编译执行）。
• GQA 不显式扩展。 同一后端用 use_gqa = layer.tp_q_head_num != layer.tp_k_head_num 判断是否分组，并把它作为 enable_gqa 传给 SDPA，由 attention 实现处理 head 分组，而不是先 expand/repeat KV head。

高性能后端（FlashInfer、FlashAttention 等）则把 kv_indices 或 page table 直接交给 kernel，在 kernel 内部按下标读取，避免先 gather 出完整 K/V；KV 池的组织方式见 Memory Pool 与 Paged Attention。

读代码时的检查顺序

1. 写出输入、输出每一维的语义，而不只看数字。
2. 判断广播后的结果规模；可以先用 torch.broadcast_shapes() 检查。
3. 区分 view、复制读取与写回原张量；看 buffer[idx] 在等号的哪一侧。
4. 检查零 stride、重复下标和动态输出长度。
5. 检查后续 kernel 是否要求物化，是否出现隐式同步。

小结

广播和索引的每一行代码都同时回答了三个问题：形状从哪来、存储是谁的、数据往哪流。广播减少输入复制，并不自动减少输出规模，也不检查维度语义；expand 省掉复制，却换来零 stride 和写入限制；同样的 buffer[idx]，在等号左边是写回原存储，在右边是 gather 出的副本，且不保证重复下标下的确定性。把 shape 推导和存储读写一起看，才能判断推理代码是否既正确又高效。