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")
|
raise SystemExit("runtime api_server overlay identity mismatch")
|
||||||
PY
|
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
|
if [[ -x /usr/local/corex-3.2.3/bin/clang++ ]]; then
|
||||||
bash ./build_corex_moe_index_combine.sh "${VLLM_ROOT}" || \
|
bash ./build_corex_moe_index_combine.sh "${VLLM_ROOT}" || \
|
||||||
echo "[WARN] moe_index_combine build failed — will use PyTorch fallback"
|
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
|
else
|
||||||
echo "[WARN] corex clang++ not found — skipping moe_index_combine build"
|
echo "[WARN] corex clang++ not found — skipping extension builds"
|
||||||
fi
|
fi
|
||||||
|
|
||||||
build_stage "compiling submission Python sources"
|
build_stage "compiling submission Python sources"
|
||||||
|
|||||||
@@ -148,6 +148,13 @@ try:
|
|||||||
except ImportError:
|
except ImportError:
|
||||||
_corex_moe_index_combine = None
|
_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,
|
from vllm.model_executor.models.interfaces import (HasInnerState, SupportsLoRA,
|
||||||
SupportsMultiModal)
|
SupportsMultiModal)
|
||||||
|
|
||||||
@@ -1105,9 +1112,14 @@ class GatedDeltaNet(nn.Module):
|
|||||||
seq_len, _DNN_CHUNK_SIZE,
|
seq_len, _DNN_CHUNK_SIZE,
|
||||||
seq_capture_offsets | seq_segment_offsets)
|
seq_capture_offsets | seq_segment_offsets)
|
||||||
sc_start = 0
|
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"):
|
with bi100_timer(f"L{self.layer_idx}.gdn.prefill"):
|
||||||
for sc_end in segment_ends:
|
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],
|
q[:, sc_start:sc_end],
|
||||||
k[:, sc_start:sc_end],
|
k[:, sc_start:sc_end],
|
||||||
v[:, sc_start:sc_end],
|
v[:, sc_start:sc_end],
|
||||||
|
|||||||
Reference in New Issue
Block a user