Added Continuous Batching (CB) Support for Subfunctions (#642)

Resolved compilation issues by adding CB support for subfunctions,
ensuring compatibility across CB and non-CB models.

Signed-off-by: abhishek-singh591 <[email protected]>

Author: abhishek-singh591

QEfficient/base/onnx_transforms.py

@@ -11,7 +11,26 @@
 import torch
 from onnx import ModelProto, external_data_helper, numpy_helper
 
-from QEfficient.customop.ctx_scatter_gather import CtxGather, CtxGatherFunc, CtxScatter, CtxScatterFunc
+from QEfficient.customop.ctx_scatter_gather import (
+    CtxGather,
+    CtxGather3D,
+    CtxGatherFunc,
+    CtxGatherFunc3D,
+    CtxScatter,
+    CtxScatter3D,
+    CtxScatterFunc,
+    CtxScatterFunc3D,
+)
+from QEfficient.customop.ctx_scatter_gather_cb import (
+    CtxGatherCB,
+    CtxGatherCB3D,
+    CtxGatherFuncCB,
+    CtxGatherFuncCB3D,
+    CtxScatterCB,
+    CtxScatterCB3D,
+    CtxScatterFuncCB,
+    CtxScatterFuncCB3D,
+)
 from QEfficient.customop.rms_norm import CustomRMSNorm, CustomRMSNormFunc
 
 
@@ -113,7 +132,13 @@ class CustomOpTransform(OnnxTransform):
     _custom_ops: Dict[str, Tuple[Any, Any]] = {
         "CustomRMSNormFunc": (CustomRMSNormFunc, CustomRMSNorm),
         "CtxScatterFunc": (CtxScatterFunc, CtxScatter),
+        "CtxScatterFunc3D": (CtxScatterFunc3D, CtxScatter3D),
         "CtxGatherFunc": (CtxGatherFunc, CtxGather),
+        "CtxGatherFunc3D": (CtxGatherFunc3D, CtxGather3D),
+        "CtxScatterFuncCB": (CtxScatterFuncCB, CtxScatterCB),
+        "CtxScatterFuncCB3D": (CtxScatterFuncCB3D, CtxScatterCB3D),
+        "CtxGatherFuncCB": (CtxGatherFuncCB, CtxGatherCB),
+        "CtxGatherFuncCB3D": (CtxGatherFuncCB3D, CtxGatherCB3D),
     }
 
     @classmethod