Skip to content

fix(ds4): correct HC shard handoff and scope runtime caches per shard#507

Draft
Graffioh wants to merge 2 commits into
Luce-Org:codex/ds4-rocmfpx-serverfrom
Graffioh:codex/fix-ds4-hc-shard-handoff
Draft

fix(ds4): correct HC shard handoff and scope runtime caches per shard#507
Graffioh wants to merge 2 commits into
Luce-Org:codex/ds4-rocmfpx-serverfrom
Graffioh:codex/fix-ds4-hc-shard-handoff

Conversation

@Graffioh

@Graffioh Graffioh commented Jul 11, 2026

Copy link
Copy Markdown
Contributor

Summary

Fixes two correctness issues in the DeepSeek4 HC-complete layer-range path:

  • pass the updated full HC state to every local shard after shard 0, instead of passing the original token embeddings again
  • move HC weights and cached decode graphs from process-global statics into runtime state owned by each shard's DeepSeek4Cache
  • validate that each shard has complete HC attention, FFN, and output weights before initializing its runtime cache
  • release the per-shard runtime cache together with the shard cache
  • add unit coverage for local-shard boundary input selection

Stack

Stacked on #503. Base is codex/ds4-rocmfpx-server, so this PR contains only the HC handoff and per-shard cache fixes. Retarget to the branch where #503 lands once it is merged.

The layer-range implementation corrected here is introduced by #503.

Handoff correctness

The layer-range input has two different contracts:

  • shard 0 consumes token embeddings: n_tokens × n_embd
  • every later shard consumes the full HC boundary state produced by the previous shard: n_tokens × n_hc × n_embd

The local multi-shard path previously passed embed.data() to every shard. Later shards therefore interpreted the original embedding buffer as a full HC tensor, instead of consuming the HC state updated by the preceding layer range.

The corrected local flow is:

token embeddings
  -> shard 0
  -> full HC state
  -> shard 1
  -> full HC state
  -> ...
  -> final shard
  -> logits

Local shards share hc_state_ in place. The layer-range step now detects that the boundary input aliases hc_state.data() and avoids an unnecessary copy or a resize that could invalidate the input pointer. IPC shard boundaries continue to copy from their separate boundary-activation buffer.

Per-shard runtime cache

The HC weights and cached decode graphs used by deepseek4_step_layer_range() were previously function-local static objects shared by every shard in the process.

That cache was keyed only through partial context checks and was reset as execution alternated between shards with different:

  • weight objects and ggml contexts
  • backends
  • layer ranges
  • output-head ownership

The runtime state now belongs to each shard's DeepSeek4Cache and is matched against the complete shard identity:

weights + context + backend + layer range + output ownership

Each shard loads only the HC weights for its assigned layer range. Only the final shard loads the HC output weights.

Commits

  • a68112a — hand off full HC state between local shards
  • b6e50aa — scope HC runtime caches per shard

Validation

  • branch built directly on the current fix/feat(ds4): ROCmFPX byte-safe decode + fused-decode path + ds4 optimizations #503 head (6c7ee19)
  • git diff --check: passed
  • diff is limited to the five expected DeepSeek4 implementation/test files
  • unit test added for local-shard boundary input selection
  • Release CUDA H200 builds passed for test_deepseek4_unit and dflash_server
  • IMPORTANT: virtual two-shard CUDA validation passed on one H200 using logical shards [0,22) and [22,43) on cuda:0
  • two sequential greedy requests completed successfully in the same server process, each producing 32 tokens with no HC boundary, forward, CUDA, SIGSEGV, or SIGABRT errors
  • the virtual split exercised in-process full-HC-state handoff and independent per-shard runtime cache reuse;

@Graffioh Graffioh changed the title fix(deepseek4): correct HC shard handoff and scope runtime caches per shard fix(ds4): correct HC shard handoff and scope runtime caches per shard Jul 11, 2026
@davide221

Copy link
Copy Markdown
Contributor

We plan to run two AMD GPUs in the same HIP process. g_scratch is currently shared across devices, so the second shard may reuse scratch memory allocated on the first GPU. Please keep one scratch buffer per HIP device, set the correct device in HC worker threads

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants