Skip to content

prefill: the fused int8 prompt experts on Turing (sm_75), opt-in (+9 % prefill on RTX 2080 Ti) #2145

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
rafatxf wants to merge 1 commit into from
+157 −29
Closed

prefill: the fused int8 prompt experts on Turing (sm_75), opt-in (+9 % prefill on RTX 2080 Ti) #2145

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 27 additions & 5 deletions src/prefill/moe_fused.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -194,6 +194,26 @@ __device__ __forceinline__ void mma_s8(int (&d)[4], const uint32_t (&a)[4], uint
: "=r"(d[0]), "=r"(d[1]), "=r"(d[2]), "=r"(d[3])
: "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1), "r"(0));
}
#elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 750
// Turing (STRATA_PF_FUSED_SM75): no cp.async - the stages are filled by plain copies (the same bytes; the barriers
// around them stay) - and no m16n8k32 with s8 operands: it is four m8n8k16, rows g / g + 8 by k 0-15 / 16-31 (the
// m16n8k32 fragments: a[0] / a[1] rows g / g + 8 at k 4 tig.., a[2] / a[3] the same at k 16 + 4 tig.., b0 / b1 the
// two k halves). The int32 sums are exact in any order: the same d as the sm_80 instruction.
__device__ __forceinline__ void cp16(void* dst, const void* src) { *(uint4*) dst = __ldg((const uint4*) src); }
__device__ __forceinline__ void cp4(void* dst, const void* src) { *(uint32_t*) dst = __ldg((const uint32_t*) src); }
__device__ __forceinline__ void cp_commit() {}
template <int N> __device__ __forceinline__ void cp_wait() {}
__device__ __forceinline__ void mma8_s8(int& c0, int& c1, uint32_t a, uint32_t b) {
asm("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 {%0,%1}, {%2}, {%3}, {%0,%1};\n"
: "+r"(c0), "+r"(c1) : "r"(a), "r"(b));
}
__device__ __forceinline__ void mma_s8(int (&d)[4], const uint32_t (&a)[4], uint32_t b0, uint32_t b1) {
d[0] = d[1] = d[2] = d[3] = 0;
mma8_s8(d[0], d[1], a[0], b0);
mma8_s8(d[0], d[1], a[2], b1);
mma8_s8(d[2], d[3], a[1], b0);
mma8_s8(d[2], d[3], a[3], b1);
}
#endif

// One work item = (tile of up to 64 routed rows of expert e, a block of 256 weight rows); a persistent grid walks
Expand All @@ -206,7 +226,7 @@ template <bool GU>
__global__ void __launch_bounds__(THREADS, 1)
expert_kernel(const Batch b, const Tables tb, const uint8_t* __restrict__ act, const int32_t* __restrict__ src,
uint8_t* __restrict__ out, float* __restrict__ dm) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 750
constexpr int NKB = GU ? 40 : 10; // 64-weight blocks along K
constexpr int NFB = GU ? 5 : 10; // weight-row blocks per tile
constexpr int CODE_LD = GU ? 640 : 160; // code bytes per weight row
Expand Down Expand Up @@ -626,10 +646,12 @@ const DevInfo& dev_info() {
cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, dev);
cudaDeviceGetAttribute(&d.sms, cudaDevAttrMultiProcessorCount, dev);
d.cc = 10 * major + minor;
if (d.cc < 80) return d;
// the device code of this card's arch must have the kernels (a build without an sm_80+ target has an empty body)
// Turing (sm_75): the same kernels with plain copies and m8n8k16 MMAs, opt-in (STRATA_PF_FUSED_SM75=1)
static const bool sm75 = [] { const char* v = std::getenv("STRATA_PF_FUSED_SM75"); return v != nullptr && v[0] == '1'; }();
if (d.cc < 75 || (d.cc < 80 && !sm75)) { d.cc = d.cc < 80 ? 0 : d.cc; return d; }
// the device code of this card's arch must have the kernels (a build without an sm_75+ target has an empty body)
cudaFuncAttributes fa{};
if (cudaFuncGetAttributes(&fa, expert_kernel<true>) != cudaSuccess || fa.ptxVersion < 80) {
if (cudaFuncGetAttributes(&fa, expert_kernel<true>) != cudaSuccess || fa.ptxVersion < (d.cc >= 80 ? 80 : 75)) {
cudaGetLastError();
d.cc = 0;
return d;
Expand All @@ -649,7 +671,7 @@ const DevInfo& dev_info() {
} // namespace

bool built() { return true; }
bool available() { return dev_info().cc >= 80; }
bool available() { return dev_info().cc >= 75; }
// 0.1.36: on by default for the Q2_0 pack (+13-23% prompts on an RTX 5070; teacher-forced against the FP16 path as
// close as MMQ at 8K and closer at 32K: phaseA-tests s18-tf); STRATA_PF_FUSED=0 keeps MMQ. The native packs' kernels
// (moe_fused_iq) stay opt-in: STRATA_PF_FUSED=1 (requested()).
Expand Down
Loading