struct hl_device *hdev = hpriv->hdev;
        struct hl_ctx *ctx = hpriv->ctx;
        bool need_soft_reset = false;
-       int rc = 0, do_ctx_switch;
+       int rc = 0, do_ctx_switch = 0;
        void __user *chunks;
        u32 num_chunks, tmp;
        u16 sob_count;
        int ret;
 
-       do_ctx_switch = atomic_cmpxchg(&ctx->thread_ctx_switch_token, 1, 0);
+       if (hdev->supports_ctx_switch)
+               do_ctx_switch = atomic_cmpxchg(&ctx->thread_ctx_switch_token, 1, 0);
 
        if (do_ctx_switch || (args->in.cs_flags & HL_CS_FLAGS_FORCE_RESTORE)) {
                mutex_lock(&hpriv->restore_phase_mutex);
                        }
                }
 
-               ctx->thread_ctx_switch_wait_token = 1;
+               if (hdev->supports_ctx_switch)
+                       ctx->thread_ctx_switch_wait_token = 1;
 
-       } else if (!ctx->thread_ctx_switch_wait_token) {
+       } else if (hdev->supports_ctx_switch && !ctx->thread_ctx_switch_wait_token) {
                rc = hl_poll_timeout_memory(hdev,
                        &ctx->thread_ctx_switch_wait_token, tmp, (tmp == 1),
                        100, jiffies_to_usecs(hdev->timeout_jiffies), false);
 
  *                         Used only for testing.
  * @heartbeat: Controls if we want to enable the heartbeat mechanism vs. the f/w, which verifies
  *             that the f/w is always alive. Used only for testing.
+ * @supports_ctx_switch: true if a ctx switch is required upon first submission.
  */
 struct hl_device {
        struct pci_dev                  *pdev;
        u8                              compute_ctx_in_release;
        u8                              supports_mmu_prefetch;
        u8                              reset_upon_device_release;
+       u8                              supports_ctx_switch;
 
        /* Parameters for bring-up */
        u64                             nic_ports_mask;
 
        hdev->asic_prop.supports_compute_reset = true;
        hdev->asic_prop.allow_inference_soft_reset = true;
        hdev->supports_wait_for_multi_cs = false;
+       hdev->supports_ctx_switch = true;
 
        hdev->asic_funcs->set_pci_memory_regions(hdev);