minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[AMD] Dsv4/pr1 fix run time issue (#25898)

Browse Source 
Co-authored-by: wunhuang <wunhuang@amd.com>
Co-authored-by: Thomas Wang <1am9trash@gmail.com>
Co-authored-by: Xinyi Song <86638975+RolaoDenthu@users.noreply.github.com>
Co-authored-by: HaiShaw <hixiao@gmail.com>
Co-authored-by: amd-danli103 <danli103@amd.com>
Co-authored-by: Lin, Soga <soga.lin@amd.com>
Co-authored-by: Raiden-Makoto <Raiden-Makoto@users.noreply.github.com>
Co-authored-by: Hubert Lu <55214931+hubertlu-tw@users.noreply.github.com>
Co-authored-by: yichiche@amd.com <jacky.cheng>
Co-authored-by: yctseng0211 <yctseng@amd.com>
Co-authored-by: Bingxu Chen <bingxche@amd.com>
This commit is contained in:
kk
2026-05-23 16:04:14 -07:00
committed by GitHub
co-authored by wunhuang Thomas Wang Xinyi Song HaiShaw amd-danli103 Lin, Soga Raiden-Makoto Hubert Lu yichiche@amd.com yctseng0211 Bingxu Chen
parent 982f67d9a6
commit af8f66940e
32 changed files with 2523 additions and 129 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/csrc/deepseek_v4/c128_v2.cuh
+2 -2
View File
@@ -348,7 +348,7 @@ struct FlashCompress128Kernel {
auto N = SymbolicSize{"batch_size"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLCUDA>();
device_.set_options<kDLGPU>();
TensorMatcher({-1, 128, Trait::kElementSize}) // kv score
.with_dtype<InFloat>()
@@ -395,7 +395,7 @@ struct FlashCompress128Kernel {
auto C = SymbolicSize{"num_c_plans"};
auto W = SymbolicSize{"num_w_plans"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLCUDA>();
device_.set_options<kDLGPU>();
TensorMatcher({-1, 128, Trait::kElementSize}) // kv score
.with_dtype<InFloat>()
python/sglang/jit_kernel/csrc/deepseek_v4/c4_v2.cuh
+2 -2
View File
@@ -309,7 +309,7 @@ struct FlashCompress4Kernel {
auto N = SymbolicSize{"batch_size"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLCUDA>();
device_.set_options<kDLGPU>();
TensorMatcher({-1, 4, Trait::kElementSize}) // kv score
.with_dtype<InFloat>()
@@ -356,7 +356,7 @@ struct FlashCompress4Kernel {
auto C = SymbolicSize{"num_c_plans"};
auto W = SymbolicSize{"num_w_plans"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLCUDA>();
device_.set_options<kDLGPU>();
TensorMatcher({-1, 4, Trait::kElementSize}) // kv score
.with_dtype<InFloat>()
python/sglang/jit_kernel/csrc/deepseek_v4/c_plan.cuh
+21 -9
View File
@@ -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<kDLCPU, kDLCUDA>();
device_.set_options<kDLCUDA>();
cpu_or_gpu.set_options<kDLCPU, kDLGPU>();
device_.set_options<kDLGPU>();
TensorMatcher({B}) //
.with_dtype<RID_T>()
@@ -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<kDLCUDA>();
device_.set_options<kDLGPU>();
TensorMatcher({B}) //
.with_dtype<RID_T>()
@@ -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<kDLCUDA>();
device_.set_options<kDLGPU>();
TensorMatcher({B}) //
.with_dtype<RID_T>()
@@ -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<kDLCUDA>();
device_.set_options<kDLGPU>();
TensorMatcher({B}) //
.with_dtype<RID_T>()
python/sglang/jit_kernel/csrc/deepseek_v4/fused_norm_rope_v2.cuh
+10 -4
View File
@@ -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<bf16x2_t*>(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<float>(cast<bf16_t>(data[0]));
const auto y = cast<float>(cast<bf16_t>(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<kDLCUDA>();
device_.set_options<kDLGPU>();
TensorMatcher({N, kHeadDim}) // input
.with_dtype<DType>()
python/sglang/jit_kernel/dsv4/attn.py
+13 -7
View File
@@ -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
python/sglang/jit_kernel/dsv4/elementwise.py
+21 -4
View File
@@ -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
python/sglang/jit_kernel/dsv4/gemm.py
+3 -3
View File
@@ -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)
python/sglang/jit_kernel/dsv4/moe.py
+32 -19
View File
@@ -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(
python/sglang/jit_kernel/dsv4/topk.py
+10 -4
View File
@@ -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
python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/fp8_utils.cuh
+71 -2
View File
@@ -5,7 +5,9 @@
#include <sgl_kernel/utils.cuh>
#include <cstdint>
#ifndef USE_ROCM
#include <cuda_fp8.h>
#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<uint8_t>((f32 >> 31) << 7);
int32_t exp32 = static_cast<int32_t>((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<uint8_t>(subnorm_mant & 0x07);
exp8 = 0;
if (subnorm_mant > 7) {
exp8 = 1;
mant3 = 0;
}
} else {
exp8 = exp32 + kBias;
mant3 = static_cast<uint8_t>(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<uint8_t>(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<uint16_t>(x8) | (static_cast<uint16_t>(y8) << 8);
}
#endif
} // namespace deepseek_v4::fp8
python/sglang/jit_kernel/include/sgl_kernel/runtime.cuh
+31
View File
@@ -10,7 +10,38 @@
#include <cstddef>
#include <cstdint>
#ifndef USE_ROCM
#include <cuda_runtime.h>
#else
#include <hip/hip_runtime.h>
#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 {
python/sglang/jit_kernel/include/sgl_kernel/tensor.h
+11
View File
@@ -33,6 +33,8 @@
#ifdef __CUDACC__
#include <sgl_kernel/utils.cuh>
#elif defined(__HIPCC__)
#include <sgl_kernel/utils.cuh>
#endif
namespace host {
@@ -79,6 +81,15 @@ template <>
struct _dtype_trait<fp8_e4m3_t> {
inline static constexpr DLDataType value = {.code = DLDataTypeCode::kDLFloat8_e4m3fn, .bits = 8, .lanes = 1};
};
#elif defined(__HIPCC__)
template <>
struct _dtype_trait<fp16_t> {
inline static constexpr DLDataType value = {.code = DLDataTypeCode::kDLFloat, .bits = 16, .lanes = 1};
};
template <>
struct _dtype_trait<bf16_t> {
inline static constexpr DLDataType value = {.code = DLDataTypeCode::kDLBfloat, .bits = 16, .lanes = 1};
};
#endif
template <DLDeviceType Code>
python/sglang/jit_kernel/include/sgl_kernel/utils.cuh
+14
View File
@@ -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.
python/sglang/jit_kernel/include/sgl_kernel/warp.cuh
+20 -23
View File
@@ -1,5 +1,5 @@
/// \file warp.cuh
/// \brief Warp-level reduction primitives using `__shfl_xor_sync`.
/// \brief Warp-level reduction primitives.
#pragma once
#include <sgl_kernel/math.cuh>
@@ -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 <uint32_t kNumThreads = kWarpThreads, typename T>
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 <uint32_t kNumThreads = kWarpThreads, typename T>
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;
}
python/sglang/jit_kernel/triton/hash_topk.py
+99
View File
@@ -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
python/sglang/jit_kernel/triton_store_cache.py
+237
View File
@@ -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)
python/sglang/jit_kernel/utils.py
+12 -1
View File
@@ -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,
python/sglang/srt/environ.py
+2
View File
@@ -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)
# ====================================================================
python/sglang/srt/layers/attention/dsv4/compress_hip.py
+91 -20
View File
@@ -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
)
python/sglang/srt/layers/fused_qk_norm_rope_store.py
+380
View File
@@ -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
python/sglang/srt/layers/layernorm.py
+6
View File
@@ -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.
python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py
+1
View File
@@ -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:
python/sglang/srt/models/deepseek_v2.py
+39 -1
View File
@@ -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))
python/sglang/srt/models/deepseek_v4.py
+153 -28
View File
@@ -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()