@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 0f67d4d962

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 14, 05:30 UTC

[Attention] Add MLA prefill backend: trtllm_ragged_attention_deepseek (#26397)

Signed-off-by: Ming Yang <minos.future@gmail.com>

Ming Yang committed Oct 24, 2025 at 10:24 UTC 0f67d4d962872767ac1fca8e98d1bb679aae762a

2 files changed +113 -2

vllm/envs.py

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

183 VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB: int | None = None

184 VLLM_NIXL_ABORT_REQUEST_TIMEOUT: int = 480

185 VLLM_USE_CUDNN_PREFILL: bool = False

186 + VLLM_USE_TRTLLM_RAGGED_DEEPSEEK_PREFILL: bool = False

187 VLLM_ENABLE_CUDAGRAPH_GC: bool = False

188 VLLM_LOOPBACK_IP: str = ""

189 VLLM_ALLOW_CHUNKED_LOCAL_ATTN_WITH_HYBRID_KV_CACHE: bool = False

@@ -1250,6 +1251,10 @@ environment_variables: dict[str, Callable[[], Any]] = {

1251 "VLLM_USE_CUDNN_PREFILL": lambda: bool(

1252 int(os.getenv("VLLM_USE_CUDNN_PREFILL", "0"))

1253 ),

1254 + # Controls whether to use TRT-LLM ragged DeepSeek prefill

1255 + "VLLM_USE_TRTLLM_RAGGED_DEEPSEEK_PREFILL": lambda: bool(

1256 + int(os.getenv("VLLM_USE_TRTLLM_RAGGED_DEEPSEEK_PREFILL", "0"))

1257 + ),

1258 # If set to 1/True, use the TRTLLM attention backend in flashinfer.

1259 # If set to 0/False, use the default attention backend in flashinfer.

1260 # If not set, auto-detect the attention backend in flashinfer.

@@ -1481,6 +1486,7 @@ def compute_hash() -> str:

1486 "VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8_CUTLASS",

1487 "VLLM_USE_FLASHINFER_MOE_MXFP4_BF16",

1488 "VLLM_USE_CUDNN_PREFILL",

1489 + "VLLM_USE_TRTLLM_RAGGED_DEEPSEEK_PREFILL",

1490 "VLLM_USE_TRTLLM_ATTENTION",

1491 "VLLM_FLASHINFER_DISABLE_Q_QUANTIZATION",

1492 "VLLM_ROCM_USE_AITER",

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

+107 -2

@@ -371,6 +371,7 @@ class MLACommonPrefillMetadata:

371 query_start_loc: torch.Tensor

372 max_query_len: int

373 chunked_context: ChunkedContextMetadata | None = None

374 + query_seq_lens: torch.Tensor | None = None

375

376

377 @dataclass

@@ -386,7 +387,6 @@ class CudnnPrefillMetadata(MLACommonPrefillMetadata):

387 class ChunkedContextMetadata(MLACommonPrefillMetadata.ChunkedContextMetadata):

388 seq_lens: torch.Tensor

389

389 - query_seq_lens: torch.Tensor | None = None

390 cudnn_workspace: torch.Tensor | None = None

391

392

@@ -457,6 +457,7 @@ def use_flashinfer_prefill() -> bool:

457 not envs.VLLM_DISABLE_FLASHINFER_PREFILL

458 and flashinfer_available

459 and not envs.VLLM_USE_CUDNN_PREFILL

460 + and not envs.VLLM_USE_TRTLLM_RAGGED_DEEPSEEK_PREFILL

461 and current_platform.is_device_capability(100)

462 )

463

@@ -470,6 +471,15 @@ def use_cudnn_prefill() -> bool:

471 )

472

473

474 +def use_trtllm_ragged_deepseek_prefill() -> bool:

475 + """Check if TRT-LLM ragged DeepSeek prefill should be used."""

476 + return (

477 + flashinfer_available

478 + and envs.VLLM_USE_TRTLLM_RAGGED_DEEPSEEK_PREFILL

479 + and current_platform.is_device_capability(100)

480 + )

481 +

482 +

483 # Currently 394MB, this can be tuned based on GEMM sizes used.

484 # Chosen to be the same as sglang:

485 # https://github.com/sgl-project/sglang/blob/766392c6bda2558b61ce6d1c1bfd8081a549e1f1/python/sglang/global_config.py#L37

@@ -593,6 +603,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

603

604 self._use_cudnn_prefill = use_cudnn_prefill()

605 self._use_fi_prefill = use_flashinfer_prefill()

606 + self._use_trtllm_ragged_prefill = use_trtllm_ragged_deepseek_prefill()

607 self.prefill_metadata_cls = (

608 FlashInferPrefillMetadata

609 if self._use_fi_prefill

@@ -613,6 +624,11 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

624 get_per_layer_parameters(vllm_config, layer_names, MLACommonImpl)

625 )

626

627 + if self._use_trtllm_ragged_prefill:

628 + self._workspace_buffer = torch.empty(

629 + FLASHINFER_WORKSPACE_BUFFER_SIZE, dtype=torch.uint8, device=device

630 + )

631 +

632 if self._use_cudnn_prefill:

633 self.cudnn_workspace = torch.empty(

634 CUDNN_WORKSPACE_SIZE * scheduler_config.max_num_seqs,

@@ -934,6 +950,11 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

950 )

951 prefill_metadata.cudnn_workspace = self.cudnn_workspace

952

953 + if self._use_trtllm_ragged_prefill:

954 + prefill_metadata.query_seq_lens = (

955 + prefill_query_start_loc[1:] - prefill_query_start_loc[:-1]

956 + )

957 +

958 decode_metadata = None

959 if num_decodes > 0:

960 decode_metadata = self._build_decode(

@@ -1230,6 +1251,13 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1251 self._run_prefill_context_chunk = self._run_prefill_context_chunk_fi

1252 self._run_prefill_new_tokens = self._run_prefill_new_tokens_fi

1253 self._pad_v = False

1254 + elif use_trtllm_ragged_deepseek_prefill():

1255 + logger.debug_once("Using TRT-LLM ragged DeepSeek prefill for MLA")

1256 + self._run_prefill_context_chunk = (

1257 + self._run_prefill_context_chunk_trtllm_ragged

1258 + )

1259 + self._run_prefill_new_tokens = self._run_prefill_new_tokens_trtllm_ragged

1260 + self._pad_v = False

1261 elif use_cudnn_prefill():

1262 logger.debug_once("Using CUDNN prefill for MLA")

1263 self._run_prefill_context_chunk = self._run_prefill_context_chunk_cudnn

@@ -1326,6 +1354,7 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1354 ):

1355 assert isinstance(prefill, FlashInferPrefillMetadata)

1356 assert prefill.prefill_main is not None

1357 +

1358 ret = prefill.prefill_main.run(

1359 q=q,

1360 k=k,

@@ -1334,7 +1363,6 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1363 )

1364

1365 if isinstance(ret, tuple):

1337 - # Convert from (q_len, num_heads) to (num_heads, q_len)

1366 return ret[0], ret[1].transpose(0, 1).contiguous()

1367 return ret

1368

@@ -1384,12 +1412,14 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1412 self, prefill: MLACommonPrefillMetadata, chunk_idx: int, q, k, v

1413 ):

1414 assert isinstance(prefill, FlashInferPrefillMetadata)

1415 +

1416 attn_out, lse = prefill.prefill_chunks[chunk_idx].run(

1417 q=q,

1418 k=k,

1419 v=v,

1420 return_lse=True,

1421 )

1422 +

1423 # Convert from (q_len, num_heads) to (num_heads, q_len)

1424 return attn_out, lse.transpose(0, 1).contiguous()

1425

@@ -1418,6 +1448,81 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1448 is_cuda_graph_compatible=True,

1449 )

1450

1451 + def _run_prefill_new_tokens_trtllm_ragged(

1452 + self, prefill: MLACommonPrefillMetadata, q, k, v, return_softmax_lse

1453 + ):

1454 + """TRT-LLM ragged attention for new tokens (causal)."""

1455 + from flashinfer.prefill import trtllm_ragged_attention_deepseek

1456 +

1457 + assert prefill.query_seq_lens is not None

1458 +

1459 + ret = trtllm_ragged_attention_deepseek(

1460 + query=q,

1461 + key=k,

1462 + value=v,

1463 + workspace_buffer=self._workspace_buffer,

1464 + seq_lens=prefill.query_seq_lens,

1465 + max_q_len=prefill.max_query_len,

1466 + max_kv_len=prefill.max_query_len,

1467 + bmm1_scale=self.scale,

1468 + bmm2_scale=1.0,

1469 + o_sf_scale=1.0,

1470 + batch_size=prefill.query_seq_lens.shape[0],

1471 + window_left=-1,

1472 + cum_seq_lens_q=prefill.query_start_loc,

1473 + cum_seq_lens_kv=prefill.query_start_loc,

1474 + enable_pdl=False,

1475 + is_causal=True,

1476 + return_lse=return_softmax_lse,

1477 + )

1478 +

1479 + if isinstance(ret, tuple):

1480 + # Convert from (q_len, num_heads) to (num_heads, q_len)

1481 + return ret[0], ret[1].transpose(0, 1).contiguous()

1482 + return ret

1483 +

1484 + def _run_prefill_context_chunk_trtllm_ragged(

1485 + self, prefill: MLACommonPrefillMetadata, chunk_idx: int, q, k, v

1486 + ):

1487 + """TRT-LLM ragged attention for context chunks (non-causal)."""

1488 + from flashinfer.prefill import trtllm_ragged_attention_deepseek

1489 +

1490 + assert prefill.chunked_context is not None

1491 + assert prefill.chunked_context.seq_lens[chunk_idx] is not None

1492 +

1493 + out = torch.zeros(

1494 + q.shape[0],

1495 + q.shape[1],

1496 + v.shape[2],

1497 + device=q.device,

1498 + dtype=q.dtype,

1499 + )

1500 + self._workspace_buffer.fill_(0)

1501 +

1502 + attn_out, lse = trtllm_ragged_attention_deepseek(

1503 + query=q,

1504 + key=k,

1505 + value=v,

1506 + workspace_buffer=self._workspace_buffer,

1507 + seq_lens=prefill.chunked_context.seq_lens[chunk_idx],

1508 + max_q_len=prefill.max_query_len,

1509 + max_kv_len=prefill.chunked_context.max_seq_lens[chunk_idx],

1510 + bmm1_scale=self.scale,

1511 + bmm2_scale=1.0,

1512 + o_sf_scale=1.0,

1513 + batch_size=prefill.chunked_context.seq_lens[chunk_idx].shape[0],

1514 + window_left=-1,

1515 + cum_seq_lens_q=prefill.query_start_loc,

1516 + cum_seq_lens_kv=prefill.chunked_context.cu_seq_lens[chunk_idx],

1517 + enable_pdl=False,

1518 + is_causal=False,

1519 + return_lse=True,

1520 + out=out,

1521 + )

1522 +

1523 + # Convert from (q_len, num_heads) to (num_heads, q_len)

1524 + return attn_out, lse.transpose(0, 1).contiguous()

1525 +

1526 def process_weights_after_loading(self, act_dtype: torch.dtype):

1527 def get_layer_weight(layer):

1528 WEIGHT_NAMES = ("weight", "qweight", "weight_packed")