@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 7c5dc571cb

[Attention][DSA] Enable W4A16 DSA (#51724)

Signed-off-by: Shiyang Chen <shiychen@nvidia.com>

sychen52 committed Aug 31, 2026 at 21:37 UTC 7c5dc571cbd1064ecc8a9b1045637ff647aa22cb
17 files changed +824 -22
CMakeLists.txt
+2 -1
@@ -1042,7 +1042,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
1042 "csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
1043 "csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu"
1044 "csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu"
1045 - "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu")
1045 + "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu"
1046 + "csrc/libtorch_stable/nvfp4_ds_mla_cache_kernels.cu")
1047 if(NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
1048 message(STATUS
1049 "Building mxfp4_experts_quant unsupported stubs because CUDA compiler version is not >= 12.9 (found ${CMAKE_CUDA_COMPILER_VERSION}).")
cmake/external_projects/flashmla.cmake
+2 -1
@@ -19,7 +19,7 @@ else()
19 FetchContent_Declare(
20 flashmla
21 GIT_REPOSITORY https://github.com/vllm-project/FlashMLA
22 - GIT_TAG 0397728d511c4e3d94ea3a01d8dda8654525a611
22 + GIT_TAG 6bc49418c5ead572ff0339191ddf3b155749e183
23 GIT_PROGRESS TRUE
24 CONFIGURE_COMMAND ""
25 BUILD_COMMAND ""
@@ -111,6 +111,7 @@ if(FLASH_MLA_ARCHS)
111 # sm100 sparse decode
112 ${flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/v32.cu
113 ${flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/model1.cu
114 + ${flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/v32_nvfp4_fp8rope.cu
115 ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu
116 )
117
csrc/libtorch_stable/cache_kernels.cu
+104
@@ -7,6 +7,10 @@
7 #include "quantization/vectorization_utils.cuh"
8 #include "concat_mla_q.cuh"
9
10 +#if !defined(USE_ROCM) && defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100
11 + #include "nvfp4_ds_mla_cache.h"
12 +#endif
13 +
14 #ifdef USE_ROCM
15 #include "../quantization/w8a8/fp8/amd/quant_utils.cuh"
16 #else
@@ -917,6 +921,7 @@ void concat_and_cache_mla(
921 int pe_dim = k_pe.size(1);
922 int block_size = kv_cache.size(1);
923
924 + const bool is_nvfp4_ds_mla = kv_cache_dtype == "nvfp4_ds_mla";
925 if (kv_cache_dtype == "fp8_ds_mla") {
926 STD_TORCH_CHECK(kv_lora_rank == 512,
927 "kv_lora_rank must be 512 for fp8_ds_mla");
@@ -927,6 +932,21 @@ void concat_and_cache_mla(
932 "kv_c.element_size() must be 2 for fp8_ds_mla");
933 STD_TORCH_CHECK(k_pe.element_size() == 2,
934 "k_pe.element_size() must be 2 for fp8_ds_mla");
935 + } else if (is_nvfp4_ds_mla) {
936 + constexpr int bytes_per_token = 352;
937 + STD_TORCH_CHECK(kv_lora_rank == 512, "kv_lora_rank must be 512 for ",
938 + kv_cache_dtype);
939 + STD_TORCH_CHECK(pe_dim == 64, "pe_dim must be 64 for ", kv_cache_dtype);
940 + STD_TORCH_CHECK(
941 + kv_cache.size(2) == bytes_per_token / kv_cache.element_size(),
942 + "kv_cache.size(2) must be ", bytes_per_token, " bytes for ",
943 + kv_cache_dtype);
944 + STD_TORCH_CHECK(
945 + kv_c.scalar_type() == torch::headeronly::ScalarType::BFloat16,
946 + "kv_c must be bfloat16 for ", kv_cache_dtype);
947 + STD_TORCH_CHECK(
948 + k_pe.scalar_type() == torch::headeronly::ScalarType::BFloat16,
949 + "k_pe must be bfloat16 for ", kv_cache_dtype);
950 } else {
951 STD_TORCH_CHECK(kv_cache.size(2) == kv_lora_rank + pe_dim);
952 }
@@ -950,6 +970,24 @@ void concat_and_cache_mla(
970 dim3 block(96);
971 DISPATCH_BY_KV_CACHE_DTYPE(kv_c.scalar_type(), kv_cache_dtype,
972 CALL_CONCAT_AND_CACHE_DS_MLA);
973 + } else if (is_nvfp4_ds_mla) {
974 + // The kernels live in nvfp4_ds_mla_cache_kernels.cu, which is only built
975 + // when the target list contains an SM100 arch.
976 +#if !defined(USE_ROCM) && defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100
977 + const cudaDeviceProp* props = get_device_prop();
978 + STD_TORCH_CHECK(props->major == 10,
979 + "nvfp4_ds_mla KV cache requires SM100 (Blackwell); "
980 + "got SM",
981 + props->major, props->minor);
982 + vllm::launch_concat_and_cache_nvfp4_ds_mla(
983 + kv_c.data_ptr(), k_pe.data_ptr(), kv_cache.data_ptr(),
984 + slot_mapping.const_data_ptr<int64_t>(), block_stride, entry_stride,
985 + kv_c_stride, k_pe_stride, block_size, num_tokens, stream);
986 +#else
987 + STD_TORCH_CHECK(false,
988 + "the nvfp4 ds_mla kv-cache format requires a CUDA build "
989 + "targeting SM100 (Blackwell)");
990 +#endif
991 } else {
992 dim3 grid(num_tokens);
993 dim3 block(std::min(kv_lora_rank, 512));
@@ -1659,6 +1697,72 @@ void cp_gather_and_upconvert_fp8_kv_cache(
1697 dst_entry_stride, seq_starts_ptr);
1698 }
1699
1700 +void cp_gather_and_upconvert_nvfp4_kv_cache(
1701 + torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE,
1702 + // 352]
1703 + torch::stable::Tensor const& dst, // [TOT_TOKENS, 576]
1704 + torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
1705 + torch::stable::Tensor const& workspace_starts, // [BATCH]
1706 + int64_t batch_size) {
1707 + // The kernel lives in nvfp4_ds_mla_cache_kernels.cu, which is only built when
1708 + // the target list contains an SM100 arch.
1709 +#if !defined(USE_ROCM) && defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100
1710 + torch::stable::accelerator::DeviceGuard device_guard(
1711 + src_cache.get_device_index());
1712 + const cudaStream_t stream = get_current_cuda_stream();
1713 +
1714 + const cudaDeviceProp* props = get_device_prop();
1715 + STD_TORCH_CHECK(props->major == 10,
1716 + "nvfp4_ds_mla KV cache requires SM100 (Blackwell); "
1717 + "got SM",
1718 + props->major, props->minor);
1719 +
1720 + int32_t block_size = src_cache.size(1);
1721 + int32_t head_dim = dst.size(1);
1722 +
1723 + STD_TORCH_CHECK(
1724 + block_table.scalar_type() == torch::headeronly::ScalarType::Int,
1725 + "block_table must be int32");
1726 + STD_TORCH_CHECK(
1727 + workspace_starts.scalar_type() == torch::headeronly::ScalarType::Int,
1728 + "workspace_starts must be int32");
1729 + STD_TORCH_CHECK(src_cache.device() == dst.device(),
1730 + "src_cache and dst must be on the same device");
1731 + STD_TORCH_CHECK(src_cache.device() == block_table.device(),
1732 + "src_cache and block_table must be on the same device");
1733 + STD_TORCH_CHECK(src_cache.device() == workspace_starts.device(),
1734 + "src_cache and workspace_starts must be on the same device");
1735 + STD_TORCH_CHECK(
1736 + src_cache.scalar_type() == torch::headeronly::ScalarType::Byte,
1737 + "src_cache must be uint8 for the nvfp4 ds_mla format, but got ",
1738 + src_cache.scalar_type());
1739 + STD_TORCH_CHECK(dst.scalar_type() == torch::headeronly::ScalarType::BFloat16,
1740 + "dst must be bfloat16");
1741 + STD_TORCH_CHECK(head_dim == 576, "head_dim must be 576 for MLA");
1742 + STD_TORCH_CHECK(src_cache.size(2) == 352,
1743 + "src_cache.size(2) must be 352 for the nvfp4 ds_mla format");
1744 +
1745 + int64_t block_table_stride = block_table.stride(0);
1746 + int64_t cache_block_stride = src_cache.stride(0);
1747 + int64_t cache_entry_stride = src_cache.stride(1);
1748 + int64_t dst_entry_stride = dst.stride(0);
1749 +
1750 + const int total_tokens = dst.size(0);
1751 +
1752 + vllm::launch_cp_gather_and_upconvert_nvfp4_kv_cache(
1753 + src_cache.const_data_ptr<uint8_t>(), dst.data_ptr(),
1754 + block_table.const_data_ptr<int32_t>(),
1755 + workspace_starts.const_data_ptr<int32_t>(),
1756 + static_cast<int32_t>(batch_size), block_size, total_tokens,
1757 + block_table_stride, cache_block_stride, cache_entry_stride,
1758 + dst_entry_stride, stream);
1759 +#else
1760 + STD_TORCH_CHECK(false,
1761 + "the nvfp4 ds_mla kv-cache format requires a CUDA build "
1762 + "targeting SM100 (Blackwell)");
1763 +#endif
1764 +}
1765 +
1766 // Macro to dispatch the kernel based on the data type.
1767 #define CALL_INDEXER_K_QUANT_AND_CACHE(KV_T, CACHE_T, KV_DTYPE) \
1768 vllm::indexer_k_quant_and_cache_kernel<KV_T, CACHE_T, KV_DTYPE> \
csrc/libtorch_stable/nvfp4_ds_mla_cache.h new
+37
@@ -0,0 +1,37 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +#pragma once
4 +
5 +#include <cuda_runtime_api.h>
6 +#include <cstdint>
7 +
8 +// Launchers for the nvfp4_ds_mla KV cache kernels. Those kernels live in
9 +// their own TU (nvfp4_ds_mla_cache_kernels.cu) so that only that TU is built
10 +// for the arch-conditional Blackwell targets their e2m1 conversions require;
11 +// these declarations are what keeps cache_kernels.cu generic.
12 +//
13 +// Only defined when the build has an SM100 target (-DENABLE_NVFP4_SM100=1);
14 +// callers must guard on the same macro.
15 +namespace vllm {
16 +
17 +// concat_and_cache_mla for the nvfp4_ds_mla layout.
18 +// kv_c: bf16 [num_tokens, 512]
19 +// k_pe: bf16 [num_tokens, 64]
20 +// kv_cache: uint8 [num_blocks, block_size, 352]
21 +void launch_concat_and_cache_nvfp4_ds_mla(const void* kv_c, const void* k_pe,
22 + void* kv_cache,
23 + const int64_t* slot_mapping,
24 + int block_stride, int entry_stride,
25 + int kv_c_stride, int k_pe_stride,
26 + int block_size, int num_tokens,
27 + cudaStream_t stream);
28 +
29 +// Gather an nvfp4_ds_mla cache into a bf16 [total_tokens, 576] workspace.
30 +void launch_cp_gather_and_upconvert_nvfp4_kv_cache(
31 + const uint8_t* src_cache, void* dst, const int32_t* block_table,
32 + const int32_t* workspace_starts, int32_t num_reqs, int32_t block_size,
33 + int32_t total_tokens, int64_t block_table_stride,
34 + int64_t cache_block_stride, int64_t cache_entry_stride,
35 + int64_t dst_entry_stride, cudaStream_t stream);
36 +
37 +} // namespace vllm
csrc/libtorch_stable/nvfp4_ds_mla_cache_kernels.cu new
+281
@@ -0,0 +1,281 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +// nvfp4_ds_mla KV cache writer and gather/upconvert reader.
5 +//
6 +// Separate TU because its e2m1 conversions are only accepted by ptxas for a
7 +// suffixed Blackwell target (sm_100a / sm_100f); only this file is compiled
8 +// with the suffixed arch list.
9 +
10 +#include "nvfp4_ds_mla_cache.h"
11 +
12 +#include "quantization/fp4/nvfp4_utils.cuh" // vllm::fp32_vec16_to_e2m1()
13 +
14 +#include <cuda_bf16.h>
15 +#include <cuda_fp8.h>
16 +
17 +#if defined(CUDART_VERSION) && CUDART_VERSION >= 12080
18 + #include <cuda_fp4.h>
19 + #define VLLM_HAS_CUDA_FP4 1
20 +#else
21 + #define VLLM_HAS_CUDA_FP4 0
22 +#endif
23 +
24 +// __CUDA_ARCH_FAMILY_SPECIFIC__ is defined only for the suffixed targets. The
25 +// non-native paths trap rather than emulate; the host entry points reject this
26 +// format off SM100, so they are unreachable.
27 +#if VLLM_HAS_CUDA_FP4 && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && \
28 + defined(__CUDA_ARCH_FAMILY_SPECIFIC__)
29 + #define VLLM_NVFP4_NATIVE_CVT 1
30 +#else
31 + #define VLLM_NVFP4_NATIVE_CVT 0
32 +#endif
33 +
34 +namespace vllm {
35 +
36 +// nvfp4_ds_mla cache format, 352 B/token:
37 +// [0, 256) 512 x e2m1 NoPE, packed 2/byte (low nibble = even element)
38 +// [256, 320) 64 x e4m3 RoPE (unscaled: e4m3's 4 exponent bits span the
39 +// RoPE range unaided)
40 +// [320, 352) 32 x e4m3 NoPE scale factors (one per 16 elements), permuted
41 +// Quantization: sf = e4m3(max(amax_16 / 6, 2^-9)), q = round(x / float(sf))
42 +// Dequantization: x = float(q) * float(sf)
43 +//
44 +// The scale-factor region is an 8x4 -> 4x8 transpose: element block s -- NoPE
45 +// dims [16s, 16s + 16) -- has its scale at byte kSfNopeOff + nvfp4_sf_byte(s),
46 +// and byte p holds the scale of element block 4 * (p & 7) + (p >> 3). Keep
47 +// nvfp4_sf_byte() in lockstep with its counterpart in FlashMLA
48 +// (csrc/sm100/decode/head64/config.h and tests/quant.py).
49 +constexpr int kSfNopeOff = 256 + 64;
50 +
51 +__host__ __device__ constexpr int nvfp4_sf_byte(int s) {
52 + return 8 * (s & 3) + (s >> 2);
53 +}
54 +
55 +// Thread q's scale groups {4c + q} must land on the contiguous bytes 8q + c.
56 +// This also makes it a permutation of [0, 32), since both sides cover that
57 +// range exactly once.
58 +constexpr bool nvfp4_sf_perm_ok() {
59 + for (int q = 0; q < 4; ++q) {
60 + for (int c = 0; c < 8; ++c) {
61 + if (nvfp4_sf_byte(4 * c + q) != 8 * q + c) {
62 + return false;
63 + }
64 + }
65 + }
66 + return true;
67 +}
68 +static_assert(nvfp4_sf_perm_ok(),
69 + "NVFP4 SF byte order must match FlashMLA's nvfp4_sf_byte()");
70 +
71 +// `cvt.rn.satfinite.e2m1x2.f32 d, a, b` puts a in d's HIGH nibble and b in its
72 +// LOW nibble, so the odd element of a pair is the first source.
73 +__device__ __forceinline__ uint64_t
74 +nvfp4_pack16_e2m1_rn(const float (&vals)[16], float inv_sf) {
75 +#if VLLM_NVFP4_NATIVE_CVT
76 + float2 pairs[8];
77 + #pragma unroll
78 + for (int i = 0; i < 8; i++) {
79 + pairs[i] = make_float2(vals[2 * i] * inv_sf, vals[2 * i + 1] * inv_sf);
80 + }
81 + const u32x2 packed = fp32_vec16_to_e2m1(pairs);
82 + return (static_cast<uint64_t>(packed.hi) << 32) | packed.lo;
83 +#else
84 + __trap();
85 + return 0;
86 +#endif
87 +}
88 +
89 +// The e2m1 data and its e4m3 scale both convert to f16 and multiply there; the
90 +// product then goes f16 -> f32 -> bf16, as no f16 -> bf16 instruction exists.
91 +__device__ __forceinline__ void nvfp4_unpack16_e2m1(uint64_t raw,
92 + uint8_t sf_byte,
93 + __nv_bfloat16 (&out)[16]) {
94 +#if VLLM_NVFP4_NATIVE_CVT
95 + const __half2_raw sf2_raw = __nv_cvt_fp8x2_to_halfraw2(
96 + static_cast<__nv_fp8x2_storage_t>(sf_byte * 0x0101u), __NV_E4M3);
97 + const __half2 sf2 = *reinterpret_cast<const __half2*>(&sf2_raw);
98 + __nv_bfloat162* out2 = reinterpret_cast<__nv_bfloat162*>(out);
99 + #pragma unroll
100 + for (int i = 0; i < 8; i++) {
101 + // low nibble -> .x, so .x = element 2i, matching the writer.
102 + const __half2_raw v_raw = __nv_cvt_fp4x2_to_halfraw2(
103 + static_cast<__nv_fp4x2_storage_t>(raw >> (8 * i)), __NV_E2M1);
104 + const __half2 v = __hmul2(*reinterpret_cast<const __half2*>(&v_raw), sf2);
105 + out2[i] = __float22bfloat162_rn(__half22float2(v));
106 + }
107 +#else
108 + __trap();
109 +#endif
110 +}
111 +
112 +// Rounded UP to the next e4m3 value so amax / sf never exceeds the payload
113 +// range: round-to-nearest can land up to 33% low near e4m3's 2^-9 floor, which
114 +// would saturate the tile's largest values.
115 +__device__ __forceinline__ float nvfp4_tile_scale(const float* vals,
116 + float max_val,
117 + uint8_t* sf_out) {
118 + float amax = 0.0f;
119 +#pragma unroll
120 + for (int i = 0; i < 16; i++) {
121 + amax = fmaxf(amax, fabsf(vals[i]));
122 + }
123 + const float sf_f = fmaxf(amax / max_val, 0.001953125f); // 2^-9
124 + __nv_fp8_e4m3 sf8(sf_f);
125 + uint8_t sf_bits = *reinterpret_cast<const uint8_t*>(&sf8);
126 + if (float(sf8) < sf_f && sf_bits < 0x7E) {
127 + // Bump to the next e4m3 value (bit patterns of positive e4m3 values are
128 + // monotonic; 0x7E = 448 is the max finite value).
129 + sf_bits += 1;
130 + sf8 = *reinterpret_cast<const __nv_fp8_e4m3*>(&sf_bits);
131 + }
132 + *sf_out = sf_bits;
133 + return float(sf8);
134 +}
135 +
136 +// One CUDA block (64 threads) per token.
137 +// Warp 0: 32 threads, one 16-element NoPE tile each.
138 +// Warp 1: 32 threads, two RoPE elements each.
139 +__global__ void concat_and_cache_nvfp4_ds_mla_kernel(
140 + const __nv_bfloat16* __restrict__ kv_c, // [num_tokens, 512]
141 + const __nv_bfloat16* __restrict__ k_pe, // [num_tokens, 64]
142 + uint8_t* __restrict__ kv_cache, // [num_blocks, block_size, 352]
143 + const int64_t* __restrict__ slot_mapping, // [num_tokens]
144 + const int block_stride, //
145 + const int entry_stride, //
146 + const int kv_c_stride, //
147 + const int k_pe_stride, //
148 + const int block_size) {
149 + const int64_t token_idx = blockIdx.x;
150 + const int64_t slot_idx = slot_mapping[token_idx];
151 + // NOTE: slot_idx can be -1 if the token is padded
152 + if (slot_idx < 0) {
153 + return;
154 + }
155 + const int64_t block_idx = slot_idx / block_size;
156 + const int64_t block_offset = slot_idx % block_size;
157 + uint8_t* token_ptr =
158 + kv_cache + block_idx * block_stride + block_offset * entry_stride;
159 +
160 + if (threadIdx.x < 32) {
161 + const int tile = threadIdx.x;
162 + const __nv_bfloat16* src = kv_c + token_idx * kv_c_stride + tile * 16;
163 + float vals[16];
164 +#pragma unroll
165 + for (int i = 0; i < 16; i++) {
166 + vals[i] = __bfloat162float(src[i]);
167 + }
168 + // `tile` is the element-block index; its scale byte is permuted.
169 + const float sf = nvfp4_tile_scale(
170 + vals, 6.0f, token_ptr + kSfNopeOff + nvfp4_sf_byte(tile));
171 + const float inv_sf = 1.0f / sf;
172 + *reinterpret_cast<uint64_t*>(token_ptr + tile * 8) =
173 + nvfp4_pack16_e2m1_rn(vals, inv_sf);
174 + } else {
175 + // Plain e4m3, unscaled. All 32 lanes take 2 elements each, matching the
176 + // reader and concat_and_cache_ds_mla_kernel.
177 + const int lane = threadIdx.x - 32;
178 + const __nv_bfloat16* src = k_pe + token_idx * k_pe_stride + lane * 2;
179 + uint8_t q[2];
180 +#pragma unroll
181 + for (int i = 0; i < 2; i++) {
182 + const __nv_fp8_e4m3 v8(__bfloat162float(src[i]));
183 + q[i] = *reinterpret_cast<const uint8_t*>(&v8);
184 + }
185 + *reinterpret_cast<uint16_t*>(token_ptr + 256 + lane * 2) =
186 + *reinterpret_cast<const uint16_t*>(q);
187 + }
188 +}
189 +
190 +// One warp per token (mirrors cp_gather_and_upconvert_fp8_kv_cache).
191 +__global__ void cp_gather_and_upconvert_nvfp4_kv_cache_kernel(
192 + const uint8_t* __restrict__ src_cache, // [NUM_BLOCKS, BLOCK_SIZE, 352]
193 + __nv_bfloat16* __restrict__ dst, // [total_tokens, 576]
194 + const int32_t* __restrict__ block_table, // [num_reqs, BLOCK_INDICES]
195 + const int32_t* __restrict__ workspace_starts, // [num_reqs]
196 + const int32_t num_reqs, const int32_t block_size,
197 + const int32_t total_tokens, const int64_t block_table_stride,
198 + const int64_t cache_block_stride, const int64_t cache_entry_stride,
199 + const int64_t dst_entry_stride) {
200 + const int flat_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) >> 5;
201 + if (flat_warp_id >= total_tokens) return;
202 + const int lane_id = threadIdx.x & 31;
203 +
204 + // Binary search to find which request owns this output token
205 + int lo = 0, hi = num_reqs - 1;
206 + while (lo < hi) {
207 + int mid = (lo + hi + 1) >> 1;
208 + if (workspace_starts[mid] <= flat_warp_id)
209 + lo = mid;
210 + else
211 + hi = mid - 1;
212 + }
213 + const int req_id = lo;
214 +
215 + const int out_token_id = flat_warp_id;
216 + const int token_offset = out_token_id - workspace_starts[req_id];
217 + const int cache_block_idx = token_offset / block_size;
218 + const int offset_in_block = token_offset % block_size;
219 + const int physical_block =
220 + block_table[req_id * block_table_stride + cache_block_idx];
221 +
222 + const uint8_t* token_ptr = src_cache + physical_block * cache_block_stride +
223 + offset_in_block * cache_entry_stride;
224 + __nv_bfloat16* dst_ptr = dst + out_token_id * dst_entry_stride;
225 +
226 + // `lane_id` is the element-block index, so its scale byte is permuted.
227 + {
228 + const uint64_t raw =
229 + *reinterpret_cast<const uint64_t*>(token_ptr + lane_id * 8);
230 + const uint8_t sf_byte = token_ptr[kSfNopeOff + nvfp4_sf_byte(lane_id)];
231 + __nv_bfloat16 out[16];
232 + nvfp4_unpack16_e2m1(raw, sf_byte, out);
233 + int4* nope_dst = reinterpret_cast<int4*>(dst_ptr) + lane_id * 2;
234 + nope_dst[0] = *reinterpret_cast<const int4*>(&out[0]);
235 + nope_dst[1] = *reinterpret_cast<const int4*>(&out[8]);
236 + }
237 +
238 + // RoPE is unscaled e4m3.
239 + {
240 + const __nv_fp8_e4m3* raw =
241 + reinterpret_cast<const __nv_fp8_e4m3*>(token_ptr + 256 + lane_id * 2);
242 + const __nv_bfloat16 out[2] = {__float2bfloat16_rn(float(raw[0])),
243 + __float2bfloat16_rn(float(raw[1]))};
244 + *reinterpret_cast<uint32_t*>(dst_ptr + 512 + lane_id * 2) =
245 + *reinterpret_cast<const uint32_t*>(out);
246 + }
247 +}
248 +
249 +void launch_concat_and_cache_nvfp4_ds_mla(const void* kv_c, const void* k_pe,
250 + void* kv_cache,
251 + const int64_t* slot_mapping,
252 + int block_stride, int entry_stride,
253 + int kv_c_stride, int k_pe_stride,
254 + int block_size, int num_tokens,
255 + cudaStream_t stream) {
256 + dim3 grid(num_tokens);
257 + dim3 block(64); // warp 0: 32 NoPE tiles; warp 1: 4 RoPE tiles
258 + concat_and_cache_nvfp4_ds_mla_kernel<<<grid, block, 0, stream>>>(
259 + reinterpret_cast<const __nv_bfloat16*>(kv_c),
260 + reinterpret_cast<const __nv_bfloat16*>(k_pe),
261 + reinterpret_cast<uint8_t*>(kv_cache), slot_mapping, block_stride,
262 + entry_stride, kv_c_stride, k_pe_stride, block_size);
263 +}
264 +
265 +void launch_cp_gather_and_upconvert_nvfp4_kv_cache(
266 + const uint8_t* src_cache, void* dst, const int32_t* block_table,
267 + const int32_t* workspace_starts, int32_t num_reqs, int32_t block_size,
268 + int32_t total_tokens, int64_t block_table_stride,
269 + int64_t cache_block_stride, int64_t cache_entry_stride,
270 + int64_t dst_entry_stride, cudaStream_t stream) {
271 + constexpr int warps_per_block = 8;
272 + const int grid_size = (total_tokens + warps_per_block - 1) / warps_per_block;
273 + const int block_size_threads = warps_per_block * 32; // 256 threads
274 + cp_gather_and_upconvert_nvfp4_kv_cache_kernel<<<grid_size, block_size_threads,
275 + 0, stream>>>(
276 + src_cache, reinterpret_cast<__nv_bfloat16*>(dst), block_table,
277 + workspace_starts, num_reqs, block_size, total_tokens, block_table_stride,
278 + cache_block_stride, cache_entry_stride, dst_entry_stride);
279 +}
280 +
281 +} // namespace vllm
csrc/libtorch_stable/ops.h
+9
@@ -672,6 +672,15 @@ void cp_gather_and_upconvert_fp8_kv_cache(
672 int64_t batch_size,
673 std::optional<torch::stable::Tensor> seq_starts = std::nullopt);
674
675 +// Gather and upconvert an nvfp4_ds_mla KV cache to a BF16 workspace
676 +void cp_gather_and_upconvert_nvfp4_kv_cache(
677 + torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE,
678 + // 352]
679 + torch::stable::Tensor const& dst, // [TOT_TOKENS, 576]
680 + torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
681 + torch::stable::Tensor const& workspace_starts, // [BATCH]
682 + int64_t batch_size);
683 +
684 // Indexer K quantization and cache function
685 void indexer_k_quant_and_cache(
686 torch::stable::Tensor& k, // [num_tokens, head_dim]
csrc/libtorch_stable/torch_bindings.cpp
+6
@@ -996,6 +996,10 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
996 "Tensor block_table, Tensor workspace_starts, int batch_size, Tensor? "
997 "seq_starts) -> ()");
998
999 + ops.def(
1000 + "cp_gather_and_upconvert_nvfp4_kv_cache(Tensor src_cache, Tensor! dst, "
1001 + "Tensor block_table, Tensor workspace_starts, int batch_size) -> ()");
1002 +
1003 ops.def(
1004 "indexer_k_quant_and_cache(Tensor k, Tensor! kv_cache, Tensor "
1005 "slot_mapping, "
@@ -1066,6 +1070,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CUDA, ops) {
1070 ops.impl("cp_gather_cache", TORCH_BOX(&cp_gather_cache));
1071 ops.impl("cp_gather_and_upconvert_fp8_kv_cache",
1072 TORCH_BOX(&cp_gather_and_upconvert_fp8_kv_cache));
1073 + ops.impl("cp_gather_and_upconvert_nvfp4_kv_cache",
1074 + TORCH_BOX(&cp_gather_and_upconvert_nvfp4_kv_cache));
1075 ops.impl("indexer_k_quant_and_cache", TORCH_BOX(&indexer_k_quant_and_cache));
1076 ops.impl("concat_mla_q", TORCH_BOX(&concat_mla_q));
1077 ops.impl("cp_gather_indexer_k_quant_cache",
tests/kernels/attention/test_cache.py
+151
@@ -907,6 +907,157 @@ def test_concat_and_cache_ds_mla(
907 torch.testing.assert_close(kv_rope, ref_rope, atol=0.001, rtol=0.1)
908
909
910 +# Bytes per token for the nvfp4_ds_mla cache layout (see flashmla_sparse.py).
911 +NVFP4_DS_MLA_ENTRY_SIZE = 352
912 +
913 +# e2m1 magnitude table; the sign lives in bit 3 of each 4-bit code.
914 +E2M1_VALUES = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]
915 +
916 +
917 +def _unpack_e2m1(packed: torch.Tensor) -> torch.Tensor:
918 + """Unpack uint8 bytes holding two e2m1 codes each into float32 values.
919 +
920 + The low nibble holds the even element, the high nibble the odd one.
921 + """
922 + low = packed & 0x0F
923 + high = packed >> 4
924 + codes = torch.stack([low, high], dim=-1).reshape(-1).long()
925 + table = torch.tensor(E2M1_VALUES, dtype=torch.float32, device=packed.device)
926 + magnitude = table[codes & 0x7]
927 + sign = torch.where((codes & 0x8) != 0, -1.0, 1.0)
928 + return sign * magnitude
929 +
930 +
931 +def _ref_nvfp4_sf_bytes(vals: torch.Tensor, divisor: float) -> torch.Tensor:
932 + """Reference e4m3 scale-factor bytes: e4m3(max(amax_per_16 / divisor, 2^-9)).
933 +
934 + In element-block order (byte s scales elements [16s, 16s + 16)).
935 + """
936 + tiles = vals.to(torch.float32).reshape(-1, 16)
937 + sf = torch.clamp(tiles.abs().amax(dim=1) / divisor, min=2.0**-9)
938 + return sf.to(torch.float8_e4m3fn).view(torch.uint8)
939 +
940 +
941 +def _nvfp4_sf_element_order(wire_sf_bytes: torch.Tensor) -> torch.Tensor:
942 + """Stored NVFP4 scale-factor byte order -> element-block order.
943 +
944 + The 32 NoPE scale-factor bytes of an nvfp4_ds_mla entry are stored
945 + permuted (an 8x4 -> 4x8 transpose): the scale for element block s lives at
946 + byte 8 * (s & 3) + (s >> 2), so that the 8 scales one FlashMLA dequant
947 + thread needs are contiguous. This undoes that permutation.
948 + """
949 + num_sf = wire_sf_bytes.shape[-1]
950 + return wire_sf_bytes.unflatten(-1, (4, num_sf // 4)).transpose(-1, -2).flatten(-2)
951 +
952 +
953 +@pytest.mark.parametrize("kv_lora_rank", KV_LORA_RANKS)
954 +@pytest.mark.parametrize("qk_rope_head_dim", QK_ROPE_HEAD_DIMS)
955 +@pytest.mark.parametrize("num_tokens", NUM_TOKENS_MLA)
956 +@pytest.mark.parametrize("block_size", BLOCK_SIZES_MLA)
957 +@pytest.mark.parametrize("num_blocks", NUM_BLOCKS_MLA)
958 +@pytest.mark.parametrize("seed", SEEDS)
959 +@pytest.mark.parametrize("device", CUDA_DEVICES)
960 +@torch.inference_mode()
961 +def test_concat_and_cache_nvfp4_ds_mla(
962 + kv_lora_rank: int,
963 + qk_rope_head_dim: int,
964 + num_tokens: int,
965 + block_size: int,
966 + num_blocks: int,
967 + seed: int,
968 + device: str,
969 +) -> None:
970 + if current_platform.is_rocm():
971 + pytest.skip("concat_and_cache_mla doesn't support NVFP4 DS-MLA on ROCm")
972 + device_capability = current_platform.get_device_capability()
973 + if device_capability is None or device_capability.major != 10:
974 + pytest.skip("The NVFP4 DS-MLA kv-cache dtype requires SM 10.x")
975 + if kv_lora_rank != 512:
976 + pytest.skip("The NVFP4 DS-MLA layout requires kv_lora_rank == 512")
977 + dtype = torch.bfloat16
978 + set_random_seed(seed)
979 + torch.set_default_device(device)
980 + torch.accelerator.set_device_index(device)
981 +
982 + kv_cache_dtype = "nvfp4_ds_mla"
983 + entry_size = NVFP4_DS_MLA_ENTRY_SIZE
984 +
985 + total_slots = num_blocks * block_size
986 + slot_mapping_lst = random.sample(range(total_slots), num_tokens)
987 + slot_mapping = torch.tensor(slot_mapping_lst, dtype=torch.long, device=device)
988 +
989 + kv_c = torch.randn(num_tokens, kv_lora_rank, dtype=dtype, device=device)
990 + k_pe = torch.randn(num_tokens, qk_rope_head_dim, dtype=dtype, device=device)
991 +
992 + scale = torch.tensor(1.0, dtype=torch.float32, device=device)
993 + kv_cache = _create_mla_cache(
994 + num_blocks,
995 + block_size,
996 + entry_size,
997 + dtype=torch.uint8,
998 + kv_cache_dtype=kv_cache_dtype,
999 + device=device,
1000 + )
1001 +
1002 + opcheck(
1003 + torch.ops._C_cache_ops.concat_and_cache_mla,
1004 + (kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale),
1005 + test_utils=DEFAULT_OPCHECK_TEST_UTILS,
1006 + )
1007 +
1008 + ops.concat_and_cache_mla(kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale)
1009 +
1010 + # Byte layout of a single cache entry:
1011 + # [0, 256) 512 e2m1 NoPE values packed 2/byte
1012 + # [256, 320) 64 unscaled e4m3 RoPE values
1013 + # [nope_sf_off, 352) 32 e4m3 NoPE SFs (one per 16 elements), stored
1014 + # permuted (see _nvfp4_sf_element_order)
1015 + nope_bytes = kv_lora_rank // 2
1016 + rope_bytes = qk_rope_head_dim
1017 + nope_sf_off = nope_bytes + rope_bytes
1018 + num_nope_sf = kv_lora_rank // 16
1019 + assert nope_sf_off + num_nope_sf == entry_size
1020 +
1021 + for i in range(num_tokens):
1022 + slot = slot_mapping[i].item()
1023 + block_idx = slot // block_size
1024 + block_offset = slot % block_size
1025 + entry = kv_cache[block_idx, block_offset]
1026 +
1027 + kv_c_ref = kv_c[i].to(torch.float32)
1028 + k_pe_ref = k_pe[i].to(torch.float32)
1029 +
1030 + # Scale-factor bytes must match e4m3(max(amax_per_16 / divisor, 2^-9))
1031 + # computed in python, allowing a 1-ulp e4m3 difference (all SFs are
1032 + # positive, so e4m3 bit patterns are monotonic and adjacent codes
1033 + # differ by 1).
1034 + nope_sf_bytes = _nvfp4_sf_element_order(
1035 + entry[nope_sf_off : nope_sf_off + num_nope_sf]
1036 + )
1037 + ref_nope_sf_bytes = _ref_nvfp4_sf_bytes(kv_c_ref, 6.0)
1038 + assert (
1039 + (nope_sf_bytes.to(torch.int16) - ref_nope_sf_bytes.to(torch.int16)).abs()
1040 + <= 1
1041 + ).all()
1042 +
1043 + # Dequantize with the *stored* scale factors: x = float(code) * float(sf).
1044 + nope_sf = nope_sf_bytes.view(torch.float8_e4m3fn).to(torch.float32)
1045 + nope_sf_per_elem = nope_sf.repeat_interleave(16)
1046 + nope_vals = _unpack_e2m1(entry[:nope_bytes]) * nope_sf_per_elem
1047 + # e2m1 error bound: the widest grid spacing is 2 (between codes 4 and
1048 + # 6), so round-to-nearest is off by at most 1.0*sf; the e4m3 rounding
1049 + # of the SF itself adds at most ~2^-4 relative, so 1.5*sf is safe for
1050 + # either encode convention (quantized or unquantized SF).
1051 + assert ((nope_vals - kv_c_ref).abs() <= 1.5 * nope_sf_per_elem).all()
1052 +
1053 + # Plain unscaled e4m3: ~2^-4 relative rounding, with a small absolute
1054 + # floor for values in the subnormal range.
1055 + rope_payload = entry[nope_bytes:nope_sf_off]
1056 + rope_vals = rope_payload.view(torch.float8_e4m3fn).to(torch.float32)
1057 + rope_tol = k_pe_ref.abs() * 2.0**-3 + 2.0**-9
1058 + assert ((rope_vals - k_pe_ref).abs() <= rope_tol).all()
1059 +
1060 +
1061 @pytest.mark.parametrize("kv_lora_rank", KV_LORA_RANKS)
1062 @pytest.mark.parametrize("qk_rope_head_dim", QK_ROPE_HEAD_DIMS)
1063 @pytest.mark.parametrize("block_size", BLOCK_SIZES_MLA)
tests/v1/attention/test_mla_backends.py
+12 -1
@@ -505,7 +505,10 @@ def create_and_prepopulate_kv_cache(
505 block_table = common_attn_metadata.block_table_tensor
506 slot_mapping = common_attn_metadata.slot_mapping
507
508 - fp8_attention = kv_cache_dtype and kv_cache_dtype.startswith("fp8")
508 + use_nvfp4_ds_mla = kv_cache_dtype == "nvfp4_ds_mla"
509 + fp8_attention = (
510 + bool(kv_cache_dtype and kv_cache_dtype.startswith("fp8")) or use_nvfp4_ds_mla
511 + )
512 use_fp8_ds_mla = kv_cache_dtype == "fp8_ds_mla"
513
514 if fp8_attention:
@@ -515,6 +518,14 @@ def create_and_prepopulate_kv_cache(
518 # 4 * 4: 4 float32 scale values for 128-element tiles
519 # 2 * rope_dim: 16-bit RoPE values
520 kv_entry_size = kv_lora_rank + 4 * 4 + 2 * rope_dim
521 + elif use_nvfp4_ds_mla:
522 + kv_lora_rank = kv_c_contexts[0].shape[-1]
523 + rope_dim = k_pe_contexts[0].shape[-1]
524 + # e2m1-packed NoPE + unscaled e4m3 rope + per-16 e4m3 NoPE SFs,
525 + # padded to 16B
526 + kv_entry_size = (
527 + cdiv(kv_lora_rank // 2 + rope_dim + kv_lora_rank // 16, 16) * 16
528 + )
529 else:
530 kv_entry_size = head_size
531
tests/v1/attention/test_sparse_mla_backends.py
+101 -1
@@ -40,6 +40,9 @@ if not current_platform.is_cuda():
40 allow_module_level=True,
41 )
42
43 +from vllm.model_executor.layers.attention.mla_attention import (
44 + _canonicalize_sparse_mla_kv_cache_dtype,
45 +)
46 from vllm.utils.math_utils import cdiv
47 from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
48 FlashInferMLASparseTRTLLMBackend,
@@ -184,13 +187,89 @@ def _quantize_dequantize_fp8_ds_mla(
187 return dequant_kv_c, dequant_k_pe
188
189
190 +_E2M1_TABLE = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]
191 +
192 +
193 +def _quantize_dequantize_nvfp4_ds_mla(
194 + kv_c: torch.Tensor,
195 + k_pe: torch.Tensor,
196 + block_size: int,
197 + kv_cache_dtype: str,
198 + scale: torch.Tensor,
199 +) -> tuple[torch.Tensor, torch.Tensor]:
200 + """Round-trip kv_c/k_pe through the nvfp4_ds_mla cache layout.
201 +
202 + Layout (per token, uint8): [e2m1-packed NoPE | unscaled e4m3 RoPE |
203 + per-16 e4m3 NoPE SFs]. The SF bytes are stored permuted (an 8x4 -> 4x8
204 + transpose: the scale for element block s lives at byte 8 * (s & 3) +
205 + (s >> 2)) so that the 8 scales one FlashMLA dequant thread needs are
206 + contiguous. Dequant: x = value * float(sf), which is exact in bf16, so
207 + this python dequant matches the kernels bit-for-bit.
208 + """
209 + if kv_c.numel() == 0:
210 + return kv_c.clone(), k_pe.clone()
211 +
212 + kv_lora_rank = kv_c.shape[-1]
213 + rope_dim = k_pe.shape[-1]
214 + num_tokens = kv_c.shape[0]
215 + num_blocks = max(1, math.ceil(num_tokens / block_size))
216 + entry_size = (
217 + math.ceil((kv_lora_rank // 2 + rope_dim + kv_lora_rank // 16) / 16) * 16
218 + )
219 + sf_nope_off = kv_lora_rank // 2 + rope_dim
220 +
221 + tmp_cache = torch.zeros(
222 + num_blocks, block_size, entry_size, dtype=torch.uint8, device=kv_c.device
223 + )
224 + slot_mapping = torch.arange(num_tokens, dtype=torch.long, device=kv_c.device)
225 + ops.concat_and_cache_mla(
226 + kv_c, k_pe, tmp_cache, slot_mapping, kv_cache_dtype=kv_cache_dtype, scale=scale
227 + )
228 +
229 + tokens = tmp_cache.view(-1, entry_size)[:num_tokens]
230 + table = torch.tensor(
231 + _E2M1_TABLE + [-v for v in _E2M1_TABLE],
232 + dtype=torch.float32,
233 + device=kv_c.device,
234 + )
235 +
236 + def unpack_e2m1(packed: torch.Tensor) -> torch.Tensor:
237 + lo = (packed & 0xF).long()
238 + hi = (packed >> 4).long()
239 + return table[torch.stack([lo, hi], dim=-1).flatten(-2)]
240 +
241 + nope_vals = unpack_e2m1(tokens[:, : kv_lora_rank // 2])
242 + num_nope_sf = kv_lora_rank // 16
243 + nope_sf = (
244 + # undo the on-wire SF permutation -> element-block order
245 + tokens[:, sf_nope_off : sf_nope_off + num_nope_sf]
246 + .unflatten(-1, (4, num_nope_sf // 4))
247 + .transpose(-1, -2)
248 + .flatten(-2)
249 + .view(torch.float8_e4m3fn)
250 + .float()
251 + )
252 + dequant_kv_c = (
253 + (nope_vals.unflatten(-1, (-1, 16)) * nope_sf.unsqueeze(-1)).flatten(-2)
254 + ).to(kv_c.dtype)
255 +
256 + # RoPE: plain e4m3, no scale factor
257 + rope_raw = tokens[:, kv_lora_rank // 2 : sf_nope_off]
258 + dequant_k_pe = rope_raw.view(torch.float8_e4m3fn).to(k_pe.dtype)
259 +
260 + return dequant_kv_c, dequant_k_pe
261 +
262 +
263 @pytest.mark.parametrize(
264 "backend_cls",
265 [FlashMLASparseBackend, FlashInferMLASparseTRTLLMBackend],
266 ids=["FlashMLA", "FlashInferTRTLLM"],
267 )
268 @pytest.mark.parametrize("batch_name", list(SPARSE_BACKEND_BATCH_SPECS.keys()))
193 -@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_ds_mla"])
269 +@pytest.mark.parametrize(
270 + "kv_cache_dtype",
271 + ["auto", "fp8", "fp8_ds_mla", "nvfp4_ds_mla"],
272 +)
273 @pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4])
274 @pytest.mark.parametrize("block_size", [32, 64])
275 @pytest.mark.parametrize(("q_scale", "k_scale"), [(1.0, 1.0), (2.0, 3.0)])
@@ -219,6 +298,17 @@ def test_sparse_backend_decode_correctness(
298 "fp8_ds_mla kv-cache dtype"
299 )
300
301 + if kv_cache_dtype == "nvfp4_ds_mla":
302 + # Unlike "fp8" (an alias canonicalized to "fp8_ds_mla" above), the
303 + # NVFP4 DS-MLA dtype must reach the backend unchanged.
304 + assert (
305 + _canonicalize_sparse_mla_kv_cache_dtype(backend_cls, kv_cache_dtype)
306 + == kv_cache_dtype
307 + )
308 + device_capability = current_platform.get_device_capability()
309 + if device_capability is None or device_capability.major != 10:
310 + pytest.skip("The NVFP4 DS-MLA kv-cache dtype requires SM 10.x")
311 +
312 supported_block_sizes = backend_cls.get_supported_kernel_block_sizes()
313 if block_size not in supported_block_sizes:
314 pytest.skip(
@@ -238,6 +328,7 @@ def test_sparse_backend_decode_correctness(
328
329 batch_spec = SPARSE_BACKEND_BATCH_SPECS[batch_name]
330 use_fp8_ds_mla_quantization = kv_cache_dtype == "fp8_ds_mla"
331 + use_nvfp4_ds_mla_quantization = kv_cache_dtype == "nvfp4_ds_mla"
332
333 device = torch.device(DEVICE_TYPE)
334 dtype = torch.bfloat16
@@ -381,6 +472,15 @@ def test_sparse_backend_decode_correctness(
472 simulate_sm100_e8m0_scales=is_sm100,
473 )
474 k_pe_full = k_pe_squeezed.unsqueeze(1)
475 + elif use_nvfp4_ds_mla_quantization:
476 + kv_c_full, k_pe_squeezed = _quantize_dequantize_nvfp4_ds_mla(
477 + kv_c_full,
478 + k_pe_full.squeeze(1),
479 + block_size=block_size,
480 + kv_cache_dtype=kv_cache_dtype,
481 + scale=kv_cache_scale,
482 + )
483 + k_pe_full = k_pe_squeezed.unsqueeze(1)
484
485 q_nope, q_pe = q_c.split([qk_nope_head_dim, qk_rope_head_dim], dim=-1)
486 ql_nope = torch.einsum("qnh,lnh->qnl", q_nope, W_UK)
vllm/_custom_ops.py
+25
@@ -3082,6 +3082,31 @@ def cp_gather_and_upconvert_fp8_kv_cache(
3082 )
3083
3084
3085 +def cp_gather_and_upconvert_nvfp4_kv_cache(
3086 + src_cache: torch.Tensor,
3087 + dst: torch.Tensor,
3088 + block_table: torch.Tensor,
3089 + workspace_starts: torch.Tensor,
3090 + batch_size: int,
3091 +) -> None:
3092 + """Gather and upconvert an nvfp4_ds_mla KV cache to a BF16 workspace.
3093 +
3094 + Args:
3095 + src_cache: NVFP4 KV cache [num_blocks, block_size, 352] (uint8)
3096 + dst: BF16 output workspace [total_tokens, 576]
3097 + block_table: Block indices [num_reqs, max_blocks]
3098 + workspace_starts: Workspace start offsets [num_reqs]
3099 + batch_size: Number of requests
3100 + """
3101 + torch.ops._C_cache_ops.cp_gather_and_upconvert_nvfp4_kv_cache(
3102 + src_cache,
3103 + dst,
3104 + block_table,
3105 + workspace_starts,
3106 + batch_size,
3107 + )
3108 +
3109 +
3110 def concat_mla_q(
3111 ql_nope: torch.Tensor,
3112 q_pe: torch.Tensor,
vllm/config/cache.py
+1
@@ -45,6 +45,7 @@ CacheDType = Literal[
45 "fp8_e5m2",
46 "fp8_inc",
47 "fp8_ds_mla",
48 + "nvfp4_ds_mla",
49 "turboquant_k8v4",
50 "turboquant_4bit_nc",
51 "turboquant_k3v4_nc",
vllm/config/vllm.py
+5 -1
@@ -2821,14 +2821,18 @@ class VllmConfig:
2821 def validate_nvfp4_kv_cache_with_mla(self) -> "VllmConfig":
2822 if self.model_config is None:
2823 return self
2824 + # The ds_mla layouts are MLA-only by construction; the plain nvfp4
2825 + # layout (head_size//2 + head_size//16) does not apply to MLA.
2826 if (
2827 self.cache_config.cache_dtype.startswith("nvfp4")
2828 + and not self.cache_config.cache_dtype.endswith("_ds_mla")
2829 and self.model_config.use_mla
2830 ):
2831 raise ValueError(
2832 "nvfp4 KV cache is not supported with MLA (Multi-head Latent "
2833 "Attention) backends. Please use a different --kv-cache-dtype "
2831 - "(e.g., 'fp8' or 'auto') for MLA models such as DeepSeek."
2834 + "(e.g., 'fp8', 'auto', or 'nvfp4_ds_mla' with a sparse MLA "
2835 + "backend) for MLA models such as DeepSeek."
2836 )
2837 return self
2838
vllm/model_executor/layers/attention/mla_attention.py
+15 -4
@@ -352,6 +352,10 @@ def _canonicalize_sparse_mla_kv_cache_dtype(
352 ) -> CacheDType:
353 backend_name = attn_backend.get_name()
354 if backend_name == "FLASHMLA_SPARSE" and is_quantized_kv_cache(kv_cache_dtype):
355 + # The NVFP4 DS-MLA format is used as-is; any other quantized dtype
356 + # (fp8, fp8_e4m3, ...) is served via the fp8_ds_mla format.
357 + if kv_cache_dtype == "nvfp4_ds_mla":
358 + return kv_cache_dtype
359 return "fp8_ds_mla"
360 if backend_name == "FLASHINFER_MLA_SPARSE_SM120" and kv_cache_dtype in (
361 "auto",
@@ -861,7 +865,11 @@ class MLAAttention(nn.Module, AttentionLayerBase):
865 k_c_normed = k_c_normed[:num_actual_toks, ...]
866 k_pe = k_pe[:num_actual_toks, ...]
867
864 - if fp8_attention and self.kv_cache_dtype != "fp8_ds_mla":
868 + if fp8_attention and self.kv_cache_dtype not in (
869 + # Opaque per-token byte formats stay as raw uint8
870 + "fp8_ds_mla",
871 + "nvfp4_ds_mla",
872 + ):
873 kv_cache = kv_cache.view(current_platform.fp8_dtype())
874
875 assert (
@@ -1262,9 +1270,12 @@ class MLAAttention(nn.Module, AttentionLayerBase):
1270 dtype=kv_cache_dtype,
1271 cache_dtype_str=self.kv_cache_dtype,
1272 kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
1265 - # fp8_ds_mla: 656-byte custom layout (kv_lora_rank=512 +
1266 - # qk_rope_head_dim=64, head_size=576). See flashmla_sparse.py.
1267 - state_content_bytes=656 if self.kv_cache_dtype == "fp8_ds_mla" else None,
1273 + # ds_mla layouts pack NoPE + RoPE + scales into one opaque per-token
1274 + # blob, so the size is not derivable from head_size.
1275 + # See flashmla_sparse.py.
1276 + state_content_bytes={"fp8_ds_mla": 656, "nvfp4_ds_mla": 352}.get(
1277 + self.kv_cache_dtype
1278 + ),
1279 )
1280 if self.sliding_window is not None:
1281 return SlidingWindowMLASpec(
vllm/utils/torch_utils.py
+1
@@ -45,6 +45,7 @@ STR_DTYPE_TO_TORCH_DTYPE = {
45 "fp8_per_token_head": torch.uint8,
46 "fp8_inc": torch.float8_e4m3fn,
47 "fp8_ds_mla": torch.uint8,
48 + "nvfp4_ds_mla": torch.uint8,
49 "turboquant_k8v4": torch.uint8,
50 "turboquant_4bit_nc": torch.uint8,
51 "turboquant_k3v4_nc": torch.uint8,
vllm/v1/attention/backends/mla/flashmla_sparse.py
+66 -13
@@ -84,8 +84,28 @@ Bytes, structured as:
84 - **Last 8 bytes:** Scale factors, containing 7 `ue8m0` values + 1B pad.
85 The first `ue8m0` is the scale for the first 64 `float8_e4m3` values,
86 the second for the next 64, and so on.
87 +
88 +In the "nvfp4_ds_mla" format (SM100 only, DeepSeek V3.2 geometry), each
89 +token's KV cache is 352 Bytes, structured as:
90 +- **First 256 bytes:** 512 `e2m1` NoPE values packed 2/byte (low nibble =
91 + even element).
92 +- **Next 64 bytes:** 64 `float8_e4m3` RoPE values. These carry no scale
93 + factor: `e4m3`'s 4 exponent bits span the RoPE magnitude range unaided.
94 +- **Last 32 bytes:** 32 `float8_e4m3` NoPE scale factors, one per 16
95 + elements, stored permuted (an 8x4 -> 4x8 transpose: the scale for element
96 + block `s` lives at byte `8 * (s & 3) + (s >> 2)`) so that the 8 scales one
97 + FlashMLA dequant thread needs are contiguous. See the layout comment in
98 + `csrc/libtorch_stable/cache_kernels.cu`.
99 +
100 """
101
102 +# Quantized DS-MLA cache formats served by the FP8/NVFP4 sparse decode kernel
103 +# path (as opposed to the plain bf16 cache). FlashMLA infers which of these the
104 +# cache holds from its bytes-per-token, so nothing else needs to be passed down.
105 +QUANTIZED_DS_MLA_CACHE_FORMATS: frozenset[str] = frozenset(
106 + {"fp8_ds_mla", "nvfp4_ds_mla"}
107 +)
108 +
109
110 class FlashMLASparseBackend(AttentionBackend):
111 supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
@@ -94,6 +114,7 @@ class FlashMLASparseBackend(AttentionBackend):
114 "bfloat16",
115 "fp8_ds_mla",
116 "fp8", # alias for fp8_ds_mla
117 + "nvfp4_ds_mla", # NVFP4 NoPE + FP8 RoPE (SM100 only)
118 ]
119
120 @staticmethod
@@ -129,6 +150,26 @@ class FlashMLASparseBackend(AttentionBackend):
150 def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
151 return capability.major in [9, 10]
152
153 + @classmethod
154 + def supports_combination(
155 + cls,
156 + head_size: int,
157 + dtype: torch.dtype,
158 + kv_cache_dtype: "CacheDType | None",
159 + block_size: int | None,
160 + use_mla: bool,
161 + has_sink: bool,
162 + use_sparse: bool,
163 + use_mm_prefix: bool,
164 + device_capability: DeviceCapability,
165 + ) -> str | None:
166 + if kv_cache_dtype == "nvfp4_ds_mla" and device_capability.major != 10:
167 + return (
168 + f"FLASHMLA_SPARSE only supports the {kv_cache_dtype} kv-cache "
169 + "dtype on SM100 (Blackwell)"
170 + )
171 + return None
172 +
173
174 @dataclass
175 class FlashMLASparseMetadata(AttentionMetadata):
@@ -258,7 +299,9 @@ class FlashMLASparseMetadataBuilder(
299 FlashMLASparseImpl._compute_fp8_decode_padded_heads(self.num_heads)
300 )
301
261 - self.use_fp8_kv_cache = cache_config.cache_dtype == "fp8_ds_mla"
302 + self.use_fp8_kv_cache = (
303 + cache_config.cache_dtype in QUANTIZED_DS_MLA_CACHE_FORMATS
304 + )
305 max_num_seqs = vllm_config.scheduler_config.max_num_seqs
306 # Shape: [max_num_seqs], all elements = topk_tokens (constant for full-CG)
307 self.topk_tokens_tensor = torch.full(
@@ -596,9 +639,10 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
639 )
640 q_concat_shape = (max_tokens, q_concat_heads, head_size)
641 if is_quantized_kv_cache(kv_cache_dtype):
599 - assert kv_cache_dtype == "fp8_ds_mla", (
600 - "FlashMLA Sparse Attention backend fp8 only supports "
601 - "fp8_ds_mla kv-cache dtype"
642 + assert kv_cache_dtype in QUANTIZED_DS_MLA_CACHE_FORMATS, (
643 + "FlashMLA Sparse Attention backend only supports the "
644 + f"{sorted(QUANTIZED_DS_MLA_CACHE_FORMATS)} quantized kv-cache "
645 + f"dtypes, got {kv_cache_dtype}"
646 )
647
648 if self.need_to_return_lse_for_decode and not is_quantized_kv_cache(
@@ -609,7 +653,7 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
653 "the bf16 sparse path is not supported under DCP."
654 )
655
612 - if kv_cache_dtype == "fp8_ds_mla":
656 + if kv_cache_dtype in QUANTIZED_DS_MLA_CACHE_FORMATS:
657 # Reserve workspace during initialization
658 assert vllm_config is not None and vllm_config.model_config is not None
659 prefill_workspace_size = get_prefill_workspace_size(
@@ -749,13 +793,22 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
793 assert fp8_metadata.prefill is not None
794 for chunk in fp8_metadata.prefill.chunks:
795 chunk_workspace = self.prefill_bf16_workspace[: chunk.chunk_tot_seqlen]
752 - ops.cp_gather_and_upconvert_fp8_kv_cache(
753 - kv_c_and_k_pe_cache,
754 - chunk_workspace,
755 - chunk.block_table,
756 - chunk.workspace_starts,
757 - len(chunk.block_table),
758 - )
796 + if self.kv_cache_dtype == "fp8_ds_mla":
797 + ops.cp_gather_and_upconvert_fp8_kv_cache(
798 + kv_c_and_k_pe_cache,
799 + chunk_workspace,
800 + chunk.block_table,
801 + chunk.workspace_starts,
802 + len(chunk.block_table),
803 + )
804 + else:
805 + ops.cp_gather_and_upconvert_nvfp4_kv_cache(
806 + kv_c_and_k_pe_cache.view(torch.uint8),
807 + chunk_workspace,
808 + chunk.block_table,
809 + chunk.workspace_starts,
810 + len(chunk.block_table),
811 + )
812
813 chunk_q = q[chunk.tokens_slice]
814 chunk_topk_indices_workspace = topk_indices[chunk.tokens_slice]
@@ -953,7 +1006,7 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
1006 assert self.topk_indices_buffer is not None
1007 topk_indices = self.topk_indices_buffer[:num_actual_toks]
1008
956 - use_fp8_cache = self.kv_cache_dtype == "fp8_ds_mla"
1009 + use_fp8_cache = self.kv_cache_dtype in QUANTIZED_DS_MLA_CACHE_FORMATS
1010
1011 lse: torch.Tensor | None = None
1012
vllm/v1/kv_cache_interface.py
+6
@@ -54,6 +54,7 @@ class KVQuantMode(IntEnum):
54 TURBOQUANT_4BIT_NC = 7
55 TURBOQUANT_K3V4_NC = 8
56 TURBOQUANT_3BIT_NC = 9
57 + NVFP4_DS_MLA = 10 # opaque-bytes NVFP4 DS-MLA layouts (FlashMLA sparse)
58
59 @property
60 def is_per_token_head(self) -> bool:
@@ -88,6 +89,11 @@ def get_kv_quant_mode(kv_cache_dtype: str) -> KVQuantMode:
89 return KVQuantMode.INT8_PER_TOKEN_HEAD
90 if kv_cache_dtype == "fp8_per_token_head":
91 return KVQuantMode.FP8_PER_TOKEN_HEAD
92 + # Must precede the ``nvfp4`` prefix test below, which would otherwise match.
93 + if kv_cache_dtype == "nvfp4_ds_mla":
94 + # Page size is keyed on cache_dtype_str in the MLA specs, not
95 + # nvfp4_kv_cache_full_dim.
96 + return KVQuantMode.NVFP4_DS_MLA
97 if kv_cache_dtype.startswith("nvfp4"):
98 return KVQuantMode.NVFP4
99 if isinstance(kv_cache_dtype, str) and kv_cache_dtype.startswith("turboquant_"):