Skip to content

[Perf] dsv41: Read Engram rows from peer device shards - #46180

Merged
hnyls2002 merged 3 commits into
mainfrom
lsyin/dsv41-engram-peer-shards
Oct 24, 2026
Merged

hnyls2002 merged 3 commits into
mainfrom
lsyin/dsv41-engram-peer-shards

Conversation

@hnyls2002

@hnyls2002 hnyls2002 commented Oct 23, 2026 •

Copy link
Copy Markdown
Collaborator

With per-rank Engram tables and DP attention, every lookup all-gathers all DP ranks' hash ids, gathers the owned rows for all of them into a [DP x tokens] bf16 buffer, and reduce-scatters it back. The traffic and the buffer grow with the DP size, so wide-EP decode pays a collective per Engram layer that is larger than the lookup itself.

This adds an opt-in path, SGLANG_OPT_DSV41_ENGRAM_PEER_SHARDS=1:

  • Each TP rank keeps its row shard in exportable device memory (CUDA VMM) and maps every peer's shard after loading: FABRIC handles across an NVLink domain (MNNVL), POSIX fds within a host.
  • engram_gather_peer: one Triton kernel reads each of a rank's own tokens' rows straight from the owning shard (owner found by binary search over the shard starts). No id gather, no global buffer, no collective; CUDA-graph safe. Covers the TP, DP and prefill-CP lookups.
  • Output is bitwise equal to the current path (-0 is written as +0, as the zero-fill + sum leaves it).
  • Takes precedence over the host table (a warning says so when both are set). The shard costs table_size / tp_size of HBM per GPU, so it suits wide DEP.
  • The VMM handle exchange (export, then FABRIC handles or POSIX fds) moves into cuda_vmm_utils.exchange_shareable_handles, shared with the DWDP transport; a group spanning hosts without FABRIC handles now fails there with a clear error. Peers are mapped with the existing import_and_map_alloc.

Off by default; nothing changes unless the variable is set.

Verification

  • test/manual/layers/test_engram_peer_shards.py at the real table size, bitwise equal to the current lookup, eager and under CUDA graph replay: 8 GB300 ranks over two nodes (FABRIC), and 8 B300 GPUs on one host (POSIX fds).
  • DeepSeek-V4.1-Flash decode with DP attention over two GB300 nodes: the ReduceScatter and AllGather are gone from the decode step, greedy outputs and accept lengths match, and the reduced KV pool still fits the tested batch.

Related: #45679 (single-node TP4 peer reads).


CI States

Latest PR Test (Base): 🚫 Run #40045998118
Latest PR Test (Extra): ❌ Run #40045998125
Latest PR Test (AMD ROCm 10): 🚫 Run #40045998260

…p_alloc; assert mapped; warn on host-table override
@hnyls2002

Copy link
Copy Markdown
Collaborator Author

/rerun-test registered/unit/test_cuda_vmm_utils.py registered/disaggregation/test_disaggregation_dwdp_gpt_oss.py registered/unit/layers/test_engram_history.py registered/unit/models/test_deepseek_v41_bcg_multimodal.py registered/unit/kernels/test_kernels_namespace.py registered/e2e/models/test_deepseek_v41_flash_gb300_dp2ep2.py

@hnyls2002

Copy link
Copy Markdown
Collaborator Author

/tag-and-rerun-ci

@github-actions

github-actions Bot commented Oct 24, 2026 •

Copy link
Copy Markdown
Contributor

🚀 2-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/unit/test_cuda_vmm_utils.py

🚀 4-gpu-b200 (1 test): ✅ View workflow run

cd test/ && python3 registered/disaggregation/test_disaggregation_dwdp_gpt_oss.py

🚀 ubuntu-latest (3 tests): ✅ View workflow run

cd test/ && python3 registered/unit/layers/test_engram_history.py
cd test/ && python3 registered/unit/models/test_deepseek_v41_bcg_multimodal.py
cd test/ && python3 registered/unit/kernels/test_kernels_namespace.py

🚀 2-gpu-gb300 (1 test): ✅ View workflow run

cd test/ && python3 registered/e2e/models/test_deepseek_v41_flash_gb300_dp2ep2.py

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

jit-kernel run-ci CI: run the baseline test suite on this PR

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants