diff --git a/cpp/tensorrt_llm/flash_mla/CMakeLists.txt b/cpp/tensorrt_llm/flash_mla/CMakeLists.txt index 84fff15b189e..9f552d82c0c0 100644 --- a/cpp/tensorrt_llm/flash_mla/CMakeLists.txt +++ b/cpp/tensorrt_llm/flash_mla/CMakeLists.txt @@ -197,8 +197,7 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9) ) endif() -# Disable LTO before creating target (similar to DeepEP) Let CMake generate -# fatbinData for CUDA separable compilation +# Disable LTO before creating target (similar to DeepEP) set(CMAKE_INTERPROCEDURAL_OPTIMIZATION FALSE) pybind11_add_module(flash_mla_cpp_tllm ${FLASH_MLA_SOURCES}) @@ -210,8 +209,10 @@ set_target_properties( CXX_SCAN_FOR_MODULES OFF CUDA_STANDARD 20 CUDA_STANDARD_REQUIRED ON - CUDA_SEPARABLE_COMPILATION ON - CUDA_RESOLVE_DEVICE_SYMBOLS ON + # Whole-program compilation, as FlashMLA's own build does: + # relocatable device code makes ptxas spill registers in the + # attention kernels. + CUDA_SEPARABLE_COMPILATION OFF LINK_DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/flash_mla_cpp_tllm.version INSTALL_RPATH "${TORCH_INSTALL_PREFIX}/lib" BUILD_WITH_INSTALL_RPATH TRUE) diff --git a/cpp/tensorrt_llm/kernels/moe/cutlass/moe_kernels.cu b/cpp/tensorrt_llm/kernels/moe/cutlass/moe_kernels.cu index 1b4959485862..909311caa19f 100644 --- a/cpp/tensorrt_llm/kernels/moe/cutlass/moe_kernels.cu +++ b/cpp/tensorrt_llm/kernels/moe/cutlass/moe_kernels.cu @@ -1862,7 +1862,17 @@ __global__ void expandInputRowsFp8BlockScaleKernel(InputActivationsType const* u int64_t const num_elems_in_col = hidden_size / ELEM_PER_THREAD; auto* const permuted_output = fp8_block_scale_out.fp8_out; - for (int64_t permuted_row = blockIdx.x; permuted_row < num_valid_tokens; permuted_row += gridDim.x) + // Each block takes a contiguous run of permuted rows: they are sorted by expert, so one search finds the + // first row's expert and the rest advance linearly instead of searching per row. + int64_t const rows_per_block = ceilDiv(num_valid_tokens, static_cast(gridDim.x)); + int64_t const row_begin = blockIdx.x * rows_per_block; + int64_t const row_end + = row_begin + rows_per_block < num_valid_tokens ? row_begin + rows_per_block : num_valid_tokens; + int64_t expert = row_begin < row_end + ? findTotalEltsLessThanTarget(expert_first_token_offset, num_experts_per_node, row_begin + 1) - 1 + : 0; + + for (int64_t permuted_row = row_begin; permuted_row < row_end; ++permuted_row) { int64_t const unpermuted_row = permuted_row_to_unpermuted_row[permuted_row]; int64_t const source_k_rank = unpermuted_row / num_tokens; @@ -1871,8 +1881,10 @@ __global__ void expandInputRowsFp8BlockScaleKernel(InputActivationsType const* u auto const* source_row_ptr = reinterpret_cast(unpermuted_input + source_row * hidden_size); auto* dest_row_ptr = permuted_output + permuted_row * hidden_size; - int64_t const expert - = findTotalEltsLessThanTarget(expert_first_token_offset, num_experts_per_node, permuted_row + 1) - 1; + while (expert + 1 < num_experts_per_node && expert_first_token_offset[expert + 1] <= permuted_row) + { + ++expert; + } int64_t const scale_row = permuted_row + fp8BlockScaleExpertPad(expert_first_token_offset[expert], expert); for (int64_t elem_index = threadIdx.x; elem_index < num_elems_in_col; elem_index += EXPAND_THREADS_PER_BLOCK) @@ -2873,8 +2885,11 @@ void doActivation(T* output, GemmOutputType const* gemm_result, float const* fp8 }(); static int32_t const sm_count = tensorrt_llm::common::getMultiProcessorCount(); - int32_t const max_blocks_per_sm - = tensorrt_llm::common::getMaxActiveBlocksPerSM(fn, ACTIVATION_THREADS_PER_BLOCK, 0); + // A row narrower than the block (e.g. a tensor-parallel slice of the intermediate size) would leave the + // remaining threads idle, so shrink the block to the row width, in whole warps. + int32_t const threads = static_cast( + std::min(ACTIVATION_THREADS_PER_BLOCK, ceilDiv(num_elems_in_col, int64_t{32}) * 32)); + int32_t const max_blocks_per_sm = tensorrt_llm::common::getMaxActiveBlocksPerSM(fn, threads, 0); // X dimension for tokens in groups of num_rows_per_cta_v // Y dimension for columns // Each CTA processes num_rows_per_cta_v * threads_x tokens @@ -2882,9 +2897,7 @@ void doActivation(T* output, GemmOutputType const* gemm_result, float const* fp8 // Divide by rows_per_cta to get the number of token blocks needed int64_t const max_token_blocks = (max_tokens + num_rows_per_cta_v - 1) / num_rows_per_cta_v; int32_t const grid_x = std::min(sm_count * max_blocks_per_sm, static_cast(max_token_blocks)); - int32_t const grid_y = static_cast( - (num_elems_in_col + ACTIVATION_THREADS_PER_BLOCK - 1) / ACTIVATION_THREADS_PER_BLOCK); - int32_t const threads = ACTIVATION_THREADS_PER_BLOCK; + int32_t const grid_y = static_cast((num_elems_in_col + threads - 1) / threads); cudaLaunchConfig_t config; config.gridDim = dim3(grid_x, grid_y, 1); 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 13937deed749..e8f963137556 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py @@ -27,11 +27,14 @@ from __future__ import annotations +import os +from collections.abc import Sequence from dataclasses import dataclass import torch -from tensorrt_llm._utils import get_sm_version +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 ...interface import ( AttentionForwardArgs, @@ -46,7 +49,9 @@ gather_fp8_kv_rows, kpool_expand, kpool_score, + kpool_score_fp8, kpool_update, + merge_split_attention, ) from .metadata import Glm5NextMamba2Metadata from .native_decode import GlmKpoolNativeDecode @@ -125,8 +130,12 @@ class GlmKpoolSparseAttention(TrtllmAttention): #: FlashMLA tiles the top-k axis in blocks of 64 (``B_TOPK``); index rows #: are padded to that multiple with ``-1`` (invalid) entries, which is - #: semantics-free by the kernel's own contract. + #: semantics-free by the kernel's own contract. Its SM90 sparse prefill + #: walks two blocks per step and requires a multiple of 128. _KERNEL_TOPK_ALIGN = 64 + _SM90_KERNEL_TOPK_ALIGN = 128 + #: Calls with at most this many query tokens split the top-k axis on SM90. + _SM90_SPLIT_TOPK_MAX_TOKENS = 16 # FlashMLA instantiates 64/128 query heads. Pad TP-local heads to the next # size and discard the extra outputs; heads attend independently. @@ -136,6 +145,64 @@ class GlmKpoolSparseAttention(TrtllmAttention): # At the published 2112-entry selection this uses 132 MiB of BF16 rows. _FP8_QUERY_CHUNK_SIZE = 64 + # FP32 elements one packed-context scoring call may hold: logits, their row + # copy and gathered pool keys grow with the visible pools, which a fresh-prompt + # profiling forward does not reach. plan_context_scores tiles calls to this + # bound and runtime_workspace_fixed_bytes reserves it in KV-cache sizing. + _CONTEXT_SCORE_ELEMENTS = 1 << 26 + + @staticmethod + def _context_score_cost( + rows: int, tables: int, pools: int, *, fp8: bool, head_dim: int, capacity_pools: int + ) -> int: + """FP32 elements a scoring call of ``rows`` rows over ``tables`` tables holds at once. + + Counts the buffers that grow with the visible ``pools``: the BF16 scorer's + capacity-wide output, or every buffer :func:`kpool_score_fp8` keeps alive + together -- the padded logits, a multi-table call's gather temporary and + aligned row copy, and per gathered key its BF16 row, FP8 row, scale and + int64 slot index. + """ + if not fp8 or pools == 0: + return rows * capacity_pools + # fp8_mqa_logits pads each row to a multiple of 256 columns, plus 256. + logits = -(-tables * pools // 256) * 256 + 256 + copy = max(1024, -(-pools // 64) * 64) + pools if tables > 1 else 0 + keys = tables * pools * (head_dim // 2 + head_dim // 4 + 3) + # int64 pool positions, pages and offsets. + return rows * (logits + copy) + keys + 6 * pools + + @classmethod + def runtime_workspace_fixed_bytes( + cls, + model_config, + mapping, + *, + max_batch_size: int, + max_num_tokens: int, + max_seq_len: int, + enable_cuda_graph: bool, + ) -> int: + """Bytes of the largest scoring call plan_context_scores can emit for these capacities.""" + if max_num_tokens <= 0 or max_seq_len <= 0: + return 0 + config = unwrap_glm5_next_text_config(model_config.pretrained_config) + defaults = GlmKpoolSparseParams() + kpool = getattr(config, "index_kpool", defaults.index_kpool) + head_dim = getattr(config, "index_head_dim", defaults.index_head_dim) + pools = -(-max_seq_len // kpool) + tables = max(1, min(max_batch_size, max_num_tokens)) + + def cost(rows: int, tables: int, fp8: bool) -> int: + return cls._context_score_cost( + rows, tables, pools, fp8=fp8, head_dim=head_dim, capacity_pools=pools + ) + + largest = max(cost(max_num_tokens, tables, True), cost(max_num_tokens, 1, False)) + # Calls fit the budget, except that a single row is never split. + bound = max(cls._CONTEXT_SCORE_ELEMENTS, cost(1, 1, True), cost(1, 1, False)) + return min(largest, bound) * 4 + def __init__( self, layer_idx: int, @@ -205,11 +272,23 @@ def __init__( #: score scale is unchanged. self.softmax_scale = float(self.qk_nope_head_dim) ** -0.5 self._native_decode = GlmKpoolNativeDecode() + self._kernel_topk_align = ( + self._SM90_KERNEL_TOPK_ALIGN if get_sm_version() == 90 else self._KERNEL_TOPK_ALIGN + ) + self._kernel_output_width = sparse_params.kernel_output_width(self._kernel_topk_align) + #: The SM90 sparse kernel runs one CTA per query token, so a decode-sized + #: call occupies a handful of SMs; split its top-k axis across CTAs. + self._split_topk = get_sm_version() == 90 #: ``softmax_scale * kv_scale_quant_orig`` for the native FP8 kernel, read to #: the host once outside CUDA graph capture (the KV scale is fixed after #: construction). A float lets the kernel take it as an argument, without a #: per-call device multiply. self._native_bmm1_scale: float | None = None + #: Packed context rows score pools in FP8 with DeepGEMM MQA logits on SM90 and + #: SM100-family GPUs; TRTLLM_GLM_INDEXER_BF16=1 keeps the BF16 Triton scorer for them. + self._indexer_fp8 = (get_sm_version() == 90 or is_sm_100f()) and os.environ.get( + "TRTLLM_GLM_INDEXER_BF16", "0" + ) != "1" @classmethod def support_fused_rope(cls) -> bool: @@ -372,6 +451,63 @@ def update_pool_keys( request_ids=request_ids, ) + def plan_context_scores( + self, metadata, spans: Sequence[tuple[int, int]] + ) -> list[tuple[int, int, int, int, int]]: + """Split packed context scoring into calls within ``_CONTEXT_SCORE_ELEMENTS``. + + spans lists ``(rows, visible bound)`` of each context request in packed + order. Each returned ``(row_start, row_end, table_start, table_end, + max_kv_len)`` scores rows ``row_start:row_end`` against context tables + ``table_start:table_end``. Consecutive requests share a call while it + fits; a request that does not fit alone is split by rows. Costs count + only the transients that grow with the visible pools. + """ + state = self._cache_state(metadata) + kpool = self.sparse_params.index_kpool + budget = self._CONTEXT_SCORE_ELEMENTS + # The BF16 scorer writes every pool column of the configured capacity. + capacity_pools = -(-state.block_tables.shape[1] * state.tokens_per_block // kpool) + + def cost(rows: int, tables: int, pools: int) -> int: + return self._context_score_cost( + rows, + tables, + pools, + fp8=self._indexer_fp8, + head_dim=self.sparse_params.index_head_dim, + capacity_pools=capacity_pools, + ) + + calls: list[tuple[int, int, int, int, int]] = [] + group: list[int] | None = None # [row_start, rows, table_start, tables, max visible] + row_start = 0 + for table, (rows, visible) in enumerate(spans): + if group is not None: + merged = max(group[4], visible) + if cost(group[1] + rows, group[3] + 1, merged // kpool) <= budget: + group[1] += rows + group[3] += 1 + group[4] = merged + row_start += rows + continue + calls.append( + (group[0], group[0] + group[1], group[2], group[2] + group[3], group[4]) + ) + group = None + if cost(rows, 1, visible // kpool) <= budget: + group = [row_start, rows, table, 1, visible] + else: + fixed = cost(0, 1, visible // kpool) + tile = max(1, (budget - fixed) // max(1, cost(1, 1, visible // kpool) - fixed)) + for start in range(row_start, row_start + rows, tile): + end = min(start + tile, row_start + rows) + calls.append((start, end, table, table + 1, visible)) + row_start += rows + if group is not None: + calls.append((group[0], group[0] + group[1], group[2], group[2] + group[3], group[4])) + return calls + def score_pools( self, q: torch.Tensor, @@ -384,6 +520,7 @@ def score_pools( rows_per_request: int = 1, request_ids: torch.Tensor | None = None, max_kv_len: int | None = None, + table_range: tuple[int, int] | None = None, ) -> torch.Tensor: """Score complete pools into FP32 [rows, pool_capacity]. @@ -393,12 +530,34 @@ def score_pools( Incomplete/invisible pools receive the FP32 minimum before top-k. Eager callers that know a host bound on the visible lengths pass max_kv_len so only reachable pools are scored; columns past it are left unwritten and - the top-k must scan only each row's complete pools. + the top-k must scan only each row's complete pools. Packed context rows + with such a bound use the FP8 scorer on SM90 and SM100-family GPUs, whose result + may be a strided view. table_range, from :meth:`plan_context_scores`, + names the context tables those rows' request_ids fall in. """ state = self._cache_state(metadata) tables, gen_lens = self._rows(state, rows_per_request, request_ids=request_ids) capacity = state.block_tables.shape[1] * state.tokens_per_block kpool = self.sparse_params.index_kpool + if self._indexer_fp8 and request_ids is not None and max_kv_len is not None: + # Packed context rows only reference the context requests' tables. + first, last = (0, state.num_contexts) if table_range is None else table_range + scores = kpool_score_fp8( + q, + weights, + state.index_pool, + tables[first:last], + kv_lens, + state.tokens_per_block, + head_dim=self.sparse_params.index_head_dim, + kpool=kpool, + q_scale=q_scale, + w_scale=w_scale, + request_ids=request_ids - first if first else request_ids, + max_kv_len=max_kv_len, + ) + if scores is not None: + return scores return kpool_score( q, weights, @@ -434,7 +593,7 @@ def expand_selection( """Selected pools ``[N, select_k]`` -> latent row ids for the kernel. Expands each valid pool into its members, appends the always-visible - tail, pads with ``-1`` to :attr:`GlmKpoolSparseParams.kernel_output_width` + tail, pads with ``-1`` to the kernel's aligned top-k width and translates positions to latent-cache row ids in one kernel. Same row addressing as :meth:`score_pools`. """ @@ -449,7 +608,7 @@ def expand_selection( base_row=base_row, rows_per_slot=rows_per_slot, kpool=self.sparse_params.index_kpool, - out_width=self.sparse_params.kernel_output_width, + out_width=self._kernel_output_width, request_ids=request_ids, ) @@ -507,7 +666,7 @@ def _dispatch_sparse_core( q_latent[start:stop], selected, local_indices ) return output - pad = (-topk_rows.shape[-1]) % self._KERNEL_TOPK_ALIGN + pad = (-topk_rows.shape[-1]) % self._kernel_topk_align if pad: topk_rows = torch.nn.functional.pad(topk_rows, (0, pad), value=INDEX_SENTINEL) local_heads = q_latent.shape[1] @@ -525,6 +684,15 @@ def _dispatch_sparse_core( q_latent = q_padded if flash_mla_sparse_fwd is None: raise RuntimeError("glm_kpool sparse attention requires FlashMLA in this build") + num_tokens, width = q_latent.shape[0], topk_rows.shape[-1] + if ( + self._split_topk + and num_tokens <= self._SM90_SPLIT_TOPK_MAX_TOKENS + and width > self._kernel_topk_align + ): + return self._split_topk_sparse( + q_latent, kv_rows, topk_rows, width // self._kernel_topk_align, local_heads + ) out, _, _ = flash_mla_sparse_fwd( q_latent, kv_rows, @@ -538,6 +706,31 @@ def _dispatch_sparse_core( out = out[:, :local_heads, :] return out + def _split_topk_sparse( + self, + q_latent: torch.Tensor, + kv_rows: torch.Tensor, + topk_rows: torch.Tensor, + splits: int, + local_heads: int, + ) -> torch.Tensor: + """Sparse attention over ``splits`` equal top-k chunks, merged by LSE. + + Each (token, chunk) pair is one kernel row, so the call spreads over + ``splits`` times as many CTAs; returns ``[T, local_heads, kv_lora]``. + """ + num_tokens, heads, dim = q_latent.shape + chunk = topk_rows.shape[-1] // splits + q_rep = q_latent.unsqueeze(1).expand(num_tokens, splits, heads, dim) + out, _, lse = flash_mla_sparse_fwd( + q_rep.reshape(num_tokens * splits, heads, dim), + kv_rows, + topk_rows.view(num_tokens * splits, 1, chunk), + self.softmax_scale, + self.kv_lora_rank, + ) + return merge_split_attention(out, lse, num_tokens, splits, local_heads) + 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 02297052194f..9f620343931c 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/kernels.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/kernels.py @@ -512,6 +512,75 @@ def kpool_score( return out +# Largest FP8 logits transient (elements) kpool_score_fp8 allocates. +_FP8_SCORE_MAX_ELEMENTS = 1 << 29 + + +def kpool_score_fp8( + q: torch.Tensor, + weights: torch.Tensor, + index_pool: torch.Tensor, + block_tables: torch.Tensor, + kv_lens: torch.Tensor, + tokens_per_block: int, + *, + head_dim: int, + kpool: int, + q_scale: float, + w_scale: float, + request_ids: torch.Tensor, + max_kv_len: int, +) -> torch.Tensor | None: + """FP8 :func:`kpool_score` for packed context rows, via DeepGEMM ``fp8_mqa_logits``. + + The pool keys of every block table are gathered into ``[tables * pools, head_dim]`` + and quantized per row; queries are quantized per (row, head), with their scales + folded into the head weights, so each score is ``sum_h w * relu(q . key)`` as in + :func:`kpool_score`. ``max_kv_len`` bounds every row's ``kv_lens``. Row ``r`` + scores only the first ``kv_lens[r] // kpool`` pools of table ``request_ids[r]``; + later columns are not meaningful, so readers must stop there. The result may be + a strided view with an aligned row stride. Returns ``None`` when the logits + transient would exceed ``_FP8_SCORE_MAX_ELEMENTS``. + """ + from tensorrt_llm.deep_gemm import fp8_mqa_logits + + n = q.shape[0] + pools = max(0, int(max_kv_len)) // kpool + tables = block_tables.shape[0] + if n == 0 or pools == 0 or n * tables * pools > _FP8_SCORE_MAX_ELEMENTS: + return None + positions = torch.arange(pools, device=q.device) * kpool + page = positions // tokens_per_block + offset = positions - page * tokens_per_block + slots = block_tables[:, page].long() # [tables, pools] + keys = index_pool[slots, offset, 2 * head_dim : 3 * head_dim].reshape(-1, head_dim) + half = head_dim // 2 + keys_fp8, key_scale = torch.ops.trtllm.fused_cat_fp8(keys[:, :half], keys[:, half:], False) + q_rows = q.reshape(-1, head_dim) + q_fp8, q_row_scale = torch.ops.trtllm.fused_cat_fp8(q_rows[:, :half], q_rows[:, half:], False) + head_weights = weights.float() * q_row_scale.view(n, -1) * (w_scale * q_scale) + start = request_ids.to(torch.int32) * pools + end = start + (kv_lens // kpool).to(torch.int32) + logits = fp8_mqa_logits( + q_fp8.view(q.shape), + (keys_fp8, key_scale.view(-1)), + head_weights, + start, + end, + clean_logits=False, + ) + if tables == 1: + # DeepGEMM pads the logits rows to an aligned stride, as the pool top-k needs. + return logits[:, :pools] + # Keep the copy's rows aligned and wider than the selection, like the logits arena. + out = torch.empty( + n, max(1024, triton.cdiv(pools, 64) * 64), dtype=torch.float32, device=q.device + ) + rows = torch.arange(n, device=q.device) + out[:, :pools] = logits[:, : tables * pools].view(n, tables, pools)[rows, request_ids.long()] + return out + + @triton.jit def _kpool_expand_kernel( SEL, @@ -683,3 +752,66 @@ def compact_sparse_rows(indices: torch.Tensor) -> tuple[torch.Tensor, torch.Tens triton.next_power_of_2(width), ) return packed, lengths, nonempty + + +@triton.jit +def _merge_split_attention_kernel( + OUT, + LSE, + DST, + stride_out_row, + stride_out_head, + stride_lse_row, + stride_dst_tok, + stride_dst_head, + SPLITS: tl.constexpr, + DIM: tl.constexpr, +): + tok = tl.program_id(0) + head = tl.program_id(1) + cols = tl.arange(0, DIM) + max_lse = float("-inf") + for s in range(SPLITS): + lse = tl.load(LSE + (tok * SPLITS + s) * stride_lse_row + head) + # A chunk without valid rows reports +inf (and a zero output). + lse = tl.where(lse == float("inf"), float("-inf"), lse) + max_lse = tl.maximum(max_lse, lse) + total = 0.0 + acc = tl.zeros([DIM], dtype=tl.float32) + for s in range(SPLITS): + row = tok * SPLITS + s + lse = tl.load(LSE + row * stride_lse_row + head) + lse = tl.where(lse == float("inf"), float("-inf"), lse) + w = tl.where(max_lse == float("-inf"), 0.0, tl.exp(lse - max_lse)) + part = tl.load(OUT + row * stride_out_row + head * stride_out_head + cols).to(tl.float32) + acc += w * part + total += w + acc = tl.where(total > 0.0, acc / total, 0.0) + tl.store( + DST + tok * stride_dst_tok + head * stride_dst_head + cols, acc.to(DST.dtype.element_ty) + ) + + +def merge_split_attention( + out: torch.Tensor, lse: torch.Tensor, num_tokens: int, splits: int, heads: int +) -> torch.Tensor: + """Merge per-chunk attention ``[tokens * splits, H, D]`` by LSE into ``[tokens, heads, D]``. + + ``lse`` is ``[tokens * splits, H]``; chunks without valid rows (LSE ``+inf``) + contribute nothing. Only the first ``heads`` query heads are merged. + """ + dst = torch.empty((num_tokens, heads, out.shape[-1]), dtype=out.dtype, device=out.device) + if num_tokens: + _merge_split_attention_kernel[(num_tokens, heads)]( + out, + lse, + dst, + out.stride(0), + out.stride(1), + lse.stride(0), + dst.stride(0), + dst.stride(1), + SPLITS=splits, + DIM=out.shape[-1], + ) + return dst diff --git a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/params.py b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/params.py index b43a9881d4a1..181b8f6bb6da 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/params.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/params.py @@ -88,10 +88,9 @@ def select_k(self) -> int: """Number of pools a query selects.""" return self.index_topk // self.index_kpool - @property - def kernel_output_width(self) -> int: - """``output_width`` padded to the FlashMLA top-k tile (64).""" - return -(-self.output_width // 64) * 64 + def kernel_output_width(self, topk_align: int) -> int: + """``output_width`` padded to the sparse kernel's top-k alignment.""" + return -(-self.output_width // topk_align) * topk_align @property def output_width(self) -> int: @@ -103,7 +102,7 @@ def output_width(self) -> int: class GlmKpoolBackendForwardArgs(SparseBackendForwardArgs): """Global latent-cache row selection produced by expand_selection. - topk_rows is int32 [tokens, kernel_output_width], with -1 marking invalid + topk_rows is int32 [tokens, kernel_output_width(align)], with -1 marking invalid entries. Request-local topk_indices from the base carrier are unsupported. """ diff --git a/tensorrt_llm/_torch/models/modeling_glm5_next.py b/tensorrt_llm/_torch/models/modeling_glm5_next.py index dcc02d2098ab..b9639016daa4 100644 --- a/tensorrt_llm/_torch/models/modeling_glm5_next.py +++ b/tensorrt_llm/_torch/models/modeling_glm5_next.py @@ -23,6 +23,7 @@ import copy import re +import weakref from collections.abc import Iterable, Sequence from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Callable @@ -58,8 +59,9 @@ from ..modules.linear import Linear, TensorParallelMode from ..modules.multi_stream_utils import maybe_execute_in_parallel from ..modules.rms_norm import RMSNorm +from ..pyexecutor.breakable_cuda_graph import eager_on_graph, is_in_breakable_cuda_graph from ..pyexecutor.config_utils import unwrap_glm5_next_text_config -from ..utils import AuxStreamType +from ..utils import AuxStreamType, get_model_extra_attrs from .checkpoints.hf.glm5_next_weight_mapper import ( Glm5NextHfWeightMapper, glm5_next_is_quantized, @@ -85,6 +87,20 @@ SPARSE_MLP = "sparse" +def _register_graph_layer( + model_config: ModelConfig, registry: str, layer_idx: int, module: nn.Module +) -> str: + """Register ``module`` under a unique key for breakable-CUDA-graph eager ops.""" + layers = model_config.extra_attrs.setdefault(registry, {}) + key = str(layer_idx) + suffix = 0 + while key in layers: + key = f"{layer_idx}_{suffix}" + suffix += 1 + layers[key] = weakref.ref(module) + return key + + @dataclass(frozen=True) class Glm5NextSchedule: """The two literal per-layer dispatch lists, validated against each other.""" @@ -879,6 +895,7 @@ def __init__( mapping: Mapping | None = None, allreduce_strategy: AllReduceStrategy = AllReduceStrategy.ONESHOT, aux_stream: torch.cuda.Stream | None = None, + model_config: ModelConfig | None = None, ) -> None: del dtype # the shared mixer is bf16 (fp32 gate parameters and pools) # The checkpoint's ``linear_attn_config`` does not name the gate rank: @@ -898,6 +915,12 @@ def __init__( allreduce_strategy=allreduce_strategy, aux_stream=aux_stream, ) + # Join the shared mixer's breakable-CUDA-graph registry. model_config is + # not passed to the mixer itself, whose projections keep this module's + # quantization setup. + if model_config is not None: + self.layer_idx_str = _register_graph_layer(model_config, "kda_layers", layer_idx, self) + self.register_to_config = True # Preserve checkpoint FP32 gate parameters. MetaInitMode rejects detach # on meta tensors, so create empty FP32 parameters on that branch. for name in ("A_log", "dt_bias"): @@ -1164,6 +1187,11 @@ def __init__( self.indexer = Glm5NextIndexer( config, layer_idx, mapping=mapping, model_config=model_config ) + self._graph_key = ( + _register_graph_layer(model_config, "glm5_sparse_layers", layer_idx, self) + if model_config is not None + else None + ) # The production sparse-MLA backend, selected through the standard # dispatch (`get_attention_backend(attn_backend, sparse_params)` via @@ -1272,13 +1300,18 @@ def _select_and_attend_paged( reduce: bool = True, request_ids: torch.Tensor | None = None, max_visible: int | None = None, + num_candidates: torch.Tensor | None = None, + context_spans: Sequence[tuple[int, int]] | None = None, ) -> torch.Tensor: """Score pools, select members and tail, then attend to paged latent rows. visible gives each query's prefix length. Packed context uses request_ids; verification uses rows_per_request. Single-token decode reads live lengths from metadata. All device work uses fixed buffer shapes for graph replay. - max_visible is an optional host bound on visible that limits scoring work. + max_visible is an optional host bound on visible that limits scoring work; + num_candidates, if given, is visible // index_kpool as int32. Packed context + passes context_spans, each request's (rows, visible bound), so scoring + runs in calls of bounded memory. """ indexer = self.indexer rows = q_resid.shape[0] @@ -1287,26 +1320,44 @@ def _select_and_attend_paged( # each request contributes several rows (context or verification). plain_decode = request_ids is None and rows_per_request == 1 kv_lens = None if plain_decode else visible - scores = self.attn_backend.score_pools( - q_index, - index_weights, - metadata, - q_scale=indexer.softmax_scale, - w_scale=indexer.head_mix_scale, - kv_lens=kv_lens, - rows_per_request=rows_per_request, - request_ids=request_ids, - max_kv_len=max_visible, - ) - num_cand = (visible // indexer.index_kpool).to(torch.int32) - selected = torch.empty(rows, indexer.select_k, dtype=torch.int32, device=scores.device) - indexer.pool_top_k( - scores, - selected, - is_prefill=False, - sequence_lengths=num_cand, - scan_lengths=num_cand, + num_cand = ( + (visible // indexer.index_kpool).to(torch.int32) + if num_candidates is None + else num_candidates ) + selected = torch.empty(rows, indexer.select_k, dtype=torch.int32, device=q_index.device) + # Each call scores and selects its own rows. + if context_spans is None: + calls = [(slice(None), max_visible, None)] + else: + calls = [ + (slice(start, end), max_kv_len, (first, last)) + for start, end, first, last, max_kv_len in self.attn_backend.plan_context_scores( + metadata, context_spans + ) + ] + for part, max_kv_len, table_range in calls: + scores = self.attn_backend.score_pools( + q_index[part], + index_weights[part], + metadata, + q_scale=indexer.softmax_scale, + w_scale=indexer.head_mix_scale, + kv_lens=None if kv_lens is None else kv_lens[part], + rows_per_request=rows_per_request, + request_ids=None if request_ids is None else request_ids[part], + max_kv_len=max_kv_len, + table_range=table_range, + ) + indexer.pool_top_k( + scores, + selected[part], + is_prefill=False, + sequence_lengths=num_cand[part], + scan_lengths=num_cand[part], + ) + # Release this call's scores before the next call allocates its own. + del scores topk_rows = self.attn_backend.expand_selection( selected, metadata, @@ -1362,10 +1413,14 @@ def forward_prefill( request_ids=rows.final_request_ids, ) query = self.q_b_proj(q_resid).view(-1, self.num_heads, self.qk_head_dim) - max_visible = max( - int(cached_lens[i]) + int(cu_seqlens[i + 1]) - int(cu_seqlens[i]) + spans = [ + ( + int(cu_seqlens[i + 1]) - int(cu_seqlens[i]), + int(cached_lens[i]) + int(cu_seqlens[i + 1]) - int(cu_seqlens[i]), + ) for i in range(len(cu_seqlens) - 1) - ) + ] + max_visible = max(visible for _, visible in spans) return self._select_and_attend_paged( q_resid, query, @@ -1376,6 +1431,7 @@ def forward_prefill( reduce=reduce, request_ids=rows.request_ids, max_visible=max_visible, + context_spans=spans, ) def forward_decode( @@ -1426,12 +1482,25 @@ def forward( Reuse the model's per-forward schedules and row cache across layers. Reduce the complete output once, including context/verify mixtures. + Inside a breakable prefill CUDA graph the phase split runs eagerly + between graph segments over the live rows of the padded batch. """ - ctx = ( - runtime_ctx - if runtime_ctx is not None - else build_glm5_next_runtime_context(attn_metadata) - ) + if is_in_breakable_cuda_graph() and self._graph_key is not None: + out = hidden_states.new_empty((hidden_states.shape[0], self.hidden_size)) + maybe_bcg_glm5_sparse_attention_inplace(hidden_states, self._graph_key, out) + else: + ctx = ( + runtime_ctx + if runtime_ctx is not None + else build_glm5_next_runtime_context(attn_metadata) + ) + out = self._forward_phases(hidden_states, ctx) + return out if self.tp_all_reduce is None else self.tp_all_reduce(out) + + def _forward_phases( + self, hidden_states: torch.Tensor, ctx: Glm5NextRuntimeContext + ) -> torch.Tensor: + """Run each batch phase and return the unreduced ``[tokens, hidden]`` output.""" parts = [] if ctx.num_contexts: parts.append( @@ -1460,8 +1529,7 @@ def forward( reduce=False, ) parts.append(out) - out = parts[0] if len(parts) == 1 else torch.cat(parts, dim=0) - return out if self.tp_all_reduce is None else self.tp_all_reduce(out) + return parts[0] if len(parts) == 1 else torch.cat(parts, dim=0) def forward_verify( self, @@ -1517,6 +1585,27 @@ def forward_verify( ) +def glm5_sparse_attention_inplace( + hidden_states: torch.Tensor, layer_key: str, output: torch.Tensor +) -> None: + """Run one sparse-attention layer's phases on the live rows into ``output``. + + Rows past the live token count (CUDA-graph padding) are zeroed. + """ + extra_attrs = get_model_extra_attrs() + metadata = extra_attrs["attention_metadata"]() + layer = extra_attrs["glm5_sparse_layers"][layer_key]() + num_tokens = metadata.num_tokens + out = layer._forward_phases( + hidden_states[:num_tokens], build_glm5_next_runtime_context(metadata) + ) + output[:num_tokens].copy_(out) + output[num_tokens:].zero_() + + +maybe_bcg_glm5_sparse_attention_inplace = eager_on_graph(glm5_sparse_attention_inplace) + + class Glm5NextGate(DeepseekV3Gate): """DeepSeek noaux_tc routing with FP32 logits and correction bias. @@ -1771,6 +1860,7 @@ def __init__( mapping=mapping, allreduce_strategy=model_config.allreduce_strategy, aux_stream=attention_aux_stream, + model_config=model_config, ) else: self.self_attn = Glm5NextSparseAttention( @@ -1916,6 +2006,10 @@ def forward( raise ValueError("specify exactly one of input_ids or inputs_embeds") if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) + if attn_metadata.padded_num_tokens is not None and not is_in_breakable_cuda_graph(): + # Eager passes over a CUDA-graph-padded batch run only the live rows; + # the speculative wrapper trims the output to the same rows. + inputs_embeds = inputs_embeds[: attn_metadata.num_tokens] if runtime_ctx is None: runtime_ctx = build_glm5_next_runtime_context(attn_metadata) # Only the decoder stack: under one-model MTP the speculative base diff --git a/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py b/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py index 8d81fb9fe6b0..6768cedf784e 100644 --- a/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py +++ b/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py @@ -793,6 +793,7 @@ def prefill_chunk_kda( cu_seqlens: Optional[torch.Tensor], chunk_indices: Optional[torch.Tensor] = None, chunk_size: int = 64, + cu_seqlens_cpu: Optional[torch.Tensor] = None, state_pool: Optional[torch.Tensor] = None, state_indices: Optional[torch.Tensor] = None, varlen_is_aligned: Optional[bool] = None, @@ -814,7 +815,8 @@ def prefill_chunk_kda( ``chunk_indices``, ``varlen_is_aligned``, and ``single_sequence_length`` are prepared once from Kimi K3 runtime metadata. They keep the optimized path from reading ``cu_seqlens`` - back from the GPU in every KDA layer. + back from the GPU in every KDA layer; ``cu_seqlens_cpu`` (the host copy + of ``cu_seqlens``) does the same for the FLA path. """ use_indexed_state = state_pool is not None if use_indexed_state and (initial_state is not None or state_indices is None): @@ -897,6 +899,7 @@ def prefill_chunk_kda( lower_bound=lower_bound, state_v_first=True, cu_seqlens=cu_seqlens, + cu_seqlens_cpu=cu_seqlens_cpu, ) def decode_kda(self, **kwargs) -> torch.Tensor: diff --git a/tensorrt_llm/_torch/modules/kimi_kda/kimi_k3_mamba_metadata.py b/tensorrt_llm/_torch/modules/kimi_kda/kimi_k3_mamba_metadata.py index 3f15c77647f0..0fd65ab9f841 100644 --- a/tensorrt_llm/_torch/modules/kimi_kda/kimi_k3_mamba_metadata.py +++ b/tensorrt_llm/_torch/modules/kimi_kda/kimi_k3_mamba_metadata.py @@ -50,6 +50,12 @@ def __init__(self, max_batch_size: int, chunk_size: int, max_num_tokens: int) -> self.kda_chunk_indices: torch.Tensor | None = None self.kda_varlen_is_aligned: bool | None = None self.kda_single_sequence_length: int | None = None + # Host copy of the context cu_seqlens: FLA's chunked prefill derives its chunk indices from it on the + # CPU instead of reading the counts back from the GPU in every KDA layer. + self.kda_cu_seqlens_cpu: torch.Tensor | None = None + # One device cu_seqlens view per forward: FLA caches its chunk metadata by tensor identity, so every + # KDA layer must pass the same object for the cache to hit. + self.kda_cu_seqlens: torch.Tensor | None = None # KDA's CuTe verify kernel requires a 16-byte-aligned DLPack pointer. # Keep the backing allocation stable so metadata preparation remains # safe for CUDA graph capture and replay. @@ -89,8 +95,14 @@ def prepare(self, attn_metadata: AttentionMetadata) -> None: if not context_lengths: self.kda_chunk_indices = None + self.kda_cu_seqlens_cpu = None + self.kda_cu_seqlens = None return + self.kda_cu_seqlens = self.query_start_loc_long[: len(context_lengths) + 1] + + self.kda_cu_seqlens_cpu = torch.tensor([0] + context_lengths, dtype=torch.long).cumsum_(0) + chunk_counts = [ (length + KDA_PREFILL_CHUNK_SIZE - 1) // KDA_PREFILL_CHUNK_SIZE for length in context_lengths diff --git a/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py b/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py index 30231196346d..926bed5b30be 100644 --- a/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py +++ b/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py @@ -35,6 +35,7 @@ from ...utils import get_model_extra_attrs from ..linear import Linear from ..mamba.causal_conv1d import causal_conv1d_fn +from ..mamba.causal_conv1d_triton import causal_conv1d_update as causal_conv1d_update_triton from ..mamba.fuse_elementwise_ops import extract_transpose_prefill_slice from ..mamba.layernorm_gated import RMSNorm, rms_norm_gated_token_major from ..mamba.recurrent_state_cache import reset_recurrent_state_rows @@ -117,6 +118,21 @@ def _kda_split_conv_sections( ) +def _adjacent_sections_view(*sections: torch.Tensor) -> Optional[torch.Tensor]: + """``torch.cat(sections, -1)`` as a view when the sections already sit back to + back in one buffer (strided slices of a fused projection), else None.""" + first = sections[0] + offset = first.storage_offset() + for section in sections: + if section.stride() != first.stride() or section.storage_offset() != offset: + return None + if section.untyped_storage().data_ptr() != first.untyped_storage().data_ptr(): + return None + offset += section.shape[-1] + shape = (*first.shape[:-1], sum(section.shape[-1] for section in sections)) + return first.as_strided(shape, first.stride(), first.storage_offset()) + + def _kda_expand_fla_conv_cache(conv_state: torch.Tensor) -> torch.Tensor: """Left-pad a production ``W - 1`` cache to FLA's full ``W`` window.""" return torch.nn.functional.pad(conv_state, (1, 0)) @@ -307,6 +323,7 @@ def projection(name: str, in_features: int, out_features: int) -> nn.Module: # Output buffer for the inplace-only ``trtllm::kda_decode`` op. self._o_dense: Optional[torch.Tensor] = None self._packed_conv_weight: Optional[torch.Tensor] = None + self._verify_index_cache: dict = {} self._mtp_conv_weights: Optional[Tuple[torch.Tensor, ...]] = None self.register_to_config = False @@ -331,10 +348,11 @@ def finalize_decode_weights(self) -> None: 2. Kernel-layout constants that ``_decode_via_optimized`` used to rebuild with ~6 device kernels per layer per decode step: transposed conv weights (bf16 ``[W, D]``) and fp32 copies of - ``A_log`` / ``dt_bias`` / ``o_norm.weight``. + ``A_log`` / ``dt_bias`` / ``o_norm.weight``; only built when the + optimized decode kernel is in use. + + The fused projections serve every decode and verify path. """ - if self._dispatch.decode_kernel_path != "optimized": - return if self.q_proj.weight.device.type != "cuda": return with torch.no_grad(): @@ -364,7 +382,8 @@ def finalize_decode_weights(self) -> None: self.qkvg_proj = self._fuse_checkpoint_projections(qkvg_modules) self.qkvg_split_sizes = [module.out_features for module in qkvg_modules] self._bfa_proj_weight = self._merge_projection_weights(bfa_modules, pad_rows_to=8) - self._build_decode_kernel_constants() + if self._dispatch.decode_kernel_path == "optimized": + self._build_decode_kernel_constants() @staticmethod def _fuse_checkpoint_projections(modules: tuple[Linear, ...]) -> Linear: @@ -469,7 +488,9 @@ def _forward_impl( num_tokens = attn_metadata.num_tokens batch_size = attn_metadata.seq_lens.shape[0] state_indices = mamba_metadata.state_indices[:batch_size] - cu_seqlens = mamba_metadata.query_start_loc_long[: num_prefills + 1] + cu_seqlens = getattr(mamba_metadata, "kda_cu_seqlens", None) + if cu_seqlens is None or cu_seqlens.shape[0] != num_prefills + 1: + cu_seqlens = mamba_metadata.query_start_loc_long[: num_prefills + 1] num_decodes = batch_size - num_prefills layer_cache = attn_metadata.kv_cache_manager.mamba_layer_cache(self.layer_idx) @@ -665,6 +686,7 @@ def forward_prefill( chunk_indices = getattr(mamba_metadata, "kda_chunk_indices", None) varlen_is_aligned = getattr(mamba_metadata, "kda_varlen_is_aligned", None) single_sequence_length = getattr(mamba_metadata, "kda_single_sequence_length", None) + cu_seqlens_cpu = getattr(mamba_metadata, "kda_cu_seqlens_cpu", None) from einops import rearrange x = x2d.unsqueeze(0) # [1, T, hidden] @@ -759,6 +781,7 @@ def forward_prefill( safe_gate=lower_bound is not None, lower_bound=lower_bound, cu_seqlens=cu_seqlens, + cu_seqlens_cpu=cu_seqlens_cpu, chunk_indices=chunk_indices, state_pool=kernel_pool if use_indexed_state else None, state_indices=kernel_indices if use_indexed_state else None, @@ -1398,6 +1421,20 @@ def forward_verify_sequential( "speculative verification requires the cache manager's " "SpeculativeState (legacy intermediate-buffer path)" ) + if self._packed_conv_weight is not None and intermediate_ssm.dtype in ( + torch.float32, + ssm_pool.dtype, + ): + return self._forward_verify_batched( + x2d, + num_steps, + intermediate_conv, + intermediate_ssm, + conv_pool, + ssm_pool, + slot_indices, + output=output, + ) d = self.proj_size num_generations = x2d.shape[0] // num_steps @@ -1473,6 +1510,115 @@ def forward_verify_sequential( # Can be removed once attention writes the core in place. return self._store_core(core, output) + def _forward_verify_batched( + self, + x2d, + num_steps, + intermediate_conv, + intermediate_ssm, + conv_pool, + ssm_pool, + slot_indices, + output: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """All verify steps in one convolution and one recurrent launch. + + Same contract as the per-step loop of :meth:`forward_verify_sequential`: + the live pools are only read, and step ``t``'s convolution window and + recurrent state land in ``intermediate_conv[:B, t]`` and + ``intermediate_ssm[:B, t]`` for ``update_mamba_states()``. + """ + from einops import rearrange + from fla.ops.kda.fused_recurrent import fused_recurrent_kda_fwd + + d = self.proj_size + num_generations = x2d.shape[0] // num_steps + num_tokens = num_generations * num_steps + device = x2d.device + x = x2d.view(num_generations, num_steps, -1) # [B, T, hidden] + + projections = self._project_verify_inputs(x, x2d.shape[0], preserve_section_strides=True) + if projections is None: + qkv = torch.cat((self.q_proj(x), self.k_proj(x), self.v_proj(x)), dim=-1) + g = self.f_b_proj(self.f_a_proj(x)) + beta = self.b_proj(x) + onorm_g = None + else: + q_proj_states, k_proj_states, v_proj_states, g, beta, onorm_g = projections + qkv = _adjacent_sections_view(q_proj_states, k_proj_states, v_proj_states) + if qkv is None: + qkv = torch.cat((q_proj_states, k_proj_states, v_proj_states), dim=-1) + if onorm_g is not None: + onorm_g = onorm_g.contiguous() + + batch_rows, state_rows, cu_seqlens = self._verify_index_tensors( + num_generations, num_steps, intermediate_ssm.shape[1], device + ) + # The convolution kernel slides the window it is given, so it runs on a + # gathered copy of the rows; per-step windows go to intermediate_conv. + conv_rows = conv_pool.index_select(0, slot_indices.long()) + conv_out = causal_conv1d_update_triton( + qkv.transpose(1, 2), # [B, 3D, T] + conv_rows, + self._packed_conv_weight, + None, + "silu", + conv_state_indices=batch_rows, + intermediate_conv_window=intermediate_conv, + intermediate_state_indices=batch_rows, + ) + q, k, v = fused_kda_post_conv( + conv_out.transpose(1, 2).reshape(num_tokens, 3 * d).t(), + num_heads=self.num_heads, + head_dim=self.head_dim, + ) + + # Step t of request b reads its state from row (b, 0), seeded with the + # pool state, and writes it to row (b, t) of intermediate_ssm. + intermediate_ssm[:num_generations, 0].copy_(ssm_pool.index_select(0, slot_indices.long())) + o, _ = fused_recurrent_kda_fwd( + q=q, + k=k, + v=v, + g=rearrange(g, "b t (h d) -> 1 (b t) h d", d=self.head_dim).contiguous(), + beta=beta.reshape(1, num_tokens, -1).contiguous(), + A_log=self.A_log, + dt_bias=self.dt_bias, + initial_state=intermediate_ssm.view(-1, *intermediate_ssm.shape[2:]), + scale=self.head_k_dim**-0.5, + inplace_final_state=True, + state_v_first=True, + cu_seqlens=cu_seqlens, + ssm_state_indices=state_rows, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + lower_bound=self.gate_lower_bound, + ) + core = self._output_gate(x, o.view(num_generations, num_steps, *o.shape[2:]), onorm_g) + # Can be removed once attention writes the core in place. + return self._store_core(core, output) + + def _verify_index_tensors( + self, num_generations: int, num_steps: int, t_max: int, device: torch.device + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Batch rows, ``[B, T]`` intermediate-state rows and ``cu_seqlens`` for + :meth:`_forward_verify_batched`, built once per shape so CUDA graph + replays do not rerun them.""" + key = (num_generations, num_steps, t_max, device) + cached = self._verify_index_cache.get(key) + if cached is None: + batch_rows = torch.arange(num_generations, device=device, dtype=torch.int32) + state_rows = ( + batch_rows[:, None] * t_max + + torch.arange(num_steps, device=device, dtype=torch.int32)[None, :] + ) + cu_seqlens = torch.arange( + 0, num_generations * num_steps + 1, num_steps, device=device, dtype=torch.long + ) + cached = (batch_rows, state_rows, cu_seqlens) + self._verify_index_cache[key] = cached + return cached + def _output_gate( self, x: torch.Tensor, o: torch.Tensor, onorm_g: Optional[torch.Tensor] = None ) -> torch.Tensor: diff --git a/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py b/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py index 5447d6b51afe..c37f905bb036 100644 --- a/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py +++ b/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py @@ -333,9 +333,7 @@ def get_valid_tactics( for bs in _BIGFUSE_BLOCK_SIZE_OPTIONS: tactics.append(("fma", 8, 2, bs)) - # V4 keeps its DeepGEMM tactics; the lagged (V4.1) entry boundary is pinned to the - # in-tree FMA kernel (no MMA variant exists for the pre-mapping op). - if _get_dg_fn() is not None and kwargs.get("pre_mix_ext") is None: + if _get_dg_fn() is not None: for bs in _BIGFUSE_BLOCK_SIZE_OPTIONS: tactics.append(("dg_splitk", 0, 0, bs)) tactics.append(("dg_nosplit", 0, 0, bs)) @@ -357,8 +355,6 @@ def forward(self, inputs: List[torch.Tensor], *, tactic: Any = -1, **kwargs) -> if tactic == -1: tactic = _FALLBACK_TACTIC backend, tile_n, tile_m, bigfuse_bs = tactic - if kwargs.get("pre_mix_ext") is not None and backend != "fma": - backend, tile_n, tile_m, bigfuse_bs = _FALLBACK_TACTIC num_splits = 1 if backend == "dg_splitk": @@ -825,11 +821,18 @@ def _alloc_fused_hc_scratch( _FUSED_HC_FALLBACK_TACTIC_FMA = ("fused_half_fma", 2, 1, 256, 1) -def _get_fused_hc_fallback_tactic(hidden_size: int | None = None): +def _get_fused_hc_fallback_tactic(hidden_size: int | None = None, num_tokens: int = 0): mma_ok = _fused_hc_mma_supported() if hidden_size is not None: mma_ok = mma_ok and hidden_size in _FUSED_HC_MMA_SUPPORTED_HIDDEN_SIZES - return _FUSED_HC_FALLBACK_TACTIC_MMA if mma_ok else _FUSED_HC_FALLBACK_TACTIC_FMA + if mma_ok: + return _FUSED_HC_FALLBACK_TACTIC_MMA + # Token counts above the tuned buckets (prefill chunks past + # _MHC_FINE_GRAINED_TUNING_LIMIT) land here; the FMA GEMM is several times + # slower than DeepGEMM's TF32 GEMM there. + if num_tokens > 32 and _get_dg_fn() is not None: + return ("split_dg", 0, 1 if num_tokens >= 4096 else _DG_NUM_SPLITS, 256, 1) + return _FUSED_HC_FALLBACK_TACTIC_FMA class MhcFusedHcRunner(TunableRunner): @@ -958,6 +961,13 @@ def add(tactic): for tn, ks in _FUSED_HC_HALF_FMA_TN_KS: if ks == 1: add(("fused_half_fma", tn, ks, bs, 1)) + # Without tcgen05, large M runs the boundary unfused: post_mapping, + # then the DeepGEMM TF32 GEMM + sqrsum (SM90 tensor cores) and + # big_fuse. num_k_splits carries the DeepGEMM split-K count. + if _get_dg_fn() is not None: + for splits in (1, _DG_NUM_SPLITS): + for bs in _FUSED_HC_BIGFUSE_BS: + add(("split_dg", 0, splits, bs, 1)) return tactics @@ -982,9 +992,8 @@ def forward(self, inputs, *, tactic=-1, **kwargs): pre_mix_ext = kwargs.get("pre_mix_ext") if tactic == -1: - tactic = _get_fused_hc_fallback_tactic(self.hidden_size) + tactic = _get_fused_hc_fallback_tactic(self.hidden_size, residual_prev.shape[0]) backend, tile_n, num_k_splits, bigfuse_bs, tile_m = tactic - backend_code = _FUSED_HC_BACKEND_CODE[backend] # Fused next-layer RMSNorm on layer_input. All four backends support # it (Path B/E inside mhcBigFuseKernel, Path D/F inside the all-in-one @@ -997,6 +1006,23 @@ def forward(self, inputs, *, tactic=-1, **kwargs): if norm_weight is not None and not norm_weight.is_contiguous(): norm_weight = norm_weight.contiguous() + if backend == "split_dg": + return self._forward_split_dg( + x_prev, + residual_prev, + post_mix_prev, + comb_mix_prev, + w_t_cur, + hc_scale_cur, + hc_base_cur, + num_k_splits, + bigfuse_bs, + norm_weight, + norm_eps, + pre_mix_ext, + ) + backend_code = _FUSED_HC_BACKEND_CODE[backend] + B = residual_prev.shape[0] residual_cur = torch.empty_like(residual_prev) post_mix_cur = torch.empty((B, self.n), dtype=torch.float32, device=x_prev.device) @@ -1055,6 +1081,74 @@ def forward(self, inputs, *, tactic=-1, **kwargs): return residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur, pre_mix_out return residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur + def _forward_split_dg( + self, + x_prev, + residual_prev, + post_mix_prev, + comb_mix_prev, + w_t_cur, + hc_scale_cur, + hc_base_cur, + num_splits, + bigfuse_bs, + norm_weight, + norm_eps, + pre_mix_ext, + ): + """post_mapping, then pre_mapping's DeepGEMM GEMM + sqrsum and big_fuse.""" + n, hidden = self.n, self.hidden_size + B = residual_prev.shape[0] + pre_mix_out = None + if pre_mix_ext is not None: + if pre_mix_ext.numel() != B * n: + # Autotuner bucket profiling (see MhcPreMappingRunner.forward). + pre_mix_ext = torch.ones((B, n), dtype=torch.float32, device=x_prev.device) + pre_mix_ext = pre_mix_ext.reshape(B, n).to(torch.float32).contiguous() + pre_mix_out = torch.empty((B, n), dtype=torch.float32, device=x_prev.device) + K = n * hidden + residual_cur = mhc_post_mapping_cuda( + residual_prev.view(B, n, hidden), + x_prev, + post_mix_prev.view(B, n), + comb_mix_prev.view(B, n, n), + n, + ) + y_acc, r_acc, num_splits = mhc_gemm_rms_dg_cuda( + residual_cur.view(B, K), w_t_cur, B, w_t_cur.shape[0], K, num_splits=num_splits + ) + post_mix_cur = torch.empty((B, n), dtype=torch.float32, device=x_prev.device) + comb_mix_cur = torch.empty((B, n * n), dtype=torch.float32, device=x_prev.device) + layer_input_cur = torch.empty_like(x_prev) + mhc_big_fuse_cuda( + y_acc, + r_acc, + residual_cur, + hc_scale_cur, + hc_base_cur, + post_mix_cur, + comb_mix_cur, + layer_input_cur, + B, + K, + hidden, + self.rms_eps, + self.hc_pre_eps, + self.hc_sinkhorn_eps, + self.hc_post_mult_value, + self.sinkhorn_repeat, + num_splits=num_splits, + block_size=bigfuse_bs, + norm_weight=norm_weight, + norm_eps=norm_eps, + pre_mix_ext=pre_mix_ext, + pre_mix_out=pre_mix_out, + ) + residual_cur = residual_cur.view(B, n, hidden) + if pre_mix_out is not None: + return residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur, pre_mix_out + return residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur + # Process-wide runner cache keyed on the mHC configuration. Avoids recreating # a MhcFusedHcRunner on every call. diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cutlass.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cutlass.py index 8d78c8c57130..b6f3b0843ab3 100755 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cutlass.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cutlass.py @@ -28,7 +28,7 @@ from ...utils import (ActivationType, AuxStreamType, EventType, Fp4QuantizedTensor) from .activation import (DEFAULT_MOE_ACTIVATION, ActivationParamShape, - MoEActivation, MoEActivationSupport) + MoEActivation, MoEActivationSupport, SwigluActivation) from .impl_base import MoEImplBase, apply_moe_impl_construction_state from .impl_contract import (MoEDeployment, MoEEligibility, MoEInputRequirement, MoEProblem, MoERejectReason, MoERunContext, @@ -65,6 +65,11 @@ def raise_moe_lora_multichunk_unsupported(num_chunks: int) -> None: f"increase `moe_max_num_tokens` so the MoE runs in a single chunk.") +# Largest batch (tokens on this rank) that runs the SM90 FP8-activation Triton +# MoE instead of the CUTLASS / DeepGEMM grouped GEMM. +_SM90_TRITON_FP8_MOE_MAX_TOKENS = 32 + + class CutlassFusedMoE(MoEImplBase): """ Fused Mixture of Experts (MoE) Layer with performance tuning. @@ -973,6 +978,32 @@ def _tuner_shapes( tuner_num_tokens = ctx.x.shape[0] * self.mapping.tp_size return tuner_num_tokens, self.routing_method.top_k + def _triton_fp8_act_swiglu( + self, + x: torch.Tensor, + x_sf: Optional[torch.Tensor], + lora_params, + moe_output: Optional[torch.Tensor], + ) -> Optional[Tuple[Optional[float], bool]]: + """SwiGLU clamp settings when this call can take the SM90 FP8-activation Triton path. + + Returns ``(limit, clamp_after_silu)``, or None to keep the CUTLASS path: + FP8 block-scale weights, a plain SwiGLU with at most a scalar clamp, BF16 + input, a decode-sized batch, and no LoRA or caller-provided output buffer. + """ + if (not self.has_deepseek_fp8_block_scales + or self.activation_type != ActivationType.Swiglu + or not isinstance(self.activation, SwigluActivation) + or x_sf is not None or x.dtype != torch.bfloat16 + or x.shape[0] > _SM90_TRITON_FP8_MOE_MAX_TOKENS + or lora_params is not None or moe_output is not None): + return None + limit = self.activation.clamp + if isinstance(limit, torch.Tensor): + return None + return (None if limit is None else float(limit), + self.activation.clamp_after_silu) + def run_moe( self, ctx: MoERunContext, @@ -1016,9 +1047,17 @@ def run_moe( # SM120 + FP8 block scales: use Triton kernel (CUTLASS TMA fails on SM120 # for large token counts due to cuTensorMapEncodeTiled limitations). - if self.has_deepseek_fp8_block_scales and get_sm_version() == 120: - from .fused_moe_triton_fp8_block_scale import \ - run_triton_fp8_block_scale_moe + # SM90 + FP8 block scales, decode-sized batches: the Triton FP8-activation + # kernels launch only the expert blocks the batch routes to, where the + # grouped DeepGEMM kernel walks every expert on one persistent CTA per SM. + sm_version = get_sm_version() + triton_fp8_act_swiglu = (self._triton_fp8_act_swiglu( + x, x_sf, lora_params, moe_output) if sm_version == 90 else None) + if self.has_deepseek_fp8_block_scales and ( + sm_version == 120 or triton_fp8_act_swiglu is not None): + from .fused_moe_triton_fp8_block_scale import ( + run_triton_fp8_block_scale_moe, + run_triton_fp8_block_scale_moe_fp8_act) # forward_chunk sets token_final_scales=None when # apply_router_weight_on_input=True (weights already folded into x); @@ -1032,7 +1071,11 @@ def run_moe( # (0 .. expert_size_per_partition-1), so remap and zero-scale any # non-local token-expert pairs to suppress their contribution. local_n = self.expert_size_per_partition - if enable_alltoall: + if self.ep_size == 1: + # Every routed expert is local: the IDs and weights need no remap. + local_ids = token_selected_experts + is_local = None + elif enable_alltoall: # After alltoall dispatch, IDs are already local; padding = local_n local_ids = token_selected_experts.clamp(0, local_n - 1) is_local = token_selected_experts < local_n @@ -1042,8 +1085,23 @@ def run_moe( 0, local_n - 1) is_local = ((token_selected_experts >= slot_start) & (token_selected_experts < slot_start + local_n)) - local_scales = token_final_scales * is_local.to( - token_final_scales.dtype) + local_scales = (token_final_scales + if is_local is None else token_final_scales * + is_local.to(token_final_scales.dtype)) + if triton_fp8_act_swiglu is not None: + swiglu_limit, swiglu_clamp_after_silu = triton_fp8_act_swiglu + return run_triton_fp8_block_scale_moe_fp8_act( + x, + local_ids, + local_scales, + self.w3_w1_weight, + self.quant_scales.fc_weight_scales, + self.w2_weight, + self.quant_scales.proj_weight_scales, + swiglu_limit=swiglu_limit, + swiglu_clamp_after_silu=swiglu_clamp_after_silu, + output_dtype=output_dtype, + ) result = run_triton_fp8_block_scale_moe( x, local_ids, diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_triton_fp8_block_scale.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_triton_fp8_block_scale.py index 7d837e9576e0..003f141e721b 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_triton_fp8_block_scale.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_triton_fp8_block_scale.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright 2026 NVIDIA Corporation. All rights reserved. """ -Triton FP8 block-scale MoE forward pass for SM120. +Triton FP8 block-scale MoE forward passes: BF16 activations for SM120, FP8 +activations for decode-sized batches on SM90. """ import torch @@ -40,15 +41,17 @@ def _moe_prefix_kernel( num_post_pad_ptr, BLOCK_SIZE: tl.constexpr, NUM_EXPERTS: tl.constexpr, + NUM_EXPERTS_POW2: tl.constexpr, ): """Compute exclusive offsets from padded expert counts.""" - e = tl.arange(0, NUM_EXPERTS) - counts = tl.load(counts_ptr + e).to(tl.int32) + e = tl.arange(0, NUM_EXPERTS_POW2) + mask = e < NUM_EXPERTS + counts = tl.load(counts_ptr + e, mask=mask, other=0).to(tl.int32) padded = ((counts + BLOCK_SIZE - 1) // BLOCK_SIZE) * BLOCK_SIZE # Inclusive scan; subtract self for exclusive prefix. inc = tl.associative_scan(padded, 0, _int32_add) out_off = inc - padded - tl.store(out_off_ptr + e, out_off) + tl.store(out_off_ptr + e, out_off, mask=mask) tl.store(num_post_pad_ptr, tl.sum(padded).to(tl.int32)) @@ -139,6 +142,7 @@ def moe_align_block_size( num_post_pad, BLOCK_SIZE=block_size, NUM_EXPERTS=num_experts, + NUM_EXPERTS_POW2=triton.next_power_of_2(num_experts), ) _moe_expert_ids_kernel[num_experts,]( @@ -601,3 +605,451 @@ def run_triton_fp8_block_scale_moe( if output_dtype is not None and output.dtype != output_dtype: output = output.to(output_dtype) return output + + +# --------------------------------------------------------------------------- +# FP8-activation path for decode-sized batches (SM90) +# --------------------------------------------------------------------------- +_FP8_GROUP_SIZE = 128 +_FP8_E4M3_MAX = 448.0 +# Largest (token, expert) pair count aligned by the single-program kernel. +_SMALL_ALIGN_MAX_PAIRS = 128 + + +@triton.jit +def _per_token_group_quant_fp8_kernel( + x_ptr, + q_ptr, + s_ptr, + stride_xm, + stride_qm, + stride_sm, + GROUP: tl.constexpr, + FP8_MAX: tl.constexpr, +): + """Quantize one 1 x GROUP slice of a row to FP8 with its own scale.""" + row = tl.program_id(0).to(tl.int64) + group = tl.program_id(1) + offs = group * GROUP + tl.arange(0, GROUP) + x = tl.load(x_ptr + row * stride_xm + offs).to(tl.float32) + scale = tl.maximum(tl.max(tl.abs(x)), 1e-10) / FP8_MAX + tl.store(q_ptr + row * stride_qm + offs, (x / scale).to(tl.float8e4nv)) + tl.store(s_ptr + row * stride_sm + group, scale) + + +@triton.jit +def _swiglu_quant_fp8_kernel( + ic1_ptr, + q_ptr, + s_ptr, + INTER, + stride_im, + stride_qm, + stride_sm, + limit, + HAS_LIMIT: tl.constexpr, + CLAMP_AFTER_SILU: tl.constexpr, + GROUP: tl.constexpr, + FP8_MAX: tl.constexpr, +): + """SwiGLU over one 1 x GROUP slice of an (up | gate) FC1 row, quantized to FP8. + + With a limit, ``up`` is clamped to [-limit, limit] and ``gate`` to at most + ``limit`` before SiLU (or the SiLU output, with CLAMP_AFTER_SILU), as the + CUTLASS MoE activation does. + """ + row = tl.program_id(0).to(tl.int64) + group = tl.program_id(1) + offs = group * GROUP + tl.arange(0, GROUP) + up = tl.load(ic1_ptr + row * stride_im + offs).to(tl.float32) + gate = tl.load(ic1_ptr + row * stride_im + INTER + offs).to(tl.float32) + if HAS_LIMIT: + up = tl.minimum(tl.maximum(up, -limit), limit) + if CLAMP_AFTER_SILU: + act = tl.minimum(gate * tl.sigmoid(gate), limit) + else: + gate = tl.minimum(gate, limit) + act = gate * tl.sigmoid(gate) + else: + act = gate * tl.sigmoid(gate) + y = act * up + scale = tl.maximum(tl.max(tl.abs(y)), 1e-10) / FP8_MAX + tl.store(q_ptr + row * stride_qm + offs, (y / scale).to(tl.float8e4nv)) + tl.store(s_ptr + row * stride_sm + group, scale) + + +@triton.jit +def _fused_moe_fp8act_fp8w_kernel( + a_ptr, + a_scale_ptr, + b_ptr, + b_scale_ptr, + c_ptr, + topk_weights_ptr, + sorted_token_ids_ptr, + expert_ids_ptr, + num_tokens_post_padded_ptr, + N, + K, + num_valid_tokens, + stride_am, + stride_ak, + stride_asm, + stride_ask, + stride_be, + stride_bn, + stride_bk, + stride_bse, + stride_bsn, + stride_bsk, + stride_cm, + stride_cn, + group_n: tl.constexpr, + group_k: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + top_k: tl.constexpr, + MUL_ROUTED_WEIGHT: tl.constexpr, +): + """Fused MoE GEMM with FP8 activations (1 x 128 scales) and FP8 weights (128 x 128 scales). + + One program per (expert block of BLOCK_SIZE_M sorted rows, BLOCK_SIZE_N + output columns); C row ``r`` holds the (token, expert) slot ``r``. + C[m, n] = sum_k dot(A_fp8[m, k], B_fp8[k, n]) * A_scale[m, k] * B_scale[n, k] + """ + pid = tl.program_id(0) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + pid_m = pid // num_pid_n + pid_n = pid % num_pid_n + if pid_m * BLOCK_SIZE_M >= tl.load(num_tokens_post_padded_ptr): + return + + offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64) + offs_token = tl.load(sorted_token_ids_ptr + offs_token_id).to(tl.int64) + token_mask = offs_token < num_valid_tokens + off_expert = tl.load(expert_ids_ptr + pid_m).to(tl.int64) + + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_row = offs_token // top_k + a_ptrs = a_ptr + a_row[:, None] * stride_am + offs_k[None, :] * stride_ak + b_ptrs = ( + b_ptr + off_expert * stride_be + offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn + ) + a_scale_ptrs = a_scale_ptr + a_row * stride_asm + b_scale_ptrs = b_scale_ptr + off_expert * stride_bse + (offs_bn // group_n) * stride_bsn + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + a = tl.load(a_ptrs, mask=token_mask[:, None], other=0.0) + b = tl.load(b_ptrs) + offs_ks = (k * BLOCK_SIZE_K) // group_k + a_scale = tl.load(a_scale_ptrs + offs_ks * stride_ask, mask=token_mask, other=0.0) + b_scale = tl.load(b_scale_ptrs + offs_ks * stride_bsk) + accumulator += tl.dot(a, b) * a_scale[:, None] * b_scale[None, :] + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + if MUL_ROUTED_WEIGHT: + moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0.0) + accumulator *= moe_weight[:, None] + + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :] + c_mask = token_mask[:, None] & (offs_cn[None, :] < N) + tl.store(c_ptrs, accumulator.to(tl.bfloat16), mask=c_mask) + + +@triton.jit +def _moe_align_small_kernel( + flat_ids_ptr, + sorted_token_ids_ptr, + expert_ids_ptr, + num_post_pad_ptr, + expert_offsets_ptr, + num_pairs, + num_slots, + NUM_PAIRS_POW2: tl.constexpr, + NUM_SLOTS_POW2: tl.constexpr, + NUM_EXPERTS_POW2: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + """moe_align_block_size in one program for small batches. + + A pair's slot is its expert's padded offset (staged in ``expert_offsets``) + plus the number of earlier pairs routed to the same expert, so the pairwise + comparison replaces the atomic scatter of the multi-kernel version. + """ + i = tl.arange(0, NUM_PAIRS_POW2) + valid = i < num_pairs + # Padding pairs go to the last bin, which no real expert uses. + ids = tl.load(flat_ids_ptr + i, mask=valid, other=NUM_EXPERTS_POW2 - 1).to(tl.int32) + e = tl.arange(0, NUM_EXPERTS_POW2) + counts = tl.histogram(ids, NUM_EXPERTS_POW2) + padded = ((counts + BLOCK_SIZE - 1) // BLOCK_SIZE) * BLOCK_SIZE + padded = tl.where(e == NUM_EXPERTS_POW2 - 1, 0, padded) + tl.store(expert_offsets_ptr + e, tl.cumsum(padded, axis=0) - padded) + tl.debug_barrier() + pair_offset = tl.load(expert_offsets_ptr + ids) + earlier = (ids[None, :] == ids[:, None]) & (i[None, :] < i[:, None]) & valid[None, :] + slot = pair_offset + tl.sum(earlier.to(tl.int32), axis=1) + + pad = tl.arange(0, NUM_SLOTS_POW2) + tl.store( + sorted_token_ids_ptr + pad, + tl.full([NUM_SLOTS_POW2], 0, tl.int32) + num_pairs, + mask=pad < num_slots, + ) + tl.debug_barrier() + tl.store(sorted_token_ids_ptr + slot, i, mask=valid) + tl.store(expert_ids_ptr + slot // BLOCK_SIZE, ids, mask=valid) + tl.store(num_post_pad_ptr, tl.sum(padded)) + + +def _moe_align_small( + topk_ids: torch.Tensor, block_size: int, num_experts: int, num_m_blocks: int +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Same outputs as :func:`moe_align_block_size` for the first ``num_m_blocks`` blocks.""" + device = topk_ids.device + flat_ids = topk_ids.reshape(-1) + num_pairs = flat_ids.shape[0] + num_slots = num_m_blocks * block_size + sorted_token_ids = torch.empty((num_slots,), dtype=torch.int32, device=device) + expert_ids = torch.empty((num_m_blocks,), dtype=torch.int32, device=device) + num_post_pad = torch.empty(1, dtype=torch.int32, device=device) + num_experts_pow2 = triton.next_power_of_2(num_experts + 1) + expert_offsets = torch.empty((num_experts_pow2,), dtype=torch.int32, device=device) + _moe_align_small_kernel[(1,)]( + flat_ids, + sorted_token_ids, + expert_ids, + num_post_pad, + expert_offsets, + num_pairs, + num_slots, + NUM_PAIRS_POW2=triton.next_power_of_2(max(num_pairs, 2)), + NUM_SLOTS_POW2=triton.next_power_of_2(max(num_slots, 2)), + NUM_EXPERTS_POW2=num_experts_pow2, + BLOCK_SIZE=block_size, + num_warps=4, + ) + return sorted_token_ids, expert_ids, num_post_pad + + +@triton.jit +def _moe_topk_sum_kernel( + in_ptr, + out_ptr, + hidden, + TOP_K: tl.constexpr, + BLOCK_H: tl.constexpr, +): + """out[t, :] = sum_k in[t * TOP_K + k, :] in FP32.""" + t = tl.program_id(0).to(tl.int64) + offs = tl.program_id(1) * BLOCK_H + tl.arange(0, BLOCK_H) + mask = offs < hidden + acc = tl.zeros([BLOCK_H], dtype=tl.float32) + for k in tl.static_range(TOP_K): + acc += tl.load(in_ptr + (t * TOP_K + k) * hidden + offs, mask=mask, other=0.0).to( + tl.float32 + ) + tl.store(out_ptr + t * hidden + offs, acc.to(out_ptr.dtype.element_ty), mask=mask) + + +def _per_token_group_quant_fp8(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """(M, K) BF16 -> (M, K) FP8 and (M, K // 128) FP32 scales.""" + M, K = x.shape + assert K % _FP8_GROUP_SIZE == 0, f"K={K} must be a multiple of {_FP8_GROUP_SIZE}" + q = torch.empty((M, K), dtype=torch.float8_e4m3fn, device=x.device) + scales = torch.empty((M, K // _FP8_GROUP_SIZE), dtype=torch.float32, device=x.device) + if M > 0: + _per_token_group_quant_fp8_kernel[(M, K // _FP8_GROUP_SIZE)]( + x, + q, + scales, + x.stride(0), + q.stride(0), + scales.stride(0), + GROUP=_FP8_GROUP_SIZE, + FP8_MAX=_FP8_E4M3_MAX, + ) + return q, scales + + +def _invoke_fp8act_fp8w_moe_kernel( + a: torch.Tensor, + a_scale: torch.Tensor, + b: torch.Tensor, + b_scale: torch.Tensor, + c: torch.Tensor, + topk_weights: torch.Tensor | None, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_padded: torch.Tensor, + num_m_blocks: int, + num_valid_tokens: int, + top_k: int, +) -> None: + """Launch `_fused_moe_fp8act_fp8w_kernel` over at most ``num_m_blocks`` expert blocks. + + b: (E, N, K) FP8, b_scale: (E, N // 128, K // 128) FP32, c: (num_valid_tokens, N) BF16. + """ + N, K = b.shape[1], b.shape[2] + # Measured on H200 at GLM-5.3-Flash TP4 shapes: narrower N tiles feed more SMs when few + # expert blocks stream a long K. + if K >= 2048 and num_m_blocks <= 32: + block_n, num_warps, num_stages = 64, 4, 4 + else: + block_n, num_warps, num_stages = 128, 4, 3 + grid = (num_m_blocks * triton.cdiv(N, block_n),) + _fused_moe_fp8act_fp8w_kernel[grid]( + a, + a_scale, + b, + b_scale, + c, + topk_weights, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + N, + K, + num_valid_tokens, + a.stride(0), + a.stride(1), + a_scale.stride(0), + a_scale.stride(1), + b.stride(0), + b.stride(1), + b.stride(2), + b_scale.stride(0), + b_scale.stride(1), + b_scale.stride(2), + c.stride(0), + c.stride(1), + group_n=_BLOCK_SHAPE[1], + group_k=_BLOCK_SHAPE[0], + BLOCK_SIZE_M=_BLOCK_SIZE_M, + BLOCK_SIZE_N=block_n, + BLOCK_SIZE_K=_BLOCK_SHAPE[0], + top_k=top_k, + MUL_ROUTED_WEIGHT=topk_weights is not None, + num_warps=num_warps, + num_stages=num_stages, + ) + + +def run_triton_fp8_block_scale_moe_fp8_act( + x: torch.Tensor, + token_selected_experts: torch.Tensor, + token_final_scales: torch.Tensor, + w3_w1: torch.Tensor, + w3_w1_scales: torch.Tensor, + w2: torch.Tensor, + w2_scales: torch.Tensor, + swiglu_limit: float | None = None, + swiglu_clamp_after_silu: bool = False, + output_dtype: torch.dtype | None = None, +) -> torch.Tensor: + """SwiGLU MoE forward with FP8 activations for decode-sized batches. + + Same weight layout as :func:`run_triton_fp8_block_scale_moe` + (DeepSeekFP8BlockScalesFusedMoEMethod without E8M0 resmoothing). Activations + are quantized per 1 x 128 group, and each GEMM launches only as many expert + blocks as the batch can occupy, so the cost follows the routed experts + rather than the expert count. + + Args: + x: (T, H) BF16 input activations. + token_selected_experts: (T, K) expert IDs local to ``w3_w1``. + token_final_scales: (T, K) F32 routing weights. + w3_w1: (E, 2I, H) FP8 up+gate projection weights. + w3_w1_scales: (E, 2I // 128, H // 128) F32 block scales. + w2: (E, H, I) FP8 down-projection weights. + w2_scales: (E, H // 128, I // 128) F32 block scales. + swiglu_limit: Optional SwiGLU clamp. + swiglu_clamp_after_silu: Clamp the SiLU output instead of the gate. + output_dtype: Optional output dtype. + + Returns: + (T, H) output tensor. + """ + num_tokens, hidden = x.shape + num_experts, gate_up_size, _ = w3_w1.shape + intermediate = gate_up_size // 2 + top_k = token_selected_experts.shape[1] + num_pairs = num_tokens * top_k + device = x.device + + # Every expert block holds at least one (token, expert) pair. + num_m_blocks = min(num_pairs, num_pairs // _BLOCK_SIZE_M + num_experts) + if num_pairs <= _SMALL_ALIGN_MAX_PAIRS: + sorted_token_ids, expert_ids, num_tokens_post_padded = _moe_align_small( + token_selected_experts, _BLOCK_SIZE_M, num_experts, num_m_blocks + ) + else: + sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( + token_selected_experts, _BLOCK_SIZE_M, num_experts + ) + + x_q, x_scales = _per_token_group_quant_fp8(x.contiguous()) + ic1 = torch.empty((num_pairs, gate_up_size), dtype=torch.bfloat16, device=device) + _invoke_fp8act_fp8w_moe_kernel( + x_q, + x_scales, + w3_w1, + w3_w1_scales, + ic1, + None, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + num_m_blocks, + num_pairs, + top_k, + ) + + assert intermediate % _FP8_GROUP_SIZE == 0 + act_q = torch.empty((num_pairs, intermediate), dtype=torch.float8_e4m3fn, device=device) + act_scales = torch.empty( + (num_pairs, intermediate // _FP8_GROUP_SIZE), dtype=torch.float32, device=device + ) + if num_pairs > 0: + _swiglu_quant_fp8_kernel[(num_pairs, intermediate // _FP8_GROUP_SIZE)]( + ic1, + act_q, + act_scales, + intermediate, + ic1.stride(0), + act_q.stride(0), + act_scales.stride(0), + float(swiglu_limit) if swiglu_limit is not None else 0.0, + HAS_LIMIT=swiglu_limit is not None, + CLAMP_AFTER_SILU=swiglu_clamp_after_silu, + GROUP=_FP8_GROUP_SIZE, + FP8_MAX=_FP8_E4M3_MAX, + ) + + ic2 = torch.empty((num_pairs, hidden), dtype=torch.bfloat16, device=device) + _invoke_fp8act_fp8w_moe_kernel( + act_q, + act_scales, + w2, + w2_scales, + ic2, + token_final_scales.contiguous(), + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + num_m_blocks, + num_pairs, + 1, + ) + output = torch.empty((num_tokens, hidden), dtype=output_dtype or torch.bfloat16, device=device) + if num_tokens > 0: + block_h = 1024 + _moe_topk_sum_kernel[(num_tokens, triton.cdiv(hidden, block_h))]( + ic2, output, hidden, TOP_K=top_k, BLOCK_H=block_h + ) + return output diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 72b0e12823ab..b1c167e375f1 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -547,9 +547,14 @@ def get_attention_workspace_fixed_bytes(model_config, mapping, **capacities) -> int: """Capacity-sized reserve declared by the model's selected backend.""" from ..attention.backends.utils import get_attention_backend - return get_attention_backend( - model_config.attn_backend).runtime_workspace_fixed_bytes( - model_config, mapping, **capacities) + backend = get_attention_backend(model_config.attn_backend) + if config_utils.is_glm5_next(model_config.pretrained_config): + # glm5_next layers select the glm_kpool sparse backend themselves, so + # model_config carries no sparse_attention_config naming it. + from ..attention.backends.sparse.glm_kpool import \ + GlmKpoolSparseAttention as backend + return backend.runtime_workspace_fixed_bytes(model_config, mapping, + **capacities) def get_overlap_headroom_penalty_workspace_bytes( diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index d7c3a419ccc0..9ac3fd4cc3a3 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -50,6 +50,7 @@ l0_h100: - unittest/_torch/moe/test_moe_routing.py - unittest/_torch/moe/fused_moe/test_deepgemm_fused_gather_finalize.py - unittest/_torch/moe/fused_moe/test_deepgemm_fused_expand_quant.py + - unittest/_torch/moe/fused_moe/test_triton_fp8_block_scale.py # ------------- MoE: test_moe_backend (by backend) --------------- # ------------- MoE: test_single_gpu (by backend) --------------- - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "CUTLASS and not None" diff --git a/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py b/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py index e989db333636..c98dd52f4994 100644 --- a/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py +++ b/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py @@ -455,12 +455,14 @@ def test_derived_geometry_follows_the_checkpoint_defaults(): assert params.select_k == 2048 // 4 # topk positions plus the always-visible incomplete tail (kpool - 1 rows). assert params.output_width == 2048 + 3 - # Padded to the FlashMLA top-k tile. - assert params.kernel_output_width == 2112 + # Padded to the FlashMLA top-k tile (64), or to two tiles for its SM90 prefill. + assert params.kernel_output_width(64) == 2112 + assert params.kernel_output_width(128) == 2176 assert params.packed_state_dim == 2 * 128 assert params.cache_row_dim == 3 * 128 small = GlmKpoolSparseParams(index_topk=64, index_kpool=8) - assert (small.select_k, small.output_width, small.kernel_output_width) == (8, 71, 128) + assert (small.select_k, small.output_width) == (8, 71) + assert small.kernel_output_width(64) == small.kernel_output_width(128) == 128 def test_paged_slot_indices_returns_page_and_offset_pairs(): @@ -499,3 +501,126 @@ def test_latent_pool_rows_rejects_layouts_without_a_uniform_row_view(): ) with pytest.raises(ValueError, match="not multiples of dim"): latent_pool_rows(odd_stride) + + +def _planner(fp8: bool, pages: int = 4096, tokens_per_block: int = 32): + backend = object.__new__(GlmKpoolSparseAttention) + backend.sparse_params = GlmKpoolSparseParams() + backend._indexer_fp8 = fp8 + state = SimpleNamespace( + block_tables=torch.zeros(8, pages, dtype=torch.long), tokens_per_block=tokens_per_block + ) + backend._cache_state = Mock(return_value=state) + return backend + + +def _call_costs(backend, calls): + """Each planned call's FP32-element cost under the backend's cost model.""" + state = backend._cache_state(None) + kpool = backend.sparse_params.index_kpool + capacity_pools = state.block_tables.shape[1] * state.tokens_per_block // kpool + return [ + GlmKpoolSparseAttention._context_score_cost( + end - start, + last - first, + max_kv_len // kpool, + fp8=backend._indexer_fp8, + head_dim=backend.sparse_params.index_head_dim, + capacity_pools=capacity_pools, + ) + for start, end, first, last, max_kv_len in calls + ] + + +def _check_plan(backend, calls, spans, budget): + """Calls cover every row in order, name each row's table, and fit the budget.""" + request_starts = [0] + for rows, _ in spans: + request_starts.append(request_starts[-1] + rows) + assert calls[0][0] == 0 and calls[-1][1] == request_starts[-1] + for (_, end, _, _, _), (start, _, _, _, _) in zip(calls, calls[1:]): + assert end == start + for start, end, first, last, max_kv_len in calls: + assert start < end and first < last + assert request_starts[first] <= start and end <= request_starts[last] + assert max_kv_len >= max(visible for _, visible in spans[first:last]) + assert max(_call_costs(backend, calls)) <= budget + + +@pytest.mark.parametrize( + "spans", + [ + [(13324, 13324)], # one fresh long prompt + [(512, 16384 + 512)] * 4, # cached prefixes + [(16384, 131072)], # last chunk of a 128K prompt + [(3, 3)] * 64 + [(4000, 4000)], # tiny prompts beside a long one + ], + ids=["fresh", "prefixes", "long-chunk", "tiny-and-long"], +) +@pytest.mark.parametrize("fp8", [False, True], ids=["bf16", "fp8"]) +def test_context_score_plan_stays_within_the_reserved_bound(spans, fp8): + backend = _planner(fp8) + budget = GlmKpoolSparseAttention._CONTEXT_SCORE_ELEMENTS + calls = backend.plan_context_scores(object(), spans) + _check_plan(backend, calls, spans, budget) + + +def test_context_score_plan_keeps_a_fresh_prompt_in_one_call(): + calls = _planner(True).plan_context_scores(object(), [(13324, 13324)]) + assert calls == [(0, 13324, 0, 1, 13324)] + + +def test_context_score_plan_groups_short_requests_and_splits_a_long_one(): + backend = _planner(True) + grouped = backend.plan_context_scores(object(), [(100, 100)] * 3) + assert grouped == [(0, 300, 0, 3, 100)] + split = backend.plan_context_scores(object(), [(16384, 131072)]) + assert len(split) > 1 and {call[2:] for call in split} == {(0, 1, 131072)} + + +@pytest.mark.parametrize( + "capacities", + [ + dict(max_batch_size=64, max_num_tokens=64, max_seq_len=4096), + dict(max_batch_size=8, max_num_tokens=2048, max_seq_len=32768), + dict(max_batch_size=64, max_num_tokens=16384, max_seq_len=131072), + ], + ids=["small-token-budget", "mid", "large"], +) +@pytest.mark.parametrize("fp8", [False, True], ids=["bf16", "fp8"]) +def test_kv_cache_reserve_covers_every_planned_call(capacities, fp8): + tokens, seq_len = capacities["max_num_tokens"], capacities["max_seq_len"] + requests = min(capacities["max_batch_size"], tokens) + config = SimpleNamespace(pretrained_config=SimpleNamespace(model_type="glm5_next_text")) + reserve = GlmKpoolSparseAttention.runtime_workspace_fixed_bytes( + config, None, enable_cuda_graph=True, **capacities + ) + backend = _planner(fp8, pages=seq_len // 32) + batches = [ + [(1, seq_len)] * requests, # one continuation row per request + [(tokens, seq_len)], # one request takes the whole token budget + [(tokens // requests, seq_len)] * requests, # the budget split across requests + ] + for spans in batches: + calls = backend.plan_context_scores(object(), spans) + assert max(_call_costs(backend, calls)) * 4 <= reserve + + +def test_kv_cache_sizing_reserves_the_context_score_bound(): + from tensorrt_llm._torch.pyexecutor._util import get_attention_workspace_fixed_bytes + + capacities = dict( + max_batch_size=64, max_num_tokens=16384, max_seq_len=131072, enable_cuda_graph=True + ) + + def reserve(model_type): + config = SimpleNamespace( + attn_backend="TRTLLM", + sparse_attention_config=None, + pretrained_config=SimpleNamespace(model_type=model_type), + ) + return get_attention_workspace_fixed_bytes(config, None, **capacities) + + bound = GlmKpoolSparseAttention._CONTEXT_SCORE_ELEMENTS * 4 + assert reserve("glm5_next") == reserve("glm5_next_text") == bound + assert reserve("llama") == 0 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 de5148df592b..9e9cd442a292 100644 --- a/tests/unittest/_torch/attention/sparse/glm_kpool/test_kernels.py +++ b/tests/unittest/_torch/attention/sparse/glm_kpool/test_kernels.py @@ -22,7 +22,9 @@ compact_sparse_rows, kpool_expand, kpool_score, + kpool_score_fp8, kpool_update, + merge_split_attention, ) from tensorrt_llm._utils import get_sm_version @@ -320,6 +322,159 @@ def test_pool_scores_broadcast_table_single_context_request(): ) +@pytest.mark.skipif( + get_sm_version() not in (90, 100, 103), reason="DeepGEMM MQA logits require SM90 or SM100/103" +) +@pytest.mark.parametrize("single_request", [True, False], ids=["one_table", "two_tables"]) +def test_fp8_pool_scores_track_bf16_selection(single_request): + """FP8 context scoring keeps the BF16 scorer's top pools on packed context rows.""" + gen = torch.Generator(device="cuda").manual_seed(10) + heads, tpb, slots = 32, 32, 400 + lens = [6000] if single_request else [6000, 2500] + pages = [(length + tpb - 1) // tpb for length in lens] + pool = torch.randn(slots, 2, tpb, 3 * HD, generator=gen, device="cuda").to(torch.bfloat16)[:, 1] + perm = torch.randperm(slots, generator=gen, device="cuda") + tables = torch.zeros(len(lens), max(pages), dtype=torch.int64, device="cuda") + start = 0 + for i, count in enumerate(pages): + tables[i, :count] = perm[start : start + count] + start += count + request_ids = torch.cat( + [torch.full((length,), i, dtype=torch.int32) for i, length in enumerate(lens)] + ).cuda() + kv_lens = torch.cat([torch.arange(1, length + 1, dtype=torch.int32) for length in lens]).cuda() + n = kv_lens.shape[0] + q = torch.randn(n, heads, HD, generator=gen, device="cuda").to(torch.bfloat16) + w = torch.randn(n, heads, generator=gen, device="cuda").to(torch.bfloat16) + kwargs = dict( + head_dim=HD, + kpool=KPOOL, + q_scale=HD**-0.5, + w_scale=heads**-0.5, + request_ids=request_ids, + max_kv_len=max(lens), + ) + want = kpool_score( + q, + w, + pool, + tables, + kv_lens, + tpb, + num_pools_max=max(pages) * tpb // KPOOL, + precision="bf16", + rows_per_program=16, + **kwargs, + ) + got = kpool_score_fp8(q, w, pool, tables, kv_lens, tpb, **kwargs) + assert got is not None and got.stride(1) == 1 and got.stride(0) % 4 == 0 + overlaps = [] + for row in torch.randint(0, n, (64,), generator=gen, device="cuda").tolist(): + visible = int(kv_lens[row]) // KPOOL + if visible < 64: + continue + k = visible // 2 + a = set(want[row, :visible].topk(k).indices.tolist()) + b = set(got[row, :visible].topk(k).indices.tolist()) + overlaps.append(len(a & b) / k) + assert overlaps and sum(overlaps) / len(overlaps) > 0.97 and min(overlaps) > 0.9 + + +@pytest.mark.skipif( + get_sm_version() not in (90, 100, 103), reason="DeepGEMM MQA logits require SM90 or SM100/103" +) +@pytest.mark.parametrize("lens", [[16384], [16000, 12000, 8000]], ids=["one-table", "three-tables"]) +def test_fp8_pool_score_peak_memory_fits_the_planned_cost(lens): + """kpool_score_fp8's peak allocation beyond its row-sized buffers fits the planner's cost.""" + from tensorrt_llm._torch.attention.backends.sparse.glm_kpool import GlmKpoolSparseAttention + + gen = torch.Generator(device="cuda").manual_seed(13) + heads, tpb, rows_per_table = 32, 32, 256 + pages = [(length + tpb - 1) // tpb for length in lens] + pool = torch.randn(sum(pages) + 8, 2, tpb, 3 * HD, generator=gen, device="cuda") + pool = pool.to(torch.bfloat16)[:, 1] + tables = torch.zeros(len(lens), max(pages), dtype=torch.int64, device="cuda") + start = 0 + for i, count in enumerate(pages): + tables[i, :count] = torch.arange(start, start + count, device="cuda") + start += count + # The last rows_per_table queries of each prompt, so most pools are visible. + request_ids = torch.cat( + [torch.full((rows_per_table,), i, dtype=torch.int32) for i in range(len(lens))] + ).cuda() + kv_lens = torch.cat( + [ + torch.arange(length - rows_per_table + 1, length + 1, dtype=torch.int32) + for length in lens + ] + ).cuda() + n = kv_lens.shape[0] + q = torch.randn(n, heads, HD, generator=gen, device="cuda").to(torch.bfloat16) + w = torch.randn(n, heads, generator=gen, device="cuda").to(torch.bfloat16) + kwargs = dict(head_dim=HD, kpool=KPOOL, q_scale=HD**-0.5, w_scale=heads**-0.5) + args = (q, w, pool, tables, kv_lens, tpb) + kpool_score_fp8(*args, request_ids=request_ids, max_kv_len=max(lens), **kwargs) # JIT warmup + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + base = torch.cuda.memory_allocated() + scores = kpool_score_fp8(*args, request_ids=request_ids, max_kv_len=max(lens), **kwargs) + torch.cuda.synchronize() + peak = torch.cuda.max_memory_allocated() - base + del scores + cost = GlmKpoolSparseAttention._context_score_cost( + n, len(lens), max(lens) // KPOOL, fp8=True, head_dim=HD, capacity_pools=0 + ) + # Quantized queries (with the quantizer's scratch), head weights and per-row + # indices grow only with rows; startup profiling measures them. + row_bytes = n * heads * (2 * HD + 16) + 64 * n + assert peak <= cost * 4 + row_bytes + (64 << 10) + + +@pytest.mark.skipif( + get_sm_version() not in (90, 100, 103), reason="DeepGEMM MQA logits require SM90 or SM100/103" +) +def test_fp8_pool_scores_of_a_table_range_match_the_packed_call(): + """Scoring one request against its own table equals its rows of a multi-table call.""" + gen = torch.Generator(device="cuda").manual_seed(11) + heads, tpb, slots = 32, 32, 400 + lens = [3000, 2500, 1200] + pages = [(length + tpb - 1) // tpb for length in lens] + pool = torch.randn(slots, 2, tpb, 3 * HD, generator=gen, device="cuda").to(torch.bfloat16)[:, 1] + perm = torch.randperm(slots, generator=gen, device="cuda") + tables = torch.zeros(len(lens), max(pages), dtype=torch.int64, device="cuda") + start = 0 + for i, count in enumerate(pages): + tables[i, :count] = perm[start : start + count] + start += count + request_ids = torch.cat( + [torch.full((length,), i, dtype=torch.int32) for i, length in enumerate(lens)] + ).cuda() + kv_lens = torch.cat([torch.arange(1, length + 1, dtype=torch.int32) for length in lens]).cuda() + q = torch.randn(kv_lens.shape[0], heads, HD, generator=gen, device="cuda").to(torch.bfloat16) + w = torch.randn(kv_lens.shape[0], heads, generator=gen, device="cuda").to(torch.bfloat16) + kwargs = dict(head_dim=HD, kpool=KPOOL, q_scale=HD**-0.5, w_scale=heads**-0.5) + packed = kpool_score_fp8( + q, w, pool, tables, kv_lens, tpb, request_ids=request_ids, max_kv_len=max(lens), **kwargs + ) + first = lens[0] + rows = slice(first, first + lens[1]) + alone = kpool_score_fp8( + q[rows], + w[rows], + pool, + tables[1:2], + kv_lens[rows], + tpb, + request_ids=request_ids[rows] - 1, + max_kv_len=lens[1], + **kwargs, + ) + assert packed is not None and alone is not None + for row in range(0, lens[1], 97): + visible = (row + 1) // KPOOL + torch.testing.assert_close(alone[row, :visible], packed[first + row, :visible]) + + def test_pool_expansion_graph_replay_tracks_request_tables_and_tails(): # Requests occupy disjoint slots, with gaps between slots from coalesced buffers. tables = torch.tensor([[3, 1], [5, 2]], dtype=torch.int64, device="cuda") @@ -382,6 +537,8 @@ def test_fp8_sparse_core_matches_dequantized_cache_and_replays_graph(heads): backend = object.__new__(GlmKpoolSparseAttention) backend.kv_lora_rank = 512 backend.softmax_scale = 256**-0.5 + backend._kernel_topk_align = 128 if get_sm_version() == 90 else 64 + backend._split_topk = False backend.kv_scale_quant_orig = torch.tensor([0.5], device="cuda") backend._FP8_QUERY_CHUNK_SIZE = 2 # Exercise multiple chunks and a short final chunk. generator = torch.Generator(device="cuda").manual_seed(71) @@ -414,6 +571,59 @@ def reference(): torch.testing.assert_close(captured, reference(), atol=2e-2, rtol=2e-2) +@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.parametrize("num_tokens", [1, 4, 16]) +def test_split_topk_sparse_core_matches_single_call(num_tokens): + """Decode-sized calls split over top-k chunks match the single-call kernel.""" + from tensorrt_llm._torch.attention.backends.sparse.glm_kpool import GlmKpoolSparseAttention + + def make_backend(split: bool) -> GlmKpoolSparseAttention: + backend = object.__new__(GlmKpoolSparseAttention) + backend.kv_lora_rank = 512 + backend.softmax_scale = 256**-0.5 + backend._kernel_topk_align = 128 if get_sm_version() == 90 else 64 + backend._split_topk = split + return backend + + generator = torch.Generator(device="cuda").manual_seed(5) + kv_rows = (torch.randn(4096, 1, 512, device="cuda", generator=generator) * 0.5).to( + torch.bfloat16 + ) + q = (torch.randn(num_tokens, 16, 512, device="cuda", generator=generator) * 0.5).to( + torch.bfloat16 + ) + indices = torch.randint( + 0, kv_rows.shape[0], (num_tokens, 2112), device="cuda", generator=generator + ).int() + indices[:, -61:] = -1 # Padding tail, as in the expanded k-pool rows. + indices[0, : 2 * 128] = -1 # Whole chunks without a valid row. + + expected = make_backend(split=False)._dispatch_sparse_core(q, kv_rows, indices) + actual = make_backend(split=True)._dispatch_sparse_core(q, kv_rows, indices) + assert actual.shape == expected.shape == (num_tokens, 16, 512) + torch.testing.assert_close(actual, expected, atol=2e-3, rtol=2e-2) + + +def test_merge_split_attention_matches_reference(): + generator = torch.Generator(device="cuda").manual_seed(9) + num_tokens, splits, heads, kernel_heads, dim = 3, 5, 16, 64, 512 + out = torch.randn( + num_tokens * splits, kernel_heads, dim, device="cuda", generator=generator + ).to(torch.bfloat16) + lse = torch.randn(num_tokens * splits, kernel_heads, device="cuda", generator=generator) * 4 + lse[1] = float("inf") # A chunk without valid rows. + + actual = merge_split_attention(out, lse, num_tokens, splits, heads) + + lse_ref = lse.view(num_tokens, splits, kernel_heads)[..., :heads] + lse_ref = torch.where(torch.isinf(lse_ref), torch.full_like(lse_ref, float("-inf")), lse_ref) + weights = torch.softmax(lse_ref, dim=1) + parts = out.view(num_tokens, splits, kernel_heads, dim)[:, :, :heads].float() + expected = (weights[..., None] * parts).sum(dim=1) + torch.testing.assert_close(actual.float(), expected, atol=2e-2, rtol=2e-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.""" diff --git a/tests/unittest/_torch/modeling/test_glm5_next_contracts.py b/tests/unittest/_torch/modeling/test_glm5_next_contracts.py index 5763a76c2ed8..3bbdf6a1cdf6 100644 --- a/tests/unittest/_torch/modeling/test_glm5_next_contracts.py +++ b/tests/unittest/_torch/modeling/test_glm5_next_contracts.py @@ -1093,6 +1093,7 @@ def attend(x, *args, reduce=True, **kwargs): create_autospec(implementation, side_effect=attend), ) attention.forward = MethodType(Glm5NextSparseAttention.forward, attention) + attention._forward_phases = MethodType(Glm5NextSparseAttention._forward_phases, attention) connection = SimpleNamespace( pre_mapping=lambda streams: (None, None, streams.mean(dim=1)), post_mapping=lambda output, residual, post, comb: residual + output.unsqueeze(1), @@ -1181,6 +1182,7 @@ def test_packed_attention_splits_phases_and_reduces_once(generation_phase, reduc forward_verify=Mock(return_value=gen), tp_all_reduce=reduction, ) + attention._forward_phases = MethodType(Glm5NextSparseAttention._forward_phases, attention) metadata = SimpleNamespace() context = Glm5NextRuntimeContext( num_contexts=1, @@ -1227,3 +1229,43 @@ def test_indexer_topk_without_cute_dsl(monkeypatch: pytest.MonkeyPatch) -> None: expected = torch.arange(max(0, count - indexer.select_k), count, dtype=torch.int32) assert torch.equal(valid, expected) assert int((selected[row] == -1).sum()) == indexer.select_k - min(count, indexer.select_k) + + +@pytest.mark.cpu_only +def test_breakable_graph_registry_keys_are_unique(): + from tensorrt_llm._torch.models import modeling_glm5_next as glm + + model_config = SimpleNamespace(extra_attrs={}) + first, second = torch.nn.Linear(1, 1), torch.nn.Linear(1, 1) + keys = [ + glm._register_graph_layer(model_config, "glm5_sparse_layers", 3, module) + for module in (first, second) + ] + layers = model_config.extra_attrs["glm5_sparse_layers"] + assert keys == ["3", "3_0"] + assert layers["3"]() is first and layers["3_0"]() is second + + +@pytest.mark.cpu_only +def test_sparse_attention_graph_op_fills_live_rows_and_zeros_padding(monkeypatch): + from tensorrt_llm._torch.models import modeling_glm5_next as glm + from tensorrt_llm._torch.utils import model_extra_attrs + + class _Layer: + def _forward_phases(self, hidden_states, ctx): + assert ctx == "ctx" and hidden_states.shape[0] == 3 + return hidden_states * 2 + + layer = _Layer() + metadata = SimpleNamespace(num_tokens=3) + monkeypatch.setattr(glm, "build_glm5_next_runtime_context", lambda m: "ctx") + hidden = torch.arange(20, dtype=torch.float32).view(5, 4) + output = torch.full((5, 4), float("nan")) + attrs = { + "attention_metadata": lambda: metadata, + "glm5_sparse_layers": {"7": lambda: layer}, + } + with model_extra_attrs(attrs): + glm.glm5_sparse_attention_inplace(hidden, "7", output) + torch.testing.assert_close(output[:3], hidden[:3] * 2) + assert torch.equal(output[3:], torch.zeros(2, 4)) diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kimi_k3_mamba_metadata.py b/tests/unittest/_torch/modules/kimi_kda/test_kimi_k3_mamba_metadata.py index 2553fe8f23ce..6582c07e4291 100644 --- a/tests/unittest/_torch/modules/kimi_kda/test_kimi_k3_mamba_metadata.py +++ b/tests/unittest/_torch/modules/kimi_kda/test_kimi_k3_mamba_metadata.py @@ -55,6 +55,10 @@ def test_k3_prepare_metadata_match_chunk_indices() -> None: cu_seqlens = metadata.query_start_loc_long[: len(context_lengths) + 1] legacy_chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size=64) torch.testing.assert_close(metadata.kda_chunk_indices, legacy_chunk_indices) + # The host copy FLA's chunked prefill reads instead of the device cu_seqlens. + assert metadata.kda_cu_seqlens_cpu.device.type == "cpu" + torch.testing.assert_close(metadata.kda_cu_seqlens_cpu, cu_seqlens.cpu()) + torch.testing.assert_close(metadata.kda_cu_seqlens, cu_seqlens) def test_k3_prepare_materializes_stable_aligned_generation_indices() -> None: diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py b/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py index 5a06177217bb..e6487ab38d31 100644 --- a/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py +++ b/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py @@ -335,7 +335,8 @@ def test_kda_decode_matches_reference(num_heads): @pytest.mark.parametrize("batch", [1, 3]) @pytest.mark.parametrize("t_steps", [2, 3]) -def test_kda_verify_matches_sequential_decode(batch, t_steps): +@pytest.mark.parametrize("batched", [False, True], ids=["per_step", "batched"]) +def test_kda_verify_matches_sequential_decode(batch, t_steps, batched): if not torch.cuda.is_available(): pytest.skip("needs a GPU") torch.manual_seed(0) @@ -351,6 +352,11 @@ def test_kda_verify_matches_sequential_decode(batch, t_steps): # load_weights(); with random weights it holds heap garbage, and a # NaN/Inf bit pattern poisons both paths identically (nvbug 6599150). torch.nn.init.normal_(runtime.dt_bias, std=0.1) + if batched: + # The packed convolution weight selects the one-launch verify path. + runtime._build_mtp_conv_weights() + else: + runtime._packed_conv_weight = None slots = batch + 2 # non-trivial slot mapping cache = _LayerCache(slots, 3 * dim, w - 1, h, lin["head_dim"], lin["head_dim"], t_steps, device) slot_indices = torch.arange(2, 2 + batch, device=device, dtype=torch.long) diff --git a/tests/unittest/_torch/modules/test_mhc.py b/tests/unittest/_torch/modules/test_mhc.py index 285495103a1f..fd5f451f6861 100644 --- a/tests/unittest/_torch/modules/test_mhc.py +++ b/tests/unittest/_torch/modules/test_mhc.py @@ -600,6 +600,35 @@ def test_mhc_fused_hc_single_wave_target_mma_ks(): assert ((m + 63) // 64) * ks <= sm_count * 5, (hidden_size, m, ks) +def test_mhc_fused_hc_split_dg_tactics_without_mma(monkeypatch): + """Without tcgen05, M > 32 also offers the unfused DeepGEMM boundary.""" + from tensorrt_llm._torch.autotuner import OptimizationProfile + from tensorrt_llm._torch.modules.mhc import mhc_cuda + + monkeypatch.setattr(mhc_cuda, "_fused_hc_mma_supported", lambda: False) + monkeypatch.setattr(mhc_cuda, "_get_dg_fn", lambda: object()) + runner = mhc_cuda.MhcFusedHcRunner( + n=4, + hidden_size=4096, + rms_eps=1e-6, + hc_pre_eps=1e-6, + hc_sinkhorn_eps=1e-6, + hc_post_mult_value=2.0, + sinkhorn_repeat=20, + ) + + def backends(num_tokens, **kwargs): + inputs = [torch.empty((num_tokens, 4096), dtype=torch.bfloat16)] + return {t[0] for t in runner.get_valid_tactics(inputs, OptimizationProfile(), **kwargs)} + + assert "split_dg" not in backends(32) + assert "split_dg" in backends(64) + assert "split_dg" in backends(64, pre_mix_ext=torch.empty(64, 4)) + # Token counts above the tuned buckets fall back to split_dg as well. + assert mhc_cuda._get_fused_hc_fallback_tactic(4096, 32)[0] == "fused_half_fma" + assert mhc_cuda._get_fused_hc_fallback_tactic(4096, 13312)[0] == "split_dg" + + def test_mhc_fused_hc_uneven_splits_unreachable_off_sm107(): """The uneven split counts must stay unselectable on non-SM107 devices.""" from tensorrt_llm._torch.modules.mhc.mhc_cuda import _fused_hc_target_mma_ks @@ -864,15 +893,19 @@ def make_runner_inputs(): ("fused_half_fma", 2, 1, 256, 1), pytest.param(("fused_all_mma", 0, 1, 0, 1), marks=skip_pre_blackwell), ("fused_all_fma", 2, 1, 0, 1), + ("split_dg", 0, 1, 256, 1), + ("split_dg", 0, 16, 256, 1), ], ) @pytest.mark.parametrize("hidden_size", [4096, 7168]) def test_mhc_fused_hc_realistic_scale_regression(tactic, hidden_size: int): """Real-scale mHC data catches fused_hc RMS normalization regressions.""" - from tensorrt_llm._torch.modules.mhc.mhc_cuda import MhcFusedHcRunner + from tensorrt_llm._torch.modules.mhc.mhc_cuda import MhcFusedHcRunner, _get_dg_fn if tactic[0].endswith("_mma") and not _mhc_fused_hc_mma_available(): pytest.skip("mHC fused-HC MMA kernels require SM100 and BUILD_DEEP_GEMM=ON") + if tactic[0] == "split_dg" and _get_dg_fn() is None: + pytest.skip("split_dg needs DeepGEMM's TF32 hc_prenorm GEMM") n = 16 hc_mult = 4 @@ -934,6 +967,84 @@ def test_mhc_fused_hc_realistic_scale_regression(tactic, hidden_size: int): torch.testing.assert_close(layer_input_ref, layer_input_cur, rtol=1e-2, atol=2e-2) +@pytest.mark.parametrize("num_splits", [1, 16]) +@pytest.mark.parametrize("num_tokens", [64, 1024]) +def test_mhc_split_dg_lagged_matches_fma(num_tokens: int, num_splits: int): + """With a lagged pre-mix, the DeepGEMM boundaries match the FMA ones.""" + from tensorrt_llm._torch.modules.mhc.mhc_cuda import ( + MhcFusedHcRunner, + MhcPreMappingRunner, + _get_dg_fn, + ) + + if _get_dg_fn() is None: + pytest.skip("split_dg needs DeepGEMM's TF32 hc_prenorm GEMM") + + hidden_size, hc_mult = 4096, 4 + pre_data = generate_realistic_pre_data(n=num_tokens, hc_mult=hc_mult, hidden_size=hidden_size) + consts = dict( + n=hc_mult, + hidden_size=hidden_size, + rms_eps=pre_data["rms_eps"], + hc_pre_eps=pre_data["hc_pre_eps"], + hc_sinkhorn_eps=pre_data["hc_sinkhorn_eps"], + hc_post_mult_value=pre_data["hc_post_mult_value"], + sinkhorn_repeat=pre_data["sinkhorn_repeat"], + ) + torch.random.manual_seed(23) + device = "cuda" + x_prev = torch.randn((num_tokens, hidden_size), device=device).bfloat16() + residual_prev = torch.randn((num_tokens, hc_mult, hidden_size), device=device).bfloat16() + post_mix_prev = torch.randn((num_tokens, hc_mult), dtype=torch.float32, device=device) * 0.1 + comb_mix_prev = ( + torch.randn((num_tokens, hc_mult, hc_mult), dtype=torch.float32, device=device) * 0.1 + ) + pre_mix = torch.rand((num_tokens, hc_mult), dtype=torch.float32, device=device) + w_t = pre_data["fn"].reshape(-1, hc_mult * hidden_size).contiguous() + norm_weight = (torch.rand(hidden_size, device=device) + 0.5).bfloat16() + kwargs = dict(pre_mix_ext=pre_mix, norm_weight=norm_weight, norm_eps=1e-6) + names = ("post_mix", "comb_mix", "layer_input", "pre_mix_out") + tols = ( + dict(rtol=3e-3, atol=5e-3), + dict(rtol=3e-3, atol=5e-3), + dict(rtol=1e-2, atol=2e-2), + dict(rtol=3e-3, atol=5e-3), + ) + + fused = MhcFusedHcRunner(**consts) + inputs = [ + x_prev, + residual_prev, + post_mix_prev, + comb_mix_prev, + w_t, + pre_data["hc_scale"], + pre_data["hc_base"], + ] + ref = fused(inputs=inputs, tactic=("fused_half_fma", 2, 1, 256, 1), **kwargs) + out = fused(inputs=inputs, tactic=("split_dg", 0, num_splits, 256, 1), **kwargs) + assert len(out) == len(ref) == 5 + torch.testing.assert_close(out[0], ref[0], rtol=1e-2, atol=2e-2) + for name, tol, o, r in zip(names, tols, out[1:], ref[1:]): + torch.testing.assert_close(o, r, **tol, msg=f"fused_hc {name} mismatch") + + pre_mapping = MhcPreMappingRunner(**consts) + residual = ref[0] + pre_inputs = [ + residual.view(num_tokens, hc_mult * hidden_size), + w_t, + residual, + pre_data["hc_scale"], + pre_data["hc_base"], + ] + dg_backend = "dg_splitk" if num_splits > 1 else "dg_nosplit" + ref = pre_mapping(inputs=pre_inputs, tactic=("fma", 8, 1, 256), **kwargs) + out = pre_mapping(inputs=pre_inputs, tactic=(dg_backend, 0, 0, 256), **kwargs) + assert len(out) == len(ref) == 4 + for name, tol, o, r in zip(names, tols, out, ref): + torch.testing.assert_close(o, r, **tol, msg=f"pre_mapping {name} mismatch") + + @pytest.mark.parametrize( ("num_tokens", "hidden_size", "num_k_splits"), [ diff --git a/tests/unittest/_torch/moe/fused_moe/test_triton_fp8_block_scale.py b/tests/unittest/_torch/moe/fused_moe/test_triton_fp8_block_scale.py index 9521c34308db..83969f053168 100644 --- a/tests/unittest/_torch/moe/fused_moe/test_triton_fp8_block_scale.py +++ b/tests/unittest/_torch/moe/fused_moe/test_triton_fp8_block_scale.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """ -Unit tests for ``run_triton_fp8_block_scale_moe``. +Unit tests for ``run_triton_fp8_block_scale_moe`` and ``run_triton_fp8_block_scale_moe_fp8_act``. Run as a pytest: pytest tests/unittest/_torch/moe/fused_moe/test_triton_fp8_block_scale.py -v @@ -17,7 +17,10 @@ from _torch.helpers import calc_diff, per_block_cast_to_fp8_e8m0 from tensorrt_llm._torch.moe.fused_moe.fused_moe_triton_fp8_block_scale import ( + _moe_align_small, + moe_align_block_size, run_triton_fp8_block_scale_moe, + run_triton_fp8_block_scale_moe_fp8_act, ) from tensorrt_llm._torch.moe.fused_moe.interface import ActivationType from tensorrt_llm._utils import get_sm_version @@ -201,3 +204,112 @@ def test_fp8_matches_bf16_reference(shape: MoeShape, activation_type: int) -> No diff = calc_diff(out_fp8, out_ref).item() assert diff < 1e-2, f"shape={shape.name} cosine-diff={diff:.4e} exceeds 1e-2" + + +skip_pre_hopper = pytest.mark.skipif( + not torch.cuda.is_available() or get_sm_version() < 90, + reason="FP8 Triton dot needs SM90 or newer", +) + + +def _dequant_blocks(w_fp8: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + num_experts, out_features, in_features = w_fp8.shape + expanded = scales.repeat_interleave(128, dim=1).repeat_interleave(128, dim=2) + return w_fp8.float() * expanded[:, :out_features, :in_features] + + +def _fake_quant_1x128(v: torch.Tensor) -> torch.Tensor: + """Round a 1-D FP32 vector through FP8 with one scale per 128 elements.""" + groups = v.view(-1, 128) + scale = groups.abs().amax(dim=1, keepdim=True).clamp(min=1e-10) / 448.0 + return ((groups / scale).to(torch.float8_e4m3fn).float() * scale).view(-1) + + +def _run_moe_fp8_act_reference( + x, + token_selected_experts, + token_final_scales, + w3_w1, + w3_w1_scales, + w2, + w2_scales, + swiglu_limit, + swiglu_clamp_after_silu, +) -> torch.Tensor: + """FP32 reference over dequantized weights and 1 x 128 FP8-rounded activations, + with the CUTLASS SwiGLU clamp.""" + intermediate = w3_w1.shape[1] // 2 + w31 = _dequant_blocks(w3_w1, w3_w1_scales) + w2f = _dequant_blocks(w2, w2_scales) + out = torch.zeros(x.shape, dtype=torch.float32, device=x.device) + for t in range(x.shape[0]): + x_q = _fake_quant_1x128(x[t].float()) + for k in range(token_selected_experts.shape[1]): + e = int(token_selected_experts[t, k]) + gate_up = (w31[e] @ x_q).bfloat16().float() + up, gate = gate_up[:intermediate], gate_up[intermediate:] + if swiglu_limit is not None: + up = up.clamp(-swiglu_limit, swiglu_limit) + if swiglu_clamp_after_silu: + act = F.silu(gate).clamp(max=swiglu_limit) + else: + act = F.silu(gate.clamp(max=swiglu_limit)) + else: + act = F.silu(gate) + out[t] += (w2f[e] @ _fake_quant_1x128(act * up)).bfloat16().float() * float( + token_final_scales[t, k] + ) + return out.bfloat16() + + +@skip_pre_hopper +@pytest.mark.parametrize("num_experts", [8, 72]) +@pytest.mark.parametrize("single_program", [False, True], ids=["multi_kernel", "single_program"]) +def test_moe_align_non_power_of_two_experts(num_experts: int, single_program: bool) -> None: + g = torch.Generator(device="cuda").manual_seed(3) + ids = torch.randint(0, num_experts, (5, 4), device="cuda", generator=g) + if single_program: + sorted_ids, expert_ids, num_post_pad = _moe_align_small(ids, 16, num_experts, ids.numel()) + else: + sorted_ids, expert_ids, num_post_pad = moe_align_block_size(ids, 16, num_experts) + counts = torch.bincount(ids.flatten(), minlength=num_experts) + padded = (counts + 15) // 16 * 16 + assert int(num_post_pad) == int(padded.sum()) + flat = ids.flatten() + for block in range(int(num_post_pad) // 16): + rows = sorted_ids[block * 16 : (block + 1) * 16] + valid = rows[rows < flat.numel()] + assert valid.numel() > 0 + assert bool((flat[valid] == expert_ids[block]).all()) + # Every pair lands in exactly one slot. + slots = sorted_ids[: int(num_post_pad)] + assert sorted(slots[slots < flat.numel()].tolist()) == list(range(flat.numel())) + + +@skip_pre_hopper +# 24 tokens x top-8 exceeds the single-program alignment and takes the multi-kernel one. +@pytest.mark.parametrize("num_tokens", [1, 4, 16, 24]) +@pytest.mark.parametrize( + "swiglu_limit, swiglu_clamp_after_silu", + [(None, False), (0.5, False), (0.5, True)], + ids=["no_clamp", "clamp", "clamp_after_silu"], +) +def test_fp8_act_matches_reference(num_tokens, swiglu_limit, swiglu_clamp_after_silu) -> None: + # A non-power-of-two expert count, as in decode-time routing of large MoEs. + shape = MoeShape( + "decode", num_tokens=num_tokens, hidden=512, intermediate=256, num_experts=72, top_k=8 + ) + inputs = _make_inputs(shape, seed=num_tokens) + del inputs["activation_type"] + + out = run_triton_fp8_block_scale_moe_fp8_act( + **inputs, swiglu_limit=swiglu_limit, swiglu_clamp_after_silu=swiglu_clamp_after_silu + ) + ref = _run_moe_fp8_act_reference( + **inputs, swiglu_limit=swiglu_limit, swiglu_clamp_after_silu=swiglu_clamp_after_silu + ) + + assert out.shape == ref.shape == (num_tokens, shape.hidden) + assert out.dtype == torch.bfloat16 + diff = calc_diff(out, ref).item() + assert diff < 1e-4, f"cosine-diff={diff:.4e}"