Fix trunk health (#4379)

Summary:
Pull Request resolved: #4379

X-link: facebookresearch/FBGEMM#1449

Some tests have long been failing in trunk.

- Fix kv_cache test  on AMD
- Fix AMD build for trt_llm

{F1979505502}

- Fix AMD build for fbcode//deeplearning/fbgemm/fbgemm_gpu/experimental/example/

  {F1979505460}

- Fix type checking in triton_gemm

Reviewed By: q10

Differential Revision: D76954296

fbshipit-source-id: ed0acb8ff50b1b27dc167a39a3dc613609ea2028

Author: spcyppt

fbgemm_gpu/experimental/gemm/triton_gemm/grouped_gemm.py

@@ -266,7 +266,7 @@ def _fbgemm_grouped_gemm(
         c_desc_ptr = None
 
     M_end_offset = 0
-    M_end_offset = M_end_offset.to(tl.int64)
+    M_end_offset = M_end_offset.to(tl.int64)  # pyre-ignore
     iterated_tiles = 0
     for g in tl.range(G):
         # Move across groups
@@ -433,7 +433,7 @@ def _fbgemm_grouped_gemm_ws(
         c_desc_ptr = None
 
     M_end_offset = 0
-    M_end_offset = M_end_offset.to(tl.int64)
+    M_end_offset = M_end_offset.to(tl.int64)  # pyre-ignore
     iterated_tiles = 0
     for g in tl.range(G):
         # Move across groups
@@ -605,7 +605,7 @@ def _fbgemm_grouped_gemm_fp8_rowwise(
         c_desc_ptr = None
 
     M_end_offset = 0
-    M_end_offset = M_end_offset.to(tl.int64)
+    M_end_offset = M_end_offset.to(tl.int64)  # pyre-ignore
     iterated_tiles = 0
     for g in tl.range(G):
         # Move across groups
@@ -786,7 +786,7 @@ def _fbgemm_grouped_gemm_fp8_rowwise_ws(
         c_desc_ptr = None
 
     M_end_offset = 0
-    M_end_offset = M_end_offset.to(tl.int64)
+    M_end_offset = M_end_offset.to(tl.int64)  # pyre-ignore
     iterated_tiles = 0
     for g in tl.range(G):
         # Move across groups

fbgemm_gpu/experimental/gen_ai/test/kv_cache/kv_cache_test.py

@@ -363,7 +363,7 @@ def test_fp8_kv_cache(self, MAX_T: int, N_KVH_L: int) -> None:
         N_KVH_L=st.sampled_from([1, 2]),
     )
     @unittest.skipIf(
-        not torch.cuda.is_available() or not HAS_XFORMERS,
+        not torch.version.hip or not HAS_XFORMERS,
         "Skip when no AMD GPU or xformers is not available",
     )
     def test_fp8_kv_e4m3fn_convert_to_e4m3fnuz(self, MAX_T: int, N_KVH_L: int) -> None: