diff --git a/python/sglang/kernels/jit/csrc/attention/fixup_zero_kv.cuh b/python/sglang/kernels/jit/csrc/attention/fixup_zero_kv.cuh index 5d1633e32..3d6026362 100644 --- a/python/sglang/kernels/jit/csrc/attention/fixup_zero_kv.cuh +++ b/python/sglang/kernels/jit/csrc/attention/fixup_zero_kv.cuh @@ -17,37 +17,52 @@ constexpr int kFixupBlockSize = 256; // -- vectorised zero-fill helpers ------------------------------------------ -// Zero-fill `n` elements of type T starting at `ptr`, using float4 stores. -// `ptr` must be 16-byte aligned (guaranteed by PyTorch allocator). +// Zero-fill `n` elements of type T starting at `ptr`. template __device__ __forceinline__ void vec_zero_fill(T* ptr, int n) { - constexpr int kVec = 16 / sizeof(T); // elements per float4 - const int n_vec = n / kVec; // full vectors - float4* dst4 = reinterpret_cast(ptr); - const float4 z4 = make_float4(0.f, 0.f, 0.f, 0.f); - for (int i = threadIdx.x; i < n_vec; i += blockDim.x) { - dst4[i] = z4; - } - // tail elements - const int tail_start = n_vec * kVec; - for (int i = tail_start + threadIdx.x; i < n; i += blockDim.x) { - ptr[i] = static_cast(0); + if ((reinterpret_cast(ptr) & 0xF) == 0) { + // 16-byte aligned -> vectorised float4 stores + constexpr int kVec = 16 / sizeof(T); // elements per float4 + const int n_vec = n / kVec; // full vectors + float4* dst4 = reinterpret_cast(ptr); + const float4 z4 = make_float4(0.f, 0.f, 0.f, 0.f); + for (int i = threadIdx.x; i < n_vec; i += blockDim.x) { + dst4[i] = z4; + } + // tail elements + const int tail_start = n_vec * kVec; + for (int i = tail_start + threadIdx.x; i < n; i += blockDim.x) { + ptr[i] = static_cast(0); + } + } else { + // misaligned row base -> scalar stores + for (int i = threadIdx.x; i < n; i += blockDim.x) { + ptr[i] = static_cast(0); + } } } -// Fill `n` float elements with -inf using float4 stores. +// Fill `n` float elements with -inf. __device__ __forceinline__ void vec_neginf_fill(float* ptr, int n) { - constexpr int kVec = 4; // float4 = 4 floats - const int n_vec = n / kVec; - float4* dst4 = reinterpret_cast(ptr); const float ninf = -INFINITY; - const float4 inf4 = make_float4(ninf, ninf, ninf, ninf); - for (int i = threadIdx.x; i < n_vec; i += blockDim.x) { - dst4[i] = inf4; - } - const int tail_start = n_vec * kVec; - for (int i = tail_start + threadIdx.x; i < n; i += blockDim.x) { - ptr[i] = ninf; + if ((reinterpret_cast(ptr) & 0xF) == 0) { + // 16-byte aligned -> vectorised float4 stores + constexpr int kVec = 4; // float4 = 4 floats + const int n_vec = n / kVec; + float4* dst4 = reinterpret_cast(ptr); + const float4 inf4 = make_float4(ninf, ninf, ninf, ninf); + for (int i = threadIdx.x; i < n_vec; i += blockDim.x) { + dst4[i] = inf4; + } + const int tail_start = n_vec * kVec; + for (int i = tail_start + threadIdx.x; i < n; i += blockDim.x) { + ptr[i] = ninf; + } + } else { + // misaligned row base -> scalar stores + for (int i = threadIdx.x; i < n; i += blockDim.x) { + ptr[i] = ninf; + } } } diff --git a/python/sglang/kernels/jit/csrc/attention/kda_fused_decode.cuh b/python/sglang/kernels/jit/csrc/attention/kda_fused_decode.cuh new file mode 100644 index 000000000..1b1944f50 --- /dev/null +++ b/python/sglang/kernels/jit/csrc/attention/kda_fused_decode.cuh @@ -0,0 +1,1064 @@ +// KDA fused decode step for Kimi K3: causal conv1d update + delta-rule +// recurrence + gated RMSNorm in a single kernel (replaces the three-kernel +// causal_conv1d_update -> kda_packed_decode -> rms_norm_gated decode chain). +// +// Kernel body vendored from the NVIDIA x Moonshot Kimi K3 optimization +// package (KDA_decode/kda_decode_fusion_kernel.cu, many-heads variant). +// +// Local integration changes vs. the NV source (assembled by script, each +// patch anchored on exact source text): +// * explicit row strides for x/g/beta/onorm_g so the fused qkvg-projection +// GEMM output slices feed the kernel without any .contiguous() copies; +// * conv state addressed through cs_slot_stride/cs_w_stride so the kernel +// updates the packed [slots, width, conv_dim] mamba pool in place +// (the pool is natively transposed on this branch); +// * ssm/temporal state addressed through state_slot_stride (state.stride(0)) +// so the kernel reads/writes envelope-strided [slots, HV, V, K] pools in +// place: the unified / page-major layouts pitch one slot across ALL layers +// (56,171,520 B on K3), NOT the dense HV*V*K pitch. int64 slot*stride math +// avoids the envelope-pitch overflow (the exact chunk_delta_h bug pattern); +// * the static-decode-layout path honors ssm_state_indices (the shipped +// config hardwired slot = blockIdx.x, valid only for dense benches); +// * padded cuda-graph slots (index < 0) zero the output row and skip all +// pool updates, mirroring kda_packed_decode; +// * tvm-ffi host wrapper (KdaFusedDecodeKernel) with TensorMatcher shape / +// stride / dtype validation replaces the pybind binding; +// * automatic 1D-TMA bulk state load for aligned recurrent-state layouts, +// ported from the standalone KDA_decode/kda_decode_fusion_kernel.cu +// experiment on chunan/kda (same mbarrier staging, same 3/4-stage +// dispatch by grid size); +// * PDL (kUsePDL, plumbed like kda_packed_decode): griddepcontrol wait at +// kernel entry, launch_dependents at every exit. + +#include // For TensorMatcher, SymbolicSize, SymbolicDevice +#include // For RuntimeCheck + +#include +#include // For bf16_t, fp32_t +#include // For LaunchKernel +#include // For device::warp::reduce_sum + +#include + +#include +#include +#include +#include + +// Local PTX primitives (cp.async / mbarrier / async-proxy fence) + +namespace ptx { + +// Generic ptr -> 32-bit `.shared` address: inline-PTX `.shared` instructions +// take a byte offset in the shared window, not a generic pointer. +template +static SGL_DEVICE uint32_t to_shared(T* ptr) { + return static_cast(__cvta_generic_to_shared(ptr)); +} + +// ---- non-bulk cp.async (PTX ISA §9.7.9.24) --------------------------------- + +// One 16-byte cache-global segment, global -> shared. Both pointers must be +// 16-byte aligned. +static SGL_DEVICE void cp_async_cg_16b(void* smem_dst, const void* gmem_src) { + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;" ::"r"(to_shared(smem_dst)), "l"(gmem_src) : "memory"); +} + +static SGL_DEVICE void cp_async_commit_group() { + asm volatile("cp.async.commit_group;"); +} + +// Wait until at most N committed cp.async groups remain pending. +template +static SGL_DEVICE void cp_async_wait_group() { + static_assert(N >= 0 && N <= 7, "cp.async wait-group count must be in [0, 7]"); + asm volatile("cp.async.wait_group %0;" ::"n"(N) : "memory"); +} + +static SGL_DEVICE void cp_async_wait_all() { + asm volatile("cp.async.wait_all;" ::: "memory"); +} + +// ---- bulk 1D TMA (PTX ISA §9.7.9.25) --------------------------------------- + +// global -> shared::cluster, completed by an smem mbarrier. Arm `bar` with +// `mbar_arrive_expect_tx(bar, bytes)` before issuing; `bytes` and both +// endpoints must be 16-byte aligned. +static SGL_DEVICE void cp_async_bulk_1d_load(void* smem_dst, const void* gmem_src, uint32_t bytes, uint64_t* bar) { + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes" + " [%0], [%1], %2, [%3];" ::"r"(to_shared(smem_dst)), + "l"(gmem_src), + "r"(bytes), + "r"(to_shared(bar)) + : "memory"); +} + +// ---- mbarrier (PTX ISA §9.7.13.15) ----------------------------------------- +// +// Only the `try_wait.parity` waiter is wrapped; the caller owns the phase +// counter and flips it at the stage wrap. After `mbar_init` the bar is at +// parity 0 and each full cycle (count arrivals -> fire -> reset) flips it, so a +// consumer-first waiter starts at 0 and a producer-first waiter (whose first +// wait must be a no-op skip) starts at 1. Getting this backwards deadlocks on +// the second wait. +static SGL_DEVICE void mbar_init(uint64_t* bar, uint32_t count) { + asm volatile("mbarrier.init.shared.b64 [%0], %1;" ::"r"(to_shared(bar)), "r"(count)); +} + +// Combined arrive + set tx-count, for TMA-load completion. +static SGL_DEVICE void mbar_arrive_expect_tx(uint64_t* bar, uint32_t bytes) { + asm volatile("mbarrier.arrive.expect_tx.shared.b64 _, [%0], %1;" ::"r"(to_shared(bar)), "r"(bytes)); +} + +// Wait for phase `parity` to complete. Looped because the spec allows spurious +// early wakeups. The default `.acquire` semantics make prior `cp.async.bulk` +// writes tracked by this mbarrier visible to later generic-proxy reads on this +// thread with no `fence.proxy.async` (spec §9.7.13.15.16 point 3). +static SGL_DEVICE void mbar_wait_parity(uint64_t* bar, uint32_t parity) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "WAIT_%=: mbarrier.try_wait.parity.shared.b64 p, [%0], %1;\n\t" + "@!p bra WAIT_%=;\n\t}\n" ::"r"(to_shared(bar)), + "r"(parity)); +} + +// ---- async-proxy fence (PTX ISA §9.7.13) ----------------------------------- + +// Make generic-proxy smem writes visible to the async proxy. Required before +// any bulk store that reads smem written by regular ld/st. +static SGL_DEVICE void fence_async_smem() { + asm volatile("fence.proxy.async.shared::cta;"); +} + +} // namespace ptx + +namespace { + +constexpr int kDimK = 128; +constexpr int kDimV = 128; +constexpr int kKernelWidth = 4; +constexpr int kConvStateWidth = kKernelWidth - 1; +constexpr int kThreads = 256; +constexpr int kWarps = kThreads / 32; +constexpr int kChunkV = 32; +constexpr int kNumChunks = kDimV / kChunkV; +constexpr int kRowsPerWarp = kChunkV / kWarps; + +SGL_DEVICE float bf16_load(const __nv_bfloat16* ptr, int idx) { + return __bfloat162float(ptr[idx]); +} + +SGL_DEVICE __nv_bfloat16 bf16_store(float value) { + return __float2bfloat16(value); +} + +template +SGL_DEVICE void store_state_float4(float* ptr, float4 value) { + if constexpr (kUseCacheGlobalStore) { + __stcg(reinterpret_cast(ptr), value); + } else { + *reinterpret_cast(ptr) = value; + } +} + +template +SGL_DEVICE void +cp_async_state_chunk_for(float* s_state, const float* state, int slot, int i_hv, int64_t state_slot_stride, int chunk) { + constexpr int kFloat4PerChunk = kChunkV * kDimK / 4; + const int tid = threadIdx.x; + const int stage = chunk & 1; + const int v_base = chunk * kChunkV; + const int64_t slot_base = static_cast(slot) * state_slot_stride; + for (int linear4 = tid; linear4 < kFloat4PerChunk; linear4 += kCopyThreads) { + const int elem = linear4 * 4; + const int row = elem / kDimK; + const int k = elem - row * kDimK; + float* dst = s_state + (stage * kChunkV + row) * kDimK + k; + const float* src = state + slot_base + ((i_hv * kDimV + v_base + row) * kDimK + k); + ptx::cp_async_cg_16b(dst, src); + } + ptx::cp_async_commit_group(); +} + +SGL_DEVICE void +cp_async_state_chunk(float* s_state, const float* state, int slot, int i_hv, int64_t state_slot_stride, int chunk) { + cp_async_state_chunk_for(s_state, state, slot, i_hv, state_slot_stride, chunk); +} + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 +#define KDA_FUSED_DECODE_HAS_TMA 1 +#else +#define KDA_FUSED_DECODE_HAS_TMA 0 +#endif + +// 1D TMA bulk copy gmem -> smem; the issuing thread arrives with expect-tx on +// the mbarrier, completion is observed via ptx::mbar_wait_parity. +template +SGL_DEVICE void tma_state_chunk_stage( + float* s_state, + const float* state, + int slot, + int i_hv, + int64_t state_slot_stride, + int chunk, + int stage, + uint64_t* bar) { + constexpr uint32_t kBytes = kStageChunkV * kDimK * sizeof(float); + float* dst = s_state + stage * kStageChunkV * kDimK; + // 16B-aligned for any slot iff state_slot_stride*sizeof(float) % 16 == 0 + // (host gates TMA off otherwise); the intra-slot chunk offset is a multiple + // of kStageChunkV*kDimK*4, always 16B-aligned. + const int64_t slot_base = static_cast(slot) * state_slot_stride; + const float* src = state + slot_base + ((i_hv * kDimV + chunk * kStageChunkV) * kDimK); +#if KDA_FUSED_DECODE_HAS_TMA + ptx::fence_async_smem(); + ptx::mbar_arrive_expect_tx(bar, kBytes); + ptx::cp_async_bulk_1d_load(dst, src, kBytes, bar); +#else + __trap(); +#endif +} + +SGL_DEVICE float block_reduce_sum(float value, float* scratch) { + const int tid = threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + + float warp_total = device::warp::reduce_sum(value); + if (lane == 0) { + scratch[warp] = warp_total; + } + __syncthreads(); + + float block_total = 0.0f; + if (warp == 0) { + block_total = lane < kWarps ? scratch[lane] : 0.0f; + block_total = device::warp::reduce_sum(block_total); + if (lane == 0) { + scratch[0] = block_total; + } + } + __syncthreads(); + return scratch[0]; +} + +struct Sum2 { + float x; + float y; +}; + +SGL_DEVICE Sum2 warp_reduce_sum_pair(float x, float y) { + return {device::warp::reduce_sum(x), device::warp::reduce_sum(y)}; +} + +template +SGL_DEVICE Sum2 block_reduce_sum2_for(float x, float y, float* scratch) { + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + + const float warp_x = device::warp::reduce_sum(x); + const float warp_y = device::warp::reduce_sum(y); + if (lane == 0) { + scratch[warp] = warp_x; + scratch[kReduceWarps + warp] = warp_y; + } + __syncthreads(); + + float block_x = 0.0f; + float block_y = 0.0f; + if (warp == 0) { + block_x = lane < kReduceWarps ? scratch[lane] : 0.0f; + block_y = lane < kReduceWarps ? scratch[kReduceWarps + lane] : 0.0f; + block_x = device::warp::reduce_sum(block_x); + block_y = device::warp::reduce_sum(block_y); + if (lane == 0) { + scratch[0] = block_x; + scratch[1] = block_y; + } + } + __syncthreads(); + return {scratch[0], scratch[1]}; +} + +SGL_DEVICE Sum2 block_reduce_sum2(float x, float y, float* scratch) { + return block_reduce_sum2_for(x, y, scratch); +} + +template +SGL_DEVICE float block_reduce_sum_active_for(float value, float* scratch) { + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + + float warp_total = 0.0f; + if (warp < kReduceWarps) { + warp_total = device::warp::reduce_sum(value); + } + if (lane == 0 && warp < kReduceWarps) { + scratch[warp] = warp_total; + } + __syncthreads(); + + float block_total = 0.0f; + if (warp == 0) { + block_total = lane < kReduceWarps ? scratch[lane] : 0.0f; + block_total = device::warp::reduce_sum(block_total); + if (lane == 0) { + scratch[0] = block_total; + } + } + __syncthreads(); + return scratch[0]; +} + +template +SGL_DEVICE Sum2 block_reduce_sum2_active_for(float x, float y, float* scratch) { + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + + float warp_x = 0.0f; + float warp_y = 0.0f; + if (warp < kReduceWarps) { + warp_x = device::warp::reduce_sum(x); + warp_y = device::warp::reduce_sum(y); + } + if (lane == 0 && warp < kReduceWarps) { + scratch[warp] = warp_x; + scratch[kReduceWarps + warp] = warp_y; + } + __syncthreads(); + + float block_x = 0.0f; + float block_y = 0.0f; + if (warp == 0) { + block_x = lane < kReduceWarps ? scratch[lane] : 0.0f; + block_y = lane < kReduceWarps ? scratch[kReduceWarps + lane] : 0.0f; + block_x = device::warp::reduce_sum(block_x); + block_y = device::warp::reduce_sum(block_y); + if (lane == 0) { + scratch[0] = block_x; + scratch[1] = block_y; + } + } + __syncthreads(); + return {scratch[0], scratch[1]}; +} + +SGL_DEVICE float fp32_at(const float* ptr, int64_t idx) { + return ptr[idx]; +} + +// One-token KDA (Kimi Delta Attention) decode step, fused end to end: +// 1. Causal conv1d update: q/k/v = SiLU(bias + w_{q,k,v}_t (dot) [conv_state, x_{q,k,v}]), +// shift-registers cs_q/cs_k/cs_v advanced in place (x_q/x_k/x_v are the raw +// per-token projections, w_*_t the depthwise conv taps). +// 2. Per-head decay gate from a_log/g/dt_bias: +// decay = exp(lower_bound * sigmoid(exp(a_log) * (g + dt_bias))) if kUseLowerBound +// = exp(-exp(a_log) * softplus(g + dt_bias)) otherwise +// 3. q, k are L2-normalized over kDimK (q additionally scaled by `scale`); beta +// is sigmoid(raw) when kApplyBetaSigmoid, else used as-is. +// 4. Delta-rule recurrent state update per value row v of the [kDimV, kDimK] state h: +// h_decay = h * decay +// v_new = (v - h_decay . k) * beta +// h' = h_decay + k (outer) v_new (written back to `state`) +// o = h' . q +// 5. Optional gated RMSNorm (onorm) over o: out = o * rsqrt(mean(o^2) + onorm_eps) +// * onorm_weight * sigmoid(onorm_g). +// Grid maps one block per (batch/token, value-head); kUseStaticDecodeLayout picks a +// fixed (B, HV) launch shape for CUDA-graph capture, cu_seqlens/ssm_state_indices +// otherwise resolve the token's batch slot and recurrent-state slot. +template < + bool kApplyOnorm, + bool kUseStaticDecodeLayout = false, + int kFixedHeads = 0, + int kFixedValueHeads = 0, + bool kUseHeadGrid = false, + bool kAccumulateOnormSumsq = false, + bool kUseActiveQkReduction = false, + bool kUseCacheGlobalStore = false, + bool kComputeOutputBeforeStore = false, + bool kSkipWarpSync = false, + bool kPreloadOnormParams = false, + bool kPrefetchNextStateChunk = false, + bool kUseActiveOnormReduction = false, + bool kUpdateConvState = false, + bool kUseLowerBound = false, + bool kApplyBetaSigmoid = true, + bool kUseTmaLoad = false, + int kTmaStages = kNumChunks, + bool kUsePDL = false> +// Shapes below use K3's per-TP-rank sizing (linear_attn_config: num_heads=96 over +// TP={8,16,32} -> H=HV={12,6,3} local heads; head_dim=128 -> kDimK=kDimV=128; kSeg = H*128; +// short_conv_kernel_size=4 -> kKernelWidth=4, kConvStateWidth=3). B is the live +// (post-padding, under kUseStaticDecodeLayout) decode batch size; slots is the +// recurrent-state / conv-cache pool capacity, addressed by ssm_state_indices, not B. +__global__ __launch_bounds__(kThreads, 2) void kda_decode_fusion_many_heads_kernel( + const __nv_bfloat16* __restrict__ x_q, // [B, H*128] row bos*x_row_stride + hk, sliced from mixed_qkv q-segment + const __nv_bfloat16* __restrict__ x_k, // [B, H*128] row bos*x_row_stride + hk, sliced from mixed_qkv k-segment + const __nv_bfloat16* __restrict__ x_v, // [B, HV*128] row bos*x_row_stride + hvv, sliced from mixed_qkv v-segment + const float* __restrict__ w_q_t, // [kKernelWidth=4, H*128] dense conv taps for q, indexed w*hkv_dim + hk + const float* __restrict__ w_k_t, // [4, H*128] dense conv taps for k + const float* __restrict__ w_v_t, // [4, HV*128] dense conv taps for v + const float* __restrict__ bias_q, // [H*128] conv bias for q, sliced from conv_bias + const float* __restrict__ bias_k, // [H*128] conv bias for k + const float* __restrict__ bias_v, // [HV*128] conv bias for v + __nv_bfloat16* __restrict__ cs_q, // [slots, kConvStateWidth=3, H*128] q shift-register, sliced from conv_states + __nv_bfloat16* __restrict__ cs_k, // [slots, 3, H*128] k shift-register + __nv_bfloat16* __restrict__ cs_v, // [slots, 3, HV*128] v shift-register + const float* __restrict__ a_log, // [H] per-head log-decay base, indexed by i_h + const __nv_bfloat16* __restrict__ g, // [B, HV*128] raw forget gate, row bos*g_row_stride + i_hv*kDimK + k + const float* __restrict__ dt_bias, // [H*128] gate bias added to g before the decay nonlinearity + const __nv_bfloat16* __restrict__ beta, // [B, HV] raw beta logit, row bos*beta_row_stride + i_hv + const __nv_bfloat16* __restrict__ onorm_g, // [B, HV*128] onorm sigmoid gate, row i_n*onormg_row_stride + i_hv*128 + // + v + const float* __restrict__ onorm_weight, // [128] onorm RMSNorm scale, shared across all heads + const int* __restrict__ ssm_state_indices, // [B] recurrent-state slot per token; <0 marks a padded cuda-graph slot + const int* __restrict__ cu_seqlens, // [B+1] token offsets into x_q/x_k/x_v/g/beta; unused under + // kUseStaticDecodeLayout + float* __restrict__ state, // [slots, HV, 128, 128] recurrent KDA state h, inner [V,K] contiguous, slot pitch = + // state_slot_stride + __nv_bfloat16* __restrict__ out, // [B, hv_count*128] fused-decode output, row i_n, col i_hv*128 + v + int B, // live decode batch size (token count for this launch) + int H, // local key/query heads on this TP rank + int HV, // local value heads on this TP rank (H==HV since KDA is MHA not GQA) + float lower_bound, // linear_attn_config.gate_lower_bound (-5.0 for K3) when kUseLowerBound + float scale, // query scale applied after L2-normalization + float onorm_eps, // onorm RMSNorm epsilon + int64_t x_row_stride, // element stride between consecutive tokens' rows in x_q/x_k/x_v (>= 3*H*128) + int64_t g_row_stride, // element stride between consecutive tokens' rows in g (>= HV*128) + int64_t beta_row_stride, // element stride between consecutive tokens' rows in beta (>= HV) + int64_t onormg_row_stride, // element stride between consecutive tokens' rows in onorm_g (>= HV*128) + int64_t cs_slot_stride, // element stride between consecutive slots in cs_q/cs_k/cs_v + int64_t cs_w_stride, // element stride between shift-register taps (w=0..2) within a slot + int64_t state_slot_stride) { // element stride between consecutive slots in state (dense HV*128*128, or larger for + // shared pools) + device::PDLWaitPrimary(); + const int tid = threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + int i_n; + int i_hv; + int i_h; + int bos; + int slot; + if constexpr (kUseStaticDecodeLayout) { + if constexpr (kUseHeadGrid) { + i_n = blockIdx.x; + i_hv = blockIdx.y; + } else { + const int nhv = blockIdx.x; + i_n = nhv / kFixedValueHeads; + i_hv = nhv - i_n * kFixedValueHeads; + } + i_h = i_hv; + bos = i_n; + slot = ssm_state_indices == nullptr ? i_n : ssm_state_indices[i_n]; + } else { + const int nhv = blockIdx.x; + i_n = nhv / HV; + i_hv = nhv - i_n * HV; + const int hv_per_h = HV / H; + i_h = i_hv / hv_per_h; + + bos = cu_seqlens == nullptr ? i_n : cu_seqlens[i_n]; + const int eos = cu_seqlens == nullptr ? i_n + 1 : cu_seqlens[i_n + 1]; + if (eos <= bos) { + device::PDLTriggerSecondary(); + return; + } + slot = ssm_state_indices == nullptr ? i_n : ssm_state_indices[i_n]; + } + + if (slot < 0) { + // Padded cuda-graph slot: zero the output row, leave the pools untouched. + const int hv_count_pad = kUseStaticDecodeLayout ? kFixedValueHeads : HV; + if (tid < kDimV) { + out[(i_n * hv_count_pad + i_hv) * kDimV + tid] = __float2bfloat16(0.0f); + } + device::PDLTriggerSecondary(); + return; + } + + const int hk_off = i_h * kDimK; + const int hv_off = i_hv * kDimV; + const int h_count = kUseStaticDecodeLayout ? kFixedHeads : H; + const int hv_count = kUseStaticDecodeLayout ? kFixedValueHeads : HV; + const int hkv_dim = h_count * kDimK; + const int hvv_dim = hv_count * kDimV; + + // Dynamic size: 2 cp.async stages (32KB) or kTmaStages TMA stages (16KB per + // stage; kTmaStages == kNumChunks means no buffer reuse). + extern __shared__ __align__(16) float s_state[]; + __shared__ float s_q[kDimK]; + __shared__ float s_k[kDimK]; + __shared__ float s_decay[kDimK]; + __shared__ float s_v[kDimV]; + __shared__ float s_o[kDimV]; + __shared__ float s_reduce[kThreads]; + __shared__ float s_beta; + __shared__ uint64_t s_tma_bar[kNumChunks]; + float pre_onorm_gate = 0.0f; + float pre_onorm_weight = 0.0f; + + if constexpr (kUseTmaLoad) { + // Chunk c lives in stage c % kTmaStages behind barrier c % kTmaStages + // (wait parity (c / kTmaStages) & 1). Chunks 0/1 are issued here; chunk+2 + // is issued at the loop top while its stage is still fresh, and chunks + // that reuse a stage are issued at the loop bottom behind a + // __syncthreads(). Barrier-init visibility for waiting threads is + // covered by the __syncthreads() before the chunk loop. + if (tid == 0) { +#pragma unroll + for (int c = 0; c < kTmaStages; ++c) { + ptx::mbar_init(&s_tma_bar[c], 1); + } + tma_state_chunk_stage(s_state, state, slot, i_hv, state_slot_stride, 0, 0, &s_tma_bar[0]); + if (kTmaStages > 1 && kNumChunks > 1) { + tma_state_chunk_stage(s_state, state, slot, i_hv, state_slot_stride, 1, 1, &s_tma_bar[1]); + } + } + } else { + cp_async_state_chunk(s_state, state, slot, i_hv, state_slot_stride, 0); + } + + if constexpr (kUpdateConvState) { + if (tid < kDimK) { + const int k = tid; + const int hk = hk_off + k; + const int64_t cs_base = slot * cs_slot_stride + hk; + const int64_t xq_idx = bos * x_row_stride + hk; + const float exp_a = __shfl_sync(0xffffffffu, lane == 0 ? __expf(a_log[i_h]) : 0.0f, 0); + + float q_acc = bias_q[hk]; + float k_acc = bias_k[hk]; + __nv_bfloat16 q_shift0 = __float2bfloat16(0.0f); + __nv_bfloat16 q_shift1 = __float2bfloat16(0.0f); + __nv_bfloat16 k_shift0 = __float2bfloat16(0.0f); + __nv_bfloat16 k_shift1 = __float2bfloat16(0.0f); +#pragma unroll + for (int w = 0; w < kConvStateWidth; ++w) { + const __nv_bfloat16 q_state = cs_q[cs_base + w * cs_w_stride]; + const __nv_bfloat16 k_state = cs_k[cs_base + w * cs_w_stride]; + q_acc += __bfloat162float(q_state) * fp32_at(w_q_t, w * hkv_dim + hk); + k_acc += __bfloat162float(k_state) * fp32_at(w_k_t, w * hkv_dim + hk); + if (w == 1) { + q_shift0 = q_state; + k_shift0 = k_state; + } else if (w == 2) { + q_shift1 = q_state; + k_shift1 = k_state; + } + } + const __nv_bfloat16 q_new = x_q[xq_idx]; + const __nv_bfloat16 k_new = x_k[xq_idx]; + q_acc += __bfloat162float(q_new) * fp32_at(w_q_t, (kKernelWidth - 1) * hkv_dim + hk); + k_acc += __bfloat162float(k_new) * fp32_at(w_k_t, (kKernelWidth - 1) * hkv_dim + hk); + + cs_q[cs_base + 0] = q_shift0; + cs_q[cs_base + cs_w_stride] = q_shift1; + cs_q[cs_base + 2 * cs_w_stride] = q_new; + cs_k[cs_base + 0] = k_shift0; + cs_k[cs_base + cs_w_stride] = k_shift1; + cs_k[cs_base + 2 * cs_w_stride] = k_new; + + s_q[k] = device::math::silu_fast(q_acc); + s_k[k] = device::math::silu_fast(k_acc); + + const float g_raw = bf16_load(g, bos * g_row_stride + i_hv * kDimK + k) + dt_bias[hk]; + if constexpr (kUseLowerBound) { + s_decay[k] = __expf(lower_bound * device::math::sigmoid_fast(exp_a * g_raw)); + } else { + s_decay[k] = __expf(-exp_a * device::math::softplus_fast(g_raw)); + } + } + } else { + if (tid < kDimK) { + const int k = tid; + const int hk = hk_off + k; + const float exp_a = __shfl_sync(0xffffffffu, lane == 0 ? __expf(a_log[i_h]) : 0.0f, 0); + + float q_acc = bias_q[hk]; + float k_acc = bias_k[hk]; +#pragma unroll + for (int w = 0; w < kConvStateWidth; ++w) { + const int64_t cs_idx = slot * cs_slot_stride + hk + w * cs_w_stride; + q_acc += bf16_load(cs_q, cs_idx) * fp32_at(w_q_t, w * hkv_dim + hk); + k_acc += bf16_load(cs_k, cs_idx) * fp32_at(w_k_t, w * hkv_dim + hk); + } + q_acc += bf16_load(x_q, bos * x_row_stride + hk) * fp32_at(w_q_t, (kKernelWidth - 1) * hkv_dim + hk); + k_acc += bf16_load(x_k, bos * x_row_stride + hk) * fp32_at(w_k_t, (kKernelWidth - 1) * hkv_dim + hk); + + s_q[k] = device::math::silu_fast(q_acc); + s_k[k] = device::math::silu_fast(k_acc); + + const float g_raw = bf16_load(g, bos * g_row_stride + i_hv * kDimK + k) + dt_bias[hk]; + if constexpr (kUseLowerBound) { + s_decay[k] = __expf(lower_bound * device::math::sigmoid_fast(exp_a * g_raw)); + } else { + s_decay[k] = __expf(-exp_a * device::math::softplus_fast(g_raw)); + } + } + } + + if constexpr (kUpdateConvState) { + if (tid < kDimV) { + const int v = tid; + const int hvv = hv_off + v; + const int64_t cs_base = slot * cs_slot_stride + hvv; + const int64_t xv_idx = bos * x_row_stride + hvv; + + float v_acc = bias_v[hvv]; + __nv_bfloat16 v_shift0 = __float2bfloat16(0.0f); + __nv_bfloat16 v_shift1 = __float2bfloat16(0.0f); +#pragma unroll + for (int w = 0; w < kConvStateWidth; ++w) { + const __nv_bfloat16 v_state = cs_v[cs_base + w * cs_w_stride]; + v_acc += __bfloat162float(v_state) * fp32_at(w_v_t, w * hvv_dim + hvv); + if (w == 1) { + v_shift0 = v_state; + } else if (w == 2) { + v_shift1 = v_state; + } + } + const __nv_bfloat16 v_new = x_v[xv_idx]; + v_acc += __bfloat162float(v_new) * fp32_at(w_v_t, (kKernelWidth - 1) * hvv_dim + hvv); + cs_v[cs_base + 0] = v_shift0; + cs_v[cs_base + cs_w_stride] = v_shift1; + cs_v[cs_base + 2 * cs_w_stride] = v_new; + s_v[v] = device::math::silu_fast(v_acc); + + if constexpr (kApplyOnorm && kPreloadOnormParams) { + const int64_t onorm_idx = i_n * onormg_row_stride + i_hv * kDimV + v; + pre_onorm_gate = device::math::sigmoid_fast(bf16_load(onorm_g, onorm_idx)); + pre_onorm_weight = onorm_weight[v]; + } + } + } else { + if (tid < kDimV) { + const int v = tid; + const int hvv = hv_off + v; + + float v_acc = bias_v[hvv]; +#pragma unroll + for (int w = 0; w < kConvStateWidth; ++w) { + const int64_t cs_idx = slot * cs_slot_stride + hvv + w * cs_w_stride; + v_acc += bf16_load(cs_v, cs_idx) * fp32_at(w_v_t, w * hvv_dim + hvv); + } + v_acc += bf16_load(x_v, bos * x_row_stride + hvv) * fp32_at(w_v_t, (kKernelWidth - 1) * hvv_dim + hvv); + s_v[v] = device::math::silu_fast(v_acc); + + if constexpr (kApplyOnorm && kPreloadOnormParams) { + const int64_t onorm_idx = i_n * onormg_row_stride + i_hv * kDimV + v; + pre_onorm_gate = device::math::sigmoid_fast(bf16_load(onorm_g, onorm_idx)); + pre_onorm_weight = onorm_weight[v]; + } + } + } + + if (tid == 0) { + const float beta_raw = bf16_load(beta, bos * beta_row_stride + i_hv); + if constexpr (kApplyBetaSigmoid) { + s_beta = device::math::sigmoid_fast(beta_raw); + } else { + s_beta = beta_raw; + } + } + __syncthreads(); + + if constexpr (!kUseTmaLoad && kPrefetchNextStateChunk && kNumChunks > 1) { + cp_async_state_chunk(s_state, state, slot, i_hv, state_slot_stride, 1); + } + + const float q_sq = tid < kDimK ? s_q[tid] * s_q[tid] : 0.0f; + const float k_sq = tid < kDimK ? s_k[tid] * s_k[tid] : 0.0f; + Sum2 qk_sum; + if constexpr (kUseActiveQkReduction) { + qk_sum = block_reduce_sum2_active_for(q_sq, k_sq, s_reduce); + } else { + qk_sum = block_reduce_sum2(q_sq, k_sq, s_reduce); + } + if (tid < kDimK) { + s_q[tid] *= rsqrtf(qk_sum.x + 1.0e-6f) * scale; + s_k[tid] *= rsqrtf(qk_sum.y + 1.0e-6f); + } + __syncthreads(); + + const int k_base = lane * 4; + const float4 q4 = *reinterpret_cast(s_q + k_base); + const float4 k4 = *reinterpret_cast(s_k + k_base); + const float4 decay4 = *reinterpret_cast(s_decay + k_base); + float r_q[4] = {q4.x, q4.y, q4.z, q4.w}; + float r_k[4] = {k4.x, k4.y, k4.z, k4.w}; + float r_decay[4] = {decay4.x, decay4.y, decay4.z, decay4.w}; + float o_sumsq = 0.0f; + +#pragma unroll + for (int chunk = 0; chunk < kNumChunks; ++chunk) { + if constexpr (kUseTmaLoad) { + if (tid == 0 && chunk + 2 < kNumChunks && chunk + 2 < kTmaStages) { + tma_state_chunk_stage( + s_state, state, slot, i_hv, state_slot_stride, chunk + 2, chunk + 2, &s_tma_bar[chunk + 2]); + } + ptx::mbar_wait_parity(&s_tma_bar[chunk % kTmaStages], (chunk / kTmaStages) & 1); + } else if constexpr (kPrefetchNextStateChunk && kNumChunks > 1) { + if (chunk + 1 < kNumChunks) { + ptx::cp_async_wait_group<1>(); + } else { + ptx::cp_async_wait_all(); + } + } else { + ptx::cp_async_wait_all(); + } + if constexpr (!kUseTmaLoad && !kSkipWarpSync) { + __syncwarp(); + } + + if constexpr (!kUseTmaLoad && !kPrefetchNextStateChunk) { + if (chunk + 1 < kNumChunks) { + cp_async_state_chunk(s_state, state, slot, i_hv, state_slot_stride, chunk + 1); + } + } + + const float* state_stage = s_state + (kUseTmaLoad ? (chunk % kTmaStages) : (chunk & 1)) * kChunkV * kDimK; + +#pragma unroll + for (int row = 0; row < kRowsPerWarp; row += 2) { + const int v_row_a = warp + row * kWarps; + const int v_row_b = warp + (row + 1) * kWarps; + const int v0 = chunk * kChunkV + v_row_a; + const int v1 = chunk * kChunkV + v_row_b; + float h_a_vals[4]; + float h_b_vals[4]; + float dot_hk_a = 0.0f; + float dot_hk_b = 0.0f; + + const float4 raw_h_a = *reinterpret_cast(state_stage + v_row_a * kDimK + k_base); + const float4 raw_h_b = *reinterpret_cast(state_stage + v_row_b * kDimK + k_base); + h_a_vals[0] = raw_h_a.x * r_decay[0]; + h_a_vals[1] = raw_h_a.y * r_decay[1]; + h_a_vals[2] = raw_h_a.z * r_decay[2]; + h_a_vals[3] = raw_h_a.w * r_decay[3]; + h_b_vals[0] = raw_h_b.x * r_decay[0]; + h_b_vals[1] = raw_h_b.y * r_decay[1]; + h_b_vals[2] = raw_h_b.z * r_decay[2]; + h_b_vals[3] = raw_h_b.w * r_decay[3]; + dot_hk_a = h_a_vals[0] * r_k[0] + h_a_vals[1] * r_k[1] + h_a_vals[2] * r_k[2] + h_a_vals[3] * r_k[3]; + dot_hk_b = h_b_vals[0] * r_k[0] + h_b_vals[1] * r_k[1] + h_b_vals[2] * r_k[2] + h_b_vals[3] * r_k[3]; + + const Sum2 dot_hk = warp_reduce_sum_pair(dot_hk_a, dot_hk_b); + const float v_new0 = (s_v[v0] - dot_hk.x) * s_beta; + const float v_new1 = (s_v[v1] - dot_hk.y) * s_beta; + + float dot_hq_a = 0.0f; + float dot_hq_b = 0.0f; + // Writeback mirrors the load addressing: slot pitch from the host stride + // (int64, envelope-safe), intra-slot offset i_hv*V*K + v*K + k contiguous. + const int64_t slot_base_wb = static_cast(slot) * state_slot_stride; + const int64_t state_idx_a = slot_base_wb + ((i_hv * kDimV + v0) * kDimK + k_base); + const int64_t state_idx_b = slot_base_wb + ((i_hv * kDimV + v1) * kDimK + k_base); + const float h_a_0 = h_a_vals[0] + r_k[0] * v_new0; + const float h_a_1 = h_a_vals[1] + r_k[1] * v_new0; + const float h_a_2 = h_a_vals[2] + r_k[2] * v_new0; + const float h_a_3 = h_a_vals[3] + r_k[3] * v_new0; + const float h_b_0 = h_b_vals[0] + r_k[0] * v_new1; + const float h_b_1 = h_b_vals[1] + r_k[1] * v_new1; + const float h_b_2 = h_b_vals[2] + r_k[2] * v_new1; + const float h_b_3 = h_b_vals[3] + r_k[3] * v_new1; + if constexpr (kComputeOutputBeforeStore) { + dot_hq_a = h_a_0 * r_q[0] + h_a_1 * r_q[1] + h_a_2 * r_q[2] + h_a_3 * r_q[3]; + dot_hq_b = h_b_0 * r_q[0] + h_b_1 * r_q[1] + h_b_2 * r_q[2] + h_b_3 * r_q[3]; + store_state_float4(state + state_idx_a, make_float4(h_a_0, h_a_1, h_a_2, h_a_3)); + store_state_float4(state + state_idx_b, make_float4(h_b_0, h_b_1, h_b_2, h_b_3)); + } else { + store_state_float4(state + state_idx_a, make_float4(h_a_0, h_a_1, h_a_2, h_a_3)); + store_state_float4(state + state_idx_b, make_float4(h_b_0, h_b_1, h_b_2, h_b_3)); + dot_hq_a = h_a_0 * r_q[0] + h_a_1 * r_q[1] + h_a_2 * r_q[2] + h_a_3 * r_q[3]; + dot_hq_b = h_b_0 * r_q[0] + h_b_1 * r_q[1] + h_b_2 * r_q[2] + h_b_3 * r_q[3]; + } + + const Sum2 dot_hq = warp_reduce_sum_pair(dot_hq_a, dot_hq_b); + if (lane == 0) { + s_o[v0] = dot_hq.x; + s_o[v1] = dot_hq.y; + if constexpr (kApplyOnorm && kAccumulateOnormSumsq) { + o_sumsq += dot_hq.x * dot_hq.x + dot_hq.y * dot_hq.y; + } + } + } + + if constexpr (kUseTmaLoad && kTmaStages < kNumChunks) { + // chunk + kTmaStages reuses the stage this chunk just read; every warp + // must be done with it before the single issuing thread overwrites it. + if (chunk + kTmaStages < kNumChunks) { + __syncthreads(); + if (tid == 0) { + tma_state_chunk_stage( + s_state, + state, + slot, + i_hv, + state_slot_stride, + chunk + kTmaStages, + chunk % kTmaStages, + &s_tma_bar[chunk % kTmaStages]); + } + } + } else if constexpr (!kUseTmaLoad && kPrefetchNextStateChunk) { + if (chunk + 2 < kNumChunks) { + cp_async_state_chunk(s_state, state, slot, i_hv, state_slot_stride, chunk + 2); + } + } + } + __syncthreads(); + + device::PDLTriggerSecondary(); + + if constexpr (kApplyOnorm) { + if constexpr (kAccumulateOnormSumsq) { + if (lane == 0) { + s_reduce[warp] = o_sumsq; + } + __syncthreads(); + + float total_sumsq = 0.0f; + if (warp == 0) { + total_sumsq = lane < kWarps ? s_reduce[lane] : 0.0f; + total_sumsq = device::warp::reduce_sum(total_sumsq); + if (lane == 0) { + s_reduce[0] = total_sumsq; + } + } + __syncthreads(); + + if (tid < kDimV) { + const int out_idx = (i_n * hv_count + i_hv) * kDimV + tid; + const float raw_o = s_o[tid]; + const float rstd = rsqrtf(s_reduce[0] / static_cast(kDimV) + onorm_eps); + float gate; + float weight; + if constexpr (kPreloadOnormParams) { + gate = pre_onorm_gate; + weight = pre_onorm_weight; + } else { + gate = device::math::sigmoid_fast(bf16_load(onorm_g, i_n * onormg_row_stride + i_hv * kDimV + tid)); + weight = onorm_weight[tid]; + } + const float y = raw_o * rstd * weight * gate; + out[out_idx] = bf16_store(y); + } + } else { + const float raw_o = tid < kDimV ? s_o[tid] : 0.0f; + const float o_sq = raw_o * raw_o; + float sumsq; + if constexpr (kUseActiveOnormReduction || kUseActiveQkReduction) { + sumsq = block_reduce_sum_active_for(o_sq, s_reduce); + } else { + sumsq = block_reduce_sum(o_sq, s_reduce); + } + + if (tid < kDimV) { + const int out_idx = (i_n * hv_count + i_hv) * kDimV + tid; + const float rstd = rsqrtf(sumsq / static_cast(kDimV) + onorm_eps); + float gate; + float weight; + if constexpr (kPreloadOnormParams) { + gate = pre_onorm_gate; + weight = pre_onorm_weight; + } else { + gate = device::math::sigmoid_fast(bf16_load(onorm_g, i_n * onormg_row_stride + i_hv * kDimV + tid)); + weight = onorm_weight[tid]; + } + const float y = raw_o * rstd * weight * gate; + out[out_idx] = bf16_store(y); + } + } + } else { + if (tid < kDimV) { + const int out_idx = (i_n * hv_count + i_hv) * kDimV + tid; + out[out_idx] = bf16_store(s_o[tid]); + } + } +} + +// K3 decode configuration of the many-heads kernel: onorm fused, static +// H = HV in {3, 6, 12} layout with a (B, HV) grid, onorm params preloaded, next state +// chunk prefetched, active onorm reduction, conv cache updated in place, +// beta sigmoid in-kernel. Both forget-gate variants are compiled (softplus +// and lower-bounded sigmoid) and selected at launch from the model config. +// kUseTmaLoad/kTmaStages select the 1D-TMA state-staging path in place of the +// cp.async fallback used for a misaligned recurrent-state slot stride. +template +constexpr auto kda_fused_decode_k3_kernel = kda_decode_fusion_many_heads_kernel< + /*kApplyOnorm=*/true, + /*kUseStaticDecodeLayout=*/true, + /*kFixedHeads=*/kFixedHeads, + /*kFixedValueHeads=*/kFixedHeads, + /*kUseHeadGrid=*/true, + /*kAccumulateOnormSumsq=*/false, + /*kUseActiveQkReduction=*/false, + /*kUseCacheGlobalStore=*/false, + /*kComputeOutputBeforeStore=*/false, + /*kSkipWarpSync=*/false, + /*kPreloadOnormParams=*/true, + /*kPrefetchNextStateChunk=*/true, + /*kUseActiveOnormReduction=*/true, + /*kUpdateConvState=*/true, + kUseLowerBound, + /*kApplyBetaSigmoid=*/true, + kUseTmaLoad, + kTmaStages, + kUsePDL>; + +template +auto select_kda_fused_decode_k3_kernel(bool use_lower_bound, int tma_stages) { + if (tma_stages == 3) { + return use_lower_bound ? kda_fused_decode_k3_kernel + : kda_fused_decode_k3_kernel; + } + if (tma_stages == 4) { + return use_lower_bound ? kda_fused_decode_k3_kernel + : kda_fused_decode_k3_kernel; + } + return use_lower_bound ? kda_fused_decode_k3_kernel + : kda_fused_decode_k3_kernel; +} + +template +struct KdaFusedDecodeKernel { + static void + run(const tvm::ffi::TensorView mixed_qkv, // [B, 3*H*128] bf16, row-strided + const tvm::ffi::TensorView a, // [B, H*128] bf16 raw forget gate + const tvm::ffi::TensorView b, // [B, H] bf16 raw beta logits + const tvm::ffi::TensorView conv_states, // [slots, 3, conv_dim] bf16 pool + const tvm::ffi::TensorView w_q_t, // [4, H*128] fp32 dense + const tvm::ffi::TensorView w_k_t, // [4, H*128] fp32 dense + const tvm::ffi::TensorView w_v_t, // [4, H*128] fp32 dense + const tvm::ffi::TensorView conv_bias, // [3*H*128] fp32 (zeros if none) + const tvm::ffi::TensorView A_log, // [H] fp32 + const tvm::ffi::TensorView dt_bias, // [H*128] fp32 + const tvm::ffi::TensorView onorm_g, // [B, H*128] bf16, row-strided + const tvm::ffi::TensorView onorm_weight, // [128] fp32 + const tvm::ffi::TensorView state, // [slots, H, 128, 128] fp32; inner-contiguous, any slot pitch + const tvm::ffi::TensorView indices, // [B] int32 (< 0 = padded slot) + const tvm::ffi::TensorView out, // [B, H*128] bf16 dense + double scale, + double onorm_eps, + double lower_bound, + bool use_lower_bound) { + using namespace host; + + const int64_t kH = A_log.shape()[0]; + RuntimeCheck(kH == 3 || kH == 6 || kH == 12, "KDA fused decode supports local head counts 3, 6, or 12, got ", kH); + const int64_t kSeg = kH * 128; // q, k and v segment width + + auto B_ = SymbolicSize{"batch"}; + auto Slots_ = SymbolicSize{"pool_slots"}; + auto device = SymbolicDevice{}; + device.set_options(); + + TensorMatcher({B_, 3 * kSeg}).with_dtype().with_device(device).with_strides({-1, 1}).verify(mixed_qkv); + TensorMatcher({B_, kSeg}).with_dtype().with_device(device).with_strides({-1, 1}).verify(a); + TensorMatcher({B_, kH}).with_dtype().with_device(device).with_strides({-1, 1}).verify(b); + TensorMatcher({Slots_, 3, 3 * kSeg}) + .with_dtype() + .with_device(device) + .with_strides({-1, -1, 1}) + .verify(conv_states); + TensorMatcher({4, kSeg}).with_dtype().with_device(device).with_strides({kSeg, 1}).verify(w_q_t); + TensorMatcher({4, kSeg}).with_dtype().with_device(device).with_strides({kSeg, 1}).verify(w_k_t); + TensorMatcher({4, kSeg}).with_dtype().with_device(device).with_strides({kSeg, 1}).verify(w_v_t); + TensorMatcher({3 * kSeg}).with_dtype().with_device(device).with_strides({1}).verify(conv_bias); + TensorMatcher({kH}).with_dtype().with_device(device).verify(A_log); + TensorMatcher({kSeg}).with_dtype().with_device(device).with_strides({1}).verify(dt_bias); + TensorMatcher({B_, kSeg}).with_dtype().with_device(device).with_strides({-1, 1}).verify(onorm_g); + TensorMatcher({128}).with_dtype().with_device(device).with_strides({1}).verify(onorm_weight); + // Slot stride (dim 0) is a wildcard: a locally-allocated pool pitches a + // slot at the dense kH*128*128, but the unified / page-major pools pitch it + // at the multi-layer envelope (56M elems on K3). The kernel reads the real + // slot pitch from state.stride(0); only the inner [HV, V, K] contiguity + // (strides {V*K, K, 1}) is load-bearing here. + TensorMatcher({Slots_, kH, 128, 128}) + .with_dtype() + .with_device(device) + .with_strides({-1, 128 * 128, 128, 1}) + .verify(state); + TensorMatcher({B_}).with_dtype().with_device(device).with_strides({1}).verify(indices); + TensorMatcher({B_, kSeg}).with_dtype().with_device(device).with_strides({kSeg, 1}).verify(out); + + const auto B = static_cast(B_.unwrap()); + if (B == 0) return; + + const auto* mixed_ptr = static_cast(mixed_qkv.data_ptr()); + auto* cs_ptr = static_cast<__nv_bfloat16*>(conv_states.data_ptr()); + const auto* bias_ptr = static_cast(conv_bias.data_ptr()); + // Real per-slot pitch of the ssm/temporal pool (elements): dense kH*128*128 + // for a local pool, the multi-layer envelope for the unified / page-major + // pools. Threaded into every ssm-state read/write; int64 avoids the + // envelope-pitch overflow. + const int64_t state_slot_stride = state.stride(0); + int tma_stages = 0; + // TMA 1D-bulk needs the per-slot source address (state + slot*stride) 16B + // aligned for every slot. state.data_ptr() is torch-aligned and each chunk + // offset is a multiple of kChunkV*kDimK*4 B, so alignment holds iff the slot + // pitch itself is 16B-aligned, i.e. state_slot_stride % 4 == 0 (fp32). The + // K3 envelope pitch (14,042,880 elems, %4==0) and the dense pitch both + // satisfy this; a pathological stride falls back to cp.async (still fully + // fused, just no TMA) rather than silently mis-addressing the descriptor. + const bool tma_slot_stride_aligned = (state_slot_stride % 4) == 0; + if (tma_slot_stride_aligned) { + // Full-state staging (4 stages, 64KB, sync-free) wins while the grid + // is small enough that occupancy isn't the limiter; past that the + // 48KB 3-stage variant (one sync for the single stage reuse) benches + // fastest. + tma_stages = B * static_cast(kH) >= 512 ? 3 : 4; + } + auto kernel = kH == 3 ? select_kda_fused_decode_k3_kernel<3, kUsePDL>(use_lower_bound, tma_stages) + : (kH == 6 ? select_kda_fused_decode_k3_kernel<6, kUsePDL>(use_lower_bound, tma_stages) + : select_kda_fused_decode_k3_kernel<12, kUsePDL>(use_lower_bound, tma_stages)); + const int smem_stages = tma_stages == 0 ? 2 : tma_stages; + const size_t smem_bytes = static_cast(smem_stages) * kChunkV * kDimK * sizeof(float); + host::RuntimeDeviceCheck( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smem_bytes))); + + LaunchKernel(dim3(B, kH), dim3(kThreads), device.unwrap(), smem_bytes) + .enable_pdl(kUsePDL)( + kernel, + /*x_q=*/mixed_ptr, + /*x_k=*/mixed_ptr + kSeg, + /*x_v=*/mixed_ptr + 2 * kSeg, + static_cast(w_q_t.data_ptr()), + static_cast(w_k_t.data_ptr()), + static_cast(w_v_t.data_ptr()), + /*bias_q=*/bias_ptr, + /*bias_k=*/bias_ptr + kSeg, + /*bias_v=*/bias_ptr + 2 * kSeg, + /*cs_q=*/cs_ptr, + /*cs_k=*/cs_ptr + kSeg, + /*cs_v=*/cs_ptr + 2 * kSeg, + static_cast(A_log.data_ptr()), + /*g=*/static_cast(a.data_ptr()), + static_cast(dt_bias.data_ptr()), + /*beta=*/static_cast(b.data_ptr()), + static_cast(onorm_g.data_ptr()), + static_cast(onorm_weight.data_ptr()), + static_cast(indices.data_ptr()), + /*cu_seqlens=*/static_cast(nullptr), + static_cast(state.data_ptr()), + static_cast<__nv_bfloat16*>(out.data_ptr()), + B, + /*H=*/static_cast(kH), + /*HV=*/static_cast(kH), + static_cast(lower_bound), + static_cast(scale), + static_cast(onorm_eps), + mixed_qkv.stride(0), + a.stride(0), + b.stride(0), + onorm_g.stride(0), + conv_states.stride(0), + conv_states.stride(1), + state_slot_stride); + } +}; + +} // namespace diff --git a/python/sglang/kernels/jit/csrc/attention/kda_packed_decode.cuh b/python/sglang/kernels/jit/csrc/attention/kda_packed_decode.cuh new file mode 100644 index 000000000..8e48f7412 --- /dev/null +++ b/python/sglang/kernels/jit/csrc/attention/kda_packed_decode.cuh @@ -0,0 +1,240 @@ +// CUDA port of the triton fused_recurrent_kda_packed_decode_kernel for +// batched decode. The triton kernel holds a whole [BV, K] fp32 state tile in +// the registers of a single warp, which caps it at ~5 TB/s of the ~9.6 TB/s +// this in-place read+write stream can reach (probe: torch inplace mul_). +// This kernel streams the state row by row instead: one warp per V-row group, +// each row is a 512B float4 load -> warp-reduced dot -> decayed delta-rule +// update -> 512B store, so loads pipeline across rows and nothing holds a +// tile. Setup (l2norm'd q/k, per-K decay, beta) is computed redundantly per +// warp - the kernel has no __syncthreads at all. +// +// Math follows the triton kernel exactly (fp32 throughout, same op order): +// h *= exp(g); t = ; v = (v - t) * sigmoid(b); h += v * k; +// o = +// with g = -exp(A_log) * softplus(a + dt_bias) (K3: no lower bound) or +// lower_bound * sigmoid(exp(A_log) * (a + dt_bias)). Warp-shuffle reduction +// order differs from tl.sum, so outputs match to ULPs, not bits (validated +// against the triton kernel with tolerance + GSM8K). + +#include // For TensorMatcher, SymbolicSize, SymbolicDevice +#include // For RuntimeCheck + +#include // For bf16_t, fp32_t, device::cast +#include // For LaunchKernel + +#include + +#include + +namespace { + +struct KdaPackedDecodeParams { + const bf16_t* __restrict__ mixed_qkv; // [B, 2*H*K + HV*V] + const bf16_t* __restrict__ a; // [B, HV*K] + const bf16_t* __restrict__ b; // [B, HV] + const fp32_t* __restrict__ A_log; // [HV] + const fp32_t* __restrict__ dt_bias; // [HV*K] + bf16_t* __restrict__ o; // [B, HV*V] (contiguous view) + fp32_t* __restrict__ state; // pool, row stride = stride_state + const int32_t* __restrict__ indices; // [B] + int64_t stride_mixed; + int64_t stride_a; + int64_t stride_b; + int64_t stride_state; // elements per pool slot + uint32_t H; + uint32_t HV; + fp32_t scale; + fp32_t lower_bound; + int32_t use_lower_bound; +}; + +__device__ __forceinline__ float warp_allreduce_sum(float v) { +#if defined(__HIP_PLATFORM_AMD__) + constexpr uint64_t kFullMask = 0xffffffffffffffffull; +#else + constexpr uint32_t kFullMask = 0xffffffffu; +#endif +#pragma unroll + for (int off = 16; off > 0; off >>= 1) { + v += __shfl_xor_sync(kFullMask, v, off); + } + return v; +} + +// K = V = 128 specialization: one lane owns 4 consecutive K-elements (16B). +template +__global__ +__launch_bounds__(kWarps * 32) void kda_packed_decode_kernel(const KdaPackedDecodeParams __grid_constant__ params) { + using namespace device; + constexpr int K = 128; + constexpr int V = 128; + constexpr int kElems = 4; // K / 32 lanes + + const uint32_t i_nh = blockIdx.x; + const uint32_t n = i_nh / params.HV; + const uint32_t hv = i_nh % params.HV; + const uint32_t i_h = hv / (params.HV / params.H); + const uint32_t warp = threadIdx.x >> 5; + const uint32_t lane = threadIdx.x & 31; + + PDLWaitPrimary(); + + bf16_t* o_ptr = params.o + (static_cast(n) * params.HV + hv) * V; + const int64_t sidx = params.indices[n]; + if (sidx < 0) { + // Padded cuda-graph slot: zero the output, leave the pool untouched. + for (uint32_t i = threadIdx.x; i < V; i += kWarps * 32) { + o_ptr[i] = cast(0.0f); + } + PDLTriggerSecondary(); + return; + } + + // --- per-warp redundant setup (no cross-warp synchronization) --- + const bf16_t* mixed = params.mixed_qkv + n * params.stride_mixed; + const uint32_t e0 = lane * kElems; + + float q[kElems], k[kElems]; + float q_sq = 0.0f, k_sq = 0.0f; +#pragma unroll + for (int e = 0; e < kElems; ++e) { + q[e] = cast(mixed[i_h * K + e0 + e]); + k[e] = cast(mixed[params.H * K + i_h * K + e0 + e]); + q_sq += q[e] * q[e]; + k_sq += k[e] * k[e]; + } + // tl: q / sqrt(sum(q*q) + 1e-6), then * scale + const float q_inv = 1.0f / sqrtf(warp_allreduce_sum(q_sq) + 1e-6f); + const float k_inv = 1.0f / sqrtf(warp_allreduce_sum(k_sq) + 1e-6f); +#pragma unroll + for (int e = 0; e < kElems; ++e) { + q[e] = q[e] * q_inv * params.scale; + k[e] = k[e] * k_inv; + } + + const float exp_A = expf(params.A_log[hv]); + float decay[kElems]; +#pragma unroll + for (int e = 0; e < kElems; ++e) { + const float x = cast(params.a[n * params.stride_a + hv * K + e0 + e]) + params.dt_bias[hv * K + e0 + e]; + float g; + if (params.use_lower_bound) { + g = params.lower_bound / (1.0f + expf(-exp_A * x)); + } else { + const float sp = (x <= 20.0f) ? logf(1.0f + expf(x)) : x; + g = -exp_A * sp; + } + decay[e] = expf(g); + } + const float beta = 1.0f / (1.0f + expf(-cast(params.b[n * params.stride_b + hv]))); + + const bf16_t* v_ptr = mixed + 2 * params.H * K + hv * V; + fp32_t* h_base = params.state + sidx * params.stride_state + static_cast(hv) * V * K; + + // --- stream this warp's V-rows: 512B load -> update -> 512B store --- + constexpr int kRowsPerWarp = V / kWarps; +#pragma unroll 4 + for (int r = warp * kRowsPerWarp; r < (int)(warp + 1) * kRowsPerWarp; ++r) { + float4 h4 = *reinterpret_cast(h_base + r * K + e0); + float h[kElems] = {h4.x, h4.y, h4.z, h4.w}; + float t = 0.0f; +#pragma unroll + for (int e = 0; e < kElems; ++e) { + h[e] *= decay[e]; + t += h[e] * k[e]; + } + t = warp_allreduce_sum(t); + const float v_new = (cast(v_ptr[r]) - t) * beta; + float o_acc = 0.0f; +#pragma unroll + for (int e = 0; e < kElems; ++e) { + h[e] += v_new * k[e]; + o_acc += h[e] * q[e]; + } + o_acc = warp_allreduce_sum(o_acc); + *reinterpret_cast(h_base + r * K + e0) = make_float4(h[0], h[1], h[2], h[3]); + if (lane == 0) { + o_ptr[r] = cast(o_acc); + } + } + + PDLTriggerSecondary(); +} + +template +struct KdaPackedDecodeKernel { + static constexpr auto kernel = kda_packed_decode_kernel; + + static void + run(const tvm::ffi::TensorView mixed_qkv, + const tvm::ffi::TensorView a, + const tvm::ffi::TensorView b, + const tvm::ffi::TensorView A_log, + const tvm::ffi::TensorView dt_bias, + const tvm::ffi::TensorView o, + const tvm::ffi::TensorView state, + const tvm::ffi::TensorView indices, + double scale, + double lower_bound, + bool use_lower_bound, + int64_t num_q_heads) { + using namespace host; + + auto B_ = SymbolicSize{"batch"}; + auto MixedDim_ = SymbolicSize{"mixed_dim"}; + auto ADim_ = SymbolicSize{"a_dim"}; + auto HV_ = SymbolicSize{"num_v_heads"}; + auto V_ = SymbolicSize{"head_v_dim"}; + auto K_ = SymbolicSize{"head_k_dim"}; + auto Slots_ = SymbolicSize{"pool_slots"}; + auto device = SymbolicDevice{}; + device.set_options(); + + TensorMatcher({B_, MixedDim_}).with_dtype().with_device(device).with_strides({-1, 1}).verify(mixed_qkv); + TensorMatcher({B_, ADim_}).with_dtype().with_device(device).with_strides({-1, 1}).verify(a); + TensorMatcher({B_, HV_}).with_dtype().with_device(device).with_strides({-1, 1}).verify(b); + TensorMatcher({HV_}).with_dtype().with_device(device).verify(A_log); + TensorMatcher({ADim_}).with_dtype().with_device(device).verify(dt_bias); + TensorMatcher({B_, HV_, V_}).with_dtype().with_device(device).verify(o); + TensorMatcher({Slots_, HV_, V_, K_}) + .with_dtype() + .with_device(device) + .with_strides({-1, -1, -1, 1}) + .verify(state); + TensorMatcher({B_}).with_dtype().with_device(device).verify(indices); + + const auto B = static_cast(B_.unwrap()); + const auto HV = static_cast(HV_.unwrap()); + const auto H = static_cast(num_q_heads); + RuntimeCheck(K_.unwrap() == 128 && V_.unwrap() == 128, "kda_packed_decode is specialized for K = V = 128"); + RuntimeCheck( + ADim_.unwrap() == HV * 128 && H > 0 && HV % H == 0, "a/dt_bias must be [*, HV*K] and HV divisible by H"); + RuntimeCheck(MixedDim_.unwrap() == 2 * H * 128 + HV * 128, "mixed_qkv last dim must be 2*H*K + HV*V"); + RuntimeCheck(state.stride(1) == 128 * 128 && state.stride(2) == 128, "state inner layout must be dense [HV, V, K]"); + if (B == 0) return; + + const auto params = KdaPackedDecodeParams{ + .mixed_qkv = static_cast(mixed_qkv.data_ptr()), + .a = static_cast(a.data_ptr()), + .b = static_cast(b.data_ptr()), + .A_log = static_cast(A_log.data_ptr()), + .dt_bias = static_cast(dt_bias.data_ptr()), + .o = static_cast(o.data_ptr()), + .state = static_cast(state.data_ptr()), + .indices = static_cast(indices.data_ptr()), + .stride_mixed = mixed_qkv.stride(0), + .stride_a = a.stride(0), + .stride_b = b.stride(0), + .stride_state = state.stride(0), + .H = H, + .HV = HV, + .scale = static_cast(scale), + .lower_bound = static_cast(lower_bound), + .use_lower_bound = use_lower_bound ? 1 : 0, + }; + + LaunchKernel(B * HV, kWarps * 32, device.unwrap()).enable_pdl(kUsePDL)(kernel, params); + } +}; + +} // namespace diff --git a/python/sglang/kernels/jit/csrc/attention/kda_prefill.cu b/python/sglang/kernels/jit/csrc/attention/kda_prefill.cu new file mode 100644 index 000000000..f1b467680 --- /dev/null +++ b/python/sglang/kernels/jit/csrc/attention/kda_prefill.cu @@ -0,0 +1,3871 @@ +// kda_prefill.cu — self-contained KDA prefill forward (inference drop-in for +// the FLA Triton chunk_kda_fwd path). Single-file artifact carrying ONLY the +// shipping default paths — two routes, picked per shape from the launch +// geometry (pick_route): +// FUSED (long single sequences) — +// eqlen (T % 64 == 0, no cu_seqlens): +// nc >= 4: kda_fused — ONE grid: NP piece-builders (k1 factors + +// in-tail running map composition) + NP trailing self-start +// chain blocks gated on per-piece flags. +// nc < 4: k1_factors_mma + k2_chain_tc (NP == 1 two-kernel path). +// varlen (cu_seqlens / ragged T): kda_fused_vl over a host-built +// per-sequence piece table (non-tail pieces run the exact eqlen bodies). +// SEQ0 (many sequences / high H): the tail-free dieted builder +// k1_tf_builder / k1_tf_builder_vl (no composition, 2 CTAs/SM), +// then ONE chain per (sequence, head) walking from h0 — k2_chain_tc / +// k2_chain_tc_vl. Its P/u0 differ from the fused route's by design (P4-lo +// and the block apply, envelope-gated) — see the seq0 section. +// Gate modes (GM): 0 = pre-transformed bf16 glog; 1 = raw softplus, +// 2 = raw safe_gate — both transformed in place off the serial path +// (eqlen AND varlen). RAW and BSIG (orthogonal to GM and to each other, same +// dispatch): RAW = q/k arrive un-normalized, so the caller's l2norm folds into +// the tiles k1 already loads (row_rnorm); BSIG = beta arrives as logits, so +// its sigmoid folds into the beta read. +// Chunk size 64, head dim K == V == 128 fixed. +// +// C++ API (pybind, see kda_prefill.py for the FLA-signature wrapper): +// kda_prefill_fwd(q, k, v, g, beta, scale, initial_state, cu_seqlens, +// use_gate_in_kernel, A_log, dt_bias, safe_gate, +// lower_bound, use_qk_l2norm_in_kernel, +// use_beta_sigmoid_in_kernel, h_per_chunk, h_v_first) +// -> (o [T,H,128] bf16, Sf [N,H,128,128] f32) +// h_per_chunk is an optional preallocated per-chunk state output. +// +// Build (torch cpp_extension JIT): +// -O3 -std=c++20 -gencode arch=compute_103a,code=sm_103a -use_fast_math +// -lineinfo, link -lcuda (cuTensorMapEncodeTiled). sm_103a (GB300) only. + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// Inlined PTX helpers: only the subset the kernels below use, kept local so this +// translation unit compiles standalone. +namespace ptx { + +template +static __device__ __forceinline__ uint32_t to_shared(T* ptr) { + return static_cast(__cvta_generic_to_shared(ptr)); +} + +// ---- mbarrier ---- +static __device__ __forceinline__ void mbar_init(uint64_t* bar, uint32_t count) { + asm volatile("mbarrier.init.shared.b64 [%0], %1;" ::"r"(to_shared(bar)), "r"(count)); +} +static __device__ __forceinline__ void mbar_arrive_expect_tx(uint64_t* bar, uint32_t bytes) { + asm volatile("mbarrier.arrive.expect_tx.shared.b64 _, [%0], %1;" ::"r"(to_shared(bar)), "r"(bytes)); +} +static __device__ __forceinline__ void mbar_wait_parity(uint64_t* bar, uint32_t parity) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "WAIT_%=: mbarrier.try_wait.parity.shared.b64 p, [%0], %1;\n\t" + "@!p bra WAIT_%=;\n\t}\n" ::"r"(to_shared(bar)), + "r"(parity)); +} + +// ---- cp.async (per-thread 16 B gmem->smem) ---- +static __device__ __forceinline__ void cp_async_16(void* smem, const void* gmem) { + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;" ::"r"(to_shared(smem)), "l"(gmem)); +} +// ignore-src predicate form: pad rows stage zeros branchlessly +static __device__ __forceinline__ void cp_async_16_zfill(void* smem, const void* gmem, int ignore_src) { + asm volatile( + "{\n\t.reg .pred pz;\n\t" + "setp.ne.b32 pz, %2, 0;\n\t" + "cp.async.cg.shared.global [%0], [%1], 16, pz;\n\t}" ::"r"(to_shared(smem)), + "l"(gmem), + "r"(ignore_src)); +} +static __device__ __forceinline__ void cp_async_commit() { + asm volatile("cp.async.commit_group;"); +} +static __device__ __forceinline__ void cp_async_wait_pending(int pending) { + switch (pending) { + case 0: + asm volatile("cp.async.wait_group 0;"); + break; + default: + asm volatile("cp.async.wait_group 1;"); + break; + } +} + +// ---- TMA (2D tiled bulk loads) ---- +static __device__ __forceinline__ void prefetch_tensormap(const void* tmap) { + asm volatile("prefetch.tensormap [%0];" ::"l"(tmap) : "memory"); +} +static __device__ __forceinline__ void +cp_async_bulk_tensor_2d_load(uint32_t dst_smem, const CUtensorMap* tmap, int32_t x, int32_t y, uint64_t* bar) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile.mbarrier::" + "complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" ::"r"(dst_smem), + "l"(tmap), + "r"(x), + "r"(y), + "r"(to_shared(bar)) + : "memory"); +} + +// ---- ldmatrix + warp mma (k1's pair-tile products) ---- +static __device__ __forceinline__ void +ldmatrix_x4_b16(uint32_t row_addr, uint32_t& r0, uint32_t& r1, uint32_t& r2, uint32_t& r3) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared::cta.b16 {%0, %1, %2, %3}, [%4];" + : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) + : "r"(row_addr)); +} +static __device__ __forceinline__ void ldmatrix_x2_b16(uint32_t row_addr, uint32_t& r0, uint32_t& r1) { + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared::cta.b16 {%0, %1}, [%2];" : "=r"(r0), "=r"(r1) : "r"(row_addr)); +} +static __device__ __forceinline__ void +mma_m16n8k16_bf16f32(float4& d, uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, uint32_t b1) { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{%0,%1,%2,%3},{%4,%5,%6,%7},{%8,%9},{%0,%1,%2,%3};" + : "+f"(d.x), "+f"(d.y), "+f"(d.z), "+f"(d.w) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); +} + +// ---- gpu-scoped release/acquire flags (fused-grid piece gating) ---- +static __device__ __forceinline__ void red_add_rel_b32(uint32_t* ptr, uint32_t value) { + asm volatile("red.release.gpu.global.add.u32 [%0], %1;" ::"l"(ptr), "r"(value)); +} +static __device__ __forceinline__ uint32_t ld_acq_b32(const uint32_t* ptr) { + uint32_t ret; + asm volatile("ld.acquire.gpu.global.b32 %0, [%1];" : "=r"(ret) : "l"(ptr)); + return ret; +} +// make generic-proxy global writes visible to the async proxy (consumers +// read the published factor tensors via TMA) +static __device__ __forceinline__ void fence_async_global() { + asm volatile("fence.proxy.async.global;"); +} + +// ---- tcgen05 (TMEM lifecycle, ld/st, MMA, fences) ---- +static __device__ __forceinline__ void tcgen05_alloc(uint32_t smem_addr_for_taddr, uint32_t n_cols) { + asm volatile( + "tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" ::"r"(smem_addr_for_taddr), "r"(n_cols)); +} +static __device__ __forceinline__ void tcgen05_dealloc(uint32_t taddr, uint32_t n_cols) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" ::"r"(taddr), "r"(n_cols)); +} +static __device__ __forceinline__ void tcgen05_relinquish() { + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); +} +static __device__ __forceinline__ void +tcgen05_st_32x32b_x4(uint32_t taddr, uint32_t r0, uint32_t r1, uint32_t r2, uint32_t r3) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32 [%0], {%1, %2, %3, %4};" ::"r"(taddr), "r"(r0), "r"(r1), "r"(r2), "r"(r3)); +} +static __device__ __forceinline__ void tcgen05_st_32x32b_x8( + uint32_t taddr, + uint32_t r0, + uint32_t r1, + uint32_t r2, + uint32_t r3, + uint32_t r4, + uint32_t r5, + uint32_t r6, + uint32_t r7) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x8.b32 [%0], " + " {%1, %2, %3, %4, %5, %6, %7, %8};" ::"r"(taddr), + "r"(r0), + "r"(r1), + "r"(r2), + "r"(r3), + "r"(r4), + "r"(r5), + "r"(r6), + "r"(r7)); +} +static __device__ __forceinline__ void tcgen05_st_32x32b_x16(uint32_t taddr, const uint32_t* r) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32 [%16], " + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15};" ::"r"(r[0]), + "r"(r[1]), + "r"(r[2]), + "r"(r[3]), + "r"(r[4]), + "r"(r[5]), + "r"(r[6]), + "r"(r[7]), + "r"(r[8]), + "r"(r[9]), + "r"(r[10]), + "r"(r[11]), + "r"(r[12]), + "r"(r[13]), + "r"(r[14]), + "r"(r[15]), + "r"(taddr)); +} +static __device__ __forceinline__ void tcgen05_ld_32x32b_x8( + uint32_t taddr, + uint32_t& r0, + uint32_t& r1, + uint32_t& r2, + uint32_t& r3, + uint32_t& r4, + uint32_t& r5, + uint32_t& r6, + uint32_t& r7) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x8.b32 " + " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3), "=r"(r4), "=r"(r5), "=r"(r6), "=r"(r7) + : "r"(taddr)); +} +static __device__ __forceinline__ void tcgen05_ld_32x32b_x16( + uint32_t taddr, + uint32_t& r0, + uint32_t& r1, + uint32_t& r2, + uint32_t& r3, + uint32_t& r4, + uint32_t& r5, + uint32_t& r6, + uint32_t& r7, + uint32_t& r8, + uint32_t& r9, + uint32_t& r10, + uint32_t& r11, + uint32_t& r12, + uint32_t& r13, + uint32_t& r14, + uint32_t& r15) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32 " + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=r"(r0), + "=r"(r1), + "=r"(r2), + "=r"(r3), + "=r"(r4), + "=r"(r5), + "=r"(r6), + "=r"(r7), + "=r"(r8), + "=r"(r9), + "=r"(r10), + "=r"(r11), + "=r"(r12), + "=r"(r13), + "=r"(r14), + "=r"(r15) + : "r"(taddr)); +} +static __device__ __forceinline__ void tcgen05_ld_32x32b_x32(uint32_t taddr, uint32_t* r) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32 " + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=r"(r[0]), + "=r"(r[1]), + "=r"(r[2]), + "=r"(r[3]), + "=r"(r[4]), + "=r"(r[5]), + "=r"(r[6]), + "=r"(r[7]), + "=r"(r[8]), + "=r"(r[9]), + "=r"(r[10]), + "=r"(r[11]), + "=r"(r[12]), + "=r"(r[13]), + "=r"(r[14]), + "=r"(r[15]), + "=r"(r[16]), + "=r"(r[17]), + "=r"(r[18]), + "=r"(r[19]), + "=r"(r[20]), + "=r"(r[21]), + "=r"(r[22]), + "=r"(r[23]), + "=r"(r[24]), + "=r"(r[25]), + "=r"(r[26]), + "=r"(r[27]), + "=r"(r[28]), + "=r"(r[29]), + "=r"(r[30]), + "=r"(r[31]) + : "r"(taddr)); +} +static __device__ __forceinline__ void tcgen05_wait_ld() { + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); +} +static __device__ __forceinline__ void tcgen05_wait_st() { + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); +} +static __device__ __forceinline__ void tcgen05_commit_arrive(uint64_t* bar) { + asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.b64 [%0];" ::"r"(to_shared(bar))); +} +static __device__ __forceinline__ void tcgen05_fence_before_thread_sync() { + asm volatile("tcgen05.fence::before_thread_sync;"); +} +static __device__ __forceinline__ void tcgen05_fence_after_thread_sync() { + asm volatile("tcgen05.fence::after_thread_sync;"); +} +static __device__ __forceinline__ void +tcgen05_mma_f16(uint32_t d, uint64_t desc_a, uint64_t desc_b, uint32_t inst_desc_high, uint32_t scale_c) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t}\n" ::"r"(d), + "l"(desc_a), + "l"(desc_b), + "r"(inst_desc_high), + "r"(scale_c)); +} +// kind::f16 with the A operand sourced from TMEM (packed-bf16 hi/lo A) +static __device__ __forceinline__ void +tcgen05_mma_f16_atmem(uint32_t d, uint32_t a_tmem, uint64_t desc_b, uint32_t inst_desc_high, uint32_t scale_c) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], [%1], %2, %3, p;\n\t}\n" ::"r"(d), + "r"(a_tmem), + "l"(desc_b), + "r"(inst_desc_high), + "r"(scale_c)); +} + +// ---- MMA descriptors (smem matrix desc + instruction desc, Table 44) ---- +enum class Major : uint8_t { K = 0, MN = 1 }; +enum class F16Type : uint8_t { F16 = 0, BF16 = 1 }; +enum class DType : uint8_t { F16 = 0, F32 = 1, S32 = 2 }; + +__host__ __device__ static __forceinline__ constexpr uint64_t +mma_smem_desc(uint32_t matrix_addr, uint32_t lbo, uint32_t sbo, uint32_t base_offset, int swizzle_bytes) { + auto enc = [](uint32_t x) -> uint64_t { return (uint64_t)((x & 0x3FFFFu) >> 4); }; + uint8_t code = (swizzle_bytes == 128) ? 2u : (swizzle_bytes == 64) ? 4u : (swizzle_bytes == 32) ? 6u : 0u; + uint64_t d = 0; + d |= enc(matrix_addr); // bits 0-13 + d |= enc(lbo) << 16; // bits 16-29 + d |= enc(sbo) << 32; // bits 32-45 + d |= uint64_t(1u) << 46; // version = 1 + d |= uint64_t(base_offset & 0x7u) << 49; // bits 49-51 + d |= uint64_t(code & 0x7u) << 61; // bits 61-63 + return d; +} +template +__host__ __device__ static __forceinline__ constexpr uint64_t +mma_smem_desc_k_major(uint32_t addr, uint32_t base_offset = 0) { + constexpr int K_BYTES = BLOCK_K * int(sizeof(T)); + static_assert(SWIZZLE_BYTES == K_BYTES, "K-major requires swizzle bytes == BLOCK_K * sizeof(T)"); + return mma_smem_desc(addr, /*lbo=*/0u, /*sbo=*/8u * uint32_t(K_BYTES), base_offset, SWIZZLE_BYTES); +} +// MN-major operand (inner = N for B) over swizzle-atom-form smem: adjacent +// K-row groups of 8 at SBO = 8 * SWZ within an MN-chunk, adjacent MN-chunks +// (one SWZ128 TMA box each) at LBO = BLOCK_K * SWZ. +template +__host__ __device__ static __forceinline__ constexpr uint64_t +mma_smem_desc_mn_major(uint32_t addr, uint32_t base_offset = 0) { + constexpr int BLOCK_MN_BYTES = BLOCK_MN * int(sizeof(T)); + static_assert( + SWIZZLE_BYTES == 32 || SWIZZLE_BYTES == 64 || SWIZZLE_BYTES == 128, "MN-major requires SWZ in {32, 64, 128}"); + static_assert(BLOCK_MN_BYTES % SWIZZLE_BYTES == 0, "MN-major: BLOCK_MN * sizeof(T) must tile the swizzle atom"); + return mma_smem_desc( + addr, + /*lbo=*/uint32_t(BLOCK_K * SWIZZLE_BYTES), + /*sbo=*/uint32_t(8 * SWIZZLE_BYTES), + base_offset, + SWIZZLE_BYTES); +} +__host__ __device__ static __forceinline__ constexpr uint32_t mma_inst_desc_f16( + uint32_t M, + uint32_t N, + F16Type a_type = F16Type::BF16, + F16Type b_type = F16Type::BF16, + DType d_type = DType::F32, + Major a_major = Major::K, + Major b_major = Major::K) { + uint32_t d = 0; + d |= (static_cast(d_type) & 0x3u) << 4; + d |= (static_cast(a_type) & 0x7u) << 7; + d |= (static_cast(b_type) & 0x7u) << 10; + d |= (static_cast(a_major) & 0x1u) << 15; + d |= (static_cast(b_major) & 0x1u) << 16; + d |= ((N >> 3) & 0x3Fu) << 17; + d |= ((M >> 4) & 0x1Fu) << 24; + return d; +} +__host__ __device__ static __forceinline__ constexpr uint32_t mma_inst_desc_tf32( + uint32_t M, uint32_t N, DType d_type = DType::F32, Major a_major = Major::K, Major b_major = Major::K) { + constexpr uint32_t TF32 = 2u; + uint32_t d = 0; + d |= (static_cast(d_type) & 0x3u) << 4; + d |= TF32 << 7; // atype = TF32 = 2 + d |= TF32 << 10; // btype = TF32 = 2 + d |= (static_cast(a_major) & 0x1u) << 15; + d |= (static_cast(b_major) & 0x1u) << 16; + d |= ((N >> 3) & 0x3Fu) << 17; + d |= ((M >> 4) & 0x1Fu) << 24; + return d; +} +// thin shim keeping the kernel bodies identical to the development source +enum class MmaDenseKind : uint8_t { F16 = 0 }; +template +__host__ __device__ static __forceinline__ constexpr uint32_t mma_inst_desc_dense( + uint32_t M, + uint32_t N, + F16Type a_type, + F16Type b_type, + DType d_type = DType::F32, + Major a_major = Major::K, + Major b_major = Major::K) { + static_assert(KIND == MmaDenseKind::F16); + return mma_inst_desc_f16(M, N, a_type, b_type, d_type, a_major, b_major); +} + +} // namespace ptx + +namespace kda { + +constexpr int BT = 64; // chunk tokens +constexpr int K = 128; // head dim (qk == v) +using bf16 = __nv_bfloat16; + +// balanced piece boundaries (k1's wall = the LONGEST piece: only a +// balanced split minimizes it) +__device__ static inline int piece_c0(int p, int nc, int NP) { + const int base = nc / NP, rem = nc % NP; + return base * p + (p < rem ? p : rem); +} +// kind::tf32 with the A operand from TMEM (probe-verified: identity (m,k)-> +// (lane,col) map — a fp32 D region reads directly as a K=128 A operand, +// +8-translatable slices, truncation read semantics). B via fp32 SWZ32 +// k-major K=8 chunks. +static __device__ __forceinline__ void +mma_tf32_atmem(uint32_t d, uint32_t a_tmem, uint64_t desc_b, uint32_t inst_desc_high, uint32_t scale_c) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::tf32 [%0], [%1], %2, %3, p;\n\t}\n" ::"r"(d), + "r"(a_tmem), + "l"(desc_b), + "r"(inst_desc_high), + "r"(scale_c)); +} +// bf16 piece/segment L maps (half the TMA bytes of the former fp32; the c +// maps stay fp32 — bf16 c accumulated past the fla o bar at 16K) land as +// two SWZ128 [128][64] tiles in the TOP half of their 64 KB fp32 slot and +// widen IN PLACE to the SWZ32 fp32 K=8 tiles the tf32 B descs read. fp32 +// chunks 8..15 overwrite the pad, so every thread stages all its reads in +// registers behind one barrier (all pad reads precede any overlapping +// write). Callers: 512 threads; thread (vc, ch) owns row vc, cols +// [ch*32, ch*32+32). +static __device__ __forceinline__ void widen_map_tf32(float* slot, const bf16* pad, int vc, int ch) { + uint32_t w[16]; +#pragma unroll + for (int g = 0; g < 4; ++g) { // one 16B SWZ128 atom per load + const int in0 = ch * 32 + g * 8; + reinterpret_cast(w)[g] = + *reinterpret_cast(pad + (in0 >> 6) * (K * 64) + vc * 64 + ((((in0 >> 3) & 7) ^ (vc & 7)) << 3)); + } + __syncthreads(); +#pragma unroll + for (int g = 0; g < 8; ++g) { // 4-col groups: one 16B half-row each + const int in0 = ch * 32 + g * 4; + const float2 f0 = __bfloat1622float2(*reinterpret_cast(&w[2 * g])); + const float2 f1 = __bfloat1622float2(*reinterpret_cast(&w[2 * g + 1])); + const int half = ((in0 >> 2) & 1) ^ ((vc >> 2) & 1); + *reinterpret_cast(slot + (in0 >> 3) * 1024 + vc * 8 + half * 4) = float4{f0.x, f0.y, f1.x, f1.y}; + } +} +// Per-kernel smem structs: bodies take them BY REFERENCE so the fused +// dispatcher can union them in one pool (nvcc SUMS per-function static variables). +struct K1Smem { + bf16 sgb[2][BT][K]; // [chunk-pair slot]: c+1's bf16 gates stage via + // cp.async during c's compute phases + float sgc[BT][K]; // P2a output: this chunk's fp32 gate cumsum + float sA[BT][BT]; + float sU[BT][2 * K + 4]; + float sb[BT]; + union SPool { + struct { + bf16 kw[BT][K + 8]; + bf16 qw[BT][K + 8]; + bf16 kz[10][16][K + 8]; + } a; + struct { + __align__(1024) bf16 kdT[K][64]; + __align__(1024) bf16 pTt[K][64]; + __align__(1024) bf16 u0Th[K][64]; + __align__(1024) bf16 u0Tl[K][64]; + } b; + struct { + __align__(1024) float sLt[16][128][8]; + } g; // compose B: + // fp32 L as SWZ32 k-major K=8 chunks (tf32 smem descs) + struct { + bf16 aC_h[BT][56]; + bf16 aC_l[BT][56]; // A cols<48 hi/lo + bf16 upT_h[2 * K][56]; + bf16 upT_l[2 * K][56]; + } f; + } sp; + uint64_t mb_f; + uint32_t s_taddr; +}; +struct ChainSmem { + __align__(1024) bf16 sPneg[2][2][64][64]; // [buf][ktile] + __align__(1024) bf16 sKdT[2][128][64]; + __align__(1024) bf16 sQd[2][2][64][64]; + __align__(1024) bf16 sAh[2][64][64]; + __align__(1024) bf16 sAl[2][64][64]; + // u0 chunk tile via TMA riding mb_tma (measured: this stage was memory- + // latency-bound on u0's 16 fp32 LDGs — a 100MB L2-spilling stream) + __align__(1024) float sU0[2][BT][K]; + // mb_pre/mb_preC: prefix L/c TMA; mb_cmL/mb_cmC: prefix mma-batch commits + uint64_t mb_tma, mb_p1, mb_p2, mb_o, mb_pre, mb_preC, mb_cmL, mb_cmC; + uint32_t s_taddr; +}; + +// ---------------- RAW input transforms (the fused pre-pass) ---------------- +// Two INDEPENDENT conventions, one per fla flag: RAW (q/k arrive +// un-normalized) folds fla's l2norm_fwd(q)/l2norm_fwd(k), BSIG (beta arrives +// as logits) folds its sigmoid(beta) — the caller's separate pre-pass launches +// — into the tiles k1 already loads. Both transforms reproduce the pre-pass's +// BYTES the way GM 1/2 reproduces its glog: fla's rstd is 1/sqrt(sum x^2 + +// 1e-6) with round-to-nearest sqrt/reciprocal (NOT the fast-math rsqrt) over +// an fp32 sum of bf16 squares, and its y lands in bf16, so every read +// re-rounds. +constexpr float L2_EPS = 1e-6f; // fla.modules.l2norm.l2norm_fwd default + +// {k, q} row reciprocal norms for a chunk's 64 token rows, ONE pass for both +// tensors: 512 threads = 8 lanes x 64 rows, 16 columns per lane (two 16 B +// loads), closed by 3 shuffles (a row's lanes share a warp); lane +// (tid & 7) == 0 publishes the row's pair to dst (the two rn sequences stay +// inside that branch — 7 of 8 lanes would throw them away). Pad rows load +// nothing and get rstd(0) — finite, and their values are selected to 0 +// downstream anyway. +template +static __device__ __forceinline__ void row_rnorm( + float2& dst, const bf16* __restrict__ q, const bf16* __restrict__ kk, size_t gbase, int H, int C_act, int tid) { + const int r = tid >> 3; + float sk = 0.f, sq = 0.f; + if (!VL || r < C_act) { + const size_t o = gbase + (size_t)r * H * K + (tid & 7) * 16; + const uint4 xb[4] = { + *reinterpret_cast(kk + o), + *reinterpret_cast(kk + o + 8), + *reinterpret_cast(q + o), + *reinterpret_cast(q + o + 8)}; + const __nv_bfloat162* x2 = reinterpret_cast(xb); +#pragma unroll + for (int e = 0; e < 8; ++e) { + const float2 kf = __bfloat1622float2(x2[e]); + const float2 qf = __bfloat1622float2(x2[e + 8]); + sk += kf.x * kf.x + kf.y * kf.y; + sq += qf.x * qf.x + qf.y * qf.y; + } + } +#pragma unroll + for (int m = 1; m < 8; m <<= 1) { + sk += __shfl_xor_sync(0xffffffffu, sk, m); + sq += __shfl_xor_sync(0xffffffffu, sq, m); + } + if ((tid & 7) == 0) dst = float2{__frcp_rn(__fsqrt_rn(sk + L2_EPS)), __frcp_rn(__fsqrt_rn(sq + L2_EPS))}; +} +// one element pair through fla's l2norm: x * rstd, rounded to bf16 +static __device__ __forceinline__ __nv_bfloat162 l2_bf16(float2 x, float rst) { + return __floats2bfloat162_rn(x.x * rst, x.y * rst); +} +static __device__ __forceinline__ float2 l2_round(float2 x, float rst) { + return __bfloat1622float2(l2_bf16(x, rst)); +} +// BSIG: beta logits -> fp32 sigmoid, bf16-rounded (the wrapper's own +// torch.sigmoid(blog.float()).to(bfloat16)); BSIG false = beta already +// activated, passed through +template +static __device__ __forceinline__ float beta_in(float b) { + return BSIG ? __bfloat162float(__float2bfloat16(1.f / (1.f + __expf(-b)))) : b; +} + +// ---------------- K1: factored-ratio chunk builds -------------------------- +// Anchored 16-token sub-blocks: kw[i]=k_i*e^{g_i-a(si)} (<=1), and for each +// ordered pair (si>=sj) kz(si)[j]=k_j*e^{a(si)-g_j} (<=1). A and Aqk become +// plain bf16 mma products; the solve stays fp32 SIMT. 512 threads/CTA. +// VL: varlen instantiation — compile-time-gates every pad-row guard so the +// eqlen (VL=false) kernels keep their exact original code. +// GM: gate mode — 0 reads pre-made bf16 glog; 1/2 (raw softplus/safe_gate) +// stage RAW graw through the same bf16 slot and transform it IN PLACE once +// landed — under the prior chunk's tail mma wait (piece-first / tail-less +// chunks: after their own P1 wait) — so P2a reads all modes as GM == 0. +// RAW: q/k un-normalized — the l2norm pre-pass folds in here (see row_rnorm); +// q/k feed P2b and P2c only. BSIG: beta as logits, its sigmoid folded in at +// the beta read. Independent of each other and of GM. +// When pieceL is non-null the piece's per-chunk affine maps compose IN-TAIL +// (tf32 TMEM chain) and the final (L, c) maps land at pieceL/piecec[pidx]. +template +__device__ static void k1_body( + K1Smem& S, + int job0, + int njobs, + const bf16* __restrict__ q, + const bf16* __restrict__ kk, + const bf16* __restrict__ v, + const bf16* __restrict__ glog, + const float* __restrict__ beta, + int T, + int H, + float scale, + bf16* __restrict__ P, + float* __restrict__ u0, + bf16* __restrict__ kdec, + bf16* __restrict__ qdec, + bf16* __restrict__ aqk_h, + bf16* __restrict__ aqk_l, + float* __restrict__ gC, + bf16* __restrict__ pieceL, + float* __restrict__ piecec, + int pidx, + // varlen piece coords (defaults = eqlen): global chunk c's tokens start + // at c*BT + tokoff; rows >= tend - t0 are pad (zero-filled on load) + int nc_tot = 0, + int tokoff = 0, + int tend = 0, + // GM != 0 gate-transform inputs (production lb = -5.0) + const float* __restrict__ a_log = nullptr, + const float* __restrict__ dtb = nullptr, + float lb = 0.f, + // fused-tail pTt staging TMA: the negated-P global (the W mma's + // MN-major B source; every pieceL-passing caller supplies it — the + // NP == 1 k1_factors_mma launch runs pieceL == nullptr, no tail) + const CUtensorMap* ptm = nullptr) { + constexpr int SB = 16, NSB = BT / SB; // 4 sub-blocks + const int tid = threadIdx.x; + const int warp = tid >> 5, lane = tid & 31; + auto& sA = S.sA; + auto& sU = S.sU; + auto& sb = S.sb; + auto& sp = S.sp; + auto& kw = sp.a.kw; + auto& qw = sp.a.qw; + auto& kz = sp.a.kz; + auto& mb_f = S.mb_f; + auto& s_taddr = S.s_taddr; + // RAW: the row reciprocal norms park in sU's dead pad columns (the rhs + // owns [0, 2K); 2K..2K+3 is the row-stride pad, written by nothing) — live + // from the norm pass to P2c at zero smem cost + auto rn2 = [&](int i) -> float2& { return *reinterpret_cast(&sU[i][2 * K]); }; + if (pieceL) { // fused-tail tmem/mbar once (re-init of a live mbar is + // UB; re-alloc after relinquish faults) + if (tid == 0) ptx::mbar_init(&mb_f, 1); + if (warp == 0) ptx::tcgen05_alloc(ptx::to_shared(&s_taddr), 512); + } + int mbph = 0; // mb_f phase counter (one wait per committed batch) + bool pend = false; // tf32 compose committed, wait deferred a chunk + float bnext = 0.f; // next chunk's beta (prefetched a chunk early) + // deferred compose wait: must land before ANY union write (the compose's + // B tiles overlay sp); by then the mmas long completed (~free) + auto pend_wait = [&] { + if (pend) { + ptx::mbar_wait_parity(&mb_f, mbph & 1); + ++mbph; + pend = false; + } + }; + const int nc1 = VL ? nc_tot : T / BT; // total chunks (varlen: sum nc_s) + const int tse = VL ? tend : T; // sequence end token + // GM != 0 hoist: h is constant across a block's jobs (h-major runs) + const float ga = GM != 0 ? expf(a_log[job0 / nc1]) : 0.f; + // GM != 0: in-place transform of a LANDED raw tile -> the same bf16 + // glog bytes GM == 0 stages (fp32 transform + bf16 round == the former + // P2a fused read, value-identical per element), so P2a reads every + // mode as GM == 0. Stable softplus subsumes the thr=20 branch + // (1+e^-20 == 1 in fp32). VL pad rows stay their zfill 0 + // (transform(0) != 0 but the pad algebra needs glog == 0 exactly). + auto gate_xform = [&](bf16(>)[BT][K], int rows) { + const float* dtr = dtb + (size_t)(job0 / nc1) * K; + for (int p = tid * 8; p < BT * K; p += blockDim.x * 8) { + if (VL && p / K >= rows) continue; // pad rows stay zfill 0 + uint4* gp = reinterpret_cast(>[p / K][p % K]); + uint4 g4 = *gp; + __nv_bfloat162* g2 = reinterpret_cast<__nv_bfloat162*>(&g4); + const float4 d0 = *reinterpret_cast(&dtr[p % K]); + const float4 d1 = *reinterpret_cast(&dtr[p % K + 4]); + const float dv[8] = {d0.x, d0.y, d0.z, d0.w, d1.x, d1.y, d1.z, d1.w}; +#pragma unroll + for (int e = 0; e < 4; ++e) { + const float2 gw = __bfloat1622float2(g2[e]); + float y[2]; +#pragma unroll + for (int x = 0; x < 2; ++x) { + const float g = (x ? gw.y : gw.x) + dv[2 * e + x]; + y[x] = + GM == 1 ? -ga * (fmaxf(g, 0.f) + __logf(1.f + __expf(-fabsf(g)))) : lb * (1.f / (1.f + __expf(-ga * g))); + } + g2[e] = __floats2bfloat162_rn(y[0], y[1]); + } + *gp = g4; + } + }; + // persistent job loop: job -> (c = job%nc1, h = job/nc1); the NEXT job's + // glog stages via cp.async under THIS job's compute (1-deep pipeline) + for (int sub = 0; sub < njobs; ++sub) { + const int job = job0 + sub; + const int c = job % nc1, h = job / nc1, t0 = c * BT + tokoff; + const int C_act = VL ? min(BT, tse - t0) : BT; // real rows (tail < BT) + bf16(&sgb)[BT][K] = S.sgb[sub & 1]; + float (&sg)[BT][K] = S.sgc; + if (sub == 0) { + for (int p = tid; p < BT * K; p += blockDim.x) + sgb[p / K][p % K] = + !VL || p / K < C_act ? glog[(size_t)(t0 + p / K) * H * K + h * K + p % K] : __float2bfloat16(0.f); + } else { + ptx::cp_async_wait_pending(0); + } + if (tid < BT) // beta: chunk 0 loads direct; later chunks read the + // register prefetched below (its P1 load was an exposed stall) + sb[tid] = sub == 0 ? (!VL || tid < C_act ? beta_in(beta[(size_t)(t0 + tid) * H + h]) : 0.f) : bnext; + if (sub + 1 < njobs) { // stage next job's gates (h-major decode) + const int jn = job + 1; + const int t1i = (jn % nc1) * BT + tokoff; + const size_t t1 = (size_t)t1i; + const int hn = jn / nc1; + const int Cn = VL ? min(BT, tse - t1i) : BT; // next chunk's rows + bf16* dst = &S.sgb[(sub + 1) & 1][0][0]; + for (int p = tid; p < BT * K / 8; p += blockDim.x) { + const bf16* src = glog + (t1 + p * 8 / K) * H * K + (size_t)hn * K + (p * 8) % K; + if constexpr (VL) // pad rows of a tail chunk stage as zeros + ptx::cp_async_16_zfill(dst + p * 8, src, p * 8 / K >= Cn); + else + ptx::cp_async_16(dst + p * 8, src); + } + ptx::cp_async_commit(); + if (tid < BT) bnext = !VL || tid < Cn ? beta_in(beta[(t1 + tid) * H + hn]) : 0.f; + } + if constexpr (RAW) // one row-norm pass per chunk-job, q and k together + // (rides the gate cp.async; P2b/P2c re-read the same lines out of L1) + row_rnorm(rn2(tid >> 3), q, kk, (size_t)t0 * H * K + (size_t)h * K, H, C_act, tid); + __syncthreads(); + if constexpr (GM != 0) // tiles no prior tail pre-transformed: the + // piece-first chunk (and every chunk when the fused tail is absent + // — no-pieceL callers) + if (sub == 0 || !pieceL) { + gate_xform(sgb, C_act); + __syncthreads(); + } + { // split cumsum: 512 threads = 4 x 16-row segments per column + // (bf16 stage widens on read; the fp32 running sums land in sgc. + // GM != 0 tiles were transformed in place on landing, so every + // mode reads pre-made bf16 glog here) + const int col = tid & (K - 1), r0 = (tid >> 7) * (BT / 4); + float acc = 0.f; + for (int r = r0; r < r0 + BT / 4; ++r) { + acc += __bfloat162float(sgb[r][col]); + sg[r][col] = acc; + } + } + __syncthreads(); + { // per-thread column is fixed: the carry-in loads once per segment + // (the in-place += stores block the compiler from hoisting it) + const int colc = tid & (K - 1), rc = tid >> 7; + for (int seg = 1; seg < 4; ++seg) { + const float carry = sg[seg * (BT / 4) - 1][colc]; +#pragma unroll + for (int it = 0; it < 4; ++it) + sg[seg * (BT / 4) + rc + it * 4][colc] += carry; + __syncthreads(); + } + } + // anchors a(s) = gamma BEFORE sub-block s (0 for s=0). A thread's + // column PAIR is fixed across every P2b/P2c position, so the anchors, + // e^{a(s)} and the chunk-end row hoist to registers; e^{sg} then + // composes as e^{sg-a(s)}*e^{a(s)} (both factors already anchored + // quantities; exact for s==0, <=1 extra fp32 rounding else) + const int colp = (tid & (K / 2 - 1)) * 2; + float a0v[NSB], a1v[NSB], ea0[NSB], ea1[NSB]; + a0v[0] = a1v[0] = 0.f; + ea0[0] = ea1[0] = 1.f; +#pragma unroll + for (int s = 1; s < NSB; ++s) { + a0v[s] = sg[s * SB - 1][colp]; + a1v[s] = sg[s * SB - 1][colp + 1]; + ea0[s] = __expf(a0v[s]); + ea1[s] = __expf(a1v[s]); + } + // build kw, qw, kappa/rhs, qdec + { // kw/qw/qdec/kappa/v + the kdec value (into kz-area staging: the + // transposed [col][i] global store below keeps both sides coalesced) + // all of a half-batch's global loads issue FIRST, then compute+store + bf16(*stg)[K + 4] = reinterpret_cast(&kz[0][0][0]); + // column-PAIRED (bf162 loads/stores: half the transactions, twice + // the in-flight bytes per MSHR) + const float gl0 = sg[BT - 1][colp], gl1 = sg[BT - 1][colp + 1]; + // running pointers: a thread's taps advance a FIXED 8 rows, so one + // add per tensor replaces the per-tap 64-bit gp rebuild (the SASS + // census's top P2b term, ~9 int ops/load; same addresses bit-exact) + const size_t gp0 = (size_t)(t0 + tid * 2 / K) * H * K + h * K + tid * 2 % K; + const size_t gst = (size_t)8 * H * K; // +8 rows per tap + const bf16* pk = &kk[gp0]; + const bf16* pq = &q[gp0]; + const bf16* pv = &v[gp0]; + bf16* qd = &qdec[((size_t)c * H + h) * BT * K + (tid >> 6) * K + colp]; + // deferred-compose wait at the last union-independent point: the + // anchor/pointer setup above reads only sgc/registers; the kw + // store below is the chunk's first union (sp) write + pend_wait(); +#pragma unroll + for (int hb = 0; hb < 2; ++hb) { + float2 kv[4], qv[4], vv[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + if (VL && tid * 2 / K + (hb * 4 + j) * 8 >= C_act) { + kv[j] = qv[j] = vv[j] = float2{0.f, 0.f}; // pad rows + } else { + kv[j] = __bfloat1622float2(*reinterpret_cast(pk)); + qv[j] = __bfloat1622float2(*reinterpret_cast(pq)); + vv[j] = __bfloat1622float2(*reinterpret_cast(pv)); + if constexpr (RAW) { // l2norm, in fla's bf16 bytes + const float2 r = rn2((tid >> 6) + (hb * 4 + j) * 8); + kv[j] = l2_round(kv[j], r.x); + qv[j] = l2_round(qv[j], r.y); + } + } + pk += gst; + pq += gst; + pv += gst; + } +#pragma unroll + for (int j = 0; j < 4; ++j) { + const int i = (tid >> 6) + (hb * 4 + j) * 8; + const int si = hb * 2 + (j >> 1); // == i / SB (tid < 512) + const float ei0 = __expf(sg[i][colp] - a0v[si]); + const float ei1 = __expf(sg[i][colp + 1] - a1v[si]); + const float kw0 = kv[j].x * ei0, kw1 = kv[j].y * ei1; + *reinterpret_cast<__nv_bfloat162*>(&kw[i][colp]) = __floats2bfloat162_rn(kw0, kw1); + const float qw0 = qv[j].x * ei0 * scale; + const float qw1 = qv[j].y * ei1 * scale; + *reinterpret_cast<__nv_bfloat162*>(&qw[i][colp]) = __floats2bfloat162_rn(qw0, qw1); + const float bi = sb[i]; // beta folded into the solve rhs + *reinterpret_cast(&sU[i][colp]) = // kappa + float2{kw0 * (bi * ea0[si]), kw1 * (bi * ea1[si])}; + *reinterpret_cast(&sU[i][K + colp]) = float2{vv[j].x * bi, vv[j].y * bi}; + *reinterpret_cast<__nv_bfloat162*>(qd) = __floats2bfloat162_rn(qw0 * ea0[si], qw1 * ea1[si]); + qd += 8 * K; // i advances 8 rows per tap + *reinterpret_cast<__nv_bfloat162*>(&stg[i][colp]) = + __floats2bfloat162_rn(kv[j].x * __expf(gl0 - sg[i][colp]), kv[j].y * __expf(gl1 - sg[i][colp + 1])); + } + } + __syncthreads(); + // b64 quads: 4 adjacent tokens of one channel ([col][i] layout). + // compile-time trip (512 threads, like the tap loop above) => the + // per-iteration %/÷ and 64-bit address rebuilds become immediates + const int rr = tid * 4 % BT, cc = tid * 4 / BT; + bf16* kd = &kdec[((size_t)c * H + h) * BT * K + tid * 4]; +#pragma unroll + for (int n = 0; n < BT * K / 2048; ++n) { + const int cn = cc + n * 32; + alignas(8) const __nv_bfloat162 p2[2] = {{stg[rr][cn], stg[rr + 1][cn]}, {stg[rr + 2][cn], stg[rr + 3][cn]}}; + *reinterpret_cast(kd + n * 2048) = *reinterpret_cast(p2); + } + } + __syncthreads(); + // kz tiles: a column's pairs (si, sj) share their exponent down si — + // one anchored base exp e^{a(sj)-sg[j]} per (sj, row), then multiply + // by the adjacent-anchor gaps f(s) = e^{a(s)-a(s-1)} (the e^{x-a}*e^{a-y} + // composition in fp32; base tiles bit-exact, derived tiles carry <=3 + // extra fp32 roundings pre-bf16). The 8 dedup'd row loads issue first. + { + const int u = tid >> 6; // row-in-block base; cols = colp pair + float f0[NSB - 1], f1[NSB - 1]; +#pragma unroll + for (int s = 1; s < NSB; ++s) { + f0[s - 1] = __expf(a0v[s] - a0v[s - 1]); + f1[s - 1] = __expf(a1v[s] - a1v[s - 1]); + } + __nv_bfloat162 kvz[NSB][2]; + // running pointer: the 8 taps advance a fixed 8 rows (sj*SB + rr*8 + // = 0,8,..,56), so one 64-bit add per tap replaces the per-tap + // address rebuild (same addresses bit-exact; the P2b diet pattern) + const bf16* pz = &kk[(size_t)(t0 + u) * H * K + h * K + colp]; +#pragma unroll + for (int sj = 0; sj < NSB; ++sj) +#pragma unroll + for (int rr = 0; rr < 2; ++rr) { + const int j = sj * SB + u + rr * 8; + kvz[sj][rr] = + !VL || j < C_act ? *reinterpret_cast(pz) : __floats2bfloat162_rn(0.f, 0.f); + if constexpr (RAW) // l2norm at the load: the tiles below + // then see the same bytes P2b's kw did + kvz[sj][rr] = l2_bf16(__bfloat1622float2(kvz[sj][rr]), rn2(j).x); + pz += (size_t)8 * H * K; + } +#pragma unroll + for (int sj = 0; sj < NSB; ++sj) +#pragma unroll + for (int rr = 0; rr < 2; ++rr) { + const int j = sj * SB + u + rr * 8, jj = u + rr * 8; + const float2 kf = __bfloat1622float2(kvz[sj][rr]); + // pad rows select 0 OUTRIGHT: their base exponent a(sj) - + // sg[j] spans block start -> seq end (up to 63 rows of gate + // mass vs <= 15 for real rows) and 0 * __expf(overflow) = NaN + const bool real = !VL || j < C_act; + float v0 = real ? kf.x * __expf(a0v[sj] - sg[j][colp]) : 0.f; + float v1 = real ? kf.y * __expf(a1v[sj] - sg[j][colp + 1]) : 0.f; +#pragma unroll + for (int si = sj; si < NSB; ++si) { + if (si > sj) { + v0 *= f0[si - 1]; + v1 *= f1[si - 1]; + } + *reinterpret_cast<__nv_bfloat162*>(&kz[si * (si + 1) / 2 + sj][jj][colp]) = __floats2bfloat162_rn(v0, v1); + } + } + } + if (tid < K / 4) { // stored as e^{gC}: compile-time trip (512 + // threads) + one float4 per thread (was a runtime striding loop) + const float4 g4 = *reinterpret_cast(&sg[BT - 1][tid * 4]); + *reinterpret_cast(&gC[((size_t)c * H + h) * K + tid * 4]) = + float4{__expf(g4.x), __expf(g4.y), __expf(g4.z), __expf(g4.w)}; + } + __syncthreads(); + // warps 10-15 are idle through the pair mma: L2-prefetch the inputs of + // the CTA one wave ahead on this SM (linear bid + 148; same h while the + // c-range allows) — its P1/P2 loads then hit L2 instead of DRAM + if (warp >= 10) { + const int jn = job + 1; // the next persistent job + if (jn < nc1 * H && sub + 1 < njobs) { + const int t1p = (jn % nc1) * BT + tokoff; + const int Cp = VL ? min(BT, tse - t1p) : BT; // stop at seq end + const size_t tb = (size_t)t1p * H * K + (size_t)(jn / nc1) * K; + const char* pq = reinterpret_cast(q + tb); + const char* pk = reinterpret_cast(kk + tb); + const char* pv = reinterpret_cast(v + tb); + const char* pg = reinterpret_cast(glog + tb); + for (int i = tid - 320; i < Cp; i += 192) { + const size_t ro = (size_t)i * H * K * 2; + for (int l = 0; l < 256; l += 128) { + asm volatile("prefetch.global.L2 [%0];" ::"l"(pq + ro + l)); + asm volatile("prefetch.global.L2 [%0];" ::"l"(pk + ro + l)); + asm volatile("prefetch.global.L2 [%0];" ::"l"(pv + ro + l)); + asm volatile("prefetch.global.L2 [%0];" ::"l"(pg + ro + l)); + } + } + } + } + // A and Aqk via mma: each warp does pairs round-robin; per pair: + // [16x16] = kw_si[16x128] @ kz_pi^T (and qw for Aqk) + for (int pi = warp; pi < NSB * (NSB + 1) / 2; pi += 16) { + int si = 0; + while ((si + 1) * (si + 2) / 2 <= pi) + ++si; + const int sj = pi - si * (si + 1) / 2; + float4 acck[2], accq[2]; + acck[0] = acck[1] = accq[0] = accq[1] = float4{0, 0, 0, 0}; + const int arow = lane & 15, aka = lane >> 4; +#pragma unroll + for (int k16 = 0; k16 < K / 16; ++k16) { + uint32_t a0, a1, a2, a3, q0, q1, q2, q3; + ptx::ldmatrix_x4_b16(ptx::to_shared(&kw[si * SB + arow][k16 * 16 + aka * 8]), a0, a1, a2, a3); + ptx::ldmatrix_x4_b16(ptx::to_shared(&qw[si * SB + arow][k16 * 16 + aka * 8]), q0, q1, q2, q3); +#pragma unroll + for (int n8 = 0; n8 < 2; ++n8) { + uint32_t b0, b1; + ptx::ldmatrix_x2_b16(ptx::to_shared(&kz[pi][(lane & 7) + n8 * 8][k16 * 16 + ((lane >> 3) & 1) * 8]), b0, b1); + ptx::mma_m16n8k16_bf16f32(acck[n8], a0, a1, a2, a3, b0, b1); + ptx::mma_m16n8k16_bf16f32(accq[n8], q0, q1, q2, q3, b0, b1); + } + } + // scatter A fragments to sA (masked, beta-scaled); Aqk fragments go + // STRAIGHT to global hi/lo as bf162 pairs (host zero-fills once; a + // diagonal-straddling pair re-writes the host zero into its + // jj+1 > ii half — bit-exact vs never-written) + const int r = lane >> 2, c2 = (lane & 3) * 2; + const size_t abase = ((size_t)c * H + h) * BT * BT; +#pragma unroll + for (int n8 = 0; n8 < 2; ++n8) { + const float vals[4] = {acck[n8].x, acck[n8].y, acck[n8].z, acck[n8].w}; + const float valq[4] = {accq[n8].x, accq[n8].y, accq[n8].z, accq[n8].w}; +#pragma unroll + for (int e2 = 0; e2 < 2; ++e2) { // fragment row: cols jj, jj+1 + const int ii = si * SB + r + e2 * 8; + const int jj = sj * SB + n8 * 8 + c2; + sA[ii][jj] = (jj < ii) ? sb[ii] * vals[e2 * 2] : 0.f; + sA[ii][jj + 1] = (jj + 1 < ii) ? sb[ii] * vals[e2 * 2 + 1] : 0.f; + if (jj <= ii) { + const float q0 = valq[e2 * 2]; + const float q1 = jj + 1 <= ii ? valq[e2 * 2 + 1] : 0.f; + const __nv_bfloat162 ah{__float2bfloat16(q0), __float2bfloat16(q1)}; + *reinterpret_cast<__nv_bfloat162*>(&aqk_h[abase + ii * BT + jj]) = ah; + *reinterpret_cast<__nv_bfloat162*>(&aqk_l[abase + ii * BT + jj]) = + __floats2bfloat162_rn(q0 - __bfloat162float(ah.x), q1 - __bfloat162float(ah.y)); + } + } + } + } + __syncthreads(); + // (beta pre-folded into the rhs fills) + pend_wait(); // the sp.f pack below is the next union write + // A hi/lo operand copy (coupling uses cols < 48 only) — oracle vet: + // hi/lo tensor-core coupling errs 8.7e-9 at the final state + for (int p2 = tid; p2 < BT * 24; p2 += blockDim.x) { + const int p = p2 * 2; + const int i = p / 48, j = p % 48; + // pairs never produce sA above the block diagonal and the coupling + // mma never reads it (A cols k16 < b): pack zeros without the sA + // round-trip + const float2 av = j / SB > i / SB ? float2{0.f, 0.f} : *reinterpret_cast(&sA[i][j]); + const bf16 ah0 = __float2bfloat16(av.x); + const bf16 ah1 = __float2bfloat16(av.y); + *reinterpret_cast<__nv_bfloat162*>(&sp.f.aC_h[i][j]) = __nv_bfloat162{ah0, ah1}; + *reinterpret_cast<__nv_bfloat162*>(&sp.f.aC_l[i][j]) = + __floats2bfloat162_rn(av.x - __bfloat162float(ah0), av.y - __bfloat162float(ah1)); + } + __syncthreads(); + // Blocked forward solve: cross-block coupling as mma (3 hi/lo products, + // every warp 2 n8-tiles), triangular block fp32 thread-per-column; + // solved rows publish transposed hi/lo as the next coupling's B. +#pragma unroll + for (int b = 0; b < NSB; ++b) { + if (b) { + const int arow = lane & 15, aka = lane >> 4; + float4 acc[2]; + acc[0] = acc[1] = float4{0, 0, 0, 0}; + for (int pr = 0; pr < 3; ++pr) { + const bf16(*ta)[56] = pr == 1 ? sp.f.aC_l : sp.f.aC_h; + const bf16(*tb)[56] = pr == 2 ? sp.f.upT_l : sp.f.upT_h; + for (int k16 = 0; k16 < b; ++k16) { + uint32_t a0, a1, a2, a3; + ptx::ldmatrix_x4_b16(ptx::to_shared(&ta[b * SB + arow][k16 * 16 + aka * 8]), a0, a1, a2, a3); +#pragma unroll + for (int n8 = 0; n8 < 2; ++n8) { + uint32_t b0, b1; + ptx::ldmatrix_x2_b16( + ptx::to_shared(&tb[warp * 16 + n8 * 8 + (lane & 7)][k16 * 16 + ((lane >> 3) & 1) * 8]), b0, b1); + ptx::mma_m16n8k16_bf16f32(acc[n8], a0, a1, a2, a3, b0, b1); + } + } + } + const int fr = lane >> 2, fc = (lane & 3) * 2; +#pragma unroll + for (int n8 = 0; n8 < 2; ++n8) { + const float vals[4] = {acc[n8].x, acc[n8].y, acc[n8].z, acc[n8].w}; +#pragma unroll + for (int e = 0; e < 4; ++e) + sU[b * SB + fr + (e >> 1) * 8][warp * 16 + n8 * 8 + fc + (e & 1)] -= vals[e]; + } + } + __syncthreads(); + if (tid < 2 * K) { + const int col = tid; + float r[SB]; +#pragma unroll + for (int i = 0; i < SB; ++i) + r[i] = sU[b * SB + i][col]; +#pragma unroll + for (int i = 1; i < SB; ++i) +#pragma unroll + for (int j = 0; j < i; ++j) + r[i] -= sA[b * SB + i][b * SB + j] * r[j]; +#pragma unroll + for (int i = 0; i < SB; ++i) + sU[b * SB + i][col] = r[i]; + if (b + 1 < NSB) // next coupling's B operand (hi/lo^T): 8 rows + // pack into ONE 16B store per array (rows are + // 112B-strided, so b*SB+i lands 16B-aligned) +#pragma unroll + for (int i = 0; i < SB; i += 8) { + alignas(16) __nv_bfloat162 hv[4], lv[4]; +#pragma unroll + for (int w = 0; w < 4; ++w) { + const float e0 = r[i + 2 * w], e1 = r[i + 2 * w + 1]; + hv[w] = __nv_bfloat162{__float2bfloat16(e0), __float2bfloat16(e1)}; + lv[w] = __floats2bfloat162_rn(e0 - __bfloat162float(hv[w].x), e1 - __bfloat162float(hv[w].y)); + } + *reinterpret_cast(&sp.f.upT_h[col][b * SB + i]) = *reinterpret_cast(hv); + *reinterpret_cast(&sp.f.upT_l[col][b * SB + i]) = *reinterpret_cast(lv); + } + } + __syncthreads(); + } + const size_t base = ((size_t)c * H + h) * BT * K; + for (int p2 = tid; p2 < BT * K / 2; p2 += blockDim.x) { + const int p = p2 * 2; + *reinterpret_cast<__nv_bfloat162*>(&P[base + p]) = + __floats2bfloat162_rn(-sU[p / K][p % K], -sU[p / K][p % K + 1]); + *reinterpret_cast(&u0[base + p]) = float2{sU[p / K][K + p % K], sU[p / K][K + p % K + 1]}; + } + if (!pieceL) continue; + // ---- fused W-form products + IN-TAIL RUNNING COMPOSITION ---- + // Per chunk: W and w0^T products (w0's operands SWAPPED vs the drained + // form — A=u0T, B=kdT — so its product lands in state orientation), then + // the piece's running maps compose as a NO-DRAIN tf32 TMEM chain: + // Lrun/crun are fp32 D pairs ping-ponging between [0,256) and [256,512). + // The chunk's W/w0 products target the FREE pair; W alone is drained + // (L = diag(e^gC) - W, staged as fp32 SWZ32 tiles); then Lrun_new = + // Lrun @ L-tiles overwrites the drained W slot and crun_new = crun @ + // L-tiles accumulates onto the completed w0 product (the c fold). + // ONE commit+wait per chunk (the products'); the compose commit is + // waited a chunk LATE (pend_wait, before its B-tile smem is rewritten); + // piece-end chunks wait in-chunk and drain the final maps to global. + __syncthreads(); // everyone past the P/u0 reads of sU + // tail pTt: TMA the L2-hot negated-P copy this CTA stored above — the + // [tok][kc] global IS the W mma's MN-major B image (probe-verified: SWZ128 + // box == the swz128b staging bytes; 2 boxes = the 2 MN atoms at LBO + // 8192). Bytes are the former in-smem pack's sign-flipped, so the + // product lands -W; the L drain folds the sign exactly. The mb_f + // phase is waited by ALL threads before the W mma. + if (tid == 0) { + ptx::fence_async_global(); // P sts -> async proxy (TMA) + ptx::mbar_arrive_expect_tx(&mb_f, uint32_t(sizeof(sp.b.pTt))); + ptx::cp_async_bulk_tensor_2d_load(ptx::to_shared(&sp.b.pTt[0][0]), ptm, 0, (c * H + h) * BT, &mb_f); + ptx::cp_async_bulk_tensor_2d_load(ptx::to_shared(&sp.b.pTt[64][0]), ptm, 64, (c * H + h) * BT, &mb_f); + } + const size_t cbase = ((size_t)c * H + h) * BT * K; +#pragma unroll + for (int hb = 0; hb < 2; ++hb) { // kdec loads issue first (MLP) + bf16 kdv[8]; +#pragma unroll + for (int j = 0; j < 8; ++j) + kdv[j] = kdec[cbase + tid + (hb * 8 + j) * 512]; +#pragma unroll + for (int j = 0; j < 8; ++j) { + const int p = tid + (hb * 8 + j) * 512; + const int kc = p / BT, tok = p % BT; + // 128B-swizzle placement (16B atoms XOR row): swz128 desc form + const int tokS = (((tok >> 3) ^ (kc & 7)) << 3) | (tok & 7); + sp.b.kdT[kc][tokS] = kdv[j]; + const float uv = sU[tok][K + kc]; + const bf16 uh = __float2bfloat16(uv); + sp.b.u0Th[kc][tokS] = uh; + sp.b.u0Tl[kc][tokS] = __float2bfloat16(uv - __bfloat162float(uh)); + } + } + ptx::tcgen05_fence_before_thread_sync(); + __syncthreads(); + const uint32_t taddr = s_taddr; + const uint32_t idw = ptx::mma_inst_desc_dense( + 128, 128, ptx::F16Type::BF16, ptx::F16Type::BF16, ptx::DType::F32, ptx::Major::K, ptx::Major::K); + // TMA-fed pTt is the [tok][kc] P image: B is MN-major (idesc bit 16) + const uint32_t idwm = ptx::mma_inst_desc_dense( + 128, 128, ptx::F16Type::BF16, ptx::F16Type::BF16, ptx::DType::F32, ptx::Major::K, ptx::Major::MN); + const uint32_t tp = (sub & 1) ? 256u : 0u; // this chunk's pair + // pTt TMA landed. ALL threads consume the phase: skipping an open + // phase would alias the next parity wait against the previous completed + ptx::mbar_wait_parity(&mb_f, mbph & 1); + ++mbph; + if (tid == 0) { + ptx::tcgen05_fence_after_thread_sync(); + // W-D at tp+[0,128): -W[kc_out][kc_in] = kdecT @ P via MN-major B +#pragma unroll + for (int k16 = 0; k16 < BT / 16; ++k16) { + const uint64_t da = ptx::mma_smem_desc_k_major(ptx::to_shared(&sp.b.kdT[0][0]) + k16 * 32); + const uint64_t db = + ptx::mma_smem_desc_mn_major(ptx::to_shared(&sp.b.pTt[0][0]) + k16 * 2048); + ptx::tcgen05_mma_f16(taddr + tp, da, db, idwm, k16 ? 1u : 0u); + } + // w0^T-D at tp+[128,256): w0^T[vc][kc_out] = u0T @ kdT^T (hi + lo + // A) — this slot IS crun's seed / c-fold addend + for (int half = 0; half < 2; ++half) +#pragma unroll + for (int k16 = 0; k16 < BT / 16; ++k16) { + const uint64_t da = ptx::mma_smem_desc_k_major( + ptx::to_shared(half ? &sp.b.u0Tl[0][0] : &sp.b.u0Th[0][0]) + k16 * 32); + const uint64_t db = ptx::mma_smem_desc_k_major(ptx::to_shared(&sp.b.kdT[0][0]) + k16 * 32); + ptx::tcgen05_mma_f16(taddr + tp + 128u, da, db, idw, (half | k16) ? 1u : 0u); + } + ptx::tcgen05_commit_arrive(&mb_f); + } + if constexpr (GM != 0) // next chunk's RAW gates landed long ago (P1 + // cp.async): transform them here, under the W/w0 mma wait, off the + // serial cumsum path. Each thread rewrites exactly the bytes it + // staged, so its own-group wait suffices; the next chunk's P2a + // read is barriers away (P7/P1 syncthreads). + if (sub + 1 < njobs) { + ptx::cp_async_wait_pending(0); + const int t1x = ((job + 1) % nc1) * BT + tokoff; + gate_xform(S.sgb[(sub + 1) & 1], VL ? min(BT, tse - t1x) : BT); + } + ptx::mbar_wait_parity(&mb_f, mbph & 1); + ++mbph; + ptx::tcgen05_fence_after_thread_sync(); + const int band = (warp & 3) * 32, ch = warp >> 2; // 16 warps: 32-col ch + const uint32_t lane_hi = uint32_t(band) << 16; + // W drain -> L = diag(e^gC) - W, staged as fp32 SWZ32 K=8-chunk tiles + // (the tf32 compose's B; b spans exactly slice ch*4+b). SWZ32 row form: + // a row's two 16B halves swap on (out>>2)&1. + // The w0 product is NOT drained — it stays in TMEM as crun's addend. + for (int b = 0; b < 4; ++b) { + uint32_t r[8]; + ptx::tcgen05_ld_32x32b_x8( + taddr + lane_hi + (tp + uint32_t(ch * 32 + b * 8)), r[0], r[1], r[2], r[3], r[4], r[5], r[6], r[7]); + ptx::tcgen05_wait_ld(); + float lt[8]; +#pragma unroll + for (int j = 0; j < 8; ++j) { + const int kin = ch * 32 + b * 8 + j; + // the mma drained -W (sign-flipped B bytes): a + (-x) == a - x + // exactly (IEEE), so L is byte-identical to the former pack's + lt[j] = (band + lane == kin ? gC[((size_t)c * H + h) * K + kin] : 0.f) + __int_as_float(r[j]); + } + const bool sw = ((band + lane) >> 2) & 1; + const float4 lo = float4{lt[0], lt[1], lt[2], lt[3]}; + const float4 hi = float4{lt[4], lt[5], lt[6], lt[7]}; + float4* dst = reinterpret_cast(&sp.g.sLt[ch * 4 + b][band + lane][0]); + dst[0] = sw ? hi : lo; + dst[1] = sw ? lo : hi; + } + ptx::tcgen05_fence_before_thread_sync(); + __syncthreads(); // sLt complete for the compose's B descs + { + const int vc = band + lane; + auto drain_reg = [&](uint32_t dc, float* r2) { // fp32 D region -> regs +#pragma unroll + for (int b2 = 0; b2 < 4; ++b2) { + uint32_t r[8]; + ptx::tcgen05_ld_32x32b_x8( + taddr + lane_hi + dc + uint32_t(ch * 32 + b2 * 8), r[0], r[1], r[2], r[3], r[4], r[5], r[6], r[7]); + ptx::tcgen05_wait_ld(); +#pragma unroll + for (int j = 0; j < 8; ++j) + r2[b2 * 8 + j] = __int_as_float(r[j]); + } + }; + auto store_maps = [&](const float* lr2, const float* cr2) { + const size_t sb2 = (size_t)pidx * K * K; + // both [out][in] row-major. L bf16 (half its prefix TMA bytes; + // single-bf16 L is the gate-proven original convention): the prefix + // TMAs SWZ128 tiles and widens to the tf32 SWZ32 B layout + // (widen_map_tf32). c stays fp32 (SWZ32-TMA'd directly): bf16 c + // accumulated past the fla o bar at 16K (1.05e-3 vs 7.2e-4). +#pragma unroll + for (int j = 0; j < 32; ++j) { + pieceL[sb2 + (size_t)(ch * 32 + j) * K + vc] = __float2bfloat16(lr2[j]); + piecec[sb2 + (size_t)(ch * 32 + j) * K + vc] = cr2[j]; + } + }; + const bool piece_end = sub + 1 == njobs; + if (sub == 0) { + // seed: crun = the w0^T product (already sitting in tp+128, state- + // oriented); Lrun = L^T via a transposed (un-swizzling) tile re-read + alignas(16) float lr[32]; +#pragma unroll + for (int j = 0; j < 32; ++j) { + const int out = ch * 32 + j; + lr[j] = sp.g.sLt[vc >> 3][out][(vc & 7) ^ (((out >> 2) & 1) << 2)]; + } + if (piece_end) { // 1-chunk piece: the seed IS its map + alignas(16) float cr[32]; + drain_reg(tp + 128u, cr); + store_maps(lr, cr); + } else { + // st over the drained W slot — safe: the products' wait above + // retired every mma (st vs in-flight A-reads is UNordered) +#pragma unroll + for (int c4 = 0; c4 < 32; c4 += 4) + ptx::tcgen05_st_32x32b_x4( + taddr + lane_hi + (tp + uint32_t(ch * 32 + c4)), + __float_as_int(lr[c4]), + __float_as_int(lr[c4 + 1]), + __float_as_int(lr[c4 + 2]), + __float_as_int(lr[c4 + 3])); + ptx::tcgen05_wait_st(); + } + } else { + // no-drain tf32 compose: D(tp) = Lrun @ L-tiles (scale 0 overwrites + // the drained W), D(tp+128) = crun @ L-tiles + w0 (scale 1 onto the + // completed product). A = the other pair's fp32 D, read as tf32; + // no wait here — the commit is consumed by pend_wait / piece end. + if (tid == 0) { + const uint32_t idt = ptx::mma_inst_desc_tf32(128, K); + const uint32_t sp2 = tp ^ 256u; + for (int st2 = 0; st2 < 2; ++st2) +#pragma unroll + for (int s = 0; s < 16; ++s) + mma_tf32_atmem( + taddr + tp + uint32_t(st2 * 128), + taddr + (sp2 + uint32_t(st2 * 128 + 8 * s)), + ptx::mma_smem_desc_k_major(ptx::to_shared(&sp.g.sLt[s][0][0])), + idt, + (st2 | s) ? 1u : 0u); + ptx::tcgen05_commit_arrive(&mb_f); + } + if (piece_end) { // the piece's ONE remaining drain: final maps + ptx::mbar_wait_parity(&mb_f, mbph & 1); + ++mbph; + ptx::tcgen05_fence_after_thread_sync(); + alignas(16) float lr[32], cr[32]; + drain_reg(tp, lr); + drain_reg(tp + 128u, cr); + store_maps(lr, cr); + } else { + pend = true; + } + } + } + if (sub + 1 < njobs) __syncthreads(); // slot reuse across jobs + } // job loop + if (pieceL) { + if (warp == 0) ptx::tcgen05_dealloc(s_taddr, 512); + ptx::tcgen05_relinquish(); + } +} + +template +__global__ void __launch_bounds__(512) k1_factors_mma( + const bf16* __restrict__ q, + const bf16* __restrict__ kk, + const bf16* __restrict__ v, + const bf16* __restrict__ glog, + const float* __restrict__ a_log, + const float* __restrict__ dtb, + float lb, + const float* __restrict__ beta, + int T, + int H, + float scale, + bf16* __restrict__ P, + float* __restrict__ u0, + bf16* __restrict__ kdec, + bf16* __restrict__ qdec, + bf16* __restrict__ aqk_h, + bf16* __restrict__ aqk_l, + float* __restrict__ gC, + bf16* __restrict__ pieceL, + float* __restrict__ piecec, + int NP) { + __shared__ K1Smem S; + // grid = NP pieces per head x H heads; piece p covers len consecutive + // chunks of ONE head (the running composition needs these runs). + const int nc1 = T / BT, bid = blockIdx.x; + const int h = bid / NP, pc = bid % NP; + const int c00 = piece_c0(pc, nc1, NP); + const int njobs = piece_c0(pc + 1, nc1, NP) - c00; + k1_body( + S, + h * nc1 + c00, + njobs, + q, + kk, + v, + glog, + beta, + T, + H, + scale, + P, + u0, + kdec, + qdec, + aqk_h, + aqk_l, + gC, + pieceL, + piecec, + pc * H + h, + 0, + 0, + 0, + a_log, + dtb, + lb); +} +// ---------------- K2: tcgen05 chain (A-from-TMEM, B via TMA) -------------- +// Per-piece CTA walks its chunks serially. State S^T lives in registers — +// thread (band = (warp&3)*32, ch = warp>>2) owns S^T[vc = band+lane][kc in +// ch*32 + 0..31] — and is re-staged per chunk as a packed SINGLE-bf16 TMEM +// A region (precision study: the bf16 B operands, not the S read, set the error +// floor — hi/lo S buys <= 1.27x in the kernel ctx). Products (A-from-TMEM, +// B = TMA 128B-swizzled K-major tiles): +// P1: U^T[vc,tok] = S^T @ Pneg^T (M=128 N=64 K=128; +u0 in drain) +// P2: S^T[vc,kc] += U^T @ kdecT^T (M=128 N=128 K=64; decay in drain) +// o is fused: P3 (o^T += S^T@qdec^T) rides phase 1; P4 (o^T += u^T@Aqk^T) +// rides phase 2. TMEM cols: [0,64) P1-D | [64,192) P2-D | [192,224) A2hi | +// [224,256) A2lo | [256,320) A1 | [384,448) o-D. +// SELF-START: the piece's start state composes from h0 through the prefix +// pieces' (L, c) maps as a NO-DRAIN tf32 TMEM chain, gated on pflags. +template // varlen: pad-row masks compile-time-gated +__device__ static void chain_body( + ChainSmem& S, + int h, + int sg, + int nseg, + const CUtensorMap& pneg_map, + const CUtensorMap& kdt_map, + const CUtensorMap& qd_map, + const CUtensorMap& aqh_map, + const CUtensorMap& aql_map, + const CUtensorMap& u0f_map, + const float* __restrict__ gC, + int n_chunks, + int H, + bf16* __restrict__ o, + float* __restrict__ Sf, + // per-chunk states (nullptr = off): row = GLOBAL chunk c, [c][h][K][V] or + // [c][h][V][K] (hpc_v_first), fp32 or bf16 (hpc_bf16) + void* __restrict__ hpc, + bool hpc_bf16, + bool hpc_v_first, + const CUtensorMap* sl_map, // fp32 piece L maps + const CUtensorMap* sc_map, // fp32 piece offsets + const float* __restrict__ h0s, + const uint32_t* __restrict__ pflags = nullptr, // fused: piece-done + // varlen (defaults = eqlen): sg/nseg/n_chunks are SEQUENCE-local; cbase/ + // pbase are the sequence's global chunk/piece bases; o rows live at + // c*BT + tokoff + tok and rows >= tend are pad (masked, never stored) + int cbase = 0, + int pbase = 0, + int tokoff = 0, + int tend = 0) { + const int tid = threadIdx.x; + const int warp = tid >> 5, lane = tid & 31; + const int band = (warp & 3) * 32, ch = warp >> 2; // 16 warps: ch = quarter + const int vc = band + lane; + auto& sPneg = S.sPneg; + auto& sKdT = S.sKdT; + auto& sQd = S.sQd; + auto& sAh = S.sAh; + auto& sAl = S.sAl; + auto& sU0 = S.sU0; + auto& mb_tma = S.mb_tma; + auto& mb_p1 = S.mb_p1; + auto& mb_p2 = S.mb_p2; + auto& mb_o = S.mb_o; + auto& s_taddr = S.s_taddr; + if (tid == 0) { + ptx::mbar_init(&mb_tma, 1); + ptx::mbar_init(&mb_p1, 1); // spine: P1 only + ptx::mbar_init(&mb_p2, 1); // spine: P2 only + ptx::mbar_init(&mb_o, 2); // o-pipe: P3 + P4 + ptx::mbar_init(&S.mb_pre, 1); + ptx::mbar_init(&S.mb_preC, 1); + ptx::mbar_init(&S.mb_cmL, 1); + ptx::mbar_init(&S.mb_cmC, 1); + ptx::prefetch_tensormap(&pneg_map); + ptx::prefetch_tensormap(&kdt_map); + } + if (warp == 0) ptx::tcgen05_alloc(ptx::to_shared(&s_taddr), 512); + __syncthreads(); + const uint32_t taddr = s_taddr; + const uint32_t lane_hi = uint32_t(band) << 16; + // ascending piece boundaries (shared with the builders) + const int c0 = cbase + piece_c0(sg, n_chunks, nseg); + const int per_seg = cbase + piece_c0(sg + 1, n_chunks, nseg) - c0; + const int tse = VL ? tend : n_chunks * BT; // sequence end token + alignas(16) float Sreg[32]; + // SELF-START: compose this piece's start from h0 through pieces + // 0..sg-1 as a NO-DRAIN tf32 TMEM chain (the fp32 D region IS the + // next A operand — identity (m,k)->(lane,col) map, same-thread D->A + // ordered without commit+wait). Per piece, two K=8-sliced mma + // batches: D_tgt = D_src @ L^T then D_tgt += I @ c^T — the c add + // rides the mma pipe (a tcgen05.st into a region in-flight mmas + // read is UNordered). L is stored bf16 (half its TMA bytes): it + // lands as two SWZ128 tiles in the TOP half of its 64 KB fp32 slot + // and all threads widen it in place to the SWZ32 fp32 tiles the + // tf32 descs read (widen_map_tf32); c stays fp32 SWZ32 TMA (bf16 c + // failed the 16K o gate). Both slots overlay the chunk-loop tiles + // (exactly 128 KB); slot reuse is gated by per-batch commits WAITED + // one piece late (mb_cmL/mb_cmC), so TMA overlaps the in-flight + // mmas. ONE drain at the end. TMEM: Da [0,128) | Db [128,256) | + // I [256,384) — all reused by the chunk loop only after the final + // drain. +#pragma unroll + for (int j = 0; j < 32; ++j) + Sreg[j] = h0s[((size_t)h * K + (ch * 32 + j)) * K + vc]; + if (sg > 0) { + float* bL = reinterpret_cast(&sPneg[0][0][0][0]); + float* bC = bL + K * K; + bf16* pL = reinterpret_cast(bL + K * K / 2); + static_assert(sizeof(S.sPneg) + sizeof(S.sKdT) + sizeof(S.sQd) + sizeof(S.sAh) + sizeof(S.sAl) == 2 * K * K * 4); +#pragma unroll + for (int c4 = 0; c4 < 32; c4 += 4) { // I diag + h0 fp32, no pack + uint32_t iv[4], sv[4]; +#pragma unroll + for (int e = 0; e < 4; ++e) { + iv[e] = __float_as_int(vc == ch * 32 + c4 + e ? 1.f : 0.f); + sv[e] = __float_as_int(Sreg[c4 + e]); + } + ptx::tcgen05_st_32x32b_x4(taddr + lane_hi + uint32_t(256 + ch * 32 + c4), iv[0], iv[1], iv[2], iv[3]); + ptx::tcgen05_st_32x32b_x4(taddr + lane_hi + uint32_t(ch * 32 + c4), sv[0], sv[1], sv[2], sv[3]); + } + ptx::tcgen05_wait_st(); + ptx::tcgen05_fence_before_thread_sync(); + __syncthreads(); + const uint32_t idt = ptx::mma_inst_desc_tf32(128, K); + for (int q2 = 0; q2 < sg; ++q2) { + if (tid == 0) { + if (pflags) // fused grid: trail piece q2's producer + while (ptx::ld_acq_b32(&pflags[(pbase + q2) * H + h]) < 1u) + __nanosleep(128); + if (q2) ptx::mbar_wait_parity(&S.mb_cmL, (q2 - 1) & 1); + ptx::mbar_arrive_expect_tx(&S.mb_pre, K * K * 2); + ptx::cp_async_bulk_tensor_2d_load(ptx::to_shared(pL), sl_map, 0, ((pbase + q2) * H + h) * K, &S.mb_pre); + ptx::cp_async_bulk_tensor_2d_load( + ptx::to_shared(pL) + uint32_t(K * 64 * 2), sl_map, 64, ((pbase + q2) * H + h) * K, &S.mb_pre); + if (q2) ptx::mbar_wait_parity(&S.mb_cmC, (q2 - 1) & 1); + ptx::mbar_arrive_expect_tx(&S.mb_preC, K * K * 4); + for (int s = 0; s < 16; ++s) + ptx::cp_async_bulk_tensor_2d_load( + ptx::to_shared(bC) + uint32_t(s * 4096), sc_map, 8 * s, ((pbase + q2) * H + h) * K, &S.mb_preC); + } + // L landed: widen (its slot writes are safe — piece q2-1's mmas + // retired before tid0 issued this TMA, and the widen is ordered + // after the TMA through mb_pre) + ptx::mbar_wait_parity(&S.mb_pre, q2 & 1); + widen_map_tf32(bL, pL, vc, ch); + ptx::tcgen05_fence_before_thread_sync(); + __syncthreads(); + if (tid == 0) { + ptx::tcgen05_fence_after_thread_sync(); + const uint32_t src = (q2 & 1) ? 128u : 0u, tgt = 128u - src; + for (int s = 0; s < 16; ++s) + mma_tf32_atmem( + taddr + tgt, + taddr + src + uint32_t(8 * s), + ptx::mma_smem_desc_k_major(ptx::to_shared(bL) + uint32_t(s * 4096)), + idt, + s ? 1u : 0u); + ptx::tcgen05_commit_arrive(&S.mb_cmL); + ptx::mbar_wait_parity(&S.mb_preC, q2 & 1); // c: direct fp32 + for (int s = 0; s < 16; ++s) + mma_tf32_atmem( + taddr + tgt, + taddr + uint32_t(256 + 8 * s), + ptx::mma_smem_desc_k_major(ptx::to_shared(bC) + uint32_t(s * 4096)), + idt, + 1u); + ptx::tcgen05_commit_arrive(&S.mb_cmC); + } + } + if (tid == 0) // last commit + ptx::mbar_wait_parity(&S.mb_cmC, (sg - 1) & 1); + __syncthreads(); + ptx::tcgen05_fence_after_thread_sync(); + const uint32_t fin = ((sg - 1) & 1) ? 0u : 128u; // last tgt region +#pragma unroll + for (int b2 = 0; b2 < 4; ++b2) { + uint32_t r[8]; + ptx::tcgen05_ld_32x32b_x8( + taddr + lane_hi + fin + uint32_t(ch * 32 + b2 * 8), r[0], r[1], r[2], r[3], r[4], r[5], r[6], r[7]); + ptx::tcgen05_wait_ld(); +#pragma unroll + for (int j = 0; j < 8; ++j) + Sreg[b2 * 8 + j] = __int_as_float(r[j]); // c folded in-pipe + } + __syncthreads(); + } // sg > 0 + // packed cvt.bf16x2 (per-half rn == the scalar pair; low half = e0) + auto pack2 = [](float e0, float e1, uint32_t& hi, uint32_t& lo) { + const __nv_bfloat162 h2 = __floats2bfloat162_rn(e0, e1); + const __nv_bfloat162 l2 = __floats2bfloat162_rn(e0 - __bfloat162float(h2.x), e1 - __bfloat162float(h2.y)); + hi = *reinterpret_cast(&h2); + lo = *reinterpret_cast(&l2); + }; + if (pflags) { // fused grid: own piece's factors must be complete + if (tid == 0) + while (ptx::ld_acq_b32(&pflags[(pbase + sg) * H + h]) < 1u) + __nanosleep(128); + __syncthreads(); + } + auto tma_issue = [&](int cc) { + const int bu2 = (cc - c0) & 1; + ptx::mbar_arrive_expect_tx(&mb_tma, 6 * 64 * 64 * 2 + 128 * 64 * 2 + BT * K * 4); + ptx::cp_async_bulk_tensor_2d_load(ptx::to_shared(&sPneg[bu2][0][0][0]), &pneg_map, 0, (cc * H + h) * BT, &mb_tma); + ptx::cp_async_bulk_tensor_2d_load(ptx::to_shared(&sPneg[bu2][1][0][0]), &pneg_map, 64, (cc * H + h) * BT, &mb_tma); + ptx::cp_async_bulk_tensor_2d_load(ptx::to_shared(&sKdT[bu2][0][0]), &kdt_map, 0, (cc * H + h) * K, &mb_tma); + ptx::cp_async_bulk_tensor_2d_load(ptx::to_shared(&sQd[bu2][0][0][0]), &qd_map, 0, (cc * H + h) * BT, &mb_tma); + ptx::cp_async_bulk_tensor_2d_load(ptx::to_shared(&sQd[bu2][1][0][0]), &qd_map, 64, (cc * H + h) * BT, &mb_tma); + ptx::cp_async_bulk_tensor_2d_load(ptx::to_shared(&sAh[bu2][0][0]), &aqh_map, 0, (cc * H + h) * BT, &mb_tma); + ptx::cp_async_bulk_tensor_2d_load(ptx::to_shared(&sAl[bu2][0][0]), &aql_map, 0, (cc * H + h) * BT, &mb_tma); + ptx::cp_async_bulk_tensor_2d_load(ptx::to_shared(&sU0[bu2][0][0]), &u0f_map, 0, (cc * H + h) * BT, &mb_tma); + }; + if (tid == 0 && (!VL || per_seg > 0)) tma_issue(c0); + // running gC chunk pointer (chunks are H*K apart — same addresses, + // minus the per-chunk 64-bit rebuild; C4 census diet) + loop-invariant + // A1 tmem column + const float* gCp = gC + ((size_t)c0 * H + h) * K + ch * 32; + const uint32_t a1st = taddr + lane_hi + (256u + ch * 16); + for (int c = c0; c < c0 + per_seg; ++c) { + const int buf_idx = (c - c0) & 1; + // prefetch this chunk's gC early: its latency hides behind A1 + // staging + P1 (u0 rides the chunk TMA into sU0 instead) + alignas(16) float gv[32]; +#pragma unroll + for (int j = 0; j < 8; ++j) + reinterpret_cast(gv)[j] = reinterpret_cast(gCp)[j]; + gCp += (size_t)H * K; + // stage A1 = S^T packed single-bf16 (o's S-term reads it too, via + // P3); no lo pack/st — vet b8 (vet_tf32_step.py): S-read precision + // is not the error floor, worst kernel-ctx cell 1.27x of hi/lo + { + uint32_t hw[16]; +#pragma unroll + for (int w = 0; w < 16; ++w) { + const __nv_bfloat162 h2 = __floats2bfloat162_rn(Sreg[2 * w], Sreg[2 * w + 1]); + hw[w] = *reinterpret_cast(&h2); + } + ptx::tcgen05_st_32x32b_x16(a1st, hw); + } + // h[c]: the register state BEFORE this chunk (store-only; the last + // boundary is final_state's). Dense — every row written; a per-sequence + // snapshot-index filter would gate this store (sparse follow-up). + if (hpc) { + const size_t hb = ((size_t)c * H + h) * K * K; + if (hpc_v_first) { // [V,K]: kc contiguous, one run per thread + const size_t off = hb + (size_t)vc * K + ch * 32; + if (hpc_bf16) { + auto* p = reinterpret_cast<__nv_bfloat162*>(reinterpret_cast(hpc) + off); +#pragma unroll + for (int j = 0; j < 16; ++j) + p[j] = __floats2bfloat162_rn(Sreg[2 * j], Sreg[2 * j + 1]); + } else { + auto* p = reinterpret_cast(reinterpret_cast(hpc) + off); +#pragma unroll + for (int j = 0; j < 8; ++j) + p[j] = reinterpret_cast(Sreg)[j]; + } + } else { // [K,V]: vc contiguous, coalesced across the warp + const size_t off = hb + (size_t)(ch * 32) * K + vc; + if (hpc_bf16) { + bf16* p = reinterpret_cast(hpc) + off; +#pragma unroll + for (int j = 0; j < 32; ++j) + p[j * K] = __float2bfloat16(Sreg[j]); + } else { + float* p = reinterpret_cast(hpc) + off; +#pragma unroll + for (int j = 0; j < 32; ++j) + p[j * K] = Sreg[j]; + } + } + } + ptx::tcgen05_wait_st(); + ptx::tcgen05_fence_before_thread_sync(); + __syncthreads(); + // ALL threads gate on the chunk TMA: the sU0 reads in the U drain + // below need the cross-proxy visibility this mbarrier's completion + // gives its waiters (non-issuer threads previously just idled into + // the mb_p1 wait from here, so no overlap is lost) + ptx::mbar_wait_parity(&mb_tma, (c - c0) & 1); + if (tid == 0 || tid == 32) { // dual issuers, disjoint D regions + // phase closed: arm the next chunk into the other buffer + if (tid == 0 && c + 1 < c0 + per_seg) tma_issue(c + 1); + ptx::tcgen05_fence_after_thread_sync(); + const uint32_t idesc1 = ptx::mma_inst_desc_dense( + 128, 64, ptx::F16Type::BF16, ptx::F16Type::BF16, ptx::DType::F32, ptx::Major::K, ptx::Major::K); + if (tid == 0) { +#pragma unroll + for (int k16 = 0; k16 < 8; ++k16) { + const uint64_t db = ptx::mma_smem_desc_k_major( + ptx::to_shared(&sPneg[buf_idx][k16 >> 2][0][0]) + (k16 & 3) * 32); + ptx::tcgen05_mma_f16_atmem(taddr, taddr + (256u + k16 * 8), db, idesc1, k16 ? 1u : 0u); + } + ptx::tcgen05_commit_arrive(&mb_p1); + } else { +#pragma unroll + for (int k16 = 0; k16 < 8; ++k16) { // P3: o^T = S^T@qdec^T + const uint64_t db = ptx::mma_smem_desc_k_major( + ptx::to_shared(&sQd[buf_idx][k16 >> 2][0][0]) + (k16 & 3) * 32); + ptx::tcgen05_mma_f16_atmem(taddr + 384u, taddr + (256u + k16 * 8), db, idesc1, k16 ? 1u : 0u); + } + ptx::tcgen05_commit_arrive(&mb_o); + } + } + ptx::mbar_wait_parity(&mb_p1, (c - c0) & 1); // U-drain: P1 alone + ptx::tcgen05_fence_after_thread_sync(); + // drain U^T (+u0), pack A2 (one x16 ld: one TMEM + // round-trip for the 16 columns instead of two x8 trips) + { + alignas(16) float Ur[16]; + { + uint32_t r[16]; + ptx::tcgen05_ld_32x32b_x16( + taddr + lane_hi + uint32_t(ch * 16), + r[0], + r[1], + r[2], + r[3], + r[4], + r[5], + r[6], + r[7], + r[8], + r[9], + r[10], + r[11], + r[12], + r[13], + r[14], + r[15]); + ptx::tcgen05_wait_ld(); +#pragma unroll + for (int j = 0; j < 16; ++j) + Ur[j] = __int_as_float(r[j]) + sU0[buf_idx][ch * 16 + j][vc]; + } + uint32_t hw[8], lw[8]; +#pragma unroll + for (int w = 0; w < 8; ++w) + pack2(Ur[2 * w], Ur[2 * w + 1], hw[w], lw[w]); + ptx::tcgen05_st_32x32b_x8( + taddr + lane_hi + (192u + ch * 8), hw[0], hw[1], hw[2], hw[3], hw[4], hw[5], hw[6], hw[7]); + ptx::tcgen05_st_32x32b_x8( + taddr + lane_hi + (224u + ch * 8), lw[0], lw[1], lw[2], lw[3], lw[4], lw[5], lw[6], lw[7]); + } + ptx::tcgen05_wait_st(); + ptx::tcgen05_fence_before_thread_sync(); + __syncthreads(); + if (tid == 0 || tid == 32) { + ptx::tcgen05_fence_after_thread_sync(); + if (tid == 0) { + const uint32_t idesc2 = ptx::mma_inst_desc_dense( + 128, 128, ptx::F16Type::BF16, ptx::F16Type::BF16, ptx::DType::F32, ptx::Major::K, ptx::Major::K); + for (int half = 0; half < 2; ++half) +#pragma unroll + for (int k16 = 0; k16 < 4; ++k16) { + const uint64_t db = + ptx::mma_smem_desc_k_major(ptx::to_shared(&sKdT[buf_idx][0][0]) + k16 * 32); + ptx::tcgen05_mma_f16_atmem( + taddr + 64u, taddr + (192u + half * 32 + k16 * 8), db, idesc2, (half | k16) ? 1u : 0u); + } + ptx::tcgen05_commit_arrive(&mb_p2); + } else { + const uint32_t idesc4 = ptx::mma_inst_desc_dense( + 128, 64, ptx::F16Type::BF16, ptx::F16Type::BF16, ptx::DType::F32, ptx::Major::K, ptx::Major::K); + // P4: o^T += u^T@Aqk^T (u-hi x A-hi, u-hi x A-lo, u-lo x A-hi) + for (int pr = 0; pr < 3; ++pr) { + const uint32_t acol = pr == 2 ? 224u : 192u; + const bf16* ab = pr == 1 ? &sAl[buf_idx][0][0] : &sAh[buf_idx][0][0]; +#pragma unroll + for (int k16 = 0; k16 < 4; ++k16) { + const uint64_t db = ptx::mma_smem_desc_k_major(ptx::to_shared(ab) + k16 * 32); + ptx::tcgen05_mma_f16_atmem(taddr + 384u, taddr + (acol + k16 * 8), db, idesc4, 1u); + } + } + ptx::tcgen05_commit_arrive(&mb_o); + } + } + ptx::mbar_wait_parity(&mb_p2, (c - c0) & 1); // S-drain: P2 alone + ptx::tcgen05_fence_after_thread_sync(); + // drain P2: S = e^{gC} o S + tmem (gv prefetched at chunk top; + // one x32 ld = the widest single-instruction drain) + { + uint32_t r[32]; + ptx::tcgen05_ld_32x32b_x32(taddr + lane_hi + uint32_t(64 + ch * 32), r); + ptx::tcgen05_wait_ld(); +#pragma unroll + for (int jj = 0; jj < 32; ++jj) + Sreg[jj] = gv[jj] * Sreg[jj] + __int_as_float(r[jj]); + } + // o-pipe drain (P3+P4) rides behind the S update, off the spine + ptx::mbar_wait_parity(&mb_o, (c - c0) & 1); + ptx::tcgen05_fence_after_thread_sync(); + { + uint32_t r[16]; + ptx::tcgen05_ld_32x32b_x16( + taddr + lane_hi + uint32_t(384 + ch * 16), + r[0], + r[1], + r[2], + r[3], + r[4], + r[5], + r[6], + r[7], + r[8], + r[9], + r[10], + r[11], + r[12], + r[13], + r[14], + r[15]); + ptx::tcgen05_wait_ld(); + // running row pointer (rows are H*K apart — same addresses, + // minus the per-element 64-bit row*H*K rebuild) + const int row0 = c * BT + tokoff + ch * 16; + bf16* op = o + ((size_t)row0 * H + h) * K + vc; + const size_t orow = (size_t)H * K; +#pragma unroll + for (int j = 0; j < 16; ++j, op += orow) + if (!VL || row0 + j < tse) // pad rows: the NEXT seq's + *op = __float2bfloat16(__int_as_float(r[j])); + } + __syncthreads(); + } + if (sg == nseg - 1) +#pragma unroll + for (int j = 0; j < 32; ++j) + Sf[((size_t)h * K + ch * 32 + j) * K + vc] = Sreg[j]; + __syncthreads(); + if (warp == 0) ptx::tcgen05_dealloc(taddr, 512); + ptx::tcgen05_relinquish(); +} +// eqlen NP == 1 fallback: whole chain from h0, no piece maps/flags (the +// sl/sc maps are valid encodes over the workspace but never dereferenced) +__global__ void __launch_bounds__(512) k2_chain_tc( + const __grid_constant__ CUtensorMap pneg_map, + const __grid_constant__ CUtensorMap kdt_map, + const __grid_constant__ CUtensorMap qd_map, + const __grid_constant__ CUtensorMap aqh_map, + const __grid_constant__ CUtensorMap aql_map, + const __grid_constant__ CUtensorMap sl_map, + const __grid_constant__ CUtensorMap sc_map, + const __grid_constant__ CUtensorMap u0f_map, + const float* __restrict__ gC, + const float* __restrict__ h0, + int n_chunks, + int H, + bf16* __restrict__ o, + float* __restrict__ Sf, + void* __restrict__ hpc, + bool hpc_bf16, + bool hpc_v_first) { + __shared__ ChainSmem S; + chain_body( + S, + blockIdx.x, + blockIdx.y, + gridDim.y, + pneg_map, + kdt_map, + qd_map, + aqh_map, + aql_map, + u0f_map, + gC, + n_chunks, + H, + o, + Sf, + hpc, + hpc_bf16, + hpc_v_first, + &sl_map, + &sc_map, + h0); +} + +// ---------------- fused grid: chain blocks trail k1's pieces -------------- +// bids [0,NP*H) = k1 piece-builders (publish pflags after their stores); +// bids [NP*H, 2*NP*H) = self-start chain blocks gated on the piece flags. +// Chains become resident as builder waves retire and their spins trail the +// (mostly complete) flags — the launch boundary and part of the chain wall +// hide under k1's tail. +template +__global__ void __launch_bounds__(512) kda_fused( + const bf16* __restrict__ q, + const bf16* __restrict__ kk, + const bf16* __restrict__ v, + const bf16* __restrict__ glog, + const float* __restrict__ a_log, + const float* __restrict__ dtb, + float lb, + const float* __restrict__ beta, + int T, + int H, + float scale, + bf16* __restrict__ P, + float* __restrict__ u0, + bf16* __restrict__ kdec, + bf16* __restrict__ qdec, + bf16* __restrict__ aqk_h, + bf16* __restrict__ aqk_l, + float* __restrict__ gC, + bf16* __restrict__ pieceL, + float* __restrict__ piecec, + int NP, + const __grid_constant__ CUtensorMap pneg_map, + const __grid_constant__ CUtensorMap kdt_map, + const __grid_constant__ CUtensorMap qd_map, + const __grid_constant__ CUtensorMap aqh_map, + const __grid_constant__ CUtensorMap aql_map, + const __grid_constant__ CUtensorMap u0f_map, + const __grid_constant__ CUtensorMap sl_map, + const __grid_constant__ CUtensorMap sc_map, + const float* __restrict__ h0, + bf16* __restrict__ o, + float* __restrict__ Sf, + void* __restrict__ hpc, + bool hpc_bf16, + bool hpc_v_first, + uint32_t* __restrict__ pflags) { + union FusedSmem { + K1Smem k1; + ChainSmem chain; + }; + __shared__ FusedSmem S; + const int tid = threadIdx.x, bid = blockIdx.x; + const int nc1 = T / BT; + if (bid < NP * H) { + const int h = bid / NP, pc = bid % NP; + const int c00 = piece_c0(pc, nc1, NP); + const int njobs = piece_c0(pc + 1, nc1, NP) - c00; + k1_body( + S.k1, + h * nc1 + c00, + njobs, + q, + kk, + v, + glog, + beta, + T, + H, + scale, + P, + u0, + kdec, + qdec, + aqk_h, + aqk_l, + gC, + pieceL, + piecec, + pc * H + h, + 0, + 0, + 0, + a_log, + dtb, + lb, + &pneg_map); + __syncthreads(); + if (tid == 0) { + ptx::fence_async_global(); // consumers read via TMA + ptx::red_add_rel_b32(&pflags[pc * H + h], 1u); + } + } else { + const int r = bid - NP * H, h = r % H, sg = r / H; + chain_body( + S.chain, + h, + sg, + NP, + pneg_map, + kdt_map, + qd_map, + aqh_map, + aql_map, + u0f_map, + gC, + nc1, + H, + o, + Sf, + hpc, + hpc_bf16, + hpc_v_first, + &sl_map, + &sc_map, + h0, + pflags); + } +} + +// ---------------- varlen fused grid (cu_seqlens + partial tails) ---------- +// kda_fused generalized over a host-built PIECE TABLE: sequences pro-rate +// a GLOBAL piece budget by nc_s (seqs already parallelize the grid, a +// per-seq NP over-fills it), capped at the per-seq split +// (min(12, max(1, nc_s/2)) — no empty pieces) and chains compose prefixes +// WITHIN their sequence only. Factor tensors index by GLOBAL chunk +// (cbase_s + local c), maps/pflags by global piece id; +// q/k/v/glog/beta/o stay flat [T,H,K] and tail-chunk pad rows are +// zero-filled on load / masked on store inside the bodies (k1_body C_act +// guards, chain_body o mask). Only TAIL pieces (the one piece holding a +// partial last chunk) need those guards: full pieces dispatch to the exact +// eqlen bodies — t0 = c*BT + tokoff and the global-chunk factor indexing +// are shared, so with T = nc_tot*BT (nc1 == nc_tot) every address matches +// the VL body's. +struct VlPiece { + int seq; // sequence index (h0 / Sf row) + int cbase; // sequence's first global chunk + int nc; // sequence chunk count ceil(len/BT) + int np; // pieces in the sequence + int sg; // piece index within the sequence + int tokoff; // tok0 - cbase*BT: global token row = c*BT + tokoff + i + int tend; // sequence end token (pad mask) + int tail; // piece holds the partial last chunk (len%BT != 0) +}; +template +__global__ void __launch_bounds__(512) kda_fused_vl( + const bf16* __restrict__ q, + const bf16* __restrict__ kk, + const bf16* __restrict__ v, + const bf16* __restrict__ glog, + const float* __restrict__ a_log, + const float* __restrict__ dtb, + float lb, + const float* __restrict__ beta, + int T, + int H, + float scale, + bf16* __restrict__ P, + float* __restrict__ u0, + bf16* __restrict__ kdec, + bf16* __restrict__ qdec, + bf16* __restrict__ aqk_h, + bf16* __restrict__ aqk_l, + float* __restrict__ gC, + bf16* __restrict__ pieceL, + float* __restrict__ piecec, + const VlPiece* __restrict__ pieces, + int npt, + int nc_tot, + const __grid_constant__ CUtensorMap pneg_map, + const __grid_constant__ CUtensorMap kdt_map, + const __grid_constant__ CUtensorMap qd_map, + const __grid_constant__ CUtensorMap aqh_map, + const __grid_constant__ CUtensorMap aql_map, + const __grid_constant__ CUtensorMap u0f_map, + const __grid_constant__ CUtensorMap sl_map, + const __grid_constant__ CUtensorMap sc_map, + const float* __restrict__ h0, + bf16* __restrict__ o, + float* __restrict__ Sf, + void* __restrict__ hpc, + bool hpc_bf16, + bool hpc_v_first, + uint32_t* __restrict__ pflags) { + union FusedSmem { + K1Smem k1; + ChainSmem chain; + }; + __shared__ FusedSmem S; + const int tid = threadIdx.x, bid = blockIdx.x; + if (bid < npt * H) { + const int h = bid / npt, p = bid % npt; + const VlPiece* pc = pieces + p; + const int c00 = pc->cbase + piece_c0(pc->sg, pc->nc, pc->np); + const int njobs = pc->cbase + piece_c0(pc->sg + 1, pc->nc, pc->np) - c00; + if (pc->tail) + k1_body( + S.k1, + h * nc_tot + c00, + njobs, + q, + kk, + v, + glog, + beta, + T, + H, + scale, + P, + u0, + kdec, + qdec, + aqk_h, + aqk_l, + gC, + pieceL, + piecec, + p * H + h, + nc_tot, + pc->tokoff, + pc->tend, + a_log, + dtb, + lb, + &pneg_map); + else // full piece: the eqlen body on T = nc_tot*BT (see VlPiece) + k1_body( + S.k1, + h * nc_tot + c00, + njobs, + q, + kk, + v, + glog, + beta, + nc_tot * BT, + H, + scale, + P, + u0, + kdec, + qdec, + aqk_h, + aqk_l, + gC, + pieceL, + piecec, + p * H + h, + nc_tot, + pc->tokoff, + pc->tend, + a_log, + dtb, + lb, + &pneg_map); + __syncthreads(); + if (tid == 0) { + ptx::fence_async_global(); // consumers read via TMA + ptx::red_add_rel_b32(&pflags[p * H + h], 1u); + } + } else { + const int r = bid - npt * H, h = r % H, p = r / H; + const VlPiece* pc = pieces + p; + const size_t so = (size_t)pc->seq * H * K * K; // h0/Sf per sequence + if (pc->tail) // only the tail piece stores a partial o chunk + chain_body( + S.chain, + h, + pc->sg, + pc->np, + pneg_map, + kdt_map, + qd_map, + aqh_map, + aql_map, + u0f_map, + gC, + pc->nc, + H, + o, + Sf + so, + hpc, + hpc_bf16, + hpc_v_first, + &sl_map, + &sc_map, + h0 + so, + pflags, + pc->cbase, + p - pc->sg, + pc->tokoff, + pc->tend); + else + chain_body( + S.chain, + h, + pc->sg, + pc->np, + pneg_map, + kdt_map, + qd_map, + aqh_map, + aql_map, + u0f_map, + gC, + pc->nc, + H, + o, + Sf + so, + hpc, + hpc_bf16, + hpc_v_first, + &sl_map, + &sc_map, + h0 + so, + pflags, + pc->cbase, + p - pc->sg, + pc->tokoff, + pc->tend); + } +} + +// ---------------- seq0 route: tail-free builder + whole-seq chains -------- +// k1_body with the fused tail DELETED — no piece-map composition, no TMEM, no +// pTt TMA. Nothing composes the maps, so chains cannot self-start: they run +// ONE CTA per (sequence, head) walking that sequence from h0 (NP_chain == 1), +// and the two halves are two launches instead of one fused grid. What the tail +// buys back is smem — every dead range unions (the gate stage over the mma +// operands, the solve pane and packed A tiles over both) — so this builder is +// 102144 B / 64 regs = 2 CTAs/SM against the fused grid's 229376 B / 1. It +// wins wherever the whole-sequence chain depth is short next to the build wall +// (many sequences, or high H); pick_route decides per shape. +// The diet is unconditional here. Every item is factor-exact vs k1_body +// (identical fp32 op sequences, order-preserving remaps, deterministic __expf +// recompute): +// - gate stage single-buffered (own-chunk arrival, no next-chunk prefetch) +// with the GM transform folded into P2a's cumsum read +// - fp32 sA -> packed strict-lower diagonal blocks, and the off-diagonal A +// fragments are held in registers and scattered straight to the coupling +// operand (so aC may alias the mma operands) +// - ONE 128-col rhs pane, rebuilt per pane inside the solve (pane 0 kappa: +// the held k tap + cumsum recompute; pane 1 v * beta), so it lives in the +// operand union +// - sgc unions with upT (upT is born in the solve, after sgc's last read) +// - the 4 diagonal kz tiles stream through slots 0..3 and the 6 off-diagonal +// tiles then overwrite them (6 slots, not 10) +// P4-lo is baked in too: the solve's cross-block coupling consumes the hi +// tiles only (the double-bf16 lo corrections are dropped, not reordered) and +// no lo tile is produced. So is the block apply: the in-block substitution is +// X_b = M_bb @ rhs_b over 4 diagonal-only inverses (below) instead of a +// 15-deep fp32 chain. P and u0 therefore DIFFER from k1_body's bytes BY +// DESIGN — this route is envelope-gated (|ours - fp64| <= |fla - fp64|), +// never byte-gated, against the fused one. +// The LSU/shared data pipe is the binding resource on this builder (the only +// counter above 70%), and four items serve it. All four are layout or +// redundancy only — same products, same accumulation order, same bf16 rounds — +// so every factor byte is unchanged: +// - the coupling shuffle and the sU swizzle below (the block apply's own +// store/read conflict fix) +// - the gate tile arrives as ONE [BT][K] bf16 TMA box (SWIZZLE_NONE, so the +// landed tile is byte-for-byte the row-major staging per-thread cp.async +// produced). TMA writes shared through its own path, so both halves of the +// LDGSTS pair leave the pipe. q/k/v are NOT converted: every element is +// scaled by a decay on the way in, so they land in registers, and staging +// 3 x 16 KB for LDS readback costs the same wavefronts for smem that does +// not exist here. +// - the k tile is loaded ONCE: P2b's taps, the kz tiles and the pane-0 rhs +// rebuild are the SAME 8 bf16x2 values per thread (row u + 8m, column +// colp, one pad mask), so two of three global reads were pure redundancy +// - the kdec staging row map (kdec_r) and the 4-wide aqk store (below) +// sU physical column for logical (row, col): the 4-column group index rotates +// by 3 on ODD rows. The apply's X store wants a row-to-bank spread of 8 and +// its B-operand read wants 4, so no row stride serves both; this map makes +// both conflict-free at zero smem cost. Group-preserving, so every 4-aligned +// float2/float4 access stays vectorized (K == 128 == 32 groups). +static __device__ __forceinline__ int su_c(int r, int c) { + return ((((c >> 2) + 3 * (r & 1)) & 31) << 2) | (c & 3); +} +// kdec staging physical row for logical row r: the 6-bit row index rotated +// right by 2. The repack's 32 lanes each read ONE column down FOUR consecutive +// rows, so lanes step the row by 4 and the [BT][K+4] pitch (132 bf16 = 66 words +// == 2 banks) gives a lane step of 8 banks: 4 banks for 32 lanes. No pitch can +// fix that (4*pitch is always 0 mod 4 banks), so the ROW index has to move. +// Under the rotation a thread's four rows sit 16 apart (16 rows == 0 mod 32 +// banks) and consecutive lanes step ONE physical row == 2 banks: one wavefront. +// Both users step r by a multiple of 4 and stay inside one field, so the map is +// affine on each (kdec_r(r + 4n) == kdec_r(r) + kdec_r(4) * n) and costs no +// instruction and no smem — the staging is a dead-window kz-area alias. +static __device__ __forceinline__ constexpr int kdec_r(int r) { + static_assert(BT == 64, "the map rotates a 6-bit staging row index"); + return (r >> 2) + ((r & 3) << 4); +} +struct K1SmemTF { + // 128 B: the gate tile's TMA destination (16 B would satisfy the box) + union __align__(128) Pool { // disjoint live ranges + bf16 sgb[BT][K]; // P1 -> P2a + struct { // P2b -> P3 mma + bf16 kw[BT][K + 8]; + bf16 qw[BT][K + 8]; + bf16 kz[6][16][K + 8]; + } a; + struct { // P3 scatter -> P4/P5 + bf16 aC_h[BT][56]; + float sU[BT][K + 4]; // the one rhs pane + } f; + } sp; + union __align__(16) GPool { // sgc dies before upT is born + float sgc[BT][K]; + bf16 upT_h[K][56]; + } g; + float tri[4][120]; // strict-lower 16x16 diag blocks, packed i(i-1)/2+j + // (I+A_bb)^-1 hi/lo, row-major mma A operands. Own members, NOT the pool: + // written in P3 while the pool still holds the mma operands, read through + // both panes. Row stride 24 (48 B, 16B-aligned) spreads the 16 ldmatrix + // row addresses over all 8 bank segments. + alignas(16) bf16 Mh[4][16][24]; // ldmatrix needs the 16B row alignment + alignas(16) bf16 Ml[4][16][24]; + float sb[BT]; + uint64_t mb_in; // gate-tile TMA arrival (one phase per chunk-job) +}; +// RAW only: {k, q} row reciprocal norms. Nothing in the pool is dead across +// their P2b -> P4 pane-0 span (that is what the diet bought), so they cost +// 512 B on top of the 102144 B — 102656 B still holds 2 CTAs/SM — and the +// pre-normed instantiations keep the struct byte-for-byte. +struct K1SmemTFRaw : K1SmemTF { + float2 rn[BT]; +}; + +// VL: varlen tail piece (the k1_body idiom) — pad rows zero-fill on load so +// every factor row past the seq end is exactly 0; eqlen keeps the original +// code (the guards compile out). +// RAW: q/k un-normalized (see row_rnorm); the one k tap is normalized where it +// is loaded, so all three consumers see the same bytes. BSIG: beta as logits. +template +__device__ static void k1_tf_body( + K1SmemTF& S, + int job0, + int njobs, + const bf16* __restrict__ q, + const bf16* __restrict__ kk, + const bf16* __restrict__ v, + const bf16* __restrict__ glog, + const CUtensorMap& glog_map, + const float* __restrict__ beta, + int T, + int H, + float scale, + bf16* __restrict__ P, + float* __restrict__ u0, + bf16* __restrict__ kdec, + bf16* __restrict__ qdec, + bf16* __restrict__ aqk_h, + bf16* __restrict__ aqk_l, + float* __restrict__ gC, + // GM != 0 gate-transform inputs (production lb = -5.0) + const float* __restrict__ a_log = nullptr, + const float* __restrict__ dtb = nullptr, + float lb = 0.f, + // varlen piece coords (defaults = eqlen): global chunk c's tokens start at + // c*BT + tokoff; rows >= tend - t0 are pad (zero-filled on load) + int tokoff = 0, + int tend = 0, + float2* __restrict__ rn = nullptr) // RAW: K1SmemTFRaw::rn +{ + constexpr int SB = 16, NSB = BT / SB; + const int tid = threadIdx.x; + const int warp = tid >> 5, lane = tid & 31; + auto& kw = S.sp.a.kw; + auto& qw = S.sp.a.qw; + auto& kz = S.sp.a.kz; + auto& aC_h = S.sp.f.aC_h; + float (&sgc)[BT][K] = S.g.sgc; + auto& upT_h = S.g.upT_h; + auto& sU = S.sp.f.sU; // the pane, in the pool union + auto& tri = S.tri; + auto& Mh = S.Mh; + auto& Ml = S.Ml; + auto& sb = S.sb; + const int nc1 = T / BT; + const int tse = VL ? tend : T; // sequence end token + // GM hoist: h is constant across a block's jobs (one piece, one head) + const float ga = GM != 0 ? expf(a_log[job0 / nc1]) : 0.f; + // P2a-fold hoist: the cumsum column is thread-fixed (tid & (K-1)) and h is + // block-constant -> ONE dt_bias scalar covers every job + const float dtv_xf = GM != 0 ? dtb[(size_t)(job0 / nc1) * K + (tid & (K - 1))] : 0.f; + if (tid == 0) { + ptx::mbar_init(&S.mb_in, 1); + ptx::prefetch_tensormap(&glog_map); + } + __syncthreads(); + for (int sub = 0; sub < njobs; ++sub) { + const int job = job0 + sub; + const int c = job % nc1, h = job / nc1, t0 = c * BT + tokoff; + const int C_act = VL ? min(BT, tse - t0) : BT; // real rows (tail < BT) + bf16(&sgb)[BT][K] = S.sp.sgb; + if (sub) __syncthreads(); // pool reuse: prior job's P4/P5 reads done + if (tid == 0) { // ONE box: [BT][K] bf16, row-major (SWIZZLE_NONE) + ptx::mbar_arrive_expect_tx(&S.mb_in, BT * K * 2); + ptx::cp_async_bulk_tensor_2d_load(ptx::to_shared(&sgb[0][0]), &glog_map, h * K, t0, &S.mb_in); + } + if (tid < BT) sb[tid] = !VL || tid < C_act ? beta_in(beta[(size_t)(t0 + tid) * H + h]) : 0.f; + if constexpr (RAW) // one row-norm pass per chunk-job, q and k together + // (rides the gate box; the reads below hit the same L1 lines) + row_rnorm(rn[tid >> 3], q, kk, (size_t)t0 * H * K + (size_t)h * K, H, C_act, tid); + ptx::mbar_wait_parity(&S.mb_in, sub & 1); // all threads consume the phase + if constexpr (VL) { + // rows past the SEQUENCE end are in-bounds for the box (they are the + // next sequence's tokens), so the zero-fill the cp.async did by + // predicate happens here; rows past the buffer end arrived as zeros + for (int p = tid; p < (BT - C_act) * K / 8; p += blockDim.x) + *reinterpret_cast(&sgb[C_act + p * 8 / K][p * 8 % K]) = uint4{0, 0, 0, 0}; + } + __syncthreads(); + { // P2a: split cumsum, 512 threads = 4 x 16-row segments per column + const int col = tid & (K - 1), r0 = (tid >> 7) * (BT / 4); + float acc = 0.f; + for (int r = r0; r < r0 + BT / 4; ++r) { + // the GM transform folds into the read: identical per-element math + // INCLUDING the bf16 round k1_body stores through its gate stage + float gv; + if constexpr (GM != 0) { + const float g = __bfloat162float(sgb[r][col]) + dtv_xf; + const float y = + GM == 1 ? -ga * (fmaxf(g, 0.f) + __logf(1.f + __expf(-fabsf(g)))) : lb * (1.f / (1.f + __expf(-ga * g))); + gv = __bfloat162float(__float2bfloat16(y)); + } else { + gv = __bfloat162float(sgb[r][col]); + } + if (VL && r >= C_act) gv = 0.f; // pad: transform(0) != 0 + acc += gv; + sgc[r][col] = acc; + } + } + __syncthreads(); + { // carry propagation + const int colc = (tid & (K / 2 - 1)) * 2, rc = tid >> 6; + for (int seg = 1; seg < 4; ++seg) { + const float2 carry = *reinterpret_cast(&sgc[seg * (BT / 4) - 1][colc]); +#pragma unroll + for (int it = 0; it < 2; ++it) { + float2& v2 = *reinterpret_cast(&sgc[seg * (BT / 4) + rc + it * 8][colc]); + v2 = float2{v2.x + carry.x, v2.y + carry.y}; + } + __syncthreads(); + } + } + const int colp = (tid & (K / 2 - 1)) * 2; + float a0v[NSB], a1v[NSB]; // pair-map anchors (the kz chains keep them) + float ea0[NSB], ea1[NSB]; + ea0[0] = ea1[0] = 1.f; + a0v[0] = a1v[0] = 0.f; +#pragma unroll + for (int s = 1; s < NSB; ++s) { + a0v[s] = sgc[s * SB - 1][colp]; + a1v[s] = sgc[s * SB - 1][colp + 1]; + ea0[s] = __expf(a0v[s]); + ea1[s] = __expf(a1v[s]); + } + // the ONE k tile: tap m == hb*4+j == 2*sj+rr is row u + 8m, column colp, + // under one pad mask — so P2b, P2c and the pane-0 rhs all read this + __nv_bfloat162 kvz[NSB][2]; + { // P2b: kw/qw/qdec/kdec (the rhs is built later, in the solve) + bf16(*stg)[K + 4] = reinterpret_cast(&kz[0][0][0]); + const float gl0 = sgc[BT - 1][colp], gl1 = sgc[BT - 1][colp + 1]; + const size_t gp0 = (size_t)(t0 + tid * 2 / K) * H * K + h * K + tid * 2 % K; + const size_t gst = (size_t)8 * H * K; // +8 rows per tap + const bf16* pk = &kk[gp0]; + const bf16* pq = &q[gp0]; + bf16* qd = &qdec[((size_t)c * H + h) * BT * K + (tid >> 6) * K + colp]; + // the staging row map on the tap layout (row == (tid>>6) + 8*tap): it + // stays affine, so each tap keeps its immediate offset — mapping the + // materialized row instead costs 3 int ops per tap + const int sr0 = kdec_r(tid >> 6), srt = kdec_r(8); +#pragma unroll + for (int hb = 0; hb < 2; ++hb) { + float2 kv[4], qv[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const int m = hb * 4 + j; + if (VL && tid * 2 / K + m * 8 >= C_act) { + kv[j] = qv[j] = float2{0.f, 0.f}; // pad rows + kvz[m >> 1][m & 1] = __floats2bfloat162_rn(0.f, 0.f); + } else { + kvz[m >> 1][m & 1] = *reinterpret_cast(pk); + qv[j] = __bfloat1622float2(*reinterpret_cast(pq)); + if constexpr (RAW) { // l2norm, in fla's bf16 bytes + const float2 r = rn[(tid >> 6) + m * 8]; + kvz[m >> 1][m & 1] = l2_bf16(__bfloat1622float2(kvz[m >> 1][m & 1]), r.x); + qv[j] = l2_round(qv[j], r.y); + } + kv[j] = __bfloat1622float2(kvz[m >> 1][m & 1]); + } + pk += gst; + pq += gst; + } +#pragma unroll + for (int j = 0; j < 4; ++j) { + const int i = (tid >> 6) + (hb * 4 + j) * 8; + const int si = hb * 2 + (j >> 1); // == i / SB (tid < 512) + const float ei0 = __expf(sgc[i][colp] - a0v[si]); + const float ei1 = __expf(sgc[i][colp + 1] - a1v[si]); + const float kw0 = kv[j].x * ei0, kw1 = kv[j].y * ei1; + *reinterpret_cast<__nv_bfloat162*>(&kw[i][colp]) = __floats2bfloat162_rn(kw0, kw1); + const float qw0 = qv[j].x * ei0 * scale; + const float qw1 = qv[j].y * ei1 * scale; + *reinterpret_cast<__nv_bfloat162*>(&qw[i][colp]) = __floats2bfloat162_rn(qw0, qw1); + *reinterpret_cast<__nv_bfloat162*>(qd) = __floats2bfloat162_rn(qw0 * ea0[si], qw1 * ea1[si]); + qd += 8 * K; // i advances 8 rows per tap + *reinterpret_cast<__nv_bfloat162*>(&stg[sr0 + srt * (hb * 4 + j)][colp]) = + __floats2bfloat162_rn(kv[j].x * __expf(gl0 - sgc[i][colp]), kv[j].y * __expf(gl1 - sgc[i][colp + 1])); + } + } + __syncthreads(); + // kdec store from the kz-area staging + const int rr = tid * 4 % BT, cc = tid * 4 / BT; + bf16* kd = &kdec[((size_t)c * H + h) * BT * K + tid * 4]; +#pragma unroll + for (int n = 0; n < BT * K / 2048; ++n) { + const int cn = cc + n * 32; + alignas(8) const __nv_bfloat162 p2[2] = { + {stg[kdec_r(rr)][cn], stg[kdec_r(rr + 1)][cn]}, {stg[kdec_r(rr + 2)][cn], stg[kdec_r(rr + 3)][cn]}}; + *reinterpret_cast(kd + n * 2048) = *reinterpret_cast(p2); + } + } + __syncthreads(); + // P2c: kz pair tiles + gC + const int u = tid >> 6; + float f0[NSB - 1], f1[NSB - 1]; +#pragma unroll + for (int s = 1; s < NSB; ++s) { + f0[s - 1] = __expf(a0v[s] - a0v[s - 1]); + f1[s - 1] = __expf(a1v[s] - a1v[s - 1]); + } + { // phase A: the 4 diagonal (base) tiles into slots 0..3 +#pragma unroll + for (int sj = 0; sj < NSB; ++sj) +#pragma unroll + for (int rr = 0; rr < 2; ++rr) { + const int j = sj * SB + u + rr * 8, jj = u + rr * 8; + const float2 kf = __bfloat1622float2(kvz[sj][rr]); + // pad rows select 0 OUTRIGHT (the k1_body law): their base + // exponent spans to the seq end and 0 * __expf(ovfl) = NaN + const bool real = !VL || j < C_act; + const float v0 = real ? kf.x * __expf(a0v[sj] - sgc[j][colp]) : 0.f; + const float v1 = real ? kf.y * __expf(a1v[sj] - sgc[j][colp + 1]) : 0.f; + *reinterpret_cast<__nv_bfloat162*>(&kz[sj][jj][colp]) = __floats2bfloat162_rn(v0, v1); + } + } + if (tid < K / 4) { // gC = e^{cumsum at chunk end} + const float4 g4 = *reinterpret_cast(&sgc[BT - 1][tid * 4]); + *reinterpret_cast(&gC[((size_t)c * H + h) * K + tid * 4]) = + float4{__expf(g4.x), __expf(g4.y), __expf(g4.z), __expf(g4.w)}; + } + __syncthreads(); + if (warp >= 10) { // L2-prefetch the next job's inputs + const int jn = job + 1; + if (jn < nc1 * H && sub + 1 < njobs) { + const int t1p = (jn % nc1) * BT + tokoff; + const int Cp = VL ? min(BT, tse - t1p) : BT; // stop at seq end + const size_t tb = (size_t)t1p * H * K + (size_t)(jn / nc1) * K; + const char* pq = reinterpret_cast(q + tb); + const char* pk = reinterpret_cast(kk + tb); + const char* pv = reinterpret_cast(v + tb); + const char* pg = reinterpret_cast(glog + tb); + for (int i = tid - 320; i < Cp; i += 192) { + const size_t ro = (size_t)i * H * K * 2; + for (int l = 0; l < 256; l += 128) { + asm volatile("prefetch.global.L2 [%0];" ::"l"(pq + ro + l)); + asm volatile("prefetch.global.L2 [%0];" ::"l"(pk + ro + l)); + asm volatile("prefetch.global.L2 [%0];" ::"l"(pv + ro + l)); + asm volatile("prefetch.global.L2 [%0];" ::"l"(pg + ro + l)); + } + } + } + } + // P3: A and Aqk via mma — per pair: [16x16] = kw_si @ kz_pi^T (qw for + // Aqk), one pair per warp. Diagonal pairs run first (slots 0..3), then the + // off-diagonal tiles overwrite those slots for the second phase. + auto pair_mma = [&](int si2, int kzslot, float4* acck, float4* accq) { + acck[0] = acck[1] = accq[0] = accq[1] = float4{0, 0, 0, 0}; + const int arow = lane & 15, aka = lane >> 4; +#pragma unroll + for (int k16 = 0; k16 < K / 16; ++k16) { + uint32_t a0, a1, a2, a3, q0, q1, q2, q3; + ptx::ldmatrix_x4_b16(ptx::to_shared(&kw[si2 * SB + arow][k16 * 16 + aka * 8]), a0, a1, a2, a3); + ptx::ldmatrix_x4_b16(ptx::to_shared(&qw[si2 * SB + arow][k16 * 16 + aka * 8]), q0, q1, q2, q3); +#pragma unroll + for (int n8 = 0; n8 < 2; ++n8) { + uint32_t b0, b1; + ptx::ldmatrix_x2_b16( + ptx::to_shared(&kz[kzslot][(lane & 7) + n8 * 8][k16 * 16 + ((lane >> 3) & 1) * 8]), b0, b1); + ptx::mma_m16n8k16_bf16f32(acck[n8], a0, a1, a2, a3, b0, b1); + ptx::mma_m16n8k16_bf16f32(accq[n8], q0, q1, q2, q3, b0, b1); + } + } + }; + // Aqk fragments go STRAIGHT to global hi/lo (host zero-fills once; a + // diagonal-straddling group re-writes the host zero past its diagonal — + // bit-exact vs never-written). + // The mma C fragment gives a lane ONE column pair, so a warp's 4 B stores + // touch 8 rows for 128 B of payload. Lanes l and l^1 hold the SAME rows, so + // one xor-1 exchange gives each lane four CONSECUTIVE columns: even lanes + // [own n8=0 | partner n8=0], odd lanes [partner n8=1 | own n8=1]. Same + // values at the same positions, half the store instructions and half the + // wavefronts; 32 B/row is the floor, since a warp owns only 16 of a row's + // 64 columns. + auto aqk_scatter = [&](int si2, int sj2, const float4* accq) { + const int r = lane >> 2, c2 = (lane & 3) * 2; + const size_t abase = ((size_t)c * H + h) * BT * BT; + const int odd = lane & 1, jj = sj2 * SB + c2 + 6 * odd; +#pragma unroll + for (int e2 = 0; e2 < 2; ++e2) { + const int ii = si2 * SB + r + e2 * 8, lim = ii - jj; + const float m0 = e2 ? accq[0].z : accq[0].x; + const float m1 = e2 ? accq[0].w : accq[0].y; + const float p0 = e2 ? accq[1].z : accq[1].x; + const float p1 = e2 ? accq[1].w : accq[1].y; + // every lane must reach the exchange, so it precedes the mask + const float s0 = __shfl_xor_sync(0xffffffffu, odd ? m0 : p0, 1); + const float s1 = __shfl_xor_sync(0xffffffffu, odd ? m1 : p1, 1); + float v[4] = {odd ? s0 : m0, odd ? s1 : m1, odd ? p0 : s0, odd ? p1 : s1}; + if (lim >= 0) { +#pragma unroll + for (int e = 1; e < 4; ++e) + if (lim < e) v[e] = 0.f; + const __nv_bfloat162 ah[2] = { + {__float2bfloat16(v[0]), __float2bfloat16(v[1])}, {__float2bfloat16(v[2]), __float2bfloat16(v[3])}}; + *reinterpret_cast(&aqk_h[abase + ii * BT + jj]) = *reinterpret_cast(ah); + const __nv_bfloat162 al[2] = { + __floats2bfloat162_rn(v[0] - __bfloat162float(ah[0].x), v[1] - __bfloat162float(ah[0].y)), + __floats2bfloat162_rn(v[2] - __bfloat162float(ah[1].x), v[3] - __bfloat162float(ah[1].y))}; + *reinterpret_cast(&aqk_l[abase + ii * BT + jj]) = *reinterpret_cast(al); + } + } + }; + auto tri_scatter = [&](int si2, const float4* acck) { // si2 == sj2 + const int r = lane >> 2, c2 = (lane & 3) * 2; +#pragma unroll + for (int n8 = 0; n8 < 2; ++n8) { + const float vals[4] = {acck[n8].x, acck[n8].y, acck[n8].z, acck[n8].w}; +#pragma unroll + for (int e2 = 0; e2 < 2; ++e2) { + const int il = r + e2 * 8, jl = n8 * 8 + c2; + const int ii = si2 * SB + il, tb = il * (il - 1) / 2; + if (jl < il) tri[si2][tb + jl] = sb[ii] * vals[e2 * 2]; + if (jl + 1 < il) tri[si2][tb + jl + 1] = sb[ii] * vals[e2 * 2 + 1]; + } + } + }; + float4 offk[2]; + int offsi = -1, offsj = -1; // held off-diag A fragments + if (warp < NSB) { // diag pairs: si == sj == warp, slot = warp + float4 acck[2], accq[2]; + pair_mma(warp, warp, acck, accq); + aqk_scatter(warp, warp, accq); + tri_scatter(warp, acck); + } + __syncthreads(); + { // phase B production: the 6 off-diag tiles (base recompute, bit-same) +#pragma unroll + for (int sj = 0; sj < NSB - 1; ++sj) +#pragma unroll + for (int rr = 0; rr < 2; ++rr) { + const int j = sj * SB + u + rr * 8, jj = u + rr * 8; + const float2 kf = __bfloat1622float2(kvz[sj][rr]); + const bool real = !VL || j < C_act; // 0*__expf(ovfl) = NaN + float v0 = real ? kf.x * __expf(a0v[sj] - sgc[j][colp]) : 0.f; + float v1 = real ? kf.y * __expf(a1v[sj] - sgc[j][colp + 1]) : 0.f; +#pragma unroll + for (int si = sj + 1; si < NSB; ++si) { + v0 *= f0[si - 1]; + v1 *= f1[si - 1]; + *reinterpret_cast<__nv_bfloat162*>(&kz[si * (si - 1) / 2 + sj][jj][colp]) = __floats2bfloat162_rn(v0, v1); + } + } + } + __syncthreads(); + // The solve's only serial pole, paid ONCE per chunk-job instead of per + // pane per block row: M_bb = (I+L_bb)^-1 by fp32 forward substitution, one + // thread per column of I, 4 blocks in parallel. tri is final at the + // barrier above and the off-diag-pair mma below occupies only 6 warps, so + // 15-deep chain rides two warps that phase leaves idle. + constexpr int MBW = NSB * (NSB - 1) / 2; // first idle pair-phase warp + if (tid >= MBW * 32 && tid < MBW * 32 + NSB * SB) { + const int td = tid - MBW * 32; + const int b = td >> 4, j = td & 15; + float m[SB]; +#pragma unroll + for (int i = 0; i < SB; ++i) { + float s = i == j ? 1.f : 0.f; + // m[k] == 0 for k < j, so those FMAs are exact no-ops +#pragma unroll + for (int k2 = 0; k2 < i; ++k2) + s -= tri[b][i * (i - 1) / 2 + k2] * m[k2]; + m[i] = s; + } +#pragma unroll + for (int i = 0; i < SB; ++i) { // column j down the A operand + const bf16 hi = __float2bfloat16(m[i]); + Mh[b][i][j] = hi; + Ml[b][i][j] = __float2bfloat16(m[i] - __bfloat162float(hi)); + } + } + if (warp < NSB * (NSB - 1) / 2) { // off-diag pairs, slot = warp + int si2 = 1; + while (si2 * (si2 + 1) / 2 <= warp) + ++si2; + const int sj2 = warp - si2 * (si2 - 1) / 2; + float4 acck[2], accq[2]; + pair_mma(si2, warp, acck, accq); + aqk_scatter(si2, sj2, accq); + offk[0] = acck[0]; + offk[1] = acck[1]; + offsi = si2; + offsj = sj2; + } + __syncthreads(); + // aC operand pack (the coupling uses cols < 48 only) + if (offsi >= 0) { // off-diag fragments -> aC direct (jj < ii always) + const int r = lane >> 2, c2 = (lane & 3) * 2; +#pragma unroll + for (int n8 = 0; n8 < 2; ++n8) { + const float vals[4] = {offk[n8].x, offk[n8].y, offk[n8].z, offk[n8].w}; +#pragma unroll + for (int e2 = 0; e2 < 2; ++e2) { + const int ii = offsi * SB + r + e2 * 8; + const int jj = offsj * SB + n8 * 8 + c2; + *reinterpret_cast<__nv_bfloat162*>(&aC_h[ii][jj]) = + __floats2bfloat162_rn(sb[ii] * vals[e2 * 2], sb[ii] * vals[e2 * 2 + 1]); + } + } + } + for (int p2 = tid; p2 < BT * 24; p2 += blockDim.x) { // diag/upper cells + const int p = p2 * 2; + const int i = p / 48, j = p % 48; + if (j / SB < i / SB) continue; // off-diag lower: fragment-scattered + float2 av{0.f, 0.f}; + if (j / SB == i / SB) { + const int il = i % SB, jl = j % SB, tb = il * (il - 1) / 2; + if (jl < il) av.x = tri[i / SB][tb + jl]; + if (jl + 1 < il) av.y = tri[i / SB][tb + jl + 1]; + } + *reinterpret_cast<__nv_bfloat162*>(&aC_h[i][j]) = __floats2bfloat162_rn(av.x, av.y); + } + __syncthreads(); + const size_t base = ((size_t)c * H + h) * BT * K; + // The rhs panes are born here, over the dead mma operands: pane 0 = kappa + // -> P, pane 1 = v * beta -> u0. Rebuilt bit-equal (deterministic __expf, + // same op sequence) and solved in k1_body's per-column order. + auto rhs_fill = [&](int pane) { // k1_body's P2b tap map (colp pair) + const size_t gp0 = (size_t)(t0 + tid * 2 / K) * H * K + h * K + tid * 2 % K; + const size_t gst = (size_t)8 * H * K; + const bf16* px = &v[gp0]; // pane 0's k taps are already in kvz +#pragma unroll + for (int hb = 0; hb < 2; ++hb) { + __nv_bfloat162 xr[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const int m = hb * 4 + j; + xr[j] = pane == 0 ? kvz[m >> 1][m & 1] + : (!VL || tid * 2 / K + m * 8 < C_act ? *reinterpret_cast(px) + : __floats2bfloat162_rn(0.f, 0.f)); // pad rows + px += gst; + } +#pragma unroll + for (int j = 0; j < 4; ++j) { + const int i = (tid >> 6) + (hb * 4 + j) * 8; + const int si = hb * 2 + (j >> 1); + const float2 xv = __bfloat1622float2(xr[j]); + const float bi = sb[i]; + if (pane == 0) { + const float ei0 = __expf(sgc[i][colp] - a0v[si]); + const float ei1 = __expf(sgc[i][colp + 1] - a1v[si]); + const float kw0 = xv.x * ei0, kw1 = xv.y * ei1; + *reinterpret_cast(&sU[i][su_c(i, colp)]) = float2{kw0 * (bi * ea0[si]), kw1 * (bi * ea1[si])}; + } else { + *reinterpret_cast(&sU[i][su_c(i, colp)]) = float2{xv.x * bi, xv.y * bi}; + } + } + } + }; +#pragma unroll 1 + for (int pane = 0; pane < 2; ++pane) { + rhs_fill(pane); + __syncthreads(); + auto pu_store4p = [&](int p) { + const float* r4 = &sU[p / K][su_c(p / K, p % K)]; + if (pane == 0) { + const __nv_bfloat162 h01 = __floats2bfloat162_rn(-r4[0], -r4[1]); + const __nv_bfloat162 h23 = __floats2bfloat162_rn(-r4[2], -r4[3]); + *reinterpret_cast(&P[base + p]) = + uint2{*reinterpret_cast(&h01), *reinterpret_cast(&h23)}; + } else { + *reinterpret_cast(&u0[base + p]) = *reinterpret_cast(r4); + } + }; + // Blocked forward solve: cross-block coupling as mma (hi only, 16 + // warps x ONE n8 tile), then the in-block apply X_b = M_bb @ rhs_b; + // solved rows publish transposed hi as the next coupling's B. The + // whole block row is warp-local (coupling, apply and the upT pack + // touch only cols 8*warp..+7), so the round syncs at warp scope. +#pragma unroll + for (int b = 0; b < NSB; ++b) { + float crr[4] = {}; // coupling correction, B-operand order + if (b) { + const int arow = lane & 15, aka = lane >> 4; + float4 acc = float4{0, 0, 0, 0}; + for (int k16 = 0; k16 < b; ++k16) { + uint32_t a0, a1, a2, a3, b0, b1; + ptx::ldmatrix_x4_b16(ptx::to_shared(&aC_h[b * SB + arow][k16 * 16 + aka * 8]), a0, a1, a2, a3); + ptx::ldmatrix_x2_b16( + ptx::to_shared(&upT_h[warp * 8 + (lane & 7)][k16 * 16 + ((lane >> 3) & 1) * 8]), b0, b1); + ptx::mma_m16n8k16_bf16f32(acc, a0, a1, a2, a3, b0, b1); + } + // C fragment (row fr(+8), col fc(+1)) -> B operand (k fc(+1, + // +8,+9), col fr): a permutation of the warp's own 16x8 tile, + // so no smem round trip — 8 shuffles, and the one fp32 + // subtraction moves into the register feeding the apply's B. + // corr[k][fr] lives in lane 8*(lane&3) + (fr>>1), regs x/y for + // k < 8 and z/w for k >= 8, picked on fr&1 == the col parity. + const int fr = lane >> 2; + const int sla = (lane & 3) * 8 + (fr >> 1), slb = sla + 4; + const unsigned fm = 0xffffffffu; + const float ax = __shfl_sync(fm, acc.x, sla); + const float ay = __shfl_sync(fm, acc.y, sla); + const float az = __shfl_sync(fm, acc.z, sla); + const float aw = __shfl_sync(fm, acc.w, sla); + const float bx = __shfl_sync(fm, acc.x, slb); + const float by = __shfl_sync(fm, acc.y, slb); + const float bz = __shfl_sync(fm, acc.z, slb); + const float bw = __shfl_sync(fm, acc.w, slb); + crr[0] = (fr & 1) ? ay : ax; // k = fc + crr[1] = (fr & 1) ? by : bx; // k = fc + 1 + crr[2] = (fr & 1) ? aw : az; // k = fc + 8 + crr[3] = (fr & 1) ? bw : bz; // k = fc + 9 + } + { // X_b = M_bb @ rhs_b, hi/lo 3-pass: ONE mma per pass per warp, + // all 16 warps. rhs_b comes straight out of fp32 sU — this + // warp's own 8 columns, which no other warp touches. + const int arow = lane & 15, aka = lane >> 4; + const int fr = lane >> 2, fc = (lane & 3) * 2; + const int bcol = warp * 8 + fr; // B operand column (n = fr) + uint32_t ah[4], al[4], bh[2], bl[2]; + ptx::ldmatrix_x4_b16(ptx::to_shared(&Mh[b][arow][aka * 8]), ah[0], ah[1], ah[2], ah[3]); + ptx::ldmatrix_x4_b16(ptx::to_shared(&Ml[b][arow][aka * 8]), al[0], al[1], al[2], al[3]); +#pragma unroll + for (int kp = 0; kp < 2; ++kp) { // k = fc(+1) and fc+8(+9) + const int rr0 = b * SB + fc + kp * 8; + float r0 = sU[rr0][su_c(rr0, bcol)]; + float r1 = sU[rr0 + 1][su_c(rr0 + 1, bcol)]; + if (b) { // the coupling's fp32 subtraction, in-register + r0 -= crr[kp * 2]; + r1 -= crr[kp * 2 + 1]; + } + const __nv_bfloat162 hv = __floats2bfloat162_rn(r0, r1); + const __nv_bfloat162 lv = __floats2bfloat162_rn(r0 - __bfloat162float(hv.x), r1 - __bfloat162float(hv.y)); + bh[kp] = *reinterpret_cast(&hv); + bl[kp] = *reinterpret_cast(&lv); + } + float4 acc = float4{0, 0, 0, 0}; + ptx::mma_m16n8k16_bf16f32(acc, ah[0], ah[1], ah[2], ah[3], bh[0], bh[1]); + ptx::mma_m16n8k16_bf16f32(acc, al[0], al[1], al[2], al[3], bh[0], bh[1]); + ptx::mma_m16n8k16_bf16f32(acc, ah[0], ah[1], ah[2], ah[3], bl[0], bl[1]); + __syncwarp(); // all lanes' rhs_b reads precede the write + const float vals[4] = {acc.x, acc.y, acc.z, acc.w}; +#pragma unroll + for (int e = 0; e < 4; ++e) { + const int xr = b * SB + fr + (e >> 1) * 8; + const int xc = warp * 8 + fc + (e & 1); + if (!(e & 1)) // the n-pair explicitly: stays 64-bit + *reinterpret_cast(&sU[xr][su_c(xr, xc)]) = float2{vals[e], vals[e + 1]}; + if (b + 1 < NSB) // next coupling's B operand (hi^T) + upT_h[xc][xr] = __float2bfloat16(vals[e]); + } + __syncwarp(); // upT_b feeds this warp's next coupling + } + } + // the apply is warp-local, so nothing streamed out under it: one + // barrier, then the whole pane publishes + __syncthreads(); + for (int p = tid * 4; p < BT * K; p += blockDim.x * 4) + pu_store4p(p); + __syncthreads(); // sU/upT reuse across panes (pool reuse at job end) + } + } // job loop +} + +// Piece-major block order: consecutive blocks are the same piece of adjacent +// heads, i.e. the same token rows — their q/k/v/g lines are contiguous, so a +// wave shares L2. h stays constant within a block (the GM hoists need it). +template +__global__ void __launch_bounds__(512, 2) k1_tf_builder( + const bf16* __restrict__ q, + const bf16* __restrict__ kk, + const bf16* __restrict__ v, + const bf16* __restrict__ glog, + const __grid_constant__ CUtensorMap glog_map, + const float* __restrict__ a_log, + const float* __restrict__ dtb, + float lb, + const float* __restrict__ beta, + int T, + int H, + float scale, + bf16* __restrict__ P, + float* __restrict__ u0, + bf16* __restrict__ kdec, + bf16* __restrict__ qdec, + bf16* __restrict__ aqk_h, + bf16* __restrict__ aqk_l, + float* __restrict__ gC, + int NP) { + __shared__ std::conditional_t S; + float2* rn = nullptr; + if constexpr (RAW) rn = S.rn; + const int nc1 = T / BT, bid = blockIdx.x; + const int h = bid % H, pc = bid / H; + const int c00 = piece_c0(pc, nc1, NP); + const int njobs = piece_c0(pc + 1, nc1, NP) - c00; + k1_tf_body( + S, + h * nc1 + c00, + njobs, + q, + kk, + v, + glog, + glog_map, + beta, + T, + H, + scale, + P, + u0, + kdec, + qdec, + aqk_h, + aqk_l, + gC, + a_log, + dtb, + lb, + 0, + 0, + rn); +} +template +__global__ void __launch_bounds__(512, 2) k1_tf_builder_vl( + const bf16* __restrict__ q, + const bf16* __restrict__ kk, + const bf16* __restrict__ v, + const bf16* __restrict__ glog, + const __grid_constant__ CUtensorMap glog_map, + const float* __restrict__ a_log, + const float* __restrict__ dtb, + float lb, + const float* __restrict__ beta, + int H, + float scale, + bf16* __restrict__ P, + float* __restrict__ u0, + bf16* __restrict__ kdec, + bf16* __restrict__ qdec, + bf16* __restrict__ aqk_h, + bf16* __restrict__ aqk_l, + float* __restrict__ gC, + const VlPiece* __restrict__ pieces, + int nc_tot) { + __shared__ std::conditional_t S; + float2* rn = nullptr; + if constexpr (RAW) rn = S.rn; + const int h = blockIdx.x % H; + const VlPiece* pc = pieces + blockIdx.x / H; + const int c00 = pc->cbase + piece_c0(pc->sg, pc->nc, pc->np); + const int njobs = pc->cbase + piece_c0(pc->sg + 1, pc->nc, pc->np) - c00; + if (pc->tail) // partial last chunk: pad-row guards live + k1_tf_body( + S, + h * nc_tot + c00, + njobs, + q, + kk, + v, + glog, + glog_map, + beta, + nc_tot * BT, + H, + scale, + P, + u0, + kdec, + qdec, + aqk_h, + aqk_l, + gC, + a_log, + dtb, + lb, + pc->tokoff, + pc->tend, + rn); + else // full piece: the eqlen body on T = nc_tot*BT (see VlPiece) + k1_tf_body( + S, + h * nc_tot + c00, + njobs, + q, + kk, + v, + glog, + glog_map, + beta, + nc_tot * BT, + H, + scale, + P, + u0, + kdec, + qdec, + aqk_h, + aqk_l, + gC, + a_log, + dtb, + lb, + pc->tokoff, + pc->tend, + rn); +} +// varlen seq0 chain: one CTA per (sequence, head) walking the whole sequence +// from h0 (nseg == 1 — no piece maps, no flags). seqs[s] is the whole-sequence +// entry the host appends after the builder piece table. +__global__ void __launch_bounds__(512) k2_chain_tc_vl( + const __grid_constant__ CUtensorMap pneg_map, + const __grid_constant__ CUtensorMap kdt_map, + const __grid_constant__ CUtensorMap qd_map, + const __grid_constant__ CUtensorMap aqh_map, + const __grid_constant__ CUtensorMap aql_map, + const __grid_constant__ CUtensorMap u0f_map, + const float* __restrict__ gC, + const float* __restrict__ h0, + int H, + bf16* __restrict__ o, + float* __restrict__ Sf, + void* __restrict__ hpc, + bool hpc_bf16, + bool hpc_v_first, + const VlPiece* __restrict__ seqs) { + __shared__ ChainSmem S; + const int s = blockIdx.x / H, h = blockIdx.x % H; + const VlPiece* sq = seqs + s; + const size_t so = (size_t)s * H * K * K; // h0/Sf per-sequence rows + if (sq->tail) // only a partial last chunk needs the o-store mask + chain_body( + S, + h, + 0, + 1, + pneg_map, + kdt_map, + qd_map, + aqh_map, + aql_map, + u0f_map, + gC, + sq->nc, + H, + o, + Sf + so, + hpc, + hpc_bf16, + hpc_v_first, + nullptr, + nullptr, + h0 + so, + nullptr, + sq->cbase, + 0, + sq->tokoff, + sq->tend); + else + chain_body( + S, + h, + 0, + 1, + pneg_map, + kdt_map, + qd_map, + aqh_map, + aql_map, + u0f_map, + gC, + sq->nc, + H, + o, + Sf + so, + hpc, + hpc_bf16, + hpc_v_first, + nullptr, + nullptr, + h0 + so, + nullptr, + sq->cbase, + 0, + sq->tokoff, + sq->tend); +} + +// Host side: tensor-map encoders, cached workspace, torch entry point. + +#define KDA_CU_CHECK(expr) \ + do { \ + CUresult _e = (expr); \ + if (_e != CUDA_SUCCESS) { \ + const char* _s = nullptr; \ + cuGetErrorString(_e, &_s); \ + TORCH_CHECK(false, "CUDA driver call failed at ", __FILE__, ":", __LINE__, ": ", _s ? _s : "?"); \ + } \ + } while (0) + +// TMA maps over K1's bf16 outputs (128B swizzle, K-major tiles) +static CUtensorMap enc2d(void* ptr, uint64_t rows, uint64_t cols, uint32_t brows, uint32_t bcols) { + cuuint64_t gdim[2] = {cols, rows}; + cuuint64_t gstr[1] = {cols * 2}; + cuuint32_t bdim[2] = {bcols, brows}; + cuuint32_t estr[2] = {1, 1}; + CUtensorMap m{}; + KDA_CU_CHECK(cuTensorMapEncodeTiled( + &m, + CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, + 2, + ptr, + gdim, + gstr, + bdim, + estr, + CU_TENSOR_MAP_INTERLEAVE_NONE, + CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE)); + return m; +} +// fp32 c-map tiles for the tf32 chain prefix: SWZ32 k-major chunks +// (box 8 fp32 x 128 rows; 32B swizzle caps the box inner extent at 8 fp32 +// -> 16 boxes per 128x128 tile) +static CUtensorMap enc2df(void* ptr, uint64_t rows, uint64_t cols) { + cuuint64_t gdim[2] = {cols, rows}; + cuuint64_t gstr[1] = {cols * 4}; + cuuint32_t bdim[2] = {8, 128}; + cuuint32_t estr[2] = {1, 1}; + CUtensorMap m{}; + KDA_CU_CHECK(cuTensorMapEncodeTiled( + &m, + CU_TENSOR_MAP_DATA_TYPE_FLOAT32, + 2, + ptr, + gdim, + gstr, + bdim, + estr, + CU_TENSOR_MAP_INTERLEAVE_NONE, + CU_TENSOR_MAP_SWIZZLE_32B, + CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE)); + return m; +} +// fp32 u0 chunk tiles for the chain: plain linear [BT x K] boxes (the +// consumer is per-thread smem loads, not an mma desc, so no swizzle — and +// only swizzled modes cap the box inner extent) +static CUtensorMap enc2dfn(void* ptr, uint64_t rows, uint64_t cols) { + cuuint64_t gdim[2] = {cols, rows}; + cuuint64_t gstr[1] = {cols * 4}; + cuuint32_t bdim[2] = {(cuuint32_t)cols, BT}; + cuuint32_t estr[2] = {1, 1}; + CUtensorMap m{}; + KDA_CU_CHECK(cuTensorMapEncodeTiled( + &m, + CU_TENSOR_MAP_DATA_TYPE_FLOAT32, + 2, + ptr, + gdim, + gstr, + bdim, + estr, + CU_TENSOR_MAP_INTERLEAVE_NONE, + CU_TENSOR_MAP_SWIZZLE_NONE, + CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE)); + return m; +} +// The tail-free builder's gate input: [BT x K] bf16 boxes out of the +// [T x H*K] gate stream. SWIZZLE_NONE, so the landed tile is row-major with a +// 256 B pitch — byte-for-byte the per-thread staging it replaces. rows == the +// ALLOCATED token count, so a tail chunk's rows past the buffer read as zeros. +static CUtensorMap enc2dgb(void* ptr, uint64_t rows, uint64_t cols) { + cuuint64_t gdim[2] = {cols, rows}; + cuuint64_t gstr[1] = {cols * 2}; + cuuint32_t bdim[2] = {K, BT}; + cuuint32_t estr[2] = {1, 1}; + CUtensorMap m{}; + KDA_CU_CHECK(cuTensorMapEncodeTiled( + &m, + CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, + 2, + ptr, + gdim, + gstr, + bdim, + estr, + CU_TENSOR_MAP_INTERLEAVE_NONE, + CU_TENSOR_MAP_SWIZZLE_NONE, + CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE)); + return m; +} + +// Cached workspace, keyed by (device, T, nc_tot, H, npieces, N). The factor +// tensors are torch allocations (caching allocator -> stream-safe reuse, +// freed with the process). The CUDA tensor maps encode over these tensors' +// data pointers, which are stable for the cache entry's lifetime — so the +// maps are encoded ONCE here rather than per call (this IS the +// (pointer, shape) tensor-map cache; a fresh encode is only host-cheap ~us, +// but per-call re-encode is pure waste when the pointers never change). +// NOTE: aqk_h/aqk_l are zero-FILLED once — k1 only ever writes the masked +// lower-triangular cells, the TMA-read upper cells must stay zero. +// The gate map is the exception: g is a CALLER tensor, so its pointer is only +// stable while the allocator hands back the same block. One-entry memo. +struct Workspace { + torch::Tensor P, u0, kdec, qdec, aqk_h, aqk_l, gC, pieceL, piecec, pflags; + torch::Tensor h0z; // zeros initial state (lazy) + torch::Tensor pieces_dev; // varlen piece table (lazy) + std::vector cu; // piece-table provenance (varlen) + CUtensorMap pneg_map, kdt_map, qd_map, aqh_map, aql_map, sl_map, sc_map, u0f_map, gin_map; + const void* gin_ptr = nullptr; // what gin_map was encoded over +}; + +// Resident builder-CTA slots on this GPU, per route: the fused grid's builder +// half is NP*H CTAs of the 2*NP*H launch and co-resides 1/SM; the tail-free +// builder is its own NP*H launch at 2/SM. This is what a builder "wave" costs. +static int64_t builder_slots(bool tf, bool varlen) { + auto query = [](auto fn) { + int occ = 0; + AT_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&occ, fn, 512, 0)); + return int64_t(at::cuda::getCurrentDeviceProperties()->multiProcessorCount) * std::max(1, occ); + }; + static const int64_t slots[4] = { + query(kda_fused<0>), query(kda_fused_vl<0>), query(k1_tf_builder<0>), query(k1_tf_builder_vl<0>)}; + return slots[2 * int(tf) + int(varlen)]; +} + +// (pieces, max chunks in any piece) for a piece budget npb +static std::pair piece_plan(const std::vector& cu, int64_t T, int npb, bool varlen) { + if (!varlen) { + const int64_t nc = T / BT; + const int64_t np = std::min(npb, std::max(1, nc / 2)); + return {np, (nc + np - 1) / np}; + } + int nc_tot = 0; + for (size_t s = 0; s + 1 < cu.size(); ++s) + nc_tot += int((cu[s + 1] - cu[s] + BT - 1) / BT); + int64_t pieces = 0, mx = 0; + for (size_t s = 0; s + 1 < cu.size(); ++s) { + const int ncs = int((cu[s + 1] - cu[s] + BT - 1) / BT); + const int np = std::min(std::min(npb, std::max(1, ncs / 2)), std::max(1, (2 * npb * ncs + nc_tot) / (2 * nc_tot))); + pieces += np; + mx = std::max(mx, (ncs + np - 1) / np); + } + return {pieces, mx}; +} + +// Pieces are persistent builders, each walking max_len chunks serially, so the +// build wall is a bin-packing makespan: ceil(builders/slots) waves of max_len, +// plus per-builder composition. Minimizing +// waves * max_len * slots + builders +// (makespan scaled into builder units + one composition unit per builder) +// reproduces every swept optimum on a 152-SM GB300: 1kx8-H12 -> 24 pieces +// (-17.1% vs the old fixed 12), 8k/16k-H12 -> 12 (one 0.95x-fill wave, already +// optimal), 8k-H96 -> 3, 1kx8-H96 -> 16, mixed6-H96 -> 14 (measured optimum 12, +// +0.3%). Ties break toward fewer pieces, which is why the H12 eqlen cells stay +// on one wave instead of splitting to two. +static int piece_budget(const std::vector& cu, int64_t T, int64_t H, bool varlen, int64_t slots) { + int64_t best_cost = -1, best_pieces = 0; + int best_npb = 1; + for (int npb = 1; npb <= 64; ++npb) { + const auto [pieces, mx] = piece_plan(cu, T, npb, varlen); + const int64_t builders = pieces * H; + const int64_t waves = (builders + slots - 1) / slots; + const int64_t cost = waves * mx * slots + builders; + if (best_cost < 0 || cost < best_cost || (cost == best_cost && pieces < best_pieces)) { + best_cost = cost; + best_pieces = pieces; + best_npb = npb; + } + } + return best_npb; +} + +// Measured per-unit costs (us) behind the route pick on this shape family: one +// builder chunk-job costs 10.2 us*SM tail-free vs 16.2 fused (that builder also +// composes its piece map in-tail), one chain chunk-step 2.1 us. The tail-free +// figure tracks that builder's measured wall as it improves: 11.7 before the +// block apply, 11.1 after it (-4.9%), 10.2 after the four data-pipe items +// (-8.1% more). Every battery shape keeps its route across that whole range — +// the crossover cell (vl-prod-H12) only gains margin. +constexpr double C_JOB_TF = 10.2, C_JOB_FUSED = 16.2, C_STEP = 2.1; + +// Route pick, derived from the launch geometry like the piece budget above: the +// seq0 pair (tail-free builders, then one chain per (sequence, head) walking +// that whole sequence) vs the fused grid. Each wall is max(throughput term, SM +// critical path): the fused grid's chains trail its builders inside one grid, +// so their step time joins its throughput term and only its longest piece's +// chain is exposed, while seq0's chain is a second launch whose depth is the +// LONGEST SEQUENCE. seq0 therefore wins on many-sequence and high-H shapes and +// loses on long single sequences (measured +50% at 8k-H12), which is what this +// reproduces. The margin is the model's own residual at the crossover: on the +// one shape that was ever near it (4 seqs 8192/8192/4096/4096, H12) the model +// priced seq0 at -2.7% against a measured -1.5%, so a flip must be modelled at +// >=2.5% to cover that 1.2-point optimism. Every other battery shape clears the +// guard by >=21 points either way. A single piece has no composition to skip, +// so it stays fused too. +struct Route { + bool seq0; + int npb; +}; +static Route pick_route(const std::vector& cu, int64_t T, int64_t H, bool varlen) { + const double sms = at::cuda::getCurrentDeviceProperties()->multiProcessorCount; + int64_t nc_tot = 0, depth = 0; + if (varlen) + for (size_t s = 0; s + 1 < cu.size(); ++s) { + const int64_t ncs = (cu[s + 1] - cu[s] + BT - 1) / BT; + nc_tot += ncs; + depth = std::max(depth, ncs); + } + else + nc_tot = depth = T / BT; + const double jobs = double(nc_tot * H) / sms; // builder jobs per SM + struct Arm { + int npb; + int64_t pieces; + double wall; + }; + auto arm = [&](bool tf) { + const int64_t slots = builder_slots(tf, varlen); + const int npb = piece_budget(cu, T, H, varlen, slots); + const auto [pieces, mx] = piece_plan(cu, T, npb, varlen); + const double c_job = tf ? C_JOB_TF : C_JOB_FUSED; + const double share = // builder CTAs sharing one SM + std::min(double(slots) / sms, std::max(1.0, double(pieces * H) / sms)); + const double build = std::max(jobs * (tf ? c_job : c_job + C_STEP), double(mx) * share * c_job); + const double chain = (tf ? std::max(double(depth), jobs) : double(mx)) * C_STEP; + return Arm{npb, pieces, build + chain}; + }; + const Arm tf = arm(true), fused = arm(false); + if (tf.pieces > 1 && tf.wall <= 0.975 * fused.wall) return {true, tf.npb}; + return {false, fused.npb}; +} + +static Workspace& +get_workspace(const torch::Device& dev, int64_t T, int64_t nc, int64_t H, int64_t npieces, int64_t N) { + // guarded by the GIL (single writer); entries live for the process + static std::map, Workspace> cache; + const std::array key{dev.index(), T, nc, H, npieces, N}; + auto it = cache.find(key); + if (it != cache.end()) return it->second; + Workspace ws; + const auto ob = torch::TensorOptions().dtype(torch::kBFloat16).device(dev); + const auto of = torch::TensorOptions().dtype(torch::kFloat).device(dev); + const auto oi = torch::TensorOptions().dtype(torch::kInt).device(dev); + const int64_t chk = nc * H * BT * K; + ws.P = torch::empty({chk}, ob); + ws.u0 = torch::empty({chk}, of); + ws.kdec = torch::empty({chk}, ob); + ws.qdec = torch::empty({chk}, ob); + ws.aqk_h = torch::zeros({nc * H * BT * BT}, ob); + ws.aqk_l = torch::zeros({nc * H * BT * BT}, ob); + ws.gC = torch::empty({nc * H * K}, of); + // L maps bf16 (the chain widens them to tf32); c maps fp32 (bf16 c + // failed the 16K o gate) + ws.pieceL = torch::empty({npieces * H * K * K}, ob); + ws.piecec = torch::empty({npieces * H * K * K}, of); + ws.pflags = torch::empty({npieces * H}, oi); + ws.pneg_map = enc2d(ws.P.data_ptr(), (uint64_t)(nc * H * BT), K, BT, 64); + ws.kdt_map = enc2d(ws.kdec.data_ptr(), (uint64_t)(nc * H * K), BT, K, 64); + ws.qd_map = enc2d(ws.qdec.data_ptr(), (uint64_t)(nc * H * BT), K, BT, 64); + ws.aqh_map = enc2d(ws.aqk_h.data_ptr(), (uint64_t)(nc * H * BT), BT, BT, 64); + ws.aql_map = enc2d(ws.aqk_l.data_ptr(), (uint64_t)(nc * H * BT), BT, BT, 64); + ws.sl_map = enc2d(ws.pieceL.data_ptr(), (uint64_t)(npieces * H * K), K, K, 64); + ws.sc_map = enc2df(ws.piecec.data_ptr(), (uint64_t)(npieces * H * K), K); + ws.u0f_map = enc2dfn(ws.u0.data_ptr(), (uint64_t)(nc * H * BT), K); + return cache.emplace(key, std::move(ws)).first->second; +} + +static_assert(sizeof(VlPiece) == 8 * sizeof(int)); + +// kda_prefill_fwd — the single forward entry point. +// q, k, v : [T, H, 128] bf16 contiguous (flat token stream); q/k l2-normed +// unless use_qk_l2norm_in_kernel +// g : [T, H, 128] bf16 — pre-transformed glog (use_gate_in_kernel = +// false; narrow fp32 -> bf16 in the wrapper) or RAW gate input +// (use_gate_in_kernel = true) +// beta : [T, H] bf16 (or fp32) — widened to fp32 internally; the +// activated beta, or its LOGITS under use_beta_sigmoid_in_kernel +// scale : q scaling (typically 128**-0.5) +// initial_state : [N, H, 128, 128] fp32 or None (zeros) +// cu_seqlens : int32/int64 [N+1] (any device; host values are needed +// for the piece table — pass a CPU tensor to avoid the D2H sync) +// or None => single sequence [0, T] +// use_gate_in_kernel + A_log [H] f32 + dt_bias [H*128] f32 + safe_gate + +// lower_bound: the production raw-gate convention (safe_gate +// false => softplus, true => lower_bound * sigmoid) +// use_qk_l2norm_in_kernel: q/k raw, l2norm(q)/l2norm(k) fused in (fla's +// l2norm eps/rounding) +// use_beta_sigmoid_in_kernel: beta raw (logits), sigmoid(beta) fused in. +// Independent of use_qk_l2norm_in_kernel, both default false — +// fla's two flags, fla's meanings +// h_per_chunk : optional PREALLOCATED [nc_tot, H, 128, 128] fp32 or bf16 +// per-chunk state output (nc_tot = sum_n ceil(len_n / 64) = the +// kernel's own chunk count), h_v_first = store it [V, K] instead +// of the native [K, V] (the wrapper's state_v_first) +// Returns (o [T, H, 128] bf16, final_state [N, H, 128, 128] fp32). +static std::tuple kda_prefill_fwd( + const torch::Tensor& q, + const torch::Tensor& k, + const torch::Tensor& v, + const torch::Tensor& g, + const torch::Tensor& beta, + double scale, + const std::optional& initial_state, + const std::optional& cu_seqlens, + bool use_gate_in_kernel, + const std::optional& A_log, + const std::optional& dt_bias, + bool safe_gate, + double lower_bound, + bool use_qk_l2norm_in_kernel, + bool use_beta_sigmoid_in_kernel, + const std::optional& h_per_chunk, + bool h_v_first) { + TORCH_CHECK(q.is_cuda() && q.dim() == 3 && q.size(2) == K, "q must be a CUDA [T, H, 128] tensor, got ", q.sizes()); + const int64_t T = q.size(0), H = q.size(1); + TORCH_CHECK(T >= 1 && H >= 1, "empty input: T=", T, " H=", H); + for (auto* t : {&q, &k, &v, &g}) { + TORCH_CHECK( + t->is_cuda() && t->is_contiguous() && t->scalar_type() == torch::kBFloat16 && t->sizes() == q.sizes(), + "q/k/v/g must be contiguous CUDA bf16 [T, H, 128] with " + "matching shapes; got ", + t->sizes(), + " ", + t->dtype()); + } + TORCH_CHECK( + beta.is_cuda() && beta.is_contiguous() && beta.dim() == 2 && beta.size(0) == T && beta.size(1) == H, + "beta must be contiguous CUDA [T, H], got ", + beta.sizes()); + TORCH_CHECK( + beta.scalar_type() == torch::kBFloat16 || beta.scalar_type() == torch::kFloat, + "beta must be bf16 or fp32, got ", + beta.dtype()); + const int gm = use_gate_in_kernel ? (safe_gate ? 2 : 1) : 0; + const float lb = float(lower_bound); + const float* alog_p = nullptr; + const float* dtb_p = nullptr; + if (use_gate_in_kernel) { + TORCH_CHECK(A_log && dt_bias, "use_gate_in_kernel=True requires A_log and dt_bias"); + TORCH_CHECK( + A_log->is_cuda() && A_log->is_contiguous() && A_log->scalar_type() == torch::kFloat && A_log->numel() == H, + "A_log must be contiguous CUDA fp32 [H], got ", + A_log->sizes(), + " ", + A_log->dtype()); + TORCH_CHECK( + dt_bias->is_cuda() && dt_bias->is_contiguous() && dt_bias->scalar_type() == torch::kFloat && + dt_bias->numel() == H * K, + "dt_bias must be contiguous CUDA fp32 with H*128 " + "elements, got ", + dt_bias->sizes(), + " ", + dt_bias->dtype()); + alog_p = A_log->data_ptr(); + dtb_p = dt_bias->data_ptr(); + } + // route: explicit cu_seqlens OR a ragged T runs the varlen grid; an + // aligned single sequence takes the eqlen fused grid + std::vector cu; + if (cu_seqlens) { + TORCH_CHECK( + cu_seqlens->dim() == 1 && cu_seqlens->numel() >= 2, "cu_seqlens must be 1-D [N+1], got ", cu_seqlens->sizes()); + const auto cu_cpu = // D2H sync when given on device + cu_seqlens->to(torch::kCPU).to(torch::kLong).contiguous(); + const int64_t* p = cu_cpu.data_ptr(); + cu.assign(p, p + cu_cpu.numel()); + TORCH_CHECK( + cu.front() == 0 && cu.back() == T, + "cu_seqlens must span [0, T=", + T, + "], got [", + cu.front(), + ", ", + cu.back(), + "]"); + for (size_t s = 0; s + 1 < cu.size(); ++s) + TORCH_CHECK( + cu[s] < cu[s + 1], + "cu_seqlens must be strictly increasing (empty " + "sequences unsupported) at index ", + s); + } else if (T % BT != 0) { + cu = {0, T}; // ragged single sequence -> varlen grid + } + const bool varlen = !cu.empty(); + const int64_t N = varlen ? int64_t(cu.size()) - 1 : 1; + TORCH_CHECK(!(varlen && N > 1) || cu_seqlens, "internal: synthetic cu is always single-sequence"); + + // auto-NP (eqlen): NP = min(npb, max(1, nc/2)). piece table (varlen): + // that piece budget is GLOBAL (sequences already parallelize the grid; + // a per-seq NP over-fills it) — pro-rated by nc_s, round-half-up, + // clamped to the per-seq cap; N=1 reduces exactly to the eqlen NP + const Route route = pick_route(cu, T, H, varlen); + const int npb = route.npb; + std::vector pieces; + int64_t nc = 0, npieces = 0; + if (varlen) { + int nc_tot = 0; + for (size_t s = 0; s + 1 < cu.size(); ++s) + nc_tot += int((cu[s + 1] - cu[s] + BT - 1) / BT); + std::vector seqs; // seq0's chains: one whole-sequence piece + int nc_run = 0; + for (size_t s = 0; s + 1 < cu.size(); ++s) { + const int len = int(cu[s + 1] - cu[s]); + const int ncs = (len + BT - 1) / BT; + const int np = + std::min(std::min(npb, std::max(1, ncs / 2)), std::max(1, (2 * npb * ncs + nc_tot) / (2 * nc_tot))); + const int tokoff = int(cu[s]) - nc_run * BT; + for (int sg = 0; sg < np; ++sg) + pieces.push_back(VlPiece{int(s), nc_run, ncs, np, sg, tokoff, int(cu[s + 1]), sg == np - 1 && len % BT != 0}); + seqs.push_back(VlPiece{int(s), nc_run, ncs, 1, 0, tokoff, int(cu[s + 1]), len % BT != 0}); + nc_run += ncs; + } + nc = nc_run; + npieces = int64_t(pieces.size()); + pieces.insert(pieces.end(), seqs.begin(), seqs.end()); + } else { + nc = T / BT; + npieces = std::min(npb, std::max(1, nc / 2)); + } + + const c10::cuda::CUDAGuard guard(q.device()); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + Workspace& ws = get_workspace(q.device(), T, nc, H, npieces, N); + + // initial state (None -> cached zeros; kernels only READ h0) + const float* h0_p; + if (initial_state) { + TORCH_CHECK( + initial_state->is_cuda() && initial_state->is_contiguous() && initial_state->scalar_type() == torch::kFloat && + initial_state->sizes() == torch::IntArrayRef({N, H, int64_t(K), int64_t(K)}), + "initial_state must be contiguous CUDA fp32 [N=", + N, + ", H=", + H, + ", 128, 128], got ", + initial_state->sizes(), + " ", + initial_state->dtype()); + h0_p = initial_state->data_ptr(); + } else { + if (!ws.h0z.defined()) + ws.h0z = + torch::zeros({N, H, int64_t(K), int64_t(K)}, torch::TensorOptions().dtype(torch::kFloat).device(q.device())); + h0_p = ws.h0z.data_ptr(); + } + + // outputs (fresh each call — everything else is workspace) + auto o = torch::empty({T, H, int64_t(K)}, torch::TensorOptions().dtype(torch::kBFloat16).device(q.device())); + auto Sf = + torch::empty({N, H, int64_t(K), int64_t(K)}, torch::TensorOptions().dtype(torch::kFloat).device(q.device())); + + // per-chunk states (caller-allocated): row (n, j) = chunk_offset[n] + j is + // sequence n's state after exactly j*BT tokens (j = 0 == initial_state), so + // a consumer can only snapshot at multiples of BT == 64. The last boundary + // is NOT here — it is final_state. + void* hpc_p = nullptr; + bool hpc_bf16 = false; + if (h_per_chunk) { + const auto& hp = *h_per_chunk; + TORCH_CHECK( + hp.is_cuda() && hp.device() == q.device() && hp.is_contiguous() && + hp.sizes() == torch::IntArrayRef({nc, H, int64_t(K), int64_t(K)}), + "h_per_chunk must be a contiguous [nc=", + nc, + ", H=", + H, + ", 128, 128] tensor on ", + q.device(), + ", got ", + hp.sizes(), + " on ", + hp.device()); + TORCH_CHECK( + hp.scalar_type() == torch::kFloat || hp.scalar_type() == torch::kBFloat16, + "h_per_chunk must be fp32 or bf16, got ", + hp.dtype()); + hpc_bf16 = hp.scalar_type() == torch::kBFloat16; + hpc_p = hp.data_ptr(); + } + + const torch::Tensor beta_f = // fp32 widen (exact; kernels read fp32) + beta.scalar_type() == torch::kFloat ? beta : beta.to(torch::kFloat); + + const bf16* q_p = reinterpret_cast(q.const_data_ptr()); + const bf16* k_p = reinterpret_cast(k.const_data_ptr()); + const bf16* v_p = reinterpret_cast(v.const_data_ptr()); + const bf16* g_p = reinterpret_cast(g.const_data_ptr()); + const float* beta_p = beta_f.data_ptr(); + bf16* o_p = reinterpret_cast(o.data_ptr()); + float* Sf_p = Sf.data_ptr(); + bf16* P_p = reinterpret_cast(ws.P.data_ptr()); + bf16* kdec_p = reinterpret_cast(ws.kdec.data_ptr()); + bf16* qdec_p = reinterpret_cast(ws.qdec.data_ptr()); + bf16* aqh_p = reinterpret_cast(ws.aqk_h.data_ptr()); + bf16* aql_p = reinterpret_cast(ws.aqk_l.data_ptr()); + float* u0_p = ws.u0.data_ptr(); + float* gC_p = ws.gC.data_ptr(); + bf16* pL_p = reinterpret_cast(ws.pieceL.data_ptr()); + float* pc_p = ws.piecec.data_ptr(); + uint32_t* pf_p = reinterpret_cast(ws.pflags.data_ptr()); + const float scl = float(scale); + + // (gate mode, q/k raw, beta logits) — the three INDEPENDENT compile-time + // input conventions the bodies fold in; f takes all three as constants + auto gm_dispatch = [&](auto&& f) { + auto bsig = [&](auto gmv, auto rawv) { + if (use_beta_sigmoid_in_kernel) + f(gmv, rawv, std::true_type{}); + else + f(gmv, rawv, std::false_type{}); + }; + auto raw = [&](auto gmv) { + if (use_qk_l2norm_in_kernel) + bsig(gmv, std::true_type{}); + else + bsig(gmv, std::false_type{}); + }; + if (gm == 2) + raw(std::integral_constant{}); + else if (gm == 1) + raw(std::integral_constant{}); + else + raw(std::integral_constant{}); + }; + + const VlPiece* pt_p = nullptr; + if (varlen) { + // per-cu piece table upload (only when cu changes for this key): the + // npieces builder pieces, then seq0's N whole-sequence chain entries + if (ws.cu != cu) { + static_assert(std::is_standard_layout_v); + const auto hp = + torch::from_blob(pieces.data(), {(npieces + N) * 8}, torch::TensorOptions().dtype(torch::kInt)).clone(); + if (!ws.pieces_dev.defined()) + ws.pieces_dev = torch::empty({(npieces + N) * 8}, torch::TensorOptions().dtype(torch::kInt).device(q.device())); + ws.pieces_dev.copy_(hp); + ws.cu = cu; + } + pt_p = reinterpret_cast(ws.pieces_dev.const_data_ptr()); + } + if (route.seq0 && g_p != ws.gin_ptr) { // only the tail-free builder reads + ws.gin_map = enc2dgb(const_cast(g_p), (uint64_t)T, (uint64_t)(H * K)); + ws.gin_ptr = g_p; + } + + // raw modes pass g (RAW graw) as glog: k1 stages it through the same bf16 + // slot and fuses the transform (GM 1/2), on every route + if (route.seq0 && varlen) { + gm_dispatch([&](auto gmv, auto rawv, auto bsv) { + k1_tf_builder_vl<<>>( + q_p, + k_p, + v_p, + g_p, + ws.gin_map, + alog_p, + dtb_p, + lb, + beta_p, + int(H), + scl, + P_p, + u0_p, + kdec_p, + qdec_p, + aqh_p, + aql_p, + gC_p, + pt_p, + int(nc)); + }); + k2_chain_tc_vl<<>>( + ws.pneg_map, + ws.kdt_map, + ws.qd_map, + ws.aqh_map, + ws.aql_map, + ws.u0f_map, + gC_p, + h0_p, + int(H), + o_p, + Sf_p, + hpc_p, + hpc_bf16, + h_v_first, + pt_p + npieces); + } else if (route.seq0) { + gm_dispatch([&](auto gmv, auto rawv, auto bsv) { + k1_tf_builder<<>>( + q_p, + k_p, + v_p, + g_p, + ws.gin_map, + alog_p, + dtb_p, + lb, + beta_p, + int(T), + int(H), + scl, + P_p, + u0_p, + kdec_p, + qdec_p, + aqh_p, + aql_p, + gC_p, + int(npieces)); + }); + k2_chain_tc<<>>( + ws.pneg_map, + ws.kdt_map, + ws.qd_map, + ws.aqh_map, + ws.aql_map, + ws.sl_map, + ws.sc_map, + ws.u0f_map, + gC_p, + h0_p, + int(nc), + int(H), + o_p, + Sf_p, + hpc_p, + hpc_bf16, + h_v_first); + } else if (varlen) { + AT_CUDA_CHECK(cudaMemsetAsync(pf_p, 0, size_t(npieces * H) * 4, stream)); + gm_dispatch([&](auto gmv, auto rawv, auto bsv) { + kda_fused_vl<<<2 * int(npieces * H), 512, 0, stream>>>( + q_p, + k_p, + v_p, + g_p, + alog_p, + dtb_p, + lb, + beta_p, + int(T), + int(H), + scl, + P_p, + u0_p, + kdec_p, + qdec_p, + aqh_p, + aql_p, + gC_p, + pL_p, + pc_p, + pt_p, + int(npieces), + int(nc), + ws.pneg_map, + ws.kdt_map, + ws.qd_map, + ws.aqh_map, + ws.aql_map, + ws.u0f_map, + ws.sl_map, + ws.sc_map, + h0_p, + o_p, + Sf_p, + hpc_p, + hpc_bf16, + h_v_first, + pf_p); + }); + } else if (npieces > 1) { // eqlen default: one fused trailing grid + AT_CUDA_CHECK(cudaMemsetAsync(pf_p, 0, size_t(npieces * H) * 4, stream)); + gm_dispatch([&](auto gmv, auto rawv, auto bsv) { + kda_fused<<<2 * int(npieces * H), 512, 0, stream>>>( + q_p, + k_p, + v_p, + g_p, + alog_p, + dtb_p, + lb, + beta_p, + int(T), + int(H), + scl, + P_p, + u0_p, + kdec_p, + qdec_p, + aqh_p, + aql_p, + gC_p, + pL_p, + pc_p, + int(npieces), + ws.pneg_map, + ws.kdt_map, + ws.qd_map, + ws.aqh_map, + ws.aql_map, + ws.u0f_map, + ws.sl_map, + ws.sc_map, + h0_p, + o_p, + Sf_p, + hpc_p, + hpc_bf16, + h_v_first, + pf_p); + }); + } else { // eqlen nc < 4: NP == 1 two-kernel path (no piece maps) + gm_dispatch([&](auto gmv, auto rawv, auto bsv) { + k1_factors_mma<<>>( + q_p, + k_p, + v_p, + g_p, + alog_p, + dtb_p, + lb, + beta_p, + int(T), + int(H), + scl, + P_p, + u0_p, + kdec_p, + qdec_p, + aqh_p, + aql_p, + gC_p, + nullptr, + nullptr, + 1); + }); + k2_chain_tc<<>>( + ws.pneg_map, + ws.kdt_map, + ws.qd_map, + ws.aqh_map, + ws.aql_map, + ws.sl_map, + ws.sc_map, + ws.u0f_map, + gC_p, + h0_p, + int(nc), + int(H), + o_p, + Sf_p, + hpc_p, + hpc_bf16, + h_v_first); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {o, Sf}; +} + +} // namespace kda + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def( + "kda_prefill_fwd", + &kda::kda_prefill_fwd, + "KDA chunked prefill forward (inference): returns (o, final_state)", + py::arg("q"), + py::arg("k"), + py::arg("v"), + py::arg("g"), + py::arg("beta"), + py::arg("scale"), + py::arg("initial_state") = std::nullopt, + py::arg("cu_seqlens") = std::nullopt, + py::arg("use_gate_in_kernel") = false, + py::arg("A_log") = std::nullopt, + py::arg("dt_bias") = std::nullopt, + py::arg("safe_gate") = false, + py::arg("lower_bound") = -5.0, + py::arg("use_qk_l2norm_in_kernel") = false, + py::arg("use_beta_sigmoid_in_kernel") = false, + py::arg("h_per_chunk") = std::nullopt, + py::arg("h_v_first") = false); +} diff --git a/python/sglang/kernels/jit/csrc/elementwise/add3.cuh b/python/sglang/kernels/jit/csrc/elementwise/add3.cuh new file mode 100644 index 000000000..d62657166 --- /dev/null +++ b/python/sglang/kernels/jit/csrc/elementwise/add3.cuh @@ -0,0 +1,93 @@ +#include +#include + +#include +#include +#include + +#include + +#include + +namespace sglang { + +struct Add3Params { + const bf16_t* __restrict__ a; // [N] contiguous + const bf16_t* __restrict__ b; // [N] contiguous + const bf16_t* __restrict__ c; // [N] contiguous + bf16_t* __restrict__ out; // [N] contiguous + int64_t n_vecs; // N / kVecElems +}; + +template +__global__ void add3_kernel(const __grid_constant__ Add3Params params) { + constexpr uint32_t kVecPairs = device::kMaxVecBytes / sizeof(bf16x2_t); + using vec_t = device::AlignedVector; + + const int64_t vid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (vid >= params.n_vecs) return; + + // Trigger early, so that the next kernel gets a chance to prefetch. + device::PDLTriggerSecondary(); + + vec_t a, b, c; + if constexpr (kPrefetchBC) { + b.load(params.b, vid); + c.load(params.c, vid); + device::PDLWaitPrimary(); + a.load(params.a, vid); + } else { + device::PDLWaitPrimary(); + a.load(params.a, vid); + b.load(params.b, vid); + c.load(params.c, vid); + } + + vec_t out; +#pragma unroll + for (uint32_t i = 0; i < kVecPairs; ++i) { + out[i] = __hadd2(__hadd2(a[i], b[i]), c[i]); + } + out.store(params.out, vid); +} + +template +struct Add3Kernel { + static constexpr int64_t kVecElems = device::kMaxVecBytes / sizeof(bf16_t); + + static void launch( + const tvm::ffi::TensorView a, + const tvm::ffi::TensorView b, + const tvm::ffi::TensorView c, + const tvm::ffi::TensorView out, + const bool prefetch_bc) { + using namespace host; + + auto N = SymbolicSize{"numel"}; + auto device = SymbolicDevice{}; + device.set_options(); + + TensorMatcher({N}).with_dtype().with_device(device).verify(a).verify(b).verify(c).verify(out); + const auto numel = N.unwrap(); + RuntimeCheck(numel % kVecElems == 0, "numel must be divisible by the vector width"); + if (numel == 0) return; + const auto params = Add3Params{ + .a = static_cast(a.data_ptr()), + .b = static_cast(b.data_ptr()), + .c = static_cast(c.data_ptr()), + .out = static_cast(out.data_ptr()), + .n_vecs = numel / kVecElems, + }; + const auto num_threads = [&]() -> int64_t { + for (int64_t n : {128, 256, 512}) { + if (params.n_vecs <= n * 256) return n; + } + return 512; + }(); + const auto grid = div_ceil(params.n_vecs, num_threads); + const auto kernel = prefetch_bc ? add3_kernel : add3_kernel; + LaunchKernel(grid, num_threads, device.unwrap()).enable_pdl(kUsePDL)(kernel, params); + } +}; + +} // namespace sglang diff --git a/python/sglang/kernels/jit/csrc/elementwise/concat_mla.cuh b/python/sglang/kernels/jit/csrc/elementwise/concat_mla.cuh index eee33318f..7141ac883 100644 --- a/python/sglang/kernels/jit/csrc/elementwise/concat_mla.cuh +++ b/python/sglang/kernels/jit/csrc/elementwise/concat_mla.cuh @@ -186,6 +186,7 @@ constexpr int A_LAST_DIM = 512; constexpr int B_LAST_DIM = 64; constexpr int OUT_LAST_DIM = A_LAST_DIM + B_LAST_DIM; +template __global__ void concat_mla_absorb_q_kernel( bf16_t* a, bf16_t* b, @@ -198,6 +199,8 @@ __global__ void concat_mla_absorb_q_kernel( const int b_stride_1, const int64_t out_stride_0, const int out_stride_1) { + device::PDLWaitPrimary(); + const int flat_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) / 32; const int lane_id = get_lane_id(); @@ -229,6 +232,8 @@ __global__ void concat_mla_absorb_q_kernel( a_buf[i] = *(base_addr + i * 32 + lane_id); } + device::PDLTriggerSecondary(); + { BBufType* base_addr = reinterpret_cast(out + idx_0 * out_stride_0 + idx_1 * out_stride_1 + A_LAST_DIM); *(base_addr + lane_id) = b_buf; @@ -241,6 +246,7 @@ __global__ void concat_mla_absorb_q_kernel( } } +template struct ConcatMlaAbsorbQKernel { static void run(tvm::ffi::TensorView a, tvm::ffi::TensorView b, tvm::ffi::TensorView out) { using namespace host; @@ -306,19 +312,20 @@ struct ConcatMlaAbsorbQKernel { const int grid_size = div_ceil(num_items, num_warps_per_block); const int block_size = num_warps_per_block * 32; - LaunchKernel(grid_size, block_size, device.unwrap())( - concat_mla_absorb_q_kernel, - static_cast(a.data_ptr()), - static_cast(b.data_ptr()), - static_cast(out.data_ptr()), - num_items, - dim_1, - S0_a.unwrap(), - static_cast(S1_a.unwrap()), - S0_b.unwrap(), - static_cast(S1_b.unwrap()), - S0_out.unwrap(), - static_cast(S1_out.unwrap())); + LaunchKernel(grid_size, block_size, device.unwrap()) + .enable_pdl(kUsePDL)( + concat_mla_absorb_q_kernel, + static_cast(a.data_ptr()), + static_cast(b.data_ptr()), + static_cast(out.data_ptr()), + num_items, + dim_1, + S0_a.unwrap(), + static_cast(S1_a.unwrap()), + S0_b.unwrap(), + static_cast(S1_b.unwrap()), + S0_out.unwrap(), + static_cast(S1_out.unwrap())); } }; diff --git a/python/sglang/kernels/jit/csrc/elementwise/set_mla_kv_buffer.cuh b/python/sglang/kernels/jit/csrc/elementwise/set_mla_kv_buffer.cuh index ce28cdc9f..682f48fac 100644 --- a/python/sglang/kernels/jit/csrc/elementwise/set_mla_kv_buffer.cuh +++ b/python/sglang/kernels/jit/csrc/elementwise/set_mla_kv_buffer.cuh @@ -27,6 +27,7 @@ #include #include #include +#include #include #include @@ -47,40 +48,6 @@ struct SetMlaKVBufferParams { uint32_t batch_size; }; -// Warp-cooperative gmem -> smem copy. Picks the widest vec width that divides -// both the per-thread share and the byte total. Caller guarantees src is -// 16-byte aligned (PyTorch tensors are) and dst is the start of a per-warp -// smem slot (also 16-byte aligned by ``alignas(16)``). -template -SGL_DEVICE void warp_g2s_copy(const void* __restrict__ src, void* __restrict__ dst) { - using namespace device; - constexpr int64_t kAlignment = (kBytes % (16 * kWarpThreads) == 0) ? 16 - : (kBytes % (8 * kWarpThreads) == 0) ? 8 - : (kBytes % (4 * kWarpThreads) == 0) ? 4 - : (kBytes % 4 == 0) ? 4 - : 0; - static_assert(kAlignment > 0, "kBytes must be a multiple of 4"); - - using vec_t = AlignedStorage; - constexpr auto kLoopBytes = sizeof(vec_t) * kWarpThreads; - constexpr auto kLoopCount = kBytes / kLoopBytes; - constexpr int64_t kTailVecs = (kBytes - kLoopCount * kLoopBytes) / sizeof(vec_t); - - const auto gmem = tile::Memory::warp(); - -#pragma unroll - for (int64_t i = 0; i < kLoopCount; ++i) { - const auto v = gmem.load(src, i); - gmem.store(dst, v, i); - } - if constexpr (kTailVecs > 0) { - if (gmem.in_bound(kLoopCount * kWarpThreads + kTailVecs, kLoopCount)) { - const auto v = gmem.load(src, kLoopCount); - gmem.store(dst, v, kLoopCount); - } - } -} - template __global__ void set_mla_kv_buffer_kernel(const __grid_constant__ SetMlaKVBufferParams params) { using namespace device; @@ -104,8 +71,8 @@ __global__ void set_mla_kv_buffer_kernel(const __grid_constant__ SetMlaKVBufferP void* const gmem_dst = pointer::offset(params.kv_buffer, loc * params.stride_buffer_bytes); // Warp-cooperative load (nope, rope) into the per-warp smem slot. - warp_g2s_copy(nope_src, &smem[warp_in_cta][0]); - warp_g2s_copy(rope_src, &smem[warp_in_cta][kNopeBytes]); + warp::copy_bytes(nope_src, &smem[warp_in_cta][0]); + warp::copy_bytes(rope_src, &smem[warp_in_cta][kNopeBytes]); // Fence required: TMA reads smem via the async proxy, normal sts writes // through the generic proxy. Without this the TMA engine can observe stale diff --git a/python/sglang/kernels/jit/csrc/elementwise/set_mla_kv_concat_q.cuh b/python/sglang/kernels/jit/csrc/elementwise/set_mla_kv_concat_q.cuh new file mode 100644 index 000000000..6a5e0d69e --- /dev/null +++ b/python/sglang/kernels/jit/csrc/elementwise/set_mla_kv_concat_q.cuh @@ -0,0 +1,607 @@ +// MLA KV-cache write fused with the Q concat, bf16 and fp8 entry points. + +#pragma once + +#include // For TensorMatcher, SymbolicSize, SymbolicDevice +#include // For RuntimeCheck, div_ceil + +#include +#include // For LaunchKernel, SGL_DEVICE, PDL helpers +#include // For AlignedVector +#include // For warp::copy_bytes, elect_one_lane, inclusive_sum + +#include +#include +#include + +#include +#include + +namespace { + +struct SetMlaKVConcatQParams { + // KV scatter side (byte-typed: dtype-agnostic row copies). + const void* __restrict__ k_nope; + const void* __restrict__ k_rope; + void* __restrict__ kv_buffer; + const void* __restrict__ loc; + int64_t stride_nope_bytes; + int64_t stride_rope_bytes; + int64_t stride_buffer_bytes; + uint32_t batch_size; + // Q concat side (bf16, element strides). + const bf16_t* __restrict__ q_nope; + const bf16_t* __restrict__ q_rope; + bf16_t* __restrict__ q_out; + uint32_t num_q_items; // batch_size * num_heads + uint32_t q_dim_1; // num_heads + int64_t qn_stride_0; + int32_t qn_stride_1; + int64_t qr_stride_0; + int32_t qr_stride_1; + int64_t qo_stride_0; + int32_t qo_stride_1; +}; + +template +__global__ void set_mla_kv_concat_q_kernel(const __grid_constant__ SetMlaKVConcatQParams params) { + using namespace device; + static_assert((kNopeBytes + kRopeBytes) % 16 == 0, "TMA bulk store requires total row to be 16-byte aligned"); + + constexpr int64_t kRowBytes = kNopeBytes + kRopeBytes; + constexpr int kQNopeDim = static_cast(kNopeBytes / sizeof(bf16_t)); + constexpr int kQRopeDim = static_cast(kRopeBytes / sizeof(bf16_t)); + + // Per-warp smem slots for the KV scatter role; concat warps leave theirs idle. + __shared__ alignas(16) uint8_t smem[kNumWarps][kRowBytes]; + + const uint32_t warp_in_cta = threadIdx.x / kWarpThreads; + const uint32_t lane_id = threadIdx.x % kWarpThreads; + const uint32_t flat_warp = blockIdx.x * kNumWarps + warp_in_cta; + + PDLWaitPrimary(); + + if (flat_warp < params.batch_size) { + // --- KV scatter role: one warp per token (smem staging + TMA bulk store) --- + const uint32_t item_id = flat_warp; + const int64_t loc = static_cast(static_cast(params.loc)[item_id]); + + const auto nope_src = pointer::offset(params.k_nope, item_id * params.stride_nope_bytes); + const auto rope_src = pointer::offset(params.k_rope, item_id * params.stride_rope_bytes); + void* const gmem_dst = pointer::offset(params.kv_buffer, loc * params.stride_buffer_bytes); + + warp::copy_bytes(nope_src, &smem[warp_in_cta][0]); + warp::copy_bytes(rope_src, &smem[warp_in_cta][kNopeBytes]); + + // TMA reads smem via the async proxy; fence so it can't observe stale sts. + __syncwarp(); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + + // elect.sync rather than `lane_id == 0`: the TMA issue must not sit + // behind a lane-index predicate (see PR review). + if (device::warp::elect_one_lane()) { + cuda::ptx::cp_async_bulk( + cuda::ptx::space_global, + cuda::ptx::space_shared, + gmem_dst, + &smem[warp_in_cta][0], + static_cast(kRowBytes)); + } + + // ``wait_group`` (not ``_read``): waits for gmem commit, not just smem reuse. + cuda::ptx::cp_async_bulk_commit_group(); + cuda::ptx::cp_async_bulk_wait_group(cuda::ptx::n32_t<0>{}); + } else if (flat_warp - params.batch_size < params.num_q_items) { + // --- Q concat role: one warp per (token, head) row --- + const uint32_t q_item = flat_warp - params.batch_size; + const uint32_t idx_0 = q_item / params.q_dim_1; + const uint32_t idx_1 = q_item % params.q_dim_1; + + using ABufType = int4; + constexpr int kAVecElems = static_cast(sizeof(ABufType) / sizeof(bf16_t)); + constexpr int kANumUnroll = kQNopeDim / (kAVecElems * kWarpThreads); + static_assert(kANumUnroll * kAVecElems * kWarpThreads == kQNopeDim, "nope dim must fill whole int4 warp rounds"); + using BBufType = int; + constexpr int kBVecElems = static_cast(sizeof(BBufType) / sizeof(bf16_t)); + static_assert(kBVecElems * kWarpThreads == kQRopeDim, "rope dim must be exactly one int warp round"); + + const bf16_t* a_row = params.q_nope + idx_0 * params.qn_stride_0 + idx_1 * params.qn_stride_1; + const bf16_t* b_row = params.q_rope + idx_0 * params.qr_stride_0 + idx_1 * params.qr_stride_1; + bf16_t* o_row = params.q_out + idx_0 * params.qo_stride_0 + idx_1 * params.qo_stride_1; + + ABufType a_buf[kANumUnroll]; +#pragma unroll + for (int i = 0; i < kANumUnroll; ++i) { + a_buf[i] = reinterpret_cast(a_row)[i * kWarpThreads + lane_id]; + } + const BBufType b_buf = reinterpret_cast(b_row)[lane_id]; + +#pragma unroll + for (int i = 0; i < kANumUnroll; ++i) { + reinterpret_cast(o_row)[i * kWarpThreads + lane_id] = a_buf[i]; + } + reinterpret_cast(o_row + kQNopeDim)[lane_id] = b_buf; + } + + PDLTriggerSecondary(); +} + +template +struct SetMlaKVConcatQKernel { + static_assert(kNopeBytes > 0 && kNopeBytes % 4 == 0, "kNopeBytes must be a positive multiple of 4"); + static_assert(kRopeBytes > 0 && kRopeBytes % 4 == 0, "kRopeBytes must be a positive multiple of 4"); + static_assert( + (kNopeBytes + kRopeBytes) % 16 == 0, "TMA bulk store requires (kNopeBytes + kRopeBytes) to be a multiple of 16"); + + static constexpr int64_t kQNopeDim = kNopeBytes / static_cast(sizeof(bf16_t)); + static constexpr int64_t kQRopeDim = kRopeBytes / static_cast(sizeof(bf16_t)); + + template + static constexpr auto kernel = set_mla_kv_concat_q_kernel; + + static void + run(tvm::ffi::TensorView kv_buffer, + tvm::ffi::TensorView loc, + tvm::ffi::TensorView k_nope, + tvm::ffi::TensorView k_rope, + tvm::ffi::TensorView q_nope, + tvm::ffi::TensorView q_rope, + tvm::ffi::TensorView q_out, + int64_t num_warps_per_block) { + using namespace host; + + auto B = SymbolicSize{"batch_size"}; + auto H = SymbolicSize{"num_heads"}; + auto D_nope = SymbolicSize{"nope_dim"}; + auto D_rope = SymbolicSize{"rope_dim"}; + auto D_buf = SymbolicSize{"buffer_last_dim"}; + auto D_qn = SymbolicSize{"q_nope_dim"}; + auto D_qr = SymbolicSize{"q_rope_dim"}; + auto D_qo = SymbolicSize{"q_out_dim"}; + auto S_nope = SymbolicSize{"nope_stride"}; + auto S_rope = SymbolicSize{"rope_stride"}; + auto S_buf = SymbolicSize{"buffer_stride"}; + auto S_loc = SymbolicSize{"loc_stride"}; + auto S0_qn = SymbolicSize{"q_nope_stride_0"}; + auto S1_qn = SymbolicSize{"q_nope_stride_1"}; + auto S0_qr = SymbolicSize{"q_rope_stride_0"}; + auto S1_qr = SymbolicSize{"q_rope_stride_1"}; + auto S0_qo = SymbolicSize{"q_out_stride_0"}; + auto S1_qo = SymbolicSize{"q_out_stride_1"}; + auto loc_dtype = SymbolicDType{}; + auto device = SymbolicDevice{}; + device.set_options(); + + D_qn.set_value(kQNopeDim); + D_qr.set_value(kQRopeDim); + D_qo.set_value(kQNopeDim + kQRopeDim); + + TensorMatcher({B, D_nope}) // + .with_strides({S_nope, 1}) + .with_dtype() + .with_device(device) + .verify(k_nope); + TensorMatcher({B, D_rope}) // + .with_strides({S_rope, 1}) + .with_dtype() + .with_device(device) + .verify(k_rope); + TensorMatcher({-1, D_buf}) // + .with_strides({S_buf, 1}) + .with_dtype() + .with_device(device) + .verify(kv_buffer); + TensorMatcher({B}) // + .with_strides({S_loc}) + .with_dtype(loc_dtype) + .with_device(device) + .verify(loc); + TensorMatcher({B, H, D_qn}) // + .with_strides({S0_qn, S1_qn, 1}) + .with_dtype() + .with_device(device) + .verify(q_nope); + TensorMatcher({B, H, D_qr}) // + .with_strides({S0_qr, S1_qr, 1}) + .with_dtype() + .with_device(device) + .verify(q_rope); + TensorMatcher({B, H, D_qo}) // + .with_strides({S0_qo, S1_qo, 1}) + .with_dtype() + .with_device(device) + .verify(q_out); + + constexpr int64_t kDtypeSize = static_cast(sizeof(bf16_t)); + CHECK_HOST(kNopeBytes == kDtypeSize * D_nope.unwrap()) + << "kNopeBytes mismatch: expected " << kNopeBytes << ", got " << kDtypeSize * D_nope.unwrap(); + CHECK_HOST(kRopeBytes == kDtypeSize * D_rope.unwrap()) + << "kRopeBytes mismatch: expected " << kRopeBytes << ", got " << kDtypeSize * D_rope.unwrap(); + CHECK_HOST(kDtypeSize * D_buf.unwrap() >= kNopeBytes + kRopeBytes) << "kv_buffer last dim too small"; + CHECK_HOST(S_loc.unwrap() == 1) << "loc must be contiguous; got stride " << S_loc.unwrap(); + + // Alignment tripwires. The device code does 16-byte vector accesses on the + // kv row / nope rows / q rows and 4-byte accesses on the rope rows; the + // python-side ``covered()`` mirrors these so uncovered layouts fall back + // instead of faulting (do NOT assume "PyTorch tensors are aligned" — views + // and odd pool pitches break that). + const auto aligned = [](const void* ptr, int64_t align) { + return reinterpret_cast(ptr) % static_cast(align) == 0; + }; + CHECK_HOST(aligned(kv_buffer.data_ptr(), 16) && (S_buf.unwrap() * kDtypeSize) % 16 == 0) + << "kv_buffer base/row-stride must be 16-byte aligned for TMA bulk store"; + CHECK_HOST(aligned(k_nope.data_ptr(), 16) && (S_nope.unwrap() * kDtypeSize) % 16 == 0) + << "k_nope base/row-stride must be 16-byte aligned"; + CHECK_HOST(aligned(k_rope.data_ptr(), 4) && (S_rope.unwrap() * kDtypeSize) % 4 == 0) + << "k_rope base/row-stride must be 4-byte aligned"; + CHECK_HOST( + aligned(q_nope.data_ptr(), 16) && (S0_qn.unwrap() * kDtypeSize) % 16 == 0 && + (S1_qn.unwrap() * kDtypeSize) % 16 == 0) + << "q_nope base/strides must be 16-byte aligned"; + CHECK_HOST( + aligned(q_rope.data_ptr(), 4) && (S0_qr.unwrap() * kDtypeSize) % 4 == 0 && + (S1_qr.unwrap() * kDtypeSize) % 4 == 0) + << "q_rope base/strides must be 4-byte aligned"; + CHECK_HOST( + aligned(q_out.data_ptr(), 16) && (S0_qo.unwrap() * kDtypeSize) % 16 == 0 && + (S1_qo.unwrap() * kDtypeSize) % 16 == 0) + << "q_out base/strides must be 16-byte aligned"; + + const uint32_t batch = static_cast(B.unwrap()); + const uint32_t num_heads = static_cast(H.unwrap()); + if (batch == 0) return; + + const auto params = SetMlaKVConcatQParams{ + .k_nope = k_nope.data_ptr(), + .k_rope = k_rope.data_ptr(), + .kv_buffer = kv_buffer.data_ptr(), + .loc = loc.data_ptr(), + .stride_nope_bytes = S_nope.unwrap() * kDtypeSize, + .stride_rope_bytes = S_rope.unwrap() * kDtypeSize, + .stride_buffer_bytes = S_buf.unwrap() * kDtypeSize, + .batch_size = batch, + .q_nope = static_cast(q_nope.data_ptr()), + .q_rope = static_cast(q_rope.data_ptr()), + .q_out = static_cast(q_out.data_ptr()), + .num_q_items = batch * num_heads, + .q_dim_1 = num_heads, + .qn_stride_0 = S0_qn.unwrap(), + .qn_stride_1 = static_cast(S1_qn.unwrap()), + .qr_stride_0 = S0_qr.unwrap(), + .qr_stride_1 = static_cast(S1_qr.unwrap()), + .qo_stride_0 = S0_qo.unwrap(), + .qo_stride_1 = static_cast(S1_qo.unwrap()), + }; + + const auto use_int32 = loc_dtype.is_type(); + const uint32_t total_warps = params.batch_size + params.num_q_items; + + auto launch = [&]() { + const auto kernel_ptr = use_int32 ? kernel : kernel; + const uint32_t num_blocks = div_ceil(total_warps, static_cast(kNW)); + const uint32_t threads_per_block = static_cast(kNW) * device::kWarpThreads; + LaunchKernel(num_blocks, threads_per_block, device.unwrap()) // + .enable_pdl(kUsePDL)(kernel_ptr, params); + }; + + switch (num_warps_per_block) { + case 1: + launch.template operator()<1>(); + break; + case 2: + launch.template operator()<2>(); + break; + case 4: + launch.template operator()<4>(); + break; + case 8: + launch.template operator()<8>(); + break; + default: + Panic("Unsupported num_warps_per_block=", num_warps_per_block); + } + } +}; + +// --------------------------------------------------------------------------- +// fp8 variant. Shares the translation unit, not the kernel: dims are runtime +// rather than template parameters, it shards DCP slots (vloc % world != rank), +// converts per lane instead of bulk-copying, and counts strides in elements. +// Only the module that instantiates it pays for it. +// --------------------------------------------------------------------------- +constexpr int kFp8NopeDim = 512; +constexpr int kFp8RopeDim = 64; +constexpr int kFp8RowBytes = kFp8NopeDim + kFp8RopeDim; // fp8: 1 byte/elem + +struct SetMlaKVConcatQFp8Params { + // KV quantize + scatter side. + const bf16_t* __restrict__ k_nope; + const bf16_t* __restrict__ k_rope; + uint8_t* __restrict__ kv_buffer; + const void* __restrict__ loc; + int64_t stride_nope; // elements + int64_t stride_rope; // elements + int64_t stride_buffer_bytes; // bytes + uint32_t batch_size; + // DCP cyclic sharding of the KV pool: ``loc`` is VIRTUAL; the physical + // row on the owner rank is loc / world, and only the owner + // (loc % world == rank) writes. world=1/rank=0 = identity (non-DCP). + int32_t dcp_world_size; + int32_t dcp_rank; + // Q quantize + concat side. + const bf16_t* __restrict__ q_nope; + const bf16_t* __restrict__ q_rope; + uint8_t* __restrict__ q_out; + uint32_t num_q_items; // batch_size * num_heads + uint32_t q_dim_1; // num_heads + int64_t qn_stride_0; + int32_t qn_stride_1; + int64_t qr_stride_0; + int32_t qr_stride_1; + int64_t qo_stride_0; // elements (== bytes for fp8) + int32_t qo_stride_1; +}; + +// 2x bf16 -> 2x fp8 e4m3, float-mediated cvt.rn NOSAT (matches aten: overflow -> NaN). +SGL_DEVICE uint16_t bf16x2_to_fp8x2(const bf16x2_t v) { + const float2 f = __bfloat1622float2(v); + return __nv_cvt_float2_to_fp8x2(f, __NV_NOSAT, __NV_E4M3); +} + +// Convert 8 bf16 (one int4 load) to 8 fp8 packed in a uint2. +SGL_DEVICE uint2 bf16x8_to_fp8x8(const int4 v) { + const bf16x2_t* p = reinterpret_cast(&v); + uint2 out; + out.x = static_cast(bf16x2_to_fp8x2(p[0])) | (static_cast(bf16x2_to_fp8x2(p[1])) << 16); + out.y = static_cast(bf16x2_to_fp8x2(p[2])) | (static_cast(bf16x2_to_fp8x2(p[3])) << 16); + return out; +} + +template +__global__ void set_mla_kv_concat_q_fp8_kernel(const __grid_constant__ SetMlaKVConcatQFp8Params params) { + using namespace device; + + // Per-warp smem slots for the KV role (fp8 rows); concat warps leave + // theirs idle. 576 % 16 == 0 satisfies the TMA bulk-store requirement. + __shared__ alignas(16) uint8_t smem[kNumWarps][kFp8RowBytes]; + + const uint32_t warp_in_cta = threadIdx.x / kWarpThreads; + const uint32_t lane_id = threadIdx.x % kWarpThreads; + const uint32_t flat_warp = blockIdx.x * kNumWarps + warp_in_cta; + + PDLWaitPrimary(); + + if (flat_warp < params.batch_size) { + // --- KV role: quantize one token's row into smem, TMA-scatter it --- + const uint32_t item_id = flat_warp; + const int64_t vloc = static_cast(static_cast(params.loc)[item_id]); + // DCP ownership: non-owner ranks write nothing for this token (mirrors + // the triton writer's is_valid mask + loc // world translation). + if (vloc % params.dcp_world_size != params.dcp_rank) { + PDLTriggerSecondary(); + return; + } + const int64_t loc = vloc / params.dcp_world_size; + const bf16_t* nope_src = params.k_nope + item_id * params.stride_nope; + const bf16_t* rope_src = params.k_rope + item_id * params.stride_rope; + + // nope: 512 bf16 -> 512 fp8; 16 elems/lane (2 int4 loads -> 1 int4 store). + { + const int4* src = reinterpret_cast(nope_src); + uint2 lo = bf16x8_to_fp8x8(src[lane_id * 2]); + uint2 hi = bf16x8_to_fp8x8(src[lane_id * 2 + 1]); + reinterpret_cast(&smem[warp_in_cta][0])[lane_id] = + make_int4(static_cast(lo.x), static_cast(lo.y), static_cast(hi.x), static_cast(hi.y)); + } + // rope: 64 bf16 -> 64 fp8; 2 elems/lane. + { + const bf16x2_t v = reinterpret_cast(rope_src)[lane_id]; + reinterpret_cast(&smem[warp_in_cta][kFp8NopeDim])[lane_id] = bf16x2_to_fp8x2(v); + } + + // TMA reads smem via the async proxy; fence so it can't observe stale sts. + __syncwarp(); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + + // elect.sync rather than `lane_id == 0`: the TMA issue must not sit + // behind a lane-index predicate (same review point as the bf16 variant). + if (device::warp::elect_one_lane()) { + cuda::ptx::cp_async_bulk( + cuda::ptx::space_global, + cuda::ptx::space_shared, + params.kv_buffer + loc * params.stride_buffer_bytes, + &smem[warp_in_cta][0], + static_cast(kFp8RowBytes)); + } + // ``wait_group`` (not ``_read``): waits for gmem commit, not just smem reuse. + cuda::ptx::cp_async_bulk_commit_group(); + cuda::ptx::cp_async_bulk_wait_group(cuda::ptx::n32_t<0>{}); + } else if (flat_warp - params.batch_size < params.num_q_items) { + // --- Q role: quantize one (token, head) row into the fp8 query --- + const uint32_t q_item = flat_warp - params.batch_size; + const uint32_t idx_0 = q_item / params.q_dim_1; + const uint32_t idx_1 = q_item % params.q_dim_1; + const bf16_t* a_row = params.q_nope + idx_0 * params.qn_stride_0 + idx_1 * params.qn_stride_1; + const bf16_t* b_row = params.q_rope + idx_0 * params.qr_stride_0 + idx_1 * params.qr_stride_1; + uint8_t* o_row = params.q_out + idx_0 * params.qo_stride_0 + idx_1 * params.qo_stride_1; + + { + const int4* src = reinterpret_cast(a_row); + uint2 lo = bf16x8_to_fp8x8(src[lane_id * 2]); + uint2 hi = bf16x8_to_fp8x8(src[lane_id * 2 + 1]); + reinterpret_cast(o_row)[lane_id] = + make_int4(static_cast(lo.x), static_cast(lo.y), static_cast(hi.x), static_cast(hi.y)); + } + { + const bf16x2_t v = reinterpret_cast(b_row)[lane_id]; + reinterpret_cast(o_row + kFp8NopeDim)[lane_id] = bf16x2_to_fp8x2(v); + } + } + + PDLTriggerSecondary(); +} + +template +struct SetMlaKVConcatQFp8Kernel { + template + static constexpr auto kernel = set_mla_kv_concat_q_fp8_kernel; + + static void + run(tvm::ffi::TensorView kv_buffer, + tvm::ffi::TensorView loc, + tvm::ffi::TensorView k_nope, + tvm::ffi::TensorView k_rope, + tvm::ffi::TensorView q_nope, + tvm::ffi::TensorView q_rope, + tvm::ffi::TensorView q_out, + int64_t num_warps_per_block, + int64_t dcp_world_size, + int64_t dcp_rank) { + using namespace host; + + auto B = SymbolicSize{"batch_size"}; + auto H = SymbolicSize{"num_heads"}; + auto D_nope = SymbolicSize{"nope_dim"}; + auto D_rope = SymbolicSize{"rope_dim"}; + auto D_buf = SymbolicSize{"buffer_last_dim"}; + auto D_qn = SymbolicSize{"q_nope_dim"}; + auto D_qr = SymbolicSize{"q_rope_dim"}; + auto D_qo = SymbolicSize{"q_out_dim"}; + auto S_nope = SymbolicSize{"nope_stride"}; + auto S_rope = SymbolicSize{"rope_stride"}; + auto S_buf = SymbolicSize{"buffer_stride"}; + auto S_loc = SymbolicSize{"loc_stride"}; + auto S0_qn = SymbolicSize{"q_nope_stride_0"}; + auto S1_qn = SymbolicSize{"q_nope_stride_1"}; + auto S0_qr = SymbolicSize{"q_rope_stride_0"}; + auto S1_qr = SymbolicSize{"q_rope_stride_1"}; + auto S0_qo = SymbolicSize{"q_out_stride_0"}; + auto S1_qo = SymbolicSize{"q_out_stride_1"}; + auto loc_dtype = SymbolicDType{}; + auto device = SymbolicDevice{}; + device.set_options(); + + D_nope.set_value(kFp8NopeDim); + D_rope.set_value(kFp8RopeDim); + D_qn.set_value(kFp8NopeDim); + D_qr.set_value(kFp8RopeDim); + D_qo.set_value(kFp8RowBytes); + + TensorMatcher({B, D_nope}) // + .with_strides({S_nope, 1}) + .with_dtype() + .with_device(device) + .verify(k_nope); + TensorMatcher({B, D_rope}) // + .with_strides({S_rope, 1}) + .with_dtype() + .with_device(device) + .verify(k_rope); + TensorMatcher({-1, D_buf}) // + .with_strides({S_buf, 1}) + .with_dtype() + .with_device(device) + .verify(kv_buffer); + TensorMatcher({B}) // + .with_strides({S_loc}) + .with_dtype(loc_dtype) + .with_device(device) + .verify(loc); + TensorMatcher({B, H, D_qn}) // + .with_strides({S0_qn, S1_qn, 1}) + .with_dtype() + .with_device(device) + .verify(q_nope); + TensorMatcher({B, H, D_qr}) // + .with_strides({S0_qr, S1_qr, 1}) + .with_dtype() + .with_device(device) + .verify(q_rope); + TensorMatcher({B, H, D_qo}) // + .with_strides({S0_qo, S1_qo, 1}) + .with_dtype() + .with_device(device) + .verify(q_out); + + CHECK_HOST(D_buf.unwrap() >= kFp8RowBytes) << "kv_buffer last dim too small"; + CHECK_HOST(dcp_world_size >= 1 && dcp_rank >= 0 && dcp_rank < dcp_world_size) + << "invalid dcp world/rank: " << dcp_world_size << "/" << dcp_rank; + CHECK_HOST(S_loc.unwrap() == 1) << "loc must be contiguous; got stride " << S_loc.unwrap(); + + // Alignment tripwires (mirrored by python covered() so uncovered layouts + // fall back instead of faulting): 16B vector loads on the bf16 nope/q + // rows, 4B on the rope rows, 16B TMA dst rows, 16B int4 stores on q_out. + const auto aligned = [](const void* ptr, int64_t align) { + return reinterpret_cast(ptr) % static_cast(align) == 0; + }; + CHECK_HOST(aligned(kv_buffer.data_ptr(), 16) && S_buf.unwrap() % 16 == 0) + << "kv_buffer base/row-stride must be 16-byte aligned for TMA bulk store"; + CHECK_HOST(aligned(k_nope.data_ptr(), 16) && (S_nope.unwrap() * 2) % 16 == 0) + << "k_nope base/row-stride must be 16-byte aligned"; + CHECK_HOST(aligned(k_rope.data_ptr(), 4) && (S_rope.unwrap() * 2) % 4 == 0) + << "k_rope base/row-stride must be 4-byte aligned"; + CHECK_HOST(aligned(q_nope.data_ptr(), 16) && (S0_qn.unwrap() * 2) % 16 == 0 && (S1_qn.unwrap() * 2) % 16 == 0) + << "q_nope base/strides must be 16-byte aligned"; + CHECK_HOST(aligned(q_rope.data_ptr(), 4) && (S0_qr.unwrap() * 2) % 4 == 0 && (S1_qr.unwrap() * 2) % 4 == 0) + << "q_rope base/strides must be 4-byte aligned"; + CHECK_HOST(aligned(q_out.data_ptr(), 16) && S0_qo.unwrap() % 16 == 0 && S1_qo.unwrap() % 16 == 0) + << "q_out base/strides must be 16-byte aligned"; + + const uint32_t batch = static_cast(B.unwrap()); + const uint32_t num_heads = static_cast(H.unwrap()); + if (batch == 0) return; + + const auto params = SetMlaKVConcatQFp8Params{ + .k_nope = static_cast(k_nope.data_ptr()), + .k_rope = static_cast(k_rope.data_ptr()), + .kv_buffer = static_cast(kv_buffer.data_ptr()), + .loc = loc.data_ptr(), + .stride_nope = S_nope.unwrap(), + .stride_rope = S_rope.unwrap(), + .stride_buffer_bytes = S_buf.unwrap(), + .batch_size = batch, + .dcp_world_size = static_cast(dcp_world_size), + .dcp_rank = static_cast(dcp_rank), + .q_nope = static_cast(q_nope.data_ptr()), + .q_rope = static_cast(q_rope.data_ptr()), + .q_out = static_cast(q_out.data_ptr()), + .num_q_items = batch * num_heads, + .q_dim_1 = num_heads, + .qn_stride_0 = S0_qn.unwrap(), + .qn_stride_1 = static_cast(S1_qn.unwrap()), + .qr_stride_0 = S0_qr.unwrap(), + .qr_stride_1 = static_cast(S1_qr.unwrap()), + .qo_stride_0 = S0_qo.unwrap(), + .qo_stride_1 = static_cast(S1_qo.unwrap()), + }; + + const auto use_int32 = loc_dtype.is_type(); + const uint32_t total_warps = params.batch_size + params.num_q_items; + + auto launch = [&]() { + const auto kernel_ptr = use_int32 ? kernel : kernel; + const uint32_t num_blocks = div_ceil(total_warps, static_cast(kNW)); + LaunchKernel(num_blocks, static_cast(kNW) * device::kWarpThreads, device.unwrap()) + .enable_pdl(kUsePDL)(kernel_ptr, params); + }; + + switch (num_warps_per_block) { + case 1: + launch.template operator()<1>(); + break; + case 2: + launch.template operator()<2>(); + break; + case 4: + launch.template operator()<4>(); + break; + case 8: + launch.template operator()<8>(); + break; + default: + Panic("Unsupported num_warps_per_block=", num_warps_per_block); + } + } +}; + +} // namespace diff --git a/python/sglang/kernels/jit/csrc/gemm/per_token_group_quant.cuh b/python/sglang/kernels/jit/csrc/gemm/per_token_group_quant.cuh index f8ecbe3f2..4b4ac822f 100644 --- a/python/sglang/kernels/jit/csrc/gemm/per_token_group_quant.cuh +++ b/python/sglang/kernels/jit/csrc/gemm/per_token_group_quant.cuh @@ -408,11 +408,22 @@ QuantHostContext build_quant_context( // TensorMatcher({E, N, -1}).with_strides({-1, -1, 1}).with_dtype().with_device(device).verify(input); TensorMatcher({E, N, H}).with_strides({-1, -1, 1}).with_dtype().with_device(device).verify(output_q); TensorMatcher({E, N, G}).with_strides({-1, -1, -1}).with_dtype().with_device(device).verify(output_s); + CHECK_HOST((input.stride(0) * sizeof(T)) % 32 == 0) + << "input expert stride must keep rows 32B-aligned for the vectorized loads"; } else { TensorMatcher({N, -1}).with_strides({-1, 1}).with_dtype().with_device(device).verify(input); TensorMatcher({N, H}).with_strides({-1, 1}).with_dtype().with_device(device).verify(output_q); TensorMatcher({N, G}).with_strides({-1, -1}).with_dtype().with_device(device).verify(output_s); } + // The 32B/lane vectorized loads need every input row to start 32B-aligned + // (kMaxVecBytes on Blackwell; over-strict but harmless on Hopper, whose + // 16B vectors only need 16). Contiguous allocations always satisfy this; it + // only bites hand-made row-strided views, which must keep rows aligned -- + // rejected loudly here rather than densified silently at the call site. + CHECK_HOST(reinterpret_cast(input.data_ptr()) % 32 == 0) + << "input base pointer must be 32B-aligned for the vectorized loads"; + CHECK_HOST((input.stride(-2) * sizeof(T)) % 32 == 0) + << "input token stride must keep rows 32B-aligned for the vectorized loads"; const uint32_t num_tokens = N.unwrap(); const uint32_t hidden_size = H.unwrap(); diff --git a/python/sglang/kernels/jit/csrc/gemm/tiny_gemm.cuh b/python/sglang/kernels/jit/csrc/gemm/tiny_gemm.cuh new file mode 100644 index 000000000..eb25e7cd3 --- /dev/null +++ b/python/sglang/kernels/jit/csrc/gemm/tiny_gemm.cuh @@ -0,0 +1,231 @@ +#include +#include + +#include +#include +#include +#include + +#include + +#include +#include +#include +#include + +namespace sglang { + +using namespace device; + +constexpr uint32_t kTinyNGemmVecSize = kMaxVecBytes / sizeof(bf16_t); + +template +__global__ __launch_bounds__(K / kTinyNGemmVecSize, 1) // 1 block per SM + void tiny_n_gemm_kernel(OutT* __restrict__ out, const bf16_t* __restrict__ x, const bf16_t* __restrict__ w) { + constexpr uint32_t kBlockSize = K / kTinyNGemmVecSize; + constexpr uint32_t kNumWarps = kBlockSize / kWarpThreads; + static_assert(M * N_SPLIT <= kBlockSize, "output tile must fit one thread each for the final reduce"); + using vec_t = AlignedVector; + + const uint32_t bx = blockIdx.x; + const uint32_t tx = threadIdx.x; + const bf16_t* w_tile = w + bx * (N_SPLIT * K); + + // Weight prefetch: address is input-independent, load before the PDL wait. + vec_t wv[N_SPLIT]; +#pragma unroll + for (uint32_t n = 0; n < N_SPLIT; ++n) { + wv[n].load(w_tile + n * K, tx); + } + + PDLWaitPrimary(); + + vec_t xv[M]; +#pragma unroll + for (uint32_t m = 0; m < M; ++m) { + xv[m].load(x + m * K, tx); + } + + __shared__ float s_acc[kNumWarps][M * N_SPLIT]; + const uint32_t warp_id = tx / kWarpThreads; + +#pragma unroll + for (uint32_t m = 0; m < M; ++m) { +#pragma unroll + for (uint32_t n = 0; n < N_SPLIT; ++n) { + float acc = 0.0f; +#if SGL_ARCH_BLACKWELL_OR_GREATER +#pragma unroll + for (uint32_t i = 0; i < kTinyNGemmVecSize; ++i) { + acc = device::math::fma_f32_bf16(xv[m][i], wv[n][i], acc); + } +#else + for (uint32_t i = 0; i < kTinyNGemmVecSize / 2; ++i) { + const auto [x0, x1] = cast(bf16x2_t{xv[m][2 * i], xv[m][2 * i + 1]}); + const auto [w0, w1] = cast(bf16x2_t{wv[n][2 * i], wv[n][2 * i + 1]}); + acc = fmaf(x0, w0, acc); + acc = fmaf(x1, w1, acc); + } +#endif + // NOTE: broadcast write (all lanes hold the reduced value), safe here. + s_acc[warp_id][m * N_SPLIT + n] = warp::reduce_sum(acc); + } + } + PDLTriggerSecondary(); + __syncthreads(); + + if (tx < M * N_SPLIT) { + float acc[kNumWarps]; +#pragma unroll + for (uint32_t i = 0; i < kNumWarps; ++i) { + acc[i] = s_acc[i][tx]; + } +#pragma unroll + for (uint32_t i = 1; i < kNumWarps; ++i) { + acc[0] += acc[i]; + } + const uint32_t m = tx / N_SPLIT; + const uint32_t n = tx % N_SPLIT; + out[m * N + bx * N_SPLIT + n] = cast(acc[0]); + } +} + +SGL_DEVICE void cp_async_cg_16(void* smem_dst, const void* gmem_src, int32_t vec_offset) { + const uint32_t offset = static_cast(vec_offset * 16); +#if defined(USE_ROCM) + *reinterpret_cast(static_cast(smem_dst) + offset) = + *reinterpret_cast(static_cast(gmem_src) + offset); +#else + const uint32_t smem_addr = static_cast(__cvta_generic_to_shared(smem_dst)) + offset; + const uint64_t gmem_addr = static_cast(__cvta_generic_to_global(gmem_src)) + offset; + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" : : "r"(smem_addr), "l"(gmem_addr) : "memory"); +#endif +} + +constexpr uint32_t kTinyKGemmVecSize = 16 / sizeof(bf16_t); // NOTE: no need to be large + +template +__global__ __launch_bounds__(N_SPLIT* K / kTinyKGemmVecSize, 1) // control the block size + void tiny_k_gemm_kernel( + OutT* __restrict__ out, const bf16_t* __restrict__ x, const bf16_t* __restrict__ w, const int64_t dx) { + using vec_t = AlignedVector; + constexpr uint32_t kNumKLanes = K / kTinyKGemmVecSize; + static_assert(std::has_single_bit(kNumKLanes), "K / vec_size must be a power of 2"); + static_assert(kNumKLanes <= kWarpThreads, "require in-warp reduction"); + static_assert((N_SPLIT * K / kTinyKGemmVecSize) % kWarpThreads == 0); + const uint32_t bx = blockIdx.x; + const uint32_t tx = threadIdx.x; + const uint32_t n_idx = bx * N_SPLIT + tx / kNumKLanes; + const uint32_t work_id = tx % kNumKLanes; + const bf16_t* w_tile = w + n_idx * K; + + // Weight prefetch: address is input-independent, load before the PDL wait. + vec_t wv; + wv.load(w_tile, work_id); + + PDLWaitPrimary(); + vec_t xv[M]; +#pragma unroll + for (uint32_t m = 0; m < M; ++m) { + xv[m].load(x + m * dx, work_id); + } + +#pragma unroll + for (uint32_t m = 0; m < M; ++m) { + float acc = 0.0f; +#pragma unroll + for (uint32_t i = 0; i < kTinyKGemmVecSize; ++i) { + acc = device::math::fma_f32_bf16(xv[m][i], wv[i], acc); + } + // Broadcast store: every lane of the group holds the reduced sum. + out[m * N + n_idx] = cast(warp::reduce_sum(acc)); + } + PDLTriggerSecondary(); +} + +} // namespace sglang + +using namespace sglang; + +template +struct TinyNGemmKernel { + static constexpr uint32_t kBlockSize = K / kTinyNGemmVecSize; + static constexpr uint32_t kNumBlocks = N / N_SPLIT; + static_assert(K % kTinyNGemmVecSize == 0, "K must be divisible by the vector width"); + static_assert(kBlockSize % kWarpThreads == 0, "K / vec_size must be a multiple of the warp size"); + static_assert(kBlockSize <= 1024, "K / vec_size exceeds the maximum block size"); + static_assert(N % N_SPLIT == 0, "N must be divisible by split_n"); + static_assert(kMaxM * N_SPLIT <= kBlockSize, "max_m * split_n must fit one thread each for the final reduce"); + + using KernelFn = void (*)(OutT*, const bf16_t*, const bf16_t*); + + template + static constexpr auto make_table(std::index_sequence) { + return std::array{nullptr, tiny_n_gemm_kernel...}; + } + static constexpr auto kTable = make_table(std::make_index_sequence{}); + + static void run(const tvm::ffi::TensorView x, const tvm::ffi::TensorView w, const tvm::ffi::TensorView out) { + using namespace host; + + auto M = SymbolicSize{"num_tokens"}; + auto device = SymbolicDevice{}; + device.set_options(); + TensorMatcher({M, K}).with_dtype().with_device(device).verify(x); + TensorMatcher({N, K}).with_dtype().with_device(device).verify(w); + TensorMatcher({M, N}).with_dtype().with_device(device).verify(out); + const auto num_tokens = static_cast(M.unwrap()); + RuntimeCheck(num_tokens >= 1 && num_tokens <= kMaxM); + LaunchKernel(kNumBlocks, kBlockSize, device.unwrap()) + .enable_pdl(kUsePDL)( + kTable[num_tokens], + static_cast(out.data_ptr()), + static_cast(x.data_ptr()), + static_cast(w.data_ptr())); + } +}; + +template +struct TinyKGemmKernel { + static constexpr uint32_t kNumKLanes = K / kTinyKGemmVecSize; + static constexpr uint32_t kBlockSize = N_SPLIT * kNumKLanes; + static constexpr uint32_t kNumBlocks = N / N_SPLIT; + static_assert(K % kTinyKGemmVecSize == 0, "K must be divisible by the vector width"); + static_assert(N % N_SPLIT == 0, "N must be divisible by split_n"); + static_assert(kBlockSize % kWarpThreads == 0, "split_n * K-lanes must fill whole warps"); + static_assert(kBlockSize <= 1024, "split_n * K-lanes exceeds the maximum block size"); + + using KernelFn = void (*)(OutT*, const bf16_t*, const bf16_t*, int64_t); + + template + static constexpr auto make_table(std::index_sequence) { + return std::array{nullptr, tiny_k_gemm_kernel...}; + } + static constexpr auto kTable = make_table(std::make_index_sequence{}); + + static void run(const tvm::ffi::TensorView x, const tvm::ffi::TensorView w, const tvm::ffi::TensorView out) { + using namespace host; + + auto M = SymbolicSize{"num_tokens"}; + auto device = SymbolicDevice{}; + device.set_options(); + // x may be a row-sliced view of a wider fused buffer: allow stride != K. + TensorMatcher({M, K}).with_dtype().with_strides({-1, 1}).with_device(device).verify(x); + TensorMatcher({N, K}).with_dtype().with_device(device).verify(w); + TensorMatcher({M, N}).with_dtype().with_device(device).verify(out); + const auto num_tokens = static_cast(M.unwrap()); + const auto x_stride = static_cast(x.stride(0)); + RuntimeCheck(num_tokens >= 1 && num_tokens <= kMaxM); + RuntimeCheck( + x_stride * sizeof(bf16_t) % (kTinyKGemmVecSize * sizeof(bf16_t)) == 0, + "x rows must stay aligned to the vector width, got stride ", + x_stride); + LaunchKernel(kNumBlocks, kBlockSize, device.unwrap()) + .enable_pdl(kUsePDL)( + kTable[num_tokens], + static_cast(out.data_ptr()), + static_cast(x.data_ptr()), + static_cast(w.data_ptr()), + x_stride); + } +}; diff --git a/python/sglang/kernels/jit/csrc/kimi_k3/attn_res/fused_tma.cuh b/python/sglang/kernels/jit/csrc/kimi_k3/attn_res/fused_tma.cuh new file mode 100644 index 000000000..3e18f3a69 --- /dev/null +++ b/python/sglang/kernels/jit/csrc/kimi_k3/attn_res/fused_tma.cuh @@ -0,0 +1,946 @@ +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "../../distributed/custom_all_reduce.cuh" +#include +#include +#include +#include +#include + +// Local PTX primitives (mbarrier / bulk TMA / tcgen05 / warp-group sync) + +namespace ptx { + +// ---- bulk 1D TMA (PTX ISA §9.7.9.25) --------------------------------------- + +// global -> shared::cluster, completed by an smem mbarrier. Arm `bar` with +// `mbar_arrive_expect_tx(bar, bytes)` before issuing; `bytes` and both +// endpoints must be 16-byte aligned. +static SGL_DEVICE void cp_async_bulk_1d_load(void* smem_dst, const void* gmem_src, uint32_t bytes, uint64_t* bar) { + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes" + " [%0], [%1], %2, [%3];" ::"r"(to_shared(smem_dst)), + "l"(gmem_src), + "r"(bytes), + "r"(to_shared(bar)) + : "memory"); +} + +// Publish mbarrier initialization from the generic proxy before an async engine +// uses the barrier. +static SGL_DEVICE void fence_mbarrier_init() { + asm volatile("fence.mbarrier_init.release.cluster;"); +} + +// ---- warp / warp-group sync (PTX ISA §9.7.4, §9.7.12.6, §9.7.13) ----------- + +// Partial-CTA rendezvous. `id` must be in [1, 15]; barrier 0 is reserved for +// the full-CTA barrier behind __syncthreads(). +static SGL_DEVICE void named_barrier_sync(uint32_t id, uint32_t num_threads) { + asm volatile("bar.sync %0, %1;" ::"r"(id), "r"(num_threads) : "memory"); +} + +// True on exactly one lane of the issuing warp — guards single-issuer sites +// (mbar init, TMA issue, MMA issue, TMEM alloc) without gating on lane_id. +static SGL_DEVICE bool elect_one() { + uint32_t pred; + asm volatile( + "{\n\t.reg .pred p;\n\t" + "elect.sync _|p, 0xffffffff;\n\t" + "selp.b32 %0, 1, 0, p;\n\t}\n" + : "=r"(pred)); + return pred != 0; +} + +// Runtime warp-group register-budget reallocation: widen the epilogue's +// per-thread budget (so it holds a larger primary array without spilling) by +// narrowing the mainloop warps, which need few registers. +// +// Both forms are `.sync.aligned`: all 128 threads of the issuing warp-group +// must execute the SAME instruction with the SAME N, from a warp-group +// boundary (issuing from only one warp of the group hangs). N in [24, 256], +// multiple of 8, per thread; the CTA total must satisfy +// sum(warp_group_threads * N) <= 64512 (the safe allocatable cap on B100/B300 +// after ~1024 reserved regs). Caller owns the budgeting — there is no +// compile-time check, since N per warp-group is orthogonal. +// +// This pays only for ASYMMETRIC budgets that ptxas cannot infer from the +// source. For a symmetric cap, `__launch_bounds__(NUM_THREADS, 1)` is cleaner +// and measured faster on B100/B300. +template +static SGL_DEVICE void setmaxnreg_dec() { + static_assert(N >= 24 && N <= 256, "setmaxnreg N must be in [24, 256]"); + static_assert((N & 7) == 0, "setmaxnreg N must be a multiple of 8"); + asm volatile("setmaxnreg.dec.sync.aligned.u32 %0;\n" ::"n"(N)); +} + +template +static SGL_DEVICE void setmaxnreg_inc() { + static_assert(N >= 24 && N <= 256, "setmaxnreg N must be in [24, 256]"); + static_assert((N & 7) == 0, "setmaxnreg N must be a multiple of 8"); + asm volatile("setmaxnreg.inc.sync.aligned.u32 %0;\n" ::"n"(N)); +} + +// ---- tcgen05 (PTX ISA §9.7.16) --------------------------------------------- +// +// Lifecycle (mandatory order, §9.7.16.7.1): alloc (one warp, n_cols a power of +// 2 in [32, 512], TMEM address written to smem) -> __syncthreads + read taddr +// -> ld/st -> dealloc -> relinquish before kernel exit. +// +// Each warp can only touch its own 32-lane TMEM band (§9.7.16.8.1): warp 0 -> +// lanes 0-31, warp 1 -> 32-63, and so on. Use `tcgen05_wait_st` / +// `tcgen05_wait_ld` before consuming the other side of a store / drain. +static SGL_DEVICE void tcgen05_alloc(uint32_t smem_addr_for_taddr, uint32_t n_cols) { + asm volatile( + "tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" ::"r"(smem_addr_for_taddr), "r"(n_cols)); +} + +static SGL_DEVICE void tcgen05_dealloc(uint32_t taddr, uint32_t n_cols) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" ::"r"(taddr), "r"(n_cols)); +} + +static SGL_DEVICE void tcgen05_relinquish() { + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); +} + +// .32x32b.x8: 8 b32 per lane = 8 TMEM columns. Per-lane 8 FP32 -> 4 bf16x2 +// packs = one int4, the natural fit for a BF16 epilogue moving a column band +// with 16-byte smem accesses. +static SGL_DEVICE void tcgen05_ld_32x32b_x8( + uint32_t taddr, + uint32_t& r0, + uint32_t& r1, + uint32_t& r2, + uint32_t& r3, + uint32_t& r4, + uint32_t& r5, + uint32_t& r6, + uint32_t& r7) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x8.b32 " + " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3), "=r"(r4), "=r"(r5), "=r"(r6), "=r"(r7) + : "r"(taddr)); +} + +static SGL_DEVICE void tcgen05_ld_32x32b_x8(uint32_t taddr, uint32_t* dst) { + tcgen05_ld_32x32b_x8(taddr, dst[0], dst[1], dst[2], dst[3], dst[4], dst[5], dst[6], dst[7]); +} + +static SGL_DEVICE void tcgen05_st_32x32b_x8(uint32_t taddr, const uint32_t* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x8.b32 " + " [%8], {%0, %1, %2, %3, %4, %5, %6, %7};" + : + : "r"(src[0]), + "r"(src[1]), + "r"(src[2]), + "r"(src[3]), + "r"(src[4]), + "r"(src[5]), + "r"(src[6]), + "r"(src[7]), + "r"(taddr)); +} + +static SGL_DEVICE void tcgen05_wait_st() { + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); +} + +} // namespace ptx + +namespace sglang { + +struct AttnResTMAParams { + const bf16_t* __restrict__ prefix_sum; // [T, H] + const bf16_t* __restrict__ bank; // [T, NB_total, H] + const bf16_t* __restrict__ cw; // [H] score norm * proj weight + const bf16_t* __restrict__ ow; // [H] out norm weight + bf16_t* __restrict__ out; // [T, H] + // Fused bank write (nullptr = off): per-token destination of the prefix + // row snapshot, bank row nvb (strided by stride_bm like the read rows). + // The kernel never reads row nvb, so the write races with nothing. + bf16_t* __restrict__ prefix_dst; + // Optional fused NVLS reduce-scatter source. When input_mc is non-null, + // the producer warp materializes this rank's reduced token shard plus the + // local residual into prefix_sum/prefix_out before TMA consumes it. + const uint8_t* input_mc; + const bf16_t* residual; + bf16_t* prefix_out; + device::distributed::Semaphore* sem_local; + uint8_t* sem_mc; + uint8_t* output_mc; + uint32_t world_size; + uint32_t rank; + int64_t stride_bm; // bank stride along T (in elements) + float eps; + uint32_t num_tokens; +}; + +template +struct KimiK3AttnResTrait { + public: + static constexpr int64_t kDim = kDim_; + static constexpr int64_t kTile = 1024; // one warp-group-wide 16B sweep + static constexpr uint32_t kNumRows = kNumBankRows_; // bank rows; +1 prefix row + static constexpr uint32_t kChunkRows = kChunkRows_; // rows per chunk (one barrier pair per chunk) + // Chunk slots in the smem ring. Frozen at 2 (double buffering): 1 stalls + // the producer behind the consumers (~10% slower), >2 gains nothing and + // costs smem at small T. + static constexpr uint32_t kNumStages = 2; + static constexpr uint32_t kNumChunks = (kNumRows + 1 + kChunkRows - 1) / kChunkRows; + static constexpr uint32_t kNumConsumerWarps = 8; + static constexpr uint32_t kConsumerRegs = kConsumerRegs_; + static constexpr uint32_t kProducerRegs = 40; + static constexpr uint32_t kNumProducerWarps = kConsumerRegs > 0 ? 4 : 1; + static constexpr uint32_t kNumWarps = kNumConsumerWarps + kNumProducerWarps; + static constexpr uint32_t kNumThreads = kNumWarps * device::kWarpThreads; + static constexpr uint32_t kNumConsumerThreads = kNumConsumerWarps * device::kWarpThreads; + static_assert( + kConsumerRegs == 0 || (kConsumerRegs % 8 == 0 && 24 <= kConsumerRegs && kConsumerRegs <= 256 && + 2 * kConsumerRegs + kProducerRegs <= 512), + "consumer register budget exceeds the SM sub-partition file"); + + // Consumer tiling (v1 layout): two 128-thread warp groups; group g owns + // tiles g, g + 2, ... of the row; each thread owns one 16B vector per tile. + static constexpr uint32_t kNumGroups = 2; + static constexpr uint32_t kGroupThreads = kNumConsumerThreads / kNumGroups; + static constexpr uint32_t kVecElems = 16 / sizeof(bf16_t); // smem ld/st are 16B max + static constexpr uint32_t kNumTiles = kDim / kTile; + static constexpr uint32_t kSlicesPerGroup = (kNumTiles + kNumGroups - 1) / kNumGroups; + static constexpr uint32_t kAccPerThread = kSlicesPerGroup * kVecElems; + + // TMEM: per group, kTmemColsPerGroup columns of cw then of ow. + static constexpr uint32_t kTmemColsPerGroup = 32; + static constexpr uint32_t kTmemCols = 2 * kNumGroups * kTmemColsPerGroup; + static constexpr uint32_t kConsumerBarId = 1; // barrier 0 stays __syncthreads' + + static_assert(kDim % kTile == 0, "kDim must be a whole number of tiles"); + static_assert(kTile == kGroupThreads * kVecElems, "a tile is one group-wide 16B sweep"); + static_assert(kNumTiles <= kNumGroups * kSlicesPerGroup, "slices must cover all tiles"); + static_assert(kSlicesPerGroup * kVecElems <= kTmemColsPerGroup, "weight slices must fit their TMEM columns"); + static_assert(kNumRows >= 1, "need at least one bank row"); + static_assert(kChunkRows >= 1, "need at least one chunk row"); + + struct Smem { + uint64_t bar_full[kNumStages]; + uint64_t bar_free[kNumStages]; + float warp_rms[kNumConsumerWarps][kChunkRows]; + float warp_dot[kNumConsumerWarps][kChunkRows]; + // The out-norm reduction gets its own buffer: it can overlap the next + // token's first score reduction. + float warp_ssq[kNumConsumerWarps]; + uint32_t tmem_base; + alignas(128) bf16_t buf[kNumStages][kChunkRows][kDim]; + }; + + static SGL_DEVICE void forward(const AttnResTMAParams& params, Smem* smem); +}; + +SGL_DEVICE float2 fma_f32x2(float2 a, float2 b, float2 c) { + const uint64_t a_bits = reinterpret_cast(a); + const uint64_t b_bits = reinterpret_cast(b); + const uint64_t c_bits = reinterpret_cast(c); + uint64_t result; + asm("fma.rn.f32x2 %0, %1, %2, %3;" : "=l"(result) : "l"(a_bits), "l"(b_bits), "l"(c_bits)); + return reinterpret_cast(result); +} + +SGL_DEVICE float2 mul_f32x2(float2 a, float2 b) { + const uint64_t a_bits = reinterpret_cast(a); + const uint64_t b_bits = reinterpret_cast(b); + uint64_t result; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(result) : "l"(a_bits), "l"(b_bits)); + return reinterpret_cast(result); +} + +template +SGL_DEVICE void KimiK3AttnResTrait::forward( + const AttnResTMAParams& params, Smem* smem) { + using namespace device; + using row_vec_t = AlignedVector; // 16 bytes + const auto tx = threadIdx.x; + const auto warp_id = tx / kWarpThreads; + const auto lane_id = tx % kWarpThreads; + + if (warp_id == 0 && lane_id < kNumStages) { + ::ptx::mbar_init(&smem->bar_full[lane_id], 1); + ::ptx::mbar_init(&smem->bar_free[lane_id], kNumConsumerWarps * kWarpThreads); + ::ptx::fence_mbarrier_init(); + } else if (warp_id == 1) { + ::ptx::tcgen05_alloc(::ptx::to_shared(&smem->tmem_base), kTmemCols); + ::ptx::tcgen05_relinquish(); + } + + __syncthreads(); + if (warp_id >= kNumConsumerWarps) { // producer warp (group); first warp works + if constexpr (kConsumerRegs > 0) ::ptx::setmaxnreg_dec(); + // TODO: reduce the register usage + if (warp_id == kNumConsumerWarps && ::ptx::elect_one()) { + uint32_t global_chunks = 0; + constexpr uint32_t kRowBytes = kDim * sizeof(bf16_t); + for (auto token = blockIdx.x; token < params.num_tokens; token += gridDim.x) { +#pragma unroll + for (uint32_t ci = 0; ci < kNumChunks; ++ci, ++global_chunks) { + const uint32_t base_row = ci * kChunkRows; + const uint32_t an = (kNumRows + 1 - base_row) < kChunkRows ? (kNumRows + 1 - base_row) : kChunkRows; + const auto slot = global_chunks % kNumStages; + const auto phase = (global_chunks / kNumStages) & 1; + if (global_chunks >= kNumStages) { + ::ptx::mbar_wait_parity(&smem->bar_free[slot], phase ^ 1); + } + // One barrier per chunk; each row still gets its own bulk copy. + ::ptx::mbar_arrive_expect_tx(&smem->bar_full[slot], an * kRowBytes); +#pragma unroll + for (uint32_t r = 0; r < an; ++r) { + const auto row = base_row + r; + const auto src = row == kNumRows ? params.prefix_sum + token * kDim // + : params.bank + token * params.stride_bm + row * kDim; + // Only prefix_sum is written by the immediately-preceding kernel; + // one wait before the first token's prefix load covers the rest. + if (token == blockIdx.x && row == kNumRows) PDLWaitPrimary(); + ::ptx::cp_async_bulk_1d_load(&smem->buf[slot][r], src, kRowBytes, &smem->bar_full[slot]); + } + } + } + PDLTriggerSecondary(); + } + } else { // 2 consumer warp groups; one chunk per rendezvous + if constexpr (kConsumerRegs > 0) ::ptx::setmaxnreg_inc(); + const auto group = warp_id / (kNumConsumerWarps / kNumGroups); + const auto tid_in_group = tx % kGroupThreads; + const auto tmem_cw = smem->tmem_base + group * kTmemColsPerGroup; + const auto tmem_ow = tmem_cw + kNumGroups * kTmemColsPerGroup; + + // Stage this thread's cw / ow slices into TMEM (read once from gmem). + { + float staged[kAccPerThread]; +#pragma unroll + for (uint32_t si = 0; si < kSlicesPerGroup; ++si) { + const auto tile = si * kNumGroups + group; + if (tile >= kNumTiles) continue; + const auto h_base = tile * kTile + tid_in_group * kVecElems; +#pragma unroll + for (uint32_t j = 0; j < kVecElems; ++j) { + staged[si * kVecElems + j] = __bfloat162float(params.cw[h_base + j]); + } + } +#pragma unroll + for (uint32_t si = 0; si < kSlicesPerGroup; ++si) { + ::ptx::tcgen05_st_32x32b_x8( + tmem_cw + si * kVecElems, reinterpret_cast(&staged[si * kVecElems])); + } +#pragma unroll + for (uint32_t si = 0; si < kSlicesPerGroup; ++si) { + const auto tile = si * kNumGroups + group; + if (tile >= kNumTiles) continue; + const auto h_base = tile * kTile + tid_in_group * kVecElems; +#pragma unroll + for (uint32_t j = 0; j < kVecElems; ++j) { + staged[si * kVecElems + j] = __bfloat162float(params.ow[h_base + j]); + } + } +#pragma unroll + for (uint32_t si = 0; si < kSlicesPerGroup; ++si) { + ::ptx::tcgen05_st_32x32b_x8( + tmem_ow + si * kVecElems, reinterpret_cast(&staged[si * kVecElems])); + } + ::ptx::tcgen05_wait_st(); + } + + uint32_t global_chunks = 0; // mirrors the producer's chunk counter + for (auto token = blockIdx.x; token < params.num_tokens; token += gridDim.x) { + float run_max = -FLT_MAX; // online-softmax state + float run_sum = 0.f; + float2 acc[kAccPerThread / 2] = {}; // packed fp32x2 accumulator + +#pragma unroll + for (uint32_t ci = 0; ci < kNumChunks; ++ci, ++global_chunks) { + const uint32_t base_row = ci * kChunkRows; + // Active rows of this chunk; folds per unrolled iteration. + const uint32_t an = (kNumRows + 1 - base_row) < kChunkRows ? (kNumRows + 1 - base_row) : kChunkRows; + const auto slot = global_chunks % kNumStages; + const auto phase = (global_chunks / kNumStages) & 1; + ::ptx::mbar_wait_parity(&smem->bar_full[slot], phase); + + // Score pass: the cw slice is loaded once and reused across the + // chunk's rows; each row's 16B slices land in registers. rms/dot + // accumulate as packed fp32x2 lanes, folded to scalars just before + // the warp reduction. + row_vec_t rows[kSlicesPerGroup][kChunkRows]; + float2 acc_rms2[kChunkRows] = {}; + float2 acc_dot2[kChunkRows] = {}; +#pragma unroll + for (uint32_t si = 0; si < kSlicesPerGroup; ++si) { + const auto tile = si * kNumGroups + group; + if (tile >= kNumTiles) continue; + float q[kVecElems]; + ::ptx::tcgen05_ld_32x32b_x8(tmem_cw + si * kVecElems, reinterpret_cast(q)); + const auto* q2 = reinterpret_cast(q); + const auto offset = tile * kTile + tid_in_group * kVecElems; +#pragma unroll + for (uint32_t r = 0; r < an; ++r) { + rows[si][r].load(&smem->buf[slot][r][offset]); + } +#pragma unroll + for (uint32_t r = 0; r < an; ++r) { +#pragma unroll + for (uint32_t j = 0; j < kVecElems / 2; ++j) { + const auto f = cast(rows[si][r][j]); + acc_rms2[r] = fma_f32x2(f, f, acc_rms2[r]); + acc_dot2[r] = fma_f32x2(f, q2[j], acc_dot2[r]); + } + } + } + ::ptx::mbar_arrive(&smem->bar_free[slot]); + + // Fused bank write: the prefix row (last row of the last chunk) is + // already in registers; snapshot it to bank row nvb with plain + // stores — the .write() copy kernel disappears. Placed after the + // arrive so the slot handoff is not delayed. + if (params.prefix_dst != nullptr && base_row + an == kNumRows + 1) { + const uint32_t pr = kNumRows - base_row; + auto* dst = params.prefix_dst + static_cast(token) * params.stride_bm; +#pragma unroll + for (uint32_t si = 0; si < kSlicesPerGroup; ++si) { + const auto tile = si * kNumGroups + group; + if (tile >= kNumTiles) continue; + rows[si][pr].store(dst, tile * (kTile / kVecElems) + tid_in_group); + } + } + + float acc_rms[kChunkRows]; + float acc_dot[kChunkRows]; +#pragma unroll + for (uint32_t r = 0; r < an; ++r) { + acc_rms[r] = acc_rms2[r].x + acc_rms2[r].y; + acc_dot[r] = acc_dot2[r].x + acc_dot2[r].y; + } + +#pragma unroll + for (int n = 0; n < an; n++) { + acc_rms[n] = warp::reduce_sum(acc_rms[n]); + acc_dot[n] = warp::reduce_sum(acc_dot[n]); + } + if (lane_id == 0) { +#pragma unroll + for (uint32_t r = 0; r < an; ++r) { + smem->warp_rms[warp_id][r] = acc_rms[r]; + smem->warp_dot[warp_id][r] = acc_dot[r]; + } + } + ::ptx::named_barrier_sync(kConsumerBarId, kNumConsumerThreads); + // Lane r totals row r, then broadcasts: an*16 smem loads per warp + // instead of per thread. + float lane_logit = 0.f; + if (lane_id < an) { + float total_rms = 0.f; + float total_dot = 0.f; +#pragma unroll + for (uint32_t w = 0; w < kNumConsumerWarps; ++w) { + total_rms += smem->warp_rms[w][lane_id]; + total_dot += smem->warp_dot[w][lane_id]; + } + constexpr float kScale = 1.f / static_cast(kDim); + lane_logit = total_dot * rsqrtf(total_rms * kScale + params.eps); + } + float logit[kChunkRows]; +#pragma unroll + for (uint32_t r = 0; r < an; ++r) { + logit[r] = __shfl_sync(0xffffffffu, lane_logit, r); + } + + // Online-softmax fold of the chunk into the running accumulator. + float chunk_max = -FLT_MAX; +#pragma unroll + for (uint32_t r = 0; r < an; ++r) { + chunk_max = fmaxf(chunk_max, logit[r]); + } + const float new_max = fmaxf(run_max, chunk_max); + const float correction = exp2f((run_max - new_max) * math::log2e); + float weight[kChunkRows]; + float weight_sum = 0.f; +#pragma unroll + for (uint32_t r = 0; r < an; ++r) { + weight[r] = exp2f((logit[r] - new_max) * math::log2e); + weight_sum += weight[r]; + } + run_sum = run_sum * correction + weight_sum; + run_max = new_max; + + // Fold the chunk into the packed accumulator (v1 loop order: scale + // once, then rows outer / vector lanes inner, all fp32x2 FMAs). + const float2 correction2 = make_float2(correction, correction); + float2 weight2[kChunkRows]; +#pragma unroll + for (uint32_t r = 0; r < an; ++r) { + weight2[r] = make_float2(weight[r], weight[r]); + } +#pragma unroll + for (uint32_t si = 0; si < kSlicesPerGroup; ++si) { + const auto tile = si * kNumGroups + group; + if (tile >= kNumTiles) continue; + float2 a[kVecElems / 2]; +#pragma unroll + for (uint32_t j = 0; j < kVecElems / 2; ++j) { + a[j] = mul_f32x2(acc[si * (kVecElems / 2) + j], correction2); + } +#pragma unroll + for (uint32_t r = 0; r < an; ++r) { +#pragma unroll + for (uint32_t j = 0; j < kVecElems / 2; ++j) { + a[j] = fma_f32x2(weight2[r], cast(rows[si][r][j]), a[j]); + } + } +#pragma unroll + for (uint32_t j = 0; j < kVecElems / 2; ++j) { + acc[si * (kVecElems / 2) + j] = a[j]; + } + } + } + + // Fused out norm: mixed = acc / run_sum, out = rmsnorm(mixed) * ow. + const float inv_sum = 1.f / run_sum; + float2 acc_sq2 = make_float2(0.f, 0.f); +#pragma unroll + for (uint32_t j = 0; j < kAccPerThread / 2; ++j) { + acc_sq2 = fma_f32x2(acc[j], acc[j], acc_sq2); + } + float acc_sq = warp::reduce_sum(acc_sq2.x + acc_sq2.y); + if (lane_id == 0) smem->warp_ssq[warp_id] = acc_sq; + ::ptx::named_barrier_sync(kConsumerBarId, kNumConsumerThreads); + float total_sq = 0.f; +#pragma unroll + for (uint32_t w = 0; w < kNumConsumerWarps; ++w) { + total_sq += smem->warp_ssq[w]; + } + const float scale = inv_sum * rsqrtf(total_sq * inv_sum * inv_sum / static_cast(kDim) + params.eps); + const float2 scale2 = make_float2(scale, scale); + + auto* out_ptr = params.out + static_cast(token) * kDim; +#pragma unroll + for (uint32_t si = 0; si < kSlicesPerGroup; ++si) { + const auto tile = si * kNumGroups + group; + if (tile >= kNumTiles) continue; + float q[kVecElems]; + ::ptx::tcgen05_ld_32x32b_x8(tmem_ow + si * kVecElems, reinterpret_cast(q)); + const auto* q2 = reinterpret_cast(q); + row_vec_t out_vec; +#pragma unroll + for (uint32_t j = 0; j < kVecElems / 2; ++j) { + const auto scaled = mul_f32x2(acc[si * (kVecElems / 2) + j], scale2); + out_vec[j] = cast(mul_f32x2(scaled, q2[j])); + } + const auto row_vid = tile * (kTile / kVecElems) + tid_in_group; + if (params.output_mc != nullptr) { + const auto global_token = static_cast(params.rank) * params.num_tokens + token; + const auto global_vid = global_token * (kDim / kVecElems) + row_vid; + st_multimem_16B(out_vec, params.output_mc, global_vid); + } else { + out_vec.store(out_ptr, row_vid); + } + } + } + ::ptx::named_barrier_sync(kConsumerBarId, kNumConsumerThreads); + if (warp_id == 1) { + ::ptx::tcgen05_dealloc(smem->tmem_base, kTmemCols); + } + } +} + +// kOccupancy > 1 caps the register budget (65536 / (kOccupancy * kNumThreads)) +// so that many CTAs actually co-reside; smem must also fit kOccupancy copies. +template +__global__ void __launch_bounds__(Trait::kNumThreads, kOccupancy) + attn_res_fused_tma_kernel(const __grid_constant__ AttnResTMAParams params) { + extern __shared__ char smem_raw[]; + Trait::forward(params, reinterpret_cast(smem_raw)); +} + +SGL_DEVICE uint32_t* attn_res_sem_mc_flag(uint8_t* sem_mc, uint32_t block) { + static_assert(sizeof(device::distributed::Semaphore) == 128); + return reinterpret_cast(sem_mc + block * sizeof(device::distributed::Semaphore)); +} + +SGL_DEVICE void attn_res_sem_arrive_relaxed(uint32_t* flag) { +#if SGL_ARCH_HOPPER_OR_GREATER + asm volatile("multimem.red.relaxed.sys.global.add.u32 [%0], 1;" ::"l"(flag) : "memory"); +#else + assert(false && "multimem red requires Hopper or later"); +#endif +} + +SGL_DEVICE void attn_res_sem_arrive_release(uint32_t* flag) { +#if SGL_ARCH_HOPPER_OR_GREATER + asm volatile("multimem.red.release.sys.global.add.u32 [%0], 1;" ::"l"(flag) : "memory"); +#else + assert(false && "multimem red requires Hopper or later"); +#endif +} + +// Fused NVLS pull RS + local residual + attention-residual aggregation. +// The entry/exit barriers make local o_proj writes visible before the +// producer's multimem reduction and preserve the shared pull-semaphore +// protocol used by the neighboring K3 collectives. +template +__global__ void __launch_bounds__(Trait::kNumThreads, kOccupancy) + attn_res_fused_pull_rs_kernel(const __grid_constant__ AttnResTMAParams params) { + __shared__ uint32_t exit_base; + if (threadIdx.x == 0) { + auto* semaphore = ¶ms.sem_local[blockIdx.x]; + const auto reserved = semaphore->counter_ptr()->inc(2 * params.world_size); + exit_base = reserved + params.world_size; + device::PDLWaitPrimary(); + attn_res_sem_arrive_relaxed(attn_res_sem_mc_flag(params.sem_mc, blockIdx.x)); + while (semaphore->get_relaxed() - reserved < params.world_size) + ; + } + __syncthreads(); + + // Cooperative NVLS materialization: all TMA producer + consumer threads + // participate, so the remote reduction exposes hundreds of outstanding + // 16-byte loads per CTA instead of serializing the row through one warp. + using pull_vec_t = device::AlignedVector; + using SumOp = device::ReductionTrait; + constexpr uint32_t kRowVecs = Trait::kDim * sizeof(bf16_t) / sizeof(pull_vec_t); + for (auto token = blockIdx.x; token < params.num_tokens; token += gridDim.x) { + auto* prefix = params.prefix_out + static_cast(token) * Trait::kDim; + const auto* input_mc = params.input_mc + static_cast(token) * Trait::kDim * sizeof(bf16_t); + const auto* residual = + params.residual == nullptr ? nullptr : params.residual + static_cast(token) * Trait::kDim; + for (uint32_t vid = threadIdx.x; vid < kRowVecs; vid += blockDim.x) { + pull_vec_t vec; + ld_multimem_16B(vec, input_mc, vid); + if (residual != nullptr) { + pull_vec_t res; + res.load(residual, vid); +#pragma unroll + for (uint32_t j = 0; j < 4; ++j) { + vec[j] = SumOp::reduce(vec[j], res[j]); + } + } + vec.store(prefix, vid); + } + } + __threadfence(); + __syncthreads(); + + extern __shared__ char smem_raw[]; + Trait::forward(params, reinterpret_cast(smem_raw)); + + __syncthreads(); + if (threadIdx.x == 0) { + auto* semaphore = ¶ms.sem_local[blockIdx.x]; + attn_res_sem_arrive_release(attn_res_sem_mc_flag(params.sem_mc, blockIdx.x)); + while (semaphore->get_acquire() - exit_base < params.world_size) + ; + } +} + +// Local attention-residual aggregation + direct AG epilogue. The consumer +// threads already hold each normalized 16B output vector in registers, so +// Trait::forward multicast-stores those vectors into every peer's symmetric +// full-token output instead of launching a separate all-gather. +template +__global__ void __launch_bounds__(Trait::kNumThreads, kOccupancy) + attn_res_fused_direct_ag_kernel(const __grid_constant__ AttnResTMAParams params) { + __shared__ uint32_t exit_base; + if (threadIdx.x == 0) { + auto* semaphore = ¶ms.sem_local[blockIdx.x]; + const auto reserved = semaphore->counter_ptr()->inc(2 * params.world_size); + exit_base = reserved + params.world_size; + attn_res_sem_arrive_relaxed(attn_res_sem_mc_flag(params.sem_mc, blockIdx.x)); + while (semaphore->get_relaxed() - reserved < params.world_size) + ; + } + __syncthreads(); + + extern __shared__ char smem_raw[]; + Trait::forward(params, reinterpret_cast(smem_raw)); + + __syncthreads(); + if (threadIdx.x == 0) { + auto* semaphore = ¶ms.sem_local[blockIdx.x]; + attn_res_sem_arrive_release(attn_res_sem_mc_flag(params.sem_mc, blockIdx.x)); + while (semaphore->get_acquire() - exit_base < params.world_size) + ; + } +} + +} // namespace sglang + +using namespace sglang; +using host::distributed::CommunicatorRef; + +// Host launcher: constexpr kernel table over nvb. + +template +struct AttnResFusedTmaKernel { + using KernelFn = void (*)(const AttnResTMAParams); + template + using Trait = KimiK3AttnResTrait; + static constexpr uint32_t kNumThreads = Trait<1>::kNumThreads; + static constexpr size_t kSmemBytes = sizeof(typename Trait<1>::Smem); + // kOccupancy copies of the smem ring must fit one SM (228KB on SM100). + static_assert(kOccupancy >= 1 && kOccupancy * kSmemBytes <= 233472 - 1024, "occupancy exceeds the smem budget"); + + template + static constexpr auto make_table(std::index_sequence) { + return std::array{nullptr, attn_res_fused_tma_kernel, kOccupancy>...}; + } + static constexpr auto kTable = make_table(std::make_index_sequence{}); + template + static constexpr auto make_pull_table(std::index_sequence) { + return std::array{nullptr, attn_res_fused_pull_rs_kernel, kOccupancy>...}; + } + static constexpr auto kPullTable = make_pull_table(std::make_index_sequence{}); + template + static constexpr auto make_ag_table(std::index_sequence) { + return std::array{ + nullptr, attn_res_fused_direct_ag_kernel, kOccupancy>...}; + } + static constexpr auto kAgTable = make_ag_table(std::make_index_sequence{}); + + static void + run(const tvm::ffi::TensorView prefix_sum, + const tvm::ffi::TensorView bank, + const tvm::ffi::TensorView cw, + const tvm::ffi::TensorView ow, + const tvm::ffi::TensorView out, + int64_t nvb, + double eps, + bool write_prefix) { + using namespace host; + + auto T_ = SymbolicSize{"num_tokens"}; + auto H_ = SymbolicSize{"hidden_size"}; + auto NB_ = SymbolicSize{"num_bank_slots"}; + auto device = SymbolicDevice{}; + device.set_options(); + + TensorMatcher({T_, H_}).with_dtype().with_device(device).verify(prefix_sum).verify(out); + TensorMatcher({T_, NB_, H_}).with_dtype().with_device(device).verify(bank); + TensorMatcher({H_}).with_dtype().with_device(device).verify(cw).verify(ow); + + const auto num_tokens = static_cast(T_.unwrap()); + const auto H = static_cast(H_.unwrap()); + const auto NB = static_cast(NB_.unwrap()); + + RuntimeCheck(H == kDim, "attn_res_fused_tma: H must be ", kDim, ", got ", H); + RuntimeCheck( + 1 <= nvb && nvb <= kMaxBankRows && nvb <= NB, + "attn_res_fused_tma: nvb must be in [1, ", + kMaxBankRows, + "] and <= NB, got nvb=", + nvb, + " NB=", + NB); + RuntimeCheck( + !write_prefix || nvb < NB, + "attn_res_fused_tma: write_prefix targets bank row nvb, needs nvb < NB, got nvb=", + nvb, + " NB=", + NB); + + if (num_tokens == 0) return; + + [[maybe_unused]] static const bool attrs_set = [] { + for (uint32_t i = 1; i <= kMaxBankRows; ++i) { + RuntimeDeviceCheck(cudaFuncSetAttribute(kTable[i], cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemBytes)); + } + return true; + }(); + + const auto num_sm = runtime::get_sm_count(device.unwrap().device_id); + const auto grid = std::min((int64_t)num_sm * kOccupancy, num_tokens); + const auto params = AttnResTMAParams{ + .prefix_sum = static_cast(prefix_sum.data_ptr()), + .bank = static_cast(bank.data_ptr()), + .cw = static_cast(cw.data_ptr()), + .ow = static_cast(ow.data_ptr()), + .out = static_cast(out.data_ptr()), + .prefix_dst = write_prefix ? static_cast(bank.data_ptr()) + nvb * H : nullptr, + .input_mc = nullptr, + .residual = nullptr, + .prefix_out = nullptr, + .sem_local = nullptr, + .sem_mc = nullptr, + .output_mc = nullptr, + .world_size = 0, + .rank = 0, + .stride_bm = NB * H, + .eps = static_cast(eps), + .num_tokens = static_cast(num_tokens), + }; + LaunchKernel(grid, kNumThreads, device.unwrap(), kSmemBytes).enable_pdl(true)(kTable[nvb], params); + } + + static void run_pull_rs( + CommunicatorRef ref, + const tvm::ffi::TensorView input, + std::optional residual, + const tvm::ffi::TensorView bank, + const tvm::ffi::TensorView cw, + const tvm::ffi::TensorView ow, + const tvm::ffi::TensorView out, + const tvm::ffi::TensorView prefix_out, + int64_t nvb, + double eps, + int64_t input_mc_ptr, + int64_t sem_mc_ptr, + int64_t max_blocks) { + using namespace host; + const auto& data = *ref.get(); + auto GT_ = SymbolicSize{"global_tokens"}; + auto T_ = SymbolicSize{"local_tokens"}; + auto H_ = SymbolicSize{"hidden_size"}; + auto NB_ = SymbolicSize{"num_bank_slots"}; + auto device = SymbolicDevice{}; + device.set_options(); + + TensorMatcher({GT_, H_}).with_dtype().with_device(device).verify(input); + TensorMatcher({T_, H_}).with_dtype().with_device(device).verify(out).verify(prefix_out); + if (residual.has_value()) { + TensorMatcher({T_, H_}).with_dtype().with_device(device).verify(residual.value()); + } + TensorMatcher({T_, NB_, H_}).with_dtype().with_device(device).verify(bank); + TensorMatcher({H_}).with_dtype().with_device(device).verify(cw).verify(ow); + + const auto global_tokens = static_cast(GT_.unwrap()); + const auto num_tokens = static_cast(T_.unwrap()); + const auto H = static_cast(H_.unwrap()); + const auto NB = static_cast(NB_.unwrap()); + RuntimeCheck(data.world_size > 1, "fused pull RS requires world_size > 1"); + RuntimeCheck(global_tokens == num_tokens * data.world_size, "global tokens must equal local tokens * world size"); + RuntimeCheck(H == kDim, "fused pull RS: H must be ", kDim, ", got ", H); + RuntimeCheck(1 <= nvb && nvb <= kMaxBankRows && nvb <= NB, "fused pull RS: invalid nvb=", nvb, " NB=", NB); + RuntimeCheck(input_mc_ptr != 0, "fused pull RS requires multicast input"); + RuntimeCheck(sem_mc_ptr != 0, "fused pull RS requires multicast semaphores"); + RuntimeCheck(max_blocks > 0, "fused pull RS requires max_blocks > 0"); + if (num_tokens == 0) return; + + [[maybe_unused]] static const bool attrs_set = [] { + for (uint32_t i = 1; i <= kMaxBankRows; ++i) { + RuntimeDeviceCheck( + cudaFuncSetAttribute(kPullTable[i], cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemBytes)); + } + return true; + }(); + + const auto num_sm = runtime::get_sm_count(device.unwrap().device_id); + const auto grid = std::min( + {static_cast(num_sm) * kOccupancy, + num_tokens, + max_blocks, + static_cast(data.num_pull_blocks)}); + const auto local_elems = num_tokens * H; + const auto params = AttnResTMAParams{ + .prefix_sum = static_cast(prefix_out.data_ptr()), + .bank = static_cast(bank.data_ptr()), + .cw = static_cast(cw.data_ptr()), + .ow = static_cast(ow.data_ptr()), + .out = static_cast(out.data_ptr()), + .prefix_dst = nullptr, + .input_mc = reinterpret_cast(static_cast(input_mc_ptr)) + + data.rank * local_elems * sizeof(bf16_t), + .residual = residual.has_value() ? static_cast(residual.value().data_ptr()) : nullptr, + .prefix_out = static_cast(prefix_out.data_ptr()), + .sem_local = data.pull_semaphores[data.rank], + .sem_mc = reinterpret_cast(static_cast(sem_mc_ptr)), + .output_mc = nullptr, + .world_size = data.world_size, + .rank = data.rank, + .stride_bm = NB * H, + .eps = static_cast(eps), + .num_tokens = static_cast(num_tokens), + }; + LaunchKernel(grid, kNumThreads, device.unwrap(), kSmemBytes).enable_pdl(true)(kPullTable[nvb], params); + } + + static void run_direct_ag( + CommunicatorRef ref, + const tvm::ffi::TensorView prefix_sum, + const tvm::ffi::TensorView bank, + const tvm::ffi::TensorView cw, + const tvm::ffi::TensorView ow, + const tvm::ffi::TensorView out, + int64_t nvb, + double eps, + int64_t output_mc_ptr, + int64_t sem_mc_ptr, + int64_t max_blocks, + bool write_prefix) { + using namespace host; + const auto& data = *ref.get(); + auto T_ = SymbolicSize{"local_tokens"}; + auto GT_ = SymbolicSize{"global_tokens"}; + auto H_ = SymbolicSize{"hidden_size"}; + auto NB_ = SymbolicSize{"num_bank_slots"}; + auto device = SymbolicDevice{}; + device.set_options(); + + TensorMatcher({T_, H_}).with_dtype().with_device(device).verify(prefix_sum); + TensorMatcher({T_, NB_, H_}).with_dtype().with_device(device).verify(bank); + TensorMatcher({H_}).with_dtype().with_device(device).verify(cw).verify(ow); + TensorMatcher({GT_, H_}).with_dtype().with_device(device).verify(out); + + const auto num_tokens = static_cast(T_.unwrap()); + const auto global_tokens = static_cast(GT_.unwrap()); + const auto H = static_cast(H_.unwrap()); + const auto NB = static_cast(NB_.unwrap()); + RuntimeCheck(data.world_size > 1, "fused direct AG requires world_size > 1"); + RuntimeCheck(global_tokens == num_tokens * data.world_size, "global tokens must equal local tokens * world size"); + RuntimeCheck(H == kDim, "fused direct AG: H must be ", kDim, ", got ", H); + RuntimeCheck(1 <= nvb && nvb <= kMaxBankRows && nvb <= NB, "fused direct AG: invalid nvb=", nvb, " NB=", NB); + RuntimeCheck(!write_prefix || nvb < NB, "fused direct AG: write_prefix targets bank row nvb, needs nvb < NB"); + RuntimeCheck(output_mc_ptr != 0, "fused direct AG requires multicast output"); + RuntimeCheck(sem_mc_ptr != 0, "fused direct AG requires multicast semaphores"); + RuntimeCheck(max_blocks > 0, "fused direct AG requires max_blocks > 0"); + if (num_tokens == 0) return; + + [[maybe_unused]] static const bool attrs_set = [] { + for (uint32_t i = 1; i <= kMaxBankRows; ++i) { + RuntimeDeviceCheck(cudaFuncSetAttribute(kAgTable[i], cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemBytes)); + } + return true; + }(); + + const auto num_sm = runtime::get_sm_count(device.unwrap().device_id); + const auto grid = std::min( + {static_cast(num_sm) * kOccupancy, + num_tokens, + max_blocks, + static_cast(data.num_pull_blocks)}); + const auto params = AttnResTMAParams{ + .prefix_sum = static_cast(prefix_sum.data_ptr()), + .bank = static_cast(bank.data_ptr()), + .cw = static_cast(cw.data_ptr()), + .ow = static_cast(ow.data_ptr()), + .out = static_cast(out.data_ptr()), + .prefix_dst = write_prefix ? static_cast(bank.data_ptr()) + nvb * H : nullptr, + .input_mc = nullptr, + .residual = nullptr, + .prefix_out = nullptr, + .sem_local = data.pull_semaphores[data.rank], + .sem_mc = reinterpret_cast(static_cast(sem_mc_ptr)), + .output_mc = reinterpret_cast(static_cast(output_mc_ptr)), + .world_size = data.world_size, + .rank = data.rank, + .stride_bm = NB * H, + .eps = static_cast(eps), + .num_tokens = static_cast(num_tokens), + }; + LaunchKernel(grid, kNumThreads, device.unwrap(), kSmemBytes).enable_pdl(true)(kAgTable[nvb], params); + } +}; diff --git a/python/sglang/kernels/jit/csrc/kimi_k3/comm/ar_fusion.cuh b/python/sglang/kernels/jit/csrc/kimi_k3/comm/ar_fusion.cuh new file mode 100644 index 000000000..a4ad28242 --- /dev/null +++ b/python/sglang/kernels/jit/csrc/kimi_k3/comm/ar_fusion.cuh @@ -0,0 +1,908 @@ +// K3 MNNVL fused all-reduce (bf16-only), two zero-copy algorithm families: +// +// - push (1shot): lamport-style push, but the data staging is a SINGLE +// multicast store into slot `rank` of every peer's push workspace +// (replacing the 7 unicast stores of the generic kernel), followed by a +// local zero-marker polling reduce. Input is read in place and the +// result is written back in place — no staging copies. Works for any +// input tensor; reuses the CustomAllReduceV2 push workspace + counter. +// Best for small messages. +// +// - pull (2shot): low-SM NVLS pull directly ON the input tensor, which +// therefore MUST live in (multicast-bound) symmetric memory: each rank +// `multimem.ld_reduce`s its shard from the input's multicast address and +// `multimem.st`s the result back — in-place, zero copies. Two +// NCCL-style ideas keep a handful of blocks (tuned 1~16, 4 at the large +// end) at the fabric limit: +// +// * deep pipelining: the copy loop manually keeps `unroll` multimem +// loads in flight per thread before the first store, cascading +// through halving widths for mid-size tails. +// * multicast barriers on the CustomAllReduceV2 pull semaphores: the +// generic pull kernels' reservation protocol with each arrival sent +// as ONE `multimem.red.add` on the semaphore flag's multicast alias +// instead of per-peer unicast reds — identical memory effects, so +// both kernel families share the slots freely (single-stream calls +// are serialized). +// +// Both families fuse either a residual add (`out = allreduce(x) + r`; the +// residual must be identical on every rank — a fully reduced tensor such as +// the attn-res prefix sum — or absent) or the RMSNorm epilogue over the +// latent of the K3 latent|shared MoE buffer (*_norm variants). +// +#include +#include +#include +#include + +#include + +// TODO: remove dependency on the custom_all_reduce, move out common utilities +#include "../../distributed/custom_all_reduce.cuh" +#include "ptx_sys.cuh" + +namespace sglang { + +// Same shape as gemm_ag / gemm_ar: pull the ptx_sys helpers in by name so the +// call sites below stay unqualified. +using device::distributed::multimem_red_add_relaxed; +using device::distributed::multimem_red_add_release; + +struct FusionParams { + uint8_t* input; // tensor pointer (in place) + const uint8_t* residual; // may be null (compile-time kHasResidual selects) + uint8_t* push_ws_mc; // multicast VA of the push workspace base + uint8_t* push_ws_local; // local push workspace base (poll side) + Counter* push_counter; // per-block phase counters (local memory) + int64_t push_buffer_stride; // per-buffer bytes (2 * world_size buffers) + uint32_t rank; + uint32_t num_vecs; // 16B vectors + // *_norm variants only: RMSNorm epilogue over the first num_norm_rows + // rows of the [num_vecs / kNormRowVecs, kNormDim] row view + const uint8_t* norm_weight; + float norm_eps; + uint32_t num_norm_rows; + uint32_t num_push_counters; // cluster variant only: full counter array size + // finalize_push_norm only: trtllm-gen deferred-finalize inputs (kimi_k3.py + // deferred finalize path); `input` is then output-only ([T, kNormDim]) + const uint8_t* fin_gemm2; // [P, kNormDim] bf16, permuted rows + const uint8_t* fin_idx; // [T * kFinTopK] int32, -1 = dropped slot + const uint8_t* fin_weights; // [T, kFinTopK] bf16 +}; + +// The *_norm variants view the input as rows of the K3 latent width (3584 +// bf16 = 448 16B vectors per row) and give the first num_norm_rows an +// RMSNorm epilogue (K3: the [N latent | 2N shared] MoE buffer with N normed +// rows, or a latent-only [N, 3584] tensor normed in full). +constexpr uint32_t kNormDim = 3584; +constexpr uint32_t kNormRowVecs = kNormDim / 8; // 448 +constexpr uint32_t kNormWarps = kNormRowVecs / device::kWarpThreads; // 14 + +template +__global__ __launch_bounds__(1024, 1) void all_reduce_push_res_kernel(const __grid_constant__ FusionParams params) { + using vec_t = device::AlignedVector; + + const auto tx = threadIdx.x; + const auto bx = blockIdx.x; + const auto global_tid = bx * blockDim.x + tx; + const auto num_threads = blockDim.x * gridDim.x; + const auto num_vecs = params.num_vecs; + + // prologue: the previous phase flip (counter inc) must be visible + device::PDLWaitPrimary(); + const auto phase = params.push_counter[bx].get() & 1; + const auto r = params.rank; + const auto stride_bytes = params.push_buffer_stride; + const auto phase_stride_bytes = (phase * kWorldSize) * stride_bytes; + // one multicast store lands this rank's data in slot r of EVERY peer + const auto push_ptr = params.push_ws_mc + r * stride_bytes + phase_stride_bytes; + const auto poll_ptr = params.push_ws_local + phase_stride_bytes; + + // stage 1: multicast-push local data, remapping all-zero bf16x2 pairs + static_assert(fp_trait::pos_zero == 0, "the empty marker is all-zero bits"); + constexpr uint32_t kNegZeroPair = 0x8000u; // {-0.0, +0.0}: sum-neutral, non-zero + for (auto vid = global_tid; vid < num_vecs; vid += num_threads) { + vec_t vec; + ld_global_16B(vec, params.input, vid); + auto& bits = *reinterpret_cast(&vec); + if (bits.x == 0) bits.x = kNegZeroPair; + if (bits.y == 0) bits.y = kNegZeroPair; + if (bits.z == 0) bits.z = kNegZeroPair; + if (bits.w == 0) bits.w = kNegZeroPair; + st_multimem_16B(vec, push_ptr, vid); + } + + // launch pdl early for low latency case + device::PDLTriggerSecondary(); + + // stage 2: poll all slots, reduce (+ residual), write back in place, + // re-establish the empty markers for the next same-phase round + vec_t zero_vec; + zero_vec.fill(bf16x2_t{get_pos_zero(), get_pos_zero()}); + for (auto vid = global_tid; vid < num_vecs; vid += num_threads) { + vec_t vec[kWorldSize + kHasResidual]; + if constexpr (kHasResidual) vec[kWorldSize].load(params.residual, vid); + do { + bool has_zero = false; +#pragma unroll + for (uint32_t i = 0; i < kWorldSize; ++i) { + ld_relaxed_16B(vec[i], poll_ptr + i * stride_bytes, vid); + // the producer remapped all-zero pairs, so a written u32 is never + // 0: u32 == 0 <=> the 4B atom still holds the empty marker + const auto bits = *reinterpret_cast(&vec[i]); + has_zero |= bits.x == 0; + has_zero |= bits.y == 0; + has_zero |= bits.z == 0; + has_zero |= bits.w == 0; + } + if (!has_zero) break; + } while (true); + const auto out_vec = reduce(vec); // fp32 accumulation over 8(+1) inputs + st_global_16B(out_vec, params.input, vid); +#pragma unroll + for (uint32_t i = 0; i < kWorldSize; ++i) { + st_global_16B(zero_vec, poll_ptr + i * stride_bytes, vid); + } + } + + // epilogue: flip this block's phase + __syncthreads(); + if (tx == 0) params.push_counter[bx].set(phase ^ 1); +} + +// --- deferred-finalize staging (finalize_push_norm) ------------------------ +// The trtllm-gen MoE with do_finalize=False hands back its finalize inputs +// (see kernels/ops/moe/trtllm_gen_moe.py); the fused kernel computes the finalize +// during the push staging pass, so the rank-local latent never materializes. + +constexpr uint32_t kFinTopK = 16; + +// One 16B vector of the deferred MoE finalize (latent width fixed to kNormDim): +// local[t] = sum_k fin_weights[t, k] * fin_gemm2[fin_idx[t*16 + k]] +// All 16 gathers issue before the FMA chain; threads of the same token +// broadcast-load the same routing rows. +SGL_DEVICE device::AlignedVector finalize_vec(const FusionParams& params, uint32_t vid) { + using namespace device; + constexpr uint32_t kIdxVecSize = kMaxVecBytes / sizeof(int32_t); + constexpr uint32_t kWVecSize = kMaxVecBytes / sizeof(bf16_t); + constexpr uint32_t kIdxVecs = kFinTopK / kIdxVecSize; // 2 on SM100+ + constexpr uint32_t kWVecs = kFinTopK / kWVecSize; // 1 on SM100+ + + const uint32_t token = vid / kNormRowVecs; + const uint32_t hvec = vid % kNormRowVecs; + + AlignedVector idx[kIdxVecs]; +#pragma unroll + for (uint32_t j = 0; j < kIdxVecs; ++j) { + idx[j].load(params.fin_idx, token * kIdxVecs + j); + } + AlignedVector weight[kWVecs]; +#pragma unroll + for (uint32_t j = 0; j < kWVecs; ++j) { + weight[j].load(params.fin_weights, token * kWVecs + j); + } + + const auto* g2 = reinterpret_cast(params.fin_gemm2); + AlignedVector in[kFinTopK]; +#pragma unroll + for (uint32_t k = 0; k < kFinTopK; ++k) { + const int32_t row = idx[k / kIdxVecSize][k % kIdxVecSize]; + if (row >= 0) { + in[k].load(g2 + static_cast(row) * kNormDim, hvec); + } + } + float acc[8] = {}; +#pragma unroll + for (uint32_t k = 0; k < kFinTopK; ++k) { + const int32_t row = idx[k / kIdxVecSize][k % kIdxVecSize]; + if (row < 0) continue; + const bf16_t w_k = weight[k / kWVecSize][k % kWVecSize]; +#pragma unroll + for (uint32_t i = 0; i < 8; ++i) { + acc[i] = device::math::fma_f32_bf16(in[k][i], w_k, acc[i]); + } + } + AlignedVector out; +#pragma unroll + for (uint32_t j = 0; j < 4; ++j) { + out[j] = cast(fp32x2_t{acc[2 * j], acc[2 * j + 1]}); + } + return out; +} + +template +SGL_DEVICE float reduce_sqr(device::AlignedVector& out_vec, device::AlignedVector (&vec)[M]) { + fp32x2_t acc_vec[N]; +#pragma unroll + for (size_t i = 0; i < M; ++i) { +#pragma unroll + for (size_t j = 0; j < N; ++j) { + const auto [x, y] = device::cast(vec[i][j]); + auto& [acc_x, acc_y] = acc_vec[j]; + acc_x = i == 0 ? x : acc_x + x; + acc_y = i == 0 ? y : acc_y + y; + } + } + float sum_eq = 0.0f; +#pragma unroll + for (size_t j = 0; j < N; ++j) { + sum_eq += acc_vec[j].x * acc_vec[j].x; + sum_eq += acc_vec[j].y * acc_vec[j].y; + out_vec[j] = device::cast(acc_vec[j]); + } + return sum_eq; +} + +// kFinalize: stage 1 computes the deferred MoE finalize per vector instead of +// reading a staged input tensor; `input` is then output-only. The host sets +// num_norm_rows to the full row count (every reduced row is normed). +template +__global__ __launch_bounds__(kNormRowVecs / kClusterSize) __cluster_dims__(kClusterSize, 1, 1) // + void all_reduce_push_norm_cluster_kernel(const __grid_constant__ FusionParams params) { + namespace cg = cooperative_groups; + using namespace device; + using vec_t = AlignedVector; + constexpr uint32_t kBlockSize = kNormRowVecs / kClusterSize; + constexpr uint32_t kNumWarps = kBlockSize / kWarpThreads; + + static_assert(kBlockSize % kWarpThreads == 0); + static_assert(kNormRowVecs % kClusterSize == 0); + static_assert(kNumWarps >= 1); + + const auto tx = threadIdx.x; + const auto bx = blockIdx.x; + const auto global_tid = bx * kBlockSize + tx; + const auto num_vecs = params.num_vecs; + const auto num_rows = num_vecs / kNormRowVecs; + + const auto row_idx = bx / kClusterSize; + const auto num_row_clusters = gridDim.x / kClusterSize - 1; // last one is the bumper + // stage-1 grid-stride is over the ROW clusters only: the bumper early-returns + // and never stages, so it must not be counted or its share of vids is dropped + const auto num_threads = kBlockSize * num_row_clusters * kClusterSize; + + PDLWaitPrimary(); + + // special case: the bumper cluster flips every remaining counter so the + // whole array stays globally in phase. Only ONE block does it (threads + // grid-stride the counters) — each counter must be inc'd exactly once, + // independent of whether kClusterSize is odd or even. + if (row_idx == num_row_clusters) { + if (bx % kClusterSize == 0) { + for (uint32_t r = num_row_clusters + tx; r < params.num_push_counters; r += kBlockSize) { + params.push_counter[r].inc(1); + } + } + return PDLTriggerSecondary(); + } + + const auto phase = params.push_counter[row_idx].get() & 1; + const auto r = params.rank; + const auto stride_bytes = params.push_buffer_stride; + const auto phase_stride_bytes = (phase * kWorldSize) * stride_bytes; + const auto push_ptr = params.push_ws_mc + r * stride_bytes + phase_stride_bytes; + const auto poll_ptr = params.push_ws_local + phase_stride_bytes; + + // stage 1: multicast staging (grid-stride); kFinalize computes each vector + // in place of the load + static_assert(fp_trait::pos_zero == 0, "the empty marker is all-zero bits"); + for (auto vid = global_tid; vid < num_vecs; vid += num_threads) { + vec_t vec; + if constexpr (kFinalize) { + vec = finalize_vec(params, vid); + } else { + ld_global_16B(vec, params.input, vid); + } + auto& bits = *reinterpret_cast(&vec); + if (bits.x == 0) bits.x = fp_trait::neg_zero; + if (bits.y == 0) bits.y = fp_trait::neg_zero; + if (bits.z == 0) bits.z = fp_trait::neg_zero; + if (bits.w == 0) bits.w = fp_trait::neg_zero; + st_multimem_16B(vec, push_ptr, vid); + } + + // stage 2: one row per cluster pass (the bumper cluster owns no rows) + const auto cluster = cg::this_cluster(); + const auto cluster_rank = bx % kClusterSize; + vec_t w; + w.load(params.norm_weight, cluster_rank * kBlockSize + tx); + vec_t zero_vec; + zero_vec.fill(bf16x2_t{get_pos_zero(), get_pos_zero()}); + __shared__ alignas(8) float smem_raw[2][kClusterSize][kNumWarps]; + uint32_t parity = 0; + + // NOTE: launch PDL earlier for low latency case + PDLTriggerSecondary(); + + for (auto row = row_idx; row < num_rows; row += num_row_clusters) { + const auto vid = row * kNormRowVecs + cluster_rank * kBlockSize + tx; + vec_t vec[kWorldSize]; + do { + bool has_zero = false; +#pragma unroll + for (uint32_t i = 0; i < kWorldSize; ++i) { + ld_relaxed_16B(vec[i], poll_ptr + i * stride_bytes, vid); + const auto bits = *reinterpret_cast(&vec[i]); + has_zero |= bits.x == 0; + has_zero |= bits.y == 0; + has_zero |= bits.z == 0; + has_zero |= bits.w == 0; + } + if (!has_zero) break; + } while (true); + + vec_t out_vec; + if (row < params.num_norm_rows) { // cluster-uniform branch + auto& smem = smem_raw[parity]; + parity ^= 1; + // push each warp's partial to EVERY peer's slot for this block: lane p + // (p < kClusterSize) targets peer p, warp w selects the [.][w] slot. So + // after the barrier every block holds all kClusterSize*kNumWarps + // partials in its own smem and reduces them locally — the read side + // never touches remote DSMEM, so no post-read barrier is needed to guard + // against a peer CTA exiting (parity double-buffers across rows). + const auto lane = tx % kWarpThreads; + const auto warp = tx / kWarpThreads; + const auto warp_sqr = warp::reduce_sum(reduce_sqr(out_vec, vec)); + if (lane < kClusterSize) { + float* dst = cluster.map_shared_rank(&smem[cluster_rank][warp], lane); + *dst = warp_sqr; + } + cluster.sync(); + // load local + float total = 0.0f; +#pragma unroll + for (uint32_t r = 0; r < kClusterSize; ++r) { + using vec_t = AlignedVector; + vec_t remote_value; + remote_value.load(smem[r]); +#pragma unroll + for (uint32_t w = 0; w < kNumWarps; ++w) { + total += remote_value[w]; + } + } + const auto norm_factor = math::rsqrt(total / kNormDim + params.norm_eps); +#pragma unroll + for (uint32_t j = 0; j < 4; ++j) { + const auto [a, b] = cast(out_vec[j]); + const auto [wa, wb] = cast(w[j]); + out_vec[j] = cast(fp32x2_t{a * norm_factor * wa, b * norm_factor * wb}); + } + } else { + out_vec = reduce(vec); + } + + st_global_16B(out_vec, params.input, vid); +#pragma unroll + for (uint32_t i = 0; i < kWorldSize; ++i) { + st_global_16B(zero_vec, poll_ptr + i * stride_bytes, vid); + } + } + + // epilogue: each row cluster flips its own counter; the bumper cluster + // flips every remaining one so the whole array stays globally uniform + if (cluster_rank == 0 && tx == 0) { + params.push_counter[row_idx].set(phase ^ 1); + } +} + +// Pull family: low-SM NVLS, reusing the CustomAllReduceV2 pull semaphores + +inline constexpr uint32_t kPullBlockSize = 512; + +struct PullParams { + uint8_t* input_mc; // multicast VA of the symmetric input + const uint8_t* residual; // may be null (compile-time kHasResidual selects) + Semaphore* sem_local; // this rank's v2 pull semaphores (poll side) + uint8_t* sem_mc; // multicast VA of the pull-semaphore region + uint32_t rank; + uint32_t world_size; + uint32_t num_vecs; // 16B vectors + // pull_norm only: RMSNorm epilogue over the first num_norm_rows rows of + // the [num_vecs / kNormRowVecs, kNormDim] row view + const uint8_t* norm_weight; + float norm_eps; + uint32_t num_norm_rows; +}; + +// K3 pull barriers reuse the v2 pull-semaphore slots with EXACTLY the +// generic kernels' reservation protocol — reserve a 2 * world_size flag +// window on the local m_counter, signal arrival on m_flag, wait for the +// window to fill — except that each arrival is ONE `multimem.red.add` on +// the m_flag's multicast alias instead of world_size per-peer unicast reds +// (same aggregate effect: every rank's flag gains world_size arrivals per +// phase). Identical memory effects per call, so both kernel families share +// the slots freely (single-stream calls are serialized). +// +// The multicast alias of Semaphore::m_flag (the struct's first member): +SGL_DEVICE uint32_t* pull_sem_mc_flag(uint8_t* sem_mc, uint32_t block) { + static_assert(sizeof(Semaphore) == 128); + return reinterpret_cast(sem_mc + block * sizeof(Semaphore)); +} + +// enter barrier (relaxed): reserve this call's flag window, signal arrival +// with one multicast red, poll the local flag until all world_size arrivals +// landed — every rank's producer has finished writing the input. The +// reservation atomicAdd sits BEFORE the PDL wait: it is safe there (the +// previous same-slot call's reservation completed at ITS enter, which +// precedes its launch_dependents and hence this kernel's start, so windows +// are handed out in stream order) and it keeps the RMW latency off the +// post-wait critical path. The red must stay AFTER the wait — it asserts +// the producer grid has flushed. Returns the window base for the exit +// barrier — meaningful in thread 0 only, the sole barrier poller. +template +SGL_DEVICE uint32_t pull_barrier_enter(const PullParams& params) { + uint32_t current = 0; + if (threadIdx.x == 0) { + const auto semaphore = ¶ms.sem_local[blockIdx.x]; + const auto reserved = semaphore->counter_ptr()->inc(2 * params.world_size); + current = reserved + params.world_size; + device::PDLWaitPrimary(); + multimem_red_add_relaxed(pull_sem_mc_flag(params.sem_mc, blockIdx.x)); + while (semaphore->get_relaxed() - reserved < params.world_size) + ; + } + __syncthreads(); + return current; +} + +// exit barrier (release/acquire): every peer has finished reading my buffer +// (and, for 2shot, its broadcast into it is visible) before my next kernel +// may touch it. Mirrors AllReducePullImpl::sync_exit_pull. +template +SGL_DEVICE void pull_barrier_exit(const PullParams& params, uint32_t current) { + device::PDLTriggerSecondary(); + __syncthreads(); + if (threadIdx.x == 0) { + const auto semaphore = ¶ms.sem_local[blockIdx.x]; + multimem_red_add_release(pull_sem_mc_flag(params.sem_mc, blockIdx.x)); + while (semaphore->get_acquire() - current < params.world_size) + ; + } +} + +// One pipelined pass at width kWidth (kWidth multimem loads in flight per +// thread), then recurse to kWidth/2 for the remainder, down to a plain +// unroll-1 tail. The cascade matters for mid sizes: a shard smaller than +// kUnroll * step would otherwise skip the main loop entirely and run the +// whole range with a single request in flight. +template +SGL_DEVICE void +pull_reduce_pass(uint32_t& vid, const uint32_t num_vecs, const uint32_t step, uint8_t* mc_ptr, const uint8_t* res_ptr) { + using vec_t = device::AlignedVector; + using SumOp = device::ReductionTrait; + for (; vid + (kWidth - 1) * step < num_vecs; vid += kWidth * step) { + vec_t vec[kWidth]; +#pragma unroll + for (uint32_t u = 0; u < kWidth; ++u) { + ld_multimem_16B(vec[u], mc_ptr, vid + u * step); + } + if constexpr (kHasResidual) { +#pragma unroll + for (uint32_t u = 0; u < kWidth; ++u) { + vec_t res_vec; + res_vec.load(res_ptr, vid + u * step); +#pragma unroll + for (uint32_t j = 0; j < 4; ++j) { + vec[u][j] = SumOp::reduce(vec[u][j], res_vec[j]); + } + } + } +#pragma unroll + for (uint32_t u = 0; u < kWidth; ++u) { + st_multimem_16B(vec[u], mc_ptr, vid + u * step); + } + } + if constexpr (kWidth > 1) { + pull_reduce_pass(vid, num_vecs, step, mc_ptr, res_ptr); + } +} + +template +__global__ +__launch_bounds__(kPullBlockSize, 1) void all_reduce_pull_res_kernel(const __grid_constant__ PullParams params) { + static_assert(1 <= kUnroll && kUnroll <= 16 && (kUnroll & (kUnroll - 1)) == 0); + + const auto tx = threadIdx.x; + const auto bx = blockIdx.x; + const auto barrier_window = pull_barrier_enter(params); + + // this rank's shard of the 16B-vector range + const auto r = params.rank; + const auto avg_vecs = params.num_vecs / params.world_size; + const auto rem_vecs = params.num_vecs % params.world_size; + const auto vec_bias = int64_t(avg_vecs) * r + min(r, rem_vecs); + const auto num_vecs = avg_vecs + (r < rem_vecs ? 1 : 0); + const auto mc_ptr = params.input_mc + vec_bias * 16; + const auto res_ptr = kHasResidual ? params.residual + vec_bias * 16 : nullptr; + + // deep-pipelined body: issue up to kUnroll multimem loads back-to-back + // before the first (residual add and) store, keeping kUnroll requests in + // flight per thread to hide the NVLink latency with very few blocks; the + // remainder cascades through halving widths down to 1. + const auto step = kPullBlockSize * gridDim.x; + auto vid = bx * kPullBlockSize + tx; + pull_reduce_pass(vid, num_vecs, step, mc_ptr, res_ptr); + + pull_barrier_exit(params, barrier_window); +} + +// Fused RMSNorm over the latent of the K3 latent|shared MoE buffer: the +// row-structured counterpart of the res kernel with kUnroll ROWS in flight +// per block. One block pass covers kUnroll consecutive rows: 448 threads +// each ld_reduce one 16B vector per row back-to-back, warp partials of the +// normed rows go to smem (one __syncthreads per group), non-norm rows store +// immediately. smem is parity double-buffered across groups so the next +// group's partial writes can't race this group's reads (the WAR pair is two +// groups apart and separated by the intervening group's barrier). Works on +// this rank's row shard, in place. +template +__global__ +__launch_bounds__(kNormRowVecs, 1) void all_reduce_pull_norm_kernel(const __grid_constant__ PullParams params) { + using vec_t = device::AlignedVector; + using namespace device; + + const auto tx = threadIdx.x; + const auto bx = blockIdx.x; + const auto barrier_window = pull_barrier_enter(params); + + // this rank's shard of the row range + const auto num_rows = params.num_vecs / kNormRowVecs; + const auto r = params.rank; + const auto avg_rows = num_rows / params.world_size; + const auto rem_rows = num_rows % params.world_size; + const auto row_bias = avg_rows * r + min(r, rem_rows); + const auto my_rows = avg_rows + (r < rem_rows ? 1 : 0); + + vec_t wvec; + wvec.load(params.norm_weight, tx); + __shared__ float smem[2][kUnroll][kNormWarps]; + const auto warp = tx / kWarpThreads; + const auto lane = tx % kWarpThreads; + uint32_t parity = 0; + + const auto row_step = gridDim.x * kUnroll; + for (auto row0 = bx * kUnroll; row0 < my_rows; row0 += row_step) { + const auto cnt = min(kUnroll, my_rows - row0); + const auto vid0 = int64_t(row_bias + row0) * kNormRowVecs + tx; + vec_t vec[kUnroll]; +#pragma unroll + for (uint32_t u = 0; u < kUnroll; ++u) { + if (u < cnt) ld_multimem_16B(vec[u], params.input_mc, vid0 + u * kNormRowVecs); + } + // norm rows: push this warp's partial sum of squares to smem; non-norm + // rows (the shared 2/3) don't wait for the barrier — store right away + auto& sm = smem[parity]; + parity ^= 1; +#pragma unroll + for (uint32_t u = 0; u < kUnroll; ++u) { + if (u >= cnt) continue; // block-uniform + if (row_bias + row0 + u < params.num_norm_rows) { + float sum_of_squares = 0.0f; +#pragma unroll + for (uint32_t j = 0; j < 4; ++j) { + const auto [a, b] = cast(vec[u][j]); + sum_of_squares += a * a + b * b; + } + sum_of_squares = warp::reduce_sum(sum_of_squares); + if (lane == 0) sm[u][warp] = sum_of_squares; + } else { + st_multimem_16B(vec[u], params.input_mc, vid0 + u * kNormRowVecs); + } + } + __syncthreads(); +#pragma unroll + for (uint32_t u = 0; u < kUnroll; ++u) { + if (u >= cnt || row_bias + row0 + u >= params.num_norm_rows) continue; // block-uniform + float total = 0.0f; +#pragma unroll + for (uint32_t w = 0; w < kNormWarps; ++w) { + total += sm[u][w]; + } + const auto norm_factor = math::rsqrt(total / kNormDim + params.norm_eps); +#pragma unroll + for (uint32_t j = 0; j < 4; ++j) { + const auto [a, b] = cast(vec[u][j]); + const auto [wa, wb] = cast(wvec[j]); + vec[u][j] = cast(fp32x2_t{a * norm_factor * wa, b * norm_factor * wb}); + } + st_multimem_16B(vec[u], params.input_mc, vid0 + u * kNormRowVecs); + } + } + + pull_barrier_exit(params, barrier_window); +} + +} // namespace sglang + +using namespace sglang; + +// Host entry points + +template +struct AllReduceFusionKernel { + private: + using TensorView = tvm::ffi::TensorView; + + template + static constexpr auto res_push_kernel = all_reduce_push_res_kernel; + + static FusionParams + make_params(const host::distributed::CommunicatorObj& data, TensorView input, std::optional residual) { + using namespace host; + SymbolicSize N = {"num_elements"}; + SymbolicDevice device; + device.set_options(); + if (residual.has_value()) { + TensorMatcher({N}) // + .with_dtype() + .with_device(device) + .verify(input) + .verify(residual.value()); + } else { + TensorMatcher({N}) // + .with_dtype() + .with_device(device) + .verify(input); + } + const auto num_elems = N.unwrap(); + CHECK_HOST(data.world_size == kWorldSize); + CHECK_HOST(num_elems > 0 && num_elems % 8 == 0); + FusionParams params{}; + params.input = static_cast(input.data_ptr()); + params.residual = residual.has_value() ? static_cast(residual.value().data_ptr()) : nullptr; + params.push_ws_mc = nullptr; + params.push_ws_local = data.push_workspaces[data.rank]; + params.push_counter = data.push_counter; + params.push_buffer_stride = data.push_bytes; + params.rank = data.rank; + params.num_vecs = static_cast(num_elems / 8); + return params; + } + + /// The input viewed as rows of the norm width (3584 bf16 each); the first + /// num_norm_rows get the RMSNorm epilogue, the rest are a plain allreduce + /// (K3 uses [N latent rows | 2N shared rows] with num_norm_rows = N, or a + /// latent-only [N, 3584] tensor with num_norm_rows = N). + static FusionParams make_params_norm( + const host::distributed::CommunicatorObj& data, + TensorView input, + TensorView weight, + float eps, + int64_t num_norm_rows) { + using namespace host; + auto params = make_params(data, input, std::nullopt); + SymbolicDevice device; + device.set_options(); + TensorMatcher({kNormDim}).with_dtype().with_device(device).verify(weight); + CHECK_HOST(params.num_vecs % kNormRowVecs == 0) + << "numel must be a multiple of " << kNormDim << ", got " << int64_t(params.num_vecs) * 8; + const auto num_rows = params.num_vecs / kNormRowVecs; + CHECK_HOST(0 <= num_norm_rows && num_norm_rows <= num_rows) + << "num_norm_rows " << num_norm_rows << " out of range [0, " << num_rows << "]"; + params.norm_weight = static_cast(weight.data_ptr()); + params.norm_eps = static_cast(eps); + params.num_norm_rows = static_cast(num_norm_rows); + params.num_push_counters = data.num_push_blocks; + return params; + } + + // Shared pull validation; the reduce is in place on the symmetric input. + static PullParams make_pull_params( + const host::distributed::CommunicatorObj& data, + TensorView input, + std::optional residual, + int64_t input_mc_ptr, + int64_t sem_mc_ptr) { + using namespace host; + SymbolicSize N = {"num_elements"}; + SymbolicDevice device; + device.set_options(); + if (residual.has_value()) { + TensorMatcher({N}) // + .with_dtype() + .with_device(device) + .verify(input) + .verify(residual.value()); + } else { + TensorMatcher({N}) // + .with_dtype() + .with_device(device) + .verify(input); + } + const auto num_elems = N.unwrap(); + CHECK_HOST(data.world_size == kWorldSize); + CHECK_HOST(num_elems > 0 && num_elems % 8 == 0) << "numel must be a positive multiple of 8, got " << num_elems; + // headroom below 2^32 so the unrolled loop's `vid + (kUnroll-1)*step` + // arithmetic can never wrap around u32 + CHECK_HOST(num_elems / 8 < (int64_t(1) << 31)) << "numel exceeds the 16B-vector limit"; + CHECK_HOST(input_mc_ptr != 0) << "pull requires the input's multicast address"; + CHECK_HOST(sem_mc_ptr != 0) << "pull requires the semaphores' multicast address"; + PullParams params{}; + params.input_mc = reinterpret_cast(static_cast(input_mc_ptr)); + params.residual = residual.has_value() ? static_cast(residual.value().data_ptr()) : nullptr; + params.sem_local = data.pull_semaphores[data.rank]; + params.sem_mc = reinterpret_cast(static_cast(sem_mc_ptr)); + params.rank = data.rank; + params.world_size = data.world_size; + params.num_vecs = static_cast(num_elems / 8); + return params; + } + + // Runtime unroll dispatch: every supported width is compiled into the + // module so the tuned per-size unroll needs no extra JIT builds. + template + static void launch_pull_res(const PullParams& params, int64_t num_blocks, int64_t unroll, DLDevice device) { + const auto run = [&](auto kernel) { + host::LaunchKernel(static_cast(num_blocks), kPullBlockSize, device).enable_pdl(kUsePDL)(kernel, params); + }; + switch (unroll) { + case 2: + return run(all_reduce_pull_res_kernel<2, kHasResidual, kUsePDL>); + case 4: + return run(all_reduce_pull_res_kernel<4, kHasResidual, kUsePDL>); + case 8: + return run(all_reduce_pull_res_kernel<8, kHasResidual, kUsePDL>); + case 16: + return run(all_reduce_pull_res_kernel<16, kHasResidual, kUsePDL>); + default: + CHECK_HOST(false) << "unsupported unroll " << unroll << " (must be 2, 4, 8, or 16)"; + } + } + + static void launch_pull_norm(const PullParams& params, int64_t num_blocks, int64_t unroll, DLDevice device) { + const auto run = [&](auto kernel) { + host::LaunchKernel(static_cast(num_blocks), kNormRowVecs, device).enable_pdl(kUsePDL)(kernel, params); + }; + switch (unroll) { + case 2: + return run(all_reduce_pull_norm_kernel<2, kUsePDL>); + case 4: + return run(all_reduce_pull_norm_kernel<4, kUsePDL>); + case 8: + return run(all_reduce_pull_norm_kernel<8, kUsePDL>); + case 16: + return run(all_reduce_pull_norm_kernel<16, kUsePDL>); + default: + CHECK_HOST(false) << "unsupported unroll " << unroll << " (must be 2, 4, 8, or 16)"; + } + } + + public: + static void push_res(CommunicatorRef ref, TensorView input, std::optional residual, int64_t ws_mc_base) { + const auto& data = *ref.get(); + auto params = make_params(data, input, residual); + CHECK_HOST(ws_mc_base != 0) << "push requires a multicast-capable workspace"; + const int64_t nbytes = int64_t(params.num_vecs) * 16; + CHECK_HOST(nbytes <= data.push_bytes) + << "input size " << nbytes << " exceeds push workspace size " << data.push_bytes; + params.push_ws_mc = reinterpret_cast(static_cast(ws_mc_base)); + const auto kernel = residual.has_value() ? res_push_kernel : res_push_kernel; + host::LaunchKernel(data.num_push_blocks, choose_block_size(params.num_vecs), input.device()) + .enable_pdl(kUsePDL)(kernel, params); + } + + static void push_norm( + CommunicatorRef ref, TensorView input, TensorView weight, float eps, int64_t num_norm_rows, int64_t ws_mc_base) { + constexpr auto kClusterSize = 7; + const auto& data = *ref.get(); + auto params = make_params_norm(data, input, weight, eps, num_norm_rows); + CHECK_HOST(ws_mc_base != 0) << "push requires a multicast-capable workspace"; + const int64_t nbytes = int64_t(params.num_vecs) * 16; + CHECK_HOST(nbytes <= data.push_bytes) + << "input size " << nbytes << " exceeds push workspace size " << data.push_bytes; + params.push_ws_mc = reinterpret_cast(static_cast(ws_mc_base)); + const auto num_rows = params.num_vecs / kNormRowVecs; + constexpr uint32_t kMaxClusters = 96; + const auto num_row_clusters = std::max(std::min(num_rows, kMaxClusters), 1); + CHECK_HOST(num_row_clusters < data.num_push_blocks); + host::LaunchKernel((num_row_clusters + 1) * kClusterSize, kNormRowVecs / kClusterSize, input.device()) + .enable_pdl(kUsePDL)(all_reduce_push_norm_cluster_kernel, params); + } + + /// Deferred MoE finalize + 1shot push all-reduce + RMSNorm over EVERY row. + /// `out` (flattened [num_tokens * kNormDim] bf16) is output-only: each + /// rank's partial latent is computed from the trtllm-gen deferred-finalize + /// triple during the staging pass and never materializes in global memory. + static void finalize_push_norm( + CommunicatorRef ref, + TensorView out, + TensorView gemm2_out, + TensorView permuted_idx, + TensorView expert_weights, + TensorView weight, + float eps, + int64_t ws_mc_base) { + using namespace host; + constexpr auto kClusterSize = 7; + const auto& data = *ref.get(); + // every row of the latent-only output is normed + auto params = make_params_norm(data, out, weight, eps, out.size(0) / kNormDim); + const auto num_tokens = params.num_vecs / kNormRowVecs; + + auto P = SymbolicSize{"num_permuted_rows"}; + auto T = SymbolicSize{"num_tokens"}; + T.set_value(num_tokens); + auto K = SymbolicSize{"top_k"}; + auto TK = SymbolicSize{"num_expanded"}; + TK.set_value(static_cast(num_tokens) * kFinTopK); + SymbolicDevice device; + device.set_options(); + TensorMatcher({P, kNormDim}).with_dtype().with_device(device).verify(gemm2_out); + TensorMatcher({T, K}).with_dtype().with_device(device).verify(expert_weights); + TensorMatcher({TK}).with_dtype().with_device(device).verify(permuted_idx); + CHECK_HOST(K.unwrap() == kFinTopK) << "finalize_push_norm is specialized for top_k = " << kFinTopK; + + CHECK_HOST(ws_mc_base != 0) << "push requires a multicast-capable workspace"; + const int64_t nbytes = int64_t(params.num_vecs) * 16; + CHECK_HOST(nbytes <= data.push_bytes) + << "output size " << nbytes << " exceeds push workspace size " << data.push_bytes; + params.push_ws_mc = reinterpret_cast(static_cast(ws_mc_base)); + params.fin_gemm2 = static_cast(gemm2_out.data_ptr()); + params.fin_idx = static_cast(permuted_idx.data_ptr()); + params.fin_weights = static_cast(expert_weights.data_ptr()); + + constexpr uint32_t kMaxClusters = 96; + const auto num_row_clusters = std::max(std::min(num_tokens, kMaxClusters), 1); + CHECK_HOST(num_row_clusters < data.num_push_blocks); + host::LaunchKernel((num_row_clusters + 1) * kClusterSize, kNormRowVecs / kClusterSize, out.device()) + .enable_pdl(kUsePDL)( + all_reduce_push_norm_cluster_kernel, params); + } + + /// Low-SM NVLS pull (+ optional residual): in-place reduce-scatter + + /// broadcast on the symmetric input. `sem_mc_ptr` is the multicast VA of + /// the v2 pull-semaphore region; num_blocks — which must be uniform + /// across ranks per call — is clamped to the semaphore capacity. + static void pull_res( + CommunicatorRef ref, + TensorView input, + std::optional residual, + int64_t input_mc_ptr, + int64_t sem_mc_ptr, + int64_t num_blocks, + int64_t unroll) { + const auto& data = *ref.get(); + const auto params = make_pull_params(data, input, residual, input_mc_ptr, sem_mc_ptr); + CHECK_HOST(num_blocks >= 1) << "invalid num_blocks: " << num_blocks; + num_blocks = std::min(num_blocks, data.num_pull_blocks); + if (residual.has_value()) { + launch_pull_res(params, num_blocks, unroll, input.device()); + } else { + launch_pull_res(params, num_blocks, unroll, input.device()); + } + } + + /// Low-SM NVLS pull + RMSNorm over the latent of the K3 latent|shared MoE + /// buffer ([num_tokens, 3584] latent then [num_tokens, 7168] shared); + /// num_tokens and the normed row range are derived from the element count. + /// Same semaphore / num_blocks semantics as pull_res. + static void pull_norm( + CommunicatorRef ref, + TensorView input, + TensorView weight, + double eps, + int64_t num_norm_rows, + int64_t input_mc_ptr, + int64_t sem_mc_ptr, + int64_t num_blocks, + int64_t unroll) { + const auto& data = *ref.get(); + auto params = make_pull_params(data, input, std::nullopt, input_mc_ptr, sem_mc_ptr); + CHECK_HOST(num_blocks >= 1) << "invalid num_blocks: " << num_blocks; + num_blocks = std::min(num_blocks, data.num_pull_blocks); + using namespace host; + SymbolicDevice device; + device.set_options(); + TensorMatcher({kNormDim}).with_dtype().with_device(device).verify(weight); + CHECK_HOST(params.num_vecs % kNormRowVecs == 0) + << "numel must be a multiple of " << kNormDim << ", got " << int64_t(params.num_vecs) * 8; + const auto num_rows = params.num_vecs / kNormRowVecs; + CHECK_HOST(0 <= num_norm_rows && num_norm_rows <= num_rows) + << "num_norm_rows " << num_norm_rows << " out of range [0, " << num_rows << "]"; + params.norm_weight = static_cast(weight.data_ptr()); + params.norm_eps = static_cast(eps); + params.num_norm_rows = static_cast(num_norm_rows); + launch_pull_norm(params, num_blocks, unroll, input.device()); + } +}; diff --git a/python/sglang/kernels/jit/csrc/kimi_k3/comm/gemm_ag.cuh b/python/sglang/kernels/jit/csrc/kimi_k3/comm/gemm_ag.cuh new file mode 100644 index 000000000..adf89784d --- /dev/null +++ b/python/sglang/kernels/jit/csrc/kimi_k3/comm/gemm_ag.cuh @@ -0,0 +1,303 @@ +#include +#include + +#include +#include +#include +#include + +#include + +#include +#include + +#include "ptx_sys.cuh" +#include +#include +#include +#include + +namespace sglang { + +namespace gemm_ag { + +using device::distributed::Counter; +using device::distributed::multimem_store_relaxed; + +constexpr uint32_t kWorld = 8; // TP world size +constexpr uint32_t kVecSize = 32 / sizeof(bf16_t); // 16 bf16 per 32B vector +constexpr uint32_t kSpinBlock = 128; // consumer threads per block +constexpr uint32_t kSpinVec = 16 / sizeof(bf16_t); // 8 bf16 (16B) per consumer thread + +// Producer: per-rank column-slice GEMV, multicast store with Lamport markers. + +struct ProducerParams { + uint8_t* ws_mc; // multicast VA of the push workspace base + Counter* counter; // per-block phase counters (READ only here) + uint32_t half_bytes; // bytes per phase half (world_size * push_bytes) + uint32_t rank; +}; + +template +__global__ __launch_bounds__(K / kVecSize) void gemm_ag_gemv_kernel( + const __grid_constant__ ProducerParams params, + const bf16_t* __restrict__ x, // [M, K] + const bf16_t* __restrict__ weight) { // [N, K] FULL replicated weight + using namespace device; + using vec_t = AlignedVector; + constexpr uint32_t kNLocal = N / kWorld; // columns computed by this rank + constexpr uint32_t kGemvBlock = K / kVecSize; + constexpr uint32_t kNumWarps = kGemvBlock / kWarpThreads; + static_assert(K % kVecSize == 0, "K must be a multiple of the 32B vector width"); + static_assert(kGemvBlock % kWarpThreads == 0, "K / vec_size must fill whole warps"); + static_assert(kGemvBlock <= 1024, "K / vec_size exceeds the maximum block size"); + static_assert(N % kWorld == 0, "N must split evenly over the TP world"); + static_assert(kNLocal % N_SPLIT == 0, "the local column slice must split into whole tiles"); + static_assert(M * N_SPLIT <= kGemvBlock, "output tile must fit one thread each for the final reduce"); + static_assert(N_SPLIT % 2 == 0, "epilogue stores adjacent column pairs"); + + const uint32_t bx = blockIdx.x; + const uint32_t tx = threadIdx.x; + // this rank's rows of the replicated [N, K] weight, sliced HERE (the + // Python side always hands the full weight) + const bf16_t* weight_tile = weight + (params.rank * kNLocal + bx * N_SPLIT) * K; + + // weight prefetch before the PDL wait (input-independent addresses) + vec_t weight_vec[N_SPLIT]; +#pragma unroll + for (uint32_t n = 0; n < N_SPLIT; ++n) { + weight_vec[n].load(weight_tile + n * K, tx); + } + + PDLWaitPrimary(); + // Every push-workspace consumer flips the WHOLE counter array each round + // (each has a tail loop up to num_counters), so all counters hold the same + // phase at this point and counter[0] is equivalent to counter[bx]. Reading a + // single counter is what frees the producer grid from num_push_blocks. + const uint32_t phase = params.counter[0].get() & 1; + + vec_t input_vec[M]; +#pragma unroll + for (uint32_t m = 0; m < M; ++m) { + input_vec[m].load(x + m * K, tx); + } + + __shared__ alignas(16) float s_acc[kNumWarps][M * N_SPLIT]; + const uint32_t warp_id = tx / kWarpThreads; + +#pragma unroll + for (uint32_t m = 0; m < M; ++m) { +#pragma unroll + for (uint32_t n = 0; n < N_SPLIT; ++n) { + float acc = 0.0f; +#pragma unroll + for (uint32_t i = 0; i < kVecSize; ++i) { + acc = device::math::fma_f32_bf16(input_vec[m][i], weight_vec[n][i], acc); + } + s_acc[warp_id][m * N_SPLIT + n] = warp::reduce_sum(acc); + } + } + __syncthreads(); + constexpr uint32_t kNumPairs = M * N_SPLIT / 2; + if (tx < kNumPairs) { + auto packed = load_as(s_acc[0], tx); +#pragma unroll + for (uint32_t i = 1; i < kNumWarps; ++i) { + const auto [lo, hi] = load_as(s_acc[i], tx); + packed.x += lo; + packed.y += hi; + } + const auto pair = cast(packed); + auto bits = *reinterpret_cast(&pair); + if (bits == 0) bits = 0x8000u; // -0.0 in the first element: never all-zero + const uint32_t m = (2 * tx) / N_SPLIT; + const uint32_t n = (2 * tx) % N_SPLIT; // even column within the tile + // bf16 index in the phase half's dense [world][M][N / world] prefix; + // one multicast store lands this rank's pair on EVERY peer + const uint32_t elem = (params.rank * M + m) * kNLocal + bx * N_SPLIT + n; + const auto base = reinterpret_cast(params.ws_mc + phase * params.half_bytes); + const auto dst = reinterpret_cast(base + elem); + multimem_store_relaxed(dst, bits); + } + PDLTriggerSecondary(); +} + +// Consumer: Lamport spin + add3, one 16B vector (8 bf16) per thread. + +struct ConsumerParams { + uint8_t* ws_local; // LOCAL VA of the push workspace base (poll + reset) + Counter* counter; // per-block phase counters (read + flip) + uint32_t num_counters; // full counter array size (num_push_blocks) + uint32_t half_bytes; // bytes per phase half (world_size * push_bytes) + const bf16_t* b; // [M, N] + const bf16_t* c; // may be null + bf16_t* out; // [M, N] + uint32_t num_rows; // M +}; + +template +__global__ void spin_add3_kernel(const __grid_constant__ ConsumerParams params) { + using namespace device; + using vec_t = AlignedVector; // 8 bf16 as 4 pairs + constexpr uint32_t kNLocal = N / kWorld; + static_assert(N % kSpinVec == 0, "rows must stay 16B aligned"); + static_assert(kNLocal % kSpinVec == 0, "a vector must never cross a rank block"); + const auto bx = blockIdx.x; + const auto tx = threadIdx.x; + const uint32_t tid = bx * kSpinBlock + tx; + const uint32_t elem = tid * kSpinVec; // first bf16 of this thread's vector + const uint32_t phase = params.counter[bx].get() & 1; + + PDLTriggerSecondary(); + + // use the last block to clean up: it flips ITS OWN counter and every one + // past the grid (work blocks flip [0, num_blocks - 1) themselves) + if (const auto num_blocks = gridDim.x; bx == num_blocks - 1) { + [[unlikely]]; + __syncthreads(); // ensure phase is ready for all threads + for (uint32_t i = num_blocks - 1 + tx; i < params.num_counters; i += kSpinBlock) { + params.counter[i].set(phase ^ 1); + } + return void(); // this block is done, no output to write + } + + // Deliberately NO PDLWaitPrimary: the dependency is carried through data + if (elem < params.num_rows * N) { + const auto row = elem / N; + const auto col = elem % N; + // out[row, col] lives at half[col / kNLocal][row][col % kNLocal] of the + // dense [world][M][N / world] prefix of the current phase half + const auto base = reinterpret_cast(params.ws_local + phase * params.half_bytes); + const auto src = base + ((col / kNLocal) * params.num_rows + row) * kNLocal + col % kNLocal; + vec_t b_vec, c_vec; + b_vec.load(params.b + elem); + if constexpr (kHasC) c_vec.load(params.c + elem); + // spin until all 4 packed pairs of the vector have landed + uint4 raw; + do { + asm volatile("ld.relaxed.gpu.global.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(raw.x), "=r"(raw.y), "=r"(raw.z), "=r"(raw.w) + : "l"(src) + : "memory"); + } while (raw.x == 0 || raw.y == 0 || raw.z == 0 || raw.w == 0); + const auto& gathered = *reinterpret_cast(&raw); + vec_t out_vec; +#pragma unroll + for (uint32_t j = 0; j < kSpinVec / 2; ++j) { + using Trait = DTypeTrait; + out_vec[j] = Trait::add(gathered[j], b_vec[j]); + if constexpr (kHasC) out_vec[j] = Trait::add(out_vec[j], c_vec[j]); + } + out_vec.store(params.out + elem); + AlignedVector zero; + zero.fill(0); + zero.store(src); + } + __syncthreads(); + if (tx == 0) params.counter[bx].set(phase ^ 1); +} + +} // namespace gemm_ag +} // namespace sglang + +using namespace sglang; +using host::distributed::CommunicatorRef; + +// Host entry point (tiny_gemm style: one GEMV instantiation per M in +// [1, kMaxM] selected through a constexpr function-pointer table, then the +// spin consumer launched with PDL right behind it). Any (K, N) that passes +// the kernels' static_asserts works; Kimi-K3 uses (3584, 7168). + +template +struct GEMMAGKernel { + using TensorView = tvm::ffi::TensorView; + + // Columns of this rank's slice per producer block; sets the grid to + // kNLocal / N_SPLIT. Measured on 2x4 GB300 TP8 at bs=1 (three reps each, + // mean TPOT): 8 -> grid 112, 8.36 ms; 4 -> 224, 8.29 ms; 2 -> 448, 8.21 ms. + // Standalone GEMV at the same shape: 4.15 / 3.20 / 2.56 us, and 16 -> 4.42 us, + // so the trend is monotonic and 2 is the floor (the epilogue stores column + // pairs, so N_SPLIT must stay even). 8 used to be the largest grid that fit + // the old kNumProducerBlocks <= num_push_blocks bound; the producer now reads + // a single phase counter, so the grid is free and 112 blocks did not even + // fill one per SM. + static constexpr uint32_t N_SPLIT = 2; + static constexpr uint32_t kNLocal = N / gemm_ag::kWorld; + static constexpr uint32_t kGemvBlock = K / gemm_ag::kVecSize; + static constexpr uint32_t kNumProducerBlocks = kNLocal / N_SPLIT; + static_assert(kNLocal % N_SPLIT == 0); + + using GemvFn = void (*)(gemm_ag::ProducerParams, const bf16_t*, const bf16_t*); + + template + static constexpr auto make_table(std::index_sequence) { + return std::array{nullptr, gemm_ag::gemm_ag_gemv_kernel...}; + } + static constexpr auto kGemvTable = make_table(std::make_index_sequence{}); + + static void + run(CommunicatorRef ref, + TensorView x, + TensorView weight, + TensorView b, + std::optional c, + TensorView out, + intptr_t ws_mc_base) { + using namespace host; + const auto& data = *ref.get(); + + auto M = SymbolicSize{"num_tokens"}; + auto device = SymbolicDevice{}; + device.set_options(); + TensorMatcher({M, K}).with_dtype().with_device(device).verify(x); + TensorMatcher({N, K}).with_dtype().with_device(device).verify(weight); + TensorMatcher({M, N}).with_dtype().with_device(device).verify(b); + if (c.has_value()) { + TensorMatcher({M, N}).with_dtype().with_device(device).verify(c.value()); + } + TensorMatcher({M, N}).with_dtype().with_device(device).verify(out); + const auto num_tokens = static_cast(M.unwrap()); + CHECK_HOST(num_tokens >= 1 && num_tokens <= kMaxM); + CHECK_HOST(data.world_size == gemm_ag::kWorld) << "the kernel is compiled for TP" << gemm_ag::kWorld; + CHECK_HOST(ws_mc_base != 0) << "requires a multicast-capable workspace"; + CHECK_HOST(int64_t(num_tokens) * kNLocal * 2 <= data.push_bytes) + << "staging slice exceeds the push slot size " << data.push_bytes; + // The producer grid is no longer bound to the counter array: it reads only + // counter[0] (see gemm_ag_gemv_kernel). The consumer grid still is. + CHECK_HOST(data.num_push_blocks > 0) << "no push blocks available"; + // producer: GEMV + const auto producer_params = gemm_ag::ProducerParams{ + .ws_mc = reinterpret_cast(ws_mc_base), + .counter = data.push_counter, + .half_bytes = static_cast(data.push_bytes * data.world_size), + .rank = data.rank, + }; + LaunchKernel(kNumProducerBlocks, kGemvBlock, device.unwrap()) + .enable_pdl(kUsePDL)( + kGemvTable[num_tokens], + producer_params, + static_cast(x.data_ptr()), + static_cast(weight.data_ptr())); + + // consumer: spin + add3 + const auto consumer_params = gemm_ag::ConsumerParams{ + .ws_local = data.push_workspaces[data.rank], + .counter = data.push_counter, + .num_counters = data.num_push_blocks, + .half_bytes = static_cast(data.push_bytes * data.world_size), + .b = static_cast(b.data_ptr()), + .c = c.has_value() ? static_cast(c.value().data_ptr()) : nullptr, + .out = static_cast(out.data_ptr()), + .num_rows = num_tokens, + }; + const auto num_vecs = num_tokens * N / gemm_ag::kSpinVec; + const auto num_consumers = host::div_ceil(num_vecs, gemm_ag::kSpinBlock); + CHECK_HOST(num_consumers + 1 <= data.num_push_blocks); + // use last block to clean up the counter + const auto num_consumer_blocks = num_consumers + 1; + using gemm_ag::spin_add3_kernel; + const auto kernel = c.has_value() ? spin_add3_kernel : spin_add3_kernel; + host::LaunchKernel(num_consumer_blocks, gemm_ag::kSpinBlock, device.unwrap()) + .enable_pdl(kUsePDL)(kernel, consumer_params); + } +}; diff --git a/python/sglang/kernels/jit/csrc/kimi_k3/comm/gemm_ar.cuh b/python/sglang/kernels/jit/csrc/kimi_k3/comm/gemm_ar.cuh new file mode 100644 index 000000000..0032a74de --- /dev/null +++ b/python/sglang/kernels/jit/csrc/kimi_k3/comm/gemm_ar.cuh @@ -0,0 +1,1569 @@ +// K3 fused o_proj GEMM + all-reduce for decode (bf16, TP row-parallel). +// +// CONTRACT (per rank r of R): out[M, 7168] = sum_r x_r[M, K] @ W_r[7168, K]^T +// bf16 in/out, fp32 accumulate, partials round to bf16 pre-sum (same +// semantics as the unfused GEMM + bf16 ring AR). +// M in [1, 512], rounded up to a tuned cell {8,16,32,64,128,256,512}; out +// must have `cell` rows — rows [M, cell) are clobbered with zeros. +// Comm plane: pure NVLink P2P (unicast pushes + per-rank flag reductions); +// one-shot AR below the two-shot threshold, two-shot RS+AG above. +// Requires SM100+ with full P2P; tuned on GB300 (sm_103a). + +#include +#include + +#include + +#include +#include + +#include "ptx_sys.cuh" +#include +#include +#include +#include +#include +#include + +// Local PTX / TMA primitives +// Only what this kernel issues, kept in-file on purpose: these are raw-ISA +// shapes (sm_100+ tcgen05, the cta_group::1 multicast TMA load) that no shared +// sglang header wraps, and splitting them out bought a dozen headers with +// exactly one consumer. Anything cute already provides goes through cute +// (`set_block_rank` below, the tensor-map driver wrapper in `w_maps`). + +namespace ptx { + +// ---- generic → shared address conversion (PTX ISA §10.4) -------------------- + +// ---- cvt: pack 2 fp32 into one bf16x2 (PTX ISA §9.7.9.21) ------------------ +// +// PACKED-PAIR ORDERING: `cvt.bf16x2.f32 d, a, b` puts cvt(a) in d's UPPER +// half and cvt(b) in the LOWER half. Read back from little-endian memory as a +// bf16 array, LOWER lands at column i and UPPER at i+1 — so for cells (c0, c1) +// destined for adjacent slots [i, i+1] pass `cvt_pack_f32x2_to(c1, c0)`. +struct bf16 { + using packed2_t = uint32_t; +}; + +template +static __device__ __forceinline__ typename Dst::packed2_t cvt_pack_f32x2_to(float a, float b); + +template <> +__device__ __forceinline__ uint32_t cvt_pack_f32x2_to(float a, float b) { + uint32_t d; + asm volatile("cvt.rn.bf16x2.f32 %0, %1, %2;" : "=r"(d) : "f"(a), "f"(b)); + return d; +} + +// ---- ldmatrix (PTX ISA §9.7.14.5.15) --------------------------------------- + +// Warp-collective load of 4 (8x8) BF16 matrices from smem into mma.sync +// fragments. `row_addr` is this lane's 16-byte-aligned row base: lanes 0-7 +// supply matrix 0, 8-15 matrix 1, 16-23 matrix 2, 24-31 matrix 3. +static SGL_DEVICE void ldmatrix_x4_b16(uint32_t row_addr, uint32_t& r0, uint32_t& r1, uint32_t& r2, uint32_t& r3) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared::cta.b16 {%0, %1, %2, %3}, [%4];" + : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) + : "r"(row_addr)); +} + +static SGL_DEVICE void ldmatrix_x2_b16(uint32_t row_addr, uint32_t& r0, uint32_t& r1) { + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared::cta.b16 {%0, %1}, [%2];" : "=r"(r0), "=r"(r1) : "r"(row_addr)); +} + +// ---- warp-level mma.sync (PTX ISA §9.7.14) --------------------------------- + +// D += A*B, bf16 x bf16 -> f32. The warp-register form: co-resides freely (no +// TMEM, no 1-CTA/SM cap), which is why the small-M members use it instead of +// tcgen05. +static __device__ __forceinline__ void +mma_m16n8k16_bf16f32(float4& d, uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, uint32_t b1) { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{%0,%1,%2,%3},{%4,%5,%6,%7},{%8,%9},{%0,%1,%2,%3};" + : "+f"(d.x), "+f"(d.y), "+f"(d.z), "+f"(d.w) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); +} + +// ---- tcgen05 matrix descriptors (PTX ISA §9.7.16.4) ------------------------ +// +// 0 = K-major (innermost dim is K, the "TN" convention), 1 = MN-major. +enum class Major : uint8_t { + K = 0, + MN = 1, +}; + +enum class F16Type : uint8_t { F16 = 0, BF16 = 1 }; // kind::f16 atype/btype +enum class DType : uint8_t { F16 = 0, F32 = 1, S32 = 2 }; + +// Smem matrix descriptor. The public PTX spec has errors at bits 46-60; the +// layout below is the one our gate tests verify: +// 0-13: start_address >> 4 16-29: leading_byte_offset >> 4 +// 32-45: stride_byte_offset >> 4 46-47: version (=1, Blackwell — larger +// shapes such as 128B-swizzle BLOCK_K>16 produce garbage at 0) +// 49-51: base_offset 52: lbo_mode (=0, relative byte offset) +// 61-63: layout_type (0=None, 2=128B, 4=64B, 6=32B) +__host__ __device__ static __forceinline__ constexpr uint64_t +mma_smem_desc(uint32_t matrix_addr, uint32_t lbo, uint32_t sbo, uint32_t base_offset, int swizzle_bytes) { + auto enc = [](uint32_t x) -> uint64_t { return (uint64_t)((x & 0x3FFFFu) >> 4); }; + uint8_t code = (swizzle_bytes == 128) ? 2u : (swizzle_bytes == 64) ? 4u : (swizzle_bytes == 32) ? 6u : 0u; + uint64_t d = 0; + d |= enc(matrix_addr); // bits 0-13 + d |= enc(lbo) << 16; // bits 16-29 + d |= enc(sbo) << 32; // bits 32-45 + d |= uint64_t(1u) << 46; // bits 46-47 = version = 1 + d |= uint64_t(base_offset & 0x7u) << 49; // bits 49-51 + d |= uint64_t(code & 0x7u) << 61; // bits 61-63 + return d; +} + +// K-major operand (A as (M, K), B as (N, K) in a TN GEMM). `T` is a size proxy +// (uint16_t for BF16/FP16); the dtype semantics live in the instruction +// descriptor. K-major mandates SWIZZLE_BYTES == BLOCK_K * sizeof(T). +template +__host__ __device__ static __forceinline__ constexpr uint64_t +mma_smem_desc_k_major(uint32_t addr, uint32_t base_offset = 0) { + constexpr int K_BYTES = BLOCK_K * int(sizeof(T)); + static_assert(SWIZZLE_BYTES == K_BYTES, "K-major requires swizzle bytes == BLOCK_K * sizeof(T)"); + return mma_smem_desc(addr, /*lbo=*/0u, /*sbo=*/8u * uint32_t(K_BYTES), base_offset, SWIZZLE_BYTES); +} + +// Instruction descriptor, kind::f16 (PTX ISA Table 44). +__host__ __device__ static __forceinline__ constexpr uint32_t mma_inst_desc_f16( + uint32_t M, + uint32_t N, + F16Type a_type = F16Type::BF16, + F16Type b_type = F16Type::BF16, + DType d_type = DType::F32, + Major a_major = Major::K, + Major b_major = Major::K, + bool negate_a = false, + bool negate_b = false) { + uint32_t d = 0; + d |= (static_cast(d_type) & 0x3u) << 4; + d |= (static_cast(a_type) & 0x7u) << 7; + d |= (static_cast(b_type) & 0x7u) << 10; + if (negate_a) d |= 1u << 13; + if (negate_b) d |= 1u << 14; + d |= (static_cast(a_major) & 0x1u) << 15; + d |= (static_cast(b_major) & 0x1u) << 16; + d |= ((N >> 3) & 0x3Fu) << 17; + d |= ((M >> 4) & 0x1Fu) << 24; + return d; +} + +// Cross-CTA arrive with `.release.cta` ordering. A plain cluster arrive carries +// NO release, so prior memory ops (notably a retired `tcgen05.ld` TMEM drain on +// the arriving warp) are not guaranteed visible before a peer warp's `.acquire` +// wait returns — and the peer then overwriting that TMEM is a real race (spec +// §8.8: a release pattern requires `mbarrier.arrive.release`). +// +// `cute::set_block_rank` is the `mapa.shared::cluster` that retargets a local +// smem offset at CTA `cta_rank`; the `.shared::cluster` qualifier on the arrive +// is what makes it cross-CTA (the plain `.shared` form hits the local mbar +// whatever the address bits say). +static SGL_DEVICE void mbar_arrive_cluster_release(uint64_t* bar, uint32_t cta_rank) { + const uint32_t mapped = cute::set_block_rank(to_shared(bar), cta_rank); + asm volatile("mbarrier.arrive.release.cta.shared::cluster.b64 _, [%0], 1;" ::"r"(mapped)); +} + +// ---- cluster / warp sync (PTX ISA §9.7.13, §9.7.4) ------------------------- + +// Cluster barrier with explicit release/acquire — the publish/observe boundary +// between `mbarrier.init` and any `.shared::cluster` use of those mbars. +static SGL_DEVICE void cluster_sync_rel_acq() { + asm volatile("barrier.cluster.arrive.release.aligned;"); + asm volatile("barrier.cluster.wait.acquire.aligned;"); +} + +// True on exactly one lane of the issuing warp — guards single-issuer sites +// (mbar init, TMA issue, MMA issue, TMEM alloc) without gating on lane_id. +static SGL_DEVICE bool elect_one() { + uint32_t pred; + asm volatile( + "{\n\t.reg .pred p;\n\t" + "elect.sync _|p, 0xffffffff;\n\t" + "selp.b32 %0, 1, 0, p;\n\t}\n" + : "=r"(pred)); + return pred != 0; +} + +static SGL_DEVICE uint32_t cluster_cta_rank() { + uint32_t rank; + asm("mov.u32 %0, %%cluster_ctarank;" : "=r"(rank)); + return rank; +} + +// ---- tcgen05 (PTX ISA §9.7.16) --------------------------------------------- +// +// Lifecycle (mandatory order, §9.7.16.7.1): alloc (one warp, n_cols a power of +// 2 in [32, 512], TMEM address written to smem) -> __syncthreads + read taddr +// -> ld/mma -> dealloc -> relinquish before kernel exit. +// +// Each warp can only touch its own 32-lane TMEM band (§9.7.16.8.1): warp 0 -> +// lanes 0-31, warp 1 -> 32-63, and so on. +// +// After an MMA, `tcgen05_commit_arrive` + `mbar_wait_parity` + +// `tcgen05_fence_after_thread_sync` before reading the result with +// `tcgen05_ld_*`; the fence is mandatory or the register reads may see stale +// values even though the mbarrier signaled "MMA done". +static SGL_DEVICE void tcgen05_alloc(uint32_t smem_addr_for_taddr, uint32_t n_cols) { + asm volatile( + "tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" ::"r"(smem_addr_for_taddr), "r"(n_cols)); +} + +static SGL_DEVICE void tcgen05_dealloc(uint32_t taddr, uint32_t n_cols) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" ::"r"(taddr), "r"(n_cols)); +} + +static SGL_DEVICE void tcgen05_relinquish() { + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); +} + +// .32x32b.x8: 8 b32 per lane = 8 TMEM columns. Per-lane 8 FP32 -> 4 bf16x2 +// packs = one int4, the natural fit for a BF16 epilogue draining a column band +// with 16-byte smem stores. +static SGL_DEVICE void tcgen05_ld_32x32b_x8( + uint32_t taddr, + uint32_t& r0, + uint32_t& r1, + uint32_t& r2, + uint32_t& r3, + uint32_t& r4, + uint32_t& r5, + uint32_t& r6, + uint32_t& r7) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x8.b32 " + " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3), "=r"(r4), "=r"(r5), "=r"(r6), "=r"(r7) + : "r"(taddr)); +} + +static SGL_DEVICE void tcgen05_ld_32x32b_x8(uint32_t taddr, uint32_t* dst) { + tcgen05_ld_32x32b_x8(taddr, dst[0], dst[1], dst[2], dst[3], dst[4], dst[5], dst[6], dst[7]); +} + +// Blocks until this thread's prior TMEM->reg drains retired. The "memory" +// clobber is load-bearing: ptxas lowers the wait to per-load scoreboard waits +// on the dependent register consumers, so an `mbarrier.arrive` that does not +// read the drained registers gets HOISTED above the last drain (observed in +// SASS) and the next tile's MMA reuses TMEM still being read. +static SGL_DEVICE void tcgen05_wait_ld() { + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); +} + +// Commit prior tcgen05.mma to an mbarrier. Spec accepts `.shared::cluster` or +// no state space, NOT `.shared::cta`. +static SGL_DEVICE void tcgen05_commit_arrive(uint64_t* bar) { + asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.b64 [%0];" ::"r"(to_shared(bar))); +} + +static SGL_DEVICE void tcgen05_fence_after_thread_sync() { + asm volatile("tcgen05.fence::after_thread_sync;"); +} + +// kind::f16 MMA, cta_group::1. Valid dense shapes: M in {64, 128}, N in +// {8, 16, ..., 256} step 8, K = 16. Off-table shapes are NOT rejected by +// ptxas — they hit cudaErrorIllegalInstruction at runtime. +static SGL_DEVICE void +tcgen05_mma_f16(uint32_t d, uint64_t desc_a, uint64_t desc_b, uint32_t inst_desc_high, uint32_t scale_c) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t}\n" ::"r"(d), + "l"(desc_a), + "l"(desc_b), + "r"(inst_desc_high), + "r"(scale_c)); +} + +// ---- TMA (PTX ISA §9.7.9.25) ----------------------------------------------- +// +// COORDINATE CONVENTION: the tensor map's globalDim is (inner, outer) — dim 0 +// is the stride-1 axis. The load calls take (x = inner offset, y = outer +// offset). Mismatch and you load the transposed tile, often with right-looking +// magnitudes but scrambled per-cell pairing. +// +// Completion is mbarrier-based for loads: arm with +// `mbar_arrive_expect_tx(bar, BYTES)` before issuing, then `mbar_wait_parity`. + +// Warm the cache line holding the tensor-map descriptor so the first load of +// the persistent loop does not pay the descriptor fetch. `tmap` is a generic +// address into the __grid_constant__ CUtensorMap param; generic addressing +// resolves it to .param (§9.7.9.15). +static SGL_DEVICE void prefetch_tensormap(const void* tmap) { + asm volatile("prefetch.tensormap [%0];" ::"l"(tmap) : "memory"); +} + +// global -> shared::cta 2D tile load. +static SGL_DEVICE void +cp_async_bulk_tensor_2d_load(uint32_t dst_smem, const CUtensorMap* tmap, int32_t x, int32_t y, uint64_t* bar) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile.mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" ::"r"(dst_smem), + "l"(tmap), + "r"(x), + "r"(y), + "r"(to_shared(bar)) + : "memory"); +} + +// MULTICAST load, cta_group::1: one leader CTA issues, and every CTA whose bit +// is set in `multicast_mask` receives the same bytes at the same CTA-relative +// smem offset AND its OWN tx-count decrement on its own local mbar at +// `bar`'s offset (spec §9.7.9.25, .cta_group::1 bullet). +// +// Contrast with the cta_group::2 form, which CONSOLIDATES all completion onto +// one CTA's mbar (bit 24 of the mbar address cleared). Here each peer keeps its +// own `expect_tx` accounting and only the leader issues, so the follower's smem +// and mbar are both served by the leader's single DRAM read — do NOT clear bit +// 24, that is the cta_group::2 trick and would mis-route the signal. +static SGL_DEVICE void cp_async_bulk_tensor_2d_load_multicast_cg1( + uint32_t dst_smem, + const CUtensorMap* tmap, + int32_t x, + int32_t y, + uint64_t* bar, + uint16_t multicast_mask = 0b11, + uint64_t cache_hint = 0x0ULL) { + const uint32_t mbar_addr = to_shared(bar); + asm volatile( + "cp.async.bulk.tensor.2d.cta_group::1.shared::cluster.global" + ".mbarrier::complete_tx::bytes.multicast::cluster.L2::cache_hint" + " [%0], [%1, {%4, %5}], [%2], %3, %6;" ::"r"(dst_smem), + "l"(tmap), + "r"(mbar_addr), + "h"(multicast_mask), + "r"(x), + "r"(y), + "l"(cache_hint) + : "memory"); +} + +} // namespace ptx + +namespace swz { + +// CPU-side 128B-swizzle math. The TMA hardware permutes the (row x col) atom +// layout when it stores a tile into smem, so reading a known cell back applies +// the same permutation: row stride = 128 B = 8 atoms = 64 BF16 cols, and +// smem_atom = logical_atom XOR (r & 7), an 8-row period (PTX ISA §5.5.7, +// Figures 23-37). Verified on B300 (sm_103a) by a load + read-back roundtrip. +// +// Returns the smem column index in BF16 units within the row. +__host__ __device__ inline uint32_t smem_col_128b_bf16(uint32_t r, uint32_t c) { + return c ^ ((r & 7u) << 3); // (r & 7) atoms shifted; atom = 8 BF16 cols +} + +} // namespace swz + +namespace tmap { + +// Thin wrapper over the tensor-map encoder. Coordinate convention: dim 0 is +// the innermost (stride-1) axis, so globalDim = {cols, rows} and globalStrides +// (length rank-1) carries the row stride in BYTES. +// +// K-major MMA feeds require swizzle == BLOCK_K bytes; the smem descriptor +// assumes equality and a smaller inner box is encoder-legal but silently +// corrupts the MMA load. +// +// The encode is the one driver-API call in this file (there is no runtime-API +// tensor-map encoder), so it goes through cutlass's dlopen-based driver wrapper +// like the other JIT kernels that encode tensor maps — that keeps the module +// off `-lcuda`. +inline CUtensorMap encode_tiled_2d( + void* global_ptr, + CUtensorMapDataType dtype, + uint64_t global_rows, + uint64_t global_cols, + uint64_t row_stride_bytes, + uint32_t box_rows, + uint32_t box_cols, + CUtensorMapSwizzle swizzle = CU_TENSOR_MAP_SWIZZLE_NONE, + CUtensorMapL2promotion promo = CU_TENSOR_MAP_L2_PROMOTION_NONE) { + cuuint64_t global_dim[2] = {global_cols, global_rows}; + cuuint64_t global_strides[1] = {row_stride_bytes}; + cuuint32_t box_dim[2] = {box_cols, box_rows}; + cuuint32_t element_strides[2] = {1, 1}; + + CUtensorMap m{}; + const CUresult res = CUTLASS_CUDA_DRIVER_WRAPPER_CALL(cuTensorMapEncodeTiled)( + &m, + dtype, + /*rank=*/2, + global_ptr, + global_dim, + global_strides, + box_dim, + element_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, + swizzle, + promo, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + if (res != CUDA_SUCCESS) { + std::fprintf( + stderr, + "gemm_ar: cuTensorMapEncodeTiled failed (%d) at %s:%d — rows=%llu cols=%llu " + "row_stride=%llu box=%ux%u swizzle=%d\n", + int(res), + __FILE__, + __LINE__, + (unsigned long long)global_rows, + (unsigned long long)global_cols, + (unsigned long long)row_stride_bytes, + box_rows, + box_cols, + int(swizzle)); + std::abort(); + } + return m; +} + +} // namespace tmap + +namespace dense_gemm_mainloop { + +// GROUP_N L2-stripe raster: linear tile index -> (bid_m, bid_n), walking the +// output grid in N-stripes of width GROUP_N so the same chunk of B stays +// resident in L2 across all M-rows of the stripe (DeepGEMM's +// get_swizzled_block_idx). `cluster_grid_m` is the M-axis tile count in CLUSTER +// units (= grid_m / CTA_GROUP); `grid_n` is NOT halved by CTA_GROUP. The tail +// stripe is shorter than GROUP_N when grid_n % GROUP_N != 0. +template +__device__ __forceinline__ int2 group_n_swizzle(int linear, int crank, int cluster_grid_m, int grid_n) { + if constexpr (CTA_GROUP == 1) { + const int num_blocks_per_group = cluster_grid_m * GROUP_N; + const int group_idx = linear / num_blocks_per_group; + const int first_n = group_idx * GROUP_N; + const int in_group = linear - group_idx * num_blocks_per_group; + const int num_n_in_group = grid_n - first_n < GROUP_N ? grid_n - first_n : GROUP_N; + const int bid_m = in_group / num_n_in_group; + const int bid_n = first_n + (in_group % num_n_in_group); + (void)crank; + return {bid_m, bid_n}; + } else { + const int c = linear / CTA_GROUP; + const int r = linear & (CTA_GROUP - 1); + const int num_blocks_per_group = cluster_grid_m * GROUP_N; + const int group_idx = c / num_blocks_per_group; + const int first_n = group_idx * GROUP_N; + const int in_group = c - group_idx * num_blocks_per_group; + const int num_n_in_group = grid_n - first_n < GROUP_N ? grid_n - first_n : GROUP_N; + const int cluster_bid_m = in_group / num_n_in_group; + const int cluster_bid_n = first_n + (in_group % num_n_in_group); + const int bid_m = cluster_bid_m * CTA_GROUP + r; + const int bid_n = cluster_bid_n; + (void)crank; // `r` already encodes the intra-cluster lane. + return {bid_m, bid_n}; + } +} + +} // namespace dense_gemm_mainloop + +namespace oproj_ar { + +using device::distributed::atomic_add_acq_rel_gpu; +using device::distributed::fence_release_sys; +using device::distributed::load_acquire_sys; +using device::distributed::red_add_relaxed_sys; + +// ---------------------------------------------------------------- constants +#ifndef OPROJ_N // output dim (columns of W / of out). The +#define OPROJ_N 7168 // default is the Kimi-K3 o_proj shape; other +#endif // shapes compile via -DOPROJ_N (see asserts). +constexpr int kN = OPROJ_N; +static_assert( + kN % 256 == 0 && kN >= 256, + "N must be a multiple of 256 (member tile table: 128-row " + "strips + BN up to 256); relaxing this needs a BN-table edit"); +constexpr int kBK = 64; // K per stage (128 B rows → swizzle-128B) +constexpr int kBNRows = 48; // B-box rows per stage (6 n8-tiles) +constexpr int kTilesMax = 6; // n8-tiles per CTA (5-tile CTAs pad, never push) +constexpr int kCWarps = 6; // consumer warp w owns n8-tile w +constexpr int kThreads = (kCWarps + 1) * 32; // +1 dedicated TMA producer warp +constexpr int kMMax = 512; // bs64 — sizes the shared slot layout +constexpr int kRing = 64; // epoch flag/gather ring (monotonic values) + +enum class Comm { kNone, kMc, kPeer, kMcPull, kTwoShot, kTwoShotPeer }; + +// Shared-region layout (BYTES, M-independent: sized at kMMax so every arm and +// every bs cell reuses one region). Parity-double-buffered payloads; the flag +// ring lives on its own 2 MB page. Slot reuse across epochs e / e+2 is safe +// with 2 parities because each rank's launch e+1 spin-waits epoch e+1 AFTER +// its own e-reduce (per-rank stream order). +// Slots are TILE-MAJOR ([n8-tile][m][8 cols]), NOT [m][n]: a warp's push for +// one tile is then a CONTIGUOUS 128 B fabric write instead of 32 scattered +// 4-16 B m-strided writes — lane-scatter runs ~0.26x on this fabric +// and the scattered form's ack-drain dominated the bs8 +// boundary (stamped 22 us at idle). The reduce un-transposes locally. +constexpr size_t kSlotBytes1 = size_t(kMMax) * kN * 2; // one [M,N] bf16 +constexpr __host__ __device__ size_t slot_off(int parity, int src, int R) { + return (size_t(parity) * R + src) * kSlotBytes1; +} +constexpr __host__ __device__ size_t pull_off(int parity, int R) { // [2][M,N] above slots + return (2 * size_t(R) + parity) * kSlotBytes1; +} +constexpr __host__ __device__ size_t flags_off(int R) { + const size_t end = (2 * size_t(R) + 2) * kSlotBytes1; + return (end + (size_t(2) << 20)) & ~((size_t(2) << 20) - 1); +} +// second flag family: "epoch e's slots fully REDUCED", at PER-CTA granularity. +// PDL's wait pairs with the prior grid's TRIGGER (not completion), and with +// 2-CTA/SM residency a fast rank's e+2 pushes can overwrite a slot replica a +// straggling rank's e-reduce still reads — so epoch e's push phase guards on +// done[e-2]. Per-CTA flags (the overwriter of tile t IS every rank's CTA-t) +// keep the publish fully parallel: no second grid-wide gather chain. +constexpr int kMaxCTA = 256; +constexpr int kFams = 7; // flag/gather/done ring FAMILY per dispatch + // CELL {8,16,32,64,128,256,512}. Monotonic ring + // targets assume every epoch of a ring came + // from the same gridDim (and only two-shot + // bumps the boundary-2 ring words), so one fam + // per cell removes the host-side ring reset a + // cell change needed under the original 3-fam + // grid-class split — the reset was a collective + // and blocked CUDA-graph capture. +constexpr __host__ __device__ size_t done_off(int R) { + return flags_off(R) + size_t(kFams) * 512; +} +constexpr __host__ __device__ size_t region_bytes(int R) { + return done_off(R) + size_t(kFams) * kRing * kMaxCTA * 4 + 512; +} + +// PDL (programmatic dependent launch): the NEXT launch on the stream may +// start its feed while THIS grid sits in the boundary spin + reduce; its +// epilogue then griddepcontrol.wait's until this grid fully completes, so +// slots/flags/out stay race-free. Serving overlaps the next layer's kernels +// the same way; an unfused cublas composite cannot cooperate across the +// vendor-kernel boundary. +SGL_DEVICE uint32_t bf2_u32(float2 f) { + const __nv_bfloat162 p = __float22bfloat162_rn(f); + return *reinterpret_cast(&p); +} + +// ------------------------------------------------------------------ params + +template +struct Params { + uint8_t* mc_base; // MC VA (null when no MC object — kPeer runs) + uint8_t* uc_base[R]; // per-rank unicast VAs of the shared region + uint32_t* gather; // device-local u32[kRing] + __nv_bfloat16* out; // [M,N] local output + const __nv_bfloat16* partial_in; // GEMM_ON=false input [M,N] + // Device-resident per-fam CTA ticket counters. Every CTA takes one ticket + // at entry and divides by the (family-stable) gridDim to recover the + // launch epoch. All CTAs have taken their tickets before this grid + // triggers a PDL successor, so successive launches receive disjoint, + // contiguous ticket ranges without a separate bump kernel. Device state + // instead of a launch arg keeps CUDA-graph replays advancing the epoch. + uint32_t* epoch_base; + int my_rank; + int fam; // ring family (dispatch cell) — see kFams +}; + +// -------------------------------------------------------------- CTA strips +// kN/8 = 896 n8-tiles over gridDim CTAs: first `rem` CTAs own base+1 tiles. +struct Strip { + int t0, nt; + static SGL_DEVICE Strip make(int cta, int ncta) { + const int kT = kN / 8; + const int base = kT / ncta, rem = kT % ncta; + Strip s; + if (cta < rem) { + s.nt = base + 1; + s.t0 = cta * (base + 1); + } else { + s.nt = base; + s.t0 = rem * (base + 1) + (cta - rem) * base; + } + return s; + } +}; + +// ------------------------------------------------------------------ kernel +// mbar contract (recipes/mbar_handshake_design): full[s] count=1, arrive = +// producer's arrive_expect_tx(B+A bytes) + TMA complete_tx; empty[s] count = +// kCWarps (one elected lane per consumer warp after its last stage read). +// Producer = warp kCWarps lane 0; it never consumes, so the ring never +// self-blocks. Ring reuse distance = S stages, issue j waits empty parity +// ((j-S)/S)&1 — both derived from S. +// C = TMA cluster size for the A feed: ONE leader multicast per stage fills +// all C CTAs' A slots (cp_async_bulk_tensor_2d_load_multicast_cg1: each CTA's +// local full-bar gets its own tx-decrement, so expect_tx is unchanged). Per-CTA +// A bytes drop C-fold — the M-scaling A-tax (LEDGER O1) is A riding every +// CTA's capped TMA pipe. Contract deltas at C>1: the LEADER's empty[s] count +// = kCWarps*C (followers' consumers cluster-arrive it — release variant, so +// their slot reads are performed-before the leader's next multicast write); +// followers skip their own A issue but keep full expect_tx. +template +__global__ void __launch_bounds__(kThreads) oproj_ar_kernel( + const __grid_constant__ CUtensorMap w_map, + const __grid_constant__ CUtensorMap x_map, + const __grid_constant__ Params prm) { + constexpr int Mp = (M + 15) & ~15; // mma m16 padding (x buffer padded) + constexpr int MT = Mp / 16; + constexpr int KSTEPS = K / (kBK * CH); // OUTER stages: CH k-chunks batched + constexpr int kBBytes = kBNRows * kBK * 2, kABytes = Mp * kBK * 2; + constexpr int kStB = CH * kBBytes, kStA = CH * kABytes; + static_assert(K % (kBK * CH) == 0); + + const int tid = threadIdx.x, warp = tid >> 5, lane = tid & 31; + const Strip strip = Strip::make(blockIdx.x, gridDim.x); + __shared__ uint32_t cta_epoch; + if (tid == 0) cta_epoch = atomicAdd(prm.epoch_base + prm.fam, 1u) / uint32_t(gridDim.x); + __syncthreads(); + const uint32_t epoch = cta_epoch; + const int parity = int(epoch & 1); + const int ring = int(epoch % kRing); + // pinned wait-set: flag VA + monotonic targets resolved before any spin + const size_t foff = flags_off(R) + size_t(prm.fam) * 512; + const size_t doff2 = done_off(R) + size_t(prm.fam) * kRing * kMaxCTA * 4; + uint32_t* const flag_local = reinterpret_cast(prm.uc_base[prm.my_rank] + foff) + ring; + uint32_t* const done_local = reinterpret_cast(prm.uc_base[prm.my_rank] + doff2) + size_t(blockIdx.x); + uint32_t* const gather_fam = prm.gather + size_t(prm.fam) * 2 * kRing; + const uint32_t wrap = epoch / kRing + 1; + const uint32_t flag_target = wrap * R; + const uint32_t gath_target = wrap * gridDim.x; + + float4 acc[MT]; +#pragma unroll + for (int i = 0; i < MT; ++i) + acc[i] = make_float4(0.f, 0.f, 0.f, 0.f); + + extern __shared__ __align__(1024) uint8_t smem[]; + uint8_t* b_st = smem; // [S][CH][kBBytes] + uint8_t* a_st = b_st + size_t(S) * kStB; // [S][CH][kABytes] + uint64_t* fullb = reinterpret_cast(a_st + size_t(S) * kStA); + uint64_t* emptyb = fullb + S; + + const uint32_t crank = C > 1 ? ptx::cluster_cta_rank() : 0; + { + if (tid == 0) { + ptx::prefetch_tensormap(&w_map); + ptx::prefetch_tensormap(&x_map); +#pragma unroll + for (int s = 0; s < S; ++s) { + ptx::mbar_init(fullb + s, 1); + ptx::mbar_init(emptyb + s, crank == 0 ? kCWarps * C : kCWarps); + } + } + __syncthreads(); + if constexpr (C > 1) ptx::cluster_sync_rel_acq(); + + if (warp == kCWarps) { + // ---- producer: the whole K stream, one thread ----------------- + // k-phase rotation: spreads B's DRAM pages across CTAs; fp32 + // accumulation order changes per CTA — a sum, gate-covered. + // CLUSTER-uniform: all members consume the same A flight. + const int phase = ((int(blockIdx.x) / C) * KSTEPS) / (int(gridDim.x) / C); + if (lane == 0) { + for (int j = 0; j < KSTEPS; ++j) { + const int slot = j % S; + const int jj = (j + phase) % KSTEPS; + if (j >= S) ptx::mbar_wait_parity(emptyb + slot, ((j - S) / S) & 1); + ptx::mbar_arrive_expect_tx(fullb + slot, kStB + kStA); +#pragma unroll + for (int c = 0; c < CH; ++c) { + ptx::cp_async_bulk_tensor_2d_load( + ptx::to_shared(b_st + size_t(slot) * kStB + c * kBBytes), + &w_map, + (jj * CH + c) * kBK, + strip.t0 * 8, + fullb + slot); + if constexpr (C > 1) { + if (crank == 0) + ptx::cp_async_bulk_tensor_2d_load_multicast_cg1( + ptx::to_shared(a_st + size_t(slot) * kStA + c * kABytes), + &x_map, + (jj * CH + c) * kBK, + 0, + fullb + slot, + uint16_t((1u << C) - 1)); + } else { + ptx::cp_async_bulk_tensor_2d_load( + ptx::to_shared(a_st + size_t(slot) * kStA + c * kABytes), + &x_map, + (jj * CH + c) * kBK, + 0, + fullb + slot); + } + } + } + } + } else { + // ---- consumers: warp w owns n8-tile (strip.t0 + w) ------------ + // A fragments load straight from gmem (x is L2-hot and tiny; the + // LSU pipe is idle here) — the TMA path stays a pure-B stream. + const int b_row = (lane & 7) + warp * 8; // row within the B box + const int b_ka = (lane >> 3) & 1; // k-atom half (x2) + const int a_row = lane & 15, a_ka = lane >> 4; + for (int s = 0; s < KSTEPS; ++s) { + const int slot = s % S; + ptx::mbar_wait_parity(fullb + slot, (s / S) & 1); +#pragma unroll + for (int c = 0; c < CH; ++c) { + const uint32_t b_base = ptx::to_shared(b_st + size_t(slot) * kStB + c * kBBytes); +#pragma unroll + for (int k16 = 0; k16 < kBK / 16; ++k16) { + uint32_t b0, b1; + ptx::ldmatrix_x2_b16( + b_base + uint32_t(b_row) * (kBK * 2) + swz::smem_col_128b_bf16(b_row, (k16 * 2 + b_ka) * 8) * 2, + b0, + b1); +#pragma unroll + for (int mt = 0; mt < MT; ++mt) { + uint32_t a0, a1, a2, a3; + ptx::ldmatrix_x4_b16( + ptx::to_shared( + a_st + size_t(slot) * kStA + c * kABytes + uint32_t(mt * 16 + a_row) * (kBK * 2) + + swz::smem_col_128b_bf16(a_row, (k16 * 2 + a_ka) * 8) * 2), + a0, + a1, + a2, + a3); + ptx::mma_m16n8k16_bf16f32(acc[mt], a0, a1, a2, a3, b0, b1); + } + } + } + if (lane == 0) { + ptx::mbar_arrive(emptyb + slot); + if constexpr (C > 1) + if (crank != 0) ptx::mbar_arrive_cluster_release(emptyb + slot, 0); + } + } + } + } + + // ---- epilogue: push ---------------------------------------------------- + __syncthreads(); // whole CTA past its smem/feed reads + device::PDLTriggerSecondary(); // next launch streams weights under our + // push+boundary+reduce (needs 2-CTA/SM + // co-residency: 100% smem carveout + this + // kernel's smem ≤ ~113 KB) + device::PDLWaitPrimary(); // prior grid reached ITS trigger (k-loop end) — NOT done + { + // guard: epoch e-2 (same parity) fully reduced everywhere before we + // overwrite its slots. Steady-state this is already set (~one hot + // acquire); it binds only when a boundary straggles. + if (tid == 0 && epoch >= 2) { + const uint32_t e2 = epoch - 2; + const uint32_t tgt = (e2 / kRing + 1) * R; + while (load_acquire_sys(done_local + size_t(e2 % kRing) * kMaxCTA) < tgt) { + } + } + __syncthreads(); + } + const int tig = lane & 3, grp = lane >> 2; + const int n0 = (strip.t0 + warp) * 8; + const bool own = warp < strip.nt; // phantom 6th tile / producer never push + + // tile-major slot offset: tile t (global n8 index), row m, byte offset + auto slot_tm = [&](uint8_t* base, size_t sb, int t, int m) { return base + sb + (size_t(t) * M + m) * 16; }; + + { + const size_t sb = slot_off(parity, prm.my_rank, R); + if (own) { + const int t = strip.t0 + warp; +#pragma unroll + for (int mt = 0; mt < MT; ++mt) { +#pragma unroll + for (int half = 0; half < 2; ++half) { + const int m = mt * 16 + grp + half * 8; + if (m < M) { + const uint32_t v = + half ? bf2_u32(make_float2(acc[mt].z, acc[mt].w)) : bf2_u32(make_float2(acc[mt].x, acc[mt].y)); + // lanes (grp,tig) of one warp land contiguous: 128 B + // per (mt,half) group + const size_t boff = size_t(m) * 16 + tig * 4; + { +#pragma unroll + for (int r = 0; r < R; ++r) + *reinterpret_cast(slot_tm(prm.uc_base[r], sb, t, 0) + boff) = v; + } + } + } + } + } + } + + // ---- boundary (gather + flag + per-CTA local-replica spin) ------------- + __syncthreads(); + if (tid == 0) { + const uint32_t old = atomic_add_acq_rel_gpu(gather_fam + ring, 1); + if (old + 1 == gath_target) { + // ONE completing-CTA fence: publishes every CTA's pushes via the + // acq_rel gather chain. Measured 0.9 us better than per-CTA + // fences at bs1 (11.04 vs 11.97) — the parallel-drain theory lost. + fence_release_sys(); + { +#pragma unroll + for (int r = 0; r < R; ++r) + red_add_relaxed_sys(reinterpret_cast(prm.uc_base[r] + foff) + ring, 1); + } + } + while (load_acquire_sys(flag_local) < flag_target) { + } + } + __syncthreads(); + + // ---- reduce: each CTA finishes its own tiles from the LOCAL replica ---- + const int units = strip.nt * M; + for (int u = tid; u < units; u += kThreads) { + const int t = u / M, m = u % M, c0 = (strip.t0 + t) * 8; + const size_t soff = (size_t(strip.t0 + t) * M + m) * 16; // tile-major + uint4 res; + { + float2 s[4] = {{0.f, 0.f}, {0.f, 0.f}, {0.f, 0.f}, {0.f, 0.f}}; +#pragma unroll + for (int r = 0; r < R; ++r) { + const uint4 v = *reinterpret_cast(prm.uc_base[prm.my_rank] + slot_off(parity, r, R) + soff); + const uint32_t w4[4] = {v.x, v.y, v.z, v.w}; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const float2 f = __bfloat1622float2(*reinterpret_cast(&w4[j])); + s[j].x += f.x; + s[j].y += f.y; + } + } + res = make_uint4(bf2_u32(s[0]), bf2_u32(s[1]), bf2_u32(s[2]), bf2_u32(s[3])); + } + *reinterpret_cast(reinterpret_cast(prm.out) + (size_t(m) * kN + c0) * 2) = res; + } + + // ---- done-publish: per-CTA, fence-free. The beacon carries no payload: + // the remote guard only needs its VALUE. This CTA's reduce loads are + // data-flow-complete (their values fed the out stores) before the + // syncthreads, so the relaxed beacon cannot pass an unfinished read; a + // fence here stamped at ~4 us draining against the co-resident feed. + __syncthreads(); + if (tid == 0) { + const size_t off = size_t(ring) * kMaxCTA + blockIdx.x; + { +#pragma unroll + for (int r = 0; r < R; ++r) + red_add_relaxed_sys(reinterpret_cast(prm.uc_base[r] + doff2) + off, 1); + } + } +} + +constexpr int kD3BM = 128, kD3BK = 64; +constexpr __host__ __device__ int d3_ns(int bn) { // ring depth by stage size + // deepest ring the smem budget fits, capped at the mbar array size. The + // old 12/11/4 step form silently ran M>=128 at NS=4 (BN64 stage 24.5 KB + // fits 9) — a 4-deep ring can't hide DRAM latency and was the dominant + // M=128 pole (21.4 us vs DG 17.2 at-shape). + const int stage = kD3BM * kD3BK * 2 + bn * kD3BK * 2; + const int ns = 220 * 1024 / stage; + return ns > 12 ? 12 : ns; +} +constexpr int kD3GroupN = 16, kD3KPer = kD3BK / 16; +constexpr int kD3Threads = 8 * 32; // warps 0,1,4-7 live +constexpr int kD3ABytes = kD3BM * kD3BK * 2; +// per-M BLOCK_N: smaller BN = more tiles = more feed pipes (grid-starvation +// at BN=256 measured M-flat ~30.7 us: 28 CTAs); bigger BN = less A-dup. +// bytes/(min(tiles,152)) + tensor-floor optimum: 64 for M<=128, 128 above. +// re-swept 2026-07-23 at floored d3_ns: M=128 BN64/NS9 16.41 vs BN128/NS6 +// 20.50 (ring depth, not BN, was the old 21.4 pole); 256 BN128/NS6 20.51; +// 512 BN256/NS4 30.9 (BN128's 224-on-152 nu-imbalance regressed to 38.9). +constexpr __host__ __device__ int d3_bn(int Mp) { + return Mp <= 128 ? 64 : (Mp == 256 ? 128 : 256); +} + +// drain-order slot offset: unit = (tile t, n-block nb, epi-warp w, lane l), +// 16 B each. Total = num_tiles * (BN/8) * 4 * 32 * 16 = Mpad*N*2 bytes. +SGL_DEVICE size_t d3_slot_off(int t, int nb, int w, int l, int nblk_per_tile) { + return ((size_t(t) * nblk_per_tile + nb) * 4u + w) * 32u * 16u + size_t(l) * 16u; +} + +// SWAP (M<=64): swapAB on the dense ring — A = W (M-slot carries the 7168 +// n-rows in 128-row strips, NO small-M padding tax), B = x^T (N-slot = Mp, +// as small as 8). Tiles = 56 n-strips; the drain's lane-rows become n and +// its cols become m — the same drain-order slots work, only the out mapping +// transposes (local scatter, cheap at these payloads). +template +__global__ void __launch_bounds__(kD3Threads) oproj_dense_ar_kernel( + const __grid_constant__ CUtensorMap x_tmap, // A = x [kMMax, K] + const __grid_constant__ CUtensorMap w_tmap, // B = W [kN, K] + const __grid_constant__ Params prm) { + constexpr bool SWAP = (M <= 64); + // two-shot planes share RS/reduce/out-copy; they differ only in the AG + + // flag transport (kTwoShot = NVLS multimem, kTwoShotPeer = pure P2P). + constexpr bool k2S = (COMM == Comm::kTwoShot || COMM == Comm::kTwoShotPeer); + // SWAP = the DG decode recipe (O9h): token-slot tiles of 16/32 m-cols — + // tiles = 56 strips × (Mp/tok) fills the grid (112 pipes at M≥32; the + // 56-tile form measured flat ~19.6-20.5, R11) — with a 12-deep ring. + // UMMA M=128 requires N ≥ 16 step 16 (bs1 pads the token slot to 16). + constexpr int Mp = SWAP ? (M < 16 ? 16 : ((M + 15) & ~15)) : (M + kD3BM - 1) / kD3BM * kD3BM; + constexpr int kD3BN = SWAP ? (Mp <= 32 ? 16 : 32) : d3_bn(Mp); + constexpr int kD3BBytes = kD3BN * kD3BK * 2; + constexpr int kGridM = SWAP ? kN / kD3BM : Mp / kD3BM; + constexpr int kGridN = SWAP ? Mp / kD3BN : kN / kD3BN; + constexpr int kTiles = kGridM * kGridN; + constexpr int kIters = K / kD3BK; + constexpr int kNBlk = kD3BN / 8; + + constexpr int kD3NS = d3_ns(kD3BN); + // tcgen05.alloc column count must be a power of two >= 32 + constexpr int kTmemCols = 2 * kD3BN <= 32 ? 32 + : 2 * kD3BN <= 64 ? 64 + : 2 * kD3BN <= 128 ? 128 + : 2 * kD3BN <= 256 ? 256 + : 512; + extern __shared__ __align__(1024) uint8_t smem_buf[]; + const uint32_t smem_base = ptx::to_shared(smem_buf); + constexpr uint32_t kSmemAOff = 0, kSmemBOff = kD3NS * kD3ABytes; + + __shared__ __align__(8) uint64_t tma_mbars[12]; + __shared__ __align__(8) uint64_t mma_mbars[12]; + __shared__ __align__(8) uint64_t mainloop_mbars[2]; + __shared__ __align__(8) uint64_t epi_mbars[2]; + __shared__ __align__(4) uint32_t s_taddr[1]; + __shared__ uint32_t cta_epoch; + + const int tid = threadIdx.x, warp_id = tid >> 5, lane_id = tid & 31; + const int bid = int(blockIdx.x), num_bids = int(gridDim.x); + if (tid == 0) cta_epoch = atomicAdd(prm.epoch_base + prm.fam, 1u) / uint32_t(gridDim.x); + __syncthreads(); + const uint32_t epoch = cta_epoch; + const int parity = int(epoch & 1); + const int ring = int(epoch % kRing); + const size_t foff = flags_off(R) + size_t(prm.fam) * 512; + const size_t doff2 = done_off(R) + size_t(prm.fam) * kRing * kMaxCTA * 4; + uint32_t* const flag_local = reinterpret_cast(prm.uc_base[prm.my_rank] + foff) + ring; + uint32_t* const done_local = reinterpret_cast(prm.uc_base[prm.my_rank] + doff2) + size_t(blockIdx.x); + uint32_t* const gather_fam = prm.gather + size_t(prm.fam) * 2 * kRing; + const uint32_t wrap = epoch / kRing + 1; + const uint32_t flag_target = wrap * R; + const uint32_t gath_target = wrap * gridDim.x; + const size_t sb = slot_off(parity, prm.my_rank, R); + + if (warp_id == 0 && ptx::elect_one()) { + for (int i = 0; i < kD3NS; ++i) { + ptx::mbar_init(&tma_mbars[i], 1); + ptx::mbar_init(&mma_mbars[i], 1); + } + for (int i = 0; i < 2; ++i) { + ptx::mbar_init(&mainloop_mbars[i], 1); + ptx::mbar_init(&epi_mbars[i], 4 * 32); + } + } else if (warp_id == 1) { + ptx::tcgen05_alloc(ptx::to_shared(s_taddr), kTmemCols); + } + __syncthreads(); + const uint32_t taddr = s_taddr[0]; + + constexpr uint32_t i_desc = ptx::mma_inst_desc_f16( + kD3BM, kD3BN, ptx::F16Type::BF16, ptx::F16Type::BF16, ptx::DType::F32, ptx::Major::K, ptx::Major::K); + auto tile_mn = [&](int linear) -> int2 { + return dense_gemm_mainloop::group_n_swizzle<1, kD3GroupN>(linear, 0, kGridM, kGridN); + }; + + // Prefetch the first ring of input-independent weight stages before the + // PDL dependency. SWAP changes which operand slot contains W, but never + // changes which tensor map is safe to touch here. + if (warp_id == 0 && ptx::elect_one()) { + constexpr int kPrefetch = kIters < kD3NS ? kIters : kD3NS; + const int2 mn = tile_mn(bid); +#pragma unroll + for (int k = 0; k < kPrefetch; ++k) { + ptx::mbar_arrive_expect_tx(&tma_mbars[k], kD3ABytes + kD3BBytes); + ptx::cp_async_bulk_tensor_2d_load( + smem_base + (SWAP ? kSmemAOff + k * kD3ABytes : kSmemBOff + k * kD3BBytes), + &w_tmap, + k * kD3BK, + (SWAP ? mn.x * kD3BM : mn.y * kD3BN), + &tma_mbars[k]); + } + } + + // x and all slot traffic remain behind the dependency. + device::PDLWaitPrimary(); + + // done-guard before any slot write (W3; PDL wait pairs with the trigger) + { + if (tid == 0 && epoch >= 2) { + const uint32_t e2 = epoch - 2; + const uint32_t tgt = (e2 / kRing + 1) * R; + while (load_acquire_sys(done_local + size_t(e2 % kRing) * kMaxCTA) < tgt) { + } + } + __syncthreads(); + } + + if (warp_id == 0 && ptx::elect_one()) { + // TMA issuer (simple persistent) — verbatim dense_1cta fp8out shape + int stage = 0, mma_phase = 1; + for (int t = bid; t < kTiles; t += num_bids) { + const int2 mn = tile_mn(t); + for (int k = 0; k < kIters; ++k) { + ptx::mbar_wait_parity(&mma_mbars[stage], mma_phase); + constexpr bool kDropA = false; + const bool prefetched = (t == bid && k < kD3NS); + if (!prefetched) ptx::mbar_arrive_expect_tx(&tma_mbars[stage], (kDropA ? 0 : kD3ABytes) + kD3BBytes); + if constexpr (!kDropA) { + // In SWAP, A is the prefetched weight; otherwise A is x. + if (!prefetched || !SWAP) + ptx::cp_async_bulk_tensor_2d_load( + smem_base + kSmemAOff + stage * kD3ABytes, + SWAP ? &w_tmap : &x_tmap, + k * kD3BK, + mn.x * kD3BM, + &tma_mbars[stage]); + } + // In SWAP, B is x; otherwise B is the prefetched weight. + if (!prefetched || SWAP) + ptx::cp_async_bulk_tensor_2d_load( + smem_base + kSmemBOff + stage * kD3BBytes, + SWAP ? &x_tmap : &w_tmap, + k * kD3BK, + mn.y * kD3BN, + &tma_mbars[stage]); + if (++stage == kD3NS) { + stage = 0; + mma_phase ^= 1; + } + } + } + } else if (warp_id == 1 && ptx::elect_one()) { + // MMA issuer with 2-stage TMEM ping-pong + int stage = 0, tma_phase = 0, ml_stage = 0, epi_phase = 1; + for (int t = bid; t < kTiles; t += num_bids) { + ptx::mbar_wait_parity(&epi_mbars[ml_stage], epi_phase); + const uint32_t tmem_d = taddr + uint32_t(ml_stage) * kD3BN; + for (int k = 0; k < kIters; ++k) { + ptx::mbar_wait_parity(&tma_mbars[stage], tma_phase); + ptx::tcgen05_fence_after_thread_sync(); + const uint32_t a_smem = smem_base + kSmemAOff + stage * kD3ABytes; + const uint32_t b_smem = smem_base + kSmemBOff + stage * kD3BBytes; +#pragma unroll + for (int k2 = 0; k2 < kD3KPer; ++k2) { + const uint64_t da = ptx::mma_smem_desc_k_major(a_smem + uint32_t(k2) * 32); + const uint64_t db = ptx::mma_smem_desc_k_major(b_smem + uint32_t(k2) * 32); + ptx::tcgen05_mma_f16(tmem_d, da, db, i_desc, (k == 0 && k2 == 0) ? 0u : 1u); + } + ptx::tcgen05_commit_arrive(&mma_mbars[stage]); + if (++stage == kD3NS) { + stage = 0; + tma_phase ^= 1; + } + } + ptx::tcgen05_commit_arrive(&mainloop_mbars[ml_stage]); + ml_stage ^= 1; + if (ml_stage == 0) epi_phase ^= 1; + } + } else if (warp_id >= 4) { + // epilogue: BF16 drain (dense_1cta idiom) → coalesced comm push + const int epi_warp = warp_id & 3; + const uint32_t taddr_lane = uint32_t(epi_warp * 32) << 16; + int ml_stage = 0, ml_phase = 0; + for (int t = bid; t < kTiles; t += num_bids) { + const int2 mn = tile_mn(t); + ptx::mbar_wait_parity(&mainloop_mbars[ml_stage], ml_phase); + ptx::tcgen05_fence_after_thread_sync(); + const uint32_t tmem_d_base = taddr + uint32_t(ml_stage) * kD3BN; + const int row = mn.x * kD3BM + epi_warp * 32 + lane_id; // SWAP: n +#pragma unroll 4 + for (int nb = 0; nb < kNBlk; ++nb) { + const uint32_t taddr_n = tmem_d_base + uint32_t(nb) * 8 + taddr_lane; + uint32_t r0, r1, r2, r3, r4, r5, r6, r7; + ptx::tcgen05_ld_32x32b_x8(taddr_n, r0, r1, r2, r3, r4, r5, r6, r7); + ptx::tcgen05_wait_ld(); + uint4 v; + v.x = ptx::cvt_pack_f32x2_to(__int_as_float(r1), __int_as_float(r0)); + v.y = ptx::cvt_pack_f32x2_to(__int_as_float(r3), __int_as_float(r2)); + v.z = ptx::cvt_pack_f32x2_to(__int_as_float(r5), __int_as_float(r4)); + v.w = ptx::cvt_pack_f32x2_to(__int_as_float(r7), __int_as_float(r6)); + if (SWAP || row < M) { // SWAP masks pad m-cols below + { + const size_t off = sb + d3_slot_off(t, nb, epi_warp, lane_id, kNBlk); + if constexpr (k2S) + // RS: unicast to tile-owner only — 1x egress and, + // spread over the tile loop, absorbed under the + // mainloop (O9b: one-shot is R x-payload INGRESS- + // bound; the composite pays its RS serially) + *reinterpret_cast(prm.uc_base[t % R] + off) = v; + else { +#pragma unroll + for (int r = 0; r < R; ++r) + *reinterpret_cast(prm.uc_base[r] + off) = v; + } + } + } + } + (void)ptx::mbar_arrive(&epi_mbars[ml_stage]); + ml_stage ^= 1; + if (ml_stage == 0) ml_phase ^= 1; + } + } + __syncthreads(); + device::PDLTriggerSecondary(); + if (warp_id == 1) { + ptx::tcgen05_dealloc(taddr, kTmemCols); + ptx::tcgen05_relinquish(); + } + + // ---- boundary (fam rings) — verbatim member-1 contract ---------------- + if (tid == 0) { + const uint32_t old = atomic_add_acq_rel_gpu(gather_fam + ring, 1); + if (old + 1 == gath_target) { + fence_release_sys(); + { +#pragma unroll + for (int r = 0; r < R; ++r) + red_add_relaxed_sys(reinterpret_cast(prm.uc_base[r] + foff) + ring, 1); + } + } + while (load_acquire_sys(flag_local) < flag_target) { + } + } + __syncthreads(); + + if constexpr (k2S) { + // ---- owner-reduce + AG: reduce MY tiles, store the result to all + // replicas' pull region (kTwoShot: one mm.st, fabric replicates; + // kTwoShotPeer: R unicast stores — (R-1)× the egress, same ingress); + // then boundary 2 gates the out-copy --------------------------------- + for (int u = tid + bid * kD3Threads; u < kTiles * kNBlk * 4 * 32; u += num_bids * kD3Threads) { + const int l = u & 31, w = (u >> 5) & 3, nb = (u >> 7) % kNBlk; + const int t = u / (kNBlk * 128); + if (t % R != prm.my_rank) continue; // not my slab + const int2 mn = tile_mn(t); + if (!SWAP && mn.x * kD3BM + w * 32 + l >= M) continue; + const size_t soff = d3_slot_off(t, nb, w, l, kNBlk); + float2 acc2[4] = {{0.f, 0.f}, {0.f, 0.f}, {0.f, 0.f}, {0.f, 0.f}}; +#pragma unroll + for (int r = 0; r < R; ++r) { + const uint4 vv = *reinterpret_cast(prm.uc_base[prm.my_rank] + slot_off(parity, r, R) + soff); + const uint32_t w4[4] = {vv.x, vv.y, vv.z, vv.w}; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const float2 f = __bfloat1622float2(*reinterpret_cast(&w4[j])); + acc2[j].x += f.x; + acc2[j].y += f.y; + } + } + const uint4 res = make_uint4(bf2_u32(acc2[0]), bf2_u32(acc2[1]), bf2_u32(acc2[2]), bf2_u32(acc2[3])); + { +#pragma unroll + for (int r = 0; r < R; ++r) + *reinterpret_cast(prm.uc_base[r] + pull_off(parity, R) + soff) = res; + } + } + // boundary 2 (second flag/gather ring words at +256 B / +kRing) + __syncthreads(); + if (tid == 0) { + uint32_t* const g2 = gather_fam + kRing; // AG gather ring + const uint32_t old = atomic_add_acq_rel_gpu(g2 + ring, 1); + if (old + 1 == gath_target) { + fence_release_sys(); + { +#pragma unroll + for (int r = 0; r < R; ++r) + red_add_relaxed_sys(reinterpret_cast(prm.uc_base[r] + foff + 256) + ring, 1); + } + } + while (load_acquire_sys(reinterpret_cast(prm.uc_base[prm.my_rank] + foff + 256) + ring) < + flag_target) { + } + } + __syncthreads(); + // out-copy: every CTA writes its grid-partition of out from the + // LOCAL reduced replica + for (int u = tid + bid * kD3Threads; u < kTiles * kNBlk * 4 * 32; u += num_bids * kD3Threads) { + const int l = u & 31, w = (u >> 5) & 3, nb = (u >> 7) % kNBlk; + const int t = u / (kNBlk * 128); + const int2 mn = tile_mn(t); + const int row = mn.x * kD3BM + w * 32 + l; // SWAP: row = n, full + if (!SWAP && row >= M) continue; // range; m masked below + const uint4 res = *reinterpret_cast( + prm.uc_base[prm.my_rank] + pull_off(parity, R) + d3_slot_off(t, nb, w, l, kNBlk)); + if constexpr (SWAP) { + const uint32_t u4[4] = {res.x, res.y, res.z, res.w}; +#pragma unroll + for (int j = 0; j < 8; ++j) { + const int mm = mn.y * kD3BN + nb * 8 + j; + if (mm < M) + *reinterpret_cast(reinterpret_cast(prm.out) + (size_t(mm) * kN + row) * 2) = + uint16_t((u4[j >> 1] >> ((j & 1) * 16)) & 0xFFFFu); + } + } else { + *reinterpret_cast( + reinterpret_cast(prm.out) + (size_t(row) * kN + size_t(mn.y) * kD3BN + nb * 8) * 2) = res; + } + } + __syncthreads(); + if (tid == 0) { + const size_t off = size_t(ring) * kMaxCTA + blockIdx.x; + { +#pragma unroll + for (int r = 0; r < R; ++r) + red_add_relaxed_sys(reinterpret_cast(prm.uc_base[r] + doff2) + off, 1); + } + } + return; + } + + // ---- reduce: descatter drain-order slots → out[m][n] ------------------- + // unit = (t, nb, w, l): row = mn.x*BM + w*32 + l; cols = mn.y*BN + nb*8. + const int units = kTiles * kNBlk * 4 * 32; + for (int u = tid + bid * kD3Threads; u < units; u += num_bids * kD3Threads) { + const int l = u & 31, w = (u >> 5) & 3, nb = (u >> 7) % kNBlk; + const int t = u / (kNBlk * 128); + const int2 mn = tile_mn(t); + const int row = mn.x * kD3BM + w * 32 + l; + if (!SWAP && row >= M) continue; + const size_t soff = d3_slot_off(t, nb, w, l, kNBlk); + uint4 res; + { + float2 acc2[4] = {{0.f, 0.f}, {0.f, 0.f}, {0.f, 0.f}, {0.f, 0.f}}; +#pragma unroll + for (int r = 0; r < R; ++r) { + const uint4 vv = *reinterpret_cast(prm.uc_base[prm.my_rank] + slot_off(parity, r, R) + soff); + const uint32_t w4[4] = {vv.x, vv.y, vv.z, vv.w}; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const float2 f = __bfloat1622float2(*reinterpret_cast(&w4[j])); + acc2[j].x += f.x; + acc2[j].y += f.y; + } + } + res = make_uint4(bf2_u32(acc2[0]), bf2_u32(acc2[1]), bf2_u32(acc2[2]), bf2_u32(acc2[3])); + } + if constexpr (SWAP) { + const uint32_t u4[4] = {res.x, res.y, res.z, res.w}; +#pragma unroll + for (int j = 0; j < 8; ++j) { + const int mm = mn.y * kD3BN + nb * 8 + j; + if (mm < M) + *reinterpret_cast(reinterpret_cast(prm.out) + (size_t(mm) * kN + row) * 2) = + uint16_t((u4[j >> 1] >> ((j & 1) * 16)) & 0xFFFFu); + } + } else { + *reinterpret_cast( + reinterpret_cast(prm.out) + (size_t(row) * kN + size_t(mn.y) * kD3BN + nb * 8) * 2) = res; + } + } + + // ---- done publish (fence-free beacon, W4) ------------------------------ + __syncthreads(); + if (tid == 0) { + const size_t off = size_t(ring) * kMaxCTA + blockIdx.x; + { +#pragma unroll + for (int r = 0; r < R; ++r) + red_add_relaxed_sys(reinterpret_cast(prm.uc_base[r] + doff2) + off, 1); + } + } +} + +template +struct Launcher3 { + static constexpr bool kSwap = (M <= 64); + static constexpr int Mp = kSwap ? (M < 16 ? 16 : ((M + 15) & ~15)) : (M + kD3BM - 1) / kD3BM * kD3BM; + static constexpr int kBN = kSwap ? (Mp <= 32 ? 16 : 32) : d3_bn(Mp); + static constexpr int kTiles = kSwap ? (kN / kD3BM) * (Mp / kBN) : (Mp / kD3BM) * (kN / kBN); + static constexpr int kGrid = kTiles < 152 ? kTiles : 152; + static constexpr size_t kSmem = size_t(d3_ns(kBN)) * (kD3ABytes + kBN * kD3BK * 2); + static void set_smem_attr() { + CHECK_CUDA(cudaFuncSetAttribute( + oproj_dense_ar_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, int(kSmem))); + CHECK_CUDA(cudaFuncSetAttribute( + oproj_dense_ar_kernel, cudaFuncAttributePreferredSharedMemoryCarveout, 100)); + } + static void + launch(const CUtensorMap& x_tmap, const CUtensorMap& w_tmap, const Params& prm, cudaStream_t stream, bool pdl) { + cudaLaunchConfig_t cfg{}; + cudaLaunchAttribute attr[1]; + int na = 0; + if (pdl) { + attr[na].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[na].val.programmaticStreamSerializationAllowed = 1; + ++na; + } + cfg.gridDim = dim3(unsigned(kGrid)); + cfg.blockDim = dim3(kD3Threads); + cfg.dynamicSmemBytes = kSmem; + cfg.stream = stream; + cfg.attrs = attr; + cfg.numAttrs = unsigned(na); + CHECK_CUDA(cudaLaunchKernelEx(&cfg, oproj_dense_ar_kernel, x_tmap, w_tmap, prm)); + } +}; + +template < + int M, + int K, + int R, + Comm COMM, + bool GEMM_ON, + int S = ((M + 15) & ~15) <= 16 ? 6 : (((M + 15) & ~15) == 32 ? 4 : 3), + int CH = 2, + int C = 1> // cluster-multicast A axis: C=2 REFUTED as-built at + // S=3/4 rings (LEDGER R9 — pair-lockstep pacing beats + // the halved A bytes); flip here to test siblings +struct Launcher { + static constexpr int Mp = (M + 15) & ~15; + static constexpr size_t kSmem = GEMM_ON + ? size_t(S) * CH * (kBNRows * kBK * 2 + Mp * kBK * 2) + 2 * S * sizeof(uint64_t) + : 4096; // AR-only path never touches the feed ring + // once per device, before first launch + static void set_smem_attr() { + CHECK_CUDA(cudaFuncSetAttribute( + oproj_ar_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, int(kSmem))); + // 100% carveout → the SM smem config fits TWO CTAs (this grid's tail + // + the next PDL grid's feed); the default config blocks dual + // residency and with it the whole tail-hiding scheme. + CHECK_CUDA(cudaFuncSetAttribute( + oproj_ar_kernel, cudaFuncAttributePreferredSharedMemoryCarveout, 100)); + } + // pdl=false = the NON-COOPERATIVE-neighbor regime: no programmatic + // serialization attribute, so successive kernels fully serialize — the + // AR tail is exposed, as it is in a serving stack whose adjacent kernels + // don't PDL-cooperate or can't co-reside. In-kernel griddepcontrol ops + // are no-ops without the attribute; the done-guard is trivially met. + static void launch( + const CUtensorMap& w_map, + const CUtensorMap& x_map, + const Params& prm, + int ncta, + cudaStream_t stream, + bool pdl) { + cudaLaunchConfig_t cfg{}; + cudaLaunchAttribute attr[2]; + int na = 0; + if (pdl) { + attr[na].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[na].val.programmaticStreamSerializationAllowed = 1; + ++na; + } + if constexpr (C > 1) { + attr[na].id = cudaLaunchAttributeClusterDimension; + attr[na].val.clusterDim.x = C; + attr[na].val.clusterDim.y = 1; + attr[na].val.clusterDim.z = 1; + ++na; + } + cfg.gridDim = dim3(unsigned(ncta)); + cfg.blockDim = dim3(kThreads); + cfg.dynamicSmemBytes = kSmem; + cfg.stream = stream; + cfg.attrs = attr; + cfg.numAttrs = unsigned(na); + CHECK_CUDA(cudaLaunchKernelEx(&cfg, oproj_ar_kernel, w_map, x_map, prm)); + } +}; + +} // namespace oproj_ar + +// ================= sglang tvm-ffi adapter ================= + +#include // For TensorMatcher, SymbolicSize, SymbolicDevice +#include // For CHECK_HOST + +#include // For bf16_t, TVMFFIEnvGetStream + +#include + +#include +#include +#include + +namespace oproj_ar_ffi { + +using namespace oproj_ar; +using tvm::ffi::TensorView; + +constexpr int kCellList[7] = {8, 16, 32, 64, 128, 256, 512}; + +inline int cell_of(int m) { + for (int c : kCellList) + if (m <= c) return c; + return -1; +} + +template +struct GemmArKernel { + static_assert(R >= 2 && R <= 8, "R outside the validated 2..8 range"); + static_assert(K % 128 == 0 && K >= 128, "K must be a multiple of 128"); + + static constexpr int kTwoShotMinM = R >= 8 ? 128 : 256; + + // one ring family per dispatch cell (see kFams) + static int64_t fam_of(int64_t m) { + const int cell = cell_of(int(m)); + CHECK_HOST(cell > 0) << "gemm_ar: M=" << m << " outside [1, " << kMMax << "]"; + for (int i = 0; i < 7; ++i) + if (kCellList[i] == cell) return i; + return -1; + } + + static int64_t cell_of_ffi(int64_t m) { + return cell_of(int(m)); + } + static int64_t region_nbytes() { + return int64_t(region_bytes(R)); + } + static int64_t gather_words() { + return int64_t(kFams) * 2 * kRing; + } + static int64_t num_fams() { + return kFams; + } + static int64_t max_tokens() { + return kMMax; + } + + // Per-weight-pointer W tensor maps (encode once; weights are static). + struct WMaps { + CUtensorMap w48, w64, w128, w256; + }; + + static const WMaps& w_maps(void* w) { + static std::unordered_map cache; + static std::mutex mu; + std::lock_guard lk(mu); + auto it = cache.find(w); + if (it == cache.end()) { + auto enc = [&](uint32_t box_rows) { + return tmap::encode_tiled_2d( + w, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, kN, K, size_t(K) * 2, box_rows, kBK, CU_TENSOR_MAP_SWIZZLE_128B); + }; + it = cache.emplace(w, WMaps{enc(kBNRows), enc(64), enc(128), enc(256)}).first; + } + return it->second; + } + + // x tensor map over the caller's [M, K] tensor: global rows = M, TMA + // zero-fills the [M, cell) padding rows out-of-bounds. + static CUtensorMap x_map(void* x, int m, uint32_t box_rows) { + return tmap::encode_tiled_2d( + x, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, uint64_t(m), K, size_t(K) * 2, box_rows, kBK, CU_TENSOR_MAP_SWIZZLE_128B); + } + + static void set_smem_attrs_once() { + static bool done = [] { + Launcher<8, K, R, Comm::kPeer, true>::set_smem_attr(); + Launcher<16, K, R, Comm::kPeer, true>::set_smem_attr(); + Launcher3<32, K, R, Comm::kPeer>::set_smem_attr(); + Launcher3<64, K, R, Comm::kPeer>::set_smem_attr(); + if constexpr (kTwoShotMinM > 128) Launcher3<128, K, R, Comm::kPeer>::set_smem_attr(); + if constexpr (kTwoShotMinM <= 128) Launcher3<128, K, R, Comm::kTwoShotPeer>::set_smem_attr(); + Launcher3<256, K, R, Comm::kTwoShotPeer>::set_smem_attr(); + Launcher3<512, K, R, Comm::kTwoShotPeer>::set_smem_attr(); + return true; + }(); + (void)done; + } + + template + static void enqueue_cell(const WMaps& wm, void* x, int m, const Params& prm, cudaStream_t stream, bool pdl) { + if constexpr (CELL <= 16) { + const CUtensorMap xm = x_map(x, m, 16); + Launcher::launch(wm.w48, xm, prm, kM1Grid_(), stream, pdl); + } else if constexpr (CELL < kTwoShotMinM) { + const CUtensorMap xm = x_map(x, m, CELL <= 32 ? 16 : (CELL <= 64 ? 32 : 128)); + const CUtensorMap& wmap = CELL <= 64 ? wm.w128 : wm.w64; + Launcher3::launch(xm, wmap, prm, stream, pdl); + } else { + const CUtensorMap xm = x_map(x, m, 128); + const CUtensorMap& wmap = CELL == 128 ? wm.w64 : (CELL == 256 ? wm.w128 : wm.w256); + Launcher3::launch(xm, wmap, prm, stream, pdl); + } + } + + static constexpr int kM1Grid_() { + return 152; + } + + // per-rank UC VAs of the comm region, stashed host-side ONCE at init: + // per-call CPU-tensor derefs from inside the op are not reliable in every + // execution context (observed dangling under the sglang scheduler). + static std::array& bases_store() { + static std::array a{}; + return a; + } + + static void set_bases(TensorView uc_bases) { + using namespace host; + TensorMatcher({R}).with_dtype().verify(uc_bases); + const int64_t* b = static_cast(uc_bases.data_ptr()); + for (int r = 0; r < R; ++r) + bases_store()[r] = reinterpret_cast(b[r]); + } + + static void + run(TensorView out, + TensorView x, + TensorView w, + TensorView gather, // [kFams * 2 * kRing] int32 CUDA, device-local + TensorView epochs, // [kFams] int32 CUDA: device-resident CTA ticket counters + int64_t my_rank) { + using namespace host; + + auto M = SymbolicSize{"num_tokens"}; + auto CellRows = SymbolicSize{"cell_rows"}; + auto device = SymbolicDevice{}; + device.set_options(); + TensorMatcher({M, K}).with_dtype().with_device(device).verify(x); + TensorMatcher({kN, K}).with_dtype().with_device(device).verify(w); + TensorMatcher({CellRows, kN}).with_dtype().with_device(device).verify(out); + + const int m = int(M.unwrap()); + const int cell = cell_of(m); + CHECK_HOST(cell > 0) << "gemm_ar: M=" << m << " outside [1, " << kMMax << "]"; + CHECK_HOST(int64_t(CellRows.unwrap()) == cell) + << "out must have cell(M)=" << cell << " rows, got " << CellRows.unwrap(); + CHECK_HOST(my_rank >= 0 && my_rank < R); + CHECK_HOST(bases_store()[0] != nullptr) << "gemm_ar: set_bases not called"; + TensorMatcher({int64_t(kFams) * 2 * kRing}).with_dtype().verify(gather); + TensorMatcher({kFams}).with_dtype().verify(epochs); + + set_smem_attrs_once(); + + const DLDevice dev = device.unwrap(); + const auto stream = static_cast(::TVMFFIEnvGetStream(dev.device_type, dev.device_id)); + + Params prm{}; + prm.mc_base = nullptr; // pure-P2P plane + for (int r = 0; r < R; ++r) + prm.uc_base[r] = bases_store()[r]; + prm.gather = static_cast(gather.data_ptr()); + prm.out = static_cast<__nv_bfloat16*>(out.data_ptr()); + prm.partial_in = nullptr; + prm.epoch_base = static_cast(epochs.data_ptr()); + prm.my_rank = int(my_rank); + prm.fam = int(fam_of(m)); + + const bool pdl = kUsePDL; + const WMaps& wm = w_maps(w.data_ptr()); + switch (cell) { + case 8: + enqueue_cell<8>(wm, x.data_ptr(), m, prm, stream, pdl); + break; + case 16: + enqueue_cell<16>(wm, x.data_ptr(), m, prm, stream, pdl); + break; + case 32: + enqueue_cell<32>(wm, x.data_ptr(), m, prm, stream, pdl); + break; + case 64: + enqueue_cell<64>(wm, x.data_ptr(), m, prm, stream, pdl); + break; + case 128: + enqueue_cell<128>(wm, x.data_ptr(), m, prm, stream, pdl); + break; + case 256: + enqueue_cell<256>(wm, x.data_ptr(), m, prm, stream, pdl); + break; + default: + enqueue_cell<512>(wm, x.data_ptr(), m, prm, stream, pdl); + break; + } + CHECK_CUDA(cudaGetLastError()) << "gemm_ar launch (cell=" << cell << ")"; + } +}; + +} // namespace oproj_ar_ffi + +using oproj_ar_ffi::GemmArKernel; diff --git a/python/sglang/kernels/jit/csrc/kimi_k3/comm/ptx_sys.cuh b/python/sglang/kernels/jit/csrc/kimi_k3/comm/ptx_sys.cuh new file mode 100644 index 000000000..a1a0c6892 --- /dev/null +++ b/python/sglang/kernels/jit/csrc/kimi_k3/comm/ptx_sys.cuh @@ -0,0 +1,68 @@ +#pragma once + +// System-scope PTX the Kimi K3 collectives need and no shared sglang header +// wraps. These live beside their only consumers (gemm_ar / gemm_ag) rather than +// in distributed/communicator.cuh: that header's Semaphore does not use any of +// them, so putting them there would grow a shared header for one caller's +// benefit. +// +// The `device::distributed` namespace is deliberate -- it is where the rest of +// the collective vocabulary lives, so call sites and `using` declarations read +// the same whichever header supplied the symbol. + +#include + +#include + +namespace device::distributed { + +// Peer-visible flag increment. `.sys` scope, relaxed: ordering is established by +// the surrounding fence, not by this store. +SGL_DEVICE void red_add_relaxed_sys(uint32_t* ptr, uint32_t val) { + asm volatile("red.relaxed.sys.global.add.u32 [%0], %1;" : : "l"(ptr), "r"(val) : "memory"); +} + +// Acquire load of a peer-written flag: everything the writer released before +// its matching store is visible to this thread afterwards. +SGL_DEVICE uint32_t load_acquire_sys(const uint32_t* ptr) { + uint32_t val; + asm volatile("ld.acquire.sys.global.u32 %0, [%1];" : "=r"(val) : "l"(ptr) : "memory"); + return val; +} + +// Publishes every prior write to system scope. Pair with a relaxed flag store so +// a peer's acquire load of that flag also observes the payload. +SGL_DEVICE void fence_release_sys() { + asm volatile("fence.release.sys;" ::: "memory"); +} + +// Device-scope arrival counter. acq_rel so the winner of the count also observes +// the losers' payload writes. +SGL_DEVICE uint32_t atomic_add_acq_rel_gpu(uint32_t* ptr, uint32_t val) { + uint32_t old; + asm volatile("atom.acq_rel.gpu.global.add.u32 %0, [%1], %2;" : "=r"(old) : "l"(ptr), "r"(val) : "memory"); + return old; +} + +// One store fanned out to every rank in the multicast team. +SGL_DEVICE void multimem_store_relaxed(uint32_t* ptr, uint32_t val) { + asm volatile("multimem.st.relaxed.sys.global.b32 [%0], %1;" : : "l"(ptr), "r"(val) : "memory"); +} + +SGL_DEVICE void multimem_red_add_relaxed(uint32_t* mc_flag) { +#if SGL_ARCH_HOPPER_OR_GREATER + asm volatile("multimem.red.relaxed.sys.global.add.u32 [%0], 1;" ::"l"(mc_flag) : "memory"); +#else + assert(false && "multimem red is only supported on Hopper or later architecture"); +#endif +} + +SGL_DEVICE void multimem_red_add_release(uint32_t* mc_flag) { +#if SGL_ARCH_HOPPER_OR_GREATER + asm volatile("multimem.red.release.sys.global.add.u32 [%0], 1;" ::"l"(mc_flag) : "memory"); +#else + assert(false && "multimem red is only supported on Hopper or later architecture"); +#endif +} + +} // namespace device::distributed diff --git a/python/sglang/kernels/jit/csrc/kimi_k3/comm/sp_collective.cuh b/python/sglang/kernels/jit/csrc/kimi_k3/comm/sp_collective.cuh new file mode 100644 index 000000000..40435807d --- /dev/null +++ b/python/sglang/kernels/jit/csrc/kimi_k3/comm/sp_collective.cuh @@ -0,0 +1,439 @@ +// K3 SP-MoE row-sharded collectives (bf16): +// +// reduce_scatter_res: +// [world * rows, hidden] -> [rows, hidden], optionally adding the +// destination rank's residual rows in the reduction epilogue. +// Every input vector is written exactly once, to the rank that owns its +// row shard; the destination polls and reduces the world producer slots. +// +// all_gather: +// [rows, hidden] -> [world * rows, hidden]. Every rank multicast-stores +// its local shard once, then every peer polls the rank slots and copies +// them into rank-concatenated row order. +// +// Both kernels reuse CustomAllReduceV2's double-buffered push workspace and +// phase counters. A bumper block advances counters outside the tuned work +// grid, so calls remain protocol-compatible with all-reduce and gemm_ag. + +#include +#include + +#include "../../distributed/custom_all_reduce.cuh" + +namespace sglang::sp_collective { + +using device::distributed::Counter; +using device::distributed::Semaphore; + +struct Params { + const uint8_t* input; + uint8_t* output; + const uint8_t* residual; + uint8_t* push_workspaces[device::distributed::kMaxWorldSize]; + uint8_t* push_ws_mc; + Counter* counter; + Semaphore* sem_local; + uint8_t* sem_mc; + uint8_t* input_mc; + uint8_t* output_mc; + int64_t stride_bytes; + uint32_t num_counters; + uint32_t rank; + uint32_t local_vecs; + uint32_t residual_is_local; +}; + +template +SGL_DEVICE void make_nonzero(Vec& vec) { + constexpr uint32_t kNegZeroPair = 0x8000u; + auto& bits = *reinterpret_cast(&vec); + if (bits.x == 0) bits.x = kNegZeroPair; + if (bits.y == 0) bits.y = kNegZeroPair; + if (bits.z == 0) bits.z = kNegZeroPair; + if (bits.w == 0) bits.w = kNegZeroPair; +} + +SGL_DEVICE uint32_t* sem_mc_flag(uint8_t* sem_mc, uint32_t block) { + static_assert(sizeof(Semaphore) == 128); + return reinterpret_cast(sem_mc + block * sizeof(Semaphore)); +} + +SGL_DEVICE void sem_arrive_relaxed(uint32_t* flag) { +#if SGL_ARCH_HOPPER_OR_GREATER + asm volatile("multimem.red.relaxed.sys.global.add.u32 [%0], 1;" ::"l"(flag) : "memory"); +#else + assert(false && "multimem red requires Hopper or later"); +#endif +} + +SGL_DEVICE void sem_arrive_release(uint32_t* flag) { +#if SGL_ARCH_HOPPER_OR_GREATER + asm volatile("multimem.red.release.sys.global.add.u32 [%0], 1;" ::"l"(flag) : "memory"); +#else + assert(false && "multimem red requires Hopper or later"); +#endif +} + +template +SGL_DEVICE bool has_empty_marker(const Vec& vec) { + const auto bits = *reinterpret_cast(&vec); + return bits.x == 0 || bits.y == 0 || bits.z == 0 || bits.w == 0; +} + +template +SGL_DEVICE Vec zero_vec() { + Vec zero; + zero.fill(bf16x2_t{get_pos_zero(), get_pos_zero()}); + return zero; +} + +template +SGL_DEVICE bool bumper_block(const Params& params) { + const auto bx = blockIdx.x; + if (bx + 1 != gridDim.x) return false; + const auto phase = params.counter[bx].get() & 1; + __syncthreads(); + for (uint32_t i = bx + threadIdx.x; i < params.num_counters; i += blockDim.x) { + params.counter[i].set(phase ^ 1); + } + return true; +} + +template +__global__ void reduce_scatter_res_kernel(const __grid_constant__ Params params) { + using vec_t = device::AlignedVector; // 16 B + + device::PDLWaitPrimary(); + if (bumper_block(params)) { + device::PDLTriggerSecondary(); + return; + } + + const uint32_t bx = blockIdx.x; + const uint32_t tid = bx * blockDim.x + threadIdx.x; + const uint32_t num_threads = (gridDim.x - 1) * blockDim.x; + const uint32_t phase = params.counter[bx].get() & 1; + const auto phase_offset = phase * kWorldSize * params.stride_bytes; + const auto producer_offset = phase_offset + params.rank * params.stride_bytes; + + // Each vector goes only to the rank that owns its row shard. + for (uint32_t vid = tid; vid < kWorldSize * params.local_vecs; vid += num_threads) { + const uint32_t dst_rank = vid / params.local_vecs; + const uint32_t local_vid = vid - dst_rank * params.local_vecs; + vec_t vec; + ld_global_16B(vec, params.input, vid); + make_nonzero(vec); + st_relaxed_16B(vec, params.push_workspaces[dst_rank] + producer_offset, local_vid); + } + + device::PDLTriggerSecondary(); + + // Poll this rank's shard from all producer slots and reduce locally. + const auto poll_base = params.push_workspaces[params.rank] + phase_offset; + const auto residual_base = params.residual + (params.residual_is_local ? 0 : params.rank * params.local_vecs * 16); + const auto zero = zero_vec(); + for (uint32_t vid = tid; vid < params.local_vecs; vid += num_threads) { + vec_t vec[kWorldSize + kHasResidual]; + if constexpr (kHasResidual) { + ld_global_16B(vec[kWorldSize], residual_base, vid); + } + do { + bool empty = false; +#pragma unroll + for (uint32_t rank = 0; rank < kWorldSize; ++rank) { + ld_relaxed_16B(vec[rank], poll_base + rank * params.stride_bytes, vid); + empty |= has_empty_marker(vec[rank]); + } + if (!empty) break; + } while (true); + const auto out = reduce(vec); + st_global_16B(out, params.output, vid); +#pragma unroll + for (uint32_t rank = 0; rank < kWorldSize; ++rank) { + st_global_16B(zero, poll_base + rank * params.stride_bytes, vid); + } + } + + __syncthreads(); + if (threadIdx.x == 0) params.counter[bx].set(phase ^ 1); +} + +template +__global__ void all_gather_kernel(const __grid_constant__ Params params) { + using vec_t = device::AlignedVector; // 16 B + + device::PDLWaitPrimary(); + if (bumper_block(params)) { + device::PDLTriggerSecondary(); + return; + } + + const uint32_t bx = blockIdx.x; + const uint32_t tid = bx * blockDim.x + threadIdx.x; + const uint32_t num_threads = (gridDim.x - 1) * blockDim.x; + const uint32_t phase = params.counter[bx].get() & 1; + const auto phase_offset = phase * kWorldSize * params.stride_bytes; + const auto producer_offset = phase_offset + params.rank * params.stride_bytes; + + // One multicast store places this rank's shard in the same slot on peers. + for (uint32_t vid = tid; vid < params.local_vecs; vid += num_threads) { + vec_t vec; + ld_global_16B(vec, params.input, vid); + make_nonzero(vec); + st_multimem_16B(vec, params.push_ws_mc + producer_offset, vid); + } + + device::PDLTriggerSecondary(); + + const auto poll_base = params.push_workspaces[params.rank] + phase_offset; + const auto zero = zero_vec(); + for (uint32_t vid = tid; vid < kWorldSize * params.local_vecs; vid += num_threads) { + const uint32_t src_rank = vid / params.local_vecs; + const uint32_t local_vid = vid - src_rank * params.local_vecs; + const auto src = poll_base + src_rank * params.stride_bytes; + vec_t vec; + do { + ld_relaxed_16B(vec, src, local_vid); + } while (has_empty_marker(vec)); + st_global_16B(vec, params.output, vid); + st_global_16B(zero, src, local_vid); + } + + __syncthreads(); + if (threadIdx.x == 0) params.counter[bx].set(phase ^ 1); +} + +// Direct variant: output is multicast-bound symmetric memory. Each producer +// writes its rank slice straight into every peer's final output, avoiding the +// staging read/copy/clear. The two pull-semaphore barriers preserve protocol +// compatibility with CustomAllReduceV2 and make the remote writes visible +// before any rank leaves the kernel. +template +__global__ void all_gather_direct_kernel(const __grid_constant__ Params params) { + using vec_t = device::AlignedVector; // 16 B + + uint32_t exit_base = 0; + if (threadIdx.x == 0) { + auto* semaphore = ¶ms.sem_local[blockIdx.x]; + const auto reserved = semaphore->counter_ptr()->inc(2 * kWorldSize); + exit_base = reserved + kWorldSize; + device::PDLWaitPrimary(); + sem_arrive_relaxed(sem_mc_flag(params.sem_mc, blockIdx.x)); + while (semaphore->get_relaxed() - reserved < kWorldSize) + ; + } + __syncthreads(); + + const uint32_t tid = blockIdx.x * blockDim.x + threadIdx.x; + const uint32_t step = gridDim.x * blockDim.x; + const uint32_t dst_bias = params.rank * params.local_vecs; + for (uint32_t vid = tid; vid < params.local_vecs; vid += step) { + vec_t vec; + ld_global_16B(vec, params.input, vid); + st_multimem_16B(vec, params.output_mc, dst_bias + vid); + } + + device::PDLTriggerSecondary(); + __syncthreads(); + if (threadIdx.x == 0) { + auto* semaphore = ¶ms.sem_local[blockIdx.x]; + sem_arrive_release(sem_mc_flag(params.sem_mc, blockIdx.x)); + while (semaphore->get_acquire() - exit_base < kWorldSize) + ; + } +} + +// NVLS pull variant: o_proj writes its TP-partial result into multicast-bound +// symmetric memory. Each rank reduces only its owned row shard directly from +// the multicast alias, so no staging or second broadcast is needed. +template +__global__ void reduce_scatter_pull_kernel(const __grid_constant__ Params params) { + using vec_t = device::AlignedVector; // 16 B + using SumOp = device::ReductionTrait; + + uint32_t exit_base = 0; + if (threadIdx.x == 0) { + auto* semaphore = ¶ms.sem_local[blockIdx.x]; + const auto reserved = semaphore->counter_ptr()->inc(2 * kWorldSize); + exit_base = reserved + kWorldSize; + device::PDLWaitPrimary(); + sem_arrive_relaxed(sem_mc_flag(params.sem_mc, blockIdx.x)); + while (semaphore->get_relaxed() - reserved < kWorldSize) + ; + } + __syncthreads(); + + const uint32_t tid = blockIdx.x * blockDim.x + threadIdx.x; + const uint32_t step = gridDim.x * blockDim.x; + const auto* input_mc = params.input_mc + params.rank * params.local_vecs * 16; + const auto* residual = + kHasResidual ? params.residual + (params.residual_is_local ? 0 : params.rank * params.local_vecs * 16) : nullptr; + for (uint32_t vid = tid; vid < params.local_vecs; vid += step) { + vec_t vec; + ld_multimem_16B(vec, input_mc, vid); + if constexpr (kHasResidual) { + vec_t res; + ld_global_16B(res, residual, vid); +#pragma unroll + for (uint32_t j = 0; j < 4; ++j) { + vec[j] = SumOp::reduce(vec[j], res[j]); + } + } + st_global_16B(vec, params.output, vid); + } + + device::PDLTriggerSecondary(); + __syncthreads(); + if (threadIdx.x == 0) { + auto* semaphore = ¶ms.sem_local[blockIdx.x]; + sem_arrive_release(sem_mc_flag(params.sem_mc, blockIdx.x)); + while (semaphore->get_acquire() - exit_base < kWorldSize) + ; + } +} + +} // namespace sglang::sp_collective + +using namespace sglang; +using host::distributed::CommunicatorRef; + +template +struct SPCollectiveKernel { + using TensorView = tvm::ffi::TensorView; + + static sp_collective::Params make_params( + const host::distributed::CommunicatorObj& data, + TensorView input, + TensorView output, + std::optional residual, + bool residual_is_local, + int64_t ws_mc_base) { + using namespace host; + auto input_elems = SymbolicSize{"input_elems"}; + auto local_elems = SymbolicSize{"local_elems"}; + auto device = SymbolicDevice{}; + device.set_options(); + TensorMatcher({input_elems}).with_dtype().with_device(device).verify(input); + TensorMatcher({local_elems}).with_dtype().with_device(device).verify(output); + if (residual.has_value()) { + if (residual_is_local) { + TensorMatcher({local_elems}).with_dtype().with_device(device).verify(residual.value()); + } else { + TensorMatcher({input_elems}).with_dtype().with_device(device).verify(residual.value()); + } + } + CHECK_HOST(data.world_size == kWorldSize); + CHECK_HOST(local_elems.unwrap() > 0); + CHECK_HOST(input_elems.unwrap() == local_elems.unwrap() * kWorldSize); + CHECK_HOST(local_elems.unwrap() % 8 == 0) << "local shard bytes must be 16B aligned"; + CHECK_HOST(local_elems.unwrap() * sizeof(bf16_t) <= data.push_bytes) << "local shard exceeds a push slot"; + + sp_collective::Params params{ + .input = static_cast(input.data_ptr()), + .output = static_cast(output.data_ptr()), + .residual = residual.has_value() ? static_cast(residual.value().data_ptr()) : nullptr, + .push_workspaces = {}, + .push_ws_mc = reinterpret_cast(ws_mc_base), + .counter = data.push_counter, + .sem_local = data.pull_semaphores[data.rank], + .sem_mc = nullptr, + .input_mc = nullptr, + .output_mc = nullptr, + .stride_bytes = data.push_bytes, + .num_counters = data.num_push_blocks, + .rank = data.rank, + .local_vecs = static_cast(local_elems.unwrap() * sizeof(bf16_t) / 16), + .residual_is_local = static_cast(residual_is_local), + }; + for (uint32_t i = 0; i < kWorldSize; ++i) { + params.push_workspaces[i] = data.push_workspaces[i]; + } + return params; + } + + static void check_launch(const host::distributed::CommunicatorObj& data, int64_t num_blocks, int64_t block_size) { + CHECK_HOST(num_blocks > 0 && num_blocks < data.num_push_blocks); + // The RS reduction keeps one 16B vector per producer in registers. + // 1024-thread CTAs exceed the GB300 launch resource limit. + CHECK_HOST(block_size >= 32 && block_size <= 512 && block_size % 32 == 0); + } + + static void reduce_scatter_res( + CommunicatorRef ref, + TensorView input, + TensorView output, + std::optional residual, + bool residual_is_local, + int64_t num_blocks, + int64_t block_size) { + const auto& data = *ref.get(); + check_launch(data, num_blocks, block_size); + auto params = make_params(data, input, output, residual, residual_is_local, 0); + const auto kernel = residual.has_value() ? sp_collective::reduce_scatter_res_kernel + : sp_collective::reduce_scatter_res_kernel; + host::LaunchKernel(num_blocks + 1, block_size, input.device()).enable_pdl(kUsePDL)(kernel, params); + } + + static void all_gather( + CommunicatorRef ref, + TensorView input, + TensorView output, + int64_t ws_mc_base, + int64_t num_blocks, + int64_t block_size) { + const auto& data = *ref.get(); + CHECK_HOST(ws_mc_base != 0) << "all-gather requires multicast workspace"; + check_launch(data, num_blocks, block_size); + // Reuse the RS matcher by swapping input/output roles conceptually. + auto params = make_params(data, output, input, std::nullopt, false, ws_mc_base); + params.input = static_cast(input.data_ptr()); + params.output = static_cast(output.data_ptr()); + host::LaunchKernel(num_blocks + 1, block_size, input.device()) + .enable_pdl(kUsePDL)(sp_collective::all_gather_kernel, params); + } + + static void all_gather_direct( + CommunicatorRef ref, + TensorView input, + TensorView output, + int64_t output_mc_ptr, + int64_t sem_mc_ptr, + int64_t num_blocks, + int64_t block_size) { + const auto& data = *ref.get(); + CHECK_HOST(output_mc_ptr != 0) << "direct all-gather needs symmetric output"; + CHECK_HOST(sem_mc_ptr != 0) << "direct all-gather needs multicast semaphores"; + CHECK_HOST(num_blocks > 0 && num_blocks <= data.num_pull_blocks); + CHECK_HOST(block_size >= 32 && block_size <= 1024 && block_size % 32 == 0); + auto params = make_params(data, output, input, std::nullopt, false, 0); + params.input = static_cast(input.data_ptr()); + params.output = static_cast(output.data_ptr()); + params.output_mc = reinterpret_cast(static_cast(output_mc_ptr)); + params.sem_mc = reinterpret_cast(static_cast(sem_mc_ptr)); + host::LaunchKernel(num_blocks, block_size, input.device()) + .enable_pdl(kUsePDL)(sp_collective::all_gather_direct_kernel, params); + } + + static void reduce_scatter_pull( + CommunicatorRef ref, + TensorView input, + TensorView output, + std::optional residual, + bool residual_is_local, + int64_t input_mc_ptr, + int64_t sem_mc_ptr, + int64_t num_blocks, + int64_t block_size) { + const auto& data = *ref.get(); + CHECK_HOST(input_mc_ptr != 0) << "pull RS needs symmetric input"; + CHECK_HOST(sem_mc_ptr != 0) << "pull RS needs multicast semaphores"; + CHECK_HOST(num_blocks > 0 && num_blocks <= data.num_pull_blocks); + CHECK_HOST(block_size >= 32 && block_size <= 1024 && block_size % 32 == 0); + auto params = make_params(data, input, output, residual, residual_is_local, 0); + params.input_mc = reinterpret_cast(static_cast(input_mc_ptr)); + params.sem_mc = reinterpret_cast(static_cast(sem_mc_ptr)); + const auto kernel = residual.has_value() ? sp_collective::reduce_scatter_pull_kernel + : sp_collective::reduce_scatter_pull_kernel; + host::LaunchKernel(num_blocks, block_size, input.device()).enable_pdl(kUsePDL)(kernel, params); + } +}; diff --git a/python/sglang/kernels/jit/csrc/kimi_k3/mla_output_gate.cuh b/python/sglang/kernels/jit/csrc/kimi_k3/mla_output_gate.cuh new file mode 100644 index 000000000..0cf73f79e --- /dev/null +++ b/python/sglang/kernels/jit/csrc/kimi_k3/mla_output_gate.cuh @@ -0,0 +1,84 @@ +// K3 MLA output gate: out = bf16(bf16(x) * bf16(sigmoid(gate))), replacing +// the torch.sigmoid + mul elementwise pair (two launches, two memory passes) +// with one kernel. sigmoid is computed in fp32 and rounded to bf16 before the +// multiply, reproducing the unfused pair's double rounding bit-for-bit. + +#include // For TensorMatcher, SymbolicSize, SymbolicDevice +#include // For RuntimeCheck + +#include // For bf16_t, fp32_t, device::cast +#include // For LaunchKernel +#include // For AlignedVector + +#include + +#include + +namespace { + +struct MlaOutputGateParams { + const bf16_t* __restrict__ x; // [N] contiguous (flattened [T, H]) + const bf16_t* __restrict__ gate; // [N] contiguous + bf16_t* __restrict__ out; // [N] contiguous + uint32_t n_vecs; +}; + +template +__global__ void mla_output_gate_kernel(const MlaOutputGateParams __grid_constant__ params) { + using namespace device; + + constexpr int kVecN = 8; + using vec_bf16_t = AlignedVector; + + const uint32_t v = blockIdx.x * kThreads + threadIdx.x; + if (v >= params.n_vecs) return; + + PDLWaitPrimary(); + + vec_bf16_t xv, gv, ov; + xv.load(params.x, v); + gv.load(params.gate, v); +#pragma unroll + for (int i = 0; i < kVecN; ++i) { + // Match torch.sigmoid(bf16): fp32 sigmoid, round to bf16, then the bf16 + // multiply upcasts both operands to fp32 and rounds once more. + const float g = cast(gv[i]); + const bf16_t s = cast(1.0f / (1.0f + expf(-g))); + ov[i] = cast(cast(xv[i]) * cast(s)); + } + ov.store(params.out, v); + + PDLTriggerSecondary(); +} + +template +struct MlaOutputGateKernel { + static constexpr auto kernel = mla_output_gate_kernel; + + static void run(const tvm::ffi::TensorView x, const tvm::ffi::TensorView gate, const tvm::ffi::TensorView out) { + using namespace host; + + auto N_ = SymbolicSize{"numel"}; + auto device = SymbolicDevice{}; + device.set_options(); + + TensorMatcher({N_}).with_dtype().with_device(device).verify(x); + TensorMatcher({N_}).with_dtype().with_device(device).verify(gate); + TensorMatcher({N_}).with_dtype().with_device(device).verify(out); + + const auto N = static_cast(N_.unwrap()); + RuntimeCheck(N % 8 == 0, "numel must be divisible by 8"); + if (N == 0) return; + + const auto params = MlaOutputGateParams{ + .x = static_cast(x.data_ptr()), + .gate = static_cast(gate.data_ptr()), + .out = static_cast(out.data_ptr()), + .n_vecs = N / 8, + }; + const uint32_t n_blocks = (params.n_vecs + kThreads - 1) / kThreads; + LaunchKernel(n_blocks, kThreads, device.unwrap()).enable_pdl(kUsePDL)(kernel, params); + } +}; + +} // namespace diff --git a/python/sglang/kernels/jit/csrc/kimi_k3/situ_and_mul.cuh b/python/sglang/kernels/jit/csrc/kimi_k3/situ_and_mul.cuh new file mode 100644 index 000000000..b2eb01106 --- /dev/null +++ b/python/sglang/kernels/jit/csrc/kimi_k3/situ_and_mul.cuh @@ -0,0 +1,450 @@ +// Kimi K3 SiTU activation kernels: plain elementwise and varlen masked with a +// grouped-quant epilogue. The shared double-softcap activation is inlined below. + +#pragma once + +#include // For TensorMatcher, SymbolicSize, SymbolicDevice +#include // For RuntimeCheck, div_ceil + +#include +#include +#include // For dtype_trait, bf16_t, fp32_t, cast +#include // For LaunchKernel, SGL_DEVICE, PDL helpers +#include // For AlignedVector +#include // For warp::copy_bytes, elect_one_lane, inclusive_sum + +#include + +#include + +#include +#include +#include +#include + +namespace sglang { + +namespace kimi_k3 { + +/// One SiTU element. `sigmoid_fast` is `1/(1+expf(-x))` (math.cuh), i.e. the +/// same expression both call sites used before they were folded together. +template +SGL_DEVICE float situ_activate(float g, float u, float beta, float inv_beta, float linear_beta, float inv_linear_beta) { + const float gate_out = beta * tanhf(g * inv_beta) * device::math::sigmoid_fast(g); + float up_out; + if constexpr (kHasLinearBeta) { + up_out = linear_beta * tanhf(u * inv_linear_beta); + } else { + up_out = u; + } + return gate_out * up_out; +} + +} // namespace kimi_k3 + +} // namespace sglang + +namespace { + +// SiTU (SoftCap-GLU) activation: +// gate_out = beta * tanh(gate / beta) * sigmoid(gate) +// up_out = linear_beta * tanh(up / linear_beta) +// output = gate_out * up_out +// +// Input: bf16 tensor [N, 2*D] (gate = [:, :D], up = [:, D:]) +// Output: bf16 tensor [N, D] + +struct SituAndMulParams { + const void* __restrict__ input; + void* __restrict__ out; + float beta; + float inv_beta; + float linear_beta; + float inv_linear_beta; + uint32_t hidden_dim; // D (output width, half of input last dim) + uint32_t num_tokens; + uint32_t stride_in_vecs; // input row stride in vector units (2*D/vec if dense) +}; + +template +__global__ void situ_and_mul_kernel(const __grid_constant__ SituAndMulParams params) { + using namespace device; + constexpr auto kVecSize = kMaxVecBytes / sizeof(T); + using vec_t = AlignedVector; + + const auto num_vecs = params.hidden_dim / kVecSize; // per token + const auto tid = blockIdx.x * blockDim.x + threadIdx.x; + const auto token_id = tid / num_vecs; + + if (token_id >= params.num_tokens) return; + + const auto offset = tid % num_vecs; + // Input rows may be strided (e.g. a slice of a wider fused-GEMM output); + // within a row: gate = [0..D-1], up = [D..2D-1]. + const auto input_offset = static_cast(token_id) * params.stride_in_vecs + offset; + const auto output_offset = tid; + + PDLWaitPrimary(); + + const auto gate = load_as(params.input, input_offset); + const auto up = load_as(params.input, input_offset + num_vecs); + + PDLTriggerSecondary(); + + const float beta = params.beta; + const float inv_beta = params.inv_beta; + const float linear_beta = params.linear_beta; + const float inv_linear_beta = params.inv_linear_beta; + + vec_t out; +#pragma unroll + for (int i = 0; i < kVecSize; ++i) { + const float g = cast(gate[i]); + const float u = cast(up[i]); + + out[i] = + cast(sglang::kimi_k3::situ_activate(g, u, beta, inv_beta, linear_beta, inv_linear_beta)); + } + + store_as(params.out, out, output_offset); +} + +// Host launcher + +template +struct SituAndMulKernel { + static constexpr auto kVecSize = device::kMaxVecBytes / sizeof(T); + static constexpr auto kBlockSize = 256u; + + static void + run(const tvm::ffi::TensorView input, + const tvm::ffi::TensorView out, + const double beta, + const double linear_beta, + const bool has_linear_beta) { + using namespace host; + + auto N = SymbolicSize{"num_tokens"}; + auto D_in = SymbolicSize{"input_width"}; + auto D_out = SymbolicSize{"output_width"}; + auto device_ = SymbolicDevice{}; + device_.set_options(); + + TensorMatcher({N, D_out}) // + .with_dtype() + .with_device(device_) + .verify(out); + TensorMatcher({N, D_in}) // + .with_dtype() + .with_device(device_) + .with_strides({-1, 1}) + .verify(input); + + const auto hidden_size = static_cast(D_out.unwrap()); + const auto num_tokens = static_cast(N.unwrap()); + const auto device = device_.unwrap(); + + if (num_tokens == 0) return; + RuntimeCheck(hidden_size * 2 == D_in.unwrap(), "invalid activation dimension: D_out * 2 != D_in"); + RuntimeCheck(hidden_size % kVecSize == 0, "hidden size must be divisible by vector size"); + RuntimeCheck(input.stride(0) % kVecSize == 0, "input row stride must be divisible by vector size"); + + const auto num_total_items = num_tokens * (hidden_size / kVecSize); + RuntimeCheck(num_total_items <= std::numeric_limits::max(), "too many items for 32-bit indexing"); + + const auto num_blocks = div_ceil(static_cast(num_total_items), kBlockSize); + const float beta_f = static_cast(beta); + const float linear_beta_f = static_cast(linear_beta); + + const auto params = SituAndMulParams{ + .input = input.data_ptr(), + .out = out.data_ptr(), + .beta = beta_f, + .inv_beta = 1.0f / beta_f, + .linear_beta = linear_beta_f, + .inv_linear_beta = linear_beta_f != 0.0f ? 1.0f / linear_beta_f : 0.0f, + .hidden_dim = hidden_size, + .num_tokens = num_tokens, + .stride_in_vecs = static_cast(input.stride(0) / kVecSize), + }; + + if (has_linear_beta) { + LaunchKernel(num_blocks, kBlockSize, device).enable_pdl(kUsePDL)(situ_and_mul_kernel, params); + } else { + LaunchKernel(num_blocks, kBlockSize, device).enable_pdl(kUsePDL)(situ_and_mul_kernel, params); + } + } +}; + +// --------------------------------------------------------------------------- +// varlen masked variant with the grouped-quant epilogue. Same activation, a +// different kernel: __launch_bounds__(1024, 2) plus a per-group scale writeback. +// --------------------------------------------------------------------------- +using deepseek_v4::fp8::cast_to_ue8m0; +using deepseek_v4::fp8::pack_fp8; + +struct SituMulQuantVarlenParams { + const bf16_t* __restrict__ input; + fp8_e4m3_t* __restrict__ output; + float* __restrict__ output_scale; + const int32_t* __restrict__ masked_m; + float beta; // gate softcap (e.g. 4.0) + float linear_beta; // up softcap (e.g. 25.0) + int64_t hidden_dim; + uint32_t num_tokens; + uint32_t num_experts; +}; + +constexpr uint32_t kMaxExperts = 256; + +struct alignas(16) CTAWork { + uint32_t expert_id; + uint32_t expert_token_id; + bool valid; +}; + +// SiTU (SoftCap-GLU) activation: +// gate_out = beta * tanh(gate / beta) * sigmoid(gate) +// up_out = linear_beta * tanh(up / linear_beta) +// output = gate_out * up_out +// Unlike SiLU, no external swiglu_limit clamp is needed: the tanh softcap +// inherently bounds the output to |beta * linear_beta| (< FP8_E4M3_MAX). +template +SGL_DEVICE fp32x2_t +situ_and_mul(DType2 gate, DType2 up, float beta, float inv_beta, float linear_beta, float inv_linear_beta) { + using namespace device; + const auto [g0, g1] = cast(gate); + const auto [u0, u1] = cast(up); + // kHasLinearBeta=true: this path always softcaps the up operand, as before. + const float val0 = sglang::kimi_k3::situ_activate(g0, u0, beta, inv_beta, linear_beta, inv_linear_beta); + const float val1 = sglang::kimi_k3::situ_activate(g1, u1, beta, inv_beta, linear_beta, inv_linear_beta); + if constexpr (kPrecise) { + return {val0, val1}; + } else { + return cast(cast(fp32x2_t{val0, val1})); + } +} + +[[maybe_unused]] +SGL_DEVICE CTAWork get_work(const SituMulQuantVarlenParams& params) { + // Preconditions: + // 1. blockDim.x >= params.num_experts + // 2. params.num_experts <= kMaxExperts + using namespace device; + static_assert(kWarpThreads == 32); + + static __shared__ uint32_t s_warp_sum[32]; + static __shared__ CTAWork result; + + result.valid = false; + + const uint32_t tx = threadIdx.x; + const uint32_t lane_id = tx % kWarpThreads; + const uint32_t warp_id = tx / kWarpThreads; + + const uint32_t val = tx < params.num_experts ? params.masked_m[tx] : 0u; + + // Per-warp inclusive scan of masked_m. + const uint32_t warp_inclusive = device::warp::inclusive_sum(lane_id, val); + const uint32_t warp_exclusive = warp_inclusive - val; + + // Write each warp total. + if (lane_id == kWarpThreads - 1) s_warp_sum[warp_id] = warp_inclusive; + __syncthreads(); + const auto tmp_val = lane_id < warp_id ? s_warp_sum[lane_id] : 0u; + const auto prefix_exclusive = warp::reduce_sum(tmp_val) + warp_exclusive; + const auto bx = blockIdx.x; + if (prefix_exclusive <= bx && bx < prefix_exclusive + val) { + result = {tx, bx - prefix_exclusive, true}; + } + __syncthreads(); + return result; +} + +template +__global__ __launch_bounds__(1024, 2) void // maximize occupancy + situ_mul_quant_varlen_kernel(const SituMulQuantVarlenParams __grid_constant__ params) { + using namespace device; + + constexpr uint32_t kGroupSize = 128u; + constexpr uint32_t kWorkThreads = 16u; + // each thread will handle 8 elements + using InputVec = AlignedVector; + using OutputVec = AlignedVector; + static_assert(8 * kWorkThreads == 128, "Invalid tiling"); + static_assert(!(kTransposed && !kScaleUE8M0), "transposed layout only supports ue8m0"); + + const auto [expert_id, token_id, valid] = get_work(params); + + if (!valid) return; + + const auto work_id = threadIdx.x / kWorkThreads; + + const auto offset = expert_id * params.num_tokens + token_id; + const auto input = params.input + offset * params.hidden_dim * 2; + const auto output = params.output + offset * params.hidden_dim; + [[maybe_unused]] + const auto output_scale = [&] { + const auto num_groups = params.hidden_dim / kGroupSize; + if constexpr (kTransposed) { + const auto base = reinterpret_cast(params.output_scale); + // Physical layout is [E, G//4, N] int32. Each int32 packs 4 consecutive + // group scales for the same token, so the byte address is: + // expert_offset + (group/4)*N*4 + token*4 + group%4 + return base + expert_id * num_groups * params.num_tokens + (work_id / 4u) * (params.num_tokens * 4u) + + token_id * 4u + (work_id % 4u); + } else { + return params.output_scale + offset * num_groups + work_id; + } + }(); + + const float beta = params.beta; + const float linear_beta = params.linear_beta; + const float inv_beta = 1.0f / beta; + const float inv_linear_beta = 1.0f / linear_beta; + + PDLWaitPrimary(); + + InputVec gate_vec, up_vec; + if constexpr (kSwizzle) { + // gran=8 interleaved: every 16-element chunk on the N axis is + // [gate[0..7], up[0..7]]. Each thread handles 8 consecutive output + // elements, so its gate chunk lives at vec index 2*threadIdx.x and its + // up chunk at 2*threadIdx.x+1. + gate_vec.load(input, threadIdx.x * 2); + up_vec.load(input, threadIdx.x * 2 + 1); + } else { + gate_vec.load(input, threadIdx.x); + up_vec.load(input, threadIdx.x + blockDim.x); + } + + float local_max = 0.0f; + float results[8]; + +#pragma unroll + for (uint32_t i = 0; i < 4; ++i) { + const auto [x, y] = situ_and_mul(gate_vec[i], up_vec[i], beta, inv_beta, linear_beta, inv_linear_beta); + results[2 * i + 0] = x; + results[2 * i + 1] = y; + local_max = fmaxf(local_max, fmaxf(fabsf(x), fabsf(y))); + } + + local_max = warp::reduce_max(local_max); + + const float absmax = fmaxf(local_max, 1e-10f); + float scale; + uint32_t ue8m0_exp; + + if constexpr (kScaleUE8M0) { + const float raw_scale = absmax / math::FP8_E4M3_MAX; + ue8m0_exp = cast_to_ue8m0(raw_scale); + scale = __uint_as_float(ue8m0_exp << 23); + } else { + scale = absmax / math::FP8_E4M3_MAX; + } + const auto inv_scale = 1.0f / scale; + + OutputVec out_vec; +#pragma unroll + for (uint32_t i = 0; i < 4; ++i) { + const float scaled_val0 = results[2 * i + 0] * inv_scale; + const float scaled_val1 = results[2 * i + 1] * inv_scale; + out_vec[i] = pack_fp8(scaled_val0, scaled_val1); + } + + PDLTriggerSecondary(); + + out_vec.store(output, threadIdx.x); + if constexpr (kTransposed) { + *output_scale = ue8m0_exp; + } else { + *output_scale = scale; + } +} + +// ---- Host wrapper + +template +struct SituAndMulMaskedPostQuantKernel { + static_assert(kGroupSize == 128); + static constexpr auto kernel_normal = situ_mul_quant_varlen_kernel; + static constexpr auto kernel_transposed = situ_mul_quant_varlen_kernel; + + static void + run(const tvm::ffi::TensorView input, + const tvm::ffi::TensorView output, + const tvm::ffi::TensorView output_scale, + const tvm::ffi::TensorView masked_m, + const uint32_t topk, + const bool transposed, + const double beta, + const double linear_beta) { + using namespace host; + + auto device = SymbolicDevice{}; + auto E = SymbolicSize{"num_experts"}; + auto T = SymbolicSize{"num_tokens_padded"}; + auto D = SymbolicSize{"hidden_dim x 2"}; + auto N = SymbolicSize{"hidden_dim"}; + auto G = SymbolicSize{"num_groups"}; + device.set_options(); + + TensorMatcher({E, T, D}) // input + .with_dtype() + .with_device(device) + .verify(input); + TensorMatcher({E, T, N}) // output + .with_dtype() + .with_device(device) + .verify(output); + if (!transposed) { + TensorMatcher({E, T, G}) // + .with_dtype() + .with_device(device) + .verify(output_scale); + } else { + RuntimeCheck(kScaleUE8M0, "transposed layout only supports scale_ue8m0=true"); + auto G_ = SymbolicSize{"G // 4"}; + TensorMatcher({E, G_, T}) // + .with_dtype() + .with_device(device) + .verify(output_scale); + G.set_value(G_.unwrap() * 4); + } + TensorMatcher({E}) // + .with_dtype() + .with_device(device) + .verify(masked_m); + + const auto num_experts = static_cast(E.unwrap()); + const auto num_tokens = static_cast(T.unwrap()); + const auto num_groups = static_cast(G.unwrap()); + const auto hidden_dim = N.unwrap(); + + RuntimeCheck(D.unwrap() == 2 * hidden_dim, "invalid dimension"); + RuntimeCheck(hidden_dim % kGroupSize == 0); + RuntimeCheck(num_experts <= kMaxExperts, "num_experts exceeds maximum (256)"); + RuntimeCheck(num_groups * kGroupSize == hidden_dim, "invalid num_groups"); + + const auto params = SituMulQuantVarlenParams{ + .input = static_cast(input.data_ptr()), + .output = static_cast(output.data_ptr()), + .output_scale = static_cast(output_scale.data_ptr()), + .masked_m = static_cast(masked_m.data_ptr()), + .beta = static_cast(beta), + .linear_beta = static_cast(linear_beta), + .hidden_dim = hidden_dim, + .num_tokens = num_tokens, + .num_experts = num_experts, + }; + + const auto num_threads = hidden_dim / 8; + RuntimeCheck(num_threads % device::kWarpThreads == 0); + RuntimeCheck(num_threads >= num_experts); + const auto kernel = transposed ? kernel_transposed : kernel_normal; + LaunchKernel(num_tokens * topk, num_threads, device.unwrap()) // + .enable_pdl(kUsePDL)(kernel, params); + } +}; + +} // namespace diff --git a/python/sglang/kernels/jit/csrc/moe/route_quant_fused.cuh b/python/sglang/kernels/jit/csrc/moe/route_quant_fused.cuh new file mode 100644 index 000000000..72b02792c --- /dev/null +++ b/python/sglang/kernels/jit/csrc/moe/route_quant_fused.cuh @@ -0,0 +1,142 @@ +// K3 MoE-front prep in one launch: radix routing (+ trtllm packed ids) on the +// first M CTAs, mxfp8 per-token-group quant of the routed activations on the +// next M. At decode batch sizes the unfused chain is three tiny back-to-back +// kernels (route 3.8us + quant 2.6us + pack 1.4us per layer) each leaving the +// SMs idle; fused, the quant CTAs run concurrently with the routing CTA and +// the pack is a 16-store epilogue. +// +// Both halves are the existing kernels verbatim: route_radix_block is the +// standalone route_radix body (same TU, same flags — no fast-math, which the +// routing bit-exactness contract requires and the quant math tolerates: its +// only transcendentals are explicit intrinsics and exact bit manipulation), +// and QuantTrait::run is the per_token_group_quant math. Specialized like +// route_radix itself: 896 experts, top-16, and a 3584-wide bf16 activation row +// (112 ue8m0 groups of 32 = 224 lanes, exactly the routing block width). + +#include "../gemm/per_token_group_quant.cuh" +#include "route_radix.cuh" + +namespace sglang { + +struct RouteQuantFusedParams { + RouteRadixParams route; + QuantKernelParams quant; +}; + +// One quant CTA covers one token row: thread pairs (2g, 2g+1) hold group g +// with lanes (0, 1) — the same subwarp layout the flat quant kernel derives +// from global_tid, so the group reduction and stores are bit-identical. +using RouteQuantTrait = QuantTrait< + bf16_t, + fp8_e4m3_t, + /*kGroupSize=*/32, + /*kUe8m0=*/true, + /*kRowMajor=*/true, + /*kAligned=*/true, + /*kFuseSiluAndMul=*/false>; + +inline constexpr uint32_t kQuantGroupsPerRow_ = LargeRouterRadixTrait::kBlockSize / RouteQuantTrait::kNumLanes; +inline constexpr uint32_t kQuantHidden_ = kQuantGroupsPerRow_ * RouteQuantTrait::kGroupSize; // 3584 + +template +__global__ __launch_bounds__(LargeRouterRadixTrait::kBlockSize) // + void route_quant_fused_kernel(const __grid_constant__ RouteQuantFusedParams params) { + const auto M = static_cast(params.route.M); + if (blockIdx.x < M) { + __shared__ typename LargeRouterRadixTrait::Smem smem; + route_radix_block(params.route, smem); + } else { + // Quant CTAs read the same primary-kernel output (the fused-front GEMM) + // as the routing CTAs, so they carry their own PDL wait/trigger. + device::PDLWaitPrimary(); + const uint32_t token_idx = blockIdx.x - M; + const uint32_t group_idx = threadIdx.x / RouteQuantTrait::kNumLanes; + const uint32_t lane_id = threadIdx.x % RouteQuantTrait::kNumLanes; + RouteQuantTrait::run(params.quant, /*expert_idx=*/0, token_idx, group_idx, lane_id); + device::PDLTriggerSecondary(); + } +} + +} // namespace sglang + +template +struct RouteQuantFusedKernel { + static void + run(const tvm::ffi::TensorView scores, + const tvm::ffi::TensorView bias, + const tvm::ffi::TensorView out_w, + const tvm::ffi::TensorView out_i, + const tvm::ffi::TensorView out_packed, + const tvm::ffi::TensorView x, + const tvm::ffi::TensorView out_q, + const tvm::ffi::TensorView out_s, + int64_t topk, + double routed_scaling_factor, + bool renormalize, + bool apply_scale) { + using namespace host; + using Trait = sglang::RouteQuantTrait; + + auto M_ = SymbolicSize{"num_tokens"}; + auto N_ = SymbolicSize{"num_experts"}; + auto K_ = SymbolicSize{"topk"}; + auto device = SymbolicDevice{}; + device.set_options(); + + auto score_dtype = SymbolicDType{}; + TensorMatcher({M_, N_}) + .with_dtype(score_dtype) + .with_device(device) + .with_strides({-1, 1}) + .verify(scores); + TensorMatcher({N_}).with_dtype().with_device(device).verify(bias); + TensorMatcher({M_, K_}).with_dtype().with_device(device).verify(out_w); + TensorMatcher({M_, K_}).with_dtype().with_device(device).verify(out_i); + TensorMatcher({M_, K_}).with_dtype().with_strides({-1, 1}).with_device(device).verify(out_packed); + + RuntimeCheck( + N_.unwrap() == sglang::kNumExperts_ && K_.unwrap() == sglang::kTopK_ && topk == sglang::kTopK_, + "route_quant_fused is specialized for N=896, K=16"); + RuntimeCheck(scores.stride(0) % 4 == 0, "route_quant_fused: scores row stride must be a multiple of 4"); + + // Quant half: shape/stride/alignment checks + byte-stride munging shared + // with the standalone flat kernel. + const auto ctx = build_quant_context(x, out_q, out_s); + RuntimeCheck( + ctx.params.hidden_size == sglang::kQuantHidden_, + "route_quant_fused is specialized for a 3584-wide activation row"); + RuntimeCheck( + ctx.params.num_tokens == static_cast(M_.unwrap()), + "route_quant_fused: scores and activations must have the same token count"); + + const auto M = static_cast(M_.unwrap()); + if (M == 0) return; + + const auto params = sglang::RouteQuantFusedParams{ + .route = + {scores.data_ptr(), + static_cast(bias.data_ptr()), + static_cast(out_w.data_ptr()), + static_cast(out_i.data_ptr()), + static_cast(out_packed.data_ptr()), + static_cast(M), + static_cast(scores.stride(0)), + static_cast(out_w.stride(0)), + static_cast(out_i.stride(0)), + static_cast(out_packed.stride(0)), + static_cast(routed_scaling_factor), + renormalize ? 1 : 0, + apply_scale ? 1 : 0, + /*sorted=*/0}, + .quant = ctx.params, + }; + + if (score_dtype.is_type()) { + LaunchKernel(2 * M, sglang::LargeRouterRadixTrait::kBlockSize, device.unwrap()) + .enable_pdl(kUsePDL)(sglang::route_quant_fused_kernel, params); + } else { + LaunchKernel(2 * M, sglang::LargeRouterRadixTrait::kBlockSize, device.unwrap()) + .enable_pdl(kUsePDL)(sglang::route_quant_fused_kernel, params); + } + } +}; diff --git a/python/sglang/kernels/jit/csrc/moe/route_radix.cuh b/python/sglang/kernels/jit/csrc/moe/route_radix.cuh new file mode 100644 index 000000000..1d9fc7f93 --- /dev/null +++ b/python/sglang/kernels/jit/csrc/moe/route_radix.cuh @@ -0,0 +1,753 @@ +// MoE routing by radix select: the standalone route kernel and the fused-gate +// front end, over one copy of the radix primitives (previously +// moe/radix_select_common.cuh, folded in once its two consumers became one file). + +#pragma once + +#include // For TensorMatcher, SymbolicSize, SymbolicDevice +#include // For RuntimeCheck, div_ceil + +#include // For dtype_trait, bf16_t, fp32_t, cast +#include // For LaunchKernel, SGL_DEVICE, PDL helpers +#include // For AlignedVector +#include // For warp::copy_bytes, elect_one_lane, inclusive_sum + +#include + +#include + +namespace sglang { + +namespace moe::radix { + +struct RadixSelectBase { + static constexpr uint32_t kRadixBits = 8; + static constexpr uint32_t kRadixSize = 1 << kRadixBits; + static constexpr uint32_t kRadixRounds = 32 / kRadixBits; + + struct alignas(16) MatchBin { + uint32_t bin; + uint32_t above_count; // active elements in bins strictly above `bin` + uint32_t equal_count; // active elements in bin `bin` + }; +}; + +inline constexpr float kNanFloor = -1e30f; + +// Monotone unsigned key: larger biased -> larger key. Caller must have floored +// biased-NaN. Canonicalizes -0.0 -> +0.0 so equal values get equal keys. +SGL_DEVICE uint32_t biased_to_key(float biased) { + if (biased == 0.0f) biased = 0.0f; + uint32_t u = __float_as_uint(biased); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} + +// tl.sigmoid(x) = 1/(1+exp(-x)). Must stay instruction-identical to v1's +// sigmoid_match so both kernels rank (and weight) identically. +SGL_DEVICE float sigmoid_match(float x) { + return __fdividef(1.0f, 1.0f + __expf(-x)); +} + +SGL_DEVICE float nan_floor(float x) { + return (x == x) ? x : kNanFloor; +} + +SGL_DEVICE void bar_sync(uint32_t id, uint32_t num_threads) { + asm volatile("bar.sync %0, %1;" ::"r"(id), "r"(num_threads) : "memory"); +} + +// Exclusive prefix (block-wide, thread-rank order) of `cnt`. Uses +// smem_warp_sum[kNumWarps]; syncs on entry (so the workspace can be reused +// across calls) and before the cross-warp read. +SGL_DEVICE uint32_t block_exclusive_sum(uint32_t cnt, uint32_t lane_id, uint32_t warp_id, uint32_t* smem_warp_sum) { + const uint32_t inc = device::warp::inclusive_sum(lane_id, cnt); + if (lane_id == 31) smem_warp_sum[warp_id] = inc; + __syncthreads(); + // TODO: replace `__reduce_add_sync` with `warp::reduce_sum` + const auto base = __reduce_add_sync(0xFFFFFFFF, lane_id < warp_id ? smem_warp_sum[lane_id] : 0u); + return base + inc - cnt; +} + +} // namespace moe::radix + +inline constexpr uint32_t kNumExperts_ = 896; +inline constexpr uint32_t kTopK_ = 16; + +struct LargeRouterRadixTrait : moe::radix::RadixSelectBase { + static constexpr uint32_t kNumExperts = kNumExperts_; + static constexpr uint32_t kTopK = kTopK_; + static constexpr uint32_t kVecSize = 4; + + static constexpr uint32_t kBlockSize = kNumExperts / kVecSize; // 224 = 7 warps + static constexpr uint32_t kNumWarps = kBlockSize / 32; + struct Smem { + uint32_t warp_sum[3][kNumWarps]; // cross-warp scan workspace + MatchBin match[kRadixRounds]; + uint32_t histogram[kRadixSize]; + // winner staging (compaction order = expert-id ascending) + int32_t wid[kTopK]; + uint32_t wkey[kTopK]; + fp32_t wact[kTopK]; + // sorted staging ((key desc, id asc) order), only used when sorted != 0 + int32_t sid[kTopK]; + fp32_t sact[kTopK]; + fp32_t norm; + }; +}; + +struct RouteRadixParams { + const void* __restrict__ scores; // bf16 or fp32, typed by the kernel template + const fp32_t* __restrict__ bias; + fp32_t* __restrict__ out_w; + int32_t* __restrict__ out_i; + // Optional trtllm-gen routed-MoE packing: (id << 16) | bf16(weight) bits, + // bit-identical to the standalone triton pack. nullptr skips the store. + int32_t* __restrict__ out_packed; + int M; + long long scores_stride; + long long out_w_stride; + long long out_i_stride; + long long out_packed_stride; + float routed_scaling_factor; + int renormalize; + int apply_scale; + int sorted; +}; + +// Whole-CTA routing body, callable from other kernels (the fused +// route+quant launch runs it on its first M CTAs). Routes row `blockIdx.x`; +// every thread of the 224-wide CTA must enter (block barriers inside). +template +SGL_DEVICE void route_radix_block(const RouteRadixParams& params, typename LargeRouterRadixTrait::Smem& smem) { + using namespace device; + using T = LargeRouterRadixTrait; + constexpr uint32_t kVecSize = T::kVecSize; + constexpr uint32_t kRadixLanes = T::kRadixSize / 2; // 128: 2 bins per thread + enum { BAR_RESERVED = 0, BAR_SUM = 1 }; + + const auto bx = blockIdx.x; + const auto tx = threadIdx.x; + const auto warp_id = tx / kWarpThreads; + const auto lane_id = tx % kWarpThreads; + // grid.x == M exactly; no row guard (an early return would deadlock the + // block-wide barriers below). + + // ---- Load + key transform: thread tx owns experts [4*tx, 4*tx+4) ---- + uint32_t keys[kVecSize]; + float act[kVecSize]; // raw sigmoid (weight source) — never NaN-sanitized + { + const auto scores = static_cast(params.scores) + bx * params.scores_stride; + AlignedVector bias_vec; + // bf16: 2x bf16x2 (8B row loads); fp32: 2x fp32x2 (16B row loads). The + // radix math below is fp32 either way — only the load width differs. + AlignedVector, kVecSize / 2> scores_vec; + + // prefetch bias (frozen weight) before the PDL wait + bias_vec.load(params.bias, tx); + PDLWaitPrimary(); + scores_vec.load(scores, tx); + +#pragma unroll + for (uint32_t i = 0; i < kVecSize / 2; ++i) { + fp32x2_t xy; + if constexpr (std::is_same_v) { + xy = scores_vec[i]; + } else { + xy = cast(scores_vec[i]); + } + const auto [x, y] = xy; + const auto sx = moe::radix::sigmoid_match(x), sy = moe::radix::sigmoid_match(y); + keys[2 * i + 0] = moe::radix::biased_to_key(moe::radix::nan_floor(sx + bias_vec[i].x)); + keys[2 * i + 1] = moe::radix::biased_to_key(moe::radix::nan_floor(sy + bias_vec[i].y)); + act[2 * i + 0] = sx; + act[2 * i + 1] = sy; + } + } + + // ---- Radix narrowing, MSB -> LSB ---- + // Invariants entering round r: + // active[i] <=> key's top 8r bits == threshold's top 8r bits + // total_active = size of the active set + // topk = winners still to take from the active set (1..total_active) + bool active[kVecSize]; +#pragma unroll + for (uint32_t i = 0; i < kVecSize; ++i) { + active[i] = true; + } + + uint32_t total_active = T::kNumExperts; + uint32_t topk = T::kTopK; + uint32_t threshold = 0; // assembled split-key prefix (unexamined low bits zero) + uint32_t examined_mask = 0; // bits of `threshold` that have been fixed + bool take_all_equals = false; + + { + AlignedVector zero; + zero.fill(0); + if (tx < kRadixLanes) zero.store(smem.histogram, tx); + +#pragma unroll + for (uint32_t round = 0; round < T::kRadixRounds; ++round) { + __syncthreads(); // histogram zeroed & previous match consumed + const uint32_t shift = 24 - round * 8; + uint32_t bin[kVecSize]; +#pragma unroll + for (uint32_t i = 0; i < kVecSize; ++i) { + bin[i] = (keys[i] >> shift) & 0xff; + } +#pragma unroll + for (uint32_t i = 0; i < kVecSize; ++i) { + if (active[i]) atomicAdd(&smem.histogram[bin[i]], 1); + } + __syncthreads(); + + // Split-bin search on 128 threads: thread t owns bins {2t, 2t+1}. + // The split bin b is the unique bin with above(b) < topk <= above(b) + hist[b]. + if (tx < kRadixLanes) { + AlignedVector hist; + hist.load(smem.histogram, tx); + const auto local_val = hist[0] + hist[1]; + const auto warp_inc = device::warp::inclusive_sum(lane_id, local_val); + if (lane_id == kWarpThreads - 1) smem.warp_sum[0][warp_id] = warp_inc; + moe::radix::bar_sync(BAR_SUM, kRadixLanes); + const auto inter = __reduce_add_sync(0xFFFFFFFF, lane_id < warp_id ? smem.warp_sum[0][lane_id] : 0u); + const auto prefix = inter + warp_inc; // active elements in bins [0, 2t+1] + const auto above_r = total_active - prefix; // in bins > 2t+1 + const auto above_m = above_r + hist[1]; // in bins > 2t + const auto above_l = above_m + hist[0]; // in bins >= 2t + if (above_r < topk && above_m >= topk) { + smem.match[round] = {tx * 2 + 1, above_r, hist[1]}; + } else if (above_m < topk && above_l >= topk) { + smem.match[round] = {tx * 2 + 0, above_m, hist[0]}; + } + } + __syncthreads(); + + const auto [threshold_bin, above_count, equal_count] = smem.match[round]; + threshold |= threshold_bin << shift; + examined_mask |= 0xffu << shift; +#pragma unroll + for (uint32_t i = 0; i < kVecSize; ++i) { + active[i] &= (bin[i] == threshold_bin); + } + total_active = equal_count; + topk -= above_count; // split condition guarantees 1 <= topk <= equal_count + if (topk == equal_count) { + // The remaining quota exactly covers the equal set: every active + // element wins, no deeper narrowing or tie-break needed. At the last + // round this is the no-full-key-tie case (the typical one). + take_all_equals = true; + break; + } + // Re-zero for the next round (synced by the loop-top barrier). Reaching + // round 3 with topk < equal_count means a full-key tie: resolved below + // by the smallest-id rank among `active`. + if (round + 1 < T::kRadixRounds && tx < kRadixLanes) zero.store(smem.histogram, tx); + } + } + + // ---- Epilogue: collect the K winners ---- + // Strict winners: examined bits compare above the split prefix (these were + // peeled off `active` in earlier rounds). Equal set (== `active`): take all + // (take_all_equals) or the `topk` smallest ids (full-key tie-break). + bool selected[kVecSize]; + if (take_all_equals) { +#pragma unroll + for (uint32_t i = 0; i < kVecSize; ++i) { + selected[i] = active[i] || (keys[i] & examined_mask) > threshold; + } + } else { // deterministic tie-break + uint32_t cnt = 0; +#pragma unroll + for (uint32_t i = 0; i < kVecSize; ++i) { + cnt += active[i] ? 1 : 0; + } + uint32_t rank = moe::radix::block_exclusive_sum(cnt, lane_id, warp_id, smem.warp_sum[1]); +#pragma unroll + for (uint32_t i = 0; i < kVecSize; ++i) { + const bool eq_win = active[i] && rank < topk; + if (active[i]) ++rank; + selected[i] = eq_win || (keys[i] & examined_mask) > threshold; + } + } + + // Compaction slots in expert-id order (deterministic). + uint32_t selected_cnt = 0; +#pragma unroll + for (uint32_t i = 0; i < kVecSize; ++i) { + selected_cnt += selected[i] ? 1 : 0; + } + uint32_t slot = moe::radix::block_exclusive_sum(selected_cnt, lane_id, warp_id, smem.warp_sum[2]); +#pragma unroll + for (uint32_t i = 0; i < kVecSize; ++i) { + if (selected[i] && slot < T::kTopK) { + smem.wid[slot] = (int32_t)(tx * kVecSize + i); + smem.wkey[slot] = keys[i]; + smem.wact[slot] = act[i]; + ++slot; + } + } + __syncthreads(); + + static_assert(T::kTopK <= kWarpThreads); + if (tx < T::kTopK) { + uint32_t rank = tx; + auto w = smem.wact[tx]; + const auto id = smem.wid[tx]; + if (params.sorted) { + const uint32_t ka = smem.wkey[tx]; + const int32_t ia = id; + rank = 0; +#pragma unroll + for (uint32_t b = 0; b < T::kTopK; ++b) { + if (smem.wkey[b] > ka || (smem.wkey[b] == ka && smem.wid[b] < ia)) ++rank; + } + } + PDLTriggerSecondary(); + float sum = 0.f; +#pragma unroll + for (uint32_t i = 0; i < T::kTopK; ++i) { + sum += smem.wact[i]; + } + const auto norm = (sum > 0.0f) ? sum : 1.0f; + if (params.renormalize) w = w / norm; + if (params.apply_scale) w = w * params.routed_scaling_factor; + params.out_w[bx * params.out_w_stride + rank] = w; + params.out_i[bx * params.out_i_stride + rank] = id; + if (params.out_packed != nullptr) { + // (id << 16) | bf16(w) bits — RN float->bf16 matches the triton pack. + const auto bits = static_cast(__bfloat16_as_ushort(__float2bfloat16_rn(w))); + params.out_packed[bx * params.out_packed_stride + rank] = + static_cast((static_cast(id) << 16) | bits); + } + } +} + +template +__global__ __launch_bounds__(LargeRouterRadixTrait::kBlockSize) // + void route_radix_kernel(const __grid_constant__ RouteRadixParams params) { + __shared__ typename LargeRouterRadixTrait::Smem smem; + route_radix_block(params, smem); +} + +// --------------------------------------------------------------------------- +// fused-gate front end. Same radix primitives above; a separate kernel because it +// folds the gate and the quant epilogue into one launch. Only the module that +// instantiates it pays for it -- an uninstantiated template costs parse time, not +// codegen, which is what lets both share this translation unit. +// --------------------------------------------------------------------------- +inline constexpr uint32_t kFGTNumExperts = 896; +inline constexpr uint32_t kFGTTopK = 16; +// 7 warps: 896 experts / 224 threads = 4 experts per thread in the epilogue, and +// one expert per warp per pass in phase 1. Keeping the block at route_radix's +// shape lets the epilogue reuse its radix-select verbatim. +/// Block size is a tunable: it sets how many experts each thread owns in the +/// radix select (kNumExperts / kBlockSize). It must be at least kRadixSize/2 = +/// 128 threads (the split-bin search puts 2 bins per thread) and must divide the +/// expert count into an even per-thread count (the loads are fp32x2 pairs), so +/// 224 (4 experts/thread) and 448 (2 experts/thread) are the legal choices for +/// 896 experts. +template +struct MoEFrontTrait : moe::radix::RadixSelectBase { + static constexpr uint32_t kNumExperts = kFGTNumExperts; + static constexpr uint32_t kTopK = kFGTTopK; + static constexpr uint32_t kBlockSize = kBlockSize_; + static constexpr uint32_t kVecSize = kNumExperts / kBlockSize; // experts per thread + static constexpr uint32_t kNumWarps = kBlockSize / 32; + + static_assert(kNumExperts % kBlockSize == 0, "block size must divide the expert count"); + static_assert(kVecSize % 2 == 0, "experts per thread must be even (fp32x2 loads)"); + static_assert(kBlockSize >= kRadixSize / 2, "block must cover the split-bin search lanes"); + + struct Smem { + uint32_t warp_sum[3][kNumWarps]; + MatchBin match[kRadixRounds]; + uint32_t histogram[kRadixSize]; + int32_t wid[kTopK]; + uint32_t wkey[kTopK]; + fp32_t wact[kTopK]; + }; +}; + +struct MoEFrontParams { + const fp32_t* __restrict__ bias; // [E] fp32 correction bias + const fp32_t* __restrict__ logits; // [M, logits_stride] fp32, gate slice first + fp32_t* __restrict__ out_w; // [M, topk] fp32 + int32_t* __restrict__ out_i; // [M, topk] int32 + int M; + int logits_stride; // E for the router-only entry, E + latent for the front + long long out_w_stride; + long long out_i_stride; + float routed_scaling_factor; + int renormalize; + int apply_scale; + // Merged-front entry only: the [M, latent] bf16 routed_input to emit. + bf16_t* __restrict__ routed_out; + int latent; + long long routed_stride; +}; + +/// Radix-select top-k over one token's fp32 logits. Lifted from +/// route_radix.cuh; the only change is the input dtype (fp32 in place of bf16). +template +SGL_DEVICE void fgt_select_topk( + const MoEFrontParams& params, typename T::Smem& smem, int m, uint32_t tx, uint32_t warp_id, uint32_t lane_id) { + constexpr uint32_t kVecSize = T::kVecSize; + constexpr uint32_t kRadixLanes = T::kRadixSize / 2; + enum { BAR_SUM = 1 }; + + uint32_t keys[kVecSize]; + float act[kVecSize]; + { + // 4 experts per thread as two fp32x2 (16B) loads, matching route_radix. + device::AlignedVector bias_vec; + bias_vec.load(params.bias, tx); + device::AlignedVector lv; + lv.load(params.logits + (long long)m * params.logits_stride, tx); + float logit[kVecSize]; +#pragma unroll + for (uint32_t i = 0; i < kVecSize / 2; ++i) { + logit[2 * i + 0] = lv[i].x; + logit[2 * i + 1] = lv[i].y; + } +#pragma unroll + for (uint32_t i = 0; i < kVecSize / 2; ++i) { + const float sx = moe::radix::sigmoid_match(logit[2 * i + 0]); + const float sy = moe::radix::sigmoid_match(logit[2 * i + 1]); + act[2 * i + 0] = sx; + act[2 * i + 1] = sy; + keys[2 * i + 0] = moe::radix::biased_to_key(moe::radix::nan_floor(sx + bias_vec[i].x)); + keys[2 * i + 1] = moe::radix::biased_to_key(moe::radix::nan_floor(sy + bias_vec[i].y)); + } + } + + bool active[kVecSize]; +#pragma unroll + for (uint32_t i = 0; i < kVecSize; ++i) { + active[i] = true; + } + + uint32_t total_active = T::kNumExperts; + uint32_t topk = T::kTopK; + uint32_t threshold = 0; + uint32_t examined_mask = 0; + bool take_all_equals = false; + + { + device::AlignedVector zero; + zero.fill(0); + if (tx < kRadixLanes) zero.store(smem.histogram, tx); + +#pragma unroll + for (uint32_t round = 0; round < T::kRadixRounds; ++round) { + __syncthreads(); + const uint32_t shift = 24 - round * 8; + uint32_t bin[kVecSize]; +#pragma unroll + for (uint32_t i = 0; i < kVecSize; ++i) { + bin[i] = (keys[i] >> shift) & 0xff; + } +#pragma unroll + for (uint32_t i = 0; i < kVecSize; ++i) { + if (active[i]) atomicAdd(&smem.histogram[bin[i]], 1); + } + __syncthreads(); + + if (tx < kRadixLanes) { + device::AlignedVector hist; + hist.load(smem.histogram, tx); + const auto local_val = hist[0] + hist[1]; + const auto warp_inc = device::warp::inclusive_sum(lane_id, local_val); + if (lane_id == 31) smem.warp_sum[0][warp_id] = warp_inc; + moe::radix::bar_sync(BAR_SUM, kRadixLanes); + const auto inter = __reduce_add_sync(0xFFFFFFFF, lane_id < warp_id ? smem.warp_sum[0][lane_id] : 0u); + const auto prefix = inter + warp_inc; + const auto above_r = total_active - prefix; + const auto above_m = above_r + hist[1]; + const auto above_l = above_m + hist[0]; + if (above_r < topk && above_m >= topk) { + smem.match[round] = {tx * 2 + 1, above_r, hist[1]}; + } else if (above_m < topk && above_l >= topk) { + smem.match[round] = {tx * 2 + 0, above_m, hist[0]}; + } + } + __syncthreads(); + + const auto [threshold_bin, above_count, equal_count] = smem.match[round]; + threshold |= threshold_bin << shift; + examined_mask |= 0xffu << shift; +#pragma unroll + for (uint32_t i = 0; i < kVecSize; ++i) { + active[i] &= (bin[i] == threshold_bin); + } + total_active = equal_count; + topk -= above_count; + if (topk == equal_count) { + take_all_equals = true; + break; + } + if (round + 1 < T::kRadixRounds && tx < kRadixLanes) zero.store(smem.histogram, tx); + } + } + + bool selected[kVecSize]; + if (take_all_equals) { +#pragma unroll + for (uint32_t i = 0; i < kVecSize; ++i) { + selected[i] = active[i] || (keys[i] & examined_mask) > threshold; + } + } else { + uint32_t cnt = 0; +#pragma unroll + for (uint32_t i = 0; i < kVecSize; ++i) { + cnt += active[i] ? 1 : 0; + } + uint32_t rank = moe::radix::block_exclusive_sum(cnt, lane_id, warp_id, smem.warp_sum[1]); +#pragma unroll + for (uint32_t i = 0; i < kVecSize; ++i) { + const bool eq_win = active[i] && rank < topk; + if (active[i]) ++rank; + selected[i] = eq_win || (keys[i] & examined_mask) > threshold; + } + } + + uint32_t selected_cnt = 0; +#pragma unroll + for (uint32_t i = 0; i < kVecSize; ++i) { + selected_cnt += selected[i] ? 1 : 0; + } + uint32_t slot = moe::radix::block_exclusive_sum(selected_cnt, lane_id, warp_id, smem.warp_sum[2]); +#pragma unroll + for (uint32_t i = 0; i < kVecSize; ++i) { + if (selected[i] && slot < T::kTopK) { + smem.wid[slot] = (int32_t)(tx * kVecSize + i); + smem.wact[slot] = act[i]; + ++slot; + } + } + __syncthreads(); + + static_assert(T::kTopK <= 32); + if (tx < T::kTopK) { + auto wv = smem.wact[tx]; + const auto id = smem.wid[tx]; + float sum = 0.f; +#pragma unroll + for (uint32_t i = 0; i < T::kTopK; ++i) { + sum += smem.wact[i]; + } + const auto norm = (sum > 0.0f) ? sum : 1.0f; + if (params.renormalize) wv = wv / norm; + if (params.apply_scale) wv = wv * params.routed_scaling_factor; + params.out_w[m * params.out_w_stride + tx] = wv; + params.out_i[m * params.out_i_stride + tx] = id; + } + __syncthreads(); // smem reuse across the token loop +} + +/// Tunables: `kBlockSize` sets the experts-per-thread of the radix select, +/// `kCastVec` the fp32 elements each thread converts per step (the cast moves +/// [T, 3584] fp32 in and bf16 out, which dominates the epilogue at large T), and +/// `kCastFirst` whether the cast is issued before the select (loads in flight +/// during the radix rounds) or after it. +template +__global__ __launch_bounds__(kBlockSize) // + void fused_front_epilogue_kernel(const __grid_constant__ MoEFrontParams params) { + using namespace device; + using T = MoEFrontTrait; + __shared__ typename T::Smem smem; + const uint32_t tx = threadIdx.x; + const int m = (int)blockIdx.x; + + PDLWaitPrimary(); + + // Cast the latent slice: [E, E + latent) fp32 -> [0, latent) bf16. + auto cast_latent = [&]() { + const fp32_t* src = params.logits + (long long)m * params.logits_stride + T::kNumExperts; + bf16_t* dst = params.routed_out + (long long)m * params.routed_stride; + for (int i = (int)tx * kCastVec; i < params.latent; i += (int)kBlockSize * kCastVec) { + AlignedVector v; + v.load(src, i / kCastVec); + AlignedVector o; +#pragma unroll + for (uint32_t j = 0; j < kCastVec / 2; ++j) { + o[2 * j + 0] = cast(v[j].x); + o[2 * j + 1] = cast(v[j].y); + } + o.store(dst, i / kCastVec); + } + }; + + if (kCastFirst) cast_latent(); + fgt_select_topk(params, smem, m, tx, tx / 32, tx % 32); + if (!kCastFirst) cast_latent(); +} + +} // namespace sglang + +template +struct RouteRadixKernel { + static void + run(const tvm::ffi::TensorView scores, + const tvm::ffi::TensorView bias, + const tvm::ffi::TensorView out_w, + const tvm::ffi::TensorView out_i, + int64_t topk, + double routed_scaling_factor, + bool renormalize, + bool apply_scale, + bool sorted) { + using namespace host; + + auto M_ = SymbolicSize{"num_tokens"}; + auto N_ = SymbolicSize{"num_experts"}; + auto K_ = SymbolicSize{"topk"}; + auto device = SymbolicDevice{}; + device.set_options(); + + auto score_dtype = SymbolicDType{}; + TensorMatcher({M_, N_}) + .with_dtype(score_dtype) + .with_device(device) + .with_strides({-1, 1}) + .verify(scores); + TensorMatcher({N_}).with_dtype().with_device(device).verify(bias); + TensorMatcher({M_, K_}).with_dtype().with_device(device).verify(out_w); + TensorMatcher({M_, K_}).with_dtype().with_device(device).verify(out_i); + + RuntimeCheck( + N_.unwrap() == sglang::kNumExperts_ && K_.unwrap() == sglang::kTopK_ && topk == sglang::kTopK_, + "route_radix is specialized for N=896, K=16"); + // Vectorized row loads (8B for bf16, 16B for fp32) need aligned row + // starts; stride % 4 elements covers both (4 x 2B = 8B / 4 x 4B = 16B). + RuntimeCheck(scores.stride(0) % 4 == 0, "route_radix: scores row stride must be a multiple of 4"); + + const auto M = static_cast(M_.unwrap()); + if (M == 0) return; + + const auto params = sglang::RouteRadixParams{ + scores.data_ptr(), + static_cast(bias.data_ptr()), + static_cast(out_w.data_ptr()), + static_cast(out_i.data_ptr()), + /*out_packed=*/nullptr, + static_cast(M), + static_cast(scores.stride(0)), + static_cast(out_w.stride(0)), + static_cast(out_i.stride(0)), + /*out_packed_stride=*/0, + static_cast(routed_scaling_factor), + renormalize ? 1 : 0, + apply_scale ? 1 : 0, + sorted ? 1 : 0}; + + if (score_dtype.is_type()) { + LaunchKernel(M, sglang::LargeRouterRadixTrait::kBlockSize, device.unwrap()) + .enable_pdl(kUsePDL)(sglang::route_radix_kernel, params); + } else { + LaunchKernel(M, sglang::LargeRouterRadixTrait::kBlockSize, device.unwrap()) + .enable_pdl(kUsePDL)(sglang::route_radix_kernel, params); + } + } +}; + +template +struct FusedFrontEpilogueKernel { + static void + run(const tvm::ffi::TensorView merged, // [M, E + latent] fp32, row-dense + const tvm::ffi::TensorView bias, // [E] fp32 + const tvm::ffi::TensorView out_w, // [M, topk] fp32 + const tvm::ffi::TensorView out_i, // [M, topk] int32 + const tvm::ffi::TensorView routed, // [M, latent] bf16 + int64_t topk, + double routed_scaling_factor, + bool renormalize, + bool apply_scale, + int64_t block_size, + int64_t cast_vec, + bool cast_first) { + using namespace host; + + auto M_ = SymbolicSize{"num_tokens"}; + auto W_ = SymbolicSize{"merged_width"}; + auto L_ = SymbolicSize{"latent"}; + auto E_ = SymbolicSize{"num_experts"}; + auto K_ = SymbolicSize{"topk"}; + auto device = SymbolicDevice{}; + device.set_options(); + + TensorMatcher({M_, W_}).with_dtype().with_device(device).with_strides({-1, 1}).verify(merged); + TensorMatcher({E_}).with_dtype().with_device(device).verify(bias); + TensorMatcher({M_, K_}).with_dtype().with_device(device).verify(out_w); + TensorMatcher({M_, K_}).with_dtype().with_device(device).verify(out_i); + TensorMatcher({M_, L_}).with_dtype().with_device(device).with_strides({-1, 1}).verify(routed); + + const auto M = static_cast(M_.unwrap()); + const auto latent = static_cast(L_.unwrap()); + RuntimeCheck( + E_.unwrap() == sglang::kFGTNumExperts && K_.unwrap() == sglang::kFGTTopK && topk == sglang::kFGTTopK, + "fused_front_epilogue is specialized for E=896, topk=16"); + RuntimeCheck( + static_cast(W_.unwrap()) == static_cast(sglang::kFGTNumExperts) + latent, + "fused_front_epilogue: merged width must be num_experts + latent"); + // 16B vectorized reads of the fp32 rows and 8B writes of the bf16 rows. + RuntimeCheck(latent % 4 == 0, "fused_front_epilogue: latent must be a multiple of 4"); + RuntimeCheck( + merged.stride(0) % 4 == 0 && routed.stride(0) % 4 == 0, + "fused_front_epilogue: row strides must be a multiple of 4"); + if (M == 0) return; + + auto params = sglang::MoEFrontParams{}; + params.bias = static_cast(bias.data_ptr()); + params.logits = static_cast(merged.data_ptr()); + params.out_w = static_cast(out_w.data_ptr()); + params.out_i = static_cast(out_i.data_ptr()); + params.M = M; + params.out_w_stride = static_cast(out_w.stride(0)); + params.out_i_stride = static_cast(out_i.stride(0)); + params.routed_scaling_factor = static_cast(routed_scaling_factor); + params.renormalize = renormalize ? 1 : 0; + params.apply_scale = apply_scale ? 1 : 0; + params.logits_stride = static_cast(merged.stride(0)); + params.routed_out = static_cast(routed.data_ptr()); + params.latent = latent; + params.routed_stride = static_cast(routed.stride(0)); + + // Tunables come from the JSON config table; see kernels/ops/moe/moe_front.py. + // cast_vec * 4 bytes per thread must stay inside the 32B vector-load limit. + RuntimeCheck(cast_vec == 2 || cast_vec == 4 || cast_vec == 8, "fused_front_epilogue: cast_vec must be 2, 4 or 8"); + RuntimeCheck(latent % cast_vec == 0, "fused_front_epilogue: cast_vec must divide latent"); + +#define SGL_FRONT_LAUNCH(BS, CV, CF) \ + LaunchKernel(M, BS, device.unwrap()) \ + .enable_pdl(kUsePDL)(sglang::fused_front_epilogue_kernel, params) +#define SGL_FRONT_DISPATCH_CV(BS, CF) \ + do { \ + if (cast_vec == 2) { \ + SGL_FRONT_LAUNCH(BS, 2, CF); \ + } else if (cast_vec == 4) { \ + SGL_FRONT_LAUNCH(BS, 4, CF); \ + } else { \ + SGL_FRONT_LAUNCH(BS, 8, CF); \ + } \ + } while (0) +#define SGL_FRONT_DISPATCH_BS(CF) \ + do { \ + if (block_size == 448) { \ + SGL_FRONT_DISPATCH_CV(448, CF); \ + } else { \ + SGL_FRONT_DISPATCH_CV(224, CF); \ + } \ + } while (0) + + RuntimeCheck(block_size == 224 || block_size == 448, "fused_front_epilogue: block_size must be 224 or 448"); + if (cast_first) { + SGL_FRONT_DISPATCH_BS(true); + } else { + SGL_FRONT_DISPATCH_BS(false); + } +#undef SGL_FRONT_DISPATCH_BS +#undef SGL_FRONT_DISPATCH_CV +#undef SGL_FRONT_LAUNCH + } +}; diff --git a/python/sglang/kernels/jit/csrc/moe/topk_sum.cuh b/python/sglang/kernels/jit/csrc/moe/topk_sum.cuh new file mode 100644 index 000000000..c8661e48c --- /dev/null +++ b/python/sglang/kernels/jit/csrc/moe/topk_sum.cuh @@ -0,0 +1,103 @@ +// Top-k expert-output sum: out[M, K] = sum_j in[M, topk, K]. +// +// Replaces sgl_kernel's moe_sum_reduce_kernel_general (~5.7us at decode +// shapes [1, 16, 3584]) with a straightforward vectorized pass (~1.5us): +// one thread per 8-element vector of K, looping the topk rows in fp32. + +#include // For TensorMatcher, SymbolicSize, SymbolicDevice +#include // For RuntimeCheck + +#include // For bf16_t, fp32_t, device::cast +#include // For LaunchKernel +#include // For AlignedVector + +#include + +#include + +namespace { + +struct TopkSumParams { + const bf16_t* __restrict__ in; // [M, topk, K] contiguous + bf16_t* __restrict__ out; // [M, K] contiguous + uint32_t K; + uint32_t topk; +}; + +template +__global__ void topk_sum_kernel(const TopkSumParams __grid_constant__ params) { + using namespace device; + + constexpr int kVecN = 8; // 8 bf16 = 128 bits + using vec_bf16_t = AlignedVector; + + const uint32_t m = blockIdx.y; + const uint32_t v = blockIdx.x * kThreads + threadIdx.x; + const uint32_t n_vecs = params.K / kVecN; + if (v >= n_vecs) return; + + const bf16_t* base = params.in + static_cast(m) * params.topk * params.K; + + PDLWaitPrimary(); + + float acc[kVecN]; +#pragma unroll + for (int i = 0; i < kVecN; ++i) { + acc[i] = 0.0f; + } + for (uint32_t j = 0; j < params.topk; ++j) { + vec_bf16_t x; + x.load(base + static_cast(j) * params.K, v); +#pragma unroll + for (int i = 0; i < kVecN; ++i) { + acc[i] += cast(x[i]); + } + } + + vec_bf16_t o; +#pragma unroll + for (int i = 0; i < kVecN; ++i) { + o[i] = cast(acc[i]); + } + o.store(params.out + static_cast(m) * params.K, v); + + PDLTriggerSecondary(); +} + +template +struct TopkSumKernel { + static constexpr auto kernel = topk_sum_kernel; + + static void run(const tvm::ffi::TensorView in, const tvm::ffi::TensorView out) { + using namespace host; + + auto M_ = SymbolicSize{"num_tokens"}; + auto T_ = SymbolicSize{"topk"}; + auto K_ = SymbolicSize{"hidden"}; + auto device = SymbolicDevice{}; + device.set_options(); + + TensorMatcher({M_, T_, K_}).with_dtype().with_device(device).verify(in); + TensorMatcher({M_, K_}).with_dtype().with_device(device).verify(out); + + const auto M = static_cast(M_.unwrap()); + const auto topk = static_cast(T_.unwrap()); + const auto K = static_cast(K_.unwrap()); + + RuntimeCheck(K % 8 == 0, "K must be divisible by 8 for vectorized loads"); + if (M == 0) return; + + const auto params = TopkSumParams{ + .in = static_cast(in.data_ptr()), + .out = static_cast(out.data_ptr()), + .K = K, + .topk = topk, + }; + + const uint32_t n_vecs = K / 8; + dim3 grid((n_vecs + kThreads - 1) / kThreads, M); + LaunchKernel(grid, kThreads, device.unwrap()).enable_pdl(kUsePDL)(kernel, params); + } +}; + +} // namespace diff --git a/python/sglang/kernels/jit/include/sgl_kernel/math.cuh b/python/sglang/kernels/jit/include/sgl_kernel/math.cuh index 92c4dad6c..d7aec1d15 100644 --- a/python/sglang/kernels/jit/include/sgl_kernel/math.cuh +++ b/python/sglang/kernels/jit/include/sgl_kernel/math.cuh @@ -58,6 +58,24 @@ SGL_DEVICE T exp(T a) { return DTypeTrait::exp(a); } +/// \brief Fast approximate sigmoid for FP32 device code. +SGL_DEVICE float sigmoid_fast(float x) { + return 1.0f / (1.0f + __expf(-x)); +} + +/// \brief Fast approximate SiLU for FP32 device code. +SGL_DEVICE float silu_fast(float x) { + return x * sigmoid_fast(x); +} + +/// \brief Fast approximate softplus for FP32 device code. +/// +/// Values above 20 use the asymptotic result directly, avoiding overflow and +/// an unnecessary exponential while matching common softplus kernels. +SGL_DEVICE float softplus_fast(float x) { + return x > 20.0f ? x : log1pf(__expf(x)); +} + /// \brief Returns sin(a). template SGL_DEVICE T sin(T a) { @@ -70,4 +88,20 @@ SGL_DEVICE T cos(T a) { return DTypeTrait::cos(a); } +// bf16 x bf16 -> fp32 fused multiply-add The mixed-precision PTX +// instruction saves the explicit converts; the fallback is bit-identical (the +// bf16 -> f32 conversion is exact, both round once). Shared by tiny_gemm, +// gemm_ag and ar_fusion. +SGL_DEVICE float fma_f32_bf16(bf16_t a, bf16_t b, float acc) { +#if SGL_ARCH_BLACKWELL_OR_GREATER + const uint16_t a_bits = __bfloat16_as_ushort(a); + const uint16_t b_bits = __bfloat16_as_ushort(b); + float result; + asm("fma.rn.f32.bf16 %0, %1, %2, %3;" : "=f"(result) : "h"(a_bits), "h"(b_bits), "f"(acc)); + return result; +#else + return fmaf(cast(a), cast(b), acc); +#endif +} + } // namespace device::math diff --git a/python/sglang/kernels/jit/include/sgl_kernel/mbarrier.cuh b/python/sglang/kernels/jit/include/sgl_kernel/mbarrier.cuh new file mode 100644 index 000000000..29a5c4ad6 --- /dev/null +++ b/python/sglang/kernels/jit/include/sgl_kernel/mbarrier.cuh @@ -0,0 +1,68 @@ +#pragma once + +// mbarrier PTX wrappers shared by the Kimi K3 kernels that drive TMA by hand: +// kimi_k3/comm/gemm_ar.cuh and kimi_k3/attn_res/fused_tma.cuh both defined these +// with identical bodies. +// +// The enclosing namespace is the same global `ptx` both files already open, so +// existing `::ptx::mbar_*` call sites need no change. +// +// gemm_ar.cuh keeps mbar_arrive_cluster_release: only it uses that one. +// attention/kda_prefill.cu duplicates a different set (MMA / ldmatrix) but is +// built without a sglang include path and cannot consume this header. + +#include + +#include + +namespace ptx { + +// Inline-PTX `.shared` instructions take a 32-bit byte offset in the shared +// window, not a generic 64-bit pointer. +template +static SGL_DEVICE uint32_t to_shared(T* ptr) { + return static_cast(__cvta_generic_to_shared(ptr)); +} + +// ---- mbarrier (PTX ISA §9.7.13.15) ----------------------------------------- +// +// Only the `try_wait.parity` waiter is wrapped: state-token waits couple +// arriver and waiter, and mixing state- and parity-tracked codepaths in one +// kernel is a deadlock risk. The caller owns the phase counter and flips it at +// the stage wrap (`phase ^= (stage == 0)`). +// +// Initial parity, the easy-to-flip part: after `mbar_init` the bar is at +// parity 0, and each full cycle (count arrivals -> fire -> reset) flips it. +// - consumer-first (waits for an external producer's first signal) -> 0. +// - producer-first (waits for a consumer to release a slot that no consumer +// has touched yet) -> 1, so the first wait is a no-op skip. +// A consumer-first wait initialized to 1 skips the producer's first signal and +// blocks forever on the second. +static SGL_DEVICE void mbar_init(uint64_t* bar, uint32_t count) { + asm volatile("mbarrier.init.shared.b64 [%0], %1;" ::"r"(to_shared(bar)), "r"(count)); +} + +static SGL_DEVICE uint64_t mbar_arrive(uint64_t* bar) { + uint64_t state; + asm volatile("mbarrier.arrive.shared.b64 %0, [%1];" : "=l"(state) : "r"(to_shared(bar))); + return state; +} + +// Combined arrive + set tx-count, for TMA-load completion. +static SGL_DEVICE void mbar_arrive_expect_tx(uint64_t* bar, uint32_t bytes) { + asm volatile("mbarrier.arrive.expect_tx.shared.b64 _, [%0], %1;" ::"r"(to_shared(bar)), "r"(bytes)); +} + +// Wait for phase `parity` to complete. Looped because the spec allows spurious +// early wakeups. Default `.acquire` semantics mean prior `cp.async.bulk` writes +// tracked by this mbarrier are visible to later generic-proxy reads on this +// thread with no `fence.proxy.async` (spec §9.7.13.15.16 point 3). +static SGL_DEVICE void mbar_wait_parity(uint64_t* bar, uint32_t parity) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "WAIT_%=: mbarrier.try_wait.parity.shared.b64 p, [%0], %1;\n\t" + "@!p bra WAIT_%=;\n\t}\n" ::"r"(to_shared(bar)), + "r"(parity)); +} + +} // namespace ptx diff --git a/python/sglang/kernels/jit/include/sgl_kernel/warp.cuh b/python/sglang/kernels/jit/include/sgl_kernel/warp.cuh index 9fafaf0d8..6631a8cc4 100644 --- a/python/sglang/kernels/jit/include/sgl_kernel/warp.cuh +++ b/python/sglang/kernels/jit/include/sgl_kernel/warp.cuh @@ -1,9 +1,11 @@ /// \file warp.cuh -/// \brief Warp-level reduction primitives. +/// \brief Warp-level reduction and cooperative-copy primitives. #pragma once #include +#include #include +#include #include #include @@ -119,4 +121,87 @@ SGL_DEVICE T reduce_min(T value, mask_t active_mask = kFullMask) { return reduce(value, active_mask); } +/// \brief Warp-cooperative gmem -> smem copy of a compile-time byte count. +/// +/// Picks the widest vector width that divides both the per-thread share and +/// the byte total. The caller guarantees ``src`` is aligned to the picked +/// width (16B for kBytes % (16*32) == 0, else 8/4) and ``dst`` is the start +/// of a 16B-aligned per-warp smem slot. +// Warp-cooperative byte copy between any two address spaces, vectorised to the +// widest unit `kBytes` allows. Named for what it does rather than where it is +// used: the MLA call sites happen to target shared memory, but nothing here is +// global->shared specific -- no cp.async, no TMA, the payload moves through +// registers. +// +// The strategy was measured against the two async alternatives on B300 (sm_103, +// 148 SMs), copying one MLA row per warp out of a 512 MB pool so every row +// streams from HBM (grid 296, 64 rows/warp, 50 launches): +// +// 1152 B/warp (bf16, nope 1024 + rope 128) 576 B/warp (fp8, 512 + 64) +// this (generic) 47.3 us 3.69 TB/s 40.8 us 2.14 TB/s +// cp.async (ldgsts) 73.9 us 2.36 TB/s 56.4 us 1.55 TB/s +// cp.async.bulk/TMA 50.0 us 3.50 TB/s 43.2 us 2.02 TB/s +// +// The generic path wins at both sizes: a ~1 KB row is too small to amortise +// cp.async's per-lane 16 B issues or TMA's fixed issue plus mbarrier round trip. +// Revisit if a call site ever copies substantially more than one row per warp. +template +SGL_DEVICE void copy_bytes(const void* __restrict__ src, void* __restrict__ dst) { + constexpr int64_t kAlignment = (kBytes % (16 * kWarpThreads) == 0) ? 16 + : (kBytes % (8 * kWarpThreads) == 0) ? 8 + : (kBytes % (4 * kWarpThreads) == 0) ? 4 + : (kBytes % 4 == 0) ? 4 + : 0; + static_assert(kAlignment > 0, "kBytes must be a multiple of 4"); + + using vec_t = AlignedStorage; + constexpr auto kLoopBytes = sizeof(vec_t) * kWarpThreads; + constexpr auto kLoopCount = kBytes / kLoopBytes; + constexpr int64_t kTailVecs = (kBytes - kLoopCount * kLoopBytes) / sizeof(vec_t); + + const auto gmem = tile::Memory::warp(); + +#pragma unroll + for (int64_t i = 0; i < kLoopCount; ++i) { + const auto v = gmem.load(src, i); + gmem.store(dst, v, i); + } + if constexpr (kTailVecs > 0) { + if (gmem.in_bound(kLoopCount * kWarpThreads + kTailVecs, kLoopCount)) { + const auto v = gmem.load(src, kLoopCount); + gmem.store(dst, v, kLoopCount); + } + } +} + +/// Inclusive prefix sum across one warp, thread-rank order. Distinct from +/// reduce_sum above: every lane keeps its own running total rather than the +/// whole-warp result. +SGL_DEVICE uint32_t inclusive_sum(uint32_t lane_id, uint32_t val) { + static_assert(kWarpThreads == 32); +#pragma unroll + for (uint32_t offset = 1; offset < 32; offset *= 2) { + uint32_t n = __shfl_up_sync(0xFFFFFFFF, val, offset); + if (lane_id >= offset) val += n; + } + return val; +} + +// One elected lane, via elect.sync. Raw PTX rather than cute::elect_one_sync, +// which would drag the whole CuTe include path into elementwise JIT modules; +// cuda::ptx has no elect_sync in CUDA 13.0. Use this to gate a single-thread +// TMA issue instead of a lane-index predicate. +SGL_DEVICE bool elect_one_lane() { + uint32_t pred; + asm volatile( + "{\n" + " .reg .pred p;\n" + " .reg .b32 r;\n" + " elect.sync r|p, 0xFFFFFFFF;\n" + " selp.b32 %0, 1, 0, p;\n" + "}\n" + : "=r"(pred)); + return pred != 0; +} + } // namespace device::warp diff --git a/python/sglang/kernels/jit/utils/compile.py b/python/sglang/kernels/jit/utils/compile.py index e44821d25..b045e731c 100644 --- a/python/sglang/kernels/jit/utils/compile.py +++ b/python/sglang/kernels/jit/utils/compile.py @@ -173,6 +173,8 @@ def load_jit( *args: str, cpp_files: List[str] | None = None, cuda_files: List[str] | None = None, + external_cpp_files: List[str] | None = None, + external_cuda_files: List[str] | None = None, cpp_wrappers: List[Tuple[str, str]] | None = None, cuda_wrappers: List[Tuple[str, str]] | None = None, extra_cflags: List[str] | None = None, @@ -195,6 +197,12 @@ def load_jit( :type cpp_files: List[str] | None :param cuda_files: A list of CUDA source files. :type cuda_files: List[str] | None + :param external_cpp_files: A list of caller-resolved C++ source paths outside + the in-tree JIT source directory. + :type external_cpp_files: List[str] | None + :param external_cuda_files: A list of caller-resolved CUDA source paths outside + the in-tree JIT source directory. + :type external_cuda_files: List[str] | None :param cpp_wrappers: A list of C++ wrappers, defining the export name and kernel name. :type cpp_wrappers: List[Tuple[str, str]] | None :param cuda_wrappers: A list of CUDA wrappers, defining the export name and kernel name. @@ -222,13 +230,26 @@ def load_jit( cpp_files = cpp_files or [] cuda_files = cuda_files or [] + external_cpp_files = external_cpp_files or [] + external_cuda_files = external_cuda_files or [] extra_cflags = extra_cflags or [] extra_cuda_cflags = extra_cuda_cflags or [] extra_ldflags = extra_ldflags or [] extra_include_paths = extra_include_paths or [] - cpp_files = [str((KERNEL_PATH / "csrc" / f).resolve()) for f in cpp_files] - cuda_files = [str((KERNEL_PATH / "csrc" / f).resolve()) for f in cuda_files] + if torch.version.hip is not None: + extra_cuda_cflags = [ + flag + for flag in extra_cuda_cflags + if flag not in ("--use_fast_math", "-use_fast_math") + ] + + cpp_files = [str((KERNEL_PATH / "csrc" / f).resolve()) for f in cpp_files] + [ + str(pathlib.Path(f).resolve()) for f in external_cpp_files + ] + cuda_files = [str((KERNEL_PATH / "csrc" / f).resolve()) for f in cuda_files] + [ + str(pathlib.Path(f).resolve()) for f in external_cuda_files + ] for dep in set(extra_dependencies or []): if dep not in REGISTERED_DEPENDENCIES: diff --git a/python/sglang/kernels/ops/__init__.py b/python/sglang/kernels/ops/__init__.py index 9574dc847..06e61082e 100644 --- a/python/sglang/kernels/ops/__init__.py +++ b/python/sglang/kernels/ops/__init__.py @@ -23,10 +23,12 @@ _GROUPS = ( "embeddings", "gemm", "grammar", + "kimi_k3", "kvcache", "layernorm", "mamba", "memory", + "mm", "moe", "quantization", "sampling", diff --git a/python/sglang/kernels/ops/attention/concat_mla.py b/python/sglang/kernels/ops/attention/concat_mla.py index 3de6e9849..7c9f9f297 100644 --- a/python/sglang/kernels/ops/attention/concat_mla.py +++ b/python/sglang/kernels/ops/attention/concat_mla.py @@ -4,7 +4,12 @@ from typing import TYPE_CHECKING import torch -from sglang.kernels.jit.utils import cache_once, load_jit +from sglang.kernels.jit.utils import ( + cache_once, + is_arch_support_pdl, + load_jit, + make_cpp_args, +) if TYPE_CHECKING: from tvm_ffi.module import Module @@ -21,10 +26,12 @@ def _jit_concat_mla_k_module() -> Module: @cache_once def _jit_concat_mla_absorb_q_module() -> Module: + args = make_cpp_args(is_arch_support_pdl()) return load_jit( "concat_mla_absorb_q", + *args, cuda_files=["elementwise/concat_mla.cuh"], - cuda_wrappers=[("concat_mla_absorb_q", "ConcatMlaAbsorbQKernel::run")], + cuda_wrappers=[("concat_mla_absorb_q", f"ConcatMlaAbsorbQKernel<{args}>::run")], ) diff --git a/python/sglang/kernels/ops/attention/cutedsl_kda.py b/python/sglang/kernels/ops/attention/cutedsl_kda.py index 7732db9ef..78c7b8495 100644 --- a/python/sglang/kernels/ops/attention/cutedsl_kda.py +++ b/python/sglang/kernels/ops/attention/cutedsl_kda.py @@ -94,6 +94,10 @@ def _define_kernels(): pool_idx = h0_indices[i_n] if pool_idx >= 0: + # State indexing in int64: envelope-strided pools (unified memory / + # page-major) put pool_idx * stride(0) far past 2^31 (the CuTe twin + # of the fla/chunk_delta_h.py stride_init_state fix). + pool_idx64 = cutlass.Int64(pool_idx) k_local = in_warp_tid // V_PER_WARP_SMALL v_local = in_warp_tid % V_PER_WARP_SMALL v_base = warp_idx * V_PER_WARP_SMALL @@ -206,7 +210,7 @@ def _define_kernels(): h_val = 0.0 if v_global_load < v.shape[3]: h_val = cutlass.Float32( - h0_source[(pool_idx, i_hv, v_global_load, k_load)] + h0_source[(pool_idx64, i_hv, v_global_load, k_load)] ) sData[(k_load, v_load, stage)] = h_val @@ -263,7 +267,7 @@ def _define_kernels(): if k_write < TILE_K: v_global_write = v_tile * TILE_V_SMALL + v_write if v_global_write < v.shape[3]: - h0_source[(pool_idx, i_hv, v_global_write, k_write)] = ( + h0_source[(pool_idx64, i_hv, v_global_write, k_write)] = ( sData[(k_write, v_write, stage)] ) @@ -311,6 +315,10 @@ def _define_kernels(): pool_idx = h0_indices[i_n] if pool_idx >= 0: + # State indexing in int64: envelope-strided pools (unified memory / + # page-major) put pool_idx * stride(0) far past 2^31 (the CuTe twin + # of the fla/chunk_delta_h.py stride_init_state fix). + pool_idx64 = cutlass.Int64(pool_idx) k_local = in_warp_tid // V_PER_WARP_SMALL v_local = in_warp_tid % V_PER_WARP_SMALL v_base = warp_idx * V_PER_WARP_SMALL @@ -423,7 +431,7 @@ def _define_kernels(): h_val = 0.0 if v_global_load < v.shape[3]: h_val = cutlass.Float32( - h0_source[(pool_idx, i_hv, v_global_load, k_load)] + h0_source[(pool_idx64, i_hv, v_global_load, k_load)] ) sData[(k_load, v_load, stage)] = h_val @@ -480,7 +488,7 @@ def _define_kernels(): if k_write < TILE_K: v_global_write = v_tile * TILE_V_SMALL + v_write if v_global_write < v.shape[3]: - h0_source[(pool_idx, i_hv, v_global_write, k_write)] = ( + h0_source[(pool_idx64, i_hv, v_global_write, k_write)] = ( sData[(k_write, v_write, stage)] ) @@ -523,6 +531,10 @@ def _define_kernels(): pool_idx = h0_indices[i_n] if pool_idx >= 0: + # State indexing in int64: envelope-strided pools (unified memory / + # page-major) put pool_idx * stride(0) far past 2^31 (the CuTe twin + # of the fla/chunk_delta_h.py stride_init_state fix). + pool_idx64 = cutlass.Int64(pool_idx) k_local = in_warp_tid // V_PER_WARP v_local = in_warp_tid % V_PER_WARP v_base = warp_idx * V_PER_WARP @@ -634,7 +646,7 @@ def _define_kernels(): h_val = 0.0 if v_global_load < v.shape[3]: h_val = cutlass.Float32( - h0_source[(pool_idx, i_hv, v_global_load, k_load)] + h0_source[(pool_idx64, i_hv, v_global_load, k_load)] ) sData[(k_load, v_load, stage)] = h_val @@ -691,7 +703,7 @@ def _define_kernels(): if k_write < TILE_K: v_global_write = v_tile * TILE_V + v_write if v_global_write < v.shape[3]: - h0_source[(pool_idx, i_hv, v_global_write, k_write)] = ( + h0_source[(pool_idx64, i_hv, v_global_write, k_write)] = ( sData[(k_write, v_write, stage)] ) @@ -734,6 +746,10 @@ def _define_kernels(): pool_idx = h0_indices[i_n] if pool_idx >= 0: + # State indexing in int64: envelope-strided pools (unified memory / + # page-major) put pool_idx * stride(0) far past 2^31 (the CuTe twin + # of the fla/chunk_delta_h.py stride_init_state fix). + pool_idx64 = cutlass.Int64(pool_idx) k_local = in_warp_tid // V_PER_WARP v_local = in_warp_tid % V_PER_WARP v_base = warp_idx * V_PER_WARP @@ -845,7 +861,7 @@ def _define_kernels(): h_val = 0.0 if v_global_load < v.shape[3]: h_val = cutlass.Float32( - h0_source[(pool_idx, i_hv, v_global_load, k_load)] + h0_source[(pool_idx64, i_hv, v_global_load, k_load)] ) sData[(k_load, v_load, stage)] = h_val @@ -902,7 +918,7 @@ def _define_kernels(): if k_write < TILE_K: v_global_write = v_tile * TILE_V + v_write if v_global_write < v.shape[3]: - h0_source[(pool_idx, i_hv, v_global_write, k_write)] = ( + h0_source[(pool_idx64, i_hv, v_global_write, k_write)] = ( sData[(k_write, v_write, stage)] ) @@ -1223,11 +1239,28 @@ def _get_jit_functions(): return _jit_functions -def _get_compiled_kernel(N, H, HV, K, V, pool_size, use_small_batch, is_varlen_decode): - """Get or compile the KDA kernel for given dimensions.""" +def _get_compiled_kernel(N, H, HV, K, V, h0_source, use_small_batch, is_varlen_decode): + """Get or compile the KDA kernel for given dimensions. + + ``h0_source`` is the caller's real state pool: ``from_dlpack`` bakes its + exact layout into the compiled kernel, so envelope-strided pools (unified + memory / page-major, slot stride(0) != HV*V*K) compile against their true + slot pitch. The cache key carries the strides alongside the shape. + """ global _compiled_kernels - key = (N, H, HV, K, V, pool_size, use_small_batch, is_varlen_decode) + pool_size = h0_source.shape[0] + key = ( + N, + H, + HV, + K, + V, + pool_size, + tuple(h0_source.stride()), + use_small_batch, + is_varlen_decode, + ) if key in _compiled_kernels: return _compiled_kernels[key] @@ -1250,7 +1283,6 @@ def _get_compiled_kernel(N, H, HV, K, V, pool_size, use_small_batch, is_varlen_d A_log = torch.zeros(HV, dtype=torch.float32, device="cuda") dt_bias = torch.zeros(HV, K, dtype=torch.bfloat16, device="cuda") - h0_source = torch.zeros(pool_size, HV, V, K, dtype=torch.float32, device="cuda") h0_indices = torch.zeros(N, dtype=torch.int32, device="cuda") cu_seqlens_tensor = from_dlpack(cu_seqlens, assumed_align=16) @@ -1261,7 +1293,7 @@ def _get_compiled_kernel(N, H, HV, K, V, pool_size, use_small_batch, is_varlen_d b_tensor = from_dlpack(b, assumed_align=16) A_log_tensor = from_dlpack(A_log, assumed_align=16) dt_bias_tensor = from_dlpack(dt_bias, assumed_align=16) - h0_source_tensor = from_dlpack(h0_source, assumed_align=16) + h0_source_tensor = from_dlpack(h0_source.detach(), assumed_align=16) h0_indices_tensor = from_dlpack(h0_indices, assumed_align=16) o_tensor = from_dlpack(o, assumed_align=16) @@ -1304,6 +1336,7 @@ def _get_compiled_kernel(N, H, HV, K, V, pool_size, use_small_batch, is_varlen_d logger.info( "CuTe DSL KDA kernel compiled: " f"N={N}, H={H}, HV={HV}, K={K}, V={V}, pool_size={pool_size}, " + f"pool_strides={tuple(h0_source.stride())}, " f"small_batch={use_small_batch}, varlen={is_varlen_decode}" ) return compiled_kernel @@ -1369,6 +1402,9 @@ def cutedsl_fused_sigmoid_gating_kda_update( State layout contract: initial_state_source.shape == (pool_size, HV, V, K) + The slot dim may be envelope-strided (stride(0) > HV*V*K under + unified-memory / page-major pools); each slot's [HV, V, K] block must + be compact. State updates are written back in place through the view. Dense decode: q/k: (N, 1, H, K) @@ -1453,7 +1489,16 @@ def cutedsl_fused_sigmoid_gating_kda_update( A_log = _normalize_A_log(A_log, HV) dt_bias = _normalize_dt_bias(dt_bias, HV, K) - h0_source = h0_source.contiguous() + # h0_source may be an envelope-strided pool view (unified memory / + # page-major): slot stride(0) is the per-slot envelope pitch, not HV*V*K. + # Never .contiguous() it — on a strided view that copies, so the kernel's + # in-place state update would land in a dropped temporary. The kernel only + # needs each slot's [HV, V, K] block itself to be compact. + assert h0_source.stride()[1:] == (V * K, K, 1), ( + "CuTe DSL KDA decode requires a compact per-slot [HV, V, K] state " + f"block; got strides {tuple(h0_source.stride())} for shape " + f"{tuple(h0_source.shape)}" + ) initial_state_indices = initial_state_indices.contiguous() if cu_seqlens is not None: @@ -1496,7 +1541,7 @@ def cutedsl_fused_sigmoid_gating_kda_update( stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) compiled_kernel = _get_compiled_kernel( - N, H, HV, K, V, pool_size, use_small_batch, is_varlen_decode + N, H, HV, K, V, h0_source, use_small_batch, is_varlen_decode ) compiled_kernel( diff --git a/python/sglang/kernels/ops/attention/fla/fused_recurrent.py b/python/sglang/kernels/ops/attention/fla/fused_recurrent.py index e2880355e..5ffdea776 100644 --- a/python/sglang/kernels/ops/attention/fla/fused_recurrent.py +++ b/python/sglang/kernels/ops/attention/fla/fused_recurrent.py @@ -414,6 +414,7 @@ def fused_recurrent_kda_packed_decode_kernel( ht, ssm_state_indices, scale, + lower_bound, stride_mixed_qkv_tok: tl.constexpr, stride_a_tok: tl.constexpr, stride_b_tok: tl.constexpr, @@ -428,6 +429,7 @@ def fused_recurrent_kda_packed_decode_kernel( BV: tl.constexpr, SOFTPLUS_THRESHOLD: tl.constexpr, USE_QK_L2NORM_IN_KERNEL: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, ): """KDA packed decode: same shape as the GDN packed decode kernel, but with a per-K gate (``a`` is ``[B, HV*K]`` and ``dt_bias`` is ``[HV*K]``), @@ -475,8 +477,14 @@ def fused_recurrent_kda_packed_decode_kernel( A_log_val = tl.load(A_log + i_hv).to(tl.float32) x = b_a + b_dt - softplus_x = tl.where(x <= SOFTPLUS_THRESHOLD, tl.log(1.0 + tl.exp(x)), x) - b_g = -tl.exp(A_log_val) * softplus_x # [BK] + if USE_LOWER_BOUND: + # KDA safe gate: lower_bound * sigmoid(exp(A_log) * (g + bias)), + # matching the chunked prefill kernel (kda.py) and FLA reference. + b_g = lower_bound * tl.sigmoid(tl.exp(A_log_val) * x) # [BK] + else: + # Standard gate: -exp(A_log) * softplus(g + bias) + softplus_x = tl.where(x <= SOFTPLUS_THRESHOLD, tl.log(1.0 + tl.exp(x)), x) + b_g = -tl.exp(A_log_val) * softplus_x # [BK] b_val = tl.load(b + i_n * stride_b_tok + i_hv).to(tl.float32) # Keep beta in fp32 (no bf16 round-trip) to match the generic decode @@ -508,6 +516,7 @@ def fused_recurrent_kda_packed_decode( out: torch.Tensor, ssm_state_indices: torch.Tensor, use_qk_l2norm_in_kernel: bool = False, + lower_bound: Optional[float] = None, ) -> tuple[torch.Tensor, torch.Tensor]: """KDA T=1 decode fast path. Mirrors ``fused_recurrent_gated_delta_rule_packed_decode`` but the gate ``g`` is a per-K vector instead of a scalar. @@ -612,6 +621,40 @@ def fused_recurrent_kda_packed_decode( f"Invalid head config inferred from mixed_qkv: H={H}, HV={HV}." ) + # Batched-decode CUDA fast path: + # row-streaming state update reaches the in-place R+W bandwidth of the + # part (~9.6 TB/s) where this triton kernel tops out at ~5 TB/s holding a + # [BV, K] register tile per warp. ULP-level output differences only + # (reduction order); small batches keep triton (launch-bound anyway). + if use_qk_l2norm_in_kernel: + from sglang.kernels.ops.attention import kda_packed_decode as kda_decode_cuda + + if kda_decode_cuda.covered( + mixed_qkv, + a, + b, + A_log, + dt_bias, + initial_state, + out, + ssm_state_indices, + H, + ): + kda_decode_cuda.kda_packed_decode( + mixed_qkv, + a, + b, + A_log, + dt_bias, + scale, + initial_state, + out, + ssm_state_indices, + H, + lower_bound, + ) + return out, initial_state + BK = triton.next_power_of_2(K) if triton.cdiv(K, BK) != 1: raise ValueError( @@ -641,6 +684,7 @@ def fused_recurrent_kda_packed_decode( ht=initial_state, ssm_state_indices=ssm_state_indices, scale=scale, + lower_bound=lower_bound if lower_bound is not None else 0.0, stride_mixed_qkv_tok=stride_mixed_qkv_tok, stride_a_tok=stride_a_tok, stride_b_tok=stride_b_tok, @@ -655,6 +699,7 @@ def fused_recurrent_kda_packed_decode( BV=BV, SOFTPLUS_THRESHOLD=20.0, USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + USE_LOWER_BOUND=lower_bound is not None, num_warps=num_warps, num_stages=num_stages, ) @@ -830,6 +875,7 @@ def fused_recurrent_gated_delta_rule_update_fwd_kernel( o, h0_source, h0_indices, + stride_h0_source, cu_seqlens, scale, intermediate_states_buffer, @@ -897,22 +943,28 @@ def fused_recurrent_gated_delta_rule_update_fwd_kernel( b_h = tl.zeros([BV, BK], dtype=tl.float32) if USE_INITIAL_STATE: - idx = tl.load(h0_indices + i_n) + # Slot stride comes from the caller (h0_source.stride(0)): the state pool + # may be an envelope-strided view (page-major / unified memory), where the + # per-slot pitch spans ALL layers' state, not HV*K*V. int64: envelope + # pitches overflow an int32 index product. + idx = tl.load(h0_indices + i_n).to(tl.int64) # Add bounds checking for idx if idx >= 0: # Assuming negative indices are invalid p_h0 = ( h0_source - + idx * HV * K * V + + idx * stride_h0_source + i_hv * K * V + o_v[:, None] * K + o_k[None, :] ) b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) - # Prepare intermediate state cache variables if enabled + # Prepare intermediate state cache variables if enabled. int64: the buffer + # is contiguous but `cache_idx * cache_steps * HV * K * V` can exceed int32 + # for large slot counts. cache_idx = -1 if CACHE_INTERMEDIATE_STATES: - cache_idx = tl.load(intermediate_state_indices + i_n) + cache_idx = tl.load(intermediate_state_indices + i_n).to(tl.int64) step_idx = 0 for _ in range(0, T): @@ -987,11 +1039,11 @@ def fused_recurrent_gated_delta_rule_update_fwd_kernel( # Store final state back to h0_source with bounds checking # ssm states if not DISABLE_STATE_UPDATE: - idx = tl.load(h0_indices + i_n) + idx = tl.load(h0_indices + i_n).to(tl.int64) if idx >= 0: # Add bounds checking p_h0 = ( h0_source - + idx * HV * K * V + + idx * stride_h0_source + i_hv * K * V + o_v[:, None] * K + o_k[None, :] @@ -1053,6 +1105,11 @@ def fused_recurrent_gated_delta_rule_update_fwd( o=o, h0_source=initial_state_source, h0_indices=initial_state_indices, + # Envelope-strided state pools (page-major / unified memory) have a + # per-slot pitch != HV*K*V; contiguous pools pass exactly HV*K*V. + stride_h0_source=( + initial_state_source.stride(0) if initial_state_source is not None else 0 + ), cu_seqlens=cu_seqlens, scale=scale, intermediate_states_buffer=intermediate_states_buffer, diff --git a/python/sglang/kernels/ops/attention/fla/fused_sigmoid_gating_recurrent.py b/python/sglang/kernels/ops/attention/fla/fused_sigmoid_gating_recurrent.py index 521b0edf4..a7ffe49e6 100644 --- a/python/sglang/kernels/ops/attention/fla/fused_sigmoid_gating_recurrent.py +++ b/python/sglang/kernels/ops/attention/fla/fused_sigmoid_gating_recurrent.py @@ -12,6 +12,7 @@ def fused_sigmoid_gating_delta_rule_update_kernel( dt_bias, softplus_beta, softplus_threshold, + lower_bound, q, k, v, @@ -19,6 +20,7 @@ def fused_sigmoid_gating_delta_rule_update_kernel( o, h0_source, h0_indices, + stride_h0_source, cu_seqlens, # Parameters for target_verify support (unused for decode) intermediate_states_buffer, @@ -47,10 +49,14 @@ def fused_sigmoid_gating_delta_rule_update_kernel( USE_QK_L2NORM_IN_KERNEL: tl.constexpr, IS_VARLEN: tl.constexpr, IS_KDA: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, # Optional flags for target_verify support (default False for decode) DISABLE_STATE_UPDATE: tl.constexpr = False, CACHE_INTERMEDIATE_STATES: tl.constexpr = False, HAS_EAGLE_TREE_CUSTOM_ATTN_MASK: tl.constexpr = False, + # ReplaySSM fused ring-write. Pointers stay None and CACHE_RING False for + # decode / flag-off -> byte-identical. The gate ring layout follows IS_KDA + # (see the store below). replayssm_rawv=None, replayssm_rawk=None, replayssm_g=None, @@ -104,11 +110,15 @@ def fused_sigmoid_gating_delta_rule_update_kernel( b_h = tl.zeros([BK, BV], dtype=tl.float32) if USE_INITIAL_STATE: - idx = tl.load(h0_indices + i_n) + # Slot stride comes from the caller (h0_source.stride(0)): the state pool + # may be an envelope-strided view (page-major / unified memory), where the + # per-slot pitch spans ALL layers' state, not HV*K*V. int64: envelope + # pitches overflow an int32 index product. + idx = tl.load(h0_indices + i_n).to(tl.int64) if idx >= 0: p_h0 = ( h0_source - + idx * HV * K * V + + idx * stride_h0_source + i_hv * K * V + o_v[None, :] * K + o_k[:, None] @@ -128,10 +138,12 @@ def fused_sigmoid_gating_delta_rule_update_kernel( retrieve_parent_token_base, mask=mask_retrieve, other=0 ) - # Prepare intermediate state cache index if enabled + # Prepare intermediate state cache index if enabled. int64: the buffer is + # contiguous but `cache_idx * cache_steps * HV * K * V` can exceed int32 for + # large slot counts. cache_idx = -1 if CACHE_INTERMEDIATE_STATES: - cache_idx = tl.load(intermediate_state_indices + i_n) + cache_idx = tl.load(intermediate_state_indices + i_n).to(tl.int64) step_idx = 0 for _ in range(0, T): @@ -169,23 +181,32 @@ def fused_sigmoid_gating_delta_rule_update_kernel( b_a = tl.load(p_a).to(tl.float32) b_dt_bias = tl.load(p_dt_bias).to(tl.float32) - # Compute g = -exp(A_log) * softplus(a + dt_bias) x = b_a + b_dt_bias - beta_x = softplus_beta * x - # Apply softplus with numerical stability - softplus_x = tl.where( - beta_x <= softplus_threshold, - (1.0 / softplus_beta) * tl.log(1.0 + tl.exp(beta_x)), - x, - ) - b_g = -tl.exp(b_A_log) * softplus_x + if USE_LOWER_BOUND: + # KDA safe gate: lower_bound * sigmoid(exp(A_log) * (a + dt_bias)) + b_g = lower_bound * tl.sigmoid(tl.exp(b_A_log) * x) + else: + # Compute g = -exp(A_log) * softplus(a + dt_bias) + beta_x = softplus_beta * x + # Apply softplus with numerical stability + softplus_x = tl.where( + beta_x <= softplus_threshold, + (1.0 / softplus_beta) * tl.log(1.0 + tl.exp(beta_x)), + x, + ) + b_g = -tl.exp(b_A_log) * softplus_x # Compute beta = sigmoid(b) b_beta = 1.0 / (1.0 + tl.exp(-b_b)) - # Stored here, pre-l2norm k / pre-delta v, so the commit fold's replay - # is bit-identical to the update below; steps >= MAX_CACHE_LEN would - # smash the next slot's ring. + # fused ring-write: stash this step's raw inputs + in-kernel gate/beta + # into the per-slot ring for the commit fold to replay. Must sit here -- + # b_k is still pre-l2norm, b_v still pre-delta, b_g/b_beta are formed, + # so the fold's replay is bit-identical to the update below. rawk uses + # the k-head i_h (shared across a GQA group); rawv/g/beta use the v-head + # i_hv. step_idx < MAX_CACHE_LEN: absorb-inflated rows can exceed the + # ring; the overflow steps are past the committable prefix, so drop them + # (writing them would smash the next slot's ring). if CACHE_RING: ring_slot = tl.load(h0_indices + i_n).to(tl.int64) if ring_slot >= 0 and step_idx < MAX_CACHE_LEN: @@ -208,13 +229,30 @@ def fused_sigmoid_gating_delta_rule_update_kernel( b_k.to(replayssm_rawk.dtype.element_ty), mask=mask_k, ) - tl.store( - replayssm_g - + ring_slot * stride_g_slot - + i_hv * MAX_CACHE_LEN - + step_idx, - b_g, - ) + # b_g follows IS_KDA: KDA loads a/dt_bias with mask_k, so the + # gate is a per-K vector and the ring row is K wide; GDN's is + # a scalar per (head, step). The two layouts are not + # interchangeable -- storing one into the other's stride is a + # shape error, not a slow path -- and memory_pool.py sizes + # replayssm_g off the same is_kda test. + if IS_KDA: + tl.store( + replayssm_g + + ring_slot * stride_g_slot + + i_hv * MAX_CACHE_LEN * K + + step_idx * K + + o_k, + b_g, + mask=mask_k, + ) + else: + tl.store( + replayssm_g + + ring_slot * stride_g_slot + + i_hv * MAX_CACHE_LEN + + step_idx, + b_g, + ) if i_k == 0: tl.store( replayssm_beta @@ -277,11 +315,11 @@ def fused_sigmoid_gating_delta_rule_update_kernel( # Store final state back to h0_source with bounds checking if not DISABLE_STATE_UPDATE: if USE_INITIAL_STATE: - idx = tl.load(h0_indices + i_n) + idx = tl.load(h0_indices + i_n).to(tl.int64) if idx >= 0: p_h0 = ( h0_source - + idx * HV * K * V + + idx * stride_h0_source + i_hv * K * V + o_v[None, :] * K + o_k[:, None] @@ -305,6 +343,7 @@ def fused_sigmoid_gating_delta_rule_update( use_qk_l2norm_in_kernel: bool = False, cu_seqlens: Optional[torch.Tensor] = None, is_kda: bool = False, + lower_bound: Optional[float] = None, # Optional parameters for target_verify support disable_state_update: bool = False, intermediate_states_buffer: Optional[torch.Tensor] = None, @@ -313,6 +352,9 @@ def fused_sigmoid_gating_delta_rule_update( int ] = None, # kept for API compat; stride is derived from ``intermediate_states_buffer.shape[1]`` retrieve_parent_token: Optional[torch.Tensor] = None, + # fused ReplaySSM ring-write (spec verify). When cache_ring, each draft step + # stores pre-norm k / raw v / gate / beta into these per-slot rings, + # replacing the eager ring-write. Off by default -> decode unchanged. cache_ring: bool = False, replayssm_rawv: Optional[torch.Tensor] = None, replayssm_rawk: Optional[torch.Tensor] = None, @@ -375,14 +417,18 @@ def fused_sigmoid_gating_delta_rule_update( else 0 ) + # ring strides (per-slot rings are contiguous [num_slots, heads, L, dim]; + # the kernel offsets within a slot with MAX_CACHE_LEN and the dim extents). if cache_ring: - assert not is_kda, "cache_ring supports GDN only (scalar gate layout)" # stride(0) is used as the slot pitch, so a tensor still carrying the - # layer dim would scribble outside its slot. + # layer dim would scribble outside its slot. The gate ring is the one + # whose rank depends on the model: per-K vector for KDA, per-head scalar + # for GDN, matching g_shape in memory_pool.py and the IS_KDA branch in + # the store above. assert ( replayssm_rawv.dim() == 4 and replayssm_rawk.dim() == 4 - and replayssm_g.dim() == 3 + and replayssm_g.dim() == (4 if is_kda else 3) and replayssm_beta.dim() == 3 ), "cache_ring expects per-layer ring views" max_cache_len = replayssm_rawv.shape[-2] @@ -400,6 +446,7 @@ def fused_sigmoid_gating_delta_rule_update( dt_bias=dt_bias, softplus_beta=softplus_beta, softplus_threshold=softplus_threshold, + lower_bound=lower_bound if lower_bound is not None else 0.0, q=q, k=k, v=v, @@ -407,6 +454,11 @@ def fused_sigmoid_gating_delta_rule_update( o=o, h0_source=initial_state_source, h0_indices=initial_state_indices, + # Envelope-strided state pools (page-major / unified memory) have a + # per-slot pitch != HV*K*V; contiguous pools pass exactly HV*K*V. + stride_h0_source=( + initial_state_source.stride(0) if initial_state_source is not None else 0 + ), cu_seqlens=cu_seqlens, intermediate_states_buffer=intermediate_states_buffer, intermediate_state_indices=intermediate_state_indices, @@ -433,6 +485,7 @@ def fused_sigmoid_gating_delta_rule_update( USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, IS_VARLEN=cu_seqlens is not None, IS_KDA=is_kda, + USE_LOWER_BOUND=lower_bound is not None, DISABLE_STATE_UPDATE=disable_state_update, CACHE_INTERMEDIATE_STATES=intermediate_states_buffer is not None, HAS_EAGLE_TREE_CUSTOM_ATTN_MASK=retrieve_parent_token is not None, diff --git a/python/sglang/kernels/ops/attention/fla/gdn_replayssm_spec_decode.py b/python/sglang/kernels/ops/attention/fla/gdn_replayssm_spec_decode.py index d578cd422..0b4ef1a69 100644 --- a/python/sglang/kernels/ops/attention/fla/gdn_replayssm_spec_decode.py +++ b/python/sglang/kernels/ops/attention/fla/gdn_replayssm_spec_decode.py @@ -38,7 +38,8 @@ Closed-loop exact fold (the state / output error split): ``"tf32"`` (~5e-4, tensor-core path; worst case through the (I+A)^{-1} amplification 2^(BS-1) still lands at the floor). ``"ieee"`` / ``"tf32x3"`` remain selectable for ablations. The committed state is untouched by any of - these dots (requires the fp32 SSM checkpoint; enforced in server_args). + these dots (fp32 SSM checkpoint is the server_args default; a 16-bit + checkpoint is allowed with a warning, unvalidated for GDN). Differences from the vLLM reference: * SGLang passes **split** ``q`` / ``k`` / ``v`` tensors (already split + post diff --git a/python/sglang/kernels/ops/attention/fla/kda.py b/python/sglang/kernels/ops/attention/fla/kda.py index 091ee5c6b..c5d2defc1 100644 --- a/python/sglang/kernels/ops/attention/fla/kda.py +++ b/python/sglang/kernels/ops/attention/fla/kda.py @@ -25,8 +25,11 @@ from sglang.kernels.ops.attention.fla.index import ( from sglang.kernels.ops.attention.fla.l2norm import l2norm_fwd from sglang.kernels.ops.attention.fla.op import exp, exp2, log from sglang.kernels.ops.attention.fla.utils import ( + autotune_cache_kwargs, check_shared_mem, is_intel, + is_nvidia, + is_tf32_supported, ) if is_intel: @@ -514,18 +517,8 @@ def chunk_kda_scaled_dot_kkt_fwd( return A, Aqk -@triton.autotune( - configs=[ - triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages) - for BK in [64, 128] - for BV in [64, 128] - for num_warps in [2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=["H", "K", "V", "BT", "IS_VARLEN"], -) @triton.jit(do_not_specialize=["T"]) -def recompute_w_u_fwd_kernel( +def _recompute_w_u_fwd_kernel( k, kg, v, @@ -645,6 +638,67 @@ def recompute_w_u_fwd_kernel( tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) +_RECOMPUTE_W_U_CONFIGS = [ + triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages) + for BK in [64, 128] + for BV in [64, 128] + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] +] + +recompute_w_u_fwd_kernel = triton.autotune( + configs=_RECOMPUTE_W_U_CONFIGS, + key=["H", "K", "V", "BT", "IS_VARLEN"], + **autotune_cache_kwargs, +)(_recompute_w_u_fwd_kernel) + +_K3_RECOMPUTE_W_U_CONFIGS = { + (9, 0): {"BK": 128, "BV": 128, "num_warps": 8, "num_stages": 2}, + (10, 3): {"BK": 64, "BV": 128, "num_warps": 8, "num_stages": 2}, +} + + +@torch.inference_mode() +def precompile_k3_recompute_w_u_kernel( + *, num_heads: int, dtype: torch.dtype, device: torch.device +) -> bool: + device = torch.device(device) + if ( + not is_nvidia + or device.type != "cuda" + or torch.cuda.get_device_capability(device) not in _K3_RECOMPUTE_W_U_CONFIGS + ): + return False + + shape = (1, 1, num_heads, 128) + k = torch.zeros(shape, dtype=dtype, device=device) + v = torch.zeros_like(k) + beta = torch.zeros((1, 1, num_heads), dtype=dtype, device=device) + A = torch.zeros((1, 1, num_heads, 64), dtype=dtype, device=device) + gk = torch.zeros(shape, dtype=torch.float32, device=device) + cu_seqlens = torch.tensor([0, 1], dtype=torch.int64, device=device) + recompute_w_u_fwd(k, v, beta, A, gk=gk, cu_seqlens=cu_seqlens) + return True + + +def _get_k3_recompute_w_u_config( + k: torch.Tensor, + gk: torch.Tensor | None, + cu_seqlens: torch.LongTensor | None, + K: int, + V: int, + BT: int, +) -> dict | None: + if ( + not is_nvidia + or gk is None + or cu_seqlens is None + or (K, V, BT) != (128, 128, 64) + ): + return None + return _K3_RECOMPUTE_W_U_CONFIGS.get(torch.cuda.get_device_capability(k.device)) + + def recompute_w_u_fwd( k: torch.Tensor, v: torch.Tensor, @@ -664,7 +718,13 @@ def recompute_w_u_fwd( w = torch.empty_like(k) u = torch.empty_like(v) kg = torch.empty_like(k) if gk is not None else None - recompute_w_u_fwd_kernel[(NT, B * H)]( + static_config = _get_k3_recompute_w_u_config(k, gk, cu_seqlens, K, V, BT) + kernel = ( + _recompute_w_u_fwd_kernel + if static_config is not None + else recompute_w_u_fwd_kernel + ) + kernel[(NT, B * H)]( k=k, kg=kg, v=v, @@ -682,7 +742,8 @@ def recompute_w_u_fwd( BT=BT, STORE_KG=kg is not None, IS_VARLEN=cu_seqlens is not None, - DOT_PRECISION="tf32", + DOT_PRECISION="tf32" if is_tf32_supported else "ieee", + **(static_config or {}), ) return w, u, kg @@ -1126,6 +1187,9 @@ def chunk_kda_fwd( del Aqk, v_new if output_intermediate_states: + # h holds the recurrent state at every chunk-size boundary + # ([1, NT, H, V, K] packed across cu_seqlens) — the mamba radix + # track path snapshots per-chunk states from it during extend. return o, h del h return o diff --git a/python/sglang/kernels/ops/attention/fla/kda_replayssm_spec_decode.py b/python/sglang/kernels/ops/attention/fla/kda_replayssm_spec_decode.py new file mode 100644 index 000000000..d5108fbc8 --- /dev/null +++ b/python/sglang/kernels/ops/attention/fla/kda_replayssm_spec_decode.py @@ -0,0 +1,386 @@ +# SPDX-License-Identifier: Apache-2.0 +"""ReplaySSM speculative-decode state commit for KDA (Kimi Delta Attention). + +KDA keeps its own recurrent verify kernel for the per-step OUTPUT (unchanged), +so — unlike the GDN spec kernel (gdn_replayssm_spec_decode.py) — this module does +NOT reconstruct the verify output. It only replaces the per-draft-token full-SSM +snapshot cache (``intermediate_ssm``, ``max_running×(γ+1)×[HV,V,K]``, the memory +hog that collapses dspark concurrency) with a small per-request input window + +an exact-fold on commit. + +Scheme (fold-every-commit, no circular ring / periodic flush): + * during verify, the KDA backend stores the draft window's raw inputs + (raw v, raw pre-norm k, per-K log-decay gate ``gk`` (fp32), beta (fp32)) into + the per-slot ring at positions ``0..spec_len``; + * on commit, :func:`commit_kda_replayssm_spec` replays the *accepted* prefix + ``0..accept_len`` from the persistent checkpoint ``h0`` into ``h0`` in place — + ``h0`` is always the current committed state, so the next verify reads it + directly as its initial state (no lag, no chunked reconstruction). + +The exact-fold is a BITWISE CLONE of ``fused_recurrent_gated_delta_rule_fwd_kernel``'s +``IS_KDA`` branch (fused_recurrent.py): same [BK, BV] fp32 tile (K rows, V cols), +same division-form L2 norm (eps inside sqrt), same per-K gate decay +``h *= exp(gk)``, same decay→delta→rank-1 op order. Given identical inputs the +folded checkpoint is bit-identical to the recurrent baseline's committed state +(the delta-rule recurrence is contractive, so no length-dependent error). Do NOT +reorder into tl.dot / reciprocal-multiply; keep num_warps=1 so the reduction +trees match. + +Linear chain only (dspark γ chain, topk<=1). +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def kda_replayssm_exact_fold_kernel( + h0, # [num_slots, HV, V, K] fp32 checkpoint (folded in place) + rawv_cache, # [num_slots, HV, L, V] raw v + rawk_cache, # [num_slots, H, L, K] raw pre-norm k + gk_cache, # [num_slots, HV, L, K] fp32 per-K log-decay gate + beta_cache, # [num_slots, HV, L] fp32 beta + ssm_state_indices, # [B] int physical slot per request + accept_lens, # [B] int committed prefix length per request + mamba_track_indices, # [B] int extra_buffer track slot (or NULL) per request + mamba_steps_to_track, # [B] int crossing step (or -1) per request + stride_state_slot: tl.constexpr, + stride_rawv_slot: tl.constexpr, + stride_rawk_slot: tl.constexpr, + stride_gk_slot: tl.constexpr, + stride_beta_slot: tl.constexpr, + stride_state_layer: tl.constexpr, + stride_rawv_layer: tl.constexpr, + stride_rawk_layer: tl.constexpr, + stride_gk_layer: tl.constexpr, + stride_beta_layer: tl.constexpr, + stride_indices: tl.constexpr, + stride_accept: tl.constexpr, + stride_track: tl.constexpr, + stride_steps: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + MAX_CACHE_LEN: tl.constexpr, + USE_QK_L2NORM_IN_KERNEL: tl.constexpr, + NULL_BLOCK_ID: tl.constexpr, + HAS_TRACK: tl.constexpr, +): + i_v = tl.program_id(0) + i_n = tl.program_id(1) + # program_id(2) packs (layer, v-head): layer-major so a single launch folds + # all KDA layers. num_layers=1 launches (per-layer entry) keep i_layer == 0. + i_hvl = tl.program_id(2) + # int64: layer stride * i_layer overflows int32 at K3 scale (69 layers x + # ~34M-element per-layer stride > 2^31). + i_layer = (i_hvl // HV).to(tl.int64) + i_hv = i_hvl % HV + i_h = i_hv // (HV // H) + # Shift the layer-indexed bases once; every pointer below is layer-relative. + h0 = h0 + i_layer * stride_state_layer + rawv_cache = rawv_cache + i_layer * stride_rawv_layer + rawk_cache = rawk_cache + i_layer * stride_rawk_layer + gk_cache = gk_cache + i_layer * stride_gk_layer + beta_cache = beta_cache + i_layer * stride_beta_layer + + state_idx = tl.load(ssm_state_indices + i_n * stride_indices).to(tl.int64) + if state_idx <= NULL_BLOCK_ID: + return + n_commit = tl.load(accept_lens + i_n * stride_accept).to(tl.int32) + if n_commit <= 0: + return + + # extra_buffer: snapshot the interval-crossing state into the track ping-pong + # slot on the step it crosses (mask-gated: -1 track step / NULL slot = skip). + if HAS_TRACK: + track_idx = tl.load(mamba_track_indices + i_n * stride_track).to(tl.int64) + track_step = tl.load(mamba_steps_to_track + i_n * stride_steps).to(tl.int32) + else: + track_idx = NULL_BLOCK_ID + track_step = -1 + + o_k = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + mask_k = o_k < K + mask_v = o_v < V + mask_h = mask_k[:, None] & mask_v[None, :] + + # [BK, BV] tile: K rows / V cols, matching the recurrent baseline's memory + # offset (v * K + k, K contiguous). + p_h0 = ( + h0 + + state_idx * stride_state_slot + + i_hv * V * K + + o_v[None, :] * K + + o_k[:, None] + ) + b_h = tl.load(p_h0, mask=mask_h, other=0.0).to(tl.float32) + + for t in range(0, n_commit): + phys = t.to(tl.int64) + b_k = tl.load( + rawk_cache + + state_idx * stride_rawk_slot + + (i_h * MAX_CACHE_LEN + phys) * K + + o_k, + mask=mask_k, + other=0.0, + ).to(tl.float32) + b_v = tl.load( + rawv_cache + + state_idx * stride_rawv_slot + + (i_hv * MAX_CACHE_LEN + phys) * V + + o_v, + mask=mask_v, + other=0.0, + ).to(tl.float32) + b_gk = tl.load( + gk_cache + + state_idx * stride_gk_slot + + (i_hv * MAX_CACHE_LEN + phys) * K + + o_k, + mask=mask_k, + other=0.0, + ).to(tl.float32) + b_beta = tl.load( + beta_cache + state_idx * stride_beta_slot + i_hv * MAX_CACHE_LEN + phys + ).to(tl.float32) + + # --- verbatim recurrent update, IS_KDA branch (see module docstring) --- + if USE_QK_L2NORM_IN_KERNEL: + b_k = b_k / (tl.sqrt(tl.sum(b_k * b_k) + 1e-6)) + b_h *= tl.exp(b_gk[:, None]) # per-K gate decay, broadcast over V + b_v -= tl.sum(b_h * b_k[:, None], 0) + b_v *= b_beta + b_h += b_k[:, None] * b_v[None, :] + + # Interval-crossing snapshot -> track slot (state AFTER step `track_step`). + if HAS_TRACK: + if (t == track_step) and (track_idx > NULL_BLOCK_ID): + tl.store( + h0 + + track_idx * stride_state_slot + + i_hv * V * K + + o_v[None, :] * K + + o_k[:, None], + b_h.to(h0.dtype.element_ty), + mask=mask_h, + ) + + tl.store(p_h0, b_h.to(p_h0.dtype.element_ty), mask=mask_h) + + +def commit_kda_replayssm_spec( + checkpoint_state: torch.Tensor, # [num_slots, HV, V, K] fp32, folded in place + rawv_cache: torch.Tensor, # [num_slots, HV, L, V] + rawk_cache: torch.Tensor, # [num_slots, H, L, K] + gk_cache: torch.Tensor, # [num_slots, HV, L, K] fp32 + beta_cache: torch.Tensor, # [num_slots, HV, L] fp32 + ssm_state_indices: torch.Tensor, # [B] int + accept_lens: torch.Tensor, # [B] int (incl. the bonus token) + max_cache_len: int, + num_k_heads: int, + mamba_track_indices: torch.Tensor | None = None, # [B] extra_buffer track slot + mamba_steps_to_track: torch.Tensor | None = None, # [B] crossing step (or -1) + use_qk_l2norm_in_kernel: bool = True, + null_block_id: int = 0, +) -> None: + """Replay each request's accepted window into its fp32 checkpoint in place. + + Tiling clones the recurrent kernel (full-K rows, BV = min(np2(V), 32) cols, + num_warps=1) so the folded checkpoint is bit-identical to the recurrent + baseline's committed state. With extra_buffer (mamba_track_indices given) the + same replay snapshots the interval-crossing state into the track slot in one + pass, so no separate track scatter / force-flush is needed. + """ + num_slots, HV, V, K = checkpoint_state.shape + B = ssm_state_indices.shape[0] + BK = triton.next_power_of_2(K) + BV = min(triton.next_power_of_2(V), 32) + grid = (triton.cdiv(V, BV), B, HV) + has_track = mamba_track_indices is not None and mamba_steps_to_track is not None + if has_track: + track_idx_t = mamba_track_indices + steps_t = mamba_steps_to_track + stride_track = track_idx_t.stride(0) + stride_steps = steps_t.stride(0) + else: + track_idx_t = ssm_state_indices # unused (HAS_TRACK False); pass a valid ptr + steps_t = accept_lens + stride_track = 0 + stride_steps = 0 + kda_replayssm_exact_fold_kernel[grid]( + checkpoint_state, + rawv_cache, + rawk_cache, + gk_cache, + beta_cache, + ssm_state_indices, + accept_lens, + track_idx_t, + steps_t, + checkpoint_state.stride(0), + rawv_cache.stride(0), + rawk_cache.stride(0), + gk_cache.stride(0), + beta_cache.stride(0), + 0, # stride_*_layer unused: single-layer entry, i_layer == 0 + 0, + 0, + 0, + 0, + ssm_state_indices.stride(0), + accept_lens.stride(0), + stride_track, + stride_steps, + H=num_k_heads, + HV=HV, + K=K, + V=V, + BK=BK, + BV=BV, + MAX_CACHE_LEN=max_cache_len, + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + NULL_BLOCK_ID=null_block_id, + HAS_TRACK=has_track, + num_warps=1, + num_stages=3, + ) + + +def commit_kda_replayssm_spec_all_layers( + checkpoint_state: torch.Tensor, # [num_layers, num_slots, HV, V, K] fp32, in place + rawv_cache: torch.Tensor, # [num_layers, num_slots, HV, L, V] + rawk_cache: torch.Tensor, # [num_layers, num_slots, H, L, K] + gk_cache: torch.Tensor, # [num_layers, num_slots, HV, L, K] fp32 + beta_cache: torch.Tensor, # [num_layers, num_slots, HV, L] fp32 + ssm_state_indices: torch.Tensor, # [B] int (shared across layers) + accept_lens: torch.Tensor, # [B] int + max_cache_len: int, + num_k_heads: int, + mamba_track_indices: torch.Tensor | None = None, + mamba_steps_to_track: torch.Tensor | None = None, + use_qk_l2norm_in_kernel: bool = True, + null_block_id: int = 0, +) -> None: + """Fold every layer's accepted window in a single launch. + + Replaces the per-layer Python loop over commit_kda_replayssm_spec (one launch + per KDA layer -> ~69 tiny eager launches at bs=1, dispatch-bound). The layer + is packed into the head grid axis (program_id(2) = layer * HV + head), so the + result is bit-identical to the loop -- each (layer, head, v-tile) block runs + the same per-slot recurrent replay. ssm_state_indices / accept_lens / track + are per-request and shared across layers. + """ + num_layers, num_slots, HV, V, K = checkpoint_state.shape + B = ssm_state_indices.shape[0] + BK = triton.next_power_of_2(K) + BV = min(triton.next_power_of_2(V), 32) + grid = (triton.cdiv(V, BV), B, HV * num_layers) + has_track = mamba_track_indices is not None and mamba_steps_to_track is not None + if has_track: + track_idx_t = mamba_track_indices + steps_t = mamba_steps_to_track + stride_track = track_idx_t.stride(0) + stride_steps = steps_t.stride(0) + else: + track_idx_t = ssm_state_indices # unused (HAS_TRACK False); valid ptr + steps_t = accept_lens + stride_track = 0 + stride_steps = 0 + kda_replayssm_exact_fold_kernel[grid]( + checkpoint_state, + rawv_cache, + rawk_cache, + gk_cache, + beta_cache, + ssm_state_indices, + accept_lens, + track_idx_t, + steps_t, + checkpoint_state.stride(1), + rawv_cache.stride(1), + rawk_cache.stride(1), + gk_cache.stride(1), + beta_cache.stride(1), + checkpoint_state.stride(0), + rawv_cache.stride(0), + rawk_cache.stride(0), + gk_cache.stride(0), + beta_cache.stride(0), + ssm_state_indices.stride(0), + accept_lens.stride(0), + stride_track, + stride_steps, + H=num_k_heads, + HV=HV, + K=K, + V=V, + BK=BK, + BV=BV, + MAX_CACHE_LEN=max_cache_len, + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + NULL_BLOCK_ID=null_block_id, + HAS_TRACK=has_track, + num_warps=1, + num_stages=3, + ) + + +def commit_kda_replayssm_after_verify( + *, + spec_state, # MambaPool.SpeculativeState (all layers) + state_batch_indices: torch.Tensor, # [B] per-req mamba slot + accept_lens: torch.Tensor, # [B] int, incl. the bonus token + last_correct_step_indices: torch.Tensor, # [B] conv rollback target step + mamba_track_indices: torch.Tensor | None = None, + mamba_steps_to_track: torch.Tensor | None = None, + null_block_id: int = -1, +) -> None: + """Fold each layer's accepted window into `temporal` and roll back conv. + + Single commit entry point shared by the generic spec_utils commit and the + dspark/dflash direct `update_mamba_state_after_mtp_verify` path. The SSM state + lives in the per-slot ring (written during verify); the fold replays the + accepted prefix into the fp32 checkpoint, so `temporal` stays current. Conv + still needs its usual accept-rollback. + """ + from sglang.kernels.ops.mamba.mamba_state_scatter_triton import ( + fused_conv_window_scatter_with_mask, + ) + + L = spec_state.replayssm_rawv.shape[-2] + num_k_heads = spec_state.replayssm_rawk.shape[2] + commit_kda_replayssm_spec_all_layers( + checkpoint_state=spec_state.temporal, + rawv_cache=spec_state.replayssm_rawv, + rawk_cache=spec_state.replayssm_rawk, + gk_cache=spec_state.replayssm_g, + beta_cache=spec_state.replayssm_beta, + ssm_state_indices=state_batch_indices, + accept_lens=accept_lens, + max_cache_len=L, + num_k_heads=num_k_heads, + mamba_track_indices=mamba_track_indices, + mamba_steps_to_track=mamba_steps_to_track, + null_block_id=null_block_id, + ) + # Conv rollback + track-slot conv snapshot, per conv group (fold already did + # the ssm side via HAS_TRACK). Loop mirrors the recurrent commit's zip; track + # scatter is mask-gated (step -1 => skip). + for conv_states, interm_conv in zip( + spec_state.conv, spec_state.intermediate_conv_window + ): + fused_conv_window_scatter_with_mask( + conv_states, interm_conv, state_batch_indices, last_correct_step_indices + ) + if mamba_track_indices is not None and mamba_steps_to_track is not None: + fused_conv_window_scatter_with_mask( + conv_states, interm_conv, mamba_track_indices, mamba_steps_to_track + ) diff --git a/python/sglang/kernels/ops/attention/flash_attention_v4.py b/python/sglang/kernels/ops/attention/flash_attention_v4.py index 456036a58..eaf935363 100644 --- a/python/sglang/kernels/ops/attention/flash_attention_v4.py +++ b/python/sglang/kernels/ops/attention/flash_attention_v4.py @@ -23,6 +23,10 @@ else: _flash_attn_import_error = None +def is_flash_attention_v4_available() -> bool: + return _flash_attn_varlen_func is not None + + def _maybe_contiguous(x: Optional[torch.Tensor]) -> Optional[torch.Tensor]: return x.contiguous() if x is not None and x.stride(-1) != 1 else x diff --git a/python/sglang/kernels/ops/attention/kda_fused_decode.py b/python/sglang/kernels/ops/attention/kda_fused_decode.py new file mode 100644 index 000000000..56579ea86 --- /dev/null +++ b/python/sglang/kernels/ops/attention/kda_fused_decode.py @@ -0,0 +1,166 @@ +"""Fully fused KDA decode step (Kimi K3 batched decode fast path). + +One kernel replaces the three-kernel decode chain +``causal_conv1d_update -> kda_packed_decode -> rms_norm_gated``: it reads the +raw (pre-conv) qkv slice straight out of the fused projection GEMM output, +does the causal conv1d update (conv state shifted in the pool in place), the +delta-rule recurrence (l2-normed q/k, softplus forget gate, sigmoid beta), +and the sigmoid-gated output RMSNorm. + +Kernel body vendored from the NVIDIA x Moonshot Kimi K3 optimization package +(see csrc/attention/kda_fused_decode.cuh for provenance and the list of +integration patches). Specialized for the K3 KDA decode regime: +K = V = 128, kernel width 4, no lower bound, T = 1 per request. +The JIT currently instantiates local head counts H = HV in {12, 6, 3} +(TP8, TP16, and TP32). + +The model must hand off the output-norm gate (attempt-and-verify stash on the +attention layer, see kimi_k3.py), and a covered() check gates supported inputs. +Everything else falls back to the unfused chain. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.kernels.jit.utils import ( + cache_once, + is_arch_support_pdl, + load_jit, + make_cpp_args, +) + +if TYPE_CHECKING: + from tvm_ffi.module import Module + +_SUPPORTED_HEADS = {3, 6, 12} +_CONV_STATE_W = 3 # kernel width 4 -> 3 cached tokens + + +@cache_once +def _jit_kda_fused_decode_module() -> Module: + args = make_cpp_args(is_arch_support_pdl()) + return load_jit( + "kda_fused_decode", + *args, + cuda_files=["attention/kda_fused_decode.cuh"], + cuda_wrappers=[("run", f"KdaFusedDecodeKernel<{args}>::run")], + extra_cuda_cflags=["-O3", "--use_fast_math"], + ) + + +def covered( + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + conv_states: torch.Tensor, + ssm_states: torch.Tensor, + cache_indices: torch.Tensor, + onorm_g: torch.Tensor, +) -> bool: + """The kernel is compiled for the K3 KDA decode regime: H heads of 128, + packed [T, 3*H*128] qkv rows, transposed [slots, 3, 3*H*128] conv pool, fp32 + [slots, H, 128, 128] ssm pool (inner-contiguous, any slot pitch — the + kernel reads the real slot stride), one token per request.""" + if ssm_states.ndim < 4: + return False + H, V, K = ssm_states.shape[-3:] + if H not in _SUPPORTED_HEADS: + return False + seg = H * 128 + conv_dim = 3 * seg + if mixed_qkv.ndim != 2 or mixed_qkv.shape[-1] != conv_dim: + return False + return ( + V == 128 + and K == 128 + and a.ndim == 2 + and a.shape[-1] == seg + and b.ndim == 2 + and b.shape[-1] == H + and onorm_g.ndim == 2 + and onorm_g.shape[-1] == seg + and conv_states.ndim == 3 + and conv_states.shape[-2:] == (_CONV_STATE_W, conv_dim) + and mixed_qkv.dtype == torch.bfloat16 + and a.dtype == torch.bfloat16 + and b.dtype == torch.bfloat16 + and onorm_g.dtype == torch.bfloat16 + and conv_states.dtype == torch.bfloat16 + and ssm_states.dtype == torch.float32 + and cache_indices.dtype == torch.int32 + and mixed_qkv.stride(-1) == 1 + and a.stride(-1) == 1 + and b.stride(-1) == 1 + and onorm_g.stride(-1) == 1 + and conv_states.stride(-1) == 1 + # Inner [HV, V, K] must be contiguous (the kernel float4-loads V*K + # chunks); the slot pitch (stride(-4)) is arbitrary — a locally + # allocated pool packs it at HV*V*K, the unified / page-major pools at + # the multi-layer envelope. The kernel reads ssm_states.stride(0), so + # any slot pitch is fine. (Do NOT use .view(-1, HV, V, K): that fails / + # copies on an envelope-strided view.) + and ssm_states.stride(-1) == 1 + and ssm_states.stride(-2) == K + and ssm_states.stride(-3) == V * K + and cache_indices.is_contiguous() + ) + + +def kda_fused_decode( + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + conv_states: torch.Tensor, + w_q_t: torch.Tensor, + w_k_t: torch.Tensor, + w_v_t: torch.Tensor, + conv_bias: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + onorm_g: torch.Tensor, + onorm_weight: torch.Tensor, + ssm_states: torch.Tensor, + cache_indices: torch.Tensor, + scale: float, + onorm_eps: float, + lower_bound: Optional[float] = None, +) -> torch.Tensor: + """In-place fused decode step: shifts `conv_states` and updates + `ssm_states` rows selected by `cache_indices` (rows < 0 are padded + cuda-graph slots and only zero their output), returns the gated-normed + attention output [1, B, HV, V] (the packed-decode output layout). + Caller must have checked covered().""" + B = mixed_qkv.shape[0] + H = ssm_states.shape[-3] + seg = H * 128 + out = torch.empty((B, seg), dtype=torch.bfloat16, device=mixed_qkv.device) + _jit_kda_fused_decode_module().run( + mixed_qkv, + a, + b, + conv_states, + w_q_t, + w_k_t, + w_v_t, + conv_bias, + A_log, + dt_bias, + onorm_g, + onorm_weight, + # Pass the pool view as-is (already [slots, HV, V, K]); the kernel + # binding reads its real slot stride via state.stride(0). A + # .view(-1, H, 128, 128) here would break on envelope-strided pools + # (unified / page-major) — the reshape can't fold a non-dense slot + # pitch and would raise / silently copy. + ssm_states, + cache_indices, + out, + float(scale), + float(onorm_eps), + float(lower_bound) if lower_bound is not None else 0.0, + lower_bound is not None, + ) + return out.view(1, B, H, 128) diff --git a/python/sglang/kernels/ops/attention/kda_packed_decode.py b/python/sglang/kernels/ops/attention/kda_packed_decode.py new file mode 100644 index 000000000..38bf4f1d3 --- /dev/null +++ b/python/sglang/kernels/ops/attention/kda_packed_decode.py @@ -0,0 +1,120 @@ +"""CUDA KDA packed-decode kernel (batched decode fast path). + +Row-streaming port of the triton fused_recurrent_kda_packed_decode_kernel: +the triton kernel keeps a [BV, K] fp32 state tile in one warp's registers and +tops out at ~5 TB/s; this kernel streams the state one 512B row at a time and +reaches the in-place read+write bandwidth of the part (~9.6 TB/s probe). +Outputs match the triton kernel to ULPs (warp-shuffle reduction order), not +bits. Unsupported inputs fall back to triton through a covered() check. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.kernels.jit.utils import ( + cache_once, + is_arch_support_pdl, + load_jit, + make_cpp_args, +) + +if TYPE_CHECKING: + from tvm_ffi.module import Module + +_WARPS: int = 8 +# The row-streaming layout needs enough (batch x head) CTAs to fill the GPU; +# below this the triton kernel's launch cost is already the floor. +_MIN_BATCH: int = 8 + + +@cache_once +def _jit_kda_packed_decode_module() -> Module: + args = make_cpp_args(_WARPS, is_arch_support_pdl()) + return load_jit( + "kda_packed_decode_" + str(_WARPS), + *args, + cuda_files=["attention/kda_packed_decode.cuh"], + cuda_wrappers=[("run", f"KdaPackedDecodeKernel<{args}>::run")], + extra_cuda_cflags=["-O3"], + ) + + +def covered( + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + initial_state: torch.Tensor, + out: torch.Tensor, + ssm_state_indices: torch.Tensor, + num_q_heads: int, +) -> bool: + B = mixed_qkv.shape[0] + HV, V, K = initial_state.shape[-3:] + return ( + B >= _MIN_BATCH + and K == 128 + and V == 128 + and HV % max(num_q_heads, 1) == 0 + # Per-K gate layout. A per-head scalar gate ([B, HV] / [HV], the GDN + # shape) is a different kernel, not a slower input for this one. + and a.dim() == 2 + and a.shape[1] == HV * K + and dt_bias.numel() == HV * K + and mixed_qkv.dtype == torch.bfloat16 + and a.dtype == torch.bfloat16 + and b.dtype == torch.bfloat16 + and A_log.dtype == torch.float32 + and dt_bias.dtype == torch.float32 + and initial_state.dtype == torch.float32 + and out.dtype == torch.bfloat16 + and ssm_state_indices.dtype == torch.int32 + and mixed_qkv.stride(-1) == 1 + and a.stride(-1) == 1 + and b.stride(-1) == 1 + and initial_state.stride(-1) == 1 + and initial_state.stride(-2) == K + and initial_state.stride(-3) == V * K + and out.is_contiguous() + and ssm_state_indices.is_contiguous() + ) + + +def kda_packed_decode( + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + out: torch.Tensor, + ssm_state_indices: torch.Tensor, + num_q_heads: int, + lower_bound: Optional[float] = None, +) -> None: + """In-place KDA decode step: updates `initial_state` rows selected by + `ssm_state_indices` and writes attention output into `out` ([B, 1, HV, V]). + Caller must have checked covered(); q/k l2-norm is always applied + (matches the production dispatch).""" + B = mixed_qkv.shape[0] + HV, V, _ = initial_state.shape[-3:] + state = initial_state.view(-1, *initial_state.shape[-3:]) + _jit_kda_packed_decode_module().run( + mixed_qkv, + a, + b, + A_log, + dt_bias, + out.view(B, HV, V), + state, + ssm_state_indices, + float(scale), + float(lower_bound) if lower_bound is not None else 0.0, + lower_bound is not None, + int(num_q_heads), + ) diff --git a/python/sglang/kernels/ops/attention/linear/kda_blackwell/kernel_h.py b/python/sglang/kernels/ops/attention/linear/kda_blackwell/kernel_h.py index 078f48e7e..0e426c0c1 100644 --- a/python/sglang/kernels/ops/attention/linear/kda_blackwell/kernel_h.py +++ b/python/sglang/kernels/ops/attention/linear/kda_blackwell/kernel_h.py @@ -698,7 +698,19 @@ def kda_h_cutedsl( ``h0``/``ht`` may be the full state pool; ``state_indices`` [N] int32 maps each sequence to its row, so state gather/scatter fuses into the kernel's TMA load/store (no per-call state intermediates). + + Envelope-strided pools (unified memory / page-major, slot stride(0) != + Hv*V*K) are supported natively: the compile-time fake tensors carry + dynamic int64 strides (``make_fake_tensor``), so the TMA descriptors pick + the real slot pitch up at launch. The fakes assume 16-element stride + divisibility (TMA also needs 16-byte global strides); guard it loudly + rather than corrupt state. """ + for name, t in (("h0", h0), ("ht", ht)): + assert t.stride(-1) == 1 and all(s % 16 == 0 for s in t.stride()[:-1]), ( + f"kda_h_cutedsl: {name} strides {tuple(t.stride())} violate the " + "16-element divisibility the kernel was compiled with" + ) _, Hv, K_dim = kg.shape _, _, V_dim = V.shape h_dtype = {torch.bfloat16: BFloat16, torch.float32: Float32}[h0.dtype] diff --git a/python/sglang/kernels/ops/attention/linear/kda_nvidia_prefill/Akk_inverse_lower_triangle_bf16.py b/python/sglang/kernels/ops/attention/linear/kda_nvidia_prefill/Akk_inverse_lower_triangle_bf16.py new file mode 100644 index 000000000..98036d5f0 --- /dev/null +++ b/python/sglang/kernels/ops/attention/linear/kda_nvidia_prefill/Akk_inverse_lower_triangle_bf16.py @@ -0,0 +1,1008 @@ +# Vendored from the NVIDIA KDA_prefill package (benchmark/ Blackwell path) +# for the Kimi-K3 chunked prefill forward. Local deltas: fla.* imports +# re-pointed to sglang's vendored fla subset, flat sibling imports made +# package-relative, RCP_LN2 inlined. +# ruff: noqa -- vendored kernel library, minimal local deltas +""" +Akk 64×64 Lower Triangular Block Inversion — Full BF16. + +BF16 input, BF16 SMEM (packed bf16x2), BF16 MMA m16n8k16 (FP32 accum), BF16 output. +Numerical precision aligned with Flash Attention / KDA kernels. + + (I+L)⁻¹ = (I-L)(I+L²)(I+L⁴)(I+L⁸) [L strictly lower triangular, L¹⁶=0] + +Architecture: + - 4 warps (128 threads) per CTA, each CTA processes one 64×64 Akk matrix + - sAkk [64,36] fp32 — each FP32 slot holds packed bf16x2 (stride 36) + Total SMEM: 64*36*4 = 9216 bytes (vs 64*72*4 = 18432 in FP32 version) + - sTemp [16, 24, 2] fp32: inter-warp FP32 accumulator communication + - cp.async: BF16 global → packed bf16x2 in FP32 SMEM (raw bytes align) + - BF16 MMA m16n8k16 with FP32 accumulator for ALL matmuls + - movmatrix.sync.aligned.m8n8.trans.b16 for A→B layout conversion + - C→A chain: cvt.rn.bf16x2.f32 to pack FP32 accum → bf16x2 A-operand + +Block layout in sAkk (4×4 sub-blocks of 16×16, packed bf16x2): + Upper tri = INPUT, Diagonal = in-place inversion, Lower tri = OUTPUT + +Stages: + 0. cp.async load bf16 64×64 → sAkk (packed bf16x2) + 1. Invert 4 diagonal blocks via Neumann series (all 4 warps, BF16 MMA) + 2. Warps 0-2: Ai10, Ai21, Ai32 via chain MMA (C→A pack) + 3. Warps 0+2 → Ai20, warps 1+3 → Ai31 (parallel pairs, sTemp) + 4. Warps 0+1+2 → Ai30 (sTemp aggregation) + 5. All warps: store bf16x2 SMEM → bf16 global + +Inputs: A_in [B, T, H, BT] bf16 +Outputs: A_out [B, T, H, BT] bf16 +""" + +import cuda.bindings.driver as cuda_drv +import cutlass +import cutlass.cute as cute +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync +from cutlass.cutlass_dsl import T, dsl_user_op + +BS = 64 +SB = 16 +THREADS = 128 +TEMP_PAD = 8 +TEMP_COLS = SB + TEMP_PAD # 24 +NUM_TEMPS = 2 +AKK_PAD = 4 +AKK_STRIDE = BS // 2 + AKK_PAD # 36 (in FP32 units = 72 bf16 elements) + + +@dsl_user_op +def mma_bf16_m16n8k16( + a0, + a1, + a2, + a3, + b0, + b1, + c0, + c1, + c2, + c3, + *, + loc=None, + ip=None, +): + """mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 + A: 4×i32 (bf16x2 pairs), B: 2×i32, C/D: 4×f32 (FP32 accum).""" + a0b = llvm.bitcast(T.i32(), a0.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + a1b = llvm.bitcast(T.i32(), a1.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + a2b = llvm.bitcast(T.i32(), a2.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + a3b = llvm.bitcast(T.i32(), a3.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + b0b = llvm.bitcast(T.i32(), b0.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + b1b = llvm.bitcast(T.i32(), b1.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + ir.Type.parse("!llvm.struct<(f32, f32, f32, f32)>"), + [ + a0b, + a1b, + a2b, + a3b, + b0b, + b1b, + c0.ir_value(loc=loc, ip=ip), + c1.ir_value(loc=loc, ip=ip), + c2.ir_value(loc=loc, ip=ip), + c3.ir_value(loc=loc, ip=ip), + ], + """{ + mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 + {$0, $1, $2, $3}, + {$4, $5, $6, $7}, + {$8, $9}, + {$10, $11, $12, $13}; + }""", + "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + d0 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)) + d1 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)) + d2 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [2], loc=loc, ip=ip)) + d3 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [3], loc=loc, ip=ip)) + return d0, d1, d2, d3 + + +# Warp-level hardware A->B transpose; works on any 16-bit format including BF16. +@dsl_user_op +def _movmatrix_trans(src, *, loc=None, ip=None): + src_b = llvm.bitcast(T.i32(), src.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + T.i32(), + [src_b], + "movmatrix.sync.aligned.m8n8.trans.b16 $0, $1;", + "=r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip)) + + +@dsl_user_op +def _pack_bf16x2(lo_f32, hi_f32, *, loc=None, ip=None): + """Pack two FP32 values into one i32 holding two BF16 values. + cvt.rn.bf16x2.f32 converts and packs in one instruction.""" + result = llvm.inline_asm( + T.i32(), + [lo_f32.ir_value(loc=loc, ip=ip), hi_f32.ir_value(loc=loc, ip=ip)], + "cvt.rn.bf16x2.f32 $0, $2, $1;", + "=r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip)) + + +@dsl_user_op +def _mask_packed_ltri(packed, row, pair, *, loc=None, ip=None): + """Apply lower-triangular mask to packed bf16x2. + Zero out elements where row < col. col0=2*pair, col1=2*pair+1.""" + p = llvm.bitcast(T.i32(), packed.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + T.i32(), + [p, row.ir_value(loc=loc, ip=ip), pair.ir_value(loc=loc, ip=ip)], + """{ + .reg .b32 %c0, %c1, %mlo, %mhi, %mask; + .reg .pred %p0, %p1; + shl.b32 %c0, $3, 1; + add.u32 %c1, %c0, 1; + setp.ge.s32 %p0, $2, %c0; + setp.ge.s32 %p1, $2, %c1; + selp.b32 %mlo, 0xFFFF, 0, %p0; + selp.b32 %mhi, 0xFFFF0000, 0, %p1; + or.b32 %mask, %mlo, %mhi; + and.b32 $0, $1, %mask; + }""", + "=r,r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip)) + + +@dsl_user_op +def _unpack_bf16x2_lo(packed, *, loc=None, ip=None): + """Unpack lower BF16 from packed i32 to FP32. + BF16 is upper 16 bits of FP32, so shift left by 16.""" + p = llvm.bitcast(T.i32(), packed.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + T.i32(), + [p], + """{ + shl.b32 $0, $1, 16; + }""", + "=r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip)) + + +@dsl_user_op +def _unpack_bf16x2_hi(packed, *, loc=None, ip=None): + """Unpack upper BF16 from packed i32 to FP32. + Upper 16 bits are already in the right position for FP32.""" + p = llvm.bitcast(T.i32(), packed.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + T.i32(), + [p], + """{ + and.b32 $0, $1, 0xFFFF0000; + }""", + "=r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip)) + + +# Neumann diagonal 16×16 inversion using BF16 MMA m16n8k16 + FP32 accum +# +# (I+L)⁻¹ = (I-L)(I+L²)(I+L⁴)(I+L⁸) +# +# All intermediate values kept in FP32. Pack to bf16x2 only at MMA boundaries. +# sAkk is packed bf16x2: sAkk[row, pair] = {bf16[row, 2*pair], bf16[row, 2*pair+1]} +@dsl_user_op +def _invert_diag_neumann(sAkk: cute.Tensor, block_idx, lane_id, *, loc=None, ip=None): + r_off = block_idx * 16 + c_off = block_idx * 8 # packed bf16x2: 16 cols → 8 pairs + gid = lane_id // 4 + tid = lane_id % 4 + + # --- Load A from packed bf16x2 SMEM → unpack to 8 FP32 values --- + packed0 = cutlass.Float32(sAkk[r_off + gid, c_off + tid]) + packed1 = cutlass.Float32(sAkk[r_off + gid + 8, c_off + tid]) + packed2 = cutlass.Float32(sAkk[r_off + gid, c_off + 4 + tid]) + packed3 = cutlass.Float32(sAkk[r_off + gid + 8, c_off + 4 + tid]) + + A_f0 = _unpack_bf16x2_lo(packed0) # A[gid, 2*tid] + A_f1 = _unpack_bf16x2_hi(packed0) # A[gid, 2*tid+1] + A_f2 = _unpack_bf16x2_lo(packed1) # A[gid+8, 2*tid] + A_f3 = _unpack_bf16x2_hi(packed1) # A[gid+8, 2*tid+1] + A_f4 = _unpack_bf16x2_lo(packed2) # A[gid, 8+2*tid] + A_f5 = _unpack_bf16x2_hi(packed2) # A[gid, 8+2*tid+1] + A_f6 = _unpack_bf16x2_lo(packed3) # A[gid+8, 8+2*tid] + A_f7 = _unpack_bf16x2_hi(packed3) # A[gid+8, 8+2*tid+1] + + # Build identity (FP32) + _one = cutlass.Float32(1.0) + _zero = cutlass.Float32(0.0) + I_f0 = _one * cutlass.Float32(gid == 2 * tid) + _zero * cutlass.Float32( + gid != 2 * tid + ) + I_f1 = _one * cutlass.Float32(gid == 2 * tid + 1) + _zero * cutlass.Float32( + gid != 2 * tid + 1 + ) + I_f2 = _one * cutlass.Float32(gid + 8 == 2 * tid) + _zero * cutlass.Float32( + gid + 8 != 2 * tid + ) + I_f3 = _one * cutlass.Float32(gid + 8 == 2 * tid + 1) + _zero * cutlass.Float32( + gid + 8 != 2 * tid + 1 + ) + I_f4 = _one * cutlass.Float32(gid == 8 + 2 * tid) + _zero * cutlass.Float32( + gid != 8 + 2 * tid + ) + I_f5 = _one * cutlass.Float32(gid == 8 + 2 * tid + 1) + _zero * cutlass.Float32( + gid != 8 + 2 * tid + 1 + ) + I_f6 = _one * cutlass.Float32(gid + 8 == 8 + 2 * tid) + _zero * cutlass.Float32( + gid + 8 != 8 + 2 * tid + ) + I_f7 = _one * cutlass.Float32(gid + 8 == 8 + 2 * tid + 1) + _zero * cutlass.Float32( + gid + 8 != 8 + 2 * tid + 1 + ) + + # L = A - I (FP32) + L_f0 = A_f0 - I_f0 + L_f1 = A_f1 - I_f1 + L_f2 = A_f2 - I_f2 + L_f3 = A_f3 - I_f3 + L_f4 = A_f4 - I_f4 + L_f5 = A_f5 - I_f5 + L_f6 = A_f6 - I_f6 + L_f7 = A_f7 - I_f7 + + # INV = I - L = 2I - A (FP32) + INV_f0 = I_f0 - L_f0 + INV_f1 = I_f1 - L_f1 + INV_f2 = I_f2 - L_f2 + INV_f3 = I_f3 - L_f3 + INV_f4 = I_f4 - L_f4 + INV_f5 = I_f5 - L_f5 + INV_f6 = I_f6 - L_f6 + INV_f7 = I_f7 - L_f7 + + _zf = cutlass.Float32(0.0) + + # Pack L and INV → bf16x2 for MMA + L_a0 = _pack_bf16x2(L_f0, L_f1) + L_a1 = _pack_bf16x2(L_f2, L_f3) + L_a2 = _pack_bf16x2(L_f4, L_f5) + L_a3 = _pack_bf16x2(L_f6, L_f7) + + INV_a0 = _pack_bf16x2(INV_f0, INV_f1) + INV_a1 = _pack_bf16x2(INV_f2, INV_f3) + INV_a2 = _pack_bf16x2(INV_f4, INV_f5) + INV_a3 = _pack_bf16x2(INV_f6, INV_f7) + + # === Iteration 1: L² = L × L, then INV = INV + INV × L² === + L_b0 = _movmatrix_trans(L_a0) + L_b1 = _movmatrix_trans(L_a1) + L_b2 = _movmatrix_trans(L_a2) + L_b3 = _movmatrix_trans(L_a3) + + Lp_c0, Lp_c1, Lp_c2, Lp_c3 = mma_bf16_m16n8k16( + L_a0, L_a1, L_a2, L_a3, L_b0, L_b1, _zf, _zf, _zf, _zf + ) + Lp_c4, Lp_c5, Lp_c6, Lp_c7 = mma_bf16_m16n8k16( + L_a0, L_a1, L_a2, L_a3, L_b2, L_b3, _zf, _zf, _zf, _zf + ) + + Lp_a0 = _pack_bf16x2(Lp_c0, Lp_c1) + Lp_a1 = _pack_bf16x2(Lp_c2, Lp_c3) + Lp_a2 = _pack_bf16x2(Lp_c4, Lp_c5) + Lp_a3 = _pack_bf16x2(Lp_c6, Lp_c7) + + Lp_b0 = _movmatrix_trans(Lp_a0) + Lp_b1 = _movmatrix_trans(Lp_a1) + Lp_b2 = _movmatrix_trans(Lp_a2) + Lp_b3 = _movmatrix_trans(Lp_a3) + + # mm = INV × L² (FP32 accum output) + mm_c0, mm_c1, mm_c2, mm_c3 = mma_bf16_m16n8k16( + INV_a0, INV_a1, INV_a2, INV_a3, Lp_b0, Lp_b1, _zf, _zf, _zf, _zf + ) + mm_c4, mm_c5, mm_c6, mm_c7 = mma_bf16_m16n8k16( + INV_a0, INV_a1, INV_a2, INV_a3, Lp_b2, Lp_b3, _zf, _zf, _zf, _zf + ) + + # INV += mm (FP32) + INV_f0 = INV_f0 + mm_c0 + INV_f1 = INV_f1 + mm_c1 + INV_f2 = INV_f2 + mm_c2 + INV_f3 = INV_f3 + mm_c3 + INV_f4 = INV_f4 + mm_c4 + INV_f5 = INV_f5 + mm_c5 + INV_f6 = INV_f6 + mm_c6 + INV_f7 = INV_f7 + mm_c7 + + INV_a0 = _pack_bf16x2(INV_f0, INV_f1) + INV_a1 = _pack_bf16x2(INV_f2, INV_f3) + INV_a2 = _pack_bf16x2(INV_f4, INV_f5) + INV_a3 = _pack_bf16x2(INV_f6, INV_f7) + + # === Iteration 2: L⁴ = L² × L², then INV = INV + INV × L⁴ === + L4_c0, L4_c1, L4_c2, L4_c3 = mma_bf16_m16n8k16( + Lp_a0, Lp_a1, Lp_a2, Lp_a3, Lp_b0, Lp_b1, _zf, _zf, _zf, _zf + ) + L4_c4, L4_c5, L4_c6, L4_c7 = mma_bf16_m16n8k16( + Lp_a0, Lp_a1, Lp_a2, Lp_a3, Lp_b2, Lp_b3, _zf, _zf, _zf, _zf + ) + + L4_a0 = _pack_bf16x2(L4_c0, L4_c1) + L4_a1 = _pack_bf16x2(L4_c2, L4_c3) + L4_a2 = _pack_bf16x2(L4_c4, L4_c5) + L4_a3 = _pack_bf16x2(L4_c6, L4_c7) + + L4_b0 = _movmatrix_trans(L4_a0) + L4_b1 = _movmatrix_trans(L4_a1) + L4_b2 = _movmatrix_trans(L4_a2) + L4_b3 = _movmatrix_trans(L4_a3) + + mm_c0, mm_c1, mm_c2, mm_c3 = mma_bf16_m16n8k16( + INV_a0, INV_a1, INV_a2, INV_a3, L4_b0, L4_b1, _zf, _zf, _zf, _zf + ) + mm_c4, mm_c5, mm_c6, mm_c7 = mma_bf16_m16n8k16( + INV_a0, INV_a1, INV_a2, INV_a3, L4_b2, L4_b3, _zf, _zf, _zf, _zf + ) + + INV_f0 = INV_f0 + mm_c0 + INV_f1 = INV_f1 + mm_c1 + INV_f2 = INV_f2 + mm_c2 + INV_f3 = INV_f3 + mm_c3 + INV_f4 = INV_f4 + mm_c4 + INV_f5 = INV_f5 + mm_c5 + INV_f6 = INV_f6 + mm_c6 + INV_f7 = INV_f7 + mm_c7 + + INV_a0 = _pack_bf16x2(INV_f0, INV_f1) + INV_a1 = _pack_bf16x2(INV_f2, INV_f3) + INV_a2 = _pack_bf16x2(INV_f4, INV_f5) + INV_a3 = _pack_bf16x2(INV_f6, INV_f7) + + # === Iteration 3: L⁸ = L⁴ × L⁴, then INV = INV + INV × L⁸ === + L8_c0, L8_c1, L8_c2, L8_c3 = mma_bf16_m16n8k16( + L4_a0, L4_a1, L4_a2, L4_a3, L4_b0, L4_b1, _zf, _zf, _zf, _zf + ) + L8_c4, L8_c5, L8_c6, L8_c7 = mma_bf16_m16n8k16( + L4_a0, L4_a1, L4_a2, L4_a3, L4_b2, L4_b3, _zf, _zf, _zf, _zf + ) + + L8_a0 = _pack_bf16x2(L8_c0, L8_c1) + L8_a1 = _pack_bf16x2(L8_c2, L8_c3) + L8_a2 = _pack_bf16x2(L8_c4, L8_c5) + L8_a3 = _pack_bf16x2(L8_c6, L8_c7) + + L8_b0 = _movmatrix_trans(L8_a0) + L8_b1 = _movmatrix_trans(L8_a1) + L8_b2 = _movmatrix_trans(L8_a2) + L8_b3 = _movmatrix_trans(L8_a3) + + mm_c0, mm_c1, mm_c2, mm_c3 = mma_bf16_m16n8k16( + INV_a0, INV_a1, INV_a2, INV_a3, L8_b0, L8_b1, _zf, _zf, _zf, _zf + ) + mm_c4, mm_c5, mm_c6, mm_c7 = mma_bf16_m16n8k16( + INV_a0, INV_a1, INV_a2, INV_a3, L8_b2, L8_b3, _zf, _zf, _zf, _zf + ) + + INV_f0 = INV_f0 + mm_c0 + INV_f1 = INV_f1 + mm_c1 + INV_f2 = INV_f2 + mm_c2 + INV_f3 = INV_f3 + mm_c3 + INV_f4 = INV_f4 + mm_c4 + INV_f5 = INV_f5 + mm_c5 + INV_f6 = INV_f6 + mm_c6 + INV_f7 = INV_f7 + mm_c7 + + # --- Store INV to packed bf16x2 SMEM --- + sAkk[r_off + gid, c_off + tid] = _pack_bf16x2(INV_f0, INV_f1) + sAkk[r_off + gid + 8, c_off + tid] = _pack_bf16x2(INV_f2, INV_f3) + sAkk[r_off + gid, c_off + 4 + tid] = _pack_bf16x2(INV_f4, INV_f5) + sAkk[r_off + gid + 8, c_off + 4 + tid] = _pack_bf16x2(INV_f6, INV_f7) + + +# 16×16 matmul: load A & B from packed bf16x2 sAkk, BF16 MMA, return FP32 C +@dsl_user_op +def _matmul_AB( + sAkk: cute.Tensor, br_A, bc_A, br_B, bc_B, lane_id, *, loc=None, ip=None +): + gid = lane_id // 4 + tid = lane_id % 4 + _zf = cutlass.Float32(0.0) + rA = br_A * 16 + cA = bc_A * 8 # packed: 16 cols → 8 pairs + rB = br_B * 16 + cB = bc_B * 8 + + # Load A-operand (packed bf16x2, direct from SMEM) + a0 = cutlass.Float32(sAkk[rA + gid, cA + tid]) + a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + tid]) + a2 = cutlass.Float32(sAkk[rA + gid, cA + 4 + tid]) + a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 4 + tid]) + + # Load B-operand: load in A-layout then movmatrix → B-layout + bA0 = cutlass.Float32(sAkk[rB + gid, cB + tid]) + bA1 = cutlass.Float32(sAkk[rB + gid + 8, cB + tid]) + bA2 = cutlass.Float32(sAkk[rB + gid, cB + 4 + tid]) + bA3 = cutlass.Float32(sAkk[rB + gid + 8, cB + 4 + tid]) + b0 = _movmatrix_trans(bA0) + b1 = _movmatrix_trans(bA1) + b2 = _movmatrix_trans(bA2) + b3 = _movmatrix_trans(bA3) + + # 16×16 = 2 × m16n8k16 + cn0_0, cn0_1, cn0_2, cn0_3 = mma_bf16_m16n8k16( + a0, a1, a2, a3, b0, b1, _zf, _zf, _zf, _zf + ) + cn1_0, cn1_1, cn1_2, cn1_3 = mma_bf16_m16n8k16( + a0, a1, a2, a3, b2, b3, _zf, _zf, _zf, _zf + ) + + # Return 8 FP32 C-registers (C-layout of m16n8k16 FP32 accum) + return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3 + + +# Chain MMA: pre-loaded A (from C→A pack), load B from sAkk (Stage 2) +# A-operand already packed as bf16x2 from previous C result. +@dsl_user_op +def _chain_mma_B( + sAkk: cute.Tensor, br_B, bc_B, a0, a1, a2, a3, lane_id, *, loc=None, ip=None +): + gid = lane_id // 4 + tid = lane_id % 4 + _zf = cutlass.Float32(0.0) + rB = br_B * 16 + cB = bc_B * 8 + + bA0 = cutlass.Float32(sAkk[rB + gid, cB + tid]) + bA1 = cutlass.Float32(sAkk[rB + gid + 8, cB + tid]) + bA2 = cutlass.Float32(sAkk[rB + gid, cB + 4 + tid]) + bA3 = cutlass.Float32(sAkk[rB + gid + 8, cB + 4 + tid]) + b0 = _movmatrix_trans(bA0) + b1 = _movmatrix_trans(bA1) + b2 = _movmatrix_trans(bA2) + b3 = _movmatrix_trans(bA3) + + cn0_0, cn0_1, cn0_2, cn0_3 = mma_bf16_m16n8k16( + a0, a1, a2, a3, b0, b1, _zf, _zf, _zf, _zf + ) + cn1_0, cn1_1, cn1_2, cn1_3 = mma_bf16_m16n8k16( + a0, a1, a2, a3, b2, b3, _zf, _zf, _zf, _zf + ) + + return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3 + + +# Chain MMA (stages 3-4): load A from sAkk, B pre-loaded. B must already be in +# B-layout, i.e. packed bf16x2 after movmatrix -- not the raw FP32 shuffle result. +@dsl_user_op +def _chain_mma_A( + sAkk: cute.Tensor, br_A, bc_A, b0, b1, b2, b3, lane_id, *, loc=None, ip=None +): + gid = lane_id // 4 + tid = lane_id % 4 + _zf = cutlass.Float32(0.0) + rA = br_A * 16 + cA = bc_A * 8 + + a0 = cutlass.Float32(sAkk[rA + gid, cA + tid]) + a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + tid]) + a2 = cutlass.Float32(sAkk[rA + gid, cA + 4 + tid]) + a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 4 + tid]) + + cn0_0, cn0_1, cn0_2, cn0_3 = mma_bf16_m16n8k16( + a0, a1, a2, a3, b0, b1, _zf, _zf, _zf, _zf + ) + cn1_0, cn1_1, cn1_2, cn1_3 = mma_bf16_m16n8k16( + a0, a1, a2, a3, b2, b3, _zf, _zf, _zf, _zf + ) + + return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3 + + +# Store negated C result (16×16, FP32 accum) to packed bf16x2 sAkk +# C-layout for m16n8k16 FP32 accum: +# cn0_0 = C[gid, 0..7 left half col0], cn0_1 = C[gid+8, left half col0] +# cn0_2 = C[gid, left half col1], cn0_3 = C[gid+8, left half col1] +# cn1_* = right half +# Packing: negate FP32 then pack pairs → bf16x2 +@dsl_user_op +def _store_neg_C( + sAkk: cute.Tensor, + br, + bc, + c0, + c1, + c2, + c3, + c4, + c5, + c6, + c7, + lane_id, + *, + loc=None, + ip=None, +): + gid = lane_id // 4 + tid = lane_id % 4 + r = br * 16 + c = bc * 8 # packed + + # Pack negated FP32 pairs → bf16x2, then store + sAkk[r + gid, c + tid] = _pack_bf16x2(-c0, -c1) + sAkk[r + gid + 8, c + tid] = _pack_bf16x2(-c2, -c3) + sAkk[r + gid, c + 4 + tid] = _pack_bf16x2(-c4, -c5) + sAkk[r + gid + 8, c + 4 + tid] = _pack_bf16x2(-c6, -c7) + + +# Pack FP32 C-accum → bf16x2 A-operand for C→A chain +# C-layout (FP32 accum, m16n8k16): 8 floats → 4 bf16x2 A-regs +# c0,c1 → a0 (rows gid/gid+8, left-half k0..7) +# c2,c3 → a1 +# c4,c5 → a2 (right-half k8..15) +# c6,c7 → a3 +@dsl_user_op +def _pack_C_to_A(c0, c1, c2, c3, c4, c5, c6, c7, *, loc=None, ip=None): + a0 = _pack_bf16x2(c0, c1) + a1 = _pack_bf16x2(c2, c3) + a2 = _pack_bf16x2(c4, c5) + a3 = _pack_bf16x2(c6, c7) + return a0, a1, a2, a3 + + +# Convert FP32 C-accum → bf16x2 B-operand via pack + movmatrix +@dsl_user_op +def _pack_C_to_B(c0, c1, c2, c3, c4, c5, c6, c7, *, loc=None, ip=None): + a0 = _pack_bf16x2(c0, c1) + a1 = _pack_bf16x2(c2, c3) + a2 = _pack_bf16x2(c4, c5) + a3 = _pack_bf16x2(c6, c7) + b0 = _movmatrix_trans(a0) + b1 = _movmatrix_trans(a1) + b2 = _movmatrix_trans(a2) + b3 = _movmatrix_trans(a3) + return b0, b1, b2, b3 + + +# sTemp helpers (FP32, non-swizzled, for inter-warp accumulator exchange) +@dsl_user_op +def _store_C_temp( + sT: cute.Tensor, + buf, + c0, + c1, + c2, + c3, + c4, + c5, + c6, + c7, + lane_id, + *, + loc=None, + ip=None, +): + gid = lane_id // 4 + tid = lane_id % 4 + sT[gid, 2 * tid, buf] = c0 + sT[gid, 2 * tid + 1, buf] = c1 + sT[gid + 8, 2 * tid, buf] = c2 + sT[gid + 8, 2 * tid + 1, buf] = c3 + sT[gid, 8 + 2 * tid, buf] = c4 + sT[gid, 8 + 2 * tid + 1, buf] = c5 + sT[gid + 8, 8 + 2 * tid, buf] = c6 + sT[gid + 8, 8 + 2 * tid + 1, buf] = c7 + + +@dsl_user_op +def _load_C_temp(sT: cute.Tensor, buf, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + c0 = cutlass.Float32(sT[gid, 2 * tid, buf]) + c1 = cutlass.Float32(sT[gid, 2 * tid + 1, buf]) + c2 = cutlass.Float32(sT[gid + 8, 2 * tid, buf]) + c3 = cutlass.Float32(sT[gid + 8, 2 * tid + 1, buf]) + c4 = cutlass.Float32(sT[gid, 8 + 2 * tid, buf]) + c5 = cutlass.Float32(sT[gid, 8 + 2 * tid + 1, buf]) + c6 = cutlass.Float32(sT[gid + 8, 8 + 2 * tid, buf]) + c7 = cutlass.Float32(sT[gid + 8, 8 + 2 * tid + 1, buf]) + return c0, c1, c2, c3, c4, c5, c6, c7 + + +@cute.kernel +def akk_inv_kernel( + g2s_copy: cute.TiledCopy, + gA_tensor: cute.Tensor, + mOut: cute.Tensor, + mBeta: cute.Tensor, + akk_smem_layout: cute.Layout, + temp_layout: cute.Layout, + NT: int, + H: int, + mCuSeqlens: cute.Tensor, + mChunkIndices: cute.Tensor, + IS_VARLEN: cutlass.Constexpr[int], +): + tidx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + lane_id = tidx % 32 + h_idx, nt_idx, b_idx = cute.arch.block_idx() + + # ===== SMEM allocation ===== + smem = cutlass.utils.SmemAllocator() + sAkk = smem.allocate_tensor(cutlass.Float32, akk_smem_layout, 128) + sTemp = smem.allocate_tensor(cutlass.Float32, temp_layout, 128) + sBeta = smem.allocate_tensor(cutlass.Float32, cute.make_layout(BS, stride=1), 128) + + # ===== Stage 0: cp.async load bf16 global → packed bf16x2 sAkk ===== + ld_bnt = nt_idx + if IS_VARLEN: + ld_seq_id = cutlass.Int32(mChunkIndices[nt_idx, 0]) + ld_local = cutlass.Int32(mChunkIndices[nt_idx, 1]) + ld_bos = cutlass.Int32(mCuSeqlens[ld_seq_id]) + ld_bnt = ld_bos + ld_local * BS + gA_batch = gA_tensor[(None, None, h_idx, ld_bnt, b_idx)] + + thr_g2s = g2s_copy.get_slice(tidx) + thr_gSrc = thr_g2s.partition_S(gA_batch) + thr_sDst = thr_g2s.partition_D(sAkk) + cute.copy(g2s_copy, thr_gSrc, thr_sDst) + cute.arch.cp_async_commit_group() + cute.arch.cp_async_wait_group(0) + cute.arch.barrier() + + # Zero out-of-bounds rows for varlen partial chunks. + if IS_VARLEN: + _z_seq = cutlass.Int32(mChunkIndices[nt_idx, 0]) + _z_local = cutlass.Int32(mChunkIndices[nt_idx, 1]) + _z_bos = cutlass.Int32(mCuSeqlens[_z_seq]) + _z_eos = cutlass.Int32(mCuSeqlens[_z_seq + 1]) + _z_cs = _z_bos + _z_local * BS + _z_vr = _z_eos - _z_cs + _zr_start = warp_idx * SB + for ri in cutlass.range_constexpr(SB): + row = _zr_start + ri + if row >= _z_vr: + # Zero packed bf16x2 slots (each slot = 2 bf16 elements) + c0 = lane_id + if c0 < AKK_STRIDE: + sAkk[row, c0] = cutlass.Float32(0.0) + cute.arch.barrier() + + # ===== Stage 0b: Load 64 per-token betas into sBeta (fp32) ===== + _b_chunk_start = nt_idx * BS + _b_eos = cutlass.Int32(_b_chunk_start + BS) + if IS_VARLEN: + _b_seq_id = cutlass.Int32(mChunkIndices[nt_idx, 0]) + _b_local = cutlass.Int32(mChunkIndices[nt_idx, 1]) + _b_bos = cutlass.Int32(mCuSeqlens[_b_seq_id]) + _b_eos = cutlass.Int32(mCuSeqlens[_b_seq_id + 1]) + _b_chunk_start = _b_bos + _b_local * BS + + if warp_idx == 0: + _bcol_lo = lane_id * 2 + _bcol_hi = _bcol_lo + 1 + if _bcol_lo < BS: + _bt_lo = _b_chunk_start + _bcol_lo + _bt_hi = _b_chunk_start + _bcol_hi + if IS_VARLEN: + if _bt_lo < _b_eos: + sBeta[_bcol_lo] = cutlass.Float32(mBeta[b_idx, _bt_lo, h_idx]) + else: + sBeta[_bcol_lo] = cutlass.Float32(0.0) + if _bt_hi < _b_eos: + sBeta[_bcol_hi] = cutlass.Float32(mBeta[b_idx, _bt_hi, h_idx]) + else: + sBeta[_bcol_hi] = cutlass.Float32(0.0) + else: + sBeta[_bcol_lo] = cutlass.Float32(mBeta[b_idx, _bt_lo, h_idx]) + sBeta[_bcol_hi] = cutlass.Float32(mBeta[b_idx, _bt_hi, h_idx]) + cute.arch.barrier() + + # ===== Stage 1: Diagonal block inversion via Neumann series (BF16 MMA) ===== + _invert_diag_neumann(sAkk, warp_idx, lane_id) + + cute.arch.barrier() + + # ===== Stage 2: First batch — Ai10, Ai21, Ai32 ===== + # C→A chain: pack FP32 C-accum → bf16x2 A-operand via _pack_C_to_A + if warp_idx == 0: + t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk, 1, 1, 0, 1, lane_id) + a0, a1, a2, a3 = _pack_C_to_A(t0, t1, t2, t3, t4, t5, t6, t7) + r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_B( + sAkk, 0, 0, a0, a1, a2, a3, lane_id + ) + _store_neg_C(sAkk, 1, 0, r0, r1, r2, r3, r4, r5, r6, r7, lane_id) + + if warp_idx == 1: + t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk, 2, 2, 1, 2, lane_id) + a0, a1, a2, a3 = _pack_C_to_A(t0, t1, t2, t3, t4, t5, t6, t7) + r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_B( + sAkk, 1, 1, a0, a1, a2, a3, lane_id + ) + _store_neg_C(sAkk, 2, 1, r0, r1, r2, r3, r4, r5, r6, r7, lane_id) + + if warp_idx == 2: + t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk, 3, 3, 2, 3, lane_id) + a0, a1, a2, a3 = _pack_C_to_A(t0, t1, t2, t3, t4, t5, t6, t7) + r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_B( + sAkk, 2, 2, a0, a1, a2, a3, lane_id + ) + _store_neg_C(sAkk, 3, 2, r0, r1, r2, r3, r4, r5, r6, r7, lane_id) + + cute.arch.barrier() + + # ===== Stage 3: Second batch — Ai20, Ai31 (warp pairs via sTemp) ===== + _z = cutlass.Float32(0.0) + t0 = _z + t1 = _z + t2 = _z + t3 = _z + t4 = _z + t5 = _z + t6 = _z + t7 = _z + + # --- Ai20 = -Ai22 @ (Akk20 @ Ai00 + Akk21 @ Ai10) --- + if warp_idx == 0: + t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk, 0, 2, 0, 0, lane_id) + + if warp_idx == 2: + s0, s1, s2, s3, s4, s5, s6, s7 = _matmul_AB(sAkk, 1, 2, 1, 0, lane_id) + _store_C_temp(sTemp, 0, s0, s1, s2, s3, s4, s5, s6, s7, lane_id) + + # --- Ai31 = -Ai33 @ (Akk31 @ Ai11 + Akk32 @ Ai21) --- + if warp_idx == 1: + t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk, 1, 3, 1, 1, lane_id) + + if warp_idx == 3: + s0, s1, s2, s3, s4, s5, s6, s7 = _matmul_AB(sAkk, 2, 3, 2, 1, lane_id) + _store_C_temp(sTemp, 1, s0, s1, s2, s3, s4, s5, s6, s7, lane_id) + + cute.arch.barrier() + + # Warp 0: accumulate T1+T2, pack→B, multiply by Ai22 + if warp_idx == 0: + e0, e1, e2, e3, e4, e5, e6, e7 = _load_C_temp(sTemp, 0, lane_id) + t0 = t0 + e0 + t1 = t1 + e1 + t2 = t2 + e2 + t3 = t3 + e3 + t4 = t4 + e4 + t5 = t5 + e5 + t6 = t6 + e6 + t7 = t7 + e7 + b0, b1, b2, b3 = _pack_C_to_B(t0, t1, t2, t3, t4, t5, t6, t7) + r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_A( + sAkk, 2, 2, b0, b1, b2, b3, lane_id + ) + _store_neg_C(sAkk, 2, 0, r0, r1, r2, r3, r4, r5, r6, r7, lane_id) + + # Warp 1: accumulate T1'+T2', pack→B, multiply by Ai33 + if warp_idx == 1: + e0, e1, e2, e3, e4, e5, e6, e7 = _load_C_temp(sTemp, 1, lane_id) + t0 = t0 + e0 + t1 = t1 + e1 + t2 = t2 + e2 + t3 = t3 + e3 + t4 = t4 + e4 + t5 = t5 + e5 + t6 = t6 + e6 + t7 = t7 + e7 + b0, b1, b2, b3 = _pack_C_to_B(t0, t1, t2, t3, t4, t5, t6, t7) + r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_A( + sAkk, 3, 3, b0, b1, b2, b3, lane_id + ) + _store_neg_C(sAkk, 3, 1, r0, r1, r2, r3, r4, r5, r6, r7, lane_id) + + cute.arch.barrier() + + # ===== Stage 4: Third batch — Ai30 ===== + t0 = _z + t1 = _z + t2 = _z + t3 = _z + t4 = _z + t5 = _z + t6 = _z + t7 = _z + + if warp_idx == 0: + t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk, 0, 3, 0, 0, lane_id) + + if warp_idx == 1: + s0, s1, s2, s3, s4, s5, s6, s7 = _matmul_AB(sAkk, 1, 3, 1, 0, lane_id) + _store_C_temp(sTemp, 0, s0, s1, s2, s3, s4, s5, s6, s7, lane_id) + + if warp_idx == 2: + s0, s1, s2, s3, s4, s5, s6, s7 = _matmul_AB(sAkk, 2, 3, 2, 0, lane_id) + _store_C_temp(sTemp, 1, s0, s1, s2, s3, s4, s5, s6, s7, lane_id) + + cute.arch.barrier() + + # Warp 0: accumulate all three, pack→B, multiply by Ai33 + if warp_idx == 0: + e0, e1, e2, e3, e4, e5, e6, e7 = _load_C_temp(sTemp, 0, lane_id) + t0 = t0 + e0 + t1 = t1 + e1 + t2 = t2 + e2 + t3 = t3 + e3 + t4 = t4 + e4 + t5 = t5 + e5 + t6 = t6 + e6 + t7 = t7 + e7 + e0, e1, e2, e3, e4, e5, e6, e7 = _load_C_temp(sTemp, 1, lane_id) + t0 = t0 + e0 + t1 = t1 + e1 + t2 = t2 + e2 + t3 = t3 + e3 + t4 = t4 + e4 + t5 = t5 + e5 + t6 = t6 + e6 + t7 = t7 + e7 + b0, b1, b2, b3 = _pack_C_to_B(t0, t1, t2, t3, t4, t5, t6, t7) + r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_A( + sAkk, 3, 3, b0, b1, b2, b3, lane_id + ) + _store_neg_C(sAkk, 3, 0, r0, r1, r2, r3, r4, r5, r6, r7, lane_id) + + cute.arch.barrier() + + # ===== Stage 5: Store packed bf16x2 sAkk → global (b32 with direct bit-mask) ===== + vl_chunk_start = nt_idx * BS + vl_eos = cutlass.Int32(vl_chunk_start + BS) + if IS_VARLEN: + vl_seq_id = cutlass.Int32(mChunkIndices[nt_idx, 0]) + vl_local = cutlass.Int32(mChunkIndices[nt_idx, 1]) + vl_bos = cutlass.Int32(mCuSeqlens[vl_seq_id]) + vl_eos = cutlass.Int32(mCuSeqlens[vl_seq_id + 1]) + vl_chunk_start = vl_bos + vl_local * BS + + row_start = warp_idx * SB + + # PDL: hint downstream K4 to pre-launch so its setup work overlaps with our + # gmem writes below. K4 has griddepcontrol_wait after its setup but before + # reading any gmem (TMA loads / gk_last_exp), so ordering is correct. + cute.arch.griddepcontrol_launch_dependents() + + for ri in cutlass.range_constexpr(SB): + row = row_start + ri + pair = lane_id + if pair < BS // 2: + packed = cutlass.Float32(sAkk[row, pair]) + masked = _mask_packed_ltri(packed, cutlass.Int32(row), cutlass.Int32(pair)) + + _beta_lo = cutlass.Float32(sBeta[2 * pair]) + _beta_hi = cutlass.Float32(sBeta[2 * pair + 1]) + _lo_f32 = _unpack_bf16x2_lo(masked) * _beta_lo + _hi_f32 = _unpack_bf16x2_hi(masked) * _beta_hi + final = _pack_bf16x2(_lo_f32, _hi_f32) + + t_row = vl_chunk_start + row + if IS_VARLEN: + if t_row < vl_eos: + mOut[b_idx, t_row, h_idx, pair] = final + else: + mOut[b_idx, t_row, h_idx, pair] = final + + +@cute.jit +def akk_inv_host( + A_in: cute.Tensor, + A_out: cute.Tensor, + Beta_in: cute.Tensor, + B: cutlass.Constexpr[int], + NT: cutlass.Constexpr[int], + H: cutlass.Constexpr[int], + mCuSeqlens: cute.Tensor, + mChunkIndices: cute.Tensor, + IS_VARLEN: cutlass.Constexpr[int], + T_VAL: cutlass.Constexpr[int], + stream: cuda_drv.CUstream, +): + # BF16 input: view as FP32 (packed bf16x2). Each FP32 element = 2 BF16. + # Original BF16 shape: [B, T, H, BS] with stride per-element in BF16. + # As FP32 packed bf16x2: shape [BS, BS//2, H, dim3, B] + # The row stride in BF16 units is H*BS. In FP32 (bf16x2) units = H*BS/2. + # The col stride in BF16 is 1. In bf16x2 pairs: pair stride = 1 (adjacent pairs in memory). + _dim3_size = IS_VARLEN * T_VAL + (1 - IS_VARLEN) * (T_VAL // BS) + _dim3_stride_bf16 = IS_VARLEN * (H * BS) + (1 - IS_VARLEN) * (BS * H * BS) + # All strides in FP32 units (each FP32 = 2 BF16): + _row_stride = H * BS // 2 # row stride (BF16 row stride = H*BS, /2 for bf16x2) + _col_stride = 1 # adjacent bf16x2 pairs + _h_stride = BS // 2 # head stride (BF16 = BS, /2 for bf16x2) + _dim3_stride = _dim3_stride_bf16 // 2 + _batch_stride = T_VAL * H * BS // 2 + + view_layout = cute.make_layout( + (BS, BS // 2, H, _dim3_size, B), + stride=(_row_stride, _col_stride, _h_stride, _dim3_stride, _batch_stride), + ) + gA_view = cute.make_tensor(A_in.iterator, view_layout) + + # sAkk: packed bf16x2, logical shape (64, 32) with stride 36 for bank-conflict-free access + akk_smem_2d = cute.make_layout((BS, BS // 2), stride=(AKK_STRIDE, 1)) + + # cp.async G→S copy: 128-bit (4×fp32 = 8 bf16) vectorised + copy_atom = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), + cutlass.Float32, + num_bits_per_copy=128, + ) + # 128 threads x 4 FP32 (128b) each = 512 elements per iteration, so the + # 64x32 source takes 4 iterations. The copy maps the source shape + # (BS, BS//2) = (64, 32) onto the SMEM shape (BS, AKK_STRIDE) = (64, 36): + # only the (64, 32) data region is written, padding cols 32-35 stay zero. + g2s_copy = cute.make_tiled_copy_tv( + copy_atom, + thr_layout=cute.make_layout((16, 8), stride=(8, 1)), + val_layout=cute.make_layout((1, 4)), + ) + + # sTemp layout (non-swizzled) + temp_layout = cute.make_layout( + (SB, TEMP_COLS, NUM_TEMPS), stride=(TEMP_COLS, 1, SB * TEMP_COLS) + ) + + smem_bytes = BS * AKK_STRIDE * 4 + SB * TEMP_COLS * NUM_TEMPS * 4 + BS * 4 + 256 + + out_layout = cute.make_layout( + (B, T_VAL, H, BS // 2), stride=(T_VAL * H * BS // 2, H * BS // 2, BS // 2, 1) + ) + out_view = cute.make_tensor(A_out.iterator, out_layout) + + akk_inv_kernel( + g2s_copy, + gA_view, + out_view, + Beta_in, + akk_smem_2d, + temp_layout, + NT, + H, + mCuSeqlens, + mChunkIndices, + IS_VARLEN, + ).launch( + grid=(H, NT, B), + block=(THREADS, 1, 1), + smem=smem_bytes, + stream=stream, + ) diff --git a/python/sglang/kernels/ops/attention/linear/kda_nvidia_prefill/__init__.py b/python/sglang/kernels/ops/attention/linear/kda_nvidia_prefill/__init__.py new file mode 100644 index 000000000..7b3627194 --- /dev/null +++ b/python/sglang/kernels/ops/attention/linear/kda_nvidia_prefill/__init__.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +# NVIDIA KDA_prefill (Blackwell): optimized chunked KDA forward, an +# FLA-compatible replacement for chunk_kda_fwd. K1 (gate+cumsum+scale, CuTe) +# + K2 (intra sub-chunk, CuTe) + K3 (inter-chunk solve, Triton) + K4 +# (W/U/v_new/O/state update, cuTile persistent). 2.3-2.9x vs the FLA +# reference on B200 in the upstream package tests. +from .chunk_fwd import chunk_kda_fwd + +__all__ = ["chunk_kda_fwd"] diff --git a/python/sglang/kernels/ops/attention/linear/kda_nvidia_prefill/chunk_fwd.py b/python/sglang/kernels/ops/attention/linear/kda_nvidia_prefill/chunk_fwd.py new file mode 100644 index 000000000..de399d33e --- /dev/null +++ b/python/sglang/kernels/ops/attention/linear/kda_nvidia_prefill/chunk_fwd.py @@ -0,0 +1,1146 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +# Vendored from the NVIDIA KDA_prefill package (benchmark/ Blackwell path) +# for the Kimi-K3 chunked prefill forward. Local deltas: fla.* imports +# re-pointed to sglang's vendored fla subset, flat sibling imports made +# package-relative, RCP_LN2 inlined. +# ruff: noqa -- vendored kernel library, minimal local deltas + +""" +KDA Chunk Forward — FLA-Compatible Interface + +Optimized implementation of chunk_kda_fwd using CuTe (K1, K2) + Triton (K3) + cuTile (K4). +Signature matches flash-linear-attention/fla/ops/kda/chunk_fwd.py exactly. + +Supported: + - Equal-length sequences (B ≥ 1, T must be multiple of 64) + - Variable-length sequences (B=1 with cu_seqlens) + - safe_gate mode (sigmoid + lower_bound) and softplus mode + - dt_bias + - use_gate_in_kernel=True with A_log + - chunk_size=64 + +Not supported: + - disable_recompute=True, return_intermediate_states=True + - cp_context (context parallel) + - chunk_size != 64 +""" + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import from_dlpack + +RCP_LN2 = 1.4426950408889634 # fla.ops.utils.constant (1/ln2) +from sglang.kernels.ops.attention.fla.index import ( + prepare_chunk_indices, +) + +try: + from .Akk_inverse_lower_triangle_bf16 import akk_inv_host as _akk_inv_host + from .fuse_k4_only_persistent import BYTES_PER_TENSORMAP as _K4P_BTM + from .fuse_k4_only_persistent import NUM_TENSORMAPS as _K4P_NTM + from .fuse_k4_only_persistent import make_host_fn as _k4p_make_host + from .fuse_kernel123_persistent import CHUNKS_PER_BLOCK as _K123_CHUNKS_PER_BLOCK + from .fuse_kernel123_persistent import make_host_function as _fused_make_host +except ImportError: + from Akk_inverse_lower_triangle_bf16 import akk_inv_host as _akk_inv_host + from fuse_k4_only_persistent import BYTES_PER_TENSORMAP as _K4P_BTM + from fuse_k4_only_persistent import NUM_TENSORMAPS as _K4P_NTM + from fuse_k4_only_persistent import make_host_fn as _k4p_make_host + from fuse_kernel123_persistent import CHUNKS_PER_BLOCK as _K123_CHUNKS_PER_BLOCK + from fuse_kernel123_persistent import make_host_function as _fused_make_host + + +def _ct(t, etype): + """Create a CuTe tensor from PyTorch tensor.""" + r = from_dlpack(t, assumed_align=16) + r.element_type = etype + return r + + +# Cached eqlen dummy cu/ci cute wrappers — shared by K123 and akk_inv. +# Avoids per-call torch.empty + from_dlpack overhead (~10-12us each). +_eqlen_dummy_cache = {} + + +def _get_eqlen_dummies(device, idx_dtype=torch.int64): + """Returns cached (cu_ct, ci_ct) cute wrappers for eqlen (B+1=2, NT+1=2).""" + key = (device.index if device.index is not None else 0, idx_dtype) + if key not in _eqlen_dummy_cache: + cu_t = torch.empty(2, dtype=idx_dtype, device=device) + ci_t = torch.empty(1, 2, dtype=idx_dtype, device=device) + cu_etype = cutlass.Int64 if idx_dtype == torch.int64 else cutlass.Int32 + _eqlen_dummy_cache[key] = (_ct(cu_t, cu_etype), _ct(ci_t, cu_etype)) + return _eqlen_dummy_cache[key] + + +def _cute_int_type(dtype): + """Map PyTorch integer dtype to CUTLASS element type.""" + if dtype == torch.int32: + return cutlass.Int32 + elif dtype == torch.int64: + return cutlass.Int64 + else: + raise ValueError(f"Unsupported integer dtype: {dtype}") + + +# ========== Fused K1+K2+K3 compilation cache ========== +_fused_k123_cache = {} +# id(cu_seqlens) -> bool. Skips per-call GPU->CPU sync on subsequent calls +# when the same cu_seqlens tensor is reused (typical training/inference loop). +_varlen_pure_cache = {} +# id(cu_seqlens) -> int seqlen, populated alongside _varlen_pure_cache for +# single-seq cu_seqlens. +_varlen_single_seqlen_cache = {} + +# id(tensor) -> cute_wrapper. The wrappers themselves are stateless views +# over the tensor's storage, so they remain valid as long as the tensor's +# data pointer / shape / strides don't change. Caller is expected to reuse +# the same tensor objects across iterations (typical PyTorch pattern). +_input_wrap_cache = {} + + +def _tkey(t): + """Serving-safe cache key for a tensor VIEW: python ids are recycled as + activation tensors churn every batch, so id-keyed entries go stale and + hand kernels freed pointers / wrong shapes. (data_ptr, shape, stride, + dtype) pins exactly what a cute wrapper captures.""" + return (t.data_ptr(), tuple(t.shape), tuple(t.stride()), t.dtype) + + +def _ct_cached(t, etype): + """`_ct(t, etype)` with id(t)-based cache. Returns the same cute wrapper + for repeated calls with the same tensor object, avoiding per-call + `from_dlpack` overhead (~5-10us each).""" + key = (_tkey(t), etype) + w = _input_wrap_cache.get(key) + if w is None: + w = _ct(t, etype) + _input_wrap_cache[key] = w + return w + + +# Cache for dt_bias `.float().contiguous().view(H, K)` + cute wrapper. +# dt_bias is typically a nn.Parameter — same object across iterations. +_dt_bias_cache = {} + + +def _get_dt_bias_ct(dt_bias, H, K): + """Returns cached cute wrapper for dt_bias.float().view(H, K).""" + key = (_tkey(dt_bias), H, K) + entry = _dt_bias_cache.get(key) + if entry is None: + bias_t = dt_bias.float().contiguous().view(H, K) + entry = (bias_t, _ct(bias_t, cutlass.Float32)) + _dt_bias_cache[key] = entry + return entry[1] + + +# Cache for the empty 1x1 fp32 bias tensor used when dt_bias is None. +_empty_bias_cache = {} + + +def _get_empty_bias_ct(device): + idx = device.index if device.index is not None else 0 + if idx not in _empty_bias_cache: + t = torch.empty(1, 1, dtype=torch.float32, device=device) + _empty_bias_cache[idx] = _ct(t, cutlass.Float32) + return _empty_bias_cache[idx] + + +# K4 varlen cu_seqlens / chunk_offsets cute wrappers. Caches the int32-cast +# tensor and its mark_layout_dynamic wrapper so they survive multiple calls +# with the same input objects. +_k4_varlen_cu_co_cache = {} + + +def _get_k4_varlen_cu_co(cu_seqlens, chunk_offsets): + key = (_tkey(cu_seqlens), _tkey(chunk_offsets)) + entry = _k4_varlen_cu_co_cache.get(key) + if entry is None: + cu_int32 = ( + cu_seqlens + if cu_seqlens.dtype == torch.int32 + else cu_seqlens.to(torch.int32) + ) + cu_int32 = cu_int32.contiguous() + co_int32 = ( + chunk_offsets + if chunk_offsets.dtype == torch.int32 + else chunk_offsets.to(torch.int32) + ) + co_int32 = co_int32.contiguous() + cu_ct = from_dlpack(cu_int32, assumed_align=4).mark_layout_dynamic() + cu_ct.element_type = cutlass.Int32 + co_ct = from_dlpack(co_int32, assumed_align=4).mark_layout_dynamic() + co_ct.element_type = cutlass.Int32 + # Hold refs to the int32 tensors so they don't get GC'd and the + # underlying storage remain valid as long as cu_ct / co_ct live. + entry = (cu_int32, co_int32, cu_ct, co_ct) + _k4_varlen_cu_co_cache[key] = entry + return entry[2], entry[3] + + +# Cache K4's reshaped v wrapper keyed by id(v_beta) — the v.reshape(-1, H, V) +# view is a fresh Python object every call but the storage is stable when the +# user reuses the same v tensor (typical benchmark / inference pattern). +_v_ct_cache = {} + + +# Cache the (cu_for_k4, chunk_offsets_for_k4) pair keyed by id(cu_seqlens). +# Both tensors are computed on-GPU with no host sync — replaces the previous +# `cu_seqlens.cpu().tolist() + Python loop + torch.tensor` chain that forced +# a 50-200us GPU->CPU stall on the K4 prep path every varlen call. +_varlen_k4_input_cache = {} + + +def _get_varlen_k4_inputs(cu_seqlens, BT): + # Recomputed every call: the outputs depend on cu_seqlens VALUES, which + # change batch-to-batch even when the tensor pointer is recycled, so no + # tensor-identity key is sound. All ops are GPU-side (no host sync). + cu_int32 = ( + cu_seqlens if cu_seqlens.dtype == torch.int32 else cu_seqlens.to(torch.int32) + ) + cu_int32 = cu_int32.contiguous() + seq_lens = cu_int32[1:] - cu_int32[:-1] + chunk_counts = (seq_lens + (BT - 1)) // BT + zero = torch.zeros(1, dtype=torch.int32, device=cu_int32.device) + co_int32 = torch.cat([zero, torch.cumsum(chunk_counts, dim=0).to(torch.int32)]) + co_int32 = co_int32.contiguous() + return cu_int32, co_int32 + + +# ========== BF16 akk_inv compilation cache ========== +_akk_inv_cache = {} +# ========== K4 persistent (varlen via cu_seqlens) compilation cache ========== +_k4p_cache = {} +_k4p_tm_ws = {} + + +# Beta absorption (K1 fusion reverted; K123 stays at 375us baseline). +# Two-stage strategy on side stream: +# 1) v*beta runs PARALLEL with K123 (no data dep on K123 outputs) +# 2) k_scaled*beta runs AFTER K123 on side stream, PARALLEL with akk_inv +@torch.compile(fullgraph=True) +def _v_absorb(v, beta): + return v * beta.unsqueeze(-1).to(v.dtype) + + +@torch.compile(fullgraph=True) +def _ks_absorb(k_scaled, beta): + return k_scaled * beta.unsqueeze(-1).to(k_scaled.dtype) + + +# Side stream cache for v_absorb || K123 overlap. +_side_streams = {} + + +def _get_side_stream(dev): + idx = dev.index or 0 + if idx not in _side_streams: + _side_streams[idx] = torch.cuda.Stream(device=dev) + return _side_streams[idx] + + +# ===== Per-section timing (K123 / K4) ===== +# Set _TIMING_ENABLED=True to record CUDA-event timings for each call. +# Times accumulate in _TIMING_STATS; use reset_timings() / get_timings() to manage. +_TIMING_ENABLED = False +_TIMING_STATS = {"k123_us": 0.0, "k4_us": 0.0, "count": 0} + + +def enable_timing(enabled=True): + """Enable/disable per-section CUDA event timing in chunk_kda_fwd. + Each call adds CUDA event records around K123 and K4 launches and a sync + at end of forward, so this slows down execution. Use only for benchmarking.""" + global _TIMING_ENABLED + _TIMING_ENABLED = enabled + + +def reset_timings(): + _TIMING_STATS["k123_us"] = 0.0 + _TIMING_STATS["k4_us"] = 0.0 + _TIMING_STATS["count"] = 0 + + +def get_timings(): + """Returns (k123_avg_us, k4_avg_us, n_calls). Avgs are over all calls since last reset.""" + n = max(_TIMING_STATS["count"], 1) + return ( + _TIMING_STATS["k123_us"] / n, + _TIMING_STATS["k4_us"] / n, + _TIMING_STATS["count"], + ) + + +# Buffer cache: avoid re-allocating ~67us of intermediate tensors per call. +# Also caches cute.Tensor wrappers (saves ~7us each call from from_dlpack). +_buf_cache = {} + +# Padded-input scratch cache for the eqlen scheduler-alignment path. Keyed by +# (B, T_padded, H, K, dtype_qkv, dtype_g, dtype_beta, device, real_T). +# real_T is part of the key so the g sentinel tail [real_T:T_padded] = -1e3 +# is set once and reused across calls with the same shape. +_padded_input_cache = {} + +# Sentinel-padded g scratch for varlen single-seq Phase 2.1 path. Keyed by +# (B, T_padded, H, K, dtype, device, real_T). The tail [real_T:T_padded] is +# pre-set to -1e3 once at cache init; subsequent calls only overwrite the +# valid prefix [0, real_T). +_g_sentinel_cache = {} + + +def _get_g_sentinel_buffer(B, T_padded, H, K, dtype_g, device, real_T): + key = ( + B, + T_padded, + H, + K, + dtype_g, + device.index if device.index is not None else 0, + real_T, + ) + e = _g_sentinel_cache.get(key) + if e is None: + e = torch.zeros(B, T_padded, H, K, dtype=dtype_g, device=device) + if real_T < T_padded: + e[:, real_T:] = -1000.0 + _g_sentinel_cache[key] = e + return e + + +def _get_padded_input_buffers( + B, T_padded, H, K, dtype_qkv, dtype_g, dtype_beta, device, real_T +): + key = ( + B, + T_padded, + H, + K, + dtype_qkv, + dtype_g, + dtype_beta, + device.index if device.index is not None else 0, + real_T, + ) + e = _padded_input_cache.get(key) + if e is None: + q_pad = torch.zeros(B, T_padded, H, K, dtype=dtype_qkv, device=device) + k_pad = torch.zeros(B, T_padded, H, K, dtype=dtype_qkv, device=device) + v_pad = torch.zeros(B, T_padded, H, K, dtype=dtype_qkv, device=device) + beta_pad = torch.zeros(B, T_padded, H, dtype=dtype_beta, device=device) + # g: zero in [0, real_T), sentinel -1e3 in [real_T, T_padded). Caller's + # data overwrites the prefix each call; the sentinel tail never moves. + g_pad = torch.zeros(B, T_padded, H, K, dtype=dtype_g, device=device) + if real_T < T_padded: + g_pad[:, real_T:] = -1000.0 + e = (q_pad, k_pad, v_pad, g_pad, beta_pad) + _padded_input_cache[key] = e + return e + + +# Multi-seq varlen repack cache for the Phase 2.2 path. Keyed by id(cu_seqlens). +# Stores: (orig_seq_lens, padded_seq_lens, new_cu_seqlens_tensor, +# new_chunk_indices_tensor, new_T_total, padded_input_buffers). +# All tensors are GPU-side and pre-allocated at cache build time. Per-call the +# kernel reads from / writes to these buffers; we copy caller's input slices in +# and output slices back (only the valid prefix of each seq). +_multiseq_repack_cache = {} + +_caller_layout_O_cache = {} + + +def _get_caller_layout_O_buffer(multiseq_info, dtype, V_dim, device): + """Per-shape cached output buffer at caller's contiguous layout (sum of + seq_lens, no padding gaps). Filled by per-seq copies from K4's padded O.""" + caller_T = multiseq_info["caller_T"] + H_x_V = multiseq_info.get("_H_V") # not strictly needed since we fix B=1 H known + key = (caller_T, V_dim, dtype, device.index if device.index is not None else 0) + e = _caller_layout_O_cache.get(key) + if e is None: + # B=1 enforced upstream when multiseq_info is built. + e = torch.empty( + 1, + caller_T, + multiseq_info["q_pad"].shape[2], + V_dim, + dtype=dtype, + device=device, + ) + _caller_layout_O_cache[key] = e + return e + + +def _get_multiseq_repack_info(cu_seqlens, q, k, v, g, beta, BT, device): + """Build (and cache) the padded layout for multi-seq varlen with non-aligned + seqs. Returns None if all seqs are already 64-aligned (caller can use the + existing varlen_pure path).""" + import weakref + + key = id(cu_seqlens) + cached = _multiseq_repack_cache.get(key) + if cached is not None: + wref, e = cached + if wref() is cu_seqlens: + return e + # id collision after GC: rebuild + del _multiseq_repack_cache[key] + cu_cpu = cu_seqlens.cpu().tolist() + seq_lens = [cu_cpu[i + 1] - cu_cpu[i] for i in range(len(cu_cpu) - 1)] + if all(sl % BT == 0 for sl in seq_lens): + _multiseq_repack_cache[key] = (weakref.ref(cu_seqlens), None) + return None + padded_lens = [((sl + BT - 1) // BT) * BT for sl in seq_lens] + new_cu = [0] + for pl in padded_lens: + new_cu.append(new_cu[-1] + pl) + new_T_total = new_cu[-1] + B = q.shape[0] + H = q.shape[2] + K = q.shape[3] + # Pre-allocated padded input buffers. q/k/v/beta tail = 0 (zero MMA), g + # tail = -1e3 sentinel (zero gate activation). The PER-SEQ tail regions + # are between (new_cu[i] + seq_lens[i], new_cu[i+1]) — pre-fill once. + q_pad = torch.zeros(B, new_T_total, H, K, dtype=q.dtype, device=device) + k_pad = torch.zeros_like(q_pad) + v_pad = torch.zeros(B, new_T_total, H, v.shape[3], dtype=v.dtype, device=device) + beta_pad = torch.zeros(B, new_T_total, H, dtype=beta.dtype, device=device) + g_pad = torch.zeros(B, new_T_total, H, K, dtype=g.dtype, device=device) + for i, (sl, pl) in enumerate(zip(seq_lens, padded_lens)): + if sl < pl: + tail_start = new_cu[i] + sl + tail_end = new_cu[i + 1] + g_pad[:, tail_start:tail_end] = -1000.0 + new_cu_tensor = torch.tensor(new_cu, dtype=cu_seqlens.dtype, device=device) + new_chunk_indices = prepare_chunk_indices(new_cu_tensor, BT) + # Build index map: dst_indices[i] = position in padded layout where orig + # row i lives. Used by index_copy_ to do the scatter in one op (instead of + # N_seqs × 5 separate slice copies, which cost ~5us each in Python). + T_total_orig = cu_cpu[-1] + dst_indices_list = [] + for i, sl in enumerate(seq_lens): + for j in range(sl): + dst_indices_list.append(new_cu[i] + j) + dst_indices = torch.tensor(dst_indices_list, dtype=torch.long, device=device) + # Mark this cu_seqlens as VARLEN_PURE eligible — every seq in the new + # layout is 64-aligned by construction. + _varlen_pure_cache[id(new_cu_tensor)] = True + e = { + "seq_lens": seq_lens, + "padded_lens": padded_lens, + "new_cu": new_cu, + "new_T_total": new_T_total, + "new_cu_tensor": new_cu_tensor, + "new_chunk_indices": new_chunk_indices, + "q_pad": q_pad, + "k_pad": k_pad, + "v_pad": v_pad, + "g_pad": g_pad, + "beta_pad": beta_pad, + "dst_indices": dst_indices, + "T_total_orig": T_total_orig, + } + _multiseq_repack_cache[key] = (weakref.ref(cu_seqlens), e) + return e + + +def _get_buffers(dev, dtype_k, B, T, H, K_dim, V_dim, NT, N_seqs, BT): + """All beta fusion lives in akk_inv kernel epilogue (post-inv column-scale).""" + key = (dev.index or 0, B, T, H, K_dim, V_dim, NT, N_seqs) + if key not in _buf_cache: + bf16 = cutlass.BFloat16 + fp32 = cutlass.Float32 + k_scaled = torch.empty( + B, T, H, K_dim, device=dev, dtype=dtype_k + ) # raw, no beta + kg = torch.empty(B, T, H, K_dim, device=dev, dtype=dtype_k) + q_scaled = torch.empty(B, T, H, K_dim, device=dev, dtype=dtype_k) + gk_last_exp = torch.empty(B, NT, H, K_dim, device=dev, dtype=torch.float32) + A_qk = torch.zeros(B, T, H, BT, device=dev, dtype=dtype_k) + A_kk = torch.zeros(B, T, H, BT, device=dev, dtype=dtype_k) + O_flat = torch.empty(B, T, H, V_dim, device=dev, dtype=dtype_k) + # K4 reads initial state from S_out (caller copies it in) and writes + # the final state back into the same buffer — no separate s_4d, no + # extra D2D memcpy at the end of the K4 launcher. Layout matches + # caller's [N_seqs, H, K, V] contig fp32 convention. + S_out = torch.empty(N_seqs, H, K_dim, V_dim, device=dev, dtype=torch.float32) + cu_eqlen = torch.arange(0, (B + 1) * T, T, dtype=torch.int32, device=dev) + co_eqlen = torch.arange( + 0, (B + 1) * (T // BT), T // BT, dtype=torch.int32, device=dev + ) + + T_total = B * T + A_kk_flat = A_kk.reshape(T_total, H, BT) + A_qk_flat = A_qk.reshape(T_total, H, BT) + KS_flat = k_scaled.reshape(T_total, H, K_dim) + QS_flat = q_scaled.reshape(T_total, H, K_dim) + KG_flat = kg.reshape(T_total, H, K_dim) + O_token = O_flat.reshape(T_total, H, V_dim) + gk_flat = gk_last_exp.reshape(-1, H, K_dim) + + def _wrap(t, etype): + r = from_dlpack(t, assumed_align=16).mark_layout_dynamic() + r.element_type = etype + return r + + a_ct = _wrap(A_kk_flat, bf16) + aqc_ct = _wrap(A_qk_flat, bf16) + ks_ct = _wrap(KS_flat, bf16) # raw k_scaled (beta absorbed in akk_inv) + qs_ct = _wrap(QS_flat, bf16) + kg_ct = _wrap(KG_flat, bf16) + o_ct = _wrap(O_token, bf16) + gk_ct = _wrap(gk_flat, fp32) + cu_eqlen_ct = from_dlpack(cu_eqlen, assumed_align=4).mark_layout_dynamic() + cu_eqlen_ct.element_type = cutlass.Int32 + co_eqlen_ct = from_dlpack(co_eqlen, assumed_align=4).mark_layout_dynamic() + co_eqlen_ct.element_type = cutlass.Int32 + + # akk_inv views: bf16 storage reinterpreted as fp32 (packed 2x bf16 -> 1x fp32). + # Built without mark_layout_dynamic to match per-call wrapper type signature + # (akk_inv kernel was compiled against this layout — must stay identical). + akk_in_view = from_dlpack(A_kk, assumed_align=16) + akk_in_view.element_type = fp32 + akk_out_view = from_dlpack(A_kk, assumed_align=16) + akk_out_view.element_type = fp32 + + # K4 state wrapper points directly at S_out — K4 reads/writes in place. + # Plain from_dlpack (no mark_layout_dynamic) keeps the type signature + # stable so the kernel only compiles once. + s_ct = from_dlpack(S_out, assumed_align=16) + s_ct.element_type = fp32 + + # Cache PyTorch streams once per buf_cache entry — torch.cuda.current_stream + # is a torch C call that costs ~2us each invocation. + main_stream_cached = torch.cuda.current_stream(dev) + side_stream_cached = _get_side_stream(dev) + + cute_wrappers = dict( + a_ct=a_ct, + aqc_ct=aqc_ct, + ks_ct=ks_ct, + qs_ct=qs_ct, + kg_ct=kg_ct, + o_ct=o_ct, + gk_ct=gk_ct, + cu_eqlen_ct=cu_eqlen_ct, + co_eqlen_ct=co_eqlen_ct, + akk_in_view=akk_in_view, + akk_out_view=akk_out_view, + s_ct=s_ct, + main_stream=main_stream_cached, + side_stream=side_stream_cached, + # Filled lazily on first launch — saves cache_key tuple build + + # outer dict lookup on subsequent calls. + _k123_fns={}, + _akk_inv_fn=None, + _k4_fn=None, + ) + + _buf_cache[key] = ( + k_scaled, + kg, + q_scaled, + gk_last_exp, + A_qk, + A_kk, + O_flat, + S_out, + cu_eqlen, + co_eqlen, + cute_wrappers, + ) + return _buf_cache[key] + + +def _launch_k4_persistent( + cute_wrappers, + v_beta, + S_in, + S_out, + cu_seqlens, + chunk_offsets, + cu_eqlen_passed=False, + num_sm=148, + H=None, + V_dim=None, + use_fast_sync=False, +): + """Persistent K4 with cached cute wrappers. + k_scaled (cached, beta-absorbed by K1 fusion) reads from cute_wrappers['ks_ct']. + v_beta: fresh per-call tensor (torch.compile output of v*beta on side stream).""" + # Launch on the CALLER's current stream (resolved per call, passed as a + # runtime arg): with sglang's overlap scheduler the forward runs on a + # non-default stream and a default-stream K4 races the whole pipeline. + launch_stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + # Fast path: cache the pre-built (k4_fn, args_tuple) keyed by + # (id(v_beta), id(cu_seqlens), id(S_in), id(S_out)). The state-copy is + # done by the caller (chunk_kda_fwd) on a side stream to overlap with + # K123 — this launcher just runs K4. + fast_key = ( + _tkey(v_beta), + 0 if cu_eqlen_passed else _tkey(cu_seqlens), + _tkey(S_in), + _tkey(S_out), + ) + fast_cache = cute_wrappers.get("_k4_fast_cache") + if fast_cache is None: + fast_cache = {} + cute_wrappers["_k4_fast_cache"] = fast_cache + fast_entry = fast_cache.get(fast_key) + if fast_entry is not None: + k4_fn, args = fast_entry + k4_fn(*args, launch_stream) + return + + bf16 = cutlass.BFloat16 + fp32 = cutlass.Float32 + N_seqs = cu_seqlens.shape[0] - 1 + BH = N_seqs * H + nsm = min(BH, num_sm) + dev = v_beta.device + dev_idx = dev.index or 0 + + s_ct = cute_wrappers["s_ct"] + a_ct = cute_wrappers["a_ct"] + b_ct = cute_wrappers["ks_ct"] # raw k_scaled (beta absorbed in akk_inv) + q_ct = cute_wrappers["qs_ct"] + aqc_ct = cute_wrappers["aqc_ct"] + kg_ct = cute_wrappers["kg_ct"] + o_ct = cute_wrappers["o_ct"] + gk_ct = cute_wrappers["gk_ct"] + + v_key = _tkey(v_beta) + v_entry = _v_ct_cache.get(v_key) + if v_entry is None: + v_view = v_beta.reshape(-1, H, V_dim) if v_beta.dim() == 4 else v_beta + v_ct = from_dlpack(v_view, assumed_align=16).mark_layout_dynamic() + v_ct.element_type = bf16 + _v_ct_cache[v_key] = (v_view, v_ct) + else: + v_ct = v_entry[1] + + if cu_eqlen_passed: + cu_ct = cute_wrappers["cu_eqlen_ct"] + co_ct = cute_wrappers["co_eqlen_ct"] + else: + cu_ct, co_ct = _get_k4_varlen_cu_co(cu_seqlens, chunk_offsets) + + tm_key = (dev_idx, nsm) + if tm_key not in _k4p_tm_ws: + tm_ws_t = torch.zeros(nsm * _K4P_NTM * _K4P_BTM, dtype=torch.uint8, device=dev) + tm_ct = from_dlpack(tm_ws_t, assumed_align=16) + tm_ct.element_type = cutlass.Uint8 + _k4p_tm_ws[tm_key] = (tm_ws_t, tm_ct) + else: + tm_ws_t, tm_ct = _k4p_tm_ws[tm_key] + + k4_fn = cute_wrappers.get("_k4_fn") + if k4_fn is None: + cache_key = (dev_idx, nsm, N_seqs, H) + k4_fn = _k4p_cache.get(cache_key) + if k4_fn is None: + host_fn = _k4p_make_host(num_sm=nsm) + k4_fn = cute.compile( + host_fn, + a_ct, + b_ct, + v_ct, + q_ct, + aqc_ct, + kg_ct, + o_ct, + gk_ct, + s_ct, + cu_ct, + co_ct, + tm_ct, + launch_stream, + ) + _k4p_cache[cache_key] = k4_fn + cute_wrappers["_k4_fn"] = k4_fn + + args = ( + a_ct, + b_ct, + v_ct, + q_ct, + aqc_ct, + kg_ct, + o_ct, + gk_ct, + s_ct, + cu_ct, + co_ct, + tm_ct, + ) + fast_cache[fast_key] = (k4_fn, args) + k4_fn(*args, launch_stream) + + +def _launch_fused_k123_inv( + q, + k, + g, + A_log, + beta, + scale, + k_scaled, + kg, + q_scaled, + gk_last_exp, + A_qk, + A_kk_inv, + cu_seqlens, + chunk_indices, + is_varlen, + NT, + dt_bias=None, + safe_gate=False, + lower_bound=None, + akk_in_view=None, + akk_out_view=None, + cute_wrappers=None, + varlen_pure_override=None, +): + """Persistent K1+K2 (writes A_kk in I+L format with diag=1) chained with + BF16 akk_inv (in-place inversion). Final A_kk_inv = (I+L)^-1.""" + + # Resolve the CALLER's stream every call: under sglang's overlap scheduler + # the forward runs on a non-default stream, and launching these kernels on + # the default stream makes the whole pipeline a cross-stream data race + # (torn states in the mamba pool, garbage outputs). Streams are passed as + # runtime launch args, so the compiled fns stay cached. + launch_stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + + # Fast path: when (q, k, g, beta, A_log, dt_bias, cu_seqlens, chunk_indices) + # are stable across calls (typical benchmark / inference), skip the per-call + # wrapper gathering and launch with a pre-built (k123_fn, args) tuple. + if cute_wrappers is not None: + fast_key = ( + _tkey(q), + _tkey(k), + _tkey(g), + _tkey(beta), + _tkey(A_log), + _tkey(dt_bias) if dt_bias is not None else 0, + _tkey(cu_seqlens) if cu_seqlens is not None else 0, + _tkey(chunk_indices) if chunk_indices is not None else 0, + bool(safe_gate), + float(lower_bound) if lower_bound is not None else 0.0, + ) + fast_cache = cute_wrappers.setdefault("_k123_fast_cache", {}) + fast_entry = fast_cache.get(fast_key) + if fast_entry is not None: + k123_fn, k123_args, akk_fn, akk_args = fast_entry + k123_fn(*k123_args, launch_stream) + akk_fn(*akk_args, launch_stream) + return + + B, T, H, K = q.shape + BT = 64 + dev = q.device.index or 0 + T_padded = T if is_varlen else None + has_bias = dt_bias is not None + + # Auto-detect VARLEN_PURE eligibility, cached by id(cu_seqlens). + # First call with a given cu_seqlens object pays one GPU->CPU sync; later + # calls with the same tensor object are a dict lookup (~100 ns). + varlen_pure = False + if varlen_pure_override is not None: + # Serving path: the caller knows the seq lengths (CPU-side, no sync) + # and the id-keyed detection below is unsafe under tensor-id reuse. + varlen_pure = bool(varlen_pure_override) if is_varlen else False + elif is_varlen and cu_seqlens is not None: + _vp_key = id(cu_seqlens) + if _vp_key not in _varlen_pure_cache: + cu_cpu = cu_seqlens.cpu().tolist() + seq_lens = [cu_cpu[i + 1] - cu_cpu[i] for i in range(len(cu_cpu) - 1)] + _varlen_pure_cache[_vp_key] = all((sl % BT) == 0 for sl in seq_lens) + varlen_pure = _varlen_pure_cache[_vp_key] + cache_key = (B, NT, H, is_varlen, T_padded, dev, has_bias, safe_gate, varlen_pure) + + # Inputs are guaranteed contiguous by upstream linear projections. + # A_log is fp32 model param; .float() is no-op when dtype already matches. + q_ct = _ct_cached(q, cutlass.BFloat16) + k_ct = _ct_cached(k, cutlass.BFloat16) + g_ct = _ct_cached(g, cutlass.BFloat16) + alog_ct = _ct_cached( + A_log if A_log.dtype == torch.float32 else A_log.float(), cutlass.Float32 + ) + beta_ct = _ct_cached(beta, cutlass.BFloat16) + + ks_ct = _ct_cached(k_scaled, cutlass.BFloat16) + kg_ct = _ct_cached(kg, cutlass.BFloat16) + qs_ct = _ct_cached(q_scaled, cutlass.BFloat16) + gk_ct = _ct_cached(gk_last_exp, cutlass.Float32) + aqk_ct = _ct_cached(A_qk, cutlass.BFloat16) + akk_ct = _ct_cached(A_kk_inv, cutlass.BFloat16) + + if is_varlen: + cu_ct = _ct_cached(cu_seqlens, _cute_int_type(cu_seqlens.dtype)) + ci_ct = _ct_cached(chunk_indices, _cute_int_type(chunk_indices.dtype)) + else: + cu_ct, ci_ct = _get_eqlen_dummies(q.device, torch.int64) + + if dt_bias is not None: + bias_ct = _get_dt_bias_ct(dt_bias, H, K) + else: + bias_ct = _get_empty_bias_ct(q.device) + lb_val = float(lower_bound) if lower_bound is not None else 0.0 + + ct_args = ( + q_ct, + k_ct, + g_ct, + alog_ct, + beta_ct, + scale, + ks_ct, + kg_ct, + qs_ct, + gk_ct, + aqk_ct, + akk_ct, + cu_ct, + ci_ct, + bias_ct, + lb_val, + ) + + if cache_key not in _fused_k123_cache: + host_fn = _fused_make_host( + B, + NT, + H, + is_varlen=is_varlen, + T_padded=T_padded, + has_bias=has_bias, + use_safe_gate=safe_gate, + varlen_pure=varlen_pure, + ) + _fused_k123_cache[cache_key] = cute.compile(host_fn, *ct_args, launch_stream) + k123_fn = _fused_k123_cache[cache_key] + k123_fn(*ct_args, launch_stream) + + # ===== Chained BF16 akk_inv (in-place: A_kk_inv = (I+L)^-1) ===== + if akk_in_view is None: + akk_in_view = from_dlpack(A_kk_inv, assumed_align=16) + akk_in_view.element_type = cutlass.Float32 + akk_out_view = from_dlpack(A_kk_inv, assumed_align=16) + akk_out_view.element_type = cutlass.Float32 + + if is_varlen: + # Reuse the cached cute wrappers from K123 (same tensor objects). + akk_cu_ct = cu_ct + akk_ci_ct = ci_ct + is_varlen_int = 1 + T_val = T + else: + akk_cu_ct, akk_ci_ct = cu_ct, ci_ct + is_varlen_int = 0 + T_val = NT * BT + + akk_cache_key = (B, NT, H, is_varlen, dev, T_val) + if akk_cache_key not in _akk_inv_cache: + _akk_inv_cache[akk_cache_key] = cute.compile( + _akk_inv_host, + akk_in_view, + akk_out_view, + beta_ct, + B, + NT, + H, + akk_cu_ct, + akk_ci_ct, + is_varlen_int, + T_val, + launch_stream, + ) + akk_fn = _akk_inv_cache[akk_cache_key] + akk_args = (akk_in_view, akk_out_view, beta_ct, akk_cu_ct, akk_ci_ct) + akk_fn(*akk_args, launch_stream) + + # Populate the fast cache so subsequent calls with the same input ids skip + # all wrapper gathering above. + if cute_wrappers is not None: + cute_wrappers["_k123_fast_cache"][fast_key] = ( + k123_fn, + ct_args, + akk_fn, + akk_args, + ) + + +def chunk_kda_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + chunk_size: int = 64, + safe_gate: bool = False, + lower_bound: float | None = None, + use_gate_in_kernel: bool = False, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + disable_recompute: bool = False, + return_intermediate_states: bool = False, + cp_context=None, + varlen_single_real_T: int | None = None, + varlen_pure: bool | None = None, +): + """KDA forward — optimized, FLA-compatible interface.""" + if safe_gate and lower_bound is None: + lower_bound = -5.0 + + is_varlen = cu_seqlens is not None + + B, T, H, K = q.shape + V_dim = v.shape[-1] + device = q.device + BT = 64 + + # Phase 1: the eqlen persistent scheduler assigns CHUNKS_PER_BLOCK chunks + # to each workgroup, so pad T to that scheduling unit even when T is + # already a 64-row chunk multiple. This prevents zero-workgroup launches + # for short inputs and avoids silently dropping a trailing chunk group. + # K123 eqlen runs unchanged and sees only full 64-row chunks; zero/sentinel + # padding handles the caller's boundary at the data level. + # + # Varlen with non-aligned seq lengths is a separate problem (Phase 2): + # multi-seq varlen can't be host-padded without repacking memory. + real_T = T + CPB_BT = _K123_CHUNKS_PER_BLOCK * BT + needs_eqlen_pad = (not is_varlen) and (T % CPB_BT != 0) + if needs_eqlen_pad: + if B != 1 and T % BT != 0: + raise NotImplementedError( + f"eqlen with B>1 and T % {BT} != 0 not supported " + f"(got B={B}, T={T})." + ) + T_padded = ((T + CPB_BT - 1) // CPB_BT) * CPB_BT + # Pre-allocated padded scratch buffers (per (B,T_padded,H,K,dtype) cache + # key). torch.cat would reallocate + copy the full 200MB q tensor every + # call — caching the destination buffer drops that to a single slice + # copy of the valid prefix (caller already lives in our buffer for + # subsequent calls reusing the same id, but we re-copy unconditionally + # since the caller may have updated the data in-place). + q_pad, k_pad, v_pad, g_pad, beta_pad = _get_padded_input_buffers( + B, T_padded, H, K, q.dtype, g.dtype, beta.dtype, q.device, real_T + ) + # q/k/v/beta zero-padded → K1/K2 MMAs naturally produce 0 for OOB rows. + # Tail [real_T:] of q_pad/k_pad/v_pad/beta_pad is pre-zeroed at cache + # init and never written, so we only copy the valid prefix. + q_pad[:, :real_T].copy_(q) + k_pad[:, :real_T].copy_(k) + v_pad[:, :real_T].copy_(v) + beta_pad[:, :real_T].copy_(beta) + # g uses a -1e3 sentinel so the gate activation saturates to 0 for OOB + # rows (both safe_gate sigmoid and softplus paths). Plain g=0 gives + # nonzero activation that would corrupt the cumsum past seq end. The + # tail is set to -1e3 once at cache init; we only copy the valid prefix. + g_pad[:, :real_T].copy_(g) + q, k, v, g, beta = q_pad, k_pad, v_pad, g_pad, beta_pad + T = T_padded # downstream buffer alloc + kernel layout use T_padded + + # Phase 2.1: varlen with a SINGLE non-aligned sequence — caller already + # zero-padded q/k/v/beta to a 64-multiple (FLA convention), but g's tail + # is also zero, which causes the gate activation to be non-zero past seq + # end and corrupts the cumsum / GkLast. We sentinel-pad g (cheap: ~5MB + # copy) and force VARLEN_PURE=1 so all 4 mask sites compile-elide. Same + # K4 "boundary at the data" principle as the eqlen path. + # + # Multi-seq varlen is NOT handled here — its OOB regions overlap with + # adjacent seqs' data, so sentinel-pad on g would corrupt the next seq. + # Multi-seq optimization needs per-seq dynamic tensormap (Phase 2.2). + needs_varlen_single_pad = False + if ( + varlen_single_real_T is not None + and is_varlen + and cu_seqlens is not None + and cu_seqlens.shape[0] == 2 + and B == 1 + ): + # Caller-supplied real length (serving path): skips the id-keyed + # alignment cache below, which is unsafe under tensor-id reuse when + # cu_seqlens churns every batch. + if varlen_single_real_T % BT != 0: + real_T = varlen_single_real_T + needs_varlen_single_pad = True + elif is_varlen and cu_seqlens is not None and cu_seqlens.shape[0] == 2 and B == 1: + _vl_key = id(cu_seqlens) + if _vl_key not in _varlen_pure_cache: + cu_cpu = cu_seqlens.cpu().tolist() + sl = cu_cpu[1] - cu_cpu[0] + _varlen_pure_cache[_vl_key] = sl % BT == 0 + _varlen_single_seqlen_cache[_vl_key] = sl + if not _varlen_pure_cache[_vl_key]: + real_T = _varlen_single_seqlen_cache[_vl_key] + needs_varlen_single_pad = True + if needs_varlen_single_pad: + # q/k/v/beta already zero-padded by caller (FLA convention). Re-build g + # with -1000 sentinel in the tail so VARLEN_PURE=1 path is correct. + # Cache the resulting g buffer so repeated calls with same input ids + # don't re-allocate. + cur_T = q.shape[1] # caller's padded T + g_pad = _get_g_sentinel_buffer(B, cur_T, H, K, g.dtype, g.device, real_T) + g_pad[:, :real_T].copy_(g[:, :real_T]) + g = g_pad + # Force VARLEN_PURE=1 — caller's cu_seqlens is reused as-is. + _varlen_pure_cache[id(cu_seqlens)] = True + + # Phase 2.2 (multi-seq via host repack) was attempted but isn't net positive: + # the scatter/gather memcpy cost (~800us GPU bandwidth per call) exceeds + # the kernel mask-elision savings (~250us). Keeping multi-seq non-pure on + # the original masked path. The right fix is kernel-level dynamic + # tensormap (K4-style per-tile bounded TMA) but that requires a major + # K123 kernel refactor — left as future work. + multiseq_info = None + + if is_varlen: + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = len(chunk_indices) + N_seqs = len(cu_seqlens) - 1 + else: + NT = T // BT + N_seqs = B + + # ===== Cached buffers + cute wrappers (avoid alloc + from_dlpack overhead per call) ===== + ( + k_scaled, + kg, + q_scaled, + gk_last_exp, + A_qk, + A_kk, + O_flat, + S_out, + cu_eqlen, + co_eqlen, + cute_wrappers, + ) = _get_buffers(device, k.dtype, B, T, H, K, V_dim, NT, N_seqs, BT) + + # ===== State copy on side stream, parallel with K123 ===== + # K4 needs S_out populated with initial_state. By doing this copy on a + # side stream BEFORE K123 launches, the D2D memcpy overlaps with K123's + # compute. Especially big win for high-N_seqs varlen where state is huge + # (192MB for N=32, ~64us memcpy) — would otherwise serialize before K4. + if initial_state is None: + S_in = torch.zeros(N_seqs, H, K, V_dim, dtype=torch.float32, device=device) + else: + S_in = initial_state + # Serving hardening: the state copy runs on the CALLER's stream, same as + # every kernel in the pipeline. The old side-stream overlap (copy || K123) + # only pays off for huge multi-seq varlen states (~192MB at N=32); for + # serving single-seq engagements the state is <1MB and the side stream + # was the last remaining cross-stream surface in this path. + main_stream = torch.cuda.current_stream(device) + needs_copy = S_in.data_ptr() != S_out.data_ptr() + if needs_copy: + S_out.copy_(S_in) + + # Beta is fused entirely in akk_inv kernel epilogue (post-inv column-scale). + # No host v*beta and no K1 k_scaled*beta any more. + if _TIMING_ENABLED: + k123_s = torch.cuda.Event(enable_timing=True) + k123_e = torch.cuda.Event(enable_timing=True) + k123_s.record(stream=main_stream) + + _launch_fused_k123_inv( + q, + k, + g, + A_log, + beta, + scale, + k_scaled, + kg, + q_scaled, + gk_last_exp, + A_qk, + A_kk, + cu_seqlens, + chunk_indices, + is_varlen, + NT, + dt_bias=dt_bias, + safe_gate=safe_gate, + lower_bound=lower_bound, + akk_in_view=cute_wrappers["akk_in_view"], + akk_out_view=cute_wrappers["akk_out_view"], + cute_wrappers=cute_wrappers, + varlen_pure_override=(True if needs_varlen_single_pad else varlen_pure), + ) + + if _TIMING_ENABLED: + k123_e.record(stream=main_stream) + + # ===== K4: persistent kernel (eqlen + varlen via cu_seqlens) ===== + if is_varlen: + # GPU-side cumsum + cache by id(cu_seqlens). No host sync. + cu_for_k4, chunk_offsets_for_k4 = _get_varlen_k4_inputs(cu_seqlens, BT) + else: + cu_for_k4 = cu_eqlen + chunk_offsets_for_k4 = co_eqlen + + # State copy now runs on the caller's stream; K4 is stream-ordered after it. + + if _TIMING_ENABLED: + k4_s = torch.cuda.Event(enable_timing=True) + k4_e = torch.cuda.Event(enable_timing=True) + k4_s.record(stream=main_stream) + + _launch_k4_persistent( + cute_wrappers, + v, + S_in, + S_out, + cu_for_k4, + chunk_offsets_for_k4, + cu_eqlen_passed=(not is_varlen), + H=H, + V_dim=V_dim, + use_fast_sync=(not is_varlen), + ) + + if _TIMING_ENABLED: + k4_e.record(stream=cute_wrappers["main_stream"]) + torch.cuda.synchronize(device) + _TIMING_STATS["k123_us"] += k123_s.elapsed_time(k123_e) * 1000.0 + _TIMING_STATS["k4_us"] += k4_s.elapsed_time(k4_e) * 1000.0 + _TIMING_STATS["count"] += 1 + + o = O_flat + A_qk_out = A_qk + A_kk_out = A_kk + if needs_eqlen_pad: + # Caller called with original T = real_T; their downstream code expects + # outputs at that shape. Slice the padded scratch tail back off. + o = o[:, :real_T] + A_qk_out = A_qk_out[:, :real_T] + A_kk_out = A_kk_out[:, :real_T] + # multiseq_info is always None here (Phase 2.2 disabled — see above) + final_state = S_out if output_final_state else None + + return ( + o, + final_state, + None, + A_qk_out, + A_kk_out, + None, + None, + None, + None, + None, + None, + initial_state, + ) diff --git a/python/sglang/kernels/ops/attention/linear/kda_nvidia_prefill/fuse_k4_only_persistent.py b/python/sglang/kernels/ops/attention/linear/kda_nvidia_prefill/fuse_k4_only_persistent.py new file mode 100644 index 000000000..289049e7e --- /dev/null +++ b/python/sglang/kernels/ops/attention/linear/kda_nvidia_prefill/fuse_k4_only_persistent.py @@ -0,0 +1,1295 @@ +# Vendored from the NVIDIA KDA_prefill package (benchmark/ Blackwell path) +# for the Kimi-K3 chunked prefill forward. Local deltas: fla.* imports +# re-pointed to sglang's vendored fla subset, flat sibling imports made +# package-relative, RCP_LN2 inlined. +# ruff: noqa -- vendored kernel library, minimal local deltas +"""K4-only fused KDA kernel: persistent variant with varlen support. + +Persistent scheduler (GDNTileScheduler) with 3D tensor layout, TensorMapManager +for per-tile TMA descriptor updates, and domain_offset+flat_divide for per-chunk +addressing. Supports variable-length sequences via cu_seqlens. + +K4 chunk loop with 6 MMAs per chunk: + MMA1: W = AB @ KS (K-MN, K=64) + MMA2: U = AB @ V (K-MN, K=64) + MMA3: NV = U + W_bf16 @ S (SS-mode: A=SMEM sO, B=SMEM sST, K=128, accumulate) + MMA4: OI = QS @ S (K-MN, K=128) + MMA5: O = OI + AQC @ NV (K-MN, K=64, accumulate) + MMA6: State += NV^T @ KG (MN-MN, K=64, accumulate on decayed state) + +Execution order: MMA1->MMA2->MMA4->MMA3->MMA5->MMA6 + +Warp assignment (3 warpgroups, 12 warps, 384 threads): + WG0 (W0-3): W0=MMA issue, W2=TMA load/store, W1/W3=idle + WG1 (W4-7): State 2-pass: TMEM->bf16->sST, gk(SMEM)->decay->TMEM + WG2 (W8-11): GDN readout: W/NV/O TMEM->bf16->SMEM + +TMA pipelines (all PipelineTmaUmma, 1-stage): + - KS, V, QS: prefetch c+1 while MMA uses c + - AB, AQC, KG: 1-stage with explicit consumer release (no prefetch) + - TMA warp decoupled from O readout (no store_nbar wait) + +State management: + - tCtState [128,128] fp32 @ TMEM offset 256: persistent + - 2-pass per chunk: bf16->sST early, gk decay->TMEM late + - gk_last preloaded to SMEM (coalesced), read in Pass 2 + - single state_ready_mbar after both passes (TMEM no-overlap constraint) +""" + +import os +import sys +import warnings + +import cutlass.cutlass_dsl as _dsl_mod + +if not hasattr(_dsl_mod, "CuteExperimentalDSL"): + + class _DummyExperimentalDSL: + jit = None + kernel = None + compile = None + + _dsl_mod.CuteExperimentalDSL = _DummyExperimentalDSL + +import cuda.bindings.driver as cuda_drv +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.cute import KeepCUBIN, PtxasOptions +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.nvgpu.tcgen05 import Field, OperandMajorMode, OperandSource +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import Int32, T + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +# GDN tile scheduler: import from the public FlashInfer release +# (`pip install flashinfer-python>=0.6.13`). Fall back to a local +# dynamic-kernel-generator checkout only if FlashInfer is unavailable. +try: + from flashinfer.gdn_kernels.blackwell.gated_delta_net_tile_scheduler import ( + GDNTileScheduler, + GDNTileSchedulerParams, + ) +except ImportError: + # Try multiple locations for GDNTileScheduler (workspace layout vs kda_optimized layout) + _GDN_PATHS = [ + # workspace/flash-linear-attention/prefill -> ../../../dynamic-kernel-generator + os.path.join( + os.path.dirname(os.path.abspath(__file__)), + "..", + "..", + "..", + "dynamic-kernel-generator", + "cutlass_ir", + "compiler", + "python", + "examples", + "blackwell", + "gated_delta_net", + ), + # scripts/kda_optimized -> ../../dynamic-kernel-generator + os.path.join( + os.path.dirname(os.path.abspath(__file__)), + "..", + "..", + "dynamic-kernel-generator", + "cutlass_ir", + "compiler", + "python", + "examples", + "blackwell", + "gated_delta_net", + ), + ] + for _p in _GDN_PATHS: + if os.path.isdir(_p): + sys.path.insert(0, _p) + break + from gated_delta_net_tile_scheduler import GDNTileSchedulerParams, GDNTileScheduler + +from cutlass.utils import TensorMapManager, TensorMapUpdateMode + + +def transform_partitioned_tensor_layout(tensor): + layout = tensor.layout + stored_layout = layout + if isinstance(stored_layout, cute.ComposedLayout): + layout = layout.outer + shape = layout.shape + stride = layout.stride + new_shape = ((shape[0][0], shape[1]), (shape[0][1], shape[2]), *shape[3:]) + new_stride = ((stride[0][0], stride[1]), (stride[0][1], stride[2]), *stride[3:]) + new_layout = cute.make_layout(shape=new_shape, stride=new_stride) + if isinstance(stored_layout, cute.ComposedLayout): + new_layout = cute.make_composed_layout( + stored_layout.inner, stored_layout.offset, new_layout + ) + return cute.make_tensor(tensor.iterator, new_layout) + + +mma_dtype = cutlass.BFloat16 +acc_dtype = cutlass.Float32 +out_dtype = cutlass.BFloat16 + +M = 64 +N = 128 +K = 64 +K3 = 128 +M6 = 128 +N6 = 128 +K6 = 64 + +threads_per_cta = 384 +warp_threads = 32 +warpgroup_threads = 128 + +BYTES_PER_TENSORMAP = 128 +NUM_TENSORMAPS = 7 # a, b, v, q, aqc, kg, o + +MMA_WARP = 0 +O_STORE_WARP = 1 +TMA_WARP = 2 +STATE_WG = 1 +READOUT_WG = 2 + +MAX_REGS = 168 + +try: + from cutlass.cutlass_dsl.cutlass import CuTeDSL as _CuTeDSL + + _orig_get_pipeline = _CuTeDSL._get_pipeline + _patch_applied = False + + def _patched_get_pipeline(self, _pipeline): + global _patch_applied + result = _orig_get_pipeline(self, _pipeline) + if result and "ptx-options=" not in result: + if "cubin-format=bin" in result: + result = result.replace( + "cubin-format=bin", "cubin-format=bin ptx-options='--uumn'" + ) + _patch_applied = True + else: + warnings.warn( + "kda_nvidia_prefill: cubin-format=bin not found in the CuTeDSL " + "pipeline, ptx-options='--uumn' was not injected" + ) + elif result and "ptx-options=" in result: + _patch_applied = True + return result + + _CuTeDSL._get_pipeline = _patched_get_pipeline +except Exception as e: + warnings.warn(f"kda_nvidia_prefill: CuTeDSL ptx-options patch failed: {e}") + _patch_applied = False + + +@cute.kernel +def k4_persistent_kernel( + tiled_mma_kmn: cute.TiledMma, + tiled_mma_mn_mn: cute.TiledMma, + tma_a, + a_sl: cute.ComposedLayout, + tma_b, + b_sl: cute.ComposedLayout, + tma_v, + v_sl: cute.ComposedLayout, + s_sl: cute.ComposedLayout, + tma_q, + q_sl: cute.ComposedLayout, + tma_aqc, + aqc_sl: cute.ComposedLayout, + tma_kg, + kg_sl: cute.ComposedLayout, + tma_o, + store_sl: cute.ComposedLayout, + readout_k_sl: cute.ComposedLayout, + nv_b_sl: cute.ComposedLayout, + nv_a_sl: cute.ComposedLayout, + kg_a_sl: cute.ComposedLayout, + nv_b_mn_sl: cute.ComposedLayout, + mGkLastExp: cute.Tensor, + mS_fp32: cute.Tensor, + cu_seqlens: cute.Tensor, + chunk_offsets: cute.Tensor, + mA: cute.Tensor, + mB: cute.Tensor, + mV_g: cute.Tensor, + mQ: cute.Tensor, + mAQC: cute.Tensor, + mKG: cute.Tensor, + mO: cute.Tensor, + tensormap_workspace: cute.Tensor, + scheduler_params: GDNTileSchedulerParams, +): + bidx, bidy, bidz = cute.arch.block_idx() + grid_dim = cute.arch.grid_dim() + tidx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.make_warp_uniform(tidx // warp_threads) + warpgroup_idx = cute.arch.make_warp_uniform(tidx // warpgroup_threads) + warpgroup_tidx = tidx % warpgroup_threads + lane_id = tidx % warp_threads + thr_kmn = tiled_mma_kmn.get_slice(0) + thr_mn = tiled_mma_mn_mn.get_slice(0) + dice = (None, None, None) + + cta_linear_idx = bidz * grid_dim[1] * grid_dim[0] + bidy * grid_dim[0] + bidx + tensormap_manager = TensorMapManager(TensorMapUpdateMode.GMEM, BYTES_PER_TENSORMAP) + tm_ws = cute.make_tensor( + tensormap_workspace.iterator, + cute.make_layout( + ( + grid_dim[0] * grid_dim[1] * grid_dim[2], + NUM_TENSORMAPS, + BYTES_PER_TENSORMAP, + ), + stride=(NUM_TENSORMAPS * BYTES_PER_TENSORMAP, BYTES_PER_TENSORMAP, 1), + ), + ) + tm_a_ptr = tensormap_manager.get_tensormap_ptr( + tm_ws[(cta_linear_idx, 0, None)].iterator + ) + tm_b_ptr = tensormap_manager.get_tensormap_ptr( + tm_ws[(cta_linear_idx, 1, None)].iterator + ) + tm_v_ptr = tensormap_manager.get_tensormap_ptr( + tm_ws[(cta_linear_idx, 2, None)].iterator + ) + tm_q_ptr = tensormap_manager.get_tensormap_ptr( + tm_ws[(cta_linear_idx, 3, None)].iterator + ) + tm_aqc_ptr = tensormap_manager.get_tensormap_ptr( + tm_ws[(cta_linear_idx, 4, None)].iterator + ) + tm_kg_ptr = tensormap_manager.get_tensormap_ptr( + tm_ws[(cta_linear_idx, 5, None)].iterator + ) + tm_o_ptr = tensormap_manager.get_tensormap_ptr( + tm_ws[(cta_linear_idx, 6, None)].iterator + ) + + smem = cutlass.utils.SmemAllocator() + AL = 128 + sA = smem.allocate_tensor(mma_dtype, a_sl.outer, AL, a_sl.inner) + sB = smem.allocate_tensor(mma_dtype, b_sl.outer, AL, b_sl.inner) + sV = smem.allocate_tensor(mma_dtype, v_sl.outer, AL, v_sl.inner) + sST = smem.allocate_tensor(mma_dtype, s_sl.outer, AL, s_sl.inner) + sQ = smem.allocate_tensor(mma_dtype, q_sl.outer, AL, q_sl.inner) + sAQC = smem.allocate_tensor(mma_dtype, aqc_sl.outer, AL, aqc_sl.inner) + sKG = smem.allocate_tensor(mma_dtype, kg_sl.outer, AL, kg_sl.inner) + sNV = smem.allocate_tensor(mma_dtype, readout_k_sl.outer, AL, readout_k_sl.inner) + sO = smem.allocate_tensor(mma_dtype, readout_k_sl.outer, AL, readout_k_sl.inner) + sO_out = smem.allocate_tensor(mma_dtype, readout_k_sl.outer, AL, readout_k_sl.inner) + sGk_buf = smem.allocate_array(acc_dtype, N) + + sNV_b = cute.make_tensor( + cute.recast_ptr(sNV.iterator, nv_b_sl.inner, mma_dtype), nv_b_sl.outer + ) + sNV_a = cute.make_tensor( + cute.recast_ptr(sNV.iterator, nv_a_sl.inner, mma_dtype), nv_a_sl.outer + ) + sKG_a = cute.make_tensor( + cute.recast_ptr(sKG.iterator, kg_a_sl.inner, mma_dtype), kg_a_sl.outer + ) + sNV_b_mn = cute.make_tensor( + cute.recast_ptr(sNV.iterator, nv_b_mn_sl.inner, mma_dtype), nv_b_mn_sl.outer + ) + sO_st = cute.make_tensor( + cute.recast_ptr(sO.iterator, store_sl.inner, out_dtype), store_sl.outer + ) + sO_out_st = cute.make_tensor( + cute.recast_ptr(sO_out.iterator, store_sl.inner, out_dtype), store_sl.outer + ) + + tmem_smem = smem.allocate_array(cutlass.Int32, 1) + if warp_idx == 0: + cute.arch.alloc_tmem(512, tmem_smem) + + sNV_ready_nbar = pipeline.NamedBarrier(3, warpgroup_threads + warp_threads) + sW_ready_nbar = pipeline.NamedBarrier(2, warpgroup_threads + warp_threads) + gk_load_nbar = pipeline.NamedBarrier(5, warpgroup_threads) + + elect_one = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1) + wg_coop = pipeline.CooperativeGroup(pipeline.Agent.Thread, warpgroup_threads) + warp_coop = pipeline.CooperativeGroup(pipeline.Agent.Thread, warp_threads) + + mma6_done_mbar = smem.allocate_array(cutlass.Int64, 1) + gmem_done_mbar = smem.allocate_array(cutlass.Int64, 1) + state_ready_mbar = smem.allocate_array(cutlass.Int64, 1) + final_state_done_mbar = smem.allocate_array(cutlass.Int64, 1) + + if warp_idx == 0: + with cute.arch.elect_one(): + cute.arch.mbarrier_init(mma6_done_mbar, warp_threads) + cute.arch.mbarrier_init(gmem_done_mbar, warpgroup_threads) + cute.arch.mbarrier_init(state_ready_mbar, warpgroup_threads) + cute.arch.mbarrier_init(final_state_done_mbar, warpgroup_threads) + + def _make_tma_pipe(total_byte_count, num_stages=1): + ptr = smem.allocate_array(cutlass.Int64, 2 * num_stages) + return pipeline.PipelineTmaUmma.create( + barrier_storage=ptr, + num_stages=num_stages, + producer_group=elect_one, + consumer_group=elect_one, + tx_count=total_byte_count // num_stages, + defer_sync=True, + ).make_participants() + + b_prod, b_cons = _make_tma_pipe(cute.size_in_bytes(mma_dtype, b_sl), num_stages=1) + v_prod, v_cons = _make_tma_pipe(cute.size_in_bytes(mma_dtype, v_sl), num_stages=1) + q_prod, q_cons = _make_tma_pipe(cute.size_in_bytes(mma_dtype, q_sl), num_stages=1) + a_prod, a_cons = _make_tma_pipe(cute.size_in_bytes(mma_dtype, a_sl), num_stages=1) + aqc_prod, aqc_cons = _make_tma_pipe( + cute.size_in_bytes(mma_dtype, aqc_sl), num_stages=1 + ) + kg_prod, kg_cons = _make_tma_pipe( + cute.size_in_bytes(mma_dtype, kg_sl), num_stages=1 + ) + + def _make_umma_pipe(): + ptr = smem.allocate_array(cutlass.Int64, 2) + return pipeline.PipelineUmmaAsync.create( + barrier_storage=ptr, + num_stages=1, + producer_group=elect_one, + consumer_group=wg_coop, + defer_sync=True, + ).make_participants() + + w_prod, w_cons = _make_umma_pipe() + nv_prod, nv_cons = _make_umma_pipe() + o_prod, o_cons = _make_umma_pipe() + + o_store_mbar = smem.allocate_array(cutlass.Int64, 2) + o_store_prod, o_store_cons = pipeline.PipelineAsync.create( + barrier_storage=o_store_mbar, + num_stages=1, + producer_group=wg_coop, + consumer_group=warp_coop, + defer_sync=True, + ).make_participants() + + cute.arch.sync_threads() + + tmem_ptr = cute.arch.retrieve_tmem_ptr(cutlass.Int32, 16, tmem_smem) + + tCtW_shape = tiled_mma_kmn.partition_shape_C((M, N)) + tCtW_fake = tiled_mma_kmn.make_fragment_C(tCtW_shape) + tCtW = cute.make_tensor( + cute.recast_ptr(tmem_ptr + 0, dtype=acc_dtype), tCtW_fake.layout + ) + + tCtNV_shape = tiled_mma_kmn.partition_shape_C((M, N)) + tCtNV_fake = tiled_mma_kmn.make_fragment_C(tCtNV_shape) + tCtNV = cute.make_tensor( + cute.recast_ptr(tmem_ptr + 128, dtype=acc_dtype), tCtNV_fake.layout + ) + tCtO = cute.make_tensor( + cute.recast_ptr(tmem_ptr + 384, dtype=acc_dtype), tCtNV_fake.layout + ) + + tCtS_shape = tiled_mma_mn_mn.partition_shape_C((M6, N6)) + tCtS_fake = tiled_mma_mn_mn.make_fragment_C(tCtS_shape) + tCtState = cute.make_tensor( + cute.recast_ptr(tmem_ptr + 256, dtype=acc_dtype), tCtS_fake.layout + ) + + mc = (0, 0, 0, 0) + ml = cute.make_layout((1, 1, 1, 1)) + + # PDL: setup is done; now wait for upstream akk_inv to commit before reading + # gmem inputs (TMA loads, gk_last_exp reads). PDL allows the entire setup phase + # above (~tmem alloc, smem layout, pipeline init) to overlap with akk_inv's tail. + cute.arch.griddepcontrol_wait() + + # ==== WG1: State readout + decay ==== + if warpgroup_idx == STATE_WG: + + cId_128 = cute.make_identity_tensor((M6, N6)) + tCtState_mn = transform_partitioned_tensor_layout(tCtState) + + atom_state_t2r = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tcgen05.Repetition(32)), acc_dtype + ) + tiled_state_t2r = tcgen05.make_tmem_copy( + atom_state_t2r, tCtState[(None, None), 0, 0] + ) + thr_state_t2r = tiled_state_t2r.get_slice(warpgroup_tidx) + tTR_tCtState = thr_state_t2r.partition_S(tCtState_mn) + tTR_tCcState = thr_state_t2r.partition_D(cId_128) + tRrState = cute.make_rmem_tensor_like(tTR_tCcState, acc_dtype) + + atom_state_r2t = cute.make_copy_atom( + tcgen05.St32x32bOp(tcgen05.Repetition(32)), acc_dtype + ) + tiled_state_r2t = tcgen05.make_tmem_copy( + atom_state_r2t, tCtState[(None, None), 0, 0] + ) + thr_state_r2t = tiled_state_r2t.get_slice(warpgroup_tidx) + tRT_tCtState = thr_state_r2t.partition_D(tCtState_mn) + + atom_state_g2r = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), acc_dtype, num_bits_per_copy=128 + ) + tiled_state_g2r = cute.make_tiled_copy_S(atom_state_g2r, tiled_state_r2t) + thr_state_g2r = tiled_state_g2r.get_slice(warpgroup_tidx) + + atom_state_r2g = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), acc_dtype, num_bits_per_copy=128 + ) + tiled_state_r2g = cute.make_tiled_copy_D(atom_state_r2g, tiled_state_t2r) + thr_state_r2g = tiled_state_r2g.get_slice(warpgroup_tidx) + + atom_state_r2s = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), mma_dtype, num_bits_per_copy=128 + ) + tiled_state_r2s = cute.make_tiled_copy_D(atom_state_r2s, tiled_state_t2r) + thr_state_r2s = tiled_state_r2s.get_slice(warpgroup_tidx) + sST_vk_view = transform_partitioned_tensor_layout(sST) + sST_kv_view = cute.make_tensor( + sST.iterator, cute.select(sST_vk_view.layout, mode=[1, 0, 2]) + ) + tCsState_inp = thr_state_r2s.partition_D(sST_kv_view) + tRrState_bf16 = cute.make_rmem_tensor_like(tTR_tCcState, mma_dtype) + tCrState_bf16 = tiled_state_r2s.retile(tRrState_bf16) + + sGk = cute.make_tensor(sGk_buf, cute.make_layout(((N,), (1,)))) + + scheduler = GDNTileScheduler.create( + scheduler_params, (bidx, bidy, bidz), grid_dim + ) + work = scheduler.initial_work_tile_info() + global_chunk = Int32(0) + + while work.is_valid_tile: + batch_idx, head_idx, _ = work.tile_idx + batch_start = cu_seqlens[batch_idx] + batch_end = cu_seqlens[batch_idx + 1] + num_chunks = cute.ceil_div(batch_end - batch_start, M) + # chunk_base from cumulative chunk-offsets array (handles non-64-aligned varlen). + # batch_start // M only works when all seq lengths are multiples of M. + chunk_base = chunk_offsets[batch_idx] + + gS_init = cute.flat_divide( + mS_fp32[batch_idx, head_idx, None, None], (M6, N6) + )[None, None, 0, 0] + tGR_tCgState_in = thr_state_g2r.partition_S(gS_init) + tGR_tCrState_in = thr_state_g2r.retile(tRrState) + cute.copy(tiled_state_g2r, tGR_tCgState_in, tGR_tCrState_in) + num_state_subs = tRrState.shape[2] + for sub in cutlass.range(num_state_subs): + cute.copy( + tiled_state_r2t, tRrState[None, 0, sub], tRT_tCtState[None, 0, sub] + ) + cute.arch.fence_view_async_tmem_store() + + sGk[warpgroup_tidx] = cutlass.Float32( + mGkLastExp[chunk_base, head_idx, warpgroup_tidx] + ) + gk_load_nbar.arrive_and_wait() + + for chunk_c in cutlass.range(num_chunks): + num_state_subs = tRrState.shape[2] + _sub_tile_size = cute.size(tRrState.shape[0]) + + if chunk_c > 0: + cute.arch.mbarrier_wait( + mma6_done_mbar, phase=(global_chunk - 1) % 2 + ) + gk_load_nbar.arrive_and_wait() + + cute.copy(tiled_state_t2r, tTR_tCtState, tRrState) + tRrState_bf16.store(tRrState.load().to(mma_dtype)) + cute.copy( + tiled_state_r2s, tCrState_bf16, tCsState_inp[None, None, None, 0] + ) + cute.arch.fence_view_async_shared() + cute.arch.mbarrier_arrive(gmem_done_mbar) + + for sub in cutlass.range(num_state_subs): + cute.copy( + tiled_state_t2r, + tTR_tCtState[None, 0, sub], + tRrState[None, 0, sub], + ) + cute.arch.fence_view_async_tmem_load() + for sub in cutlass.range(num_state_subs): + for i in cutlass.range(_sub_tile_size): + coord = tTR_tCcState[i, 0, sub] + # FIX: probe shows kernel state TMEM[m=K, n=V] (m corresponds to K-dim); + # OLD K4 / fla wants K-axis decay (state[k, v] *= gk[k]) → use coord[0]=m=K-idx. + k_idx = coord[0] + gk_val = cutlass.Float32(sGk[k_idx]) + tRrState[i, 0, sub] = tRrState[i, 0, sub] * gk_val + cute.copy( + tiled_state_r2t, + tRrState[None, 0, sub], + tRT_tCtState[None, 0, sub], + ) + cute.arch.fence_view_async_tmem_store() + cute.arch.mbarrier_arrive(state_ready_mbar) + + if chunk_c + 1 < num_chunks: + sGk[warpgroup_tidx] = cutlass.Float32( + mGkLastExp[chunk_base + chunk_c + 1, head_idx, warpgroup_tidx] + ) + global_chunk = global_chunk + 1 + + cute.arch.mbarrier_wait(mma6_done_mbar, phase=(global_chunk - 1) % 2) + gS_out = cute.flat_divide( + mS_fp32[batch_idx, head_idx, None, None], (M6, N6) + )[None, None, 0, 0] + tGR_tCgState_out = thr_state_r2g.partition_D(gS_out) + tGR_tCrState_out = thr_state_r2g.retile(tRrState) + num_state_subs_final = tRrState.shape[2] + for sub in cutlass.range(num_state_subs_final): + cute.copy( + tiled_state_t2r, tTR_tCtState[None, 0, sub], tRrState[None, 0, sub] + ) + cute.arch.fence_view_async_tmem_load() + for sub in cutlass.range(num_state_subs_final): + cute.copy( + tiled_state_r2g, + tGR_tCrState_out[None, 0, sub], + tGR_tCgState_out[None, 0, sub], + ) + + cute.arch.mbarrier_arrive(final_state_done_mbar) + + scheduler.advance_to_next_work() + work = scheduler.get_current_work() + + # ==== MMA warp (warp 0): 6 MMAs ==== + elif warp_idx == MMA_WARP: + fA_kmn = thr_kmn.make_fragment_A(sA) + fB_ks = thr_kmn.make_fragment_B(sB) + fB_v = thr_kmn.make_fragment_B(sV) + fB_s = thr_kmn.make_fragment_B(sST) + fA_q = thr_kmn.make_fragment_A(sQ) + fA_aqc = thr_kmn.make_fragment_A(sAQC) + fB_nv = thr_kmn.make_fragment_B(sNV_b) + fA_kg = thr_mn.make_fragment_A(sKG_a) + fB_nv_mn = thr_mn.make_fragment_B(sNV_b_mn) + fA_w = thr_kmn.make_fragment_A(sO) + + scheduler = GDNTileScheduler.create( + scheduler_params, (bidx, bidy, bidz), grid_dim + ) + work = scheduler.initial_work_tile_info() + global_chunk = Int32(0) + tile_count = Int32(0) + + while work.is_valid_tile: + batch_idx_m, head_idx_m, _ = work.tile_idx + num_chunks_m = cute.ceil_div( + cu_seqlens[batch_idx_m + 1] - cu_seqlens[batch_idx_m], M + ) + + for chunk_c in cutlass.range(num_chunks_m): + c_phase = global_chunk % 2 + + ah = a_cons.wait_and_advance() + bh = b_cons.wait_and_advance() + w_h = w_prod.acquire_and_advance() + tiled_mma_kmn.set(Field.ACCUMULATE, False) + for k in cutlass.range_constexpr(cute.size(sB.shape[2])): + cute.gemm( + tiled_mma_kmn, + tCtW, + fA_kmn[dice + (0,)][None, None, k], + fB_ks[dice + (bh.index,)][None, None, k], + tCtW, + ) + if k == 0: + tiled_mma_kmn.set(Field.ACCUMULATE, True) + bh.release() + w_h.commit() + + vh = v_cons.wait_and_advance() + tiled_mma_kmn.set(Field.ACCUMULATE, False) + for k in cutlass.range_constexpr(cute.size(sV.shape[2])): + cute.gemm( + tiled_mma_kmn, + tCtNV, + fA_kmn[dice + (0,)][None, None, k], + fB_v[dice + (vh.index,)][None, None, k], + tCtNV, + ) + if k == 0: + tiled_mma_kmn.set(Field.ACCUMULATE, True) + vh.release() + ah.release() + + cute.arch.mbarrier_wait(gmem_done_mbar, phase=c_phase) + sW_ready_nbar.arrive_and_wait() + + nv_h = nv_prod.acquire_and_advance() + tiled_mma_kmn.set(Field.ACCUMULATE, True) + for k in cutlass.range_constexpr(cute.size(sST.shape[2])): + cute.gemm( + tiled_mma_kmn, + tCtNV, + fA_w[dice + (0,)][None, None, k], + fB_s[dice + (0,)][None, None, k], + tCtNV, + ) + nv_h.commit() + + qh = q_cons.wait_and_advance() + tiled_mma_kmn.set(Field.ACCUMULATE, False) + for k in cutlass.range_constexpr(cute.size(sST.shape[2])): + cute.gemm( + tiled_mma_kmn, + tCtO, + fA_q[dice + (qh.index,)][None, None, k], + fB_s[dice + (0,)][None, None, k], + tCtO, + ) + if k == 0: + tiled_mma_kmn.set(Field.ACCUMULATE, True) + qh.release() + + sNV_ready_nbar.arrive_and_wait() + + aqch = aqc_cons.wait_and_advance() + o_h = o_prod.acquire_and_advance() + tiled_mma_kmn.set(Field.ACCUMULATE, True) + for k in cutlass.range_constexpr(cute.size(sNV_b.shape[2])): + cute.gemm( + tiled_mma_kmn, + tCtO, + fA_aqc[dice + (0,)][None, None, k], + fB_nv[dice + (0,)][None, None, k], + tCtO, + ) + o_h.commit() + aqch.release() + + cute.arch.mbarrier_wait(state_ready_mbar, phase=c_phase) + kgh = kg_cons.wait_and_advance() + tiled_mma_mn_mn.set(Field.ACCUMULATE, True) + for k in cutlass.range_constexpr(cute.size(sKG_a.shape[2])): + cute.gemm( + tiled_mma_mn_mn, + tCtState, + fA_kg[dice + (0,)][None, None, k], + fB_nv_mn[dice + (0,)][None, None, k], + tCtState, + ) + kgh.release() + + w_prod.tail() + nv_prod.tail() + o_prod.tail() + tcgen05.commit(mma6_done_mbar) + global_chunk = global_chunk + 1 + + cute.arch.mbarrier_wait(final_state_done_mbar, phase=tile_count % 2) + tile_count = tile_count + 1 + + scheduler.advance_to_next_work() + work = scheduler.get_current_work() + + cute.arch.relinquish_tmem_alloc_permit() + cute.arch.dealloc_tmem(tmem_ptr, 512) + + # ==== TMA warp (warp 2) ==== + elif warp_idx == TMA_WARP: + + cta_layout = cute.make_layout(1) + + scheduler = GDNTileScheduler.create( + scheduler_params, (bidx, bidy, bidz), grid_dim + ) + work = scheduler.initial_work_tile_info() + + if work.is_valid_tile: + tensormap_manager.init_tensormap_from_atom(tma_a[0], tm_a_ptr, TMA_WARP) + tensormap_manager.init_tensormap_from_atom(tma_b[0], tm_b_ptr, TMA_WARP) + tensormap_manager.init_tensormap_from_atom(tma_v[0], tm_v_ptr, TMA_WARP) + tensormap_manager.init_tensormap_from_atom(tma_q[0], tm_q_ptr, TMA_WARP) + tensormap_manager.init_tensormap_from_atom(tma_aqc[0], tm_aqc_ptr, TMA_WARP) + tensormap_manager.init_tensormap_from_atom(tma_kg[0], tm_kg_ptr, TMA_WARP) + tensormap_manager.fence_tensormap_initialization() + + while work.is_valid_tile: + batch_idx_t, head_idx_t, _ = work.tile_idx + batch_start_t = cu_seqlens[batch_idx_t] + batch_end_t = cu_seqlens[batch_idx_t + 1] + num_chunks_t = cute.ceil_div(batch_end_t - batch_start_t, M) + + bounded_a = cute.make_tensor( + mA.iterator, + cute.make_layout( + (batch_end_t, mA.shape[1], mA.shape[2]), + stride=(mA.stride[0], mA.stride[1], mA.stride[2]), + ), + ) + bounded_b = cute.make_tensor( + mB.iterator, + cute.make_layout( + (mB.shape[0], batch_end_t, mB.shape[2]), + stride=(mB.stride[0], mB.stride[1], mB.stride[2]), + ), + ) + bounded_v = cute.make_tensor( + mV_g.iterator, + cute.make_layout( + (mV_g.shape[0], batch_end_t, mV_g.shape[2]), + stride=(mV_g.stride[0], mV_g.stride[1], mV_g.stride[2]), + ), + ) + bounded_q = cute.make_tensor( + mQ.iterator, + cute.make_layout( + (batch_end_t, mQ.shape[1], mQ.shape[2]), + stride=(mQ.stride[0], mQ.stride[1], mQ.stride[2]), + ), + ) + bounded_aqc = cute.make_tensor( + mAQC.iterator, + cute.make_layout( + (batch_end_t, mAQC.shape[1], mAQC.shape[2]), + stride=(mAQC.stride[0], mAQC.stride[1], mAQC.stride[2]), + ), + ) + bounded_kg = cute.make_tensor( + mKG.iterator, + cute.make_layout( + (mKG.shape[0], batch_end_t, mKG.shape[2]), + stride=(mKG.stride[0], mKG.stride[1], mKG.stride[2]), + ), + ) + tensormap_manager.update_tensormap( + (bounded_a, bounded_b, bounded_v, bounded_q, bounded_aqc, bounded_kg), + (tma_a[0], tma_b[0], tma_v[0], tma_q[0], tma_aqc[0], tma_kg[0]), + (tm_a_ptr, tm_b_ptr, tm_v_ptr, tm_q_ptr, tm_aqc_ptr, tm_kg_ptr), + TMA_WARP, + (None, None, None, None, None, None), + ) + + for chunk_c in cutlass.range(num_chunks_t): + chunk_offset = batch_start_t + chunk_c * M + + # AB: A-operand (T, K, H) → domain_offset on (tokens, 0) + mA_c = cute.domain_offset( + (chunk_offset, Int32(0)), tma_a[1][None, None, head_idx_t] + ) + gA = cute.flat_divide(mA_c, (M, K)) + tCgA = thr_kmn.partition_A(gA) + tAsA, tAgA = cpasync.tma_partition( + tma_a[0], + 0, + cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + ah = a_prod.acquire_and_advance() + if chunk_c == 0: + tensormap_manager.fence_tensormap_update(tm_a_ptr) + cute.copy( + tma_a[0], + tAgA[(None, 0, 0)], + tAsA[(None, ah.index)], + tma_bar_ptr=ah.barrier, + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tm_a_ptr, cute.AddressSpace.generic + ), + ) + + # KS: B-operand (N, T, H) → domain_offset on (0, tokens) + mB_c = cute.domain_offset( + (Int32(0), chunk_offset), tma_b[1][None, None, head_idx_t] + ) + gB_c = cute.flat_divide(mB_c, (N, K)) + tCgB = thr_kmn.partition_B(gB_c) + tBsB, tBgB = cpasync.tma_partition( + tma_b[0], + 0, + cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + bh = b_prod.acquire_and_advance() + if chunk_c == 0: + tensormap_manager.fence_tensormap_update(tm_b_ptr) + cute.copy( + tma_b[0], + tBgB[(None, 0, 0)], + tBsB[(None, bh.index)], + tma_bar_ptr=bh.barrier, + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tm_b_ptr, cute.AddressSpace.generic + ), + ) + + # V: B-operand (N, T, H) → domain_offset on (0, tokens) + mV_c = cute.domain_offset( + (Int32(0), chunk_offset), tma_v[1][None, None, head_idx_t] + ) + gV_c = cute.flat_divide(mV_c, (N, K)) + tCgV = thr_kmn.partition_B(gV_c) + tBsV, tBgV = cpasync.tma_partition( + tma_v[0], + 0, + cta_layout, + cute.group_modes(sV, 0, 3), + cute.group_modes(tCgV, 0, 3), + ) + vh = v_prod.acquire_and_advance() + if chunk_c == 0: + tensormap_manager.fence_tensormap_update(tm_v_ptr) + cute.copy( + tma_v[0], + tBgV[(None, 0, 0)], + tBsV[(None, vh.index)], + tma_bar_ptr=vh.barrier, + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tm_v_ptr, cute.AddressSpace.generic + ), + ) + + # QS: A-operand (T, N, H) → domain_offset on (tokens, 0) + mQ_c = cute.domain_offset( + (chunk_offset, Int32(0)), tma_q[1][None, None, head_idx_t] + ) + gQ_c = cute.flat_divide(mQ_c, (M, K3)) + tCgQ = thr_kmn.partition_A(gQ_c) + tAsQ, tAgQ = cpasync.tma_partition( + tma_q[0], + 0, + cta_layout, + cute.group_modes(sQ, 0, 3), + cute.group_modes(tCgQ, 0, 3), + ) + qh = q_prod.acquire_and_advance() + if chunk_c == 0: + tensormap_manager.fence_tensormap_update(tm_q_ptr) + cute.copy( + tma_q[0], + tAgQ[(None, 0, 0)], + tAsQ[(None, qh.index)], + tma_bar_ptr=qh.barrier, + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tm_q_ptr, cute.AddressSpace.generic + ), + ) + + # AQC: A-operand (T, K, H) → domain_offset on (tokens, 0) + mAQC_c = cute.domain_offset( + (chunk_offset, Int32(0)), tma_aqc[1][None, None, head_idx_t] + ) + gAQC_c = cute.flat_divide(mAQC_c, (M, K)) + tCgAQC = thr_kmn.partition_A(gAQC_c) + tAsAQC, tAgAQC = cpasync.tma_partition( + tma_aqc[0], + 0, + cta_layout, + cute.group_modes(sAQC, 0, 3), + cute.group_modes(tCgAQC, 0, 3), + ) + aqch = aqc_prod.acquire_and_advance() + if chunk_c == 0: + tensormap_manager.fence_tensormap_update(tm_aqc_ptr) + cute.copy( + tma_aqc[0], + tAgAQC[(None, 0, 0)], + tAsAQC[(None, aqch.index)], + tma_bar_ptr=aqch.barrier, + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tm_aqc_ptr, cute.AddressSpace.generic + ), + ) + + # KG: B-operand (N6, T, H) → domain_offset on (0, tokens) + mKG_c = cute.domain_offset( + (Int32(0), chunk_offset), tma_kg[1][None, None, head_idx_t] + ) + gKG_c = cute.flat_divide(mKG_c, (N6, K6)) + tCgKG = thr_mn.partition_B(gKG_c) + tBsKG, tBgKG = cpasync.tma_partition( + tma_kg[0], + 0, + cta_layout, + cute.group_modes(sKG, 0, 3), + cute.group_modes(tCgKG, 0, 3), + ) + kgh = kg_prod.acquire_and_advance() + if chunk_c == 0: + tensormap_manager.fence_tensormap_update(tm_kg_ptr) + cute.copy( + tma_kg[0], + tBgKG[(None, 0, 0)], + tBsKG[(None, kgh.index)], + tma_bar_ptr=kgh.barrier, + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tm_kg_ptr, cute.AddressSpace.generic + ), + ) + + scheduler.advance_to_next_work() + work = scheduler.get_current_work() + + # ==== Warp 1: O TMA store ==== + elif warp_idx == O_STORE_WARP: + cta_layout_o = cute.make_layout(1) + + scheduler = GDNTileScheduler.create( + scheduler_params, (bidx, bidy, bidz), grid_dim + ) + work = scheduler.initial_work_tile_info() + + if work.is_valid_tile: + tensormap_manager.init_tensormap_from_atom(tma_o[0], tm_o_ptr, O_STORE_WARP) + tensormap_manager.fence_tensormap_initialization() + + while work.is_valid_tile: + batch_idx_o, head_idx_o, _ = work.tile_idx + batch_start_o = cu_seqlens[batch_idx_o] + batch_end_o = cu_seqlens[batch_idx_o + 1] + num_chunks_o = cute.ceil_div(batch_end_o - batch_start_o, M) + + bounded_o = cute.make_tensor( + mO.iterator, + cute.make_layout( + (batch_end_o, mO.shape[1], mO.shape[2]), + stride=(mO.stride[0], mO.stride[1], mO.stride[2]), + ), + ) + tensormap_manager.update_tensormap( + (bounded_o,), (tma_o[0],), (tm_o_ptr,), O_STORE_WARP, (None,) + ) + tensormap_manager.fence_tensormap_update(tm_o_ptr) + + for chunk_c in cutlass.range(num_chunks_o): + os_h = o_store_cons.wait_and_advance() + chunk_offset_o = batch_start_o + chunk_c * M + mO_c = cute.domain_offset( + (chunk_offset_o, Int32(0)), tma_o[1][None, None, head_idx_o] + ) + gOo = cute.flat_divide(mO_c, (M, N)) + sOt, gOt = cpasync.tma_partition( + tma_o[0], + 0, + cta_layout_o, + cute.group_modes(sO_out_st, 0, 2), + cute.group_modes(gOo, 0, 2), + ) + cute.copy( + tma_o[0], + sOt[None], + gOt[(None, 0, 0)], + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tm_o_ptr, cute.AddressSpace.generic + ), + ) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0, read=True) + os_h.release() + + scheduler.advance_to_next_work() + work = scheduler.get_current_work() + + # ==== WG2: W/NV/O readout (SS-mode, no Phase 2) ==== + elif warpgroup_idx == READOUT_WG: + + tCtW_mn = transform_partitioned_tensor_layout(tCtW) + tCtNV_mn = transform_partitioned_tensor_layout(tCtNV) + tCtO_mn = transform_partitioned_tensor_layout(tCtO) + + atom_t2r = cute.make_copy_atom( + tcgen05.Ld16x256bOp(tcgen05.Repetition(1)), acc_dtype + ) + tiled_t2r = tcgen05.make_tmem_copy(atom_t2r, tCtW[(None, None), 0, 0]) + thr_t2r = tiled_t2r.get_slice(warpgroup_tidx) + + tTR_W = thr_t2r.partition_S(tCtW_mn) + tTR_NV = thr_t2r.partition_S(tCtNV_mn) + tTR_O = thr_t2r.partition_S(tCtO_mn) + + atom_r2s_k = sm100_utils.get_smem_store_op( + utils.LayoutEnum.ROW_MAJOR, mma_dtype, acc_dtype, tiled_t2r + ) + tiled_r2s_k = cute.make_tiled_copy_D(atom_r2s_k, tiled_t2r) + thr_r2s_k = tiled_r2s_k.get_slice(warpgroup_tidx) + tCsO = thr_r2s_k.partition_D(transform_partitioned_tensor_layout(sO)) + tCsO_out = thr_r2s_k.partition_D(transform_partitioned_tensor_layout(sO_out)) + tCsNV = thr_r2s_k.partition_D(transform_partitioned_tensor_layout(sNV)) + + cId = cute.make_identity_tensor((M, N)) + tTR_cId = thr_t2r.partition_D(cId) + + scheduler = GDNTileScheduler.create( + scheduler_params, (bidx, bidy, bidz), grid_dim + ) + work = scheduler.initial_work_tile_info() + global_chunk = Int32(0) + + while work.is_valid_tile: + batch_idx_r, head_idx_r, _ = work.tile_idx + num_chunks_r = cute.ceil_div( + cu_seqlens[batch_idx_r + 1] - cu_seqlens[batch_idx_r], M + ) + + for chunk_c in cutlass.range(num_chunks_r): + c_phase = global_chunk % 2 + tRrR = cute.make_rmem_tensor_like(tTR_cId, acc_dtype) + tRrR_out = cute.make_rmem_tensor_like(tRrR, mma_dtype) + tCrR_k = tiled_r2s_k.retile(tRrR_out) + num_subs = tRrR.shape[2] + + wh = w_cons.wait_and_advance() + for sub in cutlass.range(num_subs): + cute.copy(tiled_t2r, tTR_W[None, 0, sub], tRrR[None, 0, sub]) + tRrR_out[None, 0, sub].store( + (-(tRrR[None, 0, sub].load())).to(mma_dtype) + ) + cute.copy(tiled_r2s_k, tCrR_k[None, 0, sub], tCsO[None, 0, sub, 0]) + cute.arch.fence_view_async_tmem_load() + wh.release() + cute.arch.fence_view_async_shared() + sW_ready_nbar.arrive_and_wait() + + nvh = nv_cons.wait_and_advance() + for sub in cutlass.range(num_subs): + cute.copy(tiled_t2r, tTR_NV[None, 0, sub], tRrR[None, 0, sub]) + tRrR_out[None, 0, sub].store( + tRrR[None, 0, sub].load().to(mma_dtype) + ) + cute.copy(tiled_r2s_k, tCrR_k[None, 0, sub], tCsNV[None, 0, sub, 0]) + cute.arch.fence_view_async_tmem_load() + nvh.release() + cute.arch.fence_view_async_shared() + sNV_ready_nbar.arrive_and_wait() + + os_h = o_store_prod.acquire_and_advance() + + oh = o_cons.wait_and_advance() + for sub in cutlass.range(num_subs): + cute.copy(tiled_t2r, tTR_O[None, 0, sub], tRrR[None, 0, sub]) + tRrR_out[None, 0, sub].store( + tRrR[None, 0, sub].load().to(mma_dtype) + ) + cute.copy( + tiled_r2s_k, tCrR_k[None, 0, sub], tCsO_out[None, 0, sub, 0] + ) + cute.arch.fence_view_async_tmem_load() + oh.release() + cute.arch.fence_view_async_shared() + os_h.commit() + global_chunk = global_chunk + 1 + + scheduler.advance_to_next_work() + work = scheduler.get_current_work() + o_store_prod.tail() + + +def make_host_fn(num_sm=148): + _num_sm = num_sm + + @cute.jit + def host_fn( + a_raw: cute.Tensor, + b_raw: cute.Tensor, + v_raw: cute.Tensor, + q_raw: cute.Tensor, + aqc_raw: cute.Tensor, + kg_raw: cute.Tensor, + o_raw: cute.Tensor, + gk_last_exp: cute.Tensor, + s_fp32: cute.Tensor, + cu_seqlens: cute.Tensor, + chunk_offsets: cute.Tensor, + tm_workspace: cute.Tensor, + stream: cuda_drv.CUstream, + ): + # Raw tensors from from_dlpack have PyTorch layout (T, H, dim) + # stride (H*dim, dim, 1). Reshape to 3D for TMA: + # A-operands: (T, dim, H), stride (dim*H, 1, dim) + # B-operands: (dim, T, H), stride (1, dim*H, dim) + T_tok = a_raw.shape[0] + H = a_raw.shape[1] + + a = cute.make_tensor( + a_raw.iterator, + cute.make_layout( + (T_tok, a_raw.shape[2], H), + stride=(a_raw.stride[0], a_raw.stride[2], a_raw.stride[1]), + ), + ) + q = cute.make_tensor( + q_raw.iterator, + cute.make_layout( + (T_tok, q_raw.shape[2], H), + stride=(q_raw.stride[0], q_raw.stride[2], q_raw.stride[1]), + ), + ) + aqc = cute.make_tensor( + aqc_raw.iterator, + cute.make_layout( + (T_tok, aqc_raw.shape[2], H), + stride=(aqc_raw.stride[0], aqc_raw.stride[2], aqc_raw.stride[1]), + ), + ) + o_out = cute.make_tensor( + o_raw.iterator, + cute.make_layout( + (T_tok, o_raw.shape[2], H), + stride=(o_raw.stride[0], o_raw.stride[2], o_raw.stride[1]), + ), + ) + + b = cute.make_tensor( + b_raw.iterator, + cute.make_layout( + (b_raw.shape[2], T_tok, H), + stride=(b_raw.stride[2], b_raw.stride[0], b_raw.stride[1]), + ), + ) + v = cute.make_tensor( + v_raw.iterator, + cute.make_layout( + (v_raw.shape[2], T_tok, H), + stride=(v_raw.stride[2], v_raw.stride[0], v_raw.stride[1]), + ), + ) + kg = cute.make_tensor( + kg_raw.iterator, + cute.make_layout( + (kg_raw.shape[2], T_tok, H), + stride=(kg_raw.stride[2], kg_raw.stride[0], kg_raw.stride[1]), + ), + ) + + tile1 = (M, N, K) + tile3 = (M, N, K3) + tile6 = (M6, N6, K6) + + mma_kmn = sm100_utils.make_trivial_tiled_mma( + mma_dtype, + OperandMajorMode.K, + OperandMajorMode.MN, + acc_dtype, + tcgen05.CtaGroup.ONE, + (M, N), + OperandSource.SMEM, + ) + mma_mn_mn = sm100_utils.make_trivial_tiled_mma( + mma_dtype, + OperandMajorMode.MN, + OperandMajorMode.MN, + acc_dtype, + tcgen05.CtaGroup.ONE, + (M6, N6), + OperandSource.SMEM, + ) + sl_a = sm100_utils.make_smem_layout_a(mma_kmn, tile1, mma_dtype, 1) + sl_b = sm100_utils.make_smem_layout_b(mma_kmn, tile1, mma_dtype, 1) + sl_v = sm100_utils.make_smem_layout_b(mma_kmn, tile1, mma_dtype, 1) + sl_s = sm100_utils.make_smem_layout_b(mma_kmn, tile3, mma_dtype, 1) + sl_q = sm100_utils.make_smem_layout_a(mma_kmn, tile3, mma_dtype, 1) + sl_aqc = sm100_utils.make_smem_layout_a(mma_kmn, tile1, mma_dtype, 1) + sl_kg = sm100_utils.make_smem_layout_b(mma_mn_mn, tile6, mma_dtype, 1) + + sl_readout_k = sm100_utils.make_smem_layout_a(mma_kmn, tile3, mma_dtype, 1) + sl_nv_b = sm100_utils.make_smem_layout_b(mma_kmn, tile1, mma_dtype, 1) + sl_nv_a = sm100_utils.make_smem_layout_a(mma_mn_mn, tile6, mma_dtype, 1) + sl_kg_a = sm100_utils.make_smem_layout_a(mma_mn_mn, tile6, mma_dtype, 1) + sl_nv_b_mn = sm100_utils.make_smem_layout_b(mma_mn_mn, tile6, mma_dtype, 1) + + tma_ld = cpasync.CopyBulkTensorTileG2SOp() + + ta_a = cute.nvgpu.make_tiled_tma_atom_A( + tma_ld, a, cute.select(sl_a, mode=[0, 1, 2]), tile1, mma_kmn + ) + ta_b = cute.nvgpu.make_tiled_tma_atom_B( + tma_ld, b, cute.select(sl_b, mode=[0, 1, 2]), tile1, mma_kmn + ) + ta_v = cute.nvgpu.make_tiled_tma_atom_B( + tma_ld, v, cute.select(sl_v, mode=[0, 1, 2]), tile1, mma_kmn + ) + ta_q = cute.nvgpu.make_tiled_tma_atom_A( + tma_ld, q, cute.select(sl_q, mode=[0, 1, 2]), tile3, mma_kmn + ) + ta_aqc = cute.nvgpu.make_tiled_tma_atom_A( + tma_ld, aqc, cute.select(sl_aqc, mode=[0, 1, 2]), tile1, mma_kmn + ) + ta_kg = cute.nvgpu.make_tiled_tma_atom_B( + tma_ld, kg, cute.select(sl_kg, mode=[0, 1, 2]), tile6, mma_mn_mn + ) + sk = sm100_utils.get_smem_layout_atom_ab(OperandMajorMode.K, out_dtype, (M, N)) + sl_store = cute.tile_to_shape( + sm100_utils.make_smem_layout_atom(sk, out_dtype), (M, N), order=(0, 1) + ) + tma_st = cpasync.CopyBulkTensorTileS2GOp() + ta_o = cpasync.make_tiled_tma_atom(tma_st, o_out, sl_store, (M, N)) + + n_seqs = cu_seqlens.shape[0] - 1 + scheduler_params = GDNTileSchedulerParams( + num_seqs=n_seqs, + num_q_heads=H, + num_v_heads=H, + is_GQA=False, + is_persistent=True, + ) + grid_shape = GDNTileScheduler.get_grid_shape(scheduler_params, _num_sm) + + k4_persistent_kernel( + mma_kmn, + mma_mn_mn, + ta_a, + sl_a, + ta_b, + sl_b, + ta_v, + sl_v, + sl_s, + ta_q, + sl_q, + ta_aqc, + sl_aqc, + ta_kg, + sl_kg, + ta_o, + sl_store, + sl_readout_k, + sl_nv_b, + sl_nv_a, + sl_kg_a, + sl_nv_b_mn, + gk_last_exp, + s_fp32, + cu_seqlens, + chunk_offsets, + a, + b, + v, + q, + aqc, + kg, + o_out, + tm_workspace, + scheduler_params, + ).launch( + grid=grid_shape, + block=(threads_per_cta, 1, 1), + use_pdl=True, + stream=stream, + ) + + return host_fn + + +def _compile(host_fn, *args, max_regs=None, keep_cubin=False): + mr = max_regs if max_regs is not None else MAX_REGS + if mr > 0: + opts = f"--maxrregcount={mr} --uumn" + if keep_cubin: + return cute.compile[KeepCUBIN, PtxasOptions(opts)](host_fn, *args) + return cute.compile[PtxasOptions(opts)](host_fn, *args) + if keep_cubin: + return cute.compile[KeepCUBIN](host_fn, *args) + return cute.compile(host_fn, *args) diff --git a/python/sglang/kernels/ops/attention/linear/kda_nvidia_prefill/fuse_kernel123_persistent.py b/python/sglang/kernels/ops/attention/linear/kda_nvidia_prefill/fuse_kernel123_persistent.py new file mode 100644 index 000000000..4cd8921a7 --- /dev/null +++ b/python/sglang/kernels/ops/attention/linear/kda_nvidia_prefill/fuse_kernel123_persistent.py @@ -0,0 +1,1704 @@ +# Vendored from the NVIDIA KDA_prefill package (benchmark/ Blackwell path) +# for the Kimi-K3 chunked prefill forward. Local deltas: fla.* imports +# re-pointed to sglang's vendored fla subset, flat sibling imports made +# package-relative, RCP_LN2 inlined. +# ruff: noqa -- vendored kernel library, minimal local deltas +""" +Persistent Fused K1+K2+K3 Kernel for KDA. + +Fuses gate activation + cumsum + scaling (K1), intra sub-chunk Aqk/Akk (K2), +and inter sub-chunk solve + merged inverse (K3) into a single persistent kernel. + +Grid: (NUM_SMS, 1, 1) — 148 persistent blocks, each loops over work units + Total work units = (NT/4) * H * B, distributed round-robin across SMs + Block i processes work units i, i+NUM_SMS, i+2*NUM_SMS, ... +Block: 1024 threads (32 warps), warp-specialized with setmaxnreg (all groups 4-aligned): + Warps 0-15: TMA+K1 fused (8×2, vec2, prefetch pipeline) – 4 WGs, 56 regs + Warps 16-27: K2 MMA compute (10 active + 2 idle for WG alignment) – 3 WGs, 72 regs + Warps 28-31: Store/Inversion warps – 1 WG, 24 regs + +Pipeline (single for_generate, warp groups separated by if-blocks): + per work unit: + Warps 0-15: prefetch chunk 0→stage 0 (warp 0), then loop: + TMA next chunk (warp 0), wait cur chunk, K1 compute, arrive(k1_done) + Warps 16-27: wait(k1_done)+wait(store_done), MMA, arrive(mma_done+stage_reuse) + Warps 28-31: wait(mma_done), store sAqk/sAkk→GMEM, arrive(store_done) + All warp-group invariants are computed inside each group's if-block (not hoisted) + to eliminate cross-group register pressure — same budget as the _all version. + Mbarrier phases self-reset after 4 iterations (2 stages × 2 phases). + +Mbarriers: + tma_mbars[2]: count=1, warp 0 lane 0 → K1+MMA wait for TMA data + stage_reuse_mbars[2]: count=384, MMA(12 warps) → warp 0 waits before TMA reuse + k1_done_mbars[2]: count=512, K1(16 warps) → MMA waits for g_cumsum ready + mma_done_mbars[2]: count=384, MMA(12 warps) → Store waits for sAqk/sAkk ready + store_done_mbars[2]: count=128, Store(4 warps) → MMA waits for sAqk/sAkk stage free + +SMEM: ~215KB (q+k+g × [64,128] bf16 × 2 stages + g_cumsum [64,136] fp32 × 2 stages + + sPartialLast [8,132] fp32 + sAqk [16,168,2] bf16 + + sAkk [64,72,2] fp32 block-transposed upper-tri layout) + +Inputs: + g [B,T,H,K] bf16 raw gate + k [B,T,H,K] bf16 + q [B,T,H,K] bf16 + A_log [H] fp32 per-head log decay + beta [B,T,H] bf16 used for Akk unit lower triangular + scale fp32 1/sqrt(K) + +Outputs (g_cumsum stays in SMEM, not written to GMEM): + k_scaled [B,T,H,K] bf16 + q_scaled [B,T,H,K] bf16 + kg [B,T,H,K] bf16 + gk_last_exp[B,NT,H,K] fp32 + A_qk [B,T,H,BT] bf16 full merged (diagonal + off-diagonal) + A_kk [B,T,H,BT] fp32 block-transposed upper-tri (input to akk_inv) +""" + +import cuda.bindings.driver as cuda_drv +import cutlass +import cutlass.cute as cute +from cutlass import for_generate, yield_out +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cutlass_dsl import T, dsl_user_op + +BT = 64 +BC = 16 +K_DIM = 128 +K_PAD = 8 +K_STRIDE = K_DIM + K_PAD # 136, padded row stride to avoid bank conflicts +CHUNKS_PER_BLOCK = 4 +NUM_SMS = 148 # Persistent kernel: one resident block per SM + +NUM_K1_TMA_WARPS = 16 # Warps 0-15: K1 compute (4 warpgroups, 8×2) -- TMA offloaded +NUM_MMA_WARPS = ( + 11 # Warps 16-26: MMA (10 active + 1 TMA producer, dropped idle warp 27) +) +NUM_MMA_ACTIVE = 10 # mma_warp 0..9: actual MMA work +TMA_WARP_ID = NUM_K1_TMA_WARPS + NUM_MMA_ACTIVE # warp 26 = dedicated TMA producer +NUM_STORE_WARPS = 4 # Warps 28-31: Store/Inversion (1 warpgroup) +NUM_WARPS = NUM_K1_TMA_WARPS + NUM_MMA_WARPS + NUM_STORE_WARPS # 32 +THREADS = NUM_WARPS * 32 # 1024 + +NUM_SUB_CHUNKS = BT // BC # 4 +NUM_TILES = NUM_SUB_CHUNKS * (NUM_SUB_CHUNKS + 1) // 2 # 10 lower-tri tiles +MMA_K_TILE = 16 +NUM_MMA_K_TILES = K_DIM // MMA_K_TILE # 8 (bf16 m16n8k16) +AQK_TILE_PAD = 8 +AQK_TILE_STRIDE = ( + BT + AQK_TILE_PAD +) # 72 — sAqk now 64x72 row-major (same shape as sAkk) + +AKK_PAD = 8 +AKK_STRIDE = BT + AKK_PAD # 72 + + +K1_ROW_GROUPS = 8 +K1_COL_GROUPS = 2 +ROWS_PER_K1_WARP = BT // K1_ROW_GROUPS # 8 +K1_COLS_PER_WARP = K_DIM // K1_COL_GROUPS # 64 + +VEC = K1_COLS_PER_WARP // 32 # 2 +K_VEC = K_DIM // VEC # 64 +NUM_STAGES = 2 +PARTIAL_COLS = K_DIM + 4 # 132 +PARTIAL_COLS_PER_WARP = K_DIM // NUM_K1_TMA_WARPS # 8 + +_TILE_IQ = [0, 1, 1, 2, 2, 2, 3, 3, 3, 3] +_TILE_IK = [0, 0, 1, 0, 1, 2, 0, 1, 2, 3] + +LOG2E = 1.4426950408889634 +LN2 = 0.6931471805599453 +RCP_LN2 = LOG2E + + +@dsl_user_op +def k1_internal_barrier(*, loc=None, ip=None): + """Named barrier for K1+TMA warps (0-15, 512 threads). barrier_id=2.""" + llvm.inline_asm( + T.i32(), + [], + "membar.cta; bar.sync 2, 512; mov.u32 $0, 0;", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def pack_bf16x2_f32(hi_f32, lo_f32, *, loc=None, ip=None): + """Pack two fp32 values into a bf16x2 u32 register. + Returns u32 with bf16(hi) in bits [31:16] and bf16(lo) in bits [15:0]. + PTX: cvt.rn.bf16x2.f32 d, a, b -> d[31:16]=bf16(a), d[15:0]=bf16(b) + """ + result = llvm.inline_asm( + T.i32(), + [hi_f32.ir_value(loc=loc, ip=ip), lo_f32.ir_value(loc=loc, ip=ip)], + "cvt.rn.bf16x2.f32 $0, $1, $2;", + "=r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Int32(result) + + +@dsl_user_op +def mma_bf16_m16n8k16( + a0, + a1, + a2, + a3, # 4 u32 (bf16x2 packed) A operands + b0, + b1, # 2 u32 (bf16x2 packed) B operands + c0, + c1, + c2, + c3, # 4 fp32 accumulators + *, + loc=None, + ip=None, +): + """bf16 MMA with fp32 accumulator, shape m16n8k16. + D_fp32 = A_bf16 * B_bf16 + C_fp32 + """ + # a/b already i32 (from pack_bf16x2_f32) -> no bitcast needed + result = llvm.inline_asm( + ir.Type.parse("!llvm.struct<(f32, f32, f32, f32)>"), + [ + a0.ir_value(loc=loc, ip=ip), + a1.ir_value(loc=loc, ip=ip), + a2.ir_value(loc=loc, ip=ip), + a3.ir_value(loc=loc, ip=ip), + b0.ir_value(loc=loc, ip=ip), + b1.ir_value(loc=loc, ip=ip), + c0.ir_value(loc=loc, ip=ip), + c1.ir_value(loc=loc, ip=ip), + c2.ir_value(loc=loc, ip=ip), + c3.ir_value(loc=loc, ip=ip), + ], + """{ + mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 + {$0, $1, $2, $3}, + {$4, $5, $6, $7}, + {$8, $9}, + {$10, $11, $12, $13}; + }""", + "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + d0 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)) + d1 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)) + d2 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [2], loc=loc, ip=ip)) + d3 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [3], loc=loc, ip=ip)) + return d0, d1, d2, d3 + + +SHFL_W8_CLAMP = 0x1800 + + +@dsl_user_op +def fast_rcp(x, *, loc=None, ip=None): + """Hardware fast reciprocal: rcp.approx.ftz.f32 (~2 cycles vs ~20 for div).""" + result = llvm.inline_asm( + T.f32(), + [x.ir_value(loc=loc, ip=ip)], + "rcp.approx.ftz.f32 $0, $1;", + "=f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(result) + + +@dsl_user_op +def opaque_zero_from_work_id(*, loc=None, ip=None): + """ + Return 0 via opaque side-effectful asm (no inputs needed). + + Because has_side_effects=True, MLIR LICM treats this as having memory + effects and will NOT hoist it outside the for_generate loop. Any value + computed from the result (_oz) therefore appears loop-variant to LICM, + preventing get_slice() and scalar layout-invariant computations from being + hoisted to the kernel prologue. This keeps prologue register pressure < 64 + and eliminates the 440-byte stack frame that caused ~300-cycle L2 LDL + penalties per iteration. + """ + result = llvm.inline_asm( + T.i32(), + [], + "mov.b32 $0, 0;", # output = 0 (opaque to compiler constant-folding) + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Int32(result) + + +@cute.kernel +def fused_kernel123( + tma_atom_Q: cute.CopyAtom, + tma_tensor_Q: cute.Tensor, + tma_atom_K: cute.CopyAtom, + tma_tensor_K: cute.Tensor, + tma_atom_G: cute.CopyAtom, + tma_tensor_G: cute.Tensor, + mA_log: cute.Tensor, + mBeta: cute.Tensor, + scale: cutlass.Float32, + mKscaled: cute.Tensor, + mKg: cute.Tensor, + mQscaled: cute.Tensor, + mGkLast: cute.Tensor, + mAqk: cute.Tensor, # 4D (B, T, H, BT) — used by per-tile pure path + mAkk: cute.Tensor, # 4D (B, T, H, BT) — used by per-tile pure path + mAqk_v2: cute.Tensor, # 5D (B, T, H, BT/2, 2) — used by vec autovec non-pure path + mAkk_v2: cute.Tensor, # 5D (B, T, H, BT/2, 2) — used by vec autovec non-pure path + tiled_copy_qk_k1, + tiled_mma_k2, + tiled_copy_mma_A, + tiled_copy_mma_B, + tiled_copy_Gcum_norm, + tiled_copy_Gcum_gate, + qk_smem_layout, + g_smem_layout, + g_cumsum_layout, + num_chunks: int, + num_heads: int, + batch_size: int, + mCuSeqlens: cute.Tensor, + mChunkIndices: cute.Tensor, + IS_VARLEN: cutlass.Constexpr[int], + mDtBias: cute.Tensor, + lower_bound: cutlass.Float32, + HAS_BIAS: cutlass.Constexpr[int], + USE_SAFE_GATE: cutlass.Constexpr[int], + VARLEN_PURE: cutlass.Constexpr[int] = 0, +): + block_id, _, _ = cute.arch.block_idx() + tidx = cute.arch.thread_idx()[0] + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + lane_id = tidx % 32 + + total_cgs_per_head = cutlass.Int32(0) + cgs_per_head = cutlass.Int32(0) + total_cgs = cutlass.Int32(0) + if IS_VARLEN: + total_cgs_per_head = (num_chunks + CHUNKS_PER_BLOCK - 1) // CHUNKS_PER_BLOCK + total_cgs = total_cgs_per_head * num_heads + else: + cgs_per_head = num_chunks // CHUNKS_PER_BLOCK + total_cgs = cgs_per_head * num_heads * batch_size + + smem = cutlass.utils.SmemAllocator() + sQ = smem.allocate_tensor( + cutlass.BFloat16, qk_smem_layout.outer, 128, swizzle=qk_smem_layout.inner + ) + sK = smem.allocate_tensor( + cutlass.BFloat16, qk_smem_layout.outer, 128, swizzle=qk_smem_layout.inner + ) + sG = smem.allocate_tensor(cutlass.BFloat16, g_smem_layout, 128) + sGcum = smem.allocate_tensor(cutlass.Float32, g_cumsum_layout, 128) + partial_last_layout = cute.make_layout( + (K1_ROW_GROUPS, PARTIAL_COLS), stride=(PARTIAL_COLS, 1) + ) + sPartialLast = smem.allocate_tensor(cutlass.Float32, partial_last_layout, 128) + + # sAqk: 64x72 row-major (BT rows, BT+pad cols), same shape as sAkk. + # Each sub-tile (i_q, i_k) sits at SMEM rows [i_q*BC, (i_q+1)*BC) cols + # [i_k*BC, (i_k+1)*BC) — directly mirrors the 64x64 attention matrix. + aqk_tile_layout = cute.make_layout( + (BT, AQK_TILE_STRIDE, NUM_STAGES), + stride=(AQK_TILE_STRIDE, 1, BT * AQK_TILE_STRIDE), + ) + sAqk = smem.allocate_tensor(cutlass.BFloat16, aqk_tile_layout, 128) + + akk_tile_layout = cute.make_layout( + (BT, AKK_STRIDE, NUM_STAGES), stride=(AKK_STRIDE, 1, BT * AKK_STRIDE) + ) + sAkk = smem.allocate_tensor(cutlass.BFloat16, akk_tile_layout, 128) + # sAkk_pkd / sTemp removed: akk_inv runs as a separate kernel call (chained back-to-back). + + # sBeta: chunk's 64 beta values staged in SMEM. K1 reads from here in inner loop + # to reduce register pressure / avoid repeated gmem broadcast. + beta_smem_layout = cute.make_layout((BT, NUM_STAGES), stride=(1, BT)) + sBeta = smem.allocate_tensor(cutlass.BFloat16, beta_smem_layout, 128) + + # Mbarrier allocation & init + tma_mbars = smem.allocate_array(cutlass.Int64, NUM_STAGES) + stage_reuse_mbars = smem.allocate_array(cutlass.Int64, NUM_STAGES) + k1_done_mbars = smem.allocate_array(cutlass.Int64, NUM_STAGES) + mma_done_mbars = smem.allocate_array(cutlass.Int64, NUM_STAGES) + store_done_mbars = smem.allocate_array(cutlass.Int64, NUM_STAGES) + + bytes_per_stage = BT * K_DIM * 2 * 3 + + if tidx == 0: + for s in range(NUM_STAGES): + cute.arch.mbarrier_init(tma_mbars + s, 1) + # TMA warp is waiter (not arriver) on stage_reuse; and it skips mma_done arrive + cute.arch.mbarrier_init(stage_reuse_mbars + s, (NUM_MMA_WARPS - 1) * 32) + cute.arch.mbarrier_init(k1_done_mbars + s, NUM_K1_TMA_WARPS * 32) + cute.arch.mbarrier_init(mma_done_mbars + s, (NUM_MMA_WARPS - 1) * 32) + cute.arch.mbarrier_init(store_done_mbars + s, NUM_STORE_WARPS * 32) + cute.arch.mbarrier_init_fence() + cute.arch.barrier() + + # SMEM init: zero out sAqk and sAkk valid 64x64 region (cols 64..71 are + # padding for SMEM bank-conflict avoidance, never read or written by MMA + # or store warps). Required for downstream row-major store optimizations + # — positions outside MMA-written sub-tiles stay at 0. + # + # Cooperative pattern (32 warps × 32 lanes = 1024 threads): + # - Each warp owns 2 contiguous rows (warp_id*2, warp_id*2+1) + # - Each lane owns 2 contiguous bf16 cols (lane*2, lane*2+1) + # - Per lane: 2 stages × 2 rows × 2 buffers × 2 cols = 16 bf16 stores + # - Adjacent (lane*2, lane*2+1) bf16 pairs are 4-byte aligned → + # ptxas should fuse into STS.32 (8 wide stores per lane). + _warp_id_in_cta = tidx >> 5 # tidx // 32, range 0..31 + _lane_id_warp = tidx & 31 # tidx % 32, range 0..31 + _row_base = _warp_id_in_cta * 2 # this warp owns rows [_row_base, _row_base+1] + _col_lo = _lane_id_warp * 2 # this lane owns cols [_col_lo, _col_lo+1] + _col_hi = _col_lo + 1 + for _s in cutlass.range_constexpr(NUM_STAGES): + for _ri in cutlass.range_constexpr(2): + _row = _row_base + _ri + sAqk[_row, _col_lo, _s] = cutlass.BFloat16(0.0) + sAqk[_row, _col_hi, _s] = cutlass.BFloat16(0.0) + sAkk[_row, _col_lo, _s] = cutlass.BFloat16(0.0) + sAkk[_row, _col_hi, _s] = cutlass.BFloat16(0.0) + cute.arch.barrier() + + # Pre-arrive (MMA warps only) + # stage_reuse_mbars: warp 0 waits before MMA arrives → pre-arrive all 12 MMA warps + # store_done_mbars: MMA waits before Store arrives → pre-arrive first 4 MMA warps + if ( + warp_idx >= NUM_K1_TMA_WARPS + and warp_idx < NUM_K1_TMA_WARPS + NUM_MMA_WARPS + and warp_idx != TMA_WARP_ID + ): + mma_warp_tmp = warp_idx - NUM_K1_TMA_WARPS + for s in range(NUM_STAGES): + cute.arch.mbarrier_arrive(stage_reuse_mbars + s) + if mma_warp_tmp < NUM_STORE_WARPS: + cute.arch.mbarrier_arrive(store_done_mbars + s) + + # Persistent outer loop. Single for_generate at top level (required). + # Opaque asm barrier on work_id prevents MLIR LICM from hoisting + # get_slice() and scalar layout invariants to the kernel prologue, + # keeping register pressure < 64 and eliminating prologue spill. + for work_id in for_generate(block_id, total_cgs, NUM_SMS): + i_cg = cutlass.Int32(0) + i_h = cutlass.Int32(0) + i_b = cutlass.Int32(0) + chunk_base = cutlass.Int32(0) + if IS_VARLEN: + i_cg = work_id % total_cgs_per_head + i_h = work_id // total_cgs_per_head + i_b = cutlass.Int32(0) + chunk_base = i_cg * CHUNKS_PER_BLOCK + else: + i_cg = work_id % cgs_per_head + i_h = (work_id // cgs_per_head) % num_heads + i_b = work_id // (cgs_per_head * num_heads) + chunk_base = i_cg * CHUNKS_PER_BLOCK + + # Anti-LICM barrier: _oz is always 0 but appears to depend on work_id. + # Because this asm has side_effects=True, it stays inside the loop. + # Any value computed from _oz/_lane/_warp is also loop-variant from + # LICM's perspective → get_slice() and scalar invariants stay in-loop. + _oz = opaque_zero_from_work_id() + _lane = lane_id + _oz + _warp = warp_idx + _oz + + # Warps 0-15: Fused TMA + K1 + if warp_idx < NUM_K1_TMA_WARPS: + # Warp-layout invariants (scope-local → no cross-group register spill) + k1_warp = _warp + warp_row_group = k1_warp % K1_ROW_GROUPS + warp_col_group = k1_warp // K1_ROW_GROUPS + k1_row_start = warp_row_group * ROWS_PER_K1_WARP + col_base = warp_col_group * K1_COLS_PER_WARP + _lane * VEC + col_vec_idx = warp_col_group * (K1_COLS_PER_WARP // VEC) + _lane + cumsum_scale = cutlass.Float32(RCP_LN2) + thr_copy_k1 = tiled_copy_qk_k1.get_slice(_lane) + + rAcc = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.Float32) + rPrefix = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.Float32) + rGkLast = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.Float32) + rKsOut = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.BFloat16) + rQsOut = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.BFloat16) + rKgOut = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.BFloat16) + rGkOut = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.Float32) + + # exp_A depends on i_h (changes per work unit) + exp_A = cute.exp(mA_log[i_h], fastmath=True) + + # Load dt_bias per (head, col) — broadcast across all rows + rBias = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.Float32) + if HAS_BIAS: + for vi in cutlass.range_constexpr(VEC): + rBias[vi] = mDtBias[i_h, col_base + vi] + else: + for vi in cutlass.range_constexpr(VEC): + rBias[vi] = cutlass.Float32(0.0) + + # 3D TMA head slices (fixed for this work unit's head) + gQ_head = tma_tensor_Q[(None, None, i_h)] + gK_head = tma_tensor_K[(None, None, i_h)] + gG_head = tma_tensor_G[(None, None, i_h)] + + for chunk_iter in cutlass.range_constexpr(CHUNKS_PER_BLOCK): + cur_stage = chunk_iter % NUM_STAGES + cur_phase = chunk_iter // NUM_STAGES % 2 + chunk_idx = chunk_base + chunk_iter + chunk_start = cutlass.Int32(0) + ci_eos = cutlass.Int32(0) + if IS_VARLEN: + if chunk_idx < num_chunks: + _sid = cutlass.Int32(mChunkIndices[chunk_idx, 0]) + chunk_start = ( + cutlass.Int32(mCuSeqlens[_sid]) + + cutlass.Int32(mChunkIndices[chunk_idx, 1]) * BT + ) + ci_eos = cutlass.Int32(mCuSeqlens[_sid + 1]) + else: + chunk_start = chunk_idx * BT + + cute.arch.mbarrier_wait(tma_mbars + cur_stage, cur_phase) + + csG = sG[(None, None, cur_stage)] + csGcum = sGcum[(None, None, cur_stage)] + csQ = sQ[(None, None, cur_stage)] + csK = sK[(None, None, cur_stage)] + csBeta = sBeta[(None, cur_stage)] + + # Stage chunk's 64 beta values to SMEM (warp 0 of K1 group, 32 threads × 2). + # Synced via the existing k1_internal_barrier() before Pass 2b reads it. + if k1_warp == 0: + for _bi in cutlass.range_constexpr(BT // 32): + _idx = _bi * 32 + _lane + csBeta[_idx] = mBeta[i_b, chunk_start + _idx, i_h] + + rGact = cute.make_rmem_tensor( + cute.make_layout((ROWS_PER_K1_WARP, VEC)), cutlass.Float32 + ) + for vi in cutlass.range_constexpr(VEC): + rAcc[vi] = cutlass.Float32(0.0) + + for ri in cutlass.range_constexpr(ROWS_PER_K1_WARP): + row = k1_row_start + ri + for vi in cutlass.range_constexpr(VEC): + c = col_base + vi + g_val = csG[row, c].to(cutlass.Float32) + if HAS_BIAS: + g_val = g_val + rBias[vi] + g_activated = cutlass.Float32(0.0) + if USE_SAFE_GATE: + sigmoid_g = fast_rcp( + cutlass.Float32(1.0) + + cute.exp2(-exp_A * g_val * LOG2E, fastmath=True) + ) + g_activated = lower_bound * sigmoid_g + else: + softplus_g = ( + cute.log2( + cutlass.Float32(1.0) + + cute.exp2(g_val * LOG2E, fastmath=True), + fastmath=True, + ) + * LN2 + ) + g_activated = -exp_A * softplus_g + # Varlen: zero gate for out-of-bounds rows so cumsum + # stays flat beyond the last valid position. + # VARLEN_PURE=1 elides this at compile time — caller + # guarantees all seq lengths are multiples of BT so no + # chunk has OOB rows. + if IS_VARLEN and not VARLEN_PURE: + if chunk_start + row >= ci_eos: + g_activated = cutlass.Float32(0.0) + rGact[ri, vi] = g_activated + rAcc[vi] = rAcc[vi] + g_activated + + for vi in cutlass.range_constexpr(VEC): + sPartialLast[warp_row_group, col_base + vi] = rAcc[vi] + + k1_internal_barrier() + + prefix_col_start = k1_warp * PARTIAL_COLS_PER_WARP + row_in_prefix = lane_id % K1_ROW_GROUPS + col_in_group = lane_id // K1_ROW_GROUPS + + for j in cutlass.range_constexpr(PARTIAL_COLS_PER_WARP // 4): + col = prefix_col_start + j * 4 + col_in_group + val = cutlass.Float32(sPartialLast[row_in_prefix, col]) + tmp = cute.arch.shuffle_sync_up( + val, 1, mask=-1, mask_and_clamp=SHFL_W8_CLAMP + ) + if row_in_prefix >= 1: + val = val + tmp + tmp = cute.arch.shuffle_sync_up( + val, 2, mask=-1, mask_and_clamp=SHFL_W8_CLAMP + ) + if row_in_prefix >= 2: + val = val + tmp + tmp = cute.arch.shuffle_sync_up( + val, 4, mask=-1, mask_and_clamp=SHFL_W8_CLAMP + ) + if row_in_prefix >= 4: + val = val + tmp + sPartialLast[row_in_prefix, col] = val + + k1_internal_barrier() + + for vi in cutlass.range_constexpr(VEC): + rGkLast[vi] = sPartialLast[K1_ROW_GROUPS - 1, col_base + vi] + + for vi in cutlass.range_constexpr(VEC): + rPrefix[vi] = cutlass.Float32(0.0) + if warp_row_group > 0: + for vi in cutlass.range_constexpr(VEC): + rPrefix[vi] = sPartialLast[warp_row_group - 1, col_base + vi] + + # ---- Pass 2a: ONLY cumsum + write csGcum (critical path, minimal work) ---- + for vi in cutlass.range_constexpr(VEC): + rAcc[vi] = rPrefix[vi] + + for ri in cutlass.range_constexpr(ROWS_PER_K1_WARP): + row = k1_row_start + ri + for vi in cutlass.range_constexpr(VEC): + rAcc[vi] = rAcc[vi] + rGact[ri, vi] + csGcum[row, col_base + vi] = rAcc[vi] * cumsum_scale + + # Signal MMA early: csGcum is ready + cute.arch.mbarrier_arrive(k1_done_mbars + cur_stage) + + # ---- Pass 2b: recompute + write GMEM (overlaps with MMA, off critical path) ---- + for vi in cutlass.range_constexpr(VEC): + rAcc[vi] = rPrefix[vi] + + for ri in cutlass.range_constexpr(ROWS_PER_K1_WARP): + row = k1_row_start + ri + t = chunk_start + row + + sK_tile = cute.local_tile( + csK, tiler=(1, K1_COLS_PER_WARP), coord=(row, warp_col_group) + ) + tCsK = thr_copy_k1.partition_S(sK_tile) + tCrK = cute.make_fragment_like(tCsK) + cute.copy(tiled_copy_qk_k1, tCsK, thr_copy_k1.retile(tCrK)) + + sQ_tile = cute.local_tile( + csQ, tiler=(1, K1_COLS_PER_WARP), coord=(row, warp_col_group) + ) + tCsQ = thr_copy_k1.partition_S(sQ_tile) + tCrQ = cute.make_fragment_like(tCsQ) + cute.copy(tiled_copy_qk_k1, tCsQ, thr_copy_k1.retile(tCrQ)) + + # Read beta from SMEM (staged once per chunk by warp 0). + beta_val = cutlass.Float32(csBeta[row]) + + for vi in cutlass.range_constexpr(VEC): + rAcc[vi] = rAcc[vi] + rGact[ri, vi] + cs = rAcc[vi] * cumsum_scale + + k_val = tCrK[vi].to(cutlass.Float32) + q_val = tCrQ[vi].to(cutlass.Float32) + + exp2_cs = cute.exp2(cs, fastmath=True) + gk_last_cs = rGkLast[vi] * cumsum_scale + exp2_kg = cute.exp2(gk_last_cs - cs, fastmath=True) + + rKsOut[vi] = (k_val * exp2_cs).to(cutlass.BFloat16) + rQsOut[vi] = (q_val * exp2_cs * scale).to(cutlass.BFloat16) + rKgOut[vi] = (k_val * exp2_kg).to(cutlass.BFloat16) + + if IS_VARLEN and not VARLEN_PURE: + if t < ci_eos: + cute.autovec_copy( + rKsOut, mKscaled[i_b, t, i_h, col_vec_idx, None] + ) + cute.autovec_copy( + rQsOut, mQscaled[i_b, t, i_h, col_vec_idx, None] + ) + cute.autovec_copy( + rKgOut, mKg[i_b, t, i_h, col_vec_idx, None] + ) + else: + cute.autovec_copy( + rKsOut, mKscaled[i_b, t, i_h, col_vec_idx, None] + ) + cute.autovec_copy( + rQsOut, mQscaled[i_b, t, i_h, col_vec_idx, None] + ) + cute.autovec_copy(rKgOut, mKg[i_b, t, i_h, col_vec_idx, None]) + + if warp_row_group == 0: + for vi in cutlass.range_constexpr(VEC): + rGkOut[vi] = cute.exp2( + rGkLast[vi] * cumsum_scale, fastmath=True + ) + if IS_VARLEN: + if ci_eos > cutlass.Int32(0): + cute.autovec_copy( + rGkOut, mGkLast[i_b, chunk_idx, i_h, col_vec_idx, None] + ) + else: + cute.autovec_copy( + rGkOut, mGkLast[i_b, chunk_idx, i_h, col_vec_idx, None] + ) + + # Warp 26 (TMA_WARP_ID): dedicated TMA producer. + # Waits stage_reuse (gated by MMA arrives), issues TMA for Q/K/G, + # signals tma_mbar. Decouples MMA -> TMA dependency from K1 compute. + if warp_idx == TMA_WARP_ID: + gQ_head = tma_tensor_Q[(None, None, i_h)] + gK_head = tma_tensor_K[(None, None, i_h)] + gG_head = tma_tensor_G[(None, None, i_h)] + + # Prefetch chunk 0 -> stage 0 (stage_reuse[0] pre-arrived) + pf_cs = cutlass.Int32(0) + if IS_VARLEN: + pf_seq_id_0 = cutlass.Int32(mChunkIndices[chunk_base, 0]) + pf_local_0 = cutlass.Int32(mChunkIndices[chunk_base, 1]) + pf_bos_0 = cutlass.Int32(mCuSeqlens[pf_seq_id_0]) + pf_cs = pf_bos_0 + pf_local_0 * BT + else: + pf_cs = i_b * num_chunks * BT + chunk_base * BT + cute.arch.mbarrier_wait(stage_reuse_mbars, 0) + if lane_id == 0: + cute.arch.mbarrier_expect_tx(tma_mbars, bytes_per_stage) + sQ_pf = sQ[(None, None, 0)] + gQ_pf = cute.local_tile( + cute.domain_offset((pf_cs, 0), gQ_head), (BT, K_DIM), (0, 0) + ) + ts_pf, tg_pf = cpasync.tma_partition( + tma_atom_Q, + 0, + cute.make_layout(1), + cute.group_modes(sQ_pf, 0, 2), + cute.group_modes(gQ_pf, 0, 2), + ) + cute.copy(tma_atom_Q, tg_pf, ts_pf, tma_bar_ptr=tma_mbars) + sK_pf = sK[(None, None, 0)] + gK_pf = cute.local_tile( + cute.domain_offset((pf_cs, 0), gK_head), (BT, K_DIM), (0, 0) + ) + ts_pf, tg_pf = cpasync.tma_partition( + tma_atom_K, + 0, + cute.make_layout(1), + cute.group_modes(sK_pf, 0, 2), + cute.group_modes(gK_pf, 0, 2), + ) + cute.copy(tma_atom_K, tg_pf, ts_pf, tma_bar_ptr=tma_mbars) + sG_pf = sG[(None, None, 0)] + gG_pf = cute.local_tile( + cute.domain_offset((pf_cs, 0), gG_head), (BT, K_DIM), (0, 0) + ) + ts_pf, tg_pf = cpasync.tma_partition( + tma_atom_G, + 0, + cute.make_layout(1), + cute.group_modes(sG_pf, 0, 2), + cute.group_modes(gG_pf, 0, 2), + ) + cute.copy(tma_atom_G, tg_pf, ts_pf, tma_bar_ptr=tma_mbars) + if lane_id == 0: + cute.arch.mbarrier_arrive(tma_mbars) + + # Issue TMAs for chunks 1..CHUNKS_PER_BLOCK-1 + for next_i in cutlass.range_constexpr(1, CHUNKS_PER_BLOCK): + next_stage = next_i % NUM_STAGES + next_phase = next_i // NUM_STAGES % 2 + next_cs = cutlass.Int32(0) + if IS_VARLEN: + next_chunk_idx = chunk_base + next_i + if next_chunk_idx < num_chunks: + _nsid = cutlass.Int32(mChunkIndices[next_chunk_idx, 0]) + next_cs = ( + cutlass.Int32(mCuSeqlens[_nsid]) + + cutlass.Int32(mChunkIndices[next_chunk_idx, 1]) * BT + ) + else: + next_cs = i_b * num_chunks * BT + (chunk_base + next_i) * BT + cute.arch.mbarrier_wait(stage_reuse_mbars + next_stage, next_phase) + if lane_id == 0: + cute.arch.mbarrier_expect_tx( + tma_mbars + next_stage, bytes_per_stage + ) + sQ_ns = sQ[(None, None, next_stage)] + gQ_ns = cute.local_tile( + cute.domain_offset((next_cs, 0), gQ_head), (BT, K_DIM), (0, 0) + ) + ts_ns, tg_ns = cpasync.tma_partition( + tma_atom_Q, + 0, + cute.make_layout(1), + cute.group_modes(sQ_ns, 0, 2), + cute.group_modes(gQ_ns, 0, 2), + ) + cute.copy(tma_atom_Q, tg_ns, ts_ns, tma_bar_ptr=tma_mbars + next_stage) + sK_ns = sK[(None, None, next_stage)] + gK_ns = cute.local_tile( + cute.domain_offset((next_cs, 0), gK_head), (BT, K_DIM), (0, 0) + ) + ts_ns, tg_ns = cpasync.tma_partition( + tma_atom_K, + 0, + cute.make_layout(1), + cute.group_modes(sK_ns, 0, 2), + cute.group_modes(gK_ns, 0, 2), + ) + cute.copy(tma_atom_K, tg_ns, ts_ns, tma_bar_ptr=tma_mbars + next_stage) + sG_ns = sG[(None, None, next_stage)] + gG_ns = cute.local_tile( + cute.domain_offset((next_cs, 0), gG_head), (BT, K_DIM), (0, 0) + ) + ts_ns, tg_ns = cpasync.tma_partition( + tma_atom_G, + 0, + cute.make_layout(1), + cute.group_modes(sG_ns, 0, 2), + cute.group_modes(gG_ns, 0, 2), + ) + cute.copy(tma_atom_G, tg_ns, ts_ns, tma_bar_ptr=tma_mbars + next_stage) + if lane_id == 0: + cute.arch.mbarrier_arrive(tma_mbars + next_stage) + + # Warps 16-27 (excluding TMA_WARP_ID=26): K2 MMA Compute + if ( + warp_idx >= NUM_K1_TMA_WARPS + and warp_idx < NUM_K1_TMA_WARPS + NUM_MMA_WARPS + and warp_idx != TMA_WARP_ID + ): + # Warp-layout invariants (scope-local → no cross-group register spill) + _tid_in_group = _lane % 4 + _group_id = _lane // 4 + mma_warp = _warp - NUM_K1_TMA_WARPS + my_i_q = cutlass.Int32(0) + my_i_k = cutlass.Int32(0) + if mma_warp < 1: + my_i_q = cutlass.Int32(0) + my_i_k = mma_warp + elif mma_warp < 3: + my_i_q = cutlass.Int32(1) + my_i_k = mma_warp - 1 + elif mma_warp < 6: + my_i_q = cutlass.Int32(2) + my_i_k = mma_warp - 3 + elif mma_warp < NUM_MMA_ACTIVE: + my_i_q = cutlass.Int32(3) + my_i_k = mma_warp - 6 + q_row_base = my_i_q * BC + k_row_base = my_i_k * BC + akk_row_base = k_row_base + akk_col_base = q_row_base + norm_row = q_row_base + if my_i_q == my_i_k: + norm_row = q_row_base + cutlass.Int32(BC // 2) + row0 = _group_id + row1 = _group_id + 8 + col0 = _tid_in_group * 2 + col1 = _tid_in_group * 2 + 1 + col2 = 8 + _tid_in_group * 2 + col3 = 8 + _tid_in_group * 2 + 1 + thr_mma = tiled_mma_k2.get_slice(_lane) + thr_copy_A = tiled_copy_mma_A.get_slice(_lane) + thr_copy_B = tiled_copy_mma_B.get_slice(_lane) + thr_copy_Gn = tiled_copy_Gcum_norm.get_slice(_tid_in_group) + thr_copy_Ggate = tiled_copy_Gcum_gate.get_slice(_lane) + + for chunk_iter in cutlass.range_constexpr(CHUNKS_PER_BLOCK): + s = chunk_iter % NUM_STAGES + phase = chunk_iter // NUM_STAGES % 2 + chunk_idx = chunk_base + chunk_iter + chunk_start = cutlass.Int32(0) + if IS_VARLEN: + if chunk_idx < num_chunks: + _sid = cutlass.Int32(mChunkIndices[chunk_idx, 0]) + chunk_start = ( + cutlass.Int32(mCuSeqlens[_sid]) + + cutlass.Int32(mChunkIndices[chunk_idx, 1]) * BT + ) + else: + chunk_start = chunk_idx * BT + + cute.arch.mbarrier_wait(k1_done_mbars + s, phase) + cute.arch.mbarrier_wait(store_done_mbars + s, phase) + + if mma_warp < NUM_MMA_ACTIVE: + csQ = sQ[(None, None, s)] + csK = sK[(None, None, s)] + csGcum = sGcum[(None, None, s)] + csAqk = sAqk[(None, None, s)] + csAkk = sAkk[(None, None, s)] + + _z = cutlass.Float32(0.0) + + beta_row0 = mBeta[i_b, chunk_start + q_row_base + row0, i_h].to( + cutlass.Float32 + ) + beta_row1 = mBeta[i_b, chunk_start + q_row_base + row1, i_h].to( + cutlass.Float32 + ) + + acc_aqk_n0_0, acc_aqk_n0_1, acc_aqk_n0_2, acc_aqk_n0_3 = ( + _z, + _z, + _z, + _z, + ) + acc_aqk_n1_0, acc_aqk_n1_1, acc_aqk_n1_2, acc_aqk_n1_3 = ( + _z, + _z, + _z, + _z, + ) + acc_akk_n0_0, acc_akk_n0_1, acc_akk_n0_2, acc_akk_n0_3 = ( + _z, + _z, + _z, + _z, + ) + acc_akk_n1_0, acc_akk_n1_1, acc_akk_n1_2, acc_akk_n1_3 = ( + _z, + _z, + _z, + _z, + ) + + # bf16 m16n8k16 MMA: each k_block covers k=16. + for k_block in cutlass.range_constexpr(NUM_MMA_K_TILES): + # ---- Load Q/Kq bf16 fragments (16x16, 8 bf16/thread) ---- + sQ_tile = cute.local_tile( + csQ, tiler=(16, 16), coord=(my_i_q, k_block) + ) + tCrQ = tiled_mma_k2.make_fragment_A( + thr_mma.partition_A(sQ_tile) + ) + cute.copy( + tiled_copy_mma_A, + thr_copy_A.partition_S(sQ_tile), + thr_copy_A.retile(tCrQ), + ) + + sKq_tile = cute.local_tile( + csK, tiler=(16, 16), coord=(my_i_q, k_block) + ) + tCrKq = tiled_mma_k2.make_fragment_A( + thr_mma.partition_A(sKq_tile) + ) + cute.copy( + tiled_copy_mma_A, + thr_copy_A.partition_S(sKq_tile), + thr_copy_A.retile(tCrKq), + ) + + # ---- Issue K n0/n1 LDSMs early for better ILP ---- + sK_tile_n0 = cute.local_tile( + csK, tiler=(8, 16), coord=(my_i_k * 2, k_block) + ) + tCrK_n0 = tiled_mma_k2.make_fragment_B( + thr_mma.partition_B(sK_tile_n0) + ) + cute.copy( + tiled_copy_mma_B, + thr_copy_B.partition_S(sK_tile_n0), + thr_copy_B.retile(tCrK_n0), + ) + + sK_tile_n1 = cute.local_tile( + csK, tiler=(8, 16), coord=(my_i_k * 2 + 1, k_block) + ) + tCrK_n1 = tiled_mma_k2.make_fragment_B( + thr_mma.partition_B(sK_tile_n1) + ) + cute.copy( + tiled_copy_mma_B, + thr_copy_B.partition_S(sK_tile_n1), + thr_copy_B.retile(tCrK_n1), + ) + + # ---- Gate norm (2x k=8 covers k=16) ---- + sGn_a = cute.local_tile( + csGcum, tiler=(1, 8), coord=(norm_row, k_block * 2) + ) + tCsGn_a = thr_copy_Gn.partition_S(sGn_a) + tCrGn_a = cute.make_fragment_like(tCsGn_a, cutlass.Float32) + cute.copy( + tiled_copy_Gcum_norm, tCsGn_a, thr_copy_Gn.retile(tCrGn_a) + ) + gn_a0 = tCrGn_a[0] + gn_a1 = tCrGn_a[1] + + sGn_b = cute.local_tile( + csGcum, tiler=(1, 8), coord=(norm_row, k_block * 2 + 1) + ) + tCsGn_b = thr_copy_Gn.partition_S(sGn_b) + tCrGn_b = cute.make_fragment_like(tCsGn_b, cutlass.Float32) + cute.copy( + tiled_copy_Gcum_norm, tCsGn_b, thr_copy_Gn.retile(tCrGn_b) + ) + gn_b0 = tCrGn_b[0] + gn_b1 = tCrGn_b[1] + + # ---- Gate Q (2x (16,8) partition_C covers m=16,k=16) ---- + sGq_a = cute.local_tile( + csGcum, tiler=(16, 8), coord=(my_i_q, k_block * 2) + ) + tCrGq_a = tiled_mma_k2.make_fragment_C( + thr_mma.partition_C(sGq_a) + ) + cute.copy( + tiled_copy_Gcum_gate, + thr_copy_Ggate.partition_S(sGq_a), + thr_copy_Ggate.retile(tCrGq_a), + ) + + sGq_b = cute.local_tile( + csGcum, tiler=(16, 8), coord=(my_i_q, k_block * 2 + 1) + ) + tCrGq_b = tiled_mma_k2.make_fragment_C( + thr_mma.partition_C(sGq_b) + ) + cute.copy( + tiled_copy_Gcum_gate, + thr_copy_Ggate.partition_S(sGq_b), + thr_copy_Ggate.retile(tCrGq_b), + ) + + # 8 Q gate values per thread (matching A bf16 m16n8k16 layout): + # first half k=0..7 (tCrGq_a): a0=(r0,c0) a1=(r0,c0+1) a2=(r0+8,c0) a3=(r0+8,c0+1) + # second half k=8..15 (tCrGq_b): a4=(r0,c0+8) a5=(r0,c0+9) a6=(r0+8,c0+8) a7=(r0+8,c0+9) + gate_q_0 = cute.exp2(tCrGq_a[0] - gn_a0, fastmath=True) + gate_q_1 = cute.exp2(tCrGq_a[1] - gn_a1, fastmath=True) + gate_q_2 = cute.exp2(tCrGq_a[2] - gn_a0, fastmath=True) + gate_q_3 = cute.exp2(tCrGq_a[3] - gn_a1, fastmath=True) + gate_q_4 = cute.exp2(tCrGq_b[0] - gn_b0, fastmath=True) + gate_q_5 = cute.exp2(tCrGq_b[1] - gn_b1, fastmath=True) + gate_q_6 = cute.exp2(tCrGq_b[2] - gn_b0, fastmath=True) + gate_q_7 = cute.exp2(tCrGq_b[3] - gn_b1, fastmath=True) + + # qa fp32 = Q*gate (8 per thread). tCrQ indexing assumption: + # [0..3] first k-chunk (k=0..7), [4..7] second k-chunk (k=8..15). + qa0 = tCrQ[0].to(cutlass.Float32) * gate_q_0 + qa1 = tCrQ[1].to(cutlass.Float32) * gate_q_1 + qa2 = tCrQ[2].to(cutlass.Float32) * gate_q_2 + qa3 = tCrQ[3].to(cutlass.Float32) * gate_q_3 + qa4 = tCrQ[4].to(cutlass.Float32) * gate_q_4 + qa5 = tCrQ[5].to(cutlass.Float32) * gate_q_5 + qa6 = tCrQ[6].to(cutlass.Float32) * gate_q_6 + qa7 = tCrQ[7].to(cutlass.Float32) * gate_q_7 + + # Pack fp32 qa -> 4 u32 bf16x2 for MMA A (k-adjacent pairs per u32) + qa_u32_0 = pack_bf16x2_f32(qa1, qa0) # reg0 = [qa0 lo | qa1 hi] + qa_u32_1 = pack_bf16x2_f32(qa3, qa2) + qa_u32_2 = pack_bf16x2_f32(qa5, qa4) + qa_u32_3 = pack_bf16x2_f32(qa7, qa6) + + ka0 = tCrKq[0].to(cutlass.Float32) * gate_q_0 + ka1 = tCrKq[1].to(cutlass.Float32) * gate_q_1 + ka2 = tCrKq[2].to(cutlass.Float32) * gate_q_2 + ka3 = tCrKq[3].to(cutlass.Float32) * gate_q_3 + ka4 = tCrKq[4].to(cutlass.Float32) * gate_q_4 + ka5 = tCrKq[5].to(cutlass.Float32) * gate_q_5 + ka6 = tCrKq[6].to(cutlass.Float32) * gate_q_6 + ka7 = tCrKq[7].to(cutlass.Float32) * gate_q_7 + + ka_u32_0 = pack_bf16x2_f32(ka1, ka0) + ka_u32_1 = pack_bf16x2_f32(ka3, ka2) + ka_u32_2 = pack_bf16x2_f32(ka5, ka4) + ka_u32_3 = pack_bf16x2_f32(ka7, ka6) + + # ---- Gate K (2x (16,8) partition_C covers m=16,k=16) ---- + sGk_a = cute.local_tile( + csGcum, tiler=(16, 8), coord=(my_i_k, k_block * 2) + ) + tCrGk_a = tiled_mma_k2.make_fragment_C( + thr_mma.partition_C(sGk_a) + ) + cute.copy( + tiled_copy_Gcum_gate, + thr_copy_Ggate.partition_S(sGk_a), + thr_copy_Ggate.retile(tCrGk_a), + ) + + sGk_b = cute.local_tile( + csGcum, tiler=(16, 8), coord=(my_i_k, k_block * 2 + 1) + ) + tCrGk_b = tiled_mma_k2.make_fragment_C( + thr_mma.partition_C(sGk_b) + ) + cute.copy( + tiled_copy_Gcum_gate, + thr_copy_Ggate.partition_S(sGk_b), + thr_copy_Ggate.retile(tCrGk_b), + ) + + # n0 uses rows 0..7 of (16,*) tile (tCrGk_*[0,1]) + # n1 uses rows 8..15 of (16,*) tile (tCrGk_*[2,3]) + gk_n0_0 = cute.exp2(gn_a0 - tCrGk_a[0], fastmath=True) + gk_n0_1 = cute.exp2(gn_a1 - tCrGk_a[1], fastmath=True) + gk_n0_2 = cute.exp2(gn_b0 - tCrGk_b[0], fastmath=True) + gk_n0_3 = cute.exp2(gn_b1 - tCrGk_b[1], fastmath=True) + + gk_n1_0 = cute.exp2(gn_a0 - tCrGk_a[2], fastmath=True) + gk_n1_1 = cute.exp2(gn_a1 - tCrGk_a[3], fastmath=True) + gk_n1_2 = cute.exp2(gn_b0 - tCrGk_b[2], fastmath=True) + gk_n1_3 = cute.exp2(gn_b1 - tCrGk_b[3], fastmath=True) + + # tCrK_n0 bf16 fragment: 4 elems/thread at (n_row, c0), (n_row, c0+1), + # (n_row, c0+8), (n_row, c0+9) — k-adjacent pairs + k_n0_b0 = tCrK_n0[0].to(cutlass.Float32) * gk_n0_0 + k_n0_b1 = tCrK_n0[1].to(cutlass.Float32) * gk_n0_1 + k_n0_b2 = tCrK_n0[2].to(cutlass.Float32) * gk_n0_2 + k_n0_b3 = tCrK_n0[3].to(cutlass.Float32) * gk_n0_3 + + k_n1_b0 = tCrK_n1[0].to(cutlass.Float32) * gk_n1_0 + k_n1_b1 = tCrK_n1[1].to(cutlass.Float32) * gk_n1_1 + k_n1_b2 = tCrK_n1[2].to(cutlass.Float32) * gk_n1_2 + k_n1_b3 = tCrK_n1[3].to(cutlass.Float32) * gk_n1_3 + + # Pack fp32 k -> 2 u32 bf16x2 for MMA B + k_n0_u32_0 = pack_bf16x2_f32(k_n0_b1, k_n0_b0) + k_n0_u32_1 = pack_bf16x2_f32(k_n0_b3, k_n0_b2) + k_n1_u32_0 = pack_bf16x2_f32(k_n1_b1, k_n1_b0) + k_n1_u32_1 = pack_bf16x2_f32(k_n1_b3, k_n1_b2) + + # ---- 4 bf16 MMA calls ---- + acc_aqk_n0_0, acc_aqk_n0_1, acc_aqk_n0_2, acc_aqk_n0_3 = ( + mma_bf16_m16n8k16( + qa_u32_0, + qa_u32_1, + qa_u32_2, + qa_u32_3, + k_n0_u32_0, + k_n0_u32_1, + acc_aqk_n0_0, + acc_aqk_n0_1, + acc_aqk_n0_2, + acc_aqk_n0_3, + ) + ) + acc_aqk_n1_0, acc_aqk_n1_1, acc_aqk_n1_2, acc_aqk_n1_3 = ( + mma_bf16_m16n8k16( + qa_u32_0, + qa_u32_1, + qa_u32_2, + qa_u32_3, + k_n1_u32_0, + k_n1_u32_1, + acc_aqk_n1_0, + acc_aqk_n1_1, + acc_aqk_n1_2, + acc_aqk_n1_3, + ) + ) + acc_akk_n0_0, acc_akk_n0_1, acc_akk_n0_2, acc_akk_n0_3 = ( + mma_bf16_m16n8k16( + ka_u32_0, + ka_u32_1, + ka_u32_2, + ka_u32_3, + k_n0_u32_0, + k_n0_u32_1, + acc_akk_n0_0, + acc_akk_n0_1, + acc_akk_n0_2, + acc_akk_n0_3, + ) + ) + acc_akk_n1_0, acc_akk_n1_1, acc_akk_n1_2, acc_akk_n1_3 = ( + mma_bf16_m16n8k16( + ka_u32_0, + ka_u32_1, + ka_u32_2, + ka_u32_3, + k_n1_u32_0, + k_n1_u32_1, + acc_akk_n1_0, + acc_akk_n1_1, + acc_akk_n1_2, + acc_akk_n1_3, + ) + ) + + # sQ/sK/sG reads done, signal TMA before SMEM writes + cute.arch.mbarrier_arrive(stage_reuse_mbars + s) + + # Dual-path MMA write (constexpr-gated): + # - non-pure: apply causal + diag=1 inline so SMEM matches + # final GMEM layout (pairs with row-major vec autovec + # store warp that does no causal). + # - pure: write all 16x16 unconditionally (pairs with + # per-tile store warp that applies causal + diag in + # store; this is the original baseline behavior — no + # extra MMA-write cost). + _z16 = cutlass.BFloat16(0.0) + _one16 = cutlass.BFloat16(1.0) + if IS_VARLEN and not VARLEN_PURE: + if my_i_q == my_i_k: + # Diagonal sub-tile: causal-mask sAqk and write + # diag=1 / strict-lower=MMA*beta / strict-upper=0 + # to sAkk so SMEM is in final form. + _v_q00 = (acc_aqk_n0_0 * scale).to(cutlass.BFloat16) + if row0 < col0: + _v_q00 = _z16 + csAqk[q_row_base + row0, k_row_base + col0] = _v_q00 + _v_q01 = (acc_aqk_n0_1 * scale).to(cutlass.BFloat16) + if row0 < col1: + _v_q01 = _z16 + csAqk[q_row_base + row0, k_row_base + col1] = _v_q01 + _v_q02 = (acc_aqk_n0_2 * scale).to(cutlass.BFloat16) + if row1 < col0: + _v_q02 = _z16 + csAqk[q_row_base + row1, k_row_base + col0] = _v_q02 + _v_q03 = (acc_aqk_n0_3 * scale).to(cutlass.BFloat16) + if row1 < col1: + _v_q03 = _z16 + csAqk[q_row_base + row1, k_row_base + col1] = _v_q03 + _v_q04 = (acc_aqk_n1_0 * scale).to(cutlass.BFloat16) + if row0 < col2: + _v_q04 = _z16 + csAqk[q_row_base + row0, k_row_base + col2] = _v_q04 + _v_q05 = (acc_aqk_n1_1 * scale).to(cutlass.BFloat16) + if row0 < col3: + _v_q05 = _z16 + csAqk[q_row_base + row0, k_row_base + col3] = _v_q05 + _v_q06 = (acc_aqk_n1_2 * scale).to(cutlass.BFloat16) + if row1 < col2: + _v_q06 = _z16 + csAqk[q_row_base + row1, k_row_base + col2] = _v_q06 + _v_q07 = (acc_aqk_n1_3 * scale).to(cutlass.BFloat16) + if row1 < col3: + _v_q07 = _z16 + csAqk[q_row_base + row1, k_row_base + col3] = _v_q07 + + _v_k00 = (acc_akk_n0_0 * beta_row0).to(cutlass.BFloat16) + if row0 == col0: + _v_k00 = _one16 + if row0 < col0: + _v_k00 = _z16 + csAkk[akk_row_base + row0, akk_col_base + col0] = _v_k00 + _v_k01 = (acc_akk_n0_1 * beta_row0).to(cutlass.BFloat16) + if row0 == col1: + _v_k01 = _one16 + if row0 < col1: + _v_k01 = _z16 + csAkk[akk_row_base + row0, akk_col_base + col1] = _v_k01 + _v_k02 = (acc_akk_n0_2 * beta_row1).to(cutlass.BFloat16) + if row1 == col0: + _v_k02 = _one16 + if row1 < col0: + _v_k02 = _z16 + csAkk[akk_row_base + row1, akk_col_base + col0] = _v_k02 + _v_k03 = (acc_akk_n0_3 * beta_row1).to(cutlass.BFloat16) + if row1 == col1: + _v_k03 = _one16 + if row1 < col1: + _v_k03 = _z16 + csAkk[akk_row_base + row1, akk_col_base + col1] = _v_k03 + _v_k04 = (acc_akk_n1_0 * beta_row0).to(cutlass.BFloat16) + if row0 == col2: + _v_k04 = _one16 + if row0 < col2: + _v_k04 = _z16 + csAkk[akk_row_base + row0, akk_col_base + col2] = _v_k04 + _v_k05 = (acc_akk_n1_1 * beta_row0).to(cutlass.BFloat16) + if row0 == col3: + _v_k05 = _one16 + if row0 < col3: + _v_k05 = _z16 + csAkk[akk_row_base + row0, akk_col_base + col3] = _v_k05 + _v_k06 = (acc_akk_n1_2 * beta_row1).to(cutlass.BFloat16) + if row1 == col2: + _v_k06 = _one16 + if row1 < col2: + _v_k06 = _z16 + csAkk[akk_row_base + row1, akk_col_base + col2] = _v_k06 + _v_k07 = (acc_akk_n1_3 * beta_row1).to(cutlass.BFloat16) + if row1 == col3: + _v_k07 = _one16 + if row1 < col3: + _v_k07 = _z16 + csAkk[akk_row_base + row1, akk_col_base + col3] = _v_k07 + else: + # Non-diag (i_q > i_k): write all 16x16 unchanged. + csAqk[q_row_base + row0, k_row_base + col0] = ( + acc_aqk_n0_0 * scale + ).to(cutlass.BFloat16) + csAqk[q_row_base + row0, k_row_base + col1] = ( + acc_aqk_n0_1 * scale + ).to(cutlass.BFloat16) + csAqk[q_row_base + row1, k_row_base + col0] = ( + acc_aqk_n0_2 * scale + ).to(cutlass.BFloat16) + csAqk[q_row_base + row1, k_row_base + col1] = ( + acc_aqk_n0_3 * scale + ).to(cutlass.BFloat16) + csAqk[q_row_base + row0, k_row_base + col2] = ( + acc_aqk_n1_0 * scale + ).to(cutlass.BFloat16) + csAqk[q_row_base + row0, k_row_base + col3] = ( + acc_aqk_n1_1 * scale + ).to(cutlass.BFloat16) + csAqk[q_row_base + row1, k_row_base + col2] = ( + acc_aqk_n1_2 * scale + ).to(cutlass.BFloat16) + csAqk[q_row_base + row1, k_row_base + col3] = ( + acc_aqk_n1_3 * scale + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row0, akk_col_base + col0] = ( + acc_akk_n0_0 * beta_row0 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row0, akk_col_base + col1] = ( + acc_akk_n0_1 * beta_row0 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row1, akk_col_base + col0] = ( + acc_akk_n0_2 * beta_row1 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row1, akk_col_base + col1] = ( + acc_akk_n0_3 * beta_row1 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row0, akk_col_base + col2] = ( + acc_akk_n1_0 * beta_row0 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row0, akk_col_base + col3] = ( + acc_akk_n1_1 * beta_row0 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row1, akk_col_base + col2] = ( + acc_akk_n1_2 * beta_row1 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row1, akk_col_base + col3] = ( + acc_akk_n1_3 * beta_row1 + ).to(cutlass.BFloat16) + else: + # PURE: write all 16x16 unconditionally — store warp + # applies causal+diag in its per-tile loop. This is + # the original baseline behavior (no extra MMA cost). + csAqk[q_row_base + row0, k_row_base + col0] = ( + acc_aqk_n0_0 * scale + ).to(cutlass.BFloat16) + csAqk[q_row_base + row0, k_row_base + col1] = ( + acc_aqk_n0_1 * scale + ).to(cutlass.BFloat16) + csAqk[q_row_base + row1, k_row_base + col0] = ( + acc_aqk_n0_2 * scale + ).to(cutlass.BFloat16) + csAqk[q_row_base + row1, k_row_base + col1] = ( + acc_aqk_n0_3 * scale + ).to(cutlass.BFloat16) + csAqk[q_row_base + row0, k_row_base + col2] = ( + acc_aqk_n1_0 * scale + ).to(cutlass.BFloat16) + csAqk[q_row_base + row0, k_row_base + col3] = ( + acc_aqk_n1_1 * scale + ).to(cutlass.BFloat16) + csAqk[q_row_base + row1, k_row_base + col2] = ( + acc_aqk_n1_2 * scale + ).to(cutlass.BFloat16) + csAqk[q_row_base + row1, k_row_base + col3] = ( + acc_aqk_n1_3 * scale + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row0, akk_col_base + col0] = ( + acc_akk_n0_0 * beta_row0 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row0, akk_col_base + col1] = ( + acc_akk_n0_1 * beta_row0 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row1, akk_col_base + col0] = ( + acc_akk_n0_2 * beta_row1 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row1, akk_col_base + col1] = ( + acc_akk_n0_3 * beta_row1 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row0, akk_col_base + col2] = ( + acc_akk_n1_0 * beta_row0 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row0, akk_col_base + col3] = ( + acc_akk_n1_1 * beta_row0 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row1, akk_col_base + col2] = ( + acc_akk_n1_2 * beta_row1 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row1, akk_col_base + col3] = ( + acc_akk_n1_3 * beta_row1 + ).to(cutlass.BFloat16) + else: + cute.arch.mbarrier_arrive(stage_reuse_mbars + s) + + cute.arch.mbarrier_arrive(mma_done_mbars + s) + + # Warps 28-31: Store/Inversion warps + if warp_idx >= NUM_K1_TMA_WARPS + NUM_MMA_WARPS: + store_warp = warp_idx - (NUM_K1_TMA_WARPS + NUM_MMA_WARPS) + for chunk_iter in cutlass.range_constexpr(CHUNKS_PER_BLOCK): + s = chunk_iter % NUM_STAGES + phase = chunk_iter // NUM_STAGES % 2 + chunk_idx = chunk_base + chunk_iter + chunk_start = cutlass.Int32(0) + st_eos = cutlass.Int32(0) + if IS_VARLEN: + if chunk_idx < num_chunks: + _sid = cutlass.Int32(mChunkIndices[chunk_idx, 0]) + chunk_start = ( + cutlass.Int32(mCuSeqlens[_sid]) + + cutlass.Int32(mChunkIndices[chunk_idx, 1]) * BT + ) + st_eos = cutlass.Int32(mCuSeqlens[_sid + 1]) + else: + chunk_start = chunk_idx * BT + + cute.arch.mbarrier_wait(mma_done_mbars + s, phase) + + csAqk = sAqk[(None, None, s)] + csAkk = sAkk[(None, None, s)] + + # Dual-path row-major store. SMEM is in final GMEM layout + # already (causal mask + diag=1 applied at MMA write; upper-tri + # SMEM is zero from CTA-startup init). Both paths use vec2 + # autovec_copy → STG.E.32 (4-byte coalesced). + # + # mAqk_v2 / mAkk_v2 shape (B, T, H, BT/2, 2): last dim is vec2. + # + # NON-PURE: full 64-col row-major + row mask (chunk may + # overflow seq end). 32 lanes/warp × 1 vec2 = full row. + # PURE: reduced cols. Warp s writes (s+1)*16 cols/row (no row + # mask, all rows in seq). Saves ~37% GMEM bandwidth vs full. + col_lo = lane_id * 2 + col_hi = col_lo + 1 + col_vec_idx = lane_id + rAqkOut = cute.make_rmem_tensor( + cute.make_layout((VEC,)), cutlass.BFloat16 + ) + rAkkOut = cute.make_rmem_tensor( + cute.make_layout((VEC,)), cutlass.BFloat16 + ) + + if IS_VARLEN and not VARLEN_PURE: + # NON-PURE: row-major full-row vec autovec + row mask. SMEM + # upper-tri is zero (MMA-masked at write), so writing all 64 + # cols is correct. STG.E.32 (32 lanes × 4 bytes per row). + row_base_warp = store_warp * (BT // NUM_STORE_WARPS) + for ri in cutlass.range_constexpr(BT // NUM_STORE_WARPS): + local_row = row_base_warp + ri + abs_row = chunk_start + local_row + rAqkOut[0] = csAqk[local_row, col_lo] + rAqkOut[1] = csAqk[local_row, col_hi] + rAkkOut[0] = csAkk[local_row, col_lo] + rAkkOut[1] = csAkk[local_row, col_hi] + if abs_row < st_eos: + cute.autovec_copy( + rAqkOut, mAqk_v2[i_b, abs_row, i_h, col_vec_idx, None] + ) + cute.autovec_copy( + rAkkOut, mAkk_v2[i_b, abs_row, i_h, col_vec_idx, None] + ) + else: + # PURE: per-tile loop over 10 lower-tri sub-tiles, reading + # row-major sAqk/sAkk. MMA writes 16x16 unconditional so we + # must apply causal+diag mask in store. Each warp handles + # BC/4=4 rows per sub-tile, lanes 0..15 write 1 bf16 each. + # Original baseline behavior — no MMA-write overhead. + for tile_idx in cutlass.range_constexpr(NUM_TILES): + i_q = _TILE_IQ[tile_idx] + i_k = _TILE_IK[tile_idx] + is_diag = _TILE_IQ[tile_idx] == _TILE_IK[tile_idx] + gmem_aqk_row_base = chunk_start + i_q * BC + gmem_aqk_col_base = i_k * BC + gmem_akk_row_base = chunk_start + i_k * BC + gmem_akk_col_base = i_q * BC + smem_aqk_row_base = i_q * BC + smem_aqk_col_base = i_k * BC + smem_akk_row_base = i_k * BC + smem_akk_col_base = i_q * BC + for ri in cutlass.range_constexpr(BC // NUM_STORE_WARPS): + local_row = store_warp * (BC // NUM_STORE_WARPS) + ri + if lane_id < BC: + local_col = lane_id + aqk_val = csAqk[ + smem_aqk_row_base + local_row, + smem_aqk_col_base + local_col, + ] + akk_val = csAkk[ + smem_akk_row_base + local_row, + smem_akk_col_base + local_col, + ] + if is_diag and local_row < local_col: + aqk_val = cutlass.BFloat16(0.0) + if is_diag and local_row < local_col: + akk_val = cutlass.BFloat16(0.0) + if is_diag and local_row == local_col: + akk_val = cutlass.BFloat16(1.0) + mAqk[ + i_b, + gmem_aqk_row_base + local_row, + i_h, + gmem_aqk_col_base + local_col, + ] = aqk_val + mAkk[ + i_b, + gmem_akk_row_base + local_row, + i_h, + gmem_akk_col_base + local_col, + ] = akk_val + + cute.arch.mbarrier_arrive(store_done_mbars + s) + + yield_out() + + +def make_host_function( + B, + NT, + H, + is_varlen=False, + T_padded=None, + has_bias=False, + use_safe_gate=False, + varlen_pure=False, +): + """ + `varlen_pure=True` asserts that all seq lengths in the batch are multiples + of BT (= 64). Under that assumption every chunk has 64 valid rows so the + four mask sites (K1 row mask, K1 store mask, MMA accumulator zero-fill, + Store row mask) are guaranteed to never fire and are dead-code eliminated + at compile time. Caller's data layout is unchanged — this is a hint only. + """ + _B, _NT, _H = B, NT, H + _IS_VARLEN = 1 if is_varlen else 0 + _HAS_BIAS = 1 if has_bias else 0 + _USE_SAFE_GATE = 1 if use_safe_gate else 0 + _VARLEN_PURE = 1 if (is_varlen and varlen_pure) else 0 + if is_varlen: + assert _B == 1, "Varlen requires B=1" + assert T_padded is not None, "T_padded required for varlen" + _T = T_padded + else: + _T = _NT * BT + + if is_varlen: + _total_cgs_val = ((_NT + CHUNKS_PER_BLOCK - 1) // CHUNKS_PER_BLOCK) * _H + else: + _total_cgs_val = (_NT // CHUNKS_PER_BLOCK) * _H * _B + + # 3D TMA view: (B*T, K_DIM, H) — domain_offset handles non-aligned addressing + _T_total = _B * _T + s_row = _H * K_DIM + s_col = 1 + s_h = K_DIM + + @cute.jit + def host_fn( + mQ, + mK, + mG, + mA_log, + mBeta, + scale, + mKscaled, + mKg, + mQscaled, + mGkLast, + mAqk, + mAkk, + mCuSeqlens, + mChunkIndices, + mDtBias, + lower_bound_val, + stream: cuda_drv.CUstream, + ): + # 3D TMA view: (B*T, K_DIM, H). domain_offset in kernel shifts to + # arbitrary chunk_start — no BT alignment required for varlen. + view_layout_3d = cute.make_layout( + (_T_total, K_DIM, _H), + stride=(s_row, s_col, s_h), + ) + mQ_view = cute.make_tensor(mQ.iterator, view_layout_3d) + mK_view = cute.make_tensor(mK.iterator, view_layout_3d) + mG_view = cute.make_tensor(mG.iterator, view_layout_3d) + + smem_atom_qk = tcgen05.make_smem_layout_atom( + tcgen05.SmemLayoutAtomKind.K_SW128, cutlass.BFloat16 + ) + qk_smem_2d = cute.tile_to_shape(smem_atom_qk, (BT, K_DIM), order=(0, 1)) + qk_smem_3d = cute.tile_to_shape( + smem_atom_qk, (BT, K_DIM, NUM_STAGES), order=(0, 1, 2) + ) + + g_smem_2d = cute.make_layout((BT, K_DIM), stride=(K_DIM, 1)) + g_smem_3d = cute.make_layout( + (BT, K_DIM, NUM_STAGES), stride=(K_DIM, 1, BT * K_DIM) + ) + + tma_op = cpasync.CopyBulkTensorTileG2SOp(cpasync.CtaGroup.ONE) + tma_atom_Q, tma_tensor_Q = cpasync.make_tiled_tma_atom( + tma_op, + mQ_view, + qk_smem_2d, + cute.product_each(qk_smem_2d.shape), + num_multicast=1, + ) + tma_atom_K, tma_tensor_K = cpasync.make_tiled_tma_atom( + tma_op, + mK_view, + qk_smem_2d, + cute.product_each(qk_smem_2d.shape), + num_multicast=1, + ) + tma_atom_G, tma_tensor_G = cpasync.make_tiled_tma_atom( + tma_op, + mG_view, + g_smem_2d, + cute.product_each(g_smem_2d.shape), + num_multicast=1, + ) + + g_cumsum_layout = cute.make_layout( + (BT, K_DIM, NUM_STAGES), stride=(K_STRIDE, 1, BT * K_STRIDE) + ) + + copy_atom_qk_k1 = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), cutlass.BFloat16, num_bits_per_copy=32 + ) + tiled_copy_qk_k1 = cute.make_tiled_copy_tv( + copy_atom_qk_k1, + thr_layout=cute.make_layout((1, 32)), + val_layout=cute.make_layout((1, 2)), + ) + + out_v2_layout = cute.make_layout( + (_B, _T, _H, K_VEC, VEC), + stride=(_T * _H * K_DIM, _H * K_DIM, K_DIM, VEC, 1), + ) + mKscaled_v2 = cute.make_tensor(mKscaled.iterator, out_v2_layout) + mQscaled_v2 = cute.make_tensor(mQscaled.iterator, out_v2_layout) + mKg_v2 = cute.make_tensor(mKg.iterator, out_v2_layout) + + # vec2 views for mAqk / mAkk (BT dimension instead of K_DIM). + # Shape: (B, T, H, BT/2, 2). Each VEC=2 slot = 2 contiguous bf16 = 1 + # fp32-aligned 4-byte unit. Used by store warp's autovec_copy → + # STG.E.32 with 32 lanes coalesced to 1 cache line per row. + BT_VEC = BT // VEC # 32 + akk_v2_layout = cute.make_layout( + (_B, _T, _H, BT_VEC, VEC), + stride=(_T * _H * BT, _H * BT, BT, VEC, 1), + ) + mAqk_v2 = cute.make_tensor(mAqk.iterator, akk_v2_layout) + mAkk_v2 = cute.make_tensor(mAkk.iterator, akk_v2_layout) + + gklast_v2_layout = cute.make_layout( + (_B, _NT, _H, K_VEC, VEC), + stride=(_NT * _H * K_DIM, _H * K_DIM, K_DIM, VEC, 1), + ) + mGkLast_v2 = cute.make_tensor(mGkLast.iterator, gklast_v2_layout) + + mma_op = cute.nvgpu.warp.MmaF16BF16Op( + cutlass.BFloat16, cutlass.Float32, (16, 8, 16) + ) + tiled_mma_k2 = cute.make_tiled_mma( + mma_op, cute.make_layout((1, 1, 1)), permutation_mnk=(16, 8, 16) + ) + + tiled_copy_mma_A = cute.make_tiled_copy_A( + cute.make_copy_atom( + cute.nvgpu.warp.LdMatrix8x8x16bOp(False, 4), cutlass.BFloat16 + ), + tiled_mma_k2, + ) + tiled_copy_mma_B = cute.make_tiled_copy_B( + cute.make_copy_atom( + cute.nvgpu.warp.LdMatrix8x8x16bOp(False, 2), cutlass.BFloat16 + ), + tiled_mma_k2, + ) + + copy_atom_Gcum = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), cutlass.Float32, num_bits_per_copy=64 + ) + tiled_copy_Gcum_norm = cute.make_tiled_copy_tv( + copy_atom_Gcum, + thr_layout=cute.make_layout((1, 4)), + val_layout=cute.make_layout((1, 2)), + ) + tiled_copy_Gcum_gate = cute.make_tiled_copy_C(copy_atom_Gcum, tiled_mma_k2) + + smem_size = ( + BT * K_DIM * 2 * 2 * NUM_STAGES + + BT * K_DIM * 2 * NUM_STAGES + + BT * K_STRIDE * 4 * NUM_STAGES + + K1_ROW_GROUPS * PARTIAL_COLS * 4 + + BT * AQK_TILE_STRIDE * 2 * NUM_STAGES # sAqk bf16 (64x72 row-major) + + BT * AKK_STRIDE * 2 * NUM_STAGES # sAkk bf16 + + 512 + ) + + _grid_x = min(NUM_SMS, _total_cgs_val) + + fused_kernel123( + tma_atom_Q, + tma_tensor_Q, + tma_atom_K, + tma_tensor_K, + tma_atom_G, + tma_tensor_G, + mA_log, + mBeta, + scale, + mKscaled_v2, + mKg_v2, + mQscaled_v2, + mGkLast_v2, + mAqk, + mAkk, + mAqk_v2, + mAkk_v2, + tiled_copy_qk_k1, + tiled_mma_k2, + tiled_copy_mma_A, + tiled_copy_mma_B, + tiled_copy_Gcum_norm, + tiled_copy_Gcum_gate, + qk_smem_3d, + g_smem_3d, + g_cumsum_layout, + _NT, + _H, + _B, + mCuSeqlens, + mChunkIndices, + _IS_VARLEN, + mDtBias, + lower_bound_val, + _HAS_BIAS, + _USE_SAFE_GATE, + _VARLEN_PURE, + ).launch( + grid=(_grid_x, 1, 1), + block=(THREADS, 1, 1), + smem=smem_size, + stream=stream, + ) + + return host_fn diff --git a/python/sglang/kernels/ops/attention/linear/kda_ptx_prefill/__init__.py b/python/sglang/kernels/ops/attention/linear/kda_ptx_prefill/__init__.py new file mode 100644 index 000000000..1a24cc0a7 --- /dev/null +++ b/python/sglang/kernels/ops/attention/linear/kda_ptx_prefill/__init__.py @@ -0,0 +1,207 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Hand-written PTX/tcgen05 KDA chunked-prefill kernel (GB300 / sm_103a). + +Vendored from the upstream ``kda_prefill`` artifact (commit 33583615): the CUDA +source in ``kernels/jit/csrc/attention/kda_prefill.cu`` plus the FLA-signature +``chunk_kda_fwd`` wrapper below, which is a drop-in for +``fla.ops.kda.chunk_fwd.chunk_kda_fwd`` on the inference forward path +(K = V = 128, chunk_size = 64). + +The extension is JIT-compiled with ``torch.utils.cpp_extension`` on first use +(~1-2 min, cached under ``TORCH_EXTENSIONS_DIR``); concurrent TP ranks +serialize on torch's build lock and then load the cached .so. +""" + +import os + +import torch + +K = 128 # head dim (qk == v), fixed by the kernel +CHUNK = 64 # kernel chunk size, fixed + +_EXT_NAME = "kda_prefill_ptx" +_ext = None + + +def load_ext(): + """JIT-load the CUDA extension (cached after first call).""" + global _ext + if _ext is None: + from torch.utils import cpp_extension + + src = os.path.join( + os.path.dirname(os.path.abspath(__file__)), + "../../../../jit/csrc/attention/kda_prefill.cu", + ) + # -lcuda: cuTensorMapEncodeTiled (driver API). The stubs dir covers + # boxes whose real libcuda.so lives off the default linker path; + # ld.so still binds the driver's libcuda.so.1 at import time. + stubs = os.path.join( + cpp_extension.CUDA_HOME or "/usr/local/cuda", "lib64", "stubs" + ) + _ext = cpp_extension.load( + name=_EXT_NAME, + sources=[src], + extra_cuda_cflags=[ + "-O3", + "-std=c++20", + "-use_fast_math", + "-lineinfo", + "-gencode", + "arch=compute_103a,code=sm_103a", + ], + extra_cflags=["-O3"], + extra_ldflags=[f"-L{stubs}", "-lcuda"], + ) + return _ext + + +def chunk_kda_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor | None, + output_final_state: bool, + state_v_first: bool = False, + cu_seqlens: torch.Tensor | None = None, + cu_seqlens_cpu: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = 64, + safe_gate: bool = False, + lower_bound: float | None = None, + use_gate_in_kernel: bool = False, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + disable_recompute: bool = False, + return_intermediate_states: bool = False, + cp_context=None, + use_qk_l2norm_in_kernel: bool = False, + use_beta_sigmoid_in_kernel: bool = False, + allow_neg_eigval: bool = False, +): + """Inference drop-in for fla's chunk_kda_fwd (forward only). + + Argument mapping onto the CUDA kernel: + q/k/v [B,T,H,128] bf16 -> flat [T,H,128] (B>1 folds to varlen with a + synthetic equal-length cu_seqlens; B==1 squeezes). + g [B,T,H,128]: use_gate_in_kernel=False -> pre-transformed glog, fp32 + narrowed to bf16 here (the kernel's GM=0 contract); =True -> RAW gate + input (bf16) with the transform fused in-kernel (GM=1 softplus / + GM=2 safe-gate). Following fla, the safe-gate TRANSFORM is selected + by `lower_bound is not None`; the `safe_gate` flag alone is fla's + intra-path hint and does not change the math here. + beta [B,T,H] bf16 (fp32 also accepted; widened to fp32 in the ext). + use_qk_l2norm_in_kernel=True accepts raw q/k and applies FLA-compatible + L2 normalization (eps=1e-6, bf16 rounding) in the CUDA tile loads. + use_beta_sigmoid_in_kernel=True accepts beta logits and fuses sigmoid. + cu_seqlens: host values are needed for the kernel's per-sequence piece + table -- pass cu_seqlens_cpu to avoid the D2H sync; chunk_indices is + accepted and ignored (the kernel derives its own piece table). + initial_state [N,H,128,128] fp32 or None (zeros). + return_intermediate_states=True returns dense fp32 chunk-boundary states + [1, NT, H, 128, 128] at tuple index 10. + + Returns the fla-shaped 12-tuple: (o [B,T,H,128] bf16, final_state + [N,H,128,128] fp32 or None, then Nones, ..., h, initial_state). + """ + assert ( + chunk_size == CHUNK + ), f"kda_prefill supports chunk_size={CHUNK} only, got {chunk_size}" + if cp_context is not None or disable_recompute: + raise NotImplementedError( + "kda_prefill is the inference forward path: cp_context, " + "and disable_recompute are training-side knobs it does not implement" + ) + if allow_neg_eigval and use_beta_sigmoid_in_kernel: + raise NotImplementedError( + "allow_neg_eigval=True requires 2*sigmoid(beta), which is not " + "implemented by the fused beta path; pass pre-activated beta with " + "use_beta_sigmoid_in_kernel=False" + ) + if state_v_first and initial_state is not None: + # [V,K]-layout state: pure transpose (K==V==128), exact, ~us/call + initial_state = initial_state.transpose(-1, -2).contiguous() + assert ( + q.dim() == 4 and q.shape[-1] == K and v.shape[-1] == K + ), f"expected [B,T,H,{K}] q/k/v, got q={tuple(q.shape)} v={tuple(v.shape)}" + B, T, H, _ = q.shape + + cu_cpu = None + if cu_seqlens is not None or cu_seqlens_cpu is not None: + assert B == 1, "cu_seqlens requires B == 1 (flattened varlen batch)" + src = cu_seqlens_cpu if cu_seqlens_cpu is not None else cu_seqlens + cu_cpu = torch.as_tensor(src, dtype=torch.int32).cpu() + elif B > 1: + # eqlen batch == varlen with equal lengths (host-known, no sync) + cu_cpu = torch.arange(0, (B + 1) * T, T, dtype=torch.int32) + + Tt = B * T + qf = q.reshape(Tt, H, K).contiguous() + kf = k.reshape(Tt, H, K).contiguous() + vf = v.reshape(Tt, H, K).contiguous() + betaf = beta.reshape(Tt, H).contiguous() + + if use_gate_in_kernel: + assert ( + A_log is not None and dt_bias is not None + ), "use_gate_in_kernel=True requires A_log and dt_bias" + assert g.dtype == torch.bfloat16, f"raw gate input must be bf16, got {g.dtype}" + gf = g.reshape(Tt, H, K).contiguous() + sg = lower_bound is not None # fla: lb presence selects safe-gate + a_log_flat = A_log.reshape(H).to(torch.float32).contiguous() + dtb = dt_bias.reshape(H * K).to(torch.float32).contiguous() + lb = float(lower_bound) if lower_bound is not None else 0.0 + else: + gf = g.reshape(Tt, H, K) + if gf.dtype != torch.bfloat16: # pre-transformed glog: bf16 narrow + gf = gf.to(torch.bfloat16) + gf = gf.contiguous() + sg, a_log_flat, dtb, lb = False, None, None, 0.0 + + h = None + if return_intermediate_states: + lens = (cu_cpu[1:] - cu_cpu[:-1]).tolist() if cu_cpu is not None else [Tt] + nt = sum((int(length) + CHUNK - 1) // CHUNK for length in lens) + h = torch.empty(nt, H, K, K, dtype=torch.float32, device=q.device) + + o, Sf = load_ext().kda_prefill_fwd( + qf, + kf, + vf, + gf, + betaf, + float(scale), + initial_state=initial_state, + cu_seqlens=cu_cpu, + use_gate_in_kernel=use_gate_in_kernel, + A_log=a_log_flat, + dt_bias=dtb, + safe_gate=sg, + lower_bound=lb, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + use_beta_sigmoid_in_kernel=use_beta_sigmoid_in_kernel, + h_per_chunk=h, + h_v_first=state_v_first, + ) + + o = o.view(B, T, H, K) + if state_v_first: + Sf = Sf.transpose(-1, -2).contiguous() + final = Sf if output_final_state else None + return ( + o, + final, + None, + None, + None, + None, + None, + None, + None, + None, + None if h is None else h.unsqueeze(0), + initial_state, + ) diff --git a/python/sglang/kernels/ops/attention/set_mla_kv_concat_q.py b/python/sglang/kernels/ops/attention/set_mla_kv_concat_q.py new file mode 100644 index 000000000..729b1928a --- /dev/null +++ b/python/sglang/kernels/ops/attention/set_mla_kv_concat_q.py @@ -0,0 +1,318 @@ +"""Fused MLA decode prepare tail: paged-KV scatter + absorbed-q concat. + +One launch replacing the back-to-back ``set_mla_kv_buffer`` + +``concat_mla_absorb_q`` pair on the trtllm-mla decode graph path. Both +workloads are launch-bound data movement at decode batch sizes; the fusion +saves a kernel launch per MLA layer and keeps the PDL chain to the decode +fmha kernel intact. SM90+ only (TMA bulk store) — gate via ``covered()``. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from sglang.kernels.jit.utils import ( + cache_once, + is_arch_support_pdl, + load_jit, + make_cpp_args, +) + +if TYPE_CHECKING: + from tvm_ffi.module import Module + + +@cache_once +def set_mla_kv_concat_q_module( + nope_bytes: int, rope_bytes: int, use_pdl: bool +) -> Module: + args = make_cpp_args(nope_bytes, rope_bytes, use_pdl) + return load_jit( + f"set_mla_kv_concat_q_{nope_bytes}_{rope_bytes}", + *args, + cuda_files=["elementwise/set_mla_kv_concat_q.cuh"], + cuda_wrappers=[ + ("set_mla_kv_concat_q", f"SetMlaKVConcatQKernel<{args}>::run"), + ], + ) + + +@cache_once +def can_use_set_mla_kv_concat_q(nope_bytes: int, rope_bytes: int) -> bool: + """Whether the fused kernel supports these row byte widths on this arch. + + Static gate (per process): SM90+ for the TMA bulk store, plus the + compile-time width constraints (total row 16B-aligned for TMA; nope dim + filling whole int4 warp rounds and rope dim exactly one int warp round + for the concat side — 1024/128 bytes i.e. 512/64 bf16 satisfies both). + """ + if torch.cuda.get_device_capability()[0] < 9: + return False + if nope_bytes % 4 != 0 or rope_bytes % 4 != 0: + return False + if (nope_bytes + rope_bytes) % 16 != 0: + return False + # Concat-side vector layout: int4 (8 bf16) x 32 lanes per round for nope, + # one int (2 bf16) x 32 lanes round for rope. + if (nope_bytes // 2) % (8 * 32) != 0 or rope_bytes // 2 != 2 * 32: + return False + try: + set_mla_kv_concat_q_module(nope_bytes, rope_bytes, is_arch_support_pdl()) + return True + except Exception: # pragma: no cover - compile-time only + return False + + +def _row_aligned(t: torch.Tensor, align: int) -> bool: + """Base pointer and all non-unit strides aligned to ``align`` bytes.""" + if t.data_ptr() % align != 0: + return False + esize = t.element_size() + return all(s * esize % align == 0 for s in t.stride()[:-1]) + + +def covered( + kv_buffer: torch.Tensor, + loc: torch.Tensor, + k_nope: torch.Tensor, + k_rope: torch.Tensor, + q_nope: torch.Tensor, + q_rope: torch.Tensor, +) -> bool: + """Per-call gate mirroring the launcher's alignment/layout tripwires, so + uncovered layouts fall back to the two-kernel path instead of faulting. + + Expects the already-flattened views the wrapper launches with: + kv_buffer [pages, D], k_nope/k_rope [B, d], q_nope/q_rope [B, H, d], + loc [B]. + """ + if not ( + kv_buffer.dtype == torch.bfloat16 + and k_nope.dtype == torch.bfloat16 + and k_rope.dtype == torch.bfloat16 + and q_nope.dtype == torch.bfloat16 + and q_rope.dtype == torch.bfloat16 + ): + return False + if loc.dtype not in (torch.int32, torch.int64): + return False + if loc.dim() != 1 or not loc.is_contiguous(): + return False + if not ( + k_nope.shape[0] + == k_rope.shape[0] + == loc.shape[0] + == q_nope.shape[0] + == q_rope.shape[0] + ): + return False + if q_nope.shape[1] != q_rope.shape[1]: + return False + nope_bytes = k_nope.shape[-1] * 2 + rope_bytes = k_rope.shape[-1] * 2 + if q_nope.shape[-1] * 2 != nope_bytes or q_rope.shape[-1] * 2 != rope_bytes: + return False + if not can_use_set_mla_kv_concat_q(nope_bytes, rope_bytes): + return False + # Last dims must be dense for the vectorised row accesses. + if any(t.stride(-1) != 1 for t in (kv_buffer, k_nope, k_rope, q_nope, q_rope)): + return False + return ( + _row_aligned(kv_buffer, 16) + and _row_aligned(k_nope, 16) + and _row_aligned(k_rope, 4) + and _row_aligned(q_nope, 16) + and _row_aligned(q_rope, 4) + ) + + +def _pick_num_warps(total_items: int) -> int: + # Same GB300 heuristic as set_mla_kv_buffer: small grids favour more CTAs. + return 4 if total_items <= 768 else 8 + + +def set_mla_kv_concat_q( + kv_buffer: torch.Tensor, + loc: torch.Tensor, + cache_k_nope: torch.Tensor, + cache_k_rope: torch.Tensor, + q_nope: torch.Tensor, + q_rope: torch.Tensor, + num_warps: int = 0, +) -> torch.Tensor: + """Scatter [k_nope | k_rope] rows into ``kv_buffer`` at ``loc`` and return + the concatenated query [q_nope | q_rope], all in one kernel launch. + + Shapes (leading singleton dims on the k sources are flattened away): + kv_buffer: [num_pages, total_dim] or [num_pages, 1, total_dim] + loc: [n_loc] + cache_k_nope: [n_loc, nope_dim] or [n_loc, 1, nope_dim] + cache_k_rope: [n_loc, rope_dim] or [n_loc, 1, rope_dim] + q_nope: [n_loc, num_heads, nope_dim] + q_rope: [n_loc, num_heads, rope_dim] + Returns: + query: [n_loc, num_heads, nope_dim + rope_dim] (new contiguous tensor) + """ + n_loc = loc.shape[0] + src_nope = cache_k_nope.view(n_loc, -1) if cache_k_nope.dim() != 2 else cache_k_nope + src_rope = cache_k_rope.view(n_loc, -1) if cache_k_rope.dim() != 2 else cache_k_rope + buf = kv_buffer.view(kv_buffer.shape[0], -1) if kv_buffer.dim() != 2 else kv_buffer + + q_out = torch.empty( + (*q_nope.shape[:-1], q_nope.shape[-1] + q_rope.shape[-1]), + dtype=q_nope.dtype, + device=q_nope.device, + ) + + nope_bytes = src_nope.shape[-1] * src_nope.element_size() + rope_bytes = src_rope.shape[-1] * src_rope.element_size() + if num_warps <= 0: + num_warps = _pick_num_warps(n_loc + q_nope.shape[0] * q_nope.shape[1]) + + module = set_mla_kv_concat_q_module(nope_bytes, rope_bytes, is_arch_support_pdl()) + module.set_mla_kv_concat_q( + buf, loc, src_nope, src_rope, q_nope, q_rope, q_out, num_warps + ) + return q_out + + +@cache_once +def set_mla_kv_concat_q_fp8_module(use_pdl: bool) -> Module: + args = make_cpp_args(use_pdl) + return load_jit( + "set_mla_kv_concat_q_fp8", + *args, + cuda_files=["elementwise/set_mla_kv_concat_q.cuh"], + cuda_wrappers=[ + ("set_mla_kv_concat_q_fp8", f"SetMlaKVConcatQFp8Kernel<{args}>::run"), + ], + ) + + +@cache_once +def can_use_set_mla_kv_concat_q_fp8() -> bool: + """SM90+ (TMA bulk store) and the module compiles. Row widths are fixed + at 512/64 (the MLA absorb layout) inside the kernel.""" + if torch.cuda.get_device_capability()[0] < 9: + return False + try: + set_mla_kv_concat_q_fp8_module(is_arch_support_pdl()) + return True + except Exception: # pragma: no cover - compile-time only + return False + + +def covered_fp8( + kv_buffer: torch.Tensor, + loc: torch.Tensor, + k_nope: torch.Tensor, + k_rope: torch.Tensor, + q_nope: torch.Tensor, + q_rope: torch.Tensor, +) -> bool: + """Per-call gate for the fused fp8 quantize+scatter+concat kernel, + mirroring the launcher tripwires. Expects flattened views: kv_buffer + [pages, 576] fp8/uint8, k halves [B, 512/64] bf16, q halves + [B, H, 512/64] bf16, loc [B].""" + if not ( + kv_buffer.dtype in (torch.float8_e4m3fn, torch.uint8) + and k_nope.dtype == torch.bfloat16 + and k_rope.dtype == torch.bfloat16 + and q_nope.dtype == torch.bfloat16 + and q_rope.dtype == torch.bfloat16 + ): + return False + if loc.dtype not in (torch.int32, torch.int64): + return False + if loc.dim() != 1 or not loc.is_contiguous(): + return False + if not ( + k_nope.shape[0] + == k_rope.shape[0] + == loc.shape[0] + == q_nope.shape[0] + == q_rope.shape[0] + ): + return False + if q_nope.shape[1] != q_rope.shape[1]: + return False + if ( + k_nope.shape[-1] != 512 + or k_rope.shape[-1] != 64 + or q_nope.shape[-1] != 512 + or q_rope.shape[-1] != 64 + or kv_buffer.shape[-1] < 576 + ): + return False + if not can_use_set_mla_kv_concat_q_fp8(): + return False + if any(t.stride(-1) != 1 for t in (kv_buffer, k_nope, k_rope, q_nope, q_rope)): + return False + if kv_buffer.data_ptr() % 16 != 0 or kv_buffer.stride(0) % 16 != 0: + return False + return ( + _row_aligned(k_nope, 16) + and _row_aligned(k_rope, 4) + and _row_aligned(q_nope, 16) + and _row_aligned(q_rope, 4) + ) + + +def set_mla_kv_concat_q_fp8( + kv_buffer: torch.Tensor, + loc: torch.Tensor, + cache_k_nope: torch.Tensor, + cache_k_rope: torch.Tensor, + q_nope: torch.Tensor, + q_rope: torch.Tensor, + num_warps: int = 0, + dcp_world_size: int = 1, + dcp_rank: int = 0, +) -> torch.Tensor: + """Quantize bf16 [k_nope | k_rope] rows to fp8-e4m3 and scatter them into + ``kv_buffer`` at ``loc``, and return the fp8 concatenated query + [q_nope | q_rope], all in one kernel launch (replaces concat + three + aten fp8 casts + the KV-row write on the fp8 decode path). + + Under DCP, ``loc`` is VIRTUAL: the physical row is ``loc // + dcp_world_size`` and only the owner rank (``loc % dcp_world_size == + dcp_rank``) writes its KV row (query conversion still runs for every + token). world=1/rank=0 is the non-DCP identity. + + Shapes (leading singleton dims on the k sources are flattened away): + kv_buffer: [num_pages, 576] fp8_e4m3/uint8 (or [num_pages, 1, 576]) + loc: [n_loc] + cache_k_nope: [n_loc, 512] bf16 cache_k_rope: [n_loc, 64] bf16 + q_nope: [n_loc, H, 512] bf16 q_rope: [n_loc, H, 64] bf16 + Returns: + query: [n_loc, H, 576] float8_e4m3fn (new contiguous tensor) + """ + n_loc = loc.shape[0] + src_nope = cache_k_nope.view(n_loc, -1) if cache_k_nope.dim() != 2 else cache_k_nope + src_rope = cache_k_rope.view(n_loc, -1) if cache_k_rope.dim() != 2 else cache_k_rope + buf = kv_buffer.view(kv_buffer.shape[0], -1) if kv_buffer.dim() != 2 else kv_buffer + + q_out = torch.empty( + (q_nope.shape[0], q_nope.shape[1], 576), + dtype=torch.float8_e4m3fn, + device=q_nope.device, + ) + if num_warps <= 0: + num_warps = _pick_num_warps(n_loc + q_nope.shape[0] * q_nope.shape[1]) + + module = set_mla_kv_concat_q_fp8_module(is_arch_support_pdl()) + module.set_mla_kv_concat_q_fp8( + buf, + loc, + src_nope, + src_rope, + q_nope, + q_rope, + q_out, + num_warps, + dcp_world_size, + dcp_rank, + ) + return q_out diff --git a/python/sglang/kernels/ops/attention/verify_splitkv.py b/python/sglang/kernels/ops/attention/verify_splitkv.py index dd6d3fa35..e262124c6 100644 --- a/python/sglang/kernels/ops/attention/verify_splitkv.py +++ b/python/sglang/kernels/ops/attention/verify_splitkv.py @@ -149,6 +149,8 @@ def _verify_prefix_stage1( kv_group_num: tl.constexpr, N_SPLITS: tl.constexpr, L_EXT: tl.constexpr, # padded power-of-2 row tile (>= real l_ext) + HEAD_DIM: tl.constexpr, + V_HEAD_DIM: tl.constexpr, BLOCK_DMODEL: tl.constexpr, BLOCK_DV: tl.constexpr, BLOCK_N: tl.constexpr, @@ -190,7 +192,11 @@ def _verify_prefix_stage1( + cur_head * stride_qh + offs_d[None, :] ) - q = tl.load(Q + offs_q, mask=mask_l[:, None], other=0.0) + q = tl.load( + Q + offs_q, + mask=mask_l[:, None] & (offs_d[None, :] < HEAD_DIM), + other=0.0, + ) q_k = q.to(K_Buffer.dtype.element_ty) base_offs_k = cur_kv_head * stride_buf_kh + offs_d[:, None] @@ -206,7 +212,11 @@ def _verify_prefix_stage1( ) # K block: [D, BLOCK_N] offs_buf_k = kv_loc[None, :] * stride_buf_kbs + base_offs_k - k = tl.load(K_Buffer + offs_buf_k, mask=n_mask[None, :], other=0.0) + k = tl.load( + K_Buffer + offs_buf_k, + mask=(offs_d[:, None] < HEAD_DIM) & n_mask[None, :], + other=0.0, + ) qk = tl.dot(q_k, k) # [L_EXT, BLOCK_N] qk *= sm_scale * k_scale # fp8 dequant of prefix K (k_scale==1 if bf16) # NO causal mask: full prefix is visible to all draft tokens. @@ -214,7 +224,11 @@ def _verify_prefix_stage1( # V block: [BLOCK_N, Dv] offs_buf_v = kv_loc[:, None] * stride_buf_vbs + base_offs_v - v = tl.load(V_Buffer + offs_buf_v, mask=n_mask[:, None], other=0.0) + v = tl.load( + V_Buffer + offs_buf_v, + mask=n_mask[:, None] & (offs_dv[None, :] < V_HEAD_DIM), + other=0.0, + ) n_e_max = tl.maximum(tl.max(qk, 1), e_max) re_scale = tl.exp(e_max - n_e_max) @@ -234,7 +248,11 @@ def _verify_prefix_stage1( + offs_l[:, None] * stride_ol + offs_dv[None, :] ) - tl.store(Att_Out + offs_o, acc / e_sum[:, None], mask=mask_l[:, None]) + tl.store( + Att_Out + offs_o, + acc / e_sum[:, None], + mask=mask_l[:, None] & (offs_dv[None, :] < V_HEAD_DIM), + ) offs_lse = ( cur_batch * stride_lb @@ -286,6 +304,8 @@ def _verify_combine_stage2( kv_group_num: tl.constexpr, N_SPLITS: tl.constexpr, L_EXT: tl.constexpr, + HEAD_DIM: tl.constexpr, + V_HEAD_DIM: tl.constexpr, BLOCK_DMODEL: tl.constexpr, BLOCK_DV: tl.constexpr, ): @@ -310,7 +330,11 @@ def _verify_combine_stage2( + offs_s[:, None] * stride_ls + offs_l[None, :] ) - lse = tl.load(offs_lse + Att_Lse) # [N_SPLITS, L_EXT] + lse = tl.load( + offs_lse + Att_Lse, + mask=mask_l[None, :], + other=float("-inf"), + ) # [N_SPLITS, L_EXT] m_p = tl.max(lse, 0) # [L_EXT] w = tl.exp(lse - m_p[None, :]) # [N_SPLITS, L_EXT]; -inf->0 denom_p = tl.sum(w, 0) # [L_EXT] @@ -324,7 +348,11 @@ def _verify_combine_stage2( + offs_l[None, :, None] * stride_ol + offs_dv[None, None, :] ) - ao = tl.load(offs_ao + Att_Out) # [N_SPLITS, L_EXT, Dv] + ao = tl.load( + offs_ao + Att_Out, + mask=mask_l[None, :, None] & (offs_dv[None, None, :] < V_HEAD_DIM), + other=0.0, + ) # [N_SPLITS, L_EXT, Dv] o_prefix = tl.sum(ao * w[:, :, None], 0) # [L_EXT, Dv] o_prefix = o_prefix / denom_p[:, None] lse_prefix = m_p + tl.log(denom_p) # [L_EXT] @@ -336,20 +364,32 @@ def _verify_combine_stage2( + cur_head * stride_qh + offs_d[None, :] ) - q = tl.load(Q + offs_q, mask=mask_l[:, None], other=0.0).to(tl.float32) + q = tl.load( + Q + offs_q, + mask=mask_l[:, None] & (offs_d[None, :] < HEAD_DIM), + other=0.0, + ).to(tl.float32) offs_ke = ( (cur_q_start + offs_l)[:, None] * stride_kebs + cur_kv_head * stride_keh + offs_d[None, :] ) - ke = tl.load(K_Extend + offs_ke, mask=mask_l[:, None], other=0.0).to(tl.float32) + ke = tl.load( + K_Extend + offs_ke, + mask=mask_l[:, None] & (offs_d[None, :] < HEAD_DIM), + other=0.0, + ).to(tl.float32) offs_ve = ( (cur_q_start + offs_l)[:, None] * stride_vebs + cur_kv_head * stride_veh + offs_dv[None, :] ) - ve = tl.load(V_Extend + offs_ve, mask=mask_l[:, None], other=0.0).to(tl.float32) + ve = tl.load( + V_Extend + offs_ve, + mask=mask_l[:, None] & (offs_dv[None, :] < V_HEAD_DIM), + other=0.0, + ).to(tl.float32) # scores[i,j] = q_i . k_j (i query, j key) -> [L_EXT, L_EXT] qk = tl.sum(q[:, None, :] * ke[None, :, :], 2) * sm_scale @@ -374,7 +414,11 @@ def _verify_combine_stage2( + cur_head * stride_ooh + offs_dv[None, :] ) - tl.store(O_Out + offs_oo, o.to(O_Out.dtype.element_ty), mask=mask_l[:, None]) + tl.store( + O_Out + offs_oo, + o.to(O_Out.dtype.element_ty), + mask=mask_l[:, None] & (offs_dv[None, :] < V_HEAD_DIM), + ) class VerifySplitKV: @@ -471,6 +515,8 @@ class VerifySplitKV: kv_group_num=self.group, N_SPLITS=self.n_splits, L_EXT=self.l_pad, + HEAD_DIM=self.head_dim, + V_HEAD_DIM=self.v_head_dim, BLOCK_DMODEL=triton.next_power_of_2(self.head_dim), BLOCK_DV=triton.next_power_of_2(self.v_head_dim), BLOCK_N=self.block_n, @@ -511,6 +557,8 @@ class VerifySplitKV: kv_group_num=self.group, N_SPLITS=self.n_splits, L_EXT=self.l_pad, + HEAD_DIM=self.head_dim, + V_HEAD_DIM=self.v_head_dim, BLOCK_DMODEL=triton.next_power_of_2(self.head_dim), BLOCK_DV=triton.next_power_of_2(self.v_head_dim), num_warps=1, diff --git a/python/sglang/kernels/ops/attention/vision_rope.py b/python/sglang/kernels/ops/attention/vision_rope.py new file mode 100644 index 000000000..b284075cb --- /dev/null +++ b/python/sglang/kernels/ops/attention/vision_rope.py @@ -0,0 +1,217 @@ +"""Fused interleaved complex RoPE for vision attention Q/K tensors.""" + +from __future__ import annotations + +from typing import Tuple + +import torch +import triton +import triton.language as tl + +PreparedInplaceComplexRoPE = Tuple[torch.Tensor, torch.Tensor] + + +@triton.jit(do_not_specialize=["n_pairs"]) +def _fused_qk_complex_rope_kernel( + q_ptr, + k_ptr, + freqs_ptr, + q_out_ptr, + k_out_ptr, + n_pairs, + n_heads: tl.constexpr, + head_dim: tl.constexpr, + q_stride_token: tl.constexpr, + q_stride_head: tl.constexpr, + q_stride_dim: tl.constexpr, + k_stride_token: tl.constexpr, + k_stride_head: tl.constexpr, + k_stride_dim: tl.constexpr, + freq_stride_token: tl.constexpr, + freq_stride_pair: tl.constexpr, + freq_stride_complex: tl.constexpr, + BLOCK: tl.constexpr, +) -> None: + pair_offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = pair_offsets < n_pairs + pairs_per_row = head_dim // 2 + row = pair_offsets // pairs_per_row + pair = pair_offsets - row * pairs_per_row + token = row // n_heads + head = row - token * n_heads + + q_base = token * q_stride_token + head * q_stride_head + pair * 2 * q_stride_dim + k_base = token * k_stride_token + head * k_stride_head + pair * 2 * k_stride_dim + freq_base = token * freq_stride_token + pair * freq_stride_pair + cos = tl.load(freqs_ptr + freq_base, mask=mask).to(tl.float32) + sin = tl.load(freqs_ptr + freq_base + freq_stride_complex, mask=mask).to(tl.float32) + q_real = tl.load(q_ptr + q_base, mask=mask).to(tl.float32) + q_imag = tl.load(q_ptr + q_base + q_stride_dim, mask=mask).to(tl.float32) + k_real = tl.load(k_ptr + k_base, mask=mask).to(tl.float32) + k_imag = tl.load(k_ptr + k_base + k_stride_dim, mask=mask).to(tl.float32) + + out_base = row * head_dim + pair * 2 + tl.store( + q_out_ptr + out_base, + tl.fma(-q_imag, sin, q_real * cos), + mask=mask, + ) + tl.store( + q_out_ptr + out_base + 1, + tl.fma(q_real, sin, q_imag * cos), + mask=mask, + ) + tl.store( + k_out_ptr + out_base, + tl.fma(-k_imag, sin, k_real * cos), + mask=mask, + ) + tl.store( + k_out_ptr + out_base + 1, + tl.fma(k_real, sin, k_imag * cos), + mask=mask, + ) + + +def can_use_fused_qk_complex_rope( + q: torch.Tensor, + k: torch.Tensor, + freqs_cis: torch.Tensor, +) -> bool: + """Whether the NVIDIA fused path supports these vision RoPE tensors.""" + + if not ( + q.is_cuda + and k.is_cuda + and freqs_cis.is_cuda + and q.device == k.device == freqs_cis.device + ): + return False + if q.dtype != k.dtype or q.dtype not in (torch.bfloat16, torch.float16): + return False + if freqs_cis.dtype != torch.complex64 or q.shape != k.shape or q.ndim < 3: + return False + if q.shape[-1] % 2 != 0: + return False + if freqs_cis.shape != q.shape[:-2] + (q.shape[-1] // 2,): + return False + major, _ = torch.cuda.get_device_capability(q.device) + return major >= 9 + + +def apply_fused_qk_complex_rope( + q: torch.Tensor, + k: torch.Tensor, + freqs_cis: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Rotate interleaved Q/K pairs with one kernel. + + ``q`` and ``k`` may be strided views of an interleaved QKV projection. The + output is contiguous, matching the native complex-multiply implementation. + The token count remains a runtime kernel argument so random image sizes do + not create new Triton specializations. + """ + + if not can_use_fused_qk_complex_rope(q, k, freqs_cis): + raise ValueError( + "Unsupported fused vision RoPE inputs: " + f"q={q.shape}/{q.dtype}/{q.device}, " + f"k={k.shape}/{k.dtype}/{k.device}, " + f"freqs={freqs_cis.shape}/{freqs_cis.dtype}/{freqs_cis.device}" + ) + + original_shape = q.shape + # Preserve the interleaved QKV token stride when the token dimension is 1. + # ``view(-1, ...)`` is otherwise free to collapse that singleton stride, + # producing a different Triton specialization from real image requests. + q_flat = q if q.ndim == 3 else q.view(-1, q.shape[-2], q.shape[-1]) + k_flat = k if k.ndim == 3 else k.view(-1, k.shape[-2], k.shape[-1]) + freqs = torch.view_as_real(freqs_cis).view(-1, q.shape[-1] // 2, 2) + q_out = torch.empty(q_flat.shape, dtype=q.dtype, device=q.device) + k_out = torch.empty(k_flat.shape, dtype=k.dtype, device=k.device) + + block = 128 + n_pairs = q_flat.numel() // 2 + _fused_qk_complex_rope_kernel[(triton.cdiv(n_pairs, block),)]( + q_flat, + k_flat, + freqs, + q_out, + k_out, + n_pairs, + q_flat.shape[1], + q_flat.shape[2], + q_flat.stride(0), + q_flat.stride(1), + q_flat.stride(2), + k_flat.stride(0), + k_flat.stride(1), + k_flat.stride(2), + freqs.stride(0), + freqs.stride(1), + freqs.stride(2), + BLOCK=block, + num_warps=4, + ) + return q_out.view(original_shape), k_out.view(original_shape) + + +def prepare_fused_qk_complex_rope_inplace( + freqs_cis: torch.Tensor, +) -> PreparedInplaceComplexRoPE: + """Prepare the cache and positions used by the contiguous in-place kernel.""" + + if freqs_cis.dtype != torch.complex64: + raise ValueError( + "In-place vision RoPE requires complex64 frequencies, got " + f"{freqs_cis.dtype}/{freqs_cis.device}" + ) + return ( + torch.cat((freqs_cis.real, freqs_cis.imag), dim=-1), + torch.arange( + freqs_cis.size(0), + dtype=torch.long, + device=freqs_cis.device, + ), + ) + + +def apply_fused_qk_complex_rope_inplace( + q: torch.Tensor, + k: torch.Tensor, + prepared_rope: PreparedInplaceComplexRoPE, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Rotate contiguous Q/K tensors in place with the shared JIT kernel.""" + + from sglang.kernels.ops.attention.rope import apply_rope_inplace + + cos_sin_cache, positions = prepared_rope + apply_rope_inplace( + q, + k, + cos_sin_cache, + positions, + is_neox=False, + rope_dim=cos_sin_cache.size(-1), + ) + return q, k + + +def precompile_fused_qk_complex_rope( + *, + num_heads: int, + head_dim: int, + dtype: torch.dtype, + device: torch.device, +) -> bool: + """Compile the dynamic-token QKV-view specialization before serving.""" + + if device.type != "cuda" or dtype not in (torch.bfloat16, torch.float16): + return False + qkv = torch.empty((1, 3, num_heads, head_dim), dtype=dtype, device=device) + q, k, _ = torch.unbind(qkv, dim=1) + freqs = torch.ones((1, head_dim // 2), dtype=torch.complex64, device=device) + if not can_use_fused_qk_complex_rope(q, k, freqs): + return False + apply_fused_qk_complex_rope(q, k, freqs) + return True diff --git a/python/sglang/kernels/ops/elementwise/add3.py b/python/sglang/kernels/ops/elementwise/add3.py new file mode 100644 index 000000000..791ef4023 --- /dev/null +++ b/python/sglang/kernels/ops/elementwise/add3.py @@ -0,0 +1,67 @@ +"""CUDA JIT elementwise 3-way add: out = bf16(bf16(a + b) + c).""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from sglang.kernels.jit.utils import ( + cache_once, + is_arch_support_pdl, + load_jit, + make_cpp_args, +) + +if TYPE_CHECKING: + from tvm_ffi.module import Module + +# The kernel vectorizes by device::kMaxVecBytes (16B pre-Blackwell, 32B on +# Blackwell+). Requiring divisibility by the widest case keeps covered() +# arch-independent and never looser than the compiled kernel's check. +_MAX_VEC_ELEMS: int = 16 + + +@cache_once +def _jit_add3_module() -> Module: + args = make_cpp_args(is_arch_support_pdl()) + return load_jit( + "add3_bf16", + *args, + cuda_files=["elementwise/add3.cuh"], + cuda_wrappers=[("run", f"sglang::Add3Kernel<{args}>::launch")], + extra_cuda_cflags=["-O3", "--use_fast_math"], + ) + + +def covered(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor) -> bool: + """Same-shape contiguous CUDA bf16 tensors, numel a multiple of the + widest vector (16 elements).""" + return ( + a.dtype == b.dtype == c.dtype == torch.bfloat16 + and a.shape == b.shape == c.shape + and a.is_contiguous() + and b.is_contiguous() + and c.is_contiguous() + and a.is_cuda + and a.numel() > 0 + and a.numel() % _MAX_VEC_ELEMS == 0 + ) + + +def add3( + a: torch.Tensor, + b: torch.Tensor, + c: torch.Tensor, + *, + out: torch.Tensor | None = None, + prefetch_bc: bool = False, +) -> torch.Tensor: + """out = bf16(bf16(a + b) + c); double rounding matches the unfused add + pair bit-for-bit. With prefetch_bc, b/c are loaded before the PDL wait — + only safe when their producers are at least two kernels back.""" + if out is None: + out = torch.empty_like(a) + module = _jit_add3_module() + module.run(a.view(-1), b.view(-1), c.view(-1), out.view(-1), prefetch_bc) + return out diff --git a/python/sglang/kernels/ops/gemm/cutedsl_bf16_gemm.py b/python/sglang/kernels/ops/gemm/cutedsl_bf16_gemm.py index 4b7fc70e8..f608be4fb 100644 --- a/python/sglang/kernels/ops/gemm/cutedsl_bf16_gemm.py +++ b/python/sglang/kernels/ops/gemm/cutedsl_bf16_gemm.py @@ -1294,12 +1294,39 @@ def _pick_tactic(m: int, n: int, k: int) -> int: return best +# Kimi-K3 per-rank dense-GEMM shapes (TP8). They sit in this heuristic's +# unmeasured region (hidden=7168 inputs fail k > 6144; o_proj-style k=1536 +# fails k < 2048), but TGV wins 1.04-2.43x on every one of them on GB300 +# (L2-defeating weight rotation + CUDA-graph timing; serving A/B confirmed +# e2e with GSM8K parity). kv_a (n=576) loses (0.84x) and stays out. Gated to +# small decode batches; larger m stays on the measured heuristic below. +_K3_TGV_WIN_SHAPES = frozenset( + { + (6144, 7168), # KDA fused qkvg + (6016, 7168), # merged MoE front (gate_up | router | latent down) + (7168, 1536), # KDA / MLA o_proj + (1536, 7168), # MLA q_a / shared gate_up + (2304, 1536), # MLA q_b + (3584, 7168), # MoE latent down (unfused fallback) + (7168, 3584), # MoE latent up + (7168, 768), # shared down + } +) + + def use_cutedsl_bf16_gemm(m: int, n: int, k: int) -> bool: """TGV-vs-cuBLAS (``F.linear``) decision, CUPTI-measured on B300 under CUDA graph capture (cold L2). Conservative: ties and unmeasured regions fall back to cuBLAS.""" + if m <= 0: + # Empty batch: DP-attention idle groups run a 0-token dummy forward to + # keep the mlp-sync lockstep. A 0-CTA TGV grid is a driver-level + # CUDA_ERROR_INVALID_VALUE, so leave empty inputs to cuBLAS. + return False if k % 8 != 0: # TMA requires 16B-aligned rows return False + if m <= 8 and (n, k) in _K3_TGV_WIN_SHAPES: + return True if n < 1024 or k < 2048 or k > 6144: return False ragged = m % 16 != 0 @@ -1333,6 +1360,10 @@ def _tgv_bf16_gemm_run( out = torch.empty( (x.shape[0], weight.shape[0]), dtype=torch.bfloat16, device=x.device ) + if x.shape[0] == 0: + # Match cuBLAS/F.linear semantics for empty batches; a 0-CTA launch + # would fail with CUDA_ERROR_INVALID_VALUE. + return out return _run_tgv( x, weight.t(), @@ -1343,6 +1374,33 @@ def _tgv_bf16_gemm_run( ) +def _tgv_bf16_gemm_out_run( + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + bias: Optional[torch.Tensor], +) -> None: + if get_device_sm() not in (100, 103): + raise RuntimeError("cutedsl_bf16_gemm requires SM100/SM103 (Blackwell)") + assert x.dtype == torch.bfloat16 and weight.dtype == torch.bfloat16 + assert out.dtype == torch.bfloat16 and out.device == x.device + assert x.ndim == 2 and weight.ndim == 2 and out.ndim == 2 + assert x.stride(-1) == 1, "x must be K-major [M, K]" + assert weight.stride(-1) == 1, "weight must be K-major [N, K]" + assert out.is_contiguous() and out.shape == (x.shape[0], weight.shape[0]) + if x.shape[0] == 0: + return None + _run_tgv( + x, + weight.t(), + bias, + out, + pdl=True, + tactic=_pick_tactic(x.shape[0], weight.shape[0], weight.shape[1]), + ) + return None + + def _tgv_bf16_gemm_fake( x: torch.Tensor, weight: torch.Tensor, @@ -1351,6 +1409,15 @@ def _tgv_bf16_gemm_fake( return x.new_empty((x.shape[0], weight.shape[0])) +def _tgv_bf16_gemm_out_fake( + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + bias: Optional[torch.Tensor], +) -> None: + return None + + direct_register_custom_op( op_name="cutedsl_tgv_bf16_gemm", op_func=_tgv_bf16_gemm_run, @@ -1358,6 +1425,13 @@ direct_register_custom_op( fake_impl=_tgv_bf16_gemm_fake, ) +direct_register_custom_op( + op_name="cutedsl_tgv_bf16_gemm_out", + op_func=_tgv_bf16_gemm_out_run, + mutates_args=["out"], + fake_impl=_tgv_bf16_gemm_out_fake, +) + @debug_kernel_api def cutedsl_bf16_gemm( @@ -1367,3 +1441,15 @@ def cutedsl_bf16_gemm( ) -> torch.Tensor: """out[M, N] = x[M, K] @ weight[N, K].T (+ bias[N]), all bf16, fp32 accum.""" return torch.ops.sglang.cutedsl_tgv_bf16_gemm(x, weight, bias) + + +@debug_kernel_api +def cutedsl_bf16_gemm_out( + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + bias: torch.Tensor | None = None, +) -> torch.Tensor: + """Write the BF16 GEMM directly into a caller-owned contiguous tensor.""" + torch.ops.sglang.cutedsl_tgv_bf16_gemm_out(x, weight, out, bias) + return out diff --git a/python/sglang/kernels/ops/gemm/tiny_gemm.py b/python/sglang/kernels/ops/gemm/tiny_gemm.py new file mode 100644 index 000000000..379ff3c5f --- /dev/null +++ b/python/sglang/kernels/ops/gemm/tiny_gemm.py @@ -0,0 +1,144 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.kernels.jit.utils import ( + cache_once, + is_arch_support_pdl, + load_jit, + make_cpp_args, +) + +if TYPE_CHECKING: + from tvm_ffi.module import Module + +_MAX_M_DEFAULT: int = 16 + + +@cache_once +def _jit_tiny_gemm_module( + n: int, k: int, max_m: int, split_n: int, out_dtype: torch.dtype +) -> Module: + args = make_cpp_args(n, k, max_m, split_n, out_dtype, is_arch_support_pdl()) + return load_jit( + "tiny_gemm", + *args, + cuda_files=["gemm/tiny_gemm.cuh"], + cuda_wrappers=[("run", f"TinyNGemmKernel<{args}>::run")], + extra_cuda_cflags=["-O3"], + ) + + +@cache_once +def _jit_tiny_k_gemm_module( + n: int, k: int, max_m: int, n_unroll: int, out_dtype: torch.dtype +) -> Module: + args = make_cpp_args(n, k, max_m, n_unroll, out_dtype, is_arch_support_pdl()) + return load_jit( + "tiny_k_gemm", + *args, + cuda_files=["gemm/tiny_gemm.cuh"], + cuda_wrappers=[("run", f"TinyKGemmKernel<{args}>::run")], + extra_cuda_cflags=["-O3"], + ) + + +def _vec_elems() -> int: + """bf16 elements per vectorized load; mirrors kMaxVecBytes in utils.cuh.""" + from sglang.kernels.jit.utils import get_jit_cuda_arch + + cuda = tuple(int(v) for v in (torch.version.cuda or "0.0").split(".")[:2]) + return 16 if get_jit_cuda_arch().major >= 10 and cuda >= (12, 9) else 8 + + +def _default_split_n(n: int, k: int, max_m: int, device: torch.device) -> int: + """Smallest divisor of n whose n / split_n blocks fit in one wave, subject + to the max_m * split_n <= K / vec_elems block-size constraint; falls back + to the largest split_n satisfying the constraint (multi-wave grid).""" + sm_count = torch.cuda.get_device_properties(device).multi_processor_count + split_cap = (k // _vec_elems()) // max_m + divisors = [d for d in range(1, min(n, split_cap) + 1) if n % d == 0] + if not divisors: + raise RuntimeError( + f"tiny_gemm: no valid split_n for N={n}, K={k}, max_m={max_m};" + " lower max_m" + ) + for split in divisors: + if n // split <= sm_count: + return split + return divisors[-1] + + +def tiny_n_gemm_bf16( + x: torch.Tensor, + w: torch.Tensor, + out: Optional[torch.Tensor] = None, + *, + out_dtype: Optional[torch.dtype] = None, + split_n: Optional[int] = None, + max_m: int = _MAX_M_DEFAULT, +) -> torch.Tensor: + n = w.shape[0] + k = x.shape[1] + if out is None: + out_dtype = out_dtype or torch.bfloat16 + out = torch.empty((x.shape[0], n), dtype=out_dtype, device=x.device) + else: + assert out_dtype is None or out_dtype == out.dtype + if split_n is None: + split_n = _default_split_n(n, k, max_m, x.device) + module = _jit_tiny_gemm_module(n, k, max_m, split_n, out.dtype) + module.run(x, w, out) + return out + + +def _default_k_split_n(n: int, k: int) -> int: + """Smallest divisor of n whose n / split_n blocks fit one wave, with + split_n * K-lanes whole-warp aligned and within the block-size limit.""" + lanes = k // 8 # fixed 16-byte vectors in the K variant + candidates = [ + d + for d in range(1, n + 1) + if n % d == 0 and d * lanes % 32 == 0 and d * lanes <= 1024 + ] + if not candidates: + raise RuntimeError(f"tiny_k_gemm: no valid split_n for N={n}, K={k}") + sm_count = torch.cuda.get_device_properties(0).multi_processor_count + for d in candidates: + if n // d <= sm_count: + return d + return candidates[-1] + + +def tiny_k_gemm_bf16( + x: torch.Tensor, + w: torch.Tensor, + out: Optional[torch.Tensor] = None, + *, + out_dtype: Optional[torch.dtype] = None, + split_n: Optional[int] = None, + max_m: int = _MAX_M_DEFAULT, +) -> torch.Tensor: + """Small-K / large-N variant: K / 8 lanes of one warp reduce the K + dimension for one output column; each block covers split_n columns and the + exact N / split_n grid fills the SMs (no tail). Requires K / 8 to be a + power of 2 and <= 32 (e.g. K = 128/256). x may be a row-sliced view as + long as rows stay 16-byte aligned. + + split_n trades block count for block size; the default picks the smallest + divisor of N that fits one wave (12 for [1536, 128] on B200: 128 blocks + of 6 warps).""" + n = w.shape[0] + k = x.shape[1] + if out is None: + out_dtype = out_dtype or torch.bfloat16 + out = torch.empty((x.shape[0], n), dtype=out_dtype, device=x.device) + else: + assert out_dtype is None or out_dtype == out.dtype + if split_n is None: + split_n = _default_k_split_n(n, k) + module = _jit_tiny_k_gemm_module(n, k, max_m, split_n, out.dtype) + module.run(x, w, out) + return out diff --git a/python/sglang/kernels/ops/kimi_k3/__init__.py b/python/sglang/kernels/ops/kimi_k3/__init__.py new file mode 100644 index 000000000..7e8ce029a --- /dev/null +++ b/python/sglang/kernels/ops/kimi_k3/__init__.py @@ -0,0 +1,81 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +if TYPE_CHECKING: + import torch + +_K3_N_GEMM_DISPATCH_MAP = { + (144, 7168): 16, + (896, 7168): 8, +} +_K3_K_GEMM_DISPATCH_MAP = { + (1536, 128): 12, +} + + +def situ_and_mul( + input: torch.Tensor, + out: Optional[torch.Tensor], + beta: float, + linear_beta: Optional[float], +) -> torch.Tensor: + from .activation import situ_and_mul as impl + + return impl(input, out, beta, linear_beta) + + +def situ_and_mul_masked_post_quant( + input: torch.Tensor, + output: torch.Tensor, + output_scale: torch.Tensor, + quant_group_size: int, + masked_m: torch.Tensor, + beta: float, + linear_beta: float, + scale_ue8m0: bool = False, + topk: int = 8, + transposed: bool = False, + swizzle: bool = False, +) -> None: + from .moe import situ_and_mul_masked_post_quant as impl + + return impl( + input, + output, + output_scale, + quant_group_size, + masked_m, + beta, + linear_beta, + scale_ue8m0, + topk, + transposed, + swizzle, + ) + + +def kimi_k3_tiny_gemm( + x: torch.Tensor, + w: torch.Tensor, +) -> torch.Tensor: + import torch + + from ..gemm.tiny_gemm import tiny_k_gemm_bf16, tiny_n_gemm_bf16 + + m, k = x.shape + n, _ = w.shape + if max_num_tokens := _K3_N_GEMM_DISPATCH_MAP.get((n, k)): + if 0 < m <= max_num_tokens: + return tiny_n_gemm_bf16(x, w) + if max_num_tokens := _K3_K_GEMM_DISPATCH_MAP.get((n, k)): + if 0 < m <= max_num_tokens: + return tiny_k_gemm_bf16(x, w) + return torch.nn.functional.linear(x, w) + + +__all__ = [ + "situ_and_mul", + "situ_and_mul_masked_post_quant", + "kimi_k3_tiny_gemm", +] diff --git a/python/sglang/kernels/ops/kimi_k3/activation.py b/python/sglang/kernels/ops/kimi_k3/activation.py new file mode 100644 index 000000000..20527e00e --- /dev/null +++ b/python/sglang/kernels/ops/kimi_k3/activation.py @@ -0,0 +1,87 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.kernels.jit.utils import ( + cache_once, + get_jit_cuda_arch, + is_arch_support_pdl, + is_hip_runtime, + load_jit, + make_cpp_args, +) + +if TYPE_CHECKING: + from tvm_ffi.module import Module + + +def _make_name(*args): + return "kimi_k3_" + "_".join(str(a) for a in args) + + +def _fast_math_flags() -> list[str]: + # Mirrors sgl-kernel's CMake policy: fast-math on SM90, precise on + # SM100+ (Blackwell needs bit-exact expf), off on HIP (clang rejects). + if is_hip_runtime(): + return [] + if get_jit_cuda_arch().major >= 10: + return [] + return ["--use_fast_math"] + + +@cache_once +def _jit_situ_and_mul_module(dtype: torch.dtype) -> Module: + """Compile and cache the JIT SiTU-and-mul module for a given dtype.""" + args = make_cpp_args(dtype, is_arch_support_pdl()) + return load_jit( + _make_name("situ_and_mul"), + *args, + cuda_files=["kimi_k3/situ_and_mul.cuh"], + cuda_wrappers=[("run", f"SituAndMulKernel<{args}>::run")], + extra_cuda_cflags=_fast_math_flags(), + ) + + +def situ_and_mul( + input: torch.Tensor, + out: Optional[torch.Tensor], + beta: float, + linear_beta: Optional[float], +) -> torch.Tensor: + """Fused SiTU (SoftCap-GLU) activation: bf16 -> bf16. + + gate_out = beta * tanh(gate / beta) * sigmoid(gate) + up_out = linear_beta * tanh(up / linear_beta) [if linear_beta is not None] + output = gate_out * up_out + + Parameters + ---------- + input : bf16 CUDA tensor [*, 2*D] + out : optional pre-allocated bf16 CUDA tensor [*, D] + beta : gate softcap scalar (e.g. 4.0) + linear_beta : up softcap scalar (e.g. 25.0), or None to skip + """ + hidden_size = input.shape[-1] // 2 + if out is None: + out = input.new_empty(*input.shape[:-1], hidden_size) + + # 2D inputs may be row-strided (e.g. a slice of a fused-GEMM output); + # higher-rank inputs keep the dense-view path. + if input.dim() == 2 and input.stride(1) == 1: + input_2d = input + else: + input_2d = input.contiguous().view(-1, hidden_size * 2) + out_2d = out.view(-1, hidden_size) + + has_linear_beta = linear_beta is not None + module = _jit_situ_and_mul_module(input.dtype) + module.run( + input_2d, + out_2d, + float(beta), + float(linear_beta) if has_linear_beta else 0.0, + has_linear_beta, + ) + return out diff --git a/python/sglang/kernels/ops/kimi_k3/all_reduce.py b/python/sglang/kernels/ops/kimi_k3/all_reduce.py new file mode 100644 index 000000000..7e7b7a52d --- /dev/null +++ b/python/sglang/kernels/ops/kimi_k3/all_reduce.py @@ -0,0 +1,380 @@ +"""K3 MNNVL fused all-reduce (bf16): zero-copy AR and AR+RMSNorm. + +Four entry points over ``csrc/kimi_k3/comm/ar_fusion.cuh``, spanning two +algorithm families x two epilogues: + +============ ========================= ================================== + res (+ optional residual) norm (fused RMSNorm on the latent) +============ ========================= ================================== +push (1shot) :func:`all_reduce_push_res` :func:`all_reduce_push_norm` +pull (2shot) :func:`all_reduce_pull_res` :func:`all_reduce_pull_norm` +============ ========================= ================================== + +* **push** — 1shot multicast-push. Works on ANY contiguous bf16 tensor + (input is read and written in place); reuses the CustomAllReduceV2 push + workspace, so the caller passes the workspace slab's multicast base. + Best for small messages. Needs :func:`register_comm`. +* **pull** — low-SM NVLS 2shot ON the input, which must be allocated from + multicast-bound symmetric memory (the caller passes its multicast VA): + reduce-scatter + broadcast in place. + Launch geometry defaults to :data:`RES_TUNING` / + :data:`NORM_TUNING` and can be overridden per call via ``num_blocks`` / + ``unroll`` (``num_blocks`` must be uniform across ranks per call). + Barriers reuse the CustomAllReduceV2 pull semaphores (same reservation + protocol as the generic pull kernels, signaled via one multicast red), so + :func:`register_comm` additionally needs the semaphore region's multicast + VA (``CustomAllReduceV2.pull_sem_mc_ptr``). + +Epilogue contracts: the ``res`` residual must be identical on every rank (a +fully reduced tensor such as the attn-res prefix sum) or absent; the +``norm`` input is the K3 latent|shared MoE buffer ([N, 3584] latent then +[N, 7168] shared, contiguous — the row layout is derived and hardcoded +C++-side). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, NamedTuple, Optional + +import torch + +from sglang.kernels.jit.utils import ( + cache_once, + is_arch_support_pdl, + load_jit, + make_cpp_args, +) +from sglang.srt.utils.custom_op import register_custom_op + +if TYPE_CHECKING: + from tvm_ffi.module import Module + + from sglang.kernels.ops.communication.all_reduce import Communicator + + +@cache_once +def _jit_module(world_size: int) -> Module: + args = make_cpp_args(world_size, is_arch_support_pdl()) + cls = f"AllReduceFusionKernel<{args}>" + return load_jit( + "kimi_k3_all_reduce", + *args, + cuda_files=["kimi_k3/comm/ar_fusion.cuh"], + cuda_wrappers=[ + ("push_res", f"{cls}::push_res"), + ("push_norm", f"{cls}::push_norm"), + ("pull_res", f"{cls}::pull_res"), + ("pull_norm", f"{cls}::pull_norm"), + ("finalize_push_norm", f"{cls}::finalize_push_norm"), + ], + extra_cuda_cflags=["-O3"], + ) + + +# Storage plane: the CustomAllReduceV2 Communicator + + +class _CommEntry(NamedTuple): + obj: Communicator # sgl.Communicator + pull_sem_mc_ptr: int + + +_COMM_MAP: dict[int, _CommEntry] = {} + + +def register_comm(comm: Communicator, *, pull_sem_mc_ptr: int = 0) -> None: + """Register the CustomAllReduceV2 storage plane. + + The push kernels only need ``comm``; the pull kernels additionally need + ``pull_sem_mc_ptr`` (``CustomAllReduceV2.pull_sem_mc_ptr``), the + multicast VA of the pull-semaphore region their barriers reuse. + """ + # world_size is the whole key, so at most one communicator per size can be + # registered in a process. That matches how these ops are called -- the custom + # ops below take world_size and nothing else, so a second group of the same + # size would silently inherit the first one's peer pointers and semaphores, + # and the symptom would be a hang or corruption rather than an error. Assert + # it instead of letting the overwrite happen; widening to per-group handles + # means changing the custom-op signatures, which is a separate change. + prev = _COMM_MAP.get(comm.world_size) + assert prev is None or prev.obj is comm, ( + f"a different communicator is already registered for world_size=" + f"{comm.world_size}; these ops key only on world_size, so two groups of " + f"the same size cannot coexist in one process" + ) + _COMM_MAP[comm.world_size] = _CommEntry(obj=comm, pull_sem_mc_ptr=pull_sem_mc_ptr) + + +class PullTuning(NamedTuple): + num_blocks: int + unroll: int # 2, 4, 8, or 16 (every width is compiled into the module) + + +class PullTuningTable(NamedTuple): + bands: tuple[tuple[int, PullTuning], ...] + fallback: PullTuning + + def lookup(self, nbytes: int) -> PullTuning: + for max_bytes, tuning in self.bands: + if nbytes <= max_bytes: + return tuning + return self.fallback + + +_KB, _MB = 1024, 1024 * 1024 + +RES_TUNING = PullTuningTable( + bands=( + (128 * _KB, PullTuning(num_blocks=1, unroll=8)), + (4 * _MB, PullTuning(num_blocks=16, unroll=8)), + ), + fallback=PullTuning(num_blocks=4, unroll=8), +) + +NORM_TUNING = PullTuningTable( + bands=( + (512 * _KB, PullTuning(num_blocks=2, unroll=4)), + (2 * _MB, PullTuning(num_blocks=12, unroll=2)), + ), + fallback=PullTuning(num_blocks=24, unroll=2), +) + + +def _resolve_tuning( + table: PullTuningTable, + *, + nbytes: int, + num_blocks: Optional[int], + unroll: Optional[int], +) -> PullTuning: + """Per-size tuned config, with explicit overrides taking precedence + (C++-side, the block count is clamped to the semaphore capacity).""" + tuned = table.lookup(nbytes) + return PullTuning( + num_blocks=num_blocks or tuned.num_blocks, + unroll=unroll or tuned.unroll, + ) + + +# Custom ops, one per C++ entry point + + +@register_custom_op(mutates_args=["x"]) +def _push_res_op( + world_size: int, + x: torch.Tensor, + residual: Optional[torch.Tensor], + ws_mc_base: int, +) -> None: + comm = _COMM_MAP[world_size].obj + _jit_module(world_size).push_res(comm, x.view(-1), residual, ws_mc_base) + + +@register_custom_op(mutates_args=["x"]) +def _push_norm_op( + world_size: int, + x: torch.Tensor, + weight: torch.Tensor, + eps: float, + num_norm_rows: int, + ws_mc_base: int, +) -> None: + comm = _COMM_MAP[world_size].obj + _jit_module(world_size).push_norm( + comm, x.view(-1), weight, eps, num_norm_rows, ws_mc_base + ) + + +@register_custom_op(mutates_args=["out"]) +def _finalize_push_norm_op( + world_size: int, + out: torch.Tensor, + gemm2_out: torch.Tensor, + expanded_idx_to_permuted_idx: torch.Tensor, + expert_weights: torch.Tensor, + weight: torch.Tensor, + eps: float, + ws_mc_base: int, +) -> None: + comm = _COMM_MAP[world_size].obj + _jit_module(world_size).finalize_push_norm( + comm, + out.view(-1), + gemm2_out, + expanded_idx_to_permuted_idx, + expert_weights, + weight, + eps, + ws_mc_base, + ) + + +@register_custom_op(mutates_args=["x"]) +def _pull_res_op( + world_size: int, + x: torch.Tensor, + residual: Optional[torch.Tensor], + input_mc_ptr: int, + num_blocks: int, + unroll: int, +) -> None: + entry = _COMM_MAP[world_size] + _jit_module(world_size).pull_res( + entry.obj, + x.view(-1), + residual, + input_mc_ptr, + entry.pull_sem_mc_ptr, + num_blocks, + unroll, + ) + + +@register_custom_op(mutates_args=["x"]) +def _pull_norm_op( + world_size: int, + x: torch.Tensor, + weight: torch.Tensor, + eps: float, + num_norm_rows: int, + input_mc_ptr: int, + num_blocks: int, + unroll: int, +) -> None: + entry = _COMM_MAP[world_size] + _jit_module(world_size).pull_norm( + entry.obj, + x.view(-1), + weight, + eps, + num_norm_rows, + input_mc_ptr, + entry.pull_sem_mc_ptr, + num_blocks, + unroll, + ) + + +def all_reduce_push_res( + world_size: int, + x: torch.Tensor, + residual: Optional[torch.Tensor] = None, + *, + ws_mc_base: int, +) -> torch.Tensor: + """In-place ``x = allreduce(x) [+ residual]`` via 1shot multicast push. + + ``x`` may be any contiguous bf16 CUDA tensor whose byte size fits the + registered push workspace. ``ws_mc_base`` is the multicast VA of the v2 + workspace slab base. Call :func:`register_comm` once beforehand. + """ + residual_ = residual.view(-1) if residual is not None else None + _push_res_op(world_size, x, residual_, ws_mc_base) + return x + + +def all_reduce_push_norm( + world_size: int, + x: torch.Tensor, + weight: torch.Tensor, + eps: float, + *, + num_norm_rows: int, + ws_mc_base: int, +) -> torch.Tensor: + """In-place allreduce via 1shot multicast push + RMSNorm over the first + ``num_norm_rows`` rows of ``x`` viewed as [numel / 3584, 3584].""" + _push_norm_op(world_size, x, weight, eps, num_norm_rows, ws_mc_base) + return x + + +def finalize_all_reduce_push_norm( + world_size: int, + out: torch.Tensor, + gemm2_out: torch.Tensor, + expanded_idx_to_permuted_idx: torch.Tensor, + expert_weights: torch.Tensor, + weight: torch.Tensor, + eps: float, + *, + ws_mc_base: int, +) -> torch.Tensor: + """Deferred MoE finalize + 1shot push all-reduce + RMSNorm on EVERY row. + + ``out`` ([T, 3584] bf16) is output-only; each rank's partial latent + (``sum_k expert_weights[t, k] * gemm2_out[idx[t*16 + k]]``, -1 slots + skipped) is computed during the multicast staging pass from the + trtllm-gen deferred-finalize triple (``do_finalize=False``) and never + materializes in global memory. top_k is fixed to 16 (K3).""" + _finalize_push_norm_op( + world_size, + out, + gemm2_out, + expanded_idx_to_permuted_idx, + expert_weights, + weight, + eps, + ws_mc_base, + ) + return out + + +def all_reduce_pull_res( + world_size: int, + x: torch.Tensor, + residual: Optional[torch.Tensor] = None, + *, + input_mc_ptr: int, + num_blocks: Optional[int] = None, + unroll: Optional[int] = None, +) -> torch.Tensor: + """In-place ``x = allreduce(x) [+ residual]`` via low-SM NVLS 2shot. + + ``x`` MUST be allocated from multicast-bound symmetric memory and + ``input_mc_ptr`` must be its multicast VA. Call :func:`register_comm` + (with ``pull_sem_mc_ptr``) once beforehand. + """ + tuning = _resolve_tuning( + RES_TUNING, + nbytes=x.numel() * x.element_size(), + num_blocks=num_blocks, + unroll=unroll, + ) + residual_ = residual.view(-1) if residual is not None else None + _pull_res_op( + world_size, x, residual_, input_mc_ptr, tuning.num_blocks, tuning.unroll + ) + return x + + +def all_reduce_pull_norm( + world_size: int, + x: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, + *, + num_norm_rows: int, + input_mc_ptr: int, + num_blocks: Optional[int] = None, + unroll: Optional[int] = None, +) -> torch.Tensor: + """In-place allreduce via low-SM NVLS 2shot + RMSNorm over the first + ``num_norm_rows`` rows of ``x`` viewed as [numel / 3584, 3584]; ``x`` + must live in multicast-bound symmetric memory.""" + tuning = _resolve_tuning( + NORM_TUNING, + nbytes=x.nbytes, + num_blocks=num_blocks, + unroll=unroll, + ) + _pull_norm_op( + world_size, + x, + weight, + eps, + num_norm_rows, + input_mc_ptr, + tuning.num_blocks, + tuning.unroll, + ) + return x diff --git a/python/sglang/kernels/ops/kimi_k3/attn_res.py b/python/sglang/kernels/ops/kimi_k3/attn_res.py new file mode 100644 index 000000000..76a83bac5 --- /dev/null +++ b/python/sglang/kernels/ops/kimi_k3/attn_res.py @@ -0,0 +1,301 @@ +"""CUDA JIT wrapper for the Kimi-K3 SM100 attention-residual kernel.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from sglang.kernels.jit.utils import ( + cache_once, + load_jit, + make_cpp_args, + override_jit_cuda_arch, +) +from sglang.srt.utils.custom_op import register_custom_op + +if TYPE_CHECKING: + from tvm_ffi.module import Module + + from sglang.kernels.ops.communication.all_reduce import Communicator + +_DIM: int = 7168 # K3 hidden size, template parameter of the TMA kernel +_MAX_BANK_ROWS: int = 8 # K3 has <= 8 snapshots, upper bound of the nvb dispatch tables + + +def _make_name(*args): + return "kimi_k3_attn_res_" + "_".join(str(a) for a in args) + + +@cache_once +def _jit_fused_tma_module( + chunk_rows: int, occupancy: int, consumer_regs: int +) -> Module: + """Compile and cache the warp-specialized TMA aggregation kernel (per-row + bulk copies into chunk slots; chunk_rows / occupancy / consumer_regs are + tuning knobs). The smem ring is frozen at 2 chunk slots and PDL is always + on: the kernel targets SM100+, where both are unconditional wins.""" + major, minor = torch.cuda.get_device_capability() + if major < 10: + raise RuntimeError( + "attn_res_fused_tma requires SM100+ (tcgen05, cp.async.bulk)" + ) + args = make_cpp_args( + _DIM, + _MAX_BANK_ROWS, + chunk_rows, + occupancy, + consumer_regs, + ) + with override_jit_cuda_arch(major, minor, suffix="a"): + return load_jit( + _make_name("fused_tma"), + *args, + cuda_files=["kimi_k3/attn_res/fused_tma.cuh"], + cuda_wrappers=[ + ("run", f"AttnResFusedTmaKernel<{args}>::run"), + ("run_pull_rs", f"AttnResFusedTmaKernel<{args}>::run_pull_rs"), + ("run_direct_ag", f"AttnResFusedTmaKernel<{args}>::run_direct_ag"), + ], + extra_cuda_cflags=["-O3", "--use_fast_math"], + ) + + +# Benchmarked-best (chunk_rows, occupancy, consumer_regs) per nvb +# (GB200/GB300-class, H=7168). nvb=1 is latency-bound per token, so 2 CTAs/SM +# (occupancy=2, which excludes setmaxnreg: 4*Nc + 2*Np > 512) wins large T by +# ~17%; nvb=4/8 fill 5-row chunks exactly ((nvb+1) % 5 == 0 or covers it in +# 1-2 chunks); everything else is fastest on the balanced 4-row chunk. All +# entries use the setmaxnreg producer/consumer split. consumer_regs sits on +# the 200..232 performance plateau (below 200 the consumer loop starves); we +# take 200, not the 232 budget cap, so each SMSP keeps 2K registers free for +# PDL-overlapped neighbor kernels instead of allocating them idle. +_TMA_BEST_CONFIG: dict[int, tuple[int, int, int]] = { + 1: (2, 2, 0), + 2: (4, 1, 200), + 3: (4, 1, 200), + 4: (5, 1, 200), + 5: (3, 1, 200), + 6: (4, 1, 200), + 7: (4, 1, 200), + 8: (5, 1, 200), +} + + +def _tuning(nvb: int, num_tokens: int) -> tuple[int, int, int]: + """(chunk_rows, occupancy, consumer_regs) for this aggregation point. + + Shared by all three entry points below: they run the same kernel template + and differ only in the collective fused onto it.""" + if not 1 <= nvb <= _MAX_BANK_ROWS: + raise ValueError(f"attn_res: nvb must be in [1, {_MAX_BANK_ROWS}], got {nvb}") + best = _TMA_BEST_CONFIG[nvb] + if best[1] > 1 and num_tokens < 128: + # occupancy=2 only pays off once there are enough tokens to fill both + # CTAs per SM; below that its tighter register budget just costs ~10%. + best = (4, 1, 200) + return best + + +_COMM_MAP: dict[int, Communicator] = {} +_PULL_SEM_MC_MAP: dict[int, int] = {} + + +def register_comm(comm: Communicator, *, pull_sem_mc_ptr: int) -> None: + # One communicator per world_size per process -- see the note in + # kimi_k3/all_reduce.py::register_comm. The ops key only on world_size, so an + # overwrite here would hand the old group's callers the new group's peer + # pointers. + prev = _COMM_MAP.get(comm.world_size) + assert prev is None or prev is comm, ( + f"a different communicator is already registered for world_size=" + f"{comm.world_size}" + ) + _COMM_MAP[comm.world_size] = comm + _PULL_SEM_MC_MAP[comm.world_size] = pull_sem_mc_ptr + + +@register_custom_op(mutates_args=["out", "prefix_out"]) +def _attn_res_fused_pull_rs_op( + world_size: int, + input: torch.Tensor, + residual: torch.Tensor | None, + bank: torch.Tensor, + cw: torch.Tensor, + ow: torch.Tensor, + out: torch.Tensor, + prefix_out: torch.Tensor, + nvb: int, + eps: float, + input_mc_ptr: int, + max_blocks: int, + chunk_rows: int, + occupancy: int, + consumer_regs: int, +) -> None: + _jit_fused_tma_module(chunk_rows, occupancy, consumer_regs).run_pull_rs( + _COMM_MAP[world_size], + input, + residual, + bank, + cw, + ow, + out, + prefix_out, + nvb, + eps, + input_mc_ptr, + _PULL_SEM_MC_MAP[world_size], + max_blocks, + ) + + +def attn_res_fused_tma( + prefix_sum: torch.Tensor, + bank: torch.Tensor, + cw: torch.Tensor, + ow: torch.Tensor, + out: torch.Tensor, + nvb: int, + eps: float, + *, + write_prefix: bool = False, +) -> None: + """Warp-specialized TMA aggregation (score -> online softmax -> weighted + combine -> fused output RMSNorm), one persistent CTA per SM: a producer + warp (group) fetches rows with one bulk copy each into chunk slots of + `chunk_rows` rows (one barrier pair per chunk, double-buffered ring of 2 + slots); the 8 consumer warps score and fold one chunk per rendezvous, + with cw / ow staged in TMEM. + + Restrictions: H == 7168, nvb in [1, 8], SM100a+. The launch config comes + from _TMA_BEST_CONFIG via _tuning() (each combination compiles its own + module). + + Parameters + ---------- + prefix_sum : [T, H] bf16 + bank : [T, NB, H] bf16 (rows 0..nvb-1 are aggregated) + cw : [H] bf16 — precomputed score_norm_weight * proj_weight + ow : [H] bf16 — output RMSNorm weight + out : [T, H] bf16 output buffer + nvb : number of valid bank rows (1..8) + eps : RMSNorm epsilon (shared by score and output norms) + write_prefix : also snapshot the prefix row into bank[:, nvb, :] + (bit-exact copy, fused into the score pass which already + has the row in registers); requires NB > nvb + """ + _jit_fused_tma_module(*_tuning(nvb, prefix_sum.shape[0])).run( + prefix_sum, bank, cw, ow, out, nvb, eps, write_prefix + ) + + +@register_custom_op(mutates_args=["bank", "out"]) +def _attn_res_fused_direct_ag_op( + world_size: int, + prefix_sum: torch.Tensor, + bank: torch.Tensor, + cw: torch.Tensor, + ow: torch.Tensor, + out: torch.Tensor, + nvb: int, + eps: float, + output_mc_ptr: int, + max_blocks: int, + write_prefix: bool, + chunk_rows: int, + occupancy: int, + consumer_regs: int, +) -> None: + _jit_fused_tma_module(chunk_rows, occupancy, consumer_regs).run_direct_ag( + _COMM_MAP[world_size], + prefix_sum, + bank, + cw, + ow, + out, + nvb, + eps, + output_mc_ptr, + _PULL_SEM_MC_MAP[world_size], + max_blocks, + write_prefix, + ) + + +def attn_res_fused_direct_ag( + world_size: int, + prefix_sum: torch.Tensor, + bank: torch.Tensor, + cw: torch.Tensor, + ow: torch.Tensor, + out: torch.Tensor, + nvb: int, + eps: float, + *, + output_mc_ptr: int, + max_blocks: int = 128, + write_prefix: bool = False, +) -> torch.Tensor: + """Fuse local attention-residual aggregation with direct multicast AG. + + `out` is the full symmetric [world * local_tokens, H] output. Each rank + aggregates its local token shard and multicast-stores normalized vectors + directly from consumer registers into that rank's slice on every peer. + """ + _attn_res_fused_direct_ag_op( + world_size, + prefix_sum, + bank, + cw, + ow, + out, + nvb, + eps, + output_mc_ptr, + max_blocks, + write_prefix, + *_tuning(nvb, prefix_sum.shape[0]), + ) + return out + + +def attn_res_fused_pull_rs( + world_size: int, + input: torch.Tensor, + residual: torch.Tensor | None, + bank: torch.Tensor, + cw: torch.Tensor, + ow: torch.Tensor, + out: torch.Tensor, + prefix_out: torch.Tensor, + nvb: int, + eps: float, + *, + input_mc_ptr: int, + max_blocks: int = 128, +) -> tuple[torch.Tensor, torch.Tensor]: + """Fused NVLS pull RS + local residual + TMA attention aggregation. + + `input` is the full TP-partial o_proj tensor in multicast symmetric + memory. Every rank reduces only its contiguous token shard, optionally + adds its already-local residual, materializes that prefix in `prefix_out`, + and feeds it directly to the fused attention-residual aggregation/norm. + """ + _attn_res_fused_pull_rs_op( + world_size, + input, + residual, + bank, + cw, + ow, + out, + prefix_out, + nvb, + eps, + input_mc_ptr, + max_blocks, + # prefix_out, not input: the shard this rank actually aggregates. + *_tuning(nvb, prefix_out.shape[0]), + ) + return out, prefix_out diff --git a/python/sglang/kernels/ops/kimi_k3/configs/sp_collective/world=4,H=7168,device_name=NVIDIA_GB300.json b/python/sglang/kernels/ops/kimi_k3/configs/sp_collective/world=4,H=7168,device_name=NVIDIA_GB300.json new file mode 100644 index 000000000..8317a03a0 --- /dev/null +++ b/python/sglang/kernels/ops/kimi_k3/configs/sp_collective/world=4,H=7168,device_name=NVIDIA_GB300.json @@ -0,0 +1,271 @@ +{ + "source": { + "code_commits": [ + "c4ad2e84b", + "5bf69bd32" + ], + "date": "2026-07-25", + "device": "NVIDIA GB300", + "nodes": 1, + "gpus_per_node": 4, + "push_slot_bytes": 33554432, + "raw_results": [ + "standalone-c4ad2e84b-run3.log", + "standalone-large-c4ad2e84b.json", + "standalone-8192-tiebreak-c4ad2e84b.json", + "fused-small-5bf69bd32.json", + "fused-medium-5bf69bd32.json", + "fused-8192-5bf69bd32.json" + ] + }, + "selection": { + "rule": "nearest global-token bucket at or below the workload", + "fallback": "nccl", + "note": "Standalone RS push won through T=2048 and NVLS pull won at T=4096/8192. Fused pull RS+residual+attention-residual wins through T=256; the separate path wins from T=512. Fused attention-residual+direct AG wins every measured bucket T=4..8192. The T=8192 standalone and fused AG custom wins were repeated with higher iteration counts. T=9364 is the first world=4 shape whose per-rank shard exceeds the 32 MiB symmetric slot, so it introduces explicit NCCL/separate safety boundaries." + }, + "configs": { + "reduce_scatter": { + "4": { + "strategy": "push", + "num_blocks": 128, + "block_size": 512 + }, + "8": { + "strategy": "push", + "num_blocks": 96, + "block_size": 256 + }, + "16": { + "strategy": "push", + "num_blocks": 96, + "block_size": 512 + }, + "32": { + "strategy": "push", + "num_blocks": 96, + "block_size": 512 + }, + "64": { + "strategy": "push", + "num_blocks": 64, + "block_size": 256 + }, + "128": { + "strategy": "push", + "num_blocks": 64, + "block_size": 512 + }, + "256": { + "strategy": "push", + "num_blocks": 128, + "block_size": 256 + }, + "512": { + "strategy": "push", + "num_blocks": 128, + "block_size": 256 + }, + "1024": { + "strategy": "push", + "num_blocks": 128, + "block_size": 512 + }, + "2048": { + "strategy": "push", + "num_blocks": 128, + "block_size": 512 + }, + "4096": { + "strategy": "pull", + "num_blocks": 96, + "block_size": 1024 + }, + "8192": { + "strategy": "pull", + "num_blocks": 96, + "block_size": 1024 + }, + "9364": { + "strategy": "nccl" + }, + "16384": { + "strategy": "nccl" + } + }, + "all_gather": { + "4": { + "strategy": "direct", + "num_blocks": 8, + "block_size": 1024 + }, + "8": { + "strategy": "push", + "num_blocks": 32, + "block_size": 128 + }, + "16": { + "strategy": "push", + "num_blocks": 64, + "block_size": 256 + }, + "32": { + "strategy": "direct", + "num_blocks": 64, + "block_size": 128 + }, + "64": { + "strategy": "push", + "num_blocks": 32, + "block_size": 128 + }, + "128": { + "strategy": "push", + "num_blocks": 96, + "block_size": 256 + }, + "256": { + "strategy": "direct", + "num_blocks": 64, + "block_size": 256 + }, + "512": { + "strategy": "push", + "num_blocks": 128, + "block_size": 512 + }, + "1024": { + "strategy": "direct", + "num_blocks": 96, + "block_size": 1024 + }, + "2048": { + "strategy": "direct", + "num_blocks": 64, + "block_size": 512 + }, + "4096": { + "strategy": "direct", + "num_blocks": 64, + "block_size": 256 + }, + "8192": { + "strategy": "direct", + "num_blocks": 64, + "block_size": 1024 + }, + "9364": { + "strategy": "nccl" + }, + "16384": { + "strategy": "nccl" + } + }, + "reduce_scatter_attn_res": { + "4": { + "strategy": "fused_pull", + "max_blocks": 64 + }, + "8": { + "strategy": "fused_pull", + "max_blocks": 64 + }, + "16": { + "strategy": "fused_pull", + "max_blocks": 64 + }, + "32": { + "strategy": "fused_pull", + "max_blocks": 64 + }, + "64": { + "strategy": "fused_pull", + "max_blocks": 64 + }, + "128": { + "strategy": "fused_pull", + "max_blocks": 64 + }, + "256": { + "strategy": "fused_pull", + "max_blocks": 64 + }, + "512": { + "strategy": "separate" + }, + "1024": { + "strategy": "separate" + }, + "2048": { + "strategy": "separate" + }, + "4096": { + "strategy": "separate" + }, + "8192": { + "strategy": "separate" + }, + "9364": { + "strategy": "separate" + }, + "16384": { + "strategy": "separate" + } + }, + "attn_res_all_gather": { + "4": { + "strategy": "fused_direct", + "max_blocks": 64 + }, + "8": { + "strategy": "fused_direct", + "max_blocks": 64 + }, + "16": { + "strategy": "fused_direct", + "max_blocks": 64 + }, + "32": { + "strategy": "fused_direct", + "max_blocks": 64 + }, + "64": { + "strategy": "fused_direct", + "max_blocks": 64 + }, + "128": { + "strategy": "fused_direct", + "max_blocks": 64 + }, + "256": { + "strategy": "fused_direct", + "max_blocks": 64 + }, + "512": { + "strategy": "fused_direct", + "max_blocks": 64 + }, + "1024": { + "strategy": "fused_direct", + "max_blocks": 64 + }, + "2048": { + "strategy": "fused_direct", + "max_blocks": 64 + }, + "4096": { + "strategy": "fused_direct", + "max_blocks": 64 + }, + "8192": { + "strategy": "fused_direct", + "max_blocks": 64 + }, + "9364": { + "strategy": "separate" + }, + "16384": { + "strategy": "separate" + } + } + } +} diff --git a/python/sglang/kernels/ops/kimi_k3/configs/sp_collective/world=8,H=7168,device_name=NVIDIA_GB300.json b/python/sglang/kernels/ops/kimi_k3/configs/sp_collective/world=8,H=7168,device_name=NVIDIA_GB300.json new file mode 100644 index 000000000..e1d923823 --- /dev/null +++ b/python/sglang/kernels/ops/kimi_k3/configs/sp_collective/world=8,H=7168,device_name=NVIDIA_GB300.json @@ -0,0 +1,235 @@ +{ + "source": { + "code_commits": [ + "8acabdbb7", + "e6bbf14d7", + "a74213e8c", + "7b1d23a0b", + "a132f1025", + "f574cdc8a" + ], + "date": "2026-07-25", + "device": "NVIDIA GB300", + "nodes": 2, + "gpus_per_node": 4, + "push_slot_bytes": 33554432, + "raw_results": [ + "tune-8acabdbb7.json", + "ag-strategies-e6bbf14d7.json", + "rs-strategies-a74213e8c.json", + "sp-attn-res-composite-full-7b1d23a0b.json", + "fused-pull-attn-res-full-a132f1025.json", + "fused-direct-ag-full-f574cdc8a.json" + ] + }, + "selection": { + "rule": "nearest global-token bucket at or below the workload", + "fallback": "nccl", + "note": "Standalone RS push won through T=1024 and NVLS pull won at T=2048/4096. Fused pull RS+residual+attention-residual wins through T=512; the separate path wins from T=1024. Fused attention-residual+direct AG wins every measured bucket T=8..4096. NCCL won standalone collectives at T=8192 and T=16384; fusion falls back outside its measured envelope." + }, + "configs": { + "reduce_scatter": { + "8": { + "strategy": "push", + "num_blocks": 128, + "block_size": 512 + }, + "16": { + "strategy": "push", + "num_blocks": 4, + "block_size": 512 + }, + "32": { + "strategy": "push", + "num_blocks": 96, + "block_size": 512 + }, + "64": { + "strategy": "push", + "num_blocks": 16, + "block_size": 512 + }, + "128": { + "strategy": "push", + "num_blocks": 128, + "block_size": 512 + }, + "256": { + "strategy": "push", + "num_blocks": 128, + "block_size": 512 + }, + "512": { + "strategy": "push", + "num_blocks": 128, + "block_size": 512 + }, + "1024": { + "strategy": "push", + "num_blocks": 96, + "block_size": 512 + }, + "2048": { + "strategy": "pull", + "num_blocks": 96, + "block_size": 1024 + }, + "4096": { + "strategy": "pull", + "num_blocks": 96, + "block_size": 1024 + }, + "8192": { + "strategy": "nccl" + }, + "16384": { + "strategy": "nccl" + } + }, + "all_gather": { + "8": { + "strategy": "push", + "num_blocks": 128, + "block_size": 128 + }, + "16": { + "strategy": "push", + "num_blocks": 128, + "block_size": 512 + }, + "32": { + "strategy": "push", + "num_blocks": 128, + "block_size": 512 + }, + "64": { + "strategy": "push", + "num_blocks": 32, + "block_size": 256 + }, + "128": { + "strategy": "push", + "num_blocks": 64, + "block_size": 512 + }, + "256": { + "strategy": "push", + "num_blocks": 128, + "block_size": 512 + }, + "512": { + "strategy": "push", + "num_blocks": 128, + "block_size": 512 + }, + "1024": { + "strategy": "direct", + "num_blocks": 2, + "block_size": 1024 + }, + "2048": { + "strategy": "direct", + "num_blocks": 2, + "block_size": 1024 + }, + "4096": { + "strategy": "direct", + "num_blocks": 2, + "block_size": 1024 + }, + "8192": { + "strategy": "nccl" + }, + "16384": { + "strategy": "nccl" + } + }, + "reduce_scatter_attn_res": { + "8": { + "strategy": "fused_pull", + "max_blocks": 64 + }, + "16": { + "strategy": "fused_pull", + "max_blocks": 64 + }, + "32": { + "strategy": "fused_pull", + "max_blocks": 64 + }, + "64": { + "strategy": "fused_pull", + "max_blocks": 64 + }, + "128": { + "strategy": "fused_pull", + "max_blocks": 64 + }, + "256": { + "strategy": "fused_pull", + "max_blocks": 64 + }, + "512": { + "strategy": "fused_pull", + "max_blocks": 64 + }, + "1024": { + "strategy": "separate" + }, + "2048": { + "strategy": "separate" + }, + "4096": { + "strategy": "separate" + }, + "8192": { + "strategy": "separate" + } + }, + "attn_res_all_gather": { + "8": { + "strategy": "fused_direct", + "max_blocks": 64 + }, + "16": { + "strategy": "fused_direct", + "max_blocks": 64 + }, + "32": { + "strategy": "fused_direct", + "max_blocks": 64 + }, + "64": { + "strategy": "fused_direct", + "max_blocks": 64 + }, + "128": { + "strategy": "fused_direct", + "max_blocks": 64 + }, + "256": { + "strategy": "fused_direct", + "max_blocks": 64 + }, + "512": { + "strategy": "fused_direct", + "max_blocks": 64 + }, + "1024": { + "strategy": "fused_direct", + "max_blocks": 64 + }, + "2048": { + "strategy": "fused_direct", + "max_blocks": 64 + }, + "4096": { + "strategy": "fused_direct", + "max_blocks": 64 + }, + "8192": { + "strategy": "separate" + } + } + } +} diff --git a/python/sglang/kernels/ops/kimi_k3/gemm_ag.py b/python/sglang/kernels/ops/kimi_k3/gemm_ag.py new file mode 100644 index 000000000..280cdad3a --- /dev/null +++ b/python/sglang/kernels/ops/kimi_k3/gemm_ag.py @@ -0,0 +1,88 @@ +"""K3 column-parallel up_proj + multicast all-gather + add3 (bf16, TP8). + +One entry point over ``csrc/kimi_k3/comm/gemm_ag.cuh``: for the latent MoE +up_proj ([M, 3584] x [3584, 7168]) at small decode M, every rank computes +only its 896-column slice of the replicated GEMM (the C++ side slices the +full weight itself), multicast-stores it into the CustomAllReduceV2 push +workspace (one more user of its double-buffer phase protocol), and a +Lamport-spin consumer assembles ``out = up_proj(x) + b (+ c)`` — reading +1/8 of the weight bytes per rank instead of all of them. Needs +:func:`sglang.kernels.ops.kimi_k3.all_reduce.register_comm` once beforehand +(the same registration the push all-reduce uses). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.kernels.jit.utils import ( + cache_once, + is_arch_support_pdl, + load_jit, + make_cpp_args, +) +from sglang.kernels.ops.kimi_k3.all_reduce import _COMM_MAP +from sglang.srt.utils.custom_op import register_custom_op + +if TYPE_CHECKING: + from tvm_ffi.module import Module + +# Kimi-K3 up_proj dims (the kernel template takes any K/N passing its +# static_asserts; this module instantiates the K3 shape). +K = 3584 +N = 7168 + +# Largest decode batch the kernel wins at (crossover vs the replicated +# cublas GEMM + add3 tail is ~13-14 tokens on B200x8); also the GEMV +# function-table size. +MAX_TOKENS = 12 + + +@cache_once +def _jit_module() -> Module: + args = make_cpp_args(K, N, MAX_TOKENS, is_arch_support_pdl()) + cls = f"GEMMAGKernel<{args}>" + return load_jit( + "kimi_k3_gemm_ag", + *args, + cuda_files=["kimi_k3/comm/gemm_ag.cuh"], + cuda_wrappers=[("run", f"{cls}::run")], + extra_cuda_cflags=["-O3"], + ) + + +@register_custom_op(mutates_args=["out"]) +def _gemm_ag_op( + world_size: int, + x: torch.Tensor, + weight: torch.Tensor, + b: torch.Tensor, + c: Optional[torch.Tensor], + out: torch.Tensor, + ws_mc_base: int, +) -> None: + comm = _COMM_MAP[world_size].obj + _jit_module().run(comm, x, weight, b, c, out, ws_mc_base) + + +def gemm_ag_up_proj( + world_size: int, + x: torch.Tensor, + weight: torch.Tensor, + b: torch.Tensor, + c: Optional[torch.Tensor], + out: torch.Tensor, + *, + ws_mc_base: int, +) -> torch.Tensor: + """``out = x @ weight.T (allgathered) + b (+ c)``, all bf16. + + ``x`` is [M, 3584] with M in [1, MAX_TOKENS]; ``weight`` is the FULL + replicated [7168, 3584] up_proj weight (each rank reads only its own + row block); ``b`` / ``c`` / ``out`` are [M, 7168] (``out`` is + output-only). ``ws_mc_base`` is the multicast VA of the v2 workspace + slab base (``comm.mc_base_ptr``).""" + _gemm_ag_op(world_size, x, weight, b, c, out, ws_mc_base) + return out diff --git a/python/sglang/kernels/ops/kimi_k3/gemm_ar.py b/python/sglang/kernels/ops/kimi_k3/gemm_ar.py new file mode 100644 index 000000000..9d398e504 --- /dev/null +++ b/python/sglang/kernels/ops/kimi_k3/gemm_ar.py @@ -0,0 +1,216 @@ +"""K3 fused o_proj GEMM + all-reduce for decode (bf16, TP row-parallel). + +One entry point over ``csrc/kimi_k3/comm/gemm_ar.cuh``: a single kernel per +rank computes the local ``x_r [M, K] @ W_r [7168, K]^T`` partial AND the +cross-rank sum — the epilogue pushes finished tiles straight into a +peer-mapped P2P comm region, one flag boundary, then a tile-local reduce +writes the fully reduced ``out [M, 7168]`` on every rank. Replaces the +o_proj GEMM + NCCL all-reduce pair with one launch (see GEMM_AR_README.md). + +Contracts: + +* bf16 only; ``out = sum_r bf16(x_r @ W_r^T)`` (partials round to bf16 + pre-sum — same numerics as the unfused bf16 GEMM + ring AR). +* M in [1, 512]; internally rounded up to a tuned cell {8, 16, 32, 64, + 128, 256, 512}. ``out`` is allocated with ``cell`` rows and sliced. +* SM100+ with full NVLink P2P (fabric/MNNVL across nodes); perf-tuned on + GB300 (sm_103a). +* CUDA-graph compatible: the per-cell launch epoch lives in device memory + (read at kernel entry, bumped by a trailing kernel), so replays advance + it naturally. Each dispatch cell owns its own flag-ring family — no + host-side ring reset, ever. +* All TP ranks must call :func:`o_proj_gemm_ar` with the same M in + lockstep (same stream order of cells on every rank). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, NamedTuple, Optional + +import torch + +from sglang.kernels.jit.utils import ( + cache_once, + is_arch_support_pdl, + load_jit, + make_cpp_args, +) +from sglang.srt.utils.custom_op import register_custom_op + +if TYPE_CHECKING: + from tvm_ffi.module import Module + +N = 7168 # K3 hidden size (OPROJ_N compile-time default in gemm_ar.cuh) +MAX_TOKENS = 512 # kMMax + + +@cache_once +def _jit_module(k: int, world_size: int) -> Module: + args = make_cpp_args( + k, + world_size, + is_arch_support_pdl(), + ) + cls = f"GemmArKernel<{args}>" + return load_jit( + "kimi_k3_gemm_ar", + *args, + cuda_files=["kimi_k3/comm/gemm_ar.cuh"], + cuda_wrappers=[ + ("run", f"{cls}::run"), + ("set_bases", f"{cls}::set_bases"), + ("region_nbytes", f"{cls}::region_nbytes"), + ("gather_words", f"{cls}::gather_words"), + ("num_fams", f"{cls}::num_fams"), + ], + extra_cuda_cflags=["-O3"], + extra_dependencies=["cutlass"], + ) + + +class _State(NamedTuple): + world_size: int + rank: int + region: tuple # (slab tensor, peer buffer views) — keeps mappings alive + uc_bases: torch.Tensor # [R] int64 CPU: per-rank UC VAs of the region + gather: torch.Tensor # [kFams * 2 * kRing] int32 CUDA, device-local + epochs: torch.Tensor # [kFams] int32 CUDA: device-resident CTA ticket counters + + +_STATE: Optional[_State] = None + + +def init( + *, + world_size: int, + rank: int, + group: torch.distributed.ProcessGroup, + k: int, +) -> None: + """Allocate + rendezvous the P2P comm region (collective; call once from + every TP rank, BEFORE any CUDA-graph capture). ``group`` is the TP CPU + (gloo) group used for the symm-mem rendezvous.""" + global _STATE + if _STATE is not None: + return + # the empty_strided_p2p + get_buffer path is the one that exchanges + # fabric handles and maps every peer (incl. cross-node MNNVL) into this + # process — the mem-pool rendezvous(tensor, group_name) API leaves + # remote-node (and sometimes even local) peers unmapped. + from torch._C._distributed_c10d import _SymmetricMemory + + mod = _jit_module(k, world_size) + nbytes = int(mod.region_nbytes()) + device = torch.device("cuda", torch.cuda.current_device()) + if torch.__version__ < "2.11.0": + import torch.distributed._symmetric_memory as torch_symm_mem + + torch_symm_mem.enable_symm_mem_for_group(group.group_name) + region = _SymmetricMemory.empty_strided_p2p( + (nbytes,), [1], torch.uint8, device, group.group_name + ) + symm = _SymmetricMemory.rendezvous(region) + region.zero_() + torch.cuda.synchronize() + torch.distributed.barrier(group=group) + # keep the peer buffer tensors alive alongside the region + peer_bufs = [symm.get_buffer(r, [nbytes], torch.uint8) for r in range(world_size)] + ptrs = [t.data_ptr() for t in peer_bufs] + import logging + + logging.getLogger(__name__).info( + "gemm_ar comm region: rank=%d nbytes=%d uc_bases=%s", + rank, + nbytes, + [hex(p) for p in ptrs], + ) + assert all(p != 0 for p in ptrs), f"gemm_ar: null peer pointers {ptrs}" + # explicit cpu: model build may run under a cuda default-device context, + # and a silently-cuda tensor here means the host-side deref in set_bases + # reads a device pointer (segfault) + uc_bases = torch.tensor(ptrs, dtype=torch.int64, device="cpu") + gather = torch.zeros(int(mod.gather_words()), dtype=torch.int32, device=device) + epochs = torch.zeros(int(mod.num_fams()), dtype=torch.int32, device=device) + torch.cuda.synchronize() + _STATE = _State( + world_size=world_size, + rank=rank, + region=(region, peer_bufs), + uc_bases=uc_bases, + gather=gather, + epochs=epochs, + ) + + +def initialized() -> bool: + return _STATE is not None + + +@cache_once +def _module_with_bases(k: int, world_size: int) -> Module: + """The per-K JIT module with the comm-region base addresses stashed + host-side (per-call CPU-tensor derefs from inside the custom op segfault + under the sglang scheduler, so the module holds them in a static).""" + state = _STATE + assert state is not None + mod = _jit_module(k, world_size) + mod.set_bases(state.uc_bases) + return mod + + +def fits(x: torch.Tensor) -> bool: + """Whether this o_proj input can take the fused GEMM+AR path.""" + return ( + _STATE is not None + and x.dim() == 2 + and x.dtype == torch.bfloat16 + and 0 < x.shape[0] <= MAX_TOKENS + and x.stride(1) == 1 + and x.stride(0) == x.shape[1] + ) + + +@register_custom_op(mutates_args=["out", "epochs"]) +def _gemm_ar_op( + k: int, + world_size: int, + out: torch.Tensor, + x: torch.Tensor, + weight: torch.Tensor, + gather: torch.Tensor, + epochs: torch.Tensor, + my_rank: int, +) -> None: + _module_with_bases(k, world_size).run(out, x, weight, gather, epochs, my_rank) + + +def _cell_of(m: int) -> int: + for c in (8, 16, 32, 64, 128, 256, 512): + if m <= c: + return c + raise ValueError(f"gemm_ar: M={m} outside [1, {MAX_TOKENS}]") + + +def o_proj_gemm_ar(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """Fully reduced ``sum_r x_r @ weight_r^T`` on every rank, one kernel. + + ``x`` is the TP-local [M, K] o_proj input, ``weight`` the TP-local + [7168, K] o_proj weight shard. Caller checked :func:`fits`; all ranks + call in lockstep with the same M. + """ + state = _STATE + assert state is not None + m = x.shape[0] + cell = _cell_of(m) + out = torch.empty((cell, N), dtype=torch.bfloat16, device=x.device) + _gemm_ar_op( + weight.shape[1], + state.world_size, + out, + x, + weight, + state.gather, + state.epochs, + state.rank, + ) + return out[:m] diff --git a/python/sglang/kernels/ops/kimi_k3/kda_decode_mtp.py b/python/sglang/kernels/ops/kimi_k3/kda_decode_mtp.py new file mode 100644 index 000000000..554dd753c --- /dev/null +++ b/python/sglang/kernels/ops/kimi_k3/kda_decode_mtp.py @@ -0,0 +1,1132 @@ +"""CuTe DSL device kernel KDA conv-MTP. + +conv enabled, no bias, optional fused gated RMSNorm, lower_bound gate, Q/K +L2 norm, beta sigmoid, ILP=2, W=4. Recurrent-state tiles are cp.async'd into +NUM_STATE_STAGES smem stages. Phase 2 walks the V // TILE_V state tiles in +passes of TILES_PER_PASS: a multi-token verify keeps them all register-resident +across the token chain, a single-token step streams one at a time. +""" + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +from cutlass._mlir.dialects import nvvm +from cutlass.cute.nvgpu import cpasync + +WARP_SIZE = 32 +TILE_K = 128 +KERNEL_WIDTH = 4 +# Phase 1 conv jobs: q, k, g. Each gets its own equal-sized warp group; the +# v-conv gets whatever is left over. +P1_NUM_JOBS = 3 +# Share of the block's warps handed to the v-conv, as a divisor: 4 -> a quarter. +# Tuned, not structural — v is one conv over V channels with no norm or gate, so +# it needs fewer warps than q/k/g, but the exact split is empirical. +P1_V_WARP_DIVISOR = 4 +# Recurrent-state smem stages, i.e. how many of the V // TILE_V state tiles are +# in flight at once. Three measured worse than two even for the streaming path, +# where a deeper pipeline should have helped: the transfer is not the bottleneck +# and the extra stage only costs smem. +NUM_STATE_STAGES = 2 + +# Block size is chosen per launch, not fixed: below one wave (H*N <= SM count) +# every SM gets one block regardless, so a wider block is free warps. Above it, +# two narrow blocks per SM let one issue while the other sits at a barrier, +# which a single wide block cannot do. See _block_threads. +BLOCK_THREADS_NARROW = 256 +BLOCK_THREADS_WIDE = 512 + +# Recurrence lane split: K spans a subgroup instead of all 32 lanes, so several +# v-rows reduce concurrently and each butterfly is shorter. The resident path +# uses eight lanes; the ns0 streaming path uses sixteen because halving its +# repeated q/k/decay LDS sequence outweighs the one extra shuffle. +P2_LANES_RESIDENT = 8 +P2_LANES_STREAM = 16 + +HEAD_DIM = TILE_K +VEC_SIZE = HEAD_DIM // WARP_SIZE +# Conv weights live in one smem array: [W, K] for q/k, then [W, V] for v. +V_WEIGHT_BASE = KERNEL_WIDTH * HEAD_DIM +CONV_WEIGHT_ELEMS = 2 * V_WEIGHT_BASE + + +def _stream_state(*, block_threads: int, num_spec: int) -> bool: + """Whether phase 2 streams the recurrent state instead of holding it. + + See the TILES_PER_PASS comment in the kernel: streaming trades register + residency for occupancy, which needs a single token (nothing to lose) and a + grid above one wave (something to gain, signalled by the narrow block). + """ + return num_spec == 0 and block_threads == BLOCK_THREADS_NARROW + + +def _p2_lanes(*, block_threads: int, num_spec: int) -> int: + return ( + P2_LANES_STREAM + if _stream_state(block_threads=block_threads, num_spec=num_spec) + else P2_LANES_RESIDENT + ) + + +def _qk_smem(*, block_threads: int, num_spec: int) -> tuple[tuple[int, ...], int]: + """Stride tuple and per-token element count for sQ/sK/sG. + + Both modes index the same (token, lane row, vector, element) shape; only the + strides differ. Lane-major (see P2_QK_ROW) makes a lane's P2_VEC channels + contiguous so they load as 128-bit vectors, at the cost of the bank pad. + Channel-major places element (k_grp, jv, c) at plain channel + k_grp + (jv * 4 + c) * P2_LANES_K, which is what the scalar path wants: no + pad, and the same conflict-free access as an unpermuted [token, K] tile. + """ + p2_lanes = _p2_lanes(block_threads=block_threads, num_spec=num_spec) + p2_vec = TILE_K // p2_lanes + p2_qk_row = p2_vec + 4 + if _stream_state(block_threads=block_threads, num_spec=num_spec): + return (p2_lanes * p2_qk_row, p2_qk_row, 4, 1), p2_lanes * p2_qk_row + return (TILE_K, 1, 4 * p2_lanes, p2_lanes), TILE_K + + +def _state_tile_v(*, block_threads: int, num_spec: int) -> int: + """Height of the recurrent-state tile staged in smem, in value rows. + + Never below P2_ROWS_LANE rows per warp: a warp's P2_LANES_K-lane groups + reduce that many rows concurrently, so a shorter tile leaves lanes idle. + + Streaming wants exactly that floor -- both the smem stage and the + register-resident slice scale with the tile, and a consumer is waiting per + tile, so a short tile pipelines. Residency wants the opposite: no token can + start until every tile has landed, so there is nothing to overlap and the + best shape is the fewest, largest tiles that keep all NUM_STATE_STAGES in + flight at once with no stage reuse. + """ + p2_rows_lane = WARP_SIZE // _p2_lanes( + block_threads=block_threads, num_spec=num_spec + ) + min_rows = (block_threads // WARP_SIZE) * p2_rows_lane + if _stream_state(block_threads=block_threads, num_spec=num_spec): + # Sixteen rows doubles the stage waits and loses ~12% at B32. A 32-row + # tile retains four phases while the sixteen-lane split still halves + # the q/k/decay register arrays and LDS sequence. + return max(min_rows, 32) + return max(min_rows, HEAD_DIM // NUM_STATE_STAGES) + + +def _issue_state_tile( + state_g2s_copy: cute.TiledCopy, + thr_state_copy, + gStateTiles: cute.Tensor, + sState: cute.Tensor, + i_v: int, + num_stages: int, +) -> None: + """cp.async state tile ``i_v`` into the stage it maps to, and commit it.""" + cute.copy( + state_g2s_copy, + thr_state_copy.partition_S(gStateTiles[(None, None, i_v)]), + thr_state_copy.partition_D(sState[(None, None, i_v % num_stages)]), + ) + cute.arch.cp_async_commit_group() + + +@cute.kernel +def kda_decode_mtp_kernel( + state_g2s_copy: cute.TiledCopy, + h0: cute.Tensor, + x_q: cute.Tensor, + x_k: cute.Tensor, + x_v: cute.Tensor, + w_q: cute.Tensor, + w_k: cute.Tensor, + w_v: cute.Tensor, + cs_q: cute.Tensor, + cs_k: cute.Tensor, + cs_v: cute.Tensor, + A_log: cute.Tensor, + g: cute.Tensor, + dt_bias: cute.Tensor, + beta: cute.Tensor, + o: cute.Tensor, + ht: cute.Tensor, + intermediate_state_indices: cute.Tensor, + intermediate_conv_q: cute.Tensor, + intermediate_conv_k: cute.Tensor, + intermediate_conv_v: cute.Tensor, + ring_rawv: cute.Tensor, + ring_rawk: cute.Tensor, + ring_g: cute.Tensor, + ring_beta: cute.Tensor, + onorm_g: cute.Tensor, + onorm_weight: cute.Tensor, + smem_qk_layout: cute.Layout, + smem_state_layout: cute.Layout, + ssm_state_indices: cute.Tensor, + cu_seqlens: cute.Tensor, + scale: cutlass.Constexpr[float], + NUM_SPEC: cutlass.Constexpr[int], + BLOCK_THREADS: cutlass.Constexpr[int], + P2_LANES_K: cutlass.Constexpr[int], + lower_bound: cutlass.Constexpr[float], + CACHE_RING: cutlass.Constexpr[bool], + APPLY_ONORM: cutlass.Constexpr[bool], + onorm_eps: cutlass.Constexpr[float], +): + """KDA MTP decode — SMEM pre-compute + register-resident state. + + One block owns all 128 value rows, so APPLY_ONORM can reduce the RMS + denominator without cross-block synchronization. + """ + tidx, _, _ = cute.arch.thread_idx() + in_warp_tid = tidx % WARP_SIZE + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + i_hv, i_n, _ = cute.arch.block_idx() + head_off = i_hv * HEAD_DIM + T_LOOP = 1 + NUM_SPEC + # Conv-precompute warp budget: q/k/g each get P1_JOB_WARPS warps and split + # the token dimension across them; the rest carry the v-conv. + NUM_WARPS = BLOCK_THREADS // WARP_SIZE + P1_V_WARPS = max(1, NUM_WARPS // P1_V_WARP_DIVISOR) + P1_JOB_WARPS = (NUM_WARPS - P1_V_WARPS) // P1_NUM_JOBS + P1_QKG_WARPS = P1_NUM_JOBS * P1_JOB_WARPS + V_CH_PER_THREAD = TILE_K // (P1_V_WARPS * WARP_SIZE) + P2_ROWS_LANE = WARP_SIZE // P2_LANES_K + P2_VEC = TILE_K // P2_LANES_K + TILE_V = _state_tile_v(block_threads=BLOCK_THREADS, num_spec=NUM_SPEC) + NUM_V_TILES = HEAD_DIM // TILE_V + STATE_STAGES = min(NUM_STATE_STAGES, NUM_V_TILES) + NUM_V_ROWS = TILE_V // NUM_WARPS + P2_BATCHES = NUM_V_ROWS // P2_ROWS_LANE + # How many state tiles phase 2 holds in registers at once. Holding all of + # them keeps the whole 128x128 state resident across the token chain, which + # is what makes multi-token verify cheap -- but it costs + # HEAD_DIM * TILE_K / BLOCK_THREADS registers however the lanes are split, + # and that is what caps the kernel at two blocks per SM. Streaming one tile + # at a time spends those registers on occupancy instead, which is worth it + # only when there is nothing to lose (a single token, so no reuse) and + # something to gain (a grid above one wave, which is exactly what the narrow + # block width signals -- see _block_threads). + STREAM_STATE = _stream_state(block_threads=BLOCK_THREADS, num_spec=NUM_SPEC) + TILES_PER_PASS = 1 if STREAM_STATE else NUM_V_TILES + smem = cutlass.utils.SmemAllocator() + sQ = smem.allocate_tensor(cutlass.Float32, smem_qk_layout, 16) + sK = smem.allocate_tensor(cutlass.Float32, smem_qk_layout, 16) + sG = smem.allocate_tensor(cutlass.Float32, smem_qk_layout, 16) + sBeta = smem.allocate_tensor(cutlass.Float32, cute.make_layout((T_LOOP,)), 16) + sVall = smem.allocate_tensor( + cutlass.Float32, cute.make_layout((T_LOOP * HEAD_DIM,)), 16 + ) + sConvW = smem.allocate_tensor( + cutlass.Float32, + cute.make_layout((CONV_WEIGHT_ELEMS,)), + 16, + ) + sState = smem.allocate_tensor(cutlass.Float32, smem_state_layout, 16) + if cutlass.const_expr(APPLY_ONORM): + sOall = smem.allocate_tensor( + cutlass.Float32, cute.make_layout((T_LOOP * HEAD_DIM,)), 16 + ) + else: + # Compile-time-dead placeholder; avoids charging the non-norm path + # for an output tile it never touches. + sOall = sVall + r_q = cute.make_rmem_tensor( + cute.make_layout((P2_VEC,), stride=(1,)), cutlass.Float32 + ) + r_k = cute.make_rmem_tensor( + cute.make_layout((P2_VEC,), stride=(1,)), cutlass.Float32 + ) + r_decay = cute.make_rmem_tensor( + cute.make_layout((P2_VEC,), stride=(1,)), cutlass.Float32 + ) + # Sized for whichever is larger: one pass of recurrent state (phase 2) or + # the two conv windows (phase 1). The phases never overlap. + R_STATE_ELEMS = max( + TILES_PER_PASS * P2_BATCHES * P2_VEC, 2 * (KERNEL_WIDTH - 1) * VEC_SIZE + ) + r_state = cute.make_rmem_tensor( + cute.make_layout((R_STATE_ELEMS,), stride=(1,)), cutlass.Float32 + ) + r_wq = cute.make_rmem_tensor( + cute.make_layout((KERNEL_WIDTH * VEC_SIZE,), stride=(1,)), cutlass.Float32 + ) + # One channel's KERNEL_WIDTH conv taps are contiguous in the [dim, W] + # weights, so they load as a single 16-byte vector instead of W strided + # scalars. These run once per block, i.e. entirely in the fixed cost. + r_w4 = cute.make_rmem_tensor( + cute.make_layout((KERNEL_WIDTH,), stride=(1,)), cutlass.Float32 + ) + if cutlass.const_expr(STREAM_STATE): + r_v4q = cute.make_rmem_tensor( + cute.make_layout((4,), stride=(1,)), cutlass.Float32 + ) + r_v4k = cute.make_rmem_tensor( + cute.make_layout((4,), stride=(1,)), cutlass.Float32 + ) + r_v4g = cute.make_rmem_tensor( + cute.make_layout((4,), stride=(1,)), cutlass.Float32 + ) + + slot = ssm_state_indices[i_n] + # CUDA-graph padding rows use slot == -1. + if slot < 0: + cute.arch.griddepcontrol_wait() + pad_bos = cu_seqlens[i_n] + for i_t in cutlass.range_constexpr(T_LOOP): + if tidx < HEAD_DIM: + o[0, pad_bos + i_t, i_hv, tidx] = cutlass.BFloat16(0.0) + cute.arch.griddepcontrol_launch_dependents() + # nvvm.exit, not `return`: the DSL rejects an early return out of a + # staged if (UNSUP_EARLY_EXIT). + nvvm.exit() + + # q/k/g each run on P1_JOB_WARPS warps split by token parity and the v-conv + # takes the rest. Each token's conv is an independent window over globals, + # so the split needs no cross-warp communication. + p1_job = warp_idx % P1_NUM_JOBS + p1_par = warp_idx // P1_NUM_JOBS + # Only the q-conv warps read the q window and only the k-conv warps the k + # window; the g and v warps overwrite r_state in phase 2 without reading it. + if warp_idx < P1_QKG_WARPS: + if p1_job == 0: + for i in range(VEC_SIZE): + k_idx = i * 32 + in_warp_tid + for w in range(KERNEL_WIDTH - 1): + r_state[w * VEC_SIZE + i] = cutlass.Float32( + cs_q[slot, head_off + k_idx, w] + ) + elif p1_job == 1: + for i in range(VEC_SIZE): + k_idx = i * 32 + in_warp_tid + for w in range(KERNEL_WIDTH - 1): + r_state[(KERNEL_WIDTH - 1) * VEC_SIZE + w * VEC_SIZE + i] = ( + cutlass.Float32(cs_k[slot, head_off + k_idx, w]) + ) + + if tidx < HEAD_DIM: + cute.autovec_copy(w_k[(head_off + tidx, None)], r_w4) + for w in range(KERNEL_WIDTH): + sConvW[w * HEAD_DIM + tidx] = r_w4[w] + if tidx < HEAD_DIM: + cute.autovec_copy(w_v[(head_off + tidx, None)], r_w4) + for w in range(KERNEL_WIDTH): + sConvW[V_WEIGHT_BASE + w * HEAD_DIM + tidx] = r_w4[w] + if warp_idx < P1_QKG_WARPS and warp_idx % P1_NUM_JOBS == 0: + for i in range(VEC_SIZE): + cute.autovec_copy(w_q[(head_off + i * 32 + in_warp_tid, None)], r_w4) + for w in range(KERNEL_WIDTH): + r_wq[w * VEC_SIZE + i] = r_w4[w] + + # The 128-bit copy atom maps each thread to coalesced K vectors. + gState = h0[(slot, i_hv, None, None)] + gStateTiles = cute.local_tile(gState, (TILE_V, TILE_K), (None, 0)) + thr_state_copy = state_g2s_copy.get_slice(tidx) + for i_v in cutlass.range_constexpr(STATE_STAGES): + _issue_state_tile( + state_g2s_copy, thr_state_copy, gStateTiles, sState, i_v, STATE_STAGES + ) + + cute.arch.griddepcontrol_wait() + + bos = cu_seqlens[i_n] + eos = cu_seqlens[i_n + 1] + n_tok = eos - bos + scratch_row = intermediate_state_indices[i_n] + r_exp_A = cutlass.Float32(0.0) + + cute.arch.barrier() + + if warp_idx < P1_QKG_WARPS: + if p1_job == 2: + r_exp_A = cute.math.exp(cutlass.Float32(A_log[i_hv]), fastmath=True) + # Warp starting at token p1_par needs its window advanced that + # many steps; the g path is pointwise and needs none. + if p1_job == 0: + for _pi in cutlass.range(cutlass.min(p1_par, n_tok)): + for i in range(VEC_SIZE): + _xn = cutlass.Float32(x_q[0, bos + _pi, i_hv, i * 32 + in_warp_tid]) + r_state[0 * VEC_SIZE + i] = r_state[1 * VEC_SIZE + i] + r_state[1 * VEC_SIZE + i] = r_state[2 * VEC_SIZE + i] + r_state[2 * VEC_SIZE + i] = _xn + elif p1_job == 1: + for _pi in cutlass.range(cutlass.min(p1_par, n_tok)): + for i in range(VEC_SIZE): + _xn = cutlass.Float32(x_k[0, bos + _pi, i_hv, i * 32 + in_warp_tid]) + r_state[(KERNEL_WIDTH - 1) * VEC_SIZE + 0 * VEC_SIZE + i] = r_state[ + (KERNEL_WIDTH - 1) * VEC_SIZE + 1 * VEC_SIZE + i + ] + r_state[(KERNEL_WIDTH - 1) * VEC_SIZE + 1 * VEC_SIZE + i] = r_state[ + (KERNEL_WIDTH - 1) * VEC_SIZE + 2 * VEC_SIZE + i + ] + r_state[(KERNEL_WIDTH - 1) * VEC_SIZE + 2 * VEC_SIZE + i] = _xn + for i_t in cutlass.range(p1_par, T_LOOP, P1_JOB_WARPS): + token = bos + cutlass.min(i_t, n_tok - 1) + if p1_job == 0: + for i_pair in range(VEC_SIZE // 2): + i0 = i_pair * 2 + i1 = i_pair * 2 + 1 + k_idx0 = i0 * 32 + in_warp_tid + k_idx1 = i1 * 32 + in_warp_tid + r_conv_0 = 0.0 + r_conv_1 = 0.0 + for w in range(KERNEL_WIDTH - 1): + r_conv_0 += r_state[w * VEC_SIZE + i0] * r_wq[w * VEC_SIZE + i0] + r_conv_1 += r_state[w * VEC_SIZE + i1] * r_wq[w * VEC_SIZE + i1] + r_xq_0 = cutlass.Float32(x_q[0, token, i_hv, k_idx0]) + r_xq_1 = cutlass.Float32(x_q[0, token, i_hv, k_idx1]) + _cwq_last_0 = r_wq[(KERNEL_WIDTH - 1) * VEC_SIZE + i0] + _cwq_last_1 = r_wq[(KERNEL_WIDTH - 1) * VEC_SIZE + i1] + r_conv_0 += r_xq_0 * _cwq_last_0 + r_conv_1 += r_xq_1 * _cwq_last_1 + e0 = cute.math.exp(-r_conv_0, fastmath=True) + e1 = cute.math.exp(-r_conv_1, fastmath=True) + sig_0 = cute.arch.rcp_approx(cutlass.Float32(1.0) + e0) + sig_1 = cute.arch.rcp_approx(cutlass.Float32(1.0) + e1) + r_q[i0] = r_conv_0 * sig_0 + r_q[i1] = r_conv_1 * sig_1 + r_state[0 * VEC_SIZE + i0] = r_state[1 * VEC_SIZE + i0] + r_state[0 * VEC_SIZE + i1] = r_state[1 * VEC_SIZE + i1] + r_state[1 * VEC_SIZE + i0] = r_state[2 * VEC_SIZE + i0] + r_state[1 * VEC_SIZE + i1] = r_state[2 * VEC_SIZE + i1] + r_state[2 * VEC_SIZE + i0] = r_xq_0 + r_state[2 * VEC_SIZE + i1] = r_xq_1 + sum_q = 0.0 + for i in range(VEC_SIZE): + sum_q += r_q[i] * r_q[i] + for offset in [16, 8, 4, 2, 1]: + sum_q += cute.arch.shuffle_sync_bfly( + sum_q, offset=offset, mask=-1, mask_and_clamp=31 + ) + rnorm_q_scaled = cute.math.rsqrt(sum_q + 1e-06, fastmath=True) * scale + for i in range(VEC_SIZE): + r_q[i] = r_q[i] * rnorm_q_scaled + for i in range(VEC_SIZE): + k_idx = i * WARP_SIZE + in_warp_tid + qk_grp = k_idx % P2_LANES_K + qk_j = k_idx // P2_LANES_K + sQ[ + i_t, + qk_grp, + qk_j // 4, + qk_j % 4, + ] = r_q[i] + elif p1_job == 1: + r_b_raw = cutlass.Float32(0.0) + if in_warp_tid == 0: + r_b_raw = cutlass.Float32(beta[0, token, i_hv]) + for i in range(VEC_SIZE): + k_idx = i * 32 + in_warp_tid + r_conv = ( + r_state[(KERNEL_WIDTH - 1) * VEC_SIZE + 0 * VEC_SIZE + i] + * sConvW[0 * HEAD_DIM + i * 32 + in_warp_tid] + ) + r_conv += ( + r_state[(KERNEL_WIDTH - 1) * VEC_SIZE + 1 * VEC_SIZE + i] + * sConvW[1 * HEAD_DIM + i * 32 + in_warp_tid] + ) + r_conv += ( + r_state[(KERNEL_WIDTH - 1) * VEC_SIZE + 2 * VEC_SIZE + i] + * sConvW[2 * HEAD_DIM + i * 32 + in_warp_tid] + ) + r_xk = cutlass.Float32(x_k[0, token, i_hv, k_idx]) + r_conv += ( + r_xk + * sConvW[(KERNEL_WIDTH - 1) * HEAD_DIM + i * 32 + in_warp_tid] + ) + r_conv = r_conv * cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-r_conv, fastmath=True) + ) + r_k[i] = r_conv + if cutlass.const_expr(CACHE_RING): + ring_rawk[slot, i_hv, i_t, k_idx] = cutlass.BFloat16(r_conv) + r_state[(KERNEL_WIDTH - 1) * VEC_SIZE + 0 * VEC_SIZE + i] = r_state[ + (KERNEL_WIDTH - 1) * VEC_SIZE + 1 * VEC_SIZE + i + ] + r_state[(KERNEL_WIDTH - 1) * VEC_SIZE + 1 * VEC_SIZE + i] = r_state[ + (KERNEL_WIDTH - 1) * VEC_SIZE + 2 * VEC_SIZE + i + ] + r_state[(KERNEL_WIDTH - 1) * VEC_SIZE + 2 * VEC_SIZE + i] = r_xk + sum_k = 0.0 + for i in range(VEC_SIZE): + sum_k += r_k[i] * r_k[i] + for offset in [16, 8, 4, 2, 1]: + sum_k += cute.arch.shuffle_sync_bfly( + sum_k, offset=offset, mask=-1, mask_and_clamp=31 + ) + rnorm_k = cute.math.rsqrt(sum_k + 1e-06, fastmath=True) + for i in range(VEC_SIZE): + r_k[i] = r_k[i] * rnorm_k + for i in range(VEC_SIZE): + k_idx = i * WARP_SIZE + in_warp_tid + qk_grp = k_idx % P2_LANES_K + qk_j = k_idx // P2_LANES_K + sK[ + i_t, + qk_grp, + qk_j // 4, + qk_j % 4, + ] = r_k[i] + if in_warp_tid == 0: + sBeta[i_t] = cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-r_b_raw, fastmath=True) + ) + if cutlass.const_expr(CACHE_RING): + ring_beta[slot, i_hv, i_t] = sBeta[i_t] + else: + for i in range(VEC_SIZE): + k_idx = i * 32 + in_warp_tid + r_g_raw = cutlass.Float32(g[0, token, i_hv, k_idx]) + r_g_raw = r_g_raw + cutlass.Float32( + dt_bias[i_hv * HEAD_DIM + k_idx] + ) + exp_A_x = r_exp_A * r_g_raw + sigmoid_val = cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-exp_A_x, fastmath=True) + ) + r_gk = lower_bound * sigmoid_val + qk_grp = k_idx % P2_LANES_K + qk_j = k_idx // P2_LANES_K + sG[ + i_t, + qk_grp, + qk_j // 4, + qk_j % 4, + ] = cute.math.exp(r_gk, fastmath=True) + if cutlass.const_expr(CACHE_RING): + ring_g[slot, i_hv, i_t, k_idx] = r_gk + if p1_job == 0: + for i in range(VEC_SIZE): + k_idx = i * 32 + in_warp_tid + for w in range(KERNEL_WIDTH - 1): + intermediate_conv_q[scratch_row, i_t, head_off + k_idx, w] = ( + cutlass.BFloat16(r_state[w * VEC_SIZE + i]) + ) + elif p1_job == 1: + for i in range(VEC_SIZE): + k_idx = i * 32 + in_warp_tid + for w in range(KERNEL_WIDTH - 1): + intermediate_conv_k[scratch_row, i_t, head_off + k_idx, w] = ( + cutlass.BFloat16( + r_state[ + (KERNEL_WIDTH - 1) * VEC_SIZE + w * VEC_SIZE + i + ] + ) + ) + # Reach token i_t + P1_JOB_WARPS. The conv body already + # advanced one step; clamp the rest so the final iteration's + # unread advance stays in bounds at the last request. + if p1_job == 0: + for _a in range(P1_JOB_WARPS - 1): + _nx = bos + cutlass.min(i_t + 1 + _a, T_LOOP - 1) + for i in range(VEC_SIZE): + _xn = cutlass.Float32(x_q[0, _nx, i_hv, i * 32 + in_warp_tid]) + r_state[0 * VEC_SIZE + i] = r_state[1 * VEC_SIZE + i] + r_state[1 * VEC_SIZE + i] = r_state[2 * VEC_SIZE + i] + r_state[2 * VEC_SIZE + i] = _xn + elif p1_job == 1: + for _a in range(P1_JOB_WARPS - 1): + _nx = bos + cutlass.min(i_t + 1 + _a, T_LOOP - 1) + for i in range(VEC_SIZE): + _xn = cutlass.Float32(x_k[0, _nx, i_hv, i * 32 + in_warp_tid]) + r_state[(KERNEL_WIDTH - 1) * VEC_SIZE + 0 * VEC_SIZE + i] = ( + r_state[(KERNEL_WIDTH - 1) * VEC_SIZE + 1 * VEC_SIZE + i] + ) + r_state[(KERNEL_WIDTH - 1) * VEC_SIZE + 1 * VEC_SIZE + i] = ( + r_state[(KERNEL_WIDTH - 1) * VEC_SIZE + 2 * VEC_SIZE + i] + ) + r_state[(KERNEL_WIDTH - 1) * VEC_SIZE + 2 * VEC_SIZE + i] = _xn + else: + for _c in range(V_CH_PER_THREAD): + _v_idx = (tidx - P1_QKG_WARPS * 32) + _c * (HEAD_DIM // V_CH_PER_THREAD) + _csv0 = cutlass.Float32(cs_v[slot, head_off + _v_idx, 0]) + _csv1 = cutlass.Float32(cs_v[slot, head_off + _v_idx, 1]) + _csv2 = cutlass.Float32(cs_v[slot, head_off + _v_idx, 2]) + _wv = [ + sConvW[V_WEIGHT_BASE + w * HEAD_DIM + _v_idx] + for w in range(KERNEL_WIDTH) + ] + # Sliding conv window, oldest -> newest. + _win = [_csv0, _csv1, _csv2] + for _t in cutlass.range_constexpr(T_LOOP): + _win.append( + cutlass.Float32( + x_v[0, bos + cutlass.min(_t, n_tok - 1), i_hv, _v_idx] + ) + ) + for _t in cutlass.range_constexpr(T_LOOP): + _vconv = _win[_t] * _wv[0] + for _w in cutlass.range_constexpr(1, KERNEL_WIDTH): + _vconv += _win[_t + _w] * _wv[_w] + _vconv = _vconv * cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-_vconv, fastmath=True) + ) + sVall[_t * HEAD_DIM + _v_idx] = _vconv + if cutlass.const_expr(CACHE_RING): + ring_rawv[slot, i_hv, _t, _v_idx] = cutlass.BFloat16(_vconv) + for _w in cutlass.range_constexpr(KERNEL_WIDTH - 1): + intermediate_conv_v[scratch_row, _t, head_off + _v_idx, _w] = ( + cutlass.BFloat16(_win[_t + 1 + _w]) + ) + + k_grp = in_warp_tid % P2_LANES_K + row_grp = in_warp_tid // P2_LANES_K + + staged = cutlass.const_expr(STATE_STAGES < NUM_V_TILES) + if cutlass.const_expr(not staged): + cute.arch.cp_async_wait_group(0) + cute.arch.barrier() + r_beta_val = cutlass.Float32(0.0) + for i_p in cutlass.range_constexpr(NUM_V_TILES // TILES_PER_PASS): + for i_l in cutlass.range_constexpr(TILES_PER_PASS): + i_v = i_p * TILES_PER_PASS + i_l + if cutlass.const_expr(staged): + cute.arch.cp_async_wait_group( + min(STATE_STAGES + i_v, NUM_V_TILES) - 1 - i_v + ) + cute.arch.barrier() + for b in range(P2_BATCHES): + _st = (i_l * P2_BATCHES + b) * P2_VEC + _row = warp_idx * NUM_V_ROWS + b * P2_ROWS_LANE + row_grp + for j in range(P2_VEC): + r_state[_st + j] = cutlass.Float32( + sState[ + _row, + j * P2_LANES_K + k_grp, + i_v % STATE_STAGES, + ] + ) + if cutlass.const_expr(staged) and i_v + STATE_STAGES < NUM_V_TILES: + cute.arch.barrier() + _issue_state_tile( + state_g2s_copy, + thr_state_copy, + gStateTiles, + sState, + i_v + STATE_STAGES, + STATE_STAGES, + ) + if cutlass.const_expr(i_p == NUM_V_TILES // TILES_PER_PASS - 1): + # Every global read this block makes has now landed in registers. + cute.arch.griddepcontrol_launch_dependents() + + for i_t in cutlass.range(T_LOOP): + if cutlass.const_expr(STREAM_STATE): + # ns0 uses the same q/k/decay for every state tile. Load it on + # the first pass and keep it live across the remaining three. + if cutlass.const_expr(i_p == 0): + r_beta_val = sBeta[i_t] + # Distinct staging tensors avoid making the three vector + # loads one dependent chain. + for jv in range(P2_VEC // 4): + cute.autovec_copy(sQ[(i_t, k_grp, jv, None)], r_v4q) + cute.autovec_copy(sK[(i_t, k_grp, jv, None)], r_v4k) + cute.autovec_copy(sG[(i_t, k_grp, jv, None)], r_v4g) + for c in range(4): + r_q[jv * 4 + c] = r_v4q[c] + r_k[jv * 4 + c] = r_v4k[c] + r_decay[jv * 4 + c] = r_v4g[c] + else: + # Read once per token, and the resident state already claims + # every spare register, so don't pay for vector staging. + r_beta_val = sBeta[i_t] + for j in range(P2_VEC): + r_q[j] = sQ[i_t, k_grp, j // 4, j % 4] + r_k[j] = sK[i_t, k_grp, j // 4, j % 4] + r_decay[j] = sG[i_t, k_grp, j // 4, j % 4] + for i_l in range(TILES_PER_PASS): + v_base = (i_p * TILES_PER_PASS + i_l) * TILE_V + for b in range(P2_BATCHES): + _st = (i_l * P2_BATCHES + b) * P2_VEC + v_row = warp_idx * NUM_V_ROWS + b * P2_ROWS_LANE + row_grp + # Every lane of a group wants the same v; smem broadcasts it. + r_v = sVall[i_t * HEAD_DIM + v_base + v_row] + shk_1 = 0.0 + shk_2 = 0.0 + for jp in range(P2_VEC // 2): + _p = jp * 2 + r_state[_st + _p], r_state[_st + _p + 1] = ( + cute.arch.mul_packed_f32x2( + (r_decay[_p], r_decay[_p + 1]), + (r_state[_st + _p], r_state[_st + _p + 1]), + ) + ) + shk_1, shk_2 = cute.arch.fma_packed_f32x2( + src_a=(r_state[_st + _p], r_state[_st + _p + 1]), + src_b=(r_k[_p], r_k[_p + 1]), + src_c=(shk_1, shk_2), + ) + shk = shk_1 + shk_2 + for offset in [16, 8, 4, 2, 1]: + if cutlass.const_expr(offset < P2_LANES_K): + shk += cute.arch.shuffle_sync_bfly( + shk, offset=offset, mask=-1, mask_and_clamp=31 + ) + vnb = (r_v - shk) * r_beta_val + shq_1 = 0.0 + shq_2 = 0.0 + for jp in range(P2_VEC // 2): + _p = jp * 2 + r_state[_st + _p], r_state[_st + _p + 1] = ( + cute.arch.fma_packed_f32x2( + src_a=(vnb, vnb), + src_b=(r_k[_p], r_k[_p + 1]), + src_c=(r_state[_st + _p], r_state[_st + _p + 1]), + ) + ) + shq_1, shq_2 = cute.arch.fma_packed_f32x2( + src_a=(r_state[_st + _p], r_state[_st + _p + 1]), + src_b=(r_q[_p], r_q[_p + 1]), + src_c=(shq_1, shq_2), + ) + shq = shq_1 + shq_2 + for offset in [16, 8, 4, 2, 1]: + if cutlass.const_expr(offset < P2_LANES_K): + shq += cute.arch.shuffle_sync_bfly( + shq, offset=offset, mask=-1, mask_and_clamp=31 + ) + if k_grp == 0 and i_t < n_tok: + if cutlass.const_expr(APPLY_ONORM): + sOall[i_t * HEAD_DIM + v_base + v_row] = shq + else: + o[0, bos + i_t, i_hv, v_base + v_row] = cutlass.BFloat16( + shq + ) + if cutlass.const_expr(not CACHE_RING): + for b in range(P2_BATCHES): + _st = (i_l * P2_BATCHES + b) * P2_VEC + v_row = warp_idx * NUM_V_ROWS + b * P2_ROWS_LANE + row_grp + for j in range(P2_VEC): + ht[ + scratch_row, + i_t, + i_hv, + v_base + v_row, + j * P2_LANES_K + k_grp, + ] = r_state[_st + j] + + if cutlass.const_expr(APPLY_ONORM): + cute.arch.barrier() + for i_t in cutlass.range(warp_idx, T_LOOP, NUM_WARPS): + sumsq = cutlass.Float32(0.0) + for i in range(VEC_SIZE): + _o = sOall[i_t * HEAD_DIM + i * 32 + in_warp_tid] + sumsq += _o * _o + for offset in [16, 8, 4, 2, 1]: + sumsq += cute.arch.shuffle_sync_bfly( + sumsq, offset=offset, mask=-1, mask_and_clamp=31 + ) + rms = cute.math.rsqrt( + sumsq / cutlass.Float32(HEAD_DIM) + onorm_eps, fastmath=True + ) + for i in range(VEC_SIZE): + v_idx = i * 32 + in_warp_tid + raw_o = sOall[i_t * HEAD_DIM + v_idx] + _tok = bos + cutlass.min(i_t, n_tok - 1) + gate_raw = cutlass.Float32(onorm_g[0, _tok, i_hv, v_idx]) + gate = cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-gate_raw, fastmath=True) + ) + if i_t < n_tok: + o[0, bos + i_t, i_hv, v_idx] = cutlass.BFloat16( + raw_o * rms * cutlass.Float32(onorm_weight[v_idx]) * gate + ) + + +@cute.jit +def _run_kda_decode_mtp_dspark( + h0: cute.Tensor, + x_q: cute.Tensor, + x_k: cute.Tensor, + x_v: cute.Tensor, + w_q: cute.Tensor, + w_k: cute.Tensor, + w_v: cute.Tensor, + cs_q: cute.Tensor, + cs_k: cute.Tensor, + cs_v: cute.Tensor, + A_log: cute.Tensor, + g: cute.Tensor, + dt_bias: cute.Tensor, + beta: cute.Tensor, + out: cute.Tensor, + intermediate_ssm: cute.Tensor, + intermediate_state_indices: cute.Tensor, + intermediate_conv_q: cute.Tensor, + intermediate_conv_k: cute.Tensor, + intermediate_conv_v: cute.Tensor, + ssm_state_indices: cute.Tensor, + cu_seqlens: cute.Tensor, + ring_rawv: cute.Tensor, + ring_rawk: cute.Tensor, + ring_g: cute.Tensor, + ring_beta: cute.Tensor, + onorm_g: cute.Tensor, + onorm_weight: cute.Tensor, + scale: cutlass.Constexpr[float], + H: cutlass.Constexpr[int], + N: cutlass.Constexpr[int], + NUM_SPEC: cutlass.Constexpr[int], + BLOCK_THREADS: cutlass.Constexpr[int], + lower_bound: cutlass.Constexpr[float], + CACHE_RING: cutlass.Constexpr[bool], + APPLY_ONORM: cutlass.Constexpr[bool], + onorm_eps: cutlass.Constexpr[float], + stream: cuda.CUstream, +): + """Launch the fixed Kimi-K3/DSpARK bonus + NUM_SPEC-draft specialization.""" + p2_lanes = _p2_lanes(block_threads=BLOCK_THREADS, num_spec=NUM_SPEC) + p2_vec = TILE_K // p2_lanes + # (token, lane row, vector, element within vector); see _qk_smem. + qk_stride, qk_elems = _qk_smem(block_threads=BLOCK_THREADS, num_spec=NUM_SPEC) + smem_qk_layout = cute.make_layout( + (1 + NUM_SPEC, p2_lanes, p2_vec // 4, 4), stride=qk_stride + ) + # Padding by the subgroup width shifts concurrent v-rows onto disjoint bank + # groups: +8 tiles four eight-bank windows; +16 tiles two half-warps. + smem_state_stride = TILE_K + p2_lanes + TILE_V = _state_tile_v(block_threads=BLOCK_THREADS, num_spec=NUM_SPEC) + state_stages = min(NUM_STATE_STAGES, TILE_K // TILE_V) + state_stage_elems = TILE_V * smem_state_stride + state_smem_elems = state_stages * state_stage_elems + smem_state_layout = cute.make_layout( + (TILE_V, TILE_K, state_stages), + stride=(smem_state_stride, 1, state_stage_elems), + ) + state_cache_mode = ( + cpasync.LoadCacheMode.ALWAYS + if NUM_SPEC == 0 and BLOCK_THREADS == BLOCK_THREADS_WIDE + else cpasync.LoadCacheMode.GLOBAL + ) + state_copy_atom = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=state_cache_mode), + cutlass.Float32, + num_bits_per_copy=128, + ) + state_g2s_copy = cute.make_tiled_copy_tv( + state_copy_atom, + thr_layout=cute.make_layout( + (TILE_V, BLOCK_THREADS // TILE_V), + stride=(BLOCK_THREADS // TILE_V, 1), + ), + val_layout=cute.make_layout((1, 4)), + ) + t_loop = 1 + NUM_SPEC + smem_bytes = ( + # sQ, sK, sG, sBeta, sVall, and sConvW. + (3 * t_loop * qk_elems + t_loop + t_loop * TILE_K + 8 * TILE_K) * 4 + # State stages and the raw output tile for normalization. + + state_smem_elems * 4 + + (t_loop * TILE_K * 4 if cutlass.const_expr(APPLY_ONORM) else 0) + + 256 + ) + kda_decode_mtp_kernel( + state_g2s_copy, + h0, + x_q, + x_k, + x_v, + w_q, + w_k, + w_v, + cs_q, + cs_k, + cs_v, + A_log, + g, + dt_bias, + beta, + out, + intermediate_ssm, + intermediate_state_indices, + intermediate_conv_q, + intermediate_conv_k, + intermediate_conv_v, + ring_rawv, + ring_rawk, + ring_g, + ring_beta, + onorm_g, + onorm_weight, + smem_qk_layout, + smem_state_layout, + ssm_state_indices, + cu_seqlens, + scale, + NUM_SPEC, + BLOCK_THREADS, + p2_lanes, + lower_bound, + CACHE_RING, + APPLY_ONORM, + onorm_eps, + ).launch( + grid=(H, N, 1), + block=[BLOCK_THREADS, 1, 1], + smem=smem_bytes, + stream=stream, + use_pdl=True, + ) + + +def _block_threads(*, H: int, N: int) -> int: + """Pick the block width for this grid. + + The grid is (H, N). Below one wave every SM gets a single block whatever + its width, so the wide block is free warps and roughly halves the exposed + latency of the serial token chain. At or above one wave the narrow block + wins instead, because two resident blocks let one issue while the other + waits at a barrier -- warps inside a single block all wait together. + + Measured on GB300 (152 SMs, H=12, num_spec=5): N=1 -8.8%, N=8 -11.3%, + but N=64 +47.3%. + """ + import torch + + num_sms = torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + return BLOCK_THREADS_WIDE if H * N <= num_sms else BLOCK_THREADS_NARROW + + +_DSPARK_COMPILED = {} + + +def _tensor_layout_key(tensor): + return (tensor.device, tensor.dtype, tuple(tensor.shape), tuple(tensor.stride())) + + +def _fits_32bit_stride(tensor): + max_offset = int(tensor.storage_offset()) + for size, stride in zip(tensor.shape, tensor.stride()): + if abs(int(stride)) > 2**31 - 1: + return False + if size: + max_offset += (int(size) - 1) * abs(int(stride)) + if max_offset > 2**31 - 1: + return False + return True + + +def _cute_tensor(tensor, *, dynamic=True): + from cutlass.cute.runtime import from_dlpack + + if tensor.requires_grad: + tensor = tensor.detach() + value = from_dlpack( + tensor, + assumed_align=16, + use_32bit_stride=_fits_32bit_stride(tensor), + ) + leading_dim = next( + (dim for dim, stride in enumerate(tensor.stride()) if stride == 1), None + ) + if not dynamic or leading_dim is None: + return value + return value.mark_layout_dynamic(leading_dim) + + +def fused_kda_decode_mtp_dspark( + *, + x_q, + x_k, + x_v, + w_q, + w_k, + w_v, + cs_q, + cs_k, + cs_v, + g, + beta, + A_log, + dt_bias, + recurrent_state, + intermediate_ssm, + intermediate_state_indices, + intermediate_conv_q, + intermediate_conv_k, + intermediate_conv_v, + ssm_state_indices, + cu_seqlens, + lower_bound, + scale=None, + replayssm_rawv=None, + replayssm_rawk=None, + replayssm_g=None, + replayssm_beta=None, + onorm_gate=None, + onorm_weight=None, + onorm_eps=None, +): + """Run Kimi-K3 KDA verify while preserving DSpARK rollback semantics. + + Persistent recurrent/conv states are read-only. Every post-token state and + convolution window is written to DSpARK's existing intermediate buffers. + The caller is responsible for enforcing the fixed dense token contract: + every request contributes exactly 1 + num_spec tokens (num_spec == + --speculative-dspark-block-size), inferred here from T // N - 1. + + ReplaySSM: passing the four replayssm_* rings switches the kernel to + CACHE_RING mode — per-step raw inputs go to the rings (consumed by the + commit-time exact fold, see kda_replayssm_spec_decode.py) and the per-step + intermediate_ssm state snapshots are skipped, so intermediate_ssm may be + None. + + Passing all three onorm_* arguments fuses gated RMSNorm into the recurrence + kernel. + """ + import torch + + H = x_q.shape[2] + N = cu_seqlens.numel() - 1 + T = x_q.shape[1] + expected_shape = (1, T, H, TILE_K) + if tuple(x_q.shape) != expected_shape or tuple(x_k.shape) != expected_shape: + raise ValueError(f"expected q/k shape {expected_shape}") + if tuple(x_v.shape) != expected_shape or tuple(g.shape) != expected_shape: + raise ValueError(f"expected v/g shape {expected_shape}") + if tuple(beta.shape) != (1, T, H): + raise ValueError(f"expected beta shape {(1, T, H)}") + # T // N == 1 is num_spec == 0: one token per request, i.e. a plain decode + # step. The backend never dispatches here for it (that is the dedicated + # decode kernel's job), but the layout is legal and benchmarks compare the + # two at this point, so the wrapper accepts it. + if N <= 0 or T % N != 0 or T // N < 1: + raise ValueError( + f"DSpARK KDA MTP requires a fixed 1 + num_spec dense tokens per " + f"request; got T={T}, N={N}" + ) + num_spec = T // N - 1 + if recurrent_state.shape[1:] != (H, TILE_K, TILE_K): + raise ValueError("expected recurrent state layout [pool, H, V=128, K=128]") + if ( + recurrent_state.dtype != torch.float32 + or tuple(recurrent_state.stride()[-3:]) != (TILE_K * TILE_K, TILE_K, 1) + or recurrent_state.stride(0) % 4 != 0 + or recurrent_state.storage_offset() % 4 != 0 + ): + raise ValueError( + "cp.async recurrent state requires fp32 contiguous [H, V, K] " + "inner layout and 16-byte-aligned slot offsets" + ) + rings = (replayssm_rawv, replayssm_rawk, replayssm_g, replayssm_beta) + cache_ring = all(ring is not None for ring in rings) + if any(ring is not None for ring in rings) and not cache_ring: + raise ValueError("ReplaySSM requires all four replayssm_* rings") + if cache_ring: + ring_len = replayssm_rawv.shape[2] + if ( + ring_len < 1 + num_spec + or replayssm_rawv.shape[1:] != (H, ring_len, TILE_K) + or replayssm_rawk.shape[1:] != (H, ring_len, TILE_K) + or replayssm_g.shape[1:] != (H, ring_len, TILE_K) + or replayssm_beta.shape[1:] != (H, ring_len) + ): + raise ValueError( + f"expected ReplaySSM ring layouts [slots, H={H}, L>={1 + num_spec}, " + f"{TILE_K}] / [slots, H, L]" + ) + if ( + replayssm_rawv.dtype != torch.bfloat16 + or replayssm_rawk.dtype != torch.bfloat16 + or replayssm_g.dtype != torch.float32 + or replayssm_beta.dtype != torch.float32 + ): + raise ValueError( + "expected ReplaySSM ring dtypes rawv/rawk=bf16, g/beta=fp32" + ) + if intermediate_ssm is None: + # Dead-branch placeholder: CACHE_RING skips every ht snapshot + # write, but CuTe still type-checks the rank-5 indexing. + intermediate_ssm = recurrent_state.unsqueeze(1) + if not cache_ring and ( + intermediate_ssm.shape[1] < 1 + num_spec + or intermediate_ssm.shape[2:5] != (H, TILE_K, TILE_K) + ): + raise ValueError( + f"expected intermediate SSM layout [scratch, >={1 + num_spec}, H, " + f"V=128, K=128]" + ) + if scale is None: + scale = TILE_K**-0.5 + onorm_args = (onorm_gate, onorm_weight, onorm_eps) + apply_onorm = all(value is not None for value in onorm_args) + if any(value is not None for value in onorm_args) and not apply_onorm: + raise ValueError( + "fused output norm requires onorm_gate, onorm_weight, and onorm_eps" + ) + if apply_onorm: + if tuple(onorm_gate.shape) != expected_shape: + raise ValueError(f"expected output-norm gate shape {expected_shape}") + if tuple(onorm_weight.shape) != (TILE_K,): + raise ValueError(f"expected output-norm weight shape {(TILE_K,)}") + if onorm_gate.dtype != torch.bfloat16 or onorm_weight.dtype != torch.float32: + raise ValueError("expected output-norm gate=bf16 and weight=fp32") + block_threads = _block_threads(H=H, N=N) + out = torch.empty_like(x_v) + args = ( + recurrent_state, + x_q, + x_k, + x_v, + w_q, + w_k, + w_v, + cs_q, + cs_k, + cs_v, + A_log, + g, + dt_bias, + beta, + out, + intermediate_ssm, + intermediate_state_indices, + intermediate_conv_q, + intermediate_conv_k, + intermediate_conv_v, + ssm_state_indices, + cu_seqlens, + # ReplaySSM rings; same placeholder trick when CACHE_RING is off. + replayssm_rawv if cache_ring else intermediate_conv_q, + replayssm_rawk if cache_ring else intermediate_conv_q, + replayssm_g if cache_ring else intermediate_conv_q, + replayssm_beta if cache_ring else intermediate_conv_q[..., 0], + onorm_gate if apply_onorm else x_v, + onorm_weight if apply_onorm else dt_bias, + ) + key = ( + H, + N, + num_spec, + block_threads, + cache_ring, + apply_onorm, + float(onorm_eps) if apply_onorm else 0.0, + float(scale), + float(lower_bound), + *(_tensor_layout_key(tensor) for tensor in args), + ) + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + compiled = _DSPARK_COMPILED.get(key) + if compiled is None: + cute_args = tuple(_cute_tensor(tensor, dynamic=False) for tensor in args) + compiled = cute.compile( + _run_kda_decode_mtp_dspark, + *cute_args, + scale=float(scale), + H=H, + N=N, + NUM_SPEC=num_spec, + BLOCK_THREADS=block_threads, + lower_bound=float(lower_bound), + CACHE_RING=cache_ring, + APPLY_ONORM=apply_onorm, + onorm_eps=float(onorm_eps) if apply_onorm else 0.0, + stream=stream, + ) + _DSPARK_COMPILED[key] = compiled + compiled( + *(_cute_tensor(tensor) for tensor in args), + stream, + ) + return out diff --git a/python/sglang/kernels/ops/kimi_k3/mla_output_gate.py b/python/sglang/kernels/ops/kimi_k3/mla_output_gate.py new file mode 100644 index 000000000..60a64f899 --- /dev/null +++ b/python/sglang/kernels/ops/kimi_k3/mla_output_gate.py @@ -0,0 +1,51 @@ +"""CUDA JIT K3 MLA output gate: out = x * sigmoid(gate) in one kernel.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from sglang.kernels.jit.utils import ( + cache_once, + is_arch_support_pdl, + load_jit, + make_cpp_args, +) + +if TYPE_CHECKING: + from tvm_ffi.module import Module + +_THREADS: int = 256 + + +@cache_once +def _jit_mla_output_gate_module() -> Module: + args = make_cpp_args(_THREADS, is_arch_support_pdl()) + return load_jit( + "kimi_k3_mla_output_gate_" + str(_THREADS), + *args, + cuda_files=["kimi_k3/mla_output_gate.cuh"], + cuda_wrappers=[("run", f"MlaOutputGateKernel<{args}>::run")], + extra_cuda_cflags=["-O3"], + ) + + +def covered(x: torch.Tensor, gate: torch.Tensor) -> bool: + return ( + x.dtype == torch.bfloat16 + and gate.dtype == torch.bfloat16 + and x.shape == gate.shape + and x.is_contiguous() + and gate.is_contiguous() + and x.numel() % 8 == 0 + and x.numel() > 0 + ) + + +def kimi_k3_mla_output_gate(x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + """out = bf16(x * bf16(sigmoid(gate))); double rounding matches the + unfused torch.sigmoid + mul pair bit-for-bit. Caller checks covered().""" + out = torch.empty_like(x) + _jit_mla_output_gate_module().run(x.view(-1), gate.view(-1), out.view(-1)) + return out diff --git a/python/sglang/kernels/ops/kimi_k3/moe.py b/python/sglang/kernels/ops/kimi_k3/moe.py new file mode 100644 index 000000000..fb2de4e0f --- /dev/null +++ b/python/sglang/kernels/ops/kimi_k3/moe.py @@ -0,0 +1,61 @@ +from __future__ import annotations + +import torch + +from sglang.kernels.jit.utils import ( + cache_once, + is_arch_support_pdl, + load_jit, + make_cpp_args, +) + + +def _make_name(*args): + return "kimi_k3_" + "_".join(str(a) for a in args) + + +@cache_once +def _jit_situ_mul_quant_varlen_module( + quant_group_size: int, + scale_ue8m0: bool, + swizzle: bool, +): + args = make_cpp_args( + quant_group_size, + scale_ue8m0, + swizzle, + is_arch_support_pdl(), + ) + return load_jit( + _make_name("situ_mul_quant_varlen"), + *args, + cuda_files=["kimi_k3/situ_and_mul.cuh"], + cuda_wrappers=[("run", f"SituAndMulMaskedPostQuantKernel<{args}>::run")], + extra_cuda_cflags=["-use_fast_math"], + ) + + +def situ_and_mul_masked_post_quant( + input: torch.Tensor, + output: torch.Tensor, + output_scale: torch.Tensor, + quant_group_size: int, + masked_m: torch.Tensor, + beta: float, + linear_beta: float, + scale_ue8m0: bool = False, + topk: int = 8, + transposed: bool = False, + swizzle: bool = False, +) -> None: + module = _jit_situ_mul_quant_varlen_module(quant_group_size, scale_ue8m0, swizzle) + module.run( + input, + output, + output_scale, + masked_m, + topk, + transposed, + float(beta), + float(linear_beta), + ) diff --git a/python/sglang/kernels/ops/kimi_k3/sp_collective.py b/python/sglang/kernels/ops/kimi_k3/sp_collective.py new file mode 100644 index 000000000..d3a1bcd8d --- /dev/null +++ b/python/sglang/kernels/ops/kimi_k3/sp_collective.py @@ -0,0 +1,338 @@ +"""K3 SP-MoE bf16 reduce-scatter and all-gather over MNNVL push memory.""" + +from __future__ import annotations + +import json +import os +from typing import TYPE_CHECKING, NamedTuple, Optional + +import torch + +from sglang.kernels.jit.utils import ( + cache_once, + is_arch_support_pdl, + load_jit, + make_cpp_args, +) +from sglang.srt.utils.custom_op import register_custom_op + +if TYPE_CHECKING: + from tvm_ffi.module import Module + + from sglang.kernels.ops.communication.all_reduce import Communicator + + +class Tuning(NamedTuple): + num_blocks: int + block_size: int + + +class Dispatch(NamedTuple): + strategy: str + tuning: Tuning + + +class FusionDispatch(NamedTuple): + strategy: str + max_blocks: int + + +# Safe seed values. GB300 production values are loaded by the layer glue from +# the checked-in JSON after the sweep settles. +DEFAULT_TUNING = Tuning(num_blocks=16, block_size=256) +_CONFIG_DIR = os.path.join(os.path.dirname(__file__), "configs", "sp_collective") +_TABLES: dict[str, Optional[dict]] = {} + + +def _device_name(device: torch.device) -> str: + return torch.cuda.get_device_name(device).replace(" ", "_").replace("/", "_") + + +def _table(world_size: int, hidden_size: int, device: torch.device) -> Optional[dict]: + path = os.path.join( + _CONFIG_DIR, + ( + f"world={world_size},H={hidden_size}," + f"device_name={_device_name(device)}.json" + ), + ) + if path not in _TABLES: + if os.path.exists(path): + with open(path) as f: + _TABLES[path] = json.load(f) + else: + _TABLES[path] = None + return _TABLES[path] + + +def get_dispatch( + kind: str, + world_size: int, + hidden_size: int, + num_tokens: int, + device: torch.device, +) -> Optional[Dispatch]: + """Return the tuned strategy, or None when the table selects NCCL.""" + table = _table(world_size, hidden_size, device) + if table is None: + return None + raw_configs = table["configs"].get(kind) + if raw_configs is None: + return None + configs = {int(k): v for k, v in raw_configs.items()} + bucket = min(configs) + for candidate in sorted(configs): + if candidate <= num_tokens: + bucket = candidate + else: + break + config = configs[bucket] + if config["strategy"] == "nccl": + return None + return Dispatch( + config["strategy"], + Tuning(config["num_blocks"], config["block_size"]), + ) + + +def get_tuning( + kind: str, + world_size: int, + hidden_size: int, + num_tokens: int, + device: torch.device, +) -> Optional[Tuning]: + """Compatibility helper for callers that only support staging push.""" + dispatch = get_dispatch(kind, world_size, hidden_size, num_tokens, device) + if dispatch is None or dispatch.strategy != "push": + return None + return dispatch.tuning + + +def get_fusion_dispatch( + kind: str, + world_size: int, + hidden_size: int, + num_tokens: int, + device: torch.device, +) -> Optional[FusionDispatch]: + """Return a measured fused strategy, or None for the separate path.""" + table = _table(world_size, hidden_size, device) + if table is None: + return None + raw_configs = table["configs"].get(kind) + if raw_configs is None: + return None + configs = {int(k): v for k, v in raw_configs.items()} + bucket = min(configs) + for candidate in sorted(configs): + if candidate <= num_tokens: + bucket = candidate + else: + break + config = configs[bucket] + if config["strategy"] == "separate": + return None + return FusionDispatch(config["strategy"], config["max_blocks"]) + + +@cache_once +def _jit_module(world_size: int) -> Module: + args = make_cpp_args(world_size, is_arch_support_pdl()) + cls = f"SPCollectiveKernel<{args}>" + return load_jit( + "kimi_k3_sp_collective", + *args, + cuda_files=["kimi_k3/comm/sp_collective.cuh"], + cuda_wrappers=[ + ("reduce_scatter_res", f"{cls}::reduce_scatter_res"), + ("reduce_scatter_pull", f"{cls}::reduce_scatter_pull"), + ("all_gather", f"{cls}::all_gather"), + ("all_gather_direct", f"{cls}::all_gather_direct"), + ], + extra_cuda_cflags=["-O3"], + ) + + +_COMM_MAP: dict[int, Communicator] = {} +_PULL_SEM_MC_MAP: dict[int, int] = {} + + +def register_comm(comm: Communicator, *, pull_sem_mc_ptr: int = 0) -> None: + # One communicator per world_size per process -- see the note in + # kimi_k3/all_reduce.py::register_comm. The ops key only on world_size, so an + # overwrite here would hand the old group's callers the new group's peer + # pointers. + prev = _COMM_MAP.get(comm.world_size) + assert prev is None or prev is comm, ( + f"a different communicator is already registered for world_size=" + f"{comm.world_size}" + ) + _COMM_MAP[comm.world_size] = comm + _PULL_SEM_MC_MAP[comm.world_size] = pull_sem_mc_ptr + + +@register_custom_op(mutates_args=["output"]) +def _reduce_scatter_res_op( + world_size: int, + input: torch.Tensor, + output: torch.Tensor, + residual: Optional[torch.Tensor], + residual_is_local: bool, + num_blocks: int, + block_size: int, +) -> None: + _jit_module(world_size).reduce_scatter_res( + _COMM_MAP[world_size], + input.view(-1), + output.view(-1), + None if residual is None else residual.view(-1), + residual_is_local, + num_blocks, + block_size, + ) + + +@register_custom_op(mutates_args=["output"]) +def _reduce_scatter_pull_op( + world_size: int, + input: torch.Tensor, + output: torch.Tensor, + residual: Optional[torch.Tensor], + residual_is_local: bool, + input_mc_ptr: int, + num_blocks: int, + block_size: int, +) -> None: + _jit_module(world_size).reduce_scatter_pull( + _COMM_MAP[world_size], + input.view(-1), + output.view(-1), + None if residual is None else residual.view(-1), + residual_is_local, + input_mc_ptr, + _PULL_SEM_MC_MAP[world_size], + num_blocks, + block_size, + ) + + +@register_custom_op(mutates_args=["output"]) +def _all_gather_op( + world_size: int, + input: torch.Tensor, + output: torch.Tensor, + ws_mc_base: int, + num_blocks: int, + block_size: int, +) -> None: + _jit_module(world_size).all_gather( + _COMM_MAP[world_size], + input.view(-1), + output.view(-1), + ws_mc_base, + num_blocks, + block_size, + ) + + +@register_custom_op(mutates_args=["output"]) +def _all_gather_direct_op( + world_size: int, + input: torch.Tensor, + output: torch.Tensor, + output_mc_ptr: int, + num_blocks: int, + block_size: int, +) -> None: + _jit_module(world_size).all_gather_direct( + _COMM_MAP[world_size], + input.view(-1), + output.view(-1), + output_mc_ptr, + _PULL_SEM_MC_MAP[world_size], + num_blocks, + block_size, + ) + + +def reduce_scatter_res( + world_size: int, + input: torch.Tensor, + output: torch.Tensor, + residual: Optional[torch.Tensor] = None, + *, + tuning: Tuning = DEFAULT_TUNING, +) -> torch.Tensor: + residual_is_local = residual is not None and residual.numel() == output.numel() + _reduce_scatter_res_op( + world_size, + input, + output, + residual, + residual_is_local, + tuning.num_blocks, + tuning.block_size, + ) + return output + + +def reduce_scatter_pull( + world_size: int, + input: torch.Tensor, + output: torch.Tensor, + residual: Optional[torch.Tensor] = None, + *, + input_mc_ptr: int, + tuning: Tuning = DEFAULT_TUNING, +) -> torch.Tensor: + residual_is_local = residual is not None and residual.numel() == output.numel() + _reduce_scatter_pull_op( + world_size, + input, + output, + residual, + residual_is_local, + input_mc_ptr, + tuning.num_blocks, + tuning.block_size, + ) + return output + + +def all_gather( + world_size: int, + input: torch.Tensor, + output: torch.Tensor, + *, + ws_mc_base: int, + tuning: Tuning = DEFAULT_TUNING, +) -> torch.Tensor: + _all_gather_op( + world_size, + input, + output, + ws_mc_base, + tuning.num_blocks, + tuning.block_size, + ) + return output + + +def all_gather_direct( + world_size: int, + input: torch.Tensor, + output: torch.Tensor, + *, + output_mc_ptr: int, + tuning: Tuning = DEFAULT_TUNING, +) -> torch.Tensor: + _all_gather_direct_op( + world_size, + input, + output, + output_mc_ptr, + tuning.num_blocks, + tuning.block_size, + ) + return output diff --git a/python/sglang/kernels/ops/mamba/mamba_state_scatter_triton.py b/python/sglang/kernels/ops/mamba/mamba_state_scatter_triton.py index 4a19da4c9..5c050cf6b 100644 --- a/python/sglang/kernels/ops/mamba/mamba_state_scatter_triton.py +++ b/python/sglang/kernels/ops/mamba/mamba_state_scatter_triton.py @@ -501,3 +501,101 @@ def scatter_mamba_states_after_mtp_verify( mamba_track_indices, mamba_steps_to_track, ) + + +@triton.jit +def track_mamba_states_all_layers_kernel( + conv_states_ptr, # [num_layers, pool_size, ...] full conv pool + ssm_states_ptr, # [num_layers, pool_size, ...] full ssm pool + cache_indices_ptr, + mamba_track_mask_ptr, + mamba_track_indices_ptr, + conv_layer_stride, + conv_row_stride, + ssm_layer_stride, + ssm_row_stride, + batch_size, + conv_state_numel_per_row: tl.constexpr, + ssm_state_numel_per_row: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + check_freed_slots: tl.constexpr, +): + """All-layers variant of track_mamba_state_if_needed_kernel: one launch + covers every mamba layer (grid = num_layers * batch_size) instead of one + launch per layer. The track mask / source / destination indices are shared + across layers, so a single launch at the end of the step is equivalent to + the per-layer launches (each layer's state is final by then).""" + pid = tl.program_id(0) + layer_idx = (pid // batch_size).to(tl.int64) + batch_idx = pid % batch_size + + track_mask = tl.load(mamba_track_mask_ptr + batch_idx) + if not track_mask: + return + + src_idx = tl.load(cache_indices_ptr + batch_idx).to(tl.int64) + dst_idx = tl.load(mamba_track_indices_ptr + batch_idx).to(tl.int64) + if check_freed_slots: + if src_idx < 0 or dst_idx < 0: + return + + conv_base = conv_states_ptr + layer_idx * conv_layer_stride + for offset in range(0, conv_state_numel_per_row, BLOCK_SIZE): + element_indices = offset + tl.arange(0, BLOCK_SIZE) + mask = element_indices < conv_state_numel_per_row + data = tl.load( + conv_base + src_idx * conv_row_stride + element_indices, + mask=mask, + other=0.0, + ) + tl.store( + conv_base + dst_idx * conv_row_stride + element_indices, data, mask=mask + ) + + ssm_base = ssm_states_ptr + layer_idx * ssm_layer_stride + for offset in range(0, ssm_state_numel_per_row, BLOCK_SIZE): + element_indices = offset + tl.arange(0, BLOCK_SIZE) + mask = element_indices < ssm_state_numel_per_row + data = tl.load( + ssm_base + src_idx * ssm_row_stride + element_indices, mask=mask, other=0.0 + ) + tl.store(ssm_base + dst_idx * ssm_row_stride + element_indices, data, mask=mask) + + +def track_mamba_states_all_layers( + conv_states_pool: torch.Tensor, + ssm_states_pool: torch.Tensor, + cache_indices: torch.Tensor, + mamba_track_mask: torch.Tensor, + mamba_track_indices: torch.Tensor, + batch_size: int, + check_freed_slots: bool = False, +): + """Track conv/ssm states for ALL mamba layers in one launch. + + conv_states_pool / ssm_states_pool are the full [num_layers, pool_size, + ...] pools; per-row copy semantics are identical to + track_mamba_states_if_needed applied per layer. + """ + num_layers = conv_states_pool.shape[0] + conv_state_numel_per_row = conv_states_pool[0, 0].numel() + ssm_state_numel_per_row = ssm_states_pool[0, 0].numel() + + BLOCK_SIZE = 1024 + grid = (num_layers * batch_size,) + track_mamba_states_all_layers_kernel[grid]( + conv_states_pool, + ssm_states_pool, + cache_indices, + mamba_track_mask, + mamba_track_indices, + conv_states_pool.stride(0), + conv_states_pool.stride(1), + ssm_states_pool.stride(0), + ssm_states_pool.stride(1), + batch_size, + conv_state_numel_per_row, + ssm_state_numel_per_row, + BLOCK_SIZE, + check_freed_slots, + ) diff --git a/python/sglang/kernels/ops/memory/gpu_tensor_hash.py b/python/sglang/kernels/ops/memory/gpu_tensor_hash.py index b6444f218..0622928e5 100644 --- a/python/sglang/kernels/ops/memory/gpu_tensor_hash.py +++ b/python/sglang/kernels/ops/memory/gpu_tensor_hash.py @@ -40,7 +40,7 @@ def _fmix32(x, C1: tl.constexpr, C2: tl.constexpr): return x -@triton.jit +@triton.jit(do_not_specialize=["n_u32", "seed1", "seed2"]) def hash_tiles32_kernel_blocked( in_ptr, out_ptr, @@ -100,7 +100,7 @@ def hash_tiles32_kernel_blocked( tl.store(out_ptr + pid, out) -@triton.jit +@triton.jit(do_not_specialize=["n_elems"]) def add_tree_reduce_u64_kernel(in_ptr, out_ptr, n_elems, CHUNK: tl.constexpr): pid = tl.program_id(axis=0) start = pid * CHUNK diff --git a/python/sglang/kernels/ops/mm/__init__.py b/python/sglang/kernels/ops/mm/__init__.py new file mode 100644 index 000000000..794777ad9 --- /dev/null +++ b/python/sglang/kernels/ops/mm/__init__.py @@ -0,0 +1,3 @@ +"""Multimodal kernels.""" + +__all__ = ["process"] diff --git a/python/sglang/kernels/ops/mm/process/__init__.py b/python/sglang/kernels/ops/mm/process/__init__.py new file mode 100644 index 000000000..c51e1b785 --- /dev/null +++ b/python/sglang/kernels/ops/mm/process/__init__.py @@ -0,0 +1,5 @@ +"""Multimodal input-processing kernels.""" + +from sglang.kernels.ops.mm.process.image import normalize_and_patchify + +__all__ = ["normalize_and_patchify"] diff --git a/python/sglang/kernels/ops/mm/process/image.py b/python/sglang/kernels/ops/mm/process/image.py new file mode 100644 index 000000000..a673ca2ba --- /dev/null +++ b/python/sglang/kernels/ops/mm/process/image.py @@ -0,0 +1,124 @@ +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + +_MAX_TRITON_ELEMENTS = 2**31 - 1 + + +@triton.jit +def _normalize_and_patchify_kernel( + input_ptr, + scale_ptr, + bias_ptr, + output_ptr, + channels: tl.constexpr, + input_height, + input_width, + grid_height, + grid_width, + patch_size: tl.constexpr, + element_count, + block_size: tl.constexpr, +): + offsets = tl.program_id(0) * block_size + tl.arange(0, block_size) + mask = offsets < element_count + + patch_area = patch_size * patch_size + patch_offset = offsets % patch_area + patch_x = patch_offset % patch_size + patch_y = patch_offset // patch_size + remaining = offsets // patch_area + channel = remaining % channels + remaining = remaining // channels + grid_x = remaining % grid_width + remaining = remaining // grid_width + grid_y = remaining % grid_height + batch = remaining // grid_height + + input_y = grid_y * patch_size + patch_y + input_x = grid_x * patch_size + patch_x + input_mask = mask & (input_y < input_height) & (input_x < input_width) + input_offset = ( + (batch * channels + channel) * input_height + input_y + ) * input_width + input_x + value = tl.load(input_ptr + input_offset, mask=input_mask, other=0.0) + scale = tl.load(scale_ptr + channel, mask=mask) + bias = tl.load(bias_ptr + channel, mask=mask) + tl.store(output_ptr + offsets, value * scale + bias, mask=mask) + + +def _normalize_and_patchify_torch( + image: torch.Tensor, + image_scale: torch.Tensor, + image_bias: torch.Tensor, + patch_size: int, + padded_height: int, + padded_width: int, +) -> torch.Tensor: + pad_height = padded_height - image.shape[-2] + pad_width = padded_width - image.shape[-1] + if pad_height > 0 or pad_width > 0: + image = F.pad(image, (0, pad_width, 0, pad_height), value=0.0) + image = torch.addcmul(image_bias, image, image_scale) + batch, channels, height, width = image.shape + grid_height = height // patch_size + grid_width = width // patch_size + image = image.view(batch, channels, grid_height, patch_size, grid_width, patch_size) + return image.permute(0, 2, 4, 1, 3, 5).reshape( + batch, -1, channels, patch_size, patch_size + ) + + +def normalize_and_patchify( + image: torch.Tensor, + image_scale: torch.Tensor, + image_bias: torch.Tensor, + patch_size: int, + padded_height: int, + padded_width: int, +) -> torch.Tensor: + batch, channels, _, _ = image.shape + grid_height = padded_height // patch_size + grid_width = padded_width // patch_size + element_count = ( + batch * grid_height * grid_width * channels * patch_size * patch_size + ) + if ( + not image.is_cuda + or torch.version.hip is not None + or not image.is_contiguous() + or not image_scale.is_contiguous() + or not image_bias.is_contiguous() + or element_count > _MAX_TRITON_ELEMENTS + ): + return _normalize_and_patchify_torch( + image, + image_scale, + image_bias, + patch_size, + padded_height, + padded_width, + ) + + output = torch.empty( + (batch, grid_height * grid_width, channels, patch_size, patch_size), + dtype=image.dtype, + device=image.device, + ) + block_size = 256 + _normalize_and_patchify_kernel[(triton.cdiv(element_count, block_size),)]( + image, + image_scale, + image_bias, + output, + channels, + image.shape[-2], + image.shape[-1], + grid_height, + grid_width, + patch_size, + element_count, + block_size, + ) + return output diff --git a/python/sglang/kernels/ops/moe/configs/moe_front/epilogue,E=896,topk=16,device_name=NVIDIA_GB300.json b/python/sglang/kernels/ops/moe/configs/moe_front/epilogue,E=896,topk=16,device_name=NVIDIA_GB300.json new file mode 100644 index 000000000..e46333544 --- /dev/null +++ b/python/sglang/kernels/ops/moe/configs/moe_front/epilogue,E=896,topk=16,device_name=NVIDIA_GB300.json @@ -0,0 +1,154 @@ +{ + "meta": { + "device": "NVIDIA GB300", + "device_name": "NVIDIA_GB300", + "torch": "2.11.0+cu130", + "experts": 896, + "latent": 3584, + "topk": 16, + "inner": 50, + "reps": 20 + }, + "configs": { + "1": { + "block_size": 448, + "cast_vec": 8, + "cast_first": false + }, + "2": { + "block_size": 448, + "cast_vec": 8, + "cast_first": false + }, + "3": { + "block_size": 448, + "cast_vec": 8, + "cast_first": false + }, + "4": { + "block_size": 448, + "cast_vec": 8, + "cast_first": false + }, + "6": { + "block_size": 448, + "cast_vec": 8, + "cast_first": true + }, + "8": { + "block_size": 448, + "cast_vec": 8, + "cast_first": false + }, + "12": { + "block_size": 448, + "cast_vec": 8, + "cast_first": false + }, + "16": { + "block_size": 448, + "cast_vec": 8, + "cast_first": false + }, + "24": { + "block_size": 448, + "cast_vec": 8, + "cast_first": false + }, + "32": { + "block_size": 448, + "cast_vec": 8, + "cast_first": false + }, + "48": { + "block_size": 448, + "cast_vec": 8, + "cast_first": false + }, + "64": { + "block_size": 448, + "cast_vec": 8, + "cast_first": false + }, + "96": { + "block_size": 448, + "cast_vec": 8, + "cast_first": false + }, + "128": { + "block_size": 448, + "cast_vec": 8, + "cast_first": false + }, + "192": { + "block_size": 448, + "cast_vec": 8, + "cast_first": false + }, + "256": { + "block_size": 448, + "cast_vec": 8, + "cast_first": false + }, + "384": { + "block_size": 448, + "cast_vec": 8, + "cast_first": false + }, + "512": { + "block_size": 224, + "cast_vec": 8, + "cast_first": false + }, + "768": { + "block_size": 224, + "cast_vec": 8, + "cast_first": false + }, + "1024": { + "block_size": 224, + "cast_vec": 8, + "cast_first": false + }, + "1536": { + "block_size": 224, + "cast_vec": 8, + "cast_first": false + }, + "2048": { + "block_size": 224, + "cast_vec": 8, + "cast_first": false + }, + "3072": { + "block_size": 224, + "cast_vec": 8, + "cast_first": false + }, + "4096": { + "block_size": 224, + "cast_vec": 8, + "cast_first": false + }, + "6144": { + "block_size": 224, + "cast_vec": 8, + "cast_first": false + }, + "8192": { + "block_size": 224, + "cast_vec": 8, + "cast_first": false + }, + "12288": { + "block_size": 224, + "cast_vec": 8, + "cast_first": false + }, + "16384": { + "block_size": 224, + "cast_vec": 8, + "cast_first": false + } + } +} diff --git a/python/sglang/kernels/ops/moe/configs/moe_front/strategy,E=896,topk=16,device_name=NVIDIA_GB300.json b/python/sglang/kernels/ops/moe/configs/moe_front/strategy,E=896,topk=16,device_name=NVIDIA_GB300.json new file mode 100644 index 000000000..e21a56d65 --- /dev/null +++ b/python/sglang/kernels/ops/moe/configs/moe_front/strategy,E=896,topk=16,device_name=NVIDIA_GB300.json @@ -0,0 +1,37 @@ +{ + "meta": { + "device": "NVIDIA GB300", + "hidden": 7168, + "experts": 896, + "latent": 3584, + "topk": 16, + "routing_dtype": "fp32", + "selection": "exact token count; overlap must win all 3 repeats and median speedup must exceed 1%" + }, + "configs": { + "736": "overlap", + "768": "overlap", + "896": "overlap", + "960": "overlap", + "1280": "overlap", + "1536": "overlap", + "1792": "overlap", + "2048": "overlap", + "2304": "overlap", + "2560": "overlap", + "2816": "overlap", + "3328": "overlap", + "3584": "overlap", + "3840": "overlap", + "4352": "overlap", + "4608": "overlap", + "4864": "overlap", + "5632": "overlap", + "5888": "overlap", + "6144": "overlap", + "7424": "overlap", + "8448": "overlap", + "8704": "overlap", + "8960": "overlap" + } +} diff --git a/python/sglang/kernels/ops/moe/moe_front.py b/python/sglang/kernels/ops/moe/moe_front.py new file mode 100644 index 000000000..d943d4a61 --- /dev/null +++ b/python/sglang/kernels/ops/moe/moe_front.py @@ -0,0 +1,235 @@ +"""K3 MoE front: merged gate + routed_expert_down_proj GEMM, and the fp32 router. + +The unfused MoE front -- the path every EP-a2a / WideEP deployment takes -- runs +three ops over the same `hidden_states [T, 7168]`: + + router_logits = gate(hidden_states) # [896, 7168] 12.85 MB + topk_output = topk(hidden_states, router_logits) + routed_input = routed_expert_down_proj(hidden_states) # [3584, 7168] 51.4 MB + +Plain [M, 896] fp32 logits go to route_radix. This module covers what that +cannot: the merged front. + +**fused_front** -- the two GEMMs share their input, so their weights are merged +and one cuBLAS GEMM emits `[T, 896 + 3584]` fp32; a single epilogue kernel then +runs the top-k on the gate slice and casts the latent slice to bf16. Routing +stays bit-identical to the fp32 path and routed_input comes out dense. + +**overlap** -- the original fp32 gate + top-k run on the model's side stream +while the latent down-projection runs on the main stream. The streams join +before expert dispatch, and the side stream is then reused by the existing +shared-expert overlap. + +Measured in-graph on a GB300, us per MoE layer: + + T 512 768 1024 1280 2048 2560 4096 8192 16384 + baseline 31.1 37.6 50.9 54.3 91.5 99.3 152.7 304.7 629.3 + merged 26.1 37.3 46.7 - - - - - - + overlap 29.0 33.7 46.5 52.4 90.5 96.4 151.2 307.1 637.8 + +The fastest strategy is non-monotonic because cuBLAS changes GEMM algorithms at +specific row counts. A GB300 JSON table therefore opts exact, repeatedly +measured token counts into overlap; unmeasured shapes keep the conservative +defaults (merged through 1024, unfused above it). This captures the clear 768 +and 2560 wins without regressing 512, 8192, or 16384. + +A bf16-output merged GEMM was measured too (fastest at T=1, 11.8 us). It is not +used: bf16 rounds the router logits and moves the selected expert set on 2-25% of +rows depending on T, for ~1.2 us -- and from T>=8 it loses to the fp32 variant +anyway, because its routed_input is a strided slice a dense-input runner must copy. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional, Tuple + +import torch + +from sglang.kernels.jit.utils import ( + cache_once, + is_arch_support_pdl, + load_jit, + make_cpp_args, +) + +if TYPE_CHECKING: + from tvm_ffi.module import Module + +import json +import os + +NUM_EXPERTS = 896 +TOPK = 16 + +# Above this token count the merged GEMM stops paying by default; see the table +# above. A device strategy table may override individual, measured token counts +# with the dual-stream overlap. +MERGED_FRONT_MAX_TOKENS = 1024 + +_CONFIG_DIR = os.path.join(os.path.dirname(__file__), "configs", "moe_front") + +# Kernel tunables, per token bucket, from the JSON table. +# block_size threads per CTA; sets experts-per-thread in the radix select +# (896 / block_size). 224 -> 4, 448 -> 2. +# cast_vec fp32 elements each thread converts per step in the latent cast. +# cast_first issue the cast before the select (loads in flight during the +# radix rounds) or after it. +# Fallbacks when no tuned table matches the device. Both were the sweep's most +# common winners: cast_vec 8 is 32 B/thread, the Blackwell vector-load limit, and +# it won at every one of the 28 token counts measured; issuing the cast after the +# select beat issuing it before almost everywhere. +DEFAULT_EPILOGUE_CONFIG = {"block_size": 224, "cast_vec": 8, "cast_first": False} + +_tables = {} + + +def _table(kind: str, device_name: str): + path = os.path.join( + _CONFIG_DIR, + f"{kind},E={NUM_EXPERTS},topk={TOPK},device_name={device_name}.json", + ) + if path not in _tables: + table = None + if os.path.exists(path): + with open(path) as f: + table = {int(k): v for k, v in json.load(f)["configs"].items()} + _tables[path] = table + return _tables[path] + + +def _device_name(device) -> str: + return torch.cuda.get_device_name(device).replace(" ", "_").replace("/", "_") + + +def get_config(kind: str, num_tokens: int, device, default: dict) -> dict: + """Tuned config for the nearest token bucket at or below `num_tokens`.""" + table = _table(kind, _device_name(device)) + if not table: + return dict(default) + pick = min(table) + for k in sorted(table): + if k <= num_tokens: + pick = k + else: + break + return dict(table[pick]) + + +def get_front_strategy(num_tokens: int, device) -> str: + """Return the measured front strategy for this exact workload. + + GEMM algorithm changes make the merged/overlap crossover non-monotonic in + M. Therefore strategy tables are exact-match only. Unmeasured shapes keep + the robust defaults: merged fp32 through 1024 tokens, then unfused. + """ + table = _table("strategy", _device_name(device)) + if table is not None and num_tokens in table: + return str(table[num_tokens]) + return "merged_fp32" if num_tokens <= MERGED_FRONT_MAX_TOKENS else "unfused" + + +@cache_once +def _jit_module() -> Module: + args = make_cpp_args(is_arch_support_pdl()) + return load_jit( + "moe_front", + *args, + cuda_files=["moe/route_radix.cuh"], + cuda_wrappers=[ + ("front_epilogue", f"FusedFrontEpilogueKernel<{args}>::run"), + ], + # No fast-math: scoring and expert-id selection must stay comparable to + # route_radix / the Triton router under ties and NaN. + extra_cuda_cflags=["-O3"], + ) + + +@cache_once +def available() -> bool: + import logging + + try: + _jit_module() + return True + except Exception as e: # pragma: no cover - toolchain dependent + logging.getLogger(__name__).warning( + f"Failed to load the JIT MoE front kernels: {e}" + ) + return False + + +# merged front: [gate | down] GEMM -> top-k + routed_input + + +def fused_front_covered( + hidden_states: torch.Tensor, + merged_weight: torch.Tensor, + bias: Optional[torch.Tensor], + topk: int, + latent: int, +) -> bool: + """[T<=MERGED_FRONT_MAX_TOKENS, 7168] bf16 x [896 + latent, 7168] bf16, fp32 + bias, top-16, latent a multiple of 4.""" + return ( + hidden_states.dim() == 2 + and merged_weight.dim() == 2 + and hidden_states.dtype == torch.bfloat16 + and merged_weight.dtype == torch.bfloat16 + and bias is not None + and bias.dtype == torch.float32 + and bias.numel() == NUM_EXPERTS + and int(topk) == TOPK + and merged_weight.shape[0] == NUM_EXPERTS + latent + and merged_weight.shape[1] == hidden_states.shape[1] + and latent % 4 == 0 + and 0 < hidden_states.shape[0] <= MERGED_FRONT_MAX_TOKENS + and hidden_states.stride(1) == 1 + and merged_weight.stride(1) == 1 + ) + + +def fused_front( + hidden_states: torch.Tensor, + merged_weight: torch.Tensor, + correction_bias: torch.Tensor, + latent: int, + topk: int = TOPK, + renormalize: bool = True, + routed_scaling_factor: float = 1.0, + apply_routed_scaling_factor_on_output: bool = False, + config: Optional[dict] = None, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Merged front GEMM + fused top-k/cast epilogue. + + Returns ``(topk_weights [M, topk] fp32, topk_ids [M, topk] int32, + routed_input [M, latent] bf16)``. + """ + M = hidden_states.shape[0] + device = hidden_states.device + if config is None: + config = get_config("epilogue", M, device, DEFAULT_EPILOGUE_CONFIG) + + # fp32 out keeps routing exact; the extra output traffic versus bf16 is + # M x (896 + latent) x 2 bytes, negligible against the 64 MB weight read at + # the sizes this path serves. + merged = torch.mm(hidden_states, merged_weight.t(), out_dtype=torch.float32) + + weights = torch.empty((M, topk), dtype=torch.float32, device=device) + ids = torch.empty((M, topk), dtype=torch.int32, device=device) + routed = torch.empty((M, latent), dtype=torch.bfloat16, device=device) + + _jit_module().front_epilogue( + merged, + correction_bias, + weights, + ids, + routed, + topk, + float(routed_scaling_factor if routed_scaling_factor is not None else 1.0), + bool(renormalize), + bool(apply_routed_scaling_factor_on_output), + int(config["block_size"]), + int(config["cast_vec"]), + bool(config["cast_first"]), + ) + return weights, ids, routed diff --git a/python/sglang/kernels/ops/moe/moe_fused_gate.py b/python/sglang/kernels/ops/moe/moe_fused_gate.py index 922d5e929..112f501d0 100644 --- a/python/sglang/kernels/ops/moe/moe_fused_gate.py +++ b/python/sglang/kernels/ops/moe/moe_fused_gate.py @@ -9,11 +9,11 @@ import triton.language as tl from sglang.kernel_api_logging import debug_kernel_api from sglang.kernels.jit.utils import cache_once, is_arch_support_pdl, load_jit +from sglang.kernels.ops.moe import moe_route_radix if TYPE_CHECKING: from tvm_ffi.module import Module - _SCORING_FUNC_MAP = { "sigmoid": 0, "sqrtsoftplus": 1, @@ -290,6 +290,30 @@ def moe_fused_gate( if routed_scaling_factor is None: routed_scaling_factor = 1.0 + # K3 radix-select fast path: native-CUDA radix-select replaces the 16 + # dependent argmax rounds (single CTA per token; ids bit-identical to this + # triton kernel incl. ties). + # The radix kernel keeps keys register-resident and returns winners in + # expert-id order (skipping the biased-descending sort; downstream MoE + # kernels are order-insensitive). It is 3.1-3.5x faster than the Triton + # kernel at [1..8192, 896] top-16 on B200. + if ( + scoring_func.lower() == "sigmoid" + and num_fused_shared_experts == 0 + and num_expert_group <= 1 + and moe_softcapping == 0.0 + ): + radix_args = ( + scores, + bias, + topk, + renormalize, + routed_scaling_factor, + apply_routed_scaling_factor_on_output, + ) + if moe_route_radix.covered(scores, bias, topk): + return moe_route_radix.route_radix(*radix_args, sorted=False) + M, N = scores.shape K = topk K_routed = topk - num_fused_shared_experts @@ -309,7 +333,10 @@ def moe_fused_gate( # stay occupancy-bound. Swept on H100/B200: this beats the AOT kernels across # shapes, whereas larger tiles / more warps regress (register pressure). BLOCK_M = max(1, min(4, 256 // BLOCK_N)) - num_warps = 1 + # For wide rows (e.g. Kimi K3: 896 experts, BLOCK_N 1024) the K sequential + # argmax passes dominate and benefit from more warps despite the + # cross-warp reduction cost. + num_warps = 1 if BLOCK_N <= 512 else 4 grid = (triton.cdiv(M, BLOCK_M),) use_pdl = is_arch_support_pdl() extra = {"launch_pdl": True} if use_pdl else {} diff --git a/python/sglang/kernels/ops/moe/moe_route_quant_fused.py b/python/sglang/kernels/ops/moe/moe_route_quant_fused.py new file mode 100644 index 000000000..4320db5fd --- /dev/null +++ b/python/sglang/kernels/ops/moe/moe_route_quant_fused.py @@ -0,0 +1,125 @@ +"""Fused K3 MoE-front prep: radix routing + trtllm id pack + mxfp8 quant. + +One launch replaces the three tiny kernels between the K3 fused-front GEMM and +the trtllm-gen routed-MoE op at decode batch sizes (route_radix -> triton +(id<<16|bf16(w)) pack -> per_token_group_quant, ~7.5us busy + 2 extra launches +per MoE layer with the SMs near idle). CTAs [0, M) run the route_radix body +with the pack folded into its epilogue; CTAs [M, 2M) run the +per_token_group_quant math, one CTA per token row. Both halves reuse the +standalone kernels' device code, so ids/weights, packed ids, and quantized +activations are bit-identical to the unfused chain. + +Specialized like route_radix itself: 896 experts, top-16, bf16/fp32 scores, +and a 3584-wide bf16 activation row quantized to fp8 with row-major packed +UE8M0 group-32 scales (the trtllm-gen SiTU MoE input format). Wired into +serving through sglang.srt.layers.moe.route_quant_handoff. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Tuple + +import torch + +from sglang.kernels.jit.utils import ( + cache_once, + is_arch_support_pdl, + load_jit, + make_cpp_args, +) +from sglang.kernels.ops.moe import moe_route_radix + +if TYPE_CHECKING: + from tvm_ffi.module import Module + +_HIDDEN = 3584 +_GROUP_SIZE = 32 +_NUM_GROUPS = _HIDDEN // _GROUP_SIZE +# Fusion trades the flat quant grid for one 224-thread CTA per token; that (and +# the win itself, which is launch overhead) only makes sense at small decode +# batches. Above the cap the callers run the unfused chain. +_MAX_TOKENS = 64 + + +@cache_once +def _jit_module() -> Module: + args = make_cpp_args(is_arch_support_pdl()) + return load_jit( + "moe_route_quant_fused", + *args, + cuda_files=["moe/route_quant_fused.cuh"], + cuda_wrappers=[("run", f"RouteQuantFusedKernel<{args}>::run")], + # No fast-math: the routing half must stay bit-identical to + # route_radix (see its module comment); the quant half's math is + # fast-math-independent (explicit intrinsics + bit manipulation). + extra_cuda_cflags=["-O3"], + ) + + +@cache_once +def available() -> bool: + import logging + + try: + _jit_module() + return True + except Exception as e: # pragma: no cover - toolchain dependent + logging.getLogger(__name__).warning( + f"Failed to load the JIT fused route+quant kernel: {e}" + ) + return False + + +def covered( + scores: torch.Tensor, bias: torch.Tensor, topk: int, x: torch.Tensor +) -> bool: + """route_radix coverage plus the quant half: [M<=64, 3584] bf16 rows with + 32B-aligned starts (base and stride), same token count as the scores.""" + return ( + moe_route_radix.covered(scores, bias, topk) + and x.dim() == 2 + and x.shape[0] == scores.shape[0] + and 0 < x.shape[0] <= _MAX_TOKENS + and x.shape[1] == _HIDDEN + and x.dtype == torch.bfloat16 + and x.stride(1) == 1 + and x.data_ptr() % 32 == 0 + and (x.stride(0) * x.element_size()) % 32 == 0 + ) + + +def route_quant_fused( + scores: torch.Tensor, + bias: torch.Tensor, + x: torch.Tensor, + topk: int, + renormalize: bool, + routed_scaling_factor: float, + apply_scale: bool, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Returns ``(weights [M, topk] fp32, ids [M, topk] int32, packed [M, topk] + int32, x_q [M, 3584] fp8_e4m3, x_s [M, 28] int32 row-major packed UE8M0)``. + Caller must have checked covered(). Winners come out in expert-id-ascending + order (the standalone production dispatch's sorted=False).""" + M = scores.shape[0] + device = scores.device + out_w = torch.empty((M, topk), dtype=torch.float32, device=device) + out_i = torch.empty((M, topk), dtype=torch.int32, device=device) + out_packed = torch.empty((M, topk), dtype=torch.int32, device=device) + out_q = torch.empty((M, _HIDDEN), dtype=torch.float8_e4m3fn, device=device) + out_s = torch.empty((M, _NUM_GROUPS // 4), dtype=torch.int32, device=device) + _jit_module().run( + scores, + bias, + out_w, + out_i, + out_packed, + x, + out_q, + out_s, + topk, + float(routed_scaling_factor), + bool(renormalize), + bool(apply_scale), + ) + return out_w, out_i, out_packed, out_q, out_s diff --git a/python/sglang/kernels/ops/moe/moe_route_radix.py b/python/sglang/kernels/ops/moe/moe_route_radix.py new file mode 100644 index 000000000..078658f06 --- /dev/null +++ b/python/sglang/kernels/ops/moe/moe_route_radix.py @@ -0,0 +1,98 @@ +"""Native-CUDA radix-select router for K3 routing (all batch sizes). + +Keys and activations stay in registers (224 threads, 4 experts each), the +split-bin search runs on warp scans instead of cub, rounds exit early when the +top-k separates on a byte boundary, and the (biased desc, id asc) output sort +is optional. Consumers that only gather by expert id can pass sorted=False and +skip the epilogue rank-sort entirely. + +Dispatched automatically from moe_fused_gate for covered inputs; the production +dispatch uses sorted=False. It is 3.1-3.5x faster than the Triton router at +[1..8192, 896] top-16 on B200. Correctness coverage lives in +test_kimi_k3_prerequisite_ops.py, against a pure-torch fp32 oracle rather than the +Triton router: moe_fused_gate dispatches back here for every input this kernel +covers, so using it as the reference compares the kernel with itself. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Tuple + +import torch + +from sglang.kernels.jit.utils import ( + cache_once, + is_arch_support_pdl, + load_jit, + make_cpp_args, +) + +if TYPE_CHECKING: + from tvm_ffi.module import Module + +_NUM_EXPERTS = 896 +_TOPK = 16 + + +@cache_once +def _jit_route_radix_module() -> Module: + args = make_cpp_args(is_arch_support_pdl()) + return load_jit( + "moe_route_radix", + *args, + cuda_files=["moe/route_radix.cuh"], + cuda_wrappers=[("run", f"RouteRadixKernel<{args}>::run")], + # No fast-math: expert-id selection must stay bit-identical to the + # Triton router under ties/NaN. + extra_cuda_cflags=["-O3"], + ) + + +def covered(scores: torch.Tensor, bias: torch.Tensor, topk: int) -> bool: + """Specialized for K3 decode routing: [M, 896] bf16 or fp32 + row-contiguous scores (8B/16B-aligned rows), fp32 bias, top-16.""" + return ( + scores.dim() == 2 + and scores.size(1) == _NUM_EXPERTS + and int(topk) == _TOPK + and scores.dtype in (torch.bfloat16, torch.float32) + and bias.dtype == torch.float32 + and scores.stride(1) == 1 + and scores.stride(0) % 4 == 0 + and bias.is_contiguous() + ) + + +def route_radix( + scores: torch.Tensor, + bias: torch.Tensor, + topk: int, + renormalize: bool, + routed_scaling_factor: float, + apply_scale: bool, + sorted: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Returns (weights [M, topk] fp32, ids [M, topk] int32). Caller must have + checked covered(). + + Default sorted=False: winners come out in expert-id-ascending order + (downstream MoE kernels are order-insensitive) and the epilogue rank-sort + is skipped. sorted=True restores the Triton router's (biased desc, id asc) + output order. Either way the winner set matches Triton exactly; the renorm + sum is taken in the respective output order, so weights may differ by + <= ~1 ulp.""" + M = scores.shape[0] + out_w = torch.empty((M, topk), dtype=torch.float32, device=scores.device) + out_i = torch.empty((M, topk), dtype=torch.int32, device=scores.device) + _jit_route_radix_module().run( + scores, + bias, + out_w, + out_i, + topk, + float(routed_scaling_factor), + bool(renormalize), + bool(apply_scale), + bool(sorted), + ) + return out_w, out_i diff --git a/python/sglang/kernels/ops/moe/moe_topk_sum.py b/python/sglang/kernels/ops/moe/moe_topk_sum.py new file mode 100644 index 000000000..e7b229a7f --- /dev/null +++ b/python/sglang/kernels/ops/moe/moe_topk_sum.py @@ -0,0 +1,37 @@ +"""CUDA JIT top-k expert-output sum: out[M, K] = in[M, topk, K].sum(dim=1).""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from sglang.kernels.jit.utils import ( + cache_once, + is_arch_support_pdl, + load_jit, + make_cpp_args, +) + +if TYPE_CHECKING: + from tvm_ffi.module import Module + +_THREADS: int = 256 + + +@cache_once +def _jit_topk_sum_module() -> Module: + args = make_cpp_args(_THREADS, is_arch_support_pdl()) + return load_jit( + "moe_topk_sum_" + str(_THREADS), + *args, + cuda_files=["moe/topk_sum.cuh"], + cuda_wrappers=[("run", f"TopkSumKernel<{args}>::run")], + extra_cuda_cflags=["-O3", "--use_fast_math"], + ) + + +def moe_topk_sum(x: torch.Tensor, out: torch.Tensor) -> torch.Tensor: + """out[M, K] = x[M, topk, K].sum(dim=1) for contiguous bf16 tensors.""" + _jit_topk_sum_module().run(x, out) + return out diff --git a/python/sglang/kernels/ops/moe/pack_topk_ids.py b/python/sglang/kernels/ops/moe/pack_topk_ids.py index d9d40ee63..26a9a98a6 100644 --- a/python/sglang/kernels/ops/moe/pack_topk_ids.py +++ b/python/sglang/kernels/ops/moe/pack_topk_ids.py @@ -9,6 +9,8 @@ import torch import triton import triton.language as tl +from sglang.kernels.jit.utils import is_arch_support_pdl + class PackTopkIds: @@ -51,12 +53,16 @@ class PackTopkIds: BLOCK_SIZE = 1024 grid = (triton.cdiv(numel, BLOCK_SIZE),) + pdl_kwargs = ( + {"USE_PDL": True, "launch_pdl": True} if is_arch_support_pdl() else {} + ) _pack_topk_ids_triton_kernel[grid]( topk_ids, topk_weights, out, numel, BLOCK_SIZE=BLOCK_SIZE, + **pdl_kwargs, ) return out @@ -68,14 +74,21 @@ def _pack_topk_ids_triton_kernel( out_ptr, numel, BLOCK_SIZE: tl.constexpr, + USE_PDL: tl.constexpr = False, ): pid = tl.program_id(0) offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) mask = offsets < numel + if USE_PDL: + tl.extra.cuda.gdc_wait() + ids = tl.load(topk_ids_ptr + offsets, mask=mask, other=0) w = tl.load(topk_weights_ptr + offsets, mask=mask, other=0.0) + if USE_PDL: + tl.extra.cuda.gdc_launch_dependents() + w_bf16 = w.to(tl.bfloat16) w_i16 = w_bf16.to(tl.int16, bitcast=True) w_i32 = w_i16.to(tl.int32) & 0xFFFF diff --git a/python/sglang/kernels/ops/moe/trtllm_gen_moe.py b/python/sglang/kernels/ops/moe/trtllm_gen_moe.py new file mode 100644 index 000000000..c2da4402d --- /dev/null +++ b/python/sglang/kernels/ops/moe/trtllm_gen_moe.py @@ -0,0 +1,528 @@ +"""TRT-LLM-gen fused MoE (SiTU) compiled through the sglang JIT system. + +Builds the trtllm-gen fused-MoE host/runner sources with sglang's own +tvm-ffi ``load_jit`` from a **self-contained cubin pool** +(``SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL``): a downloadable directory holding + + * the prebuilt SiTU cubins (``local/``) + ``config.json`` + + ``flashinferMetaInfo.h``, + * the flat batched-gemm ABI headers (staged into a + ``trtllmGen_bmm_export/``-shaped include tree at build time), + * an ``overlay/`` with only the sources/headers that differ from the + public ``flashinfer`` pip package. + +Every unmodified source and the CUTLASS headers come from the installed +``flashinfer`` package's ``data/`` tree (the wheel ships it for its own +JIT), so running this backend needs exactly one download and one env var — +no extra source checkout. + +This module vendors only glue: + + * header staging: the pool ships the batched-gemm ABI headers flat; they + are copied into a content-addressed include tree shaped like + ``flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/``; + * JIT build of the 12 launcher/runner/routing sources with the private + ABI defines (``TLLM_GEN_LOCAL_CUBINS_ABI`` etc.); + * the ctypes cubin-loader callback (the .so asks for cubins by absolute + path + sha256; we read them from the pool); + * a thin ``trtllm_fp4_block_scale_moe`` wrapper (FromLogits routing, + ``do_finalize=True``); kernel tile config ("tactic") defaults to the + runner's built-in heuristic — pass an explicit one for tuned setups. + +Validated for the Kimi K3 decode/prefill MoE regime: MxFP4 weights with +bf16 (w4a16) or MxFP8 (w4a8) activations, ``ActivationType.Situ`` (SiTuGlu: +``a*tanh(g/a)*sigmoid(g) * b*tanh(u/b)``), DeepSeekV3/noaux_tc routing. +""" + +from __future__ import annotations + +import ctypes +import hashlib +import logging +import os +import pathlib +import shutil +from typing import TYPE_CHECKING, Optional, Sequence + +import torch + +from sglang.kernels.jit.utils import ( + cache_once, + get_jit_cuda_arch, + load_jit, + override_jit_cuda_arch, +) +from sglang.srt.environ import envs + +if TYPE_CHECKING: + from tvm_ffi.module import Module + +# ActivationType / RoutingMethodType values from trtllm-gen's tllm_enums +# (kept as plain ints here to avoid importing anything for them). +ACTIVATION_SITU = 9 +ROUTING_DEEPSEEK_V3 = 2 +_ROUTING_TOPK = 5 +_ROUTING_INPUT_FROM_LOGITS = 0 +# NOTE: the enum VALUES start at 0; the "Mode 1/2/3" wording in upstream +# comments is documentation numbering, not the enum value. +_ROUTING_INPUT_PACKED = 1 + +# Batched-gemm ABI headers shipped flat in the cubin pool; the launcher +# includes them as flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/. +_BMM_EXPORT_HEADERS = [ + "BatchedGemmEnums.h", + "BatchedGemmInterface.h", + "BatchedGemmOptions.h", + "Enums.h", + "GemmGatedActOptions.h", + "GemmOptions.h", + "KernelParams.h", + "KernelParamsDecl.h", + "KernelTraits.h", + "TmaDescriptor.h", + "trtllm/gen/CommonUtils.h", + "trtllm/gen/CudaArchDecl.h", + "trtllm/gen/CudaKernelLauncher.h", + "trtllm/gen/DtypeDecl.h", + "trtllm/gen/MmaDecl.h", + "trtllm/gen/SfLayoutDecl.h", + "trtllm/gen/SparsityDecl.h", +] + +_SOURCES = [ + "csrc/nv_internal/cpp/kernels/quantization.cu", + "csrc/nv_internal/cpp/common/envUtils.cpp", + "csrc/nv_internal/cpp/common/logger.cpp", + "csrc/nv_internal/cpp/common/stringUtils.cpp", + "csrc/nv_internal/cpp/common/tllmException.cpp", + "csrc/nv_internal/cpp/common/memoryUtils.cu", + "csrc/trtllm_fused_moe_kernel_launcher.cu", + "csrc/trtllm_fused_moe_runner.cu", + "csrc/fused_moe/trtllm_backend/trtllm_fused_moe_routing_deepseek.cu", + "csrc/fused_moe/trtllm_backend/trtllm_fused_moe_routing_llama4.cu", + "csrc/fused_moe/trtllm_backend/trtllm_fused_moe_routing_custom.cu", + "csrc/fused_moe/trtllm_backend/trtllm_fused_moe_routing_common.cu", + "csrc/fused_moe/trtllm_backend/trtllm_fused_moe_dev_kernel.cu", + "csrc/trtllm_batched_gemm_runner.cu", +] + + +logger = logging.getLogger(__name__) + + +def cubin_pool_dir() -> Optional[pathlib.Path]: + p = envs.SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL.get() + if not p: + return None + pool = pathlib.Path(p) + return pool if pool.is_dir() else None + + +def _flashinfer_data_dir() -> Optional[pathlib.Path]: + """The installed public flashinfer package's JIT source tree (ships + csrc/, include/ and its pinned cutlass), used as the base layer under + the pool's overlay.""" + try: + import flashinfer # noqa: PLC0415 + except ImportError: + return None + data = pathlib.Path(flashinfer.__file__).parent / "data" + return data if (data / "csrc").is_dir() else None + + +def available() -> bool: + pool = cubin_pool_dir() + return ( + pool is not None + and (pool / "flashinferMetaInfo.h").is_file() + and (pool / "local").is_dir() + # Modified sources ship in the pool's overlay/, everything else + # compiles from the installed flashinfer package. + and (pool / "overlay" / "csrc").is_dir() + and _flashinfer_data_dir() is not None + ) + + +def _stage_headers(pool: pathlib.Path) -> pathlib.Path: + """Copy the pool's ABI headers into a content-addressed include tree.""" + meta = (pool / "flashinferMetaInfo.h").read_bytes() + tag = hashlib.sha256(meta).hexdigest()[:12] + cache = pathlib.Path( + os.environ.get("TVM_FFI_CACHE_DIR", "~/.cache/tvm-ffi") + ).expanduser() + root = cache / "trtllm_gen_moe_headers" / tag + dest = root / "flashinfer" / "trtllm" / "batched_gemm" / "trtllmGen_bmm_export" + stamp = root / ".staged" + if not stamp.is_file(): + for name in _BMM_EXPORT_HEADERS: + target = dest / name + target.parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(pool / name, target) + shutil.copyfile(pool / "flashinferMetaInfo.h", dest / "flashinferMetaInfo.h") + stamp.touch() + return root + + +def _cuda_home() -> pathlib.Path: + home = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH") + if not home: + nvcc = shutil.which("nvcc") + home = str(pathlib.Path(nvcc).parent.parent) if nvcc else "/usr/local/cuda" + return pathlib.Path(home) + + +def _cuda_include_dir() -> str: + return str(_cuda_home() / "include") + + +def _cuda_stub_ldflags() -> list[str]: + """-L flags for the libcuda driver stub, so -lcuda links in bare build + environments (containers without the driver lib on the default linker + path); the real driver is dlopened at runtime as usual.""" + home = _cuda_home() + stubs = [ + home / "lib64" / "stubs", + *home.glob("targets/*/lib/stubs"), + ] + return [f"-L{s}" for s in stubs if s.is_dir()] + + +_CUBIN_CB_KEEPALIVE = {} + + +def _setup_cubin_loader(so_path: str, pool_local: pathlib.Path) -> None: + """Register the ctypes callback the .so uses to fetch cubins by name. + + The runner requests ``/`` (absolute, + because the pool path is baked in at compile time); we read the bytes + and hand them back via FlashInferSetCurrentCubin. + """ + if so_path in _CUBIN_CB_KEEPALIVE: + return + lib = ctypes.CDLL(so_path) + cb_type = ctypes.CFUNCTYPE(None, ctypes.c_char_p, ctypes.c_char_p) + + def _get_cubin(name: bytes, sha256: bytes) -> None: + rel = name.decode() + path = pathlib.Path(rel) + if not path.is_absolute(): + path = pool_local / rel + if path.suffix != ".cubin": + path = path.with_name(path.name + ".cubin") + data = path.read_bytes() + want = sha256.decode() + if want: + got = hashlib.sha256(data).hexdigest() + if got != want: + raise RuntimeError( + f"cubin sha mismatch for {path}: want {want} got {got}" + ) + lib.FlashInferSetCurrentCubin( + ctypes.cast(ctypes.create_string_buffer(data, len(data)), ctypes.c_char_p), + ctypes.c_int(len(data)), + ) + + cb = cb_type(_get_cubin) + _CUBIN_CB_KEEPALIVE[so_path] = (lib, cb) + lib.FlashInferSetCubinCallback(cb) + + +@cache_once +def _jit_trtllm_gen_moe_module() -> Module: + pool = cubin_pool_dir() + fi_data = _flashinfer_data_dir() + if pool is None or not (pool / "overlay" / "csrc").is_dir() or fi_data is None: + raise RuntimeError( + "trtllm-gen MoE sources not found: point " + "SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL at an unpacked cubin pool " + "(cubins + flat ABI headers + overlay/) and install the public " + "flashinfer package." + ) + # Overlay first (its modified sources/headers shadow the public copies), + # installed flashinfer data as the base. + src_roots = [pool / "overlay", fi_data] + include_roots = [pool / "overlay", fi_data] + + def _resolve_source(rel: str) -> str: + for root in src_roots: + cand = root / rel + if cand.is_file(): + return str(cand) + raise RuntimeError(f"trtllm-gen MoE source not found in any root: {rel}") + + staged = _stage_headers(pool) + meta_tag = staged.name + cubin_path = str((pool / "local").resolve()) + + cache = pathlib.Path( + os.environ.get("TVM_FFI_CACHE_DIR", "~/.cache/tvm-ffi") + ).expanduser() + # Flags are not part of load_jit's source hash: fold the pool identity + # (meta hash + path) into the module marker so a pool change rebuilds. + path_tag = hashlib.sha256(cubin_path.encode()).hexdigest()[:8] + build_dir = cache / f"sgl_trtllm_gen_moe_{meta_tag}_{path_tag}" + + cpp_files = [_resolve_source(s) for s in _SOURCES if s.endswith(".cpp")] + cuda_files = [_resolve_source(s) for s in _SOURCES if s.endswith(".cu")] + # quantization.cu emits fp4 cvt instructions (.e2m1x2) that need the + # arch-specific feature set: compile for sm_XXXa, not plain sm_XXX. + # The trtllm-gen cubins themselves are prebuilt (sm100f) and loaded at + # runtime, unaffected by this flag. + arch = get_jit_cuda_arch() + with override_jit_cuda_arch(arch.major, arch.minor, "a"): + module = load_jit( + "trtllm_gen_moe", + meta_tag, + path_tag, + external_cpp_files=cpp_files, + external_cuda_files=cuda_files, + header_only=False, # the launcher exports its own tvm-ffi functions + extra_cflags=["-fvisibility=hidden"], + extra_cuda_cflags=[ + "-DTLLM_GEN_EXPORT_INTERFACE", + "-DTLLM_GEN_EXPORT_FLASHINFER", + "-DTLLM_ENABLE_CUDA", + "-DENABLE_BF16", + "-DENABLE_FP8", + "-DENABLE_FP4", + "-DCUTLASS_ENABLE_GDC_FOR_SM100=1", + "-DTLLM_GEN_LOCAL_CUBINS_ABI", + "-DFLASHINFER_PRIVATE_MOE_FFI_NAMES", + "-DFLASHINFER_PRIVATE_MOE_LEAN_ROUTING", + f'-DTLLM_GEN_GEMM_CUBIN_PATH=\\"{cubin_path}\\"', + "-Xcompiler=-fvisibility=hidden", + ], + extra_ldflags=[*_cuda_stub_ldflags(), "-lcuda", "-lnvrtc"], + extra_include_paths=[ + str(staged), + str( + staged + / "flashinfer" + / "trtllm" + / "batched_gemm" + / "trtllmGen_bmm_export" + ), + # Per-root include layout: include/, csrc/, csrc/nv_internal/, + # csrc/nv_internal/include/, plus the flashinfer package's + # pinned CUTLASS (data/cutlass/). The overlay root comes first + # so modified headers shadow the public copies. + *[ + str(root / sub) + for root in include_roots + for sub in ( + "include", + "csrc", + "csrc/nv_internal", + "csrc/nv_internal/include", + ) + ], + *[ + str(root / "cutlass" / "include") + for root in include_roots + if (root / "cutlass" / "include").is_dir() + ], + # Host .cpp files (g++) need the CUDA headers explicitly; nvcc + # adds them implicitly for .cu. CUDA 13's bundled CCCL is + # used as-is (mixing another pinned CCCL with the toolkit's + # explodes). + _cuda_include_dir(), + ], + build_directory=str(build_dir), + ) + so_files = sorted(build_dir.glob("*.so")) + if not so_files: + raise RuntimeError(f"no built .so under {build_dir}") + _setup_cubin_loader(str(so_files[-1]), pool / "local") + return module + + +def trtllm_fp4_block_scale_moe( + routing_logits: torch.Tensor, + routing_bias: Optional[torch.Tensor], + hidden_states: torch.Tensor, + hidden_states_scale: Optional[torch.Tensor], + gemm1_weights: torch.Tensor, + gemm1_weights_scale: torch.Tensor, + gemm1_alpha: Optional[torch.Tensor], + gemm1_beta: Optional[torch.Tensor], + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + output1_scale_scalar: Optional[torch.Tensor], + output1_scale_gate_scalar: Optional[torch.Tensor], + output2_scale_scalar: Optional[torch.Tensor], + num_experts: int, + top_k: int, + n_group: Optional[int], + topk_group: Optional[int], + intermediate_size: int, + routed_scaling_factor: Optional[float], + routing_method_type: int = ROUTING_DEEPSEEK_V3, + activation_type: int = ACTIVATION_SITU, + norm_topk_prob: bool = True, + local_expert_offset: int = 0, + local_num_experts: Optional[int] = None, + tactic: Sequence[int] = (-1, -1), + output: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """FP4 block-scale MoE with routing from logits and finalize fused. + + ``hidden_states``: bf16 ``[T, hidden]`` (w4a16) or MxFP8-packed uint8 + with ``hidden_states_scale`` (w4a8). Weights are trtllm-gen shuffled + MxFP4 (uint8 packed, fp8 block scales, MajorK). ``tactic`` is the + (gemm1, gemm2) config index pair; ``(-1, -1)`` = runner heuristic. + """ + module = _jit_trtllm_gen_moe_module() + # The FFI launcher reads these as dense row-major; a strided slice + # (e.g. a fused-GEMM split) would silently mis-route. + routing_logits = routing_logits.contiguous() + hidden_states = hidden_states.contiguous() + num_tokens = routing_logits.shape[0] + hidden_size = hidden_states.shape[-1] + if hidden_states.dtype == torch.uint8: + hidden_size *= 2 + device = hidden_states.device + topk_ids = torch.empty(num_tokens, top_k, dtype=torch.int32, device=device) + topk_weights = torch.empty( + num_tokens, top_k, dtype=routing_logits.dtype, device=device + ) + if output is None: + output = torch.empty( + num_tokens, hidden_size, dtype=torch.bfloat16, device=device + ) + module.trtllm_fp4_block_scale_moe_private( + _ROUTING_INPUT_FROM_LOGITS, + routing_logits, + topk_ids, + topk_weights, + routing_bias, + hidden_states, + hidden_states_scale, + gemm1_weights, + gemm1_weights_scale, + None, # gemm1_bias + gemm1_alpha, + gemm1_beta, + None, # gemm1_clamp_limit + gemm2_weights, + gemm2_weights_scale, + None, # gemm2_bias + output1_scale_scalar, + output1_scale_gate_scalar, + output2_scale_scalar, + None, # per_token_scale + num_experts, + top_k, + n_group, + topk_group, + intermediate_size, + local_expert_offset, + num_experts if local_num_experts is None else local_num_experts, + routed_scaling_factor, + routing_method_type, + True, # do_finalize + True, # enable_pdl + activation_type, + output, + list(tactic), + norm_topk_prob, + None, # routing_replay_out + ) + return output + + +def trtllm_fp4_block_scale_routed_moe( + packed_topk_ids: torch.Tensor, + hidden_states: torch.Tensor, + hidden_states_scale: Optional[torch.Tensor], + gemm1_weights: torch.Tensor, + gemm1_weights_scale: torch.Tensor, + gemm1_alpha: Optional[torch.Tensor], + gemm1_beta: Optional[torch.Tensor], + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + output1_scale_scalar: Optional[torch.Tensor], + output1_scale_gate_scalar: Optional[torch.Tensor], + output2_scale_scalar: Optional[torch.Tensor], + num_experts: int, + top_k: int, + intermediate_size: int, + activation_type: int = ACTIVATION_SITU, + local_expert_offset: int = 0, + local_num_experts: Optional[int] = None, + tactic: Sequence[int] = (-1, -1), + output: Optional[torch.Tensor] = None, + do_finalize: bool = True, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """FP4 block-scale MoE with PRECOMPUTED routing (PackedPrecomputed). + + ``packed_topk_ids``: int32 ``[T, top_k]`` with ``(expert_id << 16) | + bf16-weight-bits`` (PackTopkIds layout) — selection and weights come + from the caller's router, the in-op routing kernels are skipped. This + is the fast path at small T, where the in-op single-CTA routing kernel + (~22 µs at 896 experts) costs more than an external radix router. + + ``do_finalize=False`` skips the in-op finalize (top-k weighted + unpermute) and returns its inputs instead: + ``(gemm2_output [padded_rows, hidden] bf16 in permuted layout, + topk_weights [T, top_k] bf16 unpacked from packed_topk_ids, + expanded_idx_to_permuted_idx [T*top_k] int32 with -1 = dropped slot)``. + ``output`` is left unwritten in that mode. + """ + module = _jit_trtllm_gen_moe_module() + hidden_states = hidden_states.contiguous() + num_tokens = packed_topk_ids.shape[0] + hidden_size = hidden_states.shape[-1] + if hidden_states.dtype == torch.uint8: + hidden_size *= 2 + device = hidden_states.device + # Mode 2 unpacks the weights in-kernel; this is its output buffer. + topk_weights = torch.empty(num_tokens, top_k, dtype=torch.bfloat16, device=device) + if output is None: + output = torch.empty( + num_tokens, hidden_size, dtype=torch.bfloat16, device=device + ) + result = module.trtllm_fp4_block_scale_moe_private( + _ROUTING_INPUT_PACKED, + None, # routing_logits + packed_topk_ids.contiguous(), + topk_weights, + None, # routing_bias (already applied by the external router) + hidden_states, + hidden_states_scale, + gemm1_weights, + gemm1_weights_scale, + None, # gemm1_bias + gemm1_alpha, + gemm1_beta, + None, # gemm1_clamp_limit + gemm2_weights, + gemm2_weights_scale, + None, # gemm2_bias + output1_scale_scalar, + output1_scale_gate_scalar, + output2_scale_scalar, + None, # per_token_scale + num_experts, + top_k, + None, # n_group + None, # topk_group + intermediate_size, + local_expert_offset, + num_experts if local_num_experts is None else local_num_experts, + 1.0, # routed_scaling_factor (already applied by the router) + _ROUTING_TOPK, # routing_method_type (unused for precomputed) + do_finalize, + True, # enable_pdl + activation_type, + output, + list(tactic), + True, # norm_topk_prob (unused for precomputed) + None, # routing_replay_out + ) + if do_finalize: + return output + # Deferred: [gemm2_output, expert_weights (None in packed mode — the + # weights live in the topk_weights buffer mode 2 unpacked into), + # expanded_idx_to_permuted_idx]. Index access — iterating the tvm-ffi + # Array yields one-shot dlpack capsules. + return result[0], topk_weights, result[2] diff --git a/python/sglang/kernels/ops/sampling/top_p_renorm_triton.py b/python/sglang/kernels/ops/sampling/top_p_renorm_triton.py new file mode 100644 index 000000000..3730e0b5d --- /dev/null +++ b/python/sglang/kernels/ops/sampling/top_p_renorm_triton.py @@ -0,0 +1,127 @@ +"""ROCm-compatible top-p probability renormalization fallback.""" + +from __future__ import annotations + +from typing import Union + +import torch +import triton +import triton.language as tl + +_BLOCK_SIZE = 1024 + + +@triton.jit +def _mask_and_partial_sum_kernel( + probs_ptr, + pivots_ptr, + out_ptr, + partial_sums_ptr, + vocab_size: tl.constexpr, + num_chunks: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + row = tl.program_id(0) + chunk = tl.program_id(1) + offsets = chunk * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < vocab_size + row_offsets = row * vocab_size + offsets + + probs = tl.load(probs_ptr + row_offsets, mask=mask, other=0.0).to(tl.float32) + pivot = tl.load(pivots_ptr + row) + kept = tl.where(mask & (probs >= pivot), probs, 0.0) + + tl.store(out_ptr + row_offsets, kept, mask=mask) + tl.store(partial_sums_ptr + row * num_chunks + chunk, tl.sum(kept, axis=0)) + + +@triton.jit +def _normalize_kernel( + out_ptr, + row_sums_ptr, + numel, + vocab_size: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < numel + row = offsets // vocab_size + values = tl.load(out_ptr + offsets, mask=mask, other=0.0).to(tl.float32) + denominator = tl.load(row_sums_ptr + row, mask=mask, other=1.0) + tl.store(out_ptr + offsets, values / denominator, mask=mask) + + +def top_p_renorm_probs_triton( + probs: torch.Tensor, top_p: Union[torch.Tensor, float] +) -> torch.Tensor: + """Apply exact top-p thresholding and renormalize each probability row. + + Sorting and prefix sums use PyTorch's device kernels because a vocabulary-sized + in-register Triton sort does not scale to 100K+ vocabularies. Triton performs + the bandwidth-heavy masking, partial reduction, and normalization. + """ + if probs.ndim != 2: + raise ValueError(f"probs must be 2D, got shape={tuple(probs.shape)}") + if not probs.is_cuda: + raise ValueError("top_p_renorm_probs_triton requires a CUDA/HIP tensor") + + probs_fp32 = probs.float().contiguous() + batch_size, vocab_size = probs_fp32.shape + if batch_size == 0 or vocab_size == 0: + return probs_fp32 + + if isinstance(top_p, torch.Tensor): + top_ps = top_p.to(device=probs.device, dtype=torch.float32).reshape(-1) + if top_ps.numel() == 1: + top_ps = top_ps.expand(batch_size) + elif top_ps.numel() != batch_size: + raise ValueError( + f"top_p must be scalar or have one value per row, got " + f"{top_ps.numel()} values for {batch_size} rows" + ) + else: + if not 0.0 < float(top_p) <= 1.0: + raise ValueError("top_p values must be in (0, 1]") + top_ps = torch.full( + (batch_size,), float(top_p), device=probs.device, dtype=torch.float32 + ) + + # Match FlashInfer's threshold semantics: sort ascending, discard the prefix + # whose cumulative mass is below 1 - p, and retain all ties at the pivot. + sorted_probs = torch.sort(probs_fp32, dim=-1).values + cdf = torch.cumsum(sorted_probs, dim=-1) + cutoff = torch.searchsorted(cdf, (1.0 - top_ps).unsqueeze(1), right=False).squeeze( + 1 + ) + cutoff.clamp_(max=vocab_size - 1) + pivots = sorted_probs.gather(1, cutoff.unsqueeze(1)).squeeze(1).contiguous() + + num_chunks = triton.cdiv(vocab_size, _BLOCK_SIZE) + out = torch.empty_like(probs_fp32) + partial_sums = torch.empty( + (batch_size, num_chunks), device=probs.device, dtype=torch.float32 + ) + _mask_and_partial_sum_kernel[(batch_size, num_chunks)]( + probs_fp32, + pivots, + out, + partial_sums, + vocab_size=vocab_size, + num_chunks=num_chunks, + BLOCK_SIZE=_BLOCK_SIZE, + num_warps=8, + ) + + row_sums = partial_sums.sum(dim=1) + _normalize_kernel[(triton.cdiv(out.numel(), _BLOCK_SIZE),)]( + out, + row_sums, + out.numel(), + vocab_size=vocab_size, + BLOCK_SIZE=_BLOCK_SIZE, + num_warps=8, + ) + return out + + +__all__ = ["top_p_renorm_probs_triton"] diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py index 0875ee20e..7935e9a7b 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py +++ b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py @@ -211,6 +211,12 @@ class CustomAllReduceV2: multicast_ptr = int(symm_mem.multicast_ptr) can_multicast = multicast_ptr != 0 + # multicast VA of the slab base (== the push workspace, at offset 0); + # consumed by the K3 all_reduce push kernel + self.mc_base_ptr = multicast_ptr if can_multicast else 0 + # multicast VA of the pull-semaphore region; the K3 pull kernels reuse + # these semaphores (same reservation protocol, multicast-signaled) + self.pull_sem_mc_ptr = multicast_ptr + pull_sem_offset if can_multicast else 0 pull_mc_workspace = multicast_ptr + pull_ws_offset if can_multicast else None if not can_multicast or cfg.num_mc_blocks is None: self.config = self.config._replace(num_mc_blocks=None) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 625ca3465..029a917a4 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -661,6 +661,9 @@ class Envs: # Launch the TRT-LLM MoE grouped GEMMs with PDL only at or below this # token count. SGLANG_TRTLLM_MOE_PDL_MAX_TOKENS = EnvInt(8192) + # Unpacked cubin pool for the JIT-built trtllm-gen fused MoE (cubins + flat + # ABI headers + overlay/). Unset means the path is unavailable, not empty. + SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL = EnvStr(None) # SGLang needs to know FlashInfer NVFP4 4over6 config to compute the global scale factor. FLASHINFER_NVFP4_4OVER6 = EnvBool(False) FLASHINFER_NVFP4_4OVER6_E4M3_USE_256 = EnvBool(False) diff --git a/test/registered/kernels/ops/attention/test_kda_fused_decode.py b/test/registered/kernels/ops/attention/test_kda_fused_decode.py new file mode 100644 index 000000000..3c4e6d665 --- /dev/null +++ b/test/registered/kernels/ops/attention/test_kda_fused_decode.py @@ -0,0 +1,215 @@ +"""Kimi-K3 fused KDA decode must match the existing unfused decode chain. + +The fused kernel replaces: + + causal_conv1d_update -> kda_packed_decode -> sigmoid-gated RMSNorm + +This file covers the local head layouts used by Kimi-K3 TP8/TP16/TP32: +H = 12/6/3. The H=6 and H=3 cases are the branches added by the fixed-head +dispatch in ``kda_fused_decode.cuh``. +""" + +import pytest +import torch + +from sglang.kernels.ops.attention import kda_fused_decode +from sglang.kernels.ops.attention.fla.fused_norm_gate import rms_norm_gated +from sglang.kernels.ops.attention.fla.fused_recurrent import ( + fused_recurrent_kda_packed_decode, +) +from sglang.kernels.ops.mamba.causal_conv1d_triton import causal_conv1d_update +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=8, stage="base-b-kernel-unit", runner_config="1-gpu-large") + +_HEAD_DIM = 128 +_CONV_STATE_W = 3 +_SLOTS = 8 +_BATCH = 4 + + +def _randn(shape, dtype, generator, scale=1.0): + return (torch.randn(shape, device="cuda", generator=generator) * scale).to(dtype) + + +def _make_case(heads: int, seed: int): + generator = torch.Generator(device="cuda").manual_seed(seed) + seg = heads * _HEAD_DIM + conv_dim = 3 * seg + + # Keep magnitudes moderate so fp32 state updates stay in a stable range. + mixed_qkv = _randn((_BATCH, conv_dim), torch.bfloat16, generator, scale=0.2) + a = _randn((_BATCH, seg), torch.bfloat16, generator, scale=0.2) + b = _randn((_BATCH, heads), torch.bfloat16, generator, scale=0.2) + onorm_g = _randn((_BATCH, seg), torch.bfloat16, generator, scale=0.2) + + conv_states = _randn( + (_SLOTS, _CONV_STATE_W, conv_dim), torch.bfloat16, generator, scale=0.2 + ) + ssm_states = _randn( + (_SLOTS, heads, _HEAD_DIM, _HEAD_DIM), torch.float32, generator, scale=0.02 + ) + cache_indices = torch.arange(_BATCH, device="cuda", dtype=torch.int32) + + conv_weights = _randn((conv_dim, 4), torch.float32, generator, scale=0.1) + conv_bias = _randn((conv_dim,), torch.float32, generator, scale=0.05) + a_log = _randn((heads,), torch.float32, generator, scale=0.1) + dt_bias = _randn((seg,), torch.float32, generator, scale=0.1) + onorm_weight = _randn((_HEAD_DIM,), torch.float32, generator, scale=0.1) + 1.0 + + return ( + mixed_qkv, + a, + b, + onorm_g, + conv_states, + ssm_states, + cache_indices, + conv_weights, + conv_bias, + a_log, + dt_bias, + onorm_weight, + ) + + +def _run_unfused_reference( + mixed_qkv, + a, + b, + onorm_g, + conv_states, + ssm_states, + cache_indices, + conv_weights, + conv_bias, + a_log, + dt_bias, + onorm_weight, +): + heads = ssm_states.shape[-3] + qkv = causal_conv1d_update( + mixed_qkv, + conv_states.transpose(-1, -2), + conv_weights, + conv_bias, + activation="silu", + conv_state_indices=cache_indices, + ) + out = torch.empty( + (_BATCH, 1, heads, _HEAD_DIM), dtype=torch.bfloat16, device="cuda" + ) + out, _ = fused_recurrent_kda_packed_decode( + qkv, + a, + b, + a_log, + dt_bias, + _HEAD_DIM**-0.5, + ssm_states, + out, + cache_indices, + use_qk_l2norm_in_kernel=True, + ) + ref = rms_norm_gated( + out, + onorm_g.view(1, _BATCH, heads, _HEAD_DIM), + onorm_weight, + None, + activation="sigmoid", + eps=1e-6, + ) + return ref.transpose(0, 1).contiguous() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize( + "heads,tp_size", + [ + pytest.param(3, 32, id="tp32_h3"), + pytest.param(6, 16, id="tp16_h6"), + pytest.param(12, 8, id="tp8_h12"), + ], +) +def test_kda_fused_decode_matches_unfused_chain(heads: int, tp_size: int): + ( + mixed_qkv, + a, + b, + onorm_g, + conv_states, + ssm_states, + cache_indices, + conv_weights, + conv_bias, + a_log, + dt_bias, + onorm_weight, + ) = _make_case(heads=heads, seed=20260731 + tp_size) + + conv_ref = conv_states.clone() + conv_fused = conv_states.clone() + state_ref = ssm_states.clone() + state_fused = ssm_states.clone() + + w_q_t, w_k_t, w_v_t = [ + weight.t().contiguous() + for weight in conv_weights.split(heads * _HEAD_DIM, dim=0) + ] + + assert kda_fused_decode.covered( + mixed_qkv, + a, + b, + conv_fused, + state_fused, + cache_indices, + onorm_g, + ) + + ref = _run_unfused_reference( + mixed_qkv.clone(), + a, + b, + onorm_g, + conv_ref, + state_ref, + cache_indices, + conv_weights, + conv_bias, + a_log, + dt_bias, + onorm_weight, + ) + fused = kda_fused_decode.kda_fused_decode( + mixed_qkv.clone(), + a, + b, + conv_fused, + w_q_t, + w_k_t, + w_v_t, + conv_bias, + a_log, + dt_bias, + onorm_g, + onorm_weight, + state_fused, + cache_indices, + scale=_HEAD_DIM**-0.5, + onorm_eps=1e-6, + ) + torch.cuda.synchronize() + + # JIT log breadcrumb for PR/CI evidence that the fused fixed-head branch ran. + print(f"K3 fused KDA decode test used fused path: TP{tp_size}, H={heads}") + + torch.testing.assert_close(fused, ref, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(state_fused, state_ref, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(conv_fused, conv_ref, rtol=0, atol=0) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__])) diff --git a/test/registered/kernels/ops/attention/test_kda_prefill.py b/test/registered/kernels/ops/attention/test_kda_prefill.py new file mode 100644 index 000000000..73b33d1d3 --- /dev/null +++ b/test/registered/kernels/ops/attention/test_kda_prefill.py @@ -0,0 +1,166 @@ +import unittest + +import torch +import torch.nn.functional as F + +from sglang.kernels.ops.attention.fla.kda import chunk_kda +from sglang.kernels.ops.attention.linear.kda_nvidia_prefill import ( + chunk_kda_fwd as nvidia_chunk_kda_fwd, +) +from sglang.kernels.ops.attention.linear.kda_ptx_prefill import ( + chunk_kda_fwd as ptx_chunk_kda_fwd, +) +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=180, stage="base-b-kernel-unit", runner_config="4-gpu-b200") +register_cuda_ci(est_time=180, stage="base-c", runner_config="4-gpu-gb300") + + +def _inputs(seed, seq_len=128): + generator = torch.Generator(device="cuda").manual_seed(seed) + batch_size, num_heads, head_dim = 1, 2, 128 + shape = (batch_size, seq_len, num_heads, head_dim) + q = torch.randn(shape, generator=generator, device="cuda", dtype=torch.bfloat16) + k = torch.randn(shape, generator=generator, device="cuda", dtype=torch.bfloat16) + v = ( + 0.1 + * torch.randn( + shape, + generator=generator, + device="cuda", + dtype=torch.float32, + ) + ).to(torch.bfloat16) + gate = torch.randn(shape, generator=generator, device="cuda", dtype=torch.bfloat16) + beta_logits = torch.randn( + shape[:-1], + generator=generator, + device="cuda", + dtype=torch.bfloat16, + ) + a_log = torch.randn( + num_heads, generator=generator, device="cuda", dtype=torch.float32 + ) + dt_bias = torch.randn( + num_heads * head_dim, + generator=generator, + device="cuda", + dtype=torch.float32, + ) + state = torch.zeros( + batch_size, + num_heads, + head_dim, + head_dim, + device="cuda", + dtype=torch.float32, + ) + return q, k, v, gate, beta_logits, a_log, dt_bias, state + + +def _reference(q, k, v, gate, beta, a_log, dt_bias, state, fused_qk_norm): + return chunk_kda( + q=q, + k=k, + v=v, + g=gate, + beta=beta, + scale=q.shape[-1] ** -0.5, + initial_state=state, + initial_state_indices=torch.arange( + q.shape[0], device="cuda", dtype=torch.int32 + ), + use_qk_l2norm_in_kernel=fused_qk_norm, + A_log=a_log, + dt_bias=dt_bias, + lower_bound=-5.0, + ) + + +class TestKdaPrefill(CustomTestCase): + @torch.inference_mode() + def test_nvidia_prefill(self): + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10: + self.skipTest("NVIDIA KDA prefill requires datacenter Blackwell") + q, k, v, gate, beta_logits, a_log, dt_bias, state = _inputs(0) + q = F.normalize(q.float(), dim=-1).to(torch.bfloat16) + k = F.normalize(k.float(), dim=-1).to(torch.bfloat16) + beta = torch.sigmoid(beta_logits.float()).to(torch.bfloat16) + actual, actual_state = nvidia_chunk_kda_fwd( + q=q, + k=k, + v=v, + g=gate, + beta=beta, + scale=q.shape[-1] ** -0.5, + initial_state=state.transpose(-1, -2).contiguous(), + output_final_state=True, + safe_gate=True, + lower_bound=-5.0, + use_gate_in_kernel=True, + A_log=a_log, + dt_bias=dt_bias, + )[:2] + expected = _reference( + q, k, v, gate, beta, a_log, dt_bias, state, fused_qk_norm=False + ) + torch.testing.assert_close( + actual.float(), expected.float(), rtol=2e-2, atol=3e-2 + ) + torch.testing.assert_close( + actual_state.transpose(-1, -2), + state, + rtol=2e-2, + atol=3e-2, + ) + + @torch.inference_mode() + def test_ptx_prefill(self): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != ( + 10, + 3, + ): + self.skipTest("PTX KDA prefill requires GB300") + q, k, v, gate, beta_logits, a_log, dt_bias, state = _inputs(1) + actual, actual_state = ptx_chunk_kda_fwd( + q=q, + k=k, + v=v, + g=gate, + beta=beta_logits, + scale=q.shape[-1] ** -0.5, + initial_state=state.transpose(-1, -2).contiguous(), + output_final_state=True, + safe_gate=True, + lower_bound=-5.0, + use_gate_in_kernel=True, + A_log=a_log, + dt_bias=dt_bias, + use_qk_l2norm_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + )[:2] + expected = _reference( + q, + k, + v, + gate, + torch.sigmoid(beta_logits.float()).to(torch.bfloat16), + a_log, + dt_bias, + state, + fused_qk_norm=True, + ) + torch.testing.assert_close( + actual.float(), expected.float(), rtol=2e-2, atol=3e-2 + ) + torch.testing.assert_close( + actual_state.transpose(-1, -2), + state, + rtol=2e-2, + atol=3e-2, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/kernels/ops/kimi_k3/test_collectives.py b/test/registered/kernels/ops/kimi_k3/test_collectives.py new file mode 100644 index 000000000..d88892523 --- /dev/null +++ b/test/registered/kernels/ops/kimi_k3/test_collectives.py @@ -0,0 +1,359 @@ +from __future__ import annotations + +import atexit +import os + +import pytest +import torch +import torch.distributed as dist + +import sglang.srt.distributed.parallel_state as ps +from sglang.kernels.jit.utils import cache_once +from sglang.kernels.ops.communication.mp import register_comm_cleanup +from sglang.kernels.ops.kimi_k3 import ( + all_reduce, + attn_res, + gemm_ag, + gemm_ar, + sp_collective, +) +from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import ( + CustomAllReduceV2, +) +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kernels.utils import multigpu_pytest_main + +register_cuda_ci(est_time=240, stage="base-b-kernel-unit", runner_config="4-gpu-b200") +register_cuda_ci(est_time=480, suite="nightly-8-gpu-b200", nightly=True) + +_HIDDEN_SIZE = 7168 +_GEMM_AR_K_TOTAL = 12288 +_GEMM_AG_WORLD_SIZE = 8 +_MB = 1024 * 1024 +_SP_TUNING = sp_collective.Tuning(num_blocks=1, block_size=256) + + +def _device(): + return torch.device("cuda", int(os.environ["LOCAL_RANK"])) + + +def _require_sm100(): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() < (10, 0): + pytest.skip("Kimi K3 collectives require SM100+") + + +@cache_once +def _init_world(): + local_rank = int(os.environ["LOCAL_RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(local_rank) + dist.init_process_group(backend="gloo") + ps._WORLD = coord = ps.init_world_group( + ranks=list(range(world_size)), + local_rank=local_rank, + backend="nccl", + ) + atexit.register(dist.destroy_process_group) + cpu_group = coord.cpu_group + assert isinstance(cpu_group, dist.ProcessGroup) + nccl_group = dist.new_group(backend="nccl", device_id=_device()) + return cpu_group, nccl_group + + +@cache_once +def _init_comm(): + cpu_group, _ = _init_world() + comm = CustomAllReduceV2( + cpu_group, + _device(), + max_pull_size=4 * _MB, + max_push_size=4 * _MB, + ) + if comm.disabled or comm.mc_base_ptr == 0: + raise RuntimeError("Kimi K3 collectives require multicast symmetric memory") + all_reduce.register_comm(comm.obj, pull_sem_mc_ptr=comm.pull_sem_mc_ptr) + sp_collective.register_comm(comm.obj, pull_sem_mc_ptr=comm.pull_sem_mc_ptr) + attn_res.register_comm(comm.obj, pull_sem_mc_ptr=comm.pull_sem_mc_ptr) + register_comm_cleanup(comm) + return comm + + +@cache_once +def _init_gemm_ar(): + cpu_group, _ = _init_world() + world_size = dist.get_world_size() + gemm_ar.init( + world_size=world_size, + rank=dist.get_rank(), + group=cpu_group, + k=_GEMM_AR_K_TOTAL // world_size, + ) + + +def _symmetric_tensor(shape): + from torch._C._distributed_c10d import _SymmetricMemory + + cpu_group, _ = _init_world() + tensor = _SymmetricMemory.empty_strided_p2p( + shape, + torch.empty(shape).stride(), + torch.bfloat16, + _device(), + cpu_group.group_name, + ) + handle = _SymmetricMemory.rendezvous(tensor) + rank = dist.get_rank() + multicast_ptr = ( + int(handle.multicast_ptr) + tensor.data_ptr() - int(handle.buffer_ptrs[rank]) + ) + if multicast_ptr == 0: + raise RuntimeError("symmetric tensor has no multicast mapping") + return tensor, handle, multicast_ptr + + +@torch.inference_mode() +def test_all_reduce_push(): + _require_sm100() + comm = _init_comm() + rank = dist.get_rank() + generator = torch.Generator().manual_seed(10 + rank) + x = torch.randint( + 0, + 16, + (_HIDDEN_SIZE,), + generator=generator, + dtype=torch.bfloat16, + ).to(_device()) + residual = ( + torch.arange(_HIDDEN_SIZE, dtype=torch.int32, device=_device()) + .remainder_(7) + .to(torch.bfloat16) + ) + expected = x.clone() + _, nccl_group = _init_world() + dist.all_reduce(expected, group=nccl_group) + expected += residual + + all_reduce.all_reduce_push_res( + comm.world_size, + x, + residual, + ws_mc_base=comm.mc_base_ptr, + ) + torch.cuda.synchronize() + torch.testing.assert_close(x, expected, rtol=0, atol=0) + + +@torch.inference_mode() +def test_sequence_parallel_collectives(): + _require_sm100() + comm = _init_comm() + rank, world_size = dist.get_rank(), comm.world_size + local_tokens = 2 + generator = torch.Generator(device="cuda").manual_seed(20 + rank) + reduce_input = torch.randn( + world_size * local_tokens, + _HIDDEN_SIZE, + generator=generator, + device=_device(), + dtype=torch.bfloat16, + ) + residual = torch.randn( + local_tokens, + _HIDDEN_SIZE, + generator=torch.Generator(device="cuda").manual_seed(21), + device=_device(), + dtype=torch.bfloat16, + ) + expected_reduce = reduce_input.float() + _, nccl_group = _init_world() + dist.all_reduce(expected_reduce, group=nccl_group) + lo = rank * local_tokens + expected_reduce = (expected_reduce[lo : lo + local_tokens] + residual.float()).to( + torch.bfloat16 + ) + reduce_output = torch.empty_like(expected_reduce) + sp_collective.reduce_scatter_res( + world_size, + reduce_input, + reduce_output, + residual, + tuning=_SP_TUNING, + ) + + gather_input = torch.randn( + local_tokens, + _HIDDEN_SIZE, + generator=generator, + device=_device(), + dtype=torch.bfloat16, + ) + expected_gather = torch.empty( + world_size * local_tokens, + _HIDDEN_SIZE, + device=_device(), + dtype=torch.bfloat16, + ) + dist.all_gather_into_tensor( + expected_gather, + gather_input, + group=nccl_group, + ) + gather_output = torch.empty_like(expected_gather) + sp_collective.all_gather( + world_size, + gather_input, + gather_output, + ws_mc_base=comm.mc_base_ptr, + tuning=_SP_TUNING, + ) + torch.cuda.synchronize() + + torch.testing.assert_close(reduce_output, expected_reduce, rtol=2e-2, atol=3e-2) + torch.testing.assert_close(gather_output, expected_gather, rtol=0, atol=0) + + +@torch.inference_mode() +def test_gemm_all_gather(): + _require_sm100() + if int(os.environ["WORLD_SIZE"]) != _GEMM_AG_WORLD_SIZE: + pytest.skip("Kimi K3 gemm_ag is compiled for TP8") + comm = _init_comm() + generator = torch.Generator().manual_seed(30) + x = ( + (torch.randn(1, gemm_ag.K, generator=generator) * 0.05) + .to(torch.bfloat16) + .to(_device()) + ) + weight = ( + (torch.randn(gemm_ag.N, gemm_ag.K, generator=generator) * 0.05) + .to(torch.bfloat16) + .to(_device()) + ) + bias = torch.randn(1, gemm_ag.N, generator=generator).to( + device=_device(), dtype=torch.bfloat16 + ) + output = torch.empty(1, gemm_ag.N, device=_device(), dtype=torch.bfloat16) + expected = (x.float() @ weight.float().t() + bias.float()).to(torch.bfloat16) + + gemm_ag.gemm_ag_up_proj( + comm.world_size, + x, + weight, + bias, + None, + output, + ws_mc_base=comm.mc_base_ptr, + ) + torch.cuda.synchronize() + torch.testing.assert_close(output, expected, rtol=3e-2, atol=3e-2) + + +@torch.inference_mode() +def test_gemm_all_reduce(): + _require_sm100() + _init_gemm_ar() + rank, world_size = dist.get_rank(), dist.get_world_size() + local_k = _GEMM_AR_K_TOTAL // world_size + generator = torch.Generator().manual_seed(40 + rank) + x = torch.randn(1, local_k, generator=generator).to( + device=_device(), dtype=torch.bfloat16 + ) + weight = torch.randn(gemm_ar.N, local_k, generator=generator).to( + device=_device(), dtype=torch.bfloat16 + ) + expected = (x.float() @ weight.float().t()).to(torch.bfloat16).float() + _, nccl_group = _init_world() + dist.all_reduce(expected, group=nccl_group) + + output = gemm_ar.o_proj_gemm_ar(x, weight) + torch.cuda.synchronize() + bad = ((output.float() - expected).abs() > 0.05 + 0.02 * expected.abs()).sum() + assert bad.item() <= output.numel() / 1000 + + +@torch.inference_mode() +def test_attention_residual_direct_all_gather(): + _require_sm100() + comm = _init_comm() + rank, local_tokens, num_bank_rows = dist.get_rank(), 2, 3 + generator = torch.Generator(device="cuda").manual_seed(50 + rank) + prefix = torch.randn( + local_tokens, + _HIDDEN_SIZE, + generator=generator, + device=_device(), + dtype=torch.bfloat16, + ) + bank = torch.randn( + local_tokens, + num_bank_rows + 1, + _HIDDEN_SIZE, + generator=generator, + device=_device(), + dtype=torch.bfloat16, + ) + combine_weight = torch.linspace( + -0.01, 0.01, _HIDDEN_SIZE, device=_device(), dtype=torch.bfloat16 + ) + output_weight = torch.linspace( + 1.25, 0.75, _HIDDEN_SIZE, device=_device(), dtype=torch.bfloat16 + ) + local_reference = torch.empty_like(prefix) + attn_res.attn_res_fused_tma( + prefix, + bank.clone(), + combine_weight, + output_weight, + local_reference, + num_bank_rows, + 1e-6, + ) + full_reference = torch.empty( + comm.world_size * local_tokens, + _HIDDEN_SIZE, + device=_device(), + dtype=torch.bfloat16, + ) + _, nccl_group = _init_world() + dist.all_gather_into_tensor( + full_reference, + local_reference, + group=nccl_group, + ) + + output, handle, multicast_ptr = _symmetric_tensor(tuple(full_reference.shape)) + attn_res.attn_res_fused_direct_ag( + comm.world_size, + prefix, + bank, + combine_weight, + output_weight, + output, + num_bank_rows, + 1e-6, + output_mc_ptr=multicast_ptr, + max_blocks=4, + ) + torch.cuda.synchronize() + torch.testing.assert_close(output, full_reference, rtol=2e-2, atol=3e-2) + assert handle is not None + + +def _precompile(num_gpus): + for world_size in num_gpus: + all_reduce._jit_module(world_size) + sp_collective._jit_module(world_size) + gemm_ar._jit_module(_GEMM_AR_K_TOTAL // world_size, world_size) + if _GEMM_AG_WORLD_SIZE in num_gpus: + gemm_ag._jit_module() + attn_res._jit_fused_tma_module(4, 1, 200) + + +if __name__ == "__main__": + multigpu_pytest_main( + __name__, + __file__, + num_gpus=(4, 8), + pre_launch_fn=_precompile, + ) diff --git a/test/registered/kernels/ops/kimi_k3/test_compute.py b/test/registered/kernels/ops/kimi_k3/test_compute.py new file mode 100644 index 000000000..b8ccc8869 --- /dev/null +++ b/test/registered/kernels/ops/kimi_k3/test_compute.py @@ -0,0 +1,453 @@ +import unittest + +import torch + +from sglang.kernels.ops.attention.fla.kda_replayssm_spec_decode import ( + commit_kda_replayssm_spec, +) +from sglang.kernels.ops.kimi_k3 import ( + situ_and_mul, + situ_and_mul_masked_post_quant, +) +from sglang.kernels.ops.kimi_k3.attn_res import attn_res_fused_tma +from sglang.kernels.ops.kimi_k3.kda_decode_mtp import ( + fused_kda_decode_mtp_dspark, +) +from sglang.kernels.ops.kimi_k3.mla_output_gate import ( + covered, + kimi_k3_mla_output_gate, +) +from sglang.kernels.ops.moe.moe_front import ( + NUM_EXPERTS, + TOPK, + fused_front, +) +from sglang.kernels.ops.moe.moe_fused_gate import moe_fused_gate +from sglang.srt.utils import get_device_sm +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=120, stage="base-b-kernel-unit", runner_config="4-gpu-b200") + +_HIDDEN_SIZE = 7168 +_GROUP_SIZE = 128 +_BETA = 4.0 +_LINEAR_BETA = 25.0 + + +def _situ_reference(gate_up): + gate, up = gate_up.chunk(2, dim=-1) + gate = gate.float() + up = up.float() + return ( + _BETA + * torch.tanh(gate / _BETA) + * torch.sigmoid(gate) + * _LINEAR_BETA + * torch.tanh(up / _LINEAR_BETA) + ) + + +def _unpack_ue8m0_scales(packed, num_groups): + num_experts, groups_per_word, num_tokens = packed.shape + exponents = packed.contiguous().view(torch.uint8) + exponents = exponents.view(num_experts, groups_per_word, num_tokens, 4) + exponents = exponents.permute(0, 2, 1, 3).reshape( + num_experts, num_tokens, num_groups + ) + return torch.exp2(exponents.float() - 127.0) + + +class TestKimiK3ComputeKernels(CustomTestCase): + @classmethod + def setUpClass(cls): + if not torch.cuda.is_available(): + raise unittest.SkipTest("CUDA is not available") + if get_device_sm() < 100: + raise unittest.SkipTest("Kimi K3 compute kernels require SM100a+") + + def test_attn_residual_and_prefix_write(self): + generator = torch.Generator(device="cuda").manual_seed(0) + + def randn(*shape): + return torch.randn(*shape, generator=generator, device="cuda") + + num_tokens, num_bank_rows, num_valid_bank_rows = 5, 8, 5 + prefix = randn(num_tokens, _HIDDEN_SIZE).to(torch.bfloat16) + bank = randn(num_tokens, num_bank_rows, _HIDDEN_SIZE).to(torch.bfloat16) + combine_weight = (randn(_HIDDEN_SIZE) * _HIDDEN_SIZE**-0.5).to(torch.bfloat16) + output_weight = (1 + 0.1 * randn(_HIDDEN_SIZE)).to(torch.bfloat16) + output = torch.empty_like(prefix) + + rows = torch.cat( + [ + bank[:, :num_valid_bank_rows].float(), + prefix.unsqueeze(1).float(), + ], + dim=1, + ) + rms = torch.rsqrt(rows.square().mean(-1) + 1e-6) + scores = (rows * combine_weight.float()).sum(-1) * rms + mixed = (torch.softmax(scores, dim=-1).unsqueeze(-1) * rows).sum(1) + expected = ( + mixed + * torch.rsqrt(mixed.square().mean(-1, keepdim=True) + 1e-6) + * output_weight.float() + ) + + attn_res_fused_tma( + prefix, + bank, + combine_weight, + output_weight, + output, + num_valid_bank_rows, + 1e-6, + write_prefix=True, + ) + + torch.testing.assert_close(output.float(), expected, rtol=2e-2, atol=4e-2) + self.assertTrue(torch.equal(bank[:, num_valid_bank_rows], prefix)) + + def test_mla_output_gate(self): + generator = torch.Generator(device="cuda").manual_seed(1) + shape = (5, 12, 128) + x = torch.randn(shape, generator=generator, device="cuda", dtype=torch.bfloat16) + gate = torch.randn( + shape, generator=generator, device="cuda", dtype=torch.bfloat16 + ) + + self.assertTrue(covered(x, gate)) + expected = x * torch.sigmoid(gate).to(torch.bfloat16) + self.assertTrue(torch.equal(kimi_k3_mla_output_gate(x, gate), expected)) + + def test_situ_and_mul(self): + generator = torch.Generator(device="cuda").manual_seed(2) + hidden_size = 1024 + storage = torch.randn( + (7, 2 * hidden_size + 16), + generator=generator, + device="cuda", + dtype=torch.bfloat16, + ) + gate_up = storage[:, : 2 * hidden_size] + output = torch.empty( + (gate_up.shape[0], hidden_size), + device="cuda", + dtype=torch.bfloat16, + ) + + returned = situ_and_mul(gate_up, output, beta=_BETA, linear_beta=_LINEAR_BETA) + + self.assertIs(returned, output) + torch.testing.assert_close( + returned.float(), + _situ_reference(gate_up).to(torch.bfloat16).float(), + rtol=2e-2, + atol=4e-2, + ) + + def test_situ_mul_quant(self): + torch.cuda.manual_seed_all(3) + num_experts, num_tokens, hidden_size, topk = 8, 32, 1024, 16 + gate_up = ( + torch.randn( + num_experts, + num_tokens, + 2 * hidden_size, + device="cuda", + dtype=torch.float32, + ) + * 2.0 + ).to(torch.bfloat16) + masked_m = torch.randint( + 0, + num_tokens + 1, + (num_experts,), + device="cuda", + dtype=torch.int32, + ) + masked_m[0] = 0 + masked_m[-1] = num_tokens + + output = torch.full( + (num_experts, num_tokens, hidden_size), + 0x7F, + device="cuda", + dtype=torch.uint8, + ).view(torch.float8_e4m3fn) + num_groups = hidden_size // _GROUP_SIZE + output_scale = torch.zeros( + (num_experts, num_groups // 4, num_tokens), + device="cuda", + dtype=torch.int32, + ) + + situ_and_mul_masked_post_quant( + input=gate_up, + output=output, + output_scale=output_scale, + quant_group_size=_GROUP_SIZE, + masked_m=masked_m, + beta=_BETA, + linear_beta=_LINEAR_BETA, + scale_ue8m0=True, + topk=topk, + transposed=True, + ) + + scales = _unpack_ue8m0_scales(output_scale, num_groups) + expanded_scales = scales.repeat_interleave(_GROUP_SIZE, dim=-1) + dequantized = output.float() * expanded_scales + expected = _situ_reference(gate_up) + error_bound = expanded_scales * 17.0 + raw_output = output.view(torch.uint8) + for expert in range(num_experts): + valid_tokens = int(masked_m[expert].item()) + self.assertTrue( + bool( + ( + ( + dequantized[expert, :valid_tokens] + - expected[expert, :valid_tokens] + ).abs() + <= error_bound[expert, :valid_tokens] + ).all() + ) + ) + self.assertTrue(bool((raw_output[expert, valid_tokens:] == 0x7F).all())) + + def test_moe_front(self): + torch.manual_seed(4) + num_tokens, latent_dim = 1, 128 + hidden = ( + torch.randn( + num_tokens, + _HIDDEN_SIZE, + device="cuda", + dtype=torch.bfloat16, + ) + / 32 + ) + weight = ( + torch.randn( + NUM_EXPERTS + latent_dim, + _HIDDEN_SIZE, + device="cuda", + dtype=torch.bfloat16, + ) + / 32 + ) + bias = torch.randn(NUM_EXPERTS, device="cuda") + + weights, ids, routed = fused_front( + hidden, + weight, + bias, + latent_dim, + renormalize=True, + routed_scaling_factor=2.5, + apply_routed_scaling_factor_on_output=True, + ) + merged = torch.mm(hidden, weight.t(), out_dtype=torch.float32) + ref_weights, ref_ids = moe_fused_gate( + merged[:, :NUM_EXPERTS], + bias, + topk=TOPK, + scoring_func="sigmoid", + renormalize=True, + routed_scaling_factor=2.5, + apply_routed_scaling_factor_on_output=True, + ) + order = ids.argsort(dim=-1) + ref_order = ref_ids.argsort(dim=-1) + self.assertTrue( + torch.equal( + ids.gather(1, order), + ref_ids.to(torch.int32).gather(1, ref_order), + ) + ) + torch.testing.assert_close( + weights.gather(1, order), + ref_weights.gather(1, ref_order), + rtol=1e-6, + atol=0, + ) + self.assertTrue( + torch.equal( + routed, + merged[:, NUM_EXPERTS:].to(torch.bfloat16), + ) + ) + + def test_mtp_replayssm_ring(self): + num_requests, num_heads, num_spec, key_dim = 2, 2, 2, 128 + num_tokens = num_requests * (1 + num_spec) + num_slots, ring_size, conv_width = num_requests + 2, 16, 4 + + def run(cache_ring): + torch.manual_seed(5) + x_q = torch.randn( + 1, + num_tokens, + num_heads, + key_dim, + device="cuda", + dtype=torch.bfloat16, + ) + x_k = torch.randn_like(x_q) + x_v = torch.randn_like(x_q) + gate = torch.randn_like(x_q) + beta = torch.randn( + 1, + num_tokens, + num_heads, + device="cuda", + dtype=torch.bfloat16, + ) + conv_weight = [ + torch.randn( + num_heads * key_dim, + conv_width, + device="cuda", + ) + * 0.1 + for _ in range(3) + ] + conv_state = [ + torch.randn( + num_slots, + num_heads * key_dim, + conv_width - 1, + device="cuda", + dtype=torch.bfloat16, + ) + for _ in range(3) + ] + slots = torch.arange(1, num_requests + 1, device="cuda", dtype=torch.int32) + scratch = torch.arange(num_requests, device="cuda", dtype=torch.int32) + state = torch.randn( + num_slots, + num_heads, + key_dim, + key_dim, + device="cuda", + ) + intermediate_conv = torch.zeros( + num_requests, + 1 + num_spec, + num_heads * key_dim, + conv_width - 1, + device="cuda", + dtype=torch.bfloat16, + ) + kwargs = dict( + x_q=x_q, + x_k=x_k, + x_v=x_v, + w_q=conv_weight[0], + w_k=conv_weight[1], + w_v=conv_weight[2], + cs_q=conv_state[0], + cs_k=conv_state[1], + cs_v=conv_state[2], + g=gate, + beta=beta, + A_log=torch.randn(num_heads, device="cuda"), + dt_bias=torch.randn(num_heads * key_dim, device="cuda"), + recurrent_state=state, + intermediate_state_indices=scratch, + intermediate_conv_q=intermediate_conv.clone(), + intermediate_conv_k=intermediate_conv.clone(), + intermediate_conv_v=intermediate_conv.clone(), + ssm_state_indices=slots, + cu_seqlens=torch.arange( + 0, + num_tokens + 1, + 1 + num_spec, + device="cuda", + dtype=torch.int32, + ), + lower_bound=-5.0, + ) + if not cache_ring: + intermediate = torch.zeros( + num_requests, + 1 + num_spec, + num_heads, + key_dim, + key_dim, + device="cuda", + ) + output = fused_kda_decode_mtp_dspark( + intermediate_ssm=intermediate, + **kwargs, + ) + return output, intermediate, slots, scratch + + raw_v = torch.zeros( + num_slots, + num_heads, + ring_size, + key_dim, + device="cuda", + dtype=torch.bfloat16, + ) + raw_k = torch.zeros_like(raw_v) + ring_gate = torch.zeros( + num_slots, + num_heads, + ring_size, + key_dim, + device="cuda", + ) + ring_beta = torch.zeros( + num_slots, + num_heads, + ring_size, + device="cuda", + ) + output = fused_kda_decode_mtp_dspark( + intermediate_ssm=None, + replayssm_rawv=raw_v, + replayssm_rawk=raw_k, + replayssm_g=ring_gate, + replayssm_beta=ring_beta, + **kwargs, + ) + return ( + output, + state, + slots, + (raw_v, raw_k, ring_gate, ring_beta), + ) + + baseline, intermediate, slots, scratch = run(cache_ring=False) + ring_output, checkpoint, ring_slots, rings = run(cache_ring=True) + self.assertTrue(torch.equal(ring_output, baseline)) + commit_kda_replayssm_spec( + checkpoint, + *rings, + ring_slots, + torch.full( + (num_requests,), + 1 + num_spec, + device="cuda", + dtype=torch.int32, + ), + max_cache_len=ring_size, + num_k_heads=num_heads, + use_qk_l2norm_in_kernel=True, + null_block_id=-1, + ) + for request in range(num_requests): + expected = intermediate[scratch[request], num_spec] + actual = checkpoint[slots[request]] + relative_error = ( + actual - expected + ).abs().max() / expected.abs().max().clamp_min(1e-6) + self.assertLess(relative_error.item(), 2e-2) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/kernels/ops/test_kimi_k3_prerequisite_ops.py b/test/registered/kernels/ops/test_kimi_k3_prerequisite_ops.py new file mode 100644 index 000000000..51324a7e3 --- /dev/null +++ b/test/registered/kernels/ops/test_kimi_k3_prerequisite_ops.py @@ -0,0 +1,456 @@ +"""Representative parity coverage for the lightweight Kimi-K3 prerequisites.""" + +import unittest + +import torch + +from sglang.kernels.ops.attention.concat_mla import concat_mla_absorb_q +from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import ( + fused_sigmoid_gating_delta_rule_update, +) +from sglang.kernels.ops.attention.fla.kda_replayssm_spec_decode import ( + commit_kda_replayssm_spec, +) +from sglang.kernels.ops.attention.set_mla_kv_concat_q import ( + can_use_set_mla_kv_concat_q, + can_use_set_mla_kv_concat_q_fp8, + set_mla_kv_concat_q, + set_mla_kv_concat_q_fp8, +) +from sglang.kernels.ops.attention.utils import concat_mla_absorb_q_general +from sglang.kernels.ops.attention.vision_rope import ( + apply_fused_qk_complex_rope, +) +from sglang.kernels.ops.elementwise import add3 +from sglang.kernels.ops.gemm.tiny_gemm import ( + tiny_k_gemm_bf16, + tiny_n_gemm_bf16, +) +from sglang.kernels.ops.kvcache.set_mla_kv_buffer import set_mla_kv_buffer +from sglang.kernels.ops.mm.process.image import ( + _normalize_and_patchify_torch, + normalize_and_patchify, +) +from sglang.kernels.ops.moe import moe_route_quant_fused +from sglang.kernels.ops.moe.moe_route_radix import route_radix +from sglang.kernels.ops.moe.moe_topk_sum import moe_topk_sum +from sglang.kernels.ops.moe.pack_topk_ids import PackTopkIds +from sglang.kernels.ops.quantization.per_token_group_quant import ( + per_token_group_quant, +) +from sglang.kernels.ops.sampling.top_p_renorm_triton import ( + top_p_renorm_probs_triton, +) +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=120, stage="base-b-kernel-unit", runner_config="1-gpu-large") + +NUM_EXPERTS = 896 +TOPK = 16 +NOPE_DIM = 512 +ROPE_DIM = 64 +MLA_DIM = NOPE_DIM + ROPE_DIM +MLA_PAGES = 256 + + +def _route_oracle( + scores, bias, topk, renormalize, routed_scaling_factor, apply_scale, sorted +): + """Pure-torch fp32 reference for route_radix. + + Deliberately independent of moe_fused_gate: that entry dispatches back to + route_radix whenever scoring is sigmoid with no shared experts, no expert + groups and no softcapping (moe_fused_gate.py, the covered() fast path), which + is exactly the configuration under test. + + Contract, from route_radix.cuh: bias participates in RANKING only and the + emitted weight stays bias-free; NaN is floored so it can never win; ties go to + the lower expert id; renormalize divides by the winners' sum (guarded to 1 when + that sum is non-positive) and only then is routed scaling applied; sorted=True + emits (biased desc, id asc) while sorted=False emits ascending expert id. + """ + s = torch.sigmoid(scores.float()) + biased = s + bias.float() + biased = torch.where(torch.isnan(biased), torch.full_like(biased, -1e30), biased) + # stable + descending: equal biased values keep ascending-id order + ranked = torch.argsort(biased, dim=-1, descending=True, stable=True)[:, :topk] + w = s.gather(1, ranked) + total = w.sum(-1, keepdim=True) + norm = torch.where(total > 0, total, torch.ones_like(total)) + if renormalize: + w = w / norm + if apply_scale: + w = w * routed_scaling_factor + if sorted: + return w, ranked.to(torch.int32) + by_id = ranked.argsort(dim=-1) + return w.gather(1, by_id), ranked.gather(1, by_id).to(torch.int32) + + +def _make_mla_inputs(batch_size, num_heads, seed): + generator = torch.Generator(device="cuda").manual_seed(seed) + + def randn(*shape): + return ( + torch.randn(*shape, generator=generator, device="cuda", dtype=torch.float32) + .mul(0.1) + .to(torch.bfloat16) + ) + + pool = randn(MLA_PAGES, MLA_DIM) + latent = randn(batch_size, MLA_DIM) + query = randn(batch_size, num_heads, MLA_DIM) + loc = torch.randperm(MLA_PAGES, generator=generator, device="cuda")[:batch_size].to( + torch.int64 + ) + return ( + pool, + loc, + latent[:, :NOPE_DIM], + latent[:, NOPE_DIM:], + query[..., :NOPE_DIM], + query[..., NOPE_DIM:], + ) + + +class TestKimiK3PrerequisiteOps(CustomTestCase): + def test_mla_scatter_concat_bf16_and_fp8(self): + batch_size, num_heads = 64, 8 + pool, loc, k_nope, k_rope, q_nope, q_rope = _make_mla_inputs( + batch_size, num_heads, seed=0 + ) + + if not can_use_set_mla_kv_concat_q(NOPE_DIM * 2, ROPE_DIM * 2): + self.skipTest("fused MLA scatter+concat requires SM90+") + pool_ref = pool.clone() + query = set_mla_kv_concat_q(pool, loc, k_nope, k_rope, q_nope, q_rope) + set_mla_kv_buffer(pool_ref, loc, k_nope, k_rope) + query_ref = concat_mla_absorb_q(q_nope, q_rope) + self.assertTrue(torch.equal(pool, pool_ref)) + self.assertTrue(torch.equal(query, query_ref)) + + if not can_use_set_mla_kv_concat_q_fp8(): + self.skipTest("fused FP8 MLA scatter+concat requires SM90+") + fp8_pool = torch.zeros( + MLA_PAGES, MLA_DIM, device="cuda", dtype=torch.float8_e4m3fn + ) + fp8_ref = fp8_pool.clone() + fp8_query = set_mla_kv_concat_q_fp8( + fp8_pool, loc, k_nope, k_rope, q_nope, q_rope + ) + row = torch.cat([k_nope, k_rope], dim=-1).to(torch.float8_e4m3fn) + fp8_ref[loc] = row + fp8_query_ref = concat_mla_absorb_q_general(q_nope, q_rope).to( + torch.float8_e4m3fn + ) + self.assertTrue( + torch.equal(fp8_pool.view(torch.uint8), fp8_ref.view(torch.uint8)) + ) + self.assertTrue( + torch.equal( + fp8_query.view(torch.uint8), + fp8_query_ref.view(torch.uint8), + ) + ) + + def test_replayssm_ring_fold(self): + batch_size, num_steps = 8, 4 + num_value_heads, num_key_heads = 8, 2 + key_dim = value_dim = 128 + ring_size = 16 + torch.manual_seed(6) + + q = torch.randn( + batch_size, + num_steps, + num_key_heads, + key_dim, + device="cuda", + ) + k = torch.randn_like(q) + v = torch.randn( + batch_size, + num_steps, + num_value_heads, + value_dim, + device="cuda", + ) + a = torch.randn( + batch_size, + num_steps, + num_value_heads, + key_dim, + device="cuda", + ) + b = torch.randn(batch_size, num_steps, num_value_heads, device="cuda") + a_log = torch.randn(num_value_heads, device="cuda") + dt_bias = torch.randn(num_value_heads, key_dim, device="cuda") + slots = torch.arange(1, batch_size + 1, device="cuda", dtype=torch.int32) + slots[-1] = -1 + num_slots = batch_size + 1 + state = torch.randn( + num_slots, + num_value_heads, + value_dim, + key_dim, + device="cuda", + ) + intermediate = torch.zeros( + num_slots, + num_steps, + num_value_heads, + value_dim, + key_dim, + device="cuda", + ) + raw_v = torch.zeros( + num_slots, + num_value_heads, + ring_size, + value_dim, + device="cuda", + ) + raw_k = torch.zeros( + num_slots, + num_key_heads, + ring_size, + key_dim, + device="cuda", + ) + gate = torch.zeros_like(raw_v) + beta = torch.zeros( + num_slots, + num_value_heads, + ring_size, + device="cuda", + ) + + fused_sigmoid_gating_delta_rule_update( + A_log=a_log, + a=a, + dt_bias=dt_bias, + softplus_beta=1.0, + softplus_threshold=20.0, + q=q, + k=k, + v=v, + b=b, + initial_state_source=state, + initial_state_indices=slots, + scale=key_dim**-0.5, + use_qk_l2norm_in_kernel=True, + is_kda=True, + lower_bound=-5.0, + disable_state_update=True, + intermediate_states_buffer=intermediate, + intermediate_state_indices=slots, + cache_steps=num_steps, + cache_ring=True, + replayssm_rawv=raw_v, + replayssm_rawk=raw_k, + replayssm_g=gate, + replayssm_beta=beta, + ) + checkpoint = state.clone() + commit_kda_replayssm_spec( + checkpoint, + raw_v, + raw_k, + gate, + beta, + slots, + torch.full((batch_size,), num_steps, device="cuda", dtype=torch.int32), + max_cache_len=ring_size, + num_k_heads=num_key_heads, + use_qk_l2norm_in_kernel=True, + null_block_id=-1, + ) + for slot in slots[:-1].tolist(): + expected = intermediate[slot, num_steps - 1] + actual = checkpoint[slot] + relative_error = ( + actual - expected + ).abs().max() / expected.abs().max().clamp_min(1e-6) + self.assertLess(relative_error.item(), 1e-3) + + def test_add3_bit_exact(self): + torch.manual_seed(0) + tensors = [ + torch.randn(9, 112, device="cuda", dtype=torch.bfloat16) for _ in range(3) + ] + actual = add3.add3(*tensors, prefetch_bc=True) + expected = (tensors[0] + tensors[1]) + tensors[2] + self.assertTrue(torch.equal(actual, expected)) + + def test_moe_auxiliary_kernels(self): + x = torch.randn(2, TOPK, 7168, device="cuda", dtype=torch.bfloat16) + out = torch.empty(2, 7168, device="cuda", dtype=torch.bfloat16) + self.assertIs(moe_topk_sum(x, out), out) + self.assertTrue(torch.equal(out, x.float().sum(1).to(torch.bfloat16))) + + def test_moe_route_and_quant(self): + torch.manual_seed(1) + scores = torch.randn(8, NUM_EXPERTS, device="cuda", dtype=torch.bfloat16) + bias = torch.randn(NUM_EXPERTS, device="cuda", dtype=torch.float32) + args = (scores, bias, TOPK, True, 2.5, True) + weights, ids = route_radix(*args, sorted=True) + # Oracle, NOT moe_fused_gate: for this exact configuration (sigmoid, no + # shared experts, no expert groups, no softcapping) moe_fused_gate + # dispatches straight back to route_radix, so using it as the reference + # compares the kernel with itself and cannot see a selection, tie-break, + # NaN, renormalize or scaling error. + ref_weights, ref_ids = _route_oracle(*args, sorted=True) + self.assertTrue(torch.equal(ids, ref_ids)) + # rtol is not 1e-6: the kernel computes sigmoid with __fdividef/__expf, + # whose last bits differ from torch's. The old self-comparison could + # afford atol=0; a real oracle cannot. + torch.testing.assert_close(weights, ref_weights, rtol=1e-5, atol=1e-6) + + if not moe_route_quant_fused.available(): + self.skipTest("fused route+quant kernel unavailable") + hidden = torch.randn(8, 3584, device="cuda", dtype=torch.bfloat16) + ref_weights, ref_ids = route_radix(*args, sorted=False) + ref_packed = PackTopkIds.execute(ref_ids, ref_weights) + ref_q, ref_scale = per_token_group_quant( + hidden, group_size=32, scale_ue8m0=True + ) + actual = moe_route_quant_fused.route_quant_fused( + scores, + bias, + hidden, + TOPK, + renormalize=True, + routed_scaling_factor=2.5, + apply_scale=True, + ) + weights, ids, packed, quantized, scale = actual + self.assertTrue(torch.equal(ids, ref_ids)) + self.assertTrue( + torch.equal(weights.view(torch.int32), ref_weights.view(torch.int32)) + ) + self.assertTrue(torch.equal(packed, ref_packed)) + self.assertTrue( + torch.equal(quantized.view(torch.uint8), ref_q.view(torch.uint8)) + ) + torch.testing.assert_close(scale, ref_scale, rtol=0, atol=0) + + def test_route_radix_ties_and_nan(self): + """The cases the self-comparison could not see. + + Exact ties: many experts share one biased value, so the winner set is only + determined by the lowest-id rule. NaN: floored, so a NaN expert must never + be selected while enough finite ones exist. Both run with renormalize and + scaling on and off, since those are applied in a fixed order. + """ + bias = torch.zeros(NUM_EXPERTS, device="cuda", dtype=torch.float32) + + tied = torch.full((4, NUM_EXPERTS), 0.25, device="cuda", dtype=torch.bfloat16) + # a handful of strict winners above the tied plateau, the rest exactly equal + tied[:, 300] = 2.0 + tied[:, 7] = 2.0 + tied[:, 800] = 1.5 + + nan_scores = torch.randn(4, NUM_EXPERTS, device="cuda", dtype=torch.bfloat16) + nan_scores[:, 100] = float("nan") + nan_scores[:, 500] = float("nan") + # make the NaN experts the ones that would otherwise win outright + nan_scores[:, 101] = 5.0 + + for name, scores in (("ties", tied), ("nan", nan_scores)): + for renormalize in (False, True): + for apply_scale in (False, True): + for sorted_ in (False, True): + args = (scores, bias, TOPK, renormalize, 2.5, apply_scale) + ids = route_radix(*args, sorted=sorted_)[1] + ref_ids = _route_oracle(*args, sorted=sorted_)[1] + tag = ( + f"{name} renorm={renormalize} " + f"scale={apply_scale} sorted={sorted_}" + ) + self.assertTrue(torch.equal(ids, ref_ids), msg=tag) + if name == "nan": + self.assertFalse( + bool(((ids == 100) | (ids == 500)).any()), + msg=f"{tag}: a NaN expert was selected", + ) + + def test_tiny_gemm_variants(self): + torch.manual_seed(2) + x = torch.randn(2, 7168, device="cuda", dtype=torch.bfloat16) / 8 + weight = torch.randn(144, 7168, device="cuda", dtype=torch.bfloat16) / 8 + actual = tiny_n_gemm_bf16(x, weight, out_dtype=torch.float32) + torch.testing.assert_close( + actual.double(), x.double() @ weight.double().t(), rtol=1e-3, atol=1e-3 + ) + + x = torch.randn(7, 128, device="cuda", dtype=torch.bfloat16) / 4 + weight = torch.randn(1536, 128, device="cuda", dtype=torch.bfloat16) / 4 + actual = tiny_k_gemm_bf16(x, weight) + torch.testing.assert_close( + actual.double(), x.double() @ weight.double().t(), rtol=2e-2, atol=2e-2 + ) + + def test_top_p_renorm(self): + torch.manual_seed(3) + probs = torch.randn(3, 1024, device="cuda").softmax(-1) + top_p = torch.tensor([0.5, 0.8, 0.95], device="cuda") + sorted_probs = probs.sort(-1).values + cutoff = torch.searchsorted( + sorted_probs.cumsum(-1), (1 - top_p).unsqueeze(1) + ).squeeze(1) + cutoff.clamp_(max=probs.shape[1] - 1) + pivot = sorted_probs.gather(1, cutoff[:, None]) + expected = torch.where(probs >= pivot, probs, 0) + expected /= expected.sum(-1, keepdim=True) + torch.testing.assert_close( + top_p_renorm_probs_triton(probs, top_p), + expected, + rtol=2e-6, + atol=1e-8, + ) + + def test_vision_rope(self): + torch.manual_seed(4) + qkv = torch.randn(480, 3, 12, 128, device="cuda", dtype=torch.bfloat16) + q, k, _ = qkv.unbind(1) + angles = torch.randn(480, 64, device="cuda") + freqs = torch.polar(torch.ones_like(angles), angles) + freqs_expanded = freqs.unsqueeze(-2) + + def reference(x): + value = torch.view_as_complex(x.float().view(*x.shape[:-1], -1, 2)) + return torch.view_as_real(value * freqs_expanded).flatten(-2).type_as(x) + + actual_q, actual_k = apply_fused_qk_complex_rope(q, k, freqs) + atol = 2 * torch.finfo(torch.bfloat16).eps + torch.testing.assert_close(actual_q, reference(q), rtol=0, atol=atol) + torch.testing.assert_close(actual_k, reference(k), rtol=0, atol=atol) + + def test_normalize_and_patchify(self): + torch.manual_seed(5) + image = torch.randn(2, 3, 17, 19, device="cuda") + scale = torch.randn(1, 3, 1, 1, device="cuda") + bias = torch.randn(1, 3, 1, 1, device="cuda") + args = (image, scale, bias, 4, 20, 20) + actual = normalize_and_patchify( + args[0], + args[1], + args[2], + patch_size=args[3], + padded_height=args[4], + padded_width=args[5], + ) + expected = _normalize_and_patchify_torch( + args[0], + args[1], + args[2], + patch_size=args[3], + padded_height=args[4], + padded_width=args[5], + ) + torch.testing.assert_close(actual, expected, rtol=1e-2, atol=1e-2) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_kda_fused_decode_strided_state.py b/test/registered/unit/mem_cache/test_kda_fused_decode_strided_state.py new file mode 100644 index 000000000..86a48f469 --- /dev/null +++ b/test/registered/unit/mem_cache/test_kda_fused_decode_strided_state.py @@ -0,0 +1,230 @@ +"""NV KDA fused-decode kernel gate + slot addressing vs envelope-strided SSM +pools (CPU). + +Derived property under test: the fully-fused KDA decode kernel +(``kda_fused_decode``) addresses the ssm/temporal state pool by the slot pitch +the pool actually reports (``ssm_states.stride(0)``), NOT the dense ``HV*V*K`` +pitch. Under ``--enable-unified-memory`` / ``--enable-page-major-kv-layout`` the +per-layer temporal view is envelope-strided: one slot pitches across ALL layers +(56,171,520 B on K3), so a hardcoded ``slot*HV*V*K`` offset mis-addresses every +slot > 0 (the exact chunk_delta_h hardcoded-pitch bug pattern, GSM8K 0.17). + +Two things are pinned here (both CPU-checkable without the CUDA kernel): + + 1. ``covered()`` ACCEPTS the envelope-strided view. Pre-fix the gate did + ``ssm_states.view(-1, HV, V, K).is_contiguous()``, which is False on a + non-dense slot pitch, so decode silently dropped to the slower unfused + chain. Reverting to that ``.view(...)`` gate turns test (1) red. The gate + still REJECTS a view whose inner ``[HV, V, K]`` is non-contiguous (the + one contract the float4 state loads rely on). + + 2. The kernel's reconstructed slot-addressing formula + ``base + slot*stride(0) + i_hv*V*K + v*K + k`` resolves to the exact same + storage element as torch's native ``ssm_states[slot, i_hv, v, k]`` on the + strided pool, while the pre-fix dense-pitch formula + ``base + slot*(HV*V*K) + ...`` resolves ELSEWHERE for every slot > 0. + Hardcoding the dense pitch back into the ``.cuh`` turns test (2) red. + +Runs on CPU: only the Python gate and the (dtype/stride-only) addressing +arithmetic execute — no CUDA kernel is launched. + + python -m pytest test/registered/unit/mem_cache/test_kda_fused_decode_strided_state.py -v +""" + +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=6, suite="base-a-test-cpu") + +import unittest + +import torch + +from sglang.kernels.ops.attention.kda_fused_decode import ( + _CONV_STATE_W, + covered, +) +from sglang.srt.mem_cache.layout.page_major import ( + build_page_major_mamba_views, + mamba_entry_bytes, +) + +_DEV = "cpu" + +# The kernel is compiled for the K3 KDA decode regime; covered() enforces these +# supported local head counts. Multi-layer + several slots so the envelope slot +# pitch differs from the dense H*V*K pitch. +_KDA_HEADS = (3, 6, 12) +_V = 128 +_K = 128 +_LAYERS = 3 +_LAYER_UNDER_TEST = 1 +_SLOTS = 6 +_CONV_SHAPES = ((3, 8),) # tiny bf16 conv region interleaves the temporal region +_CONV_DTYPE = torch.bfloat16 +_TEMPORAL_DTYPE = torch.float32 + + +def _seg(heads: int) -> int: + return heads * _V + + +def _conv_dim(heads: int) -> int: + return 3 * _seg(heads) + + +def _make_strided_temporal_view(heads: int): + """Envelope-strided temporal (SSM) view as UnifiedMambaPool / the + page-major MambaPool serve it to the KDA backend: shape + ``(num_layers, max_slots, H, V, K)`` with the slot pitch spanning ALL + layers' state, not H*V*K.""" + entry = mamba_entry_bytes( + layer_num=_LAYERS, + conv_state_shapes=_CONV_SHAPES, + conv_dtype=_CONV_DTYPE, + temporal_state_shape=(heads, _V, _K), + temporal_dtype=_TEMPORAL_DTYPE, + ) + raw = torch.zeros(_SLOTS * entry, dtype=torch.uint8, device=_DEV) + _conv_views, temporal = build_page_major_mamba_views( + raw, + layer_num=_LAYERS, + conv_state_shapes=_CONV_SHAPES, + conv_dtype=_CONV_DTYPE, + temporal_state_shape=(heads, _V, _K), + temporal_dtype=_TEMPORAL_DTYPE, + max_slots=_SLOTS, + ) + return raw, temporal + + +def _make_covered_side_args(batch: int, heads: int): + """The non-ssm covered() arguments, in the exact K3 shapes/dtypes so the + gate turns solely on the ssm_states view under test.""" + bf16 = torch.bfloat16 + seg = _seg(heads) + conv_dim = _conv_dim(heads) + mixed_qkv = torch.zeros((batch, conv_dim), dtype=bf16, device=_DEV) + a = torch.zeros((batch, seg), dtype=bf16, device=_DEV) + b = torch.zeros((batch, heads), dtype=bf16, device=_DEV) + onorm_g = torch.zeros((batch, seg), dtype=bf16, device=_DEV) + conv_states = torch.zeros( + (_SLOTS, _CONV_STATE_W, conv_dim), dtype=bf16, device=_DEV + ) + cache_indices = torch.zeros((batch,), dtype=torch.int32, device=_DEV) + return mixed_qkv, a, b, conv_states, onorm_g, cache_indices + + +def _addressing_samples(heads: int): + return [ + (0, 0, 0, 0), + (5, heads - 1, 127, 127), + (3, heads // 2, 64, 100), + (1, 0, 0, 1), + (2, min(heads - 1, 2), 3, 7), + ] + + +class TestKdaFusedDecodeStridedState(unittest.TestCase): + def test_covered_accepts_envelope_strided_and_rejects_noncontiguous_inner(self): + for heads in _KDA_HEADS: + with self.subTest(kda_heads=heads): + _raw, temporal = _make_strided_temporal_view(heads) + ssm = temporal[_LAYER_UNDER_TEST] # what mamba2_layer_cache serves + + # Precondition: the pool really is envelope-strided (else the + # property below would be vacuous — a dense pool passes the old + # gate too). + self.assertNotEqual( + ssm.stride(0), + heads * _V * _K, + "test setup no longer produces a strided pool", + ) + # Inner [H, V, K] IS contiguous — the contract the kernel's + # float4 state loads rely on and all covered() must still require. + self.assertEqual( + (ssm.stride(-1), ssm.stride(-2), ssm.stride(-3)), + (1, _K, _V * _K), + ) + + ( + mixed_qkv, + a, + b, + conv_states, + onorm_g, + cache_indices, + ) = _make_covered_side_args(batch=2, heads=heads) + + # (1) Accept the envelope-strided view — pre-fix + # .view(...).is_contiguous() would reject this and drop decode + # to the unfused chain. + self.assertTrue( + covered(mixed_qkv, a, b, conv_states, ssm, cache_indices, onorm_g), + "covered() rejected the envelope-strided ssm pool (fused decode " + "would silently fall back to the unfused chain)", + ) + + # Still reject a view whose inner dims are NOT contiguous: the + # kernel cannot float4-load a transposed [.., V, K] state. + ssm_bad = ssm.transpose(-1, -2) # stride(-1) == K, not 1 + self.assertFalse( + covered( + mixed_qkv, + a, + b, + conv_states, + ssm_bad, + cache_indices, + onorm_g, + ), + "covered() must reject a non-inner-contiguous ssm view", + ) + + def test_slot_formula_resolves_correct_element_and_dense_pitch_misaddresses(self): + for heads in _KDA_HEADS: + with self.subTest(kda_heads=heads): + raw, temporal = _make_strided_temporal_view(heads) + ssm = temporal[_LAYER_UNDER_TEST] + + # Distinct value per storage element so an offset that lands + # elsewhere reads a provably different value. + raw_fp32 = raw.view(torch.float32) + raw_fp32.copy_(torch.arange(raw_fp32.numel(), dtype=torch.float32)) + + base = ssm.storage_offset() + # == state.stride(0) the wrapper passes the kernel. + slot_stride = ssm.stride(0) + dense_pitch = heads * _V * _K # the pre-fix hardcoded slot pitch + + for slot, i_hv, v, k in _addressing_samples(heads): + intra = ( + i_hv * (_V * _K) + v * _K + k + ) # kernel's hardcoded intra-slot offset + kernel_off = base + slot * slot_stride + intra + + # (2a) The kernel formula names exactly the element torch + # indexing names — proves slot*stride(0) + intra addresses + # the intended slot. + self.assertEqual( + raw_fp32[kernel_off].item(), + ssm[slot, i_hv, v, k].item(), + f"kernel slot formula mis-addressed " + f"(slot={slot}, i_hv={i_hv}, heads={heads})", + ) + + # (2b) The pre-fix dense-pitch formula lands on a DIFFERENT + # element (a different layer's envelope region) for every + # slot > 0. + dense_off = base + slot * dense_pitch + intra + if slot > 0: + self.assertNotEqual( + raw_fp32[dense_off].item(), + ssm[slot, i_hv, v, k].item(), + f"dense-pitch formula happened to match at " + f"slot={slot}, heads={heads}; the fix would not " + "be load-bearing", + ) + + +if __name__ == "__main__": + unittest.main()