Fix B12X NVFP4 prepare for forced W4A16 - #20
Merged
lukealonso merged 1 commit intoJun 16, 2026
Merged
Conversation
Chthonic's native B12X NVFP4 prepare path is correct for normal NVFP4, but forced W4A16 reuses the prepared ModelOpt weights through B12xExperts. That path needs the FI/B12X gated W13 reorder before b12x prepares the W4A16 weights; otherwise the gate/up halves are consumed in the wrong order and GLM emits corrupted output. Keep native B12X NVFP4 on prepare_nvfp4_moe_layer_for_b12x, switch to the FI/B12X prepare only when use_a16 or B12X_MOE_FORCE_A16=1 is active, and propagate ModelOpt's explicit use_a16 flag into the converter. Validated on the chthonic 5e83948 / b12x 465cb6e image with GLM v10 DCP1 MTP-off: B12X_MOE_FORCE_A16=1 on port 5329 and B12X_MOE_FORCE_A16=0 on port 5330 both start with B12X MLA sparse + B12X NvFp4 MoE and run /mnt/test.py -L coherently until timeout with CJK count 0. Also ran git diff --check and py_compile for the touched files.
📝 WalkthroughWalkthroughAdds a ChangesNVFP4 MoE B12X A16 path selection
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes 🚥 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 docstrings
🧪 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: 1
🤖 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/fused_moe/oracle/nvfp4.py`:
- Around line 347-390: Update the `make_nvfp4_moe_quant_config()` function to
apply the same A16 logic as the weight preparation code. The weight preparation
at line 347 checks `use_a16 or _env_flag("B12X_MOE_FORCE_A16")` to decide
between two backends, but the quant config function only checks `use_a16`,
ignoring the environment flag override. Modify `make_nvfp4_moe_quant_config()`
to include the same `_env_flag("B12X_MOE_FORCE_A16")` check when determining the
quantization method, ensuring the returned config reflects the same A16
determination as the prepared weights.
🪄 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: 0c8cc6fb-4636-4ba7-a01c-dbf72b894a84
📒 Files selected for processing (2)
vllm/model_executor/layers/fused_moe/oracle/nvfp4.pyvllm/model_executor/layers/quantization/modelopt.py
Comment on lines
+347
to
+390
| if use_a16 or _env_flag("B12X_MOE_FORCE_A16"): | ||
| ( | ||
| w13, | ||
| w13_scale, | ||
| w13_scale_2, | ||
| a13_scale, | ||
| w2, | ||
| w2_scale, | ||
| w2_scale_2, | ||
| a2_scale, | ||
| ) = prepare_nvfp4_moe_layer_for_fi_or_cutlass( | ||
| backend=NvFp4MoeBackend.FLASHINFER_B12X, | ||
| layer=layer, | ||
| w13=w13, | ||
| w13_scale=w13_scale, | ||
| w13_scale_2=w13_scale_2, | ||
| a13_scale=a13_scale, | ||
| w2=w2, | ||
| w2_scale=w2_scale, | ||
| w2_scale_2=w2_scale_2, | ||
| a2_scale=a2_scale, | ||
| is_act_and_mul=is_act_and_mul, | ||
| ) | ||
| else: | ||
| ( | ||
| w13, | ||
| w13_scale, | ||
| w13_scale_2, | ||
| a13_scale, | ||
| w2, | ||
| w2_scale, | ||
| w2_scale_2, | ||
| a2_scale, | ||
| ) = prepare_nvfp4_moe_layer_for_b12x( | ||
| w13=w13, | ||
| w13_scale=w13_scale, | ||
| w13_scale_2=w13_scale_2, | ||
| a13_scale=a13_scale, | ||
| w2=w2, | ||
| w2_scale=w2_scale, | ||
| w2_scale_2=w2_scale_2, | ||
| a2_scale=a2_scale, | ||
| is_act_and_mul=is_act_and_mul, | ||
| ) |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== Locate all A16 gating points ==="
rg -n "B12X_MOE_FORCE_A16|use_a16|make_nvfp4_moe_quant_config|nvfp4_w4a16_moe_quant_config|convert_to_nvfp4_moe_kernel_format" \
vllm/model_executor/layers/fused_moe/oracle/nvfp4.py \
vllm/model_executor/layers/quantization/modelopt.py \
vllm/model_executor/layers/fused_moe/b12x_moe.py
echo
echo "=== Inspect B12X experts implementation for force-A16 behavior ==="
sed -n '1,260p' vllm/model_executor/layers/fused_moe/b12x_moe.pyRepository: local-inference-lab/vllm
Length of output: 9754
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== Extract convert_to_nvfp4_moe_kernel_format function (lines 323-400) ==="
sed -n '323,400p' vllm/model_executor/layers/fused_moe/oracle/nvfp4.py | cat -n
echo
echo "=== Extract make_nvfp4_moe_quant_config function (lines 501-525) ==="
sed -n '501,525p' vllm/model_executor/layers/fused_moe/oracle/nvfp4.py | cat -n
echo
echo "=== Extract b12x_moe.py context around line 563 (lines 555-575) ==="
sed -n '555,575p' vllm/model_executor/layers/fused_moe/b12x_moe.py | cat -n
echo
echo "=== Search for callers of convert_to_nvfp4_moe_kernel_format ==="
rg -B 2 -A 2 "convert_to_nvfp4_moe_kernel_format\(" vllm/model_executor/layers/quantization/modelopt.pyRepository: local-inference-lab/vllm
Length of output: 5183
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== Check modelopt.py around lines 1567-1614 to see the full call sequence ==="
sed -n '1560,1620p' vllm/model_executor/layers/quantization/modelopt.py | cat -n
echo
echo "=== Search for where make_nvfp4_moe_quant_config result is used ==="
rg -A 3 "make_nvfp4_moe_quant_config" vllm/model_executor/layers/quantization/modelopt.py | head -40Repository: local-inference-lab/vllm
Length of output: 3344
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== Check if quant_config is actually used or just overridden by _quant_mode ==="
rg -B 5 -A 5 "self._quant_mode\(\)|_quant_mode|self.moe_quant_config" vllm/model_executor/layers/fused_moe/b12x_moe.py | head -80
echo
echo "=== Find where quant_mode is set in plan_b12x_moe_fp4_scratch ==="
rg -B 2 -A 2 "quant_mode=" vllm/model_executor/layers/fused_moe/b12x_moe.py | head -40Repository: local-inference-lab/vllm
Length of output: 4383
Verify A16 gating is consistent between tensor preparation and quant config.
Line 347 enables weight preparation as A16 when use_a16 or B12X_MOE_FORCE_A16=1, but make_nvfp4_moe_quant_config() at lines 514–517 only checks use_a16, ignoring the environment flag. This creates a split where:
- Weights are prepared as A16 when the env flag is set
FusedMoEQuantConfigobject returned does not reflect this override- Runtime
_quant_mode()in b12x_moe.py correctly forces w4a16
Any code path relying on quant_config.quant_method to determine behavior will miss the B12X_MOE_FORCE_A16 override, causing prepared weights and quant configuration to diverge.
🤖 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/fused_moe/oracle/nvfp4.py` around lines 347 - 390,
Update the `make_nvfp4_moe_quant_config()` function to apply the same A16 logic
as the weight preparation code. The weight preparation at line 347 checks
`use_a16 or _env_flag("B12X_MOE_FORCE_A16")` to decide between two backends, but
the quant config function only checks `use_a16`, ignoring the environment flag
override. Modify `make_nvfp4_moe_quant_config()` to include the same
`_env_flag("B12X_MOE_FORCE_A16")` check when determining the quantization
method, ensuring the returned config reflects the same A16 determination as the
prepared weights.
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
Fix GLM 5.1 NVFP4 corruption when
B12X_MOE_FORCE_A16=1forces the B12X MoE path to prepare ModelOpt NVFP4 experts as W4A16.The normal native B12X NVFP4 prepare path is correct for A16-off. The forced W4A16 path, however, needs the existing FI/B12X gated W13 reorder before
B12xExpertsprepares W4A16 weights. Without that reorder, the gate/up halves are consumed in the wrong order and GLM emits corrupted output immediately.This PR keeps native B12X NVFP4 on
prepare_nvfp4_moe_layer_for_b12x, switches to the FI/B12X prepare only whenuse_a16orB12X_MOE_FORCE_A16=1is active, and propagates ModelOpt's explicituse_a16flag into the converter.Validation
Tested on
voipmonitor/vllm:chthonic-consecration-5e83948-b12x-465cb6e-cu132with only these fix files bind-mounted:B12X_MOE_FORCE_A16=1.B12X_MOE_FORCE_A16=0.B12X_MLA_SPARSE, andUsing 'B12X' NvFp4 MoE backend./mnt/test.py -Luntil timeout with coherent output andCJK characters in output: 0.git diff --checkpassed.python3 -m py_compile vllm/model_executor/layers/fused_moe/oracle/nvfp4.py vllm/model_executor/layers/quantization/modelopt.pypassed.Summary by CodeRabbit