Niko1221/Strata
1
0 Fork 0
mirror of https://github.com/Niko1221/Strata.git synced 2026-11-07 14:21:06 -10:00
Code Issues Packages Projects Releases 1 Wiki Activity

split: build option STRATA_SPLIT_PF_JOB (64 default, 128): rows of one expert per prompt-path job, so a weight tile is dequantized once for twice the tokens

Browse Source 
Default unchanged (64). The expert prompt kernel's token-tile count follows PF_JOB (it was fixed at four).
This commit is contained in:
Niko1221 committed 2026-11-02 11:13:09 +01:00
1 parent 15fb6ffe3f
commit ea2aed12cd
3 files changed
+15 -5
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/split/CMakeLists.txt
+2 -1
View File
@@ -12,6 +12,7 @@
# the engine refuses a card of any other architecture
set(STRATA_SPLIT_CARDS "2" CACHE STRING "card counts strata-split is built for (2, 3, 4)")
set(STRATA_SPLIT_ARCH "" CACHE STRING "compute capability strata-split is built for (e.g. 86)")
set(STRATA_SPLIT_PF_JOB "64" CACHE STRING "rows of one expert per prompt-path job: 64 (default) or 128")
if(NOT STRATA_ENABLE_CUDA OR WIN32 OR APPLE)
message(STATUS "STRATA_BUILD_SPLIT: needs a Linux CUDA build, skipped")
@@ -50,7 +51,7 @@ foreach(_n ${STRATA_SPLIT_CARDS})
add_executable(${_t} ${_sp_engine} ${_sp_kparts})
target_include_directories(${_t} PRIVATE ${_sp}/src ${_sp}/tools ${STRATA_GGML_DIR}/ggml/src/ggml-cuda ${STRATA_GGML_DIR}/ggml/src
${STRATA_GGML_DIR}/ggml/include)
target_compile_definitions(${_t} PRIVATE FNX_NCARD=${_n} FNX_ARCH=${STRATA_SPLIT_ARCH} FNX_27B=0 FNX_HOSTRING=1
target_compile_definitions(${_t} PRIVATE FNX_NCARD=${_n} FNX_ARCH=${STRATA_SPLIT_ARCH} FNX_27B=0 FNX_HOSTRING=1 FNX_PF_JOB=${STRATA_SPLIT_PF_JOB}
FNX_DIR="${_sp}")
target_compile_options(${_t} PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:--use_fast_math --expt-relaxed-constexpr -Wno-deprecated-gpu-targets>)
set_target_properties(${_t} PROPERTIES CUDA_ARCHITECTURES "${STRATA_SPLIT_ARCH}-real" CUDA_STANDARD 17 CUDA_STANDARD_REQUIRED ON
src/split/src/kernels.hpp
+7 -2
View File
@@ -22,8 +22,13 @@ struct Job {
int out[MAXT]; // output rows
};
// the prompt path's expert jobs: up to PF_JOB rows of one expert
constexpr int PF_JOB = 64;
// the prompt path's expert jobs: up to PF_JOB rows of one expert (64: four 16-token tiles per job; 128: eight, the expert's weight tile is then
// dequantized once for twice the tokens; a build option, STRATA_SPLIT_PF_JOB)
#ifndef FNX_PF_JOB
#define FNX_PF_JOB 64
#endif
constexpr int PF_JOB = FNX_PF_JOB;
static_assert(PF_JOB == 64 || PF_JOB == 128, "PF_JOB: 64 or 128");
struct BigJob {
int w;
int n;
src/split/src/moe_mma.cuh
+6 -2
View File
@@ -280,8 +280,12 @@ template <int TYPE, int BM>
__global__ void __launch_bounds__(MM_THREADS) k_moe_mma(MoeW w, const BigJob* __restrict__ jobs, const int* __restrict__ njobs,
const __half* __restrict__ x, int ldx, float* __restrict__ out) {
using namespace nvcuda;
constexpr int WR = BM / 32, WT = MM_THREADS / 32 / WR, TT = 4 / WT, NP = BM / MM_PASS, LO = 16 * TT;
static_assert(WR * WT == MM_THREADS / 32 && TT * WT == 4 && NP * MM_PASS == BM, "warp layout");
constexpr int MM_NT = PF_JOB / 16; // token tiles of a job
constexpr int WR = BM / 32, WT = MM_THREADS / 32 / WR, TT = MM_NT / WT, NP = BM / MM_PASS, LO = 16 * TT;
static_assert(WR * WT == MM_THREADS / 32 && TT * WT == MM_NT && NP * MM_PASS == BM, "warp layout");
#if FNX_ARCH < 80
static_assert(PF_JOB == 64, "the m8n8k4 path is laid out for 64-row jobs");
#endif
const int ji = blockIdx.y;
if (ji >= *njobs) return;
__shared__ BigJob jb;