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
This commit is contained in:
Claude
2026-08-13 06:25:09 +00:00
parent c720cbc3a3
commit e78fa560c8
2 changed files with 17 additions and 3 deletions

View File

@@ -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"

View File

@@ -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],