Update triton_flash_blocksparse_attn.py

Add suggestion similar to https://huggingface.co/THUDM/cogagent-chat-hf/blob/d519da3b191401234f4bd86ce1c287c61bc276a3/util.py#L210 to avoid error
```ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?)```

triton_flash_blocksparse_attn.py

@@ -611,30 +611,31 @@ def _forward(ctx, q, k, v, layout_crow_indices, layout_col_indices, sm_scale, BL
     # print(f'> {q.shape=}, {k.shape=}, {layout_crow_indices.shape}, {layout_col_indices.shape}, {layout_crow_indices.stride()}, \
     #   {layout_col_indices.stride()}, {layout_crow_indices=}, {layout_col_indices=}')
 
-    _fwd_kernel[grid](
-        q, k, v, sm_scale,
-        layout_crow_indices,
-        layout_col_indices,
-        layout_crow_indices.stride(0), layout_crow_indices.stride(1),
-        layout_col_indices.stride(0), layout_col_indices.stride(1),
-        tmp, L, m,
-        o,
-        q.stride(0), q.stride(1), q.stride(2), q.stride(3),
-        k.stride(0), k.stride(1), k.stride(2), k.stride(3),
-        v.stride(0), v.stride(1), v.stride(2), v.stride(3),
-        o.stride(0), o.stride(1), o.stride(2), o.stride(3),
-        q.shape[0], q.shape[1], k.shape[2],
-        k.shape[2] - q.shape[2],
-        q_rounded_len,
-        BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N,
-        BLOCK_DMODEL=BLOCK_DMODEL,
-        EVEN_M_BLOCK=q.shape[2] % BLOCK_M == 0,
-        EVEN_N_BLOCK=k.shape[2] % BLOCK_N == 0 ,
-        INFERENCE=inference,
-        NUM_DBLOCKS=q.shape[-1] // BLOCK_DMODEL,
-        num_warps=num_warps,
-        num_stages=num_stages,
-    )
+    with torch.cuda.device(q.device.index):
+        _fwd_kernel[grid](
+            q, k, v, sm_scale,
+            layout_crow_indices,
+            layout_col_indices,
+            layout_crow_indices.stride(0), layout_crow_indices.stride(1),
+            layout_col_indices.stride(0), layout_col_indices.stride(1),
+            tmp, L, m,
+            o,
+            q.stride(0), q.stride(1), q.stride(2), q.stride(3),
+            k.stride(0), k.stride(1), k.stride(2), k.stride(3),
+            v.stride(0), v.stride(1), v.stride(2), v.stride(3),
+            o.stride(0), o.stride(1), o.stride(2), o.stride(3),
+            q.shape[0], q.shape[1], k.shape[2],
+            k.shape[2] - q.shape[2],
+            q_rounded_len,
+            BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N,
+            BLOCK_DMODEL=BLOCK_DMODEL,
+            EVEN_M_BLOCK=q.shape[2] % BLOCK_M == 0,
+            EVEN_N_BLOCK=k.shape[2] % BLOCK_N == 0 ,
+            INFERENCE=inference,
+            NUM_DBLOCKS=q.shape[-1] // BLOCK_DMODEL,
+            num_warps=num_warps,
+            num_stages=num_stages,
+        )
     if inference:
         L, m = None, None