diff --git a/atom/config.py b/atom/config.py index a651dd9299..bdaae6af4e 100644 --- a/atom/config.py +++ b/atom/config.py @@ -237,6 +237,7 @@ def compute_hash(self) -> str: factors.append(self.local_cache_dir) factors.append(self.cudagraph_capture_sizes) factors.append(self.cuda_graph_sizes) + factors.append(bool(os.getenv("ATOM_ENABLE_PREZERO"))) return hashlib.sha256(str(factors).encode()).hexdigest() diff --git a/atom/model_ops/attention_mla.py b/atom/model_ops/attention_mla.py index ea846ebed3..50d6f6f5c7 100644 --- a/atom/model_ops/attention_mla.py +++ b/atom/model_ops/attention_mla.py @@ -31,6 +31,9 @@ fused_qk_rope_concat_and_cache_mla_seg = None from aiter.dist.parallel_state import get_dp_group from aiter.mla import mla_decode_fwd, mla_prefill_fwd +from aiter.tuned_gemm import tgemm + +from atom.model_ops.prezero import prezero_active from aiter.ops.triton.attention.mla import ( mla_decode_fwd as triton_shuffle_mla_decode_fwd, ) @@ -350,7 +353,7 @@ def process_weights_after_loading(self): ) @mark_trace(prefix="v_up_proj_and_o_proj", torch_compile=False) - def _v_up_proj_and_o_proj(self, x): + def _v_up_proj_and_o_proj(self, x, oproj_prezero=None): # Convert from (B, N, L) to (N, B, L) x = x.view(-1, self.num_heads, self.kv_lora_rank).transpose(0, 1) # Multiply (N, B, L) x (N, L, V) -> (N, B, V), Convert from (N, B, V) to (B, N, V) @@ -382,14 +385,22 @@ def _v_up_proj_and_o_proj(self, x): ) # Convert from (B, N, V) to (B, N * V) x = x.reshape(-1, self.num_heads * self.v_head_dim) + if oproj_prezero is not None: + active = prezero_active(self.q_proj.prezero_n_total) + tgemm.mm(x, self.o_proj.weight, None, zero_init=not active, out=oproj_prezero) + return oproj_prezero return self.o_proj(x) @mark_trace(prefix="q_proj_and_k_up_proj", torch_compile=False) - def _q_proj_and_k_up_proj(self, x, x_scale=None): - q_nope, q_pe = ( - self.q_proj(x, x_scale) - .view(-1, self.num_heads, self.qk_head_dim) - .split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + def _q_proj_and_k_up_proj(self, x, x_scale=None, qb_prezero=None): + if qb_prezero is not None: + active = prezero_active(self.q_proj.prezero_n_base) + tgemm.mm(x, self.q_proj.weight, None, zero_init=not active, out=qb_prezero) + q_proj_out = qb_prezero + else: + q_proj_out = self.q_proj(x, x_scale) + q_nope, q_pe = q_proj_out.view(-1, self.num_heads, self.qk_head_dim).split( + [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 ) # Convert from (B, N, P) to (N, B, P) @@ -956,6 +967,7 @@ def _forward_decode( q: torch.Tensor, kv_c_and_k_pe_cache: torch.Tensor, attn_metadata: AttentionMetaData, + oproj_prezero: Optional[torch.Tensor] = None, ) -> torch.Tensor: assert kv_c_and_k_pe_cache.numel() > 0 assert attn_metadata is not None @@ -1116,7 +1128,7 @@ def _forward_decode( if self.head_repeat_factor > 1: o = o[:, :: self.head_repeat_factor, :].contiguous() - return self._v_up_proj_and_o_proj(o) + return self._v_up_proj_and_o_proj(o, oproj_prezero=oproj_prezero) def forward_impl( self, @@ -1125,6 +1137,8 @@ def forward_impl( k_rope: torch.Tensor, positions: torch.Tensor = None, q_scale: Optional[torch.Tensor] = None, + qb_prezero: Optional[torch.Tensor] = None, + oproj_prezero: Optional[torch.Tensor] = None, ) -> torch.Tensor: # kv_cache = self.kv_cache forward_context: ForwardContext = get_forward_context() @@ -1206,7 +1220,9 @@ def forward_impl( prefill_q, k_nope, k_rope, kv_cache, attn_metadata ) else: - q_nope, q_rope = self._q_proj_and_k_up_proj(q, x_scale=q_scale) + q_nope, q_rope = self._q_proj_and_k_up_proj( + q, x_scale=q_scale, qb_prezero=qb_prezero + ) if self.use_seg_mla: # Seg path: allocate q_out with a padded last dim so each head row @@ -1296,7 +1312,9 @@ def forward_impl( if context.is_prefill: output = self._forward_prefill_mla(q_out, kv_cache, attn_metadata) else: - output = self._forward_decode(q_out, kv_cache, attn_metadata) + output = self._forward_decode( + q_out, kv_cache, attn_metadata, oproj_prezero=oproj_prezero + ) return output @@ -1310,6 +1328,8 @@ def forward( positions: torch.Tensor = None, q_scale: Optional[torch.Tensor] = None, output: torch.Tensor = None, + qb_prezero: Optional[torch.Tensor] = None, + oproj_prezero: Optional[torch.Tensor] = None, **kwargs, ) -> torch.Tensor: return self.forward_impl( @@ -1318,6 +1338,8 @@ def forward( k_rope=k_rope, positions=positions, q_scale=q_scale, + qb_prezero=qb_prezero, + oproj_prezero=oproj_prezero, ) diff --git a/atom/model_ops/base_attention.py b/atom/model_ops/base_attention.py index a086a87b62..dba304ccbf 100644 --- a/atom/model_ops/base_attention.py +++ b/atom/model_ops/base_attention.py @@ -327,6 +327,8 @@ def fake_( layer_name: str, use_mla: bool, qkv: torch.Tensor, + qb_prezero: Optional[torch.Tensor] = None, + oproj_prezero: Optional[torch.Tensor] = None, ) -> torch.Tensor: output_shape = list(q.shape) # If we fusion rmsnorm and quant, the input dtype is fp8, but actually we use bf16 for output. @@ -342,7 +344,9 @@ def fake_( # Dynamo will not try to inspect any of the internal operations for prefill or decode # This way, although attention operation is complicated, # we can still capture the model's computation graph as a full-graph -@mark_spliting_op(is_custom=True, gen_fake=fake_, mutates_args=[]) +@mark_spliting_op( + is_custom=True, gen_fake=fake_, mutates_args=["qb_prezero", "oproj_prezero"] +) def unified_attention_with_output_base( q: torch.Tensor, q_scale: Optional[torch.Tensor], @@ -352,6 +356,8 @@ def unified_attention_with_output_base( layer_name: str, use_mla: bool, qkv: torch.Tensor, + qb_prezero: Optional[torch.Tensor] = None, + oproj_prezero: Optional[torch.Tensor] = None, ) -> torch.Tensor: atom_config = get_current_atom_config() self = atom_config.compilation_config.static_forward_context[layer_name] @@ -362,6 +368,8 @@ def unified_attention_with_output_base( k_rope=v, positions=positions, q_scale=q_scale, + qb_prezero=qb_prezero, + oproj_prezero=oproj_prezero, ) else: return self.impl.forward( diff --git a/atom/model_ops/layernorm.py b/atom/model_ops/layernorm.py index 0367630d0b..43c06d0584 100644 --- a/atom/model_ops/layernorm.py +++ b/atom/model_ops/layernorm.py @@ -318,6 +318,7 @@ def forward( x: torch.Tensor, residual: torch.Tensor | None = None, x_scale: Optional[torch.Tensor] = None, + zero_fill: Optional[torch.Tensor] = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: if self.x_pad_to_multiple > 0: assert ( @@ -341,6 +342,7 @@ def forward( residual, self.weight, self.eps, + zero_fill=zero_fill, ) return x, residual else: diff --git a/atom/model_ops/paged_attention.py b/atom/model_ops/paged_attention.py index 7937ee111c..02fe78d6c2 100644 --- a/atom/model_ops/paged_attention.py +++ b/atom/model_ops/paged_attention.py @@ -116,9 +116,20 @@ def forward( positions: torch.Tensor = None, q_scale: Optional[torch.Tensor] = None, qkv: torch.Tensor = None, + qb_prezero: Optional[torch.Tensor] = None, + oproj_prezero: Optional[torch.Tensor] = None, **kwargs, ): output = torch.ops.aiter.unified_attention_with_output_base( - query, q_scale, key, value, positions, self.layer_name, self.use_mla, qkv + query, + q_scale, + key, + value, + positions, + self.layer_name, + self.use_mla, + qkv, + qb_prezero, + oproj_prezero, ) return output diff --git a/atom/model_ops/prezero.py b/atom/model_ops/prezero.py new file mode 100644 index 0000000000..63baefb0b1 --- /dev/null +++ b/atom/model_ops/prezero.py @@ -0,0 +1,89 @@ +"""Split-K GEMM "prezero" helpers. + +A split-K GEMM can skip its in-kernel output zero-init (NOZINIT) and atomic-add +into a buffer the preceding fused allreduce-rmsnorm zeroed for free, when that +pays off for the current batch. The decision lives inside opaque custom ops so +it re-runs at each per-bs cudagraph capture instead of being frozen by the one +prefill-time compile. +""" + +from typing import Optional + +import torch +from aiter import is_prezero_free +from aiter.dist.communication_op import tensor_model_parallel_fused_allreduce_rmsnorm +from aiter.jit.utils.torch_guard import torch_compile_guard +from aiter.tuned_gemm import tgemm + +from atom.config import get_current_atom_config +from atom.utils.forward_context import get_forward_context + + +def _prezero_hidden() -> int: + # AR out_hidden_dim == hidden_size; sets the zero-fill CTA width. + return get_current_atom_config().hf_config.hidden_size + + +def prezero_active(n_total: int) -> bool: + import os + + if not os.getenv("ATOM_ENABLE_PREZERO"): + return False + ctx = get_forward_context().context + return (not ctx.is_prefill) and is_prezero_free( + ctx.graph_bs, n_total, _prezero_hidden() + ) + + +@torch_compile_guard( + mutates_args=["out"], + gen_fake=lambda out, a, weight, n_total, bias=None: None, +) +def mm_maybe_prezero_( + out: torch.Tensor, + a: torch.Tensor, + weight: torch.Tensor, + n_total: int, + bias: Optional[torch.Tensor] = None, +) -> None: + active = prezero_active(n_total) + tgemm.mm(a, weight, bias, zero_init=not active, out=out) + + +@torch_compile_guard( + mutates_args=["zero_buf"], + gen_fake=lambda x, residual, weight, eps, zero_buf, n_total, n_base=0: ( + torch.empty_like(x), + torch.empty_like(residual), + ), +) +def ar_rmsnorm_maybe_prezero_( + x: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float, + zero_buf: torch.Tensor, + n_total: int, + n_base: int = 0, +) -> tuple[torch.Tensor, torch.Tensor]: + # zero the largest free cumulative prefix; drop only the overflowing tail. + import os + + zf = None + if os.getenv("ATOM_ENABLE_PREZERO"): + ctx = get_forward_context().context + if not ctx.is_prefill: + bs = ctx.graph_bs + h = x.shape[-1] + if is_prezero_free(bs, n_total, h): + zf = zero_buf + elif n_base and is_prezero_free(bs, n_base, h): + m = x.shape[0] + zf = zero_buf.view(-1)[: m * n_base].view(m, n_base) + return tensor_model_parallel_fused_allreduce_rmsnorm( + x.contiguous(), + residual, + weight, + eps, + zero_fill=zf, + ) diff --git a/atom/models/deepseek_v2.py b/atom/models/deepseek_v2.py index bd72e39cb2..b943a4a524 100644 --- a/atom/models/deepseek_v2.py +++ b/atom/models/deepseek_v2.py @@ -1941,6 +1941,26 @@ def __init__( if layer_quant_dtype in (dtypes.fp8, dtypes.fp4x2): self.fuse_qknorm_quant = True + # split-K prezero widths, read by the pre-split compile pass. qkv_a + q_b + # ride the input_layernorm AR donor (base); o_proj adds an overflowing + # tail gated independently (greedy). Only bf16 (a16w16) GEMMs qualify. + if self.q_lora_rank is not None: + n_qb = ( + self.q_b_proj.weight.shape[0] + if self.q_b_proj.quant_type.value == QuantType.No.value + else 0 + ) + n_oproj = ( + self.o_proj.weight.shape[0] + if self.o_proj.quant_type.value == QuantType.No.value + else 0 + ) + self.q_b_proj.prezero_n_qkva = self.fused_qkv_a_proj.weight.shape[0] + self.q_b_proj.prezero_n_qb = n_qb + self.q_b_proj.prezero_n_oproj = n_oproj + self.q_b_proj.prezero_n_base = self.q_b_proj.prezero_n_qkva + n_qb + self.q_b_proj.prezero_n_total = self.q_b_proj.prezero_n_base + n_oproj + def forward( self, positions: torch.Tensor, diff --git a/atom/models/deepseek_v2_prezero.py b/atom/models/deepseek_v2_prezero.py new file mode 100644 index 0000000000..6d405d49b6 --- /dev/null +++ b/atom/models/deepseek_v2_prezero.py @@ -0,0 +1,168 @@ +"""Split-K prezero pass. Graph surgery on the FULL (pre-split) dynamo graph -- +the only stage where the opaque attention op is still one node, so its +qb_prezero/oproj_prezero slots (q_b, o_proj live inside it) can be threaded. + +Per input_layernorm donor it allocates one [m, n_total] buffer sliced into +qkv_a (visible gemm) | q_b | o_proj, wires the fused_allreduce_rmsnorm_ -> +ar_rmsnorm_maybe_prezero_ (greedy zeroing up to n_base), rewrites the qkv_a gemm +-> mm_maybe_prezero_, and threads the q_b/o_proj slices into the attn op. +Post_attn gate donors get the single-buffer form (no attention threading). +""" +import operator +import re + +import torch + + +def _opname(t): + try: + n = t.name() + return n.rsplit(".", 1)[0] if "." in n else n + except Exception: + return getattr(t, "__name__", str(t)) + + +def _is(target, name): + return _opname(target).replace("::", ".").split(".default")[0].endswith(name) + + +def _resolve(qualname): + ns, name = qualname.split("::") + pkt = getattr(torch.ops, ns, None) + op = getattr(pkt, name, None) if pkt else None + return getattr(op, "default", None) if op else None + + +def _weight_layer(node): + """(layer_idx, kind) from a norm-weight placeholder feeding a donor, else None.""" + if not isinstance(node, torch.fx.Node): + return None + m = re.search(r"layers_modules_(\d+)_modules_(input_layernorm|post_attention_layernorm)", node.name) + return (int(m.group(1)), m.group(2)) if m else None + + +def _attn_layer(node): + if node.op != "call_function" or not _is(node.target, "unified_attention_with_output_base"): + return None + lname = node.args[5] if len(node.args) > 5 else None + m = re.search(r"layers\.(\d+)\.", lname) if isinstance(lname, str) else None + return int(m.group(1)) if m else None + + +def deepseek_v2_prezero_pass(gm, n_total_map=None, n_base_map=None): + g = gm.graph if hasattr(gm, "graph") else gm + n_total_map = n_total_map or {} + stats = {"qkv": 0, "gate": 0, "qb": 0, "oproj": 0} + ar_pz = _resolve("aiter::ar_rmsnorm_maybe_prezero_") + mm_pz = _resolve("aiter::mm_maybe_prezero_") + if ar_pz is None or mm_pz is None: + print("[prezero] maybe_prezero ops not registered -> no-op", flush=True) + return gm + + # index attention nodes by layer + attn_by_layer = {} + for n in g.nodes: + li = _attn_layer(n) + if li is not None: + attn_by_layer[li] = n + + aten = torch.ops.aten + for donor in list(g.nodes): + if donor.op != "call_function" or not _is(donor.target, "fused_allreduce_rmsnorm_"): + continue + if len(donor.args) < 4: + continue + wl = _weight_layer(donor.args[2]) + if wl is None: + continue + layer_idx, kind = wl + + # consumer gemm: donor -> getitem[0] -> gemm_a16w16 + g0 = next((u for u in donor.users if u.target is operator.getitem and u.args[1] == 0), None) + if g0 is None: + continue + gemm = next((u for u in g0.users if u.op == "call_function" and _is(u.target, "gemm_a16w16")), None) + if gemm is None: + continue + val = gemm.meta.get("val", gemm.meta.get("example_value")) + try: + n_out = int(val.shape[-1]) + except Exception: + continue + xval = donor.args[0].meta.get("val", donor.args[0].meta.get("example_value")) \ + if isinstance(donor.args[0], torch.fx.Node) else None + if xval is None: + continue + dtype, device = xval.dtype, xval.device + + x, res, w, eps = donor.args[0], donor.args[1], donor.args[2], donor.args[3] + bias = gemm.args[2] if len(gemm.args) > 2 else gemm.kwargs.get("bias") + normed = gemm.args[0] # == g0 + + # combined input_layernorm donor = qkv_a (visible gemm) + q_b + o_proj + # (both inside the opaque attn op); post_attn gate donor = single gemm. + n_total, n_base, attn = n_out, n_out, None + if kind == "input_layernorm": + attn = attn_by_layer.get(layer_idx) + lname = attn.args[5] if attn is not None else None + nt = n_total_map.get(lname) + nb = (n_base_map or {}).get(lname) + if nt and nt > n_out: + n_total = nt + n_base = nb if (nb and n_out <= nb <= nt) else nt + n_qb = n_base - n_out + n_oproj = n_total - n_base + + with g.inserting_before(donor): + m = g.call_function(aten.sym_size.int, (x, 0)) + buf = g.call_function(aten.empty.memory_format, ([m, n_total],), + {"dtype": dtype, "device": device}) + qb = oproj = None + if n_qb > 0 or n_oproj > 0: + flat = g.call_function(aten.view.default, (buf, [-1])) + o1 = g.call_function(operator.mul, (m, n_out)) + qkva = g.call_function(aten.view.default, + (g.call_function(operator.getitem, (flat, slice(None, o1, None))), [m, n_out])) + off = o1 + if n_qb > 0: + o2 = g.call_function(operator.mul, (m, n_base)) + qb = g.call_function(aten.view.default, + (g.call_function(operator.getitem, (flat, slice(off, o2, None))), [m, n_qb])) + off = o2 + if n_oproj > 0: + oproj = g.call_function(aten.view.default, + (g.call_function(operator.getitem, (flat, slice(off, None, None))), [m, n_oproj])) + else: + qkva = buf + + # donor -> ar_rmsnorm_maybe_prezero_(x, res, w, eps, buf, n_total, n_base) + with g.inserting_before(donor): + new_donor = g.call_function(ar_pz, (x, res, w, eps, buf, n_total, n_base)) + donor.replace_all_uses_with(new_donor) + g.erase_node(donor) + + # gemm -> mm_maybe_prezero_(qkva, normed, W, n_base); consumers read qkva + with g.inserting_before(gemm): + g.call_function(mm_pz, (qkva, normed, gemm.args[1], n_base), + {} if bias is None else {"bias": bias}) + gemm.replace_all_uses_with(qkva) + g.erase_node(gemm) + stats["qkv" if kind == "input_layernorm" else "gate"] += 1 + + # thread slices into the attn op: qb_prezero=arg[8], oproj_prezero=arg[9] + if attn is not None and (qb is not None or oproj is not None): + a = list(attn.args) + while len(a) < 10: + a.append(None) + if qb is not None: + a[8] = qb + stats["qb"] += 1 + if oproj is not None: + a[9] = oproj + stats["oproj"] += 1 + attn.args = tuple(a) + + g.lint() + print(f"[prezero] rewrote qkv_a={stats['qkv']} gate={stats['gate']} " + f"qb_threaded={stats['qb']} oproj_threaded={stats['oproj']}", flush=True) + return gm diff --git a/atom/utils/backends.py b/atom/utils/backends.py index c031ca3d43..55c33d4840 100644 --- a/atom/utils/backends.py +++ b/atom/utils/backends.py @@ -645,6 +645,30 @@ def __call__(self, graph: fx.GraphModule, example_inputs) -> Callable: self.graph = graph # self.configure_post_pass() + # aiter pre-split prezero hook (env-gated by ATOM_ENABLE_PREZERO). Runs + # on the FULL graph before splitting, where the opaque attention node + # (with its qb_prezero/oproj_prezero slots) is still visible -- lets the + # pass wire the attn-internal q_b/o_proj prezero automatically. Widths + # (n_total/n_base) are static module attrs not present in the graph. + import os as _os + if _os.getenv("ATOM_ENABLE_PREZERO"): + try: + from atom.models.deepseek_v2_prezero import deepseek_v2_prezero_pass + + _nt_map, _nb_map = {}, {} + for _ln, _mod in self.compilation_config.static_forward_context.items(): + for _sub in _mod.modules(): + _nt = getattr(_sub, "prezero_n_total", None) + if _nt: + _nt_map[_ln] = int(_nt) + _nb_map[_ln] = int(getattr(_sub, "prezero_n_base", _nt)) + break + deepseek_v2_prezero_pass(graph, _nt_map, _nb_map) + graph.recompile() if hasattr(graph, "recompile") else None + except Exception as _e: + import traceback as _tb + print(f"[prezero] {_e}\n{_tb.format_exc()}", flush=True) + self.split_gm, self.piecewise_graphs = split_graph( graph, self.compilation_config.splitting_ops )