another assert

assert_tree_attn.py

@@ -87,7 +87,6 @@ def start(
     tree_out = tree_out.cpu()
     out = out.cpu()
 
-    print((tree_out - out).abs().amax())
     output_atol = 1e-2 if use_cuda else 1e-5
 
     assert torch.allclose(tree_out, out, atol = output_atol), '🟥 output is not the same'

ring_attention_pytorch/tree_attn_decoding.py

@@ -42,13 +42,15 @@ def tree_attn_decode(
     # each machine (rank) takes care of a chunk of kv sequence within the world of many machines
 
     if shard_kv_seq:
+        assert exists(k), 'keys and values must be passed if not already sharded across sequence'
+
         rank, world_size = get_rank(), get_world_size()
         k = k.chunk(world_size, dim = -2)
         v = v.chunk(world_size, dim = -2)
 
         k, v = (k[rank], v[rank]) if rank < len(k) else (None, None)
 
-    if exists(k) and exists(v):
+    if exists(k):
         # calculate local output and derive numerator and denominator
 
         use_triton = default(use_triton, q.is_cuda)