A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 1, 22:30 UTC
[DSv4] Adding TRTLLM gen attention kernel (#43827)
Yongye Zhu committed
Jun 4, 2026 at 10:35 UTC
b5235fca2eb7cf5791c94eef079d1387fb498504
20 files changed
+2971
-398
csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
+444
@@ -102,6 +102,35 @@ constexpr float NUM_TOKEN_CUTOFF = 1024;
102
constexpr int kNumLanes = 32;
103
constexpr int kElemsPerLane = kHeadDim / kNumLanes; // 16
104
105
+// Pack this lane's 16 fp32 elements into per-tensor E4M3 FP8 (one uint4 = 16
106
+// B), scaling by `scale` (a reciprocal scale) and saturating to ±448. Used by
107
+// the FlashInfer full-cache path for both the Q and KV stores.
108
+__device__ __forceinline__ uint4 packFp8E4M3x16(float const* values,
109
+ float const scale) {
110
+#ifndef USE_ROCM
111
+ uint4 out;
112
+ auto* out2 = reinterpret_cast<__nv_fp8x2_storage_t*>(&out);
113
+ #pragma unroll
114
+ for (int i = 0; i < kElemsPerLane / 2; i++) {
115
+ float2 scaled =
116
+ make_float2(values[2 * i] * scale, values[2 * i + 1] * scale);
117
+ scaled.x = fminf(fmaxf(scaled.x, -kFp8Max), kFp8Max);
118
+ scaled.y = fminf(fmaxf(scaled.y, -kFp8Max), kFp8Max);
119
+ out2[i] = __nv_cvt_float2_to_fp8x2(scaled, __NV_SATFINITE, __NV_E4M3);
120
+ }
121
+ return out;
122
+#else
123
+ uint8_t out_bytes[kElemsPerLane];
124
+ #pragma unroll
125
+ for (int i = 0; i < kElemsPerLane; i++) {
126
+ float scaled = values[i] * scale;
127
+ scaled = fminf(fmaxf(scaled, -kFp8Max), kFp8Max);
128
+ out_bytes[i] = rocm_cvt_float_to_fp8_e4m3(scaled);
129
+ }
130
+ return *reinterpret_cast<uint4 const*>(out_bytes);
131
+#endif
132
+}
133
+
134
// ────────────────────────────────────────────────────────────────────────────
135
// Small inline helpers
136
// ────────────────────────────────────────────────────────────────────────────
@@ -649,6 +678,257 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
678
#undef DISPATCH
679
}
680
681
+// ────────────────────────────────────────────────────────────────────────────
682
+// FlashInfer full-cache kernel
683
+// ────────────────────────────────────────────────────────────────────────────
684
+//
685
+// Sibling to the FlashMLA kernel above, used by the FlashInfer V4 sparse-MLA
686
+// backend. Differences from the legacy path:
687
+// * No Q head padding — output Q layout matches the input num_heads_q.
688
+// * KV is written as a *contiguous* 512-wide row per token (token-strided),
689
+// not the legacy UE8M0 paged layout with a separate scale tail.
690
+// * Q/KV are stored either as bf16 or as per-tensor E4M3 FP8 (one global
691
+// scale), selected by the STORE_Q_FP8 / STORE_KV_FP8 template flags.
692
+//
693
+// Grid: 1D, gridDim.x = ceil(num_tokens_full * (num_heads_q + 1) / warps).
694
+// Each warp handles one (token, slot): slot < num_heads_q → Q, slot ==
695
+// num_heads_q → KV.
696
+template <typename scalar_t_in, bool STORE_Q_FP8, bool STORE_KV_FP8>
697
+__global__ void fusedDeepseekV4FullCacheKernel(
698
+ scalar_t_in* __restrict__ q_inout, // [N, H, 512], in place (bf16)
699
+ uint8_t* __restrict__ q_fp8_out, // [N, H, 512] fp8, optional
700
+ int64_t const q_fp8_stride0, // elements (fp8 == bytes)
701
+ int64_t const q_fp8_stride1, // elements (fp8 == bytes)
702
+ scalar_t_in const* __restrict__ kv_in, // [N, 512] bf16
703
+ uint8_t* __restrict__ k_cache, // contiguous bf16 or fp8 cache
704
+ int64_t const* __restrict__ slot_mapping, // [num_tokens_insert] i64
705
+ int64_t const* __restrict__ position_ids, // [N] i64
706
+ float const* __restrict__ cos_sin_cache, // [max_pos, 64] fp32
707
+ float const* __restrict__ fp8_scale_ptr, // scalar, KV fp8 only
708
+ float const* __restrict__ q_fp8_scale_inv, // scalar, Q fp8 only
709
+ float const eps,
710
+ int const num_tokens_full, // = q.size(0) = kv.size(0)
711
+ int const num_tokens_insert, // = slot_mapping.size(0)
712
+ int const num_heads_q, // H (no padding)
713
+ int const cache_block_size, // tokens per cache block
714
+ int64_t const kv_block_stride, // bytes per cache block
715
+ int64_t const kv_token_stride) { // bytes per cache token
716
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
717
+ if constexpr (std::is_same_v<scalar_t_in, c10::BFloat16>) {
718
+ return;
719
+ } else {
720
+#endif
721
+ using Converter = vllm::_typeConvert<scalar_t_in>;
722
+ int const warpsPerBlock = blockDim.x / 32;
723
+ int const warpId = threadIdx.x / 32;
724
+ int const laneId = threadIdx.x % 32;
725
+ int const globalWarpIdx = blockIdx.x * warpsPerBlock + warpId;
726
+
727
+ int const slotsPerToken = num_heads_q + 1;
728
+ int const tokenIdx = globalWarpIdx / slotsPerToken;
729
+ int const slotIdx = globalWarpIdx % slotsPerToken;
730
+ if (tokenIdx >= num_tokens_full) return;
731
+ bool const isKV = (slotIdx == num_heads_q);
732
+ // KV branch: skip DP-padded tokens (no slot reserved for them).
733
+ if (isKV && tokenIdx >= num_tokens_insert) return;
734
+
735
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
736
+ cudaGridDependencySynchronize();
737
+#endif
738
+
739
+ int const dim_base = laneId * kElemsPerLane; // in [0, 512) step 16
740
+ scalar_t_in const* src_ptr;
741
+ if (isKV) {
742
+ src_ptr = kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
743
+ } else {
744
+ src_ptr = q_inout +
745
+ (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) *
746
+ kHeadDim +
747
+ dim_base;
748
+ }
749
+ uint4 const v0 = *reinterpret_cast<uint4 const*>(src_ptr);
750
+ uint4 const v1 = *reinterpret_cast<uint4 const*>(src_ptr + 8);
751
+
752
+ // ── Decode bf16 → 16 fp32 registers ───────────────────────────────────
753
+ float elements[kElemsPerLane];
754
+ {
755
+ auto const* p0 =
756
+ reinterpret_cast<typename Converter::packed_hip_type const*>(&v0);
757
+ auto const* p1 =
758
+ reinterpret_cast<typename Converter::packed_hip_type const*>(&v1);
759
+#pragma unroll
760
+ for (int i = 0; i < 4; i++) {
761
+ float2 f2 = Converter::convert(p0[i]);
762
+ elements[2 * i] = f2.x;
763
+ elements[2 * i + 1] = f2.y;
764
+ }
765
+#pragma unroll
766
+ for (int i = 0; i < 4; i++) {
767
+ float2 f2 = Converter::convert(p1[i]);
768
+ elements[8 + 2 * i] = f2.x;
769
+ elements[8 + 2 * i + 1] = f2.y;
770
+ }
771
+ }
772
+
773
+ // ── Q branch: RMSNorm (no weight) ─────────────────────────────────────
774
+ if (!isKV) {
775
+ float sumOfSquares = 0.0f;
776
+#pragma unroll
777
+ for (int i = 0; i < kElemsPerLane; i++) {
778
+ sumOfSquares += elements[i] * elements[i];
779
+ }
780
+ sumOfSquares = warpSum<float>(sumOfSquares);
781
+ float const rms_rcp =
782
+ rsqrtf(sumOfSquares / static_cast<float>(kHeadDim) + eps);
783
+#pragma unroll
784
+ for (int i = 0; i < kElemsPerLane; i++) {
785
+ elements[i] = elements[i] * rms_rcp;
786
+ }
787
+ }
788
+
789
+ // ── GPT-J RoPE on dims [NOPE_DIM, HEAD_DIM) ───────────────────────────
790
+ bool const is_rope_lane = dim_base >= kNopeDim;
791
+ if (is_rope_lane) {
792
+ int64_t const pos = position_ids[tokenIdx];
793
+ constexpr int kHalfRope = kRopeDim / 2;
794
+ float const* cos_ptr = cos_sin_cache + pos * kRopeDim;
795
+ float const* sin_ptr = cos_ptr + kHalfRope;
796
+ int const rope_local_base = dim_base - kNopeDim;
797
+ int const half_base = rope_local_base >> 1;
798
+ float4 const c0 = *reinterpret_cast<float4 const*>(cos_ptr + half_base);
799
+ float4 const c1 = *reinterpret_cast<float4 const*>(cos_ptr + half_base + 4);
800
+ float4 const s0 = *reinterpret_cast<float4 const*>(sin_ptr + half_base);
801
+ float4 const s1 = *reinterpret_cast<float4 const*>(sin_ptr + half_base + 4);
802
+ float const cos_arr[8] = {c0.x, c0.y, c0.z, c0.w, c1.x, c1.y, c1.z, c1.w};
803
+ float const sin_arr[8] = {s0.x, s0.y, s0.z, s0.w, s1.x, s1.y, s1.z, s1.w};
804
+#pragma unroll
805
+ for (int p = 0; p < kElemsPerLane / 2; p++) {
806
+ float const x_even = elements[2 * p];
807
+ float const x_odd = elements[2 * p + 1];
808
+ elements[2 * p] = x_even * cos_arr[p] - x_odd * sin_arr[p];
809
+ elements[2 * p + 1] = x_even * sin_arr[p] + x_odd * cos_arr[p];
810
+ }
811
+ }
812
+
813
+ // ── Store ─────────────────────────────────────────────────────────────
814
+ if (!isKV) {
815
+ if constexpr (STORE_Q_FP8) {
816
+ float const scale_inv = VLLM_LDG(q_fp8_scale_inv);
817
+ uint4 const out = packFp8E4M3x16(elements, scale_inv);
818
+ uint8_t* dst = q_fp8_out +
819
+ static_cast<int64_t>(tokenIdx) * q_fp8_stride0 +
820
+ static_cast<int64_t>(slotIdx) * q_fp8_stride1 + dim_base;
821
+ *reinterpret_cast<uint4*>(dst) = out;
822
+ } else {
823
+ uint4 out0, out1;
824
+ auto* po0 = reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
825
+ auto* po1 = reinterpret_cast<typename Converter::packed_hip_type*>(&out1);
826
+#pragma unroll
827
+ for (int i = 0; i < 4; i++) {
828
+ po0[i] = Converter::convert(
829
+ make_float2(elements[2 * i], elements[2 * i + 1]));
830
+ }
831
+#pragma unroll
832
+ for (int i = 0; i < 4; i++) {
833
+ po1[i] = Converter::convert(
834
+ make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
835
+ }
836
+ scalar_t_in* dst =
837
+ q_inout +
838
+ (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) * kHeadDim +
839
+ dim_base;
840
+ *reinterpret_cast<uint4*>(dst) = out0;
841
+ *reinterpret_cast<uint4*>(dst + 8) = out1;
842
+ }
843
+ } else {
844
+ int64_t const slot_id = slot_mapping[tokenIdx];
845
+ if (slot_id >= 0) {
846
+ int64_t const block_idx = slot_id / cache_block_size;
847
+ int64_t const pos_in_block = slot_id % cache_block_size;
848
+ uint8_t* cache_row =
849
+ k_cache + block_idx * kv_block_stride + pos_in_block * kv_token_stride;
850
+ if constexpr (STORE_KV_FP8) {
851
+ float const inv_scale = 1.0f / VLLM_LDG(fp8_scale_ptr);
852
+ uint4 const out = packFp8E4M3x16(elements, inv_scale);
853
+ *reinterpret_cast<uint4*>(cache_row + dim_base) = out;
854
+ } else {
855
+ uint4 out0, out1;
856
+ auto* po0 =
857
+ reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
858
+ auto* po1 =
859
+ reinterpret_cast<typename Converter::packed_hip_type*>(&out1);
860
+#pragma unroll
861
+ for (int i = 0; i < 4; i++) {
862
+ po0[i] = Converter::convert(
863
+ make_float2(elements[2 * i], elements[2 * i + 1]));
864
+ }
865
+#pragma unroll
866
+ for (int i = 0; i < 4; i++) {
867
+ po1[i] = Converter::convert(
868
+ make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
869
+ }
870
+ scalar_t_in* dst = reinterpret_cast<scalar_t_in*>(cache_row) + dim_base;
871
+ *reinterpret_cast<uint4*>(dst) = out0;
872
+ *reinterpret_cast<uint4*>(dst + 8) = out1;
873
+ }
874
+ }
875
+ }
876
+
877
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
878
+ cudaTriggerProgrammaticLaunchCompletion();
879
+#endif
880
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
881
+ }
882
+#endif
883
+}
884
+
885
+// Configure + launch helper shared by the bf16 and fp8 full-cache launchers.
886
+template <typename scalar_t_in, bool STORE_Q_FP8, bool STORE_KV_FP8>
887
+static void launchFullCacheKernel(
888
+ scalar_t_in* q_inout, uint8_t* q_fp8_out, int64_t q_fp8_stride0,
889
+ int64_t q_fp8_stride1, scalar_t_in const* kv_in, uint8_t* k_cache,
890
+ int64_t const* slot_mapping, int64_t const* position_ids,
891
+ float const* cos_sin_cache, float const* fp8_scale,
892
+ float const* q_fp8_scale_inv, float const eps, int const num_tokens_full,
893
+ int const num_tokens_insert, int const num_heads_q,
894
+ int const cache_block_size, int64_t const kv_block_stride,
895
+ int64_t const kv_token_stride, char const* op_name, cudaStream_t stream) {
896
+ constexpr int kBlockSize = 256;
897
+ constexpr int kWarpsPerBlock = kBlockSize / 32;
898
+ int64_t const total_warps =
899
+ static_cast<int64_t>(num_tokens_full) * (num_heads_q + 1);
900
+ int const grid =
901
+ static_cast<int>((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock);
902
+ auto* kernel =
903
+ fusedDeepseekV4FullCacheKernel<scalar_t_in, STORE_Q_FP8, STORE_KV_FP8>;
904
+#ifndef USE_ROCM
905
+ static int const sm_version = getSMVersion();
906
+ STD_TORCH_CHECK(sm_version >= 80, op_name,
907
+ " requires sm_80+ (Ampere or newer); got sm_", sm_version);
908
+ cudaLaunchConfig_t config;
909
+ config.gridDim = dim3(grid);
910
+ config.blockDim = dim3(kBlockSize);
911
+ config.dynamicSmemBytes = 0;
912
+ config.stream = stream;
913
+ cudaLaunchAttribute attrs[1];
914
+ attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
915
+ attrs[0].val.programmaticStreamSerializationAllowed = 1;
916
+ config.attrs = attrs;
917
+ config.numAttrs = (sm_version >= 90) ? 1 : 0;
918
+ cudaLaunchKernelEx(&config, kernel, q_inout, q_fp8_out, q_fp8_stride0,
919
+ q_fp8_stride1, kv_in, k_cache, slot_mapping, position_ids,
920
+ cos_sin_cache, fp8_scale, q_fp8_scale_inv, eps,
921
+ num_tokens_full, num_tokens_insert, num_heads_q,
922
+ cache_block_size, kv_block_stride, kv_token_stride);
923
+#else
924
+ kernel<<<grid, kBlockSize, 0, stream>>>(
925
+ q_inout, q_fp8_out, q_fp8_stride0, q_fp8_stride1, kv_in, k_cache,
926
+ slot_mapping, position_ids, cos_sin_cache, fp8_scale, q_fp8_scale_inv,
927
+ eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
928
+ kv_block_stride, kv_token_stride);
929
+#endif
930
+}
931
+
932
} // namespace deepseek_v4_fused_ops
933
} // namespace vllm
934
@@ -735,3 +1015,167 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
1015
});
1016
return q_out;
1017
}
1018
+
1019
+// ────────────────────────────────────────────────────────────────────────────
1020
+// FlashInfer full-cache torch ops
1021
+// ────────────────────────────────────────────────────────────────────────────
1022
+void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
1023
+ torch::stable::Tensor& q, // [N, H, 512] bf16, in place
1024
+ torch::stable::Tensor const& kv, // [N, 512] bf16, read-only
1025
+ torch::stable::Tensor& k_cache, // [num_blocks, bs, 512] bf16
1026
+ torch::stable::Tensor const& slot_mapping, // [num_tokens_insert] int64
1027
+ torch::stable::Tensor const& position_ids, // [N] int64
1028
+ torch::stable::Tensor const& cos_sin_cache, // [max_pos, 64] float32
1029
+ double eps, int64_t cache_block_size) {
1030
+ using torch::headeronly::ScalarType;
1031
+ STD_TORCH_CHECK(q.device().is_cuda() && q.is_contiguous(),
1032
+ "q must be contiguous CUDA");
1033
+ STD_TORCH_CHECK(kv.device().is_cuda() && kv.is_contiguous(),
1034
+ "kv must be contiguous CUDA");
1035
+ STD_TORCH_CHECK(k_cache.device().is_cuda(), "k_cache must be CUDA");
1036
+ STD_TORCH_CHECK(slot_mapping.device().is_cuda() &&
1037
+ slot_mapping.scalar_type() == ScalarType::Long,
1038
+ "slot_mapping must be int64 CUDA");
1039
+ STD_TORCH_CHECK(position_ids.device().is_cuda() &&
1040
+ position_ids.scalar_type() == ScalarType::Long,
1041
+ "position_ids must be int64 CUDA");
1042
+ STD_TORCH_CHECK(cos_sin_cache.device().is_cuda() &&
1043
+ cos_sin_cache.scalar_type() == ScalarType::Float &&
1044
+ cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64,
1045
+ "cos_sin_cache shape [max_pos, 64] float32");
1046
+ STD_TORCH_CHECK(q.dim() == 3 && q.size(2) == 512, "q shape [N, H, 512]");
1047
+ STD_TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");
1048
+ STD_TORCH_CHECK(q.scalar_type() == ScalarType::BFloat16 &&
1049
+ kv.scalar_type() == ScalarType::BFloat16,
1050
+ "q and kv must be bfloat16");
1051
+ STD_TORCH_CHECK(k_cache.dim() == 3 && k_cache.size(1) == cache_block_size &&
1052
+ k_cache.size(2) == 512 && k_cache.stride(2) == 1,
1053
+ "k_cache shape [num_blocks, cache_block_size, 512] contiguous");
1054
+ STD_TORCH_CHECK(k_cache.scalar_type() == ScalarType::BFloat16,
1055
+ "k_cache must be bfloat16");
1056
+
1057
+ int const num_tokens_full = static_cast<int>(q.size(0));
1058
+ int const num_tokens_insert = static_cast<int>(slot_mapping.size(0));
1059
+ STD_TORCH_CHECK(static_cast<int>(kv.size(0)) == num_tokens_full &&
1060
+ static_cast<int>(position_ids.size(0)) == num_tokens_full,
1061
+ "q/kv/position_ids row counts must match");
1062
+ STD_TORCH_CHECK(num_tokens_insert <= num_tokens_full,
1063
+ "slot_mapping must not exceed q row count");
1064
+ int const num_heads_q = static_cast<int>(q.size(1));
1065
+
1066
+ const torch::stable::accelerator::DeviceGuard device_guard(
1067
+ q.get_device_index());
1068
+ const cudaStream_t stream = get_current_cuda_stream(q.get_device_index());
1069
+
1070
+ // bf16 cache: 2 bytes/element -> byte strides for the uint8-addressed kernel.
1071
+ int64_t const kv_block_stride = k_cache.stride(0) * 2;
1072
+ int64_t const kv_token_stride = k_cache.stride(1) * 2;
1073
+
1074
+ VLLM_STABLE_DISPATCH_HALF_TYPES(
1075
+ q.scalar_type(),
1076
+ "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert", [&] {
1077
+ vllm::deepseek_v4_fused_ops::launchFullCacheKernel<scalar_t, false,
1078
+ false>(
1079
+ reinterpret_cast<scalar_t*>(q.mutable_data_ptr()), nullptr, 0, 0,
1080
+ reinterpret_cast<scalar_t const*>(kv.const_data_ptr()),
1081
+ reinterpret_cast<uint8_t*>(k_cache.mutable_data_ptr()),
1082
+ slot_mapping.const_data_ptr<int64_t>(),
1083
+ position_ids.const_data_ptr<int64_t>(),
1084
+ cos_sin_cache.const_data_ptr<float>(), nullptr, nullptr,
1085
+ static_cast<float>(eps), num_tokens_full, num_tokens_insert,
1086
+ num_heads_q, static_cast<int>(cache_block_size), kv_block_stride,
1087
+ kv_token_stride,
1088
+ "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert",
1089
+ stream);
1090
+ });
1091
+}
1092
+
1093
+void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
1094
+ torch::stable::Tensor const& q, // [N, H, 512] bf16, read-only
1095
+ torch::stable::Tensor const& kv, // [N, 512] bf16, read-only
1096
+ torch::stable::Tensor& q_fp8, // [N, H, 512] fp8 e4m3
1097
+ torch::stable::Tensor& k_cache, // [num_blocks, bs, 512] fp8
1098
+ torch::stable::Tensor const& slot_mapping, // [num_tokens_insert] int64
1099
+ torch::stable::Tensor const& position_ids, // [N] int64
1100
+ torch::stable::Tensor const& cos_sin_cache, // [max_pos, 64] float32
1101
+ torch::stable::Tensor const& fp8_scale, // scalar float32 (KV scale)
1102
+ torch::stable::Tensor const& q_fp8_scale_inv, // scalar float32 (1 / Q scale)
1103
+ double eps, int64_t cache_block_size) {
1104
+ using torch::headeronly::ScalarType;
1105
+ STD_TORCH_CHECK(q.device().is_cuda() && q.is_contiguous(),
1106
+ "q must be contiguous CUDA");
1107
+ STD_TORCH_CHECK(kv.device().is_cuda() && kv.is_contiguous(),
1108
+ "kv must be contiguous CUDA");
1109
+ STD_TORCH_CHECK(q_fp8.device().is_cuda() && q_fp8.is_contiguous() &&
1110
+ q_fp8.scalar_type() == ScalarType::Float8_e4m3fn &&
1111
+ q_fp8.dim() == 3 && q_fp8.size(0) == q.size(0) &&
1112
+ q_fp8.size(1) == q.size(1) && q_fp8.size(2) == q.size(2),
1113
+ "q_fp8 must be a contiguous float8_e4m3fn tensor matching q");
1114
+ STD_TORCH_CHECK(k_cache.device().is_cuda(), "k_cache must be CUDA");
1115
+ STD_TORCH_CHECK(slot_mapping.device().is_cuda() &&
1116
+ slot_mapping.scalar_type() == ScalarType::Long,
1117
+ "slot_mapping must be int64 CUDA");
1118
+ STD_TORCH_CHECK(position_ids.device().is_cuda() &&
1119
+ position_ids.scalar_type() == ScalarType::Long,
1120
+ "position_ids must be int64 CUDA");
1121
+ STD_TORCH_CHECK(cos_sin_cache.device().is_cuda() &&
1122
+ cos_sin_cache.scalar_type() == ScalarType::Float &&
1123
+ cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64,
1124
+ "cos_sin_cache shape [max_pos, 64] float32");
1125
+ STD_TORCH_CHECK(fp8_scale.device().is_cuda() &&
1126
+ fp8_scale.scalar_type() == ScalarType::Float &&
1127
+ fp8_scale.size(0) == 1,
1128
+ "fp8_scale must be a scalar float32 CUDA tensor");
1129
+ STD_TORCH_CHECK(q_fp8_scale_inv.device().is_cuda() &&
1130
+ q_fp8_scale_inv.scalar_type() == ScalarType::Float &&
1131
+ q_fp8_scale_inv.size(0) == 1,
1132
+ "q_fp8_scale_inv must be a scalar float32 CUDA tensor");
1133
+ STD_TORCH_CHECK(q.dim() == 3 && q.size(2) == 512, "q shape [N, H, 512]");
1134
+ STD_TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");
1135
+ STD_TORCH_CHECK(q.scalar_type() == kv.scalar_type(),
1136
+ "q and kv dtype must match");
1137
+ STD_TORCH_CHECK(k_cache.dim() == 3 && k_cache.size(1) == cache_block_size &&
1138
+ k_cache.size(2) == 512 && k_cache.stride(2) == 1,
1139
+ "k_cache shape [num_blocks, cache_block_size, 512] contiguous");
1140
+ STD_TORCH_CHECK(k_cache.scalar_type() == ScalarType::Float8_e4m3fn,
1141
+ "k_cache must be float8_e4m3fn");
1142
+
1143
+ int const num_tokens_full = static_cast<int>(q.size(0));
1144
+ int const num_tokens_insert = static_cast<int>(slot_mapping.size(0));
1145
+ STD_TORCH_CHECK(static_cast<int>(kv.size(0)) == num_tokens_full &&
1146
+ static_cast<int>(position_ids.size(0)) == num_tokens_full,
1147
+ "q/kv/position_ids row counts must match");
1148
+ STD_TORCH_CHECK(num_tokens_insert <= num_tokens_full,
1149
+ "slot_mapping must not exceed q row count");
1150
+ int const num_heads_q = static_cast<int>(q.size(1));
1151
+
1152
+ const torch::stable::accelerator::DeviceGuard device_guard(
1153
+ q.get_device_index());
1154
+ const cudaStream_t stream = get_current_cuda_stream(q.get_device_index());
1155
+
1156
+ VLLM_STABLE_DISPATCH_HALF_TYPES(
1157
+ q.scalar_type(),
1158
+ "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert", [&] {
1159
+ vllm::deepseek_v4_fused_ops::launchFullCacheKernel<scalar_t, true,
1160
+ true>(
1161
+ // q is read-only in the fp8 path (the kernel writes q_fp8); the
1162
+ // launcher signature is non-const, so cast away const on the ptr.
1163
+ reinterpret_cast<scalar_t*>(
1164
+ const_cast<void*>(q.const_data_ptr())),
1165
+ reinterpret_cast<uint8_t*>(q_fp8.mutable_data_ptr()),
1166
+ q_fp8.stride(0), q_fp8.stride(1),
1167
+ reinterpret_cast<scalar_t const*>(kv.const_data_ptr()),
1168
+ reinterpret_cast<uint8_t*>(k_cache.mutable_data_ptr()),
1169
+ slot_mapping.const_data_ptr<int64_t>(),
1170
+ position_ids.const_data_ptr<int64_t>(),
1171
+ cos_sin_cache.const_data_ptr<float>(),
1172
+ fp8_scale.const_data_ptr<float>(),
1173
+ q_fp8_scale_inv.const_data_ptr<float>(), static_cast<float>(eps),
1174
+ num_tokens_full, num_tokens_insert, num_heads_q,
1175
+ static_cast<int>(cache_block_size),
1176
+ // fp8 cache: 1 byte/element -> stride already in bytes.
1177
+ k_cache.stride(0), k_cache.stride(1),
1178
+ "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert",
1179
+ stream);
1180
+ });
1181
+}
csrc/libtorch_stable/ops.h
+17
@@ -238,6 +238,23 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
238
torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded,
239
double eps, int64_t cache_block_size);
240
241
+void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
242
+ torch::stable::Tensor& q, torch::stable::Tensor const& kv,
243
+ torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping,
244
+ torch::stable::Tensor const& position_ids,
245
+ torch::stable::Tensor const& cos_sin_cache, double eps,
246
+ int64_t cache_block_size);
247
+
248
+void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
249
+ torch::stable::Tensor const& q, torch::stable::Tensor const& kv,
250
+ torch::stable::Tensor& q_fp8, torch::stable::Tensor& k_cache,
251
+ torch::stable::Tensor const& slot_mapping,
252
+ torch::stable::Tensor const& position_ids,
253
+ torch::stable::Tensor const& cos_sin_cache,
254
+ torch::stable::Tensor const& fp8_scale,
255
+ torch::stable::Tensor const& q_fp8_scale_inv, double eps,
256
+ int64_t cache_block_size);
257
+
258
#ifndef USE_ROCM
259
torch::stable::Tensor minimax_allreduce_rms(
260
torch::stable::Tensor const& input,
csrc/libtorch_stable/torch_bindings.cpp
+20
@@ -343,6 +343,20 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
343
"Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
344
"int q_head_padded, float eps, int cache_block_size) -> Tensor");
345
346
+ // FlashInfer V4 full-cache variants: write Q in place (bf16) or to a separate
347
+ // FP8 tensor, and KV into a contiguous 512-wide token-strided cache.
348
+ ops.def(
349
+ "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert("
350
+ "Tensor! q, Tensor kv, Tensor! k_cache, Tensor slot_mapping, "
351
+ "Tensor position_ids, Tensor cos_sin_cache, float eps, "
352
+ "int cache_block_size) -> ()");
353
+ ops.def(
354
+ "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert("
355
+ "Tensor q, Tensor kv, Tensor! q_fp8, Tensor! k_cache, "
356
+ "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
357
+ "Tensor fp8_scale, Tensor q_fp8_scale_inv, float eps, "
358
+ "int cache_block_size) -> ()");
359
+
360
#ifndef USE_ROCM
361
ops.def(
362
"minimax_allreduce_rms("
@@ -591,6 +605,12 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
605
ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope));
606
ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert",
607
TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert));
608
+ ops.impl(
609
+ "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert",
610
+ TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert));
611
+ ops.impl(
612
+ "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert",
613
+ TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert));
614
#ifndef USE_ROCM
615
ops.impl("minimax_allreduce_rms", TORCH_BOX(&minimax_allreduce_rms));
616
ops.impl("minimax_allreduce_rms_qk", TORCH_BOX(&minimax_allreduce_rms_qk));
csrc/torch_bindings.cpp
+2
-1
@@ -55,7 +55,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
55
56
// Horizontally-fused DeepseekV4-MLA: per-head RMSNorm + GPT-J RoPE for Q, and
57
// GPT-J RoPE + UE8M0 FP8 quant + paged cache insert for KV, all in one
58
- // kernel launch. Registered in _C_stable_libtorch.
58
+ // kernel launch. Registered in _C_stable_libtorch (incl. the FlashInfer V4
59
+ // full-cache bf16/fp8 variants).
60
61
// Quantization ops
62
#ifndef USE_ROCM
docs/design/attention_backends.md
+14
@@ -228,3 +228,17 @@ MLA decode backends are selected using the standard
228
| `TOKENSPEED_MLA` | fp16, bf16 | `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
229
| `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
230
| `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | Any | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | Any |
231
+
232
+### DeepSeek V4 Decode Backends
233
+
234
+DeepSeek V4 sparse MLA uses its own decode backends, selected via
235
+`--attention-backend=<BACKEND>` (e.g., `FLASHMLA_SPARSE_DSV4`,
236
+`FLASHINFER_MLA_SPARSE_DSV4`). They share the V4 sparse-index
237
+pipeline (compressor + SWA + indexer, 256-token blocks, head 512);
238
+default on NVIDIA is `FLASHMLA_SPARSE_DSV4`.
239
+
240
+| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
241
+| ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ |
242
+| `FLASHINFER_MLA_SPARSE_DSV4` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
243
+| `FLASHMLA_SPARSE_DSV4` | fp16, bf16 | `auto` | 256 | 512 | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
244
+| `ROCM_FLASHMLA_SPARSE_DSV4` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
tests/kernels/test_compressor_kv_cache.py
+149
@@ -468,6 +468,7 @@ def _reference_kv_compress_norm_rope(
468
use_fp4: bool = False,
469
rms_eps: float = 1e-6,
470
fp8_max: float = 448.0,
471
+ return_full_cache: bool = False,
472
):
473
"""Compress → RMSNorm → GPT-J RoPE → quantize.
474
@@ -521,6 +522,12 @@ def _reference_kv_compress_norm_rope(
522
results.append(torch.cat([nope, rope]).to(state_cache.dtype))
523
result = torch.stack(results)
524
525
+ if return_full_cache:
526
+ # Contiguous 512-wide bf16 row (nope unrotated + rope rotated), matching
527
+ # the FlashInfer full-cache layout before any per-tensor fp8 quant. The
528
+ # kernel rounds the fp32 result to bf16 once at the store.
529
+ return result.to(torch.bfloat16)
530
+
531
if use_fp4:
532
return quantize_to_mxfp4(result)
533
else:
@@ -667,3 +674,145 @@ def test_fused_kv_insert_indexer(num_tokens: int, kv_block_size: int, use_fp4: b
674
assert torch.equal(actual_scale, scale[i : i + 1]), (
675
f"token {i}: scale {actual_scale.item()} != {scale[i].item()}"
676
)
677
+
678
+
679
+@pytest.mark.parametrize("compress_ratio", [4, 128])
680
+@pytest.mark.parametrize("store_fp8", [False, True])
681
+def test_cutedsl_full_cache_store(compress_ratio: int, store_fp8: bool):
682
+ """CuTeDSL compressor full-cache (FlashInfer) store parity for head=512.
683
+
684
+ Exercises the contiguous bf16 / per-tensor fp8 store branch of both the C4
685
+ fused kernel and the C128 split kernel against the PyTorch reference.
686
+ """
687
+ cutedsl = pytest.importorskip("cutlass") # noqa: F841
688
+ from vllm.models.deepseek_v4.nvidia.ops.sparse_attn_compress_cutedsl import (
689
+ fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl,
690
+ split_kv_compress_norm_rope_insert_sparse_attn_cutedsl,
691
+ )
692
+
693
+ HEAD_DIM = 512
694
+ ROPE_DIM = 64
695
+ RMS_EPS = 1e-6
696
+ FP8_MAX = 448.0
697
+ # C128 compress (Block8 kernel) requires state-cache block_size=8; C4 uses 16.
698
+ BLOCK_SIZE = 8 if compress_ratio == 128 else 16
699
+ KV_BLOCK_SIZE = 64
700
+ device = "cuda"
701
+ torch.manual_seed(7)
702
+
703
+ overlap = 1 if compress_ratio == 4 else 0
704
+ coff = 1 + overlap
705
+ num_tokens = 8
706
+
707
+ num_pages = (compress_ratio * num_tokens - 1) // BLOCK_SIZE + 2
708
+ # The production CompressorStateCache is fp32.
709
+ state_cache = torch.randn(
710
+ num_pages, BLOCK_SIZE, 2 * coff * HEAD_DIM, dtype=torch.float32, device=device
711
+ )
712
+ block_table = torch.arange(num_pages, dtype=torch.int32, device=device).unsqueeze(0)
713
+ token_to_req = torch.zeros(num_tokens, dtype=torch.int32, device=device)
714
+ slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
715
+ positions = torch.arange(
716
+ compress_ratio - 1,
717
+ compress_ratio * num_tokens,
718
+ compress_ratio,
719
+ dtype=torch.int64,
720
+ device=device,
721
+ )
722
+ rms_weight = torch.randn(HEAD_DIM, dtype=torch.bfloat16, device=device)
723
+ cos_sin_cache = torch.randn(
724
+ compress_ratio * num_tokens, ROPE_DIM, dtype=torch.float32, device=device
725
+ )
726
+
727
+ dtype = torch.float8_e4m3fn if store_fp8 else torch.bfloat16
728
+ kv_n_blocks = (num_tokens + KV_BLOCK_SIZE - 1) // KV_BLOCK_SIZE + 1
729
+ k_cache = torch.zeros(
730
+ kv_n_blocks, KV_BLOCK_SIZE, HEAD_DIM, dtype=dtype, device=device
731
+ )
732
+ fp8_scale = torch.tensor(
733
+ [0.5 if store_fp8 else 1.0], dtype=torch.float32, device=device
734
+ )
735
+
736
+ if compress_ratio == 4:
737
+ fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
738
+ state_cache,
739
+ token_to_req,
740
+ positions,
741
+ slot_mapping,
742
+ block_table,
743
+ BLOCK_SIZE,
744
+ rms_weight,
745
+ RMS_EPS,
746
+ cos_sin_cache,
747
+ k_cache,
748
+ slot_mapping,
749
+ KV_BLOCK_SIZE,
750
+ k_cache.stride(0),
751
+ head_size=HEAD_DIM,
752
+ state_width=coff * HEAD_DIM,
753
+ rope_head_dim=ROPE_DIM,
754
+ fp8_max=FP8_MAX,
755
+ quant_block=64,
756
+ token_stride=576,
757
+ scale_dim=8,
758
+ compress_ratio=compress_ratio,
759
+ overlap=True,
760
+ store_full_kv=True,
761
+ store_full_fp8=store_fp8,
762
+ fp8_scale=fp8_scale,
763
+ )
764
+ else:
765
+ compressed_kv = torch.empty(
766
+ (num_tokens, HEAD_DIM), dtype=torch.float32, device=device
767
+ )
768
+ split_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
769
+ state_cache,
770
+ token_to_req,
771
+ positions,
772
+ slot_mapping,
773
+ block_table,
774
+ BLOCK_SIZE,
775
+ compressed_kv,
776
+ rms_weight,
777
+ RMS_EPS,
778
+ cos_sin_cache,
779
+ k_cache,
780
+ slot_mapping,
781
+ KV_BLOCK_SIZE,
782
+ k_cache.stride(0),
783
+ head_size=HEAD_DIM,
784
+ state_width=coff * HEAD_DIM,
785
+ rope_head_dim=ROPE_DIM,
786
+ fp8_max=FP8_MAX,
787
+ quant_block=64,
788
+ token_stride=576,
789
+ scale_dim=8,
790
+ compress_ratio=compress_ratio,
791
+ overlap=bool(overlap),
792
+ store_full_kv=True,
793
+ store_full_fp8=store_fp8,
794
+ fp8_scale=fp8_scale,
795
+ )
796
+
797
+ ref = _reference_kv_compress_norm_rope(
798
+ state_cache,
799
+ block_table,
800
+ positions,
801
+ rms_weight,
802
+ cos_sin_cache,
803
+ compress_ratio,
804
+ overlap,
805
+ rms_eps=RMS_EPS,
806
+ return_full_cache=True,
807
+ ) # [num_tokens, HEAD_DIM] bf16
808
+
809
+ actual = torch.stack(
810
+ [k_cache[i // KV_BLOCK_SIZE, i % KV_BLOCK_SIZE] for i in range(num_tokens)]
811
+ )
812
+ if store_fp8:
813
+ ref_fp8 = torch.clamp(ref.float() / fp8_scale, -FP8_MAX, FP8_MAX).to(
814
+ torch.float8_e4m3fn
815
+ )
816
+ torch.testing.assert_close(actual.float(), ref_fp8.float(), rtol=0.0, atol=0.3)
817
+ else:
818
+ torch.testing.assert_close(actual.float(), ref.float(), rtol=3e-2, atol=3e-2)
tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
+248
-1
@@ -67,7 +67,7 @@ def apply_rope_gptj_last_k(
67
head_dim = x.shape[-1]
68
nope_dim = head_dim - rope_dim
69
70
- cs = cos_sin_cache[positions].to(torch.float32)
70
+ cs = cos_sin_cache[positions.long()].to(torch.float32)
71
cos = cs[..., :half]
72
sin = cs[..., half:]
73
@@ -114,6 +114,18 @@ def _op_available() -> bool:
114
return hasattr(torch.ops._C, "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert")
115
116
117
+def _full_cache_fp8_op_available() -> bool:
118
+ return hasattr(
119
+ torch.ops._C, "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert"
120
+ )
121
+
122
+
123
+def _full_cache_bf16_op_available() -> bool:
124
+ return hasattr(
125
+ torch.ops._C, "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert"
126
+ )
127
+
128
+
129
pytestmark = pytest.mark.skipif(
130
not torch.cuda.is_available() or not _op_available(),
131
reason="CUDA not available or fused DeepseekV4 op not built in",
@@ -415,3 +427,238 @@ def test_combined_q_and_kv(
427
"padded head slots must be exact zero"
428
)
429
torch.testing.assert_close(k_cache_fused, k_cache_ref, rtol=0, atol=0)
430
+
431
+
432
+# ── Full-cache (FlashInfer) path parity ──────────────────────────────────────
433
+
434
+
435
+def _call_full_cache_fp8_fused(
436
+ q,
437
+ kv,
438
+ q_fp8,
439
+ k_cache,
440
+ slot_mapping,
441
+ positions,
442
+ cos_sin_cache,
443
+ fp8_scale,
444
+ q_fp8_scale_inv,
445
+ eps,
446
+ bs,
447
+):
448
+ torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
449
+ q,
450
+ kv,
451
+ q_fp8,
452
+ k_cache,
453
+ slot_mapping,
454
+ positions.long(),
455
+ cos_sin_cache,
456
+ fp8_scale,
457
+ q_fp8_scale_inv,
458
+ eps,
459
+ bs,
460
+ )
461
+
462
+
463
+def _call_full_cache_bf16_fused(
464
+ q,
465
+ kv,
466
+ k_cache,
467
+ slot_mapping,
468
+ positions,
469
+ cos_sin_cache,
470
+ eps,
471
+ bs,
472
+):
473
+ torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
474
+ q,
475
+ kv,
476
+ k_cache,
477
+ slot_mapping,
478
+ positions.long(),
479
+ cos_sin_cache,
480
+ eps,
481
+ bs,
482
+ )
483
+
484
+
485
+def _fp8_full_cache_reference(
486
+ q,
487
+ kv,
488
+ k_cache,
489
+ q_fp8,
490
+ slot_mapping,
491
+ positions,
492
+ cos_sin_cache,
493
+ eps,
494
+ block_size,
495
+ fp8_scale,
496
+ q_fp8_scale_inv,
497
+):
498
+ q_ref = rmsnorm_no_weight(q, eps)
499
+ q_ref = apply_rope_gptj_last_k(q_ref, positions, cos_sin_cache)
500
+ q_fp8.copy_(
501
+ torch.clamp(q_ref.float() * q_fp8_scale_inv, -FP8_MAX, FP8_MAX).to(
502
+ torch.float8_e4m3fn
503
+ )
504
+ )
505
+
506
+ kv_ref = apply_rope_gptj_last_k(kv, positions, cos_sin_cache)
507
+ valid = slot_mapping >= 0
508
+ slots = slot_mapping[valid]
509
+ block_idx = slots // block_size
510
+ pos_in_block = slots % block_size
511
+ k_cache[block_idx, pos_in_block] = torch.clamp(
512
+ kv_ref[valid].float() / fp8_scale, -FP8_MAX, FP8_MAX
513
+ ).to(torch.float8_e4m3fn)
514
+
515
+
516
+def _bf16_full_cache_reference(
517
+ q,
518
+ kv,
519
+ k_cache,
520
+ slot_mapping,
521
+ positions,
522
+ cos_sin_cache,
523
+ eps,
524
+ block_size,
525
+):
526
+ q_ref = rmsnorm_no_weight(q, eps)
527
+ # Kernel keeps RMSNorm+RoPE in fp32 and rounds to bf16 once at the store.
528
+ q_ref = apply_rope_gptj_last_k(q_ref, positions, cos_sin_cache).to(q.dtype)
529
+
530
+ kv_ref = apply_rope_gptj_last_k(kv, positions, cos_sin_cache)
531
+ valid = slot_mapping >= 0
532
+ slots = slot_mapping[valid]
533
+ block_idx = slots // block_size
534
+ pos_in_block = slots % block_size
535
+ k_cache[block_idx, pos_in_block] = kv_ref[valid]
536
+ return q_ref
537
+
538
+
539
+@pytest.mark.skipif(
540
+ not _full_cache_fp8_op_available(),
541
+ reason="full-cache per-tensor FP8 DeepseekV4 op not built in",
542
+)
543
+@pytest.mark.parametrize("num_tokens", [4, 17])
544
+@pytest.mark.parametrize("n_heads", [8, 17])
545
+@pytest.mark.parametrize("positions_dtype", [torch.int32, torch.int64])
546
+def test_full_cache_per_tensor_fp8_matches_reference(
547
+ num_tokens: int,
548
+ n_heads: int,
549
+ positions_dtype: torch.dtype,
550
+):
551
+ torch.manual_seed(4)
552
+ device = "cuda"
553
+ dtype = torch.bfloat16
554
+ eps = 1e-6
555
+ block_size = 16
556
+ max_pos = 4096
557
+
558
+ q = torch.randn(num_tokens, n_heads, HEAD_DIM, dtype=dtype, device=device)
559
+ kv = torch.randn(num_tokens, HEAD_DIM, dtype=dtype, device=device)
560
+ positions = torch.arange(num_tokens, dtype=positions_dtype, device=device)
561
+ cos_sin_cache = make_cos_sin_cache(max_pos, ROPE_DIM, torch.float32, device)
562
+
563
+ num_blocks = (num_tokens + block_size - 1) // block_size + 1
564
+ slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
565
+ fp8_scale = torch.tensor([1.0], dtype=torch.float32, device=device)
566
+ q_fp8_scale_inv = torch.tensor([1.0], dtype=torch.float32, device=device)
567
+
568
+ q_fp8_ref = torch.empty_like(q, dtype=torch.float8_e4m3fn)
569
+ q_fp8_fused = torch.empty_like(q, dtype=torch.float8_e4m3fn)
570
+ k_cache_ref = torch.zeros(
571
+ num_blocks, block_size, HEAD_DIM, dtype=torch.float8_e4m3fn, device=device
572
+ )
573
+ k_cache_fused = torch.zeros_like(k_cache_ref)
574
+
575
+ _fp8_full_cache_reference(
576
+ q,
577
+ kv,
578
+ k_cache_ref,
579
+ q_fp8_ref,
580
+ slot_mapping,
581
+ positions,
582
+ cos_sin_cache,
583
+ eps,
584
+ block_size,
585
+ fp8_scale,
586
+ q_fp8_scale_inv,
587
+ )
588
+ _call_full_cache_fp8_fused(
589
+ q.clone(),
590
+ kv,
591
+ q_fp8_fused,
592
+ k_cache_fused,
593
+ slot_mapping,
594
+ positions,
595
+ cos_sin_cache,
596
+ fp8_scale,
597
+ q_fp8_scale_inv,
598
+ eps,
599
+ block_size,
600
+ )
601
+
602
+ torch.testing.assert_close(
603
+ q_fp8_fused.float(), q_fp8_ref.float(), rtol=0, atol=0.25
604
+ )
605
+ torch.testing.assert_close(
606
+ k_cache_fused.float(), k_cache_ref.float(), rtol=0, atol=0.25
607
+ )
608
+
609
+
610
+@pytest.mark.skipif(
611
+ not _full_cache_bf16_op_available(),
612
+ reason="full-cache BF16 DeepseekV4 op not built in",
613
+)
614
+@pytest.mark.parametrize("num_tokens", [4, 17])
615
+@pytest.mark.parametrize("n_heads", [8, 17])
616
+@pytest.mark.parametrize("positions_dtype", [torch.int32, torch.int64])
617
+def test_full_cache_bf16_matches_reference(
618
+ num_tokens: int,
619
+ n_heads: int,
620
+ positions_dtype: torch.dtype,
621
+):
622
+ torch.manual_seed(5)
623
+ device = "cuda"
624
+ dtype = torch.bfloat16
625
+ eps = 1e-6
626
+ block_size = 16
627
+ max_pos = 4096
628
+
629
+ q = torch.randn(num_tokens, n_heads, HEAD_DIM, dtype=dtype, device=device)
630
+ kv = torch.randn(num_tokens, HEAD_DIM, dtype=dtype, device=device)
631
+ positions = torch.arange(num_tokens, dtype=positions_dtype, device=device)
632
+ cos_sin_cache = make_cos_sin_cache(max_pos, ROPE_DIM, torch.float32, device)
633
+
634
+ num_blocks = (num_tokens + block_size - 1) // block_size + 1
635
+ slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
636
+
637
+ q_fused = q.clone()
638
+ k_cache_ref = torch.zeros(
639
+ num_blocks, block_size, HEAD_DIM, dtype=torch.bfloat16, device=device
640
+ )
641
+ k_cache_fused = torch.zeros_like(k_cache_ref)
642
+ q_ref = _bf16_full_cache_reference(
643
+ q,
644
+ kv,
645
+ k_cache_ref,
646
+ slot_mapping,
647
+ positions,
648
+ cos_sin_cache,
649
+ eps,
650
+ block_size,
651
+ )
652
+ _call_full_cache_bf16_fused(
653
+ q_fused,
654
+ kv,
655
+ k_cache_fused,
656
+ slot_mapping,
657
+ positions,
658
+ cos_sin_cache,
659
+ eps,
660
+ block_size,
661
+ )
662
+
663
+ torch.testing.assert_close(q_fused, q_ref, rtol=1e-2, atol=1e-2)
664
+ torch.testing.assert_close(k_cache_fused, k_cache_ref, rtol=0, atol=0)
tools/pre_commit/generate_attention_backend_docs.py
+31
-5
@@ -1562,7 +1562,9 @@ def generate_legend() -> str:
1562
1563
1564
def generate_mla_section(
1565
- prefill_backends: list[dict[str, Any]], decode_backends: list[dict[str, Any]]
1565
+ prefill_backends: list[dict[str, Any]],
1566
+ decode_backends: list[dict[str, Any]],
1567
+ v4_decode_backends: list[dict[str, Any]] | None = None,
1568
) -> str:
1569
"""Generate the complete MLA section with prefill and decode tables."""
1570
lines = [
@@ -1611,6 +1613,22 @@ def generate_mla_section(
1613
columns = _build_columns(is_mla=True, has_versions=False)
1614
lines.extend(_render_table(columns, decode_backends))
1615
1616
+ if v4_decode_backends:
1617
+ lines.extend(
1618
+ [
1619
+ "",
1620
+ "### DeepSeek V4 Decode Backends",
1621
+ "",
1622
+ "DeepSeek V4 sparse MLA uses its own decode backends, selected via",
1623
+ "`--attention-backend=<BACKEND>` (e.g., `FLASHMLA_SPARSE_DSV4`,",
1624
+ "`FLASHINFER_MLA_SPARSE_DSV4`). They share the V4 sparse-index",
1625
+ "pipeline (compressor + SWA + indexer, 256-token blocks, head 512);",
1626
+ "default on NVIDIA is `FLASHMLA_SPARSE_DSV4`.",
1627
+ "",
1628
+ ]
1629
+ )
1630
+ lines.extend(_render_table(columns, v4_decode_backends))
1631
+
1632
lines.append("")
1633
return "\n".join(lines)
1634
@@ -1651,9 +1669,15 @@ def generate_docs() -> str:
1669
if fi_features:
1670
all_backends = _expand_flashinfer_variants(all_backends, fi_features)
1671
1654
- # Split into MLA and non-MLA
1655
- mla_backends = [b for b in all_backends if b["is_mla"]]
1656
- non_mla_backends = [b for b in all_backends if not b["is_mla"]]
1672
+ # DeepSeek V4 (*_DSV4) decode backends get their own subsection rather than
1673
+ # mixing into the main MLA / standard tables (the ROCm V4 backend isn't
1674
+ # flagged is_mla by the AST heuristic, so filter purely on the name).
1675
+ def _is_v4(b: dict[str, Any]) -> bool:
1676
+ return b["name"].endswith("_DSV4")
1677
+
1678
+ v4_decode_backends = [b for b in all_backends if _is_v4(b)]
1679
+ mla_backends = [b for b in all_backends if b["is_mla"] and not _is_v4(b)]
1680
+ non_mla_backends = [b for b in all_backends if not b["is_mla"] and not _is_v4(b)]
1681
1682
# Generate documentation
1683
script_path = "tools/pre_commit/generate_attention_backend_docs.py"
@@ -1703,7 +1727,9 @@ def generate_docs() -> str:
1727
doc_lines.append("\n>\n".join(footnotes) + "\n")
1728
1729
# Add MLA section with prefill and decode backends
1706
- doc_lines.append(generate_mla_section(mla_prefill_backends, mla_backends))
1730
+ doc_lines.append(
1731
+ generate_mla_section(mla_prefill_backends, mla_backends, v4_decode_backends)
1732
+ )
1733
1734
return "\n".join(doc_lines)
1735
vllm/models/deepseek_v4/amd/rocm.py
+1
-1
@@ -576,7 +576,7 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuild
576
class DeepseekV4ROCMAiterMLASparseBackend(DeepseekV4FlashMLASparseBackend):
577
@staticmethod
578
def get_name() -> str:
579
- return "ROCM_V4_FLASHMLA_SPARSE"
579
+ return "ROCM_FLASHMLA_SPARSE_DSV4"
580
581
@staticmethod
582
def get_builder_cls() -> type["DeepseekV4ROCMAiterMLASparseMetadataBuilder"]:
vllm/models/deepseek_v4/attention.py
+141
-42
@@ -55,9 +55,6 @@ from vllm.utils.multi_stream_utils import (
55
maybe_execute_in_parallel,
56
)
57
from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
58
-from vllm.v1.attention.backends.mla.flashmla_sparse import (
59
- FlashMLASparseBackend,
60
-)
58
from vllm.v1.attention.backends.mla.indexer import (
59
DeepseekV4IndexerBackend,
60
get_max_prefill_buffer_size,
@@ -73,21 +70,82 @@ if TYPE_CHECKING:
70
logger = init_logger(__name__)
71
72
76
-def _select_v4_sparse_impl() -> "type[DeepseekV4SparseMLAAttentionImpl]":
77
- """Pick the platform-specific V4 sparse MLA impl class. Sole platform check."""
73
+def _resolve_dsv4_backend(vllm_config: VllmConfig | None):
74
+ """Return the explicitly-requested DSv4 sparse backend enum, or None."""
75
+ if vllm_config is None:
76
+ return None
77
+ attn_config = getattr(vllm_config, "attention_config", None)
78
+ return getattr(attn_config, "backend", None) if attn_config is not None else None
79
+
80
+
81
+def _select_v4_sparse_impl(
82
+ vllm_config: VllmConfig | None = None,
83
+) -> "type[DeepseekV4SparseMLAAttentionImpl]":
84
+ """Pick the V4 sparse MLA impl class.
85
+
86
+ An explicit ``--attention-backend FLASHINFER_MLA_SPARSE_DSV4`` selects the
87
+ FlashInfer TRTLLM-gen path; otherwise the platform default (FlashMLA on
88
+ NVIDIA, ROCm Aiter on AMD) is used.
89
+ """
90
+ from vllm.v1.attention.backends.registry import AttentionBackendEnum
91
+
92
+ backend = _resolve_dsv4_backend(vllm_config)
93
+ if backend == AttentionBackendEnum.FLASHINFER_MLA_SPARSE_DSV4:
94
+ from vllm.models.deepseek_v4.nvidia.flashinfer_sparse import (
95
+ DeepseekV4FlashInferMLASparseImpl,
96
+ )
97
+
98
+ logger.info_once("Using FLASHINFER_MLA_SPARSE_DSV4 backend.")
99
+ return DeepseekV4FlashInferMLASparseImpl
100
if current_platform.is_rocm():
101
from vllm.models.deepseek_v4.amd.rocm import (
102
DeepseekV4ROCMAiterMLASparseImpl,
103
)
104
105
+ logger.info_once("Using ROCM_FLASHMLA_SPARSE_DSV4 backend.")
106
return DeepseekV4ROCMAiterMLASparseImpl
107
from vllm.models.deepseek_v4.nvidia.flashmla import (
108
DeepseekV4FlashMLASparseImpl,
109
)
110
111
+ logger.info_once("Using FLASHMLA_SPARSE_DSV4 backend.")
112
return DeepseekV4FlashMLASparseImpl
113
114
115
+def _resolve_dsv4_kv_cache_dtype(
116
+ backend,
117
+ kv_cache_dtype: str,
118
+ cache_config: CacheConfig | None,
119
+) -> tuple[str, torch.dtype]:
120
+ """Map ``(backend, --kv-cache-dtype)`` to ``(cache_dtype_str, torch_dtype)``.
121
+
122
+ FlashInfer V4 reads a contiguous 512-wide KV row (bf16 or per-tensor FP8
123
+ E4M3); FlashMLA V4 reads the legacy UE8M0 paged layout (uint8 /
124
+ ``fp8_ds_mla``). For FlashMLA the canonical ``fp8_ds_mla`` string is
125
+ written back onto ``cache_config`` so the page-size specs pick the 576B
126
+ layout.
127
+ """
128
+ from vllm.v1.attention.backends.registry import AttentionBackendEnum
129
+
130
+ if backend == AttentionBackendEnum.FLASHINFER_MLA_SPARSE_DSV4:
131
+ if kv_cache_dtype.startswith("fp8"):
132
+ return kv_cache_dtype, torch.float8_e4m3fn
133
+ # auto / bfloat16 -> contiguous BF16 cache.
134
+ return kv_cache_dtype, torch.bfloat16
135
+
136
+ # FlashMLA (and ROCm Aiter): legacy UE8M0 paged uint8 cache.
137
+ assert kv_cache_dtype.startswith("fp8"), (
138
+ f"DeepseekV4 FlashMLA sparse backend only supports fp8 kv-cache, "
139
+ f"got {kv_cache_dtype}"
140
+ )
141
+ if kv_cache_dtype != "fp8_ds_mla":
142
+ if cache_config is not None:
143
+ cache_config.cache_dtype = "fp8_ds_mla"
144
+ kv_cache_dtype = "fp8_ds_mla"
145
+ logger.info_once("Using DeepSeek's fp8_ds_mla KV cache format.")
146
+ return kv_cache_dtype, torch.uint8
147
+
148
+
149
class DeepseekV4MLA(nn.Module):
150
def __init__(
151
self,
@@ -194,10 +252,17 @@ class DeepseekV4MLA(nn.Module):
252
self.ln_events = [torch.cuda.Event() for _ in range(4)]
253
254
assert cache_config is not None, "DeepseekV4 attention requires cache_config"
255
+ # Resolve the SWA cache tensor dtype from the selected backend: FlashMLA
256
+ # uses the legacy UE8M0 paged uint8 layout; FlashInfer uses a contiguous
257
+ # bf16 / per-tensor fp8 row.
258
+ backend = _resolve_dsv4_backend(vllm_config)
259
+ _, swa_cache_torch_dtype = _resolve_dsv4_kv_cache_dtype(
260
+ backend, cache_config.cache_dtype, cache_config
261
+ )
262
self.swa_cache_layer = DeepseekV4SWACache(
263
head_dim=self.head_dim,
264
window_size=self.window_size,
200
- dtype=torch.uint8,
265
+ dtype=swa_cache_torch_dtype,
266
prefix=f"{prefix}.swa_cache",
267
cache_config=cache_config,
268
)
@@ -478,25 +543,66 @@ class DeepseekV4MLA(nn.Module):
543
assert swa_metadata is not None
544
545
swa_kv_cache = self.swa_cache_layer.kv_cache
481
- swa_kv_cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1)
546
+ # The fused insert ops require int64 position_ids; the runner's positions
547
+ # buffer is already int64, so no cast is needed.
548
+ assert positions.dtype == torch.int64
549
+ cos_sin_cache = self.rotary_emb.cos_sin_cache
550
+ cache_dtype = swa_kv_cache.dtype
551
483
- # Horizontally fused:
484
- # Q side: q_head_norm (per-head RMSNorm, no weight) + GPT-J RoPE,
485
- # with zero-fill for the padding head slots. The kernel
486
- # allocates and returns the padded q tensor.
487
- # KV side: GPT-J RoPE + UE8M0 FP8 quant + paged cache insert
552
# kv is unchanged; mla_attn reads kv solely via swa_kv_cache.
489
- return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
553
+ if cache_dtype == torch.uint8:
554
+ # Legacy FlashMLA UE8M0 paged path. Horizontally fused:
555
+ # Q side: per-head RMSNorm (no weight) + GPT-J RoPE, zero-filling
556
+ # the padding head slots; the kernel allocates and returns
557
+ # the padded q tensor.
558
+ # KV side: GPT-J RoPE + UE8M0 FP8 quant + paged cache insert.
559
+ swa_kv_cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1)
560
+ return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
561
+ q,
562
+ kv,
563
+ swa_kv_cache_2d,
564
+ swa_metadata.slot_mapping,
565
+ positions,
566
+ cos_sin_cache,
567
+ self.padded_heads,
568
+ self.eps,
569
+ swa_metadata.block_size,
570
+ )
571
+
572
+ # FlashInfer full-cache path: contiguous [num_blocks, block_size, 512]
573
+ # cache (no Q padding). bf16 rewrites q in place; per-tensor fp8 writes a
574
+ # separately-allocated fp8 q and quantizes the KV row.
575
+ block_size = swa_metadata.block_size
576
+ swa_kv_cache_3d = swa_kv_cache.view(-1, block_size, self.head_dim)
577
+ if cache_dtype == torch.bfloat16:
578
+ torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
579
+ q,
580
+ kv,
581
+ swa_kv_cache_3d,
582
+ swa_metadata.slot_mapping,
583
+ positions,
584
+ cos_sin_cache,
585
+ self.eps,
586
+ block_size,
587
+ )
588
+ return q
589
+
590
+ # per-tensor fp8 (torch.float8_e4m3fn)
591
+ q_fp8 = torch.empty_like(q, dtype=torch.float8_e4m3fn)
592
+ torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
593
q,
594
kv,
492
- swa_kv_cache_2d,
595
+ q_fp8,
596
+ swa_kv_cache_3d,
597
swa_metadata.slot_mapping,
494
- positions.to(torch.int64),
495
- self.rotary_emb.cos_sin_cache,
496
- self.padded_heads,
598
+ positions,
599
+ cos_sin_cache,
600
+ self.mla_attn._flashinfer_fp8_kv_scale,
601
+ self.mla_attn._flashinfer_fp8_q_scale_inv,
602
self.eps,
498
- swa_metadata.block_size,
603
+ block_size,
604
)
605
+ return q_fp8
606
607
608
class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase):
@@ -524,7 +630,8 @@ class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase):
630
**extra_impl_args,
631
) -> None:
632
super().__init__()
527
- self.impl_cls = _select_v4_sparse_impl()
633
+ vllm_config = get_current_vllm_config()
634
+ self.impl_cls = _select_v4_sparse_impl(vllm_config)
635
self.backend_cls = self.impl_cls.backend_cls
636
self.num_heads = num_heads
637
self.num_kv_heads = 1
@@ -556,34 +663,23 @@ class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase):
663
self.swa_cache_layer: DeepseekV4SWACache = swa_cache_layer
664
665
# Get vllm config for cache setup
559
- vllm_config = get_current_vllm_config()
666
self.max_num_batched_tokens = (
667
vllm_config.scheduler_config.max_num_batched_tokens
668
)
669
self.max_model_len = vllm_config.model_config.max_model_len
564
- # DeepseekV4 only supports fp8 kv-cache format for now.
565
- kv_cache_dtype = cache_config.cache_dtype if cache_config is not None else "fp8"
670
567
- assert kv_cache_dtype.startswith("fp8"), (
568
- f"DeepseekV4 only supports fp8 kv-cache format for now, "
569
- f"got {kv_cache_dtype}"
570
- )
571
- assert issubclass(self.get_attn_backend(), FlashMLASparseBackend), (
572
- "Only FlashMLA Sparse Attention backend is supported for DeepseekV4 for now"
671
+ # Resolve the kv-cache dtype from the selected backend. FlashMLA uses
672
+ # the legacy UE8M0 paged uint8 (fp8_ds_mla) layout; FlashInfer uses a
673
+ # contiguous bf16 / per-tensor fp8 row.
674
+ backend = _resolve_dsv4_backend(vllm_config)
675
+ kv_cache_dtype = cache_config.cache_dtype if cache_config is not None else "fp8"
676
+ self.kv_cache_dtype, self.kv_cache_torch_dtype = _resolve_dsv4_kv_cache_dtype(
677
+ backend, kv_cache_dtype, cache_config
678
)
574
- # FlashMLA Sparse Attention fp8 backend uses "fp8_ds_mla" kv-cache format
575
- # Automatically convert fp8 kv-cache format to "fp8_ds_mla"
576
- if (
577
- issubclass(self.get_attn_backend(), FlashMLASparseBackend)
578
- and kv_cache_dtype.startswith("fp8")
579
- and kv_cache_dtype != "fp8_ds_mla"
580
- ):
581
- assert cache_config is not None
582
- cache_config.cache_dtype = "fp8_ds_mla"
583
- kv_cache_dtype = "fp8_ds_mla"
584
- logger.info_once("Using DeepSeek's fp8_ds_mla KV cache format.")
679
586
- self.kv_cache_dtype = kv_cache_dtype
680
+ # Per-impl layer buffers (e.g. FlashInfer FP8 scale buffers). No-op for
681
+ # the FlashMLA / ROCm impls.
682
+ self.impl_cls.init_layer_buffers(self)
683
684
# Register with compilation context for metadata lookup
685
compilation_config = vllm_config.compilation_config
@@ -602,14 +698,17 @@ class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase):
698
self.compress_ratio <= 1
699
): # SWA part. Allocated separately as DeepseekV4SWACache.
700
return None
701
+ # FlashMLA uses the UE8M0 paged uint8 layout (576B aligned); FlashInfer
702
+ # uses a contiguous bf16 / per-tensor fp8 cache with no extra alignment.
703
+ is_flashmla = self.kv_cache_dtype == "fp8_ds_mla"
704
return MLAAttentionSpec(
705
block_size=vllm_config.cache_config.block_size,
706
num_kv_heads=1,
707
head_size=self.head_dim,
609
- dtype=torch.uint8,
708
+ dtype=torch.uint8 if is_flashmla else self.kv_cache_torch_dtype,
709
compress_ratio=self.compress_ratio,
710
cache_dtype_str=self.kv_cache_dtype,
612
- alignment=576, # NOTE: FlashMLA requires 576B alignment
711
+ alignment=576 if is_flashmla else None, # FlashMLA needs 576B
712
model_version="deepseek_v4",
713
)
714
vllm/models/deepseek_v4/common/ops/__init__.py
+2
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from .cache_utils import (
5
+ build_flashinfer_mixed_sparse_indices,
6
combine_topk_swa_indices,
7
compute_global_topk_indices_and_lens,
8
dequantize_and_gather_k_cache,
@@ -15,6 +16,7 @@ from .save_partial_states import save_partial_states
16
17
__all__ = [
18
"MXFP4_BLOCK_SIZE",
19
+ "build_flashinfer_mixed_sparse_indices",
20
"combine_topk_swa_indices",
21
"compute_global_topk_indices_and_lens",
22
"dequantize_and_gather_k_cache",
vllm/models/deepseek_v4/common/ops/cache_utils.py
+305
@@ -592,3 +592,308 @@ def _combine_topk_swa_indices_kernel(
592
593
combined_len = topk_len + swa_len
594
tl.store(combined_lens_ptr + token_idx, combined_len)
595
+
596
+
597
+def build_flashinfer_mixed_sparse_indices(
598
+ decode_swa_indices: torch.Tensor,
599
+ decode_compressed_indices: torch.Tensor | None,
600
+ decode_compressed_topk_lens: torch.Tensor | None,
601
+ prefill_topk_indices: torch.Tensor,
602
+ query_start_loc: torch.Tensor,
603
+ seq_lens: torch.Tensor,
604
+ token_to_req_indices: torch.Tensor,
605
+ swa_block_table: torch.Tensor,
606
+ swa_block_size: int,
607
+ compressed_block_table: torch.Tensor | None,
608
+ compressed_block_size: int,
609
+ window_size: int,
610
+ compress_ratio: int,
611
+ topk: int,
612
+ decode_compressed_indices_are_local: bool = False,
613
+ decode_is_valid_token: torch.Tensor | None = None,
614
+) -> tuple[torch.Tensor, torch.Tensor]:
615
+ """Build the FlashInfer DSV4 sparse-index matrix for decode-first batches.
616
+
617
+ Produces ``sparse_indices`` of shape ``[num_tokens, window_size +
618
+ padded_topk]`` (the first ``window_size`` columns are SWA slot ids, the rest
619
+ are compressed/top-k slot ids) and ``sparse_topk_lens`` (active length per
620
+ token). Decode tokens read precomputed SWA/compressed indices; prefill tokens
621
+ derive their SWA window from the position and translate local compressed
622
+ indices to global slots via the block tables.
623
+ """
624
+ assert decode_swa_indices.dtype == torch.int32
625
+ assert decode_swa_indices.dim() == 2
626
+ assert decode_swa_indices.shape[-1] == window_size
627
+ if decode_compressed_topk_lens is not None:
628
+ assert decode_compressed_topk_lens.dtype == torch.int32
629
+ assert prefill_topk_indices.dtype == torch.int32
630
+ assert prefill_topk_indices.dim() == 2
631
+ assert query_start_loc.dtype == torch.int32
632
+ assert seq_lens.dtype == torch.int32
633
+ assert token_to_req_indices.dtype == torch.int32
634
+ assert swa_block_table.dtype == torch.int32
635
+
636
+ num_decode_tokens = decode_swa_indices.shape[0]
637
+ num_prefill_tokens = prefill_topk_indices.shape[0]
638
+ num_tokens = num_decode_tokens + num_prefill_tokens
639
+ assert token_to_req_indices.shape[0] >= num_tokens
640
+ if decode_compressed_topk_lens is not None:
641
+ assert decode_compressed_topk_lens.shape[0] >= num_decode_tokens
642
+
643
+ decode_compressed_topk = 0
644
+ if decode_compressed_indices is None:
645
+ decode_compressed_indices = prefill_topk_indices
646
+ else:
647
+ assert decode_compressed_indices.dtype == torch.int32
648
+ assert decode_compressed_indices.dim() == 2
649
+ assert decode_compressed_indices.shape[0] == num_decode_tokens
650
+ decode_compressed_topk = decode_compressed_indices.shape[-1]
651
+ if decode_compressed_topk > 0 and decode_compressed_indices_are_local:
652
+ assert decode_is_valid_token is not None
653
+ assert decode_is_valid_token.dtype == torch.bool
654
+ assert decode_is_valid_token.shape[0] >= num_decode_tokens
655
+ else:
656
+ decode_is_valid_token = token_to_req_indices
657
+
658
+ if compressed_block_table is None:
659
+ compressed_block_table = swa_block_table
660
+ assert compressed_block_table.dtype == torch.int32
661
+ has_decode_compressed_lens = decode_compressed_topk_lens is not None
662
+ if decode_compressed_topk_lens is None:
663
+ decode_compressed_topk_lens = token_to_req_indices
664
+
665
+ # The FlashInfer TRTLLM-gen sparse-MLA kernels require every per-token topk
666
+ # index row to start on a 16-byte boundary: the kernel loads the compressed
667
+ # indices with 128-bit (16-byte) vectorized loads, so a misaligned row would
668
+ # fault or read across rows. 16 bytes = 4 int32 indices, so round the topk
669
+ # width (and hence the row stride, since the SWA columns are fixed-width) up
670
+ # to a multiple of 4. The extra columns are filled with -1 (invalid) and bounded
671
+ # by ``sparse_topk_lens``, so padding never changes the attention result.
672
+ padded_topk = max(topk, decode_compressed_topk)
673
+ padded_topk = (padded_topk + 3) // 4 * 4
674
+ sparse_indices = torch.empty(
675
+ (num_tokens, window_size + padded_topk),
676
+ dtype=torch.int32,
677
+ device=decode_swa_indices.device,
678
+ )
679
+ sparse_topk_lens = torch.empty(
680
+ num_tokens, dtype=torch.int32, device=decode_swa_indices.device
681
+ )
682
+ if num_tokens == 0:
683
+ return sparse_indices, sparse_topk_lens
684
+
685
+ window_block_size = triton.next_power_of_2(max(window_size, 1))
686
+ topk_block_size = triton.next_power_of_2(max(padded_topk, 1))
687
+ max_block_size = max(window_block_size, topk_block_size)
688
+ num_warps = 4 if max_block_size >= 256 else 1
689
+
690
+ _build_flashinfer_mixed_sparse_indices_kernel[(num_tokens,)](
691
+ sparse_indices,
692
+ sparse_indices.stride(0),
693
+ sparse_topk_lens,
694
+ decode_swa_indices,
695
+ decode_swa_indices.stride(0),
696
+ decode_compressed_indices,
697
+ decode_compressed_indices.stride(0),
698
+ decode_compressed_topk_lens,
699
+ decode_is_valid_token,
700
+ prefill_topk_indices,
701
+ prefill_topk_indices.stride(0),
702
+ query_start_loc,
703
+ seq_lens,
704
+ token_to_req_indices,
705
+ swa_block_table,
706
+ swa_block_table.stride(0),
707
+ swa_block_size,
708
+ compressed_block_table,
709
+ compressed_block_table.stride(0),
710
+ compressed_block_size,
711
+ NUM_DECODE_TOKENS=num_decode_tokens,
712
+ WINDOW_SIZE=window_size,
713
+ COMPRESS_RATIO=compress_ratio,
714
+ TOP_K=topk,
715
+ PADDED_TOP_K=padded_topk,
716
+ PREFILL_TOPK_STRIDE=prefill_topk_indices.shape[-1],
717
+ DECODE_COMPRESSED_TOPK=decode_compressed_topk,
718
+ DECODE_COMPRESSED_INDICES_ARE_LOCAL=decode_compressed_indices_are_local,
719
+ HAS_DECODE_COMPRESSED_LENS=has_decode_compressed_lens,
720
+ WINDOW_BLOCK_SIZE=window_block_size,
721
+ TOPK_BLOCK_SIZE=topk_block_size,
722
+ num_warps=num_warps,
723
+ )
724
+ return sparse_indices, sparse_topk_lens
725
+
726
+
727
+@triton.jit(
728
+ do_not_specialize=[
729
+ "sparse_indices_stride",
730
+ "decode_swa_stride",
731
+ "decode_compressed_stride",
732
+ "prefill_topk_stride",
733
+ "swa_block_table_stride",
734
+ "swa_block_size",
735
+ "compressed_block_table_stride",
736
+ "compressed_block_size",
737
+ "NUM_DECODE_TOKENS",
738
+ "PREFILL_TOPK_STRIDE",
739
+ ]
740
+)
741
+def _build_flashinfer_mixed_sparse_indices_kernel(
742
+ sparse_indices_ptr,
743
+ sparse_indices_stride,
744
+ sparse_topk_lens_ptr,
745
+ decode_swa_indices_ptr,
746
+ decode_swa_stride,
747
+ decode_compressed_indices_ptr,
748
+ decode_compressed_stride,
749
+ decode_compressed_topk_lens_ptr,
750
+ decode_is_valid_token_ptr,
751
+ prefill_topk_indices_ptr,
752
+ prefill_topk_stride,
753
+ query_start_loc_ptr,
754
+ seq_lens_ptr,
755
+ token_to_req_indices_ptr,
756
+ swa_block_table_ptr,
757
+ swa_block_table_stride,
758
+ swa_block_size,
759
+ compressed_block_table_ptr,
760
+ compressed_block_table_stride,
761
+ compressed_block_size,
762
+ NUM_DECODE_TOKENS,
763
+ WINDOW_SIZE: tl.constexpr,
764
+ COMPRESS_RATIO: tl.constexpr,
765
+ TOP_K: tl.constexpr,
766
+ PADDED_TOP_K: tl.constexpr,
767
+ PREFILL_TOPK_STRIDE,
768
+ DECODE_COMPRESSED_TOPK: tl.constexpr,
769
+ DECODE_COMPRESSED_INDICES_ARE_LOCAL: tl.constexpr,
770
+ HAS_DECODE_COMPRESSED_LENS: tl.constexpr,
771
+ WINDOW_BLOCK_SIZE: tl.constexpr,
772
+ TOPK_BLOCK_SIZE: tl.constexpr,
773
+):
774
+ token_idx = tl.program_id(0)
775
+
776
+ if token_idx < NUM_DECODE_TOKENS:
777
+ for i in range(0, WINDOW_SIZE, WINDOW_BLOCK_SIZE):
778
+ offset = i + tl.arange(0, WINDOW_BLOCK_SIZE)
779
+ mask = offset < WINDOW_SIZE
780
+ values = tl.load(
781
+ decode_swa_indices_ptr + token_idx * decode_swa_stride + offset,
782
+ mask=mask,
783
+ other=-1,
784
+ )
785
+ tl.store(
786
+ sparse_indices_ptr + token_idx * sparse_indices_stride + offset,
787
+ values,
788
+ mask=mask,
789
+ )
790
+
791
+ compressed_len = tl.zeros((), dtype=tl.int32)
792
+ for i in range(0, PADDED_TOP_K, TOPK_BLOCK_SIZE):
793
+ offset = i + tl.arange(0, TOPK_BLOCK_SIZE)
794
+ mask = offset < PADDED_TOP_K
795
+ values = tl.load(
796
+ decode_compressed_indices_ptr
797
+ + token_idx * decode_compressed_stride
798
+ + offset,
799
+ mask=offset < DECODE_COMPRESSED_TOPK,
800
+ other=-1,
801
+ )
802
+ if DECODE_COMPRESSED_INDICES_ARE_LOCAL:
803
+ token_valid = tl.load(decode_is_valid_token_ptr + token_idx)
804
+ is_valid = values >= 0
805
+ req_idx = tl.load(token_to_req_indices_ptr + token_idx)
806
+ block_indices = values // compressed_block_size
807
+ block_numbers = tl.load(
808
+ compressed_block_table_ptr
809
+ + req_idx * compressed_block_table_stride
810
+ + block_indices,
811
+ mask=mask & is_valid,
812
+ other=-1,
813
+ )
814
+ block_offsets = values % compressed_block_size
815
+ values = block_numbers * compressed_block_size + block_offsets
816
+ values = tl.where(is_valid, values, -1)
817
+ compressed_len += tl.sum((is_valid & token_valid).to(tl.int32), axis=0)
818
+ tl.store(
819
+ sparse_indices_ptr
820
+ + token_idx * sparse_indices_stride
821
+ + WINDOW_SIZE
822
+ + offset,
823
+ values,
824
+ mask=mask,
825
+ )
826
+
827
+ if DECODE_COMPRESSED_TOPK == 0:
828
+ compressed_len = tl.zeros((), dtype=tl.int32)
829
+ elif not DECODE_COMPRESSED_INDICES_ARE_LOCAL:
830
+ if HAS_DECODE_COMPRESSED_LENS:
831
+ compressed_len = tl.load(decode_compressed_topk_lens_ptr + token_idx)
832
+ else:
833
+ compressed_len = tl.full((), DECODE_COMPRESSED_TOPK, dtype=tl.int32)
834
+
835
+ tl.store(sparse_topk_lens_ptr + token_idx, WINDOW_SIZE + compressed_len)
836
+ return
837
+
838
+ prefill_idx = token_idx - NUM_DECODE_TOKENS
839
+ req_idx = tl.load(token_to_req_indices_ptr + token_idx)
840
+ query_start = tl.load(query_start_loc_ptr + req_idx)
841
+ query_end = tl.load(query_start_loc_ptr + req_idx + 1)
842
+ query_len = query_end - query_start
843
+ seq_len = tl.load(seq_lens_ptr + req_idx)
844
+ start_pos = seq_len - query_len
845
+ token_idx_in_query = token_idx - query_start
846
+ pos = start_pos + token_idx_in_query
847
+ swa_len = tl.minimum(pos + 1, WINDOW_SIZE)
848
+ swa_start_pos = pos - swa_len + 1
849
+ topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K)
850
+
851
+ for i in range(0, WINDOW_SIZE, WINDOW_BLOCK_SIZE):
852
+ offset = i + tl.arange(0, WINDOW_BLOCK_SIZE)
853
+ mask = offset < WINDOW_SIZE
854
+ pos_offset = swa_start_pos + offset
855
+ block_indices = pos_offset // swa_block_size
856
+ block_numbers = tl.load(
857
+ swa_block_table_ptr + req_idx * swa_block_table_stride + block_indices,
858
+ mask=mask & (offset < swa_len),
859
+ other=-1,
860
+ )
861
+ block_offsets = pos_offset % swa_block_size
862
+ slot_ids = block_numbers * swa_block_size + block_offsets
863
+ slot_ids = tl.where(offset < swa_len, slot_ids, -1)
864
+ tl.store(
865
+ sparse_indices_ptr + token_idx * sparse_indices_stride + offset,
866
+ slot_ids,
867
+ mask=mask,
868
+ )
869
+
870
+ for i in range(0, PADDED_TOP_K, TOPK_BLOCK_SIZE):
871
+ offset = i + tl.arange(0, TOPK_BLOCK_SIZE)
872
+ mask = offset < PADDED_TOP_K
873
+ local_idx = tl.load(
874
+ prefill_topk_indices_ptr + prefill_idx * prefill_topk_stride + offset,
875
+ mask=(offset < PREFILL_TOPK_STRIDE) & (offset < topk_len),
876
+ other=-1,
877
+ )
878
+ is_valid = local_idx >= 0
879
+ block_indices = local_idx // compressed_block_size
880
+ block_numbers = tl.load(
881
+ compressed_block_table_ptr
882
+ + req_idx * compressed_block_table_stride
883
+ + block_indices,
884
+ mask=mask & is_valid,
885
+ other=-1,
886
+ )
887
+ block_offsets = local_idx % compressed_block_size
888
+ slot_ids = block_numbers * compressed_block_size + block_offsets
889
+ slot_ids = tl.where((offset < topk_len) & is_valid, slot_ids, -1)
890
+ tl.store(
891
+ sparse_indices_ptr
892
+ + token_idx * sparse_indices_stride
893
+ + WINDOW_SIZE
894
+ + offset,
895
+ slot_ids,
896
+ mask=mask,
897
+ )
898
+
899
+ tl.store(sparse_topk_lens_ptr + token_idx, WINDOW_SIZE + topk_len)
vllm/models/deepseek_v4/compressor.py
+38
-19
@@ -155,13 +155,17 @@ class CompressorStateCache(torch.nn.Module, AttentionLayerBase):
155
raise ValueError(f"Invalid compress ratio: {compress_ratio}")
156
157
def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
158
+ # FlashMLA's UE8M0 paged layout needs 576B alignment; the FlashInfer
159
+ # full-cache path shares state pages with contiguous KV pages, so
160
+ # padding would break page matching.
161
+ is_flashmla = vllm_config.cache_config.cache_dtype == "fp8_ds_mla"
162
return SlidingWindowMLASpec( # only has one vector instead of K + V
163
block_size=self.block_size,
164
num_kv_heads=1,
165
head_size=self.state_dim,
166
dtype=self.dtype,
167
sliding_window=self.sliding_window,
164
- alignment=576, # NOTE: FlashMLA requires 576B alignment
168
+ alignment=576 if is_flashmla else None,
169
)
170
171
def forward(self): ...
@@ -333,26 +337,40 @@ class DeepseekCompressor(nn.Module):
337
# - position used: (positions // compress_ratio) * compress_ratio
338
cos_sin_cache = rotary_emb.cos_sin_cache
339
k_cache_metadata = cast(Any, attn_metadata[self.k_cache_prefix])
336
- kv_cache = self._static_forward_context[self.k_cache_prefix].kv_cache
337
-
338
- if current_platform.is_cuda():
339
- # NVIDIA GPUs.
340
- if self.head_dim == 512:
341
- from .nvidia.ops.sparse_attn_compress_cutedsl import (
342
- compress_norm_rope_store_cutedsl,
343
- )
344
-
345
- # Main compressor path.
346
- # Use a cutedsl kernel for better performance.
347
- compress_norm_rope_store_fn = compress_norm_rope_store_cutedsl
348
- else:
349
- # Indexer path (head_dim == 128).
350
- # Use a triton kernel.
351
- compress_norm_rope_store_fn = compress_norm_rope_store_triton
340
+ k_cache_layer = self._static_forward_context[self.k_cache_prefix]
341
+ kv_cache = k_cache_layer.kv_cache
342
+
343
+ # FlashInfer V4 reads a contiguous bf16 / per-tensor fp8 cache row; the
344
+ # legacy FlashMLA path uses the UE8M0 paged uint8 layout.
345
+ store_full_kv = self.head_dim == 512 and kv_cache.dtype != torch.uint8
346
+ store_full_fp8 = kv_cache.dtype == torch.float8_e4m3fn
347
+ fp8_scale = (
348
+ getattr(k_cache_layer, "_flashinfer_fp8_kv_scale", None)
349
+ if store_full_fp8
350
+ else None
351
+ )
352
+
353
+ # cutedsl (head=512) accepts the full-cache flags; triton (indexer/AMD)
354
+ # does not, so the two callables have different signatures.
355
+ compress_norm_rope_store_fn: Any
356
+ if current_platform.is_cuda() and self.head_dim == 512:
357
+ from .nvidia.ops.sparse_attn_compress_cutedsl import (
358
+ compress_norm_rope_store_cutedsl,
359
+ )
360
+
361
+ # head=512 on CUDA always uses cutedsl, for both the legacy UE8M0
362
+ # layout and the FlashInfer full-cache layout. The full-cache flags
363
+ # are consumed only here.
364
+ compress_norm_rope_store_fn = compress_norm_rope_store_cutedsl
365
+ extra_kwargs: dict[str, Any] = dict(
366
+ store_full_kv=store_full_kv,
367
+ store_full_fp8=store_full_fp8,
368
+ fp8_scale=fp8_scale,
369
+ )
370
else:
353
- # AMD GPUs.
354
- # Always use a triton kernel.
371
+ # Indexer path (head_dim == 128) or AMD: triton, legacy UE8M0 only.
372
compress_norm_rope_store_fn = compress_norm_rope_store_triton
373
+ extra_kwargs = {}
374
375
compress_norm_rope_store_fn(
376
state_cache=state_cache,
@@ -377,4 +395,5 @@ class DeepseekCompressor(nn.Module):
395
quant_block=self._quant_block,
396
token_stride=self._token_stride,
397
scale_dim=self._scale_dim,
398
+ **extra_kwargs,
399
)
vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py
new
+407
@@ -0,0 +1,407 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""DeepSeek V4 FlashInfer TRTLLM-gen sparse MLA backend.
4
+
5
+Uses FlashInfer's public ``trtllm_batch_decode_sparse_mla_dsv4`` launcher with a
6
+contiguous bf16 / per-tensor FP8 KV cache. Shares the V4 sparse-index pipeline
7
+(SWA cache + compressor + indexer, 256-token blocks, head_size 512) with the
8
+FlashMLA V4 backend; only the attention forward differs.
9
+"""
10
+
11
+from typing import TYPE_CHECKING, cast
12
+
13
+import torch
14
+
15
+from vllm.forward_context import get_forward_context
16
+from vllm.models.deepseek_v4.common.ops import (
17
+ build_flashinfer_mixed_sparse_indices,
18
+)
19
+from vllm.models.deepseek_v4.nvidia.flashmla import (
20
+ DeepseekV4FlashMLASparseBackend,
21
+ DeepseekV4SparseMLAAttentionImpl,
22
+)
23
+from vllm.utils.flashinfer import flashinfer_trtllm_batch_decode_sparse_mla_dsv4
24
+from vllm.v1.attention.backends.mla.flashmla_sparse import FlashMLASparseMetadata
25
+
26
+if TYPE_CHECKING:
27
+ from vllm.models.deepseek_v4.attention import DeepseekV4MLAAttention
28
+ from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
29
+
30
+# 128 MB TRTLLM-gen workspace, allocated once per device and zero-initialized
31
+# (required for first use). Reused across all FlashInfer V4 layers.
32
+_FLASHINFER_DSV4_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024
33
+_flashinfer_dsv4_workspace_by_device: dict[torch.device, torch.Tensor] = {}
34
+
35
+
36
+def _get_flashinfer_dsv4_workspace(device: torch.device) -> torch.Tensor:
37
+ workspace = _flashinfer_dsv4_workspace_by_device.get(device)
38
+ if workspace is None:
39
+ workspace = torch.zeros(
40
+ _FLASHINFER_DSV4_WORKSPACE_BUFFER_SIZE,
41
+ dtype=torch.uint8,
42
+ device=device,
43
+ )
44
+ _flashinfer_dsv4_workspace_by_device[device] = workspace
45
+ return workspace
46
+
47
+
48
+class DeepseekV4FlashInferMLASparseBackend(DeepseekV4FlashMLASparseBackend):
49
+ """Shares the FlashMLA V4 metadata/cache pipeline; swaps the attention impl.
50
+
51
+ Inheriting from the FlashMLA V4 backend reuses its ``FlashMLASparseMetadata``
52
+ builder (which the V4 sparse-index pipeline needs — the V3.2 FlashInfer
53
+ builder lacks the ``c128a_*`` fields), 256-token blocks, head_size 512, and
54
+ the contiguous (num_blocks, block_size, 512) cache shape for non-``fp8_ds_mla``
55
+ dtypes.
56
+ """
57
+
58
+ @staticmethod
59
+ def get_name() -> str:
60
+ return "FLASHINFER_MLA_SPARSE_DSV4"
61
+
62
+ @staticmethod
63
+ def get_impl_cls() -> type["DeepseekV4FlashInferMLASparseImpl"]:
64
+ return DeepseekV4FlashInferMLASparseImpl
65
+
66
+
67
+class DeepseekV4FlashInferMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
68
+ """FlashInfer TRTLLM-gen sparse MLA implementation for DeepSeek V4."""
69
+
70
+ backend_cls = DeepseekV4FlashInferMLASparseBackend
71
+
72
+ @classmethod
73
+ def get_padded_num_q_heads(cls, num_heads: int) -> int:
74
+ # FP8 decode kernel only supports h_q = 64 or 128.
75
+ if num_heads > 128:
76
+ raise ValueError(
77
+ f"DeepseekV4 Flashinfer MLA Sparse does not support {num_heads} heads "
78
+ "(FP8 decode kernel requires h_q in {64, 128})."
79
+ )
80
+ return 64 if num_heads <= 64 else 128
81
+
82
+ @classmethod
83
+ def init_layer_buffers(cls, layer: "DeepseekV4MLAAttention") -> None:
84
+ # Per-tensor FP8 scale buffers + precomputed scalar BMM scales. Only the
85
+ # per-tensor FP8 cache path consumes these; bf16 reads ``layer.scale``.
86
+ if layer.kv_cache_torch_dtype != torch.float8_e4m3fn:
87
+ return
88
+ # TODO: load real per-tensor Q/KV scales from the checkpoint; unit
89
+ # scales until the scale tensor names are wired.
90
+ fp8_q_scale = 1.0
91
+ fp8_kv_scale = 1.0
92
+ layer.register_buffer(
93
+ "_flashinfer_fp8_q_scale",
94
+ torch.tensor([fp8_q_scale], dtype=torch.float32),
95
+ persistent=False,
96
+ )
97
+ layer.register_buffer(
98
+ "_flashinfer_fp8_q_scale_inv",
99
+ torch.tensor([1.0 / fp8_q_scale], dtype=torch.float32),
100
+ persistent=False,
101
+ )
102
+ layer.register_buffer(
103
+ "_flashinfer_fp8_kv_scale",
104
+ torch.tensor([fp8_kv_scale], dtype=torch.float32),
105
+ persistent=False,
106
+ )
107
+ # TRTLLM-gen takes scalar scale args on a distinct (correct) C++ path
108
+ # vs 1-elem tensors, so these are Python floats. bmm1 folds the softmax
109
+ # scale and the Q/KV per-tensor scales; bmm2 is the KV scale.
110
+ layer._flashinfer_fp8_bmm1_scale = layer.scale * fp8_q_scale * fp8_kv_scale
111
+ layer._flashinfer_fp8_bmm2_scale = fp8_kv_scale
112
+
113
+ @classmethod
114
+ def forward_mqa( # type: ignore[override]
115
+ cls,
116
+ layer: "DeepseekV4MLAAttention",
117
+ q: torch.Tensor,
118
+ kv: torch.Tensor,
119
+ positions: torch.Tensor,
120
+ output: torch.Tensor,
121
+ ) -> None:
122
+ # The TRTLLM-gen kernel requires h_q in {64, 128}, so the output buffer
123
+ # is allocated at the padded head count while q arrives at the local
124
+ # head count; _forward pads q to match before the launcher.
125
+ assert output.shape[0] == q.shape[0] and output.shape[-1] == q.shape[-1], (
126
+ f"output buffer shape {output.shape} incompatible with q shape {q.shape}"
127
+ )
128
+ assert output.shape[1] >= q.shape[1], (
129
+ f"output heads {output.shape[1]} must be >= q heads {q.shape[1]}"
130
+ )
131
+ # Per-tensor FP8 q produces a bf16 attention output.
132
+ expected_output_dtype = (
133
+ torch.bfloat16 if q.dtype == torch.float8_e4m3fn else q.dtype
134
+ )
135
+ assert output.dtype == expected_output_dtype, (
136
+ f"output dtype {output.dtype} must match expected {expected_output_dtype} "
137
+ f"for q dtype {q.dtype}"
138
+ )
139
+
140
+ forward_context = get_forward_context()
141
+ attn_metadata = forward_context.attn_metadata
142
+ if attn_metadata is None:
143
+ # Warmup dummy run: FlashInfer reads the cache directly and lazily
144
+ # allocates its workspace, so nothing to reserve here.
145
+ output.zero_()
146
+ return
147
+
148
+ assert isinstance(attn_metadata, dict)
149
+ flashmla_metadata = cast(
150
+ FlashMLASparseMetadata | None, attn_metadata.get(layer.prefix)
151
+ )
152
+ swa_metadata = cast(
153
+ "DeepseekSparseSWAMetadata | None",
154
+ attn_metadata.get(layer.swa_cache_layer.prefix),
155
+ )
156
+ assert swa_metadata is not None
157
+
158
+ swa_only = layer.compress_ratio <= 1
159
+ # SWA-only layers don't allocate their own compressed KV cache.
160
+ self_kv_cache = layer.kv_cache if not swa_only else None
161
+ swa_kv_cache = layer.swa_cache_layer.kv_cache
162
+
163
+ cls._forward(
164
+ layer=layer,
165
+ q=q,
166
+ kv_cache=self_kv_cache,
167
+ swa_k_cache=swa_kv_cache,
168
+ swa_metadata=swa_metadata,
169
+ attn_metadata=flashmla_metadata,
170
+ swa_only=swa_only,
171
+ output=output,
172
+ )
173
+
174
+ @classmethod
175
+ def _build_sparse_index_metadata(
176
+ cls,
177
+ layer: "DeepseekV4MLAAttention",
178
+ kv_cache: torch.Tensor | None,
179
+ swa_k_cache: torch.Tensor,
180
+ swa_metadata: "DeepseekSparseSWAMetadata",
181
+ attn_metadata: FlashMLASparseMetadata | None,
182
+ swa_only: bool,
183
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
184
+ """Build the combined sparse-index tensors for the mixed batch.
185
+
186
+ Returns ``(compressed_kv_cache, seq_lens, sparse_indices,
187
+ sparse_topk_lens)``.
188
+ """
189
+ num_decodes = swa_metadata.num_decodes
190
+ num_prefills = swa_metadata.num_prefills
191
+ num_decode_tokens = swa_metadata.num_decode_tokens
192
+ num_prefill_tokens = swa_metadata.num_prefill_tokens
193
+ num_reqs = num_decodes + num_prefills
194
+ num_tokens = num_decode_tokens + num_prefill_tokens
195
+
196
+ assert swa_metadata.seq_lens is not None
197
+ assert swa_metadata.query_start_loc is not None
198
+ assert swa_metadata.token_to_req_indices is not None
199
+ assert swa_metadata.decode_swa_indices is not None
200
+ assert swa_metadata.block_table is not None
201
+
202
+ decode_swa_indices = swa_metadata.decode_swa_indices.reshape(
203
+ num_decode_tokens, layer.window_size
204
+ )
205
+ decode_compressed_topk_lens = None
206
+ decode_compressed_indices_are_local = False
207
+ decode_is_valid_token = None
208
+
209
+ if swa_only:
210
+ assert layer.topk_indices_buffer is not None
211
+ compressed_kv_cache = swa_k_cache
212
+ decode_compressed_indices = None
213
+ prefill_topk_indices = layer.topk_indices_buffer[
214
+ num_decode_tokens:num_tokens, :0
215
+ ]
216
+ compressed_block_table = None
217
+ compressed_block_size = swa_metadata.block_size
218
+ top_k = 0
219
+ else:
220
+ assert kv_cache is not None
221
+ assert attn_metadata is not None
222
+ compressed_kv_cache = kv_cache
223
+ compressed_block_table = attn_metadata.block_table[:num_reqs]
224
+ compressed_block_size = attn_metadata.block_size // layer.compress_ratio
225
+
226
+ if layer.compress_ratio == 4:
227
+ assert layer.topk_indices_buffer is not None
228
+ if num_prefill_tokens > 0:
229
+ prefill_topk_indices = layer.topk_indices_buffer[
230
+ num_decode_tokens:num_tokens
231
+ ]
232
+ top_k = prefill_topk_indices.shape[-1]
233
+ else:
234
+ prefill_topk_indices = layer.topk_indices_buffer[:0, :0]
235
+ top_k = 0
236
+
237
+ decode_compressed_indices_are_local = True
238
+ assert swa_metadata.is_valid_token is not None
239
+ decode_is_valid_token = swa_metadata.is_valid_token[:num_decode_tokens]
240
+ if num_decode_tokens > 0:
241
+ decode_compressed_indices = layer.topk_indices_buffer[
242
+ :num_decode_tokens
243
+ ]
244
+ else:
245
+ # Keep the logical width aligned with the mixed-batch case so
246
+ # pure-prefill steps reuse the same Triton specialization.
247
+ decode_compressed_indices = prefill_topk_indices[:0]
248
+ else:
249
+ if num_prefill_tokens > 0:
250
+ assert attn_metadata.c128a_prefill_topk_indices is not None
251
+ prefill_topk_indices = attn_metadata.c128a_prefill_topk_indices
252
+ top_k = prefill_topk_indices.shape[-1]
253
+ else:
254
+ prefill_topk_indices = decode_swa_indices[:0, :0]
255
+ top_k = 0
256
+
257
+ if num_decode_tokens > 0:
258
+ assert attn_metadata.c128a_global_decode_topk_indices is not None
259
+ assert attn_metadata.c128a_decode_topk_lens is not None
260
+ decode_compressed_indices = (
261
+ attn_metadata.c128a_global_decode_topk_indices.view(
262
+ num_decode_tokens, -1
263
+ )
264
+ )
265
+ decode_compressed_topk_lens = attn_metadata.c128a_decode_topk_lens
266
+ if num_prefill_tokens == 0:
267
+ prefill_topk_indices = decode_compressed_indices[:0, :0]
268
+ else:
269
+ decode_compressed_indices = prefill_topk_indices[:0]
270
+ decode_compressed_topk_lens = swa_metadata.seq_lens[:0]
271
+
272
+ query_start_loc = swa_metadata.query_start_loc[: num_reqs + 1]
273
+ seq_lens = swa_metadata.seq_lens[:num_reqs]
274
+ assert seq_lens.dtype == torch.int32
275
+ sparse_indices, sparse_topk_lens = build_flashinfer_mixed_sparse_indices(
276
+ decode_swa_indices,
277
+ decode_compressed_indices,
278
+ decode_compressed_topk_lens,
279
+ prefill_topk_indices[:num_prefill_tokens],
280
+ query_start_loc,
281
+ seq_lens,
282
+ swa_metadata.token_to_req_indices[:num_tokens],
283
+ swa_metadata.block_table[:num_reqs],
284
+ swa_metadata.block_size,
285
+ compressed_block_table,
286
+ compressed_block_size,
287
+ layer.window_size,
288
+ layer.compress_ratio,
289
+ top_k,
290
+ decode_compressed_indices_are_local=decode_compressed_indices_are_local,
291
+ decode_is_valid_token=decode_is_valid_token,
292
+ )
293
+ return compressed_kv_cache, seq_lens, sparse_indices, sparse_topk_lens
294
+
295
+ @classmethod
296
+ def _forward(
297
+ cls,
298
+ layer: "DeepseekV4MLAAttention",
299
+ q: torch.Tensor,
300
+ kv_cache: torch.Tensor | None,
301
+ swa_k_cache: torch.Tensor,
302
+ swa_metadata: "DeepseekSparseSWAMetadata",
303
+ attn_metadata: FlashMLASparseMetadata | None,
304
+ swa_only: bool,
305
+ output: torch.Tensor,
306
+ ) -> None:
307
+ assert layer.kv_cache_torch_dtype in (torch.bfloat16, torch.float8_e4m3fn)
308
+ num_decodes = swa_metadata.num_decodes
309
+ num_prefills = swa_metadata.num_prefills
310
+ num_decode_tokens = swa_metadata.num_decode_tokens
311
+ num_prefill_tokens = swa_metadata.num_prefill_tokens
312
+ num_reqs = num_decodes + num_prefills
313
+ num_tokens = num_decode_tokens + num_prefill_tokens
314
+ if num_tokens == 0:
315
+ return
316
+
317
+ (
318
+ compressed_kv_cache,
319
+ seq_lens,
320
+ sparse_indices,
321
+ sparse_topk_lens,
322
+ ) = cls._build_sparse_index_metadata(
323
+ layer=layer,
324
+ kv_cache=kv_cache,
325
+ swa_k_cache=swa_k_cache,
326
+ swa_metadata=swa_metadata,
327
+ attn_metadata=attn_metadata,
328
+ swa_only=swa_only,
329
+ )
330
+
331
+ # CUDA graph execution can pad q/output past the scheduled token count;
332
+ # restrict to the real tokens (the launcher validates sparse indices).
333
+ query = q[:num_tokens]
334
+ output = output[:num_tokens]
335
+ bmm1_scale: float | torch.Tensor = layer.scale
336
+ bmm2_scale: float | torch.Tensor = 1.0
337
+ if layer.kv_cache_torch_dtype == torch.float8_e4m3fn:
338
+ assert query.dtype == torch.float8_e4m3fn
339
+ bmm1_scale = layer._flashinfer_fp8_bmm1_scale
340
+ bmm2_scale = layer._flashinfer_fp8_bmm2_scale
341
+ else:
342
+ assert query.dtype == torch.bfloat16
343
+ query = query.contiguous()
344
+
345
+ # The TRTLLM-gen sparse-MLA kernel requires h_q in {64, 128}; zero-pad
346
+ # the query heads to the allocated output head count. Padded heads attend
347
+ # to the shared KV and are sliced off downstream (output is padded too).
348
+ padded_heads = output.shape[1]
349
+ if query.shape[1] < padded_heads:
350
+ padded_query = query.new_zeros(
351
+ (query.shape[0], padded_heads, query.shape[2])
352
+ )
353
+ padded_query[:, : query.shape[1], :] = query
354
+ query = padded_query
355
+
356
+ workspace = _get_flashinfer_dsv4_workspace(q.device)
357
+ query_start_loc = swa_metadata.query_start_loc
358
+ query_start_loc_cpu = swa_metadata.query_start_loc_cpu
359
+ assert query_start_loc is not None and query_start_loc_cpu is not None
360
+
361
+ # Keep Perkz's two-call decode/prefill split: the TRTLLM-gen launcher is
362
+ # tuned for uniform-q batches, and collapsing the mixed batch into a
363
+ # single call is the suspected source of the prior IMA.
364
+ if num_decode_tokens > 0:
365
+ decode_cu = query_start_loc[: num_decodes + 1]
366
+ decode_cu_cpu = query_start_loc_cpu[: num_decodes + 1]
367
+ decode_lens_cpu = decode_cu_cpu[1:] - decode_cu_cpu[:-1]
368
+ flashinfer_trtllm_batch_decode_sparse_mla_dsv4(
369
+ query=query[:num_decode_tokens],
370
+ swa_kv_cache=swa_k_cache,
371
+ workspace_buffer=workspace,
372
+ sparse_indices=sparse_indices[:num_decode_tokens],
373
+ compressed_kv_cache=compressed_kv_cache,
374
+ sparse_topk_lens=sparse_topk_lens[:num_decode_tokens],
375
+ seq_lens=seq_lens[:num_decodes],
376
+ out=output[:num_decode_tokens],
377
+ bmm1_scale=bmm1_scale,
378
+ bmm2_scale=bmm2_scale,
379
+ sinks=layer.attn_sink,
380
+ cum_seq_lens_q=decode_cu,
381
+ max_q_len=int(decode_lens_cpu.max().item()),
382
+ )
383
+
384
+ if num_prefill_tokens > 0:
385
+ # The prefill query view re-anchors at offset 0, so rebase the
386
+ # cumulative query offsets to start at 0.
387
+ prefill_cu = (
388
+ query_start_loc[num_decodes : num_reqs + 1]
389
+ - query_start_loc[num_decodes]
390
+ )
391
+ prefill_cu_cpu = query_start_loc_cpu[num_decodes : num_reqs + 1]
392
+ prefill_lens_cpu = prefill_cu_cpu[1:] - prefill_cu_cpu[:-1]
393
+ flashinfer_trtllm_batch_decode_sparse_mla_dsv4(
394
+ query=query[num_decode_tokens:num_tokens],
395
+ swa_kv_cache=swa_k_cache,
396
+ workspace_buffer=workspace,
397
+ sparse_indices=sparse_indices[num_decode_tokens:num_tokens],
398
+ compressed_kv_cache=compressed_kv_cache,
399
+ sparse_topk_lens=sparse_topk_lens[num_decode_tokens:num_tokens],
400
+ seq_lens=seq_lens[num_decodes:num_reqs],
401
+ out=output[num_decode_tokens:num_tokens],
402
+ bmm1_scale=bmm1_scale,
403
+ bmm2_scale=bmm2_scale,
404
+ sinks=layer.attn_sink,
405
+ cum_seq_lens_q=prefill_cu,
406
+ max_q_len=int(prefill_lens_cpu.max().item()),
407
+ )
vllm/models/deepseek_v4/nvidia/flashmla.py
+10
-1
@@ -75,6 +75,15 @@ class DeepseekV4SparseMLAAttentionImpl(SparseMLAAttentionImpl[FlashMLASparseMeta
75
"""
76
raise NotImplementedError
77
78
+ @classmethod
79
+ def init_layer_buffers(cls, layer: "DeepseekV4MLAAttention") -> None:
80
+ """Register impl-specific buffers on the layer at construction.
81
+
82
+ No-op by default; FlashInfer overrides this to register its per-tensor
83
+ FP8 scale buffers.
84
+ """
85
+ return None
86
+
87
88
class DeepseekV4FlashMLASparseBackend(FlashMLASparseBackend):
89
@staticmethod
@@ -83,7 +92,7 @@ class DeepseekV4FlashMLASparseBackend(FlashMLASparseBackend):
92
93
@staticmethod
94
def get_name() -> str:
86
- return "V4_FLASHMLA_SPARSE"
95
+ return "FLASHMLA_SPARSE_DSV4"
96
97
@staticmethod
98
def get_impl_cls() -> type["DeepseekV4SparseMLAAttentionImpl"]:
vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py
+1102
-323
@@ -508,28 +508,34 @@ class SparseAttnCompressNormRopeStoreC4Kernel:
508
)
509
510
511
-class SparseAttnCompressC128Block8Kernel:
512
- head_tile = 64
513
- rows_per_warp = 16
514
- elems_per_lane = 2
515
- lanes_per_row = head_tile // elems_per_lane
516
- num_warps = 8
517
- stats_lane_stride = lanes_per_row + 1
518
- final_reduce_steps = 3
519
- final_reduce_initial_offset = 4
520
- tb_size = num_warps * 32
521
- compress_ratio = 128
522
- state_block_size = 8
523
- rcp_ln2 = 1.4426950408889634
524
-
511
+class SparseAttnCompressNormRopeStoreFullC4Kernel(
512
+ SparseAttnCompressNormRopeStoreC4Kernel
513
+):
514
def __init__(
515
self,
516
head_size: int,
517
state_width: int,
518
+ rope_head_dim: int,
519
+ fp8_max: float,
520
+ quant_block: int,
521
+ token_stride: int,
522
+ scale_dim: int,
523
+ compress_ratio: int,
524
+ overlap: bool,
525
+ store_full_fp8: bool = False,
526
):
530
- self.head_dim = head_size
531
- self.num_splits = head_size // self.head_tile
532
- self.state_width = state_width
527
+ super().__init__(
528
+ head_size,
529
+ state_width,
530
+ rope_head_dim,
531
+ fp8_max,
532
+ quant_block,
533
+ token_stride,
534
+ scale_dim,
535
+ compress_ratio,
536
+ overlap,
537
+ )
538
+ self.store_full_fp8 = store_full_fp8
539
540
@cute.jit
541
def __call__(
@@ -539,17 +545,31 @@ class SparseAttnCompressC128Block8Kernel:
545
positions: cute.Tensor,
546
slot_mapping: cute.Tensor,
547
block_table: cute.Tensor,
542
- compressed_kv: cute.Tensor,
548
+ block_size: Int64,
549
+ rms_norm_weight: cute.Tensor,
550
+ rms_norm_eps: Float32,
551
+ cos_sin_cache: cute.Tensor,
552
+ k_cache: cute.Tensor,
553
+ kv_slot_mapping: cute.Tensor,
554
+ kv_cache_block_size: Int64,
555
+ fp8_scale: cute.Tensor,
556
stream: CUstream,
557
):
545
- grid = (slot_mapping.shape[0] * self.num_splits, 1, 1)
558
+ grid = (slot_mapping.shape[0], 1, 1)
559
self.kernel(
560
state_cache,
561
token_to_req_indices,
562
positions,
563
slot_mapping,
564
block_table,
552
- compressed_kv,
565
+ block_size,
566
+ rms_norm_weight,
567
+ rms_norm_eps,
568
+ cos_sin_cache,
569
+ k_cache,
570
+ kv_slot_mapping,
571
+ kv_cache_block_size,
572
+ fp8_scale,
573
).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream)
574
575
@cute.kernel
@@ -560,95 +580,65 @@ class SparseAttnCompressC128Block8Kernel:
580
positions: cute.Tensor,
581
slot_mapping: cute.Tensor,
582
block_table: cute.Tensor,
563
- compressed_kv: cute.Tensor,
583
+ block_size: Int64,
584
+ rms_norm_weight: cute.Tensor,
585
+ rms_norm_eps: Float32,
586
+ cos_sin_cache: cute.Tensor,
587
+ k_cache: cute.Tensor,
588
+ kv_slot_mapping: cute.Tensor,
589
+ kv_cache_block_size: Int64,
590
+ fp8_scale: cute.Tensor,
591
):
565
- block_id, _, _ = cute.arch.block_idx()
592
+ token_idx, _, _ = cute.arch.block_idx()
593
tid, _, _ = cute.arch.thread_idx()
594
warp_id = cute.arch.make_warp_uniform(tid // 32)
595
lane_id = tid % 32
569
- col_group = lane_id % self.lanes_per_row
570
-
571
- token_idx = block_id // self.num_splits
572
- split_idx = block_id - token_idx * self.num_splits
573
- col_base = split_idx * self.head_tile + col_group * self.elems_per_lane
596
+ group_lane = lane_id % self.lanes_per_group
597
+ group_idx = warp_id * self.groups_per_warp + lane_id // self.lanes_per_group
598
+ elem_base = group_idx * self.quant_block + group_lane * self.elems_per_lane
599
575
- position = Int64(0)
576
- req_idx = Int32(0)
577
- slot_id = Int64(-1)
600
+ slot_id = slot_mapping[token_idx]
601
has_position = token_idx < positions.shape[0]
579
- has_req_idx = token_idx < token_to_req_indices.shape[0]
580
- if lane_id == 0:
581
- slot_id = slot_mapping[token_idx]
582
- if lane_id == 0 and has_position:
602
+ position = Int64(0)
603
+ if has_position:
604
position = positions[token_idx]
584
- if lane_id == 0 and has_req_idx:
585
- req_idx = token_to_req_indices[token_idx]
586
- slot_id = cute.arch.shuffle_sync(slot_id, offset=0)
587
- position = cute.arch.shuffle_sync(position, offset=0)
588
- req_idx = cute.arch.shuffle_sync(req_idx, offset=0)
605
boundary = has_position and (
606
(position + Int64(1)) % Int64(self.compress_ratio) == Int64(0)
607
)
592
- start = position - Int64(self.compress_ratio - 1)
593
- active = slot_id >= Int64(0) and has_req_idx and boundary
608
+ has_req_idx = token_idx < token_to_req_indices.shape[0]
609
+ has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0]
610
+ kv_slot_idx = Int64(-1)
611
+ if has_kv_slot_idx:
612
+ kv_slot_idx = kv_slot_mapping[token_idx]
613
+ active = (
614
+ slot_id >= Int64(0) and has_req_idx and boundary and kv_slot_idx >= Int64(0)
615
+ )
616
617
if active:
618
+ req_idx = token_to_req_indices[token_idx]
619
+ start = position - Int64(self.window - 1)
620
+
621
smem = cutlass.utils.SmemAllocator()
597
- s_max = smem.allocate_tensor(
598
- Float32,
599
- cute.make_layout(
600
- (
601
- self.num_warps,
602
- self.lanes_per_row,
603
- self.elems_per_lane,
604
- ),
605
- stride=(
606
- self.stats_lane_stride * self.elems_per_lane,
607
- self.elems_per_lane,
608
- 1,
609
- ),
610
- ),
611
- byte_alignment=4,
622
+ s_block_numbers = smem.allocate_tensor(
623
+ Int32, cute.make_layout((self.window,)), byte_alignment=4
624
)
613
- s_sum = smem.allocate_tensor(
614
- Float32,
615
- cute.make_layout(
616
- (
617
- self.num_warps,
618
- self.lanes_per_row,
619
- self.elems_per_lane,
620
- ),
621
- stride=(
622
- self.stats_lane_stride * self.elems_per_lane,
623
- self.elems_per_lane,
624
- 1,
625
- ),
626
- ),
627
- byte_alignment=4,
625
+ partial_sums = smem.allocate_tensor(
626
+ Float32, cute.make_layout((self.num_warps,)), byte_alignment=4
627
)
629
- s_product = smem.allocate_tensor(
630
- Float32,
631
- cute.make_layout(
632
- (
633
- self.num_warps,
634
- self.lanes_per_row,
635
- self.elems_per_lane,
636
- ),
637
- stride=(
638
- self.stats_lane_stride * self.elems_per_lane,
639
- self.elems_per_lane,
640
- 1,
641
- ),
642
- ),
643
- byte_alignment=4,
628
+ rrms_shared = smem.allocate_tensor(
629
+ Float32, cute.make_layout((1,)), byte_alignment=4
630
)
631
646
- row_layout = cute.make_layout(
647
- (self.rows_per_warp, self.elems_per_lane),
648
- stride=(self.elems_per_lane, 1),
649
- )
650
- kv_vals = cute.make_rmem_tensor(row_layout, Float32)
651
- score_vals = cute.make_rmem_tensor(row_layout, Float32)
632
+ for row in cutlass.range_constexpr(self.window):
633
+ pos = start + Int64(row)
634
+ if tid == row:
635
+ block_number_i32 = Int32(0)
636
+ if pos >= Int64(0):
637
+ block_index = pos // block_size
638
+ block_number_i32 = block_table[req_idx, block_index]
639
+ s_block_numbers[row] = block_number_i32
640
+ cute.arch.sync_threads()
641
+
642
local_max = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
643
local_sum = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
644
local_product = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
@@ -658,47 +648,477 @@ class SparseAttnCompressC128Block8Kernel:
648
local_sum[e] = Float32(0.0)
649
local_product[e] = Float32(0.0)
650
661
- first_block_index = start // Int64(self.state_block_size)
662
- warp_block_index = first_block_index + (warp_id * 2).to(Int64)
663
- block0_i32 = Int32(0)
664
- block1_i32 = Int32(0)
665
- if lane_id == 0:
666
- block0_i32 = block_table[req_idx, warp_block_index]
667
- block1_i32 = block_table[req_idx, warp_block_index + Int64(1)]
668
- block0_i32 = cute.arch.shuffle_sync(block0_i32, offset=0)
669
- block1_i32 = cute.arch.shuffle_sync(block1_i32, offset=0)
670
-
671
- cp_f32x2 = cute.make_copy_atom(
672
- cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=64
651
+ cp_f32x4 = cute.make_copy_atom(
652
+ cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128
653
)
674
- final_mask_and_clamp = const_expr(
675
- (cute.arch.WARP_SIZE - self.num_warps) << 8 | (cute.arch.WARP_SIZE - 1)
654
+ copy_layout = cute.make_layout(
655
+ (self.copy_chunks, self.copy_elems),
656
+ stride=(self.copy_elems, 1),
657
)
677
- col_tile = col_base.to(Int64) // Int64(self.elems_per_lane)
678
- score_col_tile = col_tile + Int64(self.state_width // self.elems_per_lane)
658
+ kv_vals = cute.make_rmem_tensor(copy_layout, Float32)
659
+ score_vals = cute.make_rmem_tensor(copy_layout, Float32)
660
680
- for i in cutlass.range_constexpr(self.rows_per_warp):
681
- block_number_i32 = block0_i32
682
- block_offset = Int64(i)
683
- if const_expr(i >= self.state_block_size):
684
- block_number_i32 = block1_i32
685
- block_offset = Int64(i - self.state_block_size)
686
- row_tensor = state_cache[block_number_i32.to(Int64), block_offset, None]
687
- kv_src = cute.local_tile(
688
- row_tensor,
689
- tiler=(self.elems_per_lane,),
690
- coord=(col_tile,),
691
- )
692
- score_src = cute.local_tile(
693
- row_tensor,
694
- tiler=(self.elems_per_lane,),
695
- coord=(score_col_tile,),
696
- )
697
- cute.copy(cp_f32x2, kv_src, kv_vals[i, None])
698
- cute.copy(cp_f32x2, score_src, score_vals[i, None])
661
+ for row in cutlass.range_constexpr(self.window):
662
+ pos = start + Int64(row)
663
+ if pos >= Int64(0):
664
+ block_index = pos // block_size
665
+ block_offset = pos - block_index * block_size
666
+ block_number = s_block_numbers[row].to(Int64)
667
+ head_offset = Int64((row // self.compress_ratio) * self.head_dim)
668
+ row_tensor = state_cache[block_number, block_offset, None]
669
+ for chunk in cutlass.range_constexpr(self.copy_chunks):
670
+ copy_elem = const_expr(chunk * self.copy_elems)
671
+ col_tile = (
672
+ head_offset + (elem_base + Int32(copy_elem)).to(Int64)
673
+ ) // Int64(self.copy_elems)
674
+ kv_src = cute.local_tile(
675
+ row_tensor,
676
+ tiler=(self.copy_elems,),
677
+ coord=(col_tile,),
678
+ )
679
+ score_src = cute.local_tile(
680
+ row_tensor,
681
+ tiler=(self.copy_elems,),
682
+ coord=(
683
+ col_tile + Int64(self.state_width // self.copy_elems),
684
+ ),
685
+ )
686
+ cute.copy(cp_f32x4, kv_src, kv_vals[chunk, None])
687
+ cute.copy(cp_f32x4, score_src, score_vals[chunk, None])
688
700
- for e in cutlass.range_constexpr(self.elems_per_lane):
701
- local_max[e] = cute.arch.fmax(local_max[e], score_vals[i, e])
689
+ for e in cutlass.range_constexpr(self.elems_per_lane):
690
+ chunk = const_expr(e // self.copy_elems)
691
+ copy_elem = const_expr(e % self.copy_elems)
692
+ score = score_vals[chunk, copy_elem]
693
+ kv = kv_vals[chunk, copy_elem]
694
+ new_max = cute.arch.fmax(local_max[e], score)
695
+ old_scale = cute.math.exp2(
696
+ (local_max[e] - new_max) * Float32(self.rcp_ln2),
697
+ fastmath=True,
698
+ )
699
+ new_scale = cute.math.exp2(
700
+ (score - new_max) * Float32(self.rcp_ln2),
701
+ fastmath=True,
702
+ )
703
+ local_sum[e] = local_sum[e] * old_scale + new_scale
704
+ local_product[e] = local_product[e] * old_scale + kv * new_scale
705
+ local_max[e] = new_max
706
+
707
+ x = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
708
+ local_sumsq = Float32(0.0)
709
+ for e in cutlass.range_constexpr(self.elems_per_lane):
710
+ x[e] = local_product[e] / local_sum[e]
711
+ local_sumsq += x[e] * x[e]
712
+
713
+ warp_sum = local_sumsq
714
+ for step in cutlass.range_constexpr(5):
715
+ offset = const_expr(16 >> step)
716
+ warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset)
717
+
718
+ if lane_id == 0:
719
+ partial_sums[warp_id] = warp_sum
720
+ cute.arch.sync_threads()
721
+ if tid == 0:
722
+ total = Float32(0.0)
723
+ for i in cutlass.range_constexpr(self.num_warps):
724
+ total += partial_sums[i]
725
+ rrms_shared[0] = cute.math.rsqrt(
726
+ total / Float32(self.head_dim) + rms_norm_eps, fastmath=True
727
+ )
728
+ cute.arch.sync_threads()
729
+
730
+ rrms = rrms_shared[0]
731
+ for e in cutlass.range_constexpr(self.elems_per_lane):
732
+ elem = elem_base + e
733
+ x[e] = x[e] * rrms * rms_norm_weight[elem].to(Float32)
734
+
735
+ page = kv_slot_idx // kv_cache_block_size
736
+ kv_offset = kv_slot_idx - page * kv_cache_block_size
737
+ value_base = page * k_cache.stride[0] + kv_offset * k_cache.stride[1]
738
+
739
+ if const_expr(self.store_full_fp8):
740
+ k_cache_u16 = cute.recast_tensor(k_cache, Uint16)
741
+ inv_fp8 = Float32(1.0) / fp8_scale[0]
742
+ if group_idx == self.nope_blocks:
743
+ compressed_pos = (position // Int64(self.compress_ratio)) * Int64(
744
+ self.compress_ratio
745
+ )
746
+ for pair in cutlass.range_constexpr(self.elems_per_lane // 2):
747
+ elem = const_expr(pair * 2)
748
+ pair_idx = (elem_base - self.nope_dim) // 2 + Int32(pair)
749
+ cos_v = cos_sin_cache[compressed_pos, pair_idx]
750
+ sin_v = cos_sin_cache[
751
+ compressed_pos, pair_idx + Int32(self.rope_dim // 2)
752
+ ]
753
+ real = x[elem] * cos_v - x[elem + 1] * sin_v
754
+ imag = x[elem] * sin_v + x[elem + 1] * cos_v
755
+ packed_bf16 = _fp32x2_to_bf16x2(real, imag)
756
+ b0, b1 = _bf16x2_to_fp32(packed_bf16)
757
+ y0 = cutlass.min(
758
+ cutlass.max(b0 * inv_fp8, Float32(-self.fp8_max)),
759
+ Float32(self.fp8_max),
760
+ )
761
+ y1 = cutlass.min(
762
+ cutlass.max(b1 * inv_fp8, Float32(-self.fp8_max)),
763
+ Float32(self.fp8_max),
764
+ )
765
+ packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1)
766
+ out_base = value_base + (elem_base + Int32(elem)).to(Int64)
767
+ k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8
768
+ else:
769
+ for pair in cutlass.range_constexpr(self.elems_per_lane // 2):
770
+ elem = const_expr(pair * 2)
771
+ packed_bf16 = _fp32x2_to_bf16x2(x[elem], x[elem + 1])
772
+ b0, b1 = _bf16x2_to_fp32(packed_bf16)
773
+ y0 = cutlass.min(
774
+ cutlass.max(b0 * inv_fp8, Float32(-self.fp8_max)),
775
+ Float32(self.fp8_max),
776
+ )
777
+ y1 = cutlass.min(
778
+ cutlass.max(b1 * inv_fp8, Float32(-self.fp8_max)),
779
+ Float32(self.fp8_max),
780
+ )
781
+ packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1)
782
+ out_base = value_base + (elem_base + Int32(elem)).to(Int64)
783
+ k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8
784
+ else:
785
+ k_cache_u32 = cute.recast_tensor(k_cache, Uint32)
786
+ if group_idx == self.nope_blocks:
787
+ compressed_pos = (position // Int64(self.compress_ratio)) * Int64(
788
+ self.compress_ratio
789
+ )
790
+ for pair in cutlass.range_constexpr(self.elems_per_lane // 2):
791
+ elem = const_expr(pair * 2)
792
+ pair_idx = (elem_base - self.nope_dim) // 2 + Int32(pair)
793
+ cos_v = cos_sin_cache[compressed_pos, pair_idx]
794
+ sin_v = cos_sin_cache[
795
+ compressed_pos, pair_idx + Int32(self.rope_dim // 2)
796
+ ]
797
+ real = x[elem] * cos_v - x[elem + 1] * sin_v
798
+ imag = x[elem] * sin_v + x[elem + 1] * cos_v
799
+ packed_bf16 = _fp32x2_to_bf16x2(real, imag)
800
+ out_base = value_base + ((elem_base + Int32(elem)) * 2).to(
801
+ Int64
802
+ )
803
+ k_cache_u32.iterator[out_base // Int64(4)] = packed_bf16
804
+ else:
805
+ for pair in cutlass.range_constexpr(self.elems_per_lane // 2):
806
+ elem = const_expr(pair * 2)
807
+ packed_bf16 = _fp32x2_to_bf16x2(x[elem], x[elem + 1])
808
+ out_base = value_base + ((elem_base + Int32(elem)) * 2).to(
809
+ Int64
810
+ )
811
+ k_cache_u32.iterator[out_base // Int64(4)] = packed_bf16
812
+
813
+ @cache
814
+ @staticmethod
815
+ def compile(
816
+ head_size: int = 512,
817
+ state_width: int = 1024,
818
+ rope_head_dim: int = 64,
819
+ fp8_max: float = 448.0,
820
+ quant_block: int = 64,
821
+ token_stride: int = 576,
822
+ scale_dim: int = 8,
823
+ kv_block_stride: int = 74752,
824
+ compress_ratio: int = 4,
825
+ overlap: bool = True,
826
+ store_full_fp8: bool = False,
827
+ norm_weight_dtype: type[cutlass.Numeric] = Float32,
828
+ ):
829
+ if compress_ratio != 4 or not overlap:
830
+ raise ValueError("CuTe DSL C4 fused sparse-attn requires C4 overlap.")
831
+ if head_size != 512:
832
+ raise ValueError(
833
+ "CuTe DSL C4 fused sparse-attn currently requires head_size=512."
834
+ )
835
+ if state_width != 2 * head_size:
836
+ raise ValueError(
837
+ "CuTe DSL C4 fused sparse-attn requires state_width=2*head_size."
838
+ )
839
+ if quant_block != 64:
840
+ raise ValueError(
841
+ "CuTe DSL C4 fused sparse-attn currently requires quant_block=64."
842
+ )
843
+ if rope_head_dim != 64:
844
+ raise ValueError(
845
+ "CuTe DSL C4 fused sparse-attn currently requires rope_head_dim=64."
846
+ )
847
+ num_positions = cute.sym_int()
848
+ num_slots = cute.sym_int()
849
+ num_req_indices = cute.sym_int()
850
+ num_kv_slots = cute.sym_int()
851
+ num_state_blocks = cute.sym_int()
852
+ num_kv_blocks = cute.sym_int()
853
+ state_cache_block_size = cute.sym_int()
854
+ block_table_width = cute.sym_int()
855
+ max_pos = cute.sym_int()
856
+ state_cache_width = state_width * 2
857
+
858
+ state_cache = cute.runtime.make_fake_tensor(
859
+ Float32,
860
+ (num_state_blocks, state_cache_block_size, state_cache_width),
861
+ stride=(
862
+ cute.sym_int64(divisibility=16),
863
+ cute.sym_int64(divisibility=16),
864
+ 1,
865
+ ),
866
+ assumed_align=16,
867
+ )
868
+ token_to_req_indices = make_fake_tensor(
869
+ Int32, (num_req_indices,), divisibility=4
870
+ )
871
+ positions = make_fake_tensor(Int64, (num_positions,), divisibility=8)
872
+ slot_mapping = make_fake_tensor(Int64, (num_slots,), divisibility=8)
873
+ block_table = make_fake_tensor(
874
+ Int32, (cute.sym_int(), block_table_width), divisibility=1
875
+ )
876
+ rms_norm_weight = make_fake_tensor(
877
+ norm_weight_dtype, (head_size,), divisibility=4
878
+ )
879
+ cos_sin_cache = cute.runtime.make_fake_tensor(
880
+ Float32,
881
+ (max_pos, rope_head_dim),
882
+ stride=(cute.sym_int64(divisibility=4), 1),
883
+ assumed_align=4,
884
+ )
885
+ k_cache = cute.runtime.make_fake_tensor(
886
+ Uint8,
887
+ (num_kv_blocks, cute.sym_int(), cute.sym_int()),
888
+ stride=(
889
+ cute.sym_int64(divisibility=16),
890
+ cute.sym_int64(divisibility=8),
891
+ 1,
892
+ ),
893
+ assumed_align=16,
894
+ )
895
+ kv_slot_mapping = make_fake_tensor(Int64, (num_kv_slots,), divisibility=8)
896
+ fp8_scale = make_fake_tensor(Float32, (1,), divisibility=1)
897
+
898
+ kernel = SparseAttnCompressNormRopeStoreFullC4Kernel(
899
+ head_size,
900
+ state_width,
901
+ rope_head_dim,
902
+ fp8_max,
903
+ quant_block,
904
+ token_stride,
905
+ scale_dim,
906
+ compress_ratio,
907
+ overlap,
908
+ store_full_fp8,
909
+ )
910
+ stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
911
+ return cute.compile(
912
+ kernel,
913
+ state_cache,
914
+ token_to_req_indices,
915
+ positions,
916
+ slot_mapping,
917
+ block_table,
918
+ Int64(0),
919
+ rms_norm_weight,
920
+ Float32(0.0),
921
+ cos_sin_cache,
922
+ k_cache,
923
+ kv_slot_mapping,
924
+ Int64(0),
925
+ fp8_scale,
926
+ stream,
927
+ options="--enable-tvm-ffi",
928
+ )
929
+
930
+
931
+class SparseAttnCompressC128Block8Kernel:
932
+ head_tile = 64
933
+ rows_per_warp = 16
934
+ elems_per_lane = 2
935
+ lanes_per_row = head_tile // elems_per_lane
936
+ num_warps = 8
937
+ stats_lane_stride = lanes_per_row + 1
938
+ final_reduce_steps = 3
939
+ final_reduce_initial_offset = 4
940
+ tb_size = num_warps * 32
941
+ compress_ratio = 128
942
+ state_block_size = 8
943
+ rcp_ln2 = 1.4426950408889634
944
+
945
+ def __init__(
946
+ self,
947
+ head_size: int,
948
+ state_width: int,
949
+ ):
950
+ self.head_dim = head_size
951
+ self.num_splits = head_size // self.head_tile
952
+ self.state_width = state_width
953
+
954
+ @cute.jit
955
+ def __call__(
956
+ self,
957
+ state_cache: cute.Tensor,
958
+ token_to_req_indices: cute.Tensor,
959
+ positions: cute.Tensor,
960
+ slot_mapping: cute.Tensor,
961
+ block_table: cute.Tensor,
962
+ compressed_kv: cute.Tensor,
963
+ stream: CUstream,
964
+ ):
965
+ grid = (slot_mapping.shape[0] * self.num_splits, 1, 1)
966
+ self.kernel(
967
+ state_cache,
968
+ token_to_req_indices,
969
+ positions,
970
+ slot_mapping,
971
+ block_table,
972
+ compressed_kv,
973
+ ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream)
974
+
975
+ @cute.kernel
976
+ def kernel(
977
+ self,
978
+ state_cache: cute.Tensor,
979
+ token_to_req_indices: cute.Tensor,
980
+ positions: cute.Tensor,
981
+ slot_mapping: cute.Tensor,
982
+ block_table: cute.Tensor,
983
+ compressed_kv: cute.Tensor,
984
+ ):
985
+ block_id, _, _ = cute.arch.block_idx()
986
+ tid, _, _ = cute.arch.thread_idx()
987
+ warp_id = cute.arch.make_warp_uniform(tid // 32)
988
+ lane_id = tid % 32
989
+ col_group = lane_id % self.lanes_per_row
990
+
991
+ token_idx = block_id // self.num_splits
992
+ split_idx = block_id - token_idx * self.num_splits
993
+ col_base = split_idx * self.head_tile + col_group * self.elems_per_lane
994
+
995
+ position = Int64(0)
996
+ req_idx = Int32(0)
997
+ slot_id = Int64(-1)
998
+ has_position = token_idx < positions.shape[0]
999
+ has_req_idx = token_idx < token_to_req_indices.shape[0]
1000
+ if lane_id == 0:
1001
+ slot_id = slot_mapping[token_idx]
1002
+ if lane_id == 0 and has_position:
1003
+ position = positions[token_idx]
1004
+ if lane_id == 0 and has_req_idx:
1005
+ req_idx = token_to_req_indices[token_idx]
1006
+ slot_id = cute.arch.shuffle_sync(slot_id, offset=0)
1007
+ position = cute.arch.shuffle_sync(position, offset=0)
1008
+ req_idx = cute.arch.shuffle_sync(req_idx, offset=0)
1009
+ boundary = has_position and (
1010
+ (position + Int64(1)) % Int64(self.compress_ratio) == Int64(0)
1011
+ )
1012
+ start = position - Int64(self.compress_ratio - 1)
1013
+ active = slot_id >= Int64(0) and has_req_idx and boundary
1014
+
1015
+ if active:
1016
+ smem = cutlass.utils.SmemAllocator()
1017
+ s_max = smem.allocate_tensor(
1018
+ Float32,
1019
+ cute.make_layout(
1020
+ (
1021
+ self.num_warps,
1022
+ self.lanes_per_row,
1023
+ self.elems_per_lane,
1024
+ ),
1025
+ stride=(
1026
+ self.stats_lane_stride * self.elems_per_lane,
1027
+ self.elems_per_lane,
1028
+ 1,
1029
+ ),
1030
+ ),
1031
+ byte_alignment=4,
1032
+ )
1033
+ s_sum = smem.allocate_tensor(
1034
+ Float32,
1035
+ cute.make_layout(
1036
+ (
1037
+ self.num_warps,
1038
+ self.lanes_per_row,
1039
+ self.elems_per_lane,
1040
+ ),
1041
+ stride=(
1042
+ self.stats_lane_stride * self.elems_per_lane,
1043
+ self.elems_per_lane,
1044
+ 1,
1045
+ ),
1046
+ ),
1047
+ byte_alignment=4,
1048
+ )
1049
+ s_product = smem.allocate_tensor(
1050
+ Float32,
1051
+ cute.make_layout(
1052
+ (
1053
+ self.num_warps,
1054
+ self.lanes_per_row,
1055
+ self.elems_per_lane,
1056
+ ),
1057
+ stride=(
1058
+ self.stats_lane_stride * self.elems_per_lane,
1059
+ self.elems_per_lane,
1060
+ 1,
1061
+ ),
1062
+ ),
1063
+ byte_alignment=4,
1064
+ )
1065
+
1066
+ row_layout = cute.make_layout(
1067
+ (self.rows_per_warp, self.elems_per_lane),
1068
+ stride=(self.elems_per_lane, 1),
1069
+ )
1070
+ kv_vals = cute.make_rmem_tensor(row_layout, Float32)
1071
+ score_vals = cute.make_rmem_tensor(row_layout, Float32)
1072
+ local_max = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
1073
+ local_sum = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
1074
+ local_product = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
1075
+
1076
+ for e in cutlass.range_constexpr(self.elems_per_lane):
1077
+ local_max[e] = -Float32.inf
1078
+ local_sum[e] = Float32(0.0)
1079
+ local_product[e] = Float32(0.0)
1080
+
1081
+ first_block_index = start // Int64(self.state_block_size)
1082
+ warp_block_index = first_block_index + (warp_id * 2).to(Int64)
1083
+ block0_i32 = Int32(0)
1084
+ block1_i32 = Int32(0)
1085
+ if lane_id == 0:
1086
+ block0_i32 = block_table[req_idx, warp_block_index]
1087
+ block1_i32 = block_table[req_idx, warp_block_index + Int64(1)]
1088
+ block0_i32 = cute.arch.shuffle_sync(block0_i32, offset=0)
1089
+ block1_i32 = cute.arch.shuffle_sync(block1_i32, offset=0)
1090
+
1091
+ cp_f32x2 = cute.make_copy_atom(
1092
+ cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=64
1093
+ )
1094
+ final_mask_and_clamp = const_expr(
1095
+ (cute.arch.WARP_SIZE - self.num_warps) << 8 | (cute.arch.WARP_SIZE - 1)
1096
+ )
1097
+ col_tile = col_base.to(Int64) // Int64(self.elems_per_lane)
1098
+ score_col_tile = col_tile + Int64(self.state_width // self.elems_per_lane)
1099
+
1100
+ for i in cutlass.range_constexpr(self.rows_per_warp):
1101
+ block_number_i32 = block0_i32
1102
+ block_offset = Int64(i)
1103
+ if const_expr(i >= self.state_block_size):
1104
+ block_number_i32 = block1_i32
1105
+ block_offset = Int64(i - self.state_block_size)
1106
+ row_tensor = state_cache[block_number_i32.to(Int64), block_offset, None]
1107
+ kv_src = cute.local_tile(
1108
+ row_tensor,
1109
+ tiler=(self.elems_per_lane,),
1110
+ coord=(col_tile,),
1111
+ )
1112
+ score_src = cute.local_tile(
1113
+ row_tensor,
1114
+ tiler=(self.elems_per_lane,),
1115
+ coord=(score_col_tile,),
1116
+ )
1117
+ cute.copy(cp_f32x2, kv_src, kv_vals[i, None])
1118
+ cute.copy(cp_f32x2, score_src, score_vals[i, None])
1119
+
1120
+ for e in cutlass.range_constexpr(self.elems_per_lane):
1121
+ local_max[e] = cute.arch.fmax(local_max[e], score_vals[i, e])
1122
1123
for e in cutlass.range_constexpr(self.elems_per_lane):
1124
for i in cutlass.range_constexpr(self.rows_per_warp):
@@ -709,122 +1129,407 @@ class SparseAttnCompressC128Block8Kernel:
1129
local_sum[e] += exp_score
1130
local_product[e] += kv_vals[i, e] * exp_score
1131
712
- for e in cutlass.range_constexpr(self.elems_per_lane):
713
- s_max[warp_id, col_group, e] = local_max[e]
714
- s_sum[warp_id, col_group, e] = local_sum[e]
715
- s_product[warp_id, col_group, e] = local_product[e]
716
- cute.arch.sync_threads()
1132
+ for e in cutlass.range_constexpr(self.elems_per_lane):
1133
+ s_max[warp_id, col_group, e] = local_max[e]
1134
+ s_sum[warp_id, col_group, e] = local_sum[e]
1135
+ s_product[warp_id, col_group, e] = local_product[e]
1136
+ cute.arch.sync_threads()
1137
+
1138
+ out_group = tid // self.num_warps
1139
+ final_lane = tid % self.num_warps
1140
+ final_groups_per_pass = const_expr(self.tb_size // self.num_warps)
1141
+ for pass_idx in cutlass.range_constexpr(
1142
+ self.head_tile // final_groups_per_pass
1143
+ ):
1144
+ out_idx = pass_idx * final_groups_per_pass + out_group
1145
+ out_lane = out_idx // self.elems_per_lane
1146
+ out_elem = out_idx % self.elems_per_lane
1147
+
1148
+ local_warp_max = s_max[final_lane, out_lane, out_elem]
1149
+ global_max = local_warp_max
1150
+ for step in cutlass.range_constexpr(self.final_reduce_steps):
1151
+ offset = const_expr(self.final_reduce_initial_offset >> step)
1152
+ global_max = cute.arch.fmax(
1153
+ global_max,
1154
+ cute.arch.shuffle_sync_bfly(
1155
+ global_max,
1156
+ offset=offset,
1157
+ mask_and_clamp=final_mask_and_clamp,
1158
+ ),
1159
+ )
1160
+
1161
+ scale = cute.math.exp2(
1162
+ (local_warp_max - global_max) * Float32(self.rcp_ln2),
1163
+ fastmath=True,
1164
+ )
1165
+ global_sum = s_sum[final_lane, out_lane, out_elem] * scale
1166
+ global_product = s_product[final_lane, out_lane, out_elem] * scale
1167
+ for step in cutlass.range_constexpr(self.final_reduce_steps):
1168
+ offset = const_expr(self.final_reduce_initial_offset >> step)
1169
+ global_sum += cute.arch.shuffle_sync_bfly(
1170
+ global_sum,
1171
+ offset=offset,
1172
+ mask_and_clamp=final_mask_and_clamp,
1173
+ )
1174
+ global_product += cute.arch.shuffle_sync_bfly(
1175
+ global_product,
1176
+ offset=offset,
1177
+ mask_and_clamp=final_mask_and_clamp,
1178
+ )
1179
+
1180
+ if final_lane == 0:
1181
+ compressed_kv.iterator[
1182
+ token_idx.to(Int64) * compressed_kv.stride[0]
1183
+ + (split_idx * self.head_tile + out_idx).to(Int64)
1184
+ ] = global_product / global_sum
1185
+
1186
+ @cache
1187
+ @staticmethod
1188
+ def compile(
1189
+ head_size: int = 512,
1190
+ state_width: int = 512,
1191
+ ):
1192
+ if head_size % SparseAttnCompressC128Block8Kernel.head_tile != 0:
1193
+ raise ValueError("head_size must be divisible by the 64-wide head tile.")
1194
+ num_positions = cute.sym_int()
1195
+ num_slots = cute.sym_int()
1196
+ num_req_indices = cute.sym_int()
1197
+ num_blocks = cute.sym_int()
1198
+ state_cache_block_size = cute.sym_int()
1199
+ block_table_width = cute.sym_int()
1200
+ state_cache_width = state_width * 2
1201
+
1202
+ state_cache = cute.runtime.make_fake_tensor(
1203
+ Float32,
1204
+ (num_blocks, state_cache_block_size, state_cache_width),
1205
+ stride=(
1206
+ cute.sym_int64(divisibility=16),
1207
+ cute.sym_int64(divisibility=16),
1208
+ 1,
1209
+ ),
1210
+ assumed_align=16,
1211
+ )
1212
+ token_to_req_indices = make_fake_tensor(
1213
+ Int32, (num_req_indices,), divisibility=4
1214
+ )
1215
+ positions = make_fake_tensor(Int64, (num_positions,), divisibility=8)
1216
+ slot_mapping = make_fake_tensor(Int64, (num_slots,), divisibility=8)
1217
+ block_table = make_fake_tensor(
1218
+ Int32, (cute.sym_int(), block_table_width), divisibility=1
1219
+ )
1220
+ compressed_kv = cute.runtime.make_fake_tensor(
1221
+ Float32,
1222
+ (num_slots, head_size),
1223
+ stride=(head_size, 1),
1224
+ assumed_align=4,
1225
+ )
1226
+
1227
+ kernel = SparseAttnCompressC128Block8Kernel(
1228
+ head_size,
1229
+ state_width,
1230
+ )
1231
+ stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
1232
+ return cute.compile(
1233
+ kernel,
1234
+ state_cache,
1235
+ token_to_req_indices,
1236
+ positions,
1237
+ slot_mapping,
1238
+ block_table,
1239
+ compressed_kv,
1240
+ stream,
1241
+ options="--enable-tvm-ffi",
1242
+ )
1243
+
1244
+
1245
+class SparseAttnNormRopeStoreKernel:
1246
+ min_scale = 1.0e-4
1247
+
1248
+ def __init__(
1249
+ self,
1250
+ head_size: int,
1251
+ rope_head_dim: int,
1252
+ fp8_max: float,
1253
+ quant_block: int,
1254
+ token_stride: int,
1255
+ scale_dim: int,
1256
+ compress_ratio: int,
1257
+ static_kv_cache_block_size: int,
1258
+ ):
1259
+ self.head_dim = head_size
1260
+ self.rope_dim = rope_head_dim
1261
+ self.nope_dim = head_size - rope_head_dim
1262
+ self.fp8_max = fp8_max
1263
+ self.quant_block = quant_block
1264
+ self.token_stride = token_stride
1265
+ self.scale_dim = scale_dim
1266
+ self.num_warps = head_size // quant_block
1267
+ self.nope_blocks = self.nope_dim // quant_block
1268
+ self.tb_size = head_size // 2
1269
+ self.compress_ratio = compress_ratio
1270
+ self.static_kv_cache_block_size = static_kv_cache_block_size
1271
+
1272
+ @cute.jit
1273
+ def __call__(
1274
+ self,
1275
+ compressed_kv: cute.Tensor,
1276
+ positions: cute.Tensor,
1277
+ slot_mapping: cute.Tensor,
1278
+ rms_norm_weight: cute.Tensor,
1279
+ rms_norm_eps: Float32,
1280
+ cos_sin_cache: cute.Tensor,
1281
+ k_cache: cute.Tensor,
1282
+ kv_slot_mapping: cute.Tensor,
1283
+ stream: CUstream,
1284
+ ):
1285
+ grid = (slot_mapping.shape[0], 1, 1)
1286
+ self.kernel(
1287
+ compressed_kv,
1288
+ positions,
1289
+ slot_mapping,
1290
+ rms_norm_weight,
1291
+ rms_norm_eps,
1292
+ cos_sin_cache,
1293
+ k_cache,
1294
+ kv_slot_mapping,
1295
+ ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream)
1296
718
- out_group = tid // self.num_warps
719
- final_lane = tid % self.num_warps
720
- final_groups_per_pass = const_expr(self.tb_size // self.num_warps)
721
- for pass_idx in cutlass.range_constexpr(
722
- self.head_tile // final_groups_per_pass
723
- ):
724
- out_idx = pass_idx * final_groups_per_pass + out_group
725
- out_lane = out_idx // self.elems_per_lane
726
- out_elem = out_idx % self.elems_per_lane
1297
+ @cute.kernel
1298
+ def kernel(
1299
+ self,
1300
+ compressed_kv: cute.Tensor,
1301
+ positions: cute.Tensor,
1302
+ slot_mapping: cute.Tensor,
1303
+ rms_norm_weight: cute.Tensor,
1304
+ rms_norm_eps: Float32,
1305
+ cos_sin_cache: cute.Tensor,
1306
+ k_cache: cute.Tensor,
1307
+ kv_slot_mapping: cute.Tensor,
1308
+ ):
1309
+ token_idx, _, _ = cute.arch.block_idx()
1310
+ tid, _, _ = cute.arch.thread_idx()
1311
+ warp_id = cute.arch.make_warp_uniform(tid // 32)
1312
+ lane_id = tid % 32
1313
+ elem0 = tid * 2
1314
728
- local_warp_max = s_max[final_lane, out_lane, out_elem]
729
- global_max = local_warp_max
730
- for step in cutlass.range_constexpr(self.final_reduce_steps):
731
- offset = const_expr(self.final_reduce_initial_offset >> step)
732
- global_max = cute.arch.fmax(
733
- global_max,
734
- cute.arch.shuffle_sync_bfly(
735
- global_max,
736
- offset=offset,
737
- mask_and_clamp=final_mask_and_clamp,
738
- ),
739
- )
1315
+ position = Int64(0)
1316
+ kv_slot_idx = Int64(-1)
1317
+ has_position = token_idx < positions.shape[0]
1318
+ slot_id = Int64(-1)
1319
+ if lane_id == 0:
1320
+ slot_id = slot_mapping[token_idx]
1321
+ if lane_id == 0 and has_position:
1322
+ position = positions[token_idx]
1323
+ has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0]
1324
+ if lane_id == 0 and has_kv_slot_idx:
1325
+ kv_slot_idx = kv_slot_mapping[token_idx]
1326
+ slot_id = cute.arch.shuffle_sync(slot_id, offset=0)
1327
+ position = cute.arch.shuffle_sync(position, offset=0)
1328
+ kv_slot_idx = cute.arch.shuffle_sync(kv_slot_idx, offset=0)
1329
+ boundary = has_position and (
1330
+ (position + Int64(1)) % Int64(self.compress_ratio) == Int64(0)
1331
+ )
1332
+ active = slot_id >= Int64(0) and boundary and kv_slot_idx >= Int64(0)
1333
741
- scale = cute.math.exp2(
742
- (local_warp_max - global_max) * Float32(self.rcp_ln2),
743
- fastmath=True,
1334
+ if active:
1335
+ k_cache_u16 = cute.recast_tensor(k_cache, Uint16)
1336
+ k_cache_u32 = cute.recast_tensor(k_cache, Uint32)
1337
+ static_block_size = Int64(self.static_kv_cache_block_size)
1338
+ page = kv_slot_idx // static_block_size
1339
+ kv_offset = kv_slot_idx - page * static_block_size
1340
+ scale_row_offset = static_block_size * Int64(self.token_stride)
1341
+ value_base = page * k_cache.stride[0] + kv_offset * Int64(self.token_stride)
1342
+ scale_base = (
1343
+ page * k_cache.stride[0]
1344
+ + scale_row_offset
1345
+ + kv_offset * Int64(self.scale_dim)
1346
+ )
1347
+ weight0 = rms_norm_weight[elem0].to(Float32)
1348
+ weight1 = rms_norm_weight[elem0 + 1].to(Float32)
1349
+
1350
+ base = token_idx.to(Int64) * compressed_kv.stride[0] + elem0.to(Int64)
1351
+ x0 = compressed_kv.iterator[base]
1352
+ x1 = compressed_kv.iterator[base + Int64(1)]
1353
+
1354
+ local_sumsq = x0 * x0 + x1 * x1
1355
+ warp_sum = local_sumsq
1356
+ for step in cutlass.range_constexpr(5):
1357
+ offset = const_expr(16 >> step)
1358
+ warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset)
1359
+
1360
+ smem = cutlass.utils.SmemAllocator()
1361
+ partial_sums = smem.allocate_tensor(
1362
+ Float32, cute.make_layout((self.num_warps,)), byte_alignment=4
1363
+ )
1364
+
1365
+ if lane_id == 0:
1366
+ partial_sums[warp_id] = warp_sum
1367
+ cute.arch.sync_threads()
1368
+
1369
+ total = partial_sums[lane_id % self.num_warps]
1370
+ sum_mask_and_clamp = const_expr(
1371
+ (cute.arch.WARP_SIZE - self.num_warps) << 8 | (cute.arch.WARP_SIZE - 1)
1372
+ )
1373
+ for step in cutlass.range_constexpr(3):
1374
+ offset = const_expr(4 >> step)
1375
+ total += cute.arch.shuffle_sync_bfly(
1376
+ total,
1377
+ offset,
1378
+ mask_and_clamp=sum_mask_and_clamp,
1379
)
745
- global_sum = s_sum[final_lane, out_lane, out_elem] * scale
746
- global_product = s_product[final_lane, out_lane, out_elem] * scale
747
- for step in cutlass.range_constexpr(self.final_reduce_steps):
748
- offset = const_expr(self.final_reduce_initial_offset >> step)
749
- global_sum += cute.arch.shuffle_sync_bfly(
750
- global_sum,
751
- offset=offset,
752
- mask_and_clamp=final_mask_and_clamp,
753
- )
754
- global_product += cute.arch.shuffle_sync_bfly(
755
- global_product,
756
- offset=offset,
757
- mask_and_clamp=final_mask_and_clamp,
758
- )
1380
760
- if final_lane == 0:
761
- compressed_kv.iterator[
762
- token_idx.to(Int64) * compressed_kv.stride[0]
763
- + (split_idx * self.head_tile + out_idx).to(Int64)
764
- ] = global_product / global_sum
1381
+ rrms = cute.math.rsqrt(
1382
+ total / Float32(self.head_dim) + rms_norm_eps, fastmath=True
1383
+ )
1384
+ x0 = x0 * rrms * weight0
1385
+ x1 = x1 * rrms * weight1
1386
+
1387
+ if warp_id == self.nope_blocks:
1388
+ pair_idx = lane_id
1389
+ compressed_pos = position - Int64(self.compress_ratio - 1)
1390
+ cs_base = compressed_pos * cos_sin_cache.stride[0] + pair_idx.to(Int64)
1391
+ cos_v = cos_sin_cache.iterator[cs_base]
1392
+ sin_v = cos_sin_cache.iterator[cs_base + Int64(self.rope_dim // 2)]
1393
+ real = x0 * cos_v - x1 * sin_v
1394
+ imag = x0 * sin_v + x1 * cos_v
1395
+ packed = _fp32x2_to_bf16x2(real, imag)
1396
+ out_base = value_base + Int64(self.nope_dim) + (lane_id * 4).to(Int64)
1397
+ k_cache_u32.iterator[out_base // Int64(4)] = packed
1398
+ else:
1399
+ q_packed = _fp32x2_to_bf16x2(x0, x1)
1400
+ q0, q1 = _bf16x2_to_fp32(q_packed)
1401
+ abs0 = cute.math.absf(q0)
1402
+ abs1 = cute.math.absf(q1)
1403
+ local_absmax = cute.arch.fmax(abs0, abs1)
1404
+ absmax = local_absmax
1405
+ for step in cutlass.range_constexpr(5):
1406
+ offset = const_expr(16 >> step)
1407
+ absmax = cute.arch.fmax(
1408
+ absmax, cute.arch.shuffle_sync_bfly(absmax, offset)
1409
+ )
1410
+ scale_raw = cute.arch.fmax(
1411
+ Float32(self.min_scale),
1412
+ absmax / Float32(self.fp8_max),
1413
+ )
1414
+ bits = _recast_val(scale_raw, Uint32)
1415
+ ue8m0 = ((bits + Uint32(0x7FFFFF)) >> Uint32(23)) & Uint32(0xFF)
1416
+ inv_scale = _recast_val((Uint32(254) - ue8m0) << Uint32(23), Float32)
1417
+ y0 = cutlass.min(
1418
+ cute.arch.fmax(q0 * inv_scale, Float32(-self.fp8_max)),
1419
+ Float32(self.fp8_max),
1420
+ )
1421
+ y1 = cutlass.min(
1422
+ cute.arch.fmax(q1 * inv_scale, Float32(-self.fp8_max)),
1423
+ Float32(self.fp8_max),
1424
+ )
1425
+ packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1)
1426
+ out_base = value_base + (warp_id * self.quant_block + lane_id * 2).to(
1427
+ Int64
1428
+ )
1429
+ k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8
1430
+ if lane_id == 0:
1431
+ k_cache.iterator[scale_base + warp_id.to(Int64)] = ue8m0.to(Uint8)
1432
+ if warp_id == 0:
1433
+ k_cache.iterator[scale_base + Int64(self.nope_blocks)] = Uint8(
1434
+ 0
1435
+ )
1436
1437
@cache
1438
@staticmethod
1439
def compile(
1440
head_size: int = 512,
770
- state_width: int = 512,
1441
+ rope_head_dim: int = 64,
1442
+ fp8_max: float = 448.0,
1443
+ quant_block: int = 64,
1444
+ token_stride: int = 576,
1445
+ scale_dim: int = 8,
1446
+ kv_block_stride: int = 74752,
1447
+ compress_ratio: int = 128,
1448
+ norm_weight_dtype: type[cutlass.Numeric] = Float32,
1449
+ static_kv_cache_block_size: int = 0,
1450
):
772
- if head_size % SparseAttnCompressC128Block8Kernel.head_tile != 0:
773
- raise ValueError("head_size must be divisible by the 64-wide head tile.")
1451
+ if quant_block != 64:
1452
+ raise ValueError(
1453
+ "CuTe DSL sparse-attn store currently requires quant_block=64."
1454
+ )
1455
+ if rope_head_dim != 64:
1456
+ raise ValueError(
1457
+ "CuTe DSL sparse-attn store currently requires rope_head_dim=64."
1458
+ )
1459
+ if head_size % quant_block != 0:
1460
+ raise ValueError("head_size must be divisible by quant_block.")
1461
+ if token_stride < head_size + rope_head_dim:
1462
+ raise ValueError("token_stride is too small for the packed FP8/BF16 row.")
1463
+ expected_scale_dim = (head_size - rope_head_dim) // quant_block + 1
1464
+ if scale_dim < expected_scale_dim:
1465
+ raise ValueError("scale_dim is too small for the UE8M0 scale row.")
1466
+ if static_kv_cache_block_size <= 0:
1467
+ raise ValueError(
1468
+ "CuTe DSL sparse-attn store requires a positive static "
1469
+ "kv_cache_block_size."
1470
+ )
1471
num_positions = cute.sym_int()
1472
num_slots = cute.sym_int()
776
- num_req_indices = cute.sym_int()
1473
+ num_kv_slots = cute.sym_int()
1474
+ max_pos = cute.sym_int()
1475
num_blocks = cute.sym_int()
778
- state_cache_block_size = cute.sym_int()
779
- block_table_width = cute.sym_int()
780
- state_cache_width = state_width * 2
1476
782
- state_cache = cute.runtime.make_fake_tensor(
1477
+ compressed_kv = cute.runtime.make_fake_tensor(
1478
Float32,
784
- (num_blocks, state_cache_block_size, state_cache_width),
785
- stride=(
786
- cute.sym_int64(divisibility=16),
787
- cute.sym_int64(divisibility=16),
788
- 1,
789
- ),
790
- assumed_align=16,
791
- )
792
- token_to_req_indices = make_fake_tensor(
793
- Int32, (num_req_indices,), divisibility=4
1479
+ (num_slots, head_size),
1480
+ stride=(head_size, 1),
1481
+ assumed_align=4,
1482
)
1483
positions = make_fake_tensor(Int64, (num_positions,), divisibility=8)
1484
slot_mapping = make_fake_tensor(Int64, (num_slots,), divisibility=8)
797
- block_table = make_fake_tensor(
798
- Int32, (cute.sym_int(), block_table_width), divisibility=1
1485
+ rms_norm_weight = make_fake_tensor(
1486
+ norm_weight_dtype, (head_size,), divisibility=4
1487
)
800
- compressed_kv = cute.runtime.make_fake_tensor(
1488
+ cos_sin_cache = cute.runtime.make_fake_tensor(
1489
Float32,
802
- (num_slots, head_size),
803
- stride=(head_size, 1),
1490
+ (max_pos, rope_head_dim),
1491
+ stride=(rope_head_dim, 1),
1492
assumed_align=4,
1493
)
1494
+ k_cache = cute.runtime.make_fake_tensor(
1495
+ Uint8,
1496
+ (num_blocks, cute.sym_int(), cute.sym_int()),
1497
+ stride=(
1498
+ kv_block_stride,
1499
+ cute.sym_int64(divisibility=8),
1500
+ 1,
1501
+ ),
1502
+ assumed_align=16,
1503
+ )
1504
+ kv_slot_mapping = make_fake_tensor(Int64, (num_kv_slots,), divisibility=8)
1505
807
- kernel = SparseAttnCompressC128Block8Kernel(
1506
+ kernel = SparseAttnNormRopeStoreKernel(
1507
head_size,
809
- state_width,
1508
+ rope_head_dim,
1509
+ fp8_max,
1510
+ quant_block,
1511
+ token_stride,
1512
+ scale_dim,
1513
+ compress_ratio,
1514
+ static_kv_cache_block_size,
1515
)
1516
stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
1517
return cute.compile(
1518
kernel,
814
- state_cache,
815
- token_to_req_indices,
1519
+ compressed_kv,
1520
positions,
1521
slot_mapping,
818
- block_table,
819
- compressed_kv,
1522
+ rms_norm_weight,
1523
+ Float32(0.0),
1524
+ cos_sin_cache,
1525
+ k_cache,
1526
+ kv_slot_mapping,
1527
stream,
1528
options="--enable-tvm-ffi",
1529
)
1530
1531
825
-class SparseAttnNormRopeStoreKernel:
826
- min_scale = 1.0e-4
827
-
1532
+class SparseAttnNormRopeStoreFullKernel:
1533
def __init__(
1534
self,
1535
head_size: int,
@@ -834,8 +1539,10 @@ class SparseAttnNormRopeStoreKernel:
1539
token_stride: int,
1540
scale_dim: int,
1541
compress_ratio: int,
837
- static_kv_cache_block_size: int,
1542
+ store_full_fp8: bool = False,
1543
):
1544
+ # Standalone (not inheriting the #44230-restructured legacy kernel):
1545
+ # set attrs directly so the full-cache C128 path is decoupled.
1546
self.head_dim = head_size
1547
self.rope_dim = rope_head_dim
1548
self.nope_dim = head_size - rope_head_dim
@@ -847,7 +1554,7 @@ class SparseAttnNormRopeStoreKernel:
1554
self.nope_blocks = self.nope_dim // quant_block
1555
self.tb_size = head_size // 2
1556
self.compress_ratio = compress_ratio
850
- self.static_kv_cache_block_size = static_kv_cache_block_size
1557
+ self.store_full_fp8 = store_full_fp8
1558
1559
@cute.jit
1560
def __call__(
@@ -860,6 +1567,8 @@ class SparseAttnNormRopeStoreKernel:
1567
cos_sin_cache: cute.Tensor,
1568
k_cache: cute.Tensor,
1569
kv_slot_mapping: cute.Tensor,
1570
+ kv_cache_block_size: Int64,
1571
+ fp8_scale: cute.Tensor,
1572
stream: CUstream,
1573
):
1574
grid = (slot_mapping.shape[0], 1, 1)
@@ -872,6 +1581,8 @@ class SparseAttnNormRopeStoreKernel:
1581
cos_sin_cache,
1582
k_cache,
1583
kv_slot_mapping,
1584
+ kv_cache_block_size,
1585
+ fp8_scale,
1586
).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream)
1587
1588
@cute.kernel
@@ -885,6 +1596,8 @@ class SparseAttnNormRopeStoreKernel:
1596
cos_sin_cache: cute.Tensor,
1597
k_cache: cute.Tensor,
1598
kv_slot_mapping: cute.Tensor,
1599
+ kv_cache_block_size: Int64,
1600
+ fp8_scale: cute.Tensor,
1601
):
1602
token_idx, _, _ = cute.arch.block_idx()
1603
tid, _, _ = cute.arch.thread_idx()
@@ -892,41 +1605,21 @@ class SparseAttnNormRopeStoreKernel:
1605
lane_id = tid % 32
1606
elem0 = tid * 2
1607
895
- position = Int64(0)
896
- kv_slot_idx = Int64(-1)
1608
+ slot_id = slot_mapping[token_idx]
1609
has_position = token_idx < positions.shape[0]
898
- slot_id = Int64(-1)
899
- if lane_id == 0:
900
- slot_id = slot_mapping[token_idx]
901
- if lane_id == 0 and has_position:
1610
+ position = Int64(0)
1611
+ if has_position:
1612
position = positions[token_idx]
903
- has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0]
904
- if lane_id == 0 and has_kv_slot_idx:
905
- kv_slot_idx = kv_slot_mapping[token_idx]
906
- slot_id = cute.arch.shuffle_sync(slot_id, offset=0)
907
- position = cute.arch.shuffle_sync(position, offset=0)
908
- kv_slot_idx = cute.arch.shuffle_sync(kv_slot_idx, offset=0)
1613
boundary = has_position and (
1614
(position + Int64(1)) % Int64(self.compress_ratio) == Int64(0)
1615
)
1616
+ has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0]
1617
+ kv_slot_idx = Int64(-1)
1618
+ if has_kv_slot_idx:
1619
+ kv_slot_idx = kv_slot_mapping[token_idx]
1620
active = slot_id >= Int64(0) and boundary and kv_slot_idx >= Int64(0)
1621
1622
if active:
915
- k_cache_u16 = cute.recast_tensor(k_cache, Uint16)
916
- k_cache_u32 = cute.recast_tensor(k_cache, Uint32)
917
- static_block_size = Int64(self.static_kv_cache_block_size)
918
- page = kv_slot_idx // static_block_size
919
- kv_offset = kv_slot_idx - page * static_block_size
920
- scale_row_offset = static_block_size * Int64(self.token_stride)
921
- value_base = page * k_cache.stride[0] + kv_offset * Int64(self.token_stride)
922
- scale_base = (
923
- page * k_cache.stride[0]
924
- + scale_row_offset
925
- + kv_offset * Int64(self.scale_dim)
926
- )
927
- weight0 = rms_norm_weight[elem0].to(Float32)
928
- weight1 = rms_norm_weight[elem0 + 1].to(Float32)
929
-
1623
base = token_idx.to(Int64) * compressed_kv.stride[0] + elem0.to(Int64)
1624
x0 = compressed_kv.iterator[base]
1625
x1 = compressed_kv.iterator[base + Int64(1)]
@@ -941,78 +1634,79 @@ class SparseAttnNormRopeStoreKernel:
1634
partial_sums = smem.allocate_tensor(
1635
Float32, cute.make_layout((self.num_warps,)), byte_alignment=4
1636
)
1637
+ rrms_shared = smem.allocate_tensor(
1638
+ Float32, cute.make_layout((1,)), byte_alignment=4
1639
+ )
1640
1641
if lane_id == 0:
1642
partial_sums[warp_id] = warp_sum
1643
cute.arch.sync_threads()
948
-
949
- total = partial_sums[lane_id % self.num_warps]
950
- sum_mask_and_clamp = const_expr(
951
- (cute.arch.WARP_SIZE - self.num_warps) << 8 | (cute.arch.WARP_SIZE - 1)
952
- )
953
- for step in cutlass.range_constexpr(3):
954
- offset = const_expr(4 >> step)
955
- total += cute.arch.shuffle_sync_bfly(
956
- total,
957
- offset,
958
- mask_and_clamp=sum_mask_and_clamp,
1644
+ if tid == 0:
1645
+ total = Float32(0.0)
1646
+ for i in cutlass.range_constexpr(self.num_warps):
1647
+ total += partial_sums[i]
1648
+ rrms_shared[0] = cute.math.rsqrt(
1649
+ total / Float32(self.head_dim) + rms_norm_eps, fastmath=True
1650
)
1651
+ cute.arch.sync_threads()
1652
961
- rrms = cute.math.rsqrt(
962
- total / Float32(self.head_dim) + rms_norm_eps, fastmath=True
963
- )
964
- x0 = x0 * rrms * weight0
965
- x1 = x1 * rrms * weight1
1653
+ rrms = rrms_shared[0]
1654
+ x0 = x0 * rrms * rms_norm_weight[elem0].to(Float32)
1655
+ x1 = x1 * rrms * rms_norm_weight[elem0 + 1].to(Float32)
1656
967
- if warp_id == self.nope_blocks:
968
- pair_idx = lane_id
969
- compressed_pos = position - Int64(self.compress_ratio - 1)
970
- cs_base = compressed_pos * cos_sin_cache.stride[0] + pair_idx.to(Int64)
971
- cos_v = cos_sin_cache.iterator[cs_base]
972
- sin_v = cos_sin_cache.iterator[cs_base + Int64(self.rope_dim // 2)]
973
- real = x0 * cos_v - x1 * sin_v
974
- imag = x0 * sin_v + x1 * cos_v
975
- packed = _fp32x2_to_bf16x2(real, imag)
976
- out_base = value_base + Int64(self.nope_dim) + (lane_id * 4).to(Int64)
977
- k_cache_u32.iterator[out_base // Int64(4)] = packed
978
- else:
979
- q_packed = _fp32x2_to_bf16x2(x0, x1)
980
- q0, q1 = _bf16x2_to_fp32(q_packed)
981
- abs0 = cute.math.absf(q0)
982
- abs1 = cute.math.absf(q1)
983
- local_absmax = cute.arch.fmax(abs0, abs1)
984
- absmax = local_absmax
985
- for step in cutlass.range_constexpr(5):
986
- offset = const_expr(16 >> step)
987
- absmax = cute.arch.fmax(
988
- absmax, cute.arch.shuffle_sync_bfly(absmax, offset)
1657
+ page = kv_slot_idx // kv_cache_block_size
1658
+ kv_offset = kv_slot_idx - page * kv_cache_block_size
1659
+ value_base = page * k_cache.stride[0] + kv_offset * k_cache.stride[1]
1660
+
1661
+ if const_expr(self.store_full_fp8):
1662
+ k_cache_u16 = cute.recast_tensor(k_cache, Uint16)
1663
+ inv_fp8 = Float32(1.0) / fp8_scale[0]
1664
+ fp8_v0 = x0
1665
+ fp8_v1 = x1
1666
+ if warp_id == self.nope_blocks:
1667
+ compressed_pos = (position // Int64(self.compress_ratio)) * Int64(
1668
+ self.compress_ratio
1669
)
990
- scale_raw = cute.arch.fmax(
991
- Float32(self.min_scale),
992
- absmax / Float32(self.fp8_max),
993
- )
994
- bits = _recast_val(scale_raw, Uint32)
995
- ue8m0 = ((bits + Uint32(0x7FFFFF)) >> Uint32(23)) & Uint32(0xFF)
996
- inv_scale = _recast_val((Uint32(254) - ue8m0) << Uint32(23), Float32)
1670
+ pair_idx = lane_id
1671
+ cs_base = compressed_pos * cos_sin_cache.stride[0] + pair_idx.to(
1672
+ Int64
1673
+ )
1674
+ cos_v = cos_sin_cache.iterator[cs_base]
1675
+ sin_v = cos_sin_cache.iterator[cs_base + Int64(self.rope_dim // 2)]
1676
+ fp8_v0 = x0 * cos_v - x1 * sin_v
1677
+ fp8_v1 = x0 * sin_v + x1 * cos_v
1678
+ fp8_packed_bf16 = _fp32x2_to_bf16x2(fp8_v0, fp8_v1)
1679
+ b0, b1 = _bf16x2_to_fp32(fp8_packed_bf16)
1680
y0 = cutlass.min(
998
- cute.arch.fmax(q0 * inv_scale, Float32(-self.fp8_max)),
1681
+ cutlass.max(b0 * inv_fp8, Float32(-self.fp8_max)),
1682
Float32(self.fp8_max),
1683
)
1684
y1 = cutlass.min(
1002
- cute.arch.fmax(q1 * inv_scale, Float32(-self.fp8_max)),
1685
+ cutlass.max(b1 * inv_fp8, Float32(-self.fp8_max)),
1686
Float32(self.fp8_max),
1687
)
1688
packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1)
1006
- out_base = value_base + (warp_id * self.quant_block + lane_id * 2).to(
1007
- Int64
1008
- )
1689
+ out_base = value_base + elem0.to(Int64)
1690
k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8
1010
- if lane_id == 0:
1011
- k_cache.iterator[scale_base + warp_id.to(Int64)] = ue8m0.to(Uint8)
1012
- if warp_id == 0:
1013
- k_cache.iterator[scale_base + Int64(self.nope_blocks)] = Uint8(
1014
- 0
1015
- )
1691
+ else:
1692
+ k_cache_u32 = cute.recast_tensor(k_cache, Uint32)
1693
+ bf16_v0 = x0
1694
+ bf16_v1 = x1
1695
+ if warp_id == self.nope_blocks:
1696
+ compressed_pos = (position // Int64(self.compress_ratio)) * Int64(
1697
+ self.compress_ratio
1698
+ )
1699
+ pair_idx = lane_id
1700
+ cs_base = compressed_pos * cos_sin_cache.stride[0] + pair_idx.to(
1701
+ Int64
1702
+ )
1703
+ cos_v = cos_sin_cache.iterator[cs_base]
1704
+ sin_v = cos_sin_cache.iterator[cs_base + Int64(self.rope_dim // 2)]
1705
+ bf16_v0 = x0 * cos_v - x1 * sin_v
1706
+ bf16_v1 = x0 * sin_v + x1 * cos_v
1707
+ bf16_packed = _fp32x2_to_bf16x2(bf16_v0, bf16_v1)
1708
+ out_base = value_base + (elem0 * 2).to(Int64)
1709
+ k_cache_u32.iterator[out_base // Int64(4)] = bf16_packed
1710
1711
@cache
1712
@staticmethod
@@ -1025,8 +1719,8 @@ class SparseAttnNormRopeStoreKernel:
1719
scale_dim: int = 8,
1720
kv_block_stride: int = 74752,
1721
compress_ratio: int = 128,
1722
+ store_full_fp8: bool = False,
1723
norm_weight_dtype: type[cutlass.Numeric] = Float32,
1029
- static_kv_cache_block_size: int = 0,
1724
):
1725
if quant_block != 64:
1726
raise ValueError(
@@ -1038,16 +1732,6 @@ class SparseAttnNormRopeStoreKernel:
1732
)
1733
if head_size % quant_block != 0:
1734
raise ValueError("head_size must be divisible by quant_block.")
1041
- if token_stride < head_size + rope_head_dim:
1042
- raise ValueError("token_stride is too small for the packed FP8/BF16 row.")
1043
- expected_scale_dim = (head_size - rope_head_dim) // quant_block + 1
1044
- if scale_dim < expected_scale_dim:
1045
- raise ValueError("scale_dim is too small for the UE8M0 scale row.")
1046
- if static_kv_cache_block_size <= 0:
1047
- raise ValueError(
1048
- "CuTe DSL sparse-attn store requires a positive static "
1049
- "kv_cache_block_size."
1050
- )
1735
num_positions = cute.sym_int()
1736
num_slots = cute.sym_int()
1737
num_kv_slots = cute.sym_int()
@@ -1057,7 +1741,7 @@ class SparseAttnNormRopeStoreKernel:
1741
compressed_kv = cute.runtime.make_fake_tensor(
1742
Float32,
1743
(num_slots, head_size),
1060
- stride=(head_size, 1),
1744
+ stride=(cute.sym_int64(divisibility=4), 1),
1745
assumed_align=4,
1746
)
1747
positions = make_fake_tensor(Int64, (num_positions,), divisibility=8)
@@ -1068,22 +1752,23 @@ class SparseAttnNormRopeStoreKernel:
1752
cos_sin_cache = cute.runtime.make_fake_tensor(
1753
Float32,
1754
(max_pos, rope_head_dim),
1071
- stride=(rope_head_dim, 1),
1755
+ stride=(cute.sym_int64(divisibility=4), 1),
1756
assumed_align=4,
1757
)
1758
k_cache = cute.runtime.make_fake_tensor(
1759
Uint8,
1760
(num_blocks, cute.sym_int(), cute.sym_int()),
1761
stride=(
1078
- kv_block_stride,
1762
+ cute.sym_int64(divisibility=16),
1763
cute.sym_int64(divisibility=8),
1764
1,
1765
),
1766
assumed_align=16,
1767
)
1768
kv_slot_mapping = make_fake_tensor(Int64, (num_kv_slots,), divisibility=8)
1769
+ fp8_scale = make_fake_tensor(Float32, (1,), divisibility=1)
1770
1086
- kernel = SparseAttnNormRopeStoreKernel(
1771
+ kernel = SparseAttnNormRopeStoreFullKernel(
1772
head_size,
1773
rope_head_dim,
1774
fp8_max,
@@ -1091,7 +1776,7 @@ class SparseAttnNormRopeStoreKernel:
1776
token_stride,
1777
scale_dim,
1778
compress_ratio,
1094
- static_kv_cache_block_size,
1779
+ store_full_fp8,
1780
)
1781
stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
1782
return cute.compile(
@@ -1104,6 +1789,8 @@ class SparseAttnNormRopeStoreKernel:
1789
cos_sin_cache,
1790
k_cache,
1791
kv_slot_mapping,
1792
+ Int64(0),
1793
+ fp8_scale,
1794
stream,
1795
options="--enable-tvm-ffi",
1796
)
@@ -1123,6 +1810,8 @@ def compile_split_sparse_attn_cutedsl(
1810
compress_ratio: int,
1811
overlap: bool,
1812
rms_norm_weight_dtype: torch.dtype,
1813
+ store_full_kv: bool = False,
1814
+ store_full_fp8: bool = False,
1815
):
1816
if not (
1817
head_size == 512
@@ -1141,18 +1830,33 @@ def compile_split_sparse_attn_cutedsl(
1830
state_width=state_width,
1831
)
1832
norm_weight_dtype = _TORCH_TO_CUTE[rms_norm_weight_dtype]
1144
- store = SparseAttnNormRopeStoreKernel.compile(
1145
- head_size,
1146
- rope_head_dim,
1147
- fp8_max,
1148
- quant_block,
1149
- token_stride,
1150
- scale_dim,
1151
- kv_block_stride,
1152
- compress_ratio,
1153
- norm_weight_dtype,
1154
- kv_cache_block_size,
1155
- )
1833
+ if store_full_kv:
1834
+ # FlashInfer contiguous bf16/fp8 cache: standalone full-cache store.
1835
+ store = SparseAttnNormRopeStoreFullKernel.compile(
1836
+ head_size=head_size,
1837
+ rope_head_dim=rope_head_dim,
1838
+ fp8_max=fp8_max,
1839
+ quant_block=quant_block,
1840
+ token_stride=token_stride,
1841
+ scale_dim=scale_dim,
1842
+ kv_block_stride=kv_block_stride,
1843
+ compress_ratio=compress_ratio,
1844
+ store_full_fp8=store_full_fp8,
1845
+ norm_weight_dtype=norm_weight_dtype,
1846
+ )
1847
+ else:
1848
+ store = SparseAttnNormRopeStoreKernel.compile(
1849
+ head_size,
1850
+ rope_head_dim,
1851
+ fp8_max,
1852
+ quant_block,
1853
+ token_stride,
1854
+ scale_dim,
1855
+ kv_block_stride,
1856
+ compress_ratio,
1857
+ norm_weight_dtype,
1858
+ kv_cache_block_size,
1859
+ )
1860
return compress, store
1861
1862
@@ -1180,13 +1884,16 @@ def split_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
1884
scale_dim: int = 8,
1885
compress_ratio: int = 128,
1886
overlap: bool = False,
1887
+ store_full_kv: bool = False,
1888
+ store_full_fp8: bool = False,
1889
+ fp8_scale: torch.Tensor | None = None,
1890
) -> None:
1891
if k_cache.ndim != 3:
1892
raise ValueError(
1893
"CuTe DSL sparse-attn store expects the real DeepSeek V4 "
1894
f"3D k_cache layout [num_blocks, block_size, 584], got ndim={k_cache.ndim}."
1895
)
1189
- if kv_cache_block_size != k_cache.shape[1]:
1896
+ if not store_full_kv and kv_cache_block_size != k_cache.shape[1]:
1897
raise ValueError(
1898
"CuTe DSL split sparse-attn wrapper expected kv_cache_block_size "
1899
f"to match k_cache.shape[1], got {kv_cache_block_size} and "
@@ -1199,6 +1906,8 @@ def split_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
1906
"CuTe DSL sparse-attn store supports rms_norm_weight dtype "
1907
f"bf16/fp32, got {rms_norm_weight.dtype}."
1908
)
1909
+ if store_full_fp8 and not store_full_kv:
1910
+ raise ValueError("store_full_fp8 requires store_full_kv.")
1911
compress, store = compile_split_sparse_attn_cutedsl(
1912
head_size,
1913
state_width,
@@ -1213,6 +1922,8 @@ def split_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
1922
compress_ratio,
1923
overlap,
1924
rms_norm_weight.dtype,
1925
+ store_full_kv=store_full_kv,
1926
+ store_full_fp8=store_full_fp8,
1927
)
1928
compress(
1929
state_cache,
@@ -1223,6 +1934,25 @@ def split_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
1934
compressed_kv,
1935
)
1936
1937
+ if store_full_kv:
1938
+ # Byte-addressed contiguous cache; block size + per-tensor scale are
1939
+ # passed at call time (not baked into compile).
1940
+ if fp8_scale is None:
1941
+ fp8_scale = torch.ones(1, dtype=torch.float32, device=k_cache.device)
1942
+ store(
1943
+ compressed_kv,
1944
+ positions,
1945
+ slot_mapping,
1946
+ rms_norm_weight,
1947
+ rms_norm_eps,
1948
+ cos_sin_cache,
1949
+ k_cache.view(torch.uint8),
1950
+ kv_slot_mapping,
1951
+ kv_cache_block_size,
1952
+ fp8_scale,
1953
+ )
1954
+ return
1955
+
1956
store(
1957
compressed_kv,
1958
positions,
@@ -1258,6 +1988,9 @@ def fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
1988
scale_dim: int = 8,
1989
compress_ratio: int = 4,
1990
overlap: bool = True,
1991
+ store_full_kv: bool = False,
1992
+ store_full_fp8: bool = False,
1993
+ fp8_scale: torch.Tensor | None = None,
1994
) -> None:
1995
if positions.numel() == 0:
1996
return
@@ -1272,6 +2005,43 @@ def fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
2005
"CuTe DSL sparse-attn fused store expects the real DeepSeek V4 "
2006
f"3D k_cache layout [num_blocks, block_size, 584], got ndim={k_cache.ndim}."
2007
)
2008
+ if store_full_fp8 and not store_full_kv:
2009
+ raise ValueError("store_full_fp8 requires store_full_kv.")
2010
+ if store_full_kv:
2011
+ # FlashInfer contiguous bf16/fp8 cache: byte-addressed full-cache C4 store.
2012
+ if fp8_scale is None:
2013
+ fp8_scale = torch.ones(1, dtype=torch.float32, device=k_cache.device)
2014
+ compiled = SparseAttnCompressNormRopeStoreFullC4Kernel.compile(
2015
+ head_size=head_size,
2016
+ state_width=state_width,
2017
+ rope_head_dim=rope_head_dim,
2018
+ fp8_max=fp8_max,
2019
+ quant_block=quant_block,
2020
+ token_stride=token_stride,
2021
+ scale_dim=scale_dim,
2022
+ kv_block_stride=kv_block_stride,
2023
+ compress_ratio=compress_ratio,
2024
+ overlap=overlap,
2025
+ store_full_fp8=store_full_fp8,
2026
+ norm_weight_dtype=norm_weight_dtype,
2027
+ )
2028
+ compiled(
2029
+ state_cache,
2030
+ token_to_req_indices,
2031
+ positions,
2032
+ slot_mapping,
2033
+ block_table,
2034
+ block_size,
2035
+ rms_norm_weight,
2036
+ rms_norm_eps,
2037
+ cos_sin_cache,
2038
+ k_cache.view(torch.uint8),
2039
+ kv_slot_mapping,
2040
+ kv_cache_block_size,
2041
+ fp8_scale,
2042
+ )
2043
+ return
2044
+
2045
compiled = SparseAttnCompressNormRopeStoreC4Kernel.compile(
2046
head_size=head_size,
2047
state_width=state_width,
@@ -1324,6 +2094,9 @@ def compress_norm_rope_store_cutedsl(
2094
quant_block: int,
2095
token_stride: int,
2096
scale_dim: int,
2097
+ store_full_kv: bool = False,
2098
+ store_full_fp8: bool = False,
2099
+ fp8_scale: torch.Tensor | None = None,
2100
) -> None:
2101
if compress_ratio == 4:
2102
# For C4A, the single fused kernel is faster than the two-kernel version.
@@ -1350,6 +2123,9 @@ def compress_norm_rope_store_cutedsl(
2123
scale_dim=scale_dim,
2124
compress_ratio=compress_ratio,
2125
overlap=overlap,
2126
+ store_full_kv=store_full_kv,
2127
+ store_full_fp8=store_full_fp8,
2128
+ fp8_scale=fp8_scale,
2129
)
2130
else:
2131
# For C128, the two-kernel version is faster than the single fused kernel.
@@ -1382,4 +2158,7 @@ def compress_norm_rope_store_cutedsl(
2158
scale_dim=scale_dim,
2159
compress_ratio=compress_ratio,
2160
overlap=overlap,
2161
+ store_full_kv=store_full_kv,
2162
+ store_full_fp8=store_full_fp8,
2163
+ fp8_scale=fp8_scale,
2164
)
vllm/utils/flashinfer.py
+17
@@ -72,6 +72,14 @@ def _missing(*_: Any, **__: Any) -> NoReturn:
72
)
73
74
75
+def _missing_dsv4_sparse_mla(*_: Any, **__: Any) -> NoReturn:
76
+ raise RuntimeError(
77
+ "flashinfer.mla.trtllm_batch_decode_sparse_mla_dsv4 is not available. "
78
+ "Install a FlashInfer build that includes DeepSeek V4 sparse MLA "
79
+ "TRTLLM-GEN support."
80
+ )
81
+
82
+
83
def _get_submodule(module_name: str) -> Any | None:
84
"""Safely import a submodule and return it, or None if not available."""
85
try:
@@ -141,6 +149,14 @@ flashinfer_b12x_fused_moe = _lazy_import_wrapper(
149
trtllm_fp4_block_scale_moe = _lazy_import_wrapper(
150
"flashinfer", "trtllm_fp4_block_scale_moe"
151
)
152
+# DeepSeek V4 sparse MLA TRTLLM-GEN decode launcher (public wrapper). Handles
153
+# the SWA + compressed KV pools, the concatenated sparse-index matrix, and
154
+# per-tensor FP8 / BF16 inputs with BF16 output.
155
+flashinfer_trtllm_batch_decode_sparse_mla_dsv4 = _lazy_import_wrapper(
156
+ "flashinfer.mla",
157
+ "trtllm_batch_decode_sparse_mla_dsv4",
158
+ fallback_fn=_missing_dsv4_sparse_mla,
159
+)
160
# Special case for autotune since it returns a context manager
161
autotune = _lazy_import_wrapper(
162
"flashinfer.autotuner",
@@ -965,6 +981,7 @@ __all__ = [
981
"flashinfer_b12x_fused_moe",
982
"flashinfer_convert_sf_to_mma_layout",
983
"trtllm_fp4_block_scale_moe",
984
+ "flashinfer_trtllm_batch_decode_sparse_mla_dsv4",
985
"autotune",
986
"has_flashinfer_moe",
987
"has_flashinfer_comm",
vllm/v1/attention/backends/mla/sparse_swa.py
+7
-2
@@ -73,9 +73,14 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase):
73
# determines the SWA block size of 64 tokens per block.
74
# TODO(yifan): make SWA block size automatically determined and configurable.
75
self.block_size = 64
76
- assert self.dtype == torch.uint8
76
+ # uint8: legacy FlashMLA UE8M0 paged layout. bfloat16 / float8_e4m3fn:
77
+ # FlashInfer contiguous full-cache layout.
78
+ assert self.dtype in (torch.uint8, torch.bfloat16, torch.float8_e4m3fn)
79
80
def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
81
+ # FlashMLA's UE8M0 paged layout needs 576B alignment; FlashInfer's
82
+ # contiguous bf16/fp8 cache uses the natural element-size page.
83
+ is_flashmla = self.cache_config.cache_dtype == "fp8_ds_mla"
84
return SlidingWindowMLASpec(
85
block_size=self.block_size,
86
num_kv_heads=1,
@@ -83,7 +88,7 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase):
88
dtype=self.dtype,
89
sliding_window=self.window_size,
90
cache_dtype_str=self.cache_config.cache_dtype,
86
- alignment=576, # NOTE: FlashMLA requires 576B alignment
91
+ alignment=576 if is_flashmla else None,
92
model_version="deepseek_v4",
93
)
94
vllm/v1/attention/backends/registry.py
+11
@@ -76,6 +76,17 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
76
FLASHMLA_SPARSE = (
77
"vllm.v1.attention.backends.mla.flashmla_sparse.FlashMLASparseBackend"
78
)
79
+ # DeepSeek V4 sparse MLA backends (model-driven; selected via the V4 layer).
80
+ FLASHMLA_SPARSE_DSV4 = (
81
+ "vllm.models.deepseek_v4.nvidia.flashmla.DeepseekV4FlashMLASparseBackend"
82
+ )
83
+ FLASHINFER_MLA_SPARSE_DSV4 = (
84
+ "vllm.models.deepseek_v4.nvidia.flashinfer_sparse."
85
+ "DeepseekV4FlashInferMLASparseBackend"
86
+ )
87
+ ROCM_FLASHMLA_SPARSE_DSV4 = (
88
+ "vllm.models.deepseek_v4.amd.rocm.DeepseekV4ROCMAiterMLASparseBackend"
89
+ )
90
FLASH_ATTN_MLA = "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend"
91
NO_ATTENTION = "vllm.v1.attention.backends.no_attention.NoAttentionBackend"
92
FLEX_ATTENTION = "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend"
vllm/v1/kv_cache_interface.py
+5
-3
@@ -547,10 +547,12 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
547
548
@property
549
def real_page_size_bytes(self) -> int:
550
- if self.model_version == "deepseek_v4":
551
- # DeepseekV4: 448B NoPE + 128B RoPE + 8B fp8 scale = 584B per token.
550
+ if self.model_version == "deepseek_v4" and self.cache_dtype_str == "fp8_ds_mla":
551
+ # DeepseekV4 FlashMLA: 448B NoPE + 128B RoPE + 8B fp8 scale = 584B
552
+ # per token. FlashInfer's contiguous bf16/fp8 cache falls through to
553
+ # the element-size formula below.
554
return self.storage_block_size * 584
553
- assert self.model_version is None, (
555
+ assert self.model_version in (None, "deepseek_v4"), (
556
f"Unsupported model version: {self.model_version}"
557
)
558
return (