Updates for FA3 and other changes

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

csrc/torch_bindings.cpp

@@ -478,7 +478,7 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {
       "                        Tensor! value_cache,"
       "                        Tensor slot_mapping,"
       "                        str kv_cache_dtype,"
-      "                        float k_scale, float v_scale) -> ()");
+      "                        Tensor k_scale, Tensor v_scale) -> ()");
   cache_ops.impl("reshape_and_cache_flash_full_cuda", torch::kCUDA,
                  &reshape_and_cache_flash_full_cuda);