diff --git a/arch/riscv/include/asm/vector.h b/arch/riscv/include/asm/vector.h
index 1f8dcae3af51..92192621af47 100644
--- a/arch/riscv/include/asm/vector.h
+++ b/arch/riscv/include/asm/vector.h
@@ -377,8 +377,7 @@ static inline void __switch_to_vector(struct task_struct *prev,
 	struct pt_regs *regs;
 
 	if (riscv_preempt_v_started(prev)) {
-		if (riscv_v_is_on()) {
-			WARN_ON(prev->thread.riscv_v_flags & RISCV_V_CTX_DEPTH_MASK);
+		if (!(current->thread.riscv_v_flags & RISCV_V_CTX_DEPTH_MASK)) {
 			riscv_v_disable();
 			prev->thread.riscv_v_flags |= RISCV_PREEMPT_V_IN_SCHEDULE;
 		}
diff --git a/arch/riscv/kernel/kernel_mode_vector.c b/arch/riscv/kernel/kernel_mode_vector.c
index 77e98b504485..5ad93ddf6a10 100644
--- a/arch/riscv/kernel/kernel_mode_vector.c
+++ b/arch/riscv/kernel/kernel_mode_vector.c
@@ -182,14 +182,6 @@ static int riscv_v_start_kernel_context(void)
 	get_cpu_vector_context();
 	__riscv_flush_vector_context();
 	put_cpu_vector_context();
-	/*
-	 *  A voluntary context switch caused by put_cpu_vector_context() can
-	 *  raise the NEED_RESTORE flag if preempt_v starts too early due to a
-	 *  failed risv_v_is_on() check.
-	 *
-	 *  This causes the next context_nesting_end pollute the v-reg from
-	 *  the stale context memory in kernel-mode vector.
-	 */
 	riscv_v_start(RISCV_PREEMPT_V);
 	return 0;
 }