diff --git a/qwen3_6_scripts/patch_ops.sh b/qwen3_6_scripts/patch_ops.sh index ac27d488..7604c881 100755 --- a/qwen3_6_scripts/patch_ops.sh +++ b/qwen3_6_scripts/patch_ops.sh @@ -246,12 +246,14 @@ if source != installed: raise SystemExit("runtime api_server overlay identity mismatch") PY -build_stage "compiling CoreX MoE index+combine kernel" +build_stage "compiling CoreX CUDA extensions (moe_index_combine + gdn_chunk_recurrent)" if [[ -x /usr/local/corex-3.2.3/bin/clang++ ]]; then bash ./build_corex_moe_index_combine.sh "${VLLM_ROOT}" || \ echo "[WARN] moe_index_combine build failed — will use PyTorch fallback" + bash ./build_corex_gdn_chunk_recurrent.sh "${VLLM_ROOT}" || \ + echo "[WARN] gdn_chunk_recurrent build failed — will use Python fallback" else - echo "[WARN] corex clang++ not found — skipping moe_index_combine build" + echo "[WARN] corex clang++ not found — skipping extension builds" fi build_stage "compiling submission Python sources" diff --git a/qwen3_6_scripts/qwen3_5.py b/qwen3_6_scripts/qwen3_5.py index 121076fd..93f634c6 100644 --- a/qwen3_6_scripts/qwen3_5.py +++ b/qwen3_6_scripts/qwen3_5.py @@ -148,6 +148,13 @@ try: except ImportError: _corex_moe_index_combine = None +try: + from vllm import corex_gdn_chunk_recurrent as _corex_gdn_chunk_recurrent +except ImportError: + _corex_gdn_chunk_recurrent = None + +_HAS_COREX_GDN_CHUNK = _corex_gdn_chunk_recurrent is not None + from vllm.model_executor.models.interfaces import (HasInnerState, SupportsLoRA, SupportsMultiModal) @@ -1105,9 +1112,14 @@ class GatedDeltaNet(nn.Module): seq_len, _DNN_CHUNK_SIZE, seq_capture_offsets | seq_segment_offsets) sc_start = 0 + _chunk_fn = ( + _corex_gdn_chunk_recurrent.torch_chunk_gated_delta_rule + if _HAS_COREX_GDN_CHUNK + else _torch_chunk_gated_delta_rule + ) with bi100_timer(f"L{self.layer_idx}.gdn.prefill"): for sc_end in segment_ends: - c_out, cur_state = _torch_chunk_gated_delta_rule( + c_out, cur_state = _chunk_fn( q[:, sc_start:sc_end], k[:, sc_start:sc_end], v[:, sc_start:sc_end],