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

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>

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<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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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()

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"

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
    */

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 <ATen/core/TensorBase.h>
+#include <ATen/core/TensorBody.h>
+#include <c10/cuda/CUDAStream.h>
+#include <c10/macros/Macros.h>
+#include <c10/util/Exception.h>
+#include <cuda.h>
+#include <cuda_fp16.h>
+
+#include <cstddef>
+#include <cstdint>
+#include <optional>
+
+namespace {
+
+constexpr uint32_t kMaxTopK = 1024;
+constexpr uint32_t kBlockSize = 512;
+
+#ifdef SGL_TOPK_DYNAMIC_SMEM_BYTES
+constexpr size_t kSMEM = static_cast<size_t>(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<uint16_t>(~bits) : static_cast<uint16_t>(bits | 0x8000);
+  return static_cast<uint8_t>(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<int32_t>(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<uint32_t>(length)) {
+      page_indices_out[i] = page_to_slot(page_table, i, page_bits);
+      if (raw_indices_out != nullptr) {
+        raw_indices_out[i] = static_cast<int32_t>(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<int32_t>(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<int32_t>(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<int32_t>(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<int32_t>(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<int32_t>(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<int32_t>(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<int32_t>(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<uint32_t>(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<uint32_t>(raw), params.page_bits);
+    if (raw_indices_ptr != nullptr) {
+      raw_indices_ptr[i] = raw;
+    }
+  }
+}
+
+template <auto* f, size_t kMaxDynamicSMEM>
+void setup_kernel_smem_once() {
+  [[maybe_unused]]
+  static const auto result = [] {
+#ifdef USE_ROCM
+    return ::cudaFuncSetAttribute(
+        reinterpret_cast<const void*>(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<at::Tensor> 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<int64_t>(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<uint32_t>(__builtin_ctzll(static_cast<unsigned long long>(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<int32_t>();
+  }
+
+  const TopKParams params{
+      .scores = scores.data_ptr<float>(),
+      .seq_lens = seq_lens.data_ptr<int32_t>(),
+      .page_table = page_table.data_ptr<int32_t>(),
+      .page_indices = page_indices.data_ptr<int32_t>(),
+      .raw_indices = raw_ptr,
+      .score_stride = scores.stride(0),
+      .page_table_stride = page_table.stride(0),
+      .page_bits = page_bits,
+      .topk = static_cast<uint32_t>(topk),
+      .output_stride = topk,
+  };
+
+  const auto stream = at::cuda::getCurrentCUDAStream().stream();
+  const dim3 grid(static_cast<uint32_t>(B));
+  const dim3 block(kBlockSize);
+
+  setup_kernel_smem_once<deepseek_v4_topk_transform_kernel, kSMEM>();
+  deepseek_v4_topk_transform_kernel<<<grid, block, kSMEM, stream>>>(params);
+
+  const auto err = cudaGetLastError();
+  TORCH_CHECK(err == cudaSuccess, "deepseek_v4_topk_transform kernel launch failed: ", ::cudaGetErrorString(err));
+}

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 <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda_fp8.h>
+#include <cuda_runtime.h>
+#else
+#include <hip/hip_bf16.h>
+#include <hip/hip_fp16.h>
+#include <hip/hip_runtime.h>
+#endif
+
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/all.h>
+
+#include <cstdint>
+
+#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 <uint32_t kNumThreads = kWarpSize>
+__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 <typename T, int N>
+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<const AlignedVec*>(ptr)[offset];
+  }
+  __device__ __forceinline__ void store(void* ptr, int64_t offset = 0) const {
+    reinterpret_cast<AlignedVec*>(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<int32_t>((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<uint32_t>((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<uint8_t>((f32 >> 24) & 0x80u);
+  f32 &= 0x7FFFFFFFu;
+
+  int32_t exp32 = static_cast<int32_t>((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<uint8_t>(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<uint8_t>(exp8) << 3) | static_cast<uint8_t>(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<uint16_t>(x8) | (static_cast<uint16_t>(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 <int64_t kHeadDim, int64_t kRopeDim>
+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 <int64_t kHeadDim, int64_t kRopeDim>
+__global__ __launch_bounds__(kFusedQBlockSize, 16) void fused_q_norm_rope_kernel(
+    const __grid_constant__ FusedQNormRopeParams params) {
+  using Traits = QKernelTraits<kHeadDim, kRopeDim>;
+  constexpr int64_t kVecSize = Traits::kVecSize;
+  constexpr int64_t kLocalSize = Traits::kLocalSize;
+  constexpr uint32_t kRopeSize = Traits::kRopeSize;
+
+  using Storage = AlignedVec<bf16_t, kVecSize>;
+  using Float2 = AlignedVec<float, 2>;
+
+  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<const bf16_t*>(params.q_input) + batch_id * params.q_input_stride_batch + head_id * kHeadDim;
+  const auto output_ptr =
+      static_cast<bf16_t*>(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<float>(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<bf16x2_t*>(&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<bf16x2_t*>(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 <int64_t kHeadDim, int64_t kRopeDim, int32_t kPageBits>
+__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<bf16_t, kVecSize>;
+
+  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<const bf16_t*>(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<float, kVecSize> 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<kFusedKNumWarps>(partial_sums[lane_id % kFusedKNumWarps]);
+    const float norm_factor = rsqrtf(sum_of_squares / static_cast<float>(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<bf16x2_t*>(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<fp8x2_e4m3_t*>(value_ptr)[tx] = result;
+    if (lane_id == 0) static_cast<uint8_t*>(scale_ptr)[warp_id] = static_cast<uint8_t>(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<bf16_t, kVecSize>;
+  using Float4 = AlignedVec<float, kVecSize>;
+  using OutStorage = AlignedVec<fp8x2_e4m3_t, 2>;
+
+  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<const bf16_t*>(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<const bf16_t*>(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<float>(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<uint8_t*>(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<float>(),
+      .positions = positions.data_ptr<int32_t>(),
+      .q_input_stride_batch = q_input.stride(0),
+      .q_output_stride_batch = q_output.stride(0),
+      .batch_size = static_cast<uint32_t>(B),
+      .num_q_heads = static_cast<uint32_t>(H),
+      .eps = static_cast<float>(eps),
+  };
+  const uint32_t total_works = static_cast<uint32_t>(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><<<num_blocks, kFusedQBlockSize, 0, stream>>>(params);
+      break;
+    case 192:
+      fused_q_norm_rope_kernel<192, kRopeDim><<<num_blocks, kFusedQBlockSize, 0, stream>>>(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<float>(),
+      .positions = positions.data_ptr<int32_t>(),
+      .out_loc = out_loc.data_ptr<int32_t>(),
+      .kvcache = static_cast<uint8_t*>(kvcache.data_ptr()),
+      .kv_stride_batch = kv.stride(0),
+      .batch_size = static_cast<uint32_t>(B),
+      .eps = static_cast<float>(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<kHeadDim, kRopeDim, PAGE_BITS> \
+      <<<static_cast<uint32_t>(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<float>(),
+      .weight_scale = static_cast<float>(weight_scale),
+      .freqs_cis = freqs_cis.data_ptr<float>(),
+      .positions = positions.data_ptr<int32_t>(),
+      .batch_size = static_cast<uint32_t>(B),
+      .num_heads = static_cast<uint32_t>(H),
+  };
+  const uint32_t total_works = static_cast<uint32_t>(B * H);
+  const uint32_t num_blocks = CEILDIV(total_works, kFusedQNumWarps);
+
+  fused_q_indexer_rope_hadamard_quant_kernel<<<num_blocks, kFusedQBlockSize, 0, stream>>>(params);
+}

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<at::Tensor> 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
  */

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,

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,

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",