Prhub

#35245 refactor(unified-memory): translate the KV write location once, at ForwardBatch construction

原始 PR 作者 caihuali95 合并时间 2026-08-31 14:52 文件变更 28 提交数 19 评论 4 代码增减 +1838 / -233

执行摘要

实现拆解

  1. 核心源码路径
    - python/sglang/srt/mem_cache/kv_index_translator.py(dependency-wiring):源码主路径;涉及符号 KVReadTables, KVIndexTable, sliding_window_read_ids, KVIndexTranslator;包含 导入关系调整、控制流调整、配置键调整;+358/-0;关键符号:KVReadTables, KVIndexTable, sliding_window_read_ids, KVIndexTranslator
    - python/sglang/srt/layers/attention/triton_backend.py(core-logic):源码主路径;涉及符号 _translate_cuda_graph_shared_pool_locs, _fill_cuda_graph_write_locs;包含 导入关系调整、控制流调整、配置键调整;+91/-142;关键符号:_translate_cuda_graph_shared_pool_locs, _fill_cuda_graph_write_locs
    - python/sglang/srt/mem_cache/memory_pool.py(dependency-wiring):源码主路径;包含 导入关系调整、控制流调整、配置键调整;+35/-26

  2. 测试与验证
    - test/registered/unit/model_executor/test_unified_out_cache_loc_rebind.py(test-coverage):测试配套;涉及符号 _make_fb, _armed_source, _call_names, TestForwardBatchWiring;包含 测试覆盖调整、导入关系调整、控制流调整;+181/-0;关键符号:_make_fb, _armed_source, _call_names, TestForwardBatchWiring
    - test/registered/unit/mem_cache/test_kv_index_translator.py(test-coverage):测试配套;涉及符号 _build_composite, _make_source, _reference_table, TestPassthrough;包含 测试覆盖调整、导入关系调整、控制流调整;+684/-0;关键符号:_build_composite, _make_source, _reference_table, TestPassthrough
    - test/registered/unit/server_args/test_unified_prefill_cuda_graph_gate.py(test-coverage):测试配套;涉及符号 _run_handler, TestUnifiedPrefillCudaGraphGate, test_default_prefill_capture_is_auto_disabled, test_explicit_prefill_backend_is_refused;包含 测试覆盖调整、导入关系调整、控制流调整;+92/-0;关键符号:_run_handler, TestUnifiedPrefillCudaGraphGate, test_default_prefill_capture_is_auto_disabled, test_explicit_prefill_backend_is_refused

  3. 配置与运行环境
    - python/sglang/kernels/ops/kvcache/kv_read_table.py(infrastructure):部署/基础设施;涉及符号 build_kv_read_table_kernel, build_kv_read_table;包含 导入关系调整、控制流调整、配置键调整;+140/-0;关键符号:build_kv_read_table_kernel, build_kv_read_table
    - python/sglang/kernels/ops/kvcache/kv_indices.py(infrastructure):部署/基础设施;包含 控制流调整、配置键调整;+34/-10

文件 模块 状态 重要度
python/sglang/srt/mem_cache/kv_index_translator.py sglang/srt added 9.25
python/sglang/srt/layers/attention/triton_backend.py sglang/srt modified 8.3
test/registered/unit/model_executor/test_unified_out_cache_loc_rebind.py unified/out/cache/loc/rebind added 8.27
test/registered/unit/mem_cache/test_kv_index_translator.py kv/index/translator added 8.14
test/registered/unit/server_args/test_unified_prefill_cuda_graph_gate.py unified/prefill/cuda/graph/gate added 7.68

关键符号

KVReadTables KVIndexTable sliding_window_read_ids KVIndexTranslator __init__ make_capture_tables _zeros build_index_table _translate_cuda_graph_shared_pool_locs _fill_cuda_graph_write_locs

分析完成后,这里会展示 LLM 生成的相对完整源码片段和详细注释。

评论区精华

没有提炼出高价值讨论线程

当前评论区没有形成足够清晰的争议点或结论,后续有更多讨论时会体现在这里。

风险与影响

关联 Issue

#34602 feat(unified-memory): dense KV views for uniform-row MHA/SWA models

完整报告

参与讨论