nvchenghaoz · Oct 19, 2026 · Nov 3, 2026 · chzblych · Nov 3, 2026 · coderabbitai
@@ -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,13 +47,19 @@
     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,
     kpool_expand,
     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:

@@ -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
 
 

@@ -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
@@ -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)
@@ -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

@@ -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."""