Add DSpark target-model serving (deepseek_v4_flash_dspark) - #213
Merged
Merged
Conversation
|
Important
This repository does not receive automatic reviews because it has fewer than 10 stars. ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Team Run ID: Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
ndleslx
force-pushed
the
worktree-dsv4-flash-dspark
branch
9 times, most recently
from
September 6, 2026 03:44
f530938 to
5052401
Compare
ndleslx
force-pushed
the
worktree-dsv4-flash-dspark
branch
2 times, most recently
from
September 7, 2026 05:10
9c369a8 to
618c1d8
Compare
Introduce the DSpark-specific W8A8 weight transformation while reusing the existing DeepSeek V4 checkpoint names, dtypes, expert placement, and layer grouping. - Tensor-parallel shard wo_a by assigning two of its eight projection groups to each TP rank. - Shard the matching wo_b INT8 columns across TP ranks and repeat each TP shard across the four DP groups. - Pad hc_attn_fn and hc_ffn_fn from 24 logical rows to the decode weight bank's 32-row storage layout. - Derive separate contiguous, unpadded 24-row HC slabs for prefill from the padded decode weights. - Stack all hidden-layer weights into resident device banks, including the compression-ratio-specific CSA and HCA groups. - Keep loading lazy and backed directly by safetensors; DSpark does not introduce a prepacked sidecar. - Support packing into preallocated destinations to limit temporary host memory during startup. - Accept DeepSeek configurations with multiple draft/final compression entries while consuming only the hidden-layer prefix.
ndleslx
force-pushed
the
worktree-dsv4-flash-dspark
branch
from
September 7, 2026 08:40
c4b3333 to
83dfc04
Compare
superxf
reviewed
Sep 7, 2026
| for tp_rank in range(layout.tp_size): | ||
| source_rank = source_group * layout.tp_size + tp_rank | ||
| target_rank = group * layout.tp_size + tp_rank | ||
| tensor[target_rank].copy_(tensor[source_rank]) |
Collaborator
There was a problem hiding this comment.
未参与本轮 prefill 的 TP 组，会收到活跃组的完整 block table 和写入 slot mapping，但这些页号没有重映射到独立 scratch 区域。“本轮没有 prefill”不代表这个组没有存活请求，调度器允许旧请求保留缓存、暂停 decode，再处理新请求的 prefill。
例如：A 在 group 0 使用 page 0 生成中，B 在 group 1 使用自己的 page 0 做 prefill。镜像会把 B 的数据写到 group 0 的 page 0，破坏 A 的历史 KV 和压缩状态，影响后续生成。
建议：镜像计算使用隔离的 scratch 缓存，并重映射所有相关表和写入地址；或由内核支持不写缓存的空闲组。
Collaborator
Author
There was a problem hiding this comment.
这个需要lib侧支持，无效的dp组应该跳过计算，当前decode已经做了，prefill没有
| f"DSpark prefill chunk must be in [1, {DSPARK_PREFILL_MAX_TOKENS}] tokens, " | ||
| f"got {actual_tokens}" | ||
| ) | ||
| if chunk_start + tokens > model.runtime.max_seq_len: |
Collaborator
There was a problem hiding this comment.
这里用批内最大的物理 chunk 长度 tokens 检查每个请求，而不是请求自己的有效长度。
Implement the serving runner for the DSpark target kernels on the canonical 16-rank TP4/DP4/EP16 topology. - Model the four TP groups as scheduler-visible cache partitions whose four ranks hold equivalent replicated cache pools. - Define and allocate the ori, cmp_c128, cmp_c4, idx, hca_state, csa_state, and csa_inner_state cache families. - Validate the Flash model dimensions, 32-token pages, 16K decode context ceiling, and 64-request capacity per TP group. - Materialize resident weights and prefill rings before sizing KV capacity, reserve isolated filler scratch pages, and retry allocation at reduced capacity after an OOM. - Size KV capacity from the shared L3 worker's device_memory_info, queried by logical worker ID: torch_npu is no longer a serving dependency, so torch.npu.mem_get_info must not be reached on workers. - Lower scheduler block IDs into the kernel's ring, trailing-ring, absolute, compressed, compressor-state, and SWA-window metadata. - Preserve shared prefill/decode KV and compressor-state pools while exposing the different names required by their ABIs. - Build independent base, ratio-4, ratio-128, and half-width HCA RoPE inputs, including distinct rank-local query and group-wide KV tables. - Stage up to one prefill request per TP group with TP-aligned dynamic extents, packed query boundaries, zero-padded rows, and invalid cache mappings set to -1. - Bound each packed request's context check by its own effective length: the batch-max physical padding tail carries only synthetic rows and consumes no context window. - Mirror an active prefill group into inactive partitions so all TP groups execute valid data against their partition-local cache pages. - Run decode at the fixed 64-request by 8-row tile, distribute real requests across TP owners, and isolate filler rows on scratch pages. - Commit only each request's anchor row, masking the seven unaccepted rows from raw, compressed, and recurrent-state cache publication. - Represent inactive owners with zero token counts instead of using an invalid token ID as padding. - Preserve the 16-token CSA decode transaction ring needed for the eight historical and eight eager state rows. - Return the device-greedy sampled token for each active request. - Allow each L3 dispatch to override RunConfig so prefill and decode can use their independently required ring-heap profiles.
Connect the DSpark runner to kernel compilation, command-line model selection, and the HTTP serving worker. - Add a DSpark executor that validates W8A8 model metadata, creates the layer plan and weight store, and compiles l3_prefill_fwd and l3_decode_fwd. - Import kernels with the frozen --tp 4 --ep 16 and --weight-bank-size 43 shape arguments. - Isolate generic kernel module names such as config, moe, and lm_head so DSpark imports cannot reuse modules from the MTP kernel package. - Construct the prefill and decode RoPE profiles through the kernel's own configuration and utility functions. - Declare every prefill and decode TaskArgs entry in exact pypto-lib ABI order, including weights, cache pools, metadata, scratch buffers, logits, and sampled IDs. - Advertise device embedding and device greedy sampling while rejecting speculative tokens for this target-only milestone. - Add method="dspark" model selection and route workers to PyptoDeepSeekV4DSparkExecutor. - Enforce the external 16-device DP4/TP4/EP16 contract, block size 32, maximum sequence length 16384, and the fixed decode capacity. - Collapse DP and TP to one internal 16-rank worker because those axes are already implemented by the kernels and cache partitions. - Disable prefix caching and install the DSpark cache-group definitions in the runtime configuration. - Permit DeepSeek checkpoints with additional draft compression-ratio entries without changing the existing MTP execution path. - Pin pypto-lib at f069bc78227422a1313133d1e880c2ed42c8b9fb, including the upstream DSpark fixes for RoPE separation, transaction state capacity, communication ordering, cache-line writers, prefill scratch outputs, uneven owner occupancy, the serialised CSA decode cube launches, the prefill HCA and layer ring-heap sizing with epoch-signal clearing, the HCA request-chain TaskId-array carry fix, and the coarsened compressor kv_score_proj grid.
Add focused validation and operating documentation for the DSpark target-model serving path. - Add a 16-card HTTP accuracy guard for the Palace Museum case with 64 prompt tokens and 128 greedily generated tokens. - Print the complete response for reviewer inspection and verify the response model, finish reason, and token accounting without requiring token-for-token MTP parity. - Require PYPTO_DSV4_DSPARK_MODEL_DIR in the self-hosted runner environment, validate it with the other model paths, and forward it through task-submit with the current checkout on PYTHONPATH. - Add one host-side functional test that follows multi-group prefill into uneven-occupancy decode and checks padding, per-group RoPE, inactive-group mirroring, owner counts, fixed decode tiles, and committed-row cache masking. - Extend CLI and shared DeepSeek fixtures for the DSpark chunk contract and checkpoints with multiple draft compression-ratio entries. - Document the supported topology, cache contracts, dispatch behavior, ring heaps, accuracy command, and later DSpark drafter work. - Register the guide in the MkDocs navigation so pre-commit and strict documentation builds include it.
ndleslx
force-pushed
the
worktree-dsv4-flash-dspark
branch
2 times, most recently
from
September 7, 2026 11:17
7cf264f to
5dec77c
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Serve the DSpark DeepSeek-V4-Flash target kernels end to end on the canonical 16-card TP4/DP4/EP16 topology: prefill, decode, and device-side greedy generation without speculative decoding. The DSpark drafter chain remains a subsequent milestone tracked by #211.
pypto_serving/model/deepseek_dspark/with executor routing, packed L3 prefill/decode dispatch, grouped cache metadata, device sampling, weight loading, and TP-sharded output projection.Kernel dependency
Persistent prefill/decode correctness work and remaining validation are tracked by pypto-lib#1133.
pypto-libis pinned to57e9d6a9294d9c38edd042f1edb5bdc50b4622bb. The upstream rangee3d82f0..57e9d6aprovides the corresponding DSpark decode contract fixes:Correctness result
The long-prompt generation defect was caused by decode applying the uncompressed SWA RoPE profile to CSA/HCA layers whose prefill KV state used compressed YaRN RoPE. The split kernel ABI and serving staging now keep those profiles consistent across the prefill-to-decode transition.
The Palace Museum guard uses a 64-token prompt and generates 128 tokens on 16 cards with DP4/TP4/EP16, device greedy sampling, and speculation disabled. Repeated identical runs produce different but coherent continuations, consistent with the accepted kernel numerical nondeterminism.
Verification
task_20260905_105618_16867612588, passed in 524.38s.task_20260905_111110_181110227254, passed in 378.26s.318 passed, 2 failed; both failing paths are unchanged by this delta and use stale test doubles for the current runtime/worker APIs (copy_from(src_offset=...)and missingsampleron aWorkerProcess.__new__fixture).No replay scripts, state dumps, local patch files, or diagnostic runtime switches are included in the branch.