Skip to content

[https://nvbugs/6417488][fix] Rename M3 override to _prepare_page_table_tensor(index_mapper_capacity) and… #15984

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
trtllm-agent wants to merge 3 commits into
base: main
Choose a base branch
from
+17 −11
Open

[https://nvbugs/6417488][fix] Rename M3 override to _prepare_page_table_tensor(index_mapper_capacity) and… #15984

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
09f0ae4
[nvbugs/6417488][test] Unwaive MiniMax M3 NVFP4 accuracy test
liji-nv Jul 9, 2026
09f26bd
[nvbugs/6417488][fix] Reject SWA scratch reuse for MiniMax M3
liji-nv Jul 9, 2026
27dfb12
[nvbugs/6417488][fix] Harden MiniMax M3 C++ KV cache mapping
liji-nv Jul 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -365,6 +365,14 @@ def get_buffers(self, layer_idx: int, kv_layout: str = "NHD") -> Optional[torch.
full_view = convert_to_torch_tensor(TensorWrapper(addr_key, torch_dtype, full_slot_shape))
return full_view[:, :2]

def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None:
# M3's pool coalesces INDEX_KEY with K/V for sparse layers, so the
# base's SWA scratch-reuse branch (which assumes per-layer pools)
# does not describe the M3 layout. Fail fast rather than silently
# produce a mismatched pool mapping.
assert not self.enable_swa_scratch_reuse, "MiniMax M3 does not support SWA scratch reuse"
super()._prepare_page_table_tensor(index_mapper_capacity)

def _kv_pool_mapping_offset(self, layer_id, layer_group_id, key_base_addr) -> int:
"""Pool-mapping offset from the layer's physical position in its pool.

Expand Down
19 changes: 9 additions & 10 deletions tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@
AttnLifeCycle,
BatchDesc,
BufferConfig,
BufferId,
CacheLevel,
CacheTierConfig,
CuError,
Expand Down Expand Up @@ -1924,17 +1925,15 @@ def get_index_k_buffer(
if layer_idx not in self.layer_offsets:
return None
layer_offset = self.layer_offsets[layer_idx]
try:
addr = self.impl.get_mem_pool_base_address(layer_offset, Role.INDEX_KEY)
page_stride = self.impl.get_page_stride(layer_offset, Role.INDEX_KEY)
page_upper = self.impl.get_page_index_upper_bound(layer_offset, Role.INDEX_KEY)
converter = self.impl.get_page_index_converter(layer_offset, Role.INDEX_KEY)
except KeyError:
# INDEX_KEY not registered for this layer (default V2 manager
# registers only K/V/scale; sparse subclasses register
# INDEX_KEY only on sparse layers via
# ``_extra_buffers_per_layer``).
buffer_id = BufferId(layer_offset, Role.INDEX_KEY)
if buffer_id not in self.impl.all_buffer_ids:
# The default V2 manager registers only K/V/scale buffers;
# sparse subclasses register INDEX_KEY only on sparse layers.
return None
addr = self.impl.get_mem_pool_base_address(layer_offset, Role.INDEX_KEY)
page_stride = self.impl.get_page_stride(layer_offset, Role.INDEX_KEY)
page_upper = self.impl.get_page_index_upper_bound(layer_offset, Role.INDEX_KEY)
converter = self.impl.get_page_index_converter(layer_offset, Role.INDEX_KEY)

if isinstance(dtype, DataType):
torch_dtype = binding_to_torch_dtype(dtype)
Expand Down
1 change: 0 additions & 1 deletion tests/integration/test_lists/waives.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -215,7 +215,6 @@ full:GB300/accuracy/test_llm_api_pytorch.py::TestGemma3_1BInstruct::test_fp8_pre
full:GB300/accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_bfloat16_4gpus[pp4-attn_backend=FLASHINFER-torch_compile=False] SKIP (https://nvbugs/6385771)
full:GB300/accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[pp4-fp8kv=True-attn_backend=FLASHINFER-torch_compile=False] SKIP (https://nvbugs/6385771)
full:GB300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[use_msa=False] SKIP (https://nvbugs/6422502)
full:GB300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=False] SKIP (https://nvbugs/6422502)
full:GB300/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_fp8_moe_dflash SKIP (https://nvbugs/6316985)
full:GB300/accuracy/test_llm_api_pytorch_multimodal.py::TestGemma3_27BInstruct::test_fp8_prequantized SKIP (https://nvbugs/6479708)
full:GB300/test_e2e.py::test_qwen_e2e_cpprunner_large_new_tokens[DeepSeek-R1-Distill-Qwen-1.5B-DeepSeek-R1-Distill-Qwen-1.5B] SKIP (https://nvbugs/6414760)
Expand Down
Loading