执行摘要

• 一句话：优化 TPU InputBatch swap_states 仅复制有效 token 区间
• 推荐动作：值得快速精读。该 PR 展示了固定长数组整行 swap 只需复制有效区间的经典优化，代码量极小但收益显著。值得关注的设计决策：用 num_tokens_no_spec 作为有效边界而非 max_model_len，以及测试对整行强断言与进一步裁剪优化之间的张力（作者推迟到后续 PR）。

功能与动机

PR body 指出 swap_states 原实现复制整个 max_model_len 大小的行来交换两个 batch 条目，这是浪费的，因为只有有效 token 范围需要保留。当 max_model_len 很大（如 131072）而有效 token 很少时，复制开销极大。代码中原有 # TODO(lucas): optimize this by only copying valid indices 注释，本 PR 正是落实这一项优化。

实现拆解

1. 变更入口：vllm/v1/worker/tpu_input_batch.py 的 InputBatch.swap_states(i1, i2) 方法，该方法在调度重排（preemption、prefix caching 调整等）时交换两个 batch 槽位的完整状态，是 v1 引擎的热路径。
2. 核心改造：先计算两个序列有效 token 数的最大值 n = max(self.num_tokens_no_spec[i1], self.num_tokens_no_spec[i2])，再对 token_ids_cpu 的 [:n] 区间做三段式交换（临时拷贝 + 两次赋值），跳过行尾填充数据。原实现直接对整行 [i, ...] 操作，等价于复制 3 个 max_model_len 宽度的行。
3. 正确性依据：交换后 n: 区间残留另一行旧数据，但 vLLM 对 token 的访问均以 num_tokens / num_tokens_no_spec 为界，行尾内容不会被运行时读取；condense() 等后续逻辑同样只处理有效区间。
4. 测试配套：未新增测试文件。作者运行了 tests/v1/attention/test_batch_reordering.py 与 tests/v1/worker/test_gpu_input_batch.py::test_swap_states_in_input_batch 并通过，性能对比数据在 PR body 中给出。变更不涉及配置、schema 或部署配套。

关键文件：

• vllm/v1/worker/tpu_input_batch.py（模块 输入批处理；类别 source；类型 core-logic；符号 swap_states）: 唯一的变更文件，InputBatch.swap_states() 中 token_ids_cpu 的行交换由整行复制改为仅复制两个序列有效 token 数的最大值前缀，是性能提升的核心所在。

关键符号：swap_states

关键源码片段

vllm/v1/worker/tpu_input_batch.py

唯一的变更文件，InputBatch.swap_states() 中 token_ids_cpu 的行交换由整行复制改为仅复制两个序列有效 token 数的最大值前缀，是性能提升的核心所在。

def swap_states(self, i1: int, i2: int) -> None:
    # …… 上方交换 num_prompt_tokens、采样参数等标量字段 ……
​
    # 不能直接写 `token_ids_cpu[i1], token_ids_cpu[i2] = token_ids_cpu[i2], token_ids_cpu[i1]`，
    # 因为张量切片赋值是引用语义，必须借助临时拷贝。
    # 原实现拷贝整行（宽度为 max_model_len），长上下文 + 短序列时浪费巨大；
    # 这里把交换区间裁剪为两个序列有效 token 数的最大值 n。
    n = max(self.num_tokens_no_spec[i1], self.num_tokens_no_spec[i2])
    tmp = self.token_ids_cpu[i1, :n].copy()
    self.token_ids_cpu[i1, :n] = self.token_ids_cpu[i2, :n]
    self.token_ids_cpu[i2, :n] = tmp
​
    swap_dict_values(self.generators, i1, i2)
    swap_dict_values(self.min_tokens, i1, i2)
    swap_dict_values(self.bad_words_token_ids, i1, i2)
​
    self.request_lora_mapping[i1], self.request_lora_mapping[i2] = (
        self.request_lora_mapping[i2],
        self.request_lora_mapping[i1],
    )
    self.logit_bias[i1], self.logit_bias[i2] = (
        self.logit_bias[i2],
        self.logit_bias[i1],
    )
​
    if self.allowed_token_ids_mask_cpu_tensor is not None:
        (
            self.allowed_token_ids_mask_cpu_tensor[i1],
            self.allowed_token_ids_mask_cpu_tensor[i2],
        ) = (
            self.allowed_token_ids_mask_cpu_tensor[i2],
            self.allowed_token_ids_mask_cpu_tensor[i1],
        )
    self.block_table.swap_row(i1, i2)

评论区精华

gemini-code-assist 在 GPU 和 TPU 两个文件的 swap 逻辑上都提出进一步优化建议：当两个序列长度差异大时，用短序列临时缓冲把复制量从 3 * max(len1, len2) 降到 max(len1, len2) + 2 * min(len1, len2)，避免复制较长序列行尾的垃圾数据。作者回应称该写法会让 test_swap_states_in_input_batch 失败——_compare_objs 对整行做 allclose 全量对比，但认可其逻辑，决定后续单独研究。LucasWilkinson 最终批准并留言 Thanks for the contribution! LGTM。

• 进一步裁剪垃圾数据复制（非对称长度场景） (performance): 未采纳，作者推迟到后续 PR；本 PR 维持按 max(num_tokens_no_spec) 裁剪的交换实现。
• GPU 与 TPU 两侧是否同时生效 (question): GPU 侧是否有等价改动或后续计划未在讨论中明确；作者在 GPU 评论上的回应同样表示后续单独研究。
• 整体审核与合并 (other): 已批准并合并。

风险与影响

• 风险：
  1. 正确性依赖调用约定：交换只保证有效前缀 [:n] 一致，n: 之后残留另一行旧数据。一旦有代码按整行语义读取 token_ids_cpu（如测试中的 _compare_objs），就会观察到差异；运行时路径目前均以 num_tokens 为边界，风险可控。
  2. 边界更新次序：n 取两行 num_tokens_no_spec 的最大值，若后续续写使某行有效长度超过交换时的 n，新写入的 token 会覆盖行尾旧数据，行为与原始整行交换不同，需要依赖 num_tokens_no_spec 同步增长。
  3. 投机解码假设：TPU 实现用 num_tokens_no_spec 作为边界；如果 TPU 路径引入 spec token 且其也存放在 token_ids_cpu 中，边界需要补上 spec_token_ids 长度（GPU 版本 review 中 gemini 已提示此点）。
  4. 测试缺口：本 PR 没有新增针对 TPU swap_states 的专项测试，目前覆盖主要来自 GPU 侧的 test_gpu_input_batch.py。
     - 影响：用户与系统侧：在长 max_model_len、短序列的 TPU 推理场景下，batch 重排（preemption、prefix caching 调整）开销显著下降，对 P/D 分离和调度频繁的场景尤其有利；PR 数据展示 2 到 46 倍的耗时降低。团队侧：改动极小、聚焦，易于 review 和合并，但缺少 TPU 专项回归测试；该模式对 GPU 侧同类逻辑也有参考价值，PR body 声称同时适用于 GPU 与 TPU。
     - 风险标记：仅 TPU 路径变更, 缺少 TPU 专项测试, 行尾残留数据依赖 num_tokens 边界, 投机解码场景边界需复核

关联脉络

• 暂无明显关联 PR