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 7, 00:15 UTC
[Attention][Perf] Optimize cp_gather_and_upconvert_fp8_kv_cache - DeepSeek-v3.2 (#35290)
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com> Co-authored-by: Claude <noreply@anthropic.com>
Roberto L. Castro committed
Mar 9, 2026 at 17:46 UTC
2b28b9b269e18cfe42c7e945d1da8d1c40989efa
3 files changed
+578
-69
benchmarks/kernels/bench_cp_gather_fp8.py
new
+153
@@ -0,0 +1,153 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import argparse
4
+import math
5
+
6
+import torch
7
+
8
+from vllm import _custom_ops as ops
9
+from vllm.triton_utils import triton
10
+
11
+# DeepSeek V3 MLA dimensions
12
+NOPE_DIM = 512
13
+ROPE_DIM = 64
14
+HEAD_DIM = NOPE_DIM + ROPE_DIM # 576 BF16 output elements per token
15
+ENTRY_BYTES = 656 # 512 FP8 + 16 scales + 128 BF16 RoPE
16
+BLOCK_SIZE = 64 # tokens per physical cache block - get_supported_kernel_block_sizes
17
+
18
+# Realistic prefill scenarios:
19
+# - 1 long prefill: single request, 16K-96K tokens
20
+# - 4 medium prefills: 4 requests, 4K-24K tokens each
21
+# - 16 shorter prefills: 16 requests, 1K-6K tokens each
22
+SCENARIOS = [
23
+ # (label, num_reqs, total_tokens_list)
24
+ ("1-req", 1, [8192, 16384, 32768, 65536, 98304]),
25
+ ("4-reqs", 4, [8192, 16384, 32768, 65536, 98304]),
26
+ ("16-reqs", 16, [8192, 16384, 32768, 65536, 98304]),
27
+]
28
+
29
+
30
+def make_inputs(total_tokens, num_reqs, block_size):
31
+ """Create synthetic FP8 cache, block table, and output buffer.
32
+
33
+ Fills the cache with random bytes (we only measure throughput,
34
+ not correctness). Block table maps each request to contiguous
35
+ physical blocks.
36
+ """
37
+ # Divide tokens evenly across requests
38
+ base_len = total_tokens // num_reqs
39
+ remainder = total_tokens % num_reqs
40
+ seq_lens = [base_len + (1 if r < remainder else 0) for r in range(num_reqs)]
41
+
42
+ # workspace_starts: cumulative sum of seq_lens
43
+ workspace_starts = [0] * num_reqs
44
+ for r in range(1, num_reqs):
45
+ workspace_starts[r] = workspace_starts[r - 1] + seq_lens[r - 1]
46
+
47
+ # Physical blocks needed per request
48
+ blocks_per_req = [math.ceil(s / block_size) for s in seq_lens]
49
+ total_blocks = sum(blocks_per_req)
50
+ max_blocks = max(blocks_per_req)
51
+
52
+ # Allocate cache with random data (content doesn't matter for perf)
53
+ cache = torch.randint(
54
+ 0,
55
+ 256,
56
+ (total_blocks, block_size, ENTRY_BYTES),
57
+ dtype=torch.uint8,
58
+ device="cuda",
59
+ )
60
+
61
+ # Block table: contiguous block assignments
62
+ block_table = torch.zeros(num_reqs, max_blocks, dtype=torch.int32, device="cuda")
63
+ block_idx = 0
64
+ for r in range(num_reqs):
65
+ for b in range(blocks_per_req[r]):
66
+ block_table[r, b] = block_idx
67
+ block_idx += 1
68
+
69
+ # Output workspace
70
+ dst = torch.zeros(total_tokens, HEAD_DIM, dtype=torch.bfloat16, device="cuda")
71
+
72
+ seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device="cuda")
73
+ workspace_starts_t = torch.tensor(
74
+ workspace_starts, dtype=torch.int32, device="cuda"
75
+ )
76
+
77
+ return cache, dst, block_table, seq_lens_t, workspace_starts_t
78
+
79
+
80
+def bench_scenario(label, num_reqs, total_tokens_list, save_path):
81
+ """Run benchmark for a specific (num_reqs, total_tokens) scenario."""
82
+
83
+ @triton.testing.perf_report(
84
+ triton.testing.Benchmark(
85
+ x_names=["total_tokens"],
86
+ x_vals=total_tokens_list,
87
+ line_arg="provider",
88
+ line_vals=["cuda_kernel"],
89
+ line_names=["cp_gather_fp8 (CUDA)"],
90
+ styles=[("green", "-")],
91
+ ylabel="Latency (us)",
92
+ plot_name=f"cp_gather_fp8-{label}-bs{BLOCK_SIZE}",
93
+ args={"num_reqs": num_reqs},
94
+ )
95
+ )
96
+ def bench_fn(total_tokens, provider, num_reqs):
97
+ cache, dst, block_table, seq_lens_t, ws_starts = make_inputs(
98
+ total_tokens, num_reqs, BLOCK_SIZE
99
+ )
100
+
101
+ quantiles = [0.5, 0.2, 0.8]
102
+
103
+ ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
104
+ lambda: ops.cp_gather_and_upconvert_fp8_kv_cache(
105
+ cache, dst, block_table, seq_lens_t, ws_starts, num_reqs
106
+ ),
107
+ quantiles=quantiles,
108
+ rep=500,
109
+ )
110
+
111
+ return ms * 1000, max_ms * 1000, min_ms * 1000 # us
112
+
113
+ seq_len_per_req = total_tokens_list[0] // num_reqs
114
+ seq_len_per_req_max = total_tokens_list[-1] // num_reqs
115
+ print(
116
+ f"\n--- {label}: {num_reqs} request(s), "
117
+ f"~{seq_len_per_req}-{seq_len_per_req_max} tokens/req ---"
118
+ )
119
+ bench_fn.run(print_data=True, save_path=save_path)
120
+
121
+
122
+if __name__ == "__main__":
123
+ parser = argparse.ArgumentParser(
124
+ description="Benchmark cp_gather_and_upconvert_fp8_kv_cache"
125
+ )
126
+ parser.add_argument(
127
+ "--save-path",
128
+ type=str,
129
+ default=None,
130
+ help="Path to save benchmark results as CSV",
131
+ )
132
+ args = parser.parse_args()
133
+
134
+ # Print data volume info for bandwidth analysis
135
+ read_per_token = ENTRY_BYTES # 656 bytes from cache
136
+ write_per_token = HEAD_DIM * 2 # 576 * 2 = 1152 bytes to workspace
137
+ total_per_token = read_per_token + write_per_token # 1808 bytes
138
+
139
+ print("\n" + "=" * 70)
140
+ print("CP_GATHER_AND_UPCONVERT_FP8_KV_CACHE BENCHMARKS")
141
+ print("=" * 70)
142
+ print(f"Cache entry: {ENTRY_BYTES} bytes (512 FP8 + 16 scales + 128 RoPE)")
143
+ print(f"Output row: {HEAD_DIM} BF16 = {HEAD_DIM * 2} bytes")
144
+ print(f"Per token: {total_per_token} bytes (read + write)")
145
+ print(f"Block size: {BLOCK_SIZE} tokens/block")
146
+ print("=" * 70)
147
+
148
+ for label, num_reqs, total_tokens_list in SCENARIOS:
149
+ bench_scenario(label, num_reqs, total_tokens_list, args.save_path)
150
+
151
+ print("\n" + "=" * 70)
152
+ print("Benchmarking complete!")
153
+ print("=" * 70)
csrc/cache_kernels.cu
+62
-69
@@ -995,75 +995,67 @@ namespace vllm {
995
// Similar to cp_gather_cache but specifically for FP8->BF16 conversion
996
__global__ void cp_gather_and_upconvert_fp8_kv_cache(
997
const uint8_t* __restrict__ src_cache, // [NUM_BLOCKS, BLOCK_SIZE, 656]
998
- __nv_bfloat16* __restrict__ dst, // [TOT_TOKENS, 576]
999
- const int32_t* __restrict__ block_table, // [BATCH, BLOCK_INDICES]
1000
- const int32_t* __restrict__ seq_lens, // [BATCH]
1001
- const int32_t* __restrict__ workspace_starts, // [BATCH]
1002
- const int32_t block_size, const int32_t head_dim,
1003
- const int64_t block_table_stride, const int64_t cache_block_stride,
1004
- const int64_t cache_entry_stride, const int64_t dst_entry_stride) {
1005
- const int64_t bid = blockIdx.x; // Batch ID
1006
- const int32_t num_splits = gridDim.y;
1007
- const int32_t split = blockIdx.y;
1008
- const int32_t seq_start = workspace_starts[bid];
1009
- const int32_t seq_len = seq_lens[bid];
1010
- const int32_t tot_slots = seq_len;
1011
- const int32_t split_slots = cuda_utils::ceil_div(tot_slots, num_splits);
998
+ __nv_bfloat16* __restrict__ dst, // [total_tokens, 576]
999
+ const int32_t* __restrict__ block_table, // [num_reqs, BLOCK_INDICES]
1000
+ const int32_t* __restrict__ workspace_starts, // [num_reqs]
1001
+ const int32_t num_reqs, const int32_t block_size,
1002
+ const int32_t total_tokens, const int64_t block_table_stride,
1003
+ const int64_t cache_block_stride, const int64_t cache_entry_stride,
1004
+ const int64_t dst_entry_stride) {
1005
+ const int flat_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) >> 5;
1006
+ if (flat_warp_id >= total_tokens) return;
1007
+ const int lane_id = threadIdx.x & 31;
1008
+
1009
+ // Binary search to find which request owns this output token
1010
+ int lo = 0, hi = num_reqs - 1;
1011
+ while (lo < hi) {
1012
+ int mid = (lo + hi + 1) >> 1;
1013
+ if (workspace_starts[mid] <= flat_warp_id)
1014
+ lo = mid;
1015
+ else
1016
+ hi = mid - 1;
1017
+ }
1018
+ const int req_id = lo;
1019
1013
- const int32_t split_start = split * split_slots;
1014
- const int32_t split_end = min((split + 1) * split_slots, tot_slots);
1020
+ // Compute physical token address via block table
1021
+ const int out_token_id = flat_warp_id;
1022
+ const int token_offset = out_token_id - workspace_starts[req_id];
1023
+ const int cache_block_idx = token_offset / block_size;
1024
+ const int offset_in_block = token_offset % block_size;
1025
+ const int physical_block =
1026
+ block_table[req_id * block_table_stride + cache_block_idx];
1027
1016
- const bool is_active_split = (split_start < tot_slots);
1028
+ const uint8_t* token_ptr = src_cache + physical_block * cache_block_stride +
1029
+ offset_in_block * cache_entry_stride;
1030
1018
- if (!is_active_split) return;
1031
+ const int4* nope_src = reinterpret_cast<const int4*>(token_ptr);
1032
+ const int4 fp8_data = nope_src[lane_id];
1033
1020
- // Adjust the pointer for the block_table for this batch
1021
- const int32_t batch_offset = bid * block_table_stride;
1022
- int32_t offset = split_start;
1023
- int32_t offset_div = offset / block_size;
1024
- offset = offset % block_size;
1025
- const int32_t* batch_block_table = block_table + batch_offset;
1034
+ const float* scales_ptr = reinterpret_cast<const float*>(token_ptr + 512);
1035
+ const float scale = scales_ptr[lane_id >> 3];
1036
1027
- // Adjust dst pointer based on the cumulative sequence lengths
1028
- dst += seq_start * dst_entry_stride;
1029
-
1030
- const int tid = threadIdx.x;
1037
+ const uint2 fp8_lo = make_uint2(fp8_data.x, fp8_data.y);
1038
+ const uint2 fp8_hi = make_uint2(fp8_data.z, fp8_data.w);
1039
+#ifdef USE_ROCM
1040
+ const bf16_8_t bf16_lo =
1041
+ fp8::scaled_vec_conversion<bf16_8_t, uint2>(fp8_lo, scale);
1042
+ const bf16_8_t bf16_hi =
1043
+ fp8::scaled_vec_conversion<bf16_8_t, uint2>(fp8_hi, scale);
1044
+#else
1045
+ const bf16_8_t bf16_lo =
1046
+ fp8::scaled_vec_conversion<bf16_8_t, uint2>(fp8_lo, scale, __NV_E4M3);
1047
+ const bf16_8_t bf16_hi =
1048
+ fp8::scaled_vec_conversion<bf16_8_t, uint2>(fp8_hi, scale, __NV_E4M3);
1049
+#endif
1050
1032
- // Process each token in this split
1033
- for (int pid = split_start; pid < split_end; ++pid) {
1034
- auto block_id = batch_block_table[offset_div];
1035
- const uint8_t* token_ptr =
1036
- src_cache + block_id * cache_block_stride + offset * cache_entry_stride;
1037
- __nv_bfloat16* dst_ptr = dst + pid * dst_entry_stride;
1038
-
1039
- // FP8 format: 512 bytes fp8 + 16 bytes scales + 128 bytes rope (64 bf16)
1040
- const uint8_t* no_pe_ptr = token_ptr;
1041
- const float* scales_ptr = reinterpret_cast<const float*>(token_ptr + 512);
1042
- const __nv_bfloat16* rope_ptr =
1043
- reinterpret_cast<const __nv_bfloat16*>(token_ptr + 512 + 16);
1044
-
1045
- // Parallelize fp8 dequant (512 elements) and rope copy (64 elements)
1046
- if (tid < 512) {
1047
- // FP8 dequantization
1048
- const int tile = tid >> 7; // each tile is 128 elements
1049
- const float scale = scales_ptr[tile];
1050
- const uint8_t val = no_pe_ptr[tid];
1051
- dst_ptr[tid] =
1052
- fp8::scaled_convert<__nv_bfloat16, uint8_t,
1053
- vllm::Fp8KVCacheDataType::kFp8E4M3>(val, scale);
1054
- } else if (tid < 576) {
1055
- // Rope copy (64 bf16 elements)
1056
- const int rope_idx = tid - 512;
1057
- dst_ptr[512 + rope_idx] = rope_ptr[rope_idx];
1058
- }
1051
+ __nv_bfloat16* dst_ptr = dst + out_token_id * dst_entry_stride;
1052
+ int4* nope_dst = reinterpret_cast<int4*>(dst_ptr) + lane_id * 2;
1053
+ nope_dst[0] = *reinterpret_cast<const int4*>(&bf16_lo);
1054
+ nope_dst[1] = *reinterpret_cast<const int4*>(&bf16_hi);
1055
1060
- // Move to next token
1061
- offset += 1;
1062
- if (offset == block_size) {
1063
- offset_div += 1;
1064
- offset = 0;
1065
- }
1066
- }
1056
+ const int* rope_src = reinterpret_cast<const int*>(token_ptr + 528);
1057
+ int* rope_dst = reinterpret_cast<int*>(dst_ptr + 512);
1058
+ rope_dst[lane_id] = rope_src[lane_id];
1059
}
1060
1061
template <typename scalar_t>
@@ -1257,15 +1249,16 @@ void cp_gather_and_upconvert_fp8_kv_cache(
1249
src_ptr = reinterpret_cast<const uint8_t*>(src_cache.data_ptr());
1250
}
1251
1260
- // Decide on the number of splits based on the batch size
1261
- int num_splits = batch_size > 128 ? 2 : batch_size > 64 ? 4 : 16;
1262
- dim3 grid(batch_size, num_splits);
1263
- dim3 block(576);
1252
+ const int total_tokens = dst.size(0);
1253
+ constexpr int warps_per_block = 8;
1254
+ const int grid_size = (total_tokens + warps_per_block - 1) / warps_per_block;
1255
+ const int block_size_threads = warps_per_block * 32; // 256 threads
1256
1265
- vllm::cp_gather_and_upconvert_fp8_kv_cache<<<grid, block, 0, stream>>>(
1257
+ vllm::cp_gather_and_upconvert_fp8_kv_cache<<<grid_size, block_size_threads, 0,
1258
+ stream>>>(
1259
src_ptr, reinterpret_cast<__nv_bfloat16*>(dst.data_ptr()),
1267
- block_table.data_ptr<int32_t>(), seq_lens.data_ptr<int32_t>(),
1268
- workspace_starts.data_ptr<int32_t>(), block_size, head_dim,
1260
+ block_table.data_ptr<int32_t>(), workspace_starts.data_ptr<int32_t>(),
1261
+ static_cast<int32_t>(batch_size), block_size, total_tokens,
1262
block_table_stride, cache_block_stride, cache_entry_stride,
1263
dst_entry_stride);
1264
}
tests/kernels/test_cp_gather_fp8.py
new
+363
@@ -0,0 +1,363 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import math
4
+
5
+import pytest
6
+import torch
7
+
8
+from vllm import _custom_ops as ops
9
+
10
+# DeepSeek V3 MLA dimensions
11
+NOPE_DIM = 512 # NoPE latent dimension (FP8 quantized in cache)
12
+ROPE_DIM = 64 # RoPE dimension (stored as BF16 in cache)
13
+NUM_TILES = 4 # NOPE_DIM / GROUP_SIZE = 512 / 128
14
+GROUP_SIZE = 128 # FP8 quantization group size (one scale per group)
15
+ENTRY_BYTES = 656 # 512 (FP8) + 16 (4×float32 scales) + 128 (64×BF16 RoPE)
16
+
17
+
18
+def _build_test_case(seq_lens, block_size, seed=42):
19
+ """Build a synthetic FP8 cache and compute the expected BF16 output.
20
+
21
+ This simulates what concat_and_cache_ds_mla_kernel writes into the
22
+ KV cache, then computes what cp_gather_and_upconvert should produce.
23
+
24
+ Args:
25
+ seq_lens: List of sequence lengths, one per request.
26
+ block_size: Number of tokens per physical cache block.
27
+ seed: Random seed for reproducibility.
28
+
29
+ Returns:
30
+ Tuple of (cache, block_table, seq_lens_t, workspace_starts_t,
31
+ num_reqs, total_tokens, expected_output).
32
+ """
33
+ torch.manual_seed(seed)
34
+
35
+ num_reqs = len(seq_lens)
36
+ total_tokens = sum(seq_lens)
37
+
38
+ # workspace_starts[r] = sum of seq_lens[0..r-1]
39
+ # This tells the kernel where in the output buffer each request's
40
+ # gathered tokens should be written.
41
+ workspace_starts = []
42
+ s = 0
43
+ for sl in seq_lens:
44
+ workspace_starts.append(s)
45
+ s += sl
46
+
47
+ # How many physical cache blocks each request needs
48
+ blocks_per_req = [math.ceil(s / block_size) for s in seq_lens]
49
+ total_blocks = sum(blocks_per_req)
50
+ max_blocks = max(blocks_per_req)
51
+
52
+ # Block table maps (request, logical_block_idx) -> physical_block_id.
53
+ # Here we assign blocks contiguously: request 0 gets blocks [0, 1, ...],
54
+ # request 1 gets the next set, etc.
55
+ block_table = torch.zeros(num_reqs, max_blocks, dtype=torch.int32, device="cuda")
56
+ block_idx = 0
57
+ for r in range(num_reqs):
58
+ for b in range(blocks_per_req[r]):
59
+ block_table[r, b] = block_idx
60
+ block_idx += 1
61
+
62
+ # The raw paged cache: [num_blocks, block_size, 656] as uint8
63
+ cache = torch.zeros(
64
+ total_blocks, block_size, ENTRY_BYTES, dtype=torch.uint8, device="cuda"
65
+ )
66
+ # Expected kernel output: [total_tokens, 576] as BF16
67
+ expected = torch.zeros(
68
+ total_tokens, NOPE_DIM + ROPE_DIM, dtype=torch.bfloat16, device="cuda"
69
+ )
70
+
71
+ # Fill each token's cache entry and compute expected output
72
+ for r in range(num_reqs):
73
+ for t in range(seq_lens[r]):
74
+ out_idx = workspace_starts[r] + t
75
+ # Map token position -> (physical_block, offset_within_block)
76
+ phys = block_table[r, t // block_size].item()
77
+ off = t % block_size
78
+
79
+ # --- NoPE section: 4 tiles of 128 FP8 values, each with a scale ---
80
+ for tile in range(NUM_TILES):
81
+ start = tile * GROUP_SIZE
82
+
83
+ # Generate random data and quantize to FP8 e4m3
84
+ fp8_vals = torch.randn(GROUP_SIZE, device="cuda").to(
85
+ torch.float8_e4m3fn
86
+ )
87
+ # Pack FP8 bytes into cache at bytes [start : start+128]
88
+ cache[phys, off, start : start + GROUP_SIZE] = fp8_vals.view(
89
+ torch.uint8
90
+ )
91
+
92
+ # Random positive scale in [0.1, 2.1]
93
+ scale = (torch.rand(1, device="cuda") * 2.0 + 0.1).item()
94
+ scale_t = torch.tensor([scale], dtype=torch.float32, device="cuda")
95
+ # Pack scale as 4 raw bytes at bytes [512 + tile*4 : ...]
96
+ cache[phys, off, NOPE_DIM + tile * 4 : NOPE_DIM + (tile + 1) * 4] = (
97
+ scale_t.view(torch.uint8)
98
+ )
99
+
100
+ # Reference dequant: fp8 -> float32, multiply scale, -> bf16.
101
+ # This matches the CUDA path: fp8 -> half -> float * scale -> bf16.
102
+ # (fp8 -> half is exact, half -> float is exact, so fp8 -> float
103
+ # gives the same result regardless of intermediate type.)
104
+ expected[out_idx, start : start + GROUP_SIZE] = (
105
+ fp8_vals.float() * scale
106
+ ).bfloat16()
107
+
108
+ # --- RoPE section: 64 BF16 values, direct copy (no dequant) ---
109
+ rope = torch.randn(ROPE_DIM, dtype=torch.bfloat16, device="cuda")
110
+ # Pack RoPE bytes into cache at bytes [528 : 656]
111
+ cache[phys, off, NOPE_DIM + 16 :] = rope.view(torch.uint8)
112
+ # Expected output: exact copy
113
+ expected[out_idx, NOPE_DIM:] = rope
114
+
115
+ seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device="cuda")
116
+ workspace_starts_t = torch.tensor(
117
+ workspace_starts, dtype=torch.int32, device="cuda"
118
+ )
119
+
120
+ return (
121
+ cache,
122
+ block_table,
123
+ seq_lens_t,
124
+ workspace_starts_t,
125
+ num_reqs,
126
+ total_tokens,
127
+ expected,
128
+ )
129
+
130
+
131
+def _build_test_case_fast(seq_lens, block_size, seed=42):
132
+ """Vectorized test-case builder for large sequence lengths.
133
+
134
+ Same logic as _build_test_case but uses tensor operations instead of
135
+ per-token Python loops, making it practical for seq_lens up to 128K+.
136
+ """
137
+ torch.manual_seed(seed)
138
+
139
+ num_reqs = len(seq_lens)
140
+ total_tokens = sum(seq_lens)
141
+
142
+ workspace_starts = []
143
+ s = 0
144
+ for sl in seq_lens:
145
+ workspace_starts.append(s)
146
+ s += sl
147
+
148
+ blocks_per_req = [math.ceil(sl / block_size) for sl in seq_lens]
149
+ total_blocks = sum(blocks_per_req)
150
+ max_blocks = max(blocks_per_req)
151
+
152
+ # Contiguous block allocation
153
+ block_table = torch.zeros(num_reqs, max_blocks, dtype=torch.int32, device="cuda")
154
+ block_idx = 0
155
+ for r in range(num_reqs):
156
+ for b in range(blocks_per_req[r]):
157
+ block_table[r, b] = block_idx
158
+ block_idx += 1
159
+
160
+ cache = torch.zeros(
161
+ total_blocks, block_size, ENTRY_BYTES, dtype=torch.uint8, device="cuda"
162
+ )
163
+
164
+ # Generate all data vectorized
165
+ nope_fp8 = torch.randn(total_tokens, NOPE_DIM, device="cuda").to(
166
+ torch.float8_e4m3fn
167
+ )
168
+ scales = (torch.rand(total_tokens, NUM_TILES, device="cuda") * 2.0 + 0.1).float()
169
+ rope = torch.randn(total_tokens, ROPE_DIM, dtype=torch.bfloat16, device="cuda")
170
+
171
+ # Compute expected output vectorized (same dequant logic as kernel)
172
+ expected = torch.zeros(
173
+ total_tokens, NOPE_DIM + ROPE_DIM, dtype=torch.bfloat16, device="cuda"
174
+ )
175
+ for tile in range(NUM_TILES):
176
+ start = tile * GROUP_SIZE
177
+ expected[:, start : start + GROUP_SIZE] = (
178
+ nope_fp8[:, start : start + GROUP_SIZE].float() * scales[:, tile : tile + 1]
179
+ ).bfloat16()
180
+ expected[:, NOPE_DIM:] = rope
181
+
182
+ # Build per-token cache entries as [total_tokens, 656] uint8
183
+ token_data = torch.zeros(
184
+ total_tokens, ENTRY_BYTES, dtype=torch.uint8, device="cuda"
185
+ )
186
+ token_data[:, :NOPE_DIM] = nope_fp8.view(torch.uint8)
187
+ token_data[:, NOPE_DIM : NOPE_DIM + 16] = scales.view(torch.uint8)
188
+ token_data[:, NOPE_DIM + 16 :] = rope.view(torch.uint8)
189
+
190
+ # Scatter into paged cache (loop over requests, not tokens)
191
+ block_start = 0
192
+ for r in range(num_reqs):
193
+ sl = seq_lens[r]
194
+ nb = blocks_per_req[r]
195
+ ws = workspace_starts[r]
196
+ flat_cache = cache[block_start : block_start + nb].reshape(-1, ENTRY_BYTES)
197
+ flat_cache[:sl] = token_data[ws : ws + sl]
198
+ block_start += nb
199
+
200
+ seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device="cuda")
201
+ workspace_starts_t = torch.tensor(
202
+ workspace_starts, dtype=torch.int32, device="cuda"
203
+ )
204
+
205
+ return (
206
+ cache,
207
+ block_table,
208
+ seq_lens_t,
209
+ workspace_starts_t,
210
+ num_reqs,
211
+ total_tokens,
212
+ expected,
213
+ )
214
+
215
+
216
+@pytest.mark.parametrize(
217
+ "seq_lens,block_size",
218
+ [
219
+ # Production block_size=64 (only supported value for FlashMLA sparse).
220
+ # Realistic prefill scenarios with varying request counts.
221
+ ([1], 64), # single token edge case
222
+ ([64], 64), # 1 req, exactly one block
223
+ ([128], 64), # 1 req, crosses block boundary
224
+ ([512], 64), # 1 req, longer prefill
225
+ ([256, 128, 384], 64), # 3 reqs, varying lengths
226
+ ([128] * 4, 64), # 4 reqs, equal lengths
227
+ ([64] * 16, 64), # 16 reqs, shorter prefills
228
+ ],
229
+)
230
+def test_cp_gather_and_upconvert_fp8_kv_cache(seq_lens, block_size):
231
+ """Core correctness test: build cache, run kernel, compare output."""
232
+ (
233
+ cache,
234
+ block_table,
235
+ seq_lens_t,
236
+ workspace_starts_t,
237
+ num_reqs,
238
+ total_tokens,
239
+ expected,
240
+ ) = _build_test_case(seq_lens, block_size)
241
+
242
+ dst = torch.zeros(
243
+ total_tokens, NOPE_DIM + ROPE_DIM, dtype=torch.bfloat16, device="cuda"
244
+ )
245
+
246
+ ops.cp_gather_and_upconvert_fp8_kv_cache(
247
+ cache, dst, block_table, seq_lens_t, workspace_starts_t, num_reqs
248
+ )
249
+
250
+ # NoPE: fp8 dequant has rounding error, so we allow small tolerance.
251
+ # The fp8 -> float -> bf16 path can differ by up to ~1 ULP of bf16.
252
+ torch.testing.assert_close(
253
+ dst[:, :NOPE_DIM], expected[:, :NOPE_DIM], atol=1e-3, rtol=1e-2
254
+ )
255
+
256
+ # RoPE: pure bf16 copy, must be bit-exact
257
+ assert torch.equal(dst[:, NOPE_DIM:], expected[:, NOPE_DIM:])
258
+
259
+
260
+def test_cp_gather_fp8_shuffled_blocks():
261
+ """Test that the kernel correctly follows the block table when
262
+ physical blocks are non-contiguous and out of order.
263
+
264
+ Here we allocate 4 physical blocks but map the request's 2 logical
265
+ blocks to physical blocks [3, 1] (reversed, with gaps).
266
+ """
267
+ torch.manual_seed(123)
268
+ block_size = 4
269
+ seq_lens = [8] # needs 2 blocks (tokens 0-3 in block 0, 4-7 in block 1)
270
+ total_tokens = 8
271
+
272
+ # 4 physical blocks, but only blocks 3 and 1 are used (in that order).
273
+ # Tokens 0-3 -> physical block 3, tokens 4-7 -> physical block 1.
274
+ num_phys_blocks = 4
275
+ cache = torch.zeros(
276
+ num_phys_blocks, block_size, ENTRY_BYTES, dtype=torch.uint8, device="cuda"
277
+ )
278
+ block_table = torch.tensor([[3, 1]], dtype=torch.int32, device="cuda")
279
+ workspace_starts = torch.tensor([0], dtype=torch.int32, device="cuda")
280
+ seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device="cuda")
281
+
282
+ expected = torch.zeros(
283
+ total_tokens, NOPE_DIM + ROPE_DIM, dtype=torch.bfloat16, device="cuda"
284
+ )
285
+
286
+ # Fill cache at the shuffled physical locations
287
+ for t in range(total_tokens):
288
+ # Follow the same block_table lookup the kernel will use
289
+ phys = block_table[0, t // block_size].item()
290
+ off = t % block_size
291
+
292
+ for tile in range(NUM_TILES):
293
+ start = tile * GROUP_SIZE
294
+ fp8_vals = torch.randn(GROUP_SIZE, device="cuda").to(torch.float8_e4m3fn)
295
+ cache[phys, off, start : start + GROUP_SIZE] = fp8_vals.view(torch.uint8)
296
+
297
+ # Use a fixed scale to keep this test simple
298
+ scale = 1.5
299
+ scale_t = torch.tensor([scale], dtype=torch.float32, device="cuda")
300
+ cache[phys, off, NOPE_DIM + tile * 4 : NOPE_DIM + (tile + 1) * 4] = (
301
+ scale_t.view(torch.uint8)
302
+ )
303
+
304
+ expected[t, start : start + GROUP_SIZE] = (
305
+ fp8_vals.float() * scale
306
+ ).bfloat16()
307
+
308
+ rope = torch.randn(ROPE_DIM, dtype=torch.bfloat16, device="cuda")
309
+ cache[phys, off, NOPE_DIM + 16 :] = rope.view(torch.uint8)
310
+ expected[t, NOPE_DIM:] = rope
311
+
312
+ dst = torch.zeros(
313
+ total_tokens, NOPE_DIM + ROPE_DIM, dtype=torch.bfloat16, device="cuda"
314
+ )
315
+
316
+ ops.cp_gather_and_upconvert_fp8_kv_cache(
317
+ cache, dst, block_table, seq_lens_t, workspace_starts, len(seq_lens)
318
+ )
319
+
320
+ torch.testing.assert_close(
321
+ dst[:, :NOPE_DIM], expected[:, :NOPE_DIM], atol=1e-3, rtol=1e-2
322
+ )
323
+ assert torch.equal(dst[:, NOPE_DIM:], expected[:, NOPE_DIM:])
324
+
325
+
326
+@pytest.mark.parametrize(
327
+ "seq_lens,block_size",
328
+ [
329
+ # Large sequence lengths matching end-to-end benchmark scenarios.
330
+ # Uses vectorized builder since per-token Python loops would be too slow.
331
+ ([8000], 64),
332
+ ([16000], 64),
333
+ ([32000], 64),
334
+ ([64000], 64),
335
+ ([96000], 64),
336
+ ([128000], 64),
337
+ ],
338
+)
339
+def test_cp_gather_fp8_large_seqlens(seq_lens, block_size):
340
+ """Correctness test with large sequence lengths matching benchmark
341
+ scenarios (8K-128K prefill)."""
342
+ (
343
+ cache,
344
+ block_table,
345
+ seq_lens_t,
346
+ workspace_starts_t,
347
+ num_reqs,
348
+ total_tokens,
349
+ expected,
350
+ ) = _build_test_case_fast(seq_lens, block_size)
351
+
352
+ dst = torch.zeros(
353
+ total_tokens, NOPE_DIM + ROPE_DIM, dtype=torch.bfloat16, device="cuda"
354
+ )
355
+
356
+ ops.cp_gather_and_upconvert_fp8_kv_cache(
357
+ cache, dst, block_table, seq_lens_t, workspace_starts_t, num_reqs
358
+ )
359
+
360
+ torch.testing.assert_close(
361
+ dst[:, :NOPE_DIM], expected[:, :NOPE_DIM], atol=1e-3, rtol=1e-2
362
+ )
363
+ assert torch.equal(dst[:, NOPE_DIM:], expected[:, NOPE_DIM:])