Guard DP local info when global token counts are missing - #1
Conversation
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing touches🧪 Generate unit tests (beta)
Comment |
![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
Here are some automated review suggestions for this pull request.
Reviewed commit: fe5f235769
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if forward_batch.forward_mode.is_extend( | ||
| include_draft_extend_v2=True | ||
| ): | ||
| local_num_tokens = forward_batch.seq_lens_sum | ||
| else: |
There was a problem hiding this comment.
When global_num_tokens_gpu is missing and dp_size == 1, the fallback for extend modes uses seq_lens_sum. In extend batches with cached prefixes, seq_lens_sum counts the full fill_ids length (prefix + new tokens), while the local buffers only contain the new tokens (e.g., extend_num_tokens/out_cache_loc length). In that case local_num_tokens is larger than the local buffer, so the later memcpy_triton in DP gather/scatter can read/write past the local buffer. Consider deriving the count from extend_num_tokens or out_cache_loc.shape[0] for extend modes to match the actual local token buffer size.
Useful? React with 👍 / 👎.
- hc_pre NPU branch: return 4-tuple (y, post, comb, norm_fused=False) to match caller's unpack. Was 3-tuple → ValueError on first non-empty forward. (P1, reviewer: AndyLi429 #1 + Codex) - _apply_hadamard: scale by head_dim**-0.5 to match the CUDA rotate_activation reference; otherwise indexer dot products were scaled by sqrt(n) and top-k selection drifted from reference. (P1, reviewer: AndyLi429 #2 + Codex) - DeepSeekV4TokenToKVPool: log per-req c4/c128 slab capacity at init so the average-allocator limitation is visible (low-concurrency long-context still needs a real allocator — followup). (P1, reviewer: AndyLi429 sgl-project#3) - custom_ops ImportError: fail-fast with a clear message naming the required wheel / CANN image, instead of logging "fall back" when no fallback exists. (P2, reviewer: AndyLi429 sgl-project#4) - _forward_compressed: raise on missing metadata/cmp_kv with layer_id / compress_ratio / available metadata keys, instead of silently returning zeros. (P2, reviewer: AndyLi429 sgl-project#5) - DeepSeekV4SingleKVPool.create_buffer: gate the NPU PA_ND bf16 layout on _is_npu, not on dtype alone. (P2, reviewer: AndyLi429 sgl-project#6 + zhuyijie88 inline) Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 participant
Motivation
global_num_tokens_gpuis missing, which can cause multi-rank DP to proceed with invalid metadata.Description
get_dp_local_info(python/sglang/srt/layers/dp_attention.py) requireglobal_num_tokens_gpuwhenget_attention_dp_size() > 1and raise a clearRuntimeErrorif it is missing.dp_size == 1case, compute and setdp_local_start_posanddp_local_num_tokensdirectly fromseq_lens_sumorbatch_sizeand return without constructing global token tensors.cumtokensand extracts per-rank start/length unchanged whenglobal_num_tokens_gpuis present.Testing
Codex Task