dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

fix(pybind): add py::arg + defaults to corex_gdn_chunk_recurrent

Browse Source 
Python calls: _chunk_fn(q,k,v,g,beta, initial_state=, output_final_state=, use_qk_l2norm_in_kernel=)
C++ had: positional-only (query,key,value,g,beta,chunk_size,initial_state,output_final_state,use_qk_l2norm)

Fix: py::arg() naming + chunk_size=64 default (matches Python fallback).
Re-enable _HAS_COREX_GDN_CHUNK flag.

Rebuild on real machine:
  VLLM_ROOT=/usr/local/corex/lib64/python3/dist-packages/vllm
  bash build_corex_gdn_chunk_recurrent.sh $VLLM_ROOT
Then copy .so to prebuilt/
This commit is contained in:
Claude
2026-08-14 02:02:51 +00:00
parent a4d16d36b8
commit 768d89c31a
2 changed files with 20 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
qwen3_6_scripts/corex_gdn_chunk_recurrent.cu
View File

@@ -270,7 +270,24 @@ std::tuple<torch::Tensor, torch::Tensor> torch_chunk_gated_delta_rule(
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("torch_chunk_gated_delta_rule", &torch_chunk_gated_delta_rule,
"C++ chunked gated delta rule (from xllm upstream)");
"C++ chunked gated delta rule (from xllm upstream)",
py::arg("query"),
py::arg("key"),
py::arg("value"),
py::arg("g"),
py::arg("beta"),
py::arg("chunk_size") = 64,
py::arg("initial_state") = c10::nullopt,
py::arg("output_final_state") = false,
py::arg("use_qk_l2norm_in_kernel") = false);
m.def("torch_recurrent_gated_delta_rule", &torch_recurrent_gated_delta_rule,
"C++ recurrent gated delta rule (from xllm upstream)");
"C++ recurrent gated delta rule (from xllm upstream)",
py::arg("query"),
py::arg("key"),
py::arg("value"),
py::arg("g"),
py::arg("beta"),
py::arg("initial_state") = c10::nullopt,
py::arg("output_final_state") = false,
py::arg("use_qk_l2norm_in_kernel") = false);
}

6
qwen3_6_scripts/qwen3_5.py
View File

@@ -153,11 +153,7 @@ try:
except ImportError:
_corex_gdn_chunk_recurrent = None
# DISABLED: corex_gdn_chunk_recurrent.so pybind uses positional args only,
# but the call site passes keyword args (initial_state=, output_final_state=).
# This causes "incompatible function arguments" crash during profiling.
# Force Python fallback until .so signature is fixed.
_HAS_COREX_GDN_CHUNK = False
_HAS_COREX_GDN_CHUNK = _corex_gdn_chunk_recurrent is not None
from vllm.model_executor.models.interfaces import (HasInnerState, SupportsLoRA,
SupportsMultiModal)