diff --git a/drivers/gpu/nvgpu/common/fifo/channel.c b/drivers/gpu/nvgpu/common/fifo/channel.c
index 46d76e13c..596008ddf 100644
--- a/drivers/gpu/nvgpu/common/fifo/channel.c
+++ b/drivers/gpu/nvgpu/common/fifo/channel.c
@@ -232,10 +232,12 @@ bool  gk20a_channel_check_timedout(struct channel_gk20a *ch)
 
 void gk20a_channel_abort(struct channel_gk20a *ch, bool channel_preempt)
 {
+	struct tsg_gk20a *tsg = tsg_gk20a_from_ch(ch);
+
 	nvgpu_log_fn(ch->g, " ");
 
-	if (gk20a_is_channel_marked_as_tsg(ch)) {
-		return gk20a_fifo_abort_tsg(ch->g, ch->tsgid, channel_preempt);
+	if (tsg != NULL) {
+		return gk20a_fifo_abort_tsg(ch->g, tsg, channel_preempt);
 	}
 
 	/* make sure new kickoffs are prevented */
diff --git a/drivers/gpu/nvgpu/common/fifo/tsg.c b/drivers/gpu/nvgpu/common/fifo/tsg.c
index cf527e55f..c59d1d70f 100644
--- a/drivers/gpu/nvgpu/common/fifo/tsg.c
+++ b/drivers/gpu/nvgpu/common/fifo/tsg.c
@@ -20,6 +20,7 @@
  * DEALINGS IN THE SOFTWARE.
  */
 
+#include <nvgpu/bug.h>
 #include <nvgpu/kmem.h>
 #include <nvgpu/log.h>
 #include <nvgpu/os_sched.h>
@@ -149,18 +150,21 @@ int gk20a_tsg_bind_channel(struct tsg_gk20a *tsg,
 	return 0;
 }
 
+/* The caller must ensure that channel belongs to a tsg */
 int gk20a_tsg_unbind_channel(struct channel_gk20a *ch)
 {
 	struct gk20a *g = ch->g;
-	struct tsg_gk20a *tsg = &g->fifo.tsg[ch->tsgid];
+	struct tsg_gk20a *tsg = tsg_gk20a_from_ch(ch);
 	int err;
 
+	nvgpu_assert(tsg != NULL);
+
 	err = g->ops.fifo.tsg_unbind_channel(ch);
 	if (err != 0) {
 		nvgpu_err(g, "Channel %d unbind failed, tearing down TSG %d",
 			ch->chid, tsg->tsgid);
 
-		gk20a_fifo_abort_tsg(ch->g, ch->tsgid, true);
+		gk20a_fifo_abort_tsg(ch->g, tsg, true);
 		/* If channel unbind fails, channel is still part of runlist */
 		channel_gk20a_update_runlist(ch, false);
 
diff --git a/drivers/gpu/nvgpu/gk20a/fifo_gk20a.c b/drivers/gpu/nvgpu/gk20a/fifo_gk20a.c
index f8cfbcb28..f1d3f5ec6 100644
--- a/drivers/gpu/nvgpu/gk20a/fifo_gk20a.c
+++ b/drivers/gpu/nvgpu/gk20a/fifo_gk20a.c
@@ -1453,17 +1453,16 @@ void gk20a_fifo_set_ctx_mmu_error_tsg(struct gk20a *g,
 
 }
 
-void gk20a_fifo_abort_tsg(struct gk20a *g, u32 tsgid, bool preempt)
+void gk20a_fifo_abort_tsg(struct gk20a *g, struct tsg_gk20a *tsg, bool preempt)
 {
-	struct tsg_gk20a *tsg = &g->fifo.tsg[tsgid];
-	struct channel_gk20a *ch;
+	struct channel_gk20a *ch = NULL;
 
 	nvgpu_log_fn(g, " ");
 
 	g->ops.fifo.disable_tsg(tsg);
 
 	if (preempt) {
-		g->ops.fifo.preempt_tsg(g, tsgid);
+		g->ops.fifo.preempt_tsg(g, tsg->tsgid);
 	}
 
 	nvgpu_rwsem_down_read(&tsg->ch_list_lock);
@@ -1710,7 +1709,7 @@ static bool gk20a_fifo_handle_mmu_fault_locked(
 									 tsg);
 				}
 				verbose = gk20a_fifo_error_tsg(g, tsg);
-				gk20a_fifo_abort_tsg(g, tsg->tsgid, false);
+				gk20a_fifo_abort_tsg(g, tsg, false);
 			}
 
 			/* put back the ref taken early above */
@@ -1902,7 +1901,7 @@ void gk20a_fifo_recover_tsg(struct gk20a *g, u32 tsgid, bool verbose,
 			gk20a_debug_dump(g);
 		}
 
-		gk20a_fifo_abort_tsg(g, tsgid, false);
+		gk20a_fifo_abort_tsg(g, tsg, false);
 	}
 
 	gr_gk20a_enable_ctxsw(g);
diff --git a/drivers/gpu/nvgpu/gk20a/fifo_gk20a.h b/drivers/gpu/nvgpu/gk20a/fifo_gk20a.h
index 05e740b8a..7854dadff 100644
--- a/drivers/gpu/nvgpu/gk20a/fifo_gk20a.h
+++ b/drivers/gpu/nvgpu/gk20a/fifo_gk20a.h
@@ -310,7 +310,7 @@ u32 gk20a_fifo_get_failing_engine_data(struct gk20a *g,
 		u32 *__id, bool *__is_tsg);
 void gk20a_fifo_set_ctx_mmu_error_tsg(struct gk20a *g,
 		struct tsg_gk20a *tsg);
-void gk20a_fifo_abort_tsg(struct gk20a *g, u32 tsgid, bool preempt);
+void gk20a_fifo_abort_tsg(struct gk20a *g, struct tsg_gk20a *tsg, bool preempt);
 void gk20a_fifo_set_ctx_mmu_error_ch(struct gk20a *g,
 		struct channel_gk20a *refch);
 bool gk20a_fifo_error_tsg(struct gk20a *g, struct tsg_gk20a *tsg);
diff --git a/drivers/gpu/nvgpu/gv11b/fifo_gv11b.c b/drivers/gpu/nvgpu/gv11b/fifo_gv11b.c
index a2059db8a..79736d342 100644
--- a/drivers/gpu/nvgpu/gv11b/fifo_gv11b.c
+++ b/drivers/gpu/nvgpu/gv11b/fifo_gv11b.c
@@ -1013,7 +1013,7 @@ static void gv11b_fifo_locked_abort_runlist_active_tsgs(struct gk20a *g,
 					rlid);
 			}
 
-			gk20a_fifo_abort_tsg(g, tsg->tsgid, false);
+			gk20a_fifo_abort_tsg(g, tsg, false);
 
 			nvgpu_log(g, gpu_dbg_info, "aborted tsg id %lu", tsgid);
 		}
@@ -1224,7 +1224,7 @@ void gv11b_fifo_teardown_ch_tsg(struct gk20a *g, u32 act_eng_bitmask,
 				gk20a_fifo_set_ctx_mmu_error_tsg(g, tsg);
 			}
 
-			gk20a_fifo_abort_tsg(g, tsg->tsgid, false);
+			gk20a_fifo_abort_tsg(g, tsg, false);
 		}
 	} else {
 		gv11b_fifo_locked_abort_runlist_active_tsgs(g, rc_type,