diff --git a/deepinfer/model_executor/layers/attention_backend.py b/deepinfer/model_executor/layers/attention_backend.py index 84f25a4..20b6cb5 100644 --- a/deepinfer/model_executor/layers/attention_backend.py +++ b/deepinfer/model_executor/layers/attention_backend.py @@ -88,6 +88,7 @@ def decode_mla_mqa_sparse( indices_in_kvcache_extra: torch.Tensor | None = None, index_length_extra: torch.Tensor | None = None, fused_norm_rope_cast_kwargs: dict | None = None, + neuron_dense: dict | None = None, ) -> tuple[torch.Tensor | tuple[torch.Tensor, torch.Tensor], torch.Tensor]: """ Args: @@ -117,6 +118,7 @@ def decode_mla_mqa_sparse( extra_indices_in_kvcache=indices_in_kvcache_extra, extra_topk_length=index_length_extra, **fused_norm_rope_cast_kwargs, + **({"neuron_dense": neuron_dense} if neuron_dense is not None else {}), ) return (output_fp8, output_sf), lse @@ -139,6 +141,7 @@ def decode_mla_mqa_sparse( topk_length=index_length, extra_topk_length=index_length_extra, enable_batch_invariant=batch_invariant, + **({"neuron_dense": neuron_dense} if neuron_dense is not None else {}), ) output, lse = output.squeeze(1), lse.squeeze(-1) if padded_heads is not None: diff --git a/deepinfer/model_executor/models/v4_modules/attn.py b/deepinfer/model_executor/models/v4_modules/attn.py index b5e2171..c14f88d 100644 --- a/deepinfer/model_executor/models/v4_modules/attn.py +++ b/deepinfer/model_executor/models/v4_modules/attn.py @@ -745,6 +745,9 @@ class V4Attention(V4AttentionBase): ) fused_kwargs = self._fused_norm_rope_cast_kwargs(input_metadata.positions) + dense = self._neuron_dense_chunk(q_b, kv_ca_index, kv_cache, input_metadata) + if dense is not None: + attn_kv_args["neuron_dense"] = dense with decode_part_timer("decode_core_attn-part", self): context, _ = decode_mla_mqa_sparse( q_b, @@ -758,6 +761,37 @@ class V4Attention(V4AttentionBase): ) return context + def _neuron_dense_chunk( + self, q_b: torch.Tensor, kv_ca_index: torch.Tensor | None, kv_cache: KVCache, input_metadata: InputMetadataV4 + ) -> dict | None: + """Neuron: the dense-attention inputs for a single request's prompt chunk (nki_attention.dense_chunk), or None + to keep the gathered kernel. A top-k scope qualifies only when the context bucket's compressed keys fit the + top-k width (then the indexer selects all of them).""" + if not (PLATFORM.is_neuron and self.use_swa and self.has_packed_kv and kv_ca_index is not None): + return None + from deepinfer.neuron_ops import nki_attention + + if not (nki_attention.PREFILL_DENSE and nki_attention.single_request_chunk() and q_b.shape[0] == nki_attention.PREFILL_BUCKET): + return None + ctx = getattr(input_metadata, "neuron_ctx_limit", None) or self.max_model_len + block_tables = input_metadata.block_tables + if block_tables is None or kv_cache.ca_cache is None: + return None + ratio = int(self.ratio_key) + if self.is_sparse: + # The indexer ranks the first klen compressed keys; when the bucket's keys fit the width, it keeps all. + klen = (getattr(input_metadata, "klen", None) or {}).get(self.ratio_key) + if klen is None or ctx // ratio > kv_ca_index.shape[-1]: + return None + offset_len = torch.stack([torch.zeros_like(klen), klen], dim=1) + else: + offset_len = (input_metadata.kv_ca_index_offset_len or {}).get(self.ratio_key) + if offset_len is None: + return None + block_tokens = kv_cache.ca_cache.shape[1] * ratio + blocks = min(block_tables.shape[1], -(-ctx // block_tokens)) + return {"block_table": block_tables[0], "blocks": blocks, "offset_len": offset_len} + def _store_yoco_full_kv( self, cross_layer_runtime: CrossLayerRuntime, diff --git a/deepinfer/neuron_ops/nki_attention.py b/deepinfer/neuron_ops/nki_attention.py index 1829743..4aed6b8 100644 --- a/deepinfer/neuron_ops/nki_attention.py +++ b/deepinfer/neuron_ops/nki_attention.py @@ -272,9 +272,21 @@ def flash_mla_decode( extra_indices: torch.Tensor | None, extra_topk_length: torch.Tensor | None, softmax_scale: float, + dense: dict | None = None, ) -> torch.Tensor: """The output of ``flash_mla_with_kvcache`` (decode, ``q [b, s_q, h_q, 512]``, ``indices [b, s_q, topk]``, -1 = - invalid, ``topk_length [b]`` masks the tail) -> ``[b, s_q, h_q, 512]`` bf16.""" + invalid, ``topk_length [b]`` masks the tail) -> ``[b, s_q, h_q, 512]`` bf16. ``dense``: see dense_chunk.""" + if dense is not None and extra_k_cache is not None: + if _DENSE_CHECK: + dense = {**dense, "extra_indices": extra_indices, "extra_topk_length": extra_topk_length} + out = dense_chunk(q, k_cache, indices, topk_length, extra_k_cache, dense, softmax_scale, attn_sink) + if not _DENSE_CHECK: + return out + ref = flash_mla_decode( + q, k_cache, indices, topk_length, attn_sink, extra_k_cache, extra_indices, extra_topk_length, softmax_scale + ) + _report_dense_check(out, ref, dense.get("_own")) + return ref b, s_q, h_q, d = q.shape rows = b * s_q scopes, bounds = [], [] @@ -336,6 +348,18 @@ def set_table_bounds(reqs: int | None, ctx: int | None) -> None: _TABLE_BOUNDS["reqs"], _TABLE_BOUNDS["ctx"] = reqs, ctx +#: Whether the step being replayed is a prompt chunk of at most one request (graph runner, max_num_seqs == 1). +_SINGLE = [False] + + +def set_single_request_chunk(single: bool) -> None: + _SINGLE[0] = single + + +def single_request_chunk() -> bool: + return _SINGLE[0] + + def _table_blocks(blocks: int, page: int) -> int: reqs, ctx = _TABLE_BOUNDS["reqs"], _TABLE_BOUNDS["ctx"] if not reqs or not ctx: @@ -470,10 +494,16 @@ def _dense_attention( scaled; each row's slots ``idx`` (unique per row) unmasked; the max over the slots only; the sink in the denominator; probabilities rounded to bf16 for the value product, as the kernel does. ``q [rows, heads, d]``, ``kv [T, d]`` (keys are values) -> ``[rows, heads, d]`` bf16.""" - rows, heads, d = q.shape + rows = q.shape[0] t = kv.shape[0] slot = torch.where(invalid, torch.full_like(idx, t), idx).long() # invalid slots -> a dump column allowed = torch.zeros((rows, t + 1), dtype=torch.float32, device=q.device).scatter(1, slot, 1.0)[:, :t] > 0 + return _dense_masked(q, kv, allowed, softmax_scale, sink) + + +def _dense_masked(q: torch.Tensor, kv: torch.Tensor, allowed: torch.Tensor, softmax_scale: float, sink: torch.Tensor) -> torch.Tensor: + """``_dense_attention`` given the ``[rows, T]`` mask of each row's slots.""" + heads = q.shape[1] kvf = kv.float() logits = torch.matmul(q.float(), kvf.t()) * softmax_scale # [rows, heads, t] logits = torch.where(allowed.unsqueeze(1), logits, torch.full_like(logits, EMPTY)) @@ -484,6 +514,133 @@ def _dense_attention( return out.to(torch.bfloat16) +#: DEEPINFER_NEURON_PREFILL_DENSE: a single request's prompt chunk attends densely over its own SWA ring and +#: compressed blocks, the mask built from positions (see dense_chunk), when every visible compressed key fits the +#: top-k width; otherwise (and for decode) the gathered kernel. +PREFILL_DENSE = os.environ.get("DEEPINFER_NEURON_PREFILL_DENSE", "1") == "1" + + +#: DEEPINFER_NEURON_PREFILL_DENSE_CHECK: also run the gathered kernel, log the largest difference, use the kernel's. +_DENSE_CHECK = os.environ.get("DEEPINFER_NEURON_PREFILL_DENSE_CHECK") == "1" + + +@torch._dynamo.disable +def _report_dense_check(out: torch.Tensor, ref: torch.Tensor, own: torch.Tensor | None = None) -> None: + o, r = out.float().cpu().flatten(0, 1), ref.float().cpu().flatten(0, 1) + d = (o - r).abs() + if own is not None: + keep = own.cpu().reshape(-1) + per_row = d.amax(dim=(-1, -2)) + worst = torch.argsort(per_row, descending=True)[:4] + print( + f"dense prefill check (own rows {int(keep.sum())}): max |diff| {per_row[keep].max().item():.4g}, " + f"rows differing {(per_row[keep] > 0).sum().item()}; padding rows max {per_row[~keep].max().item() if (~keep).any() else 0:.4g}; " + f"worst rows {[(int(i), bool(keep[i]), round(per_row[i].item(), 4)) for i in worst]}", + flush=True, + ) + return + print( + f"dense prefill check: max |diff| {d.max().item():.4g}, mean {d.mean().item():.3g}, ref max {r.abs().max().item():.4g}, " + f"rows differing {(d.amax(dim=(-1, -2)) > 0).sum().item()}/{d.shape[0] * d.shape[1]}", + flush=True, + ) + + +def dense_chunk( + q: torch.Tensor, + k_cache: torch.Tensor, + indices: torch.Tensor, + topk_length: torch.Tensor | None, + extra_k_cache: torch.Tensor, + hint: dict, + softmax_scale: float, + attn_sink: torch.Tensor | None, +) -> torch.Tensor: + """``flash_mla_decode`` for one request's prompt chunk, as two dense matmuls over its own keys. + + The gathered kernel attends each row to its ``window`` SWA slots and its ``topk`` compressed slots with a gather + and small matmuls per row (~9 ms a layer for a 512-row chunk). Here the tables are the request's SWA ring + (``k_cache[state]``) and its first ``hint["blocks"]`` compressed blocks in logical order (``hint["block_table"]``), + and each row's slots are given by arithmetic instead of a scatter: SWA slot ``s`` when ``(s - first) mod ring < + count`` (``get_kv_swa_index_decode``: ``count`` consecutive ring slots from the row's first index), compressed key + ``t`` when ``offset <= t < offset + len`` (``hint["offset_len"]``). For a top-k scope the caller guarantees the + visible keys fit the width, so the selection is all of them. ``q [b, s_q, h, 512]`` -> ``[b, s_q, h, 512]`` bf16. + """ + b, s_q, heads, d = q.shape + rows = b * s_q + device = q.device + ring = k_cache.shape[1] + idx = indices.reshape(rows, -1).to(torch.int32) + invalid = idx < 0 + mask = _length_mask(topk_length, b, s_q, idx.shape[1], device) + if mask is not None: + invalid = invalid | mask + count = (~invalid).to(torch.int32).sum(dim=1, keepdim=True) + first = idx[:, :1] + state = (idx.amax() // ring).reshape(1) # one request: one ring + swa = _gather_blocks(k_cache, state) # [ring, 512] + slots = torch.arange(ring, dtype=torch.int32, device=device).view(1, ring) + delta = slots - (first - first // ring * ring) + delta = torch.where(delta < 0, delta + ring, delta) + own = first // ring == state # padding rows of the chunk index a dummy slot of another ring: no keys + if _DENSE_CHECK: + hint["_own"] = own + allowed_swa = (delta < count) & own + blocks = hint["block_table"][: hint["blocks"]].to(torch.int32) + blocks = blocks.clamp(0, extra_k_cache.shape[0] - 1) # unallocated entries: masked below by the key count + ca = _gather_blocks(extra_k_cache, blocks) # [blocks * page, 512] + keys = torch.arange(ca.shape[0], dtype=torch.int32, device=device).view(1, -1) + offset_len = hint["offset_len"].to(torch.int32) + lo, n = offset_len[:rows, :1], offset_len[:rows, 1:2] + allowed_ca = (keys >= lo) & (keys < lo + n) & own + kv = torch.nan_to_num(torch.cat([swa, ca], dim=0), nan=0.0, posinf=0.0, neginf=0.0) # unwritten slots: never NaN + allowed = torch.cat([allowed_swa, allowed_ca], dim=1) + if attn_sink is None: + sink = torch.full((heads, 1), EMPTY, dtype=torch.float32, device=device) + else: + sink = attn_sink.to(torch.float32).reshape(heads, 1) + if _DENSE_CHECK and hint.get("extra_indices") is not None: + # The masks the indices give (scatter), for comparison. + ref_swa = torch.zeros((rows, ring + 1), dtype=torch.float32, device=device).scatter( + 1, torch.where(invalid | (idx // ring != state), ring, idx - state * ring).long(), 1.0)[:, :ring] > 0 + eidx = hint["extra_indices"].reshape(rows, -1).to(torch.int32) + einv = eidx < 0 + emask = _length_mask(hint.get("extra_topk_length"), b, s_q, eidx.shape[1], device) + if emask is not None: + einv = einv | emask + page = extra_k_cache.shape[1] + inv = torch.full((extra_k_cache.shape[0] + 1,), -1, dtype=torch.int32, device=device).scatter( + 0, blocks.long(), torch.arange(blocks.shape[0], dtype=torch.int32, device=device)) + safe = torch.where(einv, torch.zeros_like(eidx), eidx) + logical = inv.index_select(0, (safe // page).reshape(-1).long()).reshape(eidx.shape) * page + safe % page + bad = einv | (logical < 0) + ref_ca = torch.zeros((rows, ca.shape[0] + 1), dtype=torch.float32, device=device).scatter( + 1, torch.where(bad, ca.shape[0], logical).long(), 1.0)[:, : ca.shape[0]] > 0 + _report_masks(allowed_swa, ref_swa & own, allowed_ca, ref_ca & own, first, count, lo, n) + out = _dense_masked(q.reshape(rows, heads, d), kv, allowed, softmax_scale, sink) + out = torch.where(allowed.any(dim=1).view(rows, 1, 1), out, torch.zeros_like(out)) + return out.reshape(b, s_q, heads, d) + + +@torch._dynamo.disable +def _report_masks(a_swa, r_swa, a_ca, r_ca, first, count, lo, n) -> None: # noqa: ANN001 + a_swa, r_swa, a_ca, r_ca = (x.cpu() for x in (a_swa, r_swa, a_ca, r_ca)) + first, count, lo, n = (x.cpu().reshape(-1) for x in (first, count, lo, n)) + bs = (a_swa != r_swa).sum(1) + bc = (a_ca != r_ca).sum(1) + rows = [int(i) for i in torch.nonzero(bs + bc).reshape(-1)[:4]] + print( + f"dense mask check: swa rows wrong {(bs > 0).sum().item()}/{bs.numel()}, ca rows wrong {(bc > 0).sum().item()}; " + + "; ".join( + f"r{r}: first {first[r].item()} count {count[r].item()} ours_swa {a_swa[r].sum().item()} ref_swa {r_swa[r].sum().item()} " + f"lo {lo[r].item()} n {n[r].item()} ours_ca {a_ca[r].sum().item()} ref_ca {r_ca[r].sum().item()} " + f"ref_ca_max {int(torch.nonzero(r_ca[r]).max()) if r_ca[r].any() else -1}" + for r in rows + ), + flush=True, + ) + + #: A scope this wide is a dense compressed-KV read (every position up to the token's own), not a top-k selection. _SHARED_SCOPE_MIN = 1024 diff --git a/deepinfer/neuron_ops/vendor/flash_mla/flash_mla_interface.py b/deepinfer/neuron_ops/vendor/flash_mla/flash_mla_interface.py index dabb464..799d044 100644 --- a/deepinfer/neuron_ops/vendor/flash_mla/flash_mla_interface.py +++ b/deepinfer/neuron_ops/vendor/flash_mla/flash_mla_interface.py @@ -233,6 +233,7 @@ def _decode_core( d_v: int, softmax_scale: float, rms_norm_scale_factor: Optional[torch.Tensor], # [b, s_q, h_q] + neuron_dense: Optional[dict] = None, # DeepInfer Neuron: a prompt chunk's dense attention (nki_attention.dense_chunk) ) -> Tuple[torch.Tensor, torch.Tensor]: b, s_q, h_q, d_qk = q.shape rows = b * s_q @@ -243,7 +244,7 @@ def _decode_core( assert (extra_k_cache is None) == (extra_indices_in_kvcache is None) out = _nki.flash_mla_decode( q, k_cache, indices_in_kvcache, topk_length, attn_sink, - extra_k_cache, extra_indices_in_kvcache, extra_topk_length, softmax_scale, + extra_k_cache, extra_indices_in_kvcache, extra_topk_length, softmax_scale, neuron_dense, ) return out, torch.full((b, s_q, h_q), float("nan"), dtype=torch.float32, device=q.device) @@ -312,6 +313,7 @@ def flash_mla_with_kvcache( topk_length: Optional[torch.Tensor] = None, extra_topk_length: Optional[torch.Tensor] = None, enable_batch_invariant: bool = False, + neuron_dense: Optional[dict] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """Sparse MLA decoding with a paged quantized KV cache. @@ -414,6 +416,7 @@ def flash_mla_with_kvcache( head_dim_v, softmax_scale, None, + neuron_dense=neuron_dense, ) return out, lse.transpose(1, 2).contiguous() diff --git a/deepinfer/neuron_ops/vendor/flash_mla/fused_norm_rope_attn_rope_cast.py b/deepinfer/neuron_ops/vendor/flash_mla/fused_norm_rope_attn_rope_cast.py index d4bde87..7b2650e 100644 --- a/deepinfer/neuron_ops/vendor/flash_mla/fused_norm_rope_attn_rope_cast.py +++ b/deepinfer/neuron_ops/vendor/flash_mla/fused_norm_rope_attn_rope_cast.py @@ -144,6 +144,7 @@ def decode( extra_k_cache: Optional[torch.Tensor] = None, extra_indices_in_kvcache: Optional[torch.Tensor] = None, extra_topk_length: Optional[torch.Tensor] = None, + neuron_dense: Optional[dict] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Returns (out_fp8, out_sf, lse [s_q, h_q]). @@ -170,6 +171,7 @@ def decode( d_v, sm_scale, None if scale_factor is None else scale_factor.view(s_q, 1, h_q), + neuron_dense=neuron_dense, ) out = _rope(out.view(s_q, h_q, d_v), True, cos_sin_cache, token_positions, rope_dim) out_fp8, out_sf = _cast_out( diff --git a/deepinfer/worker/neuron_graph_runner.py b/deepinfer/worker/neuron_graph_runner.py index 1551b70..d7c8c70 100644 --- a/deepinfer/worker/neuron_graph_runner.py +++ b/deepinfer/worker/neuron_graph_runner.py @@ -200,7 +200,9 @@ class NeuronGraphRunner(CudaGraphRunner): # bucket (without a bound every prefill layer decoded the whole KV pool, issue #29). ctx = getattr(input_metadata_device, "neuron_ctx_limit", None) nki_attention.set_table_bounds(self.max_reqs if envs.DEEPINFER_NEURON_PREFILL_TABLE_BOUNDS and ctx else None, ctx) + nki_attention.set_single_request_chunk(self.max_reqs == 1) else: + nki_attention.set_single_request_chunk(False) reqs = -(-batch_size // (1 + envs.DEEPINFER_USE_NEXTN_DECODE)) nki_attention.set_table_bounds(reqs, getattr(input_metadata_device, "neuron_ctx_limit", None)) if not envs.DEEPINFER_NEURON_REGION_TIMING: