feat(mla): add decode context parallelism - #1162
Conversation
ddc5e15 to
29a340e
Compare
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
tokenspeed/python/tokenspeed/runtime/layers/attention/backends/deepseek_v4.py
Lines 1841 to 1843 in 29a340e
Restore compressed-table metadata before V4 prefill
When a DeepSeek-V4 ratio-4 layer enters prefill, this call references compressed_base_offsets and compressed_table_capacity, but the refactor moved the offset calculation into gather_compressed_history as a local base_offsets and removed both outer bindings. Python therefore raises NameError before combining the CSA and SWA indices; the same undefined names are also used by the later dense compressed-history branch.
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
bcd4823 to
d0f6e5d
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d0f6e5d2bd
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| # The above copyright notice and this permission notice shall be included in all | ||
| # copies or substantial portions of the Software. | ||
|
|
There was a problem hiding this comment.
This newly added source file stops the license notice after the inclusion clause and omits the MIT warranty and liability disclaimer. Replace the abbreviated header with the repository's full MIT header before distributing the file.
AGENTS.md reference: AGENTS.md:L15-L16
Useful? React with 👍 / 👎.
| def dsv4_compact_dcp_topk_indices( | ||
| indices: torch.Tensor, | ||
| *, | ||
| dcp_size: int, | ||
| dcp_rank: int, | ||
| ) -> tuple[torch.Tensor, torch.Tensor]: | ||
| """Stable fixed-shape compaction of one DCP rank's compressed indices.""" |
There was a problem hiding this comment.
This function is exported through both tokenspeed_kernel.ops.attention and the top-level tokenspeed_kernel package, but its one-line docstring does not explain the global-index input contract, dcp_size/dcp_rank, or the returned compacted indices and lengths. Add explicit argument and return documentation for this new public API.
AGENTS.md reference: AGENTS.md:L51-L53
Useful? React with 👍 / 👎.
| parser.add_argument( | ||
| "--decode-context-parallel-size", | ||
| "--dcp-size", |
There was a problem hiding this comment.
These new public flags introduce substantial backend, dtype, topology, disaggregation, and cache-storage constraints, but the commit adds no repository documentation explaining how to configure or use them. Add a user-facing DCP section covering supported models/backends, required options, communication modes, and incompatibilities.
AGENTS.md reference: AGENTS.md:L11-L14
Useful? React with 👍 / 👎.
kangguangli
left a comment
kangguangli
left a comment
There was a problem hiding this comment.
The overall design of virtually expanding the page size looks good, as it allows all ranks to share the same logical block-table structure. My concerns are:
- The SKIP_ZERO design is worth discussing. It requires every KV writer that may receive location 0 to skip writing to the reserved null page, while page 0 is frequently used as a placeholder in other paths.
I am concerned that this implicit contract may create maintenance issues in future development. I suggest passing an explicit write mask instead of relying on default writer behavior. For example, Kimi-K3 reuses DeepseekV3AttentionMLA and may invoke apply_rope_mla_set_kv() to update the cache, but this fused path does not currently support SKIP_ZERO.
Alternatively, page 0 could be writable, but in that case read and reconstruction paths should not rely on it always remaining zero, since that invariant is easy to break. - During MLA prefill, each rank materializes a full-length cached-prefix tensor, for example [KV0, 0, KV2, 0, KV4, 0] on rank 0. This may introduce redundant memory loads in get_mla_kv_buffer() and extra communication in reconstruct_prefix_kv(), especially when the cached prefix is long. Could we read only the rank-local rows and reconstruct them afterward?
- Enabling DCP results in a 14.3% throughput regression compared with TP. This seems relatively large considering the communication payload and the NVLink bandwidth of Blackwell. Could you provide a short analysis or profile breakdown of the main sources of this regression?
- I suggest adding accuracy tests like evaluation on GSM8K or GPQA. A direct logit-parity comparison against the TP baseline would be even better.
d0f6e5d to
f2d3731
Compare
|
hello, any update? |
f2d3731 to
8b05dd6
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8b05dd6a53
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| from typing import TYPE_CHECKING | ||
|
|
||
| import torch | ||
| from tokenspeed_kernel._triton import tl, triton |
There was a problem hiding this comment.
Importing the MLA backend now unconditionally imports this module, which reaches through the public kernel boundary to the private _triton shim and defines a @triton.jit kernel in runtime code—even when the selected communication backend is ag_rs. Move the packing kernel into tokenspeed-kernel and expose it as a public operation so the runtime does not own or directly depend on third-party kernel implementation details.
AGENTS.md reference: AGENTS.md:L47-L51
Useful? React with 👍 / 👎.
2a5adef to
2f27f52
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2f27f521e9
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| *, | ||
| world_size: int, | ||
| ) -> torch.Tensor: | ||
| """Pack output head slices and lossless FP32 LSE words for DCP A2A.""" |
There was a problem hiding this comment.
This function is newly exported through both tokenspeed_kernel.ops.attention and the top-level tokenspeed_kernel package, but its one-line docstring does not explain the required output/LSE shapes and dtypes, the world_size head-partition contract, or the packed tensor's layout and return shape. Add explicit argument and return documentation for this public API.
AGENTS.md reference: AGENTS.md:L67-L69
Useful? React with 👍 / 👎.
|
2 participants
Summary
Adding decode context parallelism (DCP), building on top of the fp8 MLA kernel introduced in #364 , i subdivide each attention tensor-parallel group and cyclically place the sequence history across its ranks. tested and implemented with dsv4 model family and kimi k3 in mind
for sequence-sharded MLA history, the per-rank physical storage is:
thought, it's importnant that we note that this is the saving for the eligible dcp history, not necessarily the entire model cache. dsv4 for example will keep SWA, the FP4 sparse indexer, and compressor state sequence-replicated, and models like kimi k3 keep kda state sequence-replicated because it is already partitioned by TP heads.
When testing with dsv4-flash with a fixed global 8,192-token capacity, dcp2 halved the compressed history rows and reduced the complete rank-local arena from 1,574,012,160 to 1,513,935,360 bytes (3.82%, about 57.3 MiB per rank).
Test Plan
dsv4 flash official checkpoint (
deepseek-ai/DeepSeek-V4-Flash) on 2×B200:ag_rs, and TP2/DCP2a2aend-to-end serving.MODEL=deepseek-ai/DeepSeek-V4-Flash
REVISION=60d8d70770c6776ff598c94bb586a859a38244f1
Baseline TP2:
tokenspeed serve "$MODEL" "${COMMON_ARGS[@]}"TP2/DCP2 with AG/RS:
TP2/DCP2 with A2A:
trace above shows dcp's decode overhead is primarily from comms, with A2A reducing the per-layer collective path from three operations and 378 µs to two operations and 292 µs.
A2A improved decode throughput by 17.4% over AG/RS, while remaining 14.3% below the no dcp baseline. dcp preserves the per-rank attention work product (DCP× gathered heads over 1/DCP context), so its primary benefit is kv capacity as the collectives add latency
kimi k3 on 8xb300, using the official checkpoint:
At a fixed 65,536-token capacity:
DCP2: 1,189,085,184 → 732,561,408 bytes/rank, saving 38.39% (435.4 MiB/rank).
DCP4: 1,189,085,184 → 508,723,200 bytes/rank, saving 57.22% (648.8 MiB/rank).
DCP8: 1,189,085,184 → 386,629,632 bytes/rank, saving 67.49% (765.3 MiB/rank).
The MLA history scales as 1/dcp, while about 263.3 MiB/rank of kda state remains replicated, which produces diminishing memory savings at higher DCP degrees.