diff --git a/drivers/gpu/nvgpu/common/gr/gr_setup.c b/drivers/gpu/nvgpu/common/gr/gr_setup.c
index 470b6933f..f836e70d3 100644
--- a/drivers/gpu/nvgpu/common/gr/gr_setup.c
+++ b/drivers/gpu/nvgpu/common/gr/gr_setup.c
@@ -242,12 +242,21 @@ int nvgpu_gr_setup_alloc_obj_ctx(struct nvgpu_channel *c, u32 class_num,
 
 #ifdef CONFIG_NVGPU_FECS_TRACE
 	if (g->ops.gr.fecs_trace.bind_channel && !c->vpr) {
+		u32 vmid = 0U;
+
 		if (nvgpu_is_enabled(g, NVGPU_SUPPORT_TSG_SUBCONTEXTS)) {
 			gr_subctx = nvgpu_tsg_subctx_get_gr_subctx(c->subctx);
 		}
 
+		if (g->ops.channel.get_vmid != NULL) {
+			err = g->ops.channel.get_vmid(c, &vmid);
+			if (err != 0) {
+				nvgpu_warn(g, "failed to get vmid from channel. chid=%u", c->chid);
+			}
+		}
+
 		err = g->ops.gr.fecs_trace.bind_channel(g, &c->inst_block,
-			gr_subctx, gr_ctx, mappings, tsg->tgid, 0);
+			gr_subctx, gr_ctx, mappings, tsg->tgid, vmid);
 		if (err != 0) {
 			nvgpu_warn(g,
 				"fail to bind channel for ctxsw trace");
diff --git a/drivers/gpu/nvgpu/include/nvgpu/gops/channel.h b/drivers/gpu/nvgpu/include/nvgpu/gops/channel.h
index 0b17175c2..07010879b 100644
--- a/drivers/gpu/nvgpu/include/nvgpu/gops/channel.h
+++ b/drivers/gpu/nvgpu/include/nvgpu/gops/channel.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2020, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2019-2022, NVIDIA CORPORATION.  All rights reserved.
  *
  * Permission is hereby granted, free of charge, to any person obtaining a
  * copy of this software and associated documentation files (the "Software"),
@@ -142,6 +142,7 @@ struct gops_channel {
 	void (*reset_faulted)(struct gk20a *g, struct nvgpu_channel *ch,
 			bool eng, bool pbdma);
 	void (*clear)(struct nvgpu_channel *ch);
+	int (*get_vmid)(struct nvgpu_channel *ch, u32 *vmid);
 
 #ifdef CONFIG_NVGPU_KERNEL_MODE_SUBMIT
 	int (*set_syncpt)(struct nvgpu_channel *ch);