perf: add fused norm + RoPE + uniform fp8 store for TRT-LLM DSv4 sparse attention

python/sglang/kernels/jit/csrc/deepseek_v4/fused_norm_rope_v2.cuh

@@ -380,10 +380,17 @@ INDEXER_KERNEL void fused_norm_rope_indexer_fp4(const __grid_constant__ FusedNor
 // Each thread loads kVecSize=2 BF16, so 256 threads cover the full 512 elems.
 // Cache layout: 584 bytes/token = 448 fp8 nope + 64 (=32 bf16x2) rope + 8 scale.
 // ----------------------------------------------------------------------------
-template <typename DType, ForwardMode kMode, int32_t kPageBits, bool kUsePDL, bool kBf16Store = false>
+template <
+    typename DType,
+    ForwardMode kMode,
+    int32_t kPageBits,
+    bool kUsePDL,
+    bool kBf16Store = false,
+    bool kUniformFp8Store = false>
 FLASHMLA_KERNEL void fused_norm_rope_flashmla(const __grid_constant__ FusedNormRopeStoreParams params) {
   using namespace device;
   using enum ForwardMode;
+  static_assert(!(kBf16Store && kUniformFp8Store));
 
   constexpr int64_t kHeadDim = 512;
   constexpr int64_t kRopeDim = 64;
@@ -393,8 +400,11 @@ FLASHMLA_KERNEL void fused_norm_rope_flashmla(const __grid_constant__ FusedNormR
   constexpr uint32_t kRopeWarp = kNumWarps - 1;
   // kBf16Store: write the whole head_dim as plain BF16 (no fp8 / no scale) into a
   // [num_slots, head_dim] bf16 cache (page_size==1) at row out_loc
-  constexpr int64_t kPageBytes =
-      kBf16Store ? ((kHeadDim * 2ll) << kPageBits) : host::div_ceil(584ll << kPageBits, 576) * 576;
+  // kUniformFp8Store: write the whole head_dim (rope tail included) as plain
+  // e4m3 at per-tensor scale 1.0 into the uniform 512-byte-per-token pool.
+  constexpr int64_t kPageBytes = kBf16Store         ? ((kHeadDim * 2ll) << kPageBits)
+                                 : kUniformFp8Store ? (kHeadDim << kPageBits)
+                                                    : host::div_ceil(584ll << kPageBits, 576) * 576;
   static_assert(kHeadDim == kBlockSize * kVecSize);
   static_assert(kRopeDim == kWarpThreads * kVecSize);
   static_assert(kHeadDim - kRopeDim == kRopeWarp * kWarpThreads * kVecSize);
@@ -465,12 +475,13 @@ FLASHMLA_KERNEL void fused_norm_rope_flashmla(const __grid_constant__ FusedNormR
   const int64_t page = out_loc >> kPageBits;
   const int64_t offset = out_loc & ((1 << kPageBits) - 1);
   const auto page_ptr = params.kvcache + page * kPageBytes;
-  const auto value_ptr = page_ptr + offset * (kBf16Store ? (kHeadDim * 2) : 576);
+  const auto value_ptr = page_ptr + offset * (kBf16Store ? (kHeadDim * 2) : kUniformFp8Store ? kHeadDim : 576);
 
   PDLTriggerSecondary<kUsePDL>();
 
-  // part 2: rope on the rope warp (BF16 store), or per-warp FP8 quant + store.
-  if constexpr (kBf16Store) {
+  // part 2: rope on the rope warp (BF16/uniform store), or per-warp FP8
+  // quant + store (packed layout).
+  if constexpr (kBf16Store || kUniformFp8Store) {
     Float2 d = data;
     if (warp_id == kRopeWarp) {
       const auto x_real = data[0];
@@ -480,7 +491,15 @@ FLASHMLA_KERNEL void fused_norm_rope_flashmla(const __grid_constant__ FusedNormR
       d[0] = x_real * freq_real - x_imag * freq_imag;
       d[1] = x_real * freq_imag + x_imag * freq_real;
     }
-    reinterpret_cast<bf16x2_t*>(value_ptr)[tx] = cast<bf16x2_t>(fp32x2_t{d[0], d[1]});
+    if constexpr (kUniformFp8Store) {
+      // BF16 round-trip to match the unfused path (Triton norm+rope emits
+      // bf16, then the pool store casts bf16 -> e4m3 at scale 1.0).
+      const auto x = cast<float>(cast<bf16_t>(d[0]));
+      const auto y = cast<float>(cast<bf16_t>(d[1]));
+      reinterpret_cast<fp8x2_e4m3_t*>(value_ptr)[tx] = pack_fp8(x, y);
+    } else {
+      reinterpret_cast<bf16x2_t*>(value_ptr)[tx] = cast<bf16x2_t>(fp32x2_t{d[0], d[1]});
+    }
   } else if (warp_id == kRopeWarp) {
     // Each rope-warp lane owns exactly one (real, imag) pair within the rope
     // tail. Apply rotation, downcast to BF16, write to the slot's rope region.
@@ -518,15 +537,21 @@ template <
     uint32_t kPageSize,
     bool kUsePDL,
     int32_t kPreshuffleSize = 0,
-    bool kBf16Store = false>
+    bool kBf16Store = false,
+    bool kUniformFp8Store = false>
 struct FusedNormRopeKernel {
   static constexpr int32_t kLogPageSize = std::countr_zero(kPageSize);
   static constexpr bool kIsIndexer = (kHeadDim == 128);
   static_assert(!(kIsIndexer && kBf16Store), "bf16 store only for flashmla head_dim=512");
+  static_assert(!(kIsIndexer && kUniformFp8Store), "uniform fp8 store only for flashmla head_dim=512");
+  static_assert(!(kBf16Store && kUniformFp8Store));
   static constexpr int64_t kIndexerBytes = 132 * kPageSize;
   static constexpr int64_t kFlashMLABytes = host::div_ceil(584 * kPageSize, 576) * 576;
   static constexpr int64_t kBf16Bytes = kHeadDim * 2 * kPageSize;  // plain bf16 cache
-  static constexpr int64_t kPageBytes = kBf16Store ? kBf16Bytes : (kIsIndexer ? kIndexerBytes : kFlashMLABytes);
+  static constexpr int64_t kUniformBytes = kHeadDim * kPageSize;   // uniform e4m3, 512 B/token
+  static constexpr int64_t kPageBytes = kBf16Store         ? kBf16Bytes
+                                        : kUniformFp8Store ? kUniformBytes
+                                                           : (kIsIndexer ? kIndexerBytes : kFlashMLABytes);
 
   /// TODO: Let's fix the config for now.
   static_assert(kRopeDim == 64 && (kHeadDim == 128 || kHeadDim == 512));
@@ -537,7 +562,7 @@ struct FusedNormRopeKernel {
     if constexpr (kIsIndexer) {
       return fused_norm_rope_indexer<DType, kMode, kLogPageSize, kUsePDL, kPreshuffleSize>;
     } else {
-      return fused_norm_rope_flashmla<DType, kMode, kLogPageSize, kUsePDL, kBf16Store>;
+      return fused_norm_rope_flashmla<DType, kMode, kLogPageSize, kUsePDL, kBf16Store, kUniformFp8Store>;
     }
   }
 

python/sglang/kernels/ops/attention/dsv4/compress.py

@@ -49,6 +49,7 @@ def _jit_compress_norm_rope_module(
     rope_dim: int,
     page_size: int,
     bf16_store: bool = False,
+    uniform_fp8_store: bool = False,
 ) -> Module:
     args = make_cpp_args(
         dtype,
@@ -58,6 +59,7 @@ def _jit_compress_norm_rope_module(
         is_arch_support_pdl(),
         INDEXER_K_CACHE_PRESHUFFLE_TILE if aiter_can_use_preshuffle_paged_mqa() else 0,
         bf16_store,
+        uniform_fp8_store,
     )
     cuda_wrappers = [("forward", f"FusedNormRopeKernel<{args}>::forward")]
     if head_dim == 128:
@@ -425,9 +427,14 @@ def compress_norm_rope_store(
     page_size: int,
     use_fp4: bool = False,
     bf16_store: bool = False,
+    uniform_fp8_store: bool = False,
 ) -> None:
     if use_fp4:
         assert kv.shape[-1] == 128
+    if uniform_fp8_store:
+        # Uniform 512-byte-per-token e4m3 pool (trtllm backend): plain cast at
+        # per-tensor scale 1.0, rope tail included; no packed scales.
+        assert kv.shape[-1] == 512 and not use_fp4 and not bf16_store
     freq_cis = torch.view_as_real(freq_cis).flatten(-2)
     if _is_xpu:
         compress_norm_rope_store_xpu(
@@ -445,7 +452,12 @@ def compress_norm_rope_store(
         )
     else:
         module = _jit_compress_norm_rope_module(
-            kv.dtype, kv.shape[-1], freq_cis.shape[-1], page_size, bf16_store
+            kv.dtype,
+            kv.shape[-1],
+            freq_cis.shape[-1],
+            page_size,
+            bf16_store,
+            uniform_fp8_store,
         )
         fn = module.forward_fp4 if use_fp4 else module.forward
         fn(

python/sglang/srt/arg_groups/deepseek_v4_hook.py

@@ -135,6 +135,36 @@ def apply_deepseek_v4_defaults(server_args: ServerArgs, model_arch: str) -> None
 
     run_post_process_pass(server_args, _deepseek_v4_kv_cache_dtype)
 
+    if server_args.dsv4_attn_backend == "trtllm":
+        from sglang.srt.utils.common import is_sm100_supported
+
+        assert (
+            server_args.device == "cuda" and is_sm100_supported()
+        ), "--dsv4-attn-backend trtllm requires an SM100/SM103 (Blackwell) GPU."
+        # "auto" is declared-but-unmaterialized here; the resolution pipeline
+        # (_deepseek_v4_kv_cache_dtype above) turns it into fp8_e4m3 on cuda.
+        assert server_args.kv_cache_dtype in ("auto", "fp8_e4m3"), (
+            "--dsv4-attn-backend trtllm requires kv_cache_dtype=fp8_e4m3, "
+            f"got {server_args.kv_cache_dtype}."
+        )
+        assert (
+            not server_args.enable_hisparse
+        ), "--dsv4-attn-backend trtllm does not support enable_hisparse."
+        assert not (
+            server_args.attn_cp_size > 1
+            or server_args.dcp_size > 1
+            or server_args.enable_prefill_cp
+            or server_args.enable_prefill_context_parallel
+            or server_args.enable_dsa_prefill_context_parallel
+        ), (
+            "--dsv4-attn-backend trtllm does not support context parallelism "
+            "(prefill CP, attention CP, or decode CP)."
+        )
+        logger.info(
+            "DeepSeek V4 attention: trtllm backend enabled "
+            "(uniform-FP8 KV pool, decode + sparse prefill)."
+        )
+
     if server_args.max_running_requests is None:
         server_args.max_running_requests = 256
         logger.warning(
@@ -151,6 +181,27 @@ def apply_deepseek_v4_defaults(server_args: ServerArgs, model_arch: str) -> None
                 server_args.speculative_eagle_topk == 1
             ), f"Only EAGLE speculative algorithm with topk == 1 is supported for {model_arch}"
 
+    # FIXME(follow-up): remove once the overlap+speculative corruption is
+    # root-caused (tracked in the PR #30805 follow-up list).
+    # TEMPORARY containment: trtllm + speculative decoding under the overlap
+    # scheduler intermittently corrupts an int32 table consumed by the
+    # trtllm-gen sparse kernel (illegal memory access in
+    # fmhaSm100fKernel...VarSeq during concurrent GSM8K-style bursts). This
+    # reproduces with both TP-only and DP-attention recipes; disabling overlap
+    # prevents the corruption while the root cause is investigated.
+    if (
+        server_args.dsv4_attn_backend == "trtllm"
+        and server_args.speculative_algorithm is not None
+        and not server_args.disable_overlap_schedule
+    ):
+        logger.warning(
+            "Disabling the overlap scheduler for the trtllm DeepSeek-V4 "
+            "backend with speculative decoding (temporary "
+            "containment for an intermittent trtllm-gen kernel memory fault; "
+            "see the dsv4 trtllm PR discussion)."
+        )
+        server_args.disable_overlap_schedule = True
+
 
 def validate_deepseek_v4_cp(server_args: ServerArgs) -> None:
     """Validate DeepSeek V4 context-parallel configuration."""

python/sglang/srt/layers/attention/attention_registry.py

@@ -166,12 +166,15 @@ def create_dsv4_backend(runner):
         )
         return DeepseekV4HipRadixBackend(runner)
     else:
-        from sglang.srt.layers.attention.deepseek_v4_backend import (
-            DeepseekV4AttnBackend,
+        from sglang.srt.layers.attention.deepseek_v4_trtllm_backend import (
+            create_deepseek_v4_attn_backend,
         )
 
-        logger.info("Using DeepseekV4AttnBackend for dsv4 attention backend (CUDA).")
-        return DeepseekV4AttnBackend(runner)
+        backend = create_deepseek_v4_attn_backend(runner)
+        logger.info(
+            f"Using {type(backend).__name__} for dsv4 attention backend (CUDA)."
+        )
+        return backend
 
 
 @register_attention_backend("triton")