diff --git a/arch/riscv/include/asm/vector.h b/arch/riscv/include/asm/vector.h
index 898b3d0036e78..398bab85ddc54 100644
--- a/arch/riscv/include/asm/vector.h
+++ b/arch/riscv/include/asm/vector.h
@@ -107,6 +107,7 @@ static inline void __riscv_m_mstate_restore(struct __riscv_m_ext_state *restore_
 extern unsigned long riscv_v_vsize;
 int riscv_v_setup_vsize(void);
 bool riscv_v_first_use_handler(struct pt_regs *regs);
+int riscv_v_thread_zalloc(struct task_struct *tsk);
 
 static __always_inline bool has_vector(void)
 {
diff --git a/arch/riscv/kernel/ptrace.c b/arch/riscv/kernel/ptrace.c
index 0735796f81e18..82114b729138c 100644
--- a/arch/riscv/kernel/ptrace.c
+++ b/arch/riscv/kernel/ptrace.c
@@ -92,8 +92,13 @@ static int riscv_vr_get(struct task_struct *target,
 	struct __riscv_v_ext_state *vstate = &target->thread.vstate;
 	struct __riscv_v_regset_state ptrace_vstate;
 
-	if (!riscv_v_vstate_query(task_pt_regs(target)))
-		return -EINVAL;
+	if (!riscv_v_vstate_query(task_pt_regs(target))) {
+		if (riscv_v_thread_zalloc(target))
+			return -EINVAL;
+		riscv_v_vstate_on(task_pt_regs(target));
+		if (target == current)
+			riscv_v_vstate_restore(current, task_pt_regs(current));
+	}
 
 	/*
 	 * Ensure the vector registers have been saved to the memory before
@@ -124,8 +129,13 @@ static int riscv_vr_set(struct task_struct *target,
 	struct __riscv_v_ext_state *vstate = &target->thread.vstate;
 	struct __riscv_v_regset_state ptrace_vstate;
 
-	if (!riscv_v_vstate_query(task_pt_regs(target)))
-		return -EINVAL;
+	if (!riscv_v_vstate_query(task_pt_regs(target))) {
+		if (riscv_v_thread_zalloc(target))
+			return -EINVAL;
+		riscv_v_vstate_on(task_pt_regs(target));
+		if (target == current)
+			riscv_v_vstate_restore(current, task_pt_regs(current));
+	}
 
 	/* Copy rest of the vstate except datap */
 	ret = user_regset_copyin(&pos, &count, &kbuf, &ubuf, &ptrace_vstate, 0,
diff --git a/arch/riscv/kernel/vector.c b/arch/riscv/kernel/vector.c
index 6e2d5d97f7468..12967f0e24ba0 100644
--- a/arch/riscv/kernel/vector.c
+++ b/arch/riscv/kernel/vector.c
@@ -113,7 +113,7 @@ static bool insn_is_matrix(u32 insn_buf)
 	return false;
 }
 
-static int riscv_v_thread_zalloc(void)
+int riscv_v_thread_zalloc(struct task_struct *tsk)
 {
 	void *datap;
 
@@ -121,9 +121,10 @@ static int riscv_v_thread_zalloc(void)
 	if (!datap)
 		return -ENOMEM;
 
-	current->thread.vstate.datap = datap;
-	memset(&current->thread.vstate, 0, offsetof(struct __riscv_v_ext_state,
-						    datap));
+	tsk->thread.vstate.datap = datap;
+	memset(&tsk->thread.vstate, 0, offsetof(struct __riscv_v_ext_state,
+						datap));
+	tsk->thread.vstate.vlenb = riscv_v_vsize/32;
 	return 0;
 }
 
@@ -222,7 +223,7 @@ static bool __riscv_v_first_use_handler(struct pt_regs *regs)
 	 * context where VS has been off. So, try to allocate the user's V
 	 * context and resume execution.
 	 */
-	if (riscv_v_thread_zalloc()) {
+	if (riscv_v_thread_zalloc(current)) {
 		force_sig(SIGBUS);
 		return true;
 	}