models: load MXFP8 sparse indexer WK scales - #33
Conversation
|
Warning Review limit reached
More reviews will be available in 17 minutes and 43 seconds. Learn how PR review limits work. Your organization has used up its prepaid credits, and credit purchases are no longer available. Enable the review add-on in the billing tab to keep reviews running — you're only billed for reviews past your plan's rate limits ($0.25/file). ⌛ How to resolve this issue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based credits. 🚦 How do rate limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan refill rate. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, the refill rate gradually slows as usage increases. The highest same-day bursts are limited more strictly. Please see our Fair Usage Limits Policy for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📥 CommitsReviewing files that changed from the base of the PR and between 0429672ce3af12d485a810376354d089ef1cb429 and ce1881c. 📒 Files selected for processing (2)
📝 WalkthroughWalkthroughTwo independent functional areas are modified: (1) ChangesFP8/MXFP8 Indexer Weight Loader
B12X DCP Global-TopK Remapping and Warmup
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (2)
vllm/model_executor/models/deepseek_v2.py (1)
825-832: 🧹 Nitpick | 🔵 Trivial | ⚡ Quick winConvert the updated docstring to Google style.
The docstring was changed but still omits the required
Args:/Returns:sections.As per coding guidelines, "Use Google-style docstrings with
Args:/Returns:/Raises:sections, not reStructuredText/Sphinx fields (:param:,:return:,:rtype:)."♻️ Proposed docstring update
- """ - We fuse the WK and weights_proj projections, but in some checkpoints WK is stored - in FP8 with a separate weight_scale_inv or MXFP8 with a separate weight_scale, - while weights_proj is stored in BF16. Upcasting to BF16 during loading enables - the fusion. This function loads the WK weights and scale, and when both are - available, dequantizes to BF16 and stores into the fused - wk_weights_proj.weight parameter. - """ + """Load isolated FP8/MXFP8 indexer WK tensors into the fused WK weight. + + Args: + name: Checkpoint tensor name. + tensor: Checkpoint tensor value. + buf: Pending WK weight/scale tensors keyed by layer prefix. + params_dict: Model parameters keyed by checkpoint name. + loaded_params: Set of fused parameter names already loaded. + pp_missing_layer_names: Pipeline-parallel layer prefixes to skip. + + Returns: + ``True`` if this tensor was handled by this loader; otherwise ``False``. + + Raises: + KeyError: If the fused WK parameter is missing after a matching tensor + pair is ready to load. + """🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vllm/model_executor/models/deepseek_v2.py` around lines 825 - 832, The docstring for the function handling WK weight loading and fusion currently only contains a description but is missing the required Google-style docstring sections. Add the `Args:` section to document any function parameters, a `Returns:` section to describe what the function returns, and a `Raises:` section if applicable to document any exceptions the function may raise. Ensure all parameters and return values are properly documented in the Google-style format as per the coding guidelines.Source: Coding guidelines
vllm/v1/attention/backends/mla/indexer.py (1)
960-965: 🧹 Nitpick | 🔵 TrivialUse the existing
_dcp_global_topk_requested()helper for consistency.The metadata builder at lines 960–965 duplicates the global-topk parsing logic from
_dcp_global_topk_requested()insparse_attn_indexer.py(lines 50–52). Both currently parseVLLM_DCP_GLOBAL_TOPKidentically, but inlining it here creates a maintenance burden: any future change to the parsing semantics must be applied in both places.Extract the helper to
envs.pyor import_dcp_global_topk_requested()into the metadata builder to keep the gates unified.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vllm/v1/attention/backends/mla/indexer.py` around lines 960 - 965, The `keep_empty_dcp_chunk` variable assignment duplicates the environment variable parsing logic that already exists in the `_dcp_global_topk_requested()` helper function. Remove the inlined `VLLM_DCP_GLOBAL_TOPK` environment variable parsing logic from the `keep_empty_dcp_chunk` assignment and replace it with a call to the existing `_dcp_global_topk_requested()` helper function. Import this helper from `sparse_attn_indexer.py` or extract it to a shared location like `envs.py` to ensure consistent behavior and reduce maintenance burden.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@vllm/model_executor/layers/sparse_attn_indexer.py`:
- Around line 1458-1475: The masked_fill_ operation on topk_indices is being
applied unconditionally after both the dcp_global_topk and local conversion
branches, which causes valid candidates merged from other DCP ranks to be erased
by this rank's local row_has_no_kv mask. Move the
topk_indices.masked_fill_(row_has_no_kv[:, None], -1) call inside the else block
(after _convert_b12x_dcp_local_topk_to_global) so that masking only occurs when
using local topk conversion, not after the global merge performed by
_merge_b12x_dcp_topk.
- Around line 435-443: The broad exception handling around get_dcp_group() and
the dcp_group.barrier() call silently swallows failures that should be
propagated. Remove the try-except block entirely or replace it with a more
targeted exception handler that re-raises the exception instead of silently
returning, so that DCP barrier failures cause the process to fail fast rather
than allowing ranks to proceed out of sync and potentially hang later collective
operations.
In `@vllm/v1/attention/backends/mla/b12x_mla_sparse.py`:
- Around line 683-699: The methods `_sync_dcp_warmup` and
`_prewarm_extend_kernels_once` are missing Google-style docstrings. Add proper
docstrings to both methods following Google style guidelines with Args, Returns,
and Raises sections as applicable. For `_sync_dcp_warmup`, document the
synchronization behavior and exception handling. For
`_prewarm_extend_kernels_once`, document the max_batched parameter and any
relevant return or exception information.
- Around line 688-694: The try-except block that calls get_dcp_group() and
dcp_group.barrier() is catching all Exception types and silently returning,
which masks synchronization failures during DCP warmup and makes distributed
startup issues difficult to diagnose. Instead of catching all exceptions
broadly, either catch only the specific exceptions you expect (such as
ImportError for the missing import or specific exceptions from get_dcp_group or
barrier operations), or at minimum log the exception details before returning to
preserve debugging information while still allowing the function to degrade
gracefully.
- Line 72: The variable declaration for _EXTEND_PREWARM_DONE exceeds the
88-character line-length limit. Split the type annotation for the
_EXTEND_PREWARM_DONE variable across multiple lines, either by creating a
separate type alias for the complex tuple type and using that alias in the
variable declaration, or by breaking the type annotation using line
continuation. Ensure the final result respects the 88-character maximum line
length requirement.
---
Nitpick comments:
In `@vllm/model_executor/models/deepseek_v2.py`:
- Around line 825-832: The docstring for the function handling WK weight loading
and fusion currently only contains a description but is missing the required
Google-style docstring sections. Add the `Args:` section to document any
function parameters, a `Returns:` section to describe what the function returns,
and a `Raises:` section if applicable to document any exceptions the function
may raise. Ensure all parameters and return values are properly documented in
the Google-style format as per the coding guidelines.
In `@vllm/v1/attention/backends/mla/indexer.py`:
- Around line 960-965: The `keep_empty_dcp_chunk` variable assignment duplicates
the environment variable parsing logic that already exists in the
`_dcp_global_topk_requested()` helper function. Remove the inlined
`VLLM_DCP_GLOBAL_TOPK` environment variable parsing logic from the
`keep_empty_dcp_chunk` assignment and replace it with a call to the existing
`_dcp_global_topk_requested()` helper function. Import this helper from
`sparse_attn_indexer.py` or extract it to a shared location like `envs.py` to
ensure consistent behavior and reduce maintenance burden.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: d1b3c3f0-bd21-4034-ac74-3a78778e7373
📥 Commits
Reviewing files that changed from the base of the PR and between 1bcacde and 0429672ce3af12d485a810376354d089ef1cb429.
📒 Files selected for processing (6)
tests/models/test_deepseek_v2_indexer_loader.pyvllm/model_executor/layers/sparse_attn_indexer.pyvllm/model_executor/models/deepseek_mtp.pyvllm/model_executor/models/deepseek_v2.pyvllm/v1/attention/backends/mla/b12x_mla_sparse.pyvllm/v1/attention/backends/mla/indexer.py
| try: | ||
| from vllm.distributed.parallel_state import get_dcp_group | ||
|
|
||
| dcp_group = get_dcp_group() | ||
| if int(dcp_group.world_size) <= 1: | ||
| return | ||
| dcp_group.barrier() | ||
| except Exception: | ||
| return |
There was a problem hiding this comment.
Don’t swallow DCP barrier failures.
Line 442 converts any get_dcp_group() or barrier failure into a silent return. During DCP warmup, that can let ranks proceed out of sync and hang later collectives; fail fast after the final CUDA sync instead.
Suggested fix
- try:
- from vllm.distributed.parallel_state import get_dcp_group
-
- dcp_group = get_dcp_group()
- if int(dcp_group.world_size) <= 1:
- return
- dcp_group.barrier()
- except Exception:
- return
- finally:
- if current_platform.is_cuda():
- torch.cuda.synchronize()
+ try:
+ from vllm.distributed.parallel_state import get_dcp_group
+
+ dcp_group = get_dcp_group()
+ if int(dcp_group.world_size) <= 1:
+ return
+ dcp_group.barrier()
+ finally:
+ if current_platform.is_cuda():
+ torch.cuda.synchronize()📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| try: | |
| from vllm.distributed.parallel_state import get_dcp_group | |
| dcp_group = get_dcp_group() | |
| if int(dcp_group.world_size) <= 1: | |
| return | |
| dcp_group.barrier() | |
| except Exception: | |
| return | |
| try: | |
| from vllm.distributed.parallel_state import get_dcp_group | |
| dcp_group = get_dcp_group() | |
| if int(dcp_group.world_size) <= 1: | |
| return | |
| dcp_group.barrier() | |
| finally: | |
| if current_platform.is_cuda(): | |
| torch.cuda.synchronize() |
🧰 Tools
🪛 Ruff (0.15.17)
[warning] 442-442: Do not catch blind exception: Exception
(BLE001)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@vllm/model_executor/layers/sparse_attn_indexer.py` around lines 435 - 443,
The broad exception handling around get_dcp_group() and the dcp_group.barrier()
call silently swallows failures that should be propagated. Remove the try-except
block entirely or replace it with a more targeted exception handler that
re-raises the exception instead of silently returning, so that DCP barrier
failures cause the process to fail fast rather than allowing ranks to proceed
out of sync and potentially hang later collective operations.
Source: Linters/SAST tools
| @@ -69,6 +69,7 @@ | |||
| # wave-balanced planner picks num_splits <= this cap. | |||
| _DECODE_SPLIT_TILE = 64 | |||
| _PREFILL_HEADS_PER_BLOCK = 16 | |||
| _EXTEND_PREWARM_DONE: set[tuple[int | None, int, int, int, int, int, bool]] = set() | |||
There was a problem hiding this comment.
Line 72 exceeds the Python 88-character limit.
Please split this declaration to keep it within the enforced line-length cap.
Proposed fix
+_ExtendPrewarmKey = tuple[int | None, int, int, int, int, int, bool]
-_EXTEND_PREWARM_DONE: set[tuple[int | None, int, int, int, int, int, bool]] = set()
+_EXTEND_PREWARM_DONE: set[_ExtendPrewarmKey] = set()As per coding guidelines, "Keep Python line length at 88 characters maximum".
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| _EXTEND_PREWARM_DONE: set[tuple[int | None, int, int, int, int, int, bool]] = set() | |
| _ExtendPrewarmKey = tuple[int | None, int, int, int, int, int, bool] | |
| _EXTEND_PREWARM_DONE: set[_ExtendPrewarmKey] = set() |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@vllm/v1/attention/backends/mla/b12x_mla_sparse.py` at line 72, The variable
declaration for _EXTEND_PREWARM_DONE exceeds the 88-character line-length limit.
Split the type annotation for the _EXTEND_PREWARM_DONE variable across multiple
lines, either by creating a separate type alias for the complex tuple type and
using that alias in the variable declaration, or by breaking the type annotation
using line continuation. Ensure the final result respects the 88-character
maximum line length requirement.
Source: Coding guidelines
| def _sync_dcp_warmup(self) -> None: | ||
| if self.device.type == "cuda": | ||
| torch.cuda.synchronize(self.device) | ||
| if self.dcp_world_size <= 1: | ||
| return | ||
| try: | ||
| from vllm.distributed.parallel_state import get_dcp_group | ||
|
|
||
| dcp_group = get_dcp_group() | ||
| dcp_group.barrier() | ||
| except Exception: | ||
| return | ||
| finally: | ||
| if self.device.type == "cuda": | ||
| torch.cuda.synchronize(self.device) | ||
|
|
||
| def _prewarm_extend_kernels_once(self, max_batched: int) -> None: |
There was a problem hiding this comment.
Add Google-style docstrings for the new helper methods.
_sync_dcp_warmup and _prewarm_extend_kernels_once were newly introduced but have no
Google-style docstrings (Args: / Returns: / Raises:).
Proposed fix
def _sync_dcp_warmup(self) -> None:
+ """Synchronize warmup across CUDA and (optionally) DCP ranks.
+
+ Returns:
+ None.
+ """
if self.device.type == "cuda":
torch.cuda.synchronize(self.device)
@@
def _prewarm_extend_kernels_once(self, max_batched: int) -> None:
+ """Run one-time extend-kernel warmup for representative row counts.
+
+ Args:
+ max_batched: Maximum batched token count used for warmup coverage.
+
+ Returns:
+ None.
+ """
if self.device.type != "cuda":
returnAs per coding guidelines, "Use Google-style docstrings with Args:/Returns:/Raises: sections, not reStructuredText/Sphinx fields (:param:, :return:, :rtype:)".
Also applies to: 699-768
🧰 Tools
🪛 Ruff (0.15.17)
[warning] 693-693: Do not catch blind exception: Exception
(BLE001)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@vllm/v1/attention/backends/mla/b12x_mla_sparse.py` around lines 683 - 699,
The methods `_sync_dcp_warmup` and `_prewarm_extend_kernels_once` are missing
Google-style docstrings. Add proper docstrings to both methods following Google
style guidelines with Args, Returns, and Raises sections as applicable. For
`_sync_dcp_warmup`, document the synchronization behavior and exception
handling. For `_prewarm_extend_kernels_once`, document the max_batched parameter
and any relevant return or exception information.
Source: Coding guidelines
| try: | ||
| from vllm.distributed.parallel_state import get_dcp_group | ||
|
|
||
| dcp_group = get_dcp_group() | ||
| dcp_group.barrier() | ||
| except Exception: | ||
| return |
There was a problem hiding this comment.
Avoid swallowing all DCP warmup failures silently.
Catching Exception and immediately returning can hide synchronization failures during
prewarm and make distributed startup issues hard to diagnose.
Proposed fix
try:
from vllm.distributed.parallel_state import get_dcp_group
dcp_group = get_dcp_group()
dcp_group.barrier()
- except Exception:
+ except AssertionError:
+ # DCP group not initialized in this runtime path.
+ return
+ except Exception:
+ logger.warning(
+ "DCP barrier failed during B12X extend-kernel warmup.",
+ exc_info=True,
+ )
return🧰 Tools
🪛 Ruff (0.15.17)
[warning] 693-693: Do not catch blind exception: Exception
(BLE001)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@vllm/v1/attention/backends/mla/b12x_mla_sparse.py` around lines 688 - 694,
The try-except block that calls get_dcp_group() and dcp_group.barrier() is
catching all Exception types and silently returning, which masks synchronization
failures during DCP warmup and makes distributed startup issues difficult to
diagnose. Instead of catching all exceptions broadly, either catch only the
specific exceptions you expect (such as ImportError for the missing import or
specific exceptions from get_dcp_group or barrier operations), or at minimum log
the exception details before returning to preserve debugging information while
still allowing the function to degrade gracefully.
Source: Linters/SAST tools
0429672 to
ce1881c
Compare
2 participants
Summary
Add the missing MXFP8 sparse-indexer WK load path for GLM/DeepSeek-style fused indexer projections.
The hybrid GLM-5.2 checkpoint stores the sparse indexer as:
indexer.wk.weight+indexer.wk.weight_scalewhere the scale tensor isuint8MXFP8 metadataindexer.wq_b.weight+indexer.wq_b.weight_scalebut the fused runtime module loads into
wk_weights_proj/wq_b_weights_proj. The existing loader handled legacy FP8weight_scale_inv, but a clean dark-devotion/PR15 image fails this hybrid checkpoint with:This patch treats
indexer.wk.weight_scalewithuint8dtype as MXFP8, dequantizes it withdequant_mxfp8_to_bf16, and loads the fusedwk_weights_proj.weight. Legacy FP8weight_scale_invremains unchanged, and unrelated MXFP8 scales are ignored.Validation
python3 -m py_compile vllm/model_executor/models/deepseek_v2.py tests/models/test_deepseek_v2_indexer_loader.pygit diff --checkindexer.wk.weight_scaleindexer.wk.weight_scale_invindexer.wq_b.weight_scaleignore path166/166shards loaded; model load51.85 GiB/GPU; KV663,872tokens.GLM-5.2-LUKE-NVFP4-PLUS-MXFP8-FP8MASK-FROM-BF16-20260618with the loader overlay:0.07206133562565563over2047positionsjs_mean=2.8724894e-06,kl_a_to_b_mean=9.8642940e-06,kl_b_to_a_mean=1.4637346e-05Summary by CodeRabbit
Tests
Bug Fixes
Performance