diff --git a/deepinfer/neuron_ops/nki_attention.py b/deepinfer/neuron_ops/nki_attention.py index 5ba7ef2..d1ada54 100644 --- a/deepinfer/neuron_ops/nki_attention.py +++ b/deepinfer/neuron_ops/nki_attention.py @@ -364,6 +364,10 @@ def _table_attention( # blocks are gathered (one contiguous run each); a per-slot gather is a DMA descriptor per record. chosen, idx = _compact_blocks(idx, blocks, page, used_blocks) table = _gather_blocks(cache, chosen) + elif cache.stride(1) == cache.shape[-1] and cache.stride(-1) == 1 and cache.stride(0) == page * cache.shape[-1]: + # The whole cache in slot order is the cache itself: decode it in place (the arange gather was an + # element-indexed GpSimd gather plus integer index math, the attention graphs' largest cost). + table = _DECODERS[cache.shape[-1]](cache.reshape(blocks * page, cache.shape[-1])) else: table = gather_paged_kv(cache, torch.arange(blocks * page, device=q.device, dtype=torch.int32)) tables.append(torch.nan_to_num(table, nan=0.0, posinf=0.0, neginf=0.0)) # unwritten slots: never NaN