diff --git a/deepinfer/neuron_ops/nki_attention.py b/deepinfer/neuron_ops/nki_attention.py index b1c2479..64a6c37 100644 --- a/deepinfer/neuron_ops/nki_attention.py +++ b/deepinfer/neuron_ops/nki_attention.py @@ -277,7 +277,7 @@ def flash_mla_decode( invalid, ``topk_length [b]`` masks the tail) -> ``[b, s_q, h_q, 512]`` bf16.""" b, s_q, h_q, d = q.shape rows = b * s_q - scopes = [] + scopes, bounds = [], [] for cache, ind, length in ((k_cache, indices, topk_length), (extra_k_cache, extra_indices, extra_topk_length)): if cache is None: continue @@ -287,7 +287,15 @@ def flash_mla_decode( mask = _length_mask(length, b, s_q, topk, q.device) if mask is not None: invalid = invalid | mask + bound = None + if cache is k_cache and extra_k_cache is not None and _SWA_SUBBLOCK: + # The SWA ring (with an extra compressed scope, k_cache is the ring cache [states, ring, ...]): a "block" + # is a whole per-request ring (window + backtracking slots), and the table decoded every state's whole + # ring. Viewed as 64-slot sub-blocks, a request's rows reference only the few its window spans, so the + # decode no longer grows with the ring (2.5 -> 6.5 ms a layer from backtracking 128 to 512, issue #29). + cache, bound = _swa_subblocks(cache, topk, rows) scopes.append((cache, torch.where(invalid, torch.zeros_like(idx), idx), invalid)) + bounds.append(bound) q_rows = q.reshape(rows, h_q, d) pad = (1 << (rows - 1).bit_length()) - rows if rows > 1 else 0 if pad: @@ -303,7 +311,7 @@ def flash_mla_decode( if rows > 1 and sum(c.shape[0] * c.shape[1] for c, _, _ in scopes) <= _TABLE_MAX_SLOTS: # Multi-row steps (prompt chunks, DSpark verification): the per-row record decode costs linearly in rows # (25 ms a layer at 8 rows), the table decode is one pass over the cache. - out = _table_attention(q_rows, scopes, softmax_scale, attn_sink) + out = _table_attention(q_rows, scopes, softmax_scale, attn_sink, bounds) elif rows == PREFILL_BUCKET and any(idx.shape[1] >= _SHARED_SCOPE_MIN for _, idx, _ in scopes): # Only a dense compressed-KV scope needs the shared-gather path. Top-k and SWA scopes are per row either way, # so the kernel takes them as in decode (the torch path was 41.6 ms a layer for a 128-token chunk, bench4). @@ -361,11 +369,32 @@ def _compact_blocks(idx: torch.Tensor, blocks: int, page: int, size: int) -> tup return chosen, remapped.to(torch.int32) +#: DEEPINFER_NEURON_SWA_SUBBLOCK: decode only the 64-slot sub-blocks of the SWA ring a step's windows reference. +_SWA_SUBBLOCK = os.environ.get("DEEPINFER_NEURON_SWA_SUBBLOCK", "1") == "1" +_SUB = 64 + + +def _swa_subblocks(cache: torch.Tensor, width: int, rows: int) -> tuple[torch.Tensor, int | None]: + """``cache`` as ``[blocks * ring / 64, 64, ...]`` (the same linear slot index) and the number of sub-blocks the + step's rows can reference: per request, its window (``width`` slots, plus the step's rows) spans + ``ceil(.. / 64) + 2`` sub-blocks (an unaligned start, and the ring's wrap).""" + blocks, page = cache.shape[0], cache.shape[1] + if page <= _SUB or page % _SUB or cache.stride(0) != page * cache.stride(1): + return cache, None + reqs = _TABLE_BOUNDS["reqs"] + if not reqs: + return cache, None + sub = cache.view(blocks * (page // _SUB), _SUB, *cache.shape[2:]) + per_req = -(-(width + -(-rows // reqs)) // _SUB) + 2 + return sub, reqs * per_req + 1 + + def _table_attention( q: torch.Tensor, scopes: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]], softmax_scale: float, attn_sink: torch.Tensor | None, + bounds: list[int | None] | None = None, ) -> torch.Tensor: """Chunk attention the vllm-neuron#71 way (``gathered_attn_idx``): decode each cache once into a bf16 table and let the kernel gather each row's slots in SBUF. The per-row path decodes ``rows x topk`` records instead (128 x 640 @@ -373,10 +402,13 @@ def _table_attention( with a pool sized to the run (--num_kv_blocks).""" rows, heads, d = q.shape tables, idxs, masks, offset = [], [], [], 0 - for cache, idx, invalid in scopes: + for i, (cache, idx, invalid) in enumerate(scopes): blocks, page = cache.shape[0], cache.shape[1] idx = idx.to(torch.int32) used_blocks = _table_blocks(blocks, page) + bound = bounds[i] if bounds else None + if bound is not None: + used_blocks = min(blocks, 1 << (bound - 1).bit_length()) # a power of two, as in _table_blocks if used_blocks < blocks: # Decode only the blocks the rows reference (the step's requests), not the whole pool: the record decode # is integer work on GpSimd and cost 6-10 ms a layer over 160 blocks (device profile, bench47). Whole