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:
@@ -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"
|
||||
|
||||
@@ -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],
|
||||
|
||||
Reference in New Issue
Block a user