[https://nvbugs/7016340][test] Remove the MLA mixed-precision KV numeric test

tests/unittest/_torch/attention/test_mixed_precision_backends.py

@@ -361,188 +361,6 @@ def _reference_dequant(packed, scales, scale):
     return table[codes] * block_scales * scale
 
 
-@pytest.mark.parametrize("backend", (TrtllmAttention, FlashInferAttention))
-@torch.no_grad()
-def test_mla_cached_context_and_decode(backend):
-    if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] not in (10, 12):
-        pytest.skip("Mixed-precision NVFP4 KV requires Blackwell")
-    from test_attention_mla import apply_mla_rope
-
-    from tensorrt_llm._torch.attention.backends.interface import (
-        PositionalEmbeddingParams,
-        RopeParams,
-    )
-    from tensorrt_llm._torch.attention.mla import MLA
-    from tensorrt_llm._torch.model_config import ModelConfig
-    from tensorrt_llm.functional import PositionEmbeddingType
-
-    torch.manual_seed(123)
-    heads, rank, nope, rope_dim, block = 16, 512, 128, 64, 128
-    length, prefix = 640, 512
-    mapping = Mapping(world_size=1, tp_size=1, rank=0)
-    rope = RopeParams(dim=rope_dim, max_positions=1024, duplicate_data=True)
-    model = MLA(
-        hidden_size=128,
-        num_attention_heads=heads,
-        num_key_value_heads=heads,
-        qk_nope_head_dim=nope,
-        qk_rope_head_dim=rope_dim,
-        v_head_dim=nope,
-        q_lora_rank=None,
-        kv_lora_rank=rank,
-        predicted_tokens_per_seq=1,
-        max_position_embeddings=1024,
-        bias=False,
-        layer_idx=0,
-        dtype=torch.bfloat16,
-        pos_embd_params=PositionalEmbeddingParams(
-            type=PositionEmbeddingType.rope_gptj, rope=rope, is_neox=False
-        ),
-        config=ModelConfig(
-            mapping=mapping,
-            attn_backend="TRTLLM" if backend is TrtllmAttention else "FLASHINFER",
-        ),
-    ).cuda()
-    for parameter in model.parameters():
-        if parameter.ndim == 1:
-            parameter.fill_(1)
-        else:
-            parameter.normal_(std=parameter.shape[-1] ** -0.5)
-    key_weight, value_weight = model.kv_b_proj.weight.split(heads * nope)
-    model.k_b_proj_trans.copy_(key_weight.view(heads, nope, rank).transpose(1, 2))
-    model.v_b_proj.copy_(value_weight.view(heads, nope, rank))
-    hidden = torch.randn(length + 1, 128, dtype=torch.bfloat16, device="cuda")
-
-    # Independent projections, RoPE and contiguous NVFP4 codec form the oracle.
-    query = torch.nn.functional.linear(hidden, model.q_b_proj.weight).view(
-        -1, heads, nope + rope_dim
-    )
-    latent, key_rope = torch.nn.functional.linear(hidden, model.kv_a_proj_with_mqa.weight).split(
-        (rank, rope_dim), dim=-1
-    )
-    latent = (
-        latent.float() * torch.rsqrt(latent.float().square().mean(-1, keepdim=True) + 1e-6)
-    ).to(hidden.dtype)
-    rope_table = rope.create_rope_const_params()[1].view(1024, -1, 2).transpose(1, 2)
-    positions = list(range(length + 1))
-    query_rope = apply_mla_rope(query[..., nope:], positions, rope_table)
-    key_rope = apply_mla_rope(key_rope, positions, rope_table)
-    if backend is FlashInferAttention:
-        # External GPT-J RoPE keeps pairs interleaved; native MLA uses split halves.
-        query_rope = query_rope.unflatten(-1, (2, -1)).transpose(-2, -1).flatten(-2)
-        key_rope = key_rope.unflatten(-1, (2, -1)).transpose(-2, -1).flatten(-2)
-    query[..., nope:] = query_rope
-    latent = torch.cat((latent, key_rope), dim=-1)
-    scale = torch.ones(1, device="cuda")
-    packed, scales = _reference_quant(latent, scale)
-    low_latent = _reference_dequant(packed, scales, scale).to(hidden.dtype)
-
-    manager = KVCacheManagerV2(
-        KvCacheConfig(
-            max_tokens=65536,
-            enable_block_reuse=True,
-            dtype="mixed_nvfp4_bfloat16",
-            mixed_precision_first_n=block,
-            mixed_precision_last_n=block,
-            pool_ratio=[0.75, 0.25],
-        ),
-        CacheType.SELFKONLY,
-        num_layers=1,
-        num_kv_heads=1,
-        head_dim=rank + rope_dim,
-        tokens_per_block=block,
-        max_seq_len=1024,
-        max_batch_size=1,
-        mapping=mapping,
-        dtype=DataType.BF16,
-    )
-    try:
-        assert manager.kv_factor == 1
-        manager.add_dummy_requests([0], [prefix], is_gen=False)
-        for begin, end, context in (
-            (0, prefix, True),
-            (prefix, length, True),
-            (length, length + 1, False),
-        ):
-            if begin:
-                cache = manager.kv_cache_map[0]
-                assert cache.resize(None, begin)
-                cache.suspend()
-                assert cache.resume(
-                    torch.cuda.current_stream().cuda_stream, is_decoding=not context
-                )
-                assert cache.resize(end)
-                manager._restore_page_index_bufs(0, cache)
-            metadata = backend.Metadata(
-                seq_lens=torch.tensor([end - begin], dtype=torch.int32),
-                num_contexts=int(context),
-                kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[begin]),
-                max_num_requests=1,
-                max_num_tokens=length + 1,
-                kv_cache_manager=manager,
-                request_ids=[0],
-                prompt_lens=[end - begin if context else length],
-                kv_layout="NHD",
-                mapping=mapping,
-                enable_context_mla_with_cached_kv=True,
-                runtime_features=AttentionRuntimeFeatures(
-                    chunked_prefill=True,
-                    cache_reuse=True,
-                    chunk_size=prefix,
-                    chunked_prefill_buffer_batch_size=1,
-                ),
-            )
-            metadata.prepare()
-            if backend is TrtllmAttention:
-                view = metadata.mixed_precision_staging
-                paged_context_fmha = metadata.use_paged_context_fmha
-                native_kv = metadata._swap_mixed_precision_kv(view)
-                assert metadata.kv_cache_manager is view
-                assert metadata.use_paged_context_fmha == paged_context_fmha
-                metadata._restore_mixed_precision_kv(native_kv)
-                assert metadata.kv_cache_manager is manager
-                assert metadata.use_paged_context_fmha == paged_context_fmha
-            _, high_pages = manager.get_mixed_precision_page_tables([0], 0)
-            if begin:
-                assert high_pages[0][0] >= 0 and high_pages[0][1] == kv.BAD_PAGE_INDEX
-            high_mask = (
-                torch.tensor(high_pages[0], device="cuda")[
-                    torch.arange(end, device="cuda") // block
-                ]
-                >= 0
-            )
-            history = low_latent[:end].clone()
-            history[high_mask] = latent[:end][high_mask]
-            history[begin:] = latent[begin:end]
-            keys = torch.nn.functional.linear(history[:, :rank], key_weight).view(end, heads, nope)
-            keys = torch.cat((keys, history[:, None, rank:].expand(-1, heads, -1)), dim=-1)
-            values = torch.nn.functional.linear(history[:, :rank], value_weight).view(
-                end, heads, nope
-            )
-            mask = torch.ones(end - begin, end, dtype=torch.bool, device="cuda").tril(begin)
-            expected = (
-                torch.nn.functional.scaled_dot_product_attention(
-                    query[begin:end].transpose(0, 1).float(),
-                    keys.transpose(0, 1).float(),
-                    values.transpose(0, 1).float(),
-                    attn_mask=mask,
-                    scale=model.softmax_scale,
-                )
-                .transpose(0, 1)
-                .reshape(end - begin, -1)
-            )
-            output = torch.empty(end - begin, heads * nope, dtype=hidden.dtype, device="cuda")
-            model.forward_impl(
-                torch.arange(begin, end, device="cuda"), hidden[begin:end], metadata, [output]
-            )
-            assert metadata.kv_cache_manager is manager
-            # BF16 MLA with a 512-wide latent accumulates more rounding than the
-            # dense cases; the TRTLLM kernel lands within 0.03 of the reference.
-            torch.testing.assert_close(output.float(), expected, atol=0.03, rtol=0.03)
-    finally:
-        manager.shutdown()
-
-
 @pytest.mark.parametrize("layout", ("HND", "NHD"))
 @pytest.mark.parametrize(
     "dtype,width", [(torch.float16, 64), (torch.bfloat16, 80), (torch.float8_e4m3fn, 96)]