Conversation
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
…-trip (SM90 deepep capture crash) Since sgl-project#30924 the masked plain-silu path emits the down-gemm activation scale already MN-major TMA-aligned. On SM90 (DEEPGEMM_NEED_TMA_ALIGNED_SCALES) the runner rebinds it through deep_gemm's get_mn_major_tma_aligned_tensor, whose short-circuit branch returns a NON-owning alias across the tvm-ffi boundary (entry-time torch::from_blob borrow, re-exported as the result); the rebind then drops the scale's only owner, freeing its storage mid-forward. The next allocation reuses the block and the down gemm converts a dangling pointer, failing decode CUDA-graph capture with "the specified pointer resides on host memory" (base-c-test-deepep-4-gpu-h100, 100% deterministic since sgl-project#30924; B200/H200 take the UE8M0 branch and never rebind). Route both masked-gemm scale rebinds through a keep-owner wrapper: call the transform, and when it returns an alias of the input (same data_ptr while both are alive), keep the owning original instead of the alias. This detects the aliasing outcome rather than mirroring deep_gemm's alignment condition, so no deep_gemm layout logic is duplicated in sglang and wheel-side condition changes cannot silently reintroduce the use-after-free. Row-major scales keep taking the owned copy path unchanged.
b6d084a to
8b6aefb
Compare
|
/tag-run-ci-label |
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
|
/tag-run-ci-label |
|
/tag-and-rerun-ci |
|
/rerun-test test_deepep_small.py |
|
Results for 🚀 |
|
/rerun-test registered/ep/test_deepep_small.py::TestTP.test_gsm8k registered/ep/test_deepep_small.py::TestTBOWithTPAttnAndDenseDP.test_gsm8k |
|
Results for 🚀 |
|
Ran a few discriminating tests on this PR (results: full file, class-level):
So this PR looks correct but not sufficient to turn base-c-test-deepep-4-gpu-h100 green; the NaN needs a separate fix. |
@hnyls2002 Thanks for verifying. Seems multiple tests are failing, do you feel reverting #30924 might be a better move? |
|
@Jialin Yes, but too many conflicts... |
|
Will close the PR for now, and expect others would follow up with a proper fix on it. IIUC, currently the unittest is skipped. |
2 participants
Motivation
Since #30924 the masked plain-silu path emits the down-gemm activation scale already MN-major TMA-aligned. On SM90 (
DEEPGEMM_NEED_TMA_ALIGNED_SCALES) the runner still rebinds it through deep_gemm'sget_mn_major_tma_aligned_tensor, whose short-circuit branch returns a NON-owning alias across the tvm-ffi boundary; the rebind then drops the scale's only owner, freeing its storage mid-forward. The next allocation reuses the block and the down gemm converts a dangling pointer, failing decode CUDA-graph capture withRuntimeError: The specified pointer resides on host memory and is not registered with any CUDA device(base-c-test-deepep-4-gpu-h100, 100% deterministic since #30924; B200/H200 take the UE8M0 branch and never rebind). On allocators that keep the freed VA mapped, the same bug silently computes the down gemm with garbage scales instead of crashing.Evidence for the diagnosis:
sgl_deep_gemm-0.1.4.post1x86_64 wheel maps the crash offsetdg_m_grouped_fp8_fp4_gemm_nt_masked+0x473to thea_sf(activation scale) argument conversion; the weights, output, andmasked_mconversions all succeed before it.get_mn_major_tma_aligned_tensoron the post-[JIT] Trait-driven per_token_group_quant: unify the quant kernel family (flat + masked) #30924 layout returns a tensor aliasing the input pointer with no owner, and dropping the Python base frees the block for the next allocation (down_output); the pre-[JIT] Trait-driven per_token_group_quant: unify the quant kernel family (flat + masked) #30924 row-major layout takes the C++ copy path whose output is owned, which is why the identical rebind was safe before.convert_to_torch_tensorrules out the empty-tensor theory: 0-size inputs hit clean shape asserts, not this error.Modifications
python/sglang/srt/layers/moe/moe_runner/deep_gemm.py: add_is_mn_major_tma_aligned, a python-side mirror of deep_gemm's short-circuit condition (stride(-2)==1,stride(-1)==ceil_align(mn,4), batch stride consistent), and guard the SM90 down-gemm rebind with it so an already-aligned scale never round-trips through the aliasing FFI path. Row-major scales keep taking deep_gemm's owned copy path, so behavior there is unchanged.test/registered/unit/layers/moe/test_deep_gemm_masked_scale_layout.py: CPU-only unit tests pinning the mirror to the layoutcreate_per_token_group_quant_fp8_output_scaleactually emits (aligned and non-multiple-of-4 token counts), plus the row-major negative and a batch-stride-mismatch negative, so either side drifting apart turns the suite red.The root defect (tvm-ffi short-circuit returns are non-owning) lives in sgl-project/DeepGEMM and affects every caller of
get_mn_major_tma_aligned_tensor/transform_sf_into_required_layout; this PR is the caller-side fix that unblocks CI, and an upstream ownership fix is being filed separately.Accuracy Tests
No numeric change: the guard only skips a transform that would have returned the same layout (deep_gemm's own short-circuit), and it also eliminates the silent-garbage-scale failure mode on allocators that keep the freed VA mapped. Validation is
test/registered/ep/test_deepep_small.pyon thebase-c-test-deepep-4-gpu-h100job, which fails deterministically on current main and passes with this change; the new CPU unit tests pass locally.Speed Tests and Profiling
Neutral to slightly positive: removes one host-side FFI round-trip per masked down-gemm on SM90. No kernel changes.
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): ❌ Run #29940954181
Latest PR Test (Extra): ❌ Run #29940953575