[GG] chore(b12x): port integration to renamed package - #246
Conversation
📝 WalkthroughWalkthroughThe change replaces legacy SparkInfer package references with B12X references across runtime integrations, distributed paths, diagnostics, tests, mocks, optional-dependency checks, and environment markers. A contract test rejects remaining legacy runtime references. ChangesB12X runtime integration
Estimated code review effort: 2 (Simple) | ~15 minutes 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 |
Signed-off-by: Martin Vit <martin@voipmonitor.org>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 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 `@tests/test_b12x_package_contract.py`:
- Around line 6-11: Update the legacy_markers check in the test to match the
standalone package token “sparkinfer” as well as the existing “SPARKINFER_”
marker, so root imports such as import sparkinfer and from sparkinfer import ...
are detected while preserving the current uppercase check.
In `@vllm/model_executor/layers/fused_moe/b12x_moe.py`:
- Around line 406-407: Update the dependency declaration for B12X used by the
imports TPMoEScratchCaps and plan_tp_moe_scratch to pin an available artifact
that actually exposes b12x.moe.fused_moe and its private paths. Remove the
unavailable 1.1.0/source assumption and ensure production installation resolves
the intended package version.
In `@vllm/v1/attention/ops/dcp_alltoall.py`:
- Line 200: Update the checkpoint function docstring beginning “Snapshot B12X
DCP pools before a disposable graph capture” to add Google-style Args and
Returns sections, documenting the cp_group parameter and the returned tuple,
while preserving the existing summary text.
🪄 Autofix
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: bb541ab0-2c61-4843-b551-da136dac6920
📒 Files selected for processing (38)
tests/distributed/test_b12x_fused_all_reduce.pytests/model_executor/kernels/test_b12x_mxfp8_linear.pytests/model_executor/layers/test_b12x_moe_warmup.pytests/model_executor/layers/test_sparse_attn_indexer_b12x.pytests/quantization/test_exl3.pytests/quantization/test_exl3_prefill_plan.pytests/test_b12x_package_contract.pytests/v1/attention/test_b12x_mla_fp8_rope_writer.pytests/v1/attention/test_mla_backends.pytests/v1/attention/test_sparse_mla_backends.pyvllm/compilation/b12x_capture.pyvllm/distributed/device_communicators/custom_all_reduce.pyvllm/distributed/parallel_state.pyvllm/envs.pyvllm/model_executor/kernels/attention/b12x_mxfp8_bmm.pyvllm/model_executor/kernels/linear/mxfp4/b12x.pyvllm/model_executor/kernels/linear/mxfp8/b12x.pyvllm/model_executor/kernels/linear/nvfp4/b12x.pyvllm/model_executor/kernels/linear/scaled_mm/b12x.pyvllm/model_executor/kernels/linear/scaled_mm/b12x_tensor.pyvllm/model_executor/layers/attention/mla_attention.pyvllm/model_executor/layers/fused_moe/b12x_ep_moe.pyvllm/model_executor/layers/fused_moe/b12x_moe.pyvllm/model_executor/layers/quantization/exl3.pyvllm/model_executor/layers/quantization/nvfp4_nf3_hybrid.pyvllm/model_executor/layers/sparse_attn_indexer.pyvllm/model_executor/warmup/b12x_sparse_indexer_warmup.pyvllm/models/deepseek_v4/attention.pyvllm/models/deepseek_v4/nvidia/b12x.pyvllm/models/deepseek_v4/nvidia/model.pyvllm/models/deepseek_v4/nvidia/mtp.pyvllm/models/minimax_m3/nvidia/sparse_attention_b12x.pyvllm/v1/attention/backends/b12x_attn.pyvllm/v1/attention/backends/mla/b12x_mla_sparse.pyvllm/v1/attention/backends/mla/indexer.pyvllm/v1/attention/ops/dcp_alltoall.pyvllm/v1/worker/gpu/model_runner.pyvllm/v1/worker/gpu_worker.py
| legacy_markers = ("sparkinfer.", "SPARKINFER_") | ||
| offenders: list[str] = [] | ||
|
|
||
| for source in runtime_root.rglob("*.py"): | ||
| text = source.read_text(encoding="utf-8") | ||
| if any(marker in text for marker in legacy_markers): |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Match the legacy package name, not only dotted submodules.
"sparkinfer." does not match import sparkinfer or from sparkinfer import .... A legacy root-package import can therefore remain in vllm without failing this contract test. Match sparkinfer as a package token while preserving the SPARKINFER_ check.
Proposed fix
+import re
from pathlib import Path
...
- legacy_markers = ("sparkinfer.", "SPARKINFER_")
+ legacy_package = re.compile(
+ r"(?<![A-Za-z0-9_])sparkinfer(?:\b|\.)"
+ )
...
- if any(marker in text for marker in legacy_markers):
+ if legacy_package.search(text) or "SPARKINFER_" in text:📝 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.
| legacy_markers = ("sparkinfer.", "SPARKINFER_") | |
| offenders: list[str] = [] | |
| for source in runtime_root.rglob("*.py"): | |
| text = source.read_text(encoding="utf-8") | |
| if any(marker in text for marker in legacy_markers): | |
| legacy_package = re.compile( | |
| r"(?<![A-Za-z0-9_])sparkinfer(?:\b|\.)" | |
| ) | |
| offenders: list[str] = [] | |
| for source in runtime_root.rglob("*.py"): | |
| text = source.read_text(encoding="utf-8") | |
| if legacy_package.search(text) or "SPARKINFER_" in text: |
🤖 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/test_b12x_package_contract.py` around lines 6 - 11, Update the
legacy_markers check in the test to match the standalone package token
“sparkinfer” as well as the existing “SPARKINFER_” marker, so root imports such
as import sparkinfer and from sparkinfer import ... are detected while
preserving the current uppercase check.
| from b12x.moe.fused_moe import Caps as TPMoEScratchCaps | ||
| from b12x.moe.fused_moe import plan as plan_tp_moe_scratch |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/usr/bin/env bash
set -euo pipefail
python - <<'PY'
from importlib import import_module
from importlib.metadata import version
expected_version = "1.1.0"
actual_version = version("b12x")
if actual_version != expected_version:
raise SystemExit(
f"expected b12x=={expected_version}, found {actual_version}"
)
exports = {
"b12x.moe.ep_moe": ("Caps", "plan", "run", "prepare_expert_map"),
"b12x.moe.fused_moe": (
"Caps",
"plan",
"plan_execution",
"run",
"plan_weights",
"prepare_weights",
),
"b12x.moe._shared.execution": ("PreparedWeightLayout",),
"b12x.attention.compressed_mla": (
"Caps",
"plan",
"run",
"split_chunks_for_contract",
),
"b12x.norm.mhc": (
"DEFAULT_BLOCK_K",
"MULT",
"Caps",
"plan",
"run_pre",
"run_post_pre",
"run_post",
),
"b12x.norm.mhc._impl": (
"MHC_GRAM_BLOCK_H",
"MHC_SOURCE_TILE_H",
"MHC_SUPPORTED_HIDDEN_SIZES",
),
"b12x.attention.paged": (
"Caps",
"compile",
"decode_graph_capacity",
"decode_graph_scratch_envelope",
"extend_graph_capacity",
"plan",
"run",
"verify_graph_capacity",
),
"b12x.attention.varlen": ("create_plan", "plan", "run"),
"b12x.attention.nsa_indexer": (
"plan_paged_schedule",
"uses_paged_schedule",
),
"b12x.comm.pcie": ("OneshotAllReducePool", "DmaAllReduce"),
}
for module_name, names in exports.items():
module = import_module(module_name)
missing = [name for name in names if not hasattr(module, name)]
if missing:
raise SystemExit(f"{module_name} is missing: {missing}")
print("B12X export checks passed")
PYRepository: local-inference-lab/vllm
Length of output: 1099
🌐 Web query:
pypi b12x package versions
💡 Result:
The b12x package is a specialized Python library containing CuTe DSL kernels designed for NVFP4 LLM inference on SM120/SM121 GPU architectures (such as Blackwell-based hardware) [1][2]. The package is available on PyPI, with version 0.30.2 being the most recent version identified [1]. It is important to note that the repository formerly associated with the b12x project has been renamed or superseded by sparkinfer [3]. While the b12x package remains available on PyPI, users should be aware that the active development and evolution of the underlying technology have transitioned to the sparkinfer package [3]. Key details: - Repository/Source: Originally hosted under lukealonso/b12x, now associated with local-inference-lab/sparkinfer [4][3]. - Purpose: Provides a focused set of high-performance kernels (NVFP4 fused MoE GEMM, dense GEMM, paged attention, etc.) for integration with frameworks like sglang and vllm [5][2]. - Dependencies: Requires Python 3.10+, torch >= 2.12, and compatible SM120/SM121 hardware [1][3]. - Installation: Can be installed via pip, though users may need to check if their specific use case is better served by the current sparkinfer package [1][3].
Citations:
- 1: https://pypi.org/project/b12x/
- 2: https://pypi.org/project/b12x/0.12.4/
- 3: https://github.com/local-inference-lab/sparkinfer
- 4: https://github.com/lukealonso/b12x
- 5: https://github.com/lukealonso/b12x/blob/master/README.md
Pin the B12X dependency to an available artifact that exposes the imported symbols.
Public PyPI only reaches b12x==0.30.2; b12x.moe.fused_moe, the private B12X paths, and a 1.1.0 artifact are not available from that release. Resolve this to the intended package version/source before relying on these imports in production.
🤖 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/b12x_moe.py` around lines 406 - 407,
Update the dependency declaration for B12X used by the imports TPMoEScratchCaps
and plan_tp_moe_scratch to pin an available artifact that actually exposes
b12x.moe.fused_moe and its private paths. Remove the unavailable 1.1.0/source
assumption and ensure production installation resolves the intended package
version.
| cp_group: GroupCoordinator, | ||
| ) -> tuple[int, dict[Any, tuple[Any, Any]]]: | ||
| """Snapshot SparkInfer DCP pools before a disposable graph capture.""" | ||
| """Snapshot B12X DCP pools before a disposable graph capture.""" |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Add Google-style sections to the checkpoint docstring.
Line [200] documents a function with a cp_group argument and a tuple return value, but it has no Args: or Returns: section. Add both sections.
Proposed docstring
def checkpoint_b12x_dcp_a2a_channels(
cp_group: GroupCoordinator,
) -> tuple[int, dict[Any, tuple[Any, Any]]]:
- """Snapshot B12X DCP pools before a disposable graph capture."""
+ """Snapshot B12X DCP pools before a disposable graph capture.
+
+ Args:
+ cp_group: DCP group whose registered pools are checkpointed.
+
+ Returns:
+ Group ID and per-pool channel checkpoints.
+ """As per coding guidelines, Python docstrings must use Google-style Args:/Returns:/Raises: sections instead of reStructuredText/Sphinx fields.
📝 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.
| """Snapshot B12X DCP pools before a disposable graph capture.""" | |
| def checkpoint_b12x_dcp_a2a_channels( | |
| cp_group: GroupCoordinator, | |
| ) -> tuple[int, dict[Any, tuple[Any, Any]]]: | |
| """Snapshot B12X DCP pools before a disposable graph capture. | |
| Args: | |
| cp_group: DCP group whose registered pools are checkpointed. | |
| Returns: | |
| Group ID and per-pool channel checkpoints. | |
| """ |
🤖 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/ops/dcp_alltoall.py` at line 200, Update the checkpoint
function docstring beginning “Snapshot B12X DCP pools before a disposable graph
capture” to add Google-style Args and Returns sections, documenting the cp_group
parameter and the returned tuple, while preserving the existing summary text.
Source: Coding guidelines
Signed-off-by: Martin Vit <martin@voipmonitor.org>
|
Superseded by dev/gilded-gnosis commit e2666d9 ( |
1 participant
Summary
B12X 1.1.0 renamed its Python package from
sparkinferback tob12xand moved its public environment prefix fromSPARKINFER_toB12X_. Port the GG runtime and tests to that API so a clean composition with currentlocal-inference-lab/b12xboots without a compatibility alias or source overlay.The change is deliberately mechanical:
b12x.*;B12X_*;This PR does not change algorithms, launch policy, tensor layouts, or numerical behavior.
Paired ownership
Two functional PRs replace files too extensively for a duplicate mechanical edit here:
vllm/model_executor/layers/quantization/exl3.pyand its EXL3 tests.tests/distributed/test_dcp_a2a.py.b12x.attention.compressed_mla.The release composition applies those PRs explicitly and performs a final source scan across both
vllm/andtests/. This keeps ownership visible while guaranteeing that the composed image contains no activesparkinfer.*import orSPARKINFER_*runtime setting.Validation
Tested against current B12X master plus PR #125 in the CUDA 13.2 release environment:
git diff --check;Multi-GPU channel and DSpark concurrency validation is performed on the full declared release composition.