@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 578977bb5e

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 8, 12:30 UTC

[SM100] Resubmit FMHA FP8 prefill for MLA (#31195)

Signed-off-by: Pavani Majety <pmajety@nvidia.com>

Pavani Majety committed Feb 10, 2026 at 13:18 UTC 578977bb5ed208c62cf9cff80d955836775e0d24

3 files changed +145 -23

tests/v1/attention/test_mla_backends.py

+4 -3

@@ -27,7 +27,7 @@ from vllm.v1.attention.backend import CommonAttentionMetadata

27 from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla

28 from vllm.v1.attention.backends.registry import AttentionBackendEnum

29 from vllm.v1.attention.ops.flashmla import is_flashmla_dense_supported

30 -from vllm.v1.kv_cache_interface import FullAttentionSpec

30 +from vllm.v1.kv_cache_interface import MLAAttentionSpec

32 BACKENDS_TO_TEST = [

33 AttentionBackendEnum.CUTLASS_MLA,

@@ -512,7 +512,7 @@ class MockMLAAttentionLayer(AttentionLayerBase):

512

513 def run_attention_backend(

514 backend: AttentionBackendEnum,

515 - kv_cache_spec: FullAttentionSpec,

515 + kv_cache_spec: MLAAttentionSpec,

516 layer_names: list[str],

517 vllm_config,

518 device: torch.device,

@@ -989,7 +989,7 @@ def test_backend_correctness(

989 kv_cache = kv_cache_per_block_size[block_size]

990

991 # Create kv_cache_spec with the correct block_size for this backend

992 - backend_kv_cache_spec = FullAttentionSpec(

992 + backend_kv_cache_spec = MLAAttentionSpec(

993 block_size=block_size,

994 num_kv_heads=vllm_config.model_config.get_num_kv_heads(

995 vllm_config.parallel_config

@@ -997,6 +997,7 @@ def test_backend_correctness(

997 head_size=vllm_config.model_config.get_head_size(),

998 dtype=vllm_config.model_config.dtype,

999 sliding_window=vllm_config.model_config.get_sliding_window(),

1000 + cache_dtype_str=vllm_config.cache_config.cache_dtype,

1001 )

1002

1003 backend_output = run_attention_backend(

vllm/config/attention.py

@@ -43,6 +43,9 @@ class AttentionConfig:

43 disable_flashinfer_q_quantization: bool = False

44 """If set, when using fp8 kv, do not quantize Q to fp8."""

46 + use_prefill_query_quantization: bool = False

47 + """If set, quantize query for attention in prefill."""

48 +

49 def compute_hash(self) -> str:

50 """

51 Provide a hash that uniquely identifies all the configs

vllm/model_executor/layers/attention/mla_attention.py

+138 -20

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

1052 query_seq_lens: torch.Tensor | None = None

1053 workspace_buffer: torch.Tensor | None = None

1054 q_data_type: torch.dtype | None = None

1055 + output_dtype: torch.dtype | None = None

1056

1057

1058 @dataclass

@@ -1145,6 +1146,7 @@ def is_deepseek_r1_mla_compatible(vllm_config: VllmConfig) -> bool:

1146 return qk_nope_head_dim == 128 and qk_rope_head_dim == 64 and v_head_dim == 128

1147

1148

1149 +@functools.cache

1150 def use_flashinfer_prefill() -> bool:

1151 # For blackwell default to flashinfer prefill if it's available since

1152 # it is faster than FA2.

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

1164 return is_deepseek_r1_mla_compatible(vllm_config)

1165

1166

1167 +@functools.cache

1168 def use_cudnn_prefill() -> bool:

1169 from vllm.config import get_current_vllm_config

1170

@@ -1174,6 +1177,7 @@ def use_cudnn_prefill() -> bool:

1177 )

1178

1179

1180 +@functools.cache

1181 def use_trtllm_ragged_deepseek_prefill() -> bool:

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

1183 from vllm.config import get_current_vllm_config

@@ -1210,6 +1214,27 @@ def get_mla_dims(model_config: ModelConfig) -> MLADims:

1214 )

1215

1216

1217 +@functools.cache

1218 +def backend_supports_prefill_query_quantization() -> bool:

1219 + """Check if the selected MLA backend supports prefill query quantization.

1220 +

1221 + Currently supported backends:

1222 + - FlashInfer prefill

1223 + - TRT-LLM ragged DeepSeek prefill

1224 +

1225 + Not supported:

1226 + - cuDNN Prefill

1227 + - FlashAttention

1228 + - Non-GB200 devices (FP8 prefill requires device capability 100)

1229 + """

1230 + # FP8 prefill query quantization requires GB200 (device capability 100)

1231 + # for the necessary FP8 kernels at the moment.

1232 + if not current_platform.is_device_capability_family(100):

1233 + return False

1234 +

1235 + return use_flashinfer_prefill() or use_trtllm_ragged_deepseek_prefill()

1236 +

1237 +

1238 class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

1239 """

1240 NOTE: Please read the comment at the top of the file before trying to

@@ -1262,6 +1287,40 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

1287

1288 return chunked_prefill_workspace_size

1289

1290 + @staticmethod

1291 + def determine_prefill_query_data_type(

1292 + vllm_config: VllmConfig,

1293 + model_dtype: torch.dtype,

1294 + ) -> torch.dtype:

1295 + """

1296 + Determine the query data type for prefill queries.

1297 + Return FP8 dtype if cache is FP8 and prefill query quantization

1298 + is enabled, else model dtype.

1299 + """

1300 + use_fp8 = (

1301 + vllm_config.cache_config.cache_dtype.startswith("fp8")

1302 + and vllm_config.attention_config.use_prefill_query_quantization

1303 + and backend_supports_prefill_query_quantization()

1304 + )

1305 +

1306 + if use_fp8:

1307 + fp8_dtype = current_platform.fp8_dtype()

1308 + logger.info_once(

1309 + "FP8 prefill attention enabled: query data type is FP8", scope="local"

1310 + )

1311 + return fp8_dtype

1312 + elif vllm_config.attention_config.use_prefill_query_quantization:

1313 + logger.info_once(

1314 + "Unable to perform FP8 prefill attention when"

1315 + " use_prefill_query_quantization is enabled. Please"

1316 + " ensure that --kv-cache-dtype is set to fp8 and your prefill"

1317 + " backend is compatible with FP8 attention.",

1318 + scope="local",

1319 + )

1320 + return model_dtype

1321 +

1322 + return model_dtype

1323 +

1324 def __init__(

1325 self,

1326 kv_cache_spec: AttentionSpec,

@@ -1285,6 +1344,12 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

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

1345 self.mla_dims = get_mla_dims(self.model_config)

1346 self.aot_schedule = current_platform.is_cuda()

1347 +

1348 + self.kv_cache_spec = kv_cache_spec

1349 + self.q_data_type = self.determine_prefill_query_data_type(

1350 + vllm_config, self.model_config.dtype

1351 + )

1352 +

1353 try:

1354 self.dcp_world_size = get_dcp_group().world_size

1355 self.dcp_rank = get_dcp_group().rank_in_group

@@ -1325,7 +1390,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

1390 self.chunked_prefill_workspace_size,

1391 self.model_config.get_head_size(),

1392 ),

1328 - dtype=self.model_config.dtype,

1393 + dtype=self.q_data_type,

1394 device=device,

1395 )

1396

@@ -1435,7 +1500,8 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

1500 sm_scale=self._global_hyperparameters.sm_scale,

1501 window_left=self._global_hyperparameters.window_left,

1502 logits_soft_cap=self._global_hyperparameters.logits_soft_cap,

1438 - q_data_type=self.model_config.dtype,

1503 + q_data_type=self.q_data_type,

1504 + o_data_type=prefill.output_dtype,

1505 )

1506

1507 # Prepare context prefills

@@ -1454,7 +1520,8 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

1520 sm_scale=self._global_hyperparameters.sm_scale,

1521 window_left=self._global_hyperparameters.window_left,

1522 logits_soft_cap=self._global_hyperparameters.logits_soft_cap,

1457 - q_data_type=self.model_config.dtype,

1523 + q_data_type=self.q_data_type,

1524 + o_data_type=prefill.output_dtype,

1525 )

1526

1527 prefill.prefill_main = self._fi_prefill_main

@@ -1709,6 +1776,8 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

1776 query_start_loc=prefill_query_start_loc,

1777 max_query_len=max_query_len,

1778 chunked_context=chunked_context_metadata,

1779 + output_dtype=self.model_config.dtype,

1780 + q_data_type=self.q_data_type,

1781 )

1782

1783 if self._use_cudnn_prefill:

@@ -1894,7 +1963,6 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

1963 self.kv_b_proj = kv_b_proj

1964 self.indexer = indexer

1965 self.q_pad_num_heads = q_pad_num_heads

1897 -

1966 self.supports_quant_query_input = True

1967

1968 # Use flashinfer's optimized concat_mla_k kernel when available.

@@ -2129,6 +2197,14 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

2197

2198 assert prefill.query_seq_lens is not None

2199 assert prefill.workspace_buffer is not None

2200 + # allocate BF16 / FP16 output tensor for TRT-LLM ragged attention

2201 + out = torch.empty(

2202 + q.shape[0],

2203 + q.shape[1],

2204 + v.shape[2],

2205 + device=q.device,

2206 + dtype=prefill.output_dtype,

2207 + )

2208

2209 ret = trtllm_ragged_attention_deepseek(

2210 query=q,

@@ -2148,6 +2224,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

2224 enable_pdl=False,

2225 is_causal=True,

2226 return_lse=return_softmax_lse,

2227 + out=out,

2228 )

2229

2230 if isinstance(ret, tuple):

@@ -2170,7 +2247,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

2247 q.shape[1],

2248 v.shape[2],

2249 device=q.device,

2173 - dtype=q.dtype,

2250 + dtype=prefill.output_dtype,

2251 )

2252 prefill.workspace_buffer.fill_(0)

2253

@@ -2240,29 +2317,59 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

2317 prefill_metadata = attn_metadata.prefill

2318 assert prefill_metadata.chunked_context is not None

2319

2320 + use_fp8_prefill = prefill_metadata.q_data_type == current_platform.fp8_dtype()

2321 +

2322 output = None

2323 iters = len(prefill_metadata.chunked_context.seq_tot)

2324 workspace = prefill_metadata.chunked_context.workspace

2325 +

2326 + if use_fp8_prefill:

2327 + q = q.to(prefill_metadata.q_data_type)

2328 +

2329 for i in range(iters):

2330 toks = prefill_metadata.chunked_context.seq_tot[i]

2248 - ops.gather_and_maybe_dequant_cache(

2249 - src_cache=kv_c_and_k_pe_cache,

2250 - dst=workspace,

2251 - block_table=prefill_metadata.block_table,

2252 - cu_seq_lens=prefill_metadata.chunked_context.cu_seq_lens[i],

2253 - token_to_seq=prefill_metadata.chunked_context.token_to_seq[i],

2254 - num_tokens=prefill_metadata.chunked_context.chunk_total_token[i],

2255 - kv_cache_dtype=self.kv_cache_dtype,

2256 - scale=k_scale,

2257 - seq_starts=prefill_metadata.chunked_context.starts[i],

2258 - )

2331 + if not use_fp8_prefill:

2332 + ops.gather_and_maybe_dequant_cache(

2333 + src_cache=kv_c_and_k_pe_cache,

2334 + dst=workspace,

2335 + block_table=prefill_metadata.block_table,

2336 + cu_seq_lens=prefill_metadata.chunked_context.cu_seq_lens[i],

2337 + token_to_seq=prefill_metadata.chunked_context.token_to_seq[i],

2338 + num_tokens=prefill_metadata.chunked_context.chunk_total_token[i],

2339 + kv_cache_dtype=self.kv_cache_dtype,

2340 + scale=k_scale,

2341 + seq_starts=prefill_metadata.chunked_context.starts[i],

2342 + )

2343 + else:

2344 + # FP8 path: gather cache without dequantization

2345 + ops.cp_gather_cache(

2346 + src_cache=kv_c_and_k_pe_cache,

2347 + dst=workspace,

2348 + block_table=prefill_metadata.block_table,

2349 + cu_seq_lens=prefill_metadata.chunked_context.cu_seq_lens[i],

2350 + batch_size=attn_metadata.num_prefills,

2351 + seq_starts=prefill_metadata.chunked_context.starts[i],

2352 + )

2353

2354 + # Extract kv_c_normed from workspace

2355 kv_c_normed = workspace[:toks][..., : self.kv_lora_rank]

2261 - k_pe = workspace[:toks][..., self.kv_lora_rank :].unsqueeze(1)

2356 + # When FP8 weights are used without FP8 prefill, kv_b_proj expects

2357 + # model dtype input and will quantize internally.

2358 + if (

2359 + use_fp8_prefill

2360 + or self.kv_b_proj.weight.dtype != current_platform.fp8_dtype()

2361 + ):

2362 + kv_c_normed = kv_c_normed.to(self.kv_b_proj.weight.dtype)

2363

2364 + k_pe = workspace[:toks][..., self.kv_lora_rank :].unsqueeze(1)

2365 kv_nope = self.kv_b_proj(kv_c_normed)[0].view(

2366 -1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim

2367 )

2368 +

2369 + # To Do: Use epilogue of kv_b_proj to generate fp8 kv_nope.

2370 + if use_fp8_prefill:

2371 + kv_nope = kv_nope.to(prefill_metadata.q_data_type)

2372 + k_pe = k_pe.to(prefill_metadata.q_data_type)

2373 k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)

2374

2375 k = self._concat_k_nope_k_pe(k_nope, k_pe)

@@ -2412,16 +2519,27 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

2519 assert attn_metadata.prefill is not None

2520 assert self.dcp_world_size != -1

2521

2415 - has_context = attn_metadata.prefill.chunked_context is not None

2522 + prefill_metadata = attn_metadata.prefill

2523 + use_fp8_prefill = prefill_metadata.q_data_type == current_platform.fp8_dtype()

2524 +

2525 + # Convert q to FP8 if FP8 prefill attention is enabled

2526 + if use_fp8_prefill:

2527 + q = q.to(prefill_metadata.q_data_type)

2528 +

2529 + has_context = prefill_metadata.chunked_context is not None

2530 +

2531 kv_nope = self.kv_b_proj(kv_c_normed)[0].view(

2532 -1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim

2533 )

2534 k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)

2420 -

2535 k = self._concat_k_nope_k_pe(k_nope, k_pe)

2536

2537 + if use_fp8_prefill:

2538 + k = k.to(prefill_metadata.q_data_type)

2539 + v = v.to(prefill_metadata.q_data_type)

2540 +

2541 output_prefill = self._run_prefill_new_tokens(

2424 - prefill=attn_metadata.prefill,

2542 + prefill=prefill_metadata,

2543 q=q,

2544 k=k,

2545 v=v,