perf: remove split partials memset and overlap decode tile loads

- alloc_split_partials now uses torch::empty: the split kernel writes every slot it owns, so the per-call zeros/full memset was pure overhead (2 kernels per layer per step)
- decode split-KV MMA kernels now run a true multi-stage cp.async pipeline (wait_group<STAGES-1> instead of wait_group<0>), keeping STAGES-1 tile loads in flight; the old wait_group<0> serialized load and compute so deeper STAGES made no difference
- add a fallback path when ntiles < STAGES to avoid a race on the last tile
This commit is contained in:
2026-07-31 22:37:44 +08:00
parent 21ddead238
commit 530d280e33
4 changed files with 64 additions and 44 deletions
-1
View File
@@ -30,7 +30,6 @@ Layout convention: all q/k/v are ``[batch, seq_len, n_heads, head_dim]``
import contextvars
import enum
import math
from abc import ABC, abstractmethod
from contextlib import contextmanager
from typing import Optional, Union