Stage 4: Sync plugin API with upstream v2.17 pybind changes

- ops.py: Add 15 new API declarations, update 4 existing signatures
- cuda backend: Add 15 new methods, update 4 signatures in cuda.py and register_ops.py
- All vendor backends (enflame/iluvatar/metax/musa/hygon): Add 19 new methods + OpImpl entries

New APIs: splits_to_offsets_multi, copy_data_ptrs_to_device, bulk_allocate,
create_empty_quantized_tensor, group_dequantize, get_grouped_gemm_setup_workspace_size,
multi_tensor_pad_last_dim, multi_tensor_swizzle_scales_for_gemm_,
multi_tensor_transpose_to_bhsd, cusolvermp_ctx_create, cusolvermp_ctx_destroy,
newton_schulz, nvfp4_quantize_with_amax, nvfp4_group_quantize_with_amax,
swizzle_scales_and_pack_ptrs_for_discrete_weights

Modified signatures: group_quantize, bgrad_group_quantize (+tensor_offsets),
clamped_swiglu, clamped_dswiglu (+glu_linear_offset)

Author: FlagScale Agent

transformer_engine/plugin/core/backends/vendor/cuda/cuda.py

@@ -212,6 +212,7 @@ def group_quantize(
         quantizer: Any,
         num_tensors: int,
         first_dims: List[int],
+        tensor_offsets: Optional[Any] = None,
     ) -> Any:
         tex = self._get_tex()
         try:
@@ -221,14 +222,15 @@ def group_quantize(
                     quantizer.dtype = tex.DType(int(qdtype))
         except Exception:
             pass
-        return tex.group_quantize(tensor, quantizer, num_tensors, first_dims)
+        return tex.group_quantize(tensor, quantizer, num_tensors, first_dims, tensor_offsets)
 
     def bgrad_group_quantize(
         self,
         tensor: torch.Tensor,
         quantizer: Any,
         num_tensors: int,
         first_dims: List[int],
+        tensor_offsets: Optional[Any] = None,
     ) -> Any:
         tex = self._get_tex()
         try:
@@ -238,7 +240,7 @@ def bgrad_group_quantize(
                     quantizer.dtype = tex.DType(int(qdtype))
         except Exception:
             pass
-        return tex.bgrad_group_quantize(tensor, quantizer, num_tensors, first_dims)
+        return tex.bgrad_group_quantize(tensor, quantizer, num_tensors, first_dims, tensor_offsets)
 
     def generic_gemm(
         self,
@@ -349,9 +351,10 @@ def clamped_swiglu(
         quantizer: Any,
         limit: float = 7.0,
         alpha: float = 1.702,
+        glu_linear_offset: float = 1.0,
     ) -> Any:
         tex = self._get_tex()
-        return tex.clamped_swiglu(input, quantizer, limit, alpha)
+        return tex.clamped_swiglu(input, quantizer, limit, alpha, glu_linear_offset)
 
     # Backward of GLU #
     def dglu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any:
@@ -408,9 +411,10 @@ def clamped_dswiglu(
         quantizer: Any,
         limit: float = 7.0,
         alpha: float = 1.702,
+        glu_linear_offset: float = 1.0,
     ) -> Any:
         tex = self._get_tex()
-        return tex.clamped_dswiglu(grad, fwd_input, quantizer, limit, alpha)
+        return tex.clamped_dswiglu(grad, fwd_input, quantizer, limit, alpha, glu_linear_offset)
 
     # DBias + DAct fusions #
     def dbias_dgelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> List[Any]:
@@ -787,6 +791,156 @@ def splits_to_offsets(
         tex = self._get_tex()
         return tex.splits_to_offsets(first_dims, logical_last_dim)
 
+    def splits_to_offsets_multi(
+        self,
+        split_sizes: List[int],
+        device: Any,
+        *,
+        strides: Any,
+        include_leading_zero: bool,
+        dtypes: Any,
+        bulk_allocate: bool = False,
+    ) -> Any:
+        tex = self._get_tex()
+        return tex.splits_to_offsets_multi(
+            split_sizes, device, strides=strides,
+            include_leading_zero=include_leading_zero,
+            dtypes=dtypes, bulk_allocate=bulk_allocate,
+        )
+
+    def copy_data_ptrs_to_device(
+        self,
+        tensors: List[torch.Tensor],
+        device: Any,
+    ) -> Any:
+        tex = self._get_tex()
+        return tex.copy_data_ptrs_to_device(tensors, device)
+
+    def bulk_allocate(
+        self,
+        shapes: List[List[int]],
+        dtypes: List[Any],
+        device: Optional[Any] = None,
+        alignments: Optional[List[int]] = None,
+    ) -> Any:
+        tex = self._get_tex()
+        return tex.bulk_allocate(shapes, dtypes, device, alignments)
+
+    def create_empty_quantized_tensor(
+        self,
+        quantizer: Any,
+        shape: List[int],
+        dtype: Any,
+        device: Any,
+        pin_memory: bool,
+    ) -> Any:
+        tex = self._get_tex()
+        return tex.create_empty_quantized_tensor(quantizer, shape, dtype, device, pin_memory)
+
+    def group_dequantize(
+        self,
+        input: Any,
+        otype: Any,
+    ) -> Any:
+        tex = self._get_tex()
+        return tex.group_dequantize(input, otype)
+
+    def get_grouped_gemm_setup_workspace_size(self) -> int:
+        tex = self._get_tex()
+        return tex.get_grouped_gemm_setup_workspace_size()
+
+    def multi_tensor_pad_last_dim(
+        self,
+        inputs: List[torch.Tensor],
+        alignment: int,
+    ) -> Any:
+        tex = self._get_tex()
+        return tex.multi_tensor_pad_last_dim(inputs, alignment)
+
+    def multi_tensor_swizzle_scales_for_gemm_(
+        self,
+        tensors: List[torch.Tensor],
+        rowwise_usage: Any,
+        columnwise_usage: Any,
+    ) -> None:
+        tex = self._get_tex()
+        return tex.multi_tensor_swizzle_scales_for_gemm_(tensors, rowwise_usage, columnwise_usage)
+
+    def multi_tensor_transpose_to_bhsd(
+        self,
+        inputs: List[torch.Tensor],
+        original_format: Any,
+        outputs: Optional[List[Optional[torch.Tensor]]] = None,
+    ) -> Any:
+        tex = self._get_tex()
+        return tex.multi_tensor_transpose_to_bhsd(inputs, original_format, outputs)
+
+    def cusolvermp_ctx_create(
+        self,
+        nccl_comm_ptr: int,
+        nranks: int,
+        rank: int,
+    ) -> Any:
+        tex = self._get_tex()
+        return tex.cusolvermp_ctx_create(nccl_comm_ptr, nranks, rank)
+
+    def cusolvermp_ctx_destroy(
+        self,
+        ctx_ptr: Any,
+    ) -> None:
+        tex = self._get_tex()
+        return tex.cusolvermp_ctx_destroy(ctx_ptr)
+
+    def newton_schulz(
+        self,
+        ctx_ptr: Any,
+        m: int,
+        n: int,
+        x: torch.Tensor,
+        num_iterations: int,
+        coefficients: Any,
+    ) -> Any:
+        tex = self._get_tex()
+        return tex.newton_schulz(ctx_ptr, m, n, x, num_iterations, coefficients)
+
+    def nvfp4_quantize_with_amax(
+        self,
+        tensor: torch.Tensor,
+        quantizer: Any,
+        rowwise_amax: torch.Tensor,
+        columnwise_amax: torch.Tensor,
+    ) -> Any:
+        tex = self._get_tex()
+        return tex.nvfp4_quantize_with_amax(tensor, quantizer, rowwise_amax, columnwise_amax)
+
+    def nvfp4_group_quantize_with_amax(
+        self,
+        tensor: torch.Tensor,
+        quantizer: Any,
+        num_tensors: int,
+        first_dims: List[int],
+        rowwise_amax: torch.Tensor,
+        columnwise_amax: torch.Tensor,
+        tensor_offsets: Optional[Any] = None,
+    ) -> Any:
+        tex = self._get_tex()
+        return tex.nvfp4_group_quantize_with_amax(
+            tensor, quantizer, num_tensors, first_dims,
+            rowwise_amax, columnwise_amax, tensor_offsets,
+        )
+
+    def swizzle_scales_and_pack_ptrs_for_discrete_weights(
+        self,
+        data_tensors: List[torch.Tensor],
+        scale_tensors: List[torch.Tensor],
+        swizzle_type: Any,
+        device: Any,
+    ) -> Any:
+        tex = self._get_tex()
+        return tex.grouped_mlp_experimental.swizzle_scales_and_pack_ptrs_for_discrete_weights(
+            data_tensors, scale_tensors, swizzle_type, device,
+        )
+
     def get_fused_attn_backend(
         self,
         is_training: bool,

transformer_engine/plugin/core/backends/vendor/cuda/register_ops.py

@@ -1168,6 +1168,127 @@ def register_builtins(registry) -> None:
             vendor="CUDA",
             priority=100,
         ),
+        # v2.17 new APIs
+        OpImpl(
+            op_name="splits_to_offsets_multi",
+            impl_id="vendor.cuda",
+            kind=BackendImplKind.VENDOR,
+            fn=_bind_is_available(backend.splits_to_offsets_multi, is_avail),
+            vendor="CUDA",
+            priority=100,
+        ),
+        OpImpl(
+            op_name="copy_data_ptrs_to_device",
+            impl_id="vendor.cuda",
+            kind=BackendImplKind.VENDOR,
+            fn=_bind_is_available(backend.copy_data_ptrs_to_device, is_avail),
+            vendor="CUDA",
+            priority=100,
+        ),
+        OpImpl(
+            op_name="bulk_allocate",
+            impl_id="vendor.cuda",
+            kind=BackendImplKind.VENDOR,
+            fn=_bind_is_available(backend.bulk_allocate, is_avail),
+            vendor="CUDA",
+            priority=100,
+        ),
+        OpImpl(
+            op_name="create_empty_quantized_tensor",
+            impl_id="vendor.cuda",
+            kind=BackendImplKind.VENDOR,
+            fn=_bind_is_available(backend.create_empty_quantized_tensor, is_avail),
+            vendor="CUDA",
+            priority=100,
+        ),
+        OpImpl(
+            op_name="group_dequantize",
+            impl_id="vendor.cuda",
+            kind=BackendImplKind.VENDOR,
+            fn=_bind_is_available(backend.group_dequantize, is_avail),
+            vendor="CUDA",
+            priority=100,
+        ),
+        OpImpl(
+            op_name="get_grouped_gemm_setup_workspace_size",
+            impl_id="vendor.cuda",
+            kind=BackendImplKind.VENDOR,
+            fn=_bind_is_available(backend.get_grouped_gemm_setup_workspace_size, is_avail),
+            vendor="CUDA",
+            priority=100,
+        ),
+        OpImpl(
+            op_name="multi_tensor_pad_last_dim",
+            impl_id="vendor.cuda",
+            kind=BackendImplKind.VENDOR,
+            fn=_bind_is_available(backend.multi_tensor_pad_last_dim, is_avail),
+            vendor="CUDA",
+            priority=100,
+        ),
+        OpImpl(
+            op_name="multi_tensor_swizzle_scales_for_gemm_",
+            impl_id="vendor.cuda",
+            kind=BackendImplKind.VENDOR,
+            fn=_bind_is_available(backend.multi_tensor_swizzle_scales_for_gemm_, is_avail),
+            vendor="CUDA",
+            priority=100,
+        ),
+        OpImpl(
+            op_name="multi_tensor_transpose_to_bhsd",
+            impl_id="vendor.cuda",
+            kind=BackendImplKind.VENDOR,
+            fn=_bind_is_available(backend.multi_tensor_transpose_to_bhsd, is_avail),
+            vendor="CUDA",
+            priority=100,
+        ),
+        OpImpl(
+            op_name="cusolvermp_ctx_create",
+            impl_id="vendor.cuda",
+            kind=BackendImplKind.VENDOR,
+            fn=_bind_is_available(backend.cusolvermp_ctx_create, is_avail),
+            vendor="CUDA",
+            priority=100,
+        ),
+        OpImpl(
+            op_name="cusolvermp_ctx_destroy",
+            impl_id="vendor.cuda",
+            kind=BackendImplKind.VENDOR,
+            fn=_bind_is_available(backend.cusolvermp_ctx_destroy, is_avail),
+            vendor="CUDA",
+            priority=100,
+        ),
+        OpImpl(
+            op_name="newton_schulz",
+            impl_id="vendor.cuda",
+            kind=BackendImplKind.VENDOR,
+            fn=_bind_is_available(backend.newton_schulz, is_avail),
+            vendor="CUDA",
+            priority=100,
+        ),
+        OpImpl(
+            op_name="nvfp4_quantize_with_amax",
+            impl_id="vendor.cuda",
+            kind=BackendImplKind.VENDOR,
+            fn=_bind_is_available(backend.nvfp4_quantize_with_amax, is_avail),
+            vendor="CUDA",
+            priority=100,
+        ),
+        OpImpl(
+            op_name="nvfp4_group_quantize_with_amax",
+            impl_id="vendor.cuda",
+            kind=BackendImplKind.VENDOR,
+            fn=_bind_is_available(backend.nvfp4_group_quantize_with_amax, is_avail),
+            vendor="CUDA",
+            priority=100,
+        ),
+        OpImpl(
+            op_name="swizzle_scales_and_pack_ptrs_for_discrete_weights",
+            impl_id="vendor.cuda",
+            kind=BackendImplKind.VENDOR,
+            fn=_bind_is_available(backend.swizzle_scales_and_pack_ptrs_for_discrete_weights, is_avail),
+            vendor="CUDA",
+            priority=100,
+        ),
     ]
 
     registry.register_many(impls)