@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 3e41992fec

Fork 0

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 11, 15:00 UTC

[Attention] Use sparse prefill kernel for fp8 kv-cache in DeepSeek-v3.2 (#27532)

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

Lucas Wilkinson committed Dec 12, 2025 at 08:57 UTC 3e41992fecdc31ee60715bb350f18fec18ed6680

30 files changed +1372 -256

csrc/cache.h

+11 -1

@@ -1,6 +1,7 @@

1 #pragma once

3 #include <torch/all.h>

4 +#include <c10/util/Optional.h>

6 #include <map>

7 #include <vector>

@@ -58,6 +59,15 @@ void cp_gather_cache(

59 torch::Tensor const& cu_seq_lens, // [BATCH+1]

60 int64_t batch_size, std::optional<torch::Tensor> seq_starts = std::nullopt);

62 +// Gather and upconvert FP8 KV cache to BF16 workspace

63 +void cp_gather_and_upconvert_fp8_kv_cache(

64 + torch::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, 656]

65 + torch::Tensor const& dst, // [TOT_TOKENS, 576]

66 + torch::Tensor const& block_table, // [BATCH, BLOCK_INDICES]

67 + torch::Tensor const& seq_lens, // [BATCH]

68 + torch::Tensor const& workspace_starts, // [BATCH]

69 + int64_t batch_size);

70 +

71 // Indexer K quantization and cache function

72 void indexer_k_quant_and_cache(

73 torch::Tensor& k, // [num_tokens, head_dim]

@@ -72,4 +82,4 @@ void cp_gather_indexer_k_quant_cache(

82 torch::Tensor& dst_k, // [num_tokens, head_dim]

83 torch::Tensor& dst_scale, // [num_tokens, head_dim / quant_block_size * 4]

84 const torch::Tensor& block_table, // [batch_size, num_blocks]

75 - const torch::Tensor& cu_seq_lens); // [batch_size + 1]

\ No newline at end of file

85 + const torch::Tensor& cu_seq_lens); // [batch_size + 1]

csrc/cache_kernels.cu

+130 -1

@@ -2,6 +2,7 @@

2 #include <ATen/cuda/CUDAContext.h>

3 #include <c10/cuda/CUDAGuard.h>

4 #include <c10/cuda/CUDAException.h>

5 +#include <c10/util/Optional.h>

7 #include "cuda_utils.h"

8 #include "cuda_compat.h"

@@ -514,7 +515,8 @@ __global__ void indexer_k_quant_and_cache_kernel(

515 const int quant_block_size, // quantization block size

516 const int cache_block_size, // cache block size

517 const int cache_stride, // stride for each token in kv_cache

517 - const bool use_ue8m0 // use ue8m0 scale format

518 +

519 + const bool use_ue8m0 // use ue8m0 scale format

520 ) {

521 constexpr int VEC_SIZE = 4;

522 const int64_t token_idx = blockIdx.x;

@@ -1061,6 +1063,82 @@ void gather_and_maybe_dequant_cache(

1063 }

1064

1065 namespace vllm {

1066 +

1067 +// Gather and upconvert FP8 KV cache tokens to BF16 workspace

1068 +// Similar to cp_gather_cache but specifically for FP8->BF16 conversion

1069 +__global__ void cp_gather_and_upconvert_fp8_kv_cache(

1070 + const uint8_t* __restrict__ src_cache, // [NUM_BLOCKS, BLOCK_SIZE, 656]

1071 + __nv_bfloat16* __restrict__ dst, // [TOT_TOKENS, 576]

1072 + const int32_t* __restrict__ block_table, // [BATCH, BLOCK_INDICES]

1073 + const int32_t* __restrict__ seq_lens, // [BATCH]

1074 + const int32_t* __restrict__ workspace_starts, // [BATCH]

1075 + const int32_t block_size, const int32_t head_dim,

1076 + const int64_t block_table_stride, const int64_t cache_block_stride,

1077 + const int64_t cache_entry_stride, const int64_t dst_entry_stride) {

1078 + const int64_t bid = blockIdx.x; // Batch ID

1079 + const int32_t num_splits = gridDim.y;

1080 + const int32_t split = blockIdx.y;

1081 + const int32_t seq_start = workspace_starts[bid];

1082 + const int32_t seq_len = seq_lens[bid];

1083 + const int32_t tot_slots = seq_len;

1084 + const int32_t split_slots = cuda_utils::ceil_div(tot_slots, num_splits);

1085 +

1086 + const int32_t split_start = split * split_slots;

1087 + const int32_t split_end = min((split + 1) * split_slots, tot_slots);

1088 +

1089 + const bool is_active_split = (split_start < tot_slots);

1090 +

1091 + if (!is_active_split) return;

1092 +

1093 + // Adjust the pointer for the block_table for this batch

1094 + const int32_t batch_offset = bid * block_table_stride;

1095 + int32_t offset = split_start;

1096 + int32_t offset_div = offset / block_size;

1097 + offset = offset % block_size;

1098 + const int32_t* batch_block_table = block_table + batch_offset;

1099 +

1100 + // Adjust dst pointer based on the cumulative sequence lengths

1101 + dst += seq_start * dst_entry_stride;

1102 +

1103 + const int tid = threadIdx.x;

1104 +

1105 + // Process each token in this split

1106 + for (int pid = split_start; pid < split_end; ++pid) {

1107 + auto block_id = batch_block_table[offset_div];

1108 + const uint8_t* token_ptr =

1109 + src_cache + block_id * cache_block_stride + offset * cache_entry_stride;

1110 + __nv_bfloat16* dst_ptr = dst + pid * dst_entry_stride;

1111 +

1112 + // FP8 format: 512 bytes fp8 + 16 bytes scales + 128 bytes rope (64 bf16)

1113 + const uint8_t* no_pe_ptr = token_ptr;

1114 + const float* scales_ptr = reinterpret_cast<const float*>(token_ptr + 512);

1115 + const __nv_bfloat16* rope_ptr =

1116 + reinterpret_cast<const __nv_bfloat16*>(token_ptr + 512 + 16);

1117 +

1118 + // Parallelize fp8 dequant (512 elements) and rope copy (64 elements)

1119 + if (tid < 512) {

1120 + // FP8 dequantization

1121 + const int tile = tid >> 7; // each tile is 128 elements

1122 + const float scale = scales_ptr[tile];

1123 + const uint8_t val = no_pe_ptr[tid];

1124 + dst_ptr[tid] =

1125 + fp8::scaled_convert<__nv_bfloat16, uint8_t,

1126 + vllm::Fp8KVCacheDataType::kFp8E4M3>(val, scale);

1127 + } else if (tid < 576) {

1128 + // Rope copy (64 bf16 elements)

1129 + const int rope_idx = tid - 512;

1130 + dst_ptr[512 + rope_idx] = rope_ptr[rope_idx];

1131 + }

1132 +

1133 + // Move to next token

1134 + offset += 1;

1135 + if (offset == block_size) {

1136 + offset_div += 1;

1137 + offset = 0;

1138 + }

1139 + }

1140 +}

1141 +

1142 template <typename scalar_t>

1143 // Note(hc): The cp_gather_cache allows seq_starts to no longer be divisible by

1144 // block_size.

@@ -1202,6 +1280,57 @@ void cp_gather_cache(

1280 }

1281 }

1282

1283 +void cp_gather_and_upconvert_fp8_kv_cache(

1284 + torch::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, 656]

1285 + torch::Tensor const& dst, // [TOT_TOKENS, 576]

1286 + torch::Tensor const& block_table, // [BATCH, BLOCK_INDICES]

1287 + torch::Tensor const& seq_lens, // [BATCH]

1288 + torch::Tensor const& workspace_starts, // [BATCH]

1289 + int64_t batch_size) {

1290 + at::cuda::OptionalCUDAGuard device_guard(src_cache.device());

1291 + const cudaStream_t stream = at::cuda::getCurrentCUDAStream();

1292 +

1293 + int32_t block_size = src_cache.size(1);

1294 + int32_t head_dim = dst.size(1);

1295 +

1296 + TORCH_CHECK(block_table.dtype() == torch::kInt32,

1297 + "block_table must be int32");

1298 + TORCH_CHECK(seq_lens.dtype() == torch::kInt32, "seq_lens must be int32");

1299 + TORCH_CHECK(workspace_starts.dtype() == torch::kInt32,

1300 + "workspace_starts must be int32");

1301 +

1302 + TORCH_CHECK(src_cache.device() == dst.device(),

1303 + "src_cache and dst must be on the same device");

1304 + TORCH_CHECK(src_cache.device() == block_table.device(),

1305 + "src_cache and block_table must be on the same device");

1306 + TORCH_CHECK(src_cache.device() == seq_lens.device(),

1307 + "src_cache and seq_lens must be on the same device");

1308 + TORCH_CHECK(src_cache.device() == workspace_starts.device(),

1309 + "src_cache and workspace_starts must be on the same device");

1310 +

1311 + TORCH_CHECK(src_cache.dtype() == torch::kUInt8, "src_cache must be uint8");

1312 + TORCH_CHECK(dst.dtype() == torch::kBFloat16, "dst must be bfloat16");

1313 + TORCH_CHECK(head_dim == 576, "head_dim must be 576 for MLA");

1314 +

1315 + int64_t block_table_stride = block_table.stride(0);

1316 + int64_t cache_block_stride = src_cache.stride(0);

1317 + int64_t cache_entry_stride = src_cache.stride(1);

1318 + int64_t dst_entry_stride = dst.stride(0);

1319 +

1320 + // Decide on the number of splits based on the batch size

1321 + int num_splits = batch_size > 128 ? 2 : batch_size > 64 ? 4 : 16;

1322 + dim3 grid(batch_size, num_splits);

1323 + dim3 block(576);

1324 +

1325 + vllm::cp_gather_and_upconvert_fp8_kv_cache<<<grid, block, 0, stream>>>(

1326 + src_cache.data_ptr<uint8_t>(),

1327 + reinterpret_cast<__nv_bfloat16*>(dst.data_ptr()),

1328 + block_table.data_ptr<int32_t>(), seq_lens.data_ptr<int32_t>(),

1329 + workspace_starts.data_ptr<int32_t>(), block_size, head_dim,

1330 + block_table_stride, cache_block_stride, cache_entry_stride,

1331 + dst_entry_stride);

1332 +}

1333 +

1334 // Macro to dispatch the kernel based on the data type.

1335 #define CALL_INDEXER_K_QUANT_AND_CACHE(KV_T, CACHE_T, KV_DTYPE) \

1336 vllm::indexer_k_quant_and_cache_kernel<KV_T, CACHE_T, KV_DTYPE> \

csrc/torch_bindings.cpp

@@ -754,6 +754,13 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {

754 "Tensor cu_seq_lens, int batch_size, Tensor? seq_starts) -> ()");

755 cache_ops.impl("cp_gather_cache", torch::kCUDA, &cp_gather_cache);

756

757 + cache_ops.def(

758 + "cp_gather_and_upconvert_fp8_kv_cache(Tensor src_cache, Tensor! dst, "

759 + "Tensor block_table, Tensor seq_lens, Tensor workspace_starts, int "

760 + "batch_size) -> ()");

761 + cache_ops.impl("cp_gather_and_upconvert_fp8_kv_cache", torch::kCUDA,

762 + &cp_gather_and_upconvert_fp8_kv_cache);

763 +

764 cache_ops.def(

765 "indexer_k_quant_and_cache(Tensor k, Tensor! kv_cache, Tensor "

766 "slot_mapping, "

tests/conftest.py

+21

@@ -202,6 +202,27 @@ def cleanup_fixture(should_do_global_cleanup_after_test: bool):

202 cleanup_dist_env_and_memory()

203

204

205 +@pytest.fixture

206 +def workspace_init():

207 + """Initialize the workspace manager for tests that need it.

208 +

209 + This fixture initializes the workspace manager with a CUDA device

210 + if available, and resets it after the test completes. Tests that

211 + create a full vLLM engine should NOT use this fixture as the engine

212 + will initialize the workspace manager itself.

213 + """

214 + from vllm.v1.worker.workspace import (

215 + init_workspace_manager,

216 + reset_workspace_manager,

217 + )

218 +

219 + if torch.cuda.is_available():

220 + device = torch.device("cuda:0")

221 + init_workspace_manager(device)

222 + yield

223 + reset_workspace_manager()

224 +

225 +

226 @pytest.fixture(autouse=True)

227 def dynamo_reset():

228 yield

tests/kernels/moe/test_batched_deepgemm.py

+1 -1

@@ -27,7 +27,7 @@ BLOCK_SIZE = [128, 128]

27 @pytest.mark.parametrize("N", [512, 1024]) # intermediate dim per expert

28 @pytest.mark.parametrize("topk", [2, 4])

29 def test_batched_deepgemm_vs_triton(

30 - E: int, T: int, K: int, N: int, topk: int, monkeypatch

30 + E: int, T: int, K: int, N: int, topk: int, monkeypatch, workspace_init

31 ):

32 """Compare BatchedDeepGemmExperts to BatchedTritonExperts."""

tests/kernels/moe/test_batched_moe.py

@@ -248,6 +248,7 @@ def test_fused_moe_batched_experts(

248 per_act_token_quant: bool,

249 block_shape: list[int] | None,

250 input_scales: bool,

251 + workspace_init,

252 ):

253 """Note: float8_e4m3fn is not supported on CUDA architecture < 89,

254 and those tests will be skipped on unsupported hardware."""

tests/kernels/moe/test_block_fp8.py

+1 -1

@@ -137,7 +137,7 @@ def setup_cuda():

137 @pytest.mark.parametrize("seed", SEEDS)

138 @torch.inference_mode()

139 def test_w8a8_block_fp8_fused_moe(

140 - M, N, K, E, topk, block_size, dtype, seed, monkeypatch

140 + M, N, K, E, topk, block_size, dtype, seed, monkeypatch, workspace_init

141 ):

142 if topk > E:

143 pytest.skip(f"Skipping test; topk={topk} > E={E}")

tests/kernels/moe/test_cutlass_moe.py

+25 -2

@@ -274,6 +274,7 @@ def test_cutlass_moe_8_bit_no_graph(

274 per_act_token: bool,

275 per_out_ch: bool,

276 monkeypatch,

277 + workspace_init,

278 ep_size: int | None = None,

279 ):

280 current_platform.seed_everything(7)

@@ -329,6 +330,7 @@ def test_cutlass_moe_8_bit_cuda_graph(

330 per_act_token: bool,

331 per_out_ch: bool,

332 monkeypatch,

333 + workspace_init,

334 ):

335 current_platform.seed_everything(7)

336 monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", "8192")

@@ -385,9 +387,19 @@ def test_cutlass_moe_8_bit_EP(

387 per_out_channel: bool,

388 ep_size: int,

389 monkeypatch,

390 + workspace_init,

391 ):

392 test_cutlass_moe_8_bit_no_graph(

390 - m, n, k, e, topk, per_act_token, per_out_channel, monkeypatch, ep_size

393 + m,

394 + n,

395 + k,

396 + e,

397 + topk,

398 + per_act_token,

399 + per_out_channel,

400 + monkeypatch,

401 + workspace_init,

402 + ep_size,

403 )

404

405

@@ -419,9 +431,19 @@ def test_cutlass_moe_8_bit_EP_large(

431 per_out_channel: bool,

432 ep_size: int,

433 monkeypatch,

434 + workspace_init,

435 ):

436 test_cutlass_moe_8_bit_no_graph(

424 - m, n, k, e, topk, per_act_token, per_out_channel, monkeypatch, ep_size

437 + m,

438 + n,

439 + k,

440 + e,

441 + topk,

442 + per_act_token,

443 + per_out_channel,

444 + monkeypatch,

445 + workspace_init,

446 + ep_size,

447 )

448

449

@@ -445,6 +467,7 @@ def test_run_cutlass_moe_fp8(

467 per_act_token: bool,

468 per_out_channel: bool,

469 ep_size: int,

470 + workspace_init,

471 ):

472 current_platform.seed_everything(7)

473 with set_current_vllm_config(vllm_config):

tests/kernels/moe/test_deepep_deepgemm_moe.py

@@ -29,6 +29,7 @@ from vllm.utils.deep_gemm import (

29 is_deep_gemm_supported,

30 )

31 from vllm.utils.import_utils import has_deep_ep, has_deep_gemm

32 +from vllm.v1.worker.workspace import init_workspace_manager

34 from ...utils import multi_gpu_test

35 from .parallel_utils import ProcessGroupInfo, parallel_launch

@@ -363,6 +364,9 @@ def _test_deepep_deepgemm_moe(

364 w1_scale: torch.Tensor,

365 w2_scale: torch.Tensor,

366 ):

367 + device = torch.device(f"cuda:{pgi.local_rank}")

368 + init_workspace_manager(device)

369 +

370 current_platform.seed_everything(pgi.rank)

371

372 w1 = w1.to(device=torch.cuda.current_device())

@@ -445,6 +449,7 @@ def test_ht_deepep_deepgemm_moe(

449 topk: int,

450 world_dp_size: tuple[int, int],

451 disable_deepgemm_ue8m0,

452 + workspace_init,

453 ):

454 """

455 Tests for High-Throughput DeepEP + DeepGemm integration.

@@ -518,6 +523,7 @@ def test_ll_deepep_deepgemm_moe(

523 block_size: list[int],

524 world_dp_size: tuple[int, int],

525 disable_deepgemm_ue8m0,

526 + workspace_init,

527 ):

528 """

529 Tests for Low-Latency DeepEP + DeepGemm integration.

tests/kernels/moe/test_deepep_moe.py

@@ -22,6 +22,7 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (

22 )

23 from vllm.platforms import current_platform

24 from vllm.utils.import_utils import has_deep_ep

25 +from vllm.v1.worker.workspace import init_workspace_manager

27 from ...utils import multi_gpu_test

28 from .parallel_utils import ProcessGroupInfo, parallel_launch

@@ -342,6 +343,9 @@ def _deep_ep_moe(

343 use_fp8_dispatch: bool,

344 per_act_token_quant: bool,

345 ):

346 + device = torch.device(f"cuda:{pgi.local_rank}")

347 + init_workspace_manager(device)

348 +

349 if not low_latency_mode:

350 assert not use_fp8_dispatch, (

351 "FP8 dispatch interface is available only in low-latency mode"

@@ -437,6 +441,7 @@ def test_deep_ep_moe(

441 topk: int,

442 world_dp_size: tuple[int, int],

443 per_act_token_quant: bool,

444 + workspace_init,

445 ):

446 low_latency_mode = False

447 use_fp8_dispatch = False

@@ -492,6 +497,7 @@ def test_low_latency_deep_ep_moe(

497 topk: int,

498 world_dp_size: tuple[int, int],

499 use_fp8_dispatch: bool,

500 + workspace_init,

501 ):

502 low_latency_mode = True

503

tests/kernels/moe/test_deepgemm.py

+1 -1

@@ -143,7 +143,7 @@ NUM_EXPERTS = [32]

143 @pytest.mark.parametrize("topk", TOPKS)

144 @pytest.mark.parametrize("num_experts", NUM_EXPERTS)

145 @pytest.mark.skipif(not is_deep_gemm_supported(), reason="Requires deep_gemm kernels")

146 -def test_deepgemm_vs_triton(m, n, k, topk, num_experts, monkeypatch):

146 +def test_deepgemm_vs_triton(m, n, k, topk, num_experts, monkeypatch, workspace_init):

147 with monkeypatch.context() as mp:

148 mp.setenv("VLLM_USE_DEEP_GEMM", "1")

149

tests/kernels/moe/test_flashinfer.py

@@ -206,6 +206,7 @@ def test_flashinfer_cutlass_moe_fp8_no_graph(

206 topk: int,

207 activation: str,

208 monkeypatch,

209 + workspace_init,

210 ):

211 current_platform.seed_everything(7)

212 monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", "8192")

tests/kernels/moe/test_flashinfer_moe.py

+8 -1

@@ -51,7 +51,14 @@ MNK_FACTORS = [

51 @pytest.mark.parametrize("activation", ["silu_and_mul", "relu2"])

52 @torch.inference_mode()

53 def test_flashinfer_fp4_moe_no_graph(

54 - m: int, n: int, k: int, e: int, topk: int, dtype: torch.dtype, activation: str

54 + m: int,

55 + n: int,

56 + k: int,

57 + e: int,

58 + topk: int,

59 + dtype: torch.dtype,

60 + activation: str,

61 + workspace_init,

62 ):

63 current_platform.seed_everything(7)

64 with set_current_vllm_config(

tests/kernels/moe/test_gpt_oss_triton_kernels.py

+1 -1

@@ -269,7 +269,7 @@ class Case:

269 )

270 @pytest.mark.parametrize("num_token", [2])

271 @pytest.mark.parametrize("tp", [1, 2, 4, 8])

272 -def test_equiv(num_token, a_dtype, w_dtype, tp):

272 +def test_equiv(num_token, a_dtype, w_dtype, tp, workspace_init):

273 from triton_kernels.tensor_details import layout

274

275 if not hasattr(layout, "make_default_matmul_mxfp4_w_layout"):

tests/kernels/moe/test_modular_kernel_combinations.py

@@ -16,6 +16,7 @@ from vllm.platforms import current_platform

16 from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe

17 from vllm.utils.import_utils import has_deep_ep, has_deep_gemm, has_pplx

18 from vllm.utils.torch_utils import cuda_device_count_stateless

19 +from vllm.v1.worker.workspace import init_workspace_manager

21 from .modular_kernel_tools.common import (

22 Config,

@@ -77,6 +78,10 @@ def rank_worker(

78 weights: WeightTensors,

79 verbose: bool,

80 ):

81 + # Initialize workspace manager in child process

82 + device = torch.device(f"cuda:{pgi.local_rank}")

83 + init_workspace_manager(device)

84 +

85 current_platform.seed_everything(pgi.rank)

87 # sanity check

@@ -300,6 +305,7 @@ def test_modular_kernel_combinations_singlegpu(

305 chunk_size: int | None,

306 world_size: int,

307 pytestconfig,

308 + workspace_init,

309 ):

310 """Note: float8_e4m3fn is not supported on CUDA architecture < 89,

311 and those tests will be skipped on unsupported hardware."""

tests/kernels/moe/test_modular_oai_triton_moe.py

@@ -209,6 +209,7 @@ def test_oai_triton_moe(

209 num_experts: int,

210 topk: int,

211 unfused: bool,

212 + workspace_init,

213 ):

214 current_platform.seed_everything(0)

215 (

tests/kernels/moe/test_moe.py

@@ -231,6 +231,7 @@ def test_fused_moe(

231 padding: bool,

232 chunk_size: int,

233 monkeypatch,

234 + workspace_init,

235 ):

236 current_platform.seed_everything(7)

237

tests/kernels/moe/test_nvfp4_moe.py

+1 -1

@@ -40,7 +40,7 @@ MNK_FACTORS = [

40 @pytest.mark.parametrize("dtype", [torch.bfloat16])

41 @torch.inference_mode()

42 def test_cutlass_fp4_moe_no_graph(

43 - m: int, n: int, k: int, e: int, topk: int, dtype: torch.dtype

43 + m: int, n: int, k: int, e: int, topk: int, dtype: torch.dtype, workspace_init

44 ):

45 current_platform.seed_everything(7)

46 with set_current_vllm_config(

tests/kernels/moe/test_pplx_moe.py

@@ -46,6 +46,7 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

46 )

47 from vllm.platforms import current_platform

48 from vllm.utils.math_utils import round_up

49 +from vllm.v1.worker.workspace import init_workspace_manager

51 from ...utils import multi_gpu_test

52 from .parallel_utils import ProcessGroupInfo, parallel_launch

@@ -181,6 +182,7 @@ def test_fused_moe_batched_experts(

182 e: int,

183 topk: int,

184 dtype: torch.dtype,

185 + workspace_init,

186 ):

187 current_platform.seed_everything(7)

188

@@ -863,6 +865,9 @@ def _pplx_test_loop(

865 make_weights: bool,

866 test_fn: Callable,

867 ):

868 + device = torch.device(f"cuda:{pgi.local_rank}")

869 + init_workspace_manager(device)

870 +

871 def format_result(msg, ex=None):

872 if ex is not None:

873 x = str(ex)

tests/v1/attention/test_sparse_mla_backends.py

+215 -36

@@ -22,10 +22,14 @@ from tests.v1.attention.utils import (

22 )

23 from vllm import _custom_ops as ops

24 from vllm.attention.ops import flashmla

25 +from vllm.config import set_current_vllm_config

26 from vllm.model_executor.layers.linear import ColumnParallelLinear

27 from vllm.utils.math_utils import cdiv

27 -from vllm.v1.attention.backends.mla.flashmla_sparse import FlashMLASparseBackend

28 -from vllm.v1.attention.backends.mla.indexer import split_prefill_chunks

28 +from vllm.v1.attention.backends.mla.flashmla_sparse import (

29 + FlashMLASparseBackend,

30 + triton_convert_req_index_to_global_index,

31 +)

32 +from vllm.v1.attention.backends.utils import split_prefill_chunks

34 SPARSE_BACKEND_BATCH_SPECS = {

35 name: BATCH_SPECS[name]

@@ -114,8 +118,12 @@ def _quantize_dequantize_fp8_ds_mla(

118 @pytest.mark.parametrize("batch_name", list(SPARSE_BACKEND_BATCH_SPECS.keys()))

119 @pytest.mark.parametrize("kv_cache_dtype", ["fp8_ds_mla", "auto"])

120 @pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4])

121 +@pytest.mark.skipif(

122 + torch.cuda.get_device_capability() < (9, 0),

123 + reason="FlashMLASparseBackend requires CUDA 9.0 or higher",

124 +)

125 def test_sparse_backend_decode_correctness(

118 - dist_init, batch_name, kv_cache_dtype, tensor_parallel_size

126 + dist_init, batch_name, kv_cache_dtype, tensor_parallel_size, workspace_init

127 ):

128 if not torch.cuda.is_available():

129 pytest.skip("CUDA is required for sparse MLA decode test")

@@ -320,28 +328,29 @@ def test_sparse_backend_decode_correctness(

328 mock_kv_b_proj.weight = torch.nn.Parameter(kv_b_proj_weight.T.contiguous())

329

330 impl_cls = FlashMLASparseBackend.get_impl_cls()

323 - impl = impl_cls(

324 - num_heads=num_heads,

325 - head_size=head_size,

326 - scale=scale,

327 - num_kv_heads=1,

328 - alibi_slopes=None,

329 - sliding_window=None,

330 - kv_cache_dtype=vllm_config.cache_config.cache_dtype,

331 - logits_soft_cap=None,

332 - attn_type="decoder",

333 - kv_sharing_target_layer_name=None,

334 - q_lora_rank=None,

335 - kv_lora_rank=kv_lora_rank,

336 - qk_nope_head_dim=qk_nope_head_dim,

337 - qk_rope_head_dim=qk_rope_head_dim,

338 - qk_head_dim=qk_nope_head_dim + qk_rope_head_dim,

339 - v_head_dim=v_head_dim,

340 - kv_b_proj=mock_kv_b_proj,

341 - indexer=mock_indexer,

342 - )

331 + with set_current_vllm_config(vllm_config):

332 + impl = impl_cls(

333 + num_heads=num_heads,

334 + head_size=head_size,

335 + scale=scale,

336 + num_kv_heads=1,

337 + alibi_slopes=None,

338 + sliding_window=None,

339 + kv_cache_dtype=vllm_config.cache_config.cache_dtype,

340 + logits_soft_cap=None,

341 + attn_type="decoder",

342 + kv_sharing_target_layer_name=None,

343 + q_lora_rank=None,

344 + kv_lora_rank=kv_lora_rank,

345 + qk_nope_head_dim=qk_nope_head_dim,

346 + qk_rope_head_dim=qk_rope_head_dim,

347 + qk_head_dim=qk_nope_head_dim + qk_rope_head_dim,

348 + v_head_dim=v_head_dim,

349 + kv_b_proj=mock_kv_b_proj,

350 + indexer=mock_indexer,

351 + )

352

344 - impl.process_weights_after_loading(dtype)

353 + impl.process_weights_after_loading(dtype)

354

355 layer = MockAttentionLayer(device)

356 out_buffer = torch.empty(

@@ -366,22 +375,192 @@ def test_sparse_backend_decode_correctness(

375 torch.testing.assert_close(backend_output, sdpa_reference, rtol=0.5, atol=0.5)

376

377

378 +def _triton_convert_reference_impl(

379 + req_ids: torch.Tensor,

380 + block_table: torch.Tensor,

381 + token_indices: torch.Tensor,

382 + block_size: int,

383 + num_topk_tokens: int,

384 + HAS_PREFILL_WORKSPACE: bool = False,

385 + prefill_workspace_request_ids: torch.Tensor | None = None,

386 + prefill_workspace_starts: torch.Tensor | None = None,

387 +) -> torch.Tensor:

388 + """Reference implementation for triton_convert_req_index_to_global_index."""

389 + num_tokens = req_ids.shape[0]

390 + max_blocks_per_req = block_table.shape[1]

391 + result = torch.empty(

392 + num_tokens, num_topk_tokens, dtype=torch.int32, device=req_ids.device

393 + )

394 +

395 + for token_id in range(num_tokens):

396 + req_id = req_ids[token_id].item()

397 +

398 + # Determine if this token uses workspace or paged cache

399 + use_prefill_workspace = False

400 + workspace_start = 0

401 + if HAS_PREFILL_WORKSPACE and prefill_workspace_request_ids is not None:

402 + assert prefill_workspace_starts is not None

403 + prefill_req_id = prefill_workspace_request_ids[token_id].item()

404 + if prefill_req_id >= 0:

405 + use_prefill_workspace = True

406 + workspace_start = prefill_workspace_starts[prefill_req_id].item()

407 +

408 + for idx_id in range(num_topk_tokens):

409 + token_idx = token_indices[token_id, idx_id].item()

410 +

411 + if token_idx == -1:

412 + result[token_id, idx_id] = -1

413 + elif use_prefill_workspace:

414 + # Prefill + using prefill workspace: map to workspace offset

415 + result[token_id, idx_id] = workspace_start + token_idx

416 + else:

417 + # Decode: map to paged cache

418 + block_id = token_idx // block_size

419 + if block_id >= max_blocks_per_req:

420 + result[token_id, idx_id] = -1

421 + else:

422 + block_num = block_table[req_id, block_id].item()

423 + offset = token_idx % block_size

424 + result[token_id, idx_id] = block_num * block_size + offset

425 +

426 + return result

427 +

428 +

429 +@pytest.mark.parametrize("block_size", [16, 64, 128])

430 +@pytest.mark.parametrize("num_topk_tokens", [128, 256, 512])

431 +@pytest.mark.skipif(

432 + torch.cuda.get_device_capability() < (9, 0),

433 + reason="FlashMLASparseBackend requires CUDA 9.0 or higher",

434 +)

435 +def test_triton_convert_req_index_to_global_index_decode_only(

436 + block_size, num_topk_tokens

437 +):

438 + device = torch.device("cuda")

439 + num_tokens = 8

440 + num_requests = 4

441 + max_blocks_per_req = 10

442 +

443 + req_id = torch.randint(

444 + 0, num_requests, (num_tokens,), dtype=torch.int32, device=device

445 + )

446 + block_table = torch.randint(

447 + 0, 100, (num_requests, max_blocks_per_req), dtype=torch.int32, device=device

448 + )

449 +

450 + token_indices = torch.randint(

451 + 0,

452 + block_size * max_blocks_per_req,

453 + (num_tokens, num_topk_tokens),

454 + dtype=torch.int32,

455 + device=device,

456 + )

457 +

458 + # Set some to -1 to test masking

459 + token_indices[0, :10] = -1

460 + token_indices[3, 50:60] = -1

461 +

462 + # Set some to out of bounds

463 + token_indices[2, 100:110] = max_blocks_per_req * block_size

464 + token_indices[6, 150:160] = max_blocks_per_req * block_size

465 +

466 + result = triton_convert_req_index_to_global_index(

467 + req_id,

468 + block_table,

469 + token_indices,

470 + BLOCK_SIZE=block_size,

471 + NUM_TOPK_TOKENS=num_topk_tokens,

472 + )

473 +

474 + reference_result = _triton_convert_reference_impl(

475 + req_id,

476 + block_table,

477 + token_indices,

478 + block_size,

479 + num_topk_tokens,

480 + )

481 +

482 + torch.testing.assert_close(result, reference_result, rtol=0, atol=0)

483 +

484 +

485 +@pytest.mark.parametrize("block_size", [16])

486 +@pytest.mark.skipif(

487 + torch.cuda.get_device_capability() < (9, 0),

488 + reason="FlashMLASparseBackend requires CUDA 9.0 or higher",

489 +)

490 +def test_triton_convert_req_index_to_global_index_with_prefill_workspace(block_size):

491 + device = torch.device("cuda")

492 + num_requests = 4

493 + max_blocks_per_req = 8

494 + num_topk_tokens = 128

495 +

496 + # First 6 tokens are decode (reqs 0, 1), last 6 are prefill (reqs 2, 3)

497 + req_id = torch.tensor(

498 + [0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3], dtype=torch.int32, device=device

499 + )

500 + prefill_workspace_request_ids = torch.tensor(

501 + [-1, -1, -1, -1, -1, -1, 0, 0, 0, 1, 1, 1], dtype=torch.int32, device=device

502 + )

503 +

504 + # Workspace starts for the 2 prefill reqs: req 2 starts at 0, req 3 starts at 100

505 + prefill_workspace_starts = torch.tensor([0, 100], dtype=torch.int32, device=device)

506 +

507 + block_table = torch.randint(

508 + 0, 50, (num_requests, max_blocks_per_req), dtype=torch.int32, device=device

509 + )

510 + token_indices = torch.randint(

511 + 0,

512 + block_size * max_blocks_per_req,

513 + (req_id.shape[0], num_topk_tokens),

514 + dtype=torch.int32,

515 + device=device,

516 + )

517 +

518 + # Set some to -1 to test masking

519 + token_indices[0, :10] = -1

520 + token_indices[3, 50:60] = -1

521 +

522 + # Set some to out of bounds

523 + token_indices[2, 100:110] = max_blocks_per_req * block_size

524 + token_indices[6, 150:160] = max_blocks_per_req * block_size

525 +

526 + result = triton_convert_req_index_to_global_index(

527 + req_id,

528 + block_table,

529 + token_indices,

530 + BLOCK_SIZE=block_size,

531 + NUM_TOPK_TOKENS=num_topk_tokens,

532 + HAS_PREFILL_WORKSPACE=True,

533 + prefill_workspace_request_ids=prefill_workspace_request_ids,

534 + prefill_workspace_starts=prefill_workspace_starts,

535 + )

536 +

537 + reference_result = _triton_convert_reference_impl(

538 + req_id,

539 + block_table,

540 + token_indices,

541 + block_size,

542 + num_topk_tokens,

543 + HAS_PREFILL_WORKSPACE=True,

544 + prefill_workspace_request_ids=prefill_workspace_request_ids,

545 + prefill_workspace_starts=prefill_workspace_starts,

546 + )

547 +

548 + torch.testing.assert_close(result, reference_result, rtol=0, atol=0)

549 +

550 +

551 @pytest.mark.parametrize(

370 - "seq_lens,max_buf,start,expected",

552 + "seq_lens,max_buf,expected",

553 [

554 # Basic split: totals per chunk ≤ max_buf

373 - (torch.tensor([2, 3, 4, 2]), 5, 0, [(0, 2), (2, 3), (3, 4)]),

374 - # Non-zero start index

375 - (torch.tensor([2, 3, 4, 2]), 5, 1, [(1, 2), (2, 3), (3, 4)]),

376 - # Exact fits should split between items when adding the next would

377 - # overflow

378 - (torch.tensor([5, 5, 5]), 5, 0, [(0, 1), (1, 2), (2, 3)]),

555 + (torch.tensor([2, 3, 4, 2]), 5, [(0, 2), (2, 3), (3, 4)]),

556 + # Exact fits should split between items when adding the next would overflow

557 + (torch.tensor([5, 5, 5]), 5, [(0, 1), (1, 2), (2, 3)]),

558 # All requests fit in a single chunk

380 - (torch.tensor([1, 1, 1]), 10, 0, [(0, 3)]),

381 - # Large buffer with non-zero start

382 - (torch.tensor([4, 4, 4]), 100, 1, [(1, 3)]),

559 + (torch.tensor([1, 1, 1]), 10, [(0, 3)]),

560 + # Large buffer

561 + (torch.tensor([4, 4, 4]), 100, [(0, 3)]),

562 ],

563 )

385 -def test_split_prefill_chunks(seq_lens, max_buf, start, expected):

386 - out = split_prefill_chunks(seq_lens, max_buf, start)

564 +def test_split_prefill_chunks(seq_lens, max_buf, expected):

565 + out = split_prefill_chunks(seq_lens, max_buf)

566 assert out == expected

vllm/_custom_ops.py

+23

@@ -2403,6 +2403,29 @@ def cp_gather_cache(

2403 )

2404

2405

2406 +def cp_gather_and_upconvert_fp8_kv_cache(

2407 + src_cache: torch.Tensor,

2408 + dst: torch.Tensor,

2409 + block_table: torch.Tensor,

2410 + seq_lens: torch.Tensor,

2411 + workspace_starts: torch.Tensor,

2412 + batch_size: int,

2413 +) -> None:

2414 + """Gather and upconvert FP8 KV cache to BF16 workspace.

2415 +

2416 + Args:

2417 + src_cache: FP8 KV cache [num_blocks, block_size, 656]

2418 + dst: BF16 output workspace [total_tokens, 576]

2419 + block_table: Block indices [num_reqs, max_blocks]

2420 + seq_lens: Sequence lengths [num_reqs]

2421 + workspace_starts: Workspace start offsets [num_reqs]

2422 + batch_size: Number of requests

2423 + """

2424 + torch.ops._C_cache_ops.cp_gather_and_upconvert_fp8_kv_cache(

2425 + src_cache, dst, block_table, seq_lens, workspace_starts, batch_size

2426 + )

2427 +

2428 +

2429 def indexer_k_quant_and_cache(

2430 k: torch.Tensor,

2431 kv_cache: torch.Tensor,

vllm/envs.py

@@ -239,6 +239,7 @@ if TYPE_CHECKING:

239 VLLM_NCCL_INCLUDE_PATH: str | None = None

240 VLLM_USE_FBGEMM: bool = False

241 VLLM_GC_DEBUG: str = ""

242 + VLLM_DEBUG_WORKSPACE: bool = False

243 VLLM_DISABLE_SHARED_EXPERTS_STREAM: bool = False

244 VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD: int = 256

245 VLLM_COMPILE_CACHE_SAVE_FORMAT: Literal["binary", "unpacked"] = "binary"

@@ -1537,6 +1538,9 @@ environment_variables: dict[str, Callable[[], Any]] = {

1538 # - VLLM_GC_DEBUG='{"top_objects":5}': enable GC debugger with

1539 # top 5 collected objects

1540 "VLLM_GC_DEBUG": lambda: os.getenv("VLLM_GC_DEBUG", ""),

1541 + # Debug workspace allocations.

1542 + # logging of workspace resize operations.

1543 + "VLLM_DEBUG_WORKSPACE": lambda: bool(int(os.getenv("VLLM_DEBUG_WORKSPACE", "0"))),

1544 # Disables parallel execution of shared_experts via separate cuda stream

1545 "VLLM_DISABLE_SHARED_EXPERTS_STREAM": lambda: bool(

1546 int(os.getenv("VLLM_DISABLE_SHARED_EXPERTS_STREAM", "0"))

vllm/model_executor/layers/fused_moe/modular_kernel.py

+16 -57

@@ -22,12 +22,12 @@ from vllm.model_executor.layers.fused_moe.utils import (

22 from vllm.platforms import current_platform

23 from vllm.utils.math_utils import cdiv

24 from vllm.v1.worker.ubatching import (

25 - dbo_current_ubatch_id,

25 dbo_enabled,

26 dbo_maybe_run_recv_hook,

27 dbo_register_recv_hook,

28 dbo_yield,

29 )

30 +from vllm.v1.worker.workspace import current_workspace_manager

32 logger = init_logger(__name__)

@@ -661,25 +661,6 @@ def _slice_scales(

661 return None

662

663

664 -class SharedResizableBuffer:

665 - def __init__(self):

666 - self.buffer = None

667 -

668 - def get(

669 - self, shape: tuple[int, ...], device: torch.device, dtype: torch.dtype

670 - ) -> torch.Tensor:

671 - assert shape != ()

672 - shape_numel = prod(shape)

673 - if (

674 - self.buffer is None

675 - or self.buffer.numel() < shape_numel

676 - or self.buffer.device != device

677 - or self.buffer.dtype != dtype

678 - ):

679 - self.buffer = torch.empty(shape_numel, device=device, dtype=dtype)

680 - return self.buffer[:shape_numel].view(*shape)

681 -

682 -

664 @final

665 class FusedMoEModularKernel(torch.nn.Module):

666 """

@@ -694,22 +675,6 @@ class FusedMoEModularKernel(torch.nn.Module):

675 objects.

676 """

677

697 - class SharedBuffers:

698 - def __init__(self) -> None:

699 - self.fused_out = SharedResizableBuffer()

700 - self.workspace13 = SharedResizableBuffer()

701 - self.workspace2 = SharedResizableBuffer()

702 -

703 - # Persistent buffers that are shared across `FusedMoEModularKernel`

704 - # instances (layers), to save memory and allocattions.

705 - #

706 - # We have two sets of buffers to support dual batch overlap (DBO) where each

707 - # microbatch (ubatch) should use its own set of buffers to avoid

708 - # cross-ubatch contimination.

709 - # NOTE that memory is lazily allocated for these buffers, meaning that if

710 - # DBO isn't being used, the second SharedBuffers will be empty.

711 - shared_buffers: list[SharedBuffers] = [SharedBuffers(), SharedBuffers()]

712 -

678 def __init__(

679 self,

680 prepare_finalize: FusedMoEPrepareAndFinalize,

@@ -806,10 +771,6 @@ class FusedMoEModularKernel(torch.nn.Module):

771 assert M_full > 0 and M_chunk > 0

772

773 num_chunks, _ = self._chunk_info(M_full)

809 -

810 - # select per-ubatch buffers to avoid cross-ubatch reuse under DBO

811 - ubatch_idx = dbo_current_ubatch_id()

812 - buffers = self.shared_buffers[ubatch_idx]

774 workspace_dtype = self.fused_experts.workspace_dtype(out_dtype)

775

776 # Force worst-case allocation in profiling run for

@@ -832,14 +793,11 @@ class FusedMoEModularKernel(torch.nn.Module):

793 expert_tokens_meta,

794 )

795 )

835 - buffers.workspace13.get(

836 - max_workspace_13, device=device, dtype=workspace_dtype

837 - )

838 - buffers.workspace2.get(

839 - max_workspace_2, device=device, dtype=workspace_dtype

840 - )

841 - buffers.fused_out.get(

842 - max_fused_out_shape, device=device, dtype=workspace_dtype

796 +

797 + current_workspace_manager().get_simultaneous(

798 + (max_workspace_13, workspace_dtype),

799 + (max_workspace_2, workspace_dtype),

800 + (max_fused_out_shape, out_dtype),

801 )

802

803 # Get intermediate workspace shapes based off the chunked M size.

@@ -866,22 +824,23 @@ class FusedMoEModularKernel(torch.nn.Module):

824

825 # We can reuse the memory between cache1 and cache3 because by the

826 # time we need cache3, we're done with cache1.

869 - workspace13 = buffers.workspace13.get(

870 - workspace13_shape, device=device, dtype=workspace_dtype

871 - )

872 - workspace2 = buffers.workspace2.get(

873 - workspace2_shape, device=device, dtype=workspace_dtype

874 - )

875 -

827 # Construct the entire output that can then be processed in chunks.

828 # Reuse workspace13 for the output in the non-chunked case as long

829 # as it is large enough. This will not always be the case for standard

830 # format experts and with experts that have empty workspaces.

831 if num_chunks == 1 and prod(workspace13_shape) >= prod(fused_out_shape):

832 + workspace13, workspace2 = current_workspace_manager().get_simultaneous(

833 + (workspace13_shape, workspace_dtype),

834 + (workspace2_shape, workspace_dtype),

835 + )

836 fused_out = _resize_cache(workspace13, fused_out_shape)

837 else:

883 - fused_out = buffers.fused_out.get(

884 - fused_out_shape, device=device, dtype=out_dtype

838 + workspace13, workspace2, fused_out = (

839 + current_workspace_manager().get_simultaneous(

840 + (workspace13_shape, workspace_dtype),

841 + (workspace2_shape, workspace_dtype),

842 + (fused_out_shape, out_dtype),

843 + )

844 )

845

846 return workspace13, workspace2, fused_out

vllm/model_executor/models/deepseek_v2.py

+18 -19

@@ -83,6 +83,7 @@ from vllm.v1.attention.backends.mla.indexer import (

83 DeepseekV32IndexerMetadata,

84 )

85 from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec

86 +from vllm.v1.worker.workspace import current_workspace_manager

88 from .interfaces import MixtureOfExperts, SupportsEagle, SupportsLoRA, SupportsPP

89 from .utils import (

@@ -616,8 +617,15 @@ def sparse_attn_indexer(

617 # careful! this will be None in dummy run

618 attn_metadata = get_forward_context().attn_metadata

619 fp8_dtype = current_platform.fp8_dtype()

620 +

621 # assert isinstance(attn_metadata, dict)

622 if not isinstance(attn_metadata, dict):

623 + # Reserve workspace for indexer during profiling run

624 + current_workspace_manager().get_simultaneous(

625 + ((total_seq_lens, head_dim), torch.float8_e4m3fn),

626 + ((total_seq_lens, 4), torch.uint8),

627 + )

628 +

629 return sparse_attn_indexer_fake(

630 hidden_states,

631 k_cache_prefix,

@@ -651,17 +659,17 @@ def sparse_attn_indexer(

659 topk_indices_buffer[: hidden_states.shape[0]] = -1

660 if has_prefill:

661 prefill_metadata = attn_metadata.prefill

662 +

663 + # Get the full shared workspace buffers once (will allocate on first use)

664 + workspace_manager = current_workspace_manager()

665 + k_fp8_full, k_scale_full = workspace_manager.get_simultaneous(

666 + ((total_seq_lens, head_dim), fp8_dtype),

667 + ((total_seq_lens, 4), torch.uint8),

668 + )

669 +

670 for chunk in prefill_metadata.chunks:

655 - k_fp8 = torch.empty(

656 - [chunk.total_seq_lens, head_dim],

657 - device=k.device,

658 - dtype=fp8_dtype,

659 - )

660 - k_scale = torch.empty(

661 - [chunk.total_seq_lens, 4],

662 - device=k.device,

663 - dtype=torch.uint8,

664 - )

671 + k_fp8 = k_fp8_full[: chunk.total_seq_lens]

672 + k_scale = k_scale_full[: chunk.total_seq_lens]

673 ops.cp_gather_indexer_k_quant_cache(

674 kv_cache,

675 k_fp8,

@@ -777,15 +785,6 @@ def sparse_attn_indexer_fake(

785 total_seq_lens: int,

786 topk_indices_buffer: torch.Tensor | None,

787 ) -> torch.Tensor:

780 - # profile run

781 - # NOTE(Chen): create the max possible flattened_kv. So that

782 - # profile_run can get correct memory usage.

783 - _flattened_kv = torch.empty(

784 - [total_seq_lens, head_dim + 4], device=k.device, dtype=torch.uint8

785 - )

786 - fp8_dtype = current_platform.fp8_dtype()

787 - _k_fp8 = _flattened_kv[..., :head_dim].view(fp8_dtype).contiguous()

788 - _k_scale = _flattened_kv[..., head_dim:].view(torch.float32).contiguous()

788 return topk_indices_buffer

789

790

vllm/v1/attention/backends/mla/flashmla_sparse.py

+567 -98

@@ -18,7 +18,7 @@ from vllm.attention.ops.flashmla import (

18 flash_mla_with_kvcache,

19 get_mla_metadata,

20 )

21 -from vllm.config import VllmConfig

21 +from vllm.config import VllmConfig, get_current_vllm_config

22 from vllm.config.cache import CacheDType

23 from vllm.logger import init_logger

24 from vllm.platforms import current_platform

@@ -30,13 +30,31 @@ from vllm.v1.attention.backends.utils import (

30 AttentionCGSupport,

31 AttentionMetadataBuilder,

32 CommonAttentionMetadata,

33 + reshape_attn_output_for_spec_decode,

34 + reshape_query_for_spec_decode,

35 + split_decodes_and_prefills,

36 + split_prefill_chunks,

37 )

38 from vllm.v1.kv_cache_interface import AttentionSpec

39 +from vllm.v1.worker.workspace import current_workspace_manager

41 if TYPE_CHECKING:

42 from vllm.model_executor.models.deepseek_v2 import Indexer

44 logger = init_logger(__name__)

45 +

46 +# For FP8 sparse attention we have two impelementations:

47 +# 1. Mixed batch mode: use the FP8 decode kernel for both prefill and decode this is

48 +# done by treating all tokens as single batch.

49 +# 2. Separate prefill and decode mode: use the BF16 prefill kernel for prefill

50 +# (upconverting the FP8 cache to BF16 then calling the prefill kernel) and using

51 +# the FP8 decode kernel for decode.

52 +# Currently we use #1 when the number of heads per rank is low (i.e. TP) since the BF16

53 +# prefill kernel requires padding the numer of heads to 128 while the decode does not

54 +# so when the per ranke head count is below MIN_HEADS_FOR_BF16_PREFILL we use the mixed

55 +# batch mode (#2).

56 +MIN_HEADS_FOR_BF16_PREFILL = 32

57 +

58 """

59 NOTE: FlashMLA Sparse uses an fp8 cache with the following format

@@ -127,19 +145,72 @@ class FlashMLASparseMetadata:

145 dummy_block_table: torch.Tensor

146 cache_lens: torch.Tensor

147

130 - fp8_extra_metadata: FP8KernelMetadata | None = None

131 -

132 -

148 + @dataclass

149 + class FP8SeperatePrefillDecode:

150 + @dataclass

151 + class Decode:

152 + kernel_metadata: "FlashMLASparseMetadata.FP8KernelMetadata"

153 + decode_query_len: int # needed for reshape in spec decode

154 +

155 + @dataclass

156 + class Prefill:

157 + # Sequence lengths (context + query) for prefill requests

158 + # Shape: [num_prefill_reqs]

159 + seq_lens: torch.Tensor

160 +

161 + # Request ID for each token: -1 for decode tokens, request index

162 + # (0, 1, 2, ...) for prefill tokens.

163 + # Shape: [num_actual_tokens]

164 + request_ids: torch.Tensor

165 +

166 + # Workspace start offsets for all prefill requests

167 + # Shape: [num_prefill_reqs], adjusted in-place per chunk to be

168 + # 0-indexed within each chunk. Used to map prefill tokens to workspace

169 + # offsets in convert_logical_index_to_physical_index

170 + workspace_starts: torch.Tensor

171 +

172 + @dataclass

173 + class Chunk:

174 + """Metadata for a chunk of prefill requests.

175 +

176 + Prefill requests may be chunked to fit within the fixed workspace size.

177 + """

178 +

179 + seq_lens: torch.Tensor

180 + tokens_slice: slice

181 + block_table: torch.Tensor

182 + req_start_idx: int

183 + workspace_starts: torch.Tensor

184 + chunk_tot_seqlen: int

185 +

186 + chunks: list[Chunk]

187 +

188 + num_prefills: int = 0

189 + num_decodes: int = 0

190 + num_prefill_tokens: int = 0

191 + num_decode_tokens: int = 0

192 +

193 + decode: Decode | None = None

194 + prefill: Prefill | None = None

195 +

196 + fp8_extra_metadata: FP8SeperatePrefillDecode | FP8KernelMetadata | None = None

197 + fp8_use_mixed_batch: bool = False

198 +

199 +

200 +# Kernel with prefill workspace support

201 @triton.jit

202 def _convert_req_index_to_global_index_kernel(

203 req_id_ptr, # int32 [num_tokens]

204 block_table_ptr, # int32 [num_requests, max_num_blocks_per_req]

205 token_indices_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS]

206 out_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS]

207 + prefill_request_id_ptr, # int32 [num_tokens], -1 for decode, >=0 for prefill

208 + workspace_starts_ptr, # int32 [num_prefill_reqs+1] or nullptr

209 # shapes (compile-time where possible)

210 max_num_blocks_per_req: tl.constexpr,

211 BLOCK_SIZE: tl.constexpr,

212 BLOCK_N: tl.constexpr, # tile width along columns

213 + HAS_PREFILL: tl.constexpr,

214 # strides (in elements)

215 bt_stride0,

216 bt_stride1,

@@ -165,7 +236,10 @@ def _convert_req_index_to_global_index_kernel(

236

237 # Only token == -1 should propagate as -1

238 is_invalid_tok = tok < 0

168 -

239 + is_prefill = False

240 + if HAS_PREFILL:

241 + prefill_req_id = tl.load(prefill_request_id_ptr + token_id)

242 + is_prefill = prefill_req_id >= 0

243 # Compute block id and in-block offset

244 block_id = tok // BLOCK_SIZE

245 inblock_off = tok % BLOCK_SIZE

@@ -173,12 +247,18 @@ def _convert_req_index_to_global_index_kernel(

247 # Guard block_table access

248 valid_block = (block_id < max_num_blocks_per_req) & (block_id >= 0)

249 bt_ptr = block_table_ptr + req * bt_stride0 + block_id * bt_stride1

176 - base = tl.load(bt_ptr, mask=valid_block, other=0)

177 -

178 - # If token == -1 OR block_id OOB, output -1; else base * BLOCK_SIZE + offset

179 - out_val = tl.where(

180 - is_invalid_tok | (~valid_block), -1, base * BLOCK_SIZE + inblock_off

181 - )

250 + is_invalid_tok |= ~valid_block

251 + base = tl.load(bt_ptr, mask=valid_block & ~is_prefill, other=0)

252 + out_val = base * BLOCK_SIZE + inblock_off

253 +

254 + # Override with prefill output if prefill is enabled

255 + if HAS_PREFILL:

256 + workspace_start = tl.load(

257 + workspace_starts_ptr + prefill_req_id, mask=is_prefill, other=0

258 + )

259 + prefill_out = workspace_start + tok

260 + out_val = tl.where(is_prefill, prefill_out, out_val)

261 + out_val = tl.where(is_invalid_tok, -1, out_val)

262

263 # Store results

264 out_ptr_ij = out_ptr + token_id * out_stride0 + indice_id * out_stride1

@@ -192,6 +272,9 @@ def triton_convert_req_index_to_global_index(

272 BLOCK_SIZE: int = 64,

273 NUM_TOPK_TOKENS: int = 2048,

274 BLOCK_N: int = 128, # tile width along columns

275 + HAS_PREFILL_WORKSPACE: bool = False,

276 + prefill_workspace_request_ids: torch.Tensor | None = None,

277 + prefill_workspace_starts: torch.Tensor | None = None,

278 ):

279 """

280 out[token_id, indice_id] =

@@ -202,17 +285,32 @@ def triton_convert_req_index_to_global_index(

285 Only when token_indices[token_id, indice_id] == -1 do we output -1.

286 For safety, we also output -1 if the derived block_id would be

287 out-of-bounds.

288 +

289 + When HAS_PREFILL_WORKSPACE is True, prefill tokens are mapped to workspace offsets

290 + instead of global cache slots. prefill_workspace_request_ids and

291 + prefill_workspace_starts must be provided.

292 +

293 + prefill_workspace_request_ids: int32 [num_tokens], -1 for decode else

294 + prefill request index (maps to prefill_workspace_starts)

295 + prefill_workspace_starts: int32 [num_prefills], 0-indexed workspace

296 + starts for each prefill request

297 """

298 assert req_id.dtype == torch.int32

299 assert block_table.dtype == torch.int32

300 assert token_indices.dtype == torch.int32

301 assert token_indices.shape[1] == NUM_TOPK_TOKENS

302 assert NUM_TOPK_TOKENS % BLOCK_N == 0, (

211 - f"NUM_TOPK_TOKENS ({NUM_TOPK_TOKENS}) must be divisible byBLOCK_N ({BLOCK_N})"

303 + f"NUM_TOPK_TOKENS ({NUM_TOPK_TOKENS}) must be divisible by BLOCK_N ({BLOCK_N})"

304 )

305

306 + if HAS_PREFILL_WORKSPACE:

307 + assert prefill_workspace_request_ids is not None

308 + assert prefill_workspace_starts is not None

309 + assert prefill_workspace_request_ids.dtype == torch.int32

310 + assert prefill_workspace_starts.dtype == torch.int32

311 +

312 num_tokens = req_id.shape[0]

215 - num_requests, max_num_blocks_per_req = block_table.shape

313 + max_num_blocks_per_req = block_table.shape[1]

314 tiles_per_row = NUM_TOPK_TOKENS // BLOCK_N

315

316 # Ensure contiguous tensors on the same device

@@ -226,6 +324,13 @@ def triton_convert_req_index_to_global_index(

324 ti_stride0, ti_stride1 = token_indices_c.stride()

325 out_stride0, out_stride1 = out.stride()

326

327 + # Prepare prefill pointers

328 + if HAS_PREFILL_WORKSPACE:

329 + assert prefill_workspace_request_ids is not None # for mypy

330 + assert prefill_workspace_starts is not None # for mypy

331 + assert prefill_workspace_request_ids.is_contiguous()

332 + assert prefill_workspace_starts.is_contiguous()

333 +

334 # Exact 2D grid: tokens × column tiles

335 grid = (num_tokens, tiles_per_row)

336

@@ -234,10 +339,13 @@ def triton_convert_req_index_to_global_index(

339 block_table_c,

340 token_indices_c,

341 out,

342 + prefill_workspace_request_ids,

343 + prefill_workspace_starts,

344 # shapes / constexprs

345 max_num_blocks_per_req,

346 BLOCK_SIZE,

347 BLOCK_N,

348 + HAS_PREFILL_WORKSPACE,

349 # strides

350 bt_stride0,

351 bt_stride1,

@@ -249,7 +357,16 @@ def triton_convert_req_index_to_global_index(

357 return out

358

359

252 -@dataclass

360 +def get_prefill_workspace_size(max_model_len: int):

361 + # NOTE(Lucas): 5 is a magic number for controlling the prefill buffer size.

362 + # May be tuned later.

363 + # Memory usage: 5 * max_model_len * 576 * 2 bytes

364 + # Example: DeepSeek-V3.2 with max_model_len=163840 ->

365 + # 5 * 163840 * 576 * 2 = ~900 MB

366 + # This fits nicely below the typical MoE workspace size of >2GB so this is "free"

367 + return max_model_len * 5

368 +

369 +

370 class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetadata]):

371 _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH

372

@@ -259,29 +376,42 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad

376 layer_names: list[str],

377 vllm_config: VllmConfig,

378 device: torch.device,

262 - ):

379 + ) -> None:

380 + self.vllm_config = vllm_config

381 + self.layer_names = layer_names

382 cache_config = vllm_config.cache_config

383 self.kv_cache_spec = kv_cache_spec

384 self.model_config = vllm_config.model_config

385 parallel_config = vllm_config.parallel_config

386 self.device = device

387

388 + # Treat requests with query length <= 1 as decodes to match the

389 + # DeepGEMM indexer constraint (fp8_paged_mqa_logits only supports next_n <= 2)

390 + self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)

391 +

392 props = torch.cuda.get_device_properties(device)

393 sm_count = props.multi_processor_count

394

395 self.num_heads = self.model_config.get_num_attention_heads(parallel_config)

396 self.mla_dims = get_mla_dims(self.model_config)

397 +

398 self.topk_tokens = vllm_config.model_config.hf_config.index_topk

399 self.use_fp8_kv_cache = cache_config.cache_dtype == "fp8_ds_mla"

276 - self.topk_tokens_tensor = torch.tensor(

277 - [self.topk_tokens], device=device, dtype=torch.int32

400 + max_num_seqs = vllm_config.scheduler_config.max_num_seqs

401 + # Shape: [max_num_seqs], all elements = topk_tokens (constant for full-CG)

402 + self.topk_tokens_tensor = torch.full(

403 + (max_num_seqs,), self.topk_tokens, device=device, dtype=torch.int32

404 )

279 - self.max_model_len_tensor = torch.tensor(

280 - [self.model_config.max_model_len], device=device, dtype=torch.int32

405 + # Shape: [max_num_seqs], all elements = max_model_len

406 + self.max_model_len_tensor = torch.full(

407 + (max_num_seqs,),

408 + self.model_config.max_model_len,

409 + device=device,

410 + dtype=torch.int32,

411 )

412 # this is ignored by `flash_mla_with_kvcache` if indices not None

413 self.dummy_block_table = torch.empty(

284 - (1, 1), dtype=torch.int32, device=self.device

414 + (max_num_seqs, 1), dtype=torch.int32, device=self.device

415 )

416

417 # Equation taken from FlashMLA/csrc/pybind.cpp

@@ -299,10 +429,9 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad

429 dtype=torch.int32,

430 device=device,

431 )

432 + # Sized for per-request batching (num_decodes + 1)

433 self.num_splits_buffer = torch.empty(

303 - # We pack all the tokens into one batch for sparse attention.

304 - # Otherwise, we can exceed the sm of `get_mla_metadata`.

305 - (2,),

434 + (max_num_seqs + 1,),

435 dtype=torch.int32,

436 device=device,

437 )

@@ -312,30 +441,171 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad

441 device=device,

442 )

443

315 - def build(

444 + def _build_fp8_mixed_decode_prefill(

445 self,

317 - common_prefix_len: int,

446 common_attn_metadata: CommonAttentionMetadata,

319 - fast_build: bool = False,

320 - ) -> FlashMLASparseMetadata:

447 + ) -> "FlashMLASparseMetadata.FP8KernelMetadata":

448 + """Build FP8 metadata treating all tokens as one mixed batch.

449 +

450 + This matches main branch's approach and avoids the BF16 prefill kernel

451 + which has head padding overhead when num_heads is small (high TP case).

452 + """

453 num_tokens = common_attn_metadata.num_actual_tokens

322 - starts = np.asarray(common_attn_metadata.query_start_loc_cpu, dtype=np.int32)

323 - seg_lengths = np.diff(starts)

324 - req_id_per_token = np.repeat(

325 - np.arange(seg_lengths.shape[0], dtype=np.int32), seg_lengths

454 +

455 + # Build metadata for all tokens as a single batch

456 + tile_scheduler_metadata, num_splits = get_mla_metadata(

457 + cache_seqlens=self.topk_tokens_tensor[:1], # Single batch

458 + num_q_tokens_per_head_k=num_tokens * self.num_heads,

459 + topk=self.topk_tokens,

460 + num_heads_q=self.num_heads,

461 + num_heads_k=1,

462 + is_fp8_kvcache=True,

463 )

327 - # Zero-fill for cudagraphs

328 - self.req_id_per_token_buffer.fill_(0)

329 - self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_(

330 - torch.from_numpy(req_id_per_token), non_blocking=True

464 +

465 + num_sm_parts = tile_scheduler_metadata.size(0)

466 + tile_scheduler_metadata_buffer = self.tile_scheduler_metadata_buffer[

467 + :num_sm_parts

468 + ]

469 + tile_scheduler_metadata_buffer.copy_(tile_scheduler_metadata)

470 + num_splits_view = self.num_splits_buffer[:2]

471 + num_splits_view.copy_(num_splits)

472 +

473 + fp8_metadata = FlashMLASparseMetadata.FP8KernelMetadata(

474 + scheduler_metadata=tile_scheduler_metadata_buffer,

475 + num_splits=num_splits_view,

476 + cache_lens=self.max_model_len_tensor[:1],

477 + dummy_block_table=self.dummy_block_table[:1],

478 )

332 - req_id_per_token = self.req_id_per_token_buffer[:num_tokens]

479

334 - fp8_extra_metadata = None

335 - if self.use_fp8_kv_cache:

480 + return fp8_metadata

481 +

482 + def _build_fp8_separate_prefill_decode(

483 + self,

484 + common_attn_metadata: CommonAttentionMetadata,

485 + ) -> "FlashMLASparseMetadata.FP8SeperatePrefillDecode":

486 + num_tokens = common_attn_metadata.num_actual_tokens

487 +

488 + (num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens) = (

489 + split_decodes_and_prefills(

490 + common_attn_metadata,

491 + decode_threshold=self.reorder_batch_threshold or 1,

492 + require_uniform=True,

493 + )

494 + )

495 +

496 + FP8Meta = FlashMLASparseMetadata.FP8SeperatePrefillDecode

497 + fp8_metadata = FP8Meta(

498 + num_decodes=num_decodes,

499 + num_prefills=num_prefills,

500 + num_decode_tokens=num_decode_tokens,

501 + num_prefill_tokens=num_prefill_tokens,

502 + )

503 +

504 + # Extract prefill sequence lengths (context + query, not just query)

505 + # Decode requests come first in the batch, prefill requests follow

506 + prefill_seq_lens = None

507 + prefill_request_id = None

508 + prefill_workspace_starts = None

509 + prefill_chunks = None

510 +

511 + # For pure decode batches, prefill_request_id will be None

512 + # For mixed batches, it will have -1 for decode and request_id for prefill

513 + if num_prefills > 0:

514 + seq_lens_cpu = common_attn_metadata.seq_lens_cpu

515 + seq_lens = common_attn_metadata.seq_lens

516 + query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu

517 +

518 + prefill_seq_lens_cpu = seq_lens_cpu[num_decodes:]

519 + prefill_seq_lens = seq_lens[num_decodes:]

520 +

521 + # Build prefill_request_id: -1 for decode, request index for

522 + # prefill. This enables a single

523 + # convert_logical_index_to_physical_index call for all tokens

524 + prefill_request_id = torch.full(

525 + (num_tokens,), -1, dtype=torch.int32, device=self.device

526 + )

527 + # Map prefill tokens to their request IDs (0, 1, 2, ...)

528 + for req_idx in range(num_prefills):

529 + # Get query token range for this prefill request

530 + global_req_idx = num_decodes + req_idx

531 + req_query_start = query_start_loc_cpu[global_req_idx]

532 + req_query_end = query_start_loc_cpu[global_req_idx + 1]

533 + prefill_request_id[req_query_start:req_query_end] = req_idx

534 +

535 + # will be adjusted by chunk loop

536 + prefill_workspace_starts_cpu = torch.zeros(

537 + num_prefills, dtype=torch.int32, pin_memory=True

538 + )

539 + prefill_workspace_starts_cpu[1:] = torch.cumsum(

540 + prefill_seq_lens_cpu[:-1], dim=0

541 + )

542 + # populated by non-blocking copy after prefill_workspace_starts_cpu is

543 + # updated by each chunk

544 + prefill_workspace_starts = torch.empty(

545 + num_prefills, dtype=torch.int32, device=self.device

546 + )

547 +

548 + # Chunk prefill requests to fit within workspace size

549 + max_prefill_buffer_size = get_prefill_workspace_size(

550 + self.vllm_config.model_config.max_model_len

551 + )

552 + chunk_bounds = split_prefill_chunks(

553 + prefill_seq_lens_cpu, max_prefill_buffer_size

554 + )

555 +

556 + prefill_chunks = []

557 + for chunk_start, chunk_end in chunk_bounds:

558 + # Adjust workspace_starts in-place per chunk to be

559 + # 0-indexed within each chunk

560 + # Example: seq_lens=[10,15,20,5], chunks=[[0,2],[2,4]]

561 + # Initial: workspace_starts=[0,10,25,45]

562 + # After: workspace_starts=[0,10,0,20]

563 + # (chunk 0 starts at 0, chunk 1 starts at 0)

564 + offset = prefill_workspace_starts_cpu[chunk_start].item()

565 + prefill_workspace_starts_cpu[chunk_start:chunk_end] -= offset

566 +

567 + chunk_seq_lens = prefill_seq_lens[chunk_start:chunk_end]

568 + chunk_tot_seqlen = prefill_seq_lens_cpu[chunk_start:chunk_end].sum()

569 + token_start = query_start_loc_cpu[num_decodes + chunk_start].item()

570 + token_end = query_start_loc_cpu[num_decodes + chunk_end].item()

571 + tokens_slice = slice(token_start, token_end)

572 +

573 + # Create chunk view of gpu tensor

574 + chunk_workspace_starts = prefill_workspace_starts[chunk_start:chunk_end]

575 + chunk_block_table = common_attn_metadata.block_table_tensor[

576 + num_decodes + chunk_start : num_decodes + chunk_end

577 + ]

578 +

579 + prefill_chunks.append(

580 + FP8Meta.Prefill.Chunk(

581 + seq_lens=chunk_seq_lens,

582 + tokens_slice=tokens_slice,

583 + block_table=chunk_block_table,

584 + req_start_idx=chunk_start,

585 + workspace_starts=chunk_workspace_starts,

586 + chunk_tot_seqlen=chunk_tot_seqlen,

587 + )

588 + )

589 +

590 + prefill_workspace_starts.copy_(

591 + prefill_workspace_starts_cpu, non_blocking=True

592 + )

593 +

594 + fp8_metadata.prefill = FP8Meta.Prefill(

595 + seq_lens=prefill_seq_lens,

596 + request_ids=prefill_request_id,

597 + workspace_starts=prefill_workspace_starts,

598 + chunks=prefill_chunks,

599 + )

600 +

601 + if num_decodes > 0:

602 + # Compute decode_query_len for spec decode (uniform due to require_uniform)

603 + query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu

604 + decode_query_len = (query_start_loc_cpu[1] - query_start_loc_cpu[0]).item()

605 +

606 tile_scheduler_metadata, num_splits = get_mla_metadata(

337 - cache_seqlens=self.topk_tokens_tensor,

338 - num_q_tokens_per_head_k=num_tokens * self.num_heads,

607 + cache_seqlens=self.topk_tokens_tensor[:num_decodes],

608 + num_q_tokens_per_head_k=decode_query_len * self.num_heads,

609 topk=self.topk_tokens,

610 num_heads_q=self.num_heads,

611 num_heads_k=1,

@@ -348,33 +618,70 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad

618 :num_sm_parts

619 ]

620 tile_scheduler_metadata_buffer.copy_(tile_scheduler_metadata)

351 - self.num_splits_buffer.copy_(num_splits)

621 + # num_splits has size [num_decodes + 1]

622 + num_splits_view = self.num_splits_buffer[: num_decodes + 1]

623 + num_splits_view.copy_(num_splits)

624

353 - fp8_extra_metadata = FlashMLASparseMetadata.FP8KernelMetadata(

625 + kernel_meta = FlashMLASparseMetadata.FP8KernelMetadata(

626 scheduler_metadata=tile_scheduler_metadata_buffer,

355 - num_splits=self.num_splits_buffer,

356 - # cache_lens and block_table are basically unused in sparse case

357 - # but the decode kernel will treat -1 and indices >= cache_lens

358 - # as invalid so we make sure cache_lens is large enough to not

359 - # accidentally mark indices invalid, we will use -1 exclusively

360 - # to mark invalid indices

361 - cache_lens=self.max_model_len_tensor,

362 - dummy_block_table=self.dummy_block_table,

627 + num_splits=num_splits_view,

628 + dummy_block_table=self.dummy_block_table[:num_decodes],

629 + cache_lens=self.max_model_len_tensor[:num_decodes],

630 + )

631 + fp8_metadata.decode = FP8Meta.Decode(

632 + kernel_metadata=kernel_meta,

633 + decode_query_len=decode_query_len,

634 )

635

636 + return fp8_metadata

637 +

638 + def build(

639 + self,

640 + common_prefix_len: int,

641 + common_attn_metadata: CommonAttentionMetadata,

642 + fast_build: bool = False,

643 + ) -> FlashMLASparseMetadata:

644 + cm = common_attn_metadata

645 + num_tokens = cm.num_actual_tokens

646 + starts = np.asarray(cm.query_start_loc_cpu, dtype=np.int32)

647 + seg_lengths = np.diff(starts)

648 + req_id_per_token = np.repeat(

649 + np.arange(seg_lengths.shape[0], dtype=np.int32), seg_lengths

650 + )

651 + # Zero-fill for cudagraphs

652 + self.req_id_per_token_buffer.fill_(0)

653 + self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_(

654 + torch.from_numpy(req_id_per_token), non_blocking=True

655 + )

656 + req_id_per_token = self.req_id_per_token_buffer[:num_tokens]

657 +

658 + fp8_extra_metadata: (

659 + FlashMLASparseMetadata.FP8SeperatePrefillDecode

660 + | FlashMLASparseMetadata.FP8KernelMetadata

661 + | None

662 + ) = None

663 + fp8_use_mixed_batch = self.num_heads < MIN_HEADS_FOR_BF16_PREFILL

664 + if self.use_fp8_kv_cache:

665 + if fp8_use_mixed_batch:

666 + fp8_extra_metadata = self._build_fp8_mixed_decode_prefill(cm)

667 + else:

668 + fp8_extra_metadata = self._build_fp8_separate_prefill_decode(cm)

669 +

670 metadata = FlashMLASparseMetadata(

366 - num_reqs=common_attn_metadata.num_reqs,

367 - max_query_len=common_attn_metadata.max_query_len,

368 - max_seq_len=common_attn_metadata.max_seq_len,

369 - num_actual_tokens=common_attn_metadata.num_actual_tokens,

370 - query_start_loc=common_attn_metadata.query_start_loc,

371 - slot_mapping=common_attn_metadata.slot_mapping,

372 - block_table=common_attn_metadata.block_table_tensor,

671 + num_reqs=cm.num_reqs,

672 + max_query_len=cm.max_query_len,

673 + max_seq_len=cm.max_seq_len,

674 + num_actual_tokens=cm.num_actual_tokens,

675 + query_start_loc=cm.query_start_loc,

676 + slot_mapping=cm.slot_mapping,

677 + block_table=cm.block_table_tensor,

678 req_id_per_token=req_id_per_token,

679 block_size=self.kv_cache_spec.block_size,

680 topk_tokens=self.topk_tokens,

681 fp8_extra_metadata=fp8_extra_metadata,

682 + fp8_use_mixed_batch=fp8_use_mixed_batch,

683 )

684 +

685 return metadata

686

687

@@ -414,12 +721,204 @@ class FlashMLASparseImpl(MLACommonBaseImpl[FlashMLASparseMetadata]):

721 self.topk_indices_buffer = indexer.topk_indices_buffer

722 self.padding = 128 if current_platform.is_device_capability(100) else 64

723

724 + if kv_cache_dtype == "fp8_ds_mla":

725 + # Reserve workspace during initialization

726 + vllm_config = get_current_vllm_config()

727 + assert vllm_config is not None and vllm_config.model_config is not None

728 + prefill_workspace_size = get_prefill_workspace_size(

729 + vllm_config.model_config.max_model_len

730 + )

731 + self.prefill_workspace_shape = (prefill_workspace_size, head_size)

732 + (self.prefill_bf16_workspace,) = (

733 + current_workspace_manager().get_simultaneous(

734 + (self.prefill_workspace_shape, torch.bfloat16)

735 + )

736 + )

737 +

738 def _forward_bf16_kv(

739 self,

740 q: torch.Tensor,

741 kv_c_and_k_pe_cache: torch.Tensor,

742 topk_indices: torch.Tensor,

743 attn_metadata: FlashMLASparseMetadata,

744 + ) -> torch.Tensor:

745 + # Convert per-request indices to global slots (decode) or workspace

746 + # offsets (prefill).

747 + topk_indices = triton_convert_req_index_to_global_index(

748 + attn_metadata.req_id_per_token,

749 + attn_metadata.block_table,

750 + topk_indices,

751 + BLOCK_SIZE=attn_metadata.block_size,

752 + NUM_TOPK_TOKENS=topk_indices.shape[1],

753 + )

754 +

755 + return self._bf16_flash_mla_kernel(q, kv_c_and_k_pe_cache, topk_indices)

756 +

757 + def _forward_fp8_kv_separate_prefill_decode(

758 + self,

759 + q: torch.Tensor,

760 + kv_c_and_k_pe_cache: torch.Tensor,

761 + topk_indices: torch.Tensor,

762 + attn_metadata: FlashMLASparseMetadata,

763 + ) -> torch.Tensor:

764 + fp8_metadata = attn_metadata.fp8_extra_metadata

765 + assert isinstance(fp8_metadata, FlashMLASparseMetadata.FP8SeperatePrefillDecode)

766 + num_decodes = fp8_metadata.num_decodes

767 +

768 + prefill_request_ids = None

769 + prefill_workspace_starts = None

770 + has_prefill_workspace = False

771 + if fp8_metadata.prefill is not None:

772 + prefill_request_ids = fp8_metadata.prefill.request_ids

773 + prefill_workspace_starts = fp8_metadata.prefill.workspace_starts

774 + has_prefill_workspace = True

775 +

776 + # Convert per-request indices to global slots (decode) or workspace

777 + # offsets (prefill).

778 + # For FP8 cache: prefill uses workspace mapping (upconverted to BF16)

779 + # For BF16 cache: always use global cache slots (no workspace)

780 + # prefill_workspace_starts has been adjusted in-place per chunk so

781 + # prefill indices automatically come out chunk-local

782 + topk_indices = triton_convert_req_index_to_global_index(

783 + attn_metadata.req_id_per_token,

784 + attn_metadata.block_table,

785 + topk_indices,

786 + BLOCK_SIZE=attn_metadata.block_size,

787 + NUM_TOPK_TOKENS=topk_indices.shape[1],

788 + HAS_PREFILL_WORKSPACE=has_prefill_workspace,

789 + prefill_workspace_request_ids=prefill_request_ids,

790 + prefill_workspace_starts=prefill_workspace_starts,

791 + )

792 +

793 + fp8_metadata = attn_metadata.fp8_extra_metadata

794 + assert isinstance(fp8_metadata, FlashMLASparseMetadata.FP8SeperatePrefillDecode)

795 +

796 + def _fp8_decode(q: torch.Tensor, topk_indices: torch.Tensor) -> torch.Tensor:

797 + # Reshape q: (num_decode_tokens, num_heads, head_dim)

798 + # -> (num_decodes, seq_len, num_heads, head_dim)

799 + q = reshape_query_for_spec_decode(q, num_decodes)

800 + seq_len = q.shape[1]

801 + # Reshape topk_indices: (num_decode_tokens, topk)

802 + # -> (num_decodes, seq_len, topk)

803 + topk_indices = topk_indices.view(num_decodes, seq_len, -1)

804 + assert fp8_metadata.decode is not None

805 + attn_out, _ = self._fp8_flash_mla_kernel(

806 + q=q,

807 + kv_c_and_k_pe_cache=kv_c_and_k_pe_cache,

808 + topk_indices=topk_indices,

809 + kernel_metadata=fp8_metadata.decode.kernel_metadata,

810 + )

811 + # Reshape output: (num_decodes, seq_len, num_heads, head_dim_v)

812 + # -> (num_decode_tokens, num_heads, head_dim_v)

813 + return reshape_attn_output_for_spec_decode(attn_out)

814 +

815 + num_decode_tokens = fp8_metadata.num_decode_tokens

816 + num_prefill_tokens = fp8_metadata.num_prefill_tokens

817 +

818 + # Pure decode: direct call without allocation

819 + if num_decode_tokens > 0 and num_prefill_tokens == 0:

820 + assert fp8_metadata.decode is not None

821 + attn_out = _fp8_decode(q, topk_indices)

822 + else:

823 + # Mixed or pure prefill: allocate output tensor

824 + attn_out = q.new_empty(

825 + (attn_metadata.num_actual_tokens, self.num_heads, self.kv_lora_rank),

826 + dtype=q.dtype,

827 + device=q.device,

828 + )

829 +

830 + if num_decode_tokens > 0:

831 + attn_out[:num_decode_tokens] = _fp8_decode(

832 + q[:num_decode_tokens], topk_indices[:num_decode_tokens]

833 + )

834 +

835 + assert fp8_metadata.prefill is not None

836 + for chunk in fp8_metadata.prefill.chunks:

837 + chunk_workspace = self.prefill_bf16_workspace[: chunk.chunk_tot_seqlen]

838 + ops.cp_gather_and_upconvert_fp8_kv_cache(

839 + kv_c_and_k_pe_cache,

840 + chunk_workspace,

841 + chunk.block_table,

842 + chunk.seq_lens,

843 + chunk.workspace_starts,

844 + len(chunk.block_table),

845 + )

846 +

847 + chunk_q = q[chunk.tokens_slice]

848 + chunk_topk_indices_workspace = topk_indices[chunk.tokens_slice]

849 +

850 + attn_out[chunk.tokens_slice] = self._bf16_flash_mla_kernel(

851 + chunk_q,

852 + chunk_workspace,

853 + chunk_topk_indices_workspace,

854 + )

855 +

856 + return attn_out

857 +

858 + def _forward_fp8_kv_mixed_batch(

859 + self,

860 + q: torch.Tensor,

861 + kv_c_and_k_pe_cache: torch.Tensor,

862 + topk_indices: torch.Tensor,

863 + attn_metadata: FlashMLASparseMetadata,

864 + ) -> torch.Tensor:

865 + """Mixed batch FP8 forward path that treats all tokens as one batch.

866 +

867 + This is equivalent to main branch's approach and avoids the BF16

868 + prefill kernel which has head padding overhead when num_heads is small.

869 + Used when use_mixed_batch is True.

870 + """

871 + # Convert per-request indices to global slots (decode) or workspace

872 + # offsets (prefill).

873 + topk_indices = triton_convert_req_index_to_global_index(

874 + attn_metadata.req_id_per_token,

875 + attn_metadata.block_table,

876 + topk_indices,

877 + BLOCK_SIZE=attn_metadata.block_size,

878 + NUM_TOPK_TOKENS=topk_indices.shape[1],

879 + )

880 +

881 + assert attn_metadata.fp8_extra_metadata is not None

882 + assert isinstance(

883 + attn_metadata.fp8_extra_metadata, FlashMLASparseMetadata.FP8KernelMetadata

884 + )

885 + fp8_metadata = attn_metadata.fp8_extra_metadata

886 +

887 + _attn_out, _ = self._fp8_flash_mla_kernel(

888 + q=q.unsqueeze(0), # unsqueeze to add batch_dim: (T, H, D) -> (1, T, H, D)

889 + kv_c_and_k_pe_cache=kv_c_and_k_pe_cache,

890 + topk_indices=topk_indices.unsqueeze(0), # (T, topk) -> (1, T, topk)

891 + kernel_metadata=fp8_metadata,

892 + )

893 +

894 + # Output is (1, T, H, D_v), squeeze back to (T, H, D_v)

895 + return _attn_out.squeeze(0)

896 +

897 + def _fp8_flash_mla_kernel(

898 + self,

899 + q: torch.Tensor,

900 + kv_c_and_k_pe_cache: torch.Tensor,

901 + topk_indices: torch.Tensor,

902 + kernel_metadata: FlashMLASparseMetadata.FP8KernelMetadata,

903 + ) -> torch.Tensor:

904 + return flash_mla_with_kvcache(

905 + q=q,

906 + k_cache=kv_c_and_k_pe_cache.view(torch.uint8).unsqueeze(-2),

907 + block_table=kernel_metadata.dummy_block_table,

908 + head_dim_v=512,

909 + cache_seqlens=kernel_metadata.cache_lens,

910 + tile_scheduler_metadata=kernel_metadata.scheduler_metadata,

911 + num_splits=kernel_metadata.num_splits,

912 + is_fp8_kvcache=True,

913 + indices=topk_indices,

914 + softmax_scale=self.softmax_scale,

915 + )

916 +

917 + def _bf16_flash_mla_kernel(

918 + self,

919 + q: torch.Tensor,

920 + kv_c_and_k_pe_cache: torch.Tensor,

921 + topk_indices: torch.Tensor,

922 ) -> torch.Tensor:

923 num_tokens = q.shape[0]

924 kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.view(

@@ -445,31 +944,6 @@ class FlashMLASparseImpl(MLACommonBaseImpl[FlashMLASparseMetadata]):

944 output = output[:, : self.num_heads, :]

945 return output

946

448 - def _forward_fp8_kv(

449 - self,

450 - q: torch.Tensor,

451 - kv_c_and_k_pe_cache: torch.Tensor,

452 - topk_indices: torch.Tensor,

453 - attn_metadata: FlashMLASparseMetadata,

454 - ) -> torch.Tensor:

455 - assert attn_metadata.fp8_extra_metadata is not None

456 - extra_metadata = attn_metadata.fp8_extra_metadata

457 -

458 - _attn_out, _ = flash_mla_with_kvcache(

459 - q=q.unsqueeze(0), # unsqueeze to add batch_dim

460 - k_cache=kv_c_and_k_pe_cache.view(torch.uint8).unsqueeze(-2),

461 - block_table=extra_metadata.dummy_block_table,

462 - head_dim_v=512,

463 - cache_seqlens=extra_metadata.cache_lens,

464 - tile_scheduler_metadata=extra_metadata.scheduler_metadata,

465 - num_splits=extra_metadata.num_splits,

466 - is_fp8_kvcache=True,

467 - indices=topk_indices.unsqueeze(0), # unsqueeze to add batch_dim

468 - softmax_scale=self.softmax_scale,

469 - )

470 -

471 - return _attn_out

472 -

947 def forward(

948 self,

949 layer: AttentionLayer,

@@ -477,7 +951,7 @@ class FlashMLASparseImpl(MLACommonBaseImpl[FlashMLASparseMetadata]):

951 k_c_normed: torch.Tensor, # key in unified attn

952 k_pe: torch.Tensor, # value in unified attn

953 kv_cache: torch.Tensor,

480 - attn_metadata: FlashMLASparseMetadata,

954 + attn_metadata: FlashMLASparseMetadata | None,

955 output: torch.Tensor | None = None,

956 output_scale: torch.Tensor | None = None,

957 output_block_scale: torch.Tensor | None = None,

@@ -493,6 +967,7 @@ class FlashMLASparseImpl(MLACommonBaseImpl[FlashMLASparseMetadata]):

967 )

968

969 if attn_metadata is None:

970 + # Dummy run - no need to allocate buffers

971 # The zero fill is required when used with DP + EP

972 # to ensure all ranks within a DP group compute the

973 # same expert outputs.

@@ -505,6 +980,7 @@ class FlashMLASparseImpl(MLACommonBaseImpl[FlashMLASparseMetadata]):

980 q = q[:num_actual_toks, ...]

981 k_c_normed = k_c_normed[:num_actual_toks, ...]

982 k_pe = k_pe[:num_actual_toks, ...]

983 + topk_indices = self.topk_indices_buffer[:num_actual_toks]

984

985 q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)

986 # Convert from (B, N, P) to (N, B, P)

@@ -514,16 +990,7 @@ class FlashMLASparseImpl(MLACommonBaseImpl[FlashMLASparseMetadata]):

990 # Convert from (N, B, L) to (B, N, L)

991 ql_nope = ql_nope.transpose(0, 1)

992

517 - topk_indices = self.topk_indices_buffer[:num_actual_toks]

518 -

519 - # TODO: handle index / kv_cache correctly

520 - topk_indices_global = triton_convert_req_index_to_global_index(

521 - attn_metadata.req_id_per_token,

522 - attn_metadata.block_table,

523 - topk_indices,

524 - BLOCK_SIZE=attn_metadata.block_size,

525 - NUM_TOPK_TOKENS=attn_metadata.topk_tokens,

526 - )

993 + use_fp8_cache = self.kv_cache_dtype == "fp8_ds_mla"

994

995 q = torch.cat([ql_nope, q_pe], dim=-1)

996

@@ -538,13 +1005,15 @@ class FlashMLASparseImpl(MLACommonBaseImpl[FlashMLASparseMetadata]):

1005 scale=layer._k_scale,

1006 )

1007

541 - if self.kv_cache_dtype != "fp8_ds_mla":

542 - attn_out = self._forward_bf16_kv(

543 - q, kv_cache, topk_indices_global, attn_metadata

1008 + if not use_fp8_cache:

1009 + attn_out = self._forward_bf16_kv(q, kv_cache, topk_indices, attn_metadata)

1010 + elif attn_metadata.fp8_use_mixed_batch:

1011 + attn_out = self._forward_fp8_kv_mixed_batch(

1012 + q, kv_cache, topk_indices, attn_metadata

1013 )

1014 else:

546 - attn_out = self._forward_fp8_kv(

547 - q, kv_cache, topk_indices_global, attn_metadata

1015 + attn_out = self._forward_fp8_kv_separate_prefill_decode(

1016 + q, kv_cache, topk_indices, attn_metadata

1017 )

1018

1019 self._v_up_proj(attn_out, out=output[:num_actual_toks])

vllm/v1/attention/backends/mla/indexer.py

+12 -36

@@ -18,6 +18,7 @@ from vllm.v1.attention.backends.utils import (

18 AttentionMetadataBuilder,

19 CommonAttentionMetadata,

20 split_decodes_and_prefills,

21 + split_prefill_chunks,

22 )

24 logger = init_logger(__name__)

@@ -176,40 +177,15 @@ def kv_spans_from_batches(

177

178 def get_max_prefill_buffer_size(vllm_config: VllmConfig):

179 max_model_len = vllm_config.model_config.max_model_len

179 - # NOTE(Chen): 2 is a magic number for controlling the prefill buffer size.

180 - # May be tuned later.

181 - return max_model_len * 2

182 -

183 -

184 -def split_prefill_chunks(

185 - seq_lens_cpu: torch.Tensor, max_prefill_buffer_size: int, reqs_start: int

186 -) -> list[tuple[int, int]]:

187 - """

188 - Split the prefill chunks into a list of tuples of (reqs_start, reqs_end)

189 - such that the total sequence length of each chunk is less than the

190 - maximum prefill buffer size.

191 -

192 - Args:

193 - seq_lens_cpu: The sequence lengths of the prefill requests.

194 - max_prefill_buffer_size: The maximum prefill buffer size.

195 - reqs_start: The start index of the prefill requests.

196 -

197 - Returns:

198 - A list of tuples of (reqs_start, reqs_end).

199 - """

200 - chunk_seq_ids = []

201 - total_seq_lens = 0

202 - for i in range(reqs_start, len(seq_lens_cpu)):

203 - cur_seq_len = seq_lens_cpu[i].item()

204 - assert cur_seq_len <= max_prefill_buffer_size

205 - total_seq_lens += cur_seq_len

206 - if total_seq_lens > max_prefill_buffer_size:

207 - chunk_seq_ids.append((reqs_start, i))

208 - reqs_start = i

209 - total_seq_lens = cur_seq_len

210 - if total_seq_lens > 0:

211 - chunk_seq_ids.append((reqs_start, len(seq_lens_cpu)))

212 - return chunk_seq_ids

180 + # NOTE(Chen): 40 is a magic number for controlling the prefill buffer size.

181 + # Each entry is 128 fp8 bytes and 4 scale bytes for a total of 132 bytes.

182 + # The flashmla_sparse backend uses a workspace size of 5 * max_model_len.

183 + # The memory usage of the workspace there is 576 * 2 bytes; so we size this as

184 + # (576 * 2 // 132) * 5 = 40 to maximize this workspace size while still fitting

185 + # within the flashmla_sparse workspace.

186 + # For DeepSeek-V3.2, the max_model_len is 163840.

187 + # 40 * 163840 * 132 = 865075200 bytes = 825 MB

188 + return max_model_len * 40

189

190

191 class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

@@ -302,9 +278,9 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

278 prefill_metadata = None

279 if num_prefills > 0:

280 chunk_seq_ids = split_prefill_chunks(

305 - common_attn_metadata.seq_lens_cpu,

281 + common_attn_metadata.seq_lens_cpu[num_decodes:],

282 self.max_prefill_buffer_size,

307 - num_decodes,

283 + request_offset=num_decodes,

284 )

285 chunks = [

286 self.build_one_prefill_chunk(

vllm/v1/attention/backends/utils.py

+27

@@ -937,6 +937,33 @@ def split_decodes_and_prefills(

937 return (num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens)

938

939

940 +def split_prefill_chunks(

941 + seq_lens_cpu: torch.Tensor, workspace_size: int, request_offset: int = 0

942 +) -> list[tuple[int, int]]:

943 + """

944 + Split the prefill requests into chunks such that the total sequence length

945 + of each chunk is less than or equal to the workspace size.

946 +

947 + Args:

948 + seq_lens_cpu: The sequence lengths of the prefill requests on CPU.

949 + workspace_size: The maximum workspace size (in tokens) per chunk.

950 + request_offset: The offset to add to the request indices.

951 + Returns:

952 + A list of tuples of (reqs_start, reqs_end) representing chunk boundaries.

953 + """

954 + chunk_bounds = []

955 + i, n = 0, len(seq_lens_cpu)

956 + assert torch.all(seq_lens_cpu <= workspace_size).item()

957 +

958 + while i < n:

959 + start, chunk_total = i, 0

960 + while i < n and (chunk_total + (s := seq_lens_cpu[i].item())) <= workspace_size:

961 + chunk_total += s

962 + i += 1

963 + chunk_bounds.append((start + request_offset, i + request_offset))

964 + return chunk_bounds

965 +

966 +

967 def reorder_batch_to_split_decodes_and_prefills(

968 input_batch: "InputBatch",

969 scheduler_output: "SchedulerOutput",

vllm/v1/worker/gpu_model_runner.py

@@ -162,6 +162,7 @@ from vllm.v1.worker.ubatch_utils import (

162 maybe_create_ubatch_slices,

163 )

164 from vllm.v1.worker.utils import is_residual_scattered_for_sp

165 +from vllm.v1.worker.workspace import lock_workspace

166

167 from .utils import (

168 AttentionGroup,

@@ -297,6 +298,7 @@ class GPUModelRunner(

298 self.device = device

299 self.pin_memory = is_pin_memory_available()

300 self.dtype = self.model_config.dtype

301 +

302 self.kv_cache_dtype = kv_cache_dtype_str_to_dtype(

303 cache_config.cache_dtype, self.model_config

304 )

@@ -4597,6 +4599,10 @@ class GPUModelRunner(

4599 # after here.

4600 set_cudagraph_capturing_enabled(False)

4601

4602 + # Lock workspace to prevent resizing during execution.

4603 + # Max workspace sizes should have been captured during warmup/profiling.

4604 + lock_workspace()

4605 +

4606 end_time = time.perf_counter()

4607 elapsed_time = end_time - start_time

4608 cuda_graph_size = start_free_gpu_memory - end_free_gpu_memory

vllm/v1/worker/gpu_worker.py

@@ -54,6 +54,7 @@ from vllm.v1.outputs import (

54 from vllm.v1.utils import report_usage_stats

55 from vllm.v1.worker.utils import is_residual_scattered_for_sp

56 from vllm.v1.worker.worker_base import WorkerBase

57 +from vllm.v1.worker.workspace import init_workspace_manager

59 logger = init_logger(__name__)

@@ -255,6 +256,10 @@ class Worker(WorkerBase):

256 else:

257 raise RuntimeError(f"Not support device type: {self.device_config.device}")

258

259 + # Initialize workspace manager

260 + num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1

261 + init_workspace_manager(self.device, num_ubatches)

262 +

263 # Construct the model runner

264 if self.use_v2_model_runner:

265 from vllm.v1.worker.gpu.model_runner import (

vllm/v1/worker/workspace.py new

+245

@@ -0,0 +1,245 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +import inspect

5 +import os

6 +from itertools import accumulate

7 +from math import prod

8 +from typing import Optional

9 +

10 +import torch

11 +

12 +import vllm.envs as envs

13 +from vllm.logger import init_logger

14 +from vllm.utils.math_utils import round_up

15 +from vllm.v1.worker.ubatching import dbo_current_ubatch_id

16 +

17 +logger = init_logger(__name__)

18 +

19 +

20 +def _compute_bytes(shape: tuple[int, ...], dtype: torch.dtype) -> int:

21 + return prod(shape) * dtype.itemsize

22 +

23 +

24 +# Constants

25 +_MB = 1024**2

26 +_GiB = 1024**3

27 +

28 +# Global workspace manager instance

29 +_manager: Optional["WorkspaceManager"] = None

30 +

31 +

32 +class WorkspaceManager:

33 + """Manager for workspace allocation.

34 +

35 + Manages workspace buffers for DBO (Dual Batch Overlap) execution.

36 + Can be locked to prevent further growth during execution.

37 + """

38 +

39 + def __init__(self, device: torch.device, num_ubatches: int | None = None):

40 + self._device = device

41 + # Cache num ubatches at init based on configuration (default to 1)

42 + self._num_ubatches = num_ubatches if num_ubatches is not None else 1

43 + self._current_workspaces: list[torch.Tensor | None] = [None, None]

44 + self._locked: bool = False

45 +

46 + @staticmethod

47 + def _workspace_size_bytes(workspace: torch.Tensor | None) -> int:

48 + """Get size of workspace in bytes."""

49 + if workspace is None:

50 + return 0

51 + return workspace.numel() * workspace.element_size()

52 +

53 + def lock(self) -> None:

54 + """Lock the workspace to prevent further growth.

55 +

56 + After locking, any attempt to allocate a larger workspace will raise

57 + an assertion error. This ensures workspace size is fixed during execution.

58 + """

59 + self._locked = True

60 + if envs.VLLM_DEBUG_WORKSPACE:

61 + logger.info(

62 + "[WORKSPACE DEBUG] Workspace locked. Current sizes: %s",

63 + [

64 + self._workspace_size_bytes(ws) / _MB

65 + for ws in self._current_workspaces

66 + if ws is not None

67 + ],

68 + )

69 +

70 + def is_locked(self) -> bool:

71 + """Check if workspace is locked."""

72 + return self._locked

73 +

74 + def get_simultaneous(

75 + self, *shapes_and_dtypes: tuple[tuple[int, ...], torch.dtype]

76 + ) -> list[torch.Tensor]:

77 + """Get multiple workspace tensors simultaneously from a single allocation.

78 +

79 + Args:

80 + *shapes_and_dtypes: One or more (shape, dtype) tuples.

81 +

82 + Returns:

83 + List of tensor views into the workspace buffer, one per shape/dtype pair.

84 + """

85 + actual_bytes = [_compute_bytes(s, d) for s, d in shapes_and_dtypes]

86 + aligned_bytes = [round_up(actual, 256) for actual in actual_bytes]

87 + total_bytes = sum(aligned_bytes)

88 +

89 + # Calculate cumulative offsets using itertools.accumulate

90 + offsets = list(accumulate([0] + aligned_bytes[:-1]))

91 +

92 + current_workspace = self._ensure_workspace_size(total_bytes)

93 +

94 + return [

95 + current_workspace[offsets[i] : offsets[i] + actual_bytes[i]]

96 + .view(shapes_and_dtypes[i][1])

97 + .reshape(shapes_and_dtypes[i][0])

98 + for i in range(len(shapes_and_dtypes))

99 + ]

100 +

101 + def _ensure_workspace_size(self, required_bytes: int) -> torch.Tensor:

102 + """Ensure workspace is allocated and large enough, return current workspace.

103 +

104 + Args:

105 + required_bytes: The number of bytes required.

106 +

107 + Returns:

108 + The current workspace tensor.

109 + """

110 + ubatch_id = dbo_current_ubatch_id()

111 + current_workspace = self._current_workspaces[ubatch_id]

112 + current_size = self._workspace_size_bytes(current_workspace)

113 +

114 + if current_size < required_bytes:

115 +

116 + def get_caller_info() -> str:

117 + """Find first frame outside WorkspaceManager."""

118 + curr_frame = inspect.currentframe()

119 + if curr_frame is None:

120 + return "unknown"

121 + # Walk up the stack skipping WorkspaceManager frames

122 + curr_frame = curr_frame.f_back

123 + while curr_frame is not None:

124 + # TODO: This only catches instance methods (self), missing

125 + # classmethods and staticmethods. Once Python 3.11+ is the

126 + # minimum supported version, use co_qualname instead:

127 + # qualname = curr_frame.f_code.co_qualname

128 + # if qualname.startswith("WorkspaceManager."):

129 + if isinstance(curr_frame.f_locals.get("self"), WorkspaceManager):

130 + curr_frame = curr_frame.f_back

131 + continue

132 + filename = os.path.basename(curr_frame.f_code.co_filename)

133 + return (

134 + f"{filename}:{curr_frame.f_lineno}:{curr_frame.f_code.co_name}"

135 + )

136 + return "unknown"

137 +

138 + if self._locked:

139 + raise AssertionError(

140 + f"Workspace is locked but allocation from '{get_caller_info()}' "

141 + f"requires {required_bytes / _MB:.2f} MB, current size is "

142 + f"{current_size / _MB:.2f} MB. "

143 + "Workspace growth is not allowed after locking."

144 + )

145 +

146 + for ubatch_id in range(self._num_ubatches):

147 + current_workspace = self._current_workspaces[ubatch_id]

148 + if current_workspace is None:

149 + self._current_workspaces[ubatch_id] = torch.empty(

150 + (required_bytes,), dtype=torch.uint8, device=self._device

151 + )

152 + elif self._workspace_size_bytes(current_workspace) < required_bytes:

153 + current_workspace.resize_(required_bytes)

154 +

155 + if envs.VLLM_DEBUG_WORKSPACE:

156 + logger.info(

157 + "[WORKSPACE DEBUG] Resized workspace from '%s': %.2f MB -> "

158 + "%.2f MB (%d ubatches, total memory %.2f MB)",

159 + get_caller_info(),

160 + current_size / _MB,

161 + required_bytes / _MB,

162 + self._num_ubatches,

163 + required_bytes * self._num_ubatches / _MB,

164 + )

165 +

166 + current_workspace = self._current_workspaces[dbo_current_ubatch_id()]

167 +

168 + return current_workspace

169 +

170 +

171 +def is_workspace_manager_initialized() -> bool:

172 + """Check if workspace manager has been initialized.

173 +

174 + Returns:

175 + True if workspace manager is initialized, False otherwise.

176 + """

177 + return _manager is not None

178 +

179 +

180 +def current_workspace_manager() -> "WorkspaceManager":

181 + """Get the current workspace manager instance.

182 +

183 + Raises:

184 + AssertionError: If workspace manager has not been initialized.

185 + """

186 + assert _manager is not None, (

187 + "WorkspaceManager not initialized. Call init_workspace_manager() "

188 + "with a device before using workspace functions."

189 + )

190 + return _manager

191 +

192 +

193 +def init_workspace_manager(

194 + device: torch.device, num_ubatches: int | None = None

195 +) -> None:

196 + """Initialize the workspace manager with a device.

197 +

198 + Must be called before using any workspace functions. Typically called

199 + from GPUModelRunner.__init__.

200 +

201 + Args:

202 + device: The device to allocate workspace on.

203 + num_ubatches: Number of micro-batches. Defaults to 1.

204 + """

205 + global _manager

206 + if _manager is not None:

207 + logger.warning(

208 + "WorkspaceManager already initialized on device %s, "

209 + "reinitializing on device %s",

210 + _manager._device,

211 + device,

212 + )

213 + _manager = WorkspaceManager(device, num_ubatches)

214 +

215 +

216 +def lock_workspace() -> None:

217 + """Lock the workspace to prevent further growth.

218 +

219 + After calling this function, any attempt to allocate a workspace larger

220 + than the current size will raise an AssertionError. This ensures that

221 + workspace size is fixed during execution and prevents unexpected memory

222 + allocations in the hot path.

223 +

224 + Example:

225 + # During initialization

226 + init_workspace_manager(device)

227 + reserve_workspace(shape1, dtype1)

228 + reserve_workspace(shape2, dtype2)

229 +

230 + # Lock after warmup/profiling

231 + lock_workspace()

232 +

233 + # Now all get_workspace calls must fit in pre-allocated size

234 + """

235 + current_workspace_manager().lock()

236 +

237 +

238 +def reset_workspace_manager() -> None:

239 + """Reset the workspace manager to uninitialized state.

240 +

241 + This is primarily intended for testing purposes to allow tests

242 + to reinitialize the workspace manager cleanly.

243 + """

244 + global _manager

245 + _manager = None