fix(mxfp8): support tail-padded linear output - #51
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (1)
Walkthrough
ChangesMXFP8 tail-padded outputs
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant mxfp8_linear
participant mxfp8_linear_fused_op
participant dense_gemm
mxfp8_linear->>mxfp8_linear_fused_op: pass tail_padding_bytes
mxfp8_linear_fused_op->>mxfp8_linear_fused_op: allocate tail-padded output
mxfp8_linear_fused_op->>dense_gemm: compute with out= output view
dense_gemm-->>mxfp8_linear_fused_op: write GEMM results
mxfp8_linear_fused_op-->>mxfp8_linear: return output view
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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 |
![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
🧹 Nitpick comments (2)
b12x/gemm/mxfp8_linear.py (1)
237-246: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueReplicate padded storage in the fake implementation.
The fake implementation currently ignores
tail_padding_bytesand returns an unpaddedtorch.emptytensor. Duringtorch.compiletracing, it is best practice for fake tensors to faithfully model the underlying physical storage size so that downstream memory planners or constraint checks behave correctly.♻️ Proposed refactor
def _mxfp8_linear_fused_fake( source_2d: torch.Tensor, weight_values: torch.Tensor, weight_scale_rows: torch.Tensor, weight_scale_mma: torch.Tensor, in_features: int, padded_in_features: int, out_features: int, expected_m: int, stream_int: int | None, tail_padding_bytes: int, ) -> torch.Tensor: - del stream_int, tail_padding_bytes + del stream_int del weight_values, weight_scale_rows, weight_scale_mma del in_features, padded_in_features, expected_m + if tail_padding_bytes > 0: + return _tail_padded_output( + (source_2d.shape[0], out_features), + dtype=source_2d.dtype, + device=source_2d.device, + tail_padding_bytes=tail_padding_bytes, + ) return torch.empty( (source_2d.shape[0], out_features), dtype=source_2d.dtype, device=source_2d.device, )b12x/distributed/pcie_oneshot.py (1)
1248-1253: 🩺 Stability & Availability | 🔵 TrivialVerify pool shutdown teardown ordering across ranks.
rollback_channelsroutes closes through_coordinated_close_channels, which barriers so every rank closes IPC imports (unmaps peer handles) before any rank frees its exports (cudaFree). Both pools'close()instead close each channel directly, freeing local exports right after closing local imports with no peer barrier — the exact ordering hazard the coordinated path prevents. This is likely intentional (invoking collectives fromclose()/__del__during GC could deadlock), so please confirm the shutdown contract guarantees peers no longer map these buffers by the timeclose()frees them.
b12x/distributed/pcie_oneshot.py#L1248-L1253: confirm uncoordinatedchannel.close()at pool shutdown is safe, or route through a coordinated path when a live exchange group exists.b12x/distributed/pcie_dcp_a2a.py#L779-L783: same confirmation for the DCP pool'sclose().
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: ce036df5-bc80-426b-92b7-4af78dbf8884
📒 Files selected for processing (11)
b12x/distributed/pcie_dcp_a2a.pyb12x/distributed/pcie_oneshot.pyb12x/gemm/mxfp8_linear.pyb12x/integration/__init__.pyb12x/integration/tp_moe.pytests/distributed/test_pcie_dcp_a2a.pytests/distributed/test_pcie_oneshot.pytests/test_gemm_mxfp8_linear.pytests/test_moe_execution_model.pytests/test_tp_moe_scratch_bindings.pytests/test_w4a8_tp_moe.py
| def _use_barrier_free_nvfp4_split( | ||
| *, | ||
| quant_mode: str, | ||
| num_tokens: int, | ||
| activation: str, | ||
| ) -> bool: | ||
| """Return whether native NVFP4 decode uses separate FC1/FC2 launches. | ||
|
|
||
| The split path remains available for diagnostics, but it loses the serving | ||
| benefit of the bounded micro launch overlapping GLM shared experts. | ||
| """ | ||
| return ( | ||
| _is_native_nvfp4_micro_decode( | ||
| quant_mode=quant_mode, | ||
| num_tokens=num_tokens, | ||
| activation=activation, | ||
| ) | ||
| and os.environ.get("B12X_NVFP4_SPLIT_DECODE", "0") == "1" | ||
| ) | ||
|
|
||
|
|
||
| def tp_moe_plan_supports_aux_stream_overlap(plan: TPMoEPlan) -> bool: | ||
| """Whether a planned MoE launch can safely overlap unrelated CUDA work.""" | ||
| if not isinstance(plan, TPMoEPlan): | ||
| raise TypeError("plan must be a TPMoEPlan") | ||
| if plan.implementation != "micro": | ||
| return False | ||
| if plan.num_topk <= 0 or plan.routed_rows % plan.num_topk != 0: | ||
| return False | ||
| # The compact native-NVFP4 micro grid is a bounded, single-wave decode | ||
| # launch. M=1..7 survives graph-replay concurrency stress; M=8 changes the | ||
| # launch geometry, so that boundary and all larger resident plans must | ||
| # remain serialized. | ||
| return plan.routed_rows // plan.num_topk <= 7 and _is_native_nvfp4_micro_decode( | ||
| quant_mode=plan.quant_mode, | ||
| num_tokens=plan.routed_rows // plan.num_topk, | ||
| activation=plan.activation, | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
tp_moe_plan_supports_aux_stream_overlap doesn't account for barrier-free split decode, contradicting its own sibling docstring.
_use_barrier_free_nvfp4_split's docstring states the split path "loses the serving benefit of the bounded micro launch overlapping GLM shared experts," implying overlap safety/benefit is tied to the fused (non-split) launch. However tp_moe_plan_supports_aux_stream_overlap (the exported, public gating API) never checks B12X_NVFP4_SPLIT_DECODE/_use_barrier_free_nvfp4_split, so it will still report True for a plan even when the opt-in split path is active for that same M-bounded band. No test exercises this combination: test_native_nvfp4_micro_allows_aux_stream_overlap and test_aux_stream_overlap_rejects_large_resident_launches both run with split decode unset/disabled, and test_native_nvfp4_fused_micro_graph_replay_with_aux_stream_work explicitly pins B12X_NVFP4_SPLIT_DECODE=0 before capturing/replaying under aux-stream load. A caller that relies on this predicate while split decode is opted in would get an overlap-safe signal for a launch pattern the code's own comments say loses that safety/benefit.
🔧 Suggested fix
def tp_moe_plan_supports_aux_stream_overlap(plan: TPMoEPlan) -> bool:
"""Whether a planned MoE launch can safely overlap unrelated CUDA work."""
if not isinstance(plan, TPMoEPlan):
raise TypeError("plan must be a TPMoEPlan")
if plan.implementation != "micro":
return False
if plan.num_topk <= 0 or plan.routed_rows % plan.num_topk != 0:
return False
+ if os.environ.get("B12X_NVFP4_SPLIT_DECODE", "0") == "1":
+ return False
# The compact native-NVFP4 micro grid is a bounded, single-wave decode
# launch. M=1..7 survives graph-replay concurrency stress; M=8 changes the
# launch geometry, so that boundary and all larger resident plans must
# remain serialized.
return plan.routed_rows // plan.num_topk <= 7 and _is_native_nvfp4_micro_decode(
quant_mode=plan.quant_mode,
num_tokens=plan.routed_rows // plan.num_topk,
activation=plan.activation,
)📝 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.
| def _use_barrier_free_nvfp4_split( | |
| *, | |
| quant_mode: str, | |
| num_tokens: int, | |
| activation: str, | |
| ) -> bool: | |
| """Return whether native NVFP4 decode uses separate FC1/FC2 launches. | |
| The split path remains available for diagnostics, but it loses the serving | |
| benefit of the bounded micro launch overlapping GLM shared experts. | |
| """ | |
| return ( | |
| _is_native_nvfp4_micro_decode( | |
| quant_mode=quant_mode, | |
| num_tokens=num_tokens, | |
| activation=activation, | |
| ) | |
| and os.environ.get("B12X_NVFP4_SPLIT_DECODE", "0") == "1" | |
| ) | |
| def tp_moe_plan_supports_aux_stream_overlap(plan: TPMoEPlan) -> bool: | |
| """Whether a planned MoE launch can safely overlap unrelated CUDA work.""" | |
| if not isinstance(plan, TPMoEPlan): | |
| raise TypeError("plan must be a TPMoEPlan") | |
| if plan.implementation != "micro": | |
| return False | |
| if plan.num_topk <= 0 or plan.routed_rows % plan.num_topk != 0: | |
| return False | |
| # The compact native-NVFP4 micro grid is a bounded, single-wave decode | |
| # launch. M=1..7 survives graph-replay concurrency stress; M=8 changes the | |
| # launch geometry, so that boundary and all larger resident plans must | |
| # remain serialized. | |
| return plan.routed_rows // plan.num_topk <= 7 and _is_native_nvfp4_micro_decode( | |
| quant_mode=plan.quant_mode, | |
| num_tokens=plan.routed_rows // plan.num_topk, | |
| activation=plan.activation, | |
| ) | |
| def _use_barrier_free_nvfp4_split( | |
| *, | |
| quant_mode: str, | |
| num_tokens: int, | |
| activation: str, | |
| ) -> bool: | |
| """Return whether native NVFP4 decode uses separate FC1/FC2 launches. | |
| The split path remains available for diagnostics, but it loses the serving | |
| benefit of the bounded micro launch overlapping GLM shared experts. | |
| """ | |
| return ( | |
| _is_native_nvfp4_micro_decode( | |
| quant_mode=quant_mode, | |
| num_tokens=num_tokens, | |
| activation=activation, | |
| ) | |
| and os.environ.get("B12X_NVFP4_SPLIT_DECODE", "0") == "1" | |
| ) | |
| def tp_moe_plan_supports_aux_stream_overlap(plan: TPMoEPlan) -> bool: | |
| """Whether a planned MoE launch can safely overlap unrelated CUDA work.""" | |
| if not isinstance(plan, TPMoEPlan): | |
| raise TypeError("plan must be a TPMoEPlan") | |
| if plan.implementation != "micro": | |
| return False | |
| if plan.num_topk <= 0 or plan.routed_rows % plan.num_topk != 0: | |
| return False | |
| if os.environ.get("B12X_NVFP4_SPLIT_DECODE", "0") == "1": | |
| return False | |
| # The compact native-NVFP4 micro grid is a bounded, single-wave decode | |
| # launch. M=1..7 survives graph-replay concurrency stress; M=8 changes the | |
| # launch geometry, so that boundary and all larger resident plans must | |
| # remain serialized. | |
| return plan.routed_rows // plan.num_topk <= 7 and _is_native_nvfp4_micro_decode( | |
| quant_mode=plan.quant_mode, | |
| num_tokens=plan.routed_rows // plan.num_topk, | |
| activation=plan.activation, | |
| ) |
b3ec4e4 to
b571d6b
Compare
|
I'm willing to add a feature where the caller can provide their own output tensor if you want to work around this in vllm, but adding a feature to this library to compensate for a cublas bug doesn't seem right. |
2 participants
Summary
Add an optional
tail_padding_bytesargument to the native MXFP8 linear path and write the dense GEMM result directly into output storage with a mapped tail.Why
The vLLM GLM MLA path feeds some linear outputs into SM120 cuBLAS strided BMMs. The guarded repro showed cuBLAS may read up to the next 64 KiB boundary past the logical BF16 tensor end for the relevant
_v_up_projshape. Normal PyTorch allocations usually hide this with allocator slack; tight custom/provided buffers can fault.This change lets vLLM request a small mapped tail for MXFP8 linear outputs without adding a separate hot-path
.clone()/.contiguous()copy.Changes
mxfp8_linear.dense_gemm(..., out=...).tail_padding_bytes=0.Validation
python -m pytest -q tests/test_gemm_mxfp8_linear.py -k "test_mxfp8_linear_writes_directly_to_tail_padded_output"in a CUDA 13.2 diagnostic container: 1 passed.[6, 64], tail bytes65536, bitwise match against the non-tail-padded path.Summary by CodeRabbit
tail_padding_bytes, reserving extra output storage.tail_padding_bytes.