[fix] Build DeepGEMM against selected torch

.claude/docs/weight_sync.md

@@ -297,16 +297,16 @@ trainer-send engines.
 The `sharded_rdt` backend is unaffected: it vendors its own trainer-side ABCs in
 `sharded_rdt_base.py`, which already match 0.28.0's `Generic[TTrainerInitInfo]` shape.
 
-### DeepGEMM is unavailable under the torch override
+### DeepGEMM follows the selected torch runtime
 
 vLLM 0.28.0's metadata pins `torch==2.13.0`; we override torch to `2.11.0` because the CUDA
 extension wheels we build against (flash-attn, causal-conv1d, mamba-ssm, transformer-engine)
 have no 2.13 builds. vLLM's main extensions are stable-libtorch-ABI and load fine, but
 `vllm/third_party/deep_gemm/_C` is a version-specific `cpython-312` build linked against torch
-2.13's `c10` and fails with an undefined-symbol `ImportError`. vLLM catches this and logs
-`Module vllm.third_party.deep_gemm was found but failed to import`, then falls back, so the
-engine runs — but the DeepGEMM fused-MoE and sparse-attention-indexer paths are gone. Revisit
-when those wheels publish torch 2.13 builds.
+2.13's `c10` and fails with an undefined-symbol `ImportError`. SkyRL therefore pins the matching
+DeepGEMM source revision as an external package and builds it against the selected torch runtime.
+vLLM prefers that package over its incompatible vendored extension, keeping the DeepGEMM
+fused-MoE and sparse-attention-indexer paths available.
 
 ## Gotchas
 

ci/gpu_ci_run_h100.sh

@@ -12,5 +12,6 @@ uv run --directory . --isolated --extra dev --extra fsdp pytest -s -vvv -m h100
 
 # Run Megatron h100 tests.
 uv run --directory . --isolated --extra dev --extra megatron pytest -s -vvv -m h100 \
+    tests/backends/skyrl_train/gpu/gpu_ci/test_deep_gemm.py \
     tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_models.py \
     tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_router_replay.py

pyproject.toml

@@ -150,6 +150,7 @@ megatron = [
     "causal-conv1d==1.6.2.post1+cu.13.0.torch.2.11; sys_platform == 'linux'",
     "mamba-ssm==2.3.2.post1+cu.13.0.torch.2.11; sys_platform == 'linux'",
     "vllm==0.28.0; sys_platform == 'linux'",
+    "deep-gemm; sys_platform == 'linux'",
     "vllm-router; sys_platform == 'linux'",
     # The `nixl` shim provides that namespace and dispatches on `torch.version.cuda`,
     # so with a cu13 torch it loads `nixl_cu13`. `nixl-cu13` ships the `nixl_cu13`
@@ -290,9 +291,9 @@ override-dependencies = [
     # `_moe_C_stable_libtorch.abi3.so`), so its main binaries load against 2.11 and
     # we hold torch there rather than churn every CUDA extension pin. The one
     # exception is `vllm/third_party/deep_gemm/_C`, a version-specific cpython-312
-    # build linked against torch 2.13's c10: it fails to import and vLLM falls back,
-    # so the DeepGEMM fused-MoE / sparse-attn-indexer paths are unavailable. See
-    # `.claude/docs/weight_sync.md`. Revisit once those wheels publish 2.13 builds.
+    # build linked against torch 2.13's c10. The external DeepGEMM source below is
+    # rebuilt against this selected torch runtime, and vLLM prefers it over the
+    # incompatible vendored extension.
     "torch==2.11.0; sys_platform == 'linux'",
     # torchvision is NOT stable-ABI (unlike vLLM), and each release hard-pins one
     # torch: 0.28.0 wants torch 2.13.0, 0.26.0 wants 2.11.0. vLLM 0.28.0 pulls
@@ -302,10 +303,12 @@ override-dependencies = [
 ]
 
 [tool.uv.extra-build-dependencies]
+deep-gemm = ["torch==2.11.0", "ninja"]
 flash-attn = [{requirement = "torch", match-runtime = true}]
 fast-hadamard-transform = ["torch==2.11.0", "ninja"]
 
 [tool.uv.extra-build-variables]
+deep-gemm = { DG_FORCE_BUILD = "1" }
 flash-attn = { FLASH_ATTENTION_SKIP_CUDA_BUILD = "TRUE"}
 fast-hadamard-transform = { FAST_HADAMARD_TRANSFORM_FORCE_BUILD = "TRUE"}
 
@@ -348,6 +351,7 @@ explicit = true
 
 [tool.uv.sources]
 skyrl-gym = { path = "./skyrl-gym", editable = true }
+deep-gemm = { git = "https://github.com/deepseek-ai/DeepGEMM", rev = "8b1392b978f5a03c828dd1711090d7fb50958b8a" }
 # Match torch's CUDA variant (cu130).
 flashinfer-jit-cache = { index = "flashinfer-cu130", marker = "sys_platform == 'linux'" }
 flashinfer-cubin = { index = "flashinfer", marker = "sys_platform == 'linux'" }

tests/backends/skyrl_train/gpu/gpu_ci/test_deep_gemm.py

@@ -0,0 +1,8 @@
+import pytest
+
+
+@pytest.mark.h100
+def test_deep_gemm_loads_against_selected_torch():
+    import deep_gemm
+
+    assert deep_gemm.get_num_sms() > 0