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:
@@ -21,11 +21,15 @@ using bf16 = __nv_bfloat16;
|
||||
" (supported: 32, 64, 128, 256)"); \
|
||||
}
|
||||
|
||||
// The split kernel unconditionally writes every (batch, q_head, split) slot it
|
||||
// owns — including empty split ranges, which store m = -FLT_MAX so the combine
|
||||
// skips them. Allocators are therefore left uninitialized (torch::empty); the
|
||||
// per-call memset (torch::zeros / torch::full) was pure overhead.
|
||||
template<typename P>
|
||||
inline void alloc_split_partials(P& p) {
|
||||
auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
|
||||
auto o_part = torch::zeros(at::IntArrayRef{p.batch, p.q_head, MAX_SPLITS, p.head_dim}, fopt);
|
||||
auto ml_part = torch::full(at::IntArrayRef{p.batch, p.q_head, MAX_SPLITS, 2}, -FLT_MAX, fopt);
|
||||
auto o_part = torch::empty(at::IntArrayRef{p.batch, p.q_head, MAX_SPLITS, p.head_dim}, fopt);
|
||||
auto ml_part = torch::empty(at::IntArrayRef{p.batch, p.q_head, MAX_SPLITS, 2}, fopt);
|
||||
p.o_part = (float*)o_part.data_ptr();
|
||||
p.ml_part = (float*)ml_part.data_ptr();
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user