diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/c128_v2.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/c128_v2.cuh index e625919e1..31353e6a1 100644 --- a/python/sglang/jit_kernel/csrc/deepseek_v4/c128_v2.cuh +++ b/python/sglang/jit_kernel/csrc/deepseek_v4/c128_v2.cuh @@ -348,7 +348,7 @@ struct FlashCompress128Kernel { auto N = SymbolicSize{"batch_size"}; auto device_ = SymbolicDevice{}; - device_.set_options(); + device_.set_options(); TensorMatcher({-1, 128, Trait::kElementSize}) // kv score .with_dtype() @@ -395,7 +395,7 @@ struct FlashCompress128Kernel { auto C = SymbolicSize{"num_c_plans"}; auto W = SymbolicSize{"num_w_plans"}; auto device_ = SymbolicDevice{}; - device_.set_options(); + device_.set_options(); TensorMatcher({-1, 128, Trait::kElementSize}) // kv score .with_dtype() diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/c4_v2.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/c4_v2.cuh index b970e3847..efa9f0510 100644 --- a/python/sglang/jit_kernel/csrc/deepseek_v4/c4_v2.cuh +++ b/python/sglang/jit_kernel/csrc/deepseek_v4/c4_v2.cuh @@ -309,7 +309,7 @@ struct FlashCompress4Kernel { auto N = SymbolicSize{"batch_size"}; auto device_ = SymbolicDevice{}; - device_.set_options(); + device_.set_options(); TensorMatcher({-1, 4, Trait::kElementSize}) // kv score .with_dtype() @@ -356,7 +356,7 @@ struct FlashCompress4Kernel { auto C = SymbolicSize{"num_c_plans"}; auto W = SymbolicSize{"num_w_plans"}; auto device_ = SymbolicDevice{}; - device_.set_options(); + device_.set_options(); TensorMatcher({-1, 4, Trait::kElementSize}) // kv score .with_dtype() diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/c_plan.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/c_plan.cuh index 4518dcffa..3e4aaaf5f 100644 --- a/python/sglang/jit_kernel/csrc/deepseek_v4/c_plan.cuh +++ b/python/sglang/jit_kernel/csrc/deepseek_v4/c_plan.cuh @@ -104,7 +104,11 @@ SGL_DEVICE uint32_t warp_inclusive_sum(uint32_t lane_id, uint32_t val) { static_assert(device::kWarpThreads == 32); #pragma unroll for (uint32_t offset = 1; offset < 32; offset *= 2) { - uint32_t n = __shfl_up_sync(0xFFFFFFFF, val, offset); +#ifndef USE_ROCM + uint32_t n = __shfl_up_sync(device::kFullMask, val, offset); +#else + uint32_t n = __shfl_up(val, offset, 32); +#endif if (lane_id >= offset) val += n; } return val; @@ -115,7 +119,11 @@ SGL_DEVICE uint32_t warp_inclusive_sum(uint32_t lane_id, uint32_t val) { SGL_DEVICE uint32_t warp_reduce_max_u32(uint32_t val) { #pragma unroll for (uint32_t mask = 16; mask > 0; mask >>= 1) { - val = max(val, __shfl_xor_sync(0xFFFFFFFF, val, mask, 32)); +#ifndef USE_ROCM + val = max(val, __shfl_xor_sync(device::kFullMask, val, mask, 32)); +#else + val = max(val, __shfl_xor(val, mask, 32)); +#endif } return val; } @@ -123,7 +131,11 @@ SGL_DEVICE uint32_t warp_reduce_max_u32(uint32_t val) { SGL_DEVICE uint32_t warp_reduce_min_u32(uint32_t val) { #pragma unroll for (uint32_t mask = 16; mask > 0; mask >>= 1) { - val = min(val, __shfl_xor_sync(0xFFFFFFFF, val, mask, 32)); +#ifndef USE_ROCM + val = min(val, __shfl_xor_sync(device::kFullMask, val, mask, 32)); +#else + val = min(val, __shfl_xor(val, mask, 32)); +#endif } return val; } @@ -452,8 +464,8 @@ inline PrefillPlan plan_compress_prefill( auto N = SymbolicSize{"num_q_tokens"}; auto cpu_or_gpu = SymbolicDevice{}; auto device_ = SymbolicDevice{}; - cpu_or_gpu.set_options(); - device_.set_options(); + cpu_or_gpu.set_options(); + device_.set_options(); TensorMatcher({B}) // .with_dtype() @@ -499,7 +511,7 @@ inline PrefillPlan plan_compress_prefill( constexpr int32_t kMaxMTPDraftTokens = 4; const auto mtp_pad = std::min(ring_size - compress_ratio, kMaxMTPDraftTokens); - if (cpu_or_gpu.unwrap().device_type == kDLCUDA) { + if (cpu_or_gpu.unwrap().device_type == kDLGPU) { // GPU input path: kernel0 builds the (CPU-loop-equivalent) plan metadata directly // on device, padding to num_q_tokens with invalid; kernel_1 then finalizes the // SWA-translated read/write locations. Used for MTP / cuda-graph capture where @@ -628,7 +640,7 @@ inline tvm::ffi::Tensor plan_compress_decode( const int32_t ring_size) { auto B = SymbolicSize{"batch_size"}; auto device_ = SymbolicDevice{}; - device_.set_options(); + device_.set_options(); TensorMatcher({B}) // .with_dtype() @@ -691,7 +703,7 @@ inline PrefillPlan plan_compress_prefill_legacy( const bool use_cuda_graph) { auto B = SymbolicSize{"batch_size"}; auto device_ = SymbolicDevice{}; - device_.set_options(); + device_.set_options(); TensorMatcher({B}) // .with_dtype() @@ -794,7 +806,7 @@ inline tvm::ffi::Tensor plan_compress_decode_legacy( const int32_t compress_ratio) { auto B = SymbolicSize{"batch_size"}; auto device_ = SymbolicDevice{}; - device_.set_options(); + device_.set_options(); TensorMatcher({B}) // .with_dtype() diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/fused_norm_rope_v2.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/fused_norm_rope_v2.cuh index 653ff8750..811dc41f1 100644 --- a/python/sglang/jit_kernel/csrc/deepseek_v4/fused_norm_rope_v2.cuh +++ b/python/sglang/jit_kernel/csrc/deepseek_v4/fused_norm_rope_v2.cuh @@ -163,7 +163,11 @@ INDEXER_KERNEL void fused_norm_rope_indexer(const __grid_constant__ FusedNormRop for (uint32_t mask = 1; mask < kWarpThreads; mask <<= 1) { #pragma unroll for (int i = 0; i < kVecSize; ++i) { - const float other = __shfl_xor_sync(0xFFFFFFFFu, data[i], mask, kWarpThreads); +#ifndef USE_ROCM + const float other = __shfl_xor_sync(kFullMask, data[i], mask, kWarpThreads); +#else + const float other = __shfl_xor(data[i], mask, kWarpThreads); +#endif data[i] = (lane_id & mask) ? (other - data[i]) : (data[i] + other); } } @@ -307,8 +311,10 @@ FLASHMLA_KERNEL void fused_norm_rope_flashmla(const __grid_constant__ FusedNormR reinterpret_cast(rope_ptr)[lane_id] = result; } else { // Non-rope warp: per-warp UE8M0 group (64 elems -> 64 fp8 + 1 scale byte). - const auto x = data[0]; - const auto y = data[1]; + // BF16 round-trip to match the precision of the non-fused path + // (which goes through quant_to_nope_fp8_rope_bf16_pack_triton with bf16 input). + const auto x = cast(cast(data[0])); + const auto y = cast(cast(data[1])); const auto abs_max = warp::reduce_max(fmaxf(fabs(x), fabs(y))); const auto scale_raw = fmaxf(1e-4f, abs_max) / math::FP8_E4M3_MAX; const auto scale_ue8m0 = cast_to_ue8m0(scale_raw); @@ -359,7 +365,7 @@ struct FusedNormRopeKernel { auto N = SymbolicSize{"num_tokens"}; auto device_ = SymbolicDevice{}; - device_.set_options(); + device_.set_options(); TensorMatcher({N, kHeadDim}) // input .with_dtype() diff --git a/python/sglang/jit_kernel/dsv4/attn.py b/python/sglang/jit_kernel/dsv4/attn.py index aa8f24b4f..87a265ecf 100644 --- a/python/sglang/jit_kernel/dsv4/attn.py +++ b/python/sglang/jit_kernel/dsv4/attn.py @@ -7,6 +7,7 @@ import triton.language as tl from sglang.jit_kernel.utils import ( cache_once, is_arch_support_pdl, + is_hip_runtime, load_jit, make_cpp_args, ) @@ -58,13 +59,18 @@ def fused_store_cache( page_size: int, type: Literal["flashmla", "indexer"], ) -> None: - module = _jit_fused_store_module( - name=type, - input_dtype=input.dtype, - index_dtype=indices.dtype, - page_size=page_size, - ) - module.run(input, cache, indices) + if is_hip_runtime(): + from sglang.jit_kernel.triton_store_cache import triton_fused_store_cache + + triton_fused_store_cache(input, cache, indices, page_size=page_size, type=type) + else: + module = _jit_fused_store_module( + name=type, + input_dtype=input.dtype, + index_dtype=indices.dtype, + page_size=page_size, + ) + module.run(input, cache, indices) @triton.jit diff --git a/python/sglang/jit_kernel/dsv4/elementwise.py b/python/sglang/jit_kernel/dsv4/elementwise.py index d2bbb8f3c..b721c841d 100644 --- a/python/sglang/jit_kernel/dsv4/elementwise.py +++ b/python/sglang/jit_kernel/dsv4/elementwise.py @@ -132,10 +132,27 @@ def fused_q_indexer_rope_hadamard_quant( weights_out = torch.empty( (*q_input.shape[:-1], 1), dtype=torch.float32, device=q_input.device ) - module = _jit_main_q_indexer_rope_hadamard_quant_module(q_input.dtype) - module.forward( - q_input, q_fp8, weight, weights_out, float(weight_scale), freqs_real, positions - ) + if _is_hip: + torch.ops.sgl_kernel.dsv4_fused_q_indexer_rope_hadamard_quant( + q_input, + q_fp8, + weight, + weights_out, + float(weight_scale), + freqs_real, + positions, + ) + else: + module = _jit_main_q_indexer_rope_hadamard_quant_module(q_input.dtype) + module.forward( + q_input, + q_fp8, + weight, + weights_out, + float(weight_scale), + freqs_real, + positions, + ) return q_fp8, weights_out diff --git a/python/sglang/jit_kernel/dsv4/gemm.py b/python/sglang/jit_kernel/dsv4/gemm.py index c89ca32a1..da60eccec 100644 --- a/python/sglang/jit_kernel/dsv4/gemm.py +++ b/python/sglang/jit_kernel/dsv4/gemm.py @@ -14,11 +14,11 @@ _linear_bf16_fp32_algo = envs.SGLANG_OPT_BF16_FP32_GEMM_ALGO.get() def linear_bf16_fp32(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: - if _linear_bf16_fp32_algo == "deep_gemm": + if _use_aiter: + return tgemm.mm(x, y, otype=x.dtype).float() + elif _linear_bf16_fp32_algo == "deep_gemm": z = torch.empty(x.size(0), y.size(0), dtype=torch.float32, device=x.device) deep_gemm_wrapper.gemm_nt_bf16bf16f32(x, y, z) return z - elif _use_aiter: - return tgemm.mm(x, y, otype=torch.float32) else: return torch.mm(x, y.t(), out_dtype=torch.float32) diff --git a/python/sglang/jit_kernel/dsv4/moe.py b/python/sglang/jit_kernel/dsv4/moe.py index a5372c781..931805e44 100644 --- a/python/sglang/jit_kernel/dsv4/moe.py +++ b/python/sglang/jit_kernel/dsv4/moe.py @@ -5,6 +5,7 @@ import torch from sglang.jit_kernel.utils import ( cache_once, is_arch_support_pdl, + is_hip_runtime, load_jit, make_cpp_args, ) @@ -114,25 +115,37 @@ def hash_topk( scoring_func: str = "sqrtsoftplus", ) -> Tuple[torch.Tensor, torch.Tensor]: assert scoring_func == "sqrtsoftplus" - num_tokens = router_logits.size(0) - topk_routed = tid2eid.size(1) - topk_fused = topk_routed + num_fused_shared_experts - topk_ids = torch.empty( - (num_tokens, topk_fused), dtype=torch.int32, device=router_logits.device - ) - topk_weights = torch.empty( - (num_tokens, topk_fused), dtype=torch.float32, device=router_logits.device - ) - module = _jit_hash_topk_module() - module.hash_topk( - router_logits, - input_ids, - tid2eid, - topk_weights, - topk_ids, - routed_scaling_factor, - ) - return topk_weights, topk_ids + if is_hip_runtime(): + from sglang.jit_kernel.triton.hash_topk import hash_topk_triton + + return hash_topk_triton( + router_logits, + input_ids, + tid2eid, + num_fused_shared_experts, + routed_scaling_factor, + scoring_func, + ) + else: + num_tokens = router_logits.size(0) + topk_routed = tid2eid.size(1) + topk_fused = topk_routed + num_fused_shared_experts + topk_ids = torch.empty( + (num_tokens, topk_fused), dtype=torch.int32, device=router_logits.device + ) + topk_weights = torch.empty( + (num_tokens, topk_fused), dtype=torch.float32, device=router_logits.device + ) + module = _jit_hash_topk_module() + module.hash_topk( + router_logits, + input_ids, + tid2eid, + topk_weights, + topk_ids, + routed_scaling_factor, + ) + return topk_weights, topk_ids def mega_moe_pre_dispatch( diff --git a/python/sglang/jit_kernel/dsv4/topk.py b/python/sglang/jit_kernel/dsv4/topk.py index 3c7d14769..a27245186 100644 --- a/python/sglang/jit_kernel/dsv4/topk.py +++ b/python/sglang/jit_kernel/dsv4/topk.py @@ -7,6 +7,7 @@ import torch from sglang.jit_kernel.utils import ( cache_once, is_arch_support_pdl, + is_hip_runtime, load_jit, make_cpp_args, ) @@ -48,10 +49,15 @@ def topk_transform_512( page_size: int, out_raw_indices: Optional[torch.Tensor] = None, ) -> None: - module = _jit_topk_v1_module(out_page_indices.shape[1]) - module.topk_transform( - scores, seq_lens, page_tables, out_page_indices, page_size, out_raw_indices - ) + if is_hip_runtime(): + torch.ops.sgl_kernel.deepseek_v4_topk_transform_512( + scores, seq_lens, page_tables, out_page_indices, page_size, out_raw_indices + ) + else: + module = _jit_topk_v1_module(out_page_indices.shape[1]) + module.topk_transform( + scores, seq_lens, page_tables, out_page_indices, page_size, out_raw_indices + ) _WORKSPACE_INTS_PER_BATCH = 2 + 1024 * 2 diff --git a/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/fp8_utils.cuh b/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/fp8_utils.cuh index 4fdbb062c..53a62755b 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/fp8_utils.cuh +++ b/python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/fp8_utils.cuh @@ -5,7 +5,9 @@ #include #include +#ifndef USE_ROCM #include +#endif // Small helpers shared by the DeepSeek-V4 FP8/UE8M0 quantization kernels // (silu_and_mul_masked_post_quant, store, mega_moe_pre_dispatch, ...). @@ -30,14 +32,81 @@ SGL_DEVICE float inv_scale_ue8m0(int32_t exp) { } // Clamp to [-FP8_E4M3_MAX, FP8_E4M3_MAX]. +// Uses platform-specific max from type.cuh (448 for E4M3FN, 224 for E4M3FNUZ). SGL_DEVICE float fp8_e4m3_clip(float val) { - namespace math = device::math; - return math::max(math::min(val, math::FP8_E4M3_MAX), -math::FP8_E4M3_MAX); + return fmaxf(fminf(val, kFP8E4M3Max), -kFP8E4M3Max); } +#ifndef USE_ROCM // Pack two fp32 values into a single fp8x2_e4m3 with clamping. SGL_DEVICE fp8x2_e4m3_t pack_fp8(float x, float y) { return fp8x2_e4m3_t{fp32x2_t{fp8_e4m3_clip(x), fp8_e4m3_clip(y)}}; } +#else +// Software float -> FP8 E4M3 conversion for ROCm/HIP. +// Supports both E4M3FN (MI350X, gfx950) and E4M3FNUZ (MI300X, gfx942). +SGL_DEVICE uint8_t cvt_float_to_fp8_e4m3(float val) { + val = fp8_e4m3_clip(val); + if (val == 0.0f) return 0; + + uint32_t f32 = __float_as_uint(val); + uint8_t sign = static_cast((f32 >> 31) << 7); + int32_t exp32 = static_cast((f32 >> 23) & 0xFF) - 127; + uint32_t mant23 = f32 & 0x7FFFFF; + +#if HIP_FP8_TYPE_FNUZ + // E4M3FNUZ: bias=8, max=240, no negative zero, NaN=0x80 + constexpr int32_t kBias = 8; + constexpr int32_t kMaxExp = 15; + constexpr int32_t kMinSubnormExp = -10; // min subnormal exponent + constexpr int32_t kMinNormExp = -7; // min normal exponent + constexpr uint8_t kSaturate = 0x7Fu; // max normal = 0_1111_111 = 240.0 +#else + // E4M3FN: bias=7, max=448, NaN=0x7F + constexpr int32_t kBias = 7; + constexpr int32_t kMaxExp = 15; + constexpr int32_t kMinSubnormExp = -9; + constexpr int32_t kMinNormExp = -6; + constexpr uint8_t kSaturate = 0x7Eu; // max normal = 0_1111_110 = 448.0 +#endif + + int32_t exp8; + uint8_t mant3; + + if (exp32 < kMinSubnormExp) { + return sign; + } else if (exp32 < kMinNormExp) { + // Subnormal range + int32_t shift = -(kBias - 1) - exp32; // 1..3 + uint32_t subnorm_mant = (0x800000 | mant23) >> (shift + 20); + uint32_t round_bit = ((0x800000 | mant23) >> (shift + 19)) & 1; + subnorm_mant += round_bit; + mant3 = static_cast(subnorm_mant & 0x07); + exp8 = 0; + if (subnorm_mant > 7) { + exp8 = 1; + mant3 = 0; + } + } else { + exp8 = exp32 + kBias; + mant3 = static_cast(mant23 >> 20); + uint32_t round_bit = (mant23 >> 19) & 1; + mant3 += round_bit; + if (mant3 > 7) { + mant3 = 0; + exp8++; + } + if (exp8 >= kMaxExp) return sign | kSaturate; + } + return sign | (static_cast(exp8) << 3) | mant3; +} + +// Pack two fp32 values into a single fp8x2_e4m3 (uint16_t on HIP). +SGL_DEVICE fp8x2_e4m3_t pack_fp8(float x, float y) { + uint8_t x8 = cvt_float_to_fp8_e4m3(x); + uint8_t y8 = cvt_float_to_fp8_e4m3(y); + return static_cast(x8) | (static_cast(y8) << 8); +} +#endif } // namespace deepseek_v4::fp8 diff --git a/python/sglang/jit_kernel/include/sgl_kernel/runtime.cuh b/python/sglang/jit_kernel/include/sgl_kernel/runtime.cuh index 2812a2f8e..4ea722a3f 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/runtime.cuh +++ b/python/sglang/jit_kernel/include/sgl_kernel/runtime.cuh @@ -10,7 +10,38 @@ #include #include +#ifndef USE_ROCM #include +#else +#include +#ifndef cudaOccupancyMaxActiveBlocksPerMultiprocessor +#define cudaOccupancyMaxActiveBlocksPerMultiprocessor hipOccupancyMaxActiveBlocksPerMultiprocessor +#endif +#ifndef cudaDeviceGetAttribute +#define cudaDeviceGetAttribute hipDeviceGetAttribute +#endif +#ifndef cudaDevAttrMultiProcessorCount +#define cudaDevAttrMultiProcessorCount hipDeviceAttributeMultiprocessorCount +#endif +#ifndef cudaDevAttrComputeCapabilityMajor +#define cudaDevAttrComputeCapabilityMajor hipDeviceAttributeComputeCapabilityMajor +#endif +#ifndef cudaRuntimeGetVersion +#define cudaRuntimeGetVersion hipRuntimeGetVersion +#endif +#ifndef cudaOccupancyAvailableDynamicSMemPerBlock +inline hipError_t +cudaOccupancyAvailableDynamicSMemPerBlock(std::size_t* smem, const void* func, int num_blocks, int block_size) { + // HIP does not expose this directly; return max shared mem as conservative estimate + hipDeviceProp_t prop; + int device; + hipGetDevice(&device); + hipGetDeviceProperties(&prop, device); + *smem = prop.sharedMemPerBlock; + return hipSuccess; +} +#endif +#endif namespace host::runtime { diff --git a/python/sglang/jit_kernel/include/sgl_kernel/tensor.h b/python/sglang/jit_kernel/include/sgl_kernel/tensor.h index 484c969b5..1ae9233a6 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/tensor.h +++ b/python/sglang/jit_kernel/include/sgl_kernel/tensor.h @@ -33,6 +33,8 @@ #ifdef __CUDACC__ #include +#elif defined(__HIPCC__) +#include #endif namespace host { @@ -79,6 +81,15 @@ template <> struct _dtype_trait { inline static constexpr DLDataType value = {.code = DLDataTypeCode::kDLFloat8_e4m3fn, .bits = 8, .lanes = 1}; }; +#elif defined(__HIPCC__) +template <> +struct _dtype_trait { + inline static constexpr DLDataType value = {.code = DLDataTypeCode::kDLFloat, .bits = 16, .lanes = 1}; +}; +template <> +struct _dtype_trait { + inline static constexpr DLDataType value = {.code = DLDataTypeCode::kDLBfloat, .bits = 16, .lanes = 1}; +}; #endif template diff --git a/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh b/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh index a5abcdd4f..2dd6f3dc9 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh +++ b/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh @@ -44,6 +44,9 @@ inline constexpr auto cudaSuccess = hipSuccess; #define cudaGetErrorString hipGetErrorString #define cudaGetLastError hipGetLastError #define cudaLaunchKernel hipLaunchKernel +#define cudaMemcpyAsync hipMemcpyAsync +#define cudaMemcpyHostToDevice hipMemcpyHostToDevice +#define cudaMemcpyDeviceToHost hipMemcpyDeviceToHost #endif #ifndef USE_ROCM @@ -83,6 +86,13 @@ using fp32x4_t = float4; #define SGLANG_LDG(arg) *(arg) #endif +// DLPack device type for the current platform +#ifndef USE_ROCM +inline constexpr auto kDLGPU = kDLCUDA; +#else +inline constexpr auto kDLGPU = kDLROCM; +#endif + namespace device { /// \brief Macro: forced-inline device function qualifier. @@ -114,7 +124,11 @@ inline constexpr std::size_t kMaxVecBytes = SGL_ARCH_BLACKWELL_OR_GREATER ? 32 : /// \brief Number of threads per warp (always 32 on NVIDIA/AMD GPUs). inline constexpr auto kWarpThreads = 32u; /// \brief Full warp active mask (all 32 lanes). +#ifndef USE_ROCM inline constexpr auto kFullMask = 0xffffffffu; +#else +inline constexpr auto kFullMask = 0xffffffffffffffffULL; +#endif /** * \brief PDL (Programmatic Dependent Launch): wait for the primary kernel. diff --git a/python/sglang/jit_kernel/include/sgl_kernel/warp.cuh b/python/sglang/jit_kernel/include/sgl_kernel/warp.cuh index 975065e03..9d82efae1 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/warp.cuh +++ b/python/sglang/jit_kernel/include/sgl_kernel/warp.cuh @@ -1,5 +1,5 @@ /// \file warp.cuh -/// \brief Warp-level reduction primitives using `__shfl_xor_sync`. +/// \brief Warp-level reduction primitives. #pragma once #include @@ -7,52 +7,49 @@ namespace device::warp { -/// \brief Full 32-thread active mask. +/// \brief Full warp active mask. +#ifndef USE_ROCM static constexpr uint32_t kFullMask = 0xffffffffu; +using mask_t = uint32_t; +#else +static constexpr uint64_t kFullMask = 0xffffffffffffffffULL; +using mask_t = uint64_t; +#endif /** * \brief Warp-level sum reduction. * - * Computes the sum of `value` across all active lanes specified by - * `active_mask` using butterfly (XOR) shuffles. The result is - * broadcast to all participating lanes. - * - * \tparam kNumThreads Group size for the reduction (defaults to a full warp). - * \tparam T Numeric type (e.g. float). - * \param value Per-lane input value. - * \param active_mask Bitmask of participating lanes (default: all 32). - * \return The sum across all active lanes. + * On CUDA: uses __shfl_xor_sync with width=32. + * On HIP: uses __shfl_xor with explicit width parameter (supports wave64 sub-groups). */ template -SGL_DEVICE T reduce_sum(T value, uint32_t active_mask = kFullMask) { +SGL_DEVICE T reduce_sum(T value, mask_t active_mask = kFullMask) { static_assert(kNumThreads >= 1 && kNumThreads <= kWarpThreads); static_assert(std::has_single_bit(kNumThreads), "must be pow of 2"); #pragma unroll for (int mask = kNumThreads / 2; mask > 0; mask >>= 1) +#ifndef USE_ROCM value = value + __shfl_xor_sync(active_mask, value, mask, 32); +#else + value = value + __shfl_xor(value, mask, kNumThreads); +#endif return value; } /** * \brief Warp-level max reduction. - * - * Computes the maximum of `value` across all active lanes using - * butterfly shuffles. The result is broadcast to all participating - * lanes. - * - * \tparam kNumThreads Group size for the reduction (defaults to a full warp). - * \tparam T Numeric type (must be supported by `math::max`). - * \param value Per-lane input value. - * \param active_mask Bitmask of participating lanes (default: all 32). - * \return The maximum across all active lanes. */ template -SGL_DEVICE T reduce_max(T value, uint32_t active_mask = kFullMask) { +SGL_DEVICE T reduce_max(T value, mask_t active_mask = kFullMask) { static_assert(kNumThreads >= 1 && kNumThreads <= kWarpThreads); static_assert(std::has_single_bit(kNumThreads), "must be pow of 2"); #pragma unroll for (int mask = kNumThreads / 2; mask > 0; mask >>= 1) +#ifndef USE_ROCM value = math::max(value, __shfl_xor_sync(active_mask, value, mask, 32)); +#else + value = math::max(value, __shfl_xor(value, mask, kNumThreads)); +#endif return value; } diff --git a/python/sglang/jit_kernel/triton/hash_topk.py b/python/sglang/jit_kernel/triton/hash_topk.py new file mode 100644 index 000000000..b4e67fe1f --- /dev/null +++ b/python/sglang/jit_kernel/triton/hash_topk.py @@ -0,0 +1,99 @@ +"""HIP fallback for ``hash_topk``: ``csrc/deepseek_v4/hash_topk.cuh`` uses +CUDA-only primitives, so on ROCm we dispatch to this Triton implementation. +""" + +from __future__ import annotations + +from typing import Tuple + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _hash_topk_triton_kernel( + router_logits_ptr, + input_ids_ptr, + tid2eid_ptr, + topk_weights_ptr, + topk_ids_ptr, + num_routed_experts: tl.constexpr, + topk_routed: tl.constexpr, + topk_fused: tl.constexpr, + routed_scaling_factor, + BLOCK_K: tl.constexpr, +): + token_pos = tl.program_id(0) + token_id = tl.load(input_ids_ptr + token_pos).to(tl.int64) + + k_off = tl.arange(0, BLOCK_K) + routed_mask = k_off < topk_routed + fused_mask = k_off < topk_fused + is_shared = k_off >= topk_routed + + expert_id = tl.load( + tid2eid_ptr + token_id * topk_routed + k_off, + mask=routed_mask, + other=0, + ).to(tl.int32) + logit = tl.load( + router_logits_ptr + token_pos * num_routed_experts + expert_id, + mask=routed_mask, + other=0.0, + ).to(tl.float32) + + softplus = tl.maximum(logit, 0.0) + tl.log(1.0 + tl.exp(-tl.abs(logit))) + weight = tl.sqrt(softplus) + weight = tl.where(routed_mask, weight, 0.0) + routed_sum = tl.sum(weight, axis=0) + + shared_weight = 1.0 / routed_scaling_factor + final_weight = tl.where(is_shared, shared_weight, weight / routed_sum) + shared_id = num_routed_experts + (k_off - topk_routed) + final_id = tl.where(is_shared, shared_id, expert_id).to(tl.int32) + + out_off = token_pos * topk_fused + k_off + tl.store(topk_weights_ptr + out_off, final_weight, mask=fused_mask) + tl.store(topk_ids_ptr + out_off, final_id, mask=fused_mask) + + +def hash_topk_triton( + router_logits: torch.Tensor, + input_ids: torch.Tensor, + tid2eid: torch.Tensor, + num_fused_shared_experts: int, + routed_scaling_factor: float, + scoring_func: str, +) -> Tuple[torch.Tensor, torch.Tensor]: + assert scoring_func == "sqrtsoftplus" + + num_tokens = router_logits.size(0) + num_routed_experts = router_logits.size(1) + topk_routed = tid2eid.size(1) + topk_fused = topk_routed + num_fused_shared_experts + + topk_weights = torch.empty( + (num_tokens, topk_fused), dtype=torch.float32, device=router_logits.device + ) + topk_ids = torch.empty( + (num_tokens, topk_fused), dtype=torch.int32, device=router_logits.device + ) + if num_tokens == 0: + return topk_weights, topk_ids + + block_k = max(triton.next_power_of_2(topk_fused), 1) + _hash_topk_triton_kernel[(num_tokens,)]( + router_logits, + input_ids, + tid2eid, + topk_weights, + topk_ids, + num_routed_experts=num_routed_experts, + topk_routed=topk_routed, + topk_fused=topk_fused, + routed_scaling_factor=float(routed_scaling_factor), + BLOCK_K=block_k, + num_warps=1, + ) + return topk_weights, topk_ids diff --git a/python/sglang/jit_kernel/triton_store_cache.py b/python/sglang/jit_kernel/triton_store_cache.py new file mode 100644 index 000000000..b42a272e7 --- /dev/null +++ b/python/sglang/jit_kernel/triton_store_cache.py @@ -0,0 +1,237 @@ +from typing import Literal + +import torch +import triton +import triton.language as tl + +from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz + +_FP8_DTYPE = torch.float8_e4m3fnuz if is_fp8_fnuz() else torch.float8_e4m3fn +_FP8_INFO = torch.finfo(_FP8_DTYPE) + +# DeepSeek-V4 MLA paged FP8 cache layout +_MLA_HEAD_DIM = 512 # full MLA token dim (elements per input row) +_MLA_NOPE_DIM = 448 # nope sub-dim (elements) +_MLA_TILE_SIZE = 64 # FP8 tile width (also rope copy stride) +_MLA_SLOT_BYTES = 576 # bytes per slot in the paged FP8 cache +_MLA_BF16_SLOT_ELEMS = _MLA_SLOT_BYTES // 2 # bf16-view slot stride (elements) +_MLA_BF16_ROPE_OFFSET = _MLA_NOPE_DIM // 2 # bf16-view rope offset (elements) +_MLA_SCALES_PER_TOKEN = 8 # UE8M0 scales per token (7 nope tiles + 1 padding) +_MLA_NUM_TILES = 8 # 7 nope quant tiles + 1 rope copy tile +_MLA_ROPE_TILE_ID = 7 # tile id reserved for the rope copy + +# C4 indexer paged FP8 cache layout +_INDEXER_HEAD_DIM = 128 + +_UE8M0_EXPONENT_BIAS = 127 + + +@triton.jit +def _triton_fused_store_flashmla_kernel( + input_ptr, + cache_fp8_ptr, + cache_bf16_ptr, + cache_u8_ptr, + indices_ptr, + N, + PAGE_SIZE: tl.constexpr, + BYTES_PER_PAGE: tl.constexpr, + BYTES_PER_PAGE_BF16: tl.constexpr, + S_OFFSET: tl.constexpr, + TILE_SIZE: tl.constexpr, + HEAD_DIM: tl.constexpr, + NOPE_DIM: tl.constexpr, + SLOT_BYTES: tl.constexpr, + BF16_SLOT_ELEMS: tl.constexpr, + BF16_ROPE_OFFSET: tl.constexpr, + SCALES_PER_TOKEN: tl.constexpr, + ROPE_TILE_ID: tl.constexpr, + UE8M0_BIAS: tl.constexpr, + FP8_MIN: tl.constexpr, + FP8_MAX: tl.constexpr, + EPS: tl.constexpr, +): + token_id = tl.program_id(0) + tile_id = tl.program_id(1) + + if token_id >= N: + return + + loc = tl.load(indices_ptr + token_id).to(tl.int32) + page = loc // PAGE_SIZE + slot = loc % PAGE_SIZE + + if tile_id == ROPE_TILE_ID: + rope_lane = tl.arange(0, TILE_SIZE) + rope_vals = tl.load(input_ptr + token_id * HEAD_DIM + NOPE_DIM + rope_lane) + rope_bf16_offset = ( + page * BYTES_PER_PAGE_BF16 + + slot * BF16_SLOT_ELEMS + + BF16_ROPE_OFFSET + + rope_lane + ) + tl.store(cache_bf16_ptr + rope_bf16_offset, rope_vals) + else: + tile_lane = tl.arange(0, TILE_SIZE) + x_bf16 = tl.load( + input_ptr + token_id * HEAD_DIM + tile_id * TILE_SIZE + tile_lane + ) + x_fp32 = x_bf16.to(tl.float32) + + abs_max = tl.max(tl.abs(x_fp32)) + scale = tl.maximum(abs_max, EPS) / FP8_MAX + + # cast scale to ue8m0 format + log2_scale = tl.log2(scale) + ceil_log2 = tl.math.ceil(log2_scale) + inv_scale = tl.exp2(-ceil_log2) + + x_fp8 = tl.clamp(x_fp32 * inv_scale, FP8_MIN, FP8_MAX).to( + cache_fp8_ptr.dtype.element_ty + ) + + nope_offset = ( + page * BYTES_PER_PAGE + slot * SLOT_BYTES + tile_id * TILE_SIZE + tile_lane + ) + tl.store(cache_fp8_ptr + nope_offset, x_fp8) + + ue8m0 = (ceil_log2.to(tl.int32) + UE8M0_BIAS).to(tl.uint8) + scale_offset = ( + page * BYTES_PER_PAGE + S_OFFSET + slot * SCALES_PER_TOKEN + tile_id + ) + tl.store(cache_u8_ptr + scale_offset, ue8m0) + + +def triton_fused_store_flashmla( + input: torch.Tensor, + cache: torch.Tensor, + indices: torch.Tensor, + page_size: int, +) -> None: + """Fused FP8 quantise + paged scatter for the SWA (flashmla) KV cache.""" + N = input.shape[0] + if N == 0: + return + + bytes_per_page = cache.shape[1] + cache_fp8 = cache.view(_FP8_DTYPE) + cache_bf16 = cache.view(torch.bfloat16) + indices_i32 = indices.to(torch.int32) if indices.dtype != torch.int32 else indices + + _triton_fused_store_flashmla_kernel[(N, _MLA_NUM_TILES)]( + input, + cache_fp8, + cache_bf16, + cache, + indices_i32, + N, + PAGE_SIZE=page_size, + BYTES_PER_PAGE=bytes_per_page, + BYTES_PER_PAGE_BF16=bytes_per_page // 2, + S_OFFSET=page_size * _MLA_SLOT_BYTES, + TILE_SIZE=_MLA_TILE_SIZE, + HEAD_DIM=_MLA_HEAD_DIM, + NOPE_DIM=_MLA_NOPE_DIM, + SLOT_BYTES=_MLA_SLOT_BYTES, + BF16_SLOT_ELEMS=_MLA_BF16_SLOT_ELEMS, + BF16_ROPE_OFFSET=_MLA_BF16_ROPE_OFFSET, + SCALES_PER_TOKEN=_MLA_SCALES_PER_TOKEN, + ROPE_TILE_ID=_MLA_ROPE_TILE_ID, + UE8M0_BIAS=_UE8M0_EXPONENT_BIAS, + FP8_MIN=_FP8_INFO.min, + FP8_MAX=_FP8_INFO.max, + EPS=1e-8, + ) + + +@triton.jit +def _triton_fused_store_indexer_kernel( + input_ptr, + cache_fp8_ptr, + cache_f32_ptr, + indices_ptr, + N, + PAGE_SIZE: tl.constexpr, + BYTES_PER_PAGE: tl.constexpr, + BYTES_PER_PAGE_F32: tl.constexpr, + SCALE_PAGE_OFFSET_F32: tl.constexpr, + HEAD_DIM: tl.constexpr, + FP8_MIN: tl.constexpr, + FP8_MAX: tl.constexpr, + EPS: tl.constexpr, +): + token_id = tl.program_id(0) + if token_id >= N: + return + + loc = tl.load(indices_ptr + token_id).to(tl.int32) + page = loc // PAGE_SIZE + slot = loc % PAGE_SIZE + + lane = tl.arange(0, HEAD_DIM) + x_fp32 = tl.load(input_ptr + token_id * HEAD_DIM + lane).to(tl.float32) + + abs_max = tl.max(tl.abs(x_fp32)) + scale = tl.maximum(abs_max, EPS) / FP8_MAX + inv_scale = 1.0 / scale + + x_fp8 = tl.clamp(x_fp32 * inv_scale, FP8_MIN, FP8_MAX).to( + cache_fp8_ptr.dtype.element_ty + ) + + fp8_offset = page * BYTES_PER_PAGE + slot * HEAD_DIM + lane + tl.store(cache_fp8_ptr + fp8_offset, x_fp8) + + f32_offset = page * BYTES_PER_PAGE_F32 + SCALE_PAGE_OFFSET_F32 + slot + tl.store(cache_f32_ptr + f32_offset, scale) + + +def triton_fused_store_indexer( + input: torch.Tensor, + cache: torch.Tensor, + indices: torch.Tensor, + page_size: int, +) -> None: + """Fused FP8 quantise + paged scatter for the C4 indexer KV cache.""" + N = input.shape[0] + if N == 0: + return + + bytes_per_page = cache.shape[1] + bytes_per_page_f32 = bytes_per_page // 4 + scale_page_offset_f32 = (_INDEXER_HEAD_DIM * page_size) // 4 + + cache_fp8 = cache.view(_FP8_DTYPE) + cache_f32 = cache.view(torch.float32) + indices_i32 = indices.to(torch.int32) if indices.dtype != torch.int32 else indices + + _triton_fused_store_indexer_kernel[(N,)]( + input, + cache_fp8, + cache_f32, + indices_i32, + N, + PAGE_SIZE=page_size, + BYTES_PER_PAGE=bytes_per_page, + BYTES_PER_PAGE_F32=bytes_per_page_f32, + SCALE_PAGE_OFFSET_F32=scale_page_offset_f32, + HEAD_DIM=_INDEXER_HEAD_DIM, + FP8_MIN=_FP8_INFO.min, + FP8_MAX=_FP8_INFO.max, + EPS=1e-8, + ) + + +def triton_fused_store_cache( + input: torch.Tensor, + cache: torch.Tensor, + indices: torch.Tensor, + *, + page_size: int, + type: Literal["flashmla", "indexer"], +) -> None: + """ROCm dispatch for fused_store_cache().""" + if type == "flashmla": + triton_fused_store_flashmla(input, cache, indices, page_size) + else: + triton_fused_store_indexer(input, cache, indices, page_size) diff --git a/python/sglang/jit_kernel/utils.py b/python/sglang/jit_kernel/utils.py index bcd42e5ce..2bd0390fb 100644 --- a/python/sglang/jit_kernel/utils.py +++ b/python/sglang/jit_kernel/utils.py @@ -277,7 +277,18 @@ def _jit_compile_context(): # NOTE: this might also be used in __main__.py for compile flags export def _get_default_target_flags() -> List[str]: if is_hip_runtime(): - return ["-DUSE_ROCM", "-std=c++20", "-O3"] + flags = ["-DUSE_ROCM", "-std=c++20", "-O3"] + # Detect FP8 type based on GPU architecture + try: + device = torch.cuda.current_device() + gcn_arch = torch.cuda.get_device_properties(device).gcnArchName + if "gfx942" in gcn_arch: + flags.append("-DHIP_FP8_TYPE_FNUZ=1") + else: + flags.append("-DHIP_FP8_TYPE_E4M3=1") + except Exception: + flags.append("-DHIP_FP8_TYPE_E4M3=1") + return flags else: return [ get_jit_cuda_arch().jit_flag, diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 2791aeec9..8f64246ae 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -625,6 +625,8 @@ class Envs: SGLANG_OPT_USE_AITER_MHC_PRE = EnvBool(True) SGLANG_OPT_USE_AITER_MHC_POST = EnvBool(True) SGLANG_OPT_USE_FUSED_COMPRESS = EnvBool(False) + SGLANG_OPT_USE_FUSED_QK_NORM_ROPE = EnvBool(True) + SGLANG_OPT_USE_FUSED_CLAMP_ACT_MUL = EnvBool(True) SGLANG_FIX_MTP_HC_HIDDEN = EnvBool(False) # ==================================================================== diff --git a/python/sglang/srt/layers/attention/dsv4/compress_hip.py b/python/sglang/srt/layers/attention/dsv4/compress_hip.py index fa70fc693..cff010c55 100644 --- a/python/sglang/srt/layers/attention/dsv4/compress_hip.py +++ b/python/sglang/srt/layers/attention/dsv4/compress_hip.py @@ -16,6 +16,11 @@ from sglang.srt.layers.deepseek_v4_rope import ( apply_rotary_emb_triton, fused_norm_rope_inplace_triton, ) + +try: + from sglang.srt.layers.deepseek_v4_rope import fused_softmax_pool_triton +except ImportError: + fused_softmax_pool_triton = None from sglang.srt.mem_cache.deepseek_v4_compress_state import ( CompressStatePool, KVAndScore, @@ -91,15 +96,24 @@ class CompressorHip(_CompressorBase): def __init__(self, *args, **kwargs) -> None: super().__init__(*args, **kwargs) self.norm = DeepseekRefRMSNorm(self.head_dim, eps=self.norm.variance_epsilon) + self._freqs_cis_real: torch.Tensor | None = None @cached_property def use_fused_compress(self) -> bool: - return False + return envs.SGLANG_OPT_USE_FUSED_COMPRESS.get() @cached_property def use_hip_fused_compress(self) -> bool: return envs.SGLANG_OPT_USE_FUSED_COMPRESS.get() + @cached_property + def use_fused_compress_triton(self) -> bool: + # The fused Triton kernel only benefits non-overlap (HCA, ratio=128) + # but HCA's K=128 loop is too sequential to outperform batched ops. + # CSA (overlap=True) has a reshape/overlap-transform semantic mismatch. + # Disabled until a tiled kernel for CSA overlap is implemented. + return False + def _get_states( self, forward_batch: ForwardBatch, @@ -247,15 +261,22 @@ class CompressorHip(_CompressorBase): pt += extend_lens[i] continue - kv_compressed = ( - kv_and_score_to_compress.kv - * kv_and_score_to_compress.score.softmax(dim=1) - ).sum(dim=1) + beg_idx = prefix_lens[i] // self.ratio * self.ratio + end_idx = (prefix_lens[i] + extend_lens[i]) // self.ratio * self.ratio + + if self.use_hip_fused_compress: + kv_compressed = fused_softmax_pool_triton( + kv_and_score_to_compress.kv_score, + kv_and_score_to_compress._item_size, + ) + else: + kv_compressed = ( + kv_and_score_to_compress.kv + * kv_and_score_to_compress.score.softmax(dim=1) + ).sum(dim=1) assert kv_compressed.dtype == torch.float32 - beg_idx = prefix_lens[i] // self.ratio * self.ratio - end_idx = (prefix_lens[i] + extend_lens[i]) // self.ratio * self.ratio freqs_cis = self.freqs_cis[beg_idx : end_idx : self.ratio] assert freqs_cis.size(0) == kv_compressed.size( 0 @@ -336,9 +357,43 @@ class CompressorHip(_CompressorBase): kv_and_score_to_compress = state_pool.get_state_by_state_loc( compress_indices_state.view(-1) ).view(-1, self.ratio, self.coff * self.head_dim) + bs = seq_lens.size(0) + + if self.use_fused_compress_triton and not self.overlap: + # Fused path for non-overlap (HCA, ratio=128, coff=1): + # APE + softmax-pool + norm + RoPE in one kernel. + # Overlap (CSA) is excluded because the overlap_transform_decode + # rearranges A/B halves across the coff dimension in a way + # that simple reshape cannot replicate correctly. + raw = kv_and_score_to_compress.kv_score + gathered = raw.reshape(bs, self.ratio, raw.shape[-1]).contiguous() + + comp_positions = (seq_lens - 1) // self.ratio * self.ratio + freqs_real_table = self._get_freqs_cis_real() + freqs_batch = freqs_real_table[comp_positions] + + from sglang.srt.layers.attention.dsv4.fused_compress_kernel import ( + fused_ape_pool_norm_rope, + ) + + kv_compressed = fused_ape_pool_norm_rope( + kv_score_gathered=gathered, + ape=self.ape, + rms_weight=self.norm.weight, + rms_eps=self.norm.eps, + freqs_cis_real=freqs_batch, + head_dim=self.head_dim, + rope_head_dim=self.rope_head_dim, + ratio=self.ratio, + overlap=self.overlap, + ) + if self.rotate: + kv_compressed = rotate_activation(kv_compressed) + return kv_compressed + + # Unfused reference path kv_and_score_to_compress.score.add_(self.ape.unsqueeze(0)) - bs = seq_lens.size(0) if self.overlap: kv_and_score_to_compress = kv_and_score_to_compress.view( bs, self.coff * self.ratio, self.coff * self.head_dim @@ -348,17 +403,20 @@ class CompressorHip(_CompressorBase): score=self.overlap_transform_decode(kv_and_score_to_compress.score), ) - self.print_tensor(kv_and_score_to_compress.kv, "kv_to_compress") - self.print_tensor(kv_and_score_to_compress.score, "score_to_compress") - kv_and_score_to_compress = kv_and_score_to_compress.view( bs, self.ratio * self.coff, self.head_dim ) - kv_compressed = ( - kv_and_score_to_compress.kv * kv_and_score_to_compress.score.softmax(dim=1) - ).sum(dim=1) - self.print_tensor(kv_compressed, "kv_before_norm") + if self.use_hip_fused_compress: + kv_compressed = fused_softmax_pool_triton( + kv_and_score_to_compress.kv_score, + kv_and_score_to_compress._item_size, + ) + else: + kv_compressed = ( + kv_and_score_to_compress.kv + * kv_and_score_to_compress.score.softmax(dim=1) + ).sum(dim=1) if self.use_hip_fused_compress: freqs_cis = self._init_freqs_cis_per_decode_step(forward_batch, seq_lens) fused_norm_rope_inplace_triton( @@ -366,17 +424,13 @@ class CompressorHip(_CompressorBase): ) else: kv_compressed = self.norm(kv_compressed) - self.print_tensor(kv_compressed, "kv_after_norm") freqs_cis = self.freqs_cis[(seq_lens - 1) // self.ratio * self.ratio] - self.print_tensor(freqs_cis, "freqs_cis") apply_rotary_emb_triton( kv_compressed[..., -self.rope_head_dim :], freqs_cis ) - self.print_tensor(kv_compressed, "kv_after_rope") if self.rotate: kv_compressed = rotate_activation(kv_compressed) - self.print_tensor(kv_compressed, "compressed_kv_output") return kv_compressed def compress_fused( @@ -404,13 +458,30 @@ class CompressorHip(_CompressorBase): is_paged=True, ) + def _get_freqs_cis_real(self) -> torch.Tensor: + """Cache the float32 view of freqs_cis (complex64 -> real interleaved).""" + if self._freqs_cis_real is None: + if self.freqs_cis.is_complex(): + self._freqs_cis_real = ( + torch.view_as_real(self.freqs_cis).flatten(-2).contiguous() + ) + else: + self._freqs_cis_real = self.freqs_cis.contiguous() + return self._freqs_cis_real + def compress_dispatch( self, kv_score: torch.Tensor, forward_batch: ForwardBatch, attn_backend: AttentionBackend, ) -> torch.Tensor: - if self.use_fused_compress: + if self.use_fused_compress and ( + envs.SGLANG_OPT_DPSK_V4_RADIX.get() + and ( + forward_batch.forward_mode.is_decode() + or forward_batch.forward_mode.is_extend_without_speculative() + ) + ): return self.compress_fused( kv_score, forward_batch, attn_backend=attn_backend ) diff --git a/python/sglang/srt/layers/fused_qk_norm_rope_store.py b/python/sglang/srt/layers/fused_qk_norm_rope_store.py new file mode 100644 index 000000000..c6499902a --- /dev/null +++ b/python/sglang/srt/layers/fused_qk_norm_rope_store.py @@ -0,0 +1,380 @@ +"""Fused Q per-head RMSNorm + KV RMSNorm + RoPE + FP8 nope quant + paged SWA store. + +Single Triton kernel replacing the 2-kernel path: + 1. fused_reduce_qk_norm_rope_swa_write (norm + RoPE) + 2. store_cache -> fused_store_cache (FP8 quant + paged scatter) + +Grid: (cdiv(M, BLOCK_SIZE_M), num_local_heads + 1). + pid_h < num_local_heads: Q head programs (split-K reduce + norm + RoPE) + pid_h == num_local_heads: KV program (norm + RoPE + FP8 quant nope + paged scatter) +""" + +from typing import Optional + +import torch +import triton +import triton.language as tl + +from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz + +_fp8_fnuz = is_fp8_fnuz() + + +# --------------------------------------------------------------------------- +# Triton JIT helpers +# --------------------------------------------------------------------------- + + +@triton.jit +def _batched_rmsnorm(row, weight, n_cols, epsilon): + row_norm = tl.sum(row * row, axis=-1) + norm_factor = tl.math.rsqrt((row_norm / n_cols) + epsilon) + if weight is not None: + return row * norm_factor[:, None] * weight[None, :] + return row * norm_factor[:, None] + + +@triton.jit +def _gptj_rotate(x, mask, BM: tl.constexpr, BD: tl.constexpr, BDH: tl.constexpr): + x_rot = tl.where(mask, x, -x) + x_rot = tl.reshape(x_rot, (BM, BDH, 2)) + x_rot = tl.flip(x_rot, 2) + return tl.reshape(x_rot, (BM, BD)) + + +@triton.jit +def _batched_rope( + x_pe, cos, sin, d_pe_offs, BM: tl.constexpr, BD: tl.constexpr, BDH: tl.constexpr +): + mask = (d_pe_offs % 2 == 0)[None, :] + x_rot = _gptj_rotate(x_pe, mask, BM, BD, BDH) + return x_pe * cos + x_rot * sin + + +# --------------------------------------------------------------------------- +# Main kernel +# --------------------------------------------------------------------------- + + +@triton.jit +def _fused_qk_norm_rope_store_kernel( + q_in_ptr, + q_out_ptr, + kv_ptr, + q_norm_weight_ptr, + kv_norm_weight_ptr, + positions_ptr, + cos_ptr, + sin_ptr, + swa_cache_ptr, + swa_loc_ptr, + M, + q_in_splitk_stride, + q_in_m_stride, + q_in_d_stride, + stride_qm, + stride_qh, + stride_qd, + stride_kv_m, + stride_kv_d, + cos_stride_t, + cos_stride_d, + swa_cache_stride_page, + q_eps, + kv_eps, + BLOCK_SIZE_M: tl.constexpr, + HEAD_DIM: tl.constexpr, + ROPE_DIM: tl.constexpr, + NUM_LOCAL_HEADS: tl.constexpr, + NUM_SPLITK: tl.constexpr, + HAS_SWA_STORE: tl.constexpr, + DIM_NOPE: tl.constexpr, + TILE_SIZE: tl.constexpr, + NUM_NOPE_TILES: tl.constexpr, + FP8_MIN: tl.constexpr, + FP8_MAX: tl.constexpr, + BYTES_PER_TOKEN: tl.constexpr, + SWA_PAGE_SIZE: tl.constexpr, +): + pid_m = tl.program_id(0).to(tl.int64) + pid_h = tl.program_id(1).to(tl.int64) + NOPE_DIM: tl.constexpr = HEAD_DIM - ROPE_DIM + NUM_PE_CHUNKS: tl.constexpr = HEAD_DIM // ROPE_DIM + + m_offs = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64) + m_mask = m_offs < M + + offs_d_full = tl.arange(0, HEAD_DIM) + nope_d_mask = offs_d_full < NOPE_DIM + + d_pe_offs = tl.arange(0, ROPE_DIM).to(tl.int64) + d_cos_offs = d_pe_offs // 2 + + # ===== Q path ===== + if pid_h < NUM_LOCAL_HEADS: + head_id = pid_h.to(tl.int32) + offs_n = head_id * HEAD_DIM + offs_d_full + + splitk_offs = tl.arange(0, NUM_SPLITK).to(tl.int64) + q_ptrs = ( + q_in_ptr + + splitk_offs[:, None, None] * q_in_splitk_stride + + m_offs[None, :, None] * q_in_m_stride + + offs_n[None, None, :] * q_in_d_stride + ) + q_tile = tl.load(q_ptrs, mask=m_mask[None, :, None], other=0.0).to(tl.float32) + q_acc = tl.sum(q_tile, axis=0) + + if q_norm_weight_ptr is not None: + w_q = tl.load(q_norm_weight_ptr + offs_d_full).to(tl.float32) + else: + w_q = None + q_normed = _batched_rmsnorm(q_acc, w_q, HEAD_DIM, q_eps) + + q_base = q_out_ptr + m_offs[:, None] * stride_qm + pid_h * stride_qh + tl.store( + q_base + offs_d_full[None, :] * stride_qd, + q_normed.to(q_out_ptr.dtype.element_ty), + mask=m_mask[:, None] & nope_d_mask[None, :], + ) + + q_pe = tl.where((offs_d_full >= NOPE_DIM)[None, :], q_normed, 0.0) + q_pe = tl.reshape(q_pe, (BLOCK_SIZE_M, NUM_PE_CHUNKS, ROPE_DIM)) + q_pe = tl.sum(q_pe, axis=1) + + pos = tl.load(positions_ptr + m_offs, mask=m_mask, other=0) + cos_o = pos[:, None] * cos_stride_t + d_cos_offs[None, :] * cos_stride_d + cos = tl.load(cos_ptr + cos_o, mask=m_mask[:, None], other=0) + sin = tl.load(sin_ptr + cos_o, mask=m_mask[:, None], other=0) + + q_pe = _batched_rope( + q_pe, cos, sin, d_pe_offs, BLOCK_SIZE_M, ROPE_DIM, ROPE_DIM // 2 + ) + tl.store( + q_base + (NOPE_DIM + d_pe_offs[None, :]) * stride_qd, + q_pe.to(q_out_ptr.dtype.element_ty), + mask=m_mask[:, None], + ) + return + + # ===== KV path ===== + src_id = m_offs.to(tl.int32) + src_mask = m_mask + + pos = tl.load(positions_ptr + src_id, mask=src_mask, other=0) + cos_o = pos[:, None] * cos_stride_t + d_cos_offs[None, :] * cos_stride_d + cos = tl.load(cos_ptr + cos_o, mask=src_mask[:, None], other=0) + sin = tl.load(sin_ptr + cos_o, mask=src_mask[:, None], other=0) + + kv_base = kv_ptr + src_id[:, None].to(tl.int64) * stride_kv_m + kv_full_ptrs = kv_base + offs_d_full[None, :] * stride_kv_d + + kv_full = tl.load(kv_full_ptrs, mask=src_mask[:, None], other=0.0).to(tl.float32) + + if kv_norm_weight_ptr is not None: + w_kv = tl.load(kv_norm_weight_ptr + offs_d_full).to(tl.float32) + else: + w_kv = None + kv_normed = _batched_rmsnorm(kv_full, w_kv, HEAD_DIM, kv_eps) + + tl.store( + kv_full_ptrs, + kv_normed.to(kv_ptr.dtype.element_ty), + mask=src_mask[:, None] & nope_d_mask[None, :], + ) + + kv_pe = tl.where((offs_d_full >= NOPE_DIM)[None, :], kv_normed, 0.0) + kv_pe = tl.reshape(kv_pe, (BLOCK_SIZE_M, NUM_PE_CHUNKS, ROPE_DIM)) + kv_pe = tl.sum(kv_pe, axis=1) + + kv_pe = _batched_rope( + kv_pe, cos, sin, d_pe_offs, BLOCK_SIZE_M, ROPE_DIM, ROPE_DIM // 2 + ) + tl.store( + kv_base + (NOPE_DIM + d_pe_offs[None, :]) * stride_kv_d, + kv_pe.to(kv_ptr.dtype.element_ty), + mask=src_mask[:, None], + ) + + # ===== Paged SWA store: FP8 quant nope + BF16 rope + scales ===== + # Layout within a page (matches fused_store_flashmla_cache CUDA kernel): + # Values region: [page_size tokens * 576 bytes/token] + # Per token: 448 bytes FP8 nope + 128 bytes BF16 rope + # Scales region: [page_size tokens * 8 bytes/token] + # Per token: 7 scale bytes + 1 pad byte + # Total per page before padding: page_size * 584 + VALUE_STRIDE: tl.constexpr = DIM_NOPE + ROPE_DIM * 2 + SCALE_BYTES: tl.constexpr = NUM_NOPE_TILES + 1 + + if HAS_SWA_STORE: + loc = tl.load(swa_loc_ptr + src_id, mask=src_mask, other=0) + page_id = loc // SWA_PAGE_SIZE + page_off = loc % SWA_PAGE_SIZE + page_base = page_id.to(tl.int64) * swa_cache_stride_page + value_base = page_base + page_off.to(tl.int64) * VALUE_STRIDE + scale_base = ( + page_base + + SWA_PAGE_SIZE * VALUE_STRIDE + + page_off.to(tl.int64) * SCALE_BYTES + ) + + EPS: tl.constexpr = 1e-8 + nope_tile_offs = tl.arange(0, TILE_SIZE) + + for tile_i in tl.static_range(NUM_NOPE_TILES): + tile_start = tile_i * TILE_SIZE + tile_data = tl.load( + kv_ptr + + src_id[:, None].to(tl.int64) * stride_kv_m + + (tile_start + nope_tile_offs[None, :]) * stride_kv_d, + mask=src_mask[:, None], + other=0.0, + ).to(tl.float32) + + abs_max = tl.max(tl.abs(tile_data), axis=-1) + abs_max_c = tl.maximum(abs_max, EPS) + scale_f = abs_max_c / FP8_MAX + log2_s = tl.log2(scale_f) + ceil_log2 = tl.math.ceil(log2_s) + scale_pow2 = tl.exp2(ceil_log2) + inv_scale = 1.0 / scale_pow2 + x_scaled = tile_data * inv_scale[:, None] + x_fp8 = tl.clamp(x_scaled, FP8_MIN, FP8_MAX) + + x_fp8_cast = x_fp8.to(tl.float8e4nv) + x_fp8_bytes = x_fp8_cast.to(tl.uint8, bitcast=True) + fp8_byte_offs = value_base[:, None] + tile_start + nope_tile_offs[None, :] + tl.store( + swa_cache_ptr + fp8_byte_offs, + x_fp8_bytes, + mask=src_mask[:, None], + ) + + scale_uint8 = (ceil_log2.to(tl.int32) + 127).to(tl.uint8) + tl.store( + swa_cache_ptr + scale_base + tile_i, + scale_uint8, + mask=src_mask, + ) + + rope_data = kv_pe.to(tl.bfloat16) + rope_offs = tl.arange(0, ROPE_DIM) + rope_byte_base = value_base[:, None] + DIM_NOPE + rope_offs[None, :] * 2 + rope_data_as_i16 = rope_data.to(tl.int16, bitcast=True) + lo = (rope_data_as_i16 & 0xFF).to(tl.uint8) + hi = ((rope_data_as_i16 >> 8) & 0xFF).to(tl.uint8) + tl.store(swa_cache_ptr + rope_byte_base, lo, mask=src_mask[:, None]) + tl.store(swa_cache_ptr + rope_byte_base + 1, hi, mask=src_mask[:, None]) + + +# --------------------------------------------------------------------------- +# Python wrapper +# --------------------------------------------------------------------------- + + +def fused_qk_norm_rope_swa_store( + q: torch.Tensor, + kv: torch.Tensor, + q_norm_weight: Optional[torch.Tensor], + kv_norm_weight: Optional[torch.Tensor], + q_rms_eps: float, + kv_rms_eps: float, + rope_head_dim: int, + cos_cache: torch.Tensor, + sin_cache: torch.Tensor, + positions: torch.Tensor, + swa_cache: Optional[torch.Tensor] = None, + swa_loc: Optional[torch.Tensor] = None, + swa_page_size: int = 128, + q_out: Optional[torch.Tensor] = None, + dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """Fused Q norm + KV norm + RoPE + optional FP8 paged SWA store. + + Args: + q: [M, N] or [splitk, M, N] where N = num_local_heads * head_dim + kv: [M, head_dim=512] mutated in-place (norm + RoPE) + swa_cache: paged SWA KV pool buffer [num_pages, bytes_per_page] uint8 + swa_loc: [M] int32 pre-translated paged indices + swa_page_size: tokens per SWA page (default 128) + """ + head_dim = kv.shape[1] + + if q.dim() == 3: + num_splitk, M, N = q.shape + q_in_splitk_stride = q.stride(0) + q_in_m_stride = q.stride(1) + q_in_d_stride = q.stride(2) + else: + M, N = q.shape + num_splitk = 1 + q_in_splitk_stride = 0 + q_in_m_stride = q.stride(0) + q_in_d_stride = q.stride(1) + + num_local_heads = N // head_dim + + if q_out is None: + q_out = torch.empty( + (M, num_local_heads, head_dim), dtype=dtype, device=q.device + ) + + HAS_SWA_STORE = swa_cache is not None and swa_loc is not None + + dim_nope = 448 + dim_rope = 64 + tile_size = 64 + num_nope_tiles = dim_nope // tile_size + scale_pad = 1 + bytes_per_token = dim_nope + dim_rope * 2 + num_nope_tiles + scale_pad + + if _fp8_fnuz: + fp8_info = torch.finfo(torch.float8_e4m3fnuz) + else: + fp8_info = torch.finfo(torch.float8_e4m3fn) + + BLOCK_SIZE_M = min(4, triton.next_power_of_2(M)) if M < 4 else 4 + num_warps = 4 + + grid = (triton.cdiv(M, BLOCK_SIZE_M), num_local_heads + 1) + _fused_qk_norm_rope_store_kernel[grid]( + q, + q_out, + kv, + q_norm_weight, + kv_norm_weight, + positions, + cos_cache, + sin_cache, + swa_cache if HAS_SWA_STORE else None, + swa_loc if HAS_SWA_STORE else None, + M, + q_in_splitk_stride, + q_in_m_stride, + q_in_d_stride, + q_out.stride(0), + q_out.stride(1), + q_out.stride(2), + kv.stride(0), + kv.stride(1), + cos_cache.stride(0), + cos_cache.stride(-1), + swa_cache.stride(0) if HAS_SWA_STORE else 0, + q_rms_eps, + kv_rms_eps, + BLOCK_SIZE_M=BLOCK_SIZE_M, + HEAD_DIM=head_dim, + ROPE_DIM=rope_head_dim, + NUM_LOCAL_HEADS=num_local_heads, + NUM_SPLITK=num_splitk, + HAS_SWA_STORE=HAS_SWA_STORE, + DIM_NOPE=dim_nope, + TILE_SIZE=tile_size, + NUM_NOPE_TILES=num_nope_tiles, + FP8_MIN=fp8_info.min, + FP8_MAX=fp8_info.max, + BYTES_PER_TOKEN=bytes_per_token, + SWA_PAGE_SIZE=swa_page_size, + num_warps=num_warps, + ) + return q_out diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index 90e42e893..6e4b7eee4 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/python/sglang/srt/layers/layernorm.py @@ -303,6 +303,12 @@ class RMSNorm(MultiPlatformOp): residual: Optional[torch.Tensor] = None, post_residual_addition: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + # Fix dsv4 dp attenton issue + # the symptom is torch.AcceleratorError: HIP error: invalid configuration argument + if x.shape[0] == 0: + if residual is not None: + return x, residual + return x # Aiter's RMSNorm kernels expect 2D contiguous inputs. Keep the # already-safe layout as a zero-copy path, and only normalize strided or # higher-rank views such as Q/K slices from packed QKV projections. diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py index 5e77fa614..e3c399cf5 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py @@ -578,6 +578,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): enable_memory_saver=enable_memory_saver, ratio=ratio, online=(ratio == 128 and ONLINE_C128), + swa_page_size=self.swa_page_size, ) if ratio == 4: diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 71e233f40..e2656c0e4 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -258,6 +258,11 @@ class DeepseekV2MLP(nn.Module): "Only silu is supported for now." ) self.act_fn = SiluAndMul() + self.use_fused_clamp_act_mul = ( + _is_hip and envs.SGLANG_OPT_USE_FUSED_CLAMP_ACT_MUL.get() + ) + self._fused_clamp_fp8_checked = False + self._fused_clamp_use_fp8 = False def forward( self, @@ -319,8 +324,41 @@ class DeepseekV2MLP(nn.Module): down_output, ) return down_output + + if self.use_fused_clamp_act_mul and self.swiglu_limit is not None: + from aiter.ops.triton.fusions.fused_clamp_act_mul import ( + fused_clamp_act_mul, + ) + + if not self._fused_clamp_fp8_checked: + from sglang.srt.layers.quantization.fp8 import Fp8LinearMethod + + qm = getattr(self.down_proj, "quant_method", None) + self._fused_clamp_use_fp8 = ( + isinstance(qm, Fp8LinearMethod) and qm.block_quant + ) + self._fused_clamp_fp8_checked = True + + if self._fused_clamp_use_fp8: + from aiter import dtypes + + x_fp8, x_scale = fused_clamp_act_mul( + gate_up, + swiglu_limit=self.swiglu_limit, + activation="silu", + dtype_quant=dtypes.fp8, + transpose_scale=False, + ) + x = (x_fp8, x_scale) + else: + x = fused_clamp_act_mul( + gate_up, + swiglu_limit=self.swiglu_limit, + activation="silu", + ) + # Fallback: fused silu+clamp kernel (still faster than unfused) - if self.swiglu_limit is not None: + elif self.swiglu_limit is not None: M, N = gate_up.shape x = gate_up.new_empty((M, N // 2)) silu_and_mul_clamp(gate_up, x, float(self.swiglu_limit)) diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index ecdfaabf2..86864dcc9 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -96,6 +96,8 @@ from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import ( LazyValue, add_prefix, + get_bool_env_var, + is_gfx95_supported, log_info_on_rank0, make_layers, ) @@ -105,6 +107,29 @@ logger = logging.getLogger(__name__) _FP8_WO_A_GEMM = envs.SGLANG_OPT_FP8_WO_A_GEMM.get() +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip +_is_gfx95_supported = is_gfx95_supported() + +if _use_aiter: + if _is_gfx95_supported: + from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_group_quant + + +def _fused_rmsnorm_fp8_quant(hidden_states, weight, eps): + x_quant, x_bf16, _, _ = fused_rms_fp8_group_quant( + hidden_states, + weight, + eps, + inp2=None, + inp2_weight=None, + inp2_epsilon=None, + group_size=128, + dtype_quant=torch.float8_e4m3fn, + res1=None, + output_unquantized_inp1=True, + ) + return x_quant, x_bf16 + if TYPE_CHECKING: from sglang.srt.layers.attention.deepseek_v4_backend import ( @@ -249,6 +274,12 @@ class MQALayer(nn.Module): self.register_buffer("freqs_cis", freqs_cis, persistent=False) self.freqs_cis: torch.Tensor + if _is_hip: + cos_cache = freqs_cis.real.to(torch.bfloat16).unsqueeze(-2).unsqueeze(-2) + sin_cache = freqs_cis.imag.to(torch.bfloat16).unsqueeze(-2).unsqueeze(-2) + self.register_buffer("cos_cache", cos_cache, persistent=False) + self.register_buffer("sin_cache", sin_cache, persistent=False) + if envs.SGLANG_OPT_USE_MULTI_STREAM_OVERLAP.get() and alt_streams is not None: self.alt_streams = alt_streams[:3] self.alt_streams_indexer = alt_streams[-2:] @@ -357,6 +388,10 @@ class MQALayer(nn.Module): prefix=add_prefix("attn_mqa", prefix), ) + self.use_fused_qk_norm_rope = ( + _is_hip and envs.SGLANG_OPT_USE_FUSED_QK_NORM_ROPE.get() + ) + # KV cache write is always fused into the K kernel # (`_compute_kv_to_cache`), so the legacy "overlap store cache" flag # has no effect here -- the fused path is on by default. @@ -443,6 +478,7 @@ class MQALayer(nn.Module): forward_batch: ForwardBatch, attn_backend, q_out: Optional[torch.Tensor] = None, + x_quant=None, ) -> torch.Tensor: assert self.alt_streams is not None assert len(self.alt_streams) >= 3 @@ -456,13 +492,14 @@ class MQALayer(nn.Module): stream_compressor.wait_stream(current_stream) stream_indexer.wait_stream(current_stream) + x_linear = x_quant if x_quant is not None else x qkv_a: Optional[torch.Tensor] = None qkv_a_ready: Optional[torch.cuda.Event] = None if self.fuse_wqa_wkv: - qkv_a, _ = self.wqkv_a(x) + qkv_a, _ = self.wqkv_a(x_linear) qkv_a_ready = current_stream.record_event() - q_lora = self._compute_q_a(x, qkv_a=qkv_a) + q_lora = self._compute_q_a(x_linear, qkv_a=qkv_a) q_lora_ready = current_stream.record_event() if self.indexer is not None: @@ -480,7 +517,7 @@ class MQALayer(nn.Module): if qkv_a_ready is not None: stream_kv.wait_event(qkv_a_ready) # Fused norm + rope + cache write -- no bf16 KV intermediate. - self._compute_kv_to_cache(x, positions, forward_batch, qkv_a=qkv_a) + self._compute_kv_to_cache(x_linear, positions, forward_batch, qkv_a=qkv_a) del qkv_a @@ -504,36 +541,100 @@ class MQALayer(nn.Module): forward_batch: ForwardBatch, attn_backend, q_out: Optional[torch.Tensor] = None, + x_quant=None, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + x_linear = x_quant if x_quant is not None else x if self.fuse_wqa_wkv: - qkv_a, _ = self.wqkv_a(x) + qkv_a, _ = self.wqkv_a(x_linear) q_lora = qkv_a[..., : self.q_lora_rank] else: - q_lora, _ = self.wq_a(x) + q_lora, _ = self.wq_a(x_linear) qkv_a = None - q_lora = self.q_norm(q_lora) - q = self._compute_q_b(q_lora, positions, q_out) use_cp = self.dsa_enable_prefill_cp and dsa_use_prefill_cp(forward_batch) kv: Optional[torch.Tensor] - if use_cp: - # DSA CP: keep bf16 kv around for the cross-rank all-gather, then - # write to the FlashMLA cache after gather. - kv = self._compute_kv_bf16(x, positions, qkv_a=qkv_a) - kv = cp_all_gather_rerange_output( - kv.contiguous(), - self.cp_size, - forward_batch, - torch.cuda.current_stream(), + + if self.use_fused_qk_norm_rope: + + if _is_gfx95_supported: + q_for_wqb, q_lora = _fused_rmsnorm_fp8_quant( + q_lora, + self.q_norm.weight, + self.q_norm.variance_epsilon, + ) + q, _ = self.wq_b(q_for_wqb) + else: + q_lora = self.q_norm(q_lora) + q, _ = self.wq_b(q_lora) + + kv = ( + qkv_a[..., self.q_lora_rank :] + if qkv_a is not None + else self.wkv(x_linear)[0] ) - attn_backend.store_cache( - layer_id=self.layer_id, - swa_k=kv, - forward_batch=forward_batch, + + from sglang.srt.layers.fused_qk_norm_rope_store import ( + fused_qk_norm_rope_swa_store, ) + + token_to_kv_pool = get_token_to_kv_pool() + swa_loc = token_to_kv_pool.translate_loc_from_full_to_swa( + forward_batch.out_cache_loc + ) + swa_cache = token_to_kv_pool.swa_kv_pool.kv_buffer[self.layer_id] + swa_page_size = token_to_kv_pool.swa_kv_pool.page_size + + q = fused_qk_norm_rope_swa_store( + q=q, + kv=kv, + q_norm_weight=None, + kv_norm_weight=self.kv_norm.weight, + q_rms_eps=self.eps, + kv_rms_eps=self.eps, + rope_head_dim=self.qk_rope_head_dim, + cos_cache=self.cos_cache, + sin_cache=self.sin_cache, + positions=positions, + swa_cache=swa_cache, + swa_loc=swa_loc, + swa_page_size=swa_page_size, + q_out=q_out, + dtype=x.dtype, + ) + + if use_cp: + # DSA CP: keep bf16 kv around for the cross-rank all-gather, then + # write to the FlashMLA cache after gather. + kv = self._compute_kv_bf16(x, positions, qkv_a=qkv_a) + kv = cp_all_gather_rerange_output( + kv.contiguous(), + self.cp_size, + forward_batch, + torch.cuda.current_stream(), + ) else: - self._compute_kv_to_cache(x, positions, forward_batch, qkv_a=qkv_a) - kv = None + q_lora = self.q_norm(q_lora) + q = self._compute_q_b(q_lora, positions, q_out) + if use_cp: + # NSA CP: keep bf16 kv around for the cross-rank all-gather, then + # write to the FlashMLA cache after gather. + kv = self._compute_kv_bf16(x_linear, positions, qkv_a=qkv_a) + kv = cp_all_gather_rerange_output( + kv.contiguous(), + self.cp_size, + forward_batch, + torch.cuda.current_stream(), + ) + attn_backend.store_cache( + layer_id=self.layer_id, + swa_k=kv, + forward_batch=forward_batch, + ) + else: + self._compute_kv_to_cache( + x_linear, positions, forward_batch, qkv_a=qkv_a + ) + kv = None del qkv_a @@ -559,6 +660,7 @@ class MQALayer(nn.Module): x: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, + x_quant=None, ) -> torch.Tensor: if not get_attn_tp_context().input_scattered and x.shape[0] == 0: assert ( @@ -592,12 +694,22 @@ class MQALayer(nn.Module): # Multi-stream path always fuses cache write into the K kernel, # so the bf16 KV intermediate is gone. q = self._forward_prepare_multi_stream( - x, positions, forward_batch, attn_backend, q_out + x, + positions, + forward_batch, + attn_backend, + q_out, + x_quant=x_quant, ) kv = None else: q, kv = self._forward_prepare( - x, positions, forward_batch, attn_backend, q_out + x, + positions, + forward_batch, + attn_backend, + q_out, + x_quant=x_quant, ) # The cache write is always fused / already done by _forward_prepare* -- @@ -924,12 +1036,23 @@ class DeepseekV4DecoderLayer(nn.Module): norm=self.input_layernorm, ) if not norm_fused: - hidden_states = self.input_layernorm(hidden_states) + if _use_aiter and _is_gfx95_supported: + x_quant, hidden_states = _fused_rmsnorm_fp8_quant( + hidden_states, + self.input_layernorm.weight, + self.rms_norm_eps, + ) + else: + hidden_states = self.input_layernorm(hidden_states) + x_quant = None + else: + x_quant = None hidden_states = self.self_attn( x=hidden_states, positions=positions, forward_batch=forward_batch, + x_quant=x_quant, ) hidden_states = self.hc_post(hidden_states, residual, post, comb) @@ -1022,9 +1145,7 @@ class DeepseekV4Model(nn.Module): else: self.embed_tokens = PPMissingLayer() self.rms_norm_eps = config.rms_norm_eps - self.alt_streams = ( - [torch.cuda.Stream() for _ in range(5)] if (_is_cuda or _is_hip) else None - ) + self.alt_streams = [torch.cuda.Stream() for _ in range(5)] if _is_cuda else None self.layers, self.start_layer, self.end_layer = make_layers( config.num_hidden_layers, lambda idx, prefix: DeepseekV4DecoderLayer( @@ -1138,6 +1259,10 @@ class DeepseekV4Model(nn.Module): hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states) positions = cp_split_and_rebuild_position(forward_batch, positions) + # Reset Compressor's per-step freqs_cis cache from any previous step. + for _attr in ("freqs_cis_c4", "freqs_cis_c128"): + if hasattr(forward_batch, _attr): + delattr(forward_batch, _attr) # Upgrade lazy raw metadata on the main stream once before any layer # forks alt-streams; later per-layer calls become no-ops. get_attn_backend()._maybe_upgrade_forward_metadata() diff --git a/sgl-kernel/CMakeLists.txt b/sgl-kernel/CMakeLists.txt index 8ede4482c..86d315783 100644 --- a/sgl-kernel/CMakeLists.txt +++ b/sgl-kernel/CMakeLists.txt @@ -258,6 +258,7 @@ set(SOURCES "csrc/elementwise/activation.cu" "csrc/elementwise/concat_mla.cu" "csrc/elementwise/copy.cu" + "csrc/elementwise/dsv4_norm_rope.cu" "csrc/elementwise/fused_add_rms_norm_kernel.cu" "csrc/elementwise/pos_enc.cu" "csrc/elementwise/topk.cu" diff --git a/sgl-kernel/csrc/common_extension_rocm.cc b/sgl-kernel/csrc/common_extension_rocm.cc index 1c7265b2b..cad9c95d9 100644 --- a/sgl-kernel/csrc/common_extension_rocm.cc +++ b/sgl-kernel/csrc/common_extension_rocm.cc @@ -47,6 +47,25 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) { "topk_indices_offset, Tensor ? row_starts) -> ()"); m.impl("fast_topk_transform_ragged_fused", torch::kCUDA, &fast_topk_transform_ragged_interface); + m.def( + "deepseek_v4_topk_transform_512(Tensor scores, Tensor seq_lens, Tensor page_table, Tensor! " + "page_indices, int page_size, Tensor!? raw_indices) -> ()"); + m.impl("deepseek_v4_topk_transform_512", torch::kCUDA, &deepseek_v4_topk_transform_512); + + m.def( + "dsv4_fused_q_norm_rope(Tensor q_input, Tensor! q_output, Tensor freqs_cis, Tensor positions, float eps) -> ()"); + m.impl("dsv4_fused_q_norm_rope", torch::kCUDA, &dsv4_fused_q_norm_rope); + + m.def( + "dsv4_fused_k_norm_rope_flashmla(Tensor kv, Tensor kv_weight, Tensor freqs_cis, Tensor positions, " + "Tensor out_loc, Tensor! kvcache, float eps, int page_size) -> ()"); + m.impl("dsv4_fused_k_norm_rope_flashmla", torch::kCUDA, &dsv4_fused_k_norm_rope_flashmla); + + m.def( + "dsv4_fused_q_indexer_rope_hadamard_quant(Tensor q_input, Tensor! q_fp8, Tensor weight, " + "Tensor! weights_out, float weight_scale, Tensor freqs_cis, Tensor positions) -> ()"); + m.impl("dsv4_fused_q_indexer_rope_hadamard_quant", torch::kCUDA, &dsv4_fused_q_indexer_rope_hadamard_quant); + /* * From csrc/allreduce */ diff --git a/sgl-kernel/csrc/elementwise/deepseek_v4_topk.cu b/sgl-kernel/csrc/elementwise/deepseek_v4_topk.cu new file mode 100644 index 000000000..5262af88e --- /dev/null +++ b/sgl-kernel/csrc/elementwise/deepseek_v4_topk.cu @@ -0,0 +1,372 @@ +/* Copyright 2025 SGLang Team. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace { + +constexpr uint32_t kMaxTopK = 1024; +constexpr uint32_t kBlockSize = 512; + +#ifdef SGL_TOPK_DYNAMIC_SMEM_BYTES +constexpr size_t kSMEM = static_cast(SGL_TOPK_DYNAMIC_SMEM_BYTES); +#else +constexpr size_t kSMEM = 48 * 1024; // bytes +#endif +static_assert(kSMEM % (2 * sizeof(int32_t)) == 0, "kSMEM must be a multiple of 8 bytes."); + +struct TopKParams { + const float* __restrict__ scores; + const int32_t* __restrict__ seq_lens; + const int32_t* __restrict__ page_table; + int32_t* __restrict__ page_indices; + int32_t* __restrict__ raw_indices; + int64_t score_stride; + int64_t page_table_stride; + uint32_t page_bits; + uint32_t topk; + int64_t output_stride; +}; + +__device__ __forceinline__ uint8_t convert_to_uint8(float x) { + __half h = __float2half_rn(x); + uint16_t bits = __half_as_ushort(h); + uint16_t key = (bits & 0x8000) ? static_cast(~bits) : static_cast(bits | 0x8000); + return static_cast(key >> 8); +} + +__device__ __forceinline__ uint32_t convert_to_uint32(float x) { + uint32_t bits = __float_as_uint(x); + return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u); +} + +__device__ __forceinline__ int32_t +page_to_slot(const int32_t* __restrict__ page_table, uint32_t i, uint32_t page_bits) { + const uint32_t mask = (1u << page_bits) - 1u; + return (page_table[i >> page_bits] << page_bits) | static_cast(i & mask); +} + +__device__ void naive_paged_transform( + int32_t length, + uint32_t topk, + uint32_t page_bits, + const int32_t* __restrict__ page_table, + int32_t* __restrict__ page_indices_out, + int32_t* __restrict__ raw_indices_out) { + for (uint32_t i = threadIdx.x; i < topk; i += kBlockSize) { + if (i < static_cast(length)) { + page_indices_out[i] = page_to_slot(page_table, i, page_bits); + if (raw_indices_out != nullptr) { + raw_indices_out[i] = static_cast(i); + } + } else { + page_indices_out[i] = -1; + if (raw_indices_out != nullptr) { + raw_indices_out[i] = -1; + } + } + } +} + +__device__ void +radix_topk(const float* __restrict__ input, int32_t* __restrict__ output, uint32_t length, uint32_t topk) { + constexpr uint32_t RADIX = 256; + constexpr uint32_t BLOCK_SIZE = kBlockSize; + constexpr uint32_t SMEM_INPUT_SIZE = kSMEM / (2 * sizeof(int32_t)); + + alignas(128) __shared__ uint32_t _s_histogram_buf[2][RADIX + 32]; + alignas(128) __shared__ uint32_t s_counter; + alignas(128) __shared__ uint32_t s_threshold_bin_id; + alignas(128) __shared__ uint32_t s_num_input[2]; + alignas(128) __shared__ int32_t s_last_remain; + + extern __shared__ uint32_t s_input_idx[][SMEM_INPUT_SIZE]; + + const uint32_t tx = threadIdx.x; + uint32_t remain_topk = topk; + auto& s_histogram = _s_histogram_buf[0]; + + const auto run_cumsum = [&] { +#pragma unroll 8 + for (int32_t i = 0; i < 8; ++i) { + static_assert(1 << 8 == RADIX); + if (tx < RADIX) { + const auto j = 1 << i; + const auto k = i & 1; + auto value = _s_histogram_buf[k][tx]; + if (tx + j < RADIX) { + value += _s_histogram_buf[k][tx + j]; + } + _s_histogram_buf[k ^ 1][tx] = value; + } + __syncthreads(); + } + }; + + // stage 1: 8bit coarse histogram + if (tx < RADIX + 1) s_histogram[tx] = 0; + __syncthreads(); + for (uint32_t idx = tx; idx < length; idx += BLOCK_SIZE) { + const auto bin = convert_to_uint8(input[idx]); + ::atomicAdd(&s_histogram[bin], 1); + } + __syncthreads(); + run_cumsum(); + if (tx < RADIX && s_histogram[tx] > remain_topk && s_histogram[tx + 1] <= remain_topk) { + s_threshold_bin_id = tx; + s_num_input[0] = 0; + s_counter = 0; + } + __syncthreads(); + + { + const auto threshold_bin = s_threshold_bin_id; + remain_topk -= s_histogram[threshold_bin + 1]; + if (remain_topk == 0) { + for (uint32_t idx = tx; idx < length; idx += BLOCK_SIZE) { + const uint32_t bin = convert_to_uint8(input[idx]); + if (bin > threshold_bin) { + const auto pos = ::atomicAdd(&s_counter, 1); + output[pos] = static_cast(idx); + } + } + __syncthreads(); + return; + } + __syncthreads(); + if (tx < RADIX + 1) s_histogram[tx] = 0; + __syncthreads(); + + for (uint32_t idx = tx; idx < length; idx += BLOCK_SIZE) { + const float raw_input = input[idx]; + const uint32_t bin = convert_to_uint8(raw_input); + if (bin > threshold_bin) { + const auto pos = ::atomicAdd(&s_counter, 1); + output[pos] = static_cast(idx); + } else if (bin == threshold_bin) { + const auto pos = ::atomicAdd(&s_num_input[0], 1); + if (C10_LIKELY(pos < SMEM_INPUT_SIZE)) { + s_input_idx[0][pos] = idx; + const auto bin32 = convert_to_uint32(raw_input); + const auto sub_bin = (bin32 >> 24) & 0xFF; + ::atomicAdd(&s_histogram[sub_bin], 1); + } + } + } + __syncthreads(); + } + + // stage 2: refine with 8bit radix passes +#pragma unroll 4 + for (int round = 0; round < 4; ++round) { + const auto r_idx = round % 2; + + const auto raw_num_input = s_num_input[r_idx]; + const auto num_input = raw_num_input < SMEM_INPUT_SIZE ? raw_num_input : SMEM_INPUT_SIZE; + + run_cumsum(); + if (tx < RADIX && s_histogram[tx] > remain_topk && s_histogram[tx + 1] <= remain_topk) { + s_threshold_bin_id = tx; + s_num_input[r_idx ^ 1] = 0; + s_last_remain = static_cast(remain_topk - s_histogram[tx + 1]); + } + __syncthreads(); + + const auto threshold_bin = s_threshold_bin_id; + remain_topk -= s_histogram[threshold_bin + 1]; + + if (remain_topk == 0) { + for (uint32_t i = tx; i < num_input; i += BLOCK_SIZE) { + const auto idx = s_input_idx[r_idx][i]; + const auto offset = 24 - round * 8; + const auto bin = (convert_to_uint32(input[idx]) >> offset) & 0xFF; + if (bin > threshold_bin) { + const auto pos = ::atomicAdd(&s_counter, 1); + output[pos] = static_cast(idx); + } + } + __syncthreads(); + break; + } + __syncthreads(); + if (tx < RADIX + 1) s_histogram[tx] = 0; + __syncthreads(); + for (uint32_t i = tx; i < num_input; i += BLOCK_SIZE) { + const auto idx = s_input_idx[r_idx][i]; + const auto raw_input = input[idx]; + const auto offset = 24 - round * 8; + const auto bin = (convert_to_uint32(raw_input) >> offset) & 0xFF; + if (bin > threshold_bin) { + const auto pos = ::atomicAdd(&s_counter, 1); + output[pos] = static_cast(idx); + } else if (bin == threshold_bin) { + if (round == 3) { + const auto pos = ::atomicAdd(&s_last_remain, -1); + if (pos > 0) { + output[topk - pos] = static_cast(idx); + } + } else { + const auto pos = ::atomicAdd(&s_num_input[r_idx ^ 1], 1); + if (C10_LIKELY(pos < SMEM_INPUT_SIZE)) { + s_input_idx[r_idx ^ 1][pos] = idx; + const auto bin32 = convert_to_uint32(raw_input); + const auto sub_bin = (bin32 >> (offset - 8)) & 0xFF; + ::atomicAdd(&s_histogram[sub_bin], 1); + } + } + } + } + __syncthreads(); + } +} + +__global__ __launch_bounds__(kBlockSize) void deepseek_v4_topk_transform_kernel(const TopKParams params) { + const auto bid = blockIdx.x; + const auto seq_len = params.seq_lens[bid]; + const auto topk = params.topk; + const auto score_ptr = params.scores + bid * params.score_stride; + const auto page_ptr = params.page_table + bid * params.page_table_stride; + const auto indices_ptr = params.page_indices + bid * params.output_stride; + const auto raw_indices_ptr = + params.raw_indices != nullptr ? params.raw_indices + bid * params.output_stride : nullptr; + + if (seq_len <= static_cast(topk)) { + naive_paged_transform(seq_len, topk, params.page_bits, page_ptr, indices_ptr, raw_indices_ptr); + return; + } + + __shared__ int32_t s_topk_indices[kMaxTopK]; + radix_topk(score_ptr, s_topk_indices, static_cast(seq_len), topk); + + __syncthreads(); + for (uint32_t i = threadIdx.x; i < topk; i += kBlockSize) { + const auto raw = s_topk_indices[i]; + indices_ptr[i] = page_to_slot(page_ptr, static_cast(raw), params.page_bits); + if (raw_indices_ptr != nullptr) { + raw_indices_ptr[i] = raw; + } + } +} + +template +void setup_kernel_smem_once() { + [[maybe_unused]] + static const auto result = [] { +#ifdef USE_ROCM + return ::cudaFuncSetAttribute( + reinterpret_cast(f), ::cudaFuncAttributeMaxDynamicSharedMemorySize, kMaxDynamicSMEM); +#else + return ::cudaFuncSetAttribute(f, ::cudaFuncAttributeMaxDynamicSharedMemorySize, kMaxDynamicSMEM); +#endif + }(); + TORCH_CHECK( + result == cudaSuccess, "deepseek_v4_topk_transform: cudaFuncSetAttribute failed: ", ::cudaGetErrorString(result)); +} + +} // namespace + +#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor") + +void deepseek_v4_topk_transform_512( + const at::Tensor& scores, + const at::Tensor& seq_lens, + const at::Tensor& page_table, + at::Tensor& page_indices, + int64_t page_size, + std::optional raw_indices_opt) { + CHECK_CUDA(scores); + CHECK_CUDA(seq_lens); + CHECK_CUDA(page_table); + CHECK_CUDA(page_indices); + if (raw_indices_opt.has_value()) { + CHECK_CUDA(raw_indices_opt.value()); + } + + TORCH_CHECK( + scores.dim() == 2 && scores.scalar_type() == at::kFloat, "scores must be float32 with shape [B, max_seq_len]"); + TORCH_CHECK(scores.stride(1) == 1, "scores must be contiguous along the last dim"); + + TORCH_CHECK( + seq_lens.dim() == 1 && seq_lens.is_contiguous() && seq_lens.scalar_type() == at::kInt, + "seq_lens must be int32 with shape [B], contiguous"); + + TORCH_CHECK( + page_table.dim() == 2 && page_table.scalar_type() == at::kInt, + "page_table must be int32 with shape [B, num_pages]"); + TORCH_CHECK(page_table.stride(1) == 1, "page_table must be contiguous along the last dim"); + + const auto topk = page_indices.size(1); + TORCH_CHECK( + page_indices.dim() == 2 && page_indices.is_contiguous() && page_indices.scalar_type() == at::kInt, + "page_indices must be int32 with shape [B, topk], contiguous"); + TORCH_CHECK( + topk > 0 && topk <= static_cast(kMaxTopK), + "page_indices last dim must be in [1, ", + kMaxTopK, + "], got ", + topk); + + const auto B = scores.size(0); + TORCH_CHECK( + seq_lens.size(0) == B && page_table.size(0) == B && page_indices.size(0) == B, + "batch sizes must match across scores, seq_lens, page_table, page_indices"); + + TORCH_CHECK( + page_size > 0 && (page_size & (page_size - 1)) == 0, "page_size must be a positive power of 2, got ", page_size); + const auto page_bits = static_cast(__builtin_ctzll(static_cast(page_size))); + + int32_t* raw_ptr = nullptr; + if (raw_indices_opt.has_value()) { + auto& raw = raw_indices_opt.value(); + TORCH_CHECK( + raw.dim() == 2 && raw.is_contiguous() && raw.scalar_type() == at::kInt, + "raw_indices must be int32 with shape [B, topk], contiguous"); + TORCH_CHECK(raw.size(0) == B && raw.size(1) == topk, "raw_indices shape must match page_indices [B, ", topk, "]"); + raw_ptr = raw.data_ptr(); + } + + const TopKParams params{ + .scores = scores.data_ptr(), + .seq_lens = seq_lens.data_ptr(), + .page_table = page_table.data_ptr(), + .page_indices = page_indices.data_ptr(), + .raw_indices = raw_ptr, + .score_stride = scores.stride(0), + .page_table_stride = page_table.stride(0), + .page_bits = page_bits, + .topk = static_cast(topk), + .output_stride = topk, + }; + + const auto stream = at::cuda::getCurrentCUDAStream().stream(); + const dim3 grid(static_cast(B)); + const dim3 block(kBlockSize); + + setup_kernel_smem_once(); + deepseek_v4_topk_transform_kernel<<>>(params); + + const auto err = cudaGetLastError(); + TORCH_CHECK(err == cudaSuccess, "deepseek_v4_topk_transform kernel launch failed: ", ::cudaGetErrorString(err)); +} diff --git a/sgl-kernel/csrc/elementwise/dsv4_norm_rope.cu b/sgl-kernel/csrc/elementwise/dsv4_norm_rope.cu new file mode 100644 index 000000000..935278d4b --- /dev/null +++ b/sgl-kernel/csrc/elementwise/dsv4_norm_rope.cu @@ -0,0 +1,700 @@ +/* Copyright 2025 SGLang Team. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +// DeepSeek-V4 fused norm + RoPE kernels, ported from JIT kernel +// python/sglang/jit_kernel/csrc/deepseek_v4/main_norm_rope.cuh +// to sgl-kernel AOT compilation with CUDA + HIP (ROCm) support. + +#ifndef USE_ROCM +#include +#include +#include +#include +#else +#include +#include +#include +#endif + +#include +#include +#include + +#include + +#include "utils.h" + +// ============================================================================ +// Platform-compatible type aliases +// ============================================================================ +#ifndef USE_ROCM +using bf16_t = __nv_bfloat16; +using bf16x2_t = __nv_bfloat162; +using fp8x2_e4m3_t = __nv_fp8x2_e4m3; +#else +using bf16_t = __hip_bfloat16; +using bf16x2_t = __hip_bfloat162; +using fp8x2_e4m3_t = uint16_t; +#ifndef __grid_constant__ +#define __grid_constant__ +#endif +#endif + +// ============================================================================ +// Utility helpers (inlined, no external header dependency) +// ============================================================================ + +static constexpr uint32_t kWarpSize = 32; + +template +__device__ __forceinline__ float warp_reduce_sum(float val) { +#pragma unroll + for (uint32_t mask = kNumThreads / 2; mask > 0; mask >>= 1) + val += SGLANG_SHFL_XOR_SYNC(FULL_MASK, val, mask); + return val; +} + +__device__ __forceinline__ float warp_reduce_max(float val) { +#pragma unroll + for (uint32_t mask = kWarpSize / 2; mask > 0; mask >>= 1) + val = fmaxf(val, SGLANG_SHFL_XOR_SYNC(FULL_MASK, val, mask)); + return val; +} + +// Aligned vector for coalesced memory access. +template +struct alignas(sizeof(T) * N) AlignedVec { + T data[N]; + __device__ __forceinline__ T& operator[](int i) { + return data[i]; + } + __device__ __forceinline__ T operator[](int i) const { + return data[i]; + } + __device__ __forceinline__ void load(const void* ptr, int64_t offset = 0) { + *this = reinterpret_cast(ptr)[offset]; + } + __device__ __forceinline__ void store(void* ptr, int64_t offset = 0) const { + reinterpret_cast(ptr)[offset] = *this; + } +}; + +__device__ __forceinline__ float bf16_to_float(bf16_t v) { + return __bfloat162float(v); +} + +__device__ __forceinline__ bf16_t float_to_bf16(float v) { +#ifndef USE_ROCM + return __float2bfloat16_rn(v); +#else + return __float2bfloat16(v); +#endif +} + +// ============================================================================ +// FP8 E4M3 helpers (portable CUDA + HIP) +// ============================================================================ + +// UE8M0 scale: round a positive float to the nearest power-of-two +// representable in UE8M0 (unsigned 8-bit exponent, no mantissa). +__device__ __forceinline__ int32_t cast_to_ue8m0(float x) { + uint32_t u = __float_as_uint(x); + int32_t exp = static_cast((u >> 23) & 0xFFu); + uint32_t mant = u & 0x7FFFFFu; + return exp + (mant != 0); +} + +__device__ __forceinline__ float inv_scale_ue8m0(int32_t exp) { + return __uint_as_float(static_cast((127 + 127 - exp) << 23)); +} + +static constexpr float kFP8Max = 448.0f; + +#ifndef USE_ROCM +__device__ __forceinline__ fp8x2_e4m3_t pack_fp8(float x, float y) { + x = fmaxf(fminf(x, kFP8Max), -kFP8Max); + y = fmaxf(fminf(y, kFP8Max), -kFP8Max); + return __nv_fp8x2_e4m3(float2{x, y}); +} +#else +// Software float -> FP8 E4M3 conversion for ROCm +__device__ __forceinline__ uint8_t cvt_float_to_fp8_e4m3(float val) { + constexpr float kMax = kFP8Max; + val = fmaxf(fminf(val, kMax), -kMax); + if (val == 0.0f) return 0; + + uint32_t f32 = __float_as_uint(val); + uint8_t sign = static_cast((f32 >> 24) & 0x80u); + f32 &= 0x7FFFFFFFu; + + int32_t exp32 = static_cast((f32 >> 23) & 0xFFu); + uint32_t mant32 = f32 & 0x7FFFFFu; + + // FP8 E4M3 bias=7, FP32 bias=127, offset=120 + int32_t exp8 = exp32 - 120; + + if (exp8 <= 0) { + mant32 |= 0x800000u; + int32_t shift = 1 - exp8; + if (shift > 24) return sign; + uint32_t shifted = mant32 >> (20 + shift); + uint32_t rbit = (shift <= 23) ? ((mant32 >> (19 + shift)) & 1u) : 0u; + uint32_t sbit = (shift <= 23) ? ((mant32 & ((1u << (19 + shift)) - 1u)) != 0) : 0u; + shifted += (rbit && (sbit || (shifted & 1u))); + return sign | static_cast(shifted & 0x7u); + } + if (exp8 >= 15) return sign | 0x7Eu; + + uint32_t mant3 = (mant32 >> 20) & 0x7u; + uint32_t rbit = (mant32 >> 19) & 1u; + uint32_t sbit = (mant32 & 0x7FFFFu) != 0; + mant3 += (rbit && (sbit || (mant3 & 1u))); + if (mant3 > 7) { + mant3 = 0; + exp8++; + if (exp8 >= 15) return sign | 0x7Eu; + } + return sign | (static_cast(exp8) << 3) | static_cast(mant3); +} + +__device__ __forceinline__ fp8x2_e4m3_t pack_fp8(float x, float y) { + uint8_t x8 = cvt_float_to_fp8_e4m3(x); + uint8_t y8 = cvt_float_to_fp8_e4m3(y); + return static_cast(x8) | (static_cast(y8) << 8); +} +#endif + +// ============================================================================ +// Kernel 1: Fused Q Norm + RoPE +// warp-per-(token, head), rmsnorm-self (no weight) + RoPE + write to q_out. +// ============================================================================ + +namespace { + +constexpr uint32_t kFusedQBlockSize = 128; +constexpr uint32_t kFusedQNumWarps = kFusedQBlockSize / kWarpSize; + +constexpr uint32_t kFusedKBlockSize = 256; +constexpr uint32_t kFusedKNumWarps = kFusedKBlockSize / kWarpSize; + +struct FusedQNormRopeParams { + const void* __restrict__ q_input; + void* __restrict__ q_output; + const float* __restrict__ freqs_cis; + const int32_t* __restrict__ positions; + int64_t q_input_stride_batch; + int64_t q_output_stride_batch; + uint32_t batch_size; + uint32_t num_q_heads; + float eps; +}; + +// Compute the largest power-of-2 vec size that divides both kHeadDim and +// fits in 16 bytes, while also dividing kRopeDim. +template +struct QKernelTraits { + static constexpr int64_t kMaxVecSize = 16 / sizeof(bf16_t); // 8 + // Use kRopeDim/kWarpSize (=2 for kRopeDim=64) as the vec size. + // This guarantees kRopeDim % kVecSize == 0 and works for all head dims + // that are multiples of kWarpSize*kVecSize. + static constexpr int64_t kVecSize = kRopeDim / kWarpSize; // 2 + static constexpr int64_t kLocalSize = kHeadDim / (kWarpSize * kVecSize); + static constexpr uint32_t kRopeSize = kRopeDim / kVecSize; + static_assert(kHeadDim % (kWarpSize * kVecSize) == 0); + static_assert(kRopeDim % kVecSize == 0); + static_assert(kRopeDim == kWarpSize * 2, "1 (real, imag) pair per lane"); +}; + +template +__global__ __launch_bounds__(kFusedQBlockSize, 16) void fused_q_norm_rope_kernel( + const __grid_constant__ FusedQNormRopeParams params) { + using Traits = QKernelTraits; + constexpr int64_t kVecSize = Traits::kVecSize; + constexpr int64_t kLocalSize = Traits::kLocalSize; + constexpr uint32_t kRopeSize = Traits::kRopeSize; + + using Storage = AlignedVec; + using Float2 = AlignedVec; + + const auto warp_id = threadIdx.x / kWarpSize; + const auto lane_id = threadIdx.x % kWarpSize; + const auto work_id = blockIdx.x * kFusedQNumWarps + warp_id; + + const uint32_t total_works = params.batch_size * params.num_q_heads; + if (work_id >= total_works) return; + + const uint32_t batch_id = work_id / params.num_q_heads; + const uint32_t head_id = work_id % params.num_q_heads; + const auto input_ptr = + static_cast(params.q_input) + batch_id * params.q_input_stride_batch + head_id * kHeadDim; + const auto output_ptr = + static_cast(params.q_output) + batch_id * params.q_output_stride_batch + head_id * kHeadDim; + const auto position = params.positions[batch_id]; + + __shared__ Storage s_rope[kFusedQNumWarps][kRopeSize]; + + // Prefetch freq pair. + Float2 freq; + freq.load(params.freqs_cis + position * kRopeDim, lane_id); + + // Part 1: rmsnorm-self (no weight). + Storage input_vec[kLocalSize]; +#pragma unroll + for (int i = 0; i < kLocalSize; ++i) { + input_vec[i].load(input_ptr, lane_id + i * kWarpSize); + } + + float sum_of_squares = 0.0f; +#pragma unroll + for (int i = 0; i < kLocalSize; ++i) { +#pragma unroll + for (int j = 0; j < kVecSize; ++j) { + float x = bf16_to_float(input_vec[i][j]); + sum_of_squares += x * x; + } + } + sum_of_squares = warp_reduce_sum(sum_of_squares); + const float norm_factor = rsqrtf(sum_of_squares / static_cast(kHeadDim) + params.eps); + +#pragma unroll + for (int i = 0; i < kLocalSize; ++i) { +#pragma unroll + for (int j = 0; j < kVecSize; ++j) { + float x = bf16_to_float(input_vec[i][j]); + input_vec[i][j] = float_to_bf16(x * norm_factor); + } + } + + // Stash rope tail into shared memory; write nope tiles to gmem. + const bool is_rope_lane = lane_id >= kWarpSize - kRopeSize; +#pragma unroll + for (int i = 0; i < kLocalSize; ++i) { + if (i == kLocalSize - 1 && is_rope_lane) { + const auto rope_id = lane_id - (kWarpSize - kRopeSize); + s_rope[warp_id][rope_id] = input_vec[i]; + } else { + input_vec[i].store(output_ptr, lane_id + i * kWarpSize); + } + } + __syncwarp(); + + // Part 2: RoPE on all 32 lanes -- one (real, imag) bf16x2 pair per lane. + auto elem_ptr = reinterpret_cast(&s_rope[warp_id][0]); + bf16x2_t elem = elem_ptr[lane_id]; +#ifndef USE_ROCM + float2 elem_f = __bfloat1622float2(elem); + float x_real = elem_f.x, x_imag = elem_f.y; +#else + float x_real = __bfloat162float(elem.x), x_imag = __bfloat162float(elem.y); +#endif + float freq_real = freq[0], freq_imag = freq[1]; + float rot_real = x_real * freq_real - x_imag * freq_imag; + float rot_imag = x_real * freq_imag + x_imag * freq_real; + bf16x2_t rotated = __float22bfloat162_rn(make_float2(rot_real, rot_imag)); + auto out_elem = reinterpret_cast(output_ptr + (kHeadDim - kRopeDim)); + out_elem[lane_id] = rotated; +} + +// ============================================================================ +// Kernel 2: Fused K Norm + RoPE + FlashMLA Store +// block-per-token, rmsnorm (with kv_weight) + RoPE + FP8 quantized store. +// ============================================================================ + +struct FusedKNormRopeFlashMLAParams { + const void* __restrict__ kv; + const void* __restrict__ kv_weight; + const float* __restrict__ freqs_cis; + const int32_t* __restrict__ positions; + const int32_t* __restrict__ out_loc; + uint8_t* __restrict__ kvcache; + int64_t kv_stride_batch; + uint32_t batch_size; + float eps; +}; + +template +__global__ __launch_bounds__(kFusedKBlockSize, 8) void fused_k_norm_rope_flashmla_kernel( + const __grid_constant__ FusedKNormRopeFlashMLAParams params) { + constexpr int64_t kVecSize = 2; + constexpr uint32_t kRopeWarp = kFusedKNumWarps - 1; + constexpr int64_t kPageBytes = ((584ll << kPageBits) + 575) / 576 * 576; + static_assert(kHeadDim == kFusedKBlockSize * kVecSize); + static_assert(kRopeDim == kWarpSize * kVecSize); + + using Storage = AlignedVec; + + const auto tx = threadIdx.x; + const auto warp_id = tx / kWarpSize; + const auto lane_id = tx % kWarpSize; + const auto work_id = blockIdx.x; + if (work_id >= params.batch_size) return; + + const auto input_ptr = static_cast(params.kv) + work_id * params.kv_stride_batch; + const auto position = params.positions[work_id]; + const auto out_loc = params.out_loc[work_id]; + const auto freqs_cis = params.freqs_cis + position * kRopeDim; + + AlignedVec data, freq; + + // Part 1: norm with block-wide reduction. + { + __shared__ float partial_sums[kFusedKNumWarps]; + + Storage input_vec, weight_vec; + input_vec.load(input_ptr, tx); + weight_vec.load(params.kv_weight, tx); + if (warp_id == kRopeWarp) freq.load(freqs_cis, lane_id); + + float sum_of_squares = 0.0f; +#pragma unroll + for (int i = 0; i < kVecSize; ++i) { + float x = bf16_to_float(input_vec[i]); + sum_of_squares += x * x; + } + const float warp_sum = warp_reduce_sum(sum_of_squares); + if (lane_id == 0) partial_sums[warp_id] = warp_sum; + __syncthreads(); + sum_of_squares = warp_reduce_sum(partial_sums[lane_id % kFusedKNumWarps]); + const float norm_factor = rsqrtf(sum_of_squares / static_cast(kHeadDim) + params.eps); + +#pragma unroll + for (int i = 0; i < kVecSize; ++i) { + float x = bf16_to_float(input_vec[i]); + float w = bf16_to_float(weight_vec[i]); + data[i] = x * norm_factor * w; + } + } + + const int32_t page = out_loc >> kPageBits; + const int32_t offset = out_loc & ((1 << kPageBits) - 1); + const auto page_ptr = params.kvcache + page * kPageBytes; + const auto value_ptr = page_ptr + offset * 576; + + // Part 2: rope on last warp (BF16 store), per-warp UE8M0 quant + store on others. + if (warp_id == kRopeWarp) { + float x_real = data[0], x_imag = data[1]; + float freq_real = freq[0], freq_imag = freq[1]; + float rot_real = x_real * freq_real - x_imag * freq_imag; + float rot_imag = x_real * freq_imag + x_imag * freq_real; + bf16x2_t result = __float22bfloat162_rn(make_float2(rot_real, rot_imag)); + auto rope_ptr = value_ptr + 448; + reinterpret_cast(rope_ptr)[lane_id] = result; + } else { + float x = data[0], y = data[1]; + float abs_max = warp_reduce_max(fmaxf(fabsf(x), fabsf(y))); + float scale_raw = fmaxf(1e-4f, abs_max) / kFP8Max; + int32_t scale_ue8m0 = cast_to_ue8m0(scale_raw); + float inv_scale = inv_scale_ue8m0(scale_ue8m0); + fp8x2_e4m3_t result = pack_fp8(x * inv_scale, y * inv_scale); + auto scale_ptr = page_ptr + (576ll << kPageBits) + offset * 8; + reinterpret_cast(value_ptr)[tx] = result; + if (lane_id == 0) static_cast(scale_ptr)[warp_id] = static_cast(scale_ue8m0); + } +} + +// ============================================================================ +// Kernel 3: Fused Q Indexer RoPE + Hadamard + FP8 Quantization +// warp-per-(token, head), no norm, RoPE + Hadamard + fp8 act-quant. +// ============================================================================ + +struct FusedQIndexerRopeHadamardQuantParams { + const void* __restrict__ q_input; + void* __restrict__ q_fp8; + const void* __restrict__ weight; + float* __restrict__ weights_out; + float weight_scale; + const float* __restrict__ freqs_cis; + const int32_t* __restrict__ positions; + uint32_t batch_size; + uint32_t num_heads; +}; + +__global__ __launch_bounds__(kFusedQBlockSize, 16) void fused_q_indexer_rope_hadamard_quant_kernel( + const __grid_constant__ FusedQIndexerRopeHadamardQuantParams params) { + constexpr int64_t kHeadDim = 128; + constexpr int64_t kRopeDim = 64; + constexpr int64_t kVecSize = 4; + constexpr uint32_t kRopeSize = kRopeDim / kVecSize; + static_assert(kHeadDim == kWarpSize * kVecSize); + + using Storage = AlignedVec; + using Float4 = AlignedVec; + using OutStorage = AlignedVec; + + const auto warp_id = threadIdx.x / kWarpSize; + const auto lane_id = threadIdx.x % kWarpSize; + const auto work_id = blockIdx.x * kFusedQNumWarps + warp_id; + const bool is_rope_lane = lane_id >= kWarpSize - kRopeSize; + + const uint32_t total_works = params.batch_size * params.num_heads; + if (work_id >= total_works) return; + + const uint32_t batch_id = work_id / params.num_heads; + const auto input_ptr = static_cast(params.q_input) + work_id * kHeadDim; + const auto position = params.positions[batch_id]; + const auto freqs_cis = params.freqs_cis + position * kRopeDim; + + Float4 data, freq; + const float weight_val = bf16_to_float(static_cast(params.weight)[work_id]); + + // Part 1: load (no norm). + { + Storage input_vec; + input_vec.load(input_ptr, lane_id); + if (is_rope_lane) freq.load(freqs_cis, lane_id - (kWarpSize - kRopeSize)); +#pragma unroll + for (int i = 0; i < kVecSize; ++i) + data[i] = bf16_to_float(input_vec[i]); + } + + // Part 2: rope on rope lanes. + if (is_rope_lane) { + float x_r = data[0], x_i = data[1], y_r = data[2], y_i = data[3]; + float fxr = freq[0], fxi = freq[1], fyr = freq[2], fyi = freq[3]; + data[0] = x_r * fxr - x_i * fxi; + data[1] = x_r * fxi + x_i * fxr; + data[2] = y_r * fyr - y_i * fyi; + data[3] = y_r * fyi + y_i * fyr; + } + + // Part 3: 128-point Hadamard (2 local + 5 cross-lane stages). + { + { + float a0 = data[0], a1 = data[1], a2 = data[2], a3 = data[3]; + data[0] = a0 + a1; + data[1] = a0 - a1; + data[2] = a2 + a3; + data[3] = a2 - a3; + } + { + float a0 = data[0], a1 = data[1], a2 = data[2], a3 = data[3]; + data[0] = a0 + a2; + data[1] = a1 + a3; + data[2] = a0 - a2; + data[3] = a1 - a3; + } +#pragma unroll + for (uint32_t mask = 1; mask < kWarpSize; mask <<= 1) { +#pragma unroll + for (int i = 0; i < kVecSize; ++i) { + float other = SGLANG_SHFL_XOR_SYNC_WIDTH(FULL_MASK, data[i], mask, kWarpSize); + data[i] = (lane_id & mask) ? (other - data[i]) : (data[i] + other); + } + } + const float kHadamardScale = rsqrtf(static_cast(kHeadDim)); +#pragma unroll + for (int i = 0; i < kVecSize; ++i) + data[i] *= kHadamardScale; + } + + // Part 4: per-warp FP8 quant + store. + { + float local_max = fabsf(data[0]); +#pragma unroll + for (int i = 1; i < kVecSize; ++i) + local_max = fmaxf(local_max, fabsf(data[i])); + float abs_max = warp_reduce_max(local_max); + float scale = fmaxf(1e-4f, abs_max) / kFP8Max; + float inv_scale = 1.0f / scale; + + OutStorage result; + result[0] = pack_fp8(data[0] * inv_scale, data[1] * inv_scale); + result[1] = pack_fp8(data[2] * inv_scale, data[3] * inv_scale); + + auto out_row = static_cast(params.q_fp8) + work_id * kHeadDim; + result.store(out_row, lane_id); + params.weights_out[work_id] = weight_val * params.weight_scale * scale; + } +} + +} // anonymous namespace + +// ============================================================================ +// Host-side launchers (PyTorch C++ extension API) +// ============================================================================ + +void dsv4_fused_q_norm_rope( + const at::Tensor& q_input, + at::Tensor& q_output, + const at::Tensor& freqs_cis, + const at::Tensor& positions, + double eps) { + TORCH_CHECK(q_input.is_cuda(), "q_input must be a CUDA tensor"); + TORCH_CHECK(q_output.is_cuda(), "q_output must be a CUDA tensor"); + TORCH_CHECK(q_input.scalar_type() == at::ScalarType::BFloat16, "q_input must be bfloat16"); + TORCH_CHECK(q_output.scalar_type() == at::ScalarType::BFloat16, "q_output must be bfloat16"); + TORCH_CHECK(q_input.dim() == 3, "q_input must be 3D: (B, H, D)"); + TORCH_CHECK(q_output.dim() == 3, "q_output must be 3D: (B, H, D)"); + TORCH_CHECK(positions.scalar_type() == at::ScalarType::Int, "positions must be int32"); + + const int64_t B = q_input.size(0); + const int64_t H = q_input.size(1); + const int64_t D = q_input.size(2); + TORCH_CHECK( + q_output.size(0) == B && q_output.size(1) == H && q_output.size(2) == D, "q_output shape must match q_input"); + TORCH_CHECK(q_input.stride(2) == 1 && q_output.stride(2) == 1, "last dim must be contiguous"); + TORCH_CHECK(q_input.stride(1) == D && q_output.stride(1) == D, "head dim must be contiguous"); + + if (B == 0) return; + + const auto stream = at::cuda::getCurrentCUDAStream(q_input.get_device()); + const auto params = FusedQNormRopeParams{ + .q_input = q_input.data_ptr(), + .q_output = q_output.data_ptr(), + .freqs_cis = freqs_cis.data_ptr(), + .positions = positions.data_ptr(), + .q_input_stride_batch = q_input.stride(0), + .q_output_stride_batch = q_output.stride(0), + .batch_size = static_cast(B), + .num_q_heads = static_cast(H), + .eps = static_cast(eps), + }; + const uint32_t total_works = static_cast(B * H); + const uint32_t num_blocks = CEILDIV(total_works, kFusedQNumWarps); + + // Dispatch on head_dim. DeepSeek V4 uses D=192 with kRopeDim=64. + constexpr int64_t kRopeDim = 64; + switch (D) { + case 128: + fused_q_norm_rope_kernel<128, kRopeDim><<>>(params); + break; + case 192: + fused_q_norm_rope_kernel<192, kRopeDim><<>>(params); + break; + default: + TORCH_CHECK(false, "Unsupported head_dim for dsv4_fused_q_norm_rope: ", D); + } +} + +void dsv4_fused_k_norm_rope_flashmla( + const at::Tensor& kv, + const at::Tensor& kv_weight, + const at::Tensor& freqs_cis, + const at::Tensor& positions, + const at::Tensor& out_loc, + at::Tensor& kvcache, + double eps, + int64_t page_size) { + TORCH_CHECK(kv.is_cuda(), "kv must be a CUDA tensor"); + TORCH_CHECK(kv.scalar_type() == at::ScalarType::BFloat16, "kv must be bfloat16"); + TORCH_CHECK(kv.dim() == 2, "kv must be 2D: (B, D)"); + TORCH_CHECK(positions.scalar_type() == at::ScalarType::Int, "positions must be int32"); + TORCH_CHECK(out_loc.scalar_type() == at::ScalarType::Int, "out_loc must be int32"); + + const int64_t B = kv.size(0); + const int64_t D = kv.size(1); + TORCH_CHECK(D == 512, "kv head_dim must be 512 for FlashMLA"); + TORCH_CHECK(kv_weight.size(0) == D, "kv_weight size must match head_dim"); + + if (B == 0) return; + + const auto stream = at::cuda::getCurrentCUDAStream(kv.get_device()); + const auto params = FusedKNormRopeFlashMLAParams{ + .kv = kv.data_ptr(), + .kv_weight = kv_weight.data_ptr(), + .freqs_cis = freqs_cis.data_ptr(), + .positions = positions.data_ptr(), + .out_loc = out_loc.data_ptr(), + .kvcache = static_cast(kvcache.data_ptr()), + .kv_stride_batch = kv.stride(0), + .batch_size = static_cast(B), + .eps = static_cast(eps), + }; + + constexpr int64_t kHeadDim = 512; + constexpr int64_t kRopeDim = 64; + + // Dispatch on page_size (must be power of 2). + TORCH_CHECK(page_size > 0 && (page_size & (page_size - 1)) == 0, "page_size must be a power of 2"); + +#define LAUNCH_K_KERNEL(PAGE_BITS) \ + fused_k_norm_rope_flashmla_kernel \ + <<(B), kFusedKBlockSize, 0, stream>>>(params) + + switch (page_size) { + case 1: + LAUNCH_K_KERNEL(0); + break; + case 2: + LAUNCH_K_KERNEL(1); + break; + case 4: + LAUNCH_K_KERNEL(2); + break; + case 8: + LAUNCH_K_KERNEL(3); + break; + case 16: + LAUNCH_K_KERNEL(4); + break; + case 32: + LAUNCH_K_KERNEL(5); + break; + case 64: + LAUNCH_K_KERNEL(6); + break; + case 128: + LAUNCH_K_KERNEL(7); + break; + case 256: + LAUNCH_K_KERNEL(8); + break; + default: + TORCH_CHECK(false, "Unsupported page_size: ", page_size); + } +#undef LAUNCH_K_KERNEL +} + +void dsv4_fused_q_indexer_rope_hadamard_quant( + const at::Tensor& q_input, + at::Tensor& q_fp8, + const at::Tensor& weight, + at::Tensor& weights_out, + double weight_scale, + const at::Tensor& freqs_cis, + const at::Tensor& positions) { + TORCH_CHECK(q_input.is_cuda(), "q_input must be a CUDA tensor"); + TORCH_CHECK(q_input.scalar_type() == at::ScalarType::BFloat16, "q_input must be bfloat16"); + TORCH_CHECK(q_input.dim() == 3, "q_input must be 3D: (B, H, D)"); + + const int64_t B = q_input.size(0); + const int64_t H = q_input.size(1); + constexpr int64_t kHeadDim = 128; + TORCH_CHECK(q_input.size(2) == kHeadDim, "q_input head_dim must be 128 for indexer"); + TORCH_CHECK( + q_input.stride(2) == 1 && q_input.stride(1) == kHeadDim, "q_input must be contiguous in (head, elem) dims"); + TORCH_CHECK(q_input.stride(0) == H * kHeadDim, "q_input must be contiguous (B, H, D)"); + TORCH_CHECK(q_fp8.stride(0) == H * kHeadDim, "q_fp8 must be contiguous (B, H, D)"); + TORCH_CHECK(positions.scalar_type() == at::ScalarType::Int, "positions must be int32"); + + if (B == 0) return; + + const auto stream = at::cuda::getCurrentCUDAStream(q_input.get_device()); + const auto params = FusedQIndexerRopeHadamardQuantParams{ + .q_input = q_input.data_ptr(), + .q_fp8 = q_fp8.data_ptr(), + .weight = weight.data_ptr(), + .weights_out = weights_out.data_ptr(), + .weight_scale = static_cast(weight_scale), + .freqs_cis = freqs_cis.data_ptr(), + .positions = positions.data_ptr(), + .batch_size = static_cast(B), + .num_heads = static_cast(H), + }; + const uint32_t total_works = static_cast(B * H); + const uint32_t num_blocks = CEILDIV(total_works, kFusedQNumWarps); + + fused_q_indexer_rope_hadamard_quant_kernel<<>>(params); +} diff --git a/sgl-kernel/include/sgl_kernel_ops.h b/sgl-kernel/include/sgl_kernel_ops.h index 5c0261b64..eb41ff5a3 100644 --- a/sgl-kernel/include/sgl_kernel_ops.h +++ b/sgl-kernel/include/sgl_kernel_ops.h @@ -172,8 +172,45 @@ void fast_topk_transform_ragged_interface( #ifdef USE_ROCM void gelu_quick(at::Tensor& out, const at::Tensor& input); + +void deepseek_v4_topk_transform_512( + const at::Tensor& scores, + const at::Tensor& seq_lens, + const at::Tensor& page_table, + at::Tensor& page_indices, + int64_t page_size, + std::optional raw_indices_opt = std::nullopt); #endif +/* + * From csrc/elementwise (DeepSeek-V4 norm + rope) + */ +void dsv4_fused_q_norm_rope( + const at::Tensor& q_input, + at::Tensor& q_output, + const at::Tensor& freqs_cis, + const at::Tensor& positions, + double eps); + +void dsv4_fused_k_norm_rope_flashmla( + const at::Tensor& kv, + const at::Tensor& kv_weight, + const at::Tensor& freqs_cis, + const at::Tensor& positions, + const at::Tensor& out_loc, + at::Tensor& kvcache, + double eps, + int64_t page_size); + +void dsv4_fused_q_indexer_rope_hadamard_quant( + const at::Tensor& q_input, + at::Tensor& q_fp8, + const at::Tensor& weight, + at::Tensor& weights_out, + double weight_scale, + const at::Tensor& freqs_cis, + const at::Tensor& positions); + /* * From csrc/gemm */ diff --git a/sgl-kernel/python/sgl_kernel/elementwise.py b/sgl-kernel/python/sgl_kernel/elementwise.py index 48fa5584c..aa325b277 100644 --- a/sgl-kernel/python/sgl_kernel/elementwise.py +++ b/sgl-kernel/python/sgl_kernel/elementwise.py @@ -324,6 +324,85 @@ if torch.version.hip is not None: return out +def dsv4_fused_q_norm_rope( + q_input: torch.Tensor, + freqs_cis: torch.Tensor, + positions: torch.Tensor, + eps: float = 1e-6, + q_output: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """DeepSeek-V4 fused Q RMSNorm (no weight) + RoPE. + + Parameters + ---------- + q_input : (B, num_q_heads, head_dim) bfloat16 + freqs_cis: (max_pos, rope_dim) float32, re/im interleaved + positions: (B,) int32 + eps : RMSNorm epsilon + q_output : optional pre-allocated output tensor + """ + if q_output is None: + q_output = torch.empty_like(q_input) + torch.ops.sgl_kernel.dsv4_fused_q_norm_rope.default( + q_input, q_output, freqs_cis, positions, eps + ) + return q_output + + +def dsv4_fused_k_norm_rope_flashmla( + kv: torch.Tensor, + kv_weight: torch.Tensor, + freqs_cis: torch.Tensor, + positions: torch.Tensor, + out_loc: torch.Tensor, + kvcache: torch.Tensor, + eps: float = 1e-6, + page_size: int = 1, +) -> None: + """DeepSeek-V4 fused K RMSNorm + RoPE + FlashMLA FP8 store. + + Parameters + ---------- + kv : (B, 512) bfloat16 + kv_weight: (512,) bfloat16 + freqs_cis: (max_pos, 64) float32 + positions: (B,) int32 + out_loc : (B,) int32 cache slot ids + kvcache : (npages, page_bytes) uint8 + eps : RMSNorm epsilon + page_size: page size (power of 2) + """ + torch.ops.sgl_kernel.dsv4_fused_k_norm_rope_flashmla.default( + kv, kv_weight, freqs_cis, positions, out_loc, kvcache, eps, page_size + ) + + +def dsv4_fused_q_indexer_rope_hadamard_quant( + q_input: torch.Tensor, + q_fp8: torch.Tensor, + weight: torch.Tensor, + weights_out: torch.Tensor, + weight_scale: float, + freqs_cis: torch.Tensor, + positions: torch.Tensor, +) -> None: + """DeepSeek-V4 fused Q indexer: RoPE + Hadamard + FP8 quant. + + Parameters + ---------- + q_input : (B, num_heads, 128) bfloat16 + q_fp8 : (B, num_heads, 128) fp8_e4m3 output + weight : (B, num_heads) bfloat16 + weights_out: (B, num_heads, 1) float32 output + weight_scale: scalar + freqs_cis : (max_pos, 64) float32 + positions : (B,) int32 + """ + torch.ops.sgl_kernel.dsv4_fused_q_indexer_rope_hadamard_quant.default( + q_input, q_fp8, weight, weights_out, weight_scale, freqs_cis, positions + ) + + def rotary_embedding( positions: torch.Tensor, query: torch.Tensor, diff --git a/sgl-kernel/python/sgl_kernel/top_k.py b/sgl-kernel/python/sgl_kernel/top_k.py index 77fd5e5e5..4b842499a 100644 --- a/sgl-kernel/python/sgl_kernel/top_k.py +++ b/sgl-kernel/python/sgl_kernel/top_k.py @@ -80,6 +80,38 @@ def fast_topk_transform_fused( return dst_page_table +def deepseek_v4_topk_transform_512( + scores: torch.Tensor, + seq_lens: torch.Tensor, + page_table: torch.Tensor, + page_indices: torch.Tensor, + page_size: int, + raw_indices: Optional[torch.Tensor] = None, +) -> None: + """ + Performs the DeepSeek-V4 indexer top-k selection and writes the paged + physical slot indices into ``page_indices``. Supports topk up to 1024. + Optionally also writes the row-relative raw token positions into + ``raw_indices`` for hisparse capture. + + Args: + scores: float32 ``[B, max_seq_len]`` indexer logits, contiguous on dim 1. + seq_lens: int32 ``[B]``, true KV length per batch row. + page_table: int32 ``[B, num_pages]``, logical->physical page table, + contiguous on dim 1. + page_indices: int32 ``[B, topk]``, output buffer, contiguous. Filled + with paged physical slots; -1 for padding entries. + page_size: power-of-2 page size. + raw_indices: optional int32 ``[B, topk]``, contiguous. If provided, + filled with raw token positions within each row. + """ + if raw_indices is not None: + assert raw_indices.dim() == 2 + torch.ops.sgl_kernel.deepseek_v4_topk_transform_512( + scores, seq_lens, page_table, page_indices, page_size, raw_indices + ) + + def fast_topk_transform_ragged_fused( score: torch.Tensor, lengths: torch.Tensor, diff --git a/sgl-kernel/setup_rocm.py b/sgl-kernel/setup_rocm.py index 8dddf027e..9762787aa 100644 --- a/sgl-kernel/setup_rocm.py +++ b/sgl-kernel/setup_rocm.py @@ -46,6 +46,8 @@ sources = [ "csrc/allreduce/quick_all_reduce.cu", "csrc/common_extension_rocm.cc", "csrc/elementwise/activation.cu", + "csrc/elementwise/deepseek_v4_topk.cu", + "csrc/elementwise/dsv4_norm_rope.cu", "csrc/elementwise/topk.cu", "csrc/grammar/apply_token_bitmask_inplace_cuda.cu", "csrc/moe/moe_align_kernel.cu",