diff --git a/docs/source/deployment-guide/deployment-guide-for-minimax-m3-on-trtllm.md b/docs/source/deployment-guide/deployment-guide-for-minimax-m3-on-trtllm.md index ef58d5ed3b83..bc23d6f39777 100644 --- a/docs/source/deployment-guide/deployment-guide-for-minimax-m3-on-trtllm.md +++ b/docs/source/deployment-guide/deployment-guide-for-minimax-m3-on-trtllm.md @@ -65,6 +65,11 @@ server (see [Launch the TensorRT LLM Server](#launch-the-tensorrt-llm-server)). * **Weight and KV precision.** BF16, MXFP8, and mixed MXFP8/NVFP4 checkpoints are supported. On SM100/SM103, the MSA backend supports FP8 and NVFP4 KV cache. NVFP4 storage is used in sparse target layers; dense target layers and the shared Eagle3 draft layer use FP8. The curated BF16-KV YAML remains a separate starting configuration. * **KV cache reuse must be disabled.** KV cache reuse is not supported on the sparse-attention path, so set `kv_cache_config.enable_block_reuse: false`. * **Eagle3.** The MSA path supports one shared GQA draft layer with linear Eagle3 drafting. Other MTP and tree-drafting configurations are unsupported. Keep NVFP4 KV and piecewise CUDA graph qualification separate: the existing PCG accuracy cases use FP8 KV. +* **Rubin (SM107) Eagle3 draft layer.** On SM107 the shared Eagle3 draft layer runs at 128-token pages through FlashInfer's TRTLLM-Gen FMHA, the only FMHA with kernels for that page size there: keep `flashinfer_trtllm_gen` enabled in `TLLM_FMHA_LIBS`. `flashinfer-python` downloads those kernels on first use. On nodes without internet access, or to keep the download off the serving path, pre-populate the cubin directory and serve offline: + 1. On one node with internet access, start the server with the production configuration (same tensor-parallel size and KV cache dtype: the kernels depend on the per-rank head layout and KV dtype), `FLASHINFER_CUBIN_DIR` set to an empty directory, and serve one Eagle3 request with a long prompt (a context and a generation step at 128-token pages). + 2. Copy that directory to each serving node, set `FLASHINFER_CUBIN_DIR` to it and `FLASHINFER_NO_DOWNLOAD=1`; see [FlashInfer JIT workspaces for MPI workers](../llm-api/index.md#flashinfer-jit-workspaces-for-mpi-workers). + + For the FP8 KV draft layer at TP2 (16 query heads and 1 KV head per rank, head size 128) the directory holds, under `/fmha/trtllm-gen/`, `include/flashInferMetaInfo.h`, `checksums.txt` and 10 cubins: the `fmhaSm107aKernel_QkvE4m3OBfloat16H128PagedKvCausalPaged` context kernel `VarSeqQ256Kv128PersistentContext` and three generation kernels (`MultiCtasKvCgaVarSeqQ64Kv256StaticGroupedKeepsAbForGen`, `MultiCtasKvVarSeqQ64Kv256StaticGroupedKeepsAbForGen`, `VarSeqQ64Kv256PersistentGroupedKeepsAbForGen`), plus six `fmhaSm100fKernel_QkvE4m3OBfloat16H128PagedKvCausalPaged` generation kernels (`Q16` and `Q32` variants). * **Use matching builds for disaggregated context and generation workers.** The Triton unquantized main KV cache uses HND physical storage. Builds using the previous NHD storage are not compatible for KV transfer; upgrade both workers together. * **Size `max_seq_len` for the active attention paths.** On SM100/SM103, supported dense layers with unquantized KV cache in the Triton configuration use native paged attention and do not gather historical KV or expand GQA heads into temporary tensors. The SDPA dense fallback still allocates those buffers when TRTLLM-Gen is unavailable, disabled, ordered after the C++ fallback in `TLLM_FMHA_LIBS`, or the Triton configuration uses FP8/NVFP4 KV. FP8 KV retains SDPA to preserve its existing Q precision. The Triton sparse decode path separately has per-query FP32 temporary buffers that grow with the attended KV length. Keep `max_seq_len` bounded for paths with these allocations: checkpoint defaults (1M for MXFP8, 512K for BF16) can cause OOM or CUDA graph capture failures. The curated YAML sets `max_seq_len` just above ISL+OSL (`2068` for the 1k/1k benchmark); longer-context workloads may require a lower `max_batch_size`. MSA retains its own dense/sparse kernels and cache layout. * **Parallelism.** MoE experts run with expert parallelism. The attention layers support both Tensor-Expert Parallelism (TEP) and Data-Expert Parallelism (DEP, via `enable_attention_dp: true`). The overlap scheduler is enabled by default. diff --git a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py index 3fa2ce645cad..e465eda7d093 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py @@ -369,6 +369,10 @@ class FlashInferTrtllmGenFmha(PhasedFmha): MAX_HEADS_RATIO_GENERATION = 32 MIN_TOKENS_PER_BLOCK = 8 SUPPORTED_TOKENS_PER_BLOCK = {16, 32, 64} + # SMs on which this library serves only a KV cache manager's extra page + # sizes (``trtllm_gen_extra_tokens_per_block``); every other request there + # stays on the fallback FMHA. + EXTRA_PAGE_SIZES_ONLY_SMS = frozenset({107}) # FlashInfer narrows key_cache.size(0) to int before constructing its TMA descriptor. MAX_NUM_PAGES_IN_MEM_POOL = (1 << 31) - 1 SUPPORTED_MLA_GENERATION_HEAD_DIMS = { @@ -407,14 +411,20 @@ def is_environment_available(cls) -> bool: return False sm = get_sm_version() - if sm not in (100, 103): + if sm not in (100, 103, 107): logger.debug( - f"FlashInfer TRTLLM-Gen FMHA is unavailable: requires SM100 or SM103, got SM{sm}." + "FlashInfer TRTLLM-Gen FMHA is unavailable: " + f"requires SM100, SM103 or SM107, got SM{sm}." ) return False return True + @classmethod + def serves_only_extra_page_sizes(cls) -> bool: + """Whether this device uses the library only for a KV cache manager's extra page sizes.""" + return get_sm_version() in cls.EXTRA_PAGE_SIZES_ONLY_SMS + @classmethod def _is_available(cls, attn: "TrtllmAttention") -> bool: if not cls.is_environment_available(): @@ -790,12 +800,21 @@ def _is_supported_with_reason( if tokens_per_block & (tokens_per_block - 1) != 0: return False, f"tokens_per_block ({tokens_per_block}) that is not a power of 2." # A KV cache manager may allow extra page sizes, e.g. MiniMax-M3 adds 128. - supported_tokens_per_block = self.SUPPORTED_TOKENS_PER_BLOCK | set( + extra_tokens_per_block = frozenset( getattr(meta.kv_cache_manager, "trtllm_gen_extra_tokens_per_block", ()) ) + supported_tokens_per_block = self.SUPPORTED_TOKENS_PER_BLOCK | extra_tokens_per_block if tokens_per_block not in supported_tokens_per_block: supported = sorted(supported_tokens_per_block) return False, f"tokens_per_block ({tokens_per_block}). Supported: {supported}." + # On SM107 this library serves only the extra page sizes, which the + # fallback thop FMHA has no paged-context kernels for. Every other SM107 + # request stays on the fallback. + if self.serves_only_extra_page_sizes() and tokens_per_block not in extra_tokens_per_block: + return False, ( + f"tokens_per_block ({tokens_per_block}) on SM107, which uses this library " + "only for the KV cache manager's extra page sizes." + ) return True, "" diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py index 37d9564664a2..c52c6d665126 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py @@ -38,6 +38,7 @@ TensorWrapper, binding_to_torch_dtype, convert_to_torch_tensor, + get_sm_version, prefer_pinned, ) from tensorrt_llm.bindings import DataType @@ -227,6 +228,28 @@ def extend_attention_op_pools_for_shared_draft_layers( ) +def _require_trtllm_gen_extra_page_route(tokens_per_block: int) -> None: + """Fail at setup, not at the first prefill, when the P128 opt-in has no FMHA (SM107). + + On SM107 only FlashInfer TRTLLM-Gen has paged-context kernels for the extra + page sizes (its cubins ship with it), so the route exists iff it is enabled. + """ + if get_sm_version() != 107: + return + # Imported here: the FMHA backends import the attention classes this module serves. + from tensorrt_llm._torch.attention.backends.fmha.flashinfer_trtllm_gen import ( + FlashInferTrtllmGenFmha, + ) + from tensorrt_llm._torch.attention.backends.fmha.registry import get_enabled_fmha_lib_classes + + if FlashInferTrtllmGenFmha not in get_enabled_fmha_lib_classes(): + raise RuntimeError( + f"SM107 with {tokens_per_block}-token KV pages needs the flashinfer_trtllm_gen FMHA " + "library, the only one with paged-context kernels for this page size here, but " + "TLLM_FMHA_LIBS excludes it. Keep flashinfer_trtllm_gen in TLLM_FMHA_LIBS (or unset it)." + ) + + class MiniMaxM3KVCacheManagerV2(KVCacheManagerV2): """KVCacheManagerV2 subclass with a V2-managed paged index-K cache per sparse layer. @@ -399,6 +422,18 @@ def __init__( self._block_scale_buffers[layer_idx, kv_layout] = self.get_block_scale_buffers( layer_idx, kv_layout ) + # The KV pools are written only at live slots, and readers that consume a whole block or page (the CSR + # sparse prefill's V dequant, the full-context decode on a last page's tail) would turn uninitialized + # bytes into 0 * NaN = NaN. Zero every layer's K/V bytes (NVFP4 sparse, FP8 dense, the drafter's in its + # own manager), the NVFP4 block scales and the paged index keys once, before any request runs: every + # byte a kernel writes is unchanged. + for view in ( + *self._kv_buffers.values(), + *self._block_scale_buffers.values(), + *self._index_k_buffers.values(), + ): + if view is not None: + view.zero_() def _validate_speculative_config(self, spec_config: DecodingBaseConfig | None) -> None: # Dynamic-tree relocation moves packed K/V bytes without their NVFP4 @@ -540,7 +575,16 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: for layer_idx in (*self._dense_native_layer_ids, *self._shared_draft_layer_ids) if layer_idx in self.layer_offsets ] - if self.is_draft or not native_layers: + if self.is_draft: + # A separate draft manager holds only draft layers, which all run + # native attention and use their own pool tables; it needs the P128 + # opt-in but no virtual pools. (MiniMax-M3 Eagle3 shares the target + # manager today; this keeps a separate one on the trtllm-gen route.) + if self.tokens_per_block == 128 and self.dtype != DataType.NVFP4: + _require_trtllm_gen_extra_page_route(128) + self.trtllm_gen_extra_tokens_per_block = frozenset({128}) + return + if not native_layers: return if self.enable_swa_scratch_reuse: raise NotImplementedError( @@ -552,8 +596,13 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: return # Native layers run at the target's 128-token pages. trtllm-gen has P128 # kernels for their dense-GQA shapes but not for every shape, so opt in - # here rather than in the global allowlist. + # here rather than in the global allowlist. On SM107 this opt-in is the + # only route for these layers: FallbackFmha has no P128 paged-context kernels. if self.tokens_per_block == 128: + # On SM107 the dense layers keep SDPA (see the model's native dense attention + # setup), so only the shared draft layers need the route there. + if any(i in self.layer_offsets for i in self._shared_draft_layer_ids): + _require_trtllm_gen_extra_page_route(128) self.trtllm_gen_extra_tokens_per_block = frozenset({128}) if self._use_per_layer_page_tables: # Per-layer page tables already give every layer its own pool. @@ -1031,6 +1080,8 @@ def __init__(self, manager: MiniMaxM3KVCacheManagerV2, draft_layer_ids: Sequence "MiniMax-M3 draft KV cache views support exactly one local shared draft layer; " f"got {len(draft_layer_ids)}. Multiple layers require separate pool roots." ) + if manager.tokens_per_block in self.trtllm_gen_extra_tokens_per_block: + _require_trtllm_gen_extra_page_route(manager.tokens_per_block) self._manager = manager layer_id = draft_layer_ids[0] local = int(manager.layer_offsets[layer_id]) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/triton_sparse_decode.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/triton_sparse_decode.py index 071f907067b4..0150347026a1 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/triton_sparse_decode.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/triton_sparse_decode.py @@ -799,8 +799,8 @@ def resolve_num_topk_chunks(total_q: int, num_kv_heads: int, max_topk: int) -> i return 1 << (target.bit_length() - 1) -# NVFP4 instructions are supported on both datacenter Blackwell targets. -_SM100F_CAPABILITIES = ((10, 0), (10, 3)) +# NVFP4 instructions are supported on the datacenter Blackwell and Rubin targets. +_SM100F_CAPABILITIES = ((10, 0), (10, 3), (10, 7)) # Shape-specific tuning was measured on GB300 only. SM100 uses the defaults. _NVFP4_TUNED_CAPABILITIES = ((10, 3),) @@ -1063,7 +1063,7 @@ def minimax_m3_sparse_attn_decode( capability = torch.cuda.get_device_capability(q.device) if capability not in _SM100F_CAPABILITIES: raise NotImplementedError( - "MiniMax-M3 NVFP4 sparse decode requires SM100/SM103; " + "MiniMax-M3 NVFP4 sparse decode requires SM100/SM103/SM107; " f"got SM{capability[0] * 10 + capability[1]}." ) _check_nvfp4_inputs(k_paged, v_paged, nvfp4_args, head_dim=head_dim, scale_cols=scale_cols) diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index 4359a6e4549b..adaaaf93682d 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -30,6 +30,7 @@ from torch.nn.attention import SDPBackend, sdpa_kernel from transformers import PretrainedConfig +from tensorrt_llm._utils import get_sm_version from tensorrt_llm.bindings import DataType from tensorrt_llm.functional import AllReduceStrategy, PositionEmbeddingType from tensorrt_llm.logger import logger @@ -1299,6 +1300,11 @@ def _init_native_dense_attention(self, model_config: ModelConfig) -> None: return if not FlashInferTrtllmGenFmha.is_environment_available(): return + if get_sm_version() == 107: + # SM107 admits the library only for the KV manager's extra page sizes + # (the Eagle3 draft layers at 128-token pages). Dense layers keep SDPA + # there, as before SM107 was admitted. + return # Q/K norm and partial RoPE run in the model before backend dispatch. native_attn = TrtllmAttention( self.layer_idx, diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 9f2f4e12a855..5b6a0514cbd3 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -56,7 +56,7 @@ maybe_create_moe_load_balancer) from ..virtual_memory import RestoreMode from ..virtual_memory import scope as virtual_memory_scope -from .config_utils import (is_hybrid_linear, is_mla, +from .config_utils import (is_hybrid_linear, is_minimax_m3, is_mla, resolve_auto_ssm_cache_dtype, supports_fp4_mla_attention, uses_fp4_mla_attention, validate_kimi_kda_state_dtype) @@ -165,9 +165,15 @@ def validate_and_set_kv_cache_quant(model_config: ModelConfig, effective_kv_cache_quant = (kv_cache_quant if pyt_kv_cache_dtype == "auto" else mapped_pyt_quant) + # MiniMax-M3 MSA keeps NVFP4 only on its sparse layers, which MSA and + # Triton serve; its dense and draft layers stay FP8 for trtllm-gen. + sparse_config = model_config.sparse_attention_config + uses_minimax_m3_msa = (is_minimax_m3(sparse_config) + and sparse_config.implementation == "msa") if (effective_kv_cache_quant in (QuantAlgo.NVFP4, QuantAlgo.NVFP4.value) and not supports_fp4_mla_attention(model_config) - and torch.cuda.is_available() and get_sm_version() == 107): + and not uses_minimax_m3_msa and torch.cuda.is_available() + and get_sm_version() == 107): logger.warning( "NVFP4 KV cache is not supported by trtllm-gen on SM107; " "using FP8 KV cache instead.") diff --git a/tests/integration/test_lists/test-db/l0_gr100.yml b/tests/integration/test_lists/test-db/l0_gr100.yml index 45cd6e6a35a8..104093738356 100644 --- a/tests/integration/test_lists/test-db/l0_gr100.yml +++ b/tests/integration/test_lists/test-db/l0_gr100.yml @@ -14,3 +14,4 @@ l0_gr100: backend: pytorch tests: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_flashinfer[enable_chunked_prefill=True] + - unittest/_torch/attention/test_flashinfer_trtllm_gen_sm107_p128.py diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_kv_pool_zero_init.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_kv_pool_zero_init.py new file mode 100644 index 000000000000..988df58473f0 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_kv_pool_zero_init.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""MiniMax-M3 KV pools start zeroed. + +Readers that consume a whole block or page (the CSR sparse prefill's V dequant, +the full-context decode on a last page's tail) also read bytes no request has +written, so every pool view must hold zeros, not uninitialized memory, before +the first request. +""" + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.cache_manager import ( + MiniMaxM3KVCacheManagerV2, +) +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.internal.batch_manager import CacheType +from tensorrt_llm.llmapi.llm_args import ( + Eagle3DecodingConfig, + KvCacheConfig, + MiniMaxM3SparseAttentionConfig, +) +from tensorrt_llm.mapping import Mapping + + +def _nonzero_bytes(view: torch.Tensor) -> int: + """Count nonzero bytes; a strided view that cannot be reinterpreted in place is copied first.""" + try: + raw = view.view(torch.uint8) + except RuntimeError: + raw = view.contiguous().view(torch.uint8) + return int(torch.count_nonzero(raw).item()) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA cache pools") +@pytest.mark.parametrize("with_draft", [False, True], ids=["target_only", "shared_eagle3"]) +def test_every_pool_view_is_zero_after_construction(with_draft: bool) -> None: + torch.cuda.init() + manager = MiniMaxM3KVCacheManagerV2( + KvCacheConfig( + max_gpu_total_bytes=16 << 20, + host_cache_size=0, + enable_block_reuse=False, + enable_swa_scratch_reuse=False, + ), + CacheType.SELF, + num_layers=4, + num_kv_heads=[2, 2, 2, 2, 1] if with_draft else [2, 2, 2, 2], + head_dim=128, + tokens_per_block=128, + max_seq_len=256, + max_batch_size=2, + max_num_tokens=256, + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + dtype=DataType.NVFP4, + vocab_size=16, + spec_config=( + Eagle3DecodingConfig(max_draft_len=3, speculative_model="unused") + if with_draft + else None + ), + sparse_attention_config=MiniMaxM3SparseAttentionConfig( + implementation="msa", indexer_kv_dtype="fp8" + ), + ) + try: + views = { + **{("kv",) + key: view for key, view in manager._kv_buffers.items()}, + **{("scale",) + key: view for key, view in manager._block_scale_buffers.items()}, + **{("index_k", key): view for key, view in manager._index_k_buffers.items()}, + } + assert any(name[0] == "kv" for name in views) + assert any(name[0] == "scale" for name in views) + assert any(name[0] == "index_k" and view is not None for name, view in views.items()) + torch.cuda.synchronize() + counts = {name: _nonzero_bytes(view) for name, view in views.items() if view is not None} + nonzero = {name: count for name, count in counts.items() if count} + assert not nonzero, f"pool views with nonzero bytes after construction: {nonzero}" + finally: + manager.shutdown() diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py index e9abf2039073..0225dea33702 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py @@ -54,6 +54,7 @@ from tensorrt_llm.runtime.kv_cache_manager_v2 import PageIndexMode DRAFT_LOCAL_LAYER = 60 +_CACHE_MANAGER_MODULE = "tensorrt_llm._torch.attention.backends.sparse.minimax_m3.cache_manager" SCALE = 179 # sub-pages per M3 mega-slot: 3 dense x 2 + 57 sparse x 3 + draft x 2 DRAFT_K_ADDR = 0x7000_0000 @@ -301,6 +302,45 @@ def fake_base_prepare(self, index_mapper_capacity): assert manager.trtllm_gen_extra_tokens_per_block == frozenset(expected_extra_pages) +@pytest.mark.parametrize( + "dtype,tokens_per_block,expected", + [ + (DataType.FP8, 128, {128}), + (DataType.BF16, 128, {128}), + (DataType.FP8, 32, set()), + (DataType.NVFP4, 128, set()), + ], +) +def test_separate_draft_manager_opts_into_p128(monkeypatch, dtype, tokens_per_block, expected): + """A separate (is_draft) manager holds only native draft layers: at 128-token pages + it lists 128 as a trtllm-gen extra page size (the only P128 route on SM107) and + builds no virtual pools.""" + pointers = torch.tensor([[0x6000_0000, 0]]) + mapping = torch.tensor([[0, 0]], dtype=torch.int32) + + def fake_base_prepare(self, index_mapper_capacity): + self._use_per_layer_page_tables = False + self.kv_cache_pool_pointers = pointers.clone() + self.kv_cache_pool_mapping = mapping.clone() + self.num_attention_op_pools = 1 + + monkeypatch.setattr(KVCacheManagerV2, "_prepare_page_table_tensor", fake_base_prepare) + manager = MiniMaxM3KVCacheManagerV2.__new__(MiniMaxM3KVCacheManagerV2) + manager.dtype = dtype + manager._shared_draft_layer_ids = [] + manager.layer_offsets = {0: 0} + manager.is_draft = True + manager.enable_swa_scratch_reuse = False + manager.tokens_per_block = tokens_per_block + + manager._prepare_page_table_tensor(8) + + assert manager.trtllm_gen_extra_tokens_per_block == frozenset(expected) + assert manager._draft_op_pools == () + assert torch.equal(manager.kv_cache_pool_pointers, pointers) + assert manager.num_attention_op_pools == (0 if dtype == DataType.NVFP4 else 1) + + def test_update_resources_refuses_tree_relocation(monkeypatch): """Linear acceptance rewinds through the base; tree acceptance is refused.""" calls = [] @@ -377,6 +417,91 @@ def test_draft_kv_cache_view_accessor_respects_local_layers( assert manager._shared_draft_layer_ids == [60, 61] +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "sm,fmha_libs,error", + [ + (107, None, None), + (107, "-flashinfer_trtllm_gen", "TLLM_FMHA_LIBS excludes it"), + (107, "fallback", "TLLM_FMHA_LIBS excludes it"), + (107, "flashinfer_trtllm_gen,fallback", None), + # Elsewhere the fallback has P128 paged-context kernels, so nothing is required. + (103, "-flashinfer_trtllm_gen", None), + ], +) +@pytest.mark.parametrize("site", ["separate_draft_manager", "nvfp4_draft_view"]) +def test_p128_opt_in_fails_fast_without_the_flashinfer_route( + monkeypatch, sm, fmha_libs, error, site +): + """On SM107 the P128 opt-in is checked when the KV manager is set up, not at the + first prefill: a missing FlashInfer route raises an error that names the fix.""" + monkeypatch.setattr(f"{_CACHE_MANAGER_MODULE}.get_sm_version", lambda: sm) + if fmha_libs is None: + monkeypatch.delenv("TLLM_FMHA_LIBS", raising=False) + else: + monkeypatch.setenv("TLLM_FMHA_LIBS", fmha_libs) + + if site == "separate_draft_manager": + monkeypatch.setattr( + KVCacheManagerV2, + "_prepare_page_table_tensor", + lambda self, capacity: setattr(self, "_use_per_layer_page_tables", False), + ) + manager = MiniMaxM3KVCacheManagerV2.__new__(MiniMaxM3KVCacheManagerV2) + manager.dtype = DataType.FP8 + manager._shared_draft_layer_ids = [] + manager.layer_offsets = {0: 0} + manager.is_draft = True + manager.enable_swa_scratch_reuse = False + manager.tokens_per_block = 128 + setup = lambda: manager._prepare_page_table_tensor(8) # noqa: E731 + else: + # Stop right after the check: the rest of the view needs real pools. + manager = SimpleNamespace(tokens_per_block=128, layer_offsets={}) + setup = lambda: MiniMaxM3DraftKVCacheView(manager, [DRAFT_LOCAL_LAYER]) # noqa: E731 + + if error is not None: + with pytest.raises(RuntimeError, match=error) as excinfo: + setup() + assert "128-token KV pages" in str(excinfo.value) + elif site == "separate_draft_manager": + setup() + assert manager.trtllm_gen_extra_tokens_per_block == frozenset({128}) + else: + # The check passed; the view then fails on the stub's missing layer. + with pytest.raises(KeyError, match=str(DRAFT_LOCAL_LAYER)): + setup() + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("shared_draft", [False, True]) +def test_p128_target_fails_fast_only_with_a_shared_draft_layer(monkeypatch, shared_draft): + """On SM107 a BF16 target's dense layers keep SDPA, so a missing FlashInfer route only + fails setup when a shared draft layer needs it.""" + monkeypatch.setattr(f"{_CACHE_MANAGER_MODULE}.get_sm_version", lambda: 107) + monkeypatch.setenv("TLLM_FMHA_LIBS", "-flashinfer_trtllm_gen") + monkeypatch.setattr( + KVCacheManagerV2, + "_prepare_page_table_tensor", + lambda self, capacity: setattr(self, "_use_per_layer_page_tables", True), + ) + manager = MiniMaxM3KVCacheManagerV2.__new__(MiniMaxM3KVCacheManagerV2) + manager.dtype = DataType.HALF + manager.is_draft = False + manager.enable_swa_scratch_reuse = False + manager.tokens_per_block = 128 + manager._dense_native_layer_ids = (0, 1, 2) + manager._shared_draft_layer_ids = [60] if shared_draft else [] + manager.layer_offsets = {0: 0, 1: 1, 2: 2, **({60: 3} if shared_draft else {})} + + if shared_draft: + with pytest.raises(RuntimeError, match="TLLM_FMHA_LIBS excludes it"): + manager._prepare_page_table_tensor(8) + else: + manager._prepare_page_table_tensor(8) + assert manager.trtllm_gen_extra_tokens_per_block == frozenset({128}) + + @pytest.mark.cpu_only @pytest.mark.parametrize("via_accessor", [False, True]) def test_draft_kv_cache_view_rejects_multiple_local_layers(via_accessor: bool) -> None: diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py index 88760a79452f..07376d86d2d7 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py @@ -1158,7 +1158,7 @@ def test_resolve_num_topk_chunks_is_shape_only_power_of_two(total_q, num_kv_head (33, {}), ], ) -@pytest.mark.parametrize("capability", [(10, 0), (10, 3)]) +@pytest.mark.parametrize("capability", [(10, 0), (10, 3), (10, 7)]) def test_sm100f_nvfp4_launch_options(local_batch, expected, capability): if capability != (10, 3): expected = {} @@ -1192,7 +1192,7 @@ def test_sm100f_nvfp4_launch_options(local_batch, expected, capability): (15, None), ], ) -@pytest.mark.parametrize("capability", [(10, 0), (10, 3)]) +@pytest.mark.parametrize("capability", [(10, 0), (10, 3), (10, 7)]) def test_sm100f_nvfp4_num_topk_chunks(local_batch, expected, capability): if capability != (10, 3): expected = None @@ -1213,7 +1213,7 @@ def test_sm100f_nvfp4_num_topk_chunks(local_batch, expected, capability): ("local_batch", "expected"), [(10, 1), (11, 2), (12, 2), (14, 2), (15, 1)], ) -@pytest.mark.parametrize("capability", [(10, 0), (10, 3)]) +@pytest.mark.parametrize("capability", [(10, 0), (10, 3), (10, 7)]) def test_sm100f_nvfp4_query_group_size_is_narrowly_scoped(local_batch, expected, capability): if capability != (10, 3): expected = 1 @@ -1259,7 +1259,7 @@ def test_sm100f_nvfp4_launch_options_rejects_unmeasured_shapes(): (8, 14, False), ], ) -@pytest.mark.parametrize("capability", [(10, 0), (10, 3)]) +@pytest.mark.parametrize("capability", [(10, 0), (10, 3), (10, 7)]) def test_sm100f_nvfp4_linear_softmax_policy(num_kv_heads, local_batch, expected, capability): if capability != (10, 3): expected = False @@ -1290,7 +1290,7 @@ def test_sm100f_nvfp4_linear_softmax_policy(num_kv_heads, local_batch, expected, (12, {}), ], ) -@pytest.mark.parametrize("capability", [(10, 0), (10, 3)]) +@pytest.mark.parametrize("capability", [(10, 0), (10, 3), (10, 7)]) def test_sm100f_nvfp4_merge_launch_options_is_narrowly_scoped(local_batch, expected, capability): if capability != (10, 3): expected = {} diff --git a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py index 94f139c52cfd..6f611923aa2c 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py @@ -258,7 +258,7 @@ def test_nvfp4_sparse_decode_rejects_unsupported_device( block_scale = torch.empty((1, 1, 128, 8), dtype=torch.uint8) global_scale = torch.ones(1) - with pytest.raises(NotImplementedError, match="NVFP4 sparse decode requires SM100/SM103"): + with pytest.raises(NotImplementedError, match="NVFP4 sparse decode requires SM100/SM103/SM107"): minimax_m3_sparse_attn_decode( q, kv, @@ -277,6 +277,34 @@ def test_nvfp4_sparse_decode_rejects_unsupported_device( get_capability.assert_called_once_with(q.device) +@pytest.mark.cpu_only +@pytest.mark.parametrize("capability", [(10, 0), (10, 3), (10, 7)]) +def test_nvfp4_sparse_decode_admits_the_sm100_family( + monkeypatch: pytest.MonkeyPatch, capability: tuple[int, int] +) -> None: + """SM100-family devices, Rubin included, pass the device check and reach + input validation, which here rejects the missing scales.""" + get_capability = Mock(spec_set=torch.cuda.get_device_capability, return_value=capability) + monkeypatch.setattr(torch.cuda, "get_device_capability", get_capability) + q = torch.empty((1, 8, 128), dtype=torch.bfloat16) + kv = torch.empty((1, 1, 128, 64), dtype=torch.uint8) + + with pytest.raises(ValueError, match="needs all four"): + minimax_m3_sparse_attn_decode( + q, + kv, + kv, + torch.zeros((1, 1, 1), dtype=torch.int32), + torch.zeros((1, 1), dtype=torch.int32), + torch.ones(1, dtype=torch.int32), + sm_scale=128**-0.5, + output=torch.empty_like(q), + decode_query_len=1, + k_block_scale=torch.empty((1, 1, 128, 8), dtype=torch.uint8), + ) + get_capability.assert_called_once_with(q.device) + + @pytest.mark.cpu_only @pytest.mark.parametrize("indexer_kv_dtype", ["bf16", "fp8"]) @pytest.mark.parametrize("implementation", ["msa", "triton"]) diff --git a/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_fmha.py b/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_fmha.py index e4642e1beeb1..fac6b0eb5359 100644 --- a/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_fmha.py +++ b/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_fmha.py @@ -34,6 +34,8 @@ from tensorrt_llm.bindings import DataType from tensorrt_llm.quantization.mode import QuantMode +_FLASHINFER_TRTLLM_GEN_MODULE = "tensorrt_llm._torch.attention.backends.fmha.flashinfer_trtllm_gen" + def test_flashinfer_fp8_mode_remains_implementation_local() -> None: attn = FakeAttention() @@ -64,6 +66,7 @@ def test_flashinfer_fp8_mode_remains_implementation_local() -> None: @pytest.mark.cpu_only @pytest.mark.parametrize("page_size,paged_context", [(32, False), (128, True)]) def test_small_context_fallback_preserves_mixed_batch_generation( + monkeypatch: pytest.MonkeyPatch, num_contexts: int, num_generations: int, dtype: torch.dtype, @@ -71,6 +74,7 @@ def test_small_context_fallback_preserves_mixed_batch_generation( page_size: int, paged_context: bool, ) -> None: + monkeypatch.setattr(f"{_FLASHINFER_TRTLLM_GEN_MODULE}.get_sm_version", lambda: 100) attn = FakeAttention() attn.head_dim = 256 attn.sparse_params = None @@ -202,3 +206,105 @@ def test_flashinfer_quantized_kv_context_avoids_fp16_bf16_fallback( else: assert supported, reason assert reason == "" + + +@pytest.mark.parametrize( + "sm_version,expected", + [(90, False), (100, True), (103, True), (107, True), (120, False)], +) +def test_flashinfer_availability_by_sm( + monkeypatch: pytest.MonkeyPatch, + sm_version: int, + expected: bool, +) -> None: + monkeypatch.setattr(f"{_FLASHINFER_TRTLLM_GEN_MODULE}.IS_FLASHINFER_AVAILABLE", True) + monkeypatch.setattr( + FlashInferTrtllmGenFmha, "_missing_fused_nanobind_ops", staticmethod(lambda: []) + ) + monkeypatch.setattr(f"{_FLASHINFER_TRTLLM_GEN_MODULE}.get_sm_version", lambda: sm_version) + attn = FakeAttention() + attn.sparse_params = None + + assert FlashInferTrtllmGenFmha._is_available(attn) is expected + + +@pytest.mark.parametrize( + "input_type", + [AttentionInputType.context_only, AttentionInputType.generation_only], + ids=["context", "generation"], +) +@pytest.mark.parametrize( + "sm_version,tokens_per_block,extra_tokens_per_block,expected", + [ + (107, 128, frozenset({128}), True), + (107, 32, frozenset({128}), False), + (107, 32, frozenset(), False), + (103, 32, frozenset(), True), + (103, 128, frozenset({128}), True), + ], + ids=[ + "sm107_extra_page_size", + "sm107_default_page_size_with_extra", + "sm107_default_page_size", + "sm103_default_page_size", + "sm103_extra_page_size", + ], +) +def test_sm107_serves_only_extra_page_sizes( + monkeypatch: pytest.MonkeyPatch, + input_type: AttentionInputType, + sm_version: int, + tokens_per_block: int, + extra_tokens_per_block: frozenset, + expected: bool, +) -> None: + """FP8-KV GQA 16:1 head-128 paged attention, the MiniMax-M3 Eagle3 draft layer shape.""" + monkeypatch.setattr(f"{_FLASHINFER_TRTLLM_GEN_MODULE}.get_sm_version", lambda: sm_version) + fmha = object.__new__(FlashInferTrtllmGenFmha) + fmha.kv_factor = 2 + monkeypatch.setattr(fmha, "_get_total_num_blocks", lambda _: 0) + attn = SimpleNamespace( + is_mla_enable=False, + sparse_params=None, + position_embedding_type=0, + head_dim=128, + num_heads=16, + num_kv_heads=1, + ) + is_context = input_type == AttentionInputType.context_only + num_tokens = 4 + q_hidden_size = attn.num_heads * attn.head_dim + qkv_hidden_size = q_hidden_size + 2 * attn.num_kv_heads * attn.head_dim + q = torch.empty((num_tokens, qkv_hidden_size), dtype=torch.bfloat16) + metadata = SimpleNamespace( + num_contexts=1 if is_context else 0, + num_generations=0 if is_context else 1, + helix_position_offsets=None, + num_sparse_topk=0, + use_spec_decoding=False, + is_spec_dec_tree=False, + kv_cache_block_offsets=object(), + kv_cache_manager=SimpleNamespace( + dtype=DataType.FP8, + trtllm_gen_extra_tokens_per_block=extra_tokens_per_block, + ), + is_cross=False, + is_spec_decoding_enabled=False, + tokens_per_block=tokens_per_block, + beam_width=1, + ) + forward_args = AttentionForwardArgs( + output=torch.empty((num_tokens, q_hidden_size), dtype=torch.bfloat16), + attention_input_type=input_type, + is_fused_qkv=True, + update_kv_cache=True, + ) + + phase = FmhaPhase.CONTEXT if is_context else FmhaPhase.GENERATION + for requested_phase in (None, phase): + supported, reason = fmha._is_supported_with_reason( + q, None, None, attn, metadata, forward_args, phase=requested_phase + ) + assert supported is expected, reason + if not expected: + assert "SM107" in reason diff --git a/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_sm107_p128.py b/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_sm107_p128.py new file mode 100644 index 000000000000..f6b06575988b --- /dev/null +++ b/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_sm107_p128.py @@ -0,0 +1,375 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""FlashInfer TRTLLM-Gen FMHA at 128-token pages on SM107 (Rubin). + +On SM107 the FallbackFmha (thop) has no paged-context kernels for 128-token pages, +so the MiniMax-M3 KV managers list 128 as a TRTLLM-Gen extra page size and +FlashInferTrtllmGenFmha serves those steps (and only those). These tests run the +real forward at the MiniMax-M3 Eagle3 draft layer shape (16 q : 1 kv head, +head dim 128, BF16 Q / output) with BF16 and FP8 KV and check, for every step: + +- FlashInfer served it (run_context / run_generation) and FallbackFmha did not; +- the output is within a tolerance of an FP32 reference built from the pool + contents (FlashInfer is precision-equivalent to the thop P32 kernels, not + bitwise, so there is no bit-exact oracle); + +over fresh, cached-prefix and chunked contexts (each chunk also matches the +same rows run unchunked and leaves the same K / V), generation (q_len 1), +linear Eagle3 verify (q_len 1 + 3) and mixed context + generation batches. + +The masked tail of a request's last page is read by the kernel: NaN bytes there +turn the rows that attend to that page into NaN (measured on SM107 for both this +route and the thop kernels). MiniMax-M3 guards against it by zeroing its KV pools +at allocation (test_minimax_m3_kv_pool_zero_init.py); +the "nan" tail cases check that rows of requests without a tail stay exact. +""" + +import pytest +import torch + +import tensorrt_llm +from tensorrt_llm._torch.attention.backends.fmha import fallback +from tensorrt_llm._torch.attention.backends.fmha import flashinfer_trtllm_gen as fi_mod +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionRuntimeFeatures, + PredefinedAttentionMask, +) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention, TrtllmAttentionMetadata +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager +from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.bindings import DataType +from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +_ON_SM107 = torch.cuda.is_available() and get_sm_version() == 107 +pytestmark = pytest.mark.skipif( + not _ON_SM107, reason="FlashInfer serves 128-token pages only on SM107 in this route" +) + +# The MiniMax-M3 Eagle3 draft attention shape on one rank. +_HEADS, _KV_HEADS, _HEAD_DIM, _PAGE = 16, 1, 128, 128 +_REF_RID = 100 # the unchunked run of request r is request _REF_RID + r +_KV_DTYPES = [torch.bfloat16, torch.float8_e4m3fn] +_KV_IDS = ["bf16_kv", "fp8_kv"] + +# Worst row's max |out - ref| / max |ref|. "float" is vs the FP32 reference (FP8 KV context +# also quantizes Q to E4M3, which the reference mirrors; FP8 KV generation keeps BF16 Q); +# "unchunked" is a chunk vs the same rows of the unchunked run (same cached K / V): 0 means +# bitwise equal. Each bound is 1.25x the worst case of a seed sweep at this 16:1 shape +# (Hecate 826693: 32 seeds per context case, 16 seeds x {zero, NaN} tails per generation +# case), rounded up. Per-case worst (BF16 / FP8 KV): +# context rows: fresh 0.0048 / 0.045, cached prefixes 0.0053 / 0.047, chunked 0.0052 / 0.046, +# mixed batches 0.0061 / 0.049; +# generation rows: q_len 1 0.0054 / 0.094, verify 1+3 0.0057 / 0.116, mixed 0.0057 / 0.143. +# FP8 grows with the past (q_len 1: 0.083 at 640 -> 0.143 at 165000); BF16 does not; +# chunks vs unchunked: 0 for every seed and case. +_TOLERANCE = { + torch.bfloat16: {"ctx": 7.7e-3, "gen": 7.2e-3, "unchunked": 0.0}, + torch.float8_e4m3fn: {"ctx": 6.2e-2, "gen": 1.8e-1, "unchunked": 0.0}, +} +# A kernel that dropped the cached prefix would match attention over the chunk alone; in the +# sweep that error was >= 1.05, at least 4x above every "ctx" bound. +_DROPPED_PREFIX_MARGIN = 4 + +# (case id, steps): each step is a batch of (request id, cached tokens, new tokens) context rows, run in order. +_CACHED_CONTEXT_CASES = [ + ("fresh_512", [[(0, 0, 512)]]), + ("prefix_1000_unaligned", [[(0, 0, 1000)], [(0, 1000, 700)]]), + ("prefix_1024_aligned", [[(0, 0, 1024)], [(0, 1024, 300)]]), + ("chunked_6096_at_2048", [[(0, 0, 2048)], [(0, 2048, 2048)], [(0, 4096, 2000)]]), + ("prefix_beside_fresh", [[(0, 0, 384)], [(0, 384, 257), (1, 0, 129)]]), +] + +_PASTS = [2047, 16383, 165000, 640] +# (case id, context rows [(cached, new)], generation rows [(past, q_len)]); contexts come first, as in the runtime. +_GENERATION_CASES = [ + ("gen_q1", [], [(p, 1) for p in _PASTS]), + ("verify_1p3", [], [(p, 4) for p in _PASTS]), + ("mixed_q1", [(0, 2048), (16384, 2048)], [(640, 1), (16383, 1), (165000, 1)]), + ("mixed_1p3", [(0, 2048), (16384, 2048)], [(640, 4), (16383, 4), (165000, 4)]), +] + + +@pytest.fixture +def served(monkeypatch): + """Count the steps each FMHA library ran.""" + counts = {"fi_ctx": 0, "fi_gen": 0, "fallback": 0} + for cls, name, key in ( + (fi_mod.FlashInferTrtllmGenFmha, "run_context", "fi_ctx"), + (fi_mod.FlashInferTrtllmGenFmha, "run_generation", "fi_gen"), + (fallback.FallbackFmha, "forward", "fallback"), + ): + original = getattr(cls, name) + + def wrapped(self, *args, _original=original, _key=key, **kwargs): + counts[_key] += 1 + return _original(self, *args, **kwargs) + + monkeypatch.setattr(cls, name, wrapped) + return counts + + +def _attention(kv_dtype: torch.dtype) -> TrtllmAttention: + assert fi_mod.FlashInferTrtllmGenFmha.is_environment_available(), ( + "SM107 serves 128-token pages only through FlashInfer TRTLLM-Gen FMHA" + ) + fp8 = kv_dtype == torch.float8_e4m3fn + return TrtllmAttention( + layer_idx=0, + num_heads=_HEADS, + head_dim=_HEAD_DIM, + num_kv_heads=_KV_HEADS, + quant_config=QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8) if fp8 else None, + ) + + +def _kv_manager(kv_dtype: torch.dtype, num_tokens: int) -> KVCacheManager: + manager = KVCacheManager( + KvCacheConfig(max_tokens=num_tokens, enable_block_reuse=False), + tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, + num_layers=1, + num_kv_heads=_KV_HEADS, + head_dim=_HEAD_DIM, + tokens_per_block=_PAGE, + max_seq_len=num_tokens, + max_batch_size=8, + mapping=Mapping(world_size=1, tp_size=1, rank=0), + dtype=DataType.FP8 if kv_dtype == torch.float8_e4m3fn else DataType.BF16, + ) + # As MiniMax-M3's KV managers list it: this is what admits 128-token pages to FlashInfer on SM107. + manager.trtllm_gen_extra_tokens_per_block = frozenset({_PAGE}) + # Zero the pool, as MiniMax-M3's managers do at allocation (see the module docstring). + _raw_pool(manager).zero_() + return manager + + +def _raw_pool(manager) -> torch.Tensor: + """The layer's pool as bytes: [pages, K/V, kv heads, page tokens, head dim bytes].""" + return manager.get_buffers(0, kv_layout="HND").view(torch.uint8) + + +def _token_slots(manager, rid: int, length: int): + """Pool (page, offset) of tokens 0..length-1 of request rid.""" + pages = manager.get_batch_cache_indices([rid], 0)[0] + blocks = torch.tensor(pages, device="cuda", dtype=torch.long) + tokens = torch.arange(length, device="cuda") + return blocks[tokens // _PAGE], tokens % _PAGE + + +def _step(attn, manager, rows, qkv: torch.Tensor, num_contexts: int) -> torch.Tensor: + """One step as the serving runner prepares it. rows = [(request id, cached, new)], contexts first. + + A context row's prompt length is its current chunk, with the prefix in num_cached_tokens_per_seq; + a generation row's prompt length is its past. + """ + metadata = TrtllmAttentionMetadata( + max_num_requests=8, + max_num_tokens=max(8192, qkv.shape[0] + 8), + kv_cache_manager=manager, + runtime_features=AttentionRuntimeFeatures(chunked_prefill=True), + num_contexts=num_contexts, + seq_lens=torch.tensor([n for _, _, n in rows], dtype=torch.int), + request_ids=[rid for rid, _, _ in rows], + prompt_lens=[n if i < num_contexts else c for i, (_, c, n) in enumerate(rows)], + kv_cache_params=KVCacheParams( + use_cache=True, num_cached_tokens_per_seq=[c for _, c, _ in rows] + ), + ) + assert metadata.tokens_per_block == _PAGE + metadata.prepare() + out = attn.forward(qkv, None, None, metadata, attention_mask=PredefinedAttentionMask.CAUSAL) + return out[0] if isinstance(out, tuple) else out + + +def _reference_from_pool( + manager, kv_dtype, rid: int, q: torch.Tensor, cached: int, quantize_q: bool +) -> torch.Tensor: + """FP32 causal attention of rows cached..cached+n over the K / V the pool holds for tokens 0..cached+n.""" + n = q.shape[0] + pages, offsets = _token_slots(manager, rid, cached + n) + kv = _raw_pool(manager)[pages, :, :, offsets, :].view(kv_dtype).float() # [T, K/V, kv heads, D] + group = _HEADS // _KV_HEADS + k = kv[:, 0].repeat_interleave(group, dim=1) # [T, heads, D] + v = kv[:, 1].repeat_interleave(group, dim=1) + if quantize_q: + q = q.to(torch.float8_e4m3fn) + q = q.float().view(n, _HEADS, _HEAD_DIM) + scores = torch.einsum("nhd,thd->nht", q, k) * _HEAD_DIM**-0.5 + q_pos = cached + torch.arange(n, device="cuda") + k_pos = torch.arange(cached + n, device="cuda") + scores = scores.masked_fill(k_pos[None, None, :] > q_pos[:, None, None], float("-inf")) + return torch.einsum("nht,thd->nhd", scores.softmax(-1), v).reshape(n, _HEADS * _HEAD_DIM) + + +def _float_reference(qkv: torch.Tensor, start: int, n: int, kv_from: int, fp8: bool): + """Causal attention of rows start..start+n over keys kv_from..start+n of one request's QKV, in FP32.""" + q_size, kv_size = _HEADS * _HEAD_DIM, _KV_HEADS * _HEAD_DIM + rows = qkv[kv_from : start + n] + if fp8: + rows = rows.to(torch.float8_e4m3fn) + rows = rows.float() + q = rows[start - kv_from :, :q_size].view(n, _HEADS, _HEAD_DIM).transpose(0, 1) + k = rows[:, q_size : q_size + kv_size].view(-1, _KV_HEADS, _HEAD_DIM).transpose(0, 1) + v = rows[:, q_size + kv_size :].view(-1, _KV_HEADS, _HEAD_DIM).transpose(0, 1) + k = k.repeat_interleave(_HEADS // _KV_HEADS, dim=0) + v = v.repeat_interleave(_HEADS // _KV_HEADS, dim=0) + scores = q @ k.transpose(1, 2) / _HEAD_DIM**0.5 + q_pos = torch.arange(start, start + n, device=qkv.device)[:, None] + k_pos = torch.arange(kv_from, start + n, device=qkv.device)[None, :] + scores = scores.masked_fill(k_pos > q_pos, float("-inf")) + return (scores.softmax(-1) @ v).transpose(0, 1).reshape(n, q_size) + + +def _relative_error(out: torch.Tensor, ref: torch.Tensor) -> float: + out, ref = out.float(), ref.float() + scale = ref.abs().amax(dim=1).clamp_min(1e-6) + return ((out - ref).abs().amax(dim=1) / scale).max().item() + + +def _cached_kv(manager, rid: int, length: int) -> torch.Tensor: + pages, offsets = _token_slots(manager, rid, length) + return _raw_pool(manager)[pages, :, :, offsets, :] + + +@pytest.mark.parametrize("kv_dtype", _KV_DTYPES, ids=_KV_IDS) +@pytest.mark.parametrize("case", _CACHED_CONTEXT_CASES, ids=[c[0] for c in _CACHED_CONTEXT_CASES]) +def test_cached_context_matches_unchunked(served, case, kv_dtype) -> None: + """FlashInfer attends to the cached prefix: each chunk's output matches the same rows of the request run + in one unchunked step and an FP32 reference, and the chunks leave the same K / V in the cache.""" + name, steps = case + fp8 = kv_dtype == torch.float8_e4m3fn + tolerance = _TOLERANCE[kv_dtype] + lengths = {} + for step in steps: + for rid, c, n in step: + assert c == lengths.get(rid, 0), "each chunk starts at its request's cached length" + lengths[rid] = c + n + + gen = torch.Generator(device="cuda").manual_seed(len(name)) + width = (_HEADS + 2 * _KV_HEADS) * _HEAD_DIM + qkv = { + rid: torch.randn(n, width, generator=gen, device="cuda").to(torch.bfloat16) + for rid, n in lengths.items() + } + num_tokens = 2 * sum(-(-n // _PAGE) * _PAGE for n in lengths.values()) + 8 * _PAGE + manager = _kv_manager(kv_dtype, num_tokens) + try: + attn = _attention(kv_dtype) + rids = sorted(lengths) + manager.add_dummy_requests( + rids + [_REF_RID + r for r in rids], [lengths[r] for r in rids] * 2 + ) + # The unchunked run: every request as one fresh context. + rows = [(_REF_RID + r, 0, lengths[r]) for r in rids] + out = _step(attn, manager, rows, torch.cat([qkv[r] for r in rids]), len(rows)) + unchunked, first = {}, 0 + for r in rids: + unchunked[r] = out[first : first + lengths[r]] + first += lengths[r] + error = _relative_error(unchunked[r], _float_reference(qkv[r], 0, lengths[r], 0, fp8)) + print(f"\n[{name}] request {r} unchunked vs FP32 reference: {error:.4f}") + assert error <= tolerance["ctx"] + + for step in steps: + out = _step( + attn, + manager, + step, + torch.cat([qkv[r][c : c + n] for r, c, n in step]), + len(step), + ) + first = 0 + for r, c, n in step: + got = out[first : first + n] + first += n + assert torch.isfinite(got.float()).all() + error = _relative_error(got, unchunked[r][c : c + n]) + print(f"[{name}] request {r} rows {c}..{c + n} vs unchunked: {error:.4f}") + assert error <= tolerance["unchunked"] + reference = _float_reference(qkv[r], c, n, 0, fp8) + assert _relative_error(got, reference) <= tolerance["ctx"] + if c: + # A kernel that dropped the cached prefix would match attention over the chunk alone. + dropped = _float_reference(qkv[r], c, n, c, fp8) + assert _relative_error(got, dropped) > _DROPPED_PREFIX_MARGIN * tolerance["ctx"] + assert served["fi_ctx"] > 0 and served["fallback"] == 0, served + torch.cuda.synchronize() + for r in rids: + assert torch.equal( + _cached_kv(manager, r, lengths[r]), _cached_kv(manager, _REF_RID + r, lengths[r]) + ), f"request {r}: the chunks' K / V differ from the unchunked run's" + finally: + manager.shutdown() + + +@pytest.mark.parametrize("tail", ["zero", "nan"]) +@pytest.mark.parametrize("kv_dtype", _KV_DTYPES, ids=_KV_IDS) +@pytest.mark.parametrize("case", _GENERATION_CASES, ids=[c[0] for c in _GENERATION_CASES]) +def test_generation_verify_and_mixed_steps(served, case, kv_dtype, tail) -> None: + """Generation (q_len 1), linear verify (1 + 3) and mixed context + generation steps over long cached + contexts: FlashInfer serves every phase and each row is within tolerance of an FP32 reference from the + pool contents. With NaN bytes in a last page's unwritten tail, rows of requests whose last page is full + stay exact; the other requests' NaN rows are what the KV-pool zero-init prevents.""" + name, ctx, gen_rows = case + fp8 = kv_dtype == torch.float8_e4m3fn + tolerance = _TOLERANCE[kv_dtype] + reqs = list(ctx) + list(gen_rows) # (cached, new) per request + num_contexts = len(ctx) + totals = [c + n for c, n in reqs] + num_tokens = sum(-(-t // _PAGE) * _PAGE for t in totals) + 8 * _PAGE + manager = _kv_manager(kv_dtype, num_tokens) + try: + attn = _attention(kv_dtype) + rids = list(range(len(reqs))) + manager.add_dummy_requests(rids, totals) + raw = _raw_pool(manager) + gen = torch.Generator(device="cuda").manual_seed(5) + for rid, (c, n) in zip(rids, reqs): + if c: + # Cached K / V written straight into the pool. + kv = torch.randn(2, c, _KV_HEADS, _HEAD_DIM, generator=gen, device="cuda") * 2 + pages, offsets = _token_slots(manager, rid, c) + for s in (0, 1): + raw[pages, s, :, offsets, :] = kv[s].to(kv_dtype).view(torch.uint8) + if tail == "nan" and (c + n) % _PAGE: + last_page = manager.get_batch_cache_indices([rid], 0)[0][(c + n - 1) // _PAGE] + raw[last_page, :, :, (c + n) % _PAGE :, :] = 0xFF # NaN in BF16 and E4M3 + width = (_HEADS + 2 * _KV_HEADS) * _HEAD_DIM + qkv = torch.cat( + [ + torch.randn(n, width, generator=gen, device="cuda").to(torch.bfloat16) + for _, n in reqs + ] + ) + rows = [(rid, c, n) for rid, (c, n) in zip(rids, reqs)] + out = _step(attn, manager, rows, qkv, num_contexts) + torch.cuda.synchronize() + + assert served["fallback"] == 0, served + assert (served["fi_ctx"] > 0) == (num_contexts > 0), served + assert served["fi_gen"] > 0, served + + first = 0 + for i, (rid, (c, n)) in enumerate(zip(rids, reqs)): + got = out[first : first + n].float() + q = qkv[first : first + n, : _HEADS * _HEAD_DIM] + first += n + is_context = i < num_contexts + has_nan_tail = tail == "nan" and (c + n) % _PAGE != 0 + nonfinite = int((~torch.isfinite(got)).any(dim=1).sum()) + if has_nan_tail: + print(f"[{name}] request {rid} (tail NaN): {nonfinite}/{n} non-finite rows") + continue + assert nonfinite == 0, f"request {rid}: {nonfinite}/{n} non-finite rows" + ref = _reference_from_pool(manager, kv_dtype, rid, q, c, quantize_q=fp8 and is_context) + error = _relative_error(got, ref) + phase = "ctx" if is_context else "gen" + print( + f"[{name}] request {rid} ({phase}, past {c}, q_len {n}) vs FP32 reference: {error:.4f}" + ) + assert error <= tolerance[phase] + finally: + manager.shutdown() diff --git a/tests/unittest/_torch/models/test_minimax_m3.py b/tests/unittest/_torch/models/test_minimax_m3.py index 25fede204ec8..5e3ce1d1292b 100644 --- a/tests/unittest/_torch/models/test_minimax_m3.py +++ b/tests/unittest/_torch/models/test_minimax_m3.py @@ -2863,15 +2863,19 @@ def test_minimax_m3_swiglu_oai_fused_matches_reference(dtype): @pytest.mark.cpu_only @pytest.mark.parametrize( - "libraries,available,use_native", + "libraries,available,extra_pages_only,use_native", [ - ("flashinfer_trtllm_gen,fallback", True, True), - ("fallback,flashinfer_trtllm_gen", True, False), - ("fallback", True, False), - ("flashinfer_trtllm_gen,fallback", False, False), + ("flashinfer_trtllm_gen,fallback", True, False, True), + ("fallback,flashinfer_trtllm_gen", True, False, False), + ("fallback", True, False, False), + ("flashinfer_trtllm_gen,fallback", False, False, False), + # SM107: the library serves only extra page sizes; dense layers keep SDPA. + ("flashinfer_trtllm_gen,fallback", True, True, False), ], ) -def test_minimax_m3_dense_backend_library_order(monkeypatch, libraries, available, use_native): +def test_minimax_m3_dense_backend_library_order( + monkeypatch, libraries, available, extra_pages_only, use_native +): from types import SimpleNamespace from unittest.mock import Mock @@ -2892,6 +2896,7 @@ def test_minimax_m3_dense_backend_library_order(monkeypatch, libraries, availabl "is_available", classmethod(lambda cls, attn: available), ) + monkeypatch.setattr(model_module, "get_sm_version", lambda: 107 if extra_pages_only else 100) layer = SimpleNamespace( attn=legacy, layer_idx=0, @@ -2909,7 +2914,7 @@ def test_minimax_m3_dense_backend_library_order(monkeypatch, libraries, availabl ) assert layer.attn is (native if use_native else legacy) assert layer.support_fused_qkv is use_native - if libraries == "flashinfer_trtllm_gen,fallback" and available: + if libraries == "flashinfer_trtllm_gen,fallback" and available and not extra_pages_only: construct.assert_called_once() assert construct.call_args.kwargs["pos_embd_params"] is None else: diff --git a/tests/unittest/_torch/test_model_config.py b/tests/unittest/_torch/test_model_config.py index 0ed31615b378..98692ef294a9 100644 --- a/tests/unittest/_torch/test_model_config.py +++ b/tests/unittest/_torch/test_model_config.py @@ -15,6 +15,7 @@ validate_and_set_kv_cache_quant, validate_encoder_decoder_kv_cache_config, ) +from tensorrt_llm.llmapi.llm_args import MiniMaxM3SparseAttentionConfig from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig @@ -130,7 +131,11 @@ def test_validate_and_set_kv_cache_quant_auto_uses_checkpoint(): assert model_config.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 -def test_validate_and_set_kv_cache_quant_explicit_dtype_overrides(): +def test_validate_and_set_kv_cache_quant_explicit_dtype_overrides( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Pin a non-SM107 device: on SM107 an explicit NVFP4 request falls back to FP8. + monkeypatch.setattr("tensorrt_llm._torch.pyexecutor.model_loader.get_sm_version", lambda: 100) model_config = _make_model_config_with_kv_quant(QuantAlgo.FP8) validate_and_set_kv_cache_quant(model_config, "nvfp4") assert model_config.quant_config.kv_cache_quant_algo == QuantAlgo.NVFP4 @@ -189,18 +194,32 @@ def test_validate_and_set_kv_cache_quant_auto_keeps_quant_config_dict(): assert layer_quant_config.kv_cache_quant_algo is None +@pytest.mark.parametrize("implementation", [None, "triton", "msa"]) +@pytest.mark.parametrize("requested_dtype", ["nvfp4", "auto"]) def test_validate_and_set_kv_cache_quant_sm107_downgrade_propagates_to_quant_config_dict( monkeypatch: pytest.MonkeyPatch, + implementation: str | None, + requested_dtype: str, ) -> None: - """The SM107 NVFP4->FP8 downgrade must update per-layer QuantConfigs too, - otherwise a mixed-precision checkpoint keeps NVFP4 in quant_config_dict - while the KV pool (sized from the global config) uses FP8.""" + """The SM107 NVFP4->FP8 decision must reach per-layer QuantConfigs too, + otherwise a mixed-precision checkpoint and the KV pool (sized from the + global config) disagree. MiniMax-M3 MSA keeps NVFP4, which its hybrid + cache serves on SM107; every other configuration falls back to FP8.""" _mock_sm107(monkeypatch) model_config = _make_mixed_precision_model_config() - validate_and_set_kv_cache_quant(model_config, "nvfp4") - assert model_config.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 + if implementation is not None: + model_config.sparse_attention_config = MiniMaxM3SparseAttentionConfig( + implementation=implementation + ) + if requested_dtype == "auto": + model_config.quant_config.kv_cache_quant_algo = QuantAlgo.NVFP4 + for layer_quant_config in model_config.quant_config_dict.values(): + layer_quant_config.kv_cache_quant_algo = QuantAlgo.NVFP4 + validate_and_set_kv_cache_quant(model_config, requested_dtype) + expected = QuantAlgo.NVFP4 if implementation == "msa" else QuantAlgo.FP8 + assert model_config.quant_config.kv_cache_quant_algo == expected for layer_quant_config in model_config.quant_config_dict.values(): - assert layer_quant_config.kv_cache_quant_algo == QuantAlgo.FP8 + assert layer_quant_config.kv_cache_quant_algo == expected def _write_safetensors_header(checkpoint_dir, tensor_dtype, tensor_shape):