aiden-dcp2.patch 68 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597
  1. --- a/vllm/distributed/communication_op.py 2026-09-12 07:48:35.828568398 +0200
  2. +++ b/vllm/distributed/communication_op.py 2026-09-11 10:13:47.000000000 +0200
  3. @@ -9,8 +9,17 @@
  4. from .parallel_state import get_tp_group
  5. +_DSV41_PREFILL_TIMING = __import__("os").environ.get("DSV41_PREFILL_TIMING", "0") == "1"
  6. +
  7. +
  8. def tensor_model_parallel_all_reduce(input_: torch.Tensor) -> torch.Tensor:
  9. """All-reduce the input tensor across model parallel group."""
  10. + if _DSV41_PREFILL_TIMING and input_.dim() >= 2 and input_.shape[0] >= 1024:
  11. + # Stage 7: phase "allreduce" of the prefill budget.
  12. + from vllm.model_executor.kernels.attention.dsa.candidate_blocks import _PrefillTimer
  13. +
  14. + with _PrefillTimer("allreduce", int(input_.shape[0])):
  15. + return get_tp_group().all_reduce(input_)
  16. return get_tp_group().all_reduce(input_)
  17. --- a/vllm/model_executor/layers/fused_moe/b12x.py 2026-09-12 07:48:35.776439710 +0200
  18. +++ b/vllm/model_executor/layers/fused_moe/b12x.py 2026-09-11 10:13:47.000000000 +0200
  19. @@ -755,16 +755,20 @@
  20. )
  21. scratch = _workspace_as_b12x_scratch(workspace2, plan)
  22. - _run_b12x_moe_plan(
  23. - plan=plan,
  24. - scratch=scratch,
  25. - hidden_states=hidden_states,
  26. - prepared=prepared,
  27. - topk_weights=topk_weights,
  28. - topk_ids=topk_ids,
  29. - output=output,
  30. - unit_scale_contract=self._quant_mode == "w4a16",
  31. - )
  32. + # Stage 7: phase "moe" of the prefill budget (inert unless DSV41_PREFILL_TIMING=1).
  33. + from vllm.model_executor.kernels.attention.dsa.candidate_blocks import _PrefillTimer
  34. +
  35. + with _PrefillTimer("moe", int(hidden_states.shape[0])):
  36. + _run_b12x_moe_plan(
  37. + plan=plan,
  38. + scratch=scratch,
  39. + hidden_states=hidden_states,
  40. + prepared=prepared,
  41. + topk_weights=topk_weights,
  42. + topk_ids=topk_ids,
  43. + output=output,
  44. + unit_scale_contract=self._quant_mode == "w4a16",
  45. + )
  46. def moe_sum(self, input: torch.Tensor, output: torch.Tensor) -> None:
  47. raise NotImplementedError("LoRA is not supported for B12xExperts")
  48. --- a/vllm/model_executor/layers/sparse_attn_indexer.py 2026-09-12 07:48:35.313947794 +0200
  49. +++ b/vllm/model_executor/layers/sparse_attn_indexer.py 2026-09-11 10:13:47.000000000 +0200
  50. @@ -13,6 +13,15 @@
  51. from vllm.forward_context import get_forward_context
  52. from vllm.logger import init_logger
  53. from vllm.model_executor.custom_op import CustomOp
  54. +from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (
  55. + apply_candidate_mask as _apply_candidate_mask,
  56. +)
  57. +from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (
  58. + _PrefillTimer as _DSV41PrefillTimer,
  59. +)
  60. +from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (
  61. + select_candidate_blocks as _select_candidate_blocks,
  62. +)
  63. from vllm.model_executor.layers.quantization.utils.quant_utils import (
  64. get_fp8_min_max,
  65. )
  66. @@ -316,6 +325,10 @@
  67. dcp_world_size: int = 1,
  68. cp_kv_cache_interleave_size: int = 1,
  69. skip_topk_buffer_clear: bool = False,
  70. + candidate_blocks: torch.Tensor | None = None,
  71. + candidate_block_size: int = 0,
  72. + candidate_write: bool = False,
  73. + compress_ratio: int = 1,
  74. ) -> torch.Tensor:
  75. # careful! this will be None in dummy run
  76. forward_context = get_forward_context()
  77. @@ -323,6 +336,26 @@
  78. fp8_dtype = current_platform.fp8_dtype()
  79. k_cache_prefix = _resolve_layer_name(k_cache_prefix)
  80. + # v4.1 two-level candidate filtering under DCP: the logits columns are this
  81. + # rank's local compressed states; with interleave 1 and world | block, local
  82. + # block b (size block // world) is exactly this rank's shard of global block
  83. + # b, so candidate ids need no translation. The producer MAX-reduces block
  84. + # scores across ranks and pins the global newest block (candidate_blocks.py).
  85. + local_cand_bs = candidate_block_size
  86. + dcp_group = None
  87. + if candidate_blocks is not None:
  88. + assert candidate_block_size > 0
  89. + if dcp_world_size > 1:
  90. + assert cp_kv_cache_interleave_size == 1, (
  91. + "v4.1 candidate filtering under DCP needs cp_kv_cache_interleave_size=1"
  92. + )
  93. + assert candidate_block_size % dcp_world_size == 0, (
  94. + f"candidate_block_size ({candidate_block_size}) must be a multiple "
  95. + f"of dcp_world_size ({dcp_world_size})"
  96. + )
  97. + local_cand_bs = candidate_block_size // dcp_world_size
  98. + dcp_group = get_dcp_group()
  99. +
  100. # assert isinstance(attn_metadata, dict)
  101. if not isinstance(attn_metadata, dict):
  102. # Reserve workspace for indexer during profiling run
  103. @@ -361,6 +394,9 @@
  104. use_pcp,
  105. dense_mha_metadata_layer_name,
  106. use_fp4_cache,
  107. + candidate_blocks=candidate_blocks,
  108. + candidate_block_size=candidate_block_size,
  109. + candidate_write=candidate_write,
  110. )
  111. attn_metadata_narrowed = attn_metadata[k_cache_prefix]
  112. assert isinstance(attn_metadata_narrowed, DeepseekV32IndexerMetadata)
  113. @@ -474,6 +510,23 @@
  114. if chunk.local_total_seq_lens == 0:
  115. logits = q_slice.new_empty((q_slice.shape[0], 0), dtype=torch.float32)
  116. topk_indices.fill_(-1)
  117. + if candidate_blocks is not None and candidate_write and dcp_group is not None:
  118. + # Other ranks may hold states for this chunk: join the MAX
  119. + # reduce with an empty local contribution.
  120. + _select_candidate_blocks(
  121. + logits,
  122. + cu_seqlen_ks,
  123. + cu_seqlen_ke,
  124. + candidate_blocks.shape[1],
  125. + local_cand_bs,
  126. + candidate_blocks[chunk.token_start : chunk.token_end],
  127. + dcp_group=dcp_group,
  128. + global_row_ke=chunk.global_ctx_lens,
  129. + global_block_size=candidate_block_size,
  130. + max_global_blocks=triton.cdiv(
  131. + chunk.max_global_ctx, candidate_block_size
  132. + ),
  133. + )
  134. else:
  135. # DeepGEMM scalar-type tags (zero-copy): MXFP4 values → int8
  136. # (kPackedFP4), scales → int32 squeezed to 1-D kv_sf / 2-D q_sf.
  137. @@ -506,6 +559,38 @@
  138. clean_logits=False,
  139. )
  140. num_rows = logits.shape[0]
  141. + if candidate_blocks is not None:
  142. + # Two-level selection (v4.1): the candidate source
  143. + # publishes its top blocks; later indexers mask their
  144. + # scores to them. Both before the row top-k.
  145. + chunk_candidates = candidate_blocks[
  146. + chunk.token_start : chunk.token_end
  147. + ]
  148. + if candidate_write:
  149. + _select_candidate_blocks(
  150. + logits,
  151. + cu_seqlen_ks,
  152. + cu_seqlen_ke,
  153. + chunk_candidates.shape[1],
  154. + local_cand_bs,
  155. + chunk_candidates,
  156. + dcp_group=dcp_group,
  157. + global_row_ke=chunk.global_ctx_lens,
  158. + global_block_size=candidate_block_size,
  159. + max_global_blocks=(
  160. + triton.cdiv(chunk.max_global_ctx, candidate_block_size)
  161. + if dcp_group is not None
  162. + else None
  163. + ),
  164. + )
  165. + else:
  166. + _apply_candidate_mask(
  167. + logits,
  168. + cu_seqlen_ks,
  169. + cu_seqlen_ke,
  170. + chunk_candidates,
  171. + local_cand_bs,
  172. + )
  173. ops.top_k_per_row_prefill(
  174. logits,
  175. cu_seqlen_ks,
  176. @@ -612,6 +697,64 @@
  177. indices=decode_metadata.indices,
  178. )
  179. num_rows = logits.shape[0]
  180. + if candidate_blocks is not None:
  181. + # Two-level selection (v4.1) on the decode logits; columns are
  182. + # request-local compressed positions. seq_lens is (B, next_n)
  183. + # for native spec decode (per-row effective lens) and (B, 1)
  184. + # otherwise.
  185. + vis = seq_lens.reshape(-1)
  186. + row_repeat = next_n if vis.numel() != num_rows else 1
  187. + vis = vis[:num_rows]
  188. + decode_candidates = candidate_blocks[:num_rows]
  189. + if candidate_write:
  190. + global_row_ke = None
  191. + max_global_blocks = None
  192. + if dcp_group is not None:
  193. + # Global (un-sharded) per-row context in STATES for the
  194. + # newest-block pin: per token when flattened / next_n == 1,
  195. + # per request (row_repeat) or expanded per row otherwise.
  196. + gl = decode_metadata.global_seq_lens
  197. + assert gl is not None
  198. + gl = gl.reshape(-1)
  199. + if row_repeat > 1:
  200. + gl = gl[:batch_size]
  201. + elif gl.numel() != num_rows:
  202. + # native spec decode: request lens -> per-row lens
  203. + gl = (
  204. + gl[:batch_size].unsqueeze(1)
  205. + - next_n
  206. + + 1
  207. + + torch.arange(next_n, device=gl.device, dtype=gl.dtype)
  208. + ).clamp_(min=0).reshape(-1)
  209. + else:
  210. + gl = gl[:num_rows]
  211. + global_row_ke = gl // compress_ratio
  212. + max_global_blocks = triton.cdiv(
  213. + attn_metadata_narrowed.max_seq_len // compress_ratio,
  214. + candidate_block_size,
  215. + )
  216. + _select_candidate_blocks(
  217. + logits,
  218. + None,
  219. + vis,
  220. + decode_candidates.shape[1],
  221. + local_cand_bs,
  222. + decode_candidates,
  223. + row_repeat,
  224. + dcp_group=dcp_group,
  225. + global_row_ke=global_row_ke,
  226. + global_block_size=candidate_block_size,
  227. + max_global_blocks=max_global_blocks,
  228. + )
  229. + else:
  230. + _apply_candidate_mask(
  231. + logits,
  232. + None,
  233. + vis,
  234. + decode_candidates,
  235. + local_cand_bs,
  236. + row_repeat,
  237. + )
  238. topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
  239. use_cooperative_topk = (
  240. @@ -622,10 +765,13 @@
  241. and current_platform.has_device_capability(90)
  242. and not current_platform.is_device_capability_family(120)
  243. )
  244. - use_persistent_topk = current_platform.is_cuda() and topk_tokens in (
  245. - 512,
  246. - 1024,
  247. - 2048,
  248. + # Tech2Wild 2026-09-10 (SM12x): GB10 has 48 SMs and 99 KB smem per block. On long
  249. + # rows persistent_topk oversubscribes and its FilteredTopK fallback needs >=128 KB
  250. + # smem, so the launch fails (32K context: 90 CTAs > 48). Use top_k_per_row_decode.
  251. + use_persistent_topk = (
  252. + current_platform.is_cuda()
  253. + and topk_tokens in (512, 1024, 2048)
  254. + and not current_platform.is_device_capability_family(120)
  255. )
  256. if use_cooperative_topk:
  257. workspace_manager = current_workspace_manager()
  258. @@ -712,6 +858,10 @@
  259. dcp_world_size: int = 1,
  260. cp_kv_cache_interleave_size: int = 1,
  261. skip_topk_buffer_clear: bool = False,
  262. + candidate_blocks: torch.Tensor | None = None,
  263. + candidate_block_size: int = 0,
  264. + candidate_write: bool = False,
  265. + compress_ratio: int = 1,
  266. ) -> torch.Tensor:
  267. return topk_indices_buffer
  268. @@ -719,7 +869,7 @@
  269. direct_register_custom_op(
  270. op_name="sparse_attn_indexer",
  271. op_func=sparse_attn_indexer,
  272. - mutates_args=["topk_indices_buffer"],
  273. + mutates_args=["topk_indices_buffer", "candidate_blocks"],
  274. fake_impl=sparse_attn_indexer_fake,
  275. dispatch_key=current_platform.dispatch_key,
  276. )
  277. @@ -751,6 +901,9 @@
  278. skip_k_cache_insert: bool = False,
  279. use_fp4_cache: bool = False,
  280. compress_ratio: int = 1,
  281. + candidate_blocks: torch.Tensor | None = None,
  282. + candidate_block_size: int = 0,
  283. + candidate_write: bool = False,
  284. ):
  285. super().__init__()
  286. self.k_cache = k_cache
  287. @@ -764,6 +917,11 @@
  288. self.skip_k_cache_insert = skip_k_cache_insert
  289. self.use_fp4_cache = use_fp4_cache
  290. self.compress_ratio = compress_ratio
  291. + # v4.1 two-level selection: the candidate source indexer writes the
  292. + # top candidate blocks here; later indexers mask their scores with it.
  293. + self.candidate_blocks = candidate_blocks
  294. + self.candidate_block_size = candidate_block_size
  295. + self.candidate_write = candidate_write
  296. self.dense_mha_metadata_layer_name = ""
  297. # DCP scalars are constant for the run; resolve them here (config is set
  298. # during model construction) and pass them into the custom op, rather
  299. @@ -818,6 +976,17 @@
  300. k: torch.Tensor | None,
  301. weights: torch.Tensor,
  302. ):
  303. + # Stage 7: phase "indexer" of the prefill timing (inert unless enabled).
  304. + with _DSV41PrefillTimer("indexer", int(hidden_states.shape[0])):
  305. + return self._forward_cuda_impl(hidden_states, q_quant, k, weights)
  306. +
  307. + def _forward_cuda_impl(
  308. + self,
  309. + hidden_states: torch.Tensor,
  310. + q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
  311. + k: torch.Tensor | None,
  312. + weights: torch.Tensor,
  313. + ):
  314. # FP8 path: single tensor (per-token scale is folded into `weights`).
  315. # FP4 path: (values, scales) tuple with scales required by the kernel.
  316. if isinstance(q_quant, tuple):
  317. @@ -846,6 +1015,10 @@
  318. self.dcp_rank,
  319. self.dcp_world_size,
  320. self.cp_kv_cache_interleave_size,
  321. + candidate_blocks=self.candidate_blocks,
  322. + candidate_block_size=self.candidate_block_size,
  323. + candidate_write=self.candidate_write,
  324. + compress_ratio=self.compress_ratio,
  325. )
  326. def forward_xpu(
  327. @@ -868,12 +1041,16 @@
  328. assert isinstance(q_quant, torch.Tensor), (
  329. "AMD sparse_attn_indexer expects a single FP8 q_quant tensor"
  330. )
  331. - from vllm.platforms.rocm import on_gfx11
  332. + from vllm.platforms.rocm import on_gfx11, on_gfx950
  333. if (
  334. rocm_aiter_ops.is_enabled()
  335. or rocm_aiter_ops.is_rdna_aiter_enabled()
  336. or on_gfx11()
  337. + # The so-called AITER sparse indexer op has a native gfx950 path:
  338. + # its cache insert, MQA logits, and top-k fallbacks are implemented
  339. + # by local Triton/C++ kernels and do not require the aiter package.
  340. + or on_gfx950()
  341. ):
  342. return torch.ops.vllm.rocm_aiter_sparse_attn_indexer(
  343. hidden_states,
  344. @@ -891,8 +1068,11 @@
  345. self.topk_indices_buffer,
  346. skip_k_cache_insert=self.skip_k_cache_insert,
  347. compress_ratio=self.compress_ratio,
  348. + candidate_blocks=self.candidate_blocks,
  349. + candidate_block_size=self.candidate_block_size,
  350. + candidate_write=self.candidate_write,
  351. )
  352. raise RuntimeError(
  353. - "Sparse attention indexer ROCm path is only supported on AITER. "
  354. - "Please enable aiter with VLLM_ROCM_USE_AITER=1"
  355. + "Sparse attention indexer ROCm path requires AITER or a supported "
  356. + "native architecture (gfx950/gfx11)."
  357. )
  358. --- a/vllm/model_executor/model_loader/weight_utils.py 2026-09-12 07:48:35.320443782 +0200
  359. +++ b/vllm/model_executor/model_loader/weight_utils.py 2026-09-11 10:13:47.000000000 +0200
  360. @@ -9,6 +9,8 @@
  361. import hashlib
  362. import json
  363. import os
  364. +
  365. +_DSV41_ENGRAM_DISK = os.environ.get("DSV41_ENGRAM_DISK", "0") == "1"
  366. import tempfile
  367. import threading
  368. import time
  369. @@ -964,6 +966,10 @@
  370. for name in f.keys(): # noqa: SIM118
  371. if should_skip_weight(name, local_expert_ids):
  372. continue
  373. + # Tech2Wild/Kai 2026-09-10: DSV41_ENGRAM_DISK=1 leaves the
  374. + # Engram tables on disk (see deepseek_v4_1/common/engram.py)
  375. + if _DSV41_ENGRAM_DISK and name.endswith((".engram.embed.weight", ".engram.embed.scale")):
  376. + continue
  377. param = f.get_tensor(name)
  378. yield name, param
  379. --- a/vllm/v1/attention/backend.py 2026-09-12 07:48:35.616608046 +0200
  380. +++ b/vllm/v1/attention/backend.py 2026-09-11 10:13:47.000000000 +0200
  381. @@ -611,6 +611,26 @@
  382. self.vllm_config = vllm_config
  383. self.device = device
  384. self.kernel_block_size: int | None = None
  385. + # The DCP geometry this builder's group actually runs at: sharded
  386. + # groups use the process DCP size; DCP-replicated groups (every rank
  387. + # holds the full KV) behave as dcp=1 for metadata purposes.
  388. + from vllm.v1.kv_cache_interface import is_dcp_replicated_spec
  389. +
  390. + process_dcp_size = vllm_config.parallel_config.decode_context_parallel_size
  391. + self.dcp_replicated = is_dcp_replicated_spec(kv_cache_spec)
  392. + if process_dcp_size > 1 and not self.dcp_replicated:
  393. + from vllm.distributed.parallel_state import get_dcp_group
  394. +
  395. + try:
  396. + self.dcp_world_size = get_dcp_group().world_size
  397. + self.dcp_rank = get_dcp_group().rank_in_group
  398. + except AssertionError:
  399. + # DCP group not initialized (unit tests).
  400. + self.dcp_world_size = process_dcp_size
  401. + self.dcp_rank = 0
  402. + else:
  403. + self.dcp_world_size = 1
  404. + self.dcp_rank = 0
  405. def set_kernel_block_size(self, kernel_block_size: int) -> None:
  406. self.kernel_block_size = kernel_block_size
  407. @@ -650,10 +670,7 @@
  408. max_num_queries_for_spec,
  409. )
  410. - if (
  411. - self.vllm_config.parallel_config.decode_context_parallel_size > 1
  412. - and not supports_dcp_with_varlen
  413. - ):
  414. + if getattr(self, "dcp_world_size", 1) > 1 and not supports_dcp_with_varlen:
  415. self.reorder_batch_threshold = 1
  416. @abstractmethod
  417. --- a/vllm/v1/attention/backends/mla/compressor_utils.py 2026-09-12 07:48:35.555787153 +0200
  418. +++ b/vllm/v1/attention/backends/mla/compressor_utils.py 2026-09-11 10:13:47.000000000 +0200
  419. @@ -38,6 +38,9 @@
  420. compress_ratio: int
  421. triton_block_size: int
  422. block_size: int
  423. + dcp_world: int = 1
  424. + dcp_rank: int = 0
  425. + dcp_interleave: int = 1
  426. @staticmethod
  427. @triton.jit(do_not_specialize=["block_table_stride"])
  428. @@ -51,10 +54,14 @@
  429. # [num_reqs, max_num_blocks]
  430. block_table_ptr,
  431. block_table_stride,
  432. + # states per page on THIS rank (kv_cache_spec.num_states)
  433. block_size,
  434. COMPRESS_RATIO: tl.constexpr,
  435. PAD_ID: tl.constexpr,
  436. TRITON_BLOCK_SIZE: tl.constexpr,
  437. + DCP_WORLD: tl.constexpr,
  438. + DCP_RANK: tl.constexpr,
  439. + DCP_INTERLEAVE: tl.constexpr,
  440. ):
  441. batch_idx = tl.program_id(0)
  442. @@ -71,14 +78,32 @@
  443. pos = start_pos + i + tl.arange(0, TRITON_BLOCK_SIZE)
  444. is_valid = (pos + 1) % COMPRESS_RATIO == 0
  445. - pos_after_compress = pos // COMPRESS_RATIO
  446. + # Global compressed state index of the group this token closes.
  447. + state = pos // COMPRESS_RATIO
  448. + if DCP_WORLD > 1:
  449. + # DCP shards the compressed cache over STATES with the runner's
  450. + # token-space formula (worker/gpu/block_table.py) applied in
  451. + # state space: one block-table entry spans block_size * DCP_WORLD
  452. + # states, this rank owns the states whose interleave group is
  453. + # congruent to DCP_RANK, and stores them contiguously per entry.
  454. + virtual_block_size = block_size * DCP_WORLD
  455. + virtual_block_idx = state // virtual_block_size
  456. + virtual_off = state % virtual_block_size
  457. + is_local = (virtual_off // DCP_INTERLEAVE) % DCP_WORLD == DCP_RANK
  458. + local_off = (virtual_off // (DCP_INTERLEAVE * DCP_WORLD)) * DCP_INTERLEAVE + (
  459. + virtual_off % DCP_INTERLEAVE
  460. + )
  461. + local_state = virtual_block_idx * block_size + local_off
  462. + is_valid = is_valid & is_local
  463. + else:
  464. + local_state = state
  465. - block_ids = pos_after_compress // block_size
  466. + block_ids = local_state // block_size
  467. block_numbers = tl.load(
  468. block_table_ptr + batch_idx * block_table_stride + block_ids,
  469. mask=mask & is_valid,
  470. )
  471. - slot_ids = block_numbers * block_size + pos_after_compress % block_size
  472. + slot_ids = block_numbers * block_size + local_state % block_size
  473. # NOTE
  474. slot_ids = tl.where(is_valid, slot_ids, PAD_ID)
  475. @@ -89,15 +114,29 @@
  476. *,
  477. compress_ratio: int,
  478. block_size: int,
  479. + dcp_world_size: int = 1,
  480. + dcp_rank: int = 0,
  481. + cp_kv_cache_interleave_size: int = 1,
  482. ) -> CompileKey:
  483. return self.CompileKey(
  484. compress_ratio=compress_ratio,
  485. triton_block_size=self.TRITON_BLOCK_SIZE,
  486. block_size=triton_scalar_specialization_rep(block_size),
  487. + dcp_world=dcp_world_size,
  488. + dcp_rank=dcp_rank,
  489. + dcp_interleave=cp_kv_cache_interleave_size,
  490. )
  491. def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
  492. hf_config = vllm_config.model_config.hf_config
  493. + parallel_config = vllm_config.parallel_config
  494. + dcp_world_size = parallel_config.decode_context_parallel_size
  495. + cp_interleave = parallel_config.cp_kv_cache_interleave_size
  496. + dcp_rank = 0
  497. + if dcp_world_size > 1:
  498. + from vllm.distributed import get_dcp_group
  499. +
  500. + dcp_rank = get_dcp_group().rank_in_group
  501. configured_ratios = (
  502. *(getattr(hf_config, "compress_ratios", None) or ()),
  503. getattr(hf_config, "index_kpool", 1) or 1,
  504. @@ -113,6 +152,9 @@
  505. dict(
  506. compress_ratio=ratio,
  507. block_size=vllm_config.cache_config.block_size // ratio,
  508. + dcp_world_size=dcp_world_size,
  509. + dcp_rank=dcp_rank,
  510. + cp_kv_cache_interleave_size=cp_interleave,
  511. )
  512. for ratio in compress_ratios
  513. )
  514. @@ -128,6 +170,9 @@
  515. block_table=int32_ptr,
  516. block_size=compile_key.block_size,
  517. compress_ratio=compile_key.compress_ratio,
  518. + dcp_world_size=compile_key.dcp_world,
  519. + dcp_rank=compile_key.dcp_rank,
  520. + cp_kv_cache_interleave_size=compile_key.dcp_interleave,
  521. )
  522. @kernel_launcher
  523. @@ -139,12 +184,22 @@
  524. block_table: torch.Tensor,
  525. block_size: int,
  526. compress_ratio: int,
  527. + dcp_world_size: int = 1,
  528. + dcp_rank: int = 0,
  529. + cp_kv_cache_interleave_size: int = 1,
  530. ) -> LaunchSpec:
  531. + assert block_size % cp_kv_cache_interleave_size == 0, (
  532. + f"states per page ({block_size}) must be a multiple of "
  533. + f"cp_kv_cache_interleave_size ({cp_kv_cache_interleave_size})"
  534. + )
  535. return (block_table.shape[0],), dict(
  536. block_table_stride=block_table.stride(0),
  537. COMPRESS_RATIO=compress_ratio,
  538. PAD_ID=-1,
  539. TRITON_BLOCK_SIZE=self.TRITON_BLOCK_SIZE,
  540. + DCP_WORLD=dcp_world_size,
  541. + DCP_RANK=dcp_rank,
  542. + DCP_INTERLEAVE=cp_kv_cache_interleave_size,
  543. )
  544. @@ -156,7 +211,17 @@
  545. block_size: int,
  546. compress_ratio: int,
  547. out: torch.Tensor | None = None,
  548. + dcp_world_size: int = 1,
  549. + dcp_rank: int = 0,
  550. + cp_kv_cache_interleave_size: int = 1,
  551. ) -> torch.Tensor:
  552. + """Slot of the compressed state each token closes (PAD otherwise).
  553. +
  554. + ``block_size`` is the number of states per page on this rank. Under DCP
  555. + (``dcp_world_size > 1``) states are sharded across ranks and only this
  556. + rank's states get a slot; the block table is the shared (manager) table
  557. + whose entries each span ``block_size * dcp_world_size`` states.
  558. + """
  559. if out is not None:
  560. # Guard: for padded / invalid sequences.
  561. # Negative positions produce bogus block indices that lead to illegal memory
  562. @@ -176,6 +241,9 @@
  563. block_table,
  564. block_size,
  565. compress_ratio,
  566. + dcp_world_size,
  567. + dcp_rank,
  568. + cp_kv_cache_interleave_size,
  569. )
  570. return slot_mapping
  571. --- a/vllm/v1/attention/backends/mla/indexer.py 2026-09-12 07:48:35.764402198 +0200
  572. +++ b/vllm/v1/attention/backends/mla/indexer.py 2026-09-11 10:13:47.000000000 +0200
  573. @@ -294,7 +294,7 @@
  574. @staticmethod
  575. def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
  576. - return [256]
  577. + return [64 if current_platform.is_device_capability_family(90) else 128]
  578. @dataclass
  579. @@ -314,6 +314,11 @@
  580. local_cu_seq_lens: torch.Tensor | None = None
  581. local_total_seq_lens: int = 0
  582. max_local_total_seq_lens: int = 0
  583. + # Per query token: GLOBAL compressed context length (states), i.e. the
  584. + # un-sharded `cu_seqlen_ke - cu_seqlen_ks`. Under DCP the candidate-block
  585. + # selection pins each row's newest global block from this.
  586. + global_ctx_lens: torch.Tensor | None = None
  587. + max_global_ctx: int = 0
  588. class BuildPrefillChunkMetadataKernel(
  589. @@ -346,6 +351,8 @@
  590. token_to_seq_ptr,
  591. cu_compressed_seq_len_ks_ptr,
  592. cu_compressed_seq_len_ke_ptr,
  593. + # [output_query_len] global (un-sharded) compressed context per token
  594. + global_ctx_len_ptr,
  595. query_slice_start,
  596. query_slice_stop,
  597. DCP_RANK,
  598. @@ -388,6 +395,7 @@
  599. # global per-token length is sharded across ranks.
  600. global_ctx = start_pos + 1 + offset
  601. len_per_token = global_ctx // COMPRESS_RATIO
  602. + tl.store(global_ctx_len_ptr + out_pos, len_per_token, mask=mask)
  603. if DCP_WORLD > 1:
  604. # Per-rank local context length under interleave-aware DCP, matching
  605. # get_dcp_local_seq_lens. K == 1 reduces to (len + world-1-rank)//world.
  606. @@ -478,6 +486,7 @@
  607. token_to_seq=int32_ptr,
  608. cu_compressed_seq_len_ks=int32_ptr,
  609. cu_compressed_seq_len_ke=int32_ptr,
  610. + global_ctx_len=int32_ptr,
  611. query_slice_start=compile_key.query_slice_start,
  612. query_slice_stop=compile_key.query_slice_stop,
  613. DCP_RANK=compile_key.dcp_rank,
  614. @@ -497,6 +506,7 @@
  615. token_to_seq: torch.Tensor,
  616. cu_compressed_seq_len_ks: torch.Tensor,
  617. cu_compressed_seq_len_ke: torch.Tensor,
  618. + global_ctx_len: torch.Tensor,
  619. query_slice_start: int,
  620. query_slice_stop: int,
  621. DCP_RANK: int,
  622. @@ -697,6 +707,9 @@
  623. # so this is None; its own split uses self.decode_threshold.
  624. reorder_batch_threshold: int | None = None
  625. requires_block_table_width = True
  626. + # DeepSeek-V4.1 DSpark under DCP: the draft is SWA-only (no indexer); the
  627. + # target's non-causal verification rows use the DCP indexer path.
  628. + supports_non_causal_multi_token_dcp = True
  629. @classmethod
  630. def get_cudagraph_support(
  631. @@ -812,9 +825,14 @@
  632. assert isinstance(self.kv_cache_spec.tokens_per_state, int)
  633. self.compress_ratio = self.kv_cache_spec.tokens_per_state
  634. if self.dcp_world_size > 1 and self.compress_ratio > 1:
  635. - raise NotImplementedError(
  636. - "DCP is not supported with sparse indexer KV compression "
  637. - f"(compress_ratio={self.compress_ratio})."
  638. + # DCP + compressed indexer cache (DeepSeek-V4.1 ratio-2 index
  639. + # sources): the K cache is sharded over compressed STATES; the
  640. + # compressed slot mapping and the decode bounds below are computed
  641. + # in state coordinates (DCP2-DESIGN-stage23.md §A.1).
  642. + logger.info_once(
  643. + "DSA indexer: DCP%d with compress_ratio=%d (state-sharded K cache)",
  644. + self.dcp_world_size,
  645. + self.compress_ratio,
  646. )
  647. # Pre-allocate buffers for CUDA graph compatibility when
  648. @@ -1057,6 +1075,10 @@
  649. and self.kv_cache_spec.block_size % kernel_block_size == 0
  650. ):
  651. factor = self.kv_cache_spec.block_size // kernel_block_size
  652. + assert self.dcp_world_size == 1, (
  653. + "DCP with a compressed indexer cache needs kernel_block_size == "
  654. + f"kv_cache_spec.block_size (got factor {factor})."
  655. + )
  656. indexer_block_table = (block_table[:, ::factor] // factor).contiguous()
  657. padded_num_tokens = num_tokens
  658. if self.pcp_world_size > 1:
  659. @@ -1069,6 +1091,9 @@
  660. self.kv_cache_spec.num_states,
  661. self.compress_ratio,
  662. out=self.compressed_slot_mapping_buffer,
  663. + dcp_world_size=self.dcp_world_size,
  664. + dcp_rank=self.dcp_rank,
  665. + cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
  666. )
  667. if self.pcp_world_size > 1:
  668. compressed_slot_mapping = get_pcp_group().all_gather(
  669. @@ -1231,20 +1256,16 @@
  670. )
  671. block_table = self.indexer_decode_block_table_buffer[:rows, :cols]
  672. - seq_lens_is_buffer_view = (use_native and next_n > 1) or (
  673. - not use_native and max_decode_len > 1
  674. - )
  675. -
  676. - # DCP: localize the now-expanded per-token global bounds to this
  677. - # rank's owned KV. Done here (after expansion) so each token's global
  678. - # causal length is localized individually; see the comment above.
  679. - if dcp_local_seq_lens is not None:
  680. - seq_lens = self._dcp_localize_decode_seq_lens(
  681. - seq_lens, num_decodes, seq_lens_is_buffer_view
  682. - )
  683. + # Flattening always returns a buffer view, including single-token
  684. + # batches. Keep its address stable across varlen graph replays.
  685. + seq_lens_is_buffer_view = not use_native or next_n > 1
  686. # For DeepseekV4 (compress_ratio > 1), the indexer KV cache stores
  687. # compressed tokens. Convert uncompressed seq_lens to compressed.
  688. + # This must happen BEFORE the DCP localization below: the cache is
  689. + # sharded over compressed states, so this rank's share is
  690. + # localize(len // R), not localize(len) // R (len=10, R=2, W=2:
  691. + # rank 0 owns states {0, 2, 4} -> 3, but 5 // 2 == 2).
  692. if self.compress_ratio > 1:
  693. if seq_lens_is_buffer_view:
  694. seq_lens //= self.compress_ratio
  695. @@ -1255,6 +1276,17 @@
  696. )
  697. self.expanded_seq_lens_buffer[num_decodes:num_decode_tokens] = 0
  698. seq_lens = self.expanded_seq_lens_buffer[:num_decode_tokens]
  699. + # The copy is private: localize it in place below.
  700. + seq_lens_is_buffer_view = True
  701. +
  702. + # DCP: localize the now-expanded (and now-compressed) per-token
  703. + # global bounds to this rank's owned KV. Done here (after expansion)
  704. + # so each token's global causal length is localized individually;
  705. + # see the comment above.
  706. + if dcp_local_seq_lens is not None:
  707. + seq_lens = self._dcp_localize_decode_seq_lens(
  708. + seq_lens, num_decodes, seq_lens_is_buffer_view
  709. + )
  710. # Non-MTP: deep_gemm paged MQA logits requires 2D context_lens
  711. # (csrc/apis/attention.hpp). Unsqueeze to (B, 1) so downstream
  712. @@ -1367,6 +1399,7 @@
  713. cu_seq_len_ks = torch.empty(output_query_len, dtype=torch.int32, device=device)
  714. cu_seq_len_ke = torch.empty(output_query_len, dtype=torch.int32, device=device)
  715. + global_ctx_len = torch.empty(output_query_len, dtype=torch.int32, device=device)
  716. # Under DCP the kernel writes this rank's local row bounds into
  717. # cu_seq_len_ks/ke; otherwise local_cu_seq_lens aliases cu_seq_lens.
  718. @@ -1378,6 +1411,7 @@
  719. token_to_seq,
  720. cu_seq_len_ks,
  721. cu_seq_len_ke,
  722. + global_ctx_len,
  723. qs_start,
  724. qs_stop,
  725. dcp_rank,
  726. @@ -1409,6 +1443,12 @@
  727. local_cu_seq_lens=local_cu_seq_lens,
  728. local_total_seq_lens=local_total_seq_lens,
  729. max_local_total_seq_lens=max_local_total_seq_lens,
  730. + global_ctx_lens=global_ctx_len if dcp_world_size > 1 else None,
  731. + max_global_ctx=(
  732. + int(compressed_seq_lens_cpu[start_idx:end_idx].max().item())
  733. + if dcp_world_size > 1
  734. + else 0
  735. + ),
  736. )
  737. --- a/vllm/v1/attention/backends/mla/sparse_swa.py 2026-09-12 07:48:35.334168175 +0200
  738. +++ b/vllm/v1/attention/backends/mla/sparse_swa.py 2026-09-11 10:13:47.000000000 +0200
  739. @@ -46,18 +46,26 @@
  740. _LAYER_TYPE_SWAONLY = "swaonly"
  741. _LAYER_TYPE_C4A = "c4a"
  742. _LAYER_TYPE_C128A = "c128a"
  743. +# v4.1 ratio-1 / ratio-2 indexer layers: same indexer-path shape as C4A, but
  744. +# their compressed page block is block_size // ratio, so each gets its own plan.
  745. +# v4.1 builders classify layer types themselves (ratio 0 is SWA-only there); see
  746. +# deepseek_v4_1/sparse_mla.py.
  747. +_LAYER_TYPE_C1A = "c1a"
  748. +_LAYER_TYPE_C2A = "c2a"
  749. def _layer_type_for(compress_ratio: int) -> str:
  750. if compress_ratio <= 1:
  751. return _LAYER_TYPE_SWAONLY
  752. + if compress_ratio == 2:
  753. + return _LAYER_TYPE_C2A
  754. if compress_ratio == 4:
  755. return _LAYER_TYPE_C4A
  756. if compress_ratio == 128:
  757. return _LAYER_TYPE_C128A
  758. raise ValueError(
  759. f"Unsupported DeepseekV4 compress_ratio={compress_ratio}; "
  760. - "expected 1, 4, or 128."
  761. + "expected 1, 2, 4, or 128."
  762. )
  763. @@ -70,6 +78,7 @@
  764. prefix: str,
  765. cache_config: CacheConfig,
  766. backend_cls: "type[AttentionBackend] | None" = None,
  767. + block_size: int = 64,
  768. ):
  769. super().__init__()
  770. self.backend_cls = backend_cls or DeepseekSparseSWABackend
  771. @@ -84,12 +93,9 @@
  772. raise ValueError(f"Duplicate layer name: {prefix}")
  773. compilation_config.static_forward_context[prefix] = self
  774. - # Block size is constrained by tensor sharing between SWA and C4A KV blocks.
  775. - # Since both block types share the same physical tensor, they must use the
  776. - # same page size. The C4A KV block shape [256//4, head_dim] = [64, head_dim]
  777. - # determines the SWA block size of 64 tokens per block.
  778. - # TODO(yifan): make SWA block size automatically determined and configurable.
  779. - self.block_size = 64
  780. + # Any multiple of 32; the sparse decode kernels take the page size at
  781. + # runtime.
  782. + self.block_size = block_size
  783. # uint8: fp8_ds_mla UE8M0 paged layout. bfloat16 / float8_e4m3fn:
  784. # contiguous full-cache layout.
  785. assert self.dtype in (torch.uint8, torch.bfloat16, torch.float8_e4m3fn)
  786. @@ -115,6 +121,9 @@
  787. alignment=576 if uses_fp8_ds_mla_layout else 512,
  788. model_version="deepseek_v4",
  789. kv_quant_mode=get_kv_quant_mode(self.cache_config.cache_dtype),
  790. + # DCP: the 128-token window is tiny; every rank keeps all of it
  791. + # (no sharding), so SWA needs no cross-rank combine.
  792. + dcp_replicated=True,
  793. )
  794. def forward(self): ...
  795. @@ -130,7 +139,14 @@
  796. @staticmethod
  797. def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
  798. - return [MultipleOf(64)]
  799. + return [MultipleOf(32)]
  800. +
  801. + @classmethod
  802. + def get_swa_block_size(cls) -> int:
  803. + """Tech2Wild/Kai 2026-09-10: tokens per SWA cache page handed to
  804. + DeepseekV4SWACache. FlashMLA / TRTLLM-gen take the page size at
  805. + runtime (32); kernels instantiated for one page size override this."""
  806. + return 32
  807. @classmethod
  808. def get_preferred_block_size(cls, default_block_size: int) -> int:
  809. @@ -207,12 +223,17 @@
  810. tile_sched_swaonly: "FlashMLASchedMeta | None" = None
  811. tile_sched_c4a: "FlashMLASchedMeta | None" = None
  812. tile_sched_c128a: "FlashMLASchedMeta | None" = None
  813. + tile_sched_c1a: "FlashMLASchedMeta | None" = None
  814. + tile_sched_c2a: "FlashMLASchedMeta | None" = None
  815. flashinfer_sparse_index_cache: dict[str, tuple[torch.Tensor, torch.Tensor]] = field(
  816. default_factory=dict
  817. )
  818. def get_prefill_chunk_plan(
  819. - self, compress_ratio: int, prefill_chunk_size: int
  820. + self,
  821. + compress_ratio: int,
  822. + prefill_chunk_size: int,
  823. + has_compressed: bool | None = None,
  824. ) -> list[tuple[int, int, int, int]]:
  825. if self.num_prefills == 0:
  826. return []
  827. @@ -220,17 +241,20 @@
  828. assert self.prefill_seq_lens_cpu is not None
  829. assert self.prefill_query_lens_cpu is not None
  830. + # Whether the layer gathers a compressed-KV region into the prefill
  831. + # workspace. v4.0 callers keep the legacy default (ratio <= 1 means
  832. + # SWA-only); v4.1 passes has_compressed explicitly because its
  833. + # compress_ratio==1 layers DO have a full-length compressed cache.
  834. + if has_compressed is None:
  835. + has_compressed = compress_ratio > 1
  836. +
  837. # query_len <= max_num_batched_tokens and
  838. # gather_len = query_len + min(prefix_len, window_size - 1), so the
  839. # worst-case gathered width is bounded by
  840. # max_num_batched_tokens + window_size - 1. The compressed prefix pool
  841. # is bounded by ceil(max_model_len / compress_ratio).
  842. max_workspace_area = prefill_chunk_size * (
  843. - (
  844. - 0
  845. - if compress_ratio <= 1
  846. - else cdiv(self.prefill_max_model_len, compress_ratio)
  847. - )
  848. + (cdiv(self.prefill_max_model_len, compress_ratio) if has_compressed else 0)
  849. + self.prefill_window_size
  850. + self.prefill_max_num_batched_tokens
  851. )
  852. @@ -239,13 +263,13 @@
  853. prefix_lens_cpu, min=0, max=self.prefill_window_size - 1
  854. )
  855. compressed_lens_cpu = (
  856. - torch.zeros_like(self.prefill_seq_lens_cpu)
  857. - if compress_ratio <= 1
  858. - else torch.div(
  859. + torch.div(
  860. self.prefill_seq_lens_cpu,
  861. compress_ratio,
  862. rounding_mode="floor",
  863. )
  864. + if has_compressed
  865. + else torch.zeros_like(self.prefill_seq_lens_cpu)
  866. )
  867. chunk_plan: list[tuple[int, int, int, int]] = []
  868. @@ -390,6 +414,9 @@
  869. reorder_batch_threshold: int | None = None
  870. _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
  871. supports_draft_decode_metadata_update = True
  872. + # DSpark's non-causal multi-token decode only touches the replicated SWA
  873. + # caches (no cross-rank combine), so it is supported under DCP.
  874. + supports_non_causal_multi_token_dcp: ClassVar[bool] = True
  875. def __init__(self, *args, **kwargs):
  876. super().__init__(*args, **kwargs)
  877. @@ -709,6 +736,8 @@
  878. tile_sched_swaonly=tile_sched[_LAYER_TYPE_SWAONLY],
  879. tile_sched_c4a=tile_sched[_LAYER_TYPE_C4A],
  880. tile_sched_c128a=tile_sched[_LAYER_TYPE_C128A],
  881. + tile_sched_c1a=tile_sched[_LAYER_TYPE_C1A],
  882. + tile_sched_c2a=tile_sched[_LAYER_TYPE_C2A],
  883. **deepseek_v4_fields, # type: ignore[arg-type]
  884. )
  885. @@ -789,6 +818,8 @@
  886. metadata.tile_sched_swaonly = tile_sched[_LAYER_TYPE_SWAONLY]
  887. metadata.tile_sched_c4a = tile_sched[_LAYER_TYPE_C4A]
  888. metadata.tile_sched_c128a = tile_sched[_LAYER_TYPE_C128A]
  889. + metadata.tile_sched_c1a = tile_sched[_LAYER_TYPE_C1A]
  890. + metadata.tile_sched_c2a = tile_sched[_LAYER_TYPE_C2A]
  891. metadata.flashinfer_sparse_index_cache.clear()
  892. def build_tile_scheduler(
  893. @@ -809,6 +840,8 @@
  894. _LAYER_TYPE_SWAONLY: None,
  895. _LAYER_TYPE_C4A: None,
  896. _LAYER_TYPE_C128A: None,
  897. + _LAYER_TYPE_C1A: None,
  898. + _LAYER_TYPE_C2A: None,
  899. }
  900. if (
  901. num_decode_tokens == 0
  902. --- a/vllm/v1/core/kv_cache_coordinator.py 2026-09-12 07:48:35.835361219 +0200
  903. +++ b/vllm/v1/core/kv_cache_coordinator.py 2026-09-11 10:13:47.000000000 +0200
  904. @@ -24,6 +24,9 @@
  905. KVCacheSpec,
  906. MambaSpec,
  907. SlidingWindowSpec,
  908. + is_dcp_replicated_spec,
  909. + is_dcp_sharded_spec,
  910. + iter_layer_specs,
  911. )
  912. from vllm.v1.request import Request
  913. @@ -619,13 +622,19 @@
  914. assert pcp_world_size == 1, "PCP not support hybrid attn now."
  915. if dcp_world_size > 1:
  916. # DCP shards full-attention KV across ranks and replicates Mamba
  917. - # state; other spec types (e.g. sliding window) have no DCP-aware
  918. - # handling yet, so reject them explicitly.
  919. + # state and attention groups marked ``dcp_replicated`` (their
  920. + # managers run at dcp=1, see dcp_world_size_for_kv_cache_spec);
  921. + # any other spec type has no DCP-aware handling, so reject it.
  922. for g in kv_cache_config.kv_cache_groups:
  923. - assert isinstance(g.kv_cache_spec, (FullAttentionSpec, MambaSpec)), (
  924. - "DCP with hybrid KV cache layouts only supports "
  925. - "full-attention and Mamba groups, got: "
  926. - f"{type(g.kv_cache_spec).__name__}."
  927. + spec = g.kv_cache_spec
  928. + assert (
  929. + is_dcp_sharded_spec(spec)
  930. + or is_dcp_replicated_spec(spec)
  931. + or all(isinstance(s, MambaSpec) for s in iter_layer_specs(spec))
  932. + ), (
  933. + "DCP with hybrid KV cache layouts only supports sharded "
  934. + "full-attention, dcp_replicated attention and Mamba groups, "
  935. + f"got: {type(spec).__name__}."
  936. )
  937. # Fine-grained hash hits require Mamba "align" and compatible cache
  938. # managers in every group. TP needs hashing finer than the Mamba block;
  939. --- a/vllm/v1/core/kv_cache_utils.py 2026-09-12 07:48:35.531214294 +0200
  940. +++ b/vllm/v1/core/kv_cache_utils.py 2026-09-11 10:13:47.000000000 +0200
  941. @@ -22,6 +22,7 @@
  942. from vllm.v1.kv_cache_interface import (
  943. AttentionSpec,
  944. ChunkedLocalAttentionSpec,
  945. + CircularBufferSpec,
  946. FullAttentionSpec,
  947. HiddenStateCacheSpec,
  948. KpoolTailSpec,
  949. @@ -36,6 +37,7 @@
  950. SlidingWindowSpec,
  951. UniformTypeKVCacheSpecs,
  952. compute_layout_strides,
  953. + is_dcp_sharded_spec,
  954. iter_layer_specs,
  955. replace_as,
  956. )
  957. @@ -650,13 +652,15 @@
  958. def resolve_dcp_kv_block_size(spec: KVCacheSpec, dcp_world_size: int) -> int:
  959. - """Return the token span of a cache block under DCP."""
  960. - layer_specs = iter_layer_specs(spec)
  961. - if len(layer_specs) > 0 and all(
  962. - isinstance(layer_spec, AttentionSpec) for layer_spec in layer_specs
  963. - ):
  964. - return spec.block_size * dcp_world_size
  965. - return spec.block_size
  966. + """Return the token span of a cache block under DCP.
  967. +
  968. + Only DCP-sharded groups span ``dcp_world_size`` blocks of tokens per
  969. + manager block; replicated attention groups (``dcp_replicated``), Mamba and
  970. + other per-rank state keep their plain block size. Must agree with
  971. + ``dcp_world_size_for_kv_cache_spec`` (the coordinator scales manager block
  972. + sizes with that helper).
  973. + """
  974. + return spec.block_size * dcp_world_size_for_kv_cache_spec(spec, dcp_world_size)
  975. def resolve_dcp_kv_cache_spec(spec: KVCacheSpec, dcp_world_size: int) -> KVCacheSpec:
  976. @@ -685,15 +689,13 @@
  977. must keep ``dcp_world_size=1`` even when the process runs with DCP > 1.
  978. Draft MLA groups on the sharded DSpark path are ``FullAttentionSpec`` /
  979. - ``MLAAttentionSpec`` and therefore keep the process DCP size. A replicated
  980. - draft group would need a different spec, not this helper.
  981. + ``MLAAttentionSpec`` and therefore keep the process DCP size. Attention
  982. + groups whose layers set ``dcp_replicated`` (DeepSeek-V4.1 SWA cache and
  983. + compressor ring) hold the full KV on every rank and keep ``1``.
  984. """
  985. if dcp_world_size <= 1:
  986. return 1
  987. - inner = spec
  988. - if isinstance(spec, UniformTypeKVCacheSpecs):
  989. - inner = next(iter(spec.kv_cache_specs.values()))
  990. - if isinstance(inner, FullAttentionSpec):
  991. + if is_dcp_sharded_spec(spec):
  992. return dcp_world_size
  993. return 1
  994. @@ -2010,15 +2012,22 @@
  995. cdiv(len(names), n) * page for page, names in page_size_layers.items()
  996. )
  997. - # Bytes a block must hold however the mamba buckets end up split: a mamba
  998. - # bucket can go down to one state per group, every other bucket's split is
  999. - # already fixed by the repeat pattern.
  1000. + # Bytes a block must hold however the state buckets end up split: mamba,
  1001. + # circular-buffer and unsplit sliding-window buckets can go down to one
  1002. + # state per group, every other bucket's split is fixed by the repeat pattern.
  1003. + def is_state_bucket(spec: UniformTypeKVCacheSpecs) -> bool:
  1004. + if isinstance(spec.first_spec, (MambaSpec, CircularBufferSpec)):
  1005. + return True
  1006. + return repeats_per_group is None and isinstance(
  1007. + spec.first_spec, SlidingWindowSpec
  1008. + )
  1009. +
  1010. anchor_bytes = max(
  1011. (
  1012. widest_group_bytes(
  1013. page_size_layers,
  1014. len(spec.kv_cache_specs)
  1015. - if isinstance(spec.first_spec, MambaSpec)
  1016. + if is_state_bucket(spec)
  1017. else num_groups_for(spec, balanced),
  1018. )
  1019. for spec, page_size_layers, balanced in bucketed
  1020. @@ -2029,10 +2038,9 @@
  1021. groups = []
  1022. for spec, page_size_layers, balanced in bucketed:
  1023. num_groups = num_groups_for(spec, balanced)
  1024. - # `_align_hybrid_block_size` pads a mamba state up to one attention
  1025. - # page, so cap a mamba group at the states a block already fits rather
  1026. - # than let it widen the block.
  1027. - if anchor_bytes and isinstance(spec.first_spec, MambaSpec):
  1028. + # Cap a state group at the states a block already fits rather than let
  1029. + # it widen the block.
  1030. + if anchor_bytes and is_state_bucket(spec):
  1031. states_per_block = max(anchor_bytes // spec.first_spec.page_size_bytes, 1)
  1032. num_groups = max(
  1033. num_groups, cdiv(len(spec.kv_cache_specs), states_per_block)
  1034. @@ -2068,8 +2076,9 @@
  1035. if spec_config is None or not spec_config.use_eagle():
  1036. return False
  1037. model_config = vllm_config.model_config
  1038. - return (
  1039. - model_config is not None and model_config.hf_config.model_type == "deepseek_v4"
  1040. + return model_config is not None and model_config.hf_config.model_type in (
  1041. + "deepseek_v4",
  1042. + "deepseek_v41",
  1043. )
  1044. @@ -2090,9 +2099,9 @@
  1045. spec merging, wherever grouping happens to land. It is sufficient but
  1046. not necessary: a drafter whose spec is indistinguishable from the
  1047. target's cannot be found this way.
  1048. - 2. Model-scoped positional fallback for DeepseekV4, whose MTP block reuses
  1049. - the target's own decoder layer and so carries no spec marker. Its draft
  1050. - attention layer is always the last registered layer, so flag whichever
  1051. + 2. Model-scoped positional fallback for DeepseekV4/V4.1, whose MTP block
  1052. + reuses the target's own decoder layer and so carries no spec marker. Its
  1053. + draft attention layer is always the last registered layer, so flag whichever
  1054. group holds it. This rule is only valid where the groups partition
  1055. exactly the layers of ``kv_cache_spec``, which is true on the packed
  1056. grouping path and not in general; other callers must leave
  1057. @@ -2105,7 +2114,8 @@
  1058. kv_cache_spec: The kv cache spec of each attention layer, in layer
  1059. registration order. Only read by rule 2.
  1060. kv_cache_groups: Groups to annotate in place.
  1061. - use_deepseek_v4_fallback: Enable rule 2 for a DeepseekV4 packed group.
  1062. + use_deepseek_v4_fallback: Enable rule 2 for a DeepseekV4/V4.1 packed
  1063. + group.
  1064. """
  1065. spec_config = vllm_config.speculative_config
  1066. if spec_config is None or not spec_config.use_eagle_block_drop():
  1067. --- a/vllm/v1/kv_cache_interface.py 2026-09-12 07:48:35.351087500 +0200
  1068. +++ b/vllm/v1/kv_cache_interface.py 2026-09-11 10:13:47.000000000 +0200
  1069. @@ -240,8 +240,11 @@
  1070. f"Unsupported KV cache spec type: {type(self)}. "
  1071. "Please register it using @register_kv_cache_spec decorator."
  1072. )
  1073. + dcp_replicated = getattr(self, "dcp_replicated", None)
  1074. return all(
  1075. - isinstance(spec, uniform_type_base_spec) for spec in kv_cache_specs.values()
  1076. + isinstance(spec, uniform_type_base_spec)
  1077. + and getattr(spec, "dcp_replicated", None) == dcp_replicated
  1078. + for spec in kv_cache_specs.values()
  1079. )
  1080. @@ -400,6 +403,13 @@
  1081. """Tokens covered by one stored state. Ints > 1 compress multiple tokens
  1082. into one state (DSv4 sparse MLA); fractions < 1 store multiple states per
  1083. token (Whisper block pooling: ``Fraction(1, block_pool_size)``)."""
  1084. + dcp_replicated: bool = False
  1085. + """DCP: every rank holds this layer's full KV instead of a 1/dcp shard.
  1086. + A replicated group keeps its plain block geometry under DCP (no block-size
  1087. + scaling, no per-rank memory divide, block table sized for the whole
  1088. + context) and every rank writes every position (DeepSeek-V4.1 SWA cache and
  1089. + compressor ring). Sharded groups (the default) round-robin positions across
  1090. + ranks at ``cp_kv_cache_interleave_size`` granularity."""
  1091. def __post_init__(self):
  1092. if self.head_size_v is None:
  1093. @@ -438,7 +448,9 @@
  1094. def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:
  1095. parallel_config = vllm_config.parallel_config
  1096. - kv_shard_count = parallel_config.decode_context_parallel_size
  1097. + kv_shard_count = (
  1098. + 1 if self.dcp_replicated else parallel_config.decode_context_parallel_size
  1099. + )
  1100. return cdiv(max_len, self.block_size * kv_shard_count)
  1101. @@ -471,7 +483,7 @@
  1102. def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
  1103. max_model_len = vllm_config.model_config.max_model_len
  1104. dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
  1105. - if dcp_world_size > 1:
  1106. + if dcp_world_size > 1 and not self.dcp_replicated:
  1107. max_model_len = cdiv(max_model_len, dcp_world_size)
  1108. return cdiv(max_model_len, self.block_size) * self.page_size_bytes
  1109. @@ -519,6 +531,7 @@
  1110. num_head_slots=specs[0].num_head_slots,
  1111. state_content_bytes=specs[0].state_content_bytes,
  1112. tokens_per_state=specs[0].tokens_per_state,
  1113. + dcp_replicated=specs[0].dcp_replicated,
  1114. sliding_window=cls.merge_window_sizes(sliding_window),
  1115. attention_chunk_size=cls.merge_window_sizes(attention_chunk_size),
  1116. # If any layer in the group is non-causal, treat the group as
  1117. @@ -602,6 +615,7 @@
  1118. state_content_bytes=specs[0].state_content_bytes,
  1119. cache_dtype_str=cache_dtype_str_set.pop(),
  1120. tokens_per_state=tokens_per_state_set.pop(),
  1121. + dcp_replicated=specs[0].dcp_replicated,
  1122. model_version=model_version_set.pop(),
  1123. storage_block_size=storage_block_size_set.pop(),
  1124. non_causal_multi_token_decode=non_causal_mtd_set.pop(),
  1125. @@ -657,6 +671,7 @@
  1126. num_head_slots=base.num_head_slots,
  1127. state_content_bytes=base.state_content_bytes,
  1128. tokens_per_state=base.tokens_per_state,
  1129. + dcp_replicated=base.dcp_replicated,
  1130. sliding_window=base.sliding_window,
  1131. attention_chunk_size=base.attention_chunk_size,
  1132. non_causal=base.non_causal,
  1133. @@ -700,6 +715,7 @@
  1134. return all(
  1135. isinstance(spec, ChunkedLocalAttentionSpec)
  1136. and spec.attention_chunk_size == self.attention_chunk_size
  1137. + and spec.dcp_replicated == self.dcp_replicated
  1138. for spec in kv_cache_specs.values()
  1139. )
  1140. @@ -743,9 +759,10 @@
  1141. return cdiv(num_tokens, self.block_size) + 1
  1142. def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
  1143. - assert vllm_config.parallel_config.decode_context_parallel_size == 1, (
  1144. - "DCP not support sliding window."
  1145. - )
  1146. + assert (
  1147. + vllm_config.parallel_config.decode_context_parallel_size == 1
  1148. + or self.dcp_replicated
  1149. + ), "DCP not support sliding window (sharded); mark the spec dcp_replicated."
  1150. max_blocks = self.max_admission_blocks_per_request(
  1151. max_in_flight_tokens=vllm_config.max_in_flight_tokens,
  1152. max_model_len=vllm_config.model_config.max_model_len,
  1153. @@ -758,6 +775,7 @@
  1154. return all(
  1155. isinstance(spec, SlidingWindowSpec)
  1156. and spec.sliding_window == self.sliding_window
  1157. + and spec.dcp_replicated == self.dcp_replicated
  1158. for spec in kv_cache_specs.values()
  1159. )
  1160. @@ -786,7 +804,9 @@
  1161. self, kv_cache_specs: dict[str, KVCacheSpec]
  1162. ) -> bool:
  1163. return all(
  1164. - isinstance(spec, CircularBufferSpec) for spec in kv_cache_specs.values()
  1165. + isinstance(spec, CircularBufferSpec)
  1166. + and spec.dcp_replicated == self.dcp_replicated
  1167. + for spec in kv_cache_specs.values()
  1168. )
  1169. @property
  1170. @@ -847,6 +867,7 @@
  1171. extra_retained_tokens=extra_retained_set.pop(),
  1172. cache_dtype_str=cache_dtype_str_set.pop(),
  1173. tokens_per_state=tokens_per_state_set.pop(),
  1174. + dcp_replicated=specs[0].dcp_replicated,
  1175. model_version=model_version_set.pop(),
  1176. )
  1177. @@ -856,6 +877,7 @@
  1178. return all(
  1179. isinstance(spec, SlidingWindowMLASpec)
  1180. and spec.sliding_window == self.sliding_window
  1181. + and spec.dcp_replicated == self.dcp_replicated
  1182. for spec in kv_cache_specs.values()
  1183. )
  1184. @@ -1013,6 +1035,7 @@
  1185. page_size_padded=specs[0].page_size_padded,
  1186. num_head_slots=specs[0].num_head_slots,
  1187. state_content_bytes=specs[0].state_content_bytes,
  1188. + dcp_replicated=specs[0].dcp_replicated,
  1189. sliding_window=cls.merge_window_sizes(sliding_window),
  1190. attention_chunk_size=cls.merge_window_sizes(attention_chunk_size),
  1191. non_causal=any(spec.non_causal for spec in specs),
  1192. @@ -1147,6 +1170,28 @@
  1193. )
  1194. +def is_dcp_replicated_spec(kv_cache_spec: KVCacheSpec) -> bool:
  1195. + """Whether a KV cache group spec is (or wraps) DCP-replicated attention
  1196. + layers: every DCP rank holds the full KV, so the group keeps its plain
  1197. + block geometry and every rank writes every position."""
  1198. + layer_specs = iter_layer_specs(kv_cache_spec)
  1199. + return len(layer_specs) > 0 and all(
  1200. + isinstance(spec, AttentionSpec) and spec.dcp_replicated
  1201. + for spec in layer_specs
  1202. + )
  1203. +
  1204. +
  1205. +def is_dcp_sharded_spec(kv_cache_spec: KVCacheSpec) -> bool:
  1206. + """Whether a KV cache group spec is (or wraps) DCP-sharded full attention
  1207. + (incl. MLA): positions are round-robined across DCP ranks, so manager block
  1208. + size, prefix hashing and per-rank memory scale by the DCP size."""
  1209. + layer_specs = iter_layer_specs(kv_cache_spec)
  1210. + return len(layer_specs) > 0 and all(
  1211. + isinstance(spec, FullAttentionSpec) and not spec.dcp_replicated
  1212. + for spec in layer_specs
  1213. + )
  1214. +
  1215. +
  1216. def get_kv_cache_spec_kind(kv_cache_spec: KVCacheSpec) -> KVCacheSpecKind:
  1217. if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs):
  1218. inner_kinds = {
  1219. --- a/vllm/v1/worker/cp_utils.py 2026-09-12 07:48:35.759086508 +0200
  1220. +++ b/vllm/v1/worker/cp_utils.py 2026-09-11 10:13:47.000000000 +0200
  1221. @@ -35,6 +35,10 @@
  1222. f"but {backend.get_name()} does not support PCP."
  1223. )
  1224. layer_impl = getattr(layer, "impl", None)
  1225. + if layer_impl is None and hasattr(layer, "need_to_return_lse_for_decode"):
  1226. + # DeepSeek-V4.1 attention runs its own forward (no
  1227. + # AttentionImpl) and exposes the DCP contract on the layer.
  1228. + layer_impl = layer
  1229. if layer_impl is None:
  1230. continue
  1231. if vllm_config.speculative_config is not None and interleave_size > 1:
  1232. --- a/vllm/v1/worker/gpu/attn_utils.py 2026-09-12 07:48:35.523821429 +0200
  1233. +++ b/vllm/v1/worker/gpu/attn_utils.py 2026-09-11 10:13:47.000000000 +0200
  1234. @@ -20,6 +20,7 @@
  1235. KVCacheConfig,
  1236. KVCacheSpec,
  1237. UniformTypeKVCacheSpecs,
  1238. + is_dcp_replicated_spec,
  1239. )
  1240. from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata
  1241. from vllm.v1.worker.utils import (
  1242. @@ -283,6 +284,13 @@
  1243. for i in range(num_kv_cache_groups):
  1244. block_table = block_tables[i]
  1245. slot_mapping = slot_mappings[i]
  1246. + # DCP-replicated groups hold every position on every rank: their
  1247. + # builders must see the full seq_lens, not this rank's shard.
  1248. + group_dcp_local_seq_lens = (
  1249. + None
  1250. + if is_dcp_replicated_spec(kv_cache_config.kv_cache_groups[i].kv_cache_spec)
  1251. + else dcp_local_seq_lens
  1252. + )
  1253. # Per-group causal for hybrid drafters (mixed SWA/full attention).
  1254. group_causal = (
  1255. causal if isinstance(causal, (bool, torch.Tensor)) else causal.get(i, True)
  1256. @@ -310,7 +318,7 @@
  1257. block_table_tensor=block_table,
  1258. slot_mapping=slot_mapping,
  1259. causal=group_causal,
  1260. - dcp_local_seq_lens=dcp_local_seq_lens,
  1261. + dcp_local_seq_lens=group_dcp_local_seq_lens,
  1262. positions=positions,
  1263. is_prefilling=group_is_prefilling,
  1264. mm_req_doc_ranges=mm_req_doc_ranges,
  1265. --- a/vllm/v1/worker/gpu/block_table.py 2026-09-12 07:48:35.514464887 +0200
  1266. +++ b/vllm/v1/worker/gpu/block_table.py 2026-09-11 10:13:47.000000000 +0200
  1267. @@ -27,6 +27,7 @@
  1268. cp_rank: int = 0,
  1269. cp_interleave: int = 1,
  1270. slot_mapping_enabled: list[bool] | None = None,
  1271. + cp_sharded: list[bool] | None = None,
  1272. ):
  1273. self.block_sizes = block_sizes
  1274. self.kernel_block_sizes = kernel_block_sizes
  1275. @@ -44,6 +45,13 @@
  1276. slot_mapping_enabled = [True] * self.num_kv_cache_groups
  1277. assert len(slot_mapping_enabled) == self.num_kv_cache_groups
  1278. self._slot_mapping_enabled = slot_mapping_enabled
  1279. + # Per group: True = positions are round-robined across the cp_size
  1280. + # ranks (this rank stores 1/cp_size of them); False = DCP-replicated
  1281. + # group, every rank stores every position at its plain block geometry.
  1282. + if cp_sharded is None:
  1283. + cp_sharded = [True] * self.num_kv_cache_groups
  1284. + assert len(cp_sharded) == self.num_kv_cache_groups
  1285. + self._cp_sharded = cp_sharded
  1286. self.blocks_per_kv_block = [
  1287. bs // kbs for bs, kbs in zip(block_sizes, kernel_block_sizes)
  1288. @@ -108,6 +116,9 @@
  1289. self.slot_mapping_enabled = torch.tensor(
  1290. self._slot_mapping_enabled, dtype=torch.bool, device=self.device
  1291. )
  1292. + self.cp_sharded = torch.tensor(
  1293. + self._cp_sharded, dtype=torch.bool, device=self.device
  1294. + )
  1295. self.input_block_table_ptrs = self._make_ptr_tensor(self.input_block_tables)
  1296. def append_block_ids(
  1297. @@ -211,6 +222,7 @@
  1298. self.block_sizes_tensor,
  1299. self.kernel_block_sizes_tensor,
  1300. self.slot_mapping_enabled,
  1301. + self.cp_sharded,
  1302. slot_mappings,
  1303. slot_mappings.stride(0),
  1304. self.cp_rank,
  1305. @@ -284,6 +296,7 @@
  1306. block_sizes, # [num_kv_cache_groups]
  1307. kernel_block_sizes, # [num_kv_cache_groups]
  1308. slot_mapping_enabled, # [num_kv_cache_groups]
  1309. + cp_sharded, # [num_kv_cache_groups]
  1310. slot_mappings_ptr, # [num_kv_cache_groups, max_num_tokens]
  1311. slot_mappings_stride,
  1312. cp_rank,
  1313. @@ -313,6 +326,7 @@
  1314. kv_block_size = tl.load(block_sizes + group_id)
  1315. kernel_block_size = tl.load(kernel_block_sizes + group_id)
  1316. mapping_enabled = tl.load(slot_mapping_enabled + group_id)
  1317. + group_sharded = tl.load(cp_sharded + group_id)
  1318. req_state_idx = tl.load(idx_mapping + batch_idx)
  1319. start_idx = tl.load(query_start_loc + batch_idx)
  1320. @@ -326,15 +340,19 @@
  1321. local_positions = positions
  1322. is_local = True
  1323. else:
  1324. - # Context parallelism is used.
  1325. + # Context parallelism is used. Sharded groups round-robin positions
  1326. + # across the CP ranks; DCP-replicated groups (cp_sharded False)
  1327. + # keep every position on every rank at the plain block geometry.
  1328. virtual_block_size = kv_block_size * CP_SIZE
  1329. virtual_block_indices = positions // virtual_block_size
  1330. virtual_block_offsets = positions % virtual_block_size
  1331. - is_local = virtual_block_offsets // CP_INTERLEAVE % CP_SIZE == cp_rank
  1332. + owner_is_me = virtual_block_offsets // CP_INTERLEAVE % CP_SIZE == cp_rank
  1333. rounds = virtual_block_offsets // (CP_INTERLEAVE * CP_SIZE)
  1334. remainder = virtual_block_offsets % CP_INTERLEAVE
  1335. local_offsets = rounds * CP_INTERLEAVE + remainder
  1336. - local_positions = virtual_block_indices * kv_block_size + local_offsets
  1337. + sharded_positions = virtual_block_indices * kv_block_size + local_offsets
  1338. + local_positions = tl.where(group_sharded, sharded_positions, positions)
  1339. + is_local = tl.where(group_sharded, owner_is_me, True)
  1340. block_indices = tl.where(
  1341. mapping_enabled, local_positions // kernel_block_size, 0
  1342. --- a/vllm/v1/worker/gpu/model_runner.py 2026-09-12 07:48:35.558729666 +0200
  1343. +++ b/vllm/v1/worker/gpu/model_runner.py 2026-09-11 10:13:47.000000000 +0200
  1344. @@ -70,6 +70,7 @@
  1345. KVCacheConfig,
  1346. MambaSpec,
  1347. UniformTypeKVCacheSpecs,
  1348. + is_dcp_replicated_spec,
  1349. )
  1350. from vllm.v1.outputs import (
  1351. DraftTokenIds,
  1352. @@ -170,6 +171,82 @@
  1353. logger = init_logger(__name__)
  1354. +# Stage 7 (DSV41_PREFILL_TIMING=1): per-phase budget of eager prefill steps.
  1355. +import os as _dsv41_os # noqa: E402
  1356. +
  1357. +_DSV41_PREFILL_TIMING = _dsv41_os.environ.get("DSV41_PREFILL_TIMING", "0") == "1"
  1358. +
  1359. +
  1360. +def _dsv41_prefill_timer_start(input_batch, batch_desc, dummy_run):
  1361. + if (
  1362. + not _DSV41_PREFILL_TIMING
  1363. + or dummy_run
  1364. + or input_batch.num_tokens < 1024
  1365. + or batch_desc.cg_mode == CUDAGraphMode.FULL
  1366. + or torch.cuda.is_current_stream_capturing()
  1367. + ):
  1368. + return None
  1369. + try: # drop phase time accumulated outside this forward (e.g. the DSpark drafter of the previous step)
  1370. + from vllm.model_executor.kernels.attention.dsa.candidate_blocks import prefill_acc_take
  1371. +
  1372. + prefill_acc_take()
  1373. + except Exception: # noqa: BLE001
  1374. + pass
  1375. + e0 = torch.cuda.Event(enable_timing=True)
  1376. + e0.record()
  1377. + return e0
  1378. +
  1379. +
  1380. +def _dsv41_prefill_timer_end(e0, input_batch, model_state):
  1381. + if e0 is None:
  1382. + return
  1383. + e1 = torch.cuda.Event(enable_timing=True)
  1384. + e1.record()
  1385. + e1.synchronize()
  1386. + fwd = e0.elapsed_time(e1)
  1387. + try:
  1388. + from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (
  1389. + prefill_acc_take,
  1390. + )
  1391. +
  1392. + acc = prefill_acc_take()
  1393. + except Exception: # noqa: BLE001
  1394. + acc = {}
  1395. + g = acc.get("attn_gather", 0.0)
  1396. + k = acc.get("attn_kernel", 0.0)
  1397. + c = acc.get("attn_combine", 0.0)
  1398. + idx = acc.get("indexer", 0.0)
  1399. + cand = acc.get("candidates", 0.0)
  1400. + stager = getattr(model_state, "engram_stager", None)
  1401. + eng = float(getattr(stager, "_t_last_total_ms", 0.0) or 0.0) if stager is not None else 0.0
  1402. + n = int(input_batch.num_tokens)
  1403. + moe = acc.get("moe", 0.0)
  1404. + ar = acc.get("allreduce", 0.0)
  1405. + egf = acc.get("engram_fwd", 0.0)
  1406. + rest = fwd - (g + k + c) - idx
  1407. + logger.info(
  1408. + "Prefill step timing: %d tokens / %d reqs: forward %.0f ms (%.0f tok/s) = "
  1409. + "attention %.0f (gather %.0f, kernel %.0f, combine %.0f) + indexer %.0f "
  1410. + "(incl. candidates %.0f) + rest[MoE/dense/other] %.0f [moe %.0f, allreduce %.0f, "
  1411. + "engram_fwd %.0f, other %.0f]; Engram staging before the step %.0f ms",
  1412. + n,
  1413. + int(input_batch.num_reqs),
  1414. + fwd,
  1415. + 1e3 * n / max(fwd, 1e-6),
  1416. + g + k + c,
  1417. + g,
  1418. + k,
  1419. + c,
  1420. + idx,
  1421. + cand,
  1422. + rest,
  1423. + moe,
  1424. + ar,
  1425. + egf,
  1426. + rest - moe - ar - egf,
  1427. + eng,
  1428. + )
  1429. +
  1430. class GPUModelRunner(LoRAModelRunnerMixin):
  1431. def __init__(self, vllm_config: VllmConfig, device: torch.device):
  1432. @@ -558,9 +635,13 @@
  1433. block_sizes = []
  1434. max_num_blocks_per_group = []
  1435. slot_mapping_enabled = []
  1436. + cp_sharded = []
  1437. for kv_cache_group in kv_cache_config.kv_cache_groups:
  1438. spec = kv_cache_group.kv_cache_spec
  1439. block_sizes.append(spec.block_size)
  1440. + # DCP-replicated groups (V4.1 SWA cache, compressor ring) store
  1441. + # every position on every rank; only sharded groups round-robin.
  1442. + cp_sharded.append(not is_dcp_replicated_spec(spec))
  1443. layer_spec = (
  1444. spec.first_spec if isinstance(spec, UniformTypeKVCacheSpecs) else spec
  1445. )
  1446. @@ -623,7 +704,15 @@
  1447. cp_size=self.dcp_size,
  1448. cp_rank=self.dcp_rank,
  1449. cp_interleave=self.cp_interleave,
  1450. + cp_sharded=cp_sharded,
  1451. )
  1452. + if self.use_dcp:
  1453. + logger.info(
  1454. + "DCP%d KV cache groups: sharded=%s (block sizes %s)",
  1455. + self.dcp_size,
  1456. + cp_sharded,
  1457. + block_sizes,
  1458. + )
  1459. self.pcp_manager = pcp.maybe_build_pcp_manager(
  1460. self.vllm_config,
  1461. self.device,
  1462. @@ -1752,6 +1841,7 @@
  1463. input_batch, batch_desc.cg_mode == CUDAGraphMode.FULL
  1464. )
  1465. self.step_timing.forward_start()
  1466. + _dsv41_pf = _dsv41_prefill_timer_start(input_batch, batch_desc, dummy_run)
  1467. # Run model.
  1468. if batch_desc.cg_mode == CUDAGraphMode.FULL:
  1469. @@ -1794,6 +1884,7 @@
  1470. else:
  1471. # Eager (NONE): call the raw model directly.
  1472. model_output = self.model(**model_inputs)
  1473. + _dsv41_prefill_timer_end(_dsv41_pf, input_batch, self.model_state)
  1474. if self.is_last_pp_rank:
  1475. if self.use_aux_hidden_state_outputs:
  1476. --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py 2026-09-12 07:48:35.759099133 +0200
  1477. +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py 2026-09-11 10:13:47.000000000 +0200
  1478. @@ -377,6 +377,19 @@
  1479. assert self.draft_kv_cache_group_id >= 0
  1480. # Support multiple draft KV cache groups by preparing inputs once for each
  1481. for i, gid in enumerate(self.draft_kv_cache_group_ids):
  1482. + # The draft's context/query slots must follow the GROUP's CP view:
  1483. + # a DCP-replicated group (DeepSeek-V4.1 draft SWA caches) stores
  1484. + # every position on every rank at the plain block geometry, so it
  1485. + # gets the DCP1 mapping; only sharded groups round-robin.
  1486. + group_sharded = getattr(self.block_tables, "_cp_sharded", None)
  1487. + if group_sharded is None or group_sharded[gid]:
  1488. + cp_rank, cp_size, cp_interleave = (
  1489. + self.block_tables.cp_rank,
  1490. + self.block_tables.cp_size,
  1491. + self.block_tables.cp_interleave,
  1492. + )
  1493. + else:
  1494. + cp_rank, cp_size, cp_interleave = 0, 1, 1
  1495. prepare_dflash_inputs(
  1496. self.input_buffers,
  1497. self.block_tables.slot_mappings[gid],
  1498. @@ -396,9 +409,9 @@
  1499. seeds,
  1500. self.block_tables.input_block_tables[gid],
  1501. self.block_tables.kernel_block_sizes[gid],
  1502. - self.block_tables.cp_rank,
  1503. - self.block_tables.cp_size,
  1504. - self.block_tables.cp_interleave,
  1505. + cp_rank,
  1506. + cp_size,
  1507. + cp_interleave,
  1508. self.parallel_drafting_token_id,
  1509. self.num_query_per_req,
  1510. self.num_speculative_steps,