执行摘要

• 一句话：修复DFlash前缀缓存因缺lookahead块的损坏
• 推荐动作：建议精读此PR及关联PR #43733，理解DFlash与EAGLE在KV写入时序上的根本差异，以及为何需要调整lookahead分配策略。设计上将条件抽取为独立方法并区分bonus token的做法值得借鉴。对于维护者，建议在合并后运行DFlash的端到端测试（如test_dflash.py）验证无回归。

功能与动机

修复DFlash在前缀缓存高并发场景下持久性MGL降解，原因是新请求的首次预填充缺少lookahead块分配，导致drafter将KV写入共享前缀块，污染其他请求的缓存。

实现拆解

1. 重构_input_fits_in_drafter方法：在GPU Model Runner中将原本内联的drafter窗口检查提取为独立方法，并针对DFlash的特性，在其drafter查询窗口大小中增加1个bonus token。
2. 替换调用点：在sample_tokens方法中用新方法self._input_fits_in_drafter(...)替换原有内联条件。
3. 新增单元测试文件tests/v1/spec_decode/test_dflash_lookahead.py，包含三个测试用例：
   • test_dflash_prefill_reserves_lookahead_blocks：验证调度器在DFlash模式下num_lookahead_tokens为num_spec_tokens + 1，且首次调度后分配的块包括一个lookahead块。
   • test_dflash_first_prefill_query_window_fits_allocated_blocks：验证首次预填充后，drafter查询位置落在已分配的块范围内。
   • test_dflash_drafter_window_reserves_bonus_token：验证_input_fits_in_drafter对DFlash和EAGLE的区分。
4. 测试通过SpeculativeConfig(method='dflash')创建真实调度器，而非直接设置属性。

关键文件：

• vllm/v1/worker/gpu_model_runner.py（模块 运行器；类别 source；类型 core-logic；符号 _input_fits_in_drafter）: 修复核心逻辑，提取_input_fits_in_drafter方法并修正DFlash的bonus token计数
• tests/v1/spec_decode/test_dflash_lookahead.py（模块 测试套件；类别 test；类型 test-coverage；符号 _dflash_speculative_config, _create_dflash_scheduler, test_dflash_prefill_reserves_lookahead_blocks, test_dflash_first_prefill_query_window_fits_allocated_blocks）: 新增全面测试，验证DFlash在预填充时的块分配和drafter窗口限制

关键符号：_input_fits_in_drafter, test_dflash_prefill_reserves_lookahead_blocks, test_dflash_first_prefill_query_window_fits_allocated_blocks, test_dflash_drafter_window_reserves_bonus_token

关键源码片段

tests/v1/spec_decode/test_dflash_lookahead.py

新增全面测试，验证DFlash在预填充时的块分配和drafter窗口限制

def test_dflash_prefill_reserves_lookahead_blocks():
    # 创建一个 DFlash 调度器（num_spec_tokens=3）
    scheduler = _create_dflash_scheduler(NUM_SPECULATIVE_TOKENS)
​
    # 断言 num_lookahead_tokens == num_spec_tokens + 1（即 4）
    assert scheduler.num_lookahead_tokens == NUM_SPECULATIVE_TOKENS + 1
​
    # 添加一个请求，具有完整块大小的令牌数（BLOCK_SIZE=16）
    (request,) = create_requests(num_requests=1, num_tokens=BLOCK_SIZE, block_size=BLOCK_SIZE)
    scheduler.add_request(request)
    output = scheduler.schedule()
​
    # 确认预填充调度了所有令牌
    assert output.num_scheduled_tokens[request.request_id] == BLOCK_SIZE
    # 确认分配了 2 个块：一个用于预填充，一个用于 lookahead
    assert len(output.scheduled_new_reqs[0].block_ids[0]) == 2
​
​
def test_dflash_first_prefill_query_window_fits_allocated_blocks():
    # 验证首次预填充后，drafter 将要查询的所有位置都在已分配的块内
    scheduler = _create_dflash_scheduler(NUM_SPECULATIVE_TOKENS)
    (request,) = create_requests(num_requests=1, num_tokens=BLOCK_SIZE, block_size=BLOCK_SIZE)
    scheduler.add_request(request)
    output = scheduler.schedule()
    block_ids = output.scheduled_new_reqs[0].block_ids[0]
    # 计算 drafter 要查询的所有位置（BLOCK_SIZE 到 BLOCK_SIZE + num_lookahead_tokens - 1）
    query_positions = range(BLOCK_SIZE, BLOCK_SIZE + scheduler.num_lookahead_tokens)
    # 确保所有查询位置都在已分配的块内
    assert all(pos // BLOCK_SIZE < len(block_ids) for pos in query_positions)
​
​
def test_dflash_drafter_window_reserves_bonus_token():
    # 测试 _input_fits_in_drafter 对 DFlash 和 EAGLE 的区分
    input_fits_in_drafter = GPUModelRunner._input_fits_in_drafter
    # 创建一个模拟的 DFlash runner
    dflash_runner = SimpleNamespace(
        num_spec_tokens=NUM_SPECULATIVE_TOKENS, # 3
        effective_drafter_max_model_len=100,
        speculative_config=_dflash_speculative_config(NUM_SPECULATIVE_TOKENS),
    )
    # 对于 DFlash，窗口大小 = 3 + 1 = 4
    # 所以 max_seq_len = 96 时，96 + 4 = 100 刚好等于上限，应返回 True
    assert input_fits_in_drafter(dflash_runner, SimpleNamespace(max_seq_len=96))
    # max_seq_len = 97 时，97 + 4 = 101 超过上限，应返回 False
    assert not input_fits_in_drafter(dflash_runner, SimpleNamespace(max_seq_len=97))

评论区精华

• Eagle vs DFlash行为差异：mgoin指出Eagle的drafter在首次预填充后不写KV到未分配位置，而DFlash的注意力是双向的，会写KV到未来位置，因此需要额外lookahead块。benchislett最初认为两者行为一致，后确认差异。
• 测试方式：mgoin建议测试应通过SpeculativeConfig创建调度器而非直接设置use_dflash，以确保调度器内部分支正确。shreyas269采纳该建议重写测试。
• P/D disaggregation考虑：mgoin担心移除num_computed_tokens == 0的守卫会影响P/D disagg。讨论后决定将调度器修复移至#43733单独处理，本PR不再改动调度器。

• 方法签名：benchislett建议_input_fits_in_drafter直接接收spec_decode_common_attn_metadata作为参数，使调用更清晰。该建议被采纳。

  • Eagle vs DFlash 在首次预填充时写 KV 的行为差异 (design): 确认 DFlash 确实需要额外 lookahead 块，因为其注意力是双向的，会在未来位置写 KV。
  • 测试中使用直接设置 use_dflash 与 SpeculativeConfig (testing): shreyas269 修改测试，使用 SpeculativeConfig(method='dflash') 创建调度器。
  • P/D disaggregation 兼容性 (design): 调度器修复部分被移至 #43733 并单独处理，本 PR 不再修改调度器；P/D disagg 兼容性由 #43733 处理。
  • _input_fits_in_drafter 方法签名 (style): shreyas269 按照建议修改。

风险与影响

• 风险：变更主要影响DFlash模式下drafter窗口检查逻辑，可能的风险包括：
  • 回归风险：_input_fits_in_drafter方法同时影响EAGLE、DraftModel等speculative方法，需要确保这些方法的行为正确。当前逻辑对非DFlash方法保持原有窗口大小，风险较低。
  • 前缀缓存协作：修复依赖于调度器已正确分配lookahead块（已在#43733修复），如果调度器部分有未发现的边缘情况，可能导致修复不完整。
  • P/D disagg未覆盖：本PR未对P/D分离部署场景进行测试，该场景由#43733整体处理。
  • 测试覆盖有限：单元测试模拟了调度器行为，但未涉及实际模型前向，无法捕捉GPU端完整交互问题。
  • 影响：直接影响：修复DFlash在前缀缓存高并发场景下的MGL退化，提升DFlash用户的推理稳定性。影响范围限于启用DFlash且使用前缀缓存的高并发部署。对其他speculative方法无影响。团队应关注该变更与#43733的协作，确保整体正确性。
  • 风险标记：核心路径变更, prefix caching 依赖

关联脉络

• PR #43733 [Bugfix][DFlash]allocate the proper number of lookahead slots: 本 PR 修复了 GPU model runner 端的 drafter 窗口检查，调度器侧的 lookahead 块分配修复已由 #43733 单独合并，本 PR 依赖该修复。