diff --git a/deepinfer/model_executor/models/v4_modules/attn.py b/deepinfer/model_executor/models/v4_modules/attn.py index e20f586..0fca8c8 100644 --- a/deepinfer/model_executor/models/v4_modules/attn.py +++ b/deepinfer/model_executor/models/v4_modules/attn.py @@ -43,14 +43,20 @@ from deepinfer.model_executor.models.v4_modules.compress import Compressor from deepinfer.model_executor.models.v4_modules.cross_layer import CrossLayerRuntime, SharedLayerTopology from deepinfer.model_executor.models.v4_modules.index_attn import MQA from deepinfer.model_executor.models.v4_modules.mla_attn import LoraRopeCompressFlashMHAAutoRecomp, runtime_attn_sink -from deepinfer.model_executor.models.v4_modules.utils import convert_rmsnorm, convert_rope, get_enable_timer_layer, merge_linear_weights +from deepinfer.model_executor.models.v4_modules.utils import ( + convert_rmsnorm, + convert_rope, + decode_part_timer, + get_enable_timer_layer, + merge_linear_weights, +) from deepinfer.model_executor.models.v4_modules.yoco import is_yoco_consumer_layer from deepinfer.neuron_ops import nki_gemm from deepinfer.neuron_ops import regions as neuron_regions from deepinfer.model_executor.parallel_utils.parallel_state import PG from deepinfer.model_executor.parallel_utils.userbuffer import userbuffer from deepinfer.model_executor.utils import Movable -from deepinfer.timer import cuda_graph_timer, cuda_timer +from deepinfer.timer import cuda_timer from deepinfer.worker.cache_engine import KVCache @@ -655,11 +661,8 @@ class V4Attention(V4AttentionBase): index_length=dense_index_len, ) - second_layer, _, same_layers = get_enable_timer_layer(self.layer_idx) fused_kwargs = self._fused_norm_rope_cast_kwargs(input_metadata.positions) - with cuda_graph_timer.record( - f"decode_core_attn-part-{self.layer_idx}", enabled=self.layer_idx == second_layer, scale=len(same_layers) - ): + with decode_part_timer("decode_core_attn-part", self): context, _ = decode_mla_mqa_sparse( q_b, head_dim_v=self.qk_head_dim, diff --git a/deepinfer/model_executor/models/v4_modules/index_attn.py b/deepinfer/model_executor/models/v4_modules/index_attn.py index 46cf987..4122290 100644 --- a/deepinfer/model_executor/models/v4_modules/index_attn.py +++ b/deepinfer/model_executor/models/v4_modules/index_attn.py @@ -33,8 +33,8 @@ from deepinfer.model_executor.layers.tilelang_ops.load_paged_kv_cache import loa from deepinfer.model_executor.models.v4_modules.compress import KProject from deepinfer.model_executor.models.v4_modules.cross_layer import CrossLayerRuntime, SharedLayerTopology from deepinfer.model_executor.models.v4_modules.mla_attn import MQAIndexer -from deepinfer.model_executor.models.v4_modules.utils import get_enable_timer_layer -from deepinfer.timer import cuda_graph_timer, cuda_timer +from deepinfer.model_executor.models.v4_modules.utils import decode_part_timer, get_enable_timer_layer +from deepinfer.timer import cuda_timer def _align_mqa_weights(weights: torch.Tensor) -> torch.Tensor: @@ -476,10 +476,7 @@ class MQA(nn.Module): cross_layer_runtime, ) - second_layer, _, same_layers = get_enable_timer_layer(self.layer_idx) - with cuda_graph_timer.record( - f"decode_index_attn-part-{self.layer_idx}", enabled=self.layer_idx == second_layer, scale=len(same_layers) - ): + with decode_part_timer("decode_index_attn-part", self): p_sum = self._compute_decode_p_sum( (q_fp8or4, q_sf), kv_cache, diff --git a/deepinfer/model_executor/models/v4_modules/utils.py b/deepinfer/model_executor/models/v4_modules/utils.py index d165a74..ae45b96 100644 --- a/deepinfer/model_executor/models/v4_modules/utils.py +++ b/deepinfer/model_executor/models/v4_modules/utils.py @@ -5,6 +5,7 @@ layers; `merge_linear_weights` and `init_dummy_weight` reshape and seed runtime The `DeepseekDecoderLayer*` registries feed the timer sampling groups. """ +import contextlib from collections.abc import Callable import torch @@ -141,6 +142,19 @@ def create_deepseek_decode_layer(layer_str: str, layer_idx: int) -> None: DeepseekDecoderLayerIdx[layer_idx] = layer_str +def decode_part_timer(name: str, module: nn.Module) -> contextlib.AbstractContextManager: + """``cuda_graph_timer`` for a decode part, sampled on the second layer of its kind. A no-op inside a compiled + Neuron region: reading ``module.layer_idx`` there makes Dynamo guard on its value (one recompile per layer + instead of one per layer kind).""" + if torch.compiler.is_compiling(): + return contextlib.nullcontext() + from deepinfer.timer import cuda_graph_timer + + layer_idx = module.layer_idx + second_layer, _, same_layers = get_enable_timer_layer(layer_idx) + return cuda_graph_timer.record(f"{name}-{layer_idx}", enabled=layer_idx == second_layer, scale=len(same_layers)) + + def get_enable_timer_layer(layer_idx: int) -> tuple[int, int, list[int]]: same_layers = DeepseekDecoderLayerType[DeepseekDecoderLayerIdx[layer_idx]] second_layer = same_layers[1] if len(same_layers) > 1 else same_layers[0] diff --git a/deepinfer/neuron_ops/nki_attention.py b/deepinfer/neuron_ops/nki_attention.py index 388c6cc..7e526d4 100644 --- a/deepinfer/neuron_ops/nki_attention.py +++ b/deepinfer/neuron_ops/nki_attention.py @@ -255,10 +255,63 @@ def flash_mla_decode( if mask is not None: invalid = invalid | mask scopes.append((cache, torch.where(invalid, torch.zeros_like(idx), idx), invalid)) - out = sparse_decode(q.reshape(rows, h_q, d), scopes, softmax_scale, attn_sink) + if rows == PREFILL_BUCKET: + out = _chunk_attention(q.reshape(rows, h_q, d), scopes, softmax_scale, attn_sink) + else: + out = sparse_decode(q.reshape(rows, h_q, d), scopes, softmax_scale, attn_sink) return out.view(b, s_q, h_q, d) +#: 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 + + +def _chunk_attention( + q: torch.Tensor, + scopes: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]], + softmax_scale: float, + attn_sink: torch.Tensor | None, +) -> torch.Tensor: + """``sparse_decode`` for a prompt chunk run as a decode step (issue #26): rows are one sequence's tokens. + + A dense compressed-KV scope lists, for every token, the positions before it: the lists are prefixes of the + longest one. Gathering every row's list (``[rows, max_model_len / ratio, 512]``) would take GBs and a NEFF of + as many DMA descriptors, so such a scope is gathered once and masked per row. Top-k and SWA scopes stay per row. + Same numerics as the kernel: fp32 logits and softmax with the sink in the denominator, bf16 P @ V. + """ + rows, heads, d = q.shape + qf = q.float() + logits, values = [], [] + for cache, idx, invalid in scopes: + if idx.shape[1] >= _SHARED_SCOPE_MIN: + longest = (~invalid).sum(dim=1).argmax() + kv = gather_paged_kv(cache, idx.index_select(0, longest.view(1))).view(-1, d).float() # [topk, 512] + s = torch.matmul(qf.view(rows * heads, d), kv.transpose(0, 1)).view(rows, heads, -1) + else: + kv = gather_paged_kv(cache, idx).float() # [rows, topk, 512] + s = torch.bmm(qf, kv.transpose(1, 2)) + logits.append(torch.where(invalid.unsqueeze(1), float("-inf"), s * softmax_scale)) + values.append(kv) + peak = torch.stack([x.amax(dim=-1) for x in logits], dim=-1).amax(dim=-1) # [rows, heads] + sink = attn_sink.float().view(1, heads) if attn_sink is not None else None + if sink is not None: + peak = torch.maximum(peak, sink.expand(rows, heads)) + peak = torch.where(torch.isinf(peak), torch.zeros_like(peak), peak) + total = torch.exp(sink - peak) if sink is not None else torch.zeros_like(peak) + out = torch.zeros((rows, heads, d), dtype=torch.float32, device=q.device) + for s, kv in zip(logits, values, strict=True): + p = torch.exp(s - peak.unsqueeze(-1)) + total = total + p.sum(dim=-1) + p = p.to(torch.bfloat16).float() + if kv.dim() == 2: + out = out + torch.matmul(p.view(rows * heads, -1), kv).view(rows, heads, d) + else: + out = out + torch.bmm(p, kv) + valid_any = torch.stack([(~inv).any(dim=1) for _, _, inv in scopes], dim=-1).any(dim=-1).view(rows, 1, 1) + out = out / total.clamp_min(1e-30).unsqueeze(-1) + return torch.where(valid_any, out, torch.zeros_like(out)).to(torch.bfloat16) + + # --------------------------------------------------------------------------- # # indexer scores (deep_gemm.fp8_fp4_paged_mqa_logits) # --------------------------------------------------------------------------- #