diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py index afad6479067d..8cd0793b8d03 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py @@ -79,6 +79,9 @@ def __init__(self, attn: "TrtllmAttention") -> None: self._context_wrappers: dict[int, "BatchPrefillPagedTSWrapper"] = {} self._decode_wrappers: dict[int, "BatchDecodePagedTSWrapper"] = {} self._mla_decode_wrappers: dict[int, "BatchMLADecodePagedTSWrapper"] = {} + # Dense MLA's quantization scales are fixed for this layer/model. PrimTS + # takes host scalars, so read them once during eager warmup, not capture. + self._mla_fp8_scales: tuple[float, float] | None = None # Decode plans retain views into the shared workspace and are invalidated # whenever its underlying allocation changes. self._workspace_allocation: Optional[tuple[object, ...]] = None @@ -288,8 +291,12 @@ def _is_supported_with_reason( quant_mode = QuantMode(attn.quant_mode) except (TypeError, ValueError): return False, "invalid KV-cache quantization mode." - if quant_mode.has_kv_cache_quant(): - return False, "quantized KV cache is not supported by the initial adapter." + is_mla = attn.is_mla_enable + is_fp8_mla = is_mla and quant_mode.has_fp8_kv_cache() + if quant_mode.has_kv_cache_quant() and ( + not is_fp8_mla or quant_mode.has_int8_kv_cache() or quant_mode.has_fp4_kv_cache() + ): + return False, "quantized KV cache is supported only for FP8 MLA decode." input_type = fwd.attention_input_type if input_type not in ( @@ -315,7 +322,6 @@ def _is_supported_with_reason( ) if attn.num_heads <= 0 or attn.num_kv_heads <= 0: return False, "query and KV head counts must be positive." - is_mla = attn.is_mla_enable if is_mla: if attn.num_kv_heads != 1: return False, "MLA decode requires one logical KV head." @@ -330,7 +336,9 @@ def _is_supported_with_reason( if q.dtype not in self.SUPPORTED_DTYPES: return False, f"query dtype {q.dtype} is unsupported." cache_dtype = binding_to_torch_dtype(meta.kv_cache_manager.dtype) - if cache_dtype != q.dtype: + if is_fp8_mla and cache_dtype != torch.float8_e4m3fn: + return False, "FP8 MLA decode requires an FP8 E4M3 KV cache." + if not is_fp8_mla and cache_dtype != q.dtype: return False, f"query and KV-cache dtypes must match, got {q.dtype} and {cache_dtype}." if output.dtype != q.dtype: return False, f"output dtype must match query dtype, got {output.dtype} and {q.dtype}." @@ -341,7 +349,7 @@ def _is_supported_with_reason( if has_context or input_type != AttentionInputType.generation_only: return False, "MLA is supported only for generation-only requests." if q.dtype != torch.bfloat16 or output.dtype != torch.bfloat16: - return False, "MLA decode requires BF16 query, cache, and output." + return False, "MLA decode requires BF16 query input and output." if attn.kv_lora_rank != 512 or attn.qk_rope_head_dim != 64: return False, ( "MLA decode requires kv_lora_rank=512 and qk_rope_head_dim=64, got " @@ -356,6 +364,17 @@ def _is_supported_with_reason( return False, f"MLA query width must be {expected_width}, got {q.shape[1]}." if output.numel() != q.shape[0] * attn.num_heads * attn.kv_lora_rank: return False, "MLA output has an incompatible extent." + if is_fp8_mla: + # The automatic 1CTA policy for <=64 query rows has no profile + # below 128 KV tokens. This is the plan's page-table capacity, + # not the live request length; short requests remain supported. + if attn.num_heads <= 64 and ( + meta.kv_cache_block_offsets.shape[-1] * tokens_per_block < 128 + ): + return False, "FP8 MLA with <=64 heads requires a paged KV capacity >=128." + error = self._mla_fp8_input_error(q, fwd) + if error is not None: + return False, error else: expected_width = (attn.num_heads + 2 * attn.num_kv_heads) * attn.head_dim if q.shape[1] != expected_width: @@ -412,6 +431,60 @@ def _is_supported_with_reason( return False, "the K-to-V page displacement could not be resolved." return True, "" + @staticmethod + def _mla_fp8_input_error(q: torch.Tensor, fwd: AttentionForwardArgs) -> str | None: + """Check the preprocessing buffers without reading device values.""" + quant_q = fwd.quant_q_buffer + if quant_q is None: + return "FP8 MLA decode requires quant_q_buffer from MLA preprocessing." + if ( + quant_q.dtype not in (torch.uint8, torch.float8_e4m3fn) + or quant_q.device != q.device + or not quant_q.is_contiguous() + or quant_q.numel() < q.numel() + ): + return ( + "FP8 MLA quant_q_buffer must be contiguous uint8 or FP8 E4M3 on the " + "query device, with at least one element per query element." + ) + for name, scale in ( + ("mla_bmm1_scale", fwd.mla_bmm1_scale), + ("mla_bmm2_scale", fwd.mla_bmm2_scale), + ): + if ( + scale is None + or scale.dtype != torch.float32 + or scale.device != q.device + or not scale.is_contiguous() + or scale.numel() < 1 + ): + return ( + f"FP8 MLA {name} must contain a contiguous float32 value on the query device." + ) + return None + + def _get_mla_fp8_inputs( + self, q: torch.Tensor, fwd: AttentionForwardArgs + ) -> tuple[torch.Tensor, float, float]: + error = self._mla_fp8_input_error(q, fwd) + if error is not None: + raise RuntimeError(error) + if self._mla_fp8_scales is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "PrimTS FP8 MLA scales must be cached before CUDA graph capture." + ) + # The producer stores the regular BMM1 scale at index 0, and its + # log2 version at index 1. PrimTS expects the regular scale. + bmm1_scale = float(fwd.mla_bmm1_scale.view(-1)[0].item()) + bmm2_scale = float(fwd.mla_bmm2_scale.view(-1)[0].item()) + if not all(math.isfinite(scale) and scale > 0 for scale in (bmm1_scale, bmm2_scale)): + raise RuntimeError("PrimTS FP8 MLA scales must be finite and positive.") + self._mla_fp8_scales = (bmm1_scale, bmm2_scale) + # Reinterpret the producer's bytes; do not quantize the BF16 query again. + query = fwd.quant_q_buffer.view(-1)[: q.numel()].view(torch.float8_e4m3fn) + return query, *self._mla_fp8_scales + @staticmethod def _get_fixed_block_tables( block_tables: torch.Tensor, @@ -699,6 +772,11 @@ def prepare_workspace( get_prims_ts_batch_mla_decode_workspace_size, ) + kernel_dtype = ( + torch.float8_e4m3fn + if QuantMode(self.attn.quant_mode).has_fp8_kv_cache() + else q.dtype + ) required_bytes = get_prims_ts_batch_mla_decode_workspace_size( batch_size, self.attn.num_heads, @@ -707,8 +785,8 @@ def prepare_workspace( int(metadata.tokens_per_block), max_seq_len, max_seq_len_q=seq_len_q, - q_dtype=q.dtype, - kv_dtype=q.dtype, + q_dtype=kernel_dtype, + kv_dtype=kernel_dtype, out_dtype=forward_args.output.dtype, mask_type=mask_type, device=q.device, @@ -1142,7 +1220,15 @@ def run_mla_generation(self, params: FmhaParams) -> None: raise RuntimeError("TRT-LLM did not return PrimTS MLA KV metadata.") fixed_block_tables = self._get_fixed_block_tables(block_tables, batch_size) seq_len_q = params.input_seq_length - query = params.qkv_input.view( + if QuantMode(attn.quant_mode).has_fp8_kv_cache(): + query, bmm1_scale, bmm2_scale = self._get_mla_fp8_inputs(params.qkv_input, params.fwd) + else: + query = params.qkv_input + bmm1_scale = 1.0 / ( + attn.q_scaling * math.sqrt(int(attn.qk_nope_head_dim) + int(attn.qk_rope_head_dim)) + ) + bmm2_scale = 1.0 + query = query.view( batch_size, seq_len_q, attn.num_heads, @@ -1155,9 +1241,6 @@ def run_mla_generation(self, params: FmhaParams) -> None: int(attn.kv_lora_rank), ) max_seq_len = int(block_tables.shape[-1]) * params.tokens_per_block - bmm1_scale = 1.0 / ( - attn.q_scaling * math.sqrt(int(attn.qk_nope_head_dim) + int(attn.qk_rope_head_dim)) - ) mask_type = self._get_prims_mask_type(params.fwd) seq_lens = self._get_sequence_lengths( params.sequence_lengths, @@ -1185,7 +1268,7 @@ def run_mla_generation(self, params: FmhaParams) -> None: seq_lens=seq_lens, out=output, bmm1_scale=bmm1_scale, - bmm2_scale=1.0, + bmm2_scale=bmm2_scale, validate=False, ) diff --git a/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py index 5be953227100..fea5d9d58dd2 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py +++ b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py @@ -15,6 +15,9 @@ import functools import inspect +import math +from types import SimpleNamespace +from unittest.mock import Mock import pytest import torch @@ -344,6 +347,193 @@ def test_prims_ts_deepseek_v3_lite_mla_generation( run_case(case) +@pytest.mark.parametrize("num_heads", [6, 12, 96]) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], ids=["v1", "v2"]) +def test_prims_ts_fp8_mla_preprocessing( + monkeypatch: pytest.MonkeyPatch, num_heads: int, use_kv_cache_manager_v2: bool +) -> None: + from test_attention_mla import RopeConfig, _run_test_for_backend + + from tensorrt_llm._torch.attention.backends.fmha.phased import FmhaParams + from tensorrt_llm._torch.attention.backends.fmha.prims_ts import PrimsTSFmha + + # Context remains on the regular backend; every decode must select PrimTS. + monkeypatch.setenv("TLLM_FMHA_LIBS", "prims_ts,fallback") + original_run = PrimsTSFmha.run_mla_generation + calls = [] + + def run_and_capture(fmha: PrimsTSFmha, params: FmhaParams) -> None: + assert params.qkv_input.dtype == torch.bfloat16 + assert params.fwd.quant_q_buffer.dtype == torch.uint8 + original_run(fmha, params) + eager = params.context_buf.clone() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + original_run(fmha, params) + params.context_buf.zero_() + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(params.context_buf, eager, atol=0, rtol=0) + calls.append(params.attn.local_layer_idx) + + monkeypatch.setattr(PrimsTSFmha, "run_mla_generation", run_and_capture) + _run_test_for_backend( + backend_name="TRTLLM", + num_heads=num_heads, + num_kv_heads=num_heads, + num_layers=2, + q_lora_rank=1536, + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + rope_config=RopeConfig( + num_attention_heads=num_heads, + max_position_embeddings=4096, + rope_scaling={ + "beta_fast": 32, + "beta_slow": 1, + "factor": 40.0, + "mscale": 1.0, + "mscale_all_dim": 1.0, + "original_max_position_embeddings": 4096, + "type": "yarn", + }, + ), + kv_cache_tokens_per_block=32, + device=torch.device("cuda"), + dtype=torch.bfloat16, + kv_cache_dtype=torch.float8_e4m3fn, + context_sequence_lengths=[129, 191], + generation_seq_len_q=1, + num_generation_steps=2, + v2_kv_cache=use_kv_cache_manager_v2, + ) + assert calls == [0, 1, 0, 1] + + +@pytest.mark.parametrize("num_heads", [6, 12, 96]) +@pytest.mark.parametrize("batch_size", [2, 65]) +def test_prims_ts_fp8_mla_scales_and_graph_replay( + monkeypatch: pytest.MonkeyPatch, num_heads: int, batch_size: int +) -> None: + import tensorrt_llm._torch.attention.backends.fmha.prims_ts as prims_ts_module + from tensorrt_llm._torch.attention.backends.fmha.phased import FmhaParams + from tensorrt_llm._torch.attention.backends.fmha.prims_ts import PrimsTSFmha + from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + AttentionInputType, + ) + from tensorrt_llm.quantization.mode import QuantMode + + torch.manual_seed(0) + device = torch.device("cuda") + attn = Mock( + is_mla_enable=True, + num_heads=num_heads, + num_kv_heads=1, + kv_lora_rank=512, + qk_rope_head_dim=64, + qk_nope_head_dim=128, + head_dim=576, + v_head_dim=128, + local_layer_idx=0, + quant_mode=QuantMode.FP8_KV_CACHE, + q_scaling=1.0, + ) + fmha = PrimsTSFmha(attn) + query = torch.randn(batch_size, num_heads, 576, device=device).to(torch.float8_e4m3fn) + kv_cache = torch.randn(batch_size * 2, 1, 32, 576, device=device).to(torch.float8_e4m3fn) + block_tables = torch.zeros((batch_size, 2, 4), dtype=torch.int32, device=device) + block_tables[:, 0, :2] = torch.arange(batch_size * 2, device=device).view(batch_size, 2) + seq_lens = torch.full((batch_size,), 33, dtype=torch.int32, device=device) + output = torch.empty((batch_size, num_heads * 512), dtype=torch.bfloat16, device=device) + # Deliberately use different, non-unit scales. BMM1[1] is log2-scaled and + # must not be passed to PrimTS; BMM2 must scale V independently of QK. + bmm1_scale, bmm2_scale = 0.025, 1.75 + fwd = AttentionForwardArgs( + attention_input_type=AttentionInputType.generation_only, + attention_window_size=128, + output=output, + quant_q_buffer=query.view(torch.uint8), + mla_bmm1_scale=torch.tensor([bmm1_scale, bmm1_scale * math.log2(math.e)], device=device), + mla_bmm2_scale=torch.tensor([bmm2_scale], device=device), + ) + metadata = SimpleNamespace( + host_kv_cache_pool_pointers=None, + host_kv_cache_pool_mapping=None, + kv_cache_block_offsets=block_tables, + num_contexts=0, + num_ctx_tokens=0, + num_generations=batch_size, + tokens_per_block=32, + ) + # Pool ownership/indexing is covered by the real v1/v2 preprocessing test. + monkeypatch.setattr( + prims_ts_module.thop, + "build_trtllm_gen_kv_cache_metadata", + lambda *args: (kv_cache, block_tables, None), + ) + # Poison the BF16 input: the adapter must consume the preprocessing bytes. + q = torch.full((batch_size, num_heads * 576), float("nan"), dtype=torch.bfloat16, device=device) + workspace = torch.empty(0, dtype=torch.uint8, device=device) + fmha.prepare_workspace(q, None, None, metadata, fwd, workspace) + params = FmhaParams( + attn=attn, + meta=metadata, + fwd=fwd, + workspace=workspace, + qkv_input=q, + context_buf=output, + sequence_lengths=seq_lens, + input_seq_length=1, + num_tokens=batch_size, + batch_size=batch_size, + num_requests=batch_size, + tokens_per_block=32, + kv_factor=1, + total_num_blocks=batch_size * 2, + ) + + def check_output() -> None: + pages = kv_cache[:, 0].float()[block_tables[:, 0, :2].long()].reshape(batch_size, 64, 576) + scores = torch.einsum("bhd,bkd->bhk", query.float(), pages) * bmm1_scale + invalid = torch.arange(64, device=device)[None, :] >= seq_lens[:, None] + scores.masked_fill_(invalid[:, None, :], float("-inf")) + ideal = torch.einsum("bhk,bkd->bhd", scores.softmax(-1), pages[..., :512]) * bmm2_scale + # This case fits one KV tile. FP8 MLA rounds the unnormalized P tile + # (scaled to E4M3's maximum 448) before PV, but keeps its sum in FP32. + # Model that intermediate rounding for the elementwise comparison. + probabilities = (scores - scores.amax(dim=-1, keepdim=True)).exp() * 448.0 + rounded_p = probabilities.to(torch.float8_e4m3fn).float() + expected = ( + torch.einsum("bhk,bkd->bhd", rounded_p, pages[..., :512]) + / probabilities.sum(dim=-1, keepdim=True) + * bmm2_scale + ) + torch.testing.assert_close( + output, expected.reshape_as(output).to(output.dtype), atol=2e-2, rtol=2e-2 + ) + # Also bound aggregate error against full-precision attention so the + # low-precision oracle does not hide a large numerical regression. + relative_error = (output.float().view_as(ideal) - ideal).norm() / ideal.norm() + assert relative_error.item() < 0.04 + + fmha.run_mla_generation(params) + check_output() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + fmha.run_mla_generation(params) + # New query bytes and page-table/length contents must be read on replay. + query.view(torch.uint8).copy_(query.view(torch.uint8).flip(0)) + block_tables[:, 0, :2].copy_(block_tables[:, 0, :2].flip(0)) + seq_lens.add_(7) + output.zero_() + graph.replay() + torch.cuda.synchronize() + check_output() + + def test_prims_ts_context_wrapper_cuda_graph_replay_with_updated_metadata( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/tests/unittest/_torch/attention/test_prims_ts_fmha.py b/tests/unittest/_torch/attention/test_prims_ts_fmha.py index 87dbd7364f4b..c2ca7658fb9d 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_fmha.py +++ b/tests/unittest/_torch/attention/test_prims_ts_fmha.py @@ -41,6 +41,7 @@ from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm.bindings import DataType +from tensorrt_llm.quantization.mode import QuantMode class _TensorSpec: @@ -126,6 +127,8 @@ def _support_result( dtype: torch.dtype = torch.bfloat16, output_dtype: torch.dtype | None = None, kv_dtype: DataType | None = None, + quant_mode: QuantMode = QuantMode(0), + mla_fp8_inputs: dict | None = None, tokens_per_block: int = 32, is_mla: bool = False, is_fused_qkv: bool = True, @@ -159,6 +162,7 @@ def _support_result( num_kv_heads=num_kv_heads, ) attn.position_embedding_type = position_embedding_type + attn.quant_mode = quant_mode attn.attention_chunk_size = attention_chunk_size if has_sparse_attention: attn.sparse_params = SimpleNamespace(algorithm="mqa_gqa") @@ -184,6 +188,12 @@ def _support_result( relative_attention_bias=torch.empty(1) if has_relative_attention_bias else None, is_fused_qkv=is_fused_qkv, ) + if is_mla and quant_mode.has_fp8_kv_cache(): + forward_args.quant_q_buffer = _TensorSpec(q.shape, torch.uint8) + forward_args.mla_bmm1_scale = _TensorSpec((2,), torch.float32) + forward_args.mla_bmm2_scale = _TensorSpec((1,), torch.float32) + if mla_fp8_inputs is not None: + forward_args = replace(forward_args, **mla_fp8_inputs) if has_sparse_runtime_metadata: forward_args.sparse_runtime_params.sparse_kv_indices = torch.empty(1) if attention_input_type == AttentionInputType.context_only: @@ -191,7 +201,7 @@ def _support_result( kv_lens = [4] elif attention_input_type == AttentionInputType.generation_only: num_contexts, num_generations, num_ctx_tokens = 0, 4, 0 - kv_lens = [128, 96, 64, 32] + kv_lens = [min(length, max_seq_len) for length in (128, 96, 64, 32)] else: num_contexts, num_generations, num_ctx_tokens = 1, 1, 3 kv_lens = [3, 128] @@ -222,7 +232,11 @@ def _support_result( is_spec_dec_tree=is_spec_dec_tree, is_spec_dec_dynamic_tree=False, is_spec_decoding_enabled=use_spec_decoding, - kv_cache_block_offsets=torch.empty(1) if has_paged_cache else None, + kv_cache_block_offsets=torch.empty( + num_contexts + num_generations, 2, math.ceil(max_seq_len / tokens_per_block) + ) + if has_paged_cache + else None, host_kv_cache_pool_pointers=torch.empty(1), host_kv_cache_pool_mapping=torch.zeros((1, 2), dtype=torch.int32), kv_cache_manager=kv_cache_manager, @@ -299,6 +313,84 @@ def test_supported_matrix(case: dict) -> None: assert supported, reason +@pytest.mark.parametrize("num_heads", [6, 12, 96, 128]) +@pytest.mark.parametrize("use_kv_cache_v2", [False, True]) +@pytest.mark.parametrize("tokens_per_block", [16, 32, 64, 128]) +def test_fp8_mla_supported(num_heads: int, use_kv_cache_v2: bool, tokens_per_block: int) -> None: + supported, reason = _support_result( + attention_input_type=AttentionInputType.generation_only, + head_dim=576, + is_mla=True, + num_heads=num_heads, + quant_mode=QuantMode.FP8_KV_CACHE, + kv_dtype=DataType.FP8, + use_kv_cache_v2=use_kv_cache_v2, + tokens_per_block=tokens_per_block, + ) + assert supported, reason + + +@pytest.mark.parametrize( + ("overrides", "expected_reason"), + [ + ({"is_mla": False, "head_dim": 128}, "only for FP8 MLA"), + ({"quant_mode": QuantMode.INT8_KV_CACHE}, "only for FP8 MLA"), + ({"quant_mode": QuantMode.NVFP4_KV_CACHE}, "only for FP8 MLA"), + ({"quant_mode": QuantMode.FP8_KV_CACHE | QuantMode.INT8_KV_CACHE}, "only for FP8 MLA"), + ({"kv_dtype": DataType.BF16}, "FP8 E4M3 KV cache"), + ({"quant_mode": QuantMode(0)}, "query and KV-cache dtypes"), + ({"dtype": torch.float16}, "BF16 query input and output"), + ({"output_dtype": torch.float8_e4m3fn}, "output dtype must match"), + ({"attention_input_type": AttentionInputType.context_only}, "generation-only"), + ({"attention_input_type": AttentionInputType.mixed}, "generation-only"), + ({"use_spec_decoding": True}, "speculative decoding"), + ({"has_sparse_attention": True}, "sparse attention"), + ({"max_seq_len": 64}, "paged KV capacity >=128"), + ], +) +def test_fp8_mla_unsupported(overrides: dict, expected_reason: str) -> None: + case = dict( + attention_input_type=AttentionInputType.generation_only, + head_dim=576, + is_mla=True, + quant_mode=QuantMode.FP8_KV_CACHE, + kv_dtype=DataType.FP8, + ) + case.update(overrides) + supported, reason = _support_result(**case) + assert not supported + assert expected_reason in reason + + +@pytest.mark.parametrize( + ("name", "buffer"), + [ + ("quant_q_buffer", None), + ("quant_q_buffer", _TensorSpec((4, 8 * 576), torch.bfloat16)), + ("quant_q_buffer", _TensorSpec((1,), torch.uint8)), + ("quant_q_buffer", _TensorSpec((4, 8 * 576), torch.uint8, device="cpu")), + ("quant_q_buffer", _TensorSpec((4, 8 * 576), torch.uint8, contiguous=False)), + ("mla_bmm1_scale", None), + ("mla_bmm1_scale", _TensorSpec((2,), torch.float16)), + ("mla_bmm1_scale", _TensorSpec((2,), torch.float32, device="cpu")), + ("mla_bmm1_scale", _TensorSpec((2,), torch.float32, contiguous=False)), + ("mla_bmm2_scale", None), + ("mla_bmm2_scale", _TensorSpec((0,), torch.float32)), + ], +) +def test_fp8_mla_requires_preprocessed_buffers(name: str, buffer: _TensorSpec | None) -> None: + supported, reason = _support_result( + attention_input_type=AttentionInputType.generation_only, + head_dim=576, + is_mla=True, + quant_mode=QuantMode.FP8_KV_CACHE, + kv_dtype=DataType.FP8, + mla_fp8_inputs={name: buffer}, + ) + assert not supported + assert name in reason + + @pytest.mark.parametrize("phase", [FmhaPhase.CONTEXT, FmhaPhase.GENERATION]) def test_phase_support_check_preserves_whole_request_semantics(phase: FmhaPhase) -> None: supported, reason = _support_result( @@ -1390,13 +1482,17 @@ def _get_test_mla_wrapper( ) +@pytest.mark.parametrize("fp8", [False, True], ids=["bf16", "fp8"]) def test_mla_eager_wrapper_plans_once_and_reads_live_fixed_metadata( monkeypatch: pytest.MonkeyPatch, + fp8: bool, ) -> None: monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) attn = _Attention(head_dim=576, is_mla=True, num_heads=4) + attn.quant_mode = QuantMode.FP8_KV_CACHE if fp8 else QuantMode(0) fmha = PrimsTSFmha(attn) - kv_cache = torch.empty((20, 1, 32, 576), dtype=torch.bfloat16) + kernel_dtype = torch.float8_e4m3fn if fp8 else torch.bfloat16 + kv_cache = torch.empty((20, 1, 32, 576), dtype=kernel_dtype) block_tables = torch.tensor( [ [[0, 1, 2], [10, 11, 12]], @@ -1436,6 +1532,11 @@ def test_mla_eager_wrapper_plans_once_and_reads_live_fixed_metadata( attention_input_type=AttentionInputType.generation_only, attention_window_size=64, is_fused_qkv=True, + quant_q_buffer=torch.zeros(2 * attn.num_heads * 576 + 16, dtype=torch.uint8) + if fp8 + else None, + mla_bmm1_scale=torch.tensor([0.125, 0.125 * math.log2(math.e)]) if fp8 else None, + mla_bmm2_scale=torch.tensor([2.0]) if fp8 else None, ) sequence_lengths = torch.tensor([33, 64], dtype=torch.int32) assert sequence_lengths.data_ptr() % 16 == 0 @@ -1481,8 +1582,8 @@ def test_mla_eager_wrapper_plans_once_and_reads_live_fixed_metadata( assert {key: value for key, value in plan_kwargs.items() if key != "workspace_buffer"} == { "max_seq_len_q": 1, "packed_query": False, - "q_data_type": torch.bfloat16, - "kv_data_type": torch.bfloat16, + "q_data_type": kernel_dtype, + "kv_data_type": kernel_dtype, "o_data_type": torch.bfloat16, "mask_type": "causal", } @@ -1490,7 +1591,9 @@ def test_mla_eager_wrapper_plans_once_and_reads_live_fixed_metadata( run_args = wrapper.run.call_args.args run_kwargs = wrapper.run.call_args.kwargs assert run_args[0].shape == (2, 1, attn.num_heads, 576) - assert run_args[0].data_ptr() == params.qkv_input.data_ptr() + query_buffer = forward_args.quant_q_buffer if fp8 else params.qkv_input + assert run_args[0].data_ptr() == query_buffer.data_ptr() + assert run_args[0].dtype == kernel_dtype assert run_args[1] is kv_cache assert run_kwargs["block_tables"].shape == (2, 3) assert run_kwargs["block_tables"].stride() == (6, 1) @@ -1503,8 +1606,8 @@ def test_mla_eager_wrapper_plans_once_and_reads_live_fixed_metadata( torch.testing.assert_close(run_kwargs["seq_lens"], sequence_lengths) assert run_kwargs["out"].shape == (2, 1, attn.num_heads, 512) assert run_kwargs["out"].data_ptr() == output.data_ptr() - assert run_kwargs["bmm1_scale"] == pytest.approx(1.0 / math.sqrt(128 + 64)) - assert run_kwargs["bmm2_scale"] == 1.0 + assert run_kwargs["bmm1_scale"] == pytest.approx(0.125 if fp8 else 1.0 / math.sqrt(128 + 64)) + assert run_kwargs["bmm2_scale"] == (2.0 if fp8 else 1.0) assert run_kwargs["validate"] is False block_tables[:, 0].add_(20) @@ -1525,6 +1628,68 @@ def test_mla_eager_wrapper_plans_once_and_reads_live_fixed_metadata( assert fmha._mla_decode_wrappers[2] is wrapper +@pytest.mark.parametrize("buffer_dtype", [torch.uint8, torch.float8_e4m3fn]) +def test_mla_fp8_inputs_reuse_scales_without_device_reads( + monkeypatch: pytest.MonkeyPatch, buffer_dtype: torch.dtype +) -> None: + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + fmha = PrimsTSFmha(_Attention(head_dim=576, is_mla=True)) + q = torch.empty((2, 8 * 576), dtype=torch.bfloat16) + quant_q = torch.arange(q.numel() + 16, dtype=torch.int32).to(torch.uint8).view(buffer_dtype) + fwd = AttentionForwardArgs( + quant_q_buffer=quant_q, + mla_bmm1_scale=torch.tensor([0.125, 0.125 * math.log2(math.e)]), + mla_bmm2_scale=torch.tensor([2.0]), + ) + query, bmm1, bmm2 = fmha._get_mla_fp8_inputs(q, fwd) + assert query.dtype == torch.float8_e4m3fn + assert query.numel() == q.numel() + assert query.data_ptr() == quant_q.data_ptr() + torch.testing.assert_close(query.view(torch.uint8), quant_q.view(torch.uint8)[: q.numel()]) + assert (bmm1, bmm2) == (0.125, 2.0) + + # Each request gets fresh preprocessing buffers with the same layer-static + # scale values. Neither eager reuse nor capture may read those values again. + fwd = replace( + fwd, + quant_q_buffer=quant_q.clone(), + mla_bmm1_scale=fwd.mla_bmm1_scale.clone(), + mla_bmm2_scale=fwd.mla_bmm2_scale.clone(), + ) + device_read = Mock(side_effect=AssertionError("unexpected device scalar read")) + monkeypatch.setattr(torch.Tensor, "item", device_read) + for capturing in (False, True): + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", Mock(return_value=capturing)) + query, bmm1, bmm2 = fmha._get_mla_fp8_inputs(q, fwd) + assert query.data_ptr() == fwd.quant_q_buffer.data_ptr() + assert (bmm1, bmm2) == (0.125, 2.0) + device_read.assert_not_called() + + fmha._mla_fp8_scales = None + with pytest.raises(RuntimeError, match="scales must be cached before CUDA graph capture"): + fmha._get_mla_fp8_inputs(q, fwd) + device_read.assert_not_called() + + +@pytest.mark.parametrize("scale_name", ["mla_bmm1_scale", "mla_bmm2_scale"]) +@pytest.mark.parametrize("value", [0.0, -1.0, float("nan"), float("inf")]) +def test_mla_fp8_rejects_invalid_scales( + monkeypatch: pytest.MonkeyPatch, scale_name: str, value: float +) -> None: + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + fmha = PrimsTSFmha(_Attention(head_dim=576, is_mla=True)) + q = torch.empty((1, 8 * 576), dtype=torch.bfloat16) + fwd = AttentionForwardArgs( + quant_q_buffer=torch.empty(q.shape, dtype=torch.uint8), + mla_bmm1_scale=torch.ones(2), + mla_bmm2_scale=torch.ones(1), + ) + fwd = replace(fwd, **{scale_name: torch.tensor([value])}) + with pytest.raises(RuntimeError, match="scales must be finite and positive"): + fmha._get_mla_fp8_inputs(q, fwd) + assert fmha._mla_fp8_scales is None + + def test_mla_wrapper_cache_plans_each_batch_once_and_reuses_a_b_a( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -1750,9 +1915,11 @@ def test_mla_wrapper_receives_v2_bound_and_shared_workspace( @pytest.mark.parametrize("is_cuda_graph", [False, True], ids=["eager", "cuda-graph"]) +@pytest.mark.parametrize("fp8", [False, True], ids=["bf16", "fp8"]) def test_mla_prepare_workspace_sizes_caller_owned_workspace( monkeypatch: pytest.MonkeyPatch, is_cuda_graph: bool, + fp8: bool, ) -> None: monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) workspace_size = Mock(return_value=48) @@ -1762,6 +1929,7 @@ def test_mla_prepare_workspace_sizes_caller_owned_workspace( workspace_size, ) attn = _Attention(head_dim=576, is_mla=True, num_heads=4) + attn.quant_mode = QuantMode.FP8_KV_CACHE if fp8 else QuantMode(0) fmha = PrimsTSFmha(attn) fmha._multi_processor_count = 1 q = torch.empty((2, attn.num_heads * 576), dtype=torch.bfloat16) @@ -1793,6 +1961,10 @@ def test_mla_prepare_workspace_sizes_caller_owned_workspace( ) workspace_size.assert_called_once() + kernel_dtype = torch.float8_e4m3fn if fp8 else torch.bfloat16 + assert workspace_size.call_args.kwargs["q_dtype"] == kernel_dtype + assert workspace_size.call_args.kwargs["kv_dtype"] == kernel_dtype + assert workspace_size.call_args.kwargs["out_dtype"] == torch.bfloat16 assert workspace.numel() == 48