From e78fa560c88d573677ca95a63a5ec8331fda27c5 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 13 Aug 2026 06:25:09 +0000 Subject: [PATCH] feat: wire corex_gdn_chunk_recurrent C++ kernel into GDN prefill path - patch_ops.sh: build corex_gdn_chunk_recurrent.so alongside moe_index_combine - qwen3_5.py: import corex_gdn_chunk_recurrent, use C++ version for prefill chunks instead of Python _torch_chunk_gated_delta_rule - C++ version from xllm upstream avoids Python loop overhead and has proper fp32 accumulation (key for NaN prevention on BI-V100) - Falls back to Python version if .so not available --- qwen3_6_scripts/patch_ops.sh | 6 ++++-- qwen3_6_scripts/qwen3_5.py | 14 +++++++++++++- 2 files changed, 17 insertions(+), 3 deletions(-) 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],