From 768d89c31a8e69698a7ed261123cc050408bfb89 Mon Sep 17 00:00:00 2001
From: Claude <claude@anthropic.com>
Date: Fri, 14 Aug 2026 02:02:51 +0000
Subject: [PATCH] fix(pybind): add py::arg + defaults to
 corex_gdn_chunk_recurrent

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/
---
 qwen3_6_scripts/corex_gdn_chunk_recurrent.cu | 21 ++++++++++++++++++--
 qwen3_6_scripts/qwen3_5.py                   |  6 +-----
 2 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/qwen3_6_scripts/corex_gdn_chunk_recurrent.cu b/qwen3_6_scripts/corex_gdn_chunk_recurrent.cu
index 8420e9f0..9e640995 100644
--- a/qwen3_6_scripts/corex_gdn_chunk_recurrent.cu
+++ b/qwen3_6_scripts/corex_gdn_chunk_recurrent.cu
@@ -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);
 }
diff --git a/qwen3_6_scripts/qwen3_5.py b/qwen3_6_scripts/qwen3_5.py
index 9e07cc0f..8103948e 100644
--- a/qwen3_6_scripts/qwen3_5.py
+++ b/qwen3_6_scripts/qwen3_5.py
@@ -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)