diff --git a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py index cf306b142f42..ecce086889ed 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py @@ -21,8 +21,10 @@ Queries have no rotary component and are absorbed into the 512-wide latent space. TRTLLM-GEN serves supported FP8 cache shapes and small BF16 TP generation batches. Other shapes use FlashMLA, gathering and dequantizing FP8 cache rows -in bounded query chunks when needed. Prepared GLM page tables and live TRTLLM -lengths provide the same cache contract for prefill, decode and verification. +in bounded query chunks when needed. On SM120/SM121, FlashInfer's GLM53_NOPE +sparse MLA reads the selected FP8 rows, packed with inline scales in the same +bounded query chunks. Prepared GLM page tables and live TRTLLM lengths provide +the same cache contract for prefill, decode and verification. """ from __future__ import annotations @@ -35,6 +37,7 @@ from tensorrt_llm._torch.pyexecutor.config_utils import unwrap_glm5_next_text_config from tensorrt_llm._utils import get_sm_version, is_sm_100f +from tensorrt_llm.math_utils import ceil_div from ...interface import ( AttentionForwardArgs, @@ -44,6 +47,11 @@ merge_attention_forward_args, ) from ...trtllm import TrtllmAttention, TrtllmAttentionMetadata +from ..flashinfer_utils import ( + SPARSE_MLA_SPLIT_KV_TILE, + allocate_sparse_mla_split_workspace, + get_sparse_mla_op, +) from .kernels import ( append_paged_state_rows, gather_fp8_kv_rows, @@ -51,6 +59,7 @@ kpool_score, kpool_score_fp8, kpool_update, + pack_glm53_kv_rows, ) from .metadata import Glm5NextMamba2Metadata from .native_decode import GlmKpoolNativeDecode @@ -136,6 +145,9 @@ class GlmKpoolSparseAttention(TrtllmAttention): # size and discard the extra outputs; heads attend independently. _KERNEL_HEAD_COUNTS = (64, 128) + # FlashInfer's SM120 GLM53_NOPE prefill instantiates these query-head counts. + _SM120_KERNEL_HEAD_COUNTS = (8, 16, 32, 64) + # Bound selected-KV staging independently of the configured context length. # At the published 2112-entry selection this uses 132 MiB of BF16 rows. _FP8_QUERY_CHUNK_SIZE = 64 @@ -638,6 +650,8 @@ def _dispatch_sparse_core( output ``[T, H, kv_lora]``; :meth:`_finalize_output` flattens it to the base-contract shape at the ``forward`` boundary. """ + if get_sm_version() in (120, 121): + return self._dispatch_sparse_core_sm120(q_latent, kv_rows, topk_rows) if kv_rows.dtype == torch.float8_e4m3fn: # Like DeepSeek-V4's BF16 context adapter, dequantize only selected # rows and reuse FlashMLA. Chunk query rows so prefill never stages @@ -684,6 +698,72 @@ def _dispatch_sparse_core( out = out[:, :local_heads, :] return out + def _dispatch_sparse_core_sm120( + self, + q_latent: torch.Tensor, + kv_rows: torch.Tensor, + topk_rows: torch.Tensor, + ) -> torch.Tensor: + """Run FlashInfer's GLM53_NOPE sparse MLA; same contract as :meth:`_dispatch_sparse_core`. + + The kernel reads 528-byte FP8 rows with inline group scales, so each + query chunk's selected rows are packed from the per-tensor-scaled FP8 + cache (:func:`pack_glm53_kv_rows`) and addressed by local row ids. + """ + if kv_rows.dtype != torch.float8_e4m3fn: + raise NotImplementedError( + "glm_kpool on SM120/SM121 runs FlashInfer's FP8 sparse MLA and needs an " + f"FP8 latent cache (kv_cache_config.dtype='fp8'); got {kv_rows.dtype}" + ) + pad = (-topk_rows.shape[-1]) % self._KERNEL_TOPK_ALIGN + if pad: + topk_rows = torch.nn.functional.pad(topk_rows, (0, pad), value=INDEX_SENTINEL) + num_tokens, local_heads, dim = q_latent.shape + kernel_heads = next((h for h in self._SM120_KERNEL_HEAD_COUNTS if h >= local_heads), None) + if kernel_heads is None: + raise ValueError( + f"glm_kpool: {local_heads} query heads exceed every FlashInfer SM120 " + f"sparse-kernel instantiation {self._SM120_KERNEL_HEAD_COUNTS}" + ) + if kernel_heads != local_heads: + q_padded = q_latent.new_zeros((num_tokens, kernel_heads, dim)) + q_padded[:, :local_heads, :] = q_latent + q_latent = q_padded + else: + q_latent = q_latent.contiguous() + output = q_latent.new_empty((num_tokens, kernel_heads, self.kv_lora_rank)) + num_splits = ceil_div(topk_rows.shape[-1], SPARSE_MLA_SPLIT_KV_TILE) + for start in range(0, num_tokens, self._FP8_QUERY_CHUNK_SIZE): + stop = min(start + self._FP8_QUERY_CHUNK_SIZE, num_tokens) + pages, local_indices = pack_glm53_kv_rows( + kv_rows, topk_rows[start:stop], self.kv_scale_quant_orig + ) + out_lse = torch.empty( + (stop - start, kernel_heads), dtype=torch.float32, device=q_latent.device + ) + mid_out, mid_lse = allocate_sparse_mla_split_workspace( + num_tokens=stop - start, + num_heads=kernel_heads, + num_splits=num_splits, + value_dim=self.kv_lora_rank, + device=q_latent.device, + ) + get_sparse_mla_op()( + q_latent[start:stop], + pages, + local_indices, + output[start:stop], + out_lse, + self.softmax_scale, + d_v=self.kv_lora_rank, + kv_scale_format="arbitrary_fp32", + mid_out=mid_out, + mid_lse=mid_lse, + ) + if kernel_heads != local_heads: + output = output[:, :local_heads, :] + return output + def _finalize_output( self, out_latent: torch.Tensor, output: torch.Tensor | None ) -> torch.Tensor: diff --git a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/kernels.py b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/kernels.py index 86037408ca0d..60679637ce59 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/kernels.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/kernels.py @@ -84,6 +84,85 @@ def gather_fp8_kv_rows( return selected, local_indices +#: FlashInfer's GLM53_NOPE sparse-MLA row: 512 FP8 latent values, then four +#: FP32 scales, one per 128-wide group. +GLM53_NOPE_ROW_BYTES = 528 +GLM53_NOPE_PAGE_ROWS = 64 + + +@triton.jit +def _pack_glm53_kv_rows_kernel( + rows, + indices, + scale, + packed, + packed_scales, + local_indices, + row_stride, + index_stride, + column_stride, + NUM_ROWS: tl.constexpr, + DIM: tl.constexpr, + ROW_BYTES: tl.constexpr, + NUM_SCALES: tl.constexpr, + TOPK: tl.constexpr, +): + token = tl.program_id(0) + index = tl.load(indices + (token // TOPK) * index_stride + (token % TOPK) * column_stride) + valid = (index >= 0) & (index < NUM_ROWS) + dim = tl.arange(0, DIM) + payload = tl.load(rows + index.to(tl.int64) * row_stride + dim, mask=valid, other=0) + tl.store(packed + token.to(tl.int64) * ROW_BYTES + dim, payload) + group_scales = tl.load(scale) + tl.zeros([NUM_SCALES], dtype=tl.float32) + scale_offset = token.to(tl.int64) * (ROW_BYTES // 4) + DIM // 4 + tl.store(packed_scales + scale_offset + tl.arange(0, NUM_SCALES), group_scales) + tl.store(local_indices + token, tl.where(valid, token, -1)) + + +def pack_glm53_kv_rows( + rows: torch.Tensor, indices: torch.Tensor, scale: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """Pack selected FP8 rows into FlashInfer GLM53_NOPE pages and remap indices. + + ``rows`` is ``[N, 1, 512]`` FP8 with contiguous features, ``indices`` is + ``[queries, topk]`` with ``topk`` a multiple of 64 and ``scale`` is the + cache's device FP32 dequantization scalar. Each packed row keeps the cached + FP8 bytes and repeats ``scale`` for all four groups, so it dequantizes to + the cached values. Returns uint8 pages ``[queries * topk / 64, 64, 528]`` + and int32 local row ids; invalid indices remain -1. + """ + dim = rows.shape[-1] + num_scales = (GLM53_NOPE_ROW_BYTES - dim) // 4 + if rows.dtype != torch.float8_e4m3fn or num_scales * 128 != dim: + raise ValueError( + f"GLM53_NOPE packing needs FP8 rows of width 512, got {rows.dtype} width {dim}" + ) + if indices.shape[1] % GLM53_NOPE_PAGE_ROWS: + raise ValueError(f"GLM53_NOPE packing needs topk % 64 == 0, got {indices.shape[1]}") + packed = torch.empty( + (indices.numel(), GLM53_NOPE_ROW_BYTES), dtype=torch.uint8, device=rows.device + ) + local_indices = torch.empty(indices.shape, dtype=torch.int32, device=indices.device) + if indices.numel(): + _pack_glm53_kv_rows_kernel[(indices.numel(),)]( + rows.view(torch.uint8), + indices, + scale, + packed, + packed.view(torch.float32), + local_indices, + rows.stride(0), + indices.stride(0), + indices.stride(1), + NUM_ROWS=rows.shape[0], + DIM=dim, + ROW_BYTES=GLM53_NOPE_ROW_BYTES, + NUM_SCALES=num_scales, + TOPK=indices.shape[1], + ) + return packed.view(-1, GLM53_NOPE_PAGE_ROWS, GLM53_NOPE_ROW_BYTES), local_indices + + _FP32_MIN = torch.finfo(torch.float32).min diff --git a/tests/integration/defs/accuracy/references/gsm8k.yaml b/tests/integration/defs/accuracy/references/gsm8k.yaml index 81c2973cd34b..e5a0cca13ef1 100644 --- a/tests/integration/defs/accuracy/references/gsm8k.yaml +++ b/tests/integration/defs/accuracy/references/gsm8k.yaml @@ -616,3 +616,7 @@ zai-org/GLM-5.3-Flash: kv_cache_quant_algo: FP8 spec_dec_algo: MTP accuracy: 92.91 + # RTX PRO 6000 TP4/EP4, NVFP4 checkpoint with FP8 KV cache, harness protocol. + - quant_algo: NVFP4 + kv_cache_quant_algo: FP8 + accuracy: 92.87 diff --git a/tests/integration/defs/accuracy/test_glm53_flash.py b/tests/integration/defs/accuracy/test_glm53_flash.py index 90372dadbeb5..ed303d277c35 100644 --- a/tests/integration/defs/accuracy/test_glm53_flash.py +++ b/tests/integration/defs/accuracy/test_glm53_flash.py @@ -467,3 +467,20 @@ def _assert_repeated_request_is_deterministic(llm: LLM) -> None: second = llm.generate([prompt], sampling_params=sampling_params, use_tqdm=False)[0] assert len(first.outputs[0].token_ids) == output_length assert second.outputs[0].token_ids == first.outputs[0].token_ids + + +class TestGLM53FlashNVFP4(LlmapiAccuracyTestHarness): + """GLM-5.3-Flash NVFP4 (FP8 KV cache) accuracy on Blackwell, including SM120.""" + + MODEL_NAME = "zai-org/GLM-5.3-Flash" + MODEL_PATH = f"{llm_models_root()}/GLM-5.3-Flash-NVFP4" + + @skip_pre_blackwell + @pytest.mark.skip_less_mpi_world_size(4) + @parametrize_with_ids("tp_size,ep_size", [(4, 4)]) + def test_tep(self, tp_size, ep_size): + with LLM(self.MODEL_PATH, **TestGLM53FlashFP8._llm_kwargs(tp_size, ep_size)) as llm: + assert llm.args.quant_config.quant_algo == QuantAlgo.NVFP4 + assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm) diff --git a/tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml b/tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml index 49f76a2b0b7c..0484aa6fe9c8 100644 --- a/tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml +++ b/tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml @@ -122,6 +122,7 @@ l0_rtx_pro_6000: - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_guided_decoding[xgrammar-mtp_nextn=0] - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_guided_decoding[xgrammar-mtp_nextn=2] - accuracy/test_llm_api_pytorch.py::TestQwen3_235B_A22B::test_nvfp4_4gpus[latency_moe_cutlass] # 20 mins + - accuracy/test_glm53_flash.py::TestGLM53FlashNVFP4::test_tep[tp_size=4-ep_size=4] # - accuracy/test_llm_api_pytorch.py::TestQwen3_235B_A22B::test_nvfp4_4gpus[latency_moe_cutlass_eagle3] # failed - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[dep4_latency_moe_cutlass-torch_compile=True] # - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[dep4_latency_moe_cutlass-torch_compile=False] # failed diff --git a/tests/unittest/_torch/attention/sparse/glm_kpool/test_kernels.py b/tests/unittest/_torch/attention/sparse/glm_kpool/test_kernels.py index 42a225841264..c62d5c6f4d47 100644 --- a/tests/unittest/_torch/attention/sparse/glm_kpool/test_kernels.py +++ b/tests/unittest/_torch/attention/sparse/glm_kpool/test_kernels.py @@ -569,6 +569,9 @@ def reference(): @pytest.mark.skipif(not HAS_FLASH_MLA, reason="FlashMLA not available") @pytest.mark.skipif(get_sm_version() < 90, reason="FlashMLA requires SM90 (Hopper) or later") +@pytest.mark.skipif( + get_sm_version() in (120, 121), reason="SM120/121 runs FlashInfer, not FlashMLA" +) @pytest.mark.parametrize("heads", [16, 64]) def test_fp8_sparse_core_matches_dequantized_cache_and_replays_graph(heads): """Exercise the selected-row core used by production prefill, decode and verify.""" @@ -612,6 +615,106 @@ def reference(): torch.testing.assert_close(captured, reference(), atol=2e-2, rtol=2e-2) +def test_glm53_packing_keeps_fp8_payload_and_cache_scale(): + """Packed rows are the cached FP8 bytes plus the per-tensor scale in every group slot.""" + from tensorrt_llm._torch.attention.backends.sparse.glm_kpool import latent_pool_rows + from tensorrt_llm._torch.attention.backends.sparse.glm_kpool.kernels import pack_glm53_kv_rows + + generator = torch.Generator(device="cuda").manual_seed(5) + storage = torch.randn(9, 8, 512, device="cuda", generator=generator).to(torch.float8_e4m3fn) + rows, base, stride = latent_pool_rows(storage[2::3]) + positions = torch.randint(0, 24, (3, 128), device="cuda", generator=generator) + indices = (base + positions // 8 * stride + positions % 8).int() + indices[:, 5] = -1 + indices[:, 7] = rows.shape[0] + scale = torch.tensor([0.0375], device="cuda") + + pages, local = pack_glm53_kv_rows(rows, indices, scale) + + assert pages.shape == (6, 64, 528) and pages.dtype == torch.uint8 + flat = pages.view(-1, 528) + valid = (indices >= 0) & (indices < rows.shape[0]) + expected_local = torch.arange(indices.numel(), device="cuda", dtype=torch.int32).view_as(local) + assert torch.equal(local, torch.where(valid, expected_local, -1)) + valid_flat = valid.flatten() + expected_payload = rows.view(torch.uint8)[indices.flatten()[valid_flat].long(), 0] + assert torch.equal(flat[valid_flat, :512], expected_payload) + assert torch.count_nonzero(flat[~valid_flat, :512]) == 0 + assert torch.equal( + flat[:, 512:].contiguous().view(torch.float32), scale.expand(flat.shape[0], 4) + ) + + +def _flashinfer_has_glm53_nope() -> bool: + try: + from flashinfer.mla._sparse_mla_sm120 import _DECODE_GLM53_NOPE_DISPATCH # noqa: F401 + except ImportError: + return False + return True + + +def _reference_sparse_mla(q, rows, indices, softmax_scale): + """Dense softmax(q . k) v over each token's valid selected rows (K = V = row).""" + valid = (indices >= 0) & (indices < rows.shape[0]) + kv = rows[indices.clamp(0, rows.shape[0] - 1).long(), 0].float() + scores = torch.einsum("thd,tkd->thk", q.float(), kv) * softmax_scale + scores = scores.masked_fill(~valid[:, None, :], float("-inf")) + return torch.einsum("thk,tkd->thd", scores.softmax(-1), kv).to(q.dtype) + + +@pytest.mark.skipif( + get_sm_version() not in (120, 121), reason="FlashInfer GLM53_NOPE sparse MLA requires SM120/121" +) +@pytest.mark.skipif( + not _flashinfer_has_glm53_nope(), reason="FlashInfer without GLM53_NOPE sparse MLA" +) +@pytest.mark.parametrize("heads", [12, 16, 64]) +def test_fp8_sparse_core_sm120_matches_dequantized_cache_and_replays_graph(heads): + """SM120 core: packed FP8 rows through FlashInfer match attention over the dequantized cache.""" + from tensorrt_llm._torch.attention.backends.sparse.glm_kpool import ( + GlmKpoolSparseAttention, + latent_pool_rows, + ) + + backend = object.__new__(GlmKpoolSparseAttention) + backend.kv_lora_rank = 512 + backend.softmax_scale = 256**-0.5 + # Scores with a standard deviation near 0.4. FlashInfer's MXFP8 MMAs round + # queries and probabilities to FP8, so sharper softmaxes need wider bounds. + backend.kv_scale_quant_orig = torch.tensor([0.0375 / 8], device="cuda") + backend._FP8_QUERY_CHUNK_SIZE = 2 # Exercise multiple chunks and a short final chunk. + generator = torch.Generator(device="cuda").manual_seed(71) + storage = ( + torch.randn(12, 64, 512, device="cuda", generator=generator).mul(64).clamp(-448, 448) + ).to(torch.float8_e4m3fn) + pool = storage[1::3] # Nonzero offset and coalesced page stride. + rows, base, stride = latent_pool_rows(pool) + q = torch.randn(5, heads, 512, device="cuda", generator=generator).to(torch.bfloat16) + positions = torch.randint(0, 256, (5, 67), device="cuda", generator=generator) + indices = (base + positions // 64 * stride + positions % 64).int() + indices[:, -2] = -1 + indices[:, -1] = rows.shape[0] # Out-of-range rows are masked like -1. + + def reference(): + decoded = (rows.float() * backend.kv_scale_quant_orig).to(torch.bfloat16) + return _reference_sparse_mla(q, decoded, indices, backend.softmax_scale) + + actual = backend._dispatch_sparse_core(q, rows, indices) + assert actual.shape == (5, heads, 512) + torch.testing.assert_close(actual, reference(), atol=5e-3, rtol=5e-2) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + backend._dispatch_sparse_core(q, rows, indices) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = backend._dispatch_sparse_core(q, rows, indices) + storage.copy_((storage.float() * 0.5).to(storage.dtype)) + graph.replay() + torch.testing.assert_close(captured, reference(), atol=5e-3, rtol=5e-2) + + @pytest.mark.parametrize("context", [False, True]) def test_fp8_cache_storage_helpers_preserve_coalesced_pages(context): """Write FP8 payloads without disturbing coalesced pages or indexer pool keys."""