fix for fp8 kv cache

IanNod · IanNod · commit c38c165430fe · 2025-04-25T10:33:19.000-05:00
diff --git a/sharktank/sharktank/utils/iree.py b/sharktank/sharktank/utils/iree.py
@@ -258,7 +258,7 @@ def torch_tensor_to_device_array(
     tensor: torch.Tensor, device: iree.runtime.HalDevice
 ) -> iree.runtime.DeviceArray:
     if tensor.dtype in halelementtype_map.keys():
-        tensor_as_int16 = tensor.view(dtype=torch.int16)
+        tensor_as_int16 = tensor.to(dtype=torch.int16)
         device_array_as_int16 = iree.runtime.asdevicearray(
             device, unbox_tensor(tensor_as_int16).to("cpu").detach().numpy()
         )

Original file line number	Diff line number	Diff line change
`@@ -258,7 +258,7 @@ def torch_tensor_to_device_array(`
`258`	`258`	`tensor: torch.Tensor, device: iree.runtime.HalDevice`
`259`	`259`	`) -> iree.runtime.DeviceArray:`
`260`	`260`	`if tensor.dtype in halelementtype_map.keys():`
`261`		`- tensor_as_int16 = tensor.view(dtype=torch.int16)`
	`261`	`+ tensor_as_int16 = tensor.to(dtype=torch.int16)`
`262`	`262`	`device_array_as_int16 = iree.runtime.asdevicearray(`
`263`	`263`	`device, unbox_tensor(tensor_as_int16).to("cpu").detach().numpy()`
`264`	`264`	`)`