TP/quantization/weight loading refactor part 1 - Simplify parallel linear logic (#1181)

2023-10-02 15:36:09 -07:00
parent 84e4e37d14
commit ba0bfd40e2
42 changed files with 819 additions and 1547 deletions
--- a/tests/kernels/test_pos_encoding.py
+++ b/tests/kernels/test_pos_encoding.py
@@ -140,7 +140,7 @@ def test_rotary_embedding(
    cos = freqs.cos()
    sin = freqs.sin()
    cos_sin_cache = torch.cat((cos, sin), dim=-1)
-    cos_sin_cache = cos_sin_cache.to(dtype=dtype, device='cuda')
+    cos_sin_cache = cos_sin_cache.to(dtype=dtype, device="cuda")

    # Run the kernel. The kernel is in-place, so we need to clone the inputs.
    out_query = query.clone()