ViperEkura
6354dbe8bc
- our 128x128 fast loop padded 26 NOPs between the 32 QMMAs while all four LDGSTS sat bunched at the loop tail: ptxas had no independent instructions to fill the tensor-pipe issue gaps, the exact structure the decompiled cuBLAS loop (166 i, NOP=0) and CUTLASS MmaMultistage avoid by issuing cp.async in small groups inside the MMA phase (copy_tiles_and_advance per warp-tile batch) - the steady-state prefetch is now a loop-carried register pair per congruous operand (PrefetchCarry: swizzled stage offset + global source, constructed once from the same (r, c0) mapping as the interior loader), whose chunks ride after the first and last k_seg MMA batches — SASS: 282 -> 110 instructions, 0 branches, 0 UIMAD.WIDE magic-divisions, LDGSTS interleaved inside the QMMA range, 26 -> 19 NOPs, still 128 regs (2 CTAs/SM) - the wait-count dispatch ladder (16 instructions of ISETP/SEL picking DEPBAR immediates) and the per-k-tile (tile % ring) * stage_bytes recomputation (UIMAD.WIDE by 0x55555555) are gone: the prologue commits unconditionally so the wait_group<kStages-1> immediate is valid for every iteration, and both read and write stage addresses advance as carried pointers with an equality wrap - cp_async.cuh splits the emitter from its policies: one raw PTX site (cp_async_16_raw) plus wrappers for unconditional/predicated and pointer/offset destinations, and the now-unused wait_group_dispatch ladder is deleted; the dispatch flip: with the stall gone the big CTA wins the whole former dip band (1280^3 kernel-level 98.2->104.4T), so prefer_small_cta keeps only the sub-5/8-wave band and the single-wave s3 variant is retired Benchmark: L20 (sm_89), kernel-level sweep 128s2ff: 1024^3 104.2->114.0T (cuBLAS 151.2), 1152^3 130.4->145.9T (157.2), 1280^3 98.2->104.7T (163.8), 1536^3 141.4->150.4T (182.3), 2048^3 178.1->189.6T (202.8), 4096^3 197.6->208.8T (223.7), 8192^3 207.6->219.1T (227.5). CUDA-graph e2e: 1024^3 108.2->118.4T, 1152^3 137.7->150.7T, 1536^3 143.4->153.7T, 2048^3 179.4->191.9T, 8192^3 199.0->208.9T, 1280^3 108.3->106.6T (old dip-band rule re-measured 106.1T — within noise). Four-layout C++ suite and 596 pytests pass.
83 lines
3.6 KiB
Plaintext
83 lines
3.6 KiB
Plaintext
// Shared cp.async primitives — pure CUDA, no torch.
|
|
//
|
|
// One header for the async-copy pipeline used by both the attention kernels
|
|
// (predicated 16-byte K/V tile staging) and the fp8 GEMM (predicated operand
|
|
// staging + the fixed-depth wait_group). The emitter is split from its
|
|
// policies: cp_async_16_raw owns the single PTX site, and each wrapper states
|
|
// one destination contract (generic pointer vs loop-carried shared offset)
|
|
// and one predication contract (unconditional vs zero-fill-when-false), so
|
|
// call sites never pass a dead `true` predicate or re-convert a carried
|
|
// offset. PTX requires wait_group's operand to be an immediate, hence the
|
|
// template form below.
|
|
|
|
#pragma once
|
|
|
|
#include <cuda_runtime.h>
|
|
|
|
namespace astrai {
|
|
|
|
// Raw emitter: read src_size bytes (<= 16) from gmem into the shared
|
|
// offset. src_size = 0 reads nothing, so a predicated-off call zero-fills
|
|
// its destination without touching the (possibly out-of-range) source.
|
|
// BypassL1 selects .cg (L2 only, default) vs .ca (L1 + L2).
|
|
template <bool BypassL1 = true>
|
|
__device__ __forceinline__ void cp_async_16_raw(unsigned smem_addr,
|
|
const void* gmem_ptr,
|
|
int src_size) {
|
|
if constexpr (BypassL1) {
|
|
asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;"
|
|
:: "r"(smem_addr), "l"(gmem_ptr), "r"(src_size));
|
|
} else {
|
|
asm volatile("cp.async.ca.shared.global [%0], [%1], 16, %2;"
|
|
:: "r"(smem_addr), "l"(gmem_ptr), "r"(src_size));
|
|
}
|
|
}
|
|
|
|
// Unconditional 16-byte copy to a generic shared pointer.
|
|
// `T` is the smem element type; only the destination pointer's type matters.
|
|
template <typename T, bool BypassL1 = true>
|
|
__device__ __forceinline__ void cp_async_16(T* smem_ptr,
|
|
const void* gmem_ptr) {
|
|
cp_async_16_raw<BypassL1>(__cvta_generic_to_shared(smem_ptr), gmem_ptr,
|
|
16);
|
|
}
|
|
|
|
// Predicated: full copy when `pred`, zero-fill otherwise.
|
|
template <typename T, bool BypassL1 = true>
|
|
__device__ __forceinline__ void cp_async_16(T* smem_ptr, const void* gmem_ptr,
|
|
bool pred) {
|
|
cp_async_16_raw<BypassL1>(__cvta_generic_to_shared(smem_ptr), gmem_ptr,
|
|
pred ? 16 : 0);
|
|
}
|
|
|
|
// Predicated raw-offset form: the destination is an already-converted
|
|
// shared-memory offset (e.g. a loop-carried swizzled stage address), so
|
|
// steady-state prefetch sites issue one LDGSTS straight from the register.
|
|
template <bool BypassL1 = true>
|
|
__device__ __forceinline__ void cp_async_16(unsigned smem_addr,
|
|
const void* gmem_ptr, bool pred) {
|
|
cp_async_16_raw<BypassL1>(smem_addr, gmem_ptr, pred ? 16 : 0);
|
|
}
|
|
|
|
// Commit all outstanding cp.async ops of this thread as one group.
|
|
__device__ __forceinline__ void cp_async_commit_group() {
|
|
asm volatile("cp.async.commit_group;");
|
|
}
|
|
|
|
// Wait for every committed group (pipeline drain).
|
|
__device__ __forceinline__ void cp_async_wait_all() {
|
|
asm volatile("cp.async.wait_all;");
|
|
}
|
|
|
|
// Wait until at most KeepGroups committed groups are still in flight.
|
|
// PTX requires an immediate operand; keep it as a template argument so the
|
|
// stage policy stays compile-time configurable.
|
|
template <int KeepGroups>
|
|
__device__ __forceinline__ void cp_async_wait_group() {
|
|
static_assert(KeepGroups >= 0 && KeepGroups <= 7,
|
|
"cp.async.wait_group supports immediates in [0, 7]");
|
|
asm volatile("cp.async.wait_group %0;" :: "n"(KeepGroups));
|
|
}
|
|
|
|
} // namespace astrai
|