glm52: calibrated NVFP4 MLA KV outer scales + serve wiring - #144
glm52: calibrated NVFP4 MLA KV outer scales + serve wiring#144MadeBy561 wants to merge 1 commit into
Conversation
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
|
Warning Review limit reached
Next review available in: 36 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?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 reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. 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, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (3)
📝 WalkthroughWalkthroughChangesThe PR adds GLM-5.2 NVFP4 MLA outer-scale configuration, B12X PCIe and DCP channel lifecycle handling, sparse MLA CUDA graph memory profiling, shared-expert auxiliary-stream gating, and lane-scoped workspace management for speculative execution. NVFP4 MLA outer scales
B12X channels and CUDA graph memory profiling
Shared-expert auxiliary-stream gating
Speculative workspace lanes
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant GPUModelRunner
participant parallel_state
participant dcp_alltoall
participant CUDA graph managers
GPUModelRunner->>parallel_state: checkpoint B12X graph channels
GPUModelRunner->>GPUModelRunner: initialize profiling KV cache
GPUModelRunner->>parallel_state: capture disposable CUDA graphs
parallel_state->>dcp_alltoall: capture B12X DCP A2A pools
GPUModelRunner->>CUDA graph managers: clear captured graph state
GPUModelRunner->>parallel_state: rollback B12X graph channels
parallel_state->>dcp_alltoall: restore DCP A2A channel pools
GPUModelRunner-->>GPUModelRunner: return estimated graph memory
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ 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 |
Per-layer outer-scale calibration for nvfp4_ds_mla KV (VLLM_NVFP4_MLA_SCALES_FILE, format v1) with an explicit default-off knob in serve-glm52.sh. Collapses the NVFP4-vs-FP8 KV KLD gap to ~+0.008-0.009 (0.1345/0.1356 vs 0.1263, 5 fresh boots each, rtx6kpro protocol) while raising max context from 373k to 550k/600k+ on 4x96GB. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
a3e7a99 to
99bf3f1
Compare
|
Superseded — reopening against the reset dev/gilded-gnosis per maintainer request. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
vllm/v1/worker/workspace.py (1)
143-160: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winAdd
Raises:section to the docstring.The method now raises a
RuntimeErrorif an unconfigured lane is selected. As per coding guidelines, please update the docstring to explicitly include aRaises:section documenting this behavior.📝 Proposed docstring update
def _ensure_workspace_size(self, required_bytes: int) -> torch.Tensor: """Ensure workspace is allocated and large enough, return current workspace. Args: required_bytes: The number of bytes required. Returns: The current workspace tensor. + + Raises: + RuntimeError: If the currently selected lane is not configured. """ ubatch_id = dbo_current_ubatch_id()🤖 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/worker/workspace.py` around lines 143 - 160, Update the _ensure_workspace_size docstring to add a Raises section documenting that it raises RuntimeError when the selected workspace lane is not configured.Source: Coding guidelines
🧹 Nitpick comments (2)
vllm/distributed/parallel_state.py (1)
1548-1569: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMinor: duplicated DCP-active condition.
_DCP is not None and get_dcp_group().world_size > 1is computed separately formaybe_dcp_captureandmaybe_b12x_dcp_capture. Hoisting it into a single local (e.g.dcp_active) would avoid the repeated attribute/property lookups and keep the two branches obviously in sync if the condition ever changes.♻️ Suggested consolidation
- maybe_dcp_capture = ( - get_dcp_group().graph_capture(context) - if _DCP is not None and get_dcp_group().world_size > 1 - else nullcontext() - ) - if _DCP is not None and get_dcp_group().world_size > 1: - # Import locally to avoid making distributed initialization depend on - # attention modules. The helper is a no-op until DCP warmup creates a - # B12X pool for this process group. - from vllm.v1.attention.ops.dcp_alltoall import capture_b12x_dcp_a2a - - maybe_b12x_dcp_capture = capture_b12x_dcp_a2a(get_dcp_group(), context.stream) - else: - maybe_b12x_dcp_capture = nullcontext() + dcp_group = _DCP + dcp_active = dcp_group is not None and dcp_group.world_size > 1 + maybe_dcp_capture = dcp_group.graph_capture(context) if dcp_active else nullcontext() + if dcp_active: + # Import locally to avoid making distributed initialization depend on + # attention modules. The helper is a no-op until DCP warmup creates a + # B12X pool for this process group. + from vllm.v1.attention.ops.dcp_alltoall import capture_b12x_dcp_a2a + + maybe_b12x_dcp_capture = capture_b12x_dcp_a2a(dcp_group, context.stream) + else: + maybe_b12x_dcp_capture = nullcontext()🤖 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/distributed/parallel_state.py` around lines 1548 - 1569, Hoist the repeated `_DCP is not None and get_dcp_group().world_size > 1` check into a local `dcp_active` boolean near `maybe_dcp_capture`. Reuse that local for both the `maybe_dcp_capture` and `maybe_b12x_dcp_capture` branches, keeping their existing behavior unchanged.tests/v1/cudagraph/test_breakable_cudagraph.py (1)
42-134: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winTest never exercises the speculator branch of
profile_cudagraph_memory.
runner.speculator = Nonehere, somanagers.extend(self.speculator.get_cudagraph_managers())and thewith use_workspace_lane(1): self.speculator.capture()path (and its pool-swap/restore) are untested. Given this PR adds the speculator-aware channel checkpoint/rollback and pool-swap logic specifically for this path, a variant with a fake speculator (exposingget_cudagraph_managers()/clear_cudagraphs()/capture()) would meaningfully increase coverage of the changed code.🤖 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 `@tests/v1/cudagraph/test_breakable_cudagraph.py` around lines 42 - 134, Extend test_memory_profile_destroys_graphs_before_restoring_pools with a fake speculator exposing get_cudagraph_managers(), clear_cudagraphs(), and capture(), and assign it to runner.speculator. Assert the speculator capture runs under the expected workspace lane and that its manager and graph pools use profile_pool during profiling, then verify they are restored afterward. Preserve the existing checkpoint, cleanup, rollback, and event-order assertions while covering the speculator-specific path.
🤖 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/attention/mla_attention.py`:
- Around line 741-751: Update _get_sparse_memory_profile_bytes so it resolves
self.impl.dcp_world_size from get_dcp_group().world_size when the value is still
unset, before evaluating the early-return condition. Ensure the resolved DCP
world size is used during the attn_metadata=None profile pass so sparse-MLA
allocation is included.
In `@vllm/v1/worker/workspace.py`:
- Around line 35-44: Update the use_workspace_lane docstring with Google-style
Args and Raises sections: document the lane integer as the workspace lane
selected for the current execution context, and document that ValueError is
raised when lane is negative.
---
Outside diff comments:
In `@vllm/v1/worker/workspace.py`:
- Around line 143-160: Update the _ensure_workspace_size docstring to add a
Raises section documenting that it raises RuntimeError when the selected
workspace lane is not configured.
---
Nitpick comments:
In `@tests/v1/cudagraph/test_breakable_cudagraph.py`:
- Around line 42-134: Extend
test_memory_profile_destroys_graphs_before_restoring_pools with a fake
speculator exposing get_cudagraph_managers(), clear_cudagraphs(), and capture(),
and assign it to runner.speculator. Assert the speculator capture runs under the
expected workspace lane and that its manager and graph pools use profile_pool
during profiling, then verify they are restored afterward. Preserve the existing
checkpoint, cleanup, rollback, and event-order assertions while covering the
speculator-specific path.
In `@vllm/distributed/parallel_state.py`:
- Around line 1548-1569: Hoist the repeated `_DCP is not None and
get_dcp_group().world_size > 1` check into a local `dcp_active` boolean near
`maybe_dcp_capture`. Reuse that local for both the `maybe_dcp_capture` and
`maybe_b12x_dcp_capture` branches, keeping their existing behavior unchanged.
🪄 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: aba75a0f-aaad-4eb9-bb2a-15c74e49cf68
📒 Files selected for processing (29)
kv-scales/README.mdkv-scales/glm52-nvfp4-nf3-hybrid_mla_outer_scales_v1.jsonserve-glm52.shtests/distributed/test_b12x_fused_all_reduce.pytests/distributed/test_dcp_a2a.pytests/model_executor/layers/fused_moe/test_shared_experts_stream.pytests/model_executor/layers/test_b12x_moe_warmup.pytests/quantization/test_nvfp4_nf3_hybrid.pytests/v1/attention/test_b12x_mla_dcp_workspace.pytests/v1/cudagraph/test_breakable_cudagraph.pytests/v1/worker/test_workspace.pyvllm/distributed/device_communicators/custom_all_reduce.pyvllm/distributed/parallel_state.pyvllm/envs.pyvllm/model_executor/layers/attention/mla_attention.pyvllm/model_executor/layers/fused_moe/b12x_moe.pyvllm/model_executor/layers/fused_moe/fused_moe_method_base.pyvllm/model_executor/layers/fused_moe/runner/moe_runner.pyvllm/model_executor/layers/fused_moe/runner/shared_experts.pyvllm/model_executor/layers/quantization/nvfp4_nf3_hybrid.pyvllm/v1/attention/ops/dcp_alltoall.pyvllm/v1/worker/gpu/cudagraph_utils.pyvllm/v1/worker/gpu/model_runner.pyvllm/v1/worker/gpu/spec_decode/autoregressive/speculator.pyvllm/v1/worker/gpu/spec_decode/dflash/speculator.pyvllm/v1/worker/gpu/spec_decode/speculator.pyvllm/v1/worker/gpu/warmup.pyvllm/v1/worker/gpu_worker.pyvllm/v1/worker/workspace.py
| def _get_sparse_memory_profile_bytes(self) -> int: | ||
| if ( | ||
| not envs.VLLM_MEMORY_PROFILE_INCLUDE_ATTN | ||
| or self.attn_backend.get_name() != "B12X_MLA_SPARSE" | ||
| or self.impl.dcp_world_size <= 1 | ||
| ): | ||
| return 0 | ||
|
|
||
| max_tokens = int(getattr(self.impl, "_max_batched", 0)) | ||
| if max_tokens <= 0: | ||
| return 0 |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file="vllm/model_executor/layers/attention/mla_attention.py"
echo "== Outline =="
ast-grep outline "$file" --view expanded || true
echo
echo "== Relevant line ranges =="
sed -n '720,980p' "$file" | cat -nRepository: local-inference-lab/vllm
Length of output: 16876
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file="vllm/model_executor/layers/attention/mla_attention.py"
echo "== Search for dcp_world_size resolution and profiling call =="
rg -n "dcp_world_size|get_dcp_group|_get_sparse_memory_profile_bytes|profile" "$file"
echo
echo "== Around lazy resolution =="
sed -n '930,980p' "$file" | cat -n
echo
echo "== Around sparse memory profile helper =="
sed -n '730,820p' "$file" | cat -nRepository: local-inference-lab/vllm
Length of output: 10187
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file="vllm/model_executor/layers/attention/mla_attention.py"
echo "== Constructor / init around dcp_world_size =="
sed -n '2980,3020p' "$file" | cat -n
echo
echo "== Configured DCP world size logic =="
sed -n '612,640p' "$file" | cat -n
echo
echo "== Profile branch and helper call =="
sed -n '820,960p' "$file" | cat -nRepository: local-inference-lab/vllm
Length of output: 10032
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file="vllm/model_executor/layers/attention/mla_attention.py"
echo "== Helper body =="
sed -n '736,790p' "$file" | cat -n
echo
echo "== Profile-related call sites =="
sed -n '868,910p' "$file" | cat -n
echo
echo "== Forward profile branch around the call site =="
sed -n '936,958p' "$file" | cat -nRepository: local-inference-lab/vllm
Length of output: 5745
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file="vllm/model_executor/layers/attention/mla_attention.py"
echo "== Exact helper with line numbers =="
sed -n '740,760p' "$file" | cat -n
echo
echo "== Exact profile call site with line numbers =="
sed -n '938,952p' "$file" | cat -n
echo
echo "== Lines right after the profile branch =="
sed -n '952,968p' "$file" | cat -nRepository: local-inference-lab/vllm
Length of output: 2766
Resolve DCP world size before the profile-pass early return _get_sparse_memory_profile_bytes() can still see the -1 sentinel during the attn_metadata is None dummy run, so the new sparse-MLA DCP allocation is skipped and the profiler underestimates peak memory. Read get_dcp_group().world_size when unset, or initialize self.impl.dcp_world_size before this helper runs.
🤖 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/attention/mla_attention.py` around lines 741 -
751, Update _get_sparse_memory_profile_bytes so it resolves
self.impl.dcp_world_size from get_dcp_group().world_size when the value is still
unset, before evaluating the early-return condition. Ensure the resolved DCP
world size is used during the attn_metadata=None profile pass so sparse-MLA
allocation is included.
| @contextmanager | ||
| def use_workspace_lane(lane: int) -> Iterator[None]: | ||
| """Select an independent workspace lane for the current execution context.""" | ||
| if lane < 0: | ||
| raise ValueError(f"Workspace lane must be non-negative, got {lane}.") | ||
| token = _workspace_lane.set(lane) | ||
| try: | ||
| yield | ||
| finally: | ||
| _workspace_lane.reset(token) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Add Args: and Raises: sections to the docstring.
As per coding guidelines, Python code must use Google-style docstrings with Args:/Returns:/Raises: sections. Please add the missing sections to document the lane argument and the ValueError.
📝 Proposed docstring update
`@contextmanager`
def use_workspace_lane(lane: int) -> Iterator[None]:
- """Select an independent workspace lane for the current execution context."""
+ """Select an independent workspace lane for the current execution context.
+
+ Args:
+ lane: The workspace lane index to select.
+
+ Raises:
+ ValueError: If the provided lane is negative.
+ """
if lane < 0:📝 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.
| @contextmanager | |
| def use_workspace_lane(lane: int) -> Iterator[None]: | |
| """Select an independent workspace lane for the current execution context.""" | |
| if lane < 0: | |
| raise ValueError(f"Workspace lane must be non-negative, got {lane}.") | |
| token = _workspace_lane.set(lane) | |
| try: | |
| yield | |
| finally: | |
| _workspace_lane.reset(token) | |
| `@contextmanager` | |
| def use_workspace_lane(lane: int) -> Iterator[None]: | |
| """Select an independent workspace lane for the current execution context. | |
| Args: | |
| lane: The workspace lane index to select. | |
| Raises: | |
| ValueError: If the provided lane is negative. | |
| """ | |
| if lane < 0: | |
| raise ValueError(f"Workspace lane must be non-negative, got {lane}.") | |
| token = _workspace_lane.set(lane) | |
| try: | |
| yield | |
| finally: | |
| _workspace_lane.reset(token) |
🤖 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/worker/workspace.py` around lines 35 - 44, Update the
use_workspace_lane docstring with Google-style Args and Raises sections:
document the lane integer as the workspace lane selected for the current
execution context, and document that ValueError is raised when lane is negative.
Source: Coding guidelines
1 participant
GLM-5.2's post-RMSNorm 512-dim kv_c latent spans ~240× across layers
(max_abs 0.02→5.2). At the writer's default outer scale of 1.0, shallow
layers quantize with E4M3-subnormal block scales; the error is amplified
downstream.
s_l = max_abs(kv_c_normed)/(6·448)puts every layer's largestblock scale at the top of the E4M3 range.
Adds
kv-scales/glm52-nvfp4-nf3-hybrid_mla_outer_scales_v1.json(format
nvfp4_ds_mla_outer_scale_v1; captured onmadeby561/GLM-5.2-MXFP8-NVFP4-NF3-Hybrid, wikitext-2 2048-ctx TP4, per-rank
agreement checked, per-layer envelope with an independent capture of the same
base for shallow-layer headroom;
max_absretained for audit) and adefault-off
VLLM_NVFP4_MLA_SCALES_FILEknob in serve-glm52.sh.KLD (rtx6kpro glm52-kld-evaluation protocol, 5 fresh boots each, 4×96GB):
Requires a b12x build with the latent_scale identity/dynamic cache fact
(local-inference-lab/sparkinfer#52) or a one-time compile-cache clear when first enabling.
🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Documentation
Tests